├── docs ├── api │ ├── data │ │ ├── base.md │ │ ├── kdv.md │ │ └── burgers.md │ ├── lib.md │ ├── deepmod.md │ ├── constraint.md │ ├── func_approx.md │ ├── sparse.md │ ├── training.md │ ├── analysis.md │ ├── convergence.md │ └── spars_sched.md ├── .DS_Store ├── figures │ ├── .DS_Store │ ├── framework.png │ ├── DeepMoD_logo.png │ └── data_workflow_for_deepymod.png ├── examples │ ├── 2DAD │ │ ├── output_7_0.png │ │ └── 2DAD.md │ ├── Burgers │ │ ├── output_6_0.png │ │ ├── output_14_0.png │ │ └── PDE_Burgers.md │ ├── PDE_KdV │ │ ├── output_6_0.png │ │ ├── output_14_0.png │ │ └── PDE_KdV.md │ └── ODE_nonlin │ │ ├── output_8_0.png │ │ └── ODE.md ├── css │ └── mkdocstrings.css ├── index.md ├── background │ └── background.md └── datasets │ └── data.md ├── src ├── deepymod │ ├── training │ │ ├── __init__.py │ │ ├── convergence.py │ │ ├── training.py │ │ └── sparsity_scheduler.py │ ├── data │ │ ├── kdv │ │ │ ├── __init__.py │ │ │ └── kdv.py │ │ ├── __init__.py │ │ ├── burgers │ │ │ ├── __init__.py │ │ │ └── burgers.py │ │ ├── diffusion │ │ │ ├── __init__.py │ │ │ └── diffusion.py │ │ ├── examples.py │ │ ├── samples.py │ │ └── base.py │ ├── utils │ │ ├── types.py │ │ ├── utilities.py │ │ └── logger.py │ ├── analysis │ │ ├── __init__.py │ │ └── load_tensorboard.py │ ├── __init__.py │ └── model │ │ ├── constraint.py │ │ ├── func_approx.py │ │ ├── sparse_estimators.py │ │ ├── library.py │ │ └── deepmod.py └── .DS_Store ├── .DS_Store ├── examples ├── data │ ├── kdv.npy │ ├── burgers.npy │ ├── keller_segel.npy │ └── advection_diffusion.mat └── VE_datagen.py ├── .github └── workflows │ ├── lint-black.yml │ ├── python-publish.yml │ └── python-test-publish.yml ├── setup.py ├── .gitignore ├── LICENSE.txt ├── tests ├── ridge.py └── burgers.py ├── mkdocs.yml ├── setup.cfg └── README.md /docs/api/data/base.md: -------------------------------------------------------------------------------- 1 | :::deepymod.data.base -------------------------------------------------------------------------------- /docs/api/lib.md: -------------------------------------------------------------------------------- 1 | :::deepymod.model.library -------------------------------------------------------------------------------- /docs/api/data/kdv.md: -------------------------------------------------------------------------------- 1 | :::deepymod.data.kdv.kdv -------------------------------------------------------------------------------- /docs/api/deepmod.md: -------------------------------------------------------------------------------- 1 | :::deepymod.model.deepmod -------------------------------------------------------------------------------- /docs/api/constraint.md: -------------------------------------------------------------------------------- 1 | :::deepymod.model.constraint -------------------------------------------------------------------------------- /docs/api/func_approx.md: -------------------------------------------------------------------------------- 1 | :::deepymod.model.func_approx -------------------------------------------------------------------------------- /docs/api/sparse.md: -------------------------------------------------------------------------------- 1 | :::deepymod.model.sparse_estimators -------------------------------------------------------------------------------- /docs/api/training.md: -------------------------------------------------------------------------------- 1 | :::deepymod.training.training -------------------------------------------------------------------------------- /docs/api/analysis.md: -------------------------------------------------------------------------------- 1 | :::deepymod.analysis.load_tensorboard -------------------------------------------------------------------------------- /docs/api/convergence.md: -------------------------------------------------------------------------------- 1 | :::deepymod.training.convergence -------------------------------------------------------------------------------- /docs/api/data/burgers.md: -------------------------------------------------------------------------------- 1 | :::deepymod.data.burgers.burgers -------------------------------------------------------------------------------- /docs/api/spars_sched.md: -------------------------------------------------------------------------------- 1 | :::deepymod.training.sparsity_scheduler -------------------------------------------------------------------------------- /src/deepymod/training/__init__.py: -------------------------------------------------------------------------------- 1 | from .training import * 2 | -------------------------------------------------------------------------------- /.DS_Store: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/PhIMaL/DeePyMoD/HEAD/.DS_Store -------------------------------------------------------------------------------- /src/.DS_Store: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/PhIMaL/DeePyMoD/HEAD/src/.DS_Store -------------------------------------------------------------------------------- /docs/.DS_Store: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/PhIMaL/DeePyMoD/HEAD/docs/.DS_Store -------------------------------------------------------------------------------- /src/deepymod/data/kdv/__init__.py: -------------------------------------------------------------------------------- 1 | from .kdv import single_soliton, double_soliton 2 | -------------------------------------------------------------------------------- /examples/data/kdv.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/PhIMaL/DeePyMoD/HEAD/examples/data/kdv.npy -------------------------------------------------------------------------------- /docs/figures/.DS_Store: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/PhIMaL/DeePyMoD/HEAD/docs/figures/.DS_Store -------------------------------------------------------------------------------- /docs/figures/framework.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/PhIMaL/DeePyMoD/HEAD/docs/figures/framework.png -------------------------------------------------------------------------------- /examples/data/burgers.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/PhIMaL/DeePyMoD/HEAD/examples/data/burgers.npy -------------------------------------------------------------------------------- /src/deepymod/data/__init__.py: -------------------------------------------------------------------------------- 1 | from deepymod.data.base import Dataset, Loader, get_train_test_loader 2 | -------------------------------------------------------------------------------- /src/deepymod/data/burgers/__init__.py: -------------------------------------------------------------------------------- 1 | from .burgers import burgers_delta, burgers_cos, burgers_sawtooth 2 | -------------------------------------------------------------------------------- /docs/figures/DeepMoD_logo.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/PhIMaL/DeePyMoD/HEAD/docs/figures/DeepMoD_logo.png -------------------------------------------------------------------------------- /examples/data/keller_segel.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/PhIMaL/DeePyMoD/HEAD/examples/data/keller_segel.npy -------------------------------------------------------------------------------- /docs/examples/2DAD/output_7_0.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/PhIMaL/DeePyMoD/HEAD/docs/examples/2DAD/output_7_0.png -------------------------------------------------------------------------------- /src/deepymod/data/diffusion/__init__.py: -------------------------------------------------------------------------------- 1 | from .diffusion import diffusion_gaussian, advection_diffusion_gaussian_2d 2 | -------------------------------------------------------------------------------- /docs/examples/Burgers/output_6_0.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/PhIMaL/DeePyMoD/HEAD/docs/examples/Burgers/output_6_0.png -------------------------------------------------------------------------------- /docs/examples/PDE_KdV/output_6_0.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/PhIMaL/DeePyMoD/HEAD/docs/examples/PDE_KdV/output_6_0.png -------------------------------------------------------------------------------- /docs/examples/Burgers/output_14_0.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/PhIMaL/DeePyMoD/HEAD/docs/examples/Burgers/output_14_0.png -------------------------------------------------------------------------------- /docs/examples/PDE_KdV/output_14_0.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/PhIMaL/DeePyMoD/HEAD/docs/examples/PDE_KdV/output_14_0.png -------------------------------------------------------------------------------- /examples/data/advection_diffusion.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/PhIMaL/DeePyMoD/HEAD/examples/data/advection_diffusion.mat -------------------------------------------------------------------------------- /docs/examples/ODE_nonlin/output_8_0.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/PhIMaL/DeePyMoD/HEAD/docs/examples/ODE_nonlin/output_8_0.png -------------------------------------------------------------------------------- /docs/figures/data_workflow_for_deepymod.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/PhIMaL/DeePyMoD/HEAD/docs/figures/data_workflow_for_deepymod.png -------------------------------------------------------------------------------- /src/deepymod/utils/types.py: -------------------------------------------------------------------------------- 1 | """ Defines Tensorlist 2 | Tensorlist (list[torch.Tensor}): a list of torch Tensors.""" 3 | 4 | from typing import List, NewType 5 | import torch 6 | 7 | TensorList = NewType("TensorList", List[torch.Tensor]) 8 | -------------------------------------------------------------------------------- /.github/workflows/lint-black.yml: -------------------------------------------------------------------------------- 1 | # Lint the code using the black linter. 2 | name: Lint 3 | 4 | on: [push, pull_request] 5 | 6 | jobs: 7 | lint: 8 | runs-on: ubuntu-latest 9 | steps: 10 | - uses: actions/checkout@v2 11 | - uses: actions/setup-python@v2 12 | - uses: psf/black@23.10.1 13 | 14 | -------------------------------------------------------------------------------- /src/deepymod/analysis/__init__.py: -------------------------------------------------------------------------------- 1 | """ 2 | Analysis library 3 | ================ 4 | 5 | Tools to interpert the results of DeePyMoD. 6 | 7 | Tools 8 | ----- 9 | 10 | load_tensorboard convert the tensorboard files into a Pandas DataFrame. 11 | plot_history plot the training history of the model. 12 | 13 | """ 14 | from .load_tensorboard import load_tensorboard 15 | from .load_tensorboard import plot_history 16 | -------------------------------------------------------------------------------- /src/deepymod/__init__.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | from pkg_resources import get_distribution, DistributionNotFound 3 | 4 | try: 5 | # Change here if project is renamed and does not equal the package name 6 | dist_name = "DeePyMoD" 7 | __version__ = get_distribution(dist_name).version 8 | except DistributionNotFound: 9 | __version__ = "unknown" 10 | finally: 11 | del get_distribution, DistributionNotFound 12 | 13 | from .model.deepmod import * 14 | -------------------------------------------------------------------------------- /src/deepymod/data/examples.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | from deepymod.data.base import Dataset 4 | 5 | 6 | def sinc_2d(): 7 | """Output: Grid[N x M x L], Data[N x M x O], 8 | N = Coordinate dimension 0 9 | M = Coordinate dimension 1 10 | L = Input data dimension 11 | O = Output data dimension 12 | """ 13 | x0 = np.linspace(0, 2 * np.pi, 100) 14 | x1 = np.linspace(-np.pi, np.pi, 100) 15 | X0, X1 = np.meshgrid(x0, x1) 16 | y = np.sinc(X0 * X1) 17 | coords = torch.tensor(np.stack((X0, X1))) 18 | data = torch.tensor(y).unsqueeze(0) 19 | return coords, data 20 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | """ 3 | Setup file for deepymod_torch. 4 | Use setup.cfg to configure your project. 5 | 6 | This file was generated with PyScaffold 3.2. 7 | PyScaffold helps you to put up the scaffold of your new Python project. 8 | Learn more under: https://pyscaffold.org/ 9 | """ 10 | import sys 11 | 12 | from pkg_resources import require, VersionConflict 13 | from setuptools import setup 14 | 15 | try: 16 | require("setuptools>=38.3") 17 | except VersionConflict: 18 | print("Error: version of setuptools is too old (<38.3)!") 19 | sys.exit(1) 20 | 21 | 22 | if __name__ == "__main__": 23 | setup(use_pyscaffold=True) 24 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Temporary and binary files 2 | *~ 3 | *.py[cod] 4 | *.so 5 | *.cfg 6 | !.isort.cfg 7 | !setup.cfg 8 | *.orig 9 | *.log 10 | *.pot 11 | __pycache__/* 12 | .cache/* 13 | .*.swp 14 | */.ipynb_checkpoints/* 15 | 16 | # Project files 17 | .ropeproject 18 | .project 19 | .pydevproject 20 | .settings 21 | .idea 22 | tags 23 | 24 | # Package files 25 | *.egg 26 | *.eggs/ 27 | .installed.cfg 28 | *.egg-info/ 29 | 30 | # Unittest and coverage 31 | htmlcov/* 32 | .coverage 33 | .tox 34 | junit.xml 35 | coverage.xml 36 | .pytest_cache/ 37 | 38 | # Build and docs folder/files 39 | build/* 40 | dist/* 41 | sdist/* 42 | docs/_rst/* 43 | docs/_build/* 44 | cover/* 45 | MANIFEST 46 | 47 | # Other remaining stuff 48 | /.venv*/ 49 | .vscode/ 50 | .DS_Store 51 | .mypy_cache/ 52 | __pycache__ 53 | src/DeePyMoD.egg/ 54 | site/ 55 | .eggs/ 56 | *events.out.tfevents.* 57 | *.pt 58 | *.pdf 59 | *.png -------------------------------------------------------------------------------- /docs/css/mkdocstrings.css: -------------------------------------------------------------------------------- 1 | /* Indentation. */ 2 | div.doc-contents:not(.first) { 3 | padding-left: 25px; 4 | border-left: 4px solid rgba(230, 230, 230); 5 | margin-bottom: 80px; 6 | } 7 | 8 | /* Don't capitalize names. */ 9 | h5.doc-heading { 10 | text-transform: none !important; 11 | } 12 | 13 | /* Don't use vertical space on hidden ToC entries. */ 14 | .hidden-toc::before { 15 | margin-top: 0 !important; 16 | padding-top: 0 !important; 17 | } 18 | 19 | /* Don't show permalink of hidden ToC entries. */ 20 | .hidden-toc a.headerlink { 21 | display: none; 22 | } 23 | 24 | /* Avoid breaking parameters name, etc. in table cells. */ 25 | td code { 26 | word-break: normal !important; 27 | } 28 | 29 | /* For pieces of Markdown rendered in table cells. */ 30 | td p { 31 | margin-top: 0 !important; 32 | margin-bottom: 0 !important; 33 | } -------------------------------------------------------------------------------- /.github/workflows/python-publish.yml: -------------------------------------------------------------------------------- 1 | # This workflows will upload a Python Package using Twine when a release is created 2 | # For more information see: https://help.github.com/en/actions/language-and-framework-guides/using-python-with-github-actions#publishing-to-package-registries 3 | 4 | name: Upload Python Package 5 | 6 | on: 7 | release: 8 | types: [created] 9 | 10 | jobs: 11 | deploy: 12 | 13 | runs-on: ubuntu-latest 14 | 15 | steps: 16 | - uses: actions/checkout@v2 17 | - name: Set up Python 18 | uses: actions/setup-python@v2 19 | with: 20 | python-version: '3.x' 21 | - name: Install dependencies 22 | run: | 23 | python -m pip install --upgrade pip 24 | pip install setuptools wheel twine 25 | - name: Build and publish 26 | env: 27 | TWINE_USERNAME: ${{ secrets.PYPI_USERNAME }} 28 | TWINE_PASSWORD: ${{ secrets.PYPI_PASSWORD }} 29 | run: | 30 | python setup.py sdist bdist_wheel 31 | twine upload dist/* 32 | -------------------------------------------------------------------------------- /docs/index.md: -------------------------------------------------------------------------------- 1 | Documentation page for the Deep learning based Model Discovery package DeepMoD. DeePyMoD is a PyTorch-based implementation of the DeepMoD algorithm for model discovery of PDEs and ODEs.[github.com/PhIMaL/DeePyMoD](https://github.com/PhIMaL/DeePyMoD). This work is based on two papers: The original DeepMoD paper [arXiv:1904.09406](http://arxiv.org/abs/1904.09406), presenting the foundation of this neural network driven model discovery and a follow-up paper [arXiv:2011.04336](https://arxiv.org/abs/2011.04336) describing a modular plug and play framework. 2 | 3 | ## Summary 4 | ![Screenshot](figures/DeepMoD_logo.png) 5 | DeepMoD is a modular model discovery framewrok aimed at discovering the ODE/PDE underlying a spatio-temporal dataset. Essentially the framework is comprised of four components: 6 | 7 | * Function approximator, e.g. a neural network to represent the dataset, 8 | * Function library on which the model discovery is performed, 9 | * Constraint function that constrains the neural network with the obtained solution 10 | * Sparsity selection algorithm. 11 | 12 | 13 | 14 | -------------------------------------------------------------------------------- /LICENSE.txt: -------------------------------------------------------------------------------- 1 | The MIT License (MIT) 2 | 3 | Copyright (c) 2019 Gert-Jan 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /.github/workflows/python-test-publish.yml: -------------------------------------------------------------------------------- 1 | # This workflows will upload a Python Package using Twine when a release is created 2 | # For more information see: https://help.github.com/en/actions/language-and-framework-guides/using-python-with-github-actions#publishing-to-package-registries 3 | 4 | name: Upload Python Package 5 | 6 | on: 7 | pull_request: 8 | branches: 9 | - master 10 | types: 11 | - closed 12 | 13 | jobs: 14 | deploy: 15 | 16 | runs-on: ubuntu-latest 17 | 18 | steps: 19 | - uses: actions/checkout@v2 20 | - name: Set up Python 21 | uses: actions/setup-python@v2 22 | with: 23 | python-version: '3.x' 24 | - name: Install dependencies 25 | run: | 26 | python -m pip install --upgrade pip 27 | pip install setuptools wheel twine 28 | - name: Build and publish 29 | env: 30 | TWINE_USERNAME: ${{ secrets.PYPI_USERNAME_TEST }} 31 | TWINE_PASSWORD: ${{ secrets.PYPI_PASSWORD_TEST }} 32 | if: github.event.pull_request.merged == true 33 | run: | 34 | python setup.py sdist bdist_wheel 35 | twine upload --repository testpypi dist/* 36 | -------------------------------------------------------------------------------- /tests/ridge.py: -------------------------------------------------------------------------------- 1 | # %% Imports 2 | import numpy as np 3 | import torch 4 | 5 | from sklearn.linear_model import Ridge as RidgeReference 6 | from deepymod.model.constraint import Ridge 7 | 8 | from deepymod.data import Dataset 9 | from deepymod.data.burgers import BurgersDelta 10 | 11 | # %% Making dataset 12 | v = 0.1 13 | A = 1.0 14 | 15 | x = np.linspace(-3, 4, 100) 16 | t = np.linspace(0.5, 5.0, 50) 17 | x_grid, t_grid = np.meshgrid(x, t, indexing="ij") 18 | dataset = Dataset(BurgersDelta, v=v, A=A) 19 | 20 | theta = dataset.library( 21 | x_grid.reshape(-1, 1), t_grid.reshape(-1, 1), poly_order=2, deriv_order=3 22 | ) 23 | dt = dataset.time_deriv(x_grid.reshape(-1, 1), t_grid.reshape(-1, 1)) 24 | 25 | theta = theta / np.linalg.norm(theta, axis=0, keepdims=True) 26 | dt += 0.1 * np.random.randn(*dt.shape) 27 | 28 | # %% Baseline 29 | reg = RidgeReference(alpha=1e-3, fit_intercept=False) 30 | coeff_ref = reg.fit(theta, dt.squeeze()).coef_[:, None] 31 | # %% 32 | constraint = Ridge(l=1e-3) 33 | coeff_constraint = constraint.fit( 34 | [torch.tensor(theta, dtype=torch.float32)], [torch.tensor(dt, dtype=torch.float32)] 35 | )[0].numpy() 36 | # %% 37 | error = np.mean(np.abs(coeff_ref - coeff_constraint)) 38 | assert error < 1e-5, f"MAE w.r.t reference is too high: {error}" 39 | -------------------------------------------------------------------------------- /mkdocs.yml: -------------------------------------------------------------------------------- 1 | site_name: DeePyMoD 2 | nav: 3 | - Home: index.md 4 | - Background: background/background.md 5 | - API: 6 | - Model: 7 | - DeepMoD: api/deepmod.md 8 | - Function Approximators: api/func_approx.md 9 | - Libraries: api/lib.md 10 | - Constraints: api/constraint.md 11 | - Sparsity: api/sparse.md 12 | - Training: 13 | - Training: api/training.md 14 | - Sparsity Scheduler: api/spars_sched.md 15 | - Convergence: api/convergence.md 16 | - Data: 17 | - Base: api/data/base.md 18 | - Equations: 19 | - Burgers: api/data/burgers.md 20 | - Korteweg-de Vries: api/data/kdv.md 21 | - Analysis: api/analysis.md 22 | - Examples: 23 | - 2D Advection Diffusion: examples/2DAD/2DAD.md 24 | - KdV: examples/PDE_KdV/PDE_KdV.md 25 | - Burgers: examples/Burgers/PDE_Burgers.md 26 | - Non-linear ODE: examples/ODE_nonlin/ODE.md 27 | 28 | theme: 29 | name: "material" 30 | palette: 31 | primary: black 32 | 33 | plugins: 34 | - search 35 | - mkdocstrings 36 | 37 | markdown_extensions: 38 | - admonition 39 | - codehilite 40 | - pymdownx.arithmatex 41 | - pymdownx.superfences 42 | - pymdownx.tabbed 43 | 44 | extra_css: 45 | - css/mkdocstrings.css 46 | 47 | repo_url: "https://github.com/PhIMaL/DeePyMoD" 48 | site_url: "https://phimal.github.io/DeePyMoD/" 49 | 50 | extra_javascript: 51 | - https://cdnjs.cloudflare.com/ajax/libs/mathjax/2.7.0/MathJax.js?config=TeX-MML-AM_CHTML 52 | 53 | 54 | 55 | -------------------------------------------------------------------------------- /src/deepymod/utils/utilities.py: -------------------------------------------------------------------------------- 1 | """ Useful tools, such as combining string matrix product and computing the terms in the library """ 2 | from itertools import product, combinations, chain 3 | import torch 4 | import torch.nn as nn 5 | 6 | 7 | def string_matmul(list_1, list_2): 8 | """Matrix multiplication with strings.""" 9 | prod = [element[0] + element[1] for element in product(list_1, list_2)] 10 | return prod 11 | 12 | 13 | def terms_definition(poly_list, deriv_list): 14 | """Calculates which terms are in the library.""" 15 | if len(poly_list) == 1: 16 | theta = string_matmul( 17 | poly_list[0], deriv_list[0] 18 | ) # If we have a single output, we simply calculate and flatten matrix product between polynomials and derivatives to get library 19 | else: 20 | theta_uv = list( 21 | chain.from_iterable( 22 | [string_matmul(u, v) for u, v in combinations(poly_list, 2)] 23 | ) 24 | ) # calculate all unique combinations between polynomials 25 | theta_dudv = list( 26 | chain.from_iterable( 27 | [string_matmul(du, dv)[1:] for du, dv in combinations(deriv_list, 2)] 28 | ) 29 | ) # calculate all unique combinations of derivatives 30 | theta_udu = list( 31 | chain.from_iterable( 32 | [ 33 | string_matmul(u[1:], du[1:]) 34 | for u, du in product(poly_list, deriv_list) 35 | ] 36 | ) 37 | ) # calculate all unique combinations of derivatives 38 | theta = theta_uv + theta_dudv + theta_udu 39 | return theta 40 | -------------------------------------------------------------------------------- /src/deepymod/training/convergence.py: -------------------------------------------------------------------------------- 1 | """This module implements convergence criteria""" 2 | import torch 3 | 4 | 5 | class Convergence: 6 | """Implements convergence criterium. Convergence is when change in patience 7 | epochs is smaller than delta. 8 | returns a bool, such that True indicates the algorithm has converged 9 | """ 10 | 11 | def __init__(self, patience: int = 200, delta: float = 1e-3) -> None: 12 | """Implements convergence criterium. Convergence is when change in patience 13 | epochs is smaller than delta. 14 | Args: 15 | patience (int): how often to check for convergence 16 | delta (float): desired accuracy 17 | """ 18 | self.patience = patience 19 | self.delta = delta 20 | self.start_iteration = None 21 | self.start_l1 = None 22 | 23 | def __call__(self, iteration: int, l1_norm: torch.Tensor) -> bool: 24 | """ 25 | 26 | Args: 27 | epoch (int): Current epoch of the optimization 28 | l1_norm (torch.Tensor): Value of the L1 norm 29 | """ 30 | converged = False # overwrite later 31 | 32 | # Initialize if doesn't exist 33 | if self.start_l1 is None: 34 | self.start_l1 = l1_norm 35 | self.start_iteration = iteration 36 | 37 | # Check if change is smaller than delta and if we've exceeded patience 38 | elif torch.abs(self.start_l1 - l1_norm).item() < self.delta: 39 | if (iteration - self.start_iteration) >= self.patience: 40 | converged = True 41 | 42 | # If not, reset and keep going 43 | else: 44 | self.start_l1 = l1_norm 45 | self.start_iteration = iteration 46 | 47 | return converged 48 | -------------------------------------------------------------------------------- /src/deepymod/data/diffusion/diffusion.py: -------------------------------------------------------------------------------- 1 | """ Contains several interactive datasets for the Diffusion equation including: 2 | - Diffusion 3 | - Advection-Diffusion in 2 dimensions 4 | """ 5 | 6 | from numpy import pi 7 | import torch 8 | 9 | 10 | def diffusion_gaussian( 11 | x: torch.tensor, t: torch.tensor, D: float, x0: float, sigma: float 12 | ) -> torch.tensor: 13 | """Function to generate the solution to the 1D diffusion equation. 14 | 15 | REFERENCE 16 | 17 | Args: 18 | x ([Tensor]): Input vector of spatial coordinates. 19 | t ([Tensor]): Input vector of temporal coordinates. 20 | D (Float): Diffusion coefficient 21 | x0 (Float): Spatial coordinate where the gaussian is centered 22 | sigma (Float): Scale parameter that adjusts the initial shape of the parameter 23 | 24 | Returns: 25 | [Tensor]: solution. 26 | """ 27 | u = (2 * pi * sigma**2 + 4 * pi * D * t) ** (-1 / 2) * torch.exp( 28 | -((x - x0) ** 2) / (2 * sigma**2 + 4 * D * t) 29 | ) 30 | coords = torch.cat((t.reshape(-1, 1), x.reshape(-1, 1)), dim=1) 31 | return coords, u.view(-1, 1) 32 | 33 | 34 | def advection_diffusion_gaussian_2d( 35 | x: torch.tensor, 36 | t: torch.tensor, 37 | D: float, 38 | x0: torch.tensor, 39 | sigma: float, 40 | v: torch.tensor, 41 | ) -> torch.tensor: 42 | """Function to generate the solution to the 2D diffusion equation. 43 | 44 | REFERENCE 45 | 46 | Args: 47 | x ([Tensor]): [N, 2] Input vector of spatial coordinates. 48 | t ([Tensor]): Input vector of temporal coordinates. 49 | D (Float): Diffusion coefficient 50 | x0 ([Tensor]): Spatial coordinate where the gaussian is centered 51 | sigma (Float): Scale parameter that adjusts the initial shape of the parameter 52 | v ([Tensor]): [2] Initial velocity of the gaussian. 53 | 54 | Returns: 55 | [Tensor]: solution. 56 | """ 57 | u = (2 * pi * sigma**2 + 4 * pi * D * t) ** (-1) * torch.exp( 58 | -((x[:, 0:1] - x0[0] - v[0] * t) ** 2 + (x[:, 1:2] - x0[1] - v[1] * t) ** 2) 59 | / (2 * sigma**2 + 4 * D * t) 60 | ) 61 | coords = torch.cat((t.reshape(-1, 1), x.reshape(-1, 2)), dim=1) 62 | return coords, u.view(-1, 1) 63 | -------------------------------------------------------------------------------- /src/deepymod/data/kdv/kdv.py: -------------------------------------------------------------------------------- 1 | """ Contains several interactive datasets for the Korteweg-de-Vries equation including: 2 | - A single soliton wave 3 | - Two soliton waves 4 | """ 5 | 6 | import torch 7 | import numpy as np 8 | 9 | 10 | def single_soliton( 11 | x: torch.tensor, t: torch.tensor, c: float, x0: float 12 | ) -> torch.tensor: 13 | """Single soliton solution of the KdV equation (u_t + u_{xxx} - 6 u u_x = 0) 14 | 15 | Args: 16 | x ([Tensor]): Input vector of spatial coordinates. 17 | t ([Tensor]): Input vector of temporal coordinates. 18 | c ([Float]): Velocity. 19 | x0 ([Float]): Offset. 20 | Returns: 21 | [Tensor]: Solution. 22 | """ 23 | xi = np.sqrt(c) / 2 * (x - c * t - x0) # switch to moving coordinate frame 24 | u = c / 2 * 1 / torch.cosh(xi) ** 2 25 | coords = torch.cat((t.reshape(-1, 1), x.reshape(-1, 1)), dim=1) 26 | return coords, u.view(-1, 1) 27 | 28 | 29 | def double_soliton( 30 | x: torch.tensor, t: torch.tensor, c: float, x0: float 31 | ) -> torch.tensor: 32 | """Single soliton solution of the KdV equation (u_t + u_{xxx} - 6 u u_x = 0) 33 | source: http://lie.math.brocku.ca/~sanco/solitons/kdv_solitons.php 34 | 35 | Args: 36 | x ([Tensor]): Input vector of spatial coordinates. 37 | t ([Tensor]): Input vector of temporal coordinates. 38 | c ([Array]): Array containing the velocities of the two solitons, note that c[0] > c[1]. 39 | x0 ([Array]): Array containing the offsets of the two solitons. 40 | 41 | Returns: 42 | [Tensor]: Solution. 43 | """ 44 | assert c[0] > c[1], "c1 has to be bigger than c[2]" 45 | 46 | xi0 = ( 47 | np.sqrt(c[0]) / 2 * (x - c[0] * t - x0[0]) 48 | ) # switch to moving coordinate frame 49 | xi1 = np.sqrt(c[1]) / 2 * (x - c[1] * t - x0[1]) 50 | 51 | part_1 = 2 * (c[0] - c[1]) 52 | numerator = c[0] * torch.cosh(xi1) ** 2 + c[1] * torch.sinh(xi0) ** 2 53 | denominator_1 = (np.sqrt(c[0]) - np.sqrt(c[1])) * torch.cosh(xi0 + xi1) 54 | denominator_2 = (np.sqrt(c[0]) + np.sqrt(c[1])) * torch.cosh(xi0 - xi1) 55 | u = part_1 * numerator / (denominator_1 + denominator_2) ** 2 56 | coords = torch.cat((t.reshape(-1, 1), x.reshape(-1, 1)), dim=1) 57 | return coords, u.view(-1, 1) 58 | -------------------------------------------------------------------------------- /tests/burgers.py: -------------------------------------------------------------------------------- 1 | # General imports 2 | import numpy as np 3 | import torch 4 | 5 | # DeepMoD stuff 6 | from deepymod import DeepMoD 7 | from deepymod.model.func_approx import NN 8 | from deepymod.model.library import Library1D 9 | from deepymod.model.constraint import LeastSquares 10 | from deepymod.model.sparse_estimators import Threshold 11 | from deepymod.training import train 12 | from deepymod.training.sparsity_scheduler import TrainTestPeriodic, Periodic, TrainTest 13 | 14 | from deepymod.data import Dataset 15 | from deepymod.data.burgers import BurgersDelta 16 | 17 | from deepymod.analysis import load_tensorboard 18 | 19 | if torch.cuda.is_available(): 20 | device = "cuda" 21 | else: 22 | device = "cpu" 23 | 24 | # Settings for reproducibility 25 | np.random.seed(42) 26 | torch.manual_seed(0) 27 | torch.backends.cudnn.deterministic = True 28 | torch.backends.cudnn.benchmark = False 29 | 30 | # Making dataset 31 | v = 0.1 32 | A = 1.0 33 | 34 | x = np.linspace(-3, 4, 100) 35 | t = np.linspace(0.5, 5.0, 50) 36 | x_grid, t_grid = np.meshgrid(x, t, indexing="ij") 37 | dataset = Dataset(BurgersDelta, v=v, A=A) 38 | X, y = dataset.create_dataset( 39 | x_grid.reshape(-1, 1), 40 | t_grid.reshape(-1, 1), 41 | n_samples=1000, 42 | noise=0.4, 43 | random=True, 44 | normalize=False, 45 | ) 46 | X, y = X.to(device), y.to(device) 47 | 48 | network = NN(2, [30, 30, 30, 30, 30], 1) 49 | library = Library1D(poly_order=2, diff_order=3) # Library function 50 | estimator = Threshold(0.1) # Sparse estimator 51 | constraint = LeastSquares() # How to constrain 52 | model = DeepMoD(network, library, estimator, constraint).to( 53 | device 54 | ) # Putting it all in the model 55 | 56 | # sparsity_scheduler = TrainTestPeriodic(periodicity=50, patience=200, delta=1e-5) # in terms of write iterations 57 | # sparsity_scheduler = Periodic(initial_iteration=1000, periodicity=25) 58 | sparsity_scheduler = TrainTest(patience=200, delta=1e-5) 59 | 60 | optimizer = torch.optim.Adam( 61 | model.parameters(), betas=(0.99, 0.999), amsgrad=True, lr=2e-3 62 | ) # Defining optimizer 63 | 64 | train( 65 | model, 66 | X, 67 | y, 68 | optimizer, 69 | sparsity_scheduler, 70 | exp_ID="Test", 71 | split=0.8, 72 | write_iterations=25, 73 | max_iterations=20000, 74 | delta=0.001, 75 | patience=200, 76 | ) 77 | -------------------------------------------------------------------------------- /src/deepymod/data/samples.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | from numpy import ndarray 4 | from numpy.random import default_rng 5 | from abc import ABC, ABCMeta, abstractmethod 6 | 7 | 8 | class Subsampler(ABC, metaclass=ABCMeta): 9 | @abstractmethod 10 | def sample(): 11 | raise NotImplementedError 12 | 13 | 14 | class Subsample_time(Subsampler): 15 | @staticmethod 16 | def sample(coords, data, number_of_slices): 17 | """Subsample on the time axis for that has shape [t,x,y,z,...,feature] for both data and features.""" 18 | # getting indices of samples 19 | x_idx = torch.linspace( 20 | 0, coords.shape[0] - 1, number_of_slices, dtype=torch.long 21 | ) # getting x locations 22 | # getting sample locations from indices 23 | return coords[x_idx], data[x_idx] 24 | 25 | 26 | class Subsample_axis(Subsampler): 27 | @staticmethod 28 | def sample(coords, data, axis, number_of_slices): 29 | """Subsample on the specified axis to the number of slices specified""" 30 | # getting indices of samples 31 | feature_idx = torch.linspace( 32 | 0, coords.shape[axis] - 1, number_of_slices, dtype=torch.long 33 | ) # getting x locations 34 | # use the indices to subsample along the axis 35 | subsampled_coords = torch.index_select(coords, axis, feature_idx) 36 | subsampled_data = torch.index_select(data, axis, feature_idx) 37 | return subsampled_coords, subsampled_data 38 | 39 | 40 | # Needs to be implemented using torch.gather 41 | class Subsample_shifted_grid(Subsampler): 42 | @staticmethod 43 | def sample(coords, data, number_of_samples): 44 | return NotImplementedError 45 | 46 | 47 | class Subsample_random(Subsampler): 48 | @staticmethod 49 | def sample(coords, data, number_of_samples): 50 | """Apply random subsampling to a dataset, if it is not already in the 51 | (number_of_samples, number_of features) format, reshape it to it.""" 52 | # Ensure that both are of the shape (number_of_samples, number_of_features) before random sampling. 53 | if len(data.shape) > 2 or len(coords.shape) > 2: 54 | coords = coords.reshape((-1, coords.shape[-1])) 55 | data = data.reshape((-1, data.shape[-1])) 56 | # getting indices of samples 57 | x_idx = torch.randperm(coords.shape[0])[ 58 | :number_of_samples 59 | ] # getting x locations 60 | # getting sample locations from indices 61 | subsampled_coords = coords[x_idx] 62 | subsampled_data = data[x_idx] 63 | return subsampled_coords, subsampled_data 64 | -------------------------------------------------------------------------------- /setup.cfg: -------------------------------------------------------------------------------- 1 | # This file is used to configure your project. 2 | # Read more about the various options under: 3 | # http://setuptools.readthedocs.io/en/latest/setuptools.html#configuring-setup-using-setup-cfg-files 4 | 5 | [metadata] 6 | name = DeePyMoD 7 | description = DeePyMoD is a PyTorch-based implementation of the DeepMoD algorithm for model discovery of PDEs. 8 | author = Gert-Jan 9 | author-email = gert-jan.both@cri-paris.org 10 | license = mit 11 | long-description = file: README.md 12 | long-description-content-type = text/markdown; charset=UTF-8 13 | url = https://github.com/phimal/deepymod 14 | project-urls = 15 | Documentation = https://github.com/phimal/deepymod 16 | # Change if running only on Windows, Mac or Linux (comma-separated) 17 | platforms = any 18 | # Add here all kinds of additional classifiers as defined under 19 | # https://pypi.python.org/pypi?%3Aaction=list_classifiers 20 | classifiers = 21 | Development Status :: 4 - Beta 22 | Programming Language :: Python 23 | 24 | [options] 25 | zip_safe = False 26 | packages = find: 27 | include_package_data = True 28 | package_dir = 29 | =src 30 | # DON'T CHANGE THE FOLLOWING LINE! IT WILL BE UPDATED BY PYSCAFFOLD! 31 | setup_requires = pyscaffold>=3.2a0,<3.3a0 32 | # Add here dependencies of your project (semicolon/line-separated), e.g. 33 | install_requires = numpy 34 | torch 35 | scikit-learn 36 | pysindy 37 | natsort 38 | tensorboard 39 | matplotlib 40 | # The usage of test_requires is discouraged, see `Dependency Management` docs 41 | # tests_require = pytest; pytest-cov 42 | # Require a specific Python version, e.g. Python 2.7 or >= 3.4 43 | python_requires = >=3.6, !=3.9 44 | [options.packages.find] 45 | where = src 46 | exclude = 47 | tests 48 | 49 | [options.entry_points] 50 | # Add here console scripts like: 51 | # console_scripts = 52 | # script_name = deepymod_torch.module:function 53 | # For example: 54 | # console_scripts = 55 | # fibonacci = deepymod_torch.skeleton:run 56 | # And any other entry points, for example: 57 | # pyscaffold.cli = 58 | # awesome = pyscaffoldext.awesome.extension:AwesomeExtension 59 | 60 | [aliases] 61 | dists = bdist_wheel 62 | 63 | [bdist_wheel] 64 | # Use this option if your package is pure-python 65 | universal = 1 66 | 67 | [devpi:upload] 68 | # Options for the devpi: PyPI server and packaging tool 69 | # VCS export must be deactivated since we are using setuptools-scm 70 | no-vcs = 1 71 | formats = bdist_wheel 72 | 73 | [flake8] 74 | # Some sane defaults for the code style checker flake8 75 | exclude = 76 | .tox 77 | build 78 | dist 79 | .eggs 80 | docs/conf.py 81 | 82 | [pyscaffold] 83 | # PyScaffold's parameters when the project was created. 84 | # This will be used when updating. Do not change! 85 | version = 3.2 86 | package = deepymod 87 | -------------------------------------------------------------------------------- /docs/background/background.md: -------------------------------------------------------------------------------- 1 | # Background 2 | 3 | ## Framework 4 | 5 | Deep learning-based model discovery typically uses a neural network to construct a noiseless surrogate $\hat{u}$ of the data $u$. A library of potential terms $\Theta$ is constructed using automatic differentiation from $\hat{u}$ and the neural network is constrained to solutions allowed by this library . The loss function of the network thus consists of two contributions, (i) a mean square error to learn the mapping $(\vec{x},t) \rightarrow \hat{u}$ and (ii) a term to constrain the network, 6 | 7 | 8 | $\mathcal{L} = \frac{1}{N}\sum_{i=1}^{N}\left( u_i - \hat{u}_i \right) ^2 +\frac{1}{N}\sum_{i=1}^{N}\left( \partial_t \hat{u}_i - \Theta_{i}\xi \right)^2 .$ 9 | 10 | 11 | The sparse coefficient vector $\xi$ is learned concurrently with the network parameters and plays two roles: 1) determining the active (i.e. non-zero) components of the underlying PDE and 2) constraining the network according to these active terms. We propose to separate these two tasks by decoupling the constraint from the sparsity selection process itself. We first calculate a sparsity mask $g$ and then constrain the network only by the active terms in the mask. Mathematically, we replace $\xi$ by $\xi \circ \ g$. The sparsity mask $g$ need not be calculated differentiably, so that any classical, non-differentiable sparse estimator can be used. Our approach has several additional advantages: i) It provides an unbiased estimate of the coefficient vector since we do not apply $l_1$ or $l_2$ regularisation on $\xi$, ii) the sparsity pattern is determined from the full library $\Theta$, rather than only from the remaining active terms, allowing dynamic addition and removal of active terms throughout training, and iii) we can use cross validation or similar methods in the sparse estimator to find the optimal hyperparameters for model selection. 12 | 13 | ![Screenshot](../figures/framework.png) 14 | 15 | Using this change, we constructed a general framework for deep learning based model discovery with any classical sparsity promoting algorithm in the above. A *function approximator* constructs a surrogate of the data, (II) from which a *Library* of possible terms and the time derivative is constructed using automatic differentiation. (III) A *sparsity estimator* selects the active terms in the library using sparse regression and (IV) the function approximator is constrained to solutions allowed by the active terms by the *constraint*. 16 | 17 | ## Training 18 | 19 | As the sparsity estimator is non-differentiable, determining the sparsity mask before the function approximator has reasonably approximated the data can adversely affect training if the wrong terms are selected. We thus split the dataset into a train- and test-set and update the sparsity mask only when the MSE on the test-set starts to increase. After updating the mask, the model needs to adjust to the tighter constraint and we hence update the sparsity pattern every 25 epochs after the first update. Final convergence is reached when the $l_1$ norm of the coefficient vector remains constant. In practice we observe that large datasets with little noise might discover the correct equation after a single sparsity update, but that highly noisy datasets typically require several updates, removing only a few terms at a time. -------------------------------------------------------------------------------- /examples/VE_datagen.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import scipy.integrate as integ 3 | 4 | 5 | # Data generation routines 6 | def calculate_strain_stress( 7 | input_type, time_array, input_expr, E_mods, viscs, D_input_lambda=None 8 | ): 9 | # In this incarnation, the kwarg is non-optional. 10 | 11 | # if D_input_lambda: 12 | input_lambda = input_expr 13 | # else: 14 | # t = sym.symbols('t', real=True) 15 | # D_input_expr = input_expr.diff(t) 16 | 17 | # input_lambda = sym.lambdify(t, input_expr) 18 | # D_input_lambda = sym.lambdify(t, D_input_expr) 19 | 20 | # The following function interprets the provided model parameters differently depending on the input_type. 21 | # If the input_type is 'Strain' then the parameters are assumed to refer to a Maxwell model, whereas 22 | # if the input_type is 'Stress' then the parameters are assumed to refer to a Kelvin model. 23 | relax_creep_lambda = relax_creep(E_mods, viscs, input_type) 24 | 25 | if relax_creep_lambda == False: 26 | return False, False 27 | 28 | start_time_point = time_array[0] 29 | 30 | integrand_lambda = lambda x, t: relax_creep_lambda(t - x) * D_input_lambda(x) 31 | integral_lambda = lambda t: integ.quad( 32 | integrand_lambda, start_time_point, t, args=(t) 33 | )[0] 34 | 35 | output_array = np.array([]) 36 | input_array = np.array([]) 37 | for time_point in time_array: 38 | first_term = input_lambda(start_time_point) * relax_creep_lambda( 39 | time_point - start_time_point 40 | ) 41 | second_term = integral_lambda(time_point) 42 | output_array = np.append(output_array, first_term + second_term) 43 | input_array = np.append(input_array, input_lambda(time_point)) 44 | 45 | if input_type == "Strain": 46 | strain_array = input_array 47 | stress_array = output_array 48 | else: 49 | strain_array = output_array 50 | stress_array = input_array 51 | 52 | strain_array = strain_array.reshape(time_array.shape) 53 | stress_array = stress_array.reshape(time_array.shape) 54 | 55 | return strain_array, stress_array 56 | 57 | 58 | def relax_creep(E_mods, viscs, input_type): 59 | # The following function interprets the provided model parameters differently depending on the input_type. 60 | # If the input_type is 'Strain' then the parameters are assumed to refer to a Maxwell model, whereas 61 | # if the input_type is 'Stress' then the parameters are assumed to refer to a Kelvin model. 62 | # The equations used thus allow the data to be generated according to the model now designated. 63 | 64 | E_mods_1plus_array = np.array(E_mods[1:]).reshape(-1, 1) 65 | viscs_array = np.array(viscs).reshape(-1, 1) 66 | 67 | taus = viscs_array / E_mods_1plus_array 68 | 69 | if input_type == "Strain": 70 | relax_creep_lambda = lambda t: E_mods[0] + np.sum( 71 | np.exp(-t / taus) * E_mods_1plus_array 72 | ) 73 | elif input_type == "Stress": 74 | relax_creep_lambda = lambda t: 1 / E_mods[0] + np.sum( 75 | (1 - np.exp(-t / taus)) / E_mods_1plus_array 76 | ) 77 | else: 78 | print("Incorrect input_type") 79 | relax_creep_lambda = False 80 | 81 | return relax_creep_lambda 82 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | ![Screenshot](docs/figures/DeepMoD_logo.png) 2 | -------------------------------------------------------------------------------- 3 | PyPI 4 | 5 | DeePyMoD is a modular framework for model discovery of PDEs and ODEs from noise data. The framework is comprised of four components, that can separately be altered: i) A function approximator to construct a surrogate of the data, ii) a function to construct the library of features, iii) a sparse regression algorithm to select the active components from the feature library and iv) a constraint on the function approximator, based on the active components. 6 | 7 | ![Screenshot](docs/figures/framework.png) 8 | 9 | More information can be found in the following two papers: , [arXiv:2011.04336](https://arxiv.org/abs/2011.04336), [arXiv:1904.09406](http://arxiv.org/abs/1904.09406) and the full documentation is available on [phimal.github.io/DeePyMoD/](https://phimal.github.io/DeePyMoD/). 10 | 11 | **What's the use case?** Classical Model Discovery methods struggle with elevated noise levels and sparse datasets due the low accuracy of numerical differentiation. DeepMoD can handle high noise and sparse datasets, making it well suited for model discovery on actual experimental data. 12 | 13 | **What types of models can you discover?** DeepMoD can discover non-linear, multi-dimensional and/or coupled ODEs and PDEs. See our paper and the examples folder for a demonstration of each. 14 | 15 | # How to install 16 | 17 | ## Dependencies and CUDA 18 | We support Python 3.6, 3.7 and 3.8. 19 | We rely on the following packages, they will be installed in the pip installation procces for you: 20 | ``` numpy, torch, sklearn, pysindy, natsort, tensorboard, matplotlib``` 21 | 22 | 23 | We also make use of the PyTorch library, which can be installed with CPU and/or GPU support. Please 24 | refer to the PyTorch ["Get Started"](https://pytorch.org/get-started/locally/) guide to get the version 25 | that is optimal for your system. 26 | 27 | ## Install 28 | To install DeePyMoD, pip can be used 29 | 30 | ``` pip install deepymod ``` 31 | 32 | in the main directory. 33 | 34 | 35 | ## Development 36 | 37 | If you wish to alter the code you can clone the package using: 38 | 39 | ``` git clone git@github.com:PhIMaL/DeePyMoD.git ``` 40 | 41 | and then install it from the cloned `DeePyMoD` directory using 42 | 43 | ``` pip install -e ./ ``` 44 | 45 | # Features 46 | 47 | * **Many example notebooks** We have implemented a variety of examples ranging from 2D Advection Diffusion, Burgers' equation to non-linear, higher order ODE's If you miss any example, don't hesitate to give us a heads-up. 48 | 49 | * **Extendable** DeePyMoD is designed to be easily extendable and modifiable. You can simply plug in your own cost function, library or training regime. 50 | 51 | * **Automatic library** The library and coefficient vectors are automatically constructed from the maximum order of polynomial and differentiation. If that doesn't cut it for your use case, it's easy to plug in your own library function. 52 | 53 | * **Extensive logging** We provide a simple command line logger to see how training is going and an extensive custom Tensorboard logger. 54 | 55 | * **Fast** Depending on the size of the data-set DeepMoD, running a model search with DeepMoD takes of the order of minutes/ tens of minutes on a standard CPU. Running the code on GPU's drastically improves performance. 56 | 57 | -------------------------------------------------------------------------------- /src/deepymod/data/burgers/burgers.py: -------------------------------------------------------------------------------- 1 | """ Contains several interactive datasets for the Burgers equation including: 2 | - Burgers with initial delta peak profile 3 | - Burgers with initial cosine profile 4 | - Burgers with initial sawtooth profile""" 5 | 6 | import torch 7 | from numpy import pi 8 | 9 | from deepymod.data import Dataset 10 | 11 | 12 | def burgers_delta(x: torch.tensor, t: torch.tensor, v: float, A: float): 13 | """Function to load the analytical solutions of Burgers equation with delta peak initial condition: u(x, 0) = A delta(x) 14 | 15 | Source: https://www.iist.ac.in/sites/default/files/people/IN08026/Burgers_equation_viscous.pdf 16 | Note that this source has an error in the erfc prefactor, should be sqrt(pi)/2, not sqrt(pi/2). 17 | 18 | Args: 19 | x ([Tensor]): Input vector of spatial coordinates. 20 | t ([Tensor]): Input vector of temporal coordinates. 21 | v (Float): Velocity. 22 | A (Float): Amplitude of the initial condition. 23 | 24 | Returns: 25 | [Tensor]: solution. 26 | """ 27 | x, t = torch.meshgrid(x, t) 28 | R = torch.tensor(A / (2 * v)) # otherwise throws error 29 | z = x / torch.sqrt(4 * v * t) 30 | 31 | u = ( 32 | torch.sqrt(v / (pi * t)) 33 | * ((torch.exp(R) - 1) * torch.exp(-(z**2))) 34 | / (1 + (torch.exp(R) - 1) / 2 * torch.erfc(z)) 35 | ) 36 | coords = torch.cat((t.reshape(-1, 1), x.reshape(-1, 1)), dim=1) 37 | return coords, u.view(-1, 1) 38 | 39 | 40 | def burgers_cos( 41 | x: torch.tensor, t: torch.tensor, v: float, a: float, b: float, k: float 42 | ): 43 | """Function to generate analytical solutions of Burgers equation with cosine initial condition: 44 | $u(x, 0) = b + a \cos(kx)$ 45 | 46 | Source: https://www.iist.ac.in/sites/default/files/people/IN08026/Burgers_equation_viscous.pdf 47 | 48 | Args: 49 | x ([Tensor]): Input vector of spatial coordinates. 50 | t ([Tensor]): Input vector of temporal coordinates. 51 | v (Float): Velocity. 52 | a ([Float]): Amplitude of the initial periodic condition. 53 | b ([Float]): Offset of the initial condition. 54 | k ([Float]): Wavenumber of the initial condition. 55 | 56 | Returns: 57 | [Tensor]: solution. 58 | """ 59 | 60 | z = v * k**2 * t 61 | 62 | u = (2 * v * a * k * torch.exp(-z) * torch.sin(k * x)) / ( 63 | b + a * torch.exp(-z) * torch.cos(k * x) 64 | ) 65 | coords = torch.cat((t.reshape(-1, 1), x.reshape(-1, 1)), dim=1) 66 | return coords, u.view(-1, 1) 67 | 68 | 69 | def burgers_sawtooth(x: torch.tensor, t: torch.tensor, v: float) -> torch.tensor: 70 | """Function to generate analytical solutions of Burgers equation with sawtooth initial condition (see soruce for exact expression). Solution only 71 | valid between for x in [0, 2pi] and t in [0, 0.5] 72 | 73 | http://www.thevisualroom.com/02_barba_projects/burgers_equation.html 74 | 75 | Args: 76 | x ([Tensor]): Input vector of spatial coordinates. 77 | t ([Tensor]): Input vector of temporal coordinates. 78 | v (Float): Velocity. 79 | 80 | Returns: 81 | [Tensor]: solution. 82 | """ 83 | 84 | z_left = x - 4 * t 85 | z_right = x - 4 * t - 2 * pi 86 | l = 4 * v * (t + 1) 87 | 88 | phi = torch.exp(-(z_left**2) / l) + torch.exp(-(z_right**2) / l) 89 | dphi_x = -2 * z_left / l * torch.exp( 90 | -(z_left**2) / l 91 | ) - 2 * z_right / l * torch.exp(-(z_right**2) / l) 92 | u = -2 * v * dphi_x / phi + 4 93 | coords = torch.cat((t.reshape(-1, 1), x.reshape(-1, 1)), dim=1) 94 | return coords, u.view(-1, 1) 95 | -------------------------------------------------------------------------------- /docs/datasets/data.md: -------------------------------------------------------------------------------- 1 | # Datasets 2 | 3 | ## The general workflow 4 | 5 | The custom DeePyMoD dataset and dataloaders are created for data that typically fits in the RAM/VRAM during training, if this is not your use case, they are interchangeable with the PyTorch general [Datasets](https://pytorch.org/docs/stable/data.html#torch.utils.data.Dataset) and 6 | and [Dataloaders](https://pytorch.org/docs/stable/data.html#torch.utils.data.Dataset). 7 | 8 | For model discovery we typically want to add some noise to our dataset, normalize certain features and ensure it is in the right place for optimal PyTorch performance. This can easily be done by using the custom `deepymod.data.Dataset` and `deepymod.data.get_train_test_loader`. An illustration of the workflow is shown below: 9 | 10 | ![Workflow](../figures/data_workflow_for_deepymod.png) 11 | 12 | ## The dataset 13 | The dataset needs a function that loads all the samples and returns it in a coordinate, data format. 14 | ```python 15 | def load_data(): 16 | # create or load your data here 17 | return coordinates, data 18 | ``` 19 | Here it is important that the last axis of the data is the number of features, even if it is just one. The dataset accepts data that is still dimensionful `(t,x,y,z,number_of_features)` as well as data that is already 20 | flattened `(number_of_samples, number_of_features)`. After returning the data tries to apply the following functions to the samples that were just loaded: `preprocessing`, `subsampling` and lastly `shuffling`. 21 | 22 | ### Preprocessing 23 | The preprocessing performs steps commonly used in the framework, normalizing the coordinates, normalizing the data and adding noise to the data. One can provide these choices via a dictionary of arguments: 24 | ```python 25 | preprocess_kwargs: dict = { 26 | "random_state": 42, 27 | "noise_level": 0.0, 28 | "normalize_coords": False, 29 | "normalize_data": False, 30 | } 31 | ``` 32 | And we can override the way we preprocess functions by defining the preprocess functions `apply_normalize`, `apply_noise` or even the way we shuffle using `apply_shuffle`. 33 | 34 | ### Subsampling 35 | Sometimes we do not wish to use the whole dataset, and as such we can subsample it. Sometimes using a subset 36 | of the time snapshots available is enough, for this we can use `deepymod.data.samples.Subsample_time` or 37 | randomly with `deepymod.data.samples.Subsample_random`. You can provide the arguments for these functions via `subsampler_kwargs` to the Dataset. 38 | 39 | ### The resulting shape of the data 40 | Since for random subsampling the (number_of_samples, number_of_features) format is better and for spatial 41 | subsampling the (t,x,y,z,number_of_features) format is best, we accept both formats. However since the trainer 42 | can only work with (number_of_samples, number_of_features), we will reshape the data to this format once 43 | the data is preprocessed and subsampled. After this we can shuffle the data. 44 | 45 | ### Shuffling 46 | If the data needs to be shuffled, `shuffle=True` can be used 47 | 48 | ## Dataloaders 49 | Dataloaders are used in the PyTorch framework to ensure that the loading of the data goes smoothly, 50 | for example with multiple workers. We however can typically fit the whole dataset into memory once, 51 | so the overhead of the PyTorch Dataloader is not needed. We thus provide the Loader, which provides 52 | a wrapper around the dataset. The loader will return the entire batch at once. 53 | 54 | ## Obtaining the dataloaders 55 | In order to create a train and test split, we can use the function `get_train_test_loader`, which divides 56 | the dataset into two pieces, and then directly passes them into the loader. 57 | 58 | 59 | -------------------------------------------------------------------------------- /src/deepymod/model/constraint.py: -------------------------------------------------------------------------------- 1 | """This module contains concrete implementations of the constraint component.""" 2 | 3 | 4 | import torch 5 | from .deepmod import Constraint 6 | from typing import List 7 | 8 | TensorList = List[torch.Tensor] 9 | 10 | 11 | class LeastSquares(Constraint): 12 | def __init__(self) -> None: 13 | """Least Squares Constraint solved by QR decomposition""" 14 | super().__init__() 15 | 16 | def fit(self, sparse_thetas: TensorList, time_derivs: TensorList) -> TensorList: 17 | """Calculates the coefficients of the constraint using the QR decomposition for every pair 18 | of sparse feature matrix and time derivative. 19 | 20 | Args: 21 | sparse_thetas (TensorList): List containing the sparse feature tensors of size (n_samples, n_active_features). 22 | time_derivs (TensorList): List containing the time derivatives of size (n_samples, n_outputs). 23 | 24 | Returns: 25 | (TensorList): List of calculated coefficients of size [(n_active_features, 1) x n_outputs]. 26 | """ 27 | coeff_vectors = [] 28 | for theta, dt in zip(sparse_thetas, time_derivs): 29 | Q, R = torch.qr(theta) # solution of lst. sq. by QR decomp. 30 | coeff_vectors.append(torch.inverse(R) @ Q.T @ dt) 31 | return coeff_vectors 32 | 33 | 34 | class GradParams(Constraint): 35 | def __init__(self, n_params: int, n_eqs: int) -> None: 36 | """Constrains the neural network by optimizing over the coefficients together with the network. 37 | Coefficient vectors are randomly initialized from a standard Gaussian. 38 | 39 | Args: 40 | n_params (int): number of features in feature matrix. 41 | n_eqs (int): number of outputs / equations to be discovered. 42 | """ 43 | super().__init__() 44 | self.coeffs = torch.nn.ParameterList( 45 | [torch.nn.Parameter(torch.randn(n_params, 1)) for _ in torch.arange(n_eqs)] 46 | ) 47 | 48 | def fit(self, sparse_thetas: TensorList, time_derivs: TensorList): 49 | """Returns the coefficients of the constraint, since we're optimizing them by 50 | gradient descent. 51 | 52 | Args: 53 | sparse_thetas (TensorList): List containing the sparse feature tensors of size (n_samples, n_active_features). 54 | time_derivs (TensorList): List containing the time derivatives of size (n_samples, n_outputs). 55 | 56 | Returns: 57 | (TensorList): Calculated coefficients of size (n_features, n_outputs). 58 | """ 59 | return self.coeffs 60 | 61 | 62 | class Ridge(Constraint): 63 | """Implements the constraint as a least squares problem solved by QR decomposition.""" 64 | 65 | def __init__(self, l=1e-3) -> None: 66 | super().__init__() 67 | self.l = torch.tensor(l) 68 | 69 | def fit(self, sparse_thetas: TensorList, time_derivs: TensorList) -> TensorList: 70 | """Calculates the coefficients of the constraint using the QR decomposition for every pair 71 | of sparse feature matrix and time derivative. 72 | 73 | Args: 74 | sparse_thetas (TensorList): List containing the sparse feature tensors. 75 | time_derivs (TensorList): List containing the time derivatives. 76 | 77 | Returns: 78 | (TensorList): Calculated coefficients. 79 | """ 80 | coeff_vectors = [] 81 | for theta, dt in zip(sparse_thetas, time_derivs): 82 | # We use the augmented data method 83 | norm = torch.norm(theta, dim=0) 84 | l_normed = torch.diag( 85 | torch.sqrt(self.l) * norm 86 | ) # we norm l rather than theta cause we want unnormed coeffs out 87 | X = torch.cat((theta, l_normed), dim=0) 88 | y = torch.cat((dt, torch.zeros((theta.shape[1], 1))), dim=0) 89 | # Now solve like normal OLS prob lem 90 | Q, R = torch.qr(X) # solution of lst. sq. by QR decomp. 91 | coeff_vectors.append(torch.inverse(R) @ Q.T @ y) 92 | 93 | return coeff_vectors 94 | -------------------------------------------------------------------------------- /src/deepymod/utils/logger.py: -------------------------------------------------------------------------------- 1 | """ Module to log performance metrics whilst training Deepmyod """ 2 | import numpy as np 3 | import torch 4 | import sys, time 5 | from torch.utils.tensorboard import SummaryWriter 6 | 7 | 8 | class Logger: 9 | def __init__(self, exp_ID, log_dir): 10 | """Log the training process of Deepymod. 11 | Args: 12 | exp_ID (str): name or ID of the this experiment 13 | log_dir (str): directory to save the log files to disk. 14 | 15 | """ 16 | self.writer = SummaryWriter( 17 | comment=exp_ID, log_dir=log_dir, max_queue=5, flush_secs=10 18 | ) 19 | self.log_dir = self.writer.get_logdir() 20 | 21 | def __call__( 22 | self, 23 | iteration, 24 | loss, 25 | MSE, 26 | Reg, 27 | constraint_coeffs, 28 | unscaled_constraint_coeffs, 29 | estimator_coeffs, 30 | **kwargs, 31 | ): 32 | l1_norm = torch.sum(torch.abs(torch.cat(constraint_coeffs, dim=1)), dim=0) 33 | 34 | self.update_tensorboard( 35 | iteration, 36 | loss, 37 | MSE, 38 | Reg, 39 | l1_norm, 40 | constraint_coeffs, 41 | unscaled_constraint_coeffs, 42 | estimator_coeffs, 43 | **kwargs, 44 | ) 45 | self.update_terminal(iteration, MSE, Reg, l1_norm) 46 | 47 | def update_tensorboard( 48 | self, 49 | iteration, 50 | loss, 51 | loss_mse, 52 | loss_reg, 53 | loss_l1, 54 | constraint_coeff_vectors, 55 | unscaled_constraint_coeff_vectors, 56 | estimator_coeff_vectors, 57 | **kwargs, 58 | ): 59 | """Write the current state of training to Tensorboard 60 | Args: 61 | iteration (int): iteration number 62 | loss (float): loss value 63 | loss_mse (float): loss of the Mean Squared Error term 64 | loss_reg (float): loss of the regularization term 65 | loss_l1 (float): loss of the L1 penalty term 66 | constraint_coeff_vectors (np.array): vector with constraint coefficients 67 | unscaled_constraint_coeff_vectors (np.array): unscaled vector with constraint coefficients 68 | estimator_coeff_vectors (np.array): coefficients as computed by the estimator. 69 | """ 70 | # Costs and coeff vectors 71 | self.writer.add_scalar("loss/loss", loss, iteration) 72 | self.writer.add_scalars( 73 | "loss/mse", 74 | {f"output_{idx}": val for idx, val in enumerate(loss_mse)}, 75 | iteration, 76 | ) 77 | self.writer.add_scalars( 78 | "loss/reg", 79 | {f"output_{idx}": val for idx, val in enumerate(loss_reg)}, 80 | iteration, 81 | ) 82 | self.writer.add_scalars( 83 | "loss/l1", 84 | {f"output_{idx}": val for idx, val in enumerate(loss_l1)}, 85 | iteration, 86 | ) 87 | 88 | for output_idx, (coeffs, unscaled_coeffs, estimator_coeffs) in enumerate( 89 | zip( 90 | constraint_coeff_vectors, 91 | unscaled_constraint_coeff_vectors, 92 | estimator_coeff_vectors, 93 | ) 94 | ): 95 | self.writer.add_scalars( 96 | f"coeffs/output_{output_idx}", 97 | {f"coeff_{idx}": val for idx, val in enumerate(coeffs.squeeze())}, 98 | iteration, 99 | ) 100 | self.writer.add_scalars( 101 | f"unscaled_coeffs/output_{output_idx}", 102 | { 103 | f"coeff_{idx}": val 104 | for idx, val in enumerate(unscaled_coeffs.squeeze()) 105 | }, 106 | iteration, 107 | ) 108 | self.writer.add_scalars( 109 | f"estimator_coeffs/output_{output_idx}", 110 | { 111 | f"coeff_{idx}": val 112 | for idx, val in enumerate(estimator_coeffs.squeeze()) 113 | }, 114 | iteration, 115 | ) 116 | 117 | # Writing remaining kwargs 118 | for key, value in kwargs.items(): 119 | if value.numel() == 1: 120 | self.writer.add_scalar(f"remaining/{key}", value, iteration) 121 | else: 122 | self.writer.add_scalars( 123 | f"remaining/{key}", 124 | { 125 | f"val_{idx}": val.squeeze() 126 | for idx, val in enumerate(value.squeeze()) 127 | }, 128 | iteration, 129 | ) 130 | 131 | def update_terminal(self, iteration, MSE, Reg, L1): 132 | """Prints and updates progress of training cycle in command line.""" 133 | sys.stdout.write( 134 | f"\r{iteration:>6} MSE: {torch.sum(MSE).item():>8.2e} Reg: {torch.sum(Reg).item():>8.2e} L1: {torch.sum(L1).item():>8.2e} " 135 | ) 136 | sys.stdout.flush() 137 | 138 | def close(self, model): 139 | """Close the Tensorboard writer""" 140 | print("Algorithm converged. Writing model to disk.") 141 | self.writer.flush() # flush remaining stuff to disk 142 | self.writer.close() # close writer 143 | 144 | # Save model 145 | model_path = self.log_dir + "model.pt" 146 | torch.save(model.state_dict(), model_path) 147 | -------------------------------------------------------------------------------- /docs/examples/2DAD/2DAD.md: -------------------------------------------------------------------------------- 1 | # 2D Advection-Diffusion equation 2 | 3 | in this notebook we provide a simple example of the DeepMoD algorithm and apply it on the 2D advection-diffusion equation. 4 | 5 | 6 | ```python 7 | # General imports 8 | import numpy as np 9 | import torch 10 | import matplotlib.pylab as plt 11 | 12 | # DeepMoD functions 13 | 14 | from deepymod import DeepMoD 15 | from deepymod.model.func_approx import NN 16 | from deepymod.model.library import Library2D 17 | from deepymod.model.constraint import LeastSquares 18 | from deepymod.model.sparse_estimators import Threshold,PDEFIND 19 | from deepymod.training import train 20 | from deepymod.training.sparsity_scheduler import TrainTestPeriodic 21 | from scipy.io import loadmat 22 | 23 | # Settings for reproducibility 24 | np.random.seed(42) 25 | torch.manual_seed(0) 26 | 27 | 28 | %load_ext autoreload 29 | %autoreload 2 30 | ``` 31 | 32 | ## Prepare the data 33 | 34 | Next, we prepare the dataset. 35 | 36 | 37 | ```python 38 | data = loadmat('data/advection_diffusion.mat') 39 | usol = np.real(data['Expression1']) 40 | usol= usol.reshape((51,51,61,4)) 41 | 42 | x_v= usol[:,:,:,0] 43 | y_v = usol[:,:,:,1] 44 | t_v = usol[:,:,:,2] 45 | u_v = usol[:,:,:,3] 46 | ``` 47 | 48 | Next we plot the dataset for three different time-points 49 | 50 | 51 | ```python 52 | fig, axes = plt.subplots(ncols=3, figsize=(15, 4)) 53 | 54 | im0 = axes[0].contourf(x_v[:,:,0], y_v[:,:,0], u_v[:,:,0], cmap='coolwarm') 55 | axes[0].set_xlabel('x') 56 | axes[0].set_ylabel('y') 57 | axes[0].set_title('t = 0') 58 | 59 | im1 = axes[1].contourf(x_v[:,:,10], y_v[:,:,10], u_v[:,:,10], cmap='coolwarm') 60 | axes[1].set_xlabel('x') 61 | axes[1].set_title('t = 10') 62 | 63 | im2 = axes[2].contourf(x_v[:,:,20], y_v[:,:,20], u_v[:,:,20], cmap='coolwarm') 64 | axes[2].set_xlabel('x') 65 | axes[2].set_title('t= 20') 66 | 67 | fig.colorbar(im1, ax=axes.ravel().tolist()) 68 | 69 | plt.show() 70 | ``` 71 | 72 | 73 | ![png](output_7_0.png) 74 | 75 | 76 | We flatten it to give it the right dimensions for feeding it to the network: 77 | 78 | 79 | ```python 80 | X = np.transpose((t_v.flatten(),x_v.flatten(), y_v.flatten())) 81 | y = np.float32(u_v.reshape((u_v.size, 1))) 82 | ``` 83 | 84 | We select the noise level we add to the data-set 85 | 86 | 87 | ```python 88 | noise_level = 0.01 89 | ``` 90 | 91 | 92 | ```python 93 | y_noisy = y + noise_level * np.std(y) * np.random.randn(y.size, 1) 94 | ``` 95 | 96 | Select the number of samples: 97 | 98 | 99 | ```python 100 | number_of_samples = 1000 101 | 102 | idx = np.random.permutation(y.shape[0]) 103 | X_train = torch.tensor(X[idx, :][:number_of_samples], dtype=torch.float32, requires_grad=True) 104 | y_train = torch.tensor(y[idx, :][:number_of_samples], dtype=torch.float32) 105 | 106 | 107 | ``` 108 | 109 | ## Configuration of DeepMoD 110 | 111 | Configuration of the function approximator: Here the first argument is the number of input and the last argument the number of output layers. 112 | 113 | 114 | ```python 115 | network = NN(3, [50, 50, 50,50], 1) 116 | ``` 117 | 118 | Configuration of the library function: We select athe library with a 2D spatial input. Note that that the max differential order has been pre-determined here out of convinience. So, for poly_order 1 the library contains the following 12 terms: 119 | * [$1, u_x, u_y, u_{xx}, u_{yy}, u_{xy}, u, u u_x, u u_y, u u_{xx}, u u_{yy}, u u_{xy}$] 120 | 121 | 122 | ```python 123 | library = Library2D(poly_order=1) 124 | ``` 125 | 126 | Configuration of the sparsity estimator and sparsity scheduler used. In this case we use the most basic threshold-based Lasso estimator and a scheduler that asseses the validation loss after a given patience. If that value is smaller than 1e-5, the algorithm is converged. 127 | 128 | 129 | ```python 130 | estimator = Threshold(0.1) 131 | sparsity_scheduler = TrainTestPeriodic(periodicity=50, patience=10, delta=1e-5) 132 | ``` 133 | 134 | Configuration of the sparsity estimator 135 | 136 | 137 | ```python 138 | constraint = LeastSquares() 139 | # Configuration of the sparsity scheduler 140 | ``` 141 | 142 | Now we instantiate the model and select the optimizer 143 | 144 | 145 | ```python 146 | model = DeepMoD(network, library, estimator, constraint) 147 | 148 | # Defining optimizer 149 | optimizer = torch.optim.Adam(model.parameters(), betas=(0.99, 0.99), amsgrad=True, lr=1e-3) 150 | 151 | ``` 152 | 153 | ## Run DeepMoD 154 | 155 | We can now run DeepMoD using all the options we have set and the training data: 156 | * The directory where the tensorboard file is written (log_dir) 157 | * The ratio of train/test set used (split) 158 | * The maximum number of iterations performed (max_iterations) 159 | * The absolute change in L1 norm considered converged (delta) 160 | * The amount of epochs over which the absolute change in L1 norm is calculated (patience) 161 | 162 | 163 | ```python 164 | train(model, X_train, y_train, optimizer,sparsity_scheduler, log_dir='runs/2DAD/', split=0.8, max_iterations=100000, delta=1e-4, patience=8) 165 | ``` 166 | 167 | | Iteration | Progress | Time remaining | Loss | MSE | Reg | L1 norm | 168 | 7000 7.00% 3733s 1.07e-04 3.60e-05 7.08e-05 1.87e+00 Algorithm converged. Stopping training. 169 | 170 | 171 | Sparsity masks provide the active and non-active terms in the PDE: 172 | 173 | 174 | ```python 175 | model.sparsity_masks 176 | ``` 177 | 178 | 179 | 180 | 181 | [tensor([False, True, True, True, True, False, False, False, False, False, 182 | False, False])] 183 | 184 | 185 | 186 | estimatior_coeffs gives the magnitude of the active terms: 187 | 188 | 189 | ```python 190 | print(model.estimator_coeffs()) 191 | ``` 192 | 193 | [array([[0. ], 194 | [0.3770935 ], 195 | [0.7139108 ], 196 | [0.389949 ], 197 | [0.32122847], 198 | [0. ], 199 | [0. ], 200 | [0. ], 201 | [0. ], 202 | [0. ], 203 | [0. ], 204 | [0. ]], dtype=float32)] 205 | 206 | 207 | 208 | ```python 209 | 210 | ``` 211 | 212 | 213 | ```python 214 | 215 | ``` 216 | -------------------------------------------------------------------------------- /src/deepymod/training/training.py: -------------------------------------------------------------------------------- 1 | """ Contains the train module that governs training Deepymod """ 2 | import torch 3 | from ..utils.logger import Logger 4 | from .convergence import Convergence 5 | from ..model.deepmod import DeepMoD 6 | from torch.utils.data import DataLoader 7 | 8 | 9 | def train( 10 | model: DeepMoD, 11 | train_dataloader: DataLoader, 12 | test_dataloader: DataLoader, 13 | optimizer, 14 | sparsity_scheduler, 15 | split: float = 0.8, 16 | exp_ID: str = None, 17 | log_dir: str = None, 18 | max_iterations: int = 10000, 19 | write_iterations: int = 25, 20 | **convergence_kwargs 21 | ) -> None: 22 | """Trains the DeepMoD model. This function automatically splits the data set in a train and test set. 23 | 24 | Args: 25 | model (DeepMoD): A DeepMoD object. 26 | data (torch.Tensor): Tensor of shape (n_samples x (n_spatial + 1)) containing the coordinates, first column should be the time coordinate. 27 | target (torch.Tensor): Tensor of shape (n_samples x n_features) containing the target data. 28 | optimizer ([type]): Pytorch optimizer. 29 | sparsity_scheduler ([type]): Decides when to update the sparsity mask. 30 | split (float, optional): Fraction of the train set, by default 0.8. 31 | exp_ID (str, optional): Unique ID to identify tensorboard file. Not used if log_dir is given, see pytorch documentation. 32 | log_dir (str, optional): Directory where tensorboard file is written, by default None. 33 | max_iterations (int, optional): [description]. Max number of epochs , by default 10000. 34 | write_iterations (int, optional): [description]. Sets how often data is written to tensorboard and checks train loss , by default 25. 35 | """ 36 | logger = Logger(exp_ID, log_dir) 37 | sparsity_scheduler.path = ( 38 | logger.log_dir 39 | ) # write checkpoint to same folder as tb output. 40 | n_features = train_dataloader[0][1].shape[-1] 41 | # n_features = model.func_approx.modules() 42 | # Training 43 | convergence = Convergence(**convergence_kwargs) # Convergence check initialization 44 | for iteration in torch.arange(0, max_iterations): # iterate over epochs 45 | # Training variables defined as: loss, mse, regularisation 46 | batch_losses = torch.zeros( 47 | (3, n_features, len(train_dataloader)), 48 | device=train_dataloader.device, 49 | ) 50 | for batch_idx, train_sample in enumerate(train_dataloader): 51 | data_train, target_train = train_sample 52 | # ================== Training Model ============================ 53 | prediction, time_derivs, thetas = model(data_train) 54 | batch_losses[1, :, batch_idx] = torch.mean( 55 | (prediction - target_train) ** 2, dim=-2 56 | ) # loss per output, equation (6) of the paper, used to be called mse_loss in DeePyMoD_torch 57 | batch_losses[2, :, batch_idx] = torch.stack( 58 | [ 59 | torch.mean( 60 | (dt - theta @ coeff_vector) ** 2 61 | ) # equation (7) of the original paper, used to be called reg_loss in DeePyMoD_torch 62 | for dt, theta, coeff_vector in zip( 63 | time_derivs, 64 | thetas, 65 | model.constraint_coeffs(scaled=False, sparse=True), 66 | ) 67 | ] 68 | ) 69 | batch_losses[0, :, batch_idx] = ( # mse_loss + reg_loss 70 | batch_losses[1, :, batch_idx] + batch_losses[2, :, batch_idx] 71 | ) # The weight optimization does not call on L1 losses, instead 72 | # L1 is included in the logging below. See the latest paper: http://arxiv.org/abs/2011.04336 73 | 74 | # Optimizer step 75 | optimizer.zero_grad() 76 | batch_losses[0, :, batch_idx].sum().backward() 77 | optimizer.step() 78 | 79 | loss, mse, reg = torch.mean(batch_losses.cpu().detach(), axis=-1) 80 | 81 | if iteration % write_iterations == 0: 82 | # ================== Validation costs ================ 83 | with torch.no_grad(): 84 | batch_mse_test = torch.zeros( 85 | (n_features, len(test_dataloader)), device=test_dataloader.device 86 | ) 87 | for batch_idx, test_sample in enumerate(test_dataloader): 88 | data_test, target_test = test_sample 89 | prediction_test = model.func_approx(data_test)[0] 90 | batch_mse_test[:, batch_idx] = torch.mean( # mse loss 91 | (prediction_test - target_test) ** 2, dim=-2 92 | ) # loss per output 93 | mse_test = batch_mse_test.cpu().detach().mean(dim=-1) 94 | # ====================== Logging ======================= 95 | _ = model.sparse_estimator( 96 | thetas, time_derivs 97 | ) # calculating estimator coeffs but not setting mask 98 | logger( 99 | iteration, 100 | loss.view(-1).mean(), 101 | mse.view(-1), 102 | reg.view(-1), 103 | model.constraint_coeffs(sparse=True, scaled=True), 104 | model.constraint_coeffs(sparse=True, scaled=False), 105 | model.estimator_coeffs(), 106 | MSE_test=mse_test, 107 | ) 108 | 109 | # ================== Sparsity update ============= 110 | # Updating sparsity 111 | update_sparsity = sparsity_scheduler( 112 | iteration, torch.sum(mse_test), model, optimizer 113 | ) 114 | if update_sparsity: 115 | model.constraint.sparsity_masks = model.sparse_estimator( 116 | thetas, time_derivs 117 | ) 118 | 119 | # ================= Checking convergence 120 | l1_norm = torch.sum( 121 | torch.abs( 122 | torch.cat(model.constraint_coeffs(sparse=True, scaled=True), dim=1) 123 | ) 124 | ) 125 | converged = convergence( 126 | iteration, l1_norm 127 | ) # Check if change is smaller than delta and if we've exceeded patience 128 | if converged: 129 | break 130 | logger.close(model) 131 | -------------------------------------------------------------------------------- /docs/examples/ODE_nonlin/ODE.md: -------------------------------------------------------------------------------- 1 | # Example ODE with custom library 2 | 3 | In this notebook we provide a simple example of the DeepMoD algorithm by applying it on the a non-linear ODE 4 | 5 | We start by importing the required DeepMoD functions: 6 | 7 | 8 | ```python 9 | # General imports 10 | import numpy as np 11 | import torch 12 | import matplotlib.pylab as plt 13 | 14 | # DeepMoD functions 15 | 16 | 17 | from deepymod import DeepMoD 18 | from deepymod.model.func_approx import NN 19 | from deepymod.model.constraint import LeastSquares 20 | from deepymod.model.sparse_estimators import Threshold, PDEFIND 21 | from deepymod.training import train 22 | from deepymod.training.sparsity_scheduler import TrainTestPeriodic 23 | from scipy.io import loadmat 24 | 25 | 26 | import torch 27 | from torch.autograd import grad 28 | from itertools import combinations 29 | from functools import reduce 30 | from typing import Tuple 31 | from deepymod.utils.types import TensorList 32 | from deepymod import Library 33 | 34 | %load_ext autoreload 35 | %autoreload 2 36 | from scipy.integrate import odeint 37 | 38 | # Settings for reproducibility 39 | np.random.seed(40) 40 | torch.manual_seed(0) 41 | 42 | ``` 43 | 44 | The autoreload extension is already loaded. To reload it, use: 45 | %reload_ext autoreload 46 | 47 | 48 | 49 | 50 | 51 | 52 | 53 | 54 | 55 | Next, we prepare the dataset. The set of ODEs we consider here are 56 | $d[y, z]/dt = [z, -z+ 5 \sin y]$ 57 | 58 | 59 | ```python 60 | def dU_dt_sin(U, t): 61 | return [U[1], -1*U[1] - 5*np.sin(U[0])] 62 | U0 = [2.5, 0.4] 63 | ts = np.linspace(0, 5, 100) 64 | Y = odeint(dU_dt_sin, U0, ts) 65 | T = ts.reshape(-1,1) 66 | ``` 67 | 68 | Here we can potentially rescale the Y and T axis and we plot the results 69 | 70 | 71 | ```python 72 | T_rs = T/np.max(np.abs(T),axis=0) 73 | Y_rs = Y/np.max(np.abs(Y),axis=0) 74 | ``` 75 | 76 | Let's plot it to get an idea of the data: 77 | 78 | 79 | ```python 80 | fig, ax = plt.subplots() 81 | ax.plot(T_rs, Y_rs[:,0]) 82 | ax.plot(T_rs, Y_rs[:,1]) 83 | ax.set_xlabel('t') 84 | 85 | plt.show() 86 | ``` 87 | 88 | 89 | ![png](output_8_0.png) 90 | 91 | 92 | 93 | ```python 94 | number_of_samples = 500 95 | 96 | idx = np.random.permutation(Y.shape[0]) 97 | X = torch.tensor(T_rs[idx, :][:number_of_samples], dtype=torch.float32, requires_grad=True) 98 | y = torch.tensor(Y_rs[idx, :][:number_of_samples], dtype=torch.float32) 99 | ``` 100 | 101 | 102 | ```python 103 | print(X.shape, y.shape) 104 | ``` 105 | 106 | torch.Size([100, 1]) torch.Size([100, 2]) 107 | 108 | 109 | # Setup a custom library 110 | 111 | In this notebook we show how the user can create a custom build library.The library function, $\theta$, in this case contains $[1,u,v, sin(u)]$ to showcase that non-linear terms can easily be added to the library 112 | 113 | 114 | ```python 115 | from torch.autograd import grad 116 | from itertools import combinations, product 117 | from functools import reduce 118 | ``` 119 | 120 | 121 | ```python 122 | class Library_nonlinear(Library): 123 | """[summary] 124 | 125 | Args: 126 | Library ([type]): [description] 127 | """ 128 | def __init__(self) -> None: 129 | super().__init__() 130 | 131 | def library(self, input: Tuple[torch.Tensor, torch.Tensor]) -> Tuple[TensorList, TensorList]: 132 | 133 | prediction, data = input 134 | samples = prediction.shape[0] 135 | poly_list = [] 136 | deriv_list = [] 137 | time_deriv_list = [] 138 | 139 | 140 | # Construct the theta matrix 141 | C = torch.ones_like(prediction[:,0]).view(samples, -1) 142 | u = prediction[:,0].view(samples, -1) 143 | v = prediction[:,1].view(samples, -1) 144 | theta = torch.cat((C, u, v, torch.sin(u)),dim=1) 145 | 146 | # Construct a list of time_derivatives 147 | time_deriv_list = [] 148 | for output in torch.arange(prediction.shape[1]): 149 | dy = grad(prediction[:,output], data, grad_outputs=torch.ones_like(prediction[:,output]), create_graph=True)[0] 150 | time_deriv = dy[:, 0:1] 151 | time_deriv_list.append(time_deriv) 152 | 153 | return time_deriv_list, [theta,theta] 154 | 155 | ``` 156 | 157 | ## Configuring DeepMoD 158 | 159 | Configuration of the function approximator: Here the first argument is the number of input and the last argument the number of output layers. 160 | 161 | 162 | ```python 163 | network = NN(1, [30, 30, 30,30], 2) 164 | ``` 165 | 166 | Configuration of the library function: We select the custom build library we created earlier 167 | 168 | 169 | ```python 170 | library = Library_nonlinear() 171 | ``` 172 | 173 | Configuration of the sparsity estimator and sparsity scheduler used. In this case we use the most basic threshold-based Lasso estimator and a scheduler that asseses the validation loss after a given patience. If that value is smaller than 1e-5, the algorithm is converged. 174 | 175 | 176 | ```python 177 | estimator = Threshold(0.5) 178 | sparsity_scheduler = TrainTestPeriodic(periodicity=50, patience=200, delta=1e-5) 179 | ``` 180 | 181 | Configuration of the sparsity estimator 182 | 183 | 184 | ```python 185 | constraint = LeastSquares() 186 | # Configuration of the sparsity scheduler 187 | ``` 188 | 189 | Now we instantiate the model and select the optimizer 190 | 191 | 192 | ```python 193 | model = DeepMoD(network, library, estimator, constraint) 194 | 195 | # Defining optimizer 196 | optimizer = torch.optim.Adam(model.parameters(), betas=(0.99, 0.99), amsgrad=True, lr=1e-3) 197 | 198 | ``` 199 | 200 | ## Run DeepMoD 201 | 202 | We can now run DeepMoD using all the options we have set and the training data. We need to slightly preprocess the input data for the derivatives: 203 | 204 | 205 | ```python 206 | train(model, X, y, optimizer,sparsity_scheduler, log_dir='runs/coupled2/', split=0.8, max_iterations=100000, delta=1e-3, patience=100) 207 | ``` 208 | 209 | 21450 MSE: 2.99e-02 Reg: 3.16e-03 L1: 2.65e+00 Algorithm converged. Writing model to disk. 210 | 211 | 212 | Now that DeepMoD has converged, it has found the following numbers: 213 | 214 | 215 | ```python 216 | model.sparsity_masks 217 | ``` 218 | 219 | 220 | 221 | 222 | [tensor([False, False, True, False]), tensor([False, False, True, True])] 223 | 224 | 225 | 226 | 227 | ```python 228 | print(model.estimator_coeffs()) 229 | ``` 230 | 231 | [array([[0. ], 232 | [0. ], 233 | [0.99987924], 234 | [0. ]], dtype=float32), array([[ 0. ], 235 | [ 0. ], 236 | [-0.56510067], 237 | [-1.076641 ]], dtype=float32)] 238 | 239 | 240 | 241 | ```python 242 | 243 | ``` 244 | -------------------------------------------------------------------------------- /docs/examples/PDE_KdV/PDE_KdV.md: -------------------------------------------------------------------------------- 1 | # Example Korteweg de Vries equation 2 | 3 | In this notebook we provide a simple example of the DeepMoD algorithm by applying it on the KdV equation. 4 | 5 | 6 | ```python 7 | # General imports 8 | import numpy as np 9 | import torch 10 | import matplotlib.pylab as plt 11 | 12 | # DeepMoD functions 13 | 14 | from deepymod import DeepMoD 15 | from deepymod.model.func_approx import NN 16 | from deepymod.model.library import Library1D 17 | from deepymod.model.constraint import LeastSquares 18 | from deepymod.model.sparse_estimators import Threshold,PDEFIND 19 | from deepymod.training import train 20 | from deepymod.training.sparsity_scheduler import TrainTestPeriodic 21 | from scipy.io import loadmat 22 | 23 | 24 | # Settings for reproducibility 25 | np.random.seed(42) 26 | torch.manual_seed(0) 27 | 28 | 29 | %load_ext autoreload 30 | %autoreload 2 31 | ``` 32 | 33 | Next, we prepare the dataset. 34 | 35 | 36 | ```python 37 | data = np.load('data/kdv.npy', allow_pickle=True).item() 38 | print('Shape of grid:', data['x'].shape) 39 | ``` 40 | 41 | Shape of grid: (512, 201) 42 | 43 | 44 | Let's plot it to get an idea of the data: 45 | 46 | 47 | ```python 48 | fig, ax = plt.subplots() 49 | im = ax.contourf(data['x'], data['t'], np.real(data['u'])) 50 | ax.set_xlabel('x') 51 | ax.set_ylabel('t') 52 | fig.colorbar(mappable=im) 53 | 54 | plt.show() 55 | ``` 56 | 57 | 58 | ![png](output_6_0.png) 59 | 60 | 61 | 62 | ```python 63 | X = np.transpose((data['t'].flatten(), data['x'].flatten())) 64 | y = np.real(data['u']).reshape((data['u'].size, 1)) 65 | print(X.shape, y.shape) 66 | ``` 67 | 68 | (102912, 2) (102912, 1) 69 | 70 | 71 | As we can see, $X$ has 2 dimensions, $\{x, t\}$, while $y$ has only one, $\{u\}$. Always explicity set the shape (i.e. $N\times 1$, not $N$) or you'll get errors. This dataset is noiseless, so let's add $2.5\%$ noise: 72 | 73 | 74 | ```python 75 | noise_level = 0.025 76 | y_noisy = y + noise_level * np.std(y) * np.random.randn(y[:,0].size, 1) 77 | ``` 78 | 79 | The dataset is also much larger than needed, so let's hussle it and pick out a 1000 samples: 80 | 81 | 82 | ```python 83 | number_of_samples = 2000 84 | 85 | idx = np.random.permutation(y.shape[0]) 86 | X_train = torch.tensor(X[idx, :][:number_of_samples], dtype=torch.float32, requires_grad=True) 87 | y_train = torch.tensor(y_noisy[idx, :][:number_of_samples], dtype=torch.float32) 88 | ``` 89 | 90 | 91 | ```python 92 | print(X_train.shape, y_train.shape) 93 | ``` 94 | 95 | torch.Size([2000, 2]) torch.Size([2000, 1]) 96 | 97 | 98 | We now have a dataset which we can use. Let's plot, for a final time, the original dataset, the noisy set and the samples points: 99 | 100 | 101 | ```python 102 | fig, axes = plt.subplots(ncols=3, figsize=(15, 4)) 103 | 104 | im0 = axes[0].contourf(data['x'], data['t'], np.real(data['u']), cmap='coolwarm') 105 | axes[0].set_xlabel('x') 106 | axes[0].set_ylabel('t') 107 | axes[0].set_title('Ground truth') 108 | 109 | im1 = axes[1].contourf(data['x'], data['t'], y_noisy.reshape(data['x'].shape), cmap='coolwarm') 110 | axes[1].set_xlabel('x') 111 | axes[1].set_title('Noisy') 112 | 113 | sampled = np.array([y_noisy[index, 0] if index in idx[:number_of_samples] else np.nan for index in np.arange(data['x'].size)]) 114 | sampled = np.rot90(sampled.reshape(data['x'].shape)) #array needs to be rotated because of imshow 115 | 116 | im2 = axes[2].imshow(sampled, aspect='auto', cmap='coolwarm') 117 | axes[2].set_xlabel('x') 118 | axes[2].set_title('Sampled') 119 | 120 | fig.colorbar(im1, ax=axes.ravel().tolist()) 121 | 122 | plt.show() 123 | ``` 124 | 125 | 126 | ![png](output_14_0.png) 127 | 128 | 129 | ## Configuring DeepMoD 130 | 131 | Configuration of the function approximator: Here the first argument is the number of input and the last argument the number of output layers. 132 | 133 | 134 | ```python 135 | network = NN(2, [30, 30, 30, 30], 1) 136 | ``` 137 | 138 | Configuration of the library function: We select athe library with a 2D spatial input. Note that that the max differential order has been pre-determined here out of convinience. So, for poly_order 1 the library contains the following 12 terms: 139 | * [$1, u_x, u_{xx}, u_{xxx}, u, u u_{x}, u u_{xx}, u u_{xxx}, u^2, u^2 u_{x}, u^2 u_{xx}, u^2 u_{xxx}$] 140 | 141 | 142 | ```python 143 | library = Library1D(poly_order=2, diff_order=3) 144 | ``` 145 | 146 | Configuration of the sparsity estimator and sparsity scheduler used. In this case we use the most basic threshold-based Lasso estimator and a scheduler that asseses the validation loss after a given patience. If that value is smaller than 1e-5, the algorithm is converged. 147 | 148 | 149 | ```python 150 | estimator = Threshold(0.1) 151 | sparsity_scheduler = TrainTestPeriodic(periodicity=50, patience=10, delta=1e-5) 152 | ``` 153 | 154 | Configuration of the sparsity estimator 155 | 156 | 157 | ```python 158 | constraint = LeastSquares() 159 | # Configuration of the sparsity scheduler 160 | ``` 161 | 162 | Now we instantiate the model and select the optimizer 163 | 164 | 165 | ```python 166 | model = DeepMoD(network, library, estimator, constraint) 167 | 168 | # Defining optimizer 169 | optimizer = torch.optim.Adam(model.parameters(), betas=(0.99, 0.99), amsgrad=True, lr=1e-3) 170 | 171 | ``` 172 | 173 | ## Run DeepMoD 174 | 175 | We can now run DeepMoD using all the options we have set and the training data: 176 | * The directory where the tensorboard file is written (log_dir) 177 | * The ratio of train/test set used (split) 178 | * The maximum number of iterations performed (max_iterations) 179 | * The absolute change in L1 norm considered converged (delta) 180 | * The amount of epochs over which the absolute change in L1 norm is calculated (patience) 181 | 182 | 183 | ```python 184 | train(model, X_train, y_train, optimizer,sparsity_scheduler, log_dir='runs/KDV/', split=0.8, max_iterations=100000) 185 | ``` 186 | 187 | 11375 MSE: 9.29e-06 Reg: 3.11e-06 L1: 2.69e+00 Algorithm converged. Writing model to disk. 188 | 189 | 190 | Sparsity masks provide the active and non-active terms in the PDE: 191 | 192 | 193 | ```python 194 | model.sparsity_masks 195 | ``` 196 | 197 | 198 | 199 | 200 | [tensor([False, False, False, True, False, True, False, False, False, False, 201 | False, False])] 202 | 203 | 204 | 205 | estimatior_coeffs gives the magnitude of the active terms: 206 | 207 | 208 | ```python 209 | print(model.estimator_coeffs()) 210 | ``` 211 | 212 | [array([[ 0. ], 213 | [ 0. ], 214 | [ 0. ], 215 | [-0.82726973], 216 | [ 0. ], 217 | [-1.5996557 ], 218 | [ 0. ], 219 | [ 0. ], 220 | [ 0. ], 221 | [ 0. ], 222 | [ 0. ], 223 | [ 0. ]], dtype=float32)] 224 | 225 | -------------------------------------------------------------------------------- /src/deepymod/training/sparsity_scheduler.py: -------------------------------------------------------------------------------- 1 | """ Contains classes that schedule when the sparsity mask should be applied """ 2 | import torch 3 | import numpy as np 4 | 5 | 6 | class Periodic: 7 | """Periodically applies sparsity every periodicity iterations 8 | after initial_epoch. 9 | """ 10 | 11 | def __init__(self, periodicity=50, initial_iteration=1000): 12 | """Periodically applies sparsity every periodicity iterations 13 | after initial_epoch. 14 | Args: 15 | periodicity (int): after initial_iterations, apply sparsity mask per periodicity epochs 16 | initial_iteration (int): wait initial_iterations before applying sparsity 17 | """ 18 | self.periodicity = periodicity 19 | self.initial_iteration = initial_iteration 20 | 21 | def __call__(self, iteration, loss, model, optimizer): 22 | # Update periodically 23 | apply_sparsity = False # we overwrite it if we need to update 24 | 25 | if (iteration - self.initial_iteration) % self.periodicity == 0: 26 | apply_sparsity = True 27 | 28 | return apply_sparsity 29 | 30 | 31 | class TrainTest: 32 | """Early stops the training if validation loss doesn't improve after a given patience. 33 | Note that periodicity should be multitude of write_iterations.""" 34 | 35 | def __init__(self, patience=200, delta=1e-5, path="checkpoint.pt"): 36 | """Early stops the training if validation loss doesn't improve after a given patience. 37 | Note that periodicity should be multitude of write_iterations. 38 | Args: 39 | patience (int): wait patience epochs before checking TrainTest 40 | delta (float): desired accuracy 41 | path (str): pathname where to store the savepoints, must have ".pt" extension 42 | """ 43 | self.path = path 44 | self.patience = patience 45 | self.delta = delta 46 | 47 | self.best_iteration = None 48 | self.best_loss = None 49 | 50 | def __call__(self, iteration, loss, model, optimizer): 51 | apply_sparsity = False # we overwrite it if we need to update 52 | 53 | # Initialize if doesnt exist yet 54 | if self.best_loss is None: 55 | self.best_loss = loss 56 | self.best_iteration = iteration 57 | self.save_checkpoint(model, optimizer) 58 | 59 | # If it didnt improve, check if we're past patience 60 | elif (self.best_loss - loss) < self.delta: 61 | if (iteration - self.best_iteration) >= self.patience: 62 | # We reload the model to the best point and reset the scheduler 63 | self.load_checkpoint(model, optimizer) # reload model to best point 64 | self.best_loss = None 65 | self.best_iteration = None 66 | apply_sparsity = True 67 | 68 | # If not, keep going 69 | else: 70 | self.best_loss = loss 71 | self.best_iteration = iteration 72 | self.save_checkpoint(model, optimizer) 73 | 74 | return apply_sparsity 75 | 76 | def save_checkpoint(self, model, optimizer): 77 | """Saves model when validation loss decrease.""" 78 | checkpoint_path = self.path + "checkpoint.pt" 79 | torch.save( 80 | { 81 | "model_state_dict": model.state_dict(), 82 | "optimizer_state_dict": optimizer.state_dict(), 83 | }, 84 | checkpoint_path, 85 | ) 86 | 87 | def load_checkpoint(self, model, optimizer): 88 | """Loads model from disk""" 89 | checkpoint_path = self.path + "checkpoint.pt" 90 | checkpoint = torch.load(checkpoint_path) 91 | model.load_state_dict(checkpoint["model_state_dict"]) 92 | optimizer.load_state_dict(checkpoint["optimizer_state_dict"]) 93 | 94 | 95 | class TrainTestPeriodic: 96 | """Early stops the training if validation loss doesn't improve after a given patience. 97 | Note that periodicity should be multitude of write_iterations.""" 98 | 99 | def __init__(self, periodicity=50, patience=200, delta=1e-5, path="checkpoint.pt"): 100 | """Early stops the training if validation loss doesn't improve after a given patience. 101 | Note that periodicity should be multitude of write_iterations. 102 | Args: 103 | periodicity (int): apply sparsity mask per periodicity epochs 104 | patience (int): wait patience epochs before checking TrainTest 105 | delta (float): desired accuracy 106 | path (str): pathname where to store the savepoints, must have ".pt" extension 107 | """ 108 | self.path = path 109 | self.patience = patience 110 | self.delta = delta 111 | self.periodicity = periodicity 112 | 113 | self.best_iteration = None 114 | self.best_loss = None 115 | self.periodic = False 116 | 117 | def __call__(self, iteration, loss, model, optimizer): 118 | # Update periodically if we have updated once 119 | apply_sparsity = False # we overwrite it if we need to update 120 | 121 | if self.periodic is True: 122 | if (iteration - self.best_iteration) % self.periodicity == 0: 123 | apply_sparsity = True 124 | 125 | # Check for improvements if we havent updated yet. 126 | # Initialize if doesnt exist yet 127 | elif self.best_loss is None: 128 | self.best_loss = loss 129 | self.best_iteration = iteration 130 | self.save_checkpoint(model, optimizer) 131 | 132 | # If it didnt improve, check if we're past patience 133 | elif (self.best_loss - loss) < self.delta: 134 | if (iteration - self.best_iteration) >= self.patience: 135 | self.load_checkpoint(model, optimizer) # reload model to best point 136 | self.periodic = True # switch to periodic regime 137 | self.best_iteration = iteration # because the iterator doesnt reset 138 | apply_sparsity = True 139 | 140 | # If not, keep going 141 | else: 142 | self.best_loss = loss 143 | self.best_iteration = iteration 144 | self.save_checkpoint(model, optimizer) 145 | 146 | return apply_sparsity 147 | 148 | def save_checkpoint(self, model, optimizer): 149 | """Saves model when validation loss decrease.""" 150 | checkpoint_path = self.path + "checkpoint.pt" 151 | torch.save( 152 | { 153 | "model_state_dict": model.state_dict(), 154 | "optimizer_state_dict": optimizer.state_dict(), 155 | }, 156 | checkpoint_path, 157 | ) 158 | 159 | def load_checkpoint(self, model, optimizer): 160 | """Loads model from disk""" 161 | checkpoint_path = self.path + "checkpoint.pt" 162 | checkpoint = torch.load(checkpoint_path) 163 | model.load_state_dict(checkpoint["model_state_dict"]) 164 | optimizer.load_state_dict(checkpoint["optimizer_state_dict"]) 165 | -------------------------------------------------------------------------------- /docs/examples/Burgers/PDE_Burgers.md: -------------------------------------------------------------------------------- 1 | # Example Burgers' equation 2 | 3 | In this notebook we provide a simple example of the DeepMoD algorithm by applying it on the Burgers' equation. 4 | 5 | We start by importing the required libraries and setting the plotting style: 6 | 7 | 8 | ```python 9 | # General imports 10 | import numpy as np 11 | import torch 12 | import matplotlib.pylab as plt 13 | 14 | # DeepMoD functions 15 | 16 | from deepymod import DeepMoD 17 | from deepymod.model.func_approx import NN 18 | from deepymod.model.library import Library1D 19 | from deepymod.model.constraint import LeastSquares 20 | from deepymod.model.sparse_estimators import Threshold,PDEFIND 21 | from deepymod.training import train 22 | from deepymod.training.sparsity_scheduler import TrainTestPeriodic 23 | from scipy.io import loadmat 24 | 25 | # Settings for reproducibility 26 | np.random.seed(42) 27 | torch.manual_seed(0) 28 | 29 | 30 | %load_ext autoreload 31 | %autoreload 2 32 | ``` 33 | 34 | Next, we prepare the dataset. 35 | 36 | 37 | ```python 38 | data = np.load('data/burgers.npy', allow_pickle=True).item() 39 | print('Shape of grid:', data['x'].shape) 40 | ``` 41 | 42 | Shape of grid: (256, 101) 43 | 44 | 45 | Let's plot it to get an idea of the data: 46 | 47 | 48 | ```python 49 | fig, ax = plt.subplots() 50 | im = ax.contourf(data['x'], data['t'], np.real(data['u'])) 51 | ax.set_xlabel('x') 52 | ax.set_ylabel('t') 53 | fig.colorbar(mappable=im) 54 | 55 | plt.show() 56 | ``` 57 | 58 | 59 | ![png](output_6_0.png) 60 | 61 | 62 | 63 | ```python 64 | X = np.transpose((data['t'].flatten(), data['x'].flatten())) 65 | y = np.real(data['u']).reshape((data['u'].size, 1)) 66 | print(X.shape, y.shape) 67 | ``` 68 | 69 | (25856, 2) (25856, 1) 70 | 71 | 72 | As we can see, $X$ has 2 dimensions, $\{x, t\}$, while $y$ has only one, $\{u\}$. Always explicity set the shape (i.e. $N\times 1$, not $N$) or you'll get errors. This dataset is noiseless, so let's add $2.5\%$ noise: 73 | 74 | 75 | ```python 76 | noise_level = 0.025 77 | y_noisy = y + noise_level * np.std(y) * np.random.randn(y[:,0].size, 1) 78 | ``` 79 | 80 | The dataset is also much larger than needed, so let's hussle it and pick out a 1000 samples: 81 | 82 | 83 | ```python 84 | number_of_samples = 2000 85 | 86 | idx = np.random.permutation(y.shape[0]) 87 | X_train = torch.tensor(X[idx, :][:number_of_samples], dtype=torch.float32, requires_grad=True) 88 | y_train = torch.tensor(y_noisy[idx, :][:number_of_samples], dtype=torch.float32) 89 | ``` 90 | 91 | 92 | ```python 93 | print(X_train.shape, y_train.shape) 94 | ``` 95 | 96 | torch.Size([2000, 2]) torch.Size([2000, 1]) 97 | 98 | 99 | We now have a dataset which we can use. Let's plot, for a final time, the original dataset, the noisy set and the samples points: 100 | 101 | 102 | ```python 103 | fig, axes = plt.subplots(ncols=3, figsize=(15, 4)) 104 | 105 | im0 = axes[0].contourf(data['x'], data['t'], np.real(data['u']), cmap='coolwarm') 106 | axes[0].set_xlabel('x') 107 | axes[0].set_ylabel('t') 108 | axes[0].set_title('Ground truth') 109 | 110 | im1 = axes[1].contourf(data['x'], data['t'], y_noisy.reshape(data['x'].shape), cmap='coolwarm') 111 | axes[1].set_xlabel('x') 112 | axes[1].set_title('Noisy') 113 | 114 | sampled = np.array([y_noisy[index, 0] if index in idx[:number_of_samples] else np.nan for index in np.arange(data['x'].size)]) 115 | sampled = np.rot90(sampled.reshape(data['x'].shape)) #array needs to be rotated because of imshow 116 | 117 | im2 = axes[2].imshow(sampled, aspect='auto', cmap='coolwarm') 118 | axes[2].set_xlabel('x') 119 | axes[2].set_title('Sampled') 120 | 121 | fig.colorbar(im1, ax=axes.ravel().tolist()) 122 | 123 | plt.show() 124 | ``` 125 | 126 | 127 | ![png](output_14_0.png) 128 | 129 | 130 | ## Configuring DeepMoD 131 | 132 | Configuration of the function approximator: Here the first argument is the number of input and the last argument the number of output layers. 133 | 134 | 135 | ```python 136 | network = NN(2, [30, 30, 30, 30], 1) 137 | ``` 138 | 139 | Configuration of the library function: We select athe library with a 2D spatial input. Note that that the max differential order has been pre-determined here out of convinience. So, for poly_order 1 the library contains the following 12 terms: 140 | * [$1, u_x, u_{xx}, u_{xxx}, u, u u_{x}, u u_{xx}, u u_{xxx}, u^2, u^2 u_{x}, u^2 u_{xx}, u^2 u_{xxx}$] 141 | 142 | 143 | ```python 144 | library = Library1D(poly_order=2, diff_order=3) 145 | ``` 146 | 147 | Configuration of the sparsity estimator and sparsity scheduler used. In this case we use the most basic threshold-based Lasso estimator and a scheduler that asseses the validation loss after a given patience. If that value is smaller than 1e-5, the algorithm is converged. 148 | 149 | 150 | ```python 151 | estimator = Threshold(0.1) 152 | sparsity_scheduler = TrainTestPeriodic(periodicity=50, patience=200, delta=1e-5) 153 | ``` 154 | 155 | Configuration of the sparsity estimator 156 | 157 | 158 | ```python 159 | constraint = LeastSquares() 160 | # Configuration of the sparsity scheduler 161 | ``` 162 | 163 | Now we instantiate the model and select the optimizer 164 | 165 | 166 | ```python 167 | model = DeepMoD(network, library, estimator, constraint) 168 | 169 | # Defining optimizer 170 | optimizer = torch.optim.Adam(model.parameters(), betas=(0.99, 0.99), amsgrad=True, lr=1e-3) 171 | 172 | ``` 173 | 174 | ## Run DeepMoD 175 | 176 | We can now run DeepMoD using all the options we have set and the training data: 177 | * The directory where the tensorboard file is written (log_dir) 178 | * The ratio of train/test set used (split) 179 | * The maximum number of iterations performed (max_iterations) 180 | * The absolute change in L1 norm considered converged (delta) 181 | * The amount of epochs over which the absolute change in L1 norm is calculated (patience) 182 | 183 | 184 | ```python 185 | train(model, X_train, y_train, optimizer,sparsity_scheduler, log_dir='runs/Burgers/', split=0.8, max_iterations=100000) 186 | ``` 187 | 188 | 13350 MSE: 2.53e-05 Reg: 1.38e-05 L1: 1.45e+00 Algorithm converged. Writing model to disk. 189 | 190 | 191 | Sparsity masks provide the active and non-active terms in the PDE: 192 | 193 | 194 | ```python 195 | model.sparsity_masks 196 | ``` 197 | 198 | 199 | 200 | 201 | [tensor([False, False, True, False, False, True, False, False, False, False, 202 | False, False])] 203 | 204 | 205 | 206 | estimatior_coeffs gives the magnitude of the active terms: 207 | 208 | 209 | ```python 210 | print(model.estimator_coeffs()) 211 | ``` 212 | 213 | [array([[ 0. ], 214 | [ 0. ], 215 | [ 0.39227325], 216 | [ 0. ], 217 | [ 0. ], 218 | [-1.001875 ], 219 | [ 0. ], 220 | [ 0. ], 221 | [ 0. ], 222 | [ 0. ], 223 | [ 0. ], 224 | [ 0. ]], dtype=float32)] 225 | 226 | 227 | So the final terms that remain are the $u_{xx}$ and $u u_{x}$ resulting in the following Burgers equation (in normalized coefficients: 228 | $u_t = 0.4 u_{xx} - u u_{x}$. 229 | 230 | 231 | ```python 232 | 233 | ``` 234 | -------------------------------------------------------------------------------- /src/deepymod/model/func_approx.py: -------------------------------------------------------------------------------- 1 | """ This files contains the function approximators that are used by DeepMoD. 2 | """ 3 | 4 | 5 | import torch 6 | import torch.nn as nn 7 | from typing import List, Tuple 8 | import numpy as np 9 | 10 | 11 | class NN(nn.Module): 12 | def __init__(self, n_in: int, n_hidden: List[int], n_out: int) -> None: 13 | """Constructs a feed-forward neural network with tanh activation. 14 | 15 | Args: 16 | n_in (int): Number of input features. 17 | n_hidden (List[int]): Number of neurons in each layer. 18 | n_out (int): Number of output features. 19 | """ 20 | super().__init__() 21 | self.network = self.build_network(n_in, n_hidden, n_out) 22 | 23 | def forward(self, input: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: 24 | """Forward pass through the network. Returns prediction and the differentiable input 25 | so we can construct the library. 26 | 27 | Args: 28 | input (torch.Tensor): Input tensor of size (n_samples, n_inputs). 29 | 30 | Returns: 31 | (torch.Tensor, torch.Tensor): prediction of size (n_samples, n_outputs) and coordinates of size (n_samples, n_inputs). 32 | """ 33 | coordinates = input.clone().detach().requires_grad_(True) 34 | return self.network(coordinates), coordinates 35 | 36 | def build_network( 37 | self, n_in: int, n_hidden: List[int], n_out: int 38 | ) -> torch.nn.Sequential: 39 | """Constructs a feed-forward neural network. 40 | 41 | Args: 42 | n_in (int): Number of input features. 43 | n_hidden (list[int]): Number of neurons in each layer. 44 | n_out (int): Number of output features. 45 | 46 | Returns: 47 | torch.Sequential: Pytorch module 48 | """ 49 | 50 | network = [] 51 | architecture = [n_in] + n_hidden + [n_out] 52 | for layer_i, layer_j in zip(architecture, architecture[1:]): 53 | network.append(nn.Linear(layer_i, layer_j)) 54 | network.append(nn.Tanh()) 55 | network.pop() # get rid of last activation function 56 | return nn.Sequential(*network) 57 | 58 | 59 | class SineLayer(nn.Module): 60 | def __init__( 61 | self, 62 | in_features: int, 63 | out_features: int, 64 | omega_0: float = 30, 65 | is_first: bool = False, 66 | ) -> None: 67 | """Sine activation function layer with omega_0 scaling. 68 | 69 | Args: 70 | in_features (int): Number of input features. 71 | out_features (int): Number of output features. 72 | omega_0 (float, optional): Scaling factor of the Sine function. Defaults to 30. 73 | is_first (bool, optional): Defaults to False. 74 | """ 75 | super().__init__() 76 | self.omega_0 = omega_0 77 | self.is_first = is_first 78 | 79 | self.in_features = in_features 80 | self.linear = nn.Linear(in_features, out_features) 81 | 82 | self.init_weights() 83 | 84 | def init_weights(self) -> None: 85 | """Initialization of the weigths.""" 86 | with torch.no_grad(): 87 | if self.is_first: 88 | self.linear.weight.uniform_(-1 / self.in_features, 1 / self.in_features) 89 | else: 90 | self.linear.weight.uniform_( 91 | -np.sqrt(6 / self.in_features) / self.omega_0, 92 | np.sqrt(6 / self.in_features) / self.omega_0, 93 | ) 94 | 95 | def forward(self, input: torch.Tensor) -> torch.Tensor: 96 | """Forward pass through the layer. 97 | 98 | Args: 99 | input (torch.Tensor): Input tensor of shape (n_samples, n_inputs). 100 | 101 | Returns: 102 | torch.Tensor: Prediction of shape (n_samples, n_outputs) 103 | """ 104 | return torch.sin(self.omega_0 * self.linear(input)) 105 | 106 | 107 | class Siren(nn.Module): 108 | def __init__( 109 | self, 110 | n_in: int, 111 | n_hidden: List[int], 112 | n_out: int, 113 | first_omega_0: float = 30.0, 114 | hidden_omega_0: float = 30.0, 115 | ) -> None: 116 | """SIREN model from the paper [Implicit Neural Representations with 117 | Periodic Activation Functions](https://arxiv.org/abs/2006.09661). 118 | 119 | Args: 120 | n_in (int): Number of input features. 121 | n_hidden (list[int]): Number of neurons in each layer. 122 | n_out (int): Number of output features. 123 | first_omega_0 (float, optional): Scaling factor of the Sine function of the first layer. Defaults to 30. 124 | hidden_omega_0 (float, optional): Scaling factor of the Sine function of the hidden layers. Defaults to 30. 125 | """ 126 | super().__init__() 127 | self.network = self.build_network( 128 | n_in, n_hidden, n_out, first_omega_0, hidden_omega_0 129 | ) 130 | 131 | def forward(self, input: torch.Tensor) -> torch.Tensor: 132 | """Forward pass through the network. 133 | 134 | Args: 135 | input (torch.Tensor): Input tensor of shape (n_samples, n_inputs). 136 | 137 | Returns: 138 | torch.Tensor: Prediction of shape (n_samples, n_outputs) 139 | """ 140 | coordinates = input.clone().detach().requires_grad_(True) 141 | return self.network(coordinates), coordinates 142 | 143 | def build_network( 144 | self, 145 | n_in: int, 146 | n_hidden: List[int], 147 | n_out: int, 148 | first_omega_0: float, 149 | hidden_omega_0: float, 150 | ) -> torch.nn.Sequential: 151 | """Constructs the Siren neural network. 152 | 153 | Args: 154 | n_in (int): Number of input features. 155 | n_hidden (list[int]): Number of neurons in each layer. 156 | n_out (int): Number of output features. 157 | first_omega_0 (float, optional): Scaling factor of the Sine function of the first layer. Defaults to 30. 158 | hidden_omega_0 (float, optional): Scaling factor of the Sine function of the hidden layers. Defaults to 30. 159 | Returns: 160 | torch.Sequential: Pytorch module 161 | """ 162 | network = [] 163 | # Input layer 164 | network.append( 165 | SineLayer(n_in, n_hidden[0], is_first=True, omega_0=first_omega_0) 166 | ) 167 | 168 | # Hidden layers 169 | for layer_i, layer_j in zip(n_hidden, n_hidden[1:]): 170 | network.append( 171 | SineLayer(layer_i, layer_j, is_first=False, omega_0=hidden_omega_0) 172 | ) 173 | 174 | # Output layer 175 | final_linear = nn.Linear(n_hidden[-1], n_out) 176 | with torch.no_grad(): 177 | final_linear.weight.uniform_( 178 | -np.sqrt(6 / n_hidden[-1]) / hidden_omega_0, 179 | np.sqrt(6 / n_hidden[-1]) / hidden_omega_0, 180 | ) 181 | network.append(final_linear) 182 | 183 | return nn.Sequential(*network) 184 | -------------------------------------------------------------------------------- /src/deepymod/analysis/load_tensorboard.py: -------------------------------------------------------------------------------- 1 | """ Constains a tool to convert from Tensorboard to Pandas DataFrame """ 2 | 3 | import pandas as pd 4 | from tensorboard.backend.event_processing.event_accumulator import EventAccumulator 5 | import os 6 | from natsort import natsorted 7 | import matplotlib.pyplot as plt 8 | 9 | 10 | def load_tensorboard(path: str) -> pd.DataFrame: 11 | """Loads tensorboard files into a pandas dataframe. Assumes one run per folder! 12 | 13 | Args: 14 | path (string): path of folder with tensorboard files. 15 | 16 | Returns: 17 | DataFrame: Pandas dataframe with all run data. 18 | """ 19 | 20 | event_paths = [ 21 | file 22 | for file in os.walk(path, topdown=True) 23 | if file[2][0][: len("events")] == "events" 24 | ] 25 | 26 | df = pd.DataFrame() 27 | steps = None # steps are the same for all files 28 | 29 | for event_idx, path in enumerate(event_paths): 30 | summary_iterator = EventAccumulator(os.path.join(path[0], path[2][0])).Reload() 31 | tags = summary_iterator.Tags()["scalars"] 32 | data = [ 33 | [event.value for event in summary_iterator.Scalars(tag)] for tag in tags 34 | ] 35 | if steps is None: 36 | steps = [event.step for event in summary_iterator.Scalars(tags[0])] 37 | 38 | # Adding to dataframe 39 | tags = [tag.replace("/", "_") for tag in tags] # for name consistency 40 | if ( 41 | event_idx > 0 42 | ): # We have one file in the top level, so after we need to use folder name 43 | tags = [path[0].split("/")[-1]] 44 | 45 | for idx, tag in enumerate(tags): 46 | try: 47 | df[tag] = data[idx] 48 | except ValueError: # more debugging info 49 | print( 50 | f"Warning: Either the {tag = } of `df` or {idx = } of `data` do not exist! Check for pre-existing saved files. " 51 | ) 52 | df.index = steps 53 | return df 54 | 55 | 56 | def plot_history(foldername: str): 57 | """Plots the training history of the model.""" 58 | history = load_tensorboard(foldername) 59 | fig, axs = plt.subplots(1, 4, figsize=(20, 5)) 60 | 61 | for history_key in history.keys(): 62 | history_key_parts = history_key.split("_") 63 | if history_key_parts[0] == "loss": 64 | if history_key_parts[-1] == "0": 65 | axs[0].semilogy( 66 | history[history_key], 67 | label=history_key_parts[1] + "_" + history_key_parts[-1], 68 | linestyle="--", 69 | ) 70 | elif history_key_parts[-1] == "1": 71 | axs[0].semilogy( 72 | history[history_key], 73 | label=history_key_parts[1] + "_" + history_key_parts[-1], 74 | linestyle=":", 75 | ) 76 | else: 77 | axs[0].semilogy( 78 | history[history_key], 79 | label=history_key_parts[1] + "_" + history_key_parts[-1], 80 | linestyle="-", 81 | ) 82 | if history_key_parts[0] == "remaining": 83 | axs[0].semilogy( 84 | history[history_key], 85 | label=history_key_parts[1] 86 | + "_" 87 | + history_key_parts[3] 88 | + "_" 89 | + history_key_parts[4], 90 | linestyle="-.", 91 | ) 92 | if history_key_parts[0] == "coeffs": 93 | if history_key_parts[2] == "0": 94 | axs[1].plot( 95 | history[history_key], 96 | label=history_key_parts[2] 97 | + "_" 98 | + history_key_parts[3] 99 | + "_" 100 | + history_key_parts[4], 101 | linestyle="--", 102 | ) 103 | elif history_key_parts[2] == "1": 104 | axs[1].plot( 105 | history[history_key], 106 | label=history_key_parts[2] 107 | + "_" 108 | + history_key_parts[3] 109 | + "_" 110 | + history_key_parts[4], 111 | linestyle=":", 112 | ) 113 | else: 114 | axs[1].plot( 115 | history[history_key], 116 | label=history_key_parts[2] 117 | + "_" 118 | + history_key_parts[3] 119 | + "_" 120 | + history_key_parts[4], 121 | linestyle="-", 122 | ) 123 | if history_key_parts[0] == "unscaled": 124 | if history_key_parts[3] == "0": 125 | axs[2].plot( 126 | history[history_key], 127 | label=history_key_parts[3] 128 | + "_" 129 | + history_key_parts[4] 130 | + "_" 131 | + history_key_parts[5], 132 | linestyle="--", 133 | ) 134 | elif history_key_parts[3] == "1": 135 | axs[2].plot( 136 | history[history_key], 137 | label=history_key_parts[3] 138 | + "_" 139 | + history_key_parts[4] 140 | + "_" 141 | + history_key_parts[5], 142 | linestyle=":", 143 | ) 144 | else: 145 | axs[2].plot( 146 | history[history_key], 147 | label=history_key_parts[3] 148 | + "_" 149 | + history_key_parts[4] 150 | + "_" 151 | + history_key_parts[5], 152 | linestyle="-", 153 | ) 154 | if history_key_parts[0] == "estimator": 155 | if history_key_parts[3] == "0": 156 | axs[3].plot( 157 | history[history_key], 158 | label=history_key_parts[3] 159 | + "_" 160 | + history_key_parts[4] 161 | + "_" 162 | + history_key_parts[5], 163 | linestyle="--", 164 | ) 165 | elif history_key_parts[3] == "1": 166 | axs[3].plot( 167 | history[history_key], 168 | label=history_key_parts[3] 169 | + "_" 170 | + history_key_parts[4] 171 | + "_" 172 | + history_key_parts[5], 173 | linestyle=":", 174 | ) 175 | else: 176 | axs[3].plot( 177 | history[history_key], 178 | label=history_key_parts[3] 179 | + "_" 180 | + history_key_parts[4] 181 | + "_" 182 | + history_key_parts[5], 183 | linestyle="-", 184 | ) 185 | 186 | # axs[0].set_ylim([-2, 2]) 187 | axs[1].set_ylim([-2, 2]) 188 | axs[2].set_ylim([-2, 2]) 189 | axs[3].set_ylim([-2, 2]) 190 | 191 | axs[0].legend() 192 | axs[1].legend() 193 | axs[2].legend() 194 | axs[3].legend() 195 | 196 | plt.show() 197 | -------------------------------------------------------------------------------- /src/deepymod/model/sparse_estimators.py: -------------------------------------------------------------------------------- 1 | """Sparsity estimators which can be plugged into deepmod. 2 | We keep the API in line with scikit learn (mostly), so scikit learn can also be plugged in. 3 | See scikitlearn.linear_models for applicable estimators.""" 4 | 5 | import numpy as np 6 | from .deepmod import Estimator 7 | from sklearn.cluster import KMeans 8 | from pysindy.optimizers import STLSQ 9 | from sklearn.linear_model import LassoCV 10 | from sklearn.model_selection import train_test_split 11 | from sklearn.base import BaseEstimator 12 | 13 | import warnings 14 | 15 | warnings.filterwarnings( 16 | "ignore", category=UserWarning 17 | ) # To silence annoying pysindy warnings 18 | 19 | 20 | class Base(Estimator): 21 | def __init__(self, estimator: BaseEstimator) -> None: 22 | """Basic sparse estimator class; simply a wrapper around the supplied sk-learn compatible estimator. 23 | 24 | Args: 25 | estimator (BaseEstimator): Sci-kit learn estimator. 26 | """ 27 | super().__init__() 28 | self.estimator = estimator 29 | self.estimator.set_params( 30 | fit_intercept=False 31 | ) # Library contains offset so turn off the intercept 32 | 33 | def fit(self, X: np.ndarray, y: np.ndarray) -> np.ndarray: 34 | """Returns an array with the coefficient verctor after sparsity estimation. 35 | 36 | Args: 37 | X (np.ndarray): Training input data of shape (n_samples, n_features). 38 | y (np.ndarray): Training target data of shape (n_samples, n_outputs). 39 | 40 | Returns: 41 | np.ndarray: Coefficient vector (n_features, n_outputs). 42 | """ 43 | coeffs = self.estimator.fit(X, y).coef_ 44 | return coeffs 45 | 46 | 47 | class Threshold(Estimator): 48 | def __init__( 49 | self, 50 | threshold: float = 0.1, 51 | estimator: BaseEstimator = LassoCV(cv=5, fit_intercept=False), 52 | ) -> None: 53 | """Performs additional thresholding on coefficient result from supplied estimator. 54 | 55 | Args: 56 | threshold (float, optional): Value of the threshold above which the terms are selected. Defaults to 0.1. 57 | estimator (BaseEstimator, optional): Sparsity estimator. Defaults to LassoCV(cv=5, fit_intercept=False). 58 | """ 59 | super().__init__() 60 | self.estimator = estimator 61 | self.threshold = threshold 62 | 63 | # Library contains offset so turn off the intercept 64 | self.estimator.set_params(fit_intercept=False) 65 | 66 | def fit(self, X: np.ndarray, y: np.ndarray) -> np.ndarray: 67 | """Returns an array with the coefficient verctor after sparsity estimation. 68 | 69 | Args: 70 | X (np.ndarray): Training input data of shape (n_samples, n_features). 71 | y (np.ndarray): Training target data of shape (n_samples, n_outputs). 72 | 73 | Returns: 74 | np.ndarray: Coefficient vector (n_features, n_outputs). 75 | """ 76 | coeffs = self.estimator.fit(X, y).coef_ 77 | coeffs[np.abs(coeffs) < self.threshold] = 0.0 78 | 79 | return coeffs 80 | 81 | 82 | class Clustering(Estimator): 83 | def __init__( 84 | self, estimator: BaseEstimator = LassoCV(cv=5, fit_intercept=False) 85 | ) -> None: 86 | """Performs additional thresholding by Kmeans-clustering on coefficient result from estimator. 87 | 88 | Args: 89 | estimator (BaseEstimator, optional): Estimator class. Defaults to LassoCV(cv=5, fit_intercept=False). 90 | """ 91 | super().__init__() 92 | self.estimator = estimator 93 | self.kmeans = KMeans(n_clusters=2) 94 | 95 | # Library contains offset so turn off the intercept 96 | self.estimator.set_params(fit_intercept=False) 97 | 98 | def fit(self, X: np.ndarray, y: np.ndarray) -> np.ndarray: 99 | """Returns an array with the coefficient verctor after sparsity estimation. 100 | 101 | Args: 102 | X (np.ndarray): Training input data of shape (n_samples, n_features). 103 | y (np.ndarray): Training target data of shape (n_samples, n_outputs). 104 | 105 | Returns: 106 | np.ndarray: Coefficient vector (n_features, n_outputs). 107 | """ 108 | coeffs = self.estimator.fit(X, y).coef_[:, None] # sklearn returns 1D 109 | clusters = self.kmeans.fit_predict(np.abs(coeffs)).astype(np.bool) 110 | 111 | # make sure terms to keep are 1 and to remove are 0 112 | max_idx = np.argmax(np.abs(coeffs)) 113 | if clusters[max_idx] != 1: 114 | clusters = ~clusters 115 | 116 | coeffs = clusters.astype(np.float32) 117 | return coeffs 118 | 119 | 120 | class PDEFIND(Estimator): 121 | def __init__(self, lam: float = 1e-3, dtol: float = 0.1) -> None: 122 | """Implements PDEFIND as a sparse estimator. 123 | 124 | Args: 125 | lam (float, optional): Magnitude of the L2 regularization. Defaults to 1e-3. 126 | dtol (float, optional): Initial stepsize for the search of the thresholdDefaults to 0.1. 127 | """ 128 | super().__init__() 129 | self.lam = lam 130 | self.dtol = dtol 131 | 132 | def fit(self, X: np.ndarray, y: np.ndarray) -> np.ndarray: 133 | """Returns an array with the coefficient verctor after sparsity estimation. 134 | 135 | Args: 136 | X (np.ndarray): Training input data of shape (n_samples, n_features). 137 | y (np.ndarray): Training target data of shape (n_samples, n_outputs). 138 | 139 | Returns: 140 | np.ndarray: Coefficient vector (n_features, n_outputs). 141 | """ 142 | 143 | coeffs = PDEFIND.TrainSTLSQ(X, y[:, None], self.lam, self.dtol) 144 | return coeffs.squeeze() 145 | 146 | @staticmethod 147 | def TrainSTLSQ( 148 | X: np.ndarray, 149 | y: np.ndarray, 150 | alpha: float, 151 | delta_threshold: float, 152 | max_iterations: int = 100, 153 | test_size: float = 0.2, 154 | random_state: int = 0, 155 | ) -> np.ndarray: 156 | """PDE-FIND sparsity selection algorithm. Based on method described by Rudy et al. (10.1126/sciadv.1602614). 157 | 158 | Args: 159 | X (np.ndarray): Training input data of shape (n_samples, n_features). 160 | y (np.ndarray): Training target data of shape (n_samples, n_outputs). 161 | alpha (float): Magnitude of the L2 regularization. 162 | delta_threshold (float): Initial stepsize for the search of the threshold 163 | max_iterations (int, optional): Maximum number of iterations. Defaults to 100. 164 | test_size (float, optional): Fraction of the data that is assigned to the test-set. Defaults to 0.2. 165 | random_state (int, optional): Defaults to 0. 166 | 167 | Returns: 168 | np.ndarray: Coefficient vector. 169 | """ 170 | # Split data 171 | X_train, X_test, y_train, y_test = train_test_split( 172 | X, y, test_size=test_size, random_state=random_state 173 | ) 174 | 175 | # Set up the initial tolerance l0 penalty and estimates 176 | l0 = 1e-3 * np.linalg.cond(X) 177 | delta_t = delta_threshold # for interal use, can be updated 178 | 179 | # Initial estimate 180 | optimizer = STLSQ( 181 | threshold=0, alpha=0.0, fit_intercept=False 182 | ) # Now similar to LSTSQ 183 | y_predict = optimizer.fit(X_train, y_train).predict(X_test) 184 | min_loss = np.linalg.norm(y_predict - y_test, 2) + l0 * np.count_nonzero( 185 | optimizer.coef_ 186 | ) 187 | 188 | # Setting alpha and tolerance 189 | best_threshold = delta_t 190 | threshold = delta_t 191 | 192 | for iteration in np.arange(max_iterations): 193 | optimizer.set_params(alpha=alpha, threshold=threshold) 194 | y_predict = optimizer.fit(X_train, y_train).predict(X_test) 195 | loss = np.linalg.norm(y_predict - y_test, 2) + l0 * np.count_nonzero( 196 | optimizer.coef_ 197 | ) 198 | 199 | if (loss <= min_loss) and not (np.all(optimizer.coef_ == 0)): 200 | min_loss = loss 201 | best_threshold = threshold 202 | threshold += delta_threshold 203 | 204 | else: # if loss increases, we need to a) lower the current threshold and/or decrease step size 205 | new_lower_threshold = np.max([0, threshold - 2 * delta_t]) 206 | delta_t = 2 * delta_t / (max_iterations - iteration) 207 | threshold = new_lower_threshold + delta_t 208 | 209 | optimizer.set_params(alpha=alpha, threshold=best_threshold) 210 | optimizer.fit(X_train, y_train) 211 | 212 | return optimizer.coef_ 213 | -------------------------------------------------------------------------------- /src/deepymod/data/base.py: -------------------------------------------------------------------------------- 1 | """ Contains the base class for the Dataset (1 and 2 dimensional) and a function 2 | that takes a Pytorch tensor and converts it to a numpy array""" 3 | 4 | import torch 5 | import numpy as np 6 | from numpy import ndarray 7 | from numpy.random import default_rng 8 | from deepymod.data.samples import Subsampler 9 | 10 | from abc import ABC, ABCMeta, abstractmethod 11 | 12 | 13 | class Dataset(torch.utils.data.Dataset): 14 | def __init__( 15 | self, 16 | load_function, 17 | apply_normalize=None, 18 | apply_noise=None, 19 | apply_shuffle=None, 20 | shuffle=True, 21 | subsampler: Subsampler = None, 22 | load_kwargs: dict = {}, 23 | preprocess_kwargs: dict = { 24 | "random_state": 42, 25 | "noise_level": 0.0, 26 | "normalize_coords": False, 27 | "normalize_data": False, 28 | }, 29 | subsampler_kwargs: dict = {}, 30 | device: str = None, 31 | ): 32 | """A dataset class that loads the data, preprocesses it and lastly applies subsampling to it 33 | 34 | Args: 35 | load_function (func): Must return torch tensors in the format: (coordinates, data) 36 | shuffle (bool, optional): Shuffle the data. Defaults to True. 37 | apply_normalize (func): if not None, apply this function to the data for normalization. Defaults to None. 38 | subsampler (Subsampler, optional): Add some subsampling function. Defaults to None. 39 | load_kwargs (dict, optional): kwargs to pass to the load_function. Defaults to {}. 40 | preprocess_kwargs (dict, optional): (optional) arguments to pass to the preprocess method. Defaults to { "random_state": 42, "noise_level": 0.0, "normalize_coords": False, "normalize_data": False, }. 41 | subsampler_kwargs (dict, optional): (optional) arguments to pass to the subsampler method. Defaults to {}. 42 | device (str, optional): which device to send the data to. Defaults to None. 43 | """ 44 | self.load = load_function 45 | self.subsampler = subsampler 46 | self.load_kwargs = load_kwargs 47 | self.preprocess_kwargs = preprocess_kwargs 48 | self.subsampler_kwargs = subsampler_kwargs # so total number of samples is size(self.t_domain) * n_samples_per_frame 49 | # If some override function is provided, use that instead of the default. 50 | if apply_normalize != None: 51 | self.apply_normalize = apply_normalize 52 | if apply_noise != None: 53 | self.apply_normalize = apply_noise 54 | if apply_shuffle != None: 55 | self.apply_shuffle = apply_shuffle 56 | self.device = device 57 | self.shuffle = shuffle 58 | self.coords, self.data = self.load(**self.load_kwargs) 59 | # Ensure the data that loaded is not 0D/1D 60 | assert ( 61 | len(self.coords.shape) >= 2 62 | ), "Please explicitely specify a feature axis for the coordinates" 63 | assert ( 64 | len(self.data.shape) >= 2 65 | ), "Please explicitely specify a feature axis for the data" 66 | # Preprocess (add noise and normalization) 67 | self.coords, self.data = self.preprocess( 68 | self.coords, self.data, **self.preprocess_kwargs 69 | ) 70 | # Apply the subsampler if there is one 71 | if self.subsampler: 72 | self.coords, self.data = self.subsampler.sample( 73 | self.coords, self.data, **self.subsampler_kwargs 74 | ) 75 | # Reshaping the data to a (number_of_samples, number_of_features) shape if needed 76 | if len(self.data.shape) != 2 or len(self.coords.shape) != 2: 77 | self.coords, self.data = self._reshape(self.coords, self.data) 78 | if self.shuffle: 79 | self.coords, self.data = self.apply_shuffle(self.coords, self.data) 80 | # Now we know the data are shape (number_of_samples, number_of_features) we can set the number_of_samples 81 | self.number_of_samples = self.data.shape[0] 82 | 83 | print("Dataset is using device: ", self.device) 84 | if self.device: 85 | self.coords = self.coords.to(self.device) 86 | self.data = self.data.to(self.device) 87 | 88 | # Pytorch methods 89 | def __len__(self) -> int: 90 | """Returns length of dataset. Required by pytorch""" 91 | return self.number_of_samples 92 | 93 | def __getitem__(self, idx: int) -> int: 94 | """Returns coordinate and value. First axis of coordinate should be time.""" 95 | return self.coords[idx], self.data[idx] 96 | 97 | # get methods 98 | def get_coords(self): 99 | """Retrieve all the coordinate features""" 100 | return self.coords 101 | 102 | def get_data(self): 103 | """Retrieve all the data features""" 104 | return self.data 105 | 106 | # Logical methods 107 | def preprocess( 108 | self, 109 | X: torch.tensor, 110 | y: torch.tensor, 111 | random_state: int = 42, 112 | noise_level: float = 0.0, 113 | normalize_coords: bool = False, 114 | normalize_data: bool = False, 115 | ): 116 | """Add noise to the data and normalize the features 117 | Arguments: 118 | X (torch.tensor) : coordinates of the dataset 119 | y (torch.tensor) : values of the dataset 120 | random_state (int) : state for random number geerator 121 | noise (float) : standard deviations of noise to add 122 | normalize_coords (bool): apply normalization to the coordinates 123 | normalize_data (bool): apply normalization to the data 124 | """ 125 | print("Preprocessing data") 126 | # add noise 127 | y_processed = self.apply_noise(y, noise_level, random_state) 128 | # normalize coordinates 129 | if normalize_coords: 130 | X_processed = self.apply_normalize(X) 131 | else: 132 | X_processed = X 133 | # normalize data 134 | if normalize_data: 135 | y_processed = self.apply_normalize(y_processed) 136 | 137 | return X_processed, y_processed 138 | 139 | @staticmethod 140 | def apply_noise(y, noise_level, random_state): 141 | """Adds gaussian white noise of noise_level standard deviation. 142 | Args: 143 | y (torch.tensor): the data to which noise should be added 144 | noise_level (float): add white noise as a function of standard deviation 145 | random_state (int): the random state used for random number generation 146 | """ 147 | noise = noise_level * torch.std(y).data 148 | y_noisy = y + torch.tensor( 149 | default_rng(random_state).normal(loc=0.0, scale=noise, size=y.shape), 150 | dtype=torch.float32, 151 | ) # TO DO: switch to pytorch rng 152 | return y_noisy 153 | 154 | @staticmethod 155 | def apply_normalize(X): 156 | """minmax Normalize the data along the zeroth axis. Per feature 157 | Args: 158 | X (torch.tensor): data to be minmax normalized 159 | Returns: 160 | (torch.tensor): minmaxed data""" 161 | X_norm = (X - X.view(-1, X.shape[-1]).min(dim=0).values) / ( 162 | X.view(-1, X.shape[-1]).max(dim=0).values 163 | - X.view(-1, X.shape[-1]).min(dim=0).values 164 | ) * 2 - 1 165 | return X_norm 166 | 167 | @staticmethod 168 | def apply_shuffle(coords, data): 169 | """Shuffle the coordinates and data""" 170 | permutation = np.random.permutation(np.arange(len(data))) 171 | return coords[permutation], data[permutation] 172 | 173 | @staticmethod 174 | def _reshape(coords, data): 175 | """Reshape the coordinates and data to the format [number_of_samples, number_of_features]""" 176 | coords = coords.reshape([-1, coords.shape[-1]]) 177 | data = data.reshape([-1, data.shape[-1]]) 178 | return coords, data 179 | 180 | 181 | class Loader: 182 | def __init__(self, dataset): 183 | """Loader created to follow the workflow of PyTorch Dataset and Dataloader 184 | Leaves all data where it currently is.""" 185 | if isinstance(dataset, torch.utils.data.Subset): 186 | self.device = dataset.dataset.device 187 | else: 188 | self.device = dataset.device 189 | self.dataset = dataset 190 | self._count = 0 191 | self._length = 1 192 | 193 | def __getitem__(self, idx): 194 | if idx < self._length: 195 | return self.dataset[:] 196 | else: 197 | raise StopIteration 198 | 199 | def __len__(self): 200 | return self._length 201 | 202 | 203 | def get_train_test_loader( 204 | dataset, train_test_split=0.8, loader=Loader, loader_kwargs={} 205 | ): 206 | """Take a dataset, shuffle it, split it into a train and test and then 207 | return two loaders that are compatible with PyTorch. 208 | 209 | Args: 210 | dataset (torch.utils.data.Dataset): The dataset to use 211 | train_test_split (float, optional): The fraction of data used for train. Defaults to 0.8. 212 | loader (torch.utils.data.Dataloader, optional): The type of Dataloader to use. Defaults to GPULoader. 213 | loader_kwargs (dict, optional): Any kwargs to be passed to the loader]. Defaults to {}. 214 | 215 | Returns: 216 | Dataloader, Dataloader: The train and test dataloader 217 | """ 218 | length = dataset.number_of_samples 219 | indices = np.arange(0, length, dtype=int) 220 | np.random.shuffle(indices) 221 | split = int(train_test_split * length) 222 | train_indices = indices[:split] 223 | test_indices = indices[split:] 224 | train_data = torch.utils.data.Subset(dataset, train_indices) 225 | test_data = torch.utils.data.Subset(dataset, test_indices) 226 | return loader(train_data, **loader_kwargs), loader(test_data, **loader_kwargs) 227 | -------------------------------------------------------------------------------- /src/deepymod/model/library.py: -------------------------------------------------------------------------------- 1 | """ Contains the library classes that store the parameters u_t, theta""" 2 | import numpy as np 3 | import torch 4 | from torch.autograd import grad 5 | from itertools import combinations, product 6 | from functools import reduce 7 | from .deepmod import Library 8 | from typing import Tuple 9 | from ..utils.types import TensorList 10 | 11 | 12 | # ==================== Library helper functions ======================= 13 | def library_poly(prediction: torch.Tensor, max_order: int) -> torch.Tensor: 14 | """Given a prediction u, returns u^n up to max_order, including ones as first column. 15 | (technically these are monomials) 16 | Args: 17 | prediction (torch.Tensor): the data u for which to evaluate the library (n_samples x 1) 18 | max_order (int): the maximum polynomial order up to which compute the library 19 | 20 | Returns: 21 | torch.Tensor: Tensor with polynomials (n_samples, max_order + 1) 22 | """ 23 | u = torch.ones_like(prediction) 24 | for order in np.arange(1, max_order + 1): 25 | u = torch.cat((u, u[:, order - 1 : order] * prediction), dim=1) 26 | 27 | return u 28 | 29 | 30 | def library_deriv( 31 | data: torch.Tensor, prediction: torch.Tensor, max_order: int 32 | ) -> Tuple[torch.Tensor, torch.Tensor]: 33 | """Given a prediction u evaluated at data (t, x), returns du/dt and du/dx up to max_order, including ones 34 | as first column. 35 | 36 | Args: 37 | data (torch.Tensor): (t, x) locations of where to evaluate derivatives (n_samples x 2) 38 | prediction (torch.Tensor): the data u for which to evaluate the library (n_samples x 1) 39 | max_order (int): maximum order of derivatives to be calculated. 40 | 41 | Returns: 42 | Tuple[torch.Tensor, torch.Tensor]: time derivative and feature library ((n_samples, 1), (n_samples, max_order + 1)) 43 | """ 44 | dy = grad( 45 | prediction, data, grad_outputs=torch.ones_like(prediction), create_graph=True 46 | )[0] 47 | time_deriv = dy[:, 0:1] # First column is time derivative 48 | 49 | if max_order == 0: # If we only want the time derivative, du is just a scalar 50 | du = torch.ones_like(time_deriv) 51 | else: # Else we calculate the spatial derivatives 52 | du = torch.cat( 53 | (torch.ones_like(time_deriv), dy[:, 1:2]), dim=1 54 | ) # second column of dy gives first order derivative 55 | if ( 56 | max_order > 1 57 | ): # If we want higher order derivatives, we calculate them successively and concatenate them to du 58 | for order in np.arange(1, max_order): 59 | du = torch.cat( 60 | ( 61 | du, 62 | grad( 63 | du[:, order : order + 1], 64 | data, 65 | grad_outputs=torch.ones_like(prediction), 66 | create_graph=True, 67 | )[0][:, 1:2], 68 | ), 69 | dim=1, 70 | ) 71 | return time_deriv, du 72 | 73 | 74 | # ========================= Actual library functions ======================== 75 | class Library1D(Library): 76 | def __init__(self, poly_order: int, diff_order: int) -> None: 77 | """Calculates the temporal derivative a library/feature matrix consisting of 78 | 1) polynomials up to order poly_order, i.e. u, u^2... 79 | 2) derivatives up to order diff_order, i.e. u_x, u_xx 80 | 3) cross terms of 1) and 2), i.e. $uu_x$, $u^2u_xx$ 81 | 82 | Order of terms is derivative first, i.e. [$1, u_x, u, uu_x, u^2, ...$] 83 | 84 | Only works for 1D+1 data. Also works for multiple outputs but in that case doesn't calculate 85 | polynomial and derivative cross terms. <- trying to go back to DeePyMoD_torch, so ignore this statement 86 | 87 | Parameters 88 | ---------- 89 | poly_order : int 90 | The maximum polynomial order to include in the library. 91 | diff_order : int 92 | The maximum derivative order to include in the library. 93 | 94 | Args: 95 | poly_order (int): maximum order of the polynomial in the library 96 | diff_order (int): maximum order of the differentials in the library 97 | """ 98 | 99 | super().__init__() 100 | self.poly_order = poly_order 101 | self.diff_order = diff_order 102 | 103 | def library( 104 | self, input: Tuple[torch.Tensor, torch.Tensor] 105 | ) -> Tuple[TensorList, TensorList]: 106 | """Compute the temporal derivative and library for the given prediction at locations given by data. 107 | Data should have t in first column, x in second. 108 | 109 | Args: 110 | input (Tuple[torch.Tensor, torch.Tensor]): A prediction u (n_samples, n_outputs) and spatiotemporal locations (n_samples, 2). 111 | 112 | Returns: 113 | Tuple[TensorList, TensorList]: 114 | The time derivatives [(n_samples, 1) x n_outputs] 115 | thetas [(n_samples, (poly_order + 1)(deriv_order + 1))] 116 | computed from the library and data. 117 | """ 118 | prediction, data = input 119 | poly_list = [] 120 | deriv_list = [] 121 | time_deriv_list = [] 122 | 123 | # Creating lists for all outputs (each degree of freedom: l.h.s. of differential equation) 124 | for output in np.arange(prediction.shape[1]): 125 | time_deriv, du = library_deriv( 126 | data, prediction[:, output : output + 1], self.diff_order 127 | ) 128 | u = library_poly(prediction[:, output : output + 1], self.poly_order) 129 | 130 | poly_list.append(u) 131 | deriv_list.append(du) 132 | time_deriv_list.append(time_deriv) 133 | 134 | samples = time_deriv_list[0].shape[0] # number of samples 135 | total_terms = ( 136 | poly_list[0].shape[1] * deriv_list[0].shape[1] 137 | ) # product of the number of possible polynomials (i.e. monomials) and the number of derivative terms 138 | 139 | # Calculating theta 140 | if len(poly_list) == 1: 141 | # If we have a single output, we simply calculate and flatten matrix product 142 | # between polynomials and derivatives to get library 143 | theta = torch.matmul( 144 | poly_list[0][:, :, None], deriv_list[0][:, None, :] 145 | ).view(samples, total_terms) 146 | # For each sample poly_list[0][each_sample, :] and deriv_list[0][each_sample, :] the above line is equivalent to np.multiply.outer(poly_list[0][each_sample, :],deriv_list[0][each_sample, :] ).reshape(-1) 147 | # so the logic of the expression can be understood by executing np.add.outer(np.array(['', 'u', 'u^2'], object),np.array(['', 'u_x', 'u_xx','u_xxx'], object)).reshape(-1) <- this is consistent with equation (4) 148 | # this means that we iterate over deriv_list first (fast index) and then over poly_list (slow index) 149 | # this gives, for example: ['', 'u_x', 'u_xx', 'u_xxx', 'u', 'uu_x', 'uu_xx', 'uu_xxx', 'u^2', 'u^2u_x', 'u^2u_xx', 'u^2u_xxx'] 150 | else: 151 | theta_uv = reduce( 152 | (lambda x, y: (x[:, :, None] @ y[:, None, :]).view(samples, -1)), 153 | poly_list, 154 | ) # TODO comment the following lines 155 | theta_dudv = torch.cat( 156 | [ 157 | torch.matmul(du[:, :, None], dv[:, None, :]).view(samples, -1)[ 158 | :, 1: 159 | ] 160 | for du, dv in combinations(deriv_list, 2) 161 | ], 162 | 1, 163 | ) # calculate all unique combinations of derivatives 164 | theta_udu = torch.cat( 165 | [ 166 | torch.matmul(u[:, 1:, None], du[:, None, 1:]).view( 167 | samples, 168 | (poly_list[0].shape[1] - 1) * (deriv_list[0].shape[1] - 1), 169 | ) 170 | for u, dv in product(poly_list, deriv_list) 171 | ], 172 | 1, 173 | ) # calculate all unique products of polynomials and derivatives. This term was absent in DeePyMoD original repo but it is necessary for identification of Keller Segel 174 | theta = torch.cat([theta_uv, theta_dudv, theta_udu], dim=1) 175 | return time_deriv_list, [theta] * len(poly_list) 176 | 177 | 178 | class Library2D(Library): 179 | def __init__(self, poly_order: int) -> None: 180 | """Create a 2D library up to given polynomial order with second order derivatives 181 | i.e. for poly_order=1: [$1, u_x, u_y, u_{xx}, u_{yy}, u_{xy}$] 182 | Args: 183 | poly_order (int): maximum order of the polynomial in the library 184 | """ 185 | super().__init__() 186 | self.poly_order = poly_order 187 | 188 | def library( 189 | self, input: Tuple[torch.Tensor, torch.Tensor] 190 | ) -> Tuple[TensorList, TensorList]: 191 | """Compute the library for the given a prediction and data 192 | 193 | Args: 194 | input (Tuple[torch.Tensor, torch.Tensor]): A prediction and its data 195 | 196 | Returns: 197 | Tuple[TensorList, TensorList]: The time derivatives and the thetas 198 | computed from the library and data. 199 | """ 200 | 201 | prediction, data = input 202 | # Polynomial 203 | 204 | u = torch.ones_like(prediction) 205 | for order in np.arange(1, self.poly_order + 1): 206 | u = torch.cat((u, u[:, order - 1 : order] * prediction), dim=1) 207 | 208 | # Gradients 209 | du = grad( 210 | prediction, 211 | data, 212 | grad_outputs=torch.ones_like(prediction), 213 | create_graph=True, 214 | )[0] 215 | u_t = du[:, 0:1] 216 | u_x = du[:, 1:2] 217 | u_y = du[:, 2:3] 218 | du2 = grad( 219 | u_x, data, grad_outputs=torch.ones_like(prediction), create_graph=True 220 | )[0] 221 | u_xx = du2[:, 1:2] 222 | u_xy = du2[:, 2:3] 223 | u_yy = grad( 224 | u_y, data, grad_outputs=torch.ones_like(prediction), create_graph=True 225 | )[0][:, 2:3] 226 | 227 | du = torch.cat((torch.ones_like(u_x), u_x, u_y, u_xx, u_yy, u_xy), dim=1) 228 | 229 | samples = du.shape[0] 230 | # Bringing it together 231 | theta = torch.matmul(u[:, :, None], du[:, None, :]).view(samples, -1) 232 | 233 | return [u_t], [theta] 234 | -------------------------------------------------------------------------------- /src/deepymod/model/deepmod.py: -------------------------------------------------------------------------------- 1 | """ This file contains building blocks for the deepmod framework: 2 | I) The constraint class that constrains the neural network with the obtained solution, 3 | II) The sparsity estimator class, 4 | III) Function library class on which the model discovery is performed. 5 | IV) The DeepMoD class integrates these seperate building blocks. 6 | These are all abstract classes and implement the flow logic, rather than the specifics. 7 | """ 8 | 9 | import torch.nn as nn 10 | import torch 11 | from typing import Tuple 12 | from ..utils.types import TensorList 13 | from abc import ABCMeta, abstractmethod 14 | import numpy as np 15 | 16 | 17 | class Constraint(nn.Module, metaclass=ABCMeta): 18 | def __init__(self) -> None: 19 | """Abstract baseclass for the constraint module. 20 | for specific use cases see deepymod.model.constraint 21 | """ 22 | super().__init__() 23 | self.sparsity_masks: TensorList = None 24 | 25 | def forward(self, input: Tuple[TensorList, TensorList]) -> TensorList: 26 | """The forward pass of the constraint module applies the sparsity mask to the 27 | feature matrix theta, and then calculates the coefficients according to the 28 | method in the child. 29 | 30 | Args: 31 | input (Tuple[TensorList, TensorList]): (time_derivs, library) tuple of size 32 | ([(n_samples, 1) X n_outputs], [(n_samples, n_features) x n_outputs]). 33 | Returns: 34 | coeff_vectors (TensorList): List with coefficient vectors of size ([(n_features, 1) x n_outputs]) 35 | """ 36 | 37 | time_derivs, thetas = input 38 | 39 | if self.sparsity_masks is None: 40 | self.sparsity_masks = [ 41 | torch.ones(theta.shape[1], dtype=torch.bool).to(theta.device) 42 | for theta in thetas 43 | ] 44 | 45 | sparse_thetas = self.apply_mask(thetas, self.sparsity_masks) 46 | 47 | # Constraint grad. desc style doesn't allow to change shape, so we return full coeff 48 | # and multiply by mask to set zeros. For least squares-style, we need to put in 49 | # zeros in the right spot to get correct shape. 50 | coeff_vectors = self.fit(sparse_thetas, time_derivs) 51 | self.coeff_vectors = [ 52 | self.map_coeffs(mask, coeff) 53 | if mask.shape[0] != coeff.shape[0] 54 | else coeff * mask[:, None] 55 | for mask, coeff in zip(self.sparsity_masks, coeff_vectors) 56 | ] 57 | 58 | return self.coeff_vectors 59 | 60 | # static method is bound to a class rather than the objects for that class. This means that a static method can be called without an object for that class. This also means that static methods cannot modify the state of an object as they are not bound to it. 61 | @staticmethod 62 | def apply_mask(thetas: TensorList, masks: TensorList) -> TensorList: 63 | """Applies the sparsity mask to the feature (library) matrix theta. 64 | 65 | Args: 66 | thetas (TensorList): List of all library matrices of size [(n_samples, n_features) x n_outputs]. 67 | masks (TensorList): List of all sparsity masks 68 | Returns: 69 | TensorList: The sparse version of the library matrices of size [(n_samples, n_active_features) x n_outputs]. 70 | """ 71 | sparse_thetas = [theta[:, mask] for theta, mask in zip(thetas, masks)] 72 | return sparse_thetas 73 | 74 | @staticmethod 75 | def map_coeffs(mask: torch.Tensor, coeff_vector: torch.Tensor) -> torch.Tensor: 76 | """Places the coeff_vector components in the true positions of the mask. 77 | i.e. maps ((0, 1, 1, 0), (0.5, 1.5)) -> (0, 0.5, 1.5, 0). 78 | 79 | Args: 80 | mask (torch.Tensor): Boolean mask describing active components. 81 | coeff_vector (torch.Tensor): Vector with active-components. 82 | 83 | Returns: 84 | mapped_coeffs (torch.Tensor): mapped coefficients. 85 | """ 86 | mapped_coeffs = ( 87 | torch.zeros((mask.shape[0], 1)) 88 | .to(coeff_vector.device) 89 | .masked_scatter_(mask[:, None], coeff_vector) 90 | ) 91 | return mapped_coeffs 92 | 93 | @abstractmethod 94 | def fit(self, sparse_thetas: TensorList, time_derivs: TensorList) -> TensorList: 95 | """Abstract method. Specific method should return the coefficients as calculated from the sparse feature 96 | matrices and temporal derivatives from dt ~ theta @ coeff_vector, equation (7) of the original paper 97 | 98 | Args: 99 | sparse_thetas (TensorList): List containing the sparse feature tensors of size (n_samples, n_active_features). 100 | time_derivs (TensorList): List containing the time derivatives of size (n_samples, n_outputs). 101 | 102 | Returns: 103 | (TensorList): Calculated coefficients of size (n_active_features, n_outputs). 104 | """ 105 | raise NotImplementedError 106 | 107 | 108 | class Estimator(nn.Module, metaclass=ABCMeta): 109 | def __init__(self) -> None: 110 | """Abstract baseclass for the sparse estimator module. For specific implementation see deepymod.model.sparse_estimators""" 111 | super().__init__() 112 | self.coeff_vectors = None 113 | 114 | def forward(self, thetas: TensorList, time_derivs: TensorList) -> TensorList: 115 | """The forward pass of the sparse estimator module first normalizes the library matrices 116 | and time derivatives by dividing each column (i.e. feature) by their l2 norm, than calculates the coefficient vectors 117 | according to the sparse estimation algorithm supplied by the child and finally returns the sparsity 118 | mask (i.e. which terms are active) based on these coefficients. 119 | 120 | Args: 121 | thetas (TensorList): List containing the sparse feature tensors of size [(n_samples, n_active_features) x n_outputs]. 122 | time_derivs (TensorList): List containing the time derivatives of size [(n_samples, 1) x n_outputs]. 123 | 124 | Returns: 125 | (TensorList): List containting the sparsity masks of a boolean type and size [(n_samples, n_features) x n_outputs]. 126 | """ 127 | 128 | # we first normalize theta and the time deriv 129 | with torch.no_grad(): 130 | normed_time_derivs = [ 131 | (time_deriv / torch.norm(time_deriv)).detach().cpu().numpy() 132 | for time_deriv in time_derivs 133 | ] 134 | normed_thetas = [ 135 | (theta / torch.norm(theta, dim=0, keepdim=True)).detach().cpu().numpy() 136 | for theta in thetas 137 | ] 138 | 139 | self.coeff_vectors = [ 140 | self.fit(theta, time_deriv.squeeze())[:, None] 141 | for theta, time_deriv in zip(normed_thetas, normed_time_derivs) 142 | ] 143 | sparsity_masks = [ 144 | torch.tensor(coeff_vector != 0.0, dtype=torch.bool) 145 | .squeeze() 146 | .to(thetas[0].device) # move to gpu if required 147 | for coeff_vector in self.coeff_vectors 148 | ] 149 | 150 | return sparsity_masks 151 | 152 | @abstractmethod 153 | def fit(self, X: np.ndarray, y: np.ndarray) -> np.ndarray: 154 | """Abstract method. Specific method should compute the coefficient based on feature matrix X and observations y. 155 | Note that we expect X and y to be numpy arrays, i.e. this module is non-differentiable. 156 | 157 | Args: 158 | x (np.ndarray): Feature matrix of size (n_samples, n_features) 159 | y (np.ndarray): observations of size (n_samples, n_outputs) 160 | 161 | Returns: 162 | (np.ndarray): Coefficients of size (n_samples, n_outputs) 163 | """ 164 | pass 165 | 166 | 167 | class Library(nn.Module): 168 | def __init__(self) -> None: 169 | """Abstract baseclass for the library module. For specific uses see 170 | deepymod.model.library 171 | """ 172 | super().__init__() 173 | self.norms = None 174 | 175 | def forward( 176 | self, input: Tuple[torch.Tensor, torch.Tensor] 177 | ) -> Tuple[TensorList, TensorList]: 178 | """Compute the library (time derivatives and thetas) from a given dataset. Also calculates the norms 179 | of these, later used to calculate the normalized coefficients. 180 | 181 | Args: 182 | input (Tuple[TensorList, TensorList]): (prediction, data) tuple of size ((n_samples, n_outputs), (n_samples, n_dims)) 183 | 184 | Returns: 185 | Tuple[TensorList, TensorList]: Temporal derivative and libraries of size ([(n_samples, 1) x n_outputs]), [(n_samples, n_features)x n_outputs]) 186 | """ 187 | time_derivs, thetas = self.library(input) 188 | self.norms = [ 189 | (torch.norm(time_deriv) / torch.norm(theta, dim=0, keepdim=True)) 190 | .detach() 191 | .squeeze() 192 | for time_deriv, theta in zip(time_derivs, thetas) 193 | ] 194 | return time_derivs, thetas 195 | 196 | @abstractmethod 197 | def library( 198 | self, input: Tuple[torch.Tensor, torch.Tensor] 199 | ) -> Tuple[TensorList, TensorList]: 200 | """Abstract method. Specific method should calculate the temporal derivative and feature matrices. 201 | These should be a list; one temporal derivative and feature matrix per output. 202 | 203 | Args: 204 | input (Tuple[TensorList, TensorList]): (prediction, data) tuple of size ((n_samples, n_outputs), (n_samples, n_dims)) 205 | 206 | Returns: 207 | Tuple[TensorList, TensorList]: Temporal derivative and libraries of size ([(n_samples, 1) x n_outputs]), [(n_samples, n_features)x n_outputs]) 208 | """ 209 | pass 210 | 211 | 212 | class DeepMoD(nn.Module): 213 | def __init__( 214 | self, 215 | function_approximator: torch.nn.Sequential, 216 | library: Library, 217 | sparsity_estimator: Estimator, 218 | constraint: Constraint, 219 | ) -> None: 220 | """The DeepMoD class integrates the various buiding blocks into one module. The function approximator approximates the data, 221 | the library than builds a feature matrix from its output and the constraint constrains these. The sparsity estimator is called 222 | during training to update the sparsity mask (i.e. which terms the constraint is allowed to use.) 223 | 224 | Args: 225 | function_approximator (torch.nn.Sequential): 226 | makes predictions about the dynamical field (variable), its value and derivatives 227 | parametrized by Neural Network (NN) so taking use of autodiff 228 | library (Library): 229 | Library of terms to be used in the model discovery process 230 | sparsity_estimator (Estimator): updates the sparsity mask 231 | Example: Threshold(0.1) would set threshold = 0.1 for the thresholding estimator of the coefficients 232 | constraint (Constraint): [description] 233 | """ 234 | super().__init__() 235 | self.func_approx = function_approximator 236 | self.library = library 237 | self.sparse_estimator = sparsity_estimator 238 | self.constraint = constraint 239 | 240 | def forward( 241 | self, input: torch.Tensor 242 | ) -> Tuple[torch.Tensor, TensorList, TensorList]: 243 | """The forward pass approximates the data, builds the time derivative and feature matrices 244 | and applies the constraint. 245 | 246 | It returns the prediction of the network, the time derivatives and the feature matrices. 247 | 248 | Args: 249 | input (torch.Tensor): Tensor of shape (n_samples, n_outputs) containing the coordinates, first column should be the time coordinate. 250 | 251 | Returns: 252 | Tuple[torch.Tensor, TensorList, TensorList]: The prediction, time derivatives and and feature matrices of respective sizes 253 | ((n_samples, n_outputs), [(n_samples, 1) x n_outputs]), [(n_samples, n_features) x n_outputs]) 254 | 255 | """ 256 | prediction, coordinates = self.func_approx( 257 | input 258 | ) # predict the dynamical field (variable), its value and derivatives 259 | time_derivs, thetas = self.library( 260 | (prediction, coordinates) 261 | ) # library function returns time_deriv and theta (equation (4) of the manuscript) 262 | coeff_vectors = self.constraint( 263 | (time_derivs, thetas) 264 | ) # used to be called `fit` in DeepMoD_torch 265 | return prediction, time_derivs, thetas 266 | 267 | @property 268 | def sparsity_masks(self): 269 | """Returns the sparsity masks which contain the active terms (array of bools). 270 | Calls on constraint object which is attribute of DeepMoD class.""" 271 | return self.constraint.sparsity_masks 272 | 273 | def estimator_coeffs(self) -> TensorList: 274 | """Calculate the coefficients as estimated by the sparse estimator. 275 | Calls on sparse_estimator object which is attribute of DeepMoD class. 276 | Returns: 277 | (TensorList): List of coefficients of size [(n_features, 1) x n_outputs] 278 | """ 279 | coeff_vectors = self.sparse_estimator.coeff_vectors 280 | return coeff_vectors 281 | 282 | def constraint_coeffs(self, scaled=False, sparse=False) -> TensorList: 283 | """Calculate the coefficients as estimated by the constraint. 284 | Calls on constraint object which is attribute of DeepMoD class to get coeff_vectors 285 | which are processed depending on: 286 | Args: 287 | scaled (bool): Determine whether or not the coefficients should be normalized 288 | sparse (bool): Whether to apply the sparsity mask to the coefficients. 289 | 290 | Returns: 291 | (TensorList): List of coefficients of size [(n_features, 1) x n_outputs] 292 | """ 293 | coeff_vectors = self.constraint.coeff_vectors 294 | if scaled: # perform normalization 295 | coeff_vectors = [ 296 | coeff / norm[:, None] 297 | for coeff, norm, mask in zip( 298 | coeff_vectors, self.library.norms, self.sparsity_masks 299 | ) 300 | ] 301 | if sparse: # apply sparsity mask 302 | coeff_vectors = [ 303 | sparsity_mask[:, None] * coeff 304 | for sparsity_mask, coeff in zip(self.sparsity_masks, coeff_vectors) 305 | ] 306 | return coeff_vectors 307 | --------------------------------------------------------------------------------