From ebbff497ec4eba3f87674cd99fff29fe14dd420e Mon Sep 17 00:00:00 2001
From: Anyang Peng <137014849+anyangml@users.noreply.github.com>
Date: Fri, 21 Feb 2025 11:35:09 +0800
Subject: [PATCH 01/14] chore: migrate nve_md

---
 lambench/models/ase_models.py                         | 8 ++++++--
 lambench/tasks/calculator/{ => nve_md}/nve_md.py      | 2 +-
 lambench/tasks/calculator/{ => nve_md}/nve_md_data.py | 0
 test/tasks/calculator/test_nve_md.py                  | 2 +-
 4 files changed, 8 insertions(+), 4 deletions(-)
 rename lambench/tasks/calculator/{ => nve_md}/nve_md.py (98%)
 rename lambench/tasks/calculator/{ => nve_md}/nve_md_data.py (100%)

diff --git a/lambench/models/ase_models.py b/lambench/models/ase_models.py
index 01b97fd6..f8e3af88 100644
--- a/lambench/models/ase_models.py
+++ b/lambench/models/ase_models.py
@@ -81,7 +81,9 @@ def evaluate(self, task) -> Optional[dict[str, float]]:
             return self.run_ase_dptest(self.calc, task.test_data)
         elif isinstance(task, CalculatorTask):
             if task.task_name == "nve_md":
-                from lambench.tasks.calculator.nve_md import run_md_nve_simulation
+                from lambench.tasks.calculator.nve_md.nve_md import (
+                    run_md_nve_simulation,
+                )
 
                 num_steps = task.calculator_params.get("num_steps", 1000)
                 timestep = task.calculator_params.get("timestep", 1.0)
@@ -133,7 +135,9 @@ def run_ase_dptest(calc: Calculator, test_data: Path) -> dict:
                         if not np.isfinite(energy_predict):
                             raise ValueError("Energy prediction is non-finite.")
                     except (ValueError, RuntimeError):
-                        file = Path(f"failed_structures/{calc.name}/{atoms.symbols}.cif")
+                        file = Path(
+                            f"failed_structures/{calc.name}/{atoms.symbols}.cif"
+                        )
                         file.parent.mkdir(parents=True, exist_ok=True)
                         write(file, atoms)
                         logging.error(
diff --git a/lambench/tasks/calculator/nve_md.py b/lambench/tasks/calculator/nve_md/nve_md.py
similarity index 98%
rename from lambench/tasks/calculator/nve_md.py
rename to lambench/tasks/calculator/nve_md/nve_md.py
index df3139a3..e5bb384b 100644
--- a/lambench/tasks/calculator/nve_md.py
+++ b/lambench/tasks/calculator/nve_md/nve_md.py
@@ -11,7 +11,7 @@
 import numpy as np
 import time
 from typing import Optional
-from lambench.tasks.calculator.nve_md_data import TEST_DATA
+from lambench.tasks.calculator.nve_md.nve_md_data import TEST_DATA
 import logging
 
 
diff --git a/lambench/tasks/calculator/nve_md_data.py b/lambench/tasks/calculator/nve_md/nve_md_data.py
similarity index 100%
rename from lambench/tasks/calculator/nve_md_data.py
rename to lambench/tasks/calculator/nve_md/nve_md_data.py
diff --git a/test/tasks/calculator/test_nve_md.py b/test/tasks/calculator/test_nve_md.py
index 9c0d3cb0..42b0afb6 100644
--- a/test/tasks/calculator/test_nve_md.py
+++ b/test/tasks/calculator/test_nve_md.py
@@ -1,4 +1,4 @@
-from lambench.tasks.calculator.nve_md import (
+from lambench.tasks.calculator.nve_md.nve_md import (
     nve_simulation_single,
     run_md_nve_simulation,
 )

From fa3c6fbb8436004b90ffe17a4cb6986b57c917e1 Mon Sep 17 00:00:00 2001
From: Anyang Peng <137014849+anyangml@users.noreply.github.com>
Date: Fri, 21 Feb 2025 13:26:38 +0800
Subject: [PATCH 02/14] feat: add ASE relaxation method with symmetry
 constraints

---
 lambench/models/ase_models.py | 31 +++++++++++++++++++++++++++++--
 1 file changed, 29 insertions(+), 2 deletions(-)

diff --git a/lambench/models/ase_models.py b/lambench/models/ase_models.py
index f8e3af88..6252514e 100644
--- a/lambench/models/ase_models.py
+++ b/lambench/models/ase_models.py
@@ -3,14 +3,17 @@
 from pathlib import Path
 from typing import Optional
 
-import ase
 import dpdata
 import numpy as np
 from ase.calculators.calculator import Calculator
+from ase import Atoms
 from ase.io import write
 from tqdm import tqdm
 
 from lambench.models.basemodel import BaseLargeAtomModel
+from ase.optimize import FIRE
+from ase.constraints import FixSymmetry
+from ase.filters import FrechetCellFilter
 
 
 class ASEModel(BaseLargeAtomModel):
@@ -126,7 +129,7 @@ def run_ase_dptest(calc: Calculator, test_data: Path) -> dict:
                 sys = dpdata.LabeledSystem(filepth, fmt="deepmd/npy")
             for ls in tqdm(sys, desc="Set", leave=False):  # type: ignore
                 for frame in tqdm(ls, desc="Frames", leave=False):
-                    atoms: ase.Atoms = frame.to_ase_structure()[0]  # type: ignore
+                    atoms: Atoms = frame.to_ase_structure()[0]  # type: ignore
                     atoms.calc = calc
 
                     # Energy
@@ -233,3 +236,27 @@ def run_ase_dptest(calc: Calculator, test_data: Path) -> dict:
                 }
             )
         return res
