ZeroSumNormal: initial commit

Author: drbenvincent

examples/generalized_linear_models/GLM-ZeroSumNormal.ipynb

@@ -1,3 +1,4 @@
+<<<<<<< HEAD
 from typing import List
 
 try:
@@ -138,3 +139,79 @@ def _distr_parameters_for_repr(self):
 
     def logcdf(self, value):
         raise NotImplementedError()
+=======
+import pymc3 as pm
+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)
+>>>>>>> 2da3052 (ZeroSumNormal: initial commit)