Merge pull request #67 from mir-group/from_to_pmg

Added to / from PMG methods, with pymatgen as an optional import

flare/struc.py

@@ -4,6 +4,12 @@
 from flare.util import element_to_Z, NumpyEncoder
 from json import dumps
 
+try:
+    # Used for to_pmg_structure method
+    import pymatgen.core.structure as pmgstruc
+    _pmg_present = True
+except ImportError:
+    _pmg_present = False
 
 class Structure:
     """
@@ -172,6 +178,52 @@ def from_dict(dictionary):
 
         return struc
 
+    def to_pmg_structure(self):
+        """
+        Returns FLARE structure as a pymatgen structure.
+        :return: Pymatgen structure corresponding to current FLARE structure
+        """
+
+        if not _pmg_present:
+            raise ModuleNotFoundError("Pymatgen is not present. Please "
+                                      "install Pymatgen and try again")
+
+        site_properties = {'force:': self.forces, 'std': self.stds}
+
+        return pmgstruc.Structure(lattice=self.cell,
+                                  species=self.species_labels,
+                                  coords=self.positions,
+                                  coords_are_cartesian=True,
+                                  site_properties=site_properties
+                                  )
+
+    @staticmethod
+    def from_pmg_structure(structure):
+        """
+        Returns Pymatgen structure as FLARE structure.
+        
+        :param structure: Pymatgen structure
+        :return: FLARE Structure
+        """
+
+        cell = structure.lattice.matrix
+        species = [str(spec) for spec in structure.species]
+        positions = structure.cart_coords
+
+        new_struc = Structure(cell=cell,species=species,
+                              positions=positions)
+
+        site_props = structure.site_properties
+
+        if 'force' in site_props.keys():
+            forces = site_props['force']
+            new_struc.forces = [np.array(force) for force in forces]
+
+        if 'std' in site_props.keys():
+            stds = site_props['std']
+            new_struc.stds = [np.array(std) for std in stds]
+
+        return new_struc
 
 def get_unique_species(species):
     unique_species = []

requirements.txt

@@ -2,3 +2,4 @@ numpy
 scipy
 memory_profiler
 numba
+pymatgen

tests/test_struc.py

@@ -5,6 +5,12 @@
 from flare.struc import Structure
 from json import loads
 
+try:
+    import pymatgen.core.structure as pmgstruc
+    _test_pmg = True
+except ImportError:
+    _test_pmg = False
+
 
 def test_random_structure_setup():
     struct, forces = get_random_structure(cell=np.eye(3),
@@ -117,3 +123,33 @@ def test_rep_methods(varied_test_struc):
 
     thestr = str(varied_test_struc)
     assert 'Structure with 10 atoms of types {' in thestr
+
+@pytest.mark.skipif(not _test_pmg,reason='Pymatgen not present in available '
+                                        'packages.')
+def test_from_pmg_structure():
+
+    pmg_struc = pmgstruc.Structure(lattice= np.eye(3),
+                                   species=['H'],
+                                   coords=[[.25, .5, 0]],
+                                   site_properties={
+                                       'force': [np.array((1., 1.,  1.))],
+                                        'std':[np.array((1., 1., 1.))]},
+                                   coords_are_cartesian=True)
+
+    new_struc = Structure.from_pmg_structure(pmg_struc)
+
+    assert len(new_struc) == 1
+
+    assert np.equal(new_struc.positions, np.array([.25, .5, 0])).all()
+    assert new_struc.coded_species == [1]
+    assert new_struc.species_labels[0] == 'H'
+    assert np.equal(new_struc.forces, np.array([1., 1., 1.])).all()
+
+@pytest.mark.skipif(not _test_pmg,reason='Pymatgen not present in available '
+                                        'packages.')
+def test_to_pmg_structure(varied_test_struc):
+
+    new_struc = Structure.to_pmg_structure(varied_test_struc)
+    assert len(varied_test_struc) == len(varied_test_struc)
+    assert np.equal(new_struc.cart_coords, varied_test_struc.positions).all()
+    assert (new_struc.atomic_numbers == varied_test_struc.coded_species).all()