From af4e08a349789579ad67be02c804fe96da0799f7 Mon Sep 17 00:00:00 2001
From: Chris Markiewicz <effigies@gmail.com>
Date: Mon, 18 Sep 2023 15:37:20 -0400
Subject: [PATCH] TEST: Test Pointset and GridIndices classes

---
 nibabel/tests/test_pointset.py | 122 +++++++++++++++++++++++++++++++++
 1 file changed, 122 insertions(+)

diff --git a/nibabel/tests/test_pointset.py b/nibabel/tests/test_pointset.py
index efea8bbd7..88001b401 100644
--- a/nibabel/tests/test_pointset.py
+++ b/nibabel/tests/test_pointset.py
@@ -2,8 +2,10 @@
 from unittest import skipUnless
 
 import numpy as np
+import pytest
 
 from nibabel import pointset as ps
+from nibabel.affines import apply_affine
 from nibabel.arrayproxy import ArrayProxy
 from nibabel.onetime import auto_attr
 from nibabel.optpkg import optional_package
@@ -14,6 +16,126 @@
 FS_DATA = Path(get_nibabel_data()) / 'nitest-freesurfer'
 
 
+class TestPointsets:
+    rng = np.random.default_rng()
+
+    @pytest.mark.parametrize('shape', [(5, 2), (5, 3), (5, 4)])
+    @pytest.mark.parametrize('homogeneous', [True, False])
+    def test_init(self, shape, homogeneous):
+        coords = self.rng.random(shape)
+
+        if homogeneous:
+            coords = np.column_stack([coords, np.ones(shape[0])])
+
+        expected_shape = (shape[0], shape[1] + homogeneous)
+
+        points = ps.Pointset(coords, homogeneous=homogeneous)
+        assert points.shape == expected_shape
+        assert np.allclose(points.affine, np.eye(shape[1] + 1))
+        assert points.homogeneous is homogeneous
+        assert points.ndim == 2
+        assert points.n_coords == shape[0]
+        assert points.dim == shape[1]
+
+        points = ps.Pointset(coords, affine=np.diag([2] * shape[1] + [1]), homogeneous=homogeneous)
+        assert points.shape == expected_shape
+        assert np.allclose(points.affine, np.diag([2] * shape[1] + [1]))
+        assert points.homogeneous is homogeneous
+        assert points.ndim == 2
+        assert points.n_coords == shape[0]
+        assert points.dim == shape[1]
+
+        # Badly shaped affine
+        with pytest.raises(ValueError):
+            ps.Pointset(coords, affine=[0, 1])
+
+        # Badly valued affine
+        with pytest.raises(ValueError):
+            ps.Pointset(coords, affine=np.ones((shape[1] + 1, shape[1] + 1)))
+
+    @pytest.mark.parametrize('shape', [(5, 2), (5, 3), (5, 4)])
+    @pytest.mark.parametrize('homogeneous', [True, False])
+    def test_affines(self, shape, homogeneous):
+        orig_coords = coords = self.rng.random(shape)
+
+        if homogeneous:
+            coords = np.column_stack([coords, np.ones(shape[0])])
+
+        points = ps.Pointset(coords, homogeneous=homogeneous)
+        assert np.allclose(points.get_coords(), orig_coords)
+
+        # Apply affines
+        scaler = np.diag([2] * shape[1] + [1])
+        scaled = scaler @ points
+        assert np.array_equal(scaled.coordinates, points.coordinates)
+        assert np.array_equal(scaled.affine, scaler)
+        assert np.allclose(scaled.get_coords(), 2 * orig_coords)
+
+        flipper = np.eye(shape[1] + 1)
+        # [[1, 0, 0], [0, 1, 0], [0, 0, 1]] becomes [[0, 1, 0], [1, 0, 0], [0, 0, 1]]
+        flipper[:-1] = flipper[-2::-1]
+        flipped = flipper @ points
+        assert np.array_equal(flipped.coordinates, points.coordinates)
+        assert np.array_equal(flipped.affine, flipper)
+        assert np.allclose(flipped.get_coords(), orig_coords[:, ::-1])
+
+        # Concatenate affines, with any associativity
+        for doubledup in [(scaler @ flipper) @ points, scaler @ (flipper @ points)]:
+            assert np.array_equal(doubledup.coordinates, points.coordinates)
+            assert np.allclose(doubledup.affine, scaler @ flipper)
+            assert np.allclose(doubledup.get_coords(), 2 * orig_coords[:, ::-1])
+
+    def test_homogeneous_coordinates(self):
+        ccoords = self.rng.random((5, 3))
+        hcoords = np.column_stack([ccoords, np.ones(5)])
+
+        cartesian = ps.Pointset(ccoords)
+        homogeneous = ps.Pointset(hcoords, homogeneous=True)
+
+        for points in (cartesian, homogeneous):
+            assert np.array_equal(points.get_coords(), ccoords)
+            assert np.array_equal(points.get_coords(as_homogeneous=True), hcoords)
+
+        affine = np.diag([2, 3, 4, 1])
+        cart2 = affine @ cartesian
+        homo2 = affine @ homogeneous
+
+        exp_c = apply_affine(affine, ccoords)
+        exp_h = (affine @ hcoords.T).T
+        for points in (cart2, homo2):
+            assert np.array_equal(points.get_coords(), exp_c)
+            assert np.array_equal(points.get_coords(as_homogeneous=True), exp_h)
+
+
+def test_GridIndices():
+    # 2D case
+    shape = (2, 3)
+    gi = ps.GridIndices(shape)
+
+    assert gi.dtype == np.dtype('u1')
+    assert gi.shape == (6, 2)
+    assert repr(gi) == '<GridIndices(2, 3)>'
+
+    gi_arr = np.asanyarray(gi)
+    assert gi_arr.dtype == np.dtype('u1')
+    assert gi_arr.shape == (6, 2)
+    # Tractable to write out
+    assert np.array_equal(gi_arr, [[0, 0], [0, 1], [0, 2], [1, 0], [1, 1], [1, 2]])
+
+    shape = (2, 3, 4)
+    gi = ps.GridIndices(shape)
+
+    assert gi.dtype == np.dtype('u1')
+    assert gi.shape == (24, 3)
+    assert repr(gi) == '<GridIndices(2, 3, 4)>'
+
+    gi_arr = np.asanyarray(gi)
+    assert gi_arr.dtype == np.dtype('u1')
+    assert gi_arr.shape == (24, 3)
+    # Separate implementation
+    assert np.array_equal(gi_arr, np.mgrid[:2, :3, :4].reshape(3, -1).T)
+
+
 class H5ArrayProxy:
     def __init__(self, file_like, dataset_name):
         self.file_like = file_like