From 7938f28973bc2298023f112a6a41ae27ebf74c1d Mon Sep 17 00:00:00 2001 From: Tennessee Leeuwenburg Date: Thu, 2 Apr 2026 14:02:11 +1100 Subject: [PATCH 1/3] Add code to tolerate exceptions of recognised (filtered) types in the pipeline controller --- packages/pipeline/src/pyearthtools/pipeline/controller.py | 4 ++++ 1 file changed, 4 insertions(+) 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(): From de0209419ba8d3fdc00cf27b2b3c340a7981e040 Mon Sep 17 00:00:00 2001 From: Tennessee Leeuwenburg Date: Thu, 2 Apr 2026 19:05:45 +1100 Subject: [PATCH 2/3] Various changes attempting to make file loading more efficient --- .../src/pyearthtools/data/indexes/_indexes.py | 12 +++++- .../data/indexes/utilities/fileload.py | 23 +++++++++++- .../data/operations/index_routines.py | 37 ++++++++++++------- .../src/site_archive_nci/_Himawari.py | 10 ++++- 4 files changed, 64 insertions(+), 18 deletions(-) diff --git a/packages/data/src/pyearthtools/data/indexes/_indexes.py b/packages/data/src/pyearthtools/data/indexes/_indexes.py index 0f6e2c49..e3b2e5ff 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,16 @@ 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(load_index, None) + result = cached_result or self.load(load_index, **kwargs) + + if getattr(self, 'cache_last_used', True): + self.object_cache = {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..dccc2adf 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: 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( From bc35498f8ebe46c291c318ee7c0139701cedf2d7 Mon Sep 17 00:00:00 2001 From: Tennessee Leeuwenburg Date: Sat, 11 Jul 2026 16:40:31 +1000 Subject: [PATCH 3/3] Refactored imports for pyearthtools/training/wrapper Added a SimpleModel class for notebooks / worked examples Added .nc4 file extension to recognised types Improved cache file finding code --- .gitignore | 3 ++ .../src/pyearthtools/data/indexes/_indexes.py | 5 ++- .../data/indexes/utilities/fileload.py | 2 +- .../pyearthtools/training/wrapper/__init__.py | 4 +- .../wrapper/{wrapper.py => _wrapper.py} | 40 +++++++++++++++++++ .../training/wrapper/predict/predict.py | 2 +- .../training/wrapper/predict/timeseries.py | 2 +- .../pyearthtools/training/wrapper/train.py | 2 +- 8 files changed, 52 insertions(+), 8 deletions(-) rename packages/training/src/pyearthtools/training/wrapper/{wrapper.py => _wrapper.py} (73%) 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 e3b2e5ff..251b8bec 100644 --- a/packages/data/src/pyearthtools/data/indexes/_indexes.py +++ b/packages/data/src/pyearthtools/data/indexes/_indexes.py @@ -201,14 +201,15 @@ def get(self, *args, **kwargs): Loaded Data """ try: + load_index = self.search(*args) cache = getattr(self, 'object_cache', {}) - cached_result = cache.get(load_index, None) + 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 = {load_index: result} + self.object_cache = {str(load_index): result} return result except Exception as e: diff --git a/packages/data/src/pyearthtools/data/indexes/utilities/fileload.py b/packages/data/src/pyearthtools/data/indexes/utilities/fileload.py index dccc2adf..ab41b8eb 100644 --- a/packages/data/src/pyearthtools/data/indexes/utilities/fileload.py +++ b/packages/data/src/pyearthtools/data/indexes/utilities/fileload.py @@ -142,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/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):