Clean up pumatgen imports so that it works with older version of

pymatgen.

m3gnet/graph/_converters.py

@@ -8,7 +8,7 @@
 import numpy as np
 import tensorflow as tf
 from ase import Atoms
-from pymatgen.core import Structure
+from pymatgen.core.structure import Structure
 
 from m3gnet.config import DataType
 from m3gnet.type import StructureOrMolecule

m3gnet/graph/_structure.py

@@ -6,7 +6,7 @@
 
 import numpy as np
 from ase import Atoms
-from pymatgen.core import Molecule, Structure
+from pymatgen.core.structure import Molecule, Structure
 from pymatgen.optimization.neighbors import find_points_in_spheres
 
 from m3gnet.config import DataType

m3gnet/graph/tests/test_batch.py

@@ -1,7 +1,7 @@
 import unittest
 
 import numpy as np
-from pymatgen.core import Lattice, Structure
+from pymatgen.core.structure import Lattice, Structure
 
 from m3gnet.graph import (
     MaterialGraph,

m3gnet/graph/tests/test_compute.py

@@ -2,7 +2,7 @@
 
 import numpy as np
 import tensorflow as tf
-from pymatgen.core import Lattice, Structure
+from pymatgen.core.structure import Lattice, Structure
 
 from m3gnet.graph import (
     Index,

m3gnet/graph/tests/test_converters.py

@@ -1,7 +1,7 @@
 import unittest
 
 import numpy as np
-from pymatgen.core import Lattice, Structure
+from pymatgen.core.structure import Lattice, Structure
 
 from m3gnet.graph import BaseGraphConverter, MaterialGraph, RadiusCutoffGraphConverter
 

m3gnet/graph/tests/test_types.py

@@ -2,7 +2,7 @@
 
 import numpy as np
 import tensorflow as tf
-from pymatgen.core import Lattice, Structure
+from pymatgen.core.structure import Lattice, Structure
 
 from m3gnet.graph import Index, MaterialGraph, RadiusCutoffGraphConverter
 

m3gnet/layers/_atom_ref.py

@@ -8,7 +8,8 @@
 import numpy as np
 import tensorflow as tf
 from ase import Atoms
-from pymatgen.core import Element, Molecule, Structure
+from pymatgen.core.periodic_table import Element
+from pymatgen.core.structure import Molecule, Structure
 
 from m3gnet.config import DataType
 from m3gnet.graph import Index

m3gnet/layers/tests/test_2_3_body.py

@@ -1,7 +1,7 @@
 import unittest
 
 import numpy as np
-from pymatgen.core import Lattice, Molecule, Structure
+from pymatgen.core.structure import Lattice, Molecule, Structure
 
 from m3gnet.graph import RadiusCutoffGraphConverter, tf_compute_distance_angle
 from m3gnet.layers import PairDistance, PairVector, SphericalBesselWithHarmonics

m3gnet/layers/tests/test_atom_ref.py

@@ -1,7 +1,7 @@
 import unittest
 
 import numpy as np
-from pymatgen.core import Lattice, Structure
+from pymatgen.core.structure import Lattice, Structure
 
 from m3gnet.layers import AtomRef
 

m3gnet/layers/tests/test_bond.py

@@ -1,7 +1,7 @@
 import unittest
 
 import numpy as np
-from pymatgen.core import Lattice, Molecule, Structure
+from pymatgen.core.structure import Lattice, Molecule, Structure
 
 from m3gnet.graph import Index, RadiusCutoffGraphConverter
 from m3gnet.layers import PairRadialBasisExpansion

m3gnet/layers/tests/test_core.py

@@ -2,7 +2,7 @@
 
 import numpy as np
 import tensorflow as tf
-from pymatgen.core import Lattice, Structure
+from pymatgen.core.structure import Lattice, Structure
 
 from m3gnet.graph import RadiusCutoffGraphConverter
 from m3gnet.layers import MLP, Embedding, GatedMLP, Pipe

m3gnet/layers/tests/test_gn.py

@@ -1,7 +1,7 @@
 import unittest
 
 import numpy as np
-from pymatgen.core import Lattice, Molecule, Structure
+from pymatgen.core.structure import Lattice, Molecule, Structure
 
 from m3gnet.graph import MaterialGraph, RadiusCutoffGraphConverter
 from m3gnet.layers import MLP

m3gnet/models/_base.py

@@ -8,7 +8,7 @@
 from typing import List, Union
 
 import tensorflow as tf
-from pymatgen.core import Molecule, Structure
+from pymatgen.core.structure import Molecule, Structure
 
 from m3gnet.config import DataType
 from m3gnet.graph import Index, MaterialGraph, assemble_material_graph

m3gnet/models/_dynamics.py

@@ -22,7 +22,7 @@
 from ase.optimize.mdmin import MDMin
 from ase.optimize.optimize import Optimizer
 from ase.optimize.sciopt import SciPyFminBFGS, SciPyFminCG
-from pymatgen.core import Molecule, Structure
+from pymatgen.core.structure import Molecule, Structure
 from pymatgen.io.ase import AseAtomsAdaptor
 
 from ._base import Potential

m3gnet/models/tests/test_model.py

@@ -3,7 +3,7 @@
 
 import numpy as np
 from monty.tempfile import ScratchDir
-from pymatgen.core import Lattice, Molecule, Structure
+from pymatgen.core.structure import Lattice, Molecule, Structure
 
 from m3gnet.models import M3GNet, MolecularDynamics, Potential, Relaxer
 

m3gnet/trainers/_potential.py

@@ -6,7 +6,7 @@
 
 import tensorflow as tf
 from ase import Atoms
-from pymatgen.core import Molecule, Structure
+from pymatgen.core.structure import Molecule, Structure
 
 from m3gnet.callbacks import ManualStop
 from m3gnet.graph import Index, MaterialGraph, MaterialGraphBatchEnergyForceStress

m3gnet/trainers/_property.py

@@ -10,7 +10,7 @@
 import numpy as np
 import tensorflow as tf
 from ase import Atoms
-from pymatgen.core import Molecule, Structure
+from pymatgen.core.structure import Molecule, Structure
 
 from m3gnet.callbacks import ManualStop
 from m3gnet.graph import MaterialGraph, MaterialGraphBatch

m3gnet/trainers/tests/test_trainer.py

@@ -5,7 +5,7 @@
 import numpy as np
 import tensorflow as tf
 from monty.tempfile import ScratchDir
-from pymatgen.core import Structure
+from pymatgen.core.structure import Structure
 
 from m3gnet.models import M3GNet, Potential
 from m3gnet.trainers import PotentialTrainer, Trainer

m3gnet/type.py

@@ -4,7 +4,7 @@
 
 import numpy as np
 from ase.atoms import Atoms
-from pymatgen.core import Molecule, Structure
+from pymatgen.core.structure import Molecule, Structure
 
 OptStrOrCallable = Optional[Union[str, Callable[..., Any]]]
 StructureOrMolecule = Union[Structure, Molecule, Atoms]

setup.py

@@ -33,7 +33,7 @@
     name="m3gnet",
     version=version,
     install_requires=[
-        "pymatgen>=2022.1.1",
+        "pymatgen",
         "tensorflow",
         "numpy",
         "monty",