# SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES.
# SPDX-FileCopyrightText: All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
import inspect
import os
import tempfile
import time
import typing
import weakref
from collections import OrderedDict
from collections.abc import Callable
from datetime import datetime, timedelta
from hashlib import sha256
from inspect import signature
from pathlib import Path
from shutil import rmtree
from typing import TYPE_CHECKING, Any, ClassVar, Literal, TypeVar
import numpy as np
import torch
import xarray as xr
from fsspec import filesystem
from fsspec.asyn import AsyncFileSystem
from fsspec.callbacks import DEFAULT_CALLBACK
from fsspec.implementations.cache_mapper import create_cache_mapper
from fsspec.implementations.cache_metadata import CacheMetadata
from fsspec.utils import isfilelike
from loguru import logger
from earth2studio.data.base import DataSource, ForecastSource
from earth2studio.utils.interp import LatLonInterpolation
from earth2studio.utils.time import (
leadtimearray_to_timedelta,
timearray_to_datetime,
to_time_array,
)
from earth2studio.utils.type import CoordSystem, LeadTimeArray, TimeArray, VariableArray
if TYPE_CHECKING:
from fsspec.implementations.cache_mapper import AbstractCacheMapper
[docs]
def fetch_data(
source: DataSource | ForecastSource,
time: TimeArray,
variable: VariableArray,
lead_time: LeadTimeArray = np.array([np.timedelta64(0, "h")]),
device: torch.device = "cpu",
interp_to: CoordSystem = None,
interp_method: str = "nearest",
) -> tuple[torch.Tensor, CoordSystem]:
"""Utility function to fetch data for models and load data on the target device.
If desired, xarray interpolation/regridding in the spatial domain can be used
by passing a target coordinate system via the optional `interp_to` argument.
Parameters
----------
source : DataSource
The data source to fetch from
time : TimeArray
Timestamps to return data for (UTC).
variable : VariableArray
Strings or list of strings that refer to variables to return
lead_time : LeadTimeArray, optional
Lead times to fetch for each provided time, by default
np.array(np.timedelta64(0, "h"))
device : torch.device, optional
Torch devive to load data tensor to, by default "cpu"
interp_to : CoordSystem, optional
If provided, the fetched data will be interpolated to the coordinates
specified by lat/lon arrays in this CoordSystem
interp_method : str
Interpolation method to use with xarray (by default 'nearest')
Returns
-------
tuple[torch.Tensor, CoordSystem]
Tuple containing output tensor and coordinate OrderedDict
"""
sig = signature(source.__call__)
if "lead_time" in sig.parameters:
# Working with a Forecast Data Source
da = source(time, lead_time, variable) # type: ignore
else:
da = []
for lead in lead_time:
adjust_times = np.array([t + lead for t in time], dtype="datetime64[ns]")
da0 = source(adjust_times, variable) # type: ignore
da0 = da0.expand_dims(dim={"lead_time": 1}, axis=1)
da0 = da0.assign_coords(lead_time=np.array([lead], dtype="timedelta64[ns]"))
da0 = da0.assign_coords(time=time)
da.append(da0)
da = xr.concat(da, "lead_time")
return prep_data_array(
da,
device=device,
interp_to=interp_to,
interp_method=interp_method,
)
[docs]
def prep_data_array(
da: xr.DataArray,
device: torch.device = "cpu",
interp_to: CoordSystem | None = None,
interp_method: str = "nearest",
) -> tuple[torch.Tensor, CoordSystem]:
"""Prepares a data array from a data source for inference workflows by converting
the data array to a torch tensor and the coordinate system to an OrderedDict.
If desired, xarray interpolation/regridding in the spatial domain can be used
by passing a target coordinate system via the optional `interp_to` argument.
Parameters
----------
da : xr.DataArray
Input data array
device : torch.device, optional
Torch devive to load data tensor to, by default "cpu"
interp_to : CoordSystem, optional
If provided, the fetched data will be interpolated to the coordinates
specified by lat/lon arrays in this CoordSystem
interp_method : str
Interpolation method to use with xarray (by default 'nearest')
Returns
-------
tuple[torch.Tensor, CoordSystem]
Tuple containing output tensor and coordinate OrderedDict
"""
# Initialize the output CoordSystem
out_coords = OrderedDict()
for dim in da.coords.dims:
if dim in da.coords:
out_coords[dim] = np.array(da.coords[dim])
# Fetch data and regrid if necessary
if interp_to is not None:
if len(interp_to["_lat"].shape) != len(interp_to["_lon"].shape):
raise ValueError(
"Discrepancy in interpolation coordinates: latitude has different number of dims than longitude"
)
if "lat" not in da.dims:
# Data source uses curvilinear coordinates
if interp_method != "linear":
raise ValueError(
"fetch_data does not support interpolation methods other than linear when data source has a curvilinear grid"
)
interp = LatLonInterpolation(
lat_in=da["_lat"].values,
lon_in=da["_lon"].values,
lat_out=interp_to["_lat"],
lon_out=interp_to["_lon"],
).to(device)
data = torch.Tensor(da.values).to(device)
out = interp(data)
# HARD CODE FOR STORMCAST
# TODO: FIX THIS BY CORRECTING STORMCAST COORDINATES
if "hrrr_y" in out_coords:
del out_coords["hrrr_y"]
if "hrrr_x" in out_coords:
del out_coords["hrrr_x"]
else:
if len(interp_to["_lat"].shape) > 1 or len(interp_to["_lon"].shape) > 1:
# Target grid uses curvilinear coordinates: define internal dims y, x
target_lat = xr.DataArray(interp_to["_lat"], dims=["y", "x"])
target_lon = xr.DataArray(interp_to["_lon"], dims=["y", "x"])
else:
target_lat = xr.DataArray(interp_to["_lat"], dims=["_lat"])
target_lon = xr.DataArray(interp_to["_lon"], dims=["_lon"])
da = da.interp(
lat=target_lat,
lon=target_lon,
method=interp_method,
)
out = torch.Tensor(da.values).to(device)
out_coords["_lat"] = interp_to["_lat"]
out_coords["_lon"] = interp_to["_lon"]
else:
out = torch.Tensor(da.values).to(device)
if "lat" in da.coords and "lat" not in da.coords.dims:
# Curvilinear grid case: lat/lon coords are 2D arrays, not in dims
out_coords["lat"] = da.coords["lat"].values
out_coords["lon"] = da.coords["lon"].values
else:
for dim in da.coords.dims:
if dim not in ["time", "lead_time", "variable"]:
out_coords[dim] = np.array(da.coords[dim])
return out, out_coords
def prep_data_inputs(
time: datetime | list[datetime] | TimeArray,
variable: str | list[str] | VariableArray,
) -> tuple[list[datetime], list[str]]:
"""Simple method to pre-process data source inputs into a common form
Parameters
----------
time : datetime | list[datetime] | TimeArray
Datetime, list of datetimes or array of np.datetime64 to fetch
variable : str | list[str] | VariableArray
String, list of strings or array of strings that refer to variables
Returns
-------
tuple[list[datetime], list[str]]
Time and variable lists
"""
if isinstance(variable, str):
variable = [variable]
if isinstance(time, datetime):
time = [time]
if isinstance(time, np.ndarray): # np.datetime64 -> datetime
time = timearray_to_datetime(time)
return time, variable
def prep_forecast_inputs(
time: datetime | list[datetime] | TimeArray,
lead_time: timedelta | list[timedelta] | LeadTimeArray,
variable: str | list[str] | VariableArray,
) -> tuple[list[datetime], list[timedelta], list[str]]:
"""Simple method to pre-process forecast source inputs into a common form
Parameters
----------
time : datetime | list[datetime] | TimeArray
Datetime, list of datetimes or array of np.datetime64 to fetch
lead_time: timedelta | list[timedelta], LeadTimeArray
Timedelta, list of timedeltas or array of np.timedelta to fetch
variable : str | list[str] | VariableArray
String, list of strings or array of strings that refer to variables
Returns
-------
tuple[list[datetime], list[timedelta], list[str]]
Time, lead time, and variable lists
"""
if isinstance(lead_time, timedelta):
lead_time = [lead_time]
if isinstance(lead_time, np.ndarray): # np.timedelta64 -> timedelta
lead_time = leadtimearray_to_timedelta(lead_time)
time, variable = prep_data_inputs(time, variable)
return time, lead_time, variable
[docs]
def datasource_to_file(
file_name: str,
source: DataSource,
time: list[str] | list[datetime] | TimeArray,
variable: VariableArray,
lead_time: LeadTimeArray = np.array([np.timedelta64(0, "h")]),
backend: Literal["netcdf", "zarr"] = "netcdf",
chunks: dict[str, int] = {"variable": 1},
dtype: np.dtype | None = None,
backend_kwargs: dict[str, Any] = {},
) -> None:
"""Utility function that can be used for building a local data store needed
for an inference request. This file can then be used with the
:py:class:`earth2studio.data.DataArrayFile` data source to load data from file.
This is useful when multiple runs of the same input data is needed.
Parameters
----------
file_name : str
File name of output NetCDF
source : DataSource
The original data source to fetch from
time : list[str] | list[datetime] | list[np.datetime64]
List of time strings, datetimes or np.datetime64 (UTC)
variable : VariableArray
Strings or list of strings that refer to variables to return
lead_time : LeadTimeArray, optional
Lead times to fetch for each provided time, by default
np.array(np.timedelta64(0, "h"))
backend : Literal["netcdf", "zarr"], optional
Storage backend to save output file as, by default "netcdf"
chunks : dict[str, int], optional
Chunk sizes along each dimension, by default {"variable": 1}
dtype : np.dtype, optional
Data type for storing data
backend_kwargs : dict[str, Any], optional
Dictionary of keyword arguments forwarded to the underlying
``xarray.DataArray.to_netcdf`` / ``xarray.DataArray.to_zarr``
call depending on the selected backend.
"""
if isinstance(time, datetime):
time = [time]
time = to_time_array(time)
# Spot check the write location is okay before pull
testfile = tempfile.TemporaryFile(dir=Path(file_name).parent.resolve())
testfile.close()
# Compile all times
for lead in lead_time:
adjust_times = np.array([t + lead for t in time], dtype="datetime64[ns]")
time = np.concatenate([time, adjust_times], axis=0)
time = np.unique(time)
# Fetch
da = source(time, variable)
da = da.assign_coords(time=time)
da = da.chunk(chunks=chunks)
if dtype is not None:
da = da.astype(dtype=dtype)
match backend:
case "netcdf":
da.to_netcdf(file_name, **backend_kwargs)
case "zarr":
da.to_zarr(file_name, **backend_kwargs)
case _:
raise ValueError(f"Unsupported backend {backend}")
def datasource_cache_root() -> str:
"""Returns the root directory for data sources"""
default_cache = os.path.join(os.path.expanduser("~"), ".cache", "earth2studio")
default_cache = os.environ.get("EARTH2STUDIO_CACHE", default_cache)
try:
os.makedirs(default_cache, exist_ok=True)
except OSError as e:
logger.error(
f"Failed to create cache folder {default_cache}, check permissions"
)
raise e
return default_cache
def get_msc_filesystem() -> filesystem | None:
"""This helper function checks if Multi-Storage Client is available and sets up
the MSC configuration if needed.
Note
----
Can force MSC to not be used with the environment variable EARTH2STUDIO_DISABLE_MSC
Returns
-------
filesystem | None
Returns multi-storage file system if installed, None otherwise
"""
if str(os.getenv("EARTH2STUDIO_DISABLE_MSC", "")).strip().lower() in ("1", "true"):
return None
try:
from multistorageclient.contrib.async_fs import MultiStorageAsyncFileSystem
config_path = Path(__file__).parent / "msc_config.yaml"
os.environ["MSC_CONFIG"] = str(config_path)
return MultiStorageAsyncFileSystem
except ImportError:
return None
T = TypeVar("T")
@typing.no_type_check
class AsyncCachingFileSystem(AsyncFileSystem):
"""Async locally caching filesystem, layer over any other FS for use with zarr 3.0.
Presently extremely limited, much is just copied from the WholeFileCache
Parameters
----------
target_protocol: str (optional)
Target filesystem protocol. Provide either this or ``fs``.
cache_storage: str or list(str)
Location to store files. If "TMP", this is a temporary directory,
and will be cleaned up by the OS when this process ends (or later).
If a list, each location will be tried in the order given, but
only the last will be considered writable.
cache_check: int
Number of seconds between reload of cache metadata
check_files: bool
Whether to explicitly see if the UID of the remote file matches
the stored one before using. Warning: some file systems such as
HTTP cannot reliably give a unique hash of the contents of some
path, so be sure to set this option to False.
expiry_time: int
The time in seconds after which a local copy is considered useless.
Set to falsy to prevent expiry. The default is equivalent to one
week.
target_options: dict or None
Passed to the instantiation of the FS, if fs is None.
fs: filesystem instance
The target filesystem to run against. Provide this or ``protocol``.
same_names: bool (optional)
By default, target URLs are hashed using a ``HashCacheMapper`` so
that files from different backends with the same basename do not
conflict. If this argument is ``true``, a ``BasenameCacheMapper``
is used instead. Other cache mapper options are available by using
the ``cache_mapper`` keyword argument. Only one of this and
``cache_mapper`` should be specified.
compression: str (optional)
To decompress on download. Can be 'infer' (guess from the URL name),
one of the entries in ``fsspec.compression.compr``, or None for no
decompression.
cache_mapper: AbstractCacheMapper (optional)
The object use to map from original filenames to cached filenames.
Only one of this and ``same_names`` should be specified.
"""
protocol: ClassVar[str | tuple[str, ...]] = ("blockcache", "cached")
@typing.no_type_check
def __init__(
self,
target_protocol=None,
cache_storage="TMP",
cache_check=10,
check_files=False,
expiry_time=604800,
target_options=None,
fs=None,
same_names: bool | None = None,
compression=None,
cache_mapper: AbstractCacheMapper | None = None,
**kwargs,
):
super().__init__(**kwargs)
if fs is None and target_protocol is None:
raise ValueError(
"Please provide filesystem instance(fs) or target_protocol"
)
if not (fs is None) ^ (target_protocol is None):
raise ValueError(
"Both filesystems (fs) and target_protocol may not be both given."
)
if cache_storage == "TMP":
tempdir = tempfile.mkdtemp()
storage = [tempdir]
weakref.finalize(self, self._remove_tempdir, tempdir)
else:
if isinstance(cache_storage, str):
storage = [cache_storage]
else:
storage = cache_storage
os.makedirs(storage[-1], exist_ok=True)
self.storage = storage
self.kwargs = target_options or {}
self.cache_check = cache_check
self.check_files = check_files
self.expiry = expiry_time
self.compression = compression
# Size of cache in bytes. If None then the size is unknown and will be
# recalculated the next time cache_size() is called. On writes to the
# cache this is reset to None.
self._cache_size = None
if same_names is not None and cache_mapper is not None:
raise ValueError(
"Cannot specify both same_names and cache_mapper in "
"CachingFileSystem.__init__"
)
if cache_mapper is not None:
self._mapper = cache_mapper
else:
self._mapper = create_cache_mapper(
same_names if same_names is not None else False
)
self.target_protocol = (
target_protocol
if isinstance(target_protocol, str)
else (fs.protocol if isinstance(fs.protocol, str) else fs.protocol[0])
)
self._metadata = CacheMetadata(self.storage)
self.load_cache()
self.fs = fs if fs is not None else filesystem(target_protocol, **self.kwargs)
if not self.fs.async_impl:
raise ValueError("Underlying filesystem needs to be async")
def _strip_protocol(path):
# acts as a method, since each instance has a difference target
return self.fs._strip_protocol(type(self)._strip_protocol(path))
self._strip_protocol: Callable = _strip_protocol
@typing.no_type_check
@staticmethod
def _remove_tempdir(tempdir):
try:
rmtree(tempdir)
except Exception: # noqa: S110
pass
@typing.no_type_check
def _mkcache(self):
os.makedirs(self.storage[-1], exist_ok=True)
@typing.no_type_check
def cache_size(self):
"""Return size of cache in bytes.
If more than one cache directory is in use, only the size of the last
one (the writable cache directory) is returned.
"""
if self._cache_size is None:
cache_dir = self.storage[-1]
self._cache_size = filesystem("file").du(cache_dir, withdirs=True)
return self._cache_size
@typing.no_type_check
def load_cache(self):
"""Read set of stored blocks from file"""
self._metadata.load()
self._mkcache()
self.last_cache = time.time()
@typing.no_type_check
def save_cache(self):
"""Save set of stored blocks from file"""
self._mkcache()
self._metadata.save()
self.last_cache = time.time()
self._cache_size = None
@typing.no_type_check
def _check_cache(self):
"""Reload caches if time elapsed or any disappeared"""
self._mkcache()
if not self.cache_check:
# explicitly told not to bother checking
return
timecond = time.time() - self.last_cache > self.cache_check
existcond = all(os.path.exists(storage) for storage in self.storage)
if timecond or not existcond:
self.load_cache()
@typing.no_type_check
def _check_file(self, path):
"""Is path in cache and still valid"""
path = self._strip_protocol(path)
self._check_cache()
return self._metadata.check_file(path, self)
@typing.no_type_check
def clear_cache(self):
"""Remove all files and metadata from the cache
In the case of multiple cache locations, this clears only the last one,
which is assumed to be the read/write one.
"""
rmtree(self.storage[-1])
self.load_cache()
self._cache_size = None
@typing.no_type_check
def clear_expired_cache(self, expiry_time=None):
"""Remove all expired files and metadata from the cache
In the case of multiple cache locations, this clears only the last one,
which is assumed to be the read/write one.
Parameters
----------
expiry_time: int
The time in seconds after which a local copy is considered useless.
If not defined the default is equivalent to the attribute from the
file caching instantiation.
"""
if not expiry_time:
expiry_time = self.expiry
self._check_cache()
expired_files, writable_cache_empty = self._metadata.clear_expired(expiry_time)
for fn in expired_files:
if os.path.exists(fn):
os.remove(fn)
if writable_cache_empty:
rmtree(self.storage[-1])
self.load_cache()
self._cache_size = None
@typing.no_type_check
def pop_from_cache(self, path):
"""Remove cached version of given file
Deletes local copy of the given (remote) path. If it is found in a cache
location which is not the last, it is assumed to be read-only, and
raises PermissionError
"""
path = self._strip_protocol(path)
fn = self._metadata.pop_file(path)
if fn is not None:
os.remove(fn)
self._cache_size = None
@typing.no_type_check
def _parent(self, path):
return self.fs._parent(path)
@typing.no_type_check
async def _ukey(self, path):
"""Hash of file properties, to tell if it has changed"""
return sha256(str(await self.fs._info(path)).encode()).hexdigest()
@typing.no_type_check
async def _make_local_details(self, path):
"""Create file detail dictionary for given file path"""
hash = self._mapper(path)
fn = os.path.join(self.storage[-1], hash)
detail = {
"original": path,
"fn": hash,
"blocks": True,
"time": time.time(),
"uid": await self._ukey(path),
}
self._metadata.update_file(path, detail)
logger.debug(f"Copying {path} to local cache")
return fn
@typing.no_type_check
async def _cat_file(
self,
path,
start=None,
end=None,
on_error="raise",
callback=DEFAULT_CALLBACK,
**kwargs,
):
"""Cat file, this is what is used inside Zarr 3.0 at the moment"""
getpath = None
try:
# Check if file is in cache
# This doesnt do anything fancy for chunked data, different chunk ranges
# are stored in different files even if there is over lap
path_chunked = path
if start:
path_chunked += f"_{start}"
if end:
path_chunked += f"_{end}"
detail = self._check_file(path_chunked)
if not detail:
storepath = await self._make_local_details(path_chunked)
getpath = path
else:
detail, storepath = (
detail if isinstance(detail, tuple) else (None, detail)
)
except Exception as e:
if on_error == "raise":
raise
if on_error == "return":
out = e
# If file was not in cache, get it using the base file system
if getpath:
resp = await self.fs._cat_file(getpath, start=start, end=end)
# Save to file
if isfilelike(storepath):
outfile = storepath
else:
outfile = open(storepath, "wb") # noqa: ASYNC101, ASYNC230
try:
outfile.write(resp)
# IDK yet
# callback.relative_update(len(chunk))
finally:
if not isfilelike(storepath):
outfile.close()
self.save_cache()
# Call back is weird here, like the progress should be on the file fetch
# but then how do we deal with the read? Maybe we times it by 2x?
if start is None:
start = 0
callback.set_size(1)
with open(storepath, "rb") as f:
f.seek(start)
if end is not None:
out = f.read(end - start)
else:
out = f.read()
callback.relative_update(1)
return out
@typing.no_type_check
def __getattribute__(self, item):
# TODO: Update
if item in {
"load_cache",
# "_open",
"save_cache",
# "close_and_update",
"__init__",
"__getattribute__",
"__reduce__",
"_make_local_details",
"_ukey",
"open",
"cat",
"_cat_file",
"cat_ranges",
"get",
"read_block",
"tail",
"head",
# "info",
# "ls",
"exists",
"isfile",
"isdir",
"_check_file",
"_check_cache",
"_mkcache",
"clear_cache",
"clear_expired_cache",
"pop_from_cache",
"local_file",
"_paths_from_path",
"get_mapper",
"open_many",
"commit_many",
"hash_name",
"__hash__",
"__eq__",
"to_json",
"to_dict",
"cache_size",
"pipe_file",
"pipe",
"start_transaction",
"end_transaction",
}:
# all the methods defined in this class. Note `open` here, since
# it calls `_open`, but is actually in superclass
return lambda *args, **kw: getattr(type(self), item).__get__(self)(
*args, **kw
)
if item in ["__reduce_ex__"]:
raise AttributeError
if item in ["transaction"]:
# property
return type(self).transaction.__get__(self)
if item in ["_cache", "transaction_type"]:
# class attributes
return getattr(type(self), item)
if item == "__class__":
return type(self)
d = object.__getattribute__(self, "__dict__")
fs = d.get("fs", None) # fs is not immediately defined
if item in d:
return d[item]
elif fs is not None:
if item in fs.__dict__:
# attribute of instance
return fs.__dict__[item]
# attributed belonging to the target filesystem
cls = type(fs)
m = getattr(cls, item)
if (inspect.isfunction(m) or inspect.isdatadescriptor(m)) and (
not hasattr(m, "__self__") or m.__self__ is None
):
# instance method
return m.__get__(fs, cls)
return m # class method or attribute
else:
# attributes of the superclass, while target is being set up
return super().__getattribute__(item)
