Merge pull request #1 from hjkgrp/fix_equivariance

Fix ordering of e3nn coefficients and add test cases

.github/workflows/CI.yaml

@@ -26,7 +26,7 @@ jobs:
     runs-on: ${{ matrix.os }}
     strategy:
       matrix:
-        os: [macOS-latest, ubuntu-latest]
+        os: [ubuntu-latest]
         python-version: [3.8]
 
     steps:
@@ -41,16 +41,14 @@ jobs:
 
 
     # More info on options: https://github.com/conda-incubator/setup-miniconda
-    - uses: conda-incubator/setup-miniconda@v2
+    - uses: conda-incubator/setup-miniconda@v3
       with:
         python-version: ${{ matrix.python-version }}
         environment-file: devtools/conda-envs/test_env.yaml
 
         channels: conda-forge,defaults
 
         activate-environment: test
-        auto-update-conda: false
-        auto-activate-base: false
         show-channel-urls: true
 
     - name: Install package

devtools/conda-envs/test_env.yaml

@@ -5,19 +5,23 @@ channels:
   - conda-forge
 dependencies:
     # Base depends
-  - python=3.8
-  - psi4
+  - python=3.8.13  
   - pip
-  - pytorch=1.10.0
-  - pandas
-  - scikit-learn
+  - pytorch=1.10.2
+  - pandas=1.4.3
+  - scikit-learn=1.1.1
+  - psi4=1.6.1
+  - pydantic=1.9.0
+  - typing-extensions=4.3.0
+  - numpy=1.23.1
 
     # Testing
   - pytest
   - pytest-cov
   - codecov
 
     # Pip-only installs
-  #- pip:
+  - pip:
+    - e3nn
   #  - codecov
 

dfa_recommender/df_class.py

@@ -6,6 +6,8 @@
 import numpy as np
 import psi4
 import os
+import torch
+from e3nn import o3
 from typing import Tuple
 
 
@@ -271,7 +273,7 @@ def pad_df_coeffs(self) -> None:
             self.max_tmplate.append((self.max_shell == ii).sum()*(2*ii + 1) + self.max_tmplate[-1])
             evenodd = "e" if ii%2==0 else "o"
             self.irreps += [str((self.max_shell == ii).sum()) + "x" + str(ii) + evenodd]
-        self.irreps = "+".join(self.irreps)
+        self.irreps = o3.Irreps("+".join(self.irreps))
 
         tmplates = []
         for jj in range(len(self.shellmap)):
