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