fix KeyError: 'weight_mode' when packing v2.1 recipes with under-specified gradients

cellpack/autopack/loaders/recipe_loader.py

@@ -17,6 +17,7 @@
 from cellpack.autopack.loaders.migrate_v1_to_v2 import convert as convert_v1_to_v2
 from cellpack.autopack.loaders.migrate_v2_to_v2_1 import convert as convert_v2_to_v2_1
 from cellpack.autopack.utils import deep_merge, expand_object_using_key
+from cellpack.autopack.validation.recipe_models import DEFAULT_GRADIENT_MODE_SETTINGS
 from cellpack.autopack.validation.recipe_validator import RecipeValidator
 from pydantic import ValidationError
 
@@ -150,6 +151,33 @@ def _sanitize_format_version(recipe_data):
             format_version = recipe_data["format_version"]
         return format_version
 
+    @staticmethod
+    def _normalize_gradients(gradients):
+        """
+        Normalize gradients to the list-of-dicts format and fill in
+        default settings for any missing keys.
+        """
+        if not gradients:
+            return gradients
+
+        # firebase recipes already store gradients as a list of dicts;
+        # file recipes store them as a dict keyed by gradient name
+        if isinstance(gradients, dict):
+            gradient_items = [
+                {**gradient_dict, "name": gradient_name}
+                for gradient_name, gradient_dict in gradients.items()
+            ]
+        else:
+            gradient_items = gradients
+
+        normalized = []
+        for gradient in gradient_items:
+            # fill any missing top-level keys from the defaults without
+            # overwriting values the recipe set
+            filled = {**copy.deepcopy(DEFAULT_GRADIENT_MODE_SETTINGS), **gradient}
+            normalized.append(filled)
+        return normalized
+
     def _migrate_version(self, old_recipe):
         converted = False
         if old_recipe["format_version"] == "1.0":
@@ -216,6 +244,11 @@ def _read(self, resolve_inheritance=True, use_docker=False):
                 recipe_data["objects"] = RecipeLoader.resolve_inheritance(
                     recipe_data["objects"]
                 )
+        if "gradients" in recipe_data:
+            recipe_data["gradients"] = RecipeLoader._normalize_gradients(
+                recipe_data["gradients"]
+            )
+
         # validate recipe after migration to v2.1 format but before transforming to class instances
         try:
             RecipeValidator.validate_recipe(recipe_data)
@@ -224,8 +257,10 @@ def _read(self, resolve_inheritance=True, use_docker=False):
             formatted_error = RecipeValidator.format_validation_error(e)
             raise ValueError(f"Recipe validation failed:\n{formatted_error}")
 
-        # keep a serializable copy before converting to class instances
-        # this ensures the original data (human-readable) is available for download in the UI
+        # keep a serializable copy after all dict-level normalization but before
+        # converting to class instances. this is the human-readable source of
+        # truth used for UI download / DB upload, and now matches both the packed
+        # recipe and the database's list-of-dicts gradient format.
         self.serializable_recipe_data = copy.deepcopy(recipe_data)
 
         if "objects" in recipe_data:
@@ -250,17 +285,6 @@ def _read(self, resolve_inheritance=True, use_docker=False):
                 if "type" in obj and not INGREDIENT_TYPE.is_member(obj["type"]):
                     raise TypeError(f"{obj['type']} is not an allowed type")
 
-        # handle gradients
-        # gradients in firebase recipes are already stored as a list of dicts
-        if "gradients" in recipe_data and not isinstance(
-            recipe_data["gradients"], list
-        ):
-            gradients = []
-            for gradient_name, gradient_dict in recipe_data["gradients"].items():
-                gradient_data = gradient_dict.copy()
-                gradient_data["name"] = gradient_name
-                gradients.append(gradient_data)
-            recipe_data["gradients"] = gradients
         return recipe_data
 
     def _load_json(self):

cellpack/tests/packing-configs/test_gradient_mixing.json

@@ -0,0 +1,21 @@
+{
+    "name": "test_gradient_mixing",
+    "out": "out",
+    "save_analyze_result": true,
+    "show_progress_bar": true,
+    "save_plot_figures": false,
+    "load_from_grid_file": true,
+    "spacing": 2,
+    "image_export_options": {
+        "hollow": false,
+        "voxel_size": [
+            1,
+            1,
+            1
+        ],
+        "projection_axis": "z"
+    },
+    "open_results_in_browser": false,
+    "upload_results": false,
+    "number_of_packings": 1
+}

cellpack/tests/recipes/v2/test_peroxisome_combined_gradient.json

