microdf-python 1.3.10__py3-none-any.whl → 1.4.1__py3-none-any.whl

microdf/_weights.py

@@ -0,0 +1,257 @@
+"""Pandas subclass hooks for copying weights and retaining row positions."""
+
+import numpy as np
+import pandas as pd
+
+
+def weight_series(values, index):
+    """Own an independent float array aligned positionally to an index."""
+    return pd.Series(np.array(values, dtype=float, copy=True), index=index)
+
+
+def aligned_weights(source, index):
+    """Align only when row identity can be established without guessing."""
+    if source.index.equals(index):
+        return weight_series(source.weights, index)
+    if source.index.is_unique:
+        positions = source.index.get_indexer(index)
+        if (positions >= 0).all():
+            return weight_series(source.weights.iloc[positions], index)
+    raise ValueError(
+        "Cannot propagate weights unambiguously to these rows. "
+        "Use positional selection, or explicitly supply weights for the result."
+    )
+
+
+def finalize_weights(result, source, method, previous_weights):
+    """Propagate row weights after pandas has finalized its own metadata."""
+    if method == "transpose" and result.ndim == 2:
+        raise ValueError(
+            "Cannot transpose row weights onto columns. "
+            "Convert to a pandas DataFrame first and supply explicit result weights."
+        )
+    if method == "concat":
+        objects = source.objs
+        if not all(isinstance(obj, WeightPropagationMixin) for obj in objects):
+            raise ValueError(
+                "Cannot concatenate weighted and unweighted objects: "
+                "provide explicit weights for every input."
+            )
+        weights = concat_weights(result, objects, getattr(source, "axis", None))
+        result.weights = weights
+        if result.ndim == 2:
+            names = [obj.__dict__.get("weights_col") for obj in objects]
+            result.weights_col = (
+                names[0] if all(name == names[0] for name in names) else None
+            )
+    elif isinstance(source, WeightPropagationMixin):
+        if method == "reset_index" and source.ndim == result.ndim == 1:
+            # Series.reset_index(drop=True) changes labels, never row order.
+            # pandas 2 constructs a result; pandas 3 relabels a shallow copy.
+            result.weights = weight_series(source.weights, result.index)
+        elif method == "rename" and previous_weights is not None:
+            result.weights = weight_series(previous_weights, result.index)
+        else:
+            result.weights = aligned_weights(source, result.index)
+    return result
+
+
+def concat_weights(result, objects, axis):
+    """Recover concat row identity without relying on pandas 2-only state."""
+    if objects[0].ndim == 1 and result.ndim == 2:
+        axis = 1
+    row_weights = None
+    if axis != 1 and len(result) == sum(len(obj) for obj in objects):
+        combined_index = objects[0].index.append([obj.index for obj in objects[1:]])
+        same_rows = result.index.equals(combined_index)
+        reset_rows = result.index.equals(pd.RangeIndex(len(result)))
+        keyed_rows = False
+        if (
+            isinstance(result.index, pd.MultiIndex)
+            and result.index.nlevels > combined_index.nlevels
+        ):
+            trailing = result.index.droplevel(
+                list(range(result.index.nlevels - combined_index.nlevels))
+            )
+            keyed_rows = trailing.equals(combined_index)
+        if same_rows or reset_rows or keyed_rows:
+            row_weights = np.concatenate([np.asarray(obj.weights) for obj in objects])
+    column_weights = None
+    if (
+        axis != 0
+        and result.ndim == 2
+        and len(result.columns)
+        == sum(obj.shape[1] if obj.ndim == 2 else 1 for obj in objects)
+    ):
+        values = np.empty(len(result))
+        known = np.zeros(len(result), dtype=bool)
+        conflict = False
+        for obj in objects:
+            if obj.index.equals(result.index):
+                positions = np.arange(len(result))
+            elif obj.index.is_unique:
+                positions = obj.index.get_indexer(result.index)
+            else:
+                conflict = True
+                break
+            present = positions >= 0
+            incoming = np.asarray(obj.weights)[positions[present]]
+            overlap = known[present]
+            if not np.array_equal(
+                values[present][overlap], incoming[overlap], equal_nan=True
+            ):
+                conflict = True
+                break
+            values[present] = incoming
+            known[present] = True
+        if not conflict and known.all():
+            column_weights = values
+    if row_weights is not None:
+        if column_weights is not None and not np.array_equal(
+            row_weights, column_weights, equal_nan=True
+        ):
+            raise ValueError(
+                "Ambiguous weights: pandas did not expose the concat axis."
+            )
+        return weight_series(row_weights, result.index)
+    if column_weights is not None:
+        return weight_series(column_weights, result.index)
+    raise ValueError(
+        "Cannot propagate concat weights: conflicting or ambiguous row weights."
+    )
+
+
+class WeightPropagationMixin:
+    """Use positional provenance for pandas operations that select rows."""
+
+    def _plain(self):
+        return (
+            pd.Series(self, copy=False)
+            if self.ndim == 1
+            else pd.DataFrame(self, copy=False)
+        )
+
+    def _weighted_result(self, plain, weights):
+        result = type(self)(plain, weights=weight_series(weights, plain.index))
+        if self.ndim == 2:
+            result.weights_col = self.__dict__.get("weights_col")
+        return result
+
+    def _finish_row_operation(self, plain, positions, inplace=False):
+        result = self._weighted_result(plain, self.weights.iloc[positions])
+        if inplace:
+            self._update_inplace(result)
+            self.weights = result.weights
+            return None
+        return result
+
+    def take(self, indices, axis=0, **kwargs):
+        """Take values and weights using the same positional indexer."""
+        axis = self._get_axis_number(axis)
+        plain = self._plain().take(indices, axis=axis, **kwargs)
+        weights = self.weights.iloc[indices] if axis == 0 else self.weights
+        return self._weighted_result(plain, weights)
+
+    def _slice(self, slobj, axis=0):
+        plain = self._plain()._slice(slobj, axis=axis)
+        weights = self.weights.iloc[slobj] if axis == 0 else self.weights
+        return self._weighted_result(plain, weights)
+
+    def _reindex_with_indexers(
+        self, reindexers, fill_value=None, copy=False, allow_dups=False
+    ):
+        plain = self._plain()._reindex_with_indexers(
+            reindexers, fill_value=fill_value, allow_dups=allow_dups
+        )
+        if copy:
+            plain = plain.copy()
+        indexer = reindexers.get(0, (None, None))[1]
+        if indexer is not None:
+            if (np.asarray(indexer) < 0).any():
+                raise ValueError(
+                    "Cannot invent weights for new rows introduced by reindex."
+                )
+            weights = self.weights.iloc[indexer]
+        else:
+            weights = self.weights
+        return self._weighted_result(plain, weights)
+
+    def drop(
+        self,
+        labels=None,
+        axis=0,
+        index=None,
+        columns=None,
+        level=None,
+        inplace=False,
+        errors="raise",
+    ):
+        plain = self._plain().drop(
+            labels=labels,
+            axis=axis,
+            index=index,
+            columns=columns,
+            level=level,
+            errors=errors,
+        )
+        positions = pd.Series(np.arange(len(self)), index=self.index)
+        row_labels = (
+            index
+            if index is not None
+            else labels
+            if self._get_axis_number(axis) == 0
+            else None
+        )
+        if row_labels is not None:
+            positions = positions.drop(row_labels, level=level, errors=errors)
+        return self._finish_row_operation(plain, np.asarray(positions), inplace)
+
+    def sample(self, *args, **kwargs):
+        """Sample weights with the selected rows, including replacements."""
+        result = super().sample(*args, **kwargs)
+        result.weights = weight_series(result.weights, result.index)
+        return result
+
+    def sort_values(self, *args, **kwargs):
+        """Sort rows and weights together even when labels are duplicated."""
+        inplace = kwargs.pop("inplace", False)
+        axis = self._get_axis_number(kwargs.get("axis", 0))
+        if self.ndim == 2:
+            plain = self._plain()
+            if axis == 1:
+                plain = plain.sort_values(*args, **kwargs)
+                return self._finish_row_operation(plain, np.arange(len(self)), inplace)
+            marker = object()
+            plain = plain.copy(deep=False)
+            plain[marker] = np.arange(len(self))
+            plain = plain.sort_values(*args, **kwargs)
+            positions = np.asarray(plain.pop(marker), dtype=int)
+        else:
+            values = self._plain()
+            positions = pd.Series(
+                np.arange(len(self)), index=self.index, name=self.name
+            )
+            key = kwargs.pop("key", None)
+            positions = positions.sort_values(
+                *args,
+                key=lambda unused: values if key is None else key(values),
+                **kwargs,
+            )
+            plain = values.iloc[np.asarray(positions)]
+            plain.index = positions.index
+            positions = np.asarray(positions, dtype=int)
+        return self._finish_row_operation(plain, positions, inplace)
+
+    def sort_index(self, *args, **kwargs):
+        """Sort the index while preserving positional weight provenance."""
+        inplace = kwargs.pop("inplace", False)
+        axis = self._get_axis_number(kwargs.get("axis", 0))
+        if axis == 1:
+            plain = self._plain().sort_index(*args, **kwargs)
+            return self._finish_row_operation(plain, np.arange(len(self)), inplace)
+        positions = pd.Series(np.arange(len(self)), index=self.index).sort_index(
+            *args, **kwargs
+        )
+        plain = self._plain().iloc[np.asarray(positions)]
+        plain.index = positions.index
+        return self._finish_row_operation(plain, np.asarray(positions), inplace)

