Sympy in combination with astropy to solve numerically with units?

user_5181181 0
Viewed 148

Is there a way to combine Sympy's solver with Astropy? For example, if I'd like to solve Kepler's third law for p for given values of all other variables:

from sympy import solveset, Poly, Eq, Function, exp
from sympy import Symbol, pi
Ms = Symbol('Ms') # mass of sun
Mp = Symbol('Mp') # mass of planet
G = Symbol('G') # Gravitational consatant
a = Symbol('a') # semi jamor axes
P = Symbol('P') # period
solveset(G * (Ms + Mp) / 4 * pi**2 - a**3/P**2, a)

How could I then use astropy to add values with corresponding units to all the variables and get a numerical output?

import astropy.units as u
from astropy.constants import G
Ms = 1 * u.M_sun
Mp = 1 * u.M_jup
s = 5 * u.pc
P = 1 * u.yr

Can this be combined with the results received by Sympy to solve for a value in a with the expected unit and numerical result?

Any suggestions are appreciated.

astropy solver sympy python
1 Answers

I've been playing around with using lambdify to be able to implement using astropy units in Sympy equations. Here's what I've come up with as a solution for the question:

from sympy import pi, Symbol, lambdify, sympify, symbols, solve

import astropy.units as u
from astropy.constants import G as astroG

Ms = Symbol('Ms') # mass of sun
Mp = Symbol('Mp') # mass of planet
G = Symbol('G') # Gravitational consatant
a = Symbol('a') # semi major axes
P = Symbol('P') # period
solutions = solve(
    (G * (Ms + Mp) / (4 * pi**2)) - (a**3/P**2),
    a
)

# set values at which to evaluate the solutions
values = {
    "P": 2.1 * u.yr,
    "Ms": 1.4 * u.M_sun,
    "Mp": 0.5 * u.M_jup, # (0.5 * u.M_jup.to("M_sun")),
    "G": astroG,
}

evalsolutions = []

reqsymbols = [Ms, Mp, P, G]

for sol in solutions:
    # lambdify solutions
    f = lambdify(
        reqsymbols,
        sol,
        modules=["numpy"],
    )
    
    res = f(**values)
    if res.dtype != "complex":
        evalsolutions.append(res.si.decompose())

print(evalsolutions)

which gives:

[<Quantity 2.74467861e+11 m>]

Old solution

Here is another solution that I posted originally. It's more complicated and turns out to be overkill, but does show how you can make use of the modules argument of the lambdify-ed function.

from sympy import pi, Symbol, lambdify, sympify, symbols, solve

import astropy.units as u
from astropy.constants import G as astroG

Ms = Symbol('Ms') # mass of sun
Mp = Symbol('Mp') # mass of planet
Munit = sympify('unit(msol)')
G = sympify('const(G)') # Gravitational consatant
a = Symbol('a') # semi major axes
P = Symbol('P') # period
Punit = Symbol('unit(yr)')
solutions = solve(
    (G * ((Ms + Mp) * Munit) / (4 * pi**2)) - (a**3/(P * Punit)**2),
    a
)

UNITS = {
    "msol": u.M_sun,
    "au": u.au,
    "yr": u.yr,
}

CONSTANTS = {
    "G": astroG,
}

def unitfunc(key):
    return 1.0 * UNITS.get(key, 1)

def constfunc(key):
    return CONSTANTS.get(key, 1)

# set values at which to evaluate the solutions
values = {
    "P": 2.1,
    "Ms": 1.4,
    "Mp": (0.5 * u.M_jup.to("M_sun")),
}

values.update({key: key for key in UNITS})
values.update({key: key for key in CONSTANTS})

evalsolutions = []

reqsymbols = [Ms, Mp, P]
reqsymbols.extend([symbols(key) for key in list(UNITS.keys()) + list(CONSTANTS.keys())])

for sol in solutions:
    # lambdify solutions
    f = lambdify(
        reqsymbols,
        sol,
        modules=[{"unit": unitfunc, "const": constfunc}, "numpy"],
    )
    
    res = f(**values)
    if res.dtype != "complex":
        evalsolutions.append(res.si.decompose())

print(evalsolutions)
user_18628610
Related