Merge branch 'stfc:main' into orb

Author: alinelena

docs/source/user_guide/command_line.rst

@@ -314,7 +314,7 @@ Calculate phonons with a 2x2x2 supercell, after geometry optimization (using the
 
 .. code-block:: bash
 
-    janus phonons --struct tests/data/NaCl.cif --supercell 2x2x2 --minimize --arch mace_mp --model-path small
+    janus phonons --struct tests/data/NaCl.cif --supercell 2 2 2 --minimize --arch mace_mp --model-path small
 
 
 This will save the Phonopy parameters, including displacements and force constants, to ``NaCl-phonopy.yml`` and ``NaCl-force_constants.hdf5``,
@@ -324,7 +324,7 @@ Additionally, the ``--bands`` option can be added to calculate the band structur
 
 .. code-block:: bash
 
-    janus phonons --struct tests/data/NaCl.cif --supercell 2x2x2 --minimize --arch mace_mp --model-path small --bands
+    janus phonons --struct tests/data/NaCl.cif --supercell 2 2 2 --minimize --arch mace_mp --model-path small --bands
 
 
 If you need eigenvectors and group velocities written, add the ``--write-full`` option. This will generate a much larger file, but can be used to visualise phonon modes.
@@ -333,7 +333,7 @@ Further calculations, including thermal properties, DOS, and PDOS, can also be c
 
 .. code-block:: bash
 
-    janus phonons --struct tests/data/NaCl.cif --supercell 2x3x4 --dos --pdos --thermal --temp-start 0 --temp-end 300 --temp-step 50
+    janus phonons --struct tests/data/NaCl.cif --supercell 2 3 4 --dos --pdos --thermal --temp-start 0 --temp-end 300 --temp-step 50
 
 
 This will create additional output files: ``NaCl-thermal.dat`` for the thermal properties (heat capacity, entropy, and free energy)

janus_core/calculations/descriptors.py

@@ -4,6 +4,8 @@
 from typing import Any, Optional
 
 from ase import Atoms
+from ase.calculators.calculator import Calculator
+from ase.calculators.mixing import SumCalculator
 import numpy as np
 
 from janus_core.calculations.base import BaseCalculation
@@ -163,12 +165,43 @@ def __init__(
         ):
             raise ValueError("Please attach a calculator to `struct`.")
 
+        if isinstance(self.struct, Atoms):
+            self._check_calculator(self.struct.calc)
+        if isinstance(self.struct, Sequence):
+            for image in self.struct:
+                self._check_calculator(image.calc)
+
         # Set output file
         self.write_kwargs.setdefault("filename", None)
         self.write_kwargs["filename"] = self._build_filename(
             "descriptors.extxyz", filename=self.write_kwargs["filename"]
         ).absolute()
 
+    @staticmethod
+    def _check_calculator(calc: Calculator) -> None:
+        """
+        Ensure calculator has ability to calculate descriptors.
+
+        Parameters
+        ----------
+        calc : Calculator
+            ASE Calculator to calculate descriptors.
+        """
+        # If dispersion added to MLIP calculator, use MLIP calculator for descriptors
+        if isinstance(calc, SumCalculator):
+            if (
+                len(calc.mixer.calcs) == 2
+                and calc.mixer.calcs[1].name == "TorchDFTD3Calculator"
+                and hasattr(calc.mixer.calcs[0], "get_descriptors")
+            ):
+                calc.get_descriptors = calc.mixer.calcs[0].get_descriptors
+
+        if not hasattr(calc, "get_descriptors") or not callable(calc.get_descriptors):
+            raise NotImplementedError(
+                "The attached calculator does not currently support calculating "
+                "descriptors"
+            )
+
     def run(self) -> None:
         """Calculate descriptors for structure(s)."""
         if self.logger:

janus_core/calculations/geom_opt.py

@@ -300,7 +300,8 @@ def run(self) -> None:
 
         if self.logger:
             self.logger.info("After optimization spacegroup: %s", s_grp)
-            self.logger.info("Max force: %.6f", max_force)
+            self.logger.info("Max force: %s", max_force)
+            self.logger.info("Final energy: %s", self.struct.get_potential_energy())
 
         if not converged:
             warnings.warn(

janus_core/calculations/md.py

@@ -509,12 +509,21 @@ def __init__(
 
         self._parse_correlations()
 
-    def _set_time_step(self):
-        """Set time in fs and current dynamics step to info."""
+    def _set_info(self):
+        """Set time in fs, current dynamics step, and density to info."""
         time = (self.offset * self.timestep + self.dyn.get_time()) / units.fs
         step = self.offset + self.dyn.nsteps
         self.dyn.atoms.info["time_fs"] = time
         self.dyn.atoms.info["step"] = step
+        try:
+            density = (
+                np.sum(self.dyn.atoms.get_masses())
+                / self.dyn.atoms.get_volume()
+                * DENS_FACT
+            )
+            self.dyn.atoms.info["density"] = density
+        except ValueError:
+            self.dyn.atoms.info["density"] = 0.0
 
     def _prepare_restart(self) -> None:
         """Prepare restart files, structure and offset."""
@@ -726,19 +735,13 @@ def get_stats(self) -> dict[str, float]:
         e_kin = self.dyn.atoms.get_kinetic_energy() / self.n_atoms
         current_temp = e_kin / (1.5 * units.kB)
 
-        self._set_time_step()
+        self._set_info()
 
         time_now = datetime.datetime.now()
         real_time = time_now - self.dyn.atoms.info["real_time"]
         self.dyn.atoms.info["real_time"] = time_now
 
         try:
-            density = (
-                np.sum(self.dyn.atoms.get_masses())
-                / self.dyn.atoms.get_volume()
-                * DENS_FACT
-            )
-            self.dyn.atoms.info["density"] = density
             volume = self.dyn.atoms.get_volume()
             pressure = (
                 -np.trace(
@@ -754,7 +757,6 @@ def get_stats(self) -> dict[str, float]:
         except ValueError:
             volume = 0.0
             pressure = 0.0
-            density = 0.0
             pressure_tensor = np.zeros(6)
 
         return {
@@ -765,7 +767,7 @@ def get_stats(self) -> dict[str, float]:
             "EKin/N": e_kin,
             "T": current_temp,
             "ETot/N": e_pot + e_kin,
-            "Density": density,
+            "Density": self.dyn.atoms.info["density"],
             "Volume": volume,
             "P": pressure,
             "Pxx": pressure_tensor[0],
@@ -874,7 +876,7 @@ def _write_traj(self) -> None:
                 self.dyn.nsteps > self.traj_start + self.traj_start % self.traj_every
             )
 
-            self._set_time_step()
+            self._set_info()
             write_kwargs = self.write_kwargs
             write_kwargs["filename"] = self.traj_file
             write_kwargs["append"] = append
@@ -895,7 +897,7 @@ def _write_final_state(self) -> None:
         # Append if final file has been created
         append = self.created_final_file
 
-        self._set_time_step()
+        self._set_info()
         write_kwargs = self.write_kwargs
         write_kwargs["filename"] = self.final_file
         write_kwargs["append"] = append
@@ -998,7 +1000,7 @@ def _write_restart(self) -> None:
         if step > 0:
             write_kwargs = self.write_kwargs
             write_kwargs["filename"] = self._restart_file
-            self._set_time_step()
+            self._set_info()
 
             output_structs(
                 images=self.struct,

janus_core/calculations/phonons.py

@@ -20,7 +20,7 @@
     PathLike,
     PhononCalcs,
 )
-from janus_core.helpers.utils import none_to_dict, write_table
+from janus_core.helpers.utils import none_to_dict, track_progress, write_table
 
 
 class Phonons(BaseCalculation):
@@ -60,6 +60,8 @@ class Phonons(BaseCalculation):
         Size of supercell for calculation. Default is 2.
     displacement : float
         Displacement for force constants calculation, in A. Default is 0.01.
+    mesh : tuple[int, int, int]
+        Mesh for sampling. Default is (10, 10, 10).
     symmetrize : bool
         Whether to symmetrize force constants after calculation.
         Default is False.
@@ -88,6 +90,8 @@ class Phonons(BaseCalculation):
     file_prefix : Optional[PathLike]
         Prefix for output filenames. Default is inferred from chemical formula of the
         structure.
+    enable_progress_bar : bool
+        Whether to show a progress bar during phonon calculations. Default is False.
 
     Attributes
     ----------
@@ -106,7 +110,7 @@ class Phonons(BaseCalculation):
         Calculate band structure and optionally write and plot results.
     write_bands(bands_file, save_plots, plot_file)
         Write results of band structure calculations.
-    calc_thermal_props(write_thermal)
+    calc_thermal_props(mesh, write_thermal)
         Calculate thermal properties and optionally write results.
     write_thermal_props(thermal_file)
         Write results of thermal properties calculations.
@@ -138,6 +142,7 @@ def __init__(
         calcs: MaybeSequence[PhononCalcs] = (),
         supercell: MaybeList[int] = 2,
         displacement: float = 0.01,
+        mesh: tuple[int, int, int] = (10, 10, 10),
         symmetrize: bool = False,
         minimize: bool = False,
         minimize_kwargs: Optional[dict[str, Any]] = None,
@@ -149,6 +154,7 @@ def __init__(
         write_results: bool = True,
         write_full: bool = True,
         file_prefix: Optional[PathLike] = None,
+        enable_progress_bar: bool = False,
     ) -> None:
         """
         Initialise Phonons class.
@@ -186,6 +192,8 @@ def __init__(
             Size of supercell for calculation. Default is 2.
         displacement : float
             Displacement for force constants calculation, in A. Default is 0.01.
+        mesh : tuple[int, int, int]
+            Mesh for sampling. Default is (10, 10, 10).
         symmetrize : bool
             Whether to symmetrize force constants after calculations.
             Default is False.
@@ -214,11 +222,14 @@ def __init__(
         file_prefix : Optional[PathLike]
             Prefix for output filenames. Default is inferred from structure name, or
             chemical formula of the structure.
+        enable_progress_bar : bool
+            Whether to show a progress bar during phonon calculations. Default is False.
         """
         (read_kwargs, minimize_kwargs) = none_to_dict((read_kwargs, minimize_kwargs))
 
         self.calcs = calcs
         self.displacement = displacement
+        self.mesh = mesh
         self.symmetrize = symmetrize
         self.minimize = minimize
         self.minimize_kwargs = minimize_kwargs
@@ -229,6 +240,7 @@ def __init__(
         self.plot_to_file = plot_to_file
         self.write_results = write_results
         self.write_full = write_full
+        self.enable_progress_bar = enable_progress_bar
 
         # Ensure supercell is a valid list
         self.supercell = [supercell] * 3 if isinstance(supercell, int) else supercell
@@ -357,6 +369,11 @@ def calc_force_constants(
         phonon.generate_displacements(distance=self.displacement)
         disp_supercells = phonon.supercells_with_displacements
 
+        if self.enable_progress_bar:
+            disp_supercells = track_progress(
+                disp_supercells, "Computing displacements..."
+            )
+
         phonon.forces = [
             self._calc_forces(supercell)
             for supercell in disp_supercells
@@ -490,13 +507,18 @@ def write_bands(
             bplt.savefig(plot_file)
 
     def calc_thermal_props(
-        self, write_thermal: Optional[bool] = None, **kwargs
+        self,
+        mesh: Optional[tuple[int, int, int]] = None,
+        write_thermal: Optional[bool] = None,
+        **kwargs,
     ) -> None:
         """
         Calculate thermal properties and optionally write results.
 
         Parameters
         ----------
+        mesh : Optional[tuple[int, int, int]]
+            Mesh for sampling. Default is self.mesh.
         write_thermal : Optional[bool]
             Whether to write out thermal properties to file. Default is
             self.write_results.
@@ -506,6 +528,9 @@ def calc_thermal_props(
         if write_thermal is None:
             write_thermal = self.write_results
 
+        if mesh is None:
+            mesh = self.mesh
+
         # Calculate phonons if not already in results
         if "phonon" not in self.results:
             # Use general (self.write_results) setting for writing force constants
@@ -515,7 +540,7 @@ def calc_thermal_props(
             self.logger.info("Starting thermal properties calculation")
             self.tracker.start_task("Thermal calculation")
 
-        self.results["phonon"].run_mesh()
+        self.results["phonon"].run_mesh(mesh)
         self.results["phonon"].run_thermal_properties(
             t_step=self.temp_step, t_max=self.temp_max, t_min=self.temp_min
         )
@@ -563,7 +588,7 @@ def write_thermal_props(self, thermal_file: Optional[PathLike] = None) -> None:
     def calc_dos(
         self,
         *,
-        mesh: MaybeList[float] = (10, 10, 10),
+        mesh: Optional[tuple[int, int, int]] = None,
         write_dos: Optional[bool] = None,
         **kwargs,
     ) -> None:
@@ -572,8 +597,8 @@ def calc_dos(
 
         Parameters
         ----------
-        mesh : MaybeList[float]
-            Mesh for sampling. Default is (10, 10, 10).
+        mesh : Optional[tuple[int, int, int]]
+            Mesh for sampling. Default is self.mesh.
         write_dos : Optional[bool]
             Whether to write out results to file. Default is True.
         **kwargs
@@ -582,6 +607,9 @@ def calc_dos(
         if write_dos is None:
             write_dos = self.write_results
 
+        if mesh is None:
+            mesh = self.mesh
+
         # Calculate phonons if not already in results
         if "phonon" not in self.results:
             # Use general (self.write_results) setting for writing force constants
@@ -665,7 +693,7 @@ def write_dos(
     def calc_pdos(
         self,
         *,
-        mesh: MaybeList[float] = (10, 10, 10),
+        mesh: Optional[tuple[int, int, int]] = None,
         write_pdos: Optional[bool] = None,
         **kwargs,
     ) -> None:
@@ -674,8 +702,8 @@ def calc_pdos(
 
         Parameters
         ----------
-        mesh : MaybeList[float]
-            Mesh for sampling. Default is (10, 10, 10).
+        mesh : Optional[tuple[int, int, int]]
+            Mesh for sampling. Default is self.mesh.
         write_pdos : Optional[bool]
             Whether to write out results to file. Default is self.write_results.
         **kwargs
@@ -684,6 +712,9 @@ def calc_pdos(
         if write_pdos is None:
             write_pdos = self.write_results
 
+        if mesh is None:
+            mesh = self.mesh
+
         # Calculate phonons if not already in results
         if "phonon" not in self.results:
             # Use general (self.write_results) setting for writing force constants