MEDS-extract 0.1__py3-none-any.whl

MEDS_extract/__init__.py

@@ -0,0 +1,50 @@
+from importlib.metadata import PackageNotFoundError, version
+from importlib.resources import files
+
+import polars as pl
+from meds import code_field, subject_id_field, time_field
+
+__package_name__ = "MEDS_extract"
+try:
+    __version__ = version(__package_name__)
+except PackageNotFoundError:  # pragma: no cover
+    __version__ = "unknown"
+
+CONFIG_YAML = files(__package_name__).joinpath("configs/_extract.yaml")
+
+MANDATORY_COLUMNS = [subject_id_field, time_field, code_field, "numeric_value"]
+
+MANDATORY_TYPES = {
+    subject_id_field: pl.Int64,
+    time_field: pl.Datetime("us"),
+    code_field: pl.String,
+    "numeric_value": pl.Float32,
+    "categorical_value": pl.String,
+    "text_value": pl.String,
+}
+
+DEPRECATED_NAMES = {
+    "numerical_value": "numeric_value",
+    "categoric_value": "categoric_value",
+    "category_value": "categoric_value",
+    "textual_value": "text_value",
+    "timestamp": "time",
+    "patient_id": subject_id_field,
+}
+
+INFERRED_STAGE_KEYS = {
+    "is_metadata",
+    "data_input_dir",
+    "metadata_input_dir",
+    "output_dir",
+    "reducer_output_dir",
+}
+
+# TODO(mmd): This should really somehow be pulled from MEDS.
+MEDS_METADATA_MANDATORY_TYPES = {
+    "code": pl.String,
+    "description": pl.String,
+    "parent_codes": pl.List(pl.String),
+}
+
+MEDS_DATA_MANDATORY_TYPES = {c: MANDATORY_TYPES[c] for c in MANDATORY_COLUMNS}

MEDS_extract/configs/_extract.yaml

@@ -0,0 +1,49 @@
+defaults:
+  - _pipeline
+  - stage_configs:
+      - shard_events
+      - split_and_shard_subjects
+      - convert_to_sharded_events
+      - merge_to_MEDS_cohort
+      - extract_code_metadata
+      - finalize_MEDS_metadata
+      - finalize_MEDS_data
+      # There is no configuration beyond the global "event_conversion_config_fp" for the
+      # convert_to_sharded_events stage, so it doesn't have a stage config block here or below.
+  - _self_
+
+etl_metadata.pipeline_name: "extract"
+
+description: |-
+  This pipeline extracts raw MEDS events in longitudinal, sparse form from an input dataset meeting select
+  criteria and converts them to the flattened, MEDS format. It can be run in its entirety, with controllable
+  levels of parallelism, or in stages. Arguments:
+    - `event_conversion_config_fp`: The path to the event conversion configuration file. This file defines
+      the events to extract from the various rows of the various input files encountered in the global input
+      directory.
+    - `input_dir`: The path to the directory containing the raw input files.
+    - `cohort_dir`: The path to the directory where the output cohort will be written. It will be written in
+      various subfolders of this dir depending on the stage, as intermediate stages cache their output during
+      computation for efficiency of re-running and distributing.
+
+# The event conversion configuration file is used throughout the pipeline to define the events to extract.
+event_conversion_config_fp: ???
+# The shards mapping is stored in the root of the final output directory.
+shards_map_fp: "${cohort_dir}/metadata/.shards.json"
+
+stages:
+  - shard_events
+  - split_and_shard_subjects
+  - convert_to_sharded_events
+  - merge_to_MEDS_cohort
+  - extract_code_metadata
+  - finalize_MEDS_metadata
+  - finalize_MEDS_data
+
+stage_configs:
+  shard_events:
+    data_input_dir: "${input_dir}"
+
+hydra:
+  searchpath:
+    - pkg://MEDS_transforms.configs

MEDS_extract/configs/stage_configs/convert_to_sharded_events.yaml

@@ -0,0 +1,2 @@
+convert_to_sharded_events:
+  do_dedup_text_and_numeric: True

MEDS_extract/configs/stage_configs/extract_code_metadata.yaml

@@ -0,0 +1,3 @@
+extract_code_metadata:
+  is_metadata: true
+  description_separator: "\n"

MEDS_extract/configs/stage_configs/finalize_MEDS_data.yaml

@@ -0,0 +1,2 @@
+finalize_MEDS_data:
+  do_retype: true

MEDS_extract/configs/stage_configs/finalize_MEDS_metadata.yaml

@@ -0,0 +1,3 @@
+finalize_MEDS_metadata:
+  is_metadata: True
+  do_retype: True

MEDS_extract/configs/stage_configs/merge_to_MEDS_cohort.yaml

@@ -0,0 +1,3 @@
+merge_to_MEDS_cohort:
+  unique_by: "*"
+  additional_sort_by: null

MEDS_extract/configs/stage_configs/shard_events.yaml

@@ -0,0 +1,3 @@
+shard_events:
+  row_chunksize: 200000000
+  infer_schema_length: 10000

MEDS_extract/configs/stage_configs/split_and_shard_subjects.yaml

@@ -0,0 +1,9 @@
+split_and_shard_subjects:
+  is_metadata: True
+  output_dir: ${cohort_dir}/metadata
+  n_subjects_per_shard: 50000
+  external_splits_json_fp: null
+  split_fracs:
+    train: 0.8
+    tuning: 0.1
+    held_out: 0.1