.. DO NOT EDIT. .. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. .. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: .. "examples/10_stormcast_ensemble_example.py" .. LINE NUMBERS ARE GIVEN BELOW. .. only:: html .. note:: :class: sphx-glr-download-link-note :ref:`Go to the end ` to download the full example code. .. rst-class:: sphx-glr-example-title .. _sphx_glr_examples_10_stormcast_ensemble_example.py: Running StormCast Ensemble Inference =============================== Ensemble StormCast inference workflow. This example will demonstrate how to run a simple inference workflow to generate a ensemble forecast using StormCast. For details about the stormcast model, see - https://arxiv.org/abs/2408.10958 .. GENERATED FROM PYTHON SOURCE LINES 32-37 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 ensemble workflow: :py:meth:`earth2studio.run.ensemble`. .. GENERATED FROM PYTHON SOURCE LINES 39-43 .. literalinclude:: ../../earth2studio/run.py :language: python :start-after: # sphinx - ensemble start :end-before: # sphinx - ensemble end .. GENERATED FROM PYTHON SOURCE LINES 45-56 Thus, we need the following: - Prognostic Model: Use the built in StormCast Model :py:class:`earth2studio.models.px.StormCast`. - perturbation_method: Use the Zero Method :py:class:`earth2studio.perturbation.Zero`. We will not perturb the initial data because StormCast has stochastic generation of ensemble members. - Datasource: Pull data from the HRRR data api :py:class:`earth2studio.data.HRRR`. - IO Backend: Let's save the outputs into a Zarr store :py:class:`earth2studio.io.ZarrBackend`. StormCast also requires a conditioning data source. We use a forecast data source here, ARCO :py:class:`earth2studio.data.ARCO`, but a forecast data source such as GFS_FX could also be used with appropriate time stamps. .. GENERATED FROM PYTHON SOURCE LINES 58-94 .. code-block:: Python import numpy as np from loguru import logger from tqdm import tqdm logger.remove() logger.add(lambda msg: tqdm.write(msg, end=""), colorize=True) 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 ARCO, HRRR from earth2studio.io import ZarrBackend from earth2studio.models.px import StormCast from earth2studio.perturbation import Zero # Load the default model package which downloads the check point from NGC package = StormCast.load_default_package() model = StormCast.load_model(package) # Instantiate the (Zero) perturbation method z = Zero() # Create the data source data = HRRR() # Create and set the conditioning data source conditioning_data_source = ARCO() model.conditioning_data_source = conditioning_data_source # Create the IO handler, store in memory io = ZarrBackend() .. rst-class:: sphx-glr-script-out .. code-block:: none 2025-03-27 08:03:55.226 | WARNING | earth2studio.data.arco:__init__:113 - Using Zarr 2.0 method for ARCO, this can be extremely slow with caching! .. GENERATED FROM PYTHON SOURCE LINES 95-103 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 4 hours .. GENERATED FROM PYTHON SOURCE LINES 105-126 .. code-block:: Python import earth2studio.run as run nsteps = 4 nensemble = 8 batch_size = 2 date = "2022-11-04T21:00:00" io = run.ensemble( [date], nsteps, nensemble, model, data, io, z, batch_size=batch_size, output_coords={"variable": np.array(["t2m", "refc"])}, ) print(io.root.tree()) .. rst-class:: sphx-glr-script-out .. code-block:: none 2025-03-27 08:04:07.612 | INFO | earth2studio.run:ensemble:315 - Running ensemble inference! 2025-03-27 08:04:07.612 | INFO | earth2studio.run:ensemble:323 - Inference device: cuda Fetching HRRR data: 0%| | 0/99 [00:00 0, x[0, 0, step], np.nan), f"{forecast} - Lead time: {step}hrs - Member: {0}", cmap, vmin=0, vmax=60, ) plot_( ax2, np.where(x[1, 0, step] > 0, x[1, 0, step], np.nan), f"{forecast} - Lead time: {step}hrs - Member: {1}", cmap, vmin=0, vmax=60, ) plot_( ax3, np.where(x[:, 0, step].mean(axis=0) > 0, x[:, 0, step].std(axis=0), np.nan), f"{forecast} - Lead time: {step}hrs - Std", cmap, vmin=0, vmax=60, ) plt.savefig(f"outputs/10_{date}_{variable}_{step}_ensemble.jpg") .. image-sg:: /examples/images/sphx_glr_10_stormcast_ensemble_example_001.png :alt: 2022-11-04T21:00:00 - Lead time: 4hrs - Member: 0, 2022-11-04T21:00:00 - Lead time: 4hrs - Member: 1, 2022-11-04T21:00:00 - Lead time: 4hrs - Std :srcset: /examples/images/sphx_glr_10_stormcast_ensemble_example_001.png, /examples/images/sphx_glr_10_stormcast_ensemble_example_001_2_00x.png 2.00x :class: sphx-glr-single-img .. rst-class:: sphx-glr-timing **Total running time of the script:** (6 minutes 49.648 seconds) .. _sphx_glr_download_examples_10_stormcast_ensemble_example.py: .. only:: html .. container:: sphx-glr-footer sphx-glr-footer-example .. container:: sphx-glr-download sphx-glr-download-jupyter :download:`Download Jupyter notebook: 10_stormcast_ensemble_example.ipynb <10_stormcast_ensemble_example.ipynb>` .. container:: sphx-glr-download sphx-glr-download-python :download:`Download Python source code: 10_stormcast_ensemble_example.py <10_stormcast_ensemble_example.py>` .. container:: sphx-glr-download sphx-glr-download-zip :download:`Download zipped: 10_stormcast_ensemble_example.zip <10_stormcast_ensemble_example.zip>` .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_