### Install xbatcher from Source

Source: https://github.com/xarray-contrib/xbatcher/blob/main/doc/index.rst

Installs the xbatcher library directly from its GitHub repository using pip.

```bash
python -m pip install git+https://github.com/xarray-contrib/xbatcher.git
```

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

### Install xbatcher from PyPI

Source: https://github.com/xarray-contrib/xbatcher/blob/main/doc/index.rst

Installs the xbatcher library using pip from the Python Package Index.

```bash
python -m pip install xbatcher
```

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

### Install xbatcher with PyTorch Support (PyPI)

Source: https://github.com/xarray-contrib/xbatcher/blob/main/doc/index.rst

Installs xbatcher with PyTorch support using pip from the Python Package Index.

```bash
python -m pip install xbatcher[torch]
```

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

### Build Documentation with Conda and Make

Source: https://github.com/xarray-contrib/xbatcher/blob/main/doc/contributing.rst

These commands set up a conda environment with the necessary documentation requirements and then build the HTML documentation. Navigate to the 'docs' directory and execute 'make html' to generate the documentation files.

```bash
conda env create --file ci/requirements/docs.yml
conda activate xbatcher-docs
cd docs
make html
```

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

### Install xbatcher with TensorFlow Support (PyPI)

Source: https://github.com/xarray-contrib/xbatcher/blob/main/doc/index.rst

Installs xbatcher with TensorFlow support using pip from the Python Package Index.

```bash
python -m pip install xbatcher[tensorflow]
```

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

### Clone and Set Up Upstream Repository (Git)

Source: https://github.com/xarray-contrib/xbatcher/blob/main/doc/contributing.rst

This snippet demonstrates how to clone your forked xbatcher repository and add the upstream repository as a remote. This is essential for keeping your fork synchronized with the main project.

```shell
git clone git@github.com:yourusername/xbatcher.git
cd xbatcher
git remote add upstream git@github.com:xarray-contrib/xbatcher.git
```

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

### Create and Checkout a New Feature Branch (Git - Simplified)

Source: https://github.com/xarray-contrib/xbatcher/blob/main/doc/contributing.rst

A more concise way to create a new branch and immediately switch to it. This is a common shortcut for starting new development work.

```shell
git checkout -b shiny-new-feature
```

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

### Generate Basic Batches with xbatcher

Source: https://github.com/xarray-contrib/xbatcher/blob/main/doc/demo.ipynb

Initializes a `BatchGenerator` from xbatcher with specified input dimensions. This example sets the 'time' dimension to 10 points per sample, flattening other dimensions into a 'sample' dimension. It then prints the total number of batches generated.

```python
n_timepoint_in_each_sample = 10

bgen = xbatcher.BatchGenerator(
    ds=ds,
    input_dims={'time': n_timepoint_in_each_sample},
)

print(f'{len(bgen)} batches')
```

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

### Install xbatcher and zarr

Source: https://github.com/xarray-contrib/xbatcher/blob/main/doc/user-guide/caching.ipynb

Installs the xbatcher library and zarr for serialization using pip or conda. These are necessary dependencies for utilizing the caching feature.

```bash
python -m pip install xbatcher zarr
```

```bash
conda install -c conda-forge xbatcher zarr
```

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

### Install Pre-commit Hooks (Shell)

Source: https://github.com/xarray-contrib/xbatcher/blob/main/doc/contributing.rst

After activating the conda environment, this command installs the pre-commit framework. Pre-commit helps manage and automate code linting and style checks before committing changes.

```shell
pre-commit install
```

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

### Install xbatcher from Conda

Source: https://github.com/xarray-contrib/xbatcher/blob/main/doc/index.rst

Installs the xbatcher library using Conda from the conda-forge channel.

```bash
conda install -c conda-forge xbatcher
```

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

### Plot Example Dataset Time Slice

Source: https://github.com/xarray-contrib/xbatcher/blob/main/doc/demo.ipynb

Visualizes the 'tasmax' (maximum air temperature) variable from the loaded dataset for the first time step. This provides a quick visual inspection of the data.

```python
# plot the first time dimension
ds.isel(time=0).tasmax.plot();
```

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

### Control Batch Size and Shape

Source: https://github.com/xarray-contrib/xbatcher/blob/main/doc/demo.ipynb

Demonstrates how to control the number of samples per batch using `batch_dims` and `concat_input_dims`. This example configures batches to have 20 time points each, resulting in fewer batches and more samples per batch compared to the previous example. `concat_input_dims=True` ensures that the input dimensions are concatenated as specified.

```python
n_timepoint_in_each_sample = 10
n_timepoint_in_each_batch = 20

bgen = xbatcher.BatchGenerator(
    ds=ds,
    input_dims={'time': n_timepoint_in_each_sample},
    batch_dims={'time': n_timepoint_in_each_batch},
    concat_input_dims=True,
)

print(f'{len(bgen)} batches')
```

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

### Manage Code Changes with Git

Source: https://github.com/xarray-contrib/xbatcher/blob/main/doc/contributing.rst

This sequence of commands demonstrates basic Git operations for contributing code changes. It includes checking the status of your repository, adding new files, committing changes with a structured message, and pushing your branch to a remote repository.

```bash
git status
git add path/to/file-to-be-added.py
git commit -m
git push origin shiny-new-feature
git remote -v
```

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

### Install xbatcher with PyTorch Support (Conda)

Source: https://github.com/xarray-contrib/xbatcher/blob/main/doc/index.rst

Installs xbatcher along with PyTorch support using Conda from the conda-forge channel.

```bash
conda install -c conda-forge xbatcher pytorch
```

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

### Load Example CMIP6 Dataset

