def assertArrayAlmostEqual(self, a, b, tol=1e-7):
a = farray(a)
b = farray(b)
self.assertEqual(a.shape, b.shape)
if isnan(a).any() or isnan(b).any():
print 'a'
print a
print 'b'
print b
self.fail('Not a number (NaN) found in array')
if a.dtype.kind != 'f':
self.assert_((a == b).all())
else:
absdiff = abs(a-b)
if absdiff.max() > tol:
loc = [x+1 for x in unravel_index(absdiff.argmax()-1, absdiff.shape) ]
print 'a'
print a
print
print 'b'
print b
print
print 'Absolute difference'
print absdiff
self.fail('Maximum abs difference between array elements is %e at location %r' % (absdiff.max(), loc))
def setUp(self):
self.z0 = fzeros(0)
self.z1 = fzeros(1)
self.z2 = fzeros((2, 2))
self.f = farray((1, 2, 3))
self.f2 = farray(((1, 2, 3), (4, 5, 6))).T
self.na = numpy.array(self.f2)
self.s = farray(('s1', 's2', 's3'))
def test_heap_sort_r(self):
a = farray([1., 4., 9., 7., -5., 3., -10., 11.])
i = farray([1, 2, 3, 4, 5, 6, 7, 8], dtype=int32)
r = farray([1., 2., 3., 4., 5., 6.6, 7., 8.])
heap_sort(a, i_data=i, r_data=r)
self.assertArrayAlmostEqual(a, [-10., -5., 1., 3., 4., 7., 9., 11.])
self.assertArrayAlmostEqual(i, [7., 5., 1., 6., 2., 4., 3., 8.])
self.assertArrayAlmostEqual(r, [7., 5., 1., 6.6, 2., 4., 3., 8.])
def write(self, at):
comment = self.comment
data = self.data
origin = self.origin
extent = self.extent
comment2 = self.comment2
if data is None and hasattr(at, 'data'):
data = at.data
if data is None:
raise ValueError("Cannot write .cube file without any volumetric data")
if comment is None and 'comment' in at.params:
comment = at.params['comment1']
if comment is None: comment = ''
comment2 = comment2
if comment2 is None and 'comment2' in at.params:
comment2 = at.params['comment2']
if comment2 is None: comment2 = ''
origin = origin
if origin is None and 'origin' in at.params:
origin = at.params['origin']
if origin is None: origin = [0., 0., 0.]
origin = farray(origin)
self.f.write(comment.strip()+'\n')
self.f.write(comment2.strip()+'\n')
origin_x, origin_y, origin_z = origin/BOHR
self.f.write('%d %f %f %f\n' % (at.n, origin_x, origin_y, origin_z))
if extent is None and 'extent' in at.params:
extent = at.params['extent']
if extent is None:
extent = at.lattice
if extent.shape == (3,):
extent = np.diag(extent)
extent = farray(extent)
for i in (1,2,3):
n = data.shape[i-1]
voxel_x, voxel_y, voxel_z = extent[:,i]/BOHR/n
n += 1 # for padding
self.f.write('%d %f %f %f\n' % (n, voxel_x, voxel_y, voxel_z))
for i in frange(at.n):
self.f.write('%d 0.0 %f %f %f\n' % (at.z[i], at.pos[1,i]/BOHR, at.pos[2,i]/BOHR, at.pos[3,i]/BOHR))
padded_data = np.zeros([s+1 for s in data.shape])
padded_data[:-1,:-1,:-1] = data
padded_data[-1,:,:] = padded_data[0,:,:]
padded_data[:,-1,:] = padded_data[:,0,:]
padded_data[:,:,-1] = padded_data[:,:,0]
for d in padded_data.flat:
self.f.write('%f\n' % d)
def test2dextindexintarray(self):
self.assert_((self.f2[farray([1])] == self.na[...,
numpy.array([0])]).all())
self.assert_(
(self.f2[farray([1, 2])] == self.na[...,
numpy.array([0, 1])]).all())
self.assert_(
(self.f2[farray([2, 1]), :] == self.na[...,
numpy.array([1, 0]), :]
).all())
def testsinglereal(self):
t = Table()
t.append(1.0)
t.append(2.0)
t.append(3.0)
self.assertArrayAlmostEqual(t.real[1, :], farray([1.0, 2.0, 3.0]))
t.append(realpart=[4.0, 5.0, 6.0])
self.assertArrayAlmostEqual(t.real[1, :],
farray([1.0, 2.0, 3.0, 4.0, 5.0, 6.0]))
self.assertArrayAlmostEqual(t.real_part(column=1, n0=t.n),
farray([1.0, 2.0, 3.0, 4.0, 5.0, 6.0]))
def testintreal(self):
t = Table()
t.append(1, 1.0)
t.append(2, 2.0)
t.append(3, 3.0)
t.append(intpart=[4, 5, 6], realpart=[4.0, 5.0, 6.0])
self.assertArrayAlmostEqual(t.int[1, :], farray([1, 2, 3, 4, 5, 6]))
self.assertArrayAlmostEqual(t.int_part(column=1, n0=t.n),
farray([1, 2, 3, 4, 5, 6]))
self.assertArrayAlmostEqual(t.real[1, :], farray([1, 2, 3, 4, 5, 6]))
self.assertArrayAlmostEqual(t.real_part(column=1, n0=t.n),
farray([1, 2, 3, 4, 5, 6]))
def testreprstr(self):
self.assertEqual(
repr(self.f2), """FortranArray([[1, 4],
[2, 5],
[3, 6]])""")
self.assertEqual(str(self.f2), """[[1 4]
[2 5]
[3 6]]""")
self.assertEqual(str(self.s), "['s1' 's2' 's3']")
self.assertEqual(str(farray('s')), 's')
s2 = farray((('s', '1'), ('s', '2'))).T
self.assertEqual([str(x) for x in s2], ["['s' '1']", "['s' '2']"])
def testmultintmultlog(self):
t = Table()
t.append(realpart=[1.0, 2.0, 3.0], logicalpart=[False, True, False])
t.append(realpart_2d=numpy.reshape([4.0, 5.0, 6.0, 7.0, 8.0, 9.0],
[3, 2],
order='F'),
logicalpart_2d=numpy.reshape(
[False, False, False, True, True, True], [3, 2],
order='F'))
self.assertArrayAlmostEqual(
t.real,
farray([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]).T)
self.assertArrayAlmostEqual(
t.logical,
farray([[False, True, False], [False, False, False],
[True, True, True]]).T)
def del_atoms(x=None):
rcut = 2.0
#x = Atoms('crack.xyz')
if x == None:
x = Atoms('1109337334_frac.xyz')
else:
pass
x.set_cutoff(3.0)
x.calc_connect()
x.calc_dists()
rem=[]
r = farray(0.0)
u = fzeros(3)
print len(x)
for i in frange(x.n):
for n in frange(x.n_neighbours(i)):
j = x.neighbour(i, n, distance=3.0, diff=u)
if x.distance_min_image(i, j) < rcut and j!=i:
rem.append(sorted([j,i]))
if i%10000==0: print i
rem = list(set([a[0] for a in rem]))
if len(rem) > 0:
print rem
x.remove_atoms(rem)
else:
print 'No duplicate atoms in list.'
x.write('crack_nodup.xyz')
return x
def testnorm(self):
self.assertAlmostEqual(farray(1.0).norm(), 1.0)
self.assertRaises(ValueError, self.f2.T.norm2)
self.assertAlmostEqual(self.f2[1].norm2(), 14.0)
n2 = self.f2.norm2()
self.assertAlmostEqual(n2[1], 14.0)
self.assertAlmostEqual(n2[2], 77.0)
self.assert_((numpy.sqrt(n2) - self.f2.norm() < 1e-6).all())
self.assertRaises(ValueError, fzeros((3, 3, 3)).norm2)
def testsingleint(self):
t = Table()
t.append(1)
t.append(2)
t.append(3)
self.assertEqual(list(t.int), [1, 2, 3])
t.append([4, 5, 6])
self.assertEqual(list(t.int), [1, 2, 3, 4, 5, 6])
self.assertArrayAlmostEqual(t.int_part(1, t.n),
farray([1, 2, 3, 4, 5, 6]))
def write(self, at):
self.f.write('''#include "colors.inc"
camera
{
right x*400/600
location <%f,%f,%f>
look_at <0,0,0>
}
background
{
colour White
}
''' % tuple(self.camerapos))
for light in self.lights:
self.f.write('''
light_source
{
<%f,%f,%f>
colour White
}
''' % tuple(light))
center = farray(self.center)
self.f.write('union {\n')
for i in frange(at.n):
if self.skip_hydrogens and str(at.species[i]) == 'H': continue
p = at.pos[i] - center
if self.colour is not None:
c = self.colour[i]
else:
c = farray((1., 1., 1.))
self.f.write(
'sphere { <%f,%f,%f>, %r pigment {color <%f,%f,%f> } finish { phong .8 } }\n'
% (p[1], -p[2], p[3], self.radius, c[1], c[2], c[3]))
if self.rotate is not None:
self.f.write('rotate <%f,%f,%f>\n' % tuple(self.rotate))
self.f.write('}')
def testtableappend(self):
t1 = Table(2, 1, 0, 0)
t1.append([1, 2], 3.0)
t1.append([4, 5], 6.0)
t1.append([7, 8], 9.0)
t2 = Table()
t2.append(t1)
t2.append(t1)
self.assertEqual(list(t2.int[1, :]), [1, 4, 7, 1, 4, 7])
self.assertEqual(list(t2.int[2, :]), [2, 5, 8, 2, 5, 8])
self.assertArrayAlmostEqual(t2.real[1, :],
farray([3.0, 6.0, 9.0, 3.0, 6.0, 9.0]))
def test_binary_search_r(self):
r = farray([1., 2., 3., 4., 5., 6.6, 7., 8.])
i = binary_search(r, 0.)
self.assertEqual(i, 0)
i = binary_search(r, 1.1)
self.assertEqual(i, 1)
i = binary_search(r, 5.)
self.assertEqual(i, 5)
i = binary_search(r, 6.6)
self.assertEqual(i, 6)
i = binary_search(r, 6.9)
self.assertEqual(i, 6)
i = binary_search(r, 8.8)
self.assertEqual(i, 8)
def get_value(self, k):
"Return a _copy_ of a value stored in Dictionary"
if not k in self:
raise KeyError('Key "%s" not found ' % k)
t, s, s2 = self.get_type_and_size(k)
if t == T_NONE:
v = None
elif t == T_INTEGER:
v,p = self._get_value_i(k)
elif t == T_REAL:
v,p = self._get_value_r(k)
elif t == T_COMPLEX:
v,p = self._get_value_c(k)
elif t == T_CHAR:
v,p = self._get_value_s(k)
v = v.strip()
elif t == T_LOGICAL:
v,p = self._get_value_l(k)
v = bool(v)
elif t == T_INTEGER_A:
v,p = self._get_value_i_a(k,s)
elif t == T_REAL_A:
v,p = self._get_value_r_a(k,s)
elif t == T_COMPLEX_A:
v,p = self._get_value_c_a(k,s)
elif t == T_CHAR_A:
v,p = self._get_value_s_a2(k,s2[1], s2[2])
v = v[...,1] # Last index is length of string, here fixed to 1
v.strides = (1, v.shape[0]) # Column-major storage
elif t == T_LOGICAL_A:
v,p = self._get_value_l_a(k,s)
v = farray(v, dtype=bool)
elif t == T_INTEGER_A2:
v,p = self._get_value_i_a2(k, s2[1], s2[2])
elif t == T_REAL_A2:
v,p = self._get_value_r_a2(k, s2[1], s2[2])
elif t == T_DICT:
v,p = self._get_value_dict(k)
else:
raise ValueError('Unsupported dictionary entry type %d' % t)
return v
def test_string_array(self):
d = Dictionary()
d['string'] = ['one', 'two', 'three']
a = FortranArray(d.get_array('string'))
a[:, 1] = farray(list('hello'), 'S1')
self.assertEqual(list(a2s(d['string'])), list(a2s(a)))
def test_binary_search(self):
a = farray([1, 2, 3, 4, 5, 6, 7, 8], dtype=int32)
i = binary_search(a, 4)
self.assertEqual(i, 4)
def test2dextindexintfail(self):
self.assertRaises(ValueError, self.f2.__getitem__, farray(0))
self.assertRaises(ValueError, self.f2.__getitem__, [[0]])
self.assertRaises(ValueError, self.f2.__getitem__, farray('str'))
self.assertRaises(ValueError, self.f2.__getitem__, [{}])
def test2dextindexsetlist(self):
self.f2[farray([1, 2]), :] = 0
self.assertEqual(list(self.f2[1]), [0, 0, 3])
self.assertEqual(list(self.f2[2]), [0, 0, 6])
def testvaluelists(self):
params = Dictionary('xyz={1.0 2.0 3.0} abc="1 2 3"')
self.assertAlmostEqual(params['xyz'][1], 1.0)
self.assertAlmostEqual(params['xyz'][2], 2.0)
self.assertAlmostEqual(params['xyz'][3], 3.0)
self.assertArrayAlmostEqual(params['abc'], farray([1, 2, 3]))
def testextsetlistbool(self):
self.f[farray([True, False, False])] = 0
self.assertEqual(list(self.f), [0, 2, 3])
def testextsetfarray(self):
self.f[farray([1])] = 0
self.assertEqual(list(self.f), [0, 2, 3])
def testrows(self):
self.assertEqual(farray(0).rows.next(), 0)
self.assertEqual(list(self.f.rows), [1, 2, 3])
def testcols(self):
self.assertEqual(farray(0).cols.next(), 0)
self.assertEqual(list(self.f.cols), [1, 2, 3])
def CubeReader(f, property_name='charge', discard_repeat=True, format=None):
def convert_line(line, *fmts):
return (f(s) for f, s in zip(fmts, line.split()))
if type(f) == type(''):
f = open(f)
opened = True
# First two lines are comments
comment1 = f.readline()
comment2 = f.readline()
# Now number of atoms and origin
n_atom, origin_x, origin_y, origin_z = convert_line(
f.readline(), int, float, float, float)
origin = farray([origin_x, origin_y, origin_z]) * BOHR
# Next three lines define number of voxels and shape of each element
shape = [0, 0, 0]
voxel = fzeros((3, 3))
for i in (1, 2, 3):
shape[i - 1], voxel[1, i], voxel[2, i], voxel[3, i] = convert_line(
f.readline(), int, float, float, float)
at = Atoms(n=n_atom, lattice=voxel * BOHR * shape)
at.add_property(property_name, 0.0)
prop_array = getattr(at, property_name)
# Now there's one line per atom
for i in frange(at.n):
at.z[i], prop_array[i], at.pos[1,
i], at.pos[2,
i], at.pos[3,
i] = convert_line(
f.readline(),
int, float,
float, float,
float)
at.pos[:, i] *= BOHR
at.set_atoms(at.z)
# Rest of file is volumetric data
data = np.fromiter((float(x) for x in f.read().split()), float, count=-1)
if data.size != shape[0] * shape[1] * shape[2]:
raise IOError("Bad array length - expected shape %r, but got size %d" %
(shape, data.size))
# Save volumetric data in at.data
data = farray(data.reshape(shape))
# Discard periodic repeats?
if discard_repeat:
at.data = data[:-1, :-1, :-1]
shape = [s - 1 for s in shape]
at.set_lattice(voxel * BOHR * shape, False)
at.params['comment1'] = comment1
at.params['comment2'] = comment2
at.params['origin'] = origin
at.params['shape'] = shape
# save grids in at.grid_x, at.grid_y, at.grid_z
if at.is_orthorhombic:
at.grid_x, at.grid_y, at.grid_z = np.mgrid[origin[1]:origin[1] +
at.lattice[1, 1]:shape[0] *
1j, origin[2]:origin[2] +
at.lattice[2, 2]:shape[1] *
1j, origin[3]:origin[3] +
at.lattice[3, 3]:shape[2] *
1j]
if opened:
f.close()
yield at
def test_heap_sort_i(self):
a = farray([1, 4, 9, 7, -5, 3, -10, 11], dtype=int32)
r = farray([1., 2., 3., 4., 5., 6.6, 7., 8.])
heap_sort(a, r_data=r)
self.assertArrayAlmostEqual(a, [-10, -5, 1, 3, 4, 7, 9, 11])
self.assertArrayAlmostEqual(r, [7., 5., 1., 6.6, 2., 4., 3., 8.])
