### Setup Conda Environment and Install JAX with CUDA

Source: https://tum-pbs.github.io/apebench

Create a new Conda environment with Python 3.12, activate it, install JAX with CUDA 12 support, and then install APEBench.

```bash
conda create -n apebench python=3.12 -y
conda activate apebench
pip install -U "jax[cuda12]"
pip install apebench
```

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

### Install APEBench

Source: https://tum-pbs.github.io/apebench

Install the APEBench package using pip. Requires Python 3.10+.

```bash
pip install apebench
```

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

### __post_init__ Method

Source: https://tum-pbs.github.io/apebench/api/scenarios/normalized/nonlinear

Performs post-initialization setup for the scenario. It includes a check to ensure the simulation is run with exactly one spatial dimension and sets the diffusion alphas. It then calls the parent class's post-initialization method.

```python
def __post_init__(self):
    if self.num_spatial_dims != 1:
        raise ValueError(
            "Conservative Kuramoto-Sivashinsky is only defined for 1 spatial dimension. Check out the non-conservative version for 2d."
        )
    self.alphas = (0.0, 0.0, self.diffusion_alpha, 0.0, self.hyp_diffusion_alpha)
    super().__post_init__()

```

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

### Quickstart: Train ConvNet for Advection Emulation

Source: https://tum-pbs.github.io/apebench

Train a ConvNet to emulate 1D advection, visualize training loss, test error metric rollout, and a sample rollout. Requires seaborn, matplotlib, and numpy.

```python
import apebench
import seaborn as sns
import matplotlib.pyplot as plt
import numpy as np

advection_scenario = apebench.scenarios.difficulty.Advection()

data, trained_nets = advection_scenario(
    task_config="predict",
    network_config="Conv;26;10;relu",
    train_config="one",
    num_seeds=3,
)

data_loss = apebench.melt_loss(data)
data_metrics = apebench.melt_metrics(data)
data_sample_rollout = apebench.melt_sample_rollouts(data)

fig, axs = plt.subplots(1, 3, figsize=(13, 3))

sns.lineplot(data_loss, x="update_step", y="train_loss", ax=axs[0])
axs[0].set_yscale("log")
axs[0].set_title("Training loss")

sns.lineplot(data_metrics, x="time_step", y="mean_nRMSE", ax=axs[1])
axs[1].set_ylim(-0.05, 1.05)
axs[1].set_title("Metric rollout")

axs[2].imshow(
    np.array(data_sample_rollout["sample_rollout"][0])[:, 0, :].T,
    origin="lower",
    aspect="auto",
    vmin=-1,
    vmax=1,
    cmap="RdBu_r",
)
axs[2].set_xlabel("time")
axs[2].set_ylabel("space")
axs[2].set_title("Sample rollout")

plt.show()
```

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

### Get Training Initial Conditions (Reference) - Python

Source: https://tum-pbs.github.io/apebench/api/scenarios/base

Produces the reference training initial condition set using the scenario's attributes and a PRNG key derived from the training seed.

```python
def get_train_ic_set(
        self,
    ) -> Float[Array, "num_train_samples num_channels *num_points"]:
        """
        Use the attributes to produce the reference training initial condition set.
        """
        return self._produce_ic_set(
            stepper=self.get_ref_stepper(),
            num_samples=self.num_train_samples,
            key=jax.random.PRNGKey(self.train_seed),
        )
```

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

### Python Configuration for APE-Bench Experiment

Source: https://tum-pbs.github.io/apebench/using_cli

Define experiment configurations, including network architectures and advection parameters, in a Python file. This example sets up training for nonlinear emulators of the 1D advection equation.

```python
"""
Trains nonlinear emulators for the (linear) 1D advection equation under varying
difficulty in terms of the `advction_gamma` (=CFL).
"""

CONFIGS = [
    {
        "scenario": "diff_adv",
        "task": "predict",
        "net": net,
        "train": "one",
        "start_seed": 0,
        "num_seeds": 10,
        "advection_gamma": advection_gamma,
    }
    for net in [
        "Conv;34;10;relu",  # 31'757 params, 11 receptive field per direction
        "UNet;12;2;relu",  # 27'193 params, 29 receptive field per direction
        "Res;26;8;relu",  # 32'943 params, 16 receptive field per direction
        "FNO;12;18;4;gelu",  # 32'527 params, inf receptive field per direction
        "Dil;2;32;2;relu",  # 31'777 params, 20 receptive field per direction
    ]
    for advection_gamma in [
        0.5,
        2.5,
        10.5,
    ]
]


```

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

### Get Training Initial Conditions (Coarse) - Python

Source: https://tum-pbs.github.io/apebench/api/scenarios/base

Produces training initial conditions using the coarse stepper, based on scenario attributes and a PRNG key from the training seed.

