diff --git a/src/harmonization_framework/primitives/__init__.py b/src/harmonization_framework/primitives/__init__.py index 9d972d9..7f43daf 100644 --- a/src/harmonization_framework/primitives/__init__.py +++ b/src/harmonization_framework/primitives/__init__.py @@ -1,6 +1,8 @@ from .base import PrimitiveOperation -from .bin_primitive import Bin +from .bin_primitive import Bin +from .case import Case from .cast import Cast +from .coalesce import Coalesce from .dates import ConvertDate from .donothing import DoNothing from .enum2enum import EnumToEnum diff --git a/src/harmonization_framework/primitives/case.py b/src/harmonization_framework/primitives/case.py new file mode 100644 index 0000000..87dce01 --- /dev/null +++ b/src/harmonization_framework/primitives/case.py @@ -0,0 +1,163 @@ +from typing import Any, Dict, List + +from .base import PrimitiveOperation, isnull + + +def _normalize_operands(branch): + """A branch produces one value from one or more 'operands'. Each operand computes + a value from one source via its own op-chain; a branch with multiple operands + combines them with `combine` (a Reduction name string, e.g. 'sum'). + + Accepts two authoring forms, normalized to the canonical operand list: + - {"source": "x", "operations": [...]} -> one operand + - {"operands": [{"source": "a", "operations": [...]}, ...], "combine": "sum"} + """ + if "operands" in branch: + operands = [{"source": o["source"], "operations": list(o.get("operations", []))} + for o in branch["operands"]] + combine = branch.get("combine") + else: + operands = [{"source": branch["source"], "operations": list(branch.get("operations", []))}] + combine = None + return operands, combine + + +def _apply_operands(operands, combine, by_name): + """Compute each operand's value, then combine. Returns the branch value.""" + from .reduce import Reduce, Reduction + + vals = [] + for o in operands: + current = by_name[o["source"]] + for op in o["operations"]: + current = op(current) + vals.append(current) + if combine is None: + return vals[0] if len(vals) == 1 else vals + return Reduce(Reduction(combine)).transform(vals) + + +class Case(PrimitiveOperation): + """ + Choose one of several branches by switching on a selector source, then compute + that branch's value from one or more source operands. + + A multi-source combinator. The rule's `sources` list names every column this + primitive may read; `transform` receives the values in that same order and + `Case` zips them back to names internally, so branches address sources by NAME, + never by fragile positional index. + + A branch is either a single source with an op-chain, or several `operands` (each + a source + op-chain) combined with a reduction. The first branch whose `when` + contains the selector value wins; null/unmatched selector -> `default`. + + Single-source branch (RADx-UP weight: unit flag selects pounds-as-is or kg->lbs): + + Case( + sources=["weight_units", "weight_lbs", "weight_kgs"], + selector="weight_units", + branches=[ + {"when": ["2"], "source": "weight_lbs", "operations": [DoNothing()]}, + {"when": ["1"], "source": "weight_kgs", + "operations": [ConvertUnits(Unit.KILOGRAM, Unit.POUNDS), Round(0)]}, + ], + ) + + Multi-operand branch (RADx-UP height: feet+inches OR meters+cm, summed in inches): + + Case( + sources=["height_units", "ft", "in", "m", "cm"], + selector="height_units", + branches=[ + {"when": ["1"], "combine": "sum", "operands": [ + {"source": "ft", "operations": [ConvertUnits(Unit.FEET, Unit.INCH)]}, + {"source": "in", "operations": [DoNothing()]}, + ]}, + {"when": ["2"], "combine": "sum", "operands": [ + {"source": "m", "operations": [ConvertUnits(Unit.METER, Unit.INCH)]}, + {"source": "cm", "operations": [ConvertUnits(Unit.CENTIMETER, Unit.INCH)]}, + ]}, + ], + ) + """ + + def __init__(self, sources, selector, branches, default=None): + self.sources = list(sources) + self.selector = selector + self.branches = [] + for b in branches: + operands, combine = _normalize_operands(b) + self.branches.append({"when": [str(w) for w in b["when"]], + "operands": operands, "combine": combine}) + self.default = default + if self.selector not in self.sources: + raise ValueError(f"Case selector {self.selector!r} not in sources {self.sources}") + for b in self.branches: + for o in b["operands"]: + if o["source"] not in self.sources: + raise ValueError(f"Case branch source {o['source']!r} not in sources {self.sources}") + + def __str__(self): + lines = [f"Switch on {self.selector}:"] + for b in self.branches: + srcs = "+".join(o["source"] for o in b["operands"]) + comb = f" ({b['combine']})" if b["combine"] else "" + lines.append(f" when {b['when']} -> {srcs}{comb}") + lines.append(f" else -> {self.default!r}") + return "\n".join(lines) + + def _by_name(self, values): + if not isinstance(values, (list, tuple)): + raise TypeError(f"Case expects a list of source values, got {type(values).__name__}") + if len(values) != len(self.sources): + raise ValueError(f"Case received {len(values)} values but has {len(self.sources)} sources") + return dict(zip(self.sources, values)) + + def transform(self, values: Any) -> Any: + by_name = self._by_name(values) + sel = by_name[self.selector] + if isnull(sel): + return self.default + sel_key = str(sel) + for b in self.branches: + if sel_key in b["when"]: + return _apply_operands(b["operands"], b["combine"], by_name) + return self.default + + def to_dict(self): + return { + "operation": "case", + "sources": list(self.sources), + "selector": self.selector, + "branches": [ + { + "when": list(b["when"]), + "combine": b["combine"], + "operands": [ + {"source": o["source"], "operations": [op.to_dict() for op in o["operations"]]} + for o in b["operands"] + ], + } + for b in self.branches + ], + "default": self.default, + } + + @classmethod + def from_serialization(cls, serialization): + from .factory import deserialize_operation + + branches = [] + for b in serialization["branches"]: + operands = [ + {"source": o["source"], + "operations": [deserialize_operation(op) for op in o.get("operations", [])]} + for o in b["operands"] + ] + branches.append({"when": b["when"], "combine": b.get("combine"), "operands": operands}) + return cls( + sources=serialization["sources"], + selector=serialization["selector"], + branches=branches, + default=serialization.get("default"), + ) diff --git a/src/harmonization_framework/primitives/coalesce.py b/src/harmonization_framework/primitives/coalesce.py new file mode 100644 index 0000000..fa51932 --- /dev/null +++ b/src/harmonization_framework/primitives/coalesce.py @@ -0,0 +1,101 @@ +from typing import Any + +from .base import PrimitiveOperation, isnull +from .case import _normalize_operands, _apply_operands + + +class Coalesce(PrimitiveOperation): + """ + Pick the first branch whose (primary) source is non-null, then compute that + branch's value from one or more source operands. + + A multi-source combinator, like Case but with the branch chosen by + populated-ness rather than an explicit selector. Use Case when the data carries + an authoritative unit/type flag; use Coalesce when "whichever field was filled + in" is the right rule. + + A branch is a single source with an op-chain, or several `operands` (each a + source + op-chain) combined with a reduction (same shape as Case branches). A + branch is considered populated when its first operand's source is non-null. + Branch order defines precedence; all sources null -> `default`. + + Coalesce( + sources=["weight_lbs", "weight_kgs"], + branches=[ + {"source": "weight_lbs", "operations": []}, + {"source": "weight_kgs", + "operations": [ConvertUnits(Unit.KILOGRAM, Unit.POUNDS), Round(0)]}, + ], + ) + """ + + def __init__(self, sources, branches, default=None): + self.sources = list(sources) + self.branches = [] + for b in branches: + operands, combine = _normalize_operands(b) + self.branches.append({"operands": operands, "combine": combine}) + self.default = default + for b in self.branches: + for o in b["operands"]: + if o["source"] not in self.sources: + raise ValueError(f"Coalesce branch source {o['source']!r} not in sources {self.sources}") + + def __str__(self): + lines = ["Coalesce (first non-null):"] + for b in self.branches: + srcs = "+".join(o["source"] for o in b["operands"]) + comb = f" ({b['combine']})" if b["combine"] else "" + lines.append(f" {srcs}{comb}") + lines.append(f" else -> {self.default!r}") + return "\n".join(lines) + + def _by_name(self, values): + if not isinstance(values, (list, tuple)): + raise TypeError(f"Coalesce expects a list of source values, got {type(values).__name__}") + if len(values) != len(self.sources): + raise ValueError(f"Coalesce received {len(values)} values but has {len(self.sources)} sources") + return dict(zip(self.sources, values)) + + def transform(self, values: Any) -> Any: + by_name = self._by_name(values) + for b in self.branches: + primary = by_name[b["operands"][0]["source"]] + if not isnull(primary): + return _apply_operands(b["operands"], b["combine"], by_name) + return self.default + + def to_dict(self): + return { + "operation": "coalesce", + "sources": list(self.sources), + "branches": [ + { + "combine": b["combine"], + "operands": [ + {"source": o["source"], "operations": [op.to_dict() for op in o["operations"]]} + for o in b["operands"] + ], + } + for b in self.branches + ], + "default": self.default, + } + + @classmethod + def from_serialization(cls, serialization): + from .factory import deserialize_operation + + branches = [] + for b in serialization["branches"]: + operands = [ + {"source": o["source"], + "operations": [deserialize_operation(op) for op in o.get("operations", [])]} + for o in b["operands"] + ] + branches.append({"combine": b.get("combine"), "operands": operands}) + return cls( + sources=serialization["sources"], + branches=branches, + default=serialization.get("default"), + ) diff --git a/src/harmonization_framework/primitives/factory.py b/src/harmonization_framework/primitives/factory.py index bf19b2f..e2943bc 100644 --- a/src/harmonization_framework/primitives/factory.py +++ b/src/harmonization_framework/primitives/factory.py @@ -11,7 +11,9 @@ from .base import PrimitiveOperation from .bin_primitive import Bin +from .case import Case from .cast import Cast +from .coalesce import Coalesce from .dates import ConvertDate from .donothing import DoNothing from .enum2enum import EnumToEnum @@ -44,8 +46,12 @@ def deserialize_operation(operation: Dict[str, Any]) -> PrimitiveOperation: match name: case PrimitiveVocabulary.BIN.value: return Bin.from_serialization(operation) + case PrimitiveVocabulary.CASE.value: + return Case.from_serialization(operation) case PrimitiveVocabulary.CAST.value: return Cast.from_serialization(operation) + case PrimitiveVocabulary.COALESCE.value: + return Coalesce.from_serialization(operation) case PrimitiveVocabulary.CONVERT_DATE.value: return ConvertDate.from_serialization(operation) case PrimitiveVocabulary.CONVERT_UNITS.value: diff --git a/src/harmonization_framework/primitives/vocabulary.py b/src/harmonization_framework/primitives/vocabulary.py index c8deb63..999001b 100644 --- a/src/harmonization_framework/primitives/vocabulary.py +++ b/src/harmonization_framework/primitives/vocabulary.py @@ -2,7 +2,9 @@ class PrimitiveVocabulary(Enum): BIN = "bin" + CASE = "case" CAST = "cast" + COALESCE = "coalesce" CONVERT_DATE = "convert_date" CONVERT_UNITS = "convert_units" DO_NOTHING = "do_nothing" diff --git a/tests/test_case_coalesce.py b/tests/test_case_coalesce.py new file mode 100644 index 0000000..2533ee5 --- /dev/null +++ b/tests/test_case_coalesce.py @@ -0,0 +1,201 @@ +"""Tests for the Case and Coalesce multi-source combinator primitives.""" +import pandas as pd + +from harmonization_framework.harmonization_rule import HarmonizationRule +from harmonization_framework.harmonize import harmonize_dataset +from harmonization_framework.primitives import ( + Case, Coalesce, ConvertUnits, DoNothing, Round, Scale) +from harmonization_framework.primitives.units import Unit +from harmonization_framework.rule_registry import RuleSet + + +# --- Case ---------------------------------------------------------------- + +def _weight_case(): + return Case( + sources=["weight_units", "weight_lbs", "weight_kgs"], + selector="weight_units", + branches=[ + {"when": ["2"], "source": "weight_lbs", "operations": [DoNothing()]}, + {"when": ["1"], "source": "weight_kgs", "operations": [Scale(2.20462), Round(0)]}, + ], + default=None, + ) + + +def test_case_selects_pounds_branch_as_is(): + c = _weight_case() + # units=2 (pounds): take weight_lbs unchanged + assert c.transform(["2", 150, 68]) == 150 + + +def test_case_selects_kg_branch_and_converts(): + c = _weight_case() + # units=1 (kilograms): convert weight_kgs -> lbs, rounded + assert c.transform(["1", 150, 68]) == round(68 * 2.20462) + + +def test_case_int_selector_matches_string_when(): + c = _weight_case() + # selector arrives as an int 1 (e.g. pandas) but `when` is "1" + assert c.transform([1, 150, 68]) == round(68 * 2.20462) + + +def test_case_null_selector_returns_default(): + c = _weight_case() + assert c.transform([None, 150, 68]) is None + + +def test_case_unmatched_selector_returns_default(): + c = _weight_case() + assert c.transform(["99", 150, 68]) is None + + +def test_case_serialization_roundtrip(): + rule = HarmonizationRule( + ["weight_units", "weight_lbs", "weight_kgs"], "nih_weight", [_weight_case()] + ) + payload = rule.serialize() + assert payload["operations"][0]["operation"] == "case" + roundtrip = HarmonizationRule.from_serialization(payload) + assert roundtrip.serialize() == payload + assert roundtrip.transform(["1", 150, 68]) == round(68 * 2.20462) + + +def test_case_in_harmonize_dataset(): + rules = RuleSet() + rules.add_rule( + HarmonizationRule( + ["weight_units", "weight_lbs", "weight_kgs"], "nih_weight", [_weight_case()] + ) + ) + df = pd.DataFrame([ + {"weight_units": "2", "weight_lbs": 150, "weight_kgs": None}, # pounds + {"weight_units": "1", "weight_lbs": None, "weight_kgs": 68}, # kilograms + ]) + out = harmonize_dataset(df, rules, dataset_name="t") + assert out["nih_weight"].tolist() == [150, round(68 * 2.20462)] + + +# --- Case with multi-operand branches (height: feet+inches OR meters+cm) ---- + +def _height_case(): + return Case( + sources=["height_units", "ft", "inch", "m", "cm"], + selector="height_units", + branches=[ + {"when": ["1"], "combine": "sum", "operands": [ + {"source": "ft", "operations": [ConvertUnits(Unit.FEET, Unit.INCH)]}, + {"source": "inch", "operations": [DoNothing()]}, + ]}, + {"when": ["2"], "combine": "sum", "operands": [ + {"source": "m", "operations": [ConvertUnits(Unit.METER, Unit.INCH)]}, + {"source": "cm", "operations": [ConvertUnits(Unit.CENTIMETER, Unit.INCH)]}, + ]}, + ], + default=None, + ) + + +def test_case_multioperand_feet_inches_sum(): + c = _height_case() + # 5 ft + 7 in -> 60 + 7 = 67 inches + assert c.transform(["1", 5, 7, None, None]) == 67 + + +def test_case_multioperand_meters_cm_sum(): + c = _height_case() + # 1 m + 70 cm -> ~39.37 + ~27.56 inches + result = c.transform(["2", None, None, 1, 70]) + assert abs(result - (1 / 0.0254 + 70 / 2.54)) < 0.01 + + +def test_case_multioperand_serialization_roundtrip(): + rule = HarmonizationRule( + ["height_units", "ft", "inch", "m", "cm"], "nih_height", [_height_case()] + ) + payload = rule.serialize() + roundtrip = HarmonizationRule.from_serialization(payload) + assert roundtrip.serialize() == payload + assert roundtrip.transform(["1", 5, 7, None, None]) == 67 + + +# --- Coalesce ------------------------------------------------------------ + +def _weight_coalesce(): + return Coalesce( + sources=["weight_lbs", "weight_kgs"], + branches=[ + {"source": "weight_lbs", "operations": [DoNothing()]}, + {"source": "weight_kgs", "operations": [Scale(2.20462), Round(0)]}, + ], + default=None, + ) + + +def test_coalesce_first_non_null_wins(): + c = _weight_coalesce() + assert c.transform([150, None]) == 150 # lbs present + assert c.transform([None, 68]) == round(68 * 2.20462) # only kg present + + +def test_coalesce_precedence_when_both_present(): + c = _weight_coalesce() + # both populated -> first branch (lbs) wins + assert c.transform([150, 68]) == 150 + + +def test_coalesce_all_null_returns_default(): + c = _weight_coalesce() + assert c.transform([None, None]) is None + + +def test_coalesce_serialization_roundtrip(): + rule = HarmonizationRule(["weight_lbs", "weight_kgs"], "nih_weight", [_weight_coalesce()]) + payload = rule.serialize() + assert payload["operations"][0]["operation"] == "coalesce" + roundtrip = HarmonizationRule.from_serialization(payload) + assert roundtrip.serialize() == payload + assert roundtrip.transform([None, 68]) == round(68 * 2.20462) + + +# --- Coalesce with multi-operand branches + combine ---------------------- +# Weight reported either as a single pounds field, OR as stone + pounds (two +# fields summed). No unit flag, so Coalesce picks whichever branch is populated; +# the stone+pounds branch uses `combine: "sum"` over two operands. + +def _weight_coalesce_combine(): + return Coalesce( + sources=["weight_lbs", "weight_stone", "weight_stone_lbs"], + branches=[ + {"source": "weight_lbs", "operations": [DoNothing()]}, + {"combine": "sum", "operands": [ + {"source": "weight_stone", "operations": [Scale(14)]}, # stone -> lbs + {"source": "weight_stone_lbs", "operations": [DoNothing()]}, # leftover pounds, as-is + ]}, + ], + default=None, + ) + + +def test_coalesce_combine_sums_multioperand_branch(): + c = _weight_coalesce_combine() + # single pounds field populated -> first branch wins, value as-is + assert c.transform([150, None, None]) == 150 + # only stone+pounds populated -> second branch: 10 st * 14 + 7 lb = 147 + assert c.transform([None, 10, 7]) == 147 + + +def test_coalesce_combine_serialization_roundtrip(): + rule = HarmonizationRule( + ["weight_lbs", "weight_stone", "weight_stone_lbs"], "nih_weight", + [_weight_coalesce_combine()], + ) + payload = rule.serialize() + # the multi-operand branch carries combine="sum" and two operands + branch = payload["operations"][0]["branches"][1] + assert branch["combine"] == "sum" + assert [o["source"] for o in branch["operands"]] == ["weight_stone", "weight_stone_lbs"] + roundtrip = HarmonizationRule.from_serialization(payload) + assert roundtrip.serialize() == payload + assert roundtrip.transform([None, 10, 7]) == 147