def __init__(self, source, frame=None, range=None, start=0, stop=None, step=1, no_compute_index=False,
zero=False, one_frame_per_file=False, indices=None, string=False, format=None):
if isinstance(source, basestring):
self.opened = True
self.source = CInOutput(source, action=INPUT, append=False, zero=zero, range=range,
no_compute_index=no_compute_index, one_frame_per_file=one_frame_per_file,
indices=indices, string=string)
self.netcdf_file = None
if self.source.string is None and source.endswith('.nc'):
try:
self.netcdf_file = netcdf_file(source)
except AssertionError:
pass
else:
self.opened = False
self.source = source
self.netcdf_file = None
if frame is not None:
self.start = frame
self.stop = frame+1
self.step = 1
else:
self.start = start
if stop is None:
try:
stop = len(self.source)
except AttributeError:
stop = None
self.stop = stop
self.step = step
def __init__(self, dest, frame=None, append=False, netcdf4=True):
self.opened = False
self.frame = frame
if isinstance(dest, str):
self.opened = True
try:
FORMAT = {True: 'NETCDF4',
False: 'NETCDF3_CLASSIC'}
MODE = {True: 'a',
False: 'w'}
self.dest = netcdf_file(dest, MODE[append], format=FORMAT[netcdf4])
except ValueError:
self.dest = netcdf_file(dest, 'w')
else:
self.dest = dest
def NetCDFReader(source, frame=None, start=0, stop=None, step=1, format=None):
opened = False
if isinstance(source, str):
opened = True
source = netcdf_file(source)
from quippy import Atoms
DEG_TO_RAD = pi/180.0
remap_names = {'coordinates': 'pos',
'velocities': 'velo',
'cell_lengths': None,
'cell_angles': None,
'cell_lattice': None,
'cell_rotated': None}
prop_type_to_value = {PROPERTY_INT: 0,
PROPERTY_REAL: 0.0,
PROPERTY_STR: "",
PROPERTY_LOGICAL: False}
prop_dim_to_ncols = {('frame','atom','spatial'): 3,
('frame','atom','label'): 1,
('frame', 'atom'): 1}
if frame is not None:
start = frame
stop = frame+1
step = 1
else:
if stop is None:
stop = source.variables['cell_lengths'].shape[0]
for frame in range(start, stop, step):
cl = source.variables['cell_lengths'][frame]
ca = source.variables['cell_angles'][frame]
lattice = make_lattice(cl[0],cl[1],cl[2],ca[0]*DEG_TO_RAD,ca[1]*DEG_TO_RAD,ca[2]*DEG_TO_RAD)
at = Atoms(n=netcdf_dimlen(source, 'atom'), lattice=lattice, properties={})
for name, var in source.variables.iteritems():
name = remap_names.get(name, name)
if name is None:
continue
name = str(name) # in case it's a unicode string
if 'frame' in var.dimensions:
if 'atom' in var.dimensions:
# It's a property
value = var[frame]
if value.dtype.kind != 'S': value = value.T
at.add_property(name, value)
else:
# It's a param
if var.dimensions == ('frame','string'):
# if it's a single string, join it and strip it
at.params[name] = ''.join(var[frame]).strip()
else:
if name == 'cutoff':
at.cutoff = var[frame]
elif name == 'cutoff_skin':
at.cutoff_skin = var[frame]
elif name == 'nneightol':
at.nneightol = var[frame]
elif name == 'pbc':
at.pbc = var[frame]
else:
at.params[name] = var[frame].T
if 'cell_rotated' in source.variables:
cell_rotated = source.variables['cell_rotated'][frame]
orig_lattice = source.variables['cell_lattice'][frame]
#if cell_rotated == 1:
at.set_lattice(orig_lattice, True)
yield at
def close(self):
if self.opened: source.close()