deepsensor.model.convnp

deepsensor.model.convnp#

class ConvNP(*args, **kwargs)[source]#

Bases: DeepSensorModel

A Convolutional Neural Process (ConvNP) regression probabilistic model (by default a ConvCNP).

Wraps around the neuralprocesses package to construct a ConvNP model. See: wesselb/neuralprocesses. Init kwargs passed to the ConvNP are passed to the neuralprocesses.construct_convgnp function and can be used to specify hyperparameters (see parameter list below). In particular, the likelihood parameter can be used to specify the likelihood of the model, which dictates whether the model outputs marginal distributions at each target point (a ConvCNP) or a joint Gaussian distribution over all target points (a ConvGNP). By default a ConvCNP with Gaussian likelihoods is constructed.

Additionally, the ConvNP can optionally be instantiated with:

a DataProcessor object to auto-unnormalise the data at inference time with the .predict method.
a TaskLoader object to infer sensible default model parameters from the data.

Multiple dispatch is implemented using plum to allow for re-using the model’s forward prediction object when computing the logpdf, entropy, etc. Alternatively, the model can be run forwards with a Task object of data from the TaskLoader.

Many of the ConvNP class methods utilise multiple dispatch so that they can either be run with a Task object or a neuralprocesses distribution object. This allows for re-using the model’s forward prediction object.

Dimension shapes are expressed in method docstrings in terms of:

N_features: number of features/dimensions in the target set.
N_targets: number of target points (1D for off-grid targets, 2D for gridded targets).
N_components: number of mixture components in the likelihood (for mixture likelihoods only).
N_samples: number of samples drawn from the distribution.

If the model has multiple target sets and the Task object has different target locations for each set, a list of arrays is returned for each target set. Otherwise, a single array is returned.

Examples

Instantiate a ConvNP with all hyperparameters set to their default values:

>>> ConvNP(data_processor, task_loader)

Instantiate a ConvNP and override some hyperparameters:

>>> ConvNP(data_processor, task_loader, internal_density=250, unet_channels=(128,) * 6)

Instantiate a ConvNP with a pre-trained model saved in the folder my_trained_model:

>>> ConvNP(data_processor, task_loader, model_ID="my_trained_model")

Instantiate a ConvNP with an existing neuralprocesses model object:

>>> ConvNP(data_processor, task_loader, neural_process=my_neural_process_model)

Parameters:

data_processor (DataProcessor, optional) – Used for unnormalising model predictions in .predict method.
task_loader (TaskLoader, optional) – Used for inferring sensible defaults for hyperparameters that are not set by the user.
model_ID (str, optional) – Folder to load the model config and weights from. This argument can only be used alongside the data_processor and task_loader arguments.
neural_process (TFModel | TorchModel, optional) – Pre-defined neural process PyTorch/TensorFlow model object. This argument can only be used alongside the data_processor and task_loader arguments.
internal_density (int, optional) – Density of the ConvNP’s internal grid (in terms of number of points per 1x1 unit square). Defaults to 100.
likelihood (str, optional) – Likelihood. Must be one of "cnp" (equivalently "het"), "gnp" (equivalently "lowrank"), "cnp-spikes-beta", (equivalently "spikes-beta") or “bernoulli-gamma”. Defaults to "cnp".
dim_x (int, optional) – Dimensionality of the inputs. Defaults to 1.
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 of the context set is not equal to the dimensionality of the outputs of the target set. You should also set this if you want to use multiple context sets. In that case, set this equal to a tuple of integers indicating the respective output dimensionalities.
dim_yt (int, optional) – Dimensionality of the outputs of the target set. You should set this if the dimensionality of the outputs of the target set is not equal to the dimensionality of the outputs of the context set.
dim_aux_t (int, optional) – Dimensionality of target-specific auxiliary variables.
conv_arch (str, optional) – Convolutional architecture to use. Must be one of "unet[-res][-sep]" or "conv[-res][-sep]". Defaults to "unet".
unet_channels (tuple[int], optional) – Number of channels in the downsampling path of the UNet (including the bottleneck). Defaults to four downsampling layers, each with 64 channels. I.e. (64, 64, 64, 64). Note: The downsampling path is followed by an upsampling path with the same number of channels in the reverse order (plus extra channels for the skip connections).
unet_kernels (int or tuple[int], optional) – Sizes of the kernels in the UNet. Defaults to 5.
unet_resize_convs (bool, optional) – Use resize convolutions rather than transposed convolutions in the UNet. Defaults to False.
unet_resize_conv_interp_method (str, optional) – Interpolation method for the resize convolutions in the UNet. Can be set to "bilinear". Defaults to “bilinear”.
num_basis_functions (int, optional) – Number of basis functions for the low-rank likelihood. Defaults to 64.
dim_lv (int, optional) – Dimensionality of the latent variable. Setting to >0 constructs a latent neural process. Defaults to 0.
encoder_scales (float or tuple[float], optional) – Initial value for the length scales of the set convolutions for the context sets embeddings. Set to a tuple equal to the number of context sets to use different values for each set. Set to a single value to use the same value for all context sets. Defaults to 1 / internal_density.
encoder_scales_learnable (bool, optional) – Whether the encoder SetConv length scale(s) are learnable. Defaults to False.
decoder_scale (float, optional) – Initial value for the length scale of the set convolution in the decoder. Defaults to 1 / internal_density.
decoder_scale_learnable (bool, optional) – Whether the decoder SetConv length scale(s) are learnable. Defaults to False.
aux_t_mlp_layers (tuple[int], optional) – Widths of the layers of the MLP for the target-specific auxiliary variable. Defaults to three layers of width 128.
epsilon (float, optional) – Epsilon added by the set convolutions before dividing by the density channel. Defaults to 1e-2.
dtype (dtype, optional) – Data type.

