diff --git a/ash/modules/module_machine_learning.py b/ash/modules/module_machine_learning.py index 023ab227a..137aa185d 100644 --- a/ash/modules/module_machine_learning.py +++ b/ash/modules/module_machine_learning.py @@ -194,7 +194,9 @@ def create_ML_training_data(xyz_dir=None, dcd_trajectory=None, xyz_trajectory=No os.chdir('..') # Remove old files if present - for f in ["train_data.xyz", "train_data.energies", "train_data.gradients", "train_data_mace.xyz"]: + for f in ["train_data.xyz", "train_data.energies", "train_data.gradients", "train_data_mace.xyz", + "train_data_theory1.energies", "train_data_theory2.energies", + "train_data_theory1.gradients", "train_data_theory2.gradients"]: try: os.remove(f) except: @@ -205,6 +207,11 @@ def create_ML_training_data(xyz_dir=None, dcd_trajectory=None, xyz_trajectory=No gradients=[] fragments=[] labels=[] + # When delta-learning (two theories), also keep the individual theory contributions + energies_theory1=[] + energies_theory2=[] + gradients_theory1=[] + gradients_theory2=[] # Removing theory_1.cleanup() @@ -243,6 +250,12 @@ def create_ML_training_data(xyz_dir=None, dcd_trajectory=None, xyz_trajectory=No energy = result_2.energy - result_1.energy if Grad is True: gradient = result_2.gradient - result_1.gradient + # Keep individual theory contributions for separate output files + energies_theory1.append(result_1.energy) + energies_theory2.append(result_2.energy) + if Grad is True: + gradients_theory1.append(result_1.gradient) + gradients_theory2.append(result_2.gradient) else: energy = result_1.energy if Grad is True: @@ -291,6 +304,9 @@ def create_ML_training_data(xyz_dir=None, dcd_trajectory=None, xyz_trajectory=No if delta is True: energy = results_theory2.energies_dict[l] - results_theory1.energies_dict[l] print("energy:", energy) + # Keep individual theory contributions for separate output files + energies_theory1.append(results_theory1.energies_dict[l]) + energies_theory2.append(results_theory2.energies_dict[l]) else: energy = results_theory1.energies_dict[l] @@ -302,6 +318,8 @@ def create_ML_training_data(xyz_dir=None, dcd_trajectory=None, xyz_trajectory=No if Grad: if delta is True: gradient = results_theory2.gradients_dict[l] - results_theory1.gradients_dict[l] + gradients_theory1.append(results_theory1.gradients_dict[l]) + gradients_theory2.append(results_theory2.gradients_dict[l]) else: gradient = results_theory1.gradients_dict[l] gradients.append(gradient) @@ -364,6 +382,27 @@ def create_ML_training_data(xyz_dir=None, dcd_trajectory=None, xyz_trajectory=No gradients_file.write(f"{g[0]:10.7f} {g[1]:10.7f} {g[2]:10.7f}\n") gradients_file.close() + # When delta-learning (two theories), also write the individual theory + # energies and gradients to separate files alongside the delta data above. + if delta is True: + # Per-theory energy files + for theoryname, theory_energies in [("theory1", energies_theory1), + ("theory2", energies_theory2)]: + with open(f"train_data_{theoryname}.energies", "w") as f_energies: + for energy in theory_energies: + f_energies.write(f"{energy}\n") + + # Per-theory gradient files + if Grad: + for theoryname, theory_gradients in [("theory1", gradients_theory1), + ("theory2", gradients_theory2)]: + with open(f"train_data_{theoryname}.gradients", "w") as f_gradients: + for frag, grad in zip(fragments, theory_gradients): + f_gradients.write(f"{frag.numatoms}\n") + f_gradients.write(f"gradient {frag.label} \n") + for g in grad: + f_gradients.write(f"{g[0]:10.7f} {g[1]:10.7f} {g[2]:10.7f}\n") + print("\nNow writing data in MACE-format with energies in units of eV and forces in eV/Å") print("Fragments labels:",[frag.label for frag in fragments]) print("energies:", energies) @@ -407,6 +446,10 @@ def create_ML_training_data(xyz_dir=None, dcd_trajectory=None, xyz_trajectory=No print("All done! Files created:\ntrain_data.xyz\ntrain_data.energies\ntrain_data_mace.xyz") if Grad: print("train_data.gradients") + if delta is True: + print("train_data_theory1.energies\ntrain_data_theory2.energies") + if Grad: + print("train_data_theory1.gradients\ntrain_data_theory2.gradients") print("Number of user-chosen snapshots:", num_snapshots) print("Number of successfully generated datapoints:", len(energies))