virola 0.0.1__py3-none-any.whl

vir/__init__.py

@@ -0,0 +1 @@
+from vir import config  # noqa: F401

vir/adapters/__init__.py

@@ -0,0 +1,46 @@
+from vir.adapters.aktin import AktinAdapter
+from vir.adapters.base import Adapter
+from vir.adapters.datasus import DatasusAdapter
+
+ADAPTERS: dict[str, Adapter] = {
+    "datasus": DatasusAdapter(),
+    "aktin": AktinAdapter(),
+}
+
+
+def known_strata() -> set[str]:
+    """Return the union of strata declared by every registered adapter."""
+    return {stratum for adapter in ADAPTERS.values() for stratum in adapter.strata}
+
+
+def adapter_for_stratum(stratum: str) -> Adapter:
+    """Return the unique adapter that declares this stratum.
+
+    Raises:
+        LookupError: if no adapter claims the stratum, or if more than one does.
+    """
+    owners = [adapter for adapter in ADAPTERS.values() if stratum in adapter.strata]
+    if not owners:
+        raise LookupError(f"unknown stratum: {stratum!r}")
+    if len(owners) > 1:
+        raise LookupError(f"ambiguous stratum: {stratum!r} is claimed by {len(owners)} adapters")
+    return owners[0]
+
+
+def stratum_metadata(stratum: str) -> dict:
+    """Resolve ``{city, region, population_range}`` for a stratum or dataset slug.
+
+    Iterates the registered adapters and asks each to interpret the value.
+    The first adapter that recognises it (returns a non-null ``city``) wins.
+    Falls back to all-None when no adapter claims the value — matching the
+    historical behaviour of the framework's ``_city_metadata`` helper.
+
+    This accepts both display-form stratum names (e.g. ``"Belo Horizonte"``)
+    and dataset-name slugs (e.g. ``"belo_horizonte"``); each adapter is
+    responsible for normalising what it owns.
+    """
+    for adapter in ADAPTERS.values():
+        metadata = adapter.stratum_metadata(stratum)
+        if metadata.get("city") is not None:
+            return metadata
+    return {"city": None, "region": None, "population_range": None}

vir/adapters/aktin/__init__.py

