Atom instances are now compatible with NumPy arrays

crystals/atom.py

@@ -79,8 +79,8 @@ def __init__(self, element, coords, displacement = (0,0,0), magmom = None, **kwa
             magmom = ELEM_TO_MAGMOM[element]
 
         self.element = element
-        self.coords = np.array(coords, dtype = np.float)
-        self.displacement = np.array(displacement, dtype = np.float)
+        self.coords = np.asfarray(coords)
+        self.displacement = np.asfarray(displacement)
         self.magmom = magmom
 
     def __repr__(self):
@@ -205,3 +205,10 @@ def transform(self, *matrices):
         """
         for matrix in matrices:
             self.coords = transform(matrix, self.coords)
+
+    def __array__(self, *args, **kwargs):
+        """ Returns an array [Z, x, y, z] """
+        arr = np.empty(shape = (4,), *args, **kwargs)
+        arr[0] = self.atomic_number
+        arr[1::] = self.coords
+        return arr

crystals/base.py

@@ -138,6 +138,5 @@ def __array__(self, *args, **kwargs):
         arr = np.empty(shape = (len(self), 4), *args, **kwargs)
         atoms = self.itersorted(key = lambda atm: atm.atomic_number)
         for row, atm in enumerate(atoms):
-            arr[row, 0] = atm.atomic_number
-            arr[row, 1:] = atm.coords
+            arr[row, :] = np.array(atm, *args, **kwargs)
         return arr

tests/test_atom.py

@@ -1,11 +1,15 @@
 # -*- coding: utf-8 -*-
 import unittest
 from copy import deepcopy
-from random import choice, randint, random, seed
+from random import choice
+from random import randint
+from random import random
+from random import seed
 
 import numpy as np
 
-from crystals import Atom, Lattice
+from crystals import Atom
+from crystals import Lattice
 from crystals.affine import rotation_matrix
 
 seed(23)
@@ -72,6 +76,14 @@ def test_transform_back_and_forth(self):
         # No assert sequence almost equal
         for x1, x2 in zip(tuple(before), tuple(after)):
             self.assertAlmostEqual(x1, x2)
+
+    def test_atom_array(self):
+        """ Test that numpy.array(Atom(...)) works as expected """
+        arr = np.array(self.atom)
+        self.assertTupleEqual(arr.shape, (4,))
+        self.assertEqual(arr[0], self.atom.atomic_number)
+        self.assertTrue(np.allclose(arr[1::], self.atom.coords))
+
 
 if __name__ == '__main__':
     unittest.main()

tests/test_base.py

@@ -7,7 +7,9 @@
 
 import numpy as np
 
-from crystals import Atom, AtomicStructure, Crystal
+from crystals import Atom
+from crystals import AtomicStructure
+from crystals import Crystal
 
 
 class TestAtomicStructure(unittest.TestCase):