Sparse fixes

.github/workflows/pytest.yml

@@ -8,6 +8,10 @@ on:
 
   pull_request:
 
+concurrency:
+  group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
+  cancel-in-progress: true
+
 jobs:
   linux:
     name: (${{ matrix.python-version }}, ${{ matrix.os }})

chainladder/core/base.py

@@ -608,6 +608,8 @@ def subtriangles(self):
         return [k for k, v in vars(self).items() if isinstance(v, TriangleBase)]
 
     def __array__(self):
+        if self.array_backend == "sparse":
+            return self.values.todense()
         return self.values
 
     def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):

chainladder/core/slice.py

@@ -541,20 +541,25 @@ def __setitem__(
             i = np.where(self.vdims == key)[0][0]
             # Case sparse backend.
             if self.array_backend == "sparse":
-                # Cast value to sparse backend.
-                value = cast("Triangle", value)
-                after = cast("COO", value.values)
+                # Unwrap a Triangle-valued assignment to its raw array. A raw
+                # COO array can also be passed directly, mirroring the numpy
+                # branch below.
+                if isinstance(value, TriangleSlicer):
+                    value = cast("Triangle", value)
+                    after = cast("COO", value.values)
+                else:
+                    after = cast("COO", value)
 
-                # Drop existing data where key matches, reassign coordinates.
-                before = self.drop(key).values
-                before = cast("COO", before)
-                bc = before.coords[1, :]
-                before.coords[1] = np.where(bc >= i, bc + 1, bc,)
+                # Filter out existing data at the target column directly from
+                # self.values' own coordinates.
+                before = cast("COO", self.values)
+                keep = before.coords[1, :] != i
 
                 # Append assigned data and new coordinates.
-                after.coords[1] = i
-                coords = np.concatenate((before.coords, after.coords), axis=1)
-                data = np.concatenate((before.data, after.data))
+                after_coords = after.coords.copy()
+                after_coords[1] = i
+                coords = np.concatenate((before.coords[:, keep], after_coords), axis=1)
+                data = np.concatenate((before.data[keep], after.data))
 
                 # Create new sparse matrix with updated coords and data, assign to backend array.
                 self.values = xp.COO(
@@ -577,7 +582,7 @@ def __setitem__(
                 value = self.iloc[:, 0] * 0 + value
             try:
                 self.values = xp.concatenate((self.values, value.values), axis=1)
-            except (ValueError, AttributeError):
+            except (ValueError, AttributeError, AssertionError):
                 # For misaligned triangle support.
                 conc = (self.values, (self.iloc[:, 0] * 0 + cast("Triangle", value)).values)
                 self.values = xp.concatenate(conc, axis=1)

chainladder/core/tests/test_slicing.py

@@ -116,7 +116,19 @@ def test_loc_ellipsis(clrd):
 
 def test_missing_first_lag(raa):
     x = raa.copy()
-    x.values[:, :, :, 0] = 0
+
+    def set_missing(values) -> None:
+        """
+        Sets the missing values for the first lag.
+        """
+        values[:, :, :, 0] = 0
+
+    if x.array_backend == "sparse":
+        # Sparse COO arrays don't support in-place item assignment.
+        with pytest.raises(TypeError, match="sparse backend"):
+            set_missing(x.values)
+        return
+    set_missing(x.values)
     x = x.sum(0)
     assert x.link_ratio.shape == (1, 1, 9, 9)
 
@@ -434,7 +446,8 @@ def test_setitem_virtual_column_numpy_backend(raa: Triangle) -> None:
     None
     """
     tri = raa.copy()
-    assert tri.array_backend == "numpy"
+    if tri.array_backend == "sparse":
+        pytest.skip("Test is specific to the numpy backend.")
     tri["double"] = lambda x: x["values"] * 2
     assert "double" in tri.columns
     assert tri["double"] == tri["values"] * 2
@@ -454,7 +467,8 @@ def test_setitem_value_backend_conversion(raa: Triangle) -> None:
     None
     """
     tri = raa.copy()
-    value = (tri["values"] * 2).set_backend("sparse")
+    other_backend = "numpy" if tri.array_backend == "sparse" else "sparse"
+    value = (tri["values"] * 2).set_backend(other_backend)
     assert tri.array_backend != value.array_backend
     tri["values"] = value
     assert tri.array_backend == raa.array_backend
@@ -475,15 +489,14 @@ def test_setitem_existing_column_triangle_value(raa: Triangle) -> None:
     None
     """
     tri = raa.copy()
-    assert tri.array_backend != "sparse"
     value = tri["values"] * 2
     tri["values"] = value
     assert tri["values"] == raa["values"] * 2
 
 
 def test_setitem_existing_column_array_value(raa: Triangle) -> None:
     """
-    Reassign an existing column to a raw array value on a non-sparse backend.
+    Reassign an existing column to a raw array value.
 
     Parameters
     ----------
@@ -495,7 +508,6 @@ def test_setitem_existing_column_array_value(raa: Triangle) -> None:
     None
     """
     tri = raa.copy()
-    assert tri.array_backend != "sparse"
     value = (tri["values"] * 3).values
     assert not isinstance(value, type(tri))
     tri["values"] = value
@@ -534,7 +546,8 @@ def test_setitem_new_column_misaligned_triangle(raa: Triangle) -> None:
     tri["misaligned"] = misaligned
     # Check the shape, new column should be added.
     assert tri.shape == (1, 2, 10, 10)
-    new_col = tri["misaligned"]
+    new_col = tri["misaligned"].set_backend("numpy")
+    misaligned = misaligned.set_backend("numpy")
     # Origin periods 1985 and prior should be nan.
     assert np.isnan(new_col.values[0, 0, :5, :]).all()
     # Origin periods 1986 and beyond should match.

chainladder/core/tests/test_triangle.py

@@ -303,7 +303,19 @@ def test_dropna_latest_diagonal(raa: Triangle) -> None:
     None
     """
     t = raa.copy()
-    t.values[:, :, 0, :] = np.nan
+
+    def set_first_origin_nan(values) -> None:
+        """
+        Sets the values for the first origin period to nan.
+        """
+        values[:, :, 0, :] = np.nan
+
+    if t.array_backend == "sparse":
+        # Sparse COO arrays don't support in-place item assignment.
+        with pytest.raises(TypeError, match="sparse backend"):
+            set_first_origin_nan(t.values)
+        return
+    set_first_origin_nan(t.values)
     result = t.latest_diagonal.dropna()
     assert result.shape == (1, 1, 9, 1)
     assert result.origin.min().year == 1982
@@ -575,8 +587,12 @@ def test_array_dunder(raa: Triangle) -> None:
     """
     arr = np.asarray(raa)
 
-    assert arr is raa.values
-    np.testing.assert_array_equal(np.array(raa), raa.values)
+    if raa.array_backend == "numpy":
+        # No conversion needed, so __array__ returns the same object.
+        assert arr is raa.values
+    else:
+        np.testing.assert_array_equal(arr, raa.values.todense())
+    np.testing.assert_array_equal(np.array(raa), raa.set_backend("numpy").values)
 
 
 def test_triangle_from_dataframe_interchange_protocol() -> None:

chainladder/utils/sparse.py

@@ -6,8 +6,18 @@
 from sparse import COO as COO
 from sparse import elemwise
 
+def _setitem_not_supported(self, key, value) -> None: # noqa
+    raise TypeError(
+        """
+        In-place item assignment (e.g. `triangle.values[...] = value`) is not 
+        supported on the sparse backend. Use numpy backend for in-place assignment.
+        """
+    )
+
+
 sp.isnan = np.isnan
 COO.nan = np.array([1.0, np.nan])[-1]
+COO.__setitem__ = _setitem_not_supported
 setattr(sp, 'testing', np.testing)
 sp.sqrt = np.sqrt
 sp.log = np.log