```python
def get_train_ic_set_coarse(
        self,
    ) -> Float[Array, "num_train_samples num_channels *num_points"]:
        """
        Use the attributes to produce training initial conditions with the coarse stepper instead.
        """
        return self._produce_ic_set(
            stepper=self.get_coarse_stepper(),
            num_samples=self.num_train_samples,
            key=jax.random.PRNGKey(self.train_seed),
        )
```

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

### Get Testing Initial Conditions (Reference) - Python

Source: https://tum-pbs.github.io/apebench/api/scenarios/base

Produces the reference testing initial condition set using the scenario's attributes and a PRNG key derived from the testing seed.

```python
def get_test_ic_set(
        self,
    ) -> Float[Array, "num_test_samples num_channels *num_points"]:
        """
        Use the attributes to produce the reference testing initial condition set.
        """
        return self._produce_ic_set(
            stepper=self.get_ref_stepper(),
            num_samples=self.num_test_samples,
            key=jax.random.PRNGKey(self.test_seed),
        )
```

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

### Initialize CNN Model, Optimizer, and Training Step

Source: https://tum-pbs.github.io/apebench/examples/benchmark_flax_models_with_linen

Initializes the CNN model, sets up the Adam optimizer, and defines the training step function. This includes model initialization, optimizer state setup, and the core logic for updating model parameters based on the one-step supervised loss.

```python
cnn = CNN(width=34, depth=10)
params = cnn.init(jax.random.PRNGKey(0), train_windows[:, 0])
optimizer = optax.adam(1e-4)
opt_state = optimizer.init(params)


def one_step_supervised_loss(params, batch):
    inputs, targets = batch[:, 0], batch[:, 1]
    predictions = cnn.apply(params, inputs)
    return jnp.mean((predictions - targets) ** 2)


@jax.jit
def train_step(params, state, batch):
    loss, grad = jax.value_and_grad(one_step_supervised_loss)(params, batch)
    updates, new_opt_state = optimizer.update(grad, state)
    new_params = optax.apply_updates(params, updates)
    return new_params, new_opt_state, loss
```

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

### Get Testing Initial Conditions (Coarse) - Python

Source: https://tum-pbs.github.io/apebench/api/scenarios/base

Produces testing initial conditions using the coarse stepper, based on scenario attributes and a PRNG key from the testing seed.

```python
def get_test_ic_set_coarse(
        self,
    ) -> Float[Array, "num_test_samples num_channels *num_points"]:
        """
        Use the attributes to produce testing initial conditions with the coarse stepper instead.
        """
        return self._produce_ic_set(
            stepper=self.get_coarse_stepper(),
            num_samples=self.num_test_samples,
            key=jax.random.PRNGKey(self.test_seed),
        )
```

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

### Get Reference Stepper

Source: https://tum-pbs.github.io/apebench/api/scenarios/base

Abstract method to produce the reference stepper for the scenario. This method must be implemented by subclasses.

```python
@abstractmethod
def get_ref_stepper(self) -> BaseStepper:
    """
    Produces the reference stepper for the scenario.
    """
    pass
```

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

### Get Coarse Stepper

Source: https://tum-pbs.github.io/apebench/api/scenarios/base

Abstract method to produce the coarse stepper for the scenario. This method must be implemented by subclasses.

```python
@abstractmethod
def get_coarse_stepper(self) -> BaseStepper:
    """
    Produces the coarse stepper for the scenario.
    """
    pass
```

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

### Import necessary libraries for Flax and APEBench

Source: https://tum-pbs.github.io/apebench/examples/benchmark_flax_models_with_linen

Imports essential libraries for numerical operations, JAX, Flax, APEBench, and visualization tools. Ensure these are installed in your environment.

```python
import apebench

import numpy as np
import jax
import jax.numpy as jnp
from flax import linen as nn
from tqdm.autonotebook import tqdm
import optax
import matplotlib.pyplot as plt
import seaborn as sns
from scipy import stats
```

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

### Load and Test JAX Rollout

Source: https://tum-pbs.github.io/apebench/examples/benchmark_pytorch_models

Loads a neural rollout from a .npy file and performs tests using JAX. Ensure RUN_PYTORCH is set to False and the environment has JAX/APEBench installed.

```python
if not RUN_PYTORCH:
    neural_rollout = jnp.load("neural_rollout.npy")
    print(neural_rollout.shape)
```

```python
if not RUN_PYTORCH:
    test_dict = advection_scenario.perform_tests_on_rollout(neural_rollout)
```

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

### Import necessary libraries for APEBench benchmarking

Source: https://tum-pbs.github.io/apebench/examples/benchmark_flax_models_with_nnx

Imports essential libraries including apebench, numpy, jax, flax.nnx, optax, and plotting tools. Ensure these are installed in your environment.

```python
import apebench

import numpy as np
import jax
import jax.numpy as jnp
from flax import nnx
from tqdm.autonotebook import tqdm
import optax
import matplotlib.pyplot as plt
import seaborn as sns
from scipy import stats
```

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

