Revived triplets tetrahedron in python

phono3py/phonon/grid.py

@@ -730,16 +730,18 @@ def can_use_std_lattice(conv_lat, tmat, std_lattice, rotations, symprec=1e-5):
     return False
 
 
-def get_grid_point_from_address_py(address, mesh):
-    """Python version of get_grid_point_from_address."""
-    # X runs first in XYZ
-    # (*In spglib, Z first is possible with MACRO setting.)
-    m = mesh
-    return (
-        address[0] % m[0]
-        + (address[1] % m[1]) * m[0]
-        + (address[2] % m[2]) * m[0] * m[1]
-    )
+def get_grid_point_from_address_py(addresses, D_diag):
+    """Return GR-grid point index from addresses.
+
+    Python version of get_grid_point_from_address.
+    X runs first in XYZ
+    In grid.c, Z first is possible with MACRO setting.
+
+    addresses :
+        shape=(..., 3)
+
+    """
+    return np.dot(addresses % D_diag, [1, D_diag[0], D_diag[0] * D_diag[1]])
 
 
 def get_grid_point_from_address(address, D_diag):

phono3py/phonon3/triplets.py

@@ -37,7 +37,7 @@ import numpy as np
 from phonopy.structure.tetrahedron_method import TetrahedronMethod
 
 from phono3py.phonon.func import gaussian
-from phono3py.phonon.grid import BZGrid
+from phono3py.phonon.grid import BZGrid, get_grid_point_from_address_py
 
 
 def get_triplets_at_q(
@@ -218,28 +218,6 @@ def get_triplets_integration_weights(
     return g, g_zero
 
 
-def get_tetrahedra_vertices(relative_address, mesh, triplets_at_q, bz_grid):
-    """Return vertices of tetrahedra used for tetrahedron method.
-
-    Equivalent function is implemented in C and this python version exists
-    only for the test and assumes q1+q2+q3=G.
-
-    """
-    bzmesh = mesh * 2
-    grid_order = [1, mesh[0], mesh[0] * mesh[1]]
-    bz_grid_order = [1, bzmesh[0], bzmesh[0] * bzmesh[1]]
-    num_triplets = len(triplets_at_q)
-    vertices = np.zeros((num_triplets, 2, 24, 4), dtype="int_")
-    for i, tp in enumerate(triplets_at_q):
-        for j, adrs_shift in enumerate((relative_address, -relative_address)):
-            adrs = bz_grid.addresses[tp[j + 1]] + adrs_shift
-            bz_gp = np.dot(adrs % bzmesh, bz_grid_order)
-            gp = np.dot(adrs % mesh, grid_order)
-            vgp = bz_grid.gp_map[bz_gp]
-            vertices[i, j] = vgp + (vgp == -1) * (gp + 1)
-    return vertices
-
-
 def _get_triplets_reciprocal_mesh_at_q(
     fixed_grid_number, D_diag, rec_rotations, is_time_reversal=True, swappable=True
 ):
@@ -382,15 +360,19 @@ def _set_triplets_integration_weights_c(g, g_zero, pp, frequency_points):
 
 
 def _set_triplets_integration_weights_py(g, pp, frequency_points):
+    """Python version of _set_triplets_integration_weights_c.
+
+    Tetrahedron method engine is that implemented in phonopy written mainly in C.
+
+    """
     thm = TetrahedronMethod(pp.bz_grid.microzone_lattice)
     triplets_at_q = pp.get_triplets_at_q()[0]
-    tetrahedra_vertices = get_tetrahedra_vertices(
+    tetrahedra_vertices = _get_tetrahedra_vertices(
         np.array(np.dot(thm.get_tetrahedra(), pp.bz_grid.P.T), dtype="int_", order="C"),
-        pp.bz_grid.D_diag,
         triplets_at_q,
         pp.bz_grid,
     )
-    pp.run_phonon_solver(np.array(np.unique(tetrahedra_vertices), dtype="int_"))
+    pp.run_phonon_solver()
     frequencies = pp.get_phonons()[0]
     num_band = frequencies.shape[1]
     for i, vertices in enumerate(tetrahedra_vertices):
@@ -410,3 +392,20 @@ def _set_triplets_integration_weights_py(g, pp, frequency_points):
             g[1, i, :, j, k] = g1 - g2
             if len(g) == 3:
                 g[2, i, :, j, k] = g0 + g1 + g2
+
+
+def _get_tetrahedra_vertices(relative_address, triplets_at_q, bz_grid: BZGrid):
+    """Return vertices of tetrahedra used for tetrahedron method.
+
+    Equivalent function is implemented in C and this python version exists
+    only for the test and assumes q1+q2+q3=G.
+
+    """
+    num_triplets = len(triplets_at_q)
+    vertices = np.zeros((num_triplets, 2, 24, 4), dtype="int_")
+    for i, tp in enumerate(triplets_at_q):
+        for j, adrs_shift in enumerate((relative_address, -relative_address)):
+            adrs = bz_grid.addresses[tp[j + 1]] + adrs_shift
+            gps = get_grid_point_from_address_py(adrs, bz_grid.D_diag)
+            vertices[i, j] = bz_grid.grg2bzg[gps]
+    return vertices

test/phonon3/test_joint_dos.py

@@ -386,7 +386,7 @@ def test_jdos_nac_NaCl_300K_C(nacl_pbe: Phono3py, gp: int, store_dense_gp_map: b
 
 @pytest.mark.parametrize("gp,store_dense_gp_map", [(103, False), (105, True)])
 def test_jdos_nac_NaCl_300K_Py(nacl_pbe: Phono3py, gp: int, store_dense_gp_map: bool):
-    """Test running JDOS of NaCl in Py mode."""
+    """Test running JDOS of NaCl in Py (JDOS) mode."""
     jdos = _get_jdos(
         nacl_pbe,
         [9, 9, 9],
@@ -404,6 +404,26 @@ def test_jdos_nac_NaCl_300K_Py(nacl_pbe: Phono3py, gp: int, store_dense_gp_map:
     )
 
 
+@pytest.mark.parametrize("gp,store_dense_gp_map", [(103, False), (105, True)])
+def test_jdos_nac_NaCl_300K_PyPy(nacl_pbe: Phono3py, gp: int, store_dense_gp_map: bool):
+    """Test running JDOS of NaCl in Py (JDOS) and Py (tetrahedron) mode."""
+    jdos = _get_jdos(
+        nacl_pbe,
+        [9, 9, 9],
+        nac_params=nacl_pbe.nac_params,
+        store_dense_gp_map=store_dense_gp_map,
+    )
+    jdos.set_grid_point(gp)
+    jdos.frequency_points = nacl_freq_points_at_300K
+    jdos.temperature = 300
+    jdos.run_phonon_solver()
+    jdos.run_integration_weights(lang="Py")
+    jdos.run_jdos(lang="Py")
+    np.testing.assert_allclose(
+        nacl_jdos_12_at_300K[2:], jdos.joint_dos.ravel()[2:], rtol=1e-2, atol=1e-5
+    )
+
+
 def _get_jdos(ph3: Phono3py, mesh, nac_params=None, store_dense_gp_map=False):
     bz_grid = BZGrid(
         mesh,