### Install and import dependencies

Source: https://github.com/google-research/neuralgcm/blob/main/docs/neuralgcm_datasets.ipynb

Install the required libraries and import them to begin working with NeuralGCM datasets.

```python
! pip install -q -U zarr xarray gcsfs
import xarray
import zarr
```

--------------------------------

### Install Dependencies

Source: https://github.com/google-research/neuralgcm/blob/main/docs/simulation_state_components_intro.ipynb

Installs necessary libraries including flax, coordax, neuralgcm, jax_datetime, and fiddle.

```bash
! pip install flax==0.12.0
! pip install --no-cache-dir git+https://github.com/neuralgcm/coordax.git
! pip install --no-cache-dir git+https://github.com/neuralgcm/neuralgcm.git@refs/pull/482/head
! pip install jax_datetime jax_solar fiddle
```

--------------------------------

### Instantiate and run a simple module

Source: https://github.com/google-research/neuralgcm/blob/main/docs/simulation_state_components_intro.ipynb

Example of setting up a SimpleModule and executing it with labeled axis inputs.

```python
simple_module = SimpleModule()
coords = cx.LabeledAxis('x', [0, 1, 2])
x_data = cx.field(jnp.array([10., 20., 30.]), coords)
inputs = {'x': x_data, 't': jdt.to_datetime('2000-01-01')}
simple_outputs = simple_module(inputs)
simple_outputs
```

--------------------------------

### Instantiate Lorenz96 model

Source: https://github.com/google-research/neuralgcm/blob/main/docs/model_api_intro.ipynb

Example of setting up coordinate axes and initializing the Lorenz96 model.

```python
# Instantiating coordinates and Lorenz96 model.
k = cx.LabeledAxis('k', np.arange(36))
j = cx.LabeledAxis('j', np.arange(8))

l96_model = Lorenz96(k, j)

nnx.display(l96_model)  # inspecting model using treescope.
```

--------------------------------

### Install Apache Beam

Source: https://github.com/google-research/neuralgcm/blob/main/docs/inference_model_api_intro.ipynb

Installs the Apache Beam SDK, which is required for running inference tasks as a Beam pipeline.

```bash
! pip install apache_beam
```

--------------------------------

### InferenceRunner Setup for HPC Array Jobs

Source: https://github.com/google-research/neuralgcm/blob/main/docs/inference_model_api_intro.ipynb

Demonstrates how to set up the InferenceRunner in a setup script for HPC environments. The runner is initialized and set up for subsequent task execution.

```python
runner = InferenceRunner(...)
runner.setup()
```

--------------------------------

### Install NeuralGCM from Source

Source: https://github.com/google-research/neuralgcm/blob/main/docs/installation.md

Use this command to install NeuralGCM and its dependencies. For best performance, ensure you have a compatible GPU or TPU attached.

```bash
pip install neuralgcm
```

--------------------------------

### Setup InferenceRunner for Direct Execution

Source: https://github.com/google-research/neuralgcm/blob/main/docs/inference_model_api_intro.ipynb

Initializes and sets up the InferenceRunner for direct execution. Ensure all parameters like output_freq and output_duration are congruent.

```python
output_path = './reforecasts.zarr'
output_query = {'slow': {'x': k}}
init_times = pd.date_range(
    pd.Timestamp('2000-01-14'),
    pd.Timestamp('2000-04-14'),
    freq=pd.Timedelta('1 day'),
    inclusive='left',
).values

zarr_chunks = {'lead_time': 1, 'init_time': 1}
runner = inference_runner.InferenceRunner(
    model=l96_inference_model,
    inputs=ground_truth_data,
    dynamic_inputs=inference_dynamic_inputs.EmptyDynamicInputs(),
    init_times=init_times,
    ensemble_size=1,
    zarr_chunks=zarr_chunks,
    output_path=output_path,
    output_query=output_query,
    output_freq=np.timedelta64(6, 'h'),
    output_duration=np.timedelta64(48, 'h'),
    write_duration=np.timedelta64(24, 'h'),
    unroll_duration=np.timedelta64(24, 'h'),
    checkpoint_duration=np.timedelta64(24, 'h'),
)
runner.setup()
```

--------------------------------

### Install NeuralGCM and Dependencies

Source: https://github.com/google-research/neuralgcm/blob/main/docs/inference_demo.ipynb

Install the NeuralGCM library and its dependencies using pip. This is a prerequisite for running the forecasting notebook.

```python
# if necessary, install NeuralGCM and dependencies
! pip install -q -U neuralgcm dinosaur gcsfs
```

--------------------------------

### Initialize DataCoupler instances

Source: https://github.com/google-research/neuralgcm/blob/main/docs/forced_parmeterized_coupled_l96.ipynb

Example of building a coupled model by wiring data between components using DataCoupler.

```python
# Building coupled model.
coupler_with_x = DataCoupler({'x': k}, {'x': 'slow'})
coupler_with_y = DataCoupler({'y': kj}, {'y': 'fast'})
```

--------------------------------

### DynamicInputSlice structure example

Source: https://github.com/google-research/neuralgcm/blob/main/docs/simulation_state_components_intro.ipynb

Example output representing the structure of DynamicInputSlice objects used for managing time-dependent data inputs.

```text
[DynamicInputSlice( # DynamicInput: 3 (20 B)
  keys_to_coords={'co2': Scalar()},
  observation_key='soil',
  time_axis=0,
  time=DynamicInput( # 2 (16 B)
    value=<Field dims=('timedelta',) shape=(1,) axes={'timedelta': TimeDelta} >
  ),
  data=DynamicInput( # 1 (4 B)
    value={'co2': <Field dims=('timedelta',) shape=(1,) axes={'timedelta': TimeDelta} >}
  )
), DynamicInputSlice( # DynamicInput: 3 (20 B)
  keys_to_coords={'co2': Scalar()},
  observation_key='atm',
  time_axis=0,
  time=DynamicInput( # 2 (16 B)
    value=<Field dims=('timedelta',) shape=(1,) axes={'timedelta': TimeDelta} >
  ),
  data=DynamicInput( # 1 (4 B)
    value={'co2': <Field dims=('timedelta',) shape=(1,) axes={'timedelta': TimeDelta} >}
  )
)]
```