### Get Metric Functions Dictionary

Source: https://tum-pbs.github.io/apebench/api/scenarios/base

Constructs and returns a dictionary of metric functions based on the `report_metrics` configuration string. Each metric is initialized using its corresponding constructor from the `metric_dict`.

```python
metric_fn_dict = {}

for metric_config in self.report_metrics.split(","):
    metric_args = metric_config.split(";")
    metric_name = metric_args[0]
    metric_constructor = metric_dict[metric_name]
    metric_fn = metric_constructor(metric_config)
    metric_fn_dict[metric_config] = metric_fn
```

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

### Get Trainer Configuration - Python

Source: https://tum-pbs.github.io/apebench/api/scenarios/base

Configures and returns a trainer based on the provided train_config string. Supports supervised and diverted chain training with various rollout steps and cut options.

```python
def get_trainer(self, *, train_config: str) -> tx.GeneralTrainer:
    """
    Expects a str of the defined interface for study. In the default
    configuration, it could for instance accept:

    'sup-03' for supervised rollout trainig with three rollout steps.

    Currently, the three major categories are available:

    - 'one' for one step supervised training
    - 'sup-XX' for supervised training with XX rollout steps
    - 'div-XX' for diverted chain training with XX rollout steps
    """
    train_trjs = self.get_train_data()

    # Needed for diverted chain training
    ref_stepper = self.get_ref_stepper()
    train_args = train_config.split(";")

    optimizer = self.get_optimizer()

    callback_fn = self.get_callback_fn()

    if train_args[0].lower() == "one":
        trainer = tx.trainer.SupervisedTrainer(
            train_trjs,
            optimizer=optimizer,
            num_training_steps=self.num_training_steps,
            batch_size=self.batch_size,
            num_rollout_steps=1,
            cut_bptt=False,
            time_level_weights=None,
            callback_fn=callback_fn,
        )
    elif train_args[0].lower() == "sup":
        num_rollout_steps = int(train_args[1])
        if len(train_args) > 2:
            cut_bptt = train_args[2].lower() == "true"
        else:
            cut_bptt = False
        trainer = tx.trainer.SupervisedTrainer(
            train_trjs,
            optimizer=optimizer,
            num_training_steps=self.num_training_steps,
            batch_size=self.batch_size,
            num_rollout_steps=num_rollout_steps,
            cut_bptt=cut_bptt,
            time_level_weights=None,
            callback_fn=callback_fn,
        )
    elif train_args[0].lower() == "div":
        num_rollout_steps = int(train_args[1])
        if len(train_args) > 2:
            cut_bptt = train_args[2].lower() == "true"
        else:
            cut_bptt = False
        if len(train_args) > 3:
            cut_div_chain = train_args[3].lower() == "true"
        else:
            cut_div_chain = False
        trainer = tx.trainer.DivertedChainBranchOneTrainer(
            train_trjs,
            ref_stepper=ref_stepper,
            optimizer=optimizer,
            num_training_steps=self.num_training_steps,
            batch_size=self.batch_size,
            num_rollout_steps=num_rollout_steps,
            cut_bptt=cut_bptt,
            cut_div_chain=cut_div_chain,
            time_level_weights=None,
            callback_fn=callback_fn,
        )
    else:
        raise ValueError(f"Unknown training argument '{train_config}'")

    return trainer
```

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

### Setup Environment Flags and Imports for PyTorch/JAX

Source: https://tum-pbs.github.io/apebench/examples/benchmark_pytorch_models

Sets up a flag to control PyTorch execution and imports necessary libraries for either PyTorch or JAX environments. This is crucial for managing dependencies and data transfer between the two frameworks.

```python
RUN_PYTORCH = False

from tqdm.autonotebook import tqdm
import matplotlib.pyplot as plt

if RUN_PYTORCH:
    import torch
    import numpy as np
else:
    import apebench
    import numpy as np
    import jax.numpy as jnp
    import seaborn as sns
    from scipy import stats
```

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

### Get Optimizer Configuration - Python

Source: https://tum-pbs.github.io/apebench/api/scenarios/base

Retrieves the optimizer and learning rate scheduler configuration based on scenario settings. Parses optimizer and scheduler arguments from a semicolon-separated string.

```python
def get_optimizer(self) -> optax.GradientTransformation:
        """
        Returns the optimizer used in the scenario.
        """
        optim_args = self.optim_config.split(";")
        optimizer_name = optim_args[0]
        num_training_steps = int(optim_args[1])
        scheduler_args = optim_args[2:]
        scheduler_name = scheduler_args[0]

        lr_scheduler = lr_scheduler_dict[scheduler_name.lower()](
            ";".join(scheduler_args), num_training_steps
        )
        optimizer = optimizer_dict[optimizer_name.lower()](self.optim_config)(
            lr_scheduler
        )

        return optimizer
```

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

### Get Activation Function - Python

