Skip to content

Commit 8cbc716

Browse files
authored
Merge pull request #1206 from esoteric-ephemera/parqtraj
New trajectory class with parquet rountrip functionality
2 parents c0b077c + 609fd24 commit 8cbc716

File tree

11 files changed

+535
-17
lines changed

11 files changed

+535
-17
lines changed

emmet-core/emmet/core/trajectory.py

Lines changed: 449 additions & 0 deletions
Large diffs are not rendered by default.

emmet-core/requirements/deployment.txt

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -74,7 +74,7 @@ pydantic-core==2.33.0
7474
# via pydantic
7575
pydantic-settings==2.8.1
7676
# via emmet-core (emmet/emmet-core/setup.py)
77-
pymatgen==2025.1.9
77+
pymatgen==2025.3.10
7878
# via emmet-core (emmet/emmet-core/setup.py)
7979
pyparsing==3.2.3
8080
# via matplotlib

emmet-core/requirements/ubuntu-latest_py3.10.txt

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -69,7 +69,7 @@ pydantic-core==2.33.0
6969
# via pydantic
7070
pydantic-settings==2.8.1
7171
# via emmet-core (setup.py)
72-
pymatgen==2025.1.9
72+
pymatgen==2025.3.10
7373
# via emmet-core (setup.py)
7474
pyparsing==3.2.3
7575
# via matplotlib

emmet-core/requirements/ubuntu-latest_py3.10_extras.txt

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -78,8 +78,6 @@ dnspython==2.7.0
7878
# via
7979
# maggma
8080
# pymongo
81-
emmet-core==0.84.6rc6
82-
# via mp-api
8381
exceptiongroup==1.2.2
8482
# via pytest
8583
filelock==3.18.0
@@ -425,6 +423,9 @@ psutil==7.0.0
425423
# dgl
426424
pubchempy==1.0.4
427425
# via robocrys
426+
pyarrow==19.0.1
427+
# via
428+
# emmet-core (setup.py)
428429
pybtex==0.24.0
429430
# via
430431
# emmet-core
@@ -459,7 +460,7 @@ pyflakes==3.3.1
459460
# via flake8
460461
pygments==2.19.1
461462
# via mkdocs-material
462-
pymatgen==2025.1.9
463+
pymatgen==2025.3.10
463464
# via
464465
# chgnet
465466
# emmet-core

emmet-core/requirements/ubuntu-latest_py3.11.txt

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -69,7 +69,7 @@ pydantic-core==2.33.0
6969
# via pydantic
7070
pydantic-settings==2.8.1
7171
# via emmet-core (setup.py)
72-
pymatgen==2025.1.9
72+
pymatgen==2025.3.10
7373
# via emmet-core (setup.py)
7474
pyparsing==3.2.3
7575
# via matplotlib

emmet-core/requirements/ubuntu-latest_py3.11_extras.txt

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -76,8 +76,6 @@ dnspython==2.7.0
7676
# via
7777
# maggma
7878
# pymongo
79-
emmet-core==0.84.6rc6
80-
# via mp-api
8179
filelock==3.18.0
8280
# via
8381
# mdanalysis
@@ -421,6 +419,9 @@ psutil==7.0.0
421419
# dgl
422420
pubchempy==1.0.4
423421
# via robocrys
422+
pyarrow==19.0.1
423+
# via
424+
# emmet-core (setup.py)
424425
pybtex==0.24.0
425426
# via
426427
# emmet-core
@@ -455,7 +456,7 @@ pyflakes==3.3.1
455456
# via flake8
456457
pygments==2.19.1
457458
# via mkdocs-material
458-
pymatgen==2025.1.9
459+
pymatgen==2025.3.10
459460
# via
460461
# chgnet
461462
# emmet-core

emmet-core/setup.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@
2828
},
2929
include_package_data=True,
3030
install_requires=[
31-
"pymatgen>=2024.6.10,<=2025.1.9", # TODO: release upper version pin when bandgap parsing issues sorted
31+
"pymatgen>=2024.6.10",
3232
"monty>=2024.2.2",
3333
"pydantic>=2.0",
3434
"pydantic-settings>=2.0",
@@ -45,6 +45,7 @@
4545
"pymatgen-analysis-alloys>=0.0.6",
4646
"solvation-analysis>=0.4.1",
4747
"MDAnalysis>=2.7.0",
48+
"pyarrow",
4849
],
4950
"ml": ["chgnet", "matgl", "dgl<=2.1"],
5051
"test": [

emmet-core/tests/test_electronic_structure.py

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -85,15 +85,15 @@ def test_from_bsdos_2(bandstructure_fs, dos_fs):
8585
)
8686

