Feature/adaptive truncation

.gitignore

@@ -129,4 +129,7 @@ dmypy.json
 .pyre/
 */data/*
 data/*
-*.DS*
+*.DS*
+
+# Claude
+.claude/

README.md

@@ -65,14 +65,71 @@ The package includes the following stages to be used in MEDS pipeline configurat
 | --- | --- |
 | `augment_event_config` | Adds shared columns, such as `row_idx`, to every event definition so they do not need to be repeated throughout the event configuration. |
 | `aggregate_numeric_metadata` | Fits per-code numeric normalization bounds and adaptive quantile bins. It supports training-only fitting, an optional date cutoff (for OOT settings), hard plausibility limits (filtering values greater or lower than biological limits), and writes reusable numeric metadata. |
-| `annotate_numeric_values` | Applies the fitted metadata (from `aggregate_numeric_metadata`) to create new columns based on the numeric values. It adds normalized values, bin indices, and binned representatives. External numeric metadata can override locally fitted metadata. |
+| `annotate_numeric_values` | Applies the fitted metadata (from `aggregate_numeric_metadata`) to create new columns based on the numeric values. It adds normalized values, bin indices, and binned representatives. If no local fit exists, falls back to external numeric metadata instead. |
+| `fit_adaptive_code_mapping` | Fits a code mapping from training-event counts by climbing character-position levels (ATC and SKS diagnosis/operation/procedure) only as far as needed to clear a minimum count. |
+| `apply_adaptive_code_mapping` | Applies the frozen local mapping (or, if none was fitted, an external one) to every data split while retaining the MEDS event namespace. |
+| `finalize_adaptive_code_metadata` | Rewrites and collapses `codes.parquet` to match the adaptively transformed data vocabulary. |
 | `join_numeric_bins` | Optionally creates the "joined representation" of numeric values, such as `LAB_CODE//bin_3`, from the numeric bin index. |
 | `bin_numeric_values_fast` | A faster, memory-efficient replacement for the standard MEDS-Transforms discrete binning stage. It rewrites codes using bin indices or interval labels. |
 
+Adaptive truncation uses raw training-event counts, not distinct-subject
+counts. Its normal configuration is deliberately small: set `minimum_count`
+and map each MEDS namespace to one of `sks_diagnosis`, `sks_operation`,
+`sks_other_procedure`, or `atc` under `namespaces`. Built-in definitions hold
+the character-position level widths for each code system, in
+`configs/MEDS/default_adaptive_code_mapping.yaml`.
+
+The built-in SKS profiles (`sks_diagnosis`, `sks_operation`,
+`sks_other_procedure`) deliberately exclude level 1 (the bare chapter
+letter): ICD-10/SKS chapter boundaries don't align with the leading letter
+(e.g. `D` spans both the tail of chapter II, neoplasms, and all of chapter
+III, blood disorders), so truncating to it would merge clinically unrelated
+codes. ATC's level 1 is kept because it's a real, official top-level tier
+(the 14 anatomical main groups), not an artifact of truncation. A pipeline
+override that reintroduces level 1 for a SKS profile is not blocked -- it
+takes effect exactly as configured.
+
+The three adaptive stages correspond to different MEDS-Transforms data flows:
+
+1. `fit_adaptive_code_mapping` globally reduces **training shards** into one
+   frozen mapping.
+2. `apply_adaptive_code_mapping` maps **every data shard** without learning
+   from tuning or test data.
+3. `finalize_adaptive_code_metadata` reconciles **code metadata** after the
+   transformed data vocabulary is known.
+
+`extract_code_metadata` is the standard upstream MEDS stage, not part of the
+adaptive implementation. The fit stage writes the full reviewable mapping as
+Parquet and a compact `*.summary.json` audit containing vocabulary sizes,
+affected event counts, and decisions by hierarchy and reason. Later per-code
+metadata stages must follow metadata reconciliation.
+
+Custom profiles remain available through the optional `hierarchies` mapping.
+For example, a fixed-length hierarchy can be defined as
+`my_codes: {levels: [2, 4, 6]}` and selected with
+`namespaces: {MY_NAMESPACE: my_codes}`. An entry named like a built-in profile
+overrides only the specified built-in fields.
+
+`mapping_filepath` can point to a frozen JSON or Parquet mapping with `code`
+and `adaptive/mapped_code` columns, sourced from an external collaborator.
+It's a fallback, not an override: whenever `fit_adaptive_code_mapping` ran
+locally, that mapping is used and `mapping_filepath` is ignored. Set it
+(and drop `fit_adaptive_code_mapping` from `stages`) to run `apply_adaptive_code_mapping`
+and `finalize_adaptive_code_metadata` purely off an externally-supplied
+mapping -- see `configs/MEDS/lymphoma_pipeline_external_mapping.yaml` for a
+worked example of this pattern (e.g. adopting a mapping produced by a
+consortium such as PHAIR instead of fitting one locally).
+
 For combined numeric encoding, use `aggregate_numeric_metadata` followed by
 `annotate_numeric_values`. Add `join_numeric_bins` afterwards only when the final
 model input should contain joined lab-and-bin codes.
 
+`annotate_numeric_values`'s `numeric_metadata_filepath` follows the same
+fallback rule as `mapping_filepath` above: it's only used when
+`aggregate_numeric_metadata` did not run locally. See
+`configs/MEDS/lymphoma_pipeline_external_mapping.yaml`, which uses this for
+both numeric metadata and code mapping together.
+
 Shared numeric column names and stage defaults are defined in
 `configs/MEDS/default_numeric_values.yaml`; pipeline configurations only
 need to specify dataset- or run-specific overrides. Any setting can be overridden

configs/MEDS/default_adaptive_code_mapping.yaml

@@ -0,0 +1,35 @@
+minimum_count: 100
+mapping_output_filepath: null
+# Defaults to <mapping_output_filepath>.summary.json.
+mapping_summary_output_filepath: null
+# A frozen mapping (JSON or Parquet, with `code` and `adaptive/mapped_code`
+# columns) from an external collaborator. Only used as a fallback when
+# `fit_adaptive_code_mapping` did not run locally -- a local fit always
+# takes precedence over this.
+mapping_filepath: null
+
+# Only namespace routing normally needs configuration; the level widths for
+# each profile are defined below under `hierarchies`.
+namespaces:
+  RD: sks_diagnosis
+  RC: sks_diagnosis
+  RC_ADM: sks_diagnosis
+  RPS: sks_operation
+  RPO: sks_other_procedure
+  RM: atc
+  RMA: atc
+
+# Character-position level widths per code system. A pipeline config can
+# override any single field of a profile (e.g. `hierarchies: {atc: {levels:
+# [...]}}`) without repeating the rest, or add a wholly new profile name.
+# See README.md ("Custom MEDS stages") for why level 1 is excluded from the
+# SKS profiles here but kept for ATC.
+hierarchies:
+  sks_diagnosis:
+    levels: [4, 3, 2]
+  sks_operation:
+    levels: [7, 6, 5, 4, 3, 2]
+  sks_other_procedure:
+    levels: [8, 7, 6, 5, 4, 3, 2]
+  atc:
+    levels: [1, 3, 4, 5, 7]

configs/MEDS/lymphoma_pipeline.yaml

@@ -41,7 +41,22 @@ stages:
   - merge_to_MEDS_cohort:
       # row_idx provides deterministic ordering for simultaneous events.
       additional_sort_by: [row_idx]
+
+  # Adaptive truncation has three explicit boundaries:
+  # 1. LEARN once from training counts (also writes a readable audit summary).
   - extract_code_metadata
+  - fit_adaptive_code_mapping:
+      train_only: true
+      minimum_count: 100
+      mapping_output_filepath: ${output_dir}/adaptive_code_mapping.parquet
+      mapping_summary_output_filepath: ${output_dir}/adaptive_code_mapping.summary.json
+
+  # 2. APPLY the frozen mapping to all train/tuning/test event shards.
+  - apply_adaptive_code_mapping
+
+  # 3. RECONCILE codes.parquet with the vocabulary now present in the data.
+  - finalize_adaptive_code_metadata
+
   - aggregate_numeric_metadata:
       train_only: true
       # Fit normalization and bins only on pre-cutoff events from training subjects.
@@ -54,9 +69,9 @@ stages:
       #   "L//NPU01566": {min: 0, max: 250}
       numeric_metadata_output_filepath: ${output_dir}/numeric_metadata.json
   - annotate_numeric_values:
-      # Optional: a fitted numeric_metadata.json, MEDS metadata directory, or codes.parquet.
-      # Matching external code transforms override the locally fitted ones.
-      # numeric_metadata_filepath: /path/to/numeric_metadata.json
+      # Optional fallback (JSON, MEDS metadata directory, or codes.parquet)
+      # used only if aggregate_numeric_metadata did not run locally.
+      numeric_metadata_filepath: null
   # Final model representation: valid lab values become LAB_CODE//BIN.
   - join_numeric_bins
   - finalize_MEDS_metadata

configs/MEDS/lymphoma_pipeline_external_mapping.yaml

@@ -0,0 +1,82 @@
+description: |-
+  Variant of `lymphoma_pipeline.yaml` for sites that adopt an adaptive code
+  mapping and/or numeric-value metadata produced elsewhere (e.g. by a
+  consortium such as PHAIR) instead of fitting their own from local training
+  counts. The only differences from the base pipeline are:
+
+  1. `fit_adaptive_code_mapping` is dropped from `stages` entirely -- no
+     local code-mapping fit is attempted. `apply_adaptive_code_mapping` and
+     `finalize_adaptive_code_metadata` both set `mapping_filepath` to the
+     externally-supplied mapping.
+  2. `aggregate_numeric_metadata` is likewise dropped -- no local numeric fit
+     is attempted. `annotate_numeric_values` sets `numeric_metadata_filepath`
+     to the externally-supplied numeric metadata instead.
+
+  In both cases, the external source is a fallback, not an override: since
+  no local fit ran, the relevant metadata columns don't exist yet at that
+  point in the pipeline, so both stages fall back to the external source
+  (see README.md, "Custom MEDS stages", for why a local fit -- when one
+  exists -- always takes precedence over this instead of being merged with
+  it).
+
+  Set the same environment variables as the base pipeline, plus:
+
+  export EXTERNAL_MAPPING_FP=/path/to/phair_code_mapping.parquet          # code + adaptive/mapped_code columns
+  export EXTERNAL_NUMERIC_METADATA_FP=/path/to/phair_numeric_metadata.json # fitted numeric-value transforms
+
+input_dir: ${oc.env:PREMEDS_INPUT_DIR}
+output_dir: ${oc.env:MEDS_OUTPUT_DIR}
+shards_map_fp: ${output_dir}/metadata/.shards.json
+source_event_conversion_config_fp: ${oc.env:EVENT_CONVERSION_CONFIG_FP}
+event_conversion_config_fp: ${output_dir}/event_conversion_config.augmented.yaml
+external_mapping_fp: ${oc.env:EXTERNAL_MAPPING_FP}
+external_numeric_metadata_fp: ${oc.env:EXTERNAL_NUMERIC_METADATA_FP}
+
+
+etl_metadata:
+  dataset_name: preMEDS
+  dataset_version: 2.0
+
+stages:
+  - augment_event_config:
+      source_event_conversion_config_fp: ${source_event_conversion_config_fp}
+      output_event_conversion_config_fp: ${event_conversion_config_fp}
+      event_columns:
+        row_idx: $row_idx
+
+  - shard_events:
+      infer_schema_length: 999999999
+      # MEDS-Extract derives the raw cohort root from this path's parent.
+      data_input_dir: ${input_dir}/data
+
+  - split_and_shard_subjects:
+      n_subjects_per_shard: 10000
+      split_fracs:
+          train: 0.9
+          tuning: 0.1
+  - convert_to_subject_sharded
+  - convert_to_MEDS_events:
+      do_dedup_text_and_numeric: true
+  - merge_to_MEDS_cohort:
+      # row_idx provides deterministic ordering for simultaneous events.
+      additional_sort_by: [row_idx]
+
+  # No local fit here -- adopt the externally-supplied mapping instead.
+  - extract_code_metadata
+
+  # APPLY the external mapping to all train/tuning/test event shards.
+  - apply_adaptive_code_mapping:
+      mapping_filepath: ${external_mapping_fp}
+
+  # RECONCILE codes.parquet with the vocabulary now present in the data.
+  - finalize_adaptive_code_metadata:
+      mapping_filepath: ${external_mapping_fp}
+
+  # No local fit here either -- adopt the externally-supplied numeric
+  # metadata instead.
+  - annotate_numeric_values:
+      numeric_metadata_filepath: ${external_numeric_metadata_fp}
+  # Final model representation: valid lab values become LAB_CODE//BIN.
+  - join_numeric_bins
+  - finalize_MEDS_metadata
+  - finalize_MEDS_data

ehr2meds/adaptive_code_mapping.py

@@ -0,0 +1,33 @@
+"""Column names and mapping I/O shared by the adaptive code-mapping stages."""
+
+from __future__ import annotations
+
+import polars as pl
+from ehr2meds.io_utils import load_frame
+from meds import DataSchema
+
+MAPPED_CODE_COLUMN = "adaptive/mapped_code"
+COUNT_COLUMN = "adaptive/count"
+MAPPED_COUNT_COLUMN = "adaptive/mapped_count"
+PROFILE_COLUMN = "adaptive/profile"
+REASON_COLUMN = "adaptive/reason"
+MEMBER_COUNT_COLUMN = "adaptive/member_count"
+
+
+def prepare_mapping(local_metadata: pl.DataFrame, external_mapping_filepath: str | None) -> pl.DataFrame:
+    """Use the local fitted mapping if one exists; otherwise fall back to an external mapping."""
+    required_local = {DataSchema.code_name, MAPPED_CODE_COLUMN}
+    if required_local.issubset(local_metadata.columns):
+        return local_metadata.select(DataSchema.code_name, MAPPED_CODE_COLUMN)
+
+    if not external_mapping_filepath:
+        missing = sorted(required_local - set(local_metadata.columns))
+        raise ValueError(f"local adaptive metadata is missing columns: {missing}")
+    external = load_frame(str(external_mapping_filepath), "external mapping")
+    missing = required_local - set(external.columns)
+    if missing:
+        raise ValueError(f"external mapping is missing columns: {sorted(missing)}")
+    external = external.select(DataSchema.code_name, MAPPED_CODE_COLUMN)
+    if external.get_column(DataSchema.code_name).n_unique() != external.height:
+        raise ValueError("external mapping must contain at most one row per code")
+    return external

ehr2meds/io_utils.py

@@ -0,0 +1,26 @@
+"""Shared file-resolution and loading helpers used across ehr2meds stages."""
+
+from __future__ import annotations
+
+import polars as pl
+from MEDS_transforms.utils import PKG_PFX, resolve_pkg_path
+from pathlib import Path
+
+
+def resolve_resource_path(filepath: str) -> Path:
+    """Resolve normal and package resource paths."""
+    return resolve_pkg_path(filepath) if filepath.startswith(PKG_PFX) else Path(filepath)
+
+
+def load_frame(filepath: str, label: str) -> pl.DataFrame:
+    """Load a JSON or Parquet resource."""
+    path = resolve_resource_path(filepath)
+    if not path.is_file():
+        raise FileNotFoundError(f"{label} filepath '{filepath}' does not exist")
+    match path.suffix.lower():
+        case ".parquet":
+            return pl.read_parquet(path)
+        case ".json":
+            return pl.read_json(path)
+        case _:
+            raise ValueError(f"{label} filepath must point to a JSON or Parquet file")