Source: https://tum-pbs.github.io/apebench/api/scenarios/base

Parses a string to a callable activation function. The function name is extracted from the part before the first semicolon.

```python
def get_activation(
    self,
    activation: str,
) -> Callable:
    """
    Parse a string to a callable activation function.
    """
    activation_fn_name = activation.split(";")[0]
    activation_fn = activation_fn_dict[activation_fn_name.lower()](activation)
    return activation_fn
```

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

### Display APE-Bench CLI Help

Source: https://tum-pbs.github.io/apebench/using_cli

Run this command to see all available options and arguments for the apebench CLI.

```bash
apebench --help

```

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

### Get Reference Stepper for Convection Scenario

Source: https://tum-pbs.github.io/apebench/api/scenarios/difficulty/generic

Returns the reference stepper for the Convection scenario, configured with the class-level gammas and convection_delta. This represents the full-fidelity simulation setup.

```python
def get_ref_stepper(self):
    return self._build_stepper(self.gammas, self.convection_delta)
```

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

### APE-Bench CLI Usage Help

Source: https://tum-pbs.github.io/apebench/using_cli

This output details the command-line arguments and options available for running an apebench experiment, including positional arguments like 'experiment' and optional flags for GPU, seeds, and metric melting.

```bash
usage: apebench [-h] [--gpu GPU] [--start_seed START_SEED] [--dont_melt_metrics] [--dont_melt_loss] [--melt_sample_rollouts] experiment

Run a study by pointing to its config list

positional arguments:
  experiment            The file containig the config list

options:
  -h, --help            show this help message and exit
  --gpu GPU             The GPU to use for the experiment
  --start_seed START_SEED
                        The seed to start the experiment from (helpful for seed-statistics across multiple GPUs)
  --dont_melt_metrics   Do not melt the metrics across the runs in the experiment
  --dont_melt_loss      Do not melt the train loss across the runs in the experiment
  --melt_sample_rollouts
                        Melt the sample rollouts across the runs in the experiment

```

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

### Scrape Data with Modified Resolution

Source: https://tum-pbs.github.io/apebench/examples/scrape_datasets

Modify the resolution of the scraped data by passing the 'num_points' keyword argument to the scrape_data_and_metadata function. This example shows how to get the shape of the first returned element (training data) with a reduced resolution.

```python
apebench.scraper.scrape_data_and_metadata(scenario="diff_ks", num_points=64)[0].shape

```

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

### Get Optimizer for Scenario

Source: https://tum-pbs.github.io/apebench/api/scenarios/base

Retrieves the configured optimizer for the scenario. It parses optimizer and learning rate scheduler configurations from a string to instantiate and return the appropriate optax gradient transformation.

```python
def get_optimizer(self) -> optax.GradientTransformation:
    """
    Returns the optimizer used in the scenario.
    """
    optim_args = self.optim_config.split(";")
    optimizer_name = optim_args[0]
    num_training_steps = int(optim_args[1])
    scheduler_args = optim_args[2:]
    scheduler_name = scheduler_args[0]

    lr_scheduler = lr_scheduler_dict[scheduler_name.lower()](
        ";".join(scheduler_args), num_training_steps
    )
    optimizer = optimizer_dict[optimizer_name.lower()](self.optim_config)(
        lr_scheduler
    )

    return optimizer

```

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

### Display APE-Bench Metrics in Jupyter Notebook

Source: https://tum-pbs.github.io/apebench/using_cli

Use pandas and seaborn to read and visualize training metrics from the 'melted/train_loss.csv' file. This example plots the mean Normalized Root Mean Square Error (nRMSE) across different advection gammas and network architectures.

```python
import apebench
import pandas as pd
import seaborn as sns

train_df = pd.read_csv("melted/train_loss.csv")
train_df = apebench.read_in_kwargs(train_df)
facet = sns.relplot(
    data=train_df,
    x="time_step",
    y="mean_nRMSE",
    hue="net",
    kind="line",
    col="advection_gamma",
    errorbar=("pi", 50),
    estimator="median",
)
for ax in facet.axes.flat:
    ax.set_ylim(-0.05, 1.05)


```

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

### Example Metadata Structure

Source: https://tum-pbs.github.io/apebench/examples/scrape_datasets

This is an example of the metadata structure returned by the scraping function.