Source: https://github.com/xarray-contrib/xbatcher/blob/main/doc/demo.ipynb

Loads a sample CMIP6 dataset from a Zarr store located in S3. This dataset represents daily maximum air temperature over 60 days. The `chunks={}` argument ensures the dataset is loaded without pre-defined chunking, and `backend_kwargs` are used for anonymous S3 access.

```python
store = 's3://carbonplan-share/xbatcher/example_cmip6_data.zarr'
ds = xr.open_dataset(
    store, engine='zarr', chunks={},
    backend_kwargs={'storage_options': {'anon': True}}
)

# inspect the dataset
ds
```

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

### Install xbatcher with TensorFlow Support (Conda)

Source: https://github.com/xarray-contrib/xbatcher/blob/main/doc/index.rst

Installs xbatcher along with TensorFlow support using Conda from the conda-forge channel.

```bash
conda install -c conda-forge xbatcher tensorflow
```

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

### Create and Checkout a New Feature Branch (Git)

Source: https://github.com/xarray-contrib/xbatcher/blob/main/doc/contributing.rst

These commands show how to create a new branch for a feature or bug fix and then switch to that branch. It's recommended to keep branches focused on a single purpose.

```shell
git branch shiny-new-feature
git checkout shiny-new-feature
```

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

### Run Continuous Benchmarks with ASV

Source: https://github.com/xarray-contrib/xbatcher/blob/main/doc/contributing.rst

This command initiates continuous benchmarking using the asv tool. It allows specifying the environment (conda or virtualenv) and optionally filtering benchmarks by name. The '-f 1.1' flag sets the threshold for significant changes.

```bash
asv continuous -f 1.1 main <my-branch>
```

```bash
asv continuous -f 1.1 -E virtualenv main <my-branch>
```

```bash
asv continuous -f 1.1 main HEAD -b benchmarks.Generator.time_batch_preload
```

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

### Run the xbatcher Test Suite (Shell)

Source: https://github.com/xarray-contrib/xbatcher/blob/main/doc/contributing.rst

This command executes the test suite for the xbatcher project using the pytest framework. Ensure your development environment is set up and activated before running this.

```shell
pytest xbatcher
```

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

### Multi-Dimensional Batching for CNNs

Source: https://context7.com/xarray-contrib/xbatcher/llms.txt

Demonstrates creating spatiotemporal batches suitable for CNNs by specifying multiple input dimensions with overlapping windows. This example uses climate model data. Requires xarray and xbatcher.

```python
import numpy as np
import xarray as xr
import xbatcher

# Climate model output: 30 days, 18x18 spatial grid
data = np.random.rand(30, 18, 18)
ds = xr.Dataset({
    'tasmax': (['time', 'lat', 'lon'], data)
})
```

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

### Inspect Overlapping Samples in a Batch

Source: https://github.com/xarray-contrib/xbatcher/blob/main/doc/demo.ipynb

Selects a specific pixel (latitude and longitude) from a batch generated with overlapping inputs. It then displays the data for this pixel and prints the start and end time points for the first two samples within that pixel's batch to demonstrate the effect of `input_overlap`.

```python
lat = -90
lon = 0
pixel = batch.sel(lat=lat, lon=lon)
display(pixel)

print(
    f'sample 1 goes from {pixel.isel(input_batch=0).time[0].values} to {pixel.isel(input_batch=0).time[-1].values}'
)
print(
    f'sample 2 goes from {pixel.isel(input_batch=1).time[0].values} to {pixel.isel(input_batch=1).time[-1].values}'
)
```

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

### Visualize a Sample Prediction

Source: https://github.com/xarray-contrib/xbatcher/blob/main/doc/user-guide/training-a-neural-network-with-keras-and-xbatcher.ipynb

Takes one batch from the `train_dataloader`, makes a prediction using the trained model, and visualizes the first image along with its true label and the predicted label. This demonstrates the model's performance on a sample.

```python
# Visualize a prediction on a sample image
for train_features, train_labels in train_dataloader.take(1):
    img = train_features[0].numpy().squeeze()
    label = train_labels[0].numpy()
    predicted_label = tf.argmax(model.predict(train_features[:1]), axis=1).numpy()[0]

    plt.imshow(img, cmap='gray')
    plt.title(f'True Label: {label}, Predicted: {predicted_label}')
    plt.show()
    break
```

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

### Instantiate Neural Network Model

Source: https://github.com/xarray-contrib/xbatcher/blob/main/doc/user-guide/training-a-neural-network-with-Pytorch-and-xbatcher.ipynb

Creates an instance of the `SimpleNN` class, initializing the neural network model ready for training.

```python
# Instantiate the model
model = SimpleNN()
model
```

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

### Create and Iterate Batch Generator in Python

Source: https://github.com/xarray-contrib/xbatcher/blob/main/doc/roadmap.rst

Demonstrates how to open an xarray dataset, preprocess it, create a BatchGenerator with specified batch dimensions, and then iterate through the generator to obtain batches for model training or prediction.

```Python
import xarray as xr
import xbatcher as xb

da = xr.open_dataset(filename, chunks=chunks)     # open a dataset and use dask
da_train = preprocess(ds)                         # perform some preprocessing
bgen = xb.BatchGenerator(da_train, {'time': 10})  # create a generator

for batch in bgen:                                # iterate through the generator
    model.fit(batch['x'], batch['y'])             # fit a deep-learning model
    # or
    model.predict(batch['x'])                     # make one batch of predictions
```

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

### Create and Activate Conda Environment for xbatcher (Shell)

Source: https://github.com/xarray-contrib/xbatcher/blob/main/doc/contributing.rst

