diff --git a/.gitignore b/.gitignore index 8b35876d..36ff0471 100644 --- a/.gitignore +++ b/.gitignore @@ -172,6 +172,9 @@ src/mypy.txt *.sublime-project *.sublime-workspace +# Zed Editor +.zed + # Ignore .DS_Store files .DS_Store diff --git a/packages/data/src/pyearthtools/data/indexes/_indexes.py b/packages/data/src/pyearthtools/data/indexes/_indexes.py index 0f6e2c49..251b8bec 100644 --- a/packages/data/src/pyearthtools/data/indexes/_indexes.py +++ b/packages/data/src/pyearthtools/data/indexes/_indexes.py @@ -130,6 +130,7 @@ def load( (Any): Loaded data """ + return open_files(files, **kwargs) def search(self, *args, **kwargs) -> Path | list[str | Path] | dict[str, str | Path]: @@ -200,7 +201,17 @@ def get(self, *args, **kwargs): Loaded Data """ try: - return self.load(self.search(*args), **kwargs) + + load_index = self.search(*args) + + cache = getattr(self, 'object_cache', {}) + cached_result = cache.get(str(load_index), None) + result = cached_result or self.load(load_index, **kwargs) + + if getattr(self, 'cache_last_used', True): + self.object_cache = {str(load_index): result} + + return result except Exception as e: raise DataNotFoundError(f"Data with args: {str(args)} could not be found.") from e diff --git a/packages/data/src/pyearthtools/data/indexes/utilities/fileload.py b/packages/data/src/pyearthtools/data/indexes/utilities/fileload.py index 0e9180a1..ab41b8eb 100644 --- a/packages/data/src/pyearthtools/data/indexes/utilities/fileload.py +++ b/packages/data/src/pyearthtools/data/indexes/utilities/fileload.py @@ -71,6 +71,8 @@ def open_dataset( return a dictionary of {path difference: dataset}. """ + # import pudb; pudb.set_trace() + def get_config(mf: bool = False): open_kwargs = copy.copy(pyearthtools.utils.config.get("data.open.xarray")) if mf: @@ -80,11 +82,28 @@ def get_config(mf: bool = False): if isinstance(location, (tuple, list)): try: + + + # open_kwargs['chunks'] = None + # # open_kwargs['compat'] = 'override' + # open_kwargs['combine_attrs'] = 'override' + # open_kwargs['data_vars'] = 'minimal' + # open_kwargs['parallel'] = False + + open_kwargs = get_config(True) + + open_kwargs['chunks'] = None + open_kwargs['compat'] = 'override' + open_kwargs['coords'] = 'minimal' + open_kwargs['combine_attrs'] = 'override' + open_kwargs['engine'] = 'h5netcdf' + open_kwargs['data_vars'] = 'minimal' + open_kwargs['parallel'] = False + return xr.open_mfdataset( filter_files(location), decode_timedelta=True, # TODO: should we raise a warning? It seems to be required for almost all our data. - compat="override", - **get_config(True), + **open_kwargs, ) except xr.MergeError as e: @@ -123,7 +142,7 @@ def get_config(mf: bool = False): return open_dataset(location, **kwargs) -NETCDF_FILE_EXTENSONS = [".nc", ".netcdf"] +NETCDF_FILE_EXTENSONS = [".nc", ".netcdf", ".nc4"] ZARR_FILE_EXTENSONS = [".zarr"] NUMPY_EXTENSIONS = [".npy", ".np", ".numpy"] CSV_EXTENSIONS = [".csv"] diff --git a/packages/data/src/pyearthtools/data/operations/index_routines.py b/packages/data/src/pyearthtools/data/operations/index_routines.py index 0966cdbe..da8bc71e 100644 --- a/packages/data/src/pyearthtools/data/operations/index_routines.py +++ b/packages/data/src/pyearthtools/data/operations/index_routines.py @@ -322,20 +322,29 @@ def get_path(file): open_kwargs.update(pyearthtools.utils.config.get("data.open.xarray_mf")) open_kwargs.update(kwargs) - try: - full_ds = xr.open_mfdataset( - list(set(dataset_paths)), - **open_kwargs, - ) - except NotImplementedError: - # Work around a bug/gap in xarray for loading NetCDF4 files and autochunking - - open_kwargs.pop("chunks") - - full_ds = xr.open_mfdataset( - list(set(dataset_paths)), - **open_kwargs, - ) + # try: + open_kwargs['chunks'] = None + open_kwargs['compat'] = 'override' + open_kwargs['coords'] = 'minimal' + open_kwargs['engine'] = 'h5netcdf' + open_kwargs['combine_attrs'] = 'override' + open_kwargs['data_vars'] = 'minimal' + open_kwargs['parallel'] = False + + full_ds = xr.open_mfdataset( + list(set(dataset_paths)), + **open_kwargs, + ) + + # except NotImplementedError: + # # Work around a bug/gap in xarray for loading NetCDF4 files and autochunking + + # open_kwargs.pop("chunks") + + # full_ds = xr.open_mfdataset( + # list(set(dataset_paths)), + # **open_kwargs, + # ) # full_ds = xr.merge( # [ diff --git a/packages/nci_site_archive/src/site_archive_nci/_Himawari.py b/packages/nci_site_archive/src/site_archive_nci/_Himawari.py index a680d6a1..36d6d287 100644 --- a/packages/nci_site_archive/src/site_archive_nci/_Himawari.py +++ b/packages/nci_site_archive/src/site_archive_nci/_Himawari.py @@ -105,6 +105,7 @@ def __init__( file_regex: str = FILE_REGEX, data_interval: tuple[int, str] = (10, "m"), transforms: Transform | TransformCollection | None = None, + cache_last_used=True, ): """ Setup Satellite Indexer @@ -118,6 +119,8 @@ def __init__( Override for data resolution. Defaults to (10, "m"). transforms (Transform | TransformCollection, optional): Base Transforms to apply. Defaults to TransformCollection(). + cache_last_used: + Will retain the last-loaded file in memory for possible re-use """ check_project(project_code="rv74") @@ -127,7 +130,12 @@ def __init__( self.file_regex = file_regex base_transform = pyearthtools.data.transforms.variables.Trim(variables) + (transforms or TransformCollection()) - super().__init__(transforms=base_transform, data_interval=data_interval or (10, "m")) + if cache_last_used: + self.object_cache = {} + + super().__init__(transforms=base_transform, + data_interval=data_interval or (10, "m"), + cache_last_used=True) self.record_initialisation() def filesystem( diff --git a/packages/pipeline/src/pyearthtools/pipeline/controller.py b/packages/pipeline/src/pyearthtools/pipeline/controller.py index 33bde25a..931e6d1e 100644 --- a/packages/pipeline/src/pyearthtools/pipeline/controller.py +++ b/packages/pipeline/src/pyearthtools/pipeline/controller.py @@ -700,6 +700,10 @@ def check(obj): except PipelineFilterException: continue + except Exception as exc: + if type(exc) in self._exceptions_to_ignore: + continue + try: if isinstance(sample, iterators.IterateResults): for sub_sample in sample.iterate_over_object(): diff --git a/packages/training/src/pyearthtools/training/wrapper/__init__.py b/packages/training/src/pyearthtools/training/wrapper/__init__.py index e5101c40..39827ad9 100644 --- a/packages/training/src/pyearthtools/training/wrapper/__init__.py +++ b/packages/training/src/pyearthtools/training/wrapper/__init__.py @@ -16,7 +16,7 @@ # ruff: noqa: F401 -from pyearthtools.training.wrapper.wrapper import ModelWrapper +from pyearthtools.training.wrapper._wrapper import ModelWrapper, SimpleModel from pyearthtools.training.wrapper import predict, train, utils @@ -35,7 +35,7 @@ except (ImportError, ModuleNotFoundError): LIGHTNING_IMPORTED = False -__all__ = ["ModelWrapper", "predict", "train", "utils", "TrainingWrapper", "Predictor"] +__all__ = ["ModelWrapper", "SimpleModel", "predict", "train", "utils", "TrainingWrapper", "Predictor"] if ONNX_IMPORTED: __all__.append("onnx") diff --git a/packages/training/src/pyearthtools/training/wrapper/wrapper.py b/packages/training/src/pyearthtools/training/wrapper/_wrapper.py similarity index 73% rename from packages/training/src/pyearthtools/training/wrapper/wrapper.py rename to packages/training/src/pyearthtools/training/wrapper/_wrapper.py index 6bb2a0c8..368da058 100644 --- a/packages/training/src/pyearthtools/training/wrapper/wrapper.py +++ b/packages/training/src/pyearthtools/training/wrapper/_wrapper.py @@ -22,6 +22,46 @@ from pyearthtools.training.data import PipelineDataModule +class SimpleModel(InitialisationRecordingMixin, metaclass=ABCMeta): + ''' + The SimpleModel base wrapper removes assumptions from the primary model wrapper class. + This provides a simpler on-ramp for early-stage model development, without imposing + as many requirements on the developer for comprehensive functionality. + + It also allows the direct use of a single Pipeline, without requiring a + train/validate split to be defined, avoiding the need for a PipelineDataModule. + + New users may wish to start here before implementing the full ModelWrapper class. + ''' + + def __init__( + self, + ): + """ + Construct Base model wrapper + + `model` will not be recorded in the initialisation by default, set `_record_model` to change + this behaviour. + """ + + pass + + @abstractmethod + def fit(self, pipeline, epochs=1): + ''' + Perform a single epoch 'fit' operation based on walking the pipeline + ''' + pass + + @abstractmethod + def predict(self, pipeline, query=None): + ''' + Perform a single prediction based either on a pipeline or a single sample + ''' + + + + class ModelWrapper(InitialisationRecordingMixin, metaclass=ABCMeta): """ Base Model Wrapper diff --git a/packages/training/src/pyearthtools/training/wrapper/predict/predict.py b/packages/training/src/pyearthtools/training/wrapper/predict/predict.py index 8069ce19..7a4d7c29 100644 --- a/packages/training/src/pyearthtools/training/wrapper/predict/predict.py +++ b/packages/training/src/pyearthtools/training/wrapper/predict/predict.py @@ -23,7 +23,7 @@ from pyearthtools.utils.initialisation import InitialisationRecordingMixin from pyearthtools.pipeline.controller import Pipeline -from pyearthtools.training.wrapper.wrapper import ModelWrapper +from pyearthtools.training.wrapper import ModelWrapper class Predictor(InitialisationRecordingMixin, metaclass=ABCMeta): diff --git a/packages/training/src/pyearthtools/training/wrapper/predict/timeseries.py b/packages/training/src/pyearthtools/training/wrapper/predict/timeseries.py index cde03cec..f68885ea 100644 --- a/packages/training/src/pyearthtools/training/wrapper/predict/timeseries.py +++ b/packages/training/src/pyearthtools/training/wrapper/predict/timeseries.py @@ -29,7 +29,7 @@ from pyearthtools.data.time import TimeDelta, Petdt, TimeRange from pyearthtools.pipeline.controller import Pipeline -from pyearthtools.training.wrapper.wrapper import ModelWrapper +from pyearthtools.training.wrapper import ModelWrapper from pyearthtools.training.wrapper.predict.predict import Predictor from pyearthtools.training.manage import Variables diff --git a/packages/training/src/pyearthtools/training/wrapper/train.py b/packages/training/src/pyearthtools/training/wrapper/train.py index 0706ed0e..6f80f72d 100644 --- a/packages/training/src/pyearthtools/training/wrapper/train.py +++ b/packages/training/src/pyearthtools/training/wrapper/train.py @@ -17,7 +17,7 @@ from abc import abstractmethod -from pyearthtools.training.wrapper.wrapper import ModelWrapper +from pyearthtools.training.wrapper import ModelWrapper class TrainingWrapper(ModelWrapper):