8787
assert str(es_doc.material_id) == "mp-25375"
88-
assert es_doc.band_gap == 0.0
89-
assert es_doc.cbm == 2.7102
90-
assert es_doc.vbm == 2.9396
91-
assert es_doc.efermi == 2.75448867
88+
assert es_doc.band_gap == pytest.approx(0.0)
89+
assert es_doc.cbm == pytest.approx(2.75448867)
90+
assert es_doc.vbm == pytest.approx(2.75448867)
91+
assert es_doc.efermi == pytest.approx(2.75448867)
9292
assert es_doc.is_gap_direct is False
9393
assert es_doc.is_metal is True
9494
assert str(es_doc.magnetic_ordering) == "Ordering.NM"
9595

9696
assert str(es_doc.bandstructure.setyawan_curtarolo.task_id) == "mp-1612487"
97-
assert es_doc.bandstructure.setyawan_curtarolo.band_gap == 1.9916
98-
assert es_doc.bandstructure.setyawan_curtarolo.efermi == 2.49084067
99-
assert es_doc.bandstructure.setyawan_curtarolo.nbands == 64.0
97+
assert es_doc.bandstructure.setyawan_curtarolo.band_gap == pytest.approx(1.9916)
98+
assert es_doc.bandstructure.setyawan_curtarolo.efermi == pytest.approx(2.49084067)
99+
assert es_doc.bandstructure.setyawan_curtarolo.nbands == pytest.approx(64.0)

emmet-core/tests/test_trajectory.py

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,66 @@
1+
import numpy as np
2+
from pathlib import Path
3+
from pytest import approx
4+
5+
from pymatgen.core.trajectory import Trajectory as PmgTraj
6+
from pymatgen.util.testing import PymatgenTest
7+
8+
from emmet.core.tasks import TaskDoc
9+
from emmet.core.trajectory import Trajectory
10+
11+
12+
class TestTrajectory(PymatgenTest):
13+
def setUp(self):
14+
self.test_dir = (
15+
Path(__file__).parent.parent.parent.joinpath("test_files").resolve()
16+
)
17+
self.task_doc = TaskDoc.from_directory(
18+
self.test_dir / "vasp" / "Si_old_double_relax"
19+
)
20+
21+
def test_task_doc(self):
22+
traj = Trajectory.from_task_doc(self.task_doc)
23+
assert traj.num_ionic_steps == sum(
24+
len(cr.output.ionic_steps) for cr in self.task_doc.calcs_reversed
25+
)
26+
istep = -1
27+
for cr in self.task_doc.calcs_reversed:
28+
for ionic_step in cr.output.ionic_steps:
29+
assert traj.energy[istep] == approx(ionic_step.e_0_energy)
30+
istep -= 1
31+
32+
def test_parquet(self):
33+
parqet_file = "test.parquet"
34+
35+
traj = Trajectory.from_task_doc(self.task_doc)
36+
traj.to(parqet_file, compression="GZIP")
37+
38+
new_traj = Trajectory.from_parquet(parqet_file)
39+
assert hash(new_traj) == hash(traj)
40+
41+
def test_pmg(self):
42+
traj = Trajectory.from_task_doc(self.task_doc)
43+
pmg_traj = traj.to(fmt="PMG")
44+
assert isinstance(pmg_traj, PmgTraj)
45+
assert len(pmg_traj) == traj.num_ionic_steps
46+
assert len(pmg_traj.frame_properties) == traj.num_ionic_steps
47+
48+
# some float-rounding that happens here on round trip
49+
roundtrip = Trajectory.from_pmg(pmg_traj)
50+
51+
for k in Trajectory.model_fields:
52+
if (
53+
k
54+
in (
55+
"num_ionic_steps",
56+
"electronic_steps",
57+
)
58+
or getattr(traj, k, None) is None
59+
):
60+
# Can't easily test for equality of electronic steps without bespoke code
61+
continue
62+
63+
assert all(
64+
np.all(np.abs(np.array(new_val) - np.array(getattr(traj, k)[i])) < 1e-6)
65+
for i, new_val in enumerate(getattr(roundtrip, k))
66+
)
10.9 KB
Binary file not shown.

0 commit comments

Comments
 (0)