This command creates a new conda environment named 'xbatcher-tests' using the specified environment file and then activates it. This isolated environment is necessary for development and testing.

```shell
conda env create --file ci/requirements/environment.yml
conda activate xbatcher-tests
```

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

### Prepare Stacked 2D Input Batches with xbatcher for ML Algorithms

Source: https://github.com/xarray-contrib/xbatcher/blob/main/doc/demo.ipynb

This code snippet shows how to prepare data for machine learning algorithms like scikit-learn or xgboost by stacking input dimensions into a 2D format (n_samples, n_features). It iterates through batches generated by xbatcher, transposes dimensions, extracts features, stacks spatial and temporal input dimensions, and then extracts labels, printing the resulting shapes.

```python
for i, batch in enumerate(bgen):
    print(f'batch {i}')
    # make sure the ordering of dimension is consistent
    batch = batch.transpose('input_batch', 'lat_input', 'lon_input', 'time_input')

    # only use the first 9 time points as features, since the last time point is the label to be predicted
    features = batch.tasmax.isel(time_input=slice(0, 9))
    features = features.stack(features=['lat_input', 'lon_input', 'time_input'])

    # select the center pixel at the last time point to be the label to be predicted
    # the actual lat/lon/time for each of the sample can be accessed in labels.coords
    labels = batch.tasmax.isel(lat_input=5, lon_input=5, time_input=9)

    print('feature shape', features.shape)
    print('label shape', labels.shape)
    print('shape of lat of each sample', labels.coords['lat'].shape, '\n')
```

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

### Import Libraries for Xbatcher, Keras, TensorFlow, and Matplotlib

Source: https://github.com/xarray-contrib/xbatcher/blob/main/doc/user-guide/training-a-neural-network-with-keras-and-xbatcher.ipynb

Imports necessary libraries for data manipulation (xarray), neural network building (Keras, TensorFlow), plotting (matplotlib), and batch generation (xbatcher).

```python
import matplotlib.pyplot as plt
import tensorflow as tf
import xarray as xr
from keras import layers, models, optimizers

import xbatcher as xb
import xbatcher.loaders.keras
```

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

### Visualize Sample Image

Source: https://github.com/xarray-contrib/xbatcher/blob/main/doc/user-guide/training-a-neural-network-with-Pytorch-and-xbatcher.ipynb

Visualizes a single image from the loaded dataset using matplotlib. This step helps in understanding the data format and content before proceeding with model training.

```python
ds.sel(sample=1).images.plot(cmap='gray');
```

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

### Initialize BatchGenerator with Caching

Source: https://github.com/xarray-contrib/xbatcher/blob/main/doc/user-guide/caching.ipynb

Demonstrates initializing xbatcher's BatchGenerator with a Zarr DirectoryStore as a cache. This sets up the generator to store and retrieve batches from a local directory.

```python
import tempfile

import xarray as xr
import zarr

import xbatcher

# create a cache using Zarr's DirectoryStore
directory = f'{tempfile.mkdtemp()}/xbatcher-cache'
print(directory)
cache = zarr.storage.DirectoryStore(directory)

# load a sample dataset
ds = xr.tutorial.open_dataset('air_temperature', chunks={})

# create a BatchGenerator with caching enabled
gen = xbatcher.BatchGenerator(ds, input_dims={'lat': 10, 'lon': 10}, cache=cache)
```

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

### Import xarray and xbatcher Libraries

Source: https://github.com/xarray-contrib/xbatcher/blob/main/doc/demo.ipynb

Imports the necessary libraries, xarray for data handling and xbatcher for batch generation. These are fundamental for using the xbatcher functionality.

```python
import xarray as xr

import xbatcher
```

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

### Update Local Branch from Upstream Main (Git)

Source: https://github.com/xarray-contrib/xbatcher/blob/main/doc/contributing.rst

This sequence fetches the latest changes from the upstream repository's main branch and merges them into your current local branch. This is crucial for resolving potential merge conflicts before submitting a pull request.

```shell
git fetch upstream
git merge upstream/main
```

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

### Convert Xarray to TensorFlow Tensor with .tf accessor

Source: https://context7.com/xarray-contrib/xbatcher/llms.txt

The `.tf` accessor allows direct conversion of Xarray DataArrays and Datasets to TensorFlow tensors. This is useful for integrating Xarray data with TensorFlow models. It requires xarray, numpy, and xbatcher to be installed.

```python
import numpy as np
import xarray as xr
import xbatcher  # Registers accessors

data = np.random.rand(10, 32, 32).astype(np.float32)
da = xr.DataArray(data, dims=['batch', 'height', 'width'])

# Convert DataArray directly to TensorFlow tensor
tensor = da.tf.to_tensor()
print(f'TensorFlow tensor shape: {tensor.shape}')
# Output: TensorFlow tensor shape: (10, 32, 32)

# Works with Datasets too (converts to single tensor)
ds = xr.Dataset({'var1': da})
tensor = ds.tf.to_tensor()
print(f'Dataset tensor shape: {tensor.shape}')
```

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

### Generate Batches with Overlapping Inputs

Source: https://github.com/xarray-contrib/xbatcher/blob/main/doc/demo.ipynb

Configures `BatchGenerator` to create batches with overlapping time points using the `input_overlap` argument. This is common in machine learning for time-series data. The example sets up 10 time points per sample, 20 time points per batch, concatenates input dimensions, and specifies an overlap of 9 time points.

```python
n_timepoint_in_each_sample = 10
n_timepoint_in_each_batch = 20
input_overlap = 9

bgen = xbatcher.BatchGenerator(
    ds=ds,
    input_dims={'time': n_timepoint_in_each_sample},
    batch_dims={'time': n_timepoint_in_each_batch},
    concat_input_dims=True,
    input_overlap={'time': input_overlap},
)

batch = bgen[0]

print(f'{len(bgen)} batches')
batch
```

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

