Source code for simulai.utilities.lorenz_solver

# (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 sys
import os
import numpy as np
from simulai.utilities.problem_classes import LorenzSystem
from simulai.math.integration import RK4
import matplotlib.pyplot as plt
from fractions import Fraction
from argparse import ArgumentParser

# Testing to solve a nonlinear oscillator problem using
# a 4th order and 4 steps Runge-Kutta


[docs]def lorenz_solver(rho, beta, sigma, T, dt, data_path, initial_state, solver='RK4', **kwargs) -> object: """ Parameters ---------- rho: float beta: sigma T: int dt: float data_path: str initial_state: List[float] Returns ------- (np.ndarray, np.ndarray) the full path to the directory that was created to store the results of the simulation """ if not os.path.isdir(data_path): os.makedirs(data_path, exist_ok=True) elif data_path=="on_memory": data_path = "no_path" beta_str = kwargs.get('beta_str').replace('/', '_') # Using the built-in RK45 implementation if solver == 'RK4': problem = LorenzSystem(rho, sigma, beta) time = np.arange(0, T, dt) solver = RK4(problem) variables_timesteps = list() derivatives_timesteps = list() current_state = initial_state iter = 0 for tt in range(time.shape[0]): variables_state, derivatives_state = solver.step(current_state, dt) variables_timesteps.append(variables_state) derivatives_timesteps.append(derivatives_state) current_state = variables_state sys.stdout.write("\rIteration {}".format(iter)) sys.stdout.flush() iter += 1 variables_matrix = np.vstack(variables_timesteps) derivatives_matrix = np.vstack(derivatives_timesteps) # Using the SciPy RK4 implementation elif solver == 'RK45': from scipy.integrate import solve_ivp from simulai.math.differentiation import CollocationDerivative problem = LorenzSystem(rho, sigma, beta, use_t=True) time = np.arange(0, T, dt) current_state = initial_state solution = solve_ivp(problem.eval, [0, T], current_state[0], max_step=dt, method='RK45') variables_matrix = solution['y'].T time = solution['t'] # After the time-integrations, it is necessary to evaluate the derivatives numerically diff = CollocationDerivative(config={}) derivatives_matrix = diff.solve(data=variables_matrix, x=time) plt.plot(time, variables_matrix[:, 0], label="x") plt.plot(time, variables_matrix[:, 1], label="y") plt.plot(time, variables_matrix[:, 2], label="z") label_string = "rho_{}_sigma_{}_beta_{}_T_{}_dt_{}/".format(rho, sigma, beta_str, T, dt) all_variables_matrix = np.hstack([variables_matrix, derivatives_matrix]) n_col = all_variables_matrix.shape[1] all_variables_record = np.core.records.fromarrays(np.split(all_variables_matrix, n_col, axis=1), names='x, y, z, dx_dt, dy_dt, dz_dt', formats='f8, f8, f8, f8, f8, f8') if data_path != "no_path": dir_path = os.path.join(data_path, label_string) if not os.path.isdir(dir_path): os.mkdir(dir_path) np.save(os.path.join(data_path, "{}Lorenz_data.npy".format(label_string)), all_variables_record) plt.xlabel("Time(s)") plt.title("Lorenz System") plt.legend() plt.grid(True) plt.close() if solver == 'RK45': return variables_matrix, derivatives_matrix, time else: return variables_matrix, derivatives_matrix
[docs]def lorenz_solver_forcing(rho, beta, sigma, T, dt, data_path, initial_state, forcing=None, **kwargs): """ Parameters ---------- rho: float beta: sigma T: int dt: float data_path: str initial_state: List[float] Returns ------- (np.ndarray, np.ndarray) the full path to the directory that was created to store the results of the simulation """ if not os.path.isdir(data_path): os.makedirs(data_path, exist_ok=True) elif data_path=="on_memory": data_path = "no_path" beta_str = kwargs.get('beta_str').replace('/', '_') problem = LorenzSystem(rho, sigma, beta, forcing=True) solver = RK4(problem) time = np.arange(0, T, dt) variables_timesteps = list() derivatives_timesteps = list() current_state = initial_state iter = 0 for tt in range(time.shape[0]): variables_state, derivatives_state = solver.step_with_forcings(current_state, forcing[tt:tt+1], dt) variables_timesteps.append(variables_state) derivatives_timesteps.append(derivatives_state) current_state = variables_state sys.stdout.write("\rIteration {}".format(iter)) sys.stdout.flush() iter += 1 variables_matrix = np.vstack(variables_timesteps) derivatives_matrix = np.vstack(derivatives_timesteps) plt.plot(time, variables_matrix[:, 0], label="x") plt.plot(time, variables_matrix[:, 1], label="y") plt.plot(time, variables_matrix[:, 2], label="z") label_string = "rho_{}_sigma_{}_beta_{}_T_{}_dt_{}/".format(rho, sigma, beta_str, T, dt) all_variables_matrix = np.hstack([variables_matrix, derivatives_matrix]) n_col = all_variables_matrix.shape[1] all_variables_record = np.core.records.fromarrays(np.split(all_variables_matrix, n_col, axis=1), names='x, y, z, dx_dt, dy_dt, dz_dt', formats='f8, f8, f8, f8, f8, f8') if data_path != "no_path": dir_path = os.path.join(data_path, label_string) if not os.path.isdir(dir_path): os.mkdir(dir_path) np.save(os.path.join(data_path, "{}Lorenz_data.npy".format(label_string)), all_variables_record) plt.xlabel("Time(s)") plt.title("Lorenz System") plt.legend() plt.grid(True) plt.close() return variables_matrix, derivatives_matrix
if __name__ == "__main__": parser = ArgumentParser(description="Reading input arguments") parser.add_argument('--data_path', type=str) parser.add_argument('--time', type=float) parser.add_argument('--dt', type=float) parser.add_argument('--rho', type=float) parser.add_argument('--sigma', type=float) parser.add_argument('--beta', type=str) args = parser.parse_args() data_path = args.data_path T = args.time dt = args.dt initial_state = np.array([1, 2, 3])[None, :] rho = args.rho #28 beta_str = args.beta beta = float(Fraction(beta_str)) #8/3 sigma = args.sigma #10 lorenz_solver(rho=rho, beta=beta, sigma=sigma, T=T, dt=dt, data_path=data_path, initial_state=initial_state, beta_str=beta_str )