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

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

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

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

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

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

@@ -2,6 +2,13 @@ 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],
@@ -13,6 +20,7 @@ def impute_offers(
     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.
@@ -28,48 +36,73 @@ def impute_offers(
     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.
+    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
     -------
-    >>> 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(
+    >>> result = 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"],
+    ...     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"],
     ... )
-    shape: (4, 9)
-    ...
     """
     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:
@@ -78,16 +111,21 @@ def impute_offers(
             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 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
 
-    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(
+        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
@@ -99,7 +137,7 @@ def impute_offers(
             & ~any_sibling_priority(sibling_priority_cols)
         )
 
-        lottery_cutoff = (
+        lottery_cutoff = maybe_over(
             pl.col(lottery_number_col)
             .filter(observed_regular_offer)
             .max()
@@ -125,7 +163,7 @@ def impute_offers(
         else:
             waitlist_cutoff_source = pl.col(original_offer) == 1
 
-        waitlist_cutoff = (
+        waitlist_cutoff = maybe_over(
             pl.col(waitlist_number_col)
             .filter(waitlist_cutoff_source)
             .max()
@@ -176,8 +214,6 @@ def impute_offers(
     )
 
 
-
-
 # mappings.py
 
 import polars as pl