latest ZeroSumNormal code, pymc3 v3, random seed for sampling

Author: drbenvincent

examples/generalized_linear_models/GLM-ZeroSumNormal.ipynb

@@ -1,74 +1,140 @@
-import pymc3 as pm
+from typing import List
+
+try:
+    import aesara.tensor as aet
+except ImportError:
+    import theano.tensor as aet
+
 import numpy as np
-import pandas as pd
-from typing import *
-import aesara
-import aesara.tensor as aet
-
-
-def ZeroSumNormal(
-    name: str,
-    sigma: Optional[float] = None,
-    *,
-    dims: Union[str, Tuple[str]],
-    model: Optional[pm.Model] = None,
-):
-    """
-    Multivariate normal, such that sum(x, axis=-1) = 0.
-
-    Parameters
-    ----------
-    name: str
-        String name representation of the PyMC variable.
-    sigma: Optional[float], defaults to None
-        Scale for the Normal distribution. If ``None``, a standard Normal is used.
-    dims: Union[str, Tuple[str]]
-        Dimension names for the shape of the distribution.
-        See https://docs.pymc.io/pymc-examples/examples/pymc3_howto/data_container.html for an example.
-    model: Optional[pm.Model], defaults to None
-        PyMC model instance. If ``None``, a model instance is created.
-    """
-    if isinstance(dims, str):
-        dims = (dims,)
-
-    model = pm.modelcontext(model)
-    *dims_pre, dim = dims
-    dim_trunc = f"{dim}_truncated_"
-    (shape,) = model.shape_from_dims((dim,))
-    assert shape >= 1
-
-    model.add_coords({f"{dim}_truncated_": pd.RangeIndex(shape - 1)})
-    raw = pm.Normal(f"{name}_truncated_", dims=tuple(dims_pre) + (dim_trunc,), sigma=sigma)
-    Q = make_sum_zero_hh(shape)
-    draws = aet.dot(raw, Q[:, 1:].T)
-
-    #if sigma is not None:
-    #    draws = sigma * draws
-
-    return pm.Deterministic(name, draws, dims=dims)
-
-
-
-def make_sum_zero_hh(N: int) -> np.ndarray:
-    """
-    Build a householder transformation matrix that maps e_1 to a vector of all 1s.
-    """
-    e_1 = np.zeros(N)
-    e_1[0] = 1
-    a = np.ones(N)
-    a /= np.sqrt(a @ a)
-    v = a + e_1
-    v /= np.sqrt(v @ v)
-    return np.eye(N) - 2 * np.outer(v, v)
-
-def make_sum_zero_hh(N: int) -> np.ndarray:
-    """
-    Build a householder transformation matrix that maps e_1 to a vector of all 1s.
-    """
-    e_1 = np.zeros(N)
-    e_1[0] = 1
-    a = np.ones(N)
-    a /= np.sqrt(a @ a)
-    v = a + e_1
-    v /= np.sqrt(v @ v)
-    return np.eye(N) - 2 * np.outer(v, v)
+import pymc3 as pm
+from scipy import stats
+from pymc3.distributions.distribution import generate_samples, draw_values
+
+def extend_axis_aet(array, axis):
+    n = array.shape[axis] + 1
+    sum_vals = array.sum(axis, keepdims=True)
+    norm = sum_vals / (np.sqrt(n) + n)
+    fill_val = norm - sum_vals / np.sqrt(n)
+    
+    out = aet.concatenate([array, fill_val.astype(str(array.dtype))], axis=axis)
+    return out - norm.astype(str(array.dtype))
+
+
+def extend_axis_rev_aet(array: np.ndarray, axis: int):
+    if axis < 0:
+        axis = axis % array.ndim
+    assert axis >= 0 and axis < array.ndim
+
+    n = array.shape[axis]
+    last = aet.take(array, [-1], axis=axis)
+    
+    sum_vals = -last * np.sqrt(n)
+    norm = sum_vals / (np.sqrt(n) + n)
+    slice_before = (slice(None, None),) * axis
+    return array[slice_before + (slice(None, -1),)] + norm.astype(str(array.dtype))
+
+
+def extend_axis(array, axis):
+    n = array.shape[axis] + 1
+    sum_vals = array.sum(axis, keepdims=True)
+    norm = sum_vals / (np.sqrt(n) + n)
+    fill_val = norm - sum_vals / np.sqrt(n)
+    
+    out = np.concatenate([array, fill_val.astype(str(array.dtype))], axis=axis)
+    return out - norm.astype(str(array.dtype))
+
+
+def extend_axis_rev(array, axis):
+    n = array.shape[axis]
+    last = np.take(array, [-1], axis=axis)
+    
+    sum_vals = -last * np.sqrt(n)
+    norm = sum_vals / (np.sqrt(n) + n)
+    slice_before = (slice(None, None),) * len(array.shape[:axis])
+    return array[slice_before + (slice(None, -1),)] + norm.astype(str(array.dtype))
+
+
+class ZeroSumTransform(pm.distributions.transforms.Transform):
+    name = "zerosum"
+    
+    _active_dims: List[int]
+    
+    def __init__(self, active_dims):
+        self._active_dims = active_dims
+    
+    def forward(self, x):
+        for axis in self._active_dims:
+            x = extend_axis_rev_aet(x, axis=axis)
+        return x
+    
+    def forward_val(self, x, point=None):
+        for axis in self._active_dims:
+            x = extend_axis_rev(x, axis=axis)
+        return x
+    
+    def backward(self, z):
+        z = aet.as_tensor_variable(z)
+        for axis in self._active_dims:
+            z = extend_axis_aet(z, axis=axis)
+        return z
+    
+    def jacobian_det(self, x):
+        return aet.constant(0.)
+    
+    
+class ZeroSumNormal(pm.Continuous):
+    def __init__(self, sigma=1, *, active_dims=None, active_axes=None, **kwargs):
+        shape = kwargs.get("shape", ())
+        dims = kwargs.get("dims", None)
+        if isinstance(shape, int):
+            shape = (shape,)
+        
+        if isinstance(dims, str):
+            dims = (dims,)
+
+        self.mu = self.median = self.mode = aet.zeros(shape)
+        self.sigma = aet.as_tensor_variable(sigma)
+        
+        if active_dims is None and active_axes is None:
+            if shape:
+                active_axes = (-1,)
+            else:
+                active_axes = ()
+        
+        if isinstance(active_axes, int):
+            active_axes = (active_axes,)
+        
+        if isinstance(active_dims, str):
+            active_dims = (active_dims,)
+        
+        if active_axes is not None and active_dims is not None:
+            raise ValueError("Only one of active_axes and active_dims can be specified.")
+        
+        if active_dims is not None:
+            model = pm.modelcontext(None)
+            print(model.RV_dims)
+            if dims is None:
+                raise ValueError("active_dims can only be used with the dims kwargs.")
+            active_axes = []
+            for dim in active_dims:
+                active_axes.append(dims.index(dim))
+        
+        super().__init__(**kwargs, transform=ZeroSumTransform(active_axes))
+
+    def logp(self, x):
+        return pm.Normal.dist(sigma=self.sigma).logp(x)
+    
+    @staticmethod
+    def _random(scale, size):
+        samples = stats.norm.rvs(loc=0, scale=scale, size=size)
+        return samples - np.mean(samples, axis=-1, keepdims=True)
+    
+    def random(self, point=None, size=None):
+        sigma, = draw_values([self.sigma], point=point, size=size)
+        return generate_samples(self._random, scale=sigma, dist_shape=self.shape, size=size)
+
+    def _distr_parameters_for_repr(self):
+        return ["sigma"]
+
+    def logcdf(self, value):
+        raise NotImplementedError()