Note
Go to the end to download the full example code.
Running Deterministic Inference#
Basic deterministic inference workflow.
This example will demonstrate how to run a simple inference workflow to generate a basic determinstic forecast using one of the built in models of Earth-2 Inference Studio.
In this example you will learn:
How to instantiate a built in prognostic model
Creating a data source and IO object
Running a simple built in workflow
Post-processing results
Set Up#
All workflows inside Earth2Studio require constructed components to be
handed to them. In this example, let’s take a look at the most basic:
earth2studio.run.deterministic().
def deterministic(
time: list[str] | list[datetime] | list[np.datetime64],
nsteps: int,
prognostic: PrognosticModel,
data: DataSource,
io: IOBackend,
output_coords: CoordSystem = OrderedDict({}),
device: torch.device | None = None,
) -> IOBackend:
"""Built in deterministic workflow.
This workflow creates a determinstic inference pipeline to produce a forecast
prediction using a prognostic model.
Parameters
----------
time : list[str] | list[datetime] | list[np.datetime64]
List of string, datetimes or np.datetime64
nsteps : int
Number of forecast steps
prognostic : PrognosticModel
Prognostic model
data : DataSource
Data source
io : IOBackend
IO object
output_coords: CoordSystem, optional
IO output coordinate system override, by default OrderedDict({})
device : torch.device, optional
Device to run inference on, by default None
Returns
-------
IOBackend
Output IO object
"""
Thus, we need the following:
Prognostic Model: Use the built in FourCastNet Model
earth2studio.models.px.FCN.Datasource: Pull data from the GFS data api
earth2studio.data.GFS.IO Backend: Let’s save the outputs into a Zarr store
earth2studio.io.ZarrBackend.
import os
os.makedirs("outputs", exist_ok=True)
from dotenv import load_dotenv
load_dotenv() # TODO: make common example prep function
from earth2studio.data import GFS
from earth2studio.io import ZarrBackend
from earth2studio.models.px import FCN
# Load the default model package which downloads the check point from NGC
package = FCN.load_default_package()
model = FCN.load_model(package)
# Create the data source
data = GFS()
# Create the IO handler, store in memory
io = ZarrBackend()
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.
For the forecast we will predict for two days (these will get executed as a batch) for 20 forecast steps which is 5 days.
import earth2studio.run as run
nsteps = 20
io = run.deterministic(["2024-01-01"], nsteps, model, data, io)
print(io.root.tree())
2024-06-25 13:54:50.151 | INFO | earth2studio.run:deterministic:75 - Running simple workflow!
2024-06-25 13:54:50.222 | INFO | earth2studio.run:deterministic:82 - Inference device: cuda
2024-06-25 13:54:50.387 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:149 - Fetching GFS index file: 2024-01-01 00:00:00
Fetching GFS for 2024-01-01 00:00:00: 0%| | 0/26 [00:00<?, ?it/s]
2024-06-25 13:54:50.816 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: u10m at 2024-01-01 00:00:00
Fetching GFS for 2024-01-01 00:00:00: 0%| | 0/26 [00:00<?, ?it/s]
Fetching GFS for 2024-01-01 00:00:00: 4%|▍ | 1/26 [00:04<01:41, 4.06s/it]
2024-06-25 13:54:54.876 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: v10m at 2024-01-01 00:00:00
Fetching GFS for 2024-01-01 00:00:00: 4%|▍ | 1/26 [00:04<01:41, 4.06s/it]
Fetching GFS for 2024-01-01 00:00:00: 8%|▊ | 2/26 [00:04<00:53, 2.21s/it]
2024-06-25 13:54:55.795 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: t2m at 2024-01-01 00:00:00
Fetching GFS for 2024-01-01 00:00:00: 8%|▊ | 2/26 [00:04<00:53, 2.21s/it]
Fetching GFS for 2024-01-01 00:00:00: 12%|█▏ | 3/26 [00:05<00:35, 1.53s/it]
2024-06-25 13:54:56.515 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: sp at 2024-01-01 00:00:00
Fetching GFS for 2024-01-01 00:00:00: 12%|█▏ | 3/26 [00:05<00:35, 1.53s/it]
Fetching GFS for 2024-01-01 00:00:00: 15%|█▌ | 4/26 [00:06<00:25, 1.15s/it]
2024-06-25 13:54:57.078 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: msl at 2024-01-01 00:00:00
Fetching GFS for 2024-01-01 00:00:00: 15%|█▌ | 4/26 [00:06<00:25, 1.15s/it]
Fetching GFS for 2024-01-01 00:00:00: 19%|█▉ | 5/26 [00:06<00:19, 1.10it/s]
2024-06-25 13:54:57.569 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: t850 at 2024-01-01 00:00:00
Fetching GFS for 2024-01-01 00:00:00: 19%|█▉ | 5/26 [00:06<00:19, 1.10it/s]
Fetching GFS for 2024-01-01 00:00:00: 23%|██▎ | 6/26 [00:07<00:15, 1.32it/s]
2024-06-25 13:54:58.032 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: u1000 at 2024-01-01 00:00:00
Fetching GFS for 2024-01-01 00:00:00: 23%|██▎ | 6/26 [00:07<00:15, 1.32it/s]
Fetching GFS for 2024-01-01 00:00:00: 27%|██▋ | 7/26 [00:07<00:12, 1.53it/s]
2024-06-25 13:54:58.475 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: v1000 at 2024-01-01 00:00:00
Fetching GFS for 2024-01-01 00:00:00: 27%|██▋ | 7/26 [00:07<00:12, 1.53it/s]
Fetching GFS for 2024-01-01 00:00:00: 31%|███ | 8/26 [00:08<00:10, 1.78it/s]
2024-06-25 13:54:58.832 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: z1000 at 2024-01-01 00:00:00
Fetching GFS for 2024-01-01 00:00:00: 31%|███ | 8/26 [00:08<00:10, 1.78it/s]
Fetching GFS for 2024-01-01 00:00:00: 35%|███▍ | 9/26 [00:08<00:08, 2.07it/s]
2024-06-25 13:54:59.148 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: u850 at 2024-01-01 00:00:00
Fetching GFS for 2024-01-01 00:00:00: 35%|███▍ | 9/26 [00:08<00:08, 2.07it/s]
Fetching GFS for 2024-01-01 00:00:00: 38%|███▊ | 10/26 [00:08<00:07, 2.08it/s]
2024-06-25 13:54:59.620 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: v850 at 2024-01-01 00:00:00
Fetching GFS for 2024-01-01 00:00:00: 38%|███▊ | 10/26 [00:08<00:07, 2.08it/s]
Fetching GFS for 2024-01-01 00:00:00: 42%|████▏ | 11/26 [00:09<00:07, 1.93it/s]
2024-06-25 13:55:00.227 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: z850 at 2024-01-01 00:00:00
Fetching GFS for 2024-01-01 00:00:00: 42%|████▏ | 11/26 [00:09<00:07, 1.93it/s]
Fetching GFS for 2024-01-01 00:00:00: 46%|████▌ | 12/26 [00:09<00:06, 2.16it/s]
2024-06-25 13:55:00.564 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: u500 at 2024-01-01 00:00:00
Fetching GFS for 2024-01-01 00:00:00: 46%|████▌ | 12/26 [00:09<00:06, 2.16it/s]
Fetching GFS for 2024-01-01 00:00:00: 50%|█████ | 13/26 [00:10<00:05, 2.18it/s]
2024-06-25 13:55:01.012 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: v500 at 2024-01-01 00:00:00
Fetching GFS for 2024-01-01 00:00:00: 50%|█████ | 13/26 [00:10<00:05, 2.18it/s]
Fetching GFS for 2024-01-01 00:00:00: 54%|█████▍ | 14/26 [00:10<00:05, 2.35it/s]
2024-06-25 13:55:01.361 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: z500 at 2024-01-01 00:00:00
Fetching GFS for 2024-01-01 00:00:00: 54%|█████▍ | 14/26 [00:10<00:05, 2.35it/s]
Fetching GFS for 2024-01-01 00:00:00: 58%|█████▊ | 15/26 [00:10<00:04, 2.41it/s]
2024-06-25 13:55:01.753 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: t500 at 2024-01-01 00:00:00
Fetching GFS for 2024-01-01 00:00:00: 58%|█████▊ | 15/26 [00:10<00:04, 2.41it/s]
Fetching GFS for 2024-01-01 00:00:00: 62%|██████▏ | 16/26 [00:11<00:03, 2.56it/s]
2024-06-25 13:55:02.088 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: z50 at 2024-01-01 00:00:00
Fetching GFS for 2024-01-01 00:00:00: 62%|██████▏ | 16/26 [00:11<00:03, 2.56it/s]
Fetching GFS for 2024-01-01 00:00:00: 65%|██████▌ | 17/26 [00:11<00:03, 2.59it/s]
2024-06-25 13:55:02.463 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: r500 at 2024-01-01 00:00:00
Fetching GFS for 2024-01-01 00:00:00: 65%|██████▌ | 17/26 [00:11<00:03, 2.59it/s]
Fetching GFS for 2024-01-01 00:00:00: 69%|██████▉ | 18/26 [00:11<00:03, 2.66it/s]
2024-06-25 13:55:02.813 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: r850 at 2024-01-01 00:00:00
Fetching GFS for 2024-01-01 00:00:00: 69%|██████▉ | 18/26 [00:11<00:03, 2.66it/s]
Fetching GFS for 2024-01-01 00:00:00: 73%|███████▎ | 19/26 [00:12<00:02, 2.80it/s]
2024-06-25 13:55:03.129 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: tcwv at 2024-01-01 00:00:00
Fetching GFS for 2024-01-01 00:00:00: 73%|███████▎ | 19/26 [00:12<00:02, 2.80it/s]
Fetching GFS for 2024-01-01 00:00:00: 77%|███████▋ | 20/26 [00:12<00:02, 2.27it/s]
2024-06-25 13:55:03.760 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: u100m at 2024-01-01 00:00:00
Fetching GFS for 2024-01-01 00:00:00: 77%|███████▋ | 20/26 [00:12<00:02, 2.27it/s]
Fetching GFS for 2024-01-01 00:00:00: 81%|████████ | 21/26 [00:13<00:02, 2.31it/s]
2024-06-25 13:55:04.179 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: v100m at 2024-01-01 00:00:00
Fetching GFS for 2024-01-01 00:00:00: 81%|████████ | 21/26 [00:13<00:02, 2.31it/s]
Fetching GFS for 2024-01-01 00:00:00: 85%|████████▍ | 22/26 [00:13<00:01, 2.48it/s]
2024-06-25 13:55:04.513 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: u250 at 2024-01-01 00:00:00
Fetching GFS for 2024-01-01 00:00:00: 85%|████████▍ | 22/26 [00:13<00:01, 2.48it/s]
Fetching GFS for 2024-01-01 00:00:00: 88%|████████▊ | 23/26 [00:14<00:01, 2.40it/s]
2024-06-25 13:55:04.959 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: v250 at 2024-01-01 00:00:00
Fetching GFS for 2024-01-01 00:00:00: 88%|████████▊ | 23/26 [00:14<00:01, 2.40it/s]
Fetching GFS for 2024-01-01 00:00:00: 92%|█████████▏| 24/26 [00:14<00:00, 2.54it/s]
2024-06-25 13:55:05.298 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: z250 at 2024-01-01 00:00:00
Fetching GFS for 2024-01-01 00:00:00: 92%|█████████▏| 24/26 [00:14<00:00, 2.54it/s]
Fetching GFS for 2024-01-01 00:00:00: 96%|█████████▌| 25/26 [00:14<00:00, 2.57it/s]
2024-06-25 13:55:05.680 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: t250 at 2024-01-01 00:00:00
Fetching GFS for 2024-01-01 00:00:00: 96%|█████████▌| 25/26 [00:14<00:00, 2.57it/s]
Fetching GFS for 2024-01-01 00:00:00: 100%|██████████| 26/26 [00:15<00:00, 2.64it/s]
Fetching GFS for 2024-01-01 00:00:00: 100%|██████████| 26/26 [00:15<00:00, 1.71it/s]
2024-06-25 13:55:06.106 | SUCCESS | earth2studio.run:deterministic:95 - Fetched data from GFS
2024-06-25 13:55:06.371 | INFO | earth2studio.run:deterministic:125 - Inference starting!
Running inference: 0%| | 0/21 [00:00<?, ?it/s]
Running inference: 5%|▍ | 1/21 [00:01<00:21, 1.08s/it]
Running inference: 10%|▉ | 2/21 [00:02<00:28, 1.50s/it]
Running inference: 14%|█▍ | 3/21 [00:04<00:27, 1.51s/it]
Running inference: 19%|█▉ | 4/21 [00:05<00:26, 1.54s/it]
Running inference: 24%|██▍ | 5/21 [00:07<00:24, 1.54s/it]
Running inference: 29%|██▊ | 6/21 [00:09<00:23, 1.55s/it]
Running inference: 33%|███▎ | 7/21 [00:10<00:21, 1.57s/it]
Running inference: 38%|███▊ | 8/21 [00:12<00:20, 1.59s/it]
Running inference: 43%|████▎ | 9/21 [00:14<00:19, 1.62s/it]
Running inference: 48%|████▊ | 10/21 [00:15<00:17, 1.64s/it]
Running inference: 52%|█████▏ | 11/21 [00:17<00:16, 1.66s/it]
Running inference: 57%|█████▋ | 12/21 [00:19<00:15, 1.69s/it]
Running inference: 62%|██████▏ | 13/21 [00:20<00:13, 1.72s/it]
Running inference: 67%|██████▋ | 14/21 [00:22<00:12, 1.75s/it]
Running inference: 71%|███████▏ | 15/21 [00:24<00:10, 1.78s/it]
Running inference: 76%|███████▌ | 16/21 [00:26<00:08, 1.75s/it]
Running inference: 81%|████████ | 17/21 [00:28<00:06, 1.73s/it]
Running inference: 86%|████████▌ | 18/21 [00:29<00:05, 1.72s/it]
Running inference: 90%|█████████ | 19/21 [00:31<00:03, 1.73s/it]
Running inference: 95%|█████████▌| 20/21 [00:33<00:01, 1.73s/it]
Running inference: 100%|██████████| 21/21 [00:35<00:00, 1.81s/it]
Running inference: 100%|██████████| 21/21 [00:35<00:00, 1.68s/it]
2024-06-25 13:55:41.557 | SUCCESS | earth2studio.run:deterministic:135 - Inference complete
/
├── lat (720,) float64
├── lead_time (21,) timedelta64[h]
├── lon (1440,) float64
├── msl (1, 21, 720, 1440) float32
├── r500 (1, 21, 720, 1440) float32
├── r850 (1, 21, 720, 1440) float32
├── sp (1, 21, 720, 1440) float32
├── t250 (1, 21, 720, 1440) float32
├── t2m (1, 21, 720, 1440) float32
├── t500 (1, 21, 720, 1440) float32
├── t850 (1, 21, 720, 1440) float32
├── tcwv (1, 21, 720, 1440) float32
├── time (1,) datetime64[ns]
├── u1000 (1, 21, 720, 1440) float32
├── u100m (1, 21, 720, 1440) float32
├── u10m (1, 21, 720, 1440) float32
├── u250 (1, 21, 720, 1440) float32
├── u500 (1, 21, 720, 1440) float32
├── u850 (1, 21, 720, 1440) float32
├── v1000 (1, 21, 720, 1440) float32
├── v100m (1, 21, 720, 1440) float32
├── v10m (1, 21, 720, 1440) float32
├── v250 (1, 21, 720, 1440) float32
├── v500 (1, 21, 720, 1440) float32
├── v850 (1, 21, 720, 1440) float32
├── z1000 (1, 21, 720, 1440) float32
├── z250 (1, 21, 720, 1440) float32
├── z50 (1, 21, 720, 1440) float32
├── z500 (1, 21, 720, 1440) float32
└── z850 (1, 21, 720, 1440) float32
Post Processing#
The last step is to post process our results. Cartopy is a great library for plotting fields on projections of a sphere. Here we will just plot the temperature at 2 meters (t2m) 1 day into the forecast.
Notice that the Zarr IO function has additional APIs to interact with the stored data.
import cartopy.crs as ccrs
import matplotlib.pyplot as plt
forecast = "2024-01-01"
variable = "t2m"
step = 4 # lead time = 24 hrs
plt.close("all")
# Create a Robinson projection
projection = ccrs.Robinson()
# Create a figure and axes with the specified projection
fig, ax = plt.subplots(subplot_kw={"projection": projection}, figsize=(10, 6))
# Plot the field using pcolormesh
im = ax.pcolormesh(
io["lon"][:],
io["lat"][:],
io[variable][0, step],
transform=ccrs.PlateCarree(),
cmap="Spectral_r",
)
# Set title
ax.set_title(f"{forecast} - Lead time: {6*step}hrs")
# Add coastlines and gridlines
ax.coastlines()
ax.gridlines()
plt.savefig("outputs/01_t2m_prediction.jpg")
/usr/local/lib/python3.10/dist-packages/cartopy/io/__init__.py:241: DownloadWarning: Downloading: https://naturalearth.s3.amazonaws.com/110m_physical/ne_110m_coastline.zip
warnings.warn(f'Downloading: {url}', DownloadWarning)
Total running time of the script: (1 minutes 10.413 seconds)