fix: normalise fold direction by its length and keep dgz shape when i…

LoopStructural/modelling/features/fold/_fold.py

@@ -104,22 +104,11 @@ def get_deformed_orientation(self, points):
         dgz[mask, :] /= np.linalg.norm(dgz[mask, :], axis=1)[:, None]
         R2 = self.rot_mat(fold_axis, self.fold_limb_rotation(gx))
         fold_direction = np.einsum("ijk,ki->kj", R2, dgx)
-        fold_direction /= np.sum(fold_direction, axis=1)[:, None]
-        # calculate dot product between fold_direction and axis
-        # if its less than 0 then inverse dgz
-        d = np.einsum("ij,ik->i", fold_direction, fold_axis)
+        fold_direction /= np.linalg.norm(fold_direction, axis=1)[:, None]
 
         if self.invert_norm:
-            new_dgz = -dgz[mask][d[mask] < 0]
-            return fold_direction, fold_axis, new_dgz
-
-        elif not self.invert_norm:
-            return fold_direction, fold_axis, dgz
-        else:
-            logger.warning("invert fold frame param not valid. Defaulting to false.")
-            return fold_direction, fold_axis, dgz
-
-        # return fold_direction, fold_axis, dgz
+            dgz = -dgz
+        return fold_direction, fold_axis, dgz
 
     # def get_regularisation_direction(self, points):
     #     self.foldframe.features[2].evaluate_gradient(points)

tests/unit/modelling/test_fold_event.py

@@ -9,3 +9,53 @@ def test_constant_fold_axis():
 
 def test_rotation_fold_axis():
     pass
+
+
+from types import SimpleNamespace
+
+import numpy as np
+
+from LoopStructural.modelling.features.fold import FoldEvent
+
+
+class _PlaneFeature:
+    """Stub fold frame coordinate with a constant gradient"""
+
+    def __init__(self, gradient):
+        self.gradient = np.array(gradient, dtype=float)
+
+    def evaluate_value(self, points):
+        return points @ self.gradient
+
+    def evaluate_gradient(self, points):
+        return np.tile(self.gradient, (points.shape[0], 1))
+
+
+def _fold_event(invert_norm=False):
+    frame = SimpleNamespace(
+        features=[_PlaneFeature([2, 0, 0]), _PlaneFeature([0, 3, 0]), _PlaneFeature([0, 0, 1])]
+    )
+    return FoldEvent(
+        frame,
+        fold_limb_rotation=lambda gx: np.full_like(gx, 30.0),
+        fold_axis=np.array([0, 1.0, 0]),
+        invert_norm=invert_norm,
+    )
+
+
+def test_deformed_orientation_is_unit_length():
+    points = np.random.default_rng(0).random((20, 3))
+    fold_direction, fold_axis, dgz = _fold_event().get_deformed_orientation(points)
+    assert fold_direction.shape == points.shape
+    assert dgz.shape == points.shape
+    assert np.allclose(np.linalg.norm(fold_direction, axis=1), 1.0)
+    assert np.allclose(np.linalg.norm(dgz, axis=1), 1.0)
+    assert np.allclose(np.einsum("ij,ij->i", fold_direction, fold_axis), 0.0)
+
+
+def test_invert_norm_flips_dgz():
+    points = np.random.default_rng(0).random((20, 3))
+    _, _, dgz = _fold_event().get_deformed_orientation(points)
+    _, _, inverted = _fold_event(invert_norm=True).get_deformed_orientation(points)
+    assert inverted.shape == points.shape
+    assert np.allclose(inverted, -dgz)