diff --git a/.gitignore b/.gitignore index 2a1488c8..1d8ce8f6 100644 --- a/.gitignore +++ b/.gitignore @@ -129,4 +129,6 @@ dmypy.json .pyre/ */data/* data/* -*.DS* \ No newline at end of file +*.DS* + +.claude/ \ No newline at end of file diff --git a/README.md b/README.md index de81807c..580e2b18 100644 --- a/README.md +++ b/README.md @@ -65,17 +65,67 @@ 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. | -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. +### Adaptive code mapping -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 -under the relevant pipeline stage. The `numeric_value_column_groups` lists control -which transform, optional bound, and derived columns are used; column names and -derived outputs are configured through `numeric_value_columns`. +Adaptive mapping uses raw training-event counts, not distinct-subject counts: + +1. `fit_adaptive_code_mapping` creates a frozen Parquet mapping and + `*.summary.json` audit from all training shards. +2. `apply_adaptive_code_mapping` maps every data shard. +3. `finalize_adaptive_code_metadata` reconciles `codes.parquet` with the mapped + vocabulary. + +Run the standard `extract_code_metadata` stage before these stages. Run later +per-code metadata stages after `finalize_adaptive_code_metadata`. + +#### Hierarchy configuration + +Set `minimum_count` and character-position widths under `hierarchies`. ATC and +SKS defaults are in +[`default_adaptive_code_mapping.yaml`](./configs/MEDS/default_adaptive_code_mapping.yaml). + +Override a built-in namespace or add a new one as needed: + +```yaml +hierarchies: + MY_NAMESPACE: + levels: [2, 4, 6] +``` + +SKS defaults exclude level 1 because an ICD-10/SKS leading letter can span +clinical chapters; for example, `D` covers parts of both neoplasm and blood +disorder chapters. ATC keeps level 1 because it represents the 14 official +anatomical groups. Either default can be overridden. + +### Numeric-value encoding + +Use `aggregate_numeric_metadata` followed by `annotate_numeric_values`. Add +`join_numeric_bins` only for joined lab-and-bin model inputs. + +Defaults are defined in +[`default_numeric_values.yaml`](./configs/MEDS/default_numeric_values.yaml). + +- `numeric_value_columns` names the source and derived columns. +- `numeric_value_column_groups` selects transforms, bounds, and derived columns. + +### Using externally fitted metadata + +External files are fallbacks, not overrides. A local fit always takes +precedence. + +| Artifact | Configuration | Used when | +| --- | --- | --- | +| Adaptive code mapping | `mapping_filepath` on the apply and finalize stages | `fit_adaptive_code_mapping` is omitted | +| Numeric metadata | `numeric_metadata_filepath` on `annotate_numeric_values` | `aggregate_numeric_metadata` is omitted | + +External code mappings may be JSON or Parquet and must contain `code` and +`adaptive/mapped_code`. See +[`lymphoma_pipeline_external_mapping.yaml`](./configs/MEDS/lymphoma_pipeline_external_mapping.yaml) +for an example using both external artifacts. diff --git a/configs/MEDS/default_adaptive_code_mapping.yaml b/configs/MEDS/default_adaptive_code_mapping.yaml new file mode 100644 index 00000000..935a1241 --- /dev/null +++ b/configs/MEDS/default_adaptive_code_mapping.yaml @@ -0,0 +1,25 @@ +minimum_count: 100 +mapping_output_filepath: null +# Defaults to .summary.json. +mapping_summary_output_filepath: null +# Optional external JSON or Parquet mapping with `code` and +# `adaptive/mapped_code` columns. A local fit takes precedence. +mapping_filepath: null + +# Character-position widths by MEDS namespace. Entries can override defaults +# or add namespaces. See README.md for the SKS and ATC level-1 rationale. +hierarchies: + RD: + levels: [4, 3, 2] + RC: + levels: [4, 3, 2] + RC_ADM: + levels: [4, 3, 2] + RPS: + levels: [7, 6, 5, 4, 3, 2] + RPO: + levels: [8, 7, 6, 5, 4, 3, 2] + RM: + levels: [1, 3, 4, 5, 7] + RMA: + levels: [1, 3, 4, 5, 7] diff --git a/configs/MEDS/lymphoma_pipeline.yaml b/configs/MEDS/lymphoma_pipeline.yaml index af5915ac..d7e01e1f 100644 --- a/configs/MEDS/lymphoma_pipeline.yaml +++ b/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 diff --git a/configs/MEDS/lymphoma_pipeline_external_mapping.yaml b/configs/MEDS/lymphoma_pipeline_external_mapping.yaml new file mode 100644 index 00000000..28e538da --- /dev/null +++ b/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 diff --git a/ehr2meds/adaptive_code_mapping.py b/ehr2meds/adaptive_code_mapping.py new file mode 100644 index 00000000..9335cb38 --- /dev/null +++ b/ehr2meds/adaptive_code_mapping.py @@ -0,0 +1,36 @@ +"""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): + local = local_metadata.select(DataSchema.code_name, MAPPED_CODE_COLUMN) + if local.get_column(DataSchema.code_name).n_unique() != local.height: + raise ValueError("local adaptive mapping must contain at most one row per code") + return local + + 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 diff --git a/ehr2meds/io_utils.py b/ehr2meds/io_utils.py new file mode 100644 index 00000000..eb391b04 --- /dev/null +++ b/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") diff --git a/ehr2meds/meds_stages/annotate_numeric_values.py b/ehr2meds/meds_stages/annotate_numeric_values.py index b25aee0d..5b447b93 100644 --- a/ehr2meds/meds_stages/annotate_numeric_values.py +++ b/ehr2meds/meds_stages/annotate_numeric_values.py @@ -4,9 +4,9 @@ import polars as pl from collections.abc import Callable +from ehr2meds.io_utils import load_frame, resolve_resource_path from meds import CodeMetadataSchema, DataSchema from MEDS_transforms.stages import Stage -from MEDS_transforms.utils import PKG_PFX, resolve_pkg_path from omegaconf import DictConfig from pathlib import Path @@ -22,20 +22,10 @@ def find_bin(row: dict[str, object]) -> int | None: def load_external_metadata(filepath: str) -> pl.DataFrame: """Load frozen numeric metadata from JSON, Parquet, or a MEDS metadata directory.""" - path = resolve_pkg_path(filepath) if filepath.startswith(PKG_PFX) else Path(filepath) + path = resolve_resource_path(filepath) if path.is_dir(): path = path / "codes.parquet" - if not path.is_file(): - raise FileNotFoundError(f"numeric_metadata_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( - "numeric_metadata_filepath must be a JSON or Parquet file, or a directory containing codes.parquet" - ) + return load_frame(str(path), "numeric_metadata_filepath") def prepare_metadata( @@ -55,34 +45,25 @@ def prepare_metadata( return metadata.select(key + transform_columns + bound_columns) -def combine_numeric_metadata( +def prepare_numeric_metadata( fitted_metadata: pl.DataFrame, - external_metadata: pl.DataFrame, + external_filepath: str | None, key: list[str], transform_columns: list[str], bound_columns: list[str], ) -> pl.DataFrame: - """Overlay external transforms on local transforms, matching by code key.""" - external_metadata = prepare_metadata( - external_metadata, - key=key, - transform_columns=transform_columns, - bound_columns=bound_columns, - label="external", - ) - fitted_metadata = prepare_metadata( - fitted_metadata, - key=key, - transform_columns=transform_columns, - bound_columns=bound_columns, - label="fitted", + """Use the locally fitted numeric metadata if it exists; otherwise fall back to an external source.""" + if set(key + transform_columns).issubset(fitted_metadata.columns): + return prepare_metadata(fitted_metadata, key=key, transform_columns=transform_columns, bound_columns=bound_columns) + + if not external_filepath: + missing = sorted(set(key + transform_columns) - set(fitted_metadata.columns)) + raise ValueError(f"fitted numeric metadata is missing columns: {missing}") + external_metadata = load_external_metadata(str(external_filepath)) + return prepare_metadata( + external_metadata, key=key, transform_columns=transform_columns, bound_columns=bound_columns, label="external" ) - # External rows come first, so they take precedence for matching codes. - combined = pl.concat([external_metadata, fitted_metadata], how="vertical_relaxed") - combined = combined.unique(subset=key, keep="first", maintain_order=True) - return combined.sort(key) - def is_usable( value: pl.Expr, @@ -186,11 +167,13 @@ def annotate_numeric_values_fntr( code_metadata: pl.DataFrame, code_modifiers: list[str] | None = None, ) -> Callable[[pl.LazyFrame], pl.LazyFrame]: - """Build the shard annotator from local and optional external metadata. + """Build the shard annotator from the locally fitted metadata, or an external source if none was fitted. ``numeric_metadata_filepath`` may point to a fitted numeric-metadata JSON, another dataset's ``metadata/codes.parquet``, or its ``metadata`` directory. - External transforms override locally fitted transforms for matching keys. + It's a fallback, not an override: whenever ``aggregate_numeric_metadata`` + ran locally, that metadata is used and ``numeric_metadata_filepath`` is + ignored. """ key = [CodeMetadataSchema.code_name] + list(code_modifiers or []) columns = stage_cfg.numeric_value_columns @@ -198,17 +181,13 @@ def annotate_numeric_values_fntr( transform_columns = [columns[role] for role in groups.transform] bound_columns = [columns[role] for role in groups.bounds] derived_roles = list(groups.derived) - metadata = code_metadata - external_filepath = stage_cfg.get("numeric_metadata_filepath") - if external_filepath: - external_metadata = load_external_metadata(str(external_filepath)) - metadata = combine_numeric_metadata( - code_metadata, - external_metadata, - key=key, - transform_columns=transform_columns, - bound_columns=bound_columns, - ) + metadata = prepare_numeric_metadata( + code_metadata, + stage_cfg.get("numeric_metadata_filepath"), + key=key, + transform_columns=transform_columns, + bound_columns=bound_columns, + ) def annotate(df: pl.LazyFrame) -> pl.LazyFrame: return annotate_numeric_values( diff --git a/ehr2meds/meds_stages/apply_adaptive_code_mapping.py b/ehr2meds/meds_stages/apply_adaptive_code_mapping.py new file mode 100644 index 00000000..e784b661 --- /dev/null +++ b/ehr2meds/meds_stages/apply_adaptive_code_mapping.py @@ -0,0 +1,41 @@ +"""Apply a frozen adaptive code mapping to every MEDS data shard.""" + +from __future__ import annotations + +import polars as pl +from collections.abc import Callable +from ehr2meds.adaptive_code_mapping import MAPPED_CODE_COLUMN, prepare_mapping +from meds import DataSchema +from MEDS_transforms.stages import Stage +from omegaconf import DictConfig +from pathlib import Path + + +def apply_mapping(data: pl.LazyFrame, mapping: pl.DataFrame) -> pl.LazyFrame: + """Rewrite codes through a frozen mapping while preserving row order and schema.""" + lookup = mapping.select(DataSchema.code_name, MAPPED_CODE_COLUMN).lazy() + return ( + data.join(lookup, on=DataSchema.code_name, how="left", maintain_order="left") + .with_columns(pl.coalesce(MAPPED_CODE_COLUMN, DataSchema.code_name).alias(DataSchema.code_name)) + .drop(MAPPED_CODE_COLUMN) + ) + + +@Stage.register( + is_metadata=False, + default_config=Path("configs/MEDS/default_adaptive_code_mapping.yaml"), +) +def apply_adaptive_code_mapping_fntr( + stage_cfg: DictConfig, + code_metadata: pl.DataFrame, +) -> Callable[[pl.LazyFrame], pl.LazyFrame]: + """Build the data transform from the local fitted mapping or an external one.""" + mapping = prepare_mapping(code_metadata, external_mapping_filepath=stage_cfg.get("mapping_filepath")) + + def transform(df: pl.LazyFrame) -> pl.LazyFrame: + return apply_mapping(df, mapping) + + return transform + + +stage = apply_adaptive_code_mapping_fntr diff --git a/ehr2meds/meds_stages/bin_numeric_values_fast.py b/ehr2meds/meds_stages/bin_numeric_values_fast.py index 0af1a2c9..bad4d494 100644 --- a/ehr2meds/meds_stages/bin_numeric_values_fast.py +++ b/ehr2meds/meds_stages/bin_numeric_values_fast.py @@ -21,11 +21,10 @@ import polars as pl import re from collections.abc import Callable +from ehr2meds.io_utils import resolve_resource_path from meds import CodeMetadataSchema, DataSchema from MEDS_transforms.stages import Stage -from MEDS_transforms.utils import PKG_PFX, resolve_pkg_path from omegaconf import DictConfig, OmegaConf -from pathlib import Path CODE = DataSchema.code_name VALUE = DataSchema.numeric_value_name @@ -192,7 +191,7 @@ def load_custom_bins(stage_cfg: DictConfig) -> dict: if not fp: return inline or {} - path = resolve_pkg_path(fp) if fp.startswith(PKG_PFX) else Path(fp) + path = resolve_resource_path(fp) if not path.is_file(): raise FileNotFoundError(f"custom_bins_filepath '{fp}' does not exist.") from_file = OmegaConf.load(path) diff --git a/ehr2meds/meds_stages/finalize_adaptive_code_metadata.py b/ehr2meds/meds_stages/finalize_adaptive_code_metadata.py new file mode 100644 index 00000000..3a817c93 --- /dev/null +++ b/ehr2meds/meds_stages/finalize_adaptive_code_metadata.py @@ -0,0 +1,125 @@ +"""Rewrite MEDS code metadata to the adaptively mapped vocabulary.""" + +from __future__ import annotations + +import polars as pl +from collections.abc import Sequence +from ehr2meds.adaptive_code_mapping import ( + COUNT_COLUMN, + MAPPED_CODE_COLUMN, + MAPPED_COUNT_COLUMN, + MEMBER_COUNT_COLUMN, + PROFILE_COLUMN, + REASON_COLUMN, + prepare_mapping, +) +from meds import DataSchema +from MEDS_transforms.stages import Stage +from omegaconf import DictConfig +from pathlib import Path + + +def collapse_code_metadata(metadata: pl.DataFrame, mapping: pl.DataFrame) -> pl.DataFrame: + """Rewrite and deterministically collapse code metadata.""" + is_exact_match = "is_exact_match" + mapped = ( + metadata.join(mapping.select(DataSchema.code_name, MAPPED_CODE_COLUMN), on=DataSchema.code_name, how="left") + .with_columns( + pl.coalesce(MAPPED_CODE_COLUMN, DataSchema.code_name).alias(MAPPED_CODE_COLUMN), + (pl.col(DataSchema.code_name) == pl.coalesce(MAPPED_CODE_COLUMN, DataSchema.code_name)).alias(is_exact_match), + ) + .sort(MAPPED_CODE_COLUMN, is_exact_match, DataSchema.code_name, descending=[False, True, False]) + ) + + technical = { + DataSchema.code_name, + is_exact_match, + MAPPED_CODE_COLUMN, + COUNT_COLUMN, + MAPPED_COUNT_COLUMN, + PROFILE_COLUMN, + REASON_COLUMN, + MEMBER_COUNT_COLUMN, + } + preserved = [column for column in metadata.columns if column not in technical] + aggregations = [pl.col(column).drop_nulls().first() for column in preserved] + if COUNT_COLUMN in metadata.columns: + aggregations.append(pl.col(COUNT_COLUMN).fill_null(0).sum().alias(COUNT_COLUMN)) + aggregations.append(pl.len().cast(pl.UInt32).alias(MEMBER_COUNT_COLUMN)) + + collapsed = ( + mapped.group_by(MAPPED_CODE_COLUMN, maintain_order=True) + .agg(*aggregations) + .rename({MAPPED_CODE_COLUMN: DataSchema.code_name}) + ) + + if "description" in collapsed.columns: + generic = pl.format("Adaptive aggregation {} ({} source codes)", DataSchema.code_name, MEMBER_COUNT_COLUMN) + collapsed = collapsed.with_columns( + pl.when(pl.col(MEMBER_COUNT_COLUMN) > 1).then(generic).otherwise(pl.col("description")).alias("description") + ) + if "parent_codes" in collapsed.columns: + collapsed = collapsed.with_columns( + pl.when(pl.col(MEMBER_COUNT_COLUMN) > 1) + .then(pl.lit(None, dtype=pl.List(pl.String))) + .otherwise(pl.col("parent_codes")) + .alias("parent_codes") + ) + return collapsed.sort(DataSchema.code_name) + + +def add_missing_observed_metadata(metadata: pl.DataFrame, observed_codes: Sequence[str]) -> pl.DataFrame: + """Ensure finalized metadata covers every code present in transformed data.""" + missing_codes = sorted(set(observed_codes) - set(metadata.get_column(DataSchema.code_name).to_list())) + if not missing_codes: + return metadata + + columns: dict[str, pl.Series] = {} + for name, dtype in metadata.schema.items(): + if name == DataSchema.code_name: + columns[name] = pl.Series(name, missing_codes, dtype=pl.String) + elif name == COUNT_COLUMN: + columns[name] = pl.Series(name, [0] * len(missing_codes), dtype=dtype) + elif name == MEMBER_COUNT_COLUMN: + columns[name] = pl.Series(name, [1] * len(missing_codes), dtype=dtype) + else: + columns[name] = pl.Series(name, [None] * len(missing_codes), dtype=dtype) + return pl.concat([metadata, pl.DataFrame(columns)]).sort(DataSchema.code_name) + + +@Stage.register( + is_metadata=True, + default_config=Path("configs/MEDS/default_adaptive_code_mapping.yaml"), +) +def main(cfg: DictConfig) -> None: + """Collapse ``codes.parquet`` using the fitted or external mapping.""" + if cfg.worker != 0: + return + + input_filepath = Path(str(cfg.stage_cfg.metadata_input_dir)) / "codes.parquet" + if not input_filepath.is_file(): + raise FileNotFoundError(f"Adaptive code metadata input does not exist: {input_filepath}") + metadata = pl.read_parquet(input_filepath) + mapping = prepare_mapping(metadata, external_mapping_filepath=cfg.stage_cfg.get("mapping_filepath")) + collapsed = collapse_code_metadata(metadata, mapping) + data_input_dir = Path(str(cfg.stage_cfg.data_input_dir)) + data_files = sorted(data_input_dir.glob("**/*.parquet")) + if not data_files: + raise FileNotFoundError(f"No transformed MEDS data shards found in {data_input_dir}") + observed_codes = ( + pl.concat([pl.scan_parquet(path).select(DataSchema.code_name) for path in data_files]) + .select(pl.col(DataSchema.code_name).unique()) + .collect() + .get_column(DataSchema.code_name) + .to_list() + ) + collapsed = add_missing_observed_metadata(collapsed, observed_codes) + + output_filepath = Path(str(cfg.stage_cfg.reducer_output_dir)) / "codes.parquet" + if output_filepath.exists() and not cfg.do_overwrite: + raise FileExistsError(f"Output file already exists: {output_filepath}") + output_filepath.parent.mkdir(parents=True, exist_ok=True) + collapsed.write_parquet(output_filepath) + + +stage = main diff --git a/ehr2meds/meds_stages/fit_adaptive_code_mapping.py b/ehr2meds/meds_stages/fit_adaptive_code_mapping.py new file mode 100644 index 00000000..fa3857c8 --- /dev/null +++ b/ehr2meds/meds_stages/fit_adaptive_code_mapping.py @@ -0,0 +1,298 @@ +"""Fit adaptive hierarchical code mappings on training event counts.""" + +from __future__ import annotations + +import json +import polars as pl +from collections import defaultdict +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from ehr2meds.adaptive_code_mapping import ( + COUNT_COLUMN, + MAPPED_CODE_COLUMN, + MAPPED_COUNT_COLUMN, + PROFILE_COLUMN, + REASON_COLUMN, +) +from meds import DataSchema +from MEDS_transforms.stages import Stage +from omegaconf import DictConfig, OmegaConf +from pathlib import Path + + +@dataclass(frozen=True) +class HierarchyProfile: + """Adaptive truncation settings for one MEDS namespace.""" + + minimum_count: int + levels: tuple[int, ...] + + +def read_profiles(stage_cfg: DictConfig) -> dict[str, HierarchyProfile]: + """Read hierarchy profiles keyed by MEDS namespace.""" + cfg = OmegaConf.to_container(stage_cfg, resolve=True) + hierarchies = cfg["hierarchies"] + default_minimum = int(cfg["minimum_count"]) + return { + str(namespace): HierarchyProfile( + minimum_count=int(hierarchy.get("minimum_count", default_minimum)), + levels=tuple(sorted({int(level) for level in hierarchy["levels"]}, reverse=True)), + ) + for namespace, hierarchy in sorted(hierarchies.items()) + } + + +def split_code(code: str) -> tuple[str, str] | None: + """Split a MEDS code into its first namespace and remaining payload.""" + namespace, separator, payload = code.partition("//") + if not separator or not namespace or not payload: + return None + return namespace, payload + + +def candidate_ancestors(payload: str, profile: HierarchyProfile) -> list[str]: + """Return nearest-to-broadest candidate ancestor payloads.""" + normalized = payload.strip().upper().replace(".", "") + return [normalized[:length] for length in profile.levels if length < len(normalized)] + + +def make_record( + code: str, + count: int, + profile_name: str | None, + reason: str, + mapped_code: str | None = None, + mapped_count: int | None = None, +) -> dict: + return { + DataSchema.code_name: code, + MAPPED_CODE_COLUMN: code if mapped_code is None else mapped_code, + COUNT_COLUMN: count, + MAPPED_COUNT_COLUMN: count if mapped_count is None else mapped_count, + PROFILE_COLUMN: profile_name, + REASON_COLUMN: reason, + } + + +def resolve_pending_codes( + counts: Mapping[str, int], + profiles: Mapping[str, HierarchyProfile], + pending_by_profile: Mapping[str, set[str]], + candidates_by_code: Mapping[str, list[str]], + records: Mapping[str, dict[str, object]], +) -> dict[str, dict[str, object]]: + """Resolve below-threshold codes against their candidate ancestors.""" + resolved_records = {code: record.copy() for code, record in records.items()} + + for profile_name, profile_pending in pending_by_profile.items(): + profile = profiles[profile_name] + pending = set(profile_pending) + + by_candidate: dict[str, list[str]] = defaultdict(list) + for code in sorted(pending): + for candidate in candidates_by_code[code]: + by_candidate[candidate].append(code) + + # Resolve longest candidates first so codes cannot later move to broader ancestors. + for candidate in sorted(by_candidate, key=lambda c: (-len(c), c)): + members = [code for code in by_candidate[candidate] if code in pending] + # Include an observed target in its own aggregate and prevent further truncation. + candidate_is_pending = candidate in pending + if candidate_is_pending: + members.append(candidate) + + mapped_count = sum(int(counts[code]) for code in members) + candidate_needs_own_count = candidate in counts and not candidate_is_pending + if candidate_needs_own_count: + mapped_count += int(counts[candidate]) + + if mapped_count < profile.minimum_count: + continue + for code in members: + resolved_records[code] = make_record( + code, int(counts[code]), profile_name, "grouped", mapped_code=candidate, mapped_count=mapped_count + ) + pending.remove(code) + if candidate_needs_own_count: + resolved_records[candidate][MAPPED_COUNT_COLUMN] = mapped_count + + for code in sorted(pending): + resolved_records[code] = make_record(code, int(counts[code]), profile_name, "below_threshold") + + return resolved_records + + +def fit_mapping( + counts: Mapping[str, int], + profiles: Mapping[str, HierarchyProfile], +) -> pl.DataFrame: + """Fit a deterministic, disjoint adaptive hierarchy mapping.""" + records: dict[str, dict[str, object]] = {} + pending_by_profile: dict[str, set[str]] = defaultdict(set) + candidates_by_code: dict[str, list[str]] = {} + + for code, count in counts.items(): + count = int(count) + parsed = split_code(code) + profile_name = parsed[0] if parsed and parsed[0] in profiles else None + if profile_name is None: + records[code] = make_record(code, count, None, "unconfigured") + continue + + profile = profiles[profile_name] + if count >= profile.minimum_count: + records[code] = make_record(code, count, profile_name, "retained") + continue + + namespace, payload = parsed + ancestors = [f"{namespace}//{candidate}" for candidate in candidate_ancestors(payload, profile)] + if not ancestors: + records[code] = make_record(code, count, profile_name, "no_hierarchy") + continue + + candidates_by_code[code] = ancestors + pending_by_profile[profile_name].add(code) + + records = resolve_pending_codes(counts, profiles, pending_by_profile, candidates_by_code, records) + + schema = { + DataSchema.code_name: pl.String, + MAPPED_CODE_COLUMN: pl.String, + COUNT_COLUMN: pl.UInt64, + MAPPED_COUNT_COLUMN: pl.UInt64, + PROFILE_COLUMN: pl.String, + REASON_COLUMN: pl.String, + } + return pl.DataFrame([records[code] for code in sorted(records)], schema=schema) + + +def combine_count_frames(*dfs: pl.DataFrame | pl.LazyFrame) -> dict[str, int]: + """Sum mapped shard counts.""" + totals: dict[str, int] = defaultdict(int) + for df in dfs: + frame = df.collect() if isinstance(df, pl.LazyFrame) else df + for code, count in frame.select(DataSchema.code_name, COUNT_COLUMN).iter_rows(): + totals[str(code)] += int(count) + return dict(totals) + + +def summarize_mapping(mapping: pl.DataFrame) -> dict[str, object]: + """Return a compact, JSON-friendly audit of a fitted mapping.""" + changed = pl.col(DataSchema.code_name) != pl.col(MAPPED_CODE_COLUMN) + training_rows = pl.col(COUNT_COLUMN) > 0 + totals = mapping.select( + metadata_source_codes=pl.len(), + training_source_codes=training_rows.sum(), + output_codes=pl.col(MAPPED_CODE_COLUMN).n_unique(), + changed_source_codes=changed.sum(), + training_events=pl.col(COUNT_COLUMN).sum(), + remapped_training_events=pl.col(COUNT_COLUMN).filter(changed).sum(), + ).to_dicts()[0] + decisions = ( + mapping.group_by( + pl.col(PROFILE_COLUMN).alias("profile"), + pl.col(REASON_COLUMN).alias("reason"), + ) + .agg( + pl.len().alias("source_codes"), + pl.col(COUNT_COLUMN).sum().alias("training_events"), + ) + .sort("profile", "reason", nulls_last=True) + .to_dicts() + ) + return { + "summary": totals, + "decisions": decisions, + "columns": { + DataSchema.code_name: "original MEDS code", + MAPPED_CODE_COLUMN: "code used after adaptive truncation", + COUNT_COLUMN: "raw events in training data", + MAPPED_COUNT_COLUMN: "training events represented by the mapped code", + PROFILE_COLUMN: "hierarchy used", + REASON_COLUMN: "mapping decision", + }, + } + + +def add_unseen_metadata_codes(mapping: pl.DataFrame, code_metadata: pl.DataFrame) -> pl.DataFrame: + """Carry forward metadata codes absent from training without fitting them.""" + unseen_codes = sorted( + set(code_metadata.get_column(DataSchema.code_name).to_list()) - set(mapping.get_column(DataSchema.code_name)) + ) + if not unseen_codes: + return mapping + unseen = pl.DataFrame( + { + DataSchema.code_name: unseen_codes, + MAPPED_CODE_COLUMN: unseen_codes, + COUNT_COLUMN: [0] * len(unseen_codes), + MAPPED_COUNT_COLUMN: [0] * len(unseen_codes), + PROFILE_COLUMN: [None] * len(unseen_codes), + REASON_COLUMN: ["unseen_training"] * len(unseen_codes), + }, + schema=mapping.schema, + ) + return pl.concat([mapping, unseen]).sort(DataSchema.code_name) + + +def mapper_fntr(stage_cfg: DictConfig) -> Callable[[pl.LazyFrame], pl.LazyFrame]: + """Count training events per code; ``train_only`` selects the shards.""" + read_profiles(stage_cfg) + + def mapper(df: pl.LazyFrame) -> pl.LazyFrame: + return ( + df.group_by(DataSchema.code_name) + .len() + .select( + pl.col(DataSchema.code_name), + pl.col("len").cast(pl.UInt64).alias(COUNT_COLUMN), + ) + .sort(DataSchema.code_name) + ) + + return mapper + + +def reducer_fntr(stage_cfg: DictConfig) -> Callable[..., pl.LazyFrame]: + """Fit one global mapping and write its reusable mapping and audit files.""" + profiles = read_profiles(stage_cfg) + configured_output = stage_cfg.get("mapping_output_filepath") + output_filepath = ( + Path(str(configured_output)) + if configured_output + else Path(str(stage_cfg.reducer_output_dir)) / "adaptive_code_mapping.parquet" + ) + configured_summary = stage_cfg.get("mapping_summary_output_filepath") + if configured_summary: + summary_filepath = Path(str(configured_summary)) + else: + summary_filepath = output_filepath.with_suffix(".summary.json") + code_metadata_filepath = Path(str(stage_cfg.metadata_input_dir)) / "codes.parquet" + code_metadata = ( + pl.read_parquet(code_metadata_filepath) + if code_metadata_filepath.is_file() + else pl.DataFrame(schema={DataSchema.code_name: pl.String}) + ) + + def reducer(*dfs: pl.DataFrame | pl.LazyFrame) -> pl.LazyFrame: + counts = combine_count_frames(*dfs) + mapping = fit_mapping(counts, profiles=profiles) + mapping = add_unseen_metadata_codes(mapping, code_metadata) + output_filepath.parent.mkdir(parents=True, exist_ok=True) + mapping.write_parquet(output_filepath) + summary_filepath.parent.mkdir(parents=True, exist_ok=True) + summary_filepath.write_text( + json.dumps(summarize_mapping(mapping), indent=2, sort_keys=True), + encoding="utf-8", + ) + # Match the lazy type expected by MEDS-Transforms' metadata merge. + return mapping.lazy() + + return reducer + + +stage = Stage.register( + map_fn=mapper_fntr, + reduce_fn=reducer_fntr, + default_config=Path("configs/MEDS/default_adaptive_code_mapping.yaml"), +) diff --git a/pyproject.toml b/pyproject.toml index 7ee2dc58..264812ab 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -50,10 +50,16 @@ line-length = 127 [project.entry-points."MEDS_transforms.stages"] aggregate_numeric_metadata = "ehr2meds.meds_stages.aggregate_numeric_metadata:stage" annotate_numeric_values = "ehr2meds.meds_stages.annotate_numeric_values:stage" +apply_adaptive_code_mapping = "ehr2meds.meds_stages.apply_adaptive_code_mapping:stage" augment_event_config = "ehr2meds.meds_stages.augment_event_config:main" bin_numeric_values_fast = "ehr2meds.meds_stages.bin_numeric_values_fast:stage" +finalize_adaptive_code_metadata = "ehr2meds.meds_stages.finalize_adaptive_code_metadata:stage" +fit_adaptive_code_mapping = "ehr2meds.meds_stages.fit_adaptive_code_mapping:stage" join_numeric_bins = "ehr2meds.meds_stages.join_numeric_bins:stage" +[tool.setuptools.package-data] +ehr2meds = ["resources/*.parquet"] + [tool.ruff.lint] # "E": style/error codes (PEP 8–type errors, e.g. if x == None). # "F": Pyflakes-style errors (unused variables, undefined names, etc.).