```python
{
 'name': '1d_diff_ks',
 'info': {
  'num_spatial_dims': 1,
  'num_points': 160,
  'num_channels': 1,
  'ic_config': 'fourier;5;true;true',
  'num_warmup_steps': 500,
  'num_train_samples': 50,
  'train_temporal_horizon': 50,
  'train_seed': 0,
  'num_test_samples': 30,
  'test_temporal_horizon': 200,
  'test_seed': 773,
  'optim_config': 'adam;10_000;warmup_cosine;0.0;1e-3;2_000',
  'batch_size': 20,
  'num_trjs_returned': 1,
  'record_loss_every': 100,
  'vlim': (-6.5, 6.5),
  'report_metrics': 'mean_nRMSE,mean_correlation',
  'callbacks': '',
  'gammas': (0.0, 0.0, -1.2, 0.0, -15.0),
  'deltas': (0.0, 0.0, -6.0),
  'num_substeps': 1,
  'coarse_proportion': 0.5,
  'order': 2,
  'dealiasing_fraction': 0.6666666666666666,
  'num_circle_points': 16,
  'circle_radius': 1.0,
  'gradient_norm_delta': -6.0,
  'diffusion_gamma': -1.2,
  'hyp_diffusion_gamma': -15.0
 }
}

```

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

### Get Coarse Training Data

Source: https://tum-pbs.github.io/apebench/api/scenarios/base

Produces training data using the coarse stepper. This is useful for training models that operate on a coarser resolution or with a simplified dynamics model.

```python
def get_train_data_coarse(
    self,
) -> Float[
    Array, "num_train_samples train_temporal_horizon+1 num_channels *num_points"
]:
    """
    Use the attributes to produce training data with the coarse stepper instead.
    """
    return self.produce_data(
        stepper=self.get_coarse_stepper(),
        num_samples=self.num_train_samples,
        num_warmup_steps=self.num_warmup_steps,
        temporal_horizon=self.train_temporal_horizon,
        key=jax.random.PRNGKey(self.train_seed),
    )
```

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

### Get Initial Condition Generator

Source: https://tum-pbs.github.io/apebench/api/scenarios/base

This method is used to generate initial conditions for simulations. It supports various generation methods like Fourier series, diffused noise, Gaussian random fields, and clamping, configurable via the `ic_config` attribute.

```python
def get_ic_generator(self) -> BaseRandomICGenerator:
    """
    Overwrite for custom initial condition generation.

    Uses the `ic_config` attribute to determine the type of initial
    condition generation.

    Allows for the following options:

    - `fourier;CUTOFF;ZERO_MEAN;MAX_ONE` for a truncated Fourier series
        with CUTOFF (int) number of modes, ZERO_MEAN (bool) for zero mean,
        and MAX_ONE (bool) for having the initial condition being at max in
        (-1, 1) but not clamped to it
    - `diffused;INTENSITY;ZERO_MEAN;MAX_ONE` for a diffused noise with
        INTENSITY (float) for the intensity, ZERO_MEAN (bool) for zero mean,
        and MAX_ONE (bool) for having the initial condition being at max in
        (-1, 1) but not clamped to it
    - `grf;POWERLAW_EXPONENT;ZERO_MEAN;MAX_ONE` for a Gaussian random
        field with POWERLAW_EXPONENT (float) for the powerlaw exponent,
        ZERO_MEAN (bool) for zero mean, and MAX_ONE (bool) for having the
        initial condition being at max in (-1, 1) but not clamped to it
    - `clamp;LOWER_BOUND;UPPER_BOUND;CONFIG` for clamping the
        configuration to the range of LOWER_BOUND (float) to UPPER_BOUND
        (float) and then using the configuration CONFIG for the generation
        of the initial condition
    """

    def _get_single_channel(config):
        ic_name = config.split(";")[0]
        ic_gen = ic_dict[ic_name.lower()](config, self.num_spatial_dims)
        return ic_gen

    ic_args = self.ic_config.split(";")
    if ic_args[0].lower() == "clamp":
        lower_bound = float(ic_args[1])
        upper_bound = float(ic_args[2])

        ic_gen = _get_single_channel(";".join(ic_args[3:]))
        ic_gen = ex.ic.ClampingICGenerator(
            ic_gen,
            limits=(lower_bound, upper_bound),
        )
    else:
        ic_gen = _get_single_channel(self.ic_config)

    multi_channel_ic_gen = ex.ic.RandomMultiChannelICGenerator(
        [
            ic_gen,
        ]
        * self.num_channels
    )

    return multi_channel_ic_gen
```

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

### FirstFour Scenario __post_init__ Method

Source: https://tum-pbs.github.io/apebench/api/scenarios/difficulty/generic

Initializes the gammas tuple for the FirstFour scenario using its specific gamma attributes.

```python
def __post_init__(self):
    self.gammas = (
        0.0,
        self.advection_gamma,
        self.diffusion_gamma,
        self.dispersion_gamma,
        self.hyp_diffusion_gamma,
    )
```

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

### Example Data Shape Output

Source: https://tum-pbs.github.io/apebench/examples/scrape_datasets

This is an example of the output shape for training and test data.

```python
((50, 51, 1, 160), (30, 201, 1, 160))

```

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

### Get Reference Training Data

Source: https://tum-pbs.github.io/apebench/api/scenarios/base

Produces reference training data using the scenario's attributes and the reference stepper. This method is used to generate the ground truth data for training.

