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

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

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

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

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

{ltc_code-0.1.5 → 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,97 +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",
-            ]
-        )
+        .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")
-    )
+    matched = stage.filter(pl.col("sid_cepr").is_not_null())
+    unmatched = stage.filter(pl.col("sid_cepr").is_null()).drop("sid_cepr")
 
-    added = len(matched)
-
-    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
 
-    current = (
-        frame.collect()
-        if is_lazy
-        else frame
-    )
+    input_columns = current.columns
 
-    #
-    # CLEAN LEFT SIDE
-    #
+    current = current.with_row_index("_row_id")
+
+    current = current.with_columns(
+        *clean_name(cols["fname"]),
+        *clean_name(cols["lname"]),
+    )
 
     current = (
         current
+        .with_columns(*clean_dob(col=cols["dob"]))
         .with_columns(
-            *clean_name(cols["fname"]),
-            *clean_name(cols["lname"]),
-            *clean_dob(col=cols["dob"]),
+            _parse_dob_expr(f"{cols['dob']}_clean").alias(f"{cols['dob']}_clean")
         )
     )
 
     matched_frames = []
-
     unmatched = current
 
-    #
-    # STAGE 1
-    # EXACT
-    #
-
     matched, unmatched = _run_match_stage(
         unmatched,
         lookup=lookups["exact"],
@@ -1750,102 +1700,144 @@ def lookup_sid_cepr(
         dob_expr=pl.col(f"{cols['dob']}_clean"),
         label="exact",
     )
-
     matched_frames.append(matched)
 
-    #
-    # STAGE 2
-    # MNAME
-    #
-
-    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=lookups["mname"],
-            fname_expr=pl.concat_str(
-                [
-                    pl.col(f"{cols['fname']}_clean"),
-                    pl.col(f"{cols['mname']}_clean"),
-                ],
-                separator=" ",
-            ),
+            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="mname",
+            label=label,
         )
-
         matched_frames.append(matched)
 
-    #
-    # STAGE 3
-    # SUFFIX
-    #
+    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",
+                pl.concat_str(
+                    [pl.col(f"{cols['fname']}_clean"), pl.col(f"{cols['mname']}_clean")],
+                    separator=" ",
+                ),
+                pl.col(f"{cols['lname']}_clean"),
+            ),
+            (
+                "left mname + lname -> right exact",
+                pl.col(f"{cols['fname']}_clean"),
+                pl.concat_str(
+                    [pl.col(f"{cols['mname']}_clean"), pl.col(f"{cols['lname']}_clean")],
+                    separator=" ",
+                ),
+            ),
+            (
+                "left fname + mname no space -> right exact",
+                pl.concat_str(
+                    [pl.col(f"{cols['fname']}_clean"), pl.col(f"{cols['mname']}_clean")],
+                    separator="",
+                ),
+                pl.col(f"{cols['lname']}_clean"),
+            ),
+        ]:
+            matched, unmatched = _run_match_stage(
+                unmatched,
+                lookup=lookups["exact"],
+                fname_expr=fname_expr,
+                lname_expr=lname_expr,
+                dob_expr=pl.col(f"{cols['dob']}_clean"),
+                label=label,
+            )
+            matched_frames.append(matched)
 
     if "suffix" in cols:
+        unmatched = unmatched.with_columns(*clean_other_name(cols["suffix"]))
 
-        unmatched = (
-            unmatched
-            .with_columns(
-                *clean_other_name(cols["suffix"])
+        for label, fname_expr, lname_expr in [
+            (
+                "left fname + suffix -> right exact",
+                pl.concat_str(
+                    [pl.col(f"{cols['fname']}_clean"), pl.col(f"{cols['suffix']}_clean")],
+                    separator=" ",
+                ),
+                pl.col(f"{cols['lname']}_clean"),
+            ),
+            (
+                "left fname + suffix no space -> right exact",
+                pl.concat_str(
+                    [pl.col(f"{cols['fname']}_clean"), pl.col(f"{cols['suffix']}_clean")],
+                    separator="",
+                ),
+                pl.col(f"{cols['lname']}_clean"),
+            ),
+            (
+                "left lname + suffix -> right exact",
+                pl.col(f"{cols['fname']}_clean"),
+                pl.concat_str(
+                    [pl.col(f"{cols['lname']}_clean"), pl.col(f"{cols['suffix']}_clean")],
+                    separator=" ",
+                ),
+            ),
+            (
+                "left lname + suffix no space -> right exact",
+                pl.col(f"{cols['fname']}_clean"),
+                pl.concat_str(
+                    [pl.col(f"{cols['lname']}_clean"), pl.col(f"{cols['suffix']}_clean")],
+                    separator="",
+                ),
+            ),
+        ]:
+            matched, unmatched = _run_match_stage(
+                unmatched,
+                lookup=lookups["exact"],
+                fname_expr=fname_expr,
+                lname_expr=lname_expr,
+                dob_expr=pl.col(f"{cols['dob']}_clean"),
+                label=label,
             )
-        )
-
+            matched_frames.append(matched)
+
+    for key in [
+        "dob_imp",
+        "dob_imp_minus_1",
+        "dob_imp_plus_1",
+        "dob_imp_minus_2",
+        "dob_imp_plus_2",
+    ]:
         matched, unmatched = _run_match_stage(
             unmatched,
-            lookup=lookups["suffix"],
+            lookup=lookups[key],
             fname_expr=pl.col(f"{cols['fname']}_clean"),
-            lname_expr=pl.concat_str(
-                [
-                    pl.col(f"{cols['lname']}_clean"),
-                    pl.col(f"{cols['suffix']}_clean"),
-                ],
-                separator=" ",
-            ),
+            lname_expr=pl.col(f"{cols['lname']}_clean"),
             dob_expr=pl.col(f"{cols['dob']}_clean"),
-            label="suffix",
+            label=key,
         )
-
         matched_frames.append(matched)
 
-    #
-    # STAGE 4
-    # DOB IMP
-    #
-
-    matched, unmatched = _run_match_stage(
-        unmatched,
-        lookup=lookups["dob_imp"],
-        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="dob_imp",
+    result = (
+        pl.concat(matched_frames + [unmatched], how="diagonal_relaxed")
+        .sort("_row_id")
+        .drop("_row_id")
     )
 
-    matched_frames.append(matched)
+    if "sid_cepr" not in input_columns:
+        input_columns = input_columns + ["sid_cepr"]
 
-    #
-    # FINAL
-    #
+    result = result.select(input_columns)
 
-    result = pl.concat(
-        matched_frames + [unmatched],
-        how="diagonal_relaxed",
-    )
-
-    print(
-        f"final matched: {result['sid_cepr'].is_not_null().sum():,}/{len(result):,}"
-    )
+    print(f"final matched: {result['sid_cepr'].is_not_null().sum():,}/{len(result):,}")
 
     return result
-
 #
 
 # EXAMPLE USAGE