[ BE ] refactored init edge rot mat (#1368)

* refactored init edge rot mat

* cleanup

* fix dtype issue

* fix test constants

* fix cuda graph version

* clean up

* test if gamma=0 fixes CI gpu test

* NCCL debug

* typo

* add gamma back in for tests

* only run the one test

* fix gpu debug

* wrong test

* try to run full mlip_unit

* test more

* run larger set of tests

* run more test

* debug memory

* warning supressed by pytest, do a print

* try empty cache

* revert yaml

* clear cuda memory between tests

Author: misko

.github/workflows/test.yml

@@ -84,13 +84,13 @@ jobs:
         env:
           HF_TOKEN: ${{ secrets.HF_TOKEN }}
         run: |
-          pytest tests --durations=0 -vv --ignore=tests/demo/ocpapi/tests/integration/ --ignore=tests/applications/ --ignore=tests/perf --cov-report=xml --cov=fairchem --junitxml=junit.xml -o junit_family=legacy -c ./packages/fairchem-core/pyproject.toml
+          NCCL_DEBUG=INFO pytest tests --durations=0 -vv --ignore=tests/demo/ocpapi/tests/integration/ --ignore=tests/applications/ --ignore=tests/perf --cov-report=xml --cov=fairchem --junitxml=junit.xml -o junit_family=legacy -c ./packages/fairchem-core/pyproject.toml
 
       - name: Core GPU tests
         env:
           HF_TOKEN: ${{ secrets.HF_TOKEN }}
         run: |
-          pytest tests/core --durations=0 -vv -m gpu -c ./packages/fairchem-core/pyproject.toml -s --junitxml=junit-gpu.xml -o junit_family=legacy --cov-report=xml:gpu-coverage.xml
+          NCCL_DEBUG=INFO  pytest tests/core --durations=0 -vv -m gpu -c ./packages/fairchem-core/pyproject.toml -s --junitxml=junit-gpu.xml -o junit_family=legacy --cov-report=xml:gpu-coverage.xml
 
       - if: ${{ matrix.python_version == '3.12' }}
         name: codecov-coverage
@@ -156,8 +156,7 @@ jobs:
         env:
           HF_TOKEN: ${{ secrets.HF_TOKEN }}
         run: |
-          pytest tests/core --durations=0 -vv -m gpu -c ./packages/fairchem-core/pyproject.toml
-
+          NCCL_DEBUG=INFO pytest tests/core --durations=0 -vv -m gpu -c ./packages/fairchem-core/pyproject.toml
       - name: Cleanup
         if: always()
         uses: ./.github/actions/multi-trigger-cleanup

src/fairchem/core/models/uma/common/rotation.py

@@ -8,15 +8,11 @@
 from __future__ import annotations
 
 import logging
-import math
 
 import torch
-from e3nn import o3
 
-YTOL = 0.999999
 
-
-def init_edge_rot_mat(edge_distance_vec, rot_clip=False):
+def init_edge_rot_euler_angles(edge_distance_vec):
     edge_vec_0 = edge_distance_vec
     edge_vec_0_distance = torch.sqrt(torch.sum(edge_vec_0**2, dim=1))
 
@@ -25,54 +21,30 @@ def init_edge_rot_mat(edge_distance_vec, rot_clip=False):
     if len(edge_vec_0_distance) > 0 and torch.min(edge_vec_0_distance) < 0.0001:
         logging.error(f"Error edge_vec_0_distance: {torch.min(edge_vec_0_distance)}")
 
-    norm_x = edge_vec_0 / (edge_vec_0_distance.view(-1, 1))
-
-    if rot_clip:
-        yprod = norm_x @ norm_x.new_tensor([0.0, 1.0, 0.0])
-        norm_x[yprod > YTOL] = norm_x.new_tensor([0.0, 1.0, 0.0])
-        norm_x[yprod < -YTOL] = norm_x.new_tensor([0.0, -1.0, 0.0])
-
-    edge_vec_2 = torch.rand_like(edge_vec_0) - 0.5
-    edge_vec_2 = edge_vec_2 / (torch.sqrt(torch.sum(edge_vec_2**2, dim=1)).view(-1, 1))
-    # Create two rotated copys of the random vectors in case the random vector is aligned with norm_x
-    # With two 90 degree rotated vectors, at least one should not be aligned with norm_x
-    edge_vec_2b = edge_vec_2.clone()
-    edge_vec_2b[:, 0] = -edge_vec_2[:, 1]
-    edge_vec_2b[:, 1] = edge_vec_2[:, 0]
-    edge_vec_2c = edge_vec_2.clone()
-    edge_vec_2c[:, 1] = -edge_vec_2[:, 2]
-    edge_vec_2c[:, 2] = edge_vec_2[:, 1]
-    vec_dot_b = torch.abs(torch.sum(edge_vec_2b * norm_x, dim=1)).view(-1, 1)
-    vec_dot_c = torch.abs(torch.sum(edge_vec_2c * norm_x, dim=1)).view(-1, 1)
-
-    vec_dot = torch.abs(torch.sum(edge_vec_2 * norm_x, dim=1)).view(-1, 1)
-    edge_vec_2 = torch.where(torch.gt(vec_dot, vec_dot_b), edge_vec_2b, edge_vec_2)
-    vec_dot = torch.abs(torch.sum(edge_vec_2 * norm_x, dim=1)).view(-1, 1)
-    edge_vec_2 = torch.where(torch.gt(vec_dot, vec_dot_c), edge_vec_2c, edge_vec_2)
-
-    vec_dot = torch.abs(torch.sum(edge_vec_2 * norm_x, dim=1))
-    # Check the vectors aren't aligned
-    if len(vec_dot) > 0:
-        assert torch.max(vec_dot) < 0.99
-
-    norm_z = torch.cross(norm_x, edge_vec_2, dim=1)
-    norm_z = norm_z / (torch.sqrt(torch.sum(norm_z**2, dim=1, keepdim=True)))
-    norm_z = norm_z / (torch.sqrt(torch.sum(norm_z**2, dim=1)).view(-1, 1))
-    norm_y = torch.cross(norm_x, norm_z, dim=1)
-    norm_y = norm_y / (torch.sqrt(torch.sum(norm_y**2, dim=1, keepdim=True)))
-
-    # Construct the 3D rotation matrix
-    norm_x = norm_x.view(-1, 3, 1)
-    norm_y = -norm_y.view(-1, 3, 1)
-    norm_z = norm_z.view(-1, 3, 1)
-
-    edge_rot_mat_inv = torch.cat([norm_z, norm_x, norm_y], dim=2)
-    edge_rot_mat = torch.transpose(edge_rot_mat_inv, 1, 2)
-
-    if rot_clip:
-        return edge_rot_mat
-    else:
-        return edge_rot_mat.detach()
+    # make unit vectors
+    xyz = edge_vec_0 / (edge_vec_0_distance.view(-1, 1))
+
+    # are we standing at the north pole
+    mask = xyz[:, 1].abs().isclose(xyz.new_ones(1))
+
+    # compute alpha and beta
+
+    # latitude (beta)
+    beta = xyz.new_zeros(xyz.shape[0])
+    beta[~mask] = torch.acos(xyz[~mask, 1])
+    beta[mask] = torch.acos(xyz[mask, 1]).detach()
+
+    # longitude (alpha)
+    alpha = torch.zeros_like(beta)
+    alpha[~mask] = torch.atan2(xyz[~mask, 0], xyz[~mask, 2])
+    alpha[mask] = torch.atan2(xyz[mask, 0], xyz[mask, 2]).detach()
+
+    # random gamma (roll)
+    gamma = torch.rand_like(alpha) * 2 * torch.pi
+    # gamma = torch.zeros_like(alpha)
+
+    # intrinsic to extrinsic swap
+    return -gamma, -beta, -alpha
 
 
 # Borrowed from e3nn @ 0.4.0:
@@ -118,58 +90,24 @@ def _z_rot_mat(angle: torch.Tensor, lv: int) -> torch.Tensor:
     return M
 
 
-def rotation_to_wigner(
-    edge_rot_mat: torch.Tensor,
+def eulers_to_wigner(
+    eulers: torch.Tensor,
     start_lmax: int,
     end_lmax: int,
     Jd: list[torch.Tensor],
-    rot_clip: bool = False,
 ) -> torch.Tensor:
     """
     set <rot_clip=True> to handle gradient instability when using gradient-based force/stress prediction.
     """
-    x = edge_rot_mat @ edge_rot_mat.new_tensor([0.0, 1.0, 0.0])
-    alpha, beta = o3.xyz_to_angles(x)
-    R = (
-        o3.angles_to_matrix(alpha, beta, torch.zeros_like(alpha)).transpose(-1, -2)
-        @ edge_rot_mat
-    )
-    gamma = torch.atan2(R[..., 0, 2], R[..., 0, 0])
-
-    if rot_clip:
-        yprod = (x @ x.new_tensor([0, 1, 0])).detach()
-        mask = (yprod > -YTOL) & (yprod < YTOL)
-        alpha_detach = alpha[~mask].clone().detach()
-        gamma_detach = gamma[~mask].clone().detach()
-        beta_detach = beta.clone().detach()
-        beta_detach[yprod > YTOL] = 0.0
-        beta_detach[yprod < -YTOL] = math.pi
-        beta_detach = beta_detach[~mask]
+    alpha, beta, gamma = eulers
 
     size = int((end_lmax + 1) ** 2) - int((start_lmax) ** 2)
-    wigner = torch.zeros(
-        len(alpha), size, size, device=edge_rot_mat.device, dtype=edge_rot_mat.dtype
-    )
+    wigner = torch.zeros(len(alpha), size, size, device=alpha.device, dtype=alpha.dtype)
     start = 0
     for lmax in range(start_lmax, end_lmax + 1):
-        if rot_clip:
-            block = wigner_D(lmax, alpha[mask], beta[mask], gamma[mask], Jd).to(
-                wigner.dtype
-            )
-            block_detach = wigner_D(
-                lmax, alpha_detach, beta_detach, gamma_detach, Jd
-            ).to(wigner.dtype)
-            end = start + block.size()[1]
-            wigner[mask, start:end, start:end] = block
-            wigner[~mask, start:end, start:end] = block_detach
-            start = end
-        else:
-            block = wigner_D(lmax, alpha, beta, gamma, Jd)
-            end = start + block.size()[1]
-            wigner[:, start:end, start:end] = block
-            start = end
-
-    if rot_clip:
-        return wigner
-    else:
-        return wigner.detach()
+        block = wigner_D(lmax, alpha, beta, gamma, Jd)
+        end = start + block.size()[1]
+        wigner[:, start:end, start:end] = block
+        start = end
+
+    return wigner