Merge branch 'dev' of github.com:theislab/batchglm into dev

Author: Hoeze

batchglm/train/tf/nb_glm/base.py

@@ -77,9 +77,15 @@ def np_clip_param(param, name):
     )
 
 
-def apply_constraints(constraints: np.ndarray, var: tf.Variable, dtype: str):
+def apply_constraints(
+        constraints: np.ndarray,
+        dtype: str,
+        var_all: tf.Variable = None,
+        var_indep: tf.Tensor = None
+):
     """ Iteratively build depend variables from other variables via constraints
 
+    :type var_all: object
     :param constraints: Array with constraints in rows and model parameters in columns.
         Each constraint contains non-zero entries for the a of parameters that
         has to sum to zero. This constraint is enforced by binding one parameter
@@ -88,16 +94,21 @@ def apply_constraints(constraints: np.ndarray, var: tf.Variable, dtype: str):
         parameter is indicated by a -1 in this array, the independent parameters
         of that constraint (which may be dependent at an earlier constraint)
         are indicated by a 1.
-    :param var: Variable tensor features x independent parameters.
+    :param var_all: Variable tensor features x independent parameters.
+        All model parameters.
+    :param var_all: Variable tensor features x independent parameters.
+        Only independent model parameters, ie. not parameters defined by constraints.
     :param dtype: Precision used in tensorflow.
 
     :return: Full model parameter matrix with dependent parameters.
     """
 
     # Find all independent variables:
-    idx_indep = np.where(np.all(constraints == -1, axis=0))[0]
+    idx_indep = np.where(np.all(constraints != -1, axis=0))[0]
+    idx_indep.astype(dtype=np.int64)
     # Relate constraints to dependent variables:
     idx_dep = np.array([np.where(constr == -1)[0] for constr in constraints])
