Physics-Informed Neural Network (PINN) for 2D Helmholtz Equation

FEM Solver | PINN Curriculum | PINN Seq2seq

Overview

This project implements a Physics-Informed Neural Network (PINN) to solve the 2D Helmholtz equation using curriculum and seq2seq methods. These methods can be further detailed in [1]. We employ Finite Element Method (FEM) simulations for generating initial condition data and comparison results and utilize PINN to learn the wave field solution.

---

Problem Definition

The 2D homogeneous Helmholtz equation is given by:

$$ \Delta U + k^2 U = 0 $$

where:

• $$\ U(x, z) $$ represents the wave field.
• $$\ k^2 $$ is the wavenumber squared.
• Here we focus on Absorbing Boundary Conditions (ABC) which are applied to model realistic wave propagation scenarios.

---

Workflow

Step 1: Geometry Setup and FEM Solver (MATLAB)

• The problem domain and wave propagation setup are defined using MATLAB-based Finite Element Method (FEM) scripts.

• The primary solver is based on:

  David Gasperini (2025), "FEM Solver for 2D Helmholtz Equation", available at:
  MATLAB Repository

• Modifications:

  • The script Diffraction2.m was implemented to model wave propagation through a slit.
  • The computed wave field data (real and imaginary components) is stored in grid_data.mat.

---

Step 2: Data Processing (Python)

• The MATLAB output (grid_data.mat) is processed in Python using Notebook.ipynb.
• The dataset is converted into CSV format for compatibility with the PINN model.

Extracted Features:

• Either real or imaginary values of ( U(x, z) ) are extracted and saved as a CSV file.
• Boundary conditions at ( z=0 ) are stored separately in bc1.csv.

These files must be placed in the boundary_conditions directory within either the curriculum or seq2seq folders.

---

PINN Implementation

The PINN is implemented using PyTorch Lightning. There is a curriculum dir applying the curriculum method and seq2seq dir applying the seq2seq method.

Key Components

1. Configuration Files (config_curriculum.yaml, config_seq2seq.yaml):

   • Define hyperparameters such as:
     • Learning rate
     • Batch size
     • Number of epochs
     • Computational domain parameters (e.g., ( L_x, z ))
     • Training error thresholds

2. Dataset Handling (dataset_curriculum.py, dataset_seq2seq.py):

   • Loads training points from a specified computational domain.

3. Model Definition (model_curriculum.py, model_seq2seq.py):

   • Implements a fully connected neural network (FCNN) with Tanh activation functions.
   • Computes Helmholtz residual loss and boundary condition loss.

4. Training Scripts (train_curriculum.py, train_seq2seq.py):

   • Loads the dataset and trains the PINN using Adam optimizer with adaptive weight balancing.
   • Uses WandB (Weights & Biases) for logging.

5. Testing & Evaluation (test_curriculum.py, test_seq2seq.py):

   • Loads trained models for evaluation.
   • Compares PINN solutions with FEM reference solutions.
   • Computes L2 error and generates visualization plots.

6. Boundary Condition Generation (bc_creator.py):

   • Generates boundary conditions dynamically for sequential training.

---

Usage Guide

1. Preprocessing (MATLAB)

Run the MATLAB script to generate grid_data.mat.

2. Convert Data (Python)

Run the Jupyter notebook to extract baundary conditions:

jupyter notebook Notebook.ipynb

3. Train the Model

To train the PINN in curriculum mode for k^2=1:

python train_curriculum.py --k_squared 1

or for seq2seq training for the first sequence:

python bc_creator.py --seq_num 1 
python train_seq2seq.py --seq_num 1

4. Test the Model

Evaluate the trained model:

python test_curriculum.py --k_squared 1

or:

python test_seq2seq.py --seq_num 1

Running the entire training

For Curriculum Training running from k^2=1 to K:

for k in {1..K}; do \
    python "/gpfs0/tamyr/users/alonfi/Adrian project/curriculum/train_curriculum.py" --k_squared $k && \
    python "/gpfs0/tamyr/users/alonfi/Adrian project/curriculum/test_curriculum.py" --k_squared $k; \
    done'

For Seq2seq training from sequence 1 to S:

for s in {1..S}; do \
    python bc_creator.py --seq_num $s && \
    python train_seq2seq.py --seq_num $s && \
    python test_seq2seq.py --seq_num $s; \
    done'

---

Results & Visualization

• The trained model predictions are compared with FEM solutions.
• Checkpoint files are saved in the checkpoints
• Visualization plots are saved in the plots directory.
• L2 error between PINN and FEM solutions is computed.

---

---

References

[1] Krishnapriyan, A.S., Gholami, A., Zhe, S., Kirby, R.M., & Mahoney, M.W. (2021). Characterizing possible failure modes in physics-informed neural networks.

[2] Gasperini, D. (2025). FEM Solver for 2D Helmholtz Equation.

---