def expanded_expression(self, expr):
"""
Return the expanded expression.
The returned expanded expression consists of the original
expression given by it basics objects (States, Parameters and
IndexedObjects)
"""
timer = Timer("Expand expression") # noqa: F841
check_arg(expr, Expression)
# First check cache
exp_expr = self._expanded_expressions.get(expr)
if exp_expr is not None:
return exp_expr
# If not recursively expand the expression
der_subs = {}
subs = {}
for obj in self.expression_dependencies[expr]:
if isinstance(obj, Derivatives):
der_subs[obj.sym] = self.expanded_expression(obj)
elif isinstance(obj, Expression):
subs[obj.sym] = self.expanded_expression(obj)
# Do the substitution
exp_expr = expr.expr.xreplace(der_subs).xreplace(subs)
self._expanded_expressions[expr] = exp_expr
return exp_expr
def store_expressions(expr, new_expr):
"Help function to store new expressions"
timer = Timer( # noqa: F841
f"Store expression while recreating body of {self.name}",
) # noqa: F841
# Update sym replace dict
if isinstance(expr, Derivatives):
der_replace_dict[expr.sym] = new_expr.sym
else:
replace_dict[expr.sym] = new_expr.sym
# Store the new expression for later references
present_ode_objects[expr.name] = new_expr
replaced_expr_map[expr] = new_expr
# Append the new expression
new_body_expressions.append(new_expr)
# Update dependency information
if expr in object_used_in:
for dep in object_used_in[expr]:
if dep in expression_dependencies:
expression_dependencies[dep].remove(expr)
expression_dependencies[dep].add(new_expr)
object_used_in[new_expr] = object_used_in.pop(expr)
if expr in expression_dependencies:
expression_dependencies[new_expr] = expression_dependencies.pop(expr)
def gotran2matlab(filename, params):
"""
Create a matlab code from a gotran model
"""
timer = Timer(f"Generate Matlab code from {filename}") # noqa: F841
# Load Gotran model
ode = load_ode(filename)
gen = MatlabCodeGenerator(params)
info("")
info(f"Generating Matlab files for the {ode.name} ode...")
# Collect monitored
if params.functions.monitored.generate:
monitored = [expr.name for expr in ode.intermediates + ode.state_expressions]
else:
monitored = None
for function_name, code in list(gen.code_dict(ode, monitored).items()):
open(f"{ode.name}_{function_name}.m", "w").write(code)
info(" done.")
def __init__(
self,
jacobian,
function_name="compute_diagonal_jacobian",
result_name="diag_jac",
params=None,
):
"""
Create a DiagonalJacobianComponent
Arguments
---------
jacobian : gotran.JacobianComponent
The Jacobian of the ODE
function_name : str
The name of the function which should be generated
result_name : str (optional)
The basename of the indexed result expression
params : dict
Parameters determining how the code should be generated
"""
check_arg(jacobian, JacobianComponent)
descr = ("Compute the diagonal jacobian of the right hand side of the "
"{0} ODE".format(jacobian.root))
super(DiagonalJacobianComponent, self).__init__(
"DiagonalJacobian",
jacobian.root,
function_name,
descr,
params=params,
)
what = "Computing diagonal jacobian"
timer = Timer(what) # noqa: F841
self.add_comment(what)
N = jacobian.jacobian.shape[0]
self.shapes[result_name] = (N, )
jacobian_name = jacobian.results[0]
# Create IndexExpressions of the diagonal Jacobian
for expr in jacobian.indexed_objects(jacobian_name):
if expr.indices[0] == expr.indices[1]:
self.add_indexed_expression(result_name, expr.indices[0],
expr.expr)
self.diagonal_jacobian = sp.Matrix(N, N, lambda i, j: 0.0)
for i in range(N):
self.diagonal_jacobian[i, i] = jacobian.jacobian[i, i]
# Call recreate body with the jacobian diagonal expressions as the
# result expressions
results = {result_name: self.indexed_objects(result_name)}
results, body_expressions = self._body_from_results(**results)
self.body_expressions = self._recreate_body(body_expressions,
**results)
def __init__(
self,
diagonal_jacobian,
with_body=True,
function_name="compute_diagonal_jacobian_action",
result_name="diag_jac_action",
params=None,
):
"""
Create a DiagonalJacobianActionComponent
Arguments
---------
jacobian : gotran.JacobianComponent
The Jacobian of the ODE
with_body : bool
If true, the body for computing the jacobian will be included
function_name : str
The name of the function which should be generated
result_name : str
The basename of the indexed result expression
params : dict
Parameters determining how the code should be generated
"""
timer = Timer("Computing jacobian action component") # noqa: F841
check_arg(diagonal_jacobian, DiagonalJacobianComponent)
descr = ("Compute the diagonal jacobian action of the right hand side "
"of the {0} ODE".format(diagonal_jacobian.root))
super(DiagonalJacobianActionComponent, self).__init__(
"DiagonalJacobianAction",
diagonal_jacobian.root,
function_name,
descr,
params=params,
)
x = self.root.full_state_vector
jac = diagonal_jacobian.diagonal_jacobian
self._action_vector = sp.Matrix(len(x), 1, lambda i, j: 0)
self.add_comment("Computing the action of the jacobian")
# Create Jacobian matrix
self.shapes[result_name] = (len(x), )
for i, expr in enumerate(jac * x):
self._action_vector[i] = self.add_indexed_expression(
result_name, i, expr)
# Call recreate body with the jacobian action expressions as the
# result expressions
results = {result_name: self.indexed_objects(result_name)}
if with_body:
results, body_expressions = self._body_from_results(**results)
else:
body_expressions = results[result_name]
self.body_expressions = self._recreate_body(body_expressions,
**results)
def _body_from_dependencies(self, **results):
timer = Timer(f"Compute dependencies for {self.name}") # noqa: F841
# Extract all result expressions
result_expressions = sum(list(results.values()), [])
# Check passed expressions
ode_expr_deps = self.root.expression_dependencies
exprs = set(result_expressions)
not_checked = set()
used_states = set()
used_parameters = set()
for expr in result_expressions:
check_arg(
expr,
(Expression, Comment),
context=CodeComponent._body_from_results,
)
if isinstance(expr, Comment):
continue
# Collect dependencies
for obj in ode_expr_deps[expr]:
if isinstance(obj, (Expression, Comment)):
not_checked.add(obj)
elif isinstance(obj, State):
used_states.add(obj)
elif isinstance(obj, Parameter):
used_parameters.add(obj)
# use list to make the order consistent
not_checked_list = sorted(not_checked)
# Collect all dependencies
while len(not_checked_list) > 0:
dep_expr = not_checked_list.pop()
exprs.add(dep_expr)
for obj in ode_expr_deps[dep_expr]:
if isinstance(obj, (Expression, Comment)):
if obj not in exprs:
if obj not in not_checked_list:
not_checked_list.append(obj)
elif isinstance(obj, State):
used_states.add(obj)
elif isinstance(obj, Parameter):
used_parameters.add(obj)
# Sort used state, parameters and expr
self.used_states = sorted(used_states)
self.used_parameters = sorted(used_parameters)
# Return a sorted list of all collected expressions
return results, sorted(exprs)
def gotran2py(filename, params):
"""
Create a c header file from a gotran model
"""
timer = Timer(f"Generate Python code from {filename}")
# Load Gotran model
ode = load_ode(filename)
# Collect monitored
if params.functions.monitored.generate:
monitored = [
expr.name for expr in ode.intermediates + ode.state_expressions
]
else:
monitored = None
# Create a Python code generator
gen = PythonCodeGenerator(
params,
ns=params.namespace,
import_inside_functions=params.import_inside_functions,
)
output = params.output
if output:
if not output.endswith(".py"):
output += ".py"
else:
output = filename.replace(".ode", "") + ".py"
info("")
info(f"Generating Python code for the {ode.name} ode...")
if params.class_code:
code = gen.class_code(ode, monitored=monitored)
else:
code = gen.module_code(ode, monitored=monitored)
info(" done.")
with open(output, "w") as f:
f.write(code)
del timer
if params.list_timings:
list_timings()
def add_derivative(self, der_expr, dep_var, expr, dependent=None):
"""
Add a derivative expression
Arguments
---------
der_expr : gotran.Expression, gotran.State, sympy.AppliedUndef
The Expression or State which is differentiated
dep_var : gotran.State, gotran.Time, gotran.Expression, sympy.AppliedUndef, sympy.Symbol
The dependent variable
expr : sympy.Basic
The expression which the differetiation should be equal
dependent : gotran.ODEObject
If given the count of this expression will follow as a
fractional count based on the count of the dependent object
"""
timer = Timer("Add derivatives") # noqa: F841
if isinstance(der_expr, AppliedUndef):
name = sympycode(der_expr)
der_expr = self.root.present_ode_objects.get(name)
if der_expr is None:
error(f"{name} is not registered in this ODE")
if isinstance(dep_var, (AppliedUndef, sp.Symbol)):
name = sympycode(dep_var)
dep_var = self.root.present_ode_objects.get(name)
if dep_var is None:
error(f"{name} is not registered in this ODE")
# Check if der_expr is a State
if isinstance(der_expr, State):
self._expect_state(der_expr)
obj = StateDerivative(der_expr, expr, dependent)
else:
# Create a DerivativeExpression in the present component
obj = DerivativeExpression(der_expr, dep_var, expr, dependent)
self._register_component_object(obj, dependent)
return obj.sym
def add_indexed_object(self, basename, indices, add_offset=False):
"""
Add an indexed object using a basename and the indices
Arguments
---------
basename : str
The basename of the indexed expression
indices : int, tuple of int
The fixed indices identifying the expression
add_offset : bool
Add offset to indices
"""
timer = Timer("Add indexed object") # noqa: F841
indices = tuplewrap(indices)
# Check that provided indices fit with the registered shape
if len(self.shapes[basename]) > len(indices):
error(
"Shape mismatch between indices {0} and registered "
"shape for {1}{2}".format(indices, basename, self.shapes[basename]),
)
for dim, (index, shape_ind) in enumerate(zip(indices, self.shapes[basename])):
if index >= shape_ind:
error(
"Indices must be smaller or equal to the shape. Mismatch "
"in dim {0}: {1}>={2}".format(dim + 1, index, shape_ind),
)
# Create IndexedObject
obj = IndexedObject(
basename,
indices,
self.shapes[basename],
self._params.array,
add_offset,
)
self._register_component_object(obj)
# Return the sympy version of the object
return obj.sym
def add_parameter(self, name, init):
"""
Add a parameter to the component
Arguments
---------
name : str
The name of the parameter
init : scalar, modelparameters.ScalarParam
The initial value of the parameter
"""
timer = Timer("Add parameters") # noqa: F841
param = Parameter(name, init)
self._register_component_object(param)
# Return the sympy version of the state
return param.sym
def add_state(self, name, init):
"""
Add a state to the component
Arguments
---------
name : str
The name of the state variable
init : scalar, modelparameters.ScalarParam
The initial value of the state
"""
timer = Timer("Add states") # noqa: F841
# Create state
state = State(name, init, self.time)
self._register_component_object(state)
# Return the sympy version of the state
return state.sym
def gotran2cpp(filename, params):
"""
Create a C++ header file from a gotran model
"""
timer = Timer(f"Generate C++ code from {filename}")
# Load Gotran model
ode = load_ode(filename)
# Create a C code generator
gen = CppCodeGenerator(params)
output = params.output
if output:
if not (output.endswith(".cpp") or output.endswith(".h")):
output += ".h"
else:
output = filename.replace(".ode", "") + ".h"
info("")
info(f"Generating C++ code for the {ode.name} ode...")
if params.class_code:
code = gen.class_code(ode)
else:
code = gen.module_code(ode)
info(" done.")
with open(output, "w") as f:
if params.system_headers:
f.write("#include <cmath>\n")
f.write("#include <cstring>\n")
f.write("#include <stdexcept>\n")
f.write(code)
del timer
if params.list_timings:
list_timings()
def add_intermediate(self, name, expr, dependent=None):
"""
Register an intermediate math expression
Arguments
---------
name : str
The name of the expression
expr : sympy.Basic, scalar
The expression
dependent : gotran.ODEObject
If given the count of this expression will follow as a
fractional count based on the count of the dependent object
"""
# Create an Intermediate in the present component
timer = Timer("Add intermediate") # noqa: F841
expr = Intermediate(name, expr, dependent)
self._register_component_object(expr, dependent)
return expr.sym
def gotran2julia(filename, params):
"""
Create a c header file from a gotran model
"""
timer = Timer(f"Generate Julia code from {filename}")
# Load Gotran model
ode = load_ode(filename)
# Collect monitored
if params.functions.monitored.generate:
monitored = [
expr.name for expr in ode.intermediates + ode.state_expressions
]
else:
monitored = None
# Create a C code generator
gen = JuliaCodeGenerator(params)
output = params.output
if output:
if not output.endswith(".jl"):
output += ".jl"
else:
output = filename.replace(".ode", "") + ".jl"
info("")
info(f"Generating Julia code for the {ode.name} ode...")
code = gen.module_code(ode, monitored=monitored)
info(" done.")
with open(output, "w") as f:
f.write(code)
del timer
if params.list_timings:
list_timings()
def add_indexed_expression(
self,
basename,
indices,
expr,
add_offset=False,
dependent=None,
enum=None,
):
"""
Add an indexed expression using a basename and the indices
Arguments
---------
basename : str
The basename of the indexed expression
indices : int, tuple of int
The fixed indices identifying the expression
expr : sympy.Basic, scalar
The expression.
add_offset : bool
Add offset to indices
dependent : gotran.ODEObject
If given the count of this expression will follow as a
fractional count based on the count of the dependent object
"""
# Create an IndexedExpression in the present component
timer = Timer("Add indexed expression") # noqa: F841
indices = tuplewrap(indices)
# Check that provided indices fit with the registered shape
if len(self.shapes[basename]) > len(indices):
error(
"Shape mismatch between indices {0} and registered "
"shape for {1}{2}".format(indices, basename, self.shapes[basename]),
)
for dim, (index, shape_ind) in enumerate(zip(indices, self.shapes[basename])):
if index >= shape_ind:
error(
"Indices must be smaller or equal to the shape. Mismatch "
"in dim {0}: {1}>={2}".format(dim + 1, index, shape_ind),
)
# Create the indexed expression
if enum is None:
expr = IndexedExpression(
basename,
indices,
expr,
self.shapes[basename],
self._params.array,
add_offset,
dependent,
enum=enum,
)
elif isinstance(enum, State):
state = enum
expr = StateIndexedExpression(
basename,
indices,
expr,
state,
self.shapes[basename],
self._params.array,
add_offset,
dependent,
)
elif isinstance(enum, Parameter):
parameter = enum
expr = ParameterIndexedExpression(
basename,
indices,
expr,
parameter,
self.shapes[basename],
self._params.array,
add_offset,
dependent,
)
else:
error(f"enum must be State or Parameter. Was {type(enum)}")
self._register_component_object(expr, dependent)
return expr.sym
def register_ode_object(self, obj, comp, dependent=None):
"""
Register an ODE object in the root ODEComponent
"""
from modelparameters.sympytools import symbols_from_expr
if self._is_finalized_ode and isinstance(obj, StateExpression):
error("Cannot register a StateExpression, the ODE is finalized")
# Check for existing object in the ODE
dup_obj = self.present_ode_objects.get(obj.name)
# If object with same name is already registered in the ode we
# need to figure out what to do
if dup_obj:
try:
dup_comp = self.object_component[dup_obj]
except KeyError:
dup_comp = None
# If a state is substituted by a state solution
if isinstance(dup_obj, State) and isinstance(obj, StateSolution):
debug(f"Reduce state '{dup_obj}' to {obj.expr}")
# If duplicated object is an ODE Parameter and the added object is
# either a State or a Parameter we replace the Parameter.
elif (isinstance(dup_obj, Parameter) and dup_comp == self
and comp != self and isinstance(obj, (State, Parameter))):
timer = Timer("Replace objects") # noqa: F841
# Remove the object
self.ode_objects.remove(dup_obj)
# FIXME: Do we need to recreate all expression the objects is used in?
# Replace the object from the object_used_in dict and update
# the correponding expressions
subs = {dup_obj.sym: obj.sym}
subs = {}
# Recursively replace object dependencies
self._replace_object(dup_obj, obj, subs)
# for expr in self.object_used_in[dup_obj]:
# updated_expr = recreate_expression(expr, subs)
# self.object_used_in[obj].add(updated_expr)
#
# # Exchange and update the dependencies
# self.expression_dependencies[expr].remove(dup_obj)
# self.expression_dependencies[expr].add(obj)
#
# # FIXME: Do not remove the dependencies
# #self.expression_dependencies[updated_expr] = \
# # self.expression_dependencies.pop(expr)
# self.expression_dependencies[updated_expr] = \
# self.expression_dependencies[expr]
#
# # Find the index of old expression and exchange it with updated
# old_comp = self.object_component[expr]
# ind = old_comp.ode_objects.index(expr)
# old_comp.ode_objects[ind] = updated_expr
#
## Remove information about the replaced objects
# self.object_used_in.pop(dup_obj)
# If duplicated object is an ODE Parameter and the added
# object is an Intermediate we raise an error.
elif (isinstance(dup_obj, Parameter) and dup_comp == self
and isinstance(obj, Expression)):
error(
"Cannot replace an ODE parameter with an Expression, "
"only with Parameters and States.", )
# If State, Parameter or DerivativeExpression we always raise an error
elif any(
isinstance(
oo,
(
State,
Parameter,
Time,
Dt,
DerivativeExpression,
AlgebraicExpression,
StateSolution,
),
) for oo in [dup_obj, obj]):
error(
"Cannot register {0}. A {1} with name '{2}' is "
"already registered in this ODE.".format(
type(obj).__name__,
type(dup_obj).__name__,
dup_obj.name,
), )
else:
# Sanity check that both obj and dup_obj are Expressions
assert all(
isinstance(oo, (Expression)) for oo in [dup_obj, obj])
# Get list of duplicated objects or an empy list
dup_objects = self.duplicated_expressions[obj.name]
if len(dup_objects) == 0:
dup_objects.append(dup_obj)
dup_objects.append(obj)
# Update global information about ode object
self.present_ode_objects[obj.name] = obj
self.object_component[obj] = comp
self.ns.update({obj.name: obj.sym})
# If Expression
if isinstance(obj, Expression):
# Append the name to the list of all ordered components with
# expressions. If the ODE is finalized we do not update components
if not self._is_finalized_ode:
self._handle_expr_component(comp, obj)
# Expand and add any derivatives in the expressions
expression_added = False
replace_dict = {}
derivative_expression_list = list(obj.expr.atoms(sp.Derivative))
derivative_expression_list.sort(key=lambda e: e.sort_key())
for der_expr in derivative_expression_list:
expression_added |= self._expand_single_derivative(
comp,
obj,
der_expr,
replace_dict,
dependent,
)
# If expressions need to be re-created
if replace_dict:
obj.replace_expr(replace_dict)
# If any expression was added we need to bump the count of the ODEObject
if expression_added:
obj._recount(dependent=dependent)
# Add dependencies between the last registered comment and
# expressions so they are carried over in Code components
if comp._local_comments:
self.object_used_in[comp._local_comments[-1]].add(obj)
self.expression_dependencies[obj].add(comp._local_comments[-1])
# Get expression dependencies
for sym in symbols_from_expr(obj.expr, include_derivatives=True):
dep_obj = self.present_ode_objects[sympycode(sym)]
if dep_obj is None:
error(
"The symbol '{0}' is not declared within the '{1}' "
"ODE.".format(sym, self.name), )
# Store object dependencies
self.expression_dependencies[obj].add(dep_obj)
self.object_used_in[dep_obj].add(obj)
# If the expression is a StateSolution the state cannot have
# been used previously
if isinstance(obj, StateSolution) and self.object_used_in.get(
obj.state):
used_in = self.object_used_in.get(obj.state)
error(
"A state solution cannot have been used in "
"any previous expressions. {0} is used in: {1}".format(
obj.state,
used_in,
), )
def save(self, basename=None):
"""
Save ODE to file
Arguments
---------
basename : str (optional)
The basename of the file which the ode will be saved to, if not
given the basename will be the same as the name of the ode.
"""
timer = Timer("Save " + self.name) # noqa: F841
if not self._is_finalized_ode:
error("ODE need to be finalized to be saved to file.")
lines = ["# Saved Gotran model"]
comp_names = dict()
basename = basename or self.name
for comp in self.components:
if comp == self:
comp_name = ""
else:
present_comp = comp
comps = [present_comp.name]
while present_comp.parent != self:
present_comp = present_comp.parent
comps.append(present_comp.name)
comp_name = ", ".join(f'"{name}"' for name in reversed(comps))
comp_names[comp] = comp_name
states = [
f"{obj.name}={obj.param.repr(include_name=False)},"
for obj in comp.ode_objects if isinstance(obj, State)
]
parameters = [
f"{obj.name}={obj.param.repr(include_name=False)},"
for obj in comp.ode_objects if isinstance(obj, Parameter)
]
if states:
lines.append("")
if comp_name:
lines.append(f"states({comp_name},")
else:
lines.append(f"states({states.pop(0)}")
for state_code in states:
lines.append(" " + state_code)
lines[-1] = lines[-1][:-1] + ")"
if parameters:
lines.append("")
if comp_name:
lines.append(f"parameters({comp_name},")
else:
lines.append(f"parameters({parameters.pop(0)}")
for param_code in parameters:
lines.append(" " + param_code)
lines[-1] = lines[-1][:-1] + ")"
# Iterate over all components
for comp_name in self.all_expr_components_ordered:
comp = self.all_components[comp_name]
comp_comment = f"Expressions for the {comp.name} component"
# Iterate over all objects of the component and save only expressions
# and comments
for obj in comp.ode_objects:
# If saving an expression
if isinstance(obj, Expression):
# If the component is a Markov model
if comp.rates:
# Do not save State derivatives
if isinstance(obj, StateDerivative):
continue
# Save rate expressions slightly different
elif isinstance(obj, RateExpression):
lines.append(
"rates[{0}, {1}] = {2}".format(
sympycode(obj.states[0]),
sympycode(obj.states[1]),
sympycode(obj.expr),
), )
continue
# All other Expressions
lines.append(f"{obj.name} = {sympycode(obj.expr)}")
# If saving a comment
elif isinstance(obj, Comment):
# If comment is component comment
if str(obj) == comp_comment:
lines.append("")
comp_name = (comp_names[comp]
if comp_names[comp] else f'"{basename}"')
lines.append(f"expressions({comp_name})")
# Just add the comment
else:
lines.append("")
lines.append(f'comment("{obj}")')
lines.append("")
# Use Python code generator to indent outputted code
# Write to file
from gotran.codegeneration.codegenerators import PythonCodeGenerator
with open(basename + ".ode", "w") as f:
f.write("\n".join(
PythonCodeGenerator.indent_and_split_lines(lines)))
def import_ode(self, ode, prefix="", components=None, **arguments):
"""
Import a Model into the present Model
Arguments
---------
ode : str, gotran.ODE
The ode which should be added. If ode is a str an
ODE stored in that file will be loaded. If it is an ODE it will be
added directly to the present ODE.
prefix : str (optional)
A prefix which all state, parameters and intermediates are
prefixed with.
components : list, tuple of str (optional)
A list of components which will either be extracted or excluded
from the imported model. If not given the whole ODE will be imported.
arguments : dict (optional)
Optional arguments which can control loading of model
"""
timer = Timer("Import ode") # noqa: F841
components = components or []
check_arg(ode, (str, Path, ODE), 0, context=ODE.import_ode)
check_arg(prefix, str, 1, context=ODE.import_ode)
check_arg(components, list, 2, context=ODE.import_ode, itemtypes=str)
# If ode is given directly
if isinstance(ode, (str, Path)):
# If not load external ODE
from gotran.model.loadmodel import load_ode
ode = load_ode(ode, **arguments)
# Postfix prefix with "_" if prefix is not ""
if prefix and prefix[-1] != "_":
prefix += "_"
# If extracting only a certain components
if components:
ode = ode.extract_components(ode.name, *components)
# Subs for expressions
subs = {}
old_new_map = {ode: self}
def add_comp_and_children(added, comp):
"Help function to recursively add components to ode"
# Add states and parameters
for obj in comp.ode_objects:
# Check if obj already excists as a ODE parameter
old_obj = self.present_ode_objects.get(str(obj))
if old_obj and self.object_component[old_obj] == self:
new_name = obj.name
else:
new_name = prefix + obj.name
if isinstance(obj, State):
subs[obj.sym] = added.add_state(new_name, obj.param)
elif isinstance(obj, Parameter):
# If adding an ODE parameter
if comp == ode:
# And parameter name already excists in the present ODE
if str(obj) in self.present_ode_objects:
# Skip it and add the registered symbol for
# substitution
subs[obj.sym] = self.present_ode_objects[str(
obj)].sym
else:
# If not already present just add unprefixed name
subs[obj.sym] = added.add_parameter(
obj.name, obj.param)
else:
subs[obj.sym] = added.add_parameter(
new_name, obj.param)
# Add child components
for child in list(comp.children.values()):
added_child = added.add_component(child.name)
# Get corresponding component in present ODE
old_new_map[child] = added_child
add_comp_and_children(added_child, child)
# Recursively add states and parameters
add_comp_and_children(self, ode)
# Iterate over all components to add expressions
for comp_name in ode.all_expr_components_ordered:
comp = ode.all_components[comp_name]
comp_comment = f"Expressions for the {comp.name} component"
# Get corresponding component in new ODE
added = old_new_map[comp]
# Iterate over all objects of the component and save only expressions
# and comments
for obj in comp.ode_objects:
# If saving an expression
if isinstance(obj, Expression):
# The new sympy expression
new_expr = obj.expr.xreplace(subs)
# If the component is a Markov model
if comp.rates:
# Do not save State derivatives
if isinstance(obj, StateDerivative):
continue
# Save rate expressions slightly different
elif isinstance(obj, RateExpression):
states = obj.states
added._add_single_rate(
subs[states[0].sym],
subs[states[1].sym],
new_expr,
)
continue
# If no prefix we just add the expression by using setattr
# magic
if prefix == "":
setattr(added, str(obj), new_expr)
subs[obj.sym] = added.ode_objects.get(obj.name).sym
elif isinstance(obj, (StateExpression, StateSolution)):
state = subs[obj.state.sym]
if isinstance(obj, AlgebraicExpression):
subs[obj.sym] = added.add_algebraic(
state, new_expr)
elif isinstance(obj, StateDerivative):
subs[obj.sym] = added.add_derivative(
state,
added.t,
new_expr,
)
elif isinstance(obj, StateSolution):
subs[obj.sym] = added.add_state_solution(
state, new_expr)
print(repr(obj), obj.sym)
else:
error("Should not reach here...")
# Derivatives are tricky. Here the der expr and dep var
# need to be registered in the ODE already. But they can
# be registered with and without prefix, so we need to
# check that.
elif isinstance(obj, Derivatives):
# Get der_expr
der_expr = self.root.present_ode_objects.get(
obj.der_expr.name)
if der_expr is None:
# If not found try prefixed version
der_expr = self.root.present_ode_objects.get(
prefix + obj.der_expr.name, )
if der_expr is None:
if prefix:
error(
"Could not find expression: "
"({0}){1} while adding "
"derivative".format(
prefix, obj.der_expr), )
else:
error(
"Could not find expression: "
"{0} while adding derivative".format(
obj.der_expr, ), )
dep_var = self.root.present_ode_objects.get(
obj.dep_var.name)
if isinstance(dep_var, Time):
dep_var = self._time
elif dep_var is None:
# If not found try prefixed version
dep_var = self.root.present_ode_objects.get(
prefix + obj.dep_var.name, )
if dep_var is None:
if prefix:
error(
"Could not find expression: "
"({0}){1} while adding "
"derivative".format(
prefix, obj.dep_var), )
else:
error(
"Could not find expression: "
"{0} while adding derivative".format(
obj.dep_var, ), )
subs[obj.sym] = added.add_derivative(
der_expr,
dep_var,
new_expr,
)
elif isinstance(obj, Intermediate):
subs[obj.sym] = added.add_intermediate(
prefix + obj.name,
new_expr,
)
else:
error("Should not reach here!")
# If saving a comment
elif isinstance(obj, Comment) and str(obj) != comp_comment:
added.add_comment(str(obj))
def __init__(
self,
ode,
function_name="forward_rush_larsen",
delta=1e-8,
params=None,
):
"""
Create a RushLarsen Solver component
Arguments
---------
ode : gotran.ODE
The parent component of this ODEComponent
function_name : str
The name of the function which should be generated
delta : float
Value to safeguard the evaluation of the Rush-Larsen step.
params : dict
Parameters determining how the code should be generated
"""
timer = Timer("Create RushLarsen expressions")
check_arg(ode, ODE)
if ode.is_dae:
error("Cannot generate a Rush-Larsen forward step for a DAE.")
# Call base class using empty result_expressions
descr = f"Compute a forward step using the Rush-Larsen scheme to the {ode} ODE"
super(RushLarsen, self).__init__(
"RushLarsen",
ode,
function_name,
descr,
params=params,
additional_arguments=["dt"],
)
# Recount the expressions if representation of states are "array" as
# then the method is not full explcit
recount = self._params.states.representation != "array"
# Gather state expressions and states
state_exprs = self.root.state_expressions
states = self.root.full_states
result_name = self._params.states.array_name
self.shapes[result_name] = (len(states), )
# Get time step and start creating the update algorithm
if self._params.states.add_offset:
offset_str = f"{result_name}_offset"
else:
offset_str = ""
might_take_time = len(states) >= 10
if might_take_time:
info(
f"Calculating derivatives of {ode.name}. Might take some time..."
)
sys.stdout.flush()
dt = self.root._dt.sym
for i, expr in enumerate(state_exprs):
dependent = expr if recount else None
# Diagonal jacobian value
time_diff = Timer("Differentiate state_expressions for RushLarsen")
expr_diff = expr.expr.diff(expr.state.sym)
del time_diff
# print expr.state.sym, expr_diff, expr_diff.args
if expr_diff and expr.state.sym not in expr_diff.args:
linearized_name = expr.name + "_linearized"
linearized = self.add_intermediate(
linearized_name,
expr_diff,
dependent=dependent,
)
need_zero_div_check = not fraction_numerator_is_nonzero(
expr_diff)
if not need_zero_div_check:
debug(
f"{linearized_name} cannot be zero. Skipping zero division check",
)
RL_term = expr.sym / linearized * (sp.exp(linearized * dt) - 1)
if need_zero_div_check:
RL_term = Conditional(
abs(linearized) > delta,
RL_term,
dt * expr.sym,
)
# Solve "exact" using exp
self.add_indexed_expression(
result_name,
(i, ),
expr.state.sym + RL_term,
offset_str,
dependent=dependent,
enum=expr.state,
)
else:
# Explicit Euler step
self.add_indexed_expression(
result_name,
(i, ),
expr.state.sym + dt * expr.sym,
offset_str,
dependent=dependent,
enum=expr.state,
)
if might_take_time:
info(" done")
# Call recreate body with the solver expressions as the result
# expressions
del timer
results = {result_name: self.indexed_objects(result_name)}
results, body_expressions = self._body_from_results(**results)
self.body_expressions = self._recreate_body(body_expressions,
**results)
def __init__(
self,
ode,
function_name="compute_jacobian",
result_name="jac",
params=None,
):
"""
Create a JacobianComponent
Arguments
---------
ode : gotran.ODE
The parent component of this ODEComponent
function_name : str
The name of the function which should be generated
result_name : str
The name of the variable storing the jacobian result
params : dict
Parameters determining how the code should be generated
"""
check_arg(ode, ODE)
# Call base class using empty result_expressions
descr = f"Compute the jacobian of the right hand side of the {ode} ODE"
super(JacobianComponent, self).__init__(
"Jacobian",
ode,
function_name,
descr,
params=params,
)
check_arg(result_name, str)
timer = Timer("Computing jacobian") # noqa:F841
# Gather state expressions and states
state_exprs = self.root.state_expressions
states = self.root.full_states
# Create Jacobian matrix
N = len(states)
self.jacobian = sp.Matrix(N, N, lambda i, j: 0.0)
self.num_nonzero = 0
self.shapes[result_name] = (N, N)
state_dict = dict((state.sym, ind) for ind, state in enumerate(states))
time_sym = states[0].time.sym
might_take_time = N >= 10
if might_take_time:
info(
f"Calculating Jacobian of {ode.name}. Might take some time...")
sys.stdout.flush()
for i, expr in enumerate(state_exprs):
states_syms = sorted(
(state_dict[sym], sym)
for sym in ode_primitives(expr.expr, time_sym)
if sym in state_dict)
self.add_comment(
f"Expressions for the sparse jacobian of state {expr.state.name}",
dependent=expr,
)
for j, sym in states_syms:
time_diff = Timer("Differentiate state_expressions")
jac_ij = expr.expr.diff(sym)
del time_diff
self.num_nonzero += 1
jac_ij = self.add_indexed_expression(
result_name,
(i, j),
jac_ij,
dependent=expr,
)
self.jacobian[i, j] = jac_ij
if might_take_time:
info(" done")
# Call recreate body with the jacobian expressions as the result
# expressions
results = {result_name: self.indexed_objects(result_name)}
results, body_expressions = self._body_from_results(**results)
self.body_expressions = self._recreate_body(body_expressions,
**results)
def __init__(self, jacobian, function_name="lu_factorize", params=None):
"""
Create a FactorizedJacobianComponent
Arguments
---------
jacobian : gotran.JacobianComponent
The Jacobian of the ODE
function_name : str
The name of the function which should be generated
params : dict
Parameters determining how the code should be generated
"""
timer = Timer("Computing factorization of jacobian") # noqa: F841
check_arg(jacobian, JacobianComponent)
descr = f"Symbolically factorize the jacobian of the {jacobian.root} ODE"
super(FactorizedJacobianComponent, self).__init__(
"FactorizedJacobian",
jacobian.root,
function_name,
descr,
params=params,
use_default_arguments=False,
additional_arguments=jacobian.results,
)
self.add_comment(f"Factorizing jacobian of {self.root.name}")
jacobian_name = jacobian.results[0]
# Recreate jacobian using only sympy Symbols
jac_orig = jacobian.jacobian
# Size of system
n = jac_orig.rows
jac = sp.Matrix(n, n, lambda i, j: sp.S.Zero)
for i in range(n):
for j in range(n):
# print jac_orig[i,j]
if not jac_orig[i, j].is_zero:
name = sympycode(jac_orig[i, j])
jac[i, j] = sp.Symbol(
name,
real=True,
imaginary=False,
commutative=True,
hermitian=True,
complex=True,
)
print(jac[i, j])
p = []
self.shapes[jacobian_name] = (n, n)
def add_intermediate_if_changed(jac, jac_ij, i, j):
# If item has changed
if jac_ij != jac[i, j]:
print("jac", i, j, jac_ij)
jac[i,
j] = self.add_indexed_expression(jacobian_name, (i, j),
jac_ij)
# Do the factorization
for j in range(n):
for i in range(j):
# Get sympy expr of A_ij
jac_ij = jac[i, j]
# Build sympy expression
for k in range(i):
jac_ij -= jac[i, k] * jac[k, j]
add_intermediate_if_changed(jac, jac_ij, i, j)
pivot = -1
for i in range(j, n):
# Get sympy expr of A_ij
jac_ij = jac[i, j]
# Build sympy expression
for k in range(j):
jac_ij -= jac[i, k] * jac[k, j]
add_intermediate_if_changed(jac, jac_ij, i, j)
# find the first non-zero pivot, includes any expression
if pivot == -1 and jac[i, j]:
pivot = i
if pivot < 0:
# this result is based on iszerofunc's analysis of the
# possible pivots, so even though the element may not be
# strictly zero, the supplied iszerofunc's evaluation gave
# True
error("No nonzero pivot found; symbolic inversion failed.")
if pivot != j: # row must be swapped
jac.row_swap(pivot, j)
p.append([pivot, j])
print("Pivoting!!")
# Scale with diagonal
if not jac[j, j]:
error("Diagonal element of the jacobian is zero. "
"Inversion failed")
scale = 1 / jac[j, j]
for i in range(j + 1, n):
# Get sympy expr of A_ij
jac_ij = jac[i, j]
jac_ij *= scale
add_intermediate_if_changed(jac, jac_ij, i, j)
# Store factorized jacobian
self.factorized_jacobian = jac
self.num_nonzero = sum(not jac[i, j].is_zero for i in range(n)
for j in range(n))
# No need to call recreate body expressions
self.body_expressions = self.ode_objects
self.used_states = set()
self.used_parameters = set()
def _body_from_cse(self, **results):
timer = Timer(f"Compute common sub expressions for {self.name}") # noqa: F841
(
orig_result_expressions,
result_names,
expanded_result_exprs,
) = self._expanded_result_expressions(**results)
state_offset = self._params["states"]["add_offset"]
# Collect results and body_expressions
body_expressions = []
new_results = defaultdict(list)
might_take_time = len(orig_result_expressions) >= 40
if might_take_time:
info(
"Computing common sub expressions for {0}. Might take "
"some time...".format(self.name),
)
sys.stdout.flush()
# Call sympy common sub expression reduction
cse_exprs, cse_result_exprs = cse(
expanded_result_exprs,
symbols=sp.numbered_symbols("cse_"),
optimizations=[],
)
# Map the cse_expr to an OrderedDict
cse_exprs = OrderedDict(cse_expr for cse_expr in cse_exprs)
# Extract the symbols into a set for fast comparison
cse_syms = set((sym for sym in cse_exprs))
# Create maps between cse_expr and result expressions trying
# to optimized the code by weaving in the result expressions
# in between the cse_expr
# A map between result expr and name and indices so we can
# construct IndexedExpressions
result_expr_map = defaultdict(list)
# A map between last cse_expr used in a particular result expr
# so that we can put the result expression right after the
# last cse_expr it uses.
last_cse_expr_used_in_result_expr = defaultdict(list)
# Result expressions that does not contain any cse_sym
result_expr_without_cse_syms = []
# A map between cse_sym and its substitutes
cse_subs = {}
for ind, (orig_result_expr, result_expr) in enumerate(
zip(orig_result_expressions, cse_result_exprs),
):
# Collect information so that we can recreate the result
# expression from
result_expr_map[result_expr].append(
(
result_names[orig_result_expr],
orig_result_expr.indices
if isinstance(orig_result_expr, IndexedExpression)
else ind,
),
)
# If result_expr does not contain any cse_sym
if not any(cse_sym in cse_syms for cse_sym in result_expr.atoms()):
result_expr_without_cse_syms.append(result_expr)
else:
# Get last cse_sym used in result expression
last_cse_sym = sorted(
(cse_sym for cse_sym in result_expr.atoms() if cse_sym in cse_syms),
key=cmp_to_key(lambda a, b: cmp(int(a.name[4:]), int(b.name[4:]))),
)[-1]
if result_expr not in last_cse_expr_used_in_result_expr[last_cse_sym]:
last_cse_expr_used_in_result_expr[last_cse_sym].append(result_expr)
debug(
"Found {0} result expressions without any cse_syms.".format(
len(result_expr_without_cse_syms),
),
)
# print ""
# print "LAST cse_syms:", last_cse_expr_used_in_result_expr.keys()
cse_cnt = 0
atoms = [state.sym for state in self.root.full_states]
atoms.extend(param.sym for param in self.root.parameters)
# Collecte what states and parameters has been used
used_states = set()
used_parameters = set()
self.add_comment(
"Common sub expressions for the body and the " "result expressions",
)
body_expressions.append(self.ode_objects[-1])
# Register the common sub expressions as Intermediates
for cse_sym, expr in list(cse_exprs.items()):
# print cse_sym, expr
# If the expression is just one of the atoms of the ODE we
# skip the cse expressions but add a subs for the atom We
# also skip Relationals and Piecewise as the type checking
# in Piecewise otherwise kicks in and destroys things for
# us.
if expr in atoms or isinstance(
expr,
(sp.Piecewise, sp.relational.Relational, sp.relational.Boolean),
):
cse_subs[cse_sym] = expr.xreplace(cse_subs)
else:
# Add body expression as an intermediate expression
sym = self.add_intermediate(f"cse_{cse_cnt}", expr.xreplace(cse_subs))
obj = self.ode_objects.get(sympycode(sym))
for dep in self.root.expression_dependencies[obj]:
if isinstance(dep, State):
used_states.add(dep)
elif isinstance(dep, Parameter):
used_parameters.add(dep)
cse_subs[cse_sym] = sym
cse_cnt += 1
body_expressions.append(obj)
# Check if we should add a result expressions
if last_cse_expr_used_in_result_expr[cse_sym]:
# Iterate over all registered result expr for this cse_sym
for result_expr in last_cse_expr_used_in_result_expr.pop(cse_sym):
for result_name, indices in result_expr_map[result_expr]:
# Replace pure state and param expressions
# print cse_subs, result_expr
exp_expr = result_expr.xreplace(cse_subs)
sym = self.add_indexed_expression(
result_name,
indices,
exp_expr,
add_offset=state_offset,
)
expr = self.ode_objects.get(sympycode(sym))
for dep in self.root.expression_dependencies[expr]:
if isinstance(dep, State):
used_states.add(dep)
elif isinstance(dep, Parameter):
used_parameters.add(dep)
# Register the new result expression
new_results[result_name].append(expr)
body_expressions.append(expr)
if might_take_time:
info(" done")
# Sort used state, parameters and expr
self.used_states = sorted(used_states)
self.used_parameters = sorted(used_parameters)
return new_results, body_expressions
def _recreate_body(self, body_expressions, **results):
"""
Create body expressions based on the given result_expressions
In this method are all expressions replaced with something that should
be used to generate code. The parameters in:
parameters["generation"]["code"]
decides how parameters, states, body expressions and indexed expressions
are represented.
"""
if not (results or body_expressions):
return
for result_name, result_expressions in list(results.items()):
check_kwarg(
result_expressions,
result_name,
list,
context=CodeComponent._recreate_body,
itemtypes=(Expression, Comment),
)
# Extract all result expressions
result_expressions = sum(list(results.values()), [])
# A map between result expression and result name
result_names = dict(
(result_expr, result_name)
for result_name, result_exprs in list(results.items())
for result_expr in result_exprs
)
timer = Timer(f"Recreate body expressions for {self.name}") # noqa: F841
# Initialize the replace_dictionaries
replace_dict = self.param_state_replace_dict
der_replace_dict = {}
# Get a copy of the map of where objects are used in and their
# present dependencies so any updates done in these dictionaries does not
# affect the original dicts
object_used_in = defaultdict(set)
for expr, used in list(self.root.object_used_in.items()):
object_used_in[expr].update(used)
expression_dependencies = defaultdict(set)
for expr, deps in list(self.root.expression_dependencies.items()):
expression_dependencies[expr].update(deps)
# Get body parameters
body_repr = self._params["body"]["representation"]
optimize_exprs = self._params["body"]["optimize_exprs"]
# Set body related variables if the body should be represented by an array
if "array" in body_repr:
body_name = self._params["body"]["array_name"]
available_indices = deque()
max_index = -1
body_ind = 0
index_available_at = defaultdict(list)
if body_name == result_name:
error("body and result cannot have the same name.")
# Initiate shapes with inf
self.shapes[body_name] = (float("inf"),)
# Iterate over body expressions and recreate the different expressions
# according to state, parameters, body and result expressions
replaced_expr_map = OrderedDict()
new_body_expressions = []
present_ode_objects = dict(
(state.name, state) for state in self.root.full_states
)
present_ode_objects.update(
(param.name, param) for param in self.root.parameters
)
old_present_ode_objects = present_ode_objects.copy()
def store_expressions(expr, new_expr):
"Help function to store new expressions"
timer = Timer( # noqa: F841
f"Store expression while recreating body of {self.name}",
) # noqa: F841
# Update sym replace dict
if isinstance(expr, Derivatives):
der_replace_dict[expr.sym] = new_expr.sym
else:
replace_dict[expr.sym] = new_expr.sym
# Store the new expression for later references
present_ode_objects[expr.name] = new_expr
replaced_expr_map[expr] = new_expr
# Append the new expression
new_body_expressions.append(new_expr)
# Update dependency information
if expr in object_used_in:
for dep in object_used_in[expr]:
if dep in expression_dependencies:
expression_dependencies[dep].remove(expr)
expression_dependencies[dep].add(new_expr)
object_used_in[new_expr] = object_used_in.pop(expr)
if expr in expression_dependencies:
expression_dependencies[new_expr] = expression_dependencies.pop(expr)
self.add_comment("Recreated body expressions")
# The main iteration over all body_expressions
for expr in body_expressions:
# 1) Comments
if isinstance(expr, Comment):
new_body_expressions.append(expr)
continue
assert isinstance(expr, Expression)
# 2) Check for expression optimizations
if not (optimize_exprs == "none" or expr in result_expressions):
timer_opt = Timer( # noqa: F841
f"Handle expression optimization for {self.name}",
) # noqa: F841
# If expr is just a number we exchange the expression with the
# number
if "numerals" in optimize_exprs and isinstance(expr.expr, sp.Number):
replace_dict[expr.sym] = expr.expr
# Remove information about this expr beeing used
for dep in object_used_in[expr]:
expression_dependencies[dep].remove(expr)
object_used_in.pop(expr)
continue
# If the expr is just a symbol (symbol multiplied with a scalar)
# we exchange the expression with the sympy expressions
elif "symbols" in optimize_exprs and (
isinstance(expr.expr, (sp.Symbol, AppliedUndef))
or isinstance(expr.expr, sp.Mul)
and len(expr.expr.args) == 2
and isinstance(expr.expr.args[1], (sp.Symbol, AppliedUndef))
and expr.expr.args[0].is_number
):
# Add a replace rule based on the stored sympy expression
sympy_expr = expr.expr.xreplace(der_replace_dict).xreplace(
replace_dict,
)
if isinstance(expr.sym, sp.Derivative):
der_replace_dict[expr.sym] = sympy_expr
else:
replace_dict[expr.sym] = sympy_expr
# Get exchanged repr
if isinstance(expr.expr, (sp.Symbol, AppliedUndef)):
name = sympycode(expr.expr)
else:
name = sympycode(expr.expr.args[1])
dep_expr = present_ode_objects[name]
# If using reused body expressions we need to update the
# index information so that the index previously available
# for this expressions gets available at the last expressions
# the present expression is used in.
if (
isinstance(dep_expr, IndexedExpression)
and dep_expr.basename == body_name
and "reused" in body_repr
):
ind = dep_expr.indices[0]
# Remove available index information
dep_used_in = sorted(object_used_in[dep_expr])
for used_expr in dep_used_in:
if ind in index_available_at[used_expr]:
index_available_at[used_expr].remove(ind)
# Update with new indices
all_used_in = object_used_in[expr].copy()
all_used_in.update(dep_used_in)
for used_expr in sorted(all_used_in, reverse=True):
if used_expr in body_expressions:
index_available_at[used_expr].append(ind)
break
# Update information about this expr beeing used
for dep in object_used_in[expr]:
expression_dependencies[dep].remove(expr)
expression_dependencies[dep].add(dep_expr)
object_used_in.pop(expr)
continue
del timer_opt
# 3) General operations for all Expressions that are kept
# Before we process the expression we check if any indices gets
# available with the expr (Only applies for the "reused" option for
# body_repr.)
if "reused" in body_repr:
# Check if any indices are available at this expression ind
available_indices.extend(index_available_at[expr])
# Store a map of old name this will preserve the ordering of
# expressions with the same name, similar to how this is treated in
# the actual ODE.
present_ode_objects[expr.name] = expr
old_present_ode_objects[expr.name] = expr
# 4) Handle result expression
if expr in result_expressions:
timer_result = Timer( # noqa: F841
f"Handle result expressions for {self.name}",
) # noqa: F841
# Get the result name
result_name = result_names[expr]
# If the expression is an IndexedExpression with the same basename
# as the result name we just recreate it
if (
isinstance(expr, IndexedExpression)
or isinstance(expr, StateIndexedExpression)
or isinstance(expr, ParameterIndexedExpression)
) and result_name == expr.basename:
new_expr = recreate_expression(expr, der_replace_dict, replace_dict)
# Not an indexed expression
else:
# Get index based on the original ordering
index = (results[result_name].index(expr),)
# Create the IndexedExpression
# NOTE: First replace any derivative expression replaces, then state and
# NOTE: params
if isinstance(expr, StateDerivative):
new_expr = StateIndexedExpression(
result_name,
index,
expr.expr.xreplace(der_replace_dict).xreplace(replace_dict),
expr.state,
(len(results[result_name]),),
array_params=self._params.array,
)
else:
new_expr = IndexedExpression(
result_name,
index,
expr.expr.xreplace(der_replace_dict).xreplace(replace_dict),
(len(results[result_name]),),
array_params=self._params.array,
)
if new_expr.basename not in self.indexed_map:
self.indexed_map[new_expr.basename] = OrderedDict()
self.indexed_map[new_expr.basename][expr] = new_expr
# Copy counter from old expression so it sort properly
new_expr._recount(expr._count)
# Store the expressions
store_expressions(expr, new_expr)
del timer_result
# 4) Handle indexed expression
# All indexed expressions are just kept but recreated with updated
# sympy expressions
elif isinstance(expr, IndexedExpression):
timer_indexed = Timer( # noqa: F841
f"Handle indexed expressions for {self.name}",
) # noqa: F841
new_expr = recreate_expression(expr, der_replace_dict, replace_dict)
# Store the expressions
store_expressions(expr, new_expr)
del timer_indexed
# 5) If replacing all body exressions with an indexed expression
elif "array" in body_repr:
timer_body = Timer( # noqa: F841
f"Handle body expressions for {self.name}",
) # noqa: F841
# 5a) If we reuse array indices
if "reused" in body_repr:
if available_indices:
ind = available_indices.popleft()
else:
max_index += 1
ind = max_index
# Check when present ind gets available again
for used_expr in sorted(object_used_in[expr], reverse=True):
if used_expr in body_expressions:
index_available_at[used_expr].append(ind)
break
else:
warning("SHOULD NOT COME HERE!")
# 5b) No reuse of array indices. Here each index corresponds to
# a distinct body expression
else:
ind = body_ind
# Increase body_ind
body_ind += 1
# Create the IndexedExpression
new_expr = IndexedExpression(
body_name,
ind,
expr.expr.xreplace(der_replace_dict).xreplace(replace_dict),
array_params=self._params.array,
enum=expr.name,
)
if body_name not in self.indexed_map:
self.indexed_map[body_name] = OrderedDict()
self.indexed_map[body_name][expr] = new_expr
# Copy counter from old expression so they sort properly
new_expr._recount(expr._count)
# Store the expressions
store_expressions(expr, new_expr)
del timer_body
# 6) If the expression is just an ordinary body expression and we
# are using named representation of body
else:
timer_expr = Timer(f"Handle expressions for {self.name}") # noqa: F841
# If the expression is a state derivative we need to add a
# replacement for the Derivative symbol
if isinstance(expr, StateDerivative):
new_expr = Intermediate(
sympycode(expr.sym),
expr.expr.xreplace(der_replace_dict).xreplace(replace_dict),
)
new_expr._recount(expr._count)
else:
new_expr = recreate_expression(expr, der_replace_dict, replace_dict)
del timer_expr
# Store the expressions
store_expressions(expr, new_expr)
# Store indices for any added arrays
if "reused_array" == body_repr:
if max_index > -1:
self.shapes[body_name] = (max_index + 1,)
else:
self.shapes.pop(body_name)
elif "array" == body_repr:
if body_ind > 0:
self.shapes[body_name] = (body_ind,)
else:
self.shapes.pop(body_name)
# Store the shape of the added result expressions
for result_name, result_expressions in list(results.items()):
if result_name not in self.shapes:
self.shapes[result_name] = (len(result_expressions),)
return new_body_expressions
def _load(filename, name, **arguments):
"""
Load an ODE or Cell from file and return the instance
The method looks for a file with .ode extension.
Arguments
---------
filename : str
Name of the ode file to load
name : str (optional)
Set the name of ODE (defaults to filename)
arguments : dict (optional)
Optional arguments which can control loading of model
"""
timer = Timer("Load ODE") # noqa: F841
filename = Path(filename).with_suffix(".ode").absolute()
# Extract name from filename
name = filename.stem
# Execute the file
if not filename.is_file():
error(f"Could not find '{filename}'")
# Copy file temporary to current directory
basename = Path(filename.name).absolute()
copyfile = False
if basename != filename:
shutil.copy(filename, basename)
filename = basename
name = filename.stem
copyfile = True
# If a Param is provided turn it into its value
for key, value in list(arguments.items()):
if isinstance(value, Param):
arguments[key] = value.getvalue()
# Dict to collect namespace
intermediate_dispatcher = IntermediateDispatcher()
# Create an ODE which will be populated with data when ode file is loaded
ode = ODE(name, intermediate_dispatcher)
intermediate_dispatcher.ode = ode
debug(f"Loading {ode.name}")
# Create namespace which the ode file will be executed in
_init_namespace(ode, arguments, intermediate_dispatcher)
# Execute file and collect
with open(filename, "r") as f:
exec(compile(f.read(), filename, "exec"), intermediate_dispatcher)
# Finalize ODE
ode.finalize()
# Check for completeness
if not ode.is_complete:
warning(f"ODE model '{ode.name}' is not complete.")
info(f"Loaded ODE model '{ode.name}' with:")
for what in ["full_states", "parameters"]:
num = getattr(ode, f"num_{what}")
info(f"{('Num ' + what.replace('_', ' ')).rjust(20)}: {num}")
if copyfile:
os.unlink(filename)
return ode
def __setitem__(self, name, value):
"""
This is only for expressions
"""
from modelparameters.sympytools import symbols_from_expr
timer = Timer("Namespace dispatcher") # noqa: F841
# Set the attr of the ODE
# If a scalar or a sympy number or it is a sympy.Basic consisting of
# sp.Symbols
if (
isinstance(value, scalars)
or isinstance(value, sp.Number)
or (isinstance(value, sp.Basic) and symbols_from_expr(value))
):
# Get source which triggers the insertion to the global namespace
frame = inspect.currentframe().f_back
lines, lnum = inspect.findsource(frame)
# Try getting the code
try:
code = lines[frame.f_lineno - 1].strip()
# Concatenate lines with line continuation symbols
prev = 2
while (
frame.f_lineno - prev >= 0
and len(lines[frame.f_lineno - prev]) >= 2
and lines[frame.f_lineno - prev][-2:] == "\\\n"
):
code = lines[frame.f_lineno - prev][:-2].strip() + code
prev += 1
except Exception:
code = ""
# Check if the line includes a for statement
# Here we strip op potiential code comments after the main for
# statement.
if re.search(_for_template, code.split("#")[0].strip()) or re.search(
_no_intermediate_template,
code,
):
debug(f"Not registering '{name}' as an intermediate.")
# If so just add the value to the namespace without
# registering the intermediate
dict.__setitem__(self, name, value)
else:
del timer
# Add obj to the present component
setattr(self.ode.present_component, name, value)
else:
debug(f"Not registering '{name}' as an intermediate.")
# If no ode attr was generated we just add the value to the
# namespace
dict.__setitem__(self, name, value)
def __init__(
self,
factorized,
function_name="forward_backward_subst",
result_name="dx",
residual_name="F",
params=None,
):
"""
Create a JacobianForwardBackwardSubstComponent
Arguments
---------
factorized : gotran.FactorizedJacobianComponent
The factorized jacobian of the ODE
function_name : str
The name of the function which should be generated
result_name : str
The name of the result (increment)
residual_name : str
The name of the residual
params : dict
Parameters determining how the code should be generated
"""
timer = Timer(
"Computing forward backward substituion component") # noqa: F841
check_arg(factorized, FactorizedJacobianComponent)
jacobian_name = list(factorized.shapes.keys())[0]
descr = ("Symbolically forward backward substitute linear system "
"of {0} ODE".format(factorized.root))
super(ForwardBackwardSubstitutionComponent, self).__init__(
"ForwardBackwardSubst",
factorized.root,
function_name,
descr,
params=params,
use_default_arguments=False,
additional_arguments=[residual_name],
)
self.add_comment(
f"Forward backward substituting factorized linear system {self.root.name}",
)
# Recreate jacobian using only sympy Symbols
jac_orig = factorized.factorized_jacobian
# Size of system
n = jac_orig.rows
jac = sp.Matrix(n, n, lambda i, j: sp.S.Zero)
for i in range(n):
for j in range(n):
# print jac_orig[i,j]
if not jac_orig[i, j].is_zero:
name = sympycode(jac_orig[i, j])
jac[i, j] = sp.Symbol(
name,
real=True,
imaginary=False,
commutative=True,
hermitian=True,
complex=True,
)
print(jac[i, j])
self.shapes[jacobian_name] = (n, n)
self.shapes[residual_name] = (n, )
self.shapes[result_name] = (n, )
F = []
dx = []
# forward substitution, all diag entries are scaled to 1
for i in range(n):
F.append(self.add_indexed_object(residual_name, i))
dx.append(self.add_indexed_expression(result_name, i, F[i]))
for j in range(i):
if jac[i, j].is_zero:
continue
dx[i] = self.add_indexed_expression(
result_name,
i,
dx[i] - dx[j] * jac[i, j],
)
# backward substitution
for i in range(n - 1, -1, -1):
for j in range(i + 1, n):
if jac[i, j].is_zero:
continue
dx[i] = self.add_indexed_expression(
result_name,
i,
dx[i] - dx[j] * jac[i, j],
)
dx[i] = self.add_indexed_expression(result_name, i,
dx[i] / jac[i, i])
# No need to call recreate body expressions
self.body_expressions = [
obj for obj in self.ode_objects
if isinstance(obj, (IndexedExpression, Comment))
]
self.results = [result_name]
self.used_states = set()
self.used_parameters = set()
def __init__(self, ode):
check_arg(ode, ODE)
assert ode.is_finalized
timer = Timer("Extract common sub expressions") # noqa: F841
newname = ode.name + "_CSE"
# Call super class
super(CommonSubExpressionODE, self).__init__(newname, ode.ns)
# Add states and parameters
atoms = []
for state in ode.full_states:
atoms.append(self.add_state(state.name, state.param))
for param in ode.parameters:
atoms.append(self.add_parameter(param.name, param.param))
# Collect all expanded state expressions
org_state_expressions = ode.state_expressions
expanded_state_exprs = [
ode.expanded_expressions[obj.name] for obj in org_state_expressions
]
# Call sympy common sub expression reduction
cse_exprs, cse_state_exprs = cse(
expanded_state_exprs,
symbols=sp.numbered_symbols("cse_"),
optimizations=[],
)
cse_cnt = 0
cse_subs = {}
# Register the common sub expressions as Intermediates
for sub, expr in cse_exprs:
# If the expression is just one of the atoms of the ODE we skip
# the cse expressions but add a subs for the atom
if expr in atoms:
cse_subs[sub] = expr
else:
cse_subs[sub] = self.add_intermediate(
f"cse_{cse_cnt}",
expr.xreplace(cse_subs),
)
cse_cnt += 1
# Register the state expressions
for org_state_expr, state_expr in zip(org_state_expressions,
cse_state_exprs):
exp_expr = state_expr.xreplace(cse_subs)
state = self.get_object(org_state_expr.state.name)[1]
# If state derivative
if isinstance(org_state_expr, StateDerivative):
self.add_derivative(state, state.time.sym, exp_expr)
# If algebraic
elif isinstance(org_state_expr, AlgebraicExpression):
self.add_algebraic(state, exp_expr)
else:
error("Should not come here!")
self.finalize()
