def rhs_expressions(ode, function_name="rhs", result_name="dy", params=None):
"""
Return a code component with body expressions for the right hand side
Arguments
---------
ode : gotran.ODE
The finalized ODE
function_name : str
The name of the function which should be generated
result_name : str
The name of the variable storing the rhs result
params : dict
Parameters determining how the code should be generated
"""
check_arg(ode, ODE)
if not ode.is_finalized:
error(
"Cannot compute right hand side expressions if the ODE is "
"not finalized", )
descr = f"Compute the right hand side of the {ode} ODE"
return CodeComponent(
"RHSComponent",
ode,
function_name,
descr,
params=params,
**{result_name: ode.state_expressions},
)
def pop(self, index):
check_arg(index, int)
if index >= len(self):
raise IndexError("pop index out of range")
obj = super(ODEObjectList, self).pop(index)
self._objects.pop(obj.name)
def forward_backward_subst_expressions(
factorized,
function_name="forward_backward_subst",
result_name="dx",
residual_name="F",
params=None,
):
"""
Return an ODEComponent holding expressions for the forward backward
substitions for a factorized jacobian
Arguments
---------
factorized : gotran.FactorizedJacobianComponent
The ODEComponent holding expressions for the factorized jacobian
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
"""
check_arg(factorized, FactorizedJacobianComponent)
return ForwardBackwardSubstitutionComponent(
factorized,
function_name=function_name,
result_name=result_name,
residual_name=residual_name,
params=params,
)
def __init__(self, state, expr, dependent=None):
"""
Create a StateDerivative
Arguments
---------
state : State
The state for which the StateDerivative should apply
expr : sympy.Basic
The expression which the differetiation should be equal
dependent : ODEObject
If given the count of this StateDerivative will follow as a
fractional count based on the count of the dependent object
"""
check_arg(state, State, 0, StateDerivative)
sym = sp.Derivative(state.sym, state.time.sym)
sym._assumptions["real"] = True
sym._assumptions["imaginary"] = False
sym._assumptions["commutative"] = True
sym._assumptions["hermitian"] = True
sym._assumptions["complex"] = True
# Call base class constructor
super(StateDerivative, self).__init__(sympycode(sym), state, expr, dependent)
self._sym = sym
def save(basename, **data):
"""
Save data using cPickle
@type basename : str
@param basename : The name of the file to save the data in, .cpickle
will be appended to the file name if not provided
@param data : The actuall data to be saved.
"""
check_arg(basename, str, 0)
# If zero data size just return
if len(data) == 0:
return
filename = basename if ".cpickle" in basename else basename + ".cpickle"
f = open(filename, "w")
p = Pickler(f)
# Dump the dictionary kwarg
p.dump(data)
f.close()
def __init__(self, ode, theta=1):
"""
Create a SymbolicNewtonSolution
Arguments
---------
ode : ODE
The ODE which the symbolc solution is created for
theta : scalar \xe2\x88\x88 [0,1]
Theta for creating the numerical integration rule of the ODE
"""
check_arg(ode, ODE, 0, SymbolicNewtonSolution)
check_arg(theta, scalars, 1, SymbolicNewtonSolution, ge=0, le=1)
# Store attributes
self.ode = ode
self.theta = theta
# Create symbolic linear system
(
self.F,
self.F_expr,
self.jacobi,
self.states,
self.jac_subs,
) = _create_newton_system(ode, theta)
# Create a simplified LU decomposition of the jacobi matrix
# FIXME: Better names!
self.x, self.new_old, self.old_new = _LU_solve(self.jacobi, self.F)
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 _expect_state(self,
state,
allow_state_solution=False,
only_local_states=False):
"""
Help function to check an argument which should be expected
to be a state
"""
if allow_state_solution:
allowed = (State, StateSolution)
else:
allowed = (State, )
if isinstance(state, AppliedUndef):
name = sympycode(state)
state = self.root.present_ode_objects.get(name)
if state is None:
error(f"{name} is not registered in this ODE")
if only_local_states and not (state in self.states or
(state in self.intermediates
and allow_state_solution)):
error(f"{name} is not registered in component {self.name}")
check_arg(state, allowed, 0)
if isinstance(state, State) and state.is_solved:
error(
"Cannot registered a state expression for a state "
"which is registered solved.", )
return state
def __init__(self, name, dependent=None):
"""
Create ODEObject instance
Arguments
---------
name : str
The name of the ODEObject
dependent : ODEObject
If given the count of this ODEObject will follow as a
fractional count based on the count of the dependent object
"""
check_arg(name, str, 0, ODEObject)
check_arg(dependent, (type(None), ODEObject))
self._name = self._check_name(name)
# Unique identifyer
if dependent is None:
self._count = ODEObject.__count
ODEObject.__count += 1
else:
ODEObject.__dependent_counts[dependent._count] += 1
# FIXME: Do not hardcode the fractional increase
self._count = (
dependent._count +
ODEObject.__dependent_counts[dependent._count] * 0.00001)
def diagonal_jacobian_action_expressions(
diagonal_jacobian,
with_body=True,
function_name="compute_diagonal_jacobian_action",
result_name="diag_jac_action",
params=None,
):
"""
Return an ODEComponent holding expressions for the diagonal jacobian action
Arguments
---------
diagonal_jacobian : gotran.DiagonalJacobianComponent
The ODEComponent holding expressions for the diagonal jacobian
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 name of the variable storing the jacobian diagonal result
params : dict
Parameters determining how the code should be generated
"""
check_arg(diagonal_jacobian, DiagonalJacobianComponent)
return DiagonalJacobianActionComponent(
diagonal_jacobian,
with_body,
function_name,
result_name,
params=params,
)
def rename(self, name):
"""
Rename the ODEObject
"""
check_arg(name, str)
self._name = self._check_name(name)
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 __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 _add_rates(self, states, rate_matrix):
"""
Use a rate matrix to set rates between states
Arguments
---------
states : list of States, tuple of two lists of States
If one list is passed the rates should be a square matrix
and the states list determines the order of the row and column of
the matrix. If two lists are passed the first determines the states
in the row and the second the states in the column of the Matrix
rates_matrix : sympy.MatrixBase
A sympy.Matrix of the rate expressions between the states given in
the states argument
"""
check_arg(states, (tuple, list), 0, ODEComponent._add_rates)
check_arg(rate_matrix, sp.MatrixBase, 1, ODEComponent._add_rates)
# If list
if isinstance(states, list):
states = (states, states)
# else tuple
elif len(states) != 2 and not all(
isinstance(list_of_states, list) for list_of_states in states):
error("expected a tuple of 2 lists with states as the "
"states argument")
# Check index arguments
# for list_of_states in states:
# print list_of_states, local_states
# if not all(state in local_states for state in list_of_states):
# error("Expected the states arguments to be States in "\
# "the Markov model")
# Check that the length of the state lists corresponds with the shape of
# the rate matrix
if rate_matrix.shape[0] != len(
states[0]) or rate_matrix.shape[1] != len(states[1], ):
error("Shape of rates does not match given states")
for i, state_i in enumerate(states[0]):
for j, state_j in enumerate(states[1]):
value = rate_matrix[i, j]
# If 0 as rate
if (isinstance(value, scalars)
and value == 0) or (isinstance(value, sp.Basic)
and value.is_zero):
continue
if state_i == state_j:
error("Cannot have a nonzero rate value between the "
"same states")
# Assign the rate
self._add_single_rate(state_i, state_j, value)
def set_state_prefix(self, prefix):
"""
Register a prefix to a state name. Used if
"""
check_arg(prefix, str)
self._state_prefix = prefix
# Reset symbol subs
self._symbol_subs = None
def set_parameter_prefix(self, prefix):
"""
Register a prefix to a parameter name. Used if
"""
check_arg(prefix, str)
self._parameter_prefix = prefix
# Reset symbol subs
self._symbol_subs = None
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 __init__(self, to_state, from_state, expr, dependent=None):
check_arg(to_state, (State, StateSolution), 0, RateExpression)
check_arg(from_state, (State, StateSolution), 1, RateExpression)
super(RateExpression, self).__init__(
f"rates_{to_state}_{from_state}",
expr,
dependent,
)
self._to_state = to_state
self._from_state = from_state
def _add_single_rate(self, to_state, from_state, expr):
"""
Add a single rate expression
"""
if self.state_expressions:
error(
"A component cannot have both state expressions (derivative "
"and algebraic expressions) and rate expressions.", )
check_arg(expr, scalars + (sp.Basic, ), 2,
ODEComponent._add_single_rate)
expr = sp.sympify(expr)
to_state = self._expect_state(
to_state,
allow_state_solution=True,
only_local_states=True,
)
from_state = self._expect_state(
from_state,
allow_state_solution=True,
only_local_states=True,
)
if to_state == from_state:
error("The two states cannot be the same.")
if (to_state.sym, from_state.sym) in self.rates:
error(
f"Rate between state {from_state} and {to_state} is already registered.",
)
# FIXME: It should also not be possible to include other
# states in the markov model, right?
syms_expr = symbols_from_expr(expr)
if to_state.sym in syms_expr or from_state.sym in syms_expr:
error(
"The rate expression cannot be dependent on the "
"states it connects.", )
# Create a RateExpression
obj = RateExpression(to_state, from_state, expr)
self._register_component_object(obj)
# Store the rate sym in the rate dict
self.rates._register_single_rate(to_state.sym, from_state.sym, obj.sym)
def monitored_expressions(
ode,
monitored,
function_name="monitored_expressions",
result_name="monitored",
params=None,
):
"""
Return a code component with body expressions to calculate monitored expressions
Arguments
---------
ode : gotran.ODE
The finalized ODE for which the monitored expression should be computed
monitored : tuple, list
A tuple/list of strings containing the name of the monitored expressions
function_name : str
The name of the function which should be generated
result_name : str
The name of the variable storing the rhs result
params : dict
Parameters determining how the code should be generated
"""
check_arg(ode, ODE)
if not ode.is_finalized:
error(
"Cannot compute right hand side expressions if the ODE is "
"not finalized", )
check_arg(monitored, (tuple, list), itemtypes=str)
monitored_exprs = []
for expr_str in monitored:
obj = ode.present_ode_objects.get(expr_str)
if not isinstance(obj, Expression):
error(f"{expr_str} is not an expression in the {ode} ODE")
monitored_exprs.append(obj)
descr = f"Computes monitored expressions of the {ode} ODE"
return CodeComponent(
"MonitoredExpressions",
ode,
function_name,
descr,
params=params,
**{result_name: monitored_exprs},
)
def __init__(self, name, parent):
"""
Create an ODEComponent
Arguments
---------
name : str
The name of the component. This str serves as the unique
identifier of the Component.
parent : gotran.ODEComponent
The parent component of this ODEComponent
"""
# Turn off magic attributes (see __setattr__ method) during
# construction
self._allow_magic_attributes = False
check_arg(name, str, 0, ODEComponent)
check_arg(parent, ODEComponent, 1, ODEComponent)
# Call super class
super(ODEComponent, self).__init__(name)
# Store parent component
self._parent = weakref.ref(parent)
# Store ODEComponent children
self.children = OrderedDict()
# Store ODEObjects of this component
self.ode_objects = ODEObjectList()
# Storage of rates
self.rates = RateDict(self)
# Store all state expressions
self._local_state_expressions = dict()
# Collect all comments
self._local_comments = []
# Flag to check if component is finalized
self._is_finalized = False
self._is_finalizing = False
# Turn on magic attributes (see __setattr__ method)
self._allow_magic_attributes = True
def __init__(self, name, init):
"""
Create a Parameter with an associated initial value
Arguments
---------
name : str
The name of the State
init : scalar, ScalarParam
The initial value of this parameter
"""
# Call super class
check_arg(init, scalars + (ScalarParam, ), 1, Parameter)
# Call super class
super(Parameter, self).__init__(name, init)
def factorized_jacobian_expressions(
jacobian,
function_name="lu_factorize",
params=None,
):
"""
Return an ODEComponent holding expressions for the factorized jacobian
Arguments
---------
jacobian : gotran.JacobianComponent
The ODEComponent holding expressions for the jacobian
params : dict
Parameters determining how the code should be generated
"""
check_arg(jacobian, JacobianComponent)
return FactorizedJacobianComponent(jacobian, function_name, params=params)
def __init__(self, name, value, dependent=None):
"""
Create ODEObject instance
Arguments
---------
name : str
The name of the ODEObject
value : scalar, ScalarParam, np.ndarray, sp. Basic
The value of this ODEObject
dependent : ODEObject
If given the count of this ODEValueObject will follow as a
fractional count based on the count of the dependent object
"""
check_arg(name, str, 0, ODEValueObject)
check_arg(value, scalars + (ScalarParam, sp.Basic), 1, ODEValueObject)
# Init super class
super(ODEValueObject, self).__init__(name, dependent)
if isinstance(value, ScalarParam):
# Re-create one with correct name
value = value.copy(include_name=False)
value.name = name
elif isinstance(value, scalars):
value = ScalarParam(value, name=name)
elif isinstance(value, str):
value = ConstParam(value, name=name)
else:
value = SlaveParam(value, name=name)
# Debug
# if get_log_level() <= DEBUG:
# if isinstance(value, SlaveParam):
# debug("{0}: {1} {2:.3f}".format(self.name, value.expr, value.value))
# else:
# debug("{0}: {1}".format(name, value.value))
# Store the Param
self._param = value
def expressions(*args):
check_arg(args, tuple, 0, itemtypes=str)
comp = ode()
args = deque(args)
while args:
comp = comp(args.popleft())
assert ode().present_component == comp
# Update the rates name so it points to the present components
# rates dictionary
namespace["rates"] = comp.rates
return comp
def load(basename, latest_timestamp=False, collect=False):
"""
Load data using cPickle
@type basename : str
@param basename : The name of the file where the data will be loaded from,
'.cpickle' will be appended to the file name if not provided
@type latest_timestamp : bool
@param latest_timestamp : If true return the data from latest version of
saved data with the same basename
@type collect : bool
@param collect : If True collect all data with the same basename and
the same parameters
"""
check_arg(basename, str, 0)
if latest_timestamp and collect:
raise TypeError("'collect' and 'latest_timestamp' cannot both be True")
# If not collect just return the data froma single data file
if not collect:
return load_single_data(basename, latest_timestamp)
filenames = get_data_filenames(basename)
# No filenames with timestamp. Try to return data file without timestamp
if not filenames:
return load_single_data(basename, False)
# Start with the latest filename and load the data and collect them if
# the data have the same parameter
data = load_single_data(filenames.pop(-1), False)
params = data["params"]
for filename in reversed(filenames):
local_data = load_single_data(filename, False)
if not compare_dicts(params, local_data["params"]):
info("Not the same parameters, skipping data from '%s'", filename)
continue
merge_data_dicts(data, local_data)
return data
def __call__(self, name):
"""
Return a child component, if the component does not excist, create
and add one
"""
check_arg(name, str)
# If the requested component is the root component
if self == self.root and name == self.root.name:
comp = self.root
else:
comp = self.children.get(name)
if comp is None:
comp = self.add_component(name)
debug(f"Adding '{name}' component to {self}")
else:
self.root._present_component = comp.name
return comp
def _recount(self, new_count=None, dependent=None):
"""
Method called when an object need to get a new count
"""
old_count = self._count
if isinstance(new_count, (int, float)):
check_arg(new_count, (int, float), ge=0, le=ODEObject.__count)
self._count = new_count
elif isinstance(dependent, ODEObject):
ODEObject.__dependent_counts[dependent._count] += 1
# FIXME: Do not hardcode the fractional increase
self._count = (
dependent._count +
ODEObject.__dependent_counts[dependent._count] * 0.00001)
else:
self._count = ODEObject.__count
ODEObject.__count += 1
debug(f"Change count of {self.name} from {old_count} to {self._count}")
def __init__(self, state, expr, dependent=None):
"""
Create a StateSolution
Arguments
---------
state : State
The state that is being solved for
expr : sympy.Basic
The expression that should equal 0 and which solves the state
dependent : ODEObject
If given the count of this StateSolution will follow as a
fractional count based on the count of the dependent object
"""
check_arg(state, State, 0, StateSolution)
super(StateSolution, self).__init__(sympycode(state.sym), expr)
# Flag solved state
state._is_solved = True
self._state = state
def import_ode(subode, prefix="", components=None, **arguments):
"""
Import an ODE into the present ode
Argument
--------
subode : str
The ode which should be imported
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 ode. If not given the whole ODE will be imported.
arguments : dict (optional)
Optional arguments which can control loading of model
"""
check_arg(subode, (str, Path), 0)
# Add the 'subode and update namespace
ode().import_ode(subode, prefix=prefix, components=components, **arguments)