OneScience-Group/PINN-TC

PINN-TC

Model Introduction

PINN-TC reconstructs three-dimensional tropical-cyclone wind and pressure-related fields from sparse flight, dropsonde, or simulated observations, using physical equations to constrain unobserved regions for vortex initialization and data-assimilation research.

Paper: Realistic tropical cyclone wind and pressure fields can be reconstructed from sparse data using deep learning
https://doi.org/10.1038/s43247-023-01144-2

Model Description

The method was proposed by research teams at Princeton University, the California Institute of Technology, Stanford University, and NOAA's Geophysical Fluid Dynamics Laboratory. The paper trains and validates the method with Hurricane Ida T-SHiELD simulations and NOAA/AOML flight-level, dropsonde, and Tail Doppler Radar observations. It is intended for sparse tropical-cyclone wind and geopotential reconstruction, physics-consistent constraints, and vortex-initialization research.

Use Cases

Use Case	Description
Tropical-cyclone reconstruction	Reconstruct wind, geopotential, and pressure-vertical-velocity fields from sparse coordinate observations.
Physics-informed constraints	Constrain unobserved regions with horizontal momentum and pressure-continuity equations.
ModelScope/OneCode execution	Validate sparse sampling, training, full-grid inference, scientific metrics, and visualization in ModelScope or OneCode environments.
Multi-GPU training	Launch distributed data-parallel training through torchrun.

Usage Instructions

1.OneCode

Experience intelligent, one-click AI4S programming through the OneCode online environment:

Try intelligent, one-click AI4S programming

2. Download and Installation

hf download OneScience-Group/PINN-TC --local-dir ./PINN-TC
cd PINN-TC

Environment Dependencies

Hardware Requirements

• A GPU or DCU is recommended.
• A CPU can be used to validate the workflow with the default small-sample configuration.
• DCU users must install DTK in advance. DTK 25.04.2 or later, or the OneScience-recommended version matching the current cluster, is recommended.

DCU Environment

# Activate DTK and Conda first
conda create -n onescience311 python=3.11 -y
conda activate onescience311
pip install onescience[earth-dcu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai
pip install numpy pyyaml matplotlib

GPU Environment

# Activate Conda first
conda create -n onescience311 python=3.11 -y libstdcxx-ng=12 libgcc-ng=12 gcc_linux-64=12 gxx_linux-64=12
conda activate onescience311
pip install onescience[earth-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai
pip install numpy pyyaml matplotlib

Training Data

Training data consist of tropical-cyclone simulation fields and sparse observations containing storm-centered coordinates, horizontal winds, and geopotential. Model inputs are [y,x,t,p] coordinates, and targets are winds and geopotential at the same locations, while vertical velocity is constrained by the physical equations. The data cover a ±400 km horizontal domain, 150-900 hPa, and observation times at -3/0/+3 h. This repository uses a small number of synthetic samples to validate training, inference, and evaluation; they do not represent the official T-SHiELD or NOAA distribution, training scale, or formal paper performance.

python scripts/fake_data.py --force

Training

For single-device training, use:

python scripts/train.py

For multi-GPU training, use:

torchrun --nproc_per_node=8 --nnodes=1 --rdzv_id=1000 --rdzv_backend=c10d --max_restarts=0 --master_addr="localhost" --master_port=29500 scripts/train.py

Training jointly optimizes observational data loss and PDE residuals from the beta-plane horizontal momentum equations and the pressure-coordinate continuity equation. The checkpoint at result/checkpoints/pinn-tc.pt contains model state, model configuration, data format version, Adam/L-BFGS states, training history, and completed-step state.

Trained Weights

This repository does not bundle official paper weights, and no confirmed official checkpoint is available in the public resources. Running the training script produces a local engineering checkpoint, but no compatibility with external weights or equivalence to paper training results is claimed.

Inference

python scripts/inference.py

Inference loads the trained checkpoint and generates complete three-dimensional wind, geopotential, and pressure-vertical-velocity fields from sparse tropical-cyclone observation coordinates and physical constraints. Complete numerical results are saved to result/output/predictions.npz.

Evaluation and Visualization

python scripts/result.py

Evaluation computes wind-speed RMSE, Pearson correlation, PDE residuals, and radial RMSE as aggregates over the three observation times; it saves neither per-class metrics nor separate per-time metrics. It also generates a four-panel comparison of sparse observations, target wind speed, predicted wind speed, and error. Synthetic-data results validate the engineering workflow only and do not represent formal paper performance.

Official OneScience Information

Platform	OneScience Main Repository	Skills Repository
Gitee	https://gitee.com/onescience-ai/onescience	https://gitee.com/onescience-ai/oneskills
GitHub	https://github.com/onescience-ai/OneScience	https://github.com/onescience-ai/oneskills

Citation and License

This repository is an independent engineering reproduction of the public PINN-TC specifications.

Use of this repository's code, official model weights, and data remains subject to the licenses and terms of their respective projects.