deepsensor 0.1.11__tar.gz → 0.2.1__tar.gz

{deepsensor-0.1.11 → deepsensor-0.2.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: deepsensor
-Version: 0.1.11
+Version: 0.2.1
 Summary: A Python package for modelling xarray and pandas data with neural processes.
 Home-page: https://github.com/tom-andersson/deepsensor
 Author: Tom R. Andersson
@@ -11,7 +11,22 @@ Platform: linux
 Platform: osx
 Requires-Python: >=3.8
 Description-Content-Type: text/markdown
+Requires-Dist: neuralprocesses
+Requires-Dist: backends-matrix
+Requires-Dist: backends
+Requires-Dist: numpy
+Requires-Dist: pandas
+Requires-Dist: xarray
+Requires-Dist: dask
+Requires-Dist: distributed
+Requires-Dist: pyshp
+Requires-Dist: tqdm
+Requires-Dist: pooch
 Provides-Extra: testing
+Requires-Dist: pytest; extra == "testing"
+Requires-Dist: pytest-cov; extra == "testing"
+Requires-Dist: mypy; extra == "testing"
+Requires-Dist: tox; extra == "testing"
 
 [//]: # (![](figs/DeepSensorLogo.png))
 <ul style="text-align: center;">
@@ -26,7 +41,7 @@ data with neural processes</p>
 
 -----------
 
-[![release](https://img.shields.io/badge/release-v0.1.11-green?logo=github)](https://github.com/tom-andersson/deepsensor/releases)
+[![release](https://img.shields.io/badge/release-v0.2.1-green?logo=github)](https://github.com/tom-andersson/deepsensor/releases)
 ![Tests](https://github.com/tom-andersson/deepsensor/actions/workflows/tests.yml/badge.svg)
 [![Coverage Status](https://coveralls.io/repos/github/tom-andersson/deepsensor/badge.svg?branch=main)](https://coveralls.io/github/tom-andersson/deepsensor?branch=main)
 [![Code style: black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black)
@@ -97,7 +112,7 @@ import numpy as np
 ds_raw = xr.tutorial.open_dataset("air_temperature")
 
 # Normalise data
-data_processor = DataProcessor(x1_name="lat", x1_map=(15, 75), x2_name="lon", x2_map=(200, 330))
+data_processor = DataProcessor(x1_name="lat", x2_name="lon")
 ds = data_processor(ds_raw)
 
 # Set up task loader
@@ -116,7 +131,7 @@ for date in pd.date_range("2013-01-01", "2014-11-30")[::7]:
 # Train model
 trainer = Trainer(model, lr=5e-5)
 for epoch in range(10):
-    trainer(model, train_tasks, progress_bar=True)
+    trainer(train_tasks, progress_bar=True)
 
 # Predict on new task with 10% of context data and a dense grid of target points
 test_task = task_loader("2014-12-31", 0.1)

{deepsensor-0.1.11 → deepsensor-0.2.1}/README.md

@@ -11,7 +11,7 @@ data with neural processes</p>
 
 -----------
 
-[![release](https://img.shields.io/badge/release-v0.1.11-green?logo=github)](https://github.com/tom-andersson/deepsensor/releases)
+[![release](https://img.shields.io/badge/release-v0.2.1-green?logo=github)](https://github.com/tom-andersson/deepsensor/releases)
 ![Tests](https://github.com/tom-andersson/deepsensor/actions/workflows/tests.yml/badge.svg)
 [![Coverage Status](https://coveralls.io/repos/github/tom-andersson/deepsensor/badge.svg?branch=main)](https://coveralls.io/github/tom-andersson/deepsensor?branch=main)
 [![Code style: black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black)
@@ -82,7 +82,7 @@ import numpy as np
 ds_raw = xr.tutorial.open_dataset("air_temperature")
 
 # Normalise data
-data_processor = DataProcessor(x1_name="lat", x1_map=(15, 75), x2_name="lon", x2_map=(200, 330))
+data_processor = DataProcessor(x1_name="lat", x2_name="lon")
 ds = data_processor(ds_raw)
 
 # Set up task loader
@@ -101,7 +101,7 @@ for date in pd.date_range("2013-01-01", "2014-11-30")[::7]:
 # Train model
 trainer = Trainer(model, lr=5e-5)
 for epoch in range(10):
-    trainer(model, train_tasks, progress_bar=True)
+    trainer(train_tasks, progress_bar=True)
 
 # Predict on new task with 10% of context data and a dense grid of target points
 test_task = task_loader("2014-12-31", 0.1)

{deepsensor-0.1.11 → deepsensor-0.2.1}/deepsensor/active_learning/algorithms.py

@@ -72,8 +72,8 @@ class GreedyAlgorithm:
         self.X_s_mask = X_s_mask
         self.X_t_mask = X_t_mask
 
-        self.x1_name = self.model.data_processor.norm_params["coords"]["x1"]["name"]
-        self.x2_name = self.model.data_processor.norm_params["coords"]["x2"]["name"]
+        self.x1_name = self.model.data_processor.config["coords"]["x1"]["name"]
+        self.x2_name = self.model.data_processor.config["coords"]["x2"]["name"]
 
         # Normalised search and target coordinates
         if not X_normalised:

{deepsensor-0.1.11 → deepsensor-0.2.1}/deepsensor/data/loader.py

@@ -1,5 +1,9 @@
 from deepsensor.data.task import Task, flatten_X
 
+import os
+import json
+import copy
+
 import numpy as np
 import xarray as xr
 import pandas as pd
@@ -10,25 +14,27 @@ from deepsensor.errors import InvalidSamplingStrategyError
 
 
 class TaskLoader:
+    config_fname = "task_loader_config.json"
+
     def __init__(
         self,
+        task_loader_ID: str | None = None,
         context: Union[
             xr.DataArray,
             xr.Dataset,
             pd.DataFrame,
-            List[Union[xr.DataArray, xr.Dataset, pd.DataFrame]],
-        ],
+            str,
+            List[Union[xr.DataArray, xr.Dataset, pd.DataFrame, str]],
+        ] = None,
         target: Union[
             xr.DataArray,
             xr.Dataset,
             pd.DataFrame,
-            List[Union[xr.DataArray, xr.Dataset, pd.DataFrame]],
-        ],
-        aux_at_contexts: Union[xr.DataArray, xr.Dataset] = None,
-        aux_at_targets: Union[
-            xr.DataArray,
-            xr.Dataset,
+            str,
+            List[Union[xr.DataArray, xr.Dataset, pd.DataFrame, str]],
         ] = None,
+        aux_at_contexts: Union[xr.DataArray, xr.Dataset, str] = None,
+        aux_at_targets: Union[xr.DataArray, xr.Dataset, str] = None,
         links: Union[Tuple, List[Tuple[int, int]], None] = None,
         context_delta_t: Union[int, List[int]] = 0,
         target_delta_t: Union[int, List[int]] = 0,
@@ -39,7 +45,15 @@ class TaskLoader:
     ) -> None:
         """Initialise a TaskLoader object
 
+        The behaviour is the following:
+        - If all data passed as paths, load the data and overwrite the paths with the loaded data
+        - Either all data is passed as paths, or all data is passed as loaded data (else ValueError)
+        - If all data passed as paths, the TaskLoader can be saved with the `save` method (using config)
+
         Args:
+            task_loader_ID: If loading a TaskLoader from a config file, this is the folder the
+                TaskLoader was saved in (using `.save`). If this argument is passed, all other
+                arguments are ignored.
             context: Context data. Can be a single xr.DataArray, xr.Dataset or pd.DataFrame,
                 or a list/tuple of these.
             target: Target data. Can be a single xr.DataArray, xr.Dataset or pd.DataFrame,
@@ -68,46 +82,64 @@ class TaskLoader:
             dtype: Data type of the data. Used to cast the data to the specified dtype.
                 Default: 'float32'.
         """
-        self.time_freq = time_freq
-        self.xarray_interp_method = xarray_interp_method
-        self.discrete_xarray_sampling = discrete_xarray_sampling
-        self.dtype = dtype
+        if task_loader_ID is not None:
+            self.task_loader_ID = task_loader_ID
+            # Load TaskLoader from config file
+            fpath = os.path.join(task_loader_ID, self.config_fname)
+            with open(fpath, "r") as f:
+                self.config = json.load(f)
+
+            self.context = self.config["context"]
+            self.target = self.config["target"]
+            self.aux_at_contexts = self.config["aux_at_contexts"]
+            self.aux_at_targets = self.config["aux_at_targets"]
+            self.time_freq = self.config["time_freq"]
+            self.xarray_interp_method = self.config["xarray_interp_method"]
+            self.discrete_xarray_sampling = self.config["discrete_xarray_sampling"]
+            self.dtype = self.config["dtype"]
+        else:
+            self.context = context
+            self.target = target
+            self.aux_at_contexts = aux_at_contexts
+            self.aux_at_targets = aux_at_targets
+            self.time_freq = time_freq
+            self.xarray_interp_method = xarray_interp_method
+            self.discrete_xarray_sampling = discrete_xarray_sampling
+            self.dtype = dtype
+
+        if not isinstance(self.context, (tuple, list)):
+            self.context = (self.context,)
+        if not isinstance(self.target, (tuple, list)):
+            self.target = (self.target,)
+
+        if isinstance(context_delta_t, int):
+            context_delta_t = (context_delta_t,) * len(self.context)
+        else:
+            assert len(context_delta_t) == len(self.context)
+        if isinstance(target_delta_t, int):
+            target_delta_t = (target_delta_t,) * len(self.target)
+        else:
+            assert len(target_delta_t) == len(self.target)
+        self.context_delta_t = context_delta_t
+        self.target_delta_t = target_delta_t
+
+        all_paths = self._check_if_all_data_passed_as_paths()
+        if all_paths:
+            self._set_config()
+            self._load_data_from_paths()
 
         if aux_at_contexts is not None:
             self._check_offgrid_aux(self._check_offgrid_aux(aux_at_contexts))
         if aux_at_targets is not None:
             self._check_offgrid_aux(aux_at_targets)
 
-        if isinstance(context, (xr.DataArray, xr.Dataset, pd.DataFrame, pd.Series)):
-            context = (context,)
-        if isinstance(target, (xr.DataArray, xr.Dataset, pd.DataFrame, pd.Series)):
-            target = (target,)
-        (
-            context,
-            target,
-            aux_at_contexts,
-            aux_at_targets,
-        ) = self._cast_context_and_target_to_dtype(
-            context, target, aux_at_contexts, aux_at_targets
-        )
-        self.context = context
-        self.target = target
-        self.aux_at_contexts = aux_at_contexts
-        self.aux_at_targets = aux_at_targets
+        self.context = self._cast_to_dtype(self.context)
+        self.target = self._cast_to_dtype(self.target)
+        self.aux_at_contexts = self._cast_to_dtype(self.aux_at_contexts)
+        self.aux_at_targets = self._cast_to_dtype(self.aux_at_targets)
 
         self.links = self.check_links(links)
 
-        if isinstance(context_delta_t, int):
-            context_delta_t = (context_delta_t,) * len(context)
-        else:
-            assert len(context_delta_t) == len(context)
-        if isinstance(target_delta_t, int):
-            target_delta_t = (target_delta_t,) * len(target)
-        else:
-            assert len(target_delta_t) == len(target)
-        self.context_delta_t = context_delta_t
-        self.target_delta_t = target_delta_t
-
         (
             self.context_dims,
             self.target_dims,
@@ -121,15 +153,128 @@ class TaskLoader:
             self.aux_at_target_var_IDs,
         ) = self.infer_context_and_target_var_IDs()
 
-    def _cast_context_and_target_to_dtype(
+    def _set_config(self):
+        """Instantiate a config dictionary for the TaskLoader object"""
+        # Take deepcopy to avoid modifying the original config
+        self.config = copy.deepcopy(
+            dict(
+                context=self.context,
+                target=self.target,
+                aux_at_contexts=self.aux_at_contexts,
+                aux_at_targets=self.aux_at_targets,
+                context_delta_t=self.context_delta_t,
+                target_delta_t=self.target_delta_t,
+                time_freq=self.time_freq,
+                xarray_interp_method=self.xarray_interp_method,
+                discrete_xarray_sampling=self.discrete_xarray_sampling,
+                dtype=self.dtype,
+            )
+        )
+
+    def _check_if_all_data_passed_as_paths(self) -> bool:
+        """If all data passed as paths, save paths to config and return True."""
+
+        def _check_if_strings(x, mode="all"):
+            if x is None:
+                return None
+            elif isinstance(x, (tuple, list)):
+                if mode == "all":
+                    return all([isinstance(x_i, str) for x_i in x])
+                elif mode == "any":
+                    return any([isinstance(x_i, str) for x_i in x])
+            else:
+                return isinstance(x, str)
+
+        all_paths = all(
+            filter(
+                lambda x: x is not None,
+                [
+                    _check_if_strings(self.context),
+                    _check_if_strings(self.target),
+                    _check_if_strings(self.aux_at_contexts),
+                    _check_if_strings(self.aux_at_targets),
+                ],
+            )
+        )
+        self._is_saveable = all_paths
+
+        any_paths = any(
+            filter(
+                lambda x: x is not None,
+                [
+                    _check_if_strings(self.context, mode="any"),
+                    _check_if_strings(self.target, mode="any"),
+                    _check_if_strings(self.aux_at_contexts, mode="any"),
+                    _check_if_strings(self.aux_at_targets, mode="any"),
+                ],
+            )
+        )
+        if any_paths and not all_paths:
+            raise ValueError(
+                "Data must be passed either all as paths or all as xarray/pandas objects (not a mix)."
+            )
+
+        return all_paths
+
+    def _load_data_from_paths(self):
+        """Load data from paths and overwrite paths with loaded data."""
+
+        def _load_pandas_or_xarray(path):
+            # Need to be careful about this. We need to ensure data gets into the right form
+            #  for TaskLoader.
+            if path is None:
+                return None
+            elif path.endswith(".nc"):
+                return xr.open_dataset(path)
+            elif path.endswith(".csv"):
+                df = pd.read_csv(path)
+                if "time" in df.columns:
+                    df["time"] = pd.to_datetime(df["time"])
+                    df = df.set_index(["time", "x1", "x2"]).sort_index()
+                else:
+                    df = df.set_index(["x1", "x2"]).sort_index()
+                return df
+            else:
+                raise ValueError(f"Unknown file extension for {path}")
+
+        def _load_data(data):
+            if isinstance(data, (tuple, list)):
+                data = tuple([_load_pandas_or_xarray(data_i) for data_i in data])
+            else:
+                data = _load_pandas_or_xarray(data)
+            return data
+
+        self.context = _load_data(self.context)
+        self.target = _load_data(self.target)
+        self.aux_at_contexts = _load_data(self.aux_at_contexts)
+        self.aux_at_targets = _load_data(self.aux_at_targets)
+
+    def save(self, folder: str):
+        """Save TaskLoader config to JSON in `folder`"""
+        if not self._is_saveable:
+            raise ValueError(
+                "TaskLoader cannot be saved because not all data was passed as paths."
+            )
+
+        os.makedirs(folder, exist_ok=True)
+        fpath = os.path.join(folder, self.config_fname)
+        with open(fpath, "w") as f:
+            json.dump(self.config, f, indent=4, sort_keys=False)
+
+    def _cast_to_dtype(
         self,
-        context: List,
-        target: List,
-        aux_at_contexts: Union[xr.DataArray, xr.Dataset] = None,
-        aux_at_targets: Union[xr.DataArray, xr.Dataset] = None,
+        var: Union[
+            xr.DataArray,
+            xr.Dataset,
+            pd.DataFrame,
+            List[Union[xr.DataArray, xr.Dataset, pd.DataFrame, str]],
+        ],
     ) -> (List, List):
         """Cast context and target data to the default dtype.
 
+        TODO unit test this by passing in a variety of data types and checking that they are
+            cast correctly.
+
         Returns
         -------
         context : tuple. Tuple of context data with specified dtype.
@@ -148,19 +293,19 @@ class TaskLoader:
             elif isinstance(var, (pd.DataFrame, pd.Series)):
                 var = var.astype(self.dtype)
                 # Note: Numeric pandas indexes are always cast to float64, so we have to cast
-                # x1/x2 coord dtypes during task sampling
+                #   x1/x2 coord dtypes during task sampling
             else:
                 raise ValueError(f"Unknown type {type(var)} for context set {var}")
             return var
 
-        context = tuple([cast_to_dtype(var) for var in context])
-        target = tuple([cast_to_dtype(var) for var in target])
-        if aux_at_contexts is not None:
-            aux_at_contexts = aux_at_contexts.astype(self.dtype)
-        if aux_at_targets is not None:
-            aux_at_targets = aux_at_targets.astype(self.dtype)
+        if var is None:
+            return var
+        elif isinstance(var, (tuple, list)):
+            var = tuple([cast_to_dtype(var_i) for var_i in var])
+        else:
+            var = cast_to_dtype(var)
 
-        return context, target, aux_at_contexts, aux_at_targets
+        return var
 
     def _check_offgrid_aux(self, offgrid_aux):
         if offgrid_aux is not None and "time" in offgrid_aux.dims:

{deepsensor-0.1.11 → deepsensor-0.2.1}/deepsensor/data/processor.py

@@ -1,4 +1,6 @@
 import numpy as np
+import os
+import json
 
 import xarray as xr
 import pandas as pd
@@ -7,50 +9,71 @@ import pprint
 
 from copy import deepcopy
 
+from plum import dispatch
 from typing import Union
 
 
 class DataProcessor:
     """Normalise xarray and pandas data for use in deepsensor models"""
 
+    config_fname = "data_processor_config.json"
+
     def __init__(
         self,
-        norm_params: dict = None,
+        folder: str | None = None,
         time_name: str = "time",
         x1_name: str = "x1",
         x2_name: str = "x2",
-        x1_map: tuple = (0, 1),
-        x2_map: tuple = (0, 1),
+        x1_map: tuple | None = None,
+        x2_map: tuple | None = None,
         deepcopy: bool = True,
         verbose: bool = False,
-    ) -> None:
+    ):
         """Initialise DataProcessor
 
         Args:
-            norm_params (dict, optional): Normalisation params. Defaults to {}.
+            folder (str, optional): Folder to load normalisation params from. Defaults to None.
             x1_name (str, optional): Name of first spatial coord (e.g. "lat"). Defaults to "x1".
             x2_name (str, optional): Name of second spatial coord (e.g. "lon"). Defaults to "x2".
             x1_map (tuple, optional): 2-tuple of raw x1 coords to linearly map to (0, 1), respectively.
-                Defaults to (0, 1) (i.e. no normalisation).
+                If None, the mapping inferred from first data object passed to DataProcessor.
+                Defaults to None.
             x2_map (tuple, optional): 2-tuple of raw x2 coords to linearly map to (0, 1), respectively.
-                Defaults to (0, 1) (i.e. no normalisation).
+                If None, the mapping inferred from first data object passed to DataProcessor.
+                Defaults to None.
             deepcopy (bool, optional): Whether to make a deepcopy of raw data to ensure it is
-                not changed by reference when normalising. Defaults to True.
+                not changed by reference when normalising. Defaults to True. TODO - test if necessary.
             verbose (bool, optional): Whether to print verbose output. Defaults to False.
         """
-        x1_map, x2_map = self._validate_coord_mappings(x1_map, x2_map)
+        if folder is not None:
+            fpath = os.path.join(folder, self.config_fname)
+            if not os.path.exists(fpath):
+                raise FileNotFoundError(
+                    f"Could not find DataProcessor config file {fpath}"
+                )
+            with open(fpath, "r") as f:
+                self.config = json.load(f)
 
-        if norm_params is None:
-            norm_params = {}
+            self.x1_none = self.config["coords"]["x1"]["map"] is None
+            self.x2_none = self.config["coords"]["x2"]["map"] is None
+        else:
+            self.config = {}
 
-        self.norm_params = norm_params
+            self.x1_none = x1_map is None
+            self.x2_none = x2_map is None
+            if (self.x1_none and not self.x2_none) or (
+                not self.x1_none and self.x2_none
+            ):
+                raise ValueError("Must provide both x1_map and x2_map, or neither.")
+            elif not self.x1_none and not self.x2_none:
+                x1_map, x2_map = self._validate_coord_mappings(x1_map, x2_map)
 
-        if "coords" not in self.norm_params:
-            # Add coordinate normalisation info to norm_params
-            self.set_coord_params(time_name, x1_name, x1_map, x2_name, x2_map)
+            if "coords" not in self.config:
+                # Add coordinate normalisation info to config
+                self.set_coord_params(time_name, x1_name, x1_map, x2_name, x2_map)
 
         self.raw_spatial_coord_names = [
-            self.norm_params["coords"][coord]["name"] for coord in ["x1", "x2"]
+            self.config["coords"][coord]["name"] for coord in ["x1", "x2"]
         ]
 
         self.deepcopy = deepcopy
@@ -59,6 +82,13 @@ class DataProcessor:
         # List of valid normalisation method names
         self.valid_methods = ["mean_std", "min_max"]
 
+    def save(self, folder: str):
+        """Save DataProcessor config to JSON in `folder`"""
+        os.makedirs(folder, exist_ok=True)
+        fpath = os.path.join(folder, self.config_fname)
+        with open(fpath, "w") as f:
+            json.dump(self.config, f, indent=4, sort_keys=False)
+
     def _validate_coord_mappings(self, x1_map, x2_map):
         """Ensure the maps are valid and of appropriate types."""
         try:
@@ -84,8 +114,7 @@ class DataProcessor:
     def _validate_xr(self, data: Union[xr.DataArray, xr.Dataset]):
         def _validate_da(da: xr.DataArray):
             coord_names = [
-                self.norm_params["coords"][coord]["name"]
-                for coord in ["time", "x1", "x2"]
+                self.config["coords"][coord]["name"] for coord in ["time", "x1", "x2"]
             ]
             if coord_names[0] not in da.dims:
                 # We don't have a time dimension.
@@ -104,7 +133,7 @@ class DataProcessor:
 
     def _validate_pandas(self, df: Union[pd.DataFrame, pd.Series]):
         coord_names = [
-            self.norm_params["coords"][coord]["name"] for coord in ["time", "x1", "x2"]
+            self.config["coords"][coord]["name"] for coord in ["time", "x1", "x2"]
         ]
 
         if coord_names[0] not in df.index.names:
@@ -122,7 +151,7 @@ class DataProcessor:
 
     def __str__(self):
         s = "DataProcessor with normalisation params:\n"
-        s += pprint.pformat(self.norm_params)
+        s += pprint.pformat(self.config)
         return s
 
     @classmethod
@@ -136,42 +165,43 @@ class DataProcessor:
 
     def set_coord_params(self, time_name, x1_name, x1_map, x2_name, x2_map):
         """Set coordinate normalisation params"""
-        self.norm_params["coords"] = {}
-        self.norm_params["coords"]["time"] = {"name": time_name}
-        self.norm_params["coords"]["x1"] = {}
-        self.norm_params["coords"]["x2"] = {}
-        self.norm_params["coords"]["x1"]["name"] = x1_name
-        self.norm_params["coords"]["x1"]["map"] = x1_map
-        self.norm_params["coords"]["x2"]["name"] = x2_name
-        self.norm_params["coords"]["x2"]["map"] = x2_map
+        self.config["coords"] = {}
+        self.config["coords"]["time"] = {"name": time_name}
+        self.config["coords"]["x1"] = {}
+        self.config["coords"]["x2"] = {}
+        self.config["coords"]["x1"]["name"] = x1_name
+        self.config["coords"]["x1"]["map"] = x1_map
+        self.config["coords"]["x2"]["name"] = x2_name
+        self.config["coords"]["x2"]["map"] = x2_map
 
     def check_params_computed(self, var_ID, method):
         """Check if normalisation params computed for a given variable"""
         if (
-            var_ID in self.norm_params
-            and self.norm_params[var_ID]["method"] == method
-            and "param1" in self.norm_params[var_ID]
-            and "param2" in self.norm_params[var_ID]
+            var_ID in self.config
+            and self.config[var_ID]["method"] == method
+            and "params" in self.config[var_ID]
         ):
             return True
         else:
             return False
 
-    def add_to_norm_params(self, var_ID, **kwargs):
-        """Add `kwargs` to `norm_params` dict for variable `var_ID`"""
-        self.norm_params[var_ID] = kwargs
+    def add_to_config(self, var_ID, **kwargs):
+        """Add `kwargs` to `config` dict for variable `var_ID`"""
+        self.config[var_ID] = kwargs
 
     def get_norm_params(self, var_ID, data, method=None):
-        """Get pre-computed normalisation params or compute them for variable `var_ID`"""
+        """Get pre-computed normalisation params or compute them for variable `var_ID`
+
+        TODO do we need to pass var_ID? Can we just use name of data?
+        """
         if method not in self.valid_methods:
             raise ValueError(
                 f"Method {method} not recognised. Must be one of {self.valid_methods}"
             )
 
         if self.check_params_computed(var_ID, method):
-            # Already have "param1" and "param2" in norm_params with `"method": method` - load them
-            param1 = self.norm_params[var_ID]["param1"]
-            param2 = self.norm_params[var_ID]["param2"]
+            # Already have "params" in config with `"method": method` - load them
+            params = self.config[var_ID]["params"]
         else:
             # Params not computed - compute them now
             if self.verbose:
@@ -182,19 +212,16 @@ class DataProcessor:
                 )
             DataProcessor.load_dask(data)
             if method == "mean_std":
-                param1 = float(data.mean())
-                param2 = float(data.std())
+                params = {"mean": float(data.mean()), "std": float(data.std())}
             elif method == "min_max":
-                param1 = float(data.min())
-                param2 = float(data.max())
+                params = {"min": float(data.min()), "max": float(data.max())}
             if self.verbose:
-                print(
-                    f"Done. {var_ID} {method} param1={param1:.3f}, param2={param2:.3f}"
-                )
-            self.add_to_norm_params(
-                var_ID, **{"method": method, "param1": param1, "param2": param2}
+                print(f"Done. {var_ID} {method} params={params}")
+            self.add_to_config(
+                var_ID,
+                **{"method": method, "params": params},
             )
-        return param1, param2
+        return params
 
     def map_coord_array(self, coord_array: np.ndarray, unnorm: bool = False):
         """Normalise or unnormalise a coordinate array
@@ -214,8 +241,8 @@ class DataProcessor:
             x1 (np.ndarray): Array of shape (N_x1,) containing spatial coords of x1
             unnorm (bool, optional): Whether to unnormalise. Defaults to False.
         """
-        x11, x12 = self.norm_params["coords"]["x1"]["map"]
-        x21, x22 = self.norm_params["coords"]["x2"]["map"]
+        x11, x12 = self.config["coords"]["x1"]["map"]
+        x21, x22 = self.config["coords"]["x2"]["map"]
 
         if not unnorm:
             new_coords_x1 = (x1 - x11) / (x12 - x11)
@@ -232,6 +259,7 @@ class DataProcessor:
         unnorm=False,
     ):
         """Normalise spatial coords in a pandas or xarray object"""
+
         if isinstance(data, (pd.DataFrame, pd.Series)):
             # Reset index to get coords as columns
             indexes = list(data.index.names)
@@ -239,21 +267,39 @@ class DataProcessor:
 
         if unnorm:
             new_coord_IDs = [
-                self.norm_params["coords"][coord_ID]["name"]
+                self.config["coords"][coord_ID]["name"]
                 for coord_ID in ["time", "x1", "x2"]
             ]
             old_coord_IDs = ["time", "x1", "x2"]
         else:
             new_coord_IDs = ["time", "x1", "x2"]
             old_coord_IDs = [
-                self.norm_params["coords"][coord_ID]["name"]
+                self.config["coords"][coord_ID]["name"]
                 for coord_ID in ["time", "x1", "x2"]
             ]
 
-        new_x1, new_x2 = self.map_x1_and_x2(
-            data[old_coord_IDs[1]], data[old_coord_IDs[2]], unnorm=unnorm
+        x1, x2 = (
+            data[old_coord_IDs[1]],
+            data[old_coord_IDs[2]],
         )
 
+        # Infer x1 and x2 mappings from min/max of data coords if not provided by user
+        if self.x1_none and self.x2_none:
+            x1_map = (x1.min(), x1.max())
+            x2_map = (x2.min(), x2.max())
+            x1_map, x2_map = self._validate_coord_mappings(x1_map, x2_map)
+            self.config["coords"]["x1"]["map"] = x1_map
+            self.config["coords"]["x2"]["map"] = x2_map
+            if self.verbose:
+                print(
+                    f"Inferring x1_map={x1_map} and x2_map={x2_map} from data min/max"
+                )
+
+            self.x2_none = False
+            self.x1_none = False
+
+        new_x1, new_x2 = self.map_x1_and_x2(x1, x2, unnorm=unnorm)
+
         if isinstance(data, (pd.DataFrame, pd.Series)):
             # Drop old spatial coord columns *before* adding new ones, in case
             # the old name is already x1.
@@ -306,46 +352,46 @@ class DataProcessor:
         """Normalise or unnormalise the data values in an xarray, pandas, or numpy object"""
         if not unnorm and method is None:
             raise ValueError("Must provide `method` if normalising data.")
-        elif (
-            unnorm
-            and method is not None
-            and self.norm_params[var_ID]["method"] != method
-        ):
+        elif unnorm and method is not None and self.config[var_ID]["method"] != method:
             # User has provided a different method to the one used for normalising
             raise ValueError(
-                f"Variable '{var_ID}' has been normalised with method '{self.norm_params[var_ID]['method']}', "
+                f"Variable '{var_ID}' has been normalised with method '{self.config[var_ID]['method']}', "
                 f"cannot unnormalise with method '{method}'. Pass `method=None` or"
-                f"`method='{self.norm_params[var_ID]['method']}'`"
+                f"`method='{self.config[var_ID]['method']}'`"
             )
         if method is None and unnorm:
-            # Determine normalisation method from norm_params for unnormalising
-            method = self.norm_params[var_ID]["method"]
+            # Determine normalisation method from config for unnormalising
+            method = self.config[var_ID]["method"]
         elif method not in self.valid_methods:
             raise ValueError(
                 f"Method {method} not recognised. Use one of {self.valid_methods}"
             )
 
-        param1, param2 = self.get_norm_params(var_ID, data, method)
+        params = self.get_norm_params(var_ID, data, method)
 
         if method == "mean_std":
+            std = params["std"]
+            mean = params["mean"]
             if unnorm:
-                scale = param2
-                offset = param1
+                scale = std
+                offset = mean
             else:
-                scale = 1 / param2
-                offset = -param1 / param2
+                scale = 1 / std
+                offset = -mean / std
             data = data * scale
             if add_offset:
                 data = data + offset
             return data
 
         elif method == "min_max":
+            minimum = params["min"]
+            maximum = params["max"]
             if unnorm:
-                scale = (param2 - param1) / 2
-                offset = (param2 + param1) / 2
+                scale = (maximum - minimum) / 2
+                offset = (maximum + minimum) / 2
             else:
-                scale = 2 / (param2 - param1)
-                offset = -(param2 + param1) / (param2 - param1)
+                scale = 2 / (maximum - minimum)
+                offset = -(maximum + minimum) / (maximum - minimum)
             data = data * scale
             if add_offset:
                 data = data + offset

{deepsensor-0.1.11 → deepsensor-0.2.1}/deepsensor/model/convnp.py

@@ -1,5 +1,7 @@
 import copy
+import os.path
 import warnings
+import json
 from typing import Union, List
 
 import lab as B
@@ -52,7 +54,7 @@ class ConvNP(DeepSensorModel):
         likelihood (str, optional): Likelihood. Must be one of `"cnp"` (equivalently `"het"`),
             `"gnp"` (equivalently `"lowrank"`), or `"cnp-spikes-beta"` (equivalently `"spikes-beta"`).
             Defaults to `"cnp"`.
-        dim_x (int, optional): Dimensionality of the inputs. Defaults to 1.
+        dim_x (int, optional): Dimensionality of the inputs. Defaults to 2.
         dim_y (int, optional): Dimensionality of the outputs. Defaults to 1.
         dim_yc (int or tuple[int], optional): Dimensionality of the outputs of the
             context set. You should set this if the dimensionality of the outputs
@@ -110,7 +112,7 @@ class ConvNP(DeepSensorModel):
         # so unnormalisation will not be performed at inference time.
         super().__init__()
 
-        self.model = construct_neural_process(*args, **kwargs)
+        self.model, self.config = construct_neural_process(*args, **kwargs)
 
     @dispatch
     def __init__(
@@ -164,7 +166,7 @@ class ConvNP(DeepSensorModel):
                 print(f"decoder_scale inferred from TaskLoader: {decoder_scale}")
             kwargs["decoder_scale"] = decoder_scale
 
-        self.model = construct_neural_process(*args, dim_x=2, **kwargs)
+        self.model, self.config = construct_neural_process(*args, **kwargs)
 
     @dispatch
     def __init__(
@@ -178,6 +180,59 @@ class ConvNP(DeepSensorModel):
 
         self.model = neural_process
 
+    @dispatch
+    def __init__(self, model_ID: str):
+        """Instantiate a model from a folder containing model weights and config."""
+        super().__init__()
+
+        self.load(model_ID)
+
+    @dispatch
+    def __init__(
+        self,
+        data_processor: DataProcessor,
+        task_loader: TaskLoader,
+        model_ID: str,
+    ):
+        """Instantiate a model from a folder containing model weights and config."""
+        super().__init__(data_processor, task_loader)
+
+        self.load(model_ID)
+
+    def save(self, model_ID: str):
+        """Save the model weights and config to a folder."""
+        os.makedirs(model_ID, exist_ok=True)
+
+        if backend.str == "torch":
+            import torch
+
+            torch.save(self.model.state_dict(), os.path.join(model_ID, "model.pt"))
+        elif backend.str == "tf":
+            self.model.save_weights(os.path.join(model_ID, "model"))
+        else:
+            raise NotImplementedError(f"Backend {backend.str} not supported.")
+
+        config_fpath = os.path.join(model_ID, "model_config.json")
+        with open(config_fpath, "w") as f:
+            json.dump(self.config, f, indent=4, sort_keys=False)
+
+    def load(self, model_ID: str):
+        """Load a model from a folder containing model weights and config."""
+        config_fpath = os.path.join(model_ID, "model_config.json")
+        with open(config_fpath, "r") as f:
+            self.config = json.load(f)
+
+        self.model, _ = construct_neural_process(**self.config)
+
+        if backend.str == "torch":
+            import torch
+
+            self.model.load_state_dict(torch.load(os.path.join(model_ID, "model.pt")))
+        elif backend.str == "tf":
+            self.model.load_weights(os.path.join(model_ID, "model"))
+        else:
+            raise NotImplementedError(f"Backend {backend.str} not supported.")
+
     @classmethod
     def modify_task(cls, task):
         """Cast numpy arrays to TensorFlow or PyTorch tensors, add batch dim, and mask NaNs"""

{deepsensor-0.1.11 → deepsensor-0.2.1}/deepsensor/model/model.py

@@ -307,7 +307,7 @@ class DeepSensorModel(ProbabilisticModel):
                 arr[i] = self.data_processor.map_array(
                     arr[i],
                     var_ID,
-                    method=self.data_processor.norm_params[var_ID]["method"],
+                    method=self.data_processor.config[var_ID]["method"],
                     unnorm=True,
                     **kwargs,
                 )

{deepsensor-0.1.11 → deepsensor-0.2.1}/deepsensor/model/nps.py

@@ -157,8 +157,8 @@ def construct_neural_process(
     elif likelihood == "cnp-spikes-beta":
         likelihood = "spikes-beta"
 
-    # Use a stride of 1 for the first layer and 2 for all other layers
-    unet_strides = (1, *(2,) * (len(unet_channels) - 1))
+    # Log the call signature for `construct_convgnp`
+    config = dict(locals())
 
     if backend.str == "torch":
         import torch
@@ -177,14 +177,15 @@ def construct_neural_process(
         dim_yt=dim_yt,
         dim_aux_t=dim_aux_t,
         dim_lv=dim_lv,
+        likelihood=likelihood,
         conv_arch=conv_arch,
-        unet_channels=unet_channels,
+        unet_channels=tuple(unet_channels),
         unet_resize_convs=unet_resize_convs,
         unet_resize_conv_interp_method=unet_resize_conv_interp_method,
         aux_t_mlp_layers=aux_t_mlp_layers,
-        likelihood=likelihood,
         unet_kernels=unet_kernels,
-        unet_strides=unet_strides,
+        # Use a stride of 1 for the first layer and 2 for all other layers
+        unet_strides=(1, *(2,) * (len(unet_channels) - 1)),
         points_per_unit=points_per_unit,
         encoder_scales=encoder_scales,
         encoder_scales_learnable=encoder_scales_learnable,
@@ -194,7 +195,8 @@ def construct_neural_process(
         epsilon=epsilon,
         dtype=dtype,
     )
-    return neural_process
+
+    return neural_process, config
 
 
 def compute_encoding_tensor(model, task: Task):

{deepsensor-0.1.11 → deepsensor-0.2.1}/deepsensor.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: deepsensor
-Version: 0.1.11
+Version: 0.2.1
 Summary: A Python package for modelling xarray and pandas data with neural processes.
 Home-page: https://github.com/tom-andersson/deepsensor
 Author: Tom R. Andersson
@@ -11,7 +11,22 @@ Platform: linux
 Platform: osx
 Requires-Python: >=3.8
 Description-Content-Type: text/markdown
+Requires-Dist: neuralprocesses
+Requires-Dist: backends-matrix
+Requires-Dist: backends
+Requires-Dist: numpy
+Requires-Dist: pandas
+Requires-Dist: xarray
+Requires-Dist: dask
+Requires-Dist: distributed
+Requires-Dist: pyshp
+Requires-Dist: tqdm
+Requires-Dist: pooch
 Provides-Extra: testing
+Requires-Dist: pytest; extra == "testing"
+Requires-Dist: pytest-cov; extra == "testing"
+Requires-Dist: mypy; extra == "testing"
+Requires-Dist: tox; extra == "testing"
 
 [//]: # (![](figs/DeepSensorLogo.png))
 <ul style="text-align: center;">
@@ -26,7 +41,7 @@ data with neural processes</p>
 
 -----------
 
-[![release](https://img.shields.io/badge/release-v0.1.11-green?logo=github)](https://github.com/tom-andersson/deepsensor/releases)
+[![release](https://img.shields.io/badge/release-v0.2.1-green?logo=github)](https://github.com/tom-andersson/deepsensor/releases)
 ![Tests](https://github.com/tom-andersson/deepsensor/actions/workflows/tests.yml/badge.svg)
 [![Coverage Status](https://coveralls.io/repos/github/tom-andersson/deepsensor/badge.svg?branch=main)](https://coveralls.io/github/tom-andersson/deepsensor?branch=main)
 [![Code style: black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black)
@@ -97,7 +112,7 @@ import numpy as np
 ds_raw = xr.tutorial.open_dataset("air_temperature")
 
 # Normalise data
-data_processor = DataProcessor(x1_name="lat", x1_map=(15, 75), x2_name="lon", x2_map=(200, 330))
+data_processor = DataProcessor(x1_name="lat", x2_name="lon")
 ds = data_processor(ds_raw)
 
 # Set up task loader
@@ -116,7 +131,7 @@ for date in pd.date_range("2013-01-01", "2014-11-30")[::7]:
 # Train model
 trainer = Trainer(model, lr=5e-5)
 for epoch in range(10):
-    trainer(model, train_tasks, progress_bar=True)
+    trainer(train_tasks, progress_bar=True)
 
 # Predict on new task with 10% of context data and a dense grid of target points
 test_task = task_loader("2014-12-31", 0.1)

{deepsensor-0.1.11 → deepsensor-0.2.1}/setup.cfg

@@ -1,6 +1,6 @@
 [metadata]
 name = deepsensor
-version = 0.1.11
+version = 0.2.1
 author = Tom R. Andersson
 author_email = tomand@bas.ac.uk
 description = A Python package for modelling xarray and pandas data with neural processes.

{deepsensor-0.1.11 → deepsensor-0.2.1}/tests/test_data_processor.py

@@ -64,30 +64,39 @@ class TestDataProcessor(unittest.TestCase):
             return False
         return True
 
+    def test_only_passing_one_x_mapping_raises_valueerror(self):
+        with self.assertRaises(ValueError):
+            DataProcessor(x1_map=(20, 40), x2_map=None)
+
     def test_unnorm_restores_data_for_each_method(self):
         """Check that the unnormalisation restores the original data for each normalisation method."""
         da_raw = _gen_data_xr()
         df_raw = _gen_data_pandas()
 
-        dp = DataProcessor(
+        dp_with_x_mappings = DataProcessor(
             x1_map=(20, 40),
             x2_map=(40, 60),
             time_name="time",
             x1_name="lat",
             x2_name="lon",
         )
-
-        for method in dp.valid_methods:
-            da_norm, df_norm = dp([da_raw, df_raw], method=method)
-            da_unnorm, df_unnorm = dp.unnormalise([da_norm, df_norm])
-            self.assertTrue(
-                self.assert_allclose_xr(da_unnorm, da_raw),
-                f"Original {type(da_raw).__name__} not restored for method {method}.",
-            )
-            self.assertTrue(
-                self.assert_allclose_pd(df_unnorm, df_raw),
-                f"Original {type(df_raw).__name__} not restored for method {method}.",
-            )
+        dp_inferred_x_mappings = DataProcessor(
+            time_name="time", x1_name="lat", x2_name="lon"
+        )
+        dps = [dp_with_x_mappings, dp_inferred_x_mappings]
+
+        for dp in dps:
+            for method in dp.valid_methods:
+                da_norm, df_norm = dp([da_raw, df_raw], method=method)
+                da_unnorm, df_unnorm = dp.unnormalise([da_norm, df_norm])
+                self.assertTrue(
+                    self.assert_allclose_xr(da_unnorm, da_raw),
+                    f"Original {type(da_raw).__name__} not restored for method {method}.",
+                )
+                self.assertTrue(
+                    self.assert_allclose_pd(df_unnorm, df_raw),
+                    f"Original {type(df_raw).__name__} not restored for method {method}.",
+                )
 
     def test_different_names_xr(self):
         """

{deepsensor-0.1.11 → deepsensor-0.2.1}/tests/test_model.py

@@ -3,7 +3,7 @@ import itertools
 
 from parameterized import parameterized
 
-
+import os
 import xarray as xr
 import numpy as np
 import pandas as pd
@@ -11,11 +11,12 @@ import unittest
 
 import lab as B
 
-import deepsensor.tensorflow as deepsensor
+import deepsensor.torch as deepsensor
 
 from deepsensor.data.processor import DataProcessor
 from deepsensor.data.loader import TaskLoader
 from deepsensor.model.convnp import ConvNP
+from deepsensor.train.train import Trainer
 
 from tests.utils import gen_random_data_xr, gen_random_data_pandas
 
@@ -390,6 +391,56 @@ class TestModel(unittest.TestCase):
             mean_df.reset_index()["longitude"], df_raw.reset_index()["longitude"]
         )
 
+    def test_saving_and_loading(self):
+        """Test saving and loading of model"""
+        folder = f"tmp_{pd.Timestamp.now().strftime('%Y%m%d_%H%M%S')}"
+        os.makedirs(folder, exist_ok=True)
+
+        ds_raw = xr.tutorial.open_dataset("air_temperature")
+
+        data_processor = DataProcessor(x1_name="lat", x2_name="lon")
+        ds = data_processor(ds_raw)
+
+        t2m_fpath = f"{folder}/air_temperature_normalised.nc"
+        ds.to_netcdf(t2m_fpath)
+
+        task_loader = TaskLoader(context=t2m_fpath, target=t2m_fpath)
+
+        model = ConvNP(
+            data_processor, task_loader, unet_channels=(32,) * 3, verbose=False
+        )
+
+        # Train the model for a few iterations to test the trained model is restored correctly later.
+        task = task_loader("2014-12-31", 40, datewise_deterministic=True)
+        trainer = Trainer(model)
+        for _ in range(10):
+            trainer([task])
+        mean_ds_before, std_ds_before = model.predict(task, X_t=ds_raw)
+        mean_ds_before["air"].plot()
+
+        data_processor.save(folder)
+        task_loader.save(folder)
+        model.save(folder)
+
+        data_processor_loaded = DataProcessor(folder)
+        task_loader_loaded = TaskLoader(folder)
+        model_loaded = ConvNP(data_processor_loaded, task_loader_loaded, folder)
+
+        task = task_loader_loaded("2014-12-31", 40, datewise_deterministic=True)
+        mean_ds_loaded, std_ds_loaded = model_loaded.predict(task, X_t=ds_raw)
+        mean_ds_loaded["air"].plot()
+
+        xr.testing.assert_allclose(mean_ds_before, mean_ds_loaded)
+        print("Means match")
+
+        xr.testing.assert_allclose(std_ds_before, std_ds_loaded)
+        print("Standard deviations match")
+
+        # Delete temporary folder
+        import shutil
+
+        shutil.rmtree(folder)
+
 
 def assert_shape(x, shape: tuple):
     """ex: assert_shape(conv_input_array, [8, 3, None, None])"""

{deepsensor-0.1.11 → deepsensor-0.2.1}/tests/test_plotting.py

@@ -23,7 +23,7 @@ class TestPlotting(unittest.TestCase):
             x2_map=(ds_raw["lon"].min(), ds_raw["lon"].max()),
         )
         ds = self.data_processor(ds_raw)
-        self.task_loader = TaskLoader(ds, ds)
+        self.task_loader = TaskLoader(context=ds, target=ds)
         self.model = ConvNP(
             self.data_processor,
             self.task_loader,

{deepsensor-0.1.11 → deepsensor-0.2.1}/tests/test_task_loader.py

@@ -133,6 +133,12 @@ class TestTaskLoader(unittest.TestCase):
 
         return None
 
+    def test_saving_taskloader_without_paths_fails(self):
+        """Saving a TaskLoader is only possible if it is instantiated with file paths"""
+        tl = TaskLoader(context=self.da, target=self.da)
+        with self.assertRaises(ValueError):
+            tl.save("foo")
+
     def test_aux_at_contexts_and_aux_at_targets(self):
         """Test the `aux_at_contexts` and `aux_at_targets` arguments"""
         context = [self.da, self.df]