Configuration

NeuralPOM uses JSON configuration files for both training and inference. All parameters are specified in a single JSON file, making experiments self-documenting and reproducible.

Training Configuration

Example: configs/train_decoder_only.json

json

{
  "lr": 0.001,
  "epochs": 100,
  "start_day": 50,
  "end_day": 800,
  "test_start_day": 700,
  "test_end_day": 800,
  "batch_size": 1,
  "sample_freq": 1,
  "train_sample_freq": 1,
  "val_sample_freq": 1,
  "target_idx": 0,
  "reuse_core_pool": true,
  "use_encoder": false,
  "use_decoder": true,
  "use_pretrain": false,
  "grad_clip_norm": 1.0,
  "core_name": "implicit",
  "core_im": 90,
  "core_jm": 82,
  "correction_scale": 0.001,
  "mean_path": "/path/to/mean_field/checkpoints/initial_state.npz",
  "supervision_mode": "pointwise_baseline",
  "multiscale_sigmas": [1.0, 2.0],
  "multiscale_weights": [0.5, 1.0],
  "multiscale_eps": 1e-6,
  "horizon_decay": 0.8,
  "loss_display_scale": 100000.0,
  "loss_weights": {
    "u": 12.0, "ub": 18.0, "v": 10.0, "vb": 12.0,
    "t": 1.0, "s": 6.0, "el": 30.0,
    "ua": 60.0, "uab": 60.0, "va": 60.0, "vab": 60.0
  }
}

Core Parameters

Parameter	Type	Description
`core_name`	str	Dynamic core: `"explicit"` or `"implicit"`
`core_im`	int	Grid size in the i-direction (east-west)
`core_jm`	int	Grid size in the j-direction (north-south)
`correction_scale`	float	Scaling factor for CorrectorNet output applied to physics state

Training Parameters

Parameter	Type	Description
`lr`	float	Learning rate for Adam optimizer
`epochs`	int	Number of training epochs
`batch_size`	int	Per-GPU batch size (total = batch_size × world_size)
`grad_clip_norm`	float	Gradient clipping norm threshold
`loss_display_scale`	float	Global scaling factor for displayed loss values

Data Parameters

Parameter	Type	Description
`start_day`	int	First day of training data (simulation day index)
`end_day`	int	Last day of training data
`test_start_day`	int	First day of validation data
`test_end_day`	int	Last day of validation data
`target_idx`	int	Target segment index in the NPZ data files
`train_sample_freq`	int	Sampling frequency for training (1 = every day)
`val_sample_freq`	int	Sampling frequency for validation
`mean_path`	str	Path to NPZ file with mean fields for normalization

Supervision Parameters

Parameter	Type	Description
`supervision_mode`	str	Loss mode: `"pointwise_baseline"`, `"multiscale_lowpass"`, or `"horizon_decay"`
`multiscale_sigmas`	list[float]	Gaussian filter sigmas for multiscale supervision
`multiscale_weights`	list[float]	Weights for each scale in multiscale loss
`multiscale_eps`	float	Epsilon for numerical stability in multiscale normalization
`horizon_decay`	float	Decay factor ( $γ \in (0, 1]$ ) for horizon-weighted supervision

Network Parameters

Parameter	Type	Description
`use_encoder`	bool	Enable encoder CorrectorNet (pre-physics)
`use_decoder`	bool	Enable decoder CorrectorNet (post-physics)
`use_pretrain`	bool	Load pretrained weights
`reuse_core_pool`	bool	Reuse a pool of core instances across batches (memory optimization)

Loss Weights

Loss weights balance the 12 prognostic fields. Higher weights emphasize fields with smaller magnitude or higher dynamical importance:

L = w_{u} L_{u} + w_{u_{b}} L_{u_{b}} + w_{v} L_{v} + w_{v_{b}} L_{v_{b}} + w_{T} L_{T} + w_{S} L_{S} + w_{η} L_{η} + w_{u_{a}} L_{u_{a}} + w_{u_{a b}} L_{u_{a b}} + w_{v_{a}} L_{v_{a}} + w_{v_{a b}} L_{v_{a b}}

For an explanation of each field, see CorrectorNet Architecture.

Inference Configuration

Example: configs/inference_default.json

json

{
  "model_path": "./checkpoints_model/corrector_best.pth",
  "mean_path": "/path/to/mean_field/checkpoints/initial_state.npz",
  "start_day": 700,
  "end_day": 701,
  "target_idx": 0,
  "use_encoder": true,
  "use_decoder": true,
  "core_name": "implicit",
  "core_im": 90,
  "core_jm": 82,
  "correction_scale": 0.001,
  "output_dir": "inference_results"
}

Parameter	Description
`model_path`	Path to the trained CorrectorNet checkpoint (`.pth`)
`start_day` / `end_day`	Simulation day range for inference
`output_dir`	Directory for saving inference outputs
`use_encoder` / `use_decoder`	Enable/disable encoder and decoder correction stages

Overriding Config via Command Line

Configuration values can be overridden from the command line:

bash

torchrun --nproc_per_node=8 scripts/training/train_npom.py \
  --config configs/train_decoder_only.json \
  --lr 0.0005 \
  --epochs 200 \
  --supervision_mode multiscale_lowpass

Any JSON key can be passed as --key value to override the file-based configuration.

Configuration ​

Training Configuration ​

Core Parameters ​

Training Parameters ​

Data Parameters ​

Supervision Parameters ​

Network Parameters ​

Loss Weights ​

Inference Configuration ​

Overriding Config via Command Line ​