microdf-python 1.2.2__py3-none-any.whl → 1.3.0__py3-none-any.whl

microdf/microseries.py

@@ -9,6 +9,46 @@ import pandas as pd
 logger = logging.getLogger(__name__)
 
 
+def _weighted_top_share(
+    values: np.ndarray, weights: np.ndarray, top_x_pct: float
+) -> float:
+    """Share of the sum held by the top ``top_x_pct`` of weight.
+
+    Sort by value ascending, cumulate the weight, pick the slice from
+    the top that covers exactly ``top_x_pct`` of total weight, and
+    distribute the tied-at-cutoff row proportionally so constant values
+    return exactly ``top_x_pct`` rather than 1.0.
+    """
+    if top_x_pct <= 0:
+        return 0.0
+    if top_x_pct >= 1:
+        return 1.0
+    total_weight = weights.sum()
+    total_sum = float((values * weights).sum())
+    if total_weight == 0 or total_sum == 0:
+        return np.nan
+    # Ascending sort; the "top" cutoff is the final ``top_x_pct`` of
+    # cumulative weight.
+    order = np.argsort(values, kind="mergesort")
+    v = values[order]
+    w = weights[order]
+    # Cumulative weight from the bottom up.
+    cum_w = np.cumsum(w)
+    target_bottom_weight = total_weight * (1.0 - top_x_pct)
+    # searchsorted(cum_w, target, side="right") gives the first index
+    # whose cumulative weight exceeds the bottom cutoff.
+    k = int(np.searchsorted(cum_w, target_bottom_weight, side="right"))
+    # Rows strictly above the cutoff contribute all of their weight.
+    if k >= len(v):
+        return 0.0
+    top_sum = float((v[k + 1 :] * w[k + 1 :]).sum())
+    # Row k straddles the cutoff; include the fraction of its weight
+    # that lies above the cutoff so ties don't double-count.
+    partial_weight = cum_w[k] - target_bottom_weight
+    top_sum += float(v[k] * partial_weight)
+    return top_sum / total_sum
+
+
 class MicroSeries(pd.Series):
     def __init__(self, *args, weights: np.array = None, **kwargs):
         """A Series-inheriting class for weighted microdata. Weights can be
@@ -64,23 +104,13 @@ class MicroSeries(pd.Series):
         )
         return super().to_numpy(*args, **kwargs)
 
-    def weighted_function(fn: Callable) -> Callable:
-        @wraps(fn)
-        def safe_fn(*args, **kwargs):
-            try:
-                return fn(*args, **kwargs)
-            except ZeroDivisionError:
-                return np.NaN
-
-        return safe_fn
-
-    @weighted_function
     def scalar_function(fn: Callable) -> Callable:
+        """Decorator marking ``fn`` as returning a scalar (float)."""
         fn._rtype = float
         return fn
 
-    @weighted_function
     def vector_function(fn: Callable) -> Callable:
+        """Decorator marking ``fn`` as returning a pandas Series."""
         fn._rtype = pd.Series
         return fn
 
@@ -97,7 +127,9 @@ class MicroSeries(pd.Series):
         if weights is None:
             if len(self) > 0:
                 self.weights = pd.Series(
-                    np.ones_like(self._values), dtype=float
+                    np.ones_like(self._values),
+                    index=self.index,
+                    dtype=float,
                 )
         else:
             if len(weights) != len(self):
@@ -109,7 +141,14 @@ class MicroSeries(pd.Series):
             if preserve_old and self.weights is not None:
                 self["old_weights"] = self.weights
 
-            self.weights = pd.Series(weights, dtype=float)
+            # Align weights to self.index so element-wise operations such
+            # as self.multiply(self.weights) (used by .sum(), .weight())
+            # don't silently produce all-NaN when the caller uses a
+            # non-default index. If a pandas Series is passed in, strip
+            # its index first so we position-align rather than label-align.
+            if isinstance(weights, pd.Series):
+                weights = weights.values
+            self.weights = pd.Series(np.asarray(weights), index=self.index, dtype=float)
 
     def nullify_weights(self) -> None:
         """Set all weights to 1, effectively making the Series unweighted.
@@ -138,12 +177,22 @@ class MicroSeries(pd.Series):
         return self.multiply(self.weights).sum()
 
     @scalar_function
