Refactored imports for pyearthtools/training/wrapper (#270)

* Refactored imports for pyearthtools/training/wrapper
* Added a SimpleModel class for notebooks / worked examples
* Improved cache file finding code

Author: tennlee

.gitignore

@@ -172,6 +172,9 @@ src/mypy.txt
 *.sublime-project
 *.sublime-workspace
 
+# Zed Editor
+.zed
+
 # Ignore .DS_Store files
 .DS_Store
 

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
 

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")

…pyearthtools/training/wrapper/wrapper.py …yearthtools/training/wrapper/_wrapper.pypackages/training/src/pyearthtools/training/wrapper/wrapper.py renamed to packages/training/src/pyearthtools/training/wrapper/_wrapper.py packages/training/src/pyearthtools/training/wrapper/wrapper.py renamed to packages/training/src/pyearthtools/training/wrapper/_wrapper.py

@@ -22,6 +22,44 @@
 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

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):

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

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):