+
+    @staticmethod
+    def run_ase_relaxation(
+        atoms: Atoms,
+        calc: Calculator,
+        fmax: float = 5e-3,
+        steps: int = 500,
+        fix_symmetry: bool = True,
+        relax_cell: bool = True,
+    ) -> Optional[Atoms]:
+        atoms.calc = calc
+        if fix_symmetry:
+            atoms.set_constraint(FixSymmetry(atoms))
+        if relax_cell:
+            atoms = FrechetCellFilter(atoms)
+        opt = FIRE(atoms, trajectory=None, logfile=None)
+        try:
+            opt.run(fmax=fmax, steps=steps)
+        except Exception as e:
+            logging.error(f"Relaxation failed: {e}")
+            return None
+        if relax_cell:
+            atoms = atoms.atoms
+        return atoms

From e427a13c67280537fd17fd4927ebb315623ec203 Mon Sep 17 00:00:00 2001
From: Anyang Peng <137014849+anyangml@users.noreply.github.com>
Date: Fri, 21 Feb 2025 13:27:12 +0800
Subject: [PATCH 03/14] feat: add phonopy dependency to pyproject.toml

---
 pyproject.toml | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/pyproject.toml b/pyproject.toml
index 9de47ee5..d26f9206 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -38,6 +38,8 @@ orb = ["orb-models","pynanoflann@git+https://github.com/dwastberg/pynanoflann#eg
 sevenn = ["sevenn"]
 test = ["pytest"]
 dflow = ["pydflow", "lbg"]
+phonopy = ["phonopy@git+https://github.com/phonopy/phonopy.git"]
+
 
 [project.urls]
 Homepage = "https://github.com/deepmodeling/LAMBench"

From 012796a8b32f0f931294001c9ba8d878e6b39443 Mon Sep 17 00:00:00 2001
From: Anyang Peng <137014849+anyangml@users.noreply.github.com>
Date: Mon, 24 Feb 2025 15:01:15 +0800
Subject: [PATCH 04/14] feat: implement phonon simulation and utility functions

---
 lambench/tasks/calculator/phonon/phonon.py    | 49 +++++++++++++++++++
 .../tasks/calculator/phonon/phonon_utils.py   | 15 ++++++
 2 files changed, 64 insertions(+)
 create mode 100644 lambench/tasks/calculator/phonon/phonon.py
 create mode 100644 lambench/tasks/calculator/phonon/phonon_utils.py

diff --git a/lambench/tasks/calculator/phonon/phonon.py b/lambench/tasks/calculator/phonon/phonon.py
new file mode 100644
index 00000000..08e5ccea
--- /dev/null
+++ b/lambench/tasks/calculator/phonon/phonon.py
@@ -0,0 +1,49 @@
+from lambench.models.ase_models import ASEModel
+from ase import Atoms
+import phonopy
+from lambench.tasks.calculator.phonon.phonon_utils import ase_to_phonopy_atoms
+
+
+def run_phonon_simulation_single(
+    atoms: Atoms,
+    model: ASEModel,
+    distance: float = 0.01,
+):
+    """
+    Run a phonon simulation using the given model and return performance metrics.
+    """
+
+    # Step 1: Run relaxation
+    atoms: Atoms = model.run_ase_relaxation(atoms, model.calc)
+    if atoms is None:
+        return None
+
+    # Step 2: Convert ASE Atoms object to PhonopyAtoms object
+    phonon_atoms = ase_to_phonopy_atoms(atoms)
+    phonon = phonopy.Phonopy(
+        phonon_atoms, supercell_matrix=[[2, 0, 0], [0, 2, 0], [0, 0, 2]]
+    )
+
+    # Step 3: Generate displacements
+    phonon.generate_displacements(distance=distance, is_diagonal=False)
+
+    # Step 4: Calculate force constants
+    forcesets = []
+
+    for frame in phonon.supercells_with_displacements:
+        frame_atom = Atoms(
+            cell=frame.cell,
+            symbols=frame.symbols,
+            scaled_positions=frame.scaled_positions,
+            pbc=True,
+        )
+        frame_atom.calc = model.calc
+        forces = frame_atom.get_forces()
+        forcesets.append(forces)
+
+    phonon.forces = forcesets
+    phonon.produce_force_constants()
+    phonon.symmetrize_force_constants()
+
+    # phonon.auto_band_structure()
+    # phonon.plot_band_structure().show()
diff --git a/lambench/tasks/calculator/phonon/phonon_utils.py b/lambench/tasks/calculator/phonon/phonon_utils.py
new file mode 100644
index 00000000..c8fd1131
--- /dev/null
+++ b/lambench/tasks/calculator/phonon/phonon_utils.py
@@ -0,0 +1,15 @@
+from ase import Atoms
+from phonopy.structure.atoms import PhonopyAtoms
+
+
+def ase_to_phonopy_atoms(atoms: Atoms) -> PhonopyAtoms:
+    """
+    Convert ASE Atoms object to PhonopyAtoms object.
+    """
+    # Extract atomic symbols and positions
+    symbols = atoms.get_chemical_symbols()
+    positions = atoms.get_positions()
+    cell = atoms.get_cell()
+    masses = atoms.get_masses()
+
+    return PhonopyAtoms(symbols=symbols, positions=positions, cell=cell, masses=masses)

From 08f73ab6dc373cacfeee47b2bb85e46f3dd3e273 Mon Sep 17 00:00:00 2001
From: Anyang Peng <137014849+anyangml@users.noreply.github.com>
Date: Tue, 25 Feb 2025 15:13:09 +0800
Subject: [PATCH 05/14] feat/support-phonon-MDR-task

---
 lambench/models/ase_models.py                 |  12 ++
 .../tasks/calculator/calculator_tasks.yml     |   4 +
 lambench/tasks/calculator/phonon/phonon.py    | 176 ++++++++++++++----
 .../tasks/calculator/phonon/phonon_utils.py   |  19 ++
 pyproject.toml                                |   1 +
 5 files changed, 180 insertions(+), 32 deletions(-)

diff --git a/lambench/models/ase_models.py b/lambench/models/ase_models.py
index 6252514e..30067657 100644
--- a/lambench/models/ase_models.py
+++ b/lambench/models/ase_models.py
@@ -96,6 +96,18 @@ def evaluate(self, task) -> Optional[dict[str, float]]:
                         self, num_steps, timestep, temperature_K
                     )
                 }
