Update cool_ds.py

Author: lhqing

ALLCools/mcds/cool_ds.py

@@ -208,27 +208,27 @@ def get_cooler(
         cooler_kwargs :
             Additional arguments to pass to create_cooler
         """
-        import subprocess
-
-        from cooler import create_cooler
-        from scipy.sparse import coo_matrix
-
-        def _chrom_iterator(
-            _samples,
-            _value_type,
-            _chrom_offset,
-            _da_name,
-            add_trans=False,
-        ):
-            """Iterate through the raw matrices and chromosomes of cells."""
-            chrom_sizes = self.chrom_sizes
-
-            def _iter_1d(_chrom1, _chrom2):
-                print(f"Saving {_chrom1} x {_chrom2 if _chrom2 is not None else _chrom1}...")
-                chrom_ds = self.fetch(chrom=_chrom1, chrom2=_chrom2)
-
-                # get chrom 2D np.array
-                with dask.config.set(**{"array.slicing.split_large_chunks": False}):
+        with dask.config.set(scheduler="sync"):
+            import subprocess
+
+            from cooler import create_cooler
+            from scipy.sparse import coo_matrix
+
+            def _chrom_iterator(
+                _samples,
+                _value_type,
+                _chrom_offset,
+                _da_name,
+                add_trans=False,
+            ):
+                """Iterate through the raw matrices and chromosomes of cells."""
+                chrom_sizes = self.chrom_sizes
+
+                def _iter_1d(_chrom1, _chrom2):
+                    print(f"Saving {_chrom1} x {_chrom2 if _chrom2 is not None else _chrom1}...")
+                    chrom_ds = self.fetch(chrom=_chrom1, chrom2=_chrom2)
+
+                    # get chrom 2D np.array
                     matrix = chrom_ds.matrix(
                         samples=_samples,
                         value_type=_value_type,
@@ -244,32 +244,31 @@ def _iter_1d(_chrom1, _chrom2):
                     if len(matrix.shape) > 2:
                         matrix = matrix.squeeze()
 
-                # to coo then to pixel
-                matrix = coo_matrix(matrix)
-                _pixel_df = pd.DataFrame({"bin1_id": matrix.row, "bin2_id": matrix.col, "count": matrix.data})
-
-                # add chrom offset
-                if _chrom2 is None:
-                    # both row and col are chrom1
-                    _pixel_df.iloc[:, :2] += _chrom_offset[_chrom1]
+                    # to coo then to pixel
+                    matrix = coo_matrix(matrix)
+                    _pixel_df = pd.DataFrame({"bin1_id": matrix.row, "bin2_id": matrix.col, "count": matrix.data})
+
+                    # add chrom offset
+                    if _chrom2 is None:
+                        # both row and col are chrom1
+                        _pixel_df.iloc[:, :2] += _chrom_offset[_chrom1]
+                    else:
+                        # row is chrom1, add chrom1 offset
+                        _pixel_df.iloc[:, 0] += _chrom_offset[_chrom1]
+                        # col is chrom2, add chrom2 offset
+                        _pixel_df.iloc[:, 1] += _chrom_offset[_chrom2]
+                    return _pixel_df
+
+                if add_trans:
+                    raise NotImplementedError
                 else:
-                    # row is chrom1, add chrom1 offset
-                    _pixel_df.iloc[:, 0] += _chrom_offset[_chrom1]
-                    # col is chrom2, add chrom2 offset
-                    _pixel_df.iloc[:, 1] += _chrom_offset[_chrom2]
-                return _pixel_df
+                    for chrom in chrom_sizes.keys():
+                        pixel_df = _iter_1d(chrom, None)
+                        yield pixel_df
 
-            if add_trans:
-                raise NotImplementedError
-            else:
-                for chrom in chrom_sizes.keys():
-                    pixel_df = _iter_1d(chrom, None)
-                    yield pixel_df
+            bins_df = binnify(self.chrom_sizes, binsize=self.bin_size)
+            chrom_offset = _get_chrom_offsets(bins_df)
 
-        bins_df = binnify(self.chrom_sizes, binsize=self.bin_size)
-        chrom_offset = _get_chrom_offsets(bins_df)
-
-        with dask.config.set(scheduler="sync"):
             create_cooler(
                 cool_uri=f"{output_prefix}.cool",
                 bins=bins_df,
@@ -285,10 +284,10 @@ def _iter_1d(_chrom1, _chrom2):
                 **(cooler_kwargs or {}),
             )
 
-        if zoomify:
-            subprocess.run(["cooler", "zoomify", f"{output_prefix}.cool", "-p", str(zoomify_cpu)], check=True)
-            # delete the original cooler file
-            subprocess.run(["rm", f"{output_prefix}.cool"], check=True)
+            if zoomify:
+                subprocess.run(["cooler", "zoomify", f"{output_prefix}.cool", "-p", str(zoomify_cpu)], check=True)
+                # delete the original cooler file
+                subprocess.run(["rm", f"{output_prefix}.cool"], check=True)
 
         return
 
@@ -366,7 +365,7 @@ def matrix(
         -------
         np.ndarray
         """
-        with dask.config.set(scheduler="sync"):
+        with dask.config.set(scheduler="synchronous"):
             sel_dict = {}
             if samples is not None:
                 sel_dict[self.sample_dim] = samples