ltc-code 0.1.6__tar.gz → 0.1.7__tar.gz

{ltc_code-0.1.6 → ltc_code-0.1.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: ltc-code
-Version: 0.1.6
+Version: 0.1.7
 Summary: Add your description here
 Requires-Python: >=3.9
 Description-Content-Type: text/markdown

{ltc_code-0.1.6 → ltc_code-0.1.7}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "ltc-code"
-version = "0.1.6"
+version = "0.1.7"
 description = "Add your description here"
 readme = "README.md"
 requires-python = ">=3.9"

{ltc_code-0.1.6 → ltc_code-0.1.7}/src/ltc_code/may27.py

@@ -1427,6 +1427,16 @@ def lookup_sid_cepr(
 
 
 
+def _parse_dob_expr(col: str) -> pl.Expr:
+    return pl.coalesce(
+        [
+            pl.col(col).cast(pl.String).str.strptime(pl.Date, "%m/%d/%Y", strict=False),
+            pl.col(col).cast(pl.String).str.strptime(pl.Date, "%m/%d/%y", strict=False),
+            pl.col(col).cast(pl.String).str.strptime(pl.Date, "%Y-%m-%d", strict=False),
+        ]
+    )
+
+
 def _build_lookup(
     census: pl.DataFrame,
     *,
@@ -1435,80 +1445,35 @@ def _build_lookup(
     dob_col: str,
     label: str,
 ) -> pl.DataFrame:
-    """
-    Build a deterministic SID lookup table.
-
-    Output schema:
-        _fname_key
-        _lname_key
-        _dob_key
-        sid_cepr
-
-    Ambiguous keys are removed.
-    """
-
     lookup = (
         census
         .select(
-            [
-                fname_expr.alias("_fname_key"),
-                lname_expr.alias("_lname_key"),
-                pl.col(dob_col).alias("_dob_key"),
-                pl.col("sid_cepr"),
-            ]
-        )
-        .drop_nulls(
-            [
-                "_fname_key",
-                "_lname_key",
-                "_dob_key",
-                "sid_cepr",
-            ]
-        )
-        .group_by(
-            [
-                "_fname_key",
-                "_lname_key",
-                "_dob_key",
-            ]
-        )
-        .agg(
-            pl.col("sid_cepr").unique().alias("_sids")
-        )
-        .with_columns(
-            pl.col("_sids").list.len().alias("_sid_count")
-        )
-        .filter(
-            pl.col("_sid_count") == 1
+            fname_expr.alias("_fname_key"),
+            lname_expr.alias("_lname_key"),
+            pl.col(dob_col).alias("_dob_key"),
+            pl.col("sid_cepr"),
         )
+        .drop_nulls(["_fname_key", "_lname_key", "_dob_key", "sid_cepr"])
+        .group_by(["_fname_key", "_lname_key", "_dob_key"])
+        .agg(pl.col("sid_cepr").unique().alias("_sids"))
+        .with_columns(pl.col("_sids").list.len().alias("_sid_count"))
+        .filter(pl.col("_sid_count") == 1)
         .select(
-            [
-                "_fname_key",
-                "_lname_key",
-                "_dob_key",
-                pl.col("_sids").list.first().alias("sid_cepr"),
-            ]
+            "_fname_key",
+            "_lname_key",
+            "_dob_key",
+            pl.col("_sids").list.first().alias("sid_cepr"),
         )
     )
 
-    print(f"built lookup: {label}")
-
+    print(f"built lookup: {label} ({len(lookup):,} usable keys)")
     return lookup
 
-def build_census_lookups(
-    *,
-    cmo_name: str,
-) -> dict[str, pl.DataFrame]:
-
-    try:
-        import mappings
-    except ImportError:
-        import mapppings as mappings
 
+def build_census_lookups(*, cmo_name: str) -> dict[str, pl.DataFrame]:
     annual_frames = []
 
     for year in range(1994, 2023):
-
         path = CENSUS_STUDENTS / f"census_student_{year}.csv"
 
         annual = (
@@ -1520,16 +1485,14 @@ def build_census_lookups(
                 ignore_errors=False,
             )
             .select(
-                [
-                    "cmo_code",
-                    "sid_cepr",
-                    "fname_clean",
-                    "lname_clean",
-                    "mname_clean",
-                    "suff_clean",
-                    "birthdate_clean",
-                    "birthdate_imp",
-                ]
+                "cmo_code",
+                "sid_cepr",
+                "fname_clean",
+                "lname_clean",
+                "mname_clean",
+                "suff_clean",
+                "birthdate_clean",
+                "birthdate_imp",
             )
             .rename(
                 {
@@ -1542,92 +1505,105 @@ def build_census_lookups(
                 }
             )
             .with_columns(
-                pl.col("cmo_code")
-                .replace(mappings.CMO_CODE_TO_NAME)
-                .alias("cmo_name")
-            )
-            .filter(
-                pl.col("cmo_name") == cmo_name
+                pl.col("cmo_code").replace(cmo_map).alias("cmo_name")
             )
+            .filter(pl.col("cmo_name") == cmo_name)
             .with_columns(
                 *clean_name("fname"),
                 *clean_name("lname"),
                 *clean_other_name("mname"),
                 *clean_other_name("suffix"),
-                *clean_dob(col="dob"),
-                *clean_dob(col="dob_imp"),
             )
-            .drop(
-                [
-                    "dob",
-                    "dob_imp",
-                ]
+            .with_columns(*clean_dob(col="dob"))
+            .with_columns(*clean_dob(col="dob_imp"))
+            .with_columns(
+                _parse_dob_expr("dob_clean").alias("dob"),
+                _parse_dob_expr("dob_imp_clean").alias("dob_imp"),
             )
-            .rename(
-                {
-                    "dob_clean": "dob",
-                    "dob_imp_clean": "dob_imp",
-                }
+            .select(
+                "sid_cepr",
+                pl.col("fname_clean").alias("fname"),
+                pl.col("lname_clean").alias("lname"),
+                pl.col("mname_clean").alias("mname"),
+                pl.col("suffix_clean").alias("suffix"),
+                "dob",
+                "dob_imp",
             )
         )
 
         annual_frames.append(annual)
 
-    #
-    # MATERIALIZE ONCE
-    #
+    census = pl.concat(annual_frames, how="vertical_relaxed").collect()
 
-    census = (
-        pl.concat(
-            annual_frames,
-            how="vertical_relaxed",
-        )
-        .collect()
+    print(f"census rows after CMO filter: {len(census):,}")
+
+    lookups = {}
+
+    lookups["exact"] = _build_lookup(
+        census,
+        fname_expr=pl.col("fname"),
+        lname_expr=pl.col("lname"),
+        dob_col="dob",
+        label="exact",
+    )
+
+    lookups["mname"] = _build_lookup(
+        census,
+        fname_expr=pl.concat_str([pl.col("fname"), pl.col("mname")], separator=" "),
+        lname_expr=pl.col("lname"),
+        dob_col="dob",
+        label="right fname + mname",
     )
 
-    print(f"census rows: {len(census):,}")
+    lookups["mname_lname"] = _build_lookup(
+        census,
+        fname_expr=pl.col("fname"),
+        lname_expr=pl.concat_str([pl.col("mname"), pl.col("lname")], separator=" "),
+        dob_col="dob",
+        label="right mname + lname",
+    )
 
-    #
-    # BUILD LOOKUPS ONCE
-    #
+    lookups["mname_nospace"] = _build_lookup(
+        census,
+        fname_expr=pl.concat_str([pl.col("fname"), pl.col("mname")], separator=""),
+        lname_expr=pl.col("lname"),
+        dob_col="dob",
+        label="right fname + mname no space",
+    )
 
-    lookup_exact = _build_lookup(
+    lookups["suffix"] = _build_lookup(
         census,
         fname_expr=pl.col("fname"),
+        lname_expr=pl.concat_str([pl.col("lname"), pl.col("suffix")], separator=" "),
+        dob_col="dob",
+        label="right lname + suffix",
+    )
+
+    lookups["suffix_fname"] = _build_lookup(
+        census,
+        fname_expr=pl.concat_str([pl.col("fname"), pl.col("suffix")], separator=" "),
         lname_expr=pl.col("lname"),
         dob_col="dob",
-        label="exact",
+        label="right fname + suffix",
     )
 
-    lookup_mname = _build_lookup(
+    lookups["suffix_fname_nospace"] = _build_lookup(
         census,
-        fname_expr=pl.concat_str(
-            [
-                pl.col("fname"),
-                pl.col("mname"),
-            ],
-            separator=" ",
-        ),
+        fname_expr=pl.concat_str([pl.col("fname"), pl.col("suffix")], separator=""),
         lname_expr=pl.col("lname"),
         dob_col="dob",
-        label="mname",
+        label="right fname + suffix no space",
     )
 
-    lookup_suffix = _build_lookup(
+    lookups["suffix_lname_nospace"] = _build_lookup(
         census,
         fname_expr=pl.col("fname"),
-        lname_expr=pl.concat_str(
-            [
-                pl.col("lname"),
-                pl.col("suffix"),
-            ],
-            separator=" ",
-        ),
+        lname_expr=pl.concat_str([pl.col("lname"), pl.col("suffix")], separator=""),
         dob_col="dob",
-        label="suffix",
+        label="right lname + suffix no space",
     )
 
-    lookup_dob_imp = _build_lookup(
+    lookups["dob_imp"] = _build_lookup(
         census,
         fname_expr=pl.col("fname"),
         lname_expr=pl.col("lname"),
@@ -1635,12 +1611,22 @@ def build_census_lookups(
         label="dob_imp",
     )
 
-    return {
-        "exact": lookup_exact,
-        "mname": lookup_mname,
-        "suffix": lookup_suffix,
-        "dob_imp": lookup_dob_imp,
-    }
+    for offset, key in [
+        ("-1y", "dob_imp_minus_1"),
+        ("1y", "dob_imp_plus_1"),
+        ("-2y", "dob_imp_minus_2"),
+        ("2y", "dob_imp_plus_2"),
+    ]:
+        lookups[key] = _build_lookup(
+            census.with_columns(pl.col("dob_imp").dt.offset_by(offset).alias(key)),
+            fname_expr=pl.col("fname"),
+            lname_expr=pl.col("lname"),
+            dob_col=key,
+            label=key,
+        )
+
+    return lookups
+
 
 def _run_match_stage(
     unmatched: pl.DataFrame,
@@ -1651,84 +1637,61 @@ def _run_match_stage(
     dob_expr: pl.Expr,
     label: str,
 ) -> tuple[pl.DataFrame, pl.DataFrame]:
-
     before = len(unmatched)
 
     stage = (
         unmatched
         .with_columns(
-            [
-                fname_expr.alias("_fname_key"),
-                lname_expr.alias("_lname_key"),
-                dob_expr.alias("_dob_key"),
-            ]
+            fname_expr.alias("_fname_key"),
+            lname_expr.alias("_lname_key"),
+            dob_expr.alias("_dob_key"),
         )
         .join(
             lookup,
-            on=[
-                "_fname_key",
-                "_lname_key",
-                "_dob_key",
-            ],
+            on=["_fname_key", "_lname_key", "_dob_key"],
             how="left",
             validate="m:1",
         )
-        .drop(
-            [
-                "_fname_key",
-                "_lname_key",
-                "_dob_key",
-            ]
-        )
-    )
-
-    matched = (
-        stage
-        .filter(
-            pl.col("sid_cepr").is_not_null()
-        )
-    )
-
-    unmatched = (
-        stage
-        .filter(
-            pl.col("sid_cepr").is_null()
-        )
-        .drop("sid_cepr")
+        .drop(["_fname_key", "_lname_key", "_dob_key"])
     )
 
-    added = len(matched)
+    matched = stage.filter(pl.col("sid_cepr").is_not_null())
+    unmatched = stage.filter(pl.col("sid_cepr").is_null()).drop("sid_cepr")
 
-    print(
-        f"{label}: matched {added:,}/{before:,}"
-    )
+    print(f"{label}: matched {len(matched):,}/{before:,}")
 
     return matched, unmatched
 
+
 def lookup_sid_cepr(
-    frame: Frame,
+    frame,
     *,
     cols: Mapping[str, str],
     lookups: dict[str, pl.DataFrame],
-) -> Frame:
-
+):
     is_lazy = isinstance(frame, pl.LazyFrame)
     current = frame.collect() if is_lazy else frame
 
+    input_columns = current.columns
+
+    current = current.with_row_index("_row_id")
+
     current = current.with_columns(
         *clean_name(cols["fname"]),
         *clean_name(cols["lname"]),
-        *clean_dob(col=cols["dob"]),
     )
 
-    current = current.with_columns(
-        _parse_dob_expr(f"{cols['dob']}_clean").alias(f"{cols['dob']}_clean")
+    current = (
+        current
+        .with_columns(*clean_dob(col=cols["dob"]))
+        .with_columns(
+            _parse_dob_expr(f"{cols['dob']}_clean").alias(f"{cols['dob']}_clean")
+        )
     )
 
     matched_frames = []
     unmatched = current
 
-    # exact
     matched, unmatched = _run_match_stage(
         unmatched,
         lookup=lookups["exact"],
@@ -1739,34 +1702,31 @@ def lookup_sid_cepr(
     )
     matched_frames.append(matched)
 
-    # middle-name variants
-    if "mname" in cols:
-        unmatched = unmatched.with_columns(
-            *clean_other_name(cols["mname"])
+    for label, lookup in [
+        ("left exact -> right fname + mname", lookups["mname"]),
+        ("left exact -> right mname + lname", lookups["mname_lname"]),
+        ("left exact -> right fname + mname no space", lookups["mname_nospace"]),
+        ("left exact -> right lname + suffix", lookups["suffix"]),
+        ("left exact -> right fname + suffix", lookups["suffix_fname"]),
+        ("left exact -> right fname + suffix no space", lookups["suffix_fname_nospace"]),
+        ("left exact -> right lname + suffix no space", lookups["suffix_lname_nospace"]),
+    ]:
+        matched, unmatched = _run_match_stage(
+            unmatched,
+            lookup=lookup,
+            fname_expr=pl.col(f"{cols['fname']}_clean"),
+            lname_expr=pl.col(f"{cols['lname']}_clean"),
+            dob_expr=pl.col(f"{cols['dob']}_clean"),
+            label=label,
         )
+        matched_frames.append(matched)
 
-        mname_stages = [
-            (
-                "left exact -> right fname + mname",
-                lookups["mname"],
-                pl.col(f"{cols['fname']}_clean"),
-                pl.col(f"{cols['lname']}_clean"),
-            ),
-            (
-                "left exact -> right mname + lname",
-                lookups["mname_lname"],
-                pl.col(f"{cols['fname']}_clean"),
-                pl.col(f"{cols['lname']}_clean"),
-            ),
-            (
-                "left exact -> right fname + mname no space",
-                lookups["mname_nospace"],
-                pl.col(f"{cols['fname']}_clean"),
-                pl.col(f"{cols['lname']}_clean"),
-            ),
+    if "mname" in cols:
+        unmatched = unmatched.with_columns(*clean_other_name(cols["mname"]))
+
+        for label, fname_expr, lname_expr in [
             (
                 "left fname + mname -> right exact",
-                lookups["exact"],
                 pl.concat_str(
                     [pl.col(f"{cols['fname']}_clean"), pl.col(f"{cols['mname']}_clean")],
                     separator=" ",
@@ -1775,7 +1735,6 @@ def lookup_sid_cepr(
             ),
             (
                 "left mname + lname -> right exact",
-                lookups["exact"],
                 pl.col(f"{cols['fname']}_clean"),
                 pl.concat_str(
                     [pl.col(f"{cols['mname']}_clean"), pl.col(f"{cols['lname']}_clean")],
@@ -1784,19 +1743,16 @@ def lookup_sid_cepr(
             ),
             (
                 "left fname + mname no space -> right exact",
-                lookups["exact"],
                 pl.concat_str(
                     [pl.col(f"{cols['fname']}_clean"), pl.col(f"{cols['mname']}_clean")],
                     separator="",
                 ),
                 pl.col(f"{cols['lname']}_clean"),
             ),
-        ]
-
-        for label, lookup, fname_expr, lname_expr in mname_stages:
+        ]:
             matched, unmatched = _run_match_stage(
                 unmatched,
-                lookup=lookup,
+                lookup=lookups["exact"],
                 fname_expr=fname_expr,
                 lname_expr=lname_expr,
                 dob_expr=pl.col(f"{cols['dob']}_clean"),
@@ -1804,40 +1760,12 @@ def lookup_sid_cepr(
             )
             matched_frames.append(matched)
 
-    # suffix variants
     if "suffix" in cols:
-        unmatched = unmatched.with_columns(
-            *clean_other_name(cols["suffix"])
-        )
+        unmatched = unmatched.with_columns(*clean_other_name(cols["suffix"]))
 
-        suffix_stages = [
-            (
-                "left exact -> right lname + suffix",
-                lookups["suffix"],
-                pl.col(f"{cols['fname']}_clean"),
-                pl.col(f"{cols['lname']}_clean"),
-            ),
-            (
-                "left exact -> right fname + suffix",
-                lookups["suffix_fname"],
-                pl.col(f"{cols['fname']}_clean"),
-                pl.col(f"{cols['lname']}_clean"),
-            ),
-            (
-                "left exact -> right fname + suffix no space",
-                lookups["suffix_fname_nospace"],
-                pl.col(f"{cols['fname']}_clean"),
-                pl.col(f"{cols['lname']}_clean"),
-            ),
-            (
-                "left exact -> right lname + suffix no space",
-                lookups["suffix_lname_nospace"],
-                pl.col(f"{cols['fname']}_clean"),
-                pl.col(f"{cols['lname']}_clean"),
-            ),
+        for label, fname_expr, lname_expr in [
             (
                 "left fname + suffix -> right exact",
-                lookups["exact"],
                 pl.concat_str(
                     [pl.col(f"{cols['fname']}_clean"), pl.col(f"{cols['suffix']}_clean")],
                     separator=" ",
@@ -1846,7 +1774,6 @@ def lookup_sid_cepr(
             ),
             (
                 "left fname + suffix no space -> right exact",
-                lookups["exact"],
                 pl.concat_str(
                     [pl.col(f"{cols['fname']}_clean"), pl.col(f"{cols['suffix']}_clean")],
                     separator="",
@@ -1855,7 +1782,6 @@ def lookup_sid_cepr(
             ),
             (
                 "left lname + suffix -> right exact",
-                lookups["exact"],
                 pl.col(f"{cols['fname']}_clean"),
                 pl.concat_str(
                     [pl.col(f"{cols['lname']}_clean"), pl.col(f"{cols['suffix']}_clean")],
@@ -1864,19 +1790,16 @@ def lookup_sid_cepr(
             ),
             (
                 "left lname + suffix no space -> right exact",
-                lookups["exact"],
                 pl.col(f"{cols['fname']}_clean"),
                 pl.concat_str(
                     [pl.col(f"{cols['lname']}_clean"), pl.col(f"{cols['suffix']}_clean")],
                     separator="",
                 ),
             ),
-        ]
-
-        for label, lookup, fname_expr, lname_expr in suffix_stages:
+        ]:
             matched, unmatched = _run_match_stage(
                 unmatched,
-                lookup=lookup,
+                lookup=lookups["exact"],
                 fname_expr=fname_expr,
                 lname_expr=lname_expr,
                 dob_expr=pl.col(f"{cols['dob']}_clean"),
@@ -1884,16 +1807,13 @@ def lookup_sid_cepr(
             )
             matched_frames.append(matched)
 
-    # dob_imp variants
-    dob_imp_stages = [
+    for key in [
         "dob_imp",
         "dob_imp_minus_1",
         "dob_imp_plus_1",
         "dob_imp_minus_2",
         "dob_imp_plus_2",
-    ]
-
-    for key in dob_imp_stages:
+    ]:
         matched, unmatched = _run_match_stage(
             unmatched,
             lookup=lookups[key],
@@ -1904,17 +1824,20 @@ def lookup_sid_cepr(
         )
         matched_frames.append(matched)
 
-    result = pl.concat(
-        matched_frames + [unmatched],
-        how="diagonal_relaxed",
+    result = (
+        pl.concat(matched_frames + [unmatched], how="diagonal_relaxed")
+        .sort("_row_id")
+        .drop("_row_id")
     )
 
-    print(
-        f"final matched: {result['sid_cepr'].is_not_null().sum():,}/{len(result):,}"
-    )
+    if "sid_cepr" not in input_columns:
+        input_columns = input_columns + ["sid_cepr"]
 
-    return result
+    result = result.select(input_columns)
 
+    print(f"final matched: {result['sid_cepr'].is_not_null().sum():,}/{len(result):,}")
+
+    return result
 #
 
 # EXAMPLE USAGE