1import struct
2
3import numpy as np
4
5# Size of basic types (in bytes)
6SIZE_LOGICAL = 4
7SIZE_INT = 4
8SIZE_DOUBLE = 8
9NAME_LEN = 16
10
11# For un-aligned data, use '=' (for aligned data set to '')
12ALIGN = "="
13
14
15def get_header(istream):
16    """Read header from an MPI-AMRVAC 2.0 snapshot.
17    istream' should be a file
18    opened in binary mode.
19    """
20    istream.seek(0)
21    h = {}
22
23    fmt = ALIGN + "i"
24    [h["datfile_version"]] = struct.unpack(fmt, istream.read(struct.calcsize(fmt)))
25
26    if h["datfile_version"] < 3:
27        raise OSError("Unsupported AMRVAC .dat file version: %d", h["datfile_version"])
28
29    # Read scalar data at beginning of file
30    fmt = ALIGN + 9 * "i" + "d"
31    hdr = struct.unpack(fmt, istream.read(struct.calcsize(fmt)))
32    [
33        h["offset_tree"],
34        h["offset_blocks"],
35        h["nw"],
36        h["ndir"],
37        h["ndim"],
38        h["levmax"],
39        h["nleafs"],
40        h["nparents"],
41        h["it"],
42        h["time"],
43    ] = hdr
44
45    # Read min/max coordinates
46    fmt = ALIGN + h["ndim"] * "d"
47    h["xmin"] = np.array(struct.unpack(fmt, istream.read(struct.calcsize(fmt))))
48    h["xmax"] = np.array(struct.unpack(fmt, istream.read(struct.calcsize(fmt))))
49
50    # Read domain and block size (in number of cells)
51    fmt = ALIGN + h["ndim"] * "i"
52    h["domain_nx"] = np.array(struct.unpack(fmt, istream.read(struct.calcsize(fmt))))
53    h["block_nx"] = np.array(struct.unpack(fmt, istream.read(struct.calcsize(fmt))))
54
55    if h["datfile_version"] >= 5:
56        # Read periodicity
57        fmt = ALIGN + h["ndim"] * "i"  # Fortran logical is 4 byte int
58        h["periodic"] = np.array(
59            struct.unpack(fmt, istream.read(struct.calcsize(fmt))), dtype=bool
60        )
61
62        # Read geometry name
63        fmt = ALIGN + NAME_LEN * "c"
64        hdr = struct.unpack(fmt, istream.read(struct.calcsize(fmt)))
65        h["geometry"] = b"".join(hdr).strip().decode()
66
67        # Read staggered flag
68        fmt = ALIGN + "i"  # Fortran logical is 4 byte int
69        h["staggered"] = bool(struct.unpack(fmt, istream.read(struct.calcsize(fmt)))[0])
70
71    # Read w_names
72    w_names = []
73    for _ in range(h["nw"]):
74        fmt = ALIGN + NAME_LEN * "c"
75        hdr = struct.unpack(fmt, istream.read(struct.calcsize(fmt)))
76        w_names.append(b"".join(hdr).strip().decode())
77    h["w_names"] = w_names
78
79    # Read physics type
80    fmt = ALIGN + NAME_LEN * "c"
81    hdr = struct.unpack(fmt, istream.read(struct.calcsize(fmt)))
82    h["physics_type"] = b"".join(hdr).strip().decode()
83
84    # Read number of physics-defined parameters
85    fmt = ALIGN + "i"
86    [n_pars] = struct.unpack(fmt, istream.read(struct.calcsize(fmt)))
87
88    # First physics-parameter values are given, then their names
89    fmt = ALIGN + n_pars * "d"
90    vals = struct.unpack(fmt, istream.read(struct.calcsize(fmt)))
91
92    fmt = ALIGN + n_pars * NAME_LEN * "c"
93    names = struct.unpack(fmt, istream.read(struct.calcsize(fmt)))
94    # Split and join the name strings (from one character array)
95    names = [
96        b"".join(names[i : i + NAME_LEN]).strip().decode()
97        for i in range(0, len(names), NAME_LEN)
98    ]
99
100    # Store the values corresponding to the names
101    for val, name in zip(vals, names):
102        h[name] = val
103    return h
104
105
106def get_tree_info(istream):
107    """
108    Read levels, morton-curve indices, and byte offsets for each block as stored in the
109    datfile istream is an open datfile buffer with 'rb' mode
110    This can be used as the "first pass" data reading required by YT's interface.
111    """
112    istream.seek(0)
113    header = get_header(istream)
114    nleafs = header["nleafs"]
115    nparents = header["nparents"]
116
117    # Read tree info. Skip 'leaf' array
118    istream.seek(header["offset_tree"] + (nleafs + nparents) * SIZE_LOGICAL)
119
120    # Read block levels
121    fmt = ALIGN + nleafs * "i"
122    block_lvls = np.array(struct.unpack(fmt, istream.read(struct.calcsize(fmt))))
123
124    # Read block indices
125    fmt = ALIGN + nleafs * header["ndim"] * "i"
126    block_ixs = np.reshape(
127        struct.unpack(fmt, istream.read(struct.calcsize(fmt))), [nleafs, header["ndim"]]
128    )
129
130    # Read block offsets (skip ghost cells !)
131    bcfmt = ALIGN + header["ndim"] * "i"
132    bcsize = struct.calcsize(bcfmt) * 2
133
134    fmt = ALIGN + nleafs * "q"
135    block_offsets = (
136        np.array(struct.unpack(fmt, istream.read(struct.calcsize(fmt)))) + bcsize
137    )
138    return block_lvls, block_ixs, block_offsets
139
140
141def get_single_block_data(istream, byte_offset, block_shape):
142    """retrieve a specific block (all fields) from a datfile"""
143    istream.seek(byte_offset)
144    # Read actual data
145    fmt = ALIGN + np.prod(block_shape) * "d"
146    d = struct.unpack(fmt, istream.read(struct.calcsize(fmt)))
147    # Fortran ordering
148    block_data = np.reshape(d, block_shape, order="F")
149    return block_data
150
151
152def get_single_block_field_data(istream, byte_offset, block_shape, field_idx):
153    """retrieve a specific block (ONE field) from a datfile"""
154    istream.seek(byte_offset)
155
156    # compute byte size of a single field
157    field_shape = block_shape[:-1]
158    fmt = ALIGN + np.prod(field_shape) * "d"
159    byte_size_field = struct.calcsize(fmt)
160
161    # Read actual data
162    istream.seek(byte_size_field * field_idx, 1)  # seek forward
163    d = struct.unpack(fmt, istream.read(struct.calcsize(fmt)))
164
165    # Fortran ordering
166    block_field_data = np.reshape(d, field_shape, order="F")
167    return block_field_data
168