```python
def get_train_data(
    self,
) -> Float[
    Array, "num_train_samples train_temporal_horizon+1 num_channels *num_points"
]:
    """
    Use the attributes to produce the reference training data.
    """
    return self.produce_data(
        stepper=self.get_ref_stepper(),
        num_samples=self.num_train_samples,
        num_warmup_steps=self.num_warmup_steps,
        temporal_horizon=self.train_temporal_horizon,
        key=jax.random.PRNGKey(self.train_seed),
    )
```

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

### Example Modified Data Shape Output

Source: https://tum-pbs.github.io/apebench/examples/scrape_datasets

This is an example of the output shape for training data when the resolution is reduced.

```python
(50, 51, 1, 64)

```

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

### Run APE-Bench Experiment from CLI

Source: https://tum-pbs.github.io/apebench/using_cli

Execute a defined APE-Bench experiment by providing the Python configuration file to the apebench CLI. This command initiates the training and evaluation process.

```bash
apebench test_study.py

```

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

### Get Test Data

Source: https://tum-pbs.github.io/apebench/api/scenarios/base

Produces reference testing data using the scenario's attributes and the reference stepper. This method is used to generate the ground truth data for testing.

```python
def get_test_data(
    self,
) -> Float[
    Array, "num_test_samples test_temporal_horizon+1 num_channels *num_points"
]:
    """
    Use the attributes to produce the reference testing data.
    """
    return self.produce_data(
        stepper=self.get_ref_stepper(),
        num_samples=self.num_test_samples,
        num_warmup_steps=self.num_warmup_steps,
        temporal_horizon=self.test_temporal_horizon,
        key=jax.random.PRNGKey(self.test_seed),
    )
```

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

### Get Nonlinear Scenario Name

Source: https://tum-pbs.github.io/apebench/api/scenarios/physical/generic

Returns the name of the nonlinear physical scenario based on the number of spatial dimensions.

```python
def get_scenario_name(self) -> str:
    return f"{self.num_spatial_dims}d_phy_nonlin"
```

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

### BurgersSingleChannel __post_init__ Method

Source: https://tum-pbs.github.io/apebench/api/scenarios/difficulty/nonlinear

Initializes gammas and deltas for the BurgersSingleChannel scenario. Uses diffusion_gamma and convection_sc_delta.

```python
def __post_init__(self):
    self.gammas = (0.0, 0.0, self.diffusion_gamma, 0.0, 0.0)
    self.deltas = (0.0, self.convection_sc_delta, 0.0)

    super().__post_init__()
```

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

### BurgersSingleChannel Scenario __post_init__ Method

Source: https://tum-pbs.github.io/apebench/api/scenarios/normalized/nonlinear

Initializes the alpha and beta coefficients for the BurgersSingleChannel scenario, incorporating diffusion_alpha and convection_sc_beta.

```python
def __post_init__(self):
    self.alphas = (0.0, 0.0, self.diffusion_alpha, 0.0, 0.0)
    self.betas = (0.0, self.convection_sc_beta, 0.0)

    super().__post_init__()
```

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

### Get Coarse Stepper

Source: https://tum-pbs.github.io/apebench/api/scenarios/physical/generic

Retrieves a coarse stepper for the nonlinear scenario, configured with a time step scaled by `coarse_proportion`.

```python
def get_coarse_stepper(self):
    return self._build_stepper(self.dt * self.coarse_proportion)
```

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

### Get Reference Stepper

Source: https://tum-pbs.github.io/apebench/api/scenarios/physical/generic

Retrieves the reference stepper for the nonlinear scenario, configured with the base time step `dt`.

```python
def get_ref_stepper(self):
    return self._build_stepper(self.dt)
```

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

### Visualize Sample Rollout

Source: https://tum-pbs.github.io/apebench/typical_workflow

Display a sample rollout using `matplotlib.pyplot.imshow`. The image represents the rollout data, with time on the x-axis and space on the y-axis. Color mapping is set to 'RdBu_r' with specific vmin and vmax values.

```python
import numpy as np

data_sample_rollout = apebench.melt_sample_rollouts(data)

plt.imshow(np.array(data_sample_rollout["sample_rollout"][0])[:, 0, :].T, origin="lower", aspect="auto", vmin=-1, vmax=1, cmap="RdBu_r")
plt.xlabel("time"); plt.ylabel("space"); plt.show()
```

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

### Get Linear Scenario Name

Source: https://tum-pbs.github.io/apebench/api/scenarios/physical/generic

Returns the name of the linear physical scenario based on the number of spatial dimensions.

```python
def get_scenario_name(self) -> str:
    return f"{self.num_spatial_dims}d_phy_four"
```

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

### __post_init__ Method for FisherKPP

Source: https://tum-pbs.github.io/apebench/api/scenarios/normalized/react

Initializes the alpha and beta parameters based on the scenario's attributes and calls the parent class's __post_init__ method.

