From b66a2416bcf512a997ef31f7c5fe2a6b4fc2974b Mon Sep 17 00:00:00 2001 From: Tennessee Leeuwenburg Date: Sat, 11 Jul 2026 16:40:31 +1000 Subject: [PATCH 1/2] 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 | 12 +++++- .../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 +- 7 files changed, 59 insertions(+), 6 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 285ddcf5..857d7ad9 100644 --- a/packages/data/src/pyearthtools/data/indexes/_indexes.py +++ b/packages/data/src/pyearthtools/data/indexes/_indexes.py @@ -200,7 +200,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/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): From 8818d2f67c55a822771c265d04d52990db01a6a6 Mon Sep 17 00:00:00 2001 From: Tennessee Leeuwenburg Date: Sun, 12 Jul 2026 13:06:07 +1000 Subject: [PATCH 2/2] Black and ruff --- .../src/pyearthtools/data/indexes/_indexes.py | 4 ++-- .../pyearthtools/training/wrapper/_wrapper.py | 16 +++++++--------- 2 files changed, 9 insertions(+), 11 deletions(-) diff --git a/packages/data/src/pyearthtools/data/indexes/_indexes.py b/packages/data/src/pyearthtools/data/indexes/_indexes.py index 857d7ad9..778adaac 100644 --- a/packages/data/src/pyearthtools/data/indexes/_indexes.py +++ b/packages/data/src/pyearthtools/data/indexes/_indexes.py @@ -203,11 +203,11 @@ def get(self, *args, **kwargs): load_index = self.search(*args) - cache = getattr(self, 'object_cache', {}) + 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): + if getattr(self, "cache_last_used", True): self.object_cache = {str(load_index): result} return result diff --git a/packages/training/src/pyearthtools/training/wrapper/_wrapper.py b/packages/training/src/pyearthtools/training/wrapper/_wrapper.py index 368da058..a53e3a23 100644 --- a/packages/training/src/pyearthtools/training/wrapper/_wrapper.py +++ b/packages/training/src/pyearthtools/training/wrapper/_wrapper.py @@ -23,16 +23,16 @@ 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 + 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, @@ -48,18 +48,16 @@ def __init__( @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):