--------------------------------

### Load Pre-trained Model from Checkpoint

Source: https://context7.com/google-research/neuralgcm/llms.txt

Load a pre-trained NeuralGCM model from a checkpoint file stored on Google Cloud Storage. This example demonstrates loading a deterministic 1.4-degree model and inspecting its properties.

```python
import gcsfs
import pickle
import neuralgcm

# Load checkpoint from Google Cloud Storage
gcs = gcsfs.GCSFileSystem(token='anon')
model_name = 'v1/deterministic_1_4_deg.pkl'

with gcs.open(f'gs://neuralgcm/models/{model_name}', 'rb') as f:
    checkpoint = pickle.load(f)

# Create model from checkpoint
model = neuralgcm.PressureLevelModel.from_checkpoint(checkpoint)

# Inspect model properties
print(f"Timestep: {model.timestep}")  # e.g., numpy.timedelta64(1, 'h')
print(f"Input variables: {model.input_variables}")
# ['geopotential', 'specific_humidity', 'temperature',
#  'u_component_of_wind', 'v_component_of_wind', ...]
print(f"Forcing variables: {model.forcing_variables}")
# ['sea_ice_cover', 'sea_surface_temperature']
```

--------------------------------

### Setup InferenceRunner for Beam Pipeline

Source: https://github.com/google-research/neuralgcm/blob/main/docs/inference_model_api_intro.ipynb

Initializes and sets up the InferenceRunner for execution within an Apache Beam pipeline. This configuration is similar to direct execution but intended for distributed processing.

```python
output_path = './reforecasts_beam.zarr'
output_query = {'slow': {'x': k}}
init_times = pd.date_range(
    pd.Timestamp('2000-01-14'),
    pd.Timestamp('2000-04-14'),
    freq=pd.Timedelta('1 day'),
    inclusive='left',
).values

zarr_chunks = {'lead_time': 1, 'init_time': 1}
runner = inference_runner.InferenceRunner(
    model=l96_inference_model,
    inputs=ground_truth_data,
    dynamic_inputs=inference_dynamic_inputs.EmptyDynamicInputs(),
    init_times=init_times,
    ensemble_size=1,
    zarr_chunks=zarr_chunks,
    output_path=output_path,
    output_query=output_query,
    output_freq=np.timedelta64(6, 'h'),
    output_duration=np.timedelta64(48, 'h'),
    write_duration=np.timedelta64(24, 'h'),
    unroll_duration=np.timedelta64(24, 'h'),
    checkpoint_duration=np.timedelta64(24, 'h'),
)
runner.setup()
```

--------------------------------

### Prepare Model Inputs and Forcings from xarray

Source: https://context7.com/google-research/neuralgcm/llms.txt

Convert xarray Datasets into the dictionary format required by NeuralGCM methods. This example shows how to prepare input variables (wind, temperature, etc.) and forcing variables (sea surface temperature, sea ice cover) from a dataset.

```python
import neuralgcm
import xarray

# Load model and data
checkpoint = neuralgcm.demo.load_checkpoint_tl63_stochastic()
model = neuralgcm.PressureLevelModel.from_checkpoint(checkpoint)
dataset = neuralgcm.demo.load_data(model.data_coords)

# Convert xarray to model input format
inputs = model.inputs_from_xarray(dataset.isel(time=0))
forcings = model.forcings_from_xarray(dataset.isel(time=0))

# Inspect shapes - arrays have shape (level, longitude, latitude)
import jax
print(jax.tree.map(lambda x: x.shape, inputs))
# {'geopotential': (37, 128, 64), 'temperature': (37, 128, 64),
#  'u_component_of_wind': (37, 128, 64), 'sim_time': (), ...}

print(jax.tree.map(lambda x: x.shape, forcings))
# {'sea_ice_cover': (1, 128, 64), 'sea_surface_temperature': (1, 128, 64),
#  'sim_time': ()}
```

--------------------------------

### Import Necessary Libraries

Source: https://github.com/google-research/neuralgcm/blob/main/docs/inference_demo.ipynb

Import essential Python libraries for data manipulation, model interaction, and numerical operations. Ensure these libraries are installed before execution.

```python
import gcsfs
import jax
import numpy as np
import pickle
import xarray

from dinosaur import horizontal_interpolation
from dinosaur import spherical_harmonic
from dinosaur import xarray_utils
import neuralgcm
```

--------------------------------

### Load ERA5 Data for NeuralGCM

Source: https://github.com/google-research/neuralgcm/blob/main/docs/data_preparation.ipynb

Loads a demo NeuralGCM model and opens a Zarr copy of ERA5 data, selecting the variables required by the model. Ensure you have the necessary libraries installed.

```python
import jax
import numpy as np
import neuralgcm
import xarray

# load demo model
checkpoint = neuralgcm.demo.load_checkpoint_tl63_stochastic()
model = neuralgcm.PressureLevelModel.from_checkpoint(checkpoint)

# create a xarray.Dataset with required variables for NeuralGCM
path = 'gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3'
full_era5 = xarray.open_zarr(
    path, chunks=None, storage_options=dict(token='anon')
)
full_era5 = full_era5[model.input_variables + model.forcing_variables]
```

--------------------------------

### Regrid ERA5 Data to Gaussian Grid

Source: https://context7.com/google-research/neuralgcm/llms.txt

Regrid ERA5 data from its native grid to NeuralGCM's Gaussian grid using utilities from the Dinosaur library. This example demonstrates opening ERA5 data from ARCO-ERA5 and selecting required variables and a time slice.

```python
import xarray
import neuralgcm
from dinosaur import horizontal_interpolation
from dinosaur import spherical_harmonic
from dinosaur import xarray_utils

# Load model
checkpoint = neuralgcm.demo.load_checkpoint_tl63_stochastic()
model = neuralgcm.PressureLevelModel.from_checkpoint(checkpoint)

# Open ERA5 data from ARCO-ERA5 project
era5_path = 'gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3'
full_era5 = xarray.open_zarr(
    era5_path, chunks=None, storage_options=dict(token='anon')
)

# Select required variables and time slice
sliced_era5 = (
    full_era5[model.input_variables + model.forcing_variables]
    .sel(time='2020-01-01T00')
    .compute()
)


```

