ltc-code 0.1.95__tar.gz → 0.1.96__tar.gz

{ltc_code-0.1.95 → ltc_code-0.1.96}/PKG-INFO

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

{ltc_code-0.1.95 → ltc_code-0.1.96}/pyproject.toml

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

ltc_code-0.1.96/src/ltc_code/june13.py

@@ -0,0 +1,294 @@
+from typing import Callable, Optional, Sequence, Union
+
+import polars as pl
+
+
+def impute_offers(
+    frame: Union[pl.DataFrame, pl.LazyFrame],
+    offer_col: str,
+    lottery_number_col: Optional[str] = None,
+    waitlist_number_col: Optional[str] = None,
+    enrollment_col: Optional[str] = None,
+    lottery_imputation: Optional[Union[pl.Expr, Callable[[Frame], pl.Expr]]] = None,
+    waitlist_imputation: Optional[Union[pl.Expr, Callable[[Frame], pl.Expr]]] = None,
+    enrollment_imputation: Optional[Union[pl.Expr, Callable[[Frame], pl.Expr]]] = None,
+    sibling_priority_cols: Optional[Sequence[str]] = None,
+) -> Union[pl.DataFrame, pl.LazyFrame]:
+    """
+    Impute offer indicators in a Polars DataFrame or LazyFrame.
+
+    The function treats null offers as 0. It sets ``offer_col`` to 1 for rows
+    that satisfy any supplied or default imputation rule, and creates:
+
+    - ``offer_imputed_ln``: newly imputed from lottery number information
+    - ``offer_imputed_wn``: newly imputed from waitlist number information
+    - ``offer_imputed_enr``: newly imputed from enrollment information
+
+    Flags are only set for rows whose original offer was 0 or null. Rows that
+    already had an observed offer are never flagged. Multiple flags may be 1
+    for the same row if multiple rules independently imply an offer.
+
+    Default rules:
+    - Lottery: find the highest lottery number among observed offers, excluding
+      sibling-priority offers when ``sibling_priority_cols`` is supplied. Impute
+      offers for all originally non-offered rows with lottery numbers at or
+      below that cutoff.
+    - Enrollment: impute offers for originally non-offered enrolled rows.
+    - Waitlist:
+        * If enrollment information is supplied, first infer offers from
+          enrollment, then use the highest waitlist number among those
+          enrollment-implied offers as the waitlist cutoff.
+        * If enrollment information is not supplied, use the highest waitlist
+          number among observed offers as the waitlist cutoff.
+      Then impute offers for all originally non-offered rows with waitlist
+      numbers at or below the cutoff.
+
+    Custom rules may be supplied as Polars expressions or as callables that
+    accept the frame and return a Polars expression.
+
+    Example
+    -------
+    >>> df = pl.DataFrame({
+    ...     "offer": [1, 0, 0, None],
+    ...     "lottery_num": [10, 11, 12, 13],
+    ...     "waitlist_num": [None, 1, 2, 3],
+    ...     "enrolled": [0, 0, 1, 0],
+    ...     "sib_1": [0, 0, 0, 0],
+    ...     "sib_2": [0, 0, 0, 0],
+    ... })
+    >>> impute_offers(
+    ...     df,
+    ...     offer_col="offer",
+    ...     lottery_number_col="lottery_num",
+    ...     waitlist_number_col="waitlist_num",
+    ...     enrollment_col="enrolled",
+    ...     sibling_priority_cols=["sib_1", "sib_2"],
+    ... )
+    shape: (4, 9)
+    ...
+    """
+    if not isinstance(frame, (pl.DataFrame, pl.LazyFrame)):
+        raise TypeError("frame must be a Polars DataFrame or LazyFrame.")
+
+    original_offer = "__original_offer_for_imputation__"
+
+    def as_rule_expr(rule: Union[pl.Expr, Callable[[Frame], pl.Expr]]) -> pl.Expr:
+        expr = rule(frame) if callable(rule) else rule
+        if not isinstance(expr, pl.Expr):
+            raise TypeError("Imputation rules must be Polars expressions.")
+        return expr.fill_null(False)
+
+    def any_sibling_priority(cols: Optional[Sequence[str]]) -> pl.Expr:
+        if not cols:
+            return pl.lit(False)
+
+        priority_expr = pl.lit(False)
+        for col in cols:
+            priority_expr = priority_expr | (pl.col(col).fill_null(0) == 1)
+        return priority_expr
+
+    work = frame.with_columns(pl.col(offer_col).fill_null(0).alias(original_offer))
+    originally_no_offer = pl.col(original_offer) == 0
+
+    lottery_rule = None
+    if lottery_imputation is not None:
+        lottery_rule = as_rule_expr(lottery_imputation)
+    elif lottery_number_col is not None:
+        observed_regular_offer = (
+            (pl.col(original_offer) == 1)
+            & ~any_sibling_priority(sibling_priority_cols)
+        )
+
+        lottery_cutoff = (
+            pl.col(lottery_number_col)
+            .filter(observed_regular_offer)
+            .max()
+        )
+
+        lottery_rule = (
+            pl.col(lottery_number_col).is_not_null()
+            & (pl.col(lottery_number_col) <= lottery_cutoff)
+        ).fill_null(False)
+
+    enrollment_rule = None
+    if enrollment_imputation is not None:
+        enrollment_rule = as_rule_expr(enrollment_imputation)
+    elif enrollment_col is not None:
+        enrollment_rule = (pl.col(enrollment_col).fill_null(0) == 1).fill_null(False)
+
+    waitlist_rule = None
+    if waitlist_imputation is not None:
+        waitlist_rule = as_rule_expr(waitlist_imputation)
+    elif waitlist_number_col is not None:
+        if enrollment_rule is not None:
+            waitlist_cutoff_source = originally_no_offer & enrollment_rule
+        else:
+            waitlist_cutoff_source = pl.col(original_offer) == 1
+
+        waitlist_cutoff = (
+            pl.col(waitlist_number_col)
+            .filter(waitlist_cutoff_source)
+            .max()
+        )
+
+        waitlist_rule = (
+            pl.col(waitlist_number_col).is_not_null()
+            & (pl.col(waitlist_number_col) <= waitlist_cutoff)
+        ).fill_null(False)
+
+    false_expr = pl.lit(False)
+    lottery_rule = lottery_rule if lottery_rule is not None else false_expr
+    waitlist_rule = waitlist_rule if waitlist_rule is not None else false_expr
+    enrollment_rule = enrollment_rule if enrollment_rule is not None else false_expr
+
+    ln_flag = originally_no_offer & lottery_rule
+    enr_flag = originally_no_offer & enrollment_rule
+    wn_flag = originally_no_offer & waitlist_rule
+
+    any_imputation = ln_flag | wn_flag | enr_flag
+
+    return (
+        work.with_columns(
+            [
+                pl.when(ln_flag)
+                .then(1)
+                .otherwise(0)
+                .cast(pl.Int8)
+                .alias("offer_imputed_ln"),
+                pl.when(wn_flag)
+                .then(1)
+                .otherwise(0)
+                .cast(pl.Int8)
+                .alias("offer_imputed_wn"),
+                pl.when(enr_flag)
+                .then(1)
+                .otherwise(0)
+                .cast(pl.Int8)
+                .alias("offer_imputed_enr"),
+                pl.when((pl.col(original_offer) == 1) | any_imputation)
+                .then(1)
+                .otherwise(0)
+                .cast(pl.Int8)
+                .alias(offer_col),
+            ]
+        )
+        .drop(original_offer)
+    )
+
+
+
+
+# mappings.py
+
+import polars as pl
+
+FINAL_SCHEMA = {
+    "cmo_name": {"dtype": pl.String, "default": None},
+    "cmo_region": {"dtype": pl.String, "default": None},
+    "ren_num": {"dtype": pl.String, "default": None},
+    "sid_cepr": {"dtype": pl.String, "default": None},
+
+    "school_name": {"dtype": pl.String, "default": None},
+    "school_year": {"dtype": pl.Int64, "default": None},
+    "entry_grade_clean": {"dtype": pl.String, "default": None},
+    "priority_group": {"dtype": pl.Int64, "default": None},
+    "risk_id_name": {"dtype": pl.String, "default": None},
+    "risk_id": {"dtype": pl.Int64, "default": None},
+
+    "offer": {"dtype": pl.Int8, "default": 0},
+    "waitlist": {"dtype": pl.Int8, "default": 0},
+    "initial_offer": {"dtype": pl.Int8, "default": 0},
+    "waitlist_offer": {"dtype": pl.Int8, "default": 0},
+    "offer_accepted": {"dtype": pl.Int8, "default": None},
+    "accept_school": {"dtype": pl.String, "default": None},
+    "offer_imputed_ln": {"dtype": pl.Int8, "default": 0},
+    "offer_imputed_wn": {"dtype": pl.Int8, "default": 0},
+    "offer_imputed_enr": {"dtype": pl.Int8, "default": 0},
+
+    "priority_group_name": {"dtype": pl.String, "default": None},
+    "priority_group_clean": {"dtype": pl.String, "default": None},
+    "sibling": {"dtype": pl.Int8, "default": 0},
+    "sibling_concur": {"dtype": pl.Int8, "default": 0},
+    "staff": {"dtype": pl.Int8, "default": 0},
+    "zoned": {"dtype": pl.Int8, "default": 0},
+    "transfer": {"dtype": pl.Int8, "default": 0},
+
+    "application_cancel": {"dtype": pl.Int8, "default": 0},
+    "late_app": {"dtype": pl.Int8, "default": 0},
+    "former_dupe": {"dtype": pl.Int8, "default": 0},
+
+    "enroll": {"dtype": pl.Int8, "default": 0},
+    "enroll_school": {"dtype": pl.Int8, "default": 0},
+    "enroll_year": {"dtype": pl.Int8, "default": 0},
+    "enroll_school_year": {"dtype": pl.Int8, "default": 0},
+    "enroll_missing": {"dtype": pl.Int8, "default": 1},
+
+    "enrollment_years": {"dtype": pl.Int64, "default": None},
+    "year_enrolled": {"dtype": pl.Int64, "default": None},
+    "withdraw": {"dtype": pl.Int8, "default": 0},
+    "waitlist_number": {"dtype": pl.Int64, "default": None},
+    "lottery_number": {"dtype": pl.Int64, "default": None},
+    "offer_date": {"dtype": pl.Date, "default": None},
+    "application_date": {"dtype": pl.Date, "default": None},
+    "birth_cohort": {"dtype": pl.Int64, "default": None},
+    "age": {"dtype": pl.Int64, "default": None},
+}
+
+
+
+from typing import Mapping, Union
+
+import polars as pl
+
+import mappings
+
+
+def apply_final_schema(
+    frame: Union[pl.DataFrame, pl.LazyFrame],
+    schema: Mapping[str, Mapping[str, object]] = mappings.FINAL_SCHEMA,
+) -> Union[pl.DataFrame, pl.LazyFrame]:
+    """
+    Restrict a Polars DataFrame/LazyFrame to the final LTC application schema.
+
+    Existing columns are kept, cast to their declared dtype, and ordered
+    according to ``schema``. Columns not listed in ``schema`` are dropped.
+    Missing schema columns are created from their declared defaults.
+
+    The schema should be a mapping like:
+
+        {
+            "offer": {"dtype": pl.Int8, "default": 0},
+            "school_name": {"dtype": pl.String, "default": None},
+        }
+
+    Defaults may be either scalar values or Polars expressions. Expression
+    defaults are useful when a missing final column should be derived from
+    other columns.
+    """
+    if not isinstance(frame, (pl.DataFrame, pl.LazyFrame)):
+        raise TypeError("frame must be a Polars DataFrame or LazyFrame.")
+
+    existing_columns = (
+        frame.collect_schema().names()
+        if isinstance(frame, pl.LazyFrame)
+        else frame.columns
+    )
+
+    final_columns = []
+
+    for column, spec in schema.items():
+        dtype = spec["dtype"]
+        default = spec.get("default")
+
+        if column in existing_columns:
+            expr = pl.col(column)
+        elif isinstance(default, pl.Expr):
+            expr = default
+        else:
+            expr = pl.lit(default)
+
+        final_columns.append(
+            expr.cast(dtype, strict=False).alias(column)
+        )
+
+    return frame.select(final_columns)