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

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

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

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

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

ltc_code-0.1.97/src/ltc_code/june13.py

@@ -0,0 +1,330 @@
+from typing import Callable, Optional, Sequence, Union
+
+import polars as pl
+
+from typing import Callable, Optional, Sequence, Union
+
+import polars as pl
+
+Frame = 
+Rule = Union[pl.Expr, Callable[[Frame], pl.Expr]]
+
+
+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,
+    within_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.
+
+    Parameters
+    ----------
+    frame:
+        Polars DataFrame or LazyFrame.
+    offer_col:
+        Name of the observed offer column. Null is treated as 0.
+    lottery_number_col:
+        Optional lottery number column.
+    waitlist_number_col:
+        Optional waitlist number column.
+    enrollment_col:
+        Optional enrollment indicator column.
+    lottery_imputation, waitlist_imputation, enrollment_imputation:
+        Optional custom imputation rules. Each may be a Polars expression or a
+        callable that accepts ``frame`` and returns a Polars expression.
+    sibling_priority_cols:
+        Optional columns indicating sibling priority. If supplied, observed
+        sibling-priority offers are excluded when computing the default lottery
+        cutoff.
+    within_cols:
+        Optional columns defining the risk set within which lottery and waitlist
+        cutoffs are computed. For example:
+
+            ["school_name", "school_year", "entry_grade_clean", "priority_group"]
+
+        If omitted, cutoffs are computed globally.
+
+    Default Rules
+    -------------
+    Lottery:
+        Within each ``within_cols`` group, find the highest lottery number among
+        observed offers, excluding sibling-priority offers if
+        ``sibling_priority_cols`` is supplied. Impute offers for 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 within each group use the highest waitlist number among
+        those enrollment-implied offers as the waitlist cutoff.
+
+        If enrollment information is not supplied, within each group use the
+        highest waitlist number among observed offers as the waitlist cutoff.
+
+        Then impute offers for originally non-offered rows with waitlist numbers
+        at or below that cutoff.
+
+    Example
+    -------
+    >>> result = impute_offers(
+    ...     df,
+    ...     offer_col="offer",
+    ...     lottery_number_col="lottery_number",
+    ...     waitlist_number_col="waitlist_number",
+    ...     enrollment_col="enroll_school_year",
+    ...     sibling_priority_cols=["sibling", "sibling_concur"],
+    ...     within_cols=["school_name", "school_year", "entry_grade_clean", "priority_group"],
+    ... )
+    """
+    if not isinstance(frame, (pl.DataFrame, pl.LazyFrame)):
+        raise TypeError("frame must be a Polars DataFrame or LazyFrame.")
+
+    within_cols = list(within_cols or [])
+    sibling_priority_cols = list(sibling_priority_cols or [])
+
+    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: Sequence[str]) -> pl.Expr:
+        priority_expr = pl.lit(False)
+        for col in cols:
+            priority_expr = priority_expr | (pl.col(col).fill_null(0) == 1)
+        return priority_expr
+
+    def maybe_over(expr: pl.Expr) -> pl.Expr:
+        if within_cols:
+            return expr.over(within_cols)
+        return expr
+
+    work = frame.with_columns(
+        pl.col(offer_col).fill_null(0).cast(pl.Int8, strict=False).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 = maybe_over(
+            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 = maybe_over(
+            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)