fixed interface to diffxpy

Author: davidsebfischer

batchglm/models/base/estimator.py

@@ -104,3 +104,15 @@ def validate_data(self, **kwargs):
     @property
     def input_data(self):
         return self._input_data
+
+    @property
+    def X(self):
+        return self.input_data.X
+
+    @property
+    def features(self):
+        return self.input_data.features
+
+    @property
+    def loss(self):
+        return self.params["loss"]

batchglm/models/base_glm/estimator.py

@@ -11,6 +11,7 @@
 ESTIMATOR_PARAMS = MODEL_PARAMS.copy()
 ESTIMATOR_PARAMS.update({
     "loss": (),
+    "full_log_likelihood": ("features",),
     "gradient": ("features",),
     "hessians": ("features", "delta_var0", "delta_var1"),
     "fisher_inv": ("features", "delta_var0", "delta_var1"),
@@ -27,12 +28,8 @@ def __init__(self):
         super(_EstimatorStore_XArray_Base, self).__init__()
 
     @property
-    def input_data(self):
-        return self._input_data
-
-    @property
-    def loss(self):
-        return self.params["loss"]
+    def log_likelihood(self):
+        return self.params["full_log_likelihood"]
 
     @property
     def gradient(self):

batchglm/models/glm_nb/estimator.py

@@ -23,7 +23,7 @@ def __init__(self, estim: AbstractEstimator):
         # causes evaluation of the properties that have not been computed during
         # training, such as the hessian.
         params = estim.to_xarray(
-            ["a_var", "b_var", "loss", "gradient", "hessians", "fisher_inv"],
+            ["a_var", "b_var", "loss", "full_log_likelihood", "gradient", "hessians", "fisher_inv"],
             coords=input_data.data
         )
 

batchglm/train/tf/base/estimator.py

@@ -1,5 +1,4 @@
 import abc
-from enum import Enum
 from typing import Dict, Any, Union, List, Iterable
 
 import os
@@ -247,7 +246,7 @@ def train(self, *args,
                 )
 
                 tf.logging.info(
-                    "Step: %d\tloss: %f\t models converged %i",
+                    "Step: \t%d\t loss: %f\t models converged %i",
                     train_step,
                     global_loss,
                     np.sum(self.model.model_vars.converged).astype("int32")

batchglm/train/tf/base_glm/estimator_graph.py

@@ -218,7 +218,7 @@ def __init__(
                 # with the Cholesky decomposition. This information is
                 # passed here with psd=True.
                 irls_update_a_full, irls_update_a_batched = self.build_updates(
-                    full_lhs=self.full_data_model.fim.fim_a,
+                    full_lhs=self.full_data_model.fim_train.fim_a,
                     batched_lhs=self.batch_fim.fim_a,
                     full_rhs=self.full_data_model.neg_jac_train_a,
                     batched_rhs=self.batch_jac.neg_jac_a,
@@ -231,7 +231,7 @@ def __init__(
 
             if train_r:
                 irls_update_b_full, irls_update_b_batched = self.build_updates(
-                    full_lhs=self.full_data_model.fim.fim_b,
+                    full_lhs=self.full_data_model.fim_train.fim_b,
                     batched_lhs=self.batch_fim.fim_b,
                     full_rhs=self.full_data_model.neg_jac_train_b,
                     batched_rhs=self.batch_jac.neg_jac_b,

batchglm/train/tf/base_glm_all/estimator.py

@@ -293,9 +293,9 @@ def train(self, *args,
         if newton_type_mode:
             if learning_rate != 1:
                 logger.warning(
-                    "Newton-rhapson or IRLS in base_glm_all is used with learing rate " +
+                    "Newton-rhapson or IRLS in base_glm_all is used with learning rate " +
                     str(learning_rate) +
-                    ". Newton-rhapson and IRLS should only be used with learing rate = 1."
+                    ". Newton-rhapson and IRLS should only be used with learning rate = 1."
                 )
 
         # Report all parameters after all defaults were imputed in settings:
@@ -360,17 +360,21 @@ def a_var(self):
     def b_var(self):
         return self.to_xarray("b_var", coords=self.input_data.data.coords)
 
+    @property
+    def loss(self):
+        return self.to_xarray("full_loss")
+
     @property
     def batch_loss(self):
         return self.to_xarray("loss")
 
     @property
-    def batch_gradient(self):
-        return self.to_xarray("gradient", coords=self.input_data.data.coords)
+    def log_likelihood(self):
+        return self.to_xarray("full_log_likelihood", coords=self.input_data.data.coords)
 
     @property
-    def loss(self):
-        return self.to_xarray("full_loss")
+    def batch_gradient(self):
+        return self.to_xarray("gradient", coords=self.input_data.data.coords)
 
     @property
     def gradient(self):
@@ -388,9 +392,9 @@ def finalize(self):
         if self.noise_model == "nb":
             from .external_nb import EstimatorStoreXArray
         else:
-            raise ValueError("noise model not rewcognized")
+            raise ValueError("noise model not recognized")
 
-        logger.debug("Collect and compute ouptut")
+        logger.debug("Collect and compute output")
         store = EstimatorStoreXArray(self)
         logger.debug("Closing session")
         self.close_session()

batchglm/train/tf/base_glm_all/estimator_graph.py

@@ -130,7 +130,7 @@ def map_model(idx, data) -> BasicModelGraph:
             else:
                 hessians_train = hessians_full
 
-            fim_train = FIM(
+            fim_full = FIM(
                 batched_data=batched_data,
                 sample_indices=sample_indices,
                 constraints_loc=constraints_loc,
@@ -139,10 +139,27 @@ def map_model(idx, data) -> BasicModelGraph:
                 mode=pkg_constants.HESSIAN_MODE,
                 noise_model=noise_model,
                 iterator=True,
-                update_a=train_a,
-                update_b=train_b,
+                update_a=True,
+                update_b=True,
                 dtype=dtype
             )
+            # Fisher information matrix of submodel which is to be trained.
+            if not train_a or not train_b:
+                fim_train = FIM(
+                    batched_data=batched_data,
+                    sample_indices=sample_indices,
+                    constraints_loc=constraints_loc,
+                    constraints_scale=constraints_scale,
+                    model_vars=model_vars,
+                    mode=pkg_constants.HESSIAN_MODE,
+                    noise_model=noise_model,
+                    iterator=True,
+                    update_a=train_a,
+                    update_b=train_b,
+                    dtype=dtype
+                )
+            else:
+                fim_train = fim_full
 
         with tf.name_scope("jacobians"):
             # Jacobian of full model for reporting.
@@ -210,7 +227,8 @@ def map_model(idx, data) -> BasicModelGraph:
         self.neg_hessian = hessians_full.neg_hessian
         self.neg_hessian_train = hessians_train.neg_hessian
 
-        self.fim = fim_train
+        self.fim_full = fim_full
+        self.fim_train = fim_train
 
 
 class EstimatorGraphAll(EstimatorGraphGLM):
@@ -439,20 +457,7 @@ def __init__(
                     noise_model=noise_model,
                     dtype=dtype
                 )
-                full_data_loss = full_data_model.loss
-                fisher_inv = op_utils.pinv(full_data_model.neg_hessian)
-
-                # with tf.name_scope("hessian_diagonal"):
-                #     hessian_diagonal = [
-                #         tf.map_fn(
-                #             # elems=tf.transpose(hess, perm=[2, 0, 1]),
-                #             elems=hess,
-                #             fn=tf.diag_part,
-                #             parallel_iterations=pkg_constants.TF_LOOP_PARALLEL_ITERATIONS
-                #         )
-                #         for hess in full_data_model.hessians
-                #     ]
-                #     fisher_a, fisher_b = hessian_diagonal
+                full_data_fisher_inv = op_utils.pinv(full_data_model.neg_hessian)  # TODO switch for fim
 
                 mu = full_data_model.mu
                 r = full_data_model.r
@@ -473,16 +478,17 @@ def __init__(
                 self.r = r
                 self.sigma2 = sigma2
 
+                self.full_loss = full_data_model.loss
+                self.full_log_likelihood = full_data_model.log_likelihood
                 self.batch_probs = batch_model.probs
                 self.batch_log_probs = batch_model.log_probs
                 self.batch_log_likelihood = batch_model.norm_log_likelihood
 
                 self.sample_selection = sample_selection
                 self.full_data_model = full_data_model
 
-                self.full_loss = full_data_loss
                 self.hessians = full_data_model.hessian
-                self.fisher_inv = fisher_inv
+                self.fisher_inv = full_data_fisher_inv
 
                 self.idx_nonconverged = idx_nonconverged
 

batchglm/train/tf/glm_nb/estimator.py

@@ -25,46 +25,27 @@ class TrainingStrategy(Enum):
         DEFAULT = [
             {
                 "convergence_criteria": "all_converged_ll",
-                "stopping_criteria": 1e-8,
+                "stopping_criteria": 1e-6,
                 "use_batching": False,
-                "optim_algo": "Newton",
+                "optim_algo": "irls",
             },
         ]
         QUICK = [
             {
                 "convergence_criteria": "all_converged_ll",
-                "stopping_criteria": 1e-6,
-                "use_batching": False,
-                "optim_algo": "Newton",
-            },
-        ]
-        PRE_INITIALIZED = [
-            {
-                "convergence_criteria": "scaled_moving_average",
-                "stopping_criteria": 1e-10,
-                "loss_window_size": 10,
+                "stopping_criteria": 1e-4,
                 "use_batching": False,
-                "optim_algo": "newton",
+                "optim_algo": "irls",
             },
         ]
-        CONSTRAINED = [  # Should not contain newton-rhapson right now.
+        EXACT = [
             {
-                "learning_rate": 0.5,
                 "convergence_criteria": "all_converged_ll",
                 "stopping_criteria": 1e-8,
-                "loss_window_size": 10,
                 "use_batching": False,
-                "optim_algo": "ADAM",
+                "optim_algo": "irls",
             },
         ]
-        CONTINUOUS = [
-            {
-                "convergence_criteria": "all_converged_ll",
-                "stopping_criteria": 1e-8,
-                "use_batching": False,
-                "optim_algo": "Newton",
-            }
-        ]
 
     def __init__(
             self,