Last active
November 7, 2019 10:08
-
-
Save versusvoid/7db4a7b4cd795ca4e6cd1122834183d4 to your computer and use it in GitHub Desktop.
SageMath expression "linearization". Incomplete, inaccurate, but works to the first order :-)
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
# 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 |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
# 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