ZeroSumNormal: initial commit

Author: drbenvincent

examples/generalized_linear_models/GLM-ZeroSumNormal.ipynb

@@ -0,0 +1,74 @@
+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)

examples/generalized_linear_models/ZeroSumNormal.py

@@ -29,6 +29,7 @@ Gallery.contents = {
     "generalized_linear_models/GLM-rolling-regression": "(Generalized) Linear and Hierarchical Linear Models",
     "generalized_linear_models/GLM-truncated-censored-regression": "(Generalized) Linear and Hierarchical Linear Models",
     "generalized_linear_models/GLM-simpsons-paradox": "(Generalized) Linear and Hierarchical Linear Models",
+    "generalized_linear_models/GLM-ZeroSumNormal": "(Generalized) Linear and Hierarchical Linear Models",
     "gaussian_processes/GP-Kron": "Gaussian Processes",
     "gaussian_processes/GP-Latent": "Gaussian Processes",
     "gaussian_processes/GP-Marginal": "Gaussian Processes",