-    def count(self) -> float:
+    def count(self, skipna: bool = True) -> float:
         """Calculates the weighted count of the MicroSeries.
 
+        By default skips NaN values (matching pandas ``Series.count``).
+
+        :param skipna: Exclude NaN values (default True). If False, the
+            weighted count of every row is returned.
+        :type skipna: bool
         :returns: The weighted count.
+        :rtype: float
         """
-        return self.weights.sum()
+        weights = np.asarray(self.weights.values, dtype=float)
+        if not skipna:
+            return float(weights.sum())
+        mask = ~pd.isna(self._values)
+        return float(weights[mask].sum())
 
     @scalar_function
     def mean(self, skipna: bool = True) -> float:
@@ -173,6 +222,89 @@ class MicroSeries(pd.Series):
 
         return np.average(values, weights=weights)
 
+    def _weighted_variance(self, ddof: int = 1, skipna: bool = True) -> float:
+        """Frequency-weighted variance.
+
+        Uses ``sum(w * (x - wmean)**2) / (sum(w) - ddof)``. With
+        ``ddof=0`` this is the population variance; with ``ddof=1`` it
+        is Bessel-corrected assuming the weights are frequency counts —
+        matching ``np.var(..., ddof=ddof)`` on a replicated sample.
+        """
+        values = np.asarray(self._values, dtype=float)
+        weights = np.asarray(self.weights.values, dtype=float)
+        if skipna:
+            mask = ~np.isnan(values)
+            values = values[mask]
+            weights = weights[mask]
+        elif np.isnan(values).any():
+            return np.nan
+        total_w = weights.sum()
+        if total_w == 0 or total_w - ddof <= 0:
+            return np.nan
+        mean = np.average(values, weights=weights)
+        return float((weights * (values - mean) ** 2).sum() / (total_w - ddof))
+
+    @scalar_function
+    def var(self, ddof: int = 1, skipna: bool = True) -> float:
+        """Calculates the weighted variance of the MicroSeries.
+
+        Treats weights as frequency counts (``sum(w) - ddof`` in the
+        denominator) so that with integer weights the result matches
+        ``np.var`` on the replicated sample.
+
+        :param ddof: Delta degrees of freedom (default 1).
+        :param skipna: Exclude NaN values (default True).
+        :returns: The weighted variance.
+        :rtype: float
+        """
+        return self._weighted_variance(ddof=ddof, skipna=skipna)
+
+    @scalar_function
+    def std(self, ddof: int = 1, skipna: bool = True) -> float:
+        """Calculates the weighted standard deviation of the MicroSeries.
+
+        :param ddof: Delta degrees of freedom (default 1).
+        :param skipna: Exclude NaN values (default True).
+        :returns: The weighted standard deviation.
+        :rtype: float
+        """
+        v = self._weighted_variance(ddof=ddof, skipna=skipna)
+        return float(np.sqrt(v)) if np.isfinite(v) else v
+
+    def cov(self, other, *args, **kwargs):
+        """Pandas ``cov`` — **unweighted**.
+
+        MicroSeries does not yet compute weighted covariance. Emits a
+        ``UserWarning`` so callers aren't silently given an unweighted
+        number after ``.sum()`` and ``.mean()`` worked as expected. See
+        issue tracker for a weighted implementation.
+        """
+        warnings.warn(
+            "MicroSeries.cov() falls through to pandas and is "
+            "unweighted. Use MicroSeries.var()/std() for weighted "
+            "second moments, or compute covariance manually with the "
+            "weights.",
+            UserWarning,
+            stacklevel=2,
+        )
+        return super().cov(other, *args, **kwargs)
+
+    def corr(self, other, *args, **kwargs):
+        """Pandas ``corr`` — **unweighted**.
+
+        MicroSeries does not yet compute weighted correlation. Emits a
+        ``UserWarning`` so callers aren't silently given an unweighted
+        number. See issue tracker for a weighted implementation.
+        """
+        warnings.warn(
+            "MicroSeries.corr() falls through to pandas and is "
+            "unweighted. Compute correlation manually with the weights "
+            "if you need the survey-weighted value.",
+            UserWarning,
+            stacklevel=2,
+        )
+        return super().corr(other, *args, **kwargs)
+
     def quantile(self, q: np.array) -> pd.Series:
         """Calculates weighted quantiles of the MicroSeries.
 
@@ -189,9 +321,23 @@ class MicroSeries(pd.Series):
         values = np.array(self._values)
         quantiles = np.atleast_1d(q)
         sample_weight = np.array(self.weights)
-        assert np.all(quantiles >= 0) and np.all(
-            quantiles <= 1
-        ), "quantiles should be in [0, 1]"
+        assert np.all(quantiles >= 0) and np.all(quantiles <= 1), (
+            "quantiles should be in [0, 1]"
+        )
+        # Drop zero-weight rows before sorting. Without this, q=0 (and
+        # internal plateaus of zero weight) picked a value with 0 weight
+        # that should have been skipped by the inverse CDF. E.g.
+        # MicroSeries([10, 20, 30], weights=[0, 1, 1]).quantile(0)
+        # returned 10 instead of 20.
+        nonzero = sample_weight > 0
+        if not nonzero.any():
+            return (
+                np.nan
+                if np.array(q).shape == ()
+                else pd.Series(np.full(len(quantiles), np.nan), index=quantiles)
+            )
+        values = values[nonzero]
+        sample_weight = sample_weight[nonzero]
         sorter = np.argsort(values)
         values = values[sorter]
         sample_weight = sample_weight[sorter]
@@ -199,11 +345,7 @@ class MicroSeries(pd.Series):
         cumsum_normalized = cumsum / cumsum[-1]
         result = np.array(
             [
-                values[
-                    min(
-                        np.searchsorted(cumsum_normalized, qi), len(values) - 1
-                    )
-                ]
+                values[min(np.searchsorted(cumsum_normalized, qi), len(values) - 1)]
                 for qi in quantiles
             ]
         )
@@ -224,58 +366,83 @@ class MicroSeries(pd.Series):
     def gini(self, negatives: Optional[str] = None) -> float:
         """Calculates Gini index.
 
