import os import sys import time from pathlib import Path import numpy as np import scipy as sp from scipy.sparse.linalg import cg from scipy.sparse.linalg import gmres from scipy.sparse.linalg import spsolve_triangular from pheromone_dispersion.advection_operator import Advection from pheromone_dispersion.diffusion_operator import Diffusion from pheromone_dispersion.identity_operator import Id from pheromone_dispersion.reaction_operator import Reaction from pheromone_dispersion.source_term import Source """ Ajouter des tests dans le cas stationnaire pour verifier que les matrices sont 3D """ class DiffusionConvectionReaction2DEquation: """ Class containing the 2D diffusion-convection-reaction PDE and its solvers """ def __init__(self, U, K, coeff_depot, S, msh, time_discretization='semi-implicit', tol_inversion=1e-14): """ Instanciation of the class. - input: * msh: object of the class MeshRect2D * U: object of the class Velocity * K: object of the class DiffusionTensor * coeff_depot: array of shape (msh.y, msh.x) containing the deposition coefficient * S: object of the class Source, contains the source term * time_discretization: string, describes the type of time discretization by default is 'semi-implicit', can also be 'implicit' - attributes: * msh: object of the class MeshRect2D * A: object of the class Advection, contains the convection linear operator * D: oject of the class Diffusion, contains the diffusion linear operator * R: object of the class Reaction, contains the reaction linear operator * ID: object of the class Id, contains the identity operator * S: object of the class Source, contains the source term * time_discretization: string, describes the type of time discretization - TO DO: * Add exceptions in case the matrix are computed/initialized """ self.msh = msh self.A = Advection(U, msh) self.D = Diffusion(K, msh) self.R = Reaction(coeff_depot, msh) self.Id = Id(msh) self.S = S self.time_discretization = time_discretization self.tol_inversion = tol_inversion implemented_solver_type = [ 'implicit', 'semi-implicit', 'implicit with matrix inversion', 'semi-implicit with matrix inversion', 'implicit with preconditionning', 'implicit with stationnary matrix inversion', ] if self.time_discretization not in implemented_solver_type: raise ValueError("The given time discretization is not implemented.") self.inv_matrix_semi_implicit_scheme = None self.inv_matrix_implicit_scheme = None self.ILU_decomp_dic = {} self.jacobi = None self.inv_matrix_implicit_scheme_t_init = None def spsolve_lu(self, b): if self.ILU_decomp_dic['perm r'] is not None: b_old = b.copy() for old_ndx, new_ndx in enumerate(self.ILU_decomp_dic['perm r']): b[new_ndx] = b_old[old_ndx] try: # unit_diagonal is a new kw c = spsolve_triangular(self.ILU_decomp_dic['L'], b, lower=True, unit_diagonal=True) except TypeError: c = spsolve_triangular(self.ILU_decomp_dic['L'], b, lower=True) px = spsolve_triangular(self.ILU_decomp_dic['U'], c, lower=False) if self.ILU_decomp_dic['perm c'] is None: return px return px[self.ILU_decomp_dic['perm c']] def init_inverse_matrix(self, path_to_matrix=None, matrix_file_name=None): if path_to_matrix is None: path_to_matrix = './data' if not os.path.isdir(path_to_matrix): os.makedirs(path_to_matrix) if matrix_file_name is None: if self.time_discretization == 'semi-implicit with matrix inversion': matrix_file_name = 'inv_matrix_semi_implicit_scheme' if self.time_discretization == 'implicit with matrix inversion': matrix_file_name = 'inv_matrix_implicit_scheme' if self.time_discretization == 'implicit with stationnary matrix inversion': matrix_file_name = 'inv_matrix_implicit_scheme' matrix_file_name += '.npy' if not (Path(path_to_matrix) / matrix_file_name).exists(): print("=== Computation of the inverse of the matrix of the implicit part of the " + self.time_discretization + " scheme ===") Identity = np.identity(self.msh.y.size * self.msh.x.size) if self.time_discretization == 'semi-implicit with matrix inversion': matrix_semi_implicit_scheme = Identity + self.msh.dt * (-self.D._matmat(Identity) + self.R._matmat(Identity)) self.inv_matrix_semi_implicit_scheme = sp.linalg.inv(matrix_semi_implicit_scheme) np.save(Path(path_to_matrix) / matrix_file_name + '.npy', self.inv_matrix_semi_implicit_scheme) if self.time_discretization == 'implicit with matrix inversion': t_i = time.time() for it, self.msh.t in enumerate(self.msh.t_array[1:]): self.at_current_time(self.msh.t) matrix_implicit_scheme = Identity + self.msh.dt * ( -self.D._matmat(Identity) + self.R._matmat(Identity) + self.A._matmat(Identity) ) self.inv_matrix_implicit_scheme = sp.linalg.inv(matrix_implicit_scheme) print("--- Computation at time t= ", self.msh.t, "in ", time.time() - t_i, " s") t_i = time.time() np.save( Path(path_to_matrix) / matrix_file_name + f'_ite{"{0:04d}".format(it+1)}' + '.npy', self.inv_matrix_implicit_scheme ) if self.time_discretization == 'implicit with stationnary matrix inversion': t_i = time.time() matrix_implicit_scheme = Identity + self.msh.dt * ( -self.D._matmat(Identity) + self.R._matmat(Identity) + self.A._matmat(Identity) ) self.inv_matrix_implicit_scheme = sp.linalg.inv(matrix_implicit_scheme) print("--- Computation at time t= ", self.msh.t, "in ", time.time() - t_i, " s") t_i = time.time() np.save(Path(path_to_matrix) / matrix_file_name + '.npy', self.inv_matrix_implicit_scheme) else: print("=== Load of the inverse of the matrix of the implicit part of the " + self.time_discretization + " scheme ===") if self.time_discretization == 'semi-implicit with matrix inversion': self.inv_matrix_semi_implicit_scheme = np.load(Path(path_to_matrix) / matrix_file_name) if self.time_discretization == 'implicit with matrix inversion': self.inv_matrix_implicit_scheme = np.load(Path(path_to_matrix) / matrix_file_name) if self.time_discretization == 'implicit with stationnary matrix inversion': self.inv_matrix_implicit_scheme = np.load(Path(path_to_matrix) / matrix_file_name) if self.time_discretization == 'implicit with preconditionning': self.inv_matrix_implicit_scheme_t_init = np.load(Path(path_to_matrix) / matrix_file_name) def set_source(self, value, t=None): """ Update the source term with the values provided as input - input: * value: numpy array of shape (t.size, msh.y.size, msh.x.size), the new values of the source term * t: numpy array of shape (t.size,), the vector containing the associated times, by default None if the source is stationnary """ setattr(self, 'S', Source(self.msh, value, t=t)) def at_current_time(self, tc, i=None): """ Update the linear operators of the PDE (attributes A, D and S of the class) at a given time. - input: * tc: the current time """ self.S.at_current_time(tc) if not self.time_discretization == 'implicit with stationnary matrix inversion': self.A.at_current_time(tc, self.A.U) self.D.at_current_time(tc) # if self.time_discretization == 'implicit with preconditionning': # #dic_name = 'ILU_decomp_dic' # path = './output' # matrix_name = 'Jacobi_matrix' # matrix_file_name = matrix_name + f'_ite{"{0:04d}".format(i+1)}.npz' # self.jacobi = sps.load_npz(Path(path) / matrix_file_name) # dic_file_name = dic_name + f'_ite{"{0:04d}".format(i+1)}.pklz' # with gzip.open(Path(path) / dic_file_name, 'rb') as f: # self.ILU_decomp_dic = pickle.load(f) def solver_one_time_step(self, c): """ solve the PDE for one time step inverse the linear system (I + dt (- D + R)).c^{n+1} = c^n + dt (- A.c^n + S) - input: * c: numpy array of shape (msh.y.size, msh.x.size), contains the concentration of pheromones at the previous time step - output: * numpy array of shape (msh.y.size, msh.x.size), resulting concentration at the current time step """ if c.shape != (self.msh.y.size, self.msh.x.size): raise ValueError( "The shape of the concentration at the previous time at the center of the cells does not match with the shape of the msh." ) # ravel the matrix of concentration into a msh.y.size*msh.x.size size vector to match with the format of the LinearOperator class c_ravel = c.ravel() # inverse the linear system using a conjugate gradient method c_ravel, tol = cg(self.Id + self.msh.dt * (-self.D + self.R), c_ravel + self.msh.dt * (-self.A.dot(c_ravel) + self.S.value.ravel())) return c_ravel.reshape((c.shape)) def solver(self, save_all=False, path_save='./save/', display_flag=True): """ solve the PDE on the whole time window - input: * save_all: boolean, False by default, the matrix of the concentration at every time step and the vector of time are saved in numpy arrays if True * path_save: string, './save/' by default, contains the path of the directory in which the ouputs are saved * display_flag: boolean, True by default, print the evolution in time of the solver if True - output: * c: numpy array of shape (msh.t_array.size, msh.y.size, msh.x.size), contains the concentration at all time steps """ # initialization of the unknown variable at the current time c = np.zeros((self.msh.y.shape[0] * self.msh.x.shape[0],)) # initialization of the outputs array to be saved if save_all: # if the save directory does not exist, then it is created if not os.path.isdir(path_save): os.makedirs(path_save) c_save = np.zeros((self.msh.t_array.shape[0], self.msh.y.shape[0], self.msh.x.shape[0])) # loop until the final time or the steady state is reached for it, self.msh.t in enumerate(self.msh.t_array[1:]): if display_flag: sys.stdout.write(f'\rt = {"{:.3f}".format(self.msh.t)} / {"{:.3f}".format(self.msh.T_final)} s') sys.stdout.flush() # update the coefficients of the equation at the current time and self.at_current_time(self.msh.t, i=it) # store the concentration at the previous time step c_old = np.copy(c) # NECESSARY??? # inverse the linear system resulting the semi-implicit time discretization # using a conjugate gradient method for the current time step if self.time_discretization == 'semi-implicit with matrix inversion': c = self.inv_matrix_semi_implicit_scheme.dot(c_old + self.msh.dt * (-self.A.matvec(c_old) + self.S.value.ravel())) info = 0 elif self.time_discretization == 'semi-implicit': c, info = cg( self.Id + self.msh.dt * (-self.D + self.R), c_old + self.msh.dt * (-self.A.matvec(c_old) + self.S.value.ravel()), x0=c_old, tol=self.tol_inversion, ) # inverse the linear system resulting the implicit time discretization using a gmres method for the current time step elif self.time_discretization == 'implicit': c, info = gmres( self.Id + self.msh.dt * (-self.D + self.R + self.A), c_old + self.msh.dt * self.S.value.ravel(), x0=c_old, tol=self.tol_inversion, ) elif self.time_discretization == 'implicit with preconditionning': precond = ( self.inv_matrix_implicit_scheme_t_init # self.jacobi # spsl.LinearOperator((self.msh.x.size*self.msh.y.size,self.msh.x.size*self.msh.y.size),matvec=self.spsolve_lu) ) c, info = gmres( self.Id + self.msh.dt * (-self.D + self.R + self.A), c_old + self.msh.dt * self.S.value.ravel(), x0=c_old, tol=self.tol_inversion, M=precond, ) elif self.time_discretization == 'implicit with matrix inversion': c = self.inv_matrix_implicit_scheme[it, :, :].dot(c_old + self.msh.dt * self.S.value.ravel()) info = 0 elif self.time_discretization == 'implicit with stationnary matrix inversion': RHS = c_old + self.msh.dt * self.S.value.ravel() LHS = (self.Id + self.msh.dt * (-self.D + self.R + self.A)).matvec(c_old) flag_residu = not np.linalg.norm(RHS - LHS) < self.tol_inversion * np.linalg.norm(RHS) if flag_residu: c = np.dot(self.inv_matrix_implicit_scheme, RHS) info = 0 if info > 0: raise ValueError( "The algorithme used to solve the linear system has not converge" + "to the expected tolerance or within the maximum number of iteration." ) if info < 0: raise ValueError("The algorithme used to solve the linear system could not proceed du to illegal input or breakdown.") c_save[it + 1, :, :] = c.reshape((self.msh.y.shape[0], self.msh.x.shape[0])) # save the ouputs if save_all: if not os.path.isdir(Path(path_save)): os.makedirs(Path(path_save)) np.save(Path(path_save) / 'c.npy', c_save) return c_save def solver_save_all(self, path_save='./save/', display_flag=True, save_rate=1000): """ solve the PDE on the whole time window - input: * save_all: boolean, False by default, the matrix of the concentration at every time step and the vector of time are saved in numpy arrays if True * path_save: string, './save/' by default, contains the path of the directory in which the ouputs are saved * display_flag: boolean, True by default, print the evolution in time of the solver if True - output: * c: numpy array of shape (msh.t_array.size, msh.y.size, msh.x.size), contains the concentration at all time steps """ # initialization of the unknown variable at the current time c = np.zeros((self.msh.y.shape[0] * self.msh.x.shape[0],)) t_save = [] c_save = [] # np.zeros((, self.msh.y.shape[0] * self.msh.x.shape[0],)) # initialization of the outputs array to be saved # if the save directory does not exist, then it is created if not os.path.isdir(path_save): os.makedirs(path_save) # loop until the final time or the steady state is reached for it, self.msh.t in enumerate(self.msh.t_array[1:]): if display_flag: # print(f"\rt = {"{:.3f}".format(self.msh.t)} / {"{:.3f}".format(self.msh.T_final)} s") sys.stdout.write(f'\n\rt = {"{:.3f}".format(self.msh.t)} / {"{:.3f}".format(self.msh.T_final)} s') sys.stdout.flush() # update the coefficients of the equation at the current time and self.at_current_time(self.msh.t, i=it) # store the concentration at the previous time step c_old = np.copy(c) # NECESSARY??? # inverse the linear system resulting the semi-implicit time discretization # using a conjugate gradient method for the current time step if self.time_discretization == 'semi-implicit with matrix inversion': c = self.inv_matrix_semi_implicit_scheme.dot(c_old + self.msh.dt * (-self.A.matvec(c_old) + self.S.value.ravel())) info = 0 elif self.time_discretization == 'semi-implicit': c, info = cg( self.Id + self.msh.dt * (-self.D + self.R), c_old + self.msh.dt * (-self.A.matvec(c_old) + self.S.value.ravel()), x0=c_old, tol=self.tol_inversion, ) # inverse the linear system resulting the implicit time discretization using a gmres method for the current time step elif self.time_discretization == 'implicit': c, info = gmres( self.Id + self.msh.dt * (-self.D + self.R + self.A), c_old + self.msh.dt * self.S.value.ravel(), x0=c_old, tol=self.tol_inversion, ) elif self.time_discretization == 'implicit with preconditionning': precond = ( self.inv_matrix_implicit_scheme_t_init # self.jacobi # spsl.LinearOperator((self.msh.x.size*self.msh.y.size,self.msh.x.size*self.msh.y.size),matvec=self.spsolve_lu) ) c, info = gmres( self.Id + self.msh.dt * (-self.D + self.R + self.A), c_old + self.msh.dt * self.S.value.ravel(), x0=c_old, tol=self.tol_inversion, M=precond, ) elif self.time_discretization == 'implicit with matrix inversion': c = self.inv_matrix_implicit_scheme[it, :, :].dot(c_old + self.msh.dt * self.S.value.ravel()) info = 0 elif self.time_discretization == 'implicit with stationnary matrix inversion': RHS = c_old + self.msh.dt * self.S.value.ravel() LHS = (self.Id + self.msh.dt * (-self.D + self.R + self.A)).matvec(c_old) flag_residu = not np.linalg.norm(RHS - LHS) < self.tol_inversion * np.linalg.norm(RHS) if flag_residu: c = np.dot(self.inv_matrix_implicit_scheme, RHS) # c = np.dot(self.inv_matrix_implicit_scheme, c_old + self.msh.dt * self.S.value.ravel()) info = 0 if info > 0: raise ValueError( "The algorithme used to solve the linear system has not converge" + "to the expected tolerance or within the maximum number of iteration." ) if info < 0: raise ValueError("The algorithme used to solve the linear system could not proceed du to illegal input or breakdown.") if it % save_rate == 0: t_save.append(self.msh.t) c_save.append(c.reshape((self.msh.y.shape[0], self.msh.x.shape[0]))) np.save(Path(path_save) / 'c_save.npy', c_save) np.save(Path(path_save) / 't_save.npy', t_save) def solver_est_at_obs_times(self, obs, display_flag=True): """ solve the PDE on the whole time window and store the resulting estimations of the concentration at the observations times - TO DO: * check that the function works - input: * obs: object of the class Obs, contains the observations and the observations times at which we aim to estimate the concentration - do: * solve the PDE on the whole time window but only store the resulting estimations at the observations times in the atrtibutes c_est of obs """ # initialization of the unknown variable at the current time c = np.zeros((self.msh.y.shape[0] * self.msh.x.shape[0],)) if 0 in obs.index_obs_to_index_time_est: c_prov = c.reshape((self.msh.y.shape[0], self.msh.x.shape[0])) for index_obs in obs.index_time_est_to_index_obs[0]: index_x_est = np.argmin(np.abs(self.msh.x - obs.X_obs[index_obs, 0])) index_y_est = np.argmin(np.abs(self.msh.y - obs.X_obs[index_obs, 1])) obs.c_est[index_obs] = c_prov[index_y_est, index_x_est] # loop until the final time or the steady state is reached for it, self.msh.t in enumerate(self.msh.t_array[1:]): if display_flag: sys.stdout.write(f'\rt = {"{:.3f}".format(self.msh.t)} / {"{:.3f}".format(self.msh.T_final)} s') sys.stdout.flush() # update the coefficients of the equation at the current time and self.at_current_time(self.msh.t, i=it) # store the concentration at the previous time step c_old = np.copy(c) # NECESSARY??? # inverse the linear system resulting the semi-implicit time discretization # using a conjugate gradient method for the current time step if self.time_discretization == 'semi-implicit with matrix inversion': c = self.inv_matrix_semi_implicit_scheme.dot(c_old + self.msh.dt * (-self.A.matvec(c_old) + self.S.value.ravel())) info = 0 elif self.time_discretization == 'semi-implicit': c, info = cg( self.Id + self.msh.dt * (-self.D + self.R), c_old + self.msh.dt * (-self.A.matvec(c_old) + self.S.value.ravel()), x0=c_old, tol=self.tol_inversion, ) # inverse the linear system resulting the implicit time discretization using a gmres method for the current time step elif self.time_discretization == 'implicit': c, info = gmres( self.Id + self.msh.dt * (-self.D + self.R + self.A), c_old + self.msh.dt * self.S.value.ravel(), x0=c_old, tol=self.tol_inversion, ) elif self.time_discretization == 'implicit with matrix inversion': c = self.inv_matrix_implicit_scheme[it, :, :].dot(c_old + self.msh.dt * self.S.value.ravel()) info = 0 elif self.time_discretization == 'implicit with stationnary matrix inversion': RHS = c_old + self.msh.dt * self.S.value.ravel() if not np.linalg.norm(RHS, ord=np.inf) < 1e-16: LHS = (self.Id + self.msh.dt * (-self.D + self.R + self.A)).matvec(c_old) flag_residu = not np.linalg.norm(RHS - LHS) < self.tol_inversion * np.linalg.norm(RHS) # print(np.linalg.norm(RHS),np.linalg.norm(RHS, ord=np.inf)) if flag_residu: c = np.dot(self.inv_matrix_implicit_scheme, RHS) else: c = np.zeros_like(c) # print("the norm is 0") # c = np.dot(self.inv_matrix_implicit_scheme, c_old + self.msh.dt * self.S.value.ravel()) info = 0 if info > 0: raise ValueError( "The algorithme used to solve the linear system has not converge" + "to the expected tolerance or within the maximum number of iteration." ) if info < 0: raise ValueError("The algorithme used to solve the linear system could not proceed du to illegal input or breakdown.") it += 1 if it in obs.index_obs_to_index_time_est: c_prov = c.reshape((self.msh.y.shape[0], self.msh.x.shape[0])) for index_obs in obs.index_time_est_to_index_obs[it]: i = np.where(obs.index_obs_to_index_time_est[index_obs, :] == it) index_x_est = np.argmin(np.abs(self.msh.x - obs.X_obs[index_obs, 0])) index_y_est = np.argmin(np.abs(self.msh.y - obs.X_obs[index_obs, 1])) obs.c_est[index_obs, i] = c_prov[index_y_est, index_x_est] def solver_steady_state(self): """ solve the steady-state version of the PDE inverse the linear system (- D + A + R) c = S using a conjugate gradient method - output: * numpy array of shape (msh.y.size, msh.x.size), resulting concentration at the steady state """ c_ravel, tol = cg(-self.D + self.A + self.R, self.S.value.ravel()) return c_ravel.reshape((self.msh.y.shape[0], self.msh.x.shape[0]))