diff --git a/pylinalg/vector.py b/pylinalg/vector.py index 2989cab..0d6a4f1 100644 --- a/pylinalg/vector.py +++ b/pylinalg/vector.py @@ -106,50 +106,30 @@ def vec_transform( """ matrix = np.asarray(matrix) - vectors = np.asarray(vectors) - - # this code has been micro-optimized for the 1D and 2D cases - vectors_ndim = vectors.ndim - if vectors_ndim > 2: - raise ValueError("vectors must be a 1D or 2D array") - - # determine if we are working with a batch of vectors - batch = vectors_ndim != 1 - - # we don't need to work in homogeneous vector space - # if matrix is purely affine and vectors is a single vector - homogeneous = projection or batch - - if homogeneous: - vectors = vec_homogeneous(vectors, w=w) - matmul_matrix = matrix - else: - # if we are not working in homogeneous space, it's - # more efficient to matmul the 3x3 (scale + rotation) - # part of the matrix with the vectors and then add - # the translation part after - matmul_matrix = matrix[:-1, :-1] - - if batch: + vectors = vec_homogeneous(vectors, w=w) + + # yes, the ndim > 2 version can also handle ndim=1 + # and ndim=2, but it's slower + if vectors.ndim == 1: + vectors = matrix @ vectors + if projection: + vectors = vectors[:-1] / vectors[-1] + else: + vectors = vectors[:-1] + elif vectors.ndim == 2: # transposing the vectors array performs better # than transposing the matrix - vectors = (matmul_matrix @ vectors.T).T + vectors = (matrix @ vectors.T).T + if projection: + vectors = vectors[:, :-1] / vectors[:, -1, None] + else: + vectors = vectors[:, :-1] else: - vectors = matmul_matrix @ vectors - if not homogeneous: - # as alluded to before, we add the translation - # part of the matrix after the matmul - # if we are not working in homogeneous space - vectors = vectors + matrix[:-1, -1] - - if projection: - # if we are projecting, we divide by the last - # element of the vectors array - vectors = vectors[..., :-1] / vectors[..., -1, None] - elif homogeneous: - # if we are NOT projecting but we are working in - # homogeneous space, just drop the last element - vectors = vectors[..., :-1] + vectors = matrix @ vectors[..., None] + if projection: + vectors = vectors[..., :-1, 0] / vectors[..., -1, :] + else: + vectors = vectors[..., :-1, 0] if out is None: out = vectors diff --git a/pyproject.toml b/pyproject.toml index 6f5850f..633985c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -2,7 +2,7 @@ [project] name = "pylinalg" -version = "0.6.6" +version = "0.6.7" description = "Linear algebra utilities for Python" readme = "README.md" license = { file = "LICENSE" } diff --git a/tests/test_vectors.py b/tests/test_vectors.py index 1212cad..bb22e8c 100644 --- a/tests/test_vectors.py +++ b/tests/test_vectors.py @@ -120,6 +120,36 @@ def test_vec_transform_projection_flag(): npt.assert_array_equal(result, expected_out) +def test_vec_transform_ndim(): + vectors_2d = np.array( + [ + [1, 0, 0], + [1, 2, 3], + [1, 1, 1], + [1, 1, -1], + [0, 0, 0], + [7, 8, -9], + ], + dtype="f8", + ) + translation = np.array([-1, 2, 2], dtype="f8") + + vectors_3d = vectors_2d.reshape((3, 2, 3)) + vectors_4d = vectors_2d.reshape((6, 1, 1, 3)) + + expected_3d = vectors_3d + translation[None, None, :] + expected_4d = vectors_4d + translation[None, None, None, :] + + matrix = la.mat_from_translation(translation) + + for projection in [True, False]: + result = la.vec_transform(vectors_3d, matrix, projection=projection) + npt.assert_array_equal(result, expected_3d) + + result = la.vec_transform(vectors_4d, matrix, projection=projection) + npt.assert_array_equal(result, expected_4d) + + @given(ct.test_spherical, none()) @example((1, 0, np.pi / 2), (0, 0, 1)) @example((1, np.pi / 2, np.pi / 2), (1, 0, 0))