.. DO NOT EDIT.
.. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY.
.. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE:
.. "examples/03_ensemble_workflow.py"
.. LINE NUMBERS ARE GIVEN BELOW.

.. only:: html

    .. note::
        :class: sphx-glr-download-link-note

        :ref:`Go to the end <sphx_glr_download_examples_03_ensemble_workflow.py>`
        to download the full example code.

.. rst-class:: sphx-glr-example-title

.. _sphx_glr_examples_03_ensemble_workflow.py:


Running Ensemble Inference
==========================

Simple ensemble inference workflow.

This example will demonstrate how to run a simple inference workflow to generate a
ensemble 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
- Select a perturbation method
- Running a simple built in workflow for ensembling
- Post-processing results

.. GENERATED FROM PYTHON SOURCE LINES 38-43

Set Up
------
All workflows inside Earth2Studio require constructed components to be
handed to them. In this example, we will use the built in ensemble workflow
:py:meth:`earth2studio.run.ensemble`.

.. GENERATED FROM PYTHON SOURCE LINES 45-49

.. literalinclude:: ../../earth2studio/run.py
   :language: python
   :start-after: # sphinx - ensemble start
   :end-before: # sphinx - ensemble end

.. GENERATED FROM PYTHON SOURCE LINES 51-57

We need the following:

- Prognostic Model: Use the built in FourCastNet model :py:class:`earth2studio.models.px.FCN`.
- perturbation_method: Use the Spherical Gaussian Method :py:class:`earth2studio.perturbation.SphericalGaussian`.
- Datasource: Pull data from the GFS data api :py:class:`earth2studio.data.GFS`.
- IO Backend: Save the outputs into a Zarr store :py:class:`earth2studio.io.ZarrBackend`.

.. GENERATED FROM PYTHON SOURCE LINES 59-92

.. code-block:: Python

    import os

    os.makedirs("outputs", exist_ok=True)
    from dotenv import load_dotenv

    load_dotenv()  # TODO: make common example prep function

    import numpy as np

    from earth2studio.data import GFS
    from earth2studio.io import ZarrBackend
    from earth2studio.models.px import FCN
    from earth2studio.perturbation import SphericalGaussian
    from earth2studio.run import ensemble

    # Load the default model package which downloads the check point from NGC
    package = FCN.load_default_package()
    model = FCN.load_model(package)

    # Instantiate the pertubation method
    sg = SphericalGaussian(noise_amplitude=0.15)

    # Create the data source
    data = GFS()

    # Create the IO handler, store in memory
    chunks = {"ensemble": 1, "time": 1}
    io = ZarrBackend(
        file_name="outputs/03_ensemble_sg.zarr",
        chunks=chunks,
        backend_kwargs={"overwrite": True},
    )


.. GENERATED FROM PYTHON SOURCE LINES 93-102

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 10 steps (for FCN, this is 60 hours) with 8 ensemble
members which will be ran in 2 batches with batch size 4.

.. GENERATED FROM PYTHON SOURCE LINES 104-120

.. code-block:: Python


    nsteps = 10
    nensemble = 8
    batch_size = 2
    io = ensemble(
        ["2024-01-01"],
        nsteps,
        nensemble,
        model,
        data,
        io,
        sg,
        batch_size=batch_size,
        output_coords={"variable": np.array(["t2m", "tcwv"])},
    )


.. GENERATED FROM PYTHON SOURCE LINES 121-127

Post Processing
---------------
The last step is to post process our results. Cartopy is a great library for plotting
fields on projections of a sphere.

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

.. GENERATED FROM PYTHON SOURCE LINES 129-185

.. code-block:: Python

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

    forecast = "2024-01-01"


    def plot_(axi, data, title, cmap):
        """Convenience function for plotting pcolormesh."""
        # Plot the field using pcolormesh
        im = axi.pcolormesh(
            io["lon"][:],
            io["lat"][:],
            data,
            transform=ccrs.PlateCarree(),
            cmap=cmap,
        )
        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()


    for variable, cmap in zip(["tcwv"], ["Blues"]):
        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, (ax1, ax2, ax3) = plt.subplots(
            nrows=1, ncols=3, subplot_kw={"projection": projection}, figsize=(16, 3)
        )

        plot_(
            ax1,
            io[variable][0, 0, step],
            f"{forecast} - Lead time: {6*step}hrs - Member: {0}",
            cmap,
        )
        plot_(
            ax2,
            io[variable][1, 0, step],
            f"{forecast} - Lead time: {6*step}hrs - Member: {1}",
            cmap,
        )
        plot_(
            ax3,
            np.std(io[variable][:, 0, step], axis=0),
            f"{forecast} - Lead time: {6*step}hrs - Std",
            cmap,
        )

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


.. _sphx_glr_download_examples_03_ensemble_workflow.py:

.. only:: html

  .. container:: sphx-glr-footer sphx-glr-footer-example

    .. container:: sphx-glr-download sphx-glr-download-jupyter

      :download:`Download Jupyter notebook: 03_ensemble_workflow.ipynb <03_ensemble_workflow.ipynb>`

    .. container:: sphx-glr-download sphx-glr-download-python

      :download:`Download Python source code: 03_ensemble_workflow.py <03_ensemble_workflow.py>`

    .. container:: sphx-glr-download sphx-glr-download-zip

      :download:`Download zipped: 03_ensemble_workflow.zip <03_ensemble_workflow.zip>`


.. only:: html

 .. rst-class:: sphx-glr-signature

    `Gallery generated by Sphinx-Gallery <https://sphinx-gallery.github.io>`_