.. 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 38-45

.. code-block:: Python

    # /// script
    # dependencies = [
    #   "earth2studio[cbottle] @ git+https://github.com/NVIDIA/earth2studio.git",
    #   "cartopy",
    # ]
    # ///








.. GENERATED FROM PYTHON SOURCE LINES 46-53

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 55-61

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 63-110

.. 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()





.. rst-class:: sphx-glr-script-out

 .. code-block:: none


    Downloading cBottle-SR.zip: 0%|          | 0.00/2.46G [00:00<?, ?B/s]
    Downloading cBottle-SR.zip: 1%|▏         | 34.5M/2.46G [00:00<00:07, 362MB/s]
    Downloading cBottle-SR.zip: 3%|▎         | 87.0M/2.46G [00:00<00:05, 471MB/s]
    Downloading cBottle-SR.zip: 5%|▌         | 133M/2.46G [00:00<00:05, 475MB/s] 
    Downloading cBottle-SR.zip: 8%|▊         | 190M/2.46G [00:00<00:04, 523MB/s]
    Downloading cBottle-SR.zip: 10%|▉         | 240M/2.46G [00:00<00:04, 484MB/s]
    Downloading cBottle-SR.zip: 11%|█▏        | 287M/2.46G [00:00<00:04, 482MB/s]
    Downloading cBottle-SR.zip: 14%|█▎        | 344M/2.46G [00:00<00:04, 517MB/s]
    Downloading cBottle-SR.zip: 16%|█▌        | 395M/2.46G [00:00<00:04, 524MB/s]
    Downloading cBottle-SR.zip: 18%|█▊        | 446M/2.46G [00:00<00:04, 522MB/s]
    Downloading cBottle-SR.zip: 20%|█▉        | 497M/2.46G [00:01<00:04, 524MB/s]
    Downloading cBottle-SR.zip: 22%|██▏       | 547M/2.46G [00:01<00:03, 520MB/s]
    Downloading cBottle-SR.zip: 24%|██▎       | 597M/2.46G [00:01<00:04, 498MB/s]
    Downloading cBottle-SR.zip: 26%|██▌       | 644M/2.46G [00:01<00:04, 491MB/s]
    Downloading cBottle-SR.zip: 27%|██▋       | 692M/2.46G [00:01<00:04, 476MB/s]
    Downloading cBottle-SR.zip: 29%|██▉       | 739M/2.46G [00:01<00:03, 480MB/s]
    Downloading cBottle-SR.zip: 31%|███       | 785M/2.46G [00:01<00:04, 449MB/s]
    Downloading cBottle-SR.zip: 33%|███▎      | 830M/2.46G [00:01<00:03, 456MB/s]
    Downloading cBottle-SR.zip: 35%|███▍      | 874M/2.46G [00:01<00:03, 440MB/s]
    Downloading cBottle-SR.zip: 36%|███▋      | 919M/2.46G [00:02<00:03, 449MB/s]
    Downloading cBottle-SR.zip: 38%|███▊      | 962M/2.46G [00:02<00:03, 426MB/s]
    Downloading cBottle-SR.zip: 40%|███▉      | 0.98G/2.46G [00:02<00:03, 436MB/s]
    Downloading cBottle-SR.zip: 42%|████▏     | 1.02G/2.46G [00:02<00:03, 406MB/s]
    Downloading cBottle-SR.zip: 44%|████▎     | 1.07G/2.46G [00:02<00:03, 435MB/s]
    Downloading cBottle-SR.zip: 45%|████▌     | 1.11G/2.46G [00:02<00:03, 442MB/s]
    Downloading cBottle-SR.zip: 47%|████▋     | 1.16G/2.46G [00:02<00:03, 429MB/s]
    Downloading cBottle-SR.zip: 49%|████▊     | 1.20G/2.46G [00:02<00:03, 408MB/s]
    Downloading cBottle-SR.zip: 50%|█████     | 1.24G/2.46G [00:02<00:03, 433MB/s]
    Downloading cBottle-SR.zip: 52%|█████▏    | 1.29G/2.46G [00:02<00:02, 447MB/s]
    Downloading cBottle-SR.zip: 54%|█████▍    | 1.34G/2.46G [00:03<00:02, 465MB/s]
    Downloading cBottle-SR.zip: 56%|█████▌    | 1.38G/2.46G [00:03<00:02, 478MB/s]
    Downloading cBottle-SR.zip: 58%|█████▊    | 1.43G/2.46G [00:03<00:02, 496MB/s]
    Downloading cBottle-SR.zip: 60%|██████    | 1.48G/2.46G [00:03<00:02, 496MB/s]
    Downloading cBottle-SR.zip: 62%|██████▏   | 1.53G/2.46G [00:03<00:02, 480MB/s]
    Downloading cBottle-SR.zip: 64%|██████▍   | 1.58G/2.46G [00:03<00:01, 499MB/s]
    Downloading cBottle-SR.zip: 66%|██████▌   | 1.62G/2.46G [00:03<00:01, 500MB/s]
    Downloading cBottle-SR.zip: 68%|██████▊   | 1.68G/2.46G [00:03<00:01, 532MB/s]
    Downloading cBottle-SR.zip: 71%|███████   | 1.74G/2.46G [00:03<00:01, 549MB/s]
    Downloading cBottle-SR.zip: 73%|███████▎  | 1.79G/2.46G [00:03<00:01, 555MB/s]
    Downloading cBottle-SR.zip: 75%|███████▍  | 1.84G/2.46G [00:04<00:01, 511MB/s]
    Downloading cBottle-SR.zip: 77%|███████▋  | 1.89G/2.46G [00:04<00:01, 515MB/s]
    Downloading cBottle-SR.zip: 79%|███████▉  | 1.94G/2.46G [00:04<00:01, 518MB/s]
    Downloading cBottle-SR.zip: 81%|████████  | 2.00G/2.46G [00:04<00:00, 550MB/s]
    Downloading cBottle-SR.zip: 83%|████████▎ | 2.05G/2.46G [00:04<00:00, 532MB/s]
    Downloading cBottle-SR.zip: 85%|████████▌ | 2.10G/2.46G [00:04<00:00, 539MB/s]
    Downloading cBottle-SR.zip: 88%|████████▊ | 2.16G/2.46G [00:04<00:00, 552MB/s]
    Downloading cBottle-SR.zip: 90%|████████▉ | 2.21G/2.46G [00:04<00:00, 520MB/s]
    Downloading cBottle-SR.zip: 92%|█████████▏| 2.26G/2.46G [00:04<00:00, 532MB/s]
    Downloading cBottle-SR.zip: 94%|█████████▍| 2.31G/2.46G [00:05<00:00, 517MB/s]
    Downloading cBottle-SR.zip: 96%|█████████▌| 2.36G/2.46G [00:05<00:00, 521MB/s]
    Downloading cBottle-SR.zip: 98%|█████████▊| 2.41G/2.46G [00:05<00:00, 533MB/s]
    Downloading cBottle-SR.zip: 100%|██████████| 2.46G/2.46G [00:05<00:00, 494MB/s]




