diff --git a/PQAnalysis/io/nep/nep_writer.py b/PQAnalysis/io/nep/nep_writer.py index 4289054c..baf9a856 100644 --- a/PQAnalysis/io/nep/nep_writer.py +++ b/PQAnalysis/io/nep/nep_writer.py @@ -467,10 +467,14 @@ def _setup_frame_splitting_for_training( exception=NEPError ) - sum_frames = n_train + n_test + n_validation + if total_ratios is not None: + sum_frames = n_train + n_test + n_validation - self.test_ratio = n_test / sum_frames - self.validation_ratio = n_validation / sum_frames + self.test_ratio = n_test / sum_frames + self.validation_ratio = n_validation / sum_frames + else: + self.test_ratio = float(test_ratio) + self.validation_ratio = 0.0 if self.test_ratio > 1.0: self.logger.error( @@ -816,7 +820,7 @@ def write_from_trajectory( self.open() for frame in trajectory: self.write_from_atomic_system( - frame, use_forces, use_stress, use_virial + frame, self.file, use_forces, use_stress, use_virial ) self.close() diff --git a/tests/io/test_nepWriter.py b/tests/io/test_nepWriter.py new file mode 100644 index 00000000..7a16297e --- /dev/null +++ b/tests/io/test_nepWriter.py @@ -0,0 +1,103 @@ +import os + +import numpy as np + +from . import pytestmark + +from PQAnalysis.io.nep.nep_writer import NEPWriter +from PQAnalysis.traj import Trajectory +from PQAnalysis.core import Atom, Cell +from PQAnalysis.atomic_system import AtomicSystem + +INFO_FILE_CONTENT = """\ +----------------------------------------------------------------------------------------- +| PQ info file | +----------------------------------------------------------------------------------------- +| SIMULATION-TIME 5.00000 ps TEMPERATURE 294.45308 K | +| PRESSURE 2875.71714 bar E(TOT) -186008.40141 kcal/mol | +| E(QM) -186197.10854 kcal/mol N(QM-ATOMS) 0.00000 - | +| E(KIN) 188.70714 kcal/mol E(INTRA) 0.00000 kcal/mol | +| VOLUME 4310.41014 A^3 DENSITY 0.83455 g/cm^3 | +| MOMENTUM 2.5e-10 amuA/fs LOOPTIME 0.90443 s | +----------------------------------------------------------------------------------------- +""" + + + +class TestNEPWriter: + + def test_write_from_files_with_test_ratio(self, tmpdir): + xyz_frame = ( + "2 10.0 10.0 10.0 90.0 90.0 90.0\n#\n" + "C 0.0 0.0 0.0\nH 1.0 1.0 1.0\n" + ) + with open("md.xyz", "w", encoding="utf-8") as file: + file.write(xyz_frame * 4) + + with open("md.info", "w", encoding="utf-8") as file: + file.write(INFO_FILE_CONTENT) + + with open("md.instant_en", "w", encoding="utf-8") as file: + for i in range(1, 5): + file.write( + f"{i} 298.0 23000.0 -185898.0 -18590{i}.0 0.0 5.6 0.0 " + "3729.1 0.83 1.1e-16 1.0\n" + ) + + writer = NEPWriter("out.xyz") + writer.write_from_files(["md"], test_ratio=0.25) + + assert np.isclose(writer.test_ratio, 0.25) + assert np.isclose(writer.validation_ratio, 0.0) + assert writer.n_train_frames == 3 + assert writer.n_test_frames == 1 + assert os.path.exists("out.xyz_train") + assert os.path.exists("out.xyz_test") + + def test_write_from_files_with_total_ratios(self, tmpdir): + xyz_frame = ( + "2 10.0 10.0 10.0 90.0 90.0 90.0\n#\n" + "C 0.0 0.0 0.0\nH 1.0 1.0 1.0\n" + ) + with open("md.xyz", "w", encoding="utf-8") as file: + file.write(xyz_frame * 4) + + with open("md.info", "w", encoding="utf-8") as file: + file.write(INFO_FILE_CONTENT) + + with open("md.instant_en", "w", encoding="utf-8") as file: + for i in range(1, 5): + file.write( + f"{i} 298.0 23000.0 -185898.0 -18590{i}.0 0.0 5.6 0.0 " + "3729.1 0.83 1.1e-16 1.0\n" + ) + + writer = NEPWriter("out.xyz") + writer.write_from_files(["md"], total_ratios="3:1") + + assert np.isclose(writer.test_ratio, 0.25) + assert np.isclose(writer.validation_ratio, 0.0) + assert writer.n_train_frames == 3 + assert writer.n_test_frames == 1 + + def test_write_from_trajectory(self, tmpdir): + system = AtomicSystem( + atoms=[Atom("C"), Atom("H")], + pos=np.array([[0.0, 0.0, 0.0], [1.0, 1.0, 1.0]]), + cell=Cell(10.0, 10.0, 10.0), + ) + system.energy = -1.5 + trajectory = Trajectory([system, system]) + + writer = NEPWriter("traj_out.xyz") + writer.write_from_trajectory(trajectory) + + with open("traj_out.xyz", "r", encoding="utf-8") as file: + lines = file.read().splitlines() + + assert len(lines) == 8 + assert lines[0] == "2" + assert lines[4] == "2" + assert lines[1].startswith("energy=") + assert lines[2].startswith("C") + assert lines[3].startswith("H")