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.])