@@ -301,11 +303,19 @@ def convert_CP2e3nn(self) -> None:
         e3nn m: [-2, -1, 0, 1, 2]
         '''
 
-        psi4_2_e3nn = [[0],[2,0,1],[4,2,0,1,3],[6,4,2,0,1,3,5],[8,6,4,2,0,1,3,5,7]]
+        psi4_2_e3nn = [
+            [0],
+            [2, 0, 1],
+            [4, 2, 0, 1, 3],
+            [6, 4, 2, 0, 1, 3, 5],
+            [8, 6, 4, 2, 0, 1, 3, 5, 7],
+            [10, 8, 6, 4, 2, 0, 1, 3, 5, 7, 9],
+            [12, 10, 8, 6, 4, 2, 0, 1, 3, 5, 7, 9, 11],
+        ]
         self.C_P_pad_e3nn = np.zeros(shape=self.C_P_pad.shape)
         for ii in range(self.C_P_pad.shape[0]):
-            for jj, ele in enumerate(self.irreps.split("+")):
-                num, l = int(ele.split("x")[0]),  int(ele.split("x")[1][0])
+            for jj, irrep in enumerate(self.irreps):
+                num, l = irrep.mul, irrep.ir.l
                 coeffs = self.C_P_pad[ii][self.max_tmplate[jj]: self.max_tmplate[jj+1]]
                 coeffs = np.array_split(coeffs, num)
                 coeffs_trans = []
@@ -316,4 +326,20 @@ def convert_CP2e3nn(self) -> None:
         mat = np.where(np.abs(self.C_P_pad_e3nn)< 1e-8)
         self.C_P_pad_e3nn[mat[0], mat[1]] = 0
 
+        # As explained in 
+        # https://docs.e3nn.org/en/latest/guide/change_of_basis.html
+        # the ordering of the basis functions was changed in release 0.2.2
+        # which necessitates a change of basis for the density fitting coefficients
+        change_of_basis = torch.tensor(
+            [
+                # this specifies the change of basis yzx -> xyz
+                [0.0, 0.0, 1.0],
+                [1.0, 0.0, 0.0],
+                [0.0, 1.0, 0.0],
+            ]
+        )
+
+        D_change = self.irreps.D_from_matrix(change_of_basis).numpy()
+        self.C_P_pad_e3nn = self.C_P_pad_e3nn @ D_change.T 
+
 

dfa_recommender/tests/inputs/rotational_equivariance/rotation0.xyz

@@ -0,0 +1,24 @@
+22
+
+Co         0.00000       -0.00000        0.00000
+S         -2.56544       -0.52966       -0.06121
+H         -2.99336        0.07099        1.06926
+H         -3.15262        0.36524       -0.88504
+S          0.50213       -2.57764       -0.08157
+H         -0.68391       -3.19224       -0.28366
+H          0.56989       -2.91299        1.22381
+S          2.58089        0.51815        0.04669
+H          2.99972       -0.01479        1.21381
+H          3.16370       -0.43448       -0.71325
+S         -0.49180        2.58341        0.10262
+H          0.70632        3.20682        0.11694
+H         -0.81067        2.99861       -1.14223
+S         -0.11866        0.13270       -2.61794
+H          1.06624        0.62828       -3.03436
+H          0.10393       -1.10167       -3.11929
+C         -0.05143        0.08009        3.27876
+N          0.00000        0.00000        2.12207
+C         -0.11499        0.18043        4.72888
+H         -0.20026       -0.81902        5.16946
+H         -0.98687        0.77657        5.02063
+H          0.79141        0.66471        5.10871

dfa_recommender/tests/inputs/rotational_equivariance/rotation1.xyz

@@ -0,0 +1,24 @@
+22
+Properties=species:S:1:pos:R:3 pbc="F F F"
+Co       0.00000000       0.00000000       0.00000000
+S       -1.93095004      -1.70502773       0.47966777
+H       -1.89000425      -1.78928128       1.82617553
+H       -3.09820596      -1.02601428       0.45180024
+S        1.16347780      -1.79497023      -1.52558757
+H        0.29769870      -2.82678264      -1.63072606
+H        1.98253713      -2.43603419      -0.66579294
+S        1.93961308       1.70809119      -0.50197427
+H        3.03262369       1.12856409       0.03755222
+H        2.32008766       1.47235577      -1.77647343
+S       -1.14630836       1.79830081       1.54223755
+H       -0.36853353       2.90163858       1.49697495
+H       -2.16162732       2.35779006       0.84954726
+S       -1.46107092       0.85993831      -2.00276739
+H       -0.87213100       1.98751397      -2.45522297
+H       -1.15215036       0.13492567      -3.09980531
+C        1.59239202      -0.98000523       2.69503321
+N        1.07351142      -0.65915148       1.70771006
+C        2.24366133      -1.38184376       3.93250243
+H        2.70836804      -2.36612708       3.80793904
+H        1.50444545      -1.43379601       4.73968056
+H        3.01479818      -0.65157800       4.20160982

dfa_recommender/tests/inputs/rotational_equivariance/rotation2.xyz

@@ -0,0 +1,24 @@
+22
+Properties=species:S:1:pos:R:3 pbc="F F F"
+Co       0.00000000       0.00000000       0.00000000
+S       -2.51336658      -0.69650307       0.25227110
+H       -2.66048244      -0.69971078       1.59398828
+H       -3.26883997       0.40750251       0.06584720
+S        0.44699214      -2.21585857      -1.33909858
+H       -0.75871541      -2.79139727      -1.53981957
+H        0.81459905      -3.08309633      -0.37274884
+S        2.52490197       0.69412927      -0.27332609
+H        3.19990688      -0.25432562       0.40961930
+H        2.90605480       0.23521366      -1.48526635
+S       -0.43198595       2.21275374       1.35776080
+H        0.74172735       2.87997776       1.39707768
+H       -1.02894768       3.10009034       0.53314508
+S       -0.72513346       1.25906000      -2.18500542
+H        0.33410985       2.00046713      -2.57399417
+H       -0.63632889       0.39965436      -3.22332286
+C        0.71578290      -1.37658059       2.88998396
+N        0.49518632      -0.93396322       1.84002291
+C        0.99323146      -1.93145195       4.20607517
+H        1.00451217      -3.02580786       4.15567764
+H        0.21867441      -1.61292723       4.91267200
+H        1.96738206      -1.57778771       4.56147792

dfa_recommender/tests/test_df.py

@@ -2,6 +2,8 @@
 test suites for DensityFitting class
 '''
 