+    idx_dep.astype(dtype=np.int64)
     # Only choose dependent variable which was not already defined above:
     idx_dep = np.concatenate([
         x[[xx not in np.concatenate(idx_dep[:i]) for xx in x]] if i > 0 else x
@@ -109,7 +120,13 @@ def apply_constraints(constraints: np.ndarray, var: tf.Variable, dtype: str):
     # tensor is initialised with the independent variables var
     # and is grown by one varibale in each iteration until
     # all variables are there.
-    x = var
+    if var_all is None:
+        x = var_indep
+    elif var_indep is None:
+        x = tf.gather(params=var_all, indices=idx_indep, axis=0)
+    else:
+        raise ValueError("only give var_all or var_indep to apply_constraints.")
+
     for i in range(constraints.shape[0]):
         idx_var_i = np.concatenate([idx_indep, idx_dep[:i]])
         constraint_model = constraints[[i], :][:, idx_var_i]
@@ -150,12 +167,11 @@ def __init__(
         # Define first layer of computation graph on identifiable variables
         # to yield dependent set of parameters of model for each location
         # and scale model.
-
         if constraints_loc is not None:
-            a = apply_constraints(constraints_loc, a, dtype=dtype)
+            a = apply_constraints(constraints=constraints_loc, var_all=a, dtype=dtype)
 
         if constraints_scale is not None:
-            b = apply_constraints(constraints_scale, b, dtype=dtype)
+            b = apply_constraints(constraints=constraints_scale, var_all=b, dtype=dtype)
 
         with tf.name_scope("mu"):
             log_mu = tf.matmul(design_loc, a, name="log_mu_obs")

batchglm/train/tf/nb_glm/estimator.py

@@ -73,18 +73,9 @@ def map_model(idx, data) -> BasicModelGraph:
         with tf.name_scope("loss"):
             loss = tf.reduce_sum(norm_neg_log_likelihood)
 
-        # TODO: remove this and decide for one implementation
-        if pkg_constants.HESSIAN_MODE == "obs":
-            # Only need iterator that yields single observations for hessian mode obs:
-            singleobs_data = dataset.map(fetch_fn, num_parallel_calls=pkg_constants.TF_NUM_THREADS)
-            singleobs_data = singleobs_data.prefetch(1)
-        else:
-            singleobs_data = None
-
         with tf.name_scope("hessians"):
             hessians = Hessians(
                 batched_data=batched_data,
-                singleobs_data=singleobs_data,
                 sample_indices=sample_indices,
                 constraints_loc=constraints_loc,
                 constraints_scale=constraints_scale,
@@ -258,20 +249,19 @@ def __init__(
                     constraints_loc=constraints_loc,
                     constraints_scale=constraints_scale,
                     model_vars=model_vars,
-                    mode="analytic",
+                    mode=pkg_constants.JACOBIAN_MODE,  #"analytic",
                     iterator=False,
                     dtype=dtype
                 )
 
                 # Define the hessian on the batched model for newton-rhapson:
                 batch_hessians = Hessians(
                     batched_data=batch_data,
-                    singleobs_data=None,
                     sample_indices=batch_sample_index,
                     constraints_loc=constraints_loc,
                     constraints_scale=constraints_scale,
                     model_vars=model_vars,
-                    mode="obs_batched",
+                    mode=pkg_constants.HESSIAN_MODE,  #"obs_batched",
                     iterator=False,
                     dtype=dtype
                 )
@@ -637,6 +627,16 @@ class TrainingStrategy(Enum):
                 "optim_algo": "newton",
             },
         ]
+        CONSTRAINED = [  # Should not contain newton-rhapson right now.
+            {
+                "learning_rate": 0.5,
+                "convergence_criteria": "scaled_moving_average",
+                "stopping_criteria": 1e-10,
+                "loss_window_size": 10,
+                "use_batching": False,
+                "optim_algo": "ADAM",
+            },
+        ]
         CONTINUOUS = [
             {
                 "learning_rate": 0.5,
@@ -810,7 +810,7 @@ def __init__(
                             design_scale=input_data.design_scale,
                             constraints=input_data.constraints_scale,
                             size_factors=size_factors_init,
-                            groupwise_means=groupwise_means,
+                            groupwise_means=None,  # Could only use groupwise_means from a init if design_loc and design_scale were the same.
                             link_fn=lambda r: np.log(np_clip_param(r, "r"))
                         )
 

batchglm/train/tf/nb_glm/hessians.py

@@ -175,7 +175,6 @@ class Hessians:
     def __init__(
             self,
             batched_data: tf.data.Dataset,
-            singleobs_data: tf.data.Dataset,
             sample_indices: tf.Tensor,
             constraints_loc,
             constraints_scale,
@@ -188,7 +187,6 @@ def __init__(
 
         :param batched_data:
             Dataset iterator over mini-batches of data (used for training) or tf.Tensors of mini-batch.
-        :param singleobs_data: Dataset iterator over single observation batches of data.
         :param sample_indices: Indices of samples to be used.
         :param constraints_loc: Constraints for location model.
             Array with constraints in rows and model parameters in columns.
@@ -218,29 +216,14 @@ def __init__(
             evaluation of the hessian via the tf.hessian function,
             which is done by feature for implementation reasons.
         :param iterator: bool
-            Whether an iterator or a tensor (single yield of an iterator) is given
-            in
+            Whether batched_data is an iterator or a tensor (such as single yield of an iterator).
         """
-        if constraints_loc != None and mode != "tf":
+        if constraints_loc is not None and mode != "tf":
             raise ValueError("closed form hessian does not work if constraints_loc is not None")
-        if constraints_scale != None and mode != "tf":
+        if constraints_scale is not None and mode != "tf":
             raise ValueError("closed form hessian does not work if constraints_scale is not None")
 
-        if mode == "obs":
-            logger.info("Performance warning for hessian mode: " +
-                        "obs_batched is strongly recommended as an alternative to obs.")
-            self.hessian = self.byobs(
-                batched_data=singleobs_data,
-                sample_indices=sample_indices,
-                constraints_loc=constraints_loc,
-                constraints_scale=constraints_scale,
-                model_vars=model_vars,
-                batched=False,
-                iterator=iterator,
-                dtype=dtype
-            )
-            self.neg_hessian = tf.negative(self.hessian)
-        elif mode == "obs_batched":
+        if mode == "obs_batched":
             self.hessian = self.byobs(
                 batched_data=batched_data,
                 sample_indices=sample_indices,
@@ -259,6 +242,7 @@ def __init__(
                 constraints_loc=constraints_loc,
                 constraints_scale=constraints_scale,
                 model_vars=model_vars,
+                iterator=iterator,
                 dtype=dtype
             )
             self.neg_hessian = tf.negative(self.hessian)
@@ -272,6 +256,7 @@ def __init__(
                 constraints_loc=constraints_loc,
                 constraints_scale=constraints_scale,
                 model_vars=model_vars,
+                iterator=iterator,
                 dtype=dtype
             )
             self.hessian = tf.negative(self.neg_hessian)
@@ -517,8 +502,8 @@ def _red(prev, cur):
             return tf.add(prev, cur)
 
         params = model_vars.params
-        p_shape_a = model_vars.a.shape[0]
-        p_shape_b = model_vars.b.shape[0]
+        p_shape_a = model_vars.a_var.shape[0]  # This has to be _var to work with constraints.
+        p_shape_b = model_vars.b_var.shape[0]  # This has to be _var to work with constraints.
 
         if iterator:
             H = op_utils.map_reduce(
@@ -542,6 +527,7 @@ def byfeature(
             constraints_loc,
             constraints_scale,
             model_vars: ModelVars,
+            iterator,
             dtype
     ):
         """
@@ -685,18 +671,24 @@ def _red(prev, cur):
             return [tf.add(p, c) for p, c in zip(prev, cur)]
 
         params = model_vars.params
-        p_shape_a = model_vars.a.shape[0]
-        p_shape_b = model_vars.b.shape[0]
-
-        H = op_utils.map_reduce(
-            last_elem=tf.gather(sample_indices, tf.size(sample_indices) - 1),
-            data=batched_data,
-            map_fn=_map,
-            reduce_fn=_red,
-            parallel_iterations=1,
-        )
-        H = H[0]
-        return H
+        p_shape_a = model_vars.a_var.shape[0]  # This has to be _var to work with constraints.
+        p_shape_b = model_vars.b_var.shape[0]  # This has to be _var to work with constraints.
+
+        if iterator:
+            H = op_utils.map_reduce(
+                last_elem=tf.gather(sample_indices, tf.size(sample_indices) - 1),
+                data=batched_data,
+                map_fn=_map,
+                reduce_fn=_red,
+                parallel_iterations=1
+            )
+        else:
+            H = _map(
+                idx=sample_indices,
+                data=batched_data
+            )
+
+        return H[0]
 
     def tf_byfeature(
             self,
@@ -705,6 +697,7 @@ def tf_byfeature(
             constraints_loc,
             constraints_scale,
             model_vars: ModelVars,
+            iterator,
             dtype
     ) -> List[tf.Tensor]:
         """
@@ -804,20 +797,27 @@ def _map(idx, data):
                 constraints_loc=constraints_loc,
                 constraints_scale=constraints_scale,
                 params=model_vars.params,
-                p_shape_a=model_vars.a.shape[0],
-                p_shape_b=model_vars.b.shape[0],
+                p_shape_a=model_vars.a_var.shape[0],  # This has to be _var to work with constraints.
+                p_shape_b=model_vars.b_var.shape[0],  # This has to be _var to work with constraints.
                 dtype=dtype,
                 size_factors=size_factors
             )
 
         def _red(prev, cur):
             return [tf.add(p, c) for p, c in zip(prev, cur)]
 
-        H = op_utils.map_reduce(
-            last_elem=tf.gather(sample_indices, tf.size(sample_indices) - 1),
-            data=batched_data,
-            map_fn=_map,
-            reduce_fn=_red,
-            parallel_iterations=1,
-        )
+        if iterator:
+            H = op_utils.map_reduce(
+                last_elem=tf.gather(sample_indices, tf.size(sample_indices) - 1),
+                data=batched_data,
+                map_fn=_map,
+                reduce_fn=_red,
+                parallel_iterations=1
+            )
+        else:
+            H = _map(
+                idx=sample_indices,
+                data=batched_data
+            )
+
         return H[0]

batchglm/train/tf/nb_glm/jacobians.py

@@ -159,9 +159,9 @@ def __init__(
             in
         """
         if constraints_loc is not None and mode != "tf":
-            raise ValueError("closed form hessian does not work if constraints_loc is not None")
+            raise ValueError("closed form jacobians do not work if constraints_loc is not None")
         if constraints_scale is not None and mode != "tf":
-            raise ValueError("closed form hessian does not work if constraints_scale is not None")
+            raise ValueError("closed form jacobians do not work if constraints_scale is not None")
 
         if mode == "analytic":
             self.jac = self.analytic(
@@ -286,8 +286,8 @@ def _red(prev, cur):
             return tf.add(prev, cur)
 
         params = model_vars.params
-        p_shape_a = model_vars.a.shape[0]
-        p_shape_b = model_vars.b.shape[0]
+        p_shape_a = model_vars.a_var.shape[0]  # This has to be _var to work with constraints.
+        p_shape_b = model_vars.b_var.shape[0]  # This has to be _var to work with constraints.
 
         if iterator:
             J = op_utils.map_reduce(
@@ -374,8 +374,8 @@ def _red(prev, cur):
             return tf.add(prev, cur)
 
         params = model_vars.params
-        p_shape_a = model_vars.a.shape[0]
-        p_shape_b = model_vars.b.shape[0]
+        p_shape_a = model_vars.a_var.shape[0]  # This has to be _var to work with constraints.
+        p_shape_b = model_vars.b_var.shape[0]  # This has to be _var to work with constraints.
 
         if iterator == True and batch_model is None:
             J = op_utils.map_reduce(