### Open Zarr Dataset with Xarray

Source: https://github.com/xarray-contrib/xbatcher/blob/main/doc/user-guide/training-a-neural-network-with-keras-and-xbatcher.ipynb

Opens a dataset stored in Zarr format from an S3 bucket using xarray. It configures the dataset to be opened anonymously and without chunking for initial loading.

```python
# Open the dataset stored in Zarr format
ds = xr.open_dataset(
    's3://carbonplan-share/xbatcher/fashion-mnist-train.zarr',
    engine='zarr',
    chunks={},
    backend_kwargs={'storage_options': {'anon': True}},
)
```

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

### Basic Batch Generation with xbatcher.BatchGenerator

Source: https://context7.com/xarray-contrib/xbatcher/llms.txt

Demonstrates the fundamental usage of BatchGenerator to create batches from an Xarray Dataset. It shows how non-specified dimensions are stacked into a 'sample' dimension by default. The code requires xarray and xbatcher libraries.

```python
import numpy as np
import xarray as xr
import xbatcher

# Create sample data: 60 time steps, 100 lat points, 100 lon points
data = np.random.rand(60, 100, 100)
ds = xr.Dataset({
    'temperature': (['time', 'lat', 'lon'], data)
})

# Basic batch generation: 10 time points per batch
# Non-specified dimensions (lat, lon) are stacked into 'sample' dimension
bgen = xbatcher.BatchGenerator(
    ds=ds,
    input_dims={'time': 10}
)

print(f'Number of batches: {len(bgen)}')  # Output: 6 batches
batch = bgen[0]
print(f'Batch shape: {dict(batch.sizes)}')
# Output: {'sample': 10000, 'time': 10}

# Iterate through all batches
for i, batch in enumerate(bgen):
    print(f'Batch {i}: {batch.temperature.shape}')
```

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

### Import Libraries and Load Dataset with Xarray

Source: https://github.com/xarray-contrib/xbatcher/blob/main/doc/user-guide/training-a-neural-network-with-Pytorch-and-xbatcher.ipynb

Imports necessary libraries for data manipulation, neural networks, and batch generation. It then opens the FashionMNIST training dataset from an S3 bucket using Xarray, specifying the Zarr engine and backend arguments for anonymous access.

```python
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import torch.optim as optim
import torch.utils.data
import xarray as xr

import xbatcher as xb
import xbatcher.loaders.torch

ds = xr.open_dataset(
    's3://carbonplan-share/xbatcher/fashion-mnist-train.zarr',
    engine='zarr',
    chunks={},
    backend_kwargs={'storage_options': {'anon': True}},
)
ds
```

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

### Visualize a Sample Batch from the DataLoader

Source: https://github.com/xarray-contrib/xbatcher/blob/main/doc/user-guide/training-a-neural-network-with-keras-and-xbatcher.ipynb

Extracts a single batch from the `train_dataloader` to print its shape and visualize the first image along with its label. This helps in verifying the data loading and batching process.

```python
# Extract a batch from the DataLoader
for train_features, train_labels in train_dataloader.take(1):
    print(f'Feature batch shape: {train_features.shape}')
    print(f'Labels batch shape: {train_labels.shape}')

    img = train_features[0].numpy().squeeze()  # Extract the first image
    label = train_labels[0].numpy()
    plt.imshow(img, cmap='gray')
    plt.title(f'Label: {label}')
    plt.show()
    break
```

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

### Custom Cache Preprocessing with BatchGenerator

Source: https://github.com/xarray-contrib/xbatcher/blob/main/doc/user-guide/caching.ipynb

Demonstrates how to use a custom preprocessing function with xbatcher's BatchGenerator before batches are cached. This allows for modifications to batches, such as adding new variables, prior to storage.

```python
import tempfile

import xarray as xr
import zarr

import xbatcher

# create a cache using Zarr's DirectoryStore
directory = f'{tempfile.mkdtemp()}/xbatcher-cache'
cache = zarr.storage.DirectoryStore(directory)


def preprocess_batch(batch):
    # example: add a new variable to each batch
    batch['new_var'] = batch['air'] * 2
    return batch


ds = xr.tutorial.open_dataset('air_temperature', chunks={})

gen_with_preprocess = xbatcher.BatchGenerator(
    ds,
    input_dims={'lat': 10, 'lon': 10},
    cache=cache,
    cache_preprocess=preprocess_batch,
)

# Now, each cached batch will include the 'new_var' variable
for batch in gen_with_preprocess:
    print(batch)
    break
```

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

### Create Xbatcher Batch Generators for Features and Labels

Source: https://github.com/xarray-contrib/xbatcher/blob/main/doc/user-guide/training-a-neural-network-with-keras-and-xbatcher.ipynb

Initializes two `BatchGenerator` instances from xbatcher: one for image features ('images') and one for corresponding labels ('labels'). It specifies input dimensions and sets `preload_batch` to False for dynamic loading.

```python
# Define batch generators for features (X) and labels (y)
X_bgen = xb.BatchGenerator(
    ds['images'],
    input_dims={'sample': 2000, 'channel': 1, 'height': 28, 'width': 28},
    preload_batch=False,  # Load each batch dynamically
)
y_bgen = xb.BatchGenerator(
    ds['labels'], input_dims={'sample': 2000}, preload_batch=False
)
```

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

### Load and Inspect a Batch

Source: https://github.com/xarray-contrib/xbatcher/blob/main/doc/user-guide/training-a-neural-network-with-Pytorch-and-xbatcher.ipynb

