catii 1.0.5.dev1__tar.gz → 1.0.5.dev2__tar.gz

{catii-1.0.5.dev1/src/catii.egg-info → catii-1.0.5.dev2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: catii
-Version: 1.0.5.dev1
+Version: 1.0.5.dev2
 Summary: A library for N-dimensional categorical data.
 Home-page: https://github.com/Crunch-io/catii
 Author: Robert Brewer, Slobodan Ilić

{catii-1.0.5.dev1 → catii-1.0.5.dev2}/setup.py

@@ -28,7 +28,7 @@ with open("README.md", "r") as fh:
 
 setup(
     name="catii",
-    version="1.0.5.dev1",
+    version="1.0.5.dev2",
     author="Robert Brewer, Slobodan Ilić",
     author_email="dev@crunch.io",
     description="A library for N-dimensional categorical data.",

{catii-1.0.5.dev1 → catii-1.0.5.dev2}/src/catii/xfuncs.py

@@ -50,6 +50,8 @@ replace NaN values in the `values` array with 0 in that case. If you prefer
 `sum([])` to return 0, for example, without a second "validity" array,
 pass `return_missing_as=0`.
 """
+import functools
+import operator
 import warnings
 
 import numpy
@@ -1105,14 +1107,14 @@ class xfunc_quantile(xfunc):
                         if self.weights.shape:
                             w = self.weights[rowmask]
                         else:
-                            w = numpy.repeat(self.weights, len(seg))
+                            w = numpy.repeat(self.weights, len(rowmask))
                         qs[i] = self.weighted_quantile(seg, self.probability, w)
 
     def weighted_quantile(self, arr, probability, weights):
         def weighted_quantile_1d(a):
             w = weights
 
-            ind = numpy.lexsort((w, a))
+            ind = a.argsort()
             a = a[ind]
             w = w[ind]
 
@@ -1490,3 +1492,518 @@ class xfunc_covariance(xfunc):
             return covs, ~missings
         else:
             return covs
+
+
+# Default quantile probabilities
+DEFAULT_QUANTILES = [0.0, 0.25, 0.5, 0.75, 1.0]
+
+
+class xfunc_numeric_summary(xfunc):
+    """Calculate multiple numeric statistics in a single pass over the data.
+
+    This xfunc computes any combination of:
+    - mean: weighted arithmetic mean
+    - stddev: weighted standard deviation (ddof=1)
+    - valid_count: unweighted count of non-missing values
+    - sum: weighted sum
+    - quantiles: weighted quantiles at specified probabilities
+
+    The goal is to reduce both the number of cube calls and the number of
+    calculations by sharing intermediate computations.
+
+    Parameters
+    ----------
+    arr : array or (values, validity) tuple
+        NumPy array of numeric values, or a tuple of (values, validity) arrays.
+        Missing values are represented by NaN or validity=False.
+
+    weights : array or (weights, validity) tuple, optional
+        NumPy array of weight values.
+
+    ignore_missing : bool, default False
+        If False, any missing value in a bin causes that bin's result to be
+        missing. If True, missing values are skipped.
+
+    return_missing_as : scalar or tuple, default NaN
+        If NaN, missing results are returned as NaN. If a tuple like (0, False),
+        results are returned as (values, validity) tuples.
+
+    include : dict, optional
+        Specifies which statistics to compute. Keys are statistic names,
+        values are True or (for quantiles) a list of probabilities.
+        Default: all statistics with quantiles=[0, 0.25, 0.5, 0.75, 1.0]
+
+        Examples:
+            include={"mean": True, "stddev": True}
+            include={"quantiles": [0.1, 0.5, 0.9]}
+            include={"mean": True, "quantiles": [0.5]}
+
+    Returns
+    -------
+    dict
+        A dict with requested statistics as keys.
+        For dimensioned cubes, values are NumPy arrays shaped like the cube.
+        The "quantiles" value is itself a dict mapping probability -> array.
+    """
+
+    def __init__(
+        self,
+        arr,
+        weights=None,
+        ignore_missing=False,
+        return_missing_as=NaN,
+        include=None,
+    ):
+        # Parse include configuration
+        if include is None:
+            include = {
+                "mean": True,
+                "stddev": True,
+                "valid_count": True,
+                "sum": True,
+                "quantiles": DEFAULT_QUANTILES,
+            }
+
+        if not include:
+            raise ValueError("include must specify at least one statistic")
+
+        self.include_mean = include.get("mean", False)
+        self.include_stddev = include.get("stddev", False)
+        self.include_valid_count = include.get("valid_count", False)
+        self.include_sum = include.get("sum", False)
+        self.include_quantiles = include.get("quantiles", False)
+
+        if self.include_quantiles:
+            if self.include_quantiles is True:
+                self.quantile_probs = numpy.array(DEFAULT_QUANTILES, dtype=float)
+            else:
+                self.quantile_probs = numpy.array(self.include_quantiles, dtype=float)
+        else:
+            self.quantile_probs = None
+
+        # Process input array
+        arr, validity = as_separate_validity(arr)
+        self.shape = arr.shape[1:]
+        self.size = numpy.prod(self.shape, 0)
+
+        # Process weights
+        if weights is not None:
+            weights, weights_validity = as_separate_validity(weights)
+            validity = (validity.T & weights_validity).T
+            # Set negative weights to 0, like SAS EXCLNPWGT option
+            with numpy.errstate(invalid="ignore"):
+                neg_weights = weights < 0
+            if neg_weights.any():
+                weights = weights.copy()
+                weights[neg_weights] = 0
+
+        self.weights = weights
+        self.validity = validity
+        self.ignore_missing = ignore_missing
+        self.return_missing_as = return_missing_as
+
+        if isinstance(self.return_missing_as, tuple):
+            self.null = self.return_missing_as[0]
+        else:
+            self.null = self.return_missing_as
+
+        # Prepare data arrays for different computations
+        # For quantiles, keep original values with NaN for missing
+        self.arr = arr.astype(float).copy()
+        self.arr[~validity] = NaN
+
+        # For mean/sum, prepare weighted summables
+        if self.include_mean or self.include_sum or self.include_stddev:
+            summables = arr.astype(float).copy()
+            if weights is not None:
+                self.wsummables = (summables.T * weights).T
+            else:
+                self.wsummables = summables
+            self.wsummables[~validity] = 0
+
+            # For valid_count computation (weighted counts for mean denominator)
+            if weights is not None:
+                self.countables = (validity.astype(float).T * weights).T
+            else:
+                self.countables = validity.astype(float)
+            self.countables[~validity] = 0
+
+            # Unweighted summables for stddev variance calculation
+            self.summables = arr.astype(float).copy()
+            self.summables[~validity] = 0
+
+    def get_initial_regions(self, cube):
+        """Return empty NumPy arrays to fill."""
+        shape = cube.shape + self.shape
+        if not shape:
+            shape = (1,)
+
+        regions = {}
+
+        # valid_count is needed for mean denominator and stddev, so always allocate
+        regions["valid_count"] = numpy.zeros(shape, dtype=int)
+
+        # For weighted mean, we need weighted counts separately
+        if self.include_mean and self.weights is not None:
+            regions["weighted_counts"] = numpy.zeros(shape, dtype=float)
+
+        if self.include_mean or self.include_sum:
+            regions["sums"] = numpy.zeros(shape, dtype=float)
+
+        if self.include_stddev:
+            regions["stddev_values"] = numpy.full(shape, NaN, dtype=float)
+
+        if self.include_quantiles:
+            regions["quantiles"] = {
+                p: numpy.full(shape, NaN, dtype=float) for p in self.quantile_probs
+            }
+
+        # Track missing for output
+        if not self.ignore_missing:
+            regions["missing_counts"] = numpy.zeros(shape, dtype=int)
+
+        return regions
+
+    def calculate(self, cube):
+        """Return a dict with requested statistics.
+
+        This overrides the base calculate to properly handle the regions dict
+        format and compute coordinates from the cube.
+        """
+        regions = self.get_initial_regions(cube)
+
+        # Compute coordinates from cube dimensions (same logic as xcube.calculate)
+        strided_dims = cube.strided_dims()
+        coordinates = (
+            functools.reduce(operator.add, strided_dims) if strided_dims else None
+        )
+
+        self.fill(coordinates, regions)
+        return self.reduce(cube, regions)
+
+    def fill(self, coordinates, regions):
+        """Fill regions with statistics computed in a single pass."""
+        if coordinates is None:
+            self._fill_no_coordinates(regions)
+        else:
+            self._fill_with_coordinates(coordinates, regions)
+
+    def _fill_no_coordinates(self, regions):
+        """Fill for scalar cube (no dimensions)."""
+        valid_count = numpy.count_nonzero(self.validity, axis=0)
+        regions["valid_count"][:] = valid_count
+
+        if not self.ignore_missing:
+            regions["missing_counts"][:] = numpy.count_nonzero(~self.validity, axis=0)
+
+        # Mean and sum share weighted sums
+        if self.include_mean or self.include_sum:
+            value_sums = numpy.nansum(self.wsummables, axis=0)
+            regions["sums"][:] = value_sums
+
+        if self.include_mean and self.weights is not None:
+            regions["weighted_counts"][:] = numpy.sum(self.countables, axis=0)
+
+        if self.include_stddev:
+            # Compute stddev
+            if self.weights is not None:
+                weighted_counts = numpy.sum(self.countables, axis=0)
+                value_sums = numpy.nansum(self.wsummables, axis=0)
+                with numpy.errstate(divide="ignore", invalid="ignore"):
+                    means = value_sums / weighted_counts
+            else:
+                value_sums = numpy.nansum(self.wsummables, axis=0)
+                with numpy.errstate(divide="ignore", invalid="ignore"):
+                    means = value_sums / valid_count
+
+            squared_variances = (self.summables - means) ** 2
+            if self.weights is not None:
+                squared_variances = (squared_variances.T * self.weights).T
+            squared_variances[~self.validity] = 0
+            varsums = numpy.sum(squared_variances, axis=0)
+
+            with numpy.errstate(divide="ignore", invalid="ignore"):
+                if self.weights is None:
+                    regions["stddev_values"][:] = numpy.sqrt(varsums / (valid_count - 1))
+                else:
+                    weight_sums = numpy.sum(
+                        numpy.where(self.validity, self.weights, 0), axis=0
+                    )
+                    regions["stddev_values"][:] = numpy.sqrt(
+                        (varsums / weight_sums) * (valid_count / (valid_count - 1))
+                    )
+
+        if self.include_quantiles:
+            if self.weights is None:
+                # Unweighted: numpy handles multiple quantiles efficiently
+                for p in self.quantile_probs:
+                    if len(self.arr) and not numpy.all(numpy.isnan(self.arr)):
+                        if self.ignore_missing:
+                            regions["quantiles"][p][:] = numpy.nanquantile(
+                                self.arr, p, axis=0
+                            )
+                        else:
+                            regions["quantiles"][p][:] = numpy.quantile(
+                                self.arr, p, axis=0
+                            )
+            else:
+                # Weighted: sort once and compute all quantiles
+                results = self._weighted_quantiles_all(
+                    self.arr, self.quantile_probs, self.weights
+                )
+                for p, val in zip(self.quantile_probs, results):
+                    regions["quantiles"][p][:] = val
+
+
+    def _fill_with_coordinates(self, coordinates, regions):
+        """Fill for cube with dimensions."""
+        # Get size from valid_count region
+        size = regions["valid_count"].size
+
+        # valid_count is the unweighted count of non-missing values
+        valid_count = numpy.bincount(
+            coordinates, weights=self.validity.astype(float), minlength=size
+        )[:size].astype(int)
+        regions["valid_count"].flat[:] = valid_count
+
+        if not self.ignore_missing:
+            missing = numpy.bincount(
+                coordinates, weights=(~self.validity).astype(float), minlength=size
+            )[:size].astype(int)
+            regions["missing_counts"].flat[:] = missing
+
+        if self.include_mean or self.include_sum or self.include_stddev:
+            value_sums = numpy.bincount(
+                coordinates, weights=self.wsummables, minlength=size
+            )[:size]
+            if self.include_mean or self.include_sum:
+                regions["sums"].flat[:] = value_sums
+
+        if self.include_mean and self.weights is not None:
+            weighted_counts = numpy.bincount(
+                coordinates, weights=self.countables, minlength=size
+            )[:size]
+            regions["weighted_counts"].flat[:] = weighted_counts
+
+        if self.include_stddev:
+            if self.weights is not None:
+                weighted_counts = numpy.bincount(
+                    coordinates, weights=self.countables, minlength=size
+                )[:size]
+                with numpy.errstate(divide="ignore", invalid="ignore"):
+                    means = value_sums / weighted_counts
+            else:
+                with numpy.errstate(divide="ignore", invalid="ignore"):
+                    means = value_sums / valid_count
+
+            # Compute variance sums
+            squared_variances = (self.summables - means[coordinates]) ** 2
+            if self.weights is not None:
+                squared_variances = squared_variances * self.weights
+            squared_variances[~self.validity] = 0
+            varsums = numpy.bincount(
+                coordinates, weights=squared_variances, minlength=size
+            )[:size]
+
+            with numpy.errstate(divide="ignore", invalid="ignore"):
+                if self.weights is None:
+                    regions["stddev_values"].flat[:] = numpy.sqrt(
+                        varsums / (valid_count - 1)
+                    )
+                else:
+                    weight_sums = numpy.bincount(
+                        coordinates,
+                        weights=numpy.where(self.validity, self.weights, 0),
+                        minlength=size,
+                    )[:size]
+                    regions["stddev_values"].flat[:] = numpy.sqrt(
+                        (varsums / weight_sums) * (valid_count / (valid_count - 1))
+                    )
+
+        if self.include_quantiles:
+            # Quantiles require per-bin iteration (sorting-based)
+            for i, rowmask in self.bins(coordinates, size):
+                seg = self.arr[rowmask]
+                if len(seg) == 0 or numpy.all(numpy.isnan(seg)):
+                    continue
+
+                if self.weights is None:
+                    for p in self.quantile_probs:
+                        if self.ignore_missing:
+                            regions["quantiles"][p].flat[i] = numpy.nanquantile(
+                                seg, p, axis=0
+                            )
+                        else:
+                            regions["quantiles"][p].flat[i] = numpy.quantile(
+                                seg, p, axis=0
+                            )
+                else:
+                    w = self.weights[rowmask]
+                    for p in self.quantile_probs:
+                        regions["quantiles"][p].flat[i] = self._weighted_quantile(
+                            seg, p, w
+                        )
+
+    def _weighted_quantiles_all(self, arr, probabilities, weights):
+        """Compute multiple weighted quantiles with a single sort.
+
+        This is much faster than calling _weighted_quantile for each probability
+        because sorting is O(n log n) and dominates the computation.
+        """
+
+        def weighted_quantiles_1d(a):
+            w = weights
+
+            # Sort once
+            ind = a.argsort()
+            a = a[ind]
+            w = w[ind]
+
+            if self.ignore_missing:
+                missing = numpy.isnan(a) | numpy.isnan(w)
+                if numpy.any(missing):
+                    valid = ~missing
+                    a = a[valid]
+                    w = w[valid]
+
+            N = len(w)
+            if N == 0:
+                return [NaN] * len(probabilities)
+
+            cs = w.cumsum(axis=0)
+            xdiff = numpy.diff(a, append=[0], axis=0)
+
+            # Compute all quantiles from the same sorted data
+            results = []
+            for prob_val in probabilities:
+                prob = prob_val * cs[-1]
+                right = numpy.digitize(prob, cs)
+                left = (right - 1).clip(min=0)
+                with numpy.errstate(divide="ignore", invalid="ignore"):
+                    frac = (prob - cs[left]).clip(min=0) / w[right.clip(max=len(w) - 1)]
+                    results.append(a[left] + frac * xdiff[left])
+            return results
+
+        if arr.ndim == 1:
+            return weighted_quantiles_1d(arr)
+
+        # For 2D arrays, apply along axis 0
+        n_probs = len(probabilities)
+        result_shape = arr.shape[1:] if arr.ndim > 1 else (1,)
+        results = [numpy.full(result_shape, NaN, dtype=float) for _ in range(n_probs)]
+
+        for col_idx in range(arr.shape[1] if arr.ndim > 1 else 1):
+            col = arr[:, col_idx] if arr.ndim > 1 else arr
+            col_results = weighted_quantiles_1d(col)
+            for i, val in enumerate(col_results):
+                if arr.ndim > 1:
+                    results[i][col_idx] = val
+                else:
+                    results[i][:] = val
+
+        return results
+
+    def _weighted_quantile(self, arr, probability, weights):
+        """Compute weighted quantile for 1D array."""
+
+        def weighted_quantile_1d(a):
+            w = weights
+
+            ind = a.argsort()
+            a = a[ind]
+            w = w[ind]
+
+            if self.ignore_missing:
+                missing = numpy.isnan(a) | numpy.isnan(w)
+                if numpy.any(missing):
+                    valid = ~missing
+                    a = a[valid]
+                    w = w[valid]
+
+            N = len(w)
+            if N == 0:
+                return NaN
+
+            cs = w.cumsum(axis=0)
+            prob = probability * cs[-1]
+            right = numpy.digitize(prob, cs)
+            left = (right - 1).clip(min=0)
+            xdiff = numpy.diff(a, append=[0], axis=0)
+            with numpy.errstate(divide="ignore", invalid="ignore"):
+                frac = (prob - cs[left]).clip(min=0) / w[right.clip(max=len(w) - 1)]
+                return a[left] + frac * xdiff[left]
+
+        if arr.ndim == 1:
+            return weighted_quantile_1d(arr)
+        return numpy.apply_along_axis(weighted_quantile_1d, 0, arr)
+
+    def reduce(self, cube, regions):
+        """Reduce regions to final output."""
+        result = {}
+
+        valid_count = regions["valid_count"]
+        if not self.ignore_missing:
+            missing_counts = regions["missing_counts"]
+            has_missing = missing_counts != 0
+        else:
+            has_missing = numpy.zeros_like(valid_count, dtype=bool)
+
+        no_data = valid_count == 0
+
+        if self.include_mean:
+            if self.weights is not None:
+                denominator = regions["weighted_counts"]
+            else:
+                denominator = valid_count
+            with numpy.errstate(divide="ignore", invalid="ignore"):
+                mean = regions["sums"] / denominator
+            output_is_missing = no_data | has_missing
+            mean = self.adjust_zeros(mean, self.null, condition=output_is_missing)
+            if isinstance(self.return_missing_as, tuple):
+                result["mean"] = (mean, ~output_is_missing)
+            else:
+                result["mean"] = mean
+
+        if self.include_stddev:
+            stddev = regions["stddev_values"]
+            # stddev needs at least 2 values
+            output_is_missing = (valid_count < 2) | has_missing
+            stddev = self.adjust_zeros(stddev, self.null, condition=output_is_missing)
+            if isinstance(self.return_missing_as, tuple):
+                result["stddev"] = (stddev, ~output_is_missing)
+            else:
+                result["stddev"] = stddev
+
+        if self.include_valid_count:
+            # valid_count is always returned, even for empty bins
+            if isinstance(self.return_missing_as, tuple):
+                result["valid_count"] = (valid_count, ~no_data)
+            else:
+                result["valid_count"] = valid_count
+
+        if self.include_sum:
+            sum_values = regions["sums"]
+            output_is_missing = no_data | has_missing
+            sum_values = self.adjust_zeros(
+                sum_values, self.null, condition=output_is_missing
+            )
+            if isinstance(self.return_missing_as, tuple):
+                result["sum"] = (sum_values, ~output_is_missing)
+            else:
+                result["sum"] = sum_values
+
+        if self.include_quantiles:
+            quantiles_result = {}
+            output_is_missing = no_data | has_missing
+            for p in self.quantile_probs:
+                q = regions["quantiles"][p]
+                if isinstance(self.return_missing_as, tuple):
+                    missings = numpy.isnan(q) | output_is_missing
+                    q = q.copy()
+                    q[missings] = self.null
+                    quantiles_result[p] = (q, ~missings)
+                else:
+                    quantiles_result[p] = q
+            result["quantiles"] = quantiles_result
+
+        return result

{catii-1.0.5.dev1 → catii-1.0.5.dev2/src/catii.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: catii
-Version: 1.0.5.dev1
+Version: 1.0.5.dev2
 Summary: A library for N-dimensional categorical data.
 Home-page: https://github.com/Crunch-io/catii
 Author: Robert Brewer, Slobodan Ilić