├── .gitignore ├── .gitlab-ci.yml ├── .readthedocs.yaml ├── LICENSE ├── README.md ├── example_scripts ├── bbm_example.ipynb ├── cahn_hilliard_example.ipynb ├── kdv_burgers_example.ipynb ├── kdv_example.ipynb ├── kdv_utils.py ├── model_evaluation.py ├── mpc_example.py ├── phnn_ode_examples.ipynb ├── phnn_pde_examples.ipynb ├── pm_example.ipynb ├── pure_kdv.gif ├── spring_example.ipynb └── train_model.py └── phlearn ├── LICENSE.txt ├── README.md ├── docs ├── Makefile ├── make.bat ├── requirements.txt └── source │ ├── conf.py │ ├── control.rst │ ├── index.rst │ ├── phnns.rst │ ├── phsystems_ode.rst │ ├── phsystems_pde.rst │ └── utils.rst ├── phlearn ├── __init__.py ├── control │ ├── __init__.py │ ├── casadiNN.py │ ├── casadiPH.py │ ├── mpc.py │ ├── phcontroller.py │ ├── pid.py │ └── reference.py ├── phnns │ ├── __init__.py │ ├── dynamic_system_neural_network.py │ ├── ode_models.py │ ├── pde_models.py │ ├── pseudo_hamiltonian_neural_network.py │ ├── pseudo_hamiltonian_pde_neural_network.py │ ├── tests │ │ ├── __init__.py │ │ └── test_phnns.py │ └── train_utils.py ├── phsystems │ ├── __init__.py │ ├── ode │ │ ├── __init__.py │ │ ├── msd_system.py │ │ ├── pseudo_hamiltonian_system.py │ │ ├── tank_system.py │ │ └── tests │ │ │ ├── __init__.py │ │ │ └── test_phsystems.py │ ├── pde │ │ ├── __init__.py │ │ ├── allen_cahn_system.py │ │ ├── bbm_system.py │ │ ├── cahn_hilliard_system.py │ │ ├── heat_system.py │ │ ├── kdv_system.py │ │ ├── perona_malik_system.py │ │ └── pseudo_hamiltonian_pde_system.py │ └── tests │ │ ├── __init__.py │ │ └── test_phsystems.py └── utils │ ├── __init__.py │ ├── derivatives.py │ └── utils.py ├── requirements-control.txt ├── requirements-dev.txt ├── requirements.txt └── setup.py /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | .eggs/ 17 | lib/ 18 | lib64/ 19 | parts/ 20 | sdist/ 21 | var/ 22 | wheels/ 23 | share/python-wheels/ 24 | *.egg-info/ 25 | .installed.cfg 26 | *.egg 27 | MANIFEST 28 | 29 | # PyInstaller 30 | # Usually these files are written by a python script from a template 31 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 32 | *.manifest 33 | *.spec 34 | 35 | # Installer logs 36 | pip-log.txt 37 | pip-delete-this-directory.txt 38 | 39 | # Unit test / coverage reports 40 | htmlcov/ 41 | .tox/ 42 | .nox/ 43 | .coverage 44 | .coverage.* 45 | .cache 46 | nosetests.xml 47 | coverage.xml 48 | *.cover 49 | *.py,cover 50 | .hypothesis/ 51 | .pytest_cache/ 52 | cover/ 53 | 54 | # Translations 55 | *.mo 56 | *.pot 57 | 58 | # Django stuff: 59 | *.log 60 | local_settings.py 61 | db.sqlite3 62 | db.sqlite3-journal 63 | 64 | # Flask stuff: 65 | instance/ 66 | .webassets-cache 67 | 68 | # Scrapy stuff: 69 | .scrapy 70 | 71 | # Sphinx documentation 72 | docs/_build/ 73 | 74 | # PyBuilder 75 | .pybuilder/ 76 | target/ 77 | 78 | # Jupyter Notebook 79 | .ipynb_checkpoints 80 | 81 | # IPython 82 | profile_default/ 83 | ipython_config.py 84 | 85 | # pyenv 86 | # For a library or package, you might want to ignore these files since the code is 87 | # intended to run in multiple environments; otherwise, check them in: 88 | # .python-version 89 | 90 | # pipenv 91 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 92 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 93 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 94 | # install all needed dependencies. 95 | #Pipfile.lock 96 | 97 | # poetry 98 | # Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. 99 | # This is especially recommended for binary packages to ensure reproducibility, and is more 100 | # commonly ignored for libraries. 101 | # https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control 102 | #poetry.lock 103 | 104 | # pdm 105 | # Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. 106 | #pdm.lock 107 | # pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it 108 | # in version control. 109 | # https://pdm.fming.dev/#use-with-ide 110 | .pdm.toml 111 | 112 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm 113 | __pypackages__/ 114 | 115 | # Celery stuff 116 | celerybeat-schedule 117 | celerybeat.pid 118 | 119 | # SageMath parsed files 120 | *.sage.py 121 | 122 | # Environments 123 | .env 124 | .venv 125 | env/ 126 | venv/ 127 | ENV/ 128 | env.bak/ 129 | venv.bak/ 130 | 131 | # Spyder project settings 132 | .spyderproject 133 | .spyproject 134 | 135 | # Rope project settings 136 | .ropeproject 137 | 138 | # mkdocs documentation 139 | /site 140 | 141 | # mypy 142 | .mypy_cache/ 143 | .dmypy.json 144 | dmypy.json 145 | 146 | # Pyre type checker 147 | .pyre/ 148 | 149 | # pytype static type analyzer 150 | .pytype/ 151 | 152 | # Cython debug symbols 153 | cython_debug/ 154 | 155 | # PyCharm 156 | # JetBrains specific template is maintained in a separate JetBrains.gitignore that can 157 | # be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore 158 | # and can be added to the global gitignore or merged into this file. For a more nuclear 159 | # option (not recommended) you can uncomment the following to ignore the entire idea folder. 160 | #.idea/ 161 | 162 | 163 | # General 164 | .DS_Store 165 | .AppleDouble 166 | .LSOverride 167 | 168 | # Icon must end with two \r 169 | Icon 170 | 171 | # Thumbnails 172 | ._* 173 | 174 | # Files that might appear in the root of a volume 175 | .DocumentRevisions-V100 176 | .fseventsd 177 | .Spotlight-V100 178 | .TemporaryItems 179 | .Trashes 180 | .VolumeIcon.icns 181 | .com.apple.timemachine.donotpresent 182 | 183 | # Directories potentially created on remote AFP share 184 | .AppleDB 185 | .AppleDesktop 186 | Network Trash Folder 187 | Temporary Items 188 | .apdisk 189 | 190 | # Windows thumbnail cache files 191 | Thumbs.db 192 | Thumbs.db:encryptable 193 | ehthumbs.db 194 | ehthumbs_vista.db 195 | 196 | # Dump file 197 | *.stackdump 198 | 199 | # Folder config file 200 | [Dd]esktop.ini 201 | 202 | # Recycle Bin used on file shares 203 | $RECYCLE.BIN/ 204 | 205 | # Windows Installer files 206 | *.cab 207 | *.msi 208 | *.msix 209 | *.msm 210 | *.msp 211 | 212 | # Windows shortcuts 213 | *.lnk 214 | 215 | # Trained models 216 | *.model 217 | models -------------------------------------------------------------------------------- /.gitlab-ci.yml: -------------------------------------------------------------------------------- 1 | stages: 2 | - test 3 | 4 | run_tests: 5 | stage: test 6 | image: python:3.9 7 | before_script: 8 | - pip install -r phlearn/requirements.txt 9 | - pip install -r phlearn/requirements-dev.txt 10 | script: 11 | - pytest -v -------------------------------------------------------------------------------- /.readthedocs.yaml: -------------------------------------------------------------------------------- 1 | version: 2 2 | 3 | build: 4 | os: "ubuntu-20.04" 5 | tools: 6 | python: "3.8" 7 | 8 | python: 9 | install: 10 | - requirements: phlearn/docs/requirements.txt 11 | - method: pip 12 | path: phlearn 13 | 14 | sphinx: 15 | fail_on_warning: true 16 | configuration: phlearn/docs/source/conf.py -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2025 SINTEF 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 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Pseudo-Hamiltonian neural networks 2 | 3 | This repository contains the package phlearn for modelling pseudo-Hamiltonian ODE and PDE systems with neural networks, and example scripts for training and simulation. 4 | 5 | See [(Eidnes et al., 2023)](https://doi.org/10.1016/j.physd.2023.133673) for a description of PHNN for ODEs. 6 | See [(Eidnes and Lye, 2024)](https://doi.org/10.1016/j.jcp.2023.112738) for a description of PHNN for PDEs. 7 | 8 | [Documentation available here](https://pseudo-Hamiltonian-neural-networks.readthedocs.io/en/latest/) 9 | 10 | ## Installation 11 | 12 | The phlearn package is available via PyPi: 13 | 14 | ``` 15 | pip install phlearn 16 | ``` 17 | 18 | 19 | Alternatively, to get the latest updates not yet available on PyPi, you can clone this repository and install directly: 20 | 21 | ``` 22 | git clone https://github.com/SINTEF/pseudo-Hamiltonian-neural-networks.git 23 | cd pseudo-Hamiltonian-neural-networks 24 | pip install -e phlearn 25 | ``` 26 | 27 | 28 | #### Authors 29 | Camilla Sterud: c.sterud@icloud.com 30 | Sølve Eidnes: solve.eidnes@sintef.no 31 | Eivind Bøhn: eivind.bohn@sintef.no 32 | Signe Riemer-Sørensen: signe.riemer-sorensen@sintef.no 33 | Alexander Stasik: alexander.stasik@sintef.no 34 | Kjetil Olsen Lye: kjetil.olsen.lye@sintef.no 35 | Ben Tapley: ben.tapley@sintef.no 36 | August Femtehjell: august.femtehjell@sintef.no 37 | -------------------------------------------------------------------------------- /example_scripts/kdv_utils.py: -------------------------------------------------------------------------------- 1 | from typing import Callable 2 | 3 | import matplotlib.pyplot as plt 4 | import numpy as np 5 | from scipy.sparse import diags_array 6 | 7 | import phlearn.phnns as phnn 8 | from phlearn.phsystems.pde import PseudoHamiltonianPDESystem 9 | 10 | 11 | def setup_KdV_system( 12 | x: np.ndarray, 13 | eta: float = 6.0, 14 | gamma: float = 1.0, 15 | nu: float = 0.0, 16 | n_solitons: int = 2, 17 | ) -> PseudoHamiltonianPDESystem: 18 | M = x.size 19 | dx = (x[-1] - x[0]) / (M - 1) 20 | 21 | # Forward difference matrix 22 | Dp = diags_array([1, -1, 1], offsets=[-M + 1, 0, 1], shape=(M, M)) 23 | Dp = 1 / dx * Dp 24 | # Central difference matrix 25 | D1 = diags_array([1, -1, 1, -1], offsets=[-M + 1, -1, 1, M - 1], shape=(M, M)) 26 | D1 = 0.5 / dx * D1 27 | # 2nd order central difference matrix 28 | D2 = diags_array([1, 1, -2, 1, 1], offsets=[-M + 1, -1, 0, 1, M - 1], shape=(M, M)) 29 | D2 = 1 / dx**2 * D2 30 | 31 | def hamiltonian(u: np.ndarray) -> np.ndarray: 32 | integrand = -1 / 6 * eta * u**3 + 0.5 * gamma**2 * ((Dp @ u.T).T ** 2) 33 | return np.sum(integrand, axis=-1) 34 | 35 | def hamiltonian_grad(u: np.ndarray) -> np.ndarray: 36 | return -0.5 * eta * u**2 - (gamma**2 * u @ D2) 37 | 38 | def initial_condition() -> Callable[[np.random.Generator], np.ndarray]: 39 | P = (x[-1] - x[0]) * M / (M - 1) 40 | sech = lambda a: 1 / np.cosh(a) 41 | 42 | def sampler_train(rng: np.random.Generator) -> np.ndarray: 43 | u0 = -np.cos(np.pi * x) 44 | u0 = np.concatenate([u0[M:], u0[:M]], axis=-1) 45 | return u0 46 | 47 | def sampler_solitons(rng: np.random.Generator) -> np.ndarray: 48 | ks = rng.uniform(0.5, 2.0, n_solitons) 49 | ds = rng.uniform(0.0, 1.0, n_solitons) 50 | u0 = np.zeros_like(x) 51 | 52 | for k, d in zip(ks, ds): 53 | center = (x + P / 2 - P * d) % P - P / 2 54 | coeff = 6.0 / eta * 2 * k**2 55 | u0 += coeff * sech(np.abs(k * center)) ** 2 56 | 57 | u0 = np.concatenate([u0[M:], u0[:M]], axis=-1) 58 | return u0 59 | 60 | if n_solitons == 0: 61 | return sampler_train 62 | 63 | return sampler_solitons 64 | 65 | KdV_system = PseudoHamiltonianPDESystem( 66 | nstates=M, 67 | skewsymmetric_matrix=D1, 68 | hamiltonian=hamiltonian, 69 | grad_hamiltonian=hamiltonian_grad, 70 | init_sampler=initial_condition(), 71 | ) 72 | 73 | return KdV_system 74 | 75 | 76 | def generate_KdV_data( 77 | x: np.ndarray, 78 | t_axis: np.ndarray, 79 | eta: float = 6.0, 80 | gamma: float = 1.0, 81 | n_solitons: int = 2, 82 | n_trajectories: int = 5, 83 | ): 84 | system = setup_KdV_system(x, eta, gamma, n_solitons=n_solitons) 85 | traindata = phnn.generate_dataset(system, n_trajectories, t_axis, xspatial=x) 86 | 87 | return system, traindata, t_axis 88 | 89 | 90 | def setup_baseline_nn(spatial_points: int, **kwargs) -> phnn.DynamicSystemNN: 91 | baseline_nn = phnn.PDEBaselineNN(spatial_points, **kwargs) 92 | basemodel = phnn.DynamicSystemNN(baseline_nn, baseline_nn) 93 | return basemodel 94 | 95 | 96 | def plot_solutions( 97 | u_exact: np.ndarray, 98 | u_model: np.ndarray, 99 | x: np.ndarray, 100 | t_axis: np.ndarray, 101 | ) -> None: 102 | N = u_exact.shape[0] 103 | fig = plt.figure(figsize=(7, 4)) 104 | lw = 2 105 | colors = [ 106 | # (0, 0.4, 1), 107 | (1, 0.7, 0.3), 108 | (0.2, 0.7, 0.2), 109 | (0.8, 0, 0.2), 110 | (0.5, 0.3, 0.9), 111 | ] 112 | 113 | timesteps = [int(i * N / 4) for i in range(1, 4)] + [-1] 114 | 115 | plt.plot(x, u_exact[0], "k--", linewidth=lw, label=f"$t = {t_axis[0]:.2f}$") 116 | for i, idx in enumerate(timesteps): 117 | plt.plot( 118 | x, 119 | u_exact[idx], 120 | color=colors[i], 121 | linestyle="--", 122 | linewidth=lw, 123 | label=f"$t = {t_axis[idx]:.2f}$, true", 124 | ) 125 | plt.plot( 126 | x, 127 | u_model[idx], 128 | color=colors[i], 129 | linestyle="-", 130 | linewidth=lw, 131 | label=f"$t = {t_axis[idx]:.2f}$, model", 132 | ) 133 | 134 | plt.xlabel("$x$", fontsize=12) 135 | plt.ylabel("$u$", fontsize=12) 136 | plt.title("PHNN vs ground truth", fontsize=14) 137 | plt.legend() 138 | plt.show() 139 | -------------------------------------------------------------------------------- /example_scripts/model_evaluation.py: -------------------------------------------------------------------------------- 1 | 2 | import argparse 3 | import glob 4 | import os 5 | 6 | import numpy as np 7 | import pandas as pd 8 | import torch 9 | 10 | from phlearn.phnns import load_dynamic_system_model, PseudoHamiltonianNN, BaselineSplitNN 11 | from phlearn.phsystems.ode import init_tanksystem, init_msdsystem 12 | 13 | ttype = torch.float32 14 | torch.set_default_dtype(ttype) 15 | 16 | if __name__ == "__main__": 17 | parser = argparse.ArgumentParser() 18 | parser.add_argument('--modelfolder', type=str, 19 | help='Path to folder containing model dicts.') 20 | parser.add_argument('--system', type=str, choices=['tank', 'msd'], required=True, 21 | help='Choose to train a tank or a forced mass spring damper.') 22 | parser.add_argument('--nrollouts', '-n', type=int, default=10, 23 | help='Number of trajectories to roll out per model.') 24 | parser.add_argument('--dt_rollout', '-d', type=float, 25 | help='Sample time of rollouts.') 26 | parser.add_argument('--t_max', '-t', type=float, 27 | help='Duration of each rollout. If not provided, duration from training of each model is used.') 28 | 29 | args = parser.parse_args() 30 | 31 | modelfolder = args.modelfolder 32 | system = args.system 33 | nrollouts = args.nrollouts 34 | seed = 100 35 | 36 | if system == 'tank': 37 | pH_system = init_tanksystem() 38 | else: 39 | pH_system = init_msdsystem() 40 | 41 | nstates = pH_system.nstates 42 | 43 | index = [] 44 | for f in glob.glob(modelfolder+"/*.model"): 45 | modelname = f.split('/')[-1] 46 | index += [modelname] 47 | 48 | df = pd.DataFrame(index=index) 49 | 50 | for modelname in df.index: 51 | modelpath = os.path.join(modelfolder, modelname) 52 | model, optimizer, metadict = load_dynamic_system_model(modelpath) 53 | 54 | pH_system.seed(seed) 55 | dt = args.dt_rollout 56 | if dt is None: 57 | dt = metadict['traininginfo']['sampling_time'] 58 | t_max = args.t_max 59 | if t_max is None: 60 | t_max = metadict['traininginfo']['t_max'] 61 | t_sample = np.arange(0, t_max, dt) 62 | nsamples = t_sample.shape[0] 63 | 64 | x_exact, dxdt, _, _ = pH_system.sample_trajectory(t_sample) 65 | x0 = x_exact[0, :] 66 | x_phnn, _ = model.simulate_trajectory(False, t_sample, x0=x0) 67 | 68 | if isinstance(model, PseudoHamiltonianNN): 69 | if (not model.external_forces_provided): 70 | F_phnn = model.external_forces(torch.tensor(x_phnn, dtype=ttype), 71 | torch.tensor(t_sample.reshape(nsamples, 1), dtype=ttype)).detach().numpy() 72 | F_phnn -= F_phnn.mean(axis=0) 73 | F_exact = pH_system.external_forces(x_exact, t_sample) 74 | df.loc[modelname, 'External force MSE'] = ((F_phnn - F_exact)**2).mean() 75 | df.loc[modelname, 'External force MAE'] = np.abs(F_phnn - F_exact).mean() 76 | if (not model.dissipation_provided): 77 | df.loc[modelname, 'R MSE'] = ((model.R(x_exact).detach().numpy() - pH_system.R(x_exact))**2).mean() 78 | df.loc[modelname, 'R MAE'] = np.abs(model.R(x_exact).detach().numpy() - pH_system.R(x_exact)).mean() 79 | if (not model.hamiltonian_provided): 80 | if pH_system.H is not None: 81 | hamiltonian_exact = pH_system.H(x_exact) 82 | hamiltonian_phnn = model.hamiltonian(torch.tensor(x_phnn, dtype=ttype)).detach().numpy() 83 | hamiltonian_exact = pH_system.H(x_exact) 84 | hamiltonian_phnn = model.hamiltonian(torch.tensor(x_phnn, dtype=ttype)).detach().numpy() 85 | df.loc[modelname, 'H MSE'] = ((hamiltonian_exact - hamiltonian_phnn)**2).mean() 86 | df.loc[modelname, 'H MAE'] = np.abs(hamiltonian_exact - hamiltonian_phnn).mean() 87 | dH_exact = pH_system.dH(x_exact) 88 | dH_phnn = model.dH(torch.tensor(x_phnn, dtype=ttype)).detach().numpy() 89 | if system == 'tank': 90 | df.loc[modelname, 'dH tanks MSE'] = ((pH_system.tanklevels(dH_exact) - pH_system.tanklevels(dH_phnn))**2).mean() 91 | df.loc[modelname, 'dH tanks MAE'] = np.abs(pH_system.tanklevels(dH_exact) - pH_system.tanklevels(dH_phnn)).mean() 92 | df.loc[modelname, 'dH pipes MSE'] = ((pH_system.pipeflows(dH_exact) - pH_system.pipeflows(dH_phnn))**2).mean() 93 | df.loc[modelname, 'dH pipes MAE'] = np.abs(pH_system.pipeflows(dH_exact) - pH_system.pipeflows(dH_phnn)).mean() 94 | else: 95 | df.loc[modelname, 'dH x1 MSE'] = ((dH_exact[:, 0] - dH_phnn[:, 0])**2).mean() 96 | df.loc[modelname, 'dH x1 MAE'] = np.abs(dH_exact[:, 0] - dH_phnn[:, 0]).mean() 97 | df.loc[modelname, 'dH x2 MSE'] = ((dH_exact[:, 1] - dH_phnn[:, 1])**2).mean() 98 | df.loc[modelname, 'dH x2 MAE'] = np.abs(dH_exact[:, 1] - dH_phnn[:, 1]).mean() 99 | elif isinstance(model.rhs_model, BaselineSplitNN): 100 | F_baseline = model.rhs_model.network_t(torch.tensor(x_phnn, dtype=ttype), 101 | torch.tensor(t_sample.reshape(nsamples, 1), dtype=ttype)).detach().numpy() 102 | F_baseline -= F_baseline.mean(axis=0) 103 | F_exact = pH_system.external_forces(x_exact, t_sample) 104 | df.loc[modelname, 'External force MSE'] = ((F_baseline - F_exact)**2).mean() 105 | df.loc[modelname, 'External force MAE'] = np.abs(F_baseline - F_exact).mean() 106 | 107 | if system == 'tank': 108 | df.loc[modelname, 'Tanks MSE'] = ((pH_system.tanklevels(x_exact) - pH_system.tanklevels(x_phnn))**2).mean() 109 | df.loc[modelname, 'Pipes MSE'] = ((pH_system.pipeflows(x_exact) - pH_system.pipeflows(x_phnn))**2).mean() 110 | df.loc[modelname, 'Tanks MAE'] = np.abs(pH_system.tanklevels(x_exact) - pH_system.tanklevels(x_phnn)).mean() 111 | df.loc[modelname, 'Pipes MAE'] = np.abs(pH_system.pipeflows(x_exact) - pH_system.pipeflows(x_phnn)).mean() 112 | else: 113 | df.loc[modelname, 'x1 MSE'] = ((x_exact[:, 0] - x_phnn[:, 0])**2).mean() 114 | df.loc[modelname, 'x1 MAE'] = np.abs(x_exact[:, 0] - x_phnn[:, 0]).mean() 115 | df.loc[modelname, 'x2 MSE'] = ((x_exact[:, 1] - x_phnn[:, 1])**2).mean() 116 | df.loc[modelname, 'x2 MAE'] = np.abs(x_exact[:, 1] - x_phnn[:, 1]).mean() 117 | 118 | for i in range(1,nrollouts): 119 | x_exact, dxdt, _, _ = pH_system.sample_trajectory(t_sample) 120 | x0 = x_exact[0, :] 121 | x_phnn, _ = model.simulate_trajectory(False, t_sample, x0=x0) 122 | 123 | if isinstance(model, PseudoHamiltonianNN): 124 | if (not model.external_forces_provided): 125 | F_phnn = model.external_forces(torch.tensor(x_phnn, dtype=ttype), 126 | torch.tensor(t_sample.reshape(nsamples, 1), dtype=ttype)).detach().numpy() 127 | F_phnn -= F_phnn.mean(axis=0) 128 | F_exact = pH_system.external_forces(x_exact, t_sample) 129 | df.loc[modelname, 'External force MSE'] += ((F_phnn - F_exact)**2).mean() 130 | df.loc[modelname, 'External force MAE'] += np.abs(F_phnn - F_exact).mean() 131 | if (not model.dissipation_provided): 132 | df.loc[modelname, 'R MSE'] += ((model.R(x_exact).detach().numpy() - pH_system.R(x_exact))**2).mean() 133 | df.loc[modelname, 'R MAE'] += np.abs(model.R(x_exact).detach().numpy() - pH_system.R(x_exact)).mean() 134 | if (not model.hamiltonian_provided): 135 | if pH_system.H is not None: 136 | hamiltonian_exact = pH_system.H(x_exact) 137 | hamiltonian_phnn = model.hamiltonian(torch.tensor(x_phnn, dtype=ttype)).detach().numpy() 138 | hamiltonian_exact = pH_system.H(x_exact) 139 | hamiltonian_phnn = model.hamiltonian(torch.tensor(x_phnn, dtype=ttype)).detach().numpy() 140 | df.loc[modelname, 'H MSE'] += ((hamiltonian_exact - hamiltonian_phnn)**2).mean() 141 | df.loc[modelname, 'H MAE'] += np.abs(hamiltonian_exact - hamiltonian_phnn).mean() 142 | dH_exact = pH_system.dH(x_exact) 143 | dH_phnn = model.dH(torch.tensor(x_phnn, dtype=ttype)).detach().numpy() 144 | if system == 'tank': 145 | df.loc[modelname, 'dH tanks MSE'] += ((pH_system.tanklevels(dH_exact) - pH_system.tanklevels(dH_phnn))**2).mean() 146 | df.loc[modelname, 'dH tanks MAE'] += np.abs(pH_system.tanklevels(dH_exact) - pH_system.tanklevels(dH_phnn)).mean() 147 | df.loc[modelname, 'dH pipes MSE'] += ((pH_system.pipeflows(dH_exact) - pH_system.pipeflows(dH_phnn))**2).mean() 148 | df.loc[modelname, 'dH pipes MAE'] += np.abs(pH_system.pipeflows(dH_exact) - pH_system.pipeflows(dH_phnn)).mean() 149 | else: 150 | df.loc[modelname, 'dH x1 MSE'] += ((dH_exact[:, 0] - dH_phnn[:, 0])**2).mean() 151 | df.loc[modelname, 'dH x1 MAE'] += np.abs(dH_exact[:, 0] - dH_phnn[:, 0]).mean() 152 | df.loc[modelname, 'dH x2 MSE'] += ((dH_exact[:, 1] - dH_phnn[:, 1])**2).mean() 153 | df.loc[modelname, 'dH x2 MAE'] += np.abs(dH_exact[:, 1] - dH_phnn[:, 1]).mean() 154 | elif isinstance(model.rhs_model, BaselineSplitNN): 155 | F_baseline = model.rhs_model.network_t(torch.tensor(x_phnn, dtype=ttype), 156 | torch.tensor(t_sample.reshape(nsamples, 1), dtype=ttype)).detach().numpy() 157 | F_baseline -= F_baseline.mean(axis=0) 158 | F_exact = pH_system.external_forces(x_exact, t_sample) 159 | df.loc[modelname, 'External force MSE'] += ((F_baseline - F_exact)**2).mean() 160 | df.loc[modelname, 'External force MAE'] += np.abs(F_baseline - F_exact).mean() 161 | 162 | if system == 'tank': 163 | df.loc[modelname, 'Tanks MSE'] += ((pH_system.tanklevels(x_exact) - pH_system.tanklevels(x_phnn))**2).mean() 164 | df.loc[modelname, 'Pipes MSE'] += ((pH_system.pipeflows(x_exact) - pH_system.pipeflows(x_phnn))**2).mean() 165 | df.loc[modelname, 'Tanks MAE'] += np.abs(pH_system.tanklevels(x_exact) - pH_system.tanklevels(x_phnn)).mean() 166 | df.loc[modelname, 'Pipes MAE'] += np.abs(pH_system.pipeflows(x_exact) - pH_system.pipeflows(x_phnn)).mean() 167 | else: 168 | df.loc[modelname, 'x1 MSE'] += ((x_exact[:, 0] - x_phnn[:, 0])**2).mean() 169 | df.loc[modelname, 'x1 MAE'] += np.abs(x_exact[:, 0] - x_phnn[:, 0]).mean() 170 | df.loc[modelname, 'x2 MSE'] += ((x_exact[:, 1] - x_phnn[:, 1])**2).mean() 171 | df.loc[modelname, 'x2 MAE'] += np.abs(x_exact[:, 1] - x_phnn[:, 1]).mean() 172 | 173 | df.loc[modelname, :] = df.loc[modelname].values / nrollouts 174 | 175 | df.to_csv(os.path.join(modelfolder, f'testresults_dt{dt:.0e}_n{int(nrollouts)}_t{int(t_max)}.csv')) -------------------------------------------------------------------------------- /example_scripts/mpc_example.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | 3 | import matplotlib.pyplot as plt 4 | import numpy as np 5 | 6 | from phlearn.phsystems import init_tanksystem 7 | from phlearn.control import StepReference, PseudoHamiltonianMPC 8 | from phlearn.phnns import load_dynamic_system_model 9 | 10 | try: 11 | import casadi 12 | except ModuleNotFoundError: 13 | raise ModuleNotFoundError("To use the phlearn.control module install via 'pip install phlearn[control]' ") 14 | 15 | 16 | if __name__ == "__main__": 17 | parser = argparse.ArgumentParser() 18 | parser.add_argument('--modelpath', type=str, 19 | help='Path to a learned model (PHNN/baseline) to use as the MPC dynamics model and to simulate the tank system. If None, an analytical tank system is initialized and used instead.') 20 | 21 | args = parser.parse_args() 22 | modelpath = args.modelpath 23 | seed = 1 24 | 25 | if modelpath is not None: 26 | pH_system, optimizer, metadict = load_dynamic_system_model(modelpath) 27 | 28 | if metadict["traininginfo"]["baseline"]: 29 | S = None 30 | H = None 31 | dH = None 32 | R = None 33 | baseline = pH_system.rhs_model 34 | external_forces = None 35 | else: 36 | S = pH_system.skewsymmetric_matrix.copy() 37 | H = pH_system.hamiltonian 38 | dH = None 39 | R = pH_system.R 40 | baseline = None 41 | external_forces = pH_system.external_forces 42 | npipes = 5 43 | ntanks = 4 44 | else: 45 | pH_system = init_tanksystem() 46 | pH_system.seed(seed) 47 | 48 | S = pH_system.skewsymmetric_matrix 49 | H = None 50 | dH = pH_system.dH 51 | baseline = None 52 | npipes = pH_system.npipes 53 | ntanks = pH_system.ntanks 54 | R = pH_system.dissipation_matrix 55 | external_forces = pH_system.external_forces 56 | 57 | control_forces_filter = np.zeros((pH_system.nstates,)) 58 | control_forces_filter[5] = 1 59 | control_reference = StepReference(0, 1, step_interval=0.75, seed=seed) 60 | 61 | mpc = PseudoHamiltonianMPC(control_forces_filter, 62 | S=S, 63 | dH=dH, 64 | H=H, 65 | F=external_forces, 66 | R=R, 67 | baseline=baseline, 68 | state_names=[f'flow_{i+1}' for i in range(npipes)] + [f'level_{i+1}' for i in range(ntanks)], 69 | control_names=['inflow'], 70 | references={'Controlled tank level': control_reference}) 71 | 72 | def mpc_setup(mpc): 73 | stage_cost = (mpc.model.x[f'level_{ntanks - 1}'] - mpc.model.tvp['Controlled tank level']) ** 2 74 | terminal_cost = casadi.DM(np.zeros((1, 1))) 75 | mpc.set_objective(lterm=stage_cost, mterm=terminal_cost) 76 | mpc.set_rterm(inflow=1e-2) 77 | mpc.bounds['lower', '_u', 'inflow'] = -10 78 | mpc.bounds['upper', '_u', 'inflow'] = 10 79 | mpc.set_param(t_step=0.025) 80 | 81 | return mpc 82 | 83 | mpc.setup(setup_callback=mpc_setup) 84 | 85 | pH_system.controller = mpc 86 | 87 | t_trajectory = np.linspace(0, 2, 201) 88 | if modelpath is not None: 89 | xs, us = pH_system.simulate_trajectory('rk4', t_trajectory) 90 | else: 91 | xs, _, _, us = pH_system.sample_trajectory(t_trajectory, noise_std=0) 92 | 93 | fig, axs = plt.subplots(3, 1, sharex=True, figsize=(8, 6)) 94 | for i in range(xs.shape[-1]): 95 | axs[0].plot(t_trajectory, xs[:, i]) 96 | axs[0].set_title('States') 97 | 98 | axs[1].plot(t_trajectory, xs[:, 5]) 99 | axs[1].plot(t_trajectory, control_reference.get_reference_data(t_trajectory)[0], linestyle='dashed') 100 | axs[1].set_title('Controlled state') 101 | 102 | axs[2].plot(t_trajectory[:-1], us[:, 5]) 103 | axs[2].axhline(mpc.mpc.bounds['lower', '_u', 'inflow'].__float__(), linestyle='dotted') 104 | axs[2].axhline(mpc.mpc.bounds['upper', '_u', 'inflow'].__float__(), linestyle='dotted') 105 | axs[2].set_title('Control inputs') 106 | axs[2].set_xlabel('Time') 107 | plt.show() 108 | -------------------------------------------------------------------------------- /example_scripts/pure_kdv.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SINTEF/pseudo-hamiltonian-neural-networks/091065ef3c1b730d56fd413b6373d0424d8114be/example_scripts/pure_kdv.gif -------------------------------------------------------------------------------- /example_scripts/spring_example.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "markdown", 5 | "id": "decent-terror", 6 | "metadata": {}, 7 | "source": [ 8 | "# Pseudo-Hamiltonian neural networks" 9 | ] 10 | }, 11 | { 12 | "cell_type": "markdown", 13 | "id": "e22c6980", 14 | "metadata": {}, 15 | "source": [ 16 | "In this notebook, we will give an example of how to setup and train a neural network on simple harmonic oscillating spring with dissipation using `phlearn`. We will also demonstrate how to add an external force to the system, which can also be learnt by the pseudo-hamiltonian framework. \n", 17 | "\n", 18 | "For details, see [\"Pseudo-Hamiltonian Neural Networks with State-Dependent\n", 19 | "External Forces\"](https://arxiv.org/abs/2206.02660)." 20 | ] 21 | }, 22 | { 23 | "cell_type": "code", 24 | "execution_count": null, 25 | "id": "8f119a43", 26 | "metadata": {}, 27 | "outputs": [], 28 | "source": [ 29 | "# Uncomment for local install: \n", 30 | "# %pip install -e ../phlearn " 31 | ] 32 | }, 33 | { 34 | "cell_type": "code", 35 | "execution_count": null, 36 | "id": "conceptual-senator", 37 | "metadata": {}, 38 | "outputs": [], 39 | "source": [ 40 | "import numpy as np\n", 41 | "import torch\n", 42 | "import phlearn.phsystems.ode as phsys\n", 43 | "import phlearn.phnns as phnn\n", 44 | "import matplotlib.pyplot as plt\n", 45 | "\n", 46 | "ttype = torch.float32\n", 47 | "torch.set_default_dtype(ttype)\n" 48 | ] 49 | }, 50 | { 51 | "cell_type": "markdown", 52 | "id": "3ddfbb35", 53 | "metadata": {}, 54 | "source": [ 55 | "#### Set up the system\n", 56 | "\n", 57 | "Below is an example of how to set up a Hamiltonian system with linear dissipation using the PseudoHamiltonianSystem() class. The below block sets up the differential equation that will be used to generate the data. Initially, we will just create a simple harmonic oscilator with damping. " 58 | ] 59 | }, 60 | { 61 | "cell_type": "code", 62 | "execution_count": null, 63 | "id": "5d97328e", 64 | "metadata": {}, 65 | "outputs": [], 66 | "source": [ 67 | "nstates = 2\n", 68 | "\n", 69 | "\n", 70 | "def setup_spring_system(\n", 71 | " mass=1.0, spring_constant=1.0, damping=0.3, external_forces=None\n", 72 | "):\n", 73 | " R = np.diag([0, damping])\n", 74 | " M = np.diag([spring_constant / 2, 1 / (2 * mass)])\n", 75 | "\n", 76 | " def hamiltonian(x):\n", 77 | " return x.T @ M @ x\n", 78 | "\n", 79 | " def hamiltonian_grad(x):\n", 80 | " return 2 * M @ x\n", 81 | "\n", 82 | " spring_system = phsys.PseudoHamiltonianSystem(\n", 83 | " nstates=nstates,\n", 84 | " hamiltonian=hamiltonian,\n", 85 | " grad_hamiltonian=hamiltonian_grad,\n", 86 | " dissipation_matrix=R,\n", 87 | " external_forces=external_forces,\n", 88 | " )\n", 89 | "\n", 90 | " return spring_system\n", 91 | "\n", 92 | "\n", 93 | "spring_system = setup_spring_system()\n" 94 | ] 95 | }, 96 | { 97 | "cell_type": "markdown", 98 | "id": "e1ce788f", 99 | "metadata": {}, 100 | "source": [ 101 | "#### Generate training data\n", 102 | "\n", 103 | "Use the `spring_system` instance to generate training data, which are numerical solutions to the exact ODE. " 104 | ] 105 | }, 106 | { 107 | "cell_type": "code", 108 | "execution_count": null, 109 | "id": "e13d2fa4", 110 | "metadata": {}, 111 | "outputs": [], 112 | "source": [ 113 | "def get_training_data(data_points=30000, dt=0.1, tmax=10):\n", 114 | " nt = round(tmax / dt)\n", 115 | " t_axis = np.linspace(0, tmax, nt + 1)\n", 116 | " ntrajectories_train = int(np.ceil(data_points / nt))\n", 117 | " traindata = phnn.generate_dataset(spring_system, ntrajectories_train, t_axis)\n", 118 | " return traindata, t_axis\n", 119 | "\n", 120 | "\n", 121 | "traindata, t_axis = get_training_data()\n" 122 | ] 123 | }, 124 | { 125 | "cell_type": "markdown", 126 | "id": "3a379aae", 127 | "metadata": {}, 128 | "source": [ 129 | "#### Set up the pseudo Hamiltonian neural network\n", 130 | "We create an instance of HamiltonianNN(), which is the untrained neural network used to approximate the Hamiltonian of the system. This is passed to the PseudoHamiltonianNN() class along with an instance of R_estimator(), which islearns the damping coefficient during training. \n", 131 | "\n", 132 | "We will allow additional keyword arguments to be passed to PseudoHamiltonianNN() so we have the option of adding an external force in the system." 133 | ] 134 | }, 135 | { 136 | "cell_type": "code", 137 | "execution_count": null, 138 | "id": "8ab6d1e1", 139 | "metadata": {}, 140 | "outputs": [], 141 | "source": [ 142 | "def setup_pseudo_hamiltonian_nn(**kwargs):\n", 143 | " states_dampened = np.diagonal(spring_system.dissipation_matrix) != 0\n", 144 | " phmodel = phnn.PseudoHamiltonianNN(\n", 145 | " nstates,\n", 146 | " dissipation_est=phnn.R_estimator(states_dampened),\n", 147 | " **kwargs,\n", 148 | " )\n", 149 | " return phmodel\n", 150 | "\n", 151 | "\n", 152 | "phmodel = setup_pseudo_hamiltonian_nn()\n" 153 | ] 154 | }, 155 | { 156 | "cell_type": "markdown", 157 | "id": "342992c4", 158 | "metadata": {}, 159 | "source": [ 160 | "#### Setup a baseline model\n", 161 | "To compare against PseudoHamiltonianNN() , we will create a baseline model which will approximate the dynamics using a standard fully connected multilayer perceptron. " 162 | ] 163 | }, 164 | { 165 | "cell_type": "code", 166 | "execution_count": null, 167 | "id": "e6e521a3", 168 | "metadata": {}, 169 | "outputs": [], 170 | "source": [ 171 | "def setup_baseline_nn(hidden_dim=100):\n", 172 | " baseline_nn = phnn.BaselineNN(nstates, hidden_dim)\n", 173 | " basemodel = phnn.DynamicSystemNN(nstates, baseline_nn)\n", 174 | " return basemodel\n", 175 | "\n", 176 | "\n", 177 | "basemodel = setup_baseline_nn()\n" 178 | ] 179 | }, 180 | { 181 | "cell_type": "markdown", 182 | "id": "18a8f4a1", 183 | "metadata": {}, 184 | "source": [ 185 | "#### Train the models" 186 | ] 187 | }, 188 | { 189 | "cell_type": "code", 190 | "execution_count": null, 191 | "id": "f3b29ddd", 192 | "metadata": {}, 193 | "outputs": [], 194 | "source": [ 195 | "def train_models(*models, epochs=30, batch_size=32, **kwargs):\n", 196 | " for model in models:\n", 197 | " model, _ = phnn.train(\n", 198 | " model,\n", 199 | " integrator=\"midpoint\",\n", 200 | " traindata=traindata,\n", 201 | " epochs=epochs,\n", 202 | " batch_size=batch_size,\n", 203 | " **kwargs\n", 204 | " )\n", 205 | " return models\n", 206 | "\n", 207 | "\n", 208 | "phmodel, basemodel = train_models(phmodel, basemodel)\n" 209 | ] 210 | }, 211 | { 212 | "cell_type": "markdown", 213 | "id": "150df002", 214 | "metadata": {}, 215 | "source": [ 216 | "#### Plot the results\n", 217 | "\n", 218 | "In this cell we compare the learned damping constant against the true value and compare some trajectories against the exact solution. " 219 | ] 220 | }, 221 | { 222 | "cell_type": "code", 223 | "execution_count": null, 224 | "id": "f328dad4", 225 | "metadata": {}, 226 | "outputs": [], 227 | "source": [ 228 | "def print_damping_constants():\n", 229 | " learned_damping_constant = phmodel.R().detach().numpy()[1, 1]\n", 230 | " true_damping_constant = spring_system.dissipation_matrix[1, 1]\n", 231 | " print(\n", 232 | " f\"true damping constant : {true_damping_constant} \\\n", 233 | " \\nlearned damping constant: {learned_damping_constant}\"\n", 234 | " )\n", 235 | "\n", 236 | "\n", 237 | "def get_trajectories(x0, t_axis):\n", 238 | " x_exact, *_ = spring_system.sample_trajectory(t_axis, x0=x0)\n", 239 | " x_phnn, _ = phmodel.simulate_trajectory(integrator=False, t_sample=t_axis, x0=x0)\n", 240 | " x_baseline, _ = basemodel.simulate_trajectory(\n", 241 | " integrator=False, t_sample=t_axis, x0=x0\n", 242 | " )\n", 243 | " return x_exact, x_phnn, x_baseline\n", 244 | "\n", 245 | "\n", 246 | "def plot_trajectories(x_exact, x_phnn, x_baseline):\n", 247 | " plt.figure(figsize=(5, 5))\n", 248 | " plt.plot(x_exact[:, 0], x_exact[:, 1], color=\"k\", linestyle=\"dashed\", label=\"Exact\")\n", 249 | " plt.plot(x_exact[0, 0], x_exact[0, 1], \"ko\")\n", 250 | " plt.plot(x_phnn[:, 0], x_phnn[:, 1], label=\"PHNN\")\n", 251 | " plt.plot(x_baseline[:, 0], x_baseline[:, 1], label=\"Baseline\")\n", 252 | " plt.legend()\n", 253 | " plt.title(\"Phase plot\")\n", 254 | " plt.show()\n", 255 | "\n", 256 | "\n", 257 | "x_exact, x_phnn, x_baseline = get_trajectories(x0=[1, 0], t_axis=t_axis)\n", 258 | "print_damping_constants()\n", 259 | "plot_trajectories(x_exact, x_phnn, x_baseline)\n" 260 | ] 261 | }, 262 | { 263 | "cell_type": "markdown", 264 | "id": "d5aab22b", 265 | "metadata": {}, 266 | "source": [ 267 | "### Learning an external force\n", 268 | "Here we will repeat the above process with the inclusion of an external force on the spring system. \n", 269 | "\n", 270 | "First we define the external force `F` that will be used to instantiate the `PseudoHamiltonianSystem` class. Next, we construct a neural network to learn this force using `ExternalForcesNN`, which is used to instantiate the `PseudoHamiltonianNN` class. " 271 | ] 272 | }, 273 | { 274 | "cell_type": "code", 275 | "execution_count": null, 276 | "id": "ec9d77fb", 277 | "metadata": {}, 278 | "outputs": [], 279 | "source": [ 280 | "def F(x, t):\n", 281 | " f0 = 0.5\n", 282 | " omega = np.pi / 2\n", 283 | " return np.array([0, f0 * np.sin(omega * t)])\n", 284 | "\n", 285 | "\n", 286 | "ext_forces_nn = phnn.ExternalForcesNN(\n", 287 | " nstates=nstates,\n", 288 | " noutputs=1, # force affects one state variable only, specified by the parameter external_forces_filter\n", 289 | " external_forces_filter=[\n", 290 | " 0,\n", 291 | " 1,\n", 292 | " ], # tells the NN to give output only in the second dimension\n", 293 | " hidden_dim=100,\n", 294 | " timedependent=True,\n", 295 | " statedependent=False, # force depends on time only\n", 296 | ")\n" 297 | ] 298 | }, 299 | { 300 | "cell_type": "markdown", 301 | "id": "edb6b9e4", 302 | "metadata": {}, 303 | "source": [ 304 | "Now we will repeat the setup, data generation and training steps with external forces. " 305 | ] 306 | }, 307 | { 308 | "cell_type": "code", 309 | "execution_count": null, 310 | "id": "fb511af6", 311 | "metadata": {}, 312 | "outputs": [], 313 | "source": [ 314 | "spring_system = setup_spring_system(external_forces=F)\n", 315 | "phmodel = setup_pseudo_hamiltonian_nn(external_forces_est=ext_forces_nn)\n", 316 | "basemodel = setup_baseline_nn()\n", 317 | "traindata, t_axis = get_training_data(dt=0.1, tmax=10, data_points=50000)\n", 318 | "phmodel, basemodel = train_models(phmodel, basemodel)\n", 319 | "x_exact, x_phnn, x_baseline = get_trajectories(x0=[1, 0], t_axis=t_axis)\n", 320 | "print_damping_constants()\n", 321 | "plot_trajectories(x_exact, x_phnn, x_baseline)\n" 322 | ] 323 | } 324 | ], 325 | "metadata": { 326 | "kernelspec": { 327 | "display_name": "Python (phnn_pde)", 328 | "language": "python", 329 | "name": "phnn_pde" 330 | }, 331 | "language_info": { 332 | "codemirror_mode": { 333 | "name": "ipython", 334 | "version": 3 335 | }, 336 | "file_extension": ".py", 337 | "mimetype": "text/x-python", 338 | "name": "python", 339 | "nbconvert_exporter": "python", 340 | "pygments_lexer": "ipython3", 341 | "version": "3.10.9" 342 | }, 343 | "vscode": { 344 | "interpreter": { 345 | "hash": "9750f35fb4e247739abdd4f1bd3b23f35c1677f637abade63436e65eb57ee7de" 346 | } 347 | } 348 | }, 349 | "nbformat": 4, 350 | "nbformat_minor": 5 351 | } 352 | -------------------------------------------------------------------------------- /example_scripts/train_model.py: -------------------------------------------------------------------------------- 1 | 2 | import argparse 3 | 4 | import numpy as np 5 | import torch 6 | 7 | from phlearn.phsystems.ode import init_tanksystem, init_msdsystem 8 | from phlearn.phnns import PseudoHamiltonianNN, DynamicSystemNN, load_dynamic_system_model 9 | from phlearn.phnns import R_estimator, BaselineNN, BaselineSplitNN, HamiltonianNN, ExternalForcesNN 10 | from phlearn.phnns import npoints_to_ntrajectories_tsample, train, generate_dataset 11 | 12 | ttype = torch.float32 13 | torch.set_default_dtype(ttype) 14 | 15 | if __name__ == "__main__": 16 | parser = argparse.ArgumentParser() 17 | parser.add_argument('--system', type=str, choices=['tank', 'msd'], 18 | required=True, 19 | help='Choose to train a tank or a ' 20 | 'forced mass spring damper.') 21 | parser.add_argument('--baseline', type=int, default=0, choices=[0, 1, 2], 22 | help='If 1 use baseline network x_dot = network(x, t). ' 23 | 'If 2 use split baseline network ' 24 | 'x_dot = network_x(x) + network_t(x).') 25 | parser.add_argument('--storedir', type=str, 26 | help='Directory for storing the best model in terms ' 27 | 'of validation loss.') 28 | parser.add_argument('--modelname', type=str, 29 | help='Name to use for the stored model.') 30 | parser.add_argument('--modelpath', type=str, 31 | help='Path to existing model to continue training.') 32 | parser.add_argument('--ntrainingpoints', '-n', type=int, default=3000, 33 | help='Number of training points.') 34 | parser.add_argument('--ntrajectories_val', type=int, default=0, 35 | help='Number of trajectories for validation.') 36 | parser.add_argument('--sampling_time', type=float, default=1/30, 37 | help='Sampling time.') 38 | parser.add_argument('--t_max', type=float, default=1, 39 | help='Length of trajectory.') 40 | parser.add_argument('--true_derivatives', action='store_true', 41 | help='Use the true derivative values for training. ' 42 | 'If not provided derivatives in the training ' 43 | 'data are estimated by the finite differences.') 44 | parser.add_argument('--integrator', type=str, 45 | choices=[False, 'euler', 'rk4', 'midpoint', 'srk4'], 46 | default='midpoint', 47 | help='Integrator used during training.') 48 | parser.add_argument('--F_timedependent', type=int, default=1, 49 | choices=[0, 1], 50 | help='If 1, make external force NN (or baseline NN) ' 51 | 'depend on time.') 52 | parser.add_argument('--F_statedependent', type=int, default=1, 53 | choices=[0, 1], 54 | help='If 1, make external force NN (or baseline NN) ' 55 | 'depend on state.') 56 | parser.add_argument('--hidden_dim', type=int, default=100, 57 | help='Hidden dimension of fully connected neural ' 58 | 'network layers.') 59 | parser.add_argument('--learning_rate', type=float, default=1e-3, 60 | help='Learning rate of Adam optimizer.') 61 | parser.add_argument('--batch_size', type=int, default=32, 62 | help='Batch size used in training.') 63 | parser.add_argument('--epochs', type=int, default=100, 64 | help='Number of training epochs.') 65 | parser.add_argument('--l1_param_forces', type=float, default=0., 66 | help='L1 penalty parameter of external force estimate.') 67 | parser.add_argument('--l1_param_dissipation', type=float, default=0., 68 | help='L1 penalty parameter of dissipation estimate.') 69 | parser.add_argument('--early_stopping_patience', type=int, 70 | help=('Number of epochs to continue training without ' 71 | 'a decrease in validation loss ' 72 | 'of at least early_stopping_delta.')) 73 | parser.add_argument('--early_stopping_delta', type=float, 74 | help='Minimum accepted decrease in validation loss to ' 75 | 'prevent early stopping.') 76 | parser.add_argument('--shuffle', action='store_true', 77 | help='Shuffle training data at every epoch.') 78 | parser.add_argument('--noise_std', type=float, default=0., 79 | help='Noise level for training.') 80 | parser.add_argument('--store_results', '-s', action='store_true', 81 | help='Store trained model and prediction results.') 82 | parser.add_argument('--seed', type=int, 83 | help='Seed for the simulated dynamic system.') 84 | parser.add_argument('--verbose', '-v', action='store_true', 85 | help='Print information while training.') 86 | args = parser.parse_args() 87 | 88 | system = args.system 89 | baseline = bool(args.baseline) 90 | storedir = args.storedir 91 | modelname = args.modelname 92 | modelpath = args.modelpath 93 | ntrainingpoints = args.ntrainingpoints 94 | sampling_time = args.sampling_time 95 | t_max = args.t_max 96 | true_derivatives = args.true_derivatives 97 | if true_derivatives: 98 | integrator = False 99 | print('Warning: As exact derivatives are used when generating ' 100 | 'training data, (true_derivatives = True) integrator' 101 | 'is set to False.') 102 | else: 103 | integrator = args.integrator 104 | F_timedependent = bool(args.F_timedependent) 105 | F_statedependent = bool(args.F_statedependent) 106 | hidden_dim = args.hidden_dim 107 | learning_rate = args.learning_rate 108 | batch_size = args.batch_size 109 | epochs = args.epochs 110 | l1_param_forces = args.l1_param_forces 111 | l1_param_dissipation = args.l1_param_dissipation 112 | shuffle = args.shuffle 113 | early_stopping_patience = args.early_stopping_patience 114 | early_stopping_delta = args.early_stopping_delta 115 | noise_std = args.noise_std 116 | store_results = args.store_results 117 | seed = args.seed 118 | verbose = args.verbose 119 | ntrajectories_val = args.ntrajectories_val 120 | 121 | ntrajectories_train, t_sample = npoints_to_ntrajectories_tsample( 122 | ntrainingpoints, t_max, sampling_time) 123 | 124 | if system == 'tank': 125 | pH_system = init_tanksystem() 126 | damped_states = np.arange(pH_system.nstates) < pH_system.npipes 127 | else: 128 | pH_system = init_msdsystem() 129 | damped_states = [False, True] 130 | 131 | pH_system.seed(seed) 132 | nstates = pH_system.nstates 133 | 134 | if modelpath is not None: 135 | model, optimizer, metadict = load_dynamic_system_model(modelpath) 136 | else: 137 | if baseline == 1: 138 | baseline_nn = BaselineNN( 139 | nstates, hidden_dim, 140 | timedependent=F_timedependent, statedependent=True) 141 | model = DynamicSystemNN(nstates, baseline_nn) 142 | elif baseline == 2: 143 | external_forces_filter_t = np.zeros(nstates) 144 | external_forces_filter_t[-1] = 1 145 | baseline_nn = BaselineSplitNN( 146 | nstates, hidden_dim, noutputs_x=nstates, 147 | noutputs_t=1, external_forces_filter_x=None, 148 | external_forces_filter_t=external_forces_filter_t, 149 | ttype=ttype) 150 | model = DynamicSystemNN(nstates, baseline_nn) 151 | else: 152 | hamiltonian_nn = HamiltonianNN(nstates, hidden_dim) 153 | external_forces_filter = np.zeros(nstates) 154 | external_forces_filter[-1] = 1 155 | ext_forces_nn = ExternalForcesNN( 156 | nstates, 1, hidden_dim=hidden_dim, 157 | timedependent=F_timedependent, 158 | statedependent=F_statedependent, 159 | external_forces_filter=external_forces_filter) 160 | 161 | r_est = R_estimator(damped_states) 162 | 163 | model = PseudoHamiltonianNN( 164 | nstates, 165 | pH_system.skewsymmetric_matrix, 166 | hamiltonian_est=hamiltonian_nn, 167 | dissipation_est=r_est, 168 | external_forces_est=ext_forces_nn) 169 | optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate, 170 | weight_decay=1e-4) 171 | 172 | traindata = generate_dataset( 173 | pH_system, ntrajectories_train, t_sample, true_derivatives, 174 | nsamples=ntrainingpoints, noise_std=noise_std) 175 | valdata = generate_dataset( 176 | pH_system, ntrajectories_val, t_sample, true_derivatives, noise_std=noise_std) 177 | 178 | bestmodel, vloss = train( 179 | model, 180 | integrator, 181 | traindata, 182 | optimizer, 183 | valdata=valdata, 184 | epochs=epochs, 185 | batch_size=batch_size, 186 | shuffle=shuffle, 187 | l1_param_forces=l1_param_forces, 188 | l1_param_dissipation=l1_param_dissipation, 189 | loss_fn=torch.nn.MSELoss(), 190 | verbose=verbose, 191 | early_stopping_patience=early_stopping_patience, 192 | early_stopping_delta=early_stopping_delta, 193 | return_best=True, 194 | store_best=store_results, 195 | store_best_dir=storedir, 196 | modelname=modelname, 197 | trainingdetails=vars(args) 198 | ) 199 | -------------------------------------------------------------------------------- /phlearn/LICENSE.txt: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) [2022] [SINTEF Digital] 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 | -------------------------------------------------------------------------------- /phlearn/README.md: -------------------------------------------------------------------------------- 1 | # phlearn 2 | Package for modelling pseudo-Hamiltonian systems with neural networks as described in [(Eidnes et al., 2022)](https://arxiv.org/pdf/2206.02660.pdf) for ODEs and [(Eidnes and Lye, 2023)](https://arxiv.org/pdf/2304.14374.pdf) for PDEs. 3 | 4 | [Documentation available here](https://pseudo-Hamiltonian-neural-networks.readthedocs.io/en/latest/) 5 | 6 | - phlearn 7 | + phsystems 8 | + phnns 9 | + utils 10 | 11 | phsystems implements classes and functionality for simulating pseudo-Hamiltonian systems. 12 | phnns implements classes and functionality for learning and inference with pseudo-Hamiltonian neural networks. 13 | -------------------------------------------------------------------------------- /phlearn/docs/Makefile: -------------------------------------------------------------------------------- 1 | # Minimal makefile for Sphinx documentation 2 | # 3 | 4 | # You can set these variables from the command line, and also 5 | # from the environment for the first two. 6 | SPHINXOPTS ?= 7 | SPHINXBUILD ?= sphinx-build 8 | SOURCEDIR = ./source 9 | BUILDDIR = _build 10 | 11 | # Put it first so that "make" without argument is like "make help". 12 | help: 13 | @$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) 14 | 15 | .PHONY: help Makefile 16 | 17 | # Catch-all target: route all unknown targets to Sphinx using the new 18 | # "make mode" option. $(O) is meant as a shortcut for $(SPHINXOPTS). 19 | %: Makefile 20 | @$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) 21 | -------------------------------------------------------------------------------- /phlearn/docs/make.bat: -------------------------------------------------------------------------------- 1 | @ECHO OFF 2 | 3 | pushd %~dp0 4 | 5 | REM Command file for Sphinx documentation 6 | 7 | if "%SPHINXBUILD%" == "" ( 8 | set SPHINXBUILD=sphinx-build 9 | ) 10 | set SOURCEDIR=./source 11 | set BUILDDIR=_build 12 | 13 | %SPHINXBUILD% >NUL 2>NUL 14 | if errorlevel 9009 ( 15 | echo. 16 | echo.The 'sphinx-build' command was not found. Make sure you have Sphinx 17 | echo.installed, then set the SPHINXBUILD environment variable to point 18 | echo.to the full path of the 'sphinx-build' executable. Alternatively you 19 | echo.may add the Sphinx directory to PATH. 20 | echo. 21 | echo.If you don't have Sphinx installed, grab it from 22 | echo.https://www.sphinx-doc.org/ 23 | exit /b 1 24 | ) 25 | 26 | if "%1" == "" goto help 27 | 28 | %SPHINXBUILD% -M %1 %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O% 29 | goto end 30 | 31 | :help 32 | %SPHINXBUILD% -M help %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O% 33 | 34 | :end 35 | popd 36 | -------------------------------------------------------------------------------- /phlearn/docs/requirements.txt: -------------------------------------------------------------------------------- 1 | 2 | sphinx==5.1.1 3 | sphinx_rtd_theme==1.0.0 4 | numpydoc==1.4.0 5 | casadi==3.5.5 6 | do-mpc==4.3.5 7 | -------------------------------------------------------------------------------- /phlearn/docs/source/conf.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.insert(0, os.path.abspath('../..')) 4 | 5 | # Configuration file for the Sphinx documentation builder. 6 | # 7 | # For the full list of built-in configuration values, see the documentation: 8 | # https://www.sphinx-doc.org/en/master/usage/configuration.html 9 | 10 | # -- Project information ----------------------------------------------------- 11 | # https://www.sphinx-doc.org/en/master/usage/configuration.html#project-information 12 | 13 | project = 'phlearn' 14 | copyright = '2022, Camilla Sterud, Sølve Eidnes, Eivind Bøhn, Signe Riemer-Sørensen, Alexander J. Stasik' 15 | author = 'Camilla Sterud, Sølve Eidnes, Eivind Bøhn, Signe Riemer-Sørensen, Alexander J. Stasik' 16 | release = '1.1.2' 17 | 18 | # -- General configuration --------------------------------------------------- 19 | # https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration 20 | 21 | 22 | extensions = [ 23 | 'sphinx.ext.duration', 24 | 'sphinx.ext.doctest', 25 | 'numpydoc' 26 | ] 27 | 28 | numpydoc_class_members_toctree = False 29 | numpydoc_show_inherited_class_members = False 30 | intersphinx_mapping = { 31 | 'python': ('https://docs.python.org/3/', None), 32 | 'sphinx': ('https://www.sphinx-doc.org/en/master/', None), 33 | } 34 | intersphinx_disabled_domains = ['std'] 35 | 36 | exclude_patterns = ['_build', 'Thumbs.db', '.DS_Store'] 37 | 38 | # -- Options for HTML output ------------------------------------------------- 39 | # https://www.sphinx-doc.org/en/master/usage/configuration.html#options-for-html-output 40 | 41 | html_theme = 'sphinx_rtd_theme' 42 | 43 | # -- Options for EPUB output 44 | epub_show_urls = 'footnote' -------------------------------------------------------------------------------- /phlearn/docs/source/control.rst: -------------------------------------------------------------------------------- 1 | phlearn.control 2 | ------------------------ 3 | 4 | .. automodule:: phlearn.control 5 | :members: 6 | :undoc-members: 7 | :show-inheritance: 8 | -------------------------------------------------------------------------------- /phlearn/docs/source/index.rst: -------------------------------------------------------------------------------- 1 | 2 | ========================================== 3 | Welcome to phlearn' documentation 4 | ========================================== 5 | 6 | phlearn is a python package for modelling pseudo-Hamiltonian systems with neural networks as decribed in `(Eidnes et al. 2022) `_. 7 | 8 | Installation 9 | ============ 10 | 11 | The phlearn package is available via PyPi: 12 | 13 | :: 14 | 15 | $ pip install phlearn 16 | 17 | Alternatively, to get the latest updates not yet available on PyPi, you can clone the repositrory from Github and install directly: 18 | 19 | :: 20 | 21 | $ git clone https://github.com/SINTEF/pseudo-Hamiltonian-neural-networks.git 22 | $ cd pseudo-Hamiltonian-neural-networks 23 | $ pip install -e phlearn 24 | 25 | 26 | Example use 27 | =========== 28 | Example scripts and notebooks can be found in `the Github repo `_. 29 | 30 | 31 | phlearn API 32 | ==================== 33 | 34 | .. automodule:: phlearn 35 | 36 | .. toctree:: 37 | :maxdepth: 4 38 | :hidden: 39 | 40 | phsystems 41 | phnns 42 | control 43 | utils 44 | -------------------------------------------------------------------------------- /phlearn/docs/source/phnns.rst: -------------------------------------------------------------------------------- 1 | phlearn.phnns 2 | ---------------------- 3 | 4 | .. automodule:: phlearn.phnns 5 | :members: 6 | :undoc-members: 7 | :show-inheritance: 8 | -------------------------------------------------------------------------------- /phlearn/docs/source/phsystems_ode.rst: -------------------------------------------------------------------------------- 1 | phlearn.phsystems.ode 2 | -------------------------- 3 | 4 | .. automodule:: phlearn.phsystems.ode 5 | :members: 6 | :undoc-members: 7 | :show-inheritance: 8 | -------------------------------------------------------------------------------- /phlearn/docs/source/phsystems_pde.rst: -------------------------------------------------------------------------------- 1 | phlearn.phsystems.pde 2 | -------------------------- 3 | 4 | .. automodule:: phlearn.phsystems.pde 5 | :members: 6 | :undoc-members: 7 | :show-inheritance: 8 | -------------------------------------------------------------------------------- /phlearn/docs/source/utils.rst: -------------------------------------------------------------------------------- 1 | phlearn.utils 2 | ---------------------- 3 | 4 | .. automodule:: phlearn.utils 5 | :members: 6 | :undoc-members: 7 | :show-inheritance: 8 | -------------------------------------------------------------------------------- /phlearn/phlearn/__init__.py: -------------------------------------------------------------------------------- 1 | """ 2 | phlearn is a package for simulating, controlling and learning 3 | dynamical systems in general, and pseudo-Hamiltonian systems in particular. 4 | phlearn contains three subpackages: 5 | 6 | - :py:mod:`~.phsystems` with tools for simulation 7 | 8 | - :py:mod:`~.phnns` with tools for constructing and training 9 | pseudo-Hamiltonian neural networks 10 | 11 | - :py:mod:`~.control` for control functionality 12 | 13 | - :py:mod:`~.utils` for convenience functions. 14 | 15 | """ 16 | try: 17 | from . import control 18 | except ModuleNotFoundError: 19 | # from warnings import warn 20 | # warn("Loading phlearn without control module") 21 | pass 22 | from . import phnns 23 | from . import phsystems 24 | from . import utils 25 | -------------------------------------------------------------------------------- /phlearn/phlearn/control/__init__.py: -------------------------------------------------------------------------------- 1 | """ 2 | The control subpackage implements PID and MPC controllers for 3 | pseudo-Hamiltonian systems and pseudo-Hamiltonian neural networks. 4 | 5 | Controller classes 6 | --------------------------------------------------- 7 | 8 | :py:class:`~.pid.PIDController` 9 | 10 | :py:class:`~.mpc.PseudoHamiltonianMPC` 11 | 12 | Reference classes 13 | ----------------- 14 | :py:class:`~.pid.Reference` 15 | 16 | :py:class:`~.pid.ConstantReference` 17 | 18 | :py:class:`~.pid.StepReference` 19 | 20 | :py:class:`~.pid.PoissonStepReference` 21 | 22 | :py:class:`~.pid.FixedReference` 23 | 24 | """ 25 | 26 | from . import reference 27 | from .reference import * 28 | from . import pid 29 | from .pid import * 30 | from . import mpc 31 | from .mpc import * 32 | from . import phcontroller 33 | from .phcontroller import * 34 | from . import casadiNN 35 | from .casadiNN import * 36 | from . import casadiPH 37 | from .casadiPH import * 38 | 39 | __all__ = mpc.__all__.copy() 40 | __all__ += pid.__all__.copy() 41 | __all__ += phcontroller.__all__.copy() 42 | __all__ += casadiNN.__all__.copy() 43 | __all__ += casadiPH.__all__.copy() -------------------------------------------------------------------------------- /phlearn/phlearn/control/casadiNN.py: -------------------------------------------------------------------------------- 1 | try: 2 | from casadi import * 3 | import casadi.tools 4 | except ModuleNotFoundError: 5 | raise ModuleNotFoundError("To use the phlearn.control module install via 'pip install phlearn[control]' ") 6 | 7 | __all__ = ['CasadiFCNN'] 8 | 9 | class CasadiFCNN: 10 | """ 11 | Casadi-implementation of a fully-connected neural network. The class 12 | takes in a specification of a neural network 13 | model and implements the specification as a Casadi function. 14 | 15 | parameters 16 | ---------- 17 | layers : list of dicts 18 | List of dictionaries where each entry describes a layer of 19 | the neural network as returned by the 20 | "get_pytorch_model_architecture" function. 21 | """ 22 | def __init__(self, layers=()): 23 | self.layers = layers 24 | self.params = None 25 | self.params_num = None 26 | 27 | self.forward_fn = None 28 | 29 | def create_forward(self, *args): 30 | input_data_cat = vertcat(*args).T 31 | blank = SX.sym('blank') # placeholder for activation function input 32 | 33 | # neuron_weights=SX.sym('neuron_weights') 34 | in_size = input_data_cat.shape[1] 35 | h_l_ps = [] 36 | 37 | act_funs = {'tanh': Function('tanh_f', [blank], [casadi.tanh(blank)]), 38 | 'relu': Function('relu_f', [blank], [casadi.fmax(0, blank)])} 39 | 40 | hidden_layer = input_data_cat 41 | for l_i, layer in enumerate(self.layers): 42 | if layer['type'] == 'activation': 43 | assert layer['name'] in act_funs 44 | hidden_layer = act_funs[layer['name']](hidden_layer) 45 | elif layer['type'] == 'linear': 46 | l_ws = casadi.tools.entry(f'hl_{l_i}_weights', sym=SX.sym(f'hl_{l_i}_weights', hidden_layer.shape[1], layer['out_features'])) 47 | h_l_ps.append(l_ws) 48 | 49 | hidden_layer = (hidden_layer @ l_ws.sym) 50 | 51 | if layer['bias']: 52 | l_bs = casadi.tools.entry(f'hl_{l_i}_bias', sym=SX.sym(f'hl_{l_i}_bias', layer['out_features'], 1)) 53 | h_l_ps.append(l_bs) 54 | hidden_layer = hidden_layer + (DM.ones(hidden_layer.shape[0], layer['out_features']) @ diag(l_bs.sym)) 55 | else: 56 | raise ValueError(f'Unsupported layer type {layer["type"]}, must be one of [linear, activation]') 57 | 58 | self.params = casadi.tools.struct_symSX(h_l_ps) 59 | self.params_num = np.zeros(self.params.shape) 60 | 61 | self.forward_fn = Function('nn_forward', list(args) + [self.params], [hidden_layer]) 62 | 63 | return hidden_layer 64 | 65 | def set_weights_and_biases(self, params, source='pytorch'): 66 | if source == 'tensorflow': 67 | self.params_num = np.concatenate([p.flatten(order='F') for p in params]).reshape(-1, 1) 68 | elif source == 'pytorch': 69 | self.params_num = np.concatenate([p.flatten() for p in params]).reshape(-1, 1) 70 | else: 71 | raise ValueError 72 | 73 | def get_parameters(self): 74 | return self.params_num 75 | 76 | def test_correct_output(self, gt_model, n=10): 77 | for i in range(n): 78 | x = np.random.normal(0, 1, size=self.params.entries[0].sym.shape[0]) 79 | assert np.isclose(self.forward_fn(x, self.params_num), gt_model(x)).all() 80 | 81 | 82 | def get_pytorch_model_parameters(model): 83 | res = [] 84 | for p_name, p in model.named_parameters(): 85 | if 'weight' in p_name: 86 | res.append(p.detach().numpy()) 87 | elif 'bias' in p_name: 88 | res.append(p.detach().numpy()) 89 | else: 90 | raise ValueError 91 | 92 | return res 93 | 94 | 95 | def get_pytorch_model_architecture(model): 96 | res = [] 97 | for l_i, layer in enumerate(model.modules()): 98 | if l_i <= 1: 99 | continue 100 | l = {'name': str(layer).split('(')[0].lower()} 101 | if not hasattr(layer, 'out_features'): # is an activation layer 102 | l['type'] = 'activation' 103 | res.append(l) 104 | else: 105 | l['type'] = 'linear' 106 | l['bias'] = layer.bias is not None 107 | l['out_features'] = layer.out_features 108 | res.append(l) 109 | 110 | return res 111 | -------------------------------------------------------------------------------- /phlearn/phlearn/control/casadiPH.py: -------------------------------------------------------------------------------- 1 | try: 2 | from casadi import * 3 | except ModuleNotFoundError: 4 | raise ModuleNotFoundError("To use the phlearn.control module install via 'pip install phlearn[control]' ") 5 | 6 | __all__ = ['CasadiPseudoHamiltonianSystem'] 7 | 8 | class CasadiPseudoHamiltonianSystem: 9 | """ 10 | Casadi implementation of the pseudo-Hamiltonian framework, modeling a 11 | system of the form:: 12 | 13 | dx/dt = (S - R)*grad(H) + F(x, t) 14 | 15 | where S is the skew-symmetric interconnection matrix, 16 | R is a diagonal positive semi-definite damping/dissipation matrix, 17 | H is the Hamiltonian of the system, F is the external interaction 18 | and x is the system state. The system dimension is denoted by nstates. 19 | 20 | """ 21 | def __init__(self, S, dH, u, R=None, F=None): 22 | self.S = S # S can possibly be dependent on x? 23 | self.nstates = S.shape[0] 24 | self.dH = dH 25 | if R is None: 26 | self.R = np.zeros((self.nstates, self.nstates)) 27 | elif callable(R): 28 | self.R = R 29 | elif len(R.shape) == 1: 30 | self.R = np.diag(R) 31 | else: 32 | self.R = R 33 | assert self.R.shape == S.shape, ('R must be of size (nstates, nstates),' 34 | f' same as S ({S.shape}), ' 35 | f'but is of size {R.shape}.') 36 | self.F = F 37 | self.u = u 38 | 39 | def create_forward(self): 40 | return (self.S - self.R) @ self.dH + self.F + self.u 41 | -------------------------------------------------------------------------------- /phlearn/phlearn/control/mpc.py: -------------------------------------------------------------------------------- 1 | from .phcontroller import PseudoHamiltonianController 2 | from ..phnns.ode_models import R_estimator 3 | from .casadiPH import CasadiPseudoHamiltonianSystem 4 | from .casadiNN import CasadiFCNN, get_pytorch_model_parameters, get_pytorch_model_architecture 5 | import numpy as np 6 | import torch 7 | 8 | try: 9 | import do_mpc 10 | import casadi 11 | except ModuleNotFoundError: 12 | raise ModuleNotFoundError("To use the phlearn.control module install via 'pip install phlearn[control]' ") 13 | 14 | __all__ = ['PseudoHamiltonianMPC'] 15 | 16 | class PseudoHamiltonianMPC(PseudoHamiltonianController): 17 | """ 18 | This class implements a model predictive controller (MPC) that 19 | solves an optimal control problem to decide control inputs, 20 | where the model is formulated as a pseudo-Hamiltonian system:: 21 | 22 | dx/dt = (S - R) * grad[H(x)] + F(x, t) + u(x, t) 23 | 24 | or if the baseline argument is provided:: 25 | 26 | dx/dt = Baseline(x, t) + u(x, t) 27 | 28 | Each component of the model can be provided as either a python 29 | function or as a pytorch neural network (NN). Note that for any 30 | component implemented as a python function mist use only operations 31 | that are compatible with casadi variable types. 32 | 33 | The MPC is based on the do-mpc python toolbox, see its documentation 34 | for configuration options and behaviour (do-mpc.com). 35 | 36 | Parameters 37 | ---------- 38 | control_forces_filter : matrix 39 | A binary matrix of (nstates, nstates) or vector of (nstates) 40 | where 1 signifies that the corresponding state has a control 41 | input in its derivative right-hand-side. 42 | S : matrix, default None 43 | (nstates, nstates) ndarray. nstates if inferred from this. 44 | dH : callable, default None 45 | Function/NN computing the gradient of the Hamiltonian. Takes one 46 | argument (state). Accepts an ndarray both of size (nstates,) and 47 | of size (nsamples, nstates), and returns either an ndarray of 48 | size (nstates,) or an ndarray of size (nsamples, nstates), 49 | correspondingly. 50 | H : callable, default None 51 | Function/NN computing the Hamiltonian of the system. Takes one 52 | argument (state,). Accepts ndarrays both of size (nstates,) and 53 | of size (nsamples, nstates), and returns either a scalar or an 54 | ndarray of size (nstates,), correspondingly. 55 | R : matrix or array, default None 56 | (nstates, nstates) ndarray or (N,) ndarray of diagonal elements 57 | or NN 58 | F : callable, default None 59 | Function/NN computing external ports taking two arguments 60 | (state and time), which can be both an ndarray of size 61 | (nstates,) + a scalar, and an ndarray of size 62 | (nsamples, nstates) + an ndarray of size (nstates,). Returns 63 | either an ndarray of size (nstates,) or an ndarray of size 64 | (nsamples, nstates), correspondingly. 65 | baseline : callable, default None 66 | Alternative model formulation 67 | state_names : list, default None 68 | List of length (nstates) where each entry sets the name of the 69 | model state variables such that states can be referenced by 70 | these names. 71 | control_names : list, default None 72 | List of size (ncontrols) where each entry sets the name of the 73 | model control input variables such that inputs can be referenced 74 | by these names. 75 | references : dict, default None 76 | Dictionary with reference names as keys and Reference as values. 77 | Note that all references must be specified using this argument 78 | on object instantiation as they must be added as variables to 79 | the model. 80 | model_callback : callable, default None 81 | Callback function for the end of model creation taking the model 82 | as argument, and returns the modified model. Can be used to add 83 | additional variables to the model. 84 | p_callback : callable, default None 85 | Callback function calles before computing the MPC solution 86 | (get_input) where values for additional parameters must be set 87 | by user. Takes as argument the parameter object and the current 88 | model time, and must return the modified parameter object. 89 | tvp_callback : callable, default None 90 | Callback function called before computing the MPC solution 91 | (get_input) where values for additional time-varying parameter 92 | must be set by user. Takes as argument the 93 | time-varying-parameter object and the current model time, and 94 | must return the modified time-varying-parameter object object. 95 | 96 | """ 97 | 98 | def __init__(self, control_forces_filter, S=None, dH=None, H=None, F=None, 99 | R=None, baseline=None, state_names=None, 100 | control_names=None, references=None, 101 | model_callback=None, p_callback=None, tvp_callback=None): 102 | self.baseline = baseline 103 | 104 | if baseline is None: 105 | assert S is not None 106 | self.S = S 107 | 108 | if dH is None: 109 | assert H is not None, 'Must provide either dH or H' 110 | self.H = H 111 | self.dH = None 112 | else: 113 | assert H is None, 'Please only provide one of dH and H' 114 | self.H = None 115 | self.dH = dH 116 | 117 | if F is None: 118 | self.F = lambda x: np.zeros_like(x) 119 | else: 120 | self.F = F 121 | 122 | if R is None: 123 | self.R = lambda x: np.zeros_like(x) 124 | else: 125 | self.R = R 126 | else: 127 | pass 128 | 129 | self.nstates = control_forces_filter.shape[0] 130 | super().__init__(control_forces_filter) 131 | 132 | self.references = references 133 | 134 | self.p_callback = p_callback 135 | self.tvp_callback = tvp_callback 136 | self._p_template = None 137 | self._tvp_template = None 138 | self._use_time_variable = False 139 | 140 | self.model_callback = model_callback 141 | 142 | self.has_been_reset = False 143 | 144 | self.dynamics_params = None 145 | 146 | model = self._get_model(state_names=state_names, control_names=control_names) 147 | self.mpc = do_mpc.controller.MPC(model) 148 | 149 | def _get_model(self, state_names=None, control_names=None): 150 | def pytorch_parameter_getter(module): 151 | return lambda: np.concatenate([p.flatten() for p in get_pytorch_model_parameters(module)]).reshape(-1, 1) 152 | 153 | model = do_mpc.model.Model('continuous') 154 | if state_names is None: 155 | state_names = [f'x{x_i}' for x_i in range(self.nstates)] 156 | else: 157 | assert len(state_names) == self.nstates 158 | 159 | if control_names is None: 160 | control_names = [f'u{u_i}' for u_i in range(self.ncontrols)] 161 | else: 162 | assert len(control_names) == self.ncontrols 163 | 164 | for x_i in range(self.nstates): 165 | _ = model.set_variable(var_type='_x', var_name=state_names[x_i], shape=(1, 1)) 166 | 167 | for u_i in range(self.ncontrols): 168 | _ = model.set_variable(var_type='_u', var_name=control_names[u_i], shape=(1, 1)) 169 | 170 | if self.references is not None: 171 | for ref_name in self.references.keys(): 172 | _ = model.set_variable(var_type='_tvp', var_name=ref_name, shape=(1, 1)) 173 | 174 | self.dynamics_params = {} 175 | u = self.control_forces_filter @ model.u.cat 176 | 177 | if self.baseline is None: 178 | if self.H is not None: 179 | if isinstance(self.H, torch.nn.Module): 180 | nn_H = CasadiFCNN(layers=get_pytorch_model_architecture(self.H)) 181 | dH = casadi.gradient(nn_H.create_forward(model.x.cat), model.x.cat) 182 | 183 | model._p["name"].append("H_params") 184 | model._p["var"].append(nn_H.params) 185 | 186 | self.dynamics_params['H'] = pytorch_parameter_getter(self.H) 187 | else: 188 | dH = casadi.gradient(self.H(model.x.cat), model.x.cat) 189 | else: 190 | if isinstance(self.dH, torch.nn.Module): 191 | nn_dH = CasadiFCNN(layers=get_pytorch_model_architecture(self.dH)) 192 | dH = nn_dH.create_forward(model.x.cat) 193 | 194 | model._p["name"].append("H_params") 195 | model._p["var"].append(nn_dH.params) 196 | 197 | self.dynamics_params['H'] = pytorch_parameter_getter(self.dH) 198 | else: 199 | dH = self.dH(model.x.cat) 200 | 201 | if isinstance(self.F, torch.nn.Module): 202 | nn_F = CasadiFCNN(layers=get_pytorch_model_architecture(self.F)) 203 | F_inputs = [] 204 | if self.F.statedependent: 205 | F_inputs.append(model.x) 206 | if self.F.timedependent: 207 | t = model.set_variable(var_type="_tvp", var_name="time", shape=(1, 1)) 208 | F_inputs.append(t) 209 | self._use_time_variable = True 210 | F = (self.F.external_forces_filter.detach().numpy() @ nn_F.create_forward(*F_inputs)).T 211 | 212 | model._p["name"].append("F_params") 213 | model._p["var"].append(nn_F.params) 214 | self.dynamics_params['F'] = pytorch_parameter_getter(self.F) 215 | else: 216 | F = self.F(model.x.cat) 217 | 218 | if isinstance(self.R, R_estimator): 219 | R_weights = casadi.SX.sym('R_weights', self.R.rs.shape[0]) 220 | R = casadi.diag(self.R.pick_rs.detach().numpy() @ R_weights) 221 | 222 | model._p["name"].append("R_params") 223 | model._p["var"].append(R_weights) 224 | self.dynamics_params['R'] = lambda: self.R.get_parameters() 225 | elif isinstance(self.R, torch.nn.Module): 226 | nn_R = CasadiFCNN(layers=get_pytorch_model_architecture(self.R)) 227 | R = nn_R.create_forward(model.x.cat) 228 | 229 | model._p["name"].append("R_params") 230 | model._p["var"].append(nn_R.params) 231 | self.dynamics_params['R'] = pytorch_parameter_getter(self.R) 232 | elif callable(self.R): 233 | R = self.R(model.x.cat) 234 | else: 235 | R = self.R 236 | 237 | dynamics = CasadiPseudoHamiltonianSystem(S=self.S, dH=dH, u=u, R=R, F=F) 238 | rhs = dynamics.create_forward() 239 | else: 240 | nn_baseline = CasadiFCNN(layers=get_pytorch_model_architecture(self.baseline)) 241 | nn_baseline.set_weights_and_biases(get_pytorch_model_parameters(self.baseline)) 242 | 243 | baseline_inputs = [] 244 | if self.baseline.statedependent: 245 | baseline_inputs.append(model.x) 246 | if self.baseline.timedependent: 247 | t = model.set_variable(var_type="_tvp", var_name="time", shape=(1, 1)) 248 | baseline_inputs.append(t) 249 | self._use_time_variable = True 250 | 251 | rhs = nn_baseline.create_forward(*baseline_inputs).T + u 252 | 253 | model._p["name"].append("baseline_params") 254 | model._p["var"].append(nn_baseline.params) 255 | 256 | self.dynamics_params['baseline'] = pytorch_parameter_getter(self.baseline) 257 | 258 | for state_i, state_name in enumerate(state_names): 259 | model.set_rhs(state_name, rhs[state_i]) 260 | 261 | if self.model_callback is not None: 262 | model = self.model_callback(model) 263 | 264 | model.setup() 265 | 266 | return model 267 | 268 | def setup(self, setup_callback): 269 | """ 270 | Function to finalize MPC creation. Must be called prior to use 271 | of the MPC for getting control inputs. Note that the objective 272 | must be set by the user. Other MPC options such as constraints 273 | and optimization horizon can also 274 | be configured here through the setup_callback argument. 275 | 276 | parameters 277 | ---------- 278 | setup_callback : callable 279 | Function for user to finalize the MPC configuration. Takes 280 | as argument the MPC object, and returns the modified mpc 281 | object. Note that set_objective must be called by user on 282 | the mpc object. 283 | 284 | """ 285 | 286 | default_settings = { 287 | 'n_horizon': 10, 288 | 't_step': 0.01, 289 | 'n_robust': 0, 290 | 'store_full_solution': True, 291 | 'nlpsol_opts': {'ipopt.max_iter': 200, 292 | 'ipopt.print_level': 0, 'ipopt.sb': 'yes', 'print_time': 0 293 | }, 294 | } 295 | self.mpc.set_param(**default_settings) 296 | 297 | self.mpc = setup_callback(self.mpc) 298 | 299 | if self.dynamics_params is not None or self.p_callback is not None: 300 | self._p_template = self.mpc.get_p_template(1) 301 | 302 | def mpc_p_fun(t_now): 303 | if self.p_callback is not None: 304 | self._p_template = self.p_callback(self._p_template, t_now) 305 | 306 | if self.dynamics_params is not None: 307 | for p_name, p_getter in self.dynamics_params.items(): 308 | self._p_template['_p', :, p_name + '_params'] = p_getter() 309 | 310 | return self._p_template 311 | 312 | self.mpc.set_p_fun(mpc_p_fun) 313 | 314 | if self.references is not None or self.tvp_callback is not None: 315 | self._tvp_template = self.mpc.get_tvp_template() 316 | 317 | def mpc_tvp_fun(t_now): 318 | if self.tvp_callback is not None: 319 | self._tvp_template = self.tvp_callback(self._tvp_template, t_now) 320 | 321 | if self.references is not None: 322 | for t_i in range(self.mpc.n_horizon + 1): 323 | for name, fun in self.references.items(): 324 | self._tvp_template['_tvp', t_i, name] = fun(float(t_now) + t_i * self.mpc.t_step) 325 | if self._use_time_variable: 326 | self._tvp_template['_tvp', :, "time"] = (self.mpc.t0 + np.arange(self.mpc.n_horizon + 1) * self.mpc.t_step).tolist() 327 | 328 | return self._tvp_template 329 | 330 | self.mpc.set_tvp_fun(mpc_tvp_fun) 331 | self.mpc.setup() 332 | 333 | def reset(self): 334 | """ 335 | Function called before starting control of a new trajectory. 336 | Resets the MPC state, and the reference object. 337 | """ 338 | self.mpc.reset_history() 339 | self.has_been_reset = True 340 | if self.references is not None: 341 | for ref_fun in self.references.values(): 342 | ref_fun.reset() 343 | 344 | def set_reference(self, references): 345 | """ 346 | Function used to change the Reference objects specified during 347 | instantiation. Note that only existing references 348 | can be updated using this function. 349 | 350 | Parameters 351 | ---------- 352 | references : dict 353 | Dictionary of reference name as keys and Reference object 354 | as values. 355 | 356 | """ 357 | 358 | assert isinstance(references, dict), 'The PseudoHamiltonianMPC class only supports references in dictionary format' 359 | assert all(name in self.references for name in references), 'New references can not be added after the MPC is created.' 360 | for name, ref_fun in references.items(): 361 | self.references[name] = ref_fun 362 | 363 | def _get_input(self, x, t=None): 364 | """ 365 | Function used to compute the MPC solution and obtain the 366 | corresponding control input. 367 | 368 | Parameters 369 | ---------- 370 | x : array 371 | Initial state of the MPC optimal control problem. 372 | t : number 373 | System time for the optimal control problem. Affects 374 | parameters and time-varying parameters such as 375 | references. 376 | 377 | """ 378 | 379 | if self.has_been_reset: 380 | self.mpc.x0 = x 381 | self.mpc.u0 = np.zeros((self.ncontrols,)) 382 | self.mpc.set_initial_guess() 383 | self.has_been_reset = False 384 | if t is not None: 385 | self.mpc.t0 = t 386 | return self.mpc.make_step(x) 387 | -------------------------------------------------------------------------------- /phlearn/phlearn/control/phcontroller.py: -------------------------------------------------------------------------------- 1 | 2 | import torch 3 | import numpy as np 4 | 5 | __all__ = ['PseudoHamiltonianController'] 6 | 7 | class PseudoHamiltonianController: 8 | """ 9 | Abstract base class for controllers of pseudo-Hamiltonian systems 10 | of the form:: 11 | 12 | dx/dt = (S - R)*grad[H(x)] + F(x, t) + u(x, t) 13 | 14 | where this class implements u(x, t), i.e. known and controlled 15 | external ports. Implementations of controllers 16 | must subclass this class and implement all of its methods. 17 | 18 | parameters 19 | ---------- 20 | control_forces_filter : (nstates, nstates) or (nstates,) ndarray 21 | A binary ndarray where 1 signifies that the corresponding 22 | state has a control input in its derivative right-hand-side. 23 | 24 | """ 25 | def __init__(self, control_forces_filter): 26 | self.ncontrols = int(np.sum(control_forces_filter)) 27 | self.control_forces_filter = self._format_control_forces_filter(control_forces_filter) 28 | 29 | def __call__(self, x, t=None): 30 | """ 31 | Control inputs are computed by calling the controller with a 32 | system state, and optionally with system time for controllers 33 | that depend on time (e.g. through time-varying references). 34 | The returned control input will have the 35 | same shape as the system state. 36 | 37 | parameters 38 | ---------- 39 | x : (nstates,) ndarray 40 | System state 41 | t : number, default None 42 | System time 43 | 44 | Returns 45 | ------- 46 | (nstates,) ndarray 47 | """ 48 | if torch.is_tensor(x): 49 | x = x.detach().numpy() 50 | if torch.is_tensor(t): 51 | t = t.detach().numpy() 52 | # TODO: what about batch operation? 53 | return (self.control_forces_filter @ self._get_input(x, t)).ravel() 54 | 55 | def _get_input(self, x, t=None): 56 | raise NotImplementedError 57 | 58 | def reset(self): 59 | """ 60 | Function called before starting control of a new trajectory. 61 | Resets the controllers' internal state, as well as 62 | any reference objects. 63 | """ 64 | raise NotImplementedError 65 | 66 | def set_reference(self, references): 67 | """ 68 | Function used to change the Reference objects specified during 69 | instantiation. 70 | 71 | Parameters 72 | ---------- 73 | references :dict of references 74 | Dictionary of reference name as keys and Reference 75 | object as values. 76 | 77 | """ 78 | raise NotImplementedError 79 | 80 | def _format_control_forces_filter(self, control_forces_filter): 81 | control_forces_filter = (np.array(control_forces_filter) > 0).astype(int) 82 | 83 | if (len(control_forces_filter.shape) == 1) or (len(control_forces_filter.shape[-1]) == 1): 84 | control_forces_filter = control_forces_filter.flatten() 85 | expanded = np.zeros((control_forces_filter.shape[-1], control_forces_filter.sum())) 86 | c = 0 87 | for i, e in enumerate(control_forces_filter): 88 | if e > 0: 89 | expanded[i, c] = 1 90 | c += 1 91 | return expanded 92 | 93 | return control_forces_filter 94 | -------------------------------------------------------------------------------- /phlearn/phlearn/control/pid.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from .phcontroller import PseudoHamiltonianController 3 | 4 | __all__ = ['PIDController'] 5 | 6 | class PIDController(PseudoHamiltonianController): 7 | """ 8 | This class implements a proportional-integral-derivative (PID) 9 | controller. The PID controller is a SISO controller, 10 | but several PIDs can be configured to run in parallel, 11 | up to one per state of the system. 12 | 13 | parameters 14 | ---------- 15 | control_forces_filter : (nstates, nstates) or (nstates,) ndarray 16 | A binary ndarray where 1 signifies that the corresponding 17 | state has a control input in its derivative right-hand-side. 18 | gains : dict 19 | Dictionary where the key is the index of the state that the 20 | PID controls, and the value is anotherdictionary of the form 21 | {"p": proportional gain, "i": integral gain", "d":derivative 22 | gain}. Missing gains are assumed to be zero. 23 | references : dict of references 24 | Dictionary where the key is the index of the state that the PID 25 | controls, and the value is a Reference object for that PID 26 | controller. 27 | input_bounds : dict 28 | Dictionary where the key is the index of the state that the PID 29 | controls, and the value is a two-element list where the first 30 | element gives a lower bound for the control input and the second 31 | element gives an upper bound. 32 | """ 33 | def __init__(self, control_forces_filter, gains, references, input_bounds=None): 34 | self.gains = {} 35 | for idx, idx_gains in gains.items(): 36 | self.gains[idx] = {} 37 | for gain in ['p', 'i', 'd']: 38 | self.gains[idx][gain] = idx_gains.get(gain, 0.0) 39 | 40 | super().__init__(control_forces_filter=control_forces_filter) 41 | 42 | self.state_idxs = [np.argwhere(self.control_forces_filter[:, i]).item() for i in range(self.ncontrols)] 43 | 44 | assert input_bounds is None or len(input_bounds.keys()) == self.ncontrols 45 | self.input_bounds = input_bounds 46 | 47 | assert len(references.keys()) == self.ncontrols 48 | self.references = references 49 | 50 | self.integrator = [0] * self.ncontrols 51 | self.prev_state = [None] * self.ncontrols 52 | 53 | self.last_t = 0 54 | self.history = {'t': [], 'input': []} 55 | 56 | def set_reference(self, references): 57 | assert isinstance(references, dict) 58 | for idx, ref_fun in references.items(): 59 | self.references[idx] = ref_fun 60 | 61 | def reset(self): 62 | self.integrator = [0] * self.ncontrols 63 | self.prev_state = [None] * self.ncontrols 64 | self.history = {'t': [], 'input': []} 65 | self.last_t = 0 66 | for ref_fun in self.references.values(): 67 | ref_fun.reset() 68 | 69 | def _get_input(self, state, t=None): 70 | us = [] 71 | if t >= self.last_t: # some solvers are not monotonic wrt. time 72 | step_dt = t - self.last_t 73 | 74 | for i, idx in enumerate(self.state_idxs): 75 | state = np.atleast_2d(state) 76 | error = self.references[idx](t) - state[..., idx].ravel() 77 | 78 | self.integrator[i] += step_dt * error 79 | 80 | u = self.gains[idx]['p'] * error + self.gains[idx]['i'] * self.integrator[i] 81 | if self.prev_state[i] is not None and self.gains[idx]['d'] > 0: 82 | u += self.gains[idx]['d'] * (state[..., idx] - self.prev_state[i]) / (step_dt + 1e-12) 83 | self.prev_state[i] = state[..., idx] 84 | 85 | # Constrain input 86 | if self.input_bounds is not None and idx in self.input_bounds: 87 | u = np.clip(u, *self.input_bounds[idx]) 88 | 89 | self.last_t = t 90 | us.append(np.squeeze(u)) 91 | self.history['t'].append(t) 92 | self.history['input'].append(np.array(us)) 93 | else: # if earlier in time, interpolate between closest previously computed inputs 94 | closest_index = np.abs(np.array(self.history['t']) - t).argmin() 95 | next_closest_index = (closest_index - 1) if self.history['t'][closest_index] - t > 0 else (closest_index + 1) 96 | t_span = np.abs(self.history['t'][closest_index] - self.history['t'][next_closest_index]) 97 | while t_span == 0: 98 | next_closest_index += 1 * (np.sign(next_closest_index - closest_index)) 99 | t_span = np.abs(self.history['t'][closest_index] - self.history['t'][next_closest_index]) 100 | interpol_factor = np.abs((self.history['t'][closest_index] - t) / t_span) 101 | assert interpol_factor <= 1, 'Something was wrong with interpolation logic' 102 | us = self.history['input'][closest_index] * (1 - interpol_factor) + self.history['input'][next_closest_index] * interpol_factor 103 | 104 | return us 105 | -------------------------------------------------------------------------------- /phlearn/phlearn/control/reference.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | class Reference: 4 | def __init__(self, seed=None): 5 | self.history = {'t': [], 'r': []} 6 | self.rng = None 7 | self.seed(seed) 8 | self.reset() 9 | 10 | def __call__(self, t): 11 | value = self._get_value(t) 12 | self.history['t'].append(t) 13 | self.history['r'].append(value) 14 | 15 | return value 16 | 17 | def _get_value(self, t): 18 | raise NotImplementedError 19 | 20 | def reset(self): 21 | self.history = {'t': [], 'r': []} 22 | 23 | def get_reference_data(self, ts=None): 24 | if ts is None: 25 | return self.history['r'], self.history['t'] 26 | else: 27 | return [self._get_value(t) for t in ts], ts 28 | 29 | def seed(self, seed): 30 | self.rng = np.random.default_rng(seed) 31 | 32 | 33 | class ConstantReference(Reference): 34 | def __init__(self, low, high, value=None, seed=None): 35 | self.low = low 36 | self.high = high 37 | self.value = value 38 | super().__init__(seed) 39 | 40 | def reset(self, value=None): 41 | super().reset() 42 | if value is not None: 43 | self.value = value 44 | else: 45 | self.value = self.rng.uniform(self.low, self.high) 46 | 47 | def _get_value(self, t): 48 | return self.value 49 | 50 | 51 | class StepReference(Reference): 52 | def __init__(self, low, high, step_interval, seed=None): 53 | self.step_interval = step_interval 54 | self.low = low 55 | self.high = high 56 | self.values = None 57 | super().__init__(seed) 58 | 59 | def _get_value(self, t): 60 | t_step_idx = int(t // self.step_interval) 61 | if t_step_idx >= len(self.values): 62 | self.values.extend([self.rng.uniform(self.low, self.high) for _ in range(t_step_idx + 1 - len(self.values))]) 63 | 64 | return self.values[t_step_idx] 65 | 66 | def reset(self): 67 | super().reset() 68 | self.values = [self.rng.uniform(self.low, self.high)] 69 | 70 | 71 | class PoissonStepReference(Reference): 72 | def __init__(self, low, high, rate, seed=None): 73 | self.rate = rate 74 | self.low = low 75 | self.high = high 76 | self.values = None 77 | super().__init__(seed) 78 | 79 | def _get_value(self, t): 80 | pass 81 | 82 | 83 | class FixedReference(Reference): 84 | def __init__(self, values, timestamps, seed=None): 85 | self.timestamps = np.array(timestamps) 86 | self.values = np.array(values) 87 | super().__init__(seed) 88 | 89 | def _get_value(self, t): 90 | closest_index = np.abs(self.timestamps - t).argmin() 91 | return self.values[closest_index] 92 | -------------------------------------------------------------------------------- /phlearn/phlearn/phnns/__init__.py: -------------------------------------------------------------------------------- 1 | """ 2 | The phnns subpackage implements neural networks and functionality for 3 | learning dynamic systems and pseudo-Hamiltonian systems. 4 | 5 | Classes present in phlearn.phnns 6 | ----------------------------------------------- 7 | 8 | :py:class:`~.dynamic_system_neural_network.DynamicSystemNN` 9 | 10 | :py:class:`~.pseudo_hamiltonian_neural_network.PseudoHamiltonianNN` 11 | 12 | :py:class:`~.pseudo_hamiltonian_pde_neural_network.PseudoHamiltonianPDENN` 13 | 14 | :py:class:`~.ode_models.BaseNN` 15 | 16 | :py:class:`~.ode_models.BaselineNN` 17 | 18 | :py:class:`~.ode_models.HamiltonianNN` 19 | 20 | :py:class:`~.ode_models.ExternalForcesNN` 21 | 22 | :py:class:`~.ode_models.BaselineSplitNN` 23 | 24 | :py:class:`~.ode_models.R_NN` 25 | 26 | :py:class:`~.ode_models.R_estimator` 27 | 28 | :py:class:`~.pde_models.CentralPadding` 29 | 30 | :py:class:`~.pde_models.ForwardPadding` 31 | 32 | :py:class:`~.pde_models.Summation` 33 | 34 | :py:class:`~.pde_models.PDEBaseNN` 35 | 36 | :py:class:`~.pde_models.PDEBaseLineNN` 37 | 38 | :py:class:`~.pde_models.PDEIntegralNN` 39 | 40 | :py:class:`~.pde_models.PDEExternalForcesNN` 41 | 42 | :py:class:`~.pde_models.PDEBaselineSplitNN` 43 | 44 | :py:class:`~.pde_models.A_estimator` 45 | 46 | :py:class:`~.pde_models.S_estimator` 47 | 48 | :py:class:`~.train_utils.EarlyStopping` 49 | 50 | 51 | 52 | Functions present in phlearn.phnns 53 | ----------------------------------------------- 54 | 55 | :py:func:`~.pseudo_hamiltonian_neural_network.load_phnn_model` 56 | 57 | :py:func:`~.pseudo_hamiltonian_neural_network.store_phnn_model` 58 | 59 | :py:func:`~.pseudo_hamiltonian_pde_neural_network.load_cdnn_model` 60 | 61 | :py:func:`~.pseudo_hamiltonian_pde_neural_network.store_cdnn_model` 62 | 63 | :py:func:`~.dynamic_system_neural_network.load_baseline_model` 64 | 65 | :py:func:`~.dynamic_system_neural_network.store_baseline_model` 66 | 67 | :py:func:`~.train_utils.generate_dataset` 68 | 69 | :py:func:`~.train_utils.train` 70 | 71 | :py:func:`~.train_utils.compute_validation_loss` 72 | 73 | :py:func:`~.train_utils.batch_data` 74 | 75 | :py:func:`~.train_utils.train_one_epoch` 76 | 77 | :py:func:`~.train_utils.l1_loss_pHnn` 78 | 79 | :py:func:`~.train_utils.npoints_to_ntrajectories_tsample` 80 | 81 | :py:func:`~.train_utils.load_dynamic_system_model` 82 | 83 | :py:func:`~.train_utils.store_dynamic_system_model` 84 | 85 | """ 86 | 87 | from .dynamic_system_neural_network import * 88 | from .pseudo_hamiltonian_neural_network import * 89 | from .pseudo_hamiltonian_pde_neural_network import * 90 | from .ode_models import * 91 | from .pde_models import * 92 | from .train_utils import * 93 | 94 | __all__ = dynamic_system_neural_network.__all__.copy() 95 | __all__ += pseudo_hamiltonian_neural_network.__all__.copy() 96 | __all__ += pseudo_hamiltonian_pde_neural_network.__all__.copy() 97 | __all__ += ode_models.__all__.copy() 98 | __all__ += pde_models.__all__.copy() 99 | __all__ += train_utils.__all__.copy() -------------------------------------------------------------------------------- /phlearn/phlearn/phnns/ode_models.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import numpy as np 4 | 5 | __all__ = ['BaseNN', 'BaselineNN', 'BaselineSplitNN', 'HamiltonianNN', 6 | 'ExternalForcesNN', 'R_NN', 'R_estimator'] 7 | 8 | class BaseNN(torch.nn.Module): 9 | """ 10 | Neural network with three hidden layers, where the first has 11 | Tanh-activation, the second has ReLU-activation and the third has 12 | linear activation. The network can take either system states or 13 | time or both as input. If it is expected to take neither states nor 14 | time as input, the network is replaced by trainable parameters. 15 | Independently of whether the network uses state and/or time or neither, 16 | it can be called with both state and time:: 17 | 18 | pred = network(x=x, t=t) 19 | 20 | 21 | Parameters 22 | ---------- 23 | nstates : int 24 | Number of states in a potential state input. 25 | noutputs : int 26 | Number of outputs from the last linear layer. 27 | hidden_dim : int 28 | Dimension of hidden layers. 29 | timedependent : bool 30 | If True, time input is expected. 31 | statedependent : bool 32 | If True, state input is expected. 33 | 34 | """ 35 | 36 | def __init__(self, nstates, noutputs, hidden_dim, 37 | timedependent, statedependent, ttype=torch.float32): 38 | super().__init__() 39 | self.nstates = nstates 40 | self.noutputs = noutputs 41 | self.hidden_dim = hidden_dim 42 | self.timedependent = timedependent 43 | self.statedependent = statedependent 44 | if not statedependent and not timedependent: 45 | self.model = nn.Parameter(torch.zeros(noutputs, dtype=ttype)) 46 | else: 47 | input_dim = int(statedependent)*nstates + int(timedependent) 48 | linear1 = nn.Linear(input_dim, hidden_dim) 49 | linear2 = nn.Linear(hidden_dim, hidden_dim) 50 | linear3 = nn.Linear(hidden_dim, noutputs) 51 | 52 | for lin in [linear1, linear2, linear3]: 53 | nn.init.orthogonal_(lin.weight) 54 | 55 | self.model = nn.Sequential( 56 | linear1, 57 | nn.Tanh(), 58 | linear2, 59 | nn.ReLU(), 60 | linear3, 61 | ) 62 | 63 | if timedependent and not statedependent: 64 | self.forward = self._forward_without_state 65 | elif statedependent and not timedependent: 66 | self.forward = self._forward_without_time 67 | elif not statedependent and not timedependent: 68 | self.forward = self._forward_without_state_or_time 69 | else: 70 | self.forward = self._forward_with_state_and_time 71 | 72 | def _forward_with_state_and_time(self, x=None, t=None): 73 | return self.model(torch.cat([x, t], dim=-1)) 74 | 75 | def _forward_without_time(self, x=None, t=None): 76 | return self.model(x) 77 | 78 | def _forward_without_state(self, x=None, t=None): 79 | return self.model(t) 80 | 81 | def _forward_without_state_or_time(self, x=None, t=None): 82 | return self.model 83 | 84 | 85 | class BaselineNN(BaseNN): 86 | """ 87 | Neural network for estimating the right hand side of a set of 88 | dynamic system equations with three hidden layers, where the first 89 | has Tanh-activation, the second has ReLU-activation and the third has 90 | linear activation. The network can take either system states or 91 | time or both as input. If it is expected to take neither states nor 92 | time as input, the network is replaced by trainable parameters. 93 | Independently of whether the network uses state and/or time or neither, 94 | it can be called with both state and time:: 95 | 96 | pred = network(x=x, t=t) 97 | 98 | The output dimension of the network is always *nstates*. 99 | 100 | Parameters 101 | ---------- 102 | nstates : int 103 | Number of states in a potential state input. 104 | hidden_dim : int 105 | Dimension of hidden layers. 106 | timedependent : bool 107 | If True, time input is expected. 108 | statedependent : bool 109 | If True, state input is expected. 110 | 111 | """ 112 | def __init__(self, nstates, hidden_dim, timedependent=True, 113 | statedependent=True): 114 | super().__init__(nstates, nstates, hidden_dim, 115 | timedependent, statedependent) 116 | 117 | 118 | class HamiltonianNN(BaseNN): 119 | """ 120 | Neural network for estimating a scalar function H(x), 121 | with three hidden layers, where the first has 122 | Tanh-activation, the second has ReLU-activation and the third has 123 | linear activation. The network takes system states as input, but 124 | can be called with both state and time:: 125 | 126 | pred = network(x=x, t=t) 127 | 128 | The output dimension of the network is always 1. 129 | 130 | Parameters 131 | ---------- 132 | nstates : int 133 | Number of states in a potential state input. 134 | hidden_dim : int 135 | Dimension of hidden layers. 136 | 137 | """ 138 | def __init__(self, nstates, hidden_dim=100): 139 | super().__init__(nstates, 1, hidden_dim, False, True) 140 | 141 | 142 | class ExternalForcesNN(BaseNN): 143 | """ 144 | Neural network for estimating external forces of a pseudo-Hamiltonian 145 | system with three hidden layers, where the first has 146 | Tanh-activation, the second has ReLU-activation and the third has 147 | linear activation. The network can take either system states or 148 | time or both as input. Independently of whether the network uses 149 | state and/or time, it can be called with both state and time:: 150 | 151 | pred = network(x=x, t=t) 152 | 153 | If neither time or state input is to be expected, the neural 154 | network is replaced by trainable parameters. 155 | 156 | Parameters 157 | ---------- 158 | nstates : int 159 | Number of states in a potential state input. 160 | noutputs : int 161 | Number of external forces to estimate. 162 | timedependent : bool 163 | If True, time input is expected. 164 | statedependent : bool 165 | If True, state input is expected. 166 | external_forces_filter : listlike of ints or None, default None 167 | If None, *noutputs* == *nstates* must be true. In this case, 168 | one external force is estimated for each state. If 169 | *noutputs* != *nstates*, *external_forces_filter* must decribe 170 | which states external forces should be estimated for. Either, 171 | *external_forces_filter* must be a 1d liststructure of length 172 | nstates filled with 0 and 1, where 1 indicates that an external 173 | force should be estimated for state corresponding to that index. 174 | Alternatively, *external_forces_filter* can be an array of shape 175 | (nstates, noutputs) of 0s and 1s, such that when it is 176 | multiplied with the network outout of shape (noutputs,), the 177 | right output is applied to the correct state. 178 | ttype : torch type, default torch.float32 179 | 180 | """ 181 | 182 | def __init__(self, nstates, noutputs, hidden_dim, timedependent, 183 | statedependent, external_forces_filter=None, 184 | ttype=torch.float32): 185 | super().__init__(nstates, noutputs, hidden_dim, 186 | timedependent, statedependent) 187 | self.nstates = nstates 188 | self.noutputs = noutputs 189 | self.ttype = ttype 190 | 191 | self.external_forces_filter = self._format_external_forces_filter( 192 | external_forces_filter) 193 | 194 | def _forward_with_state_and_time(self, x=None, t=None): 195 | return self.model(torch.cat([x, t], dim=-1))@self.external_forces_filter 196 | 197 | def _forward_without_time(self, x=None, t=None): 198 | return self.model(x)@self.external_forces_filter 199 | 200 | def _forward_without_state(self, x=None, t=None): 201 | return self.model(t)@self.external_forces_filter 202 | 203 | def _forward_without_state_or_time(self, x=None, t=None): 204 | return self.model@self.external_forces_filter 205 | 206 | def _format_external_forces_filter(self, external_forces_filter): 207 | if external_forces_filter is None: 208 | assert self.noutputs == self.nstates, ( 209 | f'noutputs ({self.noutputs}) != nstates ({self.nstates}) is ' 210 | 'not allowed when external_forces_filter is not provided.') 211 | return torch.eye(self.noutputs, dtype=self.ttype) 212 | 213 | if isinstance(external_forces_filter, (list, tuple)): 214 | external_forces_filter = np.array(external_forces_filter) 215 | 216 | if not isinstance(external_forces_filter, torch.Tensor): 217 | external_forces_filter = torch.tensor(external_forces_filter > 0) 218 | external_forces_filter = external_forces_filter.int() 219 | 220 | if ((len(external_forces_filter.shape) == 1) or 221 | (external_forces_filter.shape[-1] == 1)): 222 | 223 | external_forces_filter = external_forces_filter.flatten() 224 | assert external_forces_filter.shape[-1] == self.nstates, ( 225 | 'external_forces_filter is a vector of ' 226 | f'length {external_forces_filter.shape[-1]} != nstates, but ' 227 | f'({self.nstates}). external_forces_filter must be a ' 228 | 'vector of length nstates or a matrix of shape' 229 | '(nstates x noutputs).') 230 | expanded = torch.zeros((self.nstates, external_forces_filter.sum()), 231 | dtype=self.ttype) 232 | c = 0 233 | for i, e in enumerate(external_forces_filter): 234 | if e > 0: 235 | expanded[i, c] = 1 236 | c += 1 237 | return expanded.T 238 | 239 | assert external_forces_filter.shape == (self.nstates, self.noutputs), ( 240 | f'external_forces_filter.shape == {external_forces_filter.shape}, but ' 241 | 'external_forces_filter must be a vector of length nstates or ' 242 | 'a matrix of shape (naffected_states x noutputs).') 243 | return external_forces_filter.clone().type(self.ttype).detach().T 244 | 245 | 246 | class BaselineSplitNN(torch.nn.Module): 247 | """ 248 | Composition of two neural networks for estimating the right hand 249 | side of a set of dynamic system equations. The networks have 250 | three hidden layers, where the first has Tanh-activation, the 251 | second has ReLU-activation and the third has linear activation. 252 | One network takes system states and the other takes time as input. 253 | The output of the composition is the sum of the outputs of the 254 | two networks:: 255 | 256 | pred = network(x, t) = network_x(x) + network_t(t) 257 | 258 | Both networks are instantiated from the 259 | :py:class:`~.models.ExternalForcesNN` class, allowing adjustment 260 | of the number and location of non-zero contributions from each 261 | network. 262 | The output dimension of a BaselineSplitNN is always *nstates*. 263 | 264 | Parameters 265 | ---------- 266 | nstates : int 267 | Number of states in a potential state input. 268 | hidden_dim : int 269 | Dimension of hidden layers. Equal for both networks. 270 | noutputs_x : int or None, default None 271 | Number of non-zero outputs to estimate with network_x(x) 272 | noutputs_t : int or None, default None 273 | Number of non-zero outputs to estimate with network_t(t) 274 | external_forces_filter_x : listlike of ints or None, default None 275 | If provided, this decides to which states the output of 276 | network_x is contributing. See :py:class:`~.models.ExternalForcesNN` 277 | for fruther description. 278 | external_forces_filter_t : listlike of ints or None, default None 279 | If provided, this decides to which states the output of 280 | network_t is contributing. See :py:class:`~.models.ExternalForcesNN` 281 | for fruther description. 282 | ttype : torch type, default torch.float32 283 | 284 | """ 285 | def __init__(self, nstates, hidden_dim, noutputs_x=None, 286 | noutputs_t=None, external_forces_filter_x=None, 287 | external_forces_filter_t=None, 288 | ttype=torch.float32): 289 | super().__init__() 290 | self.nstates = nstates 291 | self.hidden_dim = hidden_dim 292 | self.noutputs_x = nstates if noutputs_x is None else noutputs_x 293 | self.noutputs_t = nstates if noutputs_t is None else noutputs_t 294 | self.network_x = ExternalForcesNN( 295 | nstates, self.noutputs_x, hidden_dim, False, 296 | True, external_forces_filter_x, ttype) 297 | self.network_t = ExternalForcesNN( 298 | nstates, self.noutputs_t, hidden_dim, True, 299 | False, external_forces_filter_t, ttype) 300 | 301 | def forward(self, x, t): 302 | return self.network_x(x, t) + self.network_t(x, t) 303 | 304 | 305 | class R_NN(BaseNN): 306 | ''' 307 | Neural network for estimating the parameters of a damping matrix. 308 | with three hidden layers, where the first has Tanh-activation, the 309 | second has ReLU-activation and the third has linear activation. 310 | The network takes system states as input. 311 | 312 | When called with a batch input, the network returns a batch of 313 | matrices of size (*nstates*, *nstates*). All damping parameters are 314 | assumed to be positive. 315 | 316 | Parameters 317 | ---------- 318 | nstates : int 319 | Number of states in a potential state input. 320 | hidden_dim : int 321 | Dimension of hidden layers. 322 | diagonal : bool 323 | If True, only damping coefficients on the diagonal 324 | are estimated. If False, all nstates**2 entries in the 325 | R matrix are estimated. 326 | 327 | ''' 328 | 329 | def __init__(self, nstates, hidden_dim, diagonal=False): 330 | if diagonal: 331 | noutputs = nstates 332 | self.forward = self._forward_diag 333 | else: 334 | noutputs = nstates**2 335 | self.forward = self._forward 336 | super().__init__(nstates, noutputs, hidden_dim, False, True) 337 | 338 | self.nstates = nstates 339 | 340 | def _forward_diag(self, x): 341 | return torch.diag_embed(self.model(x)**2).reshape( 342 | x.shape[0], self.nstates, self.nstates) 343 | 344 | def _forward(self, x): 345 | return (self.model(x)**2).reshape( 346 | x.shape[0], self.nstates, self.nstates) 347 | 348 | 349 | class R_estimator(torch.nn.Module): 350 | ''' 351 | Creates an estimator of a diagonal damping matrix of shape 352 | (nstates, nstates), where only a chosen set of states are damped. 353 | 354 | Parameters 355 | ---------- 356 | state_is_damped : listlike of bools 357 | Listlike of boolean values of length nstates. If 358 | state_is_damped[i] is True, a learnable damping parameter is 359 | created for state i. If not, the damping of state i is set to 360 | zero. 361 | ttype : torch type, default torch.float32 362 | 363 | ''' 364 | def __init__(self, state_is_damped, ttype=torch.float32): 365 | super().__init__() 366 | 367 | if not isinstance(state_is_damped, torch.Tensor): 368 | state_is_damped = torch.tensor(state_is_damped) 369 | 370 | self.state_is_damped = state_is_damped.bool() 371 | self.ttype = ttype 372 | 373 | nstates = self.state_is_damped.shape[0] 374 | 375 | self.rs = nn.Parameter(torch.zeros( 376 | torch.sum(self.state_is_damped), dtype=ttype)) 377 | 378 | self.pick_rs = torch.zeros((nstates, torch.sum(self.state_is_damped))) 379 | c = 0 380 | for i in range(nstates): 381 | if self.state_is_damped[i]: 382 | self.pick_rs[i, c] = 1 383 | c += 1 384 | 385 | def forward(self, x=None): 386 | """ 387 | Returns 388 | ------- 389 | (N, N) tensor 390 | Damping matrix 391 | """ 392 | return torch.diag(torch.abs(self.pick_rs)@(self.rs)) 393 | 394 | def get_parameters(self): 395 | """ 396 | Returns 397 | ------- 398 | ndarray 399 | Damping parameters 400 | """ 401 | return self.rs.detach().numpy() -------------------------------------------------------------------------------- /phlearn/phlearn/phnns/tests/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SINTEF/pseudo-hamiltonian-neural-networks/091065ef3c1b730d56fd413b6373d0424d8114be/phlearn/phlearn/phnns/tests/__init__.py -------------------------------------------------------------------------------- /phlearn/phlearn/phnns/tests/test_phnns.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from ..dynamic_system_neural_network import DynamicSystemNN 3 | from ..ode_models import BaseNN 4 | 5 | 6 | def test_initial_condition_sampler(): 7 | n, m = 5, 10 8 | NN = DynamicSystemNN(n) 9 | init_cond = NN._initial_condition_sampler(m) 10 | assert tuple(init_cond.size()) == (m, n) 11 | 12 | nstates=3 13 | noutputs=2 14 | hidden_dim=2 15 | bnn = BaseNN( 16 | nstates=nstates, 17 | noutputs=noutputs, 18 | hidden_dim=hidden_dim, 19 | timedependent=True, 20 | statedependent=True, 21 | ) 22 | 23 | 24 | # When developing tests, from the top level dir phlearn/ run 25 | # python3 -m phlearn.phnns.tests.test_phnns 26 | # to run this as a module, allowing use of relative imports 27 | -------------------------------------------------------------------------------- /phlearn/phlearn/phsystems/__init__.py: -------------------------------------------------------------------------------- 1 | ''' 2 | The phsystems subpackage is divided into two further subpackages, one for ODEs and one for PDEs. 3 | These implement pseudo-Hamiltonian systems and numerical integration of these to obtain data. 4 | ''' 5 | 6 | from . import ode 7 | from .ode import * 8 | from . import pde 9 | from .pde import * 10 | 11 | __all__ = ode.__all__.copy() 12 | __all__ += pde.__all__.copy() -------------------------------------------------------------------------------- /phlearn/phlearn/phsystems/ode/__init__.py: -------------------------------------------------------------------------------- 1 | """ 2 | The phsystems.ode subpackage implements pseudo-Hamiltonian ODE systems for 3 | simulation and control. 4 | 5 | Classes present in phlearn.phsystems.ode 6 | ----------------------------------------------- 7 | 8 | :py:class:`~.pseudo_hamiltonian_system.PseudoHamiltonianSystem` 9 | 10 | :py:class:`~.tank_system.TankSystem` 11 | 12 | :py:class:`~.msd_system.MassSpringDamperSystem` 13 | 14 | 15 | Functions present in phlearn.phsystems.ode 16 | ----------------------------------------------- 17 | 18 | :py:func:`~.pseudo_hamiltonian_system.zero_force` 19 | 20 | :py:func:`~.tank_system.init_tanksystem` 21 | 22 | :py:func:`~.tank_system.init_tanksystem_leaky` 23 | 24 | :py:func:`~.msd_system.init_msdsystem` 25 | 26 | :py:func:`~.msd_system.initial_condition_radial` 27 | 28 | """ 29 | 30 | from . import pseudo_hamiltonian_system 31 | from .pseudo_hamiltonian_system import * 32 | from . import tank_system 33 | from .tank_system import * 34 | from . import msd_system 35 | from .msd_system import * 36 | 37 | __all__ = pseudo_hamiltonian_system.__all__.copy() 38 | __all__ += tank_system.__all__.copy() 39 | __all__ += msd_system.__all__.copy() -------------------------------------------------------------------------------- /phlearn/phlearn/phsystems/ode/msd_system.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | from .pseudo_hamiltonian_system import PseudoHamiltonianSystem 4 | 5 | __all__ = ['MassSpringDamperSystem', 'init_msdsystem', 6 | 'initial_condition_radial'] 7 | 8 | class MassSpringDamperSystem(PseudoHamiltonianSystem): 9 | """ 10 | Implements a general forced and damped mass-spring system as a 11 | pseudo-Hamiltonian formulation:: 12 | 13 | . 14 | | q | | 0 1 | | 0 | 15 | | . | = | | * grad[H(q, p)] + | | 16 | | p | | -1 -c | | f(q, p, t) | 17 | 18 | where q is the position, p the momentum and c the damping coefficient. 19 | 20 | Parameters 21 | ---------- 22 | mass : number, default 1.0 23 | Scalar mass 24 | spring_constant : number, default 1.0 25 | Scalar spring coefficient 26 | damping : number, default 0.3 27 | Scalar damping coefficient. Corresponds to c. 28 | kwargs : any, optional 29 | Keyword arguments that are passed to PseudoHamiltonianSystem constructor. 30 | 31 | """ 32 | 33 | def __init__(self, mass=1.0, spring_constant=1.0, damping=0.3, **kwargs): 34 | R = np.diag([0, damping]) 35 | M = np.diag([spring_constant / 2, 1 / (2 * mass)]) 36 | 37 | def hamiltonian(x): 38 | return x.T @ M @ x 39 | 40 | def hamiltonian_grad(x): 41 | return 2 * M @ x 42 | 43 | super().__init__( 44 | nstates=2, 45 | hamiltonian=hamiltonian, 46 | grad_hamiltonian=hamiltonian_grad, 47 | dissipation_matrix=R, 48 | **kwargs 49 | ) 50 | 51 | 52 | def init_msdsystem(): 53 | """ 54 | Initialize a standard example of a damped mass-spring system affected by a 55 | sine force. 56 | 57 | Returns 58 | ------- 59 | MassSpringDamperSystem 60 | 61 | """ 62 | f0 = 1.0 63 | omega = 3 64 | 65 | def F(x, t): 66 | return (f0 * np.sin(omega * t)).reshape(x[..., 1:].shape) * np.array([0, 1]) 67 | 68 | return MassSpringDamperSystem( 69 | external_forces=F, init_sampler=initial_condition_radial(1, 4.5) 70 | ) 71 | 72 | 73 | def initial_condition_radial(r_min, r_max): 74 | """ 75 | Creates an initial condition sampler that draws samples uniformly 76 | from the disk r_min <= x^Tx < r_max. 77 | 78 | Returns 79 | ------- 80 | callable 81 | Function taking a numpy random generator and returning an 82 | initial state of size 2. 83 | 84 | """ 85 | 86 | def sampler(rng): 87 | r = (r_max - r_min) * np.sqrt(rng.uniform(size=1)) + r_min 88 | theta = 2.0 * np.pi * rng.uniform(size=1) 89 | q = r * np.cos(theta) 90 | p = r * np.sin(theta) 91 | return np.array([q, p]).flatten() 92 | 93 | return sampler 94 | -------------------------------------------------------------------------------- /phlearn/phlearn/phsystems/ode/pseudo_hamiltonian_system.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from scipy.integrate import solve_ivp 3 | import torch 4 | 5 | from ...utils.derivatives import time_derivative 6 | 7 | __all__ = ["PseudoHamiltonianSystem", "zero_force"] 8 | 9 | class PseudoHamiltonianSystem: 10 | """ 11 | Implements a pseudo-Hamiltonian system of the form:: 12 | 13 | dx/dt = (S(x) - R(x))*grad[H(x)] + F(x, t) 14 | 15 | where x is the system state, S is the skew-symmetric interconnection 16 | matrix, R is the positive semi-definite dissipation matrix, H is the 17 | Hamiltonian of the systemm, F is the external force(s). 18 | 19 | Parameters 20 | ---------- 21 | nstates : int 22 | Number of system states N. 23 | 24 | skewsymmetric_matrix : (N, N) ndarray or callable, default None 25 | Corresponds to the S matrix. Must either be an 26 | ndarray, or callable taking an ndarray 27 | input of shape (nsamples, nstates) and returning an ndarray 28 | of shape (nsamples, nstates, nstates). If None, 29 | the system is assumed to be canonical, and the 30 | S matrix is set to be [[0, I_N], [-I_N, 0]]. 31 | 32 | dissipation_matrix : (N, N) ndarray or callable, default None 33 | Corresponds to the R matrix. Must either be an 34 | ndarray, or callable taking an ndarray input of shape 35 | (nsamples, nstates) and returning an ndarray of shape 36 | (nsamples, nstates, nstates). If None, the R matrix is set 37 | to be the zero matrix of shape (N, N). 38 | 39 | hamiltonian : callable, default None 40 | The Hamiltonian H of the system. Callable taking a 41 | torch tensor input of shape (nsamples, nstates) and 42 | returning a torch tensor of shape (nsamples, 1). 43 | If the gradient of the Hamiltonian is not provided, 44 | the gradient of this function will be computed by torch and 45 | used instead. If this is not provided, the grad_hamiltonian 46 | must be provided. 47 | 48 | grad_hamiltonian : callable, default None 49 | The gradient of the Hamiltonian H of the system. Callable 50 | taking an ndarray input of shape (nsamples, nstates) and 51 | returning a torch tensor of shape (nsamples, nstates). 52 | If this is not provided, the hamiltonian must be provided. 53 | 54 | external_forces : callable, default None 55 | The external forces affecting system. Callable taking two 56 | ndarrays as input, x and t, of shape (nsamples, nstates), 57 | (nsamples, 1), respectively and returning an ndarray of 58 | shape (nsamples, nstates). 59 | 60 | controller : phlearn.control.PseudoHamiltonianController, 61 | default None 62 | Additional external forces set by a controller. Callable 63 | taking an ndarray x of shape (nstates,) and a scalar t as 64 | input and returning an ndarray of shape (nstates,). Note 65 | that this function should not take batch inputs, and that 66 | when calling PseudoHamiltonianSystem.sample_trajectory when a 67 | controller is provided, the Runge-Kutta 4 method will be 68 | used for integration in favor of Scipy's solve_ivp. 69 | 70 | init_sampler : callable, default None 71 | Function for sampling initial conditions. Callabale taking 72 | a numpy random generator as input and returning an ndarray 73 | of shape (nstates,) with inital conditions for the system. 74 | This sampler is used when calling 75 | PseudoHamiltonianSystem.sample_trajectory if no initial 76 | condition is provided. 77 | 78 | """ 79 | 80 | def __init__( 81 | self, 82 | nstates, 83 | skewsymmetric_matrix=None, 84 | dissipation_matrix=None, 85 | hamiltonian=None, 86 | grad_hamiltonian=None, 87 | external_forces=None, 88 | controller=None, 89 | init_sampler=None, ): 90 | self.nstates = nstates 91 | 92 | if ( 93 | skewsymmetric_matrix is not None 94 | and not callable(skewsymmetric_matrix) 95 | and not np.allclose(skewsymmetric_matrix, -skewsymmetric_matrix.T, atol=1e-15) 96 | ): 97 | raise Exception("skewsymmetric_matrix must be skew-symmetric") 98 | 99 | if hamiltonian is None and grad_hamiltonian is None: 100 | raise Exception( 101 | "Either one of hamiltonian or grad_hamiltonian must be provided" 102 | ) 103 | 104 | if skewsymmetric_matrix is None: 105 | if nstates % 2 == 1: 106 | raise Exception( 107 | "nstates must be even when skewsymmetric_matrix not provided" 108 | ) 109 | 110 | npos = nstates // 2 111 | skewsymmetric_matrix = np.block( 112 | [ 113 | [np.zeros([npos, npos]), np.eye(npos)], 114 | [-np.eye(npos), np.zeros([npos, npos])], 115 | ] 116 | ) 117 | 118 | if not callable(skewsymmetric_matrix): 119 | self.skewsymmetric_matrix = skewsymmetric_matrix 120 | self.S = lambda x: skewsymmetric_matrix 121 | else: 122 | self.skewsymmetric_matrix = None 123 | self.S = skewsymmetric_matrix 124 | 125 | if dissipation_matrix is None: 126 | dissipation_matrix = np.zeros((self.nstates, self.nstates)) 127 | 128 | if not callable(dissipation_matrix): 129 | if len(dissipation_matrix.shape) == 1: 130 | dissipation_matrix = np.diag(dissipation_matrix) 131 | self.dissipation_matrix = dissipation_matrix 132 | self.R = lambda x: dissipation_matrix 133 | else: 134 | self.dissipation_matrix = None 135 | self.R = dissipation_matrix 136 | 137 | self.H = hamiltonian 138 | self.dH = grad_hamiltonian 139 | if grad_hamiltonian is None: 140 | self.dH = self._dH 141 | 142 | self.controller = controller 143 | 144 | self.external_forces = external_forces 145 | if external_forces is None: 146 | self.external_forces = zero_force 147 | 148 | if init_sampler is not None: 149 | self._initial_condition_sampler = init_sampler 150 | 151 | self.seed(None) 152 | 153 | def seed(self, seed): 154 | """ 155 | Set the internal random state. 156 | 157 | Parameters 158 | ---------- 159 | seed : int 160 | 161 | """ 162 | 163 | self.rng = np.random.default_rng(seed) 164 | 165 | def time_derivative(self, integrator, *args, **kwargs): 166 | """ 167 | See :py:meth:~`utils.derivatives.time_derivative` 168 | """ 169 | return time_derivative(integrator, self.x_dot, *args, **kwargs) 170 | 171 | def x_dot(self, x, t, u=None): 172 | """ 173 | Computes the time derivative by the right hand side of the pseudo- 174 | Hamiltonian equation. 175 | 176 | Parameters 177 | ---------- 178 | x : (..., N) ndarray 179 | t : (..., 1) ndarray 180 | u : (..., N) ndarray or None, default None 181 | 182 | Returns 183 | ------- 184 | (..., N) ndarray 185 | 186 | """ 187 | 188 | S = self.S(x) 189 | R = self.R(x) 190 | dH = self.dH(x.T).T 191 | 192 | if (len(S.shape) == 3) or (len(R.shape) == 3): 193 | dynamics = np.matmul(S - R, np.atleast_3d(dH)).reshape( 194 | x.shape 195 | ) + self.external_forces(x, t).reshape(x.shape) 196 | else: 197 | dynamics = dH @ (S.T - R.T) + self.external_forces(x, t).reshape(x.shape) 198 | 199 | if u is not None: 200 | dynamics += u 201 | 202 | return dynamics 203 | 204 | def sample_trajectory(self, t, x0=None, noise_std=0, reference=None): 205 | """ 206 | Samples a trajectory of the system at times *t*. 207 | 208 | Parameters 209 | ---------- 210 | t : (T, 1) ndarray 211 | Times at which the trajectory is sampled. 212 | x0 : (N,) ndarray, default None 213 | Initial condition. 214 | noise_std : number, default 0. 215 | Standard deviation of Gaussian white noise added to the 216 | samples of the trajectory. 217 | reference : phlearn.control.Reference, default None 218 | If the system has a controller a reference object may be 219 | passed. 220 | 221 | Returns 222 | ------- 223 | x : (T, N) ndarray 224 | dxdt : (T, N) ndarray 225 | t : (T, 1) ndarray 226 | us : (T, N) ndarray 227 | 228 | """ 229 | 230 | if x0 is None: 231 | x0 = self._initial_condition_sampler(self.rng) 232 | 233 | if self.controller is None: 234 | x_dot = lambda t, x: self.x_dot( 235 | x.reshape(1, x.shape[-1]), np.array(t).reshape((1, 1)) 236 | ) 237 | out_ivp = solve_ivp( 238 | fun=x_dot, t_span=(t[0], t[-1]), y0=x0, t_eval=t, rtol=1e-10 239 | ) 240 | x, t = out_ivp["y"].T, out_ivp["t"].T 241 | dxdt = self.x_dot(x, t) 242 | us = None 243 | else: 244 | # Use RK4 integrator instead of solve_ivp when controller 245 | # is provided 246 | self.controller.reset() 247 | if reference is not None: 248 | self.controller.set_reference(reference) 249 | x = np.zeros([t.shape[0], x0.shape[-1]]) 250 | dxdt = np.zeros_like(x) 251 | us = np.zeros([t.shape[0] - 1, x0.shape[-1]]) 252 | x[0, :] = x0 253 | for i, t_step in enumerate(t[:-1]): 254 | dt = t[i + 1] - t[i] 255 | us[i, :] = self.controller(x[i, :], t_step) 256 | dxdt[i, :] = self.time_derivative( 257 | "rk4", 258 | x[i : i + 1, :], 259 | x[i : i + 1, :], 260 | np.array([t_step]), 261 | np.array([t_step]), 262 | dt, 263 | u=us[i : i + 1, :], 264 | ) 265 | x[i + 1, :] = x[i, :] + dt * dxdt[i, :] 266 | 267 | # Add noise: 268 | x += self.rng.normal(size=x.shape) * noise_std 269 | dxdt += self.rng.normal(size=dxdt.shape) * noise_std 270 | 271 | return x, dxdt, t, us 272 | 273 | def set_controller(self, controller): 274 | """ 275 | Set system controller. 276 | """ 277 | self.controller = controller 278 | 279 | def _dH(self, x): 280 | x = torch.tensor(x, requires_grad=True) 281 | return ( 282 | torch.autograd.grad( 283 | self.H(x).sum(), x, retain_graph=False, create_graph=False 284 | )[0] 285 | .detach() 286 | .numpy() 287 | ) 288 | 289 | def _initial_condition_sampler(self, rng=None): 290 | if rng is None: 291 | assert self.rng is not None 292 | rng = self.rng 293 | return rng.uniform(low=-1.0, high=1.0, size=self.nstates) 294 | 295 | 296 | def zero_force(x, t=None): 297 | """ 298 | A force term that is always zero. 299 | 300 | Parameters 301 | ---------- 302 | x : (..., N) ndarray 303 | t : (..., 1) ndarray, default None 304 | 305 | Returns 306 | ------- 307 | (..., N) ndarray 308 | All zero ndarray 309 | 310 | """ 311 | 312 | return np.zeros_like(x) 313 | -------------------------------------------------------------------------------- /phlearn/phlearn/phsystems/ode/tank_system.py: -------------------------------------------------------------------------------- 1 | 2 | import numbers 3 | 4 | import networkx as nx 5 | import numpy as np 6 | 7 | from .pseudo_hamiltonian_system import PseudoHamiltonianSystem 8 | 9 | __all__ = ['TankSystem', 'init_tanksystem', 10 | 'init_tanksystem_leaky'] 11 | 12 | class TankSystem(PseudoHamiltonianSystem): 13 | """ 14 | Implements a pseudo-Hamiltonian version of a coupled tanks system:: 15 | 16 | . | -R_p B^T | 17 | x = | | * grad[H(x)] + F(x, t) 18 | | -B 0 | 19 | 20 | 21 | where the state x = [phi, mu], phi and mu being proportional to 22 | pipe flows and tank levels, respectively. The interconnection of the 23 | tanks is described by a directed graph where each vertex is a tank 24 | and each edge is a pipe between two tanks. The incidence matrix of 25 | this graph corresponds to the matrix B. 26 | The number of tanks is denoted by ntanks, the number of pipes by 27 | npipes. The number of states is denoted by 28 | nstates = npipes + ntanks. 29 | 30 | External interaction can be specified for every state, but the most 31 | physically interpretable setting is to only have external 32 | interaction for the states corresponding to the tanks, which can be 33 | seen as volumetric flows into or out of the tanks. 34 | 35 | Each tank is assumed to have a uniform cross-section w.r.t. it's 36 | height. 37 | 38 | The interconnection of tanks and pipes can either be specified by a 39 | directed graph or by an interconnection matrix S. 40 | 41 | parameters 42 | ---------- 43 | incidence_matrix : (N, N) ndarray, default None 44 | Incidence matrix of the graph describing the tank system. 45 | Corresponds to the B matrix. Inferred from system_graph if 46 | system_graph is provided. 47 | 48 | system_graph : networkx.Graph, default None 49 | networkx directed graph describing the interconnection of 50 | the tanks. The graph has ntanks vertices and npipes edges. 51 | 52 | npipes : int, default None 53 | Number of pipes. Inferred from system_graph if system_graph 54 | is provided. 55 | 56 | ntanks : int, default None 57 | Number of tanks. Inferred from system_graph if system_graph 58 | is provided. 59 | dissipation_pipes : ndarray, default None 60 | ndarray of size (npipes,) or (npipes, npipes), describing 61 | energy loss in the pipes. Corresponds to R_p. Defaults to 62 | zero if not provided. 63 | J : number, default 1. 64 | Scalar or ndarray of size (npipes,) of proportionality 65 | constants relating the change in volumetric flow through 66 | each pipe to the pressure drop and potential friction 67 | through each pipe. If scalar, the same constant is 68 | used for all pipes. 69 | A : (ntanks,) ndarray or number, default 1. 70 | Scalar ndarray of size (ntanks,) of tank cross-section areas. 71 | If scalar, the same area is used for all tanks. 72 | rho : number, default 1. 73 | Positive scalar. Density of liquid. 74 | g : number, default 9.81 75 | Positive scalar. Gravitational acceleration. 76 | kwargs : any 77 | Keyword arguments passed to PseudoHamiltonianSystem 78 | constructor. 79 | 80 | """ 81 | 82 | def __init__(self, incidence_matrix=None, system_graph=None, npipes=None, 83 | ntanks=None, dissipation_pipes=None, J=1., A=1., rho=1., 84 | g=9.81, **kwargs): 85 | self.system_graph = system_graph 86 | 87 | if system_graph is not None: 88 | self.npipes = system_graph.number_of_edges() 89 | self.ntanks = system_graph.number_of_nodes() 90 | B = np.array(nx.linalg.graphmatrix.incidence_matrix( 91 | system_graph, oriented=True).todense()) 92 | else: 93 | self.npipes = npipes 94 | self.ntanks = ntanks 95 | B = incidence_matrix 96 | 97 | skewsymmetric_matrix = np.block( 98 | [[np.zeros((self.npipes, self.npipes)), B.T], 99 | [-B, np.zeros((self.ntanks, self.ntanks))]]) 100 | 101 | nstates = self.npipes + self.ntanks 102 | 103 | if dissipation_pipes is None: 104 | dissipation_pipes = np.zeros((nstates, nstates)) 105 | elif len(dissipation_pipes.shape) == 1: 106 | dissipation_pipes = np.diag(dissipation_pipes) 107 | 108 | dissipation = np.block( 109 | [[dissipation_pipes, np.zeros([self.npipes, self.ntanks])], 110 | [np.zeros([self.ntanks, self.npipes]), 111 | np.zeros([self.ntanks, self.ntanks])]]) 112 | 113 | if isinstance(J, numbers.Number): 114 | J = J*np.ones(self.npipes) 115 | if isinstance(A, numbers.Number): 116 | A = A*np.ones(self.ntanks) 117 | 118 | self.Hvec = np.concatenate((1/J, rho*g / A)) 119 | super().__init__(nstates, 120 | hamiltonian=self.H_tanksystem, 121 | grad_hamiltonian=self.dH_tanksystem, 122 | skewsymmetric_matrix=skewsymmetric_matrix, 123 | dissipation_matrix=dissipation, 124 | **kwargs) 125 | 126 | def H_tanksystem(self, x, t=None): 127 | return x**2 @ self.Hvec / 2 128 | 129 | def dH_tanksystem(self, x, t=None): 130 | return (x.T * self.Hvec).T 131 | 132 | def pipeflows(self, x): 133 | return x[..., :self.npipes] 134 | 135 | def tanklevels(self, x): 136 | return x[..., self.npipes:] 137 | 138 | def init_tanksystem(u=None): 139 | """ 140 | Initialize standard tank system. 141 | 142 | Parameters 143 | ---------- 144 | u : phlearn.control.PseudoHamiltonianController, defult None 145 | 146 | Returns 147 | ------- 148 | TankSystem 149 | 150 | """ 151 | 152 | G_s = nx.DiGraph() 153 | G_s.add_edge(1, 2) 154 | G_s.add_edge(2, 3) 155 | G_s.add_edge(3, 4) 156 | G_s.add_edge(1, 3) 157 | G_s.add_edge(1, 4) 158 | 159 | npipes = G_s.number_of_edges() 160 | ntanks = G_s.number_of_nodes() 161 | nstates = npipes + ntanks 162 | R = 1.e-2*np.diag(np.array([3., 3., 9., 3., 3.])) 163 | J = 2.e-2*np.ones(npipes) 164 | A = np.ones(ntanks) 165 | ext_filter = np.zeros(nstates) 166 | ext_filter[-1] = 1 167 | 168 | def F(x, t=None): 169 | return -1.e1*np.fmin(0.3, np.fmax(x, -0.3))*ext_filter 170 | 171 | return TankSystem(system_graph=G_s, dissipation_pipes=R, J=J, A=A, 172 | external_forces=F, controller=u) 173 | 174 | 175 | def init_tanksystem_leaky(nleaks=0): 176 | """ 177 | Initialize tank system with a leaks. 178 | 179 | Parameters 180 | ---------- 181 | nleaks : int, default 0 182 | If 0, no leaks. If 1, there is a leak on the last tank. If 2, 183 | there is a leak on the last and tank number 2. 184 | 185 | Returns 186 | ------- 187 | TankSystem 188 | 189 | """ 190 | 191 | G_s = nx.DiGraph() 192 | G_s.add_edge(1, 2) 193 | G_s.add_edge(2, 3) 194 | G_s.add_edge(3, 4) 195 | G_s.add_edge(1, 3) 196 | G_s.add_edge(1, 4) 197 | 198 | npipes = G_s.number_of_edges() 199 | ntanks = G_s.number_of_nodes() 200 | nstates = npipes + ntanks 201 | R = 1.e-2*np.diag(np.array([3., 3., 9., 3., 3.])) 202 | J = 2.e-2*np.ones(npipes) 203 | A = np.ones(ntanks) 204 | 205 | if nleaks == 0: 206 | def F(x, t=None): 207 | return np.zeros_like(x) 208 | else: 209 | if nleaks == 1: 210 | ext_filter = np.zeros(nstates) 211 | ext_filter[-1] = 3 212 | else: 213 | ext_filter = np.zeros(nstates) 214 | ext_filter[-1] = 3 215 | ext_filter[-4] = 1 216 | 217 | def F(x, t=None): 218 | return -1.e1*np.minimum(0.3, np.maximum(x, -0.3))*ext_filter 219 | 220 | return TankSystem(system_graph=G_s, dissipation_pipes=R, J=J, A=A, 221 | external_forces=F, controller=None) 222 | -------------------------------------------------------------------------------- /phlearn/phlearn/phsystems/ode/tests/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SINTEF/pseudo-hamiltonian-neural-networks/091065ef3c1b730d56fd413b6373d0424d8114be/phlearn/phlearn/phsystems/ode/tests/__init__.py -------------------------------------------------------------------------------- /phlearn/phlearn/phsystems/ode/tests/test_phsystems.py: -------------------------------------------------------------------------------- 1 | # TODO: Test case when len(S.shape) == 3 and len(R.shape) == 3 2 | # TODO: Have ignored controller for now. I.e., assumed controller = None throughout 3 | 4 | # When developing tests, from the top level dir (i.e., phlearn) run 5 | # python3 -m phlearn.phnns.tests.test_phsystems 6 | # to run this as a module, allowing use of relative imports 7 | 8 | import numpy as np 9 | from ..pseudo_hamiltonian_system import PseudoHamiltonianSystem 10 | import torch 11 | 12 | 13 | N_STATES = 10 14 | N_TIMESTEPS = 100 15 | 16 | X_RAND = np.random.rand(N_STATES) 17 | T_AXIS = np.linspace(0,10, N_TIMESTEPS) 18 | EXTERNAL_FORCE = np.random.rand(N_STATES) 19 | DISSIPATION_MATRIX = np.random.rand(N_STATES, N_STATES) 20 | DISSIPATION_MATRIX = DISSIPATION_MATRIX + DISSIPATION_MATRIX.T 21 | STRUCTURE_MATRIX = np.random.rand(N_STATES, N_STATES) 22 | STRUCTURE_MATRIX = STRUCTURE_MATRIX - STRUCTURE_MATRIX.T 23 | 24 | 25 | def H(x): 26 | x = x.flatten() if torch.is_tensor(x) else x 27 | return sum(x[i] ** 2 for i in range(N_STATES)) 28 | 29 | 30 | def dH(x): 31 | return np.array([2 * x[i] for i in range(N_STATES)]) 32 | 33 | 34 | psh_kwargs = dict( 35 | nstates=N_STATES, 36 | hamiltonian=H, 37 | ) 38 | 39 | 40 | def test_R_is_zero_matrix_when_dissipation_matrix_is_None(): 41 | phs = PseudoHamiltonianSystem(**psh_kwargs) 42 | assert not phs.R( 43 | np.random.rand(psh_kwargs["nstates"]) 44 | ).any(), "Dissipation_matrix not defaulting to zero matrix" 45 | 46 | 47 | def test_S_is_canonical_when_structure_matrix_is_None(): 48 | phs = PseudoHamiltonianSystem(**psh_kwargs) 49 | m = int(N_STATES / 2) 50 | O, I = np.zeros([m, m]), np.eye(m) 51 | S = np.block([[O, I], [-I, O]]) 52 | assert np.array_equal( 53 | S, phs.S(X_RAND) 54 | ), "strucutre_matrix not defulating to canonical form" 55 | 56 | 57 | def test_dissipation_matrix_is_returned(): 58 | phs = PseudoHamiltonianSystem(dissipation_matrix=DISSIPATION_MATRIX, **psh_kwargs) 59 | assert np.array_equal( 60 | DISSIPATION_MATRIX, phs.R(X_RAND) 61 | ), "dissipation_matrix not returned" 62 | 63 | 64 | def test_structure_matrix_is_returned(): 65 | psh_kwargs_loc = psh_kwargs 66 | n = psh_kwargs["nstates"] 67 | phs = PseudoHamiltonianSystem(skewsymmetric_matrix=STRUCTURE_MATRIX, **psh_kwargs_loc) 68 | assert np.array_equal( 69 | STRUCTURE_MATRIX, phs.S(X_RAND) 70 | ), "skewsymmetric_matrix not returned" 71 | 72 | 73 | def test_dH_calculated_correctly(): 74 | phs = PseudoHamiltonianSystem(**psh_kwargs) 75 | assert np.allclose(phs.dH(X_RAND), 2 * X_RAND), "grad_H computed incorrectly" 76 | 77 | 78 | def test_x_dot(): 79 | phs = PseudoHamiltonianSystem( 80 | dissipation_matrix=lambda x: DISSIPATION_MATRIX, 81 | skewsymmetric_matrix=lambda x: STRUCTURE_MATRIX, 82 | external_forces=lambda x, t: EXTERNAL_FORCE, 83 | **psh_kwargs 84 | ) 85 | x_dot = dH(X_RAND) @ (STRUCTURE_MATRIX.T - DISSIPATION_MATRIX.T) + EXTERNAL_FORCE 86 | assert np.allclose(phs.x_dot(X_RAND, T_AXIS), x_dot), "x_dot() returns incorrect ODE" 87 | 88 | 89 | def test_sample_trajectory_on_const_ode(): 90 | phs = PseudoHamiltonianSystem(skewsymmetric_matrix=0 * STRUCTURE_MATRIX, **psh_kwargs) 91 | x, _, _, _ = phs.sample_trajectory(t=[0, 1], x0=X_RAND, noise_std=0) 92 | assert np.allclose( 93 | x[-1], X_RAND 94 | ), "Solution of a constant ODE does not remain constant" 95 | -------------------------------------------------------------------------------- /phlearn/phlearn/phsystems/pde/__init__.py: -------------------------------------------------------------------------------- 1 | """ 2 | The phsystems.pde subpackage implements pseudo-Hamiltonian PDE systems for 3 | discretization in space and integration in time. 4 | 5 | Classes present in phlearn.phsystems.pde 6 | ----------------------------------------------- 7 | 8 | :py:class:`~.pseudo_hamiltonian_pde_system.PseudoHamiltonianPDESystem` 9 | 10 | :py:class:`~.kdv_system.KdVSystem` 11 | 12 | :py:class:`~.bbm_system.BBMSystem` 13 | 14 | :py:class:`~.perona_malik_system.PeronaMalikSystem` 15 | 16 | :py:class:`~.cahn_hilliard_system.CahnHilliardSystem` 17 | 18 | :py:class:`~.heat_system.HeatSystem` 19 | 20 | :py:class:`~.allen_cahn_system.AllenCahnSystem` 21 | 22 | 23 | Functions present in phlearn.phsystems.pde 24 | ----------------------------------------------- 25 | 26 | :py:func:`~.pseudo_hamiltonian_pde_system.zero_force` 27 | 28 | :py:func:`~.kdv_system.initial_condition_kdv` 29 | 30 | :py:func:`~.bbm_system.initial_condition_bbm` 31 | 32 | :py:func:`~.perona_malik_system.initial_condition_ac` 33 | 34 | :py:func:`~.cahn_hilliard_system.initial_condition_pm` 35 | 36 | :py:func:`~.heat_system.initial_condition_heat` 37 | 38 | :py:func:`~.allen_cahn_system.initial_condition_ac` 39 | 40 | """ 41 | 42 | from . import pseudo_hamiltonian_pde_system 43 | from .pseudo_hamiltonian_pde_system import * 44 | from . import kdv_system 45 | from .kdv_system import * 46 | from . import bbm_system 47 | from .bbm_system import * 48 | from . import perona_malik_system 49 | from .perona_malik_system import * 50 | from . import heat_system 51 | from .heat_system import * 52 | from . import allen_cahn_system 53 | from .allen_cahn_system import * 54 | from . import cahn_hilliard_system 55 | from .cahn_hilliard_system import * 56 | 57 | __all__ = pseudo_hamiltonian_pde_system.__all__.copy() 58 | __all__ += kdv_system.__all__.copy() 59 | __all__ += bbm_system.__all__.copy() 60 | __all__ += perona_malik_system.__all__.copy() 61 | __all__ += heat_system.__all__.copy() 62 | __all__ += allen_cahn_system.__all__.copy() 63 | __all__ += cahn_hilliard_system.__all__.copy() 64 | -------------------------------------------------------------------------------- /phlearn/phlearn/phsystems/pde/allen_cahn_system.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from scipy.sparse import spdiags 3 | 4 | from .pseudo_hamiltonian_pde_system import PseudoHamiltonianPDESystem 5 | 6 | __all__ = ['AllenCahnSystem', 'initial_condition_ac'] 7 | 8 | class AllenCahnSystem(PseudoHamiltonianPDESystem): 9 | """ 10 | Implements a discretization of the Allen-Cahn equation with an 11 | optional external force, as described by:: 12 | 13 | u_t - u_xx - u + u^3 = f(u,t,x) 14 | 15 | on the pseudo-Hamiltonian formulation:: 16 | 17 | du/dt = -grad[V(u)] + F(u,t,x) 18 | 19 | where u is a vector of the system states at the spatial points given by 20 | x, and t is time. 21 | 22 | The system is by default integrated in time using the implicit midpoint method. 23 | The line 'self.sample_trajectory = self.sample_trajectory_midpoint' can be 24 | commented out to instead use the RK45 method of scipy. 25 | 26 | Parameters 27 | ---------- 28 | x : nparray, default np.linspace(0, 1.0 - 1/100, 100) 29 | The spatial discretization points. 30 | nu : number, default -1.0 31 | The parameter nu in the Cahn-Hilliard equation. 32 | alpha : number, default 1.0 33 | The parameter alpha in the Cahn-Hilliard equation. 34 | mu : number, default -0.001 35 | The parameter mu in the Cahn-Hilliard equation. 36 | init_sampler : callable, default None 37 | Function for sampling initial conditions. Callable taking 38 | a numpy random generator as input and returning an ndarray 39 | of shape same as x with initial conditions for the system. 40 | This sampler is used when calling 41 | CahnHilliardSystem.sample_trajectory if no initial 42 | condition is provided. 43 | kwargs : any, optional 44 | Keyword arguments that are passed to PseudoHamiltonianPDESystem 45 | constructor. 46 | 47 | """ 48 | 49 | def __init__( 50 | self, 51 | x=np.linspace(0, 6.0 - 6 / 300, 300), 52 | nu=0.001, 53 | init_sampler=None, 54 | **kwargs 55 | ): 56 | M = x.size 57 | dx = x[-1] / (M - 1) 58 | e = np.ones(M) 59 | # Forward difference matrix: 60 | Dp = 1 / dx * spdiags([e, -e, e], np.array([-M + 1, 0, 1]), M, M).toarray() 61 | # Central difference matrix: 62 | D1 = ( 63 | 0.5 64 | / dx 65 | * spdiags([e, -e, e, -e], np.array([-M + 1, -1, 1, M - 1]), M, M).toarray() 66 | ) 67 | # 2nd order central difference matrix: 68 | D2 = ( 69 | 1 70 | / dx**2 71 | * spdiags( 72 | [e, e, -2 * e, e, e], np.array([-M + 1, -1, 0, 1, M - 1]), M, M 73 | ).toarray() 74 | ) 75 | I = np.eye(M) 76 | skewsymmetric_matrix = D1 77 | self.x = x 78 | self.M = M 79 | self.D2 = D2 80 | 81 | def ham(u): 82 | return np.sum(np.zeros_like(u), axis=1) 83 | 84 | def dissintegral(u): 85 | return ( 86 | 1 87 | / 2 88 | * np.sum( 89 | (nu * np.matmul(Dp, u.T) ** 2).T - u**2 + 1 / 2 * u**4, axis=1 90 | ) 91 | ) 92 | 93 | def ham_grad(u): 94 | return np.zeros_like(u) 95 | 96 | def dissintegral_grad(u): 97 | return -nu * u @ D2 - u + u**3 98 | 99 | def ham_hessian(u): 100 | return np.zeros_like(D1) 101 | 102 | def dissintegral_hessian(u): 103 | return -nu * D2 - I + 3 * np.diag(u**2) 104 | 105 | if init_sampler is None: 106 | init_sampler = initial_condition_ac(x) 107 | 108 | super().__init__( 109 | nstates=M, 110 | skewsymmetric_matrix=skewsymmetric_matrix, 111 | hamiltonian=ham, 112 | dissintegral=dissintegral, 113 | grad_hamiltonian=ham_grad, 114 | grad_dissintegral=dissintegral_grad, 115 | hess_hamiltonian=ham_hessian, 116 | hess_dissintegral=dissintegral_hessian, 117 | init_sampler=init_sampler, 118 | **kwargs 119 | ) 120 | 121 | self.skewsymmetric_matrix_flat = 0.5 / dx * np.array([[[-1, 0, 1]]]) 122 | 123 | 124 | def initial_condition_ac(x=np.linspace(0, 6.0 - 6 / 300, 300)): 125 | """ 126 | Creates an initial condition sampler for the Allen-Cahn eqation. 127 | 128 | Parameters 129 | ---------- 130 | x : numpy.ndarray, optional 131 | Spatial grid on which to create the initial conditions. The default is 132 | an equidistant grid between 0 and 5.98 with step size 0.02. 133 | 134 | Returns 135 | ------- 136 | callable 137 | A function that takes a numpy random generator as input and returns an 138 | initial state on the spatial grid x. 139 | """ 140 | 141 | M = x.size 142 | P = (x[-1] - x[0]) * M / (M - 1) 143 | 144 | def sampler(rng): 145 | d1, d2 = rng.uniform(0.0, 1.0, 2) 146 | k1, k2 = rng.uniform(0.5, 1.0, 2) 147 | u0 = 0 148 | u0 += k1 * np.cos(2 * np.pi / P * (x - d1 * P)) 149 | u0 += k2 * np.cos(2 * np.pi / P * (x - d2 * P)) 150 | return u0 151 | 152 | return sampler 153 | -------------------------------------------------------------------------------- /phlearn/phlearn/phsystems/pde/bbm_system.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from scipy.sparse import spdiags 3 | 4 | from .pseudo_hamiltonian_pde_system import PseudoHamiltonianPDESystem 5 | 6 | __all__ = ['BBMSystem', 'initial_condition_bbm'] 7 | 8 | class BBMSystem(PseudoHamiltonianPDESystem): 9 | """ 10 | BBMSystem class, representing a discretization of the Benjamin-Bona-Mahony 11 | (BBM) equation. 12 | 13 | Implements a discretization of the BBM equation with an optional 14 | viscosity term and external forces:: 15 | 16 | u_t - u_xxt + u_x + u u_x - nu u_xx = f(u,x,t) 17 | 18 | on the pseudo-Hamiltonian formulation:: 19 | 20 | (1-d^2/dx^2)(du/dt) = d/dx(grad[H(u)]) - grad[V(u)] + F(u, t, x) 21 | 22 | where u is a vector of the system states at the spatial points given by 23 | x, and t is time. 24 | 25 | The system is by default integrated in time using the implicit midpoint 26 | method. The line 'self.sample_trajectory = self.sample_trajectory_midpoint' 27 | can be commented out to instead use the RK45 method of scipy. 28 | 29 | Parameters 30 | ---------- 31 | x : nparray, default np.linspace(0, 20.0 - 0.2, 100) 32 | The spatial discretization points 33 | nu : number, default 0.0 34 | The damping coefficient 35 | init_sampler : callable, default None 36 | Function for sampling initial conditions. Callable taking 37 | a numpy random generator as input and returning an ndarray 38 | of shape same as x with initial conditions for the system. 39 | This sampler is used when calling BBMSystem.sample_trajectory if no 40 | initial condition is provided. 41 | kwargs : any, optional 42 | Keyword arguments that are passed to PseudoHamiltonianPDESystem 43 | constructor. 44 | 45 | """ 46 | 47 | def __init__( 48 | self, x=np.linspace(0, 20.0 - 0.2, 100), nu=0.0, init_sampler=None, **kwargs 49 | ): 50 | M = x.size 51 | dx = x[-1] / (M - 1) 52 | e = np.ones(M) 53 | # Forward difference matrix: 54 | Dp = 1 / dx * spdiags([e, -e, e], np.array([-M + 1, 0, 1]), M, M).toarray() 55 | # Central difference matrix: 56 | D1 = ( 57 | 0.5 58 | / dx 59 | * spdiags([e, -e, e, -e], np.array([-M + 1, -1, 1, M - 1]), M, M).toarray() 60 | ) 61 | # 2nd order central difference matrix: 62 | D2 = ( 63 | 1 64 | / dx**2 65 | * spdiags( 66 | [e, e, -2 * e, e, e], np.array([-M + 1, -1, 0, 1, M - 1]), M, M 67 | ).toarray() 68 | ) 69 | I = np.eye(M) 70 | lhs_matrix = I - D2 71 | self.lhs_matrix_flat = np.array([[[0, 1, 0]]]) - 1.0 / dx**2 * np.array( 72 | [[[1, -2, 1]]] 73 | ) 74 | self.lhs_matrix = lhs_matrix 75 | skewsymmetric_matrix = D1 76 | self.x = x 77 | self.M = M 78 | self.nu = nu 79 | self.D2 = D2 80 | self.sample_trajectory = self.sample_trajectory_midpoint 81 | 82 | def ham(u): 83 | return np.sum(-1 / 2 * u**2 - 1 / 6 * u**3, axis=1) 84 | 85 | def dissintegral(u): 86 | return np.sum(0.5 * nu * (np.matmul(Dp, u.T) ** 2).T, axis=1) 87 | 88 | def ham_grad(u): 89 | return -u - 0.5 * u**2 90 | 91 | def dissintegral_grad(u): 92 | return -nu * u @ D2 93 | 94 | def ham_hessian(u): 95 | return -I - np.diag(u) 96 | 97 | def dissintegral_hessian(u): 98 | return -nu * D2 99 | 100 | if init_sampler is None: 101 | init_sampler = initial_condition_bbm(x) 102 | 103 | super().__init__( 104 | nstates=M, 105 | lhs_matrix=lhs_matrix, 106 | skewsymmetric_matrix=skewsymmetric_matrix, 107 | hamiltonian=ham, 108 | dissintegral=dissintegral, 109 | grad_hamiltonian=ham_grad, 110 | grad_dissintegral=dissintegral_grad, 111 | hess_hamiltonian=ham_hessian, 112 | hess_dissintegral=dissintegral_hessian, 113 | init_sampler=init_sampler, 114 | **kwargs 115 | ) 116 | 117 | self.skewsymmetric_matrix_flat = 0.5 / dx * np.array([[[-1, 0, 1]]]) 118 | self.dissipation_matrix_flat = np.array([[[0, 1, 0]]]) 119 | 120 | 121 | def initial_condition_bbm(x=np.linspace(0, 50.0 - 0.5, 100)): 122 | """ 123 | Creates an initial condition sampler for the Benjamin-Bona-Mahony equation. 124 | 125 | Parameters 126 | ---------- 127 | x : numpy.ndarray, optional 128 | Spatial grid on which to create the initial conditions. The default is 129 | an equidistant grid between 0 and 49.5 with step size 0.5. 130 | 131 | Returns 132 | ------- 133 | callable 134 | A function that takes a numpy random generator as input and returns an 135 | initial state on the spatial grid x. 136 | """ 137 | 138 | M = x.size 139 | P = (x[-1] - x[0]) * M / (M - 1) 140 | 141 | def sech(a): 142 | return 1 / np.cosh(a) 143 | 144 | def sampler(rng): 145 | c1, c2 = rng.uniform(1.0, 4.0, 2) 146 | d1, d2 = rng.uniform(0.0, 1.0, 2) 147 | u0 = 0 148 | u0 += ( 149 | 3 150 | * (c1 - 1) 151 | * sech(1 / 2 * np.sqrt(1 - 1 / c1) * ((x + P / 2 - P * d1) % P - P / 2)) 152 | ** 2 153 | ) 154 | u0 += ( 155 | 3 156 | * (c2 - 1) 157 | * sech(1 / 2 * np.sqrt(1 - 1 / c2) * ((x + P / 2 - P * d2) % P - P / 2)) 158 | ** 2 159 | ) 160 | u0 = np.concatenate([u0[M:], u0[:M]], axis=-1) 161 | return u0 162 | 163 | return sampler 164 | -------------------------------------------------------------------------------- /phlearn/phlearn/phsystems/pde/cahn_hilliard_system.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from scipy.sparse import spdiags 3 | 4 | from .pseudo_hamiltonian_pde_system import PseudoHamiltonianPDESystem 5 | 6 | __all__ = ['CahnHilliardSystem', 'initial_condition_ch'] 7 | 8 | class CahnHilliardSystem(PseudoHamiltonianPDESystem): 9 | """ 10 | Implements a discretization of the Cahn-Hilliard equation with an 11 | optional external force, as described by:: 12 | 13 | u_t - (nu u + alpha u^3 + mu u_xx)_xx = f(u,t,x) 14 | 15 | on the pseudo-Hamiltonian formulation:: 16 | 17 | du/dt = d^2/dx^2 grad[V(u)] + F(u,t,x) 18 | 19 | where u is a vector of the system states at the spatial points given by 20 | x, and t is time. 21 | 22 | The system is by default integrated in time using the implicit midpoint method. 23 | The line 'self.sample_trajectory = self.sample_trajectory_midpoint' can be 24 | commented out to instead use the RK45 method of scipy. 25 | 26 | Parameters 27 | ---------- 28 | x : nparray, default np.linspace(0, 1.0 - 1/100, 100) 29 | The spatial discretization points. 30 | nu : number, default -1.0 31 | The parameter nu in the Cahn-Hilliard equation. 32 | alpha : number, default 1.0 33 | The parameter alpha in the Cahn-Hilliard equation. 34 | mu : number, default -0.001 35 | The parameter mu in the Cahn-Hilliard equation. 36 | init_sampler : callable, default None 37 | Function for sampling initial conditions. Callable taking 38 | a numpy random generator as input and returning an ndarray 39 | of shape same as x with initial conditions for the system. 40 | This sampler is used when calling 41 | CahnHilliardSystem.sample_trajectory if no initial 42 | condition is provided. 43 | kwargs : any, optional 44 | Keyword arguments that are passed to PseudoHamiltonianPDESystem 45 | constructor. 46 | 47 | """ 48 | 49 | def __init__( 50 | self, 51 | x=np.linspace(0, 1.0 - 1 / 100, 100), 52 | nu=-1.0, 53 | alpha=1.0, 54 | mu=-0.001, 55 | init_sampler=None, 56 | **kwargs 57 | ): 58 | M = x.size 59 | dx = x[-1] / (M - 1) 60 | e = np.ones(M) 61 | # Forward difference matrix: 62 | Dp = ( 63 | 1 / dx * spdiags([e, -e, e], np.array([-M + 1, 0, 1]), M, M).toarray() 64 | ) 65 | # Central difference matrix: 66 | D1 = ( 67 | 0.5 68 | / dx 69 | * spdiags([e, -e, e, -e], np.array([-M + 1, -1, 1, M - 1]), M, M).toarray() 70 | ) 71 | # 2nd order central difference matrix: 72 | D2 = ( 73 | 1 74 | / dx**2 75 | * spdiags( 76 | [e, e, -2 * e, e, e], np.array([-M + 1, -1, 0, 1, M - 1]), M, M 77 | ).toarray() 78 | ) 79 | I = np.eye(M) 80 | skewsymmetric_matrix = D1 81 | dissipation_matrix = -D2 82 | self.x = x 83 | self.M = M 84 | self.D2 = D2 85 | self.sample_trajectory = self.sample_trajectory_midpoint 86 | 87 | def dissintegral(u): 88 | return ( 89 | 1 90 | / 2 91 | * np.sum( 92 | nu * u**2 93 | + 1 / 2 * alpha * u**4 94 | - mu * (np.matmul(Dp, u.T) ** 2).T, 95 | axis=1, 96 | ) 97 | ) 98 | 99 | def dissintegral_grad(u): 100 | return nu * u + alpha * u**3 + mu * u @ D2 101 | 102 | def ham_hessian(u): 103 | return np.zeros_like(D1) 104 | 105 | def dissintegral_hessian(u): 106 | return nu * I + 3 * alpha * np.diag(u**2) + mu * D2 107 | 108 | if init_sampler is None: 109 | init_sampler = initial_condition_ch(x) 110 | 111 | super().__init__( 112 | nstates=M, 113 | skewsymmetric_matrix=skewsymmetric_matrix, 114 | dissipation_matrix=dissipation_matrix, 115 | dissintegral=dissintegral, 116 | grad_dissintegral=dissintegral_grad, 117 | hess_hamiltonian=ham_hessian, 118 | hess_dissintegral=dissintegral_hessian, 119 | init_sampler=init_sampler, 120 | **kwargs 121 | ) 122 | 123 | self.skewsymmetric_matrix_flat = 0.5 / dx * np.array([[[-1, 0, 1]]]) 124 | self.dissipation_matrix_flat = -1 / dx**2 * np.array([[[1, -2, 1]]]) 125 | 126 | 127 | def initial_condition_ch(x=np.linspace(0, 1.0 - 1 / 100, 100)): 128 | """ 129 | Creates an initial condition sampler for the Cahn-Hilliard eqation. 130 | 131 | Parameters 132 | ---------- 133 | x : numpy.ndarray, optional 134 | Spatial grid on which to create the initial conditions. The default is 135 | an equidistant grid between 0 and .99 with step size 0.01. 136 | 137 | Returns 138 | ------- 139 | callable 140 | A function that takes a numpy random generator as input and returns an 141 | initial state on the spatial grid x. 142 | """ 143 | 144 | M = x.size 145 | P = (x[-1] - x[0]) * M / (M - 1) 146 | 147 | def sampler(rng): 148 | a1, a2 = rng.uniform(0.0, 0.05, 2) 149 | a3, a4 = rng.uniform(0.0, 0.2, 2) 150 | k1, k2, k3, k4 = rng.integers(1, 6, 4) 151 | u0 = 0 152 | u0 += a1 * np.cos(2 * k1 * np.pi / P * x) 153 | u0 += a2 * np.cos(2 * k2 * np.pi / P * x) 154 | u0 += a3 * np.sin(2 * k3 * np.pi / P * x) 155 | u0 += a4 * np.sin(2 * k4 * np.pi / P * x) 156 | return u0 157 | 158 | return sampler 159 | -------------------------------------------------------------------------------- /phlearn/phlearn/phsystems/pde/heat_system.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from scipy.sparse import spdiags 3 | 4 | from .pseudo_hamiltonian_pde_system import PseudoHamiltonianPDESystem 5 | 6 | __all__ = ['HeatEquationSystem', 'initial_condition_heat'] 7 | 8 | class HeatEquationSystem(PseudoHamiltonianPDESystem): 9 | """ 10 | Implements a discretization of the heat equation with an 11 | optional external force:: 12 | 13 | u_t - u_xx = f(u,t,x), 14 | 15 | on the pseudo-Hamiltonian formulation:: 16 | 17 | du/dt = -grad[V(u)] + F(u,t,x) 18 | 19 | where u is a vector of the system states at the spatial points given by 20 | x, and t is time. 21 | 22 | The system is by default integrated in time using the implicit midpoint method. 23 | The line 'self.sample_trajectory = self.sample_trajectory_midpoint' can be 24 | commented out to instead use the RK45 method of scipy. 25 | 26 | Parameters 27 | ---------- 28 | x : nparray, default np.linspace(0, 1.0 - 1/100, 100) 29 | The spatial discretization points. 30 | nu : number, default -1.0 31 | The parameter nu in the Cahn-Hilliard equation. 32 | alpha : number, default 1.0 33 | The parameter alpha in the Cahn-Hilliard equation. 34 | mu : number, default -0.001 35 | The parameter mu in the Cahn-Hilliard equation. 36 | init_sampler : callable, default None 37 | Function for sampling initial conditions. Callable taking 38 | a numpy random generator as input and returning an ndarray 39 | of shape same as x with initial conditions for the system. 40 | This sampler is used when calling 41 | CahnHilliardSystem.sample_trajectory if no initial 42 | condition is provided. 43 | kwargs : any, optional 44 | Keyword arguments that are passed to PseudoHamiltonianPDESystem 45 | constructor. 46 | 47 | """ 48 | 49 | 50 | def __init__( 51 | self, 52 | x=np.linspace(0, 6.0 - 6 / 300, 300), 53 | nu=1.0, 54 | init_sampler=None, 55 | **kwargs 56 | ): 57 | M = x.size 58 | dx = x[-1] / (M - 1) 59 | e = np.ones(M) 60 | # Forward difference matrix: 61 | Dp = ( 62 | 1 / dx * spdiags([e, -e, e], np.array([-M + 1, 0, 1]), M, M).toarray() 63 | ) 64 | # Central difference matrix: 65 | D1 = ( 66 | 0.5 67 | / dx 68 | * spdiags([e, -e, e, -e], np.array([-M + 1, -1, 1, M - 1]), M, M).toarray() 69 | ) 70 | # 2nd order central difference matrix: 71 | D2 = ( 72 | 1 73 | / dx**2 74 | * spdiags( 75 | [e, e, -2 * e, e, e], np.array([-M + 1, -1, 0, 1, M - 1]), M, M 76 | ).toarray() 77 | ) 78 | skewsymmetric_matrix = D1 79 | self.x = x 80 | self.M = M 81 | self.D2 = D2 82 | 83 | def ham(u): 84 | return np.sum(np.zeros_like(u), axis=1) 85 | 86 | def dissintegral(u): 87 | return np.sum(0.5 * nu * (np.matmul(Dp, u.T) ** 2).T, axis=1) 88 | 89 | def ham_grad(u): 90 | return np.zeros_like(u) 91 | 92 | def dissintegral_grad(u): 93 | return -nu * u @ D2 94 | 95 | def ham_hessian(u): 96 | return np.zeros_like(D1) 97 | 98 | def dissintegral_hessian(u): 99 | return -nu * D2 100 | 101 | if init_sampler is None: 102 | init_sampler = initial_condition_heat(x) 103 | 104 | super().__init__( 105 | nstates=M, 106 | skewsymmetric_matrix=skewsymmetric_matrix, 107 | hamiltonian=ham, 108 | dissintegral=dissintegral, 109 | grad_hamiltonian=ham_grad, 110 | grad_dissintegral=dissintegral_grad, 111 | hess_hamiltonian=ham_hessian, 112 | hess_dissintegral=dissintegral_hessian, 113 | init_sampler=init_sampler, 114 | **kwargs 115 | ) 116 | 117 | self.skewsymmetric_matrix_flat = 0.5 / dx * np.array([[[-1, 0, 1]]]) 118 | 119 | 120 | def initial_condition_heat(x=np.linspace(0, 6.0 - 6 / 300, 300)): 121 | """ 122 | Creates an initial condition sampler for the heat eqation. 123 | 124 | Parameters 125 | ---------- 126 | x : numpy.ndarray, optional 127 | Spatial grid on which to create the initial conditions. The default is 128 | an equidistant grid between 0 and 5.98 with step size 0.02. 129 | 130 | Returns 131 | ------- 132 | callable 133 | A function that takes a numpy random generator as input and returns an 134 | initial state on the spatial grid x. 135 | """ 136 | 137 | M = x.size 138 | P = (x[-1] - x[0]) * M / (M - 1) 139 | 140 | def sampler(rng): 141 | d1, d2 = rng.uniform(0.3, 3, 2) 142 | c1, c2 = rng.uniform(0.5, 1.5, 2) 143 | k1 = rng.uniform(0.5, 3.0, 1) 144 | k2 = rng.uniform(10.0, 20.0, 1) 145 | n1 = rng.uniform(20.0, 40.0, 1) 146 | n2 = rng.uniform(0.05, 0.15, 1) 147 | u0 = 0 148 | u0 += ( 149 | rng.uniform(-5.0, 5.0, 1) 150 | - c1 * np.tanh(n1 * (x - d1)) 151 | + c1 * np.tanh(n1 * (x - P + d1)) 152 | ) 153 | u0 += -c2 * np.tanh(n1 * (x - d2)) + c2 * np.tanh(n1 * (x - P + d2)) 154 | u0 += n2 * np.sin(k1 * np.pi * x) ** 2 * np.sin(k2 * np.pi * x) 155 | return u0 156 | 157 | return sampler -------------------------------------------------------------------------------- /phlearn/phlearn/phsystems/pde/kdv_system.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from scipy.sparse import spdiags 3 | 4 | from .pseudo_hamiltonian_pde_system import PseudoHamiltonianPDESystem 5 | 6 | __all__ = ['KdVSystem', 'initial_condition_kdv'] 7 | 8 | class KdVSystem(PseudoHamiltonianPDESystem): 9 | """ 10 | Implements the a discretization of the KdV-Burgers equation with an 11 | optional external force:: 12 | 13 | u_t + eta u u_x - nu u_xx - gamma^2 u_xxx = f(u,x,t) 14 | 15 | on the pseudo-Hamiltonian formulation:: 16 | 17 | u/dt = d/dx(grad[H(u)]) - grad[V(u)] + F(u, t, x) 18 | 19 | where u is a vector of the system states at the spatial points given by 20 | x, and t is time. 21 | 22 | The system is by default integrated in time using the implicit midpoint 23 | method. The line 'self.sample_trajectory = self.sample_trajectory_midpoint' 24 | can be commented out to instead use the RK45 method of scipy. 25 | 26 | Parameters 27 | ---------- 28 | x : nparray, default np.linspace(0, 20.0 - 0.2, 100) 29 | The spatial discretization points 30 | eta : number, default 6.0 31 | The parameter eta in the KdV equation 32 | gamma : number, default 1.0 33 | The parameter gamma in the KdV equation 34 | nu : number, default 0.0 35 | The damping coefficient 36 | init_sampler : callable, default None 37 | Function for sampling initial conditions. Callabale taking 38 | a numpy random generator as input and returning an ndarray 39 | of shape same as x with inital conditions for the system. 40 | This sampler is used when calling 41 | KdVSystem.sample_trajectory if no initial 42 | condition is provided. 43 | kwargs : any, optional 44 | Keyword arguments that are passed to PseudoHamiltonianPDESystem 45 | constructor. 46 | 47 | """ 48 | 49 | def __init__( 50 | self, 51 | x=np.linspace(0, 20.0 - 0.2, 100), 52 | eta=6.0, 53 | gamma=1.0, 54 | nu=0.0, 55 | init_sampler=None, 56 | **kwargs 57 | ): 58 | M = x.size 59 | dx = x[-1] / (M - 1) 60 | e = np.ones(M) 61 | # Forward difference matrix: 62 | Dp = 1 / dx * spdiags([e, -e, e], np.array([-M + 1, 0, 1]), M, M).toarray() 63 | # Central difference matrix: 64 | D1 = ( 65 | 0.5 66 | / dx 67 | * spdiags([e, -e, e, -e], np.array([-M + 1, -1, 1, M - 1]), M, M).toarray() 68 | ) 69 | # 2nd order central difference matrix: 70 | D2 = ( 71 | 1 72 | / dx**2 73 | * spdiags( 74 | [e, e, -2 * e, e, e], np.array([-M + 1, -1, 0, 1, M - 1]), M, M 75 | ).toarray() 76 | ) 77 | skewsymmetric_matrix = D1 78 | self.x = x 79 | self.M = M 80 | self.eta = eta 81 | self.gamma = gamma 82 | self.nu = nu 83 | self.D2 = D2 84 | self.sample_trajectory = self.sample_trajectory_midpoint 85 | 86 | def ham(u): 87 | return np.sum( 88 | -1 / 6 * eta * u**3 89 | + (0.5 * gamma**2 * (np.matmul(Dp, u.T)) ** 2).T, 90 | axis=-1, 91 | ) 92 | 93 | def dissintegral(u): 94 | return np.sum(0.5 * nu * (np.matmul(Dp, u.T) ** 2).T, axis=-1) 95 | 96 | def ham_grad(u): 97 | return -0.5 * eta * u**2 - (gamma**2 * u @ D2) 98 | 99 | def dissintegral_grad(u): 100 | return -nu * u @ D2 101 | 102 | def ham_hessian(u): 103 | return -eta * np.diag(u) - gamma**2 * D2 104 | 105 | def dissintegral_hessian(u): 106 | return -nu * D2 107 | 108 | if init_sampler is None: 109 | init_sampler = initial_condition_kdv(x, eta) 110 | 111 | super().__init__( 112 | nstates=M, 113 | skewsymmetric_matrix=skewsymmetric_matrix, 114 | hamiltonian=ham, 115 | dissintegral=dissintegral, 116 | grad_hamiltonian=ham_grad, 117 | grad_dissintegral=dissintegral_grad, 118 | hess_hamiltonian=ham_hessian, 119 | hess_dissintegral=dissintegral_hessian, 120 | init_sampler=init_sampler, 121 | **kwargs 122 | ) 123 | 124 | self.skewsymmetric_matrix_flat = 0.5 / dx * np.array([[[-1, 0, 1]]]) 125 | 126 | 127 | def initial_condition_kdv(x=np.linspace(0, 20.0 - 0.2, 100), eta=6.0): 128 | """ 129 | Creates an initial condition sampler for the KdV-Burgers equation. 130 | 131 | Parameters 132 | ---------- 133 | x : numpy.ndarray, optional 134 | Spatial grid on which to create the initial conditions. The default is 135 | an equidistant grid between 0 and 19.8 with step size 0.2. 136 | eta : float, optional 137 | The parameter eta in the KdV-Burgers equation. The default is 6.0. 138 | 139 | Returns 140 | ------- 141 | callable 142 | A function that takes a numpy random generator as input and returns an 143 | initial state on the spatial grid x. 144 | """ 145 | 146 | M = x.size 147 | P = (x[-1] - x[0]) * M / (M - 1) 148 | 149 | def sech(a): 150 | return 1 / np.cosh(a) 151 | 152 | def sampler(rng): 153 | k1, k2 = rng.uniform(0.5, 2.0, 2) 154 | d1, d2 = rng.uniform(0.0, 1.0, 1), rng.uniform(0.0, 1.0, 1) 155 | u0 = 0 156 | u0 += ( 157 | (-6.0 / -eta) 158 | * 2 159 | * k1**2 160 | * sech(k1 * ((x + P / 2 - P * d1) % P - P / 2)) ** 2 161 | ) 162 | u0 += ( 163 | (-6.0 / -eta) 164 | * 2 165 | * k2**2 166 | * sech(k2 * ((x + P / 2 - P * d2) % P - P / 2)) ** 2 167 | ) 168 | u0 = np.concatenate([u0[M:], u0[:M]], axis=-1) 169 | return u0 170 | 171 | return sampler 172 | -------------------------------------------------------------------------------- /phlearn/phlearn/phsystems/pde/perona_malik_system.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from scipy.sparse import spdiags 3 | 4 | from .pseudo_hamiltonian_pde_system import PseudoHamiltonianPDESystem 5 | 6 | __all__ = ['PeronaMalikSystem', 'initial_condition_pm'] 7 | 8 | class PeronaMalikSystem(PseudoHamiltonianPDESystem): 9 | """ 10 | Implements a discretization of the Perona-Malik equation with an 11 | optional external force:: 12 | 13 | u_t + (u_x/(1+u_x^2))_x = f(u,t,x), 14 | 15 | on the pseudo-Hamiltonian formulation:: 16 | 17 | du/dt = -grad[V(u)] + F(u,t,x) 18 | 19 | where u is a vector of the system states at the spatial points given by 20 | x, and t is time. 21 | 22 | The system is by default integrated in time using the implicit midpoint method. 23 | The line 'self.sample_trajectory = self.sample_trajectory_midpoint' can be 24 | commented out to instead use the RK45 method of scipy. 25 | 26 | Parameters 27 | ---------- 28 | x : nparray, default np.linspace(0, 1.0 - 1/100, 100) 29 | The spatial discretization points. 30 | nu : number, default -1.0 31 | The parameter nu in the Cahn-Hilliard equation. 32 | alpha : number, default 1.0 33 | The parameter alpha in the Cahn-Hilliard equation. 34 | mu : number, default -0.001 35 | The parameter mu in the Cahn-Hilliard equation. 36 | init_sampler : callable, default None 37 | Function for sampling initial conditions. Callable taking 38 | a numpy random generator as input and returning an ndarray 39 | of shape same as x with initial conditions for the system. 40 | This sampler is used when calling 41 | CahnHilliardSystem.sample_trajectory if no initial 42 | condition is provided. 43 | kwargs : any, optional 44 | Keyword arguments that are passed to PseudoHamiltonianPDESystem 45 | constructor. 46 | 47 | """ 48 | 49 | def __init__( 50 | self, 51 | x=np.linspace(0, 6.0 - 6 / 300, 300), 52 | init_sampler=None, 53 | **kwargs 54 | ): 55 | M = x.size 56 | dx = x[-1] / (M - 1) 57 | e = np.ones(M) 58 | # Forward difference matrix: 59 | Dp = ( 60 | 1 / dx * spdiags([e, -e, e], np.array([-M + 1, 0, 1]), M, M).toarray() 61 | ) 62 | # Central difference matrix: 63 | D1 = ( 64 | 0.5 65 | / dx 66 | * spdiags([e, -e, e, -e], np.array([-M + 1, -1, 1, M - 1]), M, M).toarray() 67 | ) 68 | # 2nd order central difference matrix: 69 | D2 = ( 70 | 1 71 | / dx**2 72 | * spdiags( 73 | [e, e, -2 * e, e, e], np.array([-M + 1, -1, 0, 1, M - 1]), M, M 74 | ).toarray() 75 | ) 76 | skewsymmetric_matrix = D1 77 | self.x = x 78 | self.M = M 79 | self.D2 = D2 80 | 81 | def ham(u): 82 | return np.sum(np.zeros_like(u), axis=1) 83 | 84 | def dissintegral(u): 85 | return 1 / 2 * np.sum(np.log(1 + np.matmul(Dp, u.T) ** 2).T, axis=1) 86 | 87 | def ham_grad(u): 88 | return np.zeros_like(u) 89 | 90 | def dissintegral_grad(u): 91 | return -((u @ D1) / (1 + (u @ D1) ** 2)) @ D1 92 | 93 | def ham_hessian(u): 94 | return np.zeros_like(D1) 95 | 96 | # We provide a simplified approximation of the exact Hessian, 97 | # to get a more efficient pseudo-Newton solver for the implicit 98 | # midpoint method, if that was used: 99 | def dissintegral_hessian(u): 100 | return np.matmul(np.matmul(D1, 1 / (1 + np.matmul(D1, u) ** 2)), D1) 101 | 102 | if init_sampler is None: 103 | init_sampler = initial_condition_pm(x) 104 | 105 | super().__init__( 106 | nstates=M, 107 | skewsymmetric_matrix=skewsymmetric_matrix, 108 | hamiltonian=ham, 109 | dissintegral=dissintegral, 110 | grad_hamiltonian=ham_grad, 111 | grad_dissintegral=dissintegral_grad, 112 | hess_hamiltonian=ham_hessian, 113 | hess_dissintegral=dissintegral_hessian, 114 | init_sampler=init_sampler, 115 | **kwargs 116 | ) 117 | 118 | self.skewsymmetric_matrix_flat = 0.5 / dx * np.array([[[-1, 0, 1]]]) 119 | 120 | def initial_condition_pm(x=np.linspace(0, 6.0 - 6 / 300, 300)): 121 | """ 122 | Creates an initial condition sampler for the Perona-Malik eqation. 123 | 124 | Parameters 125 | ---------- 126 | x : numpy.ndarray, optional 127 | Spatial grid on which to create the initial conditions. The default is 128 | an equidistant grid between 0 and 5.98 with step size 0.02. 129 | 130 | Returns 131 | ------- 132 | callable 133 | A function that takes a numpy random generator as input and returns an 134 | initial state on the spatial grid x. 135 | """ 136 | 137 | M = x.size 138 | P = (x[-1] - x[0]) * M / (M - 1) 139 | 140 | def sampler(rng): 141 | d1, d2 = rng.uniform(0.3, 3, 2) 142 | c1, c2 = rng.uniform(0.5, 1.5, 2) 143 | k1 = rng.uniform(0.5, 3.0, 1) 144 | k2 = rng.uniform(10.0, 20.0, 1) 145 | n1 = rng.uniform(20.0, 40.0, 1) 146 | n2 = rng.uniform(0.05, 0.15, 1) 147 | u0 = 0 148 | u0 += ( 149 | rng.uniform(-5.0, 5.0, 1) 150 | - c1 * np.tanh(n1 * (x - d1)) 151 | + c1 * np.tanh(n1 * (x - P + d1)) 152 | ) 153 | u0 += -c2 * np.tanh(n1 * (x - d2)) + c2 * np.tanh(n1 * (x - P + d2)) 154 | u0 += n2 * np.sin(k1 * np.pi * x) ** 2 * np.sin(k2 * np.pi * x) 155 | return u0 156 | 157 | return sampler 158 | -------------------------------------------------------------------------------- /phlearn/phlearn/phsystems/pde/pseudo_hamiltonian_pde_system.py: -------------------------------------------------------------------------------- 1 | import autograd.numpy as np 2 | import autograd 3 | from scipy.integrate import solve_ivp 4 | 5 | from ...utils.derivatives import time_derivative 6 | from ...utils.utils import midpoint_method 7 | 8 | __all__ = ['PseudoHamiltonianPDESystem', 'zero_force'] 9 | 10 | class PseudoHamiltonianPDESystem: 11 | """ 12 | Implements a spatially discretized pseudo-Hamiltonian PDE system of the form:: 13 | 14 | dx/dt = S*grad[H(x)] - R*grad[V(x)] + F(x, t, xspatial) 15 | 16 | where x is the system state, A is a symmetric matrix, S is a skew-symmetric matrix, 17 | R is the symmetric dissipation matrix, H and V are discretized integrals of the systemm, 18 | F is the external force depending on state, time and space. 19 | 20 | What is x here is usually u in the literature, and xspatial is x. We use x for 21 | the state to be consistent with the ODE case. 22 | 23 | Parameters 24 | ---------- 25 | nstates : int 26 | Number of system states N. 27 | 28 | skewsymmetric_matrix : (N, N) ndarray or callable, default None 29 | Corresponds to the S matrix. Must either be an 30 | ndarray, or callable taking an ndarray 31 | input of shape (nsamples, nstates) and returning an ndarray 32 | of shape (nsamples, nstates, nstates). If None, 33 | the system is assumed to be canonical, and the 34 | S matrix is set ot the skew-symmetric matrix 35 | [[0, I_n], [-I_n, 0]]. 36 | 37 | hamiltonian : callable, default None 38 | The Hamiltonian H of the system. Callable taking a 39 | torch tensor input of shape (nsamples, nstates) and 40 | returning a torch tensor of shape (nsamples, 1). 41 | If the gradient of the Hamiltonian is not provided, 42 | the gradient of this function will be computed by torch and 43 | used instead. If this is not provided, the grad_hamiltonian 44 | must be provided. 45 | 46 | dissintegral : callable, default None 47 | The dissipating integral V of the system. Callable taking a 48 | torch tensor input of shape (nsamples, nstates) and 49 | returning a torch tensor of shape (nsamples, 1). 50 | If the gradient of the dissipating integral is not provided, 51 | the gradient of this function will be computed by torch and 52 | used instead. If this is not provided, the grad_dissintegral 53 | must be provided. 54 | 55 | grad_hamiltonian : callable, default None 56 | The gradient of the Hamiltonian H of the system. Callable 57 | taking an ndarray input of shape (nsamples, nstates) and 58 | returning a torch tensor of shape (nsamples, nstates). 59 | If this is not provided, the hamiltonian must be provided. 60 | 61 | grad_dissintegral : callable, default None 62 | The gradient of the dissipating integral V of the system. Callable 63 | taking an ndarray input of shape (nsamples, nstates) and 64 | returning a torch tensor of shape (nsamples, nstates). 65 | If this is not provided, the hamiltonian must be provided. 66 | 67 | external_forces : callable, default None 68 | The external forces affecting system. Callable taking two 69 | ndarrays as input, x and t, of shape (nsamples, nstates), 70 | (nsamples, 1), respectively and returning an ndarray of 71 | shape (nsamples, nstates). 72 | 73 | controller : phlearn.control.PseudoHamiltonianController, 74 | default None 75 | Additional external forces set by a controller. Callable 76 | taking an ndarray x of shape (nstates,) and a scalar t as 77 | input and returning an ndarray of shape (nstates,). Note 78 | that this function should not take batch inputs, and that 79 | when calling PseudoHamiltonianSystem.sample_trajectory when a 80 | controller is provided, the Runge-Kutta 4 method will be 81 | used for integration in favor of Scipy's solve_ivp. 82 | 83 | init_sampler : callable, default None 84 | Function for sampling initial conditions. Callabale taking 85 | a numpy random generator as input and returning an ndarray 86 | of shape (nstates,) with inital conditions for the system. 87 | This sampler is used when calling 88 | PseudoHamiltonianSystem.sample_trajectory if no initial 89 | condition is provided. 90 | 91 | """ 92 | 93 | def __init__( 94 | self, 95 | nstates, 96 | lhs_matrix=None, 97 | skewsymmetric_matrix=None, 98 | dissipation_matrix=None, 99 | hamiltonian=None, 100 | dissintegral=None, 101 | grad_hamiltonian=None, 102 | grad_dissintegral=None, 103 | hess_hamiltonian=None, 104 | hess_dissintegral=None, 105 | external_forces=None, 106 | jac_external_forces=None, 107 | controller=None, 108 | init_sampler=None, 109 | ): 110 | self.nstates = nstates 111 | 112 | self.lhs_matrix_provided = True 113 | if lhs_matrix is None: 114 | lhs_matrix = np.eye(nstates) 115 | self.lhs_matrix_provided = False 116 | 117 | self.lhs_matrix = lhs_matrix 118 | self.A = lambda x: lhs_matrix 119 | 120 | if skewsymmetric_matrix is None: 121 | npos = nstates // 2 122 | skewsymmetric_matrix = np.block( 123 | [ 124 | [np.zeros([npos, npos]), np.eye(npos)], 125 | [-np.eye(npos), np.zeros([npos, npos])], 126 | ] 127 | ) 128 | 129 | if not callable(skewsymmetric_matrix): 130 | self.skewsymmetric_matrix = skewsymmetric_matrix 131 | self.S = lambda x: skewsymmetric_matrix 132 | else: 133 | self.skewsymmetric_matrix = None 134 | self.S = skewsymmetric_matrix 135 | 136 | if dissipation_matrix is None: 137 | dissipation_matrix = np.eye(nstates) 138 | 139 | self.dissipation_matrix = dissipation_matrix 140 | self.R = lambda x: dissipation_matrix 141 | 142 | self.skewsymmetric_matrix_flat = np.array([[[-1, 0, 1]]]) 143 | self.dissipation_matrix_flat = np.array([[[1]]]) 144 | self.lhs_matrix_flat = np.array([[[1]]]) 145 | 146 | self.H = hamiltonian 147 | self.dH = grad_hamiltonian 148 | self.ddH = hess_hamiltonian 149 | if grad_hamiltonian is None: 150 | self.dH = self._dH 151 | if hess_hamiltonian is None: 152 | self.ddH = self._ddH 153 | 154 | self.V = dissintegral 155 | self.dV = grad_dissintegral 156 | self.ddV = hess_dissintegral 157 | if grad_dissintegral is None: 158 | self.dV = self._dV 159 | if hess_dissintegral is None: 160 | self.ddV = self._ddV 161 | 162 | self.controller = controller 163 | 164 | self.external_forces = external_forces 165 | self.external_forces_jacobian = jac_external_forces 166 | if external_forces is None: 167 | self.external_forces = zero_force 168 | self.external_forces_jacobian = None 169 | elif jac_external_forces is None: 170 | self.external_forces_jacobian = self._jacforce 171 | 172 | if init_sampler is not None: 173 | self._initial_condition_sampler = init_sampler 174 | 175 | self.seed(None) 176 | 177 | def seed(self, seed): 178 | """ 179 | Set the internal random state. 180 | 181 | Parameters 182 | ---------- 183 | seed : int 184 | 185 | """ 186 | 187 | self.rng = np.random.default_rng(seed) 188 | 189 | def time_derivative(self, integrator, *args, **kwargs): 190 | """ 191 | See :py:meth:~`utils.derivatives.time_derivative` 192 | """ 193 | return time_derivative(integrator, self.x_dot, *args, **kwargs) 194 | 195 | def x_dot(self, x, t, u=None): 196 | """ 197 | Computes the time derivative, the right hand side of the pseudo- 198 | Hamiltonian equation. 199 | 200 | Parameters 201 | ---------- 202 | x : (..., N) ndarray 203 | t : (..., 1) ndarray 204 | u : (..., N) ndarray or None, default None 205 | 206 | Returns 207 | ------- 208 | (..., N) ndarray 209 | 210 | """ 211 | 212 | S = self.S(x) 213 | R = self.R(x) 214 | if self.H is not None: 215 | dH = self.dH(x) 216 | else: 217 | dH = np.zeros_like(x) 218 | if self.V is not None: 219 | dV = self.dV(x) 220 | else: 221 | dV = np.zeros_like(x) 222 | if len(S.shape) == 3: 223 | dynamics = ( 224 | np.matmul(S, np.atleast_3d(dH)) - np.matmul(R, np.atleast_3d(dV)) 225 | ).reshape(x.shape) + self.external_forces(x, t) 226 | else: 227 | dynamics = dH @ (S.T) - dV @ R + self.external_forces(x, t) 228 | if u is not None: 229 | dynamics += u 230 | return dynamics 231 | 232 | def x_dot_jacobian(self, x, t, u=None): 233 | """ 234 | Computes the Jacobian of the right hand side of the pseudo- 235 | Hamiltonian equation. 236 | 237 | Parameters 238 | ---------- 239 | x : (..., N) ndarray 240 | t : (..., 1) ndarray 241 | u : (..., N) ndarray or None, default None 242 | 243 | Returns 244 | ------- 245 | (..., N) ndarray 246 | 247 | """ 248 | 249 | S = self.S(x) 250 | R = self.R(x) 251 | jacobian = np.zeros_like(S) 252 | if self.H is not None: 253 | ddH = self.ddH(x) 254 | jacobian += np.matmul(S, ddH) 255 | if self.V is not None: 256 | ddV = self.ddV(x) 257 | jacobian -= np.matmul(R, ddV) 258 | if self.external_forces_jacobian is not None: 259 | jacobian += self.external_forces_jacobian(x, t) 260 | return jacobian 261 | 262 | def sample_trajectory(self, t, x0=None, noise_std=0, reference=None): 263 | """ 264 | Samples a trajectory of the system at times *t*, found by using the 265 | solve_ivp solver for temporal integration, starting from the initial 266 | state x0. 267 | 268 | Parameters 269 | ---------- 270 | t : (T, 1) ndarray 271 | Times at which the trajectory is sampled. 272 | x0 : (N,) ndarray, default None 273 | Initial condition. 274 | noise_std : number, default 0. 275 | Standard deviation of Gaussian white noise added to the 276 | samples of the trajectory. 277 | reference : phlearn.control.Reference, default None 278 | Not in use for now. 279 | 280 | Returns 281 | ------- 282 | x : (T, N) ndarray 283 | dxdt : (T, N) ndarray 284 | t : (T, 1) ndarray 285 | us : (T, N) ndarray 286 | 287 | """ 288 | 289 | if x0 is None: 290 | x0 = self._initial_condition_sampler(self.rng) 291 | 292 | if self.lhs_matrix_provided: 293 | lhs_matrix_inv = np.linalg.inv(self.lhs_matrix) 294 | x_dot = lambda t, x: np.matmul( 295 | lhs_matrix_inv, 296 | self.x_dot(x.reshape(1, x.shape[-1]), np.array(t).reshape((1, 1))).T, 297 | ).T 298 | else: 299 | x_dot = lambda t, x: self.x_dot( 300 | x.reshape(1, x.shape[-1]), np.array(t).reshape((1, 1)) 301 | ) 302 | out_ivp = solve_ivp( 303 | fun=x_dot, t_span=(t[0], t[-1]), y0=x0, t_eval=t, rtol=1e-10 304 | ) 305 | x, t = out_ivp["y"].T, out_ivp["t"].T 306 | dxdt = self.x_dot(x, t) 307 | us = None 308 | 309 | # Add noise: 310 | x += self.rng.normal(size=x.shape) * noise_std 311 | dxdt += self.rng.normal(size=dxdt.shape) * noise_std 312 | 313 | return x, dxdt, t, us 314 | 315 | def sample_trajectory_midpoint(self, t, x0=None, noise_std=0, reference=None): 316 | """ 317 | Samples a trajectory of the system at times *t*, found by using the 318 | implicit midpoint method for temporal integration, starting from the 319 | initial state x0. Newton's method is used for solving the system of 320 | nonlinear equations at each integration step. 321 | 322 | Parameters 323 | ---------- 324 | t : (T, 1) ndarray 325 | Times at which the trajectory is sampled. 326 | x0 : (N,) ndarray, default None 327 | Initial condition. 328 | noise_std : number, default 0. 329 | Standard deviation of Gaussian white noise added to the 330 | samples of the trajectory. 331 | reference : phlearn.control.Reference, default None 332 | Not in use for now. 333 | 334 | Returns 335 | ------- 336 | x : (T, N) ndarray 337 | dxdt : (T, N) ndarray 338 | t : (T, 1) ndarray 339 | us : (T, N) ndarray 340 | 341 | """ 342 | 343 | if x0 is None: 344 | x0 = self._initial_condition_sampler(self.rng) 345 | 346 | x = np.zeros([t.shape[0], x0.shape[-1]]) 347 | dxdt = np.zeros_like(x) 348 | us = np.zeros([t.shape[0] - 1, x0.shape[-1]]) 349 | x[0, :] = x0 350 | 351 | M = x0.shape[-1] 352 | if self.lhs_matrix_provided: 353 | f = lambda u, t: np.linalg.solve(self.lhs_matrix, self.x_dot(u, t)) 354 | Df = lambda u, t: np.linalg.solve( 355 | self.lhs_matrix, self.x_dot_jacobian(u, t) 356 | ) 357 | else: 358 | f = lambda u, t: self.x_dot(u, t) 359 | Df = lambda u, t: self.x_dot_jacobian(u, t) 360 | for i, t_step in enumerate(t[:-1]): 361 | dt = t[i + 1] - t[i] 362 | dxdt[i, :] = f(x[i, :], t[i]) 363 | x[i + 1, :] = midpoint_method( 364 | x[i, :], x[i, :], t[i], f, Df, dt, M, 1e-12, 5 365 | ) 366 | 367 | # Add noise: 368 | x += self.rng.normal(size=x.shape) * noise_std 369 | dxdt += self.rng.normal(size=dxdt.shape) * noise_std 370 | 371 | return x, dxdt, t, us 372 | 373 | def _dH(self, x): 374 | H = lambda x: self.H(x).sum() 375 | return autograd.grad(H)(x) 376 | 377 | def _dV(self, x): 378 | V = lambda x: self.V(x).sum() 379 | return autograd.grad(V)(x) 380 | 381 | def _ddH(self, x): 382 | H = lambda x: self.H(x).sum() 383 | return autograd.hessian(H)(x) 384 | 385 | def _ddV(self, x): 386 | H = lambda x: self.H(x).sum() 387 | return autograd.hessian(H)(x) 388 | 389 | def _jacforce(self, x, t): 390 | external_forces_x = lambda x: self.external_forces(x, t) 391 | return autograd.jacobian(external_forces_x)(x) 392 | 393 | def _initial_condition_sampler(self, rng=None): 394 | if rng is None: 395 | assert self.rng is not None 396 | rng = self.rng 397 | return rng.uniform(low=-1.0, high=1.0, size=self.nstates) 398 | 399 | 400 | def zero_force(x, t=None): 401 | """ 402 | A force term that is always zero. 403 | 404 | Parameters 405 | ---------- 406 | x : (..., N) ndarray 407 | t : (..., 1) ndarray, default None 408 | 409 | Returns 410 | ------- 411 | (..., N) ndarray 412 | All zero ndarray 413 | 414 | """ 415 | 416 | return np.zeros_like(x) 417 | -------------------------------------------------------------------------------- /phlearn/phlearn/phsystems/tests/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SINTEF/pseudo-hamiltonian-neural-networks/091065ef3c1b730d56fd413b6373d0424d8114be/phlearn/phlearn/phsystems/tests/__init__.py -------------------------------------------------------------------------------- /phlearn/phlearn/phsystems/tests/test_phsystems.py: -------------------------------------------------------------------------------- 1 | # TODO: Test case when len(S.shape) == 3 and len(R.shape) == 3 2 | # TODO: Have ignored controller for now. I.e., assumed controller = None throughout 3 | 4 | # When developing tests, from the top level dir (i.e., phlearn) run 5 | # python3 -m phlearn.phnns.tests.test_phsystems 6 | # to run this as a module, allowing use of relative imports 7 | 8 | import numpy as np 9 | from ..ode.pseudo_hamiltonian_system import PseudoHamiltonianSystem 10 | import torch 11 | 12 | 13 | N_STATES = 10 14 | N_TIMESTEPS = 100 15 | 16 | X_RAND = np.random.rand(N_STATES) 17 | T_AXIS = np.linspace(0,10, N_TIMESTEPS) 18 | EXTERNAL_FORCE = np.random.rand(N_STATES) 19 | DISSIPATION_MATRIX = np.random.rand(N_STATES, N_STATES) 20 | DISSIPATION_MATRIX = DISSIPATION_MATRIX + DISSIPATION_MATRIX.T 21 | STRUCTURE_MATRIX = np.random.rand(N_STATES, N_STATES) 22 | STRUCTURE_MATRIX = STRUCTURE_MATRIX - STRUCTURE_MATRIX.T 23 | 24 | 25 | def H(x): 26 | x = x.flatten() if torch.is_tensor(x) else x 27 | return sum(x[i] ** 2 for i in range(N_STATES)) 28 | 29 | 30 | def dH(x): 31 | return np.array([2 * x[i] for i in range(N_STATES)]) 32 | 33 | 34 | psh_kwargs = dict( 35 | nstates=N_STATES, 36 | hamiltonian=H, 37 | ) 38 | 39 | 40 | def test_R_is_zero_matrix_when_dissipation_matrix_is_None(): 41 | phs = PseudoHamiltonianSystem(**psh_kwargs) 42 | assert not phs.R( 43 | np.random.rand(psh_kwargs["nstates"]) 44 | ).any(), "Dissipation_matrix not defaulting to zero matrix" 45 | 46 | 47 | def test_S_is_canonical_when_structure_matrix_is_None(): 48 | phs = PseudoHamiltonianSystem(**psh_kwargs) 49 | m = int(N_STATES / 2) 50 | O, I = np.zeros([m, m]), np.eye(m) 51 | S = np.block([[O, I], [-I, O]]) 52 | assert np.array_equal( 53 | S, phs.S(X_RAND) 54 | ), "strucutre_matrix not defulating to canonical form" 55 | 56 | 57 | def test_dissipation_matrix_is_returned(): 58 | phs = PseudoHamiltonianSystem(dissipation_matrix=DISSIPATION_MATRIX, **psh_kwargs) 59 | assert np.array_equal( 60 | DISSIPATION_MATRIX, phs.R(X_RAND) 61 | ), "dissipation_matrix not returned" 62 | 63 | 64 | def test_structure_matrix_is_returned(): 65 | psh_kwargs_loc = psh_kwargs 66 | n = psh_kwargs["nstates"] 67 | phs = PseudoHamiltonianSystem(skewsymmetric_matrix=STRUCTURE_MATRIX, **psh_kwargs_loc) 68 | assert np.array_equal( 69 | STRUCTURE_MATRIX, phs.S(X_RAND) 70 | ), "skewsymmetric_matrix not returned" 71 | 72 | 73 | def test_dH_calculated_correctly(): 74 | phs = PseudoHamiltonianSystem(**psh_kwargs) 75 | assert np.allclose(phs.dH(X_RAND), 2 * X_RAND), "grad_H computed incorrectly" 76 | 77 | 78 | def test_x_dot(): 79 | phs = PseudoHamiltonianSystem( 80 | dissipation_matrix=lambda x: DISSIPATION_MATRIX, 81 | skewsymmetric_matrix=lambda x: STRUCTURE_MATRIX, 82 | external_forces=lambda x, t: EXTERNAL_FORCE, 83 | **psh_kwargs 84 | ) 85 | x_dot = dH(X_RAND) @ (STRUCTURE_MATRIX.T - DISSIPATION_MATRIX.T) + EXTERNAL_FORCE 86 | assert np.allclose(phs.x_dot(X_RAND, T_AXIS), x_dot), "x_dot() returns incorrect ODE" 87 | 88 | 89 | def test_sample_trajectory_on_const_ode(): 90 | phs = PseudoHamiltonianSystem(skewsymmetric_matrix=0 * STRUCTURE_MATRIX, **psh_kwargs) 91 | x, _, _, _ = phs.sample_trajectory(t=[0, 1], x0=X_RAND, noise_std=0) 92 | assert np.allclose( 93 | x[-1], X_RAND 94 | ), "Solution of a constant ODE does not remain constant" 95 | -------------------------------------------------------------------------------- /phlearn/phlearn/utils/__init__.py: -------------------------------------------------------------------------------- 1 | """ 2 | The utils subpackage contains functionality that are handy when 3 | using the other phlearn subpackages. 4 | 5 | Functions present in phlearn.utils 6 | ------------------------------------------- 7 | 8 | :py:meth:`~.derivatives.time_derivative` 9 | 10 | :py:meth:`~.utils.to_tensor` 11 | 12 | """ 13 | 14 | from . import derivatives 15 | from .derivatives import * 16 | from . import utils 17 | from .utils import * 18 | 19 | __all__ = derivatives.__all__.copy() 20 | __all__ += utils.__all__.copy() -------------------------------------------------------------------------------- /phlearn/phlearn/utils/derivatives.py: -------------------------------------------------------------------------------- 1 | 2 | import numpy as np 3 | 4 | __all__ = ['time_derivative'] 5 | 6 | def time_derivative(integrator, x_dot, x_start, x_end, 7 | t_start, t_end, dt, u=None, xspatial=None): 8 | """ 9 | Computes the time derivative of x using the provided function *x_dot*. 10 | 11 | Parameters 12 | ---------- 13 | integrator : str or bool 14 | If 'euler' or False, the time derivative at *x_start*, *t_start* 15 | is computed. If 'midpoint', *x_start*, *x_end*, *t_start*, 16 | *t_end* are used to compute the discretized implricit midpoint 17 | estimate of the derivative. If 'rk4', *x_start*, *t_start*, *dt* 18 | are used to compute the explicit Runge-Kutta 4 estimate. 19 | If 'srk4', *x_start*, *x_end*, *t_start*, *dt* are used to 20 | compute the symmetric Runge-Kutta 4 estimate. 21 | x_dot : callable 22 | Callable taking three arguments, x, t and u, returning the time 23 | derivative at the provided points. 24 | x_start : (..., N) ndarray 25 | x_end : (..., N) ndarray 26 | t_start : number or (..., 1) ndarray 27 | t_end : number or (..., 1) ndarray 28 | dt : number 29 | u : (..., N) ndarray, default None 30 | Controlled input to provide to *x_dot*. Will only be used 31 | if *integrator* is 'srk4'. 32 | 33 | Returns 34 | ------- 35 | (..., N) ndarray 36 | The estimated time derivatives. 37 | 38 | Raises 39 | ------ 40 | ValueError 41 | If the integrator type is not recognized. 42 | 43 | """ 44 | 45 | integrator = (integrator.lower() if isinstance(integrator, str) 46 | else integrator) 47 | if integrator in (False, 'euler'): 48 | return _time_derivative_continuous(x_dot, x_start, t_start, xspatial) 49 | elif integrator == 'midpoint': 50 | x_mid = (x_end + x_start) / 2 51 | t_mid = (t_end + t_start) / 2 52 | return _time_derivative_continuous(x_dot, x_mid, t_mid, xspatial) 53 | elif integrator == 'rk4': 54 | return _discrete_time_derivative_rk4(x_dot, x_start, t_start, dt, u, xspatial) 55 | elif integrator == 'srk4': 56 | return _discrete_time_derivative_srk4(x_dot, x_start, x_end, 57 | t_start, dt, xspatial) 58 | elif integrator == 'cm4': 59 | return _discrete_time_derivative_cm4(x_dot, x_start, x_end, 60 | t_start, dt, xspatial) 61 | elif integrator == 'cs6': 62 | return _discrete_time_derivative_cs6(x_dot, x_start, x_end, 63 | t_start, dt, xspatial) 64 | else: 65 | raise ValueError(f'Unknown integrator {integrator}.') 66 | 67 | 68 | def _time_derivative_continuous(x_dot, x, t=None, xspatial=None): 69 | return x_dot(x, t, xspatial=xspatial) 70 | 71 | 72 | def _discrete_time_derivative_rk4(x_dot, x1, t1, dt, u, xspatial): 73 | k1 = x_dot(x1, t1, u, xspatial) 74 | k2 = x_dot(x1+.5*dt*k1, t1+.5*dt, u, xspatial) 75 | k3 = x_dot(x1+.5*dt*k2, t1+.5*dt, u, xspatial) 76 | k4 = x_dot(x1+dt*k3, t1+dt, u, xspatial) 77 | return 1/6*(k1+2*k2+2*k3+k4) 78 | 79 | 80 | def _discrete_time_derivative_srk4(x_dot, x1, x2, t1, dt, xspatial=None): 81 | xh = (x1+x2)/2 82 | z1 = (1/2+np.sqrt(3)/6)*x1 + (1/2-np.sqrt(3)/6)*x2 83 | z2 = (1/2-np.sqrt(3)/6)*x1 + (1/2+np.sqrt(3)/6)*x2 84 | tm = (t1 + (1/2-np.sqrt(3)/6)*dt) 85 | tp = (t1 + (1/2+np.sqrt(3)/6)*dt) 86 | z3 = xh - np.sqrt(3)/6*dt*x_dot(z2, tp, xspatial=xspatial).detach() 87 | z4 = xh + np.sqrt(3)/6*dt*x_dot(z1, tm, xspatial=xspatial).detach() 88 | return 1/2*(x_dot(z3, tm, xspatial=xspatial)+x_dot(z4, tp, xspatial=xspatial)) 89 | 90 | 91 | # Cash and Moore's 4th order scheme 92 | def _discrete_time_derivative_cm4(x_dot, x1, x2, t1, dt, xspatial=None): 93 | xh = (x1+x2)/2 94 | th = t1 + 1/2*dt 95 | t2 = t1 + dt 96 | f1 = x_dot(x1, t1, xspatial=xspatial) 97 | f2 = x_dot(x2, t2, xspatial=xspatial) 98 | z = xh - dt/8*(f2-f1).detach() 99 | return 1/6*(f1+f2)+2/3*x_dot(z, th, xspatial=xspatial) 100 | 101 | 102 | # Cash and Singhal's 6th order scheme: 103 | def _discrete_time_derivative_cs6(x_dot, x1, x2, t1, dt, xspatial=None): 104 | t14 = t1 + 1/4*dt 105 | t12 = t1 + 1/2*dt 106 | t34 = t1 + 3/4*dt 107 | t2 = t1 + dt 108 | f1 = x_dot(x1, t1, xspatial=xspatial) 109 | f2 = x_dot(x2, t2, xspatial=xspatial) 110 | x14 = 27/32*x1 + 5/32*x2 + dt*(9/64*f1-3/64*f2).detach() 111 | x34 = 5/32*x1 + 27/32*x2 + dt*(3/64*f1-9/64*f2).detach() 112 | x12 = (x1+x2)/2 + 5/24*dt*(f2-f1) - 2/3*dt*(x_dot(x34,t34, xspatial=xspatial)- 113 | x_dot(x14,t14, xspatial=xspatial)).detach() 114 | return (7/90*(f1+f2) + 16/45*(x_dot(x14,t14, xspatial=xspatial)+x_dot(x34,t34, xspatial=xspatial)) 115 | + 2/15*x_dot(x12,t12, xspatial=xspatial)) -------------------------------------------------------------------------------- /phlearn/phlearn/utils/utils.py: -------------------------------------------------------------------------------- 1 | 2 | import torch 3 | import numpy as np 4 | import numpy.linalg as la 5 | import matplotlib.pyplot as plt 6 | import imageio 7 | from IPython.display import display, Video, Image 8 | 9 | __all__ = ['to_tensor', 'midpoint_method', 'create_video'] 10 | 11 | def to_tensor(x, ttype=torch.float32): 12 | """ 13 | Converts the input to a torch tensor if the input is not None. 14 | 15 | Parameters 16 | ---------- 17 | x : listlike or None 18 | ttype : torch type, default torch.float32 19 | 20 | Returns 21 | ------- 22 | torch.tensor or None 23 | Return converted list/array/tensor unless *x* is None, 24 | in which case it returns None. 25 | 26 | """ 27 | if x is None: 28 | return x 29 | elif not isinstance(x, torch.Tensor): 30 | return torch.tensor(x, dtype=ttype) 31 | else: 32 | return x 33 | 34 | 35 | def midpoint_method(u, un, t, f, Df, dt, M, tol=1e-12, max_iter=5): 36 | """ 37 | Integrates one step of the ODE system u_t = f, 38 | from u to un, with the implicit midpoint method. 39 | Uses Newton's method to find un. 40 | 41 | Parameters 42 | ---------- 43 | u : ndarray, shape (M,) 44 | Initial state of the ODE system. 45 | un : ndarray, shape (M,) 46 | Initial guess on the state of the ODE system after one time step. 47 | t : float 48 | Time at the initial state. 49 | f : callable 50 | Function that evaluates the right-hand side of the ODE system, 51 | given the state u and the time t. 52 | It should return an array_like of shape (M,). 53 | Df : callable 54 | Function that evaluates the Jacobian matrix of f, 55 | given the state u and the time t. 56 | It should return an array_like of shape (M, M). 57 | dt : float 58 | Time step size. 59 | M : int 60 | Number of equations in the ODE system. 61 | tol : float, optional 62 | Tolerance for the Newton iteration. The iteration stops when the 63 | Euclidean norm of the residual is less than `tol`. Default is 1e-12. 64 | max_iter : int, optional 65 | Maximum number of iterations for the Newton iteration. If the 66 | iteration does not converge after `max_iter` iterations, it stops. 67 | Default is 5. 68 | 69 | 70 | Returns 71 | ------- 72 | un : array_like, shape (M,) 73 | Final state of the ODE system after one time step. 74 | """ 75 | 76 | I = np.eye(M) 77 | F = lambda u_hat: 1/dt*(u_hat-u) - f((u+u_hat)/2, t+.5*dt) 78 | J = lambda u_hat: 1/dt*I - 1/2*Df((u+u_hat)/2, t+.5*dt) 79 | err = la.norm(F(un)) 80 | it = 0 81 | while err > tol: 82 | un = un - la.solve(J(un),F(un)) 83 | err = la.norm(F(un)) 84 | it += 1 85 | if it > max_iter: 86 | break 87 | return un 88 | 89 | 90 | def create_video(arrays, labels, x_axis=None, file_name='animation.mp4', fps=10, dpi=100, output_format='MP4'): 91 | """ 92 | Creates an MP4 video or GIF showing the evaluation of a system over time, given by data in an 93 | array where the time scale is along the first dimension. 94 | 95 | Parameters 96 | ---------- 97 | arrays : list of ndarray 98 | List of numpy arrays containing the data to be plotted. 99 | labels : list of str 100 | List of labels for the data series, used in the legend of the plot, corresponding to the 101 | list of arrays. 102 | x_axis : ndarray, optional 103 | x-axis values. If provided, data will be plotted against these x-axis values. If None, 104 | data will be plotted against the array indices. 105 | file_name : str, optional 106 | Name of the output animation file. Defaults to 'animation.mp4'. 107 | fps : int, optional 108 | Frames per second for the animation. Defaults to 10. 109 | dpi : int, optional 110 | Dots per inch for the plot's resolution. Defaults to 100. 111 | output_format : str, optional 112 | Output format for the animation. Can be 'MP4' or 'GIF'. 113 | Defaults to 'MP4'. 114 | 115 | Returns 116 | ------- 117 | Video or None 118 | Returns a Video object if the output_format is 'MP4'. If output_format is 'GIF', 119 | the GIF is displayed and None is returned. 120 | """ 121 | 122 | fig, ax = plt.subplots(figsize=(7.04, 4), dpi=dpi) 123 | colors = [(0,0,0),(0,0.4,1),(1,0.7,0.3),(0.2,0.7,0.2),(0.8,0,0.2),(0.5,0.3,.9)][:len(arrays)] 124 | min_value = max(min(np.min(data) for data in arrays), -2) 125 | max_value = min(max(np.max(data) for data in arrays), 2*np.max(arrays[0])) 126 | ax.set_ylim(min_value, max_value) 127 | 128 | with imageio.get_writer(file_name, mode='I', fps=fps) as writer: 129 | for frame in range(arrays[0].shape[0]): 130 | lines = [] 131 | for (data, color, label) in zip(arrays, colors, labels): 132 | if x_axis is not None: 133 | line, = ax.plot(x_axis, data[frame, :], color=color, label=label) 134 | else: 135 | line, = ax.plot(data[frame, :], color=color, label=label) 136 | lines.append(line) 137 | ax.legend(loc='upper right') 138 | ax.set_xlabel('$x$', fontsize=12) 139 | ax.set_ylabel('$u$', fontsize=12) 140 | fig.canvas.draw() 141 | image = np.frombuffer(fig.canvas.tostring_rgb(), dtype='uint8') 142 | image = image.reshape(fig.canvas.get_width_height()[::-1] + (3,)) 143 | writer.append_data(image) 144 | ax.clear() 145 | ax.set_ylim(min_value, max_value) 146 | plt.close(fig) 147 | 148 | if output_format == 'MP4': 149 | return Video(file_name) 150 | elif output_format == 'GIF': 151 | with open(file_name,'rb') as f: 152 | display(Image(data=f.read(), format='gif')) 153 | return None -------------------------------------------------------------------------------- /phlearn/requirements-control.txt: -------------------------------------------------------------------------------- 1 | casadi==3.5.5 2 | do_mpc==4.4.0 -------------------------------------------------------------------------------- /phlearn/requirements-dev.txt: -------------------------------------------------------------------------------- 1 | pytest -------------------------------------------------------------------------------- /phlearn/requirements.txt: -------------------------------------------------------------------------------- 1 | networkx>=2.7.1 2 | numpy>=1.22.3 3 | torchvision 4 | scipy>=1.8.0 5 | torch 6 | torchaudio 7 | matplotlib 8 | imageio 9 | tqdm 10 | autograd 11 | IPython -------------------------------------------------------------------------------- /phlearn/setup.py: -------------------------------------------------------------------------------- 1 | import os 2 | import setuptools 3 | 4 | 5 | def read(fname): 6 | return open(os.path.join(os.path.dirname(__file__), fname)).read() 7 | 8 | 9 | setuptools.setup( 10 | name='phlearn', 11 | version='1.1.2', 12 | author='Sølve Eidnes', 13 | author_email='solve.eidnes@sintef.no', 14 | description=('A package for simulating and learning pseudo-Hamiltonian systems.' 15 | ' For further details, see https://arxiv.org/pdf/2206.02660.pdf and https://arxiv.org/abs/2304.14374'), 16 | keywords='pseudo-Hamiltonian neural networks', 17 | url="https://github.com/SINTEF/pseudo-hamiltonian-neural-networks", 18 | packages=setuptools.find_packages(), 19 | long_description=read('README.md'), 20 | classifiers=[ 21 | 'Development Status :: 5 - Production/Stable' 22 | ], 23 | license='MIT', 24 | install_requires=[ 25 | 'networkx>=2.7.1', 26 | 'numpy>=1.22.3', 27 | 'torchvision', 28 | 'scipy>=1.8.0', 29 | 'torch', 30 | 'torchaudio', 31 | 'matplotlib', 32 | 'imageio', 33 | 'tqdm', 34 | 'autograd', 35 | 'IPython', 36 | ], 37 | extras_require={'control': ['casadi', 'do_mpc']}, 38 | ) 39 | --------------------------------------------------------------------------------