Merge pull request #14 from wesselb/struclik

Structured likelihoods

Author: wesselb

Diff for: neuralprocesses/aggregate.py

@@ -79,14 +79,19 @@ def f(*args: Aggregate, **kw_args):
         raise ValueError(f"Invalid number of arguments {num_args}.")
 
 
+_map_f("expand_dims", 1)
+_map_f("exp", 1)
+_map_f("one", 1)
+_map_f("zero", 1)
 _map_f("mean", 1)
+_map_f("sum", 1)
+_map_f("logsumexp", 1)
 
 _map_f("add", 2)
 _map_f("subtract", 2)
 _map_f("multiply", 2)
 _map_f("divide", 2)
 
-
 _map_f("stack", "*")
 _map_f("concat", "*")
 _map_f("squeeze", "*")

Diff for: neuralprocesses/architectures/agnp.py

@@ -1,4 +1,5 @@
 import neuralprocesses as nps  # This fixes inspection below.
+
 from ..util import register_model
 
 __all__ = ["construct_agnp"]
@@ -34,7 +35,7 @@ def construct_agnp(*args, nps=nps, num_heads=8, **kw_args):
             low-rank likelihood. Defaults to 512.
         dim_lv (int, optional): Dimensionality of the latent variable. Defaults to 0.
         lv_likelihood (str, optional): Likelihood of the latent variable. Must be one of
-            `"het"` or `"dense"`. Defaults to `"het"`.
+            `"het"`, `"dense"`, or `"spikes-beta"`. Defaults to `"het"`.
         transform (str or tuple[float, float]): Bijection applied to the
             output of the model. This can help deal with positive of bounded data.
             Must be either `"positive"`, `"exp"`, `"softplus"`, or

Diff for: neuralprocesses/architectures/convgnp.py

@@ -1,10 +1,10 @@
 import lab as B
+import neuralprocesses as nps  # This fixes inspection below.
 import wbml.out as out
 from plum import convert
 
-import neuralprocesses as nps  # This fixes inspection below.
-from .util import construct_likelihood, parse_transform
 from ..util import register_model
+from .util import construct_likelihood, parse_transform
 
 __all__ = ["construct_convgnp"]
 
