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()