microdf/microdataframe.py

@@ -8,93 +8,17 @@ import numpy as np
 import pandas as pd
 
 from microdf.microseries import MicroSeries, MicroSeriesGroupBy
+from microdf._weights import (
+    WeightPropagationMixin,
+    aligned_weights,
+    finalize_weights,
+    weight_series,
+)
 
 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()
-
-    def __getattr__(self, name):
-        """Delegate unknown attributes to the parent loc indexer."""
-        return getattr(self._parent_loc, name)
-
-
-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()
-
-    def __getattr__(self, name):
-        """Delegate unknown attributes to the parent iloc indexer."""
-        return getattr(self._parent_iloc, name)
-
-
-class MicroDataFrame(pd.DataFrame):
+class MicroDataFrame(WeightPropagationMixin, pd.DataFrame):
     # Declare weight state as pandas metadata. pandas includes
     # _metadata attributes in the pickle state, so weights now survive
     # pickling, to_pickle/read_pickle and copy.deepcopy instead of
@@ -112,20 +36,54 @@ class MicroDataFrame(pd.DataFrame):
         :type weights: np.array
         """
         super().__init__(*args, **kwargs)
-        self.weights = None
+        # pandas normalizes mixed-dimensional concat inputs through this
+        # constructor, either as a Series or a one-column mapping. Preserve
+        # that Series' row weights before concat loses the original input.
+        weight_source = args[0] if args else kwargs.get("data")
+        if isinstance(weight_source, dict) and len(weight_source) == 1:
+            weight_source = next(iter(weight_source.values()))
+        if weights is None and isinstance(weight_source, MicroSeries):
+            weights = aligned_weights(weight_source, self.index)
+        self.weights = weight_series(np.ones(len(self)), self.index)
+        self.weights_col = None
         self.set_weights(weights)
         self._link_all_weights()
         self.override_df_functions()
 
-    def __finalize__(self, other, method=None, **kwargs) -> "MicroDataFrame":
-        """Retain copied weights when pandas finalizes a renamed result."""
-        copied_weights = getattr(self, "weights", None) if method == "rename" else None
+    @property
+    def _constructor(self):
+        return MicroDataFrame
+
+    # A row or a column-summary Series has no per-observation weights.
+    # _ixs wraps column selections using their unambiguous row provenance.
+    _constructor_sliced = pd.Series
+
+    def _ixs(self, i, axis=0):
+        result = pd.DataFrame(self, copy=False)._ixs(i, axis=axis)
+        if axis == 1:
+            return MicroSeries(result, weights=self.weights)
+        return result
+
+    def _get_item_cache(self, item):
+        # Weight arrays are independently mutable; cached column wrappers would
+        # retain stale copies after an in-place edit to frame.weights.
+        return self._ixs(self.columns.get_loc(item), axis=1)
+
+    def __finalize__(self, other, method=None, **kwargs):
+        previous = self.__dict__.get("weights")
         super().__finalize__(other, method=method, **kwargs)
-        if copied_weights is not None:
-            # rename already called copy(); metadata propagation must not
-            # replace those weights with the source's mutable Series.
-            self.weights = copied_weights
-        return self
+        return finalize_weights(self, other, method, previous)
+
+    @wraps(pd.DataFrame.cov)
+    def cov(self, *args, **kwargs) -> pd.DataFrame:
+        # Column summaries have no observation weights, even if labels match.
+        result = pd.DataFrame(self, copy=False).cov(*args, **kwargs)
+        return result.__finalize__(self, method="cov")
+
+    @wraps(pd.DataFrame.corr)
+    def corr(self, *args, **kwargs) -> pd.DataFrame:
+        result = pd.DataFrame(self, copy=False).corr(*args, **kwargs)
+        return result.__finalize__(self, method="corr")
 
     def __setstate__(self, state) -> None:
         """Restore a pickled MicroDataFrame.
