+
+
![\"ConvNP](\"../../figs/convnp_arch.png\")
"
- ]
+ ],
+ "metadata": {
+ "collapsed": false
+ }
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
- "The ConvCNP learns: $p_\\theta ( \\mathbf{y}^{(t)} ; \\mathbf{x}^{(t)}, \\phi(C))$,\n",
- "\n",
+ "The ConvCNP learns: $p_\\theta ( \\mathbf{y}^{(t)} ; \\mathbf{x}^{(t)}, \\phi(C))$, \n",
"where $p_\\theta$ specifies a predictive distribution over target values $\\mathbf{y}^{(t)}$ given target locations $\\mathbf{x}^{(t)}$ and the context data $C$.\n",
- "The distribution parameters $\\phi$ (such as the mean and variance of a Gaussian likelihood) *depend on the context set $C$*.\n",
+ "All the distribution parameters $\\phi$ (such as the mean and variance of a Gaussian likelihood) *depend on the entire context set $C$, including the $y$-values*!\n",
+ "One effect of this is that the ConvNP's uncertainty can hypothetically become *more uncertain* if neighbouring observations disagree.\n",
+ "This makes the ConvNP far more flexible than standard probabilistic models, like GPs.\n",
"\n",
- "This page will walk through the ConvNP architecture, explain important features, and provide some tips for using the model.\n",
- "For further background reading, check out the [](../resources.md) page.\n",
- "\n"
+ "Other benefits of the ConvNP are that it can handle:\n",
+ "* Fusing multiple context sets\n",
+ "* Off-the-grid and gridded modalities\n",
+ "* Multi-resolution data\n",
+ "* Missing data\n",
+ "* Predicting at arbitrary target locations\n",
+ "* Predicting multiple (disjoint) target sets\n",
+ "* Uncertainty quantification\n",
+ "* $O(N)$ inference cost (for most variants)\n",
+ "\n",
+ "```{caution} \n",
+ "It's important to note that the flexibility of the ConvNP makes it a *data hungry model*.\n",
+ "It has to *learn how to condition on data* from scratch, which requires\n",
+ "seeing many examples of different context sets and target sets during training.\n",
+ "The `TaskLoader` provides the necessary functionality for this.\n",
+ "However, if your data is spatially or temporally sparse (and there isn't \n",
+ "a more abundant dataset to transfer learn from), then the ConvNP may not be\n",
+ "the best model for your use case.\n",
+ "```\n",
+ "\n",
+ "```{note}\n",
+ "The 'ConvNP' is actually a *class* of model architectures, rather than a single model.\n",
+ "There are different variants of the ConvNP, such as the conditional NP (ConvCNP), the Gaussian NP (ConvGNP), and the\n",
+ "latent NP (ConvLNP). Each of these can be instantiated from the `ConvNP` class by initialising the model with\n",
+ "different hyperparameters.\n",
+ "```"
]
},
{
- "cell_type": "code",
- "execution_count": 1,
+ "cell_type": "markdown",
+ "source": [
+ "## Initialising a ConvNP\n",
+ "\n",
+ "The key arguments for initialising a `ConvNP` are:\n",
+ "* A `DataProcessor`: Used to inherit the high-level `.predict` method from the `DeepSensorModel` base class (detailed in the [](./prediction.ipynb) page).\n",
+ "* A `TaskLoader`: Under the hood, DeepSensor will use the `TaskLoader`'s context/target data and configuration to infer sensible defaults for hyperparameters that are not set by the user. "
+ ],
"metadata": {
- "ExecuteTime": {
- "end_time": "2023-10-27T16:34:23.635801716Z",
- "start_time": "2023-10-27T16:34:17.148046078Z"
- },
- "tags": [
- "hide-input"
- ]
- },
+ "collapsed": false
+ }
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
"outputs": [],
"source": [
"import logging\n",
@@ -80,17 +122,33 @@
"doy_ds = construct_circ_time_ds(dates, freq=\"D\")\n",
"land_mask_ds[\"cos_D\"] = doy_ds[\"cos_D\"]\n",
"land_mask_ds[\"sin_D\"] = doy_ds[\"sin_D\"]"
- ]
+ ],
+ "metadata": {
+ "collapsed": false,
+ "tags": [
+ "hide-cell"
+ ],
+ "is_executing": true
+ }
},
{
"cell_type": "code",
- "execution_count": 2,
+ "execution_count": 3,
"metadata": {
"ExecuteTime": {
- "end_time": "2023-10-27T16:34:23.840652157Z",
- "start_time": "2023-10-27T16:34:23.611416160Z"
+ "end_time": "2023-10-27T16:34:26.726834767Z",
+ "start_time": "2023-10-27T16:34:23.841387529Z"
}
},
+ "outputs": [],
+ "source": [
+ "import deepsensor.torch\n",
+ "from deepsensor.model import ConvNP"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 2,
"outputs": [
{
"name": "stdout",
@@ -111,22 +169,13 @@
" links=[(0, 0)],\n",
")\n",
"print(task_loader)"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": 3,
+ ],
"metadata": {
+ "collapsed": false,
"ExecuteTime": {
- "end_time": "2023-10-27T16:34:26.726834767Z",
- "start_time": "2023-10-27T16:34:23.841387529Z"
+ "start_time": "2023-10-27T16:34:23.611416160Z"
}
- },
- "outputs": [],
- "source": [
- "import deepsensor.torch\n",
- "from deepsensor.model import ConvNP"
- ]
+ }
},
{
"cell_type": "code",
@@ -347,7 +396,7 @@
{
"cell_type": "markdown",
"source": [
- "## Breaking `ConvNP` stationarity using auxiliary variables\n",
+ "## Breaking ConvNP stationarity using auxiliary variables\n",
"\n",
"The convolutional architecture of U-Net module of the `ConvNP` is translation equivariant.\n",
"This means shifting the context data results in a corresponding shift in the output prediction.\n",
@@ -439,7 +488,7 @@
{
"cell_type": "markdown",
"source": [
- "## The `likelihood` parameter\n",
+ "## The likelihood parameter\n",
"\n",
"The `likelihood` parameter specifies the output distribution of the `ConvNP`.\n",
"\n",
@@ -456,7 +505,7 @@
"cell_type": "markdown",
"metadata": {},
"source": [
- "## Saving and loading a `ConvNP`\n",
+ "## Saving and loading a ConvNP\n",
"\n",
"Similarly to the ``DataProcessor`` and ``TaskLoader`` objects, a ``ConvNP`` object can be saved using the ``save`` method, and then loaded using the ``ConvNP`` constructor."
]
diff --git a/docs/user-guide/data_processor.ipynb b/docs/user-guide/data_processor.ipynb
index e7145fd0..e37a27d6 100644
--- a/docs/user-guide/data_processor.ipynb
+++ b/docs/user-guide/data_processor.ipynb
@@ -11,10 +11,11 @@
"cell_type": "markdown",
"metadata": {},
"source": [
- "The first step in any modelling pipeline is to to normalise and standardise the data.\n",
- "In DeepSensor, this is achieved with the `DataProcessor` class.\n",
+ "The first step in any modelling pipeline is to normalise and standardise the data.\n",
+ "In DeepSensor, this is achieved with the [`DataProcessor` class](../reference/data/processor.rst).\n",
+ "Let's load some environmental data from `deepsensor.data.sources` and see how it works!\n",
"\n",
- "The `DataProcessor` does xyz..."
+ "Note, to run this yourself you will need to `pip install git+https://github.com/scott-hosking/get-station-data.git`."
]
},
{
@@ -86,6 +87,17 @@
"land_mask_raw_ds = get_gldas_land_mask(extent, cache=True, cache_dir=cache_dir)"
]
},
+ {
+ "cell_type": "markdown",
+ "source": [
+ "## Initialising a DataProcessor\n",
+ "\n",
+ "To initialise a `DataProcessor` object, provide it with the names of the spatiotemporal dimensions in your data (defaults to `time`, `x1`, `x2`)."
+ ],
+ "metadata": {
+ "collapsed": false
+ }
+ },
{
"cell_type": "code",
"execution_count": 4,
@@ -108,11 +120,21 @@
}
],
"source": [
- "# Normalise data\n",
"data_processor = DataProcessor(x1_name=\"lat\", x2_name=\"lon\")\n",
"print(data_processor)"
]
},
+ {
+ "cell_type": "markdown",
+ "source": [
+ "## Normalising data with a DataProcessor\n",
+ "\n",
+ "Calling a `DataProcessor` on an xarray or pandas object will compute normalisation parameters for each variable in the object (if not already computed) and return the normalised object/s."
+ ],
+ "metadata": {
+ "collapsed": false
+ }
+ },
{
"cell_type": "code",
"execution_count": 5,
@@ -199,7 +221,7 @@
"cell_type": "markdown",
"metadata": {},
"source": [
- "## `DataProcessor` configuration\n",
+ "## DataProcessor configuration\n",
"The `DataProcessor` keeps track of the normalisation parameters used to transform the data.\n",
"\n",
"When the `DataProcessor` is called on data with a variable ID that matches one in the `config` dictionary, those normalisation parameters are used."
@@ -314,14 +336,15 @@
"source": [
"Notice how the units of the unnormalised data are the same as the original data.\n",
"\n",
- "This functionality will come in handy later when we want to map model predictions back to the original units!\n"
+ "This functionality will be used under the hood\n",
+ "later to map model predictions back to the original units!\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
- "## Saving and loading a `DataProcessor`\n",
+ "## Saving and loading a DataProcessor\n",
"The `DataProcessor` configuration can be saved and loaded to avoid re-computing the normalisation parameters in new sessions.\n",
"\n",
"This is done using the `save` method to save the `DataProcessor` configuration to a folder, and then instantiating a new `DataProcessor` with the same folder will recover the same `DataProcessor` object."
diff --git a/docs/user-guide/deepsensor_design.md b/docs/user-guide/deepsensor_design.md
index 737cdb18..04680f64 100644
--- a/docs/user-guide/deepsensor_design.md
+++ b/docs/user-guide/deepsensor_design.md
@@ -5,8 +5,7 @@ before they begin. Others would prefer to just see some examples and
get started right away.
If you fall into the latter category,
-feel free to skip this section and go straight to the [](data_processor.ipynb)
-section. It may be useful to come back to this section later.
+feel free to jump straight to the next page ([](data_processor.ipynb)).
## Design overview
@@ -32,9 +31,9 @@ Uses the `neuralprocesses` library. This is currently the only model provided by
In addition, a [`deepsensor.plot`](../reference/plot.rst) module provides useful plotting functions for
visualising:
-* Raw data,
-* ``ConvNP`` internals (encoding and feature maps),
+* `Task` context and target sets,
* ``DeepSensorModel`` predictions,
+* ``ConvNP`` internals (encoding and feature maps),
* ``GreedyAlgorithm`` active learning outputs.
You will see examples of these `deepsensor.plot` visualisation functions
diff --git a/docs/user-guide/index.md b/docs/user-guide/index.md
index a1d6800e..5557f6b5 100644
--- a/docs/user-guide/index.md
+++ b/docs/user-guide/index.md
@@ -1,4 +1,10 @@
# User Guide
-In this user guide, you will find detailed descriptions and examples that describe many common
-tasks that you can accomplish with DeepSensor.
+The DeepSensor user guide will walk you through the core components of the package using
+code examples and visualisations.
+
+The pages of this guide are Jupyter notebooks and are fully-reproducible.
+Some of the notebooks depend on previous notebooks to be run, e.g. model training.
+However, most of the notebooks can be run in a standalone way without any modification.
+Click the download button at the top of the pages to download the .ipynb files
+if you would like to run them yourself.
diff --git a/docs/user-guide/prediction.ipynb b/docs/user-guide/prediction.ipynb
index af054fa5..01bef1ba 100644
--- a/docs/user-guide/prediction.ipynb
+++ b/docs/user-guide/prediction.ipynb
@@ -7,9 +7,12 @@
"# Prediction\n",
"\n",
"DeepSensor provides a convenient high-level interface to predict directly to `xarray` or `pandas` objects\n",
- "in the original units and coordinate system of your data. This is achieved using the `model.predict` method.\n",
+ "in the original units and coordinate system of your data. This is achieved using the [`model.predict` method](https://tom-andersson.github.io/deepsensor/reference/model/model.html#deepsensor.model.model.DeepSensorModel.predict).\n",
+ "We'll use our trained model from the [](./training.ipynb) page to demonstrate DeepSensor's prediction functionality.\n",
"\n",
- "We'll use our trained model from the [](./training.ipynb) page to demonstrate this functionality."
+ "The two key arguments of `model.predict` are 1) a `Task` (or list of `Task`s) containing context data, and 2) a set of target prediction locations, `X_t`.\n",
+ "This page will demonstrate how we can predict on-grid or off-grid based on the form of `X_t`.\n",
+ "We will also see how we can use optional extra arguments in `model.predict` for more advanced usage."
]
},
{
@@ -182,9 +185,10 @@
"cell_type": "markdown",
"metadata": {},
"source": [
- "## Predict on-grid to `xarray`\n",
+ "## Predict on-grid to xarray\n",
"\n",
- "We will use the high-level `model.predict` method, outlined in the next page"
+ "If `X_t` is an `xarray` object, `model.predict` will return `xarray` predictions on the same grid as that object\n",
+ "(with the resolution optionally scaled by `resolution_factor`)."
]
},
{
@@ -207,7 +211,9 @@
"cell_type": "markdown",
"metadata": {},
"source": [
- "After training, the model can predict directly to `xarray` in your data's original units and coordinate system:"
+ "The result is a single dictionary-like [`Prediction`](../reference/model/pred.rst) object\n",
+ "whose keys are the variable IDs of the target variables, each mapping to `xarray.Dataset` objects containing\n",
+ "the prediction parameters (in this case the `mean` and `std` of the `ConvNP`'s Gaussian likelihood)."
]
},
{
@@ -261,9 +267,13 @@
"cell_type": "markdown",
"metadata": {},
"source": [
- "## Predict off-grid to `pandas`\n",
+ "## Predict off-grid to pandas\n",
"\n",
- "Providing a `pandas` or `numpy` `X_t` object to `model.predict` will return a `pandas` object with predictions \n"
+ "Predicting at off-grid locations with `model.predict` is very similar to the on-grid case above.\n",
+ "If `X_t` is 1) a shape $(2, N)$ `numpy` array, or 2) a `pandas` object containing spatial indexes, the values of the `Prediction`\n",
+ "returned by `model.predict` will be `pandas.DataFrame`s whose columns are the prediction parameters.\n",
+ "\n",
+ "Let's see an example where we pass a *list* of `Task`s to `model.predict`, with context sets spanning the second half of 2019. Check out the indexes of the resulting `pandas.DataFrame`!"
]
},
{
@@ -277,7 +287,7 @@
},
"outputs": [],
"source": [
- "# Predict at two off-grid locations over December 2019 with 30 random, fixed context points\n",
+ "# Predict at two off-grid locations over six months of 2019 with 200 random context points (fixed across time)\n",
"test_tasks = task_loader(pd.date_range(\"2019-06-01\", \"2019-12-31\"), [200, \"all\", \"all\"], seed_override=42)\n",
"X_t = np.array([[50, -80],\n",
" [40, -110]]).T\n",
@@ -393,10 +403,9 @@
"This is achieved using by passing an integer `ar_subsample_factor > 1`.\n",
"`X_t` will be subsampled by a factor of `ar_subsample_factor` in both\n",
"spatial dimensions to obtain the AR sample locations.\n",
- "Note, this approach will result in a loss of spatial granularity in the AR samples.\n",
+ "Note, subsampling will result in a loss of spatial granularity in the AR samples.\n",
"\n",
- "For more information, check out the [Autoregressive Conditional Neural Processes\n",
- "](https://arxiv.org/abs/2303.14468) paper (ICLR, 2023)."
+ "For more information, check out the [Autoregressive Conditional Neural Processes](https://arxiv.org/abs/2303.14468) paper (ICLR, 2023)."
]
},
{
@@ -444,7 +453,7 @@
"cell_type": "markdown",
"metadata": {},
"source": [
- "Plotting the difference between each AR sample and the mean prediction can highlight spatial correlations in the samples:"
+ "Let's plot the *difference* between each AR sample and the mean prediction to highlight the spatial correlations we get in the AR samples."
]
},
{
@@ -482,6 +491,14 @@
"plt.show()"
]
},
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "The checkerboard artefacts in the samples above correspond to the subsampled grid of AR sample target locations.\n",
+ "This could be alleviated by training the model for longer and with more examples of larger/denser context sets."
+ ]
+ },
{
"cell_type": "code",
"execution_count": 16,
diff --git a/docs/user-guide/task.ipynb b/docs/user-guide/task.ipynb
index 8c1a2860..34fb8d23 100644
--- a/docs/user-guide/task.ipynb
+++ b/docs/user-guide/task.ipynb
@@ -7,54 +7,6 @@
"# Tasks"
]
},
- {
- "cell_type": "code",
- "execution_count": 1,
- "metadata": {
- "ExecuteTime": {
- "end_time": "2023-11-01T14:32:15.548664564Z",
- "start_time": "2023-11-01T14:28:15.732009455Z"
- }
- },
- "outputs": [
- {
- "name": "stderr",
- "output_type": "stream",
- "text": [
- "100%|████████████████████████████████████████████████████████████████| 3124/3124 [02:38<00:00, 19.75it/s]\n"
- ]
- }
- ],
- "source": [
- "import logging\n",
- "\n",
- "logging.captureWarnings(True)\n",
- "\n",
- "import deepsensor.torch\n",
- "from deepsensor.data import DataProcessor\n",
- "from deepsensor.data.sources import get_ghcnd_station_data, get_era5_reanalysis_data, get_earthenv_auxiliary_data, get_gldas_land_mask\n",
- "\n",
- "import matplotlib.pyplot as plt\n",
- "\n",
- "# Using the same settings allows use to use pre-downloaded cached data\n",
- "data_range = (\"2016-06-25\", \"2016-06-30\")\n",
- "extent = \"europe\"\n",
- "station_var_IDs = [\"TAVG\", \"PRCP\"]\n",
- "era5_var_IDs = [\"2m_temperature\", \"10m_u_component_of_wind\", \"10m_v_component_of_wind\"]\n",
- "auxiliary_var_IDs = [\"elevation\", \"tpi\"]\n",
- "cache_dir = \"../../.datacache\"\n",
- "\n",
- "station_raw_df = get_ghcnd_station_data(station_var_IDs, extent, date_range=data_range, cache=True, cache_dir=cache_dir)\n",
- "era5_raw_ds = get_era5_reanalysis_data(era5_var_IDs, extent, date_range=data_range, cache=True, cache_dir=cache_dir)\n",
- "auxiliary_raw_ds = get_earthenv_auxiliary_data(auxiliary_var_IDs, extent, \"10KM\", cache=True, cache_dir=cache_dir)\n",
- "land_mask_raw_ds = get_gldas_land_mask(extent, cache=True, cache_dir=cache_dir)\n",
- "\n",
- "data_processor = DataProcessor(x1_name=\"lat\", x2_name=\"lon\")\n",
- "era5_ds = data_processor(era5_raw_ds)\n",
- "aux_ds, land_mask_ds = data_processor([auxiliary_raw_ds, land_mask_raw_ds], method=\"min_max\")\n",
- "station_df = data_processor(station_raw_df)"
- ]
- },
{
"cell_type": "markdown",
"metadata": {},
@@ -74,26 +26,30 @@
"```{admonition} Click to reveal the meta-learning primer\n",
":class: dropdown\n",
"\n",
- "### Sets of observations\n",
+ "**Sets of observations**\n",
+ "\n",
"A *set* of observations is a collection of $M$ input-output pairs $\\{(\\mathbf{x}_1, \\mathbf{y}_1), (\\mathbf{x}_2, \\mathbf{y}_2), \\ldots, (\\mathbf{x}_M, \\mathbf{y}_M)\\}$.\n",
"In DeepSensor $\\mathbf{x}_i \\in \\mathbb{R}^2$ is a 2D spatial location (such as latitude-longitude)\n",
" and $\\mathbf{y}_i \\in \\mathbb{R}^N$ is an $N$-dimensional observation at that location (such as a temperature and precipitation).\n",
"Context sets may lie on scattered, off-grid locations (such as weather stations), or on a regular grid (such as a reanalysis or satellite data).\n",
"A *set* can be compactly written as $(\\mathbf{X}, \\mathbf{Y})$, where $\\mathbf{X} \\in \\mathbb{R}^{2\\times M}$ and $\\mathbf{Y} \\in \\mathbb{R}^{N\\times M}$.\n",
"\n",
- "### Context sets\n",
- "A *context set* is a set of observations that are used to make predictions for another set of observations. Following our notations above, we denote a context set as $C=(\\mathbf{X}^{(c)}, \\mathbf{Y}^{(c)})_i$.\n",
- "We may have multiple context sets, denoted as $C_i = \\{ (\\mathbf{X}^{(c)}, \\mathbf{Y}^{(c)})_i \\}_{i=1}^{N_C}$.\n",
+ "**Context sets**\n",
+ "\n",
+ "A *context set* is a set of observations that are used to make predictions for another set of observations. Following our notations above, we denote a context set as $C_j=(\\mathbf{X}^{(c)}, \\mathbf{Y}^{(c)})_j$.\n",
+ "We may have multiple context sets, denoted as $C = \\{ (\\mathbf{X}^{(c)}, \\mathbf{Y}^{(c)})_j \\}_{j=1}^{N_C}$.\n",
+ "\n",
+ "**Target sets**\n",
"\n",
- "### Target sets\n",
"A *target set* is a set of observations that we wish to predict using the context sets.\n",
- "Similarly to contet sets, we denote the collection of all target sets as $T = \\{ (\\mathbf{X}^{(t)}, \\mathbf{Y}^{(t)})_i \\}_{i=1}^{N_T}$.\n",
- "During training, the target observations are known, but at inference time will be unknown latent variables.\n",
+ "Similarly to context sets, we denote the collection of all target sets as $T = \\{ (\\mathbf{X}^{(t)}, \\mathbf{Y}^{(t)})_j \\}_{j=1}^{N_T}$.\n",
+ "During training, the target observations $\\mathbf{y}_i$ are known, but at inference time will be unknown latent variables.\n",
+ "\n",
+ "**Tasks**\n",
"\n",
- "### Tasks\n",
"A *task* is a collection of context sets and target sets.\n",
"We denote a task as $\\mathcal{D} = (C, T)$.\n",
- "The modelling goal is make probabilistic predictions for the target variables $\\mathbf{Y}^{(t)}_i$ given the context sets $C$ and target prediction locations $\\mathbf{X}^{(t)}_i$.\n",
+ "The modelling goal is make probabilistic predictions for the target variables $\\mathbf{Y}^{(t)}_j$ given the context sets $C$ and target prediction locations $\\mathbf{X}^{(t)}_j$.\n",
"```"
]
},
@@ -101,7 +57,7 @@
"cell_type": "markdown",
"metadata": {},
"source": [
- "## The DeepSensor `Task`\n",
+ "## The DeepSensor Task\n",
"\n",
"In DeepSensor, a `Task` is a `dict`-like data structure that contains context sets, target sets, and other metadata.\n",
"Before diving into the [](./task_loader) class which generates `Task` objects from `xarray` and `pandas` objects,\n",
@@ -110,6 +66,57 @@
"Printing a `Task` will print each of its entries and replace numerical arrays with their shape for convenience."
]
},
+ {
+ "cell_type": "code",
+ "execution_count": 1,
+ "outputs": [
+ {
+ "name": "stderr",
+ "output_type": "stream",
+ "text": [
+ "100%|████████████████████████████████████████████████████████████████| 3124/3124 [02:38<00:00, 19.75it/s]\n"
+ ]
+ }
+ ],
+ "source": [
+ "import logging\n",
+ "\n",
+ "logging.captureWarnings(True)\n",
+ "\n",
+ "import deepsensor.torch\n",
+ "from deepsensor.data import DataProcessor\n",
+ "from deepsensor.data.sources import get_ghcnd_station_data, get_era5_reanalysis_data, get_earthenv_auxiliary_data, get_gldas_land_mask\n",
+ "\n",
+ "import matplotlib.pyplot as plt\n",
+ "\n",
+ "# Using the same settings allows use to use pre-downloaded cached data\n",
+ "data_range = (\"2016-06-25\", \"2016-06-30\")\n",
+ "extent = \"europe\"\n",
+ "station_var_IDs = [\"TAVG\", \"PRCP\"]\n",
+ "era5_var_IDs = [\"2m_temperature\", \"10m_u_component_of_wind\", \"10m_v_component_of_wind\"]\n",
+ "auxiliary_var_IDs = [\"elevation\", \"tpi\"]\n",
+ "cache_dir = \"../../.datacache\"\n",
+ "\n",
+ "station_raw_df = get_ghcnd_station_data(station_var_IDs, extent, date_range=data_range, cache=True, cache_dir=cache_dir)\n",
+ "era5_raw_ds = get_era5_reanalysis_data(era5_var_IDs, extent, date_range=data_range, cache=True, cache_dir=cache_dir)\n",
+ "auxiliary_raw_ds = get_earthenv_auxiliary_data(auxiliary_var_IDs, extent, \"10KM\", cache=True, cache_dir=cache_dir)\n",
+ "land_mask_raw_ds = get_gldas_land_mask(extent, cache=True, cache_dir=cache_dir)\n",
+ "\n",
+ "data_processor = DataProcessor(x1_name=\"lat\", x2_name=\"lon\")\n",
+ "era5_ds = data_processor(era5_raw_ds)\n",
+ "aux_ds, land_mask_ds = data_processor([auxiliary_raw_ds, land_mask_raw_ds], method=\"min_max\")\n",
+ "station_df = data_processor(station_raw_df)"
+ ],
+ "metadata": {
+ "collapsed": false,
+ "tags": [
+ "hide-cell"
+ ],
+ "ExecuteTime": {
+ "start_time": "2023-11-01T14:28:15.732009455Z"
+ }
+ }
+ },
{
"cell_type": "code",
"execution_count": 2,
@@ -160,7 +167,7 @@
"cell_type": "markdown",
"metadata": {},
"source": [
- "## `Task` structure\n",
+ "## Task structure\n",
"\n",
"A `Task` typically contains at least the following entries:\n",
"- `\"time\"`: timestamp that was used for slicing the spatiotemporal data.\n",
@@ -200,7 +207,7 @@
{
"cell_type": "markdown",
"source": [
- "### Gridded data in `Task`s\n",
+ "### Gridded data in Tasks\n",
"\n",
"For convenience, data that lies on a regular grid is given a compact tuple representation for the `\"X\"` entries:"
],
@@ -264,7 +271,7 @@
"cell_type": "markdown",
"metadata": {},
"source": [
- "### `Task` methods\n",
+ "### Task methods\n",
"The `Task` class also contains methods for applying processing operations the data (like removing NaNs, adding batch dimensions, etc.).\n",
"These operations will be recorded in the order they were applied the `\"ops\"` entry of the `Task`.\n",
"Operations can be chained together, for example:"
diff --git a/docs/user-guide/task_loader.ipynb b/docs/user-guide/task_loader.ipynb
index 00dffa3d..124cf33b 100644
--- a/docs/user-guide/task_loader.ipynb
+++ b/docs/user-guide/task_loader.ipynb
@@ -11,7 +11,7 @@
"cell_type": "markdown",
"metadata": {},
"source": [
- "The `TaskLoader` generates `Task` objects for training, testing, and inference with DeepSensor models.\n",
+ "[The `TaskLoader`](../reference/data/loader.rst) generates `Task` objects for training, testing, and inference with DeepSensor models.\n",
"The `TaskLoader` can generate `Task`s for different kinds of predictions, such as: spatial interpolation, forecasting, downscaling, or some combination of these.\n",
"It achieves this by temporally slicing spatiotemporal data and then providing a suite of spatial sampling methods for generating context and target sets from the temporal slices.\n",
"\n",
@@ -55,6 +55,9 @@
],
"metadata": {
"collapsed": false,
+ "tags": [
+ "hide-cell"
+ ],
"ExecuteTime": {
"end_time": "2023-11-02T13:41:03.577241738Z",
"start_time": "2023-11-02T13:40:57.642318144Z"
@@ -79,7 +82,7 @@
"cell_type": "markdown",
"metadata": {},
"source": [
- "## Initialising a `TaskLoader`\n",
+ "## Initialising a TaskLoader\n",
"\n",
"A `TaskLoader` is initialised with list of `context` and `target` variables.\n",
"These variables can either be `xarray` or `pandas` objects, and are assumed to have been standardised by a `DataProcessor`.\n",
@@ -121,7 +124,7 @@
"cell_type": "markdown",
"metadata": {},
"source": [
- "## Calling a `TaskLoader`\n",
+ "## Calling a TaskLoader\n",
"\n",
"The `TaskLoader` is called with a timestamp and single entries or lists of entries for the `context_sampling` and `target_sampling` arguments.\n",
"These arguments can either be single entries (applying the same sampling strategy to each context/target set)\n",
@@ -189,7 +192,7 @@
"cell_type": "markdown",
"metadata": {},
"source": [
- "## `TaskLoader` tour\n",
+ "## TaskLoader tour\n",
"\n",
"Now that we've got the basics of the `TaskLoader` init and call signatures,\n",
"let's see some concrete examples of to generate `Task` objects for\n",
@@ -519,7 +522,7 @@
"\n",
"Optionally, we can also pass a high-resolution auxiliary `xarray` variable (e.g. local topgraphic information) via the `aux_at_targets` argument.\n",
"When calling the `TaskLoader`, the `aux_at_targets` variable will be interpolated at the target locations and added to the `Task` as the `\"Y_t_aux\"` entry.\n",
- "The `\"Y_t_aux\"` data can be modelled differently from the context data, for example, by using a pointwise MLP rather than a convolutional neural network {cite}`perez2011python`.\n",
+ "The `\"Y_t_aux\"` data can be modelled differently from the context data, for example, by using a pointwise MLP rather than a convolutional neural network.\n",
"\n",
"```{note}\n",
"The `TaskLoader` also permits an `aux_at_contexts` argument for passing high-resolution auxiliary variables at off-grid context locations.\n",
@@ -660,7 +663,7 @@
"- the context points which were just removed by the new mask become the target set.\n",
"\n",
"This may produce NaNs in the target set due to overlap between the original mask and new mask.\n",
- "`Task.remove_context_nans` and `Task.remove_target_nans` methods will these NaNs."
+ "To remove these NaNs, use the `Task.remove_context_nans` and `Task.remove_target_nans` methods."
]
},
{
@@ -732,11 +735,12 @@
"This means that the true missing data is *not representative* of the observed\n",
"data; the distribution of observed data is very different to the distribution of missing data.\n",
"A model trained with the `TaskLoader` `\"gapfill\"` strategy will never see values below -0.75\n",
- "and will therefore be heavily biased.\n",
+ "and therefore its predictions will be heavily biased.\n",
"\n",
- "This effect can occur in real data sets. For example, if missing data comes from\n",
- "clouds and the target variable is temperature, a model will not learn the\n",
- "effect of clouds on temperature. Caution is advised when using the `\"gapfill\"` strategy.\n",
+ "This is an extreme example, but the same effect can occur in real data sets.\n",
+ "For example, if missing data comes from\n",
+ "clouds and the target variable is temperature, a model trained with the `\"gapfill\"` strategy\n",
+ "will not learn the effect of clouds on temperature.\n",
"```"
]
},
@@ -847,10 +851,10 @@
{
"cell_type": "markdown",
"source": [
- "## Controlling randomness in the `TaskLoader`\n",
+ "## Controlling randomness in the TaskLoader\n",
"\n",
"There are two additional arguments in the `TaskLoader` call method for controlling randomness:\n",
- "- `datewise_determinisic`: If `True`, the same random seed will be used for a particlar seed. Useful for,\n",
+ "- `datewise_deterministic`: If `True`, the same random seed will be used for a particlar date. Useful for,\n",
"say, generating a reproducible validation set which is the same between Python sessions.\n",
"- `seed_override`: If not `None`, this seed will be used instead of the default seed.\n"
],
@@ -862,7 +866,7 @@
"cell_type": "markdown",
"metadata": {},
"source": [
- "## Saving and loading a `TaskLoader`\n",
+ "## Saving and loading a TaskLoader\n",
"\n",
"You can save the `TaskLoader` object to a file using the `.save` method.\n",
"\n",
diff --git a/docs/user-guide/training.ipynb b/docs/user-guide/training.ipynb
index 0902b0d9..fd2a4a20 100644
--- a/docs/user-guide/training.ipynb
+++ b/docs/user-guide/training.ipynb
@@ -6,7 +6,8 @@
"source": [
"# Training\n",
"\n",
- "The `ConvNP` model can be trained using the `Trainer` class, which implements stochastic gradient descent using the Adam optimiser.\n",
+ "The `ConvNP` model can be trained using [the `Trainer` class](../reference/train/train.rst), which implements stochastic gradient descent using the Adam optimiser.\n",
+ "\n",
"\n",
"We will demonstrate training a `ConvNP` model with the following set-up:\n",
"- Goal: spatially interpolate daily-average ERA5 2m temperature over North America.\n",
@@ -143,7 +144,7 @@
"cell_type": "markdown",
"metadata": {},
"source": [
- "## Initialise `TaskLoader` and `ConvNP` model"
+ "## Initialise TaskLoader and ConvNP model"
]
},
{
@@ -207,7 +208,7 @@
"cell_type": "markdown",
"metadata": {},
"source": [
- "## Define how `Task`s are generated\n",
+ "## Define how Tasks are generated\n",
"\n",
"We will generate `Task`s for a set of dates by randomly sampling between 0 and 500 ERA5 grid cells as context points,\n",
"and passing all ERA5 grid cells as target points.\n",
@@ -309,7 +310,7 @@
"cell_type": "markdown",
"metadata": {},
"source": [
- "## Training with the `Trainer` class\n",
+ "## Training with the Trainer class\n",
"\n",
"The `Trainer` class is initialised with a `ConvNP` model, and an optional learning rate.\n",
"Calling a `Trainer` object with a list of `Task`s will perform a single epoch of training, and return a list of losses for each task.\n",
diff --git a/requirements/requirements.dev.txt b/requirements/requirements.dev.txt
index 3a8aa9d8..6240a200 100644
--- a/requirements/requirements.dev.txt
+++ b/requirements/requirements.dev.txt
@@ -4,4 +4,5 @@ pytest-cov
parameterized
tox
tox-gh-actions
-coveralls
\ No newline at end of file
+coveralls
+black
diff --git a/requirements/requirements.docs.txt b/requirements/requirements.docs.txt
index 8d93b132..2b475c56 100644
--- a/requirements/requirements.docs.txt
+++ b/requirements/requirements.docs.txt
@@ -1 +1,2 @@
-jupyter-book==0.15.1
\ No newline at end of file
+jupyter-book==0.15.1
+sphinx
diff --git a/requirements/requirements.txt b/requirements/requirements.txt
index 10586315..b92f5e3b 100644
--- a/requirements/requirements.txt
+++ b/requirements/requirements.txt
@@ -7,7 +7,7 @@ dask
distributed
neuralprocesses>=0.2.2
tensorflow
-tensorflow_probability
+tensorflow_probability[tf]
torch>=2
backends-matrix
backends
diff --git a/setup.cfg b/setup.cfg
index 9bf81623..d91c5a76 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -1,8 +1,8 @@
[metadata]
name = deepsensor
-version = 0.3.2
+version = 0.3.6
author = Tom R. Andersson
-author_email = tomand@bas.ac.uk
+author_email = tomandersson3@gmail.com
description = A Python package for modelling xarray and pandas data with neural processes.
long_description = file: README.md
long_description_content_type = text/markdown
@@ -37,7 +37,6 @@ install_requires =
pooch
gcsfs
zarr
- get-station-data @ git+https://github.com/scotthosking/get-station-data@main#egg=get-station-data
[options.extras_require]
testing =
diff --git a/tests/test_active_learning.py b/tests/test_active_learning.py
index e8690e5a..a6c6b905 100644
--- a/tests/test_active_learning.py
+++ b/tests/test_active_learning.py
@@ -181,6 +181,31 @@ def test_greedy_alg_with_sequential_acquisition_fn(self):
task = self.task_loader("2014-12-31", context_sampling=10)
_ = alg(acquisition_fn, task)
+ def test_greedy_algorithm_column_names(self):
+ # Setup
+ acquisition_fn = Stddev(self.model)
+ X_s = self.ds_raw
+ alg = GreedyAlgorithm(
+ model=self.model,
+ X_t=X_s,
+ X_s=X_s,
+ N_new_context=1,
+ task_loader=self.task_loader,
+ )
+ task = self.task_loader("2014-12-31", context_sampling=10)
+
+ # Exercise
+ X_new_df, acquisition_fn_ds = alg(acquisition_fn, task)
+
+ # Assert
+ expected_columns = ["lat", "lon"] # Replace with actual expected column names
+ actual_columns = X_new_df.columns.tolist()
+ self.assertEqual(
+ expected_columns,
+ actual_columns,
+ "Column names do not match the expected names",
+ )
+
def test_greedy_alg_with_aux_at_targets_without_task_loader_raises_value_error(
self,
):
diff --git a/tests/test_model.py b/tests/test_model.py
index e6465276..80519269 100644
--- a/tests/test_model.py
+++ b/tests/test_model.py
@@ -69,8 +69,7 @@ def _gen_task_loader_call_args(self, n_context, n_target):
]:
yield [sampling_method] * n_context, [sampling_method] * n_target
- # TEMP only 1D because non-overlapping target sets are not yet supported
- @parameterized.expand(range(1, 2))
+ @parameterized.expand(range(1, 4))
def test_model_call(self, n_context_and_target):
"""Check ``ConvNP`` runs with all possible combinations of context/target
sampling methods."""
@@ -111,77 +110,135 @@ def set_list_to_data(set_list):
@parameterized.expand(range(1, 4))
def test_prediction_shapes_lowlevel(self, n_target_sets):
- """Test low-level model prediction interface over a range of number of
- target sets."""
+ """
+ Test low-level model prediction interface over a range of number of
+ target sets.
+ """
+ # Make dataset 5D for non-trivial target dimensions
+ ndim = 5
+ ds = xr.Dataset(
+ {f"var{i}": self.da for i in range(ndim)},
+ )
tl = TaskLoader(
- context=self.da,
- target=[self.da] * n_target_sets,
+ context=ds,
+ target=[ds] * n_target_sets,
)
context_sampling = 10
- model = ConvNP(self.dp, tl, unet_channels=(5, 5, 5), verbose=False)
+ likelihoods = ["cnp", "gnp", "cnp-spikes-beta"]
- for target_sampling, expected_obs_shape in (
- (10, (10,)), # expected shape is (10,) when target_sampling is 10
- (
- "all",
- self.da.shape[-2:],
- ), # expected shape is da.shape[-2:] when target_sampling is "all"
- ):
- task = tl("2020-01-01", context_sampling, target_sampling)
-
- n_targets = np.product(expected_obs_shape)
-
- # Tensors
- mean = model.mean(task)
- # TODO avoid repeated code
- if isinstance(mean, (list, tuple)):
- for m, dim_y in zip(mean, tl.target_dims):
- assert_shape(m, (dim_y, *expected_obs_shape))
- else:
- assert_shape(mean, (n_target_sets, *expected_obs_shape))
-
- variance = model.variance(task)
- if isinstance(variance, (list, tuple)):
- for v, dim_y in zip(variance, tl.target_dims):
- assert_shape(v, (dim_y, *expected_obs_shape))
- else:
- assert_shape(variance, (n_target_sets, *expected_obs_shape))
-
- stddev = model.stddev(task)
- if isinstance(stddev, (list, tuple)):
- for s, dim_y in zip(stddev, tl.target_dims):
- assert_shape(s, (dim_y, *expected_obs_shape))
- else:
- assert_shape(stddev, (n_target_sets, *expected_obs_shape))
-
- n_samples = 5
- samples = model.sample(task, n_samples)
- if isinstance(samples, (list, tuple)):
- for s, dim_y in zip(samples, tl.target_dims):
- assert_shape(s, (n_samples, dim_y, *expected_obs_shape))
- else:
- assert_shape(samples, (n_samples, n_target_sets, *expected_obs_shape))
-
- n_target_dims = np.product(tl.target_dims)
- assert_shape(
- model.covariance(task),
- (
- n_targets * n_target_sets * n_target_dims,
- n_targets * n_target_sets * n_target_dims,
- ),
+ dim_y_combined = sum(tl.target_dims)
+
+ for likelihood in likelihoods:
+ model = ConvNP(
+ self.dp,
+ tl,
+ unet_channels=(5, 5, 5),
+ likelihood=likelihood,
+ verbose=False,
)
+ assert dim_y_combined == model.config["dim_yt"]
+
+ for target_sampling, expected_obs_shape in (
+ # expected obs shape is (10,) when target_sampling is 10
+ (10, (10,)),
+ # expected obs shape is da.shape[-2:] when target_sampling is "all"
+ ("all", self.da.shape[-2:]),
+ ):
+ task = tl("2020-01-01", context_sampling, target_sampling)
- # Scalars
- x = model.logpdf(task)
- assert x.size == 1 and x.shape == ()
- x = model.joint_entropy(task)
- assert x.size == 1 and x.shape == ()
- x = model.mean_marginal_entropy(task)
- assert x.size == 1 and x.shape == ()
- x = B.to_numpy(model.loss_fn(task))
- assert x.size == 1 and x.shape == ()
+ n_targets = np.product(expected_obs_shape)
+
+ # Tensors
+ mean = model.mean(task)
+ # TODO avoid repeated code by looping over methods
+ if isinstance(mean, (list, tuple)):
+ for m, dim_y in zip(mean, tl.target_dims):
+ assert_shape(m, (dim_y, *expected_obs_shape))
+ else:
+ assert_shape(mean, (dim_y_combined, *expected_obs_shape))
+
+ variance = model.variance(task)
+ if isinstance(variance, (list, tuple)):
+ for v, dim_y in zip(variance, tl.target_dims):
+ assert_shape(v, (dim_y, *expected_obs_shape))
+ else:
+ assert_shape(variance, (dim_y_combined, *expected_obs_shape))
+
+ stddev = model.stddev(task)
+ if isinstance(stddev, (list, tuple)):
+ for s, dim_y in zip(stddev, tl.target_dims):
+ assert_shape(s, (dim_y, *expected_obs_shape))
+ else:
+ assert_shape(stddev, (dim_y_combined, *expected_obs_shape))
+
+ n_samples = 5
+ samples = model.sample(task, n_samples)
+ if isinstance(samples, (list, tuple)):
+ for s, dim_y in zip(samples, tl.target_dims):
+ assert_shape(s, (n_samples, dim_y, *expected_obs_shape))
+ else:
+ assert_shape(
+ samples, (n_samples, dim_y_combined, *expected_obs_shape)
+ )
+
+ if likelihood in ["cnp", "gnp"]:
+ n_target_dims = np.product(tl.target_dims)
+ assert_shape(
+ model.covariance(task),
+ (
+ n_targets * dim_y_combined * n_target_dims,
+ n_targets * dim_y_combined * n_target_dims,
+ ),
+ )
+ if likelihood in ["cnp-spikes-beta"]:
+ mixture_probs = model.mixture_probs(task)
+ if isinstance(mixture_probs, (list, tuple)):
+ for p, dim_y in zip(mixture_probs, tl.target_dims):
+ assert_shape(
+ p,
+ (
+ model.N_mixture_components,
+ dim_y,
+ *expected_obs_shape,
+ ),
+ )
+ else:
+ assert_shape(
+ mixture_probs,
+ (
+ model.N_mixture_components,
+ dim_y_combined,
+ *expected_obs_shape,
+ ),
+ )
+
+ x = model.alpha(task)
+ if isinstance(x, (list, tuple)):
+ for p, dim_y in zip(x, tl.target_dims):
+ assert_shape(p, (dim_y, *expected_obs_shape))
+ else:
+ assert_shape(x, (dim_y_combined, *expected_obs_shape))
+
+ x = model.beta(task)
+ if isinstance(x, (list, tuple)):
+ for p, dim_y in zip(x, tl.target_dims):
+ assert_shape(p, (dim_y, *expected_obs_shape))
+ else:
+ assert_shape(x, (dim_y_combined, *expected_obs_shape))
+
+ # Scalars
+ if likelihood in ["cnp", "gnp"]:
+ # Methods for Gaussian likelihoods only
+ x = model.joint_entropy(task)
+ assert x.size == 1 and x.shape == ()
+ x = model.mean_marginal_entropy(task)
+ assert x.size == 1 and x.shape == ()
+ x = model.logpdf(task)
+ assert x.size == 1 and x.shape == ()
+ x = B.to_numpy(model.loss_fn(task))
+ assert x.size == 1 and x.shape == ()
@parameterized.expand(range(1, 4))
def test_prediction_shapes_highlevel(self, target_dim):
@@ -212,9 +269,9 @@ def test_prediction_shapes_highlevel(self, target_dim):
tasks,
X_t=self.da,
n_samples=n_samples,
- unnormalise=True
- if target_dim == 1
- else False, # TODO fix unnormalising for multiple equally named targets
+ unnormalise=(
+ True if target_dim == 1 else False
+ ), # TODO fix unnormalising for multiple equally named targets
)
assert [isinstance(ds, xr.Dataset) for ds in pred.values()]
for var_ID in pred:
@@ -372,9 +429,7 @@ def test_highlevel_predict_coords_align_with_X_t_offgrid(self):
dp = DataProcessor(
x1_name="latitude",
- x1_map=lat_lims,
x2_name="longitude",
- x2_map=lon_lims,
)
df = dp(df_raw)
@@ -392,6 +447,81 @@ def test_highlevel_predict_coords_align_with_X_t_offgrid(self):
df_raw.reset_index()["longitude"],
)
+ def test_highlevel_predict_with_pred_params_pandas(self):
+ """
+ Test that passing ``pred_params`` to ``.predict`` works with
+ a spikes-beta likelihood for prediction to pandas.
+ """
+ tl = TaskLoader(context=self.da, target=self.da)
+ model = ConvNP(
+ self.dp,
+ tl,
+ unet_channels=(5, 5, 5),
+ verbose=False,
+ likelihood="cnp-spikes-beta",
+ )
+ task = tl("2020-01-01", context_sampling=10, target_sampling=10)
+
+ # Off-grid prediction
+ X_t = np.array([[0.0, 0.5, 1.0], [0.0, 0.5, 1.0]])
+
+ # Check that nothing breaks and the correct parameters are returned
+ pred_params = ["mean", "std", "variance", "alpha", "beta"]
+ pred = model.predict(task, X_t=X_t, pred_params=pred_params)
+ for pred_param in pred_params:
+ assert pred_param in pred["var"]
+
+ # Test mixture probs special case
+ pred_params = ["mixture_probs"]
+ pred = model.predict(task, X_t=self.da, pred_params=pred_params)
+ for component in range(model.N_mixture_components):
+ pred_param = f"mixture_probs_{component}"
+ assert pred_param in pred["var"]
+
+ def test_highlevel_predict_with_pred_params_xarray(self):
+ """
+ Test that passing ``pred_params`` to ``.predict`` works with
+ a spikes-beta likelihood for prediction to xarray.
+ """
+ tl = TaskLoader(context=self.da, target=self.da)
+ model = ConvNP(
+ self.dp,
+ tl,
+ unet_channels=(5, 5, 5),
+ verbose=False,
+ likelihood="cnp-spikes-beta",
+ )
+ task = tl("2020-01-01", context_sampling=10, target_sampling=10)
+
+ # Check that nothing breaks and the correct parameters are returned
+ pred_params = ["mean", "std", "variance", "alpha", "beta"]
+ pred = model.predict(task, X_t=self.da, pred_params=pred_params)
+ for pred_param in pred_params:
+ assert pred_param in pred["var"]
+
+ # Test mixture probs special case
+ pred_params = ["mixture_probs"]
+ pred = model.predict(task, X_t=self.da, pred_params=pred_params)
+ for component in range(model.N_mixture_components):
+ pred_param = f"mixture_probs_{component}"
+ assert pred_param in pred["var"]
+
+ def test_highlevel_predict_with_invalid_pred_params(self):
+ """Test that passing ``pred_params`` to ``.predict`` works."""
+ tl = TaskLoader(context=self.da, target=self.da)
+ model = ConvNP(
+ self.dp,
+ tl,
+ unet_channels=(5, 5, 5),
+ verbose=False,
+ likelihood="cnp-spikes-beta",
+ )
+ task = tl("2020-01-01", context_sampling=10, target_sampling=10)
+
+ # Check that passing an invalid parameter raises an AttributeError
+ with self.assertRaises(AttributeError):
+ model.predict(task, X_t=self.da, pred_params=["invalid_param"])
+
def test_saving_and_loading(self):
"""Test saving and loading of model"""
with tempfile.TemporaryDirectory() as folder:
diff --git a/tests/test_training.py b/tests/test_training.py
index 084b82ed..8b2c1f2c 100644
--- a/tests/test_training.py
+++ b/tests/test_training.py
@@ -113,3 +113,46 @@ def test_training(self):
# Check for NaNs in the loss
loss = np.mean(epoch_losses)
self.assertFalse(np.isnan(loss))
+
+ def test_training_multidim(self):
+ """A basic test of the training loop with multidimensional context sets"""
+ # Load raw data
+ ds_raw = xr.tutorial.open_dataset("air_temperature")
+
+ # Add extra dim
+ ds_raw["air2"] = ds_raw["air"].copy()
+
+ # Normalise data
+ dp = DataProcessor(x1_name="lat", x2_name="lon")
+ ds = dp(ds_raw)
+
+ # Set up task loader
+ tl = TaskLoader(context=ds, target=ds)
+
+ # Set up model
+ model = ConvNP(dp, tl)
+
+ # Generate training tasks
+ n_train_tasks = 10
+ train_tasks = []
+ for i in range(n_train_tasks):
+ date = np.random.choice(self.da.time.values)
+ task = tl(date, 10, 10)
+ task["Y_c"][0][:, 0] = np.nan # Add NaN to context
+ task["Y_t"][0][:, 0] = np.nan # Add NaN to target
+ print(task)
+ train_tasks.append(task)
+
+ # Train
+ trainer = Trainer(model, lr=5e-5)
+ # batch_size = None
+ batch_size = 5
+ n_epochs = 10
+ epoch_losses = []
+ for epoch in tqdm(range(n_epochs)):
+ batch_losses = trainer(train_tasks, batch_size=batch_size)
+ epoch_losses.append(np.mean(batch_losses))
+
+ # Check for NaNs in the loss
+ loss = np.mean(epoch_losses)
+ self.assertFalse(np.isnan(loss))