[#186590795]: Implement squared weights (iter #1)

src/cr/cube/cube.py

@@ -523,6 +523,15 @@ def weighted_valid_counts(self) -> Optional[np.ndarray]:
             self._valid_idxs
         ].astype(np.float64)
 
+    @lazyproperty
+    def weighted_squared_counts(self) -> Optional[np.ndarray]:
+        """Optional float64 ndarray of weighted_squared_counts if the measure exists."""
+        if self._measures.weighted_squared_counts is None:
+            return None
+        return self._measures.weighted_squared_counts.raw_cube_array[
+            self._valid_idxs
+        ].astype(np.float64)
+
     @lazyproperty
     def _all_dimensions(self) -> list:
         """List of all dimensions (not just user-apparent ones) for this cube."""
@@ -847,6 +856,14 @@ def weighted_valid_counts(self) -> 'Optional["_WeightedValidCountsMeasure"]':
         )
         return valid_counts if valid_counts.raw_cube_array is not None else None
 
+    @lazyproperty
+    def weighted_squared_counts(self):
+        """Return object of class for representing squared weights."""
+        squared_counts = _WeightedSquaredCountsMeasure(
+            self._cube_dict, self._all_dimensions, self._cube_idx_arg
+        )
+        return squared_counts if squared_counts.raw_cube_array is not None else None
+
 
 class _BaseMeasure:
     """Base class for measure objects."""
@@ -1100,6 +1117,20 @@ def _flat_values(self) -> Optional[np.ndarray]:
         return np.array(weighted_counts, dtype=np.float64)
 
 
+class _WeightedSquaredCountsMeasure(_BaseMeasure):
+    """Weighted squared counts for cube."""
+
+    @lazyproperty
+    def _flat_values(self) -> Optional[np.ndarray]:
+        """Optional 1D np.ndarray of np.float64 weighted squared counts."""
+        squared_counts = (
+            self._cube_dict["result"]["measures"]
+            .get("weighted_squared_count", {})
+            .get("data", [])
+        )
+        return np.array(squared_counts, dtype=np.float64) if squared_counts else None
+
+
 class _WeightedValidCountsMeasure(_BaseMeasure):
     """Weighted Valid counts for cube."""
 

src/cr/cube/enums.py

@@ -163,6 +163,7 @@ class CUBE_MEASURE(enum.Enum):
     VALID_OVERLAP = "valid_overlap"
     UNWEIGHTED_VALID_COUNT = "valid_count_unweighted"
     WEIGHTED_VALID_COUNT = "valid_count_weighted"
+    WEIGHTED_SQUARED_COUNT = "weighted_squared_count"
 
 
 NUMERIC_CUBE_MEASURES = frozenset(

tests/fixtures/squared-weights.json

@@ -0,0 +1,127 @@
+{
+    "result": {
+        "counts": [
+            2,
+            0,
+            2,
+            2,
+            2,
+            0,
+            0,
+            0,
+            1,
+            1,
+            0
+        ],
+        "dimensions": [
+            {
+                "derived": false,
+                "references": {},
+                "type": {
+                    "categories": [
+                        {
+                            "id": -1,
+                            "missing": true,
+                            "name": "No Data",
+                            "numeric_value": null
+                        },
+                        {
+                            "id": 0,
+                            "missing": false,
+                            "name": "Grade_0",
+                            "numeric_value": null
+                        },
+                        {
+                            "id": 1,
+                            "missing": false,
+                            "name": "Grade_1",
+                            "numeric_value": null
+                        },
+                        {
+                            "id": 2,
+                            "missing": false,
+                            "name": "Grade_2",
+                            "numeric_value": null
+                        },
+                        {
+                            "id": 3,
+                            "missing": false,
+                            "name": "Grade_3",
+                            "numeric_value": null
+                        },
+                        {
+                            "id": 4,
+                            "missing": false,
+                            "name": "Grade_4",
+                            "numeric_value": null
+                        },
+                        {
+                            "id": 5,
+                            "missing": false,
+                            "name": "Grade_5",
+                            "numeric_value": null
+                        },
+                        {
+                            "id": 6,
+                            "missing": false,
+                            "name": "Grade_6",
+                            "numeric_value": null
+                        },
+                        {
+                            "id": 7,
+                            "missing": false,
+                            "name": "Grade_7",
+                            "numeric_value": null
+                        },
+                        {
+                            "id": 8,
+                            "missing": false,
+                            "name": "Grade_8",
+                            "numeric_value": null
+                        },
+                        {
+                            "id": 9,
+                            "missing": false,
+                            "name": "Grade_9",
+                            "numeric_value": null
+                        }
+                    ],
+                    "class": "categorical",
+                    "ordinal": false
+                }
+            }
+        ],
+        "measures": {
+            "weighted_squared_count": {
+                "data": [
+                    324.0,
+		    0.0,
+                    148.0,
+                    212.0,
+                    292.0,
+		    0.0,
+		    0.0,
+		    0.0,
+                    64.0,
+                    100.0,
+		    0.0
+                ],
+                "metadata": {
+                    "derived": true,
+                    "references": {},
+                    "type": {
+                        "class": "numeric",
+                        "integer": false,
+                        "missing_reasons": {
+                            "NaN": -8,
+                            "No Data": -1
+                        },
+                        "missing_rules": {}
+                    }
+                },
+                "n_missing": 2
+            }
+        },
+        "n": 10
+    }
+}

tests/integration/test_cube.py

@@ -193,6 +193,21 @@ def it_does_not_get_fooled_into_single_mr_cats_dim(self):
         assert cube.dimension_types == (DT.LOGICAL,)
         assert cube.partitions[0].counts.tolist() == [200, 100]
 
+    def it_provides_squared_weights_counts(self):
+        cube = Cube(CR.SQUARED_WEIGHTS)
+        assert cube.weighted_squared_counts.tolist() == [
+            0.0,
+            148.0,
+            212.0,
+            292.0,
+            0.0,
+            0.0,
+            0.0,
+            64.0,
+            100.0,
+            0.0,
+        ]
+
 
 class DescribeIntegrated_Measures:
     """Integration-tests that exercise the `cr.cube.cube._Measures` object."""