fixed bug with computing condition number in dask in chunks of genes

Author: davidsebfischer

batchglm/train/numpy/base_glm/estimator.py

@@ -269,9 +269,10 @@ def iwls_step(
             # Have to use a workaround to solve problems in parallel in dask here. This workaround does
             # not work if there is only a single problem, ie. if the first dimension of a and b has length 1.
             if a.shape[0] != 1:
+                get_cond_number = lambda x: np.expand_dims(np.expand_dims(np.linalg.cond(x, p=None), axis=-1), axis=-1)
                 invertible = np.where(dask.array.map_blocks(
-                    np.linalg.cond, a, chunks=a.shape
-                ).compute() < 1 / sys.float_info.epsilon)[0]
+                    get_cond_number, a, chunks=a.shape
+                ).squeeze().compute() < 1 / sys.float_info.epsilon)[0]
                 delta_theta[:, idx_update[invertible]] = dask.array.map_blocks(
                     np.linalg.solve, a[invertible], b[invertible, :, None],
                     chunks=b[invertible, :, None].shape

batchglm/unit_test/test_acc_glm_all_numpy.py

@@ -36,7 +36,9 @@ def __init__(
                 design_scale_names=simulator.input_data.design_scale_names,
                 constraints_loc=simulator.input_data.constraints_loc,
                 constraints_scale=simulator.input_data.constraints_scale,
-                size_factors=simulator.input_data.size_factors
+                size_factors=simulator.input_data.size_factors,
+                chunk_size_cells=int(1e9),
+                chunk_size_genes=2
             )
         else:
             input_data = InputDataGLM(
@@ -47,7 +49,9 @@ def __init__(
                 design_scale_names=simulator.input_data.design_scale_names,
                 constraints_loc=simulator.input_data.constraints_loc,
                 constraints_scale=simulator.input_data.constraints_scale,
-                size_factors=simulator.input_data.size_factors
+                size_factors=simulator.input_data.size_factors,
+                chunk_size_cells=int(1e9),
+                chunk_size_genes=2
             )
 
         self.estimator = Estimator(