.. GENERATED FROM PYTHON SOURCE LINES 111-116

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 118-139

.. 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)





.. rst-class:: sphx-glr-script-out

 .. code-block:: none


    Generating cBottle Data:   0%|          | 0/1 [00:00<?, ?it/s]
    Generating cBottle Data: 100%|██████████| 1/1 [00:03<00:00,  3.95s/it]
    Generating cBottle Data: 100%|██████████| 1/1 [00:03<00:00,  3.95s/it]




.. GENERATED FROM PYTHON SOURCE LINES 140-145

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 147-163

.. 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)





.. rst-class:: sphx-glr-script-out

 .. code-block:: none


    Fetching WB2 data:   0%|          | 0/2 [00:00<?, ?it/s]
                                                        
    2025-11-12 19:13:54.043 | DEBUG    | earth2studio.data.wb2:fetch_array:241 - Fetching WB2 zarr array for variable: u10m at 2020-01-01T00:00:00

    Fetching WB2 data:   0%|          | 0/2 [00:00<?, ?it/s]
                                                        
    2025-11-12 19:13:54.043 | DEBUG    | earth2studio.data.wb2:fetch_array:241 - Fetching WB2 zarr array for variable: v10m at 2020-01-01T00:00:00

    Fetching WB2 data:   0%|          | 0/2 [00:00<?, ?it/s]
                                                        
    2025-11-12 19:13:54.099 | DEBUG    | earth2studio.data.utils:_make_local_details:657 - Copying /weatherbench2/datasets/era5/1959-2023_01_10-wb13-6h-1440x721_with_derived_variables.zarr/10m_v_component_of_wind/89120.0.0 to local cache

    Fetching WB2 data:   0%|          | 0/2 [00:00<?, ?it/s]
                                                        
    2025-11-12 19:13:54.106 | DEBUG    | earth2studio.data.utils:_make_local_details:657 - Copying /weatherbench2/datasets/era5/1959-2023_01_10-wb13-6h-1440x721_with_derived_variables.zarr/10m_u_component_of_wind/89120.0.0 to local cache

    Fetching WB2 data:   0%|          | 0/2 [00:00<?, ?it/s]
    Fetching WB2 data:  50%|█████     | 1/2 [00:00<00:00,  2.48it/s]
    Fetching WB2 data: 100%|██████████| 2/2 [00:00<00:00,  4.95it/s]




.. GENERATED FROM PYTHON SOURCE LINES 164-169

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 171-256

.. 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")



.. image-sg:: /examples/images/sphx_glr_16_cbottle_super_resolution_001.png
   :alt: Synthetic Data (cBottle3D, low resolution), Synthetic Data Super Resolution (cBottle3D → CBottleSR), ERA5 Data (low resolution), ERA5 Infilled Super Resolution (ERA5 → CBottleInfill → CBottleSR)
   :srcset: /examples/images/sphx_glr_16_cbottle_super_resolution_001.png
   :class: sphx-glr-single-img






.. rst-class:: sphx-glr-timing

   **Total running time of the script:** (3 minutes 46.978 seconds)


.. _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>`_