@@ -0,0 +1,184 @@
+import io
+import os
+import zipfile
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Union
+
+import polars as pl
+from loguru import logger
+
+from vir.adapters.aktin.cleaning import clean_aktin_raw
+from vir.adapters.aktin.terminology import build_icd10gm_labels
+from vir.adapters.aktin.viz import AGE_GROUP_COLOR, CODE_TYPE_COLOR, SEX_COLOR
+from vir.adapters.base import Adapter
+
+_FileSource = Union[Path, tuple[Path, str]]
+
+
+@dataclass(frozen=True)
+class _Site:
+    """One AKTIN site's pair of source streams (either filesystem paths or zip members)."""
+
+    name: str
+    case_source: _FileSource
+    diag_source: _FileSource
+
+
+def _read_aktin_tsv(source: _FileSource) -> pl.DataFrame:
+    """Read a tab-separated AKTIN file either directly or from inside a zip.
+
+    Uses ``truncate_ragged_lines=True``: AKTIN free-text columns (e.g.
+    ``isolation_reason``, ``zuweisung``) occasionally contain embedded tabs in
+    real exports, producing rows with more fields than the header declares.
+    Truncating the trailing fields lets the row survive — the dropped data is
+    in narrative columns that we don't consume downstream.
+    """
+    read_kwargs = dict(separator="\t", infer_schema_length=0, truncate_ragged_lines=True)
+    if isinstance(source, tuple):
+        zip_path, member = source
+        with zipfile.ZipFile(zip_path) as archive, archive.open(member) as stream:
+            payload = stream.read()
+        return pl.read_csv(io.BytesIO(payload), **read_kwargs)
+    return pl.read_csv(source, **read_kwargs)
+
+
+def _discover_sites(raw_path: Path) -> list[_Site]:
+    """Find every AKTIN site under ``raw_path``.
+
+    Recognises two layouts: an unpacked ``<site>/`` containing both
+    ``case_data.txt`` and ``diag_data.txt``, and a ``<site>_result.zip``
+    whose members include both.
+
+    Raises ``FileNotFoundError`` for a site that has one of the two files but
+    not the other — silent orphan-site loss is the worst failure mode (every
+    encounter from that site would disappear at the inner-join).
+    """
+    sites: list[_Site] = []
+
+    for case_path in sorted(raw_path.glob("**/case_data.txt")):
+        diag_path = case_path.parent / "diag_data.txt"
+        if not diag_path.exists():
+            raise FileNotFoundError(
+                f"Site {case_path.parent.name!r} has case_data.txt but no diag_data.txt"
+            )
+        sites.append(
+            _Site(
+                name=case_path.parent.name or "root", case_source=case_path, diag_source=diag_path
+            )
+        )
+
+    for zip_path in sorted(raw_path.glob("**/*.zip")):
+        with zipfile.ZipFile(zip_path) as archive:
+            members = set(archive.namelist())
+        has_case = "case_data.txt" in members
+        has_diag = "diag_data.txt" in members
+        if not has_case and not has_diag:
+            continue
+        if has_case != has_diag:
+            raise FileNotFoundError(
+                f"Zip {zip_path.name!r} has only one of case_data.txt / diag_data.txt"
+            )
+        sites.append(
+            _Site(
+                name=zip_path.stem,
+                case_source=(zip_path, "case_data.txt"),
+                diag_source=(zip_path, "diag_data.txt"),
+            )
+        )
+
+    return sites
+
+
+def _namespace_encounter_num(frame: pl.DataFrame, site_name: str) -> pl.DataFrame:
+    """Prefix ``a_encounter_num`` with the site name so cross-site encounter-id collisions don't merge."""
+    return frame.with_columns(
+        a_encounter_num=pl.concat_str(
+            [pl.lit(f"{site_name}:"), pl.col("a_encounter_num").cast(pl.Utf8)]
+        )
+    )
+
+
+class AktinAdapter(Adapter):
+    """Adapter for the German AKTIN emergency-department dataset.
+
+    Single stratum (``"germany"`` — the country as a whole), per-encounter
+    diagnoses (one row per encounter × confirmed diagnosis after cleaning),
+    ICD-10-GM labels from the vendored BfArM 2025 release.
+    """
+
+    name = "aktin"
+    default_model = os.getenv(
+        "VIROLA_AKTIN_MODEL", "cambridgeltl/SapBERT-UMLS-2020AB-all-lang-from-XLMR"
+    )
+    terminology_csv_filename = "icd10gm_aktin.csv"
+
+    def clean(
+        self,
+        raw_path: Path,
+        output_path: Path,
+        stratum: str | None = None,
+        **options,
+    ) -> Path:
+        """Pool every AKTIN site under ``raw_path`` into a single canonical interim parquet.
+
+        Each site (unpacked directory or ``<site>_result.zip``) must contribute
+        both ``case_data.txt`` and ``diag_data.txt`` — orphan sites raise.
+        ``a_encounter_num`` is namespaced per site so id collisions don't merge.
+        """
+        sites = _discover_sites(raw_path)
+        if not sites:
+            raise FileNotFoundError(
+                f"No AKTIN sites found under {raw_path}. Expected ``<site>/case_data.txt`` + "
+                "``diag_data.txt`` pairs or ``<site>_result.zip`` archives."
+            )
+
+        logger.info(f"AKTIN clean: pooling {len(sites)} site(s) under {raw_path}")
+
+        case_frames = [
+            _namespace_encounter_num(_read_aktin_tsv(site.case_source), site.name)
+            for site in sites
+        ]
+        diag_frames = [
+            _namespace_encounter_num(_read_aktin_tsv(site.diag_source), site.name)
+            for site in sites
+        ]
+
+        cases = pl.concat(case_frames, how="diagonal_relaxed")
+        diagnoses = pl.concat(diag_frames, how="diagonal_relaxed")
+
+        cleaned = clean_aktin_raw(cases, diagnoses)
+        cleaned.write_parquet(output_path)
+        return output_path
+
+    def build_terminology_labels(self) -> pl.DataFrame:
+        return build_icd10gm_labels()
+
+    def stratum_metadata(self, stratum: str) -> dict:
+        if "germany" in stratum.lower():
+            return {"city": "germany", "region": None, "population_range": None}
+        return {"city": None, "region": None, "population_range": None}
+
+    @property
+    def strata(self) -> list[str]:
+        return ["germany"]
+
+    @property
+    def skip_patterns(self) -> list[tuple[str, str]]:
+        return [("icd10gm", r"^Z")]
+
+    @property
+    def code_types(self) -> list[str]:
+        return ["icd10gm"]
+
+    def cleaned_filename(self, stratum: str | None) -> str:
+        slug = (stratum or "all").lower().replace(" ", "_")
+        return f"cleaned_aktin_{slug}.parquet"
+
+    @property
+    def viz_palettes(self) -> dict[str, dict[str, str]]:
+        return {
+            "sex": SEX_COLOR,
+            "age_group": AGE_GROUP_COLOR,
+            "code_type": CODE_TYPE_COLOR,
+        }