+            elif task.task_name == "phonon_mdr":
+                from lambench.tasks.calculator.phonon.phonon import (
+                    run_phonon_simulation,
+                )
+
+                task.workdir.mkdir(exist_ok=True)
+                distance = task.calculator_params.get("distance", 0.01)
+                return {
+                    "metrics": run_phonon_simulation(
+                        self, task.test_data, distance, task.workdir
+                    )
+                }
             else:
                 raise NotImplementedError(f"Task {task.task_name} is not implemented.")
 
diff --git a/lambench/tasks/calculator/calculator_tasks.yml b/lambench/tasks/calculator/calculator_tasks.yml
index 6b8d9b66..65aba576 100644
--- a/lambench/tasks/calculator/calculator_tasks.yml
+++ b/lambench/tasks/calculator/calculator_tasks.yml
@@ -4,3 +4,7 @@ nve_md:
     num_steps: 10000
     timestep: 1.0
     temperature: 300
+phonon_mdr:
+  test_data: null
+  calculator_params:
+    distance: 0.01
diff --git a/lambench/tasks/calculator/phonon/phonon.py b/lambench/tasks/calculator/phonon/phonon.py
index 08e5ccea..64878959 100644
--- a/lambench/tasks/calculator/phonon/phonon.py
+++ b/lambench/tasks/calculator/phonon/phonon.py
@@ -1,49 +1,161 @@
+from pathlib import Path
+from typing import Optional
+
+import numpy as np
 from lambench.models.ase_models import ASEModel
 from ase import Atoms
 import phonopy
