.. 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 <sphx_glr_download_examples_10_stormcast_ensemble_example.py>`
        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()


.. 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-127

.. 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())


.. GENERATED FROM PYTHON SOURCE LINES 128-135

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) 4 hours into the forecast.

Notice that the Zarr IO function has additional APIs to interact with the stored data.

.. GENERATED FROM PYTHON SOURCE LINES 137-253

.. code-block:: Python

    import cartopy
    import cartopy.crs as ccrs
    import matplotlib.pyplot as plt

    forecast = f"{date}"
    step = 4  # 4 hours, since lead_time = 1 hr

    plt.close("all")


    def plot_(axi, data, title, cmap, vmin=None, vmax=None):
        """Convenience function for plotting pcolormesh."""
        # Plot the field using pcolormesh
        im = axi.pcolormesh(
            io["lon"][:],
            io["lat"][:],
            data,
            transform=ccrs.PlateCarree(),
            cmap=cmap,
            vmin=vmin,
            vmax=vmax,
        )
        plt.colorbar(im, ax=axi, shrink=0.6, pad=0.04)
        # Set title
        axi.set_title(title)

        # Add coastlines and gridlines
        axi.coastlines()
        axi.gridlines()

        # Set state lines
        axi.add_feature(
            cartopy.feature.STATES.with_scale("50m"),
            linewidth=0.5,
            edgecolor="black",
            zorder=2,
        )


    # Create a correct Lambert Conformal projection
    projection = ccrs.LambertConformal(
        central_longitude=262.5,
        central_latitude=38.5,
        standard_parallels=(38.5, 38.5),
        globe=ccrs.Globe(semimajor_axis=6371229, semiminor_axis=6371229),
    )

    # Plot 2-meter temperature
    variable = "t2m"
    cmap = "Spectral_r"
    x = io[variable]

    plt.close("all")
    fig, (ax1, ax2, ax3) = plt.subplots(
        nrows=1, ncols=3, subplot_kw={"projection": projection}, figsize=(20, 6)
    )
    plot_(
        ax1,
        x[0, 0, step],
        f"{forecast} - Lead time: {step}hrs - Member: {0}",
        cmap,
    )

    plot_(
        ax2,
        io[variable][1, 0, step],
        f"{forecast} - Lead time: {step}hrs - Member: {1}",
        cmap,
    )

    plot_(
        ax3,
        x[:, 0, step].std(axis=0),
        f"{forecast} - Lead time: {step}hrs - Std",
        cmap,
    )

    plt.savefig(f"outputs/10_{date}_{variable}_{step}_ensemble.jpg")

    # Plot refc
    variable = "refc"
    cmap = "gist_ncar"
    x = io[variable]

    plt.close("all")
    fig, (ax1, ax2, ax3) = plt.subplots(
        nrows=1, ncols=3, subplot_kw={"projection": projection}, figsize=(20, 6)
    )
    plot_(
        ax1,
        np.where(x[0, 0, step] > 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")

**Estimated memory usage:**  0 MB


.. _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 <https://sphinx-gallery.github.io>`_