Merge pull request #22 from LaurentRDC/develop

Develop

.gitignore

@@ -8,10 +8,13 @@ __pycache__/
 
 # Visual studio cache
 *.vs/
+*.vscode/
 
 # Protein DataBank files and cache folder
 *.ent
-pdb_cache/
+
+# Structure file caches
+*_cache/
 
 # Jupyter notebooks
 notebooks/

docs/source/tutorials/simulation.rst

@@ -12,6 +12,7 @@ Contents
 ========
 
 * :ref:`powdersim`
+* :ref:`electrostatic`
 
 .. _powdersim:
 
@@ -51,4 +52,55 @@ After plot formatting:
 	plt.ylabel('Diffracted intensity (A.u.)')
 	plt.title('Polycrystalline graphite diffraction')
 
+.. _electrostatic:
+
+Electrostatic Potential Simulation
+==================================
+The scattering potential of electrons is the crystal electrostatic potential; hence
+computing such potential is a useful tool.
+
+To compute the electrostatic potential for an infinite crystal on an arbitrary 3D mesh,
+take a look at :func:`electrostatic`::
+
+    from skued import Crystal
+    from skued.simulation import electrostatic
+
+    graphite = Crystal.from_database('C')
+
+    extent = np.linspace(-10, 10, 256)
+    xx, yy, zz = np.meshgrid(extent, extent, extent)
+    potential = electrostatic(graphite, xx, yy, zz)
+
+Another possibility is to calculate the electrostatic potential for an infinite slab in the 
+x-y plane, but finite in z-direction, using :func:`pelectrostatic` (p for projected)::
+
+    from skued.simulation import pelectrostatic
+
+    extent = np.linspace(-5, 5, 256)
+    xx, yy= np.meshgrid(extent, extent)
+    potential = electrostatic(graphite, xx, yy)
+
+After plot formatting:
+
+.. plot::
+
+    import matplotlib.pyplot as plt
+    import numpy as np
+    from skued import Crystal
+    from skued.simulation import pelectrostatic
+
+    extent = np.linspace(-5, 5, 256)
+    xx, yy = np.meshgrid(extent, extent)
+
+    potential = pelectrostatic(Crystal.from_database('C'), xx, yy)
+
+    fig = plt.figure()
+    ax = fig.add_subplot(111)
+    ax.xaxis.set_visible(False)
+    ax.yaxis.set_visible(False)
+    ax.set_title('Electrostatic potential of graphite')
+    im = ax.imshow(potential)
+    cbar = plt.colorbar(im)
+    cbar.set_label('Electrostatic potential ($V \cdot \AA$)')
+
 :ref:`Return to Top <simulation_tutorial>`

docs/source/tutorials/structure.rst

@@ -37,7 +37,7 @@ constructed like so::
 Protein DataBank files are even easier to handle; simply provide the 4-letter identification code
 and the structure file will be taken care of by scikit-ued::
 	
-	hemoglobin = Crystal.from_pdb('1gxz')
+	hemoglobin = Crystal.from_pdb('1gzx')
 
 Another convenient way to construct a :class:`Crystal` is through the `Crystallography Open Database <http://www.crystallography.net/cod/>`_::
 
@@ -101,6 +101,29 @@ from the :attr:`source` attribute::
     c = Crystal.from_pdb('1gzx')
     print(c.source)
 
+:class:`Crystal` instances can be converted to NumPy arrays as well::
+
+    import numpy
+
+    arr = numpy.array(Crystal.from_database('Si'))
+
+:data:`arr` will contain one row per unit cell atom:
+
+.. table::
+    :widths: grid
+
+    +--------------+-----------------------+
+    |Atomic Number |Fractional coordinates |
+    +==============+=======+========+======+
+    |      Z1      |  x1   |   y1   |  z1  |
+    +--------------+-------+--------+------+
+    |      Z2      |  x2   |   y2   |  z2  |
+    +--------------+-------+--------+------+
+    |      Z3      |  x3   |   y3   |  z3  |
+    +--------------+-------+--------+------+
+    |            ...                       |
+    +--------------------------------------+
+
 Lattice vectors and reciprocal space
 -------------------------------------
 Once a :class:`Crystal` object is ready, you can manipulate the lattice parameters via the underlying :class:`Lattice`

skued/simulation/potential.py

@@ -32,12 +32,10 @@ def electrostatic(crystal, x, y, z):
     
     See also
     --------
-    pelectrostatic
-        Projected electrostatic potential from a crystal
+    pelectrostatic: projected electrostatic potential of an infinite crystal.
     """
     # TODO: split xx and yy into smalled non-repeating unit
     # TODO: multicore
-    # TODO: pre-load scattering params _a, _b, _c, and _d into an appropriate shape
 
     potential = np.zeros_like(x, dtype = np.float)
     r = np.zeros_like(x, dtype = np.float)
@@ -53,52 +51,51 @@ def electrostatic(crystal, x, y, z):
     potential[np.isinf(potential)] = m
     return potential
 
-def pelectrostatic(crystal, xx, yy, bounds = None):
+def pelectrostatic(crystal, x, y, bounds = None):
     """
     Projected electrostatic potential from a crystal calculated on a real-space mesh, 
