PFEM

The website of PFEM is : https://pretrainfem.streamlit.app/. If it is helpful, you could star the repo :)

This repository contains finite element surrogate modeling experiments based on physics-informed neural operators (PINOs) and Transolver. It focuses on displacement-field prediction for two-dimensional elasticity and hyperelasticity problems, with support for unstructured meshes, varying material properties, random boundary loads, and accelerated finite element iterations through warm starts.

PFEM: physics-informed neural operator pretraining followed by warm-started conventional numerical solvers

Paper

Pretrain finite element method: A pretraining and warm-start framework for PDEs via physics-informed neural operators

Yizheng Wang, Zhongkai Hao, Mohammad Sadegh Eshaghi, Cosmin Anitescu, Xiaoying Zhuang, Timon Rabczuk, Yinghua Liu. Journal of the Mechanics and Physics of Solids, 214 (2026), 106682.

• Paper PDF in this repository
• ScienceDirect article
• DOI: 10.1016/j.jmps.2026.106682

Method Overview

PFEM consists of two stages:

1. Physics-informed pretraining: Transolver operates directly on unstructured point clouds and jointly encodes coordinates, geometry, material properties, and boundary conditions. Explicit differentiation based on finite element shape functions enforces the governing equations without requiring labeled solutions.
2. Warm-start solution: The low-fidelity prediction produced by the pretrained model is used as the initial solution for conventional numerical methods such as FEM, reducing the number of iterations while retaining high solution accuracy.

Automatic differentiation versus explicit differentiation based on finite element shape functions

Benchmarks

Elastic plates with random holes	Hyperelastic beam and Cook's membrane
TPMS: 3D Solid	TPMS: 3D Sheet

• Elastic plates with random holes
  • Geometry variation: practical_problems/Elasticty_holes/geo
  • Geometry and material variation: practical_problems/Elasticty_holes/geo_material
  • Geometry, material, and boundary-load variation: practical_problems/Elasticty_holes/geo_material_boundary
• Hyperelastic beam: practical_problems/Hyper_beam
• Cook's membrane: practical_problems/Hyper_cook
• Complex 3D TPMS structures: displacement-field prediction for 3D solid and sheet structures

Results

Random Geometry, Materials, and Boundary Loads

The pretrained model predicts displacement fields for plates with random holes, spatially varying materials, and random boundary loads. Using its prediction as the initial solution, PFEM substantially reduces the number of iterations compared with FEM initialized from zero.

Hyperelastic Cook's Membrane

For the nonlinear Cook's membrane benchmark, the PFEM initial solution accelerates Newton iterations while preserving the final accuracy of the conventional finite element solver.

3D TPMS Solid and Sheet Structures

3D solid structures: reference displacement, PFEM prediction, absolute error, and the error matrix across structural classes

3D sheet structures: reference displacement, PFEM prediction, absolute error, and the error matrix across structural classes

PFEM operates directly on complex three-dimensional point clouds, learns displacement solution operators across solid and sheet structures with different topologies, and maintains low prediction errors.

Click any image in this README to view its original PDF.

Repository Structure

PFEM/
├── data/
│   ├── elasticity/          # Elastic plate data generation and processing
│   └── hyper/               # Hyperelastic beam and Cook's membrane data
├── model/                   # Transolver, Physics Attention, and FNO
├── practical_problems/      # Training, testing, fine-tuning, and visualization
├── pic/                     # Method diagrams and experimental results
├── scripts/                 # Slurm job scripts
├── utils/                   # Normalization and error utilities
├── model_dict.py            # Model registry
├── PFEM.pdf                 # Paper PDF
└── README.md

Installation

Python 3.10 or 3.11 is recommended. Install the PyTorch build that matches your CUDA environment.

conda create -n pfem python=3.11 -y
conda activate pfem

# Select the appropriate PyTorch build at:
# https://pytorch.org/get-started/locally/
pip install torch
pip install numpy scipy matplotlib h5py pandas seaborn tqdm einops timm

CPU execution is supported by adding --cpu to a training command. Full-scale training is computationally intensive, so a CUDA-enabled GPU is recommended.

Data Preparation

Option 1: Download the Data

Pre-generated datasets are available from Tsinghua Cloud:

https://cloud.tsinghua.edu.cn/d/8563cdebd4754aa58fe8/

Extract the downloaded data into data/ while preserving the directory structure in the archive. The training scripts use the following paths by default:

data/
├── elasticity/
│   ├── dataset.npz
│   ├── dataset_with_node_materials.npz
│   └── dataset_with_node_materials_and_random_force.npz
└── hyper/
    ├── training_data_beam_q4/
    │   └── hyperelastic_training_data_q4.npz
    └── training_data_cook_q8/
        └── hyperelastic_training_data_q8.npz

