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

.. only:: html

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

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

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

.. _sphx_glr_examples_12_temporal_interpolation.py:


Temporal Interpolation
======================

This example demonstrates how to use the InterpModAFNO model to interpolate
forecasts from a base model to a finer time resolution.

In this example you will learn:

- How to load a base prognostic model
- How to load the InterpModAFNO model
- How to run the interpolation model
- How to visualize the results

.. GENERATED FROM PYTHON SOURCE LINES 34-46

Set Up
------
First, import the necessary modules and set up our environment and load the models.
We will use SFNO as the base prognostic model and the InterpModAFNO model to
interpolate its output to a finer time resolution.
The prognostic model must predict the needed variables in the interpolation model.

This example needs the following:

- Interpolation Model: :py:class:`earth2studio.models.px.InterpModAFNO`.
- Prognostic Base Model: Use SFNO model :py:class:`earth2studio.models.px.SFNO`.
- Datasource: Pull data from the GFS data api :py:class:`earth2studio.data.GFS`.

.. GENERATED FROM PYTHON SOURCE LINES 48-73

.. code-block:: Python

    import os

    import matplotlib.pyplot as plt

    from earth2studio.data import GFS
    from earth2studio.io import ZarrBackend
    from earth2studio.models.px import SFNO, InterpModAFNO

    # Create output directory
    os.makedirs("outputs", exist_ok=True)

    sfno_package = SFNO.load_default_package()
    base_model = SFNO.load_model(sfno_package)

    # Load the interpolation model
    interp_package = InterpModAFNO.load_default_package()
    interp_model = InterpModAFNO.load_model(interp_package)
    interp_model.px_model = base_model  # Set the base model

    # Create the data source
    data = GFS()

    # Create the IO handler
    io = ZarrBackend()


.. GENERATED FROM PYTHON SOURCE LINES 74-79

Run the Interpolation Model
---------------------------
Now run the interpolation model to get forecasts at a finer time resolution.
The base model (SFNO) produces forecasts at 6-hour intervals, and the
interpolation model will interpolate to 1-hour intervals.

.. GENERATED FROM PYTHON SOURCE LINES 81-93

.. code-block:: Python


    # Define forecast parameters
    forecast_date = "2024-01-01"
    nsteps = 5  # Number of interpolated forecast steps

    # Run the model
    from earth2studio.run import deterministic

    io = deterministic([forecast_date], nsteps, interp_model, data, io)

    print(io.root.tree())


.. GENERATED FROM PYTHON SOURCE LINES 94-98

Visualize Results
-----------------
Let's visualize the total column water vapour (tcwv) at each time step
and save them as separate files.

.. GENERATED FROM PYTHON SOURCE LINES 100-119

.. code-block:: Python


    # Get the number of time steps
    n_steps = io["tcwv"].shape[1]

    # Create a single figure with subplots
    fig, axs = plt.subplots(2, 3, figsize=(15, 6))
    axs = axs.ravel()

    # Create plots for each time step
    for step in range(min([n_steps, 6])):
        im = axs[step].imshow(
            io["tcwv"][0, step], cmap="twilight_shifted", aspect="auto", vmin=0, vmax=85
        )
        axs[step].set_title(f"Water Vapour - Step: {step}hrs")
        fig.colorbar(im, ax=axs[step], label="kg/m^2")

    plt.tight_layout()
    # Save the figure
    plt.savefig("outputs/12_tcwv_steps.jpg")


.. _sphx_glr_download_examples_12_temporal_interpolation.py:

.. only:: html

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

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

      :download:`Download Jupyter notebook: 12_temporal_interpolation.ipynb <12_temporal_interpolation.ipynb>`

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

      :download:`Download Python source code: 12_temporal_interpolation.py <12_temporal_interpolation.py>`

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

      :download:`Download zipped: 12_temporal_interpolation.zip <12_temporal_interpolation.zip>`


.. only:: html

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

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