1"""
2Provides standard interfaces to various text and binary file formats for saving
3position and connectivity data.  Note that saving spikes, membrane potential
4and synaptic conductances is now done via Neo.
5
6Classes:
7    StandardTextFile
8    PickleFile
9    NumpyBinaryFile
10    HDF5ArrayFile - requires PyTables
11
12:copyright: Copyright 2006-2021 by the PyNN team, see AUTHORS.
13:license: CeCILL, see LICENSE for details.
14
15"""
16
17import numpy as np
18import os
19import shutil
20import pickle
21
22try:
23    import tables
24    have_hdf5 = True
25except ImportError:
26    have_hdf5 = False
27
28DEFAULT_BUFFER_SIZE = 10000
29
30
31def _savetxt(filename, data, format, delimiter):
32    """
33    Due to the lack of savetxt in older versions of numpy
34    we provide a cut-down version of that function.
35    """
36    f = open(filename, 'w')
37    for row in data:
38        f.write(delimiter.join([format % val for val in row]) + '\n')
39    f.close()
40
41
42def savez(file, *args, **kwds):
43
44    __doc__ = np.savez.__doc__
45    import zipfile
46    from numpy.lib import format
47
48    if isinstance(file, str):
49        if not file.endswith('.npz'):
50            file = file + '.npz'
51
52    namedict = kwds
53    for i, val in enumerate(args):
54        key = 'arr_%d' % i
55        if key in namedict.keys():
56            raise ValueError("Cannot use un-named variables and keyword %s" % key)
57        namedict[key] = val
58
59    zip = zipfile.ZipFile(file, mode="w")
60
61    # Place to write temporary .npy files
62    #  before storing them in the zip. We need to path this to have a working
63    # function in parallel !
64    import tempfile
65    direc = tempfile.mkdtemp()
66    for key, val in namedict.items():
67        fname = key + '.npy'
68        filename = os.path.join(direc, fname)
69        fid = open(filename, 'wb')
70        format.write_array(fid, np.asanyarray(val))
71        fid.close()
72        zip.write(filename, arcname=fname)
73    zip.close()
74    shutil.rmtree(direc)
75
76
77class BaseFile(object):
78    """
79    Base class for PyNN File classes.
80    """
81
82    def __init__(self, filename, mode='rb'):
83        """
84        Open a file with the given filename and mode.
85        """
86        self.name = filename
87        self.mode = mode
88        dir = os.path.dirname(filename)
89        if dir and not os.path.exists(dir):
90            try:  # wrapping in try...except block for MPI
91                os.makedirs(dir)
92            except IOError:
93                pass  # we assume that the directory was already created by another MPI node
94        try:  # Need this because in parallel, file names are changed
95            self.fileobj = open(self.name, mode, DEFAULT_BUFFER_SIZE)
96        except Exception as err:
97            self.open_error = err
98
99    def __del__(self):
100        self.close()
101
102    def _check_open(self):
103        if not hasattr(self, 'fileobj'):
104            raise self.open_error
105
106    def rename(self, filename):
107        self.close()
108        try:  # Need this because in parallel, only one node will delete the file with NFS
109            os.remove(self.name)
110        except Exception:
111            pass
112        self.name = filename
113        self.fileobj = open(self.name, self.mode, DEFAULT_BUFFER_SIZE)
114
115    def write(self, data, metadata):
116        """
117        Write data and metadata to file. `data` should be a NumPy array,
118        `metadata` should be a dictionary.
119        """
120        raise NotImplementedError
121
122    def read(self):
123        """
124        Read data from the file and return a NumPy array.
125        """
126        raise NotImplementedError
127
128    def get_metadata(self):
129        """
130        Read metadata from the file and return a dict.
131        """
132        raise NotImplementedError
133
134    def close(self):
135        """Close the file."""
136        if hasattr(self, 'fileobj'):
137            self.fileobj.close()
138
139
140class StandardTextFile(BaseFile):
141    """
142    Data and metadata is written as text. Metadata is written at the top of the
143    file, with each line preceded by "#". Data is written with one data point per line.
144    """
145
146    def write(self, data, metadata):
147        __doc__ = BaseFile.write.__doc__
148        self._check_open()
149        # can we write to the file more than once? In this case, should use seek,tell
150        # to always put the header information at the top?
151        # write header
152        header_lines = ["# %s = %s" % item for item in metadata.items()]
153        header = "\n".join(header_lines) + '\n'
154        self.fileobj.write(header.encode('utf-8'))
155        # write data
156        savetxt = getattr(np, 'savetxt', _savetxt)
157        savetxt(self.fileobj, data, fmt='%r', delimiter='\t')
158        self.fileobj.close()
159
160    def read(self):
161        self._check_open()
162        return np.loadtxt(self.fileobj)
163
164    def get_metadata(self):
165        self._check_open()
166        D = {}
167        for line in self.fileobj:
168            if line:
169                if line[0] != "#":
170                    break
171                name, value = line[1:].split("=")
172                name = name.strip()
173                value = eval(value)
174                if type(value) in [list, tuple]:
175                    D[name] = value
176                else:
177                    raise TypeError("Column headers must be specified using a list or tuple.")
178            else:
179                break
180        self.fileobj.seek(0)
181        return D
182
183
184class PickleFile(BaseFile):
185    """
186    Data and metadata are pickled and saved to file.
187    """
188
189    def write(self, data, metadata):
190        __doc__ = BaseFile.write.__doc__
191        self._check_open()
192        pickle.dump((data, metadata), self.fileobj)
193
194    def read(self):
195        __doc__ = BaseFile.read.__doc__
196        self._check_open()
197        data = pickle.load(self.fileobj)[0]
198        self.fileobj.seek(0)
199        return data
200
201    def get_metadata(self):
202        __doc__ = BaseFile.get_metadata.__doc__
203        self._check_open()
204        metadata = pickle.load(self.fileobj)[1]
205        self.fileobj.seek(0)
206        return metadata
207
208
209class NumpyBinaryFile(BaseFile):
210    """
211    Data and metadata are saved in .npz format, which is a zipped archive of
212    arrays.
213    """
214
215    def write(self, data, metadata):
216        __doc__ = BaseFile.write.__doc__
217        self._check_open()
218        metadata_array = np.array(list(metadata.items()), dtype=object)
219        savez(self.fileobj, data=data, metadata=metadata_array)
220
221    def read(self):
222        __doc__ = BaseFile.read.__doc__
223        self._check_open()
224        data = np.load(self.fileobj)['data']
225        self.fileobj.seek(0)
226        return data
227
228    def get_metadata(self):
229        __doc__ = BaseFile.get_metadata.__doc__
230        self._check_open()
231        D = {}
232        for name, value in np.load(self.fileobj, allow_pickle=True)['metadata']:
233            try:
234                D[name] = eval(value)
235            except Exception:
236                D[name] = value
237        self.fileobj.seek(0)
238        return D
239
240
241if have_hdf5:
242    class HDF5ArrayFile(BaseFile):
243        """
244        Data are saved as an array within a node named "data". Metadata are
245        saved as attributes of this node.
246        """
247
248        def __init__(self, filename, mode='r', title="PyNN data file"):
249            """
250            Open an HDF5 file with the given filename, mode and title.
251            """
252            self.name = filename
253            self.mode = mode
254            try:
255                self.fileobj = tables.open_file(filename, mode=mode, title=title)
256                self._new_pytables = True
257            except AttributeError:
258                self.fileobj = tables.openFile(filename, mode=mode, title=title)
259                self._new_pytables = False
260
261        # may not work with old versions of PyTables < 1.3, since they only support numarray, not numpy
262        def write(self, data, metadata):
263            __doc__ = BaseFile.write.__doc__
264            if len(data) > 0:
265                try:
266                    if self._new_pytables:
267                        node = self.fileobj.create_array(self.fileobj.root, "data", data)
268                    else:
269                        node = self.fileobj.createArray(self.fileobj.root, "data", data)
270                except tables.HDF5ExtError as e:
271                    raise tables.HDF5ExtError("%s. data.shape=%s, metadata=%s" %
272                                              (e, data.shape, metadata))
273                for name, value in metadata.items():
274                    setattr(node.attrs, name, value)
275                self.fileobj.close()
276
277        def read(self):
278            __doc__ = BaseFile.read.__doc__
279            return self.fileobj.root.data.read()
280
281        def get_metadata(self):
282            __doc__ = BaseFile.get_metadata.__doc__
283            D = {}
284            node = self.fileobj.root.data
285            for name in node._v_attrs._f_list():
286                D[name] = node.attrs.__getattr__(name)
287            return D
288