modeq.py

from dataclasses import dataclass

import sympy


@dataclass
class ModifiableEquation:
    """Track the state of a system as one value is modified

    The object is created with a sympy equality equation and the
    inital state of all or all but one of the variables.

    The state of the system can repeatedly be updated with one free
    variable and one or more updated variables. Variables that are
    not updated are held at their prior value.

    This works for any kind of `Eq` equation that is amenable to
    exact solution via `sympy.solve`.
    """

    expr: sympy.Eq
    state: dict[sympy.Symbol, float]

    def __post_init__(self):
        print("in postint", self)
        sub = self.expr.subs(self.state)
        free = sub.free_symbols
        if len(free) > 1:
            raise ValueError("Underspecified")
        elif free:
            freevar = free.pop()
            self.update_and_solve(freevar)
        else:
            assert sub  # i.e., that the result of the substitution was BooleanTrue

    def update_and_solve(
        self, freevar: sympy.Symbol, updated: dict[sympy.Symbol, float] | None = None
    ):
        state = dict(self.state)
        if updated:
            state.update(updated)
        if freevar in state:
            del state[freevar]
        sub = self.expr.subs(state)
        print()
        print(f"solving for {sub} [after applying {state}]")
        result = sympy.solve(sub)[0]
        print(f"got {freevar}={result}")
        state[freevar] = result
        self.state = state
        return result


if __name__ == "__main__":
    print("""# Example use of ModifiableEquation""")
    x, y, z = sympy.symbols("x y z")
    e = ModifiableEquation(expr=sympy.Eq(x + 2 * y + z, 27), state={x: 9, y: 7})
    print(e)
    e.update_and_solve(y, {x: 11})
    print(e)
    e.update_and_solve(z, {x: 2})
    print(e)