vir/adapters/aktin/cleaning.py

@@ -0,0 +1,87 @@
+"""AKTIN raw → canonical interim transformation.
+
+Takes the per-encounter ``case_data`` and per-diagnosis ``diag_data`` tables
+exported from AKTIN and produces one row per (encounter × confirmed diagnosis)
+in the canonical interim shape that ``prepare_dataframe`` consumes.
+"""
+
+import polars as pl
+
+AKTIN_AGE_BUCKETS: dict[str, str] = {
+    "0-4": "0-4",
+    "00-04": "0-4",
+    "5-9": "5-9",
+    "05-09": "5-9",
+    "10-14": "10-14",
+    "15-19": "15-19",
+    "20-24": "20-39",
+    "25-29": "20-39",
+    "30-34": "20-39",
+    "35-39": "20-39",
+    "40-44": "40-59",
+    "45-49": "40-59",
+    "50-54": "40-59",
+    "55-59": "40-59",
+    "60-64": "60-79",
+    "65-69": "60-79",
+    "70-74": "60-79",
+    "75-79": "60-79",
+    "80+": "80+",
+}
+
+_REQUIRED_CASE_FIELDS: tuple[str, ...] = ("geschlecht", "aufnahme_ts", "age_group")
+
+
+def _null_when_empty(column: str) -> pl.Expr:
+    return pl.when(pl.col(column).eq("")).then(None).otherwise(pl.col(column)).alias(column)
+
+
+def clean_aktin_raw(cases: pl.DataFrame, diagnoses: pl.DataFrame) -> pl.DataFrame:
+    """Transform AKTIN ``case_data`` + ``diag_data`` into the canonical interim shape.
+
+    Steps:
+
+    1. Drop case rows with empty / null ``geschlecht``, ``aufnahme_ts`` or
+       ``age_group``.
+    2. Parse ``aufnahme_ts`` accepting both ``Z`` and ``+00:00`` offsets;
+       unparseable timestamps drop the row rather than ride through as a null
+       year/week.
+    3. Collapse the native 5-year ``age_group`` to one of the 8 AKTIN buckets;
+       any value outside the bucket map drops the row.
+    4. Keep only diagnoses with ``diagnose_zusatz == 'G'`` (confirmed).
+    5. Inner-join cases × confirmed diagnoses on ``a_encounter_num``.
+    6. Add ``city = "germany"``, ``code_type = "icd10gm"``, ``quantity = 1``.
+    7. Drop ``plz_kurz``. Rename ``geschlecht`` → ``sex``. Retain
+       ``a_encounter_num`` and the case_data extras (``triage``, vitals,
+       ``verbleib``, ...) for later statistics — they are dropped later by
+       ``prepare_dataframe`` when projecting onto the canonical row.
+    """
+    cleaned_cases = (
+        cases.with_columns([_null_when_empty(column) for column in _REQUIRED_CASE_FIELDS])
+        .drop_nulls(subset=list(_REQUIRED_CASE_FIELDS))
+        .with_columns(
+            _admission=pl.col("aufnahme_ts")
+            .str.replace(r"Z$", "+00:00")
+            .str.to_datetime("%Y-%m-%dT%H:%M:%S%z", strict=False),
+            age_group=pl.col("age_group").replace_strict(AKTIN_AGE_BUCKETS, default=None),
+        )
+        .drop_nulls(subset=["_admission", "age_group"])
+        .with_columns(
+            year=pl.col("_admission").dt.iso_year(),
+            week=pl.col("_admission").dt.week(),
+        )
+        .drop(["_admission", "plz_kurz"])
+        .rename({"geschlecht": "sex"})
+    )
+
+    confirmed_diagnoses = diagnoses.filter(pl.col("diagnose_zusatz").eq("G")).select(
+        code=pl.col("icd_code"),
+        a_encounter_num=pl.col("a_encounter_num"),
+    )
+
+    joined = cleaned_cases.join(confirmed_diagnoses, on="a_encounter_num", how="inner")
+    return joined.with_columns(
+        city=pl.lit("germany"),
+        code_type=pl.lit("icd10gm"),
+        quantity=pl.lit(1),
+    )

