[ { "path": ".dockerignore", "content": ".vscode\n*.py\nwandb\nconfig\ndocs\nmodels\n*/*.py\nexp\ncheckpoints\n*/__pycache__/**" }, { "path": ".gitignore", "content": "data\nlog\n.vscode\nwandb\n**/__pycache__/**\n.idea\nfigs\ncheckpoints\n.ipynb_checkpoints\n*.ipynb\n*.pt\n*.pth\ntensordiffeq\nexp" }, { "path": "Dockerfile", "content": "FROM nvcr.io/nvidia/pytorch:22.09-py3\nRUN useradd -ms /bin/bash pino\nUSER pino\nENV PATH=/home/pino/.local/bin:$PATH\nRUN pip install --user \\\n wandb tqdm pyyaml" }, { "path": "LICENSE", "content": " Apache License\n Version 2.0, January 2004\n http://www.apache.org/licenses/\n\n TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION\n\n 1. Definitions.\n\n \"License\" shall mean the terms and conditions for use, reproduction,\n and distribution as defined by Sections 1 through 9 of this document.\n\n \"Licensor\" shall mean the copyright owner or entity authorized by\n the copyright owner that is granting the License.\n\n \"Legal Entity\" shall mean the union of the acting entity and all\n other entities that control, are controlled by, or are under common\n control with that entity. For the purposes of this definition,\n \"control\" means (i) the power, direct or indirect, to cause the\n direction or management of such entity, whether by contract or\n otherwise, or (ii) ownership of fifty percent (50%) or more of the\n outstanding shares, or (iii) beneficial ownership of such entity.\n\n \"You\" (or \"Your\") shall mean an individual or Legal Entity\n exercising permissions granted by this License.\n\n \"Source\" form shall mean the preferred form for making modifications,\n including but not limited to software source code, documentation\n source, and configuration files.\n\n \"Object\" form shall mean any form resulting from mechanical\n transformation or translation of a Source form, including but\n not limited to compiled object code, generated documentation,\n and conversions to other media types.\n\n \"Work\" shall mean the work of authorship, whether in Source or\n Object form, made available under the License, as indicated by a\n copyright notice that is included in or attached to the work\n (an example is provided in the Appendix below).\n\n \"Derivative Works\" shall mean any work, whether in Source or Object\n form, that is based on (or derived from) the Work and for which the\n editorial revisions, annotations, elaborations, or other modifications\n represent, as a whole, an original work of authorship. For the purposes\n of this License, Derivative Works shall not include works that remain\n separable from, or merely link (or bind by name) to the interfaces of,\n the Work and Derivative Works thereof.\n\n \"Contribution\" shall mean any work of authorship, including\n the original version of the Work and any modifications or additions\n to that Work or Derivative Works thereof, that is intentionally\n submitted to Licensor for inclusion in the Work by the copyright owner\n or by an individual or Legal Entity authorized to submit on behalf of\n the copyright owner. For the purposes of this definition, \"submitted\"\n means any form of electronic, verbal, or written communication sent\n to the Licensor or its representatives, including but not limited to\n communication on electronic mailing lists, source code control systems,\n and issue tracking systems that are managed by, or on behalf of, the\n Licensor for the purpose of discussing and improving the Work, but\n excluding communication that is conspicuously marked or otherwise\n designated in writing by the copyright owner as \"Not a Contribution.\"\n\n \"Contributor\" shall mean Licensor and any individual or Legal Entity\n on behalf of whom a Contribution has been received by Licensor and\n subsequently incorporated within the Work.\n\n 2. Grant of Copyright License. Subject to the terms and conditions of\n this License, each Contributor hereby grants to You a perpetual,\n worldwide, non-exclusive, no-charge, royalty-free, irrevocable\n copyright license to reproduce, prepare Derivative Works of,\n publicly display, publicly perform, sublicense, and distribute the\n Work and such Derivative Works in Source or Object form.\n\n 3. Grant of Patent License. Subject to the terms and conditions of\n this License, each Contributor hereby grants to You a perpetual,\n worldwide, non-exclusive, no-charge, royalty-free, irrevocable\n (except as stated in this section) patent license to make, have made,\n use, offer to sell, sell, import, and otherwise transfer the Work,\n where such license applies only to those patent claims licensable\n by such Contributor that are necessarily infringed by their\n Contribution(s) alone or by combination of their Contribution(s)\n with the Work to which such Contribution(s) was submitted. If You\n institute patent litigation against any entity (including a\n cross-claim or counterclaim in a lawsuit) alleging that the Work\n or a Contribution incorporated within the Work constitutes direct\n or contributory patent infringement, then any patent licenses\n granted to You under this License for that Work shall terminate\n as of the date such litigation is filed.\n\n 4. Redistribution. You may reproduce and distribute copies of the\n Work or Derivative Works thereof in any medium, with or without\n modifications, and in Source or Object form, provided that You\n meet the following conditions:\n\n (a) You must give any other recipients of the Work or\n Derivative Works a copy of this License; and\n\n (b) You must cause any modified files to carry prominent notices\n stating that You changed the files; and\n\n (c) You must retain, in the Source form of any Derivative Works\n that You distribute, all copyright, patent, trademark, and\n attribution notices from the Source form of the Work,\n excluding those notices that do not pertain to any part of\n the Derivative Works; and\n\n (d) If the Work includes a \"NOTICE\" text file as part of its\n distribution, then any Derivative Works that You distribute must\n include a readable copy of the attribution notices contained\n within such NOTICE file, excluding those notices that do not\n pertain to any part of the Derivative Works, in at least one\n of the following places: within a NOTICE text file distributed\n as part of the Derivative Works; within the Source form or\n documentation, if provided along with the Derivative Works; or,\n within a display generated by the Derivative Works, if and\n wherever such third-party notices normally appear. The contents\n of the NOTICE file are for informational purposes only and\n do not modify the License. You may add Your own attribution\n notices within Derivative Works that You distribute, alongside\n or as an addendum to the NOTICE text from the Work, provided\n that such additional attribution notices cannot be construed\n as modifying the License.\n\n You may add Your own copyright statement to Your modifications and\n may provide additional or different license terms and conditions\n for use, reproduction, or distribution of Your modifications, or\n for any such Derivative Works as a whole, provided Your use,\n reproduction, and distribution of the Work otherwise complies with\n the conditions stated in this License.\n\n 5. Submission of Contributions. Unless You explicitly state otherwise,\n any Contribution intentionally submitted for inclusion in the Work\n by You to the Licensor shall be under the terms and conditions of\n this License, without any additional terms or conditions.\n Notwithstanding the above, nothing herein shall supersede or modify\n the terms of any separate license agreement you may have executed\n with Licensor regarding such Contributions.\n\n 6. Trademarks. This License does not grant permission to use the trade\n names, trademarks, service marks, or product names of the Licensor,\n except as required for reasonable and customary use in describing the\n origin of the Work and reproducing the content of the NOTICE file.\n\n 7. Disclaimer of Warranty. Unless required by applicable law or\n agreed to in writing, Licensor provides the Work (and each\n Contributor provides its Contributions) on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or\n implied, including, without limitation, any warranties or conditions\n of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A\n PARTICULAR PURPOSE. You are solely responsible for determining the\n appropriateness of using or redistributing the Work and assume any\n risks associated with Your exercise of permissions under this License.\n\n 8. Limitation of Liability. In no event and under no legal theory,\n whether in tort (including negligence), contract, or otherwise,\n unless required by applicable law (such as deliberate and grossly\n negligent acts) or agreed to in writing, shall any Contributor be\n liable to You for damages, including any direct, indirect, special,\n incidental, or consequential damages of any character arising as a\n result of this License or out of the use or inability to use the\n Work (including but not limited to damages for loss of goodwill,\n work stoppage, computer failure or malfunction, or any and all\n other commercial damages or losses), even if such Contributor\n has been advised of the possibility of such damages.\n\n 9. Accepting Warranty or Additional Liability. While redistributing\n the Work or Derivative Works thereof, You may choose to offer,\n and charge a fee for, acceptance of support, warranty, indemnity,\n or other liability obligations and/or rights consistent with this\n License. However, in accepting such obligations, You may act only\n on Your own behalf and on Your sole responsibility, not on behalf\n of any other Contributor, and only if You agree to indemnify,\n defend, and hold each Contributor harmless for any liability\n incurred by, or claims asserted against, such Contributor by reason\n of your accepting any such warranty or additional liability.\n\n END OF TERMS AND CONDITIONS\n\n APPENDIX: How to apply the Apache License to your work.\n\n To apply the Apache License to your work, attach the following\n boilerplate notice, with the fields enclosed by brackets \"[]\"\n replaced with your own identifying information. (Don't include\n the brackets!) The text should be enclosed in the appropriate\n comment syntax for the file format. We also recommend that a\n file or class name and description of purpose be included on the\n same \"printed page\" as the copyright notice for easier\n identification within third-party archives.\n\n Copyright [yyyy] [name of copyright owner]\n\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n\n http://www.apache.org/licenses/LICENSE-2.0\n\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n" }, { "path": "README.md", "content": "# Physics-Informed Neural Operator for Learning Partial Differential Equations\n\n# 📢 DEPRECATION NOTICE 📢 \n----------------------------\n\n🚨 **This repository is no longer maintained.** 🚨 The code in this repository is **deprecated** and may not work with newer dependencies or frameworks. \nFor the most up-to-date implementation and continued development, please visit: \n\n## ➡️ **[NeuralOperator](https://github.com/neuraloperator/neuraloperator)** ⬅️\n\n🔴 We strongly recommend using the latest version to ensure compatibility, performance, and support.🔴 \n\n----------------------------\n![PINO Diagram](docs/pino-diagram4.png)\n\n[comment]: <> (![Results on Navier Stokes equation](docs/solver-pino.png))\n\n\"Results\n\n# Paper Info\n\n\nThis repo contains code for experiments from the paper [Physics-Informed Neural Operator for Learning Partial Differential Equations](https://arxiv.org/abs/2111.03794) (2021) by Zongyi Li, Hongkai Zheng, Nikola Kovachki, David Jin, Haoxuan Chen, Burigede Liu, Kamyar Azizzadenesheli, and Anima Anandkumar.\n\nAbstract: \n> Machine learning methods have recently shown promise in solving partial differential equations (PDEs). They can be classified into two broad categories: solution function approximation and operator learning. The Physics-Informed Neural Network (PINN) is an example of the former while the Fourier neural operator (FNO) is an example of the latter. Both these approaches have shortcomings. The optimization in PINN is challenging and prone to failure, especially on multi-scale dynamic systems. FNO does not suffer from this optimization issue since it carries out supervised learning on a given dataset, but obtaining such data may be too expensive or infeasible. In this work, we propose the physics-informed neural operator (PINO), where we combine the operating-learning and function-optimization frameworks, and this improves convergence rates and accuracy over both PINN and FNO models. In the operator-learning phase, PINO learns the solution operator over multiple instances of the parametric PDE family. In the test-time optimization phase, PINO optimizes the pre-trained operator ansatz for the querying instance of the PDE. Experiments show PINO outperforms previous ML methods on many popular PDE families while retaining the extraordinary speed-up of FNO compared to solvers. In particular, PINO accurately solves long temporal transient flows and Kolmogorov flows, while PINN and other methods fail to converge.\n## Requirements\n- Pytorch 1.8.0 or later\n- wandb\n- tqdm\n- scipy\n- h5py\n- numpy\n- DeepXDE:latest\n- Latest code from tensordiffeq github master branch (Not tensordiffeq 0.19)\n- tensorflow 2.4.0\n\n## Data description\n### Burgers equation\n[burgers_pino.mat](https://hkzdata.s3.us-west-2.amazonaws.com/PINO/data/burgers_pino.mat)\n\n### Darcy flow \n- spatial domain: $x\\in (0,1)^2$\n- Data file: \n - [piececonst_r421_N1024_smooth1.mat](https://hkzdata.s3.us-west-2.amazonaws.com/PINO/data/piececonst_r421_N1024_smooth1.mat)\n - [piececonst_r421_N1024_smooth2.mat](https://hkzdata.s3.us-west-2.amazonaws.com/PINO/data/piececonst_r421_N1024_smooth2.mat)\n- Raw data shape: 1024x421x421\n\n\n### Long roll out of Navier Stokes equation\n- spatial domain: $x\\in (0, 1)^2$\n- temporal domain: $t\\in \\[0, 49\\]$\n- forcing: $0.1(\\sin(2\\pi(x_1+x_2)) + \\cos(2\\pi(x_1+x_2)))$\n- viscosity = 0.001\n\nData file: `nv_V1e-3_N5000_T50.mat`, with shape 50 x 64 x 64 x 5000 \n\n- train set: -1-4799\n- test set: 4799-4999\n### Navier Stokes with Reynolds number 500\n- spatial domain: $x\\in (0, 2\\pi)^2$\n- temporal domain: $t \\in \\[0, 0.5\\]$\n- forcing: $-4\\cos(4x_2)$\n- Reynolds number: 500\n\nTrain set: data of shape (N, T, X, Y) where N is the number of instances, T is temporal resolution, X, Y are spatial resolutions. \n1. [NS_fft_Re500_T4000.npy](https://hkzdata.s3.us-west-2.amazonaws.com/PINO/data/NS_fft_Re500_T4000.npy) : 4000x64x64x65\n2. [NS_fine_Re500_T128_part0.npy](https://hkzdata.s3.us-west-2.amazonaws.com/PINO/data/NS_fine_Re500_T128_part0.npy): 100x129x128x128\n3. [NS_fine_Re500_T128_part1.npy](https://hkzdata.s3.us-west-2.amazonaws.com/PINO/data/NS_fine_Re500_T128_part1.npy): 100x129x128x128\n\nTest set: data of shape (N, T, X, Y) where N is the number of instances, T is temporal resolution, X, Y are spatial resolutions. \n1. [NS_Re500_s256_T100_test.npy](https://hkzdata.s3.us-west-2.amazonaws.com/PINO/data/NS_Re500_s256_T100_test.npy): 100x129x256x256\n2. [NS_fine_Re500_T128_part2.npy](https://hkzdata.s3.us-west-2.amazonaws.com/PINO/data/NS_fine_Re500_T128_part2.npy): 100x129x128x128\n\nConfiguration file format: see `.yaml` files under folder `configs` for detail. \n\n## Code for Burgers equation\n### Train PINO\nTo run PINO for Burgers equation, use, e.g.,\n```bash \npython3 train_burgers.py --config_path configs/pretrain/burgers-pretrain.yaml --mode train\n```\n\nTo test PINO for burgers equation, use, e.g., \n```bash\npython3 train_burgers.py --config_path configs/test/burgers.yaml --mode test\n```\n\n## Code for Darcy Flow\n\n### Operator learning\nTo run PINO for Darcy Flow, use, e.g., \n```bash\npython3 train_operator.py --config_path configs/pretrain/Darcy-pretrain.yaml\n```\nTo evaluate operator for Darcy Flow, use, e.g., \n```bash\npython3 eval_operator.py --config_path configs/test/darcy.yaml\n```\n\n### Test-time optimization\nTo do test-time optimization for Darcy Flow, use, e.g., \n```bash\npython3 run_pino2d.py --config_path configs/finetune/Darcy-finetune.yaml --start [starting index] --stop [stopping index]\n```\n\n### Baseline\nTo run DeepONet, use, e.g., \n```bash\npython3 deeponet.py --config_path configs/pretrain/Darcy-pretrain-deeponet.yaml --mode train \n```\nTo test DeepONet, use, e.g., \n```bash\npython3 deeponet.py --config_path configs/test/darcy.yaml --mode test\n```\n\n\n## Code for Navier Stokes equation\n### Run exp on new dataset\nTrain PINO with 800 low-res data and 2200 PDE. \n```bash\npython3 train_pino.py --config configs/operator/Re500-1_8-800-PINO-s.yaml\n```\nTrain FNO with 800 low-res data and 2200 PDE.\n```bash\npython3 train_pino.py --config configs/operator/Re500-1_8-800-FNO-s.yaml\n```\nRun instance-wise finetuning\n```bash\npython3 instance_opt.py --config configs/instance/Re500-1_8-PINO-s.yaml\n```\n\n### Train PINO for short time period\nTo run operator learning, use, e.g., \n```bash\npython3 train_operator.py --config_path configs/pretrain/Re500-pretrain-05s-4C0.yaml\n```\nTo evaluate trained operator, use\n```bash\npython3 eval_operator.py --config_path configs/test/Re500-05s.yaml\n```\nTo run test-time optimization, use\n```bash\npython3 train_PINO3d.py --config_path configs/***.yaml \n```\n\nTo train Navier Stokes equations sequentially without running `train_PINO3d.py` multiple times, use\n\n```bash\npython3 run_pino3d.py --config_path configs/[configuration file name].yaml --start [index of the first data] --stop [which data to stop]\n```\n\n\n### Baseline for short time period\nTo train DeepONet, use \n```bash\npython3 deeponet.py --config_path configs/[configuration file].yaml --mode train\n```\n\nTo test DeepONet, use \n```bash\npython3 deeponet.py --config_path configs/[configuration file].yaml --mode test\n```\n\nTo train and test PINNs, use, e.g., \n```bash\npython3 pinns.py --config_path configs/baseline/Re500-pinns-05s.yaml --start [starting index] --stop [stopping index]\n```\n\nTo train and test LAAF-PINN, use, e.g., \n```bash\npython3 pinns.py configs/baseline/Re500-pinns-05s-LAAF.yaml --start [starting index] --stop [stopping index]\n```\n\nTo train and test SA-PINNs, first copy the latest code of tensordiffeq under the working directory.\nThen run: \n```bash \nDDEBACKEND=pytorch python3 pinns.py configs/baseline/Re500-pinns-05s-SA.yaml --start [starting index] --stop [stopping index]\n```\n\n### Baseline for long roll out\nTo train and test PINNs, use\n```bash\npython3 pinns.py --config_path configs/baseline/NS-50s.yaml --start [starting index] --stop [stopping index]\n```\n\nTo train and test LAAF-PINN, use, e.g., \n```bash\npython3 pinns.py --config_path configs/baseline/NS-50s-LAAF.yaml --start [starting index] --stop [stopping index]\n```\n\n### Pseudospectral solver for Navier Stokes equation\nTo run solver, use \n```bash\npython3 run_solver.py --config_path configs/Re500-0.5s.yaml\n```\n" }, { "path": "baselines/__init__.py", "content": "" }, { "path": "baselines/data.py", "content": "import numpy as np\nimport torch\nfrom torch.utils.data import Dataset\nfrom .utils import get_xytgrid, get_3dboundary, get_3dboundary_points\nfrom train_utils.utils import vor2vel, torch2dgrid\nimport scipy.io\nimport h5py\n\n\nclass DarcyFlow(Dataset):\n def __init__(self,\n datapath,\n nx, sub,\n offset=0,\n num=1):\n self.S = int(nx // sub) + 1\n data = scipy.io.loadmat(datapath)\n a = data['coeff']\n u = data['sol']\n self.a = torch.tensor(a[offset: offset + num, ::sub, ::sub], dtype=torch.float)\n self.u = torch.tensor(u[offset: offset + num, ::sub, ::sub], dtype=torch.float)\n self.mesh = torch2dgrid(self.S, self.S)\n\n def __len__(self):\n return self.a.shape[0]\n\n def __getitem__(self, item):\n fa = self.a[item]\n return fa.reshape(-1), self.u[item].reshape(-1)\n\n\nclass NSLong(object):\n def __init__(self,\n datapath,\n nx, nt,\n time_scale,\n offset=0,\n num=1, vel=False):\n '''\n Load data from mat\n Args:\n datapath: path to data file\n nx: number of points in each spatial domain\n nt: number of points in temporal domain\n offset: index of the instance\n num: number of instances\n vel: compute velocity from vorticity if True\n '''\n\n self.time_scale = time_scale\n self.S = nx\n self.T = nt\n\n with h5py.File(datapath, mode='r') as file:\n raw = file['u']\n data = np.array(raw)\n vor = torch.tensor(data, dtype=torch.float).permute(3, 1, 2, 0)\n self.vor = vor[offset: offset + num, :, :, :] # num x 64 x 64 x 50\n if vel:\n self.vel_u, self.vel_v = vor2vel(self.vor, L=1.0)\n\n def get_boundary_value(self, component=0):\n '''\n Get the boundary value for component-th output\n Args:\n component: int, 0: velocity_u; 1: velocity_v; 2: vorticity;\n Returns:\n value: N by 1 array, boundary value of the component\n '''\n if component == 0:\n value = self.vel_u\n elif component == 1:\n value = self.vel_v\n elif component == 2:\n value = self.vor\n else:\n raise ValueError(f'No component {component} ')\n\n boundary = get_3dboundary(value)\n return boundary\n\n def get_boundary_points(self, num_x, num_y, num_t):\n points = get_3dboundary_points(num_x, num_y, num_t,\n bot=(0,0,0),\n top=(1, 1, self.time_scale))\n return points\n\n def get_test_xyt(self):\n '''\n\n Returns:\n points: (x, y, t) array with shape (S * S * T, 3)\n values: (u, v, w) array with shape (S * S * T, 3)\n\n '''\n points = get_xytgrid(S=self.S, T=self.T,\n bot=[0, 0, 0],\n top=[1, 1, self.time_scale])\n u_val = np.ravel(self.vel_u)\n v_val = np.ravel(self.vel_v)\n w_val = np.ravel(self.vor)\n values = np.stack([u_val, v_val, w_val], axis=0).T\n return points, values\n\n\nclass NSdata(object):\n def __init__(self, datapath1,\n nx, nt,\n offset=0, num=1,\n datapath2=None,\n sub=1, sub_t=1,\n vel=False, t_interval=1.0):\n '''\n Load data from npy and reshape to (N, X, Y, T)\n Args:\n datapath1: path to data\n nx: number of points in each spatial domain\n nt: number of points in temporal domain\n offset: index of the instance\n num: number of instances\n datapath2: path to second part of data, default None\n sub: downsample interval of spatial domain\n sub_t: downsample interval of temporal domain\n N:\n t_interval:\n '''\n self.S = nx // sub\n self.T = int(nt * t_interval) // sub_t + 1\n self.time_scale = t_interval\n data1 = np.load(datapath1)\n data1 = torch.tensor(data1, dtype=torch.float)[..., ::sub_t, ::sub, ::sub]\n\n if datapath2 is not None:\n data2 = np.load(datapath2)\n data2 = torch.tensor(data2, dtype=torch.float)[..., ::sub_t, ::sub, ::sub]\n if t_interval == 0.5:\n data1 = self.extract(data1)\n if datapath2 is not None:\n data2 = self.extract(data2)\n # transpose data into (N, S, S, T)\n part1 = data1.permute(0, 2, 3, 1)\n if datapath2 is not None:\n part2 = data2.permute(0, 2, 3, 1)\n self.data = torch.cat((part1, part2), dim=0)\n else:\n self.data = part1\n self.vor = self.data[offset: offset + num, :, :, :].cpu()\n if vel:\n self.vel_u, self.vel_v = vor2vel(self.vor) # Compute velocity from vorticity\n\n def get_init_cond(self):\n values = np.stack([self.vel_u[0, :, :, 0],\n self.vel_v[0, :, :, 0],\n self.vor[0, :, :, 0]], axis=2)\n return values\n\n def get_boundary_value(self, component=0):\n '''\n Get the boundary value for component-th output\n Args:\n component: int, 0: velocity_u; 1: velocity_v; 2: vorticity;\n Returns:\n value: N by 1 array, boundary value of the component\n '''\n if component == 0:\n value = self.vel_u\n elif component == 1:\n value = self.vel_v\n elif component == 2:\n value = self.vor\n else:\n raise ValueError(f'No component {component} ')\n\n boundary = get_3dboundary(value)\n return boundary\n\n def get_boundary_points(self, num_x, num_y, num_t):\n '''\n Args:\n num_x:\n num_y:\n\n Returns:\n points: N by 3 array\n '''\n points = get_3dboundary_points(num_x, num_y, num_t,\n bot=(0, 0, 0),\n top=(2 * np.pi, 2 * np.pi, self.time_scale))\n # x_arr = np.linspace(0, 2 * np.pi, num=num_x, endpoint=False)\n # y_arr = np.linspace(0, 2 * np.pi, num=num_y, endpoint=False)\n # xx, yy = np.meshgrid(x_arr, y_arr, indexing='ij')\n # xarr = np.ravel(xx)\n # yarr = np.ravel(yy)\n # tarr = np.zeros_like(xarr)\n # point0 = np.stack([xarr, yarr, tarr], axis=0).T # (128x128x1, 3), boundary on t=0\n #\n # # tarr = np.ones_like(xarr) * self.time_scale\n # # point1 = np.stack([xarr, yarr, tarr], axis=0).T # (128x128x1, 3), boundary on t=0.5\n #\n # t_arr = np.linspace(0, self.time_scale, num=num_t)\n # yy, tt = np.meshgrid(y_arr, t_arr, indexing='ij')\n # yarr = np.ravel(yy)\n # tarr = np.ravel(tt)\n # xarr = np.zeros_like(yarr)\n # point2 = np.stack([xarr, yarr, tarr], axis=0).T # (1x128x65, 3), boundary on x=0\n #\n # xarr = np.ones_like(yarr) * 2 * np.pi\n # point3 = np.stack([xarr, yarr, tarr], axis=0).T # (1x128x65, 3), boundary on x=2pi\n #\n # xx, tt = np.meshgrid(x_arr, t_arr, indexing='ij')\n # xarr = np.ravel(xx)\n # tarr = np.ravel(tt)\n # yarr = np.zeros_like(xarr)\n # point4 = np.stack([xarr, yarr, tarr], axis=0).T # (128x1x65, 3), boundary on y=0\n #\n # yarr = np.ones_like(xarr) * 2 * np.pi\n # point5 = np.stack([xarr, yarr, tarr], axis=0).T # (128x1x65, 3), boundary on y=2pi\n #\n # points = np.concatenate([point0,\n # point2, point3,\n # point4, point5],\n # axis=0)\n return points\n\n def get_test_xyt(self):\n '''\n\n Returns:\n points: (x, y, t) array with shape (S * S * T, 3)\n values: (u, v, w) array with shape (S * S * T, 3)\n\n '''\n points = get_xytgrid(S=self.S, T=self.T,\n bot=[0, 0, 0],\n top=[2 * np.pi, 2 * np.pi, self.time_scale])\n u_val = np.ravel(self.vel_u)\n v_val = np.ravel(self.vel_v)\n w_val = np.ravel(self.vor)\n values = np.stack([u_val, v_val, w_val], axis=0).T\n return points, values\n\n\n @staticmethod\n def extract(data):\n '''\n Extract data with time range 0-0.5, 0.25-0.75, 0.5-1.0, 0.75-1.25,...\n Args:\n data: tensor with size N x 129 x 128 x 128\n\n Returns:\n output: (4*N-1) x 65 x 128 x 128\n '''\n T = data.shape[1] // 2\n interval = data.shape[1] // 4\n N = data.shape[0]\n new_data = torch.zeros(4 * N - 1, T + 1, data.shape[2], data.shape[3])\n for i in range(N):\n for j in range(4):\n if i == N - 1 and j == 3:\n # reach boundary\n break\n if j != 3:\n new_data[i * 4 + j] = data[i, interval * j:interval * j + T + 1]\n else:\n new_data[i * 4 + j, 0: interval] = data[i, interval * j:interval * j + interval]\n new_data[i * 4 + j, interval: T + 1] = data[i + 1, 0:interval + 1]\n return new_data\n\n\nclass DeepOnetNS(Dataset):\n '''\n Dataset class customized for DeepONet's input format\n '''\n def __init__(self, datapath,\n nx, nt,\n offset=0, num=1,\n sub=1, sub_t=1,\n t_interval=1.0):\n self.S = nx // sub\n self.T = int(nt * t_interval) // sub_t + 1\n self.time_scale = t_interval\n self.N = num\n data = np.load(datapath)\n data = torch.tensor(data, dtype=torch.float)[..., ::sub_t, ::sub, ::sub]\n if t_interval == 0.5:\n data = NSdata.extract(data)\n # transpose data into (N, S, S, T)\n data = data.permute(0, 2, 3, 1)\n self.vor = data[offset: offset + num, :, :, :]\n points = get_xytgrid(S=self.S, T=self.T,\n bot=[0, 0, 0],\n top=[2 * np.pi, 2 * np.pi, self.time_scale])\n self.xyt = torch.tensor(points, dtype=torch.float)\n # (SxSxT, 3)\n\n def __len__(self):\n return self.N * self.S * self.S * self.T\n\n def __getitem__(self, idx):\n num_per_instance = self.S ** 2 * self.T\n instance_id = idx // num_per_instance\n pos_id = idx % num_per_instance\n point = self.xyt[pos_id]\n u0 = self.vor[instance_id, :, :, 0].reshape(-1)\n y = self.vor[instance_id].reshape(-1)[pos_id]\n return u0, point, y\n\n\nclass DeepONetCPNS(Dataset):\n '''\n Dataset class customized for DeepONet cartesian product's input format\n '''\n\n def __init__(self, datapath,\n nx, nt,\n offset=0, num=1,\n sub=1, sub_t=1,\n t_interval=1.0):\n self.S = nx // sub\n self.T = int(nt * t_interval) // sub_t + 1\n self.time_scale = t_interval\n self.N = num\n data = np.load(datapath)\n data = torch.tensor(data, dtype=torch.float)[..., ::sub_t, ::sub, ::sub]\n if t_interval == 0.5:\n data = NSdata.extract(data)\n # transpose data into (N, S, S, T)\n data = data.permute(0, 2, 3, 1)\n self.vor = data[offset: offset + num, :, :, :]\n points = get_xytgrid(S=self.S, T=self.T,\n bot=[0, 0, 0],\n top=[2 * np.pi, 2 * np.pi, self.time_scale])\n self.xyt = torch.tensor(points, dtype=torch.float)\n # (SxSxT, 3)\n\n def __len__(self):\n return self.N\n\n def __getitem__(self, idx):\n '''\n\n Args:\n idx:\n\n Returns:\n u0: (batchsize, u0_dim)\n y: (batchsize, SxSxT)\n '''\n u0 = self.vor[idx, :, :, 0].reshape(-1)\n y = self.vor[idx, :, :, :].reshape(-1)\n return u0, y\n" }, { "path": "baselines/deepxde_deeponet.py", "content": "import random\nimport deepxde as dde\nfrom baselines.data import NSdata\n\n'''\nTraining deepONet using deepxde implementation. \nNote that deepxde requires passing the whole dataset to Triple, which is very memory consuming. \n'''\n\n\ndef train(config):\n seed = random.randint(1, 10000)\n print(f'Random seed :{seed}')\n # construct dataloader\n data_config = config['data']\n train_set = NSdata(datapath1=data_config['datapath'],\n offset=0, num=10,\n nx=data_config['nx'], nt=data_config['nt'],\n sub=data_config['sub'], sub_t=data_config['sub_t'],\n vel=False,\n t_interval=data_config['time_interval'])\n val_set = NSdata(datapath1=data_config['data_val'],\n offset=310, num=10,\n nx=data_config['val_nx'], nt=data_config['val_nt'],\n sub=data_config['val_sub'], sub_t=data_config['val_subt'],\n vel=False,\n t_interval=data_config['time_interval'])\n # assert train_set.S == val_set.S\n dim_a = train_set.S ** 2\n dim_x = 3\n X_train, y_train = train_set.get_operator_data()\n X_val, y_val = val_set.get_operator_data()\n data = dde.data.Triple(X_train=X_train, y_train=y_train, X_test=X_val, y_test=y_val)\n\n activation = config['model']['activation']\n initializer = 'Glorot normal' # He normal or Glorot normal\n\n net = dde.maps.DeepONet([dim_a] + config['model']['layers'],\n [dim_x] + config['model']['layers'],\n activation,\n initializer,\n use_bias=True,\n stacked=False)\n model = dde.Model(data, net)\n model.compile('adam', lr=config['train']['base_lr'])\n checker = dde.callbacks.ModelCheckpoint(\n 'checkpoints/deeponet.ckpt', save_better_only=True, period=10,\n )\n model.train(epochs=config['train']['epochs'], callbacks=[checker])\n" }, { "path": "baselines/loss.py", "content": "import torch\nimport torch.autograd as autograd\nfrom train_utils.utils import set_grad\nfrom .utils import get_sample, net_NS, sub_mse\n\n\ndef boundary_loss(model, npt=100):\n device = next(model.parameters()).device\n\n bc1_x_sample, bc1_y_sample, bc1_t_sample, bc2_x_sample, bc2_y_sample, bc2_t_sample \\\n = get_sample(npt)\n\n bc1_x_sample, bc1_y_sample, bc1_t_sample, bc2_x_sample, bc2_y_sample, bc2_t_sample \\\n = bc1_x_sample.to(device), bc1_y_sample.to(device), bc1_t_sample.to(device), \\\n bc2_x_sample.to(device), bc2_y_sample.to(device), bc2_t_sample.to(device)\n set_grad([bc1_x_sample, bc1_y_sample, bc1_t_sample, bc2_x_sample, bc2_y_sample, bc2_t_sample])\n\n u1, v1, _ = net_NS(bc1_x_sample, bc1_y_sample, bc1_t_sample, model)\n u2, v2, _ = net_NS(bc2_x_sample, bc2_y_sample, bc2_t_sample, model)\n bc_loss = sub_mse(u1) + sub_mse(v1) + sub_mse(u2) + sub_mse(v2)\n return 0.5 * bc_loss # 0.5 is the normalization factor\n\n\ndef resf_NS(u, v, p, x, y, t, re=40):\n '''\n Args:\n u: x-component, tensor\n v: y-component, tensor\n x: x-dimension, tensor\n y: y-dimension, tensor\n t: time dimension, tensor\n Returns:\n Residual f error\n '''\n u_x, u_y, u_t = autograd.grad(outputs=[u.sum()], inputs=[x, y, t], create_graph=True)\n v_x, v_y, v_t = autograd.grad(outputs=[v.sum()], inputs=[x, y, t], create_graph=True)\n u_xx, = autograd.grad(outputs=[u_x.sum()], inputs=[x], create_graph=True)\n u_yy, = autograd.grad(outputs=[u_y.sum()], inputs=[y], create_graph=True)\n v_xx, = autograd.grad(outputs=[v_x.sum()], inputs=[x], create_graph=True)\n v_yy, = autograd.grad(outputs=[v_y.sum()], inputs=[y], create_graph=True)\n p_x, = autograd.grad(outputs=[p.sum()], inputs=[x], create_graph=True)\n p_y, = autograd.grad(outputs=[p.sum()], inputs=[y], create_graph=True)\n res_x = u_t + u * u_x + v * u_y + p_x - 1 / re * (u_xx + u_yy) - torch.sin(4 * y)\n res_y = v_t + u * v_x + v * v_y + p_y - 1 / re * (v_xx + v_yy)\n evp3 = u_x + v_y\n return res_x, res_y, evp3\n\n" }, { "path": "baselines/model.py", "content": "import torch\nimport torch.nn as nn\nfrom models.FCN import DenseNet\nfrom typing import List\nfrom .utils import weighted_mse\n\n\nclass DeepONet(nn.Module):\n def __init__(self, branch_layer, trunk_layer):\n super(DeepONet, self).__init__()\n self.branch = DenseNet(branch_layer, nn.ReLU)\n self.trunk = DenseNet(trunk_layer, nn.ReLU)\n\n def forward(self, u0, grid):\n a = self.branch(u0)\n b = self.trunk(grid)\n batchsize = a.shape[0]\n dim = a.shape[1]\n return torch.bmm(a.view(batchsize, 1, dim), b.view(batchsize, dim, 1))\n\n\nclass DeepONetCP(nn.Module):\n def __init__(self, branch_layer, trunk_layer):\n super(DeepONetCP, self).__init__()\n self.branch = DenseNet(branch_layer, nn.ReLU)\n self.trunk = DenseNet(trunk_layer, nn.ReLU)\n\n def forward(self, u0, grid):\n a = self.branch(u0)\n # batchsize x width\n b = self.trunk(grid)\n # N x width\n return torch.einsum('bi,ni->bn', a, b)\n\n\nclass SAWeight(nn.Module):\n def __init__(self, out_dim, num_init: List, num_bd: List, num_collo: List):\n super(SAWeight, self).__init__()\n self.init_param = nn.ParameterList(\n [nn.Parameter(100 * torch.rand(num, out_dim)) for num in num_init]\n )\n\n self.bd_param = nn.ParameterList(\n [nn.Parameter(torch.rand(num, out_dim)) for num in num_bd]\n )\n\n self.collo_param = nn.ParameterList(\n [nn.Parameter(torch.rand(num, out_dim)) for num in num_collo]\n )\n\n def forward(self, init_cond: List, bd_cond: List, residual: List):\n total_loss = 0.0\n for param, init_loss in zip(self.init_param, init_cond):\n total_loss += weighted_mse(init_loss, 0, param)\n\n for param, bd in zip(self.bd_param, bd_cond):\n total_loss += weighted_mse(bd, 0, param)\n\n for param, res in zip(self.collo_param, residual):\n total_loss += weighted_mse(res, 0, param)\n return total_loss" }, { "path": "baselines/pinns_ns_05s.py", "content": "'''\ntraining for Navier Stokes with Reynolds number 500, 0.5 second time period\n'''\nimport csv\nimport random\nfrom timeit import default_timer\nimport deepxde as dde\nfrom deepxde.optimizers.config import set_LBFGS_options\nimport numpy as np\nfrom baselines.data import NSdata\n\nimport tensorflow as tf\n\nRe = 500\n\n\ndef forcing(x):\n return - 4 * tf.math.cos(4 * x[:, 1:2])\n\n\ndef pde(x, u):\n '''\n Args:\n x: (x, y, t)\n u: (u, v, w), where (u,v) is the velocity, w is the vorticity\n Returns: list of pde loss\n\n '''\n u_vel, v_vel, w = u[:, 0:1], u[:, 1:2], u[:, 2:3]\n\n u_vel_x = dde.grad.jacobian(u, x, i=0, j=0)\n u_vel_xx = dde.grad.hessian(u, x, component=0, i=0, j=0)\n u_vel_yy = dde.grad.hessian(u, x, component=0, i=1, j=1)\n\n v_vel_y = dde.grad.jacobian(u, x, i=1, j=1)\n v_vel_xx = dde.grad.hessian(u, x, component=1, i=0, j=0)\n v_vel_yy = dde.grad.hessian(u, x, component=1, i=1, j=1)\n\n w_vor_x = dde.grad.jacobian(u, x, i=2, j=0)\n w_vor_y = dde.grad.jacobian(u, x, i=2, j=1)\n w_vor_t = dde.grad.jacobian(u, x, i=2, j=2)\n\n w_vor_xx = dde.grad.hessian(u, x, component=2, i=0, j=0)\n w_vor_yy = dde.grad.hessian(u, x, component=2, i=1, j=1)\n\n eqn1 = w_vor_t + u_vel * w_vor_x + v_vel * w_vor_y - \\\n 1 / Re * (w_vor_xx + w_vor_yy) - forcing(x)\n eqn2 = u_vel_x + v_vel_y\n eqn3 = u_vel_xx + u_vel_yy + w_vor_y\n eqn4 = v_vel_xx + v_vel_yy - w_vor_x\n return [eqn1, eqn2, eqn3, eqn4]\n\n\ndef eval(model, dataset,\n step, time_cost,\n offset, config):\n '''\n evaluate test error for the model over dataset\n '''\n test_points, test_vals = dataset.get_test_xyt()\n\n pred = model.predict(test_points)\n vel_u_truth = test_vals[:, 0]\n vel_v_truth = test_vals[:, 1]\n vor_truth = test_vals[:, 2]\n\n vel_u_pred = pred[:, 0]\n vel_v_pred = pred[:, 1]\n vor_pred = pred[:, 2]\n\n u_err = dde.metrics.l2_relative_error(vel_u_truth, vel_u_pred)\n v_err = dde.metrics.l2_relative_error(vel_v_truth, vel_v_pred)\n vor_err = dde.metrics.l2_relative_error(vor_truth, vor_pred)\n print(f'Instance index : {offset}')\n print(f'L2 relative error in u: {u_err}')\n print(f'L2 relative error in v: {v_err}')\n print(f'L2 relative error in vorticity: {vor_err}')\n with open(config['log']['logfile'], 'a') as f:\n writer = csv.writer(f)\n writer.writerow([offset, u_err, v_err, vor_err, step, time_cost])\n\n\ndef train(offset, config, args):\n seed = random.randint(1, 10000)\n print(f'Random seed :{seed}')\n np.random.seed(seed)\n # construct dataloader\n data_config = config['data']\n if 'datapath2' in data_config:\n dataset = NSdata(datapath1=data_config['datapath'],\n datapath2=data_config['datapath2'],\n offset=offset, num=1,\n nx=data_config['nx'], nt=data_config['nt'],\n sub=data_config['sub'], sub_t=data_config['sub_t'],\n vel=True,\n t_interval=data_config['time_interval'])\n else:\n dataset = NSdata(datapath1=data_config['datapath'],\n offset=offset, num=1,\n nx=data_config['nx'], nt=data_config['nt'],\n sub=data_config['sub'], sub_t=data_config['sub_t'],\n vel=True,\n t_interval=data_config['time_interval'])\n spatial_domain = dde.geometry.Rectangle(xmin=[0, 0], xmax=[2 * np.pi, 2 * np.pi])\n temporal_domain = dde.geometry.TimeDomain(0, data_config['time_interval'])\n st_domain = dde.geometry.GeometryXTime(spatial_domain, temporal_domain)\n num_boundary_points = dataset.S\n points = dataset.get_boundary_points(num_x=num_boundary_points, num_y=num_boundary_points,\n num_t=dataset.T)\n u_value = dataset.get_boundary_value(component=0)\n v_value = dataset.get_boundary_value(component=1)\n w_value = dataset.get_boundary_value(component=2)\n # u, v are velocity, w is vorticity\n boundary_u = dde.PointSetBC(points=points, values=u_value, component=0)\n boundary_v = dde.PointSetBC(points=points, values=v_value, component=1)\n boundary_w = dde.PointSetBC(points=points, values=w_value, component=2)\n\n data = dde.data.TimePDE(\n st_domain,\n pde,\n [\n boundary_u,\n boundary_v,\n boundary_w\n ],\n num_domain=config['train']['num_domain'],\n num_boundary=config['train']['num_boundary'],\n num_test=config['train']['num_test'],\n )\n\n net = dde.maps.FNN(config['model']['layers'],\n config['model']['activation'],\n 'Glorot normal')\n # net = dde.maps.STMsFFN([3] + 4 * [50] + [3], 'tanh', 'Glorot normal', [50], [50])\n model = dde.Model(data, net)\n\n model.compile('adam', lr=config['train']['base_lr'], loss_weights=[1, 1, 1, 1, 100, 100, 100])\n if 'log_step' in config['train']:\n step_size = config['train']['log_step']\n else:\n step_size = 100\n epochs = config['train']['epochs'] // step_size\n\n for i in range(epochs):\n time_start = default_timer()\n model.train(epochs=step_size, display_every=step_size)\n time_end = default_timer()\n eval(model, dataset, i * step_size,\n time_cost=time_end - time_start,\n offset=offset,\n config=config)\n print('Done!')\n # set_LBFGS_options(maxiter=10000)\n # model.compile('L-BFGS', loss_weights=[1, 1, 1, 1, 100, 100, 100])\n # model.train()\n\n # test_points, test_vals = dataset.get_test_xyt()\n #\n # pred = model.predict(test_points)\n # vel_u_truth = test_vals[:, 0]\n # vel_v_truth = test_vals[:, 1]\n # vor_truth = test_vals[:, 2]\n #\n # vel_u_pred = pred[:, 0]\n # vel_v_pred = pred[:, 1]\n # vor_pred = pred[:, 2]\n #\n # u_err = dde.metrics.l2_relative_error(vel_u_truth, vel_u_pred)\n # v_err = dde.metrics.l2_relative_error(vel_v_truth, vel_v_pred)\n # vor_err = dde.metrics.l2_relative_error(vor_truth, vor_pred)\n # print(f'Instance index : {offset}')\n # print(f'L2 relative error in u: {u_err}')\n # print(f'L2 relative error in v: {v_err}')\n # print(f'L2 relative error in vorticity: {vor_err}')\n # with open(args.logfile, 'a') as f:\n # writer = csv.writer(f)\n # writer.writerow([offset, u_err, v_err, vor_err])\n" }, { "path": "baselines/pinns_ns_50s.py", "content": "'''\ntraining for Navier Stokes with viscosity 0.001\nspatial domain: (0, 1) ** 2\ntemporal domain: [0, 49]\n'''\nimport csv\nimport random\nfrom timeit import default_timer\nimport deepxde as dde\nfrom deepxde.optimizers.config import set_LBFGS_options\nimport numpy as np\nfrom baselines.data import NSLong\n\nimport tensorflow as tf\n\n\ndef forcing(x):\n theta = x[:, 0:1] + x[:, 1:2]\n return 0.1 * (tf.math.sin(2 * np.pi * theta) + tf.math.cos(2 * np.pi * theta))\n\n\ndef pde(x, u):\n '''\n Args:\n x: (x, y, t)\n u: (u, v, w), where (u,v) is the velocity, w is the vorticity\n Returns: list of pde loss\n\n '''\n u_vel, v_vel, w = u[:, 0:1], u[:, 1:2], u[:, 2:3]\n\n u_vel_x = dde.grad.jacobian(u, x, i=0, j=0)\n u_vel_xx = dde.grad.hessian(u, x, component=0, i=0, j=0)\n u_vel_yy = dde.grad.hessian(u, x, component=0, i=1, j=1)\n\n v_vel_y = dde.grad.jacobian(u, x, i=1, j=1)\n v_vel_xx = dde.grad.hessian(u, x, component=1, i=0, j=0)\n v_vel_yy = dde.grad.hessian(u, x, component=1, i=1, j=1)\n\n w_vor_x = dde.grad.jacobian(u, x, i=2, j=0)\n w_vor_y = dde.grad.jacobian(u, x, i=2, j=1)\n w_vor_t = dde.grad.jacobian(u, x, i=2, j=2)\n\n w_vor_xx = dde.grad.hessian(u, x, component=2, i=0, j=0)\n w_vor_yy = dde.grad.hessian(u, x, component=2, i=1, j=1)\n\n eqn1 = w_vor_t + u_vel * w_vor_x + v_vel * w_vor_y - \\\n 0.001 * (w_vor_xx + w_vor_yy) - forcing(x)\n eqn2 = u_vel_x + v_vel_y\n eqn3 = u_vel_xx + u_vel_yy + w_vor_y\n eqn4 = v_vel_xx + v_vel_yy - w_vor_x\n return [eqn1, eqn2, eqn3, eqn4]\n\n\ndef eval(model, dataset,\n step, time_cost,\n offset, config):\n '''\n evaluate test error for the model over dataset\n '''\n test_points, test_vals = dataset.get_test_xyt()\n\n pred = model.predict(test_points)\n vel_u_truth = test_vals[:, 0]\n vel_v_truth = test_vals[:, 1]\n vor_truth = test_vals[:, 2]\n\n vel_u_pred = pred[:, 0]\n vel_v_pred = pred[:, 1]\n vor_pred = pred[:, 2]\n\n u_err = dde.metrics.l2_relative_error(vel_u_truth, vel_u_pred)\n v_err = dde.metrics.l2_relative_error(vel_v_truth, vel_v_pred)\n vor_err = dde.metrics.l2_relative_error(vor_truth, vor_pred)\n\n total_num = test_vals.shape[0]\n u50 = test_vals[dataset.T - 1: total_num: dataset.T, 0]\n v50 = test_vals[dataset.T - 1: total_num: dataset.T, 1]\n vor50 = test_vals[dataset.T - 1: total_num: dataset.T, 2]\n\n u50_pred = pred[dataset.T - 1: total_num: dataset.T, 0]\n v50_pred = pred[dataset.T - 1: total_num: dataset.T, 1]\n vor50_pred = pred[dataset.T - 1: total_num: dataset.T, 2]\n\n u50_err = dde.metrics.l2_relative_error(u50, u50_pred)\n v50_err = dde.metrics.l2_relative_error(v50, v50_pred)\n vor50_err = dde.metrics.l2_relative_error(vor50, vor50_pred)\n\n print(f'Instance index : {offset}')\n print(f'L2 relative error in u: {u_err}')\n print(f'L2 relative error in v: {v_err}')\n print(f'L2 relative error in vorticity: {vor_err}')\n\n print(f'Time {dataset.T - 1} L2 relative error of u : {u50_err}')\n print(f'Time {dataset.T - 1} L2 relative error of v : {v50_err}')\n print(f'Time {dataset.T - 1} L2 relative error of vor : {vor50_err}')\n with open(config['log']['logfile'], 'a') as f:\n writer = csv.writer(f)\n writer.writerow([offset, u_err, v_err, vor_err, step, time_cost, u50_err, v50_err, vor50_err])\n\n\ndef train_longtime(offset, config, args):\n seed = random.randint(1, 10000)\n print(f'Random seed :{seed}')\n np.random.seed(seed)\n # construct dataloader\n data_config = config['data']\n spatial_domain = dde.geometry.Rectangle(xmin=[0, 0], xmax=[1, 1])\n temporal_domain = dde.geometry.TimeDomain(0, data_config['time_scale'])\n st_domain = dde.geometry.GeometryXTime(spatial_domain, temporal_domain)\n\n dataset = NSLong(datapath=data_config['datapath'],\n nx=data_config['nx'], nt=data_config['nt'],\n time_scale=data_config['time_scale'],\n offset=offset, num=data_config['n_sample'],\n vel=True)\n\n points = dataset.get_boundary_points(dataset.S, dataset.S, dataset.T)\n u_value = dataset.get_boundary_value(component=0)\n v_value = dataset.get_boundary_value(component=1)\n w_value = dataset.get_boundary_value(component=2)\n # u, v are velocity, w is vorticity\n boundary_u = dde.PointSetBC(points=points, values=u_value, component=0)\n boundary_v = dde.PointSetBC(points=points, values=v_value, component=1)\n boundary_w = dde.PointSetBC(points=points, values=w_value, component=2)\n\n data = dde.data.TimePDE(\n st_domain,\n pde,\n [\n boundary_u,\n boundary_v,\n boundary_w\n ],\n num_domain=config['train']['num_domain'],\n num_boundary=config['train']['num_boundary'],\n num_test=config['train']['num_test'],\n )\n\n net = dde.maps.FNN(config['model']['layers'],\n config['model']['activation'],\n 'Glorot normal')\n # net = dde.maps.STMsFFN([3] + 4 * [50] + [3], 'tanh', 'Glorot normal', [50], [50])\n model = dde.Model(data, net)\n\n model.compile('adam', lr=config['train']['base_lr'], loss_weights=[1, 1, 1, 1, 100, 100, 100])\n if 'log_step' in config['train']:\n step_size = config['train']['log_step']\n else:\n step_size = 100\n epochs = config['train']['epochs'] // step_size\n for i in range(epochs):\n time_start = default_timer()\n model.train(epochs=step_size, display_every=step_size)\n time_end = default_timer()\n eval(model, dataset, i * step_size,\n time_cost=time_end - time_start,\n offset=offset,\n config=config)\n print('Done!')\n\n" }, { "path": "baselines/sapinns-50s.py", "content": "import csv\nimport random\nfrom timeit import default_timer\nfrom tqdm import tqdm\nimport deepxde as dde\nimport numpy as np\nfrom baselines.data import NSdata\nimport torch\nfrom torch.optim import Adam\n\nfrom tensordiffeq.boundaries import DomainND, periodicBC\nfrom .tqd_utils import PointsIC\nfrom .model import SAWeight\n\nfrom models.FCN import DenseNet\nfrom train_utils.negadam import NAdam\n\n\n\ndef forcing(x):\n theta = x[:, 0:1] + x[:, 1:2]\n return 0.1 * (torch.sin(2 * np.pi * theta) + torch.cos(2 * np.pi * theta))\n\n\ndef pde(x, u):\n '''\n Args:\n x: (x, y, t)\n u: (u, v, w), where (u,v) is the velocity, w is the vorticity\n Returns: list of pde loss\n\n '''\n u_vel, v_vel, w = u[:, 0:1], u[:, 1:2], u[:, 2:3]\n\n u_vel_x = dde.grad.jacobian(u, x, i=0, j=0)\n u_vel_xx = dde.grad.hessian(u, x, component=0, i=0, j=0)\n u_vel_yy = dde.grad.hessian(u, x, component=0, i=1, j=1)\n\n v_vel_y = dde.grad.jacobian(u, x, i=1, j=1)\n v_vel_xx = dde.grad.hessian(u, x, component=1, i=0, j=0)\n v_vel_yy = dde.grad.hessian(u, x, component=1, i=1, j=1)\n\n w_vor_x = dde.grad.jacobian(u, x, i=2, j=0)\n w_vor_y = dde.grad.jacobian(u, x, i=2, j=1)\n w_vor_t = dde.grad.jacobian(u, x, i=2, j=2)\n\n w_vor_xx = dde.grad.hessian(u, x, component=2, i=0, j=0)\n w_vor_yy = dde.grad.hessian(u, x, component=2, i=1, j=1)\n\n eqn1 = w_vor_t + u_vel * w_vor_x + v_vel * w_vor_y - \\\n 0.001 * (w_vor_xx + w_vor_yy) - forcing(x)\n eqn2 = u_vel_x + v_vel_y\n eqn3 = u_vel_xx + u_vel_yy + w_vor_y\n eqn4 = v_vel_xx + v_vel_yy - w_vor_x\n return [eqn1, eqn2, eqn3, eqn4]\n\n\n\ndef eval(model, dataset,\n step, time_cost,\n offset, config):\n '''\n evaluate test error for the model over dataset\n '''\n test_points, test_vals = dataset.get_test_xyt()\n\n test_points = torch.tensor(test_points, dtype=torch.float32)\n with torch.no_grad():\n pred = model(test_points).cpu().numpy()\n vel_u_truth = test_vals[:, 0]\n vel_v_truth = test_vals[:, 1]\n vor_truth = test_vals[:, 2]\n\n vel_u_pred = pred[:, 0]\n vel_v_pred = pred[:, 1]\n vor_pred = pred[:, 2]\n\n u_err = dde.metrics.l2_relative_error(vel_u_truth, vel_u_pred)\n v_err = dde.metrics.l2_relative_error(vel_v_truth, vel_v_pred)\n vor_err = dde.metrics.l2_relative_error(vor_truth, vor_pred)\n print(f'Instance index : {offset}')\n print(f'L2 relative error in u: {u_err}')\n print(f'L2 relative error in v: {v_err}')\n print(f'L2 relative error in vorticity: {vor_err}')\n with open(config['log']['logfile'], 'a') as f:\n writer = csv.writer(f)\n writer.writerow([offset, u_err, v_err, vor_err, step, time_cost])\n\n\ndef train_sapinn(offset, config, args):\n seed = random.randint(1, 10000)\n print(f'Random seed :{seed}')\n np.random.seed(seed)\n # construct dataloader\n data_config = config['data']\n if 'datapath2' in data_config:\n dataset = NSdata(datapath1=data_config['datapath'],\n datapath2=data_config['datapath2'],\n offset=offset, num=1,\n nx=data_config['nx'], nt=data_config['nt'],\n sub=data_config['sub'], sub_t=data_config['sub_t'],\n vel=True,\n t_interval=data_config['time_interval'])\n else:\n dataset = NSdata(datapath1=data_config['datapath'],\n offset=offset, num=1,\n nx=data_config['nx'], nt=data_config['nt'],\n sub=data_config['sub'], sub_t=data_config['sub_t'],\n vel=True,\n t_interval=data_config['time_interval'])\n domain = DomainND(['x', 'y', 't'], time_var='t')\n domain.add('x', [0.0, 2 * np.pi], dataset.S)\n domain.add('y', [0.0, 2 * np.pi], dataset.S)\n domain.add('t', [0.0, data_config['time_interval']], dataset.T)\n num_collo = config['train']['num_domain']\n domain.generate_collocation_points(num_collo)\n init_vals = dataset.get_init_cond()\n num_inits = config['train']['num_init']\n if num_inits > dataset.S ** 2:\n num_inits = dataset.S ** 2\n init_cond = PointsIC(domain, init_vals, var=['x', 'y'], n_values=num_inits)\n bd_cond = periodicBC(domain, ['x', 'y'], n_values=config['train']['num_boundary'])\n\n # prepare initial condition inputs\n init_input = torch.tensor(init_cond.input, dtype=torch.float32)\n init_val = torch.tensor(init_cond.val, dtype=torch.float32)\n\n # prepare boundary condition inputs\n upper_input0 = torch.tensor(bd_cond.upper[0], dtype=torch.float32).squeeze().t() # shape N x 3\n upper_input1 = torch.tensor(bd_cond.upper[1], dtype=torch.float32).squeeze().t()\n lower_input0 = torch.tensor(bd_cond.lower[0], dtype=torch.float32).squeeze().t()\n lower_input1 = torch.tensor(bd_cond.lower[1], dtype=torch.float32).squeeze().t()\n\n # prepare collocation points\n collo_input = torch.tensor(domain.X_f, dtype=torch.float32, requires_grad=True)\n weight_net = SAWeight(out_dim=3,\n num_init=[num_inits],\n num_bd=[upper_input0.shape[0]] * 2,\n num_collo=[num_collo] * 4)\n net = DenseNet(config['model']['layers'], config['model']['activation'])\n weight_optim = NAdam(weight_net.parameters(), lr=config['train']['base_lr'])\n net_optim = Adam(net.parameters(), lr=config['train']['base_lr'])\n\n pbar = tqdm(range(config['train']['epochs']), dynamic_ncols=True)\n\n start_time = default_timer()\n for e in pbar:\n net.zero_grad()\n weight_net.zero_grad()\n if collo_input.grad is not None:\n collo_input.grad.zero_()\n\n init_pred = net(init_input) - init_val\n\n bd_0 = net(upper_input0) - net(lower_input0)\n bd_1 = net(upper_input1) - net(lower_input1)\n\n predu = net(collo_input)\n pde_residual = pde(collo_input, predu)\n\n loss = weight_net(init_cond=[init_pred], bd_cond=[bd_0, bd_1], residual=pde_residual)\n loss.backward()\n weight_optim.step()\n net_optim.step()\n pbar.set_description(\n (\n f'Epoch: {e}, loss: {loss.item()}'\n )\n )\n if e % config['train']['log_step'] == 0:\n end_time = default_timer()\n eval(net, dataset, e, time_cost=end_time - start_time, offset=offset, config=config)\n start_time = default_timer()\n print('Done!')\n\n\n\n\n\n \n" }, { "path": "baselines/sapinns.py", "content": "import csv\nimport random\nfrom timeit import default_timer\nfrom tqdm import tqdm\nimport deepxde as dde\nimport numpy as np\nfrom baselines.data import NSdata\nimport torch\nfrom torch.optim import Adam\n\nfrom tensordiffeq.boundaries import DomainND, periodicBC\nfrom .tqd_utils import PointsIC\nfrom .model import SAWeight\n\nfrom models.FCN import DenseNet\nfrom train_utils.negadam import NAdam\n\n\nRe = 500\n\n\ndef forcing(x):\n return - 4 * torch.cos(4 * x[:, 1:2])\n\n\ndef pde(x, u):\n '''\n Args:\n x: (x, y, t)\n u: (u, v, w), where (u,v) is the velocity, w is the vorticity\n Returns: list of pde loss\n\n '''\n u_vel, v_vel, w = u[:, 0:1], u[:, 1:2], u[:, 2:3]\n\n u_vel_x = dde.grad.jacobian(u, x, i=0, j=0)\n u_vel_xx = dde.grad.hessian(u, x, component=0, i=0, j=0)\n u_vel_yy = dde.grad.hessian(u, x, component=0, i=1, j=1)\n\n v_vel_y = dde.grad.jacobian(u, x, i=1, j=1)\n v_vel_xx = dde.grad.hessian(u, x, component=1, i=0, j=0)\n v_vel_yy = dde.grad.hessian(u, x, component=1, i=1, j=1)\n\n w_vor_x = dde.grad.jacobian(u, x, i=2, j=0)\n w_vor_y = dde.grad.jacobian(u, x, i=2, j=1)\n w_vor_t = dde.grad.jacobian(u, x, i=2, j=2)\n\n w_vor_xx = dde.grad.hessian(u, x, component=2, i=0, j=0)\n w_vor_yy = dde.grad.hessian(u, x, component=2, i=1, j=1)\n\n eqn1 = w_vor_t + u_vel * w_vor_x + v_vel * w_vor_y - \\\n 1 / Re * (w_vor_xx + w_vor_yy) - forcing(x)\n eqn2 = u_vel_x + v_vel_y\n eqn3 = u_vel_xx + u_vel_yy + w_vor_y\n eqn4 = v_vel_xx + v_vel_yy - w_vor_x\n return [eqn1, eqn2, eqn3, eqn4]\n\n\ndef eval(model, dataset,\n step, time_cost,\n offset, config):\n '''\n evaluate test error for the model over dataset\n '''\n test_points, test_vals = dataset.get_test_xyt()\n\n test_points = torch.tensor(test_points, dtype=torch.float32)\n with torch.no_grad():\n pred = model(test_points).cpu().numpy()\n vel_u_truth = test_vals[:, 0]\n vel_v_truth = test_vals[:, 1]\n vor_truth = test_vals[:, 2]\n\n vel_u_pred = pred[:, 0]\n vel_v_pred = pred[:, 1]\n vor_pred = pred[:, 2]\n\n u_err = dde.metrics.l2_relative_error(vel_u_truth, vel_u_pred)\n v_err = dde.metrics.l2_relative_error(vel_v_truth, vel_v_pred)\n vor_err = dde.metrics.l2_relative_error(vor_truth, vor_pred)\n print(f'Instance index : {offset}')\n print(f'L2 relative error in u: {u_err}')\n print(f'L2 relative error in v: {v_err}')\n print(f'L2 relative error in vorticity: {vor_err}')\n with open(config['log']['logfile'], 'a') as f:\n writer = csv.writer(f)\n writer.writerow([offset, u_err, v_err, vor_err, step, time_cost])\n\n\ndef train_sapinn(offset, config, args):\n seed = random.randint(1, 10000)\n print(f'Random seed :{seed}')\n np.random.seed(seed)\n # construct dataloader\n data_config = config['data']\n if 'datapath2' in data_config:\n dataset = NSdata(datapath1=data_config['datapath'],\n datapath2=data_config['datapath2'],\n offset=offset, num=1,\n nx=data_config['nx'], nt=data_config['nt'],\n sub=data_config['sub'], sub_t=data_config['sub_t'],\n vel=True,\n t_interval=data_config['time_interval'])\n else:\n dataset = NSdata(datapath1=data_config['datapath'],\n offset=offset, num=1,\n nx=data_config['nx'], nt=data_config['nt'],\n sub=data_config['sub'], sub_t=data_config['sub_t'],\n vel=True,\n t_interval=data_config['time_interval'])\n domain = DomainND(['x', 'y', 't'], time_var='t')\n domain.add('x', [0.0, 2 * np.pi], dataset.S)\n domain.add('y', [0.0, 2 * np.pi], dataset.S)\n domain.add('t', [0.0, data_config['time_interval']], dataset.T)\n num_collo = config['train']['num_domain']\n domain.generate_collocation_points(num_collo)\n init_vals = dataset.get_init_cond()\n num_inits = config['train']['num_init']\n if num_inits > dataset.S ** 2:\n num_inits = dataset.S ** 2\n init_cond = PointsIC(domain, init_vals, var=['x', 'y'], n_values=num_inits)\n bd_cond = periodicBC(domain, ['x', 'y'], n_values=config['train']['num_boundary'])\n\n # prepare initial condition inputs\n init_input = torch.tensor(init_cond.input, dtype=torch.float32)\n init_val = torch.tensor(init_cond.val, dtype=torch.float32)\n\n # prepare boundary condition inputs\n upper_input0 = torch.tensor(bd_cond.upper[0], dtype=torch.float32).squeeze().t() # shape N x 3\n upper_input1 = torch.tensor(bd_cond.upper[1], dtype=torch.float32).squeeze().t()\n lower_input0 = torch.tensor(bd_cond.lower[0], dtype=torch.float32).squeeze().t()\n lower_input1 = torch.tensor(bd_cond.lower[1], dtype=torch.float32).squeeze().t()\n\n # prepare collocation points\n collo_input = torch.tensor(domain.X_f, dtype=torch.float32, requires_grad=True)\n\n weight_net = SAWeight(out_dim=3,\n num_init=[num_inits],\n num_bd=[upper_input0.shape[0]] * 2,\n num_collo=[num_collo] * 4)\n net = DenseNet(config['model']['layers'], config['model']['activation'])\n weight_optim = NAdam(weight_net.parameters(), lr=config['train']['base_lr'])\n net_optim = Adam(net.parameters(), lr=config['train']['base_lr'])\n\n pbar = tqdm(range(config['train']['epochs']), dynamic_ncols=True)\n\n start_time = default_timer()\n for e in pbar:\n net.zero_grad()\n weight_net.zero_grad()\n if collo_input.grad is not None:\n collo_input.grad.zero_()\n\n init_pred = net(init_input) - init_val\n\n bd_0 = net(upper_input0) - net(lower_input0)\n bd_1 = net(upper_input1) - net(lower_input1)\n\n predu = net(collo_input)\n pde_residual = pde(collo_input, predu)\n\n loss = weight_net(init_cond=[init_pred], bd_cond=[bd_0, bd_1], residual=pde_residual)\n loss.backward()\n weight_optim.step()\n net_optim.step()\n dde.gradients.clear()\n pbar.set_description(\n (\n f'Epoch: {e}, loss: {loss.item()}'\n )\n )\n\n if e % config['train']['log_step'] == 0:\n end_time = default_timer()\n eval(net, dataset, e, time_cost=end_time - start_time, offset=offset, config=config)\n start_time = default_timer()\n print('Done!')\n\n\n\n\n\n \n" }, { "path": "baselines/test.py", "content": "from tqdm import tqdm\nimport numpy as np\n\nimport torch\nfrom torch.utils.data import DataLoader\nfrom baselines.model import DeepONetCP\nfrom baselines.data import DeepONetCPNS, DarcyFlow\nfrom train_utils.losses import LpLoss\n\n\ndef test(model,\n test_loader,\n grid,\n device):\n pbar = tqdm(test_loader, dynamic_ncols=True, smoothing=0.1)\n myloss = LpLoss(size_average=True)\n model.eval()\n\n test_error = []\n with torch.no_grad():\n for x, y in pbar:\n x = x.to(device)\n y = y.to(device)\n\n grid = grid.to(device)\n\n pred = model(x, grid)\n loss = myloss(pred, y)\n\n test_error.append(loss.item())\n pbar.set_description(\n (\n f'test error: {loss.item():.5f}'\n )\n )\n\n mean = np.mean(test_error)\n std = np.std(test_error, ddof=1) / np.sqrt(len(test_error))\n print(f'Averaged test error :{mean}, standard error: {std}')\n\n\ndef test_deeponet_ns(config):\n '''\n Evaluate deeponet model on Navier Stokes equation\n Args:\n config: configurations\n\n Returns:\n\n '''\n device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n data_config = config['data']\n batch_size = config['test']['batchsize']\n dataset = DeepONetCPNS(datapath=data_config['datapath'],\n nx=data_config['nx'], nt=data_config['nt'],\n sub=data_config['sub'], sub_t=data_config['sub_t'],\n offset=data_config['offset'], num=data_config['n_sample'],\n t_interval=data_config['time_interval'])\n test_loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)\n u0_dim = dataset.S ** 2\n model = DeepONetCP(branch_layer=[u0_dim] + config['model']['branch_layers'],\n trunk_layer=[3] + config['model']['trunk_layers']).to(device)\n if 'ckpt' in config['test']:\n ckpt = torch.load(config['test']['ckpt'])\n model.load_state_dict(ckpt['model'])\n grid = test_loader.dataset.xyt\n test(model, test_loader, grid, device=device)\n\n\ndef test_deeponet_darcy(config):\n '''\n Evaluate deeponet mode on Darcy Flow\n '''\n\n device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n data_config = config['data']\n batch_size = config['test']['batchsize']\n dataset = DarcyFlow(data_config['datapath'],\n nx=data_config['nx'], sub=data_config['sub'],\n offset=data_config['offset'], num=data_config['n_sample'])\n dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)\n\n u0_dim = dataset.S ** 2\n model = DeepONetCP(branch_layer=[u0_dim] + config['model']['branch_layers'],\n trunk_layer=[2] + config['model']['trunk_layers']).to(device)\n if 'ckpt' in config['test']:\n ckpt = torch.load(config['test']['ckpt'])\n model.load_state_dict(ckpt['model'])\n print('Load model weights from %s' % config['test']['ckpt'])\n\n grid = dataset.mesh.reshape(-1, 2)\n test(model, dataloader, grid, device)" }, { "path": "baselines/tqd_sapinns.py", "content": "import random\nimport numpy as np\nimport csv\nfrom timeit import default_timer\n\nimport tensorflow as tf\nimport deepxde as dde\nimport tensordiffeq as tdq\nfrom tensordiffeq.models import CollocationSolverND\nfrom tensordiffeq.boundaries import DomainND, periodicBC\n\nfrom .tqd_utils import PointsIC\nfrom baselines.data import NSdata\n\n\nRe = 500\n\n\ndef forcing(x):\n return - 4 * tf.math.cos(4 * x)\n\n\ndef bd_model(u_model, x, y, t):\n u = u_model(tf.concat([x, y, t], 1))\n u_vel, v_vel, w = u[:, 0:1], u[:, 1:2], u[:, 2:3]\n return u_vel, v_vel, w\n\n\ndef f_model(u_model, x, y, t):\n inp = tf.concat([x, y, t], 1)\n u = u_model(inp)\n u_vel, v_vel, w = u[:, 0:1], u[:, 1:2], u[:, 2:3]\n\n u_vel_x = tf.gradients(u_vel, x)[0]\n u_vel_xx = tf.gradients(u_vel_x, x)[0]\n u_vel_y = tf.gradients(u_vel, y)[0]\n u_vel_yy = tf.gradients(u_vel_y, y)[0]\n\n v_vel_y = tf.gradients(v_vel, y)[0]\n v_vel_x = tf.gradients(v_vel, x)[0]\n v_vel_xx = tf.gradients(v_vel_x, x)[0]\n v_vel_yy = tf.gradients(v_vel_y, y)[0]\n\n w_vor_x = tf.gradients(w, x)[0]\n w_vor_y = tf.gradients(w, y)[0]\n w_vor_t = tf.gradients(w, t)[0]\n\n w_vor_xx = tf.gradients(w_vor_x, x)[0]\n w_vor_yy = tf.gradients(w_vor_y, y)[0]\n\n c1 = tdq.utils.constant(1 / Re)\n eqn1 = w_vor_t + u_vel * w_vor_x + v_vel * w_vor_y - c1 * (w_vor_xx + w_vor_yy) - forcing(x)\n eqn2 = u_vel_x + v_vel_y\n eqn3 = u_vel_xx + u_vel_yy + w_vor_y\n eqn4 = v_vel_xx + v_vel_yy - w_vor_x\n return eqn1, eqn2, eqn3, eqn4\n\n\ndef eval(model, dataset,\n step, time_cost,\n offset, config):\n '''\n evaluate test error for the model over dataset\n '''\n test_points, test_vals = dataset.get_test_xyt()\n\n pred = model.predict(test_points)\n vel_u_truth = test_vals[:, 0]\n vel_v_truth = test_vals[:, 1]\n vor_truth = test_vals[:, 2]\n\n vel_u_pred = pred[:, 0]\n vel_v_pred = pred[:, 1]\n vor_pred = pred[:, 2]\n\n u_err = dde.metrics.l2_relative_error(vel_u_truth, vel_u_pred)\n v_err = dde.metrics.l2_relative_error(vel_v_truth, vel_v_pred)\n vor_err = dde.metrics.l2_relative_error(vor_truth, vor_pred)\n print(f'Instance index : {offset}')\n print(f'L2 relative error in u: {u_err}')\n print(f'L2 relative error in v: {v_err}')\n print(f'L2 relative error in vorticity: {vor_err}')\n with open(config['log']['logfile'], 'a') as f:\n writer = csv.writer(f)\n writer.writerow([offset, u_err, v_err, vor_err, step, time_cost])\n\n\ndef train_sa(offset, config, args):\n seed = random.randint(1, 10000)\n print(f'Random seed :{seed}')\n np.random.seed(seed)\n # construct dataloader\n data_config = config['data']\n if 'datapath2' in data_config:\n dataset = NSdata(datapath1=data_config['datapath'],\n datapath2=data_config['datapath2'],\n offset=offset, num=1,\n nx=data_config['nx'], nt=data_config['nt'],\n sub=data_config['sub'], sub_t=data_config['sub_t'],\n vel=True,\n t_interval=data_config['time_interval'])\n else:\n dataset = NSdata(datapath1=data_config['datapath'],\n offset=offset, num=1,\n nx=data_config['nx'], nt=data_config['nt'],\n sub=data_config['sub'], sub_t=data_config['sub_t'],\n vel=True,\n t_interval=data_config['time_interval'])\n domain = DomainND(['x', 'y', 't'], time_var='t')\n domain.add('x', [0.0, 2 * np.pi], dataset.S)\n domain.add('y', [0.0, 2 * np.pi], dataset.S)\n domain.add('t', [0.0, data_config['time_interval']], dataset.T)\n domain.generate_collocation_points(config['train']['num_domain'])\n model = CollocationSolverND()\n init_vals = dataset.get_init_cond()\n num_inits = config['train']['num_init']\n if num_inits > dataset.S ** 2:\n num_inits = dataset.S ** 2\n init_cond = PointsIC(domain, init_vals, var=['x', 'y'], n_values=num_inits)\n bd_cond = periodicBC(domain, ['x', 'y'], [bd_model], n_values=config['train']['num_boundary'])\n BCs = [init_cond, bd_cond]\n\n dict_adaptive = {'residual': [True, True, True, True],\n 'BCs': [True, False]}\n init_weights = {\n 'residual': [tf.random.uniform([config['train']['num_domain'], 1]),\n tf.random.uniform([config['train']['num_domain'], 1]),\n tf.random.uniform([config['train']['num_domain'], 1]),\n tf.random.uniform([config['train']['num_domain'], 1])],\n 'BCs': [100 * tf.random.uniform([num_inits, 1]),\n 100 * tf.ones([config['train']['num_boundary'], 1])]\n }\n\n model.compile(config['model']['layers'], f_model, domain, BCs,\n isAdaptive=True, dict_adaptive=dict_adaptive, init_weights=init_weights)\n\n if 'log_step' in config['train']:\n step_size = config['train']['log_step']\n else:\n step_size = 100\n epochs = config['train']['epochs'] // step_size\n\n for i in range(epochs):\n time_start = default_timer()\n model.fit(tf_iter=step_size)\n time_end = default_timer()\n eval(model, dataset, i * step_size,\n time_cost=time_end - time_start,\n offset=offset,\n config=config)\n print('Done!')\n" }, { "path": "baselines/tqd_utils.py", "content": "import numpy as np\n\nfrom tensordiffeq.boundaries import BC\nfrom tensordiffeq.utils import flatten_and_stack, multimesh, MSE, convertTensor\n\n\nclass PointsIC(BC):\n '''\n Create Initial condition class from array on domain\n '''\n def __init__(self, domain, values, var, n_values=None):\n '''\n args:\n - domain:\n - values:\n '''\n super(PointsIC, self).__init__()\n self.isInit = True\n self.n_values = n_values\n self.domain = domain\n self.values = values\n self.vars = var\n self.isInit = True\n self.dicts_ = [item for item in self.domain.domaindict if item['identifier'] != self.domain.time_var]\n self.dict_ = next(item for item in self.domain.domaindict if item[\"identifier\"] == self.domain.time_var)\n self.compile()\n self.create_target(self.values)\n\n def create_input(self):\n dims = self.get_not_dims(self.domain.time_var)\n mesh = flatten_and_stack(multimesh(dims))\n t_repeat = np.repeat(0.0, len(mesh))\n\n mesh = np.concatenate((mesh, np.reshape(t_repeat, (-1, 1))), axis=1)\n if self.n_values is not None:\n self.nums = np.random.randint(0, high=len(mesh), size=self.n_values)\n mesh = mesh[self.nums]\n return mesh\n\n def create_target(self, values):\n # for i, var_ in enumerate(self.vars):\n # arg_list = []\n # for j, var in enumerate(var_):\n # var_dict = self.get_dict(var)\n # arg_list.append(get_linspace(var_dict))\n # inp = flatten_and_stack(multimesh(arg_list))\n # fun_vals.append(self.fun[i](*inp.T))\n if self.n_values is not None:\n self.val = np.reshape(values, (-1, 3))[self.nums]\n else:\n self.val = np.reshape(values, (-1, 3))\n\n def loss(self):\n return MSE(self.preds, self.val)" }, { "path": "baselines/train_darcy.py", "content": "from tqdm import tqdm\n\nimport torch\nfrom torch.utils.data import DataLoader\nfrom torch.optim import Adam\nfrom torch.optim.lr_scheduler import MultiStepLR\n\nfrom baselines.model import DeepONetCP\nfrom train_utils.losses import LpLoss\nfrom train_utils.utils import save_checkpoint\nfrom baselines.data import DarcyFlow\n\n\ndef train_deeponet_darcy(config):\n '''\n train deepONet for darcy flow\n '''\n device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n data_config = config['data']\n batch_size = config['train']['batchsize']\n dataset = DarcyFlow(data_config['datapath'],\n nx=data_config['nx'], sub=data_config['sub'],\n offset=data_config['offset'], num=data_config['n_sample'])\n dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)\n\n u0_dim = dataset.S ** 2\n model = DeepONetCP(branch_layer=[u0_dim] + config['model']['branch_layers'],\n trunk_layer=[2] + config['model']['trunk_layers']).to(device)\n optimizer = Adam(model.parameters(), lr=config['train']['base_lr'])\n scheduler = MultiStepLR(optimizer, milestones=config['train']['milestones'],\n gamma=config['train']['scheduler_gamma'])\n pbar = range(config['train']['epochs'])\n pbar = tqdm(pbar, dynamic_ncols=True, smoothing=0.1)\n myloss = LpLoss(size_average=True)\n model.train()\n grid = dataset.mesh\n grid = grid.reshape(-1, 2).to(device) # grid value, (SxS, 2)\n for e in pbar:\n train_loss = 0.0\n for x, y in dataloader:\n x = x.to(device) # initial condition, (batchsize, u0_dim)\n\n y = y.to(device) # ground truth, (batchsize, SxS)\n\n pred = model(x, grid)\n loss = myloss(pred, y)\n\n model.zero_grad()\n loss.backward()\n optimizer.step()\n\n train_loss += loss.item() * y.shape[0]\n train_loss /= len(dataset)\n scheduler.step()\n\n pbar.set_description(\n (\n f'Epoch: {e}; Averaged train loss: {train_loss:.5f}; '\n )\n )\n if e % 1000 == 0:\n print(f'Epoch: {e}, averaged train loss: {train_loss:.5f}')\n save_checkpoint(config['train']['save_dir'],\n config['train']['save_name'].replace('.pt', f'_{e}.pt'),\n model, optimizer)\n save_checkpoint(config['train']['save_dir'],\n config['train']['save_name'],\n model, optimizer)" }, { "path": "baselines/train_ns.py", "content": "from tqdm import tqdm\n\nimport torch\nfrom torch.utils.data import DataLoader\nfrom torch.optim import Adam\nfrom torch.optim.lr_scheduler import MultiStepLR\n\nfrom baselines.model import DeepONet, DeepONetCP\nfrom baselines.data import DeepOnetNS, DeepONetCPNS\nfrom train_utils.losses import LpLoss\nfrom train_utils.utils import save_checkpoint\nfrom train_utils.data_utils import sample_data\n\n\ndef train_deeponet_cp(config):\n '''\n Train Cartesian product DeepONet\n Args:\n config:\n\n Returns:\n '''\n device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n data_config = config['data']\n batch_size = config['train']['batchsize']\n dataset = DeepONetCPNS(datapath=data_config['datapath'],\n nx=data_config['nx'], nt=data_config['nt'],\n sub=data_config['sub'], sub_t=data_config['sub_t'],\n offset=data_config['offset'], num=data_config['n_sample'],\n t_interval=data_config['time_interval'])\n train_loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)\n u0_dim = dataset.S ** 2\n model = DeepONetCP(branch_layer=[u0_dim] + config['model']['branch_layers'],\n trunk_layer=[3] + config['model']['trunk_layers']).to(device)\n optimizer = Adam(model.parameters(), lr=config['train']['base_lr'])\n scheduler = MultiStepLR(optimizer, milestones=config['train']['milestones'],\n gamma=config['train']['scheduler_gamma'])\n pbar = range(config['train']['epochs'])\n pbar = tqdm(pbar, dynamic_ncols=True, smoothing=0.1)\n myloss = LpLoss(size_average=True)\n model.train()\n\n for e in pbar:\n train_loss = 0.0\n for x, y in train_loader:\n x = x.to(device) # initial condition, (batchsize, u0_dim)\n grid = dataset.xyt\n grid = grid.to(device) # grid value, (SxSxT, 3)\n y = y.to(device) # ground truth, (batchsize, SxSxT)\n\n pred = model(x, grid)\n loss = myloss(pred, y)\n\n model.zero_grad()\n loss.backward()\n optimizer.step()\n\n train_loss += loss.item() * y.shape[0]\n train_loss /= len(dataset)\n scheduler.step()\n\n pbar.set_description(\n (\n f'Epoch: {e}; Averaged train loss: {train_loss:.5f}; '\n )\n )\n if e % 500 == 0:\n print(f'Epoch: {e}, averaged train loss: {train_loss:.5f}')\n save_checkpoint(config['train']['save_dir'],\n config['train']['save_name'].replace('.pt', f'_{e}.pt'),\n model, optimizer)\n\n\ndef train_deeponet(config):\n '''\n train plain DeepOnet\n Args:\n config:\n\n Returns:\n\n '''\n device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n data_config = config['data']\n dataset = DeepOnetNS(datapath=data_config['datapath'],\n nx=data_config['nx'], nt=data_config['nt'],\n sub=data_config['sub'], sub_t=data_config['sub_t'],\n offset=data_config['offset'], num=data_config['n_sample'],\n t_interval=data_config['time_interval'])\n train_loader = DataLoader(dataset, batch_size=config['train']['batchsize'], shuffle=False)\n\n u0_dim = dataset.S ** 2\n model = DeepONet(branch_layer=[u0_dim] + config['model']['branch_layers'],\n trunk_layer=[3] + config['model']['trunk_layers']).to(device)\n optimizer = Adam(model.parameters(), lr=config['train']['base_lr'])\n scheduler = MultiStepLR(optimizer, milestones=config['train']['milestones'],\n gamma=config['train']['scheduler_gamma'])\n\n pbar = range(config['train']['epochs'])\n pbar = tqdm(pbar, dynamic_ncols=True, smoothing=0.1)\n myloss = LpLoss(size_average=True)\n model.train()\n loader = sample_data(train_loader)\n for e in pbar:\n u0, x, y = next(loader)\n u0 = u0.to(device)\n x = x.to(device)\n y = y.to(device)\n pred = model(u0, x)\n loss = myloss(pred, y)\n model.zero_grad()\n loss.backward()\n optimizer.step()\n scheduler.step()\n\n pbar.set_description(\n (\n f'Epoch: {e}; Train loss: {loss.item():.5f}; '\n )\n )\n save_checkpoint(config['train']['save_dir'],\n config['train']['save_name'],\n model, optimizer)\n print('Done!')\n" }, { "path": "baselines/unet3d.py", "content": "from functools import partial\n\nimport torch\nfrom torch import nn as nn\nfrom torch.nn import functional as F\n\n\n\n# UNet3d from https://github.com/wolny/pytorch-3dunet\n\n\nclass BaseModel(nn.Module):\n def __init__(self):\n super().__init__()\n self.device_indicator_param = nn.Parameter(torch.empty(0))\n\n @property\n def device(self):\n \"\"\"Returns the device that the model is on.\"\"\"\n return self.device_indicator_param.device\n\n def data_dict_to_input(self, data_dict, **kwargs):\n \"\"\"\n Convert data dictionary to appropriate input for the model.\n \"\"\"\n raise NotImplementedError\n\n def loss_dict(self, data_dict, **kwargs):\n \"\"\"\n Compute the loss dictionary for the model.\n \"\"\"\n raise NotImplementedError\n\n @torch.no_grad()\n def eval_dict(self, data_dict, **kwargs):\n \"\"\"\n Compute the evaluation dictionary for the model.\n \"\"\"\n raise NotImplementedError\n\ndef create_conv(\n in_channels, out_channels, kernel_size, order, num_groups, padding, is3d\n):\n \"\"\"\n Create a list of modules with together constitute a single conv layer with non-linearity\n and optional batchnorm/groupnorm.\n\n Args:\n in_channels (int): number of input channels\n out_channels (int): number of output channels\n kernel_size(int or tuple): size of the convolving kernel\n order (string): order of things, e.g.\n 'cr' -> conv + ReLU\n 'gcr' -> groupnorm + conv + ReLU\n 'cl' -> conv + LeakyReLU\n 'ce' -> conv + ELU\n 'bcr' -> batchnorm + conv + ReLU\n num_groups (int): number of groups for the GroupNorm\n padding (int or tuple): add zero-padding added to all three sides of the input\n is3d (bool): is3d (bool): if True use Conv3d, otherwise use Conv2d\n Return:\n list of tuple (name, module)\n \"\"\"\n assert \"c\" in order, \"Conv layer MUST be present\"\n assert (\n order[0] not in \"rle\"\n ), \"Non-linearity cannot be the first operation in the layer\"\n\n modules = []\n for i, char in enumerate(order):\n if char == \"r\":\n modules.append((\"ReLU\", nn.ReLU(inplace=True)))\n elif char == \"l\":\n modules.append((\"LeakyReLU\", nn.LeakyReLU(inplace=True)))\n elif char == \"e\":\n modules.append((\"ELU\", nn.ELU(inplace=True)))\n elif char == \"c\":\n # add learnable bias only in the absence of batchnorm/groupnorm\n bias = not (\"g\" in order or \"b\" in order)\n if is3d:\n conv = nn.Conv3d(\n in_channels, out_channels, kernel_size, padding=padding, bias=bias\n )\n else:\n conv = nn.Conv2d(\n in_channels, out_channels, kernel_size, padding=padding, bias=bias\n )\n\n modules.append((\"conv\", conv))\n elif char == \"g\":\n is_before_conv = i < order.index(\"c\")\n if is_before_conv:\n num_channels = in_channels\n else:\n num_channels = out_channels\n\n # use only one group if the given number of groups is greater than the number of channels\n if num_channels < num_groups:\n num_groups = 1\n\n assert (\n num_channels % num_groups == 0\n ), f\"Expected number of channels in input to be divisible by num_groups. num_channels={num_channels}, num_groups={num_groups}\"\n modules.append(\n (\n \"groupnorm\",\n nn.GroupNorm(num_groups=num_groups, num_channels=num_channels),\n )\n )\n elif char == \"b\":\n is_before_conv = i < order.index(\"c\")\n if is3d:\n bn = nn.BatchNorm3d\n else:\n bn = nn.BatchNorm2d\n\n if is_before_conv:\n modules.append((\"batchnorm\", bn(in_channels)))\n else:\n modules.append((\"batchnorm\", bn(out_channels)))\n else:\n raise ValueError(\n f\"Unsupported layer type '{char}'. MUST be one of ['b', 'g', 'r', 'l', 'e', 'c']\"\n )\n\n return modules\n\n\nclass SingleConv(nn.Sequential):\n \"\"\"\n Basic convolutional module consisting of a Conv3d, non-linearity and optional batchnorm/groupnorm. The order\n of operations can be specified via the `order` parameter\n\n Args:\n in_channels (int): number of input channels\n out_channels (int): number of output channels\n kernel_size (int or tuple): size of the convolving kernel\n order (string): determines the order of layers, e.g.\n 'cr' -> conv + ReLU\n 'crg' -> conv + ReLU + groupnorm\n 'cl' -> conv + LeakyReLU\n 'ce' -> conv + ELU\n num_groups (int): number of groups for the GroupNorm\n padding (int or tuple): add zero-padding\n is3d (bool): if True use Conv3d, otherwise use Conv2d\n \"\"\"\n\n def __init__(\n self,\n in_channels,\n out_channels,\n kernel_size=3,\n order=\"gcr\",\n num_groups=8,\n padding=1,\n is3d=True,\n ):\n super(SingleConv, self).__init__()\n\n for name, module in create_conv(\n in_channels, out_channels, kernel_size, order, num_groups, padding, is3d\n ):\n self.add_module(name, module)\n\n\nclass DoubleConv(nn.Sequential):\n \"\"\"\n A module consisting of two consecutive convolution layers (e.g. BatchNorm3d+ReLU+Conv3d).\n We use (Conv3d+ReLU+GroupNorm3d) by default.\n This can be changed however by providing the 'order' argument, e.g. in order\n to change to Conv3d+BatchNorm3d+ELU use order='cbe'.\n Use padded convolutions to make sure that the output (H_out, W_out) is the same\n as (H_in, W_in), so that you don't have to crop in the decoder path.\n\n Args:\n in_channels (int): number of input channels\n out_channels (int): number of output channels\n encoder (bool): if True we're in the encoder path, otherwise we're in the decoder\n kernel_size (int or tuple): size of the convolving kernel\n order (string): determines the order of layers, e.g.\n 'cr' -> conv + ReLU\n 'crg' -> conv + ReLU + groupnorm\n 'cl' -> conv + LeakyReLU\n 'ce' -> conv + ELU\n num_groups (int): number of groups for the GroupNorm\n padding (int or tuple): add zero-padding added to all three sides of the input\n is3d (bool): if True use Conv3d instead of Conv2d layers\n \"\"\"\n\n def __init__(\n self,\n in_channels,\n out_channels,\n encoder,\n kernel_size=3,\n order=\"gcr\",\n num_groups=8,\n padding=1,\n is3d=True,\n ):\n super(DoubleConv, self).__init__()\n if encoder:\n # we're in the encoder path\n conv1_in_channels = in_channels\n conv1_out_channels = out_channels // 2\n if conv1_out_channels < in_channels:\n conv1_out_channels = in_channels\n conv2_in_channels, conv2_out_channels = conv1_out_channels, out_channels\n else:\n # we're in the decoder path, decrease the number of channels in the 1st convolution\n conv1_in_channels, conv1_out_channels = in_channels, out_channels\n conv2_in_channels, conv2_out_channels = out_channels, out_channels\n\n # conv1\n self.add_module(\n \"SingleConv1\",\n SingleConv(\n conv1_in_channels,\n conv1_out_channels,\n kernel_size,\n order,\n num_groups,\n padding=padding,\n is3d=is3d,\n ),\n )\n # conv2\n self.add_module(\n \"SingleConv2\",\n SingleConv(\n conv2_in_channels,\n conv2_out_channels,\n kernel_size,\n order,\n num_groups,\n padding=padding,\n is3d=is3d,\n ),\n )\n\n\nclass Encoder(nn.Module):\n \"\"\"\n A single module from the encoder path consisting of the optional max\n pooling layer (one may specify the MaxPool kernel_size to be different\n from the standard (2,2,2), e.g. if the volumetric data is anisotropic\n (make sure to use complementary scale_factor in the decoder path) followed by\n a basic module (DoubleConv or ResNetBlock).\n\n Args:\n in_channels (int): number of input channels\n out_channels (int): number of output channels\n conv_kernel_size (int or tuple): size of the convolving kernel\n apply_pooling (bool): if True use MaxPool3d before DoubleConv\n pool_kernel_size (int or tuple): the size of the window\n pool_type (str): pooling layer: 'max' or 'avg'\n basic_module(nn.Module): either ResNetBlock or DoubleConv\n conv_layer_order (string): determines the order of layers\n in `DoubleConv` module. See `DoubleConv` for more info.\n num_groups (int): number of groups for the GroupNorm\n padding (int or tuple): add zero-padding added to all three sides of the input\n is3d (bool): use 3d or 2d convolutions/pooling operation\n \"\"\"\n\n def __init__(\n self,\n in_channels,\n out_channels,\n conv_kernel_size=3,\n apply_pooling=True,\n pool_kernel_size=2,\n pool_type=\"max\",\n basic_module=DoubleConv,\n conv_layer_order=\"gcr\",\n num_groups=8,\n padding=1,\n is3d=True,\n ):\n super(Encoder, self).__init__()\n assert pool_type in [\"max\", \"avg\"]\n if apply_pooling:\n if pool_type == \"max\":\n if is3d:\n self.pooling = nn.MaxPool3d(kernel_size=pool_kernel_size)\n else:\n self.pooling = nn.MaxPool2d(kernel_size=pool_kernel_size)\n else:\n if is3d:\n self.pooling = nn.AvgPool3d(kernel_size=pool_kernel_size)\n else:\n self.pooling = nn.AvgPool2d(kernel_size=pool_kernel_size)\n else:\n self.pooling = None\n\n self.basic_module = basic_module(\n in_channels,\n out_channels,\n encoder=True,\n kernel_size=conv_kernel_size,\n order=conv_layer_order,\n num_groups=num_groups,\n padding=padding,\n is3d=is3d,\n )\n\n def forward(self, x):\n if self.pooling is not None:\n x = self.pooling(x)\n x = self.basic_module(x)\n return x\n\n\nclass Decoder(nn.Module):\n \"\"\"\n A single module for decoder path consisting of the upsampling layer\n (either learned ConvTranspose3d or nearest neighbor interpolation)\n followed by a basic module (DoubleConv or ResNetBlock).\n\n Args:\n in_channels (int): number of input channels\n out_channels (int): number of output channels\n conv_kernel_size (int or tuple): size of the convolving kernel\n scale_factor (tuple): used as the multiplier for the image H/W/D in\n case of nn.Upsample or as stride in case of ConvTranspose3d, must reverse the MaxPool3d operation\n from the corresponding encoder\n basic_module(nn.Module): either ResNetBlock or DoubleConv\n conv_layer_order (string): determines the order of layers\n in `DoubleConv` module. See `DoubleConv` for more info.\n num_groups (int): number of groups for the GroupNorm\n padding (int or tuple): add zero-padding added to all three sides of the input\n upsample (bool): should the input be upsampled\n \"\"\"\n\n def __init__(\n self,\n in_channels,\n out_channels,\n conv_kernel_size=3,\n scale_factor=(2, 2, 2),\n basic_module=DoubleConv,\n conv_layer_order=\"gcr\",\n num_groups=8,\n mode=\"nearest\",\n padding=1,\n upsample=True,\n is3d=True,\n ):\n super(Decoder, self).__init__()\n\n if upsample:\n if basic_module == DoubleConv:\n # if DoubleConv is the basic_module use interpolation for upsampling and concatenation joining\n self.upsampling = InterpolateUpsampling(mode=mode)\n # concat joining\n self.joining = partial(self._joining, concat=True)\n else:\n # if basic_module=ResNetBlock use transposed convolution upsampling and summation joining\n self.upsampling = TransposeConvUpsampling(\n in_channels=in_channels,\n out_channels=out_channels,\n kernel_size=conv_kernel_size,\n scale_factor=scale_factor,\n )\n # sum joining\n self.joining = partial(self._joining, concat=False)\n # adapt the number of in_channels for the ResNetBlock\n in_channels = out_channels\n else:\n # no upsampling\n self.upsampling = NoUpsampling()\n # concat joining\n self.joining = partial(self._joining, concat=True)\n\n self.basic_module = basic_module(\n in_channels,\n out_channels,\n encoder=False,\n kernel_size=conv_kernel_size,\n order=conv_layer_order,\n num_groups=num_groups,\n padding=padding,\n is3d=is3d,\n )\n\n def forward(self, encoder_features, x):\n x = self.upsampling(encoder_features=encoder_features, x=x)\n x = self.joining(encoder_features, x)\n x = self.basic_module(x)\n return x\n\n @staticmethod\n def _joining(encoder_features, x, concat):\n if concat:\n return torch.cat((encoder_features, x), dim=1)\n else:\n return encoder_features + x\n\n\ndef create_encoders(\n in_channels,\n f_maps,\n basic_module,\n conv_kernel_size,\n conv_padding,\n layer_order,\n num_groups,\n pool_kernel_size,\n is3d,\n):\n # create encoder path consisting of Encoder modules. Depth of the encoder is equal to `len(f_maps)`\n encoders = []\n for i, out_feature_num in enumerate(f_maps):\n if i == 0:\n # apply conv_coord only in the first encoder if any\n encoder = Encoder(\n in_channels,\n out_feature_num,\n apply_pooling=False, # skip pooling in the firs encoder\n basic_module=basic_module,\n conv_layer_order=layer_order,\n conv_kernel_size=conv_kernel_size,\n num_groups=num_groups,\n padding=conv_padding,\n is3d=is3d,\n )\n else:\n encoder = Encoder(\n f_maps[i - 1],\n out_feature_num,\n basic_module=basic_module,\n conv_layer_order=layer_order,\n conv_kernel_size=conv_kernel_size,\n num_groups=num_groups,\n pool_kernel_size=pool_kernel_size,\n padding=conv_padding,\n is3d=is3d,\n )\n\n encoders.append(encoder)\n\n return nn.ModuleList(encoders)\n\n\ndef create_decoders(\n f_maps, basic_module, conv_kernel_size, conv_padding, layer_order, num_groups, is3d\n):\n # create decoder path consisting of the Decoder modules. The length of the decoder list is equal to `len(f_maps) - 1`\n decoders = []\n reversed_f_maps = list(reversed(f_maps))\n for i in range(len(reversed_f_maps) - 1):\n if basic_module == DoubleConv:\n in_feature_num = reversed_f_maps[i] + reversed_f_maps[i + 1]\n else:\n in_feature_num = reversed_f_maps[i]\n\n out_feature_num = reversed_f_maps[i + 1]\n\n decoder = Decoder(\n in_feature_num,\n out_feature_num,\n basic_module=basic_module,\n conv_layer_order=layer_order,\n conv_kernel_size=conv_kernel_size,\n num_groups=num_groups,\n padding=conv_padding,\n is3d=is3d,\n )\n decoders.append(decoder)\n return nn.ModuleList(decoders)\n\n\nclass AbstractUpsampling(nn.Module):\n \"\"\"\n Abstract class for upsampling. A given implementation should upsample a given 5D input tensor using either\n interpolation or learned transposed convolution.\n \"\"\"\n\n def __init__(self, upsample):\n super(AbstractUpsampling, self).__init__()\n self.upsample = upsample\n\n def forward(self, encoder_features, x):\n # get the spatial dimensions of the output given the encoder_features\n output_size = encoder_features.size()[2:]\n # upsample the input and return\n return self.upsample(x, output_size)\n\n\nclass InterpolateUpsampling(AbstractUpsampling):\n \"\"\"\n Args:\n mode (str): algorithm used for upsampling:\n 'nearest' | 'linear' | 'bilinear' | 'trilinear' | 'area'. Default: 'nearest'\n used only if transposed_conv is False\n \"\"\"\n\n def __init__(self, mode=\"nearest\"):\n upsample = partial(self._interpolate, mode=mode)\n super().__init__(upsample)\n\n @staticmethod\n def _interpolate(x, size, mode):\n return F.interpolate(x, size=size, mode=mode)\n\n\nclass TransposeConvUpsampling(AbstractUpsampling):\n \"\"\"\n Args:\n in_channels (int): number of input channels for transposed conv\n used only if transposed_conv is True\n out_channels (int): number of output channels for transpose conv\n used only if transposed_conv is True\n kernel_size (int or tuple): size of the convolving kernel\n used only if transposed_conv is True\n scale_factor (int or tuple): stride of the convolution\n used only if transposed_conv is True\n\n \"\"\"\n\n def __init__(\n self, in_channels=None, out_channels=None, kernel_size=3, scale_factor=(2, 2, 2)\n ):\n # make sure that the output size reverses the MaxPool3d from the corresponding encoder\n upsample = nn.ConvTranspose3d(\n in_channels,\n out_channels,\n kernel_size=kernel_size,\n stride=scale_factor,\n padding=1,\n )\n super().__init__(upsample)\n\n\nclass NoUpsampling(AbstractUpsampling):\n def __init__(self):\n super().__init__(self._no_upsampling)\n\n @staticmethod\n def _no_upsampling(x, size):\n return x\n\n\ndef number_of_features_per_level(init_channel_number, num_levels):\n return [init_channel_number * 2**k for k in range(num_levels)]\n\n\nclass AbstractUNet(BaseModel):\n \"\"\"\n Base class for standard and residual UNet.\n\n Args:\n in_channels (int): number of input channels\n out_channels (int): number of output segmentation masks;\n Note that the of out_channels might correspond to either\n different semantic classes or to different binary segmentation mask.\n It's up to the user of the class to interpret the out_channels and\n use the proper loss criterion during training (i.e. CrossEntropyLoss (multi-class)\n or BCEWithLogitsLoss (two-class) respectively)\n f_maps (int, tuple): number of feature maps at each level of the encoder; if it's an integer the number\n of feature maps is given by the geometric progression: f_maps ^ k, k=1,2,3,4\n final_sigmoid (bool): if True apply element-wise nn.Sigmoid after the final 1x1 convolution,\n otherwise apply nn.Softmax. In effect only if `self.training == False`, i.e. during validation/testing\n basic_module: basic model for the encoder/decoder (DoubleConv, ResNetBlock, ....)\n layer_order (string): determines the order of layers in `SingleConv` module.\n E.g. 'crg' stands for GroupNorm3d+Conv3d+ReLU. See `SingleConv` for more info\n num_groups (int): number of groups for the GroupNorm\n num_levels (int): number of levels in the encoder/decoder path (applied only if f_maps is an int)\n default: 4\n is_segmentation (bool): if True and the model is in eval mode, Sigmoid/Softmax normalization is applied\n after the final convolution; if False (regression problem) the normalization layer is skipped\n conv_kernel_size (int or tuple): size of the convolving kernel in the basic_module\n pool_kernel_size (int or tuple): the size of the window\n conv_padding (int or tuple): add zero-padding added to all three sides of the input\n is3d (bool): if True the model is 3D, otherwise 2D, default: True\n \"\"\"\n\n def __init__(\n self,\n in_channels,\n out_channels,\n final_sigmoid,\n basic_module,\n f_maps=64,\n layer_order=\"gcr\",\n num_groups=8,\n num_levels=4,\n is_segmentation=False,\n conv_kernel_size=3,\n pool_kernel_size=2,\n conv_padding=1,\n is3d=True,\n ):\n super(AbstractUNet, self).__init__()\n\n if isinstance(f_maps, int):\n f_maps = number_of_features_per_level(f_maps, num_levels=num_levels)\n\n assert isinstance(f_maps, list) or isinstance(f_maps, tuple)\n assert len(f_maps) > 1, \"Required at least 2 levels in the U-Net\"\n if \"g\" in layer_order:\n assert (\n num_groups is not None\n ), \"num_groups must be specified if GroupNorm is used\"\n\n # create encoder path\n self.encoders = create_encoders(\n in_channels,\n f_maps,\n basic_module,\n conv_kernel_size,\n conv_padding,\n layer_order,\n num_groups,\n pool_kernel_size,\n is3d,\n )\n\n # create decoder path\n self.decoders = create_decoders(\n f_maps,\n basic_module,\n conv_kernel_size,\n conv_padding,\n layer_order,\n num_groups,\n is3d,\n )\n\n # in the last layer a 1×1 convolution reduces the number of output channels to the number of labels\n if is3d:\n self.final_conv = nn.Conv3d(f_maps[0], out_channels, 1)\n else:\n self.final_conv = nn.Conv2d(f_maps[0], out_channels, 1)\n\n if is_segmentation:\n # semantic segmentation problem\n if final_sigmoid:\n self.final_activation = nn.Sigmoid()\n else:\n self.final_activation = nn.Softmax(dim=1)\n else:\n # regression problem\n self.final_activation = None\n\n def forward(self, x):\n # encoder part\n encoders_features = []\n for encoder in self.encoders:\n x = encoder(x)\n # reverse the encoder outputs to be aligned with the decoder\n encoders_features.insert(0, x)\n\n # remove the last encoder's output from the list\n # !!remember: it's the 1st in the list\n encoders_features = encoders_features[1:]\n\n # decoder part\n for decoder, encoder_features in zip(self.decoders, encoders_features):\n # pass the output from the corresponding encoder and the output\n # of the previous decoder\n x = decoder(encoder_features, x)\n\n x = self.final_conv(x)\n\n # apply final_activation (i.e. Sigmoid or Softmax) only during prediction.\n # During training the network outputs logits\n if not self.training and self.final_activation is not None:\n x = self.final_activation(x)\n\n return x\n\n\nclass UNet3D(AbstractUNet):\n \"\"\"\n 3DUnet model from\n `\"3D U-Net: Learning Dense Volumetric Segmentation from Sparse Annotation\"\n `.\n\n Uses `DoubleConv` as a basic_module and nearest neighbor upsampling in the decoder\n \"\"\"\n\n def __init__(\n self,\n in_channels,\n out_channels,\n final_sigmoid=False,\n f_maps=64,\n layer_order=\"gcr\",\n num_groups=8,\n num_levels=4,\n is_segmentation=False,\n conv_padding=1,\n **kwargs,\n ):\n super(UNet3D, self).__init__(\n in_channels=in_channels,\n out_channels=out_channels,\n final_sigmoid=final_sigmoid,\n basic_module=DoubleConv,\n f_maps=f_maps,\n layer_order=layer_order,\n num_groups=num_groups,\n num_levels=num_levels,\n is_segmentation=is_segmentation,\n conv_padding=conv_padding,\n is3d=True,\n )\n\n" }, { "path": "baselines/utils.py", "content": "import numpy as np\n\nimport torch\nimport torch.autograd as autograd\n\n\ndef weighted_mse(pred, target, weight=None):\n if weight is None:\n return torch.mean((pred - target) ** 2)\n else:\n return torch.mean(weight * (pred - target) ** 2)\n\n\ndef get_3dboundary_points(num_x, # number of points on x axis\n num_y, # number of points on y axis\n num_t, # number of points on t axis\n bot=(0, 0, 0), # lower bound\n top=(1.0, 1.0, 1.0) # upper bound\n ):\n x_top, y_top, t_top = top\n x_bot, y_bot, t_bot = bot\n\n x_arr = np.linspace(x_bot, x_top, num=num_x, endpoint=False)\n y_arr = np.linspace(y_bot, y_top, num=num_y, endpoint=False)\n xx, yy = np.meshgrid(x_arr, y_arr, indexing='ij')\n xarr = np.ravel(xx)\n yarr = np.ravel(yy)\n tarr = np.ones_like(xarr) * t_bot\n point0 = np.stack([xarr, yarr, tarr], axis=0).T # (SxSx1, 3), boundary on t=0\n\n t_arr = np.linspace(t_bot, t_top, num=num_t)\n yy, tt = np.meshgrid(y_arr, t_arr, indexing='ij')\n yarr = np.ravel(yy)\n tarr = np.ravel(tt)\n xarr = np.ones_like(yarr) * x_bot\n point2 = np.stack([xarr, yarr, tarr], axis=0).T # (1xSxT, 3), boundary on x=0\n\n xarr = np.ones_like(yarr) * x_top\n point3 = np.stack([xarr, yarr, tarr], axis=0).T # (1xSxT, 3), boundary on x=2pi\n\n xx, tt = np.meshgrid(x_arr, t_arr, indexing='ij')\n xarr = np.ravel(xx)\n tarr = np.ravel(tt)\n yarr = np.ones_like(xarr) * y_bot\n point4 = np.stack([xarr, yarr, tarr], axis=0).T # (128x1x65, 3), boundary on y=0\n\n yarr = np.ones_like(xarr) * y_top\n point5 = np.stack([xarr, yarr, tarr], axis=0).T # (128x1x65, 3), boundary on y=2pi\n\n points = np.concatenate([point0,\n point2, point3,\n point4, point5],\n axis=0)\n return points\n\n\ndef get_3dboundary(value):\n boundary0 = value[0, :, :, 0:1] # 128x128x1, boundary on t=0\n # boundary1 = value[0, :, :, -1:] # 128x128x1, boundary on t=0.5\n boundary2 = value[0, 0:1, :, :] # 1x128x65, boundary on x=0\n boundary3 = value[0, -1:, :, :] # 1x128x65, boundary on x=1\n boundary4 = value[0, :, 0:1, :] # 128x1x65, boundary on y=0\n boundary5 = value[0, :, -1:, :] # 128x1x65, boundary on y=1\n\n part0 = np.ravel(boundary0)\n # part1 = np.ravel(boundary1)\n part2 = np.ravel(boundary2)\n part3 = np.ravel(boundary3)\n part4 = np.ravel(boundary4)\n part5 = np.ravel(boundary5)\n boundary = np.concatenate([part0,\n part2, part3,\n part4, part5],\n axis=0)[:, np.newaxis]\n return boundary\n\n\ndef get_xytgrid(S, T, bot=[0, 0, 0], top=[1, 1, 1]):\n '''\n Args:\n S: number of points on each spatial domain\n T: number of points on temporal domain including endpoint\n bot: list or tuple, lower bound on each dimension\n top: list or tuple, upper bound on each dimension\n\n Returns:\n (S * S * T, 3) array\n '''\n x_arr = np.linspace(bot[0], top[0], num=S, endpoint=False)\n y_arr = np.linspace(bot[1], top[1], num=S, endpoint=False)\n t_arr = np.linspace(bot[2], top[2], num=T)\n xgrid, ygrid, tgrid = np.meshgrid(x_arr, y_arr, t_arr, indexing='ij')\n xaxis = np.ravel(xgrid)\n yaxis = np.ravel(ygrid)\n taxis = np.ravel(tgrid)\n points = np.stack([xaxis, yaxis, taxis], axis=0).T\n return points\n\n\ndef get_2dgird(num=31):\n x = np.linspace(-1, 1, num)\n y = np.linspace(-1, 1, num)\n gridx, gridy = np.meshgrid(x, y)\n xs = gridx.reshape(-1, 1)\n ys = gridy.reshape(-1, 1)\n result = np.hstack((xs, ys))\n return result\n\n\ndef get_3dgrid(num=11):\n x = np.linspace(-1, 1, num)\n y = np.linspace(-1, 1, num)\n z = np.linspace(-1, 1, num)\n gridx, gridy, gridz = np.meshgrid(x, y, z)\n xs = gridx.reshape(-1, 1)\n ys = gridy.reshape(-1, 1)\n zs = gridz.reshape(-1, 1)\n return np.hstack((xs, ys, zs))\n\n\ndef get_4dgrid(num=11):\n '''\n 4-D meshgrid\n Args:\n num: number of collocation points of each dimension\n\n Returns:\n (num**4, 4) tensor\n '''\n t = np.linspace(0, 1, num)\n x = np.linspace(-1, 1, num)\n y = np.linspace(-1, 1, num)\n z = np.linspace(-1, 1, num)\n gridx, gridy, gridz, gridt = np.meshgrid(x, y, z, t)\n xs = gridx.reshape(-1, 1)\n ys = gridy.reshape(-1, 1)\n zs = gridz.reshape(-1, 1)\n ts = gridt.reshape(-1, 1)\n result = np.hstack((xs, ys, zs, ts))\n return result\n\n\ndef vel2vor(u, v, x, y):\n u_y, = autograd.grad(outputs=[u.sum()], inputs=[y], create_graph=True)\n v_x, = autograd.grad(outputs=[v.sum()], inputs=[x], create_graph=True)\n vorticity = - u_y + v_x\n return vorticity\n\n\ndef sub_mse(vec):\n '''\n Compute mse of two parts of a vector\n Args:\n vec:\n\n Returns:\n\n '''\n length = vec.shape[0] // 2\n diff = (vec[:length] - vec[length: 2 * length]) ** 2\n return diff.mean()\n\n\ndef get_sample(npt=100):\n num = npt // 2\n bc1_y_sample = torch.rand(size=(num, 1)).repeat(2, 1)\n bc1_t_sample = torch.rand(size=(num, 1)).repeat(2, 1)\n\n bc1_x_sample = torch.cat([torch.zeros(num, 1), torch.ones(num, 1)], dim=0)\n\n bc2_x_sample = torch.rand(size=(num, 1)).repeat(2, 1)\n bc2_t_sample = torch.rand(size=(num, 1)).repeat(2, 1)\n\n bc2_y_sample = torch.cat([torch.zeros(num, 1), torch.ones(num, 1)], dim=0)\n return bc1_x_sample, bc1_y_sample, bc1_t_sample, \\\n bc2_x_sample, bc2_y_sample, bc2_t_sample\n\n\ndef concat(xy, z, t=0.0, offset=0):\n '''\n Args:\n xy: (N, 2)\n z: (N, 1）\n t: (N, 1)\n offset: start index of xy\n Returns:\n (N, 4) array\n '''\n output = np.zeros((z.shape[0], 4)) * t\n if offset < 2:\n output[:, offset: offset+2] = xy\n output[:, (offset+2) % 3: (offset+2) % 3 + 1] = z\n else:\n output[:, 2:] = xy[:, 0:1]\n output[:, 0:1] = xy[:, 1:]\n output[:, 1:2] = z\n return output\n\n\ndef cal_mixgrad(outputs, inputs):\n out_grad, = autograd.grad(outputs=[outputs.sum()], inputs=[inputs], create_graph=True)\n out_x2, = autograd.grad(outputs=[out_grad[:, 0].sum()], inputs=[inputs], create_graph=True)\n out_xx = out_x2[:, 0]\n out_y2, = autograd.grad(outputs=[out_grad[:, 1].sum()], inputs=[inputs], create_graph=True)\n out_yy = out_y2[:, 1]\n out_z2, = autograd.grad(outputs=[out_grad[:, 2].sum()], inputs=[inputs], create_graph=True)\n out_zz = out_z2[:, 2]\n return out_grad, out_xx, out_yy, out_zz" }, { "path": "cavity_flow.py", "content": "\"\"\"\n@author: Zongyi Li\nThis file is the Fourier Neural Operator for 3D problem such as the Navier-Stokes equation discussed in Section 5.3 in the [paper](https://arxiv.org/pdf/2010.08895.pdf),\nwhich takes the 2D spatial + 1D temporal equation directly as a 3D problem\n\"\"\"\n\nimport torch\nimport numpy as np\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nimport matplotlib.pyplot as plt\n\n\nfrom timeit import default_timer\nfrom torch.optim import Adam\nfrom train_utils.datasets import MatReader\nfrom train_utils.losses import LpLoss\nfrom train_utils.utils import count_params\n\ntorch.manual_seed(0)\nnp.random.seed(0)\n\n\n################################################################\n# 3d fourier layers\n################################################################\n\nclass SpectralConv3d(nn.Module):\n def __init__(self, in_channels, out_channels, modes1, modes2, modes3):\n super(SpectralConv3d, self).__init__()\n\n \"\"\"\n 3D Fourier layer. It does FFT, linear transform, and Inverse FFT. \n \"\"\"\n\n self.in_channels = in_channels\n self.out_channels = out_channels\n self.modes1 = modes1 # Number of Fourier modes to multiply, at most floor(N/2) + 1\n self.modes2 = modes2\n self.modes3 = modes3\n\n self.scale = (1 / (in_channels * out_channels))\n self.weights1 = nn.Parameter(\n self.scale * torch.rand(in_channels, out_channels, self.modes1, self.modes2, self.modes3,\n dtype=torch.cfloat))\n self.weights2 = nn.Parameter(\n self.scale * torch.rand(in_channels, out_channels, self.modes1, self.modes2, self.modes3,\n dtype=torch.cfloat))\n self.weights3 = nn.Parameter(\n self.scale * torch.rand(in_channels, out_channels, self.modes1, self.modes2, self.modes3,\n dtype=torch.cfloat))\n self.weights4 = nn.Parameter(\n self.scale * torch.rand(in_channels, out_channels, self.modes1, self.modes2, self.modes3,\n dtype=torch.cfloat))\n\n # Complex multiplication\n def compl_mul3d(self, input, weights):\n # (batch, in_channel, x,y,t ), (in_channel, out_channel, x,y,t) -> (batch, out_channel, x,y,t)\n return torch.einsum(\"bixyz,ioxyz->boxyz\", input, weights)\n\n def forward(self, x):\n batchsize = x.shape[0]\n # Compute Fourier coeffcients up to factor of e^(- something constant)\n x_ft = torch.fft.rfftn(x, dim=[-3, -2, -1])\n\n # Multiply relevant Fourier modes\n out_ft = torch.zeros(batchsize, self.out_channels, x.size(-3), x.size(-2), x.size(-1) // 2 + 1,\n dtype=torch.cfloat, device=x.device)\n out_ft[:, :, :self.modes1, :self.modes2, :self.modes3] = \\\n self.compl_mul3d(x_ft[:, :, :self.modes1, :self.modes2, :self.modes3], self.weights1)\n out_ft[:, :, -self.modes1:, :self.modes2, :self.modes3] = \\\n self.compl_mul3d(x_ft[:, :, -self.modes1:, :self.modes2, :self.modes3], self.weights2)\n out_ft[:, :, :self.modes1, -self.modes2:, :self.modes3] = \\\n self.compl_mul3d(x_ft[:, :, :self.modes1, -self.modes2:, :self.modes3], self.weights3)\n out_ft[:, :, -self.modes1:, -self.modes2:, :self.modes3] = \\\n self.compl_mul3d(x_ft[:, :, -self.modes1:, -self.modes2:, :self.modes3], self.weights4)\n\n # Return to physical space\n x = torch.fft.irfftn(out_ft, s=(x.size(-3), x.size(-2), x.size(-1)))\n return x\n\n\nclass FNO3d(nn.Module):\n def __init__(self, modes1, modes2, modes3, width, padding):\n super(FNO3d, self).__init__()\n\n \"\"\"\n The overall network. It contains 4 layers of the Fourier layer.\n 1. Lift the input to the desire channel dimension by self.fc0 .\n 2. 4 layers of the integral operators u' = (W + K)(u).\n W defined by self.w; K defined by self.conv .\n 3. Project from the channel space to the output space by self.fc1 and self.fc2 .\n\n input: the solution of the first 10 timesteps + 3 locations (u(1, x, y), ..., u(10, x, y), x, y, t). It's a constant function in time, except for the last index.\n input shape: (batchsize, x=64, y=64, t=40, c=13)\n output: the solution of the next 40 timesteps\n output shape: (batchsize, x=64, y=64, t=40, c=1)\n \"\"\"\n\n self.modes1 = modes1\n self.modes2 = modes2\n self.modes3 = modes3\n self.width = width\n self.padding = padding # pad the domain if input is non-periodic\n self.fc0 = nn.Linear(5, 32)\n self.fc1 = nn.Linear(32, self.width)\n # input channel is 12: the solution of the first 10 timesteps + 3 locations (u(1, x, y), ..., u(10, x, y), x, y, t)\n\n self.conv0 = SpectralConv3d(self.width, self.width, self.modes1, self.modes2, self.modes3)\n self.conv1 = SpectralConv3d(self.width, self.width, self.modes1, self.modes2, self.modes3)\n self.conv2 = SpectralConv3d(self.width, self.width, self.modes1, self.modes2, self.modes3)\n self.conv3 = SpectralConv3d(self.width, self.width, self.modes1, self.modes2, self.modes3)\n self.w0 = nn.Conv3d(self.width, self.width, 1)\n self.w1 = nn.Conv3d(self.width, self.width, 1)\n self.w2 = nn.Conv3d(self.width, self.width, 1)\n self.w3 = nn.Conv3d(self.width, self.width, 1)\n\n self.fc2 = nn.Linear(self.width, 128)\n self.fc3 = nn.Linear(128, 3)\n\n def forward(self, x):\n grid = self.get_grid(x.shape, x.device)\n x = torch.cat((x, grid), dim=-1)\n x = self.fc0(x)\n x = F.tanh(x)\n x = self.fc1(x)\n x = x.permute(0, 4, 1, 2, 3)\n x = F.pad(x, [0, self.padding, 0, self.padding, 0, self.padding]) # pad the domain if input is non-periodic\n\n x1 = self.conv0(x)\n x2 = self.w0(x)\n x = x1 + x2\n x = F.tanh(x)\n\n x1 = self.conv1(x)\n x2 = self.w1(x)\n x = x1 + x2\n x = F.tanh(x)\n\n x1 = self.conv2(x)\n x2 = self.w2(x)\n x = x1 + x2\n x = F.tanh(x)\n\n x1 = self.conv3(x)\n x2 = self.w3(x)\n x = x1 + x2\n\n # x = x[:, :, :-self.padding, :-self.padding, :-self.padding]\n x = x.permute(0, 2, 3, 4, 1) # pad the domain if input is non-periodic\n x = self.fc2(x)\n x = F.tanh(x)\n x = self.fc3(x)\n return x\n\n def get_grid(self, shape, device):\n batchsize, size_x, size_y, size_z = shape[0], shape[1], shape[2], shape[3]\n gridx = torch.tensor(np.linspace(0, 1, size_x), dtype=torch.float)\n gridx = gridx.reshape(1, size_x, 1, 1, 1).repeat([batchsize, 1, size_y, size_z, 1])\n gridy = torch.tensor(np.linspace(0, 1, size_y), dtype=torch.float)\n gridy = gridy.reshape(1, 1, size_y, 1, 1).repeat([batchsize, size_x, 1, size_z, 1])\n gridz = torch.tensor(np.linspace(0, 1, size_z), dtype=torch.float)\n gridz = gridz.reshape(1, 1, 1, size_z, 1).repeat([batchsize, size_x, size_y, 1, 1])\n return torch.cat((gridx, gridy, gridz), dim=-1).to(device)\n\n\n################################################################\n# configs\n################################################################\n\n\n\n# PATH = '../data/cavity.mat'\nPATH = '../data/lid-cavity.mat'\nntest = 1\n\nmodes = 8\nwidth = 32\n\nbatch_size = 1\n\npath = 'cavity'\npath_model = 'model/' + path\npath_train_err = 'results/' + path + 'train.txt'\npath_test_err = 'results/' + path + 'test.txt'\npath_image = 'image/' + path\n\n\n\nsub_s = 4\nsub_t = 20\nS = 256 // sub_s\nT_in = 1000 # 1000*0.005 = 5s\nT = 50 # 1000 + 50*20*0.005 = 10s\npadding = 14\n\n################################################################\n# load data\n################################################################\n\n# 15s, 3000 frames\nreader = MatReader(PATH)\ndata_u = reader.read_field('u')[T_in:T_in+T*sub_t:sub_t, ::sub_s, ::sub_s].permute(1,2,0)\ndata_v = reader.read_field('v')[T_in:T_in+T*sub_t:sub_t, ::sub_s, ::sub_s].permute(1,2,0)\n\ndata_output = torch.stack([data_u, data_v],dim=-1).reshape(batch_size,S,S,T,2)\ndata_input = data_output[:,:,:,:1,:].repeat(1,1,1,T,1).reshape(batch_size,S,S,T,2)\n\nprint(data_output.shape)\n\n\ndevice = torch.device('cuda')\n\ndef PINO_loss_Fourier_f(out, Re=500):\n pi = np.pi\n Lx = 1*(S + padding-1)/S\n Ly = 1*(S + padding-1)/S\n Lt = (0.005*sub_t*T) *(T + padding)/T\n\n nx = out.size(1)\n ny = out.size(2)\n nt = out.size(3)\n device = out.device\n\n # Wavenumbers in y-direction\n k_x = torch.cat((torch.arange(start=0, end=nx//2, step=1, device=device),\n torch.arange(start=-nx//2, end=0, step=1, device=device)), 0).reshape(nx, 1, 1).repeat(1, ny, nt).reshape(1,nx,ny,nt,1)\n k_y = torch.cat((torch.arange(start=0, end=ny//2, step=1, device=device),\n torch.arange(start=-ny//2, end=0, step=1, device=device)), 0).reshape(1, ny, 1).repeat(nx, 1, nt).reshape(1,nx,ny,nt,1)\n k_t = torch.cat((torch.arange(start=0, end=nt//2, step=1, device=device),\n torch.arange(start=-nt//2, end=0, step=1, device=device)), 0).reshape(1, 1, nt).repeat(nx, ny, 1).reshape(1,nx,ny,nt,1)\n\n out_h = torch.fft.fftn(out, dim=[1, 2, 3])\n outx_h = 1j * k_x * out_h * (2 * pi / Lx)\n outy_h = 1j * k_y * out_h * (2 * pi / Ly)\n outt_h = 1j * k_t * out_h * (2 * pi / Lt)\n outxx_h = 1j * k_x * outx_h * (2 * pi / Lx)\n outyy_h = 1j * k_y * outy_h * (2 * pi / Ly)\n\n outx = torch.fft.irfftn(outx_h[:, :, :, :nt//2+1, :], dim=[1,2,3])[:,:S,:S,:T]\n outy = torch.fft.irfftn(outy_h[:, :, :, :nt//2+1, :], dim=[1,2,3])[:,:S,:S,:T]\n outt = torch.fft.irfftn(outt_h[:, :, :, :nt//2+1, :], dim=[1,2,3])[:,:S,:S,:T]\n outxx = torch.fft.irfftn(outxx_h[:, :, :, :nt//2+1, :], dim=[1,2,3])[:,:S,:S,:T]\n outyy = torch.fft.irfftn(outyy_h[:, :, :, :nt//2+1, :], dim=[1,2,3])[:,:S,:S,:T]\n out = out[:,:S,:S,:T]\n\n\n E1 = outt[..., 0] + out[..., 0]*outx[..., 0] + out[..., 1]*outy[..., 0] + outx[..., 2] - 1/Re*(outxx[..., 0] + outyy[..., 0])\n E2 = outt[..., 1] + out[..., 0]*outx[..., 1] + out[..., 1]*outy[..., 1] + outy[..., 2] - 1/Re*(outxx[..., 1] + outyy[..., 1])\n E3 = outx[..., 0] + outy[..., 1]\n\n target = torch.zeros(E1.shape, device=E1.device)\n E1 = F.mse_loss(E1,target)\n E2 = F.mse_loss(E2,target)\n E3 = F.mse_loss(E3,target)\n\n return E1, E2, E3\n\ndef PINO_loss_FDM_f(out, Re=500):\n dx = 1 / (S+2)\n dy = 1 / (S+2)\n dt = 0.005*sub_t\n\n out = out[:,:S,:S,:T,:]\n out = F.pad(out, [0,0, 1,0, 1,1, 1,1])\n out[:, :, -1, :, 0] = 1\n\n outx = (out[:,2:,1:-1,1:-1] - out[:,:-2,1:-1,1:-1]) / (2*dx)\n outy = (out[:,1:-1,2:,1:-1] - out[:,1:-1,:-2,1:-1]) / (2*dy)\n outt = (out[:,1:-1,1:-1,2:] - out[:,1:-1,1:-1,:-2]) / (2*dt)\n outlap = (out[:,2:,1:-1,1:-1] + out[:,:-2,1:-1,1:-1] + out[:,1:-1,2:,1:-1] + out[:,1:-1,:-2,1:-1] - 4*out[:,1:-1,1:-1,1:-1]) / (dx*dy)\n\n out = out[:,1:-1,1:-1,1:-1]\n\n E1 = outt[..., 0] + out[..., 0]*outx[..., 0] + out[..., 1]*outy[..., 0] + outx[..., 2] - 1/Re*(outlap[..., 0])\n E2 = outt[..., 1] + out[..., 0]*outx[..., 1] + out[..., 1]*outy[..., 1] + outy[..., 2] - 1/Re*(outlap[..., 1])\n E3 = outx[..., 0] + outy[..., 1]\n\n target = torch.zeros(E1.shape, device=E1.device)\n E1 = F.mse_loss(E1,target)\n E2 = F.mse_loss(E2,target)\n E3 = F.mse_loss(E3,target)\n\n return E1, E2, E3\n\n\n\ndef PINO_loss_ic(out, y):\n myloss = LpLoss(size_average=True)\n # target = torch.zeros(out.shape, device=out.device)\n # target[:, :, -1, 0] = 1\n # IC = myloss(out, target)\n # return IC\n\n IC = F.mse_loss(out, y)\n return IC\n\ndef PINO_loss_bc(out, y):\n myloss = LpLoss(size_average=True)\n # target = torch.zeros((batch_size,S,T,2), device=out.device)\n # target3 = torch.zeros((batch_size,S,T,2), device=out.device)\n # target3[..., 0] = 1\n # out = torch.stack([out[:,0,:], out[:,-1,:], out[:,:,-1], out[:,:,0]], -1)\n # target = torch.stack([target, target, target3, target], -1)\n # BC = myloss(out, target)\n # return BC\n\n BC1 = F.mse_loss(out[:,0,:], y[:,0,:])\n BC2 = F.mse_loss(out[:,-1,:], y[:,-1,:])\n BC3 = F.mse_loss(out[:,:,-1], y[:,:,-1])\n BC4 = F.mse_loss(out[:,:,0], y[:,:,0])\n return (BC1+BC2+BC3+BC4)/4\n\n################################################################\n# training and evaluation\n################################################################\n\n\n\nmodel = model = FNO3d(modes, modes, modes, width, padding).cuda()\nprint(count_params(model))\n\noptimizer = Adam(model.parameters(), lr=0.0025, weight_decay=0)\nscheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=500, gamma=0.5)\n\nmyloss = LpLoss(size_average=False)\nmodel.train()\nx = data_input.cuda().reshape(batch_size,S,S,T,2)\ny = data_output.cuda().reshape(batch_size,S,S,T,2)\n\nfor ep in range(5000):\n t1 = default_timer()\n optimizer.zero_grad()\n\n out = model(x)\n\n loss_l2 = myloss(out[:,:S,:S,:T,:2], y)\n IC = PINO_loss_ic(out[:,:S,:S,0,:2], y[:,:,:,0])\n BC = PINO_loss_bc(out[:,:S,:S,:T,:2], y)\n E1, E2, E3 = PINO_loss_Fourier_f(out)\n # E1, E2, E3 = PINO_loss_FDM_f(out)\n loss_pino = IC*1 + BC*1 + E1*1 + E2*1 + E3*1\n\n loss_pino.backward()\n\n optimizer.step()\n scheduler.step()\n t2 = default_timer()\n print(ep, t2-t1, IC.item(), BC.item(), E1.item(), E2.item(), E3.item(), loss_l2.item())\n\n if ep % 1000 == 500:\n y_plot = y[0,:,:,:].cpu().numpy()\n out_plot = out[0,:S,:S,:T].detach().cpu().numpy()\n\n fig, ax = plt.subplots(2, 2)\n ax[0,0].imshow(y_plot[..., -1, 0])\n ax[0,1].imshow(y_plot[..., -1, 1])\n ax[1,0].imshow(out_plot[..., -1, 0])\n ax[1,1].imshow(out_plot[..., -1, 1])\n plt.show()" }, { "path": "configs/baseline/NS-50s-LAAF.yaml", "content": "data:\n datapath: 'data/ns_V1e-3_N5000_T50.mat'\n vis: 0.001\n total_num: 5000\n offset: 4900\n n_sample: 1\n time_scale: 49\n nx: 64\n nt: 50\n sub: 1 # not used here\n sub_t: 1 # not used here\n shuffle: True\n\nmodel:\n layers: [3, 50, 50, 50, 50, 50, 50, 3]\n activation: LAAF-10 tanh\n\ntrain:\n batchsize: 1\n epochs: 5000\n milestones: [1000, 1500, 2000]\n base_lr: 0.001\n num_domain: 10000\n num_boundary: 18000\n num_test: 100\n log_step: 100\n\nlog:\n logfile: 'log/pinns-50s-laaf.csv'" }, { "path": "configs/baseline/NS-50s.yaml", "content": "data:\n datapath: 'data/ns_V1e-3_N5000_T50.mat'\n vis: 0.001\n total_num: 5000\n offset: 4900\n n_sample: 1\n time_scale: 49\n nx: 64\n nt: 50\n sub: 1 # not used here\n sub_t: 1 # not used here\n shuffle: True\n\nmodel:\n layers: [3, 50, 50, 50, 50, 50, 50, 3]\n\ntrain:\n epochs: 15000\n base_lr: 0.001\n save_dir: 'Re500-FDM'\n save_name: 'PINO-scratch-05s.pt'\n\nlog:\n logfile: 'log/pinns-50s-best.csv'" }, { "path": "configs/baseline/Re500-05s-deeponet.yaml", "content": "data:\n datapath: '/mnt/md1/zongyi/NS_fft_Re500_T4000.npy'\n Re: 500\n total_num: 4000\n offset: 0\n n_sample: 4000\n time_interval: 0.5\n nx: 64\n nt: 64\n sub: 1\n sub_t: 1\n shuffle: False\n data_val: 'data/NS_Re500_s256_T100_test.npy'\n val_nx: 256\n val_nt: 128\n val_sub: 4\n val_subt: 2\n\nmodel:\n layers: [40, 40]\n activation: 'relu'\n\ntrain:\n batchsize: 1\n epochs: 100\n milestones: [25000, 50000, 75000]\n base_lr: 0.001\n\n ckpt: 'checkpoints/Re500-FDM/pretrain-Re500-05s-4000.pt'\n\nlog:\n project: 'PINO-None'\n group: 'eval'\n\n\n" }, { "path": "configs/baseline/Re500-pinns-05s-LAAF.yaml", "content": "data:\n datapath: 'data/NS_Re500_s256_T100_test.npy'\n Re: 500\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 0.5\n nx: 256\n nt: 128\n sub: 2\n sub_t: 1\n shuffle: True\n\nmodel:\n layers: [3, 50, 50, 50, 50, 50, 50, 3]\n activation: LAAF-10 tanh\n\ntrain:\n batchsize: 1\n epochs: 3000\n milestones: [1000, 1500, 2000]\n base_lr: 0.01\n num_domain: 5000\n num_boundary: 10000\n num_test: 100\n log_step: 100\n\nlog:\n project: 'PINO-Re500-ICLR'\n group: 'Re500-PINNs'\n logfile: 'log/pinns-plot-laaf128.csv'\n\n\n\n" }, { "path": "configs/baseline/Re500-pinns-05s-SA.yaml", "content": "data:\n datapath: 'data/NS_Re500_s256_T100_test.npy'\n Re: 500\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 0.5\n nx: 256\n nt: 128\n sub: 2\n sub_t: 1\n shuffle: True\n\nmodel:\n layers: [3, 100, 100, 100, 100, 100, 100, 3]\n activation: tanh\n\ntrain:\n batchsize: 1\n epochs: 5000\n milestones: [1000, 1500, 2000]\n base_lr: 0.005\n num_domain: 10000\n num_boundary: 10000\n num_init: 5000\n num_test: 100\n log_step: 100\n\nlog:\n project: 'PINO-Re500-ICLR'\n group: 'SA-PINNs'\n logfile: 'log/sa-pinns128-plot.csv'\n\n\n\n" }, { "path": "configs/baseline/Re500-pinns-05s.yaml", "content": "data:\n datapath: 'data/NS_Re500_s256_T100_test.npy'\n Re: 500\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 0.5\n nx: 256\n nt: 128\n sub: 2\n sub_t: 1\n shuffle: True\n\nmodel:\n layers: [3, 50, 50, 50, 50, 50, 50, 3]\n activation: tanh\n\ntrain:\n batchsize: 1\n epochs: 3000\n milestones: [1000, 1500, 2000]\n base_lr: 0.01\n num_domain: 5000\n num_boundary: 10000\n num_test: 100\n log_step: 100\n\nlog:\n project: 'PINO-Re500-ICLR'\n group: 'Re500-PINNs'\n logfile: 'log/pinns128-plot.csv'\n\n\n\n" }, { "path": "configs/baseline/Re500-pinns.yaml", "content": "data:\n datapath: 'data/NS_Re500_s256_T100_test.npy'\n Re: 500\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 0.5\n nx: 256\n nt: 128\n sub: 4\n sub_t: 1\n shuffle: True\n\ntrain:\n batchsize: 1\n epochs: 5000\n base_lr: 0.001\n num_domain: 5000\n num_boundary: 10000\n num_test: 100\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0\n log_step: 100\n\nlog:\n project: 'PINO-Re500-ICLR'\n group: 'Re500-PINNs-long'\n\n\n" }, { "path": "configs/finetune/Darcy-finetune.yaml", "content": "data:\n name: 'Darcy'\n datapath: '/mnt/md1/zongyi/piececonst_r421_N1024_smooth2.mat'\n total_num: 1024\n offset: 500\n n_sample: 1\n nx: 421\n sub: 7\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [20, 20, 20, 20]\n modes2: [20, 20, 20, 20]\n fc_dim: 128\n act: gelu\n\ntrain:\n batchsize: 1\n epochs: 500\n milestones: [100, 200, 300, 400]\n base_lr: 0.0025\n scheduler_gamma: 0.5\n f_loss: 1.0\n xy_loss: 0.0\n save_dir: 'darcy-FDM'\n save_name: 'darcy-finetune-pino.pt'\n ckpt: 'checkpoints/darcy-FDM/darcy-pretrain-pino.pt'\n\nlog:\n project: 'ICLR-Darcy-finetune'\n group: 'gelu-pino-pino'\n\n\n" }, { "path": "configs/finetune/Re100-finetune-1s.yaml", "content": "data:\n datapath: 'data/NS_fine_Re100_T128_part0.npy'\n Re: 100\n total_num: 100\n offset: 190\n n_sample: 1\n time_interval: 1.0\n nx: 128\n nt: 128\n sub: 1\n sub_t: 2\n shuffle: True\n\nmodel:\n layers: [32, 32, 32, 32, 32]\n modes1: [16, 16, 16, 16]\n modes2: [16, 16, 16, 16]\n modes3: [16, 16, 16, 16]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 7500\n milestones: [500, 1500, 3000, 4000, 5000]\n base_lr: 0.0025\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0.0\n save_dir: 'Re100-FDM'\n save_name: 'PINO-finetune-Re100-1s.pt'\n ckpt: 'checkpoints/Re100-FDM/PINO-pretrain-Re100-1s.pt'\n\nlog:\n project: 'PINO-finetune'\n group: 'Re100-finetune-1s'\n\n\n" }, { "path": "configs/finetune/Re200-finetune-1s.yaml", "content": "data:\n datapath: 'data/NS_fine_Re200_T128_part0.npy'\n Re: 200\n total_num: 100\n offset: 194\n n_sample: 1\n time_interval: 1.0\n nx: 128\n nt: 128\n sub: 1\n sub_t: 2\n shuffle: True\n\nmodel:\n layers: [32, 32, 32, 32, 32]\n modes1: [16, 16, 16, 16]\n modes2: [16, 16, 16, 16]\n modes3: [16, 16, 16, 16]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 10000\n milestones: [500, 1500, 3000, 4000, 6000]\n base_lr: 0.0025\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0.0\n save_dir: 'Re200-FDM'\n save_name: 'PINO-finetune-Re200-1s.pt'\n ckpt: 'checkpoints/Re200-FDM/PINO-pretrain-Re200-1s.pt'\n\nlog:\n project: 'PINO-finetune'\n group: 'Re200-finetune-1s'\n\n\n" }, { "path": "configs/finetune/Re250-finetune-1s.yaml", "content": "data:\n datapath: 'data/NS_fine_Re250_T128_part0.npy'\n Re: 250\n total_num: 100\n offset: 198\n n_sample: 1\n time_interval: 1.0\n nx: 128\n nt: 128\n sub: 1\n sub_t: 2\n shuffle: True\n\nmodel:\n layers: [32, 32, 32, 32, 32]\n modes1: [16, 16, 16, 16]\n modes2: [16, 16, 16, 16]\n modes3: [16, 16, 16, 16]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 10000\n milestones: [500, 1500, 3000, 4000, 6000]\n base_lr: 0.0025\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0.0\n save_dir: 'Re250-FDM'\n save_name: 'PINO-finetune-Re250-1s.pt'\n ckpt: 'checkpoints/Re250-FDM/PINO-pretrain-Re250-1s.pt'\n\nlog:\n project: 'PINO-finetune'\n group: 'Re250-finetune-1s'\n\n\n" }, { "path": "configs/finetune/Re300-finetune-1s.yaml", "content": "data:\n datapath: 'data/NS_fine_Re300_T128_part0.npy'\n Re: 300\n total_num: 100\n offset: 190\n n_sample: 1\n time_interval: 1.0\n nx: 128\n nt: 128\n sub: 1\n sub_t: 2\n shuffle: True\n\nmodel:\n layers: [32, 32, 32, 32, 32]\n modes1: [16, 16, 16, 16]\n modes2: [16, 16, 16, 16]\n modes3: [16, 16, 16, 16]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 10000\n milestones: [500, 1500, 3000, 4000, 6000]\n base_lr: 0.0025\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0.0\n save_dir: 'Re300-FDM'\n save_name: 'PINO-finetine-Re300-1s.pt'\n ckpt: 'checkpoints/Re300-FDM/PINO-pretrain-Re300-1s.pt'\n\nlog:\n project: 'PINO-finetune'\n group: 'Re300-finetune-1s'\n\n\n" }, { "path": "configs/finetune/Re350-finetune-1s.yaml", "content": "data:\n datapath: 'data/NS_fine_Re350_T128_part0.npy'\n Re: 350\n total_num: 100\n offset: 198\n n_sample: 1\n time_interval: 1.0\n nx: 128\n nt: 128\n sub: 1\n sub_t: 2\n shuffle: True\n\nmodel:\n layers: [32, 32, 32, 32, 32]\n modes1: [16, 16, 16, 16]\n modes2: [16, 16, 16, 16]\n modes3: [16, 16, 16, 16]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 10000\n milestones: [500, 1500, 3000, 4000, 6000]\n base_lr: 0.0025\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0.0\n save_dir: 'Re350-FDM'\n save_name: 'PINO-finetine-Re300-1s.pt'\n ckpt: 'checkpoints/Re350-FDM/PINO-pretrain-Re350-1s.pt'\n\nlog:\n project: 'PINO-finetune'\n group: 'Re350-finetune-1s'\n\n\n" }, { "path": "configs/finetune/Re400-finetune-1s.yaml", "content": "data:\n datapath: 'data/NS_fine_Re400_T128_part0.npy'\n Re: 400\n total_num: 100\n offset: 199\n n_sample: 1\n time_interval: 1.0\n nx: 128\n nt: 128\n sub: 1\n sub_t: 2\n shuffle: True\n\nmodel:\n layers: [32, 32, 32, 32, 32]\n modes1: [16, 16, 16, 16]\n modes2: [16, 16, 16, 16]\n modes3: [16, 16, 16, 16]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 10000\n milestones: [500, 1500, 3000, 4000, 6000]\n base_lr: 0.0025\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0.0\n save_dir: 'Re400-FDM'\n save_name: 'PINO-finetune-Re400-1s.pt'\n ckpt: 'checkpoints/Re400-FDM/PINO-pretrain-Re400-1s.pt'\n\nlog:\n project: 'PINO-finetune'\n group: 'Re400-finetune-1s'\n\n\n" }, { "path": "configs/finetune/Re500-finetune-05s-2layer.yaml", "content": "data:\n datapath: 'data/NS_Re500_s256_T100_test.npy'\n Re: 500\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 0.5\n nx: 256\n nt: 128\n sub: 2\n sub_t: 1\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 6000\n milestones: [1000, 3000, 5000]\n base_lr: 0.0025\n beta1: 0.9\n beta2: 0.999\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0\n save_dir: 'Re500'\n save_name: 'PINO-fintune-05s.pt'\n ckpt: 'checkpoints/Re500-FDM/PINO-pretrain-Re500-05s-4k1k.pt'\n twolayer: True\n\nlog:\n project: 'PINO-Re500-ICLR'\n group: 'Re500-finetune-128-4k1-2layer'\n\n\n" }, { "path": "configs/finetune/Re500-finetune-05s-eqn.yaml", "content": "data:\n datapath: 'data/NS_Re500_s256_T100_test.npy'\n Re: 500\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 0.5\n nx: 256\n nt: 128\n sub: 2\n sub_t: 1\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 2500\n milestones: [1000, 1500, 2000]\n base_lr: 0.0025\n beta1: 0.9\n beta2: 0.999\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0\n save_dir: 'Re500'\n save_name: 'PINO-fintune-05s.pt'\n ckpt: 'checkpoints/Re500-FDM/PINO-pretrain-Re500-05s-eqn.pt'\n\nlog:\n project: 'PINO-Re500-ICLR'\n group: 'Re500-finetune-128-eqn'\n\n\n" }, { "path": "configs/finetune/Re500-finetune-05s4C0.yaml", "content": "data:\n datapath: 'data/NS_Re500_s256_T100_test.npy'\n Re: 500\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 0.5\n nx: 256\n nt: 128\n sub: 2\n sub_t: 1\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 2500\n milestones: [1000, 1500, 2000]\n base_lr: 0.0025\n beta1: 0.9\n beta2: 0.999\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0\n save_dir: 'Re500'\n save_name: 'PINO-fintune-05s.pt'\n ckpt: 'checkpoints/Re500-FDM/PINO-pretrain-Re500-05s-4C0.pt'\n\nlog:\n project: 'PINO-Re500-ICLR'\n group: 'Re500-finetune-128-4C0'\n\n\n" }, { "path": "configs/finetune/Re500-finetune-05s4C1.yaml", "content": "data:\n datapath: 'data/NS_Re500_s256_T100_test.npy'\n Re: 500\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 0.5\n nx: 256\n nt: 128\n sub: 4\n sub_t: 1\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 2500\n milestones: [1000, 1500, 2000]\n base_lr: 0.0025\n beta1: 0.9\n beta2: 0.999\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0\n save_dir: 'Re500'\n save_name: 'PINO-fintune-05s.pt'\n ckpt: 'checkpoints/Re500-FDM/PINO-pretrain-Re500-05s-4C1.pt'\n profile: True\n\nlog:\n project: 'PINO-Re500-ICLR'\n group: 'Re500-finetune-128-4C1-profile-long'\n\n\n" }, { "path": "configs/finetune/Re500-finetune-05s4C4.yaml", "content": "data:\n datapath: 'data/NS_Re500_s256_T100_test.npy'\n Re: 500\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 0.5\n nx: 256\n nt: 128\n sub: 2\n sub_t: 1\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 2500\n milestones: [1000, 1500, 2000]\n base_lr: 0.0025\n beta1: 0.9\n beta2: 0.999\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0\n save_dir: 'Re500'\n save_name: 'PINO-fintune-05s.pt'\n ckpt: 'checkpoints/Re500-FDM/PINO-pretrain-Re500-05s-4C4.pt'\n\nlog:\n project: 'PINO-Re500-ICLR'\n group: 'Re500-finetune-128-4C4'\n\n\n" }, { "path": "configs/finetune/Re500-finetune-05s4k-2layer.yaml", "content": "data:\n datapath: 'data/NS_Re500_s256_T100_test.npy'\n Re: 500\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 0.5\n nx: 256\n nt: 128\n sub: 2\n sub_t: 1\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 6000\n milestones: [1000, 3000, 5000]\n base_lr: 0.0025\n beta1: 0.9\n beta2: 0.999\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0\n save_dir: 'Re500'\n save_name: 'PINO-fintune-05s.pt'\n ckpt: 'checkpoints/Re500-FDM/pretrain-Re500-05s-4000.pt'\n twolayer: True\n\nlog:\n project: 'PINO-Re500-ICLR'\n group: 'Re500-finetune-128-4k-2layer'\n\n\n" }, { "path": "configs/finetune/Re500-finetune-05s4k1k.yaml", "content": "data:\n datapath: 'data/NS_Re500_s256_T100_test.npy'\n Re: 500\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 0.5\n nx: 256\n nt: 128\n sub: 2\n sub_t: 1\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 2500\n milestones: [1000, 1500, 2000]\n base_lr: 0.0025\n beta1: 0.9\n beta2: 0.999\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0\n save_dir: 'Re500'\n save_name: 'PINO-fintune-05s.pt'\n ckpt: 'checkpoints/Re500-FDM/PINO-pretrain-Re500-05s-4k1k.pt'\n profile: True\n\nlog:\n project: 'PINO-Re500-ICLR-rebuttal'\n group: 'Re500-finetune-128-4k1'\n\n\n" }, { "path": "configs/finetune/Re500-finetune-05s4k4-2layer.yaml", "content": "data:\n datapath: 'data/NS_Re500_s256_T100_test.npy'\n Re: 500\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 0.5\n nx: 256\n nt: 128\n sub: 2\n sub_t: 1\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 6000\n milestones: [1000, 3000, 5000]\n base_lr: 0.0025\n beta1: 0.9\n beta2: 0.999\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0\n save_dir: 'Re500'\n save_name: 'PINO-fintune-05s.pt'\n ckpt: 'checkpoints/Re500-FDM/PINO-pretrain-Re500-05s-4k.pt'\n twolayer: True\n\nlog:\n project: 'PINO-Re500-ICLR'\n group: 'Re500-finetune-128-4k4-2layer'\n\n\n" }, { "path": "configs/finetune/Re500-finetune-05s4k4k.yaml", "content": "data:\n datapath: 'data/NS_Re500_s256_T100_test.npy'\n Re: 500\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 0.5\n nx: 256\n nt: 128\n sub: 4\n sub_t: 1\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 2500\n milestones: [1000, 1500, 2000]\n base_lr: 0.0025\n beta1: 0.9\n beta2: 0.999\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0\n save_dir: 'Re500'\n save_name: 'PINO-fintune-05s.pt'\n ckpt: 'checkpoints/Re500-FDM/PINO-pretrain-Re500-05s-4k4.pt'\n profile: True\n\nlog:\n project: 'PINO-Re500-ICLR'\n group: 'Re500-finetune-128-4k4-profile'\n\n\n" }, { "path": "configs/finetune/Re500-finetune-1s.yaml", "content": "data:\n datapath: 'data/NS_fine_Re500_T128_part2.npy'\n Re: 500\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 1\n nx: 256\n nt: 128\n sub: 2\n sub_t: 2\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 10000\n milestones: [500, 1500, 3000, 4000, 5000]\n base_lr: 0.0025\n scheduler_gamma: 0.5\n ckpt: 'checkpoints/Re500-FDM/PINO-pretrain-Re500-1s.pt'\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0\n save_dir: 'Re500-FDM'\n save_name: 'PINO-Re500-fintune-1s.pt'\n\nlog:\n project: 'PINO-sweep'\n group: 'Re500-finetune'\n\n\n\n\n" }, { "path": "configs/instance/Re500-1_8-FNO.yaml", "content": "data:\n name: KF\n paths: ['/raid/hongkai/NS-Re500_T300_id0-shuffle.npy']\n Re: 500\n offset: 0\n total_num: 300\n raw_res: [256, 256, 513]\n n_data_samples: 100\n data_res: [256, 256, 513] # resolution in 1 second\n pde_res: [256, 256, 513] # resolution in 1 second\n a_offset: 0\n n_a_samples: 250\n testoffset: 275\n n_test_samples: 10\n total_test_samples: 50\n t_duration: 0.125\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [12, 12, 12, 12]\n modes2: [12, 12, 12, 12]\n modes3: [12, 12, 12, 12]\n fc_dim: 128\n act: gelu\n pad_ratio: [0, 0.125]\n\ntrain:\n batchsize: 1\n epochs: 201\n num_iter: 1_001\n milestones: [400, 800]\n base_lr: 0.001\n scheduler_gamma: 0.5\n ic_loss: 10.0\n f_loss: 1.0\n save_step: 500\n\ntest:\n batchsize: 1\n data_res: [256, 256, 513]\n ckpt: model-400.pt\n\nlog:\n logdir: Re500-1_8s-800-FNO\n entity: hzzheng-pino\n project: PINO-NS-test-time-opt\n group: Re500-1_8s-800-FNO\n" }, { "path": "configs/instance/Re500-1_8-PINO-s.yaml", "content": "data:\n name: KF\n paths: ['/raid/hongkai/NS-Re500_T300_id0-shuffle.npy']\n Re: 500\n offset: 0\n total_num: 300\n raw_res: [256, 256, 513]\n n_data_samples: 100\n data_res: [256, 256, 513] # resolution in 1 second\n pde_res: [256, 256, 513] # resolution in 1 second\n a_offset: 0\n n_a_samples: 250\n testoffset: 250\n n_test_samples: 1\n t_duration: 0.125\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n act: gelu\n pad_ratio: 0.125\n\ntrain:\n batchsize: 1\n epochs: 201\n num_iter: 1_001\n milestones: [400, 800]\n base_lr: 0.001\n scheduler_gamma: 0.5\n ic_loss: 10.0\n f_loss: 1.0\n save_step: 500\n\ntest:\n batchsize: 1\n data_res: [256, 256, 513]\n ckpt: model-400.pt\n\nlog:\n logdir: Re500-1_8s-800-PINO-tto\n entity: hzzheng-pino\n project: PINO-NS-test-time-opt\n group: Re500-1_8s-800-PINO-s\n" }, { "path": "configs/instance/Re500-1_8-PINO.yaml", "content": "data:\n name: KF\n paths: ['/raid/hongkai/NS-Re500_T300_id0-shuffle.npy']\n Re: 500\n offset: 0\n total_num: 300\n raw_res: [256, 256, 513]\n n_data_samples: 100\n data_res: [256, 256, 513] # resolution in 1 second\n pde_res: [256, 256, 513] # resolution in 1 second\n a_offset: 0\n n_a_samples: 250\n testoffset: 275\n n_test_samples: 10\n total_test_samples: 50\n t_duration: 0.125\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [12, 12, 12, 12]\n modes2: [12, 12, 12, 12]\n modes3: [12, 12, 12, 12]\n fc_dim: 128\n act: gelu\n pad_ratio: [0.0, 0.125]\n\ntrain:\n batchsize: 1\n epochs: 201\n num_iter: 1_001\n milestones: [400, 800]\n base_lr: 0.001\n scheduler_gamma: 0.5\n ic_loss: 10.0\n f_loss: 1.0\n save_step: 500\n\ntest:\n batchsize: 1\n data_res: [256, 256, 513]\n ckpt: model-400.pt\n\nlog:\n logdir: Re500-1_8s-800-PINO-tto\n entity: hzzheng-pino\n project: PINO-NS-test-time-opt\n group: Re500-1_8s-800-PINO-s\n" }, { "path": "configs/ngc/Re500-1_8-dat0-PINO.yaml", "content": "data:\n name: KF\n paths: ['/mount/data/NS-Re500_T300_256x256x500.npy']\n Re: 500\n offset: 0\n total_num: 300\n raw_res: [256, 256, 513]\n n_data_samples: 5\n data_res: [64, 64, 257] # resolution in 1 second\n pde_res: [256, 256, 513] # resolution in 1 second\n a_offset: 0\n n_a_samples: 250\n testoffset: 200\n n_test_samples: 50\n t_duration: 0.125\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n act: gelu\n num_pad: 4\n pad_ratio: 0.0625\n\ntrain:\n batchsize: 1\n epochs: 201\n num_iter: 150_001\n milestones: [30_000, 60_000, 90_000, 120_000]\n base_lr: 0.001\n scheduler_gamma: 0.5\n ic_loss: 1.0\n f_loss: 1.0\n xy_loss: 0.0\n save_step: 5000\n eval_step: 5000\n\ntest:\n batchsize: 1\n data_res: [256, 256, 257]\n ckpt: model-400.pt\n\nlog:\n logdir: Re500-1_8s-dat0-PINO\n entity: hzzheng-pino\n project: PINO-NS\n group: Re500-1_8s-dat0-PINO\n" }, { "path": "configs/ngc/Re500-1_8-dat200-PINO.yaml", "content": "data:\n name: KF\n paths: ['/mount/data/NS-Re500_T300_256x256x500.npy']\n Re: 500\n offset: 0\n total_num: 300\n raw_res: [256, 256, 513]\n n_data_samples: 25\n data_res: [64, 64, 257] # resolution in 1 second\n pde_res: [256, 256, 513] # resolution in 1 second\n a_offset: 0\n n_a_samples: 250\n testoffset: 200\n n_test_samples: 50\n t_duration: 0.125\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n act: gelu\n num_pad: 4\n pad_ratio: 0.0625\n\ntrain:\n batchsize: 1\n epochs: 201\n num_iter: 150_001\n milestones: [20_000, 50_000, 80_000, 110_000]\n base_lr: 0.001\n scheduler_gamma: 0.5\n ic_loss: 1.0\n f_loss: 1.0\n xy_loss: 5.0\n save_step: 5000\n eval_step: 5000\n\ntest:\n batchsize: 1\n data_res: [256, 256, 257]\n ckpt: model-400.pt\n\nlog:\n logdir: Re500-1_8s-dat200-PINO\n entity: hzzheng-pino\n project: PINO-NS\n group: Re500-1_8s-dat200-PINO\n" }, { "path": "configs/ngc/Re500-1_8-dat40-PINO.yaml", "content": "data:\n name: KF\n paths: ['/mount/data/NS-Re500_T300_256x256x500.npy']\n Re: 500\n offset: 0\n total_num: 300\n raw_res: [256, 256, 513]\n n_data_samples: 5\n data_res: [64, 64, 257] # resolution in 1 second\n pde_res: [256, 256, 513] # resolution in 1 second\n a_offset: 0\n n_a_samples: 250\n testoffset: 200\n n_test_samples: 50\n t_duration: 0.125\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n act: gelu\n num_pad: 4\n pad_ratio: 0.0625\n\ntrain:\n batchsize: 1\n epochs: 201\n num_iter: 150_001\n milestones: [30_000, 60_000, 90_000, 120_000]\n base_lr: 0.001\n scheduler_gamma: 0.5\n ic_loss: 1.0\n f_loss: 1.0\n xy_loss: 5.0\n save_step: 5000\n eval_step: 5000\n\ntest:\n batchsize: 1\n data_res: [256, 256, 257]\n ckpt: model-400.pt\n\nlog:\n logdir: Re500-1_8s-dat40-PINO\n entity: hzzheng-pino\n project: PINO-NS\n group: Re500-1_8s-dat40-PINO\n" }, { "path": "configs/ngc/Re500-1_8-dat400-PINO.yaml", "content": "data:\n name: KF\n paths: ['/mount/data/NS-Re500_T300_256x256x500.npy']\n Re: 500\n offset: 0\n total_num: 300\n raw_res: [256, 256, 513]\n n_data_samples: 50\n data_res: [64, 64, 257] # resolution in 1 second\n pde_res: [256, 256, 513] # resolution in 1 second\n a_offset: 0\n n_a_samples: 250\n testoffset: 200\n n_test_samples: 50\n t_duration: 0.125\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n act: gelu\n num_pad: 4\n pad_ratio: 0.0625\n\ntrain:\n batchsize: 1\n epochs: 201\n num_iter: 150_001\n milestones: [20_000, 50_000, 80_000, 110_000]\n base_lr: 0.001\n scheduler_gamma: 0.5\n ic_loss: 1.0\n f_loss: 1.0\n xy_loss: 5.0\n save_step: 5000\n eval_step: 5000\n\ntest:\n batchsize: 1\n data_res: [256, 256, 257]\n ckpt: model-400.pt\n\nlog:\n logdir: Re500-1_8s-dat400-PINO\n entity: hzzheng-pino\n project: PINO-NS\n group: Re500-1_8s-dat400-PINO\n" }, { "path": "configs/ngc/Re500-1_8-dat80-PINO.yaml", "content": "data:\n name: KF\n paths: ['/mount/data/NS-Re500_T300_256x256x500.npy']\n Re: 500\n offset: 0\n total_num: 300\n raw_res: [256, 256, 513]\n n_data_samples: 10\n data_res: [64, 64, 257] # resolution in 1 second\n pde_res: [256, 256, 513] # resolution in 1 second\n a_offset: 0\n n_a_samples: 250\n testoffset: 200\n n_test_samples: 50\n t_duration: 0.125\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n act: gelu\n num_pad: 4\n pad_ratio: 0.0625\n\ntrain:\n batchsize: 1\n epochs: 201\n num_iter: 150_001\n milestones: [30_000, 60_000, 90_000, 120_000]\n base_lr: 0.001\n scheduler_gamma: 0.5\n ic_loss: 1.0\n f_loss: 1.0\n xy_loss: 5.0\n save_step: 5000\n eval_step: 5000\n\ntest:\n batchsize: 1\n data_res: [256, 256, 257]\n ckpt: model-400.pt\n\nlog:\n logdir: Re500-1_8s-dat80-PINO\n entity: hzzheng-pino\n project: PINO-NS\n group: Re500-1_8s-dat80-PINO\n" }, { "path": "configs/ngc/Re500-1_8-dat800-PINO.yaml", "content": "data:\n name: KF\n paths: ['/mount/data/NS-Re500_T300_256x256x500.npy']\n Re: 500\n offset: 0\n total_num: 300\n raw_res: [256, 256, 513]\n n_data_samples: 100\n data_res: [64, 64, 257] # resolution in 1 second\n pde_res: [256, 256, 513] # resolution in 1 second\n a_offset: 0\n n_a_samples: 250\n testoffset: 200\n n_test_samples: 50\n t_duration: 0.125\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n act: gelu\n num_pad: 4\n pad_ratio: 0.0625\n\ntrain:\n batchsize: 1\n epochs: 201\n num_iter: 150_001\n milestones: [30_000, 60_000, 90_000, 120_000]\n base_lr: 0.001\n scheduler_gamma: 0.5\n ic_loss: 1.0\n f_loss: 1.0\n xy_loss: 5.0\n save_step: 5000\n eval_step: 5000\n\ntest:\n batchsize: 1\n data_res: [256, 256, 257]\n ckpt: model-400.pt\n\nlog:\n logdir: Re500-1_8s-dat800-PINO\n entity: hzzheng-pino\n project: PINO-NS-ngc\n group: Re500-1_8s-dat800-PINO\n" }, { "path": "configs/ngc/Re500-1_8-res16-PINO.yaml", "content": "data:\n name: KF\n paths: ['/mount/data/NS-Re500_T300_256x256x500.npy']\n Re: 500\n offset: 0\n total_num: 300\n raw_res: [256, 256, 513]\n n_data_samples: 150\n data_res: [16, 16, 129] # resolution in 1 second\n pde_res: [256, 256, 513] # resolution in 1 second\n a_offset: 0\n n_a_samples: 250\n testoffset: 200\n n_test_samples: 50\n t_duration: 0.125\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n act: gelu\n num_pad: 4\n pad_ratio: 0.0625\n\ntrain:\n batchsize: 1\n epochs: 201\n num_iter: 150_001\n milestones: [30_000, 60_000, 90_000, 120_000]\n base_lr: 0.001\n scheduler_gamma: 0.5\n ic_loss: 1.0\n f_loss: 1.0\n xy_loss: 5.0\n save_step: 5000\n eval_step: 5000\n\ntest:\n batchsize: 1\n data_res: [256, 256, 257]\n ckpt: model-400.pt\n\nlog:\n logdir: Re500-1_8s-res16-PINO\n entity: hzzheng-pino\n project: PINO-NS\n group: Re500-1_8s-res16-PINO\n" }, { "path": "configs/ngc/Re500-1_8-res32-PINO.yaml", "content": "data:\n name: KF\n paths: ['/mount/data/NS-Re500_T300_256x256x500.npy']\n Re: 500\n offset: 0\n total_num: 300\n raw_res: [256, 256, 513]\n n_data_samples: 150\n data_res: [32, 32, 129] # resolution in 1 second\n pde_res: [256, 256, 513] # resolution in 1 second\n a_offset: 0\n n_a_samples: 250\n testoffset: 200\n n_test_samples: 50\n t_duration: 0.125\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n act: gelu\n num_pad: 4\n pad_ratio: 0.0625\n\ntrain:\n batchsize: 1\n epochs: 201\n num_iter: 150_001\n milestones: [30_000, 60_000, 90_000, 120_000]\n base_lr: 0.001\n scheduler_gamma: 0.5\n ic_loss: 1.0\n f_loss: 1.0\n xy_loss: 5.0\n save_step: 5000\n eval_step: 5000\n\ntest:\n batchsize: 1\n data_res: [256, 256, 257]\n ckpt: model-400.pt\n\nlog:\n logdir: Re500-1_8s-res32-PINO\n entity: hzzheng-pino\n project: PINO-NS\n group: Re500-1_8s-res32-PINO\n" }, { "path": "configs/operator/Darcy-pretrain.yaml", "content": "data:\n name: 'Darcy'\n path: '/raid/hongkai/darcy-train.mat'\n total_num: 1024\n offset: 0\n n_sample: 1000\n nx: 421\n sub: 7\n pde_sub: 2\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [20, 20, 20, 20]\n modes2: [20, 20, 20, 20]\n fc_dim: 128\n act: gelu\n pad_ratio: [0., 0.]\n\ntrain:\n batchsize: 20\n num_iter: 15_001\n milestones: [5_000, 7_500, 10_000]\n base_lr: 0.001\n scheduler_gamma: 0.5\n f_loss: 1.0\n xy_loss: 5.0\n save_step: 2_500\n eval_step: 2_500\n\ntest:\n path: '/raid/hongkai/darcy-test.mat'\n total_num: 1024\n offset: 0\n n_sample: 500\n nx: 421\n sub: 2\n batchsize: 1\n \n \nlog:\n logdir: Darcy-PINO-new\n entity: hzzheng-pino\n project: DarcyFlow\n group: PINO-1000-new\n\n\n" }, { "path": "configs/operator/Re500-05s-1000-FNO.yaml", "content": "data:\n paths: ['../data/NS-Re500Part0.npy', '../data/NS-Re500Part1.npy', '../data/NS-Re500Part2.npy']\n Re: 500\n total_num: 200\n offset: 0\n n_samples: 1000\n t_duration: 0.5\n data_res: [64, 64, 33]\n pde_res: [128, 128, 65]\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n act: gelu\n num_pad: 4\n\ntrain:\n batchsize: 2\n epochs: 501\n milestones: [300]\n base_lr: 0.001\n scheduler_gamma: 0.5\n ic_loss: 0.0\n f_loss: 0.0\n xy_loss: 1.0\n save_step: 100\n\ntest:\n batchsize: 1\n data_res: [128, 128, 65]\n ckpt: model-500.pt\n\nlog:\n logdir: Re500-05s-1000-FNO\n entity: hzzheng-pino\n project: 'PINO-NS'\n group: 'Re500-05s-1000-FNO'\n" }, { "path": "configs/operator/Re500-05s-1000-PINO.yaml", "content": "data:\n paths: ['../data/NS-Re500Part0.npy', '../data/NS-Re500Part1.npy', '../data/NS-Re500Part2.npy']\n Re: 500\n total_num: 300\n offset: 0\n n_samples: 1000\n t_duration: 0.5\n data_res: [64, 64, 33]\n pde_res: [128, 128, 65]\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n act: gelu\n num_pad: 4\n\ntrain:\n batchsize: 2\n epochs: 501\n milestones: [300]\n base_lr: 0.001\n scheduler_gamma: 0.5\n ic_loss: 1.0\n f_loss: 1.0\n xy_loss: 5.0\n save_step: 100\n\ntest:\n batchsize: 1\n data_res: [128, 128, 65]\n ckpt: model-500.pt\n\nlog:\n logdir: Re500-05s-1000-PINO\n entity: hzzheng-pino\n project: PINO-NS\n group: Re500-05s-1000-PINO\n" }, { "path": "configs/operator/Re500-05s-3000-FNO.yaml", "content": "data:\n name: KF\n paths: ['../data/NS-Re500_T3000_id0.npy']\n Re: 500\n total_num: 3000\n offset: 0\n n_samples: 300\n testoffset: 2500\n n_test_samples: 300\n t_duration: 0.5\n raw_res: [256, 256, 257]\n data_res: [64, 64, 33]\n pde_res: [64, 64, 33]\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n act: gelu\n num_pad: 4\n\ntrain:\n batchsize: 4\n epochs: 401\n milestones: [200]\n base_lr: 0.001\n scheduler_gamma: 0.5\n ic_loss: 0.0\n f_loss: 0.0\n xy_loss: 1.0\n save_step: 50\n\ntest:\n batchsize: 1\n data_res: [128, 128, 65]\n ckpt: model-400.pt\n\nlog:\n logdir: Re500-1s-3000-FNO\n entity: hzzheng-pino\n project: PINO-NS\n group: Re500-1s-3000-FNO\n" }, { "path": "configs/operator/Re500-05s-600-FNO.yaml", "content": "data:\n name: KF\n paths: ['../data/NS-Re500_T3000_id0.npy']\n Re: 500\n total_num: 3000\n offset: 0\n n_samples: 300\n testoffset: 2500\n n_test_samples: 200\n t_duration: 0.5\n raw_res: [256, 256, 257]\n data_res: [64, 64, 65] # resolution in 1 second\n pde_res: [64, 64, 65] # resolution in 1 second\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n act: gelu\n num_pad: 4\n\ntrain:\n batchsize: 2\n epochs: 401\n milestones: [200]\n base_lr: 0.001\n scheduler_gamma: 0.5\n ic_loss: 0.0\n f_loss: 0.0\n xy_loss: 1.0\n save_step: 50\n\ntest:\n batchsize: 1\n data_res: [64, 64, 65]\n ckpt: model-400.pt\n\nlog:\n logdir: Re500-05s-600-FNO\n entity: hzzheng-pino\n project: PINO-NS\n group: Re500-05s-600-FNO\n" }, { "path": "configs/operator/Re500-05s-600-PINO-xl.yaml", "content": "data:\n name: KF\n paths: ['../data/NS-Re500_T3000_id0.npy']\n Re: 500\n total_num: 3000\n offset: 0\n n_samples: 300\n testoffset: 2500\n n_test_samples: 200\n t_duration: 0.5\n raw_res: [256, 256, 257]\n data_res: [64, 64, 65] # resolution in 1 second\n pde_res: [256, 256, 257] # resolution in 1 second\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [12, 12, 12, 12]\n modes2: [12, 12, 12, 12]\n modes3: [12, 12, 12, 12]\n fc_dim: 128\n act: gelu\n num_pad: 4\n\ntrain:\n batchsize: 1\n epochs: 301\n milestones: [200]\n base_lr: 0.001\n scheduler_gamma: 0.5\n ic_loss: 1.0\n f_loss: 1.0\n xy_loss: 5.0\n save_step: 10\n\ntest:\n batchsize: 1\n data_res: [64, 64, 65]\n ckpt: model-400.pt\n\nlog:\n logdir: Re500-05s-600-PINO-xl\n entity: hzzheng-pino\n project: PINO-NS\n group: Re500-05s-600-PINO-xl\n" }, { "path": "configs/operator/Re500-05s-600-PINO.yaml", "content": "data:\n name: KF\n paths: ['../data/NS-Re500_T3000_id0.npy']\n Re: 500\n total_num: 3000\n offset: 0\n n_samples: 300\n testoffset: 2500\n n_test_samples: 200\n t_duration: 0.5\n raw_res: [256, 256, 257]\n data_res: [64, 64, 65] # resolution in 1 second\n pde_res: [256, 256, 257] # resolution in 1 second\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n act: gelu\n num_pad: 4\n\ntrain:\n batchsize: 1\n epochs: 301\n milestones: [200]\n base_lr: 0.001\n scheduler_gamma: 0.5\n ic_loss: 1.0\n f_loss: 1.0\n xy_loss: 5.0\n save_step: 10\n\ntest:\n batchsize: 1\n data_res: [64, 64, 65]\n ckpt: model-400.pt\n\nlog:\n logdir: Re500-05s-600-PINO\n entity: hzzheng-pino\n project: PINO-NS\n group: Re500-05s-600-PINO\n" }, { "path": "configs/operator/Re500-05s-FNO.yaml", "content": "data:\n paths: ['../data/NS-Re500Part0.npy', '../data/NS-Re500Part1.npy']\n Re: 500\n total_num: 200\n offset: 0\n n_samples: 700\n t_duration: 0.5\n data_res: [64, 64, 33]\n pde_res: [128, 128, 65]\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n act: gelu\n num_pad: 4\n\ntrain:\n batchsize: 2\n epochs: 501\n milestones: [300]\n base_lr: 0.001\n scheduler_gamma: 0.5\n ic_loss: 0.0\n f_loss: 0.0\n xy_loss: 1.0\n save_step: 100\n\ntest:\n batchsize: 1\n data_res: [128, 128, 65]\n ckpt: model-500.pt\n\nlog:\n logdir: Re500-05s-FNO\n entity: hzzheng-pino\n project: 'PINO-NS'\n group: 'Re500-05s-FNO'\n" }, { "path": "configs/operator/Re500-1_16-800-FNO-s.yaml", "content": "data:\n name: KF\n paths: ['/raid/hongkai/NS-Re500_T300_id0-shuffle.npy']\n Re: 500\n offset: 0\n total_num: 300\n raw_res: [256, 256, 513]\n n_data_samples: 50\n data_res: [64, 64, 257] # resolution in 1 second\n pde_res: [64, 64, 257] # resolution in 1 second\n a_offset: 0\n n_a_samples: 50\n testoffset: 275\n n_test_samples: 25\n t_duration: 0.0625\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n act: gelu\n pad_ratio: 0.125\n\ntrain:\n batchsize: 2\n start_iter: 0\n num_iter: 50_001\n milestones: [20_000, 40_000]\n base_lr: 0.001\n scheduler_gamma: 0.5\n ic_loss: 0.0\n f_loss: 0.0\n xy_loss: 1.0\n save_step: 5000\n eval_step: 5000\n\ntest:\n batchsize: 1\n data_res: [64, 64, 257]\n ckpt: model-400.pt\n\nlog:\n logdir: Re500-1_16s-800-FNO-s\n entity: hzzheng-pino\n project: PINO-KF-Re500\n group: Re500-1_16s-800-FNO-s\n" }, { "path": "configs/operator/Re500-1_16-800-PINO-s.yaml", "content": "data:\n name: KF\n paths: ['/raid/hongkai/NS-Re500_T300_id0-shuffle.npy']\n Re: 500\n offset: 0\n total_num: 300\n raw_res: [256, 256, 513]\n n_data_samples: 50\n data_res: [64, 64, 257] # resolution in 1 second\n pde_res: [256, 256, 513] # resolution in 1 second\n a_offset: 0\n n_a_samples: 200\n testoffset: 275\n n_test_samples: 25\n t_duration: 0.0625\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n act: gelu\n num_pad: 4\n pad_ratio: 0.125\n\ntrain:\n batchsize: 2\n start_iter: 0\n num_iter: 200_001\n milestones: [20_000, 60_000, 120_000]\n base_lr: 0.001\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 10.0\n save_step: 5000\n eval_step: 5000\n\ntest:\n batchsize: 1\n data_res: [128, 128, 257]\n ckpt: model-400.pt\n\nlog:\n logdir: Re500-1_16s-800-PINO-s\n entity: hzzheng-pino\n project: PINO-KF-Re500\n group: Re500-1_16s-800-PINO-s\n" }, { "path": "configs/operator/Re500-1_4-2000-FNO.yaml", "content": "data:\n name: KF\n paths: ['../data/NS-Re500_T3000_id0.npy']\n Re: 500\n total_num: 3000\n offset: 0\n n_samples: 600\n testoffset: 2500\n n_test_samples: 400\n t_duration: 0.25\n raw_res: [256, 256, 257]\n data_res: [256, 256, 257] # resolution in 1 second\n pde_res: [256, 256, 257] # resolution in 1 second\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n act: gelu\n num_pad: 4\n\ntrain:\n batchsize: 2\n epochs: 401\n milestones: [100, 300]\n base_lr: 0.001\n scheduler_gamma: 0.5\n ic_loss: 0.0\n f_loss: 0.0\n xy_loss: 1.0\n save_step: 50\n\ntest:\n batchsize: 1\n data_res: [256, 256, 257]\n ckpt: model-400.pt\n\nlog:\n logdir: Re500-1_4s-2000-FNO\n entity: hzzheng-pino\n project: PINO-NS\n group: Re500-1_4s-2000-FNO\n" }, { "path": "configs/operator/Re500-1_8-0-PINO-s.yaml", "content": "data:\n name: KF\n paths: ['/raid/hongkai/NS-Re500_T300_id0-shuffle.npy']\n Re: 500\n offset: 0\n total_num: 300\n raw_res: [256, 256, 513]\n n_data_samples: 10\n data_res: [64, 64, 129] # resolution in 1 second\n pde_res: [256, 256, 513] # resolution in 1 second\n a_offset: 0\n n_a_samples: 250\n testoffset: 275\n n_test_samples: 25\n t_duration: 0.125\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [12, 12, 12, 12]\n modes2: [12, 12, 12, 12]\n modes3: [12, 12, 12, 12]\n fc_dim: 128\n act: gelu\n num_pad: 4\n pad_ratio: 0.125\n\ntrain:\n batchsize: 2\n start_iter: 35_001\n num_iter: 200_001\n milestones: [30_000, 70_000, 110_000, 150_000]\n base_lr: 0.001\n scheduler_gamma: 0.5\n ic_loss: 10.0\n f_loss: 1.0\n xy_loss: 0.0\n save_step: 5000\n eval_step: 5000\n\ntest:\n batchsize: 1\n data_res: [256, 256, 513]\n ckpt: model-400.pt\n\nlog:\n logdir: Re500-1_8s-0-PINO-s\n entity: hzzheng-pino\n project: PINO-KF-Re500\n group: Re500-1_8s-0-PINO-s\n" }, { "path": "configs/operator/Re500-1_8-1200-FNO.yaml", "content": "data:\n name: KF\n paths: ['../data/NS-Re500_T300_id0.npy']\n Re: 500\n offset: 0\n total_num: 300\n raw_res: [256, 256, 513]\n n_data_samples: 150\n data_res: [64, 64, 129] # resolution in 1 second\n pde_res: [64, 64, 129] # resolution in 1 second\n a_offset: 0\n n_a_samples: 250\n testoffset: 250\n n_test_samples: 50\n t_duration: 0.125\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n act: gelu\n num_pad: 4\n pad_ratio: 0.125\n\ntrain:\n batchsize: 2\n epochs: 201\n num_iter: 50_001\n milestones: [20_000, 40_000]\n base_lr: 0.001\n scheduler_gamma: 0.5\n ic_loss: 0.0\n f_loss: 0.0\n xy_loss: 1.0\n save_step: 5000\n eval_step: 5000\n\ntest:\n batchsize: 1\n data_res: [64, 64, 129]\n ckpt: model-400.pt\n\nlog:\n logdir: Re500-1_8s-1200-FNO\n entity: hzzheng-pino\n project: PINO-NS\n group: Re500-1_8s-1200-FNO\n" }, { "path": "configs/operator/Re500-1_8-1200-PINO.yaml", "content": "data:\n name: KF\n paths: ['/raid/hongkai/NS-Re500_T300_id0-shuffle.npy']\n Re: 500\n offset: 0\n total_num: 300\n raw_res: [256, 256, 513]\n n_data_samples: 150\n data_res: [64, 64, 129] # resolution in 1 second\n pde_res: [256, 256, 513] # resolution in 1 second\n a_offset: 0\n n_a_samples: 250\n testoffset: 250\n n_test_samples: 50\n t_duration: 0.125\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n act: gelu\n num_pad: 4\n pad_ratio: 0.125\n\ntrain:\n batchsize: 2\n epochs: 201\n num_iter: 150_001\n milestones: [30_000, 60_000, 90_000]\n base_lr: 0.001\n scheduler_gamma: 0.5\n ic_loss: 1.0\n f_loss: 1.0\n xy_loss: 5.0\n save_step: 5000\n eval_step: 5000\n\ntest:\n batchsize: 1\n data_res: [256, 256, 513]\n ckpt: model-400.pt\n\nlog:\n logdir: Re500-1_8s-1200-PINO\n entity: hzzheng-pino\n project: PINO-NS\n group: Re500-1_8s-1200-PINO\n" }, { "path": "configs/operator/Re500-1_8-200-FNO-s.yaml", "content": "data:\n name: KF\n paths: ['/raid/hongkai/NS-Re500_T300_id0-shuffle.npy']\n Re: 500\n offset: 0\n total_num: 300\n raw_res: [256, 256, 513]\n n_data_samples: 25\n data_res: [128, 128, 257] # resolution in 1 second\n pde_res: [128, 128, 257] # resolution in 1 second\n a_offset: 0\n n_a_samples: 250\n testoffset: 250\n n_test_samples: 25\n t_duration: 0.125\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n act: gelu\n num_pad: 4\n pad_ratio: [0, 0.125]\n\ntrain:\n batchsize: 1\n epochs: 201\n num_iter: 50_001\n milestones: [20_000, 40_000]\n base_lr: 0.001\n scheduler_gamma: 0.5\n ic_loss: 0.0\n f_loss: 0.0\n xy_loss: 1.0\n save_step: 5000\n eval_step: 5000\n\ntest:\n batchsize: 1\n data_res: [256, 256, 513]\n ckpt: model-400.pt\n\nlog:\n logdir: Re500-1_8s-dat200-FNO\n entity: hzzheng-pino\n project: PINO-KF-Re500\n group: Re500-1_8s-dat200-FNO\n" }, { "path": "configs/operator/Re500-1_8-2000-FNO-s.yaml", "content": "data:\n name: KF\n paths: ['/raid/hongkai/NS-Re500_T300_id0-shuffle.npy']\n Re: 500\n offset: 0\n total_num: 300\n raw_res: [256, 256, 513]\n n_data_samples: 250\n data_res: [64, 64, 129] # resolution in 1 second\n pde_res: [64, 64, 129] # resolution in 1 second\n a_offset: 0\n n_a_samples: 250\n testoffset: 275\n n_test_samples: 25\n t_duration: 0.125\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n act: gelu\n pad_ratio: 0.125\n\ntrain:\n batchsize: 1\n epochs: 201\n num_iter: 60_001\n milestones: [20_000, 40_000]\n base_lr: 0.001\n scheduler_gamma: 0.5\n ic_loss: 0.0\n f_loss: 0.0\n xy_loss: 1.0\n save_step: 5000\n eval_step: 5000\n\ntest:\n batchsize: 1\n data_res: [256, 256, 513]\n ckpt: model-400.pt\n\nlog:\n logdir: Re500-1_8s-dat2000-FNO\n entity: hzzheng-pino\n project: PINO-KF-Re500\n group: Re500-1_8s-dat2000-FNO\n" }, { "path": "configs/operator/Re500-1_8-2000-FNO-xl.yaml", "content": "data:\n name: KF\n paths: ['../data/NS-Re500_T3000_id0.npy']\n Re: 500\n total_num: 3000\n offset: 0\n n_samples: 350\n testoffset: 2500\n n_test_samples: 400\n t_duration: 0.125\n raw_res: [256, 256, 257]\n data_res: [256, 256, 257] # resolution in 1 second\n pde_res: [256, 256, 257] # resolution in 1 second\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [12, 12, 12, 12]\n modes2: [12, 12, 12, 12]\n modes3: [12, 12, 12, 12]\n fc_dim: 128\n act: gelu\n num_pad: 4\n\ntrain:\n batchsize: 2\n epochs: 201\n milestones: [50, 100, 150]\n base_lr: 0.001\n scheduler_gamma: 0.5\n ic_loss: 0.0\n f_loss: 0.0\n xy_loss: 1.0\n save_step: 20\n\ntest:\n batchsize: 1\n data_res: [256, 256, 257]\n ckpt: model-400.pt\n\nlog:\n logdir: Re500-1_8s-2400-FNO\n entity: hzzheng-pino\n project: PINO-NS\n group: Re500-1_8s-2400-FNO\n" }, { "path": "configs/operator/Re500-1_8-2000-PINO.yaml", "content": "data:\n name: KF\n paths: ['../data/NS-Re500_T300_id0.npy']\n Re: 500\n offset: 0\n total_num: 300\n raw_res: [256, 256, 513]\n n_data_samples: 150\n data_res: [64, 64, 257] # resolution in 1 second\n pde_res: [256, 256, 513] # resolution in 1 second\n a_offset: 0\n n_a_samples: 250\n testoffset: 200\n n_test_samples: 50\n t_duration: 0.125\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n act: gelu\n num_pad: 4\n pad_ratio: 0.0625\n\ntrain:\n batchsize: 2\n epochs: 201\n num_iter: 100_001\n milestones: [10_000, 30_000, 50_000, 70_000]\n base_lr: 0.001\n scheduler_gamma: 0.5\n ic_loss: 1.0\n f_loss: 1.0\n xy_loss: 5.0\n save_step: 5000\n eval_step: 5000\n\ntest:\n batchsize: 1\n data_res: [64, 64, 257]\n ckpt: model-400.pt\n\nlog:\n logdir: Re500-1_8s-2k-PINO\n entity: hzzheng-pino\n project: PINO-NS\n group: Re500-1_8s-2k-PINO\n" }, { "path": "configs/operator/Re500-1_8-2200-FNO-s.yaml", "content": "data:\n name: KF\n paths: ['/raid/hongkai/NS-Re500_T300_id0-shuffle.npy']\n Re: 500\n offset: 0\n total_num: 300\n raw_res: [256, 256, 513]\n n_data_samples: 275\n data_res: [64, 64, 129] # resolution in 1 second\n pde_res: [64, 64, 257] # resolution in 1 second\n a_offset: 0\n n_a_samples: 250\n testoffset: 275\n n_test_samples: 25\n t_duration: 0.125\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [12, 12, 12, 12]\n modes2: [12, 12, 12, 12]\n modes3: [12, 12, 12, 12]\n fc_dim: 128\n act: gelu\n pad_ratio: 0.125\n\ntrain:\n batchsize: 1\n start_iter: 30_001\n num_iter: 60_001\n milestones: [20_000, 40_000]\n base_lr: 0.001\n scheduler_gamma: 0.5\n ic_loss: 0.0\n f_loss: 0.0\n xy_loss: 1.0\n save_step: 5000\n eval_step: 5000\n\ntest:\n batchsize: 1\n data_res: [256, 256, 513]\n ckpt: model-400.pt\n\nlog:\n logdir: Re500-1_8s-dat2200-FNO\n entity: hzzheng-pino\n project: PINO-KF-Re500\n group: Re500-1_8s-dat2200-FNO\n" }, { "path": "configs/operator/Re500-1_8-2200-PINO-s.yaml", "content": "data:\n name: KF\n paths: ['/raid/hongkai/NS-Re500_T300_id0-shuffle.npy']\n Re: 500\n offset: 0\n total_num: 300\n raw_res: [256, 256, 513]\n n_data_samples: 275\n data_res: [64, 64, 257] # resolution in 1 second\n pde_res: [256, 256, 513] # resolution in 1 second\n a_offset: 0\n n_a_samples: 275\n testoffset: 275\n n_test_samples: 25\n t_duration: 0.125\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [12, 12, 12, 12]\n modes2: [12, 12, 12, 12]\n modes3: [12, 12, 12, 12]\n fc_dim: 128\n act: gelu\n num_pad: 4\n pad_ratio: 0.125\n\ntrain:\n batchsize: 2\n start_iter: 30_001\n num_iter: 400_001\n milestones: [30_000, 90_000, 150_000, 250_000]\n base_lr: 0.001\n scheduler_gamma: 0.5\n ic_loss: 10.0\n f_loss: 1.0\n xy_loss: 50.0\n save_step: 10000\n eval_step: 10000\n\ntest:\n batchsize: 1\n data_res: [64, 64, 257]\n ckpt: model-400.pt\n\nlog:\n logdir: Re500-1_8s-2200-PINO-s\n entity: hzzheng-pino\n project: PINO-KF-Re500\n group: Re500-1_8s-2200-PINO-s\n" }, { "path": "configs/operator/Re500-1_8-800-FNO-s.yaml", "content": "data:\n name: KF\n paths: ['/raid/hongkai/NS-Re500_T300_id0-shuffle.npy']\n Re: 500\n offset: 0\n total_num: 300\n raw_res: [256, 256, 513]\n n_data_samples: 100\n data_res: [64, 64, 257] # resolution in 1 second\n pde_res: [64, 64, 129] # resolution in 1 second\n a_offset: 0\n n_a_samples: 250\n testoffset: 275\n n_test_samples: 25\n t_duration: 0.125\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [12, 12, 12, 12]\n modes2: [12, 12, 12, 12]\n modes3: [12, 12, 12, 12]\n fc_dim: 128\n act: gelu\n pad_ratio: [0, 0.125]\n\ntrain:\n batchsize: 2\n start_iter: 0\n num_iter: 50_001\n milestones: [20_000, 40_000]\n base_lr: 0.001\n scheduler_gamma: 0.5\n ic_loss: 0.0\n f_loss: 0.0\n xy_loss: 1.0\n save_step: 5000\n eval_step: 5000\n\ntest:\n batchsize: 1\n data_res: [64, 64, 257]\n ckpt: model-400.pt\n\nlog:\n logdir: Re500-1_8s-800-FNO-s\n entity: hzzheng-pino\n project: PINO-KF-Re500\n group: Re500-1_8s-800-FNO-s\n" }, { "path": "configs/operator/Re500-1_8-800-FNO-s32.yaml", "content": "data:\n name: KF\n paths: ['/raid/hongkai/NS-Re500_T300_id0-shuffle.npy']\n Re: 500\n offset: 0\n total_num: 300\n raw_res: [256, 256, 513]\n n_data_samples: 100\n data_res: [32, 32, 129] # resolution in 1 second\n pde_res: [32, 32, 129] # resolution in 1 second\n a_offset: 0\n n_a_samples: 250\n testoffset: 275\n n_test_samples: 25\n t_duration: 0.125\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [12, 12, 12, 12]\n modes2: [12, 12, 12, 12]\n modes3: [12, 12, 12, 12]\n fc_dim: 128\n act: gelu\n pad_ratio: [0.0, 0.125]\n\ntrain:\n batchsize: 2\n start_iter: 0\n num_iter: 50_001\n milestones: [20_000, 40_000]\n base_lr: 0.001\n scheduler_gamma: 0.5\n ic_loss: 0.0\n f_loss: 0.0\n xy_loss: 1.0\n save_step: 5000\n eval_step: 5000\n\ntest:\n batchsize: 1\n data_res: [256, 256, 513]\n\nlog:\n logdir: Re500-1_8s-800-FNO-s32\n entity: hzzheng-pino\n project: PINO-KF-Re500\n group: Re500-1_8s-800-FNO-s32\n" }, { "path": "configs/operator/Re500-1_8-800-PINO-s.yaml", "content": "data:\n name: KF\n paths: ['/raid/hongkai/NS-Re500_T300_id0-shuffle.npy']\n Re: 500\n offset: 0\n total_num: 300\n raw_res: [256, 256, 513]\n n_data_samples: 100\n data_res: [64, 64, 257] # resolution in 1 second\n pde_res: [256, 256, 513] # resolution in 1 second\n a_offset: 0\n n_a_samples: 275\n testoffset: 275\n n_test_samples: 25\n t_duration: 0.125\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [12, 12, 12, 12]\n modes2: [12, 12, 12, 12]\n modes3: [12, 12, 12, 12]\n fc_dim: 128\n act: gelu\n pad_ratio: [0.0, 0.125]\n\ntrain:\n batchsize: 2\n start_iter: 0\n num_iter: 200_001\n milestones: [20_000, 60_000, 120_000]\n base_lr: 0.001\n scheduler_gamma: 0.5\n ic_loss: 10.0\n f_loss: 1.0\n xy_loss: 10.0\n save_step: 5000\n eval_step: 5000\n\ntest:\n batchsize: 1\n data_res: [256, 256, 513]\n\nlog:\n logdir: Re500-1_8s-800-PINO-s\n entity: hzzheng-pino\n project: PINO-KF-Re500\n group: Re500-1_8s-800-PINO-s\n" }, { "path": "configs/operator/Re500-1_8-800-PINO-s16.yaml", "content": "data:\n name: KF\n paths: ['/raid/hongkai/NS-Re500_T300_id0-shuffle.npy']\n Re: 500\n offset: 0\n total_num: 300\n raw_res: [256, 256, 513]\n n_data_samples: 100\n data_res: [16, 16, 65] # resolution in 1 second\n pde_res: [256, 256, 513] # resolution in 1 second\n a_offset: 0\n n_a_samples: 275\n testoffset: 275\n n_test_samples: 25\n t_duration: 0.125\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [12, 12, 12, 12]\n modes2: [12, 12, 12, 12]\n modes3: [12, 12, 12, 12]\n fc_dim: 128\n act: gelu\n pad_ratio: [0.0, 0.125]\n\ntrain:\n batchsize: 1\n start_iter: 0\n num_iter: 200_001\n milestones: [20_000, 60_000, 120_000]\n base_lr: 0.001\n scheduler_gamma: 0.5\n ic_loss: 10.0\n f_loss: 1.0\n xy_loss: 10.0\n save_step: 5000\n eval_step: 5000\n\ntest:\n batchsize: 1\n data_res: [256, 256, 513]\n\nlog:\n logdir: Re500-1_8s-800-PINO-s-16\n entity: hzzheng-pino\n project: PINO-KF-Re500\n group: Re500-1_8s-800-PINO-s-16\n" }, { "path": "configs/operator/Re500-1_8-800-PINO-s32.yaml", "content": "data:\n name: KF\n paths: ['/raid/hongkai/NS-Re500_T300_id0-shuffle.npy']\n Re: 500\n offset: 0\n total_num: 300\n raw_res: [256, 256, 513]\n n_data_samples: 100\n data_res: [32, 32, 129] # resolution in 1 second\n pde_res: [256, 256, 513] # resolution in 1 second\n a_offset: 0\n n_a_samples: 275\n testoffset: 275\n n_test_samples: 25\n t_duration: 0.125\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [12, 12, 12, 12]\n modes2: [12, 12, 12, 12]\n modes3: [12, 12, 12, 12]\n fc_dim: 128\n act: gelu\n pad_ratio: [0.0, 0.125]\n\ntrain:\n batchsize: 2\n start_iter: 0\n num_iter: 200_001\n milestones: [20_000, 60_000, 120_000]\n base_lr: 0.001\n scheduler_gamma: 0.5\n ic_loss: 10.0\n f_loss: 1.0\n xy_loss: 10.0\n save_step: 5000\n eval_step: 5000\n\ntest:\n batchsize: 1\n data_res: [64, 64, 257]\n\nlog:\n logdir: Re500-1_8s-800-PINO-s-32\n entity: hzzheng-pino\n project: PINO-KF-Re500\n group: Re500-1_8s-800-PINO-s-32\n" }, { "path": "configs/operator/Re500-1_8-800-UNet.yaml", "content": "data:\n name: KF\n paths: ['/raid/hongkai/NS-Re500_T300_id0-shuffle.npy']\n Re: 500\n offset: 0\n total_num: 300\n raw_res: [256, 256, 513]\n n_data_samples: 100\n data_res: [64, 64, 257] # resolution in 1 second\n pde_res: [64, 64, 129] # resolution in 1 second\n a_offset: 0\n n_a_samples: 250\n testoffset: 275\n n_test_samples: 25\n t_duration: 0.125\n shuffle: True\n\nmodel:\n f_maps: 128\n\ntrain:\n batchsize: 2\n start_iter: 0\n num_iter: 50_001\n milestones: [20_000, 40_000]\n base_lr: 0.0002\n scheduler_gamma: 0.5\n save_step: 5000\n eval_step: 5000\n\ntest:\n batchsize: 1\n data_res: [64, 64, 257]\n ckpt: model-5000.pt\n\nlog:\n logdir: Re500-1_8s-800-UNet\n entity: hzzheng-pino\n project: PINO-KF-Re500\n group: Re500-1_8s-800-UNet\n" }, { "path": "configs/operator/Re500-1_8-dat1.6k-PINO.yaml", "content": "data:\n name: KF\n paths: ['../data/NS-Re500_T300_id0.npy']\n Re: 500\n offset: 0\n total_num: 300\n raw_res: [256, 256, 513]\n n_data_samples: 200\n data_res: [64, 64, 257] # resolution in 1 second\n pde_res: [256, 256, 513] # resolution in 1 second\n a_offset: 0\n n_a_samples: 250\n testoffset: 200\n n_test_samples: 50\n t_duration: 0.125\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n act: gelu\n num_pad: 4\n pad_ratio: 0.0625\n\ntrain:\n batchsize: 2\n epochs: 201\n num_iter: 200_001\n milestones: [20_000, 70_000, 120_000]\n base_lr: 0.001\n scheduler_gamma: 0.5\n ic_loss: 1.0\n f_loss: 1.0\n xy_loss: 5.0\n save_step: 5000\n eval_step: 5000\n\ntest:\n batchsize: 1\n data_res: [64, 64, 257]\n ckpt: model-400.pt\n\nlog:\n logdir: Re500-1_8s-pde2k-dat16-PINO\n entity: hzzheng-pino\n project: PINO-NS\n group: Re500-1_8s-pde2k-dat16-PINO\n" }, { "path": "configs/operator/Re500-1_8-dat400-FNO.yaml", "content": "data:\n name: KF\n paths: ['../data/NS-Re500_T300_id0.npy']\n Re: 500\n offset: 0\n total_num: 300\n raw_res: [256, 256, 513]\n n_data_samples: 50\n data_res: [64, 64, 129] # resolution in 1 second\n pde_res: [64, 64, 129] # resolution in 1 second\n a_offset: 0\n n_a_samples: 50\n testoffset: 250\n n_test_samples: 50\n t_duration: 0.125\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n act: gelu\n num_pad: 4\n pad_ratio: 0.0625\n\ntrain:\n batchsize: 2\n epochs: 201\n num_iter: 100_001\n milestones: [10_000, 30_000, 50_000, 70_000]\n base_lr: 0.001\n scheduler_gamma: 0.5\n ic_loss: 0.0\n f_loss: 0.0\n xy_loss: 1.0\n save_step: 5000\n eval_step: 5000\n\ntest:\n batchsize: 1\n data_res: [256, 256, 513]\n ckpt: model-400.pt\n\nlog:\n logdir: Re500-1_8s-dat400-FNO\n entity: hzzheng-pino\n project: PINO-NS\n group: Re500-1_8s-dat400-FNO\n" }, { "path": "configs/operator/Re500-1s-FNO.yaml", "content": "data:\n paths: ['../data/NS-Re500Part0.npy', '../data/NS-Re500Part1.npy', '../data/NS-Re500Part2.npy']\n Re: 500\n total_num: 300\n offset: 0\n n_samples: 200\n t_duration: 1.0\n data_res: [64, 64, 65]\n pde_res: [128, 128, 129]\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n act: gelu\n num_pad: 4\n\ntrain:\n batchsize: 2\n epochs: 501\n milestones: [200, 400]\n base_lr: 0.001\n scheduler_gamma: 0.5\n ic_loss: 0.0\n f_loss: 0.0\n xy_loss: 1.0\n save_step: 100\n\nlog:\n logdir: Re500-1s-200-FNO\n entity: hzzheng-pino\n project: 'PINO-NS'\n group: 'Re500-1s-200-FNO'\n" }, { "path": "configs/operator/Re500-3000-FNO.yaml", "content": "data:\n name: KF\n paths: ['../data/NS-Re500_T3000_id0.npy']\n Re: 500\n total_num: 3000\n offset: 0\n n_samples: 300\n testoffset: 2500\n n_test_samples: 500\n sub_x: 4\n sub_t: 4\n t_duration: 1.0\n data_res: [64, 64, 65]\n pde_res: [256, 256, 65]\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n act: gelu\n num_pad: 4\n\ntrain:\n batchsize: 4\n epochs: 401\n milestones: [200]\n base_lr: 0.001\n scheduler_gamma: 0.5\n ic_loss: 0.0\n f_loss: 0.0\n xy_loss: 1.0\n save_step: 50\n\nlog:\n logdir: Re500-1s-3000-FNO\n entity: hzzheng-pino\n project: PINO-NS\n group: Re500-1s-3000-FNO\n" }, { "path": "configs/operator/Re500-3000-PINO.yaml", "content": "data:\n name: KF\n paths: ['../data/NS-Re500_T3000_id0.npy']\n Re: 500\n total_num: 3000\n offset: 0\n n_samples: 2400\n sub_x: 4\n sub_t: 4\n pde_subx: 1\n pde_subt: 2\n t_duration: 1.0\n data_res: [64, 64, 65]\n pde_res: [256, 256, 129]\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n act: gelu\n num_pad: 4\n\ntrain:\n batchsize: 1\n epochs: 401\n milestones: [200]\n base_lr: 0.001\n scheduler_gamma: 0.5\n ic_loss: 1.0\n f_loss: 1.0\n xy_loss: 5.0\n save_step: 50\n\nlog:\n logdir: Re500-1s-3000-PINO\n entity: hzzheng-pino\n project: PINO-NS\n group: Re500-1s-3000-PINO\n" }, { "path": "configs/operator/Re500-4000-FNO.yaml", "content": "data:\n paths: ['../data/NS-T4000.npy']\n Re: 500\n total_num: 4000\n offset: 0\n n_samples: 3200\n t_duration: 1.0\n data_res: [64, 64, 65]\n pde_res: [128, 128, 65]\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n act: gelu\n num_pad: 4\n\ntrain:\n batchsize: 2\n epochs: 501\n milestones: [300]\n base_lr: 0.001\n scheduler_gamma: 0.5\n ic_loss: 0.0\n f_loss: 0.0\n xy_loss: 1.0\n save_step: 100\n\nlog:\n logdir: Re500-1s-FNO\n entity: hzzheng-pino\n project: 'PINO-NS'\n group: 'Re500-1s-FNO'\n" }, { "path": "configs/operator/Re500-FNO.yaml", "content": "data:\n paths: ['../data/NS-Re500Part0.npy', '../data/NS-Re500Part1.npy']\n Re: 500\n total_num: 200\n offset: 0\n n_samples: 700\n t_duration: 0.5\n data_res: [64, 64, 33]\n pde_res: [128, 128, 65]\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n act: gelu\n num_pad: 4\n pad_ratio: 0.03125\n\ntrain:\n batchsize: 2\n epochs: 501\n milestones: [300]\n base_lr: 0.001\n scheduler_gamma: 0.5\n ic_loss: 0.0\n f_loss: 0.0\n xy_loss: 1.0\n save_step: 100\n\ntest:\n batchsize: 1\n data_res: [128, 128, 65]\n ckpt: model-500.pt\n\nlog:\n logdir: Re500-05s-FNO\n entity: hzzheng-pino\n project: 'PINO-NS'\n group: 'Re500-05s-FNO'\n" }, { "path": "configs/operator/Re500-PINO.yaml", "content": "data:\n paths: ['../data/NS-Re500Part0.npy', '../data/NS-Re500Part1.npy']\n Re: 500\n total_num: 200\n offset: 0\n n_samples: 700\n t_duration: 0.5\n data_res: [64, 64, 33]\n pde_res: [128, 128, 65]\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n act: gelu\n num_pad: 4\n\ntrain:\n batchsize: 2\n epochs: 501\n milestones: [300]\n base_lr: 0.001\n scheduler_gamma: 0.5\n ic_loss: 1.0\n f_loss: 1.0\n xy_loss: 5.0\n save_step: 100\n\ntest:\n batchsize: 1\n data_res: [64, 64, 33]\n ckpt: model-500.pt\n\nlog:\n logdir: Re500-05s-PINO\n entity: hzzheng-pino\n project: 'PINO-NS'\n group: 'Re500-05s-PINO'\n\n\n" }, { "path": "configs/pretrain/Darcy-pretrain-deeponet.yaml", "content": "data:\n name: 'Darcy'\n datapath: '/mnt/md1/zongyi/piececonst_r421_N1024_smooth1.mat'\n total_num: 1024\n offset: 0\n n_sample: 1000\n nx: 421\n sub: 7\n\nmodel:\n branch_layers: [50, 50, 50, 50, 50, 50, 50, 50, 50, 50, 50, 50]\n trunk_layers: [50, 50, 50, 50, 50, 50, 50, 50, 50, 50, 50, 50]\n activation: tanh\n\ntrain:\n batchsize: 20\n epochs: 2000\n milestones: [400, 800, 1200]\n base_lr: 0.001\n scheduler_gamma: 0.5\n save_dir: 'darcy-deeponet'\n save_name: 'darcy-pretrain-deeponet.pt'\n\nlog:\n project: 'PINO-Darcy-pretrain'\n group: 'deeponet'\n\n\n" }, { "path": "configs/pretrain/Darcy-pretrain.yaml", "content": "data:\n name: 'Darcy'\n datapath: '/mnt/md1/zongyi/piececonst_r421_N1024_smooth1.mat'\n total_num: 1024\n offset: 0\n n_sample: 1000\n nx: 421\n sub: 7\n pde_sub: 2\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [20, 20, 20, 20]\n modes2: [20, 20, 20, 20]\n fc_dim: 128\n act: gelu\n\ntrain:\n batchsize: 20\n epochs: 300\n milestones: [100, 150, 200]\n base_lr: 0.001\n scheduler_gamma: 0.5\n f_loss: 1.0\n xy_loss: 5.0\n save_dir: 'darcy-FDM'\n save_name: 'darcy-pretrain-pino.pt'\n\nlog:\n project: 'PINO-Darcy-pretrain'\n group: 'gelu-pino'\n entity: hzzheng-pino\n\n\n" }, { "path": "configs/pretrain/Re100-pretrain-1s.yaml", "content": "data:\n datapath: 'data/NS_fine_Re100_T128_part0.npy'\n datapath2: 'data/NS_fine_Re100_T128_part1.npy'\n Re: 100\n total_num: 100\n offset: 0\n n_sample: 200\n time_interval: 1.0\n nx: 128\n nt: 128\n sub: 1\n sub_t: 2\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 150\n milestones: [25, 50, 75, 100]\n base_lr: 0.001\n scheduler_gamma: 0.5\n ic_loss: 1.0\n f_loss: 1.0\n xy_loss: 5.0\n save_dir: 'Re100-FDM'\n save_name: 'PINO-pretrain-Re100-1s.pt'\n\nlog:\n project: 'PINO-pretrain'\n group: 'Re100-1s-tanh'\n\n\n" }, { "path": "configs/pretrain/Re200-pretrain-1s.yaml", "content": "data:\n datapath: 'data/NS_fine_Re200_T128_part0.npy'\n datapath2: 'data/NS_fine_Re200_T128_part1.npy'\n Re: 200\n total_num: 100\n offset: 0\n n_sample: 200\n time_interval: 1.0\n nx: 128\n nt: 128\n sub: 1\n sub_t: 2\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 150\n milestones: [25, 50, 75, 100]\n base_lr: 0.001\n scheduler_gamma: 0.5\n ic_loss: 1.0\n f_loss: 1.0\n xy_loss: 5.0\n save_dir: 'Re200-FDM'\n save_name: 'PINO-pretrain-Re200-1s.pt'\n\nlog:\n project: 'PINO-pretrain'\n group: 'Re200-1s-tanh'\n\n\n" }, { "path": "configs/pretrain/Re250-pretrain-1s.yaml", "content": "data:\n datapath: 'data/NS_fine_Re250_T128_part0.npy'\n datapath2: 'data/NS_fine_Re250_T128_part1.npy'\n Re: 250\n total_num: 100\n offset: 0\n n_sample: 200\n time_interval: 1.0\n nx: 128\n nt: 128\n sub: 1\n sub_t: 2\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 150\n milestones: [25, 50, 75, 100]\n base_lr: 0.001\n scheduler_gamma: 0.5\n ic_loss: 1.0\n f_loss: 1.0\n xy_loss: 5.0\n save_dir: 'Re250-FDM'\n save_name: 'PINO-pretrain-Re250-1s.pt'\n\nlog:\n project: 'PINO-pretrain'\n group: 'Re250-1s-tanh'\n\n\n" }, { "path": "configs/pretrain/Re300-pretrain-1s.yaml", "content": "data:\n datapath: 'data/NS_fine_Re300_T128_part0.npy'\n datapath2: 'data/NS_fine_Re300_T128_part1.npy'\n Re: 300\n total_num: 100\n offset: 0\n n_sample: 200\n time_interval: 1.0\n nx: 128\n nt: 128\n sub: 1\n sub_t: 2\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 150\n milestones: [25, 50, 75, 100]\n base_lr: 0.001\n scheduler_gamma: 0.5\n ic_loss: 1.0\n f_loss: 1.0\n xy_loss: 5.0\n save_dir: 'Re300-FDM'\n save_name: 'PINO-pretrain-Re300-1s.pt'\n\nlog:\n project: 'PINO-pretrain'\n group: 'Re300-1s-tanh'\n\n\n" }, { "path": "configs/pretrain/Re350-pretrain-1s.yaml", "content": "data:\n datapath: 'data/NS_fine_Re350_T128_part0.npy'\n datapath2: 'data/NS_fine_Re350_T128_part1.npy'\n Re: 350\n total_num: 100\n offset: 0\n n_sample: 200\n time_interval: 1.0\n nx: 128\n nt: 128\n sub: 1\n sub_t: 2\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 150\n milestones: [25, 50, 75, 100]\n base_lr: 0.001\n scheduler_gamma: 0.5\n ic_loss: 1.0\n f_loss: 1.0\n xy_loss: 5.0\n save_dir: 'Re350-FDM'\n save_name: 'PINO-pretrain-Re350-1s.pt'\n\nlog:\n project: 'PINO-pretrain'\n group: 'Re350-1s-tanh'\n\n\n" }, { "path": "configs/pretrain/Re400-pretrain-1s.yaml", "content": "data:\n datapath: 'data/NS_fine_Re400_T128_part0.npy'\n datapath2: 'data/NS_fine_Re400_T128_part1.npy'\n Re: 400\n total_num: 100\n offset: 0\n n_sample: 200\n time_interval: 1.0\n nx: 128\n nt: 128\n sub: 1\n sub_t: 2\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 150\n milestones: [25, 50, 75, 100]\n base_lr: 0.001\n scheduler_gamma: 0.5\n ic_loss: 1.0\n f_loss: 1.0\n xy_loss: 5.0\n save_dir: 'Re400-FDM'\n save_name: 'PINO-pretrain-Re400-1s.pt'\n\nlog:\n project: 'PINO-pretrain'\n group: 'Re400-1s-tanh'\n\n\n" }, { "path": "configs/pretrain/Re500-05s-deeponet.yaml", "content": "data:\n datapath: '/mnt/md1/zongyi/NS_fft_Re500_T4000.npy'\n Re: 500\n total_num: 4000\n offset: 0\n n_sample: 400\n time_interval: 0.5\n nx: 64\n nt: 64\n sub: 1\n sub_t: 1\n shuffle: True\n\nmodel:\n branch_layers: [100, 100, 100]\n trunk_layers: [100, 100, 100]\n\ntrain:\n batchsize: 20\n epochs: 10001\n milestones: [2500, 5000, 7500]\n base_lr: 0.001\n scheduler_gamma: 0.5\n save_dir: 'Re500-deepOnet'\n save_name: 'DeepONet-pretrain-Re500.pt'\n\nlog:\n project: 'PINO-pretrain-ICLR'\n group: 'Re500-05s-deepONet'\n\n\n" }, { "path": "configs/pretrain/Re500-FNO-1s-100.yaml", "content": "data:\n datapath: '../data/NS-T4000.npy'\n Re: 500\n total_num: 4000\n offset: 0\n n_sample: 100\n time_interval: 1\n nx: 64\n nt: 64\n sub: 1\n sub_t: 1\n shuffle: True\n S2: 64\n T2: 65\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n act: gelu\n\ntrain:\n batchsize: 2\n epochs: 100_001\n milestones: [20_000, 40_000, 60_000, 80_000]\n base_lr: 0.001\n scheduler_gamma: 0.5\n ic_loss: 0.0\n f_loss: 0.0\n xy_loss: 1.0\n save_dir: Re500-FNO-100\n save_name: FNO-Re500-1s-100.pt\n data_iter: 1\n eqn_iter: 0\n\nlog:\n entity: hzzheng-pino\n project: PINO-Operator-Learning\n group: FNO-Re500-1s-100\n\n\n" }, { "path": "configs/pretrain/Re500-FNO-1s-200.yaml", "content": "data:\n datapath: '../data/NS-T4000.npy'\n Re: 500\n total_num: 4000\n offset: 0\n n_sample: 200\n time_interval: 1\n nx: 64\n nt: 64\n sub: 1\n sub_t: 1\n shuffle: True\n S2: 64\n T2: 65\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n act: gelu\n\ntrain:\n batchsize: 2\n epochs: 100_001\n milestones: [20_000, 40_000, 60_000, 80_000]\n base_lr: 0.001\n scheduler_gamma: 0.5\n ic_loss: 0.0\n f_loss: 0.0\n xy_loss: 1.0\n save_dir: Re500-FNO-200\n save_name: FNO-Re500-1s-200.pt\n data_iter: 1\n eqn_iter: 0\n\nlog:\n entity: hzzheng-pino\n project: PINO-Operator-Learning\n group: FNO-Re500-1s-200\n\n\n" }, { "path": "configs/pretrain/Re500-FNO-1s-400.yaml", "content": "data:\n datapath: '../data/NS-T4000.npy'\n Re: 500\n total_num: 4000\n offset: 0\n n_sample: 400\n time_interval: 1\n nx: 64\n nt: 64\n sub: 1\n sub_t: 1\n shuffle: True\n S2: 64\n T2: 65\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n act: gelu\n\ntrain:\n batchsize: 2\n epochs: 100_001\n milestones: [20_000, 40_000, 60_000, 80_000]\n base_lr: 0.001\n scheduler_gamma: 0.5\n ic_loss: 0.0\n f_loss: 0.0\n xy_loss: 1.0\n save_dir: Re500-FNO-400\n save_name: FNO-Re500-1s-400.pt\n data_iter: 1\n eqn_iter: 0\n\nlog:\n entity: hzzheng-pino\n project: PINO-Operator-Learning\n group: FNO-Re500-1s-400\n\n\n" }, { "path": "configs/pretrain/Re500-PINO-1s-100-4v4.yaml", "content": "data:\n datapath: '../data/NS-T4000.npy'\n Re: 500\n total_num: 4000\n offset: 0\n n_sample: 100\n time_interval: 1\n nx: 64\n nt: 64\n sub: 1\n sub_t: 1\n shuffle: True\n S2: 64\n T2: 65\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n act: gelu \n\ntrain:\n batchsize: 2\n epochs: 40_000\n milestones: [10_000, 20_000, 30_000]\n base_lr: 0.001\n scheduler_gamma: 0.5\n ic_loss: 1.0\n f_loss: 1.0\n xy_loss: 5.0\n save_dir: Re500-PINO-1s-100\n save_name: PINO-pretrain-Re500-1s-100.pt\n data_iter: 4\n eqn_iter: 4\n\nlog:\n entity: hzzheng-pino\n project: PINO-Operator-Learning\n group: PINO-Re500-1s-100-4v4\n\n\n" }, { "path": "configs/pretrain/Re500-PINO-1s-200-4v4.yaml", "content": "data:\n datapath: '../data/NS-T4000.npy'\n Re: 500\n total_num: 4000\n offset: 0\n n_sample: 200\n time_interval: 1\n nx: 64\n nt: 64\n sub: 1\n sub_t: 1\n shuffle: True\n S2: 64\n T2: 65\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n act: gelu \n\ntrain:\n batchsize: 2\n epochs: 100_000\n milestones: [20_000, 40_000, 60_000, 80_000]\n base_lr: 0.001\n scheduler_gamma: 0.5\n ic_loss: 1.0\n f_loss: 1.0\n xy_loss: 5.0\n save_dir: Re500-PINO-1s-200\n save_name: PINO-pretrain-Re500-1s-200.pt\n data_iter: 4\n eqn_iter: 4\n\nlog:\n entity: hzzheng-pino\n project: PINO-Operator-Learning\n group: PINO-Re500-1s-200-4v4\n\n\n" }, { "path": "configs/pretrain/Re500-PINO-1s-400-1v1.yaml", "content": "data:\n datapath: '../data/NS-T4000.npy'\n Re: 500\n total_num: 4000\n offset: 0\n n_sample: 400\n time_interval: 1\n nx: 64\n nt: 64\n sub: 1\n sub_t: 1\n shuffle: True\n S2: 64\n T2: 65\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n act: gelu \n\ntrain:\n batchsize: 2\n epochs: 40_000\n milestones: [10_000, 20_000, 30_000]\n base_lr: 0.001\n scheduler_gamma: 0.5\n ic_loss: 1.0\n f_loss: 1.0\n xy_loss: 5.0\n save_dir: Re500-PINO-1s-400\n save_name: PINO-pretrain-Re500-1s-400.pt\n data_iter: 4\n eqn_iter: 4\n\nlog:\n entity: hzzheng-pino\n project: PINO-Operator-Learning\n group: PINO-Re500-1s-400-4v4\n\n\n" }, { "path": "configs/pretrain/Re500-pretrain-05s-4C1.yaml", "content": "data:\n datapath: '/mnt/md1/zongyi/NS_fft_Re500_T4000.npy'\n Re: 500\n total_num: 4000\n offset: 0\n n_sample: 400\n time_interval: 0.5\n nx: 64\n nt: 64\n sub: 1\n sub_t: 1\n shuffle: True\n S2: 128\n T2: 65\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 40000\n milestones: [10000, 20000, 30000]\n base_lr: 0.001\n scheduler_gamma: 0.5\n ic_loss: 1.0\n f_loss: 1.0\n xy_loss: 5.0\n save_dir: 'Re500-FDM'\n save_name: 'PINO-pretrain-Re500-05s-4C1.pt'\n data_iter: 4 # number of update steps on data for each epoch\n eqn_iter: 1 # number of update steps on virtual PDE for each epoch\n\nlog:\n project: 'PINO-pretrain-ICLR'\n group: 'Re500-05s-4C1'\n\n\n" }, { "path": "configs/pretrain/Re500-pretrain-05s-4C4.yaml", "content": "data:\n datapath: '/mnt/md1/zongyi/NS_fft_Re500_T4000.npy'\n Re: 500\n total_num: 4000\n offset: 0\n n_sample: 400\n time_interval: 0.5\n nx: 64\n nt: 64\n sub: 1\n sub_t: 1\n shuffle: True\n num_ics: 300\n S2: 128\n T2: 65\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 40000\n milestones: [10000, 20000, 30000]\n base_lr: 0.001\n scheduler_gamma: 0.5\n ic_loss: 1.0\n f_loss: 1.0\n xy_loss: 5.0\n save_dir: 'Re500-FDM'\n save_name: 'PINO-pretrain-Re500-05s-4C4.pt'\n data_iter: 4\n eqn_iter: 4\n\nlog:\n project: 'PINO-pretrain-ICLR'\n group: 'Re500-05s-4C4'\n\n\n" }, { "path": "configs/pretrain/Re500-pretrain-05s-eqn.yaml", "content": "data:\n datapath: '/mnt/md1/zongyi/NS_fft_Re500_T4000.npy'\n Re: 500\n total_num: 4000\n offset: 0\n n_sample: 4000\n time_interval: 0.5\n nx: 64\n nt: 64\n sub: 1\n sub_t: 1\n shuffle: True\n num_ics: 300\n S2: 128\n T2: 65\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 40000\n milestones: [10000, 20000, 30000]\n base_lr: 0.001\n scheduler_gamma: 0.5\n ic_loss: 1.0\n f_loss: 1.0\n xy_loss: 0.0\n save_dir: 'Re500-FDM'\n save_name: 'PINO-pretrain-Re500-05s-eqn.pt'\n data_iter: 0\n eqn_iter: 1\n\nlog:\n project: 'PINO-pretrain-ICLR'\n group: 'Re500-05s-0C1'\n\n\n" }, { "path": "configs/pretrain/Re500-pretrain-1s.yaml", "content": "data:\n datapath: 'data/NS_fine_Re500_T128_part0.npy'\n datapath2: 'data/NS_fine_Re500_T128_part1.npy'\n Re: 500\n total_num: 100\n offset: 0\n n_sample: 200\n time_interval: 1\n nx: 128\n nt: 128\n sub: 1\n sub_t: 1\n shuffle: True\n num_ics: 200\n S2: 128\n T2: 128\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 150\n milestones: [25, 50, 75, 100]\n base_lr: 0.001\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0.0\n save_dir: 'Re500-FDM'\n save_name: 'PINO-pretrain-Re500-1s-eqn256.pt'\n\nlog:\n project: 'PINO-pretrain'\n group: 'Re500-1s-eqn'\n\n\n" }, { "path": "configs/pretrain/burgers-pretrain.yaml", "content": "data:\n name: Burgers\n datapath: '../data/burgers.mat'\n total_num: 1000\n offset: 0\n n_sample: 800\n nx: 128\n nt: 100\n sub: 1\n sub_t: 1\n\nmodel:\n layers: [16, 24, 24, 32, 32]\n modes1: [15, 12, 9, 9]\n modes2: [15, 12, 9, 9]\n fc_dim: 128\n act: gelu\n num_pad: 4\n\ntrain:\n batchsize: 20\n epochs: 500\n milestones: [150, 300, 450]\n base_lr: 0.001\n scheduler_gamma: 0.5\n ic_loss: 10.0\n f_loss: 1.0\n xy_loss: 0.0\n save_dir: 'burgers-FDM'\n save_name: 'burgers-pretrain-eqn.pt'\n\nlog:\n project: PINO-burgers-pretrain\n group: gelu-eqn\n entity: hzzheng-pino\n\n" }, { "path": "configs/scratch/Re100-scratch-1s.yaml", "content": "data:\n datapath: 'data/NS_fine_Re100_T128_part2.npy'\n Re: 100\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 1.0\n nx: 128\n nt: 128\n sub: 1\n sub_t: 2\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 8000\n milestones: [1000, 2000, 3000, 4000, 5000, 6000, 7000]\n base_lr: 0.0025\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0.0\n save_dir: 'Re100-FDM'\n save_name: 'PINO-scratch-Re100-1s.pt'\n\nlog:\n project: 'PINO-scratch-tanh'\n group: 'Re100-scratch-1s'\n\n\n" }, { "path": "configs/scratch/Re200-scratch-1s.yaml", "content": "data:\n datapath: 'data/NS_fine_Re200_T128_part2.npy'\n Re: 200\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 1.0\n nx: 128\n nt: 128\n sub: 1\n sub_t: 2\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 8000\n milestones: [1000, 2000, 3000, 4000, 5000, 6000, 7000]\n base_lr: 0.0025\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0.0\n save_dir: 'Re100-FDM'\n save_name: 'PINO-scratch-Re200-1s.pt'\n\nlog:\n project: 'PINO-scratch-tanh'\n group: 'Re200-scratch-1s'\n\n\n" }, { "path": "configs/scratch/Re250-scratch-1s.yaml", "content": "data:\n datapath: 'data/NS_fine_Re250_T128_part2.npy'\n Re: 250\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 1.0\n nx: 128\n nt: 128\n sub: 1\n sub_t: 2\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 8000\n milestones: [1000, 2000, 3000, 4000, 5000, 6000, 7000]\n base_lr: 0.0025\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0.0\n save_dir: 'Re250-FDM'\n save_name: 'PINO-scratch-Re250-1s.pt'\n\nlog:\n project: 'PINO-scratch-tanh'\n group: 'Re250-scratch-1s'\n\n\n" }, { "path": "configs/scratch/Re300-scratch-1s.yaml", "content": "data:\n datapath: 'data/NS_fine_Re300_T128_part2.npy'\n Re: 300\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 1.0\n nx: 128\n nt: 128\n sub: 1\n sub_t: 2\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 8000\n milestones: [1000, 2000, 3000, 4000, 5000, 6000, 7000]\n base_lr: 0.0025\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0.0\n save_dir: 'Re300-FDM'\n save_name: 'PINO-scratch-Re300-1s.pt'\n\nlog:\n project: 'PINO-scratch-tanh'\n group: 'Re300-scratch-1s'\n\n\n" }, { "path": "configs/scratch/Re350-scratch-1s.yaml", "content": "data:\n datapath: 'data/NS_fine_Re350_T128_part0.npy'\n Re: 350\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 1.0\n nx: 128\n nt: 128\n sub: 1\n sub_t: 2\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 8000\n milestones: [1000, 2000, 3000, 4000, 5000, 6000, 7000]\n base_lr: 0.0025\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0.0\n save_dir: 'Re350-FDM'\n save_name: 'PINO-scratch-Re350-1s.pt'\n\nlog:\n project: 'PINO-scratch-tanh'\n group: 'Re350-scratch-1s'\n\n\n" }, { "path": "configs/scratch/Re400-scratch-1s.yaml", "content": "data:\n datapath: 'data/NS_fine_Re400_T128_part0.npy'\n Re: 400\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 1.0\n nx: 128\n nt: 128\n sub: 1\n sub_t: 2\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 8000\n milestones: [1000, 2000, 3000, 4000, 5000, 6000, 7000]\n base_lr: 0.0025\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0.0\n save_dir: 'Re400-FDM'\n save_name: 'PINO-scratch-Re400-1s.pt'\n\nlog:\n project: 'PINO-scratch-tanh'\n group: 'Re400-scratch-1s'\n\n\n" }, { "path": "configs/scratch/Re500-scratch-05s-new.yaml", "content": "data:\n datapath: 'data/NS_fine_Re500_T128_part2.npy'\n Re: 500\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 0.5\n nx: 128\n nt: 128\n sub: 1\n sub_t: 1\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 6000\n milestones: [1000, 2000, 3000, 4000, 5000]\n base_lr: 0.0025\n beta1: 0.9\n beta2: 0.999\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0\n save_dir: 'Re500-FDM'\n save_name: 'PINO-scratch-05s.pt'\n\nlog:\n project: 'PINO-Re500-exp'\n group: 'Re500-scratch-128'\n\n\n" }, { "path": "configs/scratch/Re500-scratch-05s.yaml", "content": "data:\n datapath: 'data/NS_Re500_s256_T100_test.npy'\n Re: 500\n total_num: 100\n offset: 0\n n_sample: 20\n time_interval: 0.5\n nx: 256\n nt: 128\n sub: 4\n sub_t: 2\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 2500\n milestones: [1000, 1500, 2000]\n base_lr: 0.0025\n beta1: 0.9\n beta2: 0.999\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0\n save_dir: 'Re500-FDM'\n save_name: 'PINO-scratch-05s.pt'\n\nlog:\n project: 'PINO-Re500-ICLR'\n group: 'Re500-scratch-128'\n logfile: 'log/pinns-default.csv'\n\n\n\n" }, { "path": "configs/scratch/Re500-scratch-1s-progressive.yaml", "content": "data:\n datapath: 'data/NS_fine_Re500_T128_part2.npy'\n Re: 500\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 1\n nx: 128\n nt: 128\n sub: 1\n sub_t: 2\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: [200, 5800]\n milestones: [1000, 2000, 3000, 4000, 5000, 6000]\n base_lr: 0.0025\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0\n save_dir: 'Re500-FDM'\n save_name: 'PINO-scratch128-1s.pt'\n\nlog:\n project: 'PINO-default'\n group: 'Re500-scratch-1s'\n\n\n" }, { "path": "configs/scratch/Re500-scratch-1s.yaml", "content": "data:\n datapath: '../data/NS-Re500Part1.npy'\n Re: 500\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 1\n nx: 128\n nt: 128\n sub: 1\n sub_t: 2\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 8000\n milestones: [1000, 2000, 3000, 4000, 5000, 6000, 7000]\n base_lr: 0.0025\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0\n save_dir: 'Re500-FDM'\n save_name: 'PINO-scratch128-1s.pt'\n\nlog:\n entity: 'hzzheng-pino'\n project: 'PINO-NavierStokes'\n group: 'Re500-scratch-1s'\n\n\n" }, { "path": "configs/test/Re500-05s-deeponet.yaml", "content": "data:\n datapath: 'data/NS_fine_Re500_T128_part2.npy'\n Re: 500\n total_num: 100\n offset: 0\n n_sample: 300\n time_interval: 0.5\n nx: 128\n nt: 128\n sub: 2\n sub_t: 2\n shuffle: False\n\nmodel:\n branch_layers: [100, 100, 100]\n trunk_layers: [100, 100, 100]\n\ntest:\n batchsize: 20\n ckpt: 'checkpoints/Re500-deepOnet/DeepONet-pretrain-Re500_10000.pt'\n\nlog:\n project: 'PINO-None'\n group: 'eval'\n\n\n" }, { "path": "configs/test/Re500-05s-test.yaml", "content": "data:\n datapath: 'data/NS_Re500_s256_T100_test.npy'\n Re: 500\n total_num: 100\n offset: 0\n n_sample: 20\n time_interval: 0.5\n nx: 256\n nt: 128\n sub: 4\n sub_t: 1\n shuffle: False\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntest:\n batchsize: 1\n ckpt: 'checkpoints/Re500-FDM/PINO-pretrain-Re500-05s-4k1k.pt'\n\nlog:\n project: 'PINO-None'\n group: 'eval'\n\n\n" }, { "path": "configs/test/Re500-05s.yaml", "content": "data:\n datapath: 'data/NS_fine_Re500_T128_part2.npy'\n Re: 500\n total_num: 100\n offset: 0\n n_sample: 300\n time_interval: 0.5\n nx: 128\n nt: 128\n sub: 1\n sub_t: 1\n shuffle: False\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntest:\n batchsize: 1\n ckpt: 'checkpoints/Re500-FDM/PINO-pretrain-Re500-05s-4C1.pt'\n\nlog:\n project: 'PINO-None'\n group: 'eval'\n\n\n" }, { "path": "configs/test/Re500-1s-100.yaml", "content": "data:\n datapath: '../data/NS-T4000.npy'\n Re: 500\n total_num: 100\n offset: 0\n n_sample: 100\n time_interval: 1\n nx: 64\n nt: 64\n sub: 1\n sub_t: 1\n shuffle: False\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n act: gelu\n\ntest:\n batchsize: 1\n ckpt: checkpoints/Re500-FNO-100/FNO-Re500-1s-100.pt\n\nlog:\n entity: hzzheng-pino\n project: PINO-NS\n group: eval\n\n\n" }, { "path": "configs/test/burgers.yaml", "content": "data:\n name: 'Darcy'\n datapath: '../data/burgers.mat'\n total_num: 1000\n offset: 800\n n_sample: 200\n nx: 128\n nt: 100\n sub: 1\n sub_t: 1\n\nmodel:\n layers: [16, 24, 24, 32, 32]\n modes1: [15, 12, 9, 9]\n modes2: [15, 12, 9, 9]\n fc_dim: 128\n act: gelu\n\ntest:\n batchsize: 1\n ckpt: 'checkpoints/burgers-FDM/burgers-pretrain-eqn.pt'\n\nlog:\n project: 'PINO-burgers-test'\n group: 'gelu-test'\n\n\n" }, { "path": "configs/test/darcy-deeponet.yaml", "content": "data:\n name: 'Darcy'\n datapath: '/mnt/md1/zongyi/piececonst_r421_N1024_smooth2.mat'\n total_num: 1000\n offset: 0\n n_sample: 500\n nx: 421\n sub: 7\n shuffle: False\n\nmodel:\n branch_layers: [50, 50, 50, 50, 50, 50, 50, 50, 50, 50, 50, 50]\n trunk_layers: [50, 50, 50, 50, 50, 50, 50, 50, 50, 50, 50, 50]\n activation: tanh\n\ntest:\n batchsize: 1\n ckpt: 'checkpoints/darcy-deeponet/darcy-pretrain-deeponet.pt'\n\nlog:\n project: 'PINO-Darcy'\n group: 'default'\n\n\n" }, { "path": "configs/test/darcy.yaml", "content": "data:\n name: 'Darcy'\n datapath: '/mnt/md1/zongyi/piececonst_r421_N1024_smooth2.mat'\n total_num: 1000\n offset: 0\n n_sample: 500\n nx: 421\n sub: 7\n shuffle: False\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [20, 20, 20, 20]\n modes2: [20, 20, 20, 20]\n fc_dim: 128\n act: gelu\n\ntest:\n batchsize: 1\n ckpt: 'checkpoints/darcy-FDM/darcy-pretrain-eqn.pt'\n\nlog:\n project: 'PINO-Darcy'\n group: 'default'\n\n\n" }, { "path": "configs/transfer/Re100to100-1s.yaml", "content": "data:\n datapath: 'data/NS_fine_Re100_T128_part2.npy'\n Re: 100\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 1.0\n nx: 128\n nt: 128\n sub: 1\n sub_t: 2\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 8000\n milestones: [1000, 2000, 3000, 4000, 5000, 6000, 7000]\n base_lr: 0.0025\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0.0\n save_dir: 'Re100-FDM'\n save_name: 'PINO-transfer-Re100-1s.pt'\n ckpt: 'checkpoints/Re100-FDM/PINO-pretrain-Re100-1s.pt'\n\nlog:\n project: 'PINO-transfer-tanh'\n group: 'Re100to100-1s'\n\n\n" }, { "path": "configs/transfer/Re100to200-1s.yaml", "content": "data:\n datapath: 'data/NS_fine_Re200_T128_part2.npy'\n Re: 200\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 1.0\n nx: 128\n nt: 128\n sub: 1\n sub_t: 2\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 8000\n milestones: [1000, 2000, 3000, 4000, 5000, 6000, 7000]\n base_lr: 0.0025\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0.0\n save_dir: 'Re200-FDM'\n save_name: 'PINO-transfer-Re200-1s.pt'\n ckpt: 'checkpoints/Re100-FDM/PINO-pretrain-Re100-1s.pt'\n\nlog:\n project: 'PINO-transfer-tanh'\n group: 'Re100to200-1s'\n\n\n" }, { "path": "configs/transfer/Re100to250-1s.yaml", "content": "data:\n datapath: 'data/NS_fine_Re250_T128_part2.npy'\n Re: 250\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 1.0\n nx: 128\n nt: 128\n sub: 1\n sub_t: 2\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 8000\n milestones: [1000, 2000, 3000, 4000, 5000, 6000, 7000]\n base_lr: 0.0025\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0.0\n save_dir: 'Re250-FDM'\n save_name: 'PINO-transfer-Re250-1s.pt'\n ckpt: 'checkpoints/Re100-FDM/PINO-pretrain-Re100-1s.pt'\n\nlog:\n project: 'PINO-transfer-tanh'\n group: 'Re100to250-1s'\n\n\n" }, { "path": "configs/transfer/Re100to300-1s.yaml", "content": "data:\n datapath: 'data/NS_fine_Re300_T128_part2.npy'\n Re: 300\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 1.0\n nx: 128\n nt: 128\n sub: 1\n sub_t: 2\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 8000\n milestones: [1000, 2000, 3000, 4000, 5000, 6000, 7000]\n base_lr: 0.0025\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0.0\n save_dir: 'Re300-FDM'\n save_name: 'PINO-transfer-Re300-1s.pt'\n ckpt: 'checkpoints/Re100-FDM/PINO-pretrain-Re100-1s.pt'\n\nlog:\n project: 'PINO-transfer-tanh'\n group: 'Re100to300-1s'\n\n\n" }, { "path": "configs/transfer/Re100to350-1s.yaml", "content": "data:\n datapath: 'data/NS_fine_Re350_T128_part2.npy'\n Re: 350\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 1.0\n nx: 128\n nt: 128\n sub: 1\n sub_t: 2\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 8000\n milestones: [1000, 2000, 3000, 4000, 5000, 6000, 7000]\n base_lr: 0.0025\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0.0\n save_dir: 'Re350-FDM'\n save_name: 'PINO-transfer-Re350-1s.pt'\n ckpt: 'checkpoints/Re100-FDM/PINO-pretrain-Re100-1s.pt'\n\nlog:\n project: 'PINO-transfer-tanh'\n group: 'Re100to350-1s'\n\n\n" }, { "path": "configs/transfer/Re100to400-1s.yaml", "content": "data:\n datapath: 'data/NS_fine_Re400_T128_part2.npy'\n Re: 400\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 1.0\n nx: 128\n nt: 128\n sub: 1\n sub_t: 2\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 8000\n milestones: [1000, 2000, 3000, 4000, 5000, 6000, 7000]\n base_lr: 0.0025\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0.0\n save_dir: 'Re400-FDM'\n save_name: 'PINO-transfer-Re400-1s.pt'\n ckpt: 'checkpoints/Re100-FDM/PINO-pretrain-Re100-1s.pt'\n\nlog:\n project: 'PINO-transfer-tanh'\n group: 'Re100to400-1s'\n\n\n" }, { "path": "configs/transfer/Re100to500-1s.yaml", "content": "data:\n datapath: 'data/NS_fine_Re500_T128_part2.npy'\n Re: 500\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 1\n nx: 128\n nt: 128\n sub: 1\n sub_t: 2\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 8000\n milestones: [1000, 2000, 3000, 4000, 5000, 6000, 7000]\n base_lr: 0.0025\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0\n save_dir: 'Re500-FDM'\n save_name: 'PINO-Re500-1s.pt'\n ckpt: 'checkpoints/Re100-FDM/PINO-pretrain-Re100-1s.pt'\n\nlog:\n project: 'PINO-transfer-tanh'\n group: 'Re100to500-1s'\n\n\n\n\n" }, { "path": "configs/transfer/Re200to100-1s.yaml", "content": "data:\n datapath: 'data/NS_fine_Re100_T128_part2.npy'\n Re: 100\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 1.0\n nx: 128\n nt: 128\n sub: 1\n sub_t: 2\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 8000\n milestones: [1000, 2000, 3000, 4000, 5000, 6000, 7000]\n base_lr: 0.0025\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0.0\n save_dir: 'Re100-FDM'\n save_name: 'PINO-transfer-Re100-1s.pt'\n ckpt: 'checkpoints/Re200-FDM/PINO-pretrain-Re200-1s.pt'\n\nlog:\n project: 'PINO-transfer-tanh'\n group: 'Re200to100-1s'\n\n\n" }, { "path": "configs/transfer/Re200to200-1s.yaml", "content": "data:\n datapath: 'data/NS_fine_Re200_T128_part2.npy'\n Re: 200\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 1.0\n nx: 128\n nt: 128\n sub: 1\n sub_t: 2\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 8000\n milestones: [1000, 2000, 3000, 4000, 5000, 6000, 7000]\n base_lr: 0.0025\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0.0\n save_dir: 'Re200-FDM'\n save_name: 'PINO-transfer-Re200-1s.pt'\n ckpt: 'checkpoints/Re200-FDM/PINO-pretrain-Re200-1s.pt'\n\nlog:\n project: 'PINO-transfer-tanh'\n group: 'Re200to200-1s'\n\n\n" }, { "path": "configs/transfer/Re200to250-1s.yaml", "content": "data:\n datapath: 'data/NS_fine_Re250_T128_part2.npy'\n Re: 250\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 1.0\n nx: 128\n nt: 128\n sub: 1\n sub_t: 2\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 8000\n milestones: [1000, 2000, 3000, 4000, 5000, 6000, 7000]\n base_lr: 0.0025\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0.0\n save_dir: 'Re100-FDM'\n save_name: 'PINO-transfer-Re250-1s.pt'\n ckpt: 'checkpoints/Re200-FDM/PINO-pretrain-Re200-1s.pt'\n\nlog:\n project: 'PINO-transfer-tanh'\n group: 'Re200to250-1s'\n\n\n" }, { "path": "configs/transfer/Re200to300-1s.yaml", "content": "data:\n datapath: 'data/NS_fine_Re300_T128_part2.npy'\n Re: 300\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 1.0\n nx: 128\n nt: 128\n sub: 1\n sub_t: 2\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 8000\n milestones: [1000, 2000, 3000, 4000, 5000, 6000, 7000]\n base_lr: 0.0025\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0.0\n save_dir: 'Re300-FDM'\n save_name: 'PINO-transfer-Re300-1s.pt'\n ckpt: 'checkpoints/Re200-FDM/PINO-pretrain-Re200-1s.pt'\n\nlog:\n project: 'PINO-transfer-tanh'\n group: 'Re200to300-1s'\n\n\n" }, { "path": "configs/transfer/Re200to350-1s.yaml", "content": "data:\n datapath: 'data/NS_fine_Re350_T128_part2.npy'\n Re: 350\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 1.0\n nx: 128\n nt: 128\n sub: 1\n sub_t: 2\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 8000\n milestones: [1000, 2000, 3000, 4000, 5000, 6000, 7000]\n base_lr: 0.0025\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0.0\n save_dir: 'Re350-FDM'\n save_name: 'PINO-transfer-Re350-1s.pt'\n ckpt: 'checkpoints/Re200-FDM/PINO-pretrain-Re200-1s.pt'\n\nlog:\n project: 'PINO-transfer-tanh'\n group: 'Re200to350-1s'\n\n\n" }, { "path": "configs/transfer/Re200to400-1s.yaml", "content": "data:\n datapath: 'data/NS_fine_Re400_T128_part2.npy'\n Re: 400\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 1.0\n nx: 128\n nt: 128\n sub: 1\n sub_t: 2\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 8000\n milestones: [1000, 2000, 3000, 4000, 5000, 6000, 7000]\n base_lr: 0.0025\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0.0\n save_dir: 'Re400-FDM'\n save_name: 'PINO-transfer-Re400-1s.pt'\n ckpt: 'checkpoints/Re200-FDM/PINO-pretrain-Re200-1s.pt'\n\nlog:\n project: 'PINO-transfer-tanh'\n group: 'Re200to400-1s'\n\n\n" }, { "path": "configs/transfer/Re200to500-1s.yaml", "content": "data:\n datapath: 'data/NS_fine_Re500_T128_part2.npy'\n Re: 500\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 1\n nx: 128\n nt: 128\n sub: 1\n sub_t: 2\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 8000\n milestones: [1000, 2000, 3000, 4000, 5000, 6000, 7000]\n base_lr: 0.0025\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0\n save_dir: 'Re500-FDM'\n save_name: 'PINO-Re500-1s.pt'\n ckpt: 'checkpoints/Re200-FDM/PINO-pretrain-Re200-1s.pt'\n\nlog:\n project: 'PINO-transfer-tanh'\n group: 'Re200to500-1s'\n\n\n\n\n" }, { "path": "configs/transfer/Re250to100-1s.yaml", "content": "data:\n datapath: 'data/NS_fine_Re100_T128_part2.npy'\n Re: 100\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 1.0\n nx: 128\n nt: 128\n sub: 1\n sub_t: 2\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 8000\n milestones: [1000, 2000, 3000, 4000, 5000, 6000, 7000]\n base_lr: 0.0025\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0.0\n save_dir: 'Re100-FDM'\n save_name: 'PINO-Re100-1s.pt'\n ckpt: 'checkpoints/Re250-FDM/PINO-pretrain-Re250-1s.pt'\n\nlog:\n project: 'PINO-transfer-tanh'\n group: 'Re250to100-1s'\n\n\n" }, { "path": "configs/transfer/Re250to200-1s.yaml", "content": "data:\n datapath: 'data/NS_fine_Re200_T128_part2.npy'\n Re: 200\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 1.0\n nx: 128\n nt: 128\n sub: 1\n sub_t: 2\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 8000\n milestones: [1000, 2000, 3000, 4000, 5000, 6000, 7000]\n base_lr: 0.0025\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0.0\n save_dir: 'Re200-FDM'\n save_name: 'PINO-Re200-1s.pt'\n ckpt: 'checkpoints/Re250-FDM/PINO-pretrain-Re250-1s.pt'\n\nlog:\n project: 'PINO-transfer-tanh'\n group: 'Re250to200-1s'\n\n\n" }, { "path": "configs/transfer/Re250to250-1s.yaml", "content": "data:\n datapath: 'data/NS_fine_Re250_T128_part2.npy'\n Re: 250\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 1.0\n nx: 128\n nt: 128\n sub: 1\n sub_t: 2\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 8000\n milestones: [1000, 2000, 3000, 4000, 5000, 6000, 7000]\n base_lr: 0.0025\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0.0\n save_dir: 'Re250-FDM'\n save_name: 'PINO-Re250-1s.pt'\n ckpt: 'checkpoints/Re250-FDM/PINO-pretrain-Re250-1s.pt'\n\nlog:\n project: 'PINO-transfer-tanh'\n group: 'Re250to250-1s'\n\n\n" }, { "path": "configs/transfer/Re250to300-1s.yaml", "content": "data:\n datapath: 'data/NS_fine_Re300_T128_part2.npy'\n Re: 300\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 1.0\n nx: 128\n nt: 128\n sub: 1\n sub_t: 2\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 8000\n milestones: [1000, 2000, 3000, 4000, 5000, 6000, 7000]\n base_lr: 0.0025\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0.0\n save_dir: 'Re300-FDM'\n save_name: 'PINO-Re300-1s.pt'\n ckpt: 'checkpoints/Re250-FDM/PINO-pretrain-Re250-1s.pt'\n\nlog:\n project: 'PINO-transfer-tanh'\n group: 'Re250to300-1s'\n\n\n" }, { "path": "configs/transfer/Re250to350-1s.yaml", "content": "data:\n datapath: 'data/NS_fine_Re350_T128_part2.npy'\n Re: 350\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 1.0\n nx: 128\n nt: 128\n sub: 1\n sub_t: 2\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 8000\n milestones: [1000, 2000, 3000, 4000, 5000, 6000, 7000]\n base_lr: 0.0025\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0.0\n save_dir: 'Re350-FDM'\n save_name: 'PINO-Re350-1s.pt'\n ckpt: 'checkpoints/Re250-FDM/PINO-pretrain-Re250-1s.pt'\n\nlog:\n project: 'PINO-transfer-tanh'\n group: 'Re250to350-1s'\n\n\n" }, { "path": "configs/transfer/Re250to400-1s.yaml", "content": "data:\n datapath: 'data/NS_fine_Re400_T128_part2.npy'\n Re: 400\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 1.0\n nx: 128\n nt: 128\n sub: 1\n sub_t: 2\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 8000\n milestones: [1000, 2000, 3000, 4000, 5000, 6000, 7000]\n base_lr: 0.0025\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0.0\n save_dir: 'Re400-FDM'\n save_name: 'PINO-Re400-1s.pt'\n ckpt: 'checkpoints/Re250-FDM/PINO-pretrain-Re250-1s.pt'\n\nlog:\n project: 'PINO-transfer-tanh'\n group: 'Re250to400-1s'\n\n\n" }, { "path": "configs/transfer/Re250to500-1s.yaml", "content": "data:\n datapath: 'data/NS_fine_Re500_T128_part2.npy'\n Re: 500\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 1\n nx: 128\n nt: 128\n sub: 1\n sub_t: 2\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 10000\n milestones: [1000, 2000, 3000, 4000, 5000]\n base_lr: 0.0025\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0\n save_dir: 'Re500-FDM'\n save_name: 'PINO-Re500-1s.pt'\n ckpt: 'checkpoints/Re250-FDM/PINO-pretrain-Re250-1s.pt'\n\nlog:\n project: 'PINO-transfer-tanh'\n group: 'Re250to500-1s'\n\n\n\n\n" }, { "path": "configs/transfer/Re300to100-1s.yaml", "content": "data:\n datapath: 'data/NS_fine_Re100_T128_part2.npy'\n Re: 100\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 1.0\n nx: 128\n nt: 128\n sub: 1\n sub_t: 2\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 8000\n milestones: [1000, 2000, 3000, 4000, 5000, 6000, 7000]\n base_lr: 0.0025\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0.0\n save_dir: 'Re100-FDM'\n save_name: 'PINO-Re100-1s.pt'\n ckpt: 'checkpoints/Re300-FDM/PINO-pretrain-Re300-1s.pt'\n\nlog:\n project: 'PINO-transfer-tanh'\n group: 'Re300to100-1s'\n\n\n" }, { "path": "configs/transfer/Re300to200-1s.yaml", "content": "data:\n datapath: 'data/NS_fine_Re200_T128_part2.npy'\n Re: 200\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 1.0\n nx: 128\n nt: 128\n sub: 1\n sub_t: 2\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 8000\n milestones: [1000, 2000, 3000, 4000, 5000, 6000, 7000]\n base_lr: 0.0025\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0.0\n save_dir: 'Re200-FDM'\n save_name: 'PINO-Re200-1s.pt'\n ckpt: 'checkpoints/Re300-FDM/PINO-pretrain-Re300-1s.pt'\n\nlog:\n project: 'PINO-transfer-tanh'\n group: 'Re300to200-1s'\n\n\n" }, { "path": "configs/transfer/Re300to250-1s.yaml", "content": "data:\n datapath: 'data/NS_fine_Re250_T128_part2.npy'\n Re: 250\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 1.0\n nx: 128\n nt: 128\n sub: 1\n sub_t: 2\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 8000\n milestones: [1000, 2000, 3000, 4000, 5000, 6000, 7000]\n base_lr: 0.0025\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0.0\n save_dir: 'Re250-FDM'\n save_name: 'PINO-Re250-1s.pt'\n ckpt: 'checkpoints/Re300-FDM/PINO-pretrain-Re300-1s.pt'\n\nlog:\n project: 'PINO-transfer-tanh'\n group: 'Re300to250-1s'\n\n\n" }, { "path": "configs/transfer/Re300to300-1s.yaml", "content": "data:\n datapath: 'data/NS_fine_Re300_T128_part2.npy'\n Re: 300\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 1.0\n nx: 128\n nt: 128\n sub: 1\n sub_t: 2\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 8000\n milestones: [1000, 2000, 3000, 4000, 5000, 6000, 7000]\n base_lr: 0.0025\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0.0\n save_dir: 'Re300-FDM'\n save_name: 'PINO-Re300-1s.pt'\n ckpt: 'checkpoints/Re300-FDM/PINO-pretrain-Re300-1s.pt'\n\nlog:\n project: 'PINO-transfer-tanh'\n group: 'Re300to300-1s'\n\n\n" }, { "path": "configs/transfer/Re300to350-1s.yaml", "content": "data:\n datapath: 'data/NS_fine_Re350_T128_part2.npy'\n Re: 350\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 1.0\n nx: 128\n nt: 128\n sub: 1\n sub_t: 2\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 8000\n milestones: [1000, 2000, 3000, 4000, 5000, 6000, 7000]\n base_lr: 0.0025\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0.0\n save_dir: 'Re350-FDM'\n save_name: 'PINO-Re350-1s.pt'\n ckpt: 'checkpoints/Re300-FDM/PINO-pretrain-Re300-1s.pt'\n\nlog:\n project: 'PINO-transfer-tanh'\n group: 'Re300to350-1s'\n\n\n" }, { "path": "configs/transfer/Re300to400-1s.yaml", "content": "data:\n datapath: 'data/NS_fine_Re400_T128_part2.npy'\n Re: 400\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 1.0\n nx: 128\n nt: 128\n sub: 1\n sub_t: 2\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 8000\n milestones: [1000, 2000, 3000, 4000, 5000, 6000, 7000]\n base_lr: 0.0025\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0.0\n save_dir: 'Re400-FDM'\n save_name: 'PINO-Re400-1s.pt'\n ckpt: 'checkpoints/Re300-FDM/PINO-pretrain-Re300-1s.pt'\n\nlog:\n project: 'PINO-transfer-tanh'\n group: 'Re300to400-1s'\n\n\n" }, { "path": "configs/transfer/Re300to500-1s.yaml", "content": "data:\n datapath: 'data/NS_fine_Re500_T128_part2.npy'\n Re: 500\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 1\n nx: 128\n nt: 128\n sub: 1\n sub_t: 2\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 8000\n milestones: [1000, 2000, 3000, 4000, 5000, 6000, 7000]\n base_lr: 0.0025\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0\n save_dir: 'Re500-FDM'\n save_name: 'PINO-Re500-1s.pt'\n ckpt: 'checkpoints/Re300-FDM/PINO-pretrain-Re300-1s.pt'\n\nlog:\n project: 'PINO-transfer-tanh'\n group: 'Re300to500-1s'\n\n\n\n\n" }, { "path": "configs/transfer/Re350to100-1s.yaml", "content": "data:\n datapath: 'data/NS_fine_Re100_T128_part2.npy'\n Re: 100\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 1.0\n nx: 128\n nt: 128\n sub: 1\n sub_t: 2\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 8000\n milestones: [1000, 2000, 3000, 4000, 5000, 6000, 7000]\n base_lr: 0.0025\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0.0\n save_dir: 'Re100-FDM'\n save_name: 'PINO-Re100-1s.pt'\n ckpt: 'checkpoints/Re350-FDM/PINO-pretrain-Re350-1s.pt'\n\nlog:\n project: 'PINO-transfer-tanh'\n group: 'Re350to100-1s'\n\n\n" }, { "path": "configs/transfer/Re350to200-1s.yaml", "content": "data:\n datapath: 'data/NS_fine_Re200_T128_part2.npy'\n Re: 200\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 1.0\n nx: 128\n nt: 128\n sub: 1\n sub_t: 2\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 8000\n milestones: [1000, 2000, 3000, 4000, 5000, 6000, 7000]\n base_lr: 0.0025\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0.0\n save_dir: 'Re200-FDM'\n save_name: 'PINO-Re200-1s.pt'\n ckpt: 'checkpoints/Re350-FDM/PINO-pretrain-Re350-1s.pt'\n\nlog:\n project: 'PINO-transfer-tanh'\n group: 'Re350to200-1s'\n\n\n" }, { "path": "configs/transfer/Re350to250-1s.yaml", "content": "data:\n datapath: 'data/NS_fine_Re250_T128_part2.npy'\n Re: 250\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 1.0\n nx: 128\n nt: 128\n sub: 1\n sub_t: 2\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 8000\n milestones: [1000, 2000, 3000, 4000, 5000, 6000, 7000]\n base_lr: 0.0025\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0.0\n save_dir: 'Re250-FDM'\n save_name: 'PINO-Re250-1s.pt'\n ckpt: 'checkpoints/Re350-FDM/PINO-pretrain-Re350-1s.pt'\n\nlog:\n project: 'PINO-transfer-tanh'\n group: 'Re350to250-1s'\n\n\n" }, { "path": "configs/transfer/Re350to300-1s.yaml", "content": "data:\n datapath: 'data/NS_fine_Re300_T128_part2.npy'\n Re: 300\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 1.0\n nx: 128\n nt: 128\n sub: 1\n sub_t: 2\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 8000\n milestones: [1000, 2000, 3000, 4000, 5000, 6000, 7000]\n base_lr: 0.0025\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0.0\n save_dir: 'Re300-FDM'\n save_name: 'PINO-Re300-1s.pt'\n ckpt: 'checkpoints/Re350-FDM/PINO-pretrain-Re350-1s.pt'\n\nlog:\n project: 'PINO-transfer-tanh'\n group: 'Re350to300-1s'\n\n\n" }, { "path": "configs/transfer/Re350to350-1s.yaml", "content": "data:\n datapath: 'data/NS_fine_Re350_T128_part2.npy'\n Re: 350\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 1.0\n nx: 128\n nt: 128\n sub: 1\n sub_t: 2\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 8000\n milestones: [1000, 2000, 3000, 4000, 5000, 6000, 7000]\n base_lr: 0.0025\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0.0\n save_dir: 'Re350-FDM'\n save_name: 'PINO-Re300-1s.pt'\n ckpt: 'checkpoints/Re350-FDM/PINO-pretrain-Re350-1s.pt'\n\nlog:\n project: 'PINO-transfer-tanh'\n group: 'Re350to350-1s'\n\n\n" }, { "path": "configs/transfer/Re350to400-1s.yaml", "content": "data:\n datapath: 'data/NS_fine_Re400_T128_part2.npy'\n Re: 400\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 1.0\n nx: 128\n nt: 128\n sub: 1\n sub_t: 2\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 8000\n milestones: [1000, 2000, 3000, 4000, 5000, 6000, 7000]\n base_lr: 0.0025\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0.0\n save_dir: 'Re400-FDM'\n save_name: 'PINO-Re400-1s.pt'\n ckpt: 'checkpoints/Re350-FDM/PINO-pretrain-Re350-1s.pt'\n\nlog:\n project: 'PINO-transfer-tanh'\n group: 'Re350to400-1s'\n\n\n" }, { "path": "configs/transfer/Re350to500-1s.yaml", "content": "data:\n datapath: 'data/NS_fine_Re500_T128_part2.npy'\n Re: 500\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 1\n nx: 128\n nt: 128\n sub: 1\n sub_t: 2\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 8000\n milestones: [1000, 2000, 3000, 4000, 5000, 6000, 7000]\n base_lr: 0.0025\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0\n save_dir: 'Re500-FDM'\n save_name: 'PINO-Re500-1s.pt'\n ckpt: 'checkpoints/Re350-FDM/PINO-pretrain-Re350-1s.pt'\n\nlog:\n project: 'PINO-transfer-tanh'\n group: 'Re350to500-1s'\n\n\n\n\n" }, { "path": "configs/transfer/Re400to100-1s.yaml", "content": "data:\n datapath: 'data/NS_fine_Re100_T128_part2.npy'\n Re: 100\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 1.0\n nx: 128\n nt: 128\n sub: 1\n sub_t: 2\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 8000\n milestones: [1000, 2000, 3000, 4000, 5000, 6000, 7000]\n base_lr: 0.0025\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0.0\n save_dir: 'Re100-FDM'\n save_name: 'PINO-Re100-1s.pt'\n ckpt: 'checkpoints/Re400-FDM/PINO-pretrain-Re400-1s.pt'\n\nlog:\n project: 'PINO-transfer-tanh'\n group: 'Re400to100-1s'\n\n\n" }, { "path": "configs/transfer/Re400to200-1s.yaml", "content": "data:\n datapath: 'data/NS_fine_Re200_T128_part2.npy'\n Re: 200\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 1.0\n nx: 128\n nt: 128\n sub: 1\n sub_t: 2\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 8000\n milestones: [1000, 2000, 3000, 4000, 5000, 6000, 7000]\n base_lr: 0.0025\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0.0\n save_dir: 'Re200-FDM'\n save_name: 'PINO-Re200-1s.pt'\n ckpt: 'checkpoints/Re400-FDM/PINO-pretrain-Re400-1s.pt'\n\nlog:\n project: 'PINO-transfer-tanh'\n group: 'Re400to200-1s'\n\n\n" }, { "path": "configs/transfer/Re400to250-1s.yaml", "content": "data:\n datapath: 'data/NS_fine_Re250_T128_part2.npy'\n Re: 250\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 1.0\n nx: 128\n nt: 128\n sub: 1\n sub_t: 2\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 8000\n milestones: [1000, 2000, 3000, 4000, 5000, 6000, 7000]\n base_lr: 0.0025\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0.0\n save_dir: 'Re250-FDM'\n save_name: 'PINO-Re250-1s.pt'\n ckpt: 'checkpoints/Re400-FDM/PINO-pretrain-Re400-1s.pt'\n\nlog:\n project: 'PINO-transfer-tanh'\n group: 'Re400to250-1s'\n\n\n" }, { "path": "configs/transfer/Re400to300-1s.yaml", "content": "data:\n datapath: 'data/NS_fine_Re300_T128_part2.npy'\n Re: 300\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 1.0\n nx: 128\n nt: 128\n sub: 1\n sub_t: 2\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 8000\n milestones: [1000, 2000, 3000, 4000, 5000, 6000, 7000]\n base_lr: 0.0025\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0.0\n save_dir: 'Re300-FDM'\n save_name: 'PINO-Re300-1s.pt'\n ckpt: 'checkpoints/Re400-FDM/PINO-pretrain-Re400-1s.pt'\n\nlog:\n project: 'PINO-transfer-tanh'\n group: 'Re400to300-1s'\n\n\n" }, { "path": "configs/transfer/Re400to350-1s.yaml", "content": "data:\n datapath: 'data/NS_fine_Re350_T128_part2.npy'\n Re: 350\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 1.0\n nx: 128\n nt: 128\n sub: 1\n sub_t: 2\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 8000\n milestones: [1000, 2000, 3000, 4000, 5000, 6000, 7000]\n base_lr: 0.0025\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0.0\n save_dir: 'Re350-FDM'\n save_name: 'PINO-Re350-1s.pt'\n ckpt: 'checkpoints/Re400-FDM/PINO-pretrain-Re400-1s.pt'\n\nlog:\n project: 'PINO-transfer-tanh'\n group: 'Re400to350-1s'\n\n\n" }, { "path": "configs/transfer/Re400to400-1s.yaml", "content": "data:\n datapath: 'data/NS_fine_Re400_T128_part2.npy'\n Re: 400\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 1.0\n nx: 128\n nt: 128\n sub: 1\n sub_t: 2\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 8000\n milestones: [1000, 2000, 3000, 4000, 5000, 6000, 7000]\n base_lr: 0.0025\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0.0\n save_dir: 'Re400-FDM'\n save_name: 'PINO-Re400-1s.pt'\n ckpt: 'checkpoints/Re400-FDM/PINO-pretrain-Re400-1s.pt'\n\nlog:\n project: 'PINO-transfer-tanh'\n group: 'Re400to400-1s'\n\n\n" }, { "path": "configs/transfer/Re400to500-1s.yaml", "content": "data:\n datapath: 'data/NS_fine_Re500_T128_part2.npy'\n Re: 500\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 1\n nx: 128\n nt: 128\n sub: 1\n sub_t: 2\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 8000\n milestones: [1000, 2000, 3000, 4000, 5000, 6000, 7000]\n base_lr: 0.0025\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0\n save_dir: 'Re500-FDM'\n save_name: 'PINO-Re500-1s.pt'\n ckpt: 'checkpoints/Re400-FDM/PINO-pretrain-Re400-1s.pt'\n\nlog:\n project: 'PINO-transfer-tanh'\n group: 'Re400to500-1s'\n\n\n\n\n" }, { "path": "configs/transfer/Re500to100-1s.yaml", "content": "data:\n datapath: 'data/NS_fine_Re100_T128_part2.npy'\n Re: 100\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 1.0\n nx: 128\n nt: 128\n sub: 1\n sub_t: 2\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 8000\n milestones: [1000, 2000, 3000, 4000, 5000, 6000, 7000]\n base_lr: 0.0025\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0.0\n save_dir: 'Re100-FDM'\n save_name: 'PINO-Re100-1s.pt'\n ckpt: 'checkpoints/Re500-FDM/PINO-pretrain-Re500-1s.pt'\n\nlog:\n project: 'PINO-transfer-tanh'\n group: 'Re500to100-1s'\n\n\n" }, { "path": "configs/transfer/Re500to200-1s.yaml", "content": "data:\n datapath: 'data/NS_fine_Re200_T128_part2.npy'\n Re: 200\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 1.0\n nx: 128\n nt: 128\n sub: 1\n sub_t: 2\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 8000\n milestones: [1000, 2000, 3000, 4000, 5000, 6000, 7000]\n base_lr: 0.0025\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0.0\n save_dir: 'Re200-FDM'\n save_name: 'PINO-Re200-1s.pt'\n ckpt: 'checkpoints/Re500-FDM/PINO-pretrain-Re500-1s.pt'\n\nlog:\n project: 'PINO-transfer-tanh'\n group: 'Re500to200-1s'\n\n\n" }, { "path": "configs/transfer/Re500to250-1s.yaml", "content": "data:\n datapath: 'data/NS_fine_Re250_T128_part2.npy'\n Re: 250\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 1.0\n nx: 128\n nt: 128\n sub: 1\n sub_t: 2\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 8000\n milestones: [1000, 2000, 3000, 4000, 5000, 6000, 7000]\n base_lr: 0.0025\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0.0\n save_dir: 'Re250-FDM'\n save_name: 'PINO-Re250-1s.pt'\n ckpt: 'checkpoints/Re500-FDM/PINO-pretrain-Re500-1s.pt'\n\nlog:\n project: 'PINO-transfer-tanh'\n group: 'Re500to250-1s'\n\n\n" }, { "path": "configs/transfer/Re500to300-1s.yaml", "content": "data:\n datapath: 'data/NS_fine_Re300_T128_part2.npy'\n Re: 300\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 1.0\n nx: 128\n nt: 128\n sub: 1\n sub_t: 2\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 8000\n milestones: [1000, 2000, 3000, 4000, 5000, 6000, 7000]\n base_lr: 0.0025\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0.0\n save_dir: 'Re300-FDM'\n save_name: 'PINO-Re300-1s.pt'\n ckpt: 'checkpoints/Re500-FDM/PINO-pretrain-Re500-1s.pt'\n\nlog:\n project: 'PINO-transfer-tanh'\n group: 'Re500to300-1s'\n\n\n" }, { "path": "configs/transfer/Re500to350-1s.yaml", "content": "data:\n datapath: 'data/NS_fine_Re350_T128_part2.npy'\n Re: 350\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 1.0\n nx: 128\n nt: 128\n sub: 1\n sub_t: 2\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 8000\n milestones: [1000, 2000, 3000, 4000, 5000, 6000, 7000]\n base_lr: 0.0025\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0.0\n save_dir: 'Re350-FDM'\n save_name: 'PINO-Re350-1s.pt'\n ckpt: 'checkpoints/Re500-FDM/PINO-pretrain-Re500-1s.pt'\n\nlog:\n project: 'PINO-transfer-tanh'\n group: 'Re500to350-1s'\n\n\n" }, { "path": "configs/transfer/Re500to400-1s.yaml", "content": "data:\n datapath: 'data/NS_fine_Re400_T128_part2.npy'\n Re: 400\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 1.0\n nx: 128\n nt: 128\n sub: 1\n sub_t: 2\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 8000\n milestones: [1000, 2000, 3000, 4000, 5000, 6000, 7000]\n base_lr: 0.0025\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0.0\n save_dir: 'Re400-FDM'\n save_name: 'PINO-Re400-1s.pt'\n ckpt: 'checkpoints/Re500-FDM/PINO-pretrain-Re500-1s.pt'\n\nlog:\n project: 'PINO-transfer-tanh'\n group: 'Re500to400-1s'\n\n\n" }, { "path": "configs/transfer/Re500to500-05s-new.yaml", "content": "data:\n datapath: 'data/NS_fine_Re500_s2048_T100.npy'\n Re: 500\n total_num: 100\n offset: 300\n n_sample: 1\n time_interval: 0.5\n nx: 256\n nt: 128\n sub: 2\n sub_t: 1\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 8000\n milestones: [1000, 2000, 3000, 4000, 5000, 6000]\n base_lr: 0.0025\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0.0\n save_dir: 'Re500-FDM'\n save_name: 'PINO-Re500-05.pt'\n ckpt: 'checkpoints/Re500-FDM/pretrain-Re500-05s-4000.pt'\n\nlog:\n project: 'PINO-Re500-exp'\n group: 'Re500to500-128-4k-all'\n\n\n\n" }, { "path": "configs/transfer/Re500to500-05s.yaml", "content": "data:\n datapath: 'data/NS_fine_Re500_T128_part2.npy'\n Re: 500\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 0.5\n nx: 128\n nt: 128\n sub: 1\n sub_t: 1\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 5000\n milestones: [1000, 2000, 3000, 4000]\n base_lr: 0.001\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0.0\n save_dir: 'Re500-FDM'\n save_name: 'PINO-Re500-1s.pt'\n ckpt: 'checkpoints/Re500-FDM/pretrain-Re500-05s-4000.pt'\n\nlog:\n project: 'PINO-Re500-exp'\n group: 'Re500to500-05s-4layer'\n\n\n" }, { "path": "configs/transfer/Re500to500-1s.yaml", "content": "data:\n datapath: 'data/NS_fine_Re500_T128_part2.npy'\n Re: 500\n total_num: 100\n offset: 0\n n_sample: 1\n time_interval: 1.0\n nx: 128\n nt: 128\n sub: 1\n sub_t: 1\n shuffle: True\n\nmodel:\n layers: [64, 64, 64, 64, 64]\n modes1: [8, 8, 8, 8]\n modes2: [8, 8, 8, 8]\n modes3: [8, 8, 8, 8]\n fc_dim: 128\n\ntrain:\n batchsize: 1\n epochs: 8000\n milestones: [1000, 2000, 3000, 4000, 5000, 6000]\n base_lr: 0.0025\n scheduler_gamma: 0.5\n ic_loss: 5.0\n f_loss: 1.0\n xy_loss: 0.0\n save_dir: 'Re500-FDM'\n save_name: 'PINO-Re500-1s.pt'\n ckpt: 'checkpoints/Re500-FDM/PINO-pretrain-Re500-1s.pt'\n\nlog:\n project: 'PINO-transfer'\n group: 'Re500to500-1s-new'\n\n\n" }, { "path": "deeponet.py", "content": "import yaml\nfrom argparse import ArgumentParser\nfrom baselines.train_ns import train_deeponet_cp\nfrom baselines.test import test_deeponet_ns, test_deeponet_darcy\nfrom baselines.train_darcy import train_deeponet_darcy\n\n\nif __name__ == '__main__':\n parser = ArgumentParser(description='Basic paser')\n parser.add_argument('--config_path', type=str, help='Path to the configuration file')\n parser.add_argument('--mode', type=str, default='train', help='Train or test')\n args = parser.parse_args()\n\n config_file = args.config_path\n with open(config_file, 'r') as stream:\n config = yaml.load(stream, yaml.FullLoader)\n\n if args.mode == 'train':\n print('Start training DeepONet Cartesian Product')\n if 'name' in config['data'] and config['data']['name'] == 'Darcy':\n train_deeponet_darcy(config)\n else:\n train_deeponet_cp(config)\n else:\n print('Start testing DeepONet Cartesian Product')\n if 'name' in config['data'] and config['data']['name'] == 'Darcy':\n test_deeponet_darcy(config)\n else:\n test_deeponet_ns(config)\n print('Done!')" }, { "path": "download_data.py", "content": "import os\nfrom argparse import ArgumentParser\nimport requests\nfrom tqdm import tqdm\n\n\n_url_dict = {\n 'NS-T4000': 'https://hkzdata.s3.us-west-2.amazonaws.com/PINO/data/NS_fft_Re500_T4000.npy', \n 'NS-Re500Part0': 'https://hkzdata.s3.us-west-2.amazonaws.com/PINO/data/NS_fine_Re500_T128_part0.npy', \n 'NS-Re500Part1': 'https://hkzdata.s3.us-west-2.amazonaws.com/PINO/data/NS_fine_Re500_T128_part1.npy', \n 'NS-Re500Part2': 'https://hkzdata.s3.us-west-2.amazonaws.com/PINO/data/NS_fine_Re500_T128_part2.npy', \n 'NS-Re100Part0': 'https://hkzdata.s3.us-west-2.amazonaws.com/PINO/data/NS_fine_Re100_T128_part0.npy', \n 'burgers': 'https://hkzdata.s3.us-west-2.amazonaws.com/PINO/data/burgers_pino.mat', \n 'NS-Re500_T300_id0': 'https://hkzdata.s3.us-west-2.amazonaws.com/PINO/NS-Re500_T300_id0.npy',\n 'NS-Re500_T300_id0-shuffle': 'https://hkzdata.s3.us-west-2.amazonaws.com/PINO/data/NS-Re500_T300_id0-shuffle.npy',\n 'darcy-train': 'https://hkzdata.s3.us-west-2.amazonaws.com/PINO/data/piececonst_r421_N1024_smooth1.mat', \n 'darcy-test': 'https://hkzdata.s3.us-west-2.amazonaws.com/PINO/data/piececonst_r421_N1024_smooth2.mat', \n 'cavity': 'https://hkzdata.s3.us-west-2.amazonaws.com/PINO/data/cavity.mat',\n 'Re500-1_8s-800-pino-140k': 'https://hkzdata.s3.us-west-2.amazonaws.com/PINO/checkpoints/Re500-1_8s-800-PINO-140000.pt',\n}\n\n\ndef download_file(url, file_path):\n print('Start downloading...')\n with requests.get(url, stream=True) as r:\n r.raise_for_status()\n with open(file_path, 'wb') as f:\n for chunk in tqdm(r.iter_content(chunk_size=256 * 1024 * 1024)):\n f.write(chunk)\n print('Complete')\n\n\ndef main(args):\n url = _url_dict[args.name]\n file_name = url.split('/')[-1]\n os.makedirs(args.outdir, exist_ok=True)\n file_path = os.path.join(args.outdir, file_name)\n download_file(url, file_path)\n\n\nif __name__ == '__main__':\n parser = ArgumentParser(description='Parser for downloading assets')\n parser.add_argument('--name', type=str, default='NS-T4000')\n parser.add_argument('--outdir', type=str, default='../data')\n args = parser.parse_args()\n main(args)" }, { "path": "eval_operator.py", "content": "import yaml\n\nimport torch\nfrom torch.utils.data import DataLoader\nfrom models import FNO3d, FNO2d\nfrom train_utils import NSLoader, get_forcing, DarcyFlow\n\nfrom train_utils.eval_3d import eval_ns\nfrom train_utils.eval_2d import eval_darcy\n\nfrom argparse import ArgumentParser\n\n\ndef test_3d(config):\n device = 0 if torch.cuda.is_available() else 'cpu'\n data_config = config['data']\n loader = NSLoader(datapath1=data_config['datapath'],\n nx=data_config['nx'], nt=data_config['nt'],\n sub=data_config['sub'], sub_t=data_config['sub_t'],\n N=data_config['total_num'],\n t_interval=data_config['time_interval'])\n\n eval_loader = loader.make_loader(n_sample=data_config['n_sample'],\n batch_size=config['test']['batchsize'],\n start=data_config['offset'],\n train=data_config['shuffle'])\n model = FNO3d(modes1=config['model']['modes1'],\n modes2=config['model']['modes2'],\n modes3=config['model']['modes3'],\n fc_dim=config['model']['fc_dim'],\n layers=config['model']['layers']).to(device)\n\n if 'ckpt' in config['test']:\n ckpt_path = config['test']['ckpt']\n ckpt = torch.load(ckpt_path)\n model.load_state_dict(ckpt['model'])\n print('Weights loaded from %s' % ckpt_path)\n print(f'Resolution : {loader.S}x{loader.S}x{loader.T}')\n forcing = get_forcing(loader.S).to(device)\n eval_ns(model,\n loader,\n eval_loader,\n forcing,\n config,\n device=device)\n\n\ndef test_2d(config):\n device = 0 if torch.cuda.is_available() else 'cpu'\n data_config = config['data']\n dataset = DarcyFlow(data_config['datapath'],\n nx=data_config['nx'], sub=data_config['sub'],\n offset=data_config['offset'], num=data_config['n_sample'])\n dataloader = DataLoader(dataset, batch_size=config['test']['batchsize'], shuffle=False)\n print(device)\n model = FNO2d(modes1=config['model']['modes1'],\n modes2=config['model']['modes2'],\n fc_dim=config['model']['fc_dim'],\n layers=config['model']['layers'],\n act=config['model']['act']).to(device)\n # Load from checkpoint\n if 'ckpt' in config['test']:\n ckpt_path = config['test']['ckpt']\n ckpt = torch.load(ckpt_path)\n model.load_state_dict(ckpt['model'])\n print('Weights loaded from %s' % ckpt_path)\n eval_darcy(model, dataloader, config, device)\n\n\nif __name__ == '__main__':\n parser = ArgumentParser(description='Basic paser')\n parser.add_argument('--config_path', type=str, help='Path to the configuration file')\n parser.add_argument('--log', action='store_true', help='Turn on the wandb')\n options = parser.parse_args()\n config_file = options.config_path\n with open(config_file, 'r') as stream:\n config = yaml.load(stream, yaml.FullLoader)\n\n if 'name' in config['data'] and config['data']['name'] == 'Darcy':\n test_2d(config)\n else:\n test_3d(config)\n\n\n" }, { "path": "generate_data.py", "content": "import math\nimport numpy as np\nimport os\nfrom tqdm import tqdm\n\nimport torch\nfrom solver.random_fields import GaussianRF, GaussianRF2d\nfrom solver.kolmogorov_flow import KolmogorovFlow2d\nfrom solver.periodic import NavierStokes2d\nfrom timeit import default_timer\nimport argparse\n\n\n\ndef legacy_solver(args):\n save_dir = args.outdir\n os.makedirs(save_dir, exist_ok=True)\n device = torch.device('cuda:0')\n s = 1024\n sub = s // args.res_x\n \n n = 4 # forcing\n Re = args.re\n\n T_in = 100.0\n T = args.T\n t = args.t_res\n dt = 1.0 / t\n\n GRF = GaussianRF(2, s, 2 * math.pi, alpha=2.5, tau=7, device=device)\n u0 = GRF.sample(1)\n\n NS = KolmogorovFlow2d(u0, Re, n)\n NS.advance(T_in, delta_t=1e-3)\n\n sol = np.zeros((T, t + 1, s // sub, s // sub))\n sol_ini = NS.vorticity().squeeze(0).cpu().numpy()[::sub, ::sub]\n pbar = tqdm(range(T))\n for i in pbar:\n sol[i, 0, :, :] = sol_ini\n for j in range(t):\n t1 = default_timer()\n NS.advance(dt, delta_t=1e-3)\n sol[i, j + 1, :, :] = NS.vorticity().squeeze(0).cpu().numpy()[::sub, ::sub]\n t2 = default_timer()\n pbar.set_description(\n (\n f'{i}, time cost: {t2-t1}'\n )\n )\n sol_ini = sol[i, -1, :, :]\n \n save_path = os.path.join(save_dir, f'NS-Re{int(Re)}_T{t}.npy')\n # np.save('NS_fine_Re500_S512_s64_T500_t128.npy', sol)\n np.save(save_path, sol)\n\n\ndef gen_data(args):\n dtype = torch.float64\n device = torch.device('cuda:0')\n save_dir = args.outdir\n os.makedirs(save_dir, exist_ok=True)\n \n T = args.T # total time\n bsize = args.batchsize\n L = 2 * math.pi\n s =args.x_res\n x_sub = args.x_sub\n\n t_res = args.t_res\n dt = 1 / t_res\n re = args.re\n\n solver = NavierStokes2d(s,s,L,L,device=device,dtype=dtype)\n grf = GaussianRF2d(s,s,L,L,alpha=2.5,tau=3.0,sigma=None,device=device,dtype=dtype)\n\n t = torch.linspace(0, L, s+1, dtype=dtype, device=device)[0:-1]\n _, Y = torch.meshgrid(t, t, indexing='ij')\n f = -4*torch.cos(4.0*Y)\n vor = np.zeros((bsize, T, t_res + 1, s // x_sub, s // x_sub))\n\n pbar = tqdm(range(T))\n w = grf.sample(bsize)\n w = solver.advance(w, f, T=100, Re=re, adaptive=True)\n \n init_vor = w[:, ::x_sub, ::x_sub].cpu().type(torch.float32).numpy()\n for j in pbar:\n vor[:, j, 0, :, :] = init_vor\n\n for k in range(t_res):\n t1 = default_timer()\n\n w = solver.advance(w, f, T=dt, Re=re, adaptive=True)\n vor[:, j, k+1, :, :] = w[:,::x_sub,::x_sub].cpu().type(torch.float32).numpy()\n\n t2 = default_timer()\n\n pbar.set_description(\n (\n f'{j}, time cost: {t2-t1}'\n )\n )\n init_vor = vor[:, j, -1, :, :]\n\n for i in range(bsize):\n save_path = os.path.join(save_dir, f'NS-Re{int(re)}_T{T}_id{i}.npy')\n # np.save('NS_fine_Re500_S512_s64_T500_t128.npy', sol)\n np.save(save_path, vor[i])\n\n\nif __name__ == '__main__':\n parser = argparse.ArgumentParser()\n parser.add_argument('--seed', type=int, default=0)\n parser.add_argument('--re', type=float, default=40.0)\n parser.add_argument('--x_res', type=int, default=512)\n parser.add_argument('--x_sub', type=int, default=2)\n parser.add_argument('--T', type=int, default=300)\n parser.add_argument('--outdir', type=str, default='../data')\n parser.add_argument('--t_res', type=int, default=512)\n parser.add_argument('--batchsize', type=int, default=1)\n parser.add_argument('--num_batchs', type=int, default=1)\n args = parser.parse_args()\n gen_data(args)" }, { "path": "inference.py", "content": "'''\nThis code generates the prediction on one instance. \nBoth the ground truth and the prediction are saved in a .pt file.\n'''\nimport os\nimport yaml\nfrom argparse import ArgumentParser\n\nimport torch\nfrom torch.utils.data import DataLoader\n\nfrom models import FNO3d\n\nfrom train_utils.datasets import KFDataset\nfrom train_utils.losses import LpLoss\nfrom train_utils.utils import count_params\n\n\n@torch.no_grad()\ndef get_pred(args):\n with open(args.config, 'r') as stream:\n config = yaml.load(stream, yaml.FullLoader)\n basedir = os.path.join('exp', config['log']['logdir'])\n save_dir = os.path.join(basedir, 'results')\n os.makedirs(save_dir, exist_ok=True)\n save_path = os.path.join(save_dir,'fno-prediction.pt')\n device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n\n # prepare data\n dataset = KFDataset(paths=config['data']['paths'], \n raw_res=config['data']['raw_res'],\n data_res=config['data']['data_res'], \n pde_res=config['data']['data_res'], \n n_samples=config['data']['n_test_samples'],\n total_samples=config['data']['total_test_samples'],\n offset=config['data']['testoffset'], \n t_duration=config['data']['t_duration'])\n dataloader = DataLoader(dataset, batch_size=1, shuffle=False, drop_last=False)\n\n # create model\n model = FNO3d(modes1=config['model']['modes1'],\n modes2=config['model']['modes2'],\n modes3=config['model']['modes3'],\n fc_dim=config['model']['fc_dim'],\n layers=config['model']['layers'], \n act=config['model']['act'], \n pad_ratio=config['model']['pad_ratio']).to(device)\n num_params = count_params(model)\n print(f'Number of parameters: {num_params}')\n if args.ckpt_path:\n ckpt = torch.load(args.ckpt_path)\n model.load_state_dict(ckpt['model'])\n print('Weights loaded from %s' % args.ckpt_path)\n # metric\n lploss = LpLoss(size_average=True)\n model.eval()\n truth_list = []\n pred_list = []\n for u, a_in in dataloader:\n u, a_in = u.to(device), a_in.to(device)\n out = model(a_in)\n data_loss = lploss(out, u)\n print(data_loss.item())\n truth_list.append(u.cpu())\n pred_list.append(out.cpu())\n truth_arr = torch.cat(truth_list, dim=0)\n pred_arr = torch.cat(pred_list, dim=0)\n torch.save({\n 'truth': truth_arr,\n 'pred': pred_arr,\n }, save_path)\n\n\nif __name__ == \"__main__\":\n torch.backends.cudnn.benchmark = True\n parser = ArgumentParser()\n parser.add_argument('--config', type=str, default='configs/config.yaml')\n parser.add_argument('--ckpt_path', type=str, default=None)\n args = parser.parse_args()\n get_pred(args)" }, { "path": "instance_opt.py", "content": "import os\nimport yaml\nimport random\nfrom argparse import ArgumentParser\nimport math\nfrom tqdm import tqdm\n\nimport torch\n\nfrom torch.optim import Adam\nfrom torch.utils.data import DataLoader, Subset\n\nfrom models import FNO3d\n\nfrom train_utils.losses import LpLoss, PINO_loss3d, get_forcing\nfrom train_utils.datasets import KFDataset, KFaDataset, sample_data\nfrom train_utils.utils import save_ckpt, count_params, dict2str\n\ntry:\n import wandb\nexcept ImportError:\n wandb = None\n\n\ndef train_ns(model, \n u_loader, # training data\n optimizer, \n scheduler,\n device, config, args):\n\n v = 1/ config['data']['Re']\n t_duration = config['data']['t_duration']\n save_step = config['train']['save_step']\n\n ic_weight = config['train']['ic_loss']\n f_weight = config['train']['f_loss']\n # set up directory\n base_dir = os.path.join('exp', config['log']['logdir'])\n ckpt_dir = os.path.join(base_dir, 'ckpts')\n os.makedirs(ckpt_dir, exist_ok=True)\n\n # loss fn\n lploss = LpLoss(size_average=True)\n \n S = config['data']['pde_res'][0]\n forcing = get_forcing(S).to(device)\n # set up wandb\n if wandb and args.log:\n run = wandb.init(project=config['log']['project'], \n entity=config['log']['entity'], \n group=config['log']['group'], \n config=config, reinit=True, \n settings=wandb.Settings(start_method='fork'))\n \n pbar = range(config['train']['num_iter'])\n if args.tqdm:\n pbar = tqdm(pbar, dynamic_ncols=True, smoothing=0.2)\n\n u_loader = sample_data(u_loader)\n\n for e in pbar:\n log_dict = {}\n\n optimizer.zero_grad()\n # data loss\n u, a_in = next(u_loader)\n u = u.to(device)\n a_in = a_in.to(device)\n out = model(a_in)\n data_loss = lploss(out, u)\n \n u0 = a_in[:, :, :, 0, -1]\n loss_ic, loss_f = PINO_loss3d(out, u0, forcing, v, t_duration)\n log_dict['IC'] = loss_ic.item()\n log_dict['PDE'] = loss_f.item()\n loss = loss_f * f_weight + loss_ic * ic_weight\n\n loss.backward()\n optimizer.step()\n scheduler.step()\n\n log_dict['train loss'] = loss.item()\n log_dict['test error'] = data_loss.item()\n \n if args.tqdm:\n logstr = dict2str(log_dict)\n pbar.set_description(\n (\n logstr\n )\n )\n if wandb and args.log:\n wandb.log(log_dict)\n if e % save_step == 0 and e > 0:\n ckpt_path = os.path.join(ckpt_dir, f'model-{e}.pt')\n save_ckpt(ckpt_path, model, optimizer)\n\n # clean up wandb\n if wandb and args.log:\n run.finish()\n \n # save prediction and truth\n save_dir = os.path.join(base_dir, 'results')\n os.makedirs(save_dir, exist_ok=True)\n result_path = os.path.join(save_dir, f'results-{args.idx}.pt')\n\n criterion = LpLoss()\n\n model.eval()\n with torch.no_grad():\n u, a_in = next(u_loader)\n u = u.to(device)\n a_in = a_in.to(device)\n out = model(a_in)\n error = criterion(out, u)\n print(f'Test error: {error.item()}')\n torch.save({'truth': u.cpu(), 'pred': out.cpu()}, result_path)\n print(f'Results saved to {result_path}')\n\n\n\ndef subprocess(args):\n with open(args.config, 'r') as f:\n config = yaml.load(f, yaml.FullLoader)\n device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n\n # set random seed\n config['seed'] = args.seed\n seed = args.seed\n torch.manual_seed(seed)\n random.seed(seed)\n if torch.cuda.is_available():\n torch.cuda.manual_seed_all(seed)\n\n # create model \n model = FNO3d(modes1=config['model']['modes1'],\n modes2=config['model']['modes2'],\n modes3=config['model']['modes3'],\n fc_dim=config['model']['fc_dim'],\n layers=config['model']['layers'], \n act=config['model']['act'], \n pad_ratio=config['model']['pad_ratio']).to(device)\n num_params = count_params(model)\n config['num_params'] = num_params\n print(f'Number of parameters: {num_params}')\n # Load from checkpoint\n if args.ckpt:\n ckpt_path = args.ckpt\n ckpt = torch.load(ckpt_path)\n model.load_state_dict(ckpt['model'])\n print('Weights loaded from %s' % ckpt_path)\n \n # training set\n batchsize = config['train']['batchsize']\n dataset = KFDataset(paths=config['data']['paths'], \n raw_res=config['data']['raw_res'],\n data_res=config['data']['data_res'], \n pde_res=config['data']['data_res'], \n n_samples=config['data']['n_test_samples'], \n total_samples=1,\n idx=args.idx,\n offset=config['data']['testoffset'], \n t_duration=config['data']['t_duration'])\n u_loader = DataLoader(dataset, batch_size=1)\n\n optimizer = Adam(model.parameters(), lr=config['train']['base_lr'])\n scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, \n milestones=config['train']['milestones'], \n gamma=config['train']['scheduler_gamma'])\n train_ns(model, \n u_loader, \n optimizer, \n scheduler, \n device, \n config, \n args)\n print('Done!')\n \n \n\nif __name__ == '__main__':\n torch.backends.cudnn.benchmark = True\n # parse options\n parser = ArgumentParser(description='Basic paser')\n parser.add_argument('--config', type=str, help='Path to the configuration file')\n parser.add_argument('--idx', type=int, default=0, help='Index of the instance')\n parser.add_argument('--log', action='store_true', help='Turn on the wandb')\n parser.add_argument('--seed', type=int, default=None)\n parser.add_argument('--ckpt', type=str, default=None)\n parser.add_argument('--tqdm', action='store_true', help='Turn on the tqdm')\n args = parser.parse_args()\n if args.seed is None:\n args.seed = random.randint(0, 100000)\n subprocess(args)" }, { "path": "inverse-darcy-foward.py", "content": "\n\nfrom timeit import default_timer\n\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nfrom torch.optim import Adam\n\nfrom train_utils.datasets import MatReader\nfrom train_utils.losses import LpLoss\nfrom train_utils.utils import count_params\n\ntorch.manual_seed(0)\nnp.random.seed(0)\n\n\n################################################################\n# fourier layer\n################################################################\nclass SpectralConv2d(nn.Module):\n def __init__(self, in_channels, out_channels, modes1, modes2):\n super(SpectralConv2d, self).__init__()\n\n \"\"\"\n 2D Fourier layer. It does FFT, linear transform, and Inverse FFT. \n \"\"\"\n\n self.in_channels = in_channels\n self.out_channels = out_channels\n self.modes1 = modes1 # Number of Fourier modes to multiply, at most floor(N/2) + 1\n self.modes2 = modes2\n\n self.scale = (1 / (in_channels * out_channels))\n self.weights1 = nn.Parameter(\n self.scale * torch.rand(in_channels, out_channels, self.modes1, self.modes2, dtype=torch.cfloat))\n self.weights2 = nn.Parameter(\n self.scale * torch.rand(in_channels, out_channels, self.modes1, self.modes2, dtype=torch.cfloat))\n\n # Complex multiplication\n def compl_mul2d(self, input, weights):\n # (batch, in_channel, x,y ), (in_channel, out_channel, x,y) -> (batch, out_channel, x,y)\n return torch.einsum(\"bixy,ioxy->boxy\", input, weights)\n\n def forward(self, x):\n batchsize = x.shape[0]\n # Compute Fourier coeffcients up to factor of e^(- something constant)\n x_ft = torch.fft.rfft2(x)\n\n # Multiply relevant Fourier modes\n out_ft = torch.zeros(batchsize, self.out_channels, x.size(-2), x.size(-1) // 2 + 1, dtype=torch.cfloat,\n device=x.device)\n out_ft[:, :, :self.modes1, :self.modes2] = \\\n self.compl_mul2d(x_ft[:, :, :self.modes1, :self.modes2], self.weights1)\n out_ft[:, :, -self.modes1:, :self.modes2] = \\\n self.compl_mul2d(x_ft[:, :, -self.modes1:, :self.modes2], self.weights2)\n\n # Return to physical space\n x = torch.fft.irfft2(out_ft, s=(x.size(-2), x.size(-1)))\n return x\n\n\nclass FNO2d(nn.Module):\n def __init__(self, modes1, modes2, width):\n super(FNO2d, self).__init__()\n\n \"\"\"\n The overall network. It contains 4 layers of the Fourier layer.\n 1. Lift the input to the desire channel dimension by self.fc0 .\n 2. 4 layers of the integral operators u' = (W + K)(u).\n W defined by self.w; K defined by self.conv .\n 3. Project from the channel space to the output space by self.fc1 and self.fc2 .\n\n input: the solution of the coefficient function and locations (a(x, y), x, y)\n input shape: (batchsize, x=s, y=s, c=3)\n output: the solution \n output shape: (batchsize, x=s, y=s, c=1)\n \"\"\"\n\n self.modes1 = modes1\n self.modes2 = modes2\n self.width = width\n self.padding = 9 # pad the domain if input is non-periodic\n self.fc0 = nn.Linear(3, 128) # input channel is 3: (a(x, y), x, y)\n self.fc1 = nn.Linear(128, self.width)\n\n self.conv0 = SpectralConv2d(self.width, self.width, self.modes1, self.modes2)\n self.conv1 = SpectralConv2d(self.width, self.width, self.modes1, self.modes2)\n self.conv2 = SpectralConv2d(self.width, self.width, self.modes1, self.modes2)\n self.conv3 = SpectralConv2d(self.width, self.width, self.modes1, self.modes2)\n self.w0 = nn.Conv2d(self.width, self.width, 1)\n self.w1 = nn.Conv2d(self.width, self.width, 1)\n self.w2 = nn.Conv2d(self.width, self.width, 1)\n self.w3 = nn.Conv2d(self.width, self.width, 1)\n\n self.fc2 = nn.Linear(self.width, 128)\n self.fc3 = nn.Linear(128, 1)\n\n def forward(self, x):\n\n grid = self.get_grid(x.shape, x.device)\n x = torch.cat((x, grid), dim=-1)\n x = self.fc0(x)\n x = F.gelu(x)\n x = self.fc1(x)\n x = x.permute(0, 3, 1, 2)\n x = F.pad(x, [0, self.padding, 0, self.padding])\n\n x1 = self.conv0(x)\n x2 = self.w0(x)\n x = x1 + x2\n x = F.gelu(x)\n\n x1 = self.conv1(x)\n x2 = self.w1(x)\n x = x1 + x2\n x = F.gelu(x)\n\n x1 = self.conv2(x)\n x2 = self.w2(x)\n x = x1 + x2\n x = F.gelu(x)\n\n x1 = self.conv3(x)\n x2 = self.w3(x)\n x = x1 + x2\n\n x = x[..., :-self.padding, :-self.padding]\n x = x.permute(0, 2, 3, 1)\n x = self.fc2(x)\n x = F.gelu(x)\n x = self.fc3(x)\n\n return x\n\n def get_grid(self, shape, device):\n batchsize, size_x, size_y = shape[0], shape[1], shape[2]\n gridx = torch.tensor(np.linspace(0, 1, size_x), dtype=torch.float)\n gridx = gridx.reshape(1, size_x, 1, 1).repeat([batchsize, 1, size_y, 1])\n gridy = torch.tensor(np.linspace(0, 1, size_y), dtype=torch.float)\n gridy = gridy.reshape(1, 1, size_y, 1).repeat([batchsize, size_x, 1, 1])\n return torch.cat((gridx, gridy), dim=-1).to(device)\n\npretrain = False\nfinetune = not pretrain\n\nTRAIN_PATH = '../data/darcy_s61_N1200.mat'\nTEST_PATH = '../data/darcy_s61_N1200.mat'\n# TRAIN_PATH = '../data/lognormal_N1024_s61.mat'\n# TEST_PATH = '../data/lognormal_N1024_s61.mat'\n# TRAIN_PATH = '../data/piececonst_r241_N1024_smooth1.mat'\n# TEST_PATH = '../data/piececonst_r241_N1024_smooth2.mat'\n\nntrain = 1000\nntest = 1\n\nbatch_size = 1\nlearning_rate = 0.001\n\nepochs = 500\nstep_size = 100\ngamma = 0.5\n\nmodes = 12\nwidth = 32\n\nr = 1\nh = int(((61 - 1)/r) + 1)\ns = h\n\nprint(s)\n\npath = 'PINO_FDM_darcy_N'+str(ntrain)+'_ep' + str(epochs) + '_m' + str(modes) + '_w' + str(width)\npath_model = '../model/'+path\npath_pred = '../pred/'+path+'.mat'\n\nreader = MatReader(TRAIN_PATH)\n# x_train = reader.read_field('coeff')[:ntrain,::r,::r][:,:s,:s]\n# y_train = reader.read_field('sol')[:ntrain,::r,::r][:,:s,:s]\nx_train = reader.read_field('input')[:ntrain,::r,::r][:,:s,:s]\ny_train = reader.read_field('output')[:ntrain,::r,::r][:,:s,:s]\n\nreader.load_file(TEST_PATH)\n# x_test = reader.read_field('coeff')[-ntest:,::r,::r][:,:s,:s]\n# y_test = reader.read_field('sol')[-ntest:,::r,::r][:,:s,:s]\na = 1\nx_test = reader.read_field('input')[-ntest-a:-a,::r,::r][:,:s,:s]\ny_test = reader.read_field('output')[-ntest-a:-a,::r,::r][:,:s,:s]\n\n\nprint(torch.mean(x_train), torch.mean(y_train))\n\n# x_normalizer = UnitGaussianNormalizer(x_train)\n# x_train = x_normalizer.encode(x_train)\n# x_test = x_normalizer.encode(x_test)\n#\n# y_normalizer = UnitGaussianNormalizer(y_train)\n# y_train = y_normalizer.encode(y_train)\n\ngrids = []\ngrids.append(np.linspace(0, 1, s))\ngrids.append(np.linspace(0, 1, s))\ngrid = np.vstack([xx.ravel() for xx in np.meshgrid(*grids)]).T\ngrid = grid.reshape(1,s,s,2)\ngrid = torch.tensor(grid, dtype=torch.float)\n\n\nmyloss = LpLoss(size_average=False)\n\n\ndef FDM_Darcy(u, a, D=1, f=1):\n batchsize = u.size(0)\n size = u.size(1)\n u = u.reshape(batchsize, size, size)\n a = a.reshape(batchsize, size, size)\n dx = D / (size - 1)\n dy = dx\n\n # ux: (batch, size-2, size-2)\n ux = (u[:, 2:, 1:-1] - u[:, :-2, 1:-1]) / (2 * dx)\n uy = (u[:, 1:-1, 2:] - u[:, 1:-1, :-2]) / (2 * dy)\n\n ax = (a[:, 2:, 1:-1] - a[:, :-2, 1:-1]) / (2 * dx)\n ay = (a[:, 1:-1, 2:] - a[:, 1:-1, :-2]) / (2 * dy)\n uxx = (u[:, 2:, 1:-1] -2*u[:,1:-1,1:-1] +u[:, :-2, 1:-1]) / (dx**2)\n uyy = (u[:, 1:-1, 2:] -2*u[:,1:-1,1:-1] +u[:, 1:-1, :-2]) / (dy**2)\n\n a = a[:, 1:-1, 1:-1]\n u = u[:, 1:-1, 1:-1]\n # Du = -(ax*ux + ay*uy + a*uxx + a*uyy)\n\n # inner1 = torch.mean(a*(ux**2 + uy**2), dim=[1,2])\n # inner2 = torch.mean(f*u, dim=[1,2])\n # return 0.5*inner1 - inner2\n\n aux = a * ux\n auy = a * uy\n auxx = (aux[:, 2:, 1:-1] - aux[:, :-2, 1:-1]) / (2 * dx)\n auyy = (auy[:, 1:-1, 2:] - auy[:, 1:-1, :-2]) / (2 * dy)\n Du = - (auxx + auyy)\n\n return Du\n\n\ndef PINO_loss(u, a):\n batchsize = u.size(0)\n size = u.size(1)\n u = u.reshape(batchsize, size, size)\n a = a.reshape(batchsize, size, size)\n lploss = LpLoss(size_average=True)\n\n index_x = torch.cat([torch.tensor(range(0, size)), (size - 1) * torch.ones(size), torch.tensor(range(size-1, 1, -1)),\n torch.zeros(size)], dim=0).long()\n index_y = torch.cat([(size - 1) * torch.ones(size), torch.tensor(range(size-1, 1, -1)), torch.zeros(size),\n torch.tensor(range(0, size))], dim=0).long()\n\n boundary_u = u[:, index_x, index_y]\n truth_u = torch.zeros(boundary_u.shape, device=u.device)\n loss_bd = lploss.abs(boundary_u, truth_u)\n\n Du = FDM_Darcy(u, a)\n f = torch.ones(Du.shape, device=u.device)\n loss_f = lploss(Du, f)\n\n\n # im = (Du-f)[0].detach().cpu().numpy()\n # plt.imshow(im)\n # plt.show()\n\n # loss_f = FDM_Darcy(u, a)\n # loss_f = torch.mean(loss_f)\n return loss_f, loss_bd\n\nerror = np.zeros((epochs, 4))\n# x_normalizer.cuda()\n# y_normalizer.cuda()\ngrid = grid.cuda()\nmollifier = torch.sin(np.pi*grid[...,0]) * torch.sin(np.pi*grid[...,1]) * 0.001\n\nprint(mollifier.shape)\nif pretrain:\n train_loader = torch.utils.data.DataLoader(torch.utils.data.TensorDataset(x_train, y_train), batch_size=batch_size,\n shuffle=True)\n test_loader = torch.utils.data.DataLoader(torch.utils.data.TensorDataset(x_test, y_test), batch_size=batch_size,\n shuffle=False)\n\n model = FNO2d(modes, modes, width).cuda()\n num_param = count_params(model)\n print(num_param)\n optimizer = Adam(model.parameters(), lr=learning_rate, weight_decay=1e-5)\n scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=step_size, gamma=gamma)\n\n for ep in range(epochs):\n model.train()\n t1 = default_timer()\n train_pino = 0.0\n train_l2 = 0.0\n train_loss = 0\n for x, y in train_loader:\n x, y = x.cuda(), y.cuda()\n\n optimizer.zero_grad()\n out = model(x.reshape(batch_size, s, s, 1)).reshape(batch_size, s, s)\n out = out * mollifier\n\n loss_data = myloss(out.view(batch_size,-1), y.view(batch_size,-1))\n loss_f, loss_bd = PINO_loss(out, x)\n pino_loss = loss_f\n pino_loss.backward()\n\n optimizer.step()\n train_l2 += loss_data.item()\n train_pino += pino_loss.item()\n train_loss += torch.tensor([loss_bd, loss_f])\n\n scheduler.step()\n\n model.eval()\n test_l2 = 0.0\n test_pino = 0.0\n with torch.no_grad():\n for x, y in test_loader:\n x, y = x.cuda(), y.cuda()\n\n out = model(x.reshape(batch_size, s, s, 1)).reshape(batch_size, s, s)\n out = out * mollifier\n\n test_l2 += myloss(out.view(batch_size, -1), y.view(batch_size, -1)).item()\n loss_f, loss_bd = PINO_loss(out, x)\n test_pino += loss_f.item() + loss_bd.item()\n\n train_l2 /= ntrain\n test_l2 /= ntest\n train_pino /= ntrain\n test_pino /= ntest\n train_loss /= ntrain\n\n error[ep] = [train_pino, train_l2, test_pino, test_l2]\n\n t2 = default_timer()\n print(ep, t2-t1, train_pino, train_l2, test_pino, test_l2)\n print(train_loss)\n\n # torch.save(model, '../model/IP-dracy-forward')\n\ndef darcy_mask1(x):\n return 1 / (1 + torch.exp(-x)) * 9 + 3\n\ndef darcy_mask2(x):\n x = 1 / (1 + torch.exp(-x))\n x[x>0.5] = 1\n x[x<=0.5] = 0\n # x = torch.tensor(x>0.5, dtype=torch.float)\n return x * 9 + 3\n\ndef total_variance(x):\n return torch.mean(torch.abs(x[...,:-1] - x[...,1:])) + torch.mean(torch.abs(x[...,:-1,:] - x[...,1:,:]))\n\n\nif finetune:\n test_loader = torch.utils.data.DataLoader(torch.utils.data.TensorDataset(x_test, y_test), batch_size=batch_size,\n shuffle=False)\n\n model = torch.load('../model/IP-dracy-forward').cuda()\n num_param = count_params(model)\n print(num_param)\n xout = torch.rand([1,s,s,1], requires_grad=True, device=\"cuda\")\n\n optimizer = Adam([xout], lr=0.1, weight_decay=1e-5)\n scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=2000, gamma=0.5)\n # scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=step_size)\n\n for ep in range(10000):\n model.train()\n t1 = default_timer()\n\n for x, y in test_loader:\n x, y = x.cuda(), y.cuda()\n\n optimizer.zero_grad()\n out_masked = darcy_mask1(xout)\n\n yout = model(out_masked.reshape(batch_size, s, s, 1)).reshape(batch_size, s, s)\n yout = yout * mollifier\n loss_data = myloss(yout.view(batch_size, -1), y.view(batch_size, -1))\n loss_f, loss_bd = PINO_loss(y, out_masked)\n loss_TV = total_variance(xout)\n pino_loss = 0.2 * loss_f + loss_data + 0.05 * loss_TV\n # pino_loss = 0. * loss_f + loss_data + 0.05 * loss_TV\n pino_loss.backward()\n optimizer.step()\n scheduler.step()\n\n out_masked2 = darcy_mask2(xout)\n yout2 = model(out_masked2.reshape(batch_size, s, s, 1)).reshape(batch_size, s, s)\n yout2 = yout2 * mollifier\n testx_l2 = myloss(out_masked.view(batch_size, -1), x.view(batch_size, -1)).item()\n testy_l2 = myloss(yout.view(batch_size, -1), y.view(batch_size, -1)).item()\n\n\n\n t2 = default_timer()\n print(ep, t2 - t1, loss_data.item(), loss_f.item(), testx_l2, testy_l2)\n\n if ep % 2000 == 1:\n # fig, axs = plt.subplots(2, 3, figsize=(8, 8))\n # axs[0,0].imshow(x.reshape(s,s).detach().cpu().numpy())\n # axs[0,1].imshow(out_masked.reshape(s,s).detach().cpu().numpy())\n # axs[0,2].imshow(out_masked2.reshape(s,s).detach().cpu().numpy())\n # axs[1,0].imshow(y.reshape(s,s).detach().cpu().numpy())\n # axs[1,1].imshow(yout.reshape(s,s).detach().cpu().numpy())\n # axs[1,2].imshow(yout2.reshape(s,s).detach().cpu().numpy())\n # plt.show()\n name_tag = 'PINO-'\n plt.imshow(x.reshape(s,s).detach().cpu().numpy())\n plt.savefig(name_tag+'true-input.pdf',bbox_inches='tight')\n plt.imshow(out_masked.reshape(s,s).detach().cpu().numpy())\n plt.savefig(name_tag+'raw-input.pdf',bbox_inches='tight')\n plt.imshow(out_masked2.reshape(s,s).detach().cpu().numpy())\n plt.savefig(name_tag+'clip-input.pdf',bbox_inches='tight')\n\n plt.imshow(y.reshape(s,s).detach().cpu().numpy())\n plt.savefig(name_tag+'true-output.pdf',bbox_inches='tight')\n plt.imshow(yout.reshape(s,s).detach().cpu().numpy())\n plt.savefig(name_tag+'raw-output.pdf',bbox_inches='tight')\n plt.imshow(yout.reshape(s,s).detach().cpu().numpy())\n plt.savefig(name_tag+'clip-output.pdf',bbox_inches='tight')\n\n # scipy.io.savemat('../pred/IP-darcy-forward.mat', mdict={'input_truth': x.reshape(s,s).detach().cpu().numpy(),\n # 'input_pred': out_masked.reshape(s,s).detach().cpu().numpy(),\n # 'output_truth': y.reshape(s,s).detach().cpu().numpy(),\n # 'output_pred': yout.reshape(s,s).detach().cpu().numpy()})\n\n" }, { "path": "inverse-darcy.py", "content": "\nimport numpy as np\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport scipy.io\nimport matplotlib.pyplot as plt\nfrom timeit import default_timer\nfrom torch.optim import Adam\n\nfrom train_utils.datasets import MatReader\nfrom train_utils.losses import LpLoss\nfrom train_utils.utils import count_params\n\ntorch.manual_seed(0)\nnp.random.seed(0)\n\n\n################################################################\n# fourier layer\n################################################################\nclass SpectralConv2d(nn.Module):\n def __init__(self, in_channels, out_channels, modes1, modes2):\n super(SpectralConv2d, self).__init__()\n\n \"\"\"\n 2D Fourier layer. It does FFT, linear transform, and Inverse FFT. \n \"\"\"\n\n self.in_channels = in_channels\n self.out_channels = out_channels\n self.modes1 = modes1 # Number of Fourier modes to multiply, at most floor(N/2) + 1\n self.modes2 = modes2\n\n self.scale = (1 / (in_channels * out_channels))\n self.weights1 = nn.Parameter(\n self.scale * torch.rand(in_channels, out_channels, self.modes1, self.modes2, dtype=torch.cfloat))\n self.weights2 = nn.Parameter(\n self.scale * torch.rand(in_channels, out_channels, self.modes1, self.modes2, dtype=torch.cfloat))\n\n # Complex multiplication\n def compl_mul2d(self, input, weights):\n # (batch, in_channel, x,y ), (in_channel, out_channel, x,y) -> (batch, out_channel, x,y)\n return torch.einsum(\"bixy,ioxy->boxy\", input, weights)\n\n def forward(self, x):\n batchsize = x.shape[0]\n # Compute Fourier coeffcients up to factor of e^(- something constant)\n x_ft = torch.fft.rfft2(x)\n\n # Multiply relevant Fourier modes\n out_ft = torch.zeros(batchsize, self.out_channels, x.size(-2), x.size(-1) // 2 + 1, dtype=torch.cfloat,\n device=x.device)\n out_ft[:, :, :self.modes1, :self.modes2] = \\\n self.compl_mul2d(x_ft[:, :, :self.modes1, :self.modes2], self.weights1)\n out_ft[:, :, -self.modes1:, :self.modes2] = \\\n self.compl_mul2d(x_ft[:, :, -self.modes1:, :self.modes2], self.weights2)\n\n # Return to physical space\n x = torch.fft.irfft2(out_ft, s=(x.size(-2), x.size(-1)))\n return x\n\n\nclass FNO2d(nn.Module):\n def __init__(self, modes1, modes2, width):\n super(FNO2d, self).__init__()\n\n \"\"\"\n The overall network. It contains 4 layers of the Fourier layer.\n 1. Lift the input to the desire channel dimension by self.fc0 .\n 2. 4 layers of the integral operators u' = (W + K)(u).\n W defined by self.w; K defined by self.conv .\n 3. Project from the channel space to the output space by self.fc1 and self.fc2 .\n\n input: the solution of the coefficient function and locations (a(x, y), x, y)\n input shape: (batchsize, x=s, y=s, c=3)\n output: the solution \n output shape: (batchsize, x=s, y=s, c=1)\n \"\"\"\n\n self.modes1 = modes1\n self.modes2 = modes2\n self.width = width\n self.padding = 9 # pad the domain if input is non-periodic\n self.fc0 = nn.Linear(3, 128) # input channel is 3: (a(x, y), x, y)\n self.fc1 = nn.Linear(128, self.width)\n\n self.conv0 = SpectralConv2d(self.width, self.width, self.modes1, self.modes2)\n self.conv1 = SpectralConv2d(self.width, self.width, self.modes1, self.modes2)\n self.conv2 = SpectralConv2d(self.width, self.width, self.modes1, self.modes2)\n self.conv3 = SpectralConv2d(self.width, self.width, self.modes1, self.modes2)\n self.w0 = nn.Conv2d(self.width, self.width, 1)\n self.w1 = nn.Conv2d(self.width, self.width, 1)\n self.w2 = nn.Conv2d(self.width, self.width, 1)\n self.w3 = nn.Conv2d(self.width, self.width, 1)\n\n self.fc2 = nn.Linear(self.width, 128)\n self.fc3 = nn.Linear(128, 1)\n\n def forward(self, x):\n grid = self.get_grid(x.shape, x.device)\n x = torch.cat((x, grid), dim=-1)\n x = self.fc0(x)\n x = F.gelu(x)\n x = self.fc1(x)\n x = x.permute(0, 3, 1, 2)\n x = F.pad(x, [0, self.padding, 0, self.padding])\n\n x1 = self.conv0(x)\n x2 = self.w0(x)\n x = x1 + x2\n x = F.gelu(x)\n\n x1 = self.conv1(x)\n x2 = self.w1(x)\n x = x1 + x2\n x = F.gelu(x)\n\n x1 = self.conv2(x)\n x2 = self.w2(x)\n x = x1 + x2\n x = F.gelu(x)\n\n x1 = self.conv3(x)\n x2 = self.w3(x)\n x = x1 + x2\n\n x = x[..., :-self.padding, :-self.padding]\n x = x.permute(0, 2, 3, 1)\n x = self.fc2(x)\n x = F.gelu(x)\n x = self.fc3(x)\n\n return x\n\n def get_grid(self, shape, device):\n batchsize, size_x, size_y = shape[0], shape[1], shape[2]\n gridx = torch.tensor(np.linspace(0, 1, size_x), dtype=torch.float)\n gridx = gridx.reshape(1, size_x, 1, 1).repeat([batchsize, 1, size_y, 1])\n gridy = torch.tensor(np.linspace(0, 1, size_y), dtype=torch.float)\n gridy = gridy.reshape(1, 1, size_y, 1).repeat([batchsize, size_x, 1, 1])\n return torch.cat((gridx, gridy), dim=-1).to(device)\n\n\npretrain = False\nfinetune = not pretrain\n\nTRAIN_PATH = '../data/darcy_s61_N1200.mat'\nTEST_PATH = '../data/darcy_s61_N1200.mat'\n# TRAIN_PATH = '../data/lognormal_N1024_s61.mat'\n# TEST_PATH = '../data/lognormal_N1024_s61.mat'\n# TRAIN_PATH = '../data/piececonst_r241_N1024_smooth1.mat'\n# TEST_PATH = '../data/piececonst_r241_N1024_smooth2.mat'\n\nntrain = 1000\nntest = 1\n\nbatch_size = 1\nlearning_rate = 0.001\n\nepochs = 100\nstep_size = 100\ngamma = 0.5\n\nmodes = 12\nwidth = 32\n\nr = 1\nh = int(((61 - 1)/r) + 1)\ns = h\n\nprint(s)\n\npath = 'PINO_FDM_darcy_N'+str(ntrain)+'_ep' + str(epochs) + '_m' + str(modes) + '_w' + str(width)\npath_model = '../model/'+path\npath_pred = '../pred/'+path+'.mat'\n\nreader = MatReader(TRAIN_PATH)\n# x_train = reader.read_field('coeff')[:ntrain,::r,::r][:,:s,:s]\n# y_train = reader.read_field('sol')[:ntrain,::r,::r][:,:s,:s]\nx_train = reader.read_field('input')[:ntrain,::r,::r][:,:s,:s]\ny_train = reader.read_field('output')[:ntrain,::r,::r][:,:s,:s]\n\nreader.load_file(TEST_PATH)\n# x_test = reader.read_field('coeff')[-ntest:,::r,::r][:,:s,:s]\n# y_test = reader.read_field('sol')[-ntest:,::r,::r][:,:s,:s]\na = 1\nx_test = reader.read_field('input')[-ntest-a:-a,::r,::r][:,:s,:s]\ny_test = reader.read_field('output')[-ntest-a:-a,::r,::r][:,:s,:s]\n\n\nprint(torch.mean(x_train), torch.mean(y_train))\n\n# x_normalizer = UnitGaussianNormalizer(x_train)\n# x_train = x_normalizer.encode(x_train)\n# x_test = x_normalizer.encode(x_test)\n#\n# y_normalizer = UnitGaussianNormalizer(y_train)\n# y_train = y_normalizer.encode(y_train)\n\ngrids = []\ngrids.append(np.linspace(0, 1, s))\ngrids.append(np.linspace(0, 1, s))\ngrid = np.vstack([xx.ravel() for xx in np.meshgrid(*grids)]).T\ngrid = grid.reshape(1,s,s,2)\ngrid = torch.tensor(grid, dtype=torch.float)\n\nmyloss = LpLoss(size_average=False)\n\n\ndef FDM_Darcy(u, a, D=1, f=1):\n batchsize = u.size(0)\n size = u.size(1)\n u = u.reshape(batchsize, size, size)\n a = a.reshape(batchsize, size, size)\n dx = D / (size - 1)\n dy = dx\n\n # ux: (batch, size-2, size-2)\n ux = (u[:, 2:, 1:-1] - u[:, :-2, 1:-1]) / (2 * dx)\n uy = (u[:, 1:-1, 2:] - u[:, 1:-1, :-2]) / (2 * dy)\n\n ax = (a[:, 2:, 1:-1] - a[:, :-2, 1:-1]) / (2 * dx)\n ay = (a[:, 1:-1, 2:] - a[:, 1:-1, :-2]) / (2 * dy)\n uxx = (u[:, 2:, 1:-1] -2*u[:,1:-1,1:-1] +u[:, :-2, 1:-1]) / (dx**2)\n uyy = (u[:, 1:-1, 2:] -2*u[:,1:-1,1:-1] +u[:, 1:-1, :-2]) / (dy**2)\n\n a = a[:, 1:-1, 1:-1]\n u = u[:, 1:-1, 1:-1]\n # Du = -(ax*ux + ay*uy + a*uxx + a*uyy)\n\n # inner1 = torch.mean(a*(ux**2 + uy**2), dim=[1,2])\n # inner2 = torch.mean(f*u, dim=[1,2])\n # return 0.5*inner1 - inner2\n\n aux = a * ux\n auy = a * uy\n auxx = (aux[:, 2:, 1:-1] - aux[:, :-2, 1:-1]) / (2 * dx)\n auyy = (auy[:, 1:-1, 2:] - auy[:, 1:-1, :-2]) / (2 * dy)\n Du = - (auxx + auyy)\n\n return Du\n\n\ndef PINO_loss(u, a):\n batchsize = u.size(0)\n size = u.size(1)\n u = u.reshape(batchsize, size, size)\n a = a.reshape(batchsize, size, size)\n lploss = LpLoss(size_average=True)\n\n index_x = torch.cat([torch.tensor(range(0, size)), (size - 1) * torch.ones(size), torch.tensor(range(size-1, 1, -1)),\n torch.zeros(size)], dim=0).long()\n index_y = torch.cat([(size - 1) * torch.ones(size), torch.tensor(range(size-1, 1, -1)), torch.zeros(size),\n torch.tensor(range(0, size))], dim=0).long()\n\n boundary_u = u[:, index_x, index_y]\n truth_u = torch.zeros(boundary_u.shape, device=u.device)\n loss_bd = lploss.abs(boundary_u, truth_u)\n\n Du = FDM_Darcy(u, a)\n f = torch.ones(Du.shape, device=u.device)\n loss_f = lploss(Du, f)\n\n # im = (Du-f)[0].detach().cpu().numpy()\n # plt.imshow(im)\n # plt.show()\n\n # loss_f = FDM_Darcy(u, a)\n # loss_f = torch.mean(loss_f)\n return loss_f, loss_bd\n\nerror = np.zeros((epochs, 4))\n# x_normalizer.cuda()\n# y_normalizer.cuda()\ngrid = grid.cuda()\nmollifier = torch.sin(np.pi*grid[...,0]) * torch.sin(np.pi*grid[...,1]) * 0.001\n\ndef darcy_mask1(x):\n return 1 / (1 + torch.exp(-x)) * 9 + 3\n\ndef darcy_mask2(x):\n x = 1 / (1 + torch.exp(-x))\n x[x>0.5] = 1\n x[x<=0.5] = 0\n # x = torch.tensor(x>0.5, dtype=torch.float)\n return x * 9 + 3\n\ndef total_variance(x):\n return torch.mean(torch.abs(x[...,:-1] - x[...,1:])) + torch.mean(torch.abs(x[...,:-1,:] - x[...,1:,:]))\n\nif pretrain:\n train_loader = torch.utils.data.DataLoader(torch.utils.data.TensorDataset(x_train, y_train), batch_size=batch_size,\n shuffle=True)\n test_loader = torch.utils.data.DataLoader(torch.utils.data.TensorDataset(x_test, y_test), batch_size=batch_size,\n shuffle=False)\n\n model = FNO2d(modes, modes, width).cuda()\n num_param = count_params(model)\n print(num_param)\n\n optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate, weight_decay=1e-5)\n scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=step_size, gamma=gamma)\n # scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=step_size)\n\n for ep in range(epochs):\n model.train()\n t1 = default_timer()\n train_f = 0.0\n train_l2 = 0.0\n train_TV = 0.0\n for x, y in train_loader:\n x, y = x.cuda(), y.cuda()\n\n optimizer.zero_grad()\n xout = model(y.reshape(batch_size, s, s, 1)).reshape(batch_size, s, s)\n xout = darcy_mask1(xout)\n\n loss_data = myloss(xout.view(batch_size,-1), x.view(batch_size,-1))\n loss_f, loss_bd = PINO_loss(y, xout)\n loss_TV = total_variance(xout)\n pino_loss = 0.2*loss_f + loss_data + 0.01*loss_TV\n pino_loss.backward()\n\n optimizer.step()\n train_l2 += loss_data.item()\n train_f += loss_f.item()\n train_TV += loss_TV.item()\n\n\n scheduler.step()\n\n model.eval()\n test_l2 = 0.0\n test_pino = 0.0\n with torch.no_grad():\n for x, y in test_loader:\n x, y = x.cuda(), y.cuda()\n\n xout = model(y.reshape(batch_size, s, s, 1)).reshape(batch_size, s, s)\n xout = darcy_mask1(xout)\n\n test_l2 += myloss(xout.view(batch_size, -1), x.view(batch_size, -1)).item()\n\n\n train_l2 /= ntrain\n test_l2 /= ntest\n train_f /= ntrain\n test_pino /= ntest\n train_TV /= ntrain\n\n error[ep] = [train_f, train_l2, test_pino, test_l2]\n\n t2 = default_timer()\n print(ep, t2-t1, train_f, train_TV, train_l2, test_l2)\n\n torch.save(model, '../model/IP-dracy-inverse')\n\nif finetune:\n test_loader = torch.utils.data.DataLoader(torch.utils.data.TensorDataset(x_test, y_test), batch_size=batch_size,\n shuffle=False)\n\n model = torch.load('../model/IP-dracy-inverse').cuda()\n # model = FNO2d(modes, modes, width).cuda()\n model_forward = torch.load('../model/IP-dracy-forward').cuda()\n num_param = count_params(model)\n print(num_param)\n\n optimizer = torch.optim.Adam(model.parameters(), lr=0.001, weight_decay=1e-5)\n scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=2000, gamma=0.5)\n # scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=step_size)\n\n for ep in range(10000):\n model.train()\n t1 = default_timer()\n\n for x, y in test_loader:\n x, y = x.cuda(), y.cuda()\n\n optimizer.zero_grad()\n\n xout = model(y.reshape(batch_size, s, s, 1)).reshape(batch_size, s, s)\n xout1 = darcy_mask1(xout)\n loss_TV = total_variance(xout)\n\n loss_f, loss_bd = PINO_loss(y, xout1)\n pino_loss = loss_f + 0.05*loss_TV\n pino_loss.backward()\n optimizer.step()\n scheduler.step()\n\n xout2 = darcy_mask2(xout)\n testx_l2 = myloss(xout1.view(batch_size, -1), x.view(batch_size, -1)).item()\n\n t2 = default_timer()\n print(ep, t2 - t1, loss_f.item(), testx_l2)\n\n if ep % 1000 == 0:\n yout1 = model_forward(xout1.reshape(batch_size, s, s, 1)).reshape(batch_size, s, s) * mollifier\n yout2 = model_forward(xout2.reshape(batch_size, s, s, 1)).reshape(batch_size, s, s) * mollifier\n\n fig, axs = plt.subplots(2, 3, figsize=(8, 8))\n axs[0,0].imshow(x.reshape(s,s).detach().cpu().numpy())\n axs[0,1].imshow(xout1.reshape(s,s).detach().cpu().numpy())\n axs[0,2].imshow(xout2.reshape(s,s).detach().cpu().numpy())\n axs[1,0].imshow(y.reshape(s,s).detach().cpu().numpy())\n axs[1,1].imshow(yout1.reshape(s,s).detach().cpu().numpy())\n axs[1,2].imshow(yout2.reshape(s,s).detach().cpu().numpy())\n plt.show()\n\n scipy.io.savemat('../pred/IP-darcy-backward.mat', mdict={'input_truth': x.reshape(s,s).detach().cpu().numpy(),\n 'input_pred': xout1.reshape(s,s).detach().cpu().numpy(),\n 'output_truth': y.reshape(s,s).detach().cpu().numpy(),\n 'output_pred': yout1.reshape(s,s).detach().cpu().numpy()})\n\n\n\n" }, { "path": "models/FCN.py", "content": "import torch.nn as nn\n\n\ndef linear_block(in_channel, out_channel):\n block = nn.Sequential(\n nn.Linear(in_channel, out_channel),\n nn.Tanh()\n )\n return block\n\n\nclass FCNet(nn.Module):\n '''\n Fully connected layers with Tanh as nonlinearity\n Reproduced from PINNs Burger equation\n '''\n\n def __init__(self, layers=[2, 10, 1]):\n super(FCNet, self).__init__()\n\n fc_list = [linear_block(in_size, out_size)\n for in_size, out_size in zip(layers, layers[1:-1])]\n fc_list.append(nn.Linear(layers[-2], layers[-1]))\n self.fc = nn.Sequential(*fc_list)\n\n def forward(self, x):\n return self.fc(x)\n\n\nclass DenseNet(nn.Module):\n def __init__(self, layers, nonlinearity, out_nonlinearity=None, normalize=False):\n super(DenseNet, self).__init__()\n\n self.n_layers = len(layers) - 1\n assert self.n_layers >= 1\n if isinstance(nonlinearity, str):\n if nonlinearity == 'tanh':\n nonlinearity = nn.Tanh\n elif nonlinearity == 'relu':\n nonlinearity == nn.ReLU\n else:\n raise ValueError(f'{nonlinearity} is not supported')\n self.layers = nn.ModuleList()\n\n for j in range(self.n_layers):\n self.layers.append(nn.Linear(layers[j], layers[j+1]))\n\n if j != self.n_layers - 1:\n if normalize:\n self.layers.append(nn.BatchNorm1d(layers[j+1]))\n\n self.layers.append(nonlinearity())\n\n if out_nonlinearity is not None:\n self.layers.append(out_nonlinearity())\n\n def forward(self, x):\n for _, l in enumerate(self.layers):\n x = l(x)\n\n return x\n\n\n" }, { "path": "models/__init__.py", "content": "from .FCN import FCNet\nfrom .fourier1d import FNO1d\nfrom .fourier2d import FNO2d\nfrom .fourier3d import FNO3d\n" }, { "path": "models/basics.py", "content": "import numpy as np\n\nimport torch\nimport torch.nn as nn\n\n\n@torch.jit.script\ndef compl_mul1d(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:\n # (batch, in_channel, x ), (in_channel, out_channel, x) -> (batch, out_channel, x)\n res = torch.einsum(\"bix,iox->box\", a, b)\n return res\n\n\n@torch.jit.script\ndef compl_mul2d(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:\n # (batch, in_channel, x,y,t ), (in_channel, out_channel, x,y,t) -> (batch, out_channel, x,y,t)\n res = torch.einsum(\"bixy,ioxy->boxy\", a, b)\n return res\n\n\n@torch.jit.script\ndef compl_mul3d(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:\n res = torch.einsum(\"bixyz,ioxyz->boxyz\", a, b)\n return res\n\n################################################################\n# 1d fourier layer\n################################################################\n\n\nclass SpectralConv1d(nn.Module):\n def __init__(self, in_channels, out_channels, modes1):\n super(SpectralConv1d, self).__init__()\n\n \"\"\"\n 1D Fourier layer. It does FFT, linear transform, and Inverse FFT. \n \"\"\"\n\n self.in_channels = in_channels\n self.out_channels = out_channels\n # Number of Fourier modes to multiply, at most floor(N/2) + 1\n self.modes1 = modes1\n\n self.scale = (1 / (in_channels*out_channels))\n self.weights1 = nn.Parameter(\n self.scale * torch.rand(in_channels, out_channels, self.modes1, dtype=torch.cfloat))\n\n def forward(self, x):\n batchsize = x.shape[0]\n # Compute Fourier coeffcients up to factor of e^(- something constant)\n x_ft = torch.fft.rfftn(x, dim=[2])\n\n # Multiply relevant Fourier modes\n out_ft = torch.zeros(batchsize, self.in_channels, x.size(-1)//2 + 1, device=x.device, dtype=torch.cfloat)\n out_ft[:, :, :self.modes1] = compl_mul1d(x_ft[:, :, :self.modes1], self.weights1)\n\n # Return to physical space\n x = torch.fft.irfftn(out_ft, s=[x.size(-1)], dim=[2])\n return x\n\n################################################################\n# 2d fourier layer\n################################################################\n\n\nclass SpectralConv2d(nn.Module):\n def __init__(self, in_channels, out_channels, modes1, modes2):\n super(SpectralConv2d, self).__init__()\n self.in_channels = in_channels\n self.out_channels = out_channels\n # Number of Fourier modes to multiply, at most floor(N/2) + 1\n self.modes1 = modes1\n self.modes2 = modes2\n\n self.scale = (1 / (in_channels * out_channels))\n self.weights1 = nn.Parameter(\n self.scale * torch.rand(in_channels, out_channels, self.modes1, self.modes2, dtype=torch.cfloat))\n self.weights2 = nn.Parameter(\n self.scale * torch.rand(in_channels, out_channels, self.modes1, self.modes2, dtype=torch.cfloat))\n\n def forward(self, x):\n batchsize = x.shape[0]\n size1 = x.shape[-2]\n size2 = x.shape[-1]\n # Compute Fourier coeffcients up to factor of e^(- something constant)\n x_ft = torch.fft.rfftn(x, dim=[2, 3])\n\n # Multiply relevant Fourier modes\n out_ft = torch.zeros(batchsize, self.out_channels, x.size(-2), x.size(-1) // 2 + 1, device=x.device,\n dtype=torch.cfloat)\n out_ft[:, :, :self.modes1, :self.modes2] = \\\n compl_mul2d(x_ft[:, :, :self.modes1, :self.modes2], self.weights1)\n out_ft[:, :, -self.modes1:, :self.modes2] = \\\n compl_mul2d(x_ft[:, :, -self.modes1:, :self.modes2], self.weights2)\n\n # Return to physical space\n x = torch.fft.irfftn(out_ft, s=(x.size(-2), x.size(-1)), dim=[2, 3])\n return x\n\n\nclass SpectralConv3d(nn.Module):\n def __init__(self, in_channels, out_channels, modes1, modes2, modes3):\n super(SpectralConv3d, self).__init__()\n self.in_channels = in_channels\n self.out_channels = out_channels\n self.modes1 = modes1 #Number of Fourier modes to multiply, at most floor(N/2) + 1\n self.modes2 = modes2\n self.modes3 = modes3\n\n self.scale = (1 / (in_channels * out_channels))\n self.weights1 = nn.Parameter(self.scale * torch.rand(in_channels, out_channels, self.modes1, self.modes2, self.modes3, dtype=torch.cfloat))\n self.weights2 = nn.Parameter(self.scale * torch.rand(in_channels, out_channels, self.modes1, self.modes2, self.modes3, dtype=torch.cfloat))\n self.weights3 = nn.Parameter(self.scale * torch.rand(in_channels, out_channels, self.modes1, self.modes2, self.modes3, dtype=torch.cfloat))\n self.weights4 = nn.Parameter(self.scale * torch.rand(in_channels, out_channels, self.modes1, self.modes2, self.modes3, dtype=torch.cfloat))\n\n def forward(self, x):\n batchsize = x.shape[0]\n # Compute Fourier coeffcients up to factor of e^(- something constant)\n x_ft = torch.fft.rfftn(x, dim=[2,3,4])\n \n z_dim = min(x_ft.shape[4], self.modes3)\n \n # Multiply relevant Fourier modes\n out_ft = torch.zeros(batchsize, self.out_channels, x_ft.shape[2], x_ft.shape[3], self.modes3, device=x.device, dtype=torch.cfloat)\n \n # if x_ft.shape[4] > self.modes3, truncate; if x_ft.shape[4] < self.modes3, add zero padding \n coeff = torch.zeros(batchsize, self.in_channels, self.modes1, self.modes2, self.modes3, device=x.device, dtype=torch.cfloat) \n coeff[..., :z_dim] = x_ft[:, :, :self.modes1, :self.modes2, :z_dim]\n out_ft[:, :, :self.modes1, :self.modes2, :] = compl_mul3d(coeff, self.weights1)\n \n coeff = torch.zeros(batchsize, self.in_channels, self.modes1, self.modes2, self.modes3, device=x.device, dtype=torch.cfloat) \n coeff[..., :z_dim] = x_ft[:, :, -self.modes1:, :self.modes2, :z_dim]\n out_ft[:, :, -self.modes1:, :self.modes2, :] = compl_mul3d(coeff, self.weights2)\n \n coeff = torch.zeros(batchsize, self.in_channels, self.modes1, self.modes2, self.modes3, device=x.device, dtype=torch.cfloat) \n coeff[..., :z_dim] = x_ft[:, :, :self.modes1, -self.modes2:, :z_dim]\n out_ft[:, :, :self.modes1, -self.modes2:, :] = compl_mul3d(coeff, self.weights3)\n \n coeff = torch.zeros(batchsize, self.in_channels, self.modes1, self.modes2, self.modes3, device=x.device, dtype=torch.cfloat) \n coeff[..., :z_dim] = x_ft[:, :, -self.modes1:, -self.modes2:, :z_dim]\n out_ft[:, :, -self.modes1:, -self.modes2:, :] = compl_mul3d(coeff, self.weights4)\n\n #Return to physical space\n x = torch.fft.irfftn(out_ft, s=(x.size(2), x.size(3), x.size(4)), dim=[2,3,4])\n return x\n\n\nclass FourierBlock(nn.Module):\n def __init__(self, in_channels, out_channels, modes1, modes2, modes3, act='tanh'):\n super(FourierBlock, self).__init__()\n self.in_channel = in_channels\n self.out_channel = out_channels\n self.speconv = SpectralConv3d(in_channels, out_channels, modes1, modes2, modes3)\n self.linear = nn.Conv1d(in_channels, out_channels, 1)\n if act == 'tanh':\n self.act = torch.tanh_\n elif act == 'gelu':\n self.act = nn.GELU\n elif act == 'none':\n self.act = None\n else:\n raise ValueError(f'{act} is not supported')\n\n def forward(self, x):\n '''\n input x: (batchsize, channel width, x_grid, y_grid, t_grid)\n '''\n x1 = self.speconv(x)\n x2 = self.linear(x.view(x.shape[0], self.in_channel, -1))\n out = x1 + x2.view(x.shape[0], self.out_channel, x.shape[2], x.shape[3], x.shape[4])\n if self.act is not None:\n out = self.act(out)\n return out\n\n\n" }, { "path": "models/core.py", "content": "import torch\nimport torch.nn as nn\nimport tltorch\n\n\n@torch.jit.script\ndef contract_1D(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: \n res = torch.einsum(\"bix,iox->box\", a, b)\n return res\n\n\n@torch.jit.script\ndef contract_2D(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: \n res = torch.einsum(\"bixy,ioxy->boxy\", a, b)\n return res\n\n\n@torch.jit.script\ndef contract_3D(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: \n res = torch.einsum(\"bixyz,ioxyz->boxyz\", a, b)\n return res\n\n\nclass FactorizedSpectralConv3d(nn.Module):\n def __init__(self, in_channels, out_channels, modes_height, modes_width, modes_depth, n_layers=1, bias=True, scale='auto',\n fft_norm='backward', mlp=False,\n rank=0.5, factorization='cp', fixed_rank_modes=None, decomposition_kwargs=dict(), **kwargs):\n super().__init__()\n\n self.in_channels = in_channels\n self.out_channels = out_channels\n self.modes_height = modes_height\n self.modes_width = modes_width\n self.modes_depth = modes_depth\n self.rank = rank\n self.factorization = factorization\n self.n_layers = n_layers\n self.fft_norm = fft_norm\n if mlp:\n raise NotImplementedError()\n else:\n self.mlp = None\n\n if scale == 'auto':\n scale = (1 / (in_channels * out_channels))\n\n if isinstance(fixed_rank_modes, bool):\n if fixed_rank_modes:\n fixed_rank_modes=[0]\n else:\n fixed_rank_modes=None\n\n if factorization is None:\n self.weight = nn.Parameter(scale * torch.randn(4*n_layers, in_channels, out_channels, self.modes_height, self.modes_width, self.modes_depth,\n dtype=torch.cfloat))\n self._get_weight = self._get_weight_dense\n else:\n self.weight = tltorch.FactorizedTensor.new((4*n_layers, in_channels, out_channels, self.modes_height, self.modes_width, self.modes_depth),\n rank=self.rank, factorization=factorization, \n dtype=torch.cfloat, fixed_rank_modes=fixed_rank_modes,\n **decomposition_kwargs)\n self.weight = self.weight.normal_(0, scale)\n self._get_weight = self._get_weight_factorized\n\n if bias:\n self.bias = nn.Parameter(scale * torch.randn(self.out_channels, 1, 1, 1))\n else:\n self.bias = 0\n\n def _get_weight_factorized(self, layer_index, corner_index):\n \"\"\"Get the weights corresponding to a particular layer,\n corner of the Fourier coefficient (top=0 or bottom=1) -- corresponding to lower frequencies\n and complex_index (real=0 or imaginary=1)\n \"\"\"\n return self.weight()[4*layer_index + corner_index, :, :, :, :, :].to_tensor().contiguous()\n\n def _get_weight_dense(self, layer_index, corner_index):\n \"\"\"Get the weights corresponding to a particular layer,\n corner of the Fourier coefficient (top=0 or bottom=1) -- corresponding to lower frequencies\n and complex_index (real=0 or imaginary=1)\n \"\"\"\n return self.weight[4*layer_index + corner_index, :, :, :, :, :]\n\n def forward(self, x, indices=0):\n with torch.autocast(device_type='cuda', enabled=False):\n batchsize, channels, height, width, depth = x.shape\n dtype = x.dtype\n # out_fft = torch.zeros(x.shape, device=x.device) \n\n #Compute Fourier coeffcients \n x = torch.fft.rfftn(x.float(), norm=self.fft_norm, dim=[-3, -2, -1])\n\n # Multiply relevant Fourier modes \n # x = torch.view_as_real(x)\n # The output will be of size (batch_size, self.out_channels, x.size(-2), x.size(-1)//2 + 1)\n out_fft = torch.zeros([batchsize, self.out_channels, height, width, depth//2 + 1], device=x.device, dtype=torch.cfloat)\n\n out_fft[:, :, :self.modes_height, :self.modes_width, :self.modes_depth] = contract_3D(\n x[:, :, :self.modes_height, :self.modes_width, :self.modes_depth], self._get_weight(indices, 0))\n out_fft[:, :, -self.modes_height:, :self.modes_width, :self.modes_depth] = contract_3D(\n x[:, :, -self.modes_height:, :self.modes_width, :self.modes_depth], self._get_weight(indices, 1))\n out_fft[:, :, self.modes_height:, -self.modes_width:, :self.modes_depth] = contract_3D(\n x[:, :, self.modes_height:, -self.modes_width:, :self.modes_depth], self._get_weight(indices, 2))\n out_fft[:, :, -self.modes_height:, -self.modes_width:, :self.modes_depth] = contract_3D(\n x[:, :, -self.modes_height:, -self.modes_width:, :self.modes_depth], self._get_weight(indices, 3))\n\n # out_size = (int(height*super_res), int(width*super_res))\n x = torch.fft.irfftn(out_fft, s=(height, width, depth), norm=self.fft_norm).type(dtype) #(x.size(-2), x.size(-1))) +\n x = x + self.bias\n\n if self.mlp is not None:\n x = self.mlp(x)\n\n return x\n\n def get_conv(self, indices):\n \"\"\"Returns a sub-convolutional layer from the joint parametrize main-convolution\n The parametrization of sub-convolutional layers is shared with the main one.\n \"\"\"\n if self.n_layers == 1:\n raise ValueError('A single convolution is parametrized, directly use the main class.')\n \n return SubConv(self, indices)\n \n def __getitem__(self, indices):\n return self.get_conv(indices)\n\n\nclass FactorizedSpectralConv2d(nn.Module):\n def __init__(self, in_channels, out_channels, modes_height, modes_width, n_layers=1, bias=True, scale='auto',\n fft_norm='backward',\n rank=0.5, factorization='cp', fixed_rank_modes=None, decomposition_kwargs=dict(), **kwargs):\n super().__init__()\n\n self.in_channels = in_channels\n self.out_channels = out_channels\n self.modes_height = modes_height\n self.modes_width = modes_width\n self.rank = rank\n self.factorization = factorization\n self.n_layers = n_layers\n self.fft_norm = fft_norm\n\n if scale == 'auto':\n scale = (1 / (in_channels * out_channels))\n\n if isinstance(fixed_rank_modes, bool):\n if fixed_rank_modes:\n fixed_rank_modes=[0]\n else:\n fixed_rank_modes=None\n\n if factorization is None:\n self.weight = nn.Parameter(scale * torch.randn(2*n_layers, in_channels, out_channels, self.modes_height, self.modes_width,\n dtype=torch.cfloat))\n self._get_weight = self._get_weight_dense\n else:\n self.weight = tltorch.FactorizedTensor.new((2*n_layers, in_channels, out_channels, self.modes_height, self.modes_width),\n rank=self.rank, factorization=factorization, \n dtype=torch.cfloat, fixed_rank_modes=fixed_rank_modes,\n **decomposition_kwargs)\n self.weight = self.weight.normal_(0, scale)\n self._get_weight = self._get_weight_factorized\n\n if bias:\n self.bias = nn.Parameter(scale * torch.randn(self.out_channels, 1, 1))\n else:\n self.bias = 0\n\n def _get_weight_factorized(self, layer_index, corner_index):\n \"\"\"Get the weights corresponding to a particular layer,\n corner of the Fourier coefficient (top=0 or bottom=1) -- corresponding to lower frequencies\n and complex_index (real=0 or imaginary=1)\n \"\"\"\n return self.weight()[2*layer_index + corner_index, :, :, :, : ].to_tensor().contiguous()\n\n def _get_weight_dense(self, layer_index, corner_index):\n \"\"\"Get the weights corresponding to a particular layer,\n corner of the Fourier coefficient (top=0 or bottom=1) -- corresponding to lower frequencies\n and complex_index (real=0 or imaginary=1)\n \"\"\"\n return self.weight[2*layer_index + corner_index, :, :, :, :]\n\n def forward(self, x, indices=0, super_res=1):\n with torch.autocast(device_type='cuda', enabled=False):\n batchsize, channels, height, width = x.shape\n dtype = x.dtype\n # out_fft = torch.zeros(x.shape, device=x.device) \n\n #Compute Fourier coeffcients \n x = torch.fft.rfft2(x.float(), norm=self.fft_norm)\n\n # Multiply relevant Fourier modes \n # x = torch.view_as_real(x)\n # The output will be of size (batch_size, self.out_channels, x.size(-2), x.size(-1)//2 + 1)\n out_fft = torch.zeros([batchsize, self.out_channels, height, width//2 + 1], device=x.device, dtype=torch.cfloat)\n\n # upper block (truncate high freq)\n out_fft[:, :, :self.modes_height, :self.modes_width:super_res] = contract_2D(x[:, :, :self.modes_height, :self.modes_width], self._get_weight(indices, 0))\n # Lower block \n out_fft[:, :, -self.modes_height:, :self.modes_width:super_res] = contract_2D(x[:, :, -self.modes_height:, :self.modes_width], self._get_weight(indices, 1))\n\n out_size = (int(height*super_res), int(width*super_res))\n x = torch.fft.irfft2(out_fft, s=out_size, norm=self.fft_norm).type(dtype) #(x.size(-2), x.size(-1)))\n\n return x + self.bias\n\n def get_conv(self, indices):\n \"\"\"Returns a sub-convolutional layer from the joint parametrize main-convolution\n The parametrization of sub-convolutional layers is shared with the main one.\n \"\"\"\n if self.n_layers == 1:\n raise ValueError('A single convolution is parametrized, directly use the main class.')\n \n return SubConv(self, indices)\n \n def __getitem__(self, indices):\n return self.get_conv(indices)\n\n\nclass SubConv(nn.Module):\n \"\"\"Class representing one of the convolutions from the mother joint factorized convolution\n Notes\n -----\n This relies on the fact that nn.Parameters are not duplicated:\n if the same nn.Parameter is assigned to multiple modules, they all point to the same data, \n which is shared.\n \"\"\"\n def __init__(self, main_conv, indices):\n super().__init__()\n self.main_conv = main_conv\n self.indices = indices\n \n def forward(self, x, **kwargs):\n return self.main_conv.forward(x, self.indices, **kwargs)\n\n\nclass FactorizedSpectralConv1d(nn.Module):\n def __init__(self, in_channels, out_channels, modes, n_layers=1, \n bias=True, scale='auto', fft_norm='forward', rank=0.5, \n factorization='tucker', fixed_rank_modes=None, decomposition_kwargs=dict()):\n super().__init__()\n\n #Joint factorization only works for the same in and out channels\n if n_layers > 1:\n assert in_channels == out_channels\n\n self.in_channels = in_channels\n self.out_channels = out_channels\n self.modes = modes\n self.rank = rank\n self.factorization = factorization\n self.n_layers = n_layers\n self.fft_norm = fft_norm\n\n if scale == 'auto':\n scale = (1 / (in_channels * out_channels))\n\n if isinstance(fixed_rank_modes, bool):\n if fixed_rank_modes:\n fixed_rank_modes=[0]\n else:\n fixed_rank_modes=None\n\n if factorization is None:\n self.weight = nn.Parameter(scale * torch.randn(n_layers, in_channels, out_channels, self.modes,\n dtype=torch.cfloat))\n self._get_weight = self._get_weight_dense\n else:\n self.weight = tltorch.FactorizedTensor.new((n_layers, in_channels, out_channels, self.modes),\n rank=self.rank, factorization=factorization, \n dtype=torch.cfloat, fixed_rank_modes=fixed_rank_modes,\n **decomposition_kwargs)\n self.weight.normal_(0, scale)\n self._get_weight = self._get_weight_factorized\n\n if bias:\n self.bias = nn.Parameter(scale * torch.randn(1, self.out_channels, 1))\n else:\n self.bias = 0\n\n def _get_weight_factorized(self, layer_index):\n #Get the weights corresponding to a particular layer\n return self.weight()[layer_index, :, :, : ].to_tensor().contiguous()\n\n def _get_weight_dense(self, layer_index):\n #Get the weights corresponding to a particular layer\n return self.weight[layer_index, :, :, :]\n\n def forward(self, x, indices=0, s=None):\n batchsize, channels, width = x.shape\n dtype = x.dtype\n \n if s is None:\n s = width\n\n #Compute Fourier coeffcients \n x = torch.fft.rfft(x, norm=self.fft_norm)\n\n # Multiply relevant Fourier modes \n out_fft = torch.zeros([batchsize, self.out_channels, width//2 + 1], device=x.device, dtype=torch.cfloat)\n out_fft[:, :, :self.modes] = contract_1D(x[:, :, :self.modes], self._get_weight(indices))\n\n #Return to physical space\n x = torch.fft.irfft(out_fft, n=s, norm=self.fft_norm).type(dtype)\n\n return x + self.bias\n\n def get_conv(self, indices):\n \"\"\"Returns a sub-convolutional layer from the joint parametrize main-convolution\n The parametrization of sub-convolutional layers is shared with the main one.\n \"\"\"\n if self.n_layers == 1:\n raise ValueError('A single convolution is parametrized, directly use the main class.')\n \n return SubConv(self, indices)\n \n def __getitem__(self, indices):\n return self.get_conv(indices)\n \n\n\nclass JointFactorizedSpectralConv1d(nn.Module):\n def __init__(self, modes, width, n_layers=1, joint_factorization=True, in_channels=2, scale='auto',\n non_linearity=nn.GELU, rank=1.0, factorization='tucker', bias=True,\n fixed_rank_modes=False, fft_norm='forward', decomposition_kwargs=dict()):\n super().__init__()\n\n if isinstance(modes, int):\n self.modes = [modes for _ in range(n_layers)]\n else:\n self.modes = modes\n\n if isinstance(width, int):\n self.width = [width for _ in range(n_layers)]\n else:\n self.width = width\n\n assert len(self.width) == len(self.modes)\n\n self.n_layers = len(self.width)\n self.joint_factorization = joint_factorization\n\n if self.joint_factorization:\n assert self.width.count(self.width[0]) == self.n_layers and self.modes.count(self.modes[0]) == self.n_layers\n self.in_channels = self.width[0]\n else:\n self.in_channels = in_channels\n \n self.width = [self.in_channels] + self.width\n\n self.scale = scale\n self.non_linearity = non_linearity()\n self.rank = rank\n self.factorization = factorization\n self.bias = bias\n self.fixed_rank_modes = fixed_rank_modes\n self.decomposition_kwargs = decomposition_kwargs\n self.fft_norm = fft_norm\n \n if joint_factorization:\n self.convs = FactorizedSpectralConv1d(self.in_channels, self.width[0], self.modes[0],\n n_layers=self.n_layers,\n bias=self.bias,\n scale=self.scale,\n fft_norm=self.fft_norm,\n rank=self.rank,\n factorization=self.factorization,\n fixed_rank_modes=self.fixed_rank_modes,\n decomposition_kwargs=decomposition_kwargs)\n else:\n self.convs = nn.ModuleList([FactorizedSpectralConv1d(self.width[j], self.width[j+1], self.modes[j],\n n_layers=1,\n bias=self.bias,\n scale=self.scale,\n fft_norm=self.fft_norm,\n rank=self.rank,\n factorization=self.factorization,\n fixed_rank_modes=self.fixed_rank_modes,\n decomposition_kwargs=decomposition_kwargs) for j in range(self.n_layers)])\n\n self.linears = nn.ModuleList([nn.Conv1d(self.width[j], self.width[j+1], 1) for j in range(self.n_layers)])\n \n def forward(self, x, s=None):\n\n if s is None:\n s = [None for _ in range(self.n_layers)]\n\n if isinstance(s, int):\n s = [None for _ in range(self.n_layers-1)] + [s]\n\n for j in range(self.n_layers):\n x1 = self.convs[j](x, s=s[j])\n\n #Fourier interpolation\n if s[j] is not None:\n x2 = torch.fft.irfft(torch.fft.rfft(x, norm=self.fft_norm), n=s[j], norm=self.fft_norm)\n else:\n x2 = x\n \n x2 = self.linears[j](x2)\n\n x = x1 + x2\n\n if j < (self.n_layers - 1):\n x = self.non_linearity(x)\n\n return x" }, { "path": "models/fourier1d.py", "content": "import torch.nn as nn\nfrom .basics import SpectralConv1d\nfrom .utils import _get_act\n\n\nclass FNO1d(nn.Module):\n def __init__(self,\n modes, width=32,\n layers=None,\n fc_dim=128,\n in_dim=2, out_dim=1,\n act='relu'):\n super(FNO1d, self).__init__()\n\n \"\"\"\n The overall network. It contains several layers of the Fourier layer.\n 1. Lift the input to the desire channel dimension by self.fc0 .\n 2. 4 layers of the integral operators u' = (W + K)(u).\n W defined by self.w; K defined by self.conv .\n 3. Project from the channel space to the output space by self.fc1 and self.fc2 .\n\n input: the solution of the initial condition and location (a(x), x)\n input shape: (batchsize, x=s, c=2)\n output: the solution of a later timestep\n output shape: (batchsize, x=s, c=1)\n \"\"\"\n\n self.modes1 = modes\n self.width = width\n if layers is None:\n layers = [width] * 4\n\n self.fc0 = nn.Linear(in_dim, layers[0]) # input channel is 2: (a(x), x)\n\n self.sp_convs = nn.ModuleList([SpectralConv1d(\n in_size, out_size, num_modes) for in_size, out_size, num_modes in zip(layers, layers[1:], self.modes1)])\n\n self.ws = nn.ModuleList([nn.Conv1d(in_size, out_size, 1)\n for in_size, out_size in zip(layers, layers[1:])])\n\n self.fc1 = nn.Linear(layers[-1], fc_dim)\n self.fc2 = nn.Linear(fc_dim, out_dim)\n self.act = _get_act(act)\n\n def forward(self, x):\n length = len(self.ws)\n\n x = self.fc0(x)\n x = x.permute(0, 2, 1)\n\n for i, (speconv, w) in enumerate(zip(self.sp_convs, self.ws)):\n x1 = speconv(x)\n x2 = w(x)\n x = x1 + x2\n if i != length - 1:\n x = self.act(x)\n\n x = x.permute(0, 2, 1)\n x = self.fc1(x)\n x = self.act(x)\n x = self.fc2(x)\n return x\n\n" }, { "path": "models/fourier2d.py", "content": "import torch.nn as nn\nfrom .basics import SpectralConv2d\nfrom .utils import _get_act, add_padding2, remove_padding2\n\n\nclass FNO2d(nn.Module):\n def __init__(self, modes1, modes2,\n width=64, fc_dim=128,\n layers=None,\n in_dim=3, out_dim=1,\n act='gelu', \n pad_ratio=[0., 0.]):\n super(FNO2d, self).__init__()\n \"\"\"\n Args:\n - modes1: list of int, number of modes in first dimension in each layer\n - modes2: list of int, number of modes in second dimension in each layer\n - width: int, optional, if layers is None, it will be initialized as [width] * [len(modes1) + 1] \n - in_dim: number of input channels\n - out_dim: number of output channels\n - act: activation function, {tanh, gelu, relu, leaky_relu}, default: gelu\n - pad_ratio: list of float, or float; portion of domain to be extended. If float, paddings are added to the right. \n If list, paddings are added to both sides. pad_ratio[0] pads left, pad_ratio[1] pads right. \n \"\"\"\n if isinstance(pad_ratio, float):\n pad_ratio = [pad_ratio, pad_ratio]\n else:\n assert len(pad_ratio) == 2, 'Cannot add padding in more than 2 directions'\n self.modes1 = modes1\n self.modes2 = modes2\n \n self.pad_ratio = pad_ratio\n # input channel is 3: (a(x, y), x, y)\n if layers is None:\n self.layers = [width] * (len(modes1) + 1)\n else:\n self.layers = layers\n self.fc0 = nn.Linear(in_dim, layers[0])\n\n self.sp_convs = nn.ModuleList([SpectralConv2d(\n in_size, out_size, mode1_num, mode2_num)\n for in_size, out_size, mode1_num, mode2_num\n in zip(self.layers, self.layers[1:], self.modes1, self.modes2)])\n\n self.ws = nn.ModuleList([nn.Conv1d(in_size, out_size, 1)\n for in_size, out_size in zip(self.layers, self.layers[1:])])\n\n self.fc1 = nn.Linear(layers[-1], fc_dim)\n self.fc2 = nn.Linear(fc_dim, layers[-1])\n self.fc3 = nn.Linear(layers[-1], out_dim)\n self.act = _get_act(act)\n\n def forward(self, x):\n '''\n Args:\n - x : (batch size, x_grid, y_grid, 2)\n Returns:\n - x: (batch size, x_grid, y_grid, 1)\n '''\n size_1, size_2 = x.shape[1], x.shape[2]\n if max(self.pad_ratio) > 0:\n num_pad1 = [round(i * size_1) for i in self.pad_ratio]\n num_pad2 = [round(i * size_2) for i in self.pad_ratio]\n else:\n num_pad1 = num_pad2 = [0.]\n\n length = len(self.ws)\n batchsize = x.shape[0]\n x = self.fc0(x)\n x = x.permute(0, 3, 1, 2) # B, C, X, Y\n x = add_padding2(x, num_pad1, num_pad2)\n size_x, size_y = x.shape[-2], x.shape[-1]\n\n for i, (speconv, w) in enumerate(zip(self.sp_convs, self.ws)):\n x1 = speconv(x)\n x2 = w(x.view(batchsize, self.layers[i], -1)).view(batchsize, self.layers[i+1], size_x, size_y)\n x = x1 + x2\n if i != length - 1:\n x = self.act(x)\n x = remove_padding2(x, num_pad1, num_pad2)\n x = x.permute(0, 2, 3, 1)\n x = self.fc1(x)\n x = self.act(x)\n x = self.fc2(x)\n x = self.act(x)\n x = self.fc3(x)\n return x\n" }, { "path": "models/fourier3d.py", "content": "import torch.nn as nn\nfrom .basics import SpectralConv3d\nfrom .utils import add_padding, remove_padding, _get_act\n\n\nclass FNO3d(nn.Module):\n def __init__(self, \n modes1, modes2, modes3,\n width=16, \n fc_dim=128,\n layers=None,\n in_dim=4, out_dim=1,\n act='gelu', \n pad_ratio=[0., 0.]):\n '''\n Args:\n modes1: list of int, first dimension maximal modes for each layer\n modes2: list of int, second dimension maximal modes for each layer\n modes3: list of int, third dimension maximal modes for each layer\n layers: list of int, channels for each layer\n fc_dim: dimension of fully connected layers\n in_dim: int, input dimension\n out_dim: int, output dimension\n act: {tanh, gelu, relu, leaky_relu}, activation function\n pad_ratio: the ratio of the extended domain\n '''\n super(FNO3d, self).__init__()\n\n if isinstance(pad_ratio, float):\n pad_ratio = [pad_ratio, pad_ratio]\n else:\n assert len(pad_ratio) == 2, 'Cannot add padding in more than 2 directions.'\n\n self.pad_ratio = pad_ratio\n self.modes1 = modes1\n self.modes2 = modes2\n self.modes3 = modes3\n self.pad_ratio = pad_ratio\n\n if layers is None:\n self.layers = [width] * 4\n else:\n self.layers = layers\n self.fc0 = nn.Linear(in_dim, layers[0])\n\n self.sp_convs = nn.ModuleList([SpectralConv3d(\n in_size, out_size, mode1_num, mode2_num, mode3_num)\n for in_size, out_size, mode1_num, mode2_num, mode3_num\n in zip(self.layers, self.layers[1:], self.modes1, self.modes2, self.modes3)])\n\n self.ws = nn.ModuleList([nn.Conv1d(in_size, out_size, 1)\n for in_size, out_size in zip(self.layers, self.layers[1:])])\n\n self.fc1 = nn.Linear(layers[-1], fc_dim)\n self.fc2 = nn.Linear(fc_dim, out_dim)\n self.act = _get_act(act)\n\n def forward(self, x):\n '''\n Args:\n x: (batchsize, x_grid, y_grid, t_grid, 3)\n\n Returns:\n u: (batchsize, x_grid, y_grid, t_grid, 1)\n\n '''\n size_z = x.shape[-2]\n if max(self.pad_ratio) > 0:\n num_pad = [round(size_z * i) for i in self.pad_ratio]\n else:\n num_pad = [0., 0.]\n length = len(self.ws)\n batchsize = x.shape[0]\n \n x = self.fc0(x)\n x = x.permute(0, 4, 1, 2, 3)\n x = add_padding(x, num_pad=num_pad)\n size_x, size_y, size_z = x.shape[-3], x.shape[-2], x.shape[-1]\n\n for i, (speconv, w) in enumerate(zip(self.sp_convs, self.ws)):\n x1 = speconv(x)\n x2 = w(x.view(batchsize, self.layers[i], -1)).view(batchsize, self.layers[i+1], size_x, size_y, size_z)\n x = x1 + x2\n if i != length - 1:\n x = self.act(x)\n x = remove_padding(x, num_pad=num_pad)\n x = x.permute(0, 2, 3, 4, 1)\n x = self.fc1(x)\n x = self.act(x)\n x = self.fc2(x)\n return x" }, { "path": "models/lowrank2d.py", "content": "from .FCN import DenseNet\n\nimport numpy as np\nimport torch\nimport torch.nn as nn\n\n\nclass LowRank2d(nn.Module):\n def __init__(self, in_channels, out_channels):\n super(LowRank2d, self).__init__()\n self.in_channels = in_channels\n self.out_channels = out_channels\n\n self.phi = DenseNet([2, 64, 128, in_channels*out_channels], torch.nn.ReLU)\n self.psi = DenseNet([2, 64, 128, in_channels*out_channels], torch.nn.ReLU)\n\n def get_grid(self, S1, S2, batchsize, device):\n gridx = torch.tensor(np.linspace(0, 1, S1+1)[:-1], dtype=torch.float)\n gridx = gridx.reshape(1, S1, 1).repeat([batchsize, 1, S2])\n gridy = torch.tensor(np.linspace(0, 1, S2+1)[:-1], dtype=torch.float)\n gridy = gridy.reshape(1, 1, S2).repeat([batchsize, S1, 1])\n return torch.stack((gridx, gridy), dim=-1).to(device)\n\n def forward(self, x, gridy=None):\n # x (batch, channel, x, y)\n # y (batch, Ny, 2)\n batchsize, size1, size2 = x.shape[0], x.shape[2], x.shape[3]\n\n gridx = self.get_grid(S1=size1, S2=size2, batchsize=1, device=x.device).reshape(size1 * size2, 2)\n if gridy is None:\n gridy = self.get_grid(S1=size1, S2=size2, batchsize=batchsize, device=x.device).reshape(batchsize, size1 * size2, 2)\n Nx = size1 * size2\n Ny = gridy.shape[1]\n\n phi_eval = self.phi(gridx).reshape(Nx, self.out_channels, self.in_channels)\n psi_eval = self.psi(gridy).reshape(batchsize, Ny, self.out_channels, self.in_channels)\n x = x.reshape(batchsize, self.in_channels, Nx)\n\n x = torch.einsum('noi,bin,bmoi->bom', phi_eval, x, psi_eval) / Nx\n return x" }, { "path": "models/tfno.py", "content": "import torch.nn as nn\nimport torch.nn.functional as F\nfrom .core import FactorizedSpectralConv2d, JointFactorizedSpectralConv1d, FactorizedSpectralConv3d\n\n\nclass FactorizedFNO3d(nn.Module):\n def __init__(self, modes_height, modes_width, modes_depth, width, fc_channels=256, n_layers=4,\n joint_factorization=True, non_linearity=F.gelu,\n rank=1.0, factorization='cp', fixed_rank_modes=False,\n domain_padding=9, in_channels=3, Block=None,\n verbose=True, fft_contraction='complex',\n fft_norm='backward',\n mlp=False,\n decomposition_kwargs=dict()):\n super().__init__()\n self.modes_height = modes_height\n self.modes_width = modes_width\n self.modes_depth = modes_depth\n self.width = width\n self.fc_channels = fc_channels\n self.n_layers = n_layers\n self.joint_factorization = joint_factorization\n self.non_linearity = non_linearity\n self.rank = rank\n self.factorization = factorization\n self.fixed_rank_modes = fixed_rank_modes\n self.domain_padding = domain_padding # pad the domain if input is non-periodic\n self.in_channels = in_channels\n self.decomposition_kwargs = decomposition_kwargs\n self.fft_norm = fft_norm\n self.verbose = verbose\n \n if Block is None:\n Block = FactorizedSpectralConv3d\n if verbose:\n print(f'FNO Block using {Block}, fft_contraction={fft_contraction}')\n\n self.Block = Block\n\n if joint_factorization:\n self.convs = Block(self.width, self.width, self.modes_height, self.modes_width, self.modes_depth,\n rank=rank,\n fft_contraction=fft_contraction,\n fft_norm=fft_norm,\n factorization=factorization, \n fixed_rank_modes=fixed_rank_modes, \n decomposition_kwargs=decomposition_kwargs,\n mlp=mlp,\n n_layers=n_layers)\n else:\n self.convs = nn.ModuleList([Block(self.width, self.modes_height, self.modes_width, self.modes_depth,\n fft_contraction=fft_contraction,\n rank=rank,\n factorization=factorization, \n fixed_rank_modes=fixed_rank_modes, \n decomposition_kwargs=decomposition_kwargs,\n mlp=mlp,\n n_layers=1) for _ in range(n_layers)])\n self.linears = nn.ModuleList([nn.Conv3d(self.width, self.width, 1) for _ in range(n_layers)])\n \n self.fc0 = nn.Linear(in_channels, self.width) # input channel is 3: (a(x, y), x, y)\n self.fc1 = nn.Linear(self.width, fc_channels)\n self.fc2 = nn.Linear(fc_channels, 1)\n\n def forward(self, x, super_res=1):\n #grid = self.get_grid(x.shape, x.device)\n #x = torch.cat((x, grid), dim=-1)\n #x = self.fc0(x)\n #x = x.permute(0, 3, 1, 2)\n\n x = x.permute(0,2,3,4,1)\n x = self.fc0(x)\n x = x.permute(0,4,1,2,3)\n\n x = F.pad(x, [0, self.domain_padding])\n\n for i in range(self.n_layers):\n if super_res > 1 and i == (self.n_layers - 1):\n super_res = super_res\n else:\n super_res = 1\n\n x1 = self.convs[i](x) #, super_res=super_res)\n x2 = self.linears[i](x)\n x = x1 + x2\n if i < (self.n_layers - 1):\n x = self.non_linearity(x)\n\n x = x[..., :-self.domain_padding]\n x = x.permute(0, 2, 3, 4, 1)\n x = self.fc1(x)\n x = self.non_linearity(x)\n x = self.fc2(x)\n x = x.permute(0,4,1,2,3)\n return x\n\n\nclass FactorizedFNO2d(nn.Module):\n def __init__(self, modes_height, modes_width, width, fc_channels=256, n_layers=4,\n joint_factorization=True, non_linearity=F.gelu,\n rank=1.0, factorization='cp', fixed_rank_modes=False,\n domain_padding=9, in_channels=3, Block=None,\n verbose=True, fft_contraction='complex',\n fft_norm='backward',\n decomposition_kwargs=dict()):\n super().__init__()\n \"\"\"\n input: the solution of the coefficient function and locations (a(x, y), x, y)\n input shape: (batchsize, x=s, y=s, c=3)\n output: the solution \n output shape: (batchsize, x=s, y=s, c=1)\n \"\"\"\n self.modes_height = modes_height\n self.modes_width = modes_width\n self.width = width\n self.fc_channels = fc_channels\n self.n_layers = n_layers\n self.joint_factorization = joint_factorization\n self.non_linearity = non_linearity\n self.rank = rank\n self.factorization = factorization\n self.fixed_rank_modes = fixed_rank_modes\n self.domain_padding = domain_padding # pad the domain if input is non-periodic\n self.in_channels = in_channels\n self.decomposition_kwargs = decomposition_kwargs\n self.fft_norm = fft_norm\n self.verbose = verbose\n \n if Block is None:\n Block = FactorizedSpectralConv2d\n if verbose:\n print(f'FNO Block using {Block}, fft_contraction={fft_contraction}')\n\n self.Block = Block\n\n if joint_factorization:\n self.convs = Block(self.width, self.width, self.modes_height, self.modes_width, \n rank=rank,\n fft_contraction=fft_contraction,\n fft_norm=fft_norm,\n factorization=factorization, \n fixed_rank_modes=fixed_rank_modes, \n decomposition_kwargs=decomposition_kwargs,\n n_layers=n_layers)\n else:\n self.convs = nn.ModuleList([Block(self.width, self.width, self.modes_height,\n self.modes_width,\n fft_contraction=fft_contraction,\n rank=rank,\n factorization=factorization, \n fixed_rank_modes=fixed_rank_modes, \n decomposition_kwargs=decomposition_kwargs,\n n_layers=1) for _ in range(n_layers)])\n self.linears = nn.ModuleList([nn.Conv2d(self.width, self.width, 1) for _ in range(n_layers)])\n \n self.fc0 = nn.Linear(in_channels, self.width) # input channel is 3: (a(x, y), x, y)\n self.fc1 = nn.Linear(self.width, fc_channels)\n self.fc2 = nn.Linear(fc_channels, 1)\n\n def forward(self, x, super_res=1):\n #grid = self.get_grid(x.shape, x.device)\n #x = torch.cat((x, grid), dim=-1)\n #x = self.fc0(x)\n #x = x.permute(0, 3, 1, 2)\n\n x = x.permute(0,2,3,1)\n x = self.fc0(x)\n x = x.permute(0,3,1,2)\n\n x = F.pad(x, [0, self.domain_padding, 0, self.domain_padding])\n\n for i in range(self.n_layers):\n if super_res > 1 and i == (self.n_layers - 1):\n super_res = super_res\n else:\n super_res = 1\n\n x1 = self.convs[i](x) #, super_res=super_res)\n x2 = self.linears[i](x)\n x = x1 + x2\n if i < (self.n_layers - 1):\n x = self.non_linearity(x)\n\n x = x[..., :-self.domain_padding, :-self.domain_padding]\n x = x.permute(0, 2, 3, 1)\n x = self.fc1(x)\n x = self.non_linearity(x)\n x = self.fc2(x)\n x = x.permute(0,3,1,2)\n return x\n\n # def extra_repr(self):\n # s = (f'{self.modes_height=}, {self.modes_width=}, {self.width=}, {self.fc_channels=}, {self.n_layers=}, '\n # f'{self.joint_factorization=}, {self.non_linearity=}, '\n # f'{self.rank=}, {self.factorization=}, {self.fixed_rank_modes=}, '\n # f'{self.domain_padding=}, {self.in_channels=}, {self.Block=}, '\n # f'{self.verbose=}, '\n # f'{self.decomposition_kwargs=}')\n # return s\n\n \nclass FactorizedFNO1d(nn.Module):\n def __init__(self, modes, width, in_channels=2, out_channels=1, n_layers=4, \n lifting=None, projection=None, joint_factorization=True, scale='auto', \n non_linearity=nn.GELU, rank=1.0, factorization='tucker', bias=True, \n fixed_rank_modes=False, fft_norm='forward', decomposition_kwargs=dict()):\n super().__init__()\n\n if isinstance(width, int):\n init_width = width\n final_width = width\n else:\n init_width = width[0]\n final_width = width[-1]\n \n self.non_linearity = non_linearity()\n\n if lifting is None:\n self.lifting = nn.Linear(in_channels, init_width)\n \n if projection is None:\n self.projection = nn.Sequential(nn.Linear(final_width, 256),\n self.non_linearity,\n nn.Linear(256, out_channels))\n\n self.fno_layers = JointFactorizedSpectralConv1d(modes, width, n_layers=n_layers, joint_factorization=joint_factorization,\n in_channels=init_width, scale=scale, non_linearity=non_linearity,\n rank=rank, factorization=factorization, bias=bias, fixed_rank_modes=fixed_rank_modes, \n fft_norm=fft_norm, decomposition_kwargs=decomposition_kwargs)\n \n def forward(self, x, s=None):\n #Lifting\n x = x.permute(0,2,1)\n x = self.lifting(x)\n x = x.permute(0,2,1)\n\n #Fourier layers\n x = self.fno_layers(x, s=s)\n\n #Projection\n x = x.permute(0,2,1)\n x = self.projection(x)\n x = x.permute(0,2,1)\n \n return x\n\n\n" }, { "path": "models/utils.py", "content": "import torch.nn.functional as F\n\n\ndef add_padding(x, num_pad):\n if max(num_pad) > 0:\n res = F.pad(x, (num_pad[0], num_pad[1]), 'constant', 0)\n else:\n res = x\n return res\n\n\ndef add_padding2(x, num_pad1, num_pad2):\n if max(num_pad1) > 0 or max(num_pad2) > 0:\n res = F.pad(x, (num_pad2[0], num_pad2[1], num_pad1[0], num_pad1[1]), 'constant', 0.)\n else:\n res = x\n return res\n\n\ndef remove_padding(x, num_pad):\n if max(num_pad) > 0:\n res = x[..., num_pad[0]:-num_pad[1]]\n else:\n res = x\n return res\n\n\ndef remove_padding2(x, num_pad1, num_pad2):\n if max(num_pad1) > 0 or max(num_pad2) > 0:\n res = x[..., num_pad1[0]:-num_pad1[1], num_pad2[0]:-num_pad2[1]]\n else:\n res = x\n return res\n\n\ndef _get_act(act):\n if act == 'tanh':\n func = F.tanh\n elif act == 'gelu':\n func = F.gelu\n elif act == 'relu':\n func = F.relu_\n elif act == 'elu':\n func = F.elu_\n elif act == 'leaky_relu':\n func = F.leaky_relu_\n else:\n raise ValueError(f'{act} is not supported')\n return func\n\n" }, { "path": "pinns.py", "content": "from argparse import ArgumentParser\nimport yaml\n\nfrom baselines.pinns_ns_05s import train\nfrom baselines.pinns_ns_50s import train_longtime\nfrom baselines.sapinns import train_sapinn\nimport csv\n\n\nif __name__ == '__main__':\n parser = ArgumentParser(description='Basic paser')\n parser.add_argument('--config_path', type=str, help='Path to the configuration file')\n parser.add_argument('--log', action='store_true', help='log the results')\n parser.add_argument('--start', type=int, default=0, help='start index')\n parser.add_argument('--stop', type=int, default=1, help='stopping index')\n args = parser.parse_args()\n\n config_file = args.config_path\n with open(config_file, 'r') as stream:\n config = yaml.load(stream, yaml.FullLoader)\n\n with open(config['log']['logfile'], 'a') as f:\n writer = csv.writer(f)\n writer.writerow(['Index', 'Error in u', 'Error in v', 'Error in w', 'Step', 'Time cost'])\n\n for i in range(args.start, args.stop):\n print(f'Start to solve instance {i}')\n if 'time_scale' in config['data']:\n train_longtime(i, config, args)\n elif config['log']['group'] == 'SA-PINNs':\n train_sapinn(i, config, args)\n else:\n train(i, config, args)\n\n\n" }, { "path": "prepare_data.py", "content": "import numpy as np\nimport matplotlib.pyplot as plt\n\n\ndef shuffle_data(datapath):\n data = np.load(datapath)\n rng = np.random.default_rng(123)\n rng.shuffle(data, axis=0)\n savepath = datapath.replace('.npy', '-shuffle.npy')\n np.save(savepath, data)\n\n\ndef test_data(datapath):\n raw = np.load(datapath, mmap_mode='r')\n print(raw[0, 0, 0, 0:10])\n newpath = datapath.replace('.npy', '-shuffle.npy')\n new = np.load(newpath, mmap_mode='r')\n print(new[0, 0, 0, 0:10])\n\n\ndef get_slice(datapath):\n raw = np.load(datapath, mmap_mode='r')\n\n data = raw[-10:]\n print(data.shape)\n savepath = 'data/Re500-5x513x256x256.npy'\n np.save(savepath, data)\n\n\ndef plot_test(datapath):\n duration = 0.125\n raw = np.load(datapath, mmap_mode='r')\n \n\n\n\nif __name__ == '__main__':\n # datapath = '../data/NS-Re500_T300_id0.npy'\n # shuffle_data(datapath)\n # test_data(datapath)\n\n datapath = '/raid/hongkai/NS-Re500_T300_id0-shuffle.npy'\n get_slice(datapath)" }, { "path": "profile-solver-legacy.py", "content": "import math\n\nimport torch\nfrom solver.legacy_solver import navier_stokes_2d, GaussianRF\n\nimport scipy.io\nfrom timeit import default_timer\n\n\nif __name__ == '__main__':\n device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n # Resolution\n s = 2048\n sub = 1\n\n # Number of solutions to generate\n N = 1\n\n # Set up 2d GRF with covariance parameters\n GRF = GaussianRF(2, s, alpha=2.5, tau=7, device=device)\n\n # Forcing function: 0.1*(sin(2pi(x+y)) + cos(2pi(x+y)))\n t = torch.linspace(0, 1, s + 1, device=device)\n t = t[0:-1]\n\n X, Y = torch.meshgrid(t, t)\n f = 0.1 * (torch.sin(2 * math.pi * (X + Y)) + torch.cos(2 * math.pi * (X + Y)))\n\n # Number of snapshots from solution\n record_steps = 200\n\n # Inputs\n a = torch.zeros(N, s, s)\n # Solutions\n u = torch.zeros(N, s, s, record_steps)\n\n # Solve equations in batches (order of magnitude speed-up)\n\n # Batch size\n bsize = 1\n\n c = 0\n t0 = default_timer()\n for j in range(N // bsize):\n # Sample random feilds\n w0 = GRF.sample(bsize)\n\n # Solve NS\n sol, sol_t = navier_stokes_2d(w0, f, 1e-3, 50.0, 1e-4, record_steps)\n\n a[c:(c + bsize), ...] = w0\n u[c:(c + bsize), ...] = sol\n\n c += bsize\n t1 = default_timer()\n print(f'Time cost {t1 - t0} s')\n torch.save(\n {\n 'a': a.cpu(),\n 'u': u.cpu(),\n 't': sol_t.cpu()\n },\n 'data/ns_data.pt'\n )\n # scipy.io.savemat('data/ns_data.mat', mdict={'a': a.cpu().numpy(), 'u': u.cpu().numpy(), 't': sol_t.cpu().numpy()})" }, { "path": "profiler/calmacs.py", "content": "from ptflops import get_model_complexity_info\n" }, { "path": "run_pino2d.py", "content": "import yaml\nfrom argparse import ArgumentParser\nimport random\n\nimport torch\n\nfrom models import FNO2d\nfrom train_utils import Adam\nfrom torch.utils.data import DataLoader\nfrom train_utils.datasets import DarcyFlow\nfrom train_utils.train_2d import train_2d_operator\n\n\ndef train(args, config):\n seed = random.randint(1, 10000)\n print(f'Random seed :{seed}')\n torch.manual_seed(seed)\n device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n data_config = config['data']\n dataset = DarcyFlow(data_config['datapath'],\n nx=data_config['nx'], sub=data_config['sub'],\n offset=data_config['offset'], num=data_config['n_sample'])\n dataloader = DataLoader(dataset, batch_size=config['train']['batchsize'])\n model = FNO2d(modes1=config['model']['modes1'],\n modes2=config['model']['modes2'],\n fc_dim=config['model']['fc_dim'],\n layers=config['model']['layers'],\n act=config['model']['act']).to(device)\n # Load from checkpoint\n if 'ckpt' in config['train']:\n ckpt_path = config['train']['ckpt']\n ckpt = torch.load(ckpt_path)\n model.load_state_dict(ckpt['model'])\n print('Weights loaded from %s' % ckpt_path)\n optimizer = Adam(model.parameters(), betas=(0.9, 0.999),\n lr=config['train']['base_lr'])\n scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer,\n milestones=config['train']['milestones'],\n gamma=config['train']['scheduler_gamma'])\n train_2d_operator(model,\n dataloader,\n optimizer, scheduler,\n config, rank=0, log=args.log,\n project=config['log']['project'],\n group=config['log']['group'])\n\n\nif __name__ == '__main__':\n torch.backends.cudnn.benchmark = True\n parser = ArgumentParser(description='Basic paser')\n parser.add_argument('--config_path', type=str, help='Path to the configuration file')\n parser.add_argument('--start', type=int, help='Start index of test instance')\n parser.add_argument('--stop', type=int, help='Stop index of instances')\n parser.add_argument('--log', action='store_true', help='Turn on the wandb')\n args = parser.parse_args()\n\n config_file = args.config_path\n with open(config_file, 'r') as stream:\n config = yaml.load(stream, yaml.FullLoader)\n\n for i in range(args.start, args.stop):\n print(f'Start solving instance {i}')\n config['data']['offset'] = i\n train(args, config)\n print(f'{args.stop - args.start} instances are solved')" }, { "path": "run_pino3d.py", "content": "import random\nimport yaml\n\nimport torch\nfrom torch.utils.data import DataLoader\n\nfrom train_utils import Adam, NSLoader, get_forcing\nfrom train_utils.train_3d import train\n\nfrom models import FNO3d\nfrom argparse import ArgumentParser\nfrom train_utils.utils import requires_grad\n\n\ndef run_instance(loader, config, data_config):\n trainset = loader.make_dataset(data_config['n_sample'],\n start=data_config['offset'])\n train_loader = DataLoader(trainset, batch_size=config['train']['batchsize'])\n model = FNO3d(modes1=config['model']['modes1'],\n modes2=config['model']['modes2'],\n modes3=config['model']['modes3'],\n fc_dim=config['model']['fc_dim'],\n layers=config['model']['layers']).to(device)\n\n if 'ckpt' in config['train']:\n ckpt_path = config['train']['ckpt']\n ckpt = torch.load(ckpt_path)\n model.load_state_dict(ckpt['model'])\n print('Weights loaded from %s' % ckpt_path)\n\n if 'twolayer' in config['train'] and config['train']['twolayer']:\n requires_grad(model, False)\n requires_grad(model.sp_convs[-1], True)\n requires_grad(model.ws[-1], True)\n requires_grad(model.fc1, True)\n requires_grad(model.fc2, True)\n params = []\n for param in model.parameters():\n if param.requires_grad == True:\n params.append(param)\n else:\n params = model.parameters()\n\n beta1 = config['train']['beta1'] if 'beta1' in config['train'] else 0.9\n beta2 = config['train']['beta2'] if 'beta2' in config['train'] else 0.999\n optimizer = Adam(params, betas=(beta1, beta2),\n lr=config['train']['base_lr'])\n scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer,\n milestones=config['train']['milestones'],\n gamma=config['train']['scheduler_gamma'])\n forcing = get_forcing(loader.S).to(device)\n profile = config['train']['profile'] if 'profile' in config['train'] else False\n train(model,\n loader, train_loader,\n optimizer, scheduler,\n forcing, config,\n rank=0,\n log=options.log,\n project=config['log']['project'],\n group=config['log']['group'],\n use_tqdm=True,\n profile=profile)\n\n\nif __name__ == '__main__':\n torch.backends.cudnn.benchmark = True\n parser = ArgumentParser(description='Basic paser')\n parser.add_argument('--config_path', type=str, help='Path to the configuration file')\n parser.add_argument('--start', type=int, help='Start index of test instance')\n parser.add_argument('--stop', type=int, help='Stop index of instances')\n parser.add_argument('--log', action='store_true', help='Turn on the wandb')\n options = parser.parse_args()\n\n config_file = options.config_path\n with open(config_file, 'r') as stream:\n config = yaml.load(stream, yaml.FullLoader)\n\n device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n data_config = config['data']\n\n loader = NSLoader(datapath1=data_config['datapath'],\n nx=data_config['nx'], nt=data_config['nt'],\n sub=data_config['sub'], sub_t=data_config['sub_t'],\n N=data_config['total_num'],\n t_interval=data_config['time_interval'])\n for i in range(options.start, options.stop):\n print('Start training on instance %d' % i)\n config['data']['offset'] = i\n data_config['offset'] = i\n seed = random.randint(1, 10000)\n print(f'Random seed :{seed}')\n torch.manual_seed(seed)\n torch.cuda.manual_seed_all(seed)\n run_instance(loader, config, data_config)\n print('Done!')\n" }, { "path": "run_solver.py", "content": "import random\nimport math\nimport numpy as np\nimport yaml\nfrom argparse import ArgumentParser\nfrom timeit import default_timer\n\nimport torch\n\nfrom train_utils.datasets import NSLoader\nfrom train_utils.losses import LpLoss\nfrom solver.kolmogorov_flow import KolmogorovFlow2d\n\n\ndef solve(a,\n res_x,\n res_t,\n end,\n Re,\n n=4,\n delta_t=1e-3):\n '''\n Given initial condition a, solve for u in time interval [0, end]\n Args:\n a: initial condition, res_x by res_x tensor\n res_x: resolution in space\n res_t: record step in time\n end: end of the time interval\n Re: Reynolds number\n n: forcing number\n Returns:\n tensor of shape (res_x, res_x, res_t)\n '''\n dt = end / res_t\n\n solver = KolmogorovFlow2d(a, Re, n)\n sol = torch.zeros((res_x, res_x, res_t + 1), device=a.device)\n sol[:, :, 0] = a\n for j in range(res_t):\n solver.advance(dt, delta_t=delta_t)\n sol[:, :, 1 + j] = solver.vorticity().squeeze(0)\n return sol\n\n\nif __name__ == '__main__':\n torch.backends.cudnn.benchmark = True\n device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n parser = ArgumentParser(description='Basic paser')\n parser.add_argument('--config_path', type=str, help='Path to the configuration file')\n parser.add_argument('--deltat', type=float, default=1e-3, help='delta T')\n args = parser.parse_args()\n config_file = args.config_path\n with open(config_file, 'r') as stream:\n config = yaml.load(stream, yaml.FullLoader)\n\n data_config = config['data']\n loader = NSLoader(datapath1=data_config['datapath'],\n nx=data_config['nx'], nt=data_config['nt'],\n sub=data_config['sub'], sub_t=data_config['sub_t'],\n N=data_config['total_num'],\n t_interval=data_config['time_interval'])\n a_loader = loader.make_loader(data_config['n_sample'],\n batch_size=config['train']['batchsize'],\n start=data_config['offset'],\n train=data_config['shuffle'])\n\n print(f'Solver starts on device: {device}')\n myloss = LpLoss(size_average=True)\n test_err = []\n time_cost = []\n # run solver\n for _, u in a_loader:\n u = u[0].to(device)\n torch.cuda.synchronize()\n t1 = default_timer()\n pred = solve(u[:, :, 0],\n res_x=loader.S,\n res_t=loader.T - 1,\n end=data_config['time_interval'],\n Re=data_config['Re'],\n n=4,\n delta_t=args.deltat)\n torch.cuda.synchronize()\n t2 = default_timer()\n\n # report test error\n test_l2 = myloss(pred, u)\n\n test_err.append(test_l2.item())\n print(f'Test l2: {test_l2.item()}')\n time_cost.append(t2 - t1)\n\n test_err = np.array(test_err)\n time_cost = np.array(time_cost)\n\n idx = data_config['offset']\n n_sample = data_config['n_sample']\n print(f'Test instance: {idx} to {idx+n_sample}; \\n'\n f'Time cost = mean: {time_cost.mean()}s; std_err: {time_cost.std(ddof=1) / math.sqrt(len(a_loader))}s; \\n'\n f'Solver resolution: {loader.S} x {loader.S} x {loader.T}; \\n'\n f'Test L2 error = mean: {test_err.mean()}; std_err: {test_err.std(ddof=1) / math.sqrt(len(a_loader))}')\n\n\n" }, { "path": "scripts/device1-finetune.sh", "content": "#! /bin/bash\nCUDA_VISIBLE_DEVICES=3 python3 train_PINO3d.py \\\n--config_path configs/finetune/Re500-finetune-05s.yaml \\\n--start 1 \\\n--log;\nCUDA_VISIBLE_DEVICES=3 python3 train_PINO3d.py \\\n--config_path configs/finetune/Re500-finetune-05s.yaml \\\n--start 3 \\\n--log;\nCUDA_VISIBLE_DEVICES=3 python3 train_PINO3d.py \\\n--config_path configs/finetune/Re500-finetune-05s.yaml \\\n--start 4 \\\n--log;\nCUDA_VISIBLE_DEVICES=3 python3 train_PINO3d.py \\\n--config_path configs/finetune/Re500-finetune-05s.yaml \\\n--start 5 \\\n--log;\nCUDA_VISIBLE_DEVICES=3 python3 train_PINO3d.py \\\n--config_path configs/finetune/Re500-finetune-05s.yaml \\\n--start 6 \\\n--log;\nCUDA_VISIBLE_DEVICES=3 python3 train_PINO3d.py \\\n--config_path configs/finetune/Re500-finetune-05s.yaml \\\n--start 7 \\\n--log;\n" }, { "path": "scripts/device2-finetune.sh", "content": "#! /bin/bash\nCUDA_VISIBLE_DEVICES=2 python3 train_PINO3d.py \\\n--config_path configs/finetune/Re500-finetune-05s4k4.yaml \\\n--start 0 \\\n--log;\nCUDA_VISIBLE_DEVICES=2 python3 train_PINO3d.py \\\n--config_path configs/finetune/Re500-finetune-05s4k4.yaml \\\n--start 1 \\\n--log;\nCUDA_VISIBLE_DEVICES=2 python3 train_PINO3d.py \\\n--config_path configs/finetune/Re500-finetune-05s4k4.yaml \\\n--start 2 \\\n--log;\nCUDA_VISIBLE_DEVICES=2 python3 train_PINO3d.py \\\n--config_path configs/finetune/Re500-finetune-05s4k4.yaml \\\n--start 3 \\\n--log;\nCUDA_VISIBLE_DEVICES=3 python3 train_PINO3d.py \\\n--config_path configs/finetune/Re500-finetune-05s4k4.yaml \\\n--start 4 \\\n--log;\nCUDA_VISIBLE_DEVICES=3 python3 train_PINO3d.py \\\n--config_path configs/finetune/Re500-finetune-05s4k4.yaml \\\n--start 5 \\\n--log;\nCUDA_VISIBLE_DEVICES=3 python3 train_PINO3d.py \\\n--config_path configs/finetune/Re500-finetune-05s4k4.yaml \\\n--start 6 \\\n--log;\nCUDA_VISIBLE_DEVICES=3 python3 train_PINO3d.py \\\n--config_path configs/finetune/Re500-finetune-05s4k4.yaml \\\n--start 7 \\\n--log;\n\n" }, { "path": "scripts/device3.sh", "content": "#! /bin/bash\nCUDA_VISIBLE_DEVICES=3 python3 run_pino3d.py \\\n--config_path configs/transfer/Re100to300-1s.yaml \\\n--start 0 \\\n--stop 40 \\\n--log;\nCUDA_VISIBLE_DEVICES=3 python3 run_pino3d.py \\\n--config_path configs/transfer/Re200to300-1s.yaml \\\n--start 0 \\\n--stop 40 \\\n--log;\nCUDA_VISIBLE_DEVICES=3 python3 run_pino3d.py \\\n--config_path configs/transfer/Re250to300-1s.yaml \\\n--start 0 \\\n--stop 40 \\\n--log;\nCUDA_VISIBLE_DEVICES=3 python3 run_pino3d.py \\\n--config_path configs/transfer/Re300to300-1s.yaml \\\n--start 0 \\\n--stop 40 \\\n--log;\nCUDA_VISIBLE_DEVICES=3 python3 run_pino3d.py \\\n--config_path configs/transfer/Re350to300-1s.yaml \\\n--start 0 \\\n--stop 40 \\\n--log;\nCUDA_VISIBLE_DEVICES=3 python3 run_pino3d.py \\\n--config_path configs/transfer/Re400to300-1s.yaml \\\n--start 0 \\\n--stop 40 \\\n--log;\n" }, { "path": "scripts/finetune-4k-2layer.sh", "content": "#! /bin/bash\nCUDA_VISIBLE_DEVICES=1 python3 train_PINO3d.py \\\n--config_path configs/finetune/Re500-finetune-05s4k-2layer.yaml \\\n--start 9 \\\n--log;\nCUDA_VISIBLE_DEVICES=1 python3 train_PINO3d.py \\\n--config_path configs/finetune/Re500-finetune-05s4k-2layer.yaml \\\n--start 10 \\\n--log;\nCUDA_VISIBLE_DEVICES=1 python3 train_PINO3d.py \\\n--config_path configs/finetune/Re500-finetune-05s4k-2layer.yaml \\\n--start 11 \\\n--log;\nCUDA_VISIBLE_DEVICES=1 python3 train_PINO3d.py \\\n--config_path configs/finetune/Re500-finetune-05s4k-2layer.yaml \\\n--start 12 \\\n--log;\nCUDA_VISIBLE_DEVICES=1 python3 train_PINO3d.py \\\n--config_path configs/finetune/Re500-finetune-05s4k-2layer.yaml \\\n--start 13 \\\n--log;\nCUDA_VISIBLE_DEVICES=1 python3 train_PINO3d.py \\\n--config_path configs/finetune/Re500-finetune-05s4k-2layer.yaml \\\n--start 14 \\\n--log;\nCUDA_VISIBLE_DEVICES=1 python3 train_PINO3d.py \\\n--config_path configs/finetune/Re500-finetune-05s4k-2layer.yaml \\\n--start 15 \\\n--log;\nCUDA_VISIBLE_DEVICES=1 python3 train_PINO3d.py \\\n--config_path configs/finetune/Re500-finetune-05s4k-2layer.yaml \\\n--start 16 \\\n--log;\nCUDA_VISIBLE_DEVICES=1 python3 train_PINO3d.py \\\n--config_path configs/finetune/Re500-finetune-05s4k-2layer.yaml \\\n--start 17 \\\n--log;\n\n\n\n" }, { "path": "scripts/finetune-4k0.sh", "content": "#! /bin/bash\nCUDA_VISIBLE_DEVICES=2 python3 train_PINO3d.py \\\n--config_path configs/finetune/Re500-finetune-05s4k.yaml \\\n--start 9 \\\n--log;\n\n\n" }, { "path": "scripts/finetune-4k1-2layer.sh", "content": "#! /bin/bash\nCUDA_VISIBLE_DEVICES=0 python3 train_PINO3d.py \\\n--config_path configs/finetune/Re500-finetune-05s-2layer.yaml \\\n--start 9 \\\n--log;\nCUDA_VISIBLE_DEVICES=0 python3 train_PINO3d.py \\\n--config_path configs/finetune/Re500-finetune-05s-2layer.yaml \\\n--start 10 \\\n--log;\nCUDA_VISIBLE_DEVICES=0 python3 train_PINO3d.py \\\n--config_path configs/finetune/Re500-finetune-05s-2layer.yaml \\\n--start 11 \\\n--log;\nCUDA_VISIBLE_DEVICES=0 python3 train_PINO3d.py \\\n--config_path configs/finetune/Re500-finetune-05s-2layer.yaml \\\n--start 12 \\\n--log;\nCUDA_VISIBLE_DEVICES=0 python3 train_PINO3d.py \\\n--config_path configs/finetune/Re500-finetune-05s-2layer.yaml \\\n--start 13 \\\n--log;\nCUDA_VISIBLE_DEVICES=0 python3 train_PINO3d.py \\\n--config_path configs/finetune/Re500-finetune-05s-2layer.yaml \\\n--start 14 \\\n--log;\nCUDA_VISIBLE_DEVICES=0 python3 train_PINO3d.py \\\n--config_path configs/finetune/Re500-finetune-05s-2layer.yaml \\\n--start 15 \\\n--log;\nCUDA_VISIBLE_DEVICES=0 python3 train_PINO3d.py \\\n--config_path configs/finetune/Re500-finetune-05s-2layer.yaml \\\n--start 16 \\\n--log;\nCUDA_VISIBLE_DEVICES=0 python3 train_PINO3d.py \\\n--config_path configs/finetune/Re500-finetune-05s-2layer.yaml \\\n--start 17 \\\n--log;\n" }, { "path": "scripts/finetune-4k1.sh", "content": "#! /bin/bash\nCUDA_VISIBLE_DEVICES=0 python3 train_PINO3d.py \\\n--config_path configs/finetune/Re500-finetune-05s.yaml \\\n--start 9 \\\n--log;\nCUDA_VISIBLE_DEVICES=0 python3 train_PINO3d.py \\\n--config_path configs/finetune/Re500-finetune-05s.yaml \\\n--start 10 \\\n--log;\nCUDA_VISIBLE_DEVICES=0 python3 train_PINO3d.py \\\n--config_path configs/finetune/Re500-finetune-05s.yaml \\\n--start 11 \\\n--log;\nCUDA_VISIBLE_DEVICES=0 python3 train_PINO3d.py \\\n--config_path configs/finetune/Re500-finetune-05s.yaml \\\n--start 12 \\\n--log;\nCUDA_VISIBLE_DEVICES=0 python3 train_PINO3d.py \\\n--config_path configs/finetune/Re500-finetune-05s.yaml \\\n--start 13 \\\n--log;\nCUDA_VISIBLE_DEVICES=0 python3 train_PINO3d.py \\\n--config_path configs/finetune/Re500-finetune-05s.yaml \\\n--start 14 \\\n--log;\nCUDA_VISIBLE_DEVICES=0 python3 train_PINO3d.py \\\n--config_path configs/finetune/Re500-finetune-05s.yaml \\\n--start 15 \\\n--log;\nCUDA_VISIBLE_DEVICES=0 python3 train_PINO3d.py \\\n--config_path configs/finetune/Re500-finetune-05s.yaml \\\n--start 16 \\\n--log;\nCUDA_VISIBLE_DEVICES=0 python3 train_PINO3d.py \\\n--config_path configs/finetune/Re500-finetune-05s.yaml \\\n--start 17 \\\n--log;\n" }, { "path": "scripts/finetune-4k4-2layer.sh", "content": "#! /bin/bash\nCUDA_VISIBLE_DEVICES=2 python3 train_PINO3d.py \\\n--config_path configs/finetune/Re500-finetune-05s4k4-2layer.yaml \\\n--start 9 \\\n--log;\nCUDA_VISIBLE_DEVICES=2 python3 train_PINO3d.py \\\n--config_path configs/finetune/Re500-finetune-05s4k4-2layer.yaml \\\n--start 10 \\\n--log;\nCUDA_VISIBLE_DEVICES=2 python3 train_PINO3d.py \\\n--config_path configs/finetune/Re500-finetune-05s4k4-2layer.yaml \\\n--start 11 \\\n--log;\nCUDA_VISIBLE_DEVICES=2 python3 train_PINO3d.py \\\n--config_path configs/finetune/Re500-finetune-05s4k4-2layer.yaml \\\n--start 12 \\\n--log;\nCUDA_VISIBLE_DEVICES=2 python3 train_PINO3d.py \\\n--config_path configs/finetune/Re500-finetune-05s4k4-2layer.yaml \\\n--start 13 \\\n--log;\nCUDA_VISIBLE_DEVICES=2 python3 train_PINO3d.py \\\n--config_path configs/finetune/Re500-finetune-05s4k4-2layer.yaml \\\n--start 14 \\\n--log;\nCUDA_VISIBLE_DEVICES=2 python3 train_PINO3d.py \\\n--config_path configs/finetune/Re500-finetune-05s4k4-2layer.yaml \\\n--start 15 \\\n--log;\nCUDA_VISIBLE_DEVICES=2 python3 train_PINO3d.py \\\n--config_path configs/finetune/Re500-finetune-05s4k4-2layer.yaml \\\n--start 16 \\\n--log;\nCUDA_VISIBLE_DEVICES=2 python3 train_PINO3d.py \\\n--config_path configs/finetune/Re500-finetune-05s4k4-2layer.yaml \\\n--start 17 \\\n--log;\n\n" }, { "path": "scripts/fnoRe500.sh", "content": "#! /bin/bash\n#SBATCH --time=24:00:00\n#SBATCH --nodes=1\n#SBATCH --ntasks-per-node=16\n#SBATCH --gres gpu:v100:1\n#SBATCH --mem=64G\n#SBATCH --email-user=hzzheng@caltech.edu\n#SBATCH --mail-type=BEGIN\n#SBATCH --mail-type=END\n#SBATCH --mail-type=FAIL\n\npython3 train_operator.py --config_path configs/operator/Re500-FNO.yaml" }, { "path": "scripts/ngc_submit_pino.sh", "content": "ngc batch run --name 'ml-model.PINO.ns-dat400' --preempt RESUMABLE \\\n--commandline 'cd /Code/PINO; git pull; bash scripts/train_dat400.sh' \\\n--image 'nvidia/pytorch:22.08-py3' \\\n--priority HIGH \\\n--ace nv-us-west-2 \\\n--instance dgxa100.40g.1.norm \\\n--workspace QsixjfOES8uYIp5kwIDblQ:/Code \\\n--datasetid 111345:/mount/data \\\n--team nvr-aialgo \\\n--result /results\n\nngc batch run --name 'ml-model.PINO.ns-dat200' --preempt RESUMABLE \\\n--commandline 'cd /Code/PINO; git pull; bash scripts/train_dat200.sh' \\\n--image 'nvidia/pytorch:22.08-py3' \\\n--priority HIGH \\\n--ace nv-us-west-2 \\\n--instance dgxa100.40g.1.norm \\\n--workspace QsixjfOES8uYIp5kwIDblQ:/Code \\\n--datasetid 111345:/mount/data \\\n--team nvr-aialgo \\\n--result /results\n\nngc batch run --name 'ml-model.PINO.ns-dat80' --preempt RESUMABLE \\\n--commandline 'cd /Code/PINO; git pull; bash scripts/train_dat80.sh' \\\n--image 'nvidia/pytorch:22.08-py3' \\\n--priority HIGH \\\n--ace nv-us-west-2 \\\n--instance dgxa100.40g.1.norm \\\n--workspace QsixjfOES8uYIp5kwIDblQ:/Code \\\n--datasetid 111345:/mount/data \\\n--team nvr-aialgo \\\n--result /results\n\nngc batch run --name 'ml-model.PINO.ns-dat40' --preempt RESUMABLE \\\n--commandline 'cd /Code/PINO; git pull; bash scripts/train_dat40.sh' \\\n--image 'nvidia/pytorch:22.08-py3' \\\n--priority HIGH \\\n--ace nv-us-west-2 \\\n--instance dgxa100.40g.1.norm \\\n--workspace QsixjfOES8uYIp5kwIDblQ:/Code \\\n--datasetid 111345:/mount/data \\\n--team nvr-aialgo \\\n--result /results\n\nngc batch run --name 'ml-model.PINO.ns-dat0' --preempt RESUMABLE \\\n--commandline 'cd /Code/PINO; git pull; bash scripts/train_dat0.sh' \\\n--image 'nvidia/pytorch:22.08-py3' \\\n--priority HIGH \\\n--ace nv-us-west-2 \\\n--instance dgxa100.40g.1.norm \\\n--workspace QsixjfOES8uYIp5kwIDblQ:/Code \\\n--datasetid 111345:/mount/data \\\n--team nvr-aialgo \\\n--result /results\n\nngc batch run --name 'ml-model.PINO.ns-res32' --preempt RESUMABLE \\\n--commandline 'cd /Code/PINO; git pull; bash scripts/train_res32.sh' \\\n--image 'nvidia/pytorch:22.08-py3' \\\n--priority HIGH \\\n--ace nv-us-west-2 \\\n--instance dgxa100.40g.1.norm \\\n--workspace QsixjfOES8uYIp5kwIDblQ:/Code \\\n--datasetid 111345:/mount/data \\\n--team nvr-aialgo \\\n--result /results\n\nngc batch run --name 'ml-model.PINO.ns-res16' --preempt RESUMABLE \\\n--commandline 'cd /Code/PINO; git pull; bash scripts/train_res16.sh' \\\n--image 'nvidia/pytorch:22.08-py3' \\\n--priority HIGH \\\n--ace nv-us-west-2 \\\n--instance dgxa100.40g.1.norm \\\n--workspace QsixjfOES8uYIp5kwIDblQ:/Code \\\n--datasetid 111345:/mount/data \\\n--team nvr-aialgo \\\n--result /results" }, { "path": "scripts/ngc_test_submit_pino.sh", "content": "ngc batch run --name 'ml-model.PINO.ns-dat800' --preempt RESUMABLE \\\n--commandline 'cd /Code/PINO; git pull; bash scripts/train_dat800.sh' \\\n--image 'nvidia/pytorch:22.08-py3' \\\n--priority HIGH \\\n--ace nv-us-west-2 \\\n--instance dgxa100.40g.1.norm \\\n--workspace QsixjfOES8uYIp5kwIDblQ:/Code \\\n--datasetid 111345:/mount/data \\\n--team nvr-aialgo \\\n--result /results" }, { "path": "scripts/pretrain.sh", "content": "#! /bin/bash\nCUDA_VISIBLE_DEVICES=0 python3 pretrain.py \\\n--config_path configs/pretrain/Re500-pretrain-1s.yaml\n\n" }, { "path": "scripts/scratchRe500.sh", "content": "#! /bin/bash\nCUDA_VISIBLE_DEVICES=2 python3 run_pino3d.py \\\n--config_path configs/scratch/Re500-scratch-05s.yaml \\\n--start 0 \\\n--stop 10 \\\n--log" }, { "path": "scripts/test-opt/Re500-1_8.sh", "content": "#! /bin/bash\nfor i in {0..49}\ndo\nCUDA_VISIBLE_DEVICES=1 python3 instance_opt.py --config configs/instance/Re500-1_8-PINO.yaml --ckpt checkpoints/Re500-1_8s-800-PINO-140000.pt --idx $i --tqdm\ndone" }, { "path": "scripts/train_dat0.sh", "content": "pip install wandb tqdm pyyaml\nwandb login 69a3bddb4146cf76113885de5af84c7f4c165753\npython3 train_pino.py --config configs/ngc/Re500-1_8-dat40-PINO.yaml --log" }, { "path": "scripts/train_dat200.sh", "content": "pip install wandb tqdm pyyaml\nwandb login 69a3bddb4146cf76113885de5af84c7f4c165753\npython3 train_pino.py --config configs/ngc/Re500-1_8-dat200-PINO.yaml --log" }, { "path": "scripts/train_dat40.sh", "content": "pip install wandb tqdm pyyaml\nwandb login 69a3bddb4146cf76113885de5af84c7f4c165753\npython3 train_pino.py --config configs/ngc/Re500-1_8-dat40-PINO.yaml --log" }, { "path": "scripts/train_dat400.sh", "content": "pip install wandb tqdm pyyaml\nwandb login 69a3bddb4146cf76113885de5af84c7f4c165753\npython3 train_pino.py --config configs/ngc/Re500-1_8-dat400-PINO.yaml --log" }, { "path": "scripts/train_dat80.sh", "content": "pip install wandb tqdm pyyaml\nwandb login 69a3bddb4146cf76113885de5af84c7f4c165753\npython3 train_pino.py --config configs/ngc/Re500-1_8-dat80-PINO.yaml --log" }, { "path": "scripts/train_dat800.sh", "content": "pip install wandb tqdm pyyaml\nwandb login 69a3bddb4146cf76113885de5af84c7f4c165753\npython3 train_pino.py --config configs/ngc/Re500-1_8-dat800-PINO.yaml --log" }, { "path": "scripts/train_res16.sh", "content": "pip install wandb tqdm pyyaml\nwandb login 69a3bddb4146cf76113885de5af84c7f4c165753\npython3 train_pino.py --config configs/ngc/Re500-1_8-res16-PINO.yaml --log" }, { "path": "scripts/train_res32.sh", "content": "pip install wandb tqdm pyyaml\nwandb login 69a3bddb4146cf76113885de5af84c7f4c165753\npython3 train_pino.py --config configs/ngc/Re500-1_8-res32-PINO.yaml --log" }, { "path": "solver/__init__.py", "content": "" }, { "path": "solver/kolmogorov_flow.py", "content": "import torch\nimport math\n\n\nclass KolmogorovFlow2d(object):\n\n def __init__(self, w0, Re, n):\n\n # Grid size\n\n self.s = w0.size()[-1]\n\n assert self.s == w0.size()[-2], \"Grid must be uniform in both directions.\"\n\n assert math.log2(self.s).is_integer(), \"Grid size must be power of 2.\"\n\n assert n >= 0 and isinstance(n, int), \"Forcing number must be non-negative integer.\"\n\n assert n < self.s // 2 - 1, \"Forcing number too large for grid size.\"\n\n # Forcing number\n self.n = n\n\n assert Re > 0, \"Reynolds number must be positive.\"\n\n # Reynolds number\n self.Re = Re\n\n # Device\n self.device = w0.device\n\n # Current time\n self.time = 0.0\n\n # Current vorticity in Fourier space\n self.w_h = torch.fft.fft2(w0, norm=\"backward\")\n\n # Wavenumbers in y and x directions\n self.k_y = torch.cat((torch.arange(start=0, end=self.s // 2, step=1, dtype=torch.float32, device=self.device), \\\n torch.arange(start=-self.s // 2, end=0, step=1, dtype=torch.float32, device=self.device)),\n 0).repeat(self.s, 1)\n\n self.k_x = self.k_y.clone().transpose(0, 1)\n\n # Negative inverse Laplacian in Fourier space\n self.inv_lap = (self.k_x ** 2 + self.k_y ** 2)\n self.inv_lap[0, 0] = 1.0\n self.inv_lap = 1.0 / self.inv_lap\n\n # Negative scaled Laplacian\n self.G = (1.0 / self.Re) * (self.k_x ** 2 + self.k_y ** 2)\n\n # Dealiasing mask using 2/3 rule\n self.dealias = (self.k_x ** 2 + self.k_y ** 2 <= (self.s / 3.0) ** 2).float()\n # Ensure mean zero\n self.dealias[0, 0] = 0.0\n\n # Get current vorticity from stream function (Fourier space)\n def vorticity(self, stream_f=None, real_space=True):\n if stream_f is not None:\n w_h = self.Re * self.G * stream_f\n else:\n w_h = self.w_h\n\n if real_space:\n return torch.fft.irfft2(w_h, s=(self.s, self.s), norm=\"backward\")\n else:\n return w_h\n\n # Compute stream function from vorticity (Fourier space)\n def stream_function(self, w_h=None, real_space=False):\n if w_h is None:\n psi_h = self.w_h.clone()\n else:\n psi_h = w_h.clone()\n\n # Stream function in Fourier space: solve Poisson equation\n psi_h = self.inv_lap * psi_h\n\n if real_space:\n return torch.fft.irfft2(psi_h, s=(self.s, self.s), norm=\"backward\")\n else:\n return psi_h\n\n # Compute velocity field from stream function (Fourier space)\n def velocity_field(self, stream_f=None, real_space=True):\n if stream_f is None:\n stream_f = self.stream_function(real_space=False)\n\n # Velocity field in x-direction = psi_y\n q_h = stream_f * 1j * self.k_y\n\n # Velocity field in y-direction = -psi_x\n v_h = stream_f * -1j * self.k_x\n\n if real_space:\n q = torch.fft.irfft2(q_h, s=(self.s, self.s), norm=\"backward\")\n v = torch.fft.irfft2(v_h, s=(self.s, self.s), norm=\"backward\")\n return q, v\n else:\n return q_h, v_h\n\n # Compute non-linear term + forcing from given vorticity (Fourier space)\n def nonlinear_term(self, w_h):\n # Physical space vorticity\n w = torch.fft.ifft2(w_h, s=(self.s, self.s), norm=\"backward\")\n\n # Velocity field in physical space\n q, v = self.velocity_field(self.stream_function(w_h, real_space=False), real_space=True)\n\n # Compute non-linear term\n t1 = torch.fft.fft2(q * w, s=(self.s, self.s), norm=\"backward\")\n t1 = self.k_x * t1\n\n t2 = torch.fft.fft2(v * w, s=(self.s, self.s), norm=\"backward\")\n t2 = self.k_y * t2\n\n nonlin = -1j * (t1 + t2)\n\n # Apply forcing: -ncos(ny)\n if self.n > 0:\n nonlin[..., 0, self.n] -= (float(self.n) / 2.0) * (self.s ** 2)\n nonlin[..., 0, -self.n] -= (float(self.n) / 2.0) * (self.s ** 2)\n\n return nonlin\n\n def advance(self, t, delta_t=1e-3):\n\n # Final time\n T = self.time + t\n\n # Advance solution in Fourier space\n while self.time < T:\n\n if self.time + delta_t > T:\n current_delta_t = T - self.time\n else:\n current_delta_t = delta_t\n\n # Inner-step of Heun's method\n nonlin1 = self.nonlinear_term(self.w_h)\n w_h_tilde = (self.w_h + current_delta_t * (nonlin1 - 0.5 * self.G * self.w_h)) / (\n 1.0 + 0.5 * current_delta_t * self.G)\n\n # Cranck-Nicholson + Heun update\n nonlin2 = self.nonlinear_term(w_h_tilde)\n self.w_h = (self.w_h + current_delta_t * (0.5 * (nonlin1 + nonlin2) - 0.5 * self.G * self.w_h)) / (\n 1.0 + 0.5 * current_delta_t * self.G)\n\n # De-alias\n self.w_h *= self.dealias\n self.time += current_delta_t\n\n\n\n" }, { "path": "solver/legacy_solver.py", "content": "import torch\nimport math\nimport scipy.io\nfrom timeit import default_timer\nfrom tqdm import tqdm\n\n\nclass GaussianRF(object):\n\n def __init__(self, dim, size, alpha=2, tau=3, sigma=None, boundary=\"periodic\", device=None):\n\n self.dim = dim\n self.device = device\n\n if sigma is None:\n sigma = tau**(0.5*(2*alpha - self.dim))\n\n k_max = size//2\n\n if dim == 1:\n k = torch.cat((torch.arange(start=0, end=k_max, step=1, device=device), \\\n torch.arange(start=-k_max, end=0, step=1, device=device)), 0)\n\n self.sqrt_eig = size*math.sqrt(2.0)*sigma*((4*(math.pi**2)*(k**2) + tau**2)**(-alpha/2.0))\n self.sqrt_eig[0] = 0.0\n\n elif dim == 2:\n wavenumers = torch.cat((torch.arange(start=0, end=k_max, step=1, device=device), \\\n torch.arange(start=-k_max, end=0, step=1, device=device)), 0).repeat(size,1)\n\n k_x = wavenumers.transpose(0,1)\n k_y = wavenumers\n\n self.sqrt_eig = (size**2)*math.sqrt(2.0)*sigma*((4*(math.pi**2)*(k_x**2 + k_y**2) + tau**2)**(-alpha/2.0))\n self.sqrt_eig[0,0] = 0.0\n\n elif dim == 3:\n wavenumers = torch.cat((torch.arange(start=0, end=k_max, step=1, device=device), \\\n torch.arange(start=-k_max, end=0, step=1, device=device)), 0).repeat(size,size,1)\n\n k_x = wavenumers.transpose(1,2)\n k_y = wavenumers\n k_z = wavenumers.transpose(0,2)\n\n self.sqrt_eig = (size**3)*math.sqrt(2.0)*sigma*((4*(math.pi**2)*(k_x**2 + k_y**2 + k_z**2) + tau**2)**(-alpha/2.0))\n self.sqrt_eig[0,0,0] = 0.0\n\n self.size = []\n for j in range(self.dim):\n self.size.append(size)\n\n self.size = tuple(self.size)\n\n def sample(self, N):\n\n coeff = torch.randn(N, *self.size, 2, device=self.device)\n\n coeff[...,0] = self.sqrt_eig*coeff[...,0]\n coeff[...,1] = self.sqrt_eig*coeff[...,1]\n\n u = torch.ifft(coeff, self.dim, normalized=False)\n u = u[...,0]\n\n return u\n\n#w0: initial vorticity\n#f: forcing term\n#visc: viscosity (1/Re)\n#T: final time\n#delta_t: internal time-step for solve (descrease if blow-up)\n#record_steps: number of in-time snapshots to record\ndef navier_stokes_2d(w0, f, visc, T, delta_t=1e-4, record_steps=1):\n\n #Grid size - must be power of 2\n N = w0.size()[-1]\n\n #Maximum frequency\n k_max = math.floor(N/2.0)\n\n #Number of steps to final time\n steps = math.ceil(T/delta_t)\n\n #Initial vorticity to Fourier space\n w_h = torch.rfft(w0, 2, normalized=False, onesided=False)\n\n #Forcing to Fourier space\n f_h = torch.rfft(f, 2, normalized=False, onesided=False)\n\n #If same forcing for the whole batch\n if len(f_h.size()) < len(w_h.size()):\n f_h = torch.unsqueeze(f_h, 0)\n\n #Record solution every this number of steps\n record_time = math.floor(steps/record_steps)\n\n #Wavenumbers in y-direction\n k_y = torch.cat((torch.arange(start=0, end=k_max, step=1, device=w0.device), torch.arange(start=-k_max, end=0, step=1, device=w0.device)), 0).repeat(N,1)\n #Wavenumbers in x-direction\n k_x = k_y.transpose(0,1)\n #Negative Laplacian in Fourier space\n lap = 4*(math.pi**2)*(k_x**2 + k_y**2)\n lap[0,0] = 1.0\n #Dealiasing mask\n dealias = torch.unsqueeze(torch.logical_and(torch.abs(k_y) <= (2.0/3.0)*k_max, torch.abs(k_x) <= (2.0/3.0)*k_max).float(), 0)\n\n #Saving solution and time\n sol = torch.zeros(*w0.size(), record_steps, device=w0.device)\n sol_t = torch.zeros(record_steps, device=w0.device)\n\n #Record counter\n c = 0\n #Physical time\n t = 0.0\n for j in tqdm(range(steps)):\n #Stream function in Fourier space: solve Poisson equation\n psi_h = w_h.clone()\n psi_h[...,0] = psi_h[...,0]/lap\n psi_h[...,1] = psi_h[...,1]/lap\n\n #Velocity field in x-direction = psi_y\n q = psi_h.clone()\n temp = q[...,0].clone()\n q[...,0] = -2*math.pi*k_y*q[...,1]\n q[...,1] = 2*math.pi*k_y*temp\n q = torch.irfft(q, 2, normalized=False, onesided=False, signal_sizes=(N,N))\n\n #Velocity field in y-direction = -psi_x\n v = psi_h.clone()\n temp = v[...,0].clone()\n v[...,0] = 2*math.pi*k_x*v[...,1]\n v[...,1] = -2*math.pi*k_x*temp\n v = torch.irfft(v, 2, normalized=False, onesided=False, signal_sizes=(N,N))\n\n #Partial x of vorticity\n w_x = w_h.clone()\n temp = w_x[...,0].clone()\n w_x[...,0] = -2*math.pi*k_x*w_x[...,1]\n w_x[...,1] = 2*math.pi*k_x*temp\n w_x = torch.irfft(w_x, 2, normalized=False, onesided=False, signal_sizes=(N,N))\n\n #Partial y of vorticity\n w_y = w_h.clone()\n temp = w_y[...,0].clone()\n w_y[...,0] = -2*math.pi*k_y*w_y[...,1]\n w_y[...,1] = 2*math.pi*k_y*temp\n w_y = torch.irfft(w_y, 2, normalized=False, onesided=False, signal_sizes=(N,N))\n\n #Non-linear term (u.grad(w)): compute in physical space then back to Fourier space\n F_h = torch.rfft(q*w_x + v*w_y, 2, normalized=False, onesided=False)\n\n #Dealias\n F_h[...,0] = dealias* F_h[...,0]\n F_h[...,1] = dealias* F_h[...,1]\n\n #Cranck-Nicholson update\n w_h[...,0] = (-delta_t*F_h[...,0] + delta_t*f_h[...,0] + (1.0 - 0.5*delta_t*visc*lap)*w_h[...,0])/(1.0 + 0.5*delta_t*visc*lap)\n w_h[...,1] = (-delta_t*F_h[...,1] + delta_t*f_h[...,1] + (1.0 - 0.5*delta_t*visc*lap)*w_h[...,1])/(1.0 + 0.5*delta_t*visc*lap)\n\n #Update real time (used only for recording)\n t += delta_t\n\n if (j+1) % record_time == 0:\n #Solution in physical space\n w = torch.irfft(w_h, 2, normalized=False, onesided=False, signal_sizes=(N,N))\n\n #Record solution and time\n sol[...,c] = w\n sol_t[c] = t\n\n c += 1\n return sol, sol_t\n\n\nif __name__ == '__main__':\n device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n # Resolution\n # s = 2048\n # sub = 1\n #\n # # Number of solutions to generate\n # N = 10\n #\n # # Set up 2d GRF with covariance parameters\n # GRF = GaussianRF(2, s, alpha=2.5, tau=7, device=device)\n #\n # # Forcing function: 0.1*(sin(2pi(x+y)) + cos(2pi(x+y)))\n # t = torch.linspace(0, 1, s + 1, device=device)\n # t = t[0:-1]\n #\n # X, Y = torch.meshgrid(t, t)\n # f = 0.1 * (torch.sin(2 * math.pi * (X + Y)) + torch.cos(2 * math.pi * (X + Y)))\n #\n # # Number of snapshots from solution\n # record_steps = 200\n #\n # # Inputs\n # a = torch.zeros(N, s, s)\n # # Solutions\n # u = torch.zeros(N, s, s, record_steps)\n #\n # # Solve equations in batches (order of magnitude speed-up)\n #\n # # Batch size\n # bsize = 10\n #\n # c = 0\n # t0 = default_timer()\n # for j in range(N // bsize):\n # # Sample random feilds\n # w0 = GRF.sample(bsize)\n #\n # # Solve NS\n # sol, sol_t = navier_stokes_2d(w0, f, 1e-3, 50.0, 1e-4, record_steps)\n #\n # a[c:(c + bsize), ...] = w0\n # u[c:(c + bsize), ...] = sol\n #\n # c += bsize\n # t1 = default_timer()\n # print(j, c, t1 - t0)\n # torch.save(\n # {\n # 'a': a.cpu(),\n # 'u': u.cpu(),\n # 't': sol_t.cpu()\n # },\n # 'data/ns_data.pt'\n # )\n # scipy.io.savemat('data/ns_data.mat', mdict={'a': a.cpu().numpy(), 'u': u.cpu().numpy(), 't': sol_t.cpu().numpy()})" }, { "path": "solver/periodic.py", "content": "import torch\nimport torch.fft as fft\n\nimport math\n\n#Setup for indexing in the 'ij' format\n\n#Solve: -Lap(u) = f\nclass Poisson2d(object):\n\n def __init__(self, s1, s2, L1=2*math.pi, L2=2*math.pi, device=None, dtype=torch.float64):\n\n self.s1 = s1\n self.s2 = s2\n\n #Inverse negative Laplacian\n freq_list1 = torch.cat((torch.arange(start=0, end=s1//2, step=1),\\\n torch.arange(start=-s1//2, end=0, step=1)), 0)\n k1 = freq_list1.view(-1,1).repeat(1, s2//2 + 1).type(dtype).to(device)\n\n freq_list2 = torch.arange(start=0, end=s2//2 + 1, step=1)\n k2 = freq_list2.view(1,-1).repeat(s1, 1).type(dtype).to(device)\n\n self.inv_lap = ((4*math.pi**2)/(L1**2))*k1**2 + ((4*math.pi**2)/(L2**2))*k2**2\n self.inv_lap[0,0] = 1.0\n self.inv_lap = 1.0/self.inv_lap\n \n def solve(self, f):\n return fft.irfft2(fft.rfft2(f)*self.inv_lap, s=(self.s1, self.s2))\n \n def __call__(self, f):\n return self.solve(f)\n\n#Solve: w_t = - u . grad(w) + (1/Re)*Lap(w) + f\n# u = (psi_y, -psi_x)\n# -Lap(psi) = w\n#Note: Adaptive time-step takes smallest step across the batch\nclass NavierStokes2d(object):\n\n def __init__(self, s1, s2, L1=2*math.pi, L2=2*math.pi, device=None, dtype=torch.float64):\n\n self.s1 = s1\n self.s2 = s2\n\n self.L1 = L1\n self.L2 = L2\n\n self.h = 1.0/max(s1, s2)\n\n #Wavenumbers for first derivatives\n freq_list1 = torch.cat((torch.arange(start=0, end=s1//2, step=1),\\\n torch.zeros((1,)),\\\n torch.arange(start=-s1//2 + 1, end=0, step=1)), 0)\n self.k1 = freq_list1.view(-1,1).repeat(1, s2//2 + 1).type(dtype).to(device)\n\n\n freq_list2 = torch.cat((torch.arange(start=0, end=s2//2, step=1), torch.zeros((1,))), 0)\n self.k2 = freq_list2.view(1,-1).repeat(s1, 1).type(dtype).to(device)\n\n #Negative Laplacian\n freq_list1 = torch.cat((torch.arange(start=0, end=s1//2, step=1),\\\n torch.arange(start=-s1//2, end=0, step=1)), 0)\n k1 = freq_list1.view(-1,1).repeat(1, s2//2 + 1).type(dtype).to(device)\n\n freq_list2 = torch.arange(start=0, end=s2//2 + 1, step=1)\n k2 = freq_list2.view(1,-1).repeat(s1, 1).type(dtype).to(device)\n\n self.G = ((4*math.pi**2)/(L1**2))*k1**2 + ((4*math.pi**2)/(L2**2))*k2**2\n\n #Inverse of negative Laplacian\n self.inv_lap = self.G.clone()\n self.inv_lap[0,0] = 1.0\n self.inv_lap = 1.0/self.inv_lap\n\n #Dealiasing mask using 2/3 rule\n self.dealias = (self.k1**2 + self.k2**2 <= 0.6*(0.25*s1**2 + 0.25*s2**2)).type(dtype).to(device)\n #Ensure mean zero\n self.dealias[0,0] = 0.0\n\n #Compute stream function from vorticity (Fourier space)\n def stream_function(self, w_h, real_space=False):\n #-Lap(psi) = w\n psi_h = self.inv_lap*w_h\n\n if real_space:\n return fft.irfft2(psi_h, s=(self.s1, self.s2))\n else:\n return psi_h\n\n #Compute velocity field from stream function (Fourier space)\n def velocity_field(self, stream_f, real_space=True):\n #Velocity field in x-direction = psi_y\n q_h = (2*math.pi/self.L2)*1j*self.k2*stream_f\n\n #Velocity field in y-direction = -psi_x\n v_h = -(2*math.pi/self.L1)*1j*self.k1*stream_f\n\n if real_space:\n return fft.irfft2(q_h, s=(self.s1, self.s2)), fft.irfft2(v_h, s=(self.s1, self.s2))\n else:\n return q_h, v_h\n\n #Compute non-linear term + forcing from given vorticity (Fourier space)\n def nonlinear_term(self, w_h, f_h=None):\n #Physical space vorticity\n w = fft.irfft2(w_h, s=(self.s1, self.s2))\n\n #Physical space velocity\n q, v = self.velocity_field(self.stream_function(w_h, real_space=False), real_space=True)\n\n #Compute non-linear term in Fourier space\n nonlin = -1j*((2*math.pi/self.L1)*self.k1*fft.rfft2(q*w) + (2*math.pi/self.L1)*self.k2*fft.rfft2(v*w))\n\n #Add forcing function\n if f_h is not None:\n nonlin += f_h\n\n return nonlin\n \n def time_step(self, q, v, f, Re):\n #Maxixum speed\n max_speed = torch.max(torch.sqrt(q**2 + v**2)).item()\n\n #Maximum force amplitude\n if f is not None:\n xi = torch.sqrt(torch.max(torch.abs(f))).item()\n else:\n xi = 1.0\n \n #Viscosity\n mu = (1.0/Re)*xi*((self.L1/(2*math.pi))**(3.0/4.0))*(((self.L2/(2*math.pi))**(3.0/4.0)))\n\n if max_speed == 0:\n return 0.5*(self.h**2)/mu\n \n #Time step based on CFL condition\n return min(0.5*self.h/max_speed, 0.5*(self.h**2)/mu)\n\n def advance(self, w, f=None, T=1.0, Re=100, adaptive=True, delta_t=1e-3):\n\n #Rescale Laplacian by Reynolds number\n GG = (1.0/Re)*self.G\n\n #Move to Fourier space\n w_h = fft.rfft2(w)\n\n if f is not None:\n f_h = fft.rfft2(f)\n else:\n f_h = None\n \n if adaptive:\n q, v = self.velocity_field(self.stream_function(w_h, real_space=False), real_space=True)\n delta_t = self.time_step(q, v, f, Re)\n\n time = 0.0\n #Advance solution in Fourier space\n while time < T:\n if time + delta_t > T:\n current_delta_t = T - time\n else:\n current_delta_t = delta_t\n\n #Inner-step of Heun's method\n nonlin1 = self.nonlinear_term(w_h, f_h)\n w_h_tilde = (w_h + current_delta_t*(nonlin1 - 0.5*GG*w_h))/(1.0 + 0.5*current_delta_t*GG)\n\n #Cranck-Nicholson + Heun update\n nonlin2 = self.nonlinear_term(w_h_tilde, f_h)\n w_h = (w_h + current_delta_t*(0.5*(nonlin1 + nonlin2) - 0.5*GG*w_h))/(1.0 + 0.5*current_delta_t*GG)\n\n #De-alias\n w_h *= self.dealias\n\n #Update time\n time += current_delta_t\n\n #New time step\n if adaptive:\n q, v = self.velocity_field(self.stream_function(w_h, real_space=False), real_space=True)\n delta_t = self.time_step(q, v, f, Re)\n \n return fft.irfft2(w_h, s=(self.s1, self.s2))\n \n def __call__(self, w, f=None, T=1.0, Re=100, adaptive=True, delta_t=1e-3):\n return self.advance(w, f, T, Re, adaptive, delta_t)" }, { "path": "solver/random_fields.py", "content": "import torch\n\nimport math\n\ntorch.manual_seed(0)\n\n\nclass GaussianRF(object):\n def __init__(self, dim, size, length=1.0, alpha=2.0, tau=3.0, sigma=None, boundary=\"periodic\", constant_eig=False, device=None):\n\n self.dim = dim\n self.device = device\n\n if sigma is None:\n sigma = tau**(0.5*(2*alpha - self.dim))\n\n k_max = size//2\n\n const = (4*(math.pi**2))/(length**2)\n\n if dim == 1:\n k = torch.cat((torch.arange(start=0, end=k_max, step=1, device=device), \\\n torch.arange(start=-k_max, end=0, step=1, device=device)), 0)\n\n self.sqrt_eig = size*math.sqrt(2.0)*sigma*((const*(k**2) + tau**2)**(-alpha/2.0))\n\n if constant_eig:\n self.sqrt_eig[0] = size*sigma*(tau**(-alpha))\n else:\n self.sqrt_eig[0] = 0.0\n\n elif dim == 2:\n wavenumers = torch.cat((torch.arange(start=0, end=k_max, step=1, device=device), \\\n torch.arange(start=-k_max, end=0, step=1, device=device)), 0).repeat(size,1)\n\n k_x = wavenumers.transpose(0,1)\n k_y = wavenumers\n\n self.sqrt_eig = (size**2)*math.sqrt(2.0)*sigma*((const*(k_x**2 + k_y**2) + tau**2)**(-alpha/2.0))\n\n if constant_eig:\n self.sqrt_eig[0,0] = (size**2)*sigma*(tau**(-alpha))\n else:\n self.sqrt_eig[0,0] = 0.0\n\n elif dim == 3:\n wavenumers = torch.cat((torch.arange(start=0, end=k_max, step=1, device=device), \\\n torch.arange(start=-k_max, end=0, step=1, device=device)), 0).repeat(size,size,1)\n\n k_x = wavenumers.transpose(1,2)\n k_y = wavenumers\n k_z = wavenumers.transpose(0,2)\n\n self.sqrt_eig = (size**3)*math.sqrt(2.0)*sigma*((const*(k_x**2 + k_y**2 + k_z**2) + tau**2)**(-alpha/2.0))\n\n if constant_eig:\n self.sqrt_eig[0,0,0] = (size**3)*sigma*(tau**(-alpha))\n else:\n self.sqrt_eig[0,0,0] = 0.0\n\n self.size = []\n for j in range(self.dim):\n self.size.append(size)\n\n self.size = tuple(self.size)\n\n def sample(self, N):\n\n coeff = torch.randn(N, *self.size, dtype=torch.cfloat, device=self.device)\n coeff = self.sqrt_eig*coeff\n\n u = torch.fft.irfftn(coeff, self.size, norm=\"backward\")\n return u\n\n\nclass GaussianRF2d(object):\n\n def __init__(self, s1, s2, L1=2*math.pi, L2=2*math.pi, alpha=2.0, tau=3.0, sigma=None, mean=None, boundary=\"periodic\", device=None, dtype=torch.float64):\n\n self.s1 = s1\n self.s2 = s2\n\n self.mean = mean\n\n self.device = device\n self.dtype = dtype\n\n if sigma is None:\n self.sigma = tau**(0.5*(2*alpha - 2.0))\n else:\n self.sigma = sigma\n\n const1 = (4*(math.pi**2))/(L1**2)\n const2 = (4*(math.pi**2))/(L2**2)\n\n freq_list1 = torch.cat((torch.arange(start=0, end=s1//2, step=1),\\\n torch.arange(start=-s1//2, end=0, step=1)), 0)\n k1 = freq_list1.view(-1,1).repeat(1, s2//2 + 1).type(dtype).to(device)\n\n freq_list2 = torch.arange(start=0, end=s2//2 + 1, step=1)\n\n k2 = freq_list2.view(1,-1).repeat(s1, 1).type(dtype).to(device)\n\n self.sqrt_eig = s1*s2*self.sigma*((const1*k1**2 + const2*k2**2 + tau**2)**(-alpha/2.0))\n self.sqrt_eig[0,0] = 0.0\n\n def sample(self, N, xi=None):\n if xi is None:\n xi = torch.randn(N, self.s1, self.s2//2 + 1, 2, dtype=self.dtype, device=self.device)\n \n xi[...,0] = self.sqrt_eig*xi [...,0]\n xi[...,1] = self.sqrt_eig*xi [...,1]\n \n u = torch.fft.irfft2(torch.view_as_complex(xi), s=(self.s1, self.s2))\n\n if self.mean is not None:\n u += self.mean\n \n return u" }, { "path": "solver/rfsampler.py", "content": "import torch\nimport math\n\n\nclass GaussianRF(object):\n def __init__(self, dim, size, alpha=2, tau=3, sigma=None, boundary=\"periodic\", device=None):\n self.dim = dim\n self.device = device\n if sigma is None:\n sigma = tau ** (0.5 * (2 * alpha - self.dim))\n\n k_max = size // 2\n\n if dim == 1:\n k = torch.cat((torch.arange(start=0, end=k_max, step=1, device=device), \\\n torch.arange(start=-k_max, end=0, step=1, device=device)), 0)\n\n self.sqrt_eig = size * math.sqrt(2.0) * sigma * (\n (4 * (math.pi ** 2) * (k ** 2) + tau ** 2) ** (-alpha / 2.0))\n self.sqrt_eig[0] = 0.0\n\n elif dim == 2:\n wavenumers = torch.cat((torch.arange(start=0, end=k_max, step=1, device=device), \\\n torch.arange(start=-k_max, end=0, step=1, device=device)), 0).repeat(size, 1)\n\n k_x = wavenumers.transpose(0, 1)\n k_y = wavenumers\n\n self.sqrt_eig = (size ** 2) * math.sqrt(2.0) * sigma * (\n (4 * (math.pi ** 2) * (k_x ** 2 + k_y ** 2) + tau ** 2) ** (-alpha / 2.0))\n self.sqrt_eig[0, 0] = 0.0\n\n elif dim == 3:\n wavenumers = torch.cat((torch.arange(start=0, end=k_max, step=1, device=device), \\\n torch.arange(start=-k_max, end=0, step=1, device=device)), 0).repeat(size, size, 1)\n\n k_x = wavenumers.transpose(1, 2)\n k_y = wavenumers\n k_z = wavenumers.transpose(0, 2)\n\n self.sqrt_eig = (size ** 3) * math.sqrt(2.0) * sigma * (\n (4 * (math.pi ** 2) * (k_x ** 2 + k_y ** 2 + k_z ** 2) + tau ** 2) ** (-alpha / 2.0))\n self.sqrt_eig[0, 0, 0] = 0.0\n\n self.size = []\n for j in range(self.dim):\n self.size.append(size)\n\n self.size = tuple(self.size)\n\n def sample(self, N):\n\n coeff = torch.randn(N, *self.size, 2, device=self.device)\n\n coeff[..., 0] = self.sqrt_eig * coeff[..., 0]\n coeff[..., 1] = self.sqrt_eig * coeff[..., 1]\n\n u = torch.ifft(coeff, self.dim, normalized=False)\n u = u[..., 0]\n\n return u\n" }, { "path": "solver/spectrum.py", "content": "import math\n\nfrom rfsampler import GaussianRF\nimport torch\nfrom timeit import default_timer\n\nimport scipy.io\n\n\n# w0: initial vorticity\n# f: forcing term\n# visc: viscosity (1/Re)\n# T: final time\n# delta_t: internal time-step for solve (descrease if blow-up)\n# record_steps: number of in-time snapshots to record\n\ndef navier_stokes_2d(w0, f, visc, T, delta_t=1e-4, record_steps=1):\n '''\n Args:\n w0: initial vorticity\n f: forcing\n visc: 1/Re\n T: final time\n delta_t: internal time-step for solve (decrease if blow-up)\n record_steps: number of in-time snapshots to save\n Returns:\n\n '''\n # Grid size - must be power of 2\n N = w0.size()[-1]\n\n # Maximum frequency\n k_max = math.floor(N / 2.0)\n\n # Number of steps to final time\n steps = math.ceil(T / delta_t)\n\n # Initial vorticity to Fourier space\n w_h = torch.rfft(w0, 2, normalized=False, onesided=False)\n\n # Forcing to Fourier space\n f_h = torch.rfft(f, 2, normalized=False, onesided=False)\n\n # If same forcing for the whole batch\n if len(f_h.size()) < len(w_h.size()):\n f_h = torch.unsqueeze(f_h, 0)\n\n # Record solution every this number of steps\n record_time = math.floor(steps / record_steps)\n\n # Wavenumbers in y-direction\n k_y = torch.cat((torch.arange(start=0, end=k_max, step=1, device=w0.device),\n torch.arange(start=-k_max, end=0, step=1, device=w0.device)), 0).repeat(N, 1)\n # Wavenumbers in x-direction\n k_x = k_y.transpose(0, 1)\n # Negative Laplacian in Fourier space\n lap = 4 * (math.pi ** 2) * (k_x ** 2 + k_y ** 2)\n lap[0, 0] = 1.0\n # Dealiasing mask\n dealias = torch.unsqueeze(\n torch.logical_and(torch.abs(k_y) <= (2.0 / 3.0) * k_max, torch.abs(k_x) <= (2.0 / 3.0) * k_max).float(), 0)\n\n # Saving solution and time\n sol = torch.zeros(*w0.size(), record_steps, device=w0.device)\n sol_t = torch.zeros(record_steps, device=w0.device)\n\n # Record counter\n c = 0\n # Physical time\n t = 0.0\n for j in range(steps):\n # Stream function in Fourier space: solve Poisson equation\n psi_h = w_h.clone()\n psi_h[..., 0] = psi_h[..., 0] / lap\n psi_h[..., 1] = psi_h[..., 1] / lap\n\n # Velocity field in x-direction = psi_y\n q = psi_h.clone()\n temp = q[..., 0].clone()\n q[..., 0] = -2 * math.pi * k_y * q[..., 1]\n q[..., 1] = 2 * math.pi * k_y * temp\n q = torch.irfft(q, 2, normalized=False, onesided=False, signal_sizes=(N, N))\n\n # Velocity field in y-direction = -psi_x\n v = psi_h.clone()\n temp = v[..., 0].clone()\n v[..., 0] = 2 * math.pi * k_x * v[..., 1]\n v[..., 1] = -2 * math.pi * k_x * temp\n v = torch.irfft(v, 2, normalized=False, onesided=False, signal_sizes=(N, N))\n\n # Partial x of vorticity\n w_x = w_h.clone()\n temp = w_x[..., 0].clone()\n w_x[..., 0] = -2 * math.pi * k_x * w_x[..., 1]\n w_x[..., 1] = 2 * math.pi * k_x * temp\n w_x = torch.irfft(w_x, 2, normalized=False, onesided=False, signal_sizes=(N, N))\n\n # Partial y of vorticity\n w_y = w_h.clone()\n temp = w_y[..., 0].clone()\n w_y[..., 0] = -2 * math.pi * k_y * w_y[..., 1]\n w_y[..., 1] = 2 * math.pi * k_y * temp\n w_y = torch.irfft(w_y, 2, normalized=False, onesided=False, signal_sizes=(N, N))\n\n # Non-linear term (u.grad(w)): compute in physical space then back to Fourier space\n F_h = torch.rfft(q * w_x + v * w_y, 2, normalized=False, onesided=False)\n\n # Dealias\n F_h[..., 0] = dealias * F_h[..., 0]\n F_h[..., 1] = dealias * F_h[..., 1]\n\n # Cranck-Nicholson update\n w_h[..., 0] = (-delta_t * F_h[..., 0] + delta_t * f_h[..., 0] + (1.0 - 0.5 * delta_t * visc * lap) * w_h[\n ..., 0]) / (1.0 + 0.5 * delta_t * visc * lap)\n w_h[..., 1] = (-delta_t * F_h[..., 1] + delta_t * f_h[..., 1] + (1.0 - 0.5 * delta_t * visc * lap) * w_h[\n ..., 1]) / (1.0 + 0.5 * delta_t * visc * lap)\n\n # Update real time (used only for recording)\n t += delta_t\n\n if (j + 1) % record_time == 0:\n # Solution in physical space\n w = torch.irfft(w_h, 2, normalized=False, onesided=False, signal_sizes=(N, N))\n\n # Record solution and time\n sol[..., c] = w\n sol_t[c] = t\n\n c += 1\n return sol, sol_t\n\n\nif __name__ == '__main__':\n device = torch.device('cuda')\n\n # Resolution\n s = 256\n sub = 1\n\n # Number of solutions to generate\n N = 20\n\n # Set up 2d GRF with covariance parameters\n GRF = GaussianRF(2, s, alpha=2.5, tau=7, device=device)\n\n # Forcing function: 0.1*(sin(2pi(x+y)) + cos(2pi(x+y)))\n t = torch.linspace(0, 1, s + 1, device=device)\n t = t[0:-1]\n\n X, Y = torch.meshgrid(t, t)\n f = 0.1 * (torch.sin(2 * math.pi * (X + Y)) + torch.cos(2 * math.pi * (X + Y)))\n\n # Number of snapshots from solution\n record_steps = 200\n\n # Inputs\n a = torch.zeros(N, s, s)\n # Solutions\n u = torch.zeros(N, s, s, record_steps)\n\n # Solve equations in batches (order of magnitude speed-up)\n\n # Batch size\n bsize = 20\n\n c = 0\n t0 = default_timer()\n for j in range(N // bsize):\n # Sample random feilds\n w0 = GRF.sample(bsize)\n\n # Solve NS\n sol, sol_t = navier_stokes_2d(w0, f, 1e-3, 50.0, 1e-4, record_steps)\n\n a[c:(c + bsize), ...] = w0\n u[c:(c + bsize), ...] = sol\n\n c += bsize\n t1 = default_timer()\n print(j, c, t1 - t0)\n\n scipy.io.savemat('ns_data.mat', mdict={'a': a.cpu().numpy(), 'u': u.cpu().numpy(), 't': sol_t.cpu().numpy()})\n" }, { "path": "train_PINO3d.py", "content": "import yaml\nfrom argparse import ArgumentParser\nimport random\nimport torch\nfrom torch.utils.data import DataLoader\nimport torch.multiprocessing as mp\nfrom torch.nn.parallel import DistributedDataParallel as DDP\n\nfrom train_utils import Adam\nfrom train_utils.datasets import NSLoader\nfrom train_utils.data_utils import data_sampler\nfrom train_utils.losses import get_forcing\nfrom train_utils.train_3d import train\nfrom train_utils.distributed import setup, cleanup\nfrom train_utils.utils import requires_grad\n\nfrom models import FNO3d, FNO2d\n\n\ndef subprocess_fn(rank, args):\n if args.distributed:\n setup(rank, args.num_gpus)\n print(f'Running on rank {rank}')\n\n config_file = args.config_path\n with open(config_file, 'r') as stream:\n config = yaml.load(stream, yaml.FullLoader)\n\n # construct dataloader\n data_config = config['data']\n if 'datapath2' in data_config:\n loader = NSLoader(datapath1=data_config['datapath'], datapath2=data_config['datapath2'],\n nx=data_config['nx'], nt=data_config['nt'],\n sub=data_config['sub'], sub_t=data_config['sub_t'],\n N=data_config['total_num'],\n t_interval=data_config['time_interval'])\n else:\n loader = NSLoader(datapath1=data_config['datapath'],\n nx=data_config['nx'], nt=data_config['nt'],\n sub=data_config['sub'], sub_t=data_config['sub_t'],\n N=data_config['total_num'],\n t_interval=data_config['time_interval'])\n if args.start != -1:\n config['data']['offset'] = args.start\n trainset = loader.make_dataset(data_config['n_sample'],\n start=data_config['offset'])\n train_loader = DataLoader(trainset, batch_size=config['train']['batchsize'],\n sampler=data_sampler(trainset,\n shuffle=data_config['shuffle'],\n distributed=args.distributed),\n drop_last=True)\n\n # construct model\n model = FNO3d(modes1=config['model']['modes1'],\n modes2=config['model']['modes2'],\n modes3=config['model']['modes3'],\n fc_dim=config['model']['fc_dim'],\n layers=config['model']['layers']).to(rank)\n\n if 'ckpt' in config['train']:\n ckpt_path = config['train']['ckpt']\n ckpt = torch.load(ckpt_path)\n model.load_state_dict(ckpt['model'])\n print('Weights loaded from %s' % ckpt_path)\n\n if args.distributed:\n model = DDP(model, device_ids=[rank], broadcast_buffers=False)\n\n if 'twolayer' in config['train'] and config['train']['twolayer']:\n requires_grad(model, False)\n requires_grad(model.sp_convs[-1], True)\n requires_grad(model.ws[-1], True)\n requires_grad(model.fc1, True)\n requires_grad(model.fc2, True)\n params = []\n for param in model.parameters():\n if param.requires_grad == True:\n params.append(param)\n else:\n params = model.parameters()\n\n optimizer = Adam(params, betas=(0.9, 0.999),\n lr=config['train']['base_lr'])\n scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer,\n milestones=config['train']['milestones'],\n gamma=config['train']['scheduler_gamma'])\n forcing = get_forcing(loader.S).to(rank)\n train(model,\n loader, train_loader,\n optimizer, scheduler,\n forcing, config,\n rank,\n log=args.log,\n project=config['log']['project'],\n group=config['log']['group'])\n\n if args.distributed:\n cleanup()\n print(f'Process {rank} done!...')\n\n\nif __name__ == '__main__':\n seed = random.randint(1, 10000)\n print(f'Random seed :{seed}')\n torch.manual_seed(seed)\n torch.cuda.manual_seed_all(seed)\n torch.backends.cudnn.benchmark = True\n parser =ArgumentParser(description='Basic paser')\n parser.add_argument('--config_path', type=str, help='Path to the configuration file')\n parser.add_argument('--log', action='store_true', help='Turn on the wandb')\n parser.add_argument('--num_gpus', type=int, help='Number of GPUs', default=1)\n parser.add_argument('--start', type=int, default=-1, help='start index')\n args = parser.parse_args()\n args.distributed = args.num_gpus > 1\n\n if args.distributed:\n mp.spawn(subprocess_fn, args=(args, ), nprocs=args.num_gpus)\n else:\n subprocess_fn(0, args)\n\n" }, { "path": "train_burgers.py", "content": "from argparse import ArgumentParser\nimport yaml\n\nimport torch\nfrom models import FNO2d\nfrom train_utils import Adam\nfrom train_utils.datasets import BurgersLoader\nfrom train_utils.train_2d import train_2d_burger\nfrom train_utils.eval_2d import eval_burgers\n\n\ndef run(args, config):\n device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n data_config = config['data']\n dataset = BurgersLoader(data_config['datapath'],\n nx=data_config['nx'], nt=data_config['nt'],\n sub=data_config['sub'], sub_t=data_config['sub_t'], new=True)\n train_loader = dataset.make_loader(n_sample=data_config['n_sample'],\n batch_size=config['train']['batchsize'],\n start=data_config['offset'])\n\n model = FNO2d(modes1=config['model']['modes1'],\n modes2=config['model']['modes2'],\n fc_dim=config['model']['fc_dim'],\n layers=config['model']['layers'],\n act=config['model']['act']).to(device)\n # Load from checkpoint\n if 'ckpt' in config['train']:\n ckpt_path = config['train']['ckpt']\n ckpt = torch.load(ckpt_path)\n model.load_state_dict(ckpt['model'])\n print('Weights loaded from %s' % ckpt_path)\n optimizer = Adam(model.parameters(), betas=(0.9, 0.999),\n lr=config['train']['base_lr'])\n scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer,\n milestones=config['train']['milestones'],\n gamma=config['train']['scheduler_gamma'])\n train_2d_burger(model,\n train_loader,\n dataset.v,\n optimizer,\n scheduler,\n config,\n rank=0,\n log=args.log,\n project=config['log']['project'],\n group=config['log']['group'])\n\n\ndef test(config):\n device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n data_config = config['data']\n dataset = BurgersLoader(data_config['datapath'],\n nx=data_config['nx'], nt=data_config['nt'],\n sub=data_config['sub'], sub_t=data_config['sub_t'], new=True)\n dataloader = dataset.make_loader(n_sample=data_config['n_sample'],\n batch_size=config['test']['batchsize'],\n start=data_config['offset'])\n\n model = FNO2d(modes1=config['model']['modes1'],\n modes2=config['model']['modes2'],\n fc_dim=config['model']['fc_dim'],\n layers=config['model']['layers'],\n act=config['model']['act']).to(device)\n # Load from checkpoint\n if 'ckpt' in config['test']:\n ckpt_path = config['test']['ckpt']\n ckpt = torch.load(ckpt_path)\n model.load_state_dict(ckpt['model'])\n print('Weights loaded from %s' % ckpt_path)\n eval_burgers(model, dataloader, dataset.v, config, device)\n\n\nif __name__ == '__main__':\n parser = ArgumentParser(description='Basic paser')\n parser.add_argument('--config_path', type=str, help='Path to the configuration file')\n parser.add_argument('--log', action='store_true', help='Turn on the wandb')\n parser.add_argument('--mode', type=str, help='train or test')\n args = parser.parse_args()\n\n config_file = args.config_path\n with open(config_file, 'r') as stream:\n config = yaml.load(stream, yaml.FullLoader)\n if args.mode == 'train':\n run(args, config)\n else:\n test(config)\n" }, { "path": "train_darcy.py", "content": "import os\nimport yaml\nimport random\nfrom argparse import ArgumentParser\nfrom tqdm import tqdm\n\nimport numpy as np\nimport torch\n\nfrom torch.optim import Adam\nfrom torch.utils.data import DataLoader\n\nfrom models import FNO2d\n\nfrom train_utils.losses import LpLoss, darcy_loss \nfrom train_utils.datasets import DarcyFlow, DarcyIC, sample_data\nfrom train_utils.utils import save_ckpt, count_params, dict2str\n\ntry:\n import wandb\nexcept ImportError:\n wandb = None\n\n\n\ndef get_molifier(mesh, device):\n mollifier = 0.001 * torch.sin(np.pi * mesh[..., 0]) * torch.sin(np.pi * mesh[..., 1])\n return mollifier.to(device)\n\n\n@torch.no_grad()\ndef eval_darcy(model, val_loader, criterion, \n device='cpu'):\n mollifier = get_molifier(val_loader.dataset.mesh, device)\n model.eval()\n val_err = []\n for a, u in val_loader:\n a, u = a.to(device), u.to(device)\n out = model(a).squeeze(dim=-1)\n out = out * mollifier\n val_loss = criterion(out, u)\n val_err.append(val_loss.item())\n N = len(val_loader)\n\n avg_err = np.mean(val_err)\n std_err = np.std(val_err, ddof=1) / np.sqrt(N)\n return avg_err, std_err\n\n\ndef train(model, \n train_u_loader, # training data\n ic_loader, # loader for initial conditions\n val_loader, # validation data\n optimizer, \n scheduler,\n device, config, args):\n save_step = config['train']['save_step']\n eval_step = config['train']['eval_step']\n\n f_weight = config['train']['f_loss']\n xy_weight = config['train']['xy_loss']\n\n # set up directory\n base_dir = os.path.join('exp', config['log']['logdir'])\n ckpt_dir = os.path.join(base_dir, 'ckpts')\n os.makedirs(ckpt_dir, exist_ok=True)\n\n # loss fn\n lploss = LpLoss(size_average=True)\n # mollifier\n u_mol = get_molifier(train_u_loader.dataset.mesh, device)\n ic_mol = get_molifier(ic_loader.dataset.mesh, device)\n # set up wandb\n if wandb and args.log:\n run = wandb.init(project=config['log']['project'], \n entity=config['log']['entity'], \n group=config['log']['group'], \n config=config, reinit=True, \n settings=wandb.Settings(start_method='fork'))\n pbar = range(config['train']['num_iter'])\n pbar = tqdm(pbar, dynamic_ncols=True, smoothing=0.2)\n\n u_loader = sample_data(train_u_loader)\n ic_loader = sample_data(ic_loader)\n for e in pbar:\n log_dict = {}\n\n optimizer.zero_grad()\n # data loss\n if xy_weight > 0:\n ic, u = next(u_loader)\n u = u.to(device)\n ic = ic.to(device)\n out = model(ic).squeeze(dim=-1)\n out = out * u_mol\n data_loss = lploss(out, u)\n else:\n data_loss = torch.zeros(1, device=device)\n \n if f_weight > 0:\n # pde loss\n ic = next(ic_loader)\n ic = ic.to(device)\n out = model(ic).squeeze(dim=-1)\n out = out * ic_mol\n u0 = ic[..., 0]\n f_loss = darcy_loss(out, u0)\n log_dict['PDE'] = f_loss.item()\n else:\n f_loss = 0.0\n loss = data_loss * xy_weight + f_loss * f_weight\n\n loss.backward()\n optimizer.step()\n scheduler.step()\n\n log_dict['train loss'] = loss.item()\n log_dict['data'] = data_loss.item()\n if e % eval_step == 0:\n eval_err, std_err = eval_darcy(model, val_loader, lploss, device)\n log_dict['val error'] = eval_err\n logstr = dict2str(log_dict)\n pbar.set_description(\n (\n logstr\n )\n )\n if wandb and args.log:\n wandb.log(log_dict)\n if e % save_step == 0 and e > 0:\n ckpt_path = os.path.join(ckpt_dir, f'model-{e}.pt')\n save_ckpt(ckpt_path, model, optimizer, scheduler)\n\n # clean up wandb\n if wandb and args.log:\n run.finish()\n\n\ndef subprocess(args):\n with open(args.config, 'r') as f:\n config = yaml.load(f, yaml.FullLoader)\n device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n\n # set random seed\n config['seed'] = args.seed\n seed = args.seed\n torch.manual_seed(seed)\n random.seed(seed)\n if torch.cuda.is_available():\n torch.cuda.manual_seed_all(seed)\n\n # create model \n model = FNO2d(modes1=config['model']['modes1'],\n modes2=config['model']['modes2'],\n fc_dim=config['model']['fc_dim'],\n layers=config['model']['layers'], \n act=config['model']['act'], \n pad_ratio=config['model']['pad_ratio']).to(device)\n num_params = count_params(model)\n config['num_params'] = num_params\n print(f'Number of parameters: {num_params}')\n # Load from checkpoint\n if args.ckpt:\n ckpt_path = args.ckpt\n ckpt = torch.load(ckpt_path)\n model.load_state_dict(ckpt['model'])\n print('Weights loaded from %s' % ckpt_path)\n \n if args.test:\n batchsize = config['test']['batchsize']\n testset = DarcyFlow(datapath=config['test']['path'], \n nx=config['test']['nx'], \n sub=config['test']['sub'], \n offset=config['test']['offset'], \n num=config['test']['n_sample'])\n testloader = DataLoader(testset, batch_size=batchsize, num_workers=4)\n criterion = LpLoss()\n test_err, std_err = eval_darcy(model, testloader, criterion, device)\n print(f'Averaged test relative L2 error: {test_err}; Standard error: {std_err}')\n else:\n # training set\n batchsize = config['train']['batchsize']\n u_set = DarcyFlow(datapath=config['data']['path'], \n nx=config['data']['nx'], \n sub=config['data']['sub'], \n offset=config['data']['offset'], \n num=config['data']['n_sample'])\n u_loader = DataLoader(u_set, batch_size=batchsize, num_workers=4, shuffle=True)\n ic_set = DarcyIC(datapath=config['data']['path'], \n nx=config['data']['nx'], \n sub=config['data']['pde_sub'], \n offset=config['data']['offset'], \n num=config['data']['n_sample'])\n ic_loader = DataLoader(ic_set, batch_size=batchsize, num_workers=4, shuffle=True)\n # val set\n valset = DarcyFlow(datapath=config['test']['path'], \n nx=config['test']['nx'], \n sub=config['test']['sub'], \n offset=config['test']['offset'], \n num=config['test']['n_sample'])\n val_loader = DataLoader(valset, batch_size=batchsize, num_workers=4)\n print(f'Train set: {len(u_set)}; test set: {len(valset)}.')\n optimizer = Adam(model.parameters(), lr=config['train']['base_lr'])\n scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, \n milestones=config['train']['milestones'], \n gamma=config['train']['scheduler_gamma'])\n if args.ckpt:\n ckpt = torch.load(ckpt_path)\n optimizer.load_state_dict(ckpt['optim'])\n scheduler.load_state_dict(ckpt['scheduler'])\n train(model, \n u_loader,\n ic_loader, \n val_loader, \n optimizer, scheduler, \n device, \n config, args)\n \n print('Done!')\n \n \n\nif __name__ == '__main__':\n torch.backends.cudnn.benchmark = True\n # parse options\n parser = ArgumentParser(description='Basic paser')\n parser.add_argument('--config', type=str, help='Path to the configuration file')\n parser.add_argument('--log', action='store_true', help='Turn on the wandb')\n parser.add_argument('--seed', type=int, default=None)\n parser.add_argument('--ckpt', type=str, default=None)\n parser.add_argument('--test', action='store_true', help='Test')\n args = parser.parse_args()\n if args.seed is None:\n args.seed = random.randint(0, 100000)\n subprocess(args)" }, { "path": "train_no.py", "content": "import os\nimport yaml\nimport random\nfrom argparse import ArgumentParser\nimport math\nfrom tqdm import tqdm\n\nimport torch\nimport torch.nn.functional as F\nfrom torch.utils.data import DataLoader, Subset\n\n\nfrom models import FNO3d\nfrom train_utils.adam import Adam\n\nfrom train_utils.losses import LpLoss, PINO_loss3d, get_forcing\nfrom train_utils.datasets import NS3DDataset, KFDataset\nfrom train_utils.utils import save_ckpt, count_params\n\ntry:\n import wandb\nexcept ImportError:\n wandb = None\n\n\ndef pad_input(x, num_pad):\n if num_pad >0:\n res = F.pad(x, (0, 0, 0, num_pad), 'constant', 0)\n else:\n res = x\n return res\n\n\ndef train_ns(model, \n train_loader, \n val_loader,\n optimizer, \n scheduler,\n device, config, args):\n # parse configuration\n v = 1/ config['data']['Re']\n t_duration = config['data']['t_duration']\n num_pad = config['model']['num_pad']\n save_step = config['train']['save_step']\n ic_weight = config['train']['ic_loss']\n f_weight = config['train']['f_loss']\n xy_weight = config['train']['xy_loss']\n # set up directory\n base_dir = os.path.join('exp', config['log']['logdir'])\n ckpt_dir = os.path.join(base_dir, 'ckpts')\n os.makedirs(ckpt_dir, exist_ok=True)\n\n # loss fn\n lploss = LpLoss(size_average=True)\n S = config['data']['pde_res'][0]\n data_s_step = train_loader.dataset.dataset.data_s_step\n data_t_step = train_loader.dataset.dataset.data_t_step\n forcing = get_forcing(S).to(device)\n # set up wandb\n if wandb and args.log:\n run = wandb.init(project=config['log']['project'], \n entity=config['log']['entity'], \n group=config['log']['group'], \n config=config, reinit=True, \n settings=wandb.Settings(start_method='fork'))\n pbar = range(config['train']['epochs'])\n pbar = tqdm(pbar, dynamic_ncols=True, smoothing=0.2)\n zero = torch.zeros(1).to(device)\n for e in pbar:\n loss_dict = {\n 'train_loss': 0.0, \n 'ic_loss': 0.0, \n 'pde_loss': 0.0\n }\n\n # train \n model.train()\n for u, a in train_loader:\n u, a = u.to(device), a.to(device)\n optimizer.zero_grad()\n\n if ic_weight == 0.0 and f_weight == 0.0:\n # FNO\n a_in = a[:, ::data_s_step, ::data_s_step, ::data_t_step]\n out = model(a_in)\n loss_ic, loss_f = zero, zero\n loss = lploss(out, u)\n else:\n # PINO\n a_in = a\n out = model(a_in)\n # PDE loss\n u0 = a[:, :, :, 0, -1]\n loss_ic, loss_f = PINO_loss3d(out, u0, forcing, v, t_duration)\n # data loss\n # print(out.shape)\n # print(u.shape)\n data_loss = lploss(out[:, ::data_s_step, ::data_s_step, ::data_t_step], u)\n loss = data_loss * xy_weight + loss_f * f_weight + loss_ic * ic_weight\n \n loss.backward()\n optimizer.step()\n\n loss_dict['train_loss'] += loss.item()\n loss_dict['ic_loss'] += loss_ic.item()\n loss_dict['pde_loss'] += loss_f.item()\n scheduler.step()\n \n loader_size = len(train_loader)\n train_loss = loss_dict['train_loss'] / loader_size\n ic_loss = loss_dict['ic_loss'] / loader_size\n pde_loss = loss_dict['pde_loss'] / loader_size\n\n # eval\n model.eval()\n with torch.no_grad():\n val_error = 0.0\n for u, a in val_loader:\n u, a = u.to(device), a.to(device)\n\n if ic_weight == 0.0 and f_weight == 0.0:\n # FNO\n a = a[:, ::data_s_step, ::data_s_step, ::data_t_step]\n a_in = a\n out = model(a_in)\n data_loss = lploss(out, u)\n else:\n # PINO\n a_in = a\n out = model(a_in)\n # data loss\n data_loss = lploss(out[:, ::data_s_step, ::data_s_step, ::data_t_step], u)\n val_error += data_loss.item()\n avg_val_error = val_error / len(val_loader)\n\n pbar.set_description(\n (\n f'Train loss: {train_loss}. IC loss: {ic_loss}, PDE loss: {pde_loss}, val error: {avg_val_error}'\n )\n )\n log_dict = {\n 'Train loss': train_loss, \n 'IC loss': ic_loss, \n 'PDE loss': pde_loss, \n 'Val error': avg_val_error\n }\n\n if wandb and args.log:\n wandb.log(log_dict)\n if e % save_step == 0:\n ckpt_path = os.path.join(ckpt_dir, f'model-{e}.pt')\n save_ckpt(ckpt_path, model, optimizer)\n # clean up wandb\n if wandb and args.log:\n run.finish()\n\n\ndef eval_ns(model, val_loader, device, config, args):\n # parse configuration\n v = 1/ config['data']['Re']\n t_duration = config['data']['t_duration']\n num_pad = config['model']['num_pad']\n\n model.eval()\n # loss fn\n lploss = LpLoss(size_average=True)\n S = config['data']['pde_res'][0]\n data_s_step = val_loader.dataset.data_s_step\n data_t_step = val_loader.dataset.data_t_step\n\n with torch.no_grad():\n val_error = 0.0\n for u, a in tqdm(val_loader):\n u, a = u.to(device), a.to(device)\n # a = a[:, ::data_s_step, ::data_s_step, ::data_t_step]\n a_in = a\n out = model(a_in)\n out = out[:, ::data_s_step, ::data_s_step, ::data_t_step]\n data_loss = lploss(out, u)\n val_error += data_loss.item()\n avg_val_err = val_error / len(val_loader)\n\n print(f'Average relative L2 error {avg_val_err}')\n\n\ndef subprocess(args):\n with open(args.config, 'r') as f:\n config = yaml.load(f, yaml.FullLoader)\n device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n\n # set random seed\n config['seed'] = args.seed\n seed = args.seed\n torch.manual_seed(seed)\n random.seed(seed)\n if torch.cuda.is_available():\n torch.cuda.manual_seed_all(seed)\n\n # create model \n model = FNO3d(modes1=config['model']['modes1'],\n modes2=config['model']['modes2'],\n modes3=config['model']['modes3'],\n fc_dim=config['model']['fc_dim'],\n layers=config['model']['layers'], \n act=config['model']['act'], \n pad_ratio=config['model']['pad_ratio']).to(device)\n num_params = count_params(model)\n config['num_params'] = num_params\n print(f'Number of parameters: {num_params}')\n # Load from checkpoint\n if args.ckpt:\n ckpt_path = args.ckpt\n ckpt = torch.load(ckpt_path)\n model.load_state_dict(ckpt['model'])\n print('Weights loaded from %s' % ckpt_path)\n\n datasets = {\n 'KF': KFDataset, \n 'NS': NS3DDataset\n }\n if 'name' in config['data']:\n dataname = config['data']['name']\n else:\n dataname = 'NS'\n\n if args.test:\n batchsize = config['test']['batchsize']\n testset = datasets[dataname](paths=config['data']['paths'], \n raw_res=config['data']['raw_res'],\n data_res=config['test']['data_res'], \n pde_res=config['data']['pde_res'], \n n_samples=config['data']['n_test_samples'], \n offset=config['data']['testoffset'], \n t_duration=config['data']['t_duration'])\n test_loader = DataLoader(testset, batch_size=batchsize, num_workers=4, shuffle=True)\n eval_ns(model, test_loader, device, config, args)\n\n else:\n # prepare datast\n batchsize = config['train']['batchsize']\n\n dataset = datasets[dataname](paths=config['data']['paths'], \n raw_res=config['data']['raw_res'],\n data_res=config['data']['data_res'], \n pde_res=config['data']['pde_res'], \n n_samples=config['data']['n_samples'], \n offset=config['data']['offset'], \n t_duration=config['data']['t_duration'])\n idxs = torch.randperm(len(dataset))\n # setup train and test\n num_test = config['data']['n_test_samples']\n num_train = len(idxs) - num_test\n print(f'Number of training samples: {num_train};\\nNumber of test samples: {num_test}.')\n train_idx = idxs[:num_train]\n test_idx = idxs[num_train:]\n\n trainset = Subset(dataset, indices=train_idx)\n valset = Subset(dataset, indices=test_idx)\n\n train_loader = DataLoader(trainset, batch_size=batchsize, num_workers=4, shuffle=True)\n\n val_loader = DataLoader(valset, batch_size=batchsize, num_workers=4)\n optimizer = Adam(model.parameters(), lr=config['train']['base_lr'])\n scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, \n milestones=config['train']['milestones'], \n gamma=config['train']['scheduler_gamma'])\n print(dataset.data.shape)\n train_ns(model, train_loader, val_loader, \n optimizer, scheduler, device, config, args)\n print('Done!')\n\n\nif __name__ == '__main__':\n\n torch.backends.cudnn.benchmark = True\n # parse options\n parser = ArgumentParser(description='Basic paser')\n parser.add_argument('--config', type=str, help='Path to the configuration file')\n parser.add_argument('--log', action='store_true', help='Turn on the wandb')\n parser.add_argument('--seed', type=int, default=None)\n parser.add_argument('--ckpt', type=str, default=None)\n parser.add_argument('--test', action='store_true', help='Test')\n args = parser.parse_args()\n if args.seed is None:\n args.seed = random.randint(0, 100000)\n subprocess(args)" }, { "path": "train_operator.py", "content": "import yaml\nfrom argparse import ArgumentParser\nimport math\nimport torch\nfrom torch.utils.data import DataLoader\n\nfrom solver.random_fields import GaussianRF\nfrom train_utils import Adam\nfrom train_utils.datasets import NSLoader, online_loader, DarcyFlow, DarcyCombo\nfrom train_utils.train_3d import mixed_train\nfrom train_utils.train_2d import train_2d_operator\nfrom models import FNO3d, FNO2d\n\n\ndef train_3d(args, config):\n device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n data_config = config['data']\n\n # prepare dataloader for training with data\n if 'datapath2' in data_config:\n loader = NSLoader(datapath1=data_config['datapath'], datapath2=data_config['datapath2'],\n nx=data_config['nx'], nt=data_config['nt'],\n sub=data_config['sub'], sub_t=data_config['sub_t'],\n N=data_config['total_num'],\n t_interval=data_config['time_interval'])\n else:\n loader = NSLoader(datapath1=data_config['datapath'],\n nx=data_config['nx'], nt=data_config['nt'],\n sub=data_config['sub'], sub_t=data_config['sub_t'],\n N=data_config['total_num'],\n t_interval=data_config['time_interval'])\n\n train_loader = loader.make_loader(data_config['n_sample'],\n batch_size=config['train']['batchsize'],\n start=data_config['offset'],\n train=data_config['shuffle'])\n # prepare dataloader for training with only equations\n gr_sampler = GaussianRF(2, data_config['S2'], 2 * math.pi, alpha=2.5, tau=7, device=device)\n a_loader = online_loader(gr_sampler,\n S=data_config['S2'],\n T=data_config['T2'],\n time_scale=data_config['time_interval'],\n batchsize=config['train']['batchsize'])\n # create model\n print(device)\n model = FNO3d(modes1=config['model']['modes1'],\n modes2=config['model']['modes2'],\n modes3=config['model']['modes3'],\n fc_dim=config['model']['fc_dim'],\n layers=config['model']['layers'], \n act=config['model']['act']).to(device)\n # Load from checkpoint\n if 'ckpt' in config['train']:\n ckpt_path = config['train']['ckpt']\n ckpt = torch.load(ckpt_path)\n model.load_state_dict(ckpt['model'])\n print('Weights loaded from %s' % ckpt_path)\n # create optimizer and learning rate scheduler\n optimizer = Adam(model.parameters(), betas=(0.9, 0.999),\n lr=config['train']['base_lr'])\n scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer,\n milestones=config['train']['milestones'],\n gamma=config['train']['scheduler_gamma'])\n mixed_train(model,\n train_loader,\n loader.S, loader.T,\n a_loader,\n data_config['S2'], data_config['T2'],\n optimizer,\n scheduler,\n config,\n device,\n log=args.log,\n project=config['log']['project'],\n group=config['log']['group'])\n\n\ndef train_2d(args, config):\n device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n data_config = config['data']\n\n # dataset = DarcyFlow(data_config['datapath'],\n # nx=data_config['nx'], sub=data_config['sub'],\n # offset=data_config['offset'], num=data_config['n_sample'])\n\n dataset = DarcyCombo(datapath=data_config['datapath'], \n nx=data_config['nx'], \n sub=data_config['sub'], \n pde_sub=data_config['pde_sub'], \n num=data_config['n_samples'], \n offset=data_config['offset'])\n train_loader = DataLoader(dataset, batch_size=config['train']['batchsize'], shuffle=True)\n model = FNO2d(modes1=config['model']['modes1'],\n modes2=config['model']['modes2'],\n fc_dim=config['model']['fc_dim'],\n layers=config['model']['layers'],\n act=config['model']['act'], \n pad_ratio=config['model']['pad_ratio']).to(device)\n # Load from checkpoint\n if 'ckpt' in config['train']:\n ckpt_path = config['train']['ckpt']\n ckpt = torch.load(ckpt_path)\n model.load_state_dict(ckpt['model'])\n print('Weights loaded from %s' % ckpt_path)\n\n optimizer = Adam(model.parameters(), betas=(0.9, 0.999),\n lr=config['train']['base_lr'])\n scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer,\n milestones=config['train']['milestones'],\n gamma=config['train']['scheduler_gamma'])\n train_2d_operator(model,\n train_loader,\n optimizer, scheduler,\n config, rank=0, log=args.log,\n project=config['log']['project'],\n group=config['log']['group'])\n\n\nif __name__ == '__main__':\n device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n\n # parse options\n parser = ArgumentParser(description='Basic paser')\n parser.add_argument('--config_path', type=str, help='Path to the configuration file')\n parser.add_argument('--log', action='store_true', help='Turn on the wandb')\n args = parser.parse_args()\n\n config_file = args.config_path\n with open(config_file, 'r') as stream:\n config = yaml.load(stream, yaml.FullLoader)\n\n if 'name' in config['data'] and config['data']['name'] == 'Darcy':\n train_2d(args, config)\n else:\n train_3d(args, config)\n" }, { "path": "train_pino.py", "content": "from datetime import datetime\nimport os\nimport yaml\nimport random\nfrom argparse import ArgumentParser\nimport math\nfrom tqdm import tqdm\n\nimport numpy as np\n\nimport torch\nfrom torch.optim import Adam\nfrom torch.utils.data import DataLoader\n\nfrom models import FNO3d\n\nfrom train_utils.losses import LpLoss, PINO_loss3d, get_forcing\nfrom train_utils.datasets import KFDataset, KFaDataset, sample_data\nfrom train_utils.utils import save_ckpt, count_params, dict2str\n\ntry:\n import wandb\nexcept ImportError:\n wandb = None\n\n\n@torch.no_grad()\ndef eval_ns(model, val_loader, criterion, device):\n model.eval()\n val_err = []\n for u, a in val_loader:\n u, a = u.to(device), a.to(device)\n out = model(a)\n val_loss = criterion(out, u)\n val_err.append(val_loss.item())\n\n N = len(val_loader)\n\n avg_err = np.mean(val_err)\n std_err = np.std(val_err, ddof=1) / np.sqrt(N)\n return avg_err, std_err\n\n\ndef train_ns(model, \n train_u_loader, # training data\n train_a_loader, # initial conditions\n val_loader, # validation data\n optimizer, \n scheduler,\n device, config, args):\n start_iter = config['train']['start_iter']\n v = 1/ config['data']['Re']\n t_duration = config['data']['t_duration']\n save_step = config['train']['save_step']\n eval_step = config['train']['eval_step']\n\n ic_weight = config['train']['ic_loss']\n f_weight = config['train']['f_loss']\n xy_weight = config['train']['xy_loss']\n # set up directory\n base_dir = os.path.join('exp', config['log']['logdir'])\n ckpt_dir = os.path.join(base_dir, 'ckpts')\n os.makedirs(ckpt_dir, exist_ok=True)\n\n # loss fn\n lploss = LpLoss(size_average=True)\n \n S = config['data']['pde_res'][0]\n forcing = get_forcing(S).to(device)\n # set up wandb\n if wandb and args.log:\n run = wandb.init(project=config['log']['project'], \n entity=config['log']['entity'], \n group=config['log']['group'], \n config=config, reinit=True, \n settings=wandb.Settings(start_method='fork'))\n \n pbar = range(start_iter, config['train']['num_iter'])\n if args.tqdm:\n pbar = tqdm(pbar, dynamic_ncols=True, smoothing=0.2)\n\n u_loader = sample_data(train_u_loader)\n a_loader = sample_data(train_a_loader)\n\n for e in pbar:\n log_dict = {}\n\n optimizer.zero_grad()\n # data loss\n if xy_weight > 0:\n u, a_in = next(u_loader)\n u = u.to(device)\n a_in = a_in.to(device)\n out = model(a_in)\n data_loss = lploss(out, u)\n else:\n data_loss = torch.zeros(1, device=device)\n\n if f_weight != 0.0:\n # pde loss\n a = next(a_loader)\n a = a.to(device)\n out = model(a)\n \n u0 = a[:, :, :, 0, -1]\n loss_ic, loss_f = PINO_loss3d(out, u0, forcing, v, t_duration)\n log_dict['IC'] = loss_ic.item()\n log_dict['PDE'] = loss_f.item()\n else:\n loss_ic = loss_f = 0.0\n loss = data_loss * xy_weight + loss_f * f_weight + loss_ic * ic_weight\n\n loss.backward()\n optimizer.step()\n scheduler.step()\n\n log_dict['train loss'] = loss.item()\n log_dict['data'] = data_loss.item()\n if e % eval_step == 0:\n eval_err, std_err = eval_ns(model, val_loader, lploss, device)\n log_dict['val error'] = eval_err\n \n if args.tqdm:\n logstr = dict2str(log_dict)\n pbar.set_description(\n (\n logstr\n )\n )\n if wandb and args.log:\n wandb.log(log_dict)\n if e % save_step == 0 and e > 0:\n ckpt_path = os.path.join(ckpt_dir, f'model-{e}.pt')\n save_ckpt(ckpt_path, model, optimizer, scheduler)\n\n # clean up wandb\n if wandb and args.log:\n run.finish()\n\n\ndef subprocess(args):\n with open(args.config, 'r') as f:\n config = yaml.load(f, yaml.FullLoader)\n device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n\n # set random seed\n config['seed'] = args.seed\n seed = args.seed\n torch.manual_seed(seed)\n random.seed(seed)\n if torch.cuda.is_available():\n torch.cuda.manual_seed_all(seed)\n\n # create model \n model = FNO3d(modes1=config['model']['modes1'],\n modes2=config['model']['modes2'],\n modes3=config['model']['modes3'],\n fc_dim=config['model']['fc_dim'],\n layers=config['model']['layers'], \n act=config['model']['act'], \n pad_ratio=config['model']['pad_ratio']).to(device)\n num_params = count_params(model)\n config['num_params'] = num_params\n print(f'Number of parameters: {num_params}')\n # Load from checkpoint\n if args.ckpt:\n ckpt_path = args.ckpt\n ckpt = torch.load(ckpt_path)\n model.load_state_dict(ckpt['model'])\n print('Weights loaded from %s' % ckpt_path)\n \n if args.test:\n batchsize = config['test']['batchsize']\n testset = KFDataset(paths=config['data']['paths'], \n raw_res=config['data']['raw_res'],\n data_res=config['test']['data_res'], \n pde_res=config['test']['data_res'], \n n_samples=config['data']['n_test_samples'], \n offset=config['data']['testoffset'], \n t_duration=config['data']['t_duration'])\n testloader = DataLoader(testset, batch_size=batchsize, num_workers=4)\n criterion = LpLoss()\n test_err, std_err = eval_ns(model, testloader, criterion, device)\n print(f'Averaged test relative L2 error: {test_err}; Standard error: {std_err}')\n else:\n # training set\n batchsize = config['train']['batchsize']\n u_set = KFDataset(paths=config['data']['paths'], \n raw_res=config['data']['raw_res'],\n data_res=config['data']['data_res'], \n pde_res=config['data']['data_res'], \n n_samples=config['data']['n_data_samples'], \n offset=config['data']['offset'], \n t_duration=config['data']['t_duration'])\n u_loader = DataLoader(u_set, batch_size=batchsize, num_workers=4, shuffle=True)\n\n a_set = KFaDataset(paths=config['data']['paths'], \n raw_res=config['data']['raw_res'], \n pde_res=config['data']['pde_res'], \n n_samples=config['data']['n_a_samples'],\n offset=config['data']['a_offset'], \n t_duration=config['data']['t_duration'])\n a_loader = DataLoader(a_set, batch_size=batchsize, num_workers=4, shuffle=True)\n # val set\n valset = KFDataset(paths=config['data']['paths'], \n raw_res=config['data']['raw_res'],\n data_res=config['test']['data_res'], \n pde_res=config['test']['data_res'], \n n_samples=config['data']['n_test_samples'], \n offset=config['data']['testoffset'], \n t_duration=config['data']['t_duration'])\n val_loader = DataLoader(valset, batch_size=batchsize, num_workers=4)\n print(f'Train set: {len(u_set)}; Test set: {len(valset)}; IC set: {len(a_set)}')\n optimizer = Adam(model.parameters(), lr=config['train']['base_lr'])\n scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, \n milestones=config['train']['milestones'], \n gamma=config['train']['scheduler_gamma'])\n if args.ckpt:\n ckpt = torch.load(ckpt_path)\n optimizer.load_state_dict(ckpt['optim'])\n scheduler.load_state_dict(ckpt['scheduler'])\n config['train']['start_iter'] = scheduler.last_epoch\n train_ns(model, \n u_loader, a_loader, \n val_loader, \n optimizer, scheduler, \n device, \n config, args)\n print('Done!')\n \n \n\nif __name__ == '__main__':\n torch.backends.cudnn.benchmark = True\n # parse options\n parser = ArgumentParser(description='Basic paser')\n parser.add_argument('--config', type=str, help='Path to the configuration file')\n parser.add_argument('--log', action='store_true', help='Turn on the wandb')\n parser.add_argument('--seed', type=int, default=None)\n parser.add_argument('--ckpt', type=str, default=None)\n parser.add_argument('--test', action='store_true', help='Test')\n parser.add_argument('--tqdm', action='store_true', help='Turn on the tqdm')\n args = parser.parse_args()\n if args.seed is None:\n args.seed = random.randint(0, 100000)\n subprocess(args)" }, { "path": "train_unet.py", "content": "from datetime import datetime\nimport os\nimport yaml\nimport random\nfrom argparse import ArgumentParser\nimport math\nfrom tqdm import tqdm\n\nimport numpy as np\n\nimport torch\nfrom torch.optim import Adam\nfrom torch.utils.data import DataLoader\n\nfrom baselines.unet3d import UNet3D\n\nfrom train_utils.losses import LpLoss\nfrom train_utils.datasets import KFDataset, KFaDataset, sample_data\nfrom train_utils.utils import save_ckpt, count_params, dict2str\n\ntry:\n import wandb\nexcept ImportError:\n wandb = None\n\n\n@torch.no_grad()\ndef eval_ns(model, val_loader, criterion, device):\n model.eval()\n val_err = []\n for u, a in val_loader:\n u, a = u.to(device), a.to(device)\n a = a.permute(0, 4, 3, 1, 2)\n out = model(a)\n out = out.squeeze(1).permute(0, 2, 3, 1)\n val_loss = criterion(out, u)\n val_err.append(val_loss.item())\n\n N = len(val_loader)\n\n avg_err = np.mean(val_err)\n std_err = np.std(val_err, ddof=1) / np.sqrt(N)\n return avg_err, std_err\n\n\ndef train_ns(model, \n train_u_loader, # training data\n val_loader, # validation data\n optimizer, \n scheduler,\n device, config, args):\n start_iter = config['train']['start_iter']\n v = 1/ config['data']['Re']\n save_step = config['train']['save_step']\n eval_step = config['train']['eval_step']\n\n # set up directory\n base_dir = os.path.join('exp', config['log']['logdir'])\n ckpt_dir = os.path.join(base_dir, 'ckpts')\n os.makedirs(ckpt_dir, exist_ok=True)\n\n # loss fn\n lploss = LpLoss(size_average=True)\n \n S = config['data']['pde_res'][0]\n # set up wandb\n if wandb and args.log:\n run = wandb.init(project=config['log']['project'], \n entity=config['log']['entity'], \n group=config['log']['group'], \n config=config, reinit=True, \n settings=wandb.Settings(start_method='fork'))\n \n pbar = range(start_iter, config['train']['num_iter'])\n if args.tqdm:\n pbar = tqdm(pbar, dynamic_ncols=True, smoothing=0.2)\n\n u_loader = sample_data(train_u_loader)\n\n for e in pbar:\n log_dict = {}\n\n optimizer.zero_grad()\n # data loss\n u, a_in = next(u_loader)\n u = u.to(device)\n a_in = a_in.to(device).permute(0, 4, 3, 1, 2) # B, C, T, X, Y\n out = model(a_in)\n out = out.squeeze(1).permute(0, 2, 3, 1) # B, X, Y, T\n data_loss = lploss(out, u)\n\n loss = data_loss\n\n loss.backward()\n optimizer.step()\n scheduler.step()\n\n log_dict['train loss'] = loss.item()\n if e % eval_step == 0:\n eval_err, std_err = eval_ns(model, val_loader, lploss, device)\n log_dict['val error'] = eval_err\n \n if args.tqdm:\n logstr = dict2str(log_dict)\n pbar.set_description(\n (\n logstr\n )\n )\n if wandb and args.log:\n wandb.log(log_dict)\n if e % save_step == 0 and e > 0:\n ckpt_path = os.path.join(ckpt_dir, f'model-{e}.pt')\n save_ckpt(ckpt_path, model, optimizer, scheduler)\n\n # clean up wandb\n if wandb and args.log:\n run.finish()\n\n\ndef subprocess(args):\n with open(args.config, 'r') as f:\n config = yaml.load(f, yaml.FullLoader)\n device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n\n # set random seed\n config['seed'] = args.seed\n seed = args.seed\n torch.manual_seed(seed)\n random.seed(seed)\n if torch.cuda.is_available():\n torch.cuda.manual_seed_all(seed)\n\n # create model \n model = UNet3D(in_channels=4, out_channels=1, f_maps=64, final_sigmoid=False).to(device)\n num_params = count_params(model)\n config['num_params'] = num_params\n print(f'Number of parameters: {num_params}')\n # Load from checkpoint\n if args.ckpt:\n ckpt_path = args.ckpt\n ckpt = torch.load(ckpt_path)\n model.load_state_dict(ckpt['model'])\n print('Weights loaded from %s' % ckpt_path)\n \n if args.test:\n batchsize = config['test']['batchsize']\n testset = KFDataset(paths=config['data']['paths'], \n raw_res=config['data']['raw_res'],\n data_res=config['test']['data_res'], \n pde_res=config['test']['data_res'], \n n_samples=config['data']['n_test_samples'], \n offset=config['data']['testoffset'], \n t_duration=config['data']['t_duration'])\n testloader = DataLoader(testset, batch_size=batchsize, num_workers=4)\n criterion = LpLoss()\n test_err, std_err = eval_ns(model, testloader, criterion, device)\n print(f'Averaged test relative L2 error: {test_err}; Standard error: {std_err}')\n else:\n # training set\n batchsize = config['train']['batchsize']\n u_set = KFDataset(paths=config['data']['paths'], \n raw_res=config['data']['raw_res'],\n data_res=config['data']['data_res'], \n pde_res=config['data']['data_res'], \n n_samples=config['data']['n_data_samples'], \n offset=config['data']['offset'], \n t_duration=config['data']['t_duration'])\n u_loader = DataLoader(u_set, batch_size=batchsize, num_workers=4, shuffle=True)\n\n # val set\n valset = KFDataset(paths=config['data']['paths'], \n raw_res=config['data']['raw_res'],\n data_res=config['test']['data_res'], \n pde_res=config['test']['data_res'], \n n_samples=config['data']['n_test_samples'], \n offset=config['data']['testoffset'], \n t_duration=config['data']['t_duration'])\n val_loader = DataLoader(valset, batch_size=batchsize, num_workers=4)\n print(f'Train set: {len(u_set)}; Test set: {len(valset)}.')\n optimizer = Adam(model.parameters(), lr=config['train']['base_lr'])\n scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, \n milestones=config['train']['milestones'], \n gamma=config['train']['scheduler_gamma'])\n if args.ckpt:\n ckpt = torch.load(ckpt_path)\n optimizer.load_state_dict(ckpt['optim'])\n scheduler.load_state_dict(ckpt['scheduler'])\n config['train']['start_iter'] = scheduler.last_epoch\n train_ns(model, \n u_loader, \n val_loader, \n optimizer, scheduler, \n device, \n config, args)\n print('Done!')\n \n \n\nif __name__ == '__main__':\n torch.backends.cudnn.benchmark = True\n # parse options\n parser = ArgumentParser(description='Basic paser')\n parser.add_argument('--config', type=str, help='Path to the configuration file')\n parser.add_argument('--log', action='store_true', help='Turn on the wandb')\n parser.add_argument('--seed', type=int, default=None)\n parser.add_argument('--ckpt', type=str, default=None)\n parser.add_argument('--test', action='store_true', help='Test')\n parser.add_argument('--tqdm', action='store_true', help='Turn on the tqdm')\n args = parser.parse_args()\n if args.seed is None:\n args.seed = random.randint(0, 100000)\n subprocess(args)\n" }, { "path": "train_utils/__init__.py", "content": "from .adam import Adam\nfrom .datasets import NSLoader, DarcyFlow\nfrom .losses import get_forcing, LpLoss" }, { "path": "train_utils/adam.py", "content": "import math\nimport torch\nfrom torch import Tensor\nfrom typing import List, Optional\nfrom torch.optim.optimizer import Optimizer\n\n\ndef adam(params: List[Tensor],\n grads: List[Tensor],\n exp_avgs: List[Tensor],\n exp_avg_sqs: List[Tensor],\n max_exp_avg_sqs: List[Tensor],\n state_steps: List[int],\n *,\n amsgrad: bool,\n beta1: float,\n beta2: float,\n lr: float,\n weight_decay: float,\n eps: float):\n r\"\"\"Functional API that performs Adam algorithm computation.\n See :class:`~torch.optim.Adam` for details.\n \"\"\"\n\n for i, param in enumerate(params):\n\n grad = grads[i]\n exp_avg = exp_avgs[i]\n exp_avg_sq = exp_avg_sqs[i]\n step = state_steps[i]\n\n bias_correction1 = 1 - beta1 ** step\n bias_correction2 = 1 - beta2 ** step\n\n if weight_decay != 0:\n grad = grad.add(param, alpha=weight_decay)\n\n # Decay the first and second moment running average coefficient\n exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)\n exp_avg_sq.mul_(beta2).addcmul_(grad, grad.conj(), value=1 - beta2)\n if amsgrad:\n # Maintains the maximum of all 2nd moment running avg. till now\n torch.maximum(max_exp_avg_sqs[i], exp_avg_sq, out=max_exp_avg_sqs[i])\n # Use the max. for normalizing running avg. of gradient\n denom = (max_exp_avg_sqs[i].sqrt() / math.sqrt(bias_correction2)).add_(eps)\n else:\n denom = (exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(eps)\n\n step_size = lr / bias_correction1\n\n param.addcdiv_(exp_avg, denom, value=-step_size)\n\n\nclass Adam(Optimizer):\n r\"\"\"Implements Adam algorithm.\n It has been proposed in `Adam: A Method for Stochastic Optimization`_.\n The implementation of the L2 penalty follows changes proposed in\n `Decoupled Weight Decay Regularization`_.\n Args:\n params (iterable): iterable of parameters to optimize or dicts defining\n parameter groups\n lr (float, optional): learning rate (default: 1e-3)\n betas (Tuple[float, float], optional): coefficients used for computing\n running averages of gradient and its square (default: (0.9, 0.999))\n eps (float, optional): term added to the denominator to improve\n numerical stability (default: 1e-8)\n weight_decay (float, optional): weight decay (L2 penalty) (default: 0)\n amsgrad (boolean, optional): whether to use the AMSGrad variant of this\n algorithm from the paper `On the Convergence of Adam and Beyond`_\n (default: False)\n .. _Adam\\: A Method for Stochastic Optimization:\n https://arxiv.org/abs/1412.6980\n .. _Decoupled Weight Decay Regularization:\n https://arxiv.org/abs/1711.05101\n .. _On the Convergence of Adam and Beyond:\n https://openreview.net/forum?id=ryQu7f-RZ\n \"\"\"\n\n def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8,\n weight_decay=0, amsgrad=False):\n if not 0.0 <= lr:\n raise ValueError(\"Invalid learning rate: {}\".format(lr))\n if not 0.0 <= eps:\n raise ValueError(\"Invalid epsilon value: {}\".format(eps))\n if not 0.0 <= betas[0] < 1.0:\n raise ValueError(\"Invalid beta parameter at index 0: {}\".format(betas[0]))\n if not 0.0 <= betas[1] < 1.0:\n raise ValueError(\"Invalid beta parameter at index 1: {}\".format(betas[1]))\n if not 0.0 <= weight_decay:\n raise ValueError(\"Invalid weight_decay value: {}\".format(weight_decay))\n defaults = dict(lr=lr, betas=betas, eps=eps,\n weight_decay=weight_decay, amsgrad=amsgrad)\n super(Adam, self).__init__(params, defaults)\n\n def __setstate__(self, state):\n super(Adam, self).__setstate__(state)\n for group in self.param_groups:\n group.setdefault('amsgrad', False)\n\n @torch.no_grad()\n def step(self, closure=None):\n \"\"\"Performs a single optimization step.\n Args:\n closure (callable, optional): A closure that reevaluates the model\n and returns the loss.\n \"\"\"\n loss = None\n if closure is not None:\n with torch.enable_grad():\n loss = closure()\n\n for group in self.param_groups:\n params_with_grad = []\n grads = []\n exp_avgs = []\n exp_avg_sqs = []\n max_exp_avg_sqs = []\n state_steps = []\n beta1, beta2 = group['betas']\n\n for p in group['params']:\n if p.grad is not None:\n params_with_grad.append(p)\n if p.grad.is_sparse:\n raise RuntimeError('Adam does not support sparse gradients, please consider SparseAdam instead')\n grads.append(p.grad)\n\n state = self.state[p]\n # Lazy state initialization\n if len(state) == 0:\n state['step'] = 0\n # Exponential moving average of gradient values\n state['exp_avg'] = torch.zeros_like(p, memory_format=torch.preserve_format)\n # Exponential moving average of squared gradient values\n state['exp_avg_sq'] = torch.zeros_like(p, memory_format=torch.preserve_format)\n if group['amsgrad']:\n # Maintains max of all exp. moving avg. of sq. grad. values\n state['max_exp_avg_sq'] = torch.zeros_like(p, memory_format=torch.preserve_format)\n\n exp_avgs.append(state['exp_avg'])\n exp_avg_sqs.append(state['exp_avg_sq'])\n\n if group['amsgrad']:\n max_exp_avg_sqs.append(state['max_exp_avg_sq'])\n\n # update the steps for each param group update\n state['step'] += 1\n # record the step after step update\n state_steps.append(state['step'])\n\n adam(params_with_grad,\n grads,\n exp_avgs,\n exp_avg_sqs,\n max_exp_avg_sqs,\n state_steps,\n amsgrad=group['amsgrad'],\n beta1=beta1,\n beta2=beta2,\n lr=group['lr'],\n weight_decay=group['weight_decay'],\n eps=group['eps'])\n return loss\n" }, { "path": "train_utils/data_utils.py", "content": "from torch.utils import data\n\n\ndef sample_data(loader):\n while True:\n for batch in loader:\n yield batch\n\n\ndef data_sampler(dataset, shuffle, distributed):\n if distributed:\n return data.distributed.DistributedSampler(dataset, shuffle=shuffle)\n\n if shuffle:\n return data.RandomSampler(dataset)\n\n else:\n return data.SequentialSampler(dataset)" }, { "path": "train_utils/datasets.py", "content": "import scipy.io\nimport numpy as np\n\ntry:\n from pyDOE import lhs\n # Only needed for PINN's dataset\nexcept ImportError:\n lhs = None\n\nimport torch\nfrom torch.utils.data import Dataset\nfrom .utils import get_grid3d, convert_ic, torch2dgrid\n\n\ndef online_loader(sampler, S, T, time_scale, batchsize=1):\n while True:\n u0 = sampler.sample(batchsize)\n a = convert_ic(u0, batchsize,\n S, T,\n time_scale=time_scale)\n yield a\n\n\ndef sample_data(loader):\n while True:\n for batch in loader:\n yield batch\n\n\nclass MatReader(object):\n def __init__(self, file_path, to_torch=True, to_cuda=False, to_float=True):\n super(MatReader, self).__init__()\n\n self.to_torch = to_torch\n self.to_cuda = to_cuda\n self.to_float = to_float\n\n self.file_path = file_path\n\n self.data = None\n self.old_mat = None\n self._load_file()\n\n def _load_file(self):\n self.data = scipy.io.loadmat(self.file_path)\n self.old_mat = True\n\n def load_file(self, file_path):\n self.file_path = file_path\n self._load_file()\n\n def read_field(self, field):\n x = self.data[field]\n\n if not self.old_mat:\n x = x[()]\n x = np.transpose(x, axes=range(len(x.shape) - 1, -1, -1))\n\n if self.to_float:\n x = x.astype(np.float32)\n\n if self.to_torch:\n x = torch.from_numpy(x)\n\n if self.to_cuda:\n x = x.cuda()\n\n return x\n\n def set_cuda(self, to_cuda):\n self.to_cuda = to_cuda\n\n def set_torch(self, to_torch):\n self.to_torch = to_torch\n\n def set_float(self, to_float):\n self.to_float = to_float\n\n\nclass BurgersLoader(object):\n def __init__(self, datapath, nx=2 ** 10, nt=100, sub=8, sub_t=1, new=False):\n dataloader = MatReader(datapath)\n self.sub = sub\n self.sub_t = sub_t\n self.s = nx // sub\n self.T = nt // sub_t\n self.new = new\n if new:\n self.T += 1\n self.x_data = dataloader.read_field('input')[:, ::sub]\n self.y_data = dataloader.read_field('output')[:, ::sub_t, ::sub]\n self.v = dataloader.read_field('visc').item()\n\n def make_loader(self, n_sample, batch_size, start=0, train=True):\n Xs = self.x_data[start:start + n_sample]\n ys = self.y_data[start:start + n_sample]\n\n if self.new:\n gridx = torch.tensor(np.linspace(0, 1, self.s + 1)[:-1], dtype=torch.float)\n gridt = torch.tensor(np.linspace(0, 1, self.T), dtype=torch.float)\n else:\n gridx = torch.tensor(np.linspace(0, 1, self.s), dtype=torch.float)\n gridt = torch.tensor(np.linspace(0, 1, self.T + 1)[1:], dtype=torch.float)\n gridx = gridx.reshape(1, 1, self.s)\n gridt = gridt.reshape(1, self.T, 1)\n\n Xs = Xs.reshape(n_sample, 1, self.s).repeat([1, self.T, 1])\n Xs = torch.stack([Xs, gridx.repeat([n_sample, self.T, 1]), gridt.repeat([n_sample, 1, self.s])], dim=3)\n dataset = torch.utils.data.TensorDataset(Xs, ys)\n if train:\n loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=True)\n else:\n loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=False)\n return loader\n\n\nclass NSLoader(object):\n def __init__(self, datapath1,\n nx, nt,\n datapath2=None, sub=1, sub_t=1,\n N=100, t_interval=1.0):\n '''\n Load data from npy and reshape to (N, X, Y, T)\n Args:\n datapath1: path to data\n nx:\n nt:\n datapath2: path to second part of data, default None\n sub:\n sub_t:\n N:\n t_interval:\n '''\n self.S = nx // sub\n self.T = int(nt * t_interval) // sub_t + 1\n self.time_scale = t_interval\n data1 = np.load(datapath1)\n data1 = torch.tensor(data1, dtype=torch.float)[..., ::sub_t, ::sub, ::sub]\n\n if datapath2 is not None:\n data2 = np.load(datapath2)\n data2 = torch.tensor(data2, dtype=torch.float)[..., ::sub_t, ::sub, ::sub]\n if t_interval == 0.5:\n data1 = self.extract(data1)\n if datapath2 is not None:\n data2 = self.extract(data2)\n part1 = data1.permute(0, 2, 3, 1)\n if datapath2 is not None:\n part2 = data2.permute(0, 2, 3, 1)\n self.data = torch.cat((part1, part2), dim=0)\n else:\n self.data = part1\n\n def make_loader(self, n_sample, batch_size, start=0, train=True):\n if train:\n a_data = self.data[start:start + n_sample, :, :, 0].reshape(n_sample, self.S, self.S)\n u_data = self.data[start:start + n_sample].reshape(n_sample, self.S, self.S, self.T)\n else:\n a_data = self.data[-n_sample:, :, :, 0].reshape(n_sample, self.S, self.S)\n u_data = self.data[-n_sample:].reshape(n_sample, self.S, self.S, self.T)\n a_data = a_data.reshape(n_sample, self.S, self.S, 1, 1).repeat([1, 1, 1, self.T, 1])\n gridx, gridy, gridt = get_grid3d(self.S, self.T, time_scale=self.time_scale)\n a_data = torch.cat((gridx.repeat([n_sample, 1, 1, 1, 1]), gridy.repeat([n_sample, 1, 1, 1, 1]),\n gridt.repeat([n_sample, 1, 1, 1, 1]), a_data), dim=-1)\n dataset = torch.utils.data.TensorDataset(a_data, u_data)\n loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=train)\n return loader\n\n def make_dataset(self, n_sample, start=0, train=True):\n if train:\n a_data = self.data[start:start + n_sample, :, :, 0].reshape(n_sample, self.S, self.S)\n u_data = self.data[start:start + n_sample].reshape(n_sample, self.S, self.S, self.T)\n else:\n a_data = self.data[-n_sample:, :, :, 0].reshape(n_sample, self.S, self.S)\n u_data = self.data[-n_sample:].reshape(n_sample, self.S, self.S, self.T)\n a_data = a_data.reshape(n_sample, self.S, self.S, 1, 1).repeat([1, 1, 1, self.T, 1])\n gridx, gridy, gridt = get_grid3d(self.S, self.T)\n a_data = torch.cat((\n gridx.repeat([n_sample, 1, 1, 1, 1]),\n gridy.repeat([n_sample, 1, 1, 1, 1]),\n gridt.repeat([n_sample, 1, 1, 1, 1]),\n a_data), dim=-1)\n dataset = torch.utils.data.TensorDataset(a_data, u_data)\n return dataset\n\n @staticmethod\n def extract(data):\n '''\n Extract data with time range 0-0.5, 0.25-0.75, 0.5-1.0, 0.75-1.25,...\n Args:\n data: tensor with size N x 129 x 128 x 128\n\n Returns:\n output: (4*N-1) x 65 x 128 x 128\n '''\n T = data.shape[1] // 2\n interval = data.shape[1] // 4\n N = data.shape[0]\n new_data = torch.zeros(4 * N - 1, T + 1, data.shape[2], data.shape[3])\n for i in range(N):\n for j in range(4):\n if i == N - 1 and j == 3:\n # reach boundary\n break\n if j != 3:\n new_data[i * 4 + j] = data[i, interval * j:interval * j + T + 1]\n else:\n new_data[i * 4 + j, 0: interval] = data[i, interval * j:interval * j + interval]\n new_data[i * 4 + j, interval: T + 1] = data[i + 1, 0:interval + 1]\n return new_data\n\n\nclass NS3DDataset(Dataset):\n def __init__(self, paths, \n data_res, pde_res,\n n_samples=None, \n offset=0,\n t_duration=1.0, \n sub_x=1, \n sub_t=1,\n train=True):\n super().__init__()\n self.data_res = data_res\n self.pde_res = pde_res\n self.t_duration = t_duration\n self.paths = paths\n self.offset = offset\n self.n_samples = n_samples\n self.load(train=train, sub_x=sub_x, sub_t=sub_t)\n \n def load(self, train=True, sub_x=1, sub_t=1):\n data_list = []\n for datapath in self.paths:\n batch = np.load(datapath, mmap_mode='r')\n\n batch = torch.from_numpy(batch[:, ::sub_t, ::sub_x, ::sub_x]).to(torch.float32)\n if self.t_duration == 0.5:\n batch = self.extract(batch)\n data_list.append(batch.permute(0, 2, 3, 1))\n data = torch.cat(data_list, dim=0)\n if self.n_samples:\n if train:\n data = data[self.offset: self.offset + self.n_samples]\n else:\n data = data[self.offset + self.n_samples:]\n \n N = data.shape[0]\n S = data.shape[1]\n T = data.shape[-1]\n a_data = data[:, :, :, 0:1, None].repeat([1, 1, 1, T, 1])\n gridx, gridy, gridt = get_grid3d(S, T)\n a_data = torch.cat((\n gridx.repeat([N, 1, 1, 1, 1]),\n gridy.repeat([N, 1, 1, 1, 1]),\n gridt.repeat([N, 1, 1, 1, 1]),\n a_data), dim=-1)\n self.data = data # N, S, S, T, 1\n self.a_data = a_data # N, S, S, T, 4\n \n self.data_s_step = data.shape[1] // self.data_res[0]\n self.data_t_step = data.shape[3] // (self.data_res[2] - 1)\n\n def __getitem__(self, idx):\n return self.data[idx, ::self.data_s_step, ::self.data_s_step, ::self.data_t_step], self.a_data[idx]\n\n def __len__(self, ):\n return self.data.shape[0]\n\n @staticmethod\n def extract(data):\n '''\n Extract data with time range 0-0.5, 0.25-0.75, 0.5-1.0, 0.75-1.25,...\n Args:\n data: tensor with size N x 129 x 128 x 128\n\n Returns:\n output: (4*N-1) x 65 x 128 x 128\n '''\n T = data.shape[1] // 2\n interval = data.shape[1] // 4\n N = data.shape[0]\n new_data = torch.zeros(4 * N - 1, T + 1, data.shape[2], data.shape[3])\n for i in range(N):\n for j in range(4):\n if i == N - 1 and j == 3:\n # reach boundary\n break\n if j != 3:\n new_data[i * 4 + j] = data[i, interval * j:interval * j + T + 1]\n else:\n new_data[i * 4 + j, 0: interval] = data[i, interval * j:interval * j + interval]\n new_data[i * 4 + j, interval: T + 1] = data[i + 1, 0:interval + 1]\n return new_data\n\n\nclass KFDataset(Dataset):\n def __init__(self, paths, \n data_res, pde_res, \n raw_res, \n n_samples=None, \n total_samples=None,\n idx=0,\n offset=0,\n t_duration=1.0):\n super().__init__()\n self.data_res = data_res # data resolution\n self.pde_res = pde_res # pde loss resolution\n self.raw_res = raw_res # raw data resolution\n self.t_duration = t_duration\n self.paths = paths\n self.offset = offset\n self.n_samples = n_samples\n if t_duration == 1.0:\n self.T = self.pde_res[2]\n else:\n self.T = int(self.pde_res[2] * t_duration) + 1 # number of points in time dimension\n\n self.load()\n if total_samples is not None:\n print(f'Load {total_samples} samples starting from {idx}th sample')\n self.data = self.data[idx:idx + total_samples]\n self.a_data = self.a_data[idx:idx + total_samples]\n \n self.data_s_step = pde_res[0] // data_res[0]\n self.data_t_step = (pde_res[2] - 1) // (data_res[2] - 1)\n\n def load(self):\n datapath = self.paths[0]\n raw_data = np.load(datapath, mmap_mode='r')\n # subsample ratio\n sub_x = self.raw_res[0] // self.data_res[0]\n sub_t = (self.raw_res[2] - 1) // (self.data_res[2] - 1)\n \n a_sub_x = self.raw_res[0] // self.pde_res[0]\n # load data\n data = raw_data[self.offset: self.offset + self.n_samples, ::sub_t, ::sub_x, ::sub_x]\n # divide data\n if self.t_duration != 0.:\n end_t = self.raw_res[2] - 1\n K = int(1/self.t_duration)\n step = end_t // K\n data = self.partition(data)\n a_data = raw_data[self.offset: self.offset + self.n_samples, 0:end_t:step, ::a_sub_x, ::a_sub_x]\n a_data = a_data.reshape(self.n_samples * K, 1, self.pde_res[0], self.pde_res[1]) # 2N x 1 x S x S\n else:\n a_data = raw_data[self.offset: self.offset + self.n_samples, 0:1, ::a_sub_x, ::a_sub_x]\n\n # convert into torch tensor\n data = torch.from_numpy(data).to(torch.float32)\n a_data = torch.from_numpy(a_data).to(torch.float32).permute(0, 2, 3, 1)\n self.data = data.permute(0, 2, 3, 1)\n\n S = self.pde_res[1]\n \n a_data = a_data[:, :, :, :, None] # N x S x S x 1 x 1\n gridx, gridy, gridt = get_grid3d(S, self.T)\n self.grid = torch.cat((gridx[0], gridy[0], gridt[0]), dim=-1) # S x S x T x 3\n self.a_data = a_data\n\n def partition(self, data):\n '''\n Args:\n data: tensor with size N x T x S x S\n\n Returns:\n output: int(1/t_duration) *N x (T//2 + 1) x 128 x 128\n '''\n N, T, S = data.shape[:3]\n K = int(1 / self.t_duration)\n new_data = np.zeros((K * N, T // K + 1, S, S))\n step = T // K\n for i in range(N):\n for j in range(K):\n new_data[i * K + j] = data[i, j * step: (j+1) * step + 1]\n return new_data\n\n\n def __getitem__(self, idx):\n a_data = torch.cat((\n self.grid, \n self.a_data[idx].repeat(1, 1, self.T, 1)\n ), dim=-1)\n return self.data[idx], a_data\n\n def __len__(self, ):\n return self.data.shape[0]\n\n\nclass BurgerData(Dataset):\n '''\n members: \n - t, x, Exact: raw data\n - X, T: meshgrid \n - X_star, u_star: flattened (x, t), u array\n - lb, ub: lower bound and upper bound vector\n - X_u, u: boundary condition data (x, t), u\n '''\n\n def __init__(self, datapath):\n data = scipy.io.loadmat(datapath)\n\n # raw 2D data\n self.t = data['t'].flatten()[:, None] # (100,1)\n self.x = data['x'].flatten()[:, None] # (256, 1)\n self.Exact = np.real(data['usol']).T # (100, 256)\n\n # Flattened sequence\n self.get_flatten_data()\n self.get_boundary_data()\n\n def __len__(self):\n return self.Exact.shape[0]\n\n def __getitem__(self, idx):\n return self.X_star[idx], self.u_star[idx]\n\n def get_flatten_data(self):\n X, T = np.meshgrid(self.x, self.t)\n self.X, self.T = X, T\n self.X_star = np.hstack((X.flatten()[:, None], T.flatten()[:, None]))\n self.u_star = self.Exact.flatten()[:, None]\n\n # lower bound of (x, t): 2-dimensional vector\n self.lb = self.X_star.min(0)\n # upper bound of (x, t): 2-dimensional vector\n self.ub = self.X_star.max(0)\n\n def get_boundary_data(self):\n xx1 = np.hstack((self.X[0:1, :].T, self.T[0:1, :].T))\n uu1 = self.Exact[0:1, :].T\n xx2 = np.hstack((self.X[:, 0:1], self.T[:, 0:1]))\n uu2 = self.Exact[:, 0:1]\n xx3 = np.hstack((self.X[:, -1:], self.T[:, -1:]))\n uu3 = self.Exact[:, -1:]\n self.X_u = np.vstack([xx1, xx2, xx3])\n self.u = np.vstack([uu1, uu2, uu3])\n\n def sample_xt(self, N=10000):\n '''\n Sample (x, t) pairs within the boundary\n Return:\n - X_f: (N, 2) array\n '''\n X_f = self.lb + (self.ub - self.lb) * lhs(2, N)\n X_f = np.vstack((X_f, self.X_u))\n return X_f\n\n def sample_xu(self, N=100):\n '''\n Sample N points from boundary data\n Return: \n - X_u: (N, 2) array \n - u: (N, 1) array\n '''\n idx = np.random.choice(self.X_u.shape[0], N, replace=False)\n X_u = self.X_u[idx, :]\n u = self.u[idx, :]\n return X_u, u\n\n\nclass DarcyFlow(Dataset):\n def __init__(self,\n datapath,\n nx, sub,\n offset=0,\n num=1):\n self.S = int(nx // sub) + 1 if sub > 1 else nx\n data = scipy.io.loadmat(datapath)\n a = data['coeff']\n u = data['sol']\n self.a = torch.tensor(a[offset: offset + num, ::sub, ::sub], dtype=torch.float)\n self.u = torch.tensor(u[offset: offset + num, ::sub, ::sub], dtype=torch.float)\n self.mesh = torch2dgrid(self.S, self.S)\n\n def __len__(self):\n return self.a.shape[0]\n\n def __getitem__(self, item):\n fa = self.a[item]\n return torch.cat([fa.unsqueeze(2), self.mesh], dim=2), self.u[item]\n\n\nclass DarcyIC(Dataset):\n def __init__(self,\n datapath,\n nx, sub,\n offset=0,\n num=1):\n self.S = int(nx // sub) + 1 if sub > 1 else nx\n data = scipy.io.loadmat(datapath)\n a = data['coeff']\n self.a = torch.tensor(a[offset: offset + num, ::sub, ::sub], dtype=torch.float)\n self.mesh = torch2dgrid(self.S, self.S)\n data = scipy.io.loadmat(datapath)\n a = data['coeff']\n u = data['sol']\n self.a = torch.tensor(a[offset: offset + num, ::sub, ::sub], dtype=torch.float)\n self.u = torch.tensor(u[offset: offset + num, ::sub, ::sub], dtype=torch.float)\n self.mesh = torch2dgrid(self.S, self.S)\n\n def __len__(self):\n return self.a.shape[0]\n\n def __getitem__(self, item):\n fa = self.a[item]\n return torch.cat([fa.unsqueeze(2), self.mesh], dim=2) \n\n\nclass DarcyCombo(Dataset):\n def __init__(self, \n datapath, \n nx, \n sub, pde_sub, \n num=1000, offset=0) -> None:\n super().__init__()\n self.S = int(nx // sub) + 1 if sub > 1 else nx\n self.pde_S = int(nx // pde_sub) + 1 if sub > 1 else nx\n data = scipy.io.loadmat(datapath)\n a = data['coeff']\n u = data['sol']\n self.a = torch.tensor(a[offset: offset + num, ::sub, ::sub], dtype=torch.float)\n self.u = torch.tensor(u[offset: offset + num, ::sub, ::sub], dtype=torch.float)\n self.mesh = torch2dgrid(self.S, self.S)\n self.pde_a = torch.tensor(a[offset: offset + num, ::pde_sub, ::pde_sub], dtype=torch.float)\n self.pde_mesh = torch2dgrid(self.pde_S, self.pde_S)\n\n def __len__(self):\n return self.a.shape[0]\n\n def __getitem__(self, item):\n fa = self.a[item]\n pde_a = self.pde_a[item]\n data_ic = torch.cat([fa.unsqueeze(2), self.mesh], dim=2)\n pde_ic = torch.cat([pde_a.unsqueeze(2), self.pde_mesh], dim=2)\n return data_ic, self.u[item], pde_ic\n\n'''\ndataset class for loading initial conditions for Komogrov flow\n'''\nclass KFaDataset(Dataset):\n def __init__(self, paths, \n pde_res, \n raw_res, \n n_samples=None, \n offset=0,\n t_duration=1.0):\n super().__init__()\n self.pde_res = pde_res # pde loss resolution\n self.raw_res = raw_res # raw data resolution\n self.t_duration = t_duration\n self.paths = paths\n self.offset = offset\n self.n_samples = n_samples\n if t_duration == 1.0:\n self.T = self.pde_res[2]\n else:\n self.T = int(self.pde_res[2] * t_duration) + 1 # number of points in time dimension\n\n self.load()\n\n def load(self):\n datapath = self.paths[0]\n raw_data = np.load(datapath, mmap_mode='r')\n # subsample ratio\n a_sub_x = self.raw_res[0] // self.pde_res[0]\n # load data\n if self.t_duration != 0.:\n end_t = self.raw_res[2] - 1\n K = int(1/self.t_duration)\n step = end_t // K\n a_data = raw_data[self.offset: self.offset + self.n_samples, 0:end_t:step, ::a_sub_x, ::a_sub_x]\n a_data = a_data.reshape(self.n_samples * K, 1, self.pde_res[0], self.pde_res[1]) # 2N x 1 x S x S\n else:\n a_data = raw_data[self.offset: self.offset + self.n_samples, 0:1, ::a_sub_x, ::a_sub_x]\n\n # convert into torch tensor\n a_data = torch.from_numpy(a_data).to(torch.float32).permute(0, 2, 3, 1)\n S = self.pde_res[1]\n a_data = a_data[:, :, :, :, None] # N x S x S x 1 x 1\n gridx, gridy, gridt = get_grid3d(S, self.T)\n self.grid = torch.cat((gridx[0], gridy[0], gridt[0]), dim=-1) # S x S x T x 3\n self.a_data = a_data\n\n def __getitem__(self, idx):\n a_data = torch.cat((\n self.grid, \n self.a_data[idx].repeat(1, 1, self.T, 1)\n ), dim=-1)\n return a_data\n\n def __len__(self, ):\n return self.a_data.shape[0]" }, { "path": "train_utils/distributed.py", "content": "import os\nimport torch\nimport torch.distributed as dist\n\n\ndef setup(rank, world_size):\n os.environ['MASTER_ADDR'] = 'localhost'\n os.environ['MASTER_PORT'] = '7777'\n dist.init_process_group(\"nccl\", rank=rank, world_size=world_size)\n\n\ndef cleanup():\n dist.destroy_process_group()\n\n\ndef get_world_size():\n if not dist.is_available() or not dist.is_initialized():\n return 1\n\n return dist.get_world_size()\n\n\ndef all_reduce_mean(tensor):\n '''\n Reduce the tensor across all machines, the operation is in-place.\n :param tensor: tensor to reduce\n :return: reduced tensor\n '''\n if not dist.is_available() or not dist.is_initialized():\n return tensor\n\n world_size = get_world_size()\n dist.all_reduce(tensor, op=dist.ReduceOp.SUM)\n return tensor.div_(world_size)\n\n\ndef reduce_sum(tensor):\n '''\n Reduce the tensor across all machines. Only process with rank 0 will receive the final result\n Args:\n tensor: input and ouput of the collective. The function operates in-place\n Returns:\n final result\n '''\n if not dist.is_available() or not dist.is_initialized():\n return tensor\n\n dist.reduce(tensor, dst=0, op=dist.ReduceOp.SUM)\n return tensor\n\n\ndef reduce_loss_dict(loss_dict):\n if not dist.is_available() or dist.is_initialized():\n return loss_dict\n world_size = get_world_size()\n\n if world_size < 2:\n return loss_dict\n\n with torch.no_grad():\n keys = []\n losses = []\n\n for k in sorted(loss_dict.keys()):\n keys.append(k)\n losses.append(loss_dict[k])\n\n losses = torch.stack(losses, 0)\n dist.reduce(losses, dst=0)\n\n if dist.get_rank() == 0:\n losses /= world_size\n\n reduced_losses = {k: v for k, v in zip(keys, losses)}\n\n return reduced_losses" }, { "path": "train_utils/eval_2d.py", "content": "from tqdm import tqdm\nimport numpy as np\n\nimport torch\n\nfrom .losses import LpLoss, darcy_loss, PINO_loss\n\ntry:\n import wandb\nexcept ImportError:\n wandb = None\n\n\ndef eval_darcy(model,\n dataloader,\n config,\n device,\n use_tqdm=True):\n model.eval()\n myloss = LpLoss(size_average=True)\n if use_tqdm:\n pbar = tqdm(dataloader, dynamic_ncols=True, smoothing=0.05)\n else:\n pbar = dataloader\n\n mesh = dataloader.dataset.mesh\n mollifier = torch.sin(np.pi * mesh[..., 0]) * torch.sin(np.pi * mesh[..., 1]) * 0.001\n mollifier = mollifier.to(device)\n f_val = []\n test_err = []\n\n with torch.no_grad():\n for x, y in pbar:\n x, y = x.to(device), y.to(device)\n\n pred = model(x).reshape(y.shape)\n pred = pred * mollifier\n\n data_loss = myloss(pred, y)\n a = x[..., 0]\n f_loss = darcy_loss(pred, a)\n\n test_err.append(data_loss.item())\n f_val.append(f_loss.item())\n if use_tqdm:\n pbar.set_description(\n (\n f'Equation error: {f_loss.item():.5f}, test l2 error: {data_loss.item()}'\n )\n )\n mean_f_err = np.mean(f_val)\n std_f_err = np.std(f_val, ddof=1) / np.sqrt(len(f_val))\n\n mean_err = np.mean(test_err)\n std_err = np.std(test_err, ddof=1) / np.sqrt(len(test_err))\n\n print(f'==Averaged relative L2 error mean: {mean_err}, std error: {std_err}==\\n'\n f'==Averaged equation error mean: {mean_f_err}, std error: {std_f_err}==')\n\n\ndef eval_burgers(model,\n dataloader,\n v,\n config,\n device,\n use_tqdm=True):\n model.eval()\n myloss = LpLoss(size_average=True)\n if use_tqdm:\n pbar = tqdm(dataloader, dynamic_ncols=True, smoothing=0.05)\n else:\n pbar = dataloader\n\n test_err = []\n f_err = []\n\n for x, y in pbar:\n x, y = x.to(device), y.to(device)\n out = model(x).reshape(y.shape)\n data_loss = myloss(out, y)\n\n loss_u, f_loss = PINO_loss(out, x[:, 0, :, 0], v)\n test_err.append(data_loss.item())\n f_err.append(f_loss.item())\n\n mean_f_err = np.mean(f_err)\n std_f_err = np.std(f_err, ddof=1) / np.sqrt(len(f_err))\n\n mean_err = np.mean(test_err)\n std_err = np.std(test_err, ddof=1) / np.sqrt(len(test_err))\n\n print(f'==Averaged relative L2 error mean: {mean_err}, std error: {std_err}==\\n'\n f'==Averaged equation error mean: {mean_f_err}, std error: {std_f_err}==')\n\n" }, { "path": "train_utils/eval_3d.py", "content": "import torch\nimport torch.nn.functional as F\n\nfrom tqdm import tqdm\nfrom timeit import default_timer\n\nfrom .losses import LpLoss, PINO_loss3d\n\ntry:\n import wandb\nexcept ImportError:\n wandb = None\n\n\ndef eval_ns(model, # model\n loader, # dataset instance\n dataloader, # dataloader\n forcing, # forcing\n config, # configuration dict\n device, # device id\n log=False,\n project='PINO-default',\n group='FDM',\n tags=['Nan'],\n use_tqdm=True):\n '''\n Evaluate the model for Navier Stokes equation\n '''\n if wandb and log:\n run = wandb.init(project=project,\n entity=config['log']['entity'],\n group=group,\n config=config,\n tags=tags, reinit=True,\n settings=wandb.Settings(start_method=\"fork\"))\n # data parameters\n v = 1 / config['data']['Re']\n S, T = loader.S, loader.T\n t_interval = config['data']['time_interval']\n # eval settings\n batch_size = config['test']['batchsize']\n\n model.eval()\n myloss = LpLoss(size_average=True)\n if use_tqdm:\n pbar = tqdm(dataloader, dynamic_ncols=True, smoothing=0.05)\n else:\n pbar = dataloader\n loss_dict = {'f_error': 0.0,\n 'test_l2': 0.0}\n start_time = default_timer()\n with torch.no_grad():\n for x, y in pbar:\n x, y = x.to(device), y.to(device)\n x_in = F.pad(x, (0, 0, 0, 5), \"constant\", 0)\n out = model(x_in).reshape(batch_size, S, S, T + 5)\n out = out[..., :-5]\n x = x[:, :, :, 0, -1]\n loss_l2 = myloss(out.view(batch_size, S, S, T), y.view(batch_size, S, S, T))\n loss_ic, loss_f = PINO_loss3d(out.view(batch_size, S, S, T), x, forcing, v, t_interval)\n\n loss_dict['f_error'] += loss_f\n loss_dict['test_l2'] += loss_l2\n if device == 0 and use_tqdm:\n pbar.set_description(\n (\n f'Train f error: {loss_f.item():.5f}; Test l2 error: {loss_l2.item():.5f}'\n )\n )\n end_time = default_timer()\n test_l2 = loss_dict['test_l2'].item() / len(dataloader)\n loss_f = loss_dict['f_error'].item() / len(dataloader)\n print(f'==Averaged relative L2 error is: {test_l2}==\\n'\n f'==Averaged equation error is: {loss_f}==')\n print(f'Time cost: {end_time - start_time} s')\n if device == 0:\n if wandb and log:\n wandb.log(\n {\n 'Train f error': loss_f,\n 'Test L2 error': test_l2,\n }\n )\n run.finish()\n" }, { "path": "train_utils/losses.py", "content": "import numpy as np\nimport torch\nimport torch.nn.functional as F\n\n\ndef FDM_Darcy(u, a, D=1):\n batchsize = u.size(0)\n size = u.size(1)\n u = u.reshape(batchsize, size, size)\n a = a.reshape(batchsize, size, size)\n dx = D / (size - 1)\n dy = dx\n\n # ux: (batch, size-2, size-2)\n ux = (u[:, 2:, 1:-1] - u[:, :-2, 1:-1]) / (2 * dx)\n uy = (u[:, 1:-1, 2:] - u[:, 1:-1, :-2]) / (2 * dy)\n\n # ax = (a[:, 2:, 1:-1] - a[:, :-2, 1:-1]) / (2 * dx)\n # ay = (a[:, 1:-1, 2:] - a[:, 1:-1, :-2]) / (2 * dy)\n # uxx = (u[:, 2:, 1:-1] -2*u[:,1:-1,1:-1] +u[:, :-2, 1:-1]) / (dx**2)\n # uyy = (u[:, 1:-1, 2:] -2*u[:,1:-1,1:-1] +u[:, 1:-1, :-2]) / (dy**2)\n\n a = a[:, 1:-1, 1:-1]\n # u = u[:, 1:-1, 1:-1]\n # Du = -(ax*ux + ay*uy + a*uxx + a*uyy)\n\n # inner1 = torch.mean(a*(ux**2 + uy**2), dim=[1,2])\n # inner2 = torch.mean(f*u, dim=[1,2])\n # return 0.5*inner1 - inner2\n\n aux = a * ux\n auy = a * uy\n auxx = (aux[:, 2:, 1:-1] - aux[:, :-2, 1:-1]) / (2 * dx)\n auyy = (auy[:, 1:-1, 2:] - auy[:, 1:-1, :-2]) / (2 * dy)\n Du = - (auxx + auyy)\n return Du\n\n\ndef darcy_loss(u, a):\n batchsize = u.size(0)\n size = u.size(1)\n u = u.reshape(batchsize, size, size)\n a = a.reshape(batchsize, size, size)\n lploss = LpLoss(size_average=True)\n\n # index_x = torch.cat([torch.tensor(range(0, size)), (size - 1) * torch.ones(size), torch.tensor(range(size-1, 1, -1)),\n # torch.zeros(size)], dim=0).long()\n # index_y = torch.cat([(size - 1) * torch.ones(size), torch.tensor(range(size-1, 1, -1)), torch.zeros(size),\n # torch.tensor(range(0, size))], dim=0).long()\n\n # boundary_u = u[:, index_x, index_y]\n # truth_u = torch.zeros(boundary_u.shape, device=u.device)\n # loss_u = lploss.abs(boundary_u, truth_u)\n\n Du = FDM_Darcy(u, a)\n f = torch.ones(Du.shape, device=u.device)\n loss_f = lploss.rel(Du, f)\n\n # im = (Du-f)[0].detach().cpu().numpy()\n # plt.imshow(im)\n # plt.show()\n\n # loss_f = FDM_Darcy(u, a)\n # loss_f = torch.mean(loss_f)\n return loss_f\n\n\ndef FDM_NS_vorticity(w, v=1/40, t_interval=1.0):\n batchsize = w.size(0)\n nx = w.size(1)\n ny = w.size(2)\n nt = w.size(3)\n device = w.device\n w = w.reshape(batchsize, nx, ny, nt)\n\n w_h = torch.fft.fft2(w, dim=[1, 2])\n # Wavenumbers in y-direction\n k_max = nx//2\n N = nx\n k_x = torch.cat((torch.arange(start=0, end=k_max, step=1, device=device),\n torch.arange(start=-k_max, end=0, step=1, device=device)), 0).reshape(N, 1).repeat(1, N).reshape(1,N,N,1)\n k_y = torch.cat((torch.arange(start=0, end=k_max, step=1, device=device),\n torch.arange(start=-k_max, end=0, step=1, device=device)), 0).reshape(1, N).repeat(N, 1).reshape(1,N,N,1)\n # Negative Laplacian in Fourier space\n lap = (k_x ** 2 + k_y ** 2)\n lap[0, 0, 0, 0] = 1.0\n f_h = w_h / lap\n\n ux_h = 1j * k_y * f_h\n uy_h = -1j * k_x * f_h\n wx_h = 1j * k_x * w_h\n wy_h = 1j * k_y * w_h\n wlap_h = -lap * w_h\n\n ux = torch.fft.irfft2(ux_h[:, :, :k_max + 1], dim=[1, 2])\n uy = torch.fft.irfft2(uy_h[:, :, :k_max + 1], dim=[1, 2])\n wx = torch.fft.irfft2(wx_h[:, :, :k_max+1], dim=[1,2])\n wy = torch.fft.irfft2(wy_h[:, :, :k_max+1], dim=[1,2])\n wlap = torch.fft.irfft2(wlap_h[:, :, :k_max+1], dim=[1,2])\n\n dt = t_interval / (nt-1)\n wt = (w[:, :, :, 2:] - w[:, :, :, :-2]) / (2 * dt)\n\n Du1 = wt + (ux*wx + uy*wy - v*wlap)[...,1:-1] #- forcing\n return Du1\n\n\ndef Autograd_Burgers(u, grid, v=1/100):\n from torch.autograd import grad\n gridt, gridx = grid\n\n ut = grad(u.sum(), gridt, create_graph=True)[0]\n ux = grad(u.sum(), gridx, create_graph=True)[0]\n uxx = grad(ux.sum(), gridx, create_graph=True)[0]\n Du = ut + ux*u - v*uxx\n return Du, ux, uxx, ut\n\n\ndef AD_loss(u, u0, grid, index_ic=None, p=None, q=None):\n batchsize = u.size(0)\n # lploss = LpLoss(size_average=True)\n\n Du, ux, uxx, ut = Autograd_Burgers(u, grid)\n\n if index_ic is None:\n # u in on a uniform grid\n nt = u.size(1)\n nx = u.size(2)\n u = u.reshape(batchsize, nt, nx)\n\n index_t = torch.zeros(nx,).long()\n index_x = torch.tensor(range(nx)).long()\n boundary_u = u[:, index_t, index_x]\n\n # loss_bc0 = F.mse_loss(u[:, :, 0], u[:, :, -1])\n # loss_bc1 = F.mse_loss(ux[:, :, 0], ux[:, :, -1])\n else:\n # u is randomly sampled, 0:p are BC, p:2p are ic, 2p:2p+q are interior\n boundary_u = u[:, :p]\n batch_index = torch.tensor(range(batchsize)).reshape(batchsize, 1).repeat(1, p)\n u0 = u0[batch_index, index_ic]\n\n # loss_bc0 = F.mse_loss(u[:, p:p+p//2], u[:, p+p//2:2*p])\n # loss_bc1 = F.mse_loss(ux[:, p:p+p//2], ux[:, p+p//2:2*p])\n\n loss_ic = F.mse_loss(boundary_u, u0)\n f = torch.zeros(Du.shape, device=u.device)\n loss_f = F.mse_loss(Du, f)\n return loss_ic, loss_f\n\n\nclass LpLoss(object):\n '''\n loss function with rel/abs Lp loss\n '''\n def __init__(self, d=2, p=2, size_average=True, reduction=True):\n super(LpLoss, self).__init__()\n\n #Dimension and Lp-norm type are postive\n assert d > 0 and p > 0\n\n self.d = d\n self.p = p\n self.reduction = reduction\n self.size_average = size_average\n\n def abs(self, x, y):\n num_examples = x.size()[0]\n\n #Assume uniform mesh\n h = 1.0 / (x.size()[1] - 1.0)\n\n all_norms = (h**(self.d/self.p))*torch.norm(x.view(num_examples,-1) - y.view(num_examples,-1), self.p, 1)\n\n if self.reduction:\n if self.size_average:\n return torch.mean(all_norms)\n else:\n return torch.sum(all_norms)\n\n return all_norms\n\n def rel(self, x, y):\n num_examples = x.size()[0]\n\n diff_norms = torch.norm(x.reshape(num_examples,-1) - y.reshape(num_examples,-1), self.p, 1)\n y_norms = torch.norm(y.reshape(num_examples,-1), self.p, 1)\n\n if self.reduction:\n if self.size_average:\n return torch.mean(diff_norms/y_norms)\n else:\n return torch.sum(diff_norms/y_norms)\n\n return diff_norms/y_norms\n\n def __call__(self, x, y):\n return self.rel(x, y)\n\n\ndef FDM_Burgers(u, v, D=1):\n batchsize = u.size(0)\n nt = u.size(1)\n nx = u.size(2)\n\n u = u.reshape(batchsize, nt, nx)\n dt = D / (nt-1)\n dx = D / (nx)\n\n u_h = torch.fft.fft(u, dim=2)\n # Wavenumbers in y-direction\n k_max = nx//2\n k_x = torch.cat((torch.arange(start=0, end=k_max, step=1, device=u.device),\n torch.arange(start=-k_max, end=0, step=1, device=u.device)), 0).reshape(1,1,nx)\n ux_h = 2j *np.pi*k_x*u_h\n uxx_h = 2j *np.pi*k_x*ux_h\n ux = torch.fft.irfft(ux_h[:, :, :k_max+1], dim=2, n=nx)\n uxx = torch.fft.irfft(uxx_h[:, :, :k_max+1], dim=2, n=nx)\n ut = (u[:, 2:, :] - u[:, :-2, :]) / (2 * dt)\n Du = ut + (ux*u - v*uxx)[:,1:-1,:]\n return Du\n\n\ndef PINO_loss(u, u0, v):\n batchsize = u.size(0)\n nt = u.size(1)\n nx = u.size(2)\n\n u = u.reshape(batchsize, nt, nx)\n # lploss = LpLoss(size_average=True)\n\n index_t = torch.zeros(nx,).long()\n index_x = torch.tensor(range(nx)).long()\n boundary_u = u[:, index_t, index_x]\n loss_u = F.mse_loss(boundary_u, u0)\n\n Du = FDM_Burgers(u, v)[:, :, :]\n f = torch.zeros(Du.shape, device=u.device)\n loss_f = F.mse_loss(Du, f)\n\n # loss_bc0 = F.mse_loss(u[:, :, 0], u[:, :, -1])\n # loss_bc1 = F.mse_loss((u[:, :, 1] - u[:, :, -1]) /\n # (2/(nx)), (u[:, :, 0] - u[:, :, -2])/(2/(nx)))\n return loss_u, loss_f\n\n\ndef PINO_loss3d(u, u0, forcing, v=1/40, t_interval=1.0):\n batchsize = u.size(0)\n nx = u.size(1)\n ny = u.size(2)\n nt = u.size(3)\n\n u = u.reshape(batchsize, nx, ny, nt)\n lploss = LpLoss(size_average=True)\n\n u_in = u[:, :, :, 0]\n loss_ic = lploss(u_in, u0)\n\n Du = FDM_NS_vorticity(u, v, t_interval)\n f = forcing.repeat(batchsize, 1, 1, nt-2)\n loss_f = lploss(Du, f)\n\n return loss_ic, loss_f\n\n\ndef PDELoss(model, x, t, nu):\n '''\n Compute the residual of PDE:\n residual = u_t + u * u_x - nu * u_{xx} : (N,1)\n\n Params:\n - model\n - x, t: (x, t) pairs, (N, 2) tensor\n - nu: constant of PDE\n Return:\n - mean of residual : scalar\n '''\n u = model(torch.cat([x, t], dim=1))\n # First backward to compute u_x (shape: N x 1), u_t (shape: N x 1)\n grad_x, grad_t = torch.autograd.grad(outputs=[u.sum()], inputs=[x, t], create_graph=True)\n # Second backward to compute u_{xx} (shape N x 1)\n\n gradgrad_x, = torch.autograd.grad(outputs=[grad_x.sum()], inputs=[x], create_graph=True)\n\n residual = grad_t + u * grad_x - nu * gradgrad_x\n return residual\n\n\ndef get_forcing(S):\n x1 = torch.tensor(np.linspace(0, 2*np.pi, S, endpoint=False), dtype=torch.float).reshape(S, 1).repeat(1, S)\n x2 = torch.tensor(np.linspace(0, 2*np.pi, S, endpoint=False), dtype=torch.float).reshape(1, S).repeat(S, 1)\n return -4 * (torch.cos(4*(x2))).reshape(1,S,S,1)" }, { "path": "train_utils/negadam.py", "content": "import math\nimport torch\nfrom torch import Tensor\nfrom typing import List, Optional\nfrom torch.optim.optimizer import Optimizer\n\n\ndef adam(params: List[Tensor],\n grads: List[Tensor],\n exp_avgs: List[Tensor],\n exp_avg_sqs: List[Tensor],\n max_exp_avg_sqs: List[Tensor],\n state_steps: List[int],\n *,\n amsgrad: bool,\n beta1: float,\n beta2: float,\n lr: float,\n weight_decay: float,\n eps: float):\n r\"\"\"Functional API that performs Adam algorithm computation.\n See :class:`~torch.optim.Adam` for details.\n \"\"\"\n\n for i, param in enumerate(params):\n\n grad = grads[i]\n exp_avg = exp_avgs[i]\n exp_avg_sq = exp_avg_sqs[i]\n step = state_steps[i]\n\n bias_correction1 = 1 - beta1 ** step\n bias_correction2 = 1 - beta2 ** step\n\n if weight_decay != 0:\n grad = grad.add(param, alpha=weight_decay)\n\n # Decay the first and second moment running average coefficient\n exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)\n exp_avg_sq.mul_(beta2).addcmul_(grad, grad.conj(), value=1 - beta2)\n if amsgrad:\n # Maintains the maximum of all 2nd moment running avg. till now\n torch.maximum(max_exp_avg_sqs[i], exp_avg_sq, out=max_exp_avg_sqs[i])\n # Use the max. for normalizing running avg. of gradient\n denom = (max_exp_avg_sqs[i].sqrt() / math.sqrt(bias_correction2)).add_(eps)\n else:\n denom = (exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(eps)\n\n step_size = lr / bias_correction1\n\n param.addcdiv_(exp_avg, denom, value=step_size)\n\n\nclass NAdam(Optimizer):\n r\"\"\"Implements Adam algorithm.\n It has been proposed in `Adam: A Method for Stochastic Optimization`_.\n The implementation of the L2 penalty follows changes proposed in\n `Decoupled Weight Decay Regularization`_.\n Args:\n params (iterable): iterable of parameters to optimize or dicts defining\n parameter groups\n lr (float, optional): learning rate (default: 1e-3)\n betas (Tuple[float, float], optional): coefficients used for computing\n running averages of gradient and its square (default: (0.9, 0.999))\n eps (float, optional): term added to the denominator to improve\n numerical stability (default: 1e-8)\n weight_decay (float, optional): weight decay (L2 penalty) (default: 0)\n amsgrad (boolean, optional): whether to use the AMSGrad variant of this\n algorithm from the paper `On the Convergence of Adam and Beyond`_\n (default: False)\n .. _Adam\\: A Method for Stochastic Optimization:\n https://arxiv.org/abs/1412.6980\n .. _Decoupled Weight Decay Regularization:\n https://arxiv.org/abs/1711.05101\n .. _On the Convergence of Adam and Beyond:\n https://openreview.net/forum?id=ryQu7f-RZ\n \"\"\"\n\n def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8,\n weight_decay=0, amsgrad=False):\n if not 0.0 <= lr:\n raise ValueError(\"Invalid learning rate: {}\".format(lr))\n if not 0.0 <= eps:\n raise ValueError(\"Invalid epsilon value: {}\".format(eps))\n if not 0.0 <= betas[0] < 1.0:\n raise ValueError(\"Invalid beta parameter at index 0: {}\".format(betas[0]))\n if not 0.0 <= betas[1] < 1.0:\n raise ValueError(\"Invalid beta parameter at index 1: {}\".format(betas[1]))\n if not 0.0 <= weight_decay:\n raise ValueError(\"Invalid weight_decay value: {}\".format(weight_decay))\n defaults = dict(lr=lr, betas=betas, eps=eps,\n weight_decay=weight_decay, amsgrad=amsgrad)\n super(NAdam, self).__init__(params, defaults)\n\n def __setstate__(self, state):\n super(NAdam, self).__setstate__(state)\n for group in self.param_groups:\n group.setdefault('amsgrad', False)\n\n @torch.no_grad()\n def step(self, closure=None):\n \"\"\"Performs a single optimization step.\n Args:\n closure (callable, optional): A closure that reevaluates the model\n and returns the loss.\n \"\"\"\n loss = None\n if closure is not None:\n with torch.enable_grad():\n loss = closure()\n\n for group in self.param_groups:\n params_with_grad = []\n grads = []\n exp_avgs = []\n exp_avg_sqs = []\n max_exp_avg_sqs = []\n state_steps = []\n beta1, beta2 = group['betas']\n\n for p in group['params']:\n if p.grad is not None:\n params_with_grad.append(p)\n if p.grad.is_sparse:\n raise RuntimeError('Adam does not support sparse gradients, please consider SparseAdam instead')\n grads.append(p.grad)\n\n state = self.state[p]\n # Lazy state initialization\n if len(state) == 0:\n state['step'] = 0\n # Exponential moving average of gradient values\n state['exp_avg'] = torch.zeros_like(p, memory_format=torch.preserve_format)\n # Exponential moving average of squared gradient values\n state['exp_avg_sq'] = torch.zeros_like(p, memory_format=torch.preserve_format)\n if group['amsgrad']:\n # Maintains max of all exp. moving avg. of sq. grad. values\n state['max_exp_avg_sq'] = torch.zeros_like(p, memory_format=torch.preserve_format)\n\n exp_avgs.append(state['exp_avg'])\n exp_avg_sqs.append(state['exp_avg_sq'])\n\n if group['amsgrad']:\n max_exp_avg_sqs.append(state['max_exp_avg_sq'])\n\n # update the steps for each param group update\n state['step'] += 1\n # record the step after step update\n state_steps.append(state['step'])\n\n adam(params_with_grad,\n grads,\n exp_avgs,\n exp_avg_sqs,\n max_exp_avg_sqs,\n state_steps,\n amsgrad=group['amsgrad'],\n beta1=beta1,\n beta2=beta2,\n lr=group['lr'],\n weight_decay=group['weight_decay'],\n eps=group['eps'])\n return loss\n" }, { "path": "train_utils/train_2d.py", "content": "import numpy as np\nimport torch\nfrom tqdm import tqdm\nfrom .utils import save_checkpoint\nfrom .losses import LpLoss, darcy_loss, PINO_loss\n\ntry:\n import wandb\nexcept ImportError:\n wandb = None\n\n\ndef train_2d_operator(model,\n train_loader,\n optimizer, scheduler,\n config,\n rank=0, log=False,\n project='PINO-2d-default',\n group='default',\n tags=['default'],\n use_tqdm=True,\n profile=False):\n '''\n train PINO on Darcy Flow\n Args:\n model:\n train_loader:\n optimizer:\n scheduler:\n config:\n rank:\n log:\n project:\n group:\n tags:\n use_tqdm:\n profile:\n\n Returns:\n\n '''\n if rank == 0 and wandb and log:\n run = wandb.init(project=project,\n entity=config['log']['entity'],\n group=group,\n config=config,\n tags=tags, reinit=True,\n settings=wandb.Settings(start_method=\"fork\"))\n\n data_weight = config['train']['xy_loss']\n f_weight = config['train']['f_loss']\n model.train()\n myloss = LpLoss(size_average=True)\n pbar = range(config['train']['epochs'])\n if use_tqdm:\n pbar = tqdm(pbar, dynamic_ncols=True, smoothing=0.1)\n mesh = train_loader.dataset.mesh\n mollifier = torch.sin(np.pi * mesh[..., 0]) * torch.sin(np.pi * mesh[..., 1]) * 0.001\n mollifier = mollifier.to(rank)\n pde_mesh = train_loader.dataset.pde_mesh\n pde_mol = torch.sin(np.pi * pde_mesh[..., 0]) * torch.sin(np.pi * pde_mesh[..., 1]) * 0.001\n pde_mol = pde_mol.to(rank)\n for e in pbar:\n loss_dict = {'train_loss': 0.0,\n 'data_loss': 0.0,\n 'f_loss': 0.0,\n 'test_error': 0.0}\n for data_ic, u, pde_ic in train_loader:\n data_ic, u, pde_ic = data_ic.to(rank), u.to(rank), pde_ic.to(rank)\n\n optimizer.zero_grad()\n\n # data loss\n if data_weight > 0:\n pred = model(data_ic).squeeze(dim=-1)\n pred = pred * mollifier\n data_loss = myloss(pred, y)\n\n a = x[..., 0]\n f_loss = darcy_loss(pred, a)\n\n loss = data_weight * data_loss + f_weight * f_loss\n loss.backward()\n optimizer.step()\n\n loss_dict['train_loss'] += loss.item() * y.shape[0]\n loss_dict['f_loss'] += f_loss.item() * y.shape[0]\n loss_dict['data_loss'] += data_loss.item() * y.shape[0]\n\n scheduler.step()\n train_loss_val = loss_dict['train_loss'] / len(train_loader.dataset)\n f_loss_val = loss_dict['f_loss'] / len(train_loader.dataset)\n data_loss_val = loss_dict['data_loss'] / len(train_loader.dataset)\n\n if use_tqdm:\n pbar.set_description(\n (\n f'Epoch: {e}, train loss: {train_loss_val:.5f}, '\n f'f_loss: {f_loss_val:.5f}, '\n f'data loss: {data_loss_val:.5f}'\n )\n )\n if wandb and log:\n wandb.log(\n {\n 'train loss': train_loss_val,\n 'f loss': f_loss_val,\n 'data loss': data_loss_val\n }\n )\n save_checkpoint(config['train']['save_dir'],\n config['train']['save_name'],\n model, optimizer)\n if wandb and log:\n run.finish()\n print('Done!')\n\n\ndef train_2d_burger(model,\n train_loader, v,\n optimizer, scheduler,\n config,\n rank=0, log=False,\n project='PINO-2d-default',\n group='default',\n tags=['default'],\n use_tqdm=True):\n if rank == 0 and wandb and log:\n run = wandb.init(project=project,\n entity=config['log']['entity'],\n group=group,\n config=config,\n tags=tags, reinit=True,\n settings=wandb.Settings(start_method=\"fork\"))\n\n data_weight = config['train']['xy_loss']\n f_weight = config['train']['f_loss']\n ic_weight = config['train']['ic_loss']\n model.train()\n myloss = LpLoss(size_average=True)\n pbar = range(config['train']['epochs'])\n if use_tqdm:\n pbar = tqdm(pbar, dynamic_ncols=True, smoothing=0.1)\n\n for e in pbar:\n model.train()\n train_pino = 0.0\n data_l2 = 0.0\n train_loss = 0.0\n\n for x, y in train_loader:\n x, y = x.to(rank), y.to(rank)\n out = model(x).reshape(y.shape)\n data_loss = myloss(out, y)\n\n loss_u, loss_f = PINO_loss(out, x[:, 0, :, 0], v)\n total_loss = loss_u * ic_weight + loss_f * f_weight + data_loss * data_weight\n\n optimizer.zero_grad()\n total_loss.backward()\n optimizer.step()\n\n data_l2 += data_loss.item()\n train_pino += loss_f.item()\n train_loss += total_loss.item()\n scheduler.step()\n data_l2 /= len(train_loader)\n train_pino /= len(train_loader)\n train_loss /= len(train_loader)\n if use_tqdm:\n pbar.set_description(\n (\n f'Epoch {e}, train loss: {train_loss:.5f} '\n f'train f error: {train_pino:.5f}; '\n f'data l2 error: {data_l2:.5f}'\n )\n )\n if wandb and log:\n wandb.log(\n {\n 'Train f error': train_pino,\n 'Train L2 error': data_l2,\n 'Train loss': train_loss,\n }\n )\n\n if e % 100 == 0:\n save_checkpoint(config['train']['save_dir'],\n config['train']['save_name'].replace('.pt', f'_{e}.pt'),\n model, optimizer)\n save_checkpoint(config['train']['save_dir'],\n config['train']['save_name'],\n model, optimizer)\n print('Done!')" }, { "path": "train_utils/train_3d.py", "content": "import torch\nfrom tqdm import tqdm\nfrom timeit import default_timer\nimport torch.nn.functional as F\nfrom .utils import save_checkpoint\nfrom .losses import LpLoss, PINO_loss3d, get_forcing\nfrom .distributed import reduce_loss_dict\nfrom .data_utils import sample_data\n\ntry:\n import wandb\nexcept ImportError:\n wandb = None\n \n\ndef train(model,\n loader, train_loader,\n optimizer, scheduler,\n forcing, config,\n rank=0,\n log=False,\n project='PINO-default',\n group='FDM',\n tags=['Nan'],\n use_tqdm=True,\n profile=False):\n if rank == 0 and wandb and log:\n run = wandb.init(project=project,\n entity=config['log']['entity'],\n group=group,\n config=config,\n tags=tags, reinit=True,\n settings=wandb.Settings(start_method=\"fork\"))\n\n # data parameters\n v = 1 / config['data']['Re']\n S, T = loader.S, loader.T\n t_interval = config['data']['time_interval']\n\n # training settings\n batch_size = config['train']['batchsize']\n ic_weight = config['train']['ic_loss']\n f_weight = config['train']['f_loss']\n xy_weight = config['train']['xy_loss']\n\n model.train()\n myloss = LpLoss(size_average=True)\n pbar = range(config['train']['epochs'])\n if use_tqdm:\n pbar = tqdm(pbar, dynamic_ncols=True, smoothing=0.05)\n zero = torch.zeros(1).to(rank)\n\n for ep in pbar:\n loss_dict = {'train_loss': 0.0,\n 'train_ic': 0.0,\n 'train_f': 0.0,\n 'test_l2': 0.0}\n log_dict = {}\n if rank == 0 and profile:\n torch.cuda.synchronize()\n t1 = default_timer()\n # start solving\n for x, y in train_loader:\n x, y = x.to(rank), y.to(rank)\n\n optimizer.zero_grad()\n x_in = F.pad(x, (0, 0, 0, 5), \"constant\", 0)\n out = model(x_in).reshape(batch_size, S, S, T + 5)\n out = out[..., :-5]\n x = x[:, :, :, 0, -1]\n\n loss_l2 = myloss(out.view(batch_size, S, S, T), y.view(batch_size, S, S, T))\n\n if ic_weight != 0 or f_weight != 0:\n loss_ic, loss_f = PINO_loss3d(out.view(batch_size, S, S, T), x, forcing, v, t_interval)\n else:\n loss_ic, loss_f = zero, zero\n\n total_loss = loss_l2 * xy_weight + loss_f * f_weight + loss_ic * ic_weight\n\n total_loss.backward()\n\n optimizer.step()\n loss_dict['train_ic'] += loss_ic\n loss_dict['test_l2'] += loss_l2\n loss_dict['train_loss'] += total_loss\n loss_dict['train_f'] += loss_f\n\n if rank == 0 and profile:\n torch.cuda.synchronize()\n t2 = default_timer()\n log_dict['Time cost'] = t2 - t1\n scheduler.step()\n loss_reduced = reduce_loss_dict(loss_dict)\n train_ic = loss_reduced['train_ic'].item() / len(train_loader)\n train_f = loss_reduced['train_f'].item() / len(train_loader)\n train_loss = loss_reduced['train_loss'].item() / len(train_loader)\n test_l2 = loss_reduced['test_l2'].item() / len(train_loader)\n log_dict = {\n 'Train f error': train_f,\n 'Train L2 error': train_ic,\n 'Train loss': train_loss,\n 'Test L2 error': test_l2\n }\n\n if rank == 0:\n if use_tqdm:\n pbar.set_description(\n (\n f'Train f error: {train_f:.5f}; Train ic l2 error: {train_ic:.5f}. '\n f'Train loss: {train_loss:.5f}; Test l2 error: {test_l2:.5f}'\n )\n )\n if wandb and log:\n wandb.log(log_dict)\n\n if rank == 0:\n save_checkpoint(config['train']['save_dir'],\n config['train']['save_name'],\n model, optimizer)\n if wandb and log:\n run.finish()\n\n\ndef mixed_train(model, # model of neural operator\n train_loader, # dataloader for training with data\n S1, T1, # spacial and time dimension for training with data\n a_loader, # generator for ICs\n S2, T2, # spacial and time dimension for training with equation only\n optimizer, # optimizer\n scheduler, # learning rate scheduler\n config, # configuration dict\n device=torch.device('cpu'),\n log=False, # turn on the wandb\n project='PINO-default', # project name\n group='FDM', # group name\n tags=['Nan'], # tags\n use_tqdm=True): # turn on tqdm\n if wandb and log:\n run = wandb.init(project=project,\n entity=config['log']['entity'],\n group=group,\n config=config,\n tags=tags, reinit=True,\n settings=wandb.Settings(start_method=\"fork\"))\n\n # data parameters\n v = 1 / config['data']['Re']\n t_interval = config['data']['time_interval']\n forcing_1 = get_forcing(S1).to(device)\n forcing_2 = get_forcing(S2).to(device)\n # training settings\n batch_size = config['train']['batchsize']\n ic_weight = config['train']['ic_loss']\n f_weight = config['train']['f_loss']\n xy_weight = config['train']['xy_loss']\n num_data_iter = config['train']['data_iter']\n num_eqn_iter = config['train']['eqn_iter']\n\n model.train()\n myloss = LpLoss(size_average=True)\n pbar = range(config['train']['epochs'])\n if use_tqdm:\n pbar = tqdm(pbar, dynamic_ncols=True, smoothing=0.05)\n zero = torch.zeros(1).to(device)\n train_loader = sample_data(train_loader)\n for ep in pbar:\n model.train()\n t1 = default_timer()\n train_loss = 0.0\n train_ic = 0.0\n train_f = 0.0\n test_l2 = 0.0\n err_eqn = 0.0\n # train with data\n for _ in range(num_data_iter):\n x, y = next(train_loader)\n x, y = x.to(device), y.to(device)\n optimizer.zero_grad()\n x_in = F.pad(x, (0, 0, 0, 5), \"constant\", 0)\n out = model(x_in).reshape(batch_size, S1, S1, T1 + 5)\n out = out[..., :-5]\n x = x[:, :, :, 0, -1]\n\n loss_l2 = myloss(out.view(batch_size, S1, S1, T1),\n y.view(batch_size, S1, S1, T1))\n\n if ic_weight != 0 or f_weight != 0:\n loss_ic, loss_f = PINO_loss3d(out.view(batch_size, S1, S1, T1),\n x, forcing_1,\n v, t_interval)\n else:\n loss_ic, loss_f = zero, zero\n\n total_loss = loss_l2 * xy_weight + loss_f * f_weight + loss_ic * ic_weight\n\n total_loss.backward()\n optimizer.step()\n\n train_ic = loss_ic.item()\n test_l2 += loss_l2.item()\n train_loss += total_loss.item()\n train_f += loss_f.item()\n if num_data_iter != 0:\n train_ic /= num_data_iter\n train_f /= num_data_iter\n train_loss /= num_data_iter\n test_l2 /= num_data_iter\n # train with random ICs\n for _ in range(num_eqn_iter):\n new_a = next(a_loader)\n new_a = new_a.to(device)\n optimizer.zero_grad()\n x_in = F.pad(new_a, (0, 0, 0, 5), \"constant\", 0)\n out = model(x_in).reshape(batch_size, S2, S2, T2 + 5)\n out = out[..., :-5]\n new_a = new_a[:, :, :, 0, -1]\n loss_ic, loss_f = PINO_loss3d(out.view(batch_size, S2, S2, T2),\n new_a, forcing_2,\n v, t_interval)\n eqn_loss = loss_f * f_weight + loss_ic * ic_weight\n eqn_loss.backward()\n optimizer.step()\n\n err_eqn += eqn_loss.item()\n\n scheduler.step()\n t2 = default_timer()\n if num_eqn_iter != 0:\n err_eqn /= num_eqn_iter\n if use_tqdm:\n pbar.set_description(\n (\n f'Data f error: {train_f:.5f}; Data ic l2 error: {train_ic:.5f}. '\n f'Data train loss: {train_loss:.5f}; Data l2 error: {test_l2:.5f}'\n f'Eqn loss: {err_eqn:.5f}'\n )\n )\n if wandb and log:\n wandb.log(\n {\n 'Data f error': train_f,\n 'Data IC L2 error': train_ic,\n 'Data train loss': train_loss,\n 'Data L2 error': test_l2,\n 'Random IC Train equation loss': err_eqn,\n 'Time cost': t2 - t1\n }\n )\n\n save_checkpoint(config['train']['save_dir'],\n config['train']['save_name'],\n model, optimizer)\n if wandb and log:\n run.finish()\n\n\ndef progressive_train(model,\n loader, train_loader,\n optimizer, scheduler,\n milestones, config,\n device=torch.device('cpu'),\n log=False,\n project='PINO-default',\n group='FDM',\n tags=['Nan'],\n use_tqdm=True):\n if wandb and log:\n run = wandb.init(project=project,\n entity=config['log']['entity'],\n group=group,\n config=config,\n tags=tags, reinit=True,\n settings=wandb.Settings(start_method=\"fork\"))\n\n # data parameters\n v = 1 / config['data']['Re']\n T = loader.T\n t_interval = config['data']['time_interval']\n\n # training settings\n batch_size = config['train']['batchsize']\n ic_weight = config['train']['ic_loss']\n f_weight = config['train']['f_loss']\n xy_weight = config['train']['xy_loss']\n\n model.train()\n myloss = LpLoss(size_average=True)\n zero = torch.zeros(1).to(device)\n for milestone, epochs in zip(milestones, config['train']['epochs']):\n pbar = range(epochs)\n if use_tqdm:\n pbar = tqdm(pbar, dynamic_ncols=True, smoothing=0.05)\n S = loader.S // milestone\n print(f'Resolution :{S}')\n forcing = get_forcing(S).to(device)\n for ep in pbar:\n model.train()\n t1 = default_timer()\n train_loss = 0.0\n train_ic = 0.0\n train_f = 0.0\n test_l2 = 0.0\n for x, y in train_loader:\n x, y = x.to(device), y.to(device)\n x = x[:, ::milestone, ::milestone, :, :]\n y = y[:, ::milestone, ::milestone, :]\n optimizer.zero_grad()\n x_in = F.pad(x, (0, 0, 0, 5), \"constant\", 0)\n out = model(x_in).reshape(batch_size, S, S, T + 5)\n out = out[..., :-5]\n x = x[:, :, :, 0, -1]\n\n loss_l2 = myloss(out.view(batch_size, S, S, T), y.view(batch_size, S, S, T))\n\n if ic_weight != 0 or f_weight != 0:\n loss_ic, loss_f = PINO_loss3d(out.view(batch_size, S, S, T),\n x, forcing, v, t_interval)\n else:\n loss_ic, loss_f = zero, zero\n\n total_loss = loss_l2 * xy_weight + loss_f * f_weight + loss_ic * ic_weight\n\n total_loss.backward()\n\n optimizer.step()\n train_ic = loss_ic.item()\n test_l2 += loss_l2.item()\n train_loss += total_loss.item()\n train_f += loss_f.item()\n scheduler.step()\n\n train_ic /= len(train_loader)\n train_f /= len(train_loader)\n train_loss /= len(train_loader)\n test_l2 /= len(train_loader)\n t2 = default_timer()\n if use_tqdm:\n pbar.set_description(\n (\n f'Train f error: {train_f:.5f}; Train ic l2 error: {train_ic:.5f}. '\n f'Train loss: {train_loss:.5f}; Test l2 error: {test_l2:.5f}'\n )\n )\n if wandb and log:\n wandb.log(\n {\n 'Train f error': train_f,\n 'Train L2 error': train_ic,\n 'Train loss': train_loss,\n 'Test L2 error': test_l2,\n 'Time cost': t2 - t1\n }\n )\n\n save_checkpoint(config['train']['save_dir'],\n config['train']['save_name'],\n model, optimizer)\n if wandb and log:\n run.finish()\n\n\n" }, { "path": "train_utils/utils.py", "content": "import os\nimport numpy as np\nimport torch\n\n\ndef vor2vel(w, L=2 * np.pi):\n '''\n Convert vorticity into velocity\n Args:\n w: vorticity with shape (batchsize, num_x, num_y, num_t)\n\n Returns:\n ux, uy with the same shape\n '''\n batchsize = w.size(0)\n nx = w.size(1)\n ny = w.size(2)\n nt = w.size(3)\n device = w.device\n w = w.reshape(batchsize, nx, ny, nt)\n\n w_h = torch.fft.fft2(w, dim=[1, 2])\n # Wavenumbers in y-direction\n k_max = nx // 2\n N = nx\n k_x = torch.cat((torch.arange(start=0, end=k_max, step=1, device=device),\n torch.arange(start=-k_max, end=0, step=1, device=device)), 0) \\\n .reshape(N, 1).repeat(1, N).reshape(1, N, N, 1)\n k_y = torch.cat((torch.arange(start=0, end=k_max, step=1, device=device),\n torch.arange(start=-k_max, end=0, step=1, device=device)), 0) \\\n .reshape(1, N).repeat(N, 1).reshape(1, N, N, 1)\n # Negative Laplacian in Fourier space\n lap = (k_x ** 2 + k_y ** 2)\n lap[0, 0, 0, 0] = 1.0\n f_h = w_h / lap\n\n ux_h = 2 * np.pi / L * 1j * k_y * f_h\n uy_h = -2 * np.pi / L * 1j * k_x * f_h\n\n ux = torch.fft.irfft2(ux_h[:, :, :k_max + 1], dim=[1, 2])\n uy = torch.fft.irfft2(uy_h[:, :, :k_max + 1], dim=[1, 2])\n return ux, uy\n\n\ndef get_sample(N, T, s, p, q):\n # sample p nodes from Initial Condition, p nodes from Boundary Condition, q nodes from Interior\n\n # sample IC\n index_ic = torch.randint(s, size=(N, p))\n sample_ic_t = torch.zeros(N, p)\n sample_ic_x = index_ic/s\n\n # sample BC\n sample_bc = torch.rand(size=(N, p//2))\n sample_bc_t = torch.cat([sample_bc, sample_bc],dim=1)\n sample_bc_x = torch.cat([torch.zeros(N, p//2), torch.ones(N, p//2)],dim=1)\n\n # sample I\n # sample_i_t = torch.rand(size=(N,q))\n # sample_i_t = torch.rand(size=(N,q))**2\n sample_i_t = -torch.cos(torch.rand(size=(N, q))*np.pi/2) + 1\n sample_i_x = torch.rand(size=(N,q))\n\n sample_t = torch.cat([sample_ic_t, sample_bc_t, sample_i_t], dim=1).cuda()\n sample_t.requires_grad = True\n sample_x = torch.cat([sample_ic_x, sample_bc_x, sample_i_x], dim=1).cuda()\n sample_x.requires_grad = True\n sample = torch.stack([sample_t, sample_x], dim=-1).reshape(N, (p+p+q), 2)\n return sample, sample_t, sample_x, index_ic.long()\n\n\ndef get_grid(N, T, s):\n gridt = torch.tensor(np.linspace(0, 1, T), dtype=torch.float).reshape(1, T, 1).repeat(N, 1, s).cuda()\n gridt.requires_grad = True\n gridx = torch.tensor(np.linspace(0, 1, s+1)[:-1], dtype=torch.float).reshape(1, 1, s).repeat(N, T, 1).cuda()\n gridx.requires_grad = True\n grid = torch.stack([gridt, gridx], dim=-1).reshape(N, T*s, 2)\n return grid, gridt, gridx\n\n\ndef get_2dgrid(S):\n '''\n get array of points on 2d grid in (0,1)^2\n Args:\n S: resolution\n\n Returns:\n points: flattened grid, ndarray (N, 2)\n '''\n xarr = np.linspace(0, 1, S)\n yarr = np.linspace(0, 1, S)\n xx, yy = np.meshgrid(xarr, yarr, indexing='ij')\n points = np.stack([xx.ravel(), yy.ravel()], axis=0).T\n return points\n\n\ndef torch2dgrid(num_x, num_y, bot=(0,0), top=(1,1)):\n x_bot, y_bot = bot\n x_top, y_top = top\n x_arr = torch.linspace(x_bot, x_top, steps=num_x)\n y_arr = torch.linspace(y_bot, y_top, steps=num_y)\n xx, yy = torch.meshgrid(x_arr, y_arr, indexing='ij')\n mesh = torch.stack([xx, yy], dim=2)\n return mesh\n\n\ndef get_grid3d(S, T, time_scale=1.0, device='cpu'):\n gridx = torch.tensor(np.linspace(0, 1, S + 1)[:-1], dtype=torch.float, device=device)\n gridx = gridx.reshape(1, S, 1, 1, 1).repeat([1, 1, S, T, 1])\n gridy = torch.tensor(np.linspace(0, 1, S + 1)[:-1], dtype=torch.float, device=device)\n gridy = gridy.reshape(1, 1, S, 1, 1).repeat([1, S, 1, T, 1])\n gridt = torch.tensor(np.linspace(0, 1 * time_scale, T), dtype=torch.float, device=device)\n gridt = gridt.reshape(1, 1, 1, T, 1).repeat([1, S, S, 1, 1])\n return gridx, gridy, gridt\n\n\ndef convert_ic(u0, N, S, T, time_scale=1.0):\n u0 = u0.reshape(N, S, S, 1, 1).repeat([1, 1, 1, T, 1])\n gridx, gridy, gridt = get_grid3d(S, T, time_scale=time_scale, device=u0.device)\n a_data = torch.cat((gridx.repeat([N, 1, 1, 1, 1]), gridy.repeat([N, 1, 1, 1, 1]),\n gridt.repeat([N, 1, 1, 1, 1]), u0), dim=-1)\n return a_data\n\n\n\ndef requires_grad(model, flag=True):\n for p in model.parameters():\n p.requires_grad = flag\n\n\ndef set_grad(tensors, flag=True):\n for p in tensors:\n p.requires_grad = flag\n\n\ndef zero_grad(params):\n '''\n set grad field to 0\n '''\n if isinstance(params, torch.Tensor):\n if params.grad is not None:\n params.grad.zero_()\n else:\n for p in params:\n if p.grad is not None:\n p.grad.zero_()\n\n\ndef count_params(net):\n count = 0\n for p in net.parameters():\n count += p.numel()\n return count\n\n\ndef save_checkpoint(path, name, model, optimizer=None):\n ckpt_dir = 'checkpoints/%s/' % path\n if not os.path.exists(ckpt_dir):\n os.makedirs(ckpt_dir)\n try:\n model_state_dict = model.module.state_dict()\n except AttributeError:\n model_state_dict = model.state_dict()\n\n if optimizer is not None:\n optim_dict = optimizer.state_dict()\n else:\n optim_dict = 0.0\n\n torch.save({\n 'model': model_state_dict,\n 'optim': optim_dict\n }, ckpt_dir + name)\n print('Checkpoint is saved at %s' % ckpt_dir + name)\n\n\n\ndef save_ckpt(path, model, optimizer=None, scheduler=None):\n model_state = model.state_dict()\n if optimizer:\n optim_state = optimizer.state_dict()\n else:\n optim_state = None\n \n if scheduler:\n scheduler_state = scheduler.state_dict()\n else:\n scheduler_state = None\n torch.save({\n 'model': model_state, \n 'optim': optim_state, \n 'scheduler': scheduler_state\n }, path)\n print(f'Checkpoint is saved to {path}')\n\n\ndef dict2str(log_dict):\n res = ''\n for key, value in log_dict.items():\n res += f'{key}: {value}|'\n return res" } ]