```python
def __post_init__(self):
    self.alphas = (self.drag_alpha, 0.0, self.diffusion_alpha, 0.0, 0.0)
    self.betas = (self.quadratic_beta, 0.0, 0.0)

    super().__post_init__()

```

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

### Get Swift-Hohenberg Scenario Name

Source: https://tum-pbs.github.io/apebench/api/scenarios/physical/react

Returns the name of the Swift-Hohenberg scenario, formatted based on the number of spatial dimensions.

```python
def get_scenario_name(self) -> str:
    return f"{self.num_spatial_dims}d_phy_sh"
```

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

### Generate and Get Data Shapes

Source: https://tum-pbs.github.io/apebench/examples/scrape_datasets

Generate training and testing data using the instantiated scenario and display their shapes.

```python
train_data = advection_1d_difficulty.get_train_data()
test_data = advection_1d_difficulty.get_test_data()

train_data.shape, test_data.shape

```

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

### Get Coarse Stepper

Source: https://tum-pbs.github.io/apebench/api/scenarios/normalized/generic

Retrieves a coarse version of the normalized convection stepper. This is achieved by scaling down the convection and alphas.

```python
def get_coarse_stepper(self) -> ex.BaseStepper:
    return self._build_stepper(
        self.coarse_proportion * self.convection_beta,
        tuple(f * self.coarse_proportion for f in self.alphas),
    )
```

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

### Korteweg-De Vries __post_init__ Method

Source: https://tum-pbs.github.io/apebench/api/scenarios/physical/nonlinear

Initializes the `a_coefs` and `b_coefs` for the KdV scenario using dispersion and convection coefficients. This method is called after the instance is initialized.

```python
def __post_init__(self):
    self.a_coefs = (0.0, 0.0, 0.0, self.dispersion_coef, self.hyp_diffusion_coef)
    self.b_coefs = (0.0, self.convection_sc_coef, 0.0)
```

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

### Initialize SupervisedTrainer with Rollout Steps

Source: https://tum-pbs.github.io/apebench/api/scenarios/base

Configures supervised training with a specified number of rollout steps and optional BPTT cutting. Accepts training trajectories, optimizer, and rollout step count.

```python
num_rollout_steps = int(train_args[1])
            if len(train_args) > 2:
                cut_bptt = train_args[2].lower() == "true"
            else:
                cut_bptt = False
            trainer = tx.trainer.SupervisedTrainer(
                train_trjs,
                optimizer=optimizer,
                num_training_steps=self.num_training_steps,
                batch_size=self.batch_size,
                num_rollout_steps=num_rollout_steps,
                cut_bptt=cut_bptt,
                time_level_weights=None,
                callback_fn=callback_fn,
            )
```

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

### Get Normalized Linear Scenario Name

Source: https://tum-pbs.github.io/apebench/api/scenarios/normalized/generic

Returns the name for a normalized linear scenario based on spatial dimensions.

```python
def get_scenario_name(self) -> str:
    return f"{self.num_spatial_dims}d_norm_four"
```

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

### Generate Configurations with List Comprehension

Source: https://tum-pbs.github.io/apebench/typical_workflow

Dynamically create the experiment configurations list using a list comprehension for conciseness, especially when only a few parameters differ between experiments.

```python
CONFIGS = [
    {
        "scenario": "diff_adv",
        "task": "predict",
        "net": net,
        "train": "one",
        "start_seed": 0,
        "num_seeds": 10,
    }
    for net in ["Conv;26;10;relu", "FNO;12;8;4;gelu"]
]
```

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

### Get Subset of Test Data

Source: https://tum-pbs.github.io/apebench/api/scenarios/base

Returns a subset of the testing data, with the number of samples determined by the `num_trjs_returned` attribute.

```python
def get_ref_sample_data(
        self,
    ) -> Float[
        Array, "num_trjs_returned test_temporal_horizon+1 num_channels *num_points"
    ]:
        """
        Return a subset of the testing data, the number of samples is defined by
        the attribute `num_trjs_returned`
        """
        test_trj_set = self.get_test_data()
        test_trj_subset = test_trj_set[: self.num_trjs_returned]
        return test_trj_subset
```

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

### Perform Test Rollout and Calculate Error

Source: https://tum-pbs.github.io/apebench/api/scenarios/base

Executes a rollout simulation using a neural stepper starting from the test initial condition and compares the predictions against the reference trajectory. It computes the mean error using a provided or default mean error function (e.g., nRMSE).

```python
test_trjs = self.get_test_data()
test_ics = test_trjs[:, 0]
ref_trjs = test_trjs[:, 1:]
pred_trjs = jax.vmap(
    ex.rollout(
        neural_stepper,
        self.test_temporal_horizon,
        include_init=False,
    )
)(test_ics)

error_rollout = jax.vmap(
    mean_error_fn,
    in_axes=1,  # over the temporal axis
)(pred_trjs, ref_trjs)

return error_rollout
```

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

### __post_init__