--------------------------------

### Adding Point Measurement Observation Operator

Source: https://github.com/google-research/neuralgcm/blob/main/docs/model_api_intro.ipynb

This example shows how to add a simple observation operator that selects a specific value (k=4) from the state, mimicking a point measurement. It then unrolls the model and plots the trajectory of this measurement.

```python
sel_point = transforms.Sel({'k': 4})  # will select k == 4 from the X field.
point_obs = observation_operators.TransformObservationOperator(sel_point)
l96_model.operators['k=4'] = point_obs
k4_trajectory = unroll_model(l96_model, 100, {'k=4': {'x': cx.Scalar()}}, 25)
k4_trajectory['k=4']['x'].to_xarray().plot(x='timedelta')
```

--------------------------------

### Initialize and Advance Model Simulation

Source: https://github.com/google-research/neuralgcm/blob/main/docs/model_api_intro.ipynb

Demonstrates preparing inputs, assimilating data, observing state, and advancing the simulation over time.

```python
# Using models' core methods.

fig, ax = plt.subplots()  # will be used to plot X_k state.

# Preparing inputs for model initialization
kj = cx.compose_coordinates(k, j)
rng = np.random.RandomState(0)
x_init = abs(10 * np.sin(np.linspace(0, 13 * 2*np.pi, k.sizes['k'])))
x_init[0] += 0.01
t0 = jdt.Datetime.from_isoformat('2000-01-01')
inputs = {
    'slow': {'x': cx.field(x_init, k), 'time': cx.field(t0)},
    'fast': {'y': cx.field(rng.uniform(size=kj.shape, low=-0.5, high=0.5), kj)},
}
# assimilate expects at least a dummy timedelta dimension, we add it here.
t_del = coordinates.TimeDelta(np.zeros(1) * np.timedelta64(1, 'h'))
add_td = lambda f: f.broadcast_like(cx.compose_coordinates(t_del, f.coordinate))
inputs = jax.tree.map(add_td, inputs, is_leaf=cx.is_field)

# Assimilate updates the simulation state in place using inputs.
l96_model.assimilate(inputs)

# Observe lets the user to "query" the simulation state.
observations = l96_model.observe({'slow': {'x': k}})  # observe X.
observations['slow']['x'].to_xarray().plot(x='k', ax=ax, label='i=0')  # plot X.

# Advance updates the simulation state in place.
for _ in range(500):
  l96_model.advance()

observations = l96_model.observe({'slow': {'x': k}})  # observe evolved state.
observations['slow']['x'].to_xarray().plot(x='k', ax=ax, label='i=500')
plt.legend()
```

--------------------------------

### Use DataObservationOperator

Source: https://github.com/google-research/neuralgcm/blob/main/docs/model_api_intro.ipynb

Example of using DataObservationOperator to return precomputed predictions based on a query.

```python
# DataObservationOperator example.

x = cx.SizedAxis('x', 5)
precomputed_predictions = {'u': cx.field(np.arange(5), x), 't': cx.field(0.0)}
obs_op = observation_operators.DataObservationOperator(precomputed_predictions)
ignored_inputs = {}
print(obs_op.observe(ignored_inputs, {'u': x}))
print(obs_op.observe(ignored_inputs, {'t': cx.Scalar()}))

with print_error():
 obs_op.observe(ignored_inputs, {'v': x})
```

--------------------------------

### Demo Utilities

Source: https://github.com/google-research/neuralgcm/blob/main/docs/api.md

Helper functions for loading test datasets and models without requiring external cloud storage.

```APIDOC
## Demo Utilities

### Description
Functions provided for testing purposes to load small, packaged datasets and model checkpoints.

### Functions
- **neuralgcm.demo.load_data**: Loads a small test dataset.
- **neuralgcm.demo.load_checkpoint_tl63_stochastic**: Loads a small test model checkpoint.
```

--------------------------------

### Build and run Lorenz96 with two scales

Source: https://github.com/google-research/neuralgcm/blob/main/docs/forced_parmeterized_coupled_l96.ipynb

Initializes and runs the `Lorenz96WithTwoScales` model to generate a steady-state rollout. This involves setting up coordinates, initial conditions, and then performing a burn-in phase using `assimilate` and `unroll_model`.

```python
#@title Building model and running Lorenz96 with two scales

k, j = cx.LabeledAxis('k', np.arange(36)), cx.LabeledAxis('j', np.arange(8))
l96_model_two_scales = lorenz96.Lorenz96WithTwoScales(k, j)

# Sample two scale model inputs.
kj = cx.compose_coordinates(k, j)
t0 = jdt.Datetime.from_isoformat('2000-01-01')
dt = coordinates.TimeDelta(np.zeros(1) * np.timedelta64(1, 'h'))
rng = np.random.RandomState(0)

x_init = abs(10 * np.sin(np.linspace(0, 13 * 2 * np.pi, k.sizes['k'])))
inputs = {
    'slow': {
        'x': cx.field(x_init[None, :], dt, k),
        'time': cx.field(t0[None], dt),
    },
    'fast': {
        'y': cx.field(rng.uniform(-0.5, 0.5, dt.shape + kj.shape), dt, kj),
        'time': cx.field(t0[None], dt),
    },
}

# Running burn-in phase to get to a steady state.
l96_model_two_scales.assimilate(inputs)
_ = unroll_model(l96_model_two_scales, 500, {}, inner_steps=25)

```

--------------------------------

### Visualize model state with Treescope

Source: https://github.com/google-research/neuralgcm/blob/main/docs/deepdive_into_models.ipynb

Installs and uses Treescope to inspect the nested structure of the encoded model state.

```python
! pip install -q treescope
```

```python
import treescope

treescope.display(encoded)
```

--------------------------------

### Perform end-to-end weather forecasting

Source: https://context7.com/google-research/neuralgcm/llms.txt

Demonstrates loading ERA5 data, regridding to model resolution, and initializing a forecast.

