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

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

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

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

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

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

@@ -2,6 +2,18 @@ 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]]
+
+
+from typing import Callable, Optional, Sequence, Union
+
+import polars as pl
+
 
 def impute_offers(
     frame: Union[pl.DataFrame, pl.LazyFrame],
@@ -9,131 +21,110 @@ def impute_offers(
     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.
-
-    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)
-    ...
+    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.")
 
+    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:
+    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 any_sibling_priority(cols: Optional[Sequence[str]]) -> pl.Expr:
-        if not cols:
-            return pl.lit(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
 
-    work = frame.with_columns(pl.col(offer_col).fill_null(0).alias(original_offer))
+    def maybe_over(expr: pl.Expr) -> pl.Expr:
+        if within_cols:
+            return expr.over(within_cols)
+        return expr
+
+    work = frame.with_columns(
+        binary_col(offer_col).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:
+        lottery_number = number_col(lottery_number_col)
+
         observed_regular_offer = (
             (pl.col(original_offer) == 1)
             & ~any_sibling_priority(sibling_priority_cols)
         )
 
-        lottery_cutoff = (
-            pl.col(lottery_number_col)
+        lottery_cutoff = maybe_over(
+            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 = (
-            pl.col(waitlist_number_col)
+        waitlist_cutoff = maybe_over(
+            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)
@@ -175,9 +166,6 @@ def impute_offers(
         .drop(original_offer)
     )
 
-
-
-
 # mappings.py
 
 import polars as pl