Skip to content

Instantly share code, notes, and snippets.

@versusvoid
Last active November 7, 2019 10:08
Show Gist options
  • Save versusvoid/7db4a7b4cd795ca4e6cd1122834183d4 to your computer and use it in GitHub Desktop.
Save versusvoid/7db4a7b4cd795ca4e6cd1122834183d4 to your computer and use it in GitHub Desktop.
SageMath expression "linearization". Incomplete, inaccurate, but works to the first order :-)
# vi: ft=python
from sage.symbolic.operators import add_vararg, mul_vararg
from sage.symbolic.integration.integral import definite_integral, indefinite_integral
w0, w1, w2, w3 = map(SR.wild, range(4))
definite_integral_pattern = integrate(w0, w1, w2, w3)
indefinite_integral_pattern = integrate(w0, w1)
def _substitute_definite_integral(d, integral_subs, processor):
res0 = _process_integrals(d[w0], integral_subs, processor)
res0 = processor(res0)
res2 = _process_integrals(d[w2], integral_subs, processor)
res2 = processor(res2)
res3 = _process_integrals(d[w3], integral_subs, processor)
res3 = processor(res3)
free_vars = set(res0.free_variables())
free_vars.discard(d[w1])
free_vars.update(res2.free_variables())
free_vars.update(res3.free_variables())
if not free_vars:
return integrate(res0, d[w1], res2, res3)
sub = function(f'__int{len(integral_subs)}__')(*free_vars)
integral_subs.append((integrate(res0, d[w1], res2, res3), sub))
return sub
def _substitute_indefinite_integral(d, integral_subs, processor):
res0 = _process_integrals(d[w0], integral_subs, processor)
res0 = processor(res0)
free_vars = set(res0.free_variables())
free_vars.add(d[w1])
sub = function(f'__int{len(integral_subs)}__')(*free_vars)
integral_subs.append((integrate(res0, d[w1]), sub))
return sub
def _process_integrals(expr, integral_subs, processor):
expr = expr.substitution_delayed(
definite_integral_pattern,
lambda d: _substitute_definite_integral(d, integral_subs, processor),
)
expr = expr.substitution_delayed(
indefinite_integral_pattern,
lambda d: _substitute_indefinite_integral(d, integral_subs, processor),
)
return expr
def integrals_subprocess(expr, processor):
integral_subs = []
expr = _process_integrals(expr, integral_subs, processor)
for original, sub in integral_subs[::-1]:
expr = expr.subs({sub: original})
return expr
# vi: ft=python
from sage.symbolic.operators import add_vararg, mul_vararg
from sage.symbolic.integration.integral import definite_integral, indefinite_integral
attach('integrals.sage')
def _try_subs_mutiplicative(expr, multiplicative_quantities):
# FIXME will break on non-integer powers (including symbolic expressions as powers)
expr_factors = dict(_split_factors(expr))
for quantity_factors, sub in multiplicative_quantities:
powers = [
expr_factors.get(factor, 0) // factor_power
for factor, factor_power in quantity_factors
]
if min(powers) < 1:
continue
power = max(powers)
for factor, factor_power in quantity_factors:
expr_factors[factor] -= factor_power * power
expr_factors[sub] = power
return mul_vararg(*(factor ** power for factor, power in expr_factors.items()))
def _expand_to_order(expr, order, q2s, multiplicative_quantities):
op = expr.operator()
if not op or op in q2s:
return expr
expr = op(*[
_expand_to_order(operand, order, q2s, multiplicative_quantities)
for operand in expr.operands()
])
op = expr.operator()
if not op:
return expr
if op == mul_vararg:
expr = _try_subs_mutiplicative(expr, multiplicative_quantities)
expr = expr.subs(q2s)
expr = expr.taylor(*[(v, 0) for v in q2s.values()], order)
return expr
def _get_order(expr, small_quantities):
if expr in small_quantities:
return 1
op = expr.operator()
assert op != add_vararg
operands = expr.operands()
if op == pow:
assert len(operands) == 2
return _get_order(operands[0], small_quantities) * operands[1]
elif op == mul_vararg:
return sum(_get_order(operand, small_quantities) for operand in operands)
elif op in [definite_integral, indefinite_integral]:
return _get_order(operands[0], small_quantities)
else:
return 0
def _split_factors(expr):
factors = []
for i, op in enumerate(expr.operands()):
if op.operator() == pow:
factors.append(op.operands())
else:
factors.append((op, 1))
return factors
def perturbate(expr, order, *small_quantities):
# TODO need two passes with overlapping multiplicative quantities,
# e.g. a*b/(1 + c*d) with quantities c*d and b*c
if callable(getattr(expr, 'expr', None)):
# TODO rewrap?
expr = expr.expr()
subs = SR.var('__subs__', n=len(small_quantities))
q2s = dict(zip(small_quantities, subs))
multiplicative_quantities = [
(_split_factors(key), value)
for key, value in q2s.items()
if key.operator() == mul_vararg
]
expr = integrals_subprocess(
expr,
lambda e: _expand_to_order(e, order, q2s, multiplicative_quantities),
)
expr = _expand_to_order(expr, order, q2s, multiplicative_quantities)
if expr.operator() != add_vararg:
operands = [expr]
else:
operands = expr.operands()
s2q = dict(zip(subs, small_quantities))
expr = sum(
(0 if _get_order(operand, s2q) > order else operand)
for operand in operands
)
return expr.subs(s2q)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment