Added distances as separate functions

crystals/__init__.py

@@ -10,6 +10,8 @@
 from .atom import Atom
 from .atom import frac_coords
 from .atom import real_coords
+from .atom import distance_fractional
+from .atom import distance_cartesian
 from .atom_data import ELEM_TO_MAGMOM
 from .atom_data import ELEM_TO_MASS
 from .atom_data import ELEM_TO_NAME

crystals/atom.py

@@ -13,7 +13,6 @@
 from .atom_data import NUM_TO_ELEM
 from .lattice import Lattice
 
-
 # TODO: store atomic data as class attributes?
 class Atom(object):
     """
@@ -74,17 +73,13 @@ def __repr__(self):
             self.element, *tuple(self.coords_fractional)
         )
 
-    # TODO: add `distance_from` function for atoms on a lattice
-    def __sub__(self, other):
-        """ Distance between two atoms in **fractional** coordinates. """
-        return np.linalg.norm(self.coords_fractional - other.coords_fractional)
-
     def __eq__(self, other):
         return (
             isinstance(other, self.__class__)
             and (self.element == other.element)
             and (self.magmom == other.magmom)
             and np.allclose(self.coords_fractional, other.coords_fractional, atol=1e-3)
+            and np.allclose(self.coords_cartesian, other.coords_cartesian, atol=1e-3)
             and np.allclose(self.displacement, other.displacement, atol=1e-3)
         )
 
@@ -94,6 +89,9 @@ def __hash__(self):
                 self.element,
                 self.magmom,
                 tuple(np.round(self.coords_fractional, 3)),
+                tuple(
+                    np.round(self.coords_cartesian, 3)
+                ),  # effect of lattice is encoded in cartesian coordinates
                 tuple(np.round(self.displacement, 3)),
             )
         )
@@ -139,50 +137,8 @@ def coords_cartesian(self):
         ------
         RuntimeError : if this atom is not place on a lattice
         """
-        if not self.lattice:
-            return real_coords(self.coords_fractional, np.eye(3))
         return real_coords(self.coords_fractional, self.lattice.lattice_vectors)
 
-    def distance_fractional(self, atm):
-        """
-        Calculate the distance between atoms in fractional coordinates.
-
-        Parameters
-        ----------
-        atm : ``crystals.Atom`` instance
-        
-        Returns
-        -------
-        dist : float
-            Distance in fractional coordinates
-        """
-        return np.linalg.norm(self.coords_fractional - atm.coords_fractional)
-
-    def distance_cartesian(self, atm):
-        """
-        Calculate the distance between atoms in cartesian coordinates.
-        If the parent lattices are different, an error is raised.
-
-        Parameters
-        ----------
-        atm : ``crystals.Atom`` instance
-        
-        Returns
-        -------
-        dist : float
-            Distance in Angstroms.
-        
-        Raises
-        ------
-        RuntimeError : if atoms are not defined on the same lattice.
-        """
-        if self.lattice != atm.lattice:
-            raise RuntimeError(
-                "Cartesian distance is undefined if atoms are sitting on different lattices."
-            )
-
-        return np.linalg.norm(self.coords_cartesian - atm.coords_cartesian)
-
     def transform(self, *matrices):
         """
         Transforms the real space coordinates according to a matrix.
@@ -240,3 +196,53 @@ def frac_coords(real_coords, lattice_vectors):
     """
     COB = change_of_basis(np.eye(3), np.array(lattice_vectors))
     return np.mod(transform(COB, real_coords), 1)
+
+
+def distance_fractional(atm1, atm2):
+    """ 
+    Calculate the distance between two atoms in fractional coordinates.
+    
+    Parameters
+    ----------
+    atm1, atm2 : ``crystals.Atom``
+
+    Returns
+    -------
+    dist : float
+        Fractional distance between atoms.
+    
+    Raises
+    ------
+    RuntimeError : if atoms are not associated with the same lattice.
+    """
+    if atm1.lattice != atm2.lattice:
+        raise RuntimeError(
+            "Distance is undefined if atoms are sitting on different lattices."
+        )
+
+    return np.linalg.norm(atm1.coords_fractional - atm2.coords_fractional)
+
+
+def distance_cartesian(atm1, atm2):
+    """ 
+    Calculate the distance between two atoms in cartesian coordinates.
+    
+    Parameters
+    ----------
+    atm1, atm2 : ``crystals.Atom``
+
+    Returns
+    -------
+    dist : float
+        Cartesian distance between atoms in Angstroms..
+    
+    Raises
+    ------
+    RuntimeError : if atoms are not associated with the same lattice.
+    """
+    if atm1.lattice != atm2.lattice:
+        raise RuntimeError(
+            "Distance is undefined if atoms are sitting on different lattices."
+        )
+
+    return np.linalg.norm(atm1.coords_cartesian - atm2.coords_cartesian)

crystals/crystal.py

@@ -85,10 +85,11 @@ class Crystal(AtomicStructure, Lattice):
     builtins = frozenset(map(lambda fn: fn.stem, CIF_ENTRIES))
 
     def __init__(self, unitcell, lattice_vectors, source=None, **kwargs):
-        unitcell = list(unitcell)
-        for atom in unitcell:
-            atom.lattice = self
         super().__init__(atoms=unitcell, lattice_vectors=lattice_vectors, **kwargs)
+
+        for atom in iter(self):
+            atom.lattice = self
+
         self.source = source
 
     @classmethod

docs/guide.rst

@@ -325,25 +325,20 @@ or :class:`Lattice) can be accessed using the :meth:`Atom.coords_cartesian` meth
 Atomic distances
 ----------------
 
-The distance between two atoms can be calculated by taking their difference::
+The fractional/cartesian distance between two atoms sitting *on the same lattice* is possible:
 
-    >>> copper = Atom('Cu', coords = [0,0,0])
-    >>> silver = Atom('Ag', coords = [1,0,0])
-    >>> silver.distance_fractional(copper)			# distance in fractional coordinates
-    1.0
-
-The cartesian distance between two atoms sitting *on the same lattice* is also possible:
-
-    >>> from crystals import Crystal
+    >>> from crystals import Crystal, distance_fractional, distance_cartesian
     >>> graphite = Crystal.from_database('C')
     >>> 
     >>> carbon1, carbon2, *_ = tuple(graphite)
-    >>> carbon1.coords_cartesian
-    array([0.000, 0.000, 5.033])
-    >>> carbon2.coords_cartesian
-    array([1.232, 0.711, 5.033])
-    >>> carbon1.distance_cartesian(carbon2)
-    1.4225981762919013
+    >>> carbon1
+    < Atom C  @ (0.00, 0.00, 0.25) >
+    >>> carbon2
+    < Atom C  @ (0.67, 0.33, 0.75) >
+    >>> distance_fractional(carbon1, carbon2)
+    0.8975324197487241
+    >>> distance_cartesian(carbon1, carbon2)    # in Angstroms
+    3.6446077732986644
 
 If atoms are not sitting on the same lattice, calculating the distance should not be defined. In this case, an exception is raised:
 
@@ -354,11 +349,11 @@ If atoms are not sitting on the same lattice, calculating the distance should no
 >>> gold1, *_ = tuple(gold)
 >>> silver1, *_ = tuple(silver)
 >>>
->>> gold1.distance_cartesian(silver1)
+>>> distance_cartesian(gold1, silver1)
 Traceback (most recent call last):
   File "<stdin>", line 1, in <module>
   File "(...omitted...)\crystals\atom.py", line 181, in distance_cartesian
     "Cartesian distance is undefined if atoms are sitting on different lattices."
-RuntimeError: Cartesian distance is undefined if atoms are sitting on different lattices.
+RuntimeError: Distance is undefined if atoms are sitting on different lattices.
 
 :ref:`Return to Top <user_guide>`

docs/reference.rst

@@ -49,6 +49,13 @@ To help with fleshing out unit cell atoms from symmetry operators:
     symmetry_expansion
     lattice_system
 
+.. autosummary::
+    :toctree: functions/
+    :nosignatures:
+
+    distance_fractional
+    distance_cartesian
+
 Conversion between ``Crystal`` and other packages
 
 .. autosummary::

tests/test_atom.py

@@ -8,6 +8,8 @@
 import numpy as np
 
 from crystals import Atom
+from crystals import distance_fractional
+from crystals import distance_cartesian
 from crystals import Lattice
 from crystals.affine import rotation_matrix
 
@@ -66,12 +68,42 @@ def test_atom_array(self):
         self.assertEqual(arr[0], self.atom.atomic_number)
         self.assertTrue(np.allclose(arr[1::], self.atom.coords_fractional))
 
-    def test_distance(self):
-        """ Test the distance between atoms """
+
+class TestAtomicDistances(unittest.TestCase):
+    def test_distance_fractional(self):
+        """ Test the fractional distance between atoms """
         atm1 = Atom("He", [0, 0, 0])
         atm2 = Atom("He", [1, 0, 0])
-        self.assertEqual(atm1.distance_fractional(atm2), 1)
-        self.assertEqual(atm1.distance_fractional(atm2), atm2.distance_fractional(atm1))
+        self.assertEqual(distance_fractional(atm1, atm2), 1)
+        self.assertEqual(
+            distance_fractional(atm1, atm2), distance_fractional(atm2, atm1)
+        )
+
+    def test_distance_cartesian(self):
+        """ Test the cartesian distance between atom """
+        lattice = Lattice(4 * np.eye(3))  # Cubic lattice side length 4 angs
+
+        atm1 = Atom("He", [0, 0, 0], lattice=lattice)
+        atm2 = Atom("He", [1, 0, 0], lattice=lattice)
+
+        self.assertEqual(distance_cartesian(atm1, atm2), 4.0)
+
+    def test_distance_different_lattice(self):
+        """ Test that fractional and cartesian distances 
+        between atoms in different lattices raises an error. """
+        lattice1 = Lattice(np.eye(3))
+        lattice2 = Lattice(2 * np.eye(3))
+
+        atm1 = Atom("He", [0, 0, 0], lattice=lattice1)
+        atm2 = Atom("He", [1, 0, 0], lattice=lattice2)
+
+        with self.subTest("Fractional distance"):
+            with self.assertRaises(RuntimeError):
+                distance_fractional(atm1, atm2)
+
+        with self.subTest("Cartesian distance"):
+            with self.assertRaises(RuntimeError):
+                distance_cartesian(atm1, atm2)
 
 
 if __name__ == "__main__":

tests/test_crystal.py

@@ -5,7 +5,7 @@
 import unittest
 from contextlib import suppress
 from copy import copy
-from copy import deepcopy
+from copy import copy
 from math import radians
 from pathlib import Path
 
@@ -87,7 +87,7 @@ def test_crystal_equality(self):
         """ Tests that Crystal.__eq__ is working properly """
         self.assertEqual(self.crystal, self.crystal)
 
-        cryst2 = deepcopy(self.crystal)
+        cryst2 = copy(self.crystal)
         cryst2.transform(
             2 * np.eye(3)
         )  # This stretches lattice vectors, symmetry operators
@@ -99,29 +99,29 @@ def test_crystal_equality(self):
 
     def test_trivial_rotation(self):
         """ Test rotation by 360 deg around all axes. """
-        unrotated = deepcopy(self.crystal)
+        unrotated = copy(self.crystal)
         r = rotation_matrix(radians(360), [0, 0, 1])
         self.crystal.transform(r)
 
         self.assertEqual(self.crystal, unrotated)
 
     def test_identity_transform(self):
         """ Tests the trivial identity transform """
-        transf = deepcopy(self.crystal)
+        transf = copy(self.crystal)
         transf.transform(np.eye(3))
         self.assertEqual(self.crystal, transf)
 
     def test_one_axis_rotation(self):
         """ Tests the crystal orientation after rotations. """
-        unrotated = deepcopy(self.crystal)
+        unrotated = copy(self.crystal)
         self.crystal.transform(rotation_matrix(radians(37), [0, 1, 0]))
         self.assertNotEqual(unrotated, self.crystal)
         self.crystal.transform(rotation_matrix(radians(-37), [0, 1, 0]))
         self.assertEqual(unrotated, self.crystal)
 
     def test_wraparound_rotation(self):
-        cryst1 = deepcopy(self.crystal)
-        cryst2 = deepcopy(self.crystal)
+        cryst1 = copy(self.crystal)
+        cryst2 = copy(self.crystal)
 
         cryst1.transform(rotation_matrix(radians(22.3), [0, 0, 1]))
         cryst2.transform(rotation_matrix(radians(22.3 - 360), [0, 0, 1]))