-        :param negatives: An optional string indicating how to treat negative
-            values of x:
+        :param negatives: An optional string indicating how to treat
+            negative values of x:
             'zero' replaces negative values with zeroes.
             'shift' subtracts the minimum value from all values of x,
             when this minimum is negative. That is, it adds the absolute
             minimum value.
             Defaults to None, which leaves negative values as they are.
-        :type q: str
+        :type negatives: str
         :returns: Gini index.
         :rtype: float
         """
         x = np.array(self).astype("float")
+        w = np.asarray(self.weights.values, dtype=float)
         if negatives == "zero":
-            x[x < 0] = 0
-        if negatives == "shift" and np.amin(x) < 0:
-            x -= np.amin(x)
-        if (self.weights != np.ones(len(self))).any():  # Varying weights.
-            sorted_indices = np.argsort(self)
-            sorted_x = np.array(self[sorted_indices])
-            sorted_w = np.array(self.weights[sorted_indices])
-            cumw = np.cumsum(sorted_w)
-            cumxw = np.cumsum(sorted_x * sorted_w)
-            return np.sum(cumxw[1:] * cumw[:-1] - cumxw[:-1] * cumw[1:]) / (
-                cumxw[-1] * cumw[-1]
+            x = np.where(x < 0, 0.0, x)
+        elif negatives == "shift" and len(x) > 0 and np.amin(x) < 0:
+            x = x - np.amin(x)
+        elif negatives is not None:
+            raise ValueError(
+                f"Unknown negatives option {negatives!r}; expected "
+                "'zero', 'shift', or None."
             )
-        else:
-            sorted_x = np.sort(self)
-            n = len(x)
-            cumxw = np.cumsum(sorted_x)
-            # The above formula, with all weights equal to 1 simplifies to:
-            return (n + 1 - 2 * np.sum(cumxw) / cumxw[-1]) / n
+
+        if len(x) == 0:
+            return np.nan
+        if np.any(x < 0):
+            # The Lorenz-based formula assumes non-negative values; with
+            # negatives it can return values outside [0, 1].
+            warnings.warn(
+                "gini() called on data containing negative values; the "
+                "result is not guaranteed to lie in [0, 1]. Pass "
+                "negatives='zero' or negatives='shift' to handle them.",
+                UserWarning,
+                stacklevel=2,
+            )
+
+        # Short-circuit degenerate cases so we don't divide by zero.
+        total = float((x * w).sum())
+        if total == 0:
+            return 0.0
+
+        sorter = np.argsort(x, kind="mergesort")
+        sorted_x = x[sorter]
+        sorted_w = w[sorter]
+        cumw = np.cumsum(sorted_w)
+        cumxw = np.cumsum(sorted_x * sorted_w)
+        # Trapezoidal approximation of the area under the Lorenz curve.
+        return float(
+            np.sum(cumxw[1:] * cumw[:-1] - cumxw[:-1] * cumw[1:])
+            / (cumxw[-1] * cumw[-1])
+        )
 
     @scalar_function
     def top_x_pct_share(self, top_x_pct: float) -> float:
         """Calculates top x% share.
 
