From c87f67cd41202915a5eeda93cb1be74997fbfb04 Mon Sep 17 00:00:00 2001
From: CalCraven <nicholas.c.craven@vanderbilt.edu>
Date: Mon, 5 Dec 2022 11:07:03 -0600
Subject: [PATCH] Patches for element reading using CIF reader and other
 inference of elements form compound_dicts

---
 mbuild/lattice.py            | 21 ++++++++++++++++++++-
 mbuild/tests/test_cif.py     | 11 +++++++++++
 mbuild/tests/test_lattice.py | 30 ++++++++++++++++++++++++++++++
 3 files changed, 61 insertions(+), 1 deletion(-)

diff --git a/mbuild/lattice.py b/mbuild/lattice.py
index c6f860bfc..fcc9ca9d3 100644
--- a/mbuild/lattice.py
+++ b/mbuild/lattice.py
@@ -1,9 +1,12 @@
 """mBuild lattice module for working with crystalline systems."""
 import itertools as it
 import pathlib
+import warnings
 from collections import defaultdict
 
 import numpy as np
+from ele.element import element_from_name, element_from_symbol
+from ele.exceptions import ElementError
 
 import mbuild as mb
 from mbuild.utils.io import import_
@@ -631,9 +634,19 @@ def populate(self, compound_dict=None, x=1, y=1, z=1):
         ret_lattice = mb.Compound()
 
         # Create (clone) a mb.Compound for the newly generate positions
+        elementsSet = set()
         if compound_dict is None:
             for key_id, all_pos in cell.items():
-                particle = mb.Compound(name=key_id, pos=[0, 0, 0])
+                for idElement in [element_from_symbol, element_from_name]:
+                    try:  # populate element info if it's there
+                        element = idElement(key_id)
+                        elementsSet.add(element)
+                        break
+                    except ElementError:
+                        element = None
+                particle = mb.Compound(
+                    name=key_id, pos=[0, 0, 0], element=element
+                )
                 for pos in all_pos:
                     particle_to_add = mb.clone(particle)
                     particle_to_add.translate_to(list(pos))
@@ -654,6 +667,12 @@ def populate(self, compound_dict=None, x=1, y=1, z=1):
                             key_id, err_type
                         )
                     )
+        # Raise warnings about assumed elements
+        for element in elementsSet:
+            warnings.warn(
+                f"Element assumed from cif file to be {element}.", UserWarning
+            )
+
         # Create mbuild.box
         ret_lattice.box = mb.Box(
             lengths=[a * x, b * y, c * z], angles=self.angles
diff --git a/mbuild/tests/test_cif.py b/mbuild/tests/test_cif.py
index 4dd0a0f28..e1cd7e124 100644
--- a/mbuild/tests/test_cif.py
+++ b/mbuild/tests/test_cif.py
@@ -170,3 +170,14 @@ def test_cif_triclinic_box_properties(self):
         assert np.all(np.isclose(manual_lengths, list(periodic_box.lengths)))
         assert np.all(np.isclose(manual_angles, list(periodic_box.angles)))
         assert len(periodic_boxed_molecule.children) == manual_num_atoms
+        assert None not in list(
+            map(lambda x: x.element, periodic_boxed_molecule.particles())
+        )
+
+    def test_cif_raise_warnings(self):
+        with pytest.warns(
+            UserWarning,
+            match=r"Element assumed from cif file to be Element: silicon, symbol: Si, atomic number: 14, mass: 28.085.",
+        ):
+            lattice_cif = load_cif(file_or_path=get_fn("ETV_triclinic.cif"))
+            periodic_boxed_molecule = lattice_cif.populate(x=1, y=1, z=1)
diff --git a/mbuild/tests/test_lattice.py b/mbuild/tests/test_lattice.py
index b669e4473..d3e9acde5 100644
--- a/mbuild/tests/test_lattice.py
+++ b/mbuild/tests/test_lattice.py
@@ -310,3 +310,33 @@ def test_get_box_non_rectangular(self):
         assert isinstance(mylat.box, mb.Box)
         np.testing.assert_allclose([90, 90, 120], mylat.box.angles)
         np.testing.assert_allclose(expected_lengths, mylat.box.lengths)
+
+    def test_populate_with_element_symbol(self):
+        lattice = mb.Lattice(
+            lattice_spacing=[0.5, 0.5, 1],
+            angles=[90, 90, 120],
+            lattice_points={"O": [[0, 0, 0]]},
+        )
+        cpd_lat = lattice.populate(x=1, y=1, z=1)
+        for part in cpd_lat:
+            assert part.element.name == "oxygen"
+
+    def test_populate_with_element_name(self):
+        lattice = mb.Lattice(
+            lattice_spacing=[0.5, 0.5, 1],
+            angles=[90, 90, 120],
+            lattice_points={"Oxygen": [[0, 0, 0]]},
+        )
+        cpd_lat = lattice.populate(x=1, y=1, z=1)
+        for part in cpd_lat:
+            assert part.element.name == "oxygen"
+
+    def test_populate_no_element(self):
+        lattice = mb.Lattice(
+            lattice_spacing=[0.5, 0.5, 1],
+            angles=[90, 90, 120],
+            lattice_points={"A": [[0, 0, 0]]},
+        )
+        cpd_lat = lattice.populate(x=1, y=1, z=1)
+        for part in cpd_lat:
+            assert part.element is None