@@ -141,23 +99,6 @@ 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:
@@ -378,7 +319,7 @@ class MicroDataFrame(pd.DataFrame):
         pass
 
     def _link_all_weights(self) -> None:
-        if self.weights is None:
+        if self.weights is None or len(self.weights) == 0:
             if len(self) > 0:
                 self.set_weights(np.ones((len(self))))
         # In pandas 3.0+, columns are wrapped as MicroSeries on access via
@@ -472,33 +413,19 @@ class MicroDataFrame(pd.DataFrame):
         # that treats it as a Series (equals(), reindex() in __getitem__).
         self.set_weights(np.ones(len(self)))
 
-    def __getitem__(
-        self, key: Union[str, List]
-    ) -> Union[MicroSeries, "MicroDataFrame"]:
-        # Let pandas handle the initial slicing
-        result = super().__getitem__(key)
-
-        # If the result is a DataFrame, re-synchronize the weights
-        if isinstance(result, pd.DataFrame):
-            new_weights = self.weights.reindex(result.index)
-            return MicroDataFrame(result, weights=new_weights)
-
-        # If the result is a Series (single column), wrap as MicroSeries
-        if isinstance(result, pd.Series):
-            return MicroSeries(result, weights=self.weights)
-
-        # Otherwise, the result is a scalar, so just return it
-        return result
+    def __getitem__(self, key):
+        return super().__getitem__(key)
 
     def catch_series_relapse(self) -> None:
         # In pandas 3.0+, we don't need to track series class changes since
         # __getitem__ always wraps columns as MicroSeries on access.
         pass
 
