WSINDy4Atmos

Weak SINDy model discovery for atmospheric data.

Python code accompanying the manuscript "Learning Physically Interpretable Atmospheric Models from Data with WSINDy" published in the Journal of Geophysical Research: Machine Learning and Computation.

@article{https://doi.org/10.1029/2025JH000602,
         author = {Minor, Seth and Messenger, Daniel A. and Dukic, Vanja and Bortz, David M.},
         title = {Learning Physically Interpretable Atmospheric Models From Data With WSINDy},
         journal = {Journal of Geophysical Research: Machine Learning and Computation},
         volume = {2},
         number = {3},
         pages = {e2025JH000602},
         doi = {https://doi.org/10.1029/2025JH000602},
         url = {https://agupubs.onlinelibrary.wiley.com/doi/abs/10.1029/2025JH000602},
         eprint = {https://agupubs.onlinelibrary.wiley.com/doi/pdf/10.1029/2025JH000602},
         note = {e2025JH000602 2025JH000602},
         year = {2025}}

Also see:

• Zenodo: Code for "Learning Physically Interpretable Atmospheric Models from Data with WSINDy"
• ArXiV pre-print: "Learning Physically Interpretable Atmospheric Models from Data with WSINDy"
• The tutorials and examples located here for instructions on how to use the codebase.
• To recreate results in the paper, see wsindy_for_weather_examples.ipynb and full_globe_wsindy.py (uses older code).

Note: unfortunately, our examples data files are too large host on GitHub. Luckily, they are all publicly-available! Interested readers are directed towards the following locations:

• barotropic: this PyQG example
• swe_wsindy: this Dedalus example
• stratified: this JHTDB simulation
• full_globe: this ERA5 data store

This algorithm uses the following dependencies:

import torch
import scipy
import numpy as np
import itertools
import symengine as sp

import torch.linalg as la
from scipy.signal import convolve
from scipy.special import factorial
import matplotlib.pyplot as plt
from tqdm import tqdm

from wsindy import *
from helper_fcns import *

Install PyWSINDy for PDEs a Bash environment:

wget -q https://raw.githubusercontent.com/SethMinor/PyWSINDy-for-PDEs/main/wsindy.py
wget -q https://raw.githubusercontent.com/SethMinor/PyWSINDy-for-PDEs/main/helper_fcns.py