Implemented rasterization utilities

Author: mtortora

polykit/renderers/viewers.py

@@ -1,25 +1,28 @@
 import numpy as np
 import matplotlib.pyplot as plt
 
+from scipy.ndimage import gaussian_filter
+
+from matplotlib.cm import ScalarMappable
 from matplotlib.ticker import AutoLocator
 from matplotlib.colorbar import ColorbarBase
-from matplotlib.colors import ListedColormap, NoNorm
+from matplotlib.colors import ListedColormap, Normalize, NoNorm
 
 
-def chromosome_viewer(chrom_lengths, 
+def chromosome_viewer(chrom_lengths,
                       colors=None,
                       height=0.5,
                       width=10.):
     """
-    Simple matplotlib-based visualization of individual chromosomes 
+    Simple matplotlib-based visualization of individual chromosomes
     
     Parameters
         ----------
         chrom_lengths : Cx1 int array
             List of the sizes of each chromosome
         colors : Nx4 float array or None
-            Colors to be assigned to each monomer. 
-            If omitted, random colors will be attributed to each chromosome 
+            Colors to be assigned to each monomer.
+            If omitted, random colors will be attributed to each chromosome
         height : float
             Per-chromosome height of the output figure (in inches)
         width : float
@@ -63,3 +66,90 @@ def chromosome_viewer(chrom_lengths,
         cb.update_ticks()
 
     plt.show()
+
+
+def rasterize(positions,
+              box_size,
+              resolution=100,
+              length_unit=50,
+              gaussian_width=250):
+    """
+    3D positions rasterizer for numerical microscopy analysis
+    
+    Parameters
+        ----------
+        positions : Nx3 float array
+            List of X,Y,Z particle positions to be rasterize.
+            Assumes particle coordinates are wrapped in PBCs within the range [-box_size/2,+box_size/2]
+        box_size : float
+            Linear dimension of periodic box
+        resolution : float
+			Linear dimension of output voxels (in nm)
+        length_unit : float
+            Model unit of length (in nm)
+        gaussian_width : float
+            Width of Gaussian point-spread function to be used.
+            If gaussian_width<=0, returns the raw (undiffracted) raster
+    """
+    
+    n_voxels = int(box_size / (resolution/length_unit))
+    bins = np.linspace(-box_size/2., box_size/2., num=n_voxels+1)
+																			
+    digitized_x = np.digitize(positions[:,0], bins) - 1
+    digitized_y = np.digitize(positions[:,1], bins) - 1
+    digitized_z = np.digitize(positions[:,2], bins) - 1
+					
+    digitized = digitized_x + digitized_y*n_voxels + digitized_z*n_voxels**2
+    bincounts = np.bincount(digitized, minlength=n_voxels**3)
+					
+    raster = bincounts.reshape((n_voxels,n_voxels,n_voxels))
+	
+    if gaussian_width > 0:
+        smoothed_raster = gaussian_filter(raster/raster.max(), sigma=gaussian_width/resolution, mode='wrap')
+
+        return smoothed_raster
+        
+    return raster/raster.max()
+
+
+def raster_map_2D(raster,
+                  cmap,
+                  vmin=None,
+                  vmax=None,
+                  mode='max',
+                  axis=2):
+    """
+    RGB mapping of 2D raster projection
+    
+    Parameters
+        ----------
+        raster : Mx3 float array
+            3D normalized raster to visualize.
+        cmap : matplotlib colormap object or str
+            Colormap to be used
+        vmin, vmax : float or None
+            Color dynamic range.
+            If set to None, use full data range
+        mode : str
+            Projection mode to be used.
+            Set to either 'max' for maximum intensity or 'sum' for summed intensity projection
+        axis : int
+            Index of projection axis
+    """
+
+    if mode == 'max':
+        raster_2D = raster.max(axis=axis)
+    elif mode == 'sum':
+        raster_2D = raster.sum(axis=axis)
+    else:
+        raise RuntimeError(f"Unsupported projection mode {mode}")
+		
+    vmin = vmin if vmin else raster_2D.min()
+    vmax = vmax if vmax else raster_2D.max()
+
+    norm = Normalize(vmin=vmin, vmax=vmax)
+    cm = ScalarMappable(norm=norm, cmap=cmap)
+
+    im = cm.to_rgba(raster_2D)
+	
+    return im[:, :, :3]