Merge pull request #4 from mtortora/coolify

Implemented the contact_maps.coolify method

Author: mtortora

polykit/analysis/contact_maps.py

@@ -43,12 +43,14 @@
 """
 
 import ctypes
+import cooler
 import multiprocessing as mp
 import random
 import warnings
 from contextlib import closing
 
 import numpy as np
+import pandas as pd
 
 from . import polymer_analyses
 
@@ -585,3 +587,49 @@ def monomerResolutionContactMapSubchains(
         uniqueContacts=True,
         nproc=n,
     )
+
+
+def coolify(contact_matrix,
+            cool_uri,
+            chrom_dict={},
+            binsize=2500,
+            chunksize=10000000):
+    """
+    Save a simulated contact map to .cool, and return corresponding Cooler
+    
+    Parameters
+        ----------
+        contact_matrix : NxN int array
+            Simulated contact map (in dense numpy.ndarray format)
+        cool_uri : str
+            Name of .cool file to be created (excluding extension)
+        chrom_dict : dict
+            Dictionary of chromosome lengths, in the form {'chr1':size1, ...}
+            If unspecified, assumes the map spans a single chromosome
+        binsize : int
+            Map resolution in bp (i.e., genomic size of each simulated monomer)
+        chunksize : int
+            Number of pixels handled by each worker process
+            
+        Returns
+        -------
+            Associated cooler.Cooler object
+    """
+    
+    nbins = contact_matrix.shape[0]
+
+    chrom_dict = chrom_dict if chrom_dict else {'chr_sim': binsize*nbins}
+    chrom_sizes = pd.Series(chrom_dict, name='length', dtype='int64')
+
+    assert binsize*nbins == sum(chrom_sizes), "Chromosome sizes do not match map dimensions"
+
+    bins = cooler.binnify(chrom_sizes, binsize)
+    bins['weight'] = np.ones(nbins) * np.sqrt(2/nbins)
+
+    pixels = cooler.create.ArrayLoader(bins, contact_matrix, chunksize=chunksize)
+    
+    cool_uri = f"{cool_uri}.{binsize}.cool"
+    cooler.create_cooler(cool_uri, bins, pixels, ordered=True)
+
+    return cooler.Cooler(cool_uri)
+    

requirements.txt

@@ -3,6 +3,7 @@ scipy>=0.16
 pandas>=0.19
 matplotlib>=2.0
 rowan
+cooler
 joblib
 pyside2
 ipython