CBottle Data Generation and Infilling¶
Climate in a Bottle (cBottle) inference workflows for global weather data synthesis.
This example will demonstrate the cBottle diffusion model data source and infilling diagnostic model for generating global climate and weather data. Both the cBottle data source and infilling diagnostic use the same diffusion model but the sampling procedure is different enabling two unique modes of interaction.
For more information on cBottle see:
In this example you will learn:
- Generating synthetic climate data with cBottle data source
- Instantiating cBottle infill diagnostic model
- Creating a simple infilling inference workflow
Set Up¶
For this example we will use the cBottle data source and infill diagnostic. Unlike other data sources the cBottle3D data source needs to be loaded similar to a prognostic or diagnostic model.
Thus, we need the following:
- Datasource: Generate data from the CBottle3D data api
earth2studio.data.CBottle3D. - Datasource: Pull data from the WeatherBench2 data api
earth2studio.data.WB2ERA5. - Diagnostic Model: Use the built in CBottle Infill Model
earth2studio.models.dx.CBottleInfill.
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
from earth2studio.models.dx import CBottleInfill
# Load the default model package which downloads the check point from NGC
package = CBottle3D.load_default_package()
cbottle_ds = CBottle3D.load_model(package, seed=None)
# This is an AI data source, so also move it to device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
cbottle_ds = cbottle_ds.to(device)
# Create the ground truth data source
era5_ds = WB2ERA5()
Console output9 lines
/__w/earth2studio/earth2studio/.venv/lib/python3.13/site-packages/torch/cuda/__init__.py:64: FutureWarning: The pynvml package is deprecated. Please install nvidia-ml-py instead. If you did not install pynvml directly, please report this to the maintainers of the package that installed pynvml for you.
import pynvml # type: ignore[import]
WARNING[XFORMERS]: xFormers can't load C++/CUDA extensions. xFormers was built for:
PyTorch 2.10.0+cu128 with CUDA 1208 (you have 2.13.0+cu130)
Python 3.10.19 (you have 3.13.13)
Please reinstall xformers (see https://github.com/facebookresearch/xformers#installing-xformers)
Memory-efficient attention, SwiGLU, sparse and more won't be available.
Set XFORMERS_MORE_DETAILS=1 for more details
CuPy distance computation test failed with error: cuVS >= 24.12 or pylibraft < 24.12 should be installed to use this featureGenerating Synthetic Weather Data¶
Once loaded, generating data from cBottle is as easy as any other data source. Under the hood the model is conditioned on the timestamp requested as well as a mid-month SST field which is internally handle for users but limits the range of the data source to years between 1970 and 2022.
Note that this diffusion model is stochastic, so querying the same timestamp will generate different fields that are reflective of the requested time and SST state.
from datetime import datetime
n_samples = 5
timestamp = datetime(2022, 9, 5)
# Fetch the ground truth
era5_da = era5_ds([timestamp], ["msl", "tcwv"])
# Generate some samples from cBottle
cbottle_da = cbottle_ds([timestamp for i in range(n_samples)], ["msl", "tcwv"])
print(era5_da)
print(cbottle_da)
Console output88 lines
Fetching WB2 data: 0%| | 0/2 [00:00<?, ?it/s]
Fetching WB2 data: 50%|█████ | 1/2 [00:00<00:00, 9.18it/s]
Fetching WB2 data: 100%|██████████| 2/2 [00:00<00:00, 12.40it/s]
Generating cBottle Data: 0%| | 0/2 [00:00<?, ?it/s]
Generating cBottle Data: 50%|█████ | 1/2 [00:09<00:09, 9.98s/it]
Generating cBottle Data: 100%|██████████| 2/2 [00:15<00:00, 7.37s/it]
Generating cBottle Data: 100%|██████████| 2/2 [00:15<00:00, 7.76s/it]
<xarray.DataArray (time: 1, variable: 2, lat: 721, lon: 1440)> Size: 17MB
array([[[[1.01743422e+05, 1.01743422e+05, 1.01743422e+05, ...,
1.01743422e+05, 1.01743422e+05, 1.01743422e+05],
[1.01793445e+05, 1.01793289e+05, 1.01793289e+05, ...,
1.01794055e+05, 1.01793750e+05, 1.01793750e+05],
[1.01839055e+05, 1.01838445e+05, 1.01838297e+05, ...,
1.01839969e+05, 1.01839508e+05, 1.01839203e+05],
...,
[1.00018555e+05, 1.00018555e+05, 1.00018555e+05, ...,
1.00017945e+05, 1.00017945e+05, 1.00018250e+05],
[9.98792891e+04, 9.98789844e+04, 9.98789844e+04, ...,
9.98786797e+04, 9.98789844e+04, 9.98789844e+04],
[9.98712266e+04, 9.98712266e+04, 9.98712266e+04, ...,
9.98712266e+04, 9.98712266e+04, 9.98712266e+04]],
[[1.33166103e+01, 1.33166103e+01, 1.33166103e+01, ...,
1.33166103e+01, 1.33166103e+01, 1.33166103e+01],
[1.31231613e+01, 1.31231613e+01, 1.31231613e+01, ...,
1.31190758e+01, 1.31217995e+01, 1.31217995e+01],
[1.28479767e+01, 1.28479767e+01, 1.28493385e+01, ...,
1.28357162e+01, 1.28384399e+01, 1.28425274e+01],
...,
[2.75253296e-01, 2.75253296e-01, 2.75253296e-01, ...,
2.75253296e-01, 2.75253296e-01, 2.75253296e-01],
[2.75253296e-01, 2.75253296e-01, 2.75253296e-01, ...,
2.75253296e-01, 2.75253296e-01, 2.75253296e-01],
[2.73891449e-01, 2.73891449e-01, 2.73891449e-01, ...,
2.73891449e-01, 2.73891449e-01, 2.73891449e-01]]]])
Coordinates:
* time (time) datetime64[ns] 8B 2022-09-05
* variable (variable) <U4 32B 'msl' 'tcwv'
* lat (lat) float64 6kB 90.0 89.75 89.5 89.25 ... -89.5 -89.75 -90.0
* lon (lon) float64 12kB 0.0 0.25 0.5 0.75 ... 359.0 359.2 359.5 359.8
<xarray.DataArray (time: 5, variable: 2, lat: 721, lon: 1440)> Size: 83MB
array([[[[ 1.00616930e+05, 1.00616930e+05, 1.00616930e+05, ...,
1.00616930e+05, 1.00616930e+05, 1.00616930e+05],
[ 1.00632123e+05, 1.00632189e+05, 1.00632254e+05, ...,
1.00631928e+05, 1.00631993e+05, 1.00632058e+05],
[ 1.00647317e+05, 1.00647448e+05, 1.00647578e+05, ...,
1.00646927e+05, 1.00647057e+05, 1.00647187e+05],
...,
[ 1.00454323e+05, 1.00453190e+05, 1.00452057e+05, ...,
1.00457722e+05, 1.00456589e+05, 1.00455456e+05],
[ 1.00481323e+05, 1.00480757e+05, 1.00480190e+05, ...,
1.00483023e+05, 1.00482456e+05, 1.00481890e+05],
[ 1.00508323e+05, 1.00508323e+05, 1.00508323e+05, ...,
1.00508323e+05, 1.00508323e+05, 1.00508323e+05]],
[[ 7.47921266e+00, 7.47921266e+00, 7.47921266e+00, ...,
7.47921266e+00, 7.47921266e+00, 7.47921266e+00],
[ 7.40058556e+00, 7.40120249e+00, 7.40181941e+00, ...,
7.39873479e+00, 7.39935171e+00, 7.39996864e+00],
[ 7.32195847e+00, 7.32319232e+00, 7.32442617e+00, ...,
7.31825691e+00, 7.31949076e+00, 7.32072461e+00],
...
[ 9.96921631e+04, 9.96926044e+04, 9.96930457e+04, ...,
9.96908392e+04, 9.96912805e+04, 9.96917218e+04],
[ 9.96236235e+04, 9.96238441e+04, 9.96240648e+04, ...,
9.96229615e+04, 9.96231822e+04, 9.96234028e+04],
[ 9.95550838e+04, 9.95550838e+04, 9.95550838e+04, ...,
9.95550838e+04, 9.95550838e+04, 9.95550838e+04]],
[[ 3.96856767e+00, 3.96856767e+00, 3.96856767e+00, ...,
3.96856767e+00, 3.96856767e+00, 3.96856767e+00],
[ 4.11879121e+00, 4.11724743e+00, 4.11570365e+00, ...,
4.12342254e+00, 4.12187876e+00, 4.12033499e+00],
[ 4.26901475e+00, 4.26592720e+00, 4.26283964e+00, ...,
4.27827742e+00, 4.27518986e+00, 4.27210231e+00],
...,
[-2.09590869e+00, -2.09076396e+00, -2.08561923e+00, ...,
-2.11134289e+00, -2.10619816e+00, -2.10105343e+00],
[-2.24401055e+00, -2.24143819e+00, -2.23886582e+00, ...,
-2.25172765e+00, -2.24915528e+00, -2.24658292e+00],
[-2.39211241e+00, -2.39211241e+00, -2.39211241e+00, ...,
-2.39211241e+00, -2.39211241e+00, -2.39211241e+00]]]])
Coordinates:
* time (time) datetime64[ns] 40B 2022-09-05 2022-09-05 ... 2022-09-05
* variable (variable) <U4 32B 'msl' 'tcwv'
* lat (lat) float64 6kB 90.0 89.75 89.5 89.25 ... -89.5 -89.75 -90.0
* lon (lon) float64 12kB 0.0 0.25 0.5 0.75 ... 359.0 359.2 359.5 359.8Post Processing CBottle Data¶
Let's visualize this data to better understand what the cBottle data source is able to provide. It is clear that each sample is indeed unique, yet remains physically realizable. In other words the cBottle data source can be used to create climates that do not exist but could based on the conditional distribution learned from the training data.
import cartopy.crs as ccrs
import matplotlib.pyplot as plt
variable = "tcwv"
plt.close("all")
projection = ccrs.Orthographic(central_longitude=300.0)
# Create a figure and axes with the specified projection
fig, ax = plt.subplots(2, 3, subplot_kw={"projection": projection}, figsize=(11, 6))
ax = ax.flatten()
ax[0].pcolormesh(
era5_da.coords["lon"],
era5_da.coords["lat"],
era5_da.sel(variable=variable).isel(time=0),
transform=ccrs.PlateCarree(),
cmap="cubehelix",
)
ax[0].set_title("ERA5")
for i in range(n_samples):
ax[i + 1].pcolormesh(
cbottle_da.coords["lon"],
cbottle_da.coords["lat"],
cbottle_da.sel(variable=variable).isel(time=i),
transform=ccrs.PlateCarree(),
cmap="cubehelix",
vmin=0,
vmax=90,
)
ax[i + 1].set_title(f"CBottle Sample {i}")
for ax0 in ax:
ax0.coastlines()
ax0.gridlines()
plt.tight_layout()
plt.savefig("outputs/15_tcwv_cbottle_datasource.jpg")
Variable Infilling with CBottleInfill Diagnostic¶
Next lets look at using the same model but for variable infilling. CBottleInfill allows users to generate global weather fields like the data source but condition it on a set of input fields that can be configured. This means that this diagnostic is extremely flexible and can be used with all types of data sources and models.
To demonstrate this lets consider two instances of the infilling diagnostic with a different set of inputs and then compare the resulting infilled variables. Note that the outputs of both configurations are the same size with the same variables.
import numpy as np
from earth2studio.data.utils import fetch_data
# Input variables
input_variables = ["u10m", "v10m"]
# Load the default model package which downloads the check point from NGC
package = CBottleInfill.load_default_package()
model = CBottleInfill.load_model(package, input_variables=input_variables)
model = model.to(device)
torch.manual_seed(0)
torch.cuda.manual_seed(0)
times = np.array([timestamp] * n_samples, dtype="datetime64[ns]")
x, coords = fetch_data(era5_ds, times, input_variables, device=device)
output_0, output_coords = model(x, coords)
print(output_0.shape)
Console output4 lines
Fetching WB2 data: 0%| | 0/10 [00:00<?, ?it/s]
Fetching WB2 data: 10%|█ | 1/10 [00:00<00:01, 7.62it/s]
Fetching WB2 data: 100%|██████████| 10/10 [00:00<00:00, 61.38it/s]
torch.Size([5, 1, 45, 721, 1440])Now repeat the process above but with an expanded set of variables. In this instance we provide a lot more data to the model to condition it with more information.
input_variables = [
"u10m",
"v10m",
"t2m",
"msl",
"z50",
"u50",
"v50",
"z500",
"u500",
"v500",
"z1000",
"u1000",
"v1000",
]
# Load the default model package which downloads the check point from NGC
package = CBottleInfill.load_default_package()
model = CBottleInfill.load_model(package, input_variables=input_variables)
model = model.to(device)
torch.manual_seed(0)
torch.cuda.manual_seed(0)
x, coords = fetch_data(era5_ds, times, input_variables, device=device)
output_1, output_coords = model(x, coords)
print(output_1.shape)
Console output12 lines
Fetching WB2 data: 0%| | 0/65 [00:00<?, ?it/s]
Fetching WB2 data: 3%|▎ | 2/65 [00:00<00:03, 18.63it/s]
Fetching WB2 data: 32%|███▏ | 21/65 [00:00<00:01, 25.80it/s]
Fetching WB2 data: 42%|████▏ | 27/65 [00:00<00:01, 30.25it/s]
Fetching WB2 data: 51%|█████ | 33/65 [00:01<00:00, 34.85it/s]
Fetching WB2 data: 62%|██████▏ | 40/65 [00:01<00:00, 41.62it/s]
Fetching WB2 data: 71%|███████ | 46/65 [00:01<00:00, 45.33it/s]
Fetching WB2 data: 80%|████████ | 52/65 [00:01<00:00, 48.42it/s]
Fetching WB2 data: 89%|████████▉ | 58/65 [00:01<00:00, 50.70it/s]
Fetching WB2 data: 98%|█████████▊| 64/65 [00:01<00:00, 52.92it/s]
Fetching WB2 data: 100%|██████████| 65/65 [00:01<00:00, 40.87it/s]
torch.Size([5, 1, 45, 721, 1440])Post Processing CBottleInfill¶
To post process the results, we take a look at a infilled variable, total column water vapour. Compared to the samples from the CBottle3D, the results are much more aligned with the ground truth since the infill model is sampling a conditional distribution. Additionally, the model provided more variables is better aligned with the ground truth due the additional information provided.
variable = "tcwv"
var_idx = np.where(output_coords["variable"] == "tcwv")[0][0]
era5_data, _ = fetch_data(era5_ds, times[:1], [variable], device=device)
plt.close("all")
projection = ccrs.Mollweide(central_longitude=0)
# Create a figure and axes with the specified projection
fig, ax = plt.subplots(2, 3, subplot_kw={"projection": projection}, figsize=(10, 6))
def plot_contour(
ax0: plt.axes,
data: torch.Tensor,
cmap: str = "jet",
vrange: tuple[int, int] = (0, 90),
) -> None:
"""Contour helper"""
ax0.contourf(
output_coords["lon"],
output_coords["lat"],
data.cpu(),
vmin=vrange[0],
vmax=vrange[1],
transform=ccrs.PlateCarree(),
levels=12,
cmap=cmap,
)
ax0.coastlines()
ax0.gridlines()
plot_contour(ax[0, 0], era5_data[0, 0, 0])
plot_contour(ax[0, 1], torch.mean(output_0[:, 0, var_idx], axis=0))
plot_contour(ax[0, 2], torch.mean(output_1[:, 0, var_idx], axis=0))
plot_contour(
ax[1, 1], torch.std(output_0[:, 0, var_idx], axis=0), cmap="inferno", vrange=(0, 10)
)
plot_contour(
ax[1, 2], torch.std(output_1[:, 0, var_idx], axis=0), cmap="inferno", vrange=(0, 10)
)
ax[0, 0].set_title("ERA5")
ax[0, 1].set_title("3 Input Variables Mean")
ax[0, 2].set_title("13 Input Variables Mean")
ax[1, 1].set_title("3 Input Variables Std")
ax[1, 2].set_title("13 Input Variables Std")
plt.tight_layout()
plt.savefig("outputs/15_tcwv_cbottle_infill.jpg")
Console output2 lines
Fetching WB2 data: 0%| | 0/1 [00:00<?, ?it/s]
Fetching WB2 data: 100%|██████████| 1/1 [00:00<00:00, 213.95it/s]
Execution profile
Runtime telemetry
Total runtime2m 11s
Execution environment
CPUAMD EPYC 7313P 16-Core Processor
GPUNVIDIA H100 PCIe · 79.6 GiB
System RAM58.5 GiB
PlatformLinux 6.8.0-136-generic
Python3.13.13
GPU driver / CUDADriver 595.84 · CUDA support 13.2