__call__(task, n_samples=10, requires_grad=False)[source]#

Compute ConvNP distribution.

Parameters:

task (Task) –
…
n_samples (int, optional) – Number of samples to draw from the distribution, by default 10.
requires_grad (bool, optional) – Whether to compute gradients, by default False.

Returns:

… – The ConvNP distribution.

alpha(dist)[source]#

alpha(self, task: deepsensor.data.task.Task) → Union[numpy.ndarray, List[numpy.ndarray]][source]

Alpha parameter values of model’s distribution at target locations in task.

Returned numpy arrays have shape (N_features, *N_targets).

Note

This method only works for models that return a distribution with a dist.slab.alpha attribute, e.g. models with a Beta or Bernoulli-Gamma likelihood, where it returns the alpha values of the slab component of the mixture model.

Parameters:: task (Task) – The task containing the context and target data.
Returns:: numpy.ndarray | List[numpy.ndarray] – Alpha values.

ar_sample(task, n_samples=1, X_target_AR=None, ar_subsample_factor=1, fill_type='mean')[source]#

Autoregressive sampling from the model.

AR sampling with optional functionality to only draw AR samples over a subset of the target set and then infill the rest of the sample with the model mean or joint sample conditioned on the AR samples.

Returned numpy arrays have shape (N_samples, N_features, *N_targets),

Note

AR sampling only works for 0th context/target set, and only for models with a single target set.

Parameters:

task (Task) – The task to sample from.
n_samples (int, optional) – The number of samples to draw from the distribution, by default 1.
X_target_AR (numpy.ndarray, optional) – Locations to draw AR samples over. If None, AR samples will be drawn over the target locations in the task. Defaults to None.
ar_subsample_factor (int, optional) – Subsample target locations to draw AR samples over. Defaults to 1.
fill_type (Literal["mean", "sample"], optional) – How to infill the rest of the sample. Must be one of “mean” or “sample”. Defaults to “mean”.

Returns:

numpy.ndarray: The samples.

beta(dist)[source]#

beta(self, task: deepsensor.data.task.Task) → Union[numpy.ndarray, List[numpy.ndarray]][source]

Beta values of model’s distribution at target locations in task.

Returned numpy arrays have shape (N_features, *N_targets).

Note

This method only works for models that return a distribution with a dist.slab.beta attribute, e.g. models with a Beta or Bernoulli-Gamma likelihood.

Parameters:: task (Task) – The task containing the context and target data.
Returns:: numpy.ndarray | List[numpy.ndarray] – Beta values.

covariance(dist)[source]#

covariance(self, task: deepsensor.data.task.Task)[source]

…

Parameters:

task (Task) –

…

Returns:

… – …

joint_entropy(dist)[source]#

Model entropy over target points given context points.

Parameters:: dist (neuralprocesses.dist.AbstractMultiOutputDistribution) – The distribution to compute the entropy of.
Returns:: float – The model entropy.

joint_entropy(self, task: deepsensor.data.task.Task)[source]

Model entropy over target points given context points.

Parameters:: task (Task) – The task to compute the entropy of.
Returns:: float – The model entropy.

k(dist)[source]#

k(self, task: deepsensor.data.task.Task) → Union[numpy.ndarray, List[numpy.ndarray]][source]

k parameter values of model’s distribution at target locations in task.

Returned numpy arrays have shape (N_features, *N_targets).

Note

This method only works for models that return a distribution with a dist.slab.k attribute, e.g. models with a Beta or Bernoulli-Gamma likelihood, where it returns the k values of the slab component of the mixture model.

Parameters:: task (Task) – The task containing the context and target data.
Returns:: numpy.ndarray | List[numpy.ndarray] – k values.

load(model_ID)[source]#

Load a model from a folder containing model weights and config.

Parameters:: model_ID (str) – Folder to load the model from.
Returns:: None.

logpdf(dist, task)[source]#

Joint logpdf over all target sets.

Note

