init

FengWu

@@ -0,0 +1 @@
+Subproject commit 3bdf35bd6a84600e95c8d534fec727c69e4e7982

Pangu-Weather

@@ -0,0 +1 @@
+Subproject commit 72bdd99096721e1a1f8912c37a9a3aff9ff0a4f2

Prithvi.py

@@ -0,0 +1,2505 @@
+import streamlit as st
+import torch
+import random
+import numpy as np
+import yaml
+from pathlib import Path
+from io import BytesIO
+import random
+from pathlib import Path
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+from huggingface_hub import hf_hub_download, snapshot_download
+import tempfile
+import traceback
+import functools as ft
+import os
+import random
+import re
+from collections import defaultdict
+from datetime import datetime, timedelta
+from pathlib import Path
+import h5py
+import numpy as np
+import pandas as pd
+import torch
+from torch import Tensor
+from torch.utils.data import Dataset
+import logging
+from Prithvi import *
+
+def preproc(batch: list[dict], padding: dict[tuple[int]]) -> dict[str, Tensor]:
+    """Prepressing function for MERRA2 Dataset
+
+    Args:
+        batch (dict): List of training samples, each sample should be a
+            dictionary with the following keys::
+
+            'sur_static': Numpy array of shape (3, lat, lon). For each pixel (lat, lon), the first dimension indexes sin(lat), cos(lon), sin(lon).
+            'sur_vals': Torch tensor of shape (parameter, time, lat, lon).
+            'sur_tars': Torch tensor of shape (parameter, time, lat, lon).
+            'ulv_vals': Torch tensor of shape (parameter, level, time, lat, lon).
+            'ulv_tars': Torch tensor of shape (parameter, level, time, lat, lon).
+            'sur_climate': Torch tensor of shape (parameter, lat, lon)
+            'ulv_climate': Torch tensor of shape (parameter, level, lat, lon)
+            'lead_time': Integer.
+            'input_time': Integer.
+
+        padding: Dictionary with keys 'level', 'lat', 'lon', each of dim 2.
+
+    Returns:
+        Dictionary with the following keys::
+
+            'x': [batch, time, parameter, lat, lon]
+            'y': [batch, parameter, lat, lon]
+            'static': [batch, parameter, lat, lon]
+            'lead_time': [batch]
+            'input_time': [batch]
+            'climate (Optional)': [batch, parameter, lat, lon]
+
+    Note:
+        Here, for x and y, 'parameter' is [surface parameter, upper level,
+        parameter x level]. Similarly for the static information we have
+        [sin(lat), cos(lon), sin(lon), cos(doy), sin(doy), cos(hod), sin(hod),
+        ...].
+    """  # noqa: E501
+    b0 = batch[0]
+    nbatch = len(batch)
+    data_keys = set(b0.keys())
+
+    essential_keys = {
+        "sur_static",
+        "sur_vals",
+        "sur_tars",
+        "ulv_vals",
+        "ulv_tars",
+        "input_time",
+        "lead_time",
+    }
+
+    climate_keys = {
+        "sur_climate",
+        "ulv_climate",
+    }
+
+    all_keys = essential_keys | climate_keys
+
+    if not essential_keys.issubset(data_keys):
+        raise ValueError("Missing essential keys.")
+
+    if not data_keys.issubset(all_keys):
+        raise ValueError("Unexpected keys in batch.")
+
+    # Bring all tensors from the batch into a single tensor
+    upl_x = torch.empty((nbatch, *b0["ulv_vals"].shape))
+    upl_y = torch.empty((nbatch, *b0["ulv_tars"].shape))
+
+    sur_x = torch.empty((nbatch, *b0["sur_vals"].shape))
+    sur_y = torch.empty((nbatch, *b0["sur_tars"].shape))
+
+    sur_sta = torch.empty((nbatch, *b0["sur_static"].shape))
+
+    lead_time = torch.empty((nbatch,), dtype=torch.float32)
+    input_time = torch.empty((nbatch,), dtype=torch.float32)
+
+    for i, rec in enumerate(batch):
+        sur_x[i] = rec["sur_vals"]
+        sur_y[i] = rec["sur_tars"]
+
+        upl_x[i] = rec["ulv_vals"]
+        upl_y[i] = rec["ulv_tars"]
+
+        sur_sta[i] = rec["sur_static"]
+
+        lead_time[i] = rec["lead_time"]
+        input_time[i] = rec["input_time"]
+
+    return_value = {
+        "lead_time": lead_time,
+        "input_time": input_time,
+    }
+
+    # Reshape (batch, parameter, level, time, lat, lon) ->
+    #  (batch, time, parameter, level, lat, lon)
+    upl_x = upl_x.permute((0, 3, 1, 2, 4, 5))
+    upl_y = upl_y.permute((0, 3, 1, 2, 4, 5))
+    # Reshape (batch, parameter, time, lat, lon) ->
+    #  (batch, time, parameter, lat, lon)
+    sur_x = sur_x.permute((0, 2, 1, 3, 4))
+    sur_y = sur_y.permute((0, 2, 1, 3, 4))
+
+    # Pad
+    padding_2d = (*padding["lon"], *padding["lat"])
+
+    def pad2d(x):
+        return torch.nn.functional.pad(x, padding_2d, mode="constant", value=0)
+
+    padding_3d = (*padding["lon"], *padding["lat"], *padding["level"])
+
+    def pad3d(x):
+        return torch.nn.functional.pad(x, padding_3d, mode="constant", value=0)
+
+    sur_x = pad2d(sur_x).contiguous()
+    upl_x = pad3d(upl_x).contiguous()
+    sur_y = pad2d(sur_y).contiguous()
+    upl_y = pad3d(upl_y).contiguous()
+    return_value["static"] = pad2d(sur_sta).contiguous()
+
+    # Remove time for targets
+    upl_y = torch.squeeze(upl_y, 1)
+    sur_y = torch.squeeze(sur_y, 1)
+
+    # We stack along the combined parameter x level dimension
+    return_value["x"] = torch.cat(
+        (sur_x, upl_x.view(*upl_x.shape[:2], -1, *upl_x.shape[4:])), dim=2
+    )
+    return_value["y"] = torch.cat(
+        (sur_y, upl_y.view(upl_y.shape[0], -1, *upl_y.shape[3:])), dim=1
+    )
+
+    if climate_keys.issubset(data_keys):
+        sur_climate = torch.empty((nbatch, *b0["sur_climate"].shape))
+        ulv_climate = torch.empty((nbatch, *b0["ulv_climate"].shape))
+        for i, rec in enumerate(batch):
+            sur_climate[i] = rec["sur_climate"]
+            ulv_climate[i] = rec["ulv_climate"]
+        sur_climate = pad2d(sur_climate)
+        ulv_climate = pad3d(ulv_climate)
+
+        return_value["climate"] = torch.cat(
+            (
+                sur_climate,
+                ulv_climate.view(nbatch, -1, *ulv_climate.shape[3:]),
+            ),
+            dim=1,
+        )
+
+    return return_value
+
+
+def input_scalers(
+    surf_vars: list[str],
+    vert_vars: list[str],
+    levels: list[float],
+    surf_path: str | Path,
+    vert_path: str | Path,
+) -> tuple[Tensor, Tensor]:
+    """Reads the input scalers
+
+    Args:
+        surf_vars: surface variables to be used.
+        vert_vars: vertical variables to be used.
+        levels: MERRA2 levels to use.
+        surf_path: path to surface scalers file.
+        vert_path: path to vertical level scalers file.
+
+    Returns:
+        mu (Tensor): mean values
+        var (Tensor): varience values
+    """
+    with h5py.File(Path(surf_path), "r", libver="latest") as surf_file:
+        stats = [x.decode().lower() for x in surf_file["statistic"][()]]
+        mu_idx = stats.index("mu")
+        sig_idx = stats.index("sigma")
+
+        s_mu = torch.tensor([surf_file[k][()][mu_idx] for k in surf_vars])
+        s_sig = torch.tensor([surf_file[k][()][sig_idx] for k in surf_vars])
+
+    with h5py.File(Path(vert_path), "r", libver="latest") as vert_file:
+        stats = [x.decode().lower() for x in vert_file["statistic"][()]]
+        mu_idx = stats.index("mu")
+        sig_idx = stats.index("sigma")
+
+        lvl = vert_file["lev"][()]
+        l_idx = [np.where(lvl == v)[0].item() for v in levels]
+
+        v_mu = np.array([vert_file[k][()][mu_idx, l_idx] for k in vert_vars])
+        v_sig = np.array([vert_file[k][()][sig_idx, l_idx] for k in vert_vars])
+
+    v_mu = torch.from_numpy(v_mu).view(-1)
+    v_sig = torch.from_numpy(v_sig).view(-1)
+
+    mu = torch.cat((s_mu, v_mu), dim=0).to(torch.float32)
+    sig = torch.cat((s_sig, v_sig), dim=0).to(torch.float32).clamp(1e-4, 1e4)
+    return mu, sig
+
+
+def static_input_scalers(
+    scalar_path: str | Path, stat_vars: list[str], unscaled_params: int = 7
+) -> tuple[Tensor, Tensor]:
+    scalar_path = Path(scalar_path)
+
+    with h5py.File(scalar_path, "r", libver="latest") as scaler_file:
+        stats = [x.decode().lower() for x in scaler_file["statistic"][()]]
+        mu_idx = stats.index("mu")
+        sig_idx = stats.index("sigma")
+
+        mu = torch.tensor([scaler_file[k][()][mu_idx] for k in stat_vars])
+        sig = torch.tensor([scaler_file[k][()][sig_idx] for k in stat_vars])
+
+    z = torch.zeros(unscaled_params, dtype=mu.dtype, device=mu.device)
+    o = torch.ones(unscaled_params, dtype=sig.dtype, device=sig.device)
+    mu = torch.cat((z, mu), dim=0).to(torch.float32)
+    sig = torch.cat((o, sig), dim=0).to(torch.float32)
+
+    return mu, sig.clamp(1e-4, 1e4)
+
+
+def output_scalers(
+    surf_vars: list[str],
+    vert_vars: list[str],
+    levels: list[float],
+    surf_path: str | Path,
+    vert_path: str | Path,
+) -> Tensor:
+    surf_path = Path(surf_path)
+    vert_path = Path(vert_path)
+
+    with h5py.File(surf_path, "r", libver="latest") as surf_file:
+        svars = torch.tensor([surf_file[k][()] for k in surf_vars])
+
+    with h5py.File(vert_path, "r", libver="latest") as vert_file:
+        lvl = vert_file["lev"][()]
+        l_idx = [np.where(lvl == v)[0].item() for v in levels]
+        vvars = np.array([vert_file[k][()][l_idx] for k in vert_vars])
+    vvars = torch.from_numpy(vvars).view(-1)
+
+    var = torch.cat((svars, vvars), dim=0).to(torch.float32).clamp(1e-7, 1e7)
+
+    return var
+
+
+class SampleSpec:
+    """
+    A data class to collect the information used to define a sample.
+    """
+
+    def __init__(
+        self,
+        inputs: tuple[pd.Timestamp, pd.Timestamp],
+        lead_time: int,
+        target: pd.Timestamp | list[pd.Timestamp],
+    ):
+        """
+        Args:
+            inputs: Tuple of timestamps. In ascending order.
+            lead_time: Lead time. In hours.
+            target: Timestamp of the target. Can be before or after the inputs.
+        """
+        if not inputs[0] < inputs[1]:
+            raise ValueError(
+                "Timestamps in `inputs` should be in strictly ascending order."
+            )
+
+        self.inputs = inputs
+        self.input_time = (inputs[1] - inputs[0]).total_seconds() / 3600
+        self.lead_time = lead_time
+        self.target = target
+
+        self.times = [*inputs, target]
+        self.stat_times = [inputs[-1]]
+
+    @property
+    def climatology_info(self) -> tuple[int, int]:
+        """Get the required climatology info.
+
+        :return: information required to obtain climatology data. Essentially
+            this is the day of the year and hour of the day of the target
+            timestamp, with the former restricted to the interval [1, 365].
+        :rtype: tuple
+        """
+        return (min(self.target.dayofyear, 365), self.target.hour)
+
+    @property
+    def year(self) -> int:
+        return self.inputs[1].year
+
+    @property
+    def dayofyear(self) -> int:
+        return self.inputs[1].dayofyear
+
+    @property
+    def hourofday(self) -> int:
+        return self.inputs[1].hour
+
+    def _info_str(self) -> str:
+        iso_8601 = "%Y-%m-%dT%H:%M:%S"
+
+        return (
+            f"Issue time: {self.inputs[1].strftime(iso_8601)}\n"
+            f"Lead time: {self.lead_time} hours ahead\n"
+            f"Input delta: {self.input_time} hours\n"
+            f"Target time: {self.target.strftime(iso_8601)}"
+        )
+
+    @classmethod
+    def get(cls, timestamp: pd.Timestamp, dt: int, lead_time: int):
+        """Given a timestamp and lead time, generates a SampleSpec object
+        describing the sample further.
+
+        Args:
+            timestamp: Timstamp of the sample, Ie this is the larger of the two
+                input timstamps.
+            dt: Time between input samples, in hours.
+            lead_time: Lead time. In hours.
+
+        Returns:
+            SampleSpec
+        """  # noqa: E501
+        assert dt > 0, "dt should be possitive"
+        lt = pd.to_timedelta(lead_time, unit="h")
+        dt = pd.to_timedelta(dt, unit="h")
+
+        if lead_time >= 0:
+            timestamp_target = timestamp + lt
+        else:
+            timestamp_target = timestamp - dt + lt
+
+        spec = cls(
+            inputs=(timestamp - dt, timestamp),
+            lead_time=lead_time,
+            target=timestamp_target,
+        )
+
+        return spec
+
+    def __repr__(self) -> str:
+        return self._info_str()
+
+    def __str__(self) -> str:
+        return self._info_str()
+
+
+class Merra2Dataset(Dataset):
+    """MERRA2 dataset. The dataset unifies surface and vertical data as well as
+    optional climatology.
+
+    Samples come in the form of a dictionary. Not all keys support all
+    variables, yet the general ordering of dimensions is
+    parameter, level, time, lat, lon
+
+    Note:
+        Data is assumed to be in NetCDF files containing daily data at 3-hourly
+        intervals. These follow the naming patterns
+        MERRA2_sfc_YYYYMMHH.nc and MERRA_pres_YYYYMMHH.nc and can be located in
+        two different locations. Optional climatology data comes from files
+        climate_surface_doyDOY_hourHOD.nc and
+        climate_vertical_doyDOY_hourHOD.nc.
+
+
+    Note:
+        `_get_valid_timestamps` assembles a set of all timestamps for which
+        there is data (with hourly resolutions). The result is stored in
+        `_valid_timestamps`. `_get_valid_climate_timestamps` does the same with
+        climatology data and stores it in `_valid_climate_timestamps`.
+
+        Based on this information, `samples` generates a list of valid samples,
+        stored in `samples`. Here the format is::
+
+            [
+                [
+                    (timestamp 1, lead time A),
+                    (timestamp 1, lead time B),
+                    (timestamp 1, lead time C),
+                ],
+                [
+                    (timestamp 2, lead time D),
+                    (timestamp 2, lead time E),
+                ]
+            ]
+
+        That is, the outer list iterates over timestamps (init times), the
+        inner over lead times. Only valid entries are stored.
+    """
+
+    valid_vertical_vars = [
+        "CLOUD",
+        "H",
+        "OMEGA",
+        "PL",
+        "QI",
+        "QL",
+        "QV",
+        "T",
+        "U",
+        "V",
+    ]
+    valid_surface_vars = [
+        "EFLUX",
+        "GWETROOT",
+        "HFLUX",
+        "LAI",
+        "LWGAB",
+        "LWGEM",
+        "LWTUP",
+        "PRECTOT",
+        "PS",
+        "QV2M",
+        "SLP",
+        "SWGNT",
+        "SWTNT",
+        "T2M",
+        "TQI",
+        "TQL",
+        "TQV",
+        "TS",
+        "U10M",
+        "V10M",
+        "Z0M",
+    ]
+    valid_static_surface_vars = ["FRACI", "FRLAND", "FROCEAN", "PHIS"]
+
+    valid_levels = [
+        34.0,
+        39.0,
+        41.0,
+        43.0,
+        44.0,
+        45.0,
+        48.0,
+        51.0,
+        53.0,
+        56.0,
+        63.0,
+        68.0,
+        71.0,
+        72.0,
+    ]
+
+    timedelta_input = pd.to_timedelta(3, unit="h")
+
+    def __init__(
+        self,
+        time_range: tuple[str | pd.Timestamp, str | pd.Timestamp],
+        lead_times: list[int],
+        input_times: list[int],
+        data_path_surface: str | Path,
+        data_path_vertical: str | Path,
+        climatology_path_surface: str | Path | None = None,
+        climatology_path_vertical: str | Path | None = None,
+        surface_vars: list[str] | None = None,
+        static_surface_vars: list[str] | None = None,
+        vertical_vars: list[str] | None = None,
+        levels: list[float] | None = None,
+        roll_longitudes: int = 0,
+        positional_encoding: str = "absolute",
+        rtype: type = np.float32,
+        dtype: torch.dtype = torch.float32,
+    ) -> None:
+        """
+        Args:
+            data_path_surface: Location of surface data.
+            data_path_vertical: Location of vertical data.
+            climatology_path_surface: Location of (optional) surface
+                climatology.
+            climatology_path_vertical: Location of (optional) vertical
+                climatology.
+            surface_vars: Surface variables.
+            static_surface_vars: Static surface variables.
+            vertical_vars: Vertical variables.
+            levels: Levels.
+            time_range: Used to subset data.
+            lead_times: Lead times for generalized forecasting.
+            roll_longitudes: Set to non-zero value to data by random amount
+                along longitude dimension.
+            position_encoding: possible values are
+              ['absolute' (default), 'fourier'].
+                'absolute' returns lat lon encoded in 3 dimensions using sine
+                  and cosine
+                'fourier' returns lat/lon to be encoded by model
+                <any other key> returns lat/lon to be encoded by model
+            rtype: numpy data type used during read
+            dtype: torch data type of data output
+        """
+
+        self.time_range = (
+            pd.to_datetime(time_range[0]),
+            pd.to_datetime(time_range[1]),
+        )
+        self.lead_times = lead_times
+        self.input_times = input_times
+        self._roll_longitudes = list(range(roll_longitudes + 1))
+
+        self._uvars = vertical_vars or self.valid_vertical_vars
+        self._level = levels or self.valid_levels
+        self._svars = surface_vars or self.valid_surface_vars
+        self._sstat = static_surface_vars or self.valid_static_surface_vars
+        self._nuvars = len(self._uvars)
+        self._nlevel = len(self._level)
+        self._nsvars = len(self._svars)
+        self._nsstat = len(self._sstat)
+
+        self.rtype = rtype
+        self.dtype = dtype
+
+        self.positional_encoding = positional_encoding
+
+        self._data_path_surface = Path(data_path_surface)
+        self._data_path_vertical = Path(data_path_vertical)
+
+        self.dir_exists(self._data_path_surface)
+        self.dir_exists(self._data_path_vertical)
+
+        self._get_coordinates()
+
+        self._climatology_path_surface = Path(climatology_path_surface) or None
+        self._climatology_path_vertical = (
+            Path(climatology_path_vertical) or None
+        )
+        self._require_clim = (
+            self._climatology_path_surface is not None
+            and self._climatology_path_vertical is not None
+        )
+
+        if self._require_clim:
+            self.dir_exists(self._climatology_path_surface)
+            self.dir_exists(self._climatology_path_vertical)
+        elif (
+            climatology_path_surface is None
+            and climatology_path_vertical is None
+        ):
+            self._climatology_path_surface = None
+            self._climatology_path_vertical = None
+        else:
+            raise ValueError(
+                "Either both or neither of"
+                "`climatology_path_surface` and"
+                "`climatology_path_vertical` should be None."
+            )
+
+        if not set(self._svars).issubset(set(self.valid_surface_vars)):
+            raise ValueError("Invalid surface variable.")
+
+        if not set(self._sstat).issubset(set(self.valid_static_surface_vars)):
+            raise ValueError("Invalid static surface variable.")
+
+        if not set(self._uvars).issubset(set(self.valid_vertical_vars)):
+            raise ValueError("Inalid vertical variable.")
+
+        if not set(self._level).issubset(set(self.valid_levels)):
+            raise ValueError("Invalid level.")
+
+    @staticmethod
+    def dir_exists(path: Path) -> None:
+        if not path.is_dir():
+            raise ValueError(f"Directory {path} does not exist.")
+
+    @property
+    def upper_shape(self) -> tuple:
+        """Returns the vertical variables shape
+        Returns:
+            tuple: vertical variable shape in the following order::
+
+                [VAR, LEV, TIME, LAT, LON]
+        """
+        return self._nuvars, self._nlevel, 2, 361, 576
+
+    @property
+    def surface_shape(self) -> tuple:
+        """Returns the surface variables shape
+
+        Returns:
+            tuple: surafce shape in the following order::
+
+                [VAR, LEV, TIME, LAT, LON]
+        """
+        return self._nsvars, 2, 361, 576
+
+    def data_file_surface(self, timestamp: pd.Timestamp) -> Path:
+        """Build the surfcae data file name based on timestamp
+
+        Args:
+            timestamp: a timestamp
+
+        Returns:
+            Path: constructed path
+        """
+        pattern = "MERRA2_sfc_%Y%m%d.nc"
+        data_file = self._data_path_surface / timestamp.strftime(pattern)
+        return data_file
+
+    def data_file_vertical(self, timestamp: pd.Timestamp) -> Path:
+        """Build the vertical data file name based on timestamp
+
+        Args:
+            timestamp: a timestamp
+
+        Returns:
+            Path: constructed path
+        """
+        pattern = "MERRA_pres_%Y%m%d.nc"
+        data_file = self._data_path_vertical / timestamp.strftime(pattern)
+        return data_file
+
+    def data_file_surface_climate(
+        self,
+        timestamp: pd.Timestamp | None = None,
+        dayofyear: int | None = None,
+        hourofday: int | None = None,
+    ) -> Path:
+        """
+        Returns the path to a climatology file based either on a timestamp or
+        the dayofyear / hourofday combination.
+        Args:
+            timestamp: A timestamp.
+            dayofyear: Day of the year. 1 to 366.
+            hourofday: Hour of the day. 0 to 23.
+        Returns:
+            Path: Path to climatology file.
+        """
+        if timestamp is not None and (
+            (dayofyear is not None) or (hourofday is not None)
+        ):
+            raise ValueError(
+                "Provide either timestamp or both dayofyear and hourofday."
+            )
+
+        if timestamp is not None:
+            dayofyear = min(timestamp.dayofyear, 365)
+            hourofday = timestamp.hour
+
+        file_name = f"climate_surface_doy{dayofyear:03}_hour{hourofday:02}.nc"
+        data_file = self._climatology_path_surface / file_name
+        return data_file
+
+    def data_file_vertical_climate(
+        self,
+        timestamp: pd.Timestamp | None = None,
+        dayofyear: int | None = None,
+        hourofday: int | None = None,
+    ) -> Path:
+        """Returns the path to a climatology file based either on a timestamp
+        or the dayofyear / hourofday combination.
+
+        Args:
+            timestamp: A timestamp. dayofyear: Day of the year. 1 to 366.
+            hourofday: Hour of the day. 0 to 23.
+        Returns:
+            Path: Path to climatology file.
+        """
+        if timestamp is not None and (
+            (dayofyear is not None) or (hourofday is not None)
+        ):
+            raise ValueError(
+                "Provide either timestamp or both dayofyear and hourofday."
+            )
+
+        if timestamp is not None:
+            dayofyear = min(timestamp.dayofyear, 365)
+            hourofday = timestamp.hour
+
+        file_name = f"climate_vertical_doy{dayofyear:03}_hour{hourofday:02}.nc"
+        data_file = self._climatology_path_vertical / file_name
+        return data_file
+
+    def _get_coordinates(self) -> None:
+        """
+        Obtains the coordiantes (latitudes and longitudes) from a single data
+        file.
+        """
+        timestamp = next(iter(self.valid_timestamps))
+
+        file = self.data_file_surface(timestamp)
+        with h5py.File(file, "r", libver="latest") as handle:
+            self.lats = lats = handle["lat"][()].astype(self.rtype)
+            self.lons = lons = handle["lon"][()].astype(self.rtype)
+
+        deg_to_rad = np.pi / 180
+        self._embed_lat = np.sin(lats * deg_to_rad).reshape(-1, 1)
+
+        self._embed_lon = np.empty((2, 1, len(lons)), dtype=self.rtype)
+        self._embed_lon[0, 0] = np.cos(lons * deg_to_rad)
+        self._embed_lon[1, 0] = np.sin(lons * deg_to_rad)
+
+    @ft.cached_property
+    def lats(self) -> np.ndarray:
+        timestamp = next(iter(self.valid_timestamps))
+
+        file = self.data_file_surface(timestamp)
+        with h5py.File(file, "r", libver="latest") as handle:
+            return handle["lat"][()].astype(self.rtype)
+
+    @ft.cached_property
+    def lons(self) -> np.ndarray:
+        timestamp = next(iter(self.valid_timestamps))
+
+        file = self.data_file_surface(timestamp)
+        with h5py.File(file, "r", libver="latest") as handle:
+            return handle["lon"][()].astype(self.rtype)
+
+    @ft.cached_property
+    def position_signal(self) -> np.ndarray:
+        """Generates the "position signal" that is part of the static
+        features.
+
+        Returns:
+            Tensor: Torch tensor of dimension (parameter, lat, lon) containing
+            sin(lat), cos(lon), sin(lon).
+        """
+
+        latitudes, longitudes = np.meshgrid(
+            self.lats, self.lons, indexing="ij"
+        )
+
+        if self.positional_encoding == "absolute":
+            latitudes = latitudes / 360 * 2.0 * np.pi
+            longitudes = longitudes / 360 * 2.0 * np.pi
+            sur_static = np.stack(
+                [np.sin(latitudes), np.cos(longitudes), np.sin(longitudes)],
+                axis=0,
+            )
+        else:
+            sur_static = np.stack([latitudes, longitudes], axis=0)
+
+        sur_static = sur_static.astype(self.rtype)
+
+        return sur_static
+
+    @ft.cached_property
+    def valid_timestamps(self) -> set[pd.Timestamp]:
+        """Generates list of valid timestamps based on available files. Only
+        timestamps for which both surface and vertical information is available
+        are considered valid.
+        Returns:
+            list: list of timestamps
+        """
+
+        s_glob = self._data_path_surface.glob("MERRA2_sfc_????????.nc")
+        s_files = [os.path.basename(f) for f in s_glob]
+        v_glob = self._data_path_surface.glob("MERRA_pres_????????.nc")
+        v_files = [os.path.basename(f) for f in v_glob]
+
+        s_re = re.compile(r"MERRA2_sfc_(\d{8}).nc\Z")
+        v_re = re.compile(r"MERRA_pres_(\d{8}).nc\Z")
+        fmt = "%Y%m%d"
+
+        s_times = {
+            (datetime.strptime(m[1], fmt))
+            for f in s_files
+            if (m := s_re.match(f))
+        }
+        v_times = {
+            (datetime.strptime(m[1], fmt))
+            for f in v_files
+            if (m := v_re.match(f))
+        }
+
+        times = s_times.intersection(v_times)
+
+        # Each file contains a day at 3 hour intervals
+        times = {
+            t + timedelta(hours=i) for i in range(0, 24, 3) for t in times
+        }
+
+        start_time, end_time = self.time_range
+        times = {pd.Timestamp(t) for t in times if start_time <= t <= end_time}
+
+        return times
+
+    @ft.cached_property
+    def valid_climate_timestamps(self) -> set[tuple[int, int]]:
+        """Generates list of "timestamps" (dayofyear, hourofday) for which
+        climatology data is present. Only instances for which surface and
+        vertical data is available are considered valid.
+        Returns:
+            list: List of tuples describing valid climatology instances.
+        """
+        if not self._require_clim:
+            return set()
+
+        s_glob = self._climatology_path_surface.glob(
+            "climate_surface_doy???_hour??.nc"
+        )
+        s_files = [os.path.basename(f) for f in s_glob]
+
+        v_glob = self._climatology_path_vertical.glob(
+            "climate_vertical_doy???_hour??.nc"
+        )
+        v_files = [os.path.basename(f) for f in v_glob]
+
+        s_re = re.compile(r"climate_surface_doy(\d{3})_hour(\d{2}).nc\Z")
+        v_re = re.compile(r"climate_vertical_doy(\d{3})_hour(\d{2}).nc\Z")
+
+        s_times = {
+            (int(m[1]), int(m[2])) for f in s_files if (m := s_re.match(f))
+        }
+        v_times = {
+            (int(m[1]), int(m[2])) for f in v_files if (m := v_re.match(f))
+        }
+
+        times = s_times.intersection(v_times)
+
+        return times
+
+    def _data_available(self, spec: SampleSpec) -> bool:
+        """
+        Checks whether data is available for a given SampleSpec object. Does so
+        using the internal sets with available data previously constructed. Not
+        by checking the file system.
+        Args:
+            spec: SampleSpec object as returned by SampleSpec.get
+        Returns:
+            bool: if data is availability.
+        """
+        valid = set(spec.times).issubset(self.valid_timestamps)
+
+        if self._require_clim:
+            sci = spec.climatology_info
+            ci = set(sci) if isinstance(sci, list) else set([sci])  # noqa: C405
+            valid &= ci.issubset(self.valid_climate_timestamps)
+
+        return valid
+
+    @ft.cached_property
+    def samples(self) -> list[tuple[pd.Timestamp, int, int]]:
+        """
+        Generates list of all valid samlpes.
+        Returns:
+            list: List of tuples (timestamp, input time, lead time).
+        """
+        valid_samples = []
+        dts = [(it, lt) for it in self.input_times for lt in self.lead_times]
+
+        for timestamp in sorted(self.valid_timestamps):
+            timestamp_samples = []
+            for it, lt in dts:
+                spec = SampleSpec.get(timestamp, -it, lt)
+
+                if self._data_available(spec):
+                    timestamp_samples.append((timestamp, it, lt))
+
+            if timestamp_samples:
+                valid_samples.append(timestamp_samples)
+
+        return valid_samples
+
+    def _to_torch(
+        self,
+        data: dict[str, Tensor | list[Tensor]],
+        dtype: torch.dtype = torch.float32,
+    ) -> dict[str, Tensor | list[Tensor]]:
+        out = {}
+        for k, v in data.items():
+            if isinstance(v, list):
+                out[k] = [torch.from_numpy(x).to(dtype) for x in v]
+            else:
+                out[k] = torch.from_numpy(v).to(dtype)
+
+        return out
+
+    def _lat_roll(
+        self, data: dict[str, Tensor | list[Tensor]], n: int
+    ) -> dict[str, Tensor | list[Tensor]]:
+        out = {}
+        for k, v in data.items():
+            if isinstance(v, list):
+                out[k] = [torch.roll(x, shifts=n, dims=-1) for x in v]
+            else:
+                out[k] = torch.roll(v, shifts=n, dims=-1)
+
+        return out
+
+    def _read_static_data(
+        self, file: str | Path, doy: int, hod: int
+    ) -> np.ndarray:
+        with h5py.File(file, "r", libver="latest") as handle:
+            lats_surf = handle["lat"]
+            lons_surf = handle["lon"]
+
+            nll = (len(lats_surf), len(lons_surf))
+
+            npos = len(self.position_signal)
+            ntime = 4
+
+            nstat = npos + ntime + self._nsstat
+            data = np.empty((nstat, *nll), dtype=self.rtype)
+
+            for i, key in enumerate(self._sstat, start=npos + ntime):
+                data[i] = handle[key][()].astype(dtype=self.rtype)
+
+        # [possition signal], cos(doy), sin(doy), cos(hod), sin(hod)
+        data[0:npos] = self.position_signal
+        data[npos + 0] = np.cos(2 * np.pi * doy / 366)
+        data[npos + 1] = np.sin(2 * np.pi * doy / 366)
+        data[npos + 2] = np.cos(2 * np.pi * hod / 24)
+        data[npos + 3] = np.sin(2 * np.pi * hod / 24)
+
+        return data
+
+    def _read_surface(
+        self, tidx: int, nll: tuple[int, int], handle: h5py.File
+    ) -> np.ndarray:
+        data = np.empty((self._nsvars, *nll), dtype=self.rtype)
+
+        for i, key in enumerate(self._svars):
+            data[i] = handle[key][tidx][()].astype(dtype=self.rtype)
+
+        return data
+
+    def _read_levels(
+        self, tidx: int, nll: tuple[int, int], handle: h5py.File
+    ) -> np.ndarray:
+        lvls = handle["lev"][()]
+        lidx = self._level_idxs(lvls)
+
+        data = np.empty((self._nuvars, self._nlevel, *nll), dtype=self.rtype)
+
+        for i, key in enumerate(self._uvars):
+            data[i] = handle[key][tidx, lidx][()].astype(dtype=self.rtype)
+
+        return np.ascontiguousarray(np.flip(data, axis=1))
+
+    def _level_idxs(self, lvls):
+        lidx = [np.argwhere(lvls == int(lvl)).item() for lvl in self._level]
+        return sorted(lidx)
+
+    @staticmethod
+    def _date_to_tidx(date: datetime | pd.Timestamp, handle: h5py.File) -> int:
+        if isinstance(date, pd.Timestamp):
+            date = date.to_pydatetime()
+
+        time = handle["time"]
+
+        t0 = time.attrs["begin_time"][()].item()
+        d0 = f"{time.attrs['begin_date'][()].item()}"
+
+        offset = datetime.strptime(d0, "%Y%m%d")
+
+        times = [offset + timedelta(minutes=int(t + t0)) for t in time[()]]
+        return times.index(date)
+
+    def _read_data(
+        self, file_pair: tuple[str, str], date: datetime
+    ) -> dict[str, np.ndarray]:
+        s_file, v_file = file_pair
+
+        with h5py.File(s_file, "r", libver="latest") as shandle:
+            lats_surf = shandle["lat"]
+            lons_surf = shandle["lon"]
+
+            nll = (len(lats_surf), len(lons_surf))
+
+            tidx = self._date_to_tidx(date, shandle)
+
+            sdata = self._read_surface(tidx, nll, shandle)
+
+        with h5py.File(v_file, "r", libver="latest") as vhandle:
+            lats_vert = vhandle["lat"]
+            lons_vert = vhandle["lon"]
+
+            nll = (len(lats_vert), len(lons_vert))
+
+            tidx = self._date_to_tidx(date, vhandle)
+
+            vdata = self._read_levels(tidx, nll, vhandle)
+
+        data = {"vert": vdata, "surf": sdata}
+
+        return data
+
+    def _read_climate(
+        self, file_pair: tuple[str, str]
+    ) -> dict[str, np.ndarray]:
+        s_file, v_file = file_pair
+
+        with h5py.File(s_file, "r", libver="latest") as shandle:
+            lats_surf = shandle["lat"]
+            lons_surf = shandle["lon"]
+
+            nll = (len(lats_surf), len(lons_surf))
+
+            sdata = np.empty((self._nsvars, *nll), dtype=self.rtype)
+
+            for i, key in enumerate(self._svars):
+                sdata[i] = shandle[key][()].astype(dtype=self.rtype)
+
+        with h5py.File(v_file, "r", libver="latest") as vhandle:
+            lats_vert = vhandle["lat"]
+            lons_vert = vhandle["lon"]
+
+            nll = (len(lats_vert), len(lons_vert))
+
+            lvls = vhandle["lev"][()]
+            lidx = self._level_idxs(lvls)
+
+            vdata = np.empty(
+                (self._nuvars, self._nlevel, *nll), dtype=self.rtype
+            )
+
+            for i, key in enumerate(self._uvars):
+                vdata[i] = vhandle[key][lidx][()].astype(dtype=self.rtype)
+
+        data = {
+            "vert": np.ascontiguousarray(np.flip(vdata, axis=1)),
+            "surf": sdata,
+        }
+
+        return data
+
+    def get_data_from_sample_spec(
+        self, spec: SampleSpec
+    ) -> dict[str, Tensor | int | float]:
+        """Loads and assembles sample data given a SampleSpec object.
+
+        Args:
+            spec (SampleSpec): Full details regarding the data to be loaded
+        Returns:
+            dict: Dictionary with the following keys::
+
+                'sur_static': Torch tensor of shape [parameter, lat, lon]. For
+                each pixel (lat, lon), the first 7 dimensions index sin(lat),
+                cos(lon), sin(lon), cos(doy), sin(doy), cos(hod), sin(hod).
+                Where doy is the day of the year [1, 366] and hod the hour of
+                the day [0, 23].
+                'sur_vals': Torch tensor of shape [parameter, time, lat, lon].
+                'sur_tars': Torch tensor of shape [parameter, time, lat, lon].
+                'ulv_vals': Torch tensor of shape [parameter, level, time, lat, lon].
+                'ulv_tars': Torch tensor of shape [parameter, level, time, lat, lon].
+                'sur_climate': Torch tensor of shape [parameter, lat, lon].
+                'ulv_climate': Torch tensor of shape [paramter, level, lat, lon].
+                'lead_time': Float.
+                'input_time': Float.
+
+        """  # noqa: E501
+
+        # We assemble the unique timestamps for which we need data.
+        vals_required = {*spec.times}
+        stat_required = {*spec.stat_times}
+
+        # We assemble the unique data files from which we need value data
+        vals_file_map = defaultdict(list)
+        for t in vals_required:
+            data_files = (
+                self.data_file_surface(t),
+                self.data_file_vertical(t),
+            )
+            vals_file_map[data_files].append(t)
+
+        # We assemble the unique data files from which we need static data
+        stat_file_map = defaultdict(list)
+        for t in stat_required:
+            data_files = (
+                self.data_file_surface(t),
+                self.data_file_vertical(t),
+            )
+            stat_file_map[data_files].append(t)
+
+        # Load the value data
+        data = {}
+        for data_files, times in vals_file_map.items():
+            for time in times:
+                data[time] = self._read_data(data_files, time)
+
+        # Combine times
+        sample_data = {}
+
+        input_upl = np.stack([data[t]["vert"] for t in spec.inputs], axis=2)
+        sample_data["ulv_vals"] = input_upl
+
+        target_upl = data[spec.target]["vert"]
+        sample_data["ulv_tars"] = target_upl[:, :, None]
+
+        input_sur = np.stack([data[t]["surf"] for t in spec.inputs], axis=1)
+        sample_data["sur_vals"] = input_sur
+
+        target_sur = data[spec.target]["surf"]
+        sample_data["sur_tars"] = target_sur[:, None]
+
+        # Load the static data
+        data_files, times = stat_file_map.popitem()
+        time = times[0].dayofyear, times[0].hour
+        sample_data["sur_static"] = self._read_static_data(
+            data_files[0], *time
+        )
+
+        # If required load the surface data
+        if self._require_clim:
+            ci_year, ci_hour = spec.climatology_info
+
+            surf_file = self.data_file_surface_climate(
+                dayofyear=ci_year,
+                hourofday=ci_hour,
+            )
+
+            vert_file = self.data_file_vertical_climate(
+                dayofyear=ci_year,
+                hourofday=ci_hour,
+            )
+
+            clim_data = self._read_climate((surf_file, vert_file))
+
+            sample_data["sur_climate"] = clim_data["surf"]
+            sample_data["ulv_climate"] = clim_data["vert"]
+
+        # Move the data from numpy to torch
+        sample_data = self._to_torch(sample_data, dtype=self.dtype)
+
+        # Optionally roll
+        if len(self._roll_longitudes) > 0:
+            roll_by = random.choice(self._roll_longitudes)
+            sample_data = self._lat_roll(sample_data, roll_by)
+
+        # Now that we have rolled, we can add the static data
+        sample_data["lead_time"] = spec.lead_time
+        sample_data["input_time"] = spec.input_time
+
+        return sample_data
+
+    def get_data(
+        self, timestamp: pd.Timestamp, input_time: int, lead_time: int
+    ) -> dict[str, Tensor | int]:
+        """
+        Loads data based on timestamp and lead time.
+        Args:
+            timestamp: Timestamp.
+            input_time: time between input samples.
+            lead_time: lead time.
+         Returns:
+            Dictionary with keys 'sur_static', 'sur_vals', 'sur_tars',
+                'ulv_vals', 'ulv_tars', 'sur_climate', 'ulv_climate',
+                'lead_time'.
+        """
+        spec = SampleSpec.get(timestamp, -input_time, lead_time)
+        sample_data = self.get_data_from_sample_spec(spec)
+        return sample_data
+
+    def __getitem__(self, idx: int) -> dict[str, Tensor | int]:
+        """
+        Loads data based on sample index and random choice of sample.
+        Args:
+            idx: Sample index.
+         Returns:
+            Dictionary with keys 'sur_static', 'sur_vals', 'sur_tars',
+                'ulv_vals', 'ulv_tars', 'sur_climate', 'ulv_climate',
+                'lead_time', 'input_time'.
+        """
+        sample_set = self.samples[idx]
+        timestamp, input_time, lead_time, *nsteps = random.choice(sample_set)
+        sample_data = self.get_data(timestamp, input_time, lead_time)
+        return sample_data
+
+    def __len__(self):
+        return len(self.samples)
+
+from functools import cached_property
+from importlib.metadata import version
+
+from torch import Tensor
+from torch.utils.checkpoint import checkpoint
+
+if version("torch") > "2.3.0":
+    from torch.nn.attention import SDPBackend, sdpa_kernel
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+# DropPath code is straight from timm
+# (https://huggingface.co/spaces/Roll20/pet_score/blame/main/lib/timm/models/layers/drop.py)
+def drop_path(
+    x: Tensor,
+    drop_prob: float = 0.0,
+    training: bool = False,
+    scale_by_keep: bool = True,
+) -> Tensor:
+    """Drop paths (Stochastic Depth) per sample (when applied in main path of
+    residual blocks). Taken form timm.
+
+    Args:
+        x (Tensor): Input tensor.
+        drop_prob (float): Probability of dropping `x`, defaults to 0.
+        training (bool): Whether model is in in traingin of eval mode,
+            defaults to False.
+        scale_by_keep (bool): Whether the output should scaled by
+            (`1 - drop_prob`), defaults to True.
+    Returns:
+        Tensor: Tensor that may have randomly dropped with proability
+            `drop_path`
+    """
+    if drop_prob == 0.0 or not training:
+        return x
+    keep_prob = 1 - drop_prob
+    shape = (x.shape[0],) + (1,) * (x.ndim - 1)
+    random_tensor = x.new_empty(shape).bernoulli_(keep_prob)
+    if keep_prob > 0.0 and scale_by_keep:
+        random_tensor.div_(keep_prob)
+    return x * random_tensor
+
+
+class DropPath(nn.Module):
+    """
+    Drop paths (Stochastic Depth) per sample (when applied in main path of
+    residual blocks).
+    """
+
+    def __init__(
+        self, drop_prob: float | None = None, scale_by_keep: bool = True
+    ) -> None:
+        super(DropPath, self).__init__()
+        self.drop_prob = drop_prob
+        self.scale_by_keep = scale_by_keep
+
+    def forward(self, x: Tensor) -> Tensor:
+        """Runs drop path on input tensor
+
+        Args:
+            x: input
+
+        Returns:
+            tensor: output after drop_path
+        """
+        return drop_path(x, self.drop_prob, self.training, self.scale_by_keep)
+
+
+class Mlp(nn.Module):
+    """
+    Multi layer perceptron.
+    """
+
+    def __init__(
+        self, features: int, hidden_features: int, dropout: float = 0.0
+    ) -> None:
+        """
+        Args:
+            features: Input/output dimension.
+            hidden_features: Hidden dimension.
+            dropout: Dropout.
+        """
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(features, hidden_features),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(hidden_features, features),
+            nn.Dropout(dropout),
+        )
+
+    def forward(self, x: Tensor) -> Tensor:
+        """
+        Args:
+            x (Tesnor): Tensor of shape [..., channel]
+        Returns:
+            Tenosr: Tensor of same shape as x.
+        """
+        return self.net(x)
+
+
+class LayerNormPassThrough(nn.LayerNorm):
+    """Normalising layer that allows the attention mask to be passed through"""
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def forward(
+        self, d: tuple[Tensor, Tensor | None]
+    ) -> tuple[Tensor, Tensor | None]:
+        """Forwards function
+
+        Args:
+            d (tuple): tuple of the data tensor and the attention mask
+        Returns:
+            output (Tensor): normalised output data
+            attn_mask (Tensor): the attention mask that was passed in
+        """
+        input, attn_mask = d
+        output = F.layer_norm(
+            input, self.normalized_shape, self.weight, self.bias, self.eps
+        )
+        return output, attn_mask
+
+
+class MultiheadAttention(nn.Module):
+    """Multihead attention layer for inputs of shape
+    [..., sequence, features].
+    """
+
+    def __init__(self, features: int, n_heads: int, dropout: float) -> None:
+        """
+        Args:
+            features: Number of features for inputs to the layer.
+            n_heads: Number of attention heads. Should be a factor of features.
+                (I.e. the layer uses features // n_heads.)
+            dropout: Dropout.
+        """  # noqa: E501
+        super().__init__()
+
+        if (features % n_heads) != 0:
+            raise ValueError(
+                f"Features '{features}' is not divisible by heads '{n_heads}'."
+            )
+
+        self.features = features
+        self.n_heads = n_heads
+        self.dropout = dropout
+
+        self.qkv_layer = torch.nn.Linear(features, features * 3, bias=False)
+        self.w_layer = torch.nn.Linear(features, features, bias=False)
+
+    def forward(self, d: tuple[Tensor, Tensor | None]) -> Tensor:
+        """
+        Args:
+            d (tuple): tuple containing Tensor of shape [..., sequence, features] and the attention mask
+        Returns:
+            Tensor: Tensor of shape [..., sequence, features]
+        """  # noqa: E501
+        x, attn_mask = d
+
+        if not x.shape[-1] == self.features:
+            raise ValueError(
+                f"Expecting tensor with last dimension size {self.features}."
+            )
+
+        passenger_dims = x.shape[:-2]
+        B = passenger_dims.numel()
+        S = x.shape[-2]
+        C = x.shape[-1]
+        x = x.reshape(B, S, C)
+
+        # x [B, S, C]
+        # q, k, v [B, H, S, C/H]
+        q, k, v = (
+            self.qkv_layer(x)
+            .view(B, S, self.n_heads, 3 * (C // self.n_heads))
+            .transpose(1, 2)
+            .chunk(chunks=3, dim=3)
+        )
+
+        # Let us enforce either flash (A100+) or memory efficient attention.
+        if version("torch") > "2.3.0":
+            with sdpa_kernel(
+                [SDPBackend.FLASH_ATTENTION, SDPBackend.EFFICIENT_ATTENTION]
+            ):
+                # x [B, H, S, C//H]
+                x = F.scaled_dot_product_attention(
+                    q, k, v, attn_mask=attn_mask, dropout_p=self.dropout
+                )
+        else:
+            with torch.backends.cuda.sdp_kernel(
+                enable_flash=True, enable_math=False, enable_mem_efficient=True
+            ):
+                # x [B, H, S, C//H]
+                x = F.scaled_dot_product_attention(
+                    q, k, v, dropout_p=self.dropout
+                )
+
+        # x [B, S, C]
+        x = x.transpose(1, 2).view(B, S, C)
+
+        # x [B, S, C]
+        x = self.w_layer(x)
+
+        # Back to input shape
+        x = x.view(*passenger_dims, S, self.features)
+        return x
+
+
+class Transformer(nn.Module):
+    """
+    Transformer for inputs of shape [..., S, features].
+    """
+
+    def __init__(
+        self,
+        features: int,
+        mlp_multiplier: int,
+        n_heads: int,
+        dropout: float,
+        drop_path: float,
+    ) -> None:
+        """
+        Args:
+            features: Number of features for inputs to the layer.
+            mlp_multiplier: Model uses features*mlp_multiplier hidden units.
+            n_heads: Number of attention heads. Should be a factor of features.
+            (I.e. the layer uses features // n_heads.) dropout: Dropout.
+            drop_path: DropPath.
+        """
+        super().__init__()
+
+        self.features = features
+        self.mlp_multiplier = mlp_multiplier
+        self.n_heads = n_heads
+        self.dropout = dropout
+        self.drop_path = (
+            DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+        )
+
+        self.attention = nn.Sequential(
+            LayerNormPassThrough(features),
+            MultiheadAttention(features, n_heads, dropout),
+        )
+
+        self.ff = nn.Sequential(
+            nn.LayerNorm(features),
+            Mlp(
+                features=features,
+                hidden_features=features * mlp_multiplier,
+                dropout=dropout,
+            ),
+        )
+
+    def forward(self, d: tuple[Tensor, Tensor | None]) -> Tensor:
+        """
+        Args:
+            x: Tensor of shape [..., sequence, features]
+        Returns:
+            Tensor: Tensor of shape [..., sequence, features]
+        """
+        x, attn_mask = d
+        if not x.shape[-1] == self.features:
+            raise ValueError(
+                f"Expecting tensor with last dimension size {self.features}."
+            )
+
+        attention_x = self.attention(d)
+
+        x = x + self.drop_path(attention_x)
+        x = x + self.drop_path(self.ff(x))
+
+        return x
+
+
+class _Shift(nn.Module):
+    """Private base class for the shifter. This allows some behaviour to be
+    easily handled when the shifter isn't used.
+    """
+
+    def __init__(self):
+        super().__init__()
+
+        self._shifted = False
+
+    @torch.no_grad()
+    def reset(self) -> None:
+        """
+        Resets the bool tracking whether the data is shifted
+        """
+        self._shifted: bool = False
+
+    def forward(self, data: Tensor) -> tuple[Tensor, dict[bool, None]]:
+        return data, {True: None, False: None}
+
+
+class SWINShift(_Shift):
+    """
+    Handles the shifting of patches similar to how SWIN works. However if we
+    shift the latitudes then the poles will wrap and potentially that might be
+    problematic. The possition tokens should handle it but masking is safer.
+    """
+
+    def __init__(
+        self,
+        mu_shape: tuple[int, int],
+        global_shape: tuple[int, int],
+        local_shape: tuple[int, int],
+        patch_shape: tuple[int, int],
+        n_context_tokens: int = 2,
+    ) -> None:
+        """
+        Args:
+            mu_shape: the shape to the masking units
+            global_shape: number of global patches in lat and lon
+            local_shape: size of the local patches
+            patch_shape: patch size
+            n_context_token: number of additional context tokens at start of
+            _each_ local sequence
+        """
+        super().__init__()
+
+        self._mu_shape = ms = mu_shape
+        self._g_shape = gs = global_shape
+        self._l_shape = ls = local_shape
+        self._p_shape = ps = patch_shape
+        self._lat_patch = (gs[0], ls[0], gs[1], ls[1])
+        self._n_context_tokens = n_context_tokens
+
+        self._g_shift_to = tuple(
+            int(0.5 * x / p) for x, p in zip(ms, ps, strict=False)
+        )
+        self._g_shift_from = tuple(
+            -int(0.5 * x / p) for x, p in zip(ms, ps, strict=False)
+        )
+
+        # Define the attention masks for the shifted MaxViT.
+        nglobal = global_shape[0] * global_shape[1]
+        nlocal = (
+            local_shape[0] * local_shape[1] + self._n_context_tokens
+        )  # "+ 1" for leadtime
+
+        lm = torch.ones((nglobal, 1, nlocal, nlocal), dtype=bool)
+        mwidth = int(0.5 * local_shape[1]) * local_shape[0]
+        lm[
+            : gs[1],
+            :,
+            self._n_context_tokens : mwidth + self._n_context_tokens,
+            self._n_context_tokens : mwidth + self._n_context_tokens,
+        ] = False
+        self.register_buffer("local_mask", lm)
+
+        gm = torch.ones((nlocal, 1, nglobal, nglobal), dtype=bool)
+        gm[: int(0.5 * ls[1]) * ls[0], :, : gs[1], : gs[1]] = False
+        self.register_buffer("global_mask", gm)
+
+    def _to_grid_global(self, x: Tensor) -> Tensor:
+        """
+        Shuffle and reshape the data from the global/local setting back to the
+        lat/lon grid setting
+        Args:
+            x: the data tensor to be shuffled.
+        Returns:
+            x: data in the global/local setting
+        """
+        nbatch, *other = x.shape
+
+        y1 = x.view(nbatch, *self._g_shape, *self._l_shape, -1)
+        y2 = y1.permute(0, 5, 1, 3, 2, 4).contiguous()
+
+        s = y2.shape
+        return y2.view((nbatch, -1, s[2] * s[3], s[4] * s[5]))
+
+    def _to_grid_local(self, x: Tensor) -> Tensor:
+        """
+        Shuffle and reshape the data from the local/global setting to the
+        lat/lon grid setting
+        Args:
+            x: the data tensor to be shuffled.
+        Returns:
+            x: data in the lat/lon setting.
+        """
+        x = x.transpose(2, 1).contiguous()
+        return self._to_grid_global(x)
+
+    def _from_grid_global(self, x: Tensor) -> Tensor:
+        """
+        Shuffle and reshape the data from the lat/lon grid to the global/local
+        setting
+        Args:
+            x: the data tensor to be shuffled.
+        Returns:
+            x: data in the global/local setting
+        """
+        nbatch, *other = x.shape
+
+        z1 = x.view(nbatch, -1, *self._lat_patch)
+        z2 = z1.permute(0, 2, 4, 3, 5, 1).contiguous()
+
+        s = z2.shape
+        return z2.view(nbatch, s[1] * s[2], s[3] * s[4], -1)
+
+    def _from_grid_local(self, x: Tensor) -> Tensor:
+        """
+        Shuffle and reshape the data from the lat/lon grid to the local/global
+        setting
+        Args:
+            x: the data tensor to be shuffled.
+        Returns:
+            x: data in the local/global setting
+        """
+        x = self._from_grid_global(x)
+        return x.transpose(2, 1).contiguous()
+
+    def _shift(self, x: Tensor) -> Tensor:
+        """
+        Shifts data in the gridded lat/lon setting by half the mask unit shape
+        Args:
+            x: data to be shifted
+        Returns:
+            x: either the hsifted or unshifted data
+        """
+        shift = self._g_shift_from if self._shifted else self._g_shift_to
+        x_shifted = torch.roll(x, shift, (-2, -1))
+
+        self._shifted = not self._shifted
+        return x_shifted
+
+    def _sep_lt(self, x: Tensor) -> tuple[Tensor, Tensor]:
+        """
+        Seperate off the leadtime from the local patches
+        Args:
+            x: data to have leadtime removed from
+        Returns:
+            lt: leadtime
+            x: data without the lead time in the local patch
+        """
+        lt_it = x[:, : self._n_context_tokens, :, :]
+        x_stripped = x[:, self._n_context_tokens :, :, :]
+
+        return lt_it, x_stripped
+
+    def forward(self, data: Tensor) -> tuple[Tensor, Tensor]:
+        """Shift or unshift the the data depending on whether the data is
+        already shifted, as defined by self._shifte.
+
+        Args:
+            data: data to be shifted
+        Returns:
+            Tensor: shifted data Tensor
+        """
+        lt, x = self._sep_lt(data)
+
+        x_grid = self._to_grid_local(x)
+        x_shifted = self._shift(x_grid)
+        x_patched = self._from_grid_local(x_shifted)
+
+        # Mask has to be repeated based on batch size
+        n_batch = x_grid.shape[0]
+        local_rep = [n_batch] + [1] * (self.local_mask.ndim - 1)
+        global_rep = [n_batch] + [1] * (self.global_mask.ndim - 1)
+
+        if self._shifted:
+            attn_mask = {
+                True: self.local_mask.repeat(local_rep),
+                False: self.global_mask.repeat(global_rep),
+            }
+        else:
+            attn_mask = {True: None, False: None}
+
+        return torch.cat((lt, x_patched), axis=1), attn_mask
+
+
+class LocalGlobalLocalBlock(nn.Module):
+    """
+    Applies alternating block and grid attention. Given a parameter n_blocks,
+    the entire module contains 2*n_blocks+1 transformer blocks. The first,
+    third, ..., last apply local (block) attention. The second, fourth, ...
+    global (grid) attention.
+
+    This is heavily inspired by
+    Tu et al. "MaxViT: Multi-Axis Vision Transformer"
+    (https://arxiv.org/abs/2204.01697).
+    """
+
+    def __init__(
+        self,
+        features: int,
+        mlp_multiplier: int,
+        n_heads: int,
+        dropout: float,
+        n_blocks: int,
+        drop_path: float,
+        shifter: nn.Module | None = None,
+        checkpoint: list[int] | None = None,
+    ) -> None:
+        """
+        Args:
+            features: Number of features for inputs to the layer.
+            mlp_multiplier: Model uses features*mlp_multiplier hidden units.
+            n_heads: Number of attention heads. Should be a factor of features.
+            (I.e. the layer uses features // n_heads.)
+            dropout: Dropout.
+            drop_path: DropPath.
+            n_blocks: Number of local-global transformer pairs.
+        """
+        super().__init__()
+
+        self.features = features
+        self.mlp_multiplier = mlp_multiplier
+        self.n_heads = n_heads
+        self.dropout = dropout
+        self.drop_path = drop_path
+        self.n_blocks = n_blocks
+        self._checkpoint = checkpoint or []
+
+        if not all(0 <= c < 2 * n_blocks + 1 for c in self._checkpoint):
+            raise ValueError(
+                "Checkpoints should be 0 <= i < 2*n_blocks+1. "
+                f"{self._checkpoint=}."
+            )
+
+        self.transformers = nn.ModuleList(
+            [
+                Transformer(
+                    features=features,
+                    mlp_multiplier=mlp_multiplier,
+                    n_heads=n_heads,
+                    dropout=dropout,
+                    drop_path=drop_path,
+                )
+                for _ in range(2 * n_blocks + 1)
+            ]
+        )
+
+        self.evaluator = [
+            self._checkpoint_wrapper
+            if i in self._checkpoint
+            else lambda m, x: m(x)
+            for i, _ in enumerate(self.transformers)
+        ]
+
+        self.shifter = shifter or _Shift()
+
+    @staticmethod
+    def _checkpoint_wrapper(
+        model: nn.Module, data: tuple[Tensor, Tensor | None]
+    ) -> Tensor:
+        return checkpoint(model, data, use_reentrant=False)
+
+    def forward(self, x: Tensor) -> Tensor:
+        """
+        Args:
+            x: Tensor of shape::
+
+                [batch, global_sequence, local_sequence, features]
+
+        Returns:
+            Tensor: Tensor of shape::
+
+                [batch, global_sequence, local_sequence, features]
+        """
+        if x.shape[-1] != self.features:
+            raise ValueError(
+                f"Expecting tensor with last dimension size {self.features}."
+            )
+        if x.ndim != 4:
+            raise ValueError(
+                f"Expecting tensor with exactly four dimensions. {x.shape=}."
+            )
+
+        self.shifter.reset()
+        local: bool = True
+        attn_mask = {True: None, False: None}
+
+        transformer_iter = zip(self.evaluator, self.transformers, strict=False)
+
+        # First local block
+        evaluator, transformer = next(transformer_iter)
+        x = evaluator(transformer, (x, attn_mask[local]))
+
+        for evaluator, transformer in transformer_iter:
+            local = not local
+            # We are making exactly 2*n_blocks transposes.
+            # So the output has the same shape as input.
+            x = x.transpose(1, 2)
+
+            x = evaluator(transformer, (x, attn_mask[local]))
+
+            if not local:
+                x, attn_mask = self.shifter(x)
+
+        return x
+
+
+class PatchEmbed(nn.Module):
+    """
+    Patch embedding via 2D convolution.
+    """
+
+    def __init__(
+        self, patch_size: int | tuple[int, ...], channels: int, embed_dim: int
+    ):
+        super().__init__()
+
+        self.patch_size = patch_size
+        self.channels = channels
+        self.embed_dim = embed_dim
+
+        self.proj = nn.Conv2d(
+            channels,
+            embed_dim,
+            kernel_size=patch_size,
+            stride=patch_size,
+            bias=True,
+        )
+
+    def forward(self, x: Tensor) -> Tensor:
+        """
+        Args:
+            x: Tensor of shape [batch, channels, lat, lon].
+        Returns:
+            Tensor: Tensor with shape
+                [batch, embed_dim, lat//patch_size, lon//patch_size]
+        """
+
+        H, W = x.shape[-2:]
+
+        if W % self.patch_size[1] != 0:
+            raise ValueError(
+                f"Cannot do patch embedding for tensor of shape {x.size()}"
+                " with patch size {self.patch_size}. (Dimensions are BSCHW.)"
+            )
+        if H % self.patch_size[0] != 0:
+            raise ValueError(
+                f"Cannot do patch embedding for tensor of shape {x.size()}"
+                f" with patch size {self.patch_size}. (Dimensions are BSCHW.)"
+            )
+
+        x = self.proj(x)
+
+        return x
+
+
+class PrithviWxCEncoderDecoder(nn.Module):
+    """
+    Hiera-MaxViT encoder/decoder code.
+    """
+
+    def __init__(
+        self,
+        embed_dim: int,
+        n_blocks: int,
+        mlp_multiplier: float,
+        n_heads: int,
+        dropout: float,
+        drop_path: float,
+        shifter: nn.Module | None = None,
+        transformer_cp: list[int] | None = None,
+    ) -> None:
+        """
+        Args:
+            embed_dim: Embedding dimension
+            n_blocks: Number of local-global transformer pairs.
+            mlp_multiplier: MLP multiplier for hidden features in feed forward
+                networks.
+            n_heads: Number of attention heads.
+            dropout: Dropout.
+            drop_path: DropPath.
+        """
+        super().__init__()
+
+        self.embed_dim = embed_dim
+        self.n_blocks = n_blocks
+        self.mlp_multiplier = mlp_multiplier
+        self.n_heads = n_heads
+        self.dropout = dropout
+        self._transformer_cp = transformer_cp
+
+        self.lgl_block = LocalGlobalLocalBlock(
+            features=embed_dim,
+            mlp_multiplier=mlp_multiplier,
+            n_heads=n_heads,
+            dropout=dropout,
+            drop_path=drop_path,
+            n_blocks=n_blocks,
+            shifter=shifter,
+            checkpoint=transformer_cp,
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x: Tensor of shape
+              [batch, global sequence, local sequence, embed_dim]
+        Returns:
+            Tensor of shape
+                [batch, mask_unit_sequence, local_sequence, embed_dim].
+                Identical in shape to the input x.
+        """
+
+        x = self.lgl_block(x)
+
+        return x
+
+
+class PrithviWxC(nn.Module):
+    """Encoder-decoder fusing Hiera with MaxViT. See
+    - Ryali et al. "Hiera: A Hierarchical Vision Transformer without the
+    Bells-and-Whistles" (https://arxiv.org/abs/2306.00989)
+    - Tu et al. "MaxViT: Multi-Axis Vision Transformer"
+    (https://arxiv.org/abs/2204.01697)
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        input_size_time: int,
+        in_channels_static: int,
+        input_scalers_mu: Tensor,
+        input_scalers_sigma: Tensor,
+        input_scalers_epsilon: float,
+        static_input_scalers_mu: Tensor,
+        static_input_scalers_sigma: Tensor,
+        static_input_scalers_epsilon: float,
+        output_scalers: Tensor,
+        n_lats_px: int,
+        n_lons_px: int,
+        patch_size_px: tuple[int],
+        mask_unit_size_px: tuple[int],
+        mask_ratio_inputs: float,
+        embed_dim: int,
+        n_blocks_encoder: int,
+        n_blocks_decoder: int,
+        mlp_multiplier: float,
+        n_heads: int,
+        dropout: float,
+        drop_path: float,
+        parameter_dropout: float,
+        residual: str,
+        masking_mode: str,
+        positional_encoding: str,
+        decoder_shifting: bool = False,
+        checkpoint_encoder: list[int] | None = None,
+        checkpoint_decoder: list[int] | None = None,
+    ) -> None:
+        """
+        Args:
+            in_channels: Number of input channels.
+            input_size_time: Number of timestamps in input.
+            in_channels_static: Number of input channels for static data.
+            input_scalers_mu: Tensor of size (in_channels,). Used to rescale
+                input.
+            input_scalers_sigma: Tensor of size (in_channels,). Used to rescale
+                input.
+            input_scalers_epsilon: Float. Used to rescale input.
+            static_input_scalers_mu: Tensor of size (in_channels_static). Used
+                to rescale static inputs.
+            static_input_scalers_sigma: Tensor of size (in_channels_static).
+                Used to rescale static inputs.
+            static_input_scalers_epsilon: Float. Used to rescale static inputs.
+            output_scalers: Tensor of shape (in_channels,). Used to rescale
+                output.
+            n_lats_px: Total latitudes in data. In pixels.
+            n_lons_px: Total longitudes in data. In pixels.
+            patch_size_px: Patch size for tokenization. In pixels lat/lon.
+            mask_unit_size_px: Size of each mask unit. In pixels lat/lon.
+            mask_ratio_inputs: Masking ratio for inputs. 0 to 1.
+            embed_dim: Embedding dimension
+            n_blocks_encoder: Number of local-global transformer pairs in
+                encoder.
+            n_blocks_decoder: Number of local-global transformer pairs in
+                decoder.
+            mlp_multiplier: MLP multiplier for hidden features in feed forward
+                networks.
+            n_heads: Number of attention heads.
+            dropout: Dropout.
+            drop_path: DropPath.
+            parameter_dropout: Dropout applied to parameters.
+            residual: Indicates whether and how model should work as residual
+                model. Accepted values are 'climate', 'temporal' and 'none'
+            positional_encoding: possible values are
+              ['absolute' (default), 'fourier'].
+                'absolute'  lat lon encoded in 3 dimensions using sine and
+                  cosine
+                'fourier' lat/lon to be encoded using various frequencies
+            masking_mode: String ['local', 'global', 'both'] that controls the
+                type of masking used.
+            checkpoint_encoder: List of integers controlling if gradient
+              checkpointing is used on encoder.
+                Format: [] for no gradient checkpointing. [3, 7] for
+                  checkpointing after 4th and 8th layer etc.
+            checkpoint_decoder: List of integers controlling if gradient
+              checkpointing is used on decoder.
+                Format: See `checkpoint_encoder`.
+            masking_mode: The type of masking to use
+              {'global', 'local', 'both'}
+            decoder_shifting: Whether to use swin shifting in the decoder.
+        """
+        super().__init__()
+
+        self.in_channels = in_channels
+        self.input_size_time = input_size_time
+        self.in_channels_static = in_channels_static
+        self.n_lats_px = n_lats_px
+        self.n_lons_px = n_lons_px
+        self.patch_size_px = patch_size_px
+        self.mask_unit_size_px = mask_unit_size_px
+        self.mask_ratio_inputs = mask_ratio_inputs
+        self.embed_dim = embed_dim
+        self.n_blocks_encoder = n_blocks_encoder
+        self.n_blocks_decoder = n_blocks_decoder
+        self.mlp_multiplier = mlp_multiplier
+        self.n_heads = n_heads
+        self.dropout = dropout
+        self.drop_path = drop_path
+        self.residual = residual
+        self._decoder_shift = decoder_shifting
+        self.positional_encoding = positional_encoding
+        self._checkpoint_encoder = checkpoint_encoder
+        self._checkpoint_decoder = checkpoint_decoder
+
+        assert self.n_lats_px % self.mask_unit_size_px[0] == 0
+        assert self.n_lons_px % self.mask_unit_size_px[1] == 0
+        assert self.mask_unit_size_px[0] % self.patch_size_px[0] == 0
+        assert self.mask_unit_size_px[1] % self.patch_size_px[1] == 0
+
+        if self.patch_size_px[0] != self.patch_size_px[1]:
+            raise NotImplementedError(
+                "Current pixel shuffle symmetric patches."
+            )
+
+        self.local_shape_mu = (
+            self.mask_unit_size_px[0] // self.patch_size_px[0],
+            self.mask_unit_size_px[1] // self.patch_size_px[1],
+        )
+        self.global_shape_mu = (
+            self.n_lats_px // self.mask_unit_size_px[0],
+            self.n_lons_px // self.mask_unit_size_px[1],
+        )
+
+        assert input_scalers_mu.shape == (in_channels,)
+        assert input_scalers_sigma.shape == (in_channels,)
+        assert output_scalers.shape == (in_channels,)
+
+        if self.positional_encoding != "fourier":
+            assert static_input_scalers_mu.shape == (in_channels_static,)
+            assert static_input_scalers_sigma.shape == (in_channels_static,)
+
+        # Input shape [batch, time, parameter, lat, lon]
+        self.input_scalers_epsilon = input_scalers_epsilon
+        self.register_buffer(
+            "input_scalers_mu", input_scalers_mu.reshape(1, 1, -1, 1, 1)
+        )
+        self.register_buffer(
+            "input_scalers_sigma", input_scalers_sigma.reshape(1, 1, -1, 1, 1)
+        )
+
+        # Static inputs shape [batch, parameter, lat, lon]
+        self.static_input_scalers_epsilon = static_input_scalers_epsilon
+        self.register_buffer(
+            "static_input_scalers_mu",
+            static_input_scalers_mu.reshape(1, -1, 1, 1),
+        )
+        self.register_buffer(
+            "static_input_scalers_sigma",
+            static_input_scalers_sigma.reshape(1, -1, 1, 1),
+        )
+
+        # Output shape [batch, parameter, lat, lon]
+        self.register_buffer(
+            "output_scalers", output_scalers.reshape(1, -1, 1, 1)
+        )
+
+        self.parameter_dropout = nn.Dropout2d(p=parameter_dropout)
+
+        self.patch_embedding = PatchEmbed(
+            patch_size=patch_size_px,
+            channels=in_channels * input_size_time,
+            embed_dim=embed_dim,
+        )
+
+        if self.residual == "climate":
+            self.patch_embedding_static = PatchEmbed(
+                patch_size=patch_size_px,
+                channels=in_channels + in_channels_static,
+                embed_dim=embed_dim,
+            )
+        else:
+            self.patch_embedding_static = PatchEmbed(
+                patch_size=patch_size_px,
+                channels=in_channels_static,
+                embed_dim=embed_dim,
+            )
+
+        self.input_time_embedding = nn.Linear(1, embed_dim // 4, bias=True)
+        self.lead_time_embedding = nn.Linear(1, embed_dim // 4, bias=True)
+
+        self.mask_token = nn.Parameter(torch.randn(1, 1, 1, self.embed_dim))
+        self._nglobal_mu = np.prod(self.global_shape_mu)
+        self._global_idx = torch.arange(self._nglobal_mu)
+
+        self._nlocal_mu = np.prod(self.local_shape_mu)
+        self._local_idx = torch.arange(self._nlocal_mu)
+
+        self.encoder = PrithviWxCEncoderDecoder(
+            embed_dim=embed_dim,
+            n_blocks=n_blocks_encoder,
+            mlp_multiplier=mlp_multiplier,
+            n_heads=n_heads,
+            dropout=dropout,
+            drop_path=drop_path,
+            transformer_cp=checkpoint_encoder,
+        )
+
+        if n_blocks_decoder != 0:
+            if self._decoder_shift:
+                self.decoder_shifter = d_shifter = SWINShift(
+                    self.mask_unit_size_px,
+                    self.global_shape_mu,
+                    self.local_shape_mu,
+                    self.patch_size_px,
+                    n_context_tokens=0,
+                )
+            else:
+                self.decoder_shifter = d_shifter = None
+
+            self.decoder = PrithviWxCEncoderDecoder(
+                embed_dim=embed_dim,
+                n_blocks=n_blocks_decoder,
+                mlp_multiplier=mlp_multiplier,
+                n_heads=n_heads,
+                dropout=dropout,
+                drop_path=0.0,
+                shifter=d_shifter,
+                transformer_cp=checkpoint_decoder,
+            )
+
+            self.unembed = nn.Linear(
+                self.embed_dim,
+                self.in_channels
+                * self.patch_size_px[0]
+                * self.patch_size_px[1],
+                bias=True,
+            )
+
+        self.masking_mode = masking_mode.lower()
+        match self.masking_mode:
+            case "local":
+                self.generate_mask = self._gen_mask_local
+            case "global":
+                self.generate_mask = self._gen_mask_global
+            case "both":
+                self._mask_both_local: bool = True
+                self.generate_mask = self._gen_mask_both
+            case _:
+                raise ValueError(
+                    f"Masking mode '{masking_mode}' not supported"
+                )
+
+    def swap_masking(self) -> None:
+        self._mask_both_local = not self._mask_both_local
+
+    @cached_property
+    def n_masked_global(self):
+        return int(self.mask_ratio_inputs * np.prod(self.global_shape_mu))
+
+    @cached_property
+    def n_masked_local(self):
+        return int(self.mask_ratio_inputs * np.prod(self.local_shape_mu))
+
+    @staticmethod
+    def _shuffle_along_axis(a, axis):
+        idx = torch.argsort(input=torch.rand(*a.shape), dim=axis)
+        return torch.gather(a, dim=axis, index=idx)
+
+    def _gen_mask_local(self, sizes: tuple[int]) -> tuple[Tensor]:
+        """
+        Args:
+            batch_size: Number of elements in batch
+        Returns:
+            Tuple of torch tensors. [indices masked, indices unmasked].
+            Each of these is a tensor of shape (batch, global sequene)
+        """
+        # Identify which indices (values) should be masked
+
+        maskable_indices = self._local_idx.view(1, -1).expand(*sizes[:2], -1)
+
+        maskable_indices = self._shuffle_along_axis(maskable_indices, 2)
+
+        indices_masked = maskable_indices[:, :, : self.n_masked_local]
+        indices_unmasked = maskable_indices[:, :, self.n_masked_local :]
+
+        return indices_masked, indices_unmasked
+
+    def _gen_mask_global(self, sizes: tuple[int]) -> tuple[Tensor]:
+        """
+        Args:
+            batch_size: Number of elements in batch
+        Returns:
+            Tuple of torch tensors. [indices masked, indices unmasked].
+            Each of these is a tensor of shape (batch, global sequene)
+        """
+        # Identify which indices (values) should be masked
+
+        maskable_indices = self._global_idx.view(1, -1).expand(*sizes[:1], -1)
+
+        maskable_indices = self._shuffle_along_axis(maskable_indices, 1)
+
+        indices_masked = maskable_indices[:, : self.n_masked_global]
+        indices_unmasked = maskable_indices[:, self.n_masked_global :]
+
+        return indices_masked, indices_unmasked
+
+    def _gen_mask_both(self, sizes: tuple[int]) -> tuple[Tensor]:
+        if self._mask_both_local:
+            return self._gen_mask_local(sizes)
+        else:
+            return self._gen_mask_global(sizes)
+
+    @staticmethod
+    def reconstruct_batch(
+        idx_masked: Tensor,
+        idx_unmasked: Tensor,
+        data_masked: Tensor,
+        data_unmasked: Tensor,
+    ) -> Tensor:
+        """Reconstructs a tensor along the mask unit dimension. Batched
+        version.
+
+        Args:
+            idx_masked: Tensor of shape `batch, mask unit sequence`.
+            idx_unmasked: Tensor of shape `batch, mask unit sequence`.
+            data_masked: Tensor of shape `batch, mask unit sequence, ...`.
+                Should have same size along mask unit sequence dimension as
+                idx_masked. Dimensions beyond the first two, marked here as ...
+                will typically be `local_sequence, channel` or
+                `channel, lat, lon`. These dimensions should agree with
+                data_unmasked.
+            data_unmasked: Tensor of shape `batch, mask unit sequence, ...`.
+                Should have same size along mask unit sequence dimension as
+                idx_unmasked. Dimensions beyond the first two, marked here as
+                ... will typically be `local_sequence, channel` or `channel,
+                lat, lon`. These dimensions should agree with data_masked.
+        Returns:
+            Tensor: Tensor of same shape as inputs data_masked and
+                data_unmasked. I.e. `batch, mask unit sequence, ...`. Index for
+                the total data composed of the masked and the unmasked part.
+        """
+        dim: int = idx_masked.ndim
+
+        idx_total = torch.argsort(
+            torch.cat([idx_masked, idx_unmasked], dim=-1), dim=-1
+        )
+        idx_total = idx_total.view(
+            *idx_total.shape, *[1] * (data_unmasked.ndim - dim)
+        )
+        idx_total = idx_total.expand(
+            *idx_total.shape[:dim], *data_unmasked.shape[dim:]
+        )
+
+        data = torch.cat([data_masked, data_unmasked], dim=dim - 1)
+        data = torch.gather(data, dim=dim - 1, index=idx_total)
+
+        return data, idx_total
+
+    def fourier_pos_encoding(self, x_static: Tensor) -> Tensor:
+        """
+        Args
+            x_static: B x C x H x W. first two channels are lat, and lon
+        Returns
+            Tensor: Tensor of shape B x E x H x W where E is the embedding
+                dimension.
+        """
+
+        # B x C x H x W -> B x 1 x H/P x W/P
+        latitudes_patch = F.avg_pool2d(
+            x_static[:, [0]],
+            kernel_size=self.patch_size_px,
+            stride=self.patch_size_px,
+        )
+        longitudes_patch = F.avg_pool2d(
+            x_static[:, [1]],
+            kernel_size=self.patch_size_px,
+            stride=self.patch_size_px,
+        )
+
+        modes = (
+            torch.arange(self.embed_dim // 4, device=x_static.device).view(
+                1, -1, 1, 1
+            )
+            + 1.0
+        )
+        pos_encoding = torch.cat(
+            (
+                torch.sin(latitudes_patch * modes),
+                torch.sin(longitudes_patch * modes),
+                torch.cos(latitudes_patch * modes),
+                torch.cos(longitudes_patch * modes),
+            ),
+            axis=1,
+        )
+
+        return pos_encoding  # B x E x H/P x W/P
+
+    def time_encoding(self, input_time, lead_time):
+        """
+        Args:
+            input_time: Tensor of shape [batch].
+            lead_time: Tensor of shape [batch].
+        Returns:
+            Tensor: Tensor of shape [batch, embed_dim, 1, 1]
+        """
+        input_time = self.input_time_embedding(input_time.view(-1, 1, 1, 1))
+        lead_time = self.lead_time_embedding(lead_time.view(-1, 1, 1, 1))
+
+        time_encoding = torch.cat(
+            (
+                torch.cos(input_time),
+                torch.cos(lead_time),
+                torch.sin(input_time),
+                torch.sin(lead_time),
+            ),
+            axis=3,
+        )
+        return time_encoding
+
+    def to_patching(self, x: Tensor) -> Tensor:
+        """Transform data from lat/lon space to two axis patching
+
+        Args: ->
+            x: Tesnor in lat/lon space (N, C, Nlat//P_0, Nlon//P_1)
+
+        Returns:
+            Tensor in patch space (N, G, L, C)
+        """
+        n_batch = x.shape[0]
+
+        x = x.view(
+            n_batch,
+            -1,
+            self.global_shape_mu[0],
+            self.local_shape_mu[0],
+            self.global_shape_mu[1],
+            self.local_shape_mu[1],
+        )
+        x = x.permute(0, 2, 4, 3, 5, 1).contiguous()
+
+        s = x.shape
+        return x.view(n_batch, s[1] * s[2], s[3] * s[4], -1)
+
+    def from_patching(self, x: Tensor) -> Tensor:
+        """Transform data from two axis patching to lat/lon space
+
+        Args:
+            x: Tensor in patch space with shape (N, G, L, C*P_0*P_1)
+
+        Returns:
+            Tensor: Tensor in lat/lon space
+                (N, C*P_0*P_1, Nlat//P_0, Nlon // P_1)
+        """
+        n_batch = x.shape[0]
+
+        x = x.view(
+            n_batch,
+            self.global_shape_mu[0],
+            self.global_shape_mu[1],
+            self.local_shape_mu[0],
+            self.local_shape_mu[1],
+            -1,
+        )
+        x = x.permute(0, 5, 1, 3, 2, 4).contiguous()
+
+        s = x.shape
+        return x.view(n_batch, -1, s[2] * s[3], s[4] * s[5])
+
+    def forward(self, batch: dict[str, torch.Tensor]) -> torch.Tensor:
+        """
+        Args:
+            batch: Dictionary the following keys::
+
+                'x': Tensor of shape [batch, time, parameter, lat, lon]
+                'y': Tensor of shape [batch, parameter, lat, lon]
+                'static': Tensor of shape [batch, channel_static, lat, lon]
+                'climate': Optional tensor of shape [batch, parameter, lat, lon]
+                'input_time': Tensor of shape [batch]. Or none.
+                'lead_time': Tensor of shape [batch]. Or none.
+
+        Returns:
+            Tensor: Tensor of shape [batch, parameter, lat, lon].
+        """  # noqa: E501
+        x_rescaled = (batch["x"] - self.input_scalers_mu) / (
+            self.input_scalers_sigma + self.input_scalers_epsilon
+        )
+        batch_size = x_rescaled.shape[0]
+
+        if self.positional_encoding == "fourier":
+            x_static_pos = self.fourier_pos_encoding(batch["static"])
+            x_static = (
+                batch["static"][:, 2:] - self.static_input_scalers_mu[:, 3:]
+            ) / (
+                self.static_input_scalers_sigma[:, 3:]
+                + self.static_input_scalers_epsilon
+            )
+        else:
+            x_static = (batch["static"] - self.static_input_scalers_mu) / (
+                self.static_input_scalers_sigma
+                + self.static_input_scalers_epsilon
+            )
+
+        if self.residual == "temporal":
+            # We create a residual of same shape as y
+            index = torch.where(
+                batch["lead_time"] > 0, batch["x"].shape[1] - 1, 0
+            )
+            index = index.view(-1, 1, 1, 1, 1)
+            index = index.expand(batch_size, 1, *batch["x"].shape[2:])
+            x_hat = torch.gather(batch["x"], dim=1, index=index)
+            x_hat = x_hat.squeeze(1)
+        elif self.residual == "climate":
+            climate_scaled = (
+                batch["climate"] - self.input_scalers_mu.view(1, -1, 1, 1)
+            ) / (
+                self.input_scalers_sigma.view(1, -1, 1, 1)
+                + self.input_scalers_epsilon
+            )
+
+        # [batch, time, parameter, lat, lon]
+        # -> [batch, time x parameter, lat, lon]
+        x_rescaled = x_rescaled.flatten(1, 2)
+        # Parameter dropout
+        x_rescaled = self.parameter_dropout(x_rescaled)
+
+        x_embedded = self.patch_embedding(x_rescaled)
+
+        if self.residual == "climate":
+            static_embedded = self.patch_embedding_static(
+                torch.cat((x_static, climate_scaled), dim=1)
+            )
+        else:
+            static_embedded = self.patch_embedding_static(x_static)
+
+        if self.positional_encoding == "fourier":
+            static_embedded += x_static_pos
+
+        x_embedded = self.to_patching(x_embedded)
+        static_embedded = self.to_patching(static_embedded)
+
+        time_encoding = self.time_encoding(
+            batch["input_time"], batch["lead_time"]
+        )
+
+        tokens = x_embedded + static_embedded + time_encoding
+
+        # Now we generate masks based on masking_mode
+        indices_masked, indices_unmasked = self.generate_mask(
+            (batch_size, self._nglobal_mu)
+        )
+        indices_masked = indices_masked.to(device=tokens.device)
+        indices_unmasked = indices_unmasked.to(device=tokens.device)
+        maskdim: int = indices_masked.ndim
+
+        # Unmasking
+        unmask_view = (*indices_unmasked.shape, *[1] * (tokens.ndim - maskdim))
+        unmasked = torch.gather(
+            tokens,
+            dim=maskdim - 1,
+            index=indices_unmasked.view(*unmask_view).expand(
+                *indices_unmasked.shape, *tokens.shape[maskdim:]
+            ),
+        )
+
+        # Encoder
+        x_encoded = self.encoder(unmasked)
+
+        # Generate and position encode the mask tokens
+        # [1, 1, 1, embed_dim]
+        # -> [batch, global_seq_masked, local seq, embed_dim]
+        mask_view = (*indices_masked.shape, *[1] * (tokens.ndim - maskdim))
+        masking = self.mask_token.repeat(*static_embedded.shape[:3], 1)
+        masked = masking + static_embedded
+        masked = torch.gather(
+            masked,
+            dim=maskdim - 1,
+            index=indices_masked.view(*mask_view).expand(
+                *indices_masked.shape, *tokens.shape[maskdim:]
+            ),
+        )
+
+        recon, _ = self.reconstruct_batch(
+            indices_masked, indices_unmasked, masked, x_encoded
+        )
+
+        x_decoded = self.decoder(recon)
+
+        # Output: [batch, global sequence, local sequence,
+        #          in_channels * patch_size[0] * patch_size[1]]
+        x_unembed = self.unembed(x_decoded)
+
+        # Reshape to [batch, global_lat, global_lon, local_lat, local_lon,
+        #             in_channels * patch_size[0] * patch_size[1]]
+        x_out = self.from_patching(x_unembed)
+
+        # Pixel shuffle to [batch, in_channels, lat, lon]
+        x_out = F.pixel_shuffle(x_out, self.patch_size_px[0])
+
+        if self.residual == "temporal":
+            x_out = self.output_scalers * x_out + x_hat
+        elif self.residual == "climate":
+            x_out = self.output_scalers * x_out + batch["climate"]
+        elif self.residual == "none":
+            x_out = (
+                self.output_scalers * x_out
+                + self.input_scalers_mu.reshape(1, -1, 1, 1)
+            )
+
+        return x_out

app.py

@@ -0,0 +1,285 @@
+import streamlit as st
+import torch
+import random
+import numpy as np
+import yaml
+import logging
+import os
+import matplotlib.pyplot as plt
+from pathlib import Path
+import tempfile
+import traceback
+
+from data_utils import (
+    save_uploaded_files,
+    load_dataset,
+)
+
+from inference_utils import run_inference
+from config_utils import load_config
+from plot_utils import plot_prithvi_output, plot_aurora_output
+from prithvi_utils import (
+    prithvi_config_ui,
+    initialize_prithvi_model,
+    prepare_prithvi_batch
+)
+from aurora_utils import aurora_config_ui, prepare_aurora_batch, initialize_aurora_model
+
+from pangu_utils import (
+    pangu_config_data,
+    inference_1hr,
+    inference_3hrs,
+    inference_6hrs,
+    inference_24hrs,
+    inference_custom_hrs,
+    plot_pangu_output,
+)
+
+from fengwu_utils import (fengwu_config_data, inference_6hrs_fengwu, inference_12hrs_fengwu, inference_custom_hrs_fengwu, plot_fengwu_output)
+
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# Set page configuration
+st.set_page_config(
+    page_title="Weather Data Processor",
+    layout="wide",
+    initial_sidebar_state="expanded",
+)
+
+header_col1, header_col2 = st.columns([4, 1]) 
+with header_col1:
+    st.title("🌦️ Weather & Climate Data Processor and Forecaster")
+
+with header_col2:
+    st.markdown("### Select a Model")
+    selected_model = st.selectbox(
+        "",
+        options=["Pangu-Weather", "FengWu", "Aurora", "Climax", "Prithvi", "LSTM"],
+        index=0,
+        key="model_selector",
+        help="Select the model you want to use."
+    )
+
+st.write("---")
+
+# --- Layout: Two Columns ---
+left_col, right_col = st.columns([1, 2])
+
+with left_col:
+    st.header("🔧 Configuration")
+
+    # Dynamically show configuration UI based on selected model
+    if selected_model == "Prithvi":
+        (config, uploaded_surface_files, uploaded_vertical_files, 
+         clim_surf_path, clim_vert_path, config_path, weights_path) = prithvi_config_ui()
+    elif selected_model == "Aurora":
+        uploaded_files = aurora_config_ui()
+    elif selected_model == "Pangu-Weather":
+        input_surface_file, input_upper_file = pangu_config_data()
+    elif selected_model == "FengWu":
+        input_file1_fengwu, input_file2_fengwu = fengwu_config_data()
+    else:
+        # Generic data upload for other models
+        st.subheader(f"{selected_model} Model Data Upload")
+        st.markdown("### Drag and Drop Your Data Files Here")
+        uploaded_files = st.file_uploader(
+            f"Upload Data Files for {selected_model}",
+            accept_multiple_files=True,
+            key=f"{selected_model.lower()}_uploader",
+            type=["nc", "netcdf", "nc4"],
+        )
+
+    st.write("---")
+
+    # --- Forecast Duration Selection ---
+    st.subheader("Forecast Duration")
+    forecast_options = ["1 hour", "3 hours", "6 hours", "24 hours", "Custom"]
+    selected_duration = st.selectbox(
+        "Select forecast duration",
+        forecast_options,
+        index=3,  # Default to 24 hours
+        help="Select how many hours to forecast."
+    )
+
+    custom_hours = None
+    if selected_duration == "Custom":
+        custom_hours = st.number_input(
+            "Enter custom forecast hours",
+            min_value=24,
+            max_value=480,
+            value=48,
+            step=24,
+            help="Enter the number of hours you want to forecast."
+        )
+
+    st.write("---")
+
+    # Run Inference button
+    if st.button("🚀 Run Inference"):
+        with right_col:
+            st.header("📈 Inference Progress & Visualization")
+
+            # Set seeds and device
+            try:
+                torch.jit.enable_onednn_fusion(True)
+                if torch.cuda.is_available():
+                    device = torch.device("cuda")
+                    st.write(f"Using device: **{torch.cuda.get_device_name()}**")
+                    torch.backends.cudnn.benchmark = True
+                    torch.backends.cudnn.deterministic = True
+                else:
+                    device = torch.device("cpu")
+                    st.write("Using device: **CPU**")
+
+                random.seed(42)
+                if torch.cuda.is_available():
+                    torch.cuda.manual_seed(42)
+                torch.manual_seed(42)
+                np.random.seed(42)
+            except Exception:
+                st.error("Error initializing device:")
+                st.error(traceback.format_exc())
+                st.stop()
+
+            # Use a spinner while running inference
+            with st.spinner("Running inference, please wait..."):
+                # Initialize and run inference for selected model
+                if selected_model == "Prithvi":
+                    model, in_mu, in_sig, output_sig, static_mu, static_sig = initialize_prithvi_model(
+                        config, config_path, weights_path, device
+                    )
+                    batch = prepare_prithvi_batch(
+                        uploaded_surface_files, uploaded_vertical_files, clim_surf_path, clim_vert_path, device
+                    )
+                    out = run_inference(selected_model, model, batch, device)
+                    # Store results
+                    st.session_state['prithvi_out'] = out
+                    st.session_state['prithvi_done'] = True
+
+                elif selected_model == "Aurora":
+                    if uploaded_files:
+                        save_uploaded_files(uploaded_files)
+                        ds = load_dataset(st.session_state.temp_file_paths)
+                        if ds is not None:
+                            batch = prepare_aurora_batch(ds)
+                            model = initialize_aurora_model(device)
+                            out = run_inference(selected_model, model, batch, device)
+                            st.session_state['aurora_out'] = out
+                            st.session_state['aurora_ds_subset'] = ds
+                            st.session_state['aurora_done'] = True
+                        else:
+                            st.error("Failed to load dataset for Aurora.")
+                            st.stop()
+                    else:
+                        st.error("Please upload data files for Aurora.")
+                        st.stop()
+                
+                elif selected_model == "FengWu":
+                    if input_file1_fengwu and input_file2_fengwu:
+                        try:
+                            input1 = np.load(input_file1_fengwu)
+                            input2 = np.load(input_file2_fengwu)
+                            if selected_duration == "1 hour":
+                                st.warning("1hr inference is not yet available on this model.")
+                            elif selected_duration == "3 hours":
+                                st.warning("3hrs inference is not yet available on this model.")
+                            elif selected_duration == "6 hours":
+                                output_fengwu = inference_6hrs_fengwu(input1, input2)
+                            elif selected_duration == "12 hours":
+                                output_fengwu = inference_12hrs_fengwu(input1, input2)
+                            else:
+                                output_fengwu = inference_custom_hrs_fengwu(input1, input2, custom_hours)
+
+                            st.session_state['output_fengwu'] = output_fengwu
+                            st.session_state['fengwu_done'] = True
+                            st.session_state['input_fengwu'] = input_file2_fengwu
+                        except Exception as e:
+                            st.error(f"An error occurred: {e}")
+                    else:
+                        st.error("Please upload data files for Aurora.")
+                        st.stop()
+
+                elif selected_model == "Pangu-Weather":
+                    if input_surface_file and input_upper_file:
+                        try:
+                            surface_data = np.load(input_surface_file)
+                            upper_data = np.load(input_upper_file)
+
+                            # Decide which inference function to use based on selection
+                            if selected_duration == "1 hour":
+                                out_upper, out_surface = inference_1hr(upper_data, surface_data)
+                            elif selected_duration == "3 hours":
+                                out_upper, out_surface = inference_3hrs(upper_data, surface_data)
+                            elif selected_duration == "6 hours":
+                                out_upper, out_surface = inference_6hrs(upper_data, surface_data)
+                            elif selected_duration == "24 hours":
+                                out_upper, out_surface = inference_24hrs(upper_data, surface_data)
+                            else:
+                                out_upper, out_surface = inference_custom_hrs(upper_data, surface_data, custom_hours)
+
+                            # Store results in session_state
+                            st.session_state['pangu_upper_data'] = upper_data
+                            st.session_state['pangu_surface_data'] = surface_data
+                            st.session_state['pangu_out_upper'] = out_upper
+                            st.session_state['pangu_out_surface'] = out_surface
+                            st.session_state['pangu_done'] = True
+
+                            st.write("**Forecast Results:**")
+                            st.write("Upper Data Forecast Shape:", out_upper.shape)
+                            st.write("Surface Data Forecast Shape:", out_surface.shape)
+
+                        except Exception as e:
+                            st.error(f"An error occurred: {e}")
+                    else:
+                        st.error("Please upload data files for Pangu-Weather.")
+                        st.stop()
+
+                else:
+                    st.warning("Inference not implemented for this model.")
+                    st.stop()
+
+            # Visualization after inference is done
+            if selected_model == "Prithvi":
+                if 'prithvi_done' in st.session_state and st.session_state['prithvi_done']:
+                    plot_prithvi_output(st.session_state['prithvi_out'])
+            elif selected_model == "Aurora":
+                if 'aurora_done' in st.session_state and st.session_state['aurora_done']:
+                    plot_aurora_output(st.session_state['aurora_out'], st.session_state['aurora_ds_subset'])
+            elif selected_model == "FengWu":
+                if 'fengwu_done' in st.session_state and st.session_state['fengwu_done']:
+                    plot_fengwu_output(st.session_state['input_fengwu'], st.session_state['output_fengwu'])
+            elif selected_model == "Pangu-Weather":
+                if 'pangu_done' in st.session_state and st.session_state['pangu_done']:
+                    plot_pangu_output(
+                        st.session_state['pangu_upper_data'], 
+                        st.session_state['pangu_surface_data'], 
+                        st.session_state['pangu_out_upper'], 
+                        st.session_state['pangu_out_surface']
+                    )
+            else:
+                st.info("No visualization implemented for this model.")
+
+    else:
+        # If not running inference now, but we have previously computed results, show them
+        with right_col:
+            st.header("🖥️ Visualization & Progress")
+
+            # Check which model was selected and if we have done inference before
+            if selected_model == "Prithvi" and 'prithvi_done' in st.session_state and st.session_state['prithvi_done']:
+                plot_prithvi_output(st.session_state['prithvi_out'])
+            elif selected_model == "Aurora" and 'aurora_done' in st.session_state and st.session_state['aurora_done']:
+                plot_aurora_output(st.session_state['aurora_out'], st.session_state['aurora_ds_subset'])
+            elif selected_model == "Pangu-Weather" and 'pangu_done' in st.session_state and st.session_state['pangu_done']:
+                plot_pangu_output(
+                    st.session_state['pangu_upper_data'], 
+                    st.session_state['pangu_surface_data'], 
+                    st.session_state['pangu_out_upper'], 
+                    st.session_state['pangu_out_surface']
+                )
+            elif selected_model == "FengWu" and 'output_fengwu' in st.session_state and st.session_state['fengwu_done']:
+                plot_fengwu_output(st.session_state['input_fengwu'], st.session_state['output_fengwu'])
+            else:
+                st.info("Awaiting inference to display results.")
+

app1.py

@@ -0,0 +1,477 @@
+import streamlit as st
+import torch
+import random
+import numpy as np
+import yaml
+from pathlib import Path
+from io import BytesIO
+import random
+from pathlib import Path
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+from huggingface_hub import hf_hub_download, snapshot_download
+import tempfile
+import traceback
+import functools as ft
+import os
+import random
+import re
+from collections import defaultdict
+from datetime import datetime, timedelta
+from pathlib import Path
+import h5py
+import numpy as np
+import pandas as pd
+import torch
+from torch import Tensor
+from torch.utils.data import Dataset
+import logging
+from Prithvi import *
+
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+# Set page configuration
+st.set_page_config(
+    page_title="MERRA2 Data Processor",
+    layout="wide",
+    initial_sidebar_state="expanded",
+)
+dataset_type = st.sidebar.selectbox(
+    "Select Dataset Type",
+    options=["MERRA2", "GEOS5"],
+    index=0
+)
+st.title("MERRA2 Data Processor with PrithviWxC Model")
+
+# Sidebar for file uploads
+st.sidebar.header("Upload MERRA2 Data Files")
+
+# File uploader for surface data
+uploaded_surface_files = st.sidebar.file_uploader(
+    "Upload Surface Data Files",
+    type=["nc", "netcdf"],
+    accept_multiple_files=True,
+    key="surface_uploader",
+)
+
+# File uploader for vertical data
+uploaded_vertical_files = st.sidebar.file_uploader(
+    "Upload Vertical Data Files",
+    type=["nc", "netcdf"],
+    accept_multiple_files=True,
+    key="vertical_uploader",
+)
+
+# Optional: Upload config.yaml
+uploaded_config = st.sidebar.file_uploader(
+    "Upload config.yaml",
+    type=["yaml", "yml"],
+    key="config_uploader",
+)
+
+# Optional: Upload model weights
+uploaded_weights = st.sidebar.file_uploader(
+    "Upload Model Weights (.pt)",
+    type=["pt"],
+    key="weights_uploader",
+)
+
+# Other configurations
+st.sidebar.header("Task Configuration")
+
+lead_times = st.sidebar.multiselect(
+    "Select Lead Times",
+    options=[12, 24, 36, 48],
+    default=[12],
+)
+
+input_times = st.sidebar.multiselect(
+    "Select Input Times",
+    options=[-6, -12, -18, -24],
+    default=[-6],
+)
+
+time_range_start = st.sidebar.text_input(
+    "Start Time (e.g., 2020-01-01T00:00:00)",
+    value="2020-01-01T00:00:00",
+)
+
+time_range_end = st.sidebar.text_input(
+    "End Time (e.g., 2020-01-01T23:59:59)",
+    value="2020-01-01T23:59:59",
+)
+
+time_range = (time_range_start, time_range_end)
+
+# Function to save uploaded files
+def save_uploaded_files(uploaded_files, folder_name, max_size_mb=1024):
+    if not uploaded_files:
+        st.warning(f"No {folder_name} files uploaded.")
+        return None
+    # Validate file sizes
+    for file in uploaded_files:
+        if file.size > max_size_mb * 1024 * 1024:
+            st.error(f"File {file.name} exceeds the maximum size of {max_size_mb} MB.")
+            return None
+    temp_dir = tempfile.mkdtemp()
+    with st.spinner(f"Saving {folder_name} files..."):
+        for uploaded_file in uploaded_files:
+            file_path = Path(temp_dir) / uploaded_file.name
+            with open(file_path, "wb") as f:
+                f.write(uploaded_file.getbuffer())
+    st.success(f"Saved {len(uploaded_files)} {folder_name} files.")
+    return Path(temp_dir)
+
+# Save uploaded files
+surf_dir = save_uploaded_files(uploaded_surface_files, "surface")
+vert_dir = save_uploaded_files(uploaded_vertical_files, "vertical")
+
+# Display uploaded files
+if surf_dir:
+    st.sidebar.subheader("Surface Files Uploaded:")
+    for file in surf_dir.iterdir():
+        st.sidebar.write(file.name)
+
+if vert_dir:
+    st.sidebar.subheader("Vertical Files Uploaded:")
+    for file in vert_dir.iterdir():
+        st.sidebar.write(file.name)
+
+# Handle Climatology Files
+st.sidebar.header("Upload Climatology Files (If Missing)")
+
+# Climatology files paths
+default_clim_dir = Path("Prithvi-WxC/examples/climatology")
+surf_in_scal_path = default_clim_dir / "musigma_surface.nc"
+vert_in_scal_path = default_clim_dir / "musigma_vertical.nc"
+surf_out_scal_path = default_clim_dir / "anomaly_variance_surface.nc"
+vert_out_scal_path = default_clim_dir / "anomaly_variance_vertical.nc"
+
+# Check if climatology files exist
+clim_files_exist = all(
+    [
+        surf_in_scal_path.exists(),
+        vert_in_scal_path.exists(),
+        surf_out_scal_path.exists(),
+        vert_out_scal_path.exists(),
+    ]
+)
+
+if not clim_files_exist:
+    st.sidebar.warning("Climatology files are missing.")
+    uploaded_clim_surface = st.sidebar.file_uploader(
+        "Upload Climatology Surface File",
+        type=["nc", "netcdf"],
+        key="clim_surface_uploader",
+    )
+    uploaded_clim_vertical = st.sidebar.file_uploader(
+        "Upload Climatology Vertical File",
+        type=["nc", "netcdf"],
+        key="clim_vertical_uploader",
+    )
+    
+    if uploaded_clim_surface and uploaded_clim_vertical:
+        clim_temp_dir = tempfile.mkdtemp()
+        clim_surf_path = Path(clim_temp_dir) / uploaded_clim_surface.name
+        with open(clim_surf_path, "wb") as f:
+            f.write(uploaded_clim_surface.getbuffer())
+        clim_vert_path = Path(clim_temp_dir) / uploaded_clim_vertical.name
+        with open(clim_vert_path, "wb") as f:
+            f.write(uploaded_clim_vertical.getbuffer())
+        st.success("Climatology files uploaded and saved.")
+    else:
+        if not (uploaded_clim_surface and uploaded_clim_vertical):
+            st.warning("Please upload both climatology surface and vertical files.")
+else:
+    clim_surf_path = surf_in_scal_path
+    clim_vert_path = vert_in_scal_path
+
+# Save uploaded config.yaml
+if uploaded_config:
+    temp_config = tempfile.mktemp(suffix=".yaml")
+    with open(temp_config, "wb") as f:
+        f.write(uploaded_config.getbuffer())
+    config_path = Path(temp_config)
+    st.sidebar.success("Config.yaml uploaded and saved.")
+else:
+    # Use default config.yaml path
+    config_path = Path("Prithvi-WxC/examples/config.yaml")
+    if not config_path.exists():
+        st.sidebar.error("Default config.yaml not found. Please upload a config file.")
+        st.stop()
+
+# Save uploaded model weights
+if uploaded_weights:
+    temp_weights = tempfile.mktemp(suffix=".pt")
+    with open(temp_weights, "wb") as f:
+        f.write(uploaded_weights.getbuffer())
+    weights_path = Path(temp_weights)
+    st.sidebar.success("Model weights uploaded and saved.")
+else:
+    # Use default weights path
+    weights_path = Path("Prithvi-WxC/examples/weights/prithvi.wxc.2300m.v1.pt")
+    if not weights_path.exists():
+        st.sidebar.error("Default model weights not found. Please upload model weights.")
+        st.stop()
+
+# Button to run inference
+if st.sidebar.button("Run Inference"):
+
+    # Initialize device
+    torch.jit.enable_onednn_fusion(True)
+    if torch.cuda.is_available():
+        device = torch.device("cuda")
+        st.write(f"Using device: {torch.cuda.get_device_name()}")
+        torch.backends.cudnn.benchmark = True
+        torch.backends.cudnn.deterministic = True
+    else:
+        device = torch.device("cpu")
+        st.write("Using device: CPU")
+
+    # Set random seeds
+    random.seed(42)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed(42)
+    torch.manual_seed(42)
+    np.random.seed(42)
+
+    # Define variables and parameters
+    surface_vars = [
+        "EFLUX",
+        "GWETROOT",
+        "HFLUX",
+        "LAI",
+        "LWGAB",
+        "LWGEM",
+        "LWTUP",
+        "PS",
+        "QV2M",
+        "SLP",
+        "SWGNT",
+        "SWTNT",
+        "T2M",
+        "TQI",
+        "TQL",
+        "TQV",
+        "TS",
+        "U10M",
+        "V10M",
+        "Z0M",
+    ]
+    static_surface_vars = ["FRACI", "FRLAND", "FROCEAN", "PHIS"]
+    vertical_vars = ["CLOUD", "H", "OMEGA", "PL", "QI", "QL", "QV", "T", "U", "V"]
+    levels = [
+        34.0,
+        39.0,
+        41.0,
+        43.0,
+        44.0,
+        45.0,
+        48.0,
+        51.0,
+        53.0,
+        56.0,
+        63.0,
+        68.0,
+        71.0,
+        72.0,
+    ]
+    padding = {"level": [0, 0], "lat": [0, -1], "lon": [0, 0]}
+
+    residual = "climate"
+    masking_mode = "local"
+    decoder_shifting = True
+    masking_ratio = 0.99
+
+    positional_encoding = "fourier"
+
+    # Initialize Dataset
+    try:
+        with st.spinner("Initializing dataset..."):
+            # Validate climatology files
+            if not clim_files_exist and not (uploaded_clim_surface and uploaded_clim_vertical):
+                st.error("Climatology files are missing. Please upload both surface and vertical climatology files.")
+                st.stop()
+
+            dataset = Merra2Dataset(
+                time_range=time_range,
+                lead_times=lead_times,
+                input_times=input_times,
+                data_path_surface=Path("Prithvi-WxC/examples/merra-2"),
+                data_path_vertical=Path("Prithvi-WxC/examples/merra-2"),
+                climatology_path_surface=Path("Prithvi-WxC/examples/climatology"),
+                climatology_path_vertical=Path("Prithvi-WxC/examples/climatology"),
+                surface_vars=surface_vars,
+                static_surface_vars=static_surface_vars,
+                vertical_vars=vertical_vars,
+                levels=levels,
+                positional_encoding=positional_encoding,
+            )
+            assert len(dataset) > 0, "There doesn't seem to be any valid data."
+        st.success("Dataset initialized successfully.")
+    except Exception as e:
+        st.error("Error initializing dataset:")
+        st.error(traceback.format_exc())
+        st.stop()
+
+    # Load scalers
+    try:
+        with st.spinner("Loading scalers..."):
+            # Assuming the scaler paths are the same as climatology paths
+            surf_in_scal_path = clim_surf_path
+            vert_in_scal_path = clim_vert_path
+            surf_out_scal_path = Path(clim_surf_path.parent) / "anomaly_variance_surface.nc"
+            vert_out_scal_path = Path(clim_vert_path.parent) / "anomaly_variance_vertical.nc"
+
+            # Check if output scaler files exist
+            if not surf_out_scal_path.exists() or not vert_out_scal_path.exists():
+                st.error("Anomaly variance scaler files are missing.")
+                st.stop()
+
+            in_mu, in_sig = input_scalers(
+                surface_vars,
+                vertical_vars,
+                levels,
+                surf_in_scal_path,
+                vert_in_scal_path,
+            )
+
+            output_sig = output_scalers(
+                surface_vars,
+                vertical_vars,
+                levels,
+                surf_out_scal_path,
+                vert_out_scal_path,
+            )
+
+            static_mu, static_sig = static_input_scalers(
+                surf_in_scal_path,
+                static_surface_vars,
+            )
+        st.success("Scalers loaded successfully.")
+    except Exception as e:
+        st.error("Error loading scalers:")
+        st.error(traceback.format_exc())
+        st.stop()
+
+    # Load configuration
+    try:
+        with st.spinner("Loading configuration..."):
+            with open(config_path, "r") as f:
+                config = yaml.safe_load(f)
+            # Validate config
+            required_params = [
+                "in_channels", "input_size_time", "in_channels_static",
+                "input_scalers_epsilon", "static_input_scalers_epsilon",
+                "n_lats_px", "n_lons_px", "patch_size_px",
+                "mask_unit_size_px", "embed_dim", "n_blocks_encoder",
+                "n_blocks_decoder", "mlp_multiplier", "n_heads",
+                "dropout", "drop_path", "parameter_dropout"
+            ]
+            missing_params = [param for param in required_params if param not in config.get("params", {})]
+            if missing_params:
+                st.error(f"Missing configuration parameters: {missing_params}")
+                st.stop()
+        st.success("Configuration loaded successfully.")
+    except Exception as e:
+        st.error("Error loading configuration:")
+        st.error(traceback.format_exc())
+        st.stop()
+
+    # Initialize the model
+    try:
+        with st.spinner("Initializing model..."):
+            model = PrithviWxC(
+                in_channels=config["params"]["in_channels"],
+                input_size_time=config["params"]["input_size_time"],
+                in_channels_static=config["params"]["in_channels_static"],
+                input_scalers_mu=in_mu,
+                input_scalers_sigma=in_sig,
+                input_scalers_epsilon=config["params"]["input_scalers_epsilon"],
+                static_input_scalers_mu=static_mu,
+                static_input_scalers_sigma=static_sig,
+                static_input_scalers_epsilon=config["params"]["static_input_scalers_epsilon"],
+                output_scalers=output_sig**0.5,
+                n_lats_px=config["params"]["n_lats_px"],
+                n_lons_px=config["params"]["n_lons_px"],
+                patch_size_px=config["params"]["patch_size_px"],
+                mask_unit_size_px=config["params"]["mask_unit_size_px"],
+                mask_ratio_inputs=masking_ratio,
+                embed_dim=config["params"]["embed_dim"],
+                n_blocks_encoder=config["params"]["n_blocks_encoder"],
+                n_blocks_decoder=config["params"]["n_blocks_decoder"],
+                mlp_multiplier=config["params"]["mlp_multiplier"],
+                n_heads=config["params"]["n_heads"],
+                dropout=config["params"]["dropout"],
+                drop_path=config["params"]["drop_path"],
+                parameter_dropout=config["params"]["parameter_dropout"],
+                residual=residual,
+                masking_mode=masking_mode,
+                decoder_shifting=decoder_shifting,
+                positional_encoding=positional_encoding,
+                checkpoint_encoder=[],
+                checkpoint_decoder=[],
+            )
+        st.success("Model initialized successfully.")
+    except Exception as e:
+        st.error("Error initializing model:")
+        st.error(traceback.format_exc())
+        st.stop()
+
+    # Load model weights
+    try:
+        with st.spinner("Loading model weights..."):
+            state_dict = torch.load(weights_path, map_location=device)
+            if "model_state" in state_dict:
+                state_dict = state_dict["model_state"]
+            model.load_state_dict(state_dict, strict=True)
+            model.to(device)
+        st.success("Model weights loaded successfully.")
+    except Exception as e:
+        st.error("Error loading model weights:")
+        st.error(traceback.format_exc())
+        st.stop()
+
+    # Prepare data batch
+    try:
+        with st.spinner("Preparing data batch..."):
+            data = next(iter(dataset))
+            batch = preproc([data], padding)
+
+            for k, v in batch.items():
+                if isinstance(v, torch.Tensor):
+                    batch[k] = v.to(device)
+        st.success("Data batch prepared successfully.")
+    except Exception as e:
+        st.error("Error preparing data batch:")
+        st.error(traceback.format_exc())
+        st.stop()
+
+    # Run inference
+    try:
+        with st.spinner("Running model inference..."):
+            rng_state_1 = torch.get_rng_state()
+            with torch.no_grad():
+                model.eval()
+                out = model(batch)
+        st.success("Model inference completed successfully.")
+    except Exception as e:
+        st.error("Error during model inference:")
+        st.error(traceback.format_exc())
+        st.stop()
+
+    # Display output
+    st.header("Inference Results")
+    st.write(out)  # Adjust based on the structure of 'out'
+
+    # Optionally, provide download links or visualizations
+    # For example, if 'out' contains tensors or dataframes:
+    # st.write("Output Tensor:", out["some_key"].cpu().numpy())
+
+else:
+    st.info("Please upload the necessary files and click 'Run Inference' to start.")

app2.py

@@ -0,0 +1,959 @@
+import streamlit as st
+import torch
+import random
+import numpy as np
+import yaml
+from pathlib import Path
+import tempfile
+import traceback
+import matplotlib.pyplot as plt
+import plotly.graph_objects as go
+from Prithvi import *  # Ensure this import includes your model and dataset classes
+import xarray as xr
+from aurora import Batch, Metadata
+from aurora import Aurora, rollout
+import logging
+import matplotlib.pyplot as plt
+import numpy as np
+import cartopy.crs as ccrs
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# Function to save uploaded files to temporary files and store paths in session_state
+def save_uploaded_files(uploaded_files):
+    if 'temp_file_paths' not in st.session_state:
+        st.session_state.temp_file_paths = []
+        for uploaded_file in uploaded_files:
+            suffix = os.path.splitext(uploaded_file.name)[1]
+            temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=suffix)
+            temp_file.write(uploaded_file.read())
+            temp_file.close()
+            st.session_state.temp_file_paths.append(temp_file.name)
+# Cached function to load dataset
+@st.cache_resource
+def load_dataset(file_paths):
+    try:
+        ds = xr.open_mfdataset(file_paths, combine='by_coords').load()
+        return ds
+    except Exception as e:
+        st.error("Error loading dataset:")
+        st.error(traceback.format_exc())
+        return None
+
+# Set page configuration
+st.set_page_config(
+    page_title="Weather Data Processor",
+    layout="wide",
+    initial_sidebar_state="expanded",
+)
+
+
+
+# Create a header with two columns: one for the title and one for the model selector
+header_col1, header_col2 = st.columns([4, 1])  # Adjust the ratio as needed
+
+with header_col1:
+    st.title("🌦️ Weather & Climate Data Processor and Forecaster")
+
+with header_col2:
+    st.markdown("### Select a Model")
+    selected_model = st.selectbox(
+        "",
+        options=["Aurora", "Climax", "Prithvi", "LSTM"],
+        index=0,
+        key="model_selector",
+        help="Select the model you want to use for processing the data."
+    )
+
+st.write("---")  # Horizontal separator
+
+# --- Layout: Two Columns ---
+left_col, right_col = st.columns([1, 2])  # Adjust column ratios as needed
+
+with left_col:
+    st.header("🔧 Configuration")
+
+    # --- Dynamic Configuration Based on Selected Model ---
+    def get_model_configuration(model_name):
+        if model_name == "Prithvi":
+            st.subheader("Prithvi Model Configuration")
+
+            # Prithvi-specific configuration inputs
+            param1 = st.number_input("Prithvi Parameter 1", value=10, step=1)
+            param2 = st.text_input("Prithvi Parameter 2", value="default_prithvi")
+            # Add other Prithvi-specific parameters here
+
+            config = {
+                "param1": param1,
+                "param2": param2,
+                # Include other parameters as needed
+            }
+
+            # --- Prithvi-Specific File Uploads ---
+            st.markdown("### Upload Data Files for Prithvi Model")
+
+            # File uploader for surface data
+            uploaded_surface_files = st.file_uploader(
+                "Upload Surface Data Files",
+                type=["nc", "netcdf"],
+                accept_multiple_files=True,
+                key="surface_uploader",
+            )
+
+            # File uploader for vertical data
+            uploaded_vertical_files = st.file_uploader(
+                "Upload Vertical Data Files",
+                type=["nc", "netcdf"],
+                accept_multiple_files=True,
+                key="vertical_uploader",
+            )
+
+            # Handle Climatology Files
+            st.markdown("### Upload Climatology Files (If Missing)")
+
+            # Climatology files paths
+            default_clim_dir = Path("Prithvi-WxC/examples/climatology")
+            surf_in_scal_path = default_clim_dir / "musigma_surface.nc"
+            vert_in_scal_path = default_clim_dir / "musigma_vertical.nc"
+            surf_out_scal_path = default_clim_dir / "anomaly_variance_surface.nc"
+            vert_out_scal_path = default_clim_dir / "anomaly_variance_vertical.nc"
+
+            # Check if climatology files exist
+            clim_files_exist = all(
+                [
+                    surf_in_scal_path.exists(),
+                    vert_in_scal_path.exists(),
+                    surf_out_scal_path.exists(),
+                    vert_out_scal_path.exists(),
+                ]
+            )
+
+            if not clim_files_exist:
+                st.warning("Climatology files are missing.")
+                uploaded_clim_surface = st.file_uploader(
+                    "Upload Climatology Surface File",
+                    type=["nc", "netcdf"],
+                    key="clim_surface_uploader",
+                )
+                uploaded_clim_vertical = st.file_uploader(
+                    "Upload Climatology Vertical File",
+                    type=["nc", "netcdf"],
+                    key="clim_vertical_uploader",
+                )
+
+                # Process uploaded climatology files
+                if uploaded_clim_surface and uploaded_clim_vertical:
+                    clim_temp_dir = tempfile.mkdtemp()
+                    clim_surf_path = Path(clim_temp_dir) / uploaded_clim_surface.name
+                    with open(clim_surf_path, "wb") as f:
+                        f.write(uploaded_clim_surface.getbuffer())
+                    clim_vert_path = Path(clim_temp_dir) / uploaded_clim_vertical.name
+                    with open(clim_vert_path, "wb") as f:
+                        f.write(uploaded_clim_vertical.getbuffer())
+                    st.success("Climatology files uploaded and saved.")
+                else:
+                    st.warning("Please upload both climatology surface and vertical files.")
+            else:
+                clim_surf_path = surf_in_scal_path
+                clim_vert_path = vert_in_scal_path
+
+            # Optional: Upload config.yaml
+            uploaded_config = st.file_uploader(
+                "Upload config.yaml",
+                type=["yaml", "yml"],
+                key="config_uploader",
+            )
+
+            if uploaded_config:
+                temp_config = tempfile.mktemp(suffix=".yaml")
+                with open(temp_config, "wb") as f:
+                    f.write(uploaded_config.getbuffer())
+                config_path = Path(temp_config)
+                st.success("Config.yaml uploaded and saved.")
+            else:
+                # Use default config.yaml path
+                config_path = Path("Prithvi-WxC/examples/config.yaml")
+                if not config_path.exists():
+                    st.error("Default config.yaml not found. Please upload a config file.")
+                    st.stop()
+
+            # Optional: Upload model weights
+            uploaded_weights = st.file_uploader(
+                "Upload Model Weights (.pt)",
+                type=["pt"],
+                key="weights_uploader",
+            )
+
+            if uploaded_weights:
+                temp_weights = tempfile.mktemp(suffix=".pt")
+                with open(temp_weights, "wb") as f:
+                    f.write(uploaded_weights.getbuffer())
+                weights_path = Path(temp_weights)
+                st.success("Model weights uploaded and saved.")
+            else:
+                # Use default weights path
+                weights_path = Path("Prithvi-WxC/examples/weights/prithvi.wxc.2300m.v1.pt")
+                if not weights_path.exists():
+                    st.error("Default model weights not found. Please upload model weights.")
+                    st.stop()
+
+            return config, uploaded_surface_files, uploaded_vertical_files, clim_surf_path, clim_vert_path, config_path, weights_path
+
+        else:
+            # For other models, provide a simple file uploader
+            st.subheader(f"{model_name} Model Data Upload")
+            st.markdown("### Drag and Drop Your Data Files Here")
+            uploaded_files = st.file_uploader(
+                f"Upload Data Files for {model_name}",
+                accept_multiple_files=True,
+                key=f"{model_name.lower()}_uploader",
+                type=["nc", "netcdf", "nc4"],
+            )
+            return uploaded_files
+
+    # Retrieve model-specific configuration and files
+    if selected_model == "Prithvi":
+        config, uploaded_surface_files, uploaded_vertical_files, clim_surf_path, clim_vert_path, config_path, weights_path = get_model_configuration(selected_model)
+    else:
+        uploaded_files = get_model_configuration(selected_model)
+
+    st.write("---")  # Horizontal separator
+
+    # --- Run Inference Button ---
+    if st.button("🚀 Run Inference"):
+        with right_col:
+            st.header("📈 Inference Progress & Visualization")
+
+            # Initialize device
+            try:
+                torch.jit.enable_onednn_fusion(True)
+                if torch.cuda.is_available():
+                    device = torch.device("cuda")
+                    st.write(f"Using device: **{torch.cuda.get_device_name()}**")
+                    torch.backends.cudnn.benchmark = True
+                    torch.backends.cudnn.deterministic = True
+                else:
+                    device = torch.device("cpu")
+                    st.write("Using device: **CPU**")
+            except Exception as e:
+                st.error("Error initializing device:")
+                st.error(traceback.format_exc())
+                st.stop()
+
+            # Set random seeds
+            try:
+                random.seed(42)
+                if torch.cuda.is_available():
+                    torch.cuda.manual_seed(42)
+                torch.manual_seed(42)
+                np.random.seed(42)
+            except Exception as e:
+                st.error("Error setting random seeds:")
+                st.error(traceback.format_exc())
+                st.stop()
+
+            # # Define variables and parameters based on dataset type
+            # if dataset_type == "MERRA2":
+            #     surface_vars = [
+            #         "EFLUX",
+            #         "GWETROOT",
+            #         "HFLUX",
+            #         "LAI",
+            #         "LWGAB",
+            #         "LWGEM",
+            #         "LWTUP",
+            #         "PS",
+            #         "QV2M",
+            #         "SLP",
+            #         "SWGNT",
+            #         "SWTNT",
+            #         "T2M",
+            #         "TQI",
+            #         "TQL",
+            #         "TQV",
+            #         "TS",
+            #         "U10M",
+            #         "V10M",
+            #         "Z0M",
+            #     ]
+            #     static_surface_vars = ["FRACI", "FRLAND", "FROCEAN", "PHIS"]
+            #     vertical_vars = ["CLOUD", "H", "OMEGA", "PL", "QI", "QL", "QV", "T", "U", "V"]
+            #     levels = [
+            #         34.0,
+            #         39.0,
+            #         41.0,
+            #         43.0,
+            #         44.0,
+            #         45.0,
+            #         48.0,
+            #         51.0,
+            #         53.0,
+            #         56.0,
+            #         63.0,
+            #         68.0,
+            #         71.0,
+            #         72.0,
+            #     ]
+            # elif dataset_type == "GEOS5":
+            #     # Define GEOS5 specific variables
+            #     surface_vars = [
+            #         "GEOS5_EFLUX",
+            #         "GEOS5_GWETROOT",
+            #         "GEOS5_HFLUX",
+            #         "GEOS5_LAI",
+            #         "GEOS5_LWGAB",
+            #         "GEOS5_LWGEM",
+            #         "GEOS5_LWTUP",
+            #         "GEOS5_PS",
+            #         "GEOS5_QV2M",
+            #         "GEOS5_SLP",
+            #         "GEOS5_SWGNT",
+            #         "GEOS5_SWTNT",
+            #         "GEOS5_T2M",
+            #         "GEOS5_TQI",
+            #         "GEOS5_TQL",
+            #         "GEOS5_TQV",
+            #         "GEOS5_TS",
+            #         "GEOS5_U10M",
+            #         "GEOS5_V10M",
+            #         "GEOS5_Z0M",
+            #     ]
+            #     static_surface_vars = ["GEOS5_FRACI", "GEOS5_FRLAND", "GEOS5_FROCEAN", "GEOS5_PHIS"]
+            #     vertical_vars = ["GEOS5_CLOUD", "GEOS5_H", "GEOS5_OMEGA", "GEOS5_PL", "GEOS5_QI", "GEOS5_QL", "GEOS5_QV", "GEOS5_T", "GEOS5_U", "GEOS5_V"]
+            #     levels = [
+            #         # Define levels specific to GEOS5 if different
+            #         10.0,
+            #         20.0,
+            #         30.0,
+            #         40.0,
+            #         50.0,
+            #         60.0,
+            #         70.0,
+            #         80.0,
+            #     ]
+            # else:
+            #     st.error("Unsupported dataset type selected.")
+            #     st.stop()
+
+            padding = {"level": [0, 0], "lat": [0, -1], "lon": [0, 0]}
+
+            residual = "climate"
+            masking_mode = "local"
+            decoder_shifting = True
+            masking_ratio = 0.99
+
+            positional_encoding = "fourier"
+
+            # --- Initialize Dataset ---
+            try:
+                with st.spinner("Initializing dataset..."):
+                    if selected_model == "Prithvi":
+                        pass
+                        # # Validate climatology files
+                        # if not clim_files_exist and not (uploaded_clim_surface and uploaded_clim_vertical):
+                        #     st.error("Climatology files are missing. Please upload both climatology surface and vertical files.")
+                        #     st.stop()
+
+                        # dataset = Merra2Dataset(
+                        #     time_range=time_range,
+                        #     lead_times=lead_times,
+                        #     input_times=input_times,
+                        #     data_path_surface=surf_dir,
+                        #     data_path_vertical=vert_dir,
+                        #     climatology_path_surface=clim_surf_path,
+                        #     climatology_path_vertical=clim_vert_path,
+                        #     surface_vars=surface_vars,
+                        #     static_surface_vars=static_surface_vars,
+                        #     vertical_vars=vertical_vars,
+                        #     levels=levels,
+                        #     positional_encoding=positional_encoding,
+                        # )
+                        # assert len(dataset) > 0, "There doesn't seem to be any valid data."
+                    elif selected_model == "Aurora":
+                        # TODO just temporary, replace this
+                        if uploaded_files:
+                            temp_file_paths = []  # List to store paths of temporary files
+                            try:
+                                # Save each uploaded file to a temporary file
+                                save_uploaded_files(uploaded_files)
+                                ds = load_dataset(st.session_state.temp_file_paths)
+
+                                # Now, use xarray to open the multiple files
+                                if ds:
+                                    st.success("Files successfully loaded!")
+                                    st.session_state.ds_subset = ds
+
+                                    
+                                    # print(ds)
+                                    ds = ds.fillna(ds.mean())
+
+                                    desired_levels = [50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000]
+
+                                    # Ensure that the 'lev' dimension exists
+                                    if 'lev' not in ds.dims:
+                                        raise ValueError("The dataset does not contain a 'lev' (pressure level) dimension.")
+
+                                    # Define the _prepare function
+                                    def _prepare(x: np.ndarray, i: int) -> torch.Tensor:
+                                        # Select previous and current time steps
+                                        selected = x[[i - 6, i]]
+                                        
+                                        # Add a batch dimension
+                                        selected = selected[None]
+                                        
+                                        # Ensure data is contiguous
+                                        selected = selected.copy()
+                                        
+                                        # Convert to PyTorch tensor
+                                        return torch.from_numpy(selected)
+
+                                    # Adjust latitudes and longitudes
+                                    lat = ds.lat.values * -1
+                                    lon = ds.lon.values + 180
+
+                                    # Subset the dataset to only include the desired pressure levels
+                                    ds_subset = ds.sel(lev=desired_levels, method="nearest")
+
+                                    # Verify that all desired levels are present
+                                    present_levels = ds_subset.lev.values
+                                    missing_levels = set(desired_levels) - set(present_levels)
+                                    if missing_levels:
+                                        raise ValueError(f"The following desired pressure levels are missing in the dataset: {missing_levels}")
+
+                                    # Extract pressure levels after subsetting
+                                    lev = ds_subset.lev.values  # Pressure levels in hPa
+
+                                    # Prepare surface variables at 1000 hPa
+                                    try:
+                                        lev_index_1000 = np.where(lev == 1000)[0][0]
+                                    except IndexError:
+                                        raise ValueError("1000 hPa level not found in the 'lev' dimension after subsetting.")
+
+                                    T_surface = ds_subset.T.isel(lev=lev_index_1000).compute()
+                                    U_surface = ds_subset.U.isel(lev=lev_index_1000).compute()
+                                    V_surface = ds_subset.V.isel(lev=lev_index_1000).compute()
+                                    SLP = ds_subset.SLP.compute()
+
+                                    # Reorder static variables (selecting the first time index to remove the time dimension)
+                                    PHIS = ds_subset.PHIS.isel(time=0).compute()
+
+                                    # Prepare atmospheric variables for the desired pressure levels excluding 1000 hPa
+                                    atmos_levels = [int(level) for level in lev if level != 1000]
+
+                                    T_atm = (ds_subset.T.sel(lev=atmos_levels)).compute()
+                                    U_atm = (ds_subset.U.sel(lev=atmos_levels)).compute()
+                                    V_atm = (ds_subset.V.sel(lev=atmos_levels)).compute()
+
+                                    # Select time index
+                                    num_times = ds_subset.time.size
+                                    i = 6  # Adjust as needed (1 <= i < num_times)
+
+                                    if i >= num_times or i < 1:
+                                        raise IndexError("Time index i is out of bounds.")
+
+                                    time_values = ds_subset.time.values
+                                    current_time = np.datetime64(time_values[i]).astype('datetime64[s]').astype(datetime)
+
+                                    # Prepare surface variables
+                                    surf_vars = {
+                                        "2t": _prepare(T_surface.values, i),   # Two-meter temperature
+                                        "10u": _prepare(U_surface.values, i),  # Ten-meter eastward wind
+                                        "10v": _prepare(V_surface.values, i),  # Ten-meter northward wind
+                                        "msl": _prepare(SLP.values, i),        # Mean sea-level pressure
+                                    }
+
+                                    # Prepare static variables (now 2D tensors)
+                                    static_vars = {
+                                        "z": torch.from_numpy(PHIS.values.copy()),  # Geopotential (h, w)
+                                        # Add 'lsm' and 'slt' if available and needed
+                                    }
+
+                                    # Prepare atmospheric variables
+                                    atmos_vars = {
+                                        "t": _prepare(T_atm.values, i),  # Temperature at desired levels
+                                        "u": _prepare(U_atm.values, i),  # Eastward wind at desired levels
+                                        "v": _prepare(V_atm.values, i),  # Southward wind at desired levels
+                                    }
+
+                                    # Define metadata
+                                    metadata = Metadata(
+                                        lat=torch.from_numpy(lat.copy()),
+                                        lon=torch.from_numpy(lon.copy()),
+                                        time=(current_time,),
+                                        atmos_levels=tuple(atmos_levels),  # Only the desired atmospheric levels
+                                    )
+
+                                    # Create the Batch object
+                                    batch = Batch(
+                                        surf_vars=surf_vars,
+                                        static_vars=static_vars,
+                                        atmos_vars=atmos_vars,
+                                        metadata=metadata
+                                    ) # Display the dataset or perform further processing
+
+                                    st.session_state['batch'] = batch
+
+                            except Exception as e:
+                                st.error(f"An error occurred: {e}")
+
+                            # finally:
+                            #     # Clean up: Remove temporary files
+                            #     for path in temp_file_paths:
+                            #         try:
+                            #             os.remove(path)
+                            #         except Exception as e:
+                            #             st.warning(f"Could not delete temp file {path}: {e}")
+                    else:
+                        # For other models, implement their specific dataset initialization
+                        # Placeholder: Replace with actual dataset initialization for other models
+                        dataset = None  # Replace with actual dataset
+                        st.warning("Dataset initialization for this model is not implemented yet.")
+                        st.stop()
+                st.success("Dataset initialized successfully.")
+            except Exception as e:
+                st.error("Error initializing dataset:")
+                st.error(traceback.format_exc())
+                st.stop()
+
+            # --- Load Scalers ---
+            try:
+                with st.spinner("Loading scalers..."):
+                    if selected_model == "Prithvi":
+                        pass
+                        # # Assuming the scaler paths are the same as climatology paths
+                        # surf_in_scal_path = clim_surf_path
+                        # vert_in_scal_path = clim_vert_path
+                        # surf_out_scal_path = Path(clim_surf_path.parent) / "anomaly_variance_surface.nc"
+                        # vert_out_scal_path = Path(clim_vert_path.parent) / "anomaly_variance_vertical.nc"
+
+                        # # Check if output scaler files exist
+                        # if not surf_out_scal_path.exists() or not vert_out_scal_path.exists():
+                        #     st.error("Anomaly variance scaler files are missing.")
+                        #     st.stop()
+
+                        # in_mu, in_sig = input_scalers(
+                        #     surface_vars,
+                        #     vertical_vars,
+                        #     levels,
+                        #     surf_in_scal_path,
+                        #     vert_in_scal_path,
+                        # )
+
+                        # output_sig = output_scalers(
+                        #     surface_vars,
+                        #     vertical_vars,
+                        #     levels,
+                        #     surf_out_scal_path,
+                        #     vert_out_scal_path,
+                        # )
+
+                        # static_mu, static_sig = static_input_scalers(
+                        #     surf_in_scal_path,
+                        #     static_surface_vars,
+                        # )
+                    else:
+                        # Load scalers for other models if applicable
+                        # Placeholder: Replace with actual scaler loading for other models
+                        in_mu, in_sig = None, None
+                        output_sig = None
+                        static_mu, static_sig = None, None
+                st.success("Scalers loaded successfully.")
+            except Exception as e:
+                st.error("Error loading scalers:")
+                st.error(traceback.format_exc())
+                st.stop()
+
+            # --- Load Configuration ---
+            try:
+                with st.spinner("Loading configuration..."):
+                    if selected_model == "Prithvi":
+                        with open(config_path, "r") as f:
+                            config = yaml.safe_load(f)
+                        # Validate config
+                        required_params = [
+                            "in_channels", "input_size_time", "in_channels_static",
+                            "input_scalers_epsilon", "static_input_scalers_epsilon",
+                            "n_lats_px", "n_lons_px", "patch_size_px",
+                            "mask_unit_size_px", "embed_dim", "n_blocks_encoder",
+                            "n_blocks_decoder", "mlp_multiplier", "n_heads",
+                            "dropout", "drop_path", "parameter_dropout"
+                        ]
+                        missing_params = [param for param in required_params if param not in config.get("params", {})]
+                        if missing_params:
+                            st.error(f"Missing configuration parameters: {missing_params}")
+                            st.stop()
+                    else:
+                        # Load configuration for other models if applicable
+                        # Placeholder: Replace with actual configuration loading for other models
+                        config = {}
+                st.success("Configuration loaded successfully.")
+            except Exception as e:
+                st.error("Error loading configuration:")
+                st.error(traceback.format_exc())
+                st.stop()
+
+            # --- Initialize the Model ---
+            try:
+                with st.spinner("Initializing model..."):
+                    if selected_model == "Prithvi":
+                        model = PrithviWxC(
+                            in_channels=config["params"]["in_channels"],
+                            input_size_time=config["params"]["input_size_time"],
+                            in_channels_static=config["params"]["in_channels_static"],
+                            input_scalers_mu=in_mu,
+                            input_scalers_sigma=in_sig,
+                            input_scalers_epsilon=config["params"]["input_scalers_epsilon"],
+                            static_input_scalers_mu=static_mu,
+                            static_input_scalers_sigma=static_sig,
+                            static_input_scalers_epsilon=config["params"]["static_input_scalers_epsilon"],
+                            output_scalers=output_sig**0.5,
+                            n_lats_px=config["params"]["n_lats_px"],
+                            n_lons_px=config["params"]["n_lons_px"],
+                            patch_size_px=config["params"]["patch_size_px"],
+                            mask_unit_size_px=config["params"]["mask_unit_size_px"],
+                            mask_ratio_inputs=masking_ratio,
+                            embed_dim=config["params"]["embed_dim"],
+                            n_blocks_encoder=config["params"]["n_blocks_encoder"],
+                            n_blocks_decoder=config["params"]["n_blocks_decoder"],
+                            mlp_multiplier=config["params"]["mlp_multiplier"],
+                            n_heads=config["params"]["n_heads"],
+                            dropout=config["params"]["dropout"],
+                            drop_path=config["params"]["drop_path"],
+                            parameter_dropout=config["params"]["parameter_dropout"],
+                            residual=residual,
+                            masking_mode=masking_mode,
+                            decoder_shifting=decoder_shifting,
+                            positional_encoding=positional_encoding,
+                            checkpoint_encoder=[],
+                            checkpoint_decoder=[],
+                        )
+                    elif selected_model == "Aurora":
+                        pass
+                        
+                    else:
+                        
+                        # Initialize other models here
+                        # Placeholder: Replace with actual model initialization for other models
+                        model = None
+                        st.warning("Model initialization for this model is not implemented yet.")
+                        st.stop()
+                    # model.to(device)
+                st.success("Model initialized successfully.")
+            except Exception as e:
+                st.error("Error initializing model:")
+                st.error(traceback.format_exc())
+                st.stop()
+
+            # --- Load Model Weights ---
+            try:
+                with st.spinner("Loading model weights..."):
+                    if selected_model == "Prithvi":
+                        state_dict = torch.load(weights_path, map_location=device)
+                        if "model_state" in state_dict:
+                            state_dict = state_dict["model_state"]
+                        model.load_state_dict(state_dict, strict=True)
+                        model.to(device)
+                    else:
+                        # Load weights for other models if applicable
+                        # Placeholder: Replace with actual weight loading for other models
+                        pass
+                st.success("Model weights loaded successfully.")
+            except Exception as e:
+                st.error("Error loading model weights:")
+                st.error(traceback.format_exc())
+                st.stop()
+
+            # --- Prepare Data Batch ---
+            try:
+                with st.spinner("Preparing data batch..."):
+                    if selected_model == "Prithvi":
+                        data = next(iter(dataset))
+                        batch = preproc([data], padding)
+                        for k, v in batch.items():
+                            if isinstance(v, torch.Tensor):
+                                batch[k] = v.to(device)
+                    elif selected_model == "Aurora":
+                        batch = batch.regrid(res=0.25)
+                        
+                    else:
+                        # Prepare data batch for other models
+                        # Placeholder: Replace with actual data preparation for other models
+                        batch = None
+                st.success("Data batch prepared successfully.")
+            except Exception as e:
+                st.error("Error preparing data batch:")
+                st.error(traceback.format_exc())
+                st.stop()
+
+            # --- Run Inference ---
+            try:
+                with st.spinner("Running model inference..."):
+                    if selected_model == "Prithvi":
+                        model.eval()
+                        with torch.no_grad():
+                            out = model(batch)
+                    elif selected_model == "Aurora":
+                        
+                        model = Aurora(use_lora=False)
+                        # model = Aurora()
+                        model.load_checkpoint("microsoft/aurora", "aurora-0.25-pretrained.ckpt")
+                        # model.load_checkpoint("microsoft/aurora", "aurora-0.25-pretrained.ckpt")
+
+                        model.eval()
+                        # model = model.to("cuda")  # Uncomment if using a GPU
+
+                        with torch.inference_mode():
+                            out = [pred.to("cpu") for pred in rollout(model, batch, steps=2)]
+
+                        model = model.to("cpu")
+                        st.session_state.model = model
+                    else:
+                        # Run inference for other models
+                        # Placeholder: Replace with actual inference code for other models
+                        out = torch.randn(1, 10, 180, 360)  # Dummy tensor
+                st.success("Model inference completed successfully.")
+                st.session_state['out'] = out
+            except Exception as e:
+                st.error("Error during model inference:")
+                st.error(traceback.format_exc())
+                st.stop()
+
+            # --- Visualization Settings ---
+            st.markdown("## 📊 Visualization Settings")
+
+            if 'out' in st.session_state and 'batch' in st.session_state and selected_model == "Prithvi":
+                # Display the shape of the output tensor
+                out_tensor = st.session_state['out']
+                st.write(f"**Output tensor shape:** {out_tensor.shape}")
+
+                # Ensure the output tensor has at least 4 dimensions (batch, variables, lat, lon)
+                if out_tensor.ndim < 4:
+                    st.error("The output tensor does not have the expected number of dimensions (batch, variables, lat, lon).")
+                    st.stop()
+
+                # Get the number of variables
+                num_variables = out_tensor.shape[1]
+
+                # Define variable names (update with your actual variable names)
+                variable_names = [f"Variable_{i}" for i in range(num_variables)]
+
+                # Visualization settings
+                col1, col2 = st.columns(2)
+
+                with col1:
+                    # Select variable to plot
+                    selected_variable_name = st.selectbox(
+                        "Select Variable to Plot",
+                        options=variable_names,
+                        index=0,
+                        help="Choose the variable you want to visualize."
+                    )
+
+                    # Select plot type
+                    plot_type = st.selectbox(
+                        "Select Plot Type",
+                        options=["Contour", "Heatmap"],
+                        index=0,
+                        help="Choose the type of plot to display."
+                    )
+
+                with col2:
+                    # Select color map
+                    cmap = st.selectbox(
+                        "Select Color Map",
+                        options=plt.colormaps(),
+                        index=plt.colormaps().index("viridis"),
+                        help="Choose the color map for the plot."
+                    )
+
+                    # Set number of levels (for contour plot)
+                    if plot_type == "Contour":
+                        num_levels = st.slider(
+                            "Number of Contour Levels",
+                            min_value=5,
+                            max_value=100,
+                            value=20,
+                            step=5,
+                            help="Set the number of contour levels."
+                        )
+                    else:
+                        num_levels = None
+
+                # Find the index based on the selected name
+                variable_index = variable_names.index(selected_variable_name)
+
+                # Extract the selected variable
+                selected_variable = out_tensor[0, variable_index].cpu().numpy()
+
+                # Generate latitude and longitude arrays
+                lat = np.linspace(-90, 90, selected_variable.shape[0])
+                lon = np.linspace(-180, 180, selected_variable.shape[1])
+                X, Y = np.meshgrid(lon, lat)
+
+                # Plot the selected variable
+                st.markdown(f"### Plot of {selected_variable_name}")
+
+                # Matplotlib figure
+                fig, ax = plt.subplots(figsize=(10, 6))
+
+                if plot_type == "Contour":
+                    # Generate the contour plot
+                    contour = ax.contourf(X, Y, selected_variable, levels=num_levels, cmap=cmap)
+                elif plot_type == "Heatmap":
+                    # Generate the heatmap
+                    contour = ax.imshow(selected_variable, extent=[-180, 180, -90, 90], cmap=cmap, origin='lower', aspect='auto')
+
+                # Add a color bar
+                cbar = plt.colorbar(contour, ax=ax)
+                cbar.set_label(f'{selected_variable_name}', fontsize=12)
+
+                # Set aspect ratio and labels
+                ax.set_xlabel("Longitude", fontsize=12)
+                ax.set_ylabel("Latitude", fontsize=12)
+                ax.set_title(f"{selected_variable_name}", fontsize=14)
+
+                # Display the plot in Streamlit
+                st.pyplot(fig)
+
+                # Optional: Provide interactive Plotly plot
+                st.markdown("#### Interactive Plot")
+                if plot_type == "Contour":
+                    fig_plotly = go.Figure(data=go.Contour(
+                        z=selected_variable,
+                        x=lon,
+                        y=lat,
+                        colorscale=cmap,
+                        contours=dict(
+                            coloring='fill',
+                            showlabels=True,
+                            labelfont=dict(size=12, color='white'),
+                            ncontours=num_levels
+                        )
+                    ))
+                elif plot_type == "Heatmap":
+                    fig_plotly = go.Figure(data=go.Heatmap(
+                        z=selected_variable,
+                        x=lon,
+                        y=lat,
+                        colorscale=cmap
+                    ))
+
+                fig_plotly.update_layout(
+                    xaxis_title="Longitude",
+                    yaxis_title="Latitude",
+                    autosize=False,
+                    width=800,
+                    height=600,
+                )
+
+                st.plotly_chart(fig_plotly)
+
+            elif 'out' in st.session_state and selected_model == "Aurora" and st.session_state['out'] is not None:
+                preds = st.session_state['out']
+                ds_subset = st.session_state.get('ds_subset', None)
+                batch = st.session_state.get('batch', None)
+
+                # **Determine Available Levels**
+                # For example, let's assume levels range from 0 to max_level_index
+                # You need to replace 'max_level_index' with the actual maximum level index in your data
+                try:
+                    # Assuming 'lev' dimension exists and is 1D
+                    levels = preds[0].atmos_vars["t"].shape[2]  # Adjust based on your data structure
+                    level_indices = list(range(levels))
+                except Exception as e:
+                    st.error("Error determining available levels:")
+                    st.error(traceback.format_exc())
+                    levels = None  # Set to None if levels cannot be determined
+
+                if levels is not None:
+                    # **Add a Slider for Level Selection**
+                    selected_level = st.slider(
+                        'Select Level',
+                        min_value=0,
+                        max_value=levels - 1,
+                        value=11,  # Default level index
+                        step=1,
+                        help="Select the vertical level for plotting."
+                    )
+
+                    # Loop through predictions and ground truths
+                    for idx in range(len(preds)):
+                        pred = preds[idx]
+                        pred_time = pred.metadata.time[0]
+
+                        # Display prediction time
+                        st.write(f"### Prediction Time: {pred_time}")
+
+                        # **Extract Data at Selected Level**
+                        try:
+                            # Update indices with the selected_level
+                            pred_data = pred.atmos_vars["t"][0][0][selected_level].numpy() - 273.15
+                            truth_data = ds_subset.T.isel(lev=selected_level)[idx].values - 273.15
+                            
+                        except Exception as e:
+                            st.error("Error extracting data for plotting:")
+                            st.error(traceback.format_exc())
+                            continue
+
+                        # Extract latitude and longitude
+                        try:
+                            lat = np.array(pred.metadata.lat)  # Assuming 'lat' is 1D
+                            lon = np.array(pred.metadata.lon)  # Assuming 'lon' is 1D
+                        except Exception as e:
+                            st.error("Error extracting latitude and longitude:")
+                            st.error(traceback.format_exc())
+                            continue
+
+                        # Create a meshgrid for plotting
+                        lon_grid, lat_grid = np.meshgrid(lon, lat)
+
+                        # Create a Matplotlib figure with Cartopy projection
+                        fig, axes = plt.subplots(
+                            1, 3, figsize=(18, 6),
+                            subplot_kw={'projection': ccrs.PlateCarree()}
+                        )
+
+                        # **Ground Truth Plot**
+                        im1 = axes[0].imshow(
+                            truth_data,
+                            extent=[lon.min(), lon.max(), lat.min(), lat.max()],
+                            origin='lower',
+                            cmap='coolwarm',
+                            transform=ccrs.PlateCarree()
+                        )
+                        axes[0].set_title(f"Ground Truth at Level {selected_level} - {pred_time}")
+                        axes[0].set_xlabel('Longitude')
+                        axes[0].set_ylabel('Latitude')
+                        plt.colorbar(im1, ax=axes[0], orientation='horizontal', pad=0.05)
+
+                        # **Prediction Plot**
+                        im2 = axes[1].imshow(
+                            pred_data,
+                            extent=[lon.min(), lon.max(), lat.min(), lat.max()],
+                            origin='lower',
+                            cmap='coolwarm',
+                            transform=ccrs.PlateCarree()
+                        )
+                        axes[1].set_title(f"Prediction at Level {selected_level} - {pred_time}")
+                        axes[1].set_xlabel('Longitude')
+                        axes[1].set_ylabel('Latitude')
+                        plt.colorbar(im2, ax=axes[1], orientation='horizontal', pad=0.05)
+
+                        plt.tight_layout()
+
+                        # Display the plot in Streamlit
+                        st.pyplot(fig)
+                else:
+                    st.error("Could not determine the available levels in the data.")
+
+
+            else:
+                st.warning("No output available to display or visualization is not implemented for this model.")
+
+    # --- End of Inference Button ---
+    else:
+        with right_col:
+            st.header("🖥️ Visualization & Progress")
+            st.info("Awaiting inference to display results.")
+
+

aurora

@@ -0,0 +1 @@
+Subproject commit 8b11659b91d06f87c2d22e541dbcd0092baf2157

aurora_utils.py

@@ -0,0 +1,129 @@
+import streamlit as st
+import torch
+from aurora import Aurora, Batch, Metadata
+import numpy as np
+from datetime import datetime
+
+def aurora_config_ui():
+    st.subheader("Aurora Model Data Upload")
+    st.markdown("### Drag and Drop Your Data Files Here")
+    uploaded_files = st.file_uploader(
+        "Upload Data Files for Aurora",
+        accept_multiple_files=True,
+        key="aurora_uploader",
+        type=["nc", "netcdf", "nc4"]
+    )
+    return uploaded_files
+
+def prepare_aurora_batch(ds):
+    desired_levels = [50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000]
+
+    # Ensure that the 'lev' dimension exists
+    if 'lev' not in ds.dims:
+        raise ValueError("The dataset does not contain a 'lev' (pressure level) dimension.")
+
+    # Define the _prepare function
+    def _prepare(x: np.ndarray, i: int) -> torch.Tensor:
+                                        # Select previous and current time steps
+        selected = x[[i - 6, i]]
+                                        
+        # Add a batch dimension
+        selected = selected[None]
+                                        
+        # Ensure data is contiguous
+        selected = selected.copy()
+                                        
+        # Convert to PyTorch tensor
+        return torch.from_numpy(selected)
+
+                                    # Adjust latitudes and longitudes
+    lat = ds.lat.values * -1
+    lon = ds.lon.values + 180
+
+                                    # Subset the dataset to only include the desired pressure levels
+    ds_subset = ds.sel(lev=desired_levels, method="nearest")
+
+                                    # Verify that all desired levels are present
+    present_levels = ds_subset.lev.values
+    missing_levels = set(desired_levels) - set(present_levels)
+    if missing_levels:
+        raise ValueError(f"The following desired pressure levels are missing in the dataset: {missing_levels}")
+
+    # Extract pressure levels after subsetting
+    lev = ds_subset.lev.values  # Pressure levels in hPa
+
+    # Prepare surface variables at 1000 hPa
+    try:
+        lev_index_1000 = np.where(lev == 1000)[0][0]
+    except IndexError:
+        raise ValueError("1000 hPa level not found in the 'lev' dimension after subsetting.")
+
+    T_surface = ds_subset.T.isel(lev=lev_index_1000).compute()
+    U_surface = ds_subset.U.isel(lev=lev_index_1000).compute()
+    V_surface = ds_subset.V.isel(lev=lev_index_1000).compute()
+    SLP = ds_subset.SLP.compute()
+
+                                    # Reorder static variables (selecting the first time index to remove the time dimension)
+    PHIS = ds_subset.PHIS.isel(time=0).compute()
+
+                                    # Prepare atmospheric variables for the desired pressure levels excluding 1000 hPa
+    atmos_levels = [int(level) for level in lev if level != 1000]
+
+    T_atm = (ds_subset.T.sel(lev=atmos_levels)).compute()
+    U_atm = (ds_subset.U.sel(lev=atmos_levels)).compute()
+    V_atm = (ds_subset.V.sel(lev=atmos_levels)).compute()
+
+                                    # Select time index
+    num_times = ds_subset.time.size
+    i = 6  # Adjust as needed (1 <= i < num_times)
+
+    if i >= num_times or i < 1:
+        raise IndexError("Time index i is out of bounds.")
+
+    time_values = ds_subset.time.values
+    current_time = np.datetime64(time_values[i]).astype('datetime64[s]').astype(datetime)
+
+                                    # Prepare surface variables
+    surf_vars = {
+                                        "2t": _prepare(T_surface.values, i),   # Two-meter temperature
+                                        "10u": _prepare(U_surface.values, i),  # Ten-meter eastward wind
+                                        "10v": _prepare(V_surface.values, i),  # Ten-meter northward wind
+                                        "msl": _prepare(SLP.values, i),        # Mean sea-level pressure
+                                    }
+
+                                    # Prepare static variables (now 2D tensors)
+    static_vars = {
+                                        "z": torch.from_numpy(PHIS.values.copy()),  # Geopotential (h, w)
+                                        # Add 'lsm' and 'slt' if available and needed
+                                    }
+
+                                    # Prepare atmospheric variables
+    atmos_vars = {
+                                        "t": _prepare(T_atm.values, i),  # Temperature at desired levels
+                                        "u": _prepare(U_atm.values, i),  # Eastward wind at desired levels
+                                        "v": _prepare(V_atm.values, i),  # Southward wind at desired levels
+                                    }
+
+                                    # Define metadata
+    metadata = Metadata(
+                                        lat=torch.from_numpy(lat.copy()),
+                                        lon=torch.from_numpy(lon.copy()),
+                                        time=(current_time,),
+                                        atmos_levels=tuple(atmos_levels),  # Only the desired atmospheric levels
+                                    )
+
+                                    # Create the Batch object
+    batch = Batch(
+                                        surf_vars=surf_vars,
+                                        static_vars=static_vars,
+                                        atmos_vars=atmos_vars,
+                                        metadata=metadata
+                                    ) # Display the dataset or perform further processing
+    return batch
+
+def initialize_aurora_model(device):
+    model = Aurora(use_lora=False)
+    # Load pretrained checkpoint if available
+    model.load_checkpoint("microsoft/aurora", "aurora-0.25-pretrained.ckpt")
+    model = model.to(device)
+    return model

config_utils.py

@@ -0,0 +1,8 @@
+import yaml
+import streamlit as st
+from pathlib import Path
+
+def load_config(config_path: Path):
+    with open(config_path, "r") as f:
+        config = yaml.safe_load(f)
+    return config

data_utils.py

@@ -0,0 +1,39 @@
+import os
+import tempfile
+import traceback
+import streamlit as st
+import xarray as xr
+from typing import List
+import numpy as np
+
+@st.cache_resource
+def save_uploaded_files(uploaded_files):
+    if 'temp_file_paths' not in st.session_state:
+        st.session_state.temp_file_paths = []
+    for uploaded_file in uploaded_files:
+        suffix = os.path.splitext(uploaded_file.name)[1]
+        temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=suffix)
+        temp_file.write(uploaded_file.read())
+        temp_file.close()
+        st.session_state.temp_file_paths.append(temp_file.name)
+        
+
+@st.cache_resource
+def load_dataset(file_paths: List[str]):
+    try:
+        ds = xr.open_mfdataset(file_paths, combine='by_coords').load()
+        return ds
+    except Exception:
+        st.error("Error loading dataset:")
+        st.error(traceback.format_exc())
+        return None
+    
+@st.cache_resource
+def load_dataset_pangu(file_path: str):
+    try:
+        ds = np.load(file_path)
+        return ds
+    except Exception:
+        st.error("Error loading dataset:")
+        st.error(traceback.format_exc())
+        return None

fengwu_utils.py

@@ -0,0 +1,311 @@
+import streamlit as st
+import torch
+# from Pangu-Weather import *
+import numpy as np
+from datetime import datetime
+import numpy as np
+import onnx
+import onnxruntime as ort
+import matplotlib.pyplot as plt
+import cartopy.crs as ccrs
+import io
+
+def fengwu_config_data():
+    st.subheader("FengWu Model Data Input")
+
+    # Detailed data description section
+    st.markdown("""
+    **Input Data Requirements (FengWu):**  
+    FengWu takes **two consecutive six-hour atmospheric states** as input:
+    1. **First Input (input1.npy)**: Atmospheric data at the initial time.  
+    2. **Second Input (input2.npy)**: Atmospheric data 6 hours later.
+
+    **Shape & Variables:**  
+    Each input is a NumPy array with shape `(69, 721, 1440)`:
+    - **Dimension 0 (69 features):**  
+      The first 4 features are surface variables:  
+      1. U10 (10-meter Eastward Wind)  
+      2. V10 (10-meter Northward Wind)  
+      3. T2M (2-meter Temperature)  
+      4. MSL (Mean Sea Level Pressure)
+
+      These are followed by non-surface variables, each with 13 pressure levels:  
+      - Z (Geopotential)  
+      - Q (Specific Humidity)  
+      - U (Eastward Wind)  
+      - V (Northward Wind)  
+      - T (Temperature)  
+
+      The 13 vertical levels are: [50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000] hPa
+
+      The total count is:  
+      - Surface vars: 4  
+      - For each non-surface var (Z, Q, U, V, T): 13 levels = 65 vars  
+      4 (surface) + 65 (5 vars * 13 levels) = 69 total features.
+
+    **Spatial & Coordinate Details:**  
+    - Latitude dimension (721 points) ranges from 90°N to -90°S with ~0.25° spacing.
+    - Longitude dimension (1440 points) ranges from 0° to 360°E with ~0.25° spacing.
+    - Ensure data is single precision floats (`.astype(np.float32)`).
+
+    **Data Frequency & Forecasting Scheme:**  
+    - `input1.npy` corresponds to a given time (e.g., 06:00 UTC Jan 1, 2018).
+    - `input2.npy` corresponds to 6 hours later (e.g., 12:00 UTC Jan 1, 2018).
+    - The model predicts future states at subsequent 6-hour intervals.
+
+    **Converting Your Data:**  
+    - ERA5 `.nc` files or ECMWF `.grib` files can be converted to `.npy` using appropriate Python packages (`netCDF4` or `pygrib`).
+    - Ensure you follow the exact variable and level ordering as described.
+
+
+    """)
+
+    # File uploaders for FengWu input data (two consecutive time steps)
+    st.markdown("### Upload Your FengWu Input Data Files")
+    input1_file = st.file_uploader(
+        "Upload input1.npy (Initial Time)",
+        type=["npy"],
+        key="fengwu_input1"
+    )
+
+    input2_file = st.file_uploader(
+        "Upload input2.npy (6 Hours Later)",
+        type=["npy"],
+        key="fengwu_input2"
+    )
+
+    st.markdown("---")
+    st.markdown("### References & Resources")
+    st.markdown("""
+    - **Research Paper:** [FengWu: Pushing the Skillful Global Medium-range Weather Forecast beyond 10 Days Lead](https://arxiv.org/abs/2304.02948)
+                
+    - **GitHub Source Code:** [Fengwu on GitHub](https://github.com/OpenEarthLab/FengWu?tab=readme-ov-file)
+    """)
+
+    return input1_file, input2_file
+
+
+@st.cache_resource
+def inference_6hrs_fengwu(input1, input2):
+    model_6 = onnx.load('FengWu/fengwu_v2.onnx')
+
+    # Set the behavier of onnxruntime
+    options = ort.SessionOptions()
+    options.enable_cpu_mem_arena=False
+    options.enable_mem_pattern = False
+    options.enable_mem_reuse = False
+    # Increase the number for faster inference and more memory consumption
+    options.intra_op_num_threads = 1
+
+    # Set the behavier of cuda provider
+    cuda_provider_options = {'arena_extend_strategy':'kSameAsRequested',}
+
+    # Initialize onnxruntime session for Pangu-Weather Models
+    ort_session_6 = ort.InferenceSession('FengWu/fengwu_v2.onnx', sess_options=options, providers=[('CUDAExecutionProvider', cuda_provider_options)])
+
+
+    data_mean = np.load("FengWu/data_mean.npy")[:, np.newaxis, np.newaxis]
+    data_std = np.load("FengWu/data_std.npy")[:, np.newaxis, np.newaxis]
+
+    input1_after_norm = (input1 - data_mean) / data_std
+    input2_after_norm = (input2 - data_mean) / data_std
+    input = np.concatenate((input1_after_norm, input2_after_norm), axis=0)[np.newaxis, :, :, :]
+    input = input.astype(np.float32)
+
+    output = ort_session_6.run(None, {'input':input})[0]
+    output = (output[0, :69] * data_std) + data_mean
+
+    return output
+
+
+@st.cache_resource
+def inference_12hrs_fengwu(input1, input2):
+    model_6 = onnx.load('FengWu/fengwu_v2.onnx')
+
+    # Set the behavier of onnxruntime
+    options = ort.SessionOptions()
+    options.enable_cpu_mem_arena=False
+    options.enable_mem_pattern = False
+    options.enable_mem_reuse = False
+    # Increase the number for faster inference and more memory consumption
+    options.intra_op_num_threads = 1
+
+    # Set the behavier of cuda provider
+    cuda_provider_options = {'arena_extend_strategy':'kSameAsRequested',}
+
+    # Initialize onnxruntime session for Pangu-Weather Models
+    ort_session_6 = ort.InferenceSession('FengWu/fengwu_v2.onnx', sess_options=options, providers=[('CUDAExecutionProvider', cuda_provider_options)])
+
+
+    data_mean = np.load("FengWu/data_mean.npy")[:, np.newaxis, np.newaxis]
+    data_std = np.load("FengWu/data_std.npy")[:, np.newaxis, np.newaxis]
+
+    input1_after_norm = (input1 - data_mean) / data_std
+    input2_after_norm = (input2 - data_mean) / data_std
+    input = np.concatenate((input1_after_norm, input2_after_norm), axis=0)[np.newaxis, :, :, :]
+    input = input.astype(np.float32)
+
+    for i in range(2):
+        output = ort_session_6.run(None, {'input':input})[0]
+        input = np.concatenate((input[:, 69:], output[:, :69]), axis=1)
+        output = (output[0, :69] * data_std) + data_mean
+        # print(output.shape)
+
+    return output
+
+@st.cache_resource
+def inference_custom_hrs_fengwu(input1, input2, forecast_hours):
+    model_6 = onnx.load('FengWu/fengwu_v2.onnx')
+
+    # Set the behavier of onnxruntime
+    options = ort.SessionOptions()
+    options.enable_cpu_mem_arena=False
+    options.enable_mem_pattern = False
+    options.enable_mem_reuse = False
+    # Increase the number for faster inference and more memory consumption
+    options.intra_op_num_threads = 1
+
+    # Set the behavier of cuda provider
+    cuda_provider_options = {'arena_extend_strategy':'kSameAsRequested',}
+
+    # Initialize onnxruntime session for Pangu-Weather Models
+    ort_session_6 = ort.InferenceSession('FengWu/fengwu_v2.onnx', sess_options=options, providers=[('CUDAExecutionProvider', cuda_provider_options)])
+
+
+    data_mean = np.load("FengWu/data_mean.npy")[:, np.newaxis, np.newaxis]
+    data_std = np.load("FengWu/data_std.npy")[:, np.newaxis, np.newaxis]
+
+    input1_after_norm = (input1 - data_mean) / data_std
+    input2_after_norm = (input2 - data_mean) / data_std
+    input = np.concatenate((input1_after_norm, input2_after_norm), axis=0)[np.newaxis, :, :, :]
+    input = input.astype(np.float32)
+
+    for i in range(forecast_hours/6):
+        output = ort_session_6.run(None, {'input':input})[0]
+        input = np.concatenate((input[:, 69:], output[:, :69]), axis=1)
+        output = (output[0, :69] * data_std) + data_mean
+        # print(output.shape)
+
+    return output
+
+def plot_fengwu_output(initial_data, predicted_data):
+    """
+    Plot initial and predicted Fengwu model outputs.
+    
+    Parameters:
+    - initial_data: np.ndarray of shape (69, 721, 1440) representing the initial or input state.
+    - predicted_data: np.ndarray of shape (69, 721, 1440) representing the predicted state by Fengwu.
+    """
+    # Coordinate setup
+    lat = np.linspace(90, -90, 721)   # Latitude from 90N to 90S
+    lon = np.linspace(0, 360, 1440)   # Longitude from 0E to 360E
+
+    # Surface and upper-level variable definitions
+    surface_vars = ["U10", "V10", "T2M", "MSL"]
+    upper_vars = ["Z (Geopotential)", "Q (Specific Humidity)", "U (Eastward Wind)", "V (Northward Wind)", "T (Temperature)"]
+    upper_levels = [50,100,150,200,250,300,400,500,600,700,850,925,1000]
+
+    # Mapping of upper variable groups to their starting indices
+    # Each group has 13 levels, so indices shift by 13 for each subsequent group.
+    var_group_start = {
+        "Z (Geopotential)": 4,   # Z starts at index 4
+        "Q (Specific Humidity)": 17,  # Q = 4+13=17
+        "U (Eastward Wind)": 30, # U = 17+13=30
+        "V (Northward Wind)": 43,# V = 30+13=43
+        "T (Temperature)": 56    # T = 43+13=56
+    }
+
+    # --- Initial Data Visualization ---
+    st.subheader("Initial Data Visualization (Fengwu)")
+    init_col1, init_col2 = st.columns([1,1])
+
+    with init_col1:
+        init_data_choice = st.selectbox("Data Source", ["Upper-Air Data", "Surface Data"], key="fengwu_init_data_choice")
+    with init_col2:
+        if init_data_choice == "Upper-Air Data":
+            init_var = st.selectbox("Variable", upper_vars, key="fengwu_init_upper_var")
+        else:
+            init_var = st.selectbox("Variable", surface_vars, key="fengwu_init_surface_var")
+
+    # Select the data slice for initial data
+    if init_data_choice == "Upper-Air Data":
+        selected_level_hpa_init = st.select_slider(
+            "Select Pressure Level (hPa)",
+            options=upper_levels,
+            value=850,  # Default to 850hPa
+            help="Select the pressure level in hPa.",
+            key="fengwu_init_level_hpa_slider"
+        )
+        level_index_init = upper_levels.index(selected_level_hpa_init)
+        start_index_init = var_group_start[init_var]
+        data_index_init = start_index_init + level_index_init
+        data_to_plot_init = initial_data[data_index_init, :, :]
+        title_init = f"Initial Upper-Air: {init_var} at {selected_level_hpa_init}hPa"
+    else:
+        # Surface variable
+        var_index_init = surface_vars.index(init_var)
+        data_to_plot_init = initial_data[var_index_init, :, :]
+        title_init = f"Initial Surface: {init_var}"
+
+    # Plot initial data
+    fig_init, ax_init = plt.subplots(figsize=(10, 5), subplot_kw={'projection': ccrs.PlateCarree()})
+    ax_init.set_title(title_init)
+    im_init = ax_init.imshow(data_to_plot_init, extent=[lon.min(), lon.max(), lat.min(), lat.max()],
+                             origin='lower', cmap='coolwarm', transform=ccrs.PlateCarree())
+    ax_init.coastlines()
+    plt.colorbar(im_init, ax=ax_init, orientation='horizontal', pad=0.05)
+    st.pyplot(fig_init)
+
+    # --- Predicted Data Visualization ---
+    st.subheader("Predicted Data Visualization (Fengwu)")
+    pred_col1, pred_col2 = st.columns([1,1])
+
+    with pred_col1:
+        pred_data_choice = st.selectbox("Data Source", ["Upper-Air Data", "Surface Data"], key="fengwu_pred_data_choice")
+    with pred_col2:
+        if pred_data_choice == "Upper-Air Data":
+            pred_var = st.selectbox("Variable", upper_vars, key="fengwu_pred_upper_var")
+        else:
+            pred_var = st.selectbox("Variable", surface_vars, key="fengwu_pred_surface_var")
+
+    # Select the data slice for predicted data
+    if pred_data_choice == "Upper-Air Data":
+        selected_level_hpa_pred = st.select_slider(
+            "Select Pressure Level (hPa)",
+            options=upper_levels,
+            value=850,  # Default to 850hPa
+            help="Select the pressure level in hPa.",
+            key="fengwu_pred_level_hpa_slider"
+        )
+        level_index_pred = upper_levels.index(selected_level_hpa_pred)
+        start_index_pred = var_group_start[pred_var]
+        data_index_pred = start_index_pred + level_index_pred
+        data_to_plot_pred = predicted_data[data_index_pred, :, :]
+        title_pred = f"Predicted Upper-Air: {pred_var} at {selected_level_hpa_pred}hPa"
+    else:
+        # Surface variable for predicted data
+        var_index_pred = surface_vars.index(pred_var)
+        data_to_plot_pred = predicted_data[var_index_pred, :, :]
+        title_pred = f"Predicted Surface: {pred_var}"
+
+    # Plot predicted data
+    fig_pred, ax_pred = plt.subplots(figsize=(10, 5), subplot_kw={'projection': ccrs.PlateCarree()})
+    ax_pred.set_title(title_pred)
+    im_pred = ax_pred.imshow(data_to_plot_pred, extent=[lon.min(), lon.max(), lat.min(), lat.max()],
+                             origin='lower', cmap='coolwarm', transform=ccrs.PlateCarree())
+    ax_pred.coastlines()
+    plt.colorbar(im_pred, ax=ax_pred, orientation='horizontal', pad=0.05)
+    st.pyplot(fig_pred)
+
+    # --- Download Buttons ---
+    st.subheader("Download Predicted Fengwu Data")
+
+    # Convert predicted_data to binary format for download
+    buffer_pred = io.BytesIO()
+    np.save(buffer_pred, predicted_data)
+    buffer_pred.seek(0)
+
+    st.download_button(label="Download Predicted Fengwu Data",
+                       data=buffer_pred,
+                       file_name="predicted_fengwu.npy",
+                       mime="application/octet-stream")

inference_utils.py

@@ -0,0 +1,19 @@
+import torch
+import streamlit as st
+from aurora import rollout, Aurora
+
+def run_inference(selected_model, model, batch, device):
+    if selected_model == "Prithvi":
+        model.eval()
+        with torch.no_grad():
+            out = model(batch)
+        return out
+    elif selected_model == "Aurora":
+        model.eval()
+        with torch.inference_mode():
+            # Example: Predict 2 steps ahead
+            out = [pred.to("cpu") for pred in rollout(model, batch, steps=2)]
+        return out
+    else:
+        st.error("Inference not implemented for this model.")
+        return None

pangu_utils.py

@@ -0,0 +1,311 @@
+import streamlit as st
+import torch
+# from Pangu-Weather import *
+import numpy as np
+from datetime import datetime
+import numpy as np
+import onnx
+import onnxruntime as ort
+import matplotlib.pyplot as plt
+import cartopy.crs as ccrs
+import io
+
+
+def pangu_config_data():
+    st.subheader("Pangu-Weather Model Data Input")
+
+    # Detailed data description section
+    st.markdown("""
+    **Input Data Requirements:**  
+    Pangu-Weather uses two NumPy arrays to represent initial atmospheric conditions:  
+    1. **Surface Data (input_surface.npy)**  
+       - Shape: `(4, 721, 1440)`  
+       - Variables: MSLP, U10, V10, T2M in this exact order.
+         - **MSLP:** Mean Sea Level Pressure  
+         - **U10:** 10-meter Eastward Wind  
+         - **V10:** 10-meter Northward Wind  
+         - **T2M:** 2-meter Temperature  
+    2. **Upper-Air Data (input_upper.npy)**  
+       - Shape: `(5, 13, 721, 1440)`  
+       - Variables (first dim): Z, Q, T, U, V in this exact order
+         - **Z:** Geopotential (Note: if your source provides geopotential height, multiply by 9.80665 to get geopotential)
+         - **Q:** Specific Humidity  
+         - **T:** Temperature  
+         - **U:** Eastward Wind  
+         - **V:** Northward Wind  
+       - Pressure Levels (second dim): 1000hPa, 925hPa, 850hPa, 700hPa, 600hPa, 500hPa, 400hPa, 300hPa, 250hPa, 200hPa, 150hPa, 100hPa, 50hPa in this exact order.
+
+    **Spatial & Coordinate Details:**  
+    - Latitude dimension (721 points) ranges from 90°N to -90°S with a 0.25° spacing.  
+    - Longitude dimension (1440 points) ranges from 0° to 359.75°E with a 0.25° spacing.
+    - Data should be single precision floats (`.astype(np.float32)`).
+
+    **Supported Data Sources:**  
+    - ERA5 initial fields (strongly recommended).  
+    - ECMWF initial fields (e.g., HRES forecast) can be used, but may result in a slight accuracy drop.  
+    - Other types of initial fields are not currently supported due to potentially large discrepancies in data fields.
+
+    **Converting Your Data:**  
+    - ERA5 `.nc` files can be converted to `.npy` using the `netCDF4` Python package.
+    - ECMWF `.grib` files can be converted to `.npy` using the `pygrib` Python package.
+    - Ensure the order of variables and pressure levels is exactly as described above.
+    """)
+
+    # File uploaders for surface and upper data separately
+    st.markdown("### Upload Your Input Data Files")
+    input_surface_file = st.file_uploader(
+        "Upload input_surface.npy",
+        type=["npy"],
+        key="pangu_input_surface"
+    )
+
+    input_upper_file = st.file_uploader(
+        "Upload input_upper.npy",
+        type=["npy"],
+        key="pangu_input_upper"
+    )
+
+    st.markdown("---")
+    st.markdown("### References & Resources")
+    st.markdown("""
+    - **Research Paper:** [Accurate medium-range global weather forecasting with 3D neural networks](https://www.nature.com/articles/s41586-023-06185-3)
+                - [Pangu-Weather: A 3D High-Resolution Model for Fast and Accurate Global Weather Forecast](https://arxiv.org/abs/2211.02556)
+    - **GitHub Source Code:** [Pangu-Weather on GitHub](https://github.com/198808xc/Pangu-Weather?tab=readme-ov-file)
+    """)
+
+    return input_surface_file, input_upper_file
+
+def inference_24hrs(input, input_surface):
+    model_24 = onnx.load('Pangu-Weather/pangu_weather_24.onnx')
+
+    # Set the behavier of onnxruntime
+    options = ort.SessionOptions()
+    options.enable_cpu_mem_arena=False
+    options.enable_mem_pattern = False
+    options.enable_mem_reuse = False
+    # Increase the number for faster inference and more memory consumption
+    options.intra_op_num_threads = 1
+
+    # Set the behavier of cuda provider
+    cuda_provider_options = {'arena_extend_strategy':'kSameAsRequested',}
+
+    # Initialize onnxruntime session for Pangu-Weather Models
+    ort_session_24 = ort.InferenceSession('Pangu-Weather/pangu_weather_24.onnx', sess_options=options, providers=['CPUExecutionProvider'])
+
+    # Run the inference session
+    output, output_surface = ort_session_24.run(None, {'input':input, 'input_surface':input_surface})
+
+    return output, output_surface
+
+@st.cache_resource
+def inference_6hrs(input, input_surface):
+    model_6 = onnx.load('Pangu-Weather/pangu_weather_6.onnx')
+
+    # Set the behavier of onnxruntime
+    options = ort.SessionOptions()
+    options.enable_cpu_mem_arena=False
+    options.enable_mem_pattern = False
+    options.enable_mem_reuse = False
+    # Increase the number for faster inference and more memory consumption
+    options.intra_op_num_threads = 1
+
+    # Set the behavier of cuda provider
+    cuda_provider_options = {'arena_extend_strategy':'kSameAsRequested',}
+
+    # Initialize onnxruntime session for Pangu-Weather Models
+    ort_session_6 = ort.InferenceSession('Pangu-Weather/pangu_weather_6.onnx', sess_options=options, providers=['CPUExecutionProvider'])
+
+    # Run the inference session
+    output, output_surface = ort_session_6.run(None, {'input':input, 'input_surface':input_surface})
+
+    return output, output_surface
+
+@st.cache_resource
+def inference_1hr(input, input_surface):
+    model_1 = onnx.load('Pangu-Weather/pangu_weather_1.onnx')
+
+    # Set the behavier of onnxruntime
+    options = ort.SessionOptions()
+    options.enable_cpu_mem_arena=False
+    options.enable_mem_pattern = False
+    options.enable_mem_reuse = False
+    # Increase the number for faster inference and more memory consumption
+    options.intra_op_num_threads = 1
+
+    # Set the behavier of cuda provider
+    cuda_provider_options = {'arena_extend_strategy':'kSameAsRequested',}
+
+    # Initialize onnxruntime session for Pangu-Weather Models
+    ort_session_1 = ort.InferenceSession('Pangu-Weather/pangu_weather_1.onnx', sess_options=options, providers=['CPUExecutionProvider'])
+
+    # Run the inference session
+    output, output_surface = ort_session_1.run(None, {'input':input, 'input_surface':input_surface})
+
+    return output, output_surface
+
+@st.cache_resource
+def inference_3hrs(input, input_surface):
+    model_3 = onnx.load('Pangu-Weather/pangu_weather_3.onnx')
+
+    # Set the behavier of onnxruntime
+    options = ort.SessionOptions()
+    options.enable_cpu_mem_arena=False
+    options.enable_mem_pattern = False
+    options.enable_mem_reuse = False
+    # Increase the number for faster inference and more memory consumption
+    options.intra_op_num_threads = 1
+
+    # Set the behavier of cuda provider
+    cuda_provider_options = {'arena_extend_strategy':'kSameAsRequested',}
+
+    # Initialize onnxruntime session for Pangu-Weather Models
+    ort_session_3 = ort.InferenceSession('Pangu-Weather/pangu_weather_3.onnx', sess_options=options, providers=['CPUExecutionProvider'])
+
+    # Run the inference session
+    output, output_surface = ort_session_3.run(None, {'input':input, 'input_surface':input_surface})
+
+    return output, output_surface
+
+@st.cache_resource
+def inference_custom_hrs(input, input_surface, forecast_hours):
+    # Ensure forecast_hours is a multiple of 24
+    if forecast_hours % 24 != 0:
+        raise ValueError("forecast_hours must be a multiple of 24.")
+
+    # Load the 24-hour model
+    model_24 = onnx.load('Pangu-Weather/pangu_weather_24.onnx')
+
+    # Configure ONNX Runtime session
+    options = ort.SessionOptions()
+    options.enable_cpu_mem_arena = False
+    options.enable_mem_pattern = False
+    options.enable_mem_reuse = False
+    options.intra_op_num_threads = 1
+
+    # Using CPUExecutionProvider for simplicity
+    ort_session_24 = ort.InferenceSession('Pangu-Weather/pangu_weather_24.onnx', sess_options=options, providers=['CPUExecutionProvider'])
+
+    # Calculate how many 24-hour steps we need
+    steps = forecast_hours // 24
+
+    # Run the 24-hour model repeatedly
+    for i in range(steps):
+        output, output_surface = ort_session_24.run(None, {'input': input, 'input_surface': input_surface})
+        input, input_surface = output, output_surface
+
+    # Return the final predictions after completing all steps
+    return input, input_surface
+
+
+def plot_pangu_output(upper_data, surface_data, out_upper, out_surface):
+    # Coordinate setup
+    lat = np.linspace(90, -90, 721)   # Latitude grid
+    lon = np.linspace(0, 360, 1440)   # Longitude grid
+
+    # Variable and level names
+    upper_vars = ["Z (Geopotential)", "Q (Specific Humidity)", "T (Temperature)", "U (Eastward Wind)", "V (Northward Wind)"]
+    upper_levels = ["1000hPa", "925hPa", "850hPa", "700hPa", "600hPa", "500hPa",
+                    "400hPa", "300hPa", "250hPa", "200hPa", "150hPa", "100hPa", "50hPa"]
+    # Extract numeric hPa values for selection
+    upper_hpa_values = [int(l.replace("hPa", "")) for l in upper_levels]
+
+    surface_vars = ["MSLP", "U10", "V10", "T2M"]
+
+    # --- Initial Data Visualization ---
+    st.subheader("Initial Data Visualization")
+    init_col1, init_col2 = st.columns([1,1])
+
+    with init_col1:
+        init_data_choice = st.selectbox("Data Source", ["Upper-Air Data", "Surface Data"], key="init_data_choice")
+    with init_col2:
+        if init_data_choice == "Upper-Air Data":
+            init_var = st.selectbox("Variable", upper_vars, key="init_upper_var")
+        else:
+            init_var = st.selectbox("Variable", surface_vars, key="init_surface_var")
+
+    if init_data_choice == "Upper-Air Data":
+        selected_level_hpa_init = st.select_slider(
+            "Select Pressure Level (hPa)",
+            options=upper_hpa_values,
+            value=850,  # Default to 850hPa
+            help="Select the pressure level in hPa.",
+            key="init_level_hpa_slider"
+        )
+        # Find the corresponding index from the selected hPa value
+        selected_level_index_init = upper_hpa_values.index(selected_level_hpa_init)
+        selected_var_index_init = upper_vars.index(init_var)
+        data_to_plot_init = upper_data[selected_var_index_init, selected_level_index_init, :, :]
+        title_init = f"Initial Upper-Air: {init_var} at {selected_level_hpa_init}hPa"
+    else:
+        selected_var_index_init = surface_vars.index(init_var)
+        data_to_plot_init = surface_data[selected_var_index_init, :, :]
+        title_init = f"Initial Surface: {init_var}"
+
+    # Plot initial data
+    fig_init, ax_init = plt.subplots(figsize=(10, 5), subplot_kw={'projection': ccrs.PlateCarree()})
+    ax_init.set_title(title_init)
+    im_init = ax_init.imshow(data_to_plot_init, extent=[lon.min(), lon.max(), lat.min(), lat.max()],
+                             origin='lower', cmap='coolwarm', transform=ccrs.PlateCarree())
+    ax_init.coastlines()
+    plt.colorbar(im_init, ax=ax_init, orientation='horizontal', pad=0.05)
+    st.pyplot(fig_init)
+
+    # --- Predicted Data Visualization ---
+    st.subheader("Predicted Data Visualization")
+    pred_col1, pred_col2 = st.columns([1,1])
+
+    with pred_col1:
+        pred_data_choice = st.selectbox("Data Source", ["Upper-Air Data", "Surface Data"], key="pred_data_choice")
+    with pred_col2:
+        if pred_data_choice == "Upper-Air Data":
+            pred_var = st.selectbox("Variable", upper_vars, key="pred_upper_var")
+        else:
+            pred_var = st.selectbox("Variable", surface_vars, key="pred_surface_var")
+
+    if pred_data_choice == "Upper-Air Data":
+        selected_level_hpa_pred = st.select_slider(
+            "Select Pressure Level (hPa)",
+            options=upper_hpa_values,
+            value=850,  # Default to 850hPa
+            help="Select the pressure level in hPa.",
+            key="pred_level_hpa_slider"
+        )
+        selected_level_index_pred = upper_hpa_values.index(selected_level_hpa_pred)
+        selected_var_index_pred = upper_vars.index(pred_var)
+        data_to_plot_pred = out_upper[selected_var_index_pred, selected_level_index_pred, :, :]
+        title_pred = f"Predicted Upper-Air: {pred_var} at {selected_level_hpa_pred}hPa"
+    else:
+        selected_var_index_pred = surface_vars.index(pred_var)
+        data_to_plot_pred = out_surface[selected_var_index_pred, :, :]
+        title_pred = f"Predicted Surface: {pred_var}"
+
+    # Plot predicted data
+    fig_pred, ax_pred = plt.subplots(figsize=(10, 5), subplot_kw={'projection': ccrs.PlateCarree()})
+    ax_pred.set_title(title_pred)
+    im_pred = ax_pred.imshow(data_to_plot_pred, extent=[lon.min(), lon.max(), lat.min(), lat.max()],
+                             origin='lower', cmap='coolwarm', transform=ccrs.PlateCarree())
+    ax_pred.coastlines()
+    plt.colorbar(im_pred, ax=ax_pred, orientation='horizontal', pad=0.05)
+    st.pyplot(fig_pred)
+
+    # --- Download Buttons ---
+    st.subheader("Download Predicted Data")
+
+    # Convert out_upper and out_surface to binary format for download
+    buffer_upper = io.BytesIO()
+    np.save(buffer_upper, out_upper)
+    buffer_upper.seek(0)
+
+    buffer_surface = io.BytesIO()
+    np.save(buffer_surface, out_surface)
+    buffer_surface.seek(0)
+
+    st.download_button(label="Download Predicted Upper-Air Data",
+                       data=buffer_upper,
+                       file_name="predicted_upper.npy",
+                       mime="application/octet-stream")
+
+    st.download_button(label="Download Predicted Surface Data",
+                       data=buffer_surface,
+                       file_name="predicted_surface.npy",
+                       mime="application/octet-stream")

plot_utils.py

@@ -0,0 +1,127 @@
+import streamlit as st
+import matplotlib.pyplot as plt
+import numpy as np
+import plotly.graph_objects as go
+import cartopy.crs as ccrs
+
+def plot_prithvi_output(out_tensor):
+    if out_tensor is None:
+        st.warning("No output available for plotting.")
+        return
+
+    # Example visualization UI for Prithvi:
+    st.markdown("## 📊 Visualization Settings")
+    # Extract shape and variable names as needed
+    if out_tensor.ndim < 4:
+        st.error("The output tensor does not have the expected dimensions.")
+        return
+
+    num_variables = out_tensor.shape[1]
+    variable_names = [f"Variable_{i}" for i in range(num_variables)]
+
+    col1, col2 = st.columns(2)
+    with col1:
+        selected_variable_name = st.selectbox(
+            "Select Variable to Plot",
+            options=variable_names,
+            index=0,
+            help="Choose the variable to visualize."
+        )
+        plot_type = st.selectbox("Select Plot Type", ["Contour", "Heatmap"], index=0)
+
+    with col2:
+        cmap = st.selectbox("Select Color Map", options=plt.colormaps(), index=plt.colormaps().index("viridis"))
+        if plot_type == "Contour":
+            num_levels = st.slider("Number of Contour Levels", 5, 100, 20, 5)
+        else:
+            num_levels = None
+
+    variable_index = variable_names.index(selected_variable_name)
+    selected_variable = out_tensor[0, variable_index].cpu().numpy()
+    lat = np.linspace(-90, 90, selected_variable.shape[0])
+    lon = np.linspace(-180, 180, selected_variable.shape[1])
+    X, Y = np.meshgrid(lon, lat)
+
+    st.markdown(f"### Plot of {selected_variable_name}")
+    fig, ax = plt.subplots(figsize=(10, 6))
+    if plot_type == "Contour":
+        contour = ax.contourf(X, Y, selected_variable, levels=num_levels, cmap=cmap)
+    else:
+        contour = ax.imshow(selected_variable, extent=[-180, 180, -90, 90], cmap=cmap, origin='lower', aspect='auto')
+
+    cbar = plt.colorbar(contour, ax=ax)
+    cbar.set_label(f'{selected_variable_name}', fontsize=12)
+    ax.set_xlabel("Longitude", fontsize=12)
+    ax.set_ylabel("Latitude", fontsize=12)
+    ax.set_title(selected_variable_name, fontsize=14)
+    st.pyplot(fig)
+
+    # Plotly interactive plot
+    st.markdown("#### Interactive Plot")
+    if plot_type == "Contour":
+        fig_plotly = go.Figure(data=go.Contour(
+            z=selected_variable,
+            x=lon,
+            y=lat,
+            colorscale=cmap,
+            contours=dict(coloring='fill', showlabels=True, labelfont=dict(size=12, color='white'), ncontours=num_levels)
+        ))
+    else:
+        fig_plotly = go.Figure(data=go.Heatmap(z=selected_variable, x=lon, y=lat, colorscale=cmap))
+
+    fig_plotly.update_layout(
+        xaxis_title="Longitude",
+        yaxis_title="Latitude",
+        width=800,
+        height=600,
+    )
+    st.plotly_chart(fig_plotly)
+
+
+def plot_aurora_output(preds, ds_subset):
+    if preds is None or ds_subset is None:
+        st.error("No predictions or dataset subset available for visualization.")
+        return
+
+    try:
+        levels = preds[0].atmos_vars["t"].shape[2]
+    except:
+        st.error("Could not determine available levels in the data.")
+        return
+
+    selected_level = st.slider('Select Level', 0, levels - 1, 11, 1)
+
+    for idx, pred in enumerate(preds):
+        pred_time = pred.metadata.time[0]
+
+        try:
+            pred_data = pred.atmos_vars["t"][0][0][selected_level].numpy() - 273.15
+            truth_data = ds_subset.T.isel(lev=selected_level)[idx].values - 273.15
+        except Exception as e:
+            st.error("Error extracting data for plotting:")
+            st.error(e)
+            continue
+
+        lat = np.array(pred.metadata.lat)
+        lon = np.array(pred.metadata.lon)
+
+        fig, axes = plt.subplots(1, 3, figsize=(18, 6), subplot_kw={'projection': ccrs.PlateCarree()})
+
+        # Ground Truth
+        im1 = axes[0].imshow(
+            truth_data, extent=[lon.min(), lon.max(), lat.min(), lat.max()],
+            origin='lower', cmap='coolwarm', transform=ccrs.PlateCarree()
+        )
+        axes[0].set_title(f"Ground Truth at Level {selected_level} - {pred_time}")
+        plt.colorbar(im1, ax=axes[0], orientation='horizontal', pad=0.05)
+
+        # Prediction
+        im2 = axes[1].imshow(
+            pred_data, extent=[lon.min(), lon.max(), lat.min(), lat.max()],
+            origin='lower', cmap='coolwarm', transform=ccrs.PlateCarree()
+        )
+        axes[1].set_title(f"Prediction at Level {selected_level} - {pred_time}")
+        plt.colorbar(im2, ax=axes[1], orientation='horizontal', pad=0.05)
+
+        plt.tight_layout()
+        st.pyplot(fig)

prithvi_utils.py

@@ -0,0 +1,145 @@
+import streamlit as st
+import tempfile
+from pathlib import Path
+import torch
+import traceback
+import yaml
+# from Prithvi import PrithviWxC, Merra2Dataset, input_scalers, output_scalers, static_input_scalers, preproc
+
+def prithvi_config_ui():
+    st.subheader("Prithvi Model Configuration")
+    param1 = st.number_input("Prithvi Parameter 1", value=10, step=1)
+    param2 = st.text_input("Prithvi Parameter 2", value="default_prithvi")
+
+    config = {"param1": param1, "param2": param2}
+
+    st.markdown("### Upload Data Files for Prithvi Model")
+    uploaded_surface_files = st.file_uploader(
+        "Upload Surface Data Files",
+        type=["nc", "netcdf"],
+        accept_multiple_files=True,
+        key="surface_uploader",
+    )
+
+    uploaded_vertical_files = st.file_uploader(
+        "Upload Vertical Data Files",
+        type=["nc", "netcdf"],
+        accept_multiple_files=True,
+        key="vertical_uploader",
+    )
+
+    st.markdown("### Upload Climatology Files (If Missing)")
+    default_clim_dir = Path("Prithvi-WxC/examples/climatology")
+    surf_in_scal_path = default_clim_dir / "musigma_surface.nc"
+    vert_in_scal_path = default_clim_dir / "musigma_vertical.nc"
+    surf_out_scal_path = default_clim_dir / "anomaly_variance_surface.nc"
+    vert_out_scal_path = default_clim_dir / "anomaly_variance_vertical.nc"
+    clim_files_exist = all([
+        surf_in_scal_path.exists(),
+        vert_in_scal_path.exists(),
+        surf_out_scal_path.exists(),
+        vert_out_scal_path.exists(),
+    ])
+
+    if not clim_files_exist:
+        st.warning("Climatology files are missing.")
+        uploaded_clim_surface = st.file_uploader(
+            "Upload Climatology Surface File",
+            type=["nc", "netcdf"],
+            key="clim_surface_uploader",
+        )
+        uploaded_clim_vertical = st.file_uploader(
+            "Upload Climatology Vertical File",
+            type=["nc", "netcdf"],
+            key="clim_vertical_uploader",
+        )
+        if uploaded_clim_surface and uploaded_clim_vertical:
+            clim_temp_dir = tempfile.mkdtemp()
+            clim_surf_path = Path(clim_temp_dir) / uploaded_clim_surface.name
+            with open(clim_surf_path, "wb") as f:
+                f.write(uploaded_clim_surface.getbuffer())
+            clim_vert_path = Path(clim_temp_dir) / uploaded_clim_vertical.name
+            with open(clim_vert_path, "wb") as f:
+                f.write(uploaded_clim_vertical.getbuffer())
+            st.success("Climatology files uploaded and saved.")
+        else:
+            st.warning("Please upload both climatology surface and vertical files.")
+            clim_surf_path, clim_vert_path = None, None
+    else:
+        clim_surf_path = surf_in_scal_path
+        clim_vert_path = vert_in_scal_path
+
+    uploaded_config = st.file_uploader(
+        "Upload config.yaml",
+        type=["yaml", "yml"],
+        key="config_uploader",
+    )
+
+    if uploaded_config:
+        temp_config = tempfile.mktemp(suffix=".yaml")
+        with open(temp_config, "wb") as f:
+            f.write(uploaded_config.getbuffer())
+        config_path = Path(temp_config)
+        st.success("Config.yaml uploaded and saved.")
+    else:
+        config_path = Path("Prithvi-WxC/examples/config.yaml")
+        if not config_path.exists():
+            st.error("Default config.yaml not found. Please upload a config file.")
+            st.stop()
+
+    uploaded_weights = st.file_uploader(
+        "Upload Model Weights (.pt)",
+        type=["pt"],
+        key="weights_uploader",
+    )
+
+    if uploaded_weights:
+        temp_weights = tempfile.mktemp(suffix=".pt")
+        with open(temp_weights, "wb") as f:
+            f.write(uploaded_weights.getbuffer())
+        weights_path = Path(temp_weights)
+        st.success("Model weights uploaded and saved.")
+    else:
+        weights_path = Path("Prithvi-WxC/examples/weights/prithvi.wxc.2300m.v1.pt")
+        if not weights_path.exists():
+            st.error("Default model weights not found. Please upload model weights.")
+            st.stop()
+
+    return config, uploaded_surface_files, uploaded_vertical_files, clim_surf_path, clim_vert_path, config_path, weights_path
+
+
+def initialize_prithvi_model(config, config_path, weights_path, device):
+    # Load the configuration
+    with open(config_path, "r") as f:
+        cfg = yaml.safe_load(f)
+
+    # Validate and load scalers, etc.
+    # Insert your logic here (loading scalers, etc.)
+    # Example (pseudo-code):
+    # in_mu, in_sig = input_scalers(...)
+    # output_sig = output_scalers(...)
+    # static_mu, static_sig = static_input_scalers(...)
+
+    # from Prithvi import PrithviWxC
+    # model = PrithviWxC(**cfg["params"], ...)
+    # state_dict = torch.load(weights_path, map_location=device)
+    # model.load_state_dict(state_dict["model_state"] if "model_state" in state_dict else state_dict, strict=True)
+    # model.to(device)
+
+    # Placeholder returns until actual logic is implemented
+    model = None
+    in_mu, in_sig, output_sig, static_mu, static_sig = None, None, None, None, None
+    return model, in_mu, in_sig, output_sig, static_mu, static_sig
+
+
+def prepare_prithvi_batch(uploaded_surface_files, uploaded_vertical_files, clim_surf_path, clim_vert_path, device):
+    # Prepare your dataset and batch for Prithvi inference
+    # dataset = Merra2Dataset(...)
+    # data = next(iter(dataset))
+    # batch = preproc([data], padding={...})
+    # for k,v in batch.items():
+    #     if isinstance(v, torch.Tensor):
+    #         batch[k] = v.to(device)
+
+    # Placeholder until implemented
+    return None