import types import Numeric, LinearAlgebra valuetypes = (types.IntType, types.LongType, types.FloatType) class var: # this class represents a scalar variable def __init__(self, varname="(no variable name provided)"): self.varname = varname def __add__(self, other): if isinstance(other, valuetypes): return term([1], [self], other) elif isinstance(other, var): return term([1, 1], [self, other], 0) else: assert isinstance(other, term) return term([1], [self], 0) + other __radd__ = __add__ def __sub__(self, other): if isinstance(other, valuetypes): return term([1], [self], -other) elif isinstance(other, var): return term([1, -1], [self, other], 0) else: assert isinstance(other, term) return term([1], [self], 0) - other def __rsub__(self, other): return -self+other def __neg__(self): return term([-1], [self], 0) def __mul__(self, other): assert isinstance(other, valuetypes) return term([other], [self], 0) __rmul__ = __mul__ def __div__(self, other): assert isinstance(other, valuetypes) return term([1/other], [self], 0) def __str__(self): return self.varname class term: # this class represents the linear term: # sum([p*v.value for p, v in zip(self.prefactors, self.vars]) + self.const def __init__(self, prefactors, vars, const): self.prefactors = prefactors self.vars = vars self.const = const def __add__(self, other): if isinstance(other, valuetypes): return term(self.prefactors, self.vars, self.const + other) elif isinstance(other, var): vars = self.vars[:] prefactors = self.prefactors[:] try: prefactors[vars.index(other)] += 1 except ValueError: vars.append(other) prefactors.append(+1) return term(prefactors, vars, self.const) else: assert isinstance(other, term) vars = self.vars[:] prefactors = self.prefactors[:] for p, v in zip(other.prefactors, other.vars): try: prefactors[vars.index(v)] += p except ValueError: vars.append(v) prefactors.append(p) return term(prefactors, vars, self.const + other.const) __radd__ = __add__ def __sub__(self, other): if isinstance(other, valuetypes): return term(self.prefactors, self.vars, self.const - other) elif isinstance(other, var): vars = self.vars[:] prefactors = self.prefactors[:] try: prefactors[vars.index(other)] -= 1 except ValueError: vars.append(other) prefactors.append(-1) return term(prefactors, vars, self.const) else: assert isinstance(other, term) vars = self.vars[:] prefactors = self.prefactors[:] for p, v in zip(other.prefactors, other.vars): try: prefactors[vars.index(v)] -= p except ValueError: vars.append(v) prefactors.append(-p) return term(prefactors, vars, self.const - other.const) def __neg__(self): return term([-p for p in self.prefactors], self.vars, -self.const) def __rsub__(self, other): return -self+other def __mul__(self, other): assert isinstance(other, valuetypes) return term([p*other for p in self.prefactors], self.vars, self.const*other) __rmul__ = __mul__ def __div__(self, other): assert isinstance(other, valuetypes) return term([p/other for p in self.prefactors], self.vars, self.const/other) def __str__(self): return "+".join(["%s*%s" % pv for pv in zip(self.prefactors, self.vars)]) + "+" + str(self.const) def solve(*eqs): # this function solves a list of equations, where each equation is a tuple of terms (lhs, rhs) vars = [] l = len(eqs) a = Numeric.zeros((l, l)) b = Numeric.zeros((l, )) for i, eq in enumerate(eqs): lhs, rhs = eq if isinstance(lhs, valuetypes): lhs = term([], [], lhs) elif isinstance(lhs, var): lhs = term([1], [lhs], 0) else: assert isinstance(lhs, term) if isinstance(rhs, valuetypes): rhs = term([], [], rhs) elif isinstance(rhs, var): rhs = term([1], [rhs], 0) else: assert isinstance(rhs, term) for p, v in zip(lhs.prefactors, lhs.vars): try: j = vars.index(v) except ValueError: j = len(vars) assert j < l vars.append(v) a[i, j] += p for p, v in zip(rhs.prefactors, rhs.vars): try: j = vars.index(v) except ValueError: j = len(vars) assert j < l vars.append(v) a[i, j] -= p b[i] = rhs.const - lhs.const assert len(vars) == l for v, value in zip(vars, LinearAlgebra.solve_linear_equations(a, b)): v.value = value x = var() y = var() z = var() solve((x + y, 2*x - y + 3), (x - y, z), (5, z)) print "x=%s" % x.value print "y=%s" % y.value print "z=%s" % z.value