Merge pull request #54 from theislab/feature_batching

Feature-wise fitting termination

Author: davidsebfischer

batchglm/train/tf/base.py

@@ -287,6 +287,26 @@ def train(self, *args,
                 )
 
                 tf.logging.info("Step: %d\tloss: %f", train_step, global_loss)
+        elif convergence_criteria == "all_converged":
+            theta_current = self.session.run(self.model.model_vars.params)
+            while np.any(self.model.model_vars.converged == False):
+                theta_prev = theta_current
+                train_step, global_loss, _ = self.session.run(
+                    (self.model.global_step, loss, train_op),
+                    feed_dict=feed_dict
+                )
+                theta_current = self.session.run(self.model.model_vars.params)
+                theta_delta = np.abs(theta_prev - theta_current)
+                self.model.model_vars.converged = np.logical_or(  # Only update non-converged.
+                    self.model.model_vars.converged,
+                    np.max(theta_delta, axis=0) < stopping_criteria
+                )
+                tf.logging.info(
+                    "Step: %d\tloss: %f\t models converged %i",
+                    train_step,
+                    global_loss,
+                    np.sum(self.model.model_vars.converged).astype("int32")
+                )
         else:
             self._train_to_convergence(
                 loss=loss,

batchglm/train/tf/nb_glm/base.py

@@ -225,6 +225,8 @@ class ModelVars:
     a_var: tf.Variable
     b_var: tf.Variable
     params: tf.Variable
+    converged: np.ndarray
+
     """ Build tf.Variables to be optimzed and their constraints.
 
     a_var and b_var slices of the tf.Variable params which contains
@@ -309,8 +311,13 @@ def __init__(
             axis=0
         ), name="params")
 
-        a_var = params[0:init_a.shape[0]]
-        b_var = params[init_a.shape[0]:]
+        params_by_gene = [tf.expand_dims(params[:, i], axis=-1) for i in range(params.shape[1])]
+        a_by_gene = [x[0:init_a.shape[0],:] for x in params_by_gene]
+        b_by_gene = [x[init_a.shape[0]:, :] for x in params_by_gene]
+        a_var = tf.concat(a_by_gene, axis=1)
+        b_var = tf.concat(b_by_gene, axis=1)
+        #a_var = params[0:init_a.shape[0]]
+        #b_var = params[init_a.shape[0]:]
 
         # Define first layer of computation graph on identifiable variables
         # to yield dependent set of parameters of model for each location
@@ -334,3 +341,10 @@ def __init__(
         self.a_var = a_var
         self.b_var = b_var
         self.params = params
+        # Properties to follow gene-wise convergence.
+        self.params_by_gene = params_by_gene
+        self.a_by_gene = a_by_gene
+        self.b_by_gene = b_by_gene
+        self.converged = np.repeat(a=False, repeats=self.params.shape[1])  # Initialise to non-converged.
+        self.n_features = self.params.shape[1]
+