feat: split training targets by task type

2026-09-11 11:27:49 +08:00
parent 272542bad5
commit 947fd12b3a
3 changed files with 164 additions and 0 deletions
@@ -0,0 +1,82 @@
+import pytest
+from pydantic import ValidationError
+
+from lmpm.training.dataset import TrainingDataset, TrainingRecord
+from lmpm.training.targets import TargetSpec, split_targets, target_specs
+
+
+def make_record(number: int = 1, *, target_overrides=None) -> TrainingRecord:
+    targets = {
+        "is_cut_through": True,
+        "carbonized_edge_width": 0.2,
+        "etching_depth": 30.0,
+        "is_fire_smolder": False,
+        "pattern_clarity_score": 9.0,
+        "presentation_balance_score": 8.0,
+    }
+    if target_overrides:
+        targets.update(target_overrides)
+
+    return TrainingRecord(
+        metadata={"experiment_id": f"LAS-2026-{number:04d}"},
+        features={"actual_output_power": 8.5},
+        targets=targets,
+    )
+
+
+def test_target_specs_describe_all_platform_targets():
+    specs = target_specs()
+
+    assert [(spec.name, spec.task) for spec in specs] == [
+        ("is_cut_through", "classification"),
+        ("carbonized_edge_width", "regression"),
+        ("etching_depth", "regression"),
+        ("is_fire_smolder", "classification"),
+        ("pattern_clarity_score", "regression"),
+        ("presentation_balance_score", "regression"),
+    ]
+
+
+def test_target_spec_rejects_unknown_task_type():
+    with pytest.raises(ValidationError):
+        TargetSpec(name="unsupported_target", task="unknown")
+
+
+def test_split_targets_groups_by_task_type():
+    dataset = TrainingDataset(
+        records=(
+            make_record(1),
+            make_record(2, target_overrides={"is_cut_through": False}),
+        )
+    )
+
+    split = split_targets(dataset)
+
+    assert split.classification["is_cut_through"] == (True, False)
+    assert split.classification["is_fire_smolder"] == (False, False)
+    assert split.regression["carbonized_edge_width"] == pytest.approx((0.2, 0.2))
+    assert split.regression["etching_depth"] == pytest.approx((30.0, 30.0))
+    assert split.regression["pattern_clarity_score"] == pytest.approx((9.0, 9.0))
+    assert split.regression["presentation_balance_score"] == pytest.approx((8.0, 8.0))
+
+
+def test_empty_dataset_has_empty_target_groups():
+    split = split_targets(TrainingDataset(records=()))
+
+    assert split.classification == {"is_cut_through": (), "is_fire_smolder": ()}
+    assert split.regression == {
+        "carbonized_edge_width": (),
+        "etching_depth": (),
+        "pattern_clarity_score": (),
+        "presentation_balance_score": (),
+    }
+
+
+@pytest.mark.parametrize("target_name", ["is_cut_through", "etching_depth"])
+def test_missing_target_value_is_rejected(target_name):
+    dataset = TrainingDataset(
+        records=(make_record(target_overrides={target_name: None}),)
+    )
+
+    with pytest.raises(ValueError, match=target_name):
+        split_targets(dataset)