Fetches the first batch of features and labels from the `train_dataloader` and prints their shapes. It then visualizes the first image in the batch and displays its corresponding label.

```python
train_features, train_labels = next(iter(train_dataloader))

print(f'Feature batch shape: {train_features.size()}')
print(f'Labels batch shape: {train_labels.size()}')
img = train_features[0].squeeze()
label = train_labels[0]
plt.imshow(img, cmap='gray')
plt.show()
print(f'Label: {label}')
```

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

### Build and Compile a Simple Keras Neural Network

Source: https://github.com/xarray-contrib/xbatcher/blob/main/doc/user-guide/training-a-neural-network-with-keras-and-xbatcher.ipynb

Defines a sequential Keras model with a Flatten layer, a Dense hidden layer with ReLU activation, and a Dense output layer with softmax activation for classification. The model is compiled with the Adam optimizer, sparse categorical crossentropy loss, and accuracy as a metric.

```python
# Define a simple feedforward neural network
model = models.Sequential(
    [
        layers.Flatten(input_shape=(1, 28, 28)),  # Flatten input images
        layers.Dense(128, activation='relu'),  # Fully connected layer with 128 units
        layers.Dense(10, activation='softmax'),  # Output layer for 10 classes
    ]
)

# Compile the model
model.compile(
    optimizer=optimizers.Adam(learning_rate=0.001),
    loss='sparse_categorical_crossentropy',
    metrics=['accuracy'],
)

# Display model summary
model.summary()
```

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

### Train the Keras Model using the DataLoader

Source: https://github.com/xarray-contrib/xbatcher/blob/main/doc/user-guide/training-a-neural-network-with-keras-and-xbatcher.ipynb

Trains the compiled Keras model for a specified number of epochs using the `train_dataloader`. The `%%time` magic command is used to measure the training duration.

```python
%%time

# Train the model for 5 epochs
epochs = 5

model.fit(
    train_dataloader,  # Pass the DataLoader directly
    epochs=epochs,
    verbose=1,  # Print progress during training
)
```

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

### Visualize Sample Prediction

Source: https://github.com/xarray-contrib/xbatcher/blob/main/doc/user-guide/training-a-neural-network-with-Pytorch-and-xbatcher.ipynb

Evaluates the trained model on a sample image from the training data. It displays the image, its true label, and the predicted label generated by the model.

```python
# Visualize a sample prediction
img = train_features[0].squeeze()
label = train_labels[0]
predicted_label = torch.argmax(model(train_features[0:1]), dim=1).item()

plt.imshow(img, cmap='gray')
plt.title(f'True Label: {label}, Predicted: {predicted_label}')
plt.show()
```

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

### PyTorch MapDataset Loader for Supervised Learning

Source: https://context7.com/xarray-contrib/xbatcher/llms.txt

Explains how to use the `MapDataset` class to wrap xbatcher generators into a PyTorch-compatible Dataset. It supports separate generators for features (X) and targets (y), suitable for supervised learning tasks.

```python
import numpy as np
import xarray as xr
import xbatcher
from xbatcher.loaders.torch import MapDataset
import torch
from torch.utils.data import DataLoader

# Create feature and target data
X_data = np.random.rand(100, 32, 32).astype(np.float32)
y_data = np.random.rand(100, 1).astype(np.float32)

X_ds = xr.DataArray(X_data, dims=['time', 'lat', 'lon'])
y_ds = xr.DataArray(y_data, dims=['time', 'feature'])

# Create generators for features and targets
X_gen = xbatcher.BatchGenerator(X_ds, input_dims={'lat': 32, 'lon': 32})
y_gen = xbatcher.BatchGenerator(y_ds, input_dims={'feature': 1})

# Create PyTorch dataset
dataset = MapDataset(
    X_generator=X_gen,
    y_generator=y_gen
)

# Use with PyTorch DataLoader
dataloader = DataLoader(dataset, batch_size=8, shuffle=True)

for X_batch, y_batch in dataloader:
    print(f'X shape: {X_batch.shape}, y shape: {y_batch.shape}')
    # X shape: torch.Size([8, 32, 32]), y shape: torch.Size([8, 1])
    break
```

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

### Create Batches using Xarray's Batch Accessor

Source: https://context7.com/xarray-contrib/xbatcher/llms.txt

Shows how to use the `.batch` accessor registered on Xarray DataArrays and Datasets. This provides a more concise way to create batch generators directly from data objects, simplifying the workflow.

```python
import numpy as np
import xarray as xr
import xbatcher  # Registers the accessor

data = np.random.rand(100, 50, 50)
da = xr.DataArray(
    data,
    dims=['time', 'lat', 'lon'],
    name='precipitation'
)

# Use the batch accessor instead of BatchGenerator directly
bgen = da.batch.generator(
    input_dims={'time': 10, 'lat': 25, 'lon': 25}
)

print(f'Number of batches: {len(bgen)}')
for batch in bgen:
    print(f'Batch shape: {batch.shape}')
```

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

### Map Xbatcher Generators to a Keras-Compatible TensorFlow Dataset

Source: https://github.com/xarray-contrib/xbatcher/blob/main/doc/user-guide/training-a-neural-network-with-keras-and-xbatcher.ipynb

Wraps the xbatcher `BatchGenerator` instances using `xbatcher.loaders.keras.CustomTFDataset`. This is then converted into a `tf.data.Dataset` for use with Keras, including prefetching for performance.