@@ -140,8 +140,8 @@ def construct_convgnp(
             Defaults to 64.
         margin (float, optional): Margin of the internal discretisation. Defaults to
             0.1.
-        likelihood (str, optional): Likelihood. Must be one of `"het"` or `"lowrank".
-            Defaults to `"lowrank"`.
+        likelihood (str, optional): Likelihood. Must be one of `"het"`, `"lowrank"`,
+            or `"spikes-beta"`. Defaults to `"lowrank"`.
         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): Channels of every layer of the UNet.

Diff for: neuralprocesses/architectures/gnp.py

@@ -1,8 +1,8 @@
+import neuralprocesses as nps  # This fixes inspection below.
 from plum import convert
 
-import neuralprocesses as nps  # This fixes inspection below.
-from .util import construct_likelihood, parse_transform
 from ..util import register_model
+from .util import construct_likelihood, parse_transform
 
 __all__ = ["construct_gnp"]
 
@@ -58,7 +58,7 @@ def construct_gnp(
             low-rank likelihood. Defaults to 512.
         dim_lv (int, optional): Dimensionality of the latent variable. Defaults to 0.
         lv_likelihood (str, optional): Likelihood of the latent variable. Must be one of
-            `"het"` or `"dense"`. Defaults to `"het"`.
+            `"het"`, `"dense"`, or `"spikes-beta"`. Defaults to `"het"`.
         transform (str or tuple[float, float]): Bijection applied to the
             output of the model. This can help deal with positive of bounded data.
             Must be either `"positive"`, `"exp"`, `"softplus"`, or

Diff for: neuralprocesses/architectures/util.py

@@ -11,8 +11,9 @@ def construct_likelihood(nps=nps, *, spec, dim_y, num_basis_functions, dtype):
 
     Args:
         nps (module): Appropriate backend-specific module.
-        spec (str, optional): Specification. Must be one of `"het"`, `"lowrank"`, or
-            `"dense"`. Defaults to `"lowrank"`. Must be given as a keyword argument.
+        spec (str, optional): Specification. Must be one of `"het"`, `"lowrank"`,
+            `"dense"`, or `"spikes-beta"`. Defaults to `"lowrank"`. Must be given as
+            a keyword argument.
         dim_y (int): Dimensionality of the outputs. Must be given as a keyword argument.
         num_basis_functions (int): Number of basis functions for the low-rank
             likelihood. Must be given as a keyword argument.
@@ -47,6 +48,10 @@ def construct_likelihood(nps=nps, *, spec, dim_y, num_basis_functions, dtype):
             ),
             nps.DenseGaussianLikelihood(),
         )
+    elif spec == "spikes-beta":
+        num_channels = (2 + 3) * dim_y  # Alpha, beta, and three log-probabilities
+        selector = nps.SelectFromChannels(dim_y, dim_y, dim_y, dim_y, dim_y)
+        lik = nps.SpikesBetaLikelihood()
 
     else:
         raise ValueError(f'Incorrect likelihood specification "{spec}".')

Diff for: neuralprocesses/coding.py

@@ -2,7 +2,7 @@
 from plum import isinstance, issubclass
 
 from . import _dispatch
-from .dist import AbstractMultiOutputDistribution, Dirac
+from .dist import AbstractDistribution, Dirac
 from .parallel import Parallel
 from .util import is_composite_coder
 
@@ -70,7 +70,7 @@ def code_track(coder, xz, z, x, h, **kw_args):
     if is_composite_coder(coder):
         raise RuntimeError(
             f"Dispatched to fallback implementation of `code_track` for "
-            f"`{ptype(type(coder))}`, but the coder is composite."
+            f"`{type(coder)}`, but the coder is composite."
         )
     xz, z = code(coder, xz, z, x, **kw_args)
     return xz, z, h + [x]
@@ -95,7 +95,7 @@ def recode(coder, xz, z, h, **kw_args):
     if is_composite_coder(coder):
         raise RuntimeError(
             f"Dispatched to fallback implementation of `recode` for "
-            f"`{ptype(type(coder))}`, but the coder is composite."
+            f"`{type(coder)}`, but the coder is composite."
         )
     xz, z = code(coder, xz, z, h[0], **kw_args)
     return xz, z, h[1:]
@@ -151,6 +151,6 @@ def _choose(new: Dirac, old: Dirac):
 
 
 @_dispatch
-def _choose(new: AbstractMultiOutputDistribution, old: AbstractMultiOutputDistribution):
+def _choose(new: AbstractDistribution, old: AbstractDistribution):
     # Do recode other distributions.
     return new

Diff for: neuralprocesses/data/bimodal.py

@@ -0,0 +1,65 @@
+import lab as B
+import numpy as np
+from lab.shape import Dimension
+
+from neuralprocesses.data import SyntheticGenerator, new_batch
+
+__all__ = ["BiModalGenerator"]
+
+
+class BiModalGenerator(SyntheticGenerator):
+    """Bi-modal distribution generator.
+
+    Further takes in arguments and keyword arguments from the constructor of
+    :class:`.data.SyntheticGenerator`. Moreover, also has the attributes of
+    :class:`.data.SyntheticGenerator`.
+    """
+
+    def __init__(self, *args, **kw_args):
+        super().__init__(*args, **kw_args)
+
+    def generate_batch(self):
+        with B.on_device(self.device):
+            set_batch, xcs, xc, nc, xts, xt, nt = new_batch(self, self.dim_y)
+            x = B.concat(xc, xt, axis=1)
+
+            # Draw a different random phase, amplitude, and period for every task in
+            # the batch.
+            self.state, rand = B.rand(
+                self.state,
+                self.float64,
+                3,
+                self.batch_size,
+                1,  # Broadcast over `n`.
+                1,  # There is only one input dimension.
+            )
+            phase = 2 * B.pi * rand[0]
+            amplitude = 1 + rand[1]
+            period = 1 + rand[2]
+
+            # Construct the noiseless function.
+            f = amplitude * B.sin(phase + (2 * B.pi / period) * x)
+
+            # Add noise with variance.
+            probs = B.cast(self.float64, np.array([0.5, 0.5]))
+            means = B.cast(self.float64, np.array([-0.1, 0.1]))
+            variance = 1
+            # Randomly choose from `means` with probabilities `probs`.
+            self.state, mean = B.choice(self.state, means, self.batch_size, p=probs)
+            self.state, randn = B.randn(
+                self.state,
+                self.float64,
+                self.batch_size,
+                # `nc` and `nt` are tensors rather than plain integers. Tell dispatch
+                # that they can be interpreted as dimensions of a shape.
+                Dimension(nc + nt),
+                1,
+            )
+            noise = B.sqrt(variance) * randn + mean[:, None, None]
+
+            # Construct the noisy function.
+            y = f + noise
+
+            batch = {}
+            set_batch(batch, y[:, :nc], y[:, nc:])
+            return batch

Diff for: neuralprocesses/dist/__init__.py

@@ -1,6 +1,8 @@
+from .beta import *
 from .dirac import *
 from .dist import *
 from .geom import *
 from .normal import *
+from .spikeslab import *
 from .transformed import *
 from .uniform import *

Diff for: neuralprocesses/dist/beta.py

@@ -0,0 +1,119 @@
+import lab as B
+from matrix.shape import broadcast
+from plum import parametric
+
+from .. import _dispatch
+from ..aggregate import Aggregate
+from ..mask import Masked
+from .dist import AbstractDistribution, shape_batch
+
+__all__ = ["Beta"]
+
+
+@parametric
+class Beta(AbstractDistribution):
+    """Beta distribution.
+
+    Args:
+        alpha (tensor): Shape parameter `alpha`.
+        beta (tensor): Shape parameter `beta`.
+        d (int): Dimensionality of the data.
+
+    Attributes:
+        alpha (tensor): Shape parameter `alpha`.
+        beta (tensor): Shape parameter `beta`.
+        d (int): Dimensionality of the data.
+    """
+
+    def __init__(self, alpha, beta, d):
+        self.alpha = alpha
+        self.beta = beta
+        self.d = d
+
+    @property
+    def mean(self):
+        return B.divide(self.alpha, B.add(self.alpha, self.beta))
+
+    @property
+    def var(self):
+        sum = B.add(self.alpha, self.beta)
+        with B.on_device(sum):
+            one = B.one(sum)
+        return B.divide(
+            B.multiply(self.alpha, self.beta),
+            B.multiply(B.multiply(sum, sum), B.add(sum, one)),
+        )
+
+    @_dispatch
+    def sample(
+        self: "Beta[Aggregate, Aggregate, Aggregate]",
+        state: B.RandomState,
+        dtype: B.DType,
+        *shape,
+    ):
+        samples = []
+        for ai, bi, di in zip(self.alpha, self.beta, self.d):
+            state, sample = Beta(ai, bi, di).sample(state, dtype, *shape)
+            samples.append(sample)
+        return state, Aggregate(*samples)
+
+    @_dispatch
+    def sample(
+        self: "Beta[B.Numeric, B.Numeric, B.Int]",
+        state: B.RandomState,
+        dtype: B.DType,
+        *shape,
+    ):
+        return B.randbeta(state, dtype, *shape, alpha=self.alpha, beta=self.beta)
+
+    @_dispatch
+    def logpdf(self: "Beta[Aggregate, Aggregate, Aggregate]", x: Aggregate):
+        return sum(
+            [
+                Beta(ai, bi, di).logpdf(xi)
+                for ai, bi, di, xi in zip(self.alpha, self.beta, self.d, x)
+            ],
+            0,
+        )
+
+    @_dispatch
+    def logpdf(self: "Beta[B.Numeric, B.Numeric, B.Int]", x: Masked):
+        x, mask = x.y, x.mask
+        with B.on_device(self.alpha):
+            safe = B.to_active_device(B.cast(B.dtype(self.alpha), 0.5))
+        # Make inputs safe.
+        x = mask * x + (1 - mask) * safe
+        # Run with safe inputs, and filter out the right logpdfs.
+        return self.logpdf(x, mask=mask)
+
+    @_dispatch
+    def logpdf(self: "Beta[B.Numeric, B.Numeric, B.Int]", x: B.Numeric, *, mask=1):
+        logz = B.logbeta(self.alpha, self.beta)
+        logpdf = (self.alpha - 1) * B.log(x) + (self.beta - 1) * B.log(1 - x) - logz
+        return B.sum(mask * logpdf, axis=tuple(range(B.rank(logpdf)))[-self.d :])
+
+    def __str__(self):
+        return f"Beta({self.alpha}, {self.beta})"
+
+    def __repr__(self):
+        return f"Beta({self.alpha!r}, {self.beta!r})"
+
+
+@B.dtype.dispatch
+def dtype(dist: Beta):
+    return B.dtype(dist.alpha, dist.beta)
+
+
+@shape_batch.dispatch
+def shape_batch(dist: "Beta[B.Numeric, B.Numeric, B.Int]"):
+    return B.shape_broadcast(dist.alpha, dist.beta)[: -dist.d]
+
+
+@shape_batch.dispatch
+def shape_batch(dist: "Beta[Aggregate, Aggregate, Aggregate]"):
+    return broadcast(
+        *(
+            shape_batch(Beta(ai, bi, di))
+            for ai, bi, di in zip(dist.alpha, dist.beta, dist.d)
+        )
+    )