Note
Go to the end to download the full example code.
Statistical Inference#
This example will demonstrate how to run a simple inference workflow to generate a forecast and then to save a statistic of that data. There are a handful of built-in statistics available in earth2studio.statistics, but here we will demonstrate how to define a custom statistic and run inference.
In this example you will learn:
How to instantiate a built in prognostic model
Creating a data source and IO object
Create a custom statistic
Running a simple built in workflow
Post-processing results
Creating a statistical workflow#
Start with creating a simple inference workflow to use. We encourage
users to explore and experiment with their own custom workflows that borrow ideas from
built in workflows inside earth2studio.run or the examples.
Creating our own generalizable workflow to use with statistics is easy when we rely on the component interfaces defined in Earth2Studio (use dependency injection). Here we create a run method that accepts the following:
time: Input list of datetimes / strings to run inference for
nsteps: Number of forecast steps to predict
prognostic: Our initialized prognostic model
statistic: our custom statistic
data: Initialized data source to fetch initial conditions from
io: IOBackend
We do not run an ensemble inference workflow here, even though it is common for statistical inference. See ensemble examples for details on how to extend this example for that purpose.
import os
os.makedirs("outputs", exist_ok=True)
from dotenv import load_dotenv
load_dotenv() # TODO: make common example prep function
from datetime import datetime
import numpy as np
import pandas as pd
from loguru import logger
from tqdm import tqdm
from earth2studio.data import DataSource, fetch_data
from earth2studio.io import IOBackend
from earth2studio.models.px import PrognosticModel
from earth2studio.statistics import Statistic
from earth2studio.utils.coords import map_coords
from earth2studio.utils.time import to_time_array
logger.remove()
logger.add(lambda msg: tqdm.write(msg, end=""), colorize=True)
def run_stats(
time: list[str] | list[datetime] | list[np.datetime64],
nsteps: int,
nensemble: int,
prognostic: PrognosticModel,
statistic: Statistic,
data: DataSource,
io: IOBackend,
) -> IOBackend:
"""Simple statistics workflow
Parameters
----------
time : list[str] | list[datetime] | list[np.datetime64]
List of string, datetimes or np.datetime64
nsteps : int
Number of forecast steps
nensemble : int
Number of ensemble members to run inference for.
prognostic : PrognosticModel
Prognostic models
statistic : Statistic
Custom statistic to compute and write to IO.
data : DataSource
Data source
io : IOBackend
IO object
Returns
-------
IOBackend
Output IO object
"""
logger.info("Running simple statistics workflow!")
# Load model onto the device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
logger.info(f"Inference device: {device}")
prognostic = prognostic.to(device)
# Fetch data from data source and load onto device
time = to_time_array(time)
x, coords = fetch_data(
source=data,
time=time,
lead_time=prognostic.input_coords()["lead_time"],
variable=prognostic.input_coords()["variable"],
device=device,
)
logger.success(f"Fetched data from {data.__class__.__name__}")
# Set up IO backend
total_coords = coords.copy()
output_coords = prognostic.output_coords(prognostic.input_coords())
total_coords["lead_time"] = np.asarray(
[output_coords["lead_time"] * i for i in range(nsteps + 1)]
).flatten()
# Remove reduced dimensions from statistic
for d in statistic.reduction_dimensions:
total_coords.pop(d, None)
io.add_array(total_coords, str(statistic))
# Map lat and lon if needed
x, coords = map_coords(x, coords, prognostic.input_coords())
# Create prognostic iterator
model = prognostic.create_iterator(x, coords)
logger.info("Inference starting!")
with tqdm(total=nsteps + 1, desc="Running inference") as pbar:
for step, (x, coords) in enumerate(model):
s, coords = statistic(x, coords)
io.write(s, coords, str(statistic))
pbar.update(1)
if step == nsteps:
break
logger.success("Inference complete")
return io
Set Up#
With the statistical workflow defined, we now need to create the individual components.
We need the following:
Prognostic Model: Use the built in Pangu 24 hour model
earth2studio.models.px.Pangu24.statistic: We define our own statistic: the Southern Oscillation Index (SOI).
Datasource: Pull data from the GFS data api
earth2studio.data.GFS.IO Backend: Save the outputs into a NetCDF4 store
earth2studio.io.NetCDF4Backend.
from collections import OrderedDict
import numpy as np
import torch
from earth2studio.data import GFS
from earth2studio.io import NetCDF4Backend
from earth2studio.models.px import Pangu24
from earth2studio.utils.type import CoordSystem
# Load the default model package which downloads the check point from NGC
package = Pangu24.load_default_package()
model = Pangu24.load_model(package)
# Create the data source
data = GFS()
# Create the IO handler, store in memory
io = NetCDF4Backend(
file_name="outputs/soi.nc",
)
# Create the custom statistic
class SOI:
"""Custom metric calculation the Southern Oscillation Index.
SOI = ( standardized_tahiti_slp - standardized_darwin_slp ) / soi_normalization
soi_normalization = std( historical ( standardized_tahiti_slp - standardized_darwin_slp ) )
standardized_*_slp = (*_slp - climatological_mean_*_slp) / climatological_std_*_slp
Note
----
__str__
Name that will be applied to the output of this statistic, primarily for IO purposes.
reduction_dimensions
Dimensions that this statistic reduces over. This is used to help automatically determine
the output coordinates, primarily used for IO purposes.
"""
def __str__(self) -> str:
return "soi"
def __init__(
self,
):
# Read in Tahiti and Darwin SLP data
from modulus.utils.filesystem import _download_cached
file_path = _download_cached(
"http://data.longpaddock.qld.gov.au/SeasonalClimateOutlook/SouthernOscillationIndex/SOIDataFiles/DailySOI1933-1992Base.txt"
)
ds = pd.read_csv(file_path, sep=r"\s+")
dates = pd.date_range("1999-01-01", freq="d", periods=len(ds))
ds["date"] = dates
ds = ds.set_index("date")
ds = ds.drop(["Year", "Day", "SOI"], axis=1)
ds = ds.rolling(30, min_periods=1).mean().dropna()
self.climatological_means = torch.tensor(
ds.groupby(ds.index.month).mean().to_numpy(), dtype=torch.float32
)
self.climatological_std = torch.tensor(
ds.groupby(ds.index.month).std().to_numpy(), dtype=torch.float32
)
standardized = ds.groupby(ds.index.month).transform(
lambda x: (x - x.mean()) / x.std()
)
diff = standardized["Tahiti"] - standardized["Darwin"]
self.normalization = torch.tensor(
diff.groupby(ds.index.month).std().to_numpy(), dtype=torch.float32
)
self.tahiti_coords = {
"variable": np.array(["msl"]),
"lat": np.array([-17.65]),
"lon": np.array([210.57]),
}
self.darwin_coords = {
"variable": np.array(["msl"]),
"lat": np.array([-12.46]),
"lon": np.array([130.84]),
}
self.reduction_dimensions = list(self.tahiti_coords)
def __call__(
self, x: torch.Tensor, coords: CoordSystem
) -> tuple[torch.Tensor, CoordSystem]:
"""Computes the SOI given an input.
coords must be a superset of both
tahiti_coords = {
'variable': np.array(['msl']),
'lat': np.array([-17.65]),
'lon': np.array([210.57])
}
and
darwin_coords = {
'variable': np.array(['msl']),
'lat': np.array([-12.46]),
'lon': np.array([130.84])
}
So make sure that the model chosen predicts the `msl` variable.
Parameters
----------
x : torch.Tensor
Input tensor
coords : CoordSystem
coordinate system belonging to the input tensor.
Returns
-------
tuple[torch.Tensor, CoordSystem]
Returns the SOI and appropriate coordinate system.
"""
tahiti, _ = map_coords(x, coords, self.tahiti_coords)
darwin, _ = map_coords(x, coords, self.darwin_coords)
tahiti = tahiti.squeeze(-3, -2, -1) / 100.0
darwin = darwin.squeeze(-3, -2, -1) / 100.0
output_coords = OrderedDict(
{k: v for k, v in coords.items() if k not in self.reduction_dimensions}
)
# Get time coordinates
times = coords["time"].reshape(-1, 1) + coords["lead_time"].reshape(1, -1)
months = torch.broadcast_to(
torch.as_tensor(
[pd.Timestamp(t).month for t in times.flatten()],
device=tahiti.device,
dtype=torch.int32,
).reshape(times.shape),
tahiti.shape,
)
cm = self.climatological_means.to(tahiti.device)
cs = self.climatological_std.to(tahiti.device)
norm = self.normalization.to(tahiti.device)
tahiti_std_anomaly = (tahiti - cm[months, 0]) / cs[months, 0]
darwin_std_anomaly = (tahiti - cm[months, 1]) / cs[months, 1]
return (tahiti_std_anomaly - darwin_std_anomaly) / norm[months], output_coords
soi = SOI()
Execute the Workflow#
With all components initialized, running the workflow is a single line of Python code. Workflow will return the provided IO object back to the user, which can be used to then post process. Some have additional APIs that can be handy for post-processing or saving to file. Check the API docs for more information. We simulate a trajectory of 60 time steps, or 2 months using Pangu24
nsteps = 60
nensemble = 1
io = run_stats(["2022-01-01"], nsteps, nensemble, model, soi, data, io)
2024-06-25 14:00:32.686 | INFO | __main__:run_stats:117 - Running simple statistics workflow!
2024-06-25 14:00:32.687 | INFO | __main__:run_stats:120 - Inference device: cuda
2024-06-25 14:00:49.187 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:149 - Fetching GFS index file: 2022-01-01 00:00:00
Fetching GFS for 2022-01-01 00:00:00: 0%| | 0/69 [00:00<?, ?it/s]
2024-06-25 14:00:49.527 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: z1000 at 2022-01-01 00:00:00
Fetching GFS for 2022-01-01 00:00:00: 0%| | 0/69 [00:00<?, ?it/s]
Fetching GFS for 2022-01-01 00:00:00: 1%|▏ | 1/69 [00:00<00:47, 1.43it/s]
2024-06-25 14:00:50.225 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: z925 at 2022-01-01 00:00:00
Fetching GFS for 2022-01-01 00:00:00: 1%|▏ | 1/69 [00:00<00:47, 1.43it/s]
Fetching GFS for 2022-01-01 00:00:00: 3%|▎ | 2/69 [00:01<00:34, 1.97it/s]
2024-06-25 14:00:50.601 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: z850 at 2022-01-01 00:00:00
Fetching GFS for 2022-01-01 00:00:00: 3%|▎ | 2/69 [00:01<00:34, 1.97it/s]
Fetching GFS for 2022-01-01 00:00:00: 4%|▍ | 3/69 [00:01<00:29, 2.22it/s]
2024-06-25 14:00:50.984 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: z700 at 2022-01-01 00:00:00
Fetching GFS for 2022-01-01 00:00:00: 4%|▍ | 3/69 [00:01<00:29, 2.22it/s]
Fetching GFS for 2022-01-01 00:00:00: 6%|▌ | 4/69 [00:01<00:28, 2.29it/s]
2024-06-25 14:00:51.396 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: z600 at 2022-01-01 00:00:00
Fetching GFS for 2022-01-01 00:00:00: 6%|▌ | 4/69 [00:01<00:28, 2.29it/s]
Fetching GFS for 2022-01-01 00:00:00: 7%|▋ | 5/69 [00:02<00:25, 2.49it/s]
2024-06-25 14:00:51.736 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: z500 at 2022-01-01 00:00:00
Fetching GFS for 2022-01-01 00:00:00: 7%|▋ | 5/69 [00:02<00:25, 2.49it/s]
Fetching GFS for 2022-01-01 00:00:00: 9%|▊ | 6/69 [00:02<00:23, 2.65it/s]
2024-06-25 14:00:52.067 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: z400 at 2022-01-01 00:00:00
Fetching GFS for 2022-01-01 00:00:00: 9%|▊ | 6/69 [00:02<00:23, 2.65it/s]
Fetching GFS for 2022-01-01 00:00:00: 10%|█ | 7/69 [00:02<00:23, 2.60it/s]
2024-06-25 14:00:52.469 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: z300 at 2022-01-01 00:00:00
Fetching GFS for 2022-01-01 00:00:00: 10%|█ | 7/69 [00:02<00:23, 2.60it/s]
Fetching GFS for 2022-01-01 00:00:00: 12%|█▏ | 8/69 [00:03<00:22, 2.67it/s]
2024-06-25 14:00:52.821 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: z250 at 2022-01-01 00:00:00
Fetching GFS for 2022-01-01 00:00:00: 12%|█▏ | 8/69 [00:03<00:22, 2.67it/s]
Fetching GFS for 2022-01-01 00:00:00: 13%|█▎ | 9/69 [00:03<00:23, 2.60it/s]
2024-06-25 14:00:53.226 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: z200 at 2022-01-01 00:00:00
Fetching GFS for 2022-01-01 00:00:00: 13%|█▎ | 9/69 [00:03<00:23, 2.60it/s]
Fetching GFS for 2022-01-01 00:00:00: 14%|█▍ | 10/69 [00:04<00:26, 2.23it/s]
2024-06-25 14:00:53.817 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: z150 at 2022-01-01 00:00:00
Fetching GFS for 2022-01-01 00:00:00: 14%|█▍ | 10/69 [00:04<00:26, 2.23it/s]
Fetching GFS for 2022-01-01 00:00:00: 16%|█▌ | 11/69 [00:04<00:26, 2.19it/s]
2024-06-25 14:00:54.292 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: z100 at 2022-01-01 00:00:00
Fetching GFS for 2022-01-01 00:00:00: 16%|█▌ | 11/69 [00:04<00:26, 2.19it/s]
Fetching GFS for 2022-01-01 00:00:00: 17%|█▋ | 12/69 [00:05<00:26, 2.16it/s]
2024-06-25 14:00:54.769 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: z50 at 2022-01-01 00:00:00
Fetching GFS for 2022-01-01 00:00:00: 17%|█▋ | 12/69 [00:05<00:26, 2.16it/s]
Fetching GFS for 2022-01-01 00:00:00: 19%|█▉ | 13/69 [00:05<00:24, 2.32it/s]
2024-06-25 14:00:55.129 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: q1000 at 2022-01-01 00:00:00
Fetching GFS for 2022-01-01 00:00:00: 19%|█▉ | 13/69 [00:05<00:24, 2.32it/s]
Fetching GFS for 2022-01-01 00:00:00: 20%|██ | 14/69 [00:05<00:23, 2.37it/s]
2024-06-25 14:00:55.527 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: q925 at 2022-01-01 00:00:00
Fetching GFS for 2022-01-01 00:00:00: 20%|██ | 14/69 [00:06<00:23, 2.37it/s]
Fetching GFS for 2022-01-01 00:00:00: 22%|██▏ | 15/69 [00:06<00:22, 2.39it/s]
2024-06-25 14:00:55.937 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: q850 at 2022-01-01 00:00:00
Fetching GFS for 2022-01-01 00:00:00: 22%|██▏ | 15/69 [00:06<00:22, 2.39it/s]
Fetching GFS for 2022-01-01 00:00:00: 23%|██▎ | 16/69 [00:06<00:21, 2.41it/s]
2024-06-25 14:00:56.344 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: q700 at 2022-01-01 00:00:00
Fetching GFS for 2022-01-01 00:00:00: 23%|██▎ | 16/69 [00:06<00:21, 2.41it/s]
Fetching GFS for 2022-01-01 00:00:00: 25%|██▍ | 17/69 [00:07<00:22, 2.36it/s]
2024-06-25 14:00:56.788 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: q600 at 2022-01-01 00:00:00
Fetching GFS for 2022-01-01 00:00:00: 25%|██▍ | 17/69 [00:07<00:22, 2.36it/s]
Fetching GFS for 2022-01-01 00:00:00: 26%|██▌ | 18/69 [00:07<00:22, 2.29it/s]
2024-06-25 14:00:57.258 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: q500 at 2022-01-01 00:00:00
Fetching GFS for 2022-01-01 00:00:00: 26%|██▌ | 18/69 [00:07<00:22, 2.29it/s]
Fetching GFS for 2022-01-01 00:00:00: 28%|██▊ | 19/69 [00:08<00:20, 2.38it/s]
2024-06-25 14:00:57.636 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: q400 at 2022-01-01 00:00:00
Fetching GFS for 2022-01-01 00:00:00: 28%|██▊ | 19/69 [00:08<00:20, 2.38it/s]
Fetching GFS for 2022-01-01 00:00:00: 29%|██▉ | 20/69 [00:08<00:19, 2.49it/s]
2024-06-25 14:00:57.997 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: q300 at 2022-01-01 00:00:00
Fetching GFS for 2022-01-01 00:00:00: 29%|██▉ | 20/69 [00:08<00:19, 2.49it/s]
Fetching GFS for 2022-01-01 00:00:00: 30%|███ | 21/69 [00:08<00:19, 2.41it/s]
2024-06-25 14:00:58.441 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: q250 at 2022-01-01 00:00:00
Fetching GFS for 2022-01-01 00:00:00: 30%|███ | 21/69 [00:08<00:19, 2.41it/s]
Fetching GFS for 2022-01-01 00:00:00: 32%|███▏ | 22/69 [00:09<00:18, 2.55it/s]
2024-06-25 14:00:58.780 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: q200 at 2022-01-01 00:00:00
Fetching GFS for 2022-01-01 00:00:00: 32%|███▏ | 22/69 [00:09<00:18, 2.55it/s]
Fetching GFS for 2022-01-01 00:00:00: 33%|███▎ | 23/69 [00:09<00:16, 2.74it/s]
2024-06-25 14:00:59.081 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: q150 at 2022-01-01 00:00:00
Fetching GFS for 2022-01-01 00:00:00: 33%|███▎ | 23/69 [00:09<00:16, 2.74it/s]
Fetching GFS for 2022-01-01 00:00:00: 35%|███▍ | 24/69 [00:10<00:18, 2.50it/s]
2024-06-25 14:00:59.564 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: q100 at 2022-01-01 00:00:00
Fetching GFS for 2022-01-01 00:00:00: 35%|███▍ | 24/69 [00:10<00:18, 2.50it/s]
Fetching GFS for 2022-01-01 00:00:00: 36%|███▌ | 25/69 [00:10<00:18, 2.32it/s]
2024-06-25 14:01:00.066 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: q50 at 2022-01-01 00:00:00
Fetching GFS for 2022-01-01 00:00:00: 36%|███▌ | 25/69 [00:10<00:18, 2.32it/s]
Fetching GFS for 2022-01-01 00:00:00: 38%|███▊ | 26/69 [00:10<00:18, 2.35it/s]
2024-06-25 14:01:00.477 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: t1000 at 2022-01-01 00:00:00
Fetching GFS for 2022-01-01 00:00:00: 38%|███▊ | 26/69 [00:10<00:18, 2.35it/s]
Fetching GFS for 2022-01-01 00:00:00: 39%|███▉ | 27/69 [00:11<00:18, 2.27it/s]
2024-06-25 14:01:00.954 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: t925 at 2022-01-01 00:00:00
Fetching GFS for 2022-01-01 00:00:00: 39%|███▉ | 27/69 [00:11<00:18, 2.27it/s]
Fetching GFS for 2022-01-01 00:00:00: 41%|████ | 28/69 [00:11<00:17, 2.38it/s]
2024-06-25 14:01:01.325 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: t850 at 2022-01-01 00:00:00
Fetching GFS for 2022-01-01 00:00:00: 41%|████ | 28/69 [00:11<00:17, 2.38it/s]
Fetching GFS for 2022-01-01 00:00:00: 42%|████▏ | 29/69 [00:12<00:17, 2.33it/s]
2024-06-25 14:01:01.777 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: t700 at 2022-01-01 00:00:00
Fetching GFS for 2022-01-01 00:00:00: 42%|████▏ | 29/69 [00:12<00:17, 2.33it/s]
Fetching GFS for 2022-01-01 00:00:00: 43%|████▎ | 30/69 [00:12<00:16, 2.33it/s]
2024-06-25 14:01:02.204 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: t600 at 2022-01-01 00:00:00
Fetching GFS for 2022-01-01 00:00:00: 43%|████▎ | 30/69 [00:12<00:16, 2.33it/s]
Fetching GFS for 2022-01-01 00:00:00: 45%|████▍ | 31/69 [00:13<00:15, 2.42it/s]
2024-06-25 14:01:02.583 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: t500 at 2022-01-01 00:00:00
Fetching GFS for 2022-01-01 00:00:00: 45%|████▍ | 31/69 [00:13<00:15, 2.42it/s]
Fetching GFS for 2022-01-01 00:00:00: 46%|████▋ | 32/69 [00:13<00:14, 2.49it/s]
2024-06-25 14:01:02.957 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: t400 at 2022-01-01 00:00:00
Fetching GFS for 2022-01-01 00:00:00: 46%|████▋ | 32/69 [00:13<00:14, 2.49it/s]
Fetching GFS for 2022-01-01 00:00:00: 48%|████▊ | 33/69 [00:13<00:14, 2.52it/s]
2024-06-25 14:01:03.343 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: t300 at 2022-01-01 00:00:00
Fetching GFS for 2022-01-01 00:00:00: 48%|████▊ | 33/69 [00:13<00:14, 2.52it/s]
Fetching GFS for 2022-01-01 00:00:00: 49%|████▉ | 34/69 [00:14<00:14, 2.40it/s]
2024-06-25 14:01:03.803 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: t250 at 2022-01-01 00:00:00
Fetching GFS for 2022-01-01 00:00:00: 49%|████▉ | 34/69 [00:14<00:14, 2.40it/s]
Fetching GFS for 2022-01-01 00:00:00: 51%|█████ | 35/69 [00:14<00:13, 2.46it/s]
2024-06-25 14:01:04.187 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: t200 at 2022-01-01 00:00:00
Fetching GFS for 2022-01-01 00:00:00: 51%|█████ | 35/69 [00:14<00:13, 2.46it/s]
Fetching GFS for 2022-01-01 00:00:00: 52%|█████▏ | 36/69 [00:15<00:13, 2.52it/s]
2024-06-25 14:01:04.564 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: t150 at 2022-01-01 00:00:00
Fetching GFS for 2022-01-01 00:00:00: 52%|█████▏ | 36/69 [00:15<00:13, 2.52it/s]
Fetching GFS for 2022-01-01 00:00:00: 54%|█████▎ | 37/69 [00:15<00:12, 2.58it/s]
2024-06-25 14:01:04.930 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: t100 at 2022-01-01 00:00:00
Fetching GFS for 2022-01-01 00:00:00: 54%|█████▎ | 37/69 [00:15<00:12, 2.58it/s]
Fetching GFS for 2022-01-01 00:00:00: 55%|█████▌ | 38/69 [00:15<00:11, 2.63it/s]
2024-06-25 14:01:05.291 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: t50 at 2022-01-01 00:00:00
Fetching GFS for 2022-01-01 00:00:00: 55%|█████▌ | 38/69 [00:15<00:11, 2.63it/s]
Fetching GFS for 2022-01-01 00:00:00: 57%|█████▋ | 39/69 [00:16<00:11, 2.54it/s]
2024-06-25 14:01:05.716 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: u1000 at 2022-01-01 00:00:00
Fetching GFS for 2022-01-01 00:00:00: 57%|█████▋ | 39/69 [00:16<00:11, 2.54it/s]
Fetching GFS for 2022-01-01 00:00:00: 58%|█████▊ | 40/69 [00:16<00:11, 2.59it/s]
2024-06-25 14:01:06.085 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: u925 at 2022-01-01 00:00:00
Fetching GFS for 2022-01-01 00:00:00: 58%|█████▊ | 40/69 [00:16<00:11, 2.59it/s]
Fetching GFS for 2022-01-01 00:00:00: 59%|█████▉ | 41/69 [00:17<00:11, 2.43it/s]
2024-06-25 14:01:06.559 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: u850 at 2022-01-01 00:00:00
Fetching GFS for 2022-01-01 00:00:00: 59%|█████▉ | 41/69 [00:17<00:11, 2.43it/s]
Fetching GFS for 2022-01-01 00:00:00: 61%|██████ | 42/69 [00:17<00:12, 2.11it/s]
2024-06-25 14:01:07.179 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: u700 at 2022-01-01 00:00:00
Fetching GFS for 2022-01-01 00:00:00: 61%|██████ | 42/69 [00:17<00:12, 2.11it/s]
Fetching GFS for 2022-01-01 00:00:00: 62%|██████▏ | 43/69 [00:18<00:13, 2.00it/s]
2024-06-25 14:01:07.739 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: u600 at 2022-01-01 00:00:00
Fetching GFS for 2022-01-01 00:00:00: 62%|██████▏ | 43/69 [00:18<00:13, 2.00it/s]
Fetching GFS for 2022-01-01 00:00:00: 64%|██████▍ | 44/69 [00:18<00:12, 1.97it/s]
2024-06-25 14:01:08.262 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: u500 at 2022-01-01 00:00:00
Fetching GFS for 2022-01-01 00:00:00: 64%|██████▍ | 44/69 [00:18<00:12, 1.97it/s]
Fetching GFS for 2022-01-01 00:00:00: 65%|██████▌ | 45/69 [00:19<00:12, 1.98it/s]
2024-06-25 14:01:08.762 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: u400 at 2022-01-01 00:00:00
Fetching GFS for 2022-01-01 00:00:00: 65%|██████▌ | 45/69 [00:19<00:12, 1.98it/s]
Fetching GFS for 2022-01-01 00:00:00: 67%|██████▋ | 46/69 [00:19<00:12, 1.90it/s]
2024-06-25 14:01:09.340 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: u300 at 2022-01-01 00:00:00
Fetching GFS for 2022-01-01 00:00:00: 67%|██████▋ | 46/69 [00:19<00:12, 1.90it/s]
Fetching GFS for 2022-01-01 00:00:00: 68%|██████▊ | 47/69 [00:20<00:11, 1.93it/s]
2024-06-25 14:01:09.841 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: u250 at 2022-01-01 00:00:00
Fetching GFS for 2022-01-01 00:00:00: 68%|██████▊ | 47/69 [00:20<00:11, 1.93it/s]
Fetching GFS for 2022-01-01 00:00:00: 70%|██████▉ | 48/69 [00:20<00:10, 1.96it/s]
2024-06-25 14:01:10.332 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: u200 at 2022-01-01 00:00:00
Fetching GFS for 2022-01-01 00:00:00: 70%|██████▉ | 48/69 [00:20<00:10, 1.96it/s]
Fetching GFS for 2022-01-01 00:00:00: 71%|███████ | 49/69 [00:21<00:10, 1.91it/s]
2024-06-25 14:01:10.889 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: u150 at 2022-01-01 00:00:00
Fetching GFS for 2022-01-01 00:00:00: 71%|███████ | 49/69 [00:21<00:10, 1.91it/s]
Fetching GFS for 2022-01-01 00:00:00: 72%|███████▏ | 50/69 [00:22<00:11, 1.60it/s]
2024-06-25 14:01:11.745 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: u100 at 2022-01-01 00:00:00
Fetching GFS for 2022-01-01 00:00:00: 72%|███████▏ | 50/69 [00:22<00:11, 1.60it/s]
Fetching GFS for 2022-01-01 00:00:00: 74%|███████▍ | 51/69 [00:22<00:09, 1.88it/s]
2024-06-25 14:01:12.060 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: u50 at 2022-01-01 00:00:00
Fetching GFS for 2022-01-01 00:00:00: 74%|███████▍ | 51/69 [00:22<00:09, 1.88it/s]
Fetching GFS for 2022-01-01 00:00:00: 75%|███████▌ | 52/69 [00:22<00:08, 1.98it/s]
2024-06-25 14:01:12.507 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: v1000 at 2022-01-01 00:00:00
Fetching GFS for 2022-01-01 00:00:00: 75%|███████▌ | 52/69 [00:22<00:08, 1.98it/s]
Fetching GFS for 2022-01-01 00:00:00: 77%|███████▋ | 53/69 [00:23<00:07, 2.27it/s]
2024-06-25 14:01:12.792 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: v925 at 2022-01-01 00:00:00
Fetching GFS for 2022-01-01 00:00:00: 77%|███████▋ | 53/69 [00:23<00:07, 2.27it/s]
Fetching GFS for 2022-01-01 00:00:00: 78%|███████▊ | 54/69 [00:23<00:06, 2.43it/s]
2024-06-25 14:01:13.136 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: v850 at 2022-01-01 00:00:00
Fetching GFS for 2022-01-01 00:00:00: 78%|███████▊ | 54/69 [00:23<00:06, 2.43it/s]
Fetching GFS for 2022-01-01 00:00:00: 80%|███████▉ | 55/69 [00:23<00:05, 2.72it/s]
2024-06-25 14:01:13.404 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: v700 at 2022-01-01 00:00:00
Fetching GFS for 2022-01-01 00:00:00: 80%|███████▉ | 55/69 [00:23<00:05, 2.72it/s]
Fetching GFS for 2022-01-01 00:00:00: 81%|████████ | 56/69 [00:24<00:04, 2.86it/s]
2024-06-25 14:01:13.711 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: v600 at 2022-01-01 00:00:00
Fetching GFS for 2022-01-01 00:00:00: 81%|████████ | 56/69 [00:24<00:04, 2.86it/s]
Fetching GFS for 2022-01-01 00:00:00: 83%|████████▎ | 57/69 [00:24<00:03, 3.03it/s]
2024-06-25 14:01:13.995 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: v500 at 2022-01-01 00:00:00
Fetching GFS for 2022-01-01 00:00:00: 83%|████████▎ | 57/69 [00:24<00:03, 3.03it/s]
Fetching GFS for 2022-01-01 00:00:00: 84%|████████▍ | 58/69 [00:24<00:03, 3.05it/s]
2024-06-25 14:01:14.317 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: v400 at 2022-01-01 00:00:00
Fetching GFS for 2022-01-01 00:00:00: 84%|████████▍ | 58/69 [00:24<00:03, 3.05it/s]
Fetching GFS for 2022-01-01 00:00:00: 86%|████████▌ | 59/69 [00:25<00:03, 3.01it/s]
2024-06-25 14:01:14.661 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: v300 at 2022-01-01 00:00:00
Fetching GFS for 2022-01-01 00:00:00: 86%|████████▌ | 59/69 [00:25<00:03, 3.01it/s]
Fetching GFS for 2022-01-01 00:00:00: 87%|████████▋ | 60/69 [00:25<00:02, 3.17it/s]
2024-06-25 14:01:14.935 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: v250 at 2022-01-01 00:00:00
Fetching GFS for 2022-01-01 00:00:00: 87%|████████▋ | 60/69 [00:25<00:02, 3.17it/s]
Fetching GFS for 2022-01-01 00:00:00: 88%|████████▊ | 61/69 [00:25<00:02, 3.33it/s]
2024-06-25 14:01:15.202 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: v200 at 2022-01-01 00:00:00
Fetching GFS for 2022-01-01 00:00:00: 88%|████████▊ | 61/69 [00:25<00:02, 3.33it/s]
Fetching GFS for 2022-01-01 00:00:00: 90%|████████▉ | 62/69 [00:25<00:02, 3.45it/s]
2024-06-25 14:01:15.466 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: v150 at 2022-01-01 00:00:00
Fetching GFS for 2022-01-01 00:00:00: 90%|████████▉ | 62/69 [00:25<00:02, 3.45it/s]
Fetching GFS for 2022-01-01 00:00:00: 91%|█████████▏| 63/69 [00:26<00:01, 3.29it/s]
2024-06-25 14:01:15.803 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: v100 at 2022-01-01 00:00:00
Fetching GFS for 2022-01-01 00:00:00: 91%|█████████▏| 63/69 [00:26<00:01, 3.29it/s]
Fetching GFS for 2022-01-01 00:00:00: 93%|█████████▎| 64/69 [00:26<00:01, 3.16it/s]
2024-06-25 14:01:16.149 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: v50 at 2022-01-01 00:00:00
Fetching GFS for 2022-01-01 00:00:00: 93%|█████████▎| 64/69 [00:26<00:01, 3.16it/s]
Fetching GFS for 2022-01-01 00:00:00: 94%|█████████▍| 65/69 [00:26<00:01, 3.33it/s]
2024-06-25 14:01:16.412 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: msl at 2022-01-01 00:00:00
Fetching GFS for 2022-01-01 00:00:00: 94%|█████████▍| 65/69 [00:26<00:01, 3.33it/s]
Fetching GFS for 2022-01-01 00:00:00: 96%|█████████▌| 66/69 [00:27<00:01, 2.96it/s]
2024-06-25 14:01:16.835 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: u10m at 2022-01-01 00:00:00
Fetching GFS for 2022-01-01 00:00:00: 96%|█████████▌| 66/69 [00:27<00:01, 2.96it/s]
Fetching GFS for 2022-01-01 00:00:00: 97%|█████████▋| 67/69 [00:27<00:00, 2.82it/s]
2024-06-25 14:01:17.229 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: v10m at 2022-01-01 00:00:00
Fetching GFS for 2022-01-01 00:00:00: 97%|█████████▋| 67/69 [00:27<00:00, 2.82it/s]
Fetching GFS for 2022-01-01 00:00:00: 99%|█████████▊| 68/69 [00:27<00:00, 3.00it/s]
2024-06-25 14:01:17.512 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: t2m at 2022-01-01 00:00:00
Fetching GFS for 2022-01-01 00:00:00: 99%|█████████▊| 68/69 [00:27<00:00, 3.00it/s]
Fetching GFS for 2022-01-01 00:00:00: 100%|██████████| 69/69 [00:28<00:00, 2.65it/s]
Fetching GFS for 2022-01-01 00:00:00: 100%|██████████| 69/69 [00:28<00:00, 2.42it/s]
2024-06-25 14:01:18.174 | SUCCESS | __main__:run_stats:131 - Fetched data from GFS
2024-06-25 14:01:18.192 | INFO | __main__:run_stats:151 - Inference starting!
Running inference: 0%| | 0/61 [00:00<?, ?it/s]
Running inference: 3%|▎ | 2/61 [00:01<00:38, 1.54it/s]
Running inference: 5%|▍ | 3/61 [00:02<00:44, 1.31it/s]
Running inference: 7%|▋ | 4/61 [00:03<00:46, 1.21it/s]
Running inference: 8%|▊ | 5/61 [00:04<00:48, 1.16it/s]
Running inference: 10%|▉ | 6/61 [00:04<00:48, 1.13it/s]
Running inference: 11%|█▏ | 7/61 [00:05<00:48, 1.12it/s]
Running inference: 13%|█▎ | 8/61 [00:06<00:47, 1.10it/s]
Running inference: 15%|█▍ | 9/61 [00:07<00:47, 1.10it/s]
Running inference: 16%|█▋ | 10/61 [00:08<00:46, 1.09it/s]
Running inference: 18%|█▊ | 11/61 [00:09<00:45, 1.09it/s]
Running inference: 20%|█▉ | 12/61 [00:10<00:45, 1.09it/s]
Running inference: 21%|██▏ | 13/61 [00:11<00:44, 1.08it/s]
Running inference: 23%|██▎ | 14/61 [00:12<00:43, 1.08it/s]
Running inference: 25%|██▍ | 15/61 [00:13<00:42, 1.08it/s]
Running inference: 26%|██▌ | 16/61 [00:14<00:41, 1.08it/s]
Running inference: 28%|██▊ | 17/61 [00:15<00:40, 1.08it/s]
Running inference: 30%|██▉ | 18/61 [00:16<00:39, 1.08it/s]
Running inference: 31%|███ | 19/61 [00:17<00:38, 1.08it/s]
Running inference: 33%|███▎ | 20/61 [00:17<00:37, 1.08it/s]
Running inference: 34%|███▍ | 21/61 [00:18<00:37, 1.08it/s]
Running inference: 36%|███▌ | 22/61 [00:19<00:36, 1.08it/s]
Running inference: 38%|███▊ | 23/61 [00:20<00:35, 1.08it/s]
Running inference: 39%|███▉ | 24/61 [00:21<00:34, 1.08it/s]
Running inference: 41%|████ | 25/61 [00:22<00:33, 1.08it/s]
Running inference: 43%|████▎ | 26/61 [00:23<00:32, 1.08it/s]
Running inference: 44%|████▍ | 27/61 [00:24<00:31, 1.08it/s]
Running inference: 46%|████▌ | 28/61 [00:25<00:30, 1.08it/s]
Running inference: 48%|████▊ | 29/61 [00:26<00:29, 1.08it/s]
Running inference: 49%|████▉ | 30/61 [00:27<00:28, 1.08it/s]
Running inference: 51%|█████ | 31/61 [00:28<00:27, 1.08it/s]
Running inference: 52%|█████▏ | 32/61 [00:29<00:26, 1.08it/s]
Running inference: 54%|█████▍ | 33/61 [00:30<00:25, 1.08it/s]
Running inference: 56%|█████▌ | 34/61 [00:30<00:25, 1.08it/s]
Running inference: 57%|█████▋ | 35/61 [00:31<00:24, 1.08it/s]
Running inference: 59%|█████▉ | 36/61 [00:32<00:23, 1.08it/s]
Running inference: 61%|██████ | 37/61 [00:33<00:22, 1.08it/s]
Running inference: 62%|██████▏ | 38/61 [00:34<00:21, 1.08it/s]
Running inference: 64%|██████▍ | 39/61 [00:35<00:20, 1.08it/s]
Running inference: 66%|██████▌ | 40/61 [00:36<00:19, 1.08it/s]
Running inference: 67%|██████▋ | 41/61 [00:37<00:18, 1.08it/s]
Running inference: 69%|██████▉ | 42/61 [00:38<00:17, 1.08it/s]
Running inference: 70%|███████ | 43/61 [00:39<00:16, 1.08it/s]
Running inference: 72%|███████▏ | 44/61 [00:40<00:15, 1.08it/s]
Running inference: 74%|███████▍ | 45/61 [00:41<00:14, 1.08it/s]
Running inference: 75%|███████▌ | 46/61 [00:42<00:13, 1.08it/s]
Running inference: 77%|███████▋ | 47/61 [00:43<00:12, 1.08it/s]
Running inference: 79%|███████▊ | 48/61 [00:43<00:12, 1.08it/s]
Running inference: 80%|████████ | 49/61 [00:44<00:11, 1.08it/s]
Running inference: 82%|████████▏ | 50/61 [00:45<00:10, 1.08it/s]
Running inference: 84%|████████▎ | 51/61 [00:46<00:09, 1.08it/s]
Running inference: 85%|████████▌ | 52/61 [00:47<00:08, 1.08it/s]
Running inference: 87%|████████▋ | 53/61 [00:48<00:07, 1.08it/s]
Running inference: 89%|████████▊ | 54/61 [00:49<00:06, 1.08it/s]
Running inference: 90%|█████████ | 55/61 [00:50<00:05, 1.08it/s]
Running inference: 92%|█████████▏| 56/61 [00:51<00:04, 1.08it/s]
Running inference: 93%|█████████▎| 57/61 [00:52<00:03, 1.08it/s]
Running inference: 95%|█████████▌| 58/61 [00:53<00:02, 1.08it/s]
Running inference: 97%|█████████▋| 59/61 [00:54<00:01, 1.08it/s]
Running inference: 98%|█████████▊| 60/61 [00:55<00:00, 1.08it/s]
Running inference: 100%|██████████| 61/61 [00:56<00:00, 1.08it/s]
Running inference: 100%|██████████| 61/61 [00:56<00:00, 1.09it/s]
2024-06-25 14:02:14.218 | SUCCESS | __main__:run_stats:160 - Inference complete
Post Processing#
The last step is to post process our results.
Notice that the NetCDF IO function has additional APIs to interact with the stored data.
import matplotlib.pyplot as plt
times = io["time"][:].flatten() + io["lead_time"][:].flatten()
fig = plt.figure(figsize=(12, 4))
ax = fig.add_subplot(1, 1, 1)
ax.plot(times, io["soi"][:].flatten())
ax.set_title("Southern Oscillation Index")
ax.grid("on")
plt.savefig("outputs/southern_oscillation_index_prediction_2022.png")
io.close()
Total running time of the script: (2 minutes 6.477 seconds)