Metadata-Version: 2.4
Name: folax
Version: 0.0.2
Summary: A framework for solving and optimizing PDEs by integrating machine learning with numerical methods in computational mechanics.
Author: Reza Najian Asl, email = <reza.najian-asl@tum.de>
License: BSD 4-Clause License
        
        Copyright (c) 2024, Reza Najian Asl, Shahed Rezaei
        All rights reserved.
        
        Redistribution and use in source and binary forms, with or without
        modification, are permitted provided that the following conditions are met:
        
        1. Redistributions of source code must retain the above copyright notice, this
           list of conditions and the following disclaimer.
        
        2. Redistributions in binary form must reproduce the above copyright notice,
           this list of conditions and the following disclaimer in the documentation
           and/or other materials provided with the distribution.
        
        3. All advertising materials mentioning features or use of this software must
           display the following acknowledgement:
             This product includes software developed by Reza Najian Asl, Shahed Rezaei.
        
        4. Neither the name of the copyright holder nor the names of its
           contributors may be used to endorse or promote products derived from
           this software without specific prior written permission.
        
        THIS SOFTWARE IS PROVIDED BY COPYRIGHT HOLDER "AS IS" AND ANY EXPRESS OR
        IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
        MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO
        EVENT SHALL COPYRIGHT HOLDER BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
        SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
        PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
        OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
        WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR
        OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
        ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
License-File: LICENSE
Keywords: Numerical-Simulation-Optimization,PDE-Solver,Physics-Informed-Neural-Networks (PINNs),Scientific Machine Learning (SciML)
Requires-Python: >=3.10
Requires-Dist: dask
Requires-Dist: flax~=0.10.7
Requires-Dist: gmsh
Requires-Dist: h5py
Requires-Dist: matplotlib
Requires-Dist: meshio
Requires-Dist: numpy
Requires-Dist: optax
Requires-Dist: scipy
Requires-Dist: tqdm
Requires-Dist: zarr
Provides-Extra: cpu
Requires-Dist: jax~=0.6.2; extra == 'cpu'
Provides-Extra: cuda
Requires-Dist: jax[cuda12]~=0.6.2; extra == 'cuda'
Provides-Extra: dev
Requires-Dist: pytest; extra == 'dev'
Requires-Dist: pytest-cov; extra == 'dev'
Requires-Dist: pytest-xdist; extra == 'dev'
Description-Content-Type: text/markdown

<p align=center><img height="54.125%" width="54.125%" src="https://github.com/RezaNajian/eFOL/assets/62375973/0e1ca4e0-0658-4f5d-aad9-1ae7c9f67574"></p>

[![License][license-image]][license] 
[![CI](https://github.com/RezaNajian/folax/actions/workflows/CI.yml/badge.svg)](https://github.com/RezaNajian/folax/actions/workflows/CI.yml)
[![PyPI version](https://badge.fury.io/py/folax.svg)](https://pypi.org/project/folax/)


[license-image]: https://img.shields.io/badge/license-BSD-green.svg?style=flat
[license]: https://github.com/RezaNajian/FOL/LICENSE

# Folax: Solution and Optimization of parameterized PDEs
**F**inite **O**perator **L**earning (FOL) with [**JAX**](https://github.com/jax-ml/jax) constitutes a unified numerical framework that seamlessly integrates established numerical methods with advanced scientific machine learning techniques for solving and optimizing parametrized partial differential equations (PDEs).  In constructing a physics-informed operator learning approach, FOL formulates a purely physics-based loss function derived from the Method of Weighted Residuals, allowing discrete residuals—computed using classical PDE solution techniques—to be directly incorporated into backpropagation during network training. This approach ensures that the learned operators rigorously satisfy the underlying governing equations while maintaining consistency with established numerical discretizations. Importantly, this loss formulation is agnostic to the network architecture and has been successfully applied to architectures such as Conditional Neural Fields, Fourier Neural Operators (FNO), and DeepONets. 

FOL has been applied in the following scientific studies:
- A Physics-Informed Meta-Learning Framework for the Continuous Solution of Parametric PDEs on Arbitrary Geometries [[arXiv](https://arxiv.org/abs/2504.02459)].
- Finite Operator Learning: Bridging Neural Operators and Numerical Methods for Efficient Parametric Solution and Optimization of PDEs [[arXiv](https://arxiv.org/abs/2407.04157)].
- Digitalizing metallic materials from image segmentation to multiscale solutions via physics informed operator learning [[npj Computational Materials](https://www.nature.com/articles/s41524-025-01718-y)].
- A Finite Operator Learning Technique for Mapping the Elastic Properties of Microstructures to Their Mechanical Deformations [[Numerical Methods in Eng.](https://onlinelibrary.wiley.com/doi/full/10.1002/nme.7637)].
- SPiFOL: A Spectral-based physics-informed finite operator learning for prediction of mechanical behavior of microstructures [[J. Mechanics and Physics of Solids](https://www.sciencedirect.com/science/article/pii/S0022509625001954)].

We built upon several widely adopted Python packages, including [JAX](https://github.com/jax-ml/jax) for high-performance array computations on CPUs and GPUs, [PETSc](https://petsc.org/release/) for the efficient solution of large-scale linear systems, [Metis](https://github.com/KarypisLab/METIS) for mesh partitioning (integration forthcoming), [Flax](https://github.com/google/flax?tab=readme-ov-file) for constructing modular and flexible neural networks, [Optax](https://github.com/google-deepmind/optax) for applying state-of-the-art gradient-based optimization algorithms, and [Orbax](https://github.com/google/orbax) for efficient checkpointing and serialization. This foundation ensures scalability, computational efficiency, and ease of use in large-scale training and simulation workflows.

## Installation
### CPU installation 
To install folax using pip (recommended) for CPU usage you can type the following command

``pip install folax[cpu]``

### GPU installation
To install folax using pip (recommended) for GPU usage you can type the following command

``pip install folax[cuda]``

### Developer installation
If you would like to do development in folax, please first clone the repo and in the folax folder, run the following command

``pip install -e .[cuda,dev]``

## Contributing
If you would like to contribute to the project, please open a pull request with small changes. If you would like to see big changes in the source code, please open an issue or discussion so we can start a conversation.