```python
# Use xbatcher's MapDataset to wrap the generators
dataset = xbatcher.loaders.keras.CustomTFDataset(X_bgen, y_bgen)

# Create a DataLoader using tf.data.Dataset
train_dataloader = tf.data.Dataset.from_generator(
    lambda: iter(dataset),
    output_signature=(
        tf.TensorSpec(shape=(2000, 1, 28, 28), dtype=tf.float32),  # Images
        tf.TensorSpec(shape=(2000,), dtype=tf.int64),  # Labels
    ),
).prefetch(3)  # Prefetch 3 batches to improve performance
```

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

### Using S3 as a Cache Backend for BatchGenerator

Source: https://github.com/xarray-contrib/xbatcher/blob/main/doc/user-guide/caching.ipynb

Illustrates configuring xbatcher's BatchGenerator to use an S3 bucket as the cache storage backend. This requires the s3fs library and appropriate AWS credentials.

```python
import s3fs
import zarr

import xarray as xr
import xbatcher

# Set up S3 filesystem (you'll need appropriate credentials)
s3 = s3fs.S3FileSystem(anon=False)
cache = s3.get_mapper('s3://my-bucket/my-cache.zarr')

# load a sample dataset
ds = xr.tutorial.open_dataset('air_temperature', chunks={})

# Use this cache with BatchGenerator
gen_s3 = xbatcher.BatchGenerator(ds, input_dims={'lat': 10, 'lon': 10}, cache=cache)
```

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

### Access and Inspect a Batch

Source: https://github.com/xarray-contrib/xbatcher/blob/main/doc/demo.ipynb

Retrieves the first batch (index 0) from the `BatchGenerator` and displays its content. This allows for inspection of the batch's structure and data.

```python
batch = bgen[0]
batch
```

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

### Demonstrate Last Batch Behavior

Source: https://github.com/xarray-contrib/xbatcher/blob/main/doc/demo.ipynb

Illustrates how xbatcher handles cases where the dataset size is not perfectly divisible by the specified `input_dims`. The remainder is discarded, ensuring only full batches are produced. The loop prints the last time point of the dataset and the last time point of the current batch to highlight this behavior.

```python
n_timepoint_in_batch = 31

bgen = xbatcher.BatchGenerator(ds=ds, input_dims={'time': n_timepoint_in_batch})

for batch in bgen:
    print(f'last time point in ds    is {ds.time[-1].values}')
    print(f'last time point in batch is {batch.time[-1].values}')
batch
```

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

### Batch Generation with batch_dims and concat_input_dims

Source: https://context7.com/xarray-contrib/xbatcher/llms.txt

Illustrates advanced BatchGenerator configuration using `batch_dims` to control the number of samples per batch and `concat_input_dims=True` to combine input dimension chunks into the sample dimension. This is useful for creating batches with multiple windows. Requires xarray and xbatcher.

```python
import numpy as np
import xarray as xr
import xbatcher

data = np.random.rand(60, 100, 100)
ds = xr.Dataset({
    'temperature': (['time', 'lat', 'lon'], data)
})

# 10 time points per sample, 20 time points per batch
# This creates 2 samples per batch (20/10 = 2)
bgen = xbatcher.BatchGenerator(
    ds=ds,
    input_dims={'time': 10},
    batch_dims={'time': 20},
    concat_input_dims=True
)

print(f'Number of batches: {len(bgen)}')  # Output: 3 batches
batch = bgen[0]
print(f'Batch dimensions: {dict(batch.sizes)}')
# Output: {'sample': 20000, 'time_input': 10}
# Note: 'time' becomes 'time_input' when concat_input_dims=True
```

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

### Verify Batch Size Calculation

Source: https://github.com/xarray-contrib/xbatcher/blob/main/doc/demo.ipynb

Calculates and prints the expected number of samples within a batch based on the spatial dimensions (latitude and longitude) of the dataset. This is compared against the actual number of samples in the retrieved batch.

```python
expected_batch_size = len(ds.lat) * len(ds.lon)
print(
    f'Expecting {expected_batch_size} samples per batch, getting {len(batch.sample)} samples per batch'
)
```

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

### Prepare CNN Input Batches with xbatcher

Source: https://github.com/xarray-contrib/xbatcher/blob/main/doc/demo.ipynb

This code snippet demonstrates how to initialize and use xbatcher.BatchGenerator to create batches suitable for training a CNN. It configures input and batch dimensions, including overlap, and then iterates through the generated batches to extract features and labels, printing their shapes.

```python
bgen = xbatcher.BatchGenerator(
    ds=ds[['tasmax']].isel(lat=slice(0, 18), lon=slice(0, 18), time=slice(0, 30)),
    input_dims={'lat': 9, 'lon': 9, 'time': 10},
    batch_dims={'lat': 18, 'lon': 18, 'time': 15},
    concat_input_dims=True,
    input_overlap={'lat': 8, 'lon': 8, 'time': 9},
)

for i, batch in enumerate(bgen):
    print(f'batch {i}')
    # make sure the ordering of dimension is consistent
    batch = batch.transpose('input_batch', 'lat_input', 'lon_input', 'time_input')

    # only use the first 9 time points as features, since the last time point is the label to be predicted
    features = batch.tasmax.isel(time_input=slice(0, 9))
    # select the center pixel at the last time point to be the label to be predicted
    # the actual lat/lon/time for each of the sample can be accessed in labels.coords
    labels = batch.tasmax.isel(lat_input=5, lon_input=5, time_input=9)

    print('feature shape', features.shape)
    print('label shape', labels.shape)
    print('shape of lat of each sample', labels.coords['lat'].shape)
    print('')
```

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

### Configure PyTorch DataLoader

Source: https://github.com/xarray-contrib/xbatcher/blob/main/doc/user-guide/training-a-neural-network-with-Pytorch-and-xbatcher.ipynb

