.. DO NOT EDIT.
.. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY.
.. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE:
.. "examples/16_cbottle_super_resolution.py"
.. LINE NUMBERS ARE GIVEN BELOW.

.. only:: html

    .. note::
        :class: sphx-glr-download-link-note

        :ref:`Go to the end <sphx_glr_download_examples_16_cbottle_super_resolution.py>`
        to download the full example code.

.. rst-class:: sphx-glr-example-title

.. _sphx_glr_examples_16_cbottle_super_resolution.py:


CBottle Super Resolution
========================

Climate in a Bottle (cBottle) super resolution workflows for global weather.

This example will demonstrate the cBottle diffusion model for super resolution of global
weather data. The CBottleSR model takes low-resolution climate data and generates
high-resolution outputs using a diffusion-based approach.

For more information on cBottle see:

- https://arxiv.org/abs/2505.06474v1

In this example you will learn:

- Performing super resolution on synthetic data from the cBottle data source
- Performing super resolution on ERA5 data after infilling with cBottle
- Post-processing and visualizing super-resolution results

.. GENERATED FROM PYTHON SOURCE LINES 40-47

Set Up
------
For this example we will use the cBottle data source, infill diagnostic, and the
CBottleSR super resolution model. The workflow demonstrates two approaches:

1. Super resolution on synthetic data generated by cBottle3D
2. Super resolution on real ERA5 data after variable infilling

.. GENERATED FROM PYTHON SOURCE LINES 49-55

We need the following components:

- Datasource: Generate data from the CBottle3D data api :py:class:`earth2studio.data.CBottle3D`.
- Datasource: Pull data from the WeatherBench2 data api :py:class:`earth2studio.data.WB2ERA5`.
- Diagnostic Model: Use the built in CBottle Infill Model :py:class:`earth2studio.models.dx.CBottleInfill`.
- Super Resolution Model: Use the CBottleSR super resolution model :py:class:`earth2studio.models.dx.CBottleSR`.

.. GENERATED FROM PYTHON SOURCE LINES 57-104

.. code-block:: Python


    import os

    os.makedirs("outputs", exist_ok=True)
    from dotenv import load_dotenv

    load_dotenv()  # TODO: make common example prep function

    import torch

    from earth2studio.data import WB2ERA5, CBottle3D, fetch_data
    from earth2studio.models.dx import CBottleInfill, CBottleSR

    # Get the device
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # Load the cBottle data source
    package = CBottle3D.load_default_package()
    cbottle_ds = CBottle3D.load_model(package)
    cbottle_ds = cbottle_ds.to(device)

    # Load the super resolution model
    super_resolution_window = (
        0,
        -120,
        50,
        -40,
    )  # (lat south, lon west, lat north, lon east)
    package = CBottleSR.load_default_package()
    cbottle_sr = CBottleSR.load_model(
        package,
        output_resolution=(1024, 1024),
        super_resolution_window=super_resolution_window,
    )
    cbottle_sr = cbottle_sr.to(device)

    # Load the infill model
    input_variables = ["u10m", "v10m"]
    package = CBottleInfill.load_default_package()
    cbottle_infill = CBottleInfill.load_model(
        package, input_variables=input_variables, sampler_steps=18
    )
    cbottle_infill = cbottle_infill.to(device)

    # Load the ERA5 data source
    era5_ds = WB2ERA5()


.. GENERATED FROM PYTHON SOURCE LINES 105-110

Super Resolution on Synthetic Data
----------------------------------
First, let's generate synthetic climate data using the cBottle3D data source
and then perform super resolution on it. This demonstrates the full cBottle
pipeline from data generation to super-resolution enhancement.

.. GENERATED FROM PYTHON SOURCE LINES 112-133

.. code-block:: Python


    import datetime

    import numpy as np

    from earth2studio.utils.coords import map_coords

    times = np.array([datetime.datetime(2020, 1, 1)], dtype="datetime64[ns]")

    # Generate some samples from cBottle
    synth_x, synth_coords = fetch_data(
        cbottle_ds,
        times,
        cbottle_sr.input_coords()["variable"],
        device=device,
    )

    # Perform super resolution on synthetic data
    synth_x, synth_coords = map_coords(synth_x, synth_coords, cbottle_sr.input_coords())
    sr_synth_x, sr_synth_coords = cbottle_sr(synth_x, synth_coords)


.. GENERATED FROM PYTHON SOURCE LINES 134-139

Super Resolution on ERA5 Data
-----------------------------
Next, we'll demonstrate super resolution on real ERA5 data. Since ERA5 doesn't
contain all the variables needed by CBottleSR, we first use the CBottleInfill
model to predict the missing variables, then perform super resolution.

.. GENERATED FROM PYTHON SOURCE LINES 141-157

.. code-block:: Python


    # Get the ERA5 data (only u10m and v10m available)
    era5_x, era5_coords = fetch_data(
        era5_ds,
        times,
        input_variables,
        device=device,
    )

    # Perform infilling to get all required variables
    infill_x, infill_coords = cbottle_infill(era5_x, era5_coords)

    # Select the required variables and reshape for super resolution
    infill_x, infill_coords = map_coords(infill_x, infill_coords, cbottle_sr.input_coords())
    sr_infill_x, sr_infill_coords = cbottle_sr(infill_x, infill_coords)


