how to use sympy solver to solve equation generated by corresponding elements from two numpy.array?

Question

As I know, if the parameters in the equation are scalars, solving the equation above is implemented as:

from sympy.solvers import solve
from sympy import Symbol
from math import sin, cos, sqrt

# ep, t, a are in short for epsilon, theta and alpha, respectively
ep = Symbol('ep')
a = 0.4798212489019141 
t = 39.603733
solve(abs(a)-abs((ep-1)(sin(t)**2 - ep*(1+sin(t)**2)/(ep*cos(t)+sqrt(ep-sin(t)**2))**2)), ep)

The question is, if these parameters are numpy.ndarray in 3 dimensions, how to implement this calculation? In my case, the shape of theta and alpha is (5,277,232). So the output, ep, should also have the same shape.

Answer 1

When I pick the problem expression apart, the minimal segment that gives your error is:

In [13]: sqrt(ep-sin(t)**2)                                                                    
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-13-1d595e0b7cc7> in <module>
----> 1 sqrt(ep-sin(t)**2)

/usr/local/lib/python3.6/dist-packages/sympy/core/expr.py in __float__(self)
    323         if result.is_number and result.as_real_imag()[1]:
    324             raise TypeError("can't convert complex to float")
--> 325         raise TypeError("can't convert expression to float")
    326 
    327     def __complex__(self):

TypeError: can't convert expression to float

The argument to sqrt is a sympy expression (because ep is a symbol):

In [14]: ep-sin(t)**2                                                                          
Out[14]: ep - 0.892639513815489

math.sqrt requires a float input; it can't do symbolic sqrt.

If I replace it with the sympy sqrt:

In [16]: from sympy import sqrt                                                                

In [17]: sqrt(ep-sin(t)**2)                                                                    
Out[17]: 
  ________________________
╲╱ ep - 0.892639513815489 

Now if I try the whole expression:

In [18]: abs(a)-abs((ep-1)(sin(t)**2 - ep*(1+sin(t)**2)/(ep*cos(t)+sqrt(ep-sin(t)**2))**2))    
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-18-7d1c0e01daaf> in <module>
----> 1 abs(a)-abs((ep-1)(sin(t)**2 - ep*(1+sin(t)**2)/(ep*cos(t)+sqrt(ep-sin(t)**2))**2))

TypeError: 'Add' object is not callable

(ep-1) is a sympy Add expression. Due to python syntax rules you can't do (...)(...) implying multiplication. You have to be explicit:

In [19]: abs(a)-abs((ep-1)*(sin(t)**2 - ep*(1+sin(t)**2)/(ep*cos(t)+sqrt(ep-sin(t)**2))**2))   
Out[19]: 
                    │         ⎛                1.89263951381549⋅ep                           
0.479821248901914 - │(ep - 1)⋅⎜──────────────────────────────────────────────────── - 0.89263
                    │         ⎜                                                   2          
                    │         ⎜⎛                         ________________________⎞           
                    │         ⎝⎝0.327659100567207⋅ep - ╲╱ ep - 0.892639513815489 ⎠           

          ⎞│
9513815489⎟│
          ⎟│
          ⎟│
          ⎠│

In [20]: solve(abs(a)-abs((ep-1)*(sin(t)**2 - ep*(1+sin(t)**2)/(ep*cos(t)+sqrt(ep-sin(t)**2))**
    ...: 2)), ep)                                                                              
---------------------------------------------------------------------------
NotImplementedError                       Traceback (most recent call last)
<ipython-input-20-944cce37e63b> in <module>
----> 1 solve(abs(a)-abs((ep-1)*(sin(t)**2 - ep*(1+sin(t)**2)/(ep*cos(t)+sqrt(ep-sin(t)**2))**2)), ep)

/usr/local/lib/python3.6/dist-packages/sympy/solvers/solvers.py in solve(f, *symbols, **flags)
   1032             if e.has(*symbols):
   1033                 raise NotImplementedError('solving %s when the argument '
-> 1034                     'is not real or imaginary.' % e)
   1035 
   1036         # arg

NotImplementedError: solving Abs((ep - 1)*(1.89263951381549*ep/(0.327659100567207*ep - sqrt(ep - 0.892639513815489))**2 - 0.892639513815489)) when the argument is not real or imaginary.

Looks like you need to assign some constraints to the variable ep. (I'm not enough of an sympy expert to do that without reading the docs.)

===

If I remove the abs from the expression, and refine ep, I get a numeric solution:

In [57]: expr =a-(ep-1)*(sin(t)**2 - ep*(1+sin(t)**2)/(ep*cos(t)+sqrt(ep-sin(t)**2))**2)       

In [58]: ep =  Symbol('ep', real=True)                                                         

In [59]: solve(expr, ep)                                                                       
Out[59]: [56.1511793662046]

In [60]: expr.subs(ep, _[0])                                                                   
Out[60]: -1.65423230669148e-13

presumably that could be repeated for other values of a and t. The iteration method shouldn't matter, since the solve step will dominate that time.

Answer 2

If you want to solve the above equation for pairs of a and t, you can vectorize your function. This allows you put in arrays with arbitrary shapes and get the function evaluated for all pairs and sorted into the right shape.

Example:

import numpy as np

def eps_solve(a, t):
    return a + t
eps_solve_vec = np.vectorize(eps_solve)

a = np.array([[[1, 1], [2, 2]], [[3, 4], [7, 8]]])
t = np.array([[[1, 2], [3, 4]], [[7, 0], [3, 1]]])
res = eps_solve_vec(a, t)
print(res, res.shape)