Save gradient data as image

cellpack/autopack/Analysis.py

@@ -2169,6 +2169,7 @@ def pack_one_seed(
         image_export_options=None,
         show_grid=False,
         plot_figures=False,
+        save_gradient_data_as_image=False,
     ):
         """
         Packs one seed of a recipe and returns the recipe object
@@ -2309,13 +2310,26 @@ def pack_one_seed(
                 plt.close()  # closes the current figure
 
         if image_export_options is not None:
-            image_writer = ImageWriter(
+            env_image_writer = ImageWriter(
                 env=self.env,
                 name=seed_basename,
                 output_path=self.figures_path,
                 **image_export_options,
             )
-            image_writer.export_image()
+            env_image_writer = self.env.create_voxelization(env_image_writer)
+            env_image_writer.export_image()
+
+        if save_gradient_data_as_image:
+            gradient_data_figure_path = self.figures_path / "gradient_data_figures"
+            gradient_data_figure_path.mkdir(exist_ok=True)
+            for _, gradient in self.env.gradients.items():
+                grid_image_writer = ImageWriter(
+                    env=self.env,
+                    name=f"{seed_basename}_grid_data",
+                    output_path=gradient_data_figure_path,
+                )
+                grid_image_writer = gradient.create_voxelization(grid_image_writer)
+                grid_image_writer.export_image()
 
         return (
             center_distance_dict,
@@ -2338,6 +2352,7 @@ def doloop(
         recipe_version="1.0.0",
         image_export_options=None,
         parallel=False,
+        save_gradient_data_as_image=False,
     ):
         """
         Runs multiple packings of the same recipe in a loop. This workflow
@@ -2437,6 +2452,7 @@ def doloop(
                             ingredient_key_dict=ingredient_key_dict,
                             get_distance_distribution=get_distance_distribution,
                             image_export_options=image_export_options,
+                            save_gradient_data_as_image=save_gradient_data_as_image,
                         )
                     )
                 for future in concurrent.futures.as_completed(futures):
@@ -2478,6 +2494,7 @@ def doloop(
                     image_export_options=image_export_options,
                     show_grid=show_grid,
                     plot_figures=plot_figures,
+                    save_gradient_data_as_image=save_gradient_data_as_image,
                 )
 
         self.writeJSON(center_distance_file, center_distance_dict)

cellpack/autopack/Environment.py

@@ -3129,50 +3129,47 @@ def applyStep(self, step):
         self.traj.applyState_primitive_name(self, step)
         # ho can we apply to parent instance the rotatiotn?
 
-    def create_voxelization(self, image_data, image_size, voxel_size, hollow=False):
+    def create_voxelization(self, image_writer):
         """
         Update the image data for all molecules in the recipe by creating voxelized
         representations.
 
         Parameters
         ----------
-        image_data: numpy.ndarray
-            The image data to update.
-        image_size: list
-            The size of the image data.
-        voxel_size: float
-            The size of a voxel in the image data.
-        hollow: bool
-            If True, the voxelization will be hollow.
+        image_writer: ImageWriter
+            The image writer to use for writing the voxelized representations.
 
         Returns
         ----------
         image_data: numpy.ndarray
             The updated image data.
         """
-        channel_colors = []
+        channel_colors = {}
 
         for obj in self.packed_objects.get_all():
             mesh_store = None
-            if obj.name not in image_data:
-                image_data[obj.name] = numpy.zeros(image_size, dtype=numpy.uint8)
+            if obj.name not in image_writer.image_data:
+                image_writer.image_data[obj.name] = numpy.zeros(
+                    image_writer.image_size, dtype=numpy.uint8
+                )
                 if obj.color is not None:
                     color = obj.color
                     if all([x <= 1 for x in obj.color]):
                         color = [int(col * 255) for col in obj.color]
-                    channel_colors.append(color)
+                    channel_colors[obj.name] = color
             if obj.is_compartment:
                 mesh_store = self.mesh_store
 
