adding ASE compatibility to Crystal class

py4DSTEM/io/google_drive_downloader.py

@@ -123,6 +123,30 @@
         "downsample_Si_SiGe_analysis_braggdisks_cal.h5",
         "1bYgDdAlnWHyFmY-SwN3KVpMutWBI5MhP",
     ),
+    "test_Si_cif": (
+        "Si.cif",
+        "1b5ZonO0Upcp6eoiE-UEzGe1ymSj4QBaO"
+    ),
+    "test_Si_cif_sym": (
+        "Si-sym.cif",
+        "1Q26anzuq9FRgjNPGHUyrQ3d1pmP1B--T"
+    ),
+    "test_Si_prismatic": (
+        "Si.prismatic",
+        "1h_tnmygDGin3rMMBYtV2-VCJbO9_4Q_F"
+    ),
+    "test_NaCl_cif": (
+        "NaCl.cif",
+        "1nfUTyHTwYM6S5mCxIgwzrtN9gS0ReOYb"
+    ),
+    "test_NaCl_prismatic": (
+        "NaCl.prismatic",
+        "1_fWx6C3PkT47fBJufYjM3DN7eeMIQqYj"
+    ),
+    "test_DRX_cif": (
+        "DRX.cif",
+        "1hiGqFW6BJMFl718h63jQybpWfY-lNmsr"
+    ),
 }
 
 # collections of files
@@ -163,6 +187,14 @@
     ),
     "test_braggvectors": ("Au_sim",),
     "strain": ("test_strain",),
+    "test_atomic_coords": (
+        "test_Si_cif",
+        "test_Si_cif_sym",
+        "test_Si_prismatic",
+        "test_NaCl_cif",
+        "test_NaCl_prismatic",
+        "test_DRX_cif",
+    ),
 }
 
 

py4DSTEM/process/diffraction/crystal.py

@@ -14,12 +14,19 @@
 from py4DSTEM.process.diffraction.crystal_viz import plot_ring_pattern
 from py4DSTEM.process.diffraction.utils import Orientation, calc_1D_profile
 
+from importlib.util import find_spec
+
 try:
     from pymatgen.symmetry.analyzer import SpacegroupAnalyzer
     from pymatgen.core.structure import Structure
 except ImportError:
     pass
 
+try:
+    from ase.io import read
+except ImportError:
+    pass
+
 
 class Crystal:
     """
@@ -147,6 +154,9 @@ def __init__(
         # Calculate lattice parameters
         self.calculate_lattice()
 
+        # Set occupancy attribute to None
+        self.occupancy = None
+
     def calculate_lattice(self):
         if not hasattr(self, "lat_real"):
             # calculate unit cell lattice vectors
@@ -265,27 +275,27 @@ def get_strained_crystal(
         else:
             return crystal_strained
 
-    def from_CIF(CIF, conventional_standard_structure=True):
-        """
-        Create a Crystal object from a CIF file, using pymatgen to import the CIF
+    # def from_CIF(CIF, conventional_standard_structure=True):
+    #     """
+    #     Create a Crystal object from a CIF file, using pymatgen to import the CIF
 
-        Note that pymatgen typically prefers to return primitive unit cells,
-        which can be overridden by setting conventional_standard_structure=True.
+    #     Note that pymatgen typically prefers to return primitive unit cells,
+    #     which can be overridden by setting conventional_standard_structure=True.
 
-        Args:
-            CIF: (str or Path) path to the CIF File
-            conventional_standard_structure: (bool) if True, conventional standard unit cell will be returned
-                instead of the primitive unit cell pymatgen typically returns
-        """
-        from pymatgen.io.cif import CifParser
+    #     Args:
+    #         CIF: (str or Path) path to the CIF File
+    #         conventional_standard_structure: (bool) if True, conventional standard unit cell will be returned
+    #             instead of the primitive unit cell pymatgen typically returns
+    #     """
+    #     from pymatgen.io.cif import CifParser
 
-        parser = CifParser(CIF)
+    #     parser = CifParser(CIF)
 
-        structure = parser.get_structures()[0]
+    #     structure = parser.get_structures()[0]
 