If the model has multiple target sets, the returned logpdf is the mean logpdf over all target sets.

Parameters:

dist (neuralprocesses.dist.AbstractMultiOutputDistribution) – The distribution to compute the logpdf of.
task (Task) – The task to compute the logpdf of.

Returns:

float – The logpdf.

logpdf(self, task: deepsensor.data.task.Task)[source]

Joint logpdf over all target sets.

Note

If the model has multiple target sets, the returned logpdf is the mean logpdf over all target sets.

Parameters:: task (Task) – The task to compute the logpdf of.
Returns:: float – The logpdf.

loss_fn(task, fix_noise=None, num_lv_samples=8, normalise=False)[source]#

Compute the loss of a task.

Parameters:

task (Task) – The task to compute the loss of.
fix_noise – Whether to fix the noise to the value specified in the model config.
num_lv_samples (int, optional) – If latent variable model, number of lv samples for evaluating the loss, by default 8.
normalise (bool, optional) – Whether to normalise the loss by the number of target points, by default False.

Returns:

float – The loss.

mean(dist)[source]#

mean(self, task: deepsensor.data.task.Task)[source]

Mean values of model’s distribution at target locations in task.

Returned numpy arrays have shape (N_features, *N_targets).

Parameters:: task (Task) – The task containing the context and target data.
Returns:: numpy.ndarray | List[numpy.ndarray] – Mean values.

mean_marginal_entropy(dist)[source]#

Mean marginal entropy over target points given context points.

Parameters:: dist (neuralprocesses.dist.AbstractMultiOutputDistribution) – The distribution to compute the entropy of.
Returns:: float – The mean marginal entropy.

mean_marginal_entropy(self, task: deepsensor.data.task.Task)[source]

Mean marginal entropy over target points given context points.

Parameters:: task (Task) – The task to compute the entropy of.
Returns:: float – The mean marginal entropy.

mixture_probs(dist)[source]#

mixture_probs(self, task: deepsensor.data.task.Task)[source]

Mixture probabilities of model’s distribution at target locations in task.

Returned numpy arrays have shape (N_components, N_features, *N_targets).

Parameters:: task (Task) – The task containing the context and target data.
Returns:: numpy.ndarray | List[numpy.ndarray] – Mixture probabilities.

classmethod modify_task(task)[source]#

Cast numpy arrays to TensorFlow or PyTorch tensors, add batch dim, and mask NaNs.

Parameters:

task (Task) –

…

Returns:

… – …

sample(dist, n_samples=1)[source]#

sample(self, task: deepsensor.data.task.Task, n_samples: int = 1)[source]

Create samples from a ConvNP distribution.

Returned numpy arrays have shape (N_samples, N_features, *N_targets),

Parameters:

dist (neuralprocesses.dist.AbstractMultiOutputDistribution) – The distribution to sample from.
n_samples (int, optional) – The number of samples to draw from the distribution, by default 1.

Returns:

numpy.ndarray | List[numpy.ndarray] – The samples as an array or list of arrays.

save(model_ID)[source]#

Save the model weights and config to a folder.

Parameters:: model_ID (str) – Folder to save the model to.
Returns:: None.

scale(dist)[source]#

scale(self, task: deepsensor.data.task.Task) → Union[numpy.ndarray, List[numpy.ndarray]][source]

Scale parameter values of model’s distribution at target locations in task.

Returned numpy arrays have shape (N_features, *N_targets).

Note

This method only works for models that return a distribution with a dist.slab.scale attribute, e.g. models with a Beta or Bernoulli-Gamma likelihood, where it returns the scale values of the slab component of the mixture model.

Parameters:: task (Task) – The task containing the context and target data.
Returns:: numpy.ndarray | List[numpy.ndarray] – Scale values.

slice_diag(task)[source]#

Slice out the ConvCNP part of the ConvNP distribution.

Parameters:: task (Task) – The task to slice.
Returns:: … – …

slice_diag(self, dist: neuralprocesses.dist.dist.AbstractDistribution)[source]

Slice out the ConvCNP part of the ConvNP distribution.

Parameters:: dist (neuralprocesses.dist.AbstractMultiOutputDistribution) – The distribution to slice.
Returns:: … – …

std(dist)[source]#

std(self, task: deepsensor.data.task.Task)[source]

Standard deviation values of model’s distribution at target locations in task.

Returned numpy arrays have shape (N_features, *N_targets).

Parameters:: task (Task) – The task containing the context and target data.
Returns:: numpy.ndarray | List[numpy.ndarray] – Standard deviation values.

variance(dist)[source]#

variance(self, task: deepsensor.data.task.Task)[source]

Variance values of model’s distribution at target locations in task.

Returned numpy arrays have shape (N_features, *N_targets).

Parameters:: task (Task) – The task containing the context and target data.
Returns:: numpy.ndarray | List[numpy.ndarray] – Variance values.

concat_tasks(tasks, multiple=1)[source]#