add canonical poisson model

Author: zhengp0

src/regmod/models/__init__.py

@@ -7,6 +7,6 @@
 from .model import Model
 from .negativebinomial import NegativeBinomialModel
 from .pogit import PogitModel
-from .poisson import PoissonModel
+from .poisson import CanonicalPoissonModel, PoissonModel, create_poisson_model
 from .tobit import TobitModel
 from .weibull import WeibullModel

src/regmod/models/poisson.py

@@ -5,12 +5,12 @@
 import numpy as np
 from scipy.stats import poisson
 
-from regmod._typing import Callable, DataFrame, NDArray
+from regmod._typing import Callable, DataFrame, Matrix, NDArray
 from regmod.data import Data
 from regmod.optimizer import msca_optimize
 
 from .model import Model
-from .utils import model_post_init
+from .utils import get_params, model_post_init
 
 
 class PoissonModel(Model):
@@ -34,6 +34,13 @@ def attach_df(self, df: DataFrame):
             self.linear_gvec,
         )
 
+    def get_lin_param(self, coefs: NDArray) -> NDArray:
+        mat = self.mat[0]
+        lin_param = mat.dot(coefs)
+        if self.params[0].offset is not None:
+            lin_param += self.data.get_cols(self.params[0].offset)
+        return lin_param
+
     def hessian_from_gprior(self):
         return self.hmat
 
@@ -157,3 +164,62 @@ def d2nll(self, params: list[NDArray]) -> list[list[NDArray]]:
     def get_ui(self, params: list[NDArray], bounds: tuple[float, float]) -> NDArray:
         mean = params[0]
         return [poisson.ppf(bounds[0], mu=mean), poisson.ppf(bounds[1], mu=mean)]
+
+
+class CanonicalPoissonModel(PoissonModel):
+    def __init__(self, data: Data, **kwargs):
+        super().__init__(data, **kwargs)
+        if self.params[0].inv_link.name != "exp":
+            raise ValueError("Canonical Poisson model requires inverse link to be exp.")
+
+    def objective(self, coefs: NDArray) -> float:
+        weights = self.data.weights * self.data.trim_weights
+        y = self.get_lin_param(coefs)
+        z = np.exp(y)
+
+        prior_obj = self.objective_from_gprior(coefs)
+        likli_obj = weights.dot(z - self.data.obs * y)
+        return prior_obj + likli_obj
+
+    def gradient(self, coefs: NDArray) -> NDArray:
+        mat = self.mat[0]
+        weights = self.data.weights * self.data.trim_weights
+        z = np.exp(self.get_lin_param(coefs))
+
+        prior_grad = self.gradient_from_gprior(coefs)
+        likli_grad = mat.T.dot(weights * (z - self.data.obs))
+        return prior_grad + likli_grad
+
+    def hessian(self, coefs: NDArray) -> Matrix:
+        mat = self.mat[0]
+        weights = self.data.weights * self.data.trim_weights
+        z = np.exp(self.get_lin_param(coefs))
+        likli_hess_scale = weights * z
+
+        prior_hess = self.hessian_from_gprior()
+        likli_hess_right = mat.scale_rows(likli_hess_scale)
+        likli_hess = mat.T.dot(likli_hess_right)
+
+        return prior_hess + likli_hess
+
+    def jacobian2(self, coefs: NDArray) -> NDArray:
+        mat = self.mat[0]
+        weights = self.data.weights * self.data.trim_weights
+        z = np.exp(self.get_lin_param(coefs))
+        likli_jac_scale = weights * (z - self.data.obs)
+
+        likli_jac = mat.T.scale_cols(likli_jac_scale)
+        likli_jac2 = likli_jac.dot(likli_jac.T)
+        return self.hessian_from_gprior() + likli_jac2
+
+
+def create_poisson_model(data: Data, **kwargs) -> PoissonModel:
+    params = get_params(
+        params=kwargs.get("params"),
+        param_specs=kwargs.get("param_specs"),
+        default_param_specs=PoissonModel.default_param_specs,
+    )
+
+    if params[0].inv_link.name == "exp":
+        return CanonicalPoissonModel(data, params=params)
+    return PoissonModel(data, params=params)

tests/test_poissonmodel.py

@@ -1,15 +1,19 @@
 """
 Test Poisson Model
 """
+
 import numpy as np
 import pandas as pd
 import pytest
-
 from regmod.data import Data
 from regmod.function import fun_dict
-from regmod.models import PoissonModel
-from regmod.prior import (GaussianPrior, SplineGaussianPrior,
-                          SplineUniformPrior, UniformPrior)
+from regmod.models import create_poisson_model
+from regmod.prior import (
+    GaussianPrior,
+    SplineGaussianPrior,
+    SplineUniformPrior,
+    UniformPrior,
+)
 from regmod.utils import SplineSpecs
 from regmod.variable import SplineVariable, Variable
 
@@ -19,27 +23,27 @@
 @pytest.fixture
 def data():
     num_obs = 5