```python
import gcsfs
import jax
import numpy as np
import pickle
import xarray
from dinosaur import horizontal_interpolation, spherical_harmonic, xarray_utils
import neuralgcm

# Load pre-trained model
gcs = gcsfs.GCSFileSystem(token='anon')
with gcs.open('gs://neuralgcm/models/v1/deterministic_2_8_deg.pkl', 'rb') as f:
    checkpoint = pickle.load(f)
model = neuralgcm.PressureLevelModel.from_checkpoint(checkpoint)

# Load and prepare ERA5 data
era5_path = 'gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3'
full_era5 = xarray.open_zarr(era5_path, chunks=None, storage_options=dict(token='anon'))

demo_start_time = '2020-02-14'
demo_end_time = '2020-02-18'
data_inner_steps = 24

sliced_era5 = (
    full_era5[model.input_variables + model.forcing_variables]
    .pipe(xarray_utils.selective_temporal_shift,
          variables=model.forcing_variables, time_shift='24 hours')
    .sel(time=slice(demo_start_time, demo_end_time, data_inner_steps))
    .compute()
)

# Regrid to model resolution
era5_grid = spherical_harmonic.Grid(
    latitude_nodes=full_era5.sizes['latitude'],
    longitude_nodes=full_era5.sizes['longitude'],
    latitude_spacing=xarray_utils.infer_latitude_spacing(full_era5.latitude),
    longitude_offset=xarray_utils.infer_longitude_offset(full_era5.longitude),
)
regridder = horizontal_interpolation.ConservativeRegridder(
    era5_grid, model.data_coords.horizontal, skipna=True
)
eval_era5 = xarray_utils.regrid(sliced_era5, regridder)
eval_era5 = xarray_utils.fill_nan_with_nearest(eval_era5)

# Initialize model
inputs = model.inputs_from_xarray(eval_era5.isel(time=0))
input_forcings = model.forcings_from_xarray(eval_era5.isel(time=0))
initial_state = model.encode(inputs, input_forcings, jax.random.key(42))

# Use persistence forcings
all_forcings = model.forcings_from_xarray(eval_era5.head(time=1))
```

--------------------------------

### Vectorize Model for Ensemble and Batch Processing

Source: https://github.com/google-research/neuralgcm/blob/main/docs/model_vectorization_tutorial.ipynb

Vectorizes the model for both batch and ensemble dimensions. This setup is necessary for mimicking model training with proper scoring rules.

```python
life = build_new_model()
vectorized_life = life.to_vectorized({
    typing.Prognostic: b,
    typing.Diagnostic: b,
    typing.Randomness: b,
    typing.DynamicInput: b,
})
vectorized_life.update_dynamic_inputs(batched_dynamic_inputs)
vectorized_life.assimilate(batched_inputs)
vectorized_life = vectorized_life.to_vectorized({
    typing.Prognostic: e,
    typing.Diagnostic: e,
    typing.Randomness: e,
})
vectorized_life.initialize_random_processes(eb_rngs)
for _ in range(4):
  vectorized_life.advance()

post_advance_observations = vectorized_life.observe(observe_query)
xarray_utils.fields_to_xarray(post_advance_observations['game']).u.plot(
    x='x', y='y', row='ensemble', col='batch'
)
```

--------------------------------

### Initialize NeuralGCM Environment

Source: https://github.com/google-research/neuralgcm/blob/main/docs/model_api_intro.ipynb

Imports necessary libraries and configures the environment for NeuralGCM model development.

```python
import dataclasses
from functools import partial
from typing import Callable

from flax import nnx
import jax
import jax.numpy as jnp
import numpy as np
import matplotlib.pyplot as plt
import graphviz

import jax_datetime as jdt
import coordax as cx
from neuralgcm.experimental.core import api
from neuralgcm.experimental.core import coordinates
from neuralgcm.experimental.core import observation_operators
from neuralgcm.experimental.core import time_integrators
from neuralgcm.experimental.core import parallelism
from neuralgcm.experimental.core import transforms
from neuralgcm.experimental.core import spherical_transforms
from neuralgcm.experimental.core import typing
from neuralgcm.experimental.core import module_utils


# small formatting improvements
import contextlib
np.set_printoptions(threshold=100)

@contextlib.contextmanager
def print_error():
  try:
    yield
  except Exception as e:
    print(f'{type(e).__name__}: {e}')


import logging
import warnings
warnings.simplefilter('ignore')
logging.getLogger('jax._src.lib.xla_bridge').addFilter(lambda _: False)
```

--------------------------------

### Open NeuralGCM Zarr Dataset

Source: https://github.com/google-research/neuralgcm/blob/main/docs/neuralgcm_datasets.ipynb

Opens a NeuralGCM ensemble dataset from a Zarr store on Google Cloud Storage. Ensure you have the 'gcsfs' library installed for anonymous access.

```python
path = 'gs://weatherbench2/datasets/neuralgcm_ens/2020-256x128.zarr'
ngcm_ens = xarray.open_zarr(path, storage_options=dict(token='anon'))
ngcm_ens
```

--------------------------------

### Configure Dynamic Inputs and Forced Parameterization

Source: https://github.com/google-research/neuralgcm/blob/main/docs/forced_parmeterized_coupled_l96.ipynb

Sets up dynamic input slices and applies forced parameterization to a Lorenz96 model.

```python
dynamic_x = dynamic_io.DynamicInputSlice({'x': k}, observation_key='slow')
x_features = feature_transforms.DynamicInputFeatures(['x'], dynamic_x)

h, c, b = 1.0, 10.0, 10.0  # will use expected parameters.
x_to_coupling_fn = lambda x: (h * c / b) * x
x_to_coupling = transforms.ApplyFnToKeys(x_to_coupling_fn, ['x'])
parameterization_transform = transforms.Sequential([x_features, x_to_coupling])
forced_parameterization = EulerParameterization(parameterization_transform, 'x')

forced_y_l96 = lorenz96.Lorenz96FastMode(k, j, [forced_parameterization])
forced_y_l96.assimilate(fast_mode_inputs)
forced_y_l96.update_dynamic_inputs({'slow': all_targets['slow']})
forced_y_predictions = unroll_model(
    forced_y_l96, n_steps, {'fast': {'y': kj}}
)
forced_y_fast_ds = xarray_utils.fields_to_xarray(forced_y_predictions['fast'])
forced_y_fast_ds = forced_y_fast_ds.stack(k_prime=('k', 'j'))
forced_y_fast_ds.coords['k_prime'] = np.linspace(0, k.sizes['k'], math.prod(kj.shape))

# Also include model without forcing for comparison.
default_y_l96 = lorenz96.Lorenz96FastMode(k, j)

default_y_l96.assimilate(fast_mode_inputs)
default_y_predictions = unroll_model(
    default_y_l96, n_steps, {'fast': {'y': kj}}
)
default_y_fast_ds = xarray_utils.fields_to_xarray(default_y_predictions['fast'])
default_y_fast_ds = default_y_fast_ds.stack(k_prime=('k', 'j'))
default_y_fast_ds.coords['k_prime'] = np.linspace(0, k.sizes['k'], math.prod(kj.shape))
```

