microdf-python 1.0.2__tar.gz → 1.1.1__tar.gz

{microdf_python-1.0.2/microdf_python.egg-info → microdf_python-1.1.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: microdf-python
-Version: 1.0.2
+Version: 1.1.1
 Summary: Weighted pandas DataFrames and Series for survey microdata
 Author-email: Max Ghenis <max@policyengine.org>
 License: MIT
@@ -20,8 +20,6 @@ Requires-Dist: linecheck; extra == "dev"
 Requires-Dist: pytest; extra == "dev"
 Requires-Dist: pytest-cov; extra == "dev"
 Requires-Dist: setuptools; extra == "dev"
-Provides-Extra: docs
-Requires-Dist: jupyter_book; extra == "docs"
 Dynamic: license-file
 
 [![Build](https://github.com/PolicyEngine/microdf/workflows/Pull%20request/badge.svg)](https://github.com/PolicyEngine/microdf/actions)

{microdf_python-1.0.2 → microdf_python-1.1.1}/microdf/microdataframe.py

@@ -12,6 +12,80 @@ from microdf.microseries import MicroSeries, MicroSeriesGroupBy
 logger = logging.getLogger(__name__)
 
 
+class _MicroLocIndexer:
+    """Custom loc indexer that returns MicroDataFrame with proper weights."""
+
+    def __init__(self, mdf: "MicroDataFrame"):
+        self._mdf = mdf
+        # Get the parent's loc indexer
+        self._parent_loc = pd.DataFrame.loc.fget(mdf)
+
+    def __getitem__(self, key):
+        # Use the parent DataFrame's loc indexer
+        result = self._parent_loc[key]
+
+        if isinstance(result, pd.DataFrame):
+            # Get the filtered weights based on the result's index
+            new_weights = self._mdf.weights.reindex(result.index)
+            return MicroDataFrame(result, weights=new_weights)
+        elif isinstance(result, pd.Series):
+            # Single row or column selected
+            if result.name in self._mdf.columns:
+                # Column was selected - return MicroSeries with all weights
+                return MicroSeries(result, weights=self._mdf.weights)
+            else:
+                # Row was selected - return as-is (scalar values for each col)
+                return result
+        else:
+            # Scalar value
+            return result
+
+    def __setitem__(self, key, value):
+        self._parent_loc[key] = value
+        self._mdf._link_all_weights()
+
+
+class _MicroILocIndexer:
+    """Custom iloc indexer that returns MicroDataFrame with proper weights."""
+
+    def __init__(self, mdf: "MicroDataFrame"):
+        self._mdf = mdf
+        # Get the parent's iloc indexer
+        self._parent_iloc = pd.DataFrame.iloc.fget(mdf)
+
+    def __getitem__(self, key):
+        # Use the parent DataFrame's iloc indexer
+        result = self._parent_iloc[key]
+
+        if isinstance(result, pd.DataFrame):
+            # Get the filtered weights based on the result's index
+            new_weights = self._mdf.weights.iloc[
+                self._mdf.index.get_indexer(result.index)
+            ]
+            new_weights = pd.Series(new_weights.values, index=result.index)
+            return MicroDataFrame(result, weights=new_weights)
+        elif isinstance(result, pd.Series):
+            # Single row or column selected
+            if isinstance(key, tuple) and len(key) == 2:
+                # df.iloc[:, col_idx] - column selection
+                row_key = key[0]
+                if isinstance(row_key, slice) and row_key == slice(None):
+                    # All rows selected for a column
+                    return MicroSeries(result, weights=self._mdf.weights)
+            # Check if this is a column (result index matches mdf index)
+            if result.index.equals(self._mdf.index):
+                return MicroSeries(result, weights=self._mdf.weights)
+            # Row selection - return as-is
+            return result
+        else:
+            # Scalar value
+            return result
+
+    def __setitem__(self, key, value):
+        self._parent_iloc[key] = value
+        self._mdf._link_all_weights()
+
+
 class MicroDataFrame(pd.DataFrame):
     def __init__(self, *args, weights=None, **kwargs):
         """A DataFrame-inheriting class for weighted microdata. Weights can be
@@ -26,6 +100,23 @@ class MicroDataFrame(pd.DataFrame):
         self._link_all_weights()
         self.override_df_functions()
 
+    @property
+    def loc(self) -> _MicroLocIndexer:
+        """Label-based indexer that preserves MicroDataFrame type and weights.
+
+        :return: Custom loc indexer for MicroDataFrame
+        """
+        return _MicroLocIndexer(self)
+
+    @property
+    def iloc(self) -> _MicroILocIndexer:
+        """Integer-based indexer that preserves MicroDataFrame type and
+        weights.
+
+        :return: Custom iloc indexer for MicroDataFrame
+        """
+        return _MicroILocIndexer(self)
+
     def override_df_functions(self) -> None:
         """Override DataFrame functions to work with weighted operations."""
         for name in MicroSeries.FUNCTIONS:
@@ -643,6 +734,7 @@ class MicroDataFrame(pd.DataFrame):
 
 class MicroDataFrameGroupBy(pd.core.groupby.generic.DataFrameGroupBy):
     def _init(self, by: Union[str, List]):
+        self._by = by
         self.columns = list(self.obj.columns)
         if isinstance(by, list):
             for column in by:
@@ -656,6 +748,10 @@ class MicroDataFrameGroupBy(pd.core.groupby.generic.DataFrameGroupBy):
             for col in self.columns
             if pd.api.types.is_numeric_dtype(self.obj[col])
         ]
+        # Store reference to weights groupby for column selection
+        self._weights_groupby = copy.deepcopy(
+            super().__getitem__("__tmp_weights")
+        )
         for fn_name in MicroSeries.SCALAR_FUNCTIONS:
 
             def get_fn(name):
@@ -669,11 +765,9 @@ class MicroDataFrameGroupBy(pd.core.groupby.generic.DataFrameGroupBy):
                         except Exception:
                             # Skip columns that can't be aggregated
                             pass
-                    return (
-                        MicroDataFrame(results)
-                        if results
-                        else MicroDataFrame()
-                    )
+                    # Return plain DataFrame - aggregated results don't have
+                    # per-row weights (weights were already applied)
+                    return pd.DataFrame(results) if results else pd.DataFrame()
 
                 return fn
 
@@ -691,12 +785,99 @@ class MicroDataFrameGroupBy(pd.core.groupby.generic.DataFrameGroupBy):
                         except Exception:
                             # Skip columns that can't be aggregated
                             pass
-                    return (
-                        MicroDataFrame(results)
-                        if results
-                        else MicroDataFrame()
-                    )
+                    # Return plain DataFrame - aggregated results don't have
+                    # per-row weights (weights were already applied)
+                    return pd.DataFrame(results) if results else pd.DataFrame()
 
                 return fn
 
             setattr(self, fn_name, get_fn(fn_name))
+
+    def __getitem__(
+        self, key: Union[str, List]
+    ) -> Union["MicroSeriesGroupBy", "MicroDataFrameGroupBy"]:
+        """Select columns from the groupby object while preserving weights.
+
+        This ensures that operations like groupby(col)["y"].sum() or
+        groupby(col)[["y"]].sum() use weighted aggregation.
+
+        :param key: Column name or list of column names
+        :return: MicroSeriesGroupBy for single column, MicroDataFrameGroupBy
+            for multiple columns
+        """
+        if isinstance(key, str):
+            # Single column - return MicroSeriesGroupBy
+            result = super().__getitem__(key)
+            result.__class__ = MicroSeriesGroupBy
+            result._init()
+            result.weights = self._weights_groupby
+            return result
+        else:
+            # Multiple columns - return a new MicroDataFrameGroupBy
+            # with only the selected columns
+            result = super().__getitem__(key)
+            result.__class__ = MicroDataFrameGroupBy
+            # Re-initialize with the subset of columns
+            result._by = self._by
+            result.columns = list(key) if hasattr(key, "__iter__") else [key]
+            result.numeric_columns = [
+                col
+                for col in result.columns
+                if pd.api.types.is_numeric_dtype(result.obj[col])
+            ]
+            result._weights_groupby = self._weights_groupby
+            # Set up the column attributes as MicroSeriesGroupBy
+            for col in result.columns:
+                col_gb = super().__getitem__(col)
+                col_gb.__class__ = MicroSeriesGroupBy
+                col_gb._init()
+                col_gb.weights = self._weights_groupby
+                setattr(result, col, col_gb)
+            # Set up the scalar and vector functions
+            for fn_name in MicroSeries.SCALAR_FUNCTIONS:
+
+                def get_scalar_fn(name, res):
+                    def fn(*args, **kwargs):
+                        results = {}
+                        for col in res.numeric_columns:
+                            try:
+                                results[col] = getattr(
+                                    getattr(res, col), name
+                                )(*args, **kwargs)
+                            except Exception:
+                                pass
+                        # Return plain DataFrame - aggregated results don't
+                        # have per-row weights (weights were already applied)
+                        return (
+                            pd.DataFrame(results)
+                            if results
+                            else pd.DataFrame()
+                        )
+
+                    return fn
+
+                setattr(result, fn_name, get_scalar_fn(fn_name, result))
+            for fn_name in MicroSeries.VECTOR_FUNCTIONS:
+
+                def get_vector_fn(name, res):
+                    def fn(*args, **kwargs):
+                        results = {}
+                        for col in res.numeric_columns:
+                            try:
+                                results[col] = getattr(
+                                    getattr(res, col), name
+                                )(*args, **kwargs)
+                            except Exception:
+                                pass
+                        # Return plain DataFrame - aggregated results don't
+                        # have per-row weights (weights were already applied)
+                        return (
+                            pd.DataFrame(results)
+                            if results
+                            else pd.DataFrame()
+                        )
+
+                    return fn
+
+                setattr(result, fn_name, get_vector_fn(fn_name, result))
+            return result

{microdf_python-1.0.2 → microdf_python-1.1.1}/microdf/microseries.py

@@ -112,17 +112,18 @@ class MicroSeries(pd.Series):
     def quantile(self, q: np.array) -> pd.Series:
         """Calculates weighted quantiles of the MicroSeries.
 
-        Doesn't exactly match unweighted quantiles of stacked values.
-        See stackoverflow.com/q/21844024#comment102342137_29677616.
+        Uses the inverse CDF method: the q-th quantile is the smallest
+        value where the cumulative weight proportion >= q. This matches
+        the default behavior of R's survey::svyquantile.
 
-        :param q: Array of quantiles to calculate.
-        :type q: np.array
+        :param q: Quantile(s) to calculate, must be in [0, 1].
+        :type q: float or np.array
 
-        :return: Array of weighted quantiles.
-        :rtype: pd.Series
+        :return: Weighted quantile value(s).
+        :rtype: float or pd.Series
         """
         values = np.array(self.values)
-        quantiles = np.array(q)
+        quantiles = np.atleast_1d(q)
         sample_weight = np.array(self.weights)
         assert np.all(quantiles >= 0) and np.all(
             quantiles <= 1
@@ -130,11 +131,20 @@ class MicroSeries(pd.Series):
         sorter = np.argsort(values)
         values = values[sorter]
         sample_weight = sample_weight[sorter]
-        weighted_quantiles = np.cumsum(sample_weight) - 0.5 * sample_weight
-        weighted_quantiles /= np.sum(sample_weight)
-        result = np.interp(quantiles, weighted_quantiles, values)
-        if quantiles.shape == ():
-            return result
+        cumsum = np.cumsum(sample_weight)
+        cumsum_normalized = cumsum / cumsum[-1]
+        result = np.array(
+            [
+                values[
+                    min(
+                        np.searchsorted(cumsum_normalized, qi), len(values) - 1
+                    )
+                ]
+                for qi in quantiles
+            ]
+        )
+        if np.array(q).shape == ():
+            return result[0]
         return pd.Series(result, index=quantiles)
 
     @scalar_function

{microdf_python-1.0.2 → microdf_python-1.1.1}/microdf/tests/test_microseries_dataframe.py

@@ -112,6 +112,55 @@ def test_median() -> None:
     assert series.median() == 4
 
 
+def test_weighted_quantile_skewed() -> None:
+    # 99% of the population has 0 income, 1% has 1M
+    # The median should be 0, not an interpolated value
+    series = mdf.MicroSeries([0, 1_000_000], weights=[99, 1])
+    assert series.median() == 0
+    assert series.quantile(0.5) == 0
+    # 99th percentile is still 0 since exactly 99% have 0
+    assert series.quantile(0.99) == 0
+    # Only quantile > 0.99 gives 1M
+    assert series.quantile(1.0) == 1_000_000
+    # Test multiple quantiles
+    result = series.quantile([0.1, 0.5, 0.99, 1.0])
+    assert result[0.1] == 0
+    assert result[0.5] == 0
+    assert result[0.99] == 0
+    assert result[1.0] == 1_000_000
+
+
+def test_weighted_quantile_boundaries() -> None:
+    # Test q=0 returns minimum, q=1 returns maximum
+    series = mdf.MicroSeries([10, 20, 30], weights=[1, 1, 1])
+    assert series.quantile(0.0) == 10
+    assert series.quantile(1.0) == 30
+
+
+def test_weighted_quantile_equal_weights() -> None:
+    # With equal weights, should match "replicated" interpretation
+    # Values: 1, 2, 3 each with weight 2 -> like [1,1,2,2,3,3]
+    series = mdf.MicroSeries([1, 2, 3], weights=[2, 2, 2])
+    # cumsum_normalized = [2/6, 4/6, 6/6] = [0.333, 0.667, 1.0]
+    # median (0.5): smallest where cumsum >= 0.5 -> index 1 -> value 2
+    assert series.median() == 2
+    # 0.25 quantile: smallest where cumsum >= 0.25 -> index 0 -> value 1
+    assert series.quantile(0.25) == 1
+    # 0.75 quantile: smallest where cumsum >= 0.75 -> index 2 -> value 3
+    assert series.quantile(0.75) == 3
+
+
+def test_weighted_quantile_unsorted_input() -> None:
+    # Ensure sorting works correctly
+    series = mdf.MicroSeries([30, 10, 20], weights=[1, 2, 1])
+    # Sorted: values [10, 20, 30], weights [2, 1, 1]
+    # cumsum_normalized = [0.5, 0.75, 1.0]
+    assert series.quantile(0.0) == 10
+    assert series.quantile(0.5) == 10  # cumsum[0]=0.5 >= 0.5
+    assert series.quantile(0.6) == 20  # cumsum[1]=0.75 >= 0.6
+    assert series.quantile(1.0) == 30
+
+
 def test_unweighted_groupby() -> None:
     df = mdf.MicroDataFrame({"x": [1, 2], "y": [3, 4], "z": [5, 6]})
     assert (df.groupby("x").z.sum().values == np.array([5.0, 6.0])).all()
@@ -287,3 +336,60 @@ def test_reset_index_inplace() -> None:
     assert "second" in mdf_multi.columns
     assert list(mdf_multi.index) == [0, 1, 2, 3]
     np.testing.assert_array_equal(mdf_multi.weights.values, weights)
+
+
+def test_loc_preserves_weights() -> None:
+    """Test that .loc[] returns MicroDataFrame with proper weights (issue
+    #265)."""
+    df = mdf.MicroDataFrame(
+        {"one": [1, 1, 1, 1, 1]}, weights=[10, 20, 30, 40, 50]
+    )
+
+    # Filter all rows (should get same weights)
+    filtered = df.loc[df.one == 1]
+    assert isinstance(filtered, MicroDataFrame)
+    assert filtered.one.sum() == 150.0  # Weighted sum
+
+    # Partial filter
+    df2 = mdf.MicroDataFrame(
+        {"x": [1, 2, 3, 4, 5]}, weights=[10, 20, 30, 40, 50]
+    )
+    subset = df2.loc[df2.x > 2]
+    assert isinstance(subset, MicroDataFrame)
+    assert subset.x.sum() == 500.0  # 3*30 + 4*40 + 5*50 = 500
+    np.testing.assert_array_equal(subset.weights.values, [30.0, 40.0, 50.0])
+
+
+def test_iloc_preserves_weights() -> None:
+    """Test that .iloc[] returns MicroDataFrame with proper weights."""
+    df = mdf.MicroDataFrame(
+        {"x": [1, 2, 3, 4, 5]}, weights=[10, 20, 30, 40, 50]
+    )
+
+    # Select rows by position
+    subset = df.iloc[2:5]
+    assert isinstance(subset, MicroDataFrame)
+    assert subset.x.sum() == 500.0  # 3*30 + 4*40 + 5*50 = 500
+    np.testing.assert_array_equal(subset.weights.values, [30.0, 40.0, 50.0])
+
+
+def test_groupby_column_selection() -> None:
+    """Test that groupby column selection preserves weights (issue #193)."""
+    d = mdf.MicroDataFrame(
+        dict(g=["a", "a", "b"], y=[1, 2, 3]), weights=[4, 5, 6]
+    )
+
+    # Test single column string selection
+    result_str = d.groupby("g")["y"].sum()
+    assert result_str["a"] == 14.0  # 1*4 + 2*5 = 14
+    assert result_str["b"] == 18.0  # 3*6 = 18
+
+    # Test list column selection
+    result_list = d.groupby("g")[["y"]].sum()
+    assert result_list.loc["a", "y"] == 14.0
+    assert result_list.loc["b", "y"] == 18.0
+
+    # Aggregated results should be plain DataFrame (no spurious weight column)
+    result_all = d.groupby("g").sum()
+    assert "weight" not in result_all.columns
+    assert list(result_all.columns) == ["y"]

{microdf_python-1.0.2 → microdf_python-1.1.1/microdf_python.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: microdf-python
-Version: 1.0.2
+Version: 1.1.1
 Summary: Weighted pandas DataFrames and Series for survey microdata
 Author-email: Max Ghenis <max@policyengine.org>
 License: MIT
@@ -20,8 +20,6 @@ Requires-Dist: linecheck; extra == "dev"
 Requires-Dist: pytest; extra == "dev"
 Requires-Dist: pytest-cov; extra == "dev"
 Requires-Dist: setuptools; extra == "dev"
-Provides-Extra: docs
-Requires-Dist: jupyter_book; extra == "docs"
 Dynamic: license-file
 
 [![Build](https://github.com/PolicyEngine/microdf/workflows/Pull%20request/badge.svg)](https://github.com/PolicyEngine/microdf/actions)

{microdf_python-1.0.2 → microdf_python-1.1.1}/microdf_python.egg-info/requires.txt

@@ -12,6 +12,3 @@ linecheck
 pytest
 pytest-cov
 setuptools
-
-[docs]
-jupyter_book

{microdf_python-1.0.2 → microdf_python-1.1.1}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "microdf-python"
-version = "1.0.2"
+version = "1.1.1"
 description = "Weighted pandas DataFrames and Series for survey microdata"
 readme = "README.md"
 authors = [
@@ -30,9 +30,8 @@ dev = [
     "pytest-cov",
     "setuptools",
 ]
-docs = [
-    "jupyter_book",
-]
+# Note: Documentation uses MyST (Jupyter Book 2.0) which is installed via npm
+# Run: cd docs && myst build --html
 
 [tool.setuptools.packages.find]
 where = ["."]