pycontrails 0.58.0__cp314-cp314-macosx_11_0_arm64.whl

pycontrails/datalib/ecmwf/hres.py

@@ -0,0 +1,804 @@
+"""ECWMF HRES forecast data access."""
+
+from __future__ import annotations
+
+import hashlib
+import logging
+import pathlib
+import sys
+from contextlib import ExitStack
+from datetime import datetime
+from typing import TYPE_CHECKING, Any
+
+if sys.version_info >= (3, 12):
+    from typing import override
+else:
+    from typing_extensions import override
+
+LOG = logging.getLogger(__name__)
+
+import pandas as pd
+import xarray as xr
+
+import pycontrails
+from pycontrails.core import cache
+from pycontrails.core.met import MetDataset, MetVariable
+from pycontrails.datalib._met_utils import metsource
+from pycontrails.datalib.ecmwf.common import ECMWFAPI
+from pycontrails.datalib.ecmwf.variables import PRESSURE_LEVEL_VARIABLES, SURFACE_VARIABLES
+from pycontrails.utils import dependencies, iteration, temp
+from pycontrails.utils.types import DatetimeLike
+
+if TYPE_CHECKING:
+    from ecmwfapi import ECMWFService
+
+
+def get_forecast_filename(
+    forecast_time: datetime,
+    timestep: datetime,
+    cc: str = "A1",
+    S: str | None = None,
+    E: str = "1",
+) -> str:
+    """Create forecast filename from ECMWF dissemination products.
+
+    See `ECMWF Dissemination <https://confluence.ecmwf.int/x/0EykDQ>`_ for more information::
+
+        The following dissemination filename convention is used for the
+        transmission of ECMWF dissemination products:
+
+        ccSMMDDHHIImmddhhiiE where:
+
+            cc is Dissemination stream name
+            S is dissemination data stream indicator
+            MMDDHHII is month, day, hour and minute on which the products are based
+            mmddhhii is month, day, hour and minute on which the products are valid at
+                ddhhii is set to “______” for Seasonal Forecasting System products
+                ii is set to 01 for high resolution forecast time step zero, type=fc, step=0
+            E is the Experiment Version Number' (as EXPVER in MARS, normally 1)
+
+    Parameters
+    ----------
+    forecast_time : datetime
+        Forecast time to stage
+    timestep : datetime
+        Time within forecast
+    cc : str, optional
+        Dissemination stream name.
+        Defaults to "A1"
+    S : str, optional
+        Dissemination data stream indicator. If None, S is set to "S" for 6 or 18 hour forecast
+        and "D" for 0 or 12 hour forecast.
+    E : str, optional
+        Experiment Version Number.
+        Defaults to "1"
+
+    Returns
+    -------
+    str
+        Filename to forecast file
+
+    Raises
+    ------
+    ValueError
+        If ``forecast_time`` is not on a synoptic hour (00, 06, 12, 18)
+        If ``timestep`` is before ``forecast_time``
+    """
+
+    if forecast_time.hour % 6 != 0:
+        raise ValueError("Forecast time must have hour 0, 6, 12, or 18")
+
+    if timestep < forecast_time:
+        raise ValueError("Forecast timestep must be on or after forecast time")
+
+    if S is None:
+        S = "D" if forecast_time.hour in (0, 12) else "S"
+
+    forecast_time_str = forecast_time.strftime("%m%d%H%M")
+    timestep_str = timestep.strftime("%m%d%H")
+
+    # for some reason "ii" is set to 01 for the first forecast timestep
+    ii = "01" if forecast_time == timestep else "00"
+
+    return f"{cc}{S}{forecast_time_str}{timestep_str}{ii}{E}"
+
+
+class HRES(ECMWFAPI):
+    """Class to support HRES data access, download, and organization.
+
+    Requires account with ECMWF and API key.
+
+    API credentials set in local ``~/.ecmwfapirc`` file:
+
+    .. code:: json
+
+        {
+            "url": "https://api.ecmwf.int/v1",
+            "email": "<email>",
+            "key": "<key>"
+        }
+
+    Credentials can also be provided directly ``url`` ``key``, and ``email`` keyword args.
+    A third option is to set the environment variables ``ECMWF_API_URL``, ``ECMWF_API_KEY``,
+    and ``ECMWF_API_EMAIL``.
+
+    See `ecmwf-api-client <https://github.com/ecmwf/ecmwf-api-client>`_ documentation
+    for more information.
+
+    Parameters
+    ----------
+    time : metsource.TimeInput | None
+        The time range for data retrieval, either a single datetime or (start, end) datetime range.
+        Input must be a datetime-like or tuple of datetime-like
+        (datetime, :class:`pandas.Timestamp`, :class:`numpy.datetime64`)
+        specifying the (start, end) of the date range, inclusive.
+        If ``forecast_time`` is unspecified, the forecast reference time will
+        be assumed to be the nearest synoptic hour: 00, 06, 12, 18.
+        All subsequent times will be downloaded for relative to :attr:`forecast_time`.
+        If None, ``paths`` must be defined and all time coordinates will be loaded from files.
+    variables : metsource.VariableInput
+        Variable name (i.e. "air_temperature", ["air_temperature, relative_humidity"])
+        See :attr:`pressure_level_variables` for the list of available variables.
+    pressure_levels : metsource.PressureLevelInput, optional
+        Pressure levels for data, in hPa (mbar)
+        Set to -1 for to download surface level parameters.
+        Defaults to -1.
+    paths : str | list[str] | pathlib.Path | list[pathlib.Path] | None, optional
+        Path to CDS NetCDF files to load manually.
+        Can include glob patterns to load specific files.
+        Defaults to None, which looks for files in the :attr:`cachestore` or CDS.
+    grid : float, optional
+        Specify latitude/longitude grid spacing in data.
+        Defaults to 0.25.
+    stream : str, optional
+        - "oper" = high resolution forecast, atmospheric fields, run at hours 00Z and 12Z
+        - "scda" = short cut-off high resolution forecast, atmospheric fields,
+          run at hours 06Z and 18Z
+        - "enfo" = ensemble forecast, atmospheric fields, run at hours 00Z, 06Z, 12Z, and 18Z
+
+        Defaults to "oper" (HRES).
+        If the stream is incompatible with a provided forecast_time, a ``ValueError`` is raised.
+        See the `ECMWF documentation <https://confluence.ecmwf.int/display/DAC/ECMWF+open+data%3A+real-time+forecasts+from+IFS+and+AIFS>`_
+        for additional information.
+    field_type : str, optional
+        Field type can be e.g. forecast (fc), perturbed forecast (pf),
+        control forecast (cf), analysis (an).
+        Defaults to "fc".
+    forecast_time : DatetimeLike, optional
+        Specify forecast reference time (the time at which the forecast was initialized).
+        Defaults to None.
+    cachestore : cache.CacheStore | None, optional
+        Cache data store for staging data files.
+        Defaults to :class:`cache.DiskCacheStore`.
+        If None, cache is turned off.
+    url : str
+        Override `ecmwf-api-client <https://github.com/ecmwf/ecmwf-api-client>`_ url
+    key : str
+        Override `ecmwf-api-client <https://github.com/ecmwf/ecmwf-api-client>`_ key
+    email : str
+        Override `ecmwf-api-client <https://github.com/ecmwf/ecmwf-api-client>`_ email
+
+    Notes
+    -----
+    `MARS key word definitions <https://confluence.ecmwf.int/display/UDOC/Identification+keywords>`_
+
+    - `class <https://apps.ecmwf.int/codes/grib/format/mars/class/>`_:
+      in most cases this will be operational data, or "od"
+    - `stream <https://apps.ecmwf.int/codes/grib/format/mars/stream/>`_:
+      "enfo" = ensemble forecast, "oper" = atmospheric model/HRES
+    - `expver <https://confluence.ecmwf.int/pages/viewpage.action?pageId=124752178>`_:
+      experimental version, production data is 1 or 2
+    - `date <https://confluence.ecmwf.int/pages/viewpage.action?pageId=118817289>`_:
+      there are numerous acceptible date formats
+    - `time <https://confluence.ecmwf.int/pages/viewpage.action?pageId=118817378>`_:
+      forecast base time, always in synoptic time (0,6,12,18 UTC)
+    - `type <https://confluence.ecmwf.int/pages/viewpage.action?pageId=127315300>`_:
+      forecast (oper), perturbed or control forecast (enfo only), or analysis
+    - `levtype <https://confluence.ecmwf.int/pages/viewpage.action?pageId=149335319>`_:
+      options include surface, pressure levels, or model levels
+    - `levelist <https://confluence.ecmwf.int/pages/viewpage.action?pageId=149335403>`_:
+      list of levels in format specified by **levtype** `levelist`_
+    - `param <https://confluence.ecmwf.int/pages/viewpage.action?pageId=149335858>`_:
+      list of variables in catalog number, long name or short name
+    - `step <https://confluence.ecmwf.int/pages/viewpage.action?pageId=118820050>`_:
+      hourly time steps from base forecast time
+    - `number <https://confluence.ecmwf.int/pages/viewpage.action?pageId=149335478>`_:
+      for ensemble forecasts, ensemble numbers
+    - `format <https://confluence.ecmwf.int/pages/viewpage.action?pageId=116970058>`_:
+      specify netcdf instead of default grib, DEPRECATED `format`_
+    - `grid <https://confluence.ecmwf.int/pages/viewpage.action?pageId=123799065>`_:
+      specify model return grid spacing
+
+    Local ``paths`` are loaded using :func:`xarray.open_mfdataset`.
+    Pass ``xr_kwargs`` inputs to :meth:`open_metdataset` to customize file loading.
+
+
+    Examples
+    --------
+    >>> from datetime import datetime
+    >>> from pycontrails import GCPCacheStore
+    >>> from pycontrails.datalib.ecmwf import HRES
+
+    >>> # Store data files to local disk (default behavior)
+    >>> times = (datetime(2021, 5, 1, 2), datetime(2021, 5, 1, 3))
+    >>> hres = HRES(times, variables="air_temperature", pressure_levels=[300, 250])
+
+    >>> # Cache files to google cloud storage
+    >>> gcp_cache = GCPCacheStore(
+    ...     bucket="contrails-301217-unit-test",
+    ...     cache_dir="ecmwf",
+    ... )
+    >>> hres = HRES(
+    ...     times,
+    ...     variables="air_temperature",
+    ...     pressure_levels=[300, 250],
+    ...     cachestore=gcp_cache
+    ... )
+    """
+
+    __slots__ = ("email", "field_type", "forecast_time", "key", "server", "stream", "url")
+
+    #: stream type, "oper" or "scda" for atmospheric model/HRES, "enfo" for ensemble forecast.
+    stream: str
+
+    #: Field type, forecast ("fc"), perturbed forecast ("pf"),
+    #: control forecast ("cf"), analysis ("an").
+    field_type: str
+
+    #: Handle to ECMWFService client
+    server: ECMWFService
+
+    #: Forecast run time, either specified or assigned by the closest previous forecast run
+    forecast_time: datetime
+
+    __marker = object()
+
+    def __init__(
+        self,
+        time: metsource.TimeInput | None,
+        variables: metsource.VariableInput,
+        pressure_levels: metsource.PressureLevelInput = -1,
+        paths: str | list[str] | pathlib.Path | list[pathlib.Path] | None = None,
+        grid: float = 0.25,
+        stream: str = "oper",
+        field_type: str = "fc",
+        forecast_time: DatetimeLike | None = None,
+        cachestore: cache.CacheStore | None = __marker,  # type: ignore[assignment]
+        url: str | None = None,
+        key: str | None = None,
+        email: str | None = None,
+    ) -> None:
+        try:
+            from ecmwfapi import ECMWFService
+        except ModuleNotFoundError as e:
+            dependencies.raise_module_not_found_error(
+                name="HRES class",
+                package_name="ecmwf-api-client",
+                module_not_found_error=e,
+                pycontrails_optional_package="ecmwf",
+            )
+
+        # ERA5 now delays creating the server attribute until it is needed to download
+        # from CDS. We could do the same here.
+        self.server = ECMWFService("mars", url=url, key=key, email=email)
+        self.paths = paths
+
+        self.cachestore = cache.DiskCacheStore() if cachestore is self.__marker else cachestore
+
+        if time is None and paths is None:
+            raise ValueError("Time input is required when paths is None")
+
+        self.timesteps = metsource.parse_timesteps(time, freq="1h")
+        self.pressure_levels = metsource.parse_pressure_levels(
+            pressure_levels, self.supported_pressure_levels
+        )
+        self.variables = metsource.parse_variables(variables, self.supported_variables)
+
+        self.grid = metsource.parse_grid(grid, [0.1, 0.25, 0.5, 1])  # lat/lon degree resolution
+
+        # "fc" = forecast
+        # "pf" = perturbed forecast
+        # "cf" = control forecast
+        # "an" = analysis
+        if field_type not in ("fc", "pf", "cf", "an"):
+            msg = "Parameter field_type must be 'fc', 'pf', 'cf', or 'an'"
+            raise ValueError(msg)
+
+        self.field_type = field_type
+
+        # set specific forecast time is requested
+        if forecast_time is not None:
+            forecast_time_pd = pd.to_datetime(forecast_time)
+            if forecast_time_pd.hour % 6:
+                raise ValueError("Forecast hour must be on one of 00, 06, 12, 18")
+
+            self.forecast_time = metsource.round_hour(forecast_time_pd.to_pydatetime(), 6)
+
+        # if no specific forecast is requested, set the forecast time using timesteps
+        elif self.timesteps:
+            # round first element to the nearest 6 hour time (00, 06, 12, 18 UTC) for forecast_time
+            self.forecast_time = metsource.round_hour(self.timesteps[0], 6)
+
+        # NOTE: when no forecast_time or time input, forecast_time is defined in _open_and_cache
+        # This could occur when only the paths parameter is provided
+
+        # "enfo" = ensemble forecast
+        # "oper" = atmospheric model/HRES for 00 and 12 model runs
+        # "scda" = atmospheric model/HRES for 06 and 18 model runs
+        available_streams = ("oper", "enfo", "scda")
+        if stream not in available_streams:
+            msg = f"Parameter stream must be one of {available_streams}"
+            raise ValueError(msg)
+
+        if self.forecast_time.hour in (0, 12) and stream == "scda":
+            raise ValueError(
+                f"Stream {stream} is not compatible with forecast_time {self.forecast_time}. "
+                "Set stream='oper' for 00 and 12 UTC forecast times."
+            )
+
+        if self.forecast_time.hour in (6, 18) and stream == "oper":
+            raise ValueError(
+                f"Stream {stream} is not compatible with forecast_time {self.forecast_time}. "
+                "Set stream='scda' for 06 and 18 UTC forecast times."
+            )
+        self.stream = stream
+
+    def __repr__(self) -> str:
+        base = super().__repr__()
+        return (
+            f"{base}\n\tForecast time: {getattr(self, 'forecast_time', '')}\n\tSteps:"
+            f" {getattr(self, 'steps', '')}"
+        )
+
+    @staticmethod
+    def create_synoptic_time_ranges(
+        timesteps: list[pd.Timestamp],
+    ) -> list[tuple[pd.Timestamp, pd.Timestamp]]:
+        """Create synoptic time bounds encompassing date range.
+
+        Extracts time bounds for synoptic time range ([00:00, 11:59], [12:00, 23:59])
+        for a list of input timesteps.
+
+        Parameters
+        ----------
+        timesteps : list[pd.Timestamp]
+            List of timesteps formatted as :class:`pd.Timestamps`.
+            Often this it the output from `pd.date_range()`
+
+        Returns
+        -------
+        list[tuple[pd.Timestamp, pd.Timestamp]]
+            List of tuple time bounds that can be used as inputs to :class:`HRES(time=...)`
+        """
+        time_ranges = sorted({t.floor("12h") for t in timesteps})
+
+        if len(time_ranges) == 1:
+            return [(timesteps[0], timesteps[-1])]
+
+        time_ranges[0] = (timesteps[0], time_ranges[1] - pd.Timedelta(hours=1))
+        time_ranges[1:-1] = [(t, t + pd.Timedelta(hours=11)) for t in time_ranges[1:-1]]
+        time_ranges[-1] = (time_ranges[-1], timesteps[-1])
+
+        return time_ranges
+
+    @property
+    def hash(self) -> str:
+        """Generate a unique hash for this datasource.
+
+        Returns
+        -------
+        str
+            Unique hash for met instance (sha1)
+        """
+        hashstr = (
+            f"{self.__class__.__name__}{self.timesteps}{self.variable_shortnames}"
+            f"{self.pressure_levels}{self.grid}{self.forecast_time}{self.field_type}{self.stream}"
+        )
+        return hashlib.sha1(bytes(hashstr, "utf-8")).hexdigest()
+
+    @property
+    def pressure_level_variables(self) -> list[MetVariable]:
+        """ECMWF pressure level parameters.
+
+        Returns
+        -------
+        list[MetVariable] | None
+            List of MetVariable available in datasource
+        """
+        return PRESSURE_LEVEL_VARIABLES
+
+    @property
+    def single_level_variables(self) -> list[MetVariable]:
+        """ECMWF surface level parameters.
+
+        Returns
+        -------
+        list[MetVariable] | None
+            List of MetVariable available in datasource
+        """
+        return SURFACE_VARIABLES
+
+    @property
+    def supported_pressure_levels(self) -> list[int]:
+        """Get pressure levels available from MARS.
+
+        Returns
+        -------
+        list[int]
+            List of integer pressure level values
+        """
+        return [
+            1000,
+            950,
+            925,
+            900,
+            850,
+            800,
+            700,
+            600,
+            500,
+            400,
+            300,
+            250,
+            200,
+            150,
+            100,
+            70,
+            50,
+            30,
+            20,
+            10,
+            7,
+            5,
+            3,
+            2,
+            1,
+            -1,
+        ]
+
+    @property
+    def step_offset(self) -> int:
+        """Difference between :attr:`forecast_time` and first timestep.
+
+        Returns
+        -------
+        int
+            Number of steps to offset in order to retrieve data starting from input time.
+            Returns 0 if :attr:`timesteps` is empty when loading from :attr:`paths`.
+        """
+        if self.timesteps:
+            return int((self.timesteps[0] - self.forecast_time).total_seconds() // 3600)
+
+        return 0
+
+    @property
+    def steps(self) -> list[int]:
+        """Forecast steps from :attr:`forecast_time` corresponding within input :attr:`time`.
+
+        Returns
+        -------
+        list[int]
+            List of forecast steps relative to :attr:`forecast_time`
+        """
+        return [self.step_offset + i for i in range(len(self.timesteps))]
+
+    def list_from_mars(self) -> str:
+        """List metadata on query from MARS.
+
+        Returns
+        -------
+        str
+            Metadata for MARS request.
+            Note this is queued the same as data requests.
+        """
+        request = self.generate_mars_request(self.forecast_time, self.steps, request_type="list")
+
+        # hold downloaded file in named temp file
+        with temp.temp_file() as mars_temp_filename:
+            LOG.debug("Performing MARS request: %s", request)
+            self.server.execute(request, target=mars_temp_filename)
+
+            with open(mars_temp_filename) as f:
+                return f.read()
+
+    def generate_mars_request(
+        self,
+        forecast_time: datetime | None = None,
+        steps: list[int] | None = None,
+        request_type: str = "retrieve",
+        request_format: str = "mars",
+    ) -> str | dict[str, Any]:
+        """Generate MARS request in MARS request syntax.
+
+        Parameters
+        ----------
+        forecast_time : :class:`datetime`, optional
+            Base datetime for the forecast.
+            Defaults to :attr:`forecast_time`.
+        steps : list[int], optional
+            list of steps.
+            Defaults to :attr:`steps`.
+        request_type : str, optional
+            "retrieve" for download request or "list" for metadata request.
+            Defaults to "retrieve".
+        request_format : str, optional
+            "mars" for MARS string format, or "dict" for dict version.
+            Defaults to "mars".
+
+        Returns
+        -------
+        str | dict[str, Any]
+            Returns MARS query string if ``request_format`` is "mars".
+            Returns dict query if ``request_format`` is "dict"
+
+        Notes
+        -----
+        Brief overview of `MARS request syntax
+        <https://confluence.ecmwf.int/display/WEBAPI/Brief+MARS+request+syntax>`_
+        """
+
+        if forecast_time is None:
+            forecast_time = self.forecast_time
+
+        if steps is None:
+            steps = self.steps
+
+        # set date/time for file
+        date = forecast_time.strftime("%Y%m%d")
+        time = forecast_time.strftime("%H")
+
+        # make request of mars
+        request: dict[str, Any] = {
+            "class": "od",  # operational data
+            "stream": self.stream,
+            "expver": "1",  # production data only
+            "date": date,
+            "time": time,
+            "type": self.field_type,
+            "param": f"{'/'.join(self.variable_shortnames)}",
+            "step": f"{'/'.join([str(s) for s in steps])}",
+            "grid": f"{self.grid}/{self.grid}",
+        }
+
+        if self.pressure_levels != [-1]:
+            request["levtype"] = "pl"
+            request["levelist"] = f"{'/'.join([str(pl) for pl in self.pressure_levels])}"
+        else:
+            request["levtype"] = "sfc"
+
+        if request_format == "dict":
+            return request
+
+        levelist = f",\n\tlevelist={request['levelist']}" if self.pressure_levels != [-1] else ""
+        return (
+            f"{request_type},\n\tclass={request['class']},\n\tstream={request['stream']},"
+            f"\n\texpver={request['expver']},\n\tdate={request['date']},"
+            f"\n\ttime={request['time']},\n\ttype={request['type']},"
+            f"\n\tparam={request['param']},\n\tstep={request['step']},"
+            f"\n\tgrid={request['grid']},\n\tlevtype={request['levtype']}{levelist}"
+        )
+
+    @override
+    def create_cachepath(self, t: datetime) -> str:
+        if self.cachestore is None:
+            raise ValueError("self.cachestore attribute must be defined to create cache path")
+
+        # get forecast_time and step for specific file
+        datestr = self.forecast_time.strftime("%Y%m%d-%H")
+
+        # get step relative to forecast forecast_time
+        step = self.step_offset + self.timesteps.index(t)
+
+        # single level or pressure level
+        levtype = "sl" if self.pressure_levels == [-1] else "pl"
+        suffix = f"hres{levtype}{self.grid}{self.stream}{self.field_type}"
+
+        # return cache path
+        return self.cachestore.path(f"{datestr}-{step}-{suffix}.nc")
+
+    @override
+    def download_dataset(self, times: list[datetime]) -> None:
+        """Download data from data source for input times.
+
+        Parameters
+        ----------
+        times : list[:class:`datetime`]
+            List of datetimes to download and store in cache datastore
+        """
+
+        # get step relative to forecast forecast_time
+        steps = [self.step_offset + self.timesteps.index(t) for t in times]
+        LOG.debug(f"Downloading HRES dataset for base time {self.forecast_time} and steps {steps}")
+
+        # download in sets of 24
+        if len(steps) > 24:
+            for _steps in iteration.chunk_list(steps, 24):
+                self._download_file(_steps)
+        elif len(steps) > 0:
+            self._download_file(steps)
+
+    @override
+    def open_metdataset(
+        self,
+        dataset: xr.Dataset | None = None,
+        xr_kwargs: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> MetDataset:
+        xr_kwargs = xr_kwargs or {}
+
+        #  short-circuit dataset or file paths if provided
+        if dataset is not None:
+            ds = self._preprocess_hres_dataset(dataset)
+
+        # load from local paths
+        elif self.paths is not None:
+            ds = self._open_and_cache(xr_kwargs)
+
+        # download from MARS
+        else:
+            if self.cachestore is None:
+                raise ValueError("Cachestore is required to download data")
+
+            # confirm files are downloaded from CDS or MARS
+            self.download(**xr_kwargs)
+
+            # ensure all files are guaranteed to be available locally here
+            # this would download a file from a remote (e.g. GCP) cache
+            disk_cachepaths = [self.cachestore.get(f) for f in self._cachepaths]
+
+            # open cache files as xr.Dataset
+            ds = self.open_dataset(disk_cachepaths, **xr_kwargs)
+
+            ds.attrs.setdefault("pycontrails_version", pycontrails.__version__)
+
+        # run the same ECMWF-specific processing on the dataset
+        mds = self._process_dataset(ds, **kwargs)
+
+        self.set_metadata(mds)
+        return mds
+
+    @override
+    def set_metadata(self, ds: xr.Dataset | MetDataset) -> None:
+        if self.stream in ("oper", "scda"):
+            product = "forecast"
+        elif self.stream == "enfo":
+            product = "ensemble"
+        else:
+            msg = f"Unknown stream type {self.stream}"
+            raise ValueError(msg)
+
+        ds.attrs.update(
+            provider="ECMWF",
+            dataset="HRES",
+            product=product,
+            radiation_accumulated=True,
+        )
+
+    def _open_and_cache(self, xr_kwargs: dict[str, Any]) -> xr.Dataset:
+        """Open and cache :class:`xr.Dataset` from :attr:`self.paths`.
+
+        Parameters
+        ----------
+        xr_kwargs : dict[str, Any]
+            Additional kwargs passed directly to :func:`xarray.open_mfdataset`.
+            See :meth:`open_metdataset`.
+
+        Returns
+        -------
+        xr.Dataset
+            Dataset opened from local paths.
+        """
+
+        if self.paths is None:
+            raise ValueError("Attribute `self.paths` must be defined to open and cache")
+
+        # if timesteps are defined and all timesteps are cached already
+        # then we can skip loading
+        if self.timesteps and self.cachestore and not self.list_timesteps_not_cached(**xr_kwargs):
+            LOG.debug("All timesteps already in cache store")
+            disk_cachepaths = [self.cachestore.get(f) for f in self._cachepaths]
+            return self.open_dataset(disk_cachepaths, **xr_kwargs)
+
+        # set default parameters for loading grib files
+        xr_kwargs.setdefault("engine", "cfgrib")
+        xr_kwargs.setdefault("combine", "nested")
+        xr_kwargs.setdefault("concat_dim", "step")
+        xr_kwargs.setdefault("parallel", False)
+        ds = self.open_dataset(self.paths, **xr_kwargs)
+
+        # set forecast time if it's not defined (this occurs when only the paths param is provided)
+        if not hasattr(self, "forecast_time"):
+            self.forecast_time = ds["time"].values.astype("datetime64[s]").tolist()
+
+        # check that forecast_time is correct if defined
+        # note the "time" coordinate here is the HRES forecast_time
+        elif self.forecast_time != ds["time"].values.astype("datetime64[s]").tolist():
+            raise ValueError(
+                f"HRES.forecast_time {self.forecast_time} is not the same forecast time listed"
+                " in file"
+            )
+
+        ds = self._preprocess_hres_dataset(ds)
+
+        # set timesteps if not defined
+        # note that "time" is now the actual timestep coordinates
+        if not self.timesteps:
+            self.timesteps = ds["time"].values.astype("datetime64[s]").tolist()
+
+        self.cache_dataset(ds)
+
+        return ds
+
+    def _download_file(self, steps: list[int]) -> None:
+        """Download data file for base datetime and timesteps.
+
+        Overwrites files if they already exists.
+
+        Parameters
+        ----------
+        steps : list[int]
+            Steps to download relative to base date
+        """
+        request = self.generate_mars_request(self.forecast_time, steps)
+
+        # Open ExitStack to control temp_file context manager
+        with ExitStack() as stack:
+            # hold downloaded file in named temp file
+            mars_temp_grib_filename = stack.enter_context(temp.temp_file())
+
+            # retrieve data from MARS
+            LOG.debug(f"Performing MARS request: {request}")
+            self.server.execute(request, mars_temp_grib_filename)
+
+            # translate into netcdf from grib
+            LOG.debug("Translating file into netcdf")
+            ds = stack.enter_context(xr.open_dataset(mars_temp_grib_filename, engine="cfgrib"))
+
+            # run preprocessing before cache
+            ds = self._preprocess_hres_dataset(ds)
+
+            self.cache_dataset(ds)
+
+    def _preprocess_hres_dataset(self, ds: xr.Dataset) -> xr.Dataset:
+        """Process HRES data before caching.
+
+        Parameters
+        ----------
+        ds : xr.Dataset
+            Loaded :class:`xr.Dataset`
+
+        Returns
+        -------
+        xr.Dataset
+            Processed :class:`xr.Dataset`
+        """
+
+        if "pycontrails_version" in ds.attrs:
+            LOG.debug("Input dataset processed with pycontrails > 0.29")
+            return ds
+
+        # for pressure levels, need to rename "level" field
+        if self.pressure_levels != [-1]:
+            ds = ds.rename({"isobaricInhPa": "level"})
+
+        # for single level, and singular pressure levels, add the level dimension
+        if len(self.pressure_levels) == 1:
+            ds = ds.expand_dims({"level": self.pressure_levels})
+
+        # for single time, add the step dimension and assign time coords to step
+        if ds["step"].size == 1:
+            if "step" not in ds.dims:
+                ds = ds.expand_dims({"step": [ds["step"].values]})
+
+            ds = ds.assign_coords({"valid_time": ("step", [ds["valid_time"].values])})
+
+        # rename fields and swap time dimension for step
+        ds = ds.rename({"time": "forecast_time"})
+        ds = ds.rename({"valid_time": "time"})
+        ds = ds.swap_dims({"step": "time"})
+
+        # drop step/number
+        ds = ds.drop_vars(["step", "number"], errors="ignore")
+
+        ds.attrs["pycontrails_version"] = pycontrails.__version__
+        return ds