Refactor JDOS

phono3py/api_jointdos.py

@@ -34,6 +34,7 @@
 # POSSIBILITY OF SUCH DAMAGE.
 
 import numpy as np
+from phonopy.structure.cells import Primitive, Supercell
 from phonopy.structure.symmetry import Symmetry
 from phonopy.units import VaspToTHz
 
@@ -51,8 +52,8 @@ class Phono3pyJointDos:
 
     def __init__(
         self,
-        supercell,
-        primitive,
+        supercell: Supercell,
+        primitive: Primitive,
         fc2,
         mesh=None,
         nac_params=None,
@@ -65,58 +66,59 @@ class Phono3pyJointDos:
         frequency_factor_to_THz=VaspToTHz,
         frequency_scale_factor=None,
         use_grg=False,
+        SNF_coordinates="reciprocal",
         is_mesh_symmetry=True,
+        is_symmetry=True,
         store_dense_gp_map=False,
         symprec=1e-5,
         output_filename=None,
         log_level=0,
     ):
         """Init method."""
+        self._primitive = primitive
+        self._supercell = supercell
+        self._fc2 = fc2
+        self._temperatures = temperatures
+        self._nac_params = nac_params
+        self._nac_q_direction = nac_q_direction
         if sigmas is None:
             self._sigmas = [None]
         else:
             self._sigmas = sigmas
-        self._temperatures = temperatures
-        self._nac_q_direction = nac_q_direction
+        self._cutoff_frequency = cutoff_frequency
+        self._frequency_factor_to_THz = frequency_factor_to_THz
+        self._frequency_scale_factor = frequency_scale_factor
         self._is_mesh_symmetry = is_mesh_symmetry
+        self._is_symmetry = is_symmetry
+        self._store_dense_gp_map = store_dense_gp_map
+        self._use_grg = use_grg
+        self._SNF_coordinates = SNF_coordinates
+        self._symprec = symprec
         self._filename = output_filename
         self._log_level = log_level
 
-        symmetry = Symmetry(primitive, symprec)
-        self._bz_grid = BZGrid(
-            mesh,
-            lattice=primitive.cell,
-            symmetry_dataset=symmetry.dataset,
-            use_grg=use_grg,
-            store_dense_gp_map=store_dense_gp_map,
-        )
-
-        self._jdos = JointDos(
-            primitive,
-            supercell,
-            self._bz_grid,
-            fc2,
-            nac_params=nac_params,
-            cutoff_frequency=cutoff_frequency,
-            frequency_factor_to_THz=frequency_factor_to_THz,
-            frequency_scale_factor=frequency_scale_factor,
-            is_mesh_symmetry=self._is_mesh_symmetry,
-            store_dense_gp_map=store_dense_gp_map,
-            symprec=symprec,
-            filename=output_filename,
-            log_level=self._log_level,
-        )
-
         self._joint_dos = None
         self._num_frequency_points_in_batch = None
         self._frequency_step = frequency_step
         self._num_frequency_points = num_frequency_points
 
+        self._primitive_symmetry = Symmetry(
+            self._primitive, self._symprec, self._is_symmetry
+        )
+
+        if mesh is not None:
+            self.initialize(mesh)
+
     @property
     def grid(self):
         """Return BZGrid class instance."""
         return self._bz_grid
 