-    def __setattr__(self, key, value) -> None:
+    def __setattr__(self, key, value):
+        weights = self.__dict__.get("weights") if key == "index" else None
         super().__setattr__(key, value)
-        # No need to call catch_series_relapse in pandas 3.0+ since we wrap
-        # on access rather than store MicroSeries internally.
+        if weights is not None and len(weights) == len(self.index):
+            self.weights = weight_series(weights, self.index)
 
     def reset_index(
         self,
@@ -567,13 +494,7 @@ class MicroDataFrame(pd.DataFrame):
             return out
 
     def copy(self, deep: Optional[bool] = True) -> "MicroDataFrame":
-        res = super().copy(deep)
-        # super().copy() corrupts self's column types to plain Series.
-        # Restore them in O(N) instead of O(N²) by calling
-        # _link_all_weights once rather than per-column __setitem__.
-        self._link_all_weights()
-        res = MicroDataFrame(res, weights=self.weights.copy(deep))
-        return res
+        return super().copy(deep)
 
     def drop(
         self,
@@ -605,55 +526,15 @@ class MicroDataFrame(pd.DataFrame):
             dropped.
         :return: MicroDataFrame or None if inplace=True.
         """
-        row_drop = axis in (0, "index") or index is not None
-        if inplace:
-            # Snapshot the pre-drop weights keyed by the pre-drop index so
-            # we can reindex to the surviving rows after the drop.
-            pre_drop_weights = pd.Series(self.weights.values, index=self.index.copy())
-            # Perform in-place drop on the parent DataFrame
-            super().drop(
-                labels=labels,
-                axis=axis,
-                index=index,
-                columns=columns,
-                level=level,
-                inplace=True,
-                errors=errors,
-            )
-            if row_drop:
-                surviving = pre_drop_weights.reindex(self.index)
-                self.weights = pd.Series(
-                    surviving.values, index=self.index, dtype=float
-                )
-            else:
-                self.weights = pd.Series(
-                    pre_drop_weights.values, index=self.index, dtype=float
-                )
-            self._link_all_weights()
-            return None
-        else:
-            res = super().drop(
-                labels=labels,
-                axis=axis,
-                index=index,
-                columns=columns,
-                level=level,
-                inplace=False,
-                errors=errors,
-            )
-            if row_drop:
-                # Row drop: keep only the weights for surviving rows,
-                # in the order of the resulting DataFrame.
-                pre_drop_weights = pd.Series(self.weights.values, index=self.index)
-                new_weights = pre_drop_weights.reindex(res.index).values
-            else:
-                new_weights = self.weights.values
-            out = MicroDataFrame(res, weights=new_weights)
-            # Guard against the set_weights path building weights with a
-            # default RangeIndex, which would misalign against res.index
-            # and silently zero weighted aggregations.
-            out.weights = pd.Series(new_weights, index=out.index, dtype=float)
-            return out
+        return super().drop(
+            labels=labels,
+            axis=axis,
+            index=index,
+            columns=columns,
+            level=level,
+            inplace=inplace,
+            errors=errors,
+        )
 
     def merge(
         self,