-        return Crystal.from_pymatgen_structure(
-            structure, conventional_standard_structure=conventional_standard_structure
-        )
+    #     return Crystal.from_pymatgen_structure(
+    #         structure, conventional_standard_structure=conventional_standard_structure
+    #     )
 
     def from_pymatgen_structure(
         structure=None,
@@ -400,6 +410,149 @@ def from_pymatgen_structure(
 
         return Crystal(positions, numbers, cell)
 
+    def from_ase(
+        atoms,
+    ):
+        """
+        Create a py4DSTEM Crystal object from an ASE atoms object
+
+        Args:
+            atoms (ase.Atoms): an ASE atoms object
+
+        """
+
+        xtal = Crystal(
+            positions=atoms.get_scaled_positions(),  # fractional coords
+            numbers=atoms.numbers,
+            cell=atoms.cell.array,
+        )
+
+        # get occupancies
+        # ASE seems to have different ways of storing occupancies
+        # if ASE object created from prismatic file
+        if "occupancies" in atoms.arrays.keys():
+            # np.array with length number of atoms
+            xtal.occupancy = atoms.arrays["occupancies"]
+            # if created from cif file
+        elif "occupancy" in atoms.info.keys():
+            # python dict with occupancy per site
+            xtal.occupancy = atoms.info["occupancy"]
+        # TODO add in elif statement if other ways appear
+        else:
+            print(Warning("Could not find occupancies of crystal"))
+        return xtal
+
+    def from_prismatic(filepath):
+        """
+        Create a py4DSTEM Crystal object from an prismatic style xyz co-ordinate file
+
+        Args:
+            filepath (str|Pathlib.Path): path to the prismatic format xyz file
+
+        """
+
+        # check if ase is installed
+        if find_spec("ase") is None:
+            raise ImportWarning("Could not import ASE, please install, and try again")
+        else:
+            from ase.io import read
+
+        atoms = read(filepath, format="prismatic")
+
+        # create the crystal object
+        xtal = Crystal(
+            positions=atoms.get_scaled_positions(),  # fractional coords
+            numbers=atoms.numbers,
+            cell=atoms.cell.array,
+        )
+        # add occupancies
+        # It should be this one but keeping other method in case
+        if "occupancies" in atoms.arrays.keys():
+            # np.array with length number of atoms
+            xtal.occupancy = atoms.arrays["occupancies"]
+            # if created from cif file
+        elif "occupancy" in atoms.info.keys():
+            # python dict with occupancy per site
+            xtal.occupancy = atoms.info["occupancy"]
+        # TODO add in elif statement if other ways appear
+        else:
+            print(Warning("Could not find occupancies of crystal"))
+        return xtal
+
+    def from_cif(
+        filepath,
+    ):
+        """
+        Create a py4DSTEM Crystal object from a cif file using ase.io.read function
+
+        Args:
+            filepath (str|Pathlib.Path): path to the file
+        """
+
+        # check if ase is installed
+        if find_spec("ase") is None:
+            raise ImportWarning(
+                "Could not import ASE, please install, restart and try again"
+            )
+        else:
+            from ase.io import read
+        # try loading the file using ase read and get required properties
+        atoms = read(filepath, format="cif")
+        xtal = Crystal(
+            positions=atoms.get_scaled_positions(),  # fractional coords
+            numbers=atoms.numbers,
+            cell=atoms.cell.array,
+        )
+
+        # add occupancies
+        # It should be this one but keeping other method in case
+        if "occupancies" in atoms.arrays.keys():
+            # np.array with length number of atoms
+            xtal.occupancy = atoms.arrays["occupancies"]
+            # if created from cif file
+        elif "occupancy" in atoms.info.keys():
+            # python dict with occupancy per site
+            xtal.occupancy = atoms.info["occupancy"]
+        # TODO add in elif statement if other ways appear
+        # TODO when nested python dict, it needs to be unpacked. into numpy array like from_prismatic
+        else:
+            print(Warning("Could not find occupancies of crystal"))
+        return xtal
+
+    # def from_generic_file(filepath, **kwargs):
+    #     """
+    #     Create a py4DSTEM Crystal from a wide range of generic file types using
+    #     `ase.io.read`,  kwargs are passed to `ase.io.read` function. For more details
+    #     and potentially compatible filetypes please see https://wiki.fysik.dtu.dk/ase/ase/io/io.html.
+    #     Note this has not been tested extensively. The loaded file must have these three properties:
+    #         .get_scaled_positions()
+    #         .numbers
+    #         .cell.array,
+
+    #     Args:
+    #         filepath (str|Pathlib.Path): path to the file
+    #         kwargs: key word arguments to be passed to `ase.io.read`
+
+    #     """
+
+    #     # check if ase is installed
+    #     if find_spec("ase") is None:
+    #         raise ImportWarning(
+    #             "Could not import ASE, please install, restart and try again"
+    #         )
+    #     else:
+    #         from ase.io import read
+    #     # try loading the file using ase read and get required properties
+    #     try:
+    #         atoms = read(filepath, **kwargs)
+    #         return Crystal(
+    #             positions=atoms.get_scaled_positions(),  # fractional coords
+    #             numbers=atoms.numbers,
+    #             cell=atoms.cell.array,
+    #         )
+    #     except Exception as e:
+    #         raise e
+
     def from_unitcell_parameters(
         latt_params,
         elements,

setup.py

@@ -40,6 +40,7 @@
         "emdfile >= 0.0.13",
         "mpire >= 2.7.1",
         "threadpoolctl >= 3.1.0",
+        "ase >= 3.22.0",
     ],
     extras_require={
         "ipyparallel": ["ipyparallel >= 6.2.4", "dill >= 0.3.3"],

test/gettestdata.py

@@ -52,7 +52,7 @@
 if args.data == "tutorials":
     data = ["tutorials"]
 elif args.data == "io":
-    data = ["test_io", "test_arina"]
+    data = ["test_io", "test_arina", "test_atomic_coords"]
 elif args.data == "basic":
     data = ["small_datacube"]
 elif args.data == "strain":