Sets up a `torch.utils.data.DataLoader` to efficiently load data in batches for training. It utilizes the custom `MapDataset`, sets `batch_size` to None (as batching is handled by xbatcher), and configures prefetching, multiprocessing workers, and persistent workers for optimized data loading performance.

```python
# Create a DataLoader
train_dataloader = torch.utils.data.DataLoader(
    dataset,
    batch_size=None,  # Using batches defined by the dataset itself (via xbatcher)
    prefetch_factor=3,  # Prefetch up to 3 batches in advance to reduce data loading latency
    num_workers=4,  # Use 4 parallel worker processes to load data concurrently
    persistent_workers=True,  # Keep workers alive between epochs for faster subsequent epochs
    multiprocessing_context='forkserver',  # Use "forkserver" to spawn subprocesses, ensuring stability in multiprocessing
)
```

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

### Train the Neural Network

Source: https://github.com/xarray-contrib/xbatcher/blob/main/doc/user-guide/training-a-neural-network-with-Pytorch-and-xbatcher.ipynb

Implements the training loop for the neural network over a specified number of epochs. In each epoch, it iterates through the `train_dataloader`, performs a forward pass, calculates the loss, performs backpropagation, and updates the model weights using the optimizer. Loss is printed every 10 batches.

```python
%%time

epochs = 5

for epoch in range(epochs):
    print(f'Epoch {epoch+1}/{epochs}')
    for batch, (X, y) in enumerate(train_dataloader):
        # Forward pass
        predictions = model(X)
        loss = loss_fn(predictions, y)

        # Backward pass
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if batch % 10 == 0:
            print(f'Batch {batch}: Loss = {loss.item():.4f}')

print('Training completed!')
```

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

### Compare BatchGenerator Performance With and Without Cache

Source: https://github.com/xarray-contrib/xbatcher/blob/main/doc/user-guide/caching.ipynb

Compares the time taken to iterate through batches generated by xbatcher.BatchGenerator with and without the caching feature enabled. It highlights the performance gains on subsequent runs when caching is active.

```python
import time
import tempfile

import xarray as xr
import zarr

import xbatcher

def time_iteration(gen):
    start = time.time()
    for batch in gen:
        pass
    end = time.time()
    return end - start


# load a sample dataset
ds = xr.tutorial.open_dataset('air_temperature', chunks={})


# Without cache
gen_no_cache = xbatcher.BatchGenerator(ds, input_dims={'lat': 10, 'lon': 10})
time_no_cache = time_iteration(gen_no_cache)
print(f'Time without cache: {time_no_cache:.2f} seconds')


# With cache
directory = f'{tempfile.mkdtemp()}/xbatcher-cache'
cache = zarr.storage.DirectoryStore(directory)

gen_with_cache = xbatcher.BatchGenerator(
    ds, input_dims={'lat': 10, 'lon': 10}, cache=cache
)
time_first_run = time_iteration(gen_with_cache)
print(f'Time with cache (first run): {time_first_run:.2f} seconds')


time_second_run = time_iteration(gen_with_cache)
print(f'Time with cache (second run): {time_second_run:.2f} seconds')
```

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

### Implement batch caching with Zarr for performance

Source: https://context7.com/xarray-contrib/xbatcher/llms.txt

xbatcher supports batch caching using Zarr stores to avoid redundant computations, which is particularly beneficial for large datasets or computationally intensive preprocessing steps. It allows for optional preprocessing functions to be applied before caching. Requires numpy, xarray, xbatcher, zarr, tempfile, and os.

```python
import numpy as np
import xarray as xr
import xbatcher
import zarr
import tempfile
import os

data = np.random.rand(50, 64, 64).astype(np.float32)
ds = xr.Dataset({
    'temperature': (['time', 'lat', 'lon'], data)
})

# Create a Zarr store for caching
cache_dir = tempfile.mkdtemp()
cache = zarr.DirectoryStore(os.path.join(cache_dir, 'batch_cache.zarr'))

# Optional preprocessing function applied before caching
def preprocess(batch):
    # Normalize data
    return (batch - batch.mean()) / batch.std()

bgen = xbatcher.BatchGenerator(
    ds=ds,
    input_dims={'time': 10, 'lat': 32, 'lon': 32},
    cache=cache,
    cache_preprocess=preprocess
)

# First access computes and caches
batch0 = bgen[0]
print(f'First batch computed and cached')

# Subsequent access loads from cache
batch0_cached = bgen[0]
print(f'Batch retrieved from cache')
```

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

### Create and manage reusable batch configurations with BatchSchema

Source: https://context7.com/xarray-contrib/xbatcher/llms.txt

The `BatchSchema` class allows for the separation and reuse of batch configuration parameters, independent of the data itself. It supports serialization to JSON and saving to files, making it ideal for managing complex batching strategies. Requires xarray, numpy, and xbatcher.

```python
import numpy as np
import xarray as xr
from xbatcher import BatchSchema, BatchGenerator

data = np.random.rand(60, 100, 100)
ds = xr.Dataset({
    'temperature': (['time', 'lat', 'lon'], data)
})

# Create a schema (stores configuration without retaining data)
schema = BatchSchema(
    ds=ds,
    input_dims={'time': 10, 'lat': 50, 'lon': 50},
    input_overlap={'time': 5},
    batch_dims={'time': 20},
    concat_input_bins=True,
    preload_batch=True
)

# Access schema properties
print(f'Input dims: {schema.input_dims}')
print(f'Batch dims: {schema.batch_dims}')
print(f'Number of selectors: {len(schema.selectors)}')

# Export schema to JSON for persistence
json_str = schema.to_json()
print(f'Schema JSON: {json_str[:100]}...')

# Save schema to file
schema.to_file('batch_schema.json')
```

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