--------------------------------

### Load Toy Stochastic Model and Demo Data

Source: https://context7.com/google-research/neuralgcm/llms.txt

Load a small toy stochastic model for testing purposes, included within the package. This also demonstrates loading corresponding demo data that matches the model's coordinate system.

```python
import neuralgcm

# Load toy TL63 stochastic model (included in package)
checkpoint = neuralgcm.demo.load_checkpoint_tl63_stochastic()
model = neuralgcm.PressureLevelModel.from_checkpoint(checkpoint)

# Load demo data matching the model's coordinate system
demo_data = neuralgcm.demo.load_data(model.data_coords)
print(demo_data)
# <xarray.Dataset> with variables: geopotential, temperature,
# u_component_of_wind, v_component_of_wind, specific_humidity, etc.
```

--------------------------------

### Run Simulation and Plot Trajectory

Source: https://github.com/google-research/neuralgcm/blob/main/docs/model_api_intro.ipynb

Execute the `unroll_model` function to generate a simulation trajectory and visualize the results using `matplotlib`. This example plots the 'slow' component of the trajectory.

```python
trajectory = unroll_model(l96_model, 100, {'slow': {'x': k}}, 25)
trajectory['slow']['x'].to_xarray().plot.imshow(x='k', y='timedelta')
```

--------------------------------

### Initialize and Run Lorenz96 Two-Scale Simulation

Source: https://github.com/google-research/neuralgcm/blob/main/docs/forced_parmeterized_coupled_l96.ipynb

Extracts current state, prepares inputs for slow and fast modes, and performs a long rollout to generate target datasets.

```python
x1, y1 = l96_model_two_scales.x.value.data, l96_model_two_scales.y.value.data
t1 = l96_model_two_scales.time.value.data
slow_mode_inputs = {
    'slow': {
        'x': cx.field(x1[None, :], dt, k), 'time': cx.field(t1[None], dt)
    },
}
fast_mode_inputs = {
    'fast': {
        'y': cx.field(y1[None, :], dt, kj), 'time': cx.field(t1[None], dt)
    },
}

# Producing a long rollout outputting both slow and fast modes + time.
n_steps = 2000
queries = {
    'slow': {'x': k, 'time': cx.Scalar()},
    'fast': {'y': kj, 'time': cx.Scalar()},
}
all_targets = unroll_model(l96_model_two_scales, n_steps, queries)

# Prep datasets for plotting the ground truth.
targets_slow_ds = xarray_utils.fields_to_xarray(all_targets['slow'])
targets_fast_ds = xarray_utils.fields_to_xarray(all_targets['fast'])
targets_fast_ds = targets_fast_ds.stack(k_prime=('k', 'j'))
targets_fast_ds.coords['k_prime'] = np.linspace(0, k.sizes['k'], math.prod(kj.shape))
```

--------------------------------

### Initialize and execute module diagnostics

Source: https://github.com/google-research/neuralgcm/blob/main/docs/simulation_state_components_intro.ipynb

Instantiate the diagnostic with extraction functions and retrieve values after the module execution.

```python
diagnostic = diagnostics.InstantDiagnostic(
    extract=extract_fn,
    extract_coords=diagnostic_coords,
)
simple_with_diagnostic = SimpleModuleWithDiagnostic(diagnostic=diagnostic)
_ = simple_with_diagnostic(inputs)
```

```python
diagnostic_values = diagnostic.diagnostic_values()
print(f"Diagnosed value: {diagnostic_values['intermediate'].data}")
```

--------------------------------

### Create Vectorized Model Instance

Source: https://github.com/google-research/neuralgcm/blob/main/docs/model_vectorization_tutorial.ipynb

Demonstrates how to instantiate a vectorized model using the `from_model_with_vectorization` classmethod. This method wraps a model instance and handles state vectorization.

```python
# First we create a `Model` instance
```

--------------------------------

### Load 1.4 Degree Deterministic NeuralGCM Model

Source: https://github.com/google-research/neuralgcm/blob/main/docs/neuralgcm_datasets.ipynb

Use this code to open the Zarr archive for the 2-year, 1.4-degree deterministic NeuralGCM model. Ensure you have xarray installed. The 'anon' token is used for public access.

```python
path = 'gs://neuralgcm/amip_runs/v1_deterministic_1_4_deg/2019-to-2021_256x128_gauss_37-level_stride12h.zarr/'
ngcm_1_4 = xarray.open_zarr(path, storage_options=dict(token='anon'))
ngcm_1_4
```

--------------------------------

### Execute Batched Rollout

Source: https://github.com/google-research/neuralgcm/blob/main/docs/model_vectorization_tutorial.ipynb

Initializes a model and performs a batched rollout by specifying vectorization axes for prognostic and diagnostic variables.

```python
life = build_new_model()
life.initialize_random_processes(rngs)
life.update_dynamic_inputs(dynamic_inputs)

# Note how we use more granular types here to specify vectorization axes, e.g.
# (Prognostic, Diagnostic) rather than using a wider scope parent type
# SimulationVariable that also includes Randomness.
vectorized_life = life.to_vectorized({typing.Prognostic: b, typing.Diagnostic: b})
```

--------------------------------

### Instantiate Parameterizations with Distinct Modules

Source: https://github.com/google-research/neuralgcm/blob/main/docs/simulation_state_components_intro.ipynb

