def scale(in_vars: Tuple[TegVar], scale: Tuple[ITeg]):
out_vars = [TegVar(f'{in_var.name}_t') for in_var in in_vars]
return (partial(BiMap,
targets=out_vars,
target_exprs=[in_var * s for (in_var, s) in zip(in_vars, scale)],
sources=in_vars,
source_exprs=[out_var / s for (out_var, s) in zip(out_vars, scale)],
inv_jacobian=teg_abs(Invert(reduce(operator.mul, scale))),
target_upper_bounds=[s * IfElse(s > 0, in_var.ub(), in_var.lb())
for (in_var, s) in zip(in_vars, scale)],
target_lower_bounds=[s * IfElse(s > 0, in_var.lb(), in_var.ub())
for (in_var, s) in zip(in_vars, scale)]),
out_vars)
def transfer_bounds_general(expr: BiMap, source_lower: Dict[TegVar, ITeg],
source_upper: Dict[TegVar, ITeg]):
"""Implements a derivative-based pessimistic bounds computation for
continuous monotonic maps. """
lb_lets = {}
ub_lets = {}
for tegvar in source_lower:
deriv_expr = fwd_deriv(expr, {tegvar: Const(1)})
lb_lets[tegvar] = (IfElse(deriv_expr > 0, source_upper[tegvar],
source_lower[tegvar]))
ub_lets[tegvar] = (IfElse(deriv_expr > 0, source_lower[tegvar],
source_upper[tegvar]))
return LetIn(lb_lets.keys(), lb_lets.values(),
expr), LetIn(ub_lets.keys(), ub_lets.values(), expr)
def outer_fn(e, ctx):
ctx['has_expr'] = any(ctx['has_exprs'])
# Check if we need to handle other such cases.
assert not (ctx['has_expr'] and isinstance(e, SmoothFunc)),\
f'expr is contained in a non-linear function {type(e)}'
if isinstance(e, Add):
if ctx['has_expr']:
ctx['expr'] = sum([
child
for child, has_expr in zip(ctx['exprs'], ctx['has_exprs'])
if has_expr
])
return ctx['expr'], ctx
else:
ctx['expr'] = e
return e, ctx
elif isinstance(e, Tup):
if ctx['has_expr']:
ctx['expr'] = Tup(*[
ctx['exprs'][idx] if has_expr else Const(0)
for idx, has_expr in enumerate(ctx['has_exprs'])
])
return ctx['expr'], ctx
elif isinstance(e, IfElse):
if ctx['has_expr']:
ctx['expr'] = IfElse(
e.cond,
ctx['exprs'][1] if ctx['has_exprs'][1] else Const(0),
ctx['exprs'][2] if ctx['has_exprs'][2] else Const(0))
return ctx, e
elif isinstance(e, LetIn):
if any(ctx['has_exprs'][1:]):
# Let expressions contain exprs.
new_exprs = [
let_var for let_var, has_expr in zip(
e.new_vars, ctx['has_exprs'][1:]) if has_expr
]
# Recursively split the body with the new expressions.
s_expr = split_exprs(new_exprs, ctx['let_body'])
let_body = (s_expr if s_expr else Const(0)) +\
(e.expr if ctx['has_exprs'][0] else Const(0))
try:
vs, es = zip(*[(v, e)
for v, e in zip(e.new_vars, e.new_exprs)
if v in let_body])
ctx['expr'] = LetIn(vs, es, let_body)
except ValueError:
# No need for a let expr.
ctx['expr'] = let_body
return ctx['expr'], ctx
ctx['expr'] = e
return ctx['expr'], ctx
def bounds_of(linear: Dict[Tuple[str, int], ITeg],
source_vars: List[TegVar]) -> List[ITeg]:
"""Generates the bounds of integration after rotation (i.e., it's the bounds transfer function). """
lower_bounds, upper_bounds = [], []
num_vars = len(source_vars)
exprs = [linear[(s_var.name, s_var.uid)] for s_var in source_vars]
for target_index in range(num_vars):
if target_index == 0:
lower = sum(exprs[i] *
IfElse(exprs[i] > 0, source_vars[i].lower_bound(),
source_vars[i].upper_bound())
for i in range(num_vars))
upper = sum(exprs[i] *
IfElse(exprs[i] > 0, source_vars[i].upper_bound(),
source_vars[i].lower_bound())
for i in range(num_vars))
elif target_index < len(linear):
def coeff(u, v):
if v == 0:
return -exprs[u]
else:
return ((Const(1) if u == v else Const(0)) -
(exprs[u] * exprs[v]) / (Const(1) + exprs[0]))
i = target_index
lower = upper = Const(0)
for j in range(num_vars):
placeholder_lb = source_vars[j].lower_bound()
placeholder_ub = source_vars[j].upper_bound()
lower += coeff(i, j) * IfElse(
coeff(i, j) > 0, placeholder_lb, placeholder_ub)
upper += coeff(i, j) * IfElse(
coeff(i, j) > 0, placeholder_ub, placeholder_lb)
else:
raise ValueError(
f'Requested target coordinate index: {target_index} is out of bounds.'
)
lower_bounds.append(lower)
upper_bounds.append(upper)
return lower_bounds, upper_bounds
def teg_cases(exprs, conditions):
assert len(exprs) - len(conditions) == 1, 'Need one additional expression for the default step'
# Build ladder in reverse order.s
exprs = exprs[::-1]
conditions = conditions[::-1]
if_else_ladder = exprs[0]
for expr, condition in zip(exprs[1:], conditions):
if_else_ladder = IfElse(condition, expr, if_else_ladder)
return if_else_ladder
def negate_degenerate_coeffs(affine: Dict[Tuple[str, int], ITeg],
source_vars: List[TegVar]):
"""Flips all the coefficients if expr[0] < 0 to avoid degeneracies"""
exprs = [affine[(s_var.name, s_var.uid)] for s_var in source_vars]
flip_condition = exprs[0] < 0
robust_affine_set = {
var: IfElse(flip_condition, -coeff, coeff)
for (var, coeff) in affine.items()
}
return robust_affine_set, flip_condition
def rotate_2d(x, y, theta):
x_ = TegVar('x_')
y_ = TegVar('y_')
return (partial(BiMap,
targets=[x_, y_],
target_exprs=[x * Cos(theta) + y * Sin(theta), -x * Sin(theta) + y * Cos(theta)],
sources=[x, y],
source_exprs=[x_ * Cos(theta) - y_ * Sin(theta), x_ * Sin(theta) + y_ * Cos(theta)],
inv_jacobian=Const(1),
target_lower_bounds=[Cos(theta) * IfElse(Cos(theta) > 0, x.lb(), x.ub()) +
Sin(theta) * IfElse(Sin(theta) > 0, y.lb(), y.ub()),
-Sin(theta) * IfElse(Sin(theta) > 0, x.ub(), x.lb()) +
Cos(theta) * IfElse(Cos(theta) > 0, y.lb(), y.ub())],
target_upper_bounds=[Cos(theta) * IfElse(Cos(theta) > 0, x.ub(), x.lb()) +
Sin(theta) * IfElse(Sin(theta) > 0, y.ub(), y.lb()),
-Sin(theta) * IfElse(Sin(theta) > 0, x.lb(), x.ub()) +
Cos(theta) * IfElse(Cos(theta) > 0, y.ub(), y.lb())]),
[x_, y_])
def outer_fn(e, ctx):
if isinstance(e, Delta) and (ctx['search_expr'] is e):
assert is_delta_normal(
e
), f'Delta {e} is not in normal form. Call normalize_delta() first'
if e.expr not in ctx['upper_tegvars']:
return Const(0), ctx
else:
return Const(1), {
**ctx, 'eliminate_tegs': {
**ctx['eliminate_tegs'], e.expr: Const(0)
}
}
elif isinstance(e, Teg):
if e.dvar in ctx['eliminate_tegs']:
value = ctx['eliminate_tegs'][e.dvar]
bounds_check = (e.lower < value) & (e.upper > value)
return (LetIn([e.dvar], [value],
IfElse(bounds_check, e.body, Const(0))), ctx)
return e, ctx
def outer_fn(e, ctx):
if isinstance(e, BiMap) and (bimap is e):
if not all([k in ctx['upper_tegvars'] for k in e.sources]):
# BiMap is invalid, null everything.
print(
f'WARNING: Attempting to map non-Teg vars {e.sources}, {ctx["upper_tegvars"]}'
)
return Const(0), ctx
bounds_checks = reduce(
operator.and_,
[(lb < dvar) & (ub > dvar)
for (dvar, (lb, ub)) in ctx['source_bounds'].items()])
reparamaterized_expr = IfElse(bounds_checks,
e.expr * e.inv_jacobian, Const(0))
return (reparamaterized_expr, {
**ctx, 'teg_sources': list(e.sources),
'teg_targets': list(e.targets),
'let_mappings':
{s: sexpr
for s, sexpr in zip(e.sources, e.source_exprs)},
'target_lower_bounds':
{t: tlb
for t, tlb in zip(e.targets, e.target_lower_bounds)},
'target_upper_bounds':
{t: tub
for t, tub in zip(e.targets, e.target_upper_bounds)}
})
elif isinstance(e, Teg):
if e.dvar in ctx.get('teg_sources', {}):
ctx['teg_sources'].remove(e.dvar)
target_dvar = ctx['teg_targets'].pop()
placeholders = {
**{
f'{svar.uid}_ub': upper
for svar, (lower, upper) in ctx['source_bounds'].items(
)
},
**{
f'{svar.uid}_lb': lower
for svar, (lower, upper) in ctx['source_bounds'].items(
)
}
}
target_lower_bounds = resolve_placeholders(
ctx['target_lower_bounds'][target_dvar], placeholders)
target_upper_bounds = resolve_placeholders(
ctx['target_upper_bounds'][target_dvar], placeholders)
# Add new teg to list.
ctx['new_tegs'] = [
*ctx.get('new_tegs', []),
(target_dvar, (target_lower_bounds, target_upper_bounds))
]
# Remove old teg.
e = e.body
if len(ctx['teg_sources']) == 0:
# Add let mappings here.
source_vars, source_exprs = zip(
*list(ctx['let_mappings'].items()))
e = LetIn(source_vars, source_exprs, e)
# Add new tegs here.
for (new_dvar, (new_lb, new_ub)) in ctx['new_tegs']:
e = Teg(new_lb, new_ub, e, new_dvar)
# Add dependent mappings here.
for new_vars, new_exprs in ctx.get('dependent_mappings',
[]):
e = LetIn(new_vars, new_exprs, e)
return e, ctx
elif isinstance(e, LetIn):
if len(ctx.get('teg_sources', {})) > 0:
if (any([
new_var in map_expr for new_var in e.new_vars
for map_vars, map_exprs in ctx.get(
'dependent_mappings', []) for map_expr in map_exprs
]) or any([
new_var in map_expr for new_var in e.new_vars
for map_var, map_expr in ctx.get('let_mappings',
{}).items()
])):
# reparametrization is dependent on this let_map. lift this map.
ctx['dependent_mappings'] = [
*ctx.get('dependent_mappings', []),
(e.new_vars, e.new_exprs)
]
return e.expr, ctx
return e, ctx
x = TegVar('x')
y = TegVar('y')
t1 = Var('t1')
t2 = Var('t2')
x_lb = Var('x_lb')
x_ub = Var('x_ub')
y_lb = Var('y_lb')
y_ub = Var('y_ub')
translate_map, (x_, y_) = translate([x, y], [t1, t2])
# Derivative of threshold only.
integral = Teg(x_lb, x_ub,
Teg(y_lb, y_ub,
translate_map(IfElse(x_ + y_ > 0.75, 2, 1)), y
), x
)
integral = reduce_to_base(integral)
image = render_image(integral, variables=((x_lb, x_ub), (y_lb, y_ub)),
bindings={t1: 0, t2: 0}, res=(args.res_x, args.res_y))
save_image(image, filename=f'{args.testname}.png')
elif args.testname == 'circle':
x = TegVar('x')
y = TegVar('y')
r = TegVar('r')
x_lb = Var('x_lb')
x_ub = Var('x_ub')
def teg_max(a, b):
return IfElse(a > b, a, b)
def teg_smoothstep(x):
return IfElse(x > 0, IfElse(x < 1, 3 * Sqr(x) - 2 * Sqr(x) * x, Const(1)), Const(0))
def polar_2d_map(expr, x, y, r):
"""
Create a polar 2D map with x=0, y=0 as center and negative y axis as 0 & 2PI
"""
theta = TegVar('theta')
distance_to_origin = Sqrt(
Sqr((y.lb() + y.ub()) / 2) + Sqr((x.lb() + x.ub()) / 2))
box_radius = Sqrt(Sqr((y.ub() - y.lb()) / 2) + Sqr((x.ub() - x.lb()) / 2))
# Manual interval arithmetic for conservative polar bounds.
# These are not strictly necessary.
# Using (0,2pi) still produces the correct unbiased integrals.
# However, they will have terrible sample behaviour)
box_upper_right = ATan2(Tup(x.ub(), y.ub()))
box_lower_right = ATan2(Tup(x.ub(), y.lb()))
box_upper_left = ATan2(Tup(x.lb(), y.ub()))
box_lower_left = ATan2(Tup(x.lb(), y.lb()))
right_theta_lower = IfElse(y.ub() > 0, IfElse(x.lb() > 0, box_upper_left,
0), box_upper_right)
right_theta_upper = IfElse(y.lb() < 0,
IfElse(x.lb() > 0, box_lower_left, TEG_PI),
box_lower_right)
left_theta_upper = IfElse(y.ub() > 0, IfElse(x.ub() < 0, box_upper_right,
0), box_upper_left)
left_theta_lower = IfElse(
y.lb() < 0, IfElse(x.ub() < 0, box_lower_right, TEG_NEGATIVE_PI),
box_lower_left)
return IfElse(
x > 0,
BiMap(expr,
sources=[x, y],
source_exprs=[r * Sin(theta), r * Cos(theta)],
targets=[r, theta],
target_exprs=[Sqrt(Sqr(x) + Sqr(y)),
ATan2(Tup(x, y))],
inv_jacobian=r,
target_lower_bounds=[
teg_max(distance_to_origin - box_radius, 0),
right_theta_lower
],
target_upper_bounds=[
distance_to_origin + box_radius, right_theta_upper
]),
BiMap(expr,
sources=[x, y],
source_exprs=[r * Sin(theta), r * Cos(theta)],
targets=[r, theta],
target_exprs=[Sqrt(Sqr(x) + Sqr(y)),
ATan2(Tup(x, y))],
inv_jacobian=r,
target_lower_bounds=[
teg_max(distance_to_origin - box_radius, 0), left_theta_lower
],
target_upper_bounds=[
distance_to_origin + box_radius, left_theta_upper
]))
def teg_abs(a):
return IfElse(a > 0, a, -a)
def teg_min(a, b):
return IfElse(a > b, b, a)
def reverse_deriv_transform(
expr: ITeg, out_deriv_vals: Tuple, not_ctx: Set[Tuple[str, int]],
deps: Dict[TegVar, Set[Var]],
args: Dict[str, Any]) -> Iterable[Tuple[Tuple[str, int], ITeg]]:
if isinstance(expr, TegVar):
if (((expr.name, expr.uid) not in not_ctx)
or {(v.name, v.uid)
for v in extend_dependencies({expr}, deps)} - not_ctx):
yield ((f'd{expr.name}', expr.uid), out_deriv_vals)
elif isinstance(expr, (Const, Delta)):
pass
elif isinstance(expr, Var):
if (expr.name, expr.uid) not in not_ctx:
yield ((f'd{expr.name}', expr.uid), out_deriv_vals)
elif isinstance(expr, Add):
left, right = expr.children
# yield from reverse_deriv_transform(left, out_deriv_vals, not_ctx, teg_list)
# yield from reverse_deriv_transform(right, out_deriv_vals, not_ctx, teg_list)
left_list = list(
reverse_deriv_transform(left, Const(1), not_ctx, deps, args))
right_list = list(
reverse_deriv_transform(right, Const(1), not_ctx, deps, args))
yield from merge(left_list, right_list, out_deriv_vals)
elif isinstance(expr, Mul):
left, right = expr.children
# yield from reverse_deriv_transform(left, out_deriv_vals * right, not_ctx, deps)
# yield from reverse_deriv_transform(right, out_deriv_vals * left, not_ctx, deps)
left_list = list(
reverse_deriv_transform(left, right, not_ctx, deps, args))
right_list = list(
reverse_deriv_transform(right, left, not_ctx, deps, args))
yield from merge(left_list, right_list, out_deriv_vals)
elif isinstance(expr, Invert):
child = expr.child
yield from reverse_deriv_transform(child,
-out_deriv_vals * expr * expr,
not_ctx, deps, args)
elif isinstance(expr, SmoothFunc):
child = expr.expr
yield from reverse_deriv_transform(
child, expr.rev_deriv(out_deriv_expr=out_deriv_vals), not_ctx,
deps, args)
elif isinstance(expr, IfElse):
derivs_if = reverse_deriv_transform(expr.if_body, Const(1), not_ctx,
deps, args)
derivs_else = reverse_deriv_transform(expr.else_body, Const(1),
not_ctx, deps, args)
yield from ((name_uid,
out_deriv_vals * IfElse(expr.cond, deriv_if, Const(0)))
for name_uid, deriv_if in derivs_if)
yield from ((name_uid,
out_deriv_vals * IfElse(expr.cond, Const(0), deriv_else))
for name_uid, deriv_else in derivs_else)
if not args.get('ignore_deltas', False):
for boolean in primitive_booleans_in(expr.cond, not_ctx, deps):
jump = substitute(expr, boolean, true) - substitute(
expr, boolean, false)
delta_expr = boolean.right_expr - boolean.left_expr
derivs_delta_expr = reverse_deriv_transform(
delta_expr, Const(1), not_ctx, deps, args)
yield from (
(name_uid, out_deriv_vals * deriv_delta_expr * jump *
Delta(delta_expr))
for name_uid, deriv_delta_expr in derivs_delta_expr)
elif isinstance(expr, Teg):
not_ctx.discard((expr.dvar.name, expr.dvar.uid))
# Apply Leibniz rule directly for moving boundaries
if not args.get('ignore_bounds', False):
lower_derivs = reverse_deriv_transform(expr.lower, out_deriv_vals,
not_ctx, deps, args)
upper_derivs = reverse_deriv_transform(expr.upper, out_deriv_vals,
not_ctx, deps, args)
yield from ((name_uid, upper_deriv *
substitute(expr.body, expr.dvar, expr.upper))
for name_uid, upper_deriv in upper_derivs)
yield from ((name_uid, -lower_deriv *
substitute(expr.body, expr.dvar, expr.lower))
for name_uid, lower_deriv in lower_derivs)
not_ctx.add((expr.dvar.name, expr.dvar.uid))
deriv_body_traces = reverse_deriv_transform(expr.body, Const(1),
not_ctx, deps, args)
yield from ((name_uid, out_deriv_vals *
Teg(expr.lower, expr.upper, deriv_body, expr.dvar))
for name_uid, deriv_body in deriv_body_traces)
elif isinstance(expr, Tup):
yield [
reverse_deriv_transform(child, out_deriv_vals, not_ctx, deps, args)
for child in expr
]
elif isinstance(expr, LetIn):
# Include derivatives of each expression to the let body
dnew_vars, body_derivs = set(), {}
for var, e in zip(expr.new_vars, expr.new_exprs):
# print(not_ctx)
# print(var, e)
if any(
Var(name=ctx_name, uid=ctx_uid) in e
for ctx_name, ctx_uid in not_ctx):
# Add dependent variables.
assert isinstance(var, TegVar), f'{var} is dependent on TegVar(s):'\
f'({[ctx_var for ctx_var in not_ctx if ctx_var in e]}).'\
f'{var} must also be declared as a TegVar and not a Var'
# print(not_ctx)
not_ctx = not_ctx | {(var.name, var.uid)}
# print(var)
if var not in expr.expr:
# print('Not in expression')
continue
# print('In expression')
dname = f'd{var.name}'
dnew_vars.add((dname, var.uid))
body_derivs[(dname, var.uid)] = list(
reverse_deriv_transform(e, Const(1), not_ctx, deps, args))
# Thread through derivatives of each subexpression
for (name, uid), dname_expr in reverse_deriv_transform(
expr.expr, out_deriv_vals, not_ctx, deps, args):
dvar_with_ctx = LetIn(expr.new_vars, expr.new_exprs, dname_expr)
if (name, uid) in dnew_vars:
yield from ((n, d * dvar_with_ctx)
for n, d in body_derivs[(name, uid)])
else:
yield ((name, uid), dvar_with_ctx)
elif isinstance(expr, BiMap):
# Include derivatives of each expression to the let body
dnew_vars, body_derivs = set(), {}
new_deps = {}
for var, e in zip(expr.targets, expr.target_exprs):
if any(
Var(name=ctx_name, uid=ctx_uid) in e
for ctx_name, ctx_uid in not_ctx):
# Add dependent variables.
assert isinstance(var, TegVar), f'{var} is dependent on TegVar(s):'\
f'({[ctx_var for ctx_var in not_ctx if ctx_var in e]}).'\
f'{var} must also be declared as a TegVar and not a Var'
not_ctx = not_ctx | {(var.name, var.uid)}
if var in expr.expr:
new_deps[var] = extract_vars(e)
dname = f'd{var.name}'
dnew_vars.add((dname, var.uid))
body_derivs[(dname, var.uid)] = list(
reverse_deriv_transform(e, Const(1), not_ctx, deps, args))
deps = {**deps, **new_deps}
# Thread through derivatives of each subexpression
for (name, uid), dname_expr in reverse_deriv_transform(
expr.expr, out_deriv_vals, not_ctx, deps, args):
dvar_with_ctx = BiMap(dname_expr,
expr.targets,
expr.target_exprs,
expr.sources,
expr.source_exprs,
inv_jacobian=expr.inv_jacobian,
target_lower_bounds=expr.target_lower_bounds,
target_upper_bounds=expr.target_upper_bounds)
if (name, uid) in dnew_vars:
yield from ((n, d * dvar_with_ctx)
for n, d in body_derivs[(name, uid)])
else:
yield ((name, uid), dvar_with_ctx)
else:
raise ValueError(
f'The type of the expr "{type(expr)}" does not have a supported derivative.'
)
def do_pass(expr: ITeg, context, inner_fn, outer_fn,
context_combine) -> Tuple[ITeg, Dict]:
"""Substitute this_var with that_var in expr."""
if isinstance(expr, Const):
expr, out_context = outer_fn(expr, context_combine([], context))
return expr, out_context
elif isinstance(expr, (Var, TegVar, Placeholder)):
expr, out_context = outer_fn(expr, context_combine([], context))
return expr, out_context
elif isinstance(expr, Add):
left_expr, left_ctx = do_pass(*inner_fn(expr.children[0], context),
inner_fn, outer_fn, context_combine)
right_expr, right_ctx = do_pass(*inner_fn(expr.children[1], context),
inner_fn, outer_fn, context_combine)
return outer_fn(
expr if (left_expr is expr.children[0]) and
(right_expr is expr.children[1]) else left_expr + right_expr,
context_combine([left_ctx, right_ctx], context))
elif isinstance(expr, Mul):
left_expr, left_ctx = do_pass(*inner_fn(expr.children[0], context),
inner_fn, outer_fn, context_combine)
right_expr, right_ctx = do_pass(*inner_fn(expr.children[1], context),
inner_fn, outer_fn, context_combine)
return outer_fn(
expr if (left_expr is expr.children[0]) and
(right_expr is expr.children[1]) else left_expr * right_expr,
context_combine([left_ctx, right_ctx], context))
elif isinstance(expr, Invert):
child, child_ctx = do_pass(*inner_fn(expr.child, context), inner_fn,
outer_fn, context_combine)
return outer_fn(
Invert(child) if expr.child is not child else expr,
context_combine([child_ctx], context))
elif isinstance(expr, SmoothFunc):
child, child_ctx = do_pass(*inner_fn(expr.expr, context), inner_fn,
outer_fn, context_combine)
return outer_fn(
type(expr)(child) if expr.expr is not child else expr,
context_combine([child_ctx], context))
elif isinstance(expr, IfElse):
cond, cond_ctx = do_pass(*inner_fn(expr.cond, context), inner_fn,
outer_fn, context_combine)
if_body, if_ctx = do_pass(*inner_fn(expr.if_body, context), inner_fn,
outer_fn, context_combine)
else_body, else_ctx = do_pass(*inner_fn(expr.else_body, context),
inner_fn, outer_fn, context_combine)
expr = IfElse(cond, if_body, else_body) if (
cond is not expr.cond or if_body is not expr.if_body
or else_body is not expr.else_body) else expr
return outer_fn(expr,
context_combine([cond_ctx, if_ctx, else_ctx], context))
elif isinstance(expr, Teg):
# dvar, dvar_ctx = do_pass(*inner_fn(expr.dvar, context), inner_fn, outer_fn, context_combine)
body, body_ctx = do_pass(*inner_fn(expr.body, context), inner_fn,
outer_fn, context_combine)
lower, lower_ctx = do_pass(*inner_fn(expr.lower, context), inner_fn,
outer_fn, context_combine)
upper, upper_ctx = do_pass(*inner_fn(expr.upper, context), inner_fn,
outer_fn, context_combine)
expr = expr if (body is expr.body and lower is expr.lower
and upper is expr.upper) else Teg(
lower, upper, body, expr.dvar)
return outer_fn(
expr, context_combine([lower_ctx, upper_ctx, body_ctx], context))
elif isinstance(expr, Tup):
exprs, expr_contexts = zip(*[(do_pass(*inner_fn(
child, context), inner_fn, outer_fn, context_combine))
for child in expr])
expr = expr if all([
new_child is old_child
for new_child, old_child in zip(exprs, expr)
]) else Tup(*exprs)
return outer_fn(expr, context_combine(expr_contexts, context))
elif isinstance(expr, LetIn):
body_expr, body_context = do_pass(*inner_fn(expr.expr, context),
inner_fn, outer_fn, context_combine)
let_exprs, let_contexts = zip(*[
do_pass(*inner_fn(child, context), inner_fn, outer_fn,
context_combine) for child in expr.new_exprs
])
expr = expr if (all([
new_child is old_child
for new_child, old_child in zip(let_exprs, expr.new_exprs)
]) and body_expr is expr.expr) else LetIn(expr.new_vars, let_exprs,
body_expr)
return outer_fn(
expr, context_combine([body_context, *let_contexts], context))
elif isinstance(expr, Bool):
left_expr, left_ctx = do_pass(*inner_fn(expr.left_expr, context),
inner_fn, outer_fn, context_combine)
right_expr, right_ctx = do_pass(*inner_fn(expr.right_expr, context),
inner_fn, outer_fn, context_combine)
expr = expr if (left_expr is expr.left_expr
and right_expr is expr.right_expr) else Bool(
left_expr, right_expr, allow_eq=expr.allow_eq)
return outer_fn(expr, context_combine([left_ctx, right_ctx], context))
elif isinstance(expr, (And, Or)):
left_expr, left_ctx = do_pass(*inner_fn(expr.left_expr, context),
inner_fn, outer_fn, context_combine)
right_expr, right_ctx = do_pass(*inner_fn(expr.right_expr, context),
inner_fn, outer_fn, context_combine)
expr = expr if (left_expr is expr.left_expr
and right_expr is expr.right_expr) else type(expr)(
left_expr, right_expr)
return outer_fn(expr, context_combine([left_ctx, right_ctx], context))
elif isinstance(expr, BiMap):
body_expr, body_context = do_pass(*inner_fn(expr.expr, context),
inner_fn, outer_fn, context_combine)
source_exprs, source_contexts = zip(*[
do_pass(*inner_fn(child, context), inner_fn, outer_fn,
context_combine) for child in expr.source_exprs
])
target_exprs, target_contexts = zip(*[
do_pass(*inner_fn(child, context), inner_fn, outer_fn,
context_combine) for child in expr.target_exprs
])
jacobian_expr, jacobian_context = do_pass(
*inner_fn(expr.inv_jacobian, context), inner_fn, outer_fn,
context_combine)
target_upper_bounds, ub_contexts = zip(*[
do_pass(*inner_fn(child, context), inner_fn, outer_fn,
context_combine) for child in expr.target_upper_bounds
])
target_lower_bounds, lb_contexts = zip(*[
do_pass(*inner_fn(child, context), inner_fn, outer_fn,
context_combine) for child in expr.target_lower_bounds
])
expr = expr if (all([
new_child is old_child
for new_child, old_child in zip(source_exprs, expr.source_exprs)
]) and all([
new_child is old_child
for new_child, old_child in zip(target_exprs, expr.target_exprs)
]) and all([
new_child is old_child for new_child, old_child in zip(
target_upper_bounds, expr.target_upper_bounds)
]) and all([
new_child is old_child for new_child, old_child in zip(
target_lower_bounds, expr.target_lower_bounds)
]) and body_expr is expr.expr
and jacobian_expr is expr.inv_jacobian) else BiMap(
body_expr, expr.targets, target_exprs,
expr.sources, source_exprs, jacobian_expr,
target_upper_bounds, target_lower_bounds)
return outer_fn(
expr,
context_combine([
*source_contexts, *target_contexts, jacobian_context,
*ub_contexts, *lb_contexts, body_context
], context))
elif isinstance(expr, Delta):
body_expr, body_context = do_pass(*inner_fn(expr.expr, context),
inner_fn, outer_fn, context_combine)
expr = expr if body_expr is expr.expr else Delta(body_expr)
return outer_fn(expr, context_combine([body_context], context))
else:
raise ValueError(
f'The type of the expr "{type(expr)}" is not supported by substitute.'
)
def teg_abs(x):
return IfElse(x > 0, x, -x)
def fwd_deriv_transform(
expr: ITeg, ctx: Dict[Tuple[str, int], ITeg],
not_ctx: Set[Tuple[str, int]], deps: Dict[TegVar, Set[Var]]
) -> Tuple[ITeg, Dict[Tuple[str, int], str], Set[Tuple[str, int]]]:
"""Compute the source-to-source foward derivative of the given expression."""
if isinstance(expr, TegVar):
if (((expr.name, expr.uid) not in not_ctx
or {(v.name, v.uid)
for v in extend_dependencies({expr}, deps)} - not_ctx)
and (expr.name, expr.uid) in ctx):
expr = ctx[(expr.name, expr.uid)]
else:
expr = Const(0)
elif isinstance(expr, (Const, Placeholder, Delta)):
expr = Const(0)
elif isinstance(expr, Var):
if (expr.name, expr.uid) not in not_ctx and (expr.name,
expr.uid) in ctx:
expr = ctx[(expr.name, expr.uid)]
else:
expr = Const(0)
elif isinstance(expr, SmoothFunc):
in_deriv_expr, ctx, not_ctx, deps = fwd_deriv_transform(
expr.expr, ctx, not_ctx, deps)
deriv_expr = expr.fwd_deriv(in_deriv_expr=in_deriv_expr)
expr = deriv_expr
elif isinstance(expr, Add):
sum_of_derivs = Const(0)
for child in expr.children:
deriv_child, ctx, not_ctx, deps = fwd_deriv_transform(
child, ctx, not_ctx, deps)
sum_of_derivs += deriv_child
expr = sum_of_derivs
elif isinstance(expr, Mul):
# NOTE: Consider n-ary multiplication.
assert len(
expr.children
) == 2, 'fwd_deriv only supports binary multiplication not n-ary.'
expr1, expr2 = [child for child in expr.children]
(deriv_expr1, ctx1, not_ctx1,
_) = fwd_deriv_transform(expr1, ctx, not_ctx, deps)
(deriv_expr2, ctx2, not_ctx2,
_) = fwd_deriv_transform(expr2, ctx, not_ctx, deps)
expr = expr1 * deriv_expr2 + expr2 * deriv_expr1
ctx = {**ctx1, **ctx2}
not_ctx = not_ctx1 | not_ctx2
elif isinstance(expr, Invert):
deriv_expr, ctx, not_ctx, deps = fwd_deriv_transform(
expr.child, ctx, not_ctx, deps)
expr = -expr * expr * deriv_expr
elif isinstance(expr, IfElse):
if_body, ctx, not_ctx1, _ = fwd_deriv_transform(
expr.if_body, ctx, not_ctx, deps)
else_body, ctx, not_ctx2, _ = fwd_deriv_transform(
expr.else_body, ctx, not_ctx, deps)
not_ctx = not_ctx1 | not_ctx2
deltas = Const(0)
for boolean in primitive_booleans_in(expr.cond, not_ctx, deps):
jump = substitute(expr, boolean, true) - substitute(
expr, boolean, false)
delta_expr = boolean.right_expr - boolean.left_expr
delta_deriv, ctx, _ignore_not_ctx, _ = fwd_deriv_transform(
delta_expr, ctx, not_ctx, deps)
deltas = deltas + delta_deriv * jump * Delta(delta_expr)
expr = IfElse(expr.cond, if_body, else_body) + deltas
elif isinstance(expr, Teg):
assert expr.dvar not in ctx, f'Names of infinitesimal "{expr.dvar}" are distinct from context "{ctx}"'
# In int_x f(x), the variable x is in scope for the integrand f(x)
not_ctx.discard(expr.dvar.name)
# Include derivative contribution from moving boundaries of integration
boundary_val, new_ctx, new_not_ctx = boundary_contribution(
expr, ctx, not_ctx, deps)
not_ctx.add((expr.dvar.name, expr.dvar.uid))
body, ctx, not_ctx, _ = fwd_deriv_transform(expr.body, ctx, not_ctx,
deps)
ctx.update(new_ctx)
not_ctx |= new_not_ctx
expr = Teg(expr.lower, expr.upper, body, expr.dvar) + boundary_val
elif isinstance(expr, Tup):
new_expr_list, new_ctx, new_not_ctx = [], Ctx(), set()
for child in expr:
child, ctx, not_ctx, _ = fwd_deriv_transform(
child, ctx, not_ctx, deps)
new_expr_list.append(child)
new_ctx.update(ctx)
new_not_ctx |= not_ctx
ctx, not_ctx = new_ctx, new_not_ctx
expr = Tup(*new_expr_list)
elif isinstance(expr, LetIn):
# Compute derivatives of each expression and bind them to the corresponding dvar
new_vars_with_derivs, new_exprs_with_derivs = list(
expr.new_vars), list(expr.new_exprs)
new_deps = {}
for v, e in zip(expr.new_vars, expr.new_exprs):
if v in expr.expr:
# By not passing in the updated contexts,
# we require that assignments in let expressions are independent
de, ctx, not_ctx, _ = fwd_deriv_transform(
e, ctx, not_ctx, deps)
ctx[(v.name, v.uid)] = Var(f'd{v.name}')
new_vars_with_derivs.append(ctx[(v.name, v.uid)])
new_exprs_with_derivs.append(de)
new_deps[v] = extract_vars(e)
deps = {**deps, **new_deps}
# We want an expression in terms of f'd{var_in_let_body}'
# This means that they are erroniously added to ctx, so we
# remove them from ctx!
dexpr, ctx, not_ctx, _ = fwd_deriv_transform(expr.expr, ctx, not_ctx,
deps)
[ctx.pop((c.name, c.uid), None) for c in expr.new_vars]
expr = LetIn(Tup(*new_vars_with_derivs), Tup(*new_exprs_with_derivs),
dexpr)
elif isinstance(expr, BiMap):
# TODO: is it possible to not repeat this code and make another recursive call instead?
# Compute derivatives of each expression and bind them to the corresponding dvar
new_vars_with_derivs, new_exprs_with_derivs = [], []
for v, e in zip(expr.targets, expr.target_exprs):
if v in expr.expr:
# By not passing in the updated contexts, require independence of exprs in the body of the let expression
de, ctx, not_ctx, _ = fwd_deriv_transform(
e, ctx, not_ctx, deps)
ctx[(v.name, v.uid)] = Var(f'd{v.name}')
new_vars_with_derivs.append(ctx[(v.name, v.uid)])
new_exprs_with_derivs.append(de)
not_ctx = not_ctx | {(v.name, v.uid)}
# We want an expression in terms of f'd{var_in_let_body}'
# This means that they are erroniously added to ctx, so we
# remove them from ctx!
dexpr, ctx, not_ctx, _ = fwd_deriv_transform(expr.expr, ctx, not_ctx,
deps)
[ctx.pop((c.name, c.uid), None) for c in expr.targets]
expr = LetIn(
Tup(*new_vars_with_derivs), Tup(*new_exprs_with_derivs),
BiMap(dexpr,
targets=expr.targets,
target_exprs=expr.target_exprs,
sources=expr.sources,
source_exprs=expr.source_exprs,
inv_jacobian=expr.inv_jacobian,
target_lower_bounds=expr.target_lower_bounds,
target_upper_bounds=expr.target_upper_bounds))
else:
raise ValueError(
f'The type of the expr "{type(expr)}" does not have a supported fwd_derivative.'
)
return expr, ctx, not_ctx, deps
def rewrite(delta, not_ctx=set()):
# Extract polynomial coefficients.
poly_set = extract_coefficients_from_polynomial(
delta.expr, {(var.name, var.uid)
for var in not_ctx})
unique_vars = []
for term in poly_set:
for var, _ in term:
if var is not CONST_VAR:
unique_vars.append(var)
x = unique_vars[0]
y = unique_vars[1]
c_xy = get_poly_term(poly_set, {x: 1, y: 1})
c_x = get_poly_term(poly_set, {x: 1})
c_y = get_poly_term(poly_set, {y: 1})
c_1 = get_poly_term(poly_set, {})
c_xy_var = Var(f'c_{x.name}_{y.name}')
c_x_var = Var(f'c_{x.name}')
c_y_var = Var(f'c_{y.name}')
c_1_var = Var('c_1')
coeff_vars = [c_xy_var, c_x_var, c_y_var, c_1_var]
coeff_exprs = [
teg_abs(c_xy),
IfElse(c_xy > 0, c_x, -c_x),
IfElse(c_xy > 0, c_y, -c_y),
IfElse(c_xy > 0, c_1, -c_1)
]
sqrt_c_xy = Sqrt(c_xy_var)
sqrt_c_xy_var = Var(f'{x.name}_{y.name}_sqrt')
needs_transforms = (c_x != Const(0) or c_y != Const(0))
if needs_transforms:
scale_map = partial(scale, scale=[sqrt_c_xy_var, sqrt_c_xy_var])
translate_map = partial(
translate,
translate=[c_y_var / sqrt_c_xy_var, c_x_var / sqrt_c_xy_var])
scaler, (x_s, y_s) = scale_map([x, y])
translater, (x_st, y_st) = translate_map([x_s, y_s])
sqr_constant = (c_x_var * c_y_var) / (c_xy_var) - c_1_var
scale_jacobian = Const(1)
else:
x_st, y_st = x, y
sqr_constant = -c_1_var / c_xy_var
scale_jacobian = c_xy_var
# If threshold is negative, the hyperbola is in the second and fourth quadrants.
# Inverting either one of x or y automatically handles this.
conditional_inverter, (x_st, ) = scale(
[x_st], scale=[IfElse(sqr_constant > 0, 1, -1)])
adjusted_sqr_constant = teg_abs(sqr_constant)
constant = Sqrt(adjusted_sqr_constant)
# Hyperbolic transform
hyp_a, hyp_t = TegVar('hyp_a'), TegVar('hyp_t')
# Build bounds transfer expressions.
pos_a_lb = teg_cases([Sqrt(x_st.lb() * y_st.lb()),
Const(0)], [(x_st.lb() > 0) & (y_st.lb() > 0)])
pos_a_ub = teg_cases([Sqrt(x_st.ub() * y_st.ub()),
Const(0)], [(x_st.ub() > 0) & (y_st.ub() > 0)])
neg_a_lb = teg_cases([-Sqrt(x_st.lb() * y_st.lb()),
Const(0)], [(x_st.lb() < 0) & (y_st.lb() < 0)])
neg_a_ub = teg_cases([-Sqrt(x_st.ub() * y_st.ub()),
Const(0)], [(x_st.ub() < 0) & (y_st.ub() < 0)])
pos_t_lb = teg_max(
teg_cases([
teg_max(x_st.lb() / hyp_a, hyp_a / y_st.ub()),
hyp_a / y_st.ub(), MIN_T
], [(y_st.ub() > 0) & (x_st.lb() > 0),
y_st.ub() > 0]), MIN_T)
pos_t_ub = teg_min(
teg_cases([
teg_min(x_st.ub() / hyp_a, hyp_a / y_st.lb()),
x_st.ub() / hyp_a, MAX_T
], [(x_st.ub() > 0) & (y_st.lb() > 0),
x_st.ub() > 0]), MAX_T)
neg_t_lb = teg_max(
teg_cases([
teg_max(x_st.ub() / hyp_a, hyp_a / y_st.lb()),
hyp_a / y_st.lb(), MIN_T
], [(y_st.lb() < 0) & (x_st.ub() < 0),
y_st.lb() < 0]), MIN_T)
neg_t_ub = teg_min(
teg_cases([
teg_min(x_st.lb() / hyp_a, hyp_a / y_st.ub()),
x_st.lb() / hyp_a, MAX_T
], [(x_st.lb() < 0) & (y_st.ub() < 0),
x_st.lb() < 0]), MAX_T)
pos_curve = BiMap(Delta(hyp_a - constant),
sources=[x_st, y_st],
source_exprs=[hyp_a * hyp_t, hyp_a / hyp_t],
targets=[hyp_a, hyp_t],
target_exprs=[Sqrt(x_st * y_st),
Sqrt(x_st / y_st)],
inv_jacobian=(hyp_a / hyp_t) *
(1 / (constant * scale_jacobian)),
target_lower_bounds=[pos_a_lb, pos_t_lb],
target_upper_bounds=[pos_a_ub, pos_t_ub])
neg_curve = BiMap(Delta(hyp_a + constant),
sources=[x_st, y_st],
source_exprs=[hyp_a * hyp_t, hyp_a / hyp_t],
targets=[hyp_a, hyp_t],
target_exprs=[-Sqrt(x_st * y_st),
Sqrt(x_st / y_st)],
inv_jacobian=(-1 * hyp_a / hyp_t) *
(1 / (constant * scale_jacobian)),
target_lower_bounds=[neg_a_lb, neg_t_lb],
target_upper_bounds=[neg_a_ub, neg_t_ub])
if needs_transforms:
return LetIn(
coeff_vars, coeff_exprs,
LetIn([sqrt_c_xy_var], [sqrt_c_xy],
scaler(
translater(
conditional_inverter(pos_curve + neg_curve)))))
else:
return LetIn(coeff_vars, coeff_exprs,
conditional_inverter(pos_curve + neg_curve))
def simplify(expr: ITeg) -> ITeg:
if isinstance(expr, Var):
return expr
elif isinstance(expr, Add):
expr1, expr2 = expr.children
simple1, simple2 = simplify(expr1), simplify(expr2)
if isinstance(simple1, Const) and simple1.value == 0:
return simple2
if isinstance(simple2, Const) and simple2.value == 0:
return simple1
if isinstance(simple1, Const) and isinstance(simple2, Const):
return Const(evaluate(simple1 + simple2))
# Associative reordering.
if isinstance(simple1,
(Add, Const)) and isinstance(simple2, (Add, Const)):
nodes1 = [
simple1,
] if isinstance(simple1, Const) else simple1.children
nodes2 = [
simple2,
] if isinstance(simple2, Const) else simple2.children
all_nodes = nodes1 + nodes2
assert 2 <= len(
all_nodes
) <= 4, 'Unexpected number of nodes in Add-associative tree'
const_nodes = [
node for node in all_nodes if isinstance(node, Const)
]
other_nodes = [
node for node in all_nodes if not isinstance(node, Const)
]
# No const nodes -> Reordering is pointless.
if len(other_nodes) == len(all_nodes):
return simple1 + simple2
# Compress const nodes.
const_node = Const(evaluate(reduce(operator.add, const_nodes)))
# Re-order to front.
if const_node == Const(0):
simplified_nodes = other_nodes
else:
simplified_nodes = other_nodes + [const_node]
# Build tree in reverse (so const node is at top level)
return reduce(operator.add, simplified_nodes)
if isinstance(simple1, LetIn) and isinstance(simple2, LetIn):
if simple1.new_vars == simple2.new_vars and simple1.new_exprs == simple2.new_exprs:
return LetIn(new_vars=simple1.new_vars,
new_exprs=simple1.new_exprs,
expr=simplify(simple1.expr + simple2.expr))
else:
return simple1 + simple2
if isinstance(simple1, Teg) and isinstance(simple2, Teg):
if (simple1.dvar == simple2.dvar and simple1.lower == simple2.lower
and simple1.upper == simple2.upper):
return simplify(
Teg(simple1.lower, simple1.upper,
simplify(simple1.body + simple2.body), simple1.dvar))
else:
return simple1 + simple2
if isinstance(simple1, IfElse) and isinstance(simple2, IfElse):
if simple1.cond == simple2.cond:
return IfElse(simple1.cond,
simplify(simple1.if_body + simple2.if_body),
simplify(simple1.else_body + simple2.else_body))
else:
return simple1 + simple2
if isinstance(simple1, Mul) and isinstance(simple2, Mul):
# Distribution.
exprLL, exprLR = simple1.children
exprRL, exprRR = simple2.children
if exprLL == exprRR:
return simplify(exprLL * (simplify(exprLR + exprRL)))
if exprLL == exprRL:
return simplify(exprLL * (simplify(exprLR + exprRR)))
if exprLR == exprRL:
return simplify(exprLR * (simplify(exprLL + exprRR)))
if exprLR == exprRR:
return simplify(exprLR * (simplify(exprLL + exprRL)))
return simple1 + simple2
elif isinstance(expr, Mul):
expr1, expr2 = expr.children
simple1, simple2 = simplify(expr1), simplify(expr2)
# 0-elimination
if ((isinstance(simple1, Const) and simple1.value == 0)
or (isinstance(simple2, Const) and hasattr(simple2, 'value')
and simple2.value == 0)):
return Const(0)
# Multiplicative inverse.
if isinstance(simple1, Const) and simple1.value == 1.0:
return simple2
if isinstance(simple2, Const) and simple2.value == 1.0:
return simple1
# Local constant compression.
if isinstance(simple1, Const) and isinstance(simple2, Const):
return Const(evaluate(simple1 * simple2))
# Associative reordering.
if isinstance(simple1,
(Mul, Const)) and isinstance(simple2, (Mul, Const)):
nodes1 = [simple1] if isinstance(simple1,
Const) else simple1.children
nodes2 = [simple2] if isinstance(simple2,
Const) else simple2.children
all_nodes = nodes1 + nodes2
assert 2 <= len(
all_nodes
) <= 4, 'Unexpected number of nodes in Mul-associative tree'
const_nodes = [
node for node in all_nodes if isinstance(node, Const)
]
other_nodes = [
node for node in all_nodes if not isinstance(node, Const)
]
# No const nodes -> Reordering is pointless.
if len(other_nodes) == len(all_nodes):
return simple1 * simple2
# Compress const nodes.
const_node = Const(evaluate(reduce(operator.mul, const_nodes)))
# Re-order to front.
if not (const_node == Const(1)):
simplified_nodes = other_nodes + [const_node]
else:
simplified_nodes = other_nodes
# Build tree in reverse (so const node is at top level)
return reduce(operator.mul, simplified_nodes)
return simple1 * simple2
elif isinstance(expr, Invert):
simple = simplify(expr.child)
if isinstance(simple, Const):
return Const(evaluate(Invert(simple)))
return Invert(simple)
elif isinstance(expr, SmoothFunc):
simple = simplify(expr.expr)
if isinstance(simple, Const):
return Const(evaluate(type(expr)(simple)))
return type(expr)(simplify(expr.expr))
elif isinstance(expr, IfElse):
cond, if_body, else_body = simplify(expr.cond), simplify(
expr.if_body), simplify(expr.else_body)
if (isinstance(if_body, Const) and isinstance(else_body, Const)
and if_body.value == 0 and else_body.value == 0):
return if_body
if cond == true:
return if_body
if cond == false:
return else_body
return IfElse(cond, if_body, else_body)
elif isinstance(expr, Teg):
body = simplify(expr.body)
if isinstance(body, Const) and hasattr(body,
'value') and body.value == 0:
return Const(0)
return Teg(simplify(expr.lower), simplify(expr.upper), body, expr.dvar)
elif isinstance(expr, Tup):
return Tup(*(simplify(child) for child in expr))
elif isinstance(expr, LetIn):
simplified_exprs = Tup(*(simplify(e) for e in expr.new_exprs))
child_expr = simplify(expr.expr)
vars_list = expr.new_vars
for s_var, s_expr in zip(vars_list, simplified_exprs):
if isinstance(s_expr, Const):
child_expr = substitute(child_expr, s_var, s_expr)
non_const_bindings = [
(s_var, s_expr)
for s_var, s_expr in zip(vars_list, simplified_exprs)
if not isinstance(s_expr, Const)
]
child_expr = simplify(child_expr)
if non_const_bindings:
non_const_vars, non_const_exprs = zip(*list(non_const_bindings))
return (LetIn(non_const_vars, non_const_exprs, child_expr)
if not isinstance(child_expr, Const) else child_expr)
else:
return child_expr
elif isinstance(expr, BiMap):
simplified_target_exprs = list(simplify(e) for e in expr.target_exprs)
simplified_source_exprs = list(simplify(e) for e in expr.source_exprs)
simplified_ubs = list(simplify(e) for e in expr.target_upper_bounds)
simplified_lbs = list(simplify(e) for e in expr.target_lower_bounds)
child_expr = simplify(expr.expr)
return BiMap(expr=child_expr,
targets=expr.targets,
target_exprs=simplified_target_exprs,
sources=expr.sources,
source_exprs=simplified_source_exprs,
inv_jacobian=simplify(expr.inv_jacobian),
target_lower_bounds=simplified_lbs,
target_upper_bounds=simplified_ubs)
elif isinstance(expr, Delta):
return Delta(simplify(expr.expr))
elif {'FwdDeriv', 'RevDeriv'} & {t.__name__ for t in type(expr).__mro__}:
return simplify(expr.__getattribute__('deriv_expr'))
elif isinstance(expr, Bool):
left_expr, right_expr = simplify(expr.left_expr), simplify(
expr.right_expr)
if isinstance(left_expr, Const) and isinstance(right_expr, Const):
return false if evaluate(Bool(left_expr,
right_expr)) == 0.0 else true
return Bool(left_expr, right_expr)
elif isinstance(expr, And):
left_expr, right_expr = simplify(expr.left_expr), simplify(
expr.right_expr)
if left_expr == true:
return right_expr
if right_expr == true:
return left_expr
if left_expr == false or right_expr == false:
return false
return And(left_expr, right_expr)
elif isinstance(expr, Or):
left_expr, right_expr = simplify(expr.left_expr), simplify(
expr.right_expr)
if left_expr == false:
return right_expr
if right_expr == false:
return left_expr
if left_expr == true or right_expr == true:
return true
return Or(left_expr, right_expr)
else:
raise ValueError(
f'The type of the expr "{type(expr)}" does not have a supported simplify rule'
)
def outer_fn(e, ctx):
ctx['has_expr'] = any(ctx['has_exprs'])
# Check if we need to handle other such cases.
assert not (ctx['has_expr'] and isinstance(e, SmoothFunc)),\
f'expr is contained in a non-linear function {type(e)}'
if isinstance(e, Add):
if ctx['has_expr']:
assert sum(
ctx['has_exprs']) == 1, 'More than one branch with expr'
ctx['expr'] = ctx['exprs'][ctx['has_exprs'].index(True)]
return ctx['expr'], ctx
else:
ctx['expr'] = e
return e, ctx
elif isinstance(e, IfElse):
if ctx['has_expr']:
assert sum(
ctx['has_exprs']) == 1, 'More than one branch with expr'
if ctx['has_exprs'][1]:
# If block contains expr.
ctx['expr'] = IfElse(e.cond, ctx['exprs'][1], Const(0))
elif ctx['has_exprs'][2]:
ctx['expr'] = IfElse(e.cond, Const(0), ctx['exprs'][2])
else:
assert False, 'condition must not contain expr. expr is not linear'
return ctx['expr'], ctx
elif isinstance(e, Tup):
if ctx['has_expr']:
assert sum(
ctx['has_exprs']) == 1, 'More than one branch with expr'
ctx['expr'] = Tup(*[
ctx['exprs'][idx] if has_expr else Const(0)
for idx, has_expr in enumerate(ctx['has_exprs'])
])
return ctx['expr'], ctx
elif isinstance(e, LetIn):
assert sum(ctx['has_exprs']) <= 1, 'More than one branch with expr'
if any(ctx['has_exprs'][1:]):
# Let expressions contain exprs.
new_exprs = [
let_var for let_var, has_expr in zip(
e.new_vars, ctx['has_exprs'][1:]) if has_expr
]
# Recursively split the body with the new expressions.
s_expr = split_exprs(new_exprs, ctx['let_body'])
let_body = (s_expr if s_expr else Const(0)) +\
(e.expr if ctx['has_exprs'][0] else Const(0))
try:
vs, es = zip(*[(v, e)
for v, e in zip(e.new_vars, e.new_exprs)
if v in let_body])
ctx['expr'] = LetIn(vs, es, let_body)
except ValueError:
# No need for a let expr.
ctx['expr'] = let_body
return ctx['expr'], ctx
ctx['expr'] = e
return ctx['expr'], ctx