.. GENERATED FROM PYTHON SOURCE LINES 158-163

Post Processing CBottle Super Resolution Data
---------------------------------------------
Let's visualize the super resolution results to compare the synthetic data
approach with the ERA5 infilling approach. We'll plot the total cloud liquid
water (tclw) variable as an example.

.. GENERATED FROM PYTHON SOURCE LINES 165-250

.. code-block:: Python

    import cartopy.crs as ccrs
    import matplotlib.pyplot as plt

    plt.close("all")

    # Create projection focused on a region of interest (North Atlantic/Europe)
    projection = ccrs.PlateCarree()
    fig = plt.figure(figsize=(24, 24))

    # Plot synthetic data
    ax0 = fig.add_subplot(2, 2, 1, projection=projection)
    extent = [
        super_resolution_window[1] - 10,
        super_resolution_window[3] + 10,
        super_resolution_window[0] - 10,
        super_resolution_window[2] + 10,
    ]
    ax0.set_extent(extent, crs=ccrs.PlateCarree())
    c = ax0.pcolormesh(
        synth_coords["lon"],
        synth_coords["lat"],
        synth_x[0, 0, 3, :, :].cpu().numpy(),  # u10m (variable index 3)
        transform=ccrs.PlateCarree(),
        cmap="RdBu_r",
        vmin=-20,
        vmax=20,
    )
    plt.colorbar(c, ax=ax0, shrink=0.6, label="u10m (m/s)")
    ax0.coastlines()
    ax0.gridlines(draw_labels=True)
    ax0.set_title("Synthetic Data (cBottle3D, low resolution)")

    # Plot the synthetic super resolution data
    ax1 = fig.add_subplot(2, 2, 2, projection=projection)
    ax1.set_extent(extent, crs=ccrs.PlateCarree())
    c = ax1.pcolormesh(
        sr_synth_coords["lon"],
        sr_synth_coords["lat"],
        sr_synth_x[0, 0, 3, :, :].cpu().numpy(),  # u10m (variable index 3)
        transform=ccrs.PlateCarree(),
        cmap="RdBu_r",
        vmin=-20,
        vmax=20,
    )
    plt.colorbar(c, ax=ax1, shrink=0.6, label="u10m (m/s)")
    ax1.coastlines()
    ax1.gridlines(draw_labels=True)
    ax1.set_title("Synthetic Data Super Resolution (cBottle3D → CBottleSR)")

    # Plot the ERA5 data
    ax2 = fig.add_subplot(2, 2, 3, projection=projection)
    ax2.set_extent(extent, crs=ccrs.PlateCarree())
    c = ax2.pcolormesh(
        era5_coords["lon"],
        era5_coords["lat"],
        era5_x[0, 0, 0, :, :].cpu().numpy(),  # u10m (variable index 0)
        transform=ccrs.PlateCarree(),
        cmap="RdBu_r",
        vmin=-20,
        vmax=20,
    )
    plt.colorbar(c, ax=ax2, shrink=0.6, label="u10m (m/s)")
    ax2.coastlines()
    ax2.gridlines(draw_labels=True)
    ax2.set_title("ERA5 Data (low resolution)")

    # Plot the ERA5 infilled super resolution data
    ax3 = fig.add_subplot(2, 2, 4, projection=projection)
    ax3.set_extent(extent, crs=ccrs.PlateCarree())
    c = ax3.pcolormesh(
        sr_infill_coords["lon"],
        sr_infill_coords["lat"],
        sr_infill_x[0, 0, 3, :, :].cpu().numpy(),  # u10m (variable index 3)
        transform=ccrs.PlateCarree(),
        cmap="RdBu_r",
        vmin=-20,
        vmax=20,
    )
    plt.colorbar(c, ax=ax3, shrink=0.6, label="u10m (m/s)")
    ax3.coastlines()
    ax3.gridlines(draw_labels=True)
    ax3.set_title("ERA5 Infilled Super Resolution (ERA5 → CBottleInfill → CBottleSR)")

    plt.tight_layout()
    plt.savefig("outputs/16_cbottle_super_resolution.jpg", dpi=150, bbox_inches="tight")


.. _sphx_glr_download_examples_16_cbottle_super_resolution.py:

.. only:: html

  .. container:: sphx-glr-footer sphx-glr-footer-example

    .. container:: sphx-glr-download sphx-glr-download-jupyter

      :download:`Download Jupyter notebook: 16_cbottle_super_resolution.ipynb <16_cbottle_super_resolution.ipynb>`

    .. container:: sphx-glr-download sphx-glr-download-python

      :download:`Download Python source code: 16_cbottle_super_resolution.py <16_cbottle_super_resolution.py>`

    .. container:: sphx-glr-download sphx-glr-download-zip

      :download:`Download zipped: 16_cbottle_super_resolution.zip <16_cbottle_super_resolution.zip>`


.. only:: html

 .. rst-class:: sphx-glr-signature

    `Gallery generated by Sphinx-Gallery <https://sphinx-gallery.github.io>`_