@@ -0,0 +1,153 @@
+{
+    "version": "interpolated_743916",
+    "format_version": "2.1",
+    "name": "peroxisome",
+    "bounding_box": [
+        [
+            35.325,
+            35.325,
+            7.125
+        ],
+        [
+            210.675,
+            210.675,
+            106.875
+        ]
+    ],
+    "objects": {
+        "membrane_mesh": {
+            "type": "mesh",
+            "color": [
+                1,
+                0,
+                1
+            ],
+            "representations": {
+                "mesh": {
+                    "path": "s3://cellpack-analysis-data/structure_data/SLC25A17/meshes",
+                    "name": "mem_mesh_743916.obj",
+                    "format": "obj"
+                }
+            }
+        },
+        "nucleus_mesh": {
+            "type": "mesh",
+            "color": [
+                0,
+                1,
+                1
+            ],
+            "representations": {
+                "mesh": {
+                    "path": "s3://cellpack-analysis-data/structure_data/SLC25A17/meshes",
+                    "name": "nuc_mesh_743916.obj",
+                    "format": "obj"
+                }
+            }
+        },
+        "peroxisome": {
+            "jitter_attempts": 300,
+            "type": "single_sphere",
+            "color": [
+                0.2,
+                0.7,
+                0.1
+            ],
+            "radius": 0.2677546285806425,
+            "packing_mode": "gradient",
+            "gradient": [
+                "uniform",
+                "nucleus_gradient",
+                "membrane_gradient",
+                "apical_gradient"
+            ],
+            "gradient_weights": {
+                "uniform": 0.0,
+                "nucleus_gradient": 0.3323145062439923,
+                "membrane_gradient": 0.0,
+                "apical_gradient": 0.6676854937560077
+            }
+        }
+    },
+    "composition": {
+        "bounding_area": {
+            "regions": {
+                "interior": [
+                    "membrane"
+                ]
+            }
+        },
+        "membrane": {
+            "object": "membrane_mesh",
+            "count": 1,
+            "regions": {
+                "interior": [
+                    "nucleus",
+                    {
+                        "object": "peroxisome",
+                        "count": 59
+                    }
+                ]
+            }
+        },
+        "nucleus": {
+            "object": "nucleus_mesh",
+            "count": 1,
+            "regions": {
+                "interior": []
+            }
+        }
+    },
+    "randomness_seed": 743916,
+    "grid_file_path": "s3://cellpack-analysis-data/structure_data/SLC25A17/grids/743916_grid.dat",
+    "gradients": {
+        "uniform": {
+            "description": "uniform distribution throughout the cell volume",
+            "mode": "uniform"
+        },
+        "nucleus_gradient": {
+            "description": "gradient based on distance from the surface of the nucleus mesh",
+            "mode": "surface",
+            "mode_settings": {
+                "object": "nucleus",
+                "scale_to_next_surface": false
+            },
+            "weight_mode": "exponential",
+            "weight_mode_settings": {
+                "decay_length": 0.1
+            }
+        },
+        "membrane_gradient": {
+            "description": "gradient based on distance from the surface of the membrane mesh",
+            "mode": "surface",
+            "mode_settings": {
+                "object": "membrane",
+                "scale_to_next_surface": false
+            },
+            "weight_mode": "exponential",
+            "weight_mode_settings": {
+                "decay_length": 0.01
+            }
+        },
+        "apical_gradient": {
+            "description": "gradient based on distance from a plane",
+            "mode": "vector",
+            "mode_settings": {
+                "direction": [
+                    0,
+                    0,
+                    1
+                ],
+                "center": [
+                    123.0,
+                    123.0,
+                    106.875
+                ]
+            },
+            "weight_mode": "exponential",
+            "weight_mode_settings": {
+                "decay_length": 0.7
+            }
+        }
+    }
+}

cellpack/tests/test_recipe_loader.py

@@ -7,6 +7,7 @@
 
 import pytest
 from cellpack.autopack.loaders.recipe_loader import RecipeLoader
+from cellpack.autopack.validation.recipe_models import DEFAULT_GRADIENT_MODE_SETTINGS
 
 test_objects = {
     "sphere_25": {
@@ -94,3 +95,34 @@ def test_resolve_objects():
 )
 def test_sanitize_format_version(expected_result, input_recipe_data):
     assert expected_result == RecipeLoader._sanitize_format_version(input_recipe_data)
+
+
+def test_normalize_gradients_fills_defaults_for_under_specified_v2_1_gradient():
+    # a v2.1 recipe authored directly skips migration, so an under-specified
+    # gradient reaches normalization missing the default keys
+    gradients = {"my_gradient": {"mode": "surface"}}
+
+    normalized = RecipeLoader._normalize_gradients(gradients)
+
+    assert isinstance(normalized, list)
+    gradient = normalized[0]
+    assert gradient["name"] == "my_gradient"
+    assert gradient["mode"] == "surface"
+    for key in ("weight_mode", "pick_mode", "mode_settings", "weight_mode_settings"):
+        assert key in gradient
+    assert gradient["weight_mode"] == DEFAULT_GRADIENT_MODE_SETTINGS["weight_mode"]
+
+
+def test_normalize_gradients_preserves_authored_values_and_normalizes_list_input():
+    # firebase recipes already store gradients as a list of dicts
+    gradients = [{"name": "test", "mode": "surface", "weight_mode": "exponential"}]
+
+    normalized = RecipeLoader._normalize_gradients(gradients)
+
+    assert normalized[0]["weight_mode"] == "exponential"  # not overwritten by default
+    assert normalized[0]["pick_mode"] == DEFAULT_GRADIENT_MODE_SETTINGS["pick_mode"]
+
+
+def test_normalize_gradients_passes_through_empty():
+    assert RecipeLoader._normalize_gradients(None) is None
+    assert RecipeLoader._normalize_gradients({}) == {}