+        Uses a cumulative-weight sort so that rows tied at the cutoff
+        contribute proportionally rather than all-or-nothing. With
+        constant values this correctly returns ``top_x_pct`` itself.
+
         :param top_x_pct: Decimal between 0 and 1 of the top %, e.g. 0.1,
             0.001.
         :type top_x_pct: float
         :returns: The weighted share held by the top x%.
         :rtype: float
         """
-        threshold = self.quantile(1 - top_x_pct)
-        top_x_pct_sum = self[self >= threshold].sum()
-        total_sum = self.sum()
-        return top_x_pct_sum / total_sum
+        return _weighted_top_share(
+            np.asarray(self._values, dtype=float),
+            np.asarray(self.weights.values, dtype=float),
+            float(top_x_pct),
+        )
 
     @scalar_function
     def bottom_x_pct_share(self, bottom_x_pct: float) -> float:
         """Calculates bottom x% share.
 
-        :param bottom_x_pct: Decimal between 0 and 1 of the top %, e.g. 0.1,
+        :param bottom_x_pct: Decimal between 0 and 1 of the bottom %, e.g. 0.1,
             0.001.
         :type bottom_x_pct: float
         :returns: The weighted share held by the bottom x%.
@@ -350,19 +517,44 @@ class MicroSeries(pd.Series):
 
     @vector_function
     def rank(self, pct: Optional[bool] = False) -> pd.Series:
-        weights_sum = self.weights.values.sum()
+        """Weighted rank of each element.
+
+        Each element's rank is the cumulative weight of all values that
+        are less than or equal to it. Tied values therefore share the
+        same rank, so downstream bucketing (``decile_rank``,
+        ``quintile_rank``, etc.) lands tied rows in the same bucket.
+
+        :param pct: If True, divide ranks by the total weight so they
+            lie in ``(0, 1]``.
+        :type pct: bool
+        :returns: MicroSeries of ranks aligned to ``self``.
+        :rtype: MicroSeries
+        """
+        weights_sum = np.asarray(self.weights.values, dtype=float).sum()
         if weights_sum == 0:
             raise ZeroDivisionError(
                 "Cannot calculate rank with zero total weight. "
-                "All weights in the MicroSeries are zero, which would result "
-                "in division by zero."
+                "All weights in the MicroSeries are zero, which would "
+                "result in division by zero."
             )
 
-        order = np.argsort(self._values)
-        inverse_order = np.argsort(order)
-        ranks = np.array(self.weights.values)[order].cumsum()[inverse_order]
+        values = np.asarray(self._values)
+        weights = np.asarray(self.weights.values, dtype=float)
+        order = np.argsort(values, kind="mergesort")
+        sorted_values = values[order]
+        sorted_weights = weights[order]
+        cum_w = np.cumsum(sorted_weights)
+        # Max rank semantics: every tied group gets the cumulative
+        # weight at the *end* of the group, so ties share one rank.
+        # searchsorted(side='right') on the sorted values finds the
+        # index just past each tied block in sort order.
+        group_end = np.searchsorted(sorted_values, sorted_values, side="right") - 1
+        sorted_ranks = cum_w[group_end]
+        # Invert the sort to put ranks back into the caller's order.
+        inverse_order = np.argsort(order, kind="mergesort")
+        ranks = sorted_ranks[inverse_order]
         if pct:
-            ranks /= weights_sum
+            ranks = ranks / weights_sum
             ranks = np.where(ranks > 1.0, 1.0, ranks)
         return MicroSeries(ranks, index=self.index, weights=self.weights)
 
@@ -454,9 +646,7 @@ class MicroSeries(pd.Series):
             return MicroSeries(res, weights=self.weights)
         return self
 