Create separate `DynamicInputSlice` and `DynamicInputFeatures` for radiation and carbon cycle parameterizations to demonstrate a scenario where they do not share modules.

```python
co2_forcing_radiation = dynamic_io.DynamicInputSlice({'co2': cx.Scalar()}, observation_key='atm')
co2_features_radiation = feature_transforms.DynamicInputFeatures(['co2'], co2_forcing_radiation)

co2_forcing_carbon = dynamic_io.DynamicInputSlice({'co2': cx.Scalar()}, observation_key='soil')
co2_features_carbon = feature_transforms.DynamicInputFeatures(['co2'], co2_forcing_carbon)

radiation = RadiationParameterization(co2_features_radiation)
carbon_cycle = CarbonCycleParameterization(co2_features_carbon)

```

--------------------------------

### Helper function for generating model rollouts

Source: https://github.com/google-research/neuralgcm/blob/main/docs/forced_parmeterized_coupled_l96.ipynb

A helper function to generate model rollouts by repeatedly calling `model.observe(query)` and `model.advance()`. It uses `nnx.jit` for performance and `nnx.scan` for efficient iteration over simulation steps.

```python
# Helper function for generating model rollous.

@nnx.jit(static_argnums=(1, 3))
def unroll_model(
    model,
    outer_steps: int,
    query,
    inner_steps: int = 1,
) -> dict[str, dict[str, cx.Field]]:
  """Runs simulation calling `model.observe(query)` every `innfer_steps`."""
  def _inner(model):
    model.advance()

  in_axes = nnx.StateAxes({typing.SimulationVariable: nnx.Carry, ...: None})
  inner_step = nnx.scan(_inner, length=inner_steps, in_axes=in_axes, out_axes=0)

  def _step(model):
    observation = model.observe(query)
    inner_step(model)
    return observation

  unroll = nnx.scan(_step, length=outer_steps, in_axes=in_axes, out_axes=0)
  observations = unroll(model)
  # time = cx.LabeledAxis('timedelta', np.arange(outer_steps) * inner_steps)
  time = cx.LabeledAxis('timedelta', np.arange(outer_steps))
  observations = cx.tag(observations, time)
  return observations
```

--------------------------------

### Initialize simulation state using Coordax Fields

Source: https://github.com/google-research/neuralgcm/blob/main/docs/data_model_intro.ipynb

Sets up coordinate systems (SigmaLevels, LonLatGrid, SphericalHarmonicGrid) and a mapping for spherical harmonics. It then generates a perturbed atmospheric state using `idealized_states.perturbed_jw` and stores it in the `prognostics` Fields object.

```python
sigma = coordinates.SigmaLevels.equidistant(8)
grid = coordinates.LonLatGrid.T42()
ylm_grid = coordinates.SphericalHarmonicGrid.T42()
ylm_map = spherical_transforms.FixedYlmMapping(
    grid, ylm_grid, partition_schema_key=None, mesh=parallelism.Mesh()
)
rng = jax.random.key(3)
prognostics = idealized_states.perturbed_jw(
    ylm_map, sigma, rng, as_nodal=True, temperature_format='absolute'
)
```

--------------------------------

### Perform Batched Rollout

Source: https://github.com/google-research/neuralgcm/blob/main/docs/model_vectorization_tutorial.ipynb

Execute a rollout with dynamic inputs across a batch dimension.

```python
life = build_new_model()
life.initialize_random_processes(rngs)

# Include dynamic input to vectorization specs.
vectorized_life = life.to_vectorized(
    {typing.Prognostic: b, typing.Diagnostic: b, typing.DynamicInput: b}
)

# update_dynamic_inputs now expects batched inputs.
vectorized_life.update_dynamic_inputs(batched_dynamic_inputs)
vectorized_life.assimilate(batched_inputs)
init_observations = vectorized_life.observe(observe_query)

vectorized_life.advance()
vectorized_life.advance()
vectorized_life.advance()

post_advance_observations = vectorized_life.observe(observe_query)
```

```python
ds = xarray.concat(
    [
        init_observations['game']['u'].to_xarray(),
        post_advance_observations['game']['u'].to_xarray(),
    ],
    'step'
)
ds.coords['step'] = np.arange(2) * 3
ds.plot(x='x', y='y', row='step', col='batch')
```

--------------------------------

### Import Libraries

Source: https://github.com/google-research/neuralgcm/blob/main/docs/simulation_state_components_intro.ipynb

Imports core libraries for simulation state management, including JAX, Flax, Coordax, and NeuralGCM modules.

```python
import dataclasses

import coordax as cx
from flax import nnx
import jax
import jax.numpy as jnp
import jax_datetime as jdt
import matplotlib.pyplot as plt
import numpy as np

from neuralgcm.experimental.atmosphere import idealized_states
from neuralgcm.experimental.core import coordinates
from neuralgcm.experimental.core import dynamic_io
from neuralgcm.experimental.core import diagnostics
from neuralgcm.experimental.core import random_processes
from neuralgcm.experimental.core import feature_transforms
from neuralgcm.experimental.core import parallelism
from neuralgcm.experimental.core import module_utils
from neuralgcm.experimental.core import transforms
from neuralgcm.experimental.core import spherical_transforms
from neuralgcm.experimental.core import typing
from neuralgcm.experimental.core import units

import logging
import warnings
warnings.simplefilter('ignore')
logging.getLogger('jax._src.lib.xla_bridge').addFilter(lambda _: False)
```

--------------------------------

### Run Vectorized Model Methods

Source: https://github.com/google-research/neuralgcm/blob/main/docs/model_vectorization_tutorial.ipynb

Basic workflow for assimilating data, observing state, and advancing a vectorized model.

```python
vectorized_life.assimilate(batched_inputs)
init_observations = vectorized_life.observe(observe_query)
vectorized_life.advance()
post_advance_observations = vectorized_life.observe(observe_query)
```

```python
ds = xarray.concat(
    [
        init_observations['game']['u'].to_xarray(),
        post_advance_observations['game']['u'].to_xarray(),
    ],
    'step'
)
ds.plot(x='x', y='y', row='step', col='batch')
```

