diff --git a/examples/atmospheric/baroclinic_one_layer.py b/examples/atmospheric/baroclinic_one_layer.py index 9c327d8..b85ce41 100644 --- a/examples/atmospheric/baroclinic_one_layer.py +++ b/examples/atmospheric/baroclinic_one_layer.py @@ -122,7 +122,7 @@ # defining the LHS as the time derivative of the vorticity vorticity = OperatorTerm(psi, Laplacian, atmospheric_basis.coordinate_system) -barotropic_equation = Equation(psi, lhs_term=vorticity) +barotropic_equation = Equation(psi, lhs_terms=vorticity) # Defining the advection term advection_term1 = vorticity_advection(psi, psi, atmospheric_basis.coordinate_system, sign=-1) @@ -165,7 +165,7 @@ lin_lhs = LinearTerm(theta, prefactor=a, sign=-1) lhs = AdditionOfTerms(lin_lhs, vorticity) -baroclinic_equation = Equation(theta, lhs_term=lhs) +baroclinic_equation = Equation(theta, lhs_terms=lhs) # Defining the advection terms advection_term1 = vorticity_advection(psi, theta, atmospheric_basis.coordinate_system, sign=-1) diff --git a/examples/atmospheric/barotropic_one_layer.py b/examples/atmospheric/barotropic_one_layer.py index 70c8cfc..8d2c52e 100644 --- a/examples/atmospheric/barotropic_one_layer.py +++ b/examples/atmospheric/barotropic_one_layer.py @@ -75,7 +75,7 @@ # defining the LHS as the time derivative of the vorticity vorticity = OperatorTerm(psi, Laplacian, basis.coordinate_system) -barotropic_equation = Equation(psi, lhs_term=vorticity) +barotropic_equation = Equation(psi, lhs_terms=vorticity) # defining the advection term advection_term = vorticity_advection(psi, psi, basis.coordinate_system, sign=-1) diff --git a/examples/atmospheric/stratospheric_planetary_flow.py b/examples/atmospheric/stratospheric_planetary_flow.py index e702401..da7bec6 100644 --- a/examples/atmospheric/stratospheric_planetary_flow.py +++ b/examples/atmospheric/stratospheric_planetary_flow.py @@ -61,7 +61,7 @@ # defining the LHS as the time derivative of the vorticity vorticity = OperatorTerm(psi, Laplacian, basis.coordinate_system) -planetary_equation = Equation(psi, lhs_term=vorticity) +planetary_equation = Equation(psi, lhs_terms=vorticity) # defining the advection term advection_term = vorticity_advection(psi, psi, basis.coordinate_system, sign=-1) diff --git a/examples/ocean-atmosphere/maooam.py b/examples/ocean-atmosphere/maooam.py index 4145691..5352973 100644 --- a/examples/ocean-atmosphere/maooam.py +++ b/examples/ocean-atmosphere/maooam.py @@ -155,7 +155,7 @@ # defining LHS as the time derivative of the vorticity vorticity = OperatorTerm(psi_a, Laplacian, atmospheric_basis.coordinate_system) -barotropic_equation = Equation(psi_a, lhs_term=vorticity) +barotropic_equation = Equation(psi_a, lhs_terms=vorticity) # Defining the advection term advection_term1 = vorticity_advection(psi_a, psi_a, atmospheric_basis.coordinate_system, sign=-1) @@ -190,7 +190,7 @@ lin_lhs = LinearTerm(theta_a, prefactor=a, sign=-1) lhs = AdditionOfTerms(lin_lhs, vorticity) -baroclinic_equation = Equation(theta_a, lhs_term=lhs) +baroclinic_equation = Equation(theta_a, lhs_terms=lhs) # Defining the advection terms advection_term1 = vorticity_advection(psi_a, theta_a, atmospheric_basis.coordinate_system, sign=-1) @@ -250,7 +250,7 @@ vorticity = OperatorTerm(psi_o, Laplacian, oceanic_basis.coordinate_system) lin_lhs = LinearTerm(psi_o, prefactor=G) lhs = AdditionOfTerms(lin_lhs, vorticity) -oceanic_equation = Equation(psi_o, lhs_term=lhs) +oceanic_equation = Equation(psi_o, lhs_terms=lhs) # Defining the advection term advection_term = vorticity_advection(psi_o, psi_o, oceanic_basis.coordinate_system, sign=-1) @@ -283,7 +283,7 @@ # Defining the LHS lhs = LinearTerm(deltaT_o) -ocean_temperature_equation = Equation(deltaT_o, lhs_term=lhs) +ocean_temperature_equation = Equation(deltaT_o, lhs_terms=lhs) # Defining the advection term advection_term = Jacobian(psi_o, deltaT_o, oceanic_basis.coordinate_system, sign=-1) diff --git a/layercake/arithmetic/equation.py b/layercake/arithmetic/equation.py index 1229991..4c17e0e 100644 --- a/layercake/arithmetic/equation.py +++ b/layercake/arithmetic/equation.py @@ -16,6 +16,8 @@ from layercake.utils import isin +# TODO: deal with the cases where lists of terms are empty - maybe nothing is needed but to check + class Equation(object): """Main class to define and specify partial differential equations. @@ -28,20 +30,18 @@ class Equation(object): ---------- field: ~field.Field The spatial field over which the partial differential equation. - lhs_term: ~arithmetic.terms.base.ArithmeticTerms - Term on the left-hand side of the equation. - Must be a single term, possibly a combination - through :class:`~layercake.arithmetic.terms.base.OperationOnTerms` operations. + lhs_terms: ~arithmetic.terms.base.ArithmeticTerms or list(~arithmetic.terms.base.ArithmeticTerms), optional + Terms on the left-hand side of the equation. At least one must involve the field defined above. name: str, optional - Optional name for the equation. + Name for the equation. Attributes ---------- field: ~field.Field The spatial field over which the partial differential equation. - terms: list(~arithmetic.terms.base.ArithmeticTerms) + rhs_terms: list(~arithmetic.terms.base.ArithmeticTerms) List of additive terms in the right-hand side of the equation. - lhs_term: ~arithmetic.terms.base.ArithmeticTerms + lhs_terms: ListOfArithmeticTerms(~arithmetic.terms.base.ArithmeticTerms) Term on the left-hand side of the equation. name: str Optional name for the equation. @@ -49,13 +49,17 @@ class Equation(object): _t = Symbol('t') - def __init__(self, field, lhs_term, name=''): + def __init__(self, field, lhs_terms=None, name=''): self.field = field self.field._equation = self - self.terms = list() - self.lhs_term = lhs_term - self.lhs_term.field = self.field + self.rhs_terms = ListOfAdditiveArithmeticTerms() + if lhs_terms is None: + self.lhs_terms = ListOfAdditiveArithmeticTerms() + elif isinstance(lhs_terms, list): + self.lhs_terms = ListOfAdditiveArithmeticTerms(lhs_terms) + else: + self.lhs_terms = ListOfAdditiveArithmeticTerms([lhs_terms]) self.name = name self._layer = None self._cake = None @@ -83,6 +87,34 @@ def add_rhs_terms(self, terms): for t in terms: self.add_rhs_term(t) + def add_lhs_term(self, term): + """Add a term to the left-hand side of the equation. + + Parameters + ---------- + term: ~arithmetic.terms.base.ArithmeticTerms + Term to be added to the left-hand side of the equation. + """ + if not issubclass(term.__class__, ArithmeticTerms): + raise ValueError('Provided term must be a valid ArithmeticTerm object.') + self.lhs_terms.append(term) + + def add_lhs_terms(self, terms): + """Add multiple terms to the left-hand side of the equation. + + Parameters + ---------- + terms: list(~arithmetic.terms.base.ArithmeticTerms) + Terms to be added to the left-hand side of the equation. + """ + for t in terms: + self.add_lhs_term(t) + + @property + def terms(self): + """Alias for the list of RHS arithmetic terms.""" + return self.rhs_terms + @property def other_fields(self): """list(~field.Field): List of additional fields present in the equation.""" @@ -91,6 +123,17 @@ def other_fields(self): for term in equation_term.terms: if term.field is not self.field and term.field.dynamical and term.field not in other_fields: other_fields.append(term.field) + other_fields = other_fields + self.other_fields_in_lhs + return other_fields + + @property + def other_fields_in_lhs(self): + """list(~field.Field): List of additional fields present in the LHS of the equation.""" + other_fields = list() + for equation_term in self.lhs_terms: + for term in equation_term.terms: + if term.field is not self.field and term.field.dynamical and term.field not in other_fields: + other_fields.append(term.field) return other_fields @property @@ -101,16 +144,17 @@ def parameter_fields(self): for term in equation_term.terms: if term.field is not self.field and not term.field.dynamical and term.field not in parameter_fields: parameter_fields.append(term.field) - for term in self.lhs_term.terms: - if term.field is not self.field and not term.field.dynamical and term.field not in parameter_fields: - parameter_fields.append(term.field) + for equation_term in self.lhs_terms: + for term in equation_term.terms: + if term.field is not self.field and not term.field.dynamical and term.field not in parameter_fields: + parameter_fields.append(term.field) return parameter_fields @property def parameters(self): """list(~parameter.Parameter): List of parameters present in the equation.""" parameters_list = list() - for term in self.terms + [self.lhs_term]: + for term in self.terms + self.lhs_terms: params_list = term.parameters for param in params_list: if not isin(param, parameters_list): @@ -131,64 +175,58 @@ def parameters_symbols(self): @property def symbolic_expression(self): """~sympy.core.expr.Expr: Symbolic expression of the equation.""" - rterm = S.Zero - for term in self.terms: - rterm += term.symbolic_expression - return Eq(self.symbolic_lhs.diff(self._t, evaluate=False), rterm) + return Eq(self.symbolic_lhs.diff(self._t, evaluate=False), self.rhs_terms.symbolic_expression) @property def numerical_expression(self): """~sympy.core.expr.Expr: Expression of the equation with parameters replaced by their configured values.""" - rterm = S.Zero - for term in self.terms: - rterm += term.numerical_expression - return Eq(self.numerical_lhs.diff(self._t, evaluate=False), rterm) + return Eq(self.numerical_lhs.diff(self._t, evaluate=False), self.rhs_terms.symbolic_expression) @property def symbolic_rhs(self): """~sympy.core.expr.Expr: Symbolic expression of the right-hand side of the equation.""" - rterm = S.Zero - for term in self.terms: - rterm += term.symbolic_expression - return rterm + return self.rhs_terms.symbolic_expression @property def numerical_rhs(self): """~sympy.core.expr.Expr: Expression of the right-hand side of the equation with parameters replaced by their configured values.""" - rterm = S.Zero - for term in self.terms: - rterm += term.numerical_expression - return rterm + return self.rhs_terms.numerical_expression @property def symbolic_lhs(self): """~sympy.core.expr.Expr: Symbolic expression of the left-hand side of the equation.""" - return self.lhs_term.symbolic_expression + return self.lhs_terms.symbolic_expression @property def numerical_lhs(self): """~sympy.core.expr.Expr: Expression of the left-hand side of the equation with parameters replaced by their configured values.""" - return self.lhs_term.numerical_expression + return self.lhs_terms.numerical_expression @property def lhs_inner_products(self): - """~sympy.matrices.immutable.ImmutableSparseMatrix or ~sympy.tensor.array.ImmutableSparseNDimArray or sparse.COO(float): Left-hand - side inner products of the equation, if available.""" - return self.lhs_term.inner_products + """list(~sympy.matrices.immutable.ImmutableSparseMatrix or ~sympy.tensor.array.ImmutableSparseNDimArray) or list(sparse.COO(float)): Inner products of each term of + the left-hand side of the equation, if available.""" + return [term.inner_products for term in self.lhs_terms] + + @property + def lhs_inner_products_addition(self): + """~sympy.matrices.immutable.ImmutableSparseMatrix or ~sympy.tensor.array.ImmutableSparseNDimArray or sparse.COO(float): Added left-hand + side inner products of the equation, if available. Might raise an error if not all terms are compatible.""" + result = self.lhs_terms[0].inner_products.copy() + for term in self.lhs_terms[1:]: + result = result + term.inner_products + return result @property def maximum_rank(self): """int: Maximum rank of the right-hand side terms tensors.""" - rhs_max_rank = 0 - for term in self.terms: - rhs_max_rank = max(rhs_max_rank, term.rank) - return max(rhs_max_rank, self.lhs_term.rank) + return max(self.terms.maximum_rank, self.lhs_terms.maximum_rank) - def compute_lhs_inner_products(self, basis, numerical=False, timeout=None, num_threads=None, permute=False): - """Compute the inner products tensor of the left-hand side term. + def compute_inner_products(self, basis, numerical=False, timeout=None, num_threads=None, permute=False): + """Compute the inner products tensor of the left-hand and right-hand side terms. Parameters ---------- @@ -201,17 +239,24 @@ def compute_lhs_inner_products(self, basis, numerical=False, timeout=None, num_t Number of threads to use to compute the inner products. If `None` use all the cpus available. Default to `None`. timeout: int or float or bool or None, optional - The timeout for the numerical computation of each inner product. - After the timeout, compute the inner product with a quadrature instead of symbolic integration. - Does not apply to symbolic output computations. - If `None` or `False`, no timeout occurs. - Default to `None`. + Control the switch from symbolic to numerical integration. By default, `parallel_integration` workers will try to integrate + |Sympy| expressions symbolically, but a fallback to numerical integration can be enforced. + The options are: + + * `None`: This is the "full-symbolic" mode. No timeout will be applied, and the switch to numerical integration will never happen. + Can result in very long and improbable computation time. + * `True`: This is the "full-numerical" mode. Symbolic computations do not occur, and the workers try directly to integrate + numerically. + * `False`: Same as `None`. + * An integer: defines a timeout after which, if a symbolic integration have not completed, the worker switch to the + numerical integration. permute: bool, optional If `True`, applies all the possible permutations to the tensor indices from 1 to the rank of the tensor. Default to `False`, i.e. no permutation is applied. """ - self.lhs_term.compute_inner_products(basis, numerical, timeout, num_threads, permute) + self.lhs_terms.compute_inner_products(basis, numerical, timeout, num_threads, permute) + self.rhs_terms.compute_inner_products(basis, numerical, timeout, num_threads, permute) def to_latex(self, enclose_lhs=True, drop_first_lhs_char=True, drop_first_rhs_char=False): """Generate a LaTeX string representing the equation mathematically. @@ -235,7 +280,7 @@ def to_latex(self, enclose_lhs=True, drop_first_lhs_char=True, drop_first_rhs_ch str The LaTeX string representing the equation. """ - lhs = self.lhs_term.terms[0].latex + lhs = self.lhs_terms[0].latex if drop_first_lhs_char: lhs = lhs[2:] if enclose_lhs: @@ -243,10 +288,10 @@ def to_latex(self, enclose_lhs=True, drop_first_lhs_char=True, drop_first_rhs_ch else: latex_string = r'\frac{\partial}{\partial t} ' + lhs - fterm = self.terms[0].latex + first_term = self.terms[0].latex if drop_first_rhs_char: - fterm = fterm[2:] - latex_string += ' = ' + fterm + first_term = first_term[2:] + latex_string += ' = ' + first_term for term in self.terms[1:]: latex_string += term.latex @@ -287,3 +332,77 @@ def __repr__(self): def __str__(self): return self.__repr__() + + +class ListOfAdditiveArithmeticTerms(list): + """Class holding list of additive arithmetic terms in equations.""" + + def compute_inner_products(self, basis, numerical=False, timeout=None, num_threads=None, permute=False): + """Compute the inner products tensor of the all the terms of the list. + + Parameters + ---------- + basis: SymbolicBasis + Basis with which to compute the inner products. + numerical: bool, optional + Whether the resulting computed inner products must be numerical or symbolic. + Default to `False`, i.e. symbolic output. + num_threads: int or None, optional + Number of threads to use to compute the inner products. If `None` use all the cpus available. + Default to `None`. + timeout: int or float or bool or None, optional + Control the switch from symbolic to numerical integration. By default, `parallel_integration` workers will try to integrate + |Sympy| expressions symbolically, but a fallback to numerical integration can be enforced. + The options are: + + * `None`: This is the "full-symbolic" mode. No timeout will be applied, and the switch to numerical integration will never happen. + Can result in very long and improbable computation time. + * `True`: This is the "full-numerical" mode. Symbolic computations do not occur, and the workers try directly to integrate + numerically. + * `False`: Same as `None`. + * An integer: defines a timeout after which, if a symbolic integration have not completed, the worker switch to the + numerical integration. + permute: bool, optional + If `True`, applies all the possible permutations to the tensor indices + from 1 to the rank of the tensor. + Default to `False`, i.e. no permutation is applied. + """ + + for term in self: + term.compute_inner_products(basis, numerical, timeout, num_threads, permute) + + @property + def maximum_rank(self): + """int: Maximum rank of the right-hand side terms tensors.""" + max_rank = 0 + for term in self: + max_rank = max(max_rank, term.rank) + return max_rank + + @property + def same_rank(self): + """bool: Check if all terms have the same rank.""" + rank = self[0].rank + for term in self[1:]: + if term.rank != rank: + return False + return True + + @property + def symbolic_expression(self): + """~sympy.core.expr.Expr: Symbolic expression of the collection of additive arithmetic terms.""" + + rterm = S.Zero + for term in self: + rterm += term.symbolic_expression + return rterm + + @property + def numerical_expression(self): + """~sympy.core.expr.Expr: Expression of the collection of additive arithmetic terms with + parameters replaced by their configured values.""" + + rterm = S.Zero + for term in self: + rterm += term.numerical_expression + return rterm diff --git a/layercake/bakery/cake.py b/layercake/bakery/cake.py index 5ec9cfe..7b9e4e4 100644 --- a/layercake/bakery/cake.py +++ b/layercake/bakery/cake.py @@ -15,18 +15,25 @@ from contextlib import redirect_stdout import numpy as np +from numpy.linalg import LinAlgError import matplotlib.pyplot as plt import sparse as sp from numba import njit + +from sympy import ImmutableSparseNDimArray, MutableSparseNDimArray +from sympy import MutableSparseMatrix +from sympy import zeros as sympy_zeros +from sympy import simplify, N +from sympy.tensor.array import permutedims +from sympy.matrices.exceptions import NonInvertibleMatrixError + from layercake.utils.tensor import sparse_mul, jsparse_mul from layercake.utils.symbolic_tensor import get_coords_and_values_from_tensor, compute_jacobian_permutations from layercake.formatters.fortran import FortranJacobianEquationFormatter, FortranEquationFormatter from layercake.formatters.python import PythonJacobianEquationFormatter, PythonEquationFormatter from layercake.formatters.julia import JuliaJacobianEquationFormatter, JuliaEquationFormatter from layercake.utils import isin -from sympy import ImmutableSparseNDimArray, MutableSparseNDimArray -from sympy import simplify, N -from sympy.tensor.array import permutedims +from layercake.utils.symbolic_tensor import symbolic_tensordot real_eps = np.finfo(np.float64).eps small_number = 1.e-12 @@ -45,6 +52,7 @@ class Cake(object): def __init__(self): self.layers = list() + self._lhs_inversion_in_layer = True def add_layer(self, layer): """Add a layer object to the cake. @@ -120,7 +128,7 @@ def number_of_layers(self): return self.layers.__len__() def compute_tensor(self, numerical=True, compute_inner_products=False, compute_inner_products_kwargs=None, - substitutions=None, basis_subs=False, parameters_subs=None): + substitutions=None, basis_subs=False, parameters_subs=None, lhs_inversion_in_layer=True): """Compute the tensor of the symbolic or numerical representation of the ordinary differential equations tendencies of all the layers. Arguments are passed to the layer :meth:`~Layer.compute_tensor` method. @@ -148,8 +156,19 @@ def compute_tensor(self, numerical=True, compute_inner_products=False, compute_i parameters_subs: list(~parameter.Parameter), optional List of model's parameters to substitute in the symbolic tendencies' tensor. Only applies for the symbolic tendencies. + lhs_inversion_in_layer: bool, optional + Try to inverse the LHS matrix and take the matricial product with the RHS in each layer if possible. Default to `True`. + If `False`, it forces the inversion of the LHS at the cake level. """ + if lhs_inversion_in_layer: + self._lhs_inversion_in_layer = True + for layer in self.layers: + if layer.other_fields_in_lhs: + self._lhs_inversion_in_layer = False + break + else: + self._lhs_inversion_in_layer = False for layer in self.layers: layer.compute_tensor(numerical=numerical, @@ -157,7 +176,8 @@ def compute_tensor(self, numerical=True, compute_inner_products=False, compute_i compute_inner_products_kwargs=compute_inner_products_kwargs, substitutions=substitutions, basis_subs=basis_subs, - parameters_subs=parameters_subs + parameters_subs=parameters_subs, + lhs_inversion_in_layer=self._lhs_inversion_in_layer ) @property @@ -205,8 +225,12 @@ def tensor(self): if numerical: tensor = sp.zeros(shape, dtype=np.float64, format='dok') + if not self._lhs_inversion_in_layer: + lhs_mat = sp.zeros((self.ndim + 1, self.ndim + 1), dtype=np.float64, format='dok') else: tensor = MutableSparseNDimArray(iterable={}, shape=shape) + if not self._lhs_inversion_in_layer: + lhs_mat = MutableSparseMatrix(sympy_zeros(self.ndim + 1, self.ndim + 1)) for i, layer in enumerate(self.layers): if (numerical and not isinstance(layer.tensor, sp.COO) or @@ -225,13 +249,33 @@ def tensor(self): args = tuple(slices + zeros) if numerical: tensor[args] = tensor[args] + layer.tensor.todense()[1:] + if not self._lhs_inversion_in_layer: + lhs_mat[args[:2]] = lhs_mat[args[:1]] + layer._lhs_mat.todense()[1:] else: tensor[args] = tensor[args] + layer.tensor[1:] + if not self._lhs_inversion_in_layer: + lhs_mat[args[:2]] = lhs_mat[args[:2]] + layer._lhs_mat[1:, :] if numerical: - tensor = tensor.to_coo() + if not self._lhs_inversion_in_layer: + try: + lhs_mat_inverted = np.zeros((self.ndim + 1, self.ndim + 1)) + lhs_mat_inverted[1:, 1:] = np.linalg.inv(lhs_mat.todense()[1:, 1:]) + tensor = sp.COO(np.tensordot(lhs_mat_inverted, tensor.to_coo(), 1)) + except LinAlgError: + raise LinAlgError(f'The left-hand side of the cake is not invertible with the provided basis.') + else: + tensor = tensor.to_coo() else: - tensor = ImmutableSparseNDimArray(tensor) + if not self._lhs_inversion_in_layer: + try: + lhs_mat_inverted = MutableSparseMatrix(sympy_zeros(self.ndim + 1, self.ndim + 1)) + lhs_mat_inverted[1:, 1:] = lhs_mat[1:, 1:].inv().simplify() + tensor = ImmutableSparseNDimArray(symbolic_tensordot(lhs_mat_inverted, tensor, 1)) + except NonInvertibleMatrixError: + raise NonInvertibleMatrixError(f'The left-hand side of the cake is not invertible with the provided basis.') + else: + tensor = ImmutableSparseNDimArray(tensor) tensor = self.simplify_tensor(tensor, small_number) diff --git a/layercake/bakery/layers.py b/layercake/bakery/layers.py index 3ea036d..90f2f6e 100644 --- a/layercake/bakery/layers.py +++ b/layercake/bakery/layers.py @@ -18,6 +18,7 @@ import matplotlib.pyplot as plt import sparse as sp from copy import deepcopy +import warnings from sympy import MutableSparseNDimArray, MutableSparseMatrix, ImmutableMatrix, ImmutableSparseNDimArray from sympy import zeros as sympy_zeros @@ -37,7 +38,7 @@ class Layer(object): Parameters ---------- name: str, optional - Optional name for the layer. + Name for the layer. Attributes ---------- @@ -57,6 +58,9 @@ def __init__(self, name=''): self.name = name self._cake = None self._cake_order = 0 + self._lhs_inversion = True + self._lhs_inverted = False + self._lhs_mat = None @property def _cake_first_index(self): @@ -93,6 +97,30 @@ def fields(self): fields_list.append(eq.field) return fields_list + @property + def other_fields(self): + """list(~field.Field): Returns the list of dynamical fields of other layers, i.e. the fields whose time + evolution is provided by the partial differential equations of other layers.""" + layer_fields = self.fields + other_fields = list() + for eq in self.equations: + for field in eq.other_fields: + if field not in layer_fields: + other_fields.append(field) + return other_fields + + @property + def other_fields_in_lhs(self): + """list(~field.Field): Returns the list of dynamical fields of other layers present in the LHS of the layer's equations, + i.e. the fields whose time evolution is provided by the partial differential equations of other layers.""" + layer_fields = self.fields + other_fields = list() + for eq in self.equations: + for field in eq.other_fields_in_lhs: + if field not in layer_fields: + other_fields.append(field) + return other_fields + @property def parameters(self): """list(~parameter.Parameter): Returns the list of parameters of the layer, i.e. the explicit parameters @@ -144,7 +172,7 @@ def maximum_rank(self): def compute_inner_products(self, numerical=True, timeout=None, num_threads=None): """Compute the inner products tensors, either symbolic or numerical ones, of all the terms - of the layer equations, including the left-hand side term. + of the layer equations, including the left-hand side terms. Computations are parallelized on multiple CPUs. Parameters @@ -153,18 +181,26 @@ def compute_inner_products(self, numerical=True, timeout=None, num_threads=None) Whether to compute numerical or symbolic inner products. Default to `True` (numerical inner products as output). timeout: int or bool or None, optional - TODO + Control the switch from symbolic to numerical integration. By default, `parallel_integration` workers will try to integrate + |Sympy| expressions symbolically, but a fallback to numerical integration can be enforced. + The options are: + + * `None`: This is the "full-symbolic" mode. No timeout will be applied, and the switch to numerical integration will never happen. + Can result in very long and improbable computation time. + * `True`: This is the "full-numerical" mode. Symbolic computations do not occur, and the workers try directly to integrate + numerically. + * `False`: Same as `None`. + * An integer: defines a timeout after which, if a symbolic integration have not completed, the worker switch to the + numerical integration. num_threads: None or int, optional Number of CPUs to use in parallel for the computations. If `None`, use all the CPUs available. Default to `None`. """ for field, eq in zip(self.fields, self.equations): - eq.lhs_term.compute_inner_products(field.basis, numerical=numerical, timeout=timeout, num_threads=num_threads) - for term in eq.terms: - term.compute_inner_products(field.basis, numerical=numerical, timeout=timeout, num_threads=num_threads) + eq.compute_inner_products(field.basis, numerical=numerical, timeout=timeout, num_threads=num_threads) def compute_tensor(self, numerical=True, compute_inner_products=False, compute_inner_products_kwargs=None, - substitutions=None, basis_subs=False, parameters_subs=None): + substitutions=None, basis_subs=False, parameters_subs=None, lhs_inversion_in_layer=True): """Compute the tensor of the symbolic or numerical representation of the ordinary differential equations tendencies of the layer. Results are stored in the :attr:`~Layer.tensor` attribute. @@ -176,7 +212,7 @@ def compute_tensor(self, numerical=True, compute_inner_products=False, compute_i Whether to compute the numerical or the symbolic tensor. Default to `True` (numerical tensor as output). compute_inner_products: bool, optional - Whether the inner products tensors of the layer equations' terms must be compute first. + Whether the inner products tensors of the layer equations' terms must be computed first. Default to `False`. Please note that if the inner products are not computed firsthand, the tensor computation will fail. compute_inner_products_kwargs: dict, optional @@ -193,6 +229,8 @@ def compute_tensor(self, numerical=True, compute_inner_products=False, compute_i parameters_subs: list(~parameter.Parameter), optional List of model's parameters to substitute in the symbolic tendencies' tensor. Only applies for the symbolic tendencies. + lhs_inversion_in_layer: bool, optional + Try to inverse the LHS matrix and take the matricial product with the RHS. Default to `True`. """ @@ -218,15 +256,52 @@ def compute_tensor(self, numerical=True, compute_inner_products=False, compute_i if numerical: self.tensor = sp.zeros(shape, dtype=np.float64, format='dok') - lhs_mat = np.zeros((self.ndim+1, self.ndim+1)) + self._lhs_mat = sp.zeros((self.ndim+1, self.ndim+1), dtype=np.float64, format='dok') + lhs_mat_inverted = np.zeros((self.ndim+1, self.ndim+1)) lhs_order = 1 for field, eq in zip(self.fields, self.equations): ndim = field.state.__len__() - try: - lhs_mat[lhs_order:lhs_order + ndim, lhs_order:lhs_order + ndim] = np.linalg.inv(eq.lhs_term.inner_products.todense()) - except LinAlgError: - raise LinAlgError(f'The left-hand side of the equation {eq} is not invertible with the provided basis.') - for equation_term in eq.terms: + if self.other_fields_in_lhs or not lhs_inversion_in_layer: + self._lhs_inversion = False + if self._cake is None: + raise ValueError('Field(s) in the LHS are not found in the layer and there is no cake.') + if self._lhs_mat.shape[1] != self._cake.ndim + 1: + self._lhs_mat = sp.zeros((self.ndim + 1, self._cake.ndim + 1), dtype=np.float64, format='dok') + + for lhs_term in eq.lhs_terms: + ofield = lhs_term.field + ofield_order = 1 + ondim = ofield.state.__len__() + for tfield in self._cake.fields: + if ofield is tfield: + break + tndim = tfield.state.__len__() + ofield_order += tndim + else: + raise ValueError(f'Field {ofield} not found in the cake.') + self._lhs_mat[lhs_order:lhs_order + ndim, ofield_order:ofield_order + ondim] = \ + self._lhs_mat[lhs_order:lhs_order + ndim, ofield_order:ofield_order + ondim] + lhs_term.inner_products.todense() + + elif eq.other_fields_in_lhs: + for lhs_term in eq.lhs_terms: + ofield = lhs_term.field + ofield_order = 1 + ondim = ofield.state.__len__() + for tfield in self.fields: + if ofield is tfield: + break + tndim = tfield.state.__len__() + ofield_order += tndim + self._lhs_mat[lhs_order:lhs_order + ndim, ofield_order:ofield_order + ondim] = \ + self._lhs_mat[lhs_order:lhs_order + ndim, ofield_order:ofield_order + ondim] + lhs_term.inner_products.todense() + else: + try: + lhs_mat_inverted[lhs_order:lhs_order + ndim, lhs_order:lhs_order + ndim] = np.linalg.inv(eq.lhs_inner_products_addition.todense()) + self._lhs_inverted = True + except LinAlgError: + raise LinAlgError(f'The left-hand side of the equation {eq} is not invertible with the provided basis.') + + for equation_term in eq.rhs_terms: slices = [slice(lhs_order, lhs_order + ndim)] for term in equation_term.terms: term_field = term.field @@ -266,7 +341,17 @@ def compute_tensor(self, numerical=True, compute_inner_products=False, compute_i if np.any(increment != 0): self.tensor[args] = self.tensor[args] + increment lhs_order += ndim - self.tensor = sp.COO(np.tensordot(lhs_mat, self.tensor.to_coo(), 1)) + if self._lhs_inversion and not self._lhs_inverted: + try: + lhs_mat_inverted[1:, 1:] = np.linalg.inv(self._lhs_mat.todense()[1:, 1:]) + self._lhs_inverted = True + except LinAlgError: + raise LinAlgError(f'The left-hand side of the layer {self} is not invertible with the provided basis.') + if self._lhs_inverted: + self.tensor = sp.COO(np.tensordot(lhs_mat_inverted, self.tensor.to_coo(), 1)) + else: + warnings.warn(f'LHS of layer {self} has not been inverted. Its tensor hold the RHS tendencies alone.') + self.tensor = self.tensor.to_coo() else: b_subs = list() @@ -277,7 +362,8 @@ def compute_tensor(self, numerical=True, compute_inner_products=False, compute_i else: p_subs = list() self.tensor = MutableSparseNDimArray(iterable={}, shape=shape) - lhs_mat = MutableSparseMatrix(sympy_zeros(self.ndim + 1, self.ndim + 1)) + self._lhs_mat = MutableSparseMatrix(sympy_zeros(self.ndim + 1, self.ndim + 1)) + lhs_mat_inverted = MutableSparseMatrix(sympy_zeros(self.ndim + 1, self.ndim + 1)) lhs_order = 1 for field, eq in zip(self.fields, self.equations): bsb = field.basis.substitutions @@ -289,11 +375,44 @@ def compute_tensor(self, numerical=True, compute_inner_products=False, compute_i else: b_subs.append(sbsb) ndim = field.state.__len__() - try: - lhs_mat[lhs_order:lhs_order + ndim, lhs_order:lhs_order + ndim] = eq.lhs_term.inner_products.inverse().simplify() - except NonInvertibleMatrixError: - raise NonInvertibleMatrixError(f'The left-hand side of the equation {eq} is not invertible with the provided basis.') - for equation_term in eq.terms: + if self.other_fields_in_lhs or not lhs_inversion_in_layer: + self._lhs_inversion = False + if self._cake is None: + raise ValueError('Field(s) in the LHS are not found in the layer or in the cake.') + if self._lhs_mat.shape[1] != self._cake.ndim + 1: + self._lhs_mat = MutableSparseMatrix(sympy_zeros(self.ndim + 1, self._cake.ndim + 1)) + + for lhs_term in eq.lhs_terms: + ofield = lhs_term.field + ofield_order = 1 + ondim = ofield.state.__len__() + for tfield in self._cake.fields: + if ofield is tfield: + break + tndim = tfield.state.__len__() + ofield_order += tndim + self._lhs_mat[lhs_order:lhs_order + ndim, ofield_order:ofield_order + ondim] += lhs_term.inner_products + elif eq.other_fields_in_lhs: + for lhs_term in eq.lhs_terms: + ofield = lhs_term.field + ofield_order = 1 + ondim = ofield.state.__len__() + for tfield in self.fields: + if ofield is tfield: + break + tndim = tfield.state.__len__() + ofield_order += tndim + else: + raise ValueError(f'Field {ofield} not found in the cake.') + self._lhs_mat[lhs_order:lhs_order + ndim, ofield_order:ofield_order + ondim] += lhs_term.inner_products + else: + try: + lhs_mat_inverted[lhs_order:lhs_order + ndim, lhs_order:lhs_order + ndim] = eq.lhs_inner_products_addition.inv().simplify() + self._lhs_inverted = True + except NonInvertibleMatrixError: + raise NonInvertibleMatrixError(f'The left-hand side of the equation {eq} is not invertible with the provided basis.') + + for equation_term in eq.rhs_terms: slices = [slice(lhs_order, lhs_order + ndim)] for term in equation_term.terms: term_field = term.field @@ -356,8 +475,20 @@ def compute_tensor(self, numerical=True, compute_inner_products=False, compute_i increment = ImmutableSparseNDimArray(increment) self.tensor[args] = self.tensor[args] + increment lhs_order += ndim - self.tensor = (ImmutableSparseNDimArray(symbolic_tensordot(lhs_mat, self.tensor, 1)) - .subs(b_subs).subs(p_subs).subs(substitutions)) + + if self._lhs_inversion and not self._lhs_inverted: + try: + lhs_mat_inverted[1:, 1:] = self._lhs_mat[1:, 1:].inv().simplify() + self._lhs_inverted = True + except NonInvertibleMatrixError: + raise NonInvertibleMatrixError(f'The left-hand side of the layer {self} is not invertible with the provided basis.') + if self._lhs_inverted: + self.tensor = (ImmutableSparseNDimArray(symbolic_tensordot(lhs_mat_inverted, self.tensor, 1)) + .subs(b_subs).subs(p_subs).subs(substitutions)) + else: + warnings.warn(f'LHS of layer {self} has not been inverted. Its tensor hold the RHS tendencies alone.') + self.tensor = ImmutableSparseNDimArray(self.tensor).subs(b_subs).subs(p_subs).subs(substitutions) + self._lhs_mat = self._lhs_mat.subs(b_subs).subs(p_subs).subs(substitutions) def to_latex(self, enclose_lhs=True, drop_first_lhs_char=True, drop_first_rhs_char=False): """Generate the LaTeX strings representing the layer's equations mathematically. diff --git a/test/test_RP1982_qgs_numerical.py b/test/test_RP1982_qgs_numerical.py index 00f679e..5b5d5f1 100644 --- a/test/test_RP1982_qgs_numerical.py +++ b/test/test_RP1982_qgs_numerical.py @@ -184,7 +184,7 @@ def layercake_outputs(self, output_func=None): # Defining the equation and LHS # Laplacian vorticity = OperatorTerm(psi, Laplacian, atmospheric_basis.coordinate_system) - barotropic_equation = Equation(psi, lhs_term=vorticity) + barotropic_equation = Equation(psi, lhs_terms=vorticity) # Defining the advection term advection_term1 = vorticity_advection(psi, psi, atmospheric_basis.coordinate_system, sign=-1) @@ -228,7 +228,7 @@ def layercake_outputs(self, output_func=None): lin_lhs = LinearTerm(theta, prefactor=a, sign=-1) lhs = AdditionOfTerms(lin_lhs, vorticity) - baroclinic_equation = Equation(theta, lhs_term=lhs) + baroclinic_equation = Equation(theta, lhs_terms=lhs) # Defining the advection terms advection_term1 = vorticity_advection(psi, theta, atmospheric_basis.coordinate_system, sign=-1) diff --git a/test/test_RP1982_qgs_symbolic.py b/test/test_RP1982_qgs_symbolic.py index 46fd544..5e66354 100644 --- a/test/test_RP1982_qgs_symbolic.py +++ b/test/test_RP1982_qgs_symbolic.py @@ -185,7 +185,7 @@ def layercake_outputs(self, output_func=None): # Defining the equation and LHS # Laplacian vorticity = OperatorTerm(psi, Laplacian, atmospheric_basis.coordinate_system) - barotropic_equation = Equation(psi, lhs_term=vorticity) + barotropic_equation = Equation(psi, lhs_terms=vorticity) # Defining the advection term advection_term1 = vorticity_advection(psi, psi, atmospheric_basis.coordinate_system, sign=-1) @@ -229,7 +229,7 @@ def layercake_outputs(self, output_func=None): lin_lhs = LinearTerm(theta, prefactor=a, sign=-1) lhs = AdditionOfTerms(lin_lhs, vorticity) - baroclinic_equation = Equation(theta, lhs_term=lhs) + baroclinic_equation = Equation(theta, lhs_terms=lhs) # Defining the advection terms advection_term1 = vorticity_advection(psi, theta, atmospheric_basis.coordinate_system, sign=-1) diff --git a/test/test_maooam_qgs_numerical.py b/test/test_maooam_qgs_numerical.py index 821e9a4..efe4ade 100644 --- a/test/test_maooam_qgs_numerical.py +++ b/test/test_maooam_qgs_numerical.py @@ -227,7 +227,7 @@ def layercake_outputs(self, output_func=None): # Defining the equation and LHS # Laplacian vorticity = OperatorTerm(psi_a, Laplacian, atmospheric_basis.coordinate_system) - barotropic_equation = Equation(psi_a, lhs_term=vorticity) + barotropic_equation = Equation(psi_a, lhs_terms=vorticity) # Defining the advection term advection_term1 = vorticity_advection(psi_a, psi_a, atmospheric_basis.coordinate_system, sign=-1) @@ -262,7 +262,7 @@ def layercake_outputs(self, output_func=None): lin_lhs = LinearTerm(theta_a, prefactor=a, sign=-1) lhs = AdditionOfTerms(lin_lhs, vorticity) - baroclinic_equation = Equation(theta_a, lhs_term=lhs) + baroclinic_equation = Equation(theta_a, lhs_terms=lhs) # Defining the advection terms advection_term1 = vorticity_advection(psi_a, theta_a, atmospheric_basis.coordinate_system, sign=-1) @@ -322,7 +322,7 @@ def layercake_outputs(self, output_func=None): vorticity = OperatorTerm(psi_o, Laplacian, oceanic_basis.coordinate_system) lin_lhs = LinearTerm(psi_o, prefactor=G) lhs = AdditionOfTerms(lin_lhs, vorticity) - oceanic_equation = Equation(psi_o, lhs_term=lhs) + oceanic_equation = Equation(psi_o, lhs_terms=lhs) # Defining the advection term advection_term = vorticity_advection(psi_o, psi_o, oceanic_basis.coordinate_system, sign=-1) @@ -355,7 +355,7 @@ def layercake_outputs(self, output_func=None): # Defining the equation and LHS # Laplacian lhs = LinearTerm(deltaT_o) - ocean_temperature_equation = Equation(deltaT_o, lhs_term=lhs) + ocean_temperature_equation = Equation(deltaT_o, lhs_terms=lhs) # Defining the advection term advection_term = Jacobian(psi_o, deltaT_o, oceanic_basis.coordinate_system, sign=-1) diff --git a/test/test_maooam_qgs_symbolic.py b/test/test_maooam_qgs_symbolic.py index 515f19f..817e8b8 100644 --- a/test/test_maooam_qgs_symbolic.py +++ b/test/test_maooam_qgs_symbolic.py @@ -229,7 +229,7 @@ def layercake_outputs(self, output_func=None): # Defining the equation and LHS # Laplacian vorticity = OperatorTerm(psi_a, Laplacian, atmospheric_basis.coordinate_system) - barotropic_equation = Equation(psi_a, lhs_term=vorticity) + barotropic_equation = Equation(psi_a, lhs_terms=vorticity) # Defining the advection term advection_term1 = vorticity_advection(psi_a, psi_a, atmospheric_basis.coordinate_system, sign=-1) @@ -264,7 +264,7 @@ def layercake_outputs(self, output_func=None): lin_lhs = LinearTerm(theta_a, prefactor=a, sign=-1) lhs = AdditionOfTerms(lin_lhs, vorticity) - baroclinic_equation = Equation(theta_a, lhs_term=lhs) + baroclinic_equation = Equation(theta_a, lhs_terms=lhs) # Defining the advection terms advection_term1 = vorticity_advection(psi_a, theta_a, atmospheric_basis.coordinate_system, sign=-1) @@ -324,7 +324,7 @@ def layercake_outputs(self, output_func=None): vorticity = OperatorTerm(psi_o, Laplacian, oceanic_basis.coordinate_system) lin_lhs = LinearTerm(psi_o, prefactor=G) lhs = AdditionOfTerms(lin_lhs, vorticity) - oceanic_equation = Equation(psi_o, lhs_term=lhs) + oceanic_equation = Equation(psi_o, lhs_terms=lhs) # Defining the advection term advection_term = vorticity_advection(psi_o, psi_o, oceanic_basis.coordinate_system, sign=-1) @@ -357,7 +357,7 @@ def layercake_outputs(self, output_func=None): # Defining the equation and LHS # Laplacian lhs = LinearTerm(deltaT_o) - ocean_temperature_equation = Equation(deltaT_o, lhs_term=lhs) + ocean_temperature_equation = Equation(deltaT_o, lhs_terms=lhs) # Defining the advection term advection_term = Jacobian(psi_o, deltaT_o, oceanic_basis.coordinate_system, sign=-1)