Add get_full_source

Author: ValerianRey

src/torchjd/sparse/_structured_sparse_tensor.py

@@ -243,6 +243,43 @@ def unwrap_to_dense(t: Tensor):
         return t
 
 
+def get_full_source(source: list[int], destination: list[int], ndim: int) -> list[int]:
+    """
+    Doing a movedim with source and destination is always equivalent to doing a movedim with
+    [0, 1, ..., ndim-1] (aka "full_destination") as destination, and the "full_source" as source.
+
+    This function computes the full_source based on a source and destination.
+
+    Example:
+    source=[2, 4]
+    destination=[0, 3]
+    ndim=5
+
+    full_source = [2, 0, 1, 4, 3]
+    full_destination = [0, 1, 2, 3, 4]
+    """
+
+    idx = torch.full((ndim,), -1, dtype=torch.int64)
+    idx[destination] = tensor(source)
+    source_set = set(source)
+    idx[idx.eq(-1)] = tensor([i for i in range(ndim) if i not in source_set])
+
+    # source_mask = torch.zeros(ndim, dtype=torch.bool)
+    # destination_mask = torch.zeros(ndim, dtype=torch.bool)
+    # source_mask[source] = True
+    # destination_mask[destination] = True
+    #
+    # destination_cumsum = torch.cumsum(destination_mask, dim=0)
+    # source_cumsum = torch.cumsum(source_mask, dim=0)
+    # base = arange(ndim, dtype=torch.int64)
+    #
+    # idx = torch.empty((ndim,), dtype=torch.int64)
+    # idx[destination_mask] = tensor(source)
+    # idx[~destination_mask] = base[~destination_mask] - destination_cumsum[~destination_mask] + source_cumsum[:ndim - len(source)]
+
+    return idx.tolist()
+
+
 def fix_dim_of_size_1(physical: Tensor, strides: Tensor) -> tuple[Tensor, Tensor]:
     is_of_size_1 = torch.tensor([s == 1 for s in physical.shape])
     return physical.squeeze(), strides[:, ~is_of_size_1]

tests/unit/sparse/test_structured_sparse_tensor.py

@@ -15,6 +15,7 @@
     StructuredSparseTensor,
     fix_ungrouped_dims,
     fix_zero_stride_columns,
+    get_full_source,
     get_groupings,
 )
 
@@ -349,6 +350,26 @@ def test_unsquash_pdim(
     assert torch.equal(new_strides, expected_strides)
 
 
+@mark.parametrize(
+    [
+        "source",
+        "destination",
+        "ndim",
+    ],
+    [
+        ([2, 4], [0, 3], 5),
+        ([5, 3, 6], [2, 0, 5], 8),
+    ],
+)
+def test_get_column_indices(source: list[int], destination: list[int], ndim: int):
+    # TODO: this test should be improved / removed. It creates quite big tensors for nothing.
+
+    t = randn_(list(torch.randint(3, 8, size=(ndim,))))
+    full_destination = list(range(ndim))
+    full_source = get_full_source(source, destination, ndim)
+    assert torch.equal(t.movedim(full_source, full_destination), t.movedim(source, destination))
+
+
 @mark.parametrize(
     ["sst_args", "dim"],
     [