def __init__(self, rule_mapping_context, diff_context, diff_inames,
allowed_nonsmoothness=None):
RuleAwareIdentityMapper.__init__(self, rule_mapping_context)
DifferentiationMapper.__init__(
self,
# This is actually ignored because we
# override map_variable below.
variable=None,
allowed_nonsmoothness=None)
self.diff_context = diff_context
self.diff_inames = diff_inames
self.diff_iname_exprs = tuple(var(diname) for diname in diff_inames)
self.function_map = func_map
def map_call(self, expr, *args):
dc = self.diff_context
if expr.function.name in dc.kernel.substitutions:
# FIXME: Deal with subsitution rules
# Need to use chain rule here, too.
raise NotImplementedError("substitution rules in differentiation")
else:
return DifferentiationMapper.map_call(self, expr, *args)
def map_call(self, expr, *args):
dc = self.diff_context
if expr.function.name in dc.kernel.substitutions:
# FIXME: Deal with subsitution rules
# Need to use chain rule here, too.
raise NotImplementedError("substitution rules in differentiation")
else:
return DifferentiationMapper.map_call(self, expr, *args)
def get_phase_step_matrix(self, phase_name, shapes=None, sparse=False):
"""
`shapes` maps variable names to vector lengths.
`sparse` controls whether the output is sparse or dense. When
`sparse=True`, returns a SparseStepMatrix.
Otherwise returns a numpy object array.
"""
if shapes is None:
shapes = {}
components, initial_vals = self.run_symbolic_step(phase_name, shapes)
from pymbolic.mapper.differentiator import DifferentiationMapper
from pymbolic.mapper.dependency import DependencyMapper
dependencies = DependencyMapper()
nv = len(components)
shape = (nv, nv)
if not sparse:
step_matrix = np.zeros(shape, dtype=np.object)
else:
indices = []
data = []
iv_to_index = {iv: i for i, iv in enumerate(initial_vals)}
for i, v in enumerate(components):
# Get the expression for v.
if isinstance(v, self.VectorComponent):
expr = self.context[v.name][v.index]
else:
expr = self.context[v]
# Selectively evaluate the derivative only for components that are
# actually present in the expression. This takes advantage of
# sparsity.
component_vars = dependencies(expr)
for iv in component_vars:
if iv not in iv_to_index:
continue
j = iv_to_index[iv]
entry = DifferentiationMapper(iv)(expr)
if not sparse:
step_matrix[i, j] = entry
else:
indices.append((i, j))
data.append(entry)
if not sparse:
return step_matrix
else:
return SparseStepMatrix(shape, indices, data)
def diff(var):
"""Return the symbolic derivative operator with respect to *var*."""
from pymbolic.mapper.differentiator import DifferentiationMapper
def func_map(arg_index, func, arg, allowed_nonsmoothness):
if func == pmbl.var("sin"):
return pmbl.var("cos")(*arg)
elif func == pmbl.var("cos"):
return -pmbl.var("sin")(*arg)
else:
raise ValueError("Unrecognized function")
return DifferentiationMapper(var, func_map=func_map)
def partial(self, expr):
exprs = []
for var in self.indvars:
exprs.append(simplify(DifferentiationMapper(var)(expr)))
return np.array(exprs)