From 9a2c3ceceafab1ce0314e3208d497173d7331860 Mon Sep 17 00:00:00 2001 From: Janosh Riebesell Date: Mon, 22 Jan 2024 16:11:43 +0100 Subject: [PATCH 1/3] lazy import pandas in Structure.as_dataframe() to improve startup speed --- pymatgen/core/structure.py | 12 +++++++----- tests/core/test_structure.py | 9 ++++++--- 2 files changed, 13 insertions(+), 8 deletions(-) diff --git a/pymatgen/core/structure.py b/pymatgen/core/structure.py index f4e57d7563f..a19866b9474 100644 --- a/pymatgen/core/structure.py +++ b/pymatgen/core/structure.py @@ -24,7 +24,6 @@ from typing import TYPE_CHECKING, Any, Callable, Literal, SupportsIndex, cast, get_args import numpy as np -import pandas as pd from monty.dev import deprecated from monty.io import zopen from monty.json import MSONable @@ -2625,16 +2624,19 @@ def as_dataframe(self): 0 (Si) 0.0 0.0 0.000000e+00 0.0 0.000000e+00 0.000000e+00 5 1 (Si) 0.0 0.0 1.000000e-7 0.0 -2.217138e-7 3.135509e-7 -5 """ - data = [] + # pandas lazy imported for speed (https://github.com/materialsproject/pymatgen/issues/3563) + import pandas as pd + + data: list[list[str | float]] = [] site_properties = self.site_properties prop_keys = list(site_properties) for site in self: row = [site.species, *site.frac_coords, *site.coords] - for k in prop_keys: - row.append(site.properties.get(k)) + for key in prop_keys: + row.append(site.properties.get(key)) data.append(row) - df = pd.DataFrame(data, columns=["Species", "a", "b", "c", "x", "y", "z", *prop_keys]) + df = pd.DataFrame(data, columns=["Species", *"abcxyz", *prop_keys]) df.attrs["Reduced Formula"] = self.composition.reduced_formula df.attrs["Lattice"] = self.lattice return df diff --git a/tests/core/test_structure.py b/tests/core/test_structure.py index 4d9dfbbf835..864cf460992 100644 --- a/tests/core/test_structure.py +++ b/tests/core/test_structure.py @@ -105,9 +105,12 @@ def test_get_orderings(self): assert sqs[0].formula == "Mn8 Fe8" def test_as_dataframe(self): - df = self.propertied_structure.as_dataframe() - assert df.attrs["Reduced Formula"] == self.propertied_structure.composition.reduced_formula - assert df.shape == (2, 8) + df_struct = self.propertied_structure.as_dataframe() + assert df_struct.attrs["Reduced Formula"] == self.propertied_structure.composition.reduced_formula + assert df_struct.shape == (2, 8) + assert list(df_struct) == ["Species", *"abcxyz", "magmom"] + assert list(df_struct["magmom"]) == [5, -5] + assert list(map(str, df_struct["Species"])) == ["Si1", "Si1"] def test_equal(self): struct = self.struct From c8ab49d031a6edcf98a6d1b52dfc4a3c58ac8c91 Mon Sep 17 00:00:00 2001 From: Janosh Riebesell Date: Mon, 22 Jan 2024 16:19:18 +0100 Subject: [PATCH 2/3] fix allclose removal in scipy 1.12 --- tests/electronic_structure/test_plotter.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/tests/electronic_structure/test_plotter.py b/tests/electronic_structure/test_plotter.py index 30462b5e5e7..f2d4bef66dd 100644 --- a/tests/electronic_structure/test_plotter.py +++ b/tests/electronic_structure/test_plotter.py @@ -7,7 +7,6 @@ import matplotlib.pyplot as plt import numpy as np -import scipy from matplotlib import rc from numpy.testing import assert_allclose from pytest import approx @@ -269,11 +268,11 @@ def test_bz_plot(self): ) def test_fold_point(self): - assert scipy.allclose( + assert_allclose( fold_point([0.0, -0.5, 0.5], lattice=self.rec_latt), self.rec_latt.get_cartesian_coords([0.0, 0.5, 0.5]), ) - assert scipy.allclose( + assert_allclose( fold_point([0.1, -0.6, 0.2], lattice=self.rec_latt), self.rec_latt.get_cartesian_coords([0.1, 0.4, 0.2]), ) From 86e86f3fa3b884f087bad471a6c91abacee4e974 Mon Sep 17 00:00:00 2001 From: Janosh Riebesell Date: Mon, 22 Jan 2024 16:57:49 +0100 Subject: [PATCH 3/3] increase atol=0 -> 1e-8 --- tests/electronic_structure/test_plotter.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/electronic_structure/test_plotter.py b/tests/electronic_structure/test_plotter.py index f2d4bef66dd..74e36766926 100644 --- a/tests/electronic_structure/test_plotter.py +++ b/tests/electronic_structure/test_plotter.py @@ -271,10 +271,12 @@ def test_fold_point(self): assert_allclose( fold_point([0.0, -0.5, 0.5], lattice=self.rec_latt), self.rec_latt.get_cartesian_coords([0.0, 0.5, 0.5]), + atol=1e-8, ) assert_allclose( fold_point([0.1, -0.6, 0.2], lattice=self.rec_latt), self.rec_latt.get_cartesian_coords([0.1, 0.4, 0.2]), + atol=1e-8, )