Run on a mac the climax inference

[sakomws]

This pull request focuses on improving compatibility and usability for running ClimaX regional forecasting on Apple Silicon (macOS), adds predict-only demo scripts, and updates configuration and documentation to make experimentation easier. The main changes include adjustments to the training configuration for macOS, the addition of demo scripts for inference, enhancements to training script compatibility across Lightning versions, and expanded documentation for setup and usage.

Apple Silicon/macOS compatibility & configuration updates:

• Updated configs/regional_forecast_climax.yaml to use "mps" accelerator, full precision (precision: 32-true), and macOS-friendly data paths and batch sizes; also disables problematic features for Apple Silicon and streamlines checkpoint and logging settings. [1] [2] [3] [4] [5] [6]

Predict-only and demo scripts:

• Added predict_only_climax.py and demo_compare_resolutions.py scripts to perform inference with pretrained ClimaX models, including side-by-side resolution comparison and support for macOS/MPS. These scripts provide dummy input generation, flexible variable selection, and output visualization. [1] [2]

Training script compatibility and robustness:

• Refactored src/climax/regional_forecast/train.py to support both old and new Lightning import paths, handle config saving robustly, and ensure model/data module wiring works across Lightning versions. [1] [2]

Documentation and reproducibility:

• Added output/README.md with step-by-step setup instructions, dependency list, and usage examples for both demo and predict-only scripts, improving onboarding and reproducibility for new users.

Minor code quality improvements:

• Updated type usage in src/climax/utils/pos_embed.py to use standard Python types (float instead of np.float) for better compatibility and future-proofing. [1] [2]

Run

Download Weights: https://huggingface.co/microsoft/ClimaX/tree/main

python predict_only_climax.py \
  --ckpt ./data-ch/5.625deg.ckpt \   
  --res 5.625 \
  --device mps

python demo_compare_resolutions.py \
  --ckpt_5625 ./data-ch/5.625deg.ckpt \
  --ckpt_1406 ./data-ch/1.40625deg.ckpt \
  --device mps \
  --out_vars geopotential_500 temperature_850 2m_temperature

python src/climax/regional_forecast/train.py \
  --config configs/regional_forecast_climax.yaml \
  --trainer.callbacks=null \
  --trainer.accelerator=mps --trainer.devices=1 --trainer.precision=32-true \
  --data.root_dir=/mnt/data/5.625deg_npz \
  --data.region="NorthAmerica" \
  --data.predict_range=72 \
  --data.out_variables="['z_500','t_850','t2m']" \
  --data.batch_size=1 --data.num_workers=0 \
  --trainer.max_epochs=1 \
  --model.pretrained_path="./data-ch/1.40625deg.ckpt" \
  --model.lr=5e-7 --model.beta_1=0.9 --model.beta_2=0.99 \
  --model.weight_decay=1e-5

[sakomws]

@microsoft-github-policy-service agree company="Gladly"

[microsoft-github-policy-service]

@sakomws the command you issued was incorrect. Please try again.

Examples are:

@microsoft-github-policy-service agree

and

@microsoft-github-policy-service agree company="your company"