#!/usr/bin/env python
# coding=utf8
import logging
import sys
from functools import partial
from itertools import product
from pickle import dump, load
import numpy as np
import theano as th
import theano.sparse as ths
from sympy import Derivative, Function, Symbol, SympifyError, symbols, sympify
from sympy.printing.theanocode import theano_code
from theano import tensor as T
from theano.ifelse import ifelse
from triflow.core.fields import BaseFields
from triflow.core.routines import F_Routine, J_Routine
logging.getLogger(__name__).addHandler(logging.NullHandler())
logging = logging.getLogger(__name__)
sys.setrecursionlimit(40000)
def _generate_sympify_namespace(independent_variable,
dependent_variables,
helper_functions):
"""Generate the link between the symbols of the derivatives and the
sympy Derivative operation.
Parameters
----------
independent_variable : str
name of the independant variable ("x")
dependent_variables : iterable of str
names of the dependent variables
helper_functions : iterable of str
names of the helper functions
Returns
-------
dict
dictionnary containing the symbol to parse as keys and the sympy expression to evaluate instead as values.
""" # noqa
symbolic_independent_variable = Symbol(independent_variable)
def partial_derivative(symbolic_independent_variable,
i, expr):
return Derivative(expr, symbolic_independent_variable,
i)
namespace = {independent_variable: symbolic_independent_variable}
namespace.update({'d%s' % (independent_variable * i):
partial(partial_derivative,
symbolic_independent_variable, i)
for i in range(1, 10)})
namespace.update({'d%s%s' % (independent_variable * order, var):
Derivative(Function(var)(independent_variable),
independent_variable, order)
for order, var in product(range(1, 10),
dependent_variables +
helper_functions)})
logging.debug(f"sympy namespace: {namespace}")
return namespace
def _reduce_model(eq_diffs, dep_vars, pars,
help_functions):
model = Model(eq_diffs, dep_vars, pars,
help_functions)
return model
[docs]class Model:
"""Contain finite difference approximation and routine of the dynamical system
Take a mathematical form as input, use Sympy to transform it as a symbolic
expression, perform the finite difference approximation and expose theano
optimized routine for both the right hand side of the dynamical system and
Jacobian matrix approximation.
Parameters
----------
differential_equations : iterable of str or str
the right hand sides of the partial differential equations written as :math:`\\frac{\partial U}{\partial t} = F(U)`, where the spatial derivative can be written as `dxxU` or `dx(U, 2)` or with the sympy notation `Derivative(U, x, x)`
dependent_variables : iterable of str or str
the dependent variables with the same order as the temporal derivative of the previous arg.
parameters : iterable of str or str, optional, default None
list of the parameters. Can be feed with a scalar of an array with the same size
help_functions : None, optional
All fields which have not to be solved with the time derivative but will be derived in space.
Attributes
----------
F : triflow.F_Routine
Callable used to compute the right hand side of the dynamical system
F_array : numpy.ndarray of sympy.Expr
Symbolic expressions of the right hand side of the dynamical system
J : triflow.J_Routine
Callable used to compute the Jacobian of the dynamical system
J_array : numpy.ndarray of sympy.Expr
Symbolic expressions of the Jacobian side of the dynamical system
Properties
----------
fields_template: Model specific Fields container used to store and access to the model variables in an efficient way.
Methods
-------
save: Save a binary of the Model with pre-optimized F and J routines
Examples
--------
A simple diffusion equation:
>>> from triflow import Model
>>> model = Model("k * dxxU", "U", "k")
A coupled system of convection-diffusion equation:
>>> from triflow import Model
>>> model = Model(["k1 * dxxU - c1 * dxV",
... "k2 * dxxV - c2 * dxU",],
... ["U", "V"], ["k1", "k2", "c1", "c2"])
""" # noqa
def __init__(self,
differential_equations,
dependent_variables,
parameters=None,
help_functions=None):
logging.debug('enter __init__ Model')
# coerce the inputs the way to have coherent types
# TODO: be more pythonic, it should not be necessary with coherent
# duck typing..
(self._diff_eqs,
self._dep_vars,
self._pars,
self._help_funcs) = self._coerce_inputs(differential_equations,
dependent_variables,
parameters,
help_functions)
self._nvar = len(self._dep_vars)
# generate the sympy namespace which will connect the math input into
# sympy operation.
sympify_namespace = {}
sympify_namespace.update(_generate_sympify_namespace(
'x',
self._dep_vars,
self._help_funcs))
# parse the inputs in order to have Sympy symbols and expressions
(self._symb_diff_eqs,
self._symb_dep_vars,
self._symb_pars,
self._symb_help_funcs) = self._sympify_model(self._diff_eqs,
self._dep_vars,
self._pars,
self._help_funcs,
sympify_namespace)
# we will need to extract the order of the != spatial derivative
# in order to know the size of the ghost area at the limit of
self._symb_vars_with_spatial_diff_order = {str(svar.func):
{(svar.func, 0)}
for svar
in (self._symb_dep_vars +
self._symb_help_funcs)}
# Use finite difference scheme to generate a spatial approximation and
# have a rhs which will only be a function of the discretized dependent
# variables and help functions.
approximated_diff_eqs = self._approximate_derivative(
self._symb_diff_eqs,
self._symb_dep_vars,
self._symb_help_funcs)
logging.debug(f"approximated equations: {approximated_diff_eqs}")
# We compute the size of the the ghost area at the limit of
# the spatial domain
self._bounds = self._extract_bounds(
self._dep_vars,
self._symb_vars_with_spatial_diff_order)
self._window_range = self._bounds[-1] - self._bounds[0] + 1
# We obtain a Fortran like flattened vector containing all the discrete
# dependent variable needed for the spatial approximation around the
# local node i
U = self._extract_unknowns(
self._dep_vars,
self._bounds,
self._symb_vars_with_spatial_diff_order).flatten('F')
# We do the same as for the U variable but with all the discrete
# variables, dependent variables and help functions.
self._discrete_variables = self._extract_unknowns(
self._dep_vars + self._help_funcs,
self._bounds,
self._symb_vars_with_spatial_diff_order).flatten('F')
# We expose a numpy.ndarray filled with the rhs of our approximated
# dynamical system
self.F_array = np.array(approximated_diff_eqs)
# We compute the jacobian as the partial derivative of all equation of
# our system according to all the discrete variable in U.
self.J_array = np.array([
[diff_eq.diff(u).expand()
for u in U]
for diff_eq in approximated_diff_eqs]).flatten('F')
# We flag and store the null entry of the Jacobian matrix
self._sparse_indices = np.where(self.J_array != 0)
# We drop all the null term of the Jacobian matrix, because we target
# a sparse matrix storage for memory saving and efficient linalg ops.
self._J_sparse_array = self.J_array[self._sparse_indices]
# We convert the sympy description of the system into a theano
# graph compilation
F, J, args, self._map_extended = self._theano_convert(
self.F_array,
self._J_sparse_array)
# We compile the previous graphs and obtain optimized C based
# routines
F_theano_function = th.function(inputs=args,
outputs=F,
on_unused_input='ignore')
J_theano_function = th.function(inputs=args,
outputs=J,
on_unused_input='ignore')
self._th_args = args
self._th_vectors = [F, J]
self._theano_routines = [F_theano_function, J_theano_function]
self._compile(self.F_array, self._J_sparse_array,
F_theano_function, J_theano_function)
def _theano_convert(self, F_array, J_array):
th_args = list(
map(
partial(theano_code,
broadcastables={arg: (False,)
for arg
in [Symbol('x'),
*self._discrete_variables,
*self._symb_pars]}),
self._symbolic_args))
ins = th.gof.graph.inputs(th_args)
mapargs = {inp.name: inp
for inp in ins if isinstance(inp, T.TensorVariable)}
x_th = T.dvector('x')
periodic = T.bscalar('periodic')
middle_point = int((self._window_range - 1) / 2)
N = x_th.size
L = x_th[-1]
computed_dx = L / (N - 1)
subs = {mapargs['dx']: computed_dx}
map_extended = {}
for (varname, discretisation_tree) in \
self._symb_vars_with_spatial_diff_order.items():
pad_left, pad_right = self._bounds
per_extended_var = T.concatenate([mapargs[varname][pad_left:],
mapargs[varname],
mapargs[varname][:pad_right]])
edge_extended_var = T.concatenate([T.repeat(mapargs[varname][:1],
middle_point),
mapargs[varname],
T.repeat(mapargs[varname][-1:],
middle_point)])
extended_var = ifelse(periodic,
per_extended_var,
edge_extended_var)
map_extended[varname] = extended_var
for order in range(pad_left, pad_right + 1):
if order != 0:
var = (f"{varname}_{'m' if order < 0 else 'p'}"
f"{np.abs(order)}")
new_var = extended_var[order - pad_left:
extended_var.size +
order - pad_right]
subs.update({mapargs[var]:
new_var})
F = list(map(partial(theano_code,
broadcastables={arg: (False,)
for arg
in [Symbol('x'),
*self._discrete_variables,
*self._symb_pars]}),
F_array.flatten().tolist()))
J = list(map(partial(theano_code,
broadcastables={arg: (False,)
for arg
in [Symbol('x'),
*self._discrete_variables,
*self._symb_pars]}),
J_array.flatten().tolist()))
J = [Ji if not isinstance(Ji, (float, int)) else th.shared(Ji)
for Ji in J]
F = (T.concatenate(th.clone(F, replace=subs))
.reshape((self._nvar, N)).T.flatten())
J = th.clone(J, replace=subs)
J = T.stack([T.repeat(j, N) if j.ndim == 0 else j for j in J]).T
i = T.arange(N).dimshuffle([0, 'x'])
idx = T.arange(N * self._nvar).reshape((N, self._nvar)).T
edge_extended_idx = T.concatenate([T.repeat(idx[:, :1],
middle_point,
axis=1),
idx,
T.repeat(idx[:, -1:],
middle_point,
axis=1)],
axis=1).T.flatten()
per_extended_idx = T.concatenate([idx[:, -middle_point:],
idx,
idx[:, :middle_point]],
axis=1).T.flatten()
extended_idx = ifelse(periodic,
per_extended_idx, edge_extended_idx)
rows = T.tile(T.arange(self._nvar),
self._window_range * self._nvar) + i * self._nvar
cols = T.repeat(T.arange(self._window_range * self._nvar),
self._nvar) + i * self._nvar
rows = rows[:, self._sparse_indices].reshape(J.shape).flatten()
cols = extended_idx[cols][:, self._sparse_indices] \
.reshape(J.shape).flatten()
permutation = T.argsort(cols)
J = J.flatten()[permutation]
rows = rows[permutation]
cols = cols[permutation]
count = T.zeros((N * self._nvar + 1,), dtype=int)
uq, cnt = T.extra_ops.Unique(False, False, True)(cols)
count = T.set_subtensor(count[uq + 1], cnt)
indptr = T.cumsum(count)
shape = T.stack([N * self._nvar, N * self._nvar])
sparse_J = ths.CSC(J, rows, indptr, shape)
th_args = [x_th, *[mapargs[key]
for key
in (self._dep_vars +
self._help_funcs +
self._pars)],
periodic]
return F, sparse_J, th_args, map_extended
def _compile(self, F_array, J_array, F_theano_function, J_theano_function):
logging.debug('compile F')
self.F = F_Routine(F_array,
(self._dep_vars +
self._help_funcs),
self._pars, F_theano_function)
logging.debug('compile J')
self.J = J_Routine(J_array,
(self._dep_vars +
self._help_funcs),
self._pars, J_theano_function)
@property
def fields_template(self):
return BaseFields.factory(self._dep_vars,
self._help_funcs)
@property
def _args(self):
return map(str, self._symbolic_args)
@property
def _symbolic_args(self):
return ([Symbol('x'),
*list(self._discrete_variables),
*list(self._symb_pars),
Symbol('dx')])
[docs] def save(self, filename):
"""Save the model as a binary pickle file.
Parameters
----------
filename : str
name of the file where the model is saved.
Returns
-------
None
"""
with open(filename, 'wb') as f:
dump(self, f)
@staticmethod
[docs] def load(filename):
"""load a pre-compiled triflow model. The internal of theano allow a
caching of the model. Will be slow if it is the first time the model is
loaded on the system.
Parameters
----------
filename : str
path of the pre-compiled model
Returns
-------
triflow.core.Model
triflow pre-compiled model
"""
with open(filename, 'rb') as f:
return load(f)
def _extract_bounds(self, variables, dict_symbol):
bounds = (0, 0)
for var in variables:
dvars, orders = zip(*dict_symbol[var])
bounds = (min(bounds[0],
min(orders)),
max(bounds[1],
max(orders))
)
return bounds
def _extract_unknowns(self, vars, bounds, dict_symbol):
unknowns = np.zeros((len(vars), bounds[-1] - bounds[0] + 1),
dtype=object)
for i, var in enumerate(vars):
dvars, orders = zip(*dict_symbol[var])
for j, order in enumerate(range(bounds[0], bounds[1] + 1)):
if order == 0:
unknowns[i, j] = Symbol(var)
if order < 0:
unknowns[i, j] = Symbol(f'{var}_m{np.abs(order)}')
if order > 0:
unknowns[i, j] = Symbol(f'{var}_p{np.abs(order)}')
return unknowns
def _coerce_inputs(self, diff_eqs, dep_vars, pars, helper_functions):
pars = tuple(pars) if pars is not None else tuple()
helper_functions = (tuple(helper_functions)
if helper_functions is not None else tuple())
if isinstance(diff_eqs, (str, )):
diff_eqs = [diff_eqs]
if isinstance(dep_vars, (str, )):
dep_vars = [dep_vars]
diff_eqs = tuple(diff_eqs)
dep_vars = tuple(dep_vars)
helper_functions = tuple(helper_functions)
return diff_eqs, dep_vars, pars, helper_functions
def _finite_diff_scheme(self, U, order):
logging.debug("finite diff approximation %i, %s" % (order, U))
dx = Symbol('dx')
var_label = str(U)
if order == 1:
Um1 = Symbol('%s_m1' % var_label)
Up1 = Symbol('%s_p1' % var_label)
self._symb_vars_with_spatial_diff_order[var_label].add((Um1, -1))
self._symb_vars_with_spatial_diff_order[var_label].add((Up1, 1))
return (1 / 2 * Up1 - 1 / 2 * Um1) / dx
if order == 2:
Um1 = Symbol('%s_m1' % var_label)
Up1 = Symbol('%s_p1' % var_label)
self._symb_vars_with_spatial_diff_order[var_label].add((Um1, -1))
self._symb_vars_with_spatial_diff_order[var_label].add((Up1, 1))
return (Up1 - 2 * U + Um1) / dx ** 2
if order == 3:
Um1 = Symbol('%s_m1' % var_label)
Um2 = Symbol('%s_m2' % var_label)
Up1 = Symbol('%s_p1' % var_label)
Up2 = Symbol('%s_p2' % var_label)
self._symb_vars_with_spatial_diff_order[var_label].add((Um1, -1))
self._symb_vars_with_spatial_diff_order[var_label].add((Up1, 1))
self._symb_vars_with_spatial_diff_order[var_label].add((Um2, -2))
self._symb_vars_with_spatial_diff_order[var_label].add((Up2, 2))
return (-1 / 2 * Um2 + Um1 - Up1 + 1 / 2 * Up2) / dx ** 3
if order == 4:
Um1 = Symbol('%s_m1' % var_label)
Um2 = Symbol('%s_m2' % var_label)
Up1 = Symbol('%s_p1' % var_label)
Up2 = Symbol('%s_p2' % var_label)
self._symb_vars_with_spatial_diff_order[var_label].add((Um1, -1))
self._symb_vars_with_spatial_diff_order[var_label].add((Up1, 1))
self._symb_vars_with_spatial_diff_order[var_label].add((Um2, -2))
self._symb_vars_with_spatial_diff_order[var_label].add((Up2, 2))
return (Um2 - 4 * Um1 + 6 * U - 4 * Up1 + Up2) / dx ** 4
raise NotImplementedError('Finite difference up '
'to 5th order not implemented yet')
def _sympify_model(self,
diff_eqs: tuple,
dep_vars: tuple,
pars: tuple,
help_functions: tuple,
sympify_namespace: dict) -> tuple:
logging.debug('enter _sympify_model')
symbolic_dep_vars = tuple([Function(dep_var)(Symbol('x'))
for dep_var in dep_vars])
symbolic_help_functions = tuple([Function(help_function)(Symbol('x'))
for help_function in help_functions])
symbolic_pars = symbols(pars)
try:
symbolic_diff_eqs = tuple([sympify(func, locals=sympify_namespace)
.subs(zip(map(Symbol, dep_vars),
(symbolic_dep_vars +
symbolic_help_functions)))
for func
in diff_eqs])
except (TypeError, SympifyError):
raise ValueError("badly formated differential equations")
return (symbolic_diff_eqs, symbolic_dep_vars,
symbolic_pars, symbolic_help_functions)
def _approximate_derivative(self,
symbolic_diff_eqs: tuple,
symbolic_vars: tuple,
symbolic_fields: tuple) -> tuple:
logging.debug('enter _approximate_derivative')
approximated_diff_eqs = []
for func in symbolic_diff_eqs:
afunc = func
for derivative in func.find(Derivative):
var = Symbol(str(derivative.args[0].func))
logging.debug(f"{derivative}, {var}")
order = len(derivative.args) - 1
afunc = afunc.replace(
derivative,
self._finite_diff_scheme(var,
order))
afunc = afunc.subs([(var, Symbol(str(var.func)))
for var in symbolic_vars + symbolic_fields])
approximated_diff_eqs.append(afunc.expand())
return tuple(approximated_diff_eqs)
def __reduce__(self):
return (_reduce_model, (self._diff_eqs, self._dep_vars,
self._pars, self._help_funcs))