xtuner/tools/plugins/solve.py

# Copyright (c) OpenMMLab. All rights reserved.
import math
import re
from math import *  # noqa: F401, F403

from sympy import Eq, solve, symbols

from .calculate import Calculate


def Solve(equations_str):
    try:
        equations_str = equations_str.replace(' ', '')
        equations_ori = re.split(r'[,;]+', equations_str)
        equations_str = equations_str.replace('^', '**')
        equations_str = re.sub(r'(\(.*\))([a-zA-Z])', r'\1 * \2',
                               equations_str)
        equations_str = re.sub(r'(\d+)([a-zA-Z])', r'\1 * \2', equations_str)
        equations_str = equations_str.replace('pi', str(math.pi))
        equations = re.split(r'[,;]+', equations_str)
        vars_list = list(set(re.findall(r'[a-zA-Z]+', equations_str)))
        vars = {var: symbols(var) for var in vars_list}

        output = ''
        eqs = []
        for eq in equations:
            if '=' in eq:
                left, right = eq.split('=')
                eqs.append(
                    Eq(
                        eval(left.strip(), {}, vars),
                        eval(right.strip(), {}, vars)))
        solutions = solve(eqs, vars, dict=True)

        vars_values = {var: [] for var in vars_list}
        if isinstance(solutions, list):
            for idx, solution in enumerate(solutions):
                for var, sol in solution.items():
                    output += f'{var}_{idx} = {sol}\n'
                    vars_values[str(var)].append(sol)
        else:
            for var, sol in solutions.items():
                output += f'{var} = {sol}\n'
                vars_values[str(var)].append(sol)
        for eq, eq_o in zip(equations, equations_ori):
            if '=' not in eq:
                for var in vars_list:
                    need_note = True if len(vars_values[var]) > 1 else False
                    for idx, value in enumerate(vars_values[var]):
                        eq_to_calc = eq.replace(var, str(value))
                        calc_result = Calculate(eq_to_calc)
                        if need_note:
                            eq_name = eq_o.replace(var, f'{var}_{idx}')
                        else:
                            eq_name = eq_o
                        if calc_result != 'No results.':
                            output += f'{eq_name} = {calc_result}\n'

        return output.strip()
    except Exception:
        return 'No result.'