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

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

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

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

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

{ltc_code-0.1.97 → ltc_code-0.1.98}/src/ltc_code/june13.py

@@ -10,92 +10,35 @@ Frame =
 Rule = Union[pl.Expr, Callable[[Frame], pl.Expr]]
 
 
+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,
+    lottery_imputation: Optional[Union[pl.Expr, Callable[[Union[pl.DataFrame, pl.LazyFrame]], pl.Expr]]] = None,
+    waitlist_imputation: Optional[Union[pl.Expr, Callable[[Union[pl.DataFrame, pl.LazyFrame]], pl.Expr]]] = None,
+    enrollment_imputation: Optional[Union[pl.Expr, Callable[[Union[pl.DataFrame, pl.LazyFrame]], 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"],
-    ... )
+    The function safely casts provided 0/1 and numeric order columns before
+    comparing them, so columns stored as strings like ``"0"``, ``"1"``, or
+    ``"12"`` can still be used.
+
+    Binary columns are interpreted as 1 only when they cast cleanly to numeric
+    value 1. Nulls and non-numeric strings are treated as 0.
+
+    Lottery and waitlist number columns are cast to Float64 for cutoff
+    comparisons. Nulls and non-numeric strings are ignored.
     """
     if not isinstance(frame, (pl.DataFrame, pl.LazyFrame)):
         raise TypeError("frame must be a Polars DataFrame or LazyFrame.")
@@ -105,16 +48,22 @@ def impute_offers(
 
     original_offer = "__original_offer_for_imputation__"
 
-    def as_rule_expr(rule: Union[pl.Expr, Callable[[Frame], pl.Expr]]) -> pl.Expr:
+    def as_rule_expr(rule: Union[pl.Expr, Callable[[Union[pl.DataFrame, pl.LazyFrame]], 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 binary_col(col: str) -> pl.Expr:
+        return pl.col(col).cast(pl.Float64, strict=False).fill_null(0)
+
+    def number_col(col: str) -> pl.Expr:
+        return pl.col(col).cast(pl.Float64, strict=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)
+            priority_expr = priority_expr | (binary_col(col) == 1)
         return priority_expr
 
     def maybe_over(expr: pl.Expr) -> pl.Expr:
@@ -123,7 +72,7 @@ def impute_offers(
         return expr
 
     work = frame.with_columns(
-        pl.col(offer_col).fill_null(0).cast(pl.Int8, strict=False).alias(original_offer)
+        binary_col(offer_col).cast(pl.Int8, strict=False).alias(original_offer)
     )
 
     originally_no_offer = pl.col(original_offer) == 0
@@ -132,46 +81,50 @@ def impute_offers(
     if lottery_imputation is not None:
         lottery_rule = as_rule_expr(lottery_imputation)
     elif lottery_number_col is not None:
+        lottery_number = number_col(lottery_number_col)
+
         observed_regular_offer = (
             (pl.col(original_offer) == 1)
             & ~any_sibling_priority(sibling_priority_cols)
         )
 
         lottery_cutoff = maybe_over(
-            pl.col(lottery_number_col)
+            lottery_number
             .filter(observed_regular_offer)
             .max()
         )
 
         lottery_rule = (
-            pl.col(lottery_number_col).is_not_null()
-            & (pl.col(lottery_number_col) <= lottery_cutoff)
+            lottery_number.is_not_null()
+            & (lottery_number <= 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)
+        enrollment_rule = (binary_col(enrollment_col) == 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:
+        waitlist_number = number_col(waitlist_number_col)
+
         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)
+            waitlist_number
             .filter(waitlist_cutoff_source)
             .max()
         )
 
         waitlist_rule = (
-            pl.col(waitlist_number_col).is_not_null()
-            & (pl.col(waitlist_number_col) <= waitlist_cutoff)
+            waitlist_number.is_not_null()
+            & (waitlist_number <= waitlist_cutoff)
         ).fill_null(False)
 
     false_expr = pl.lit(False)
@@ -213,7 +166,6 @@ def impute_offers(
         .drop(original_offer)
     )
 
-
 # mappings.py
 
 import polars as pl