-
Notifications
You must be signed in to change notification settings - Fork 4
/
Copy pathflusi_tools.py
175 lines (130 loc) · 5.3 KB
/
flusi_tools.py
1
2
3
4
5
6
7
8
9
10
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
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Tue Mar 5 15:31:42 2019
@author: engels
"""
import numpy as np
def get_dset_name( fname ):
from os.path import basename
dset_name = basename(fname)
dset_name = dset_name[0:dset_name.find('_')]
return dset_name
def get_timestamp_name( fname ):
from os.path import basename
dset_name = basename(fname)
dset_name = dset_name[dset_name.find('_')+1:dset_name.find('.')]
return dset_name
def read_flusi_HDF5( fname, dtype=np.float32, twoD=False):
""" Read HDF5 file generated by FLUSI.
Returns: time, box, origin, data
"""
import h5py
f = h5py.File(fname, 'r')
# list all hdf5 datasets in the file - usually, we expect
# to find only one.
datasets = list(f.keys())
# if we find more than one dset we warn that this is unusual
if (len(datasets) != 1):
raise ValueError("we found more than one dset in the file (problemo)"+fname)
else:
# as there should be only one, this should be our dataset:
dset_name = datasets[0]
# get the dataset handle
dset_id = f.get(dset_name)
# from the dset handle, read the attributes
time = dset_id.attrs.get('time')
res = dset_id.attrs.get('nxyz')
box = dset_id.attrs.get('domain_size')
origin = dset_id.attrs.get('origin')
if origin is None:
origin = np.array([0,0,0])
b = f[dset_name][:]
data = np.array(b, dtype=dtype)
if len(data.shape) == 3:
# its a funny flusi convention that we have to swap axes here, and I
# never understood why it is this way.
# NB: this applies to 3D data only (even though if running in 2D mode,
# flusi stores 3d array with 1 index length 1, but other softwares
# such as UP2D produce real 2D arrays)
data = np.swapaxes(data, 0, 2)
if (np.max(res-data.shape)>0):
print('WARNING!!!!!!')
print('read_flusi_HDF5: array dimensions look funny')
f.close()
if twoD and data.shape[0] == 1:
data = data[0,:,:].copy()
data = data.transpose()
print("We read FLUSI file %s at time=%f \nResolution :" % (fname, time), data.shape)
return time, box, origin, data
def write_flusi_HDF5( fname, time, box, data, viscosity=0.0, origin=np.array([0.0,0.0,0.0]), dtype=np.float32 ):
import h5py
dset_name = get_dset_name( fname )
if len(data.shape)==3:
#3d data
nx, ny, nz = data.shape
print( "Writing to file=%s dset=%s max=%e min=%e size=%i %i %i " % (fname, dset_name, np.max(data), np.min(data), nx,ny,nz) )
# i dont really know why, but there is a messup in fortran vs c ordering, so here we have to swap
# axis
data = np.swapaxes(data, 0, 2)
nxyz = np.array([nz,nx,ny])
else:
#2d data
nx, ny = data.shape
print( "Writing to file=%s dset=%s max=%e min=%e size=%i %i" % (fname, dset_name, np.max(data), np.min(data), nx,ny) )
data = np.swapaxes(data, 0, 1)
nxyz = np.array([1, nx,ny])
fid = h5py.File( fname, 'w')
fid.create_dataset( dset_name, data=data, dtype=dtype )#, shape=data.shape[::-1] )
fid.close()
fid = h5py.File(fname,'a')
dset_id = fid.get( dset_name )
dset_id.attrs.create('time', time)
dset_id.attrs.create('viscosity', viscosity)
dset_id.attrs.create('domain_size', box )
dset_id.attrs.create('origin', origin )
dset_id.attrs.create('nxyz', nxyz )
fid.close()
def crop_flusi_HDF5(file, Nxcut=[0, 0], Nycut=[0, 0], Nzcut=[0, 0]):
"""
Crop the data matrix.
Input:
N[x|y|t]_cut: 1d-array of size 2
Array will be croped from Nx_cut[0]:-Nx_cut[1] in x dimension,
Ny_cut[0]:-Ny_cut[1] in y dimension,
Nz_cut[0]:-Nz_cut[1] in z dimension
"""
time, box, origin, data = read_flusi_HDF5( file )
data = np.squeeze(data)
if len(data.shape)==2:
# this is the lazy variant:
y = np.linspace(0,box[-1],np.size(data,-1)) + origin[-1]
x = np.linspace(0,box[-2],np.size(data,-2)) + origin[-2]
x_cut= x[Nxcut[0]:-Nxcut[1]] - x[Nxcut[0]]
y_cut= y[Nycut[0]:-Nycut[1]] - y[Nycut[0]]
box[-2] = max(x_cut)- min(x_cut)
box[-1] = max(y_cut)- min(y_cut)
origin[-2] = x_cut[0]
origin[-1] = y_cut[0]
data_cut = data[ Nxcut[0] : -Nxcut[1], Nycut[0] : -Nycut[1]]
data = np.expand_dims(data_cut,2) # we have to add the z dimension again
else:
print("crop_flusi_hdf5: 3d not implemented")
return
write_flusi_HDF5( file, time, box, data, origin )
def resample_flusi_HDF5(file, N):
"""
Sample the data matrix up
Input:
shape
"""
import fourier_tools
time, box, origin, data = read_flusi_HDF5( file )
data = np.squeeze(data)
if len(data.shape)==2:
data = fourier_tools.fft2_resample( data, N )
data = np.expand_dims(data,2) # we have to add the z dimension again
else:
print("crop_flusi_hdf5: 3d not implemented")
return
write_flusi_HDF5( file, time, box, data, origin )