-    def round(
-        self, decimals: Optional[int] = 0, *args, **kwargs
-    ) -> "MicroSeries":
+    def round(self, decimals: Optional[int] = 0, *args, **kwargs) -> "MicroSeries":
         res = super().round(decimals=decimals, *args, **kwargs)
         return MicroSeries(res, weights=self.weights)
 
@@ -488,14 +678,10 @@ class MicroSeries(pd.Series):
     def __mul__(self, other: Union[int, float, pd.Series]) -> "MicroSeries":
         return MicroSeries(super().__mul__(other), weights=self.weights)
 
-    def __floordiv__(
-        self, other: Union[int, float, pd.Series]
-    ) -> "MicroSeries":
+    def __floordiv__(self, other: Union[int, float, pd.Series]) -> "MicroSeries":
         return MicroSeries(super().__floordiv__(other), weights=self.weights)
 
-    def __truediv__(
-        self, other: Union[int, float, pd.Series]
-    ) -> "MicroSeries":
+    def __truediv__(self, other: Union[int, float, pd.Series]) -> "MicroSeries":
         return MicroSeries(super().__truediv__(other), weights=self.weights)
 
     def __mod__(self, other: Union[int, float, pd.Series]) -> "MicroSeries":
@@ -525,14 +711,10 @@ class MicroSeries(pd.Series):
     def __rmul__(self, other: Union[int, float, pd.Series]) -> "MicroSeries":
         return MicroSeries(super().__rmul__(other), weights=self.weights)
 
-    def __rfloordiv__(
-        self, other: Union[int, float, pd.Series]
-    ) -> "MicroSeries":
+    def __rfloordiv__(self, other: Union[int, float, pd.Series]) -> "MicroSeries":
         return MicroSeries(super().__rfloordiv__(other), weights=self.weights)
 
-    def __rtruediv__(
-        self, other: Union[int, float, pd.Series]
-    ) -> "MicroSeries":
+    def __rtruediv__(self, other: Union[int, float, pd.Series]) -> "MicroSeries":
         return MicroSeries(super().__rtruediv__(other), weights=self.weights)
 
     def __rmod__(self, other: Union[int, float, pd.Series]) -> "MicroSeries":
@@ -585,17 +767,13 @@ class MicroSeries(pd.Series):
     def __imul__(self, other: Union[int, float, pd.Series]) -> "MicroSeries":
         return MicroSeries(super().__imul__(other), weights=self.weights)
 
-    def __ifloordiv__(
-        self, other: Union[int, float, pd.Series]
-    ) -> "MicroSeries":
+    def __ifloordiv__(self, other: Union[int, float, pd.Series]) -> "MicroSeries":
         return MicroSeries(super().__ifloordiv__(other), weights=self.weights)
 
     def __idiv__(self, other: Union[int, float, pd.Series]) -> "MicroSeries":
         return MicroSeries(super().__idiv__(other), weights=self.weights)
 
-    def __itruediv__(
-        self, other: Union[int, float, pd.Series]
-    ) -> "MicroSeries":
+    def __itruediv__(self, other: Union[int, float, pd.Series]) -> "MicroSeries":
         return MicroSeries(super().__itruediv__(other), weights=self.weights)
 
     def __imod__(self, other: Union[int, float, pd.Series]) -> "MicroSeries":
@@ -666,16 +844,12 @@ class MicroSeriesGroupBy(pd.core.groupby.generic.SeriesGroupBy):
     def _init(self):
         def _weighted_agg(name) -> Callable:
             def via_micro_series(row, *args, **kwargs):
-                return getattr(MicroSeries(row.a, weights=row.w), name)(
-                    *args, **kwargs
-                )
+                return getattr(MicroSeries(row.a, weights=row.w), name)(*args, **kwargs)
 
             fn = getattr(MicroSeries, name)
 
             @wraps(fn)
-            def _weighted_agg_fn(
-                *args, **kwargs
-            ) -> Union[pd.Series, pd.DataFrame]:
+            def _weighted_agg_fn(*args, **kwargs) -> Union[pd.Series, pd.DataFrame]:
                 arrays = self.apply(np.array)
                 weights = self.weights.apply(np.array)
                 df = pd.DataFrame(dict(a=arrays, w=weights))