Skip to content

Commit 7672d73

Browse files
rugelimogres
rugeli
and
mogres
authored
fix KeyError: 'weight_mode' when packing v2.1 recipes with under-specified gradients (#470)
* Add recipe and config to test gradient mixing in peroxisome data * normalize grad in early process * add tests --------- Co-authored-by: mogres <saurabh.mogre@alleninstitute.org>
1 parent 643fabb commit 7672d73

4 files changed

Lines changed: 243 additions & 13 deletions

File tree

‎cellpack/autopack/loaders/recipe_loader.py‎

Lines changed: 37 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
from cellpack.autopack.loaders.migrate_v1_to_v2 import convert as convert_v1_to_v2
1818
from cellpack.autopack.loaders.migrate_v2_to_v2_1 import convert as convert_v2_to_v2_1
1919
from cellpack.autopack.utils import deep_merge, expand_object_using_key
20+
from cellpack.autopack.validation.recipe_models import DEFAULT_GRADIENT_MODE_SETTINGS
2021
from cellpack.autopack.validation.recipe_validator import RecipeValidator
2122
from pydantic import ValidationError
2223

@@ -150,6 +151,33 @@ def _sanitize_format_version(recipe_data):
150151
format_version = recipe_data["format_version"]
151152
return format_version
152153

154+
@staticmethod
155+
def _normalize_gradients(gradients):
156+
"""
157+
Normalize gradients to the list-of-dicts format and fill in
158+
default settings for any missing keys.
159+
"""
160+
if not gradients:
161+
return gradients
162+
163+
# firebase recipes already store gradients as a list of dicts;
164+
# file recipes store them as a dict keyed by gradient name
165+
if isinstance(gradients, dict):
166+
gradient_items = [
167+
{**gradient_dict, "name": gradient_name}
168+
for gradient_name, gradient_dict in gradients.items()
169+
]
170+
else:
171+
gradient_items = gradients
172+
173+
normalized = []
174+
for gradient in gradient_items:
175+
# fill any missing top-level keys from the defaults without
176+
# overwriting values the recipe set
177+
filled = {**copy.deepcopy(DEFAULT_GRADIENT_MODE_SETTINGS), **gradient}
178+
normalized.append(filled)
179+
return normalized
180+
153181
def _migrate_version(self, old_recipe):
154182
converted = False
155183
if old_recipe["format_version"] == "1.0":
@@ -216,6 +244,11 @@ def _read(self, resolve_inheritance=True, use_docker=False):
216244
recipe_data["objects"] = RecipeLoader.resolve_inheritance(
217245
recipe_data["objects"]
218246
)
247+
if "gradients" in recipe_data:
248+
recipe_data["gradients"] = RecipeLoader._normalize_gradients(
249+
recipe_data["gradients"]
250+
)
251+
219252
# validate recipe after migration to v2.1 format but before transforming to class instances
220253
try:
221254
RecipeValidator.validate_recipe(recipe_data)
@@ -224,8 +257,10 @@ def _read(self, resolve_inheritance=True, use_docker=False):
224257
formatted_error = RecipeValidator.format_validation_error(e)
225258
raise ValueError(f"Recipe validation failed:\n{formatted_error}")
226259

227-
# keep a serializable copy before converting to class instances
228-
# this ensures the original data (human-readable) is available for download in the UI
260+
# keep a serializable copy after all dict-level normalization but before
261+
# converting to class instances. this is the human-readable source of
262+
# truth used for UI download / DB upload, and now matches both the packed
263+
# recipe and the database's list-of-dicts gradient format.
229264
self.serializable_recipe_data = copy.deepcopy(recipe_data)
230265

231266
if "objects" in recipe_data:
@@ -250,17 +285,6 @@ def _read(self, resolve_inheritance=True, use_docker=False):
250285
if "type" in obj and not INGREDIENT_TYPE.is_member(obj["type"]):
251286
raise TypeError(f"{obj['type']} is not an allowed type")
252287

253-
# handle gradients
254-
# gradients in firebase recipes are already stored as a list of dicts
255-
if "gradients" in recipe_data and not isinstance(
256-
recipe_data["gradients"], list
257-
):
258-
gradients = []
259-
for gradient_name, gradient_dict in recipe_data["gradients"].items():
260-
gradient_data = gradient_dict.copy()
261-
gradient_data["name"] = gradient_name
262-
gradients.append(gradient_data)
263-
recipe_data["gradients"] = gradients
264288
return recipe_data
265289

266290
def _load_json(self):

‎cellpack/tests/packing-configs/test_gradient_mixing.json‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
{
2+
"name": "test_gradient_mixing",
3+
"out": "out",
4+
"save_analyze_result": true,
5+
"show_progress_bar": true,
6+
"save_plot_figures": false,
7+
"load_from_grid_file": true,
8+
"spacing": 2,
9+
"image_export_options": {
10+
"hollow": false,
11+
"voxel_size": [
12+
1,
13+
1,
14+
1
15+
],
16+
"projection_axis": "z"
17+
},
18+
"open_results_in_browser": false,
19+
"upload_results": false,
20+
"number_of_packings": 1
21+
}

‎cellpack/tests/recipes/v2/test_peroxisome_combined_gradient.json‎

Lines changed: 153 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,153 @@
1+
{
2+
"version": "interpolated_743916",
3+
"format_version": "2.1",
4+
"name": "peroxisome",
5+
"bounding_box": [
6+
[
7+
35.325,
8+
35.325,
9+
7.125
10+
],
11+
[
12+
210.675,
13+
210.675,
14+
106.875
15+
]
16+
],
17+
"objects": {
18+
"membrane_mesh": {
19+
"type": "mesh",
20+
"color": [
21+
1,
22+
0,
23+
1
24+
],
25+
"representations": {
26+
"mesh": {
27+
"path": "s3://cellpack-analysis-data/structure_data/SLC25A17/meshes",
28+
"name": "mem_mesh_743916.obj",
29+
"format": "obj"
30+
}
31+
}
32+
},
33+
"nucleus_mesh": {
34+
"type": "mesh",
35+
"color": [
36+
0,
37+
1,
38+
1
39+
],
40+
"representations": {
41+
"mesh": {
42+
"path": "s3://cellpack-analysis-data/structure_data/SLC25A17/meshes",
43+
"name": "nuc_mesh_743916.obj",
44+
"format": "obj"
45+
}
46+
}
47+
},
48+
"peroxisome": {
49+
"jitter_attempts": 300,
50+
"type": "single_sphere",
51+
"color": [
52+
0.2,
53+
0.7,
54+
0.1
55+
],
56+
"radius": 0.2677546285806425,
57+
"packing_mode": "gradient",
58+
"gradient": [
59+
"uniform",
60+
"nucleus_gradient",
61+
"membrane_gradient",
62+
"apical_gradient"
63+
],
64+
"gradient_weights": {
65+
"uniform": 0.0,
66+
"nucleus_gradient": 0.3323145062439923,
67+
"membrane_gradient": 0.0,
68+
"apical_gradient": 0.6676854937560077
69+
}
70+
}
71+
},
72+
"composition": {
73+
"bounding_area": {
74+
"regions": {
75+
"interior": [
76+
"membrane"
77+
]
78+
}
79+
},
80+
"membrane": {
81+
"object": "membrane_mesh",
82+
"count": 1,
83+
"regions": {
84+
"interior": [
85+
"nucleus",
86+
{
87+
"object": "peroxisome",
88+
"count": 59
89+
}
90+
]
91+
}
92+
},
93+
"nucleus": {
94+
"object": "nucleus_mesh",
95+
"count": 1,
96+
"regions": {
97+
"interior": []
98+
}
99+
}
100+
},
101+
"randomness_seed": 743916,
102+
"grid_file_path": "s3://cellpack-analysis-data/structure_data/SLC25A17/grids/743916_grid.dat",
103+
"gradients": {
104+
"uniform": {
105+
"description": "uniform distribution throughout the cell volume",
106+
"mode": "uniform"
107+
},
108+
"nucleus_gradient": {
109+
"description": "gradient based on distance from the surface of the nucleus mesh",
110+
"mode": "surface",
111+
"mode_settings": {
112+
"object": "nucleus",
113+
"scale_to_next_surface": false
114+
},
115+
"weight_mode": "exponential",
116+
"weight_mode_settings": {
117+
"decay_length": 0.1
118+
}
119+
},
120+
"membrane_gradient": {
121+
"description": "gradient based on distance from the surface of the membrane mesh",
122+
"mode": "surface",
123+
"mode_settings": {
124+
"object": "membrane",
125+
"scale_to_next_surface": false
126+
},
127+
"weight_mode": "exponential",
128+
"weight_mode_settings": {
129+
"decay_length": 0.01
130+
}
131+
},
132+
"apical_gradient": {
133+
"description": "gradient based on distance from a plane",
134+
"mode": "vector",
135+
"mode_settings": {
136+
"direction": [
137+
0,
138+
0,
139+
1
140+
],
141+
"center": [
142+
123.0,
143+
123.0,
144+
106.875
145+
]
146+
},
147+
"weight_mode": "exponential",
148+
"weight_mode_settings": {
149+
"decay_length": 0.7
150+
}
151+
}
152+
}
153+
}

‎cellpack/tests/test_recipe_loader.py‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77

88
import pytest
99
from cellpack.autopack.loaders.recipe_loader import RecipeLoader
10+
from cellpack.autopack.validation.recipe_models import DEFAULT_GRADIENT_MODE_SETTINGS
1011

1112
test_objects = {
1213
"sphere_25": {
@@ -94,3 +95,34 @@ def test_resolve_objects():
9495
)
9596
def test_sanitize_format_version(expected_result, input_recipe_data):
9697
assert expected_result == RecipeLoader._sanitize_format_version(input_recipe_data)
98+
99+
100+
def test_normalize_gradients_fills_defaults_for_under_specified_v2_1_gradient():
101+
# a v2.1 recipe authored directly skips migration, so an under-specified
102+
# gradient reaches normalization missing the default keys
103+
gradients = {"my_gradient": {"mode": "surface"}}
104+
105+
normalized = RecipeLoader._normalize_gradients(gradients)
106+
107+
assert isinstance(normalized, list)
108+
gradient = normalized[0]
109+
assert gradient["name"] == "my_gradient"
110+
assert gradient["mode"] == "surface"
111+
for key in ("weight_mode", "pick_mode", "mode_settings", "weight_mode_settings"):
112+
assert key in gradient
113+
assert gradient["weight_mode"] == DEFAULT_GRADIENT_MODE_SETTINGS["weight_mode"]
114+
115+
116+
def test_normalize_gradients_preserves_authored_values_and_normalizes_list_input():
117+
# firebase recipes already store gradients as a list of dicts
118+
gradients = [{"name": "test", "mode": "surface", "weight_mode": "exponential"}]
119+
120+
normalized = RecipeLoader._normalize_gradients(gradients)
121+
122+
assert normalized[0]["weight_mode"] == "exponential" # not overwritten by default
123+
assert normalized[0]["pick_mode"] == DEFAULT_GRADIENT_MODE_SETTINGS["pick_mode"]
124+
125+
126+
def test_normalize_gradients_passes_through_empty():
127+
assert RecipeLoader._normalize_gradients(None) is None
128+
assert RecipeLoader._normalize_gradients({}) == {}

0 commit comments

Comments
 (0)