Improving format of object returned by `DeepSensorModel.predict`

[tom-andersson]

I'd like to get feedback on potential breaking changes in the objects returned by model.predict in DeepSensor. As a reminder, model.predict allows you to take a bunch of Tasks (containing context data) and get unnormalised model predictions directly in xarray (for gridded predictions) or pandas (for non-gridded predictions). This is one of the key selling points of DeepSensor because it is very convenient for downstream analysis of model predictions.

There are two key challenges that make it tricky to decide what to return from .predict:

• We aim to support non-Gaussian distributions in future, so the model's prediction parameters will depend on the model's likelihood/distribution
• There can be multiple target variables, leading to N-D predictions. Not all models will support this, but the default (and only) model in DeepSensor right now, the ConvNP, does.

Therefore, we can have a varying number of distribution parameters, and also a varying number of variables.

Currently, .predict by default returns two objects: a mean and a stddev of the model's Gaussian distributions. If the user passes the n_samples kwarg, there is also a third samples object returned. One can obtain the predictions for a specific variable like mean["var1"] and mean["var2"] (assuming a 2D target), for example.

When we start to support non-Gaussian distributions in future, it would be silly to then start returning these as additional return objects. For example, if we used a Bernoulli-Gamma distribution for a positive semi-definite variable like precip, it would be silly to return extra bernoulli_prob and gamma_alpha and gamma_beta objects. This would risk ValueError: too many values to unpack in user code.

Here is my proposal: model.predict instead returns a single dictionary where keys are variable IDs (e.g. "t2m") and values are single xarray/pandas objects containing distribution parameters. For example:

pred = model.predict(tasks, ..., n_samples=5)
print(pred["t2m"]["mean"])  # the model mean for the t2m var (xarray.DataArray or pandas.Series)
print(pred["t2m"]["sample1"])  # the first sample for the t2m var (xarray.DataArray or pandas.Series)
print(pred["sst"]["std"])  # the model std dev for the sst var (xarray.DataArray or pandas.Series)

This allows the number of target variables and number of distribution parameters to change, while wrapping everything in a single return object.

You can see in the above snippet that it gets a little awkward when drawing model samples with n_samples>0 because samples have a 'sample number' dimension as well as spatial and temporal dimensions. Forcing the samples into the same xarray/pandas objects means we need to have sample<n>-like variables. Perhaps we can get away with putting the samples in a separate object, eg pred, samples = model.predict(tasks, ..., n_samples=5), and assume the user will be happy to handle the second returned object in this case.

[acocac]

Not an immediate solution, but there's a WIP hierarchical data structure building upon xarray objects, see datatree.DataTree. The same development might be interesting for data loaders.

[tom-andersson]

@patel-zeel's idea: give the pred object an informative __str__ method so that it can be printed like:

>>> print(pred)
Prediction object with variables:
't2m' with prediction attributes ['mean', 'std']
'precip' with prediction attributes ['mean', 'std', 'bernoulli_prob',  ...]

[tom-andersson]

I'm starting to think it's much cleaner to put the model samples in the same xarray/pandas object as the other prediction parameters (eg mean/std). We can always have a utility that separates out the sample data vars/columns and adds a sample dimension if that is useful.

I also don't like model.predict returning either 1 or 2 objects - IMO it's nicer if it's only ever one object.

[tom-andersson]

Closed by cbdc731. Now we have:

>>> test_task = task_loader("2014-12-31", 0.1)
>>> pred = model.predict(test_task, X_t=ds_raw)
>>> pred["air"]
<xarray.Dataset>
Dimensions:  (time: 1, lat: 25, lon: 53)
Coordinates:
  * time     (time) datetime64[ns] 2014-12-31
  * lat      (lat) float32 75.0 72.5 70.0 67.5 65.0 ... 25.0 22.5 20.0 17.5 15.0
  * lon      (lon) float32 200.0 202.5 205.0 207.5 ... 322.5 325.0 327.5 330.0
Data variables:
    mean     (time, lat, lon) float32 267.7 267.2 266.4 ... 297.5 297.8 297.9
    std      (time, lat, lon) float32 9.855 9.845 9.848 ... 1.356 1.36 1.487