separable-operator-networks


Nameseparable-operator-networks JSON
Version 0.0.2 PyPI version JSON
download
home_pageNone
SummarySeparable operator models for extreme-scale learning of parametric PDEs
upload_time2024-11-01 23:54:41
maintainerNone
docs_urlNone
authorNone
requires_python>=3.8
licenseCopyright (c) 2024 Hewlett Packard Development Enterprise LP Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
keywords deeponet pde deep operator network operator learning partial differential equations phyics-informed neural network scientific machine learning
VCS
bugtrack_url
requirements No requirements were recorded.
Travis-CI No Travis.
coveralls test coverage No coveralls.
            # Separable Operator Networks (SepONet)
[![Static Badge](https://img.shields.io/badge/arXiv-2407.11253-blue?link=https%3A%2F%2Farxiv.org%2Fabs%2F2407.11253)](https://arxiv.org/abs/2407.11253)
[![Static Badge](https://img.shields.io/badge/pypi-v0.0.1-green?link=https%3A%2F%2Fpypi.org%2Fproject%2Fseparable-operator-networks%2F)](https://pypi.org/project/separable-operator-networks)

This is the official repository for separable operator networks (SepONet) originally introduced in [this preprint](https://arxiv.org/abs/2407.11253) [1]. 

## Installation
This code uses JAX as a dependency. It is recommended to [install with GPU/TPU compatibility](https://jax.readthedocs.io/en/latest/installation.html) prior to installing this library. JAX CPU is provided as the default dependency.

Please install with pip:
```bash
pip install separable-operator-networks
```
Alternatively, you may specify the `[cuda12]` extra to install `jax[cuda12]` automatically:
```bash
pip install separable-operator-networks[cuda12]
```

## Description

Operator learning has become a powerful tool in machine learning for modeling complex physical systems governed by partial differential equations (PDEs). Although Deep Operator Networks (DeepONet) show promise, they require extensive data acquisition. Physics-informed DeepONets (PI-DeepONet) mitigate data scarcity but suffer from inefficient training processes. We introduce Separable Operator Networks (SepONet), a novel framework that significantly enhances the efficiency of physics-informed operator learning. SepONet uses independent trunk networks to learn basis functions separately for different coordinate axes, enabling faster and more memory-efficient training via forward-mode automatic differentiation. The SepONet architecture for a $d=2$ dimensional coordinate grid is depicted below. The architecture is inspired by the method of separation of variables and recent exploration of separable physics-informed neural networks [2] for single instance PDE solutions.

Our [preprint](https://arxiv.org/abs/2407.11253) provides a universal approximation theorem for SepONet proving that it generalizes to arbitrary operator learning problems. For a variety of 1D time-dependent PDEs, SepONet has similar accuracy scaling to PI-DeepONet, but with as much as 112x faster training time and 82x reduction in GPU memory usage. For 2D time-dependent PDEs, SepONet is capable of accurate predictions at scales where PI-DeepONet fails. The full test scaling results as a function of the number of collocation points and number of input functions is shown below. These results may be reproduced using our [scripts](https://github.com/HewlettPackard/separable-operator-networks/tree/main/scripts).

![SepONet architecture for 2 dimensional coordinate grid](docs/assets/SepONet_Architecture.png?raw=true)
![Comparing SepONet to PI-DeepONet when varying number of collocation points](docs/assets/figure1_varying_Nc.png?raw=true)
![Comparing SepONet to PI-DeepONet when varying number of input functions](docs/assets/figure2_varying_Nf.png?raw=true)

## Code Overview

A SepONet model can be imported using:
```python
import jax
import separable_operator_networks as sepop
d = ... # replace with problem dimension
branch_dim = ... # replace with input shape for branch network (MLP by default)
key = jax.random.key(0)

model = sepop.models.SepONet(d, branch_dim, key=key)
```
Other model classes such as `PINN`, `SPINN`, `DeepONet` are implemented in the `sepop.models` submodule. These models are implemented as subclasses of `eqx.Module` (see [equinox](https://github.com/patrick-kidger/equinox)), enabling `eqx.filter_vmap` and `eqx.filter_grad`, along with easily customizable training routines via [optax](https://github.com/google-deepmind/optax) (see `sepop.train.train_loop(...)` for a simple `optax` training loop). PDE instances, loss functions, and other helper functions can be imported from the corresponding examples in the `sepop.pde` submodule (such as `sepop.pde.advection`).

Test data can be generated using the Python scripts in `/scripts/generate_test_data`. Test cases can be ran using the scripts in `/scripts/main_scripts` and `/scripts/scale_tests`.

## Citation

```tex
@misc{yu2024separableoperatornetworks,
title={Separable Operator Networks}, 
author={Xinling Yu and Sean Hooten and Ziyue Liu and Yequan Zhao and Marco Fiorentino and Thomas Van Vaerenbergh and Zheng Zhang},
year={2024},
eprint={2407.11253},
archivePrefix={arXiv},
primaryClass={cs.LG},
url={https://arxiv.org/abs/2407.11253}, 
}
```

## Authors

Sean Hooten (sean dot hooten at hpe dot com)  
Xinling Yu (xyu644 at ucsb dot edu)

## License

MIT (see LICENSE.md)

## References

[1] X. Yu, S. Hooten, Z. Liu, Y. Zhao, M. Fiorentino, T. Van Vaerenbergh, and Z. Zhang. Separable Operator Networks. arXiv preprint arXiv:2407.11253 (2024).  
[2] J. Cho, S. Nam, H. Yang, S.-B. Yun, Y. Hong, E. Park. Separable PINN: Mitigating the Curse of Dimensionality in Physics-Informed Neural Networks. arXiv preprint arXiv: 2211.08761 (2023).
            

Raw data

            {
    "_id": null,
    "home_page": null,
    "name": "separable-operator-networks",
    "maintainer": null,
    "docs_url": null,
    "requires_python": ">=3.8",
    "maintainer_email": "Sean Hooten <sean.hooten@hpe.com>, Xinling Yu <xyu644@ucsb.edu>",
    "keywords": "DeepONet, PDE, deep operator network, operator learning, partial differential equations, phyics-informed neural network, scientific machine learning",
    "author": null,
    "author_email": "Sean Hooten <sean.hooten@hpe.com>, Xinling Yu <xyu644@ucsb.edu>",
    "download_url": "https://files.pythonhosted.org/packages/de/ef/b808e739f042f50c24bba057aad21bed013ed0f13d1044a89f0ffe6abf3d/separable_operator_networks-0.0.2.tar.gz",
    "platform": null,
    "description": "# Separable Operator Networks (SepONet)\n[![Static Badge](https://img.shields.io/badge/arXiv-2407.11253-blue?link=https%3A%2F%2Farxiv.org%2Fabs%2F2407.11253)](https://arxiv.org/abs/2407.11253)\n[![Static Badge](https://img.shields.io/badge/pypi-v0.0.1-green?link=https%3A%2F%2Fpypi.org%2Fproject%2Fseparable-operator-networks%2F)](https://pypi.org/project/separable-operator-networks)\n\nThis is the official repository for separable operator networks (SepONet) originally introduced in [this preprint](https://arxiv.org/abs/2407.11253) [1]. \n\n## Installation\nThis code uses JAX as a dependency. It is recommended to [install with GPU/TPU compatibility](https://jax.readthedocs.io/en/latest/installation.html) prior to installing this library. JAX CPU is provided as the default dependency.\n\nPlease install with pip:\n```bash\npip install separable-operator-networks\n```\nAlternatively, you may specify the `[cuda12]` extra to install `jax[cuda12]` automatically:\n```bash\npip install separable-operator-networks[cuda12]\n```\n\n## Description\n\nOperator learning has become a powerful tool in machine learning for modeling complex physical systems governed by partial differential equations (PDEs). Although Deep Operator Networks (DeepONet) show promise, they require extensive data acquisition. Physics-informed DeepONets (PI-DeepONet) mitigate data scarcity but suffer from inefficient training processes. We introduce Separable Operator Networks (SepONet), a novel framework that significantly enhances the efficiency of physics-informed operator learning. SepONet uses independent trunk networks to learn basis functions separately for different coordinate axes, enabling faster and more memory-efficient training via forward-mode automatic differentiation. The SepONet architecture for a $d=2$ dimensional coordinate grid is depicted below. The architecture is inspired by the method of separation of variables and recent exploration of separable physics-informed neural networks [2] for single instance PDE solutions.\n\nOur [preprint](https://arxiv.org/abs/2407.11253) provides a universal approximation theorem for SepONet proving that it generalizes to arbitrary operator learning problems. For a variety of 1D time-dependent PDEs, SepONet has similar accuracy scaling to PI-DeepONet, but with as much as 112x faster training time and 82x reduction in GPU memory usage. For 2D time-dependent PDEs, SepONet is capable of accurate predictions at scales where PI-DeepONet fails. The full test scaling results as a function of the number of collocation points and number of input functions is shown below. These results may be reproduced using our [scripts](https://github.com/HewlettPackard/separable-operator-networks/tree/main/scripts).\n\n![SepONet architecture for 2 dimensional coordinate grid](docs/assets/SepONet_Architecture.png?raw=true)\n![Comparing SepONet to PI-DeepONet when varying number of collocation points](docs/assets/figure1_varying_Nc.png?raw=true)\n![Comparing SepONet to PI-DeepONet when varying number of input functions](docs/assets/figure2_varying_Nf.png?raw=true)\n\n## Code Overview\n\nA SepONet model can be imported using:\n```python\nimport jax\nimport separable_operator_networks as sepop\nd = ... # replace with problem dimension\nbranch_dim = ... # replace with input shape for branch network (MLP by default)\nkey = jax.random.key(0)\n\nmodel = sepop.models.SepONet(d, branch_dim, key=key)\n```\nOther model classes such as `PINN`, `SPINN`, `DeepONet` are implemented in the `sepop.models` submodule. These models are implemented as subclasses of `eqx.Module` (see [equinox](https://github.com/patrick-kidger/equinox)), enabling `eqx.filter_vmap` and `eqx.filter_grad`, along with easily customizable training routines via [optax](https://github.com/google-deepmind/optax) (see `sepop.train.train_loop(...)` for a simple `optax` training loop). PDE instances, loss functions, and other helper functions can be imported from the corresponding examples in the `sepop.pde` submodule (such as `sepop.pde.advection`).\n\nTest data can be generated using the Python scripts in `/scripts/generate_test_data`. Test cases can be ran using the scripts in `/scripts/main_scripts` and `/scripts/scale_tests`.\n\n## Citation\n\n```tex\n@misc{yu2024separableoperatornetworks,\ntitle={Separable Operator Networks}, \nauthor={Xinling Yu and Sean Hooten and Ziyue Liu and Yequan Zhao and Marco Fiorentino and Thomas Van Vaerenbergh and Zheng Zhang},\nyear={2024},\neprint={2407.11253},\narchivePrefix={arXiv},\nprimaryClass={cs.LG},\nurl={https://arxiv.org/abs/2407.11253}, \n}\n```\n\n## Authors\n\nSean Hooten (sean dot hooten at hpe dot com)  \nXinling Yu (xyu644 at ucsb dot edu)\n\n## License\n\nMIT (see LICENSE.md)\n\n## References\n\n[1] X. Yu, S. Hooten, Z. Liu, Y. Zhao, M. Fiorentino, T. Van Vaerenbergh, and Z. Zhang. Separable Operator Networks. arXiv preprint arXiv:2407.11253 (2024).  \n[2] J. Cho, S. Nam, H. Yang, S.-B. Yun, Y. Hong, E. Park. Separable PINN: Mitigating the Curse of Dimensionality in Physics-Informed Neural Networks. arXiv preprint arXiv: 2211.08761 (2023).",
    "bugtrack_url": null,
    "license": "Copyright (c) 2024 Hewlett Packard Development Enterprise LP  Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the \"Software\"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:  The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.  THE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.",
    "summary": "Separable operator models for extreme-scale learning of parametric PDEs",
    "version": "0.0.2",
    "project_urls": {
        "Homepage": "https://github.com/hewlettpackard/separable-operator-networks",
        "Issues": "https://github.com/hewlettpackard/separable-operator-networks/issues",
        "Repository": "https://github.com/hewlettpackard/separable-operator-networks.git"
    },
    "split_keywords": [
        "deeponet",
        " pde",
        " deep operator network",
        " operator learning",
        " partial differential equations",
        " phyics-informed neural network",
        " scientific machine learning"
    ],
    "urls": [
        {
            "comment_text": "",
            "digests": {
                "blake2b_256": "20f2891b29e523780491d60e950cf7ab4826ab11599e8ed27e7b0e74c444a689",
                "md5": "dc1e1edf3ca5932b16c6c66ee0cd1181",
                "sha256": "01b00d6fc8541173b98c8d548e0684f6457bd3139da1af7b84daae3eaee0e9b5"
            },
            "downloads": -1,
            "filename": "separable_operator_networks-0.0.2-py3-none-any.whl",
            "has_sig": false,
            "md5_digest": "dc1e1edf3ca5932b16c6c66ee0cd1181",
            "packagetype": "bdist_wheel",
            "python_version": "py3",
            "requires_python": ">=3.8",
            "size": 24676,
            "upload_time": "2024-11-01T23:54:40",
            "upload_time_iso_8601": "2024-11-01T23:54:40.355677Z",
            "url": "https://files.pythonhosted.org/packages/20/f2/891b29e523780491d60e950cf7ab4826ab11599e8ed27e7b0e74c444a689/separable_operator_networks-0.0.2-py3-none-any.whl",
            "yanked": false,
            "yanked_reason": null
        },
        {
            "comment_text": "",
            "digests": {
                "blake2b_256": "deefb808e739f042f50c24bba057aad21bed013ed0f13d1044a89f0ffe6abf3d",
                "md5": "57d60ce1242ecac2feb8a71421d33a2a",
                "sha256": "83a097fe694f2f70db67de84fcaa0e134e88d04d5707524a2ebf75b15ef5d6fa"
            },
            "downloads": -1,
            "filename": "separable_operator_networks-0.0.2.tar.gz",
            "has_sig": false,
            "md5_digest": "57d60ce1242ecac2feb8a71421d33a2a",
            "packagetype": "sdist",
            "python_version": "source",
            "requires_python": ">=3.8",
            "size": 15794,
            "upload_time": "2024-11-01T23:54:41",
            "upload_time_iso_8601": "2024-11-01T23:54:41.739507Z",
            "url": "https://files.pythonhosted.org/packages/de/ef/b808e739f042f50c24bba057aad21bed013ed0f13d1044a89f0ffe6abf3d/separable_operator_networks-0.0.2.tar.gz",
            "yanked": false,
            "yanked_reason": null
        }
    ],
    "upload_time": "2024-11-01 23:54:41",
    "github": true,
    "gitlab": false,
    "bitbucket": false,
    "codeberg": false,
    "github_user": "hewlettpackard",
    "github_project": "separable-operator-networks",
    "travis_ci": false,
    "coveralls": false,
    "github_actions": false,
    "lcname": "separable-operator-networks"
}
        
Elapsed time: 0.51879s