Normal Model (#139)

* [WIP] The initializer...does everything?

* [WIP] Refactor to fit model, accuracy test passes.

* [WIP] Rempve unecessary tests.

* Seems right?

* Fill in functions that might be needed for testing.

* Make tests more consistent.

* Clean up model a bit.

* [WIP] Fix log likelihood?

* Remove unfinished methods

* [WIP] Add normal scale modeling

* [WIP] Fix normal model ll

* Small fixes and clean up.

* Fix missing export.

* Fix init_par function

* Fix init_par for scale.

* Add docs.

* Move to completed.

* Fix typing issue.

* Fix typing issue.

* Small test cleanups.

* Move ll function.

* Move ll function.

* Ok remove!

* Add changes from update_api

* pre-commit

* Missed that variable!

Author: ilan-gold

batchglm/models/glm_norm/model.py

@@ -12,95 +12,42 @@ class Model(_ModelGLM, metaclass=abc.ABCMeta):
     """Generalized Linear Model (GLM) with normal noise."""
 
     def link_loc(self, data) -> Union[np.ndarray, dask.array.core.Array]:
-        """Short summary.
-
-        :param type data: Description of parameter `data`.
-        :return: Description of returned object.
-        :rtype: type
-
-        """
         return data
 
     def inverse_link_loc(self, data) -> Union[np.ndarray, dask.array.core.Array]:
-        """Short summary.
-
-        :param type data: Description of parameter `data`.
-        :return: Description of returned object.
-        :rtype: type
-
-        """
         return data
 
     def link_scale(self, data) -> Union[np.ndarray, dask.array.core.Array]:
-        """Short summary.
-
-        :param type data: Description of parameter `data`.
-        :return: Description of returned object.
-        :rtype: type
-
-        """
         return np.log(data)
 
     def inverse_link_scale(self, data) -> Union[np.ndarray, dask.array.core.Array]:
-        """Short summary.
-
-        :param type data: Description of parameter `data`.
-        :return: Description of returned object.
-        :rtype: type
-
-        """
         return np.exp(data)
 
     @property
     def eta_loc(self) -> Union[np.ndarray, dask.array.core.Array]:
-        """Short summary.
-
-        :return: Description of returned object.
-        :rtype: np.ndarray
-
-        """
         eta = np.matmul(self.design_loc, self.theta_location_constrained)
         if self.size_factors is not None:
-            eta *= np.expand_dims(self.size_factors, axis=1)
+            eta *= self.size_factors
         return eta
 
     def eta_loc_j(self, j) -> Union[np.ndarray, dask.array.core.Array]:
-        """Short summary.
-
-        :param type j: Description of parameter `j`.
-        :return: Description of returned object.
-        :rtype: np.ndarray
-
-        """
         # Make sure that dimensionality of sliced array is kept:
         if isinstance(j, int) or isinstance(j, np.int32) or isinstance(j, np.int64):
             j = [j]
         eta = np.matmul(self.design_loc, self.theta_location_constrained[:, j])
         if self.size_factors is not None:
-            eta *= np.expand_dims(self.size_factors, axis=1)
+            eta *= self.size_factors
         eta = self.np_clip_param(eta, "eta_loc")
         return eta
 
     # Re-parameterizations:
 
     @property
     def mean(self) -> Union[np.ndarray, dask.array.core.Array]:
-        """Short summary.
-
-        :return: Description of returned object.
-        :rtype: np.ndarray
-
-        """
         return self.location
 
     @property
     def sd(self) -> Union[np.ndarray, dask.array.core.Array]:
-        """Short summary.
-
-        :return: Description of returned object.
-        :rtype: np.ndarray
-
-        """
         return self.scale
 
     # param constraints:

batchglm/models/glm_norm/utils.py

@@ -1,47 +1,18 @@
 import logging
-from typing import Union
+from typing import Tuple, Union
 
+import dask
 import numpy as np
 import scipy.sparse
 
-from .external import closedform_glm_mean, closedform_glm_scale
+from .external import closedform_glm_scale
 
 logger = logging.getLogger("batchglm")
 
 
-def closedform_norm_glm_mean(
-    x: Union[np.ndarray, scipy.sparse.csr_matrix],
-    design_loc: np.ndarray,
-    constraints_loc,
-    size_factors=None,
-    link_fn=lambda x: x,
-    inv_link_fn=lambda x: x,
-):
-    r"""
-    Calculates a closed-form solution for the `mean` parameters of normal GLMs.
-
-    :param x: The sample data
-    :param design_loc: design matrix for location
-    :param constraints_loc: tensor (all parameters x dependent parameters)
-        Tensor that encodes how complete parameter set which includes dependent
-        parameters arises from indepedent parameters: all = <constraints, indep>.
-        This form of constraints is used in vector generalized linear models (VGLMs).
-    :param size_factors: size factors for X
-    :return: tuple: (groupwise_means, mean, rmsd)
-    """
-    return closedform_glm_mean(
-        x=x,
-        dmat=design_loc,
-        constraints=constraints_loc,
-        size_factors=size_factors,
-        link_fn=link_fn,
-        inv_link_fn=inv_link_fn,
-    )
-
-
 def closedform_norm_glm_logsd(
-    x: Union[np.ndarray, scipy.sparse.csr_matrix],
-    design_scale: np.ndarray,
+    x: Union[np.ndarray, scipy.sparse.csr_matrix, dask.array.core.Array],
+    design_scale: Union[np.ndarray, dask.array.core.Array],
     constraints=None,
     size_factors=None,
     groupwise_means=None,
@@ -71,3 +42,73 @@ def compute_scales_fun(variance, mean):
         link_fn=link_fn,
         compute_scales_fun=compute_scales_fun,
     )
+
+
+def init_par(model, init_location: str, init_scale: str) -> Tuple[np.ndarray, np.ndarray, bool, bool]:
+    r"""
+    standard:
+    Only initialise intercept and keep other coefficients as zero.
+
+    closed-form:
+    Initialize with Maximum Likelihood / Maximum of Momentum estimators
+
+    Idea:
+    $$
+        \theta &= f(x) \\
+        \Rightarrow f^{-1}(\theta) &= x \\
+            &= (D \cdot D^{+}) \cdot x \\
+            &= D \cdot (D^{+} \cdot x) \\
+            &= D \cdot x' = f^{-1}(\theta)
+    $$
+    """
+
+    groupwise_means = None
+
+    init_location_str = init_location.lower()
+    # Chose option if auto was chosen
+    auto_or_closed_form = init_location_str == "auto" or init_location_str == "closed_form"
+    if auto_or_closed_form or init_location_str == "all_zero":
+        if auto_or_closed_form:
+            logger.warning(
+                (
+                    "There is no need for closed form location model initialization"
+                    "because it is already closed form - falling back to zeros"
+                )
+            )
+        init_theta_location = np.zeros([model.num_loc_params, model.num_features])
+    elif init_location_str == "standard":
+        overall_means = np.mean(model.x, axis=0)  # directly calculate the mean
+        init_theta_location = np.zeros([model.num_loc_params, model.num_features])
+        init_theta_location[0, :] = np.log(overall_means)
+    else:
+        raise ValueError("init_location string %s not recognized" % init_location)
+
+    init_scale_str = init_scale.lower()
+    if init_scale_str == "auto":
+        init_scale_str = "standard"
+
+    if init_scale_str == "standard":
+        groupwise_scales, init_scale_intercept, rmsd_b = closedform_norm_glm_logsd(
+            x=model.x,
+            design_scale=model.design_scale[:, [0]],
+            constraints=model.constraints_scale[[0], :][:, [0]],
+            size_factors=model.size_factors,
+            groupwise_means=None,
+            link_fn=lambda r: np.log(r + np.nextafter(0, 1, dtype=r.dtype)),
+        )
+        init_theta_scale = np.zeros([model.num_scale_params, model.num_features])
+        init_theta_scale[0, :] = init_scale_intercept
+    elif init_scale_str == "closed_form":
+        groupwise_scales, init_theta_scale, rmsd_b = closedform_norm_glm_logsd(
+            x=model.x,
+            design_scale=model.design_scale,
+            constraints=model.constraints_scale,
+            size_factors=model.size_factors,
+            groupwise_means=groupwise_means,
+        )
+    elif init_scale_str == "all_zero":
+        init_theta_scale = np.zeros([model.num_scale_params, model.x.shape[1]])
+    else:
+        raise ValueError("init_scale string %s not recognized" % init_scale_str)
+
+    return init_theta_location, init_theta_scale, True, True

batchglm/train/numpy/base_glm/model_container.py

@@ -327,7 +327,7 @@ def fim(self) -> Union[np.ndarray, dask.array.core.Array]:
         fim_scale_scale = self.fim_scale_scale
         fim_location_scale = self.fim_location_scale
         fim_ba = np.transpose(fim_location_scale, axes=[0, 2, 1])
-        return -np.concatenate(
+        return np.concatenate(
             [
                 np.concatenate([fim_location_location, fim_location_scale], axis=2),
                 np.concatenate([fim_ba, fim_scale_scale], axis=2),

batchglm/train/numpy/glm_nb/model_container.py

@@ -18,7 +18,7 @@ def fim_weight_location_location(self) -> Union[np.ndarray, dask.array.core.Arra
         Fisher inverse matrix weights
         :return: observations x features
         """
-        return -self.location * self.scale / (self.scale + self.location)
+        return self.location * self.scale / (self.scale + self.location)
 
     @property
     def ybar(self) -> Union[np.ndarray, dask.array.core.Array]:

batchglm/train/numpy/glm_norm/__init__.py

@@ -0,0 +1,2 @@
+from .estimator import Estimator
+from .model_container import ModelContainer

batchglm/train/numpy/glm_norm/estimator.py

@@ -0,0 +1,75 @@
+import logging
+
+import numpy as np
+
+from .external import EstimatorGlm, Model, init_par
+from .model_container import ModelContainer
+
+logger = logging.getLogger("batchglm")
+
+
+class Estimator(EstimatorGlm):
+    def __init__(
+        self,
+        model: Model,
+        init_location: str = "AUTO",
+        init_scale: str = "AUTO",
+        # batch_size: Optional[Union[Tuple[int, int], int]] = None,
+        quick_scale: bool = False,
+        dtype: str = "float64",
+    ):
+        """
+        Performs initialisation and creates a new estimator.
+        :param model:
+            The GLM model to be fit
+        :param init_location: (Optional)
+            Low-level initial values for a. Can be:
+
+            - str:
+                * "auto": automatically choose best initialization
+                * "standard": initialize intercept with observed mean
+                * "closed_form": try to initialize with closed form
+            - np.ndarray: direct initialization of 'a'
+        :param init_scale: (Optional)
+            Low-level initial values for b. Can be:
+
+            - str:
+                * "auto": automatically choose best initialization
+                * "random": initialize with random values
+                * "standard": initialize with zeros
+                * "closed_form": try to initialize with closed form
+            - np.ndarray: direct initialization of 'b'
+        :param quick_scale: bool
+            Whether `scale` will be fitted faster and maybe less accurate.
+            Useful in scenarios where fitting the exact `scale` is not absolutely necessary.
+        :param dtype: Numerical precision.
+        """
+        init_theta_location, init_theta_scale, train_loc, train_scale = init_par(
+            model=model, init_location=init_location, init_scale=init_scale
+        )
+        init_theta_location = init_theta_location.astype(dtype)
+        init_theta_scale = init_theta_scale.astype(dtype)
+        self._train_scale = train_scale
+        self._train_loc = train_loc
+        if quick_scale:
+            self._train_scale = False
+        _model_container = ModelContainer(
+            model=model,
+            init_theta_location=init_theta_location,
+            init_theta_scale=init_theta_scale,
+            chunk_size_genes=model.chunk_size_genes,
+            dtype=dtype,
+        )
+        super(Estimator, self).__init__(model_container=_model_container, dtype=dtype)
+
+    def train(
+        self,
+        **kwargs,
+    ):
+        model = self._model_container.model
+        if self._train_loc:
+            theta_location, _, _, _ = np.linalg.lstsq(model.design_loc, model.x)
+            self._model_container.theta_location = theta_location
+        self._train_loc = False
+        super().train(**kwargs)
+        self._train_loc = True

batchglm/train/numpy/glm_norm/external.py

@@ -0,0 +1,9 @@
+import batchglm.utils.data as data_utils
+from batchglm import pkg_constants
+from batchglm.models.base_glm.utils import closedform_glm_scale
+from batchglm.models.glm_norm.model import Model
+from batchglm.models.glm_norm.utils import closedform_norm_glm_logsd, init_par
+
+# import necessary base_glm layers
+from batchglm.train.numpy.base_glm import BaseModelContainer, EstimatorGlm
+from batchglm.utils.linalg import groupwise_solve_lm