-from lambench.tasks.calculator.phonon.phonon_utils import ase_to_phonopy_atoms
+from phonopy.harmonic.dynmat_to_fc import get_commensurate_points
+from lambench.tasks.calculator.phonon.phonon_utils import (
+    ase_to_phonopy_atoms,
+    phonopy_to_ase_atoms,
+    THz_TO_K,
+)
+import logging
+import yaml
+import pandas as pd
+from sklearn.metrics import mean_absolute_error
 
 
 def run_phonon_simulation_single(
-    atoms: Atoms,
     model: ASEModel,
-    distance: float = 0.01,
-):
+    phonon_file: Path,
+    distance: float,
+    workdir: Path,
+) -> Optional[dict[str, float]]:
     """
-    Run a phonon simulation using the given model and return performance metrics.
+    Run phonon related calculations for a single given phonon file.
+
+    Parameters:
+        model: ASEModel object.
+        phonon_file: Path to the phonon file.
+        distance: Distance for displacements.
+        workdir: Path to the working directory.
     """
+    try:
+        # Step 1: Run relaxation
+        atoms: Atoms = phonopy_to_ase_atoms(phonon_file)
+        atoms = model.run_ase_relaxation(atoms, model.calc)
 
-    # Step 1: Run relaxation
-    atoms: Atoms = model.run_ase_relaxation(atoms, model.calc)
-    if atoms is None:
-        return None
+        # Step 2: Convert ASE Atoms object to PhonopyAtoms object
+        phonon_atoms = ase_to_phonopy_atoms(atoms)
+        phonon = phonopy.Phonopy(
+            phonon_atoms, supercell_matrix=[[2, 0, 0], [0, 2, 0], [0, 0, 2]]
+        )
+
+        # Step 3: Generate displacements
+        phonon.generate_displacements(distance=distance, is_diagonal=False)
+
+        # Step 4: Calculate force constants
+        forcesets = []
+
+        for frame in phonon.supercells_with_displacements:
+            frame_atom = Atoms(
+                cell=frame.cell,
+                symbols=frame.symbols,
+                scaled_positions=frame.scaled_positions,
+                pbc=True,
+            )
+            frame_atom.calc = model.calc
+            forces = frame_atom.get_forces()
+            forcesets.append(forces)
+
+        phonon.forces = forcesets
+        phonon.produce_force_constants()
+        phonon.symmetrize_force_constants()
+
+        # Step 5: save output files
+
+        phonon.save(workdir / phonon_file.name, settings={"force_constants": True})
+
+        # Step 6: Calculate thermal properties
+        phonon.init_mesh()
+        phonon.run_mesh()
+        phonon.run_thermal_properties(temperatures=(300,))
+        thermal_dict = phonon.get_thermal_properties_dict()
 
-    # Step 2: Convert ASE Atoms object to PhonopyAtoms object
-    phonon_atoms = ase_to_phonopy_atoms(atoms)
-    phonon = phonopy.Phonopy(
-        phonon_atoms, supercell_matrix=[[2, 0, 0], [0, 2, 0], [0, 0, 2]]
-    )
+        commensurate_q = get_commensurate_points(phonon.supercell_matrix)
+        phonon_freqs = np.array([phonon.get_frequencies(q) for q in commensurate_q])
 
-    # Step 3: Generate displacements
-    phonon.generate_displacements(distance=distance, is_diagonal=False)
+        # Step 7: Updata output files
+        with open(workdir / phonon_file.name, "r") as f:
+            output = yaml.load(f, yaml.FullLoader)
 
-    # Step 4: Calculate force constants
-    forcesets = []
+        output["free_e"] = thermal_dict["free_energy"].tolist()
+        output["entropy"] = thermal_dict["entropy"].tolist()
+        output["heat_capacity"] = thermal_dict["heat_capacity"].tolist()
+        output["phonon_freq"] = phonon_freqs.tolist()
 
