Type decoding tuning schema

Author: Donglai Wei

connectomics/config/schema/__init__.py

@@ -36,7 +36,11 @@
     DecodeBinaryContourDistanceWatershedConfig,
     DecodeModeConfig,
     DecodingConfig,
+    DecodingTuningConfig,
+    DecodingTuningParameterSpaceConfig,
     PostprocessingConfig,
+    TuningFunctionSpaceConfig,
+    TuningParameterConfig,
 )
 from .evaluation import EvaluationConfig
 from .inference import (
@@ -133,12 +137,16 @@
     "SavePredictionConfig",
     "InferenceMemoryCleanupConfig",
     "DecodingConfig",
+    "DecodingTuningConfig",
+    "DecodingTuningParameterSpaceConfig",
     "PostprocessingConfig",
     "BinaryPostprocessingConfig",
     "ConnectedComponentsConfig",
     "EvaluationConfig",
     "DecodeModeConfig",
     "DecodeBinaryContourDistanceWatershedConfig",
+    "TuningFunctionSpaceConfig",
+    "TuningParameterConfig",
     # Test configuration
     "TestConfig",
     # Tuning configuration

connectomics/config/schema/decoding.py

@@ -56,6 +56,48 @@ class DecodeModeConfig:
     kwargs: Dict[str, Any] = field(default_factory=dict)
 
 
+@dataclass
+class TuningParameterConfig:
+    """Single tunable decoding/postprocessing parameter."""
+
+    type: str = "float"
+    range: List[Any] = field(default_factory=list)
+    choices: List[Any] = field(default_factory=list)
+    step: Optional[float] = None
+    log: bool = False
+    param_group: Optional[str] = None
+    tuple_index: Optional[int] = None
+    description: Optional[str] = None
+
+
+@dataclass
+class TuningFunctionSpaceConfig:
+    """Search space for one decoding or postprocessing function."""
+
+    enabled: bool = False
+    function_name: str = ""
+    defaults: Dict[str, Any] = field(default_factory=dict)
+    parameters: Dict[str, TuningParameterConfig] = field(default_factory=dict)
+
+
+@dataclass
+class DecodingTuningParameterSpaceConfig:
+    """Parameter spaces attached to a decoding pipeline."""
+
+    decoding: TuningFunctionSpaceConfig = field(default_factory=TuningFunctionSpaceConfig)
+    postprocessing: TuningFunctionSpaceConfig = field(default_factory=TuningFunctionSpaceConfig)
+
+
+@dataclass
+class DecodingTuningConfig:
+    """Structured tuning metadata for decoded-output workflows."""
+
+    enabled: bool = False
+    parameter_space: DecodingTuningParameterSpaceConfig = field(
+        default_factory=DecodingTuningParameterSpaceConfig
+    )
+
+
 @dataclass
 class DecodingConfig:
     """Decoded-output orchestration configuration."""
@@ -64,4 +106,4 @@ class DecodingConfig:
     postprocessing: PostprocessingConfig = field(default_factory=PostprocessingConfig)
     output_path: str = ""
     input_prediction_path: str = ""
-    tuning: Optional[Dict[str, Any]] = None
+    tuning: Optional[DecodingTuningConfig] = None

tests/unit/test_hydra_config.py

@@ -875,6 +875,46 @@ def test_top_level_decoding_template_list_ref(tmp_path):
     assert cfg.decoding.steps[0].kwargs["min_instance_size"] == 7
 
 
+def test_decoding_tuning_uses_structured_schema(tmp_path):
+    config_yaml = tmp_path / "decoding_tuning.yaml"
+    config_yaml.write_text("""
+decoding:
+  tuning:
+    enabled: true
+    parameter_space:
+      decoding:
+        enabled: true
+        function_name: decode_waterz
+        parameters:
+          thresholds:
+            range: [0.1, 0.9]
+            step: 0.1
+""".strip())
+
+    cfg = load_config(config_yaml)
+
+    assert cfg.decoding.tuning is not None
+    assert cfg.decoding.tuning.enabled is True
+    decoding_space = cfg.decoding.tuning.parameter_space.decoding
+    assert decoding_space.function_name == "decode_waterz"
+    threshold_param = decoding_space.parameters["thresholds"]
+    assert threshold_param.range == [0.1, 0.9]
+    assert threshold_param.step == 0.1
+
+
+def test_decoding_tuning_rejects_unknown_nested_keys(tmp_path):
+    config_yaml = tmp_path / "bad_decoding_tuning.yaml"
+    config_yaml.write_text("""
+decoding:
+  tuning:
+    enabled: true
+    unknown_key: true
+""".strip())
+
+    with pytest.raises(Exception, match="unknown_key"):
+        load_config(config_yaml)
+
+
 def test_loss_profile_positional_overrides(tmp_path):
     """Loss profile + overrides dict patches individual list entries by index."""
     base_yaml = tmp_path / "base.yaml"