Add base class for distributions with non-conjugate priors and an example use

cxx/distributions/BUILD

@@ -63,6 +63,15 @@ cc_library(
     ],
 )
 
+cc_library(
+    name = "nonconjugate",
+    hdrs = ["nonconjugate.hh"],
+    visibility = ["//:__subpackages__"],
+    deps = [
+        ":base",
+    ],
+)
+
 cc_library(
     name = "normal",
     srcs = ["normal.cc"],
@@ -73,6 +82,16 @@ cc_library(
         "//:util_math",
     ],
 )
+
+cc_library(
+    name = "skellam",
+    hdrs = ["skellam.hh"],
+    deps = [
+        ":nonconjugate",
+        "//:util_math",
+    ],
+)
+
 cc_library(
     name = "zero_mean_normal",
     srcs = ["zero_mean_normal.cc"],
@@ -137,6 +156,17 @@ cc_test(
     ],
 )
 
+cc_test(
+    name = "skellam_test",
+    srcs = ["skellam_test.cc"],
+    deps = [
+        ":skellam",
+        "@boost//:algorithm",
+        "@boost//:math",
+        "@boost//:test",
+    ],
+)
+
 cc_test(
     name = "zero_mean_normal_test",
     srcs = ["zero_mean_normal_test.cc"],

cxx/distributions/base.hh

@@ -39,5 +39,13 @@ class Distribution {
   // e^logp_score() under those hyperparameters.
   virtual void transition_hyperparameters(std::mt19937* prng) = 0;
 
+  // Set the current latent values to a sample from the parameter prior.
+  // Only children of NonconjugateDistribution need define this.
+  virtual void init_theta(std::mt19937* prng) {};
+
+  // Transition the current latent values.  Only children of
+  // NonconjugateDistribution need define this.
+  virtual void transition_theta(std::mt19937* prng) {};
+
   virtual ~Distribution() = default;
 };

cxx/distributions/nonconjugate.hh

@@ -0,0 +1,42 @@
+#pragma once
+
+#include <map>
+#include <random>
+#include "distributions/base.hh"
+
+template <typename T>
+class NonconjugateDistribution : public Distribution<T> {
+ public:
+  // Abstract base class for Distributions that don't have conjugate priors.
+
+  // The log probability of x given the current latent values.
+  virtual double logp(const T& x) const = 0;
+
+  // Sample a value from the distribution given the current latent values.
+  virtual T sample(std::mt19937* prng) = 0;
+
+  // Transition hyperparameters given the current latent values.
+  virtual void transition_hyperparameters(std::mt19937* prng) = 0;
+
+  // Set the current latent values to a sample from the parameter prior.
+  virtual void init_theta(std::mt19937* prng) = 0;
+
+  // Transition the current latent values.
+  virtual void transition_theta(std::mt19937* prng) = 0;
+
+  double cumulative_logp = 0.0;
+
+  void incorporate(const T& x) {
+    (this->N)++;
+    cumulative_logp += logp(x);
+  };
+
+  void unincorporate(const T& x) {
+    --(this->N);
+    cumulative_logp -= logp(x);
+  };
+
+  double logp_score() const {
+    return cumulative_logp;
+  }
+};

cxx/distributions/skellam.hh

@@ -0,0 +1,73 @@
+#pragma once
+
+#include <cmath>
+
+#include "distributions/nonconjugate.hh"
+#include "util_math.hh"
+
+#define MEAN_GRID { -10.0, 0.0, 10.0 }
+#define STDDEV_GRID { 0.1, 1.0, 10.0 }
+
+double lognormal_logp(double x, double mean, double stddev) {
+  double y = (std::log(x) - mean) / stddev;
+  return - y*y / 2.0
+      - std::log(x * stddev) - 0.5 * std::log(2.0 * std::numbers::pi);
+}
+
+class Skellam : public NonconjugateDistribution<int> {
+ public:
+  // Skellam distribution with log Normal hyperprior of latent rates.
+  double mean1, mean2, stddev1, stddev2;  // Hyperparameters
+  // Skellam distribution with log-normal priors on mu1 and mu2.
+  double mu1, mu2;   // Latent values.
+
+  Skellam(): mean1(0.0), mean2(0.0), stddev1(1.0), stddev2(1.0),
+             mu1(1.0), mu2(1.0) {}
+
+  double logp(const int& x) const {
+    return -mu1 - mu2 + (x / 2.0) * std::log(mu1 / mu2)
+        // TODO(thomaswc): Replace this with something more numerically stable.
+        + std::log(std::cyl_bessel_i(x, 2.0 * std::sqrt(mu1 * mu2)));
+  }
+
+  int sample(std::mt19937* prng) {
+    std::poisson_distribution<int> d1(mu1);
+    std::poisson_distribution<int> d2(mu2);
+    return d1(*prng) - d2(*prng);
+  }
+
+  void transition_hyperparameters(std::mt19937* prng) {
+    std::vector<double> logps;
+    std::vector<std::tuple<double, double, double, double>> hypers;
+    for (double tmean1 : MEAN_GRID) {
+      for (double tstddev1 : STDDEV_GRID) {
+        for (double tmean2 : MEAN_GRID) {
+          for (double tstddev2 : STDDEV_GRID) {
+            double lp = lognormal_logp(mu1, tmean1, tstddev1)
+                        + lognormal_logp(mu2, tmean2, tstddev2);
+            logps.push_back(lp);
+            hypers.push_back(
+                std::make_tuple(tmean1, tstddev1, tmean2, tstddev2));
+          }
+        }
+      }
+    }
+    int i = sample_from_logps(logps, prng);
+    mean1 = std::get<0>(hypers[i]);
+    stddev1 = std::get<1>(hypers[i]);
+    mean2 = std::get<2>(hypers[i]);
+    stddev2 = std::get<3>(hypers[i]);
+  }
+
+  void init_theta(std::mt19937* prng) {
+    std::normal_distribution<double> d1(mean1, stddev1);
+    std::normal_distribution<double> d2(mean2, stddev2);
+    mu1 = std::exp(d1(*prng));
+    mu2 = std::exp(d2(*prng));
+  }
+
+  void transition_theta(std::mt19937* prng) {
+    // TODO(thomaswc): This
+  }
+
+};

cxx/distributions/skellam_test.cc

@@ -0,0 +1,33 @@
+// Apache License, Version 2.0, refer to LICENSE.txt
+
+#define BOOST_TEST_MODULE test Skellam
+
+#include "distributions/skellam.hh"
+
+#include <boost/test/included/unit_test.hpp>
+
+namespace tt = boost::test_tools;
+
+BOOST_AUTO_TEST_CASE(simple) {
+  Skellam sd;
+  std::mt19937 prng;
+
+  sd.init_theta(&prng);
+
+  BOOST_TEST(sd.logp_score() == 0.0, tt::tolerance(1e-6));
+  BOOST_TEST(sd.logp(6) == -8.2461659399497425, tt::tolerance(1e-6));
+
+  sd.incorporate(5);
+  sd.incorporate(2);
+  BOOST_TEST(sd.N == 2);
+
+  sd.unincorporate(5);
+  sd.incorporate(7);
+  BOOST_TEST(sd.N == 2);
+
+  BOOST_TEST(sd.logp_score() == -12.676907210873877, tt::tolerance(1e-6));
+  BOOST_TEST(sd.logp(6) == -8.2461659399497425, tt::tolerance(1e-6));
+
+  int s = sd.sample(&prng);
+  BOOST_TEST(s < 100.0);
+}