-    for frame in phonon.supercells_with_displacements:
-        frame_atom = Atoms(
-            cell=frame.cell,
-            symbols=frame.symbols,
-            scaled_positions=frame.scaled_positions,
-            pbc=True,
+        # TODO: optional: update and save output files
+        return {
+            "mp_id": phonon_file.name.split(".")[0],
+            "entropy": output["entropy"][0],
+            "heat_capacity": output["heat_capacity"][0],
+            "free_energy": output["free_e"][0],
+            "max_freq": np.max(np.array(phonon_freqs)) * THz_TO_K,
+        }
+
+    except Exception as e:
+        logging.error(f"Error occured for {str(phonon_file.name)}: {e}")
+        return None
+
+
+def run_phonon_simulation(
+    model: ASEModel,
+    test_data: Path,
+    distance: float,
+    workdir: Path,
+) -> dict[str, float]:
+    """
+    This function runs phonon simulations for a list of test systems using the given model.
+    """
+    test_files = list(test_data.glob("*.yaml.bz2"))
+    if len(test_files) == 0:
+        logging.error("No test files found.")
+        return {}
+    logging.info(f"Running phonon simulations for {len(test_files)} files...")
+
+    dataframe_rows = []
+    for test_file in test_files:
+        result = run_phonon_simulation_single(
+            model,
+            test_file,
+            distance,
+            workdir,
         )
-        frame_atom.calc = model.calc
-        forces = frame_atom.get_forces()
-        forcesets.append(forces)
+        logging.info(f"Simulation completed for system {str(test_file.name)}.\n")
+
+    if result is not None:
+        dataframe_rows.append(result)
+    preds = pd.DataFrame(dataframe_rows)
+
+    # Post-processing
+    results = {}
+    try:
+        labels = pd.read_csv(test_data / "pbe.csv")
+        TOTAL_RECORDS = len(labels)
+        preds.sort_values("mp_id", inplace=True)
+        labels.sort_values("mp_id", inplace=True)
+
+        preds = preds[np.isfinite(preds["free_energy"])]
+        preds = preds[np.isfinite(preds["heat_capacity"])]
 
-    phonon.forces = forcesets
-    phonon.produce_force_constants()
-    phonon.symmetrize_force_constants()
+        valid_mp_ids = set(preds["mp_id"])
+        labels = labels[labels["mp_id"].isin(valid_mp_ids)]
+        preds = preds[preds["mp_id"].isin(valid_mp_ids)]
 
-    # phonon.auto_band_structure()
-    # phonon.plot_band_structure().show()
+        success_rate = len(preds) / TOTAL_RECORDS
+        mae_wmax = mean_absolute_error(labels["max_freq"], preds["max_freq"])
+        mae_s = mean_absolute_error(labels["entropy"], preds["entropy"])
+        mae_f = mean_absolute_error(labels["free_energy"], preds["free_energy"])
+        mae_c = mean_absolute_error(labels["heat_capacity"], preds["heat_capacity"])
+        results = {
+            "success_rate": success_rate,
+            "mae_max_freq": mae_wmax,
+            "mae_entropy": mae_s,
+            "mae_free_energy": mae_f,
+            "mae_heat_capacity": mae_c,
+        }
+    except Exception as e:
+        logging.error(f"Error occured during post-processing: {e}")
+    return results
diff --git a/lambench/tasks/calculator/phonon/phonon_utils.py b/lambench/tasks/calculator/phonon/phonon_utils.py
index c8fd1131..fcf1206b 100644
--- a/lambench/tasks/calculator/phonon/phonon_utils.py
+++ b/lambench/tasks/calculator/phonon/phonon_utils.py
@@ -1,5 +1,7 @@
 from ase import Atoms
 from phonopy.structure.atoms import PhonopyAtoms
+from pathlib import Path
+import phonopy
 
 
 def ase_to_phonopy_atoms(atoms: Atoms) -> PhonopyAtoms:
@@ -13,3 +15,20 @@ def ase_to_phonopy_atoms(atoms: Atoms) -> PhonopyAtoms:
     masses = atoms.get_masses()
 
     return PhonopyAtoms(symbols=symbols, positions=positions, cell=cell, masses=masses)
+
+
+def phonopy_to_ase_atoms(phonon_file: Path) -> Atoms:
+    """
+    Convert PhonopyAtoms object to ASE Atoms object.
+    """
+    phonon = phonopy.load(phonon_file)
+    return Atoms(
+        cell=phonon.unitcell.cell,
+        symbols=phonon.unitcell.symbols,
+        scaled_positions=phonon.unitcell.scaled_positions,
+        pbc=True,
+    )
+
+
+# Constants unit conversion
+THz_TO_K = 47.9924
diff --git a/pyproject.toml b/pyproject.toml
index d26f9206..ad786ad3 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -39,6 +39,7 @@ sevenn = ["sevenn"]
 test = ["pytest"]
 dflow = ["pydflow", "lbg"]
 phonopy = ["phonopy@git+https://github.com/phonopy/phonopy.git"]
+sklearn = ["scikit-learn"]
 
 
 [project.urls]

From 3a68268760753fda43db7cdc88f018e90179d09e Mon Sep 17 00:00:00 2001
From: Anyang Peng <137014849+anyangml@users.noreply.github.com>
Date: Tue, 25 Feb 2025 16:04:12 +0800
Subject: [PATCH 06/14] fxi: indentation

---
 lambench/tasks/calculator/phonon/phonon.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/lambench/tasks/calculator/phonon/phonon.py b/lambench/tasks/calculator/phonon/phonon.py
index 64878959..f4dab095 100644
--- a/lambench/tasks/calculator/phonon/phonon.py
+++ b/lambench/tasks/calculator/phonon/phonon.py
@@ -125,8 +125,8 @@ def run_phonon_simulation(
         )
         logging.info(f"Simulation completed for system {str(test_file.name)}.\n")
 
-    if result is not None:
-        dataframe_rows.append(result)
+        if result is not None:
+            dataframe_rows.append(result)
     preds = pd.DataFrame(dataframe_rows)
 
     # Post-processing

From 03e2df4b7f237654e823c978a65b38e28bf10d58 Mon Sep 17 00:00:00 2001
From: Anyang Peng <137014849+anyangml@users.noreply.github.com>
Date: Wed, 26 Feb 2025 10:46:31 +0800
Subject: [PATCH 07/14] fix: remove supercell matrix

---
 lambench/tasks/calculator/phonon/phonon.py      | 17 ++++++++---------
 .../tasks/calculator/phonon/phonon_utils.py     |  8 ++++----
 2 files changed, 12 insertions(+), 13 deletions(-)

diff --git a/lambench/tasks/calculator/phonon/phonon.py b/lambench/tasks/calculator/phonon/phonon.py
index f4dab095..273ab6dd 100644
--- a/lambench/tasks/calculator/phonon/phonon.py
+++ b/lambench/tasks/calculator/phonon/phonon.py
@@ -35,13 +35,11 @@ def run_phonon_simulation_single(
     try:
         # Step 1: Run relaxation
         atoms: Atoms = phonopy_to_ase_atoms(phonon_file)
-        atoms = model.run_ase_relaxation(atoms, model.calc)
+        atoms = model.run_ase_relaxation(atoms, model.calc, fmax=1e-3)
 
         # Step 2: Convert ASE Atoms object to PhonopyAtoms object
         phonon_atoms = ase_to_phonopy_atoms(atoms)
-        phonon = phonopy.Phonopy(
-            phonon_atoms, supercell_matrix=[[2, 0, 0], [0, 2, 0], [0, 0, 2]]
-        )
+        phonon = phonopy.Phonopy(phonon_atoms)
 
         # Step 3: Generate displacements
         phonon.generate_displacements(distance=distance, is_diagonal=False)
@@ -137,12 +135,13 @@ def run_phonon_simulation(
         preds.sort_values("mp_id", inplace=True)
         labels.sort_values("mp_id", inplace=True)
 
-        preds = preds[np.isfinite(preds["free_energy"])]
-        preds = preds[np.isfinite(preds["heat_capacity"])]
-
-        valid_mp_ids = set(preds["mp_id"])
+        # Filter predictions and labels based on valid mp_ids
+        valid_preds = preds[
+            np.isfinite(preds[["free_energy", "heat_capacity"]]).all(axis=1)
+        ]
+        valid_mp_ids = set(valid_preds["mp_id"])
         labels = labels[labels["mp_id"].isin(valid_mp_ids)]
-        preds = preds[preds["mp_id"].isin(valid_mp_ids)]
+        preds = valid_preds
 
         success_rate = len(preds) / TOTAL_RECORDS
         mae_wmax = mean_absolute_error(labels["max_freq"], preds["max_freq"])
diff --git a/lambench/tasks/calculator/phonon/phonon_utils.py b/lambench/tasks/calculator/phonon/phonon_utils.py
index fcf1206b..b2414b42 100644
--- a/lambench/tasks/calculator/phonon/phonon_utils.py
+++ b/lambench/tasks/calculator/phonon/phonon_utils.py
@@ -4,6 +4,10 @@
 import phonopy
 
 
+# Constants unit conversion
+THz_TO_K = 47.9924
+
+
 def ase_to_phonopy_atoms(atoms: Atoms) -> PhonopyAtoms:
     """
     Convert ASE Atoms object to PhonopyAtoms object.
@@ -28,7 +32,3 @@ def phonopy_to_ase_atoms(phonon_file: Path) -> Atoms:
         scaled_positions=phonon.unitcell.scaled_positions,
         pbc=True,
     )
-
-
-# Constants unit conversion
-THz_TO_K = 47.9924

From 394c92c15f226c695fd162009e96852b2729a1a9 Mon Sep 17 00:00:00 2001
From: Anyang Peng <137014849+anyangml@users.noreply.github.com>
Date: Wed, 26 Feb 2025 11:04:25 +0800
Subject: [PATCH 08/14] chore: add tqdm

---
 lambench/tasks/calculator/phonon/phonon.py | 25 +++++++++++++++++++++-
 1 file changed, 24 insertions(+), 1 deletion(-)

diff --git a/lambench/tasks/calculator/phonon/phonon.py b/lambench/tasks/calculator/phonon/phonon.py
index 273ab6dd..65cfb76b 100644
--- a/lambench/tasks/calculator/phonon/phonon.py
+++ b/lambench/tasks/calculator/phonon/phonon.py
@@ -15,6 +15,7 @@
 import yaml
 import pandas as pd
 from sklearn.metrics import mean_absolute_error
+from tqdm import tqdm
 
 
 def run_phonon_simulation_single(
@@ -114,7 +115,7 @@ def run_phonon_simulation(
     logging.info(f"Running phonon simulations for {len(test_files)} files...")
 
     dataframe_rows = []
-    for test_file in test_files:
+    for test_file in tqdm(test_files):
         result = run_phonon_simulation_single(
             model,
             test_file,
@@ -158,3 +159,25 @@ def run_phonon_simulation(
     except Exception as e:
         logging.error(f"Error occured during post-processing: {e}")
     return results
+
+
+if __name__ == "__main__":
+    dp = ASEModel(
+        model_name="test",
+        model_type="ASE",
+        model_family="DP",
+        model_path=Path("/Users/aisi_ap/Downloads/salex_ft_110w.pth"),
+        model_metadata={
+            "author": "alex",
+        },
+        virtualenv="test",
+    )
+
+    print(
+        run_phonon_simulation(
+            dp,
+            Path("/Users/aisi_ap/Downloads/"),
+            0.01,
+            Path("/Users/aisi_ap/Desktop/lamtest"),
+        )
+    )

From 797c4db3bbf7674daa2dfc541e17dc89d31f99eb Mon Sep 17 00:00:00 2001
From: Anyang Peng <137014849+anyangml@users.noreply.github.com>
Date: Wed, 26 Feb 2025 11:08:17 +0800
Subject: [PATCH 09/14] chore:add reference

---
 lambench/tasks/calculator/phonon/phonon.py | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/lambench/tasks/calculator/phonon/phonon.py b/lambench/tasks/calculator/phonon/phonon.py
index 65cfb76b..1c95b35f 100644
--- a/lambench/tasks/calculator/phonon/phonon.py
+++ b/lambench/tasks/calculator/phonon/phonon.py
@@ -1,3 +1,17 @@
+"""
+Code adapted from the following paper and code:
+
+@misc{loew2024universalmachinelearninginteratomic,
+      title={Universal Machine Learning Interatomic Potentials are Ready for Phonons},
+      author={Antoine Loew and Dewen Sun and Hai-Chen Wang and Silvana Botti and Miguel A. L. Marques},
+      year={2024},
+      eprint={2412.16551},
+      archivePrefix={arXiv},
+      primaryClass={cond-mat.mtrl-sci},
+      url={https://arxiv.org/abs/2412.16551},
+}
+"""
+
 from pathlib import Path
 from typing import Optional
 

From c41df1197f6c76becd11f86d6e3396ded5d6b137 Mon Sep 17 00:00:00 2001
From: Anyang Peng <137014849+anyangml@users.noreply.github.com>
Date: Wed, 26 Feb 2025 11:19:09 +0800
Subject: [PATCH 10/14] refactor: remove unused main execution block

---
 lambench/tasks/calculator/phonon/phonon.py | 22 ----------------------
 1 file changed, 22 deletions(-)

diff --git a/lambench/tasks/calculator/phonon/phonon.py b/lambench/tasks/calculator/phonon/phonon.py
index 1c95b35f..5a61f51c 100644
--- a/lambench/tasks/calculator/phonon/phonon.py
+++ b/lambench/tasks/calculator/phonon/phonon.py
@@ -173,25 +173,3 @@ def run_phonon_simulation(
     except Exception as e:
         logging.error(f"Error occured during post-processing: {e}")
     return results
-
-
-if __name__ == "__main__":
-    dp = ASEModel(
-        model_name="test",
-        model_type="ASE",
-        model_family="DP",
-        model_path=Path("/Users/aisi_ap/Downloads/salex_ft_110w.pth"),
-        model_metadata={
-            "author": "alex",
-        },
-        virtualenv="test",
-    )
-
-    print(
-        run_phonon_simulation(
-            dp,
-            Path("/Users/aisi_ap/Downloads/"),
-            0.01,
-            Path("/Users/aisi_ap/Desktop/lamtest"),
-        )
-    )

From c40ee2b9d0b993643032f0cd659446c293e0cd85 Mon Sep 17 00:00:00 2001
From: Anyang Peng <137014849+anyangml@users.noreply.github.com>
Date: Wed, 26 Feb 2025 12:34:00 +0800
Subject: [PATCH 11/14] chore: update data path

---
 lambench/tasks/calculator/calculator_tasks.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/lambench/tasks/calculator/calculator_tasks.yml b/lambench/tasks/calculator/calculator_tasks.yml
index 65aba576..2643cc2b 100644
--- a/lambench/tasks/calculator/calculator_tasks.yml
+++ b/lambench/tasks/calculator/calculator_tasks.yml
@@ -5,6 +5,6 @@ nve_md:
     timestep: 1.0
     temperature: 300
 phonon_mdr:
-  test_data: null
+  test_data: /bohr/lambench-phonon-y7vk/v1/MDR_PBE_phonon
   calculator_params:
     distance: 0.01

From dd95c7b63b238b5c7f11b1405f56f092a69b4b76 Mon Sep 17 00:00:00 2001
From: Anyang Peng <137014849+anyangml@users.noreply.github.com>
Date: Wed, 26 Feb 2025 17:51:20 +0800
Subject: [PATCH 12/14] fix: add supercell info

---
 lambench/tasks/calculator/phonon/phonon.py       | 4 +++-
 lambench/tasks/calculator/phonon/phonon_utils.py | 1 +
 2 files changed, 4 insertions(+), 1 deletion(-)

diff --git a/lambench/tasks/calculator/phonon/phonon.py b/lambench/tasks/calculator/phonon/phonon.py
index 5a61f51c..b69d1458 100644
--- a/lambench/tasks/calculator/phonon/phonon.py
+++ b/lambench/tasks/calculator/phonon/phonon.py
@@ -54,7 +54,9 @@ def run_phonon_simulation_single(
 
         # Step 2: Convert ASE Atoms object to PhonopyAtoms object
         phonon_atoms = ase_to_phonopy_atoms(atoms)
-        phonon = phonopy.Phonopy(phonon_atoms)
+        phonon = phonopy.Phonopy(
+            phonon_atoms, supercell_matrix=atoms.info["supercell_matrix"]
+        )
 
         # Step 3: Generate displacements
         phonon.generate_displacements(distance=distance, is_diagonal=False)
diff --git a/lambench/tasks/calculator/phonon/phonon_utils.py b/lambench/tasks/calculator/phonon/phonon_utils.py
index b2414b42..e96d3bd9 100644
--- a/lambench/tasks/calculator/phonon/phonon_utils.py
+++ b/lambench/tasks/calculator/phonon/phonon_utils.py
@@ -31,4 +31,5 @@ def phonopy_to_ase_atoms(phonon_file: Path) -> Atoms:
         symbols=phonon.unitcell.symbols,
         scaled_positions=phonon.unitcell.scaled_positions,
         pbc=True,
+        info={"supercell_matrix": phonon.supercell_matrix},
     )

From 61fbc9efdfe963229196e90bfe77449f0d27d1b2 Mon Sep 17 00:00:00 2001
From: Chun Cai <amoycaic@gmail.com>
Date: Wed, 26 Feb 2025 18:00:06 +0800
Subject: [PATCH 13/14] sort imports

---
 lambench/tasks/calculator/phonon/phonon.py | 17 +++++++++--------
 1 file changed, 9 insertions(+), 8 deletions(-)

diff --git a/lambench/tasks/calculator/phonon/phonon.py b/lambench/tasks/calculator/phonon/phonon.py
index b69d1458..a83ae348 100644
--- a/lambench/tasks/calculator/phonon/phonon.py
+++ b/lambench/tasks/calculator/phonon/phonon.py
@@ -12,24 +12,25 @@
 }
 """
 
+import logging
 from pathlib import Path
 from typing import Optional
 
 import numpy as np
-from lambench.models.ase_models import ASEModel
-from ase import Atoms
+import pandas as pd
 import phonopy
+import yaml
+from ase import Atoms
 from phonopy.harmonic.dynmat_to_fc import get_commensurate_points
+from sklearn.metrics import mean_absolute_error
+from tqdm import tqdm
+
+from lambench.models.ase_models import ASEModel
 from lambench.tasks.calculator.phonon.phonon_utils import (
+    THz_TO_K,
     ase_to_phonopy_atoms,
     phonopy_to_ase_atoms,
-    THz_TO_K,
 )
-import logging
-import yaml
-import pandas as pd
-from sklearn.metrics import mean_absolute_error
-from tqdm import tqdm
 
 
 def run_phonon_simulation_single(

From 511789f029c0b876bb26e3cebcd46826b71c240e Mon Sep 17 00:00:00 2001
From: Anyang Peng <137014849+anyangml@users.noreply.github.com>
Date: Wed, 26 Feb 2025 18:02:50 +0800
Subject: [PATCH 14/14] chore: update phonon task dependency

Co-authored-by: Chun Cai <amoycaic@gmail.com>
---
 pyproject.toml | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index ad786ad3..3c01e517 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -38,8 +38,7 @@ orb = ["orb-models","pynanoflann@git+https://github.com/dwastberg/pynanoflann#eg
 sevenn = ["sevenn"]
 test = ["pytest"]
 dflow = ["pydflow", "lbg"]
-phonopy = ["phonopy@git+https://github.com/phonopy/phonopy.git"]
-sklearn = ["scikit-learn"]
+phonopy = ["phonopy@git+https://github.com/phonopy/phonopy.git", "scikit-learn"]
 
 
 [project.urls]