-    assuming an infinite crystal. Projection axis is defined as the z-axis. To project the potential
-    along a different axis, the crystal can be rotated with ``Crystal.transform``.
+    assuming an infinite crystal in x and y. Projection axis is defined as the z-axis. 
+    To project the potential along a different axis, the crystal can be rotated with ``Crystal.transform``.
 
     Parameters
     ----------
     crystal : skued.Crystal
         
-    xx, yy:  `~numpy.ndarray`
-        Real space coordinates mesh. 
+    x, y:  `~numpy.ndarray`
+        Real-space coordinates. 
+    bounds : iterable or None, optional
+        Bounds of atom inclusion. Atoms with real-space z-position outside [ min(bounds), max(bounds) )
+        are not counted in the computation.
     
     Returns
     -------
     potential : `~numpy.ndarray`, dtype float
-        Linear superposition of atomic potential [V*Angs]
+        Linear superposition of electrostatic potential [V*Angs]
+    
+    See also
+    --------
+    electrostatic: three-dimensional electrostatic potential of an infinite crystal.
     """
     # TODO: split xx and yy into smalled non-repeating unit
+    #       np.unique(np.mod(xx, per_x))
     # TODO: multicore
-    # TODO: pre-load scattering params _a, _b, _c, and _d into an appropriate shape
 
     if bounds:
         min_z, max_z = min(bounds), max(bounds)
         atoms = (atom for atom in iter(crystal) if min_z <= atom.xyz(crystal)[2] < max_z)
     else:
         atoms = iter(crystal)
 
-    potential = np.zeros_like(xx)
-    zz = np.zeros_like(xx)
+    potential = np.zeros_like(x, dtype = np.float)
+    z = np.zeros_like(x)
     for atom in atoms:
         xa, ya, _ = atom.xyz(crystal)
-        r = minimum_image_distance(xx - xa, yy - ya, zz, lattice = np.array(crystal.lattice_vectors)).reshape(-1,1)
-        potential += np.sum( 2*atom._a*bessel(2*pi*r*np.sqrt(atom._b)) + (atom._c/atom._d) * np.exp( -(r*pi)**2 / atom._d), axis = -1).reshape(xx.shape)
+        r = minimum_image_distance(x - xa, y - ya, z, lattice = np.array(crystal.lattice_vectors)).reshape(-1,1)
+        potential += np.sum( 2*atom._a*bessel(2*pi*r*np.sqrt(atom._b)) + (atom._c/atom._d) * np.exp( -(r*pi)**2 / atom._d), axis = -1).reshape(x.shape)
     potential *= 2 * a0 * e * (pi**2)
 
     # Due to sampling, x,y, and z might pass through the center of atoms
     # Replace n.inf by the next largest value
     potential[np.isinf(potential)] = np.nan
     potential[np.isnan(potential)] = np.nanmax(potential)
-    return potential
-
-def _proj_elec_atm(atom, xx, yy, lattice):
-    potential = np.zeros_like(xx, dtype = np.float)[:,:,None]
-    xa, ya, _ = tuple(atom.xyz(lattice))
-    r = np.zeros_like(xx)[:,:,None, None]
-    r[:,:, 0, 0] = minimum_image_distance(xx - xa, yy - ya, np.zeros_like(xx), lattice)
-    potential = np.sum( 2*atom._a*bessel(2*pi*r*np.sqrt(atom._b)) + (atom._c/atom._d) * np.exp( -(r*pi)**2 / atom._d), axis = -1)
-    return 2*a0*e*(pi**2)*np.squeeze(potential)   
+    return potential

skued/structure/crystal.py

@@ -84,7 +84,7 @@ class Crystal(Lattice):
         Provenance, e.g. filename.
     """
 
-    builtins = set(map(lambda fn: os.path.basename(fn).split('.')[0], CIF_ENTRIES))
+    builtins = frozenset(map(lambda fn: os.path.basename(fn).split('.')[0], CIF_ENTRIES))
 
     def __init__(self, unitcell, lattice_vectors, source = None, **kwargs):
         self.unitcell = frozenset(unitcell)
@@ -104,6 +104,14 @@ def __eq__(self, other):
         return (isinstance(other, self.__class__) 
                 and (set(self) == set(other))
                 and super().__eq__(other))
+
+    def __array__(self):
+        """ Returns an array in which each row represents a unit cell atom """
+        arr = np.empty(shape = (len(self), 4), dtype = np.float)
+        for row, atm in enumerate(self):
+            arr[row, 0] = atm.atomic_number
+            arr[row, 1:] = atm.coords
+        return arr
 
     @classmethod
     def from_cif(cls, path):
