diff --git a/dagrt/language.py b/dagrt/language.py index b85031ea6006697d36c8c3b43fa98ddfe7fcc0c3..c1257f4f1f06f11860ba80ba7c5498c5b6cb7c55 100644 --- a/dagrt/language.py +++ b/dagrt/language.py @@ -294,23 +294,25 @@ class AssignImplicit(AssignmentBase): """ .. attribute:: assignees - A tuple of strings. The names of the variables to assign with the - results of solving for `solve_variables` + A tuple of strings, the names of the variables that are assigned by this + instruction. The assignees represent solutions to a single + equation or a system of (implicit) equations. - .. attribute:: solve_variables + .. attribute:: input_expressions - A tuple of strings, the names of the variables being solved for + A dictionary mapping strings to (symbolic) expressions that are + evaluated (at runtime) as the inputs to an equation solver. The expressions + in this dictionary are tracked for dependency and transformation purposes. - .. attribute:: expressions + To help organize structured data such as a Butcher tableau, the values in + this dictionary are allowed to contain nested tuples of expressions, though + tuples are not otherwise generally supported in :mod:`dagrt` expressions. - A tuple of expressions, which represent the left hand sides of a - system of (possibly nonlinear) equations. The solver will attempt - to find a simultaneous root of the system. + .. attribute:: params - .. attribute:: other_params - - A dictionary used to pass extra arguments to the solver, for instance - a starting guess + A dictionary mapping strings to (non-symbolic expression) parameters + that describe the system of equations solved by the implementation of + this instruction. .. attribute:: solver_id @@ -318,16 +320,46 @@ class AssignImplicit(AssignmentBase): system. This identifier is intended to match information about solvers which is supplied by the user. + The contents of the dictionaries ``input_expressions`` and ``params`` are + deliberately under-specified and the details are expected to be documented + by the scheme designer. They should contain enough information to be able to + reconstruct the system of equations to be solved. + + As an example, consider a system consisting of a single implicit equation + arising from the solution of the ODE :math:`y' = f(t, y)` using an + Adams-Bashforth method + + .. math:: + + y_{n+1} = y_n - 1/2 h f(t_n, y_n) + 3/2 h f(t_{n+1}, y_{n+1}), + + where :math:`y_{n+1}` is the solved-for variable. + + There are many ways to describe this system in an :class:`AssignImplicit` + instruction, but the following interface is a possibility. ``input_expressions`` + could contain keys *state*, *times*, and *coeffs*, where + + - the value corresponding to *state* is a symbolic expression that evaluates + to :math:`y_n`, + - the value corresponding to *times* is a tuple of symbolic expressions + evaluating to :math:`t_{n}` and :math:`t_{n+1}`, + - the value corresponding to *coeffs* is a tuple of symbolic expressions + evaluating to :math:`-1/2 h` and :math:`3/2h`. + + ``params`` could contain the key *component_id*, whose value associates the + right-hand side function :math:`f` to the solution component (useful in the + case where the method being generated contains different right-hand sides). + """ - def __init__(self, assignees, solve_variables, expressions, other_params, - solver_id, **kwargs): - Statement.__init__(self, assignees=assignees, - solve_variables=solve_variables, - expressions=expressions, - other_params=other_params, - solver_id=solver_id, - **kwargs) + def __init__(self, assignees, input_expressions, params, solver_id, + **kwargs): + Statement.__init__(self, + assignees=assignees, + input_expressions=input_expressions, + params=params, + solver_id=solver_id, + **kwargs) exec_method = six.moves.intern("exec_AssignImplicit") @@ -336,19 +368,13 @@ class AssignImplicit(AssignmentBase): def get_read_variables(self): # Variables can be read by: - # 1. expressions (except for those in solve_variables) - # 2. values in other_params - # 3. condition - from itertools import chain - - def flatten(iter_arg): - return chain(*list(iter_arg)) + # 1. values in input_expressions + # 2. superclass attributes (i.e., condition) variables = super(AssignImplicit, self).get_read_variables() - variables |= set(flatten(get_variables(expr) for expr in self.expressions)) - variables -= set(self.solve_variables) - variables |= set(flatten(get_variables(expr) for expr - in self.other_params.values())) + variables |= set(var + for expr in self.input_expressions.values() + for var in get_variables(expr)) return variables def map_expressions(self, mapper, include_lhs=True): @@ -361,26 +387,52 @@ class AssignImplicit(AssignmentBase): else: assignees = self.assignees + input_expressions = {key: mapper(val) + for key, val in self.input_expressions.items()} + return (super(AssignImplicit, self) .map_expressions(mapper, include_lhs=include_lhs) .copy( assignees=assignees, - expressions=mapper(self.expressions))) + input_expressions=input_expressions)) def __str__(self): + def stringify_tuple(tuple_maybe): + # This differs from str() in that the elements of the tuple are + # stringified with str() rather than repr(). This makes the + # stringified output of nested tuples of symbolic expressions look + # nicer. + if not isinstance(tuple_maybe, tuple): + return str(tuple_maybe) + result = ["("] + for elem in tuple_maybe: + if len(result) >= 2: + result.append(", ") + result.append(stringify_tuple(elem)) + if len(result) == 2: + result.append(",)") + else: + result.append(")") + return "".join(result) + lines = [] + + def format_params(dict_, stringifier): + nonlocal lines + for key, val in dict_.items(): + lines.append(" %s = %s" % (key, stringifier(val))) + lines.append("AssignImplicit") - lines.append("solver_id = " + str(self.solver_id)) - for assignee_index, assignee in enumerate(self.assignees): - lines.append(assignee + " <- " + self.solve_variables[assignee_index]) - lines.append("where") - for expression in self.expressions: - lines.append(" " + str(expression) + " = 0") - if self.other_params: - lines.append("with parameters") - for param_name, param_value in self.other_params.items(): - lines.append(param_name + ": " + str(param_value)) - lines.append(self._condition_printing_suffix()) + lines.append("assignees = %s" % ", ".join(self.assignees)) + lines.append("solver_id = %s" % self.solver_id) + if self.input_expressions: + lines.append("with input expressions:") + format_params(self.input_expressions, stringify_tuple) + if self.params: + lines.append("with params:") + format_params(self.params, str) + if self.condition is not True: + lines.append(self._condition_printing_suffix()) return "\n".join(lines) @@ -1096,18 +1148,50 @@ class CodeBuilder(object): __call__ = assign - def assign_implicit_1(self, assignee, solve_component, expression, guess, - solver_id=None): - """Special case of AssignImplicit when there is 1 component to solve for.""" - self.assign_implicit( - (assignee.name,), (solve_component.name,), (expression,), - {"guess": guess}, solver_id) + def assign_implicit(self, + assignees, input_expressions=None, params=None, solver_id=None): + """Generate code for an implicit assignment. See :class:`AssignImplicit`. + + *assignees* may be a variable, a variable name, or a tuple of either. + """ + assignees_tuple = [] + + if input_expressions is None: + input_expressions = {} + + if params is None: + params = {} + + from pymbolic.primitives import Variable + + if isinstance(assignees, str): + from pymbolic import parse + assignees = parse(assignees) + + for assignee in assignees: + if isinstance(assignee, str): + assignees_tuple.append(assignee) + elif isinstance(assignee, Variable): + assignees_tuple.append(assignee.name) + else: + raise ValueError( + "could not determine assignee name: %s" % assignee) - def assign_implicit(self, assignees, solve_components, expressions, - other_params, solver_id): self._add_statement(AssignImplicit( - assignees, solve_components, - expressions, other_params, solver_id)) + assignees=tuple(assignees_tuple), + input_expressions=input_expressions, + params=params, + solver_id=solver_id)) + + def assign_implicit_1(self, + assignee, input_expressions=None, params=None, solver_id=None): + """Special case of :meth:`CodeBuilder.assign_implicit` for a single assignee. + """ + return self.assign_implicit( + assignees=(assignee,), + input_expressions=input_expressions, + params=params, + solver_id=solver_id) def yield_state(self, expression, component_id, time, time_id): """Yield a value.""" diff --git a/dagrt/version.py b/dagrt/version.py index 5cea8857d9550183c2caabd0948032faac48e0ae..82ea6fb128de4adc8cddd1928ee3eb063bd70aa0 100644 --- a/dagrt/version.py +++ b/dagrt/version.py @@ -1,2 +1,2 @@ -VERSION = (2020, 1) +VERSION = (2020, 2) VERSION_TEXT = ".".join(str(i) for i in VERSION)