```python
observe_fn = lambda model: model.observe(observe_query | diagnostic_query)
unroll_fn = functools.partial(unroll_model, observe_fn=observe_fn)
observations = vectorized_life.run_vectorized(unroll_fn, b, 6, life.timestep * 2)
xarray_utils.fields_to_xarray(observations['game']).u.plot(x='x', y='y', col='timedelta', row='batch')
```

```python
xarray_utils.fields_to_xarray(observations['diagnostic']).n_alive.plot(x='timedelta', hue='batch', marker='o')
```

--------------------------------

### Initialize random process modules

Source: https://github.com/google-research/neuralgcm/blob/main/docs/simulation_state_components_intro.ipynb

Configure NormalUncorrelated or GaussianRandomField processes using grid definitions and random number generator seeds.

```python
sim_units = units.get_si_units()
grid = coordinates.LonLatGrid.T42()
ylm_grid = coordinates.SphericalHarmonicGrid.T42()
ylm_map = spherical_transforms.FixedYlmMapping(
    grid,
    ylm_grid,
    partition_schema_key=None,
    mesh=parallelism.Mesh(),
    radius=sim_units.radius,
)
uncorrelated = random_processes.NormalUncorrelated(
    grid, mean=0.0, std=1.0, rngs=nnx.Rngs(0)
)
gaussian_random_process = random_processes.GaussianRandomField(
    ylm_map, 3600, sim_units, '10 hours', '500 km', 1.0, rngs=nnx.Rngs(1)
)
```

--------------------------------

### Fine-tune models with JAX gradients

Source: https://context7.com/google-research/neuralgcm/llms.txt

Demonstrates defining a differentiable loss function and using Optax to update model parameters via JAX transformations.

```python
import jax
import optax
import neuralgcm

checkpoint = neuralgcm.demo.load_checkpoint_tl63_stochastic()
model = neuralgcm.PressureLevelModel.from_checkpoint(checkpoint)
dataset = neuralgcm.demo.load_data(model.data_coords)

inputs = model.inputs_from_xarray(dataset.isel(time=0))
forcings = model.forcings_from_xarray(dataset.isel(time=0))

# Define a differentiable loss function
@jax.jit
def compute_loss(model, inputs, forcings, rng):
    encoded = model.encode(inputs, forcings, rng_key=rng)
    restored = model.decode(encoded, forcings)
    return abs(inputs['temperature'] - restored['temperature']).mean()

# Compute loss and gradients
loss = compute_loss(model, inputs, forcings, jax.random.key(0))
print(f"Initial loss: {loss}")

# Fine-tune with Optax optimizer
optimizer = optax.adam(1e-3)
opt_model = model
opt_state = optimizer.init(model)

for i in range(5):
    loss, grads = jax.value_and_grad(compute_loss)(
        opt_model, inputs, forcings, jax.random.key(0)
    )
    updates, opt_state = optimizer.update(grads, opt_state)
    opt_model = optax.apply_updates(opt_model, updates)
    print(f"Step {i}, loss: {loss}")
```

--------------------------------

### Perform forecast with forecast_steps

Source: https://github.com/google-research/neuralgcm/blob/main/docs/inference_model_api_intro.ipynb

Combines initial assimilation and rollout steps into a single forecast operation.

```python
another_final_state, another_trajectory = api.forecast_steps(
    l96_inference_model,
    inputs=inputs,
    timedelta=timedelta_between_saves,
    steps=100,
    query={'slow': {'x': k}}
)
another_datasets = xarray_utils.nested_fields_to_xarray(another_trajectory)
another_datasets['slow'].coords['timedelta'] = another_datasets['slow'].timedelta / np.timedelta64(1, 's')
another_datasets['slow'].x.plot.imshow(x='k', y='timedelta')
```

--------------------------------

### Initialize and run NeuralGCM forecast

Source: https://github.com/google-research/neuralgcm/blob/main/docs/checkpoint_modifications.ipynb

Initializes the NeuralGCM model state and runs a forecast for a specified number of steps, using persistence for forcing variables.

```python
inner_steps = 1  # save model outputs once every 24 hours
outer_steps = 30 * 1 // inner_steps  # total of 30 days
timedelta = np.timedelta64(24, 'h') * inner_steps
times = (np.arange(outer_steps) * inner_steps)  # time axis in hours

# initialize model state
inputs = model.inputs_from_xarray(eval_era5.isel(time=0))
input_forcings = model.forcings_from_xarray(eval_era5.isel(time=0))
rng_key = jax.random.key(42)  # optional for deterministic models
initial_state = model.encode(inputs, input_forcings, rng_key)

# use persistence for forcing variables (SST and sea ice cover)
all_forcings = model.forcings_from_xarray(eval_era5.head(time=1))

# make forecast
final_state, predictions = model.unroll(
    initial_state,
    all_forcings,
    steps=outer_steps,
    timedelta=timedelta,
    start_with_input=True,
)
predictions_ds_fix_lsp = model.data_to_xarray(predictions, times=times)
```

--------------------------------

### Import Apache Beam Libraries

Source: https://github.com/google-research/neuralgcm/blob/main/docs/inference_model_api_intro.ipynb

Imports the necessary Apache Beam modules for pipeline creation and option configuration.

```python
import apache_beam as beam
from apache_beam.options.pipeline_options import PipelineOptions
```

--------------------------------

### Instantiate InferenceModel

Source: https://github.com/google-research/neuralgcm/blob/main/docs/inference_model_api_intro.ipynb

Initializes an inference model from a Lorenz-96 model instance using the from_model_api classmethod.

```python
k = cx.LabeledAxis('k', np.arange(36))
j = cx.LabeledAxis('j', np.arange(8))

l96_model = lorenz96.Lorenz96WithTwoScales(k, j, forcing=nnx.Param(10.0))
l96_inference_model = api.InferenceModel.from_model_api(l96_model)
```

--------------------------------

### Configure dynamic input modules

Source: https://github.com/google-research/neuralgcm/blob/main/docs/simulation_state_components_intro.ipynb

Initialize a DynamicInputSlice to manage time-indexed external forcings for model components.

