Merge pull request #423 from anyangml2nd/main

fix: separate spin charge

Author: anyangml

lambench/models/ase_models.py

@@ -398,8 +398,12 @@ def run_ase_dptest(
                                 }
                             )
                         elif model.model_family == "DP":
-                            atoms.info.update({"charge_spin": frame.data["fparam"]})
-                            atoms.info.update({"fparam": frame.data["fparam"]})
+                            if getattr(
+                                model.model_metadata, "use_separate_spin_charge", False
+                            ):
+                                atoms.info.update({"charge_spin": frame.data["fparam"]})
+                            else:
+                                atoms.info.update({"fparam": frame.data["fparam"]})
                     atoms.calc = calc
 
                     # Energy

lambench/tasks/calculator/binding/binding.py

@@ -38,16 +38,18 @@ def run_inference(
         try:
             for atoms in (site, drug, combo):
                 atoms.calc = calc
-                atoms.info.update(
-                    {
-                        "charge_spin": np.array(
-                            [atoms.info["charge"], atoms.info["spin"]]
-                        )
-                    }
-                )
-                atoms.info.update(
-                    {"fparam": np.array([atoms.info["charge"], atoms.info["spin"]])}
-                )
+                if getattr(model.model_metadata, "use_separate_spin_charge", False):
+                    atoms.info.update(
+                        {
+                            "charge_spin": np.array(
+                                [atoms.info["charge"], atoms.info["spin"]]
+                            )
+                        }
+                    )
+                else:
+                    atoms.info.update(
+                        {"fparam": np.array([atoms.info["charge"], atoms.info["spin"]])}
+                    )
 
             site_energy = site.get_potential_energy()
             drug_energy = drug.get_potential_energy()

lambench/tasks/calculator/rxn_barrier/barrier.py

@@ -57,16 +57,18 @@ def run_inference(
                     lookup_table["ID"] == reactant
                 ].index.values[0]
                 atoms = traj[structure_index]
-                atoms.info.update(
-                    {
-                        "charge_spin": np.array(
-                            [atoms.info["charge"], atoms.info["spin"]]
-                        )
-                    }
-                )
-                atoms.info.update(
-                    {"fparam": np.array([atoms.info["charge"], atoms.info["spin"]])}
-                )
+                if getattr(model.model_metadata, "use_separate_spin_charge", False):
+                    atoms.info.update(
+                        {
+                            "charge_spin": np.array(
+                                [atoms.info["charge"], atoms.info["spin"]]
+                            )
+                        }
+                    )
+                else:
+                    atoms.info.update(
+                        {"fparam": np.array([atoms.info["charge"], atoms.info["spin"]])}
+                    )
                 atoms.calc = calc
                 energy = atoms.get_potential_energy()
                 pred += stoi * energy