Each experiment only requires its corresponding dataset. If the extracted files are stored elsewhere, pass their locations through the training script's --path argument.

Option 2: Generate the Data

Data generation uses SciPy-based finite element solvers. The default sample counts and mesh resolutions are large. For an initial test, reduce Nsamples, Nx, and Ny near the bottom of the relevant script.

Elastic Plate

cd data/elasticity

# Random holes with fixed material properties and loads
python data_generate.py

# Random holes and spatially varying materials
python data_generate_E_nu.py

# Random holes, materials, and boundary loads
python data_generate_E_nu_boundary.py

cd ../..

The scripts generate dataset.npz, dataset_with_node_materials.npz, and dataset_with_node_materials_and_random_force.npz, respectively.

Hyperelastic Beam

Generate an HDF5 dataset first, then convert it to the NPZ format used for training:

cd data/hyper
python data_generate_beam.py --num_index 1

# The conversion script currently reads physics_training_data_beam
mv physics_training_data_1 physics_training_data_beam
python data_generate_beam_convert_quad.py
cd ../..

The converted dataset is saved to:

data/hyper/training_data_beam_q4/hyperelastic_training_data_q4.npz

Cook's Membrane

The following commands generate a small example dataset and convert it to the training format:

cd data/hyper
python data_generate_cook.py \
  --out cook_random_dataset \
  --num_samples 20 \
  --Nx_elem 10 \
  --Ny_elem 10

python data_generate_cook_convert_quad.py \
  --data_dir cook_random_dataset \
  --out training_data_cook_q8/hyperelastic_training_data_q8.npz
cd ../..

Quick Start

Run all training and evaluation commands from the PFEM/ project root. Copy the contents of practical_problems/ into the PFEM root directory before running the commands below.

1. Elastic Plate with Random Holes

Start with a small run to verify the environment and dataset:

python -m Elasticty_holes.geo.exp_elas_PINO_transolver \
  --path ./data/elasticity/dataset.npz \
  --ntrain 10 \
  --ntest 2 \
  --epochs 5 \
  --out_dir ./results_pino_transolver_debug

Use the default parameters for full training:

python -m Elasticty_holes.geo.exp_elas_PINO_transolver

Run the material and boundary-load experiments with:

python -m Elasticty_holes.geo_material.exp_elas_PINO_transolver_g_m

python -m Elasticty_holes.geo_material_boundary.exp_elas_PINO_transolver_all

2. Hyperelastic Beam

python -m Hyper_beam.exp_hyper_PINO_transolver_quad \
  --path ./data/hyper/training_data_beam_q4/hyperelastic_training_data_q4.npz \
  --ntrain 10 \
  --ntest 2 \
  --epochs 5 \
  --out_dir ./results_hyper_beam_debug

3. Cook's Membrane

python -m Hyper_cook.exp_cook_PINO_transolver_quad \
  --path ./data/hyper/training_data_cook_q8/hyperelastic_training_data_q8.npz \
  --ntrain 10 \
  --ntest 2 \
  --epochs 5 \
  --out_dir ./results_hyper_cook_debug

If the dataset contains fewer than ntrain + ntest samples, reduce these two arguments accordingly.

Outputs and Visualization

Training outputs are saved to the directory specified by --out_dir and typically include:

• model.pt or model_final.pt: trained model parameters
• model_epoch*.pt: intermediate checkpoints
• metrics_history.json: training and testing metrics
• Mesh, displacement, error-field, and training-curve figures

For example, visualize predictions from a trained elastic plate model with:

python -m Elasticty_holes.geo.PFEM_contourf_prediction \
  --path ./data/elasticity/dataset.npz \
  --model_path ./results_pino_transolver/model.pt \
  --out_dir ./test_results_contourf

Use --help to inspect all arguments supported by a script:

python -m Elasticty_holes.geo.exp_elas_PINO_transolver --help

Slurm

The scripts/ directory contains example Slurm job scripts. Update the partition, GPU resources, log paths, and Python entry point to match your cluster before submission:

sbatch scripts/b_transolver.python

Notes

• Some .npz datasets contain NumPy object arrays and must be loaded with allow_pickle=True; the included scripts already do this.
• Generating a complete dataset can require substantial runtime and memory. Validate the workflow using a small number of samples first.
• Model architecture settings such as n_hidden, n_layers, and slice_num must match the checkpoint being loaded.
• This repository primarily contains research code. Some data-generation parameters are defined directly near the bottom of each script and should be adjusted before large-scale generation.