add VonMises family

bambi/backend/links.py

@@ -37,3 +37,7 @@ def identity(x):
 
 def inverse_squared(x):
     return tt.inv(tt.sqrt(x))
+
+
+def arctan_2(x):
+    return 2 * tt.arctan(x)

bambi/backend/pymc.py

@@ -8,7 +8,7 @@
 
 from bambi import version
 
-from bambi.backend.links import cloglog, identity, inverse_squared, logit, probit
+from bambi.backend.links import cloglog, identity, inverse_squared, logit, probit, arctan_2
 from bambi.backend.terms import CommonTerm, GroupSpecificTerm, InterceptTerm, ResponseTerm
 
 _log = logging.getLogger("bambi")
@@ -25,6 +25,7 @@ class PyMC3Model:
         "log": tt.exp,
         "logit": logit,
         "probit": probit,
+        "tan_2": arctan_2,
         "softmax": tt.nnet.softmax,
     }
 

bambi/defaults/defaults.py

@@ -10,6 +10,7 @@
     NegativeBinomial,
     Poisson,
     StudentT,
+    VonMises,
     Wald,
 )
 
@@ -31,6 +32,7 @@
     "NegativeBinomial": {"alpha": 1, "mu": 1},
     "Poisson": {"mu": 1},
     "StudentT": {"lam": 1, "nu": 1},
+    "VonMises": {"mu": 0, "kappa": 1},
     "Wald": {"mu": 1, "lam": 1},
 }
 
@@ -128,6 +130,17 @@
         "link": "identity",
         "family": StudentT,
     },
+    "vonmises": {
+        "likelihood": {
+            "name": "VonMises",
+            "args": {
+                "kappa": "HalfNormal"
+            },
+            "parent": "mu",
+        },
+        "link": "tan_2",
+        "family": VonMises,
+    },
     "wald": {
         "likelihood": {
             "name": "Wald",

bambi/families/family.py

@@ -63,6 +63,7 @@ class Family:
         "logit",
         "probit",
         "softmax",
+        "tan_2",
     ]
 
     def __init__(self, name, likelihood, link):

bambi/families/likelihood.py

@@ -11,6 +11,7 @@
     "NegativeBinomial": {"params": ("mu", "alpha"), "parent": "mu", "args": ("alpha",)},
     "Poisson": {"params": ("mu",), "parent": "mu", "args": None},
     "StudentT": {"params": ("mu", "sigma"), "args": ("sigma", "nu")},
+    "VonMises": {"params": ("mu", "kappa"), "parent": "mu", "args": ("kappa",)},
     "Wald": {"params": ("mu", "lam"), "parent": "mu", "args": ("lam",)},
 }
 

bambi/families/link.py

@@ -75,6 +75,14 @@ def inv_inverse_squared(eta):
     return 1 / np.sqrt(eta)
 
 
+def arctan_2(x):
+    return 2 * np.arctan(x)
+
+
+def tan_2(x):
+    return np.tan(x / 2)
+
+
 def link_not_implemented(*args, **kwargs):
     raise NotImplementedError("link not implemented")
 
@@ -114,7 +122,11 @@ def link_not_implemented(*args, **kwargs):
     "softmax": {
         "link": link_not_implemented,
         "linkinv": softmax
-    }
+    },
+    "tan_2": {
+        "link": tan_2,
+        "linkinv": arctan_2
+    },
 }
 # fmt: on
 

bambi/families/univariate.py

@@ -136,6 +136,20 @@ def posterior_predictive(self, model, posterior, linear_predictor, draws, draw_n
         return stats.t.rvs(nu, mean, sigma)
 
 
+class VonMises(UnivariateFamily):
+    SUPPORTED_LINKS = ["identity", "tan_2"]
+
+    def posterior_predictive(self, model, posterior, linear_predictor, draws, draw_n):
+        mean = self.link.linkinv(linear_predictor)
+        kappa = posterior[model.response.name + "_kappa"].values
+
+        idxs = np.random.randint(low=0, high=draw_n, size=draws)
+        mean = mean[:, idxs, :]
+        kappa = kappa[:, idxs, np.newaxis]
+
+        return np.random.vonmises(mean, kappa)
+
+
 class Wald(UnivariateFamily):
     SUPPORTED_LINKS = ["inverse", "inverse_squared", "identity", "log"]
 

bambi/families/utils.py

@@ -7,6 +7,7 @@
     "gaussian": ("sigma",),
     "negativebinomial": ("alpha",),
     "t": ("sigma", "nu"),
+    "vonmises": ("kappa",),
     "wald": ("lam",),
 }
 

bambi/tests/test_built_models.py

@@ -464,6 +464,11 @@ def test_t_regression():
     Model("y ~ x", data, family="t").fit(draws=10, tune=10)
 
 
+def test_vonmises_regression():
+    data = pd.DataFrame({"y": np.random.vonmises(0, 1, size=100), "x": np.random.normal(size=100)})
+    Model("y ~ x", data, family="vonmises").fit(draws=10, tune=10)
+
+
 def test_plot_priors(crossed_data):
     model = Model("Y ~ 0 + threecats", crossed_data)
     # Priors cannot be plotted until model is built.

bambi/tests/test_model_construction.py

@@ -328,6 +328,7 @@ def test_sparse_fails():
         "bernoulli",
         "poisson",
         "gamma",
+        "vonmises",
         "wald",
     ],
 )
@@ -353,6 +354,7 @@ def test_links():
         "gaussian": ["identity", "log", "inverse"],
         "negativebinomial": ["identity", "log", "cloglog"],
         "poisson": ["identity", "log"],
+        "vonmises": ["identity", "tan_2"],
         "wald": ["inverse", "inverse_squared", "identity", "log"],
     }
     for family, links in FAMILIES.items():
@@ -379,6 +381,7 @@ def test_bad_links():
         "gaussian": ["logit", "probit", "cloglog"],
         "negativebinomial": ["logit", "probit", "inverse", "inverse_squared"],
         "poisson": ["logit", "probit", "cloglog", "inverse", "inverse_squared"],
+        "vonmises": ["logit", "probit", "cloglog"],
         "wald": ["logit", "probit", "cloglog"],
     }