+import torch
+from e3nn import o3
 from dfa_recommender.df_class import DensityFitting
 from dfa_recommender.df_utils import get_spectra, get_subtracted_spectra
 import numpy as np
@@ -37,6 +39,53 @@ def test_df_water():
     assert np.max([np.max(np.abs(np.array(ref_ps[ii]) - np.array(ps[ii]))) for ii in range(ps.shape[0])]) < 1e-6
 
 
+@pytest.mark.parametrize("name, axis, angle",
+    [
+        ("rotation1", [1.0, 1.0, 1.0], np.pi / 4),
+        ("rotation2", [1.0, 0.5, 0.1], np.pi / 6),
+    ]
+)
+def test_df_rotation_equivariance(name, axis, angle):
+
+    charge = 2
+    spin = 2
+    basis = "def2-tzvp"
+
+    basepath = resource_filename(Requirement.parse("dfa_recommender"), "/dfa_recommender/tests/")
+    densfit = DensityFitting(
+        wfnpath=basepath + "inputs/rotational_equivariance/rotation0_LS_HF_wfn.npy",
+        xyzfile=basepath + "inputs/rotational_equivariance/rotation0.xyz",
+        charge=charge,
+        spin=spin,
+        basis=basis,
+    )
+    densfit.get_df_coeffs(densfit.wfn.Da())
+    densfit.pad_df_coeffs()
+    densfit.convert_CP2e3nn()
+
+    densfit_ref = DensityFitting(
+        wfnpath=basepath + f"inputs/rotational_equivariance/{name}_LS_HF_wfn.npy",
+        xyzfile=basepath + f"inputs/rotational_equivariance/{name}.xyz",
+        charge=charge,
+        spin=spin,
+        basis=basis,
+    )
+    densfit_ref.get_df_coeffs(densfit_ref.wfn.Da())
+    densfit_ref.pad_df_coeffs()
+    densfit_ref.convert_CP2e3nn()
+
+    rot = o3.axis_angle_to_matrix(torch.tensor(axis), torch.tensor(angle))
+    # Construct wigner D matrix from rotation matrix
+    D = densfit.irreps.D_from_matrix(rot).numpy()
+
+    np.testing.assert_allclose(
+        densfit.C_P_pad_e3nn @ D.T,
+        densfit_ref.C_P_pad_e3nn,
+        atol=1e-6,
+        rtol=1e-5,
+    )
+
+
 def test_df_vert_spin_splitting():
 
     basepath = resource_filename(Requirement.parse("dfa_recommender"), "/dfa_recommender/tests/")