Accelerating 3D Photoacoustic Computed Tomography with End-to-End Physics-Aware Neural Operators

Overview

This repository contains the official implementation of PANO proposed in "LAccelerating 3D Photoacoustic Computed Tomography with End-to-End Physics-Aware Neural Operators", an end-to-edn approach for accelerating 3D Photoacoustic Computed Tomography (PACT) reconstruction using physics-aware neural operators. Our method achieves significant speedup while maintaining high reconstruction quality compared to traditional iterative methods.

Installation

Prerequisites

• Python 3.8+
• MATLAB R2020a+ (for data generation and baseline methods only)
• CUDA-compatible GPU (recommended)

Python Dependencies

Install the required Python packages:

pip install torch torchvision torchaudio
pip install numpy scipy matplotlib
pip install opencv-python-headless
pip install h5py scikit-image
pip install wandb
pip install neuralop
pip install triton
pip install pyyaml
pip install tqdm

Or install from the requirements file:

pip install -r requirements.txt

MATLAB Dependencies

For the baseline reconstruction methods and data generation:

1. VascuSynth: Clone and install VascuSynth following their installation instructions
2. MATLAB Toolboxes: The code requires the following MATLAB toolboxes:
   • Parallel Computing Toolbox (for GPU acceleration)
   • Image Processing Toolbox
   • Signal Processing Toolbox

Repository Structure

pact3d/
├── data_gen/                    # Synthetic 3D vessel data generation
│   ├── VascuSynth.m            # Main data generation script
│   └── README.md               # Data generation instructions
├── data_baseline/              # Baseline reconstruction methods
│   ├── sim/                    # Simulation and baseline methods
│   │   ├── Generate_Training_Data_NeuroOp_1k-test.m
│   │   ├── run_sim.sh
│   │   └── ImageRec.p          # Compiled reconstruction class
│   └── real/                   # Real data processing
│       ├── Generate_Training_Data_NeuroOp_1k_Recon.m
│       ├── run_real.sh
│       └── ImageRec.p
└── pano/                       # PANO neural operator implementation
    ├── main.py                 # Main training/inference script
    ├── models.py               # Neural operator architectures
    ├── unet3d/                 # 3D U-Net components
    ├── torch_harmonics_local/  # Custom spherical harmonics implementation
    └── pytorch_ssim/           # SSIM loss implementation

Quick Start

1. Data Generation and Simulation

1. Generate synthetic data:

   cd data_gen
   # Follow instructions in data_gen/README.md
   # This requires VascuSynth installation

2. Run simulation baseline:

   cd data_baseline/sim
   ./run_sim.sh 1  # Process data point 1

3. Run real data processing:

   cd data_baseline/real
   ./run_real.sh 1  # Process data point 1

2. Train PANO Model

cd pano
python main.py

3. Run Inference

cd pano
# Set EVALUATE=True in main.py or modify the script
python main.py

Usage

Training

The main training script pano/main.py supports various configuration options:

• Batch size: Modify batch_size in the script
• Learning rate: Adjust learning_rate parameter
• Epochs: Set epochs for training duration
• Model architecture: Configure modes and width parameters

Configuration

Key parameters can be modified in pano/cfg_file_morepino.py:

config_file = {
    'opt': {
        'epoch': 80+45,
        'btz': 2,           # Batch size
        'lr': 1e-3,         # Learning rate
        'l1w': 0.7,         # L1 loss weight
    },
    'dsp': {
        'dataset': [43,3],  # Dataset configuration
        'sensor': 30,       # Number of sensor points
        'mode': 15,         # Number of modes
    }
}

Model Architecture

PANO uses a novel neural operator architecture featuring:

• DISCO layers for local feature learning on spherical coordinates
• FNO (Fourier Neural Operator) for global feature learning
• 3D U-Net components for spatial feature and final image refinement
• Physics-aware loss functions incorporating wave propagation in the forward processing (traininig only, optional)

Evaluation

The model is evaluated using multiple metrics:

• PSNR (Peak Signal-to-Noise Ratio)
• SSIM (Structural Similarity Index)
• NSME (Normalized Squared Mean Error)

Citation

If you find this work useful, please cite our paper:

@article{wang2025accelerating,
  title={Accelerating 3D Photoacoustic Computed Tomography with End-to-End Physics-Aware Neural Operators},
  author={Wang, Jiayun and Aborahama, Yousuf and Khokhar, Arya and Zhang, Yang and Wang, Chuwei and Sastry, Karteekeya and Berner, Julius and Luo, Yilin and Bonev, Boris and Li, Zongyi and others},
  journal={arXiv preprint arXiv:2509.09894},
  year={2025}
}

License

This project is licensed under the MIT License - see the LICENSE file for details.

Acknowledgments

• VascuSynth for synthetic vessel data generation
• NeuralOperators library for FNO implementations
• PyTorch3DUNet for 3D U-Net components