Note
Go to the end to download the full example code.
Extending Data Sources#
Implementing a custom data source
This example will demonstrate how to extend Earth2Studio by implementing a custom data source to use in a built in workflow.
In this example you will learn:
API requirements of data soruces
Implementing a custom data soruce
Custom Data Source#
Earth2Studio defines the required APIs for data sources in
earth2studio.data.base.DataSource which requires just a call function.
For this example, we will consider extending an existing remote data source with
another atmospheric field we can calculate.
The earth2studio.data.ARCO data source provides the ERA5 dataset in a cloud
optimized format, however it only provides specific humidity. This is a problem for
models that may use relative humidity as an input. Based on ECMWF documentation we can
calculate the relative humidity based on temperature and geo-potential.
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 xarray as xr
from earth2studio.data import ARCO, GFS
from earth2studio.data.utils import prep_data_inputs
from earth2studio.utils.type import TimeArray, VariableArray
class CustomDataSource:
"""Custom ARCO datasource"""
relative_humidity_ids = [
"r50",
"r100",
"r150",
"r200",
"r250",
"r300",
"r400",
"r500",
"r600",
"r700",
"r850",
"r925",
"r1000",
]
def __init__(self, cache: bool = True, verbose: bool = True):
self.arco = ARCO(cache, verbose)
def __call__(
self,
time: datetime | list[datetime] | TimeArray,
variable: str | list[str] | VariableArray,
) -> xr.DataArray:
"""Function to get data.
Parameters
----------
time : datetime | list[datetime] | TimeArray
Timestamps to return data for (UTC).
variable : str | list[str] | VariableArray
String, list of strings or array of strings that refer to variables to
return. Must be in IFS lexicon.
Returns
-------
xr.DataArray
"""
time, variable = prep_data_inputs(time, variable)
# Replace relative humidity with respective temperature
# and specifc humidity fields
variable_expanded = []
for v in variable:
if v in self.relative_humidity_ids:
level = int(v[1:])
variable_expanded.extend([f"t{level}", f"q{level}"])
else:
variable_expanded.append(v)
variable_expanded = list(set(variable_expanded))
# Fetch from ARCO
da_exp = self.arco(time, variable_expanded)
# Calculate relative humidity when needed
arrays = []
for v in variable:
if v in self.relative_humidity_ids:
level = int(v[1:])
t = da_exp.sel(variable=f"t{level}").values
q = da_exp.sel(variable=f"q{level}").values
rh = self.calc_relative_humdity(t, q, 100 * level)
arrays.append(rh)
else:
arrays.append(da_exp.sel(variable=v).values)
da = xr.DataArray(
data=np.stack(arrays, axis=1),
dims=["time", "variable", "lat", "lon"],
coords=dict(
time=da_exp.coords["time"].values,
variable=np.array(variable),
lat=da_exp.coords["lat"].values,
lon=da_exp.coords["lon"].values,
),
)
return da
def calc_relative_humdity(
self, temperature: np.array, specific_humidity: np.array, pressure: float
) -> np.array:
"""Relative humidity calculation
Parameters
----------
temperature : np.array
Temperature field (K)
specific_humidity : np.array
Specific humidity field (g.kg-1)
pressure : float
Pressure (Pa)
Returns
-------
np.array
"""
epsilon = 0.621981
p = pressure
q = specific_humidity
t = temperature
e = (p * q * (1.0 / epsilon)) / (1 + q * (1.0 / (epsilon) - 1))
es_w = 611.21 * np.exp(17.502 * (t - 273.16) / (t - 32.19))
es_i = 611.21 * np.exp(22.587 * (t - 273.16) / (t + 0.7))
alpha = np.clip((t - 250.16) / (273.16 - 250.16), 0, 1.2) ** 2
es = alpha * es_w + (1 - alpha) * es_i
rh = 100 * e / es
return rh
__call__() API#
The call function is the main API of data source which return the Xarray data array with the requested data. For this custom data source we intercept relative humidity variables, replace them with temperature and specific humidity requests then calculate the relative humidity from these fields. Note that the ARCO data source is handling the remote complexity, we are just manipulating Numpy arrays
calc_relative_humdity()#
Based on the calculations ECMWF uses in their IFS numerical simulator which accounts for estimating the water vapor and ice present in the atmosphere.
Note
See reference, equation 7.98 onwards: https://www.ecmwf.int/en/elibrary/81370-ifs-documentation-cy48r1-part-iv-physical-processes
Verification#
Before plugging this into our workflow, let’s quickly verify our data source is consistent with when GFS provides for relative humidity.
ds = CustomDataSource()
da_custom = ds(time=datetime(2022, 1, 1, hour=0), variable=["r500"])
ds_gfs = GFS()
da_gfs = ds_gfs(time=datetime(2022, 1, 1, hour=0), variable=["r500"])
print(da_custom)
Fetching ARCO for 2022-01-01 00:00:00: 0%| | 0/2 [00:00<?, ?it/s]
2024-06-25 14:05:38.018 | DEBUG | earth2studio.data.arco:fetch_arco_dataarray:166 - Fetching ARCO zarr array for variable: t500 at 2022-01-01T00:00:00
Fetching ARCO for 2022-01-01 00:00:00: 0%| | 0/2 [00:00<?, ?it/s]
Fetching ARCO for 2022-01-01 00:00:00: 50%|█████ | 1/2 [00:03<00:03, 3.47s/it]
2024-06-25 14:05:41.488 | DEBUG | earth2studio.data.arco:fetch_arco_dataarray:166 - Fetching ARCO zarr array for variable: q500 at 2022-01-01T00:00:00
Fetching ARCO for 2022-01-01 00:00:00: 50%|█████ | 1/2 [00:03<00:03, 3.47s/it]
Fetching ARCO for 2022-01-01 00:00:00: 100%|██████████| 2/2 [00:07<00:00, 3.89s/it]
Fetching ARCO for 2022-01-01 00:00:00: 100%|██████████| 2/2 [00:07<00:00, 3.83s/it]
2024-06-25 14:05:45.705 | 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/1 [00:00<?, ?it/s]
2024-06-25 14:05:46.056 | DEBUG | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: r500 at 2022-01-01 00:00:00
Fetching GFS for 2022-01-01 00:00:00: 0%| | 0/1 [00:00<?, ?it/s]
Fetching GFS for 2022-01-01 00:00:00: 100%|██████████| 1/1 [00:00<00:00, 1.44it/s]
Fetching GFS for 2022-01-01 00:00:00: 100%|██████████| 1/1 [00:00<00:00, 1.44it/s]
<xarray.DataArray (time: 1, variable: 1, lat: 721, lon: 1440)>
array([[[[ 28.01413468, 28.01413468, 28.01413468, ..., 28.01413468,
28.01413468, 28.01413468],
[ 29.75579658, 29.75314258, 29.81191126, ..., 29.77167483,
29.76636404, 29.76376009],
[ 31.28908409, 31.28075462, 31.27237381, ..., 31.31699177,
31.30860053, 31.29746972],
...,
[ 97.36539026, 97.35685903, 97.39891442, ..., 97.34011022,
97.32321665, 97.31468906],
[ 97.50236648, 97.49380369, 97.54603948, ..., 97.5280597 ,
97.51949448, 97.51949448],
[102.58417544, 102.58417544, 102.58417544, ..., 102.58417544,
102.58417544, 102.58417544]]]])
Coordinates:
* time (time) datetime64[ns] 2022-01-01
* variable (variable) <U4 'r500'
* lat (lat) float64 90.0 89.75 89.5 89.25 ... -89.25 -89.5 -89.75 -90.0
* lon (lon) float64 0.0 0.25 0.5 0.75 1.0 ... 359.0 359.2 359.5 359.8
import cartopy.crs as ccrs
import matplotlib.pyplot as plt
fig, ax = plt.subplots(
1,
2,
figsize=(10, 3),
subplot_kw={"projection": ccrs.Mollweide()},
constrained_layout=True,
)
ax[0].imshow(
da_custom.sel(variable="r500")[0], transform=ccrs.PlateCarree(), vmin=0, vmax=100
)
ax[1].imshow(
da_gfs.sel(variable="r500")[0], transform=ccrs.PlateCarree(), vmin=0, vmax=100
)
ax[0].set_title("Custom ARCO")
ax[1].set_title("GFS")
plt.suptitle("r500", fontsize=24)
cbar = plt.cm.ScalarMappable()
cbar.set_array(da_custom.sel(variable="r500")[0])
cbar.set_clim(0, 100)
cbar = fig.colorbar(cbar, ax=ax[-1], orientation="vertical", shrink=0.8)
plt.savefig("outputs/custom_datasource_gfs_versus_custom.jpg")
Execute Workflow#
We will use this custom data source to run deterministic inference with a model that
requires relative humidity. earth2studio.models.px.FCN is one such model. Since
we are using ARCO, we can run inference for a time quite far back in time.
Let’s instantiate the components needed.
Prognostic Model: Use the built in FourCastNet Model
earth2studio.models.px.FCN.Datasource: Custom data source above
IO Backend: Save the outputs into a Zarr store
earth2studio.io.ZarrBackend.
from dotenv import load_dotenv
load_dotenv() # TODO: make common example prep function
import earth2studio.run as run
from earth2studio.io import ZarrBackend
from earth2studio.models.px import FCN
package = FCN.load_default_package()
model = FCN.load_model(package)
# Create the data source
data = CustomDataSource()
# Create the IO handler, store in memory
io = ZarrBackend()
nsteps = 4
io = run.deterministic(["1993-04-05"], nsteps, model, data, io)
print(io.root.tree())
2024-06-25 14:05:54.507 | INFO | earth2studio.run:deterministic:75 - Running simple workflow!
2024-06-25 14:05:54.507 | INFO | earth2studio.run:deterministic:82 - Inference device: cuda
Fetching ARCO for 1993-04-05 00:00:00: 0%| | 0/26 [00:00<?, ?it/s]
2024-06-25 14:05:54.668 | DEBUG | earth2studio.data.arco:fetch_arco_dataarray:166 - Fetching ARCO zarr array for variable: u10m at 1993-04-05T00:00:00
Fetching ARCO for 1993-04-05 00:00:00: 0%| | 0/26 [00:00<?, ?it/s]
Fetching ARCO for 1993-04-05 00:00:00: 4%|▍ | 1/26 [00:01<00:29, 1.17s/it]
2024-06-25 14:05:55.834 | DEBUG | earth2studio.data.arco:fetch_arco_dataarray:166 - Fetching ARCO zarr array for variable: v100m at 1993-04-05T00:00:00
Fetching ARCO for 1993-04-05 00:00:00: 4%|▍ | 1/26 [00:01<00:29, 1.17s/it]
Fetching ARCO for 1993-04-05 00:00:00: 8%|▊ | 2/26 [00:02<00:25, 1.07s/it]
2024-06-25 14:05:56.844 | DEBUG | earth2studio.data.arco:fetch_arco_dataarray:166 - Fetching ARCO zarr array for variable: q500 at 1993-04-05T00:00:00
Fetching ARCO for 1993-04-05 00:00:00: 8%|▊ | 2/26 [00:02<00:25, 1.07s/it]
Fetching ARCO for 1993-04-05 00:00:00: 12%|█▏ | 3/26 [00:05<00:47, 2.05s/it]
2024-06-25 14:06:00.047 | DEBUG | earth2studio.data.arco:fetch_arco_dataarray:166 - Fetching ARCO zarr array for variable: v250 at 1993-04-05T00:00:00
Fetching ARCO for 1993-04-05 00:00:00: 12%|█▏ | 3/26 [00:05<00:47, 2.05s/it]
Fetching ARCO for 1993-04-05 00:00:00: 15%|█▌ | 4/26 [00:10<01:10, 3.20s/it]
2024-06-25 14:06:05.024 | DEBUG | earth2studio.data.arco:fetch_arco_dataarray:166 - Fetching ARCO zarr array for variable: z1000 at 1993-04-05T00:00:00
Fetching ARCO for 1993-04-05 00:00:00: 15%|█▌ | 4/26 [00:10<01:10, 3.20s/it]
Fetching ARCO for 1993-04-05 00:00:00: 19%|█▉ | 5/26 [00:12<01:02, 2.98s/it]
2024-06-25 14:06:07.602 | DEBUG | earth2studio.data.arco:fetch_arco_dataarray:166 - Fetching ARCO zarr array for variable: z500 at 1993-04-05T00:00:00
Fetching ARCO for 1993-04-05 00:00:00: 19%|█▉ | 5/26 [00:12<01:02, 2.98s/it]
Fetching ARCO for 1993-04-05 00:00:00: 23%|██▎ | 6/26 [00:14<00:50, 2.53s/it]
2024-06-25 14:06:09.271 | DEBUG | earth2studio.data.arco:fetch_arco_dataarray:166 - Fetching ARCO zarr array for variable: z50 at 1993-04-05T00:00:00
Fetching ARCO for 1993-04-05 00:00:00: 23%|██▎ | 6/26 [00:14<00:50, 2.53s/it]
Fetching ARCO for 1993-04-05 00:00:00: 27%|██▋ | 7/26 [00:15<00:40, 2.14s/it]
2024-06-25 14:06:10.592 | DEBUG | earth2studio.data.arco:fetch_arco_dataarray:166 - Fetching ARCO zarr array for variable: t500 at 1993-04-05T00:00:00
Fetching ARCO for 1993-04-05 00:00:00: 27%|██▋ | 7/26 [00:15<00:40, 2.14s/it]
Fetching ARCO for 1993-04-05 00:00:00: 31%|███ | 8/26 [00:18<00:40, 2.26s/it]
2024-06-25 14:06:13.128 | DEBUG | earth2studio.data.arco:fetch_arco_dataarray:166 - Fetching ARCO zarr array for variable: t250 at 1993-04-05T00:00:00
Fetching ARCO for 1993-04-05 00:00:00: 31%|███ | 8/26 [00:18<00:40, 2.26s/it]
Fetching ARCO for 1993-04-05 00:00:00: 35%|███▍ | 9/26 [00:20<00:35, 2.10s/it]
2024-06-25 14:06:14.853 | DEBUG | earth2studio.data.arco:fetch_arco_dataarray:166 - Fetching ARCO zarr array for variable: u100m at 1993-04-05T00:00:00
Fetching ARCO for 1993-04-05 00:00:00: 35%|███▍ | 9/26 [00:20<00:35, 2.10s/it]
Fetching ARCO for 1993-04-05 00:00:00: 38%|███▊ | 10/26 [00:21<00:28, 1.77s/it]
2024-06-25 14:06:15.883 | DEBUG | earth2studio.data.arco:fetch_arco_dataarray:166 - Fetching ARCO zarr array for variable: z250 at 1993-04-05T00:00:00
Fetching ARCO for 1993-04-05 00:00:00: 38%|███▊ | 10/26 [00:21<00:28, 1.77s/it]
Fetching ARCO for 1993-04-05 00:00:00: 42%|████▏ | 11/26 [00:22<00:26, 1.74s/it]
2024-06-25 14:06:17.554 | DEBUG | earth2studio.data.arco:fetch_arco_dataarray:166 - Fetching ARCO zarr array for variable: u500 at 1993-04-05T00:00:00
Fetching ARCO for 1993-04-05 00:00:00: 42%|████▏ | 11/26 [00:22<00:26, 1.74s/it]
Fetching ARCO for 1993-04-05 00:00:00: 46%|████▌ | 12/26 [00:26<00:33, 2.42s/it]
2024-06-25 14:06:21.539 | DEBUG | earth2studio.data.arco:fetch_arco_dataarray:166 - Fetching ARCO zarr array for variable: msl at 1993-04-05T00:00:00
Fetching ARCO for 1993-04-05 00:00:00: 46%|████▌ | 12/26 [00:26<00:33, 2.42s/it]
Fetching ARCO for 1993-04-05 00:00:00: 50%|█████ | 13/26 [00:27<00:25, 1.97s/it]
2024-06-25 14:06:22.462 | DEBUG | earth2studio.data.arco:fetch_arco_dataarray:166 - Fetching ARCO zarr array for variable: q850 at 1993-04-05T00:00:00
Fetching ARCO for 1993-04-05 00:00:00: 50%|█████ | 13/26 [00:27<00:25, 1.97s/it]
Fetching ARCO for 1993-04-05 00:00:00: 54%|█████▍ | 14/26 [00:30<00:25, 2.14s/it]
2024-06-25 14:06:25.003 | DEBUG | earth2studio.data.arco:fetch_arco_dataarray:166 - Fetching ARCO zarr array for variable: t850 at 1993-04-05T00:00:00
Fetching ARCO for 1993-04-05 00:00:00: 54%|█████▍ | 14/26 [00:30<00:25, 2.14s/it]
Fetching ARCO for 1993-04-05 00:00:00: 58%|█████▊ | 15/26 [00:32<00:23, 2.17s/it]
2024-06-25 14:06:27.250 | DEBUG | earth2studio.data.arco:fetch_arco_dataarray:166 - Fetching ARCO zarr array for variable: sp at 1993-04-05T00:00:00
Fetching ARCO for 1993-04-05 00:00:00: 58%|█████▊ | 15/26 [00:32<00:23, 2.17s/it]
Fetching ARCO for 1993-04-05 00:00:00: 62%|██████▏ | 16/26 [00:33<00:18, 1.81s/it]
2024-06-25 14:06:28.214 | DEBUG | earth2studio.data.arco:fetch_arco_dataarray:166 - Fetching ARCO zarr array for variable: tcwv at 1993-04-05T00:00:00
Fetching ARCO for 1993-04-05 00:00:00: 62%|██████▏ | 16/26 [00:33<00:18, 1.81s/it]
Fetching ARCO for 1993-04-05 00:00:00: 65%|██████▌ | 17/26 [00:34<00:13, 1.53s/it]
2024-06-25 14:06:29.108 | DEBUG | earth2studio.data.arco:fetch_arco_dataarray:166 - Fetching ARCO zarr array for variable: v1000 at 1993-04-05T00:00:00
Fetching ARCO for 1993-04-05 00:00:00: 65%|██████▌ | 17/26 [00:34<00:13, 1.53s/it]
Fetching ARCO for 1993-04-05 00:00:00: 69%|██████▉ | 18/26 [00:37<00:14, 1.85s/it]
2024-06-25 14:06:31.704 | DEBUG | earth2studio.data.arco:fetch_arco_dataarray:166 - Fetching ARCO zarr array for variable: u1000 at 1993-04-05T00:00:00
Fetching ARCO for 1993-04-05 00:00:00: 69%|██████▉ | 18/26 [00:37<00:14, 1.85s/it]
Fetching ARCO for 1993-04-05 00:00:00: 73%|███████▎ | 19/26 [00:39<00:15, 2.15s/it]
2024-06-25 14:06:34.544 | DEBUG | earth2studio.data.arco:fetch_arco_dataarray:166 - Fetching ARCO zarr array for variable: u850 at 1993-04-05T00:00:00
Fetching ARCO for 1993-04-05 00:00:00: 73%|███████▎ | 19/26 [00:39<00:15, 2.15s/it]
Fetching ARCO for 1993-04-05 00:00:00: 77%|███████▋ | 20/26 [00:42<00:12, 2.16s/it]
2024-06-25 14:06:36.723 | DEBUG | earth2studio.data.arco:fetch_arco_dataarray:166 - Fetching ARCO zarr array for variable: v850 at 1993-04-05T00:00:00
Fetching ARCO for 1993-04-05 00:00:00: 77%|███████▋ | 20/26 [00:42<00:12, 2.16s/it]
Fetching ARCO for 1993-04-05 00:00:00: 81%|████████ | 21/26 [00:44<00:10, 2.17s/it]
2024-06-25 14:06:38.935 | DEBUG | earth2studio.data.arco:fetch_arco_dataarray:166 - Fetching ARCO zarr array for variable: z850 at 1993-04-05T00:00:00
Fetching ARCO for 1993-04-05 00:00:00: 81%|████████ | 21/26 [00:44<00:10, 2.17s/it]
Fetching ARCO for 1993-04-05 00:00:00: 85%|████████▍ | 22/26 [00:46<00:08, 2.16s/it]
2024-06-25 14:06:41.072 | DEBUG | earth2studio.data.arco:fetch_arco_dataarray:166 - Fetching ARCO zarr array for variable: t2m at 1993-04-05T00:00:00
Fetching ARCO for 1993-04-05 00:00:00: 85%|████████▍ | 22/26 [00:46<00:08, 2.16s/it]
Fetching ARCO for 1993-04-05 00:00:00: 88%|████████▊ | 23/26 [00:46<00:04, 1.64s/it]
2024-06-25 14:06:41.485 | DEBUG | earth2studio.data.arco:fetch_arco_dataarray:166 - Fetching ARCO zarr array for variable: v500 at 1993-04-05T00:00:00
Fetching ARCO for 1993-04-05 00:00:00: 88%|████████▊ | 23/26 [00:46<00:04, 1.64s/it]
Fetching ARCO for 1993-04-05 00:00:00: 92%|█████████▏| 24/26 [00:48<00:03, 1.70s/it]
2024-06-25 14:06:43.317 | DEBUG | earth2studio.data.arco:fetch_arco_dataarray:166 - Fetching ARCO zarr array for variable: u250 at 1993-04-05T00:00:00
Fetching ARCO for 1993-04-05 00:00:00: 92%|█████████▏| 24/26 [00:48<00:03, 1.70s/it]
Fetching ARCO for 1993-04-05 00:00:00: 96%|█████████▌| 25/26 [00:50<00:01, 1.84s/it]
2024-06-25 14:06:45.509 | DEBUG | earth2studio.data.arco:fetch_arco_dataarray:166 - Fetching ARCO zarr array for variable: v10m at 1993-04-05T00:00:00
Fetching ARCO for 1993-04-05 00:00:00: 96%|█████████▌| 25/26 [00:50<00:01, 1.84s/it]
Fetching ARCO for 1993-04-05 00:00:00: 100%|██████████| 26/26 [00:51<00:00, 1.56s/it]
Fetching ARCO for 1993-04-05 00:00:00: 100%|██████████| 26/26 [00:51<00:00, 1.99s/it]
2024-06-25 14:06:46.541 | SUCCESS | earth2studio.run:deterministic:95 - Fetched data from CustomDataSource
2024-06-25 14:06:46.567 | INFO | earth2studio.run:deterministic:125 - Inference starting!
Running inference: 0%| | 0/5 [00:00<?, ?it/s]
Running inference: 20%|██ | 1/5 [00:00<00:02, 1.60it/s]
Running inference: 40%|████ | 2/5 [00:01<00:02, 1.45it/s]
Running inference: 60%|██████ | 3/5 [00:02<00:01, 1.36it/s]
Running inference: 80%|████████ | 4/5 [00:02<00:00, 1.32it/s]
Running inference: 100%|██████████| 5/5 [00:03<00:00, 1.28it/s]
Running inference: 100%|██████████| 5/5 [00:03<00:00, 1.33it/s]
2024-06-25 14:06:50.335 | SUCCESS | earth2studio.run:deterministic:135 - Inference complete
/
├── lat (720,) float64
├── lead_time (5,) timedelta64[h]
├── lon (1440,) float64
├── msl (1, 5, 720, 1440) float32
├── r500 (1, 5, 720, 1440) float32
├── r850 (1, 5, 720, 1440) float32
├── sp (1, 5, 720, 1440) float32
├── t250 (1, 5, 720, 1440) float32
├── t2m (1, 5, 720, 1440) float32
├── t500 (1, 5, 720, 1440) float32
├── t850 (1, 5, 720, 1440) float32
├── tcwv (1, 5, 720, 1440) float32
├── time (1,) datetime64[ns]
├── u1000 (1, 5, 720, 1440) float32
├── u100m (1, 5, 720, 1440) float32
├── u10m (1, 5, 720, 1440) float32
├── u250 (1, 5, 720, 1440) float32
├── u500 (1, 5, 720, 1440) float32
├── u850 (1, 5, 720, 1440) float32
├── v1000 (1, 5, 720, 1440) float32
├── v100m (1, 5, 720, 1440) float32
├── v10m (1, 5, 720, 1440) float32
├── v250 (1, 5, 720, 1440) float32
├── v500 (1, 5, 720, 1440) float32
├── v850 (1, 5, 720, 1440) float32
├── z1000 (1, 5, 720, 1440) float32
├── z250 (1, 5, 720, 1440) float32
├── z50 (1, 5, 720, 1440) float32
├── z500 (1, 5, 720, 1440) float32
└── z850 (1, 5, 720, 1440) float32
Post Processing#
To confirm that our model is working as expected, we will plot the total column water vapor field for a few time-steps.
forecast = "1993-04-05"
variable = "tcwv"
plt.close("all")
# Create a figure and axes with the specified projection
fig, ax = plt.subplots(2, 2, figsize=(6, 4))
# Plot tcwv every 6 hours
ax[0, 0].imshow(io[variable][0, 0], vmin=0, vmax=80, cmap="magma")
ax[0, 1].imshow(io[variable][0, 1], vmin=0, vmax=80, cmap="magma")
ax[1, 0].imshow(io[variable][0, 2], vmin=0, vmax=80, cmap="magma")
ax[1, 1].imshow(io[variable][0, 3], vmin=0, vmax=80, cmap="magma")
# Set title
plt.suptitle(f"{variable} - {forecast}")
times = io["lead_time"].astype("timedelta64[h]").astype(int)
ax[0, 0].set_title(f"Lead time: {times[0]}hrs")
ax[0, 1].set_title(f"Lead time: {times[1]}hrs")
ax[1, 0].set_title(f"Lead time: {times[2]}hrs")
ax[1, 1].set_title(f"Lead time: {times[3]}hrs")
plt.savefig("outputs/custom_datasource_prediction.jpg", bbox_inches="tight")
Total running time of the script: (1 minutes 14.206 seconds)