Updated function names & added example rasterization notebook

Author: mtortora

data/hoomd_snapshot.gsd

@@ -1,13 +1,13 @@
 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, Normalize, NoNorm
 
+from scipy.ndimage import gaussian_filter
+
 
 def chromosome_viewer(chrom_lengths,
                       colors=None,
@@ -17,7 +17,7 @@ def chromosome_viewer(chrom_lengths,
     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
@@ -72,24 +72,31 @@ def rasterize(positions,
               box_size,
               resolution=100,
               length_unit=50,
-              gaussian_width=250):
+              gaussian_width=250,
+              normalize=False):
     """
     3D positions rasterizer for numerical microscopy analysis
     
     Parameters
-        ----------
+    ----------
         positions : Nx3 float array
-            List of X,Y,Z particle positions to be rasterize.
+            List of X,Y,Z particle positions to be rasterized (in model units).
             Assumes particle coordinates are wrapped in PBCs within the range [-box_size/2,+box_size/2]
         box_size : float
-            Linear dimension of periodic box
+            Linear dimension of periodic box (in model units)
         resolution : float
-			Linear dimension of output voxels (in nm)
+            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.
+            Width of Gaussian point-spread function to be used (in nm).
             If gaussian_width<=0, returns the raw (undiffracted) raster
+		normalize : bool
+            Set to True to scale maximum voxel intensity to 1
+            
+    Returns
+    -------
+        MxMxM raster array of 3D voxels
     """
 
     n_voxels = int(box_size / (resolution/length_unit))
@@ -105,36 +112,41 @@ def rasterize(positions,
     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
+        raster = gaussian_filter(raster/raster.max(), sigma=gaussian_width/resolution, mode='wrap')
 
-    return raster/raster.max()
+    if normalize:
+        raster /= raster.max()
+        
+    return raster
 
 
-def raster_map_2D(raster,
-                  cmap,
-                  vmin=None,
-                  vmax=None,
-                  mode='max',
-                  axis=2):
+def voxels_to_pixels_RGB(raster,
+                         cmap,
+                         vmin=None,
+                         vmax=None,
+                         mode='max',
+                         axis=2):
     """
-    RGB mapping of 2D raster projection
+    Project voxels onto 2D pixels and map values to RGB
     
     Parameters
-        ----------
-        raster : Mx3 float array
-            3D normalized raster to visualize.
+    ----------
+        raster : MxMxM float array
+            3D voxel array to visualize
         cmap : matplotlib colormap object or str
             Colormap to be used
         vmin, vmax : float or None
-            Color dynamic range.
+            RGB 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
+            
+    Returns
+    -------
+        MxMx3 image array of RGB pixels
     """
 
     if mode == 'max':
@@ -144,12 +156,12 @@ def raster_map_2D(raster,
     else:
         raise RuntimeError(f"Unsupported projection mode {mode}")
 
-    vmin = vmin if vmin else raster_2D.min()
-    vmax = vmax if vmax else raster_2D.max()
+    vmin = raster_2D.min() if vmin is None else vmin
+    vmax = raster_2D.max() if vmax is None else vmax
 
     norm = Normalize(vmin=vmin, vmax=vmax)
     cm = ScalarMappable(norm=norm, cmap=cmap)
 
     im = cm.to_rgba(raster_2D)
 
-    return im[:, :, :3]
+    return im[..., :3]

examples/rasterize.ipynb

@@ -2,6 +2,7 @@ numpy<1.24
 scipy>=0.16
 pandas>=0.19
 matplotlib>=2.0
+gsd
 rowan
 cooler
 joblib