solve Sympy functions in Scipy (optimize.root)

Question

I'm trying to find the roots with the root function in Scipy for functions/variables created in Sympy.

sympy:

AllEquations = [x1 - 5, y1 - 5, z1 - 5, ((x1 - x2)**2 + (y1 - y2)**2 + (z1 - z2)**2)**0.5 - 150, y2 - 100, z2, ((x1 - x2)**2 + (y1 - y2)**2 + (z1 - z2)**2)**0.5 - 150]

AllVariables = [x1, y1, z1, x2, y2, z2]

Below works in Scipy, but is done manually:

def equation(p):
    x1, y1, z1, x2, y2, z2 = p
    return x1 - 5, y1 - 5, z1 - 5, ((x1 - x2)**2 + (y1 - y2)**2 + (z1 - z2)**2)**0.5 - 150, y2 - 100, z2, ((x1 - x2)**2 + (y1 - y2)**2 + (z1 - z2)**2)**0.5 - 150

initial_guess = [5,5,5,100,100,0]

from scipy.optimize import root
sol = root(equation, initial_guess, method = 'lm")

>>>sol.x
array([ 5.00000000e+00,  5.00000000e+00,  5.00000000e+00,  1.20974135e+02,
1.00000000e+02, -2.00332805e-25])

Question is now, how can I automate this? The equation function above is static and should be made dynamically as the number of parameters and functions will change continuously.

Answer 1

how can I automate this?

The concept is ill-defined. To a certain extent, you can't fully automate this, because you need sane initial estimates for a well-behaved root-find. There is not a one-size-fits-all. If you write a sympy-to-scipy wrapper, you still need to be able to adjust parameters like the method and tolerance.

Since you're in Sympy, there's a good chance that it will be able to calculate a Jacobian for you, and you should definitely do that.

Don't call lambdify once per equation; just call it once on a vector (Matrix) of all of your equations.

import typing

import numpy as np
import scipy.optimize
import sympy
from scipy.optimize import root, check_grad


def numerical_roots(
    symbolic_eqs: typing.Sequence[sympy.Eq],
    x0: dict[str, float],
    method: typing.Literal[
        'hybr', 'lm', 'broyden1', 'broyden2', 'anderson', 'linearmixing', 'diagbroyden',
        'excitingmixing', 'krylov', 'df-sane',
    ] = 'hybr',
    tol: float | None = None,
    jac_tol: float | None = 1e-5,
) -> tuple[scipy.optimize.OptimizeResult, dict[str, float]]:
    matrix = sympy.Matrix([
        eq.args[0] - eq.args[1]
        for eq in symbolic_eqs
    ])
    syms = tuple(matrix.free_symbols)
    jacmat = matrix.jacobian(syms)
    eq_unpacked = sympy.lambdify(args=syms, expr=matrix)
    jac_unpacked = sympy.lambdify(args=syms, expr=jacmat)
    numeric_x0 = [x0[s.name] for s in syms]

    def eq_packed(x: np.ndarray) -> np.ndarray:
        return eq_unpacked(*x).ravel()
    def jac_packed(x: np.ndarray) -> np.ndarray:
        return jac_unpacked(*x)

    if jac_tol is not None:
        jerr = check_grad(func=eq_packed, grad=jac_packed, x0=numeric_x0)
        if jerr > jac_tol:
            raise ValueError('Jacobian is incorrect')

    result = root(
        fun=eq_packed, jac=jac_packed,
        x0=numeric_x0, method=method, tol=tol,
    )
    if not result.success:
        raise ValueError(result.message)
    return result, {
        s.name: x
        for s, x in zip(syms, result.x)
    }


def test() -> None:
    a = {'real': True, 'finite': True}
    x1 = sympy.Symbol(name='x1', **a)
    y1 = sympy.Symbol(name='y1', **a)
    z1 = sympy.Symbol(name='z1', **a)
    x2 = sympy.Symbol(name='x2', **a)
    y2 = sympy.Symbol(name='y2', **a)
    z2 = sympy.Symbol(name='z2', **a)
    equations = (
        sympy.Eq(5, x1),
        sympy.Eq(5, y1),
        sympy.Eq(5, z1),
        sympy.Eq(150, sympy.sqrt((x1 - x2)**2 + (y1 - y2)**2 + (z1 - z2)**2)),
        sympy.Eq(100, y2),
        sympy.Eq(0, z2),
    )
    for method in (
        'hybr', 'lm', 'broyden2', 'anderson', 'krylov',
    ):
        result, sol = numerical_roots(
            symbolic_eqs=equations, method=method,
            x0={'x1': 1, 'y1': 1, 'z1': 1, 'x2': 1, 'y2': -1, 'z2': -1},
        )
        assert np.isclose(sol['x1'], 5, rtol=0, atol=1e-5)
        assert np.isclose(sol['y1'], 5, rtol=0, atol=1e-5)
        assert np.isclose(sol['z1'], 5, rtol=0, atol=1e-5)
        if sol['x2'] > 0:
            assert np.isclose(sol['x2'], 120.97413504743209, rtol=0, atol=1e-5)
        else:
            assert np.isclose(sol['x2'], -110.97413504743206, rtol=0, atol=1e-5)
        assert np.isclose(sol['y2'], 100, rtol=0, atol=1e-5)
        assert np.isclose(sol['z2'], 0, rtol=0, atol=1e-5)
        print(f'method={method} nfev={result.nfev} njev={result.get("njev")} status={result.status}')


if __name__ == '__main__':
    test()

method=hybr nfev=13 njev=1 status=1

method=lm nfev=12 njev=8 status=2

C:\Program Files\Python312\Lib\site-packages\scipy\optimize\_root.py:255: RuntimeWarning: Method broyden2 does not use the jacobian (jac).
  _warn_jac_unused(jac, method)
method=broyden2 nfev=100 njev=None status=1

C:\Program Files\Python312\Lib\site-packages\scipy\optimize\_root.py:255: RuntimeWarning: Method anderson does not use the jacobian (jac).
  _warn_jac_unused(jac, method)
method=anderson nfev=344 njev=None status=1

C:\Program Files\Python312\Lib\site-packages\scipy\optimize\_root.py:255: RuntimeWarning: Method krylov does not use the jacobian (jac).
  _warn_jac_unused(jac, method)
method=krylov nfev=418 njev=None status=1

Answer 2

This can be done using sympy's lambdify function to turn the equation into a Python function that can be evaluated numerically.

import sympy as smp
from scipy.optimize import root

x1, x2, y1, y2, z1, z2 = smp.symbols("x1 x2 y1 y2 z1 z2")
variables = [x1, x2, y1, y2, z1, z2]

equations = [x1 - 5,
             y1 - 5,
             z1 - 5,
             ((x1 - x2)**2 + (y1 - y2)**2 + (z1 - z2)**2)**0.5 - 150,
             y2 - 100,
             z2,    
    ]

# this list will contain all the functions you want to find the roots of
equations_lamb = [smp.lambdify(variables, eq) for eq in equations]

def combined_function(x, function_list):
    # evaluates the provided list of functions
    res = []
    for function in function_list:
        res.append(function(*x))
    return res

initial_guess = [0]*len(variables)
sol = root(combined_function, initial_guess, args=(equations_lamb,))
print(sol)