├── .gitignore ├── README.md ├── analysis └── n-ode_cpu-vs-gpu.py ├── data ├── allen_cahn │ ├── data_generation.py │ └── simulator.py ├── burger │ ├── data_generation.py │ └── simulator.py ├── burger_2d │ ├── data_generation.py │ └── simulator.py ├── diffusion_reaction │ ├── data_generation.py │ └── simulator.py └── diffusion_sorption │ ├── data_generation.py │ └── simulator.py ├── environment.yml ├── models ├── aphynity │ ├── aphynity.py │ ├── config.json │ ├── experiment.py │ ├── test.py │ └── train.py ├── cnn_node │ ├── cnn_node.py │ ├── config.json │ ├── experiment.py │ ├── test.py │ └── train.py ├── convlstm │ ├── config.json │ ├── conv_lstm.py │ ├── experiment.py │ ├── test.py │ └── train.py ├── cvpinn │ ├── cvpinn.py │ └── cvpinns.py ├── distana │ ├── config.json │ ├── conv_lstm.py │ ├── distana.py │ ├── experiment.py │ ├── test.py │ └── train.py ├── finn │ ├── config.json │ ├── experiment.py │ ├── finn.py │ ├── test.py │ └── train.py ├── finn_poly │ ├── config.json │ ├── finn.py │ ├── test.py │ └── train.py ├── fno │ ├── config.json │ ├── experiment.py │ ├── fno.py │ ├── test.py │ └── train.py ├── phydnet │ ├── config.json │ ├── constrain_moments.py │ ├── experiment.py │ ├── phydnet.py │ ├── test.py │ └── train.py ├── pinn │ ├── config.json │ ├── experiment.py │ ├── pinn.py │ ├── test.py │ └── train.py ├── polynomial │ ├── allen_cahn.py │ ├── burger.py │ ├── diffusion_reaction.py │ └── diffusion_sorption.py ├── tcn │ ├── config.json │ ├── experiment.py │ ├── tcn.py │ ├── test.py │ └── train.py └── utils │ ├── configuration.py │ ├── datasets.py │ └── helper_functions.py └── sampling ├── checkpoints └── diff-sorp_00 │ ├── cfg.json │ └── diff-sorp_00.pt ├── config.json ├── data_core1.xlsx ├── data_core2.xlsx ├── data_core2_long.xlsx ├── finn_exp.py ├── init.py ├── main.py ├── retardation_phys.txt ├── sampler.py ├── set_module.py ├── train.py └── var_inf.py /.gitignore: -------------------------------------------------------------------------------- 1 | __pycache__/ 2 | *.synctex.gz 3 | *.aux 4 | *.log 5 | *.toc 6 | *.out 7 | *.blg 8 | *.bbl 9 | *.DS_Store 10 | *.npy 11 | *.egg-info 12 | *.png 13 | *.sw? 14 | /bin 15 | /models/*/out/* 16 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | 2 | # FInite volume Neural Network (FINN) 3 | 4 | This repository contains the PyTorch code for models, training, and testing, and Python code for data generation to conduct the experiments as reported in the work [Composing Partial Differential Equations with Physics-Aware Neural Networks](https://arxiv.org/abs/2111.11798) 5 | 6 | If you find this repository helpful, please cite our work: 7 | 8 | ``` 9 | @inproceedings{karlbauer2021composing, 10 | address = {Baltimore, USA}, 11 | author = {Karlbauer, Matthias and Praditia, Timothy and Otte, Sebastian and Oladyshkin, Sergey and Nowak, Wolfgang and Butz, Martin V}, 12 | booktitle = {Proceedings of the 39th International Conference on Machine Learning}, 13 | month = {16--23 Jul}, 14 | series = {Proceedings of Machine Learning Research}, 15 | title = {Composing Partial Differential Equations with Physics-Aware Neural Networks}, 16 | year = {2022} 17 | } 18 | ``` 19 | 20 | ## Dependencies 21 | 22 | We recommend setting up an (e.g. [conda](https://docs.conda.io/projects/conda/en/latest/user-guide/tasks/manage-environments.html)) environment with python 3.7 (i.e. `conda create -n finn python=3.7`). The required packages for data generation and model evaluation are 23 | 24 | - `conda install -c anaconda numpy scipy` 25 | - `conda install -c pytorch pytorch==1.9.0` 26 | - `conda install -c jmcmurray json` 27 | - `conda install -c conda-forge matplotlib torchdiffeq jsmin` 28 | 29 | Alternatively, the `environment.yml` can be used to create an according conda environment with `conda env create -f environment.yml`. 30 | 31 | ## Models & Experiments 32 | 33 | The code of the different pure machine learning models ([Temporal Convolutional Network (TCN)](https://arxiv.org/pdf/2111.07470.pdf), [Convolutional Long-Short Term Memory (ConvLSTM)](https://arxiv.org/pdf/1506.04214.pdf)), physics-motivated models ([DISTANA](https://arxiv.org/pdf/1912.11141.pdf), CNN-NODE, [Fourier Neural Operator (FNO)](https://arxiv.org/pdf/2010.08895.pdf)) and physics-aware models ([Physics-Informed Neural Networks (PINN)](https://www.sciencedirect.com/science/article/pii/S0021999118307125), [PhyDNet](https://arxiv.org/pdf/2003.01460.pdf), [APHYNITY](https://arxiv.org/pdf/2010.04456.pdf), FINN) can be found in the `models` directory. 34 | 35 | Each model directory contains a `config.json` file to specify model parameters, data, etc. Please modify the sections in the respective `config.json` files as detailed below (further information about data and model architectures is reported in the according data sections of the paper's appendices): 36 | 37 | ``` 38 | "training": { 39 | "t_stop": 150 // burger and allen-cahn 150, diff-sorp 400, diff-react 70 40 | }, 41 | 42 | "validation": { 43 | "t_start": 150, // burger, burger_2d, and allen-cahn 150, diff-sorp 400, diff-react 70 44 | "t_stop": 200 // burger, burger2d, and allen-cahn 200, diff-sorp 500, diff-react 100 45 | }, 46 | 47 | "data": { 48 | "type": "burger", // "burger", "burger_2d", diffusion_sorption", "diffusion_reaction", "allen_cahn" 49 | "name": "data_ext", // "data_train", "data_ext", "data_test" 50 | } 51 | 52 | "model": { 53 | "name": "burger" // "burger", "burger_2d", diff-sorp", "diff-react", "allen-cahn" 54 | "field_size": [49], // burger and allen-cahn [49], diff-sorp [26], burger_2d and diff-react [49, 49] 55 | ... other settings to be specified according to the model architectures section in the paper's appendix 56 | } 57 | ``` 58 | 59 | 60 | The actual models can be trained and tested by calling the according `python train.py` or `python test.py` scripts. Alternatively, `python experiment.py` can be used to either train or test n models (please consider the settings in the `experiment.py` script). 61 | 62 | ## Data generation 63 | 64 | The Python scripts to generate the burger, burger_2d, diffusion-sorption, diffusion-reaction, and allen-cahn data can be found in the `data` directory. 65 | 66 | In each of the `burger`, `burger_2d`, `diffusion_sorption`, `diffusion_reaction`, and `allen-cahn` directories, a `data_generation.py` and `simulator.py` script can be found. The former is used to generate train, extrapolation (ext), or test data. For details about the according data generation settings of each dataset, please refer to the corresponding data sections in the paper's appendix. 67 | 68 | ## Uncertainty quantification 69 | 70 | The scripts required for the uncertainty quantification part can be found in the `sampling` directory. 71 | 72 | Here, we provide the experimental data in the `data_core1.xlsx`, `data_core2.xlsx`, and `data_core2_long.xlsx` files. Additionally, the retardation factor values of the fitted physical model are provided in `retardation_phys.txt`. 73 | 74 | There are two options to perform uncertainty quantification in this work, namely [Variational Inference](https://arxiv.org/pdf/1505.05424.pdf) and Markov Chain Monte Carlo (MCMC). 75 | 76 | Training and sampling using the Variational Inference method is provided in the `var_inf.py` script, and the collections of MCMC methods are provided in the `sampler.py` script, including the [Metropolis-Hastings](https://www.jstor.org/stable/2684568?seq=1), [Metropolis-Adjusted Langevin Algorithm (MALA)](https://projecteuclid.org/journals/bernoulli/volume-2/issue-4/Exponential-convergence-of-Langevin-distributions-and-their-discrete-approximations/bj/1178291835.full), and [Barker proposal](https://arxiv.org/pdf/1908.11812.pdf). 77 | 78 | Sampling with MCMC can be performed by running the `main.py` script, and selecting the desired sampler in the `config.json` file. In the config file, user can also choose whether to start with a random initial point or using a pre-trained model provided in the `checkpoints` directory (by setting `sampling.random_init` true or false). 79 | 80 | # Citations 81 | 82 | ``` 83 | 84 | @article{espeholt2021skillful, 85 | title={Skillful Twelve Hour Precipitation Forecasts using Large Context Neural Networks}, 86 | author={Espeholt, Lasse and Agrawal, Shreya and S{\o}nderby, Casper and Kumar, Manoj and Heek, Jonathan and Bromberg, Carla and Gazen, Cenk and Hickey, Jason and Bell, Aaron and Kalchbrenner, Nal}, 87 | journal={arXiv preprint arXiv:2111.07470}, 88 | year={2021} 89 | } 90 | 91 | @article{shi2015convolutional, 92 | title={Convolutional LSTM network: A machine learning approach for precipitation nowcasting}, 93 | author={Shi, X. and Chen, Z. and Wang, H. and Yeung, D.Y. and Wong, W.K. and Woo, W.C.}, 94 | journal={arXiv preprint arXiv:1506.04214}, 95 | year={2015} 96 | } 97 | 98 | @article{Karlbauer2019, 99 | title={A Distributed Neural Network Architecture for Robust Non-Linear Spatio-Temporal Prediction}, 100 | author={Karlbauer, M. and Otte, S. and Lensch, H.P.A. and Scholten, T. and Wulfmeyer, V. and Butz, M.V.}, 101 | year={2019}, 102 | journal={arXiv preprint arXiv:1912.11141}, 103 | } 104 | 105 | @article{li2020fourier, 106 | title={Fourier neural operator for parametric partial differential equations}, 107 | author={Li, Z. and Kovachki, N. and Azizzadenesheli, K. and Liu, B. and Bhattacharya, K. and Stuart, A. and Anandkumar, A.}, 108 | journal={arXiv preprint arXiv:2010.08895}, 109 | year={2020} 110 | } 111 | 112 | @article{Raissi2019, 113 | author={Raissi, M. and Perdikaris, P. and Karniadakis, G.E.}, 114 | journal={Journal of Computational Physics}, 115 | title={Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations}, 116 | year={2019}, 117 | volume={378}, 118 | pages={686-707}, 119 | doi={10.1016/j.jcp.2018.10.045} 120 | } 121 | 122 | @inproceedings{guen2020disentangling, 123 | title={Disentangling physical dynamics from unknown factors for unsupervised video prediction}, 124 | author={Guen, V.L. and Thome, N.}, 125 | booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition}, 126 | pages={11474--11484}, 127 | year={2020} 128 | } 129 | 130 | @article{guen2020augmenting, 131 | title={Augmenting physical models with deep networks for complex dynamics forecasting}, 132 | author={Yin, Y. and Guen, V.L. and Dona, J. and Ayed, I. and de B{\'e}zenac, E. and Thome, N. and Gallinari, P.}, 133 | journal={arXiv preprint arXiv:2010.04456}, 134 | year={2020} 135 | } 136 | 137 | @misc{https://doi.org/10.48550/arxiv.1505.05424, 138 | doi = {10.48550/ARXIV.1505.05424}, 139 | url = {https://arxiv.org/abs/1505.05424}, 140 | author = {Blundell, Charles and Cornebise, Julien and Kavukcuoglu, Koray and Wierstra, Daan}, 141 | title = {Weight Uncertainty in Neural Networks}, 142 | year = {2015} 143 | } 144 | 145 | @article{Chib1995MH, 146 | url = {http://www.jstor.org/stable/2684568}, 147 | author = {Siddhartha Chib and Edward Greenberg}, 148 | journal = {The American Statistician}, 149 | number = {4}, 150 | pages = {327--335}, 151 | publisher = {[American Statistical Association, Taylor & Francis, Ltd.]}, 152 | title = {Understanding the Metropolis-Hastings Algorithm}, 153 | volume = {49}, 154 | year = {1995} 155 | } 156 | 157 | @article{Roberts1996MALA, 158 | author = {Gareth O. Roberts and Richard L. Tweedie}, 159 | title = {{Exponential convergence of Langevin distributions and their discrete approximations}}, 160 | volume = {2}, 161 | journal = {Bernoulli}, 162 | number = {4}, 163 | publisher = {Bernoulli Society for Mathematical Statistics and Probability}, 164 | pages = {341 -- 363}, 165 | year = {1996}, 166 | doi = {bj/1178291835} 167 | } 168 | 169 | @misc{Livingstone2019Barker, 170 | doi = {10.48550/ARXIV.1908.11812}, 171 | url = {https://arxiv.org/abs/1908.11812}, 172 | author = {Livingstone, Samuel and Zanella, Giacomo}, 173 | title = {The Barker proposal: combining robustness and efficiency in gradient-based MCMC}, 174 | year = {2019} 175 | } 176 | 177 | 178 | ``` 179 | 180 | # Code contributors 181 | * [Matthias Karlbauer](https://github.com/MatKbauer) 182 | * [Timothy Praditia](https://github.com/timothypraditia) -------------------------------------------------------------------------------- /analysis/n-ode_cpu-vs-gpu.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | """ 3 | This script contrasts the runtime of training a simple feed forward neural 4 | network under incorporation of the Neural ODE method from 5 | https://github.com/rtqichen/torchdiffeq 6 | """ 7 | 8 | import torch.nn as nn 9 | import torch as th 10 | from torchdiffeq import odeint 11 | import time 12 | import matplotlib.pyplot as plt 13 | 14 | class Model(nn.Module): 15 | def __init__(self): 16 | super(Model, self).__init__() 17 | self.net = nn.Sequential( 18 | nn.Linear(1, 50), 19 | nn.Tanh(), 20 | nn.Linear(50, 1) 21 | ) 22 | 23 | def forward(self, t, y): 24 | return self.net(y) 25 | 26 | 27 | device_cpu = th.device("cpu") 28 | device_cuda = th.device("cuda") 29 | 30 | t_cpu = th.linspace(0,1,1001).to(device=device_cpu) 31 | t_cuda = th.linspace(0,1,1001).to(device=device_cuda) 32 | 33 | model_cpu = Model().to(device=device_cpu) 34 | model_cuda = Model().to(device=device_cuda) 35 | 36 | runtime_cpu = th.empty(240) 37 | runtime_cuda = th.empty(240) 38 | num_batches = th.empty(240) 39 | 40 | with th.no_grad(): 41 | for i in range(10,250): 42 | num_batches[i-10] = i**2 43 | inp_cpu = th.zeros(i**2,1).to(device=device_cpu) 44 | inp_cuda = th.zeros(i**2,1).to(device=device_cuda) 45 | start_time_cpu = time.time() 46 | out_cpu = odeint(model_cpu,inp_cpu,t_cpu) 47 | runtime_cpu[i-10] = time.time() - start_time_cpu 48 | start_time_cuda = time.time() 49 | out_cuda = odeint(model_cuda,inp_cuda,t_cuda) 50 | runtime_cuda[i-10] = time.time() - start_time_cuda 51 | 52 | print(i-10, runtime_cpu[i-10], runtime_cuda[i-10]) 53 | 54 | plt.figure() 55 | plt.plot(num_batches, runtime_cpu, label="CPU") 56 | plt.plot(num_batches, runtime_cuda, label="GPU") 57 | plt.xlabel("Number of batches") 58 | plt.ylabel("Runtime [s]") 59 | plt.title("Neural ODE runtime comparison") 60 | plt.legend() 61 | plt.tight_layout() 62 | plt.show() 63 | -------------------------------------------------------------------------------- /data/allen_cahn/simulator.py: -------------------------------------------------------------------------------- 1 | """ 2 | This script provides a class solving the Burger's equation via numerical 3 | integration using scipy's solve_ivp method. It can be used to generate data 4 | samples of the Burger's equation with Dirichlet boundary condition on both 5 | sides (u = 0). 6 | """ 7 | 8 | import numpy as np 9 | from scipy.integrate import solve_ivp 10 | 11 | 12 | class Simulator: 13 | 14 | def __init__(self, diffusion_coefficient, t_max, t_steps, x_left, x_right, 15 | x_steps, train_data): 16 | """ 17 | Constructor method initializing the parameters for the Burger's 18 | equation. 19 | :param diffusion_coefficient: The diffusion coefficient 20 | :param t_max: Stop time of the simulation 21 | :param t_steps: Number of simulation steps 22 | :param x_left: Left end of the 1D simulation field 23 | :param x_right: Right end of the 1D simulation field 24 | :param x_steps: Number of spatial steps between x_left and x_right 25 | """ 26 | 27 | # Set class parameters 28 | self.D = diffusion_coefficient 29 | 30 | self.T = t_max 31 | self.X0 = x_left 32 | self.X1 = x_right 33 | 34 | self.Nx = x_steps 35 | self.Nt = t_steps 36 | 37 | self.dx = (self.X1 - self.X0)/(self.Nx - 1) 38 | 39 | self.x = np.linspace(self.X0 + self.dx, self.X1 - self.dx, self.Nx - 2) 40 | self.t = np.linspace(0, self.T, self.Nt) 41 | 42 | self.train_data = train_data 43 | 44 | def generate_sample(self): 45 | """ 46 | Single sample generation using the parameters of this simulator. 47 | :return: The generated sample as numpy array(t, x) 48 | """ 49 | 50 | # Initialize the simulation field 51 | if self.train_data: 52 | u0 = self.x**2 * np.cos(np.pi*self.x) 53 | else: 54 | u0 = np.sin(np.pi*self.x/2) 55 | 56 | nx_minus_2 = np.diag(-2*np.ones(self.Nx-2), k=0) 57 | nx_minus_3 = np.diag(np.ones(self.Nx-3), k=-1) 58 | nx_plus_3 = np.diag(np.ones(self.Nx-3), k=1) 59 | 60 | self.lap = nx_minus_2 + nx_minus_3 + nx_plus_3 61 | self.lap /= self.dx**2 62 | # Periodic BC 63 | self.lap[0,-1] = 1/self.dx**2 64 | self.lap[-1,0] = 1/self.dx**2 65 | 66 | # Solve Burger's equation 67 | prob = solve_ivp(self.rc_ode, (0, self.T), u0, t_eval=self.t) 68 | ode_data = prob.y 69 | 70 | self.sample = np.transpose(ode_data) 71 | 72 | return self.sample 73 | 74 | def rc_ode(self, t, u): 75 | """ 76 | Solves a given equation for a particular time step. 77 | :param t: The current time step 78 | :param u: The equation values to solve 79 | :return: A finite difference solution 80 | """ 81 | # Return finite difference 82 | return self.D*np.matmul(self.lap, u) - 5*(u**3-u) 83 | -------------------------------------------------------------------------------- /data/burger/simulator.py: -------------------------------------------------------------------------------- 1 | """ 2 | This script provides a class solving the Burger's equation via numerical 3 | integration using scipy's solve_ivp method. It can be used to generate data 4 | samples of the Burger's equation with Dirichlet boundary condition on both 5 | sides (u = 0). 6 | """ 7 | 8 | import numpy as np 9 | from scipy.integrate import solve_ivp 10 | 11 | 12 | class Simulator: 13 | 14 | def __init__(self, diffusion_coefficient, t_max, t_steps, x_left, x_right, 15 | x_steps, train_data): 16 | """ 17 | Constructor method initializing the parameters for the Burger's 18 | equation. 19 | :param diffusion_coefficient: The diffusion coefficient 20 | :param t_max: Stop time of the simulation 21 | :param t_steps: Number of simulation steps 22 | :param x_left: Left end of the 1D simulation field 23 | :param x_right: Right end of the 1D simulation field 24 | :param x_steps: Number of spatial steps between x_left and x_right 25 | """ 26 | 27 | # Set class parameters 28 | self.D = diffusion_coefficient 29 | 30 | self.T = t_max 31 | self.X0 = x_left 32 | self.X1 = x_right 33 | 34 | self.Nx = x_steps 35 | self.Nt = t_steps 36 | 37 | self.dx = (self.X1 - self.X0)/(self.Nx - 1) 38 | 39 | self.x = np.linspace(self.X0 + self.dx, self.X1 - self.dx, self.Nx - 2) 40 | self.t = np.linspace(0, self.T, self.Nt) 41 | 42 | self.train_data = train_data 43 | 44 | def generate_sample(self): 45 | """ 46 | Single sample generation using the parameters of this simulator. 47 | :return: The generated sample as numpy array(t, x) 48 | """ 49 | 50 | # Initialize the simulation field 51 | if self.train_data: 52 | u0 = -np.sin(np.pi*self.x) 53 | else: 54 | u0 = np.sin(np.pi*self.x) 55 | 56 | nx_minus_2 = np.diag(-2*np.ones(self.Nx-2), k=0) 57 | nx_minus_3 = np.diag(np.ones(self.Nx-3), k=-1) 58 | nx_plus_3 = np.diag(np.ones(self.Nx-3), k=1) 59 | 60 | self.lap = nx_minus_2 + nx_minus_3 + nx_plus_3 61 | self.lap /= self.dx**2 62 | # Periodic BC 63 | # lap[0,-1] = 1/dx**2 64 | # lap[-1,0] = 1/dx**2 65 | 66 | # Solve Burger's equation 67 | prob = solve_ivp(self.rc_ode, (0, self.T), u0, t_eval=self.t) 68 | ode_data = prob.y 69 | 70 | self.sample = np.transpose(ode_data) 71 | 72 | return self.sample 73 | 74 | def rc_ode(self, t, u): 75 | """ 76 | Solves a given equation for a particular time step. 77 | :param t: The current time step 78 | :param u: The equation values to solve 79 | :return: A finite difference solution 80 | """ 81 | a_plus = np.maximum(u,0) 82 | a_min = np.minimum(u,0) 83 | 84 | u_left = np.concatenate((np.array([0]),u[:-1])) 85 | u_right = np.concatenate((u[1:],np.array([0]))) 86 | 87 | # Return finite difference 88 | return self.D*np.matmul(self.lap, u) - (a_plus*(u - u_left)/self.dx + 89 | a_min*(u_right - u)/self.dx) 90 | -------------------------------------------------------------------------------- /data/burger_2d/simulator.py: -------------------------------------------------------------------------------- 1 | """ 2 | This script provides a class solving the Burger's equation via numerical 3 | integration using scipy's solve_ivp method. It can be used to generate data 4 | samples of the Burger's equation with Dirichlet boundary condition on both 5 | sides (u = 0). 6 | """ 7 | 8 | import numpy as np 9 | from scipy.integrate import solve_ivp 10 | from scipy.sparse import diags 11 | 12 | 13 | class Simulator: 14 | 15 | def __init__(self, diffusion_coefficient, t_max, t_steps, x_left, x_right, 16 | x_steps, y_bottom, y_top, y_steps, train_data): 17 | """ 18 | Constructor method initializing the parameters for the Burger's 19 | equation. 20 | :param diffusion_coefficient: The diffusion coefficient 21 | :param t_max: Stop time of the simulation 22 | :param t_steps: Number of simulation steps 23 | :param x_left: Left end of the 1D simulation field 24 | :param x_right: Right end of the 1D simulation field 25 | :param x_steps: Number of spatial steps between x_left and x_right 26 | """ 27 | 28 | # Set class parameters 29 | self.D = diffusion_coefficient 30 | 31 | self.T = t_max 32 | self.X0 = x_left 33 | self.X1 = x_right 34 | self.Y0 = y_bottom 35 | self.Y1 = y_top 36 | 37 | self.Nx = x_steps 38 | self.Ny = y_steps 39 | self.Nt = t_steps 40 | 41 | self.dx = (self.X1 - self.X0)/(self.Nx - 1) 42 | self.dy = (self.Y1 - self.Y0)/(self.Ny - 1) 43 | 44 | self.x = np.linspace(self.X0 + self.dx, self.X1 - self.dx, self.Nx) 45 | self.y = np.linspace(self.Y0 + self.dy, self.Y1 - self.dy, self.Ny) 46 | self.t = np.linspace(0, self.T, self.Nt) 47 | 48 | self.train_data = train_data 49 | 50 | def generate_sample(self): 51 | """ 52 | Single sample generation using the parameters of this simulator. 53 | :return: The generated sample as numpy array(t, x) 54 | """ 55 | 56 | # Initialize the simulation field 57 | X, Y = np.meshgrid(self.x, self.y) 58 | if self.train_data: 59 | u0 = -np.sin(np.pi*(X+Y)) 60 | else: 61 | u0 = -np.sin(np.pi*(X-Y)) 62 | u0 = u0.reshape(self.Nx*self.Ny) 63 | 64 | # Generate arrays as diagonal inputs to the Laplacian matrix 65 | main_diag = -2*np.ones(self.Nx)/self.dx**2 -2*np.ones(self.Nx)/self.dy**2 66 | main_diag = np.tile(main_diag, self.Ny) 67 | 68 | left_diag = np.ones(self.Nx) 69 | left_diag[0] = 0 70 | left_diag = np.tile(left_diag, self.Ny) 71 | left_diag = left_diag[1:]/self.dx**2 72 | 73 | right_diag = np.ones(self.Nx) 74 | right_diag[-1] = 0 75 | right_diag = np.tile(right_diag, self.Ny) 76 | right_diag = right_diag[:-1]/self.dx**2 77 | 78 | bottom_diag = np.ones(self.Nx*(self.Ny-1))/self.dy**2 79 | 80 | top_diag = np.ones(self.Nx*(self.Ny-1))/self.dy**2 81 | 82 | # Generate the sparse Laplacian matrix 83 | diagonals = [main_diag, left_diag, right_diag, bottom_diag, top_diag] 84 | offsets = [0, -1, 1, -self.Nx, self.Nx] 85 | self.lap = diags(diagonals, offsets) 86 | 87 | # Solve Burger's equation 88 | prob = solve_ivp(self.rc_ode, (0, self.T), u0, t_eval=self.t) 89 | ode_data = prob.y 90 | 91 | self.sample = np.transpose(ode_data).reshape(-1,self.Ny,self.Nx) 92 | 93 | return self.sample 94 | 95 | def rc_ode(self, t, u): 96 | """ 97 | Solves a given equation for a particular time step. 98 | :param t: The current time step 99 | :param u: The equation values to solve 100 | :return: A finite difference solution 101 | """ 102 | 103 | u = u.reshape(self.Ny,self.Nx) 104 | 105 | a_plus = np.maximum(u,0) 106 | a_min = np.minimum(u,0) 107 | 108 | u_left = np.concatenate((np.zeros((u.shape[0],1)),u[:,:-1]),axis=1) 109 | u_right = np.concatenate((u[:,1:],np.zeros((u.shape[0],1))),axis=1) 110 | 111 | u_bottom = np.concatenate((np.zeros((1,u.shape[1])),u[:-1]),axis=0) 112 | u_top = np.concatenate((u[1:],np.zeros((1,u.shape[1]))),axis=0) 113 | 114 | adv = a_plus*(u-u_left)/self.dx + a_min*(u_right-u)/self.dx + \ 115 | a_plus*(u-u_bottom)/self.dy + a_min*(u_top-u)/self.dy 116 | u = u.reshape(self.Nx*self.Ny) 117 | 118 | return self.D*(self.lap@u) - adv.reshape(self.Nx*self.Ny) 119 | -------------------------------------------------------------------------------- /data/diffusion_reaction/simulator.py: -------------------------------------------------------------------------------- 1 | """ 2 | This script provides a class solving a diffusion reaction problem via numerical 3 | integration using scipy's solve_ivp method. It can be used to generate data 4 | samples of the diffusion reaction problem with Fitzhugh-Nagumo equation. 5 | """ 6 | 7 | import numpy as np 8 | from scipy.integrate import solve_ivp 9 | 10 | 11 | class Simulator: 12 | 13 | def __init__(self, diffusion_coefficient_u, diffusion_coefficient_v, k, 14 | t_max, t_steps, x_left, x_right, x_steps, y_bottom, y_top, 15 | y_steps, train_data): 16 | """ 17 | Constructor method initializing the parameters for the diffusion 18 | sorption problem. 19 | :param diffusion_coefficient_u: The diffusion coefficient of u 20 | :param diffusion_coefficient_u: The diffusion coefficient of v 21 | :param k: The reaction parameter 22 | :param t_max: Stop time of the simulation 23 | :param t_steps: Number of simulation steps 24 | :param x_left: Left end of the 2D simulation field 25 | :param x_right: Right end of the 2D simulation field 26 | :param x_steps: Number of spatial steps between x_left and x_right 27 | :param y_bottom: bottom end of the 2D simulation field 28 | :param y_top: top end of the 2D simulation field 29 | :param y_steps: Number of spatial steps between y_bottom and y_top 30 | """ 31 | 32 | # Set class parameters 33 | self.Du = diffusion_coefficient_u 34 | self.Dv = diffusion_coefficient_v 35 | self.k = k 36 | 37 | self.T = t_max 38 | self.X0 = x_left 39 | self.X1 = x_right 40 | self.Y0 = y_bottom 41 | self.Y1 = y_top 42 | 43 | self.Nx = x_steps 44 | self.Ny = y_steps 45 | self.Nt = t_steps 46 | 47 | self.dx = (self.X1 - self.X0)/(self.Nx - 1) 48 | self.dy = (self.Y1 - self.Y0)/(self.Ny - 1) 49 | 50 | self.x = np.linspace(self.X0 + self.dx, self.X1 - self.dx, self.Nx) 51 | self.y = np.linspace(self.Y0 + self.dy, self.Y1 - self.dy, self.Ny) 52 | self.t = np.linspace(0, self.T, self.Nt) 53 | 54 | self.train_data = train_data 55 | 56 | 57 | def generate_sample(self): 58 | """ 59 | Single sample generation using the parameters of this simulator. 60 | :return: The generated sample as numpy array(t, x, y) 61 | """ 62 | 63 | # Initialize the simulation field 64 | X, Y = np.meshgrid(self.x, self.y) 65 | if self.train_data: 66 | u0 = np.sin(np.pi*(X+1)/2) * np.sin(np.pi*(Y+1)/2) 67 | else: 68 | u0 = np.sin(np.pi*(X+1)/2) * np.sin(np.pi*(Y+1)/2) - 0.5 69 | u0 = u0.reshape(self.Nx*self.Ny) 70 | u0 = np.concatenate((u0,u0)) 71 | 72 | # 73 | # Laplacian matrix 74 | main_diag = -2*np.ones(self.Nx)/self.dx**2 -2*np.ones(self.Nx)/self.dy**2 75 | main_diag[0] = -1/self.dx**2 -2/self.dy**2 76 | main_diag[-1] = -1/self.dx**2 -2/self.dy**2 77 | main_diag = np.tile(main_diag, self.Ny) 78 | main_diag[:self.Nx] = -2/self.dx**2 -1/self.dy**2 79 | main_diag[self.Nx*(self.Ny-1):] = -2/self.dx**2 -1/self.dy**2 80 | main_diag[0] = -1/self.dx**2 -1/self.dy**2 81 | main_diag[self.Nx-1] = -1/self.dx**2 -1/self.dy**2 82 | main_diag[self.Nx*(self.Ny-1)] = -1/self.dx**2 -1/self.dy**2 83 | main_diag[-1] = -1/self.dx**2 -1/self.dy**2 84 | 85 | left_diag = np.ones(self.Nx) 86 | left_diag[0] = 0 87 | left_diag = np.tile(left_diag, self.Ny) 88 | left_diag = left_diag[1:]/self.dx**2 89 | 90 | right_diag = np.ones(self.Nx) 91 | right_diag[-1] = 0 92 | right_diag = np.tile(right_diag, self.Ny) 93 | right_diag = right_diag[:-1]/self.dx**2 94 | 95 | bottom_diag = np.ones(self.Nx*(self.Ny-1))/self.dy**2 96 | 97 | top_diag = np.ones(self.Nx*(self.Ny-1))/self.dy**2 98 | 99 | self.lap = np.diag(main_diag, k=0) +\ 100 | np.diag(left_diag, k=-1) +\ 101 | np.diag(right_diag, k=1) +\ 102 | np.diag(bottom_diag, k=-self.Nx) +\ 103 | np.diag(top_diag, k=self.Nx) 104 | 105 | 106 | # self.q_u = np.zeros(self.Nx*self.Ny) 107 | # self.q_v = np.zeros(self.Nx*self.Ny) 108 | 109 | # Solve the diffusion sorption problem 110 | prob = solve_ivp(self.rc_ode, (0, self.T), u0, t_eval=self.t) 111 | ode_data = prob.y 112 | 113 | sample_u = np.transpose(ode_data[:self.Nx*self.Ny]).reshape(-1,self.Ny,self.Nx) 114 | sample_v = np.transpose(ode_data[self.Nx*self.Ny:]).reshape(-1,self.Ny,self.Nx) 115 | 116 | return sample_u, sample_v 117 | 118 | def rc_ode(self, t, y): 119 | """ 120 | Solves a given equation for a particular time step. 121 | :param t: The current time step 122 | :param u: The equation values to solve 123 | :return: A finite difference solution 124 | """ 125 | 126 | # Separate y into u and v 127 | u = y[:self.Nx*self.Ny] 128 | v = y[self.Nx*self.Ny:] 129 | 130 | # Calculate reaction function for each unknown 131 | react_u = u - u**3 - self.k - v 132 | react_v = u - v 133 | 134 | # Calculate time derivative for each unknown 135 | u_t = react_u + self.Du*np.matmul(self.lap,u) 136 | v_t = react_v + self.Dv*np.matmul(self.lap,v) 137 | 138 | # Stack the time derivative into a single array u_t 139 | u_t = np.concatenate((u_t,v_t)) 140 | 141 | return u_t 142 | -------------------------------------------------------------------------------- /data/diffusion_sorption/simulator.py: -------------------------------------------------------------------------------- 1 | """ 2 | This script provides a class solving a diffusion sorption problem via numerical 3 | integration using scipy's solve_ivp method. It can be used to generate data 4 | samples of the diffusion sorption problem with different boundary conditions 5 | (left Dirichlet, that is constant, right Cauchy). 6 | """ 7 | 8 | import numpy as np 9 | from scipy.integrate import solve_ivp 10 | 11 | 12 | class Simulator: 13 | 14 | def __init__(self, ret_factor_fun, diffusion_coefficient, porosity, rho_s, 15 | k_f, n_f, s_max, kl, kd, solubility, t_max, t_steps, x_left, 16 | x_right, x_steps, train_data): 17 | """ 18 | Constructor method initializing the parameters for the diffusion 19 | sorption problem. 20 | :param ret_factor_fun: The retardation factor function as string 21 | :param diffusion_coefficient: The diffusion coefficient [m^2/day] 22 | :param porosity: The porosity of the medium [-] 23 | :param rho_s: Dry bulk density [kg/m^3] 24 | :param k_f: Freundlich K 25 | :param n_f: Freundlich exponent 26 | :param s_max: Sorption capacity [m^3/kg] 27 | :param kl: Half-concentration [kg/m^3] 28 | :param kd: Partitioning coefficient [m^3/kg] 29 | :param solubility: The solubility of the quantity 30 | :param t_max: Stop time of the simulation 31 | :param t_steps: Number of simulation steps 32 | :param x_left: Left end of the 1D simulation field 33 | :param x_right: Right end of the 1D simulation field 34 | :param x_steps: Number of spatial steps between x_left and x_right 35 | """ 36 | 37 | # Set class parameters 38 | self.D = diffusion_coefficient 39 | 40 | self.por = porosity 41 | self.rho_s = rho_s 42 | self.k_f = k_f 43 | self.n_f = n_f 44 | self.s_max = s_max 45 | self.kl = kl 46 | self.kd = kd 47 | self.solubility = solubility 48 | 49 | self.T = t_max 50 | self.X0 = x_left 51 | self.X1 = x_right 52 | 53 | self.Nx = x_steps 54 | self.Nt = t_steps 55 | 56 | self.dx = (self.X1 - self.X0)/(self.Nx - 1) 57 | 58 | self.x = np.linspace(0, self.X1, self.Nx) 59 | self.t = np.linspace(0, self.T, self.Nt) 60 | 61 | self.train_data = train_data 62 | 63 | # Specify the retardation function according to the user input 64 | if ret_factor_fun == "linear": 65 | self.retardation = self.retardation_linear 66 | elif ret_factor_fun == "langmuir": 67 | self.retardation = self.retardation_langmuir 68 | elif ret_factor_fun == "freundlich": 69 | self.retardation = self.retardation_freundlich 70 | 71 | def retardation_linear(self, u): 72 | """ 73 | Linear retardation factor function. 74 | :param u: The simulation field 75 | """ 76 | return 1 + ((1 - self.por)/self.por)*self.rho_s\ 77 | *self.kd 78 | 79 | def retardation_freundlich(self, u): 80 | """ 81 | Langmuir retardation factor function. 82 | :param u: The simulation field 83 | """ 84 | return 1 + ((1 - self.por)/self.por)*self.rho_s\ 85 | *self.k_f*self.n_f*(u + 1e-6)**(self.n_f-1) 86 | 87 | def retardation_langmuir(self, u): 88 | """ 89 | Freundlich retardation factor function. 90 | :param u: The simulation field 91 | """ 92 | return 1 + ((1 - self.por)/self.por)*self.rho_s\ 93 | *((self.s_max*self.kl)/(u + self.kl)**2) 94 | 95 | def generate_sample(self): 96 | """ 97 | Single sample generation using the parameters of this simulator. 98 | :return: The generated sample as numpy array(t, x) 99 | """ 100 | 101 | # Initialize the simulation field 102 | u0 = np.zeros(self.Nx) 103 | u0 = np.concatenate((u0, u0)) 104 | 105 | # 106 | # Laplacian matrix 107 | nx = np.diag(-2*np.ones(self.Nx), k=0) 108 | nx_minus_1 = np.diag(np.ones(self.Nx-1), k=-1) 109 | nx_plus_1 = np.diag(np.ones(self.Nx-1), k=1) 110 | 111 | self.lap = nx + nx_minus_1 + nx_plus_1 112 | self.lap /= self.dx**2 113 | 114 | self.q = np.zeros(self.Nx) 115 | self.q_tot = np.zeros(self.Nx) 116 | 117 | # Solve the diffusion sorption problem 118 | prob = solve_ivp(self.rc_ode, (0, self.T), u0, t_eval=self.t, method="BDF") 119 | ode_data = prob.y 120 | 121 | sample_c = np.transpose(ode_data[:self.Nx]) 122 | sample_c_tot = np.transpose(ode_data[self.Nx:]) 123 | 124 | return sample_c, sample_c_tot 125 | 126 | def rc_ode(self, t, u): 127 | """ 128 | Solves a given equation for a particular time step. 129 | :param t: The current time step 130 | :param u: The equation values to solve 131 | :return: A finite difference solution 132 | """ 133 | 134 | # Separate u into c and c_tot 135 | c = u[:self.Nx] 136 | c_tot = u[self.Nx:] 137 | 138 | # Calculate left and right BC 139 | left_BC = self.solubility 140 | right_BC = (c[-2]-c[-1])/self.dx * self.D 141 | 142 | # Calculate time derivative of each unknown 143 | self.q[0] = self.D/self.retardation(c[0])/(self.dx**2)*left_BC 144 | self.q[-1] = self.D/self.retardation(c[-1])/(self.dx**2)*right_BC 145 | c_t = self.D/self.retardation(c)*np.matmul(self.lap, c) + self.q 146 | 147 | self.q_tot[0] = self.D*self.por/(self.rho_s/1000)\ 148 | /(self.dx**2)*left_BC 149 | self.q_tot[-1] = self.D*self.por/(self.rho_s/1000)\ 150 | /(self.dx**2)*right_BC 151 | c_tot_t = self.D*self.por/(self.rho_s/1000)\ 152 | *np.matmul(self.lap, c) + self.q_tot 153 | 154 | # Stack the time derivative into a single array u_t 155 | u_t = np.concatenate((c_t, c_tot_t)) 156 | 157 | return u_t 158 | -------------------------------------------------------------------------------- /environment.yml: -------------------------------------------------------------------------------- 1 | name: finn 2 | channels: 3 | - pytorch 4 | - jmcmurray 5 | - anaconda 6 | - bioconda 7 | - r 8 | - defaults 9 | - conda-forge 10 | dependencies: 11 | - _libgcc_mutex=0.1=main 12 | - _openmp_mutex=4.5=1_gnu 13 | - blas=1.0=mkl 14 | - brotli=1.0.9=he6710b0_2 15 | - ca-certificates=2021.5.30=ha878542_0 16 | - certifi=2021.5.30=py37h89c1867_0 17 | - cudatoolkit=10.2.89=hfd86e86_1 18 | - cycler=0.10.0=py37_0 19 | - dbus=1.13.18=hb2f20db_0 20 | - expat=2.4.1=h2531618_2 21 | - fontconfig=2.13.1=h6c09931_0 22 | - fonttools=4.25.0=pyhd3eb1b0_0 23 | - freetype=2.10.4=h5ab3b9f_0 24 | - glib=2.69.1=h5202010_0 25 | - gst-plugins-base=1.14.0=h8213a91_2 26 | - gstreamer=1.14.0=h28cd5cc_2 27 | - icu=58.2=he6710b0_3 28 | - intel-openmp=2020.2=254 29 | - jpeg=9d=h7f8727e_0 30 | - jsmin=2.2.2=py37hc8dfbb8_1002 31 | - json=0.1.1=0 32 | - kiwisolver=1.3.1=py37h2531618_0 33 | - lcms2=2.12=h3be6417_0 34 | - ld_impl_linux-64=2.35.1=h7274673_9 35 | - libffi=3.3=he6710b0_2 36 | - libgcc-ng=9.3.0=h5101ec6_17 37 | - libgfortran-ng=7.3.0=hdf63c60_0 38 | - libgomp=9.3.0=h5101ec6_17 39 | - libpng=1.6.37=hbc83047_0 40 | - libstdcxx-ng=9.3.0=hd4cf53a_17 41 | - libtiff=4.2.0=h85742a9_0 42 | - libuuid=1.0.3=h1bed415_2 43 | - libuv=1.40.0=h7b6447c_0 44 | - libwebp-base=1.2.0=h27cfd23_0 45 | - libxcb=1.14=h7b6447c_0 46 | - libxml2=2.9.12=h03d6c58_0 47 | - lz4-c=1.9.3=h295c915_1 48 | - matplotlib=3.4.2=py37h89c1867_0 49 | - matplotlib-base=3.4.2=py37hab158f2_0 50 | - mkl=2019.4=243 51 | - mkl-service=2.3.0=py37he904b0f_0 52 | - mkl_fft=1.2.0=py37h23d657b_0 53 | - mkl_random=1.1.0=py37hd6b4f25_0 54 | - munkres=1.0.7=py_1 55 | - ncurses=6.2=he6710b0_1 56 | - ninja=1.10.2=hff7bd54_1 57 | - numpy=1.19.1=py37hbc911f0_0 58 | - numpy-base=1.19.1=py37hfa32c7d_0 59 | - olefile=0.46=py37_0 60 | - openjpeg=2.4.0=h3ad879b_0 61 | - openssl=1.1.1k=h7f98852_0 62 | - pcre=8.45=h295c915_0 63 | - pillow=8.3.1=py37h2c7a002_0 64 | - pip=21.2.2=py37h06a4308_0 65 | - pyparsing=2.4.7=pyhd3eb1b0_0 66 | - pyqt=5.9.2=py37h05f1152_2 67 | - python=3.7.11=h12debd9_0 68 | - python-dateutil=2.8.2=pyhd3eb1b0_0 69 | - python_abi=3.7=2_cp37m 70 | - pytorch=1.9.0=py3.7_cuda10.2_cudnn7.6.5_0 71 | - qt=5.9.7=h5867ecd_1 72 | - qutil=3.2.1=6 73 | - readline=8.1=h27cfd23_0 74 | - scipy=1.5.2=py37h0b6359f_0 75 | - setuptools=58.0.4=py37h06a4308_0 76 | - sip=4.19.8=py37hf484d3e_0 77 | - six=1.15.0=py_0 78 | - sqlite=3.36.0=hc218d9a_0 79 | - tk=8.6.11=h1ccaba5_0 80 | - torchdiffeq=0.2.2=pyhd8ed1ab_0 81 | - tornado=6.1=py37h27cfd23_0 82 | - typing_extensions=3.10.0.2=pyh06a4308_0 83 | - wheel=0.37.0=pyhd3eb1b0_1 84 | - xz=5.2.5=h7b6447c_0 85 | - zlib=1.2.11=h7b6447c_3 86 | - zstd=1.4.9=haebb681_0 87 | prefix: /home/karlbau/anaconda3/envs/finn 88 | -------------------------------------------------------------------------------- /models/aphynity/config.json: -------------------------------------------------------------------------------- 1 | { 2 | "paths": { 3 | }, 4 | 5 | "general": { 6 | "device": "cpu" // "cpu" or "cuda" 7 | }, 8 | 9 | "training": { 10 | "save_model": true, 11 | "continue_training": false, 12 | "epochs": 100, 13 | "learning_rate": 1.0 14 | }, 15 | 16 | "data": { 17 | "type": "burger_2d", // "burger", "burger_2d", diffusion_sorption", "diffusion_reaction", "allen_cahn" 18 | "name": "data_train", // "data_train", "data_ext", "data_test" 19 | "noise": 0.0 20 | }, 21 | 22 | "model": { 23 | "name": "burger_2d", // "burger", "burger_2d", diff-sorp", "diff-react", "allen-cahn" 24 | "number": 0 // The i-th model 25 | } 26 | } -------------------------------------------------------------------------------- /models/aphynity/experiment.py: -------------------------------------------------------------------------------- 1 | """ 2 | This script either trains n models or tests them as determined in the 3 | config.json file. 4 | """ 5 | 6 | import multiprocessing as mp 7 | import numpy as np 8 | import train 9 | import test 10 | 11 | # 12 | # Parameters 13 | 14 | TRAIN = False 15 | MODEL_NUMBERS = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] 16 | 17 | # Number of models that can be trained in parallel 18 | POOLSIZE = 1 19 | 20 | 21 | # 22 | # SCRIPT 23 | 24 | if TRAIN: 25 | 26 | # Train the MODEL_NUMBERS using POOLSIZE threads in parallel 27 | if POOLSIZE > 1: 28 | arguments = list() 29 | for model_number in MODEL_NUMBERS: 30 | arguments.append([False, model_number]) 31 | 32 | with mp.Pool(POOLSIZE) as pool: 33 | pool.starmap(train.run_training, arguments) 34 | else: 35 | for model_number in MODEL_NUMBERS: 36 | train.run_training(print_progress=False, model_number=model_number) 37 | 38 | else: 39 | # Test the MODEL_NUMBERS iteratively 40 | error_list = list() 41 | for model_number in MODEL_NUMBERS: 42 | mse = test.run_testing(model_number=model_number) 43 | error_list.append(mse) 44 | 45 | print(error_list) 46 | print(f"Average MSE: {np.mean(error_list)}, STD: {np.std(error_list)}") 47 | -------------------------------------------------------------------------------- /models/cnn_node/cnn_node.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | 3 | """ 4 | Physics-derivatives learning layer. Implementation taken and modified from 5 | https://github.com/vincent-leguen/PhyDNet 6 | """ 7 | 8 | import torch 9 | import torch.nn as nn 10 | import numpy as np 11 | from torchdiffeq import odeint_adjoint, odeint 12 | import torch.nn.functional as F 13 | 14 | class ConvNet1D(nn.Module): 15 | def __init__(self, bc, state_c, hidden=16, sigmoid=False, device=torch.device('cpu')): 16 | super().__init__() 17 | kernel_size = 3 18 | padding = kernel_size // 2 19 | self.bc = bc 20 | self.state_c = state_c 21 | self.device = device 22 | self.net = nn.Sequential( 23 | nn.Conv1d(state_c, hidden, kernel_size=kernel_size, padding=0, bias=False), 24 | nn.BatchNorm1d(hidden, track_running_stats=False), 25 | nn.Tanh(), 26 | nn.Conv1d(hidden, hidden, kernel_size=kernel_size, padding=padding, bias=False), 27 | nn.BatchNorm1d(hidden, track_running_stats=False), 28 | nn.Tanh(), 29 | nn.Conv1d(hidden, state_c, kernel_size=kernel_size, padding=padding), 30 | ).to(device=self.device) 31 | if sigmoid: 32 | self.net.add_module("Sigmoid", nn.Sigmoid()) 33 | 34 | def forward(self, t, x): 35 | x = torch.cat((self.bc[:1,:,:1],x,self.bc[:1,:,1:]),dim=2) 36 | return self.net(x) 37 | 38 | class ConvNet2D(nn.Module): 39 | def __init__(self, bc, state_c, hidden=16, sigmoid=False, device=torch.device('cpu')): 40 | super().__init__() 41 | kernel_size = 3 42 | padding = kernel_size // 2 43 | self.bc = bc 44 | self.state_c = state_c 45 | self.device = device 46 | self.net = nn.Sequential( 47 | nn.Conv2d(state_c, hidden, kernel_size=kernel_size, padding=0, bias=False), 48 | nn.BatchNorm2d(hidden, track_running_stats=False), 49 | nn.Tanh(), 50 | nn.Conv2d(hidden, hidden, kernel_size=kernel_size, padding=padding, bias=False), 51 | nn.BatchNorm2d(hidden, track_running_stats=False), 52 | nn.Tanh(), 53 | nn.Conv2d(hidden, state_c, kernel_size=kernel_size, padding=padding), 54 | ).to(device=self.device) 55 | if sigmoid: 56 | self.net.add_module("Sigmoid", nn.Sigmoid()) 57 | 58 | def forward(self, t, x): 59 | Nx = x.size(-2) 60 | x = torch.cat((self.bc[:1,:,:1].repeat(1,1,Nx).unsqueeze(-1),x, 61 | self.bc[:1,:,1:2].repeat(1,1,Nx).unsqueeze(-1)),dim=3) 62 | Ny = x.size(-1) 63 | x = torch.cat((self.bc[:1,:,2:3].repeat(1,1,Ny).unsqueeze(-2),x, 64 | self.bc[:1,:,3:4].repeat(1,1,Ny).unsqueeze(-2)),dim=2) 65 | return self.net(x) 66 | 67 | 68 | class NeuralODE(nn.Module): 69 | def __init__(self, convnet): 70 | super().__init__() 71 | self.convnet = convnet 72 | 73 | def forward(self, t, u0): 74 | u = odeint(self.convnet, y0=u0, t=t) 75 | # u: T x batch_size x n_c (x h x w) 76 | dim_seq = u0.dim() + 1 77 | dims = [1, 0, 2] + list(range(dim_seq))[3:] 78 | return u.permute(*dims) # batch_size x T x n_c (x h x w) -------------------------------------------------------------------------------- /models/cnn_node/config.json: -------------------------------------------------------------------------------- 1 | { 2 | "paths": { 3 | }, 4 | 5 | "general": { 6 | "device": "cpu" // "cpu" or "cuda" 7 | }, 8 | 9 | "training": { 10 | "save_model": true, 11 | "continue_training": false, 12 | "epochs": 100, 13 | "learning_rate": 1 14 | }, 15 | 16 | "data": { 17 | "type": "burger_2d", // "burger", "burger_2d", "diffusion_sorption", "diffusion_reaction", "allen_cahn" 18 | "name": "data_train", // "data_train", "data_ext", "data_test" 19 | "noise": 0.0 20 | }, 21 | 22 | "model": { 23 | "name": "burger_2d", // "burger", "burger_2d", "diff-sorp", "diff-react", "allen-cahn" 24 | "number": 0 // The i-th model 25 | } 26 | } -------------------------------------------------------------------------------- /models/cnn_node/experiment.py: -------------------------------------------------------------------------------- 1 | """ 2 | This script either trains n models or tests them as determined in the 3 | config.json file. 4 | """ 5 | 6 | import multiprocessing as mp 7 | import numpy as np 8 | import train 9 | import test 10 | 11 | # 12 | # Parameters 13 | 14 | TRAIN = False 15 | MODEL_NUMBERS = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] 16 | 17 | # Number of models that can be trained in parallel 18 | POOLSIZE = 1 19 | 20 | 21 | # 22 | # SCRIPT 23 | 24 | if TRAIN: 25 | 26 | # Train the MODEL_NUMBERS using POOLSIZE threads in parallel 27 | if POOLSIZE > 1: 28 | arguments = list() 29 | for model_number in MODEL_NUMBERS: 30 | arguments.append([False, model_number]) 31 | 32 | with mp.Pool(POOLSIZE) as pool: 33 | pool.starmap(train.run_training, arguments) 34 | else: 35 | for model_number in MODEL_NUMBERS: 36 | train.run_training(print_progress=False, model_number=model_number) 37 | 38 | else: 39 | # Test the MODEL_NUMBERS iteratively 40 | error_list = list() 41 | for model_number in MODEL_NUMBERS: 42 | mse = test.run_testing(model_number=model_number) 43 | error_list.append(mse) 44 | 45 | print(error_list) 46 | print(f"Average MSE: {np.mean(error_list)}, STD: {np.std(error_list)}") 47 | -------------------------------------------------------------------------------- /models/convlstm/config.json: -------------------------------------------------------------------------------- 1 | { 2 | "paths": { 3 | }, 4 | 5 | "general": { 6 | "device": "GPU" 7 | }, 8 | 9 | "training": { 10 | "save_model": true, 11 | "continue_training": false, 12 | "epochs": 200, 13 | "batch_size": 1, 14 | "learning_rate": 1, 15 | "t_start": 0, 16 | "t_stop": 150 // burger, burger_2d, allen-cahn 150, diff-sorp 400, diff-react 70 17 | }, 18 | 19 | "validation": { 20 | "batch_size": 1, 21 | "t_start": 150, // burger, burger_2d, allen-cahn 150, diff-sorp 400, diff-react 70 22 | "t_stop": 200 // burger, burger_2d, allen-cahn 200, diff-sorp 500, diff-react 100 23 | }, 24 | 25 | "testing": { 26 | "batch_size": 1, 27 | "teacher_forcing_steps": 20, 28 | "closed_loop_steps": 1990, 29 | "feed_boundary_data": true 30 | }, 31 | 32 | "data": { 33 | "type": "burger_2d", // "burger", "burger_2d", diffusion_sorption", "diffusion_reaction", "allen_cahn" 34 | "name": "data_train", // "data_train", "data_ext", "data_test" 35 | "noise": 0.0 36 | }, 37 | 38 | "model": { 39 | "name": "burger_2d", // "burger", "burger_2d", diff-sorp", "diff-react", "allen-cahn" 40 | "number": 0, // The i-th model 41 | "field_size": [49, 49], // burger and allen-cahn [49], diff-sorp [26], burger_2d and diff-react [49, 49] 42 | "channels": [1, 24, 1], // burger, burger_2d, and allen-cahn [1, ..., 1], diff-sorp and diff-react [2, ..., 2] 43 | "kernel_size": 3 44 | } 45 | } 46 | -------------------------------------------------------------------------------- /models/convlstm/conv_lstm.py: -------------------------------------------------------------------------------- 1 | from torch.autograd import Variable 2 | import torch as th 3 | 4 | 5 | class ConvLSTMCell(th.nn.Module): 6 | """ 7 | A Convolutional LSTM implementation taken from 8 | https://github.com/ndrplz/ConvLSTM_pytorch/blob/master/convlstm.py 9 | """ 10 | 11 | def __init__(self, input_channels, hidden_channels, kernel_size, bias, 12 | dimensions): 13 | """ 14 | Initialize ConvLSTM cell. 15 | :param input_channels: Number of channels of input tensor. 16 | :param hidden_channels: Number of channels of hidden state. 17 | :param kernel_size: Size of the convolutional kernel. 18 | :param bias: Whether or not to add the bias. 19 | :param dimensions: The spatial dimensions of the data 20 | """ 21 | 22 | super(ConvLSTMCell, self).__init__() 23 | 24 | self.input_channels = input_channels 25 | self.hidden_channels = hidden_channels 26 | 27 | self.kernel_size = kernel_size 28 | self.padding = kernel_size // 2 29 | self.bias = bias 30 | 31 | conv_function = th.nn.Conv1d if dimensions == 1 else th.nn.Conv2d 32 | self.conv = conv_function( 33 | in_channels=self.input_channels + self.hidden_channels, 34 | out_channels=4 * self.hidden_channels, 35 | kernel_size=self.kernel_size, 36 | padding=self.padding, 37 | bias=self.bias 38 | ) 39 | 40 | 41 | def forward(self, input_tensor, cur_state): 42 | h_cur, c_cur = cur_state 43 | 44 | # concatenate along channel axis 45 | combined = th.cat([input_tensor, h_cur], dim=1) 46 | 47 | combined_conv = self.conv(combined) 48 | cc_i, cc_f, cc_o, cc_g = th.split( 49 | combined_conv, self.hidden_channels, dim=1 50 | ) 51 | i = th.sigmoid(cc_i) 52 | f = th.sigmoid(cc_f) 53 | o = th.sigmoid(cc_o) 54 | g = th.tanh(cc_g) 55 | 56 | c_next = f * c_cur + i * g 57 | h_next = o * th.tanh(c_next) 58 | 59 | return h_next, c_next 60 | 61 | 62 | class ConvLSTM(th.nn.Module): 63 | """ 64 | ConvLSTM class to model the 1D Burgers equation, the two 1D diffusion 65 | sorption equations or the 2d diffusion reaction equations. 66 | """ 67 | 68 | def __init__(self, config, device): 69 | """ 70 | Constructor method to initialize the TCN instance. 71 | :param config: The configuration class of the model 72 | :param device: The device (GPU or CPU) which processes the tensors 73 | """ 74 | super(ConvLSTM, self).__init__() 75 | 76 | self.layers = th.nn.ModuleList() 77 | self.state_list = [] 78 | 79 | for ch_idx in range(1, len(config.model.channels)): 80 | self.layers.append(ConvLSTMCell( 81 | input_channels=config.model.channels[ch_idx - 1], 82 | hidden_channels=config.model.channels[ch_idx], 83 | kernel_size=config.model.kernel_size, 84 | bias=True, 85 | dimensions=len(config.model.field_size) 86 | )) 87 | 88 | # Initialize an exemplary state 89 | s = th.zeros(config.training.batch_size, 90 | config.model.channels[ch_idx], 91 | *config.model.field_size, 92 | device=device) 93 | 94 | # Append a pair of h and c states to the state list 95 | self.state_list.append([s.clone(), s.clone()]) 96 | 97 | def forward(self, input_tensor, cur_state_list=None): 98 | """ 99 | Forward pass of the ConvLSTMModel class 100 | """ 101 | 102 | if cur_state_list is None: 103 | cur_state_list = self.state_list 104 | 105 | next_state_list = [] 106 | 107 | for layer_idx, layer in enumerate(self.layers): 108 | h, c = layer.forward(input_tensor=input_tensor, 109 | cur_state=cur_state_list[layer_idx]) 110 | 111 | next_state_list.append([h, c]) 112 | input_tensor = h 113 | 114 | self.state_list = next_state_list 115 | 116 | return h, next_state_list 117 | 118 | def reset(self, batch_size, device): 119 | """ 120 | Resets the ConvLSTM's states to zero by potentially changing the 121 | batch_size. 122 | """ 123 | for layer_idx, state_pair in enumerate(self.state_list): 124 | s = th.zeros(batch_size, *state_pair[0][0].shape, device=device) 125 | self.state_list[layer_idx] = [s.clone(), s.clone()] 126 | -------------------------------------------------------------------------------- /models/convlstm/experiment.py: -------------------------------------------------------------------------------- 1 | """ 2 | This script either trains n models or tests them as determined in the 3 | config.json file. 4 | """ 5 | 6 | import multiprocessing as mp 7 | import numpy as np 8 | import train 9 | import test 10 | 11 | # 12 | # Parameters 13 | 14 | TRAIN = False 15 | MODEL_NUMBERS = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] 16 | 17 | # Number of models that can be trained in parallel 18 | POOLSIZE = 5 19 | 20 | 21 | # 22 | # SCRIPT 23 | 24 | if TRAIN: 25 | 26 | # Train the MODEL_NUMBERS using POOLSIZE threads in parallel 27 | if POOLSIZE > 1: 28 | arguments = list() 29 | for model_number in MODEL_NUMBERS: 30 | arguments.append([False, model_number]) 31 | 32 | with mp.Pool(POOLSIZE) as pool: 33 | pool.starmap(train.run_training, arguments) 34 | else: 35 | for model_number in MODEL_NUMBERS: 36 | train.run_training(print_progress=False, model_number=model_number) 37 | 38 | else: 39 | # Test the MODEL_NUMBERS iteratively 40 | error_list = list() 41 | for model_number in MODEL_NUMBERS: 42 | mse = test.run_testing(model_number=model_number) 43 | error_list.append(mse) 44 | 45 | print(error_list) 46 | print(f"Average MSE: {np.mean(error_list)}, STD: {np.std(error_list)}") 47 | -------------------------------------------------------------------------------- /models/convlstm/test.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch as th 3 | import time 4 | import glob 5 | import os 6 | import matplotlib.pyplot as plt 7 | import matplotlib.animation as animation 8 | import sys 9 | 10 | sys.path.append("..") 11 | from utils.configuration import Configuration 12 | import utils.helper_functions as helpers 13 | from conv_lstm import ConvLSTM 14 | 15 | 16 | def run_testing(print_progress=False, visualize=False, model_number=None): 17 | 18 | th.set_num_threads(1) 19 | 20 | # Load the user configurations 21 | config = Configuration("config.json") 22 | 23 | # Append the model number to the name of the model 24 | if model_number is None: 25 | model_number = config.model.number 26 | config.model.name = config.model.name + "_" + str(model_number).zfill(2) 27 | 28 | # Hide the GPU(s) in case the user specified to use the CPU in the config 29 | # file 30 | if config.general.device == "CPU": 31 | os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" 32 | os.environ["CUDA_VISIBLE_DEVICES"] = "" 33 | 34 | # Set device on GPU if specified in the configuration file, else CPU 35 | device = helpers.determine_device() 36 | 37 | # Initialize and set up the network 38 | model = ConvLSTM(config=config, device=device).to(device=device) 39 | 40 | if print_progress: 41 | # Count and print number of trainable parameters 42 | pytorch_total_params = sum( 43 | p.numel() for p in model.parameters() if p.requires_grad 44 | ) 45 | print("Trainable model parameters:", pytorch_total_params) 46 | 47 | # Restore the network by loading the weights saved in the .pt file 48 | print("Restoring model (that is the network\"s weights) from file...") 49 | 50 | model.load_state_dict(th.load(os.path.join(os.path.abspath(""), 51 | "checkpoints", 52 | config.model.name, 53 | config.model.name + ".pt"))) 54 | model.eval() 55 | 56 | """ 57 | TESTING 58 | """ 59 | 60 | # 61 | # Load data depending on the task 62 | if config.data.type == "burger": 63 | data_path = os.path.join("../../data/", 64 | config.data.type, 65 | config.data.name, 66 | "sample.npy") 67 | data = np.array(np.load(data_path), dtype=np.float32) 68 | data = np.expand_dims(data, axis=1) 69 | 70 | elif config.data.type == "diffusion_sorption": 71 | data_path_base = os.path.join("../../data/", 72 | config.data.type, 73 | config.data.name) 74 | data_path_c = os.path.join(data_path_base, "sample_c.npy") 75 | data_path_ct = os.path.join(data_path_base, "sample_ct.npy") 76 | data_c = np.array(np.load(data_path_c), dtype=np.float32) 77 | data_ct = np.array(np.load(data_path_ct), dtype=np.float32) 78 | data = np.stack((data_c, data_ct), axis=1) 79 | 80 | elif config.data.type == "diffusion_reaction": 81 | data_path_base = os.path.join("../../data/", 82 | config.data.type, 83 | config.data.name) 84 | data_path_u = os.path.join(data_path_base, "sample_u.npy") 85 | data_path_v = os.path.join(data_path_base, "sample_v.npy") 86 | data_u = np.array(np.load(data_path_u), dtype=np.float32) 87 | data_v = np.array(np.load(data_path_v), dtype=np.float32) 88 | data = np.stack((data_u, data_v), axis=1) 89 | 90 | elif config.data.type == "allen_cahn": 91 | data_path = os.path.join("../../data/", 92 | config.data.type, 93 | config.data.name, 94 | "sample.npy") 95 | data = np.array(np.load(data_path), dtype=np.float32) 96 | data = np.expand_dims(data, axis=1) 97 | 98 | elif config.data.type == "burger_2d": 99 | data_path = os.path.join("../../data/", 100 | config.data.type, 101 | config.data.name, 102 | "sample.npy") 103 | data = np.array(np.load(data_path), dtype=np.float32) 104 | data = np.expand_dims(data, axis=1) 105 | 106 | # 107 | # Set up the test data 108 | data_test = th.tensor(data, device=device).unsqueeze(1) 109 | sequence_length = len(data_test) - 1 110 | 111 | # Evaluate the network for the given test data 112 | 113 | # Separate the data into network inputs and labels 114 | net_inputs = th.clone(data_test[:-1]) 115 | net_labels = th.clone(data_test[1:]) 116 | 117 | # Set up an array of zeros to store the network outputs 118 | net_outputs = th.zeros(size=(sequence_length, 119 | config.testing.batch_size, 120 | config.model.channels[-1], 121 | *config.model.field_size), 122 | device=device) 123 | state_list = None 124 | 125 | # Iterate over the remaining sequence of the training example and perform a 126 | # forward pass 127 | time_start = time.time() 128 | for t in range(len(net_inputs)): 129 | 130 | if t < config.testing.teacher_forcing_steps: 131 | # Teacher forcing 132 | net_input = net_inputs[t] 133 | else: 134 | # Closed loop 135 | net_input = net_outputs[t - 1] 136 | 137 | # Feed the boundary data also in closed loop if desired 138 | if config.testing.feed_boundary_data: 139 | net_input[:, :, 0] = net_inputs[t, :, :, 0] 140 | net_input[:, :, -1] = net_inputs[t, :, :, -1] 141 | 142 | net_output, state_list = model.forward(input_tensor=net_input, 143 | cur_state_list=state_list) 144 | net_outputs[t] = net_output[-1] 145 | 146 | if print_progress: 147 | forward_pass_duration = time.time() - time_start 148 | print("Forward pass took:", forward_pass_duration, "seconds.") 149 | 150 | # Convert the PyTorch network output tensor into a numpy array 151 | net_outputs = net_outputs.cpu().detach().numpy()[:, 0, 0] 152 | net_labels = net_labels.cpu().detach().numpy()[:, 0, 0] 153 | 154 | # 155 | # Visualize the data 156 | if visualize: 157 | plt.style.use("dark_background") 158 | 159 | # Plot over space and time 160 | fig, ax = plt.subplots(1, 2, figsize=(16, 6), sharey=True) 161 | 162 | if config.data.type == "burger" or\ 163 | config.data.type == "diffusion_sorption" or\ 164 | config.data.type == "allen_cahn": 165 | 166 | im1 = ax[0].imshow( 167 | np.transpose(net_labels), interpolation='nearest', 168 | cmap='rainbow', origin='lower', aspect='auto', vmin=-0.4, 169 | vmax=0.4 170 | ) 171 | fig.colorbar(im1, ax=ax[0]) 172 | im2 = ax[1].imshow( 173 | np.transpose(net_outputs), interpolation='nearest', 174 | cmap='rainbow', origin='lower', aspect='auto', vmin=-0.4, 175 | vmax=0.4 176 | ) 177 | fig.colorbar(im2, ax=ax[1]) 178 | 179 | ax[0].set_xlabel("t") 180 | ax[0].set_ylabel("x") 181 | ax[1].set_xlabel("t") 182 | 183 | elif config.data.type == "diffusion_reaction": 184 | 185 | im1 = ax[0].imshow( 186 | np.transpose(net_labels[..., 0]), interpolation='nearest', 187 | cmap='rainbow', origin='lower', aspect='auto', vmin=-0.4, 188 | vmax=0.4 189 | ) 190 | fig.colorbar(im1, ax=ax[0]) 191 | im2 = ax[1].imshow( 192 | np.transpose(net_outputs[..., 0]), interpolation='nearest', 193 | cmap='rainbow', origin='lower', aspect='auto', vmin=-0.4, 194 | vmax=0.4 195 | ) 196 | fig.colorbar(im2, ax=ax[1]) 197 | 198 | ax[0].set_xlabel("x") 199 | ax[0].set_ylabel("y") 200 | ax[1].set_xlabel("x") 201 | 202 | elif config.data.type == "burger_2d": 203 | 204 | im1 = ax[0].imshow( 205 | np.transpose(net_labels[-1, ...]), interpolation='nearest', 206 | cmap='rainbow', origin='lower', aspect='auto', vmin=-0.4, 207 | vmax=0.4 208 | ) 209 | fig.colorbar(im1, ax=ax[0]) 210 | im2 = ax[1].imshow( 211 | np.transpose(net_outputs[-1, ...]), interpolation='nearest', 212 | cmap='rainbow', origin='lower', aspect='auto', vmin=-0.4, 213 | vmax=0.4 214 | ) 215 | fig.colorbar(im2, ax=ax[1]) 216 | 217 | ax[0].set_xlabel("x") 218 | ax[0].set_ylabel("y") 219 | ax[1].set_xlabel("x") 220 | 221 | 222 | ax[0].set_title("Ground Truth") 223 | ax[1].set_title("Network Output") 224 | 225 | 226 | if config.data.type == "diffusion_reaction"\ 227 | or config.data.type == "burger_2d": 228 | anim = animation.FuncAnimation( 229 | fig, 230 | animate, 231 | frames=sequence_length, 232 | fargs=(im1, im2, net_labels, net_outputs), 233 | interval=20 234 | ) 235 | 236 | plt.show() 237 | 238 | # Compute error 239 | mse = np.mean(np.square(net_outputs - net_labels)) 240 | 241 | return mse 242 | 243 | 244 | def animate(t, im1, im2, net_label, net_outputs): 245 | """ 246 | Data animation function animating an image over time. 247 | :param t: The current time step 248 | :param axis: The matplotlib image object 249 | :param field: The data field 250 | :return: The matplotlib image object updated with the current time step's 251 | image date 252 | """ 253 | im1.set_array(net_label[t]) 254 | im2.set_array(net_outputs[t]) 255 | 256 | 257 | if __name__ == "__main__": 258 | mse = run_testing(print_progress=True, visualize=True) 259 | print(f"MSE: {mse}") 260 | 261 | print("Done.") -------------------------------------------------------------------------------- /models/convlstm/train.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch as th 3 | import torch.nn as nn 4 | import time 5 | import glob 6 | import os 7 | import matplotlib.pyplot as plt 8 | from threading import Thread 9 | import sys 10 | 11 | sys.path.append("..") 12 | from utils.configuration import Configuration 13 | import utils.helper_functions as helpers 14 | from conv_lstm import ConvLSTM 15 | 16 | 17 | def process_sequence(model, criterion, data, batch_size, epoch, config, 18 | device): 19 | 20 | # Separate the data into network inputs and labels 21 | net_inputs = data[:-1] 22 | net_labels = data[1:] 23 | 24 | # Set up an array of zeros to store the network outputs 25 | net_outputs = th.zeros(size=(len(net_labels), 26 | batch_size, 27 | config.model.channels[-1], 28 | *config.model.field_size), 29 | device=device) 30 | 31 | model.reset(batch_size=batch_size, device=device) 32 | 33 | # Initial network input and forward pass 34 | net_output, state_list = model.forward(input_tensor=net_inputs[0], 35 | cur_state_list=None) 36 | net_outputs[0] = net_output 37 | 38 | # Probability of closed loop in current epoch 39 | p_cl = epoch / (config.training.epochs * 1.5) 40 | 41 | # Iterate over the whole sequence of the training example and 42 | # perform a forward pass 43 | for t in range(1, len(net_inputs)): 44 | 45 | # Draw random number to determine teacher forcing or closed loop 46 | # (scheduled sampling) 47 | p_uniform = np.random.uniform(0, 1) 48 | closed_loop = True if p_uniform < p_cl else False 49 | 50 | if not closed_loop: 51 | # Teacher forcing 52 | net_input = net_inputs[t] 53 | else: 54 | # Closed loop 55 | net_input = net_outputs[t - 1] 56 | 57 | # Forward the input through the network 58 | #net_output, _ = model.forward(input_tensor=net_inputs[t]) 59 | net_output, state_list = model.forward(input_tensor=net_input, 60 | cur_state_list=state_list) 61 | 62 | # Store the output of the network for this sequence step 63 | #net_outputs[t] = net_output 64 | net_outputs[t] = net_output[-1] 65 | 66 | # Compute the mean squared error 67 | mse = criterion(net_outputs, net_labels) 68 | 69 | return net_outputs, mse 70 | 71 | 72 | def run_training(print_progress=False, model_number=None): 73 | 74 | # Set a random seed for varying weight initializations 75 | th.seed() 76 | 77 | th.set_num_threads(1) 78 | 79 | # Load the user configurations 80 | config = Configuration("config.json") 81 | 82 | # Append the model number to the name of the model 83 | if model_number is None: 84 | model_number = config.model.number 85 | config.model.name = config.model.name + "_" + str(model_number).zfill(2) 86 | 87 | print("Model name:", config.model.name) 88 | 89 | # Hide the GPU(s) in case the user specified to use the CPU in the config 90 | # file 91 | if config.general.device == "CPU": 92 | os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" # see issue #152 93 | os.environ["CUDA_VISIBLE_DEVICES"] = "" 94 | 95 | time_start = time.time() 96 | 97 | # setting device on GPU if available, else CPU 98 | device = helpers.determine_device() 99 | 100 | # Initialize and set up the network 101 | model = ConvLSTM(config=config, device=device).to(device=device) 102 | 103 | if print_progress: 104 | # Count and print number of trainable parameters 105 | pytorch_total_params = sum( 106 | p.numel() for p in model.parameters() if p.requires_grad 107 | ) 108 | print("Trainable model parameters:", pytorch_total_params) 109 | 110 | # If desired, restore the network by loading the weights saved in the .pt 111 | # file 112 | if config.training.continue_training: 113 | print("Restoring model (that is the network\"s weights) from file...") 114 | model.load_state_dict(th.load(os.path.join(os.path.abspath(""), 115 | "checkpoints", 116 | config.model.name, 117 | config.model.name + ".pt"))) 118 | model.train() 119 | 120 | # 121 | # Set up the optimizer and the criterion (loss) 122 | optimizer = th.optim.LBFGS(model.parameters(), 123 | lr=config.training.learning_rate) 124 | criterion = nn.MSELoss() 125 | 126 | # 127 | # Load data depending on the task 128 | if config.data.type == "burger": 129 | data_path = os.path.join("../../data/", 130 | config.data.type, 131 | config.data.name, 132 | "sample.npy") 133 | data = np.array(np.load(data_path), dtype=np.float32) 134 | data = np.expand_dims(data, axis=1) 135 | 136 | elif config.data.type == "diffusion_sorption": 137 | data_path_base = os.path.join("../../data/", 138 | config.data.type, 139 | config.data.name) 140 | data_path_c = os.path.join(data_path_base, "sample_c.npy") 141 | data_path_ct = os.path.join(data_path_base, "sample_ct.npy") 142 | data_c = np.array(np.load(data_path_c), dtype=np.float32) 143 | data_ct = np.array(np.load(data_path_ct), dtype=np.float32) 144 | data = np.stack((data_c, data_ct), axis=1) 145 | 146 | elif config.data.type == "diffusion_reaction": 147 | data_path_base = os.path.join("../../data/", 148 | config.data.type, 149 | config.data.name) 150 | data_path_u = os.path.join(data_path_base, "sample_u.npy") 151 | data_path_v = os.path.join(data_path_base, "sample_v.npy") 152 | data_u = np.array(np.load(data_path_u), dtype=np.float32) 153 | data_v = np.array(np.load(data_path_v), dtype=np.float32) 154 | data = np.stack((data_u, data_v), axis=1) 155 | 156 | elif config.data.type == "allen_cahn": 157 | data_path = os.path.join("../../data/", 158 | config.data.type, 159 | config.data.name, 160 | "sample.npy") 161 | data = np.array(np.load(data_path), dtype=np.float32) 162 | data = np.expand_dims(data, axis=1) 163 | 164 | elif config.data.type == "burger_2d": 165 | data_path = os.path.join("../../data/", 166 | config.data.type, 167 | config.data.name, 168 | "sample.npy") 169 | data = np.array(np.load(data_path), dtype=np.float32) 170 | data = np.expand_dims(data, axis=1) 171 | 172 | # Set up the training and validation datasets and -loaders 173 | data_train = th.tensor( 174 | data[:config.training.t_stop], 175 | device=device 176 | ).unsqueeze(1) 177 | data_valid = th.tensor( 178 | data[config.validation.t_start:config.validation.t_stop], 179 | device=device 180 | ).unsqueeze(1) 181 | 182 | # 183 | # Set up lists to save and store the epoch errors 184 | epoch_errors_train = [] 185 | epoch_errors_valid = [] 186 | best_train = np.infty 187 | best_valid = np.infty 188 | 189 | 190 | """ 191 | TRAINING 192 | """ 193 | 194 | a = time.time() 195 | 196 | # 197 | # Start the training and iterate over all epochs 198 | for epoch in range(config.training.epochs): 199 | 200 | epoch_start_time = time.time() 201 | 202 | # 203 | # Train the network on the given data 204 | 205 | def closure(): 206 | # Set the model to train mode 207 | model.train() 208 | 209 | # Reset the optimizer to clear data from previous iterations 210 | optimizer.zero_grad() 211 | 212 | # Forward the input through the network 213 | _, mse = process_sequence( 214 | model=model, 215 | criterion=criterion, 216 | data=data_train, 217 | batch_size=config.training.batch_size, 218 | epoch=epoch, 219 | config=config, 220 | device=device 221 | ) 222 | 223 | mse.backward() 224 | 225 | return mse 226 | 227 | # Backpropagate the error and perform a weight update 228 | optimizer.step(closure) 229 | 230 | # Extract the MSE value from the closure function 231 | mse = closure() 232 | 233 | # Append the error to the error list 234 | epoch_errors_train.append(mse.item()) 235 | 236 | train_sign = "(-)" 237 | if epoch_errors_train[-1] < best_train: 238 | best_train = epoch_errors_train[-1] 239 | train_sign = "(+)" 240 | 241 | # 242 | # Validate the network 243 | 244 | net_output, mse = process_sequence( 245 | model=model, 246 | criterion=criterion, 247 | data=data_valid, 248 | batch_size=config.validation.batch_size, 249 | epoch=epoch, 250 | config=config, 251 | device=device 252 | ) 253 | epoch_errors_valid.append(mse.item()) 254 | 255 | # Create a plus or minus sign for the validation error 256 | valid_sign = "(-)" 257 | if epoch_errors_valid[-1] < best_valid: 258 | best_valid = epoch_errors_valid[-1] 259 | valid_sign = "(+)" 260 | 261 | if config.training.save_model: 262 | # Start a separate thread to save the model 263 | thread = Thread(target=helpers.save_model_to_file( 264 | model_src_path=os.path.abspath(""), 265 | config=config, 266 | epoch=epoch, 267 | epoch_errors_train=epoch_errors_train, 268 | epoch_errors_valid=epoch_errors_valid, 269 | net=model)) 270 | thread.start() 271 | 272 | # 273 | # Print progress to the console 274 | if print_progress: 275 | print(f"Epoch {str(epoch+1).zfill(int(np.log10(config.training.epochs))+1)}/{str(config.training.epochs)} took {str(np.round(time.time() - epoch_start_time, 2)).ljust(5, '0')} seconds. \t\tAverage epoch training error: {train_sign}{str(np.round(epoch_errors_train[-1], 10)).ljust(12, ' ')} \t\tValidation error: {valid_sign}{str(np.round(epoch_errors_valid[-1], 10)).ljust(12, ' ')}") 276 | 277 | 278 | b = time.time() 279 | if print_progress: 280 | print("\nTraining took " + str(np.round(b - a, 2)) + " seconds.\n\n") 281 | 282 | 283 | if __name__ == "__main__": 284 | run_training(print_progress=True) 285 | 286 | print("Done.") -------------------------------------------------------------------------------- /models/cvpinn/cvpinns.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | 3 | """ 4 | https://github.com/rgp62/cvpinns/blob/master/cvpinns.py 5 | """ 6 | 7 | import numpy as np 8 | import tensorflow as tf 9 | 10 | class PDE: 11 | def __init__(self,L,nx,quad,F0,F1,IC,BCl,BCr,u): 12 | """ 13 | Inputs: 14 | L: 2d list containing the lower and upper bounds of the domain in each direction, 15 | [[x0min,x0max],[x1min,x1max]] 16 | nx: 1d list containing the number of cells in each direction 17 | quad: dictionary specifying quadratures, {'0':(wi0,xi0),'1',(wi1,xi1)}. (wi0,xi0) 18 | and (wi1,xi1) are the weights and biases used to integrate fluxes in the 0 and 1 19 | directions, respectively. Abscissas \in [-1,1] and \sum {weights} = 2 20 | F0: function that computes flux in 0 direction, given u 21 | F1: function that computes flux in 1 direction, given u 22 | IC: function that computes flux in 0 direction along x0=x0min 23 | BCl: function that computes flux in 1 direction along x1=x1min 24 | BCr: function that computes flux in 1 direction along x1=x1max 25 | u: neural network for PDE solution 26 | """ 27 | 28 | 29 | self.F0 = F0 30 | self.F1 = F1 31 | self.IC = IC 32 | self.BCl = BCl 33 | self.BCr = BCr 34 | self.u = u 35 | 36 | self.L = L 37 | self.nx = nx 38 | self.quad = quad 39 | 40 | 41 | xi0 = quad['0'] [0] 42 | xi1 = quad['1'] [0] 43 | 44 | self.wi0 = quad['0'] [1] 45 | self.wi1 = quad['1'] [1] 46 | 47 | 48 | 49 | x0 = np.linspace(L[0][0],L[0][1],nx[0]+1) 50 | x1 = np.linspace(L[1][0],L[1][1],nx[1]+1) 51 | 52 | x0c = (x0[0:-1]+x0[1:])/2 53 | x1c = (x1[0:-1]+x1[1:])/2 54 | xc = np.transpose(np.meshgrid(x0c,x1c,indexing='ij'),(1,2,0)) 55 | 56 | 57 | dx = xc[1,1] - xc[0,0] 58 | self.dx = dx 59 | 60 | self.x0 = tf.constant(np.expand_dims(np.transpose(np.meshgrid(x0,x1c,indexing='ij'),(1,2,0)),2) \ 61 | + np.stack([np.zeros(len(xi0)),dx[1]/2*xi0],1)) 62 | self.x1 = tf.constant(np.expand_dims(np.transpose(np.meshgrid(x0c,x1,indexing='ij'),(1,2,0)),1) \ 63 | + np.expand_dims(np.stack([dx[0]/2*xi1,np.zeros(len(xi1))],1),1)) 64 | 65 | self.x = xc 66 | 67 | x0r = tf.reshape(self.x0,(nx[0]+1,nx[1]*len(xi0),2)) 68 | 69 | jrmax = nx[1]*len(xi0) 70 | 71 | if xi1[0]<-1.+1e-10: 72 | jl = 1 73 | else: 74 | jl = 0 75 | if xi1[-1]>1-1e-10: 76 | jr = jrmax-1 77 | else: 78 | jr = jrmax 79 | 80 | self.x0_I = x0r[1::,jl:jr] 81 | self.x0_1L = x0r[:,0:jl] 82 | self.x0_1R = x0r[:,jr:jrmax] 83 | self.x0_0L = tf.expand_dims(x0r[0,jl:jr],0) 84 | 85 | 86 | x1r = tf.reshape(self.x1,(nx[0]*len(xi1),nx[1]+1,2)) 87 | 88 | jrmax = nx[0]*len(xi1) 89 | 90 | if xi0[0]<-1.+1e-10: 91 | jl = 1 92 | else: 93 | jl = 0 94 | if xi0[-1]>1-1e-10: 95 | jr = jrmax-1 96 | else: 97 | jr = jrmax 98 | 99 | 100 | self.xl_I = x1r[1::,1:-1] 101 | self.x1_1L = tf.expand_dims(x1r[:,0],1) 102 | self.x1_1R = tf.expand_dims(x1r[:,-1],1) 103 | 104 | self.xl_0L = tf.expand_dims(x1r[0,1:-1],0) 105 | 106 | 107 | 108 | @tf.function 109 | def getRES(self): 110 | """ 111 | Outputs: 112 | RES: CVPINNs residual 113 | """ 114 | 115 | F0h = tf.concat([self.F0(self.BCl(self.x0_1L)), 116 | tf.concat([self.F0(self.IC(self.x0_0L)), 117 | self.F0(self.u(self.x0_I))],0), 118 | self.F0(self.BCr(self.x0_1R))],1) 119 | F0i = tf.reshape(F0h,(self.nx[0]+1,self.nx[1],len(self.quad['0'][0]),F0h.shape[-1])) 120 | 121 | F1h = tf.concat([self.F1(self.BCl(self.x1_1L)), 122 | tf.concat([self.F1(self.IC(self.xl_0L)), 123 | self.F1(self.u(self.xl_I))],0), 124 | self.F1(self.BCr(self.x1_1R))],1) 125 | F1i = tf.reshape(F1h,(self.nx[0],len(self.quad['1'][0]),self.nx[1]+1,F0h.shape[-1])) 126 | 127 | F0int = self.dx[1]/2*tf.tensordot(F0i,self.wi0,[2,0]) 128 | F1int = self.dx[0]/2*tf.tensordot(F1i,self.wi1,[1,0]) 129 | RES = (F0int[1::] - F0int[0:-1]) \ 130 | +(F1int[:,1::] - F1int[:,0:-1]) 131 | return RES 132 | -------------------------------------------------------------------------------- /models/distana/config.json: -------------------------------------------------------------------------------- 1 | { 2 | "paths": { 3 | }, 4 | 5 | "general": { 6 | "device": "CPU" // "CPU" or "GPU" 7 | }, 8 | 9 | "training": { 10 | "save_model": true, 11 | "continue_training": false, 12 | "epochs": 200, 13 | "batch_size": 1, 14 | "learning_rate": 1, 15 | "t_start": 0, 16 | "t_stop": 150 // burger, burger_2d, allen-cahn 150, diff-sorp 400, diff-react 70 17 | }, 18 | 19 | "validation": { 20 | "batch_size": 1, 21 | "t_start": 150, // burger, burger_2d, allen-cahn 150, diff-sorp 400, diff-react 70 22 | "t_stop": 200 // burger, burger_2d, allen-cahn 200, diff-sorp 500, diff-react 100 23 | }, 24 | 25 | "testing": { 26 | "batch_size": 1, 27 | "teacher_forcing_steps": 20, 28 | "closed_loop_steps": 1990, 29 | "feed_boundary_data": true 30 | }, 31 | 32 | "data": { 33 | "type": "burger_2d", // "burger", "burger_2d", "diffusion_sorption", "diffusion_reaction", "allen_cahn" 34 | "name": "data_test", // "data_train", "data_ext", "data_test" 35 | "noise": 0.0 36 | }, 37 | 38 | "model": { 39 | "name": "burger_2d", // "burger", "burger_2d", "diff-sorp", "diff-react", "allen-cahn" 40 | "number": 0, // i-th trained model 41 | "field_size": [49, 49], // burger and allen-cahn [49], diff-sorp [26], burger_2d and diff-react [49, 49] 42 | "dynamic_channels": [1], // burger, burger_2d, allen-cahn [1], diff-sorp [2], diff-react [2] 43 | "lateral_channels": [1], 44 | "hidden_channels": [16], 45 | "kernel_size": 3, 46 | "lateral_size": 1 47 | } 48 | } 49 | -------------------------------------------------------------------------------- /models/distana/conv_lstm.py: -------------------------------------------------------------------------------- 1 | from torch.autograd import Variable 2 | import torch as th 3 | 4 | 5 | class ConvLSTMCell(th.nn.Module): 6 | """ 7 | A Convolutional LSTM implementation taken from 8 | https://github.com/ndrplz/ConvLSTM_pytorch/blob/master/convlstm.py 9 | """ 10 | 11 | def __init__(self, input_channels, hidden_channels, kernel_size, bias, 12 | dimensions): 13 | """ 14 | Initialize ConvLSTM cell. 15 | :param input_channels: Number of channels of input tensor. 16 | :param hidden_channels: Number of channels of hidden state. 17 | :param kernel_size: Size of the convolutional kernel. 18 | :param bias: Whether or not to add the bias. 19 | :param dimensions: The spatial dimensions of the data 20 | """ 21 | 22 | super(ConvLSTMCell, self).__init__() 23 | 24 | self.input_channels = input_channels 25 | self.hidden_channels = hidden_channels 26 | 27 | self.kernel_size = kernel_size 28 | self.padding = kernel_size // 2 29 | self.bias = bias 30 | 31 | conv_function = th.nn.Conv1d if dimensions == 1 else th.nn.Conv2d 32 | self.conv = conv_function( 33 | in_channels=self.input_channels + self.hidden_channels, 34 | out_channels=4 * self.hidden_channels, 35 | kernel_size=self.kernel_size, 36 | padding=self.padding, 37 | bias=self.bias 38 | ) 39 | 40 | 41 | def forward(self, input_tensor, cur_state): 42 | h_cur, c_cur = cur_state 43 | 44 | # concatenate along channel axis 45 | combined = th.cat([input_tensor, h_cur], dim=1) 46 | 47 | combined_conv = self.conv(combined) 48 | cc_i, cc_f, cc_o, cc_g = th.split( 49 | combined_conv, self.hidden_channels, dim=1 50 | ) 51 | i = th.sigmoid(cc_i) 52 | f = th.sigmoid(cc_f) 53 | o = th.sigmoid(cc_o) 54 | g = th.tanh(cc_g) 55 | 56 | c_next = f * c_cur + i * g 57 | h_next = o * th.tanh(c_next) 58 | 59 | return h_next, c_next 60 | -------------------------------------------------------------------------------- /models/distana/distana.py: -------------------------------------------------------------------------------- 1 | import torch as th 2 | import torch.nn as nn 3 | from conv_lstm import ConvLSTMCell 4 | 5 | 6 | class DISTANACell(nn.Module): 7 | """ 8 | This class contains the kernelized network topology for the spatio-temporal 9 | propagation of information 10 | """ 11 | 12 | def __init__(self, dyn_channels, lat_channels, hidden_channels, kernel_size, 13 | batch_size, data_size, bias, device): 14 | 15 | super(DISTANACell, self).__init__() 16 | 17 | # Initialize DISTANA's states for lateral data exchange and lstm h,c 18 | #s = th.zeros(batch_size, hidden_channels, *data_size, device=device) 19 | self.l = th.zeros(batch_size, lat_channels, *data_size, device=device) 20 | self.h = th.zeros(batch_size, hidden_channels, *data_size, device=device) 21 | self.c = th.zeros(batch_size, hidden_channels, *data_size, device=device) 22 | 23 | self.device = device 24 | self.dyn_channels = dyn_channels 25 | 26 | dimensions = len(data_size) 27 | conv_function = th.nn.Conv1d if dimensions == 1 else th.nn.Conv2d 28 | 29 | # Lateral input convolution layer 30 | self.lat_in_conv_layer = conv_function( 31 | in_channels=lat_channels, 32 | out_channels=lat_channels, 33 | kernel_size=3, 34 | stride=1, 35 | padding=1, 36 | bias=bias 37 | ).to(device=device) 38 | 39 | # Dynamic and lateral input preprocessing layer 40 | self.pre_layer = conv_function( 41 | in_channels=dyn_channels + lat_channels, 42 | out_channels=hidden_channels, 43 | kernel_size=3, 44 | stride=1, 45 | padding=1, 46 | bias=bias 47 | ).to(device=device) 48 | 49 | # Central LSTM layer 50 | self.clstm = ConvLSTMCell( 51 | input_channels=hidden_channels, 52 | hidden_channels=hidden_channels, 53 | kernel_size=kernel_size, 54 | bias=bias, 55 | dimensions=dimensions 56 | ).to(device=device) 57 | 58 | # Postprocessing layer 59 | self.post_layer = conv_function( 60 | in_channels=hidden_channels, 61 | out_channels=dyn_channels + lat_channels, 62 | kernel_size=3, 63 | stride=1, 64 | padding=1, 65 | bias=bias 66 | ).to(device=device) 67 | 68 | def forward(self, dyn_in, state=None): 69 | """ 70 | Runs the forward pass of all PKs and TKs, respectively, in parallel for 71 | a given input 72 | :param dyn_in: The dynamic input for the PKs 73 | :param state: The latent state of the module [l, h, c] 74 | """ 75 | 76 | if state is None: 77 | lat_in, lstm_h, lstm_c = self.l, self.h, self.c 78 | else: 79 | lat_in, lstm_h, lstm_c = state 80 | 81 | # Compute the lateral input as convolution of the latent lateral 82 | # outputs from the previous timestep 83 | lat_in = self.lat_in_conv_layer(lat_in) 84 | 85 | # Forward the dynamic and lateral inputs through the preprocessing 86 | # layer 87 | dynlat_in = th.cat(tensors=(dyn_in, lat_in), dim=1) 88 | pre_act = th.tanh(self.pre_layer(dynlat_in)) 89 | 90 | # Feed the preprocessed data through the lstm 91 | lstm_h, lstm_c = self.clstm(input_tensor=pre_act, 92 | cur_state=[lstm_h, lstm_c]) 93 | 94 | # Pass the lstm output through the postprocessing layer 95 | post_act = self.post_layer(lstm_h) 96 | 97 | # Split the post activation into dynamic and latent lateral outputs 98 | dyn_out = post_act[:, :self.dyn_channels] 99 | 100 | # Lateral output 101 | lat_out = th.tanh(post_act[:, self.dyn_channels:]) 102 | 103 | # State update 104 | self.l = lat_out 105 | self.h = lstm_h 106 | self.c = lstm_c 107 | 108 | return dyn_out 109 | 110 | def reset(self, batch_size): 111 | #s = th.zeros(batch_size, *self.l.shape[1:], device=self.device) 112 | self.l = th.zeros(batch_size, *self.l.shape[1:], device=self.device) 113 | self.h = th.zeros(batch_size, *self.h.shape[1:], device=self.device) 114 | self.c = th.zeros(batch_size, *self.c.shape[1:], device=self.device) 115 | 116 | 117 | class DISTANA(nn.Module): 118 | """ 119 | DISTANA class to model the 1D Burgers equation, the two 1D diffusion 120 | sorption equations or the 2d diffusion reaction equations. 121 | """ 122 | def __init__(self, config, device): 123 | """ 124 | Constructor 125 | """ 126 | super(DISTANA, self).__init__() 127 | 128 | self.layers = th.nn.ModuleList() 129 | 130 | for ch_idx in range(len(config.model.dynamic_channels)): 131 | layer = DISTANACell( 132 | dyn_channels=config.model.dynamic_channels[ch_idx], 133 | lat_channels=config.model.lateral_channels[ch_idx], 134 | hidden_channels=config.model.hidden_channels[ch_idx], 135 | kernel_size=config.model.kernel_size, 136 | batch_size=config.training.batch_size, 137 | data_size=config.model.field_size, 138 | bias=True, 139 | device=device 140 | ) 141 | self.layers.append(layer) 142 | 143 | def forward(self, input_tensor, cur_state_list=None): 144 | """ 145 | Forward pass of the ConvLSTM_Burger class 146 | """ 147 | 148 | next_state_list = [] 149 | 150 | for layer_idx, layer in enumerate(self.layers): 151 | dyn_out = layer.forward( 152 | dyn_in=input_tensor, 153 | state=None if cur_state_list is None else cur_state_list[layer_idx] 154 | ) 155 | 156 | next_state_list.append([layer.l, layer.h, layer.c]) 157 | input_tensor = dyn_out 158 | 159 | self.state_list = next_state_list 160 | 161 | return dyn_out, next_state_list 162 | 163 | def reset(self, batch_size): 164 | """ 165 | Set all states back to zero using a given batch_size 166 | """ 167 | for layer in self.layers: 168 | layer.reset(batch_size=batch_size) -------------------------------------------------------------------------------- /models/distana/experiment.py: -------------------------------------------------------------------------------- 1 | """ 2 | This script either trains n models or tests them as determined in the 3 | config.json file. 4 | """ 5 | 6 | import multiprocessing as mp 7 | import numpy as np 8 | import train 9 | import test 10 | 11 | # 12 | # Parameters 13 | 14 | TRAIN = False 15 | MODEL_NUMBERS = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] 16 | 17 | # Number of models that can be trained in parallel 18 | POOLSIZE = 5 19 | 20 | 21 | # 22 | # SCRIPT 23 | 24 | if TRAIN: 25 | 26 | # Train the MODEL_NUMBERS using POOLSIZE threads in parallel 27 | if POOLSIZE > 1: 28 | arguments = list() 29 | for model_number in MODEL_NUMBERS: 30 | arguments.append([False, model_number]) 31 | 32 | with mp.Pool(POOLSIZE) as pool: 33 | pool.starmap(train.run_training, arguments) 34 | else: 35 | for model_number in MODEL_NUMBERS: 36 | train.run_training(print_progress=False, model_number=model_number) 37 | 38 | else: 39 | # Test the MODEL_NUMBERS iteratively 40 | error_list = list() 41 | for model_number in MODEL_NUMBERS: 42 | mse = test.run_testing(model_number=model_number) 43 | error_list.append(mse) 44 | 45 | print(error_list) 46 | print(f"Average MSE: {np.mean(error_list)}, STD: {np.std(error_list)}") 47 | -------------------------------------------------------------------------------- /models/distana/test.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch as th 3 | import time 4 | import glob 5 | import os 6 | import matplotlib.pyplot as plt 7 | import matplotlib.animation as animation 8 | import sys 9 | 10 | sys.path.append("..") 11 | from utils.configuration import Configuration 12 | import utils.helper_functions as helpers 13 | from distana import DISTANA 14 | 15 | 16 | def run_testing(print_progress=False, visualize=False, model_number=None): 17 | 18 | th.set_num_threads(1) 19 | 20 | # Load the user configurations 21 | config = Configuration("config.json") 22 | 23 | # Append the model number to the name of the model 24 | if model_number is None: 25 | model_number = config.model.number 26 | config.model.name = config.model.name + "_" + str(model_number).zfill(2) 27 | 28 | # Hide the GPU(s) in case the user specified to use the CPU in the config 29 | # file 30 | if config.general.device == "CPU": 31 | os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" 32 | os.environ["CUDA_VISIBLE_DEVICES"] = "" 33 | 34 | # Set device on GPU if specified in the configuration file, else CPU 35 | device = helpers.determine_device() 36 | 37 | # Initialize and set up the network 38 | model = DISTANA(config=config, device=device).to(device=device) 39 | 40 | if print_progress: 41 | # Count number of trainable parameters 42 | pytorch_total_params = sum( 43 | p.numel() for p in model.parameters() if p.requires_grad 44 | ) 45 | print("Trainable model parameters:", pytorch_total_params) 46 | 47 | # Restore the network by loading the weights saved in the .pt file 48 | print("Restoring model (that is the network\"s weights) from file...") 49 | 50 | model.load_state_dict(th.load(os.path.join(os.path.abspath(""), 51 | "checkpoints", 52 | config.model.name, 53 | config.model.name + ".pt"))) 54 | model.eval() 55 | 56 | """ 57 | TESTING 58 | """ 59 | 60 | # 61 | # Load the data depending on the task 62 | if config.data.type == "burger": 63 | data_path = os.path.join("../../data/", 64 | config.data.type, 65 | config.data.name, 66 | "sample.npy") 67 | data = np.array(np.load(data_path), dtype=np.float32) 68 | data = np.expand_dims(data, axis=1) 69 | 70 | elif config.data.type == "diffusion_sorption": 71 | data_path_base = os.path.join("../../data/", 72 | config.data.type, 73 | config.data.name) 74 | data_path_c = os.path.join(data_path_base, "sample_c.npy") 75 | data_path_ct = os.path.join(data_path_base, "sample_ct.npy") 76 | data_c = np.array(np.load(data_path_c), dtype=np.float32) 77 | data_ct = np.array(np.load(data_path_ct), dtype=np.float32) 78 | data = np.stack((data_c, data_ct), axis=1) 79 | 80 | elif config.data.type == "diffusion_reaction": 81 | data_path_base = os.path.join("../../data/", 82 | config.data.type, 83 | config.data.name) 84 | data_path_u = os.path.join(data_path_base, "sample_u.npy") 85 | data_path_v = os.path.join(data_path_base, "sample_v.npy") 86 | data_u = np.array(np.load(data_path_u), dtype=np.float32) 87 | data_v = np.array(np.load(data_path_v), dtype=np.float32) 88 | data = np.stack((data_u, data_v), axis=1) 89 | 90 | elif config.data.type == "allen_cahn": 91 | data_path = os.path.join("../../data/", 92 | config.data.type, 93 | config.data.name, 94 | "sample.npy") 95 | data = np.array(np.load(data_path), dtype=np.float32) 96 | data = np.expand_dims(data, axis=1) 97 | 98 | elif config.data.type == "burger_2d": 99 | data_path = os.path.join("../../data/", 100 | config.data.type, 101 | config.data.name, 102 | "sample.npy") 103 | data = np.array(np.load(data_path), dtype=np.float32) 104 | data = np.expand_dims(data, axis=1) 105 | 106 | # Set up the training and validation datasets and -loaders 107 | data_test = th.tensor(data, 108 | device=device).unsqueeze(1) 109 | sequence_length = len(data_test) - 1 110 | 111 | # Evaluate the network for the given test data 112 | 113 | # Separate the data into network inputs and labels 114 | net_inputs = th.clone(data_test[:-1]) 115 | net_labels = th.clone(data_test[1:]) 116 | 117 | # Set up an array of zeros to store the network outputs 118 | net_outputs = th.zeros(size=(sequence_length, 119 | config.testing.batch_size, 120 | config.model.dynamic_channels[-1], 121 | *config.model.field_size), 122 | device=device) 123 | state_list = None 124 | 125 | # Iterate over the remaining sequence of the training example and perform a 126 | # forward pass 127 | time_start = time.time() 128 | for t in range(len(net_inputs)): 129 | 130 | if t < config.testing.teacher_forcing_steps: 131 | # Teacher forcing 132 | net_input = net_inputs[t] 133 | else: 134 | # Closed loop 135 | net_input = net_outputs[t - 1] 136 | 137 | # Feed the boundary data also in closed loop if desired 138 | if config.testing.feed_boundary_data: 139 | net_input[:, :, 0] = net_inputs[t, :, :, 0] 140 | net_input[:, :, -1] = net_inputs[t, :, :, -1] 141 | 142 | net_output, state_list = model.forward(input_tensor=net_input, 143 | cur_state_list=state_list) 144 | net_outputs[t] = net_output[-1] 145 | 146 | if print_progress: 147 | forward_pass_duration = time.time() - time_start 148 | print("Forward pass took:", forward_pass_duration, "seconds.") 149 | 150 | # Convert the PyTorch network output tensor into a numpy array 151 | net_outputs = net_outputs.cpu().detach().numpy()[:, 0, 0] 152 | net_labels = net_labels.cpu().detach().numpy()[:, 0, 0] 153 | 154 | # 155 | # Visualize the data 156 | if visualize: 157 | plt.style.use("dark_background") 158 | 159 | # Plot over space and time 160 | fig, ax = plt.subplots(1, 2, figsize=(16, 6), sharey=True) 161 | 162 | if config.data.type == "burger" or\ 163 | config.data.type == "diffusion_sorption" or\ 164 | config.data.type == "allen_cahn": 165 | 166 | im1 = ax[0].imshow( 167 | np.transpose(net_labels), interpolation='nearest', 168 | cmap='rainbow', origin='lower', aspect='auto', vmin=-0.4, 169 | vmax=0.4 170 | ) 171 | fig.colorbar(im1, ax=ax[0]) 172 | im2 = ax[1].imshow( 173 | np.transpose(net_outputs), interpolation='nearest', 174 | cmap='rainbow', origin='lower', aspect='auto', vmin=-0.4, 175 | vmax=0.4 176 | ) 177 | fig.colorbar(im2, ax=ax[1]) 178 | 179 | ax[0].set_xlabel("t") 180 | ax[0].set_ylabel("x") 181 | ax[1].set_xlabel("t") 182 | 183 | elif config.data.type == "diffusion_reaction": 184 | 185 | im1 = ax[0].imshow( 186 | np.transpose(net_labels[..., 0]), interpolation='nearest', 187 | cmap='rainbow', origin='lower', aspect='auto', vmin=-0.4, 188 | vmax=0.4 189 | ) 190 | fig.colorbar(im1, ax=ax[0]) 191 | im2 = ax[1].imshow( 192 | np.transpose(net_outputs[..., 0]), interpolation='nearest', 193 | cmap='rainbow', origin='lower', aspect='auto', vmin=-0.4, 194 | vmax=0.4 195 | ) 196 | fig.colorbar(im2, ax=ax[1]) 197 | 198 | ax[0].set_xlabel("x") 199 | ax[0].set_ylabel("y") 200 | ax[1].set_xlabel("x") 201 | 202 | elif config.data.type == "burger_2d": 203 | 204 | im1 = ax[0].imshow( 205 | np.transpose(net_labels[-1, ...]), interpolation='nearest', 206 | cmap='rainbow', origin='lower', aspect='auto', vmin=-0.4, 207 | vmax=0.4 208 | ) 209 | fig.colorbar(im1, ax=ax[0]) 210 | im2 = ax[1].imshow( 211 | np.transpose(net_outputs[-1, ...]), interpolation='nearest', 212 | cmap='rainbow', origin='lower', aspect='auto', vmin=-0.4, 213 | vmax=0.4 214 | ) 215 | fig.colorbar(im2, ax=ax[1]) 216 | 217 | ax[0].set_xlabel("x") 218 | ax[0].set_ylabel("y") 219 | ax[1].set_xlabel("x") 220 | 221 | 222 | ax[0].set_title("Ground Truth") 223 | ax[1].set_title("Network Output") 224 | 225 | 226 | if config.data.type == "diffusion_reaction"\ 227 | or config.data.type == "burger_2d": 228 | anim = animation.FuncAnimation( 229 | fig, 230 | animate, 231 | frames=sequence_length, 232 | fargs=(im1, im2, net_labels, net_outputs), 233 | interval=20 234 | ) 235 | 236 | plt.show() 237 | 238 | # 239 | # Compute and return statistics 240 | mse = np.mean(np.square(net_outputs - net_labels)) 241 | 242 | return net_outputs, net_labels 243 | 244 | 245 | def animate(t, im1, im2, net_labels, net_outputs): 246 | """ 247 | Data animation function animating an image over time. 248 | :param t: The current time step 249 | :param axis: The matplotlib image object 250 | :param field: The data field 251 | :return: The matplotlib image object updated with the current time step's 252 | image date 253 | """ 254 | im1.set_array(net_labels[t]) 255 | im2.set_array(net_outputs[t]) 256 | 257 | 258 | if __name__ == "__main__": 259 | mse = run_testing(print_progress=True, visualize=True) 260 | 261 | print(f"MSE: {mse}") 262 | exit() 263 | 264 | u_hat = [] 265 | 266 | for i in range(10): 267 | pred, u = run_testing(visualize=False, model_number=i) 268 | pred = np.expand_dims(pred,-1) 269 | u_hat.append(pred) 270 | 271 | u_hat = np.stack(u_hat) 272 | 273 | np.save("distana_test.npy", u_hat) 274 | 275 | print("Done.") -------------------------------------------------------------------------------- /models/distana/train.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch as th 3 | import torch.nn as nn 4 | import time 5 | import glob 6 | import os 7 | import matplotlib.pyplot as plt 8 | from threading import Thread 9 | import sys 10 | 11 | sys.path.append("..") 12 | from utils.configuration import Configuration 13 | import utils.helper_functions as helpers 14 | from distana import DISTANA 15 | 16 | 17 | def process_sequence(model, criterion, data, batch_size, epoch, config, 18 | device): 19 | 20 | # Separate the data into network inputs and labels 21 | net_inputs = data[:-1] 22 | net_labels = data[1:] 23 | 24 | # Set up an array of zeros to store the network outputs 25 | net_outputs = th.zeros(size=(len(net_labels), 26 | config.training.batch_size, 27 | config.model.dynamic_channels[-1], 28 | *config.model.field_size), 29 | device=device) 30 | 31 | model.reset(batch_size=batch_size) 32 | 33 | # Initial network input and forward pass 34 | net_output, state_list = model.forward(input_tensor=net_inputs[0], 35 | cur_state_list=None) 36 | net_outputs[0] = net_output 37 | 38 | # Probability of closed loop in current epoch 39 | p_cl = epoch / (config.training.epochs * 1.5) 40 | 41 | # Iterate over the whole sequence of the training example and 42 | # perform a forward pass 43 | for t in range(1, len(net_inputs)): 44 | 45 | # Draw random number to determine teacher forcing or closed loop 46 | # (scheduled sampling) 47 | p_uniform = np.random.uniform(0, 1) 48 | closed_loop = True if p_uniform < p_cl else False 49 | 50 | if not closed_loop: 51 | # Teacher forcing 52 | net_input = net_inputs[t] 53 | else: 54 | # Closed loop 55 | net_input = net_outputs[t - 1] 56 | 57 | # Forward the input through the network 58 | #net_output, _ = model.forward(input_tensor=net_inputs[t]) 59 | net_output, state_list = model.forward(input_tensor=net_input, 60 | cur_state_list=state_list) 61 | 62 | # Store the output of the network for this sequence step 63 | #net_outputs[t] = net_output 64 | net_outputs[t] = net_output[-1] 65 | 66 | # Compute the mean squared error 67 | mse = criterion(net_outputs, net_labels) 68 | 69 | return net_outputs, mse 70 | 71 | 72 | def run_training(print_progress=False, model_number=None): 73 | 74 | # Set a random seed for varying weight initializations 75 | th.seed() 76 | 77 | th.set_num_threads(1) 78 | 79 | # Load the user configurations 80 | config = Configuration("config.json") 81 | 82 | # Append the model number to the name of the model 83 | if model_number is None: 84 | model_number = config.model.number 85 | config.model.name = config.model.name + "_" + str(model_number).zfill(2) 86 | 87 | print("Model name:", config.model.name) 88 | 89 | # Hide the GPU(s) in case the user specified to use the CPU in the config 90 | # file 91 | if config.general.device == "CPU": 92 | os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" # see issue #152 93 | os.environ["CUDA_VISIBLE_DEVICES"] = "" 94 | 95 | time_start = time.time() 96 | 97 | # setting device on GPU if available, else CPU 98 | device = helpers.determine_device(print_progress=False) 99 | 100 | # Initialize and set up the network 101 | model = DISTANA(config=config, device=device).to(device=device) 102 | 103 | if print_progress: 104 | # Count number of trainable parameters 105 | pytorch_total_params = sum( 106 | p.numel() for p in model.parameters() if p.requires_grad 107 | ) 108 | print("Trainable model parameters:", pytorch_total_params) 109 | 110 | # If desired, restore the network by loading the weights saved in the .pt 111 | # file 112 | if config.training.continue_training: 113 | print("Restoring model (that is the network\"s weights) from file...") 114 | model.load_state_dict(th.load(os.path.join(os.path.abspath(""), 115 | "checkpoints", 116 | config.model.name, 117 | config.model.name + ".pt"))) 118 | model.train() 119 | 120 | # 121 | # Set up the optimizer and the criterion (loss) 122 | optimizer = th.optim.LBFGS(model.parameters(), 123 | lr=config.training.learning_rate) 124 | 125 | criterion = nn.MSELoss() 126 | 127 | # 128 | # Load the data depending on the task 129 | if config.data.type == "burger": 130 | data_path = os.path.join("../../data/", 131 | config.data.type, 132 | config.data.name, 133 | "sample.npy") 134 | data = np.array(np.load(data_path), dtype=np.float32) 135 | data = np.expand_dims(data, axis=1) 136 | 137 | elif config.data.type == "diffusion_sorption": 138 | data_path_base = os.path.join("../../data/", 139 | config.data.type, 140 | config.data.name) 141 | data_path_c = os.path.join(data_path_base, "sample_c.npy") 142 | data_path_ct = os.path.join(data_path_base, "sample_ct.npy") 143 | data_c = np.array(np.load(data_path_c), dtype=np.float32) 144 | data_ct = np.array(np.load(data_path_ct), dtype=np.float32) 145 | data = np.stack((data_c, data_ct), axis=1) 146 | 147 | elif config.data.type == "diffusion_reaction": 148 | data_path_base = os.path.join("../../data/", 149 | config.data.type, 150 | config.data.name) 151 | data_path_u = os.path.join(data_path_base, "sample_u.npy") 152 | data_path_v = os.path.join(data_path_base, "sample_v.npy") 153 | data_u = np.array(np.load(data_path_u), dtype=np.float32) 154 | data_v = np.array(np.load(data_path_v), dtype=np.float32) 155 | data = np.stack((data_u, data_v), axis=1) 156 | 157 | elif config.data.type == "allen_cahn": 158 | data_path = os.path.join("../../data/", 159 | config.data.type, 160 | config.data.name, 161 | "sample.npy") 162 | data = np.array(np.load(data_path), dtype=np.float32) 163 | data = np.expand_dims(data, axis=1) 164 | 165 | elif config.data.type == "burger_2d": 166 | data_path = os.path.join("../../data/", 167 | config.data.type, 168 | config.data.name, 169 | "sample.npy") 170 | data = np.array(np.load(data_path), dtype=np.float32) 171 | data = np.expand_dims(data, axis=1) 172 | 173 | # 174 | # Set up the training and validation datasets and -loaders 175 | data_train = th.tensor( 176 | data[:config.training.t_stop], 177 | device=device 178 | ).unsqueeze(1) 179 | data_valid = th.tensor( 180 | data[config.validation.t_start:config.validation.t_stop], 181 | device=device 182 | ).unsqueeze(1) 183 | 184 | # 185 | # Set up lists to save and store the epoch errors 186 | epoch_errors_train = [] 187 | epoch_errors_valid = [] 188 | best_train = np.infty 189 | best_valid = np.infty 190 | 191 | """ 192 | TRAINING 193 | """ 194 | 195 | a = time.time() 196 | 197 | # 198 | # Start the training and iterate over all epochs 199 | for epoch in range(config.training.epochs): 200 | 201 | epoch_start_time = time.time() 202 | 203 | # 204 | # Train the network on the given data 205 | 206 | # Define the closure function that consists of resetting the 207 | # gradient buffer, loss function calculation, and backpropagation 208 | # It is necessary for LBFGS optimizer, because it requires multiple 209 | # function evaluations 210 | def closure(): 211 | # Set the model to train mode 212 | model.train() 213 | 214 | # Reset the optimizer to clear data from previous iterations 215 | optimizer.zero_grad() 216 | 217 | _, mse = process_sequence( 218 | model=model, 219 | criterion=criterion, 220 | data=data_train, 221 | batch_size=config.training.batch_size, 222 | epoch=epoch, 223 | config=config, 224 | device=device 225 | ) 226 | 227 | mse.backward() 228 | 229 | return mse 230 | 231 | optimizer.step(closure) 232 | 233 | # Extract the MSE value from the closure function 234 | mse = closure() 235 | # 236 | 237 | # Append the error to the error list 238 | epoch_errors_train.append(mse.item()) 239 | 240 | train_sign = "(-)" 241 | if epoch_errors_train[-1] < best_train: 242 | best_train = epoch_errors_train[-1] 243 | train_sign = "(+)" 244 | 245 | # 246 | # Validate the network 247 | 248 | net_output, mse = process_sequence( 249 | model=model, 250 | criterion=criterion, 251 | data=data_valid, 252 | batch_size=config.training.batch_size, 253 | epoch=epoch, 254 | config=config, 255 | device=device 256 | ) 257 | epoch_errors_valid.append(mse.item()) 258 | 259 | # Create a plus or minus sign for the validation error 260 | valid_sign = "(-)" 261 | if epoch_errors_valid[-1] < best_valid: 262 | best_valid = epoch_errors_valid[-1] 263 | valid_sign = "(+)" 264 | 265 | if config.training.save_model: 266 | # Start a separate thread to save the model 267 | thread = Thread(target=helpers.save_model_to_file( 268 | model_src_path=os.path.abspath(""), 269 | config=config, 270 | epoch=epoch, 271 | epoch_errors_train=epoch_errors_train, 272 | epoch_errors_valid=epoch_errors_valid, 273 | net=model)) 274 | thread.start() 275 | 276 | # 277 | # Print progress to the console 278 | if print_progress: 279 | print(f"Epoch {str(epoch+1).zfill(int(np.log10(config.training.epochs))+1)}/{str(config.training.epochs)} took {str(np.round(time.time() - epoch_start_time, 2)).ljust(5, '0')} seconds. \t\tAverage epoch training error: {train_sign}{str(np.round(epoch_errors_train[-1], 10)).ljust(12, ' ')} \t\tValidation error: {valid_sign}{str(np.round(epoch_errors_valid[-1], 10)).ljust(12, ' ')}") 280 | 281 | b = time.time() 282 | if print_progress: 283 | print("\nTraining took " + str(np.round(b - a, 2)) + " seconds.\n\n") 284 | 285 | 286 | if __name__ == "__main__": 287 | run_training(print_progress=True) 288 | 289 | print("Done.") -------------------------------------------------------------------------------- /models/finn/config.json: -------------------------------------------------------------------------------- 1 | { 2 | "paths": { 3 | }, 4 | 5 | "general": { 6 | "device": "cpu" // "cpu" or "cuda" 7 | }, 8 | 9 | "training": { 10 | "save_model": true, 11 | "continue_training": false, 12 | "epochs": 100, 13 | "learning_rate": 0.1 14 | }, 15 | 16 | "data": { 17 | "type": "burger_2d", // "burger", "burger_2d" "diffusion_sorption", "diffusion_reaction", "allen_cahn" 18 | "name": "data_train", // "data_train", "data_ext", "data_test" 19 | "noise": 0.0 20 | }, 21 | 22 | "model": { 23 | "name": "burger_2d", // "burger", "burger_2d", "diff-sorp", "diff-react", "allen-cahn" 24 | "number": 0, // The i-th model 25 | "layer_sizes": [1, 10, 20, 10, 1] // [1, 10, 20, 10, 1] for burger, burger_2d, diffusion_sorption, and allen_cahn, [2, 20, 20, 20, 2] for diffusion_reaction 26 | } 27 | } -------------------------------------------------------------------------------- /models/finn/experiment.py: -------------------------------------------------------------------------------- 1 | """ 2 | This script either trains n models or tests them as determined in the 3 | config.json file. 4 | """ 5 | 6 | import multiprocessing as mp 7 | import numpy as np 8 | import train 9 | import test 10 | 11 | # 12 | # Parameters 13 | 14 | TRAIN = False 15 | MODEL_NUMBERS = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] 16 | 17 | # Number of models that can be trained in parallel 18 | POOLSIZE = 5 19 | 20 | 21 | # 22 | # SCRIPT 23 | 24 | if TRAIN: 25 | 26 | # Train the MODEL_NUMBERS using POOLSIZE threads in parallel 27 | if POOLSIZE > 1: 28 | arguments = list() 29 | for model_number in MODEL_NUMBERS: 30 | arguments.append([False, model_number]) 31 | 32 | with mp.Pool(POOLSIZE) as pool: 33 | pool.starmap(train.run_training, arguments) 34 | else: 35 | for model_number in MODEL_NUMBERS: 36 | train.run_training(print_progress=False, model_number=model_number) 37 | 38 | else: 39 | # Test the MODEL_NUMBERS iteratively 40 | error_list = list() 41 | for model_number in MODEL_NUMBERS: 42 | mse = test.run_testing(model_number=model_number) 43 | error_list.append(mse) 44 | 45 | print(error_list) 46 | print(f"Average MSE: {np.mean(error_list)}, STD: {np.std(error_list)}") 47 | -------------------------------------------------------------------------------- /models/finn_poly/config.json: -------------------------------------------------------------------------------- 1 | { 2 | "paths": { 3 | }, 4 | 5 | "general": { 6 | "device": "cuda" 7 | }, 8 | 9 | "training": { 10 | "save_model": true, 11 | "continue_training": false, 12 | "epochs": 100, 13 | "learning_rate": 0.1 14 | }, 15 | 16 | "data": { 17 | "type": "diffusion_reaction", // "burger", "diffusion_sorption", "diffusion_reaction", "allen_cahn" 18 | "name": "data_train", // "data_train", "data_ext", "data_test" 19 | "noise": 0.0 20 | }, 21 | 22 | "model": { 23 | "name": "diff-react", // "burger", "diff-sorp", "diff-react", "allen-cahn" 24 | "number": 0 // The i-th model 25 | } 26 | } -------------------------------------------------------------------------------- /models/finn_poly/train.py: -------------------------------------------------------------------------------- 1 | #! env/bin/python3 2 | 3 | """ 4 | Main file for training a model with FINN 5 | """ 6 | 7 | import numpy as np 8 | import torch as th 9 | import torch.nn as nn 10 | import os 11 | import time 12 | from threading import Thread 13 | import sys 14 | 15 | sys.path.append("..") 16 | from utils.configuration import Configuration 17 | import utils.helper_functions as helpers 18 | from finn import FINN_Burger, FINN_DiffSorp, FINN_DiffReact, FINN_AllenCahn 19 | 20 | _author_ = "Matthias Karlbauer, Timothy Praditia" 21 | 22 | 23 | def run_training(model_number=None): 24 | 25 | # Load the user configurations 26 | config = Configuration("config.json") 27 | 28 | # Append the model number to the name of the model 29 | if model_number is None: 30 | model_number = config.model.number 31 | config.model.name = config.model.name + "_" + str(model_number).zfill(2) 32 | 33 | # Print some information to console 34 | print("Model name:", config.model.name) 35 | 36 | # Hide the GPU(s) in case the user specified to use the CPU in the config 37 | # file 38 | if config.general.device == "CPU": 39 | os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" 40 | os.environ["CUDA_VISIBLE_DEVICES"] = "" 41 | 42 | root_path = os.path.abspath("../../data") 43 | data_path = os.path.join(root_path, config.data.type, config.data.name) 44 | 45 | # Set device on GPU if specified in the configuration file, else CPU 46 | # device = helpers.determine_device() 47 | device = th.device(config.general.device) 48 | 49 | if config.data.type == "burger": 50 | # Load samples, together with x, y, and t series 51 | t = th.tensor(np.load(os.path.join(data_path, "t_series.npy")), 52 | dtype=th.float).to(device=device) 53 | x = np.load(os.path.join(data_path, "x_series.npy")) 54 | u = th.tensor(np.load(os.path.join(data_path, "sample.npy")), 55 | dtype=th.float).to(device=device) 56 | 57 | u[1:] = u[1:] + th.normal(th.zeros_like(u[1:]),th.ones_like(u[1:])*config.data.noise) 58 | 59 | dx = x[1]-x[0] 60 | 61 | # Initialize and set up the model 62 | model = FINN_Burger( 63 | u = u, 64 | D = np.array([1.0]), 65 | BC = np.array([[0.0], [0.0]]), 66 | dx = dx, 67 | device = device, 68 | mode="train", 69 | learn_coeff=True 70 | ).to(device=device) 71 | 72 | 73 | elif config.data.type == "diffusion_sorption": 74 | # Load samples, together with x, y, and t series 75 | t = th.tensor(np.load(os.path.join(data_path, "t_series.npy")), 76 | dtype=th.float).to(device=device) 77 | x = np.load(os.path.join(data_path, "x_series.npy")) 78 | sample_c = th.tensor(np.load(os.path.join(data_path, "sample_c.npy")), 79 | dtype=th.float).to(device=device) 80 | sample_ct = th.tensor(np.load(os.path.join(data_path, "sample_ct.npy")), 81 | dtype=th.float).to(device=device) 82 | 83 | dx = x[1]-x[0] 84 | u = th.stack((sample_c, sample_ct), dim=len(sample_c.shape)) 85 | 86 | u[1:] = u[1:] + th.normal(th.zeros_like(u[1:]),th.ones_like(u[1:])*config.data.noise) 87 | 88 | # Initialize and set up the model 89 | model = FINN_DiffSorp( 90 | u = u, 91 | D = np.array([0.5, 0.1]), 92 | BC = np.array([[1.0, 1.0], [0.0, 0.0]]), 93 | dx = dx, 94 | device = device, 95 | mode="train", 96 | learn_coeff=True 97 | ).to(device=device) 98 | 99 | elif config.data.type == "diffusion_reaction": 100 | # Load samples, together with x, y, and t series 101 | t = th.tensor(np.load(os.path.join(data_path, "t_series.npy")), 102 | dtype=th.float).to(device=device) 103 | x = np.load(os.path.join(data_path, "x_series.npy")) 104 | y = np.load(os.path.join(data_path, "y_series.npy")) 105 | sample_u = th.tensor(np.load(os.path.join(data_path, "sample_u.npy")), 106 | dtype=th.float).to(device=device) 107 | sample_v = th.tensor(np.load(os.path.join(data_path, "sample_v.npy")), 108 | dtype=th.float).to(device=device) 109 | 110 | dx = x[1]-x[0] 111 | dy = y[1]-y[0] 112 | 113 | u = th.stack((sample_u, sample_v), dim=len(sample_u.shape)) 114 | u[1:] = u[1:] + th.normal(th.zeros_like(u[1:]),th.ones_like(u[1:])*config.data.noise) 115 | 116 | # Initialize and set up the model 117 | model = FINN_DiffReact( 118 | u = u, 119 | D = np.array([5E-4/(dx**2), 1E-3/(dx**2)]), 120 | BC = np.zeros((4,2)), 121 | dx = dx, 122 | dy = dy, 123 | device = device, 124 | mode="train", 125 | learn_coeff=True 126 | ).to(device=device) 127 | 128 | elif config.data.type == "allen_cahn": 129 | # Load samples, together with x, y, and t series 130 | t = th.tensor(np.load(os.path.join(data_path, "t_series.npy")), 131 | dtype=th.float).to(device=device) 132 | x = np.load(os.path.join(data_path, "x_series.npy")) 133 | u = th.tensor(np.load(os.path.join(data_path, "sample.npy")), 134 | dtype=th.float).to(device=device) 135 | 136 | u[1:] = u[1:] + th.normal(th.zeros_like(u[1:]),th.ones_like(u[1:])*config.data.noise) 137 | 138 | dx = x[1]-x[0] 139 | 140 | # Initialize and set up the model 141 | model = FINN_AllenCahn( 142 | u = u, 143 | D = np.array([0.6]), 144 | BC = np.array([[0.0], [0.0]]), 145 | dx = dx, 146 | device = device, 147 | mode="train", 148 | learn_coeff=True 149 | ).to(device=device) 150 | 151 | # Count number of trainable parameters 152 | pytorch_total_params = sum( 153 | p.numel() for p in model.parameters() if p.requires_grad 154 | ) 155 | 156 | print("Trainable model parameters:", pytorch_total_params) 157 | 158 | # If desired, restore the network by loading the weights saved in the .pt 159 | # file 160 | if config.training.continue_training: 161 | print('Restoring model (that is the network\'s weights) from file...') 162 | model.load_state_dict(th.load(os.path.join(os.path.abspath(""), 163 | "checkpoints", 164 | config.model.name, 165 | config.model.name + ".pt"))) 166 | model.train() 167 | 168 | # 169 | # Set up an optimizer and the criterion (loss) 170 | optimizer = th.optim.LBFGS(model.parameters(), 171 | lr=config.training.learning_rate) 172 | 173 | criterion = nn.MSELoss(reduction="mean") 174 | 175 | # 176 | # Set up lists to save and store the epoch errors 177 | epoch_errors_train = [] 178 | best_train = np.infty 179 | 180 | """ 181 | TRAINING 182 | """ 183 | 184 | a = time.time() 185 | 186 | # 187 | # Start the training and iterate over all epochs 188 | for epoch in range(config.training.epochs): 189 | 190 | epoch_start_time = time.time() 191 | 192 | # Define the closure function that consists of resetting the 193 | # gradient buffer, loss function calculation, and backpropagation 194 | # It is necessary for LBFGS optimizer, because it requires multiple 195 | # function evaluations 196 | def closure(): 197 | # Set the model to train mode 198 | model.train() 199 | 200 | # Reset the optimizer to clear data from previous iterations 201 | optimizer.zero_grad() 202 | 203 | # Forward propagate and calculate loss function 204 | u_hat = model(t=t, u=u) 205 | 206 | mse = criterion(u_hat, u) 207 | 208 | mse.backward() 209 | 210 | print(mse.item()) 211 | print(model.D) 212 | 213 | return mse 214 | 215 | optimizer.step(closure) 216 | 217 | # Extract the MSE value from the closure function 218 | mse = closure() 219 | 220 | epoch_errors_train.append(mse.item()) 221 | 222 | # Create a plus or minus sign for the training error 223 | train_sign = "(-)" 224 | if epoch_errors_train[-1] < best_train: 225 | train_sign = "(+)" 226 | best_train = epoch_errors_train[-1] 227 | # Save the model to file (if desired) 228 | if config.training.save_model: 229 | # Start a separate thread to save the model 230 | thread = Thread(target=helpers.save_model_to_file( 231 | model_src_path=os.path.abspath(""), 232 | config=config, 233 | epoch=epoch, 234 | epoch_errors_train=epoch_errors_train, 235 | epoch_errors_valid=epoch_errors_train, 236 | net=model)) 237 | thread.start() 238 | 239 | 240 | 241 | # 242 | # Print progress to the console 243 | print(f"Epoch {str(epoch+1).zfill(int(np.log10(config.training.epochs))+1)}/{str(config.training.epochs)} took {str(np.round(time.time() - epoch_start_time, 2)).ljust(5, '0')} seconds. \t\tAverage epoch training error: {train_sign}{str(np.round(epoch_errors_train[-1], 10)).ljust(12, ' ')}") 244 | 245 | b = time.time() 246 | print('\nTraining took ' + str(np.round(b - a, 2)) + ' seconds.\n\n') 247 | 248 | 249 | if __name__ == "__main__": 250 | th.set_num_threads(1) 251 | run_training() 252 | 253 | print("Done.") -------------------------------------------------------------------------------- /models/fno/config.json: -------------------------------------------------------------------------------- 1 | { 2 | "paths": { 3 | }, 4 | 5 | "general": { 6 | "device": "cpu" // "cpu" or "cuda" 7 | }, 8 | 9 | "training": { 10 | "save_model": true, 11 | "continue_training": false, 12 | "epochs": 2000, 13 | "learning_rate": 0.001 14 | }, 15 | 16 | "data": { 17 | "type": "burger_2d", // "burger", "burger_2d", "diffusion_sorption", "diffusion_reaction", "allen_cahn" 18 | "name": "data_train", // "data_train", "data_ext", "data_test" 19 | "noise": 0.0 20 | }, 21 | 22 | "model": { 23 | "name": "burger_2d", // "burger", "burger_2d", "diff-sorp", "diff-react", "allen-cahn" 24 | "number": 0 // The i-th model 25 | } 26 | } -------------------------------------------------------------------------------- /models/fno/experiment.py: -------------------------------------------------------------------------------- 1 | """ 2 | This script either trains n models or tests them as determined in the 3 | config.json file. 4 | """ 5 | 6 | import multiprocessing as mp 7 | import numpy as np 8 | import train 9 | import test 10 | 11 | # 12 | # Parameters 13 | 14 | TRAIN = False 15 | MODEL_NUMBERS = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] 16 | 17 | # Number of models that can be trained in parallel 18 | POOLSIZE = 1 19 | 20 | 21 | # 22 | # SCRIPT 23 | 24 | if TRAIN: 25 | 26 | # Train the MODEL_NUMBERS using POOLSIZE threads in parallel 27 | if POOLSIZE > 1: 28 | arguments = list() 29 | for model_number in MODEL_NUMBERS: 30 | arguments.append([False, model_number]) 31 | 32 | with mp.Pool(POOLSIZE) as pool: 33 | pool.starmap(train.run_training, arguments) 34 | else: 35 | for model_number in MODEL_NUMBERS: 36 | train.run_training(print_progress=False, model_number=model_number) 37 | 38 | else: 39 | # Test the MODEL_NUMBERS iteratively 40 | error_list = list() 41 | for model_number in MODEL_NUMBERS: 42 | mse = test.run_testing(model_number=model_number) 43 | error_list.append(mse) 44 | 45 | print(error_list) 46 | print(f"Average MSE: {np.mean(error_list)}, STD: {np.std(error_list)}") 47 | -------------------------------------------------------------------------------- /models/fno/fno.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | 3 | """ 4 | Physics-derivatives learning layer. Implementation taken and modified from 5 | https://github.com/vincent-leguen/PhyDNet 6 | """ 7 | 8 | import torch 9 | import torch.nn as nn 10 | import numpy as np 11 | import torch.nn.functional as F 12 | 13 | 14 | class SpectralConv1d(nn.Module): 15 | def __init__(self, in_channels, out_channels, modes1): 16 | super(SpectralConv1d, self).__init__() 17 | 18 | """ 19 | 1D Fourier layer. It does FFT, linear transform, and Inverse FFT. 20 | """ 21 | 22 | self.in_channels = in_channels 23 | self.out_channels = out_channels 24 | self.modes1 = modes1 #Number of Fourier modes to multiply, at most floor(N/2) + 1 25 | 26 | self.scale = (1 / (in_channels*out_channels)) 27 | self.weights1 = nn.Parameter(self.scale * torch.rand(in_channels, out_channels, self.modes1, dtype=torch.cfloat)) 28 | 29 | # Complex multiplication 30 | def compl_mul1d(self, input, weights): 31 | # (batch, in_channel, x ), (in_channel, out_channel, x) -> (batch, out_channel, x) 32 | return torch.einsum("bix,iox->box", input, weights) 33 | 34 | def forward(self, x): 35 | batchsize = x.shape[0] 36 | #Compute Fourier coeffcients up to factor of e^(- something constant) 37 | x_ft = torch.fft.rfft(x) 38 | 39 | # Multiply relevant Fourier modes 40 | out_ft = torch.zeros(batchsize, self.out_channels, x.size(-1)//2 + 1, device=x.device, dtype=torch.cfloat) 41 | out_ft[:, :, :self.modes1] = self.compl_mul1d(x_ft[:, :, :self.modes1], self.weights1) 42 | 43 | #Return to physical space 44 | x = torch.fft.irfft(out_ft, n=x.size(-1)) 45 | return x 46 | 47 | class FNO1d(nn.Module): 48 | def __init__(self, num_channels, bc, modes=12, width=10): 49 | super(FNO1d, self).__init__() 50 | 51 | """ 52 | The overall network. It contains 4 layers of the Fourier layer. 53 | 1. Lift the input to the desire channel dimension by self.fc0 . 54 | 2. 4 layers of the integral operators u' = (W + K)(u). 55 | W defined by self.w; K defined by self.conv . 56 | 3. Project from the channel space to the output space by self.fc1 and self.fc2 . 57 | 58 | input: the solution of the initial condition and location (a(x), x) 59 | input shape: (batchsize, x=s, c=2) 60 | output: the solution of a later timestep 61 | output shape: (batchsize, x=s, c=1) 62 | """ 63 | 64 | self.modes1 = modes 65 | self.width = width 66 | self.padding = 2 # pad the domain if input is non-periodic 67 | self.fc0 = nn.Linear(10*num_channels+1, self.width) # input channel is 2: (a(x), x) 68 | 69 | self.conv0 = SpectralConv1d(self.width, self.width, self.modes1) 70 | self.conv1 = SpectralConv1d(self.width, self.width, self.modes1) 71 | self.conv2 = SpectralConv1d(self.width, self.width, self.modes1) 72 | self.conv3 = SpectralConv1d(self.width, self.width, self.modes1) 73 | self.w0 = nn.Conv1d(self.width, self.width, 1) 74 | self.w1 = nn.Conv1d(self.width, self.width, 1) 75 | self.w2 = nn.Conv1d(self.width, self.width, 1) 76 | self.w3 = nn.Conv1d(self.width, self.width, 1) 77 | 78 | self.fc1 = nn.Linear(self.width, 32) 79 | self.fc2 = nn.Linear(32, num_channels) 80 | 81 | self.bc = bc 82 | 83 | def forward(self, x, grid): 84 | x = torch.cat((x, grid), dim=-1) 85 | x = self.fc0(x) 86 | x = x.permute(0, 2, 1) 87 | 88 | mult = x.size(1) // self.bc.size(1) 89 | x = torch.cat((self.bc[:1,:,:1].repeat(x.size(0),mult,1),x,self.bc[:1,:,1:].repeat(x.size(0),mult,1)),dim=2) 90 | # x = F.pad(x, [0,self.padding]) # pad the domain if input is non-periodic 91 | 92 | x1 = self.conv0(x) 93 | x2 = self.w0(x) 94 | x = x1 + x2 95 | x = F.gelu(x) 96 | 97 | x1 = self.conv1(x) 98 | x2 = self.w1(x) 99 | x = x1 + x2 100 | x = F.gelu(x) 101 | 102 | x1 = self.conv2(x) 103 | x2 = self.w2(x) 104 | x = x1 + x2 105 | x = F.gelu(x) 106 | 107 | x1 = self.conv3(x) 108 | x2 = self.w3(x) 109 | x = x1 + x2 110 | 111 | # x = x[..., :-self.padding] # pad the domain if input is non-periodic 112 | x = x[..., 1:-1] 113 | x = x.permute(0, 2, 1) 114 | x = self.fc1(x) 115 | x = F.gelu(x) 116 | x = self.fc2(x) 117 | return x.unsqueeze(-2) 118 | 119 | 120 | class SpectralConv2d_fast(nn.Module): 121 | def __init__(self, in_channels, out_channels, modes1, modes2): 122 | super(SpectralConv2d_fast, self).__init__() 123 | 124 | """ 125 | 2D Fourier layer. It does FFT, linear transform, and Inverse FFT. 126 | """ 127 | 128 | self.in_channels = in_channels 129 | self.out_channels = out_channels 130 | self.modes1 = modes1 #Number of Fourier modes to multiply, at most floor(N/2) + 1 131 | self.modes2 = modes2 132 | 133 | self.scale = (1 / (in_channels * out_channels)) 134 | self.weights1 = nn.Parameter(self.scale * torch.rand(in_channels, out_channels, self.modes1, self.modes2, dtype=torch.cfloat)) 135 | self.weights2 = nn.Parameter(self.scale * torch.rand(in_channels, out_channels, self.modes1, self.modes2, dtype=torch.cfloat)) 136 | 137 | # Complex multiplication 138 | def compl_mul2d(self, input, weights): 139 | # (batch, in_channel, x,y ), (in_channel, out_channel, x,y) -> (batch, out_channel, x,y) 140 | return torch.einsum("bixy,ioxy->boxy", input, weights) 141 | 142 | def forward(self, x): 143 | batchsize = x.shape[0] 144 | #Compute Fourier coeffcients up to factor of e^(- something constant) 145 | x_ft = torch.fft.rfft2(x) 146 | 147 | # Multiply relevant Fourier modes 148 | out_ft = torch.zeros(batchsize, self.out_channels, x.size(-2), x.size(-1)//2 + 1, dtype=torch.cfloat, device=x.device) 149 | out_ft[:, :, :self.modes1, :self.modes2] = \ 150 | self.compl_mul2d(x_ft[:, :, :self.modes1, :self.modes2], self.weights1) 151 | out_ft[:, :, -self.modes1:, :self.modes2] = \ 152 | self.compl_mul2d(x_ft[:, :, -self.modes1:, :self.modes2], self.weights2) 153 | 154 | #Return to physical space 155 | x = torch.fft.irfft2(out_ft, s=(x.size(-2), x.size(-1))) 156 | return x 157 | 158 | class FNO2d(nn.Module): 159 | def __init__(self, num_channels, bc, modes1=12, modes2=12, width=10): 160 | super(FNO2d, self).__init__() 161 | 162 | """ 163 | The overall network. It contains 4 layers of the Fourier layer. 164 | 1. Lift the input to the desire channel dimension by self.fc0 . 165 | 2. 4 layers of the integral operators u' = (W + K)(u). 166 | W defined by self.w; K defined by self.conv . 167 | 3. Project from the channel space to the output space by self.fc1 and self.fc2 . 168 | 169 | input: the solution of the previous 10 timesteps + 2 locations (u(t-10, x, y), ..., u(t-1, x, y), x, y) 170 | input shape: (batchsize, x, y, c) 171 | output: the solution of the next timestep 172 | output shape: (batchsize, x, y, c) 173 | """ 174 | 175 | self.modes1 = modes1 176 | self.modes2 = modes2 177 | self.width = width 178 | self.padding = 2 # pad the domain if input is non-periodic 179 | self.bc = bc 180 | self.fc0 = nn.Linear(10*num_channels+2, self.width) 181 | # input channel is 12: the solution of the previous 10 timesteps + 2 locations (u(t-10, x, y), ..., u(t-1, x, y), x, y) 182 | 183 | self.conv0 = SpectralConv2d_fast(self.width, self.width, self.modes1, self.modes2) 184 | self.conv1 = SpectralConv2d_fast(self.width, self.width, self.modes1, self.modes2) 185 | self.conv2 = SpectralConv2d_fast(self.width, self.width, self.modes1, self.modes2) 186 | self.conv3 = SpectralConv2d_fast(self.width, self.width, self.modes1, self.modes2) 187 | self.w0 = nn.Conv2d(self.width, self.width, 1) 188 | self.w1 = nn.Conv2d(self.width, self.width, 1) 189 | self.w2 = nn.Conv2d(self.width, self.width, 1) 190 | self.w3 = nn.Conv2d(self.width, self.width, 1) 191 | 192 | self.fc1 = nn.Linear(self.width, 32) 193 | self.fc2 = nn.Linear(32, num_channels) 194 | 195 | def forward(self, x, grid): 196 | x = torch.cat((x, grid), dim=-1) 197 | x = self.fc0(x) 198 | x = x.permute(0, 3, 1, 2) 199 | 200 | # Pad tensor with boundary condition 201 | mult = x.size(1)//self.bc.size(1) 202 | Nx = x.size(-2) 203 | x = torch.cat((self.bc[:1,:,:1].repeat(x.size(0),mult,Nx).unsqueeze(-1),x, 204 | self.bc[:1,:,1:2].repeat(x.size(0),mult,Nx).unsqueeze(-1)),dim=3) 205 | Ny = x.size(-1) 206 | x = torch.cat((self.bc[:1,:,2:3].repeat(x.size(0),mult,Ny).unsqueeze(-2),x, 207 | self.bc[:1,:,3:4].repeat(x.size(0),mult,Ny).unsqueeze(-2)),dim=2) 208 | 209 | x1 = self.conv0(x) 210 | x2 = self.w0(x) 211 | x = x1 + x2 212 | x = F.gelu(x) 213 | 214 | x1 = self.conv1(x) 215 | x2 = self.w1(x) 216 | x = x1 + x2 217 | x = F.gelu(x) 218 | 219 | x1 = self.conv2(x) 220 | x2 = self.w2(x) 221 | x = x1 + x2 222 | x = F.gelu(x) 223 | 224 | x1 = self.conv3(x) 225 | x2 = self.w3(x) 226 | x = x1 + x2 227 | 228 | x = x[..., 1:-1, 1:-1] # Unpad the tensor 229 | x = x.permute(0, 2, 3, 1) 230 | x = self.fc1(x) 231 | x = F.gelu(x) 232 | x = self.fc2(x) 233 | 234 | return x.unsqueeze(-2) -------------------------------------------------------------------------------- /models/phydnet/config.json: -------------------------------------------------------------------------------- 1 | { 2 | "paths": { 3 | }, 4 | 5 | "general": { 6 | "device": "cpu" // "cpu" or "cuda" 7 | }, 8 | 9 | "training": { 10 | "save_model": true, 11 | "continue_training": false, 12 | "epochs": 2001, 13 | "learning_rate": 0.001 14 | }, 15 | 16 | "data": { 17 | "type": "burger_2d", // "burger", "burger_2d", "diffusion_sorption", "diffusion_reaction", "allen_cahn" 18 | "name": "data_train", // "data_train", "data_ext", "data_test" 19 | "noise": 0.0 20 | }, 21 | 22 | "model": { 23 | "name": "burger_2d", // "burger", "burger_2d", "diff-sorp", "diff-react", "allen-cahn" 24 | "number": 0, // The i-th model 25 | "small" : 1 // 0 if original size, 1 if smaller network 26 | } 27 | } -------------------------------------------------------------------------------- /models/phydnet/constrain_moments.py: -------------------------------------------------------------------------------- 1 | from numpy import * 2 | from numpy.linalg import * 3 | from scipy.special import factorial 4 | from functools import reduce 5 | import torch 6 | import torch.nn as nn 7 | 8 | from functools import reduce 9 | 10 | __all__ = ['M2K','K2M'] 11 | 12 | def _apply_axis_left_dot(x, mats): 13 | assert x.dim() == len(mats)+1 14 | sizex = x.size() 15 | k = x.dim()-1 16 | for i in range(k): 17 | x = tensordot(mats[k-i-1], x, dim=[1,k]) 18 | x = x.permute([k,]+list(range(k))).contiguous() 19 | x = x.view(sizex) 20 | return x 21 | 22 | def _apply_axis_right_dot(x, mats): 23 | assert x.dim() == len(mats)+1 24 | sizex = x.size() 25 | k = x.dim()-1 26 | x = x.permute(list(range(1,k+1))+[0,]) 27 | for i in range(k): 28 | x = tensordot(x, mats[i], dim=[0,0]) 29 | x = x.contiguous() 30 | x = x.view(sizex) 31 | return x 32 | 33 | class _MK(nn.Module): 34 | def __init__(self, shape): 35 | super(_MK, self).__init__() 36 | self._size = torch.Size(shape) 37 | self._dim = len(shape) 38 | M = [] 39 | invM = [] 40 | assert len(shape) > 0 41 | j = 0 42 | for l in shape: 43 | M.append(zeros((l,l))) 44 | for i in range(l): 45 | M[-1][i] = ((arange(l)-(l-1)//2)**i)/factorial(i) 46 | invM.append(inv(M[-1])) 47 | self.register_buffer('_M'+str(j), torch.from_numpy(M[-1])) 48 | self.register_buffer('_invM'+str(j), torch.from_numpy(invM[-1])) 49 | j += 1 50 | 51 | @property 52 | def M(self): 53 | return list(self._buffers['_M'+str(j)] for j in range(self.dim())) 54 | @property 55 | def invM(self): 56 | return list(self._buffers['_invM'+str(j)] for j in range(self.dim())) 57 | 58 | def size(self): 59 | return self._size 60 | def dim(self): 61 | return self._dim 62 | def _packdim(self, x): 63 | assert x.dim() >= self.dim() 64 | if x.dim() == self.dim(): 65 | x = x[newaxis,:] 66 | x = x.contiguous() 67 | x = x.view([-1,]+list(x.size()[-self.dim():])) 68 | return x 69 | 70 | def forward(self): 71 | pass 72 | 73 | class M2K(_MK): 74 | """ 75 | convert moment matrix to convolution kernel 76 | Arguments: 77 | shape (tuple of int): kernel shape 78 | Usage: 79 | m2k = M2K([5,5]) 80 | m = torch.randn(5,5,dtype=torch.float64) 81 | k = m2k(m) 82 | """ 83 | def __init__(self, shape): 84 | super(M2K, self).__init__(shape) 85 | def forward(self, m): 86 | """ 87 | m (Tensor): torch.size=[...,*self.shape] 88 | """ 89 | sizem = m.size() 90 | m = self._packdim(m) 91 | m = _apply_axis_left_dot(m, self.invM) 92 | m = m.view(sizem) 93 | return m 94 | 95 | class K2M(_MK): 96 | """ 97 | convert convolution kernel to moment matrix 98 | Arguments: 99 | shape (tuple of int): kernel shape 100 | Usage: 101 | k2m = K2M([5,5]) 102 | k = torch.randn(5,5,dtype=torch.float64) 103 | m = k2m(k) 104 | """ 105 | def __init__(self, shape): 106 | super(K2M, self).__init__(shape) 107 | def forward(self, k): 108 | """ 109 | k (Tensor): torch.size=[...,*self.shape] 110 | """ 111 | sizek = k.size() 112 | k = self._packdim(k) 113 | k = _apply_axis_left_dot(k, self.M) 114 | k = k.view(sizek) 115 | return k 116 | 117 | 118 | 119 | def tensordot(a,b,dim): 120 | """ 121 | tensordot in PyTorch, see numpy.tensordot? 122 | """ 123 | l = lambda x,y:x*y 124 | if isinstance(dim,int): 125 | a = a.contiguous() 126 | b = b.contiguous() 127 | sizea = a.size() 128 | sizeb = b.size() 129 | sizea0 = sizea[:-dim] 130 | sizea1 = sizea[-dim:] 131 | sizeb0 = sizeb[:dim] 132 | sizeb1 = sizeb[dim:] 133 | N = reduce(l, sizea1, 1) 134 | assert reduce(l, sizeb0, 1) == N 135 | else: 136 | adims = dim[0] 137 | bdims = dim[1] 138 | adims = [adims,] if isinstance(adims, int) else adims 139 | bdims = [bdims,] if isinstance(bdims, int) else bdims 140 | adims_ = set(range(a.dim())).difference(set(adims)) 141 | adims_ = list(adims_) 142 | adims_.sort() 143 | perma = adims_+adims 144 | bdims_ = set(range(b.dim())).difference(set(bdims)) 145 | bdims_ = list(bdims_) 146 | bdims_.sort() 147 | permb = bdims+bdims_ 148 | a = a.permute(*perma).contiguous() 149 | b = b.permute(*permb).contiguous() 150 | 151 | sizea = a.size() 152 | sizeb = b.size() 153 | sizea0 = sizea[:-len(adims)] 154 | sizea1 = sizea[-len(adims):] 155 | sizeb0 = sizeb[:len(bdims)] 156 | sizeb1 = sizeb[len(bdims):] 157 | N = reduce(l, sizea1, 1) 158 | assert reduce(l, sizeb0, 1) == N 159 | a = a.view([-1,N]) 160 | b = b.view([N,-1]) 161 | c = a@b 162 | return c.view(sizea0+sizeb1) -------------------------------------------------------------------------------- /models/phydnet/experiment.py: -------------------------------------------------------------------------------- 1 | """ 2 | This script either trains n models or tests them as determined in the 3 | config.json file. 4 | """ 5 | 6 | import multiprocessing as mp 7 | import numpy as np 8 | import train 9 | import test 10 | 11 | # 12 | # Parameters 13 | 14 | TRAIN = False 15 | MODEL_NUMBERS = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] 16 | 17 | # Number of models that can be trained in parallel 18 | POOLSIZE = 5 19 | 20 | 21 | # 22 | # SCRIPT 23 | 24 | if TRAIN: 25 | 26 | # Train the MODEL_NUMBERS using POOLSIZE threads in parallel 27 | if POOLSIZE > 1: 28 | arguments = list() 29 | for model_number in MODEL_NUMBERS: 30 | arguments.append([False, model_number]) 31 | 32 | with mp.Pool(POOLSIZE) as pool: 33 | pool.starmap(train.run_training, arguments) 34 | else: 35 | for model_number in MODEL_NUMBERS: 36 | train.run_training(print_progress=False, model_number=model_number) 37 | 38 | else: 39 | # Test the MODEL_NUMBERS iteratively 40 | error_list = list() 41 | for model_number in MODEL_NUMBERS: 42 | mse = test.run_testing(model_number=model_number) 43 | error_list.append(mse) 44 | 45 | print(error_list) 46 | print(f"Average MSE: {np.mean(error_list)}, STD: {np.std(error_list)}") 47 | -------------------------------------------------------------------------------- /models/pinn/config.json: -------------------------------------------------------------------------------- 1 | { 2 | "paths": { 3 | }, 4 | 5 | "general": { 6 | "device": "cpu" // "cpu", or "cuda" 7 | }, 8 | 9 | "training": { 10 | "save_model": true, 11 | "continue_training": false, 12 | "epochs": 500, 13 | "batch_size": 20250, // burger 7565, burger_2d 20250, diff-sorp 10065, diff-react 20250, allen_cahn 5065 14 | "learning_rate": 1.0 15 | }, 16 | 17 | "validation": { 18 | "batch_size": 20250 // burger 7565, diff-sorp 10065, diff-react 20250, allen_cahn 5065 19 | }, 20 | 21 | "data": { 22 | "type": "burger_2d", // "burger", "burger_2d", "diffusion_sorption", "diffusion_reaction", "allen_cahn" 23 | "name": "data_train", // "data_train", "data_ext", "data_test" 24 | "burger": { 25 | }, 26 | "diffusion_sorption": { 27 | "D": 0.0005, 28 | "porosity": 0.29, 29 | "rho_s": 2880, 30 | "k_f_nominator": 1.016, // will become k_f = k_f_nominator/rho_s 31 | "n_f": 0.874, 32 | "s_max": 0.000588235, // 1/1700 33 | "kl": 1, 34 | "kd": 0.000429, // 0.429/1000 35 | "solubility": 1.0 36 | }, 37 | "diffusion_reaction": { 38 | "k": 5E-3, 39 | "D_u": 1E-3, 40 | "D_v": 5E-3 41 | }, 42 | "noise": 0.0 43 | }, 44 | 45 | "model": { 46 | "name": "burger_2d", // "burger", "burger_2d", "diff-sorp", "diff-react", "allen-cahn" 47 | "number": 0, // The i-th model 48 | "layer_sizes": [3, 20, 20, 20, 20, 20, 20, 20, 20, 1] // burger and allen_cahn [2, ..., 1], diff-sorp [2, ..., 2], diff-react [3, ..., 2], burger_2d [3, ..., 1] 49 | } 50 | } 51 | -------------------------------------------------------------------------------- /models/pinn/experiment.py: -------------------------------------------------------------------------------- 1 | """ 2 | This script either trains n models or tests them as determined in the 3 | config.json file. 4 | """ 5 | 6 | import multiprocessing as mp 7 | import numpy as np 8 | import train 9 | import test 10 | 11 | # 12 | # Parameters 13 | 14 | TRAIN = False 15 | MODEL_NUMBERS = [1, 2, 3, 4, 5, 6, 7, 8, 9] 16 | 17 | # Number of models that can be trained in parallel 18 | POOLSIZE = 1 19 | 20 | 21 | # 22 | # SCRIPT 23 | 24 | if TRAIN: 25 | 26 | # Train the MODEL_NUMBERS using POOLSIZE threads in parallel 27 | if POOLSIZE > 1: 28 | arguments = list() 29 | for model_number in MODEL_NUMBERS: 30 | arguments.append([False, model_number]) 31 | 32 | with mp.Pool(POOLSIZE) as pool: 33 | pool.starmap(train.run_training, arguments) 34 | else: 35 | for model_number in MODEL_NUMBERS: 36 | train.run_training(print_progress=False, model_number=model_number) 37 | 38 | else: 39 | # Test the MODEL_NUMBERS iteratively 40 | error_list = list() 41 | for model_number in MODEL_NUMBERS: 42 | mse, pred, labels = test.run_testing(model_number=model_number) 43 | error_list.append(mse) 44 | 45 | print(error_list) 46 | print(f"Average MSE: {np.mean(error_list)}, STD: {np.std(error_list)}") 47 | -------------------------------------------------------------------------------- /models/polynomial/allen_cahn.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | """ 3 | Created on Wed Nov 10 09:15:03 2021 4 | """ 5 | 6 | import numpy as np 7 | import os 8 | import matplotlib.pyplot as plt 9 | from mpl_toolkits.axes_grid1 import make_axes_locatable 10 | from sklearn.linear_model import Ridge, Lasso 11 | from sklearn.preprocessing import PolynomialFeatures 12 | from sklearn.pipeline import make_pipeline 13 | 14 | # Data path 15 | root_path = os.path.abspath("../../data") 16 | data_type = "allen_cahn" 17 | data_name = "data_train" 18 | data_name_ext = "data_ext" 19 | data_path = os.path.join(root_path, data_type, data_name) 20 | 21 | # Polynomial configuration 22 | max_degree = 10 23 | 24 | # Load data and reshape 25 | t = np.load(os.path.join(data_path, "t_series.npy")) 26 | x = np.load(os.path.join(data_path, "x_series.npy")) 27 | u = np.load(os.path.join(data_path, "sample.npy")) 28 | 29 | t_mesh, x_mesh = np.meshgrid(t,x) 30 | t_mesh = t_mesh.flatten() 31 | x_mesh = x_mesh.flatten() 32 | 33 | inp_mesh = np.stack((t_mesh,x_mesh), -1) 34 | out_mesh = u.transpose().flatten() 35 | out_mesh = np.expand_dims(out_mesh,-1) 36 | 37 | # Fit polynomial to data 38 | model = make_pipeline(PolynomialFeatures(max_degree), Lasso(alpha=0)) 39 | model.fit(inp_mesh,out_mesh) 40 | 41 | # Predict training data and plot 42 | pred = model.predict(inp_mesh) 43 | pred = pred.reshape((len(x),len(t))).transpose() 44 | 45 | fig, ax = plt.subplots(2, 2, figsize=(4.5,4)) 46 | h = ax[0,0].imshow(pred.transpose(), interpolation='nearest', 47 | extent=[t.min(), t.max(), 48 | x.min(), x.max()], 49 | origin='lower', aspect='auto') 50 | divider = make_axes_locatable(ax[0,0]) 51 | cax = divider.append_axes("right", size="5%", pad=0.05) 52 | cbar = fig.colorbar(h, cax=cax) 53 | h.set_clim(u.min(), u.max()) 54 | ax[0,0].set_title("Training") 55 | ax[0,0].set_ylabel("$x$") 56 | ax[0,0].set_xlabel("$t$") 57 | ax[0,0].axes.set_xticks([0, 0.25, 0.5]) 58 | ax[0,0].axes.set_yticks([-1, 0, 1]) 59 | 60 | h = ax[1,0].plot(x, u[-1], "ro-", markersize=2, linewidth = 0.5) 61 | h = ax[1,0].plot(x, pred[-1]) 62 | ax[1,0].set_title("t = 0.5") 63 | ax[1,0].set_ylabel("$u$") 64 | ax[1,0].set_xlabel("$x$") 65 | ax[1,0].axes.set_xticks([-1, 0, 1]) 66 | ax[1,0].axes.set_yticks([-1.0, -0.5, 0, 0.5]) 67 | 68 | print("MSE train = ", np.mean((u-pred)**2)) 69 | 70 | # # Extrapolation 71 | 72 | # Load data and reshape 73 | data_path = os.path.join(root_path, data_type, data_name_ext) 74 | t = np.load(os.path.join(data_path, "t_series.npy")) 75 | x = np.load(os.path.join(data_path, "x_series.npy")) 76 | u = np.load(os.path.join(data_path, "sample.npy")) 77 | 78 | t_mesh, x_mesh = np.meshgrid(t,x) 79 | t_mesh = t_mesh.flatten() 80 | x_mesh = x_mesh.flatten() 81 | 82 | inp_mesh = np.stack((t_mesh,x_mesh),-1) 83 | out_mesh = u.transpose().flatten() 84 | out_mesh = np.expand_dims(out_mesh,-1) 85 | 86 | # Predict extrapolation and plot 87 | 88 | pred = model.predict(inp_mesh) 89 | pred = pred.reshape((len(x),len(t))).transpose() 90 | 91 | # fig, ax = plt.subplots(2, 1, figsize=(3,5)) 92 | h = ax[0,1].imshow(pred.transpose(), interpolation='nearest', 93 | extent=[t.min(), t.max(), 94 | x.min(), x.max()], 95 | origin='lower', aspect='auto') 96 | divider = make_axes_locatable(ax[0,1]) 97 | cax = divider.append_axes("right", size="5%", pad=0.05) 98 | cbar = fig.colorbar(h, cax=cax) 99 | h.set_clim(u.min(), u.max()) 100 | ax[0,1].set_title("Extrapolation") 101 | ax[0,1].set_ylabel("$x$") 102 | ax[0,1].set_xlabel("$t$") 103 | ax[0,1].axes.set_xticks([0, 0.5, 1.0]) 104 | ax[0,1].axes.set_yticks([-1, 0, 1]) 105 | 106 | h = ax[1,1].plot(x, u[-1], "ro-", markersize=2, linewidth = 0.5, label="Data") 107 | h = ax[1,1].plot(x, pred[-1],label="Prediction") 108 | ax[1,1].set_title("t = 1") 109 | ax[1,1].set_ylabel("$u$") 110 | ax[1,1].set_xlabel("$x$") 111 | ax[1,1].axes.set_xticks([-1, 0, 1]) 112 | ax[1,1].axes.set_yticks([0, 150, 300]) 113 | 114 | fig.legend(loc=8, ncol=2) 115 | fig.suptitle("Allen-Cahn, polynomial (d=10)") 116 | plt.tight_layout(rect=[0,0.05,1,0.95]) # [left, bottom, right, top] 117 | plt.savefig('poly_allen_cahn.pdf') 118 | 119 | print("MSE ext = ", np.mean((u-pred)**2)) -------------------------------------------------------------------------------- /models/polynomial/burger.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | """ 3 | Created on Wed Nov 10 09:15:03 2021 4 | """ 5 | 6 | import numpy as np 7 | import os 8 | import matplotlib.pyplot as plt 9 | from mpl_toolkits.axes_grid1 import make_axes_locatable 10 | from sklearn.linear_model import Ridge, Lasso 11 | from sklearn.preprocessing import PolynomialFeatures 12 | from sklearn.pipeline import make_pipeline 13 | 14 | # Data path 15 | root_path = os.path.abspath("../../data") 16 | data_type = "burger" 17 | data_name = "data_train" 18 | data_name_ext = "data_ext" 19 | data_path = os.path.join(root_path, data_type, data_name) 20 | 21 | # Polynomial configuration 22 | max_degree = 7 23 | 24 | # Load data and reshape 25 | t = np.load(os.path.join(data_path, "t_series.npy")) 26 | x = np.load(os.path.join(data_path, "x_series.npy")) 27 | u = np.load(os.path.join(data_path, "sample.npy")) 28 | 29 | t_mesh, x_mesh = np.meshgrid(t,x) 30 | t_mesh = t_mesh.flatten() 31 | x_mesh = x_mesh.flatten() 32 | 33 | inp_mesh = np.stack((t_mesh,x_mesh), -1) 34 | out_mesh = u.transpose().flatten() 35 | out_mesh = np.expand_dims(out_mesh,-1) 36 | 37 | # Fit polynomial to data 38 | model = make_pipeline(PolynomialFeatures(max_degree), Lasso(alpha=1e-3)) 39 | model.fit(inp_mesh,out_mesh) 40 | 41 | # Predict training data and plot 42 | pred = model.predict(inp_mesh) 43 | pred = pred.reshape((len(x),len(t))).transpose() 44 | 45 | fig, ax = plt.subplots(2, 2, figsize=(4.5,4)) 46 | h = ax[0,0].imshow(pred.transpose(), interpolation='nearest', 47 | extent=[t.min(), t.max(), 48 | x.min(), x.max()], 49 | origin='lower', aspect='auto') 50 | divider = make_axes_locatable(ax[0,0]) 51 | cax = divider.append_axes("right", size="5%", pad=0.05) 52 | cbar = fig.colorbar(h, cax=cax) 53 | h.set_clim(u.min(), u.max()) 54 | ax[0,0].set_title("Training") 55 | ax[0,0].set_ylabel("$x$") 56 | ax[0,0].set_xlabel("$t$") 57 | ax[0,0].axes.set_xticks([0, 0.5, 1.0]) 58 | ax[0,0].axes.set_yticks([-1, 0, 1]) 59 | 60 | h = ax[1,0].plot(x, u[-1], "ro-", markersize=2, linewidth = 0.5) 61 | h = ax[1,0].plot(x, pred[-1]) 62 | ax[1,0].set_title("t = 1") 63 | ax[1,0].set_ylabel("$u$") 64 | ax[1,0].set_xlabel("$x$") 65 | ax[1,0].axes.set_xticks([-1, 0, 1]) 66 | ax[1,0].axes.set_yticks([-0.5, 0, 0.5]) 67 | 68 | print("MSE train = ", np.mean((u-pred)**2)) 69 | 70 | 71 | # # Extrapolation 72 | 73 | # Load data and reshape 74 | data_path = os.path.join(root_path, data_type, data_name_ext) 75 | t = np.load(os.path.join(data_path, "t_series.npy")) 76 | x = np.load(os.path.join(data_path, "x_series.npy")) 77 | u = np.load(os.path.join(data_path, "sample.npy")) 78 | 79 | t_mesh, x_mesh = np.meshgrid(t,x) 80 | t_mesh = t_mesh.flatten() 81 | x_mesh = x_mesh.flatten() 82 | 83 | inp_mesh = np.stack((t_mesh,x_mesh),-1) 84 | out_mesh = u.transpose().flatten() 85 | out_mesh = np.expand_dims(out_mesh,-1) 86 | 87 | # Predict extrapolation and plot 88 | 89 | pred = model.predict(inp_mesh) 90 | pred = pred.reshape((len(x),len(t))).transpose() 91 | 92 | # fig, ax = plt.subplots(2, 1, figsize=(3,5)) 93 | h = ax[0,1].imshow(pred.transpose(), interpolation='nearest', 94 | extent=[t.min(), t.max(), 95 | x.min(), x.max()], 96 | origin='lower', aspect='auto') 97 | divider = make_axes_locatable(ax[0,1]) 98 | cax = divider.append_axes("right", size="5%", pad=0.05) 99 | cbar = fig.colorbar(h, cax=cax) 100 | h.set_clim(u.min(), u.max()) 101 | ax[0,1].set_title("Extrapolation") 102 | ax[0,1].set_ylabel("$x$") 103 | ax[0,1].set_xlabel("$t$") 104 | ax[0,1].axes.set_xticks([0, 1.0, 2.0]) 105 | ax[0,1].axes.set_yticks([-1, 0, 1]) 106 | 107 | h = ax[1,1].plot(x, u[-1], "ro-", markersize=2, linewidth = 0.5, label="Data") 108 | h = ax[1,1].plot(x, pred[-1],label="Prediction") 109 | ax[1,1].set_title("t = 2") 110 | ax[1,1].set_ylabel("$u$") 111 | ax[1,1].set_xlabel("$x$") 112 | ax[1,1].axes.set_xticks([-1, 0, 1]) 113 | ax[1,1].axes.set_yticks([-0.5, 0, 0.5]) 114 | 115 | fig.legend(loc=8, ncol=2) 116 | fig.suptitle(f"Burgers', polynomial (d=5)") 117 | plt.tight_layout(rect=[0,0.05,1,0.95]) # [left, bottom, right, top] 118 | plt.savefig('poly_burger.pdf') 119 | 120 | print("MSE ext = ", np.mean((u-pred)**2)) -------------------------------------------------------------------------------- /models/polynomial/diffusion_reaction.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | """ 3 | Created on Wed Nov 10 09:15:03 2021 4 | """ 5 | 6 | import numpy as np 7 | import os 8 | import matplotlib.pyplot as plt 9 | from mpl_toolkits.axes_grid1 import make_axes_locatable 10 | from sklearn.linear_model import Ridge, Lasso 11 | from sklearn.preprocessing import PolynomialFeatures 12 | from sklearn.pipeline import make_pipeline 13 | 14 | # Data path 15 | root_path = os.path.abspath("../../data") 16 | data_type = "diffusion_reaction" 17 | data_name = "data_train" 18 | data_name_ext = "data_ext" 19 | data_path = os.path.join(root_path, data_type, data_name) 20 | 21 | # Polynomial configuration 22 | max_degree = 5 23 | 24 | # Load data and reshape 25 | t = np.load(os.path.join(data_path, "t_series.npy")) 26 | x = np.load(os.path.join(data_path, "x_series.npy")) 27 | y = np.load(os.path.join(data_path, "y_series.npy")) 28 | u = np.load(os.path.join(data_path, "sample_u.npy")) 29 | v = np.load(os.path.join(data_path, "sample_v.npy")) 30 | 31 | x_mesh, t_mesh, y_mesh = np.meshgrid(x,t,y) 32 | t_mesh = t_mesh.flatten() 33 | x_mesh = x_mesh.flatten() 34 | y_mesh = y_mesh.flatten() 35 | 36 | inp_mesh = np.stack((t_mesh,x_mesh,y_mesh),-1) 37 | u_flat = u.transpose(0,2,1).flatten() 38 | v_flat = v.transpose(0,2,1).flatten() 39 | out_mesh = np.stack((u_flat,v_flat),-1) 40 | 41 | # Fit polynomial to data 42 | model = make_pipeline(PolynomialFeatures(max_degree), Lasso(alpha=1e-3)) 43 | model.fit(inp_mesh,out_mesh) 44 | 45 | # Predict training data and plot 46 | pred = model.predict(inp_mesh)[...,0] 47 | pred = pred.reshape((len(t),len(y),len(x))) 48 | 49 | fig, ax = plt.subplots(2, 2, figsize=(4.5,4)) 50 | 51 | h = ax[0,0].imshow(pred[-1].transpose(), interpolation='nearest', 52 | extent=[x.min(), x.max(), 53 | y.min(), y.max()], 54 | origin='lower', aspect='auto') 55 | divider = make_axes_locatable(ax[0,0]) 56 | cax = divider.append_axes("right", size="5%", pad=0.05) 57 | cbar = fig.colorbar(h, cax=cax, ticks=[-0.5,0,0.5]) 58 | h.set_clim(u[-1].min(), u[-1].max()) 59 | ax[0,0].set_title("Training") 60 | ax[0,0].set_xlabel("$x$") 61 | ax[0,0].set_ylabel("$y$") 62 | ax[0,0].axes.set_xticks([-1, 0, 1]) 63 | ax[0,0].axes.set_yticks([-1, 0, 1]) 64 | ax[0,0].set_xticklabels([-1, 0, 1]) 65 | ax[0,0].set_yticklabels([-1, 0, 1]) 66 | 67 | h = ax[1,0].plot(x, u[-1,...,49//2], "ro-", markersize=2, linewidth = 0.5) 68 | h = ax[1,0].plot(x, pred[-1,...,49//2]) 69 | ax[1,0].set_title("t = 10, y = 0") 70 | ax[1,0].set_xlabel("$x$") 71 | ax[1,0].set_ylabel("$u_1$") 72 | ax[1,0].axes.set_xticks([-1,0,1]) 73 | ax[1,0].axes.set_yticks([-0.5,0,0.5]) 74 | ax[1,0].set_xticklabels([-1, 0, 1]) 75 | ax[1,0].set_yticklabels([-0.5, 0, 0.5]) 76 | 77 | print("MSE train = ", np.mean((u-pred)**2)) 78 | 79 | # Extrapolation 80 | 81 | # Load data 82 | data_path = os.path.join(root_path, data_type, data_name_ext) 83 | t = np.load(os.path.join(data_path, "t_series.npy")) 84 | x = np.load(os.path.join(data_path, "x_series.npy")) 85 | y = np.load(os.path.join(data_path, "y_series.npy")) 86 | u = np.load(os.path.join(data_path, "sample_u.npy")) 87 | v = np.load(os.path.join(data_path, "sample_v.npy")) 88 | 89 | x_mesh, t_mesh, y_mesh = np.meshgrid(x,t,y) 90 | t_mesh = t_mesh.flatten() 91 | x_mesh = x_mesh.flatten() 92 | y_mesh = y_mesh.flatten() 93 | 94 | inp_mesh = np.stack((t_mesh,x_mesh,y_mesh),-1) 95 | u_flat = u.transpose(0,2,1).flatten() 96 | v_flat = v.transpose(0,2,1).flatten() 97 | out_mesh = np.stack((u_flat,v_flat),-1) 98 | 99 | # Predict extrapolation and plot 100 | pred = model.predict(inp_mesh)[...,0] 101 | pred = pred.reshape((len(t),len(y),len(x))) 102 | 103 | h = ax[0,1].imshow(pred[-1].transpose(), interpolation='nearest', 104 | extent=[x.min(), x.max(), 105 | y.min(), y.max()], 106 | origin='lower', aspect='auto') 107 | divider = make_axes_locatable(ax[0,1]) 108 | cax = divider.append_axes("right", size="5%", pad=0.05) 109 | cbar = fig.colorbar(h, cax=cax, ticks=[-0.5,0,0.5]) 110 | h.set_clim(u[-1].min(), u[-1].max()) 111 | ax[0,1].set_title("Extrapolation") 112 | ax[0,1].set_xlabel("$x$") 113 | ax[0,1].set_ylabel("$y$") 114 | ax[0,1].axes.set_xticks([-1, 0, 1]) 115 | ax[0,1].axes.set_yticks([-1, 0, 1]) 116 | ax[0,1].set_xticklabels([-1, 0, 1]) 117 | ax[0,1].set_yticklabels([-1, 0, 1]) 118 | 119 | h = ax[1,1].plot(x, u[-1,...,49//2], "ro-", markersize=2, linewidth = 0.5, label="Data") 120 | h = ax[1,1].plot(x, pred[-1,...,49//2], label="Prediction") 121 | ax[1,1].set_title("t = 50, y = 0") 122 | ax[1,1].set_xlabel("$x$") 123 | ax[1,1].set_ylabel("$u_1$") 124 | ax[1,1].axes.set_xticks([-1,0,1]) 125 | ax[1,1].axes.set_yticks([0,350,700]) 126 | ax[1,1].set_xticklabels([-1, 0, 1]) 127 | ax[1,1].set_yticklabels([0,350,700]) 128 | 129 | fig.legend(loc=8, ncol=2) 130 | plt.suptitle("Diffusion reaction, polynomial (d=5)") 131 | plt.tight_layout(rect=[0,0.05,1,0.95]) # [left, bottom, right, top] 132 | plt.savefig('poly_diff_react.pdf') 133 | 134 | print("MSE ext = ", np.mean((u-pred)**2)) -------------------------------------------------------------------------------- /models/polynomial/diffusion_sorption.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | """ 3 | Created on Wed Nov 10 09:15:03 2021 4 | """ 5 | 6 | import numpy as np 7 | import os 8 | import matplotlib.pyplot as plt 9 | from mpl_toolkits.axes_grid1 import make_axes_locatable 10 | from sklearn.linear_model import Ridge, Lasso 11 | from sklearn.preprocessing import PolynomialFeatures 12 | from sklearn.pipeline import make_pipeline 13 | 14 | # Data path 15 | root_path = os.path.abspath("../../data") 16 | data_type = "diffusion_sorption" 17 | data_name = "data_train" 18 | data_name_ext = "data_ext" 19 | data_path = os.path.join(root_path, data_type, data_name) 20 | 21 | # Polynomial configuration 22 | max_degree = 3 23 | 24 | # Load data and reshape 25 | t = np.load(os.path.join(data_path, "t_series.npy")) 26 | x = np.load(os.path.join(data_path, "x_series.npy")) 27 | c = np.load(os.path.join(data_path, "sample_c.npy")) 28 | ct = np.load(os.path.join(data_path, "sample_ct.npy")) 29 | 30 | t_mesh, x_mesh = np.meshgrid(t,x) 31 | t_mesh = t_mesh.flatten() 32 | x_mesh = x_mesh.flatten() 33 | 34 | inp_mesh = np.stack((t_mesh,x_mesh),-1) 35 | c_flat = c.transpose().flatten() 36 | ct_flat = ct.transpose().flatten() 37 | out_mesh = np.stack((c_flat,ct_flat),-1) 38 | 39 | # Fit polynomial to data 40 | model = make_pipeline(PolynomialFeatures(max_degree), Lasso(alpha=1e-3)) 41 | model.fit(inp_mesh,out_mesh) 42 | 43 | # Predict training data and plot 44 | pred = model.predict(inp_mesh)[...,0] 45 | pred = pred.reshape((len(x),len(t))).transpose() 46 | 47 | fig, ax = plt.subplots(2, 2, figsize=(4.5,4)) 48 | h = ax[0,0].imshow(pred.transpose(), interpolation='nearest', 49 | extent=[t.min(), t.max(), 50 | x.min(), x.max()], 51 | origin='lower', aspect='auto') 52 | divider = make_axes_locatable(ax[0,0]) 53 | cax = divider.append_axes("right", size="5%", pad=0.05) 54 | cbar = fig.colorbar(h, cax=cax) 55 | h.set_clim(c.min(), c.max()) 56 | ax[0,0].set_title("Training") 57 | ax[0,0].set_ylabel("$x$") 58 | ax[0,0].set_xlabel("$t$") 59 | ax[0,0].axes.set_xticks([0, 2500]) 60 | ax[0,0].axes.set_yticks([0.0, 0.5, 1.0]) 61 | ax[0,0].set_xticklabels([0, 2500]) 62 | ax[0,0].set_yticklabels([0.0, 0.5, 1.0]) 63 | 64 | h = ax[1,0].plot(x, c[-1], "ro-", markersize=2, linewidth = 0.5) 65 | h = ax[1,0].plot(x, pred[-1]) 66 | ax[1,0].set_title("t = 2500") 67 | ax[1,0].set_ylabel("$u$") 68 | ax[1,0].set_xlabel("$x$") 69 | ax[1,0].axes.set_xticks([0, 0.5, 1]) 70 | ax[1,0].axes.set_yticks([0, 0.5, 1]) 71 | ax[1,0].set_xticklabels([0, 0.5, 1]) 72 | ax[1,0].set_yticklabels([0, 0.5, 1]) 73 | 74 | print("MSE train = ", np.mean((c-pred)**2)) 75 | 76 | # Extrapolation 77 | 78 | # Load data and reshape 79 | data_path = os.path.join(root_path, data_type, data_name_ext) 80 | t = np.load(os.path.join(data_path, "t_series.npy")) 81 | x = np.load(os.path.join(data_path, "x_series.npy")) 82 | c = np.load(os.path.join(data_path, "sample_c.npy")) 83 | ct = np.load(os.path.join(data_path, "sample_ct.npy")) 84 | 85 | t_mesh, x_mesh = np.meshgrid(t,x) 86 | t_mesh = t_mesh.flatten() 87 | x_mesh = x_mesh.flatten() 88 | 89 | inp_mesh = np.stack((t_mesh,x_mesh),-1) 90 | c_flat = c.transpose().flatten() 91 | ct_flat = ct.transpose().flatten() 92 | out_mesh = np.stack((c_flat,ct_flat),-1) 93 | 94 | # Predict extrapolation and plot 95 | pred = model.predict(inp_mesh)[...,0] 96 | pred = pred.reshape((len(x),len(t))).transpose() 97 | 98 | h = ax[0,1].imshow(pred.transpose(), interpolation='nearest', 99 | extent=[t.min(), t.max(), 100 | x.min(), x.max()], 101 | origin='lower', aspect='auto') 102 | divider = make_axes_locatable(ax[0,1]) 103 | cax = divider.append_axes("right", size="5%", pad=0.05) 104 | cbar = fig.colorbar(h, cax=cax) 105 | h.set_clim(c.min(), c.max()) 106 | ax[0,1].set_title("Extrapolation") 107 | ax[0,1].set_ylabel("$x$") 108 | ax[0,1].set_xlabel("$t$") 109 | ax[0,1].axes.set_xticks([0, 10000]) 110 | ax[0,1].axes.set_yticks([0.0, 0.5, 1.0]) 111 | ax[0,1].set_xticklabels([0, 10000]) 112 | ax[0,1].set_yticklabels([0.0, 0.5, 1.0]) 113 | 114 | h = ax[1,1].plot(x, c[-1], "ro-", markersize=2, linewidth = 0.5, label="Data") 115 | h = ax[1,1].plot(x, pred[-1],label="Prediction") 116 | ax[1,1].set_title("t = 10000") 117 | ax[1,1].set_ylabel("$u$") 118 | ax[1,1].set_xlabel("$x$") 119 | ax[1,1].axes.set_xticks([0, 0.5, 1]) 120 | ax[1,1].axes.set_yticks([0, 10, 20]) 121 | ax[1,1].set_xticklabels([0, 0.5, 1]) 122 | ax[1,1].set_yticklabels([0, 10, 20]) 123 | 124 | fig.legend(loc=8, ncol=2) 125 | plt.suptitle("Diffusion sorption, polynomial (d=3)") 126 | plt.tight_layout(rect=[0,0.05,1,0.95]) # [left, bottom, right, top] 127 | plt.savefig('poly_diff_sorp.pdf') 128 | 129 | print("MSE ext = ", np.mean((c-pred)**2)) -------------------------------------------------------------------------------- /models/tcn/config.json: -------------------------------------------------------------------------------- 1 | { 2 | "paths": { 3 | }, 4 | 5 | "general": { 6 | "device": "CPU" // "CPU" or "GPU" 7 | }, 8 | 9 | "training": { 10 | "save_model": true, 11 | "continue_training": false, 12 | "epochs": 200, 13 | "batch_size": 1, 14 | "learning_rate": 1, 15 | "t_start": 0, 16 | "t_stop": 150 // burger, burger_2d, allen-cahn 150, diff-sorp 400, diff-react 70 17 | }, 18 | 19 | "validation": { 20 | "batch_size": 1, 21 | "t_start": 150, // burger, burger_2d, allen-cahn 150, diff-sorp 400, diff-react 70 22 | "t_stop": 200 // burger, burger_2d, allen-cahn 200, diff-sorp 500, diff-react 100 23 | }, 24 | 25 | "testing": { 26 | "batch_size": 1, 27 | "teacher_forcing_steps": 20, 28 | "closed_loop_steps": 1990, 29 | "feed_boundary_data": true 30 | }, 31 | 32 | "data": { 33 | "type": "burger_2d", // "burger", "burger_2d", "diffusion_sorption", "diffusion_reaction", "allen_cahn" 34 | "name": "data_train", // "data_train", "data_ext", "data_test" 35 | "noise": 0.0 36 | }, 37 | 38 | "model": { 39 | "name": "burger_2d", // "burger", "burger_2d", "diff-sorp", "diff-react", "allen-cahn" 40 | "number": 0, // The i-th model 41 | "field_size": [49, 49], // burger and allen-cahn [49], diff-sorp [26], burger_2d and diff-react [49, 49] 42 | "channels": [1, 32, 1], // burger, buger_2d, and allen-cahn [1, ..., 1], diff-sorp and diff-react [2, ..., 2] 43 | "kernel_size": 3, 44 | "horizon": 250 45 | } 46 | } 47 | -------------------------------------------------------------------------------- /models/tcn/experiment.py: -------------------------------------------------------------------------------- 1 | """ 2 | This script either trains n models or tests them as determined in the 3 | config.json file. 4 | """ 5 | 6 | import multiprocessing as mp 7 | import numpy as np 8 | import train 9 | import test 10 | 11 | # 12 | # Parameters 13 | 14 | TRAIN = False 15 | MODEL_NUMBERS = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] 16 | 17 | # Number of models that can be trained in parallel 18 | POOLSIZE = 1 19 | 20 | # 21 | # SCRIPT 22 | 23 | if TRAIN: 24 | 25 | # Train the MODEL_NUMBERS using POOLSIZE threads in parallel 26 | if POOLSIZE > 1: 27 | arguments = list() 28 | for model_number in MODEL_NUMBERS: 29 | arguments.append([False, model_number]) 30 | 31 | with mp.Pool(POOLSIZE) as pool: 32 | pool.starmap(train.run_training, arguments) 33 | else: 34 | for model_number in MODEL_NUMBERS: 35 | train.run_training(print_progress=False, model_number=model_number) 36 | 37 | else: 38 | # Test the MODEL_NUMBERS iteratively 39 | error_list = list() 40 | for model_number in MODEL_NUMBERS: 41 | mse = test.run_testing(model_number=model_number) 42 | error_list.append(mse) 43 | 44 | print(error_list) 45 | print(f"Average MSE: {np.mean(error_list)}, STD: {np.std(error_list)}") 46 | -------------------------------------------------------------------------------- /models/tcn/tcn.py: -------------------------------------------------------------------------------- 1 | """ 2 | This code is taken from https://github.com/locuslab/TCN/blob/master/TCN/tcn.py 3 | """ 4 | 5 | 6 | import torch 7 | import torch.nn as nn 8 | from torch.nn.utils import weight_norm 9 | 10 | 11 | class Chomp(nn.Module): 12 | def __init__(self, chomp_size): 13 | super(Chomp, self).__init__() 14 | self.chomp_size = chomp_size 15 | 16 | def forward(self, x): 17 | return x[:, :, :-self.chomp_size].contiguous() 18 | """ 19 | class Chomp1d(nn.Module): 20 | def __init__(self, chomp_size): 21 | super(Chomp1d, self).__init__() 22 | self.chomp_size = chomp_size 23 | 24 | def forward(self, x): 25 | return x[:, :, :-self.chomp_size].contiguous() 26 | 27 | class Chomp2d(nn.Module): 28 | def __init__(self, chomp_size): 29 | super(Chomp2d, self).__init__() 30 | self.chomp_size = chomp_size 31 | 32 | def forward(self, x): 33 | return x[:, :, :-self.chomp_size, :].contiguous() 34 | 35 | 36 | class Chomp3d(nn.Module): 37 | def __init__(self, chomp_size): 38 | super(Chomp3d, self).__init__() 39 | self.chomp_size = chomp_size 40 | 41 | def forward(self, x): 42 | return x[:, :, :-self.chomp_size].contiguous() 43 | """ 44 | 45 | 46 | class TemporalBlock(nn.Module): 47 | def __init__(self, n_inputs, n_outputs, kernel_size, stride, dilation, padding_size, dimensions, dropout=0.2): 48 | super(TemporalBlock, self).__init__() 49 | 50 | #""" 51 | if dimensions == 1: 52 | conv_function = nn.Conv2d 53 | padding = ((kernel_size-1) * dilation, padding_size) 54 | dilation = (dilation, 1) 55 | elif dimensions == 2: 56 | conv_function = nn.Conv3d 57 | padding = ((kernel_size-1) * dilation, padding_size, padding_size) 58 | dilation = (dilation, 1, 1) 59 | 60 | self.conv1 = weight_norm(conv_function( 61 | n_inputs, n_outputs, kernel_size, stride=stride, padding=padding, 62 | dilation=dilation, bias=True 63 | )) 64 | self.chomp1 = Chomp(padding[0]) 65 | self.dropout1 = nn.Dropout(dropout) 66 | 67 | self.conv2 = weight_norm(conv_function( 68 | n_outputs, n_outputs, kernel_size, stride=stride, padding=padding, 69 | dilation=dilation, bias=True 70 | )) 71 | self.chomp2 = Chomp(padding[0]) 72 | self.dropout2 = nn.Dropout(dropout) 73 | 74 | self.net = nn.Sequential(self.conv1, self.chomp1, self.dropout1, 75 | self.conv2, self.chomp2, self.dropout2) 76 | 77 | self.downsample = conv_function(n_inputs, n_outputs, 1, bias=True) if n_inputs != n_outputs else None 78 | 79 | 80 | """ 81 | if dimensions == 1: 82 | padding = (kernel_size-1) * dilation 83 | self.conv1 = weight_norm(nn.Conv1d( 84 | n_inputs, n_outputs, kernel_size, stride=stride, padding=padding, 85 | dilation=dilation, bias=True 86 | )) 87 | self.chomp1 = Chomp1d(padding) 88 | self.dropout1 = nn.Dropout(dropout) 89 | 90 | self.conv2 = weight_norm(nn.Conv1d( 91 | n_outputs, n_outputs, kernel_size, stride=stride, padding=padding, 92 | dilation=dilation, bias=True 93 | )) 94 | self.chomp2 = Chomp1d(padding) 95 | self.dropout2 = nn.Dropout(dropout) 96 | 97 | self.net = nn.Sequential(self.conv1, self.chomp1, self.dropout1, 98 | self.conv2, self.chomp2, self.dropout2) 99 | 100 | self.downsample = nn.Conv1d(n_inputs, n_outputs, 1, bias=True) if n_inputs != n_outputs else None 101 | 102 | elif dimensions == 2: 103 | padding = ((kernel_size-1) * dilation, padding_size) 104 | dilation = (dilation, 1) 105 | self.conv1 = weight_norm(nn.Conv2d( 106 | n_inputs, n_outputs, kernel_size, stride=stride, padding=padding, 107 | dilation=dilation, bias=True 108 | )) 109 | self.chomp1 = Chomp2d(padding[0]) 110 | self.dropout1 = nn.Dropout(dropout) 111 | 112 | self.conv2 = weight_norm(nn.Conv2d( 113 | n_outputs, n_outputs, kernel_size, stride=stride, padding=padding, 114 | dilation=dilation, bias=True 115 | )) 116 | self.chomp2 = Chomp2d(padding[0]) 117 | self.dropout2 = nn.Dropout(dropout) 118 | 119 | self.net = nn.Sequential(self.conv1, self.chomp1, self.dropout1, 120 | self.conv2, self.chomp2, self.dropout2) 121 | 122 | self.downsample = nn.Conv2d(n_inputs, n_outputs, 1, bias=True) if n_inputs != n_outputs else None 123 | #""" 124 | 125 | self.init_weights() 126 | 127 | def init_weights(self): 128 | self.conv1.weight.data.normal_(0, 0.01) 129 | self.conv2.weight.data.normal_(0, 0.01) 130 | if self.downsample is not None: 131 | self.downsample.weight.data.normal_(0, 0.01) 132 | 133 | def forward(self, x): 134 | out = self.net(x) 135 | res = x if self.downsample is None else self.downsample(x) 136 | return torch.tanh(out + res) 137 | 138 | 139 | class TemporalConvNet(nn.Module): 140 | def __init__(self, config): 141 | super(TemporalConvNet, self).__init__() 142 | layers = [] 143 | num_levels = len(config.model.channels) 144 | for i in range(num_levels): 145 | dilation_size = 2 ** i 146 | in_channels = config.model.channels[i-1] 147 | out_channels = config.model.channels[i] 148 | padding_size = config.model.kernel_size // 2 149 | layers += [TemporalBlock( 150 | in_channels, out_channels, config.model.kernel_size, 151 | stride=1, 152 | dilation=dilation_size, 153 | padding_size=padding_size, 154 | dropout=0.2, 155 | dimensions=len(config.model.field_size) 156 | )] 157 | 158 | self.network = nn.Sequential(*layers) 159 | 160 | def forward(self, x): 161 | return self.network(x) 162 | -------------------------------------------------------------------------------- /models/tcn/test.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch as th 3 | import time 4 | import glob 5 | import os 6 | import matplotlib.pyplot as plt 7 | import matplotlib.animation as animation 8 | import sys 9 | 10 | sys.path.append("..") 11 | from utils.configuration import Configuration 12 | import utils.helper_functions as helpers 13 | from tcn import TemporalConvNet 14 | 15 | 16 | def run_testing(print_progress=False, visualize=False, model_number=None): 17 | 18 | th.set_num_threads(1) 19 | 20 | # Load the user configurations 21 | config = Configuration("config.json") 22 | 23 | # Append the model number to the name of the model 24 | if model_number is None: 25 | model_number = config.model.number 26 | config.model.name = config.model.name + "_" + str(model_number).zfill(2) 27 | 28 | # Hide the GPU(s) in case the user specified to use the CPU in the config 29 | # file 30 | if config.general.device == "CPU": 31 | os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" 32 | os.environ["CUDA_VISIBLE_DEVICES"] = "" 33 | 34 | # Set device on GPU if specified in the configuration file, else CPU 35 | device = helpers.determine_device() 36 | 37 | # Initialize and set up the network 38 | model = TemporalConvNet(config=config).to(device=device) 39 | 40 | if print_progress: 41 | # Count number of trainable parameters 42 | pytorch_total_params = sum( 43 | p.numel() for p in model.parameters() if p.requires_grad 44 | ) 45 | print("Trainable model parameters:", pytorch_total_params) 46 | 47 | # Restore the network by loading the weights saved in the .pt file 48 | print("Restoring model (that is the network\"s weights) from file...") 49 | 50 | model.load_state_dict(th.load(os.path.join(os.path.abspath(""), 51 | "checkpoints", 52 | config.model.name, 53 | config.model.name + ".pt"), 54 | map_location=device)) 55 | model.eval() 56 | 57 | """ 58 | TESTING 59 | """ 60 | 61 | # 62 | # Load data depending on the task 63 | if config.data.type == "burger": 64 | data_path = os.path.join("../../data/", 65 | config.data.type, 66 | config.data.name, 67 | "sample.npy") 68 | data = np.array(np.load(data_path), dtype=np.float32) 69 | data = np.expand_dims(data, axis=0) 70 | 71 | elif config.data.type == "diffusion_sorption": 72 | data_path_base = os.path.join("../../data/", 73 | config.data.type, 74 | config.data.name) 75 | data_path_c = os.path.join(data_path_base, "sample_c.npy") 76 | data_path_ct = os.path.join(data_path_base, "sample_ct.npy") 77 | data_c = np.array(np.load(data_path_c), dtype=np.float32) 78 | data_ct = np.array(np.load(data_path_ct), dtype=np.float32) 79 | data = np.stack((data_c, data_ct), axis=0) 80 | 81 | elif config.data.type == "diffusion_reaction": 82 | data_path_base = os.path.join("../../data/", 83 | config.data.type, 84 | config.data.name) 85 | data_path_u = os.path.join(data_path_base, "sample_u.npy") 86 | data_path_v = os.path.join(data_path_base, "sample_v.npy") 87 | data_u = np.array(np.load(data_path_u), dtype=np.float32) 88 | data_v = np.array(np.load(data_path_v), dtype=np.float32) 89 | data = np.stack((data_u, data_v), axis=0) 90 | 91 | elif config.data.type == "allen_cahn": 92 | data_path = os.path.join("../../data/", 93 | config.data.type, 94 | config.data.name, 95 | "sample.npy") 96 | data = np.array(np.load(data_path), dtype=np.float32) 97 | data = np.expand_dims(data, axis=0) 98 | 99 | elif config.data.type == "burger_2d": 100 | data_path = os.path.join("../../data/", 101 | config.data.type, 102 | config.data.name, 103 | "sample.npy") 104 | data = np.array(np.load(data_path), dtype=np.float32) 105 | data = np.expand_dims(data, axis=0) 106 | 107 | # Set up the training and validation datasets and -loaders 108 | data_test = th.tensor(data[:], device=device).unsqueeze(0) 109 | sequence_length = len(data_test[0, 0]) - 1 110 | 111 | # Evaluate the network for the given test data 112 | 113 | # Separate the data into network inputs and labels 114 | net_inputs = th.clone(data_test[:, :, :-1]) 115 | net_labels = th.clone(data_test[:, :, 1:]) 116 | 117 | # Set up an array of zeros to store the network outputs 118 | net_outputs = th.zeros(size=(config.testing.batch_size, 119 | config.model.channels[-1], 120 | sequence_length, 121 | *config.model.field_size), 122 | device=device) 123 | 124 | net_input_steps = net_inputs[:, :, :config.testing.teacher_forcing_steps] 125 | 126 | time_start = time.time() 127 | with th.no_grad(): 128 | # Forward the input through the network 129 | net_output = model.forward(x=net_input_steps) 130 | 131 | # Store the output of the network for this sequence step 132 | net_outputs[:, :, :config.testing.teacher_forcing_steps] = net_output 133 | 134 | horizon = config.model.horizon 135 | tf_steps = config.testing.teacher_forcing_steps 136 | 137 | # Iterate over the remaining sequence of the training example and 138 | # perform a forward pass 139 | for t in range(tf_steps, sequence_length): 140 | 141 | t_start = max(t - horizon, 0) 142 | 143 | # Prepare the network input for this sequence step 144 | # Closed loop - receiving the output of the last time step as input 145 | net_input_steps = net_outputs[:, :, t_start:t] 146 | 147 | # Teacher forcing for time steps before the switch to closed loop 148 | t_diff = tf_steps - t_start 149 | if t_diff > 0: 150 | net_input_steps[:, :, :t_diff] = th.clone( 151 | net_inputs[:, :, t_start:tf_steps] 152 | ) 153 | 154 | # Feed the boundary data also in closed loop if desired 155 | if config.testing.feed_boundary_data: 156 | net_input_steps[:, :, :, 0] = net_inputs[:, :, t_start:t, 0] 157 | net_input_steps[:, :, :, -1] = net_inputs[:, :, t_start:t, -1] 158 | 159 | net_output = model.forward(x=net_input_steps) 160 | net_outputs[:, :, t] = net_output[:, :, -1] 161 | 162 | if print_progress: 163 | forward_pass_duration = time.time() - time_start 164 | print("Forward pass took:", forward_pass_duration, "seconds.") 165 | 166 | # Convert the PyTorch network output tensor into a numpy array 167 | net_outputs = net_outputs.cpu().detach().numpy()[0, 0] 168 | net_labels = net_labels.cpu().detach().numpy()[0, 0] 169 | 170 | # 171 | # Visualize the data 172 | if visualize: 173 | plt.style.use("dark_background") 174 | 175 | # Plot over space and time 176 | fig, ax = plt.subplots(1, 2, figsize=(16, 6), sharey=True) 177 | 178 | if config.data.type == "burger" or\ 179 | config.data.type == "diffusion_sorption" or\ 180 | config.data.type == "allen_cahn": 181 | 182 | im1 = ax[0].imshow( 183 | np.transpose(net_labels), interpolation='nearest', 184 | cmap='rainbow', origin='lower', aspect='auto', vmin=-0.4, 185 | vmax=0.4 186 | ) 187 | fig.colorbar(im1, ax=ax[0]) 188 | im2 = ax[1].imshow( 189 | np.transpose(net_outputs), interpolation='nearest', 190 | cmap='rainbow', origin='lower', aspect='auto', vmin=-0.4, 191 | vmax=0.4 192 | ) 193 | fig.colorbar(im2, ax=ax[1]) 194 | 195 | ax[0].set_xlabel("t") 196 | ax[0].set_ylabel("x") 197 | ax[1].set_xlabel("t") 198 | 199 | elif config.data.type == "diffusion_reaction": 200 | 201 | im1 = ax[0].imshow( 202 | np.transpose(net_labels[..., 0]), interpolation='nearest', 203 | cmap='rainbow', origin='lower', aspect='auto', vmin=-0.4, 204 | vmax=0.4 205 | ) 206 | fig.colorbar(im1, ax=ax[0]) 207 | im2 = ax[1].imshow( 208 | np.transpose(net_outputs[..., 0]), interpolation='nearest', 209 | cmap='rainbow', origin='lower', aspect='auto', vmin=-0.4, 210 | vmax=0.4 211 | ) 212 | fig.colorbar(im2, ax=ax[1]) 213 | 214 | ax[0].set_xlabel("x") 215 | ax[0].set_ylabel("y") 216 | ax[1].set_xlabel("x") 217 | 218 | elif config.data.type == "burger_2d": 219 | 220 | im1 = ax[0].imshow( 221 | np.transpose(net_labels[-1, ...]), interpolation='nearest', 222 | cmap='rainbow', origin='lower', aspect='auto', vmin=-0.4, 223 | vmax=0.4 224 | ) 225 | fig.colorbar(im1, ax=ax[0]) 226 | im2 = ax[1].imshow( 227 | np.transpose(net_outputs[-1, ...]), interpolation='nearest', 228 | cmap='rainbow', origin='lower', aspect='auto', vmin=-0.4, 229 | vmax=0.4 230 | ) 231 | fig.colorbar(im2, ax=ax[1]) 232 | 233 | ax[0].set_xlabel("x") 234 | ax[0].set_ylabel("y") 235 | ax[1].set_xlabel("x") 236 | 237 | 238 | ax[0].set_title("Ground Truth") 239 | ax[1].set_title("Network Output") 240 | 241 | 242 | if config.data.type == "diffusion_reaction"\ 243 | or config.data.type == "burger_2d": 244 | anim = animation.FuncAnimation( 245 | fig, 246 | animate, 247 | frames=sequence_length, 248 | fargs=(im1, im2, net_labels, net_outputs), 249 | interval=20 250 | ) 251 | 252 | plt.show() 253 | 254 | # Compute error 255 | mse = np.mean(np.square(net_outputs - net_labels)) 256 | 257 | return mse 258 | 259 | 260 | def animate(t, im1, im2, net_label, net_outputs): 261 | """ 262 | Data animation function animating an image over time. 263 | :param t: The current time step 264 | :param axis: The matplotlib image object 265 | :param field: The data field 266 | :return: The matplotlib image object updated with the current time step's 267 | image date 268 | """ 269 | im1.set_array(net_label[t]) 270 | im2.set_array(net_outputs[t]) 271 | 272 | 273 | if __name__ == "__main__": 274 | mse = run_testing(print_progress=True, visualize=True) 275 | print(f"MSE: {mse}") 276 | 277 | print("Done.") -------------------------------------------------------------------------------- /models/tcn/train.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch as th 3 | import torch.nn as nn 4 | import time 5 | import glob 6 | import os 7 | import matplotlib.pyplot as plt 8 | from threading import Thread 9 | import sys 10 | 11 | sys.path.append("..") 12 | from utils.configuration import Configuration 13 | import utils.helper_functions as helpers 14 | from tcn import TemporalConvNet 15 | 16 | 17 | def run_training(print_progress=False, model_number=None): 18 | 19 | # Set a random seed for varying weight initializations 20 | th.seed() 21 | 22 | th.set_num_threads(1) 23 | 24 | # Load the user configurations 25 | config = Configuration("config.json") 26 | 27 | # Append the model number to the name of the model 28 | if model_number is None: 29 | model_number = config.model.number 30 | config.model.name = config.model.name + "_" + str(model_number).zfill(2) 31 | 32 | print("Model name:", config.model.name) 33 | 34 | # Hide the GPU(s) in case the user specified to use the CPU in the config 35 | # file 36 | if config.general.device == "CPU": 37 | os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" # see issue #152 38 | os.environ["CUDA_VISIBLE_DEVICES"] = "" 39 | 40 | time_start = time.time() 41 | 42 | # setting device on GPU if available, else CPU 43 | device = helpers.determine_device() 44 | 45 | # Initialize and set up the network 46 | model = TemporalConvNet(config=config).to(device=device) 47 | 48 | if print_progress: 49 | # Count number of trainable parameters 50 | pytorch_total_params = sum( 51 | p.numel() for p in model.parameters() if p.requires_grad 52 | ) 53 | print("Trainable model parameters:", pytorch_total_params) 54 | 55 | # If desired, restore the network by loading the weights saved in the .pt 56 | # file 57 | if config.training.continue_training: 58 | print("Restoring model (that is the network\"s weights) from file...") 59 | model.load_state_dict(th.load(os.path.join(os.path.abspath(""), 60 | "checkpoints", 61 | config.model.name, 62 | config.model.name + ".pt"))) 63 | model.train() 64 | 65 | # 66 | # Set up the optimizer and the criterion (loss) 67 | optimizer = th.optim.LBFGS(model.parameters(), 68 | lr=config.training.learning_rate) 69 | criterion = nn.MSELoss(reduction="mean") 70 | 71 | 72 | # 73 | # Load data depending on the task 74 | if config.data.type == "burger": 75 | data_path = os.path.join("../../data/", 76 | config.data.type, 77 | config.data.name, 78 | "sample.npy") 79 | data = np.array(np.load(data_path), dtype=np.float32) 80 | data = np.expand_dims(data, axis=0) 81 | 82 | elif config.data.type == "diffusion_sorption": 83 | data_path_base = os.path.join("../../data/", 84 | config.data.type, 85 | config.data.name) 86 | data_path_c = os.path.join(data_path_base, "sample_c.npy") 87 | data_path_ct = os.path.join(data_path_base, "sample_ct.npy") 88 | data_c = np.array(np.load(data_path_c), dtype=np.float32) 89 | data_ct = np.array(np.load(data_path_ct), dtype=np.float32) 90 | data = np.stack((data_c, data_ct), axis=0) 91 | 92 | elif config.data.type == "diffusion_reaction": 93 | data_path_base = os.path.join("../../data/", 94 | config.data.type, 95 | config.data.name) 96 | data_path_u = os.path.join(data_path_base, "sample_u.npy") 97 | data_path_v = os.path.join(data_path_base, "sample_v.npy") 98 | data_u = np.array(np.load(data_path_u), dtype=np.float32) 99 | data_v = np.array(np.load(data_path_v), dtype=np.float32) 100 | data = np.stack((data_u, data_v), axis=0) 101 | 102 | elif config.data.type == "allen_cahn": 103 | data_path = os.path.join("../../data/", 104 | config.data.type, 105 | config.data.name, 106 | "sample.npy") 107 | data = np.array(np.load(data_path), dtype=np.float32) 108 | data = np.expand_dims(data, axis=0) 109 | 110 | elif config.data.type == "burger_2d": 111 | data_path = os.path.join("../../data/", 112 | config.data.type, 113 | config.data.name, 114 | "sample.npy") 115 | data = np.array(np.load(data_path), dtype=np.float32) 116 | data = np.expand_dims(data, axis=0) 117 | 118 | # Set up the training and validation datasets and -loaders 119 | data_train = th.tensor( 120 | data[:, :config.training.t_stop], 121 | device=device 122 | ).unsqueeze(0) 123 | data_valid = th.tensor( 124 | data[:, config.validation.t_start:config.validation.t_stop], 125 | device=device 126 | ).unsqueeze(0) 127 | 128 | # 129 | # Set up lists to save and store the epoch errors 130 | epoch_errors_train = [] 131 | epoch_errors_valid = [] 132 | best_train = np.infty 133 | best_valid = np.infty 134 | 135 | """ 136 | TRAINING 137 | """ 138 | 139 | a = time.time() 140 | 141 | # 142 | # Start the training and iterate over all epochs 143 | for epoch in range(config.training.epochs): 144 | 145 | epoch_start_time = time.time() 146 | 147 | # Separate the data into network inputs and labels 148 | net_input = data_train[:, :, :-1] 149 | net_label = data_train[:, :, 1:] 150 | 151 | def closure(): 152 | # Set the model to train mode 153 | model.train() 154 | 155 | # Reset the optimizer to clear data from previous iterations 156 | optimizer.zero_grad() 157 | 158 | # Forward the input through the network 159 | net_outputs = model.forward(x=net_input) 160 | 161 | # Comput the error 162 | mse = criterion(net_outputs, net_label) 163 | 164 | mse.backward() 165 | 166 | return mse 167 | 168 | # Backpropagate the error and perform a weight update 169 | optimizer.step(closure) 170 | 171 | # Extract the MSE value from the closure function 172 | mse = closure() 173 | 174 | epoch_errors_train.append(mse.item()) 175 | 176 | train_sign = "(-)" 177 | if epoch_errors_train[-1] < best_train: 178 | best_train = epoch_errors_train[-1] 179 | train_sign = "(+)" 180 | 181 | # 182 | # Validation 183 | 184 | # Separate the data into network inputs and labels 185 | net_input = data_valid[:, :, :-1] 186 | net_label = data_valid[:, :, 1:] 187 | 188 | # Forward the input through the network 189 | net_outputs = model.forward(x=net_input) 190 | 191 | # Comput the error 192 | mse = criterion(net_outputs, net_label) 193 | epoch_errors_valid.append(mse.item()) 194 | 195 | # Save the model to file (if desired) 196 | if config.training.save_model \ 197 | and mse.item() < best_valid: 198 | 199 | # Start a separate thread to save the model 200 | thread = Thread(target=helpers.save_model_to_file( 201 | model_src_path=os.path.abspath(""), 202 | config=config, 203 | epoch=epoch, 204 | epoch_errors_train=epoch_errors_train, 205 | epoch_errors_valid=epoch_errors_valid, 206 | net=model)) 207 | thread.start() 208 | 209 | # Create a plus or minus sign for the validation error 210 | valid_sign = "(-)" 211 | if epoch_errors_valid[-1] < best_valid: 212 | best_valid = epoch_errors_valid[-1] 213 | valid_sign = "(+)" 214 | 215 | # 216 | # Print progress to the console 217 | if print_progress: 218 | print(f"Epoch {str(epoch+1).zfill(int(np.log10(config.training.epochs))+1)}/{str(config.training.epochs)} took {str(np.round(time.time() - epoch_start_time, 2)).ljust(5, '0')} seconds. \t\tAverage epoch training error: {train_sign}{str(np.round(epoch_errors_train[-1], 10)).ljust(12, ' ')} \t\tValidation error: {valid_sign}{str(np.round(epoch_errors_valid[-1], 10)).ljust(12, ' ')}") 219 | 220 | if print_progress: 221 | b = time.time() 222 | print("\nTraining took " + str(np.round(b - a, 2)) + " seconds.\n\n") 223 | 224 | 225 | if __name__ == "__main__": 226 | run_training(print_progress=True) 227 | 228 | print("Done.") 229 | -------------------------------------------------------------------------------- /models/utils/configuration.py: -------------------------------------------------------------------------------- 1 | import json 2 | import jsmin 3 | import os 4 | 5 | 6 | class Dict(dict): 7 | """ 8 | Dictionary that allows to access per attributes and to except names from being loaded 9 | """ 10 | def __init__(self, dictionary: dict = None): 11 | super(Dict, self).__init__() 12 | 13 | if dictionary is not None: 14 | self.load(dictionary) 15 | 16 | def __getattr__(self, item): 17 | try: 18 | return self[item] if item in self else getattr(super(Dict, self), item) 19 | except AttributeError: 20 | raise AttributeError(f'This dictionary has no attribute "{item}"') 21 | 22 | def load(self, dictionary: dict, name_list: list = None): 23 | """ 24 | Loads a dictionary 25 | :param dictionary: Dictionary to be loaded 26 | :param name_list: List of names to be updated 27 | """ 28 | for name in dictionary: 29 | data = dictionary[name] 30 | if name_list is None or name in name_list: 31 | if isinstance(data, dict): 32 | if name in self: 33 | self[name].load(data) 34 | else: 35 | self[name] = Dict(data) 36 | elif isinstance(data, list): 37 | self[name] = list() 38 | for item in data: 39 | if isinstance(item, dict): 40 | self[name].append(Dict(item)) 41 | else: 42 | self[name].append(item) 43 | else: 44 | self[name] = data 45 | 46 | def save(self, path): 47 | """ 48 | Saves the dictionary into a json file 49 | :param path: Path of the json file 50 | """ 51 | if not os.path.exists(path): 52 | os.makedirs(path) 53 | 54 | path = os.path.join(path, 'cfg.json') 55 | 56 | with open(path, 'w') as file: 57 | json.dump(self, file, indent=True) 58 | 59 | 60 | class Configuration(Dict): 61 | """ 62 | Configuration loaded from a json file 63 | """ 64 | def __init__(self, path: str, default_path=None): 65 | super(Configuration, self).__init__() 66 | 67 | if default_path is not None: 68 | self.load(default_path) 69 | 70 | self.load(path) 71 | 72 | def load_model(self, path: str): 73 | self.load(path, name_list=["model"]) 74 | 75 | def load(self, path: str, name_list: list = None): 76 | """ 77 | Loads attributes from a json file 78 | :param path: Path of the json file 79 | :param name_list: List of names to be updated 80 | :return: 81 | """ 82 | with open(path) as file: 83 | data = json.loads(jsmin.jsmin(file.read())) 84 | 85 | super(Configuration, self).load(data, name_list) 86 | 87 | -------------------------------------------------------------------------------- /models/utils/datasets.py: -------------------------------------------------------------------------------- 1 | """ 2 | This file contains the organization of the custom datasets such that it can be 3 | read efficiently in combination with the DataLoader from PyTorch to prevent that 4 | data reading and preparing becomes the bottleneck. 5 | 6 | This script was inspired by 7 | https://stanford.edu/~shervine/blog/pytorch-how-to-generate-data-parallel 8 | """ 9 | 10 | from torch.utils import data 11 | import numpy as np 12 | import os 13 | import glob 14 | 15 | 16 | class DatasetBurger(data.Dataset): 17 | """ 18 | The dataset class which can be used with PyTorch's DataLoader. 19 | """ 20 | 21 | def __init__(self, root_path, dataset_type, dataset_name, mode): 22 | """ 23 | Constructor class setting up the data loader 24 | :param root_path: The root path of the project 25 | :param dataset_type: The type of the dataset (e.g. "burgers") 26 | :param dataset_name: The name of the dataset (e.g. "d0.01_pi") 27 | :param mode: Any of "train", "val" or "test" 28 | """ 29 | 30 | # Determine the path to the data and the file paths to all samples 31 | data_root_path = os.path.join( 32 | root_path, dataset_type, dataset_name, mode 33 | ) 34 | self.data_paths = np.sort( 35 | glob.glob(os.path.join(data_root_path, "*.npy")) 36 | ) 37 | 38 | def __len__(self): 39 | return len(self.data_paths) 40 | 41 | def __getitem__(self, index): 42 | """ 43 | Generates a sample batch in the form [batch_size, dim], where dim is 44 | the number of features, here (t, x). 45 | :param index: The index of the sample in the path array 46 | :return: One batch of data as np.array 47 | """ 48 | return np.load(self.data_paths[index]) 49 | 50 | 51 | class DatasetDiffSorp(data.Dataset): 52 | """ 53 | The dataset class which can be used with PyTorch's DataLoader. 54 | """ 55 | 56 | def __init__(self, root_path, dataset_type, dataset_name, mode): 57 | """ 58 | Constructor class setting up the data loader 59 | :param root_path: The root path of the project 60 | :param dataset_type: The type of the dataset (e.g. "burgers") 61 | :param dataset_name: The name of the dataset (e.g. "d0.01_pi") 62 | :param mode: Any of "train", "val" or "test" 63 | """ 64 | 65 | # Determine the path to the data and the file paths to all samples 66 | data_root_path = os.path.join( 67 | root_path, dataset_type, dataset_name, mode 68 | ) 69 | self.data_paths = np.sort( 70 | glob.glob(os.path.join(data_root_path, "*.npy")) 71 | ) 72 | 73 | def __len__(self): 74 | return len(self.data_paths) 75 | 76 | def __getitem__(self, index): 77 | """ 78 | Generates a sample batch in the form [batch_size, dim], where dim is 79 | the number of features, here (t, x). 80 | :param index: The index of the sample in the path array 81 | :return: One batch of data as np.array 82 | """ 83 | return np.load(self.data_paths[index]) 84 | 85 | 86 | class DatasetDiffReact(data.Dataset): 87 | """ 88 | The dataset class which can be used with PyTorch's DataLoader. 89 | """ 90 | 91 | def __init__(self, root_path, dataset_type, dataset_name, mode): 92 | """ 93 | Constructor class setting up the data loader 94 | :param root_path: The root path of the project 95 | :param dataset_type: The type of the dataset (e.g. "burgers") 96 | :param dataset_name: The name of the dataset (e.g. "d0.01_pi") 97 | :param mode: Any of "train", "val" or "test" 98 | """ 99 | 100 | # Determine the path to the data and the file paths to all samples 101 | data_root_path = os.path.join( 102 | root_path, dataset_type, dataset_name, mode 103 | ) 104 | self.data_paths = np.sort( 105 | glob.glob(os.path.join(data_root_path, "*.npy")) 106 | ) 107 | 108 | def __len__(self): 109 | return len(self.data_paths) 110 | 111 | def __getitem__(self, index): 112 | """ 113 | Generates a sample batch in the form [batch_size, dim], where dim is 114 | the number of features, here (t, x). 115 | :param index: The index of the sample in the path array 116 | :return: One batch of data as np.array 117 | """ 118 | return np.load(self.data_paths[index]) 119 | 120 | class DatasetAllenCahn(data.Dataset): 121 | """ 122 | The dataset class which can be used with PyTorch's DataLoader. 123 | """ 124 | 125 | def __init__(self, root_path, dataset_type, dataset_name, mode): 126 | """ 127 | Constructor class setting up the data loader 128 | :param root_path: The root path of the project 129 | :param dataset_type: The type of the dataset (e.g. "burgers") 130 | :param dataset_name: The name of the dataset (e.g. "d0.01_pi") 131 | :param mode: Any of "train", "val" or "test" 132 | """ 133 | 134 | # Determine the path to the data and the file paths to all samples 135 | data_root_path = os.path.join( 136 | root_path, dataset_type, dataset_name, mode 137 | ) 138 | self.data_paths = np.sort( 139 | glob.glob(os.path.join(data_root_path, "*.npy")) 140 | ) 141 | 142 | def __len__(self): 143 | return len(self.data_paths) 144 | 145 | def __getitem__(self, index): 146 | """ 147 | Generates a sample batch in the form [batch_size, dim], where dim is 148 | the number of features, here (t, x). 149 | :param index: The index of the sample in the path array 150 | :return: One batch of data as np.array 151 | """ 152 | return np.load(self.data_paths[index]) 153 | 154 | 155 | class DatasetBurger2D(data.Dataset): 156 | """ 157 | The dataset class which can be used with PyTorch's DataLoader. 158 | """ 159 | 160 | def __init__(self, root_path, dataset_type, dataset_name, mode): 161 | """ 162 | Constructor class setting up the data loader 163 | :param root_path: The root path of the project 164 | :param dataset_type: The type of the dataset (e.g. "burgers") 165 | :param dataset_name: The name of the dataset (e.g. "d0.01_pi") 166 | :param mode: Any of "train", "val" or "test" 167 | """ 168 | 169 | # Determine the path to the data and the file paths to all samples 170 | data_root_path = os.path.join( 171 | root_path, dataset_type, dataset_name, mode 172 | ) 173 | self.data_paths = np.sort( 174 | glob.glob(os.path.join(data_root_path, "*.npy")) 175 | ) 176 | 177 | def __len__(self): 178 | return len(self.data_paths) 179 | 180 | def __getitem__(self, index): 181 | """ 182 | Generates a sample batch in the form [batch_size, dim], where dim is 183 | the number of features, here (t, x). 184 | :param index: The index of the sample in the path array 185 | :return: One batch of data as np.array 186 | """ 187 | return np.load(self.data_paths[index]) 188 | -------------------------------------------------------------------------------- /models/utils/helper_functions.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch as th 3 | import os 4 | import time 5 | import matplotlib.pyplot as plt 6 | import matplotlib.animation as animation 7 | 8 | import utils.datasets as datasets 9 | 10 | 11 | def determine_device(print_progress=False): 12 | """ 13 | This function evaluates whether a GPU is accessible at the system and 14 | returns it as device to calculate on, otherwise it returns the CPU. 15 | :return: The device where tensor calculations shall be made on 16 | """ 17 | device = th.device("cuda" if th.cuda.is_available() else "cpu") 18 | if print_progress: 19 | print("Using device:", device, "\n") 20 | 21 | # Additional Info when using cuda 22 | if device.type == "cuda" and print_progress: 23 | print(th.cuda.get_device_name(0)) 24 | print("Memory Usage:") 25 | print("\tAllocated:", 26 | round(th.cuda.memory_allocated(0) / 1024 ** 3, 1), "GB") 27 | print("\tCached: ", round(th.cuda.memory_reserved(0) / 1024 ** 3, 1), 28 | "GB") 29 | print() 30 | 31 | return device 32 | 33 | 34 | def build_dataloader(config, mode, batch_size): 35 | """ 36 | This function creates a dataset and return the appropriate dataloader to 37 | iterate over this dataset 38 | :param config: The general configurations of the model 39 | :param mode: Any of "train", "val", or "test" 40 | :param batch_size: The number of samples per batch 41 | :return: A PyTorch dataloader object 42 | """ 43 | 44 | # Set up a dataset and dataloader 45 | if config.data.type == "burger": 46 | dataset = datasets.DatasetBurger( 47 | root_path=os.path.abspath("../../data"), 48 | dataset_type=config.data.type, 49 | dataset_name=config.data.name, 50 | mode=mode 51 | ) 52 | elif config.data.type == "diffusion_sorption": 53 | dataset = datasets.DatasetDiffSorp( 54 | root_path=os.path.abspath("../../data"), 55 | dataset_type=config.data.type, 56 | dataset_name=config.data.name, 57 | mode=mode 58 | ) 59 | elif config.data.type == "diffusion_reaction": 60 | dataset = datasets.DatasetDiffReact( 61 | root_path=os.path.abspath("../../data"), 62 | dataset_type=config.data.type, 63 | dataset_name=config.data.name, 64 | mode=mode 65 | ) 66 | elif config.data.type == "allen_cahn": 67 | dataset = datasets.DatasetAllenCahn( 68 | root_path=os.path.abspath("../../data"), 69 | dataset_type=config.data.type, 70 | dataset_name=config.data.name, 71 | mode=mode 72 | ) 73 | elif config.data.type == "burger_2d": 74 | dataset = datasets.DatasetBurger2D( 75 | root_path=os.path.abspath("../../data"), 76 | dataset_type=config.data.type, 77 | dataset_name=config.data.name, 78 | mode=mode 79 | ) 80 | 81 | dataloader = th.utils.data.DataLoader( 82 | dataset=dataset, 83 | batch_size=batch_size, 84 | shuffle=True, 85 | num_workers=2, 86 | pin_memory=True 87 | ) 88 | 89 | return dataloader 90 | 91 | 92 | def save_model_to_file(model_src_path, config, epoch, epoch_errors_train, 93 | epoch_errors_valid, net): 94 | """ 95 | This function writes the model weights along with the network configuration 96 | and current performance to file. 97 | :param model_src_path: The source path where the model will be saved to 98 | :param config: The configurations of the model 99 | :param epoch: The current epoch 100 | :param epoch_errors_train: The training epoch errors 101 | :param epoch_errors_valid: The validation epoch errors, 102 | :param net: The actual model 103 | :return: Nothing 104 | """ 105 | 106 | model_save_path = os.path.join( 107 | model_src_path, "checkpoints", config.model.name 108 | ) 109 | 110 | os.makedirs(model_save_path, exist_ok="True") 111 | 112 | # Save model weights to file 113 | th.save(net.state_dict(), 114 | os.path.join(model_save_path, config.model.name + ".pt")) 115 | 116 | # Copy the configurations and add a results entry 117 | config["results"] = { 118 | "current_epoch": epoch + 1, 119 | "current_training_error": epoch_errors_train[-1], 120 | "lowest_train_error": min(epoch_errors_train), 121 | "current_validation_error": epoch_errors_valid[-1], 122 | "lowest_validation_error": min(epoch_errors_valid) 123 | } 124 | 125 | # Save the configuration and current performance to file 126 | config.save(model_save_path) 127 | -------------------------------------------------------------------------------- /sampling/checkpoints/diff-sorp_00/cfg.json: -------------------------------------------------------------------------------- 1 | { 2 | "paths": {}, 3 | "general": { 4 | "device": "cuda" 5 | }, 6 | "training": { 7 | "save_model": true, 8 | "continue_training": false, 9 | "epochs": 100, 10 | "learning_rate": 0.1 11 | }, 12 | "data": { 13 | "type": "diffusion_sorption", 14 | "name": "experiment", 15 | "noise": 0.0 16 | }, 17 | "model": { 18 | "name": "diff-sorp", 19 | "number": 0, 20 | "layer_sizes": [ 21 | 1, 22 | 10, 23 | 20, 24 | 10, 25 | 1 26 | ] 27 | }, 28 | "results": { 29 | "current_epoch": 4, 30 | "current_training_error": 0.31605613231658936, 31 | "lowest_train_error": 0.31605613231658936, 32 | "current_validation_error": 0.31605613231658936, 33 | "lowest_validation_error": 0.31605613231658936 34 | } 35 | } -------------------------------------------------------------------------------- /sampling/checkpoints/diff-sorp_00/diff-sorp_00.pt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/CognitiveModeling/finn/82b4eac6b55be3b44e736f139ed50256c4415f0c/sampling/checkpoints/diff-sorp_00/diff-sorp_00.pt -------------------------------------------------------------------------------- /sampling/config.json: -------------------------------------------------------------------------------- 1 | { 2 | "paths": { 3 | }, 4 | 5 | "general": { 6 | "device": "cpu" 7 | }, 8 | 9 | "training": { 10 | "save_model": true, 11 | "continue_training": false, 12 | "epochs": 100, 13 | "learning_rate": 0.1 14 | }, 15 | 16 | "sampling": { 17 | "train": false, 18 | "sample": true, 19 | "random_init": false, 20 | "sampler": "mala", 21 | "name": "mala_sampler", 22 | "step_size": 0.00003, 23 | "noise_tol": 0.05, 24 | "num_sample": 1000 25 | }, 26 | 27 | "data": { 28 | "type": "diffusion_sorption", // "burger", "diffusion_sorption", "diffusion_reaction", "allen_cahn" 29 | "name": "experiment", // "data_train", "data_ext", "data_test" 30 | "noise": 0.0 31 | }, 32 | 33 | "model": { 34 | "name": "diff-sorp", 35 | "number": 0, // The i-th model 36 | "layer_sizes": [1, 10, 20, 10, 1] // [1, 10, 20, 10, 1] for burger, diffusion_sorption, and allen_cahn, [2, 20, 20, 20, 2] for diffusion_reaction 37 | } 38 | } -------------------------------------------------------------------------------- /sampling/data_core1.xlsx: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/CognitiveModeling/finn/82b4eac6b55be3b44e736f139ed50256c4415f0c/sampling/data_core1.xlsx -------------------------------------------------------------------------------- /sampling/data_core2.xlsx: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/CognitiveModeling/finn/82b4eac6b55be3b44e736f139ed50256c4415f0c/sampling/data_core2.xlsx -------------------------------------------------------------------------------- /sampling/data_core2_long.xlsx: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/CognitiveModeling/finn/82b4eac6b55be3b44e736f139ed50256c4415f0c/sampling/data_core2_long.xlsx -------------------------------------------------------------------------------- /sampling/init.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | """ 3 | Created on March 2021 4 | @author: Timothy Praditia 5 | 6 | This script contains the Initialize class that reads the configuration file 7 | and construct an object with the corresponding configuration parameters 8 | """ 9 | 10 | import torch as th 11 | import os 12 | import shutil 13 | import sys 14 | sys.path.append("../models") 15 | from utils.configuration import Configuration 16 | import utils.helper_functions as helpers 17 | 18 | class Initialize: 19 | 20 | def __init__(self, config): 21 | """ 22 | Constructor 23 | 24 | """ 25 | 26 | # Load the user configurations 27 | self.config = config 28 | 29 | # Append the model number to the name of the model 30 | model_number = self.config.model.number 31 | self.config.model.name = self.config.model.name + "_" + str(model_number).zfill(2) 32 | 33 | # Print some information to console 34 | print("Model name:", self.config.model.name) 35 | 36 | # Hide the GPU(s) in case the user specified to use the CPU in the config 37 | # file 38 | if self.config.general.device == "CPU": 39 | os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" 40 | os.environ["CUDA_VISIBLE_DEVICES"] = "" 41 | 42 | # Set device on GPU if specified in the configuration file, else CPU 43 | # device = helpers.determine_device() 44 | self.device = th.device(self.config.general.device) 45 | 46 | 47 | # # SET WORKING PATH 48 | # self.main_path = os.getcwd() 49 | 50 | # # MODEL NAME & SETTING 51 | # self.model_name = params.model_name 52 | 53 | # self.model_path = self.main_path + '\\' + self.model_name 54 | # self.check_dir(self.model_path) 55 | 56 | # self.log_path = self.main_path + '\\runs\\' + self.model_name 57 | # # Remove old log files to prevent unclear visualization in tensorboard 58 | # self.check_dir(self.log_path, remove=True) 59 | 60 | # self.save_model = params.save_model 61 | # self.continue_training = params.continue_training 62 | # self.device_name = params.device_name 63 | # self.device = self.determine_device() 64 | 65 | # # NETWORK HYPER-PARAMETERS 66 | # self.flux_layers = params.flux_layers 67 | # self.state_layers = params.state_layers 68 | # self.flux_nodes = params.flux_nodes 69 | # self.state_nodes = params.state_nodes 70 | # self.learning_rate = params.learning_rate 71 | # self.error_mult = params.error_mult 72 | # self.phys_mult = params.phys_mult 73 | # self.epochs = params.epochs 74 | # self.lbfgs_optim = params.lbfgs_optim 75 | 76 | # # SIMULATION-RELATED INPUTS 77 | # self.num_vars = params.num_vars -------------------------------------------------------------------------------- /sampling/set_module.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | """ 3 | Created on March 2021 4 | @author: Timothy Praditia 5 | 6 | This script contains the Set_Module class that constructs separate objects for 7 | different core samples, based on the parameters read from the Excel input file 8 | """ 9 | 10 | import torch 11 | import numpy as np 12 | import pandas as pd 13 | 14 | 15 | class Set_Module: 16 | 17 | def __init__(self, filename, params): 18 | """ 19 | Constructor 20 | 21 | Inputs: 22 | filename : the corresponding filename for the core sample 23 | params : the configuration object containing the model settings 24 | 25 | """ 26 | 27 | # Load parameters from the Excel file 28 | in_params = pd.read_excel(filename, sheet_name=1, index_col=0, header=None) 29 | 30 | # Determine the device on which the training takes place 31 | self.device = params.device 32 | 33 | self.config = params.config 34 | 35 | # Soil Parameters 36 | self.D = in_params[1]['D'] 37 | self.por = in_params[1]['por'] 38 | self.rho_s = in_params[1]['rho_s'] 39 | 40 | 41 | # Simulation Domain 42 | self.X = in_params[1]['X'] 43 | self.Nx = int(in_params[1]['Nx']) 44 | self.dx = self.X / (self.Nx+1) 45 | self.T = in_params[1]['T'] 46 | self.Nt = int(in_params[1]['Nt']) 47 | self.r = in_params[1]['sample_radius'] 48 | self.A = np.pi * self.r**2 49 | self.Q = in_params[1]['Q'] 50 | self.solubility = in_params[1]['solubility'] 51 | self.cauchy_mult = self.por * self.A / self.Q 52 | 53 | ## Effective diffusion coefficient for each variable 54 | self.D_eff = np.array([self.D / (self.dx**2), 55 | self.D * self.por / (self.rho_s/1000) / (self.dx**2)]) 56 | -------------------------------------------------------------------------------- /sampling/train.py: -------------------------------------------------------------------------------- 1 | #! env/bin/python3 2 | 3 | """ 4 | Main file for training a model with FINN 5 | """ 6 | 7 | import numpy as np 8 | import torch as th 9 | import torch.nn as nn 10 | import os 11 | import time 12 | from threading import Thread 13 | import sys 14 | 15 | sys.path.append("../models") 16 | from utils.configuration import Configuration 17 | import utils.helper_functions as helpers 18 | 19 | 20 | 21 | def run_training(x, t, u0, data, model, cfg, print_progress=True): 22 | 23 | # Set device on GPU if specified in the configuration file, else CPU 24 | # device = helpers.determine_device() 25 | config = cfg.config 26 | device = config.general.device 27 | model = model.to(device) 28 | x = x.to(device) 29 | t = t.to(device) 30 | u0 = u0.to(device) 31 | data = data.to(device) 32 | 33 | # Count number of trainable parameters 34 | pytorch_total_params = sum( 35 | p.numel() for p in model.parameters() if p.requires_grad 36 | ) 37 | 38 | if print_progress: 39 | print("Trainable model parameters:", pytorch_total_params) 40 | 41 | # If desired, restore the network by loading the weights saved in the .pt 42 | # file 43 | if config.training.continue_training: 44 | if print_progress: 45 | print('Restoring model (that is the network\'s weights) from file...') 46 | model.load_state_dict(th.load(os.path.join(os.path.abspath(""), 47 | "checkpoints", 48 | config.model.name, 49 | config.model.name + ".pt"))) 50 | model.train() 51 | 52 | # 53 | # Set up an optimizer and the criterion (loss) 54 | optimizer = th.optim.LBFGS(model.parameters(), 55 | lr=config.training.learning_rate) 56 | 57 | criterion = nn.MSELoss(reduction="sum") 58 | 59 | # 60 | # Set up lists to save and store the epoch errors 61 | epoch_errors_train = [] 62 | 63 | best_train = np.infty 64 | 65 | """ 66 | TRAINING 67 | """ 68 | 69 | a = time.time() 70 | 71 | # 72 | # Start the training and iterate over all epochs 73 | for epoch in range(config.training.epochs): 74 | 75 | epoch_start_time = time.time() 76 | 77 | # Define the closure function that consists of resetting the 78 | # gradient buffer, loss function calculation, and backpropagation 79 | # It is necessary for LBFGS optimizer, because it requires multiple 80 | # function evaluations 81 | def closure(): 82 | 83 | # Set the model to train mode 84 | model.train() 85 | 86 | # Reset the optimizer to clear data from previous iterations 87 | optimizer.zero_grad() 88 | 89 | # Forward propagate and calculate loss function 90 | u_hat = model(t=t, u=u0) 91 | 92 | # Calculate predicted breakthrough curve 93 | cauchy_mult = cfg.cauchy_mult * cfg.D_eff[0] * cfg.dx 94 | pred = ((u_hat[:,-2,0] - u_hat[:,-1,0]) * cauchy_mult).squeeze() 95 | 96 | loss_data = criterion(1e3 * pred, 1e3 * data) 97 | 98 | # Calculate physical loss based on retardation factor monotonicity 99 | inp_temp = th.linspace(0,2,501).unsqueeze(-1).to(device) 100 | ret_temp = 1/(model.func_nn(inp_temp))[...,0] 101 | 102 | loss_phys = 100*th.mean(th.relu(ret_temp[1:]-ret_temp[:-1])) 103 | 104 | loss = loss_data + loss_phys 105 | 106 | loss.backward() 107 | 108 | if print_progress: 109 | print(loss_data.item(), loss_phys.item(), loss.item()) 110 | 111 | return loss 112 | 113 | optimizer.step(closure) 114 | 115 | # Extract the MSE value from the closure function 116 | loss = closure() 117 | 118 | epoch_errors_train.append(loss.item()) 119 | 120 | # Create a plus or minus sign for the training error 121 | train_sign = "(-)" 122 | if epoch_errors_train[-1] < best_train: 123 | train_sign = "(+)" 124 | best_train = epoch_errors_train[-1] 125 | # Save the model to file (if desired) 126 | if config.training.save_model: 127 | # Start a separate thread to save the model 128 | thread = Thread(target=helpers.save_model_to_file( 129 | model_src_path=os.path.abspath(""), 130 | config=config, 131 | epoch=epoch, 132 | epoch_errors_train=epoch_errors_train, 133 | epoch_errors_valid=epoch_errors_train, 134 | net=model)) 135 | thread.start() 136 | 137 | # 138 | # Print progress to the console 139 | if print_progress: 140 | print(f"Epoch {str(epoch+1).zfill(int(np.log10(config.training.epochs))+1)}/{str(config.training.epochs)} took {str(np.round(time.time() - epoch_start_time, 2)).ljust(5, '0')} seconds. \t\tAverage epoch training error: {train_sign}{str(np.round(epoch_errors_train[-1], 10)).ljust(12, ' ')}") 141 | 142 | b = time.time() 143 | if print_progress: 144 | print('\nTraining took ' + str(np.round(b - a, 2)) + ' seconds.\n\n') --------------------------------------------------------------------------------