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()
Exemplo n.º 3
0
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)
Exemplo n.º 4
0
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