Source: https://tum-pbs.github.io/apebench/api/scenarios/physical/react

Performs post-initialization checks for the Swift-Hohenberg scenario. It ensures that the scenario is configured for 2D or 3D spatial dimensions, as 1D is not supported.

```APIDOC
## __post_init__

### Description
Performs post-initialization checks for the Swift-Hohenberg scenario. It ensures that the scenario is configured for 2D or 3D spatial dimensions, as 1D is not supported.

### Method
```
__post_init__(self)
```

### Parameters
This method does not take any parameters.

### Raises
- **ValueError**: If `self.num_spatial_dims` is 1.
```

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

### Display Sample Rollout

Source: https://tum-pbs.github.io/apebench/examples/simple_advection_1d_learning

Displays a single sample rollout as an image. It extracts the spatial data from the first sample, transposes it, and sets visualization parameters like colormap and aspect ratio.

```python
plt.imshow(
    np.array(sample_rollout_data["sample_rollout"][0])[:, 0, :].T,
    vmin=-1,
    vmax=1,
    cmap="RdBu_r",
    aspect="auto",
    origin="lower",
)
```

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

### Run Advection scenario (Run 0)

Source: https://tum-pbs.github.io/apebench/examples/sources_of_randomness_and_reproducibility

Executes the Advection scenario for the first time to collect data and network states.

```python
data_run_0, nets_run_0 = advection_scenario()

```

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

### Get Training and Test Data

Source: https://tum-pbs.github.io/apebench/examples/benchmark_flax_models_with_nnx

Retrieves training data trajectories and test initial conditions from the defined APEBench scenario.

```python
train_data = advection_scenario.get_train_data()
test_ic_set = advection_scenario.get_test_ic_set()
train_data.shape, test_ic_set.shape

```

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

### FisherKPP __post_init__ Method

Source: https://tum-pbs.github.io/apebench/api/scenarios/difficulty/react

Initializes the gammas and deltas based on class attributes and calls the parent class's __post_init__.

```python
def __post_init__(self):
    self.gammas = (self.drag_gamma, 0.0, self.diffusion_gamma)
    self.deltas = (self.quadratic_delta, 0.0, 0.0)

    super().__post_init__()

```

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

### Get Coarse Stepper Configuration

Source: https://tum-pbs.github.io/apebench/api/scenarios/physical/generic

Retrieves the stepper configuration for coarse time stepping. This is used to set up simulation parameters.

```python
def get_coarse_stepper(self) -> ex.BaseStepper:
    return self._build_stepper(self.dt * self.coarse_proportion)
```

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

### Instantiate and configure an APEBench scenario

Source: https://tum-pbs.github.io/apebench/examples/component_extensions

Instantiate the Advection scenario with specific optimizer configuration and report only the custom metric 'my_crazy_metric'. The number of training steps is reduced for quicker execution.

```python
adv_scene = apebench.scenarios.difficulty.Advection(
    optim_config="adam;100;constant;1e-3",
    report_metrics="my_crazy_metric",
)
```

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

### Generate and Get Modified Data Shapes

Source: https://tum-pbs.github.io/apebench/examples/scrape_datasets

Generate training and testing data using the modified scenario and display their shapes.

```python
modified_train_data = modified_advection_1d_difficulty.get_train_data()
modified_test_data = modified_advection_1d_difficulty.get_test_data()

modified_train_data.shape, modified_test_data.shape

```

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

### Get Reference Stepper

Source: https://tum-pbs.github.io/apebench/api/scenarios/normalized/generic

Retrieves the reference stepper for normalized convection. It uses the base convection and alphas defined in the class.

```python
def get_ref_stepper(self):
    return self._build_stepper(self.convection_beta, self.alphas)
```

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

### BurgersSingleChannel __post_init__ Method

Source: https://tum-pbs.github.io/apebench/api/scenarios/physical/nonlinear

Initializes the BurgersSingleChannel scenario by setting diffusion and convection coefficients for both a_coefs and b_coefs. This method is called after the instance is created.

```python
def __post_init__(self):
    self.a_coefs = (0.0, 0.0, self.diffusion_coef, 0.0, 0.0)
    self.b_coefs = (0.0, self.convection_sc_coef, 0.0)
```

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

### Get Reference Stepper for Nonlinear Scenario

Source: https://tum-pbs.github.io/apebench/api/scenarios/normalized/generic

Returns the reference stepper for the nonlinear scenario using the class-defined `alphas` and `betas`.

```python
def get_ref_stepper(self):
    return self._build_stepper(self.alphas, self.betas)
```

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

### Get Reference Stepper for Nonlinear Scenario

Source: https://tum-pbs.github.io/apebench/api/scenarios/difficulty/generic

Retrieves the reference stepper for the nonlinear scenario using the default `gammas` and `deltas`.

```python
def get_ref_stepper(self):
    return self._build_stepper(self.gammas, self.deltas)

```