Skip to content

Instantly share code, notes, and snippets.

@BenPortner
Last active January 20, 2023 12:25
Show Gist options
  • Save BenPortner/4b6032e58de170ba56f6d7f947cf7e66 to your computer and use it in GitHub Desktop.
Save BenPortner/4b6032e58de170ba56f6d7f947cf7e66 to your computer and use it in GitHub Desktop.
Solve equation systems including units with pint and sympy
"""
Note: the most recent version of pint contains [a bug](https://github.com/hgrecco/pint/pull/1701).
Install [this version of pint](https://github.com/BenPortner/pint/tree/fix-subtraction-error) to
run below code.
"""
from pint.facets.plain.registry import build_eval_tree, tokenizer
from pint.util import string_preprocessor
from pint import UnitRegistry
from sympy.parsing.sympy_parser import parse_expr
from sympy import Symbol, solve
# define system of equations (Note: 0 = 1 m - A instead of A = 1 m)
equations = [
"1 m - A",
"A + 200 mm - B",
"B * kg/m - C",
"A * B * C - D"
]
# define unknowns
unknowns = ["A", "B", "C", "D"]
# parse equations
# note: parsing first with pint to convert 200 mm -> 200 * mm etc.
parsed = [
parse_expr(build_eval_tree(tokenizer(string_preprocessor(eq))).to_string())
for eq in equations
]
# solve equation system
solution = solve(parsed, [Symbol(s) for s in unknowns])
# parse units in solution
ureg = UnitRegistry()
solution = [ureg(str(r)).to_base_units() for r in solution[0]]
# for reference
expected = [
ureg("1 m"), # A
ureg("1.2 m"), # B
ureg("1.2 kg"), # C
ureg("1.44 kg * m^2"), # D
]
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment