Source code for yggdrasil.communication.NetCDFFileComm

import sys
import pprint
import numpy as np
try:
    from scipy.io import _netcdf, netcdf_file
    REVERSE = _netcdf.REVERSE  # pragma: debug
except ImportError:
    from scipy.io import netcdf
    REVERSE = netcdf.REVERSE
    netcdf_file = netcdf.netcdf_file
from yggdrasil import units
from yggdrasil.communication.DedicatedFileBase import DedicatedFileBase


[docs]class NetCDFFileComm(DedicatedFileBase): r"""Class for handling I/O from/to an netCDF file. Args: read_attributes (bool, optional): If True, the attributes are read in as well as the variables. Defaults to False. variables (list, optional): List of variables to read in. If not provided, all variables will be read. version (int, optional): Version of netCDF format that should be used. Defaults to 1. Options are 1 (classic format) and 2 (64-bit offset format). **kwargs: Additional keywords arguments are passed to parent class. """ _filetype = 'netcdf' _schema_subtype_description = ('The file is read/written as netCDF.') _schema_properties = { 'read_attributes': {'type': 'boolean', 'default': False}, 'variables': {'type': 'array', 'items': {'type': 'string'}}, 'version': {'type': 'integer', 'enum': [1, 2], 'default': 1}} _extensions = ['.nc'] _synchronous_read = True _stores_fd = True
[docs] def transform_type_recv(self, x): x_dtype = np.dtype(x.dtype) typecode, size = x_dtype.char, x_dtype.itemsize if (typecode == 'c') and (x.ndim > 1): size = x.shape[0] new_shape = x.shape[1:] x_str = np.zeros(new_shape, 'S%d' % size) for index in np.ndindex(*new_shape): for i in range(size): x_str[index] += x[tuple([i, *index])] x = x_str else: x = x.copy() if ((((sys.byteorder == 'little') and (x.dtype.byteorder == '>')) or ((sys.byteorder == 'big') and (x.dtype.byteorder == '<')))): x = x.astype(x.dtype.name) return x
[docs] def transform_type_send(self, x): x_dtype = np.dtype(x.dtype) typecode, size = x_dtype.char, x_dtype.itemsize typecode_map = {'l': 'i', 'q': 'i'} if (typecode, size) not in REVERSE: REVERSE_keys = list(REVERSE.keys()) REVERSE_typecode = [k[0] for k in REVERSE_keys] typecode = typecode_map.get(typecode, typecode) if typecode == 'S': x_str = np.zeros(tuple([size] + list(x.shape)), 'S1') for index in np.ndindex(*x.shape): for i in range(size): if i == len(x[index]): x_str[tuple([i, *index])] = '\0' elif i < len(x[index]): x_str[tuple([i, *index])] = x[index][i:(i + 1)] x = x_str elif typecode in REVERSE_typecode: x = x.astype(np.dtype(*REVERSE_keys[ REVERSE_typecode.index(typecode)])) else: # pragma: debug raise RuntimeError( ("Type (%s, %d) is not in set accepted by " "netCDF %s.") % (typecode, size, pprint.pformat(REVERSE))) return x
@property def fd(self): r"""Associated file identifier.""" if self._external_fd: return self._external_fd.fp return super(NetCDFFileComm, self).fd def _dedicated_open(self, address, mode): self._external_fd = netcdf_file(address, mode, mmap=True, version=self.version) return self._external_fd.fp def _dedicated_close(self): self._external_fd.fp.close() self._external_fd.close() self._external_fd = None def _dedicated_send(self, msg): assert isinstance(msg, dict) for k, v in msg.items(): if isinstance(v, np.ndarray): dims = [] v = self.transform_type_send(v) for i, d in enumerate(v.shape): if i == 0: idim = k else: idim = '%s%d' % (k, i) dims.append(idim) self._external_fd.createDimension(idim, d) var = self._external_fd.createVariable(k, v.dtype, dims) var[:] = v if units.has_units(v): var.units = units.get_units(v) else: # pragma: debug raise TypeError("Type '%s' no supported." % type(msg)) def _dedicated_recv(self): out = {} variables = self.variables if not variables: variables = list(self._external_fd.variables.keys()) for v in variables: out[v] = self.transform_type_recv( self._external_fd.variables[v][:]) if hasattr(self._external_fd.variables[v], 'units'): out[v] = units.add_units( out[v], self._external_fd.variables[v].units) return out
[docs] @classmethod def get_testing_options(cls, **kwargs): r"""Method to return a dictionary of testing options for this class. Returns: dict: Dictionary of variables to use for testing. Items: kwargs (dict): Keyword arguments for comms tested with the provided content. send (list): List of objects to send to test file. recv (list): List of objects that will be received from a test file that was sent the messages in 'send'. contents (bytes): Bytes contents of test file created by sending the messages in 'send'. """ data = { 'time': units.add_units(np.arange(10).astype('float32'), 's'), 'x': np.array(['a', 'hello', 'c'], 'S5')} data_add = { 'space': units.add_units(np.ones((5, 5), 'int64'), 'mol')} out = {'kwargs': {}, 'exact_contents': False, 'msg': data, 'dict': False, 'objects': [data, data_add], 'send': [data, data_add], 'recv': [dict(data, **data_add)], 'recv_partial': [[data], [dict(data, **data_add)]], 'contents': ( b'CDF\x01\x00\x00\x00\x00\x00\x00\x00\n\x00\x00' b'\x00\x05\x00\x00\x00\x04time\x00\x00\x00\n\x00' b'\x00\x00\x01x\x00\x00\x00\x00\x00\x00\x05\x00' b'\x00\x00\x02x1\x00\x00\x00\x00\x00\x03\x00\x00' b'\x00\x05space\x00\x00\x00\x00\x00\x00\x05\x00' b'\x00\x00\x06space1\x00\x00\x00\x00\x00\x05\x00' b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x0b\x00' b'\x00\x00\x03\x00\x00\x00\x04time\x00\x00\x00\x01' b'\x00\x00\x00\x00\x00\x00\x00\x0c\x00\x00\x00\x01' b'\x00\x00\x00\x05units\x00\x00\x00\x00\x00\x00' b'\x02\x00\x00\x00\x01s\x00\x00\x00\x00\x00\x00' b'\x05\x00\x00\x00(\x00\x00\x01\x0c\x00\x00\x00' b'\x05space\x00\x00\x00\x00\x00\x00\x02\x00\x00' b'\x00\x03\x00\x00\x00\x04\x00\x00\x00\x0c\x00\x00' b'\x00\x01\x00\x00\x00\x05units\x00\x00\x00\x00' b'\x00\x00\x02\x00\x00\x00\x03mol\x00\x00\x00\x00' b'\x04\x00\x00\x00d\x00\x00\x014\x00\x00\x00\x01x' b'\x00\x00\x00\x00\x00\x00\x02\x00\x00\x00\x01\x00' b'\x00\x00\x02\x00\x00\x00\x00\x00\x00\x00\x00\x00' b'\x00\x00\x02\x00\x00\x00\x10\x00\x00\x01\x98\x00' b'\x00\x00\x00?\x80\x00\x00@\x00\x00\x00@@\x00\x00' b'@\x80\x00\x00@\xa0\x00\x00@\xc0\x00\x00@\xe0\x00' b'\x00A\x00\x00\x00A\x10\x00\x00\x00\x00\x00\x01' b'\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x01' b'\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x01' b'\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x01' b'\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x01' b'\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x01' b'\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x01' b'\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x01' b'\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x01' b'ahc\x00e\x00\x00l\x00\x00l\x00\x00o\x00\x00')} return out