@@ -256,9 +264,8 @@ def primitive(self, symprec = 1e-2):
     def spglib_cell(self):
         """ 3-tuple of ndarrays properly formatted for spglib's routines """
         lattice = np.array(self.lattice_vectors)
-        positions = np.array([atom.coords for atom in iter(self)])
-        numbers = np.array(tuple(atom.atomic_number for atom in iter(self)))
-        return (lattice, positions, numbers)
+        arr = np.asarray(self)
+        return (lattice, arr[:, 1:], arr[:,0])
 
     def ase_atoms(self, **kwargs):
         """ 
@@ -346,66 +353,6 @@ def spacegroup_info(self, symprec = 1e-2, angle_tolerance = -1.0):
         if err_msg:
             raise RuntimeError('Symmetry-determination has returned the following error: {}'.format(err_msg))
 
-    def periodicity(self):
-        """
-        Crystal periodicity in x, y and z direction from the lattice constants.
-        This is effectively a bounding cube for the unit cell, which is itself a unit cell.
-
-        Parameters
-        ----------
-        lattice : Lattice
-
-        Returns
-        -------
-        out : tuple
-            Periodicity in x, y and z directions [angstroms]
-        
-        Notes
-        -----
-        Warning: the periodicity of the lattice depends on its orientation in real-space.
-        """
-        # By definition of a lattice, moving by the projection of all Lattice
-        # vectors on an axis should return you to an equivalent lattice position
-        e1, e2, e3 = np.eye(3)
-        per_x = sum( (abs(np.vdot(e1,a)) for a in self.lattice_vectors) )
-        per_y = sum( (abs(np.vdot(e2,a)) for a in self.lattice_vectors) )
-        per_z = sum( (abs(np.vdot(e3,a)) for a in self.lattice_vectors) )
-        return per_x, per_y, per_z
-
-    def potential(self, x, y, z):
-        """
-        Scattering potential calculated on a real-space mesh, assuming an
-        infinite crystal.
-
-        Parameters
-        ----------
-        x, y, z : `~numpy.ndarray`
-            Real space coordinates mesh. 
-        
-        Returns
-        -------
-        potential : `~numpy.ndarray`, dtype float
-            Linear superposition of atomic potential [V*Angs]
-
-        See also
-        --------
-        skued.minimum_image_distance
-        """
-        # TODO: multicore
-        potential = np.zeros_like(x, dtype = np.float)
-        r = np.zeros_like(x, dtype = np.float)
-        for atom in iter(self):
-            ax, ay, az = atom.xyz(self)
-            r[:] = minimum_image_distance(x - ax, y - ay, z - az, 
-                                          lattice = self.lattice_vectors)
-            potential += atom.potential(r)
-
-        # Due to sampling, x,y, and z might pass through the center of atoms
-        # Replace np.inf by the next largest value
-        m = potential[np.isfinite(potential)].max()
-        potential[np.isinf(potential)] = m
-        return potential
-
     def scattering_vector(self, h, k, l):
         """
         Returns the scattering vector G from Miller indices.

skued/structure/lattice.py

@@ -2,7 +2,7 @@
 from math import sin, cos, tan, sqrt, radians
 import numpy as np
 from numpy.linalg import norm
-from .. import transform
+from .. import transform, change_basis_mesh
 
 e1, e2, e3 = np.eye(3) # Euclidian basis
 
@@ -113,6 +113,16 @@ def reciprocal_vectors(self):
         b3 = 2*np.pi*np.cross(self.a1, self.a2)/cell_volume
         return b1, b2, b3
 
+    @property
+    def periodicity(self):
+        """ Crystal periodicity in x, y and z direction from the lattice constants.
+        This is effectively a bounding cube for the unit cell, which is itself a unit cell. """
+        e1, e2, e3 = np.eye(3)
+        per_x = sum( (abs(np.vdot(e1,a)) for a in self.lattice_vectors) )
+        per_y = sum( (abs(np.vdot(e2,a)) for a in self.lattice_vectors) )
+        per_z = sum( (abs(np.vdot(e3,a)) for a in self.lattice_vectors) )
+        return per_x, per_y, per_z
+
     def transform(self, *matrices):
         """
         Transforms the real space coordinates according to a matrix.

skued/structure/tests/test_crystal.py

@@ -38,6 +38,17 @@ def test_back_and_forth(self):
         # self.assertSetEqual(set(self.crystal), set(crystal2))
         self.assertEqual(len(self.crystal), len(crystal2))
 
+class TestCrystalMethods(unittest.TestCase):
+
+    def setUp(self):
+        name = choice(list(Crystal.builtins))
+        self.crystal = Crystal.from_database(name)
+
+    def test_array(self):
+        """ Test Crystal.__array__ """
+        arr = np.array(self.crystal)
+        self.assertSequenceEqual(arr.shape, (len(self.crystal), 4))
+
 class TestBoundedReflections(unittest.TestCase):
 
     def setUp(self):