Skip to content

Commit 6e9a033

Browse files
authored
Lazy import pandas in Structure.as_dataframe() to improve startup speed (#3568)
* lazy import pandas in Structure.as_dataframe() to improve startup speed * fix allclose removal in scipy 1.12
1 parent 20475b6 commit 6e9a033

File tree

3 files changed

+17
-11
lines changed

3 files changed

+17
-11
lines changed

pymatgen/core/structure.py

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,6 @@
2424
from typing import TYPE_CHECKING, Any, Callable, Literal, SupportsIndex, cast, get_args
2525

2626
import numpy as np
27-
import pandas as pd
2827
from monty.dev import deprecated
2928
from monty.io import zopen
3029
from monty.json import MSONable
@@ -2625,16 +2624,19 @@ def as_dataframe(self):
26252624
0 (Si) 0.0 0.0 0.000000e+00 0.0 0.000000e+00 0.000000e+00 5
26262625
1 (Si) 0.0 0.0 1.000000e-7 0.0 -2.217138e-7 3.135509e-7 -5
26272626
"""
2628-
data = []
2627+
# pandas lazy imported for speed (https://github.com/materialsproject/pymatgen/issues/3563)
2628+
import pandas as pd
2629+
2630+
data: list[list[str | float]] = []
26292631
site_properties = self.site_properties
26302632
prop_keys = list(site_properties)
26312633
for site in self:
26322634
row = [site.species, *site.frac_coords, *site.coords]
2633-
for k in prop_keys:
2634-
row.append(site.properties.get(k))
2635+
for key in prop_keys:
2636+
row.append(site.properties.get(key))
26352637
data.append(row)
26362638

2637-
df = pd.DataFrame(data, columns=["Species", "a", "b", "c", "x", "y", "z", *prop_keys])
2639+
df = pd.DataFrame(data, columns=["Species", *"abcxyz", *prop_keys])
26382640
df.attrs["Reduced Formula"] = self.composition.reduced_formula
26392641
df.attrs["Lattice"] = self.lattice
26402642
return df

tests/core/test_structure.py

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -105,9 +105,12 @@ def test_get_orderings(self):
105105
assert sqs[0].formula == "Mn8 Fe8"
106106

107107
def test_as_dataframe(self):
108-
df = self.propertied_structure.as_dataframe()
109-
assert df.attrs["Reduced Formula"] == self.propertied_structure.composition.reduced_formula
110-
assert df.shape == (2, 8)
108+
df_struct = self.propertied_structure.as_dataframe()
109+
assert df_struct.attrs["Reduced Formula"] == self.propertied_structure.composition.reduced_formula
110+
assert df_struct.shape == (2, 8)
111+
assert list(df_struct) == ["Species", *"abcxyz", "magmom"]
112+
assert list(df_struct["magmom"]) == [5, -5]
113+
assert list(map(str, df_struct["Species"])) == ["Si1", "Si1"]
111114

112115
def test_equal(self):
113116
struct = self.struct

tests/electronic_structure/test_plotter.py

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,6 @@
77

88
import matplotlib.pyplot as plt
99
import numpy as np
10-
import scipy
1110
from matplotlib import rc
1211
from numpy.testing import assert_allclose
1312
from pytest import approx
@@ -269,13 +268,15 @@ def test_bz_plot(self):
269268
)
270269

271270
def test_fold_point(self):
272-
assert scipy.allclose(
271+
assert_allclose(
273272
fold_point([0.0, -0.5, 0.5], lattice=self.rec_latt),
274273
self.rec_latt.get_cartesian_coords([0.0, 0.5, 0.5]),
274+
atol=1e-8,
275275
)
276-
assert scipy.allclose(
276+
assert_allclose(
277277
fold_point([0.1, -0.6, 0.2], lattice=self.rec_latt),
278278
self.rec_latt.get_cartesian_coords([0.1, 0.4, 0.2]),
279+
atol=1e-8,
279280
)
280281

281282

0 commit comments

Comments
 (0)