-    df = pd.DataFrame({
-        "obs": np.random.rand(num_obs)*10,
-        "cov0": np.random.randn(num_obs),
-        "cov1": np.random.randn(num_obs)
-    })
-    return Data(col_obs="obs",
-                col_covs=["cov0", "cov1"],
-                df=df)
+    df = pd.DataFrame(
+        {
+            "obs": np.random.rand(num_obs) * 10,
+            "cov0": np.random.randn(num_obs),
+            "cov1": np.random.randn(num_obs),
+        }
+    )
+    return Data(col_obs="obs", col_covs=["cov0", "cov1"], df=df)
 
 
 @pytest.fixture
 def wrong_data():
     num_obs = 5
-    df = pd.DataFrame({
-        "obs": np.random.randn(num_obs),
-        "cov0": np.random.randn(num_obs),
-        "cov1": np.random.randn(num_obs)
-    })
-    return Data(col_obs="obs",
-                col_covs=["cov0", "cov1"],
-                df=df)
+    df = pd.DataFrame(
+        {
+            "obs": np.random.randn(num_obs),
+            "cov0": np.random.randn(num_obs),
+            "cov1": np.random.randn(num_obs),
+        }
+    )
+    return Data(col_obs="obs", col_covs=["cov0", "cov1"], df=df)
 
 
 @pytest.fixture
@@ -54,9 +58,9 @@ def uprior():
 
 @pytest.fixture
 def spline_specs():
-    return SplineSpecs(knots=np.linspace(0.0, 1.0, 5),
-                       degree=3,
-                       knots_type="rel_domain")
+    return SplineSpecs(
+        knots=np.linspace(0.0, 1.0, 5), degree=3, knots_type="rel_domain"
+    )
 
 
 @pytest.fixture
@@ -71,20 +75,21 @@ def spline_uprior():
 
 @pytest.fixture
 def var_cov0(gprior, uprior):
-    return Variable(name="cov0",
-                    priors=[gprior, uprior])
+    return Variable(name="cov0", priors=[gprior, uprior])
 
 
 @pytest.fixture
 def var_cov1(spline_gprior, spline_uprior, spline_specs):
-    return SplineVariable(name="cov1",
-                          spline_specs=spline_specs,
-                          priors=[spline_gprior, spline_uprior])
+    return SplineVariable(
+        name="cov1", spline_specs=spline_specs, priors=[spline_gprior, spline_uprior]
+    )
 
 
 @pytest.fixture
 def model(data, var_cov0, var_cov1):
-    return PoissonModel(data, param_specs={"lam": {"variables": [var_cov0, var_cov1]}})
+    return create_poisson_model(
+        data, param_specs={"lam": {"variables": [var_cov0, var_cov1]}}
+    )
 
 
 def test_model_size(model, var_cov0, var_cov1):
@@ -124,7 +129,7 @@ def test_model_gradient(model, inv_link):
     tr_grad = np.zeros(model.size)
     for i in range(model.size):
         coefs_c[i] += 1e-16j
-        tr_grad[i] = model.objective(coefs_c).imag/1e-16
+        tr_grad[i] = model.objective(coefs_c).imag / 1e-16
         coefs_c[i] -= 1e-16j
     assert np.allclose(my_grad, tr_grad)
 
@@ -139,15 +144,17 @@ def test_model_hessian(model, inv_link):
     for i in range(model.size):
         for j in range(model.size):
             coefs_c[j] += 1e-16j
-            tr_hess[i][j] = model.gradient(coefs_c).imag[i]/1e-16
+            tr_hess[i][j] = model.gradient(coefs_c).imag[i] / 1e-16
             coefs_c[j] -= 1e-16j
 
     assert np.allclose(my_hess, tr_hess)
 
 
 def test_wrong_data(wrong_data, var_cov0, var_cov1):
     with pytest.raises(ValueError):
-        PoissonModel(wrong_data, param_specs={"lam": {"variables": [var_cov0, var_cov1]}})
+        create_poisson_model(
+            wrong_data, param_specs={"lam": {"variables": [var_cov0, var_cov1]}}
+        )
 
 
 def test_get_ui(model):
@@ -160,16 +167,18 @@ def test_get_ui(model):
 
 def test_model_no_variables():
     num_obs = 5
-    df = pd.DataFrame({
-        "obs": np.random.rand(num_obs)*10,
-        "offset": np.ones(num_obs),
-    })
+    df = pd.DataFrame(
+        {
+            "obs": np.random.rand(num_obs) * 10,
+            "offset": np.ones(num_obs),
+        }
+    )
     data = Data(
         col_obs="obs",
         col_offset="offset",
         df=df,
     )
-    model = PoissonModel(data, param_specs={"lam": {"offset": "offset"}})
+    model = create_poisson_model(data, param_specs={"lam": {"offset": "offset"}})
     coefs = np.array([])
     grad = model.gradient(coefs)
     hessian = model.hessian(coefs)