-            image_data[obj.name] = obj.ingredient.create_voxelization(
-                image_data=image_data[obj.name],
+            image_writer.image_data[obj.name] = obj.ingredient.create_voxelization(
+                image_data=image_writer.image_data[obj.name],
                 bounding_box=self.boundingBox,
-                voxel_size=voxel_size,
-                image_size=image_size,
+                voxel_size=image_writer.voxel_size,
+                image_size=image_writer.image_size,
                 position=obj.position,
                 rotation=obj.rotation,
-                hollow=hollow,
+                hollow=image_writer.hollow,
                 mesh_store=mesh_store,
             )
+        image_writer.channel_colors = channel_colors
 
-        return image_data, channel_colors
+        return image_writer

cellpack/autopack/Gradient.py

@@ -363,3 +363,26 @@ def getSubWeighted(self, listPts):
             rnd -= w
             if rnd < 0:
                 return listPts[i]
+
+    def create_voxelization(self, image_writer):
+        """
+        Creates a voxelized representation of the gradient distances and weights
+        """
+        image_writer.channel_colors = {"distances": [255, 0, 0], "weight": [0, 255, 0]}
+
+        for channel in ["distances", "weight"]:
+            channel_values = getattr(self, channel, None)
+            image_writer.image_data[channel] = numpy.zeros(
+                image_writer.image_size, dtype=numpy.uint8
+            )
+            if channel_values is None:
+                continue
+            normalized_values = self.get_normalized_values(channel_values)
+            reshaped_values = numpy.reshape(
+                normalized_values, image_writer.image_size, order="F"
+            )
+            image_writer.image_data[channel] = numpy.where(
+                reshaped_values > 0, numpy.rint(reshaped_values * 255), 0
+            )
+
+        return image_writer

cellpack/autopack/loaders/config_loader.py

@@ -40,6 +40,7 @@ class ConfigLoader(object):
         "randomness_seed": None,
         "save_analyze_result": False,
         "save_converted_recipe": False,
+        "save_gradient_data_as_image": False,
         "show_grid_plot": False,
         "show_sphere_trees": False,
         "show_progress_bar": False,

cellpack/autopack/writers/ImageWriter.py

@@ -15,28 +15,29 @@ def __init__(
         name=None,
         output_path=None,
         voxel_size=None,
-        num_voxels=None,
         hollow=False,
         convolution_options=None,
         projection_axis="z",
     ):
         self.env = env
 
-        self.name = "default"
         if name is not None:
             self.name = name
+        else:
+            self.name = "default"
 
-        self.output_path = Path(self.env.out_folder)
         if output_path is not None:
             self.output_path = Path(output_path)
+        else:
+            self.output_path = Path(self.env.out_folder)
 
-        self.voxel_size = numpy.array([1, 1, 1])  # units of grid points per voxel
         if voxel_size is not None:
             self.voxel_size = numpy.array(voxel_size)
-        elif num_voxels is not None:
+        else:
+            grid_spacing = self.env.grid.gridSpacing
             self.voxel_size = (
-                self.env.boundingBox[1] - self.env.boundingBox[0]
-            ) / numpy.array(num_voxels)
+                numpy.array([1, 1, 1]) * grid_spacing
+            )  # units of grid points per voxel
 
         self.hollow = hollow
 
@@ -47,6 +48,7 @@ def __init__(
             ).astype(int)
         )
         self.image_data = {}
+        self.channel_colors = {}
 
         self.convolution_options = convolution_options
 
@@ -174,22 +176,22 @@ def convolve_image(self, image, psf="gaussian", psf_parameters=None):
             conv_img[channel] = self.convolve_channel(image[channel], psf)
         return conv_img
 
-    def create_voxelization(self):
+    def concatenate_image_data(self):
         """
         Creates a voxelized representation of the current scene
         """
-
-        self.image_data, channel_colors = self.env.create_voxelization(
-            self.image_data, self.image_size, self.voxel_size, self.hollow
-        )
-
         concatenated_image = numpy.zeros(
             (len(self.image_data), *self.image_size), dtype=numpy.uint8
         )
+
         channel_names = []
+        channel_colors = []
         for ct, (channel_name, channel_image) in enumerate(self.image_data.items()):
             concatenated_image[ct] = channel_image
             channel_names.append(channel_name)
+            channel_colors.append(
+                self.channel_colors.get(channel_name, [255, 255, 255])
+            )
 
         if self.convolution_options is not None:
             concatenated_image = self.convolve_image(
@@ -207,8 +209,12 @@ def export_image(self):
         Saves the results as a tiff file
         """
         print(f"Exporting image to {self.output_path}")
-        concatenated_image, channel_names, channel_colors = self.create_voxelization()
-        if len(channel_names) != 0:
+        (
+            concatenated_image,
+            channel_names,
+            channel_colors,
+        ) = self.concatenate_image_data()
+        if len(self.channel_colors) != 0:
             filepath = self.output_path / f"voxelized_image_{self.name}.ome.tiff"
             OmeTiffWriter.save(
                 concatenated_image,

cellpack/bin/pack.py

@@ -70,6 +70,9 @@ def pack(recipe, config_path=None, analysis_config_path=None):
             recipe_version=recipe_data["version"],
             image_export_options=packing_config_data.get("image_export_options"),
             parallel=packing_config_data.get("parallel", False),
+            save_gradient_data_as_image=packing_config_data.get(
+                "save_gradient_data_as_image", False
+            ),
         )
         if analysis_config_path is not None:
             analyze.run_analysis_workflow(

cellpack/tests/packing-configs/test_gradient_data_image.json

@@ -0,0 +1,14 @@
+{
+    "name": "test_gradient_data_image",
+    "save_analyze_result": true,
+    "show_grid_plot": true,
+    "spacing": 0.25,
+    "show_sphere_trees": false,
+    "load_from_grid_file": true,
+    "save_gradient_data_as_image": true,
+    "image_export_options": {
+        "hollow": false,
+        "voxel_size": [1,1,1],
+        "projection_axis": "z"
+    }
+}