+    @property
+    def nac_params(self):
+        """Setter and getter of parameters for non-analytical term correction."""
+        return self._nac_params
+
     @property
     def num_frequency_points_in_batch(self):
         """Getter and setter of num_frequency_points_in_batch.
@@ -132,6 +134,63 @@ class Phono3pyJointDos:
     def num_frequency_points_in_batch(self, nelems_in_batch):
         self._num_frequency_points_in_batch = nelems_in_batch
 
+    @property
+    def mesh_numbers(self):
+        """Setter and getter of sampling mesh numbers in reciprocal space."""
+        if self._bz_grid is None:
+            return None
+        else:
+            return self._bz_grid.D_diag
+
+    @mesh_numbers.setter
+    def mesh_numbers(self, mesh_numbers):
+        self._bz_grid = BZGrid(
+            mesh_numbers,
+            lattice=self._primitive.cell,
+            symmetry_dataset=self._primitive_symmetry.dataset,
+            is_time_reversal=self._is_symmetry,
+            use_grg=self._use_grg,
+            force_SNF=False,
+            SNF_coordinates=self._SNF_coordinates,
+            store_dense_gp_map=self._store_dense_gp_map,
+        )
+
+    def initialize(self, mesh_numbers):
+        """Initialize JointDos."""
+        self._jdos = JointDos(
+            self._primitive,
+            self._supercell,
+            self._fc2,
+            nac_params=self._nac_params,
+            cutoff_frequency=self._cutoff_frequency,
+            frequency_factor_to_THz=self._frequency_factor_to_THz,
+            frequency_scale_factor=self._frequency_scale_factor,
+            is_mesh_symmetry=self._is_mesh_symmetry,
+            store_dense_gp_map=self._store_dense_gp_map,
+            symprec=self._symprec,
+            filename=self._filename,
+            log_level=self._log_level,
+        )
+        if self._log_level:
+            print("Generating grid system ... ", end="", flush=True)
+
+        self.mesh_numbers = mesh_numbers
+        self._jdos.bz_grid = self._bz_grid
+
+        if self._log_level:
+            if self._bz_grid.grid_matrix is None:
+                print("[ %d %d %d ]" % tuple(self._bz_grid.D_diag))
+            else:
+                print("")
+                print(
+                    "Generalized regular grid: [ %d %d %d ]"
+                    % tuple(self._bz_grid.D_diag)
+                )
+                print("Grid generation matrix:")
+                print("  [ %d %d %d ]" % tuple(self._bz_grid.grid_matrix[0]))
+                print("  [ %d %d %d ]" % tuple(self._bz_grid.grid_matrix[1]))
+                print("  [ %d %d %d ]" % tuple(self._bz_grid.grid_matrix[2]))
+
     def run(self, grid_points, write_jdos=False):
         """Calculate joint-density-of-states."""
         if self._log_level:
@@ -139,11 +198,9 @@ class Phono3pyJointDos:
                 "--------------------------------- Joint DOS "
                 "---------------------------------"
             )
-            print("Sampling mesh: [ %d %d %d ]" % tuple(self._bz_grid.D_diag))
+            print("Running harmonic phonon calculations...", flush=True)
 
-        self._jdos.run_phonon_solver(
-            np.arange(len(self._bz_grid.addresses), dtype="int_")
-        )
+        self._jdos.run_phonon_solver()
         frequencies, _, _ = self._jdos.get_phonons()
         max_phonon_freq = np.max(frequencies)
         self._frequency_points = get_frequency_points(
@@ -211,7 +268,7 @@ class Phono3pyJointDos:
                     self._jdos.temperature = temperature
 
                     for ib, freq_indices in enumerate(batches):
-                        print(f"{ib + 1}/{len(batches)}: {freq_indices}")
+                        print(f"{ib + 1}/{len(batches)}: {freq_indices}", flush=True)
                         self._jdos.frequency_points = self._frequency_points[
                             freq_indices
                         ]

phono3py/phonon3/joint_dos.py

@@ -38,7 +38,7 @@ import warnings
 
 import numpy as np
 from phonopy.harmonic.dynamical_matrix import DynamicalMatrix, get_dynamical_matrix
-from phonopy.structure.cells import Primitive
+from phonopy.structure.cells import Primitive, Supercell
 from phonopy.units import VaspToTHz
 
 from phono3py.phonon3.triplets import (
@@ -58,7 +58,6 @@ class JointDos:
         self,
         primitive,
         supercell,
-        bz_grid: BZGrid,
         fc2,
         nac_params=None,
         nac_q_direction=None,
@@ -78,7 +77,6 @@ class JointDos:
         self._grid_point = None
         self._primitive = primitive
         self._supercell = supercell
-        self._bz_grid = bz_grid
         self._fc2 = fc2
         self._nac_params = nac_params
         self.nac_q_direction = nac_q_direction
@@ -116,6 +114,7 @@ class JointDos:
         self._g_zero = None
         self._ones_pp_strength = None
         self._temperature = None
+        self._bz_grid = None
 
     @property
     def dynamical_matrix(self) -> DynamicalMatrix:
@@ -160,6 +159,11 @@ class JointDos:
         warnings.warn("Use attribute, primitive", DeprecationWarning)
         return self.primitive
 
+    @property
+    def supercell(self) -> Supercell:
+        """Return supercell."""
+        return self._supercell
+
     @property
     def mesh_numbers(self):
         """Return mesh numbers by three integer values."""
@@ -209,9 +213,13 @@ class JointDos:
 
     @property
     def bz_grid(self) -> BZGrid:
-        """Return BZGrid."""
+        """Setter and getter of BZGrid."""
         return self._bz_grid
 
+    @bz_grid.setter
+    def bz_grid(self, bz_grid: BZGrid):
+        self._bz_grid = bz_grid
+
     @property
     def temperature(self):
         """Setter and getter of temperature."""

test/phonon3/test_joint_dos.py

@@ -411,14 +411,13 @@ def _get_jdos(ph3: Phono3py, mesh, nac_params=None, store_dense_gp_map=False):
         symmetry_dataset=ph3.primitive_symmetry.dataset,
         store_dense_gp_map=store_dense_gp_map,
     )
-    ph3.mesh_numbers = mesh
     jdos = JointDos(
         ph3.primitive,
         ph3.supercell,
-        bz_grid,
         ph3.fc2,
         nac_params=nac_params,
         store_dense_gp_map=store_dense_gp_map,
         cutoff_frequency=1e-4,
     )
+    jdos.bz_grid = bz_grid
     return jdos