```python
grid = coordinates.LonLatGrid.T42()
dynamic_input = dynamic_io.DynamicInputSlice(
    keys_to_coords={'a': grid},  # specifies data coords (modulo the time axis).
    observation_key='abc',  # specifies data_source key associated with data.
)
```

--------------------------------

### Visualize global mean log surface pressure

Source: https://github.com/google-research/neuralgcm/blob/main/docs/checkpoint_modifications.ipynb

Plots the comparison of global mean log surface pressure between fixed and non-fixed model runs.

```python
import matplotlib.pyplot as plt
plt.plot(log_sp_spherical_harmonics_fix_lsp[:,0,0], label='Fixed global mean')
plt.plot(log_sp_spherical_harmonics[:,0,0], label='No Fix')
plt.legend()
```

--------------------------------

### Import necessary libraries for NeuralGCM tutorial

Source: https://github.com/google-research/neuralgcm/blob/main/docs/data_model_intro.ipynb

Imports required libraries including `flax.nnx`, `jax`, `matplotlib`, `numpy`, `coordax`, and specific modules from `neuralgcm` for atmosphere, core functionalities, and utilities.

```python
import warnings
warnings.simplefilter('ignore')


from flax import nnx
import jax
import matplotlib.pyplot as plt
import numpy as np

import coordax as cx
from neuralgcm.experimental.atmosphere import idealized_states
from neuralgcm.experimental.core import coordinates
from neuralgcm.experimental.core import parallelism
from neuralgcm.experimental.core import transforms
from neuralgcm.experimental.core import feature_transforms
from neuralgcm.experimental.core import spherical_transforms
from neuralgcm.experimental.core import xarray_utils
```

--------------------------------

### Assimilate Inputs

Source: https://github.com/google-research/neuralgcm/blob/main/docs/inference_model_api_intro.ipynb

Prepares inputs and assimilates them into the inference model to retrieve the initial simulation state.

```python
# Preparing inputs for model initialization
kj = cx.compose_coordinates(k, j)
rng = np.random.RandomState(0)
x_init = abs(10 * np.sin(np.linspace(0, 13 * 2*np.pi, k.sizes['k'])))
x_init[0] += 0.01
t0 = jdt.Datetime.from_isoformat('2000-01-01')
inputs = {
    'slow': {'x': cx.field(x_init, k), 'time': cx.field(t0)},
    'fast': {'y': cx.field(rng.uniform(size=kj.shape, low=-0.5, high=0.5), kj)},
}
# assimilate expects at least a dummy timedelta dimension, we add it here.
t_del = coordinates.TimeDelta(np.zeros(1) * np.timedelta64(1, 'h'))
add_td = lambda f: f.broadcast_like(cx.compose_coordinates(t_del, f.coordinate))
inputs = jax.tree.map(add_td, inputs, is_leaf=cx.is_field)

# Assimilating inputs returns current simulation state.
sim_state = l96_inference_model.assimilate(inputs)
```

```python
sim_state
```

--------------------------------

### Instantiate and Share Dynamic Input Features

Source: https://github.com/google-research/neuralgcm/blob/main/docs/simulation_state_components_intro.ipynb

Instantiate `DynamicInputSlice` and `DynamicInputFeatures`, then share the `co2_features` instance across `RadiationParameterization` and `CarbonCycleParameterization` when creating the `CombinedParameterization`.

```python
t0 = jdt.to_datetime('2000-01-01')
timedelta = coordinates.TimeDelta(np.arange(12) * np.timedelta64(30, 'D'))
time = t0 + timedelta.fields['timedelta']
values = cx.field(
    370*np.exp(0.0063 * timedelta.deltas / np.timedelta64(365, 'D')), timedelta
)
co2_forcing_values = {'global': {'co2': values, 'time': time}}

co2_forcing = dynamic_io.DynamicInputSlice({'co2': cx.Scalar()}, observation_key='global')
co2_features = feature_transforms.DynamicInputFeatures(['co2'], co2_forcing)
radiation = RadiationParameterization(co2_features)
carbon_cycle = CarbonCycleParameterization(co2_features)  # same co2_features.
combined = CombinedParameterization(radiation, carbon_cycle)

```

--------------------------------

### Initialize Model State and Make Forecast

Source: https://github.com/google-research/neuralgcm/blob/main/docs/inference_demo.ipynb

Initialize the NeuralGCM model state using the first time step of the prepared ERA5 data and generate a forecast for a specified number of steps. Persistence is used for forcing variables.

```python
inner_steps = 24  # save model outputs once every 24 hours
outer_steps = 4 * 24 // inner_steps  # total of 4 days
timedelta = np.timedelta64(1, 'h') * inner_steps
times = (np.arange(outer_steps) * inner_steps)  # time axis in hours

# initialize model state
inputs = model.inputs_from_xarray(eval_era5.isel(time=0))
input_forcings = model.forcings_from_xarray(eval_era5.isel(time=0))
rng_key = jax.random.key(42)  # optional for deterministic models
initial_state = model.encode(inputs, input_forcings, rng_key)

# use persistence for forcing variables (SST and sea ice cover)
all_forcings = model.forcings_from_xarray(eval_era5.head(time=1))

# make forecast
final_state, predictions = model.unroll(
    initial_state,
    all_forcings,
    steps=outer_steps,
    timedelta=timedelta,
    start_with_input=True,
)
predictions_ds = model.data_to_xarray(predictions, times=times)
```

--------------------------------

### Initialize model state for forecast

Source: https://github.com/google-research/neuralgcm/blob/main/docs/checkpoint_modifications.ipynb

Initializes the model state using data from the first time step of the preprocessed ERA5 data. This is a prerequisite for running the forecast.

```python
inner_steps = 1  # save model outputs once every 24 hours
outer_steps = 30 * 1 // inner_steps  # total of 30 days
timedelta = np.timedelta64(24, 'h') * inner_steps
times = (np.arange(outer_steps) * inner_steps)  # time axis in hours

# initialize model state
inputs = model.inputs_from_xarray(eval_era5.isel(time=0))
input_forcings = model.forcings_from_xarray(eval_era5.isel(time=0))
rng_key = jax.random.key(42)  # optional for deterministic models
initial_state = model.encode(inputs, input_forcings, rng_key)
```