[ENH] GLM with multiple distributions and link function support

skpro/regression/linear/_glm.py

@@ -1,6 +1,8 @@
-"""Interface adapter for the Generalized Linear Model Regressor with Gaussian Link."""
+"""Interface adapter for the Generalized Linear Model Regressor."""
 # copyright: skpro developers, BSD-3-Clause License (see LICENSE file)
 
+__author__ = ["ShreeshaM07", "julian-fong"]
+
 import pandas as pd
 
 from skpro.regression.base import BaseProbaRegressor
@@ -18,6 +20,16 @@ class GLMRegressor(BaseProbaRegressor):
 
     Parameters
     ----------
+    family : str
+        Available options are
+        1.Normal
+        2.Poisson
+        3.Gamma
+    link : str
+        Available safe options are
+        Normal : Log, Identity, InversePower
+        Poisson : Log, Identity, Sqrt
+        Gamma : Log, Identity, InversePower
     missing : str
         Available options are 'none', 'drop' and 'raise'. If 'none', no nan
         checking is done. If 'drop', any observations with nans are dropped.
@@ -157,8 +169,8 @@ class GLMRegressor(BaseProbaRegressor):
     """
 
     _tags = {
-        "authors": ["julian-fong"],
-        "maintainers": ["julian-fong"],
+        "authors": ["ShreeshaM07", "julian-fong"],
+        "maintainers": ["ShreeshaM07", "julian-fong"],
         "python_version": None,
         "python_dependencies": "statsmodels",
         "capability:multioutput": False,
@@ -167,8 +179,44 @@ class GLMRegressor(BaseProbaRegressor):
         "y_inner_mtype": "pd_DataFrame_Table",
     }
 
+    def _str_to_sm_family(self, family, link):
+        """Convert the string to a statsmodel object.
+
+        If the link function is also explcitly mentioned then include then
+        that must be passed to the family/distribution object.
+        """
+        from warnings import warn
+
+        from statsmodels.genmod.families.family import Gamma, Gaussian, Poisson
+        from statsmodels.genmod.families.links import Identity, InversePower, Log, Sqrt
+
+        sm_fmly = {
+            "Normal": Gaussian,
+            "Poisson": Poisson,
+            "Gamma": Gamma,
+        }
+
+        links = {
+            "Log": Log,
+            "Identity": Identity,
+            "InversePower": InversePower,
+            "Sqrt": Sqrt,
+        }
+
+        if link in links:
+            link_function = links[link]()
+            try:
+                return sm_fmly[family](link_function)
+            except Exception:
+                msg = "Invalid link for family, default link will be used"
+                warn(msg)
+
+        return sm_fmly[family]()
+
     def __init__(
         self,
+        family="Normal",
+        link=None,
         missing="none",
         start_params=None,
         maxiter=100,
@@ -184,9 +232,11 @@ def __init__(
         add_constant=False,
     ):
         super().__init__()
-        from statsmodels.genmod.families.family import Gaussian
 
-        self._family = Gaussian()
+        if family is None:
+            family = "Normal"
+        self.family = family
+        self.link = link
         self.missing = missing
         self.start_params = start_params
         self.maxiter = maxiter
@@ -231,10 +281,14 @@ def _fit(self, X, y):
 
         y_col = y.columns
 
+        family = self.family
+        link = self.link
+        sm_family = self._str_to_sm_family(family=family, link=link)
+
         glm_estimator = GLM(
             endog=y,
             exog=X_,
-            family=self._family,
+            family=sm_family,
             missing=self.missing,
         )
 
@@ -313,6 +367,47 @@ def _predict(self, X):
 
         return y_pred
 
+    def _params_sm_to_skpro(self, y_predictions_df, index, columns, family):
+        """Convert the statsmodels output to equivalent skpro distribution."""
+        from skpro.distributions.gamma import Gamma
+        from skpro.distributions.normal import Normal
+        from skpro.distributions.poisson import Poisson
+
+        skpro_distr = {
+            "Normal": Normal,
+            "Poisson": Poisson,
+            "Gamma": Gamma,
+        }
+
+        params = {}
+        skp_dist = Normal
+
+        if family in skpro_distr:
+            skp_dist = skpro_distr[family]
+
+        if skp_dist == Normal:
+            y_mu = y_predictions_df["mean"].rename("mu").to_frame()
+            y_sigma = y_predictions_df["mean_se"].rename("sigma").to_frame()
+            params["mu"] = y_mu
+            params["sigma"] = y_sigma
+        elif skp_dist == Poisson:
+            y_mu = y_predictions_df["mean"].rename("mu").to_frame()
+            params["mu"] = y_mu
+        elif skp_dist == Gamma:
+            y_mean = y_predictions_df["mean"]
+            y_sd = y_predictions_df["mean_se"]
+            y_alpha = (y_mean / y_sd) ** 2
+            y_beta = (y_alpha / y_mean).rename("beta").to_frame()
+            y_alpha = y_alpha.rename("alpha").to_frame()
+            params["alpha"] = y_alpha
+            params["beta"] = y_beta
+
+        params["index"] = index
+        params["columns"] = columns
+
+        y_pred = skp_dist(**params)
+        return y_pred
+
     def _predict_proba(self, X):
         """Predict distribution over labels for data from features.
 
@@ -332,24 +427,20 @@ def _predict_proba(self, X):
         y_pred : skpro BaseDistribution, same length as `X`
             labels predicted for `X`
         """
-        from skpro.distributions.normal import Normal
-
         X_ = self._prep_x(X)
 
         # instead of using the conventional predict() method, we use statsmodels
         # get_prediction method, which returns a pandas df that contains
         # the prediction and prediction variance i.e mu and sigma
         y_column = self.y_col
         y_predictions_df = self.glm_fit_.get_prediction(X_).summary_frame()
-        y_mu = y_predictions_df["mean"].rename("mu").to_frame()
-        y_sigma = y_predictions_df["mean_se"].rename("sigma").to_frame()
-        params = {
-            "mu": y_mu,
-            "sigma": y_sigma,
-            "index": X_.index,
-            "columns": y_column,
-        }
-        y_pred = Normal(**params)
+
+        # convert the returned values to skpro equivalent distribution
+        family = self.family
+        index = X_.index
+        columns = y_column
+
+        y_pred = self._params_sm_to_skpro(y_predictions_df, index, columns, family)
         return y_pred
 
     def _prep_x(self, X):
@@ -395,5 +486,19 @@ def get_test_params(cls, parameter_set="default"):
         """
         params1 = {}
         params2 = {"add_constant": True}
+        params3 = {
+            "family": "Poisson",
+            "add_constant": True,
+        }
+        params4 = {"family": "Gamma"}
+        params5 = {
+            "family": "Normal",
+            "link": "InversePower",
+        }
+        params6 = {
+            "family": "Poisson",
+            "link": "Log",
+            "add_constant": True,
+        }
 
-        return [params1, params2]
+        return [params1, params2, params3, params4, params5, params6]