def test_dummyprop(self):
# Define a dummy derived property and test it
class DummyProperty(AtomsProperty):
default_name = 'dummy'
default_params = {'mul': 2.0}
@staticmethod
def extract(s, mul):
return s.positions.shape[0] * mul
# Now two atoms objects to test it on
a1 = Atoms('C')
a2 = Atoms('CC')
dummyDoubled = DummyProperty(name='doubledummy', mul=4)
self.assertEqual(DummyProperty.get(a1), 2.0)
self.assertEqual(dummyDoubled(a2), 8.0)
# Also check that improper parameters are rejected
self.assertRaises(ValueError, DummyProperty, wrong='this is wrong')
# Test behaviour on a collection instead
c1 = AtomsCollection([a1, a2])
c2 = AtomsCollection([a2, a1])
DummyProperty.get(c1, store_array=True)
dummyDoubled(c2, store_array=True)
self.assertTrue(np.all(c1.get_array('dummy') == [2, 4]))
self.assertTrue(np.all(c2.get_array('doubledummy') == [8, 4]))
def test_save(self):
# Test saving and loading collection
testcoll = AtomsCollection([Atoms('H')])
testcoll.set_array('test', np.array([2]))
outf = os.path.join(_TEST_DIR, 'collection.pkl')
testcoll.save(outf)
# Reload
testcoll = AtomsCollection.load(outf)
self.assertEqual(testcoll.get_array('test')[0], 2)
os.remove(outf)
def test_sorting(self):
# Generate a few structures
struct_n = 5
aselist = []
for n in range(struct_n):
aselist.append(Atoms())
testcoll = AtomsCollection(aselist)
testcoll.set_array('sorter', np.array(range(struct_n, 0, -1)))
testcoll.set_array('sorted', np.array(range(1, struct_n + 1)))
testcoll = testcoll.sorted_byarray('sorter')
self.assertTrue(
np.all(testcoll.get_array('sorted') == range(struct_n, 0, -1)))
def test_arrays(self):
# Generate a few random structures
elrnd = ['H', 'C', 'O', 'N']
asernd = []
for n in range(4):
aselen = np.random.randint(1, 10)
asernd.append(
Atoms(symbols=np.random.choice(elrnd, aselen),
positions=np.random.random((aselen, 3))))
testcoll = AtomsCollection(asernd)
# Now try assigning some arrays
arr = np.arange(testcoll.length)
testcoll.set_array('testarr', arr, shape=(1, ))
testcoll.set_array('testarr_2', list(zip(arr, arr)), shape=(2, ))
testcoll.set_array('testarr_func',
lambda a: len(a.get_positions()),
shape=(1, ))
self.assertTrue(np.all(testcoll.get_array('testarr') == arr))
print aColl.all.info, '\n\n'
# Collections can also be sliced like Numpy arrays for convenience
aColl02 = aColl[0:2]
aColl25 = aColl[2:5]
# Then join them together
aColl05 = aColl02 + aColl25
print "---- Collection slice lengths ---- \n"
print "aColl02 = {0}\taColl25 = {1}\taColl05 = {2}\n\n".format(
aColl02.length, aColl25.length, aColl05.length)
# Collections can also store "arrays" of data, similarly to Atoms objects in ase
# These arrays' elements are tied each to one structure, and can be used to sort them
arr = range(10, 0,
-1) # Let's use this array to reverse the order of a collection
aColl.set_array('reversed_range', arr)
aCollSorted = aColl.sorted_byarray('reversed_range')
print "---- Getting an array from a collection ---- \n"
print "Unsorted: ", aColl.get_array('reversed_range'), "\n"
print "Sorted: ", aCollSorted.get_array('reversed_range'), "\n\n"
# And to make sure
print "---- First vs. last elements ---- \n"
print aColl.structures[0].get_positions(), "\n"
print aCollSorted.structures[-1].get_positions()
