def test_loadarray(self):
g0 = Gene('latt_abc_len', 1.0, {})
c1 = AtomsCollection([Atoms('C'), Atoms('C')])
c1.set_array('latt_abc_len', [1, 2])
p1 = PhylogenCluster(c1, [])
p1.set_genes([g0], load_arrays=True)
self.assertTrue(np.all(p1.get_genome_vectors()[0] == [[1], [2]]))
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_sum(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))))
testcoll1 = AtomsCollection(asernd[:2])
testcoll2 = AtomsCollection(asernd[2:])
testcoll1.set_array('joint', ['t', 'e'])
testcoll2.set_array('joint', ['s', 't'])
testcoll = testcoll1
testcoll += testcoll2
self.assertTrue(''.join(testcoll.get_array('joint')) == 'test')
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()