vir/adapters/aktin/notebooks/eda_aktin.py

@@ -0,0 +1,272 @@
+import marimo
+
+__generated_with = "0.23.4"
+app = marimo.App(width="medium")
+
+
+@app.cell
+def _():
+    import marimo as mo
+    import plotly.express as px
+    import polars as pl
+
+    from vir.config import INTERIM_DATA_DIR
+
+    return INTERIM_DATA_DIR, mo, pl, px
+
+
+@app.cell
+def _(mo):
+    mo.md(r"""
+    # AKTIN — Exploratory data analysis
+
+    Basic shape, demographic distribution, code coverage, and temporal distribution
+    of the cleaned AKTIN interim parquet (one row per encounter × confirmed
+    diagnosis).
+    """)
+    return
+
+
+@app.cell
+def _(INTERIM_DATA_DIR, mo):
+    available_files = sorted(INTERIM_DATA_DIR.glob("cleaned_aktin_*.parquet"))
+    file_dropdown = mo.ui.dropdown(
+        options={path.name: str(path) for path in available_files},
+        value=available_files[0].name if available_files else None,
+        label="Cleaned AKTIN parquet",
+    )
+    file_dropdown
+    return (file_dropdown,)
+
+
+@app.cell
+def _(file_dropdown, mo, pl):
+    mo.stop(
+        not file_dropdown.value,
+        mo.callout(
+            mo.md("No `cleaned_aktin_*.parquet` files found under `INTERIM_DATA_DIR`."),
+            kind="warn",
+        ),
+    )
+    df = pl.read_parquet(file_dropdown.value)
+    return (df,)
+
+
+@app.cell
+def _(mo):
+    mo.md(r"""
+    ## Shape
+    """)
+    return
+
+
+@app.cell
+def _(df, mo):
+    n_rows = df.height
+    n_encounters = df["a_encounter_num"].n_unique()
+    n_codes = df["code"].n_unique()
+    year_min = df["year"].min()
+    year_max = df["year"].max()
+    week_min = df["week"].min()
+    week_max = df["week"].max()
+    diagnoses_per_encounter = n_rows / n_encounters if n_encounters else 0
+
+    mo.md(f"""
+    | metric | value |
+    |---|---|
+    | rows (encounter × diagnosis) | {n_rows:,} |
+    | distinct encounters | {n_encounters:,} |
+    | distinct ICD-10-GM codes | {n_codes:,} |
+    | mean diagnoses per encounter | {diagnoses_per_encounter:.2f} |
+    | ISO year range | {year_min} – {year_max} |
+    | ISO week range | {week_min} – {week_max} |
+    """)
+    return
+
+
+@app.cell
+def _(mo):
+    mo.md(r"""
+    ## Demographics
+    """)
+    return
+
+
+@app.cell
+def _():
+    age_order = ["0-4", "5-9", "10-14", "15-19", "20-39", "40-59", "60-79", "80+"]
+    return (age_order,)
+
+
+@app.cell
+def _(df, mo, px):
+    sex_per_encounter = df.unique(subset=["a_encounter_num"]).group_by("sex").len().sort("sex")
+    sex_chart = px.bar(
+        sex_per_encounter,
+        x="sex",
+        y="len",
+        title="Encounters by sex",
+        labels={"len": "encounters"},
+    )
+    mo.ui.plotly(sex_chart)
+    return
+
+
+@app.cell
+def _(age_order, df, mo, pl, px):
+    age_per_encounter = (
+        df.unique(subset=["a_encounter_num"])
+        .group_by("age_group")
+        .len()
+        .with_columns(pl.col("age_group").cast(pl.Enum(age_order)).alias("age_group_ordered"))
+        .sort("age_group_ordered")
+    )
+    age_chart = px.bar(
+        age_per_encounter,
+        x="age_group",
+        y="len",
+        title="Encounters by age group",
+        category_orders={"age_group": age_order},
+        labels={"len": "encounters"},
+    )
+    mo.ui.plotly(age_chart)
+    return
+
+
+@app.cell
+def _(age_order, df, mo, px):
+    sex_age = df.unique(subset=["a_encounter_num"]).group_by(["age_group", "sex"]).len()
+    sex_age_chart = px.bar(
+        sex_age,
+        x="age_group",
+        y="len",
+        color="sex",
+        barmode="group",
+        title="Encounters by age group × sex",
+        category_orders={"age_group": age_order},
+        labels={"len": "encounters"},
+    )
+    mo.ui.plotly(sex_age_chart)
+    return
+
+
+@app.cell
+def _(mo):
+    mo.md(r"""
+    ## ICD-10-GM coverage
+    """)
+    return
+
+
+@app.cell
+def _(df, mo, pl, px):
+    chapter_distribution = (
+        df.with_columns(pl.col("code").str.slice(0, 1).alias("chapter_letter"))
+        .group_by("chapter_letter")
+        .len()
+        .sort("chapter_letter")
+    )
+    chapter_chart = px.bar(
+        chapter_distribution,
+        x="chapter_letter",
+        y="len",
+        title="Diagnoses by ICD-10-GM chapter letter",
+        labels={"len": "diagnoses", "chapter_letter": "chapter (A–Z, U)"},
+    )
+    mo.ui.plotly(chapter_chart)
+    return
+
+
+@app.cell
+def _(df, mo, px):
+    top_codes = df.group_by("code").len().sort("len", descending=True).head(20)
+    top_codes_chart = px.bar(
+        top_codes,
+        x="code",
+        y="len",
+        title="Top-20 ICD-10-GM codes by diagnosis count",
+        labels={"len": "diagnoses"},
+    )
+    mo.ui.plotly(top_codes_chart)
+    return (top_codes,)
+
+
+@app.cell
+def _(mo, top_codes):
+    mo.ui.table(top_codes, label="Top-20 codes")
+    return
+
+
+@app.cell
+def _(mo):
+    mo.md(r"""
+    ## Temporal distribution
+    """)
+    return
+
+
+@app.cell
+def _(df, pl):
+    has_aufnahme_ts = "aufnahme_ts" in df.columns
+    timestamps_df = (
+        df.with_columns(
+            pl.col("aufnahme_ts").str.to_datetime("%Y-%m-%dT%H:%M:%SZ").alias("_admission")
+        )
+        if has_aufnahme_ts
+        else df
+    )
+    return has_aufnahme_ts, timestamps_df
+
+
+@app.cell
+def _(has_aufnahme_ts, mo, pl, px, timestamps_df):
+    mo.stop(
+        not has_aufnahme_ts,
+        mo.callout(
+            mo.md(
+                "``aufnahme_ts`` was dropped from the interim parquet — month-level chart unavailable."
+            ),
+            kind="warn",
+        ),
+    )
+    monthly = (
+        timestamps_df.unique(subset=["a_encounter_num"])
+        .with_columns(pl.col("_admission").dt.strftime("%Y-%m").alias("year_month"))
+        .group_by("year_month")
+        .len()
+        .sort("year_month")
+    )
+    monthly_chart = px.bar(
+        monthly,
+        x="year_month",
+        y="len",
+        title="Encounters per month",
+        labels={"len": "encounters", "year_month": "year-month"},
+    )
+    mo.ui.plotly(monthly_chart)
+    return
+
+
+@app.cell
+def _(df, mo, px):
+    weekly = (
+        df.unique(subset=["a_encounter_num"])
+        .group_by(["year", "week"])
+        .len()
+        .sort(["year", "week"])
+    )
+    weekly_chart = px.scatter(
+        weekly,
+        x="week",
+        y="len",
+        color="year",
+        title="Encounters per ISO week, by year",
+        labels={"len": "encounters", "week": "ISO week"},
+    )
+    weekly_chart.update_layout(xaxis=dict(tickmode="linear", dtick=4))
+    mo.ui.plotly(weekly_chart)
+    return
+
+
+if __name__ == "__main__":
+    app.run()

vir/adapters/aktin/references/README.md

@@ -0,0 +1,18 @@
+# AKTIN adapter references
+
+## icd10gm2025syst_kodes.txt
+
+Source: BfArM (Bundesinstitut für Arzneimittel und Medizinprodukte), ICD-10-GM
+Version 2025 — Systematisches Verzeichnis (classification file
+`icd10gm2025syst-meta.zip` → `Klassifikationsdateien/icd10gm2025syst_kodes.txt`).
+
+Stand der Klassifikation: 2024-09-13. OID `1.2.276.0.76.5.548`.
+
+Format: semicolon-separated, no header, one row per code. ~16,800 rows (4 MB).
+
+Redistribution / attribution: subject to BfArM download conditions
+(`downloadbedingungen-2024.pdf`). License terms must be confirmed before
+making the public repo containing this file generally available — tracked
+in `AKTIN-PLAN.md § Stakeholder check-ins`.
+
+Parsing of this file lives in `vir/adapters/aktin/terminology.py`.