Source code for simulai.math.integration

# (C) Copyright IBM Corp. 2019, 2020, 2021, 2022.

#    Licensed under the Apache License, Version 2.0 (the "License");
#    you may not use this file except in compliance with the License.
#    You may obtain a copy of the License at

#           http://www.apache.org/licenses/LICENSE-2.0

#     Unless required by applicable law or agreed to in writing, software
#     distributed under the License is distributed on an "AS IS" BASIS,
#     WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#     See the License for the specific language governing permissions and
#     limitations under the License.

import numpy as np
import sys
from scipy.integrate import odeint

from simulai.abstract import Integral

# Parent class for explicit time-integrators
[docs]class ExplicitIntegrator(Integral): name = "int" def __init__(self, coeffs:np.ndarray, weights:np.ndarray, right_operator:callable) -> None: """ Explicit time-integrator parent class :param coeffs: coefficients to shift time during integration :type coeffs: np.ndarray :param weights: the weights for penalizing each shifted state :type weights: np.ndarray :param right_operator: the callable for evaluating the residual in each iteration :returns: nothing """ super().__init__() self.coeffs = coeffs self.weights = weights self.right_operator = right_operator self.n_stages = len(self.coeffs) self.log_phrase = "Executing integrator "
[docs] def step(self, variables_state_initial:np.ndarray, dt:float) -> (np.ndarray, np.ndarray): """It marches a step in the time-integration :param variables_state_initial: initial state :type variables_state_initial: np.ndarray :param dt: timestep size :type dt: float :returns: the integrated state and its time-derivative :rtype: (np.ndarray, np.ndarray) """ variables_state = variables_state_initial residuals_list = np.zeros((self.n_stages,) + variables_state.shape[1:]) k_weighted = None for stage in range(self.n_stages): k = self.right_operator(variables_state) residuals_list[stage, :] = k k_weighted = self.weights[stage].dot(residuals_list) variables_state = variables_state_initial + self.coeffs[stage] * dt * k_weighted return variables_state, k_weighted
[docs] def step_with_forcings(self, variables_state_initial:np.ndarray, forcing_state:np.ndarray, dt:float) -> (np.ndarray, np.ndarray): """It marches a step in the time-integration using a concatenated [variables, forcings] state :param variables_state_initial: initial state :type variables_state_initial: np.ndarray :param forcing_state: the state of the forcing terms :type forcing_state: np.ndarray :param dt: timestep size :type dt: float :returns: the integrated state and its time-derivative :rtype: (np.ndarray, np.ndarray) """ variables_state = np.concatenate([variables_state_initial, forcing_state], axis=-1) residuals_list = np.zeros((self.n_stages,) + variables_state_initial.shape[1:]) k_weighted = None for stage in range(self.n_stages): k = self.right_operator(variables_state) residuals_list[stage, :] = k k_weighted = self.weights[stage].dot(residuals_list) variables_state_ = variables_state_initial + self.coeffs[stage] * dt * k_weighted variables_state = np.concatenate([variables_state_, forcing_state], axis=1) return variables_state_, k_weighted
[docs] def step_with_forcings_separated(self, variables_state_initial:np.ndarray, forcing_state:np.ndarray, dt:float) -> (np.ndarray, np.ndarray): """It marches a step in the time-integration for variables and forcings being operated separately :param variables_state_initial: initial state :type variables_state_initial: np.ndarray :param forcing_state: the state of the forcing terms :type forcing_state: np.ndarray :param dt: timestep size :type dt: float :returns: the integrated state and its time-derivative :rtype: (np.ndarray, np.ndarray) """ variables_state = {'input_data': variables_state_initial, 'forcing_data': forcing_state} residuals_list = np.zeros((self.n_stages,) + variables_state_initial.shape[1:]) k_weighted = None for stage in range(self.n_stages): k = self.right_operator(**variables_state) residuals_list[stage, :] = k k_weighted = self.weights[stage].dot(residuals_list) variables_state_ = variables_state_initial + self.coeffs[stage] * dt * k_weighted variables_state = {'input_data': variables_state_, 'forcing_data': forcing_state} return variables_state_, k_weighted
# Looping over multiple steps without using forcings def _loop(self, initial_state:np.ndarray, epochs:int, dt:float) -> list: """No forced time-integration loop :param initial_state: the initial state :type initial_state: np.ndarray :param epochs: number of steps to be used in the time-integration :type epochs: int :param dt: the timestep size :type dt: float :returns: the list of integrated states :rtype: list """ ii = 0 integrated_variables = list() while ii < epochs: state, derivative_state = self.step(initial_state, dt) integrated_variables.append(state) initial_state = state sys.stdout.write("\r {}, iteration: {}/{}".format(self.log_phrase, ii+1, epochs)) sys.stdout.flush() ii += 1 return integrated_variables # Looping over multiple steps using forcings def _loop_forcings(self, initial_state:np.ndarray, forcings:np.ndarray, epochs:int, dt:float) -> list: """Forced time-integration loop :param initial_state: the initial state :type initial_state: np.ndarray :param forcings: the array containing all the forcing states :type forcings: np.ndarray :param epochs: number of steps to be used in the time-integration :type epochs: int :param dt: the timestep size :type dt: float :returns: the list of integrated states :rtype: list """ ii = 0 integrated_variables = list() while ii < epochs: forcings_state = forcings[ii][None, :] state, derivative_state = self.step_with_forcings(initial_state, forcings_state, dt) integrated_variables.append(state) initial_state = state sys.stdout.write("\r {}, iteration: {}/{}".format(self.log_phrase, ii + 1, epochs)) sys.stdout.flush() ii += 1 return integrated_variables def __call__(self, initial_state:np.ndarray=None, epochs:int=None, dt:float=None, resolution:float=None, forcings:np.ndarray=None) -> np.ndarray: """It determined the proper time-integration loop to be used and executes it. :param initial_state: the initial state :type initial_state: np.ndarray :param forcings: the array containing all the forcing states :type forcings: np.ndarray :param epochs: number of steps to be used in the time-integration :type epochs: int :param dt: the timestep size :type dt: float :returns: the list of integrated states :rtype: list """ if forcings is None: integrated_variables = self._loop(initial_state, epochs, dt) else: integrated_variables = self._loop_forcings(initial_state, forcings, epochs, dt) if resolution is None: resolution_step = 1 elif resolution >= dt: resolution_step = int(resolution/dt) else: raise Exception("Resolution is lower than the discretization step.") return np.vstack(integrated_variables)[::resolution_step]
# Built-in Runge-Kutta 4th order
[docs]class RK4(ExplicitIntegrator): name = "rk4_int" def __init__(self, right_operator): """ Runge-Kutta 4th-order time-integrator :param right_operator: an operator for representing the right-hand side of a dynamic system :type right_operator: callable (def or lambda) """ weights = np.array( [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [1/6, 2/6, 2/6, 1/6]] ) coeffs = np.array([0.5, 0.5, 1, 1]) super().__init__(coeffs, weights, right_operator) self.log_phrase += "Runge-Kutta 4th order "
# Wrapper for using the SciPy LSODA (LSODA itself is a wrapper for ODEPACK)
[docs]class LSODA: def __init__(self, right_operator:callable) -> None: self.right_operator = right_operator
[docs] def run(self, current_state:np.ndarray=None, t:np.ndarray=None) -> np.ndarray: if hasattr(self.right_operator, "jacobian"): Jacobian = self.right_operator.jacobian else: Jacobian = None solution = odeint(self.right_operator.eval, current_state, t, Dfun=Jacobian) return solution
# Wrapper for handling function objects in time-integrators
[docs]class FunctionWrapper: def __init__(self, function:callable, extra_dim:bool=True) -> None: self.function = function if extra_dim is True: self.prepare_input = self._extra_dim_prepare_input else: self.prepare_input = self._no_extra_dim_prepare_input if extra_dim is True: self.prepare_output = self._extra_dim_prepare_output else: self.prepare_output = self._no_extra_dim_prepare_output def _extra_dim_prepare_input(self, input_data:np.ndarray) -> np.ndarray: return input_data[None, :] def _no_extra_dim_prepare_input(self, input_data:np.ndarray) -> np.ndarray: return input_data def _extra_dim_prepare_output(self, output_data:np.ndarray) -> np.ndarray: return output_data[0, :] def _no_extra_dim_prepare_output(self, output_data:np.ndarray) -> np.ndarray: return output_data def __call__(self, input_data:np.ndarray) -> np.ndarray: input_data = self.prepare_input(input_data) return self.prepare_output(self.function(input_data))
# Wrapper for handling class objects in time-integrators
[docs]class ClassWrapper: def __init__(self, class_instance:callable) -> None: assert hasattr(class_instance, 'eval'), f"The objet class_instance={class_instance} has no attribute eval." self.class_instance = class_instance if hasattr(self.class_instance, 'jacobian'): self.jacobian = self._jacobian else: pass def __call__(self, input_data:np.ndarray) -> np.ndarray: return self.class_instance(input_data)[0, :]
[docs] def eval(self, input_data:np.ndarray, t:float) -> np.ndarray: input_data = input_data return self.class_instance.eval(input_data[None, :])[0, :]
def _jacobian(self, input_data:np.ndarray, t:float) -> np.ndarray: print("Using jacobian: stiffness alert.") return self.class_instance.jacobian(input_data)