def saveNetcdf(fileName,data): # NetCDF package try: from ase.io.pupynere import NetCDFFile except ImportError: print "ase.io.pupynere not installed on this computer." # get input data shape nx,ny,nz = data.shape print 'data type '+str(data.dtype) # open file f = NetCDFFile(fileName,'w') # create netcdf dimensions f.createDimension('x',nx) f.createDimension('y',ny) f.createDimension('z',nz) rho = f.createVariable('rho','d',('x','y','z')) rho[:] = np.ones((nx,ny,nz)) E = f.createVariable('E','d',('x','y','z')) E[:] = np.zeros((nx,ny,nz)) rho_vx = f.createVariable('rho_vx','d',('x','y','z')) rho_vx[:] = data rho_vy = f.createVariable('rho_vy','d',('x','y','z')) rho_vy[:] = np.zeros((nx,ny,nz)) rho_vz = f.createVariable('rho_vz','d',('x','y','z')) rho_vz[:] = np.zeros((nx,ny,nz)) Bx = f.createVariable('Bx','d',('x','y','z')) Bx[:] = np.zeros((nx,ny,nz)) By = f.createVariable('By','d',('x','y','z')) By[:] = np.zeros((nx,ny,nz)) Bz = f.createVariable('Bz','d',('x','y','z')) Bz[:] = np.zeros((nx,ny,nz)) f.sync() f.close()
def saveNetcdf(fileName, data): # NetCDF package try: from ase.io.pupynere import NetCDFFile except ImportError: print "ase.io.pupynere not installed on this computer." # get input data shape nx, ny, nz = data.shape print 'data type ' + str(data.dtype) # open file f = NetCDFFile(fileName, 'w') # create netcdf dimensions f.createDimension('x', nx) f.createDimension('y', ny) f.createDimension('z', nz) rho = f.createVariable('rho', 'd', ('x', 'y', 'z')) rho[:] = np.ones((nx, ny, nz)) E = f.createVariable('E', 'd', ('x', 'y', 'z')) E[:] = np.zeros((nx, ny, nz)) rho_vx = f.createVariable('rho_vx', 'd', ('x', 'y', 'z')) rho_vx[:] = data rho_vy = f.createVariable('rho_vy', 'd', ('x', 'y', 'z')) rho_vy[:] = np.zeros((nx, ny, nz)) rho_vz = f.createVariable('rho_vz', 'd', ('x', 'y', 'z')) rho_vz[:] = np.zeros((nx, ny, nz)) Bx = f.createVariable('Bx', 'd', ('x', 'y', 'z')) Bx[:] = np.zeros((nx, ny, nz)) By = f.createVariable('By', 'd', ('x', 'y', 'z')) By[:] = np.zeros((nx, ny, nz)) Bz = f.createVariable('Bz', 'd', ('x', 'y', 'z')) Bz[:] = np.zeros((nx, ny, nz)) f.sync() f.close()
class Reader: def __init__(self, filename, comm): self.nc = NetCDFFile(filename) def dimension(self, name): return self.nc.dimensions[name] def __getitem__(self, name): value = getattr(self.nc, name) if isinstance(value, str): try: value = eval(value) except (SyntaxError, NameError): pass return value else: return value[0] def has_array(self, name): return name in self.nc.variables def get(self, name, *indices): var = self.nc.variables[name] if var.shape == (): return var.getValue() else: if var.dimensions[-1] == 'two': x = var[indices] array = np.empty(x.shape[:-1], complex) array.real = x[..., 0] array.imag = x[..., 1] return array else: return var[indices] def get_reference(self, name, *indices): return NetCDFReference(self.nc.variables[name], indices) def close(self): self.nc.close()
import numpy as np from ase.io.pupynere import NetCDFFile # Write array a1 = np.random.rand(5, 5) a2 = a1 * 2 - 5 nc = NetCDFFile('test.nc', 'w') nc.createDimension('dimx', a1.shape[0]) nc.createDimension('dimy', a1.shape[1]) nc.createVariable('matrix1', 'd', ('dimx', 'dimy'))[:] = a1 nc.createVariable('matrix2', 'd', ('dimx', 'dimy'))[:] = a2 nc.sync() nc.close() # Read array nc = NetCDFFile('test.nc', 'r') b1 = nc.variables['matrix1'][:] b2 = nc.variables['matrix2'][:] assert np.all(a1 == b1) and np.all(a2 == b2)
class ETSFWriter: def __init__(self, filename): from Scientific.IO.NetCDF import NetCDFFile self.nc = NetCDFFile(filename, 'w') self.nc.file_format = 'ETSF Nanoquanta' self.nc.file_format_version = np.array([3.3], dtype=np.float32) self.nc.Conventions = 'http://www.etsf.eu/fileformats/' self.nc.history = 'File generated by ASE' def write_atoms(self, atoms): specie_a = np.empty(len(atoms), np.int32) nspecies = 0 species = {} numbers = [] for a, Z in enumerate(atoms.get_atomic_numbers()): if Z not in species: species[Z] = nspecies nspecies += 1 numbers.append(Z) specie_a[a] = species[Z] dimensions = [ ('character_string_length', 80), ('number_of_atoms', len(atoms)), ('number_of_atom_species', nspecies), ('number_of_cartesian_directions', 3), ('number_of_reduced_dimensions', 3), ('number_of_vectors', 3)] for name, size in dimensions: self.nc.createDimension(name, size) var = self.add_variable var('primitive_vectors', ('number_of_vectors', 'number_of_cartesian_directions'), atoms.cell / Bohr, units='atomic units') var('atom_species', ('number_of_atoms',), specie_a + 1) var('reduced_atom_positions', ('number_of_atoms', 'number_of_reduced_dimensions'), atoms.get_scaled_positions()) var('atomic_numbers', ('number_of_atom_species',), np.array(numbers, dtype=float)) def close(self): self.nc.close() def add_variable(self, name, dims, data=None, **kwargs): if data is None: char = 'd' else: if isinstance(data, np.ndarray): char = data.dtype.char elif isinstance(data, float): char = 'd' elif isinstance(data, int): char = 'i' else: char = 'c' var = self.nc.createVariable(name, char, dims) for attr, value in kwargs.items(): setattr(var, attr, value) if data is not None: if len(dims) == 0: var.assignValue(data) else: if char == 'c': if len(dims) == 1: var[:len(data)] = data else: for i, x in enumerate(data): var[i, :len(x)] = x else: var[:] = data return var