### Utilize xbatcher testing utilities for validation

Source: https://context7.com/xarray-contrib/xbatcher/llms.txt

xbatcher provides a suite of testing utilities, including functions to retrieve expected batch dimensions, validate generator lengths, and confirm the dimensions of generated batches. These tools aid in ensuring the correctness of batch generator configurations. Requires numpy, xarray, xbatcher, and specific testing modules from xbatcher.testing.

```python
import numpy as np
import xarray as xr
import xbatcher
from xbatcher.testing import (
    get_batch_dimensions,
    validate_batch_dimensions,
    validate_generator_length
)

data = np.random.rand(60, 100, 100)
ds = xr.Dataset({
    'temperature': (['time', 'lat', 'lon'], data)
})

bgen = xbatcher.BatchGenerator(
    ds=ds,
    input_dims={'time': 10, 'lat': 50, 'lon': 50}
)

# Get expected batch dimensions
expected_dims = get_batch_dimensions(bgen)
print(f'Expected dimensions: {expected_dims}')

# Validate generator length matches expectations
validate_generator_length(bgen)
print('Generator length validated')

# Validate batch dimensions
batch = bgen[0]
validate_batch_dimensions(expected_dims=expected_dims, batch=batch)
print('Batch dimensions validated')
```

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

### Batch Generation with Overlapping Windows

Source: https://context7.com/xarray-contrib/xbatcher/llms.txt

Shows how to create overlapping windows using the `input_overlap` parameter in BatchGenerator. This is common for time series tasks where adjacent samples share data points. Requires xarray and xbatcher.

```python
import numpy as np
import xarray as xr
import xbatcher

data = np.random.rand(60, 50, 50)
ds = xr.Dataset({
    'temperature': (['time', 'lat', 'lon'], data)
})

# 10 time points per sample with 9-point overlap
# Sample 1: time[0:10], Sample 2: time[1:11], etc.
bgen = xbatcher.BatchGenerator(
    ds=ds,
    input_dims={'time': 10},
    batch_dims={'time': 20},
    concat_input_dims=True,
    input_overlap={'time': 9}
)

print(f'Number of batches: {len(bgen)}')
batch = bgen[0]
print(f'Samples per batch: {batch.sizes["sample"]}')

# Verify overlap by checking time values for a specific pixel
pixel_data = batch.sel(lat=batch.lat[0], lon=batch.lon[0])
print(f'First sample time range: {pixel_data.time.isel(input_batch=0).values}')
print(f'Second sample time range: {pixel_data.time.isel(input_batch=1).values}')
```

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

### Verify Batch Count Calculation

Source: https://github.com/xarray-contrib/xbatcher/blob/main/doc/demo.ipynb

Calculates and prints the expected number of batches based on the total number of time points in the dataset and the number of time points per batch. This is compared against the actual number of batches generated by `BatchGenerator`.

```python
expected_n_batch = len(ds.time) / n_timepoint_in_each_sample
print(f'Expecting {expected_n_batch} batches, getting {len(bgen)} batches')
```

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

### Create Batch Generators with Xbatcher

Source: https://github.com/xarray-contrib/xbatcher/blob/main/doc/user-guide/training-a-neural-network-with-Pytorch-and-xbatcher.ipynb

Initializes two `BatchGenerator` instances from xbatcher: one for image data ('images') and another for corresponding labels ('labels'). These generators are configured with specific input dimensions to create batches of a defined size, with `preload_batch` set to False for on-demand loading.

```python
# Define batch generators
X_bgen = xb.BatchGenerator(
    ds['images'],
    input_dims={'sample': 2000, 'channel': 1, 'height': 28, 'width': 28},
    preload_batch=False,
)
y_bgen = xb.BatchGenerator(
    ds['labels'], input_dims={'sample': 2000}, preload_batch=False
)
X_bgen[0]
```

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

### Create Patches with Overlap using BatchGenerator

Source: https://context7.com/xarray-contrib/xbatcher/llms.txt

Demonstrates how to configure BatchGenerator to create spatial and temporal patches with specified overlap. This is useful for tasks like time series forecasting or image segmentation where context from neighboring data points is important.

```python
import xarray as xr
import xbatcher

# Assume ds is an xarray Dataset
data = xr.tutorial.load_dataset('air_temperature')
ds = data.air

# Create 9x9x10 patches with 8-point spatial overlap and 9-point temporal overlap
bgen = xbatcher.BatchGenerator(
    ds=ds,
    input_dims={'lat': 9, 'lon': 9, 'time': 10},
    batch_dims={'lat': 18, 'lon': 18, 'time': 15},
    concat_input_dims=True,
    input_overlap={'lat': 8, 'lon': 8, 'time': 9}
)

for i, batch in enumerate(bgen):
    # Reorder dimensions for ML frameworks: (batch, height, width, channels/time)
    batch = batch.transpose('input_batch', 'lat_input', 'lon_input', 'time_input')

    # Use first 9 time steps as features, last as label
    features = batch.tasmax.isel(time_input=slice(0, 9))
    labels = batch.tasmax.isel(lat_input=4, lon_input=4, time_input=9)  # center pixel

    print(f'Batch {i}: features {features.shape}, labels {labels.shape}')
```

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

### Inspect Batch with Controlled Dimensions

Source: https://github.com/xarray-contrib/xbatcher/blob/main/doc/demo.ipynb

Retrieves and displays the first batch generated with the specified `batch_dims` and `concat_input_dims`. This allows for visual confirmation of the altered batch structure.

```python
bgen[0]
```