neuralgcm/neuralgcm
Hybrid ML + physics model of the Earth's atmosphere
13 hidden assumptions · 7-stage pipeline · 9 components
Like any codebase, this scientific computing makes assumptions it never checks — most are routine. The ones worth your attention are below.
Hybrid ML/physics model that simulates Earth's atmosphere by combining neural networks with traditional physics equations
Atmospheric data enters as observational records (temperature, wind, pressure on lat/lon grids), gets transformed to spectral harmonic coefficients for efficient physics calculations, passes through both traditional equation solvers and neural network parameterizations that compute tendency updates, then integrates forward in time before being converted back to observational variables for comparison with targets. The core loop alternates between physics-based tendencies (computed in spectral space) and ML-based corrections (learned from data residuals).
Under the hood, the system uses 3 feedback loops, 3 data pools, 4 control points to manage its runtime behavior.
A 9-component scientific computing. 178 files analyzed. Data flows through 7 distinct pipeline stages.
Hidden Assumptions
Most of what this code assumes is routine. These 3 are the ones most likely to cause trouble here. The rest are minor; they're under "Show everything".
If pressure levels are not monotonic (e.g., due to numerical errors or incorrect data), jnp.searchsorted produces wrong indices leading to interpolation to wrong pressure levels, causing physically impossible temperature/wind profiles
If training data uses different variable names (e.g., 'u_wind' vs 'u_component_of_wind'), the operator silently produces wrong observation mappings, causing the model to train on incorrect variable associations
If the model is configured with a different number of vertical levels or pressure units differ from the hardcoded reference, array indexing fails or physical calculations use wrong pressure values, corrupting thermodynamic computations
Show everything (10 more)
The neural parameterization expects spherical harmonic coefficients to be ordered consistently with the ylm_map coordinate system, particularly that wavenumber indices match between input data and the mapping
If this fails: If input coefficients have different wavenumber ordering than expected by ylm_map, the neural network processes wrong spectral components, leading to physically inconsistent tendency predictions that violate conservation laws
neuralgcm/experimental/atmosphere/parameterizations.py:ModalNeuralDivCurlParameterization
FixerModule protocol assumes that diagnostic fields and target fields have identical coordinate systems and that the diagnostic calculation succeeded without NaN/inf values
If this fails: If diagnostics contain NaN/inf due to numerical instability or coordinate mismatches, the fixer propagates these bad values into tendency corrections, causing the entire simulation to diverge
neuralgcm/experimental/atmosphere/fixers.py:FixerModule
Energy balance adjustments assume that energy imbalances represent steady-state errors rather than transient physical processes, and that redistributing energy preserves the model's temporal stability
If this fails: If energy imbalances arise from legitimate transient processes (e.g., during rapid weather transitions), forcing balance creates artificial energy redistributions that suppress realistic weather variability
neuralgcm/experimental/atmosphere/diagnostics.py:EnergyBalanceModule
Linear interpolation with extrapolation assumes that arrays fit in memory and that unlimited extrapolation won't produce values outside the valid physical range for atmospheric variables
If this fails: Extrapolation can produce unphysical values like negative absolute temperatures or supersonic wind speeds that cause downstream numerical instability in physics computations
neuralgcm/experimental/atmosphere/interpolators.py:_linear_interp_with_linear_extrap
The regridding function assumes input datasets have coordinate dimensions named exactly 'longitude' and 'latitude', and that coordinate values are in the same units as expected by the regridder
If this fails: If input data uses different coordinate names or units (degrees vs radians), xarray.apply_ufunc fails to find the core dimensions or regridding operates on wrong coordinate values, producing geographically misplaced data
neuralgcm/demo.py:_horizontal_regrid
Demo functions assume that importlib.resources and pickle can successfully load packaged model data, and that the dinosaur library's coordinate systems are compatible with the loaded model configuration
If this fails: If package installation is incomplete or dinosaur version mismatches, demo fails with unclear import errors, making it impossible for users to verify basic functionality
neuralgcm/demo.py
Pressure interpolation assumes that the pressure range in input data covers the target pressure levels, but linear extrapolation is used without bounds checking for extreme values
If this fails: Extrapolating to pressure levels far outside the data range (e.g., stratosphere when data only covers troposphere) produces unphysical atmospheric profiles that violate hydrostatic balance
neuralgcm/experimental/atmosphere/interpolators.py:LinearOnPressure
The primitive equations module assumes that input atmospheric state fields have consistent physical units and that the coordinate system metadata correctly describes the spatial grid
If this fails: Unit inconsistencies or wrong coordinate metadata cause physics calculations to use incorrect scaling factors, leading to wrong magnitude tendencies that can destabilize the simulation
neuralgcm/experimental/atmosphere/equations.py:PrimitiveEquationsModule
The observation operator assumes that vertical levels in the atmospheric state are ordered consistently with the target observation data (either surface-to-top or top-to-surface)
If this fails: If vertical ordering differs between model and observations, computed observation variables have flipped vertical profiles, causing training to learn inverted atmospheric relationships
neuralgcm/experimental/atmosphere/observation_operators.py:StandardVariablesObservationOperator
Atmospheric diagnostics assume that all required state variables are present and physically meaningful, with no validation that computed diagnostics fall within realistic ranges for Earth's atmosphere
If this fails: Diagnostic calculations proceed even with unphysical input states, producing misleading diagnostic values that hide model problems and make it difficult to detect when the simulation has gone off track
neuralgcm/experimental/atmosphere/diagnostics.py
Open the standalone hidden-assumptions report for neuralgcm →
How Data Flows Through the System
Atmospheric data enters as observational records (temperature, wind, pressure on lat/lon grids), gets transformed to spectral harmonic coefficients for efficient physics calculations, passes through both traditional equation solvers and neural network parameterizations that compute tendency updates, then integrates forward in time before being converted back to observational variables for comparison with targets. The core loop alternates between physics-based tendencies (computed in spectral space) and ML-based corrections (learned from data residuals).
- Initialize atmospheric state — Functions in idealized_states.py convert initial conditions or observations into the model's internal atmospheric state representation using coordinate system transformations and field mapping [ObservedVariables → AtmosphericState]
- Transform to spectral space — FixedYlmMapping.nodal_to_modal converts physical grid fields to spherical harmonic coefficients, enabling efficient computation of derivatives and global operations required by physics equations [AtmosphericState → SphericalHarmonicCoefficients]
- Compute physics tendencies — PrimitiveEquationsModule calculates time derivatives using traditional atmospheric physics equations - momentum advection, pressure gradients, Coriolis forces, and thermodynamic processes [SphericalHarmonicCoefficients → Tendencies]
- Apply neural parameterizations — ModalNeuralDivCurlParameterization and similar neural modules predict sub-grid scale processes by learning from data residuals, adding ML-computed tendencies to physics-based ones [SphericalHarmonicCoefficients → Tendencies]
- Enforce conservation constraints — EnergyFixerModule adjusts total tendencies to maintain physical conservation laws like energy and mass balance, redistributing any imbalances according to physical constraints [Tendencies → Tendencies]
- Integrate forward in time — ImplicitExplicitTimeIntegrator advances the atmospheric state using the combined tendencies, handling fast gravity waves implicitly and advection explicitly for numerical stability [AtmosphericState → AtmosphericState]
- Convert to observations — StandardVariablesObservationOperator transforms the updated model state back to standard meteorological variables (u,v,t,z) for comparison with targets during training or for output [AtmosphericState → ObservedVariables] (config: u_name, v_name, vector_name)
Data Models
The data structures that flow between stages — the contracts that hold the system together.
neuralgcm/experimental/atmosphere/equations.pydict[str, cx.Field] with keys like 'temperature', 'u_component_of_wind', 'v_component_of_wind', 'surface_pressure' where cx.Field contains arrays with coordinate system metadata
Created from initial conditions or observations, transformed between coordinate systems, updated by physics and ML components, and converted to observational variables
neuralgcm/experimental/core/spherical_harmonics.pycomplex arrays with dimensions (wavenumber_m, wavenumber_l, vertical_levels) representing atmospheric fields in spectral space
Generated by transforming gridded data to spectral space, operated on by physics equations and neural networks, then transformed back to physical grid
neuralgcm/experimental/core/coordinates.pycx.Coordinate objects defining spatial grids like LonLatGrid.T42() with longitude/latitude arrays or SigmaLevels with vertical level definitions
Defined at model initialization, used throughout simulation to define where data lives spatially, enables transformations between different grid resolutions
neuralgcm/experimental/atmosphere/equations.pydict[str, cx.Field] with same structure as AtmosphericState but containing time derivatives (rates of change) for each variable
Computed by physics equations and neural network parameterizations, accumulated from multiple sources, fed into time integrator to update atmospheric state
neuralgcm/experimental/xreader/iterators.pyPyTree with 'block' containing atmospheric data arrays and 'sub_slices' list defining temporal/spatial windows for training sequences
Loaded from zarr datasets, chunked into training sequences, fed through model for prediction and loss computation
neuralgcm/experimental/training/trainer.pydataclass with train: dict[str, float], eval: dict[str, float], eval_ema: dict[str, float] containing loss and accuracy metrics
Computed during training steps, accumulated over batches, logged for monitoring training progress and model performance
System Behavior
How the system operates at runtime — where data accumulates, what loops, what waits, and what controls what.
Data Pools
Maintains current atmospheric state (temperature, winds, pressure) across integration steps, preserving coordinate system metadata and field relationships
Temporarily stores spherical harmonic coefficients during physics computations, enabling efficient derivative calculations and spectral filtering
Zarr-based storage of atmospheric training data chunked for efficient streaming during model training and evaluation
Feedback Loops
- Atmospheric Integration Loop (simulation, reinforcing) — Trigger: Time step advancement. Action: Each timestep: compute tendencies from current state, integrate forward using implicit-explicit scheme, update state, repeat for next timestep. Exit: Simulation end time reached or convergence criteria met.
- Training Loss Minimization (training-loop, balancing) — Trigger: Batch of training data. Action: Forward pass through model predicts atmospheric state, compute loss against observations, backpropagate gradients, update neural network parameters. Exit: Maximum epochs reached or validation loss stops improving.
- Energy Conservation Adjustment (self-correction, balancing) — Trigger: Energy imbalance detected in tendencies. Action: EnergyFixerModule redistributes energy imbalance across state variables according to physical constraints to maintain conservation. Exit: Energy balance restored within tolerance.
Delays
- Spectral Transform Computation (computation, ~varies with grid resolution) — Forward and backward spherical harmonic transforms require FFTs and Legendre polynomial evaluations, creating computational bottlenecks at high resolution
- Data Loading from Zarr (async-processing, ~depends on chunk size and I/O bandwidth) — Training data must be read from compressed zarr archives, potentially causing GPU idle time if data pipeline doesn't keep up
- Neural Network Forward Pass (computation, ~proportional to model size and batch size) — ML parameterizations add computational overhead to each physics timestep, trading speed for improved accuracy
Control Points
- Grid Resolution Selection (architecture-switch) — Controls: Spatial resolution of the atmospheric simulation through LonLatGrid.T21(), T42(), etc. affecting accuracy and computational cost. Default: T21, T42 common choices
- Integration Time Step (hyperparameter) — Controls: Temporal resolution of the simulation, balancing numerical stability with computational efficiency
- Vertical Coordinate System (architecture-switch) — Controls: Choice between sigma levels, pressure levels, or hybrid coordinates affecting how vertical structure is represented. Default: SigmaLevels.equidistant() commonly used
- Observation Variable Names (feature-flag) — Controls: Which meteorological variables to extract for training and evaluation (u_component_of_wind, v_component_of_wind, etc.). Default: u_name='u_component_of_wind', v_name='v_component_of_wind'
Technology Stack
Provides automatic differentiation and JIT compilation for both physics calculations and neural network training in the atmospheric model
Builds neural network components that parameterize sub-grid atmospheric processes, integrating with JAX for gradient-based training
Manages coordinate system abstractions and field transformations, enabling seamless operations between different spatial representations
Provides core atmospheric physics equations and primitive equation solvers that form the physics backbone of the hybrid model
Handles labeled multi-dimensional atmospheric datasets with metadata preservation for input/output operations
Stores compressed atmospheric training data in chunked format for efficient streaming during model training
Implements gradient-based optimizers for training the neural network components of the atmospheric model
Manages configuration and hyperparameter injection for experiments and model architecture choices
Key Components
- PrimitiveEquationsModule (processor) — Implements the primitive equations of atmospheric dynamics, computing tendencies from current state using physics-based differential equations for momentum, temperature, and moisture
neuralgcm/experimental/atmosphere/equations.py - FixedYlmMapping (transformer) — Transforms atmospheric fields between physical grid coordinates and spherical harmonic spectral coefficients, enabling efficient computation of derivatives and global operations
neuralgcm/experimental/core/spherical_harmonics.py - ModalNeuralDivCurlParameterization (processor) — Neural network module that learns to parameterize unresolved atmospheric processes by predicting divergence and vorticity tendencies in spectral space
neuralgcm/experimental/atmosphere/parameterizations.py - StandardVariablesObservationOperator (adapter) — Converts model's internal atmospheric state representation to standard meteorological variables (u,v,t,z) that match observational data formats for training and evaluation
neuralgcm/experimental/atmosphere/observation_operators.py - LinearOnPressure (transformer) — Interpolates atmospheric fields between different vertical coordinate systems (sigma levels to pressure levels) using linear interpolation with extrapolation
neuralgcm/experimental/atmosphere/interpolators.py - ImplicitExplicitTimeIntegrator (processor) — Advances atmospheric state forward in time using implicit-explicit schemes that handle fast gravity waves implicitly and advection explicitly for numerical stability
neuralgcm/experimental/core/time_integrators.py - BlockIterator (loader) — Streams training data from zarr archives, yielding blocks of atmospheric data with configurable temporal and spatial chunking for efficient batch processing
neuralgcm/experimental/xreader/iterators.py - AtmosphereDiagnostics (processor) — Computes derived atmospheric quantities like energy budgets, mass conservation diagnostics, and physical consistency checks from model state
neuralgcm/experimental/atmosphere/diagnostics.py - EnergyFixerModule (processor) — Adjusts model tendencies to conserve total energy by redistributing energy imbalances across atmospheric state variables according to physical constraints
neuralgcm/experimental/atmosphere/fixers.py
Explore the interactive analysis
See the full architecture map, data flow, and code patterns visualization.
Analyze on CodeSeaRelated Scientific Computing Repositories
Frequently Asked Questions
What is neuralgcm used for?
Hybrid ML/physics model that simulates Earth's atmosphere by combining neural networks with traditional physics equations neuralgcm/neuralgcm is a 9-component scientific computing written in Python. Data flows through 7 distinct pipeline stages. The codebase contains 178 files.
How is neuralgcm architected?
neuralgcm is organized into 5 architecture layers: Atmospheric Physics, Core Components, Training Infrastructure, Data Pipeline, and 1 more. Data flows through 7 distinct pipeline stages. This layered structure keeps concerns separated and modules independent.
How does data flow through neuralgcm?
Data moves through 7 stages: Initialize atmospheric state → Transform to spectral space → Compute physics tendencies → Apply neural parameterizations → Enforce conservation constraints → .... Atmospheric data enters as observational records (temperature, wind, pressure on lat/lon grids), gets transformed to spectral harmonic coefficients for efficient physics calculations, passes through both traditional equation solvers and neural network parameterizations that compute tendency updates, then integrates forward in time before being converted back to observational variables for comparison with targets. The core loop alternates between physics-based tendencies (computed in spectral space) and ML-based corrections (learned from data residuals). This pipeline design reflects a complex multi-stage processing system.
What technologies does neuralgcm use?
The core stack includes JAX (Provides automatic differentiation and JIT compilation for both physics calculations and neural network training in the atmospheric model), Flax (Builds neural network components that parameterize sub-grid atmospheric processes, integrating with JAX for gradient-based training), coordax (Manages coordinate system abstractions and field transformations, enabling seamless operations between different spatial representations), dinosaur (Provides core atmospheric physics equations and primitive equation solvers that form the physics backbone of the hybrid model), xarray (Handles labeled multi-dimensional atmospheric datasets with metadata preservation for input/output operations), zarr (Stores compressed atmospheric training data in chunked format for efficient streaming during model training), and 2 more. A focused set of dependencies that keeps the build manageable.
What system dynamics does neuralgcm have?
neuralgcm exhibits 3 data pools (Atmospheric State Memory, Spectral Coefficient Cache), 3 feedback loops, 4 control points, 3 delays. The feedback loops handle simulation and training-loop. These runtime behaviors shape how the system responds to load, failures, and configuration changes.
What design patterns does neuralgcm use?
5 design patterns detected: Coordinate System Abstraction, Physics-ML Hybrid Architecture, Spectral Method Implementation, Conservation Law Enforcement, Field-Coordinate Binding.
Analyzed on April 20, 2026 by CodeSea. Written by Karolina Sarna.