tar.py 2.13 KB
Newer Older
1
from __future__ import print_function
2
3
from builtins import str
from builtins import object
Jens J. Mortensen's avatar
Jens J. Mortensen committed
4
5
6
7
8
9
import tarfile
import xml.etree.ElementTree as ElementTree

import numpy as np


10
class Reader(object):
Jens J. Mortensen's avatar
Jens J. Mortensen committed
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
    def __init__(self, name):
        self.dims = {}
        self.shapes = {}
        self.dtypes = {}
        self.parameters = {}
        self.tar = tarfile.open(name, 'r')
        f = self.tar.extractfile('info.xml').read()
        root = ElementTree.fromstring(f)
        self.byteswap = ((root.attrib['endianness'] == 'little') !=
                         np.little_endian)
        for child in root:
            a = child.attrib
            if child.tag == 'parameter':
                try:
                    value = eval(a['value'], {})
                except (SyntaxError, NameError):
                    value = str(a['value'])
                self.parameters[a['name']] = value
            elif child.tag == 'array':
                name = a['name']
                self.dtypes[name] = a['type']
                self.shapes[name] = []
                for dim in child:
                    n = int(dim.attrib['length'])
                    self.shapes[name].append(n)
                    self.dims[dim.attrib['name']] = n

Jens J. Mortensen's avatar
Jens J. Mortensen committed
38
39
    def __contains__(self, name):
        return name in self.parameters or name in self.shapes
40

Jens J. Mortensen's avatar
Jens J. Mortensen committed
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
    def __getattr__(self, name):
        if name in self.parameters:
            return self.parameters[name]
        else:
            return self.get_array(name)

    def get_array(self, name):
        fileobj = self.tar.extractfile(name)
        shape = self.shapes[name]
        dtype = self.dtypes[name]
        dtype = np.dtype({'int': 'int32'}.get(dtype, dtype))
        size = np.prod(shape) * dtype.itemsize
        array = np.fromstring(fileobj.read(size), dtype)
        if self.byteswap:
            array = array.byteswap()
        if dtype == np.int32:
            array = np.asarray(array, int)
        array.shape = shape
        return array
60

Jens J. Mortensen's avatar
Jens J. Mortensen committed
61
62
63
    def close(self):
        self.tar.close()

64

Jens J. Mortensen's avatar
Jens J. Mortensen committed
65
66
67
if __name__ == '__main__':
    r = Reader('H2.gpw')
    print(r.Eigenvalues, r.FermiLevel, r.UnitCell)
Jens J. Mortensen's avatar
Jens J. Mortensen committed
68
    print('UnitCell' in r, 'sdgsdgsdfg' in r)