stable-NN-ROM/data_utils.py at main · lbalicki/stable-NN-ROM · GitHub

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
import json
import pyvista as pv
import os
import re

import scipy.io as spio
import numpy as np

def get_meta_data(data_str='data'):
    # Get directory names and extract Re, u1 and u2 values for the setup.
    data_dirs = [entry.name for entry in os.scandir(data_str) if entry.is_dir()]

    combined_data = []

    pattern = r"Re_([0-9.]+)_u1_([0-9.]+)_u2_([0-9.]+)"

    for d in data_dirs:
        match = re.match(pattern, d)
        if match:
            Re, u1, u2 = match.groups()
            combined_data.append([d, float(Re), float(u1), float(u2)])

    return combined_data


def get_v_data(data_str='data', include_meta_data=True):
    Re_vals = []
    u1_vals = []
    u2_vals = []
    v_blocks = []

    meta_data = get_meta_data(data_str=data_str)

    # add transient simulation data
    for md in meta_data:
        v_file = os.path.join(data_str, md[0], 'v_mat.npy')
        v_blocks.append(
            np.load(v_file)
        )
        Re_vals.append(md[1])
        u1_vals.append(md[2])
        u2_vals.append(md[3])

    # add steady state data
    for md in meta_data:
        v_file = os.path.join(data_str, md[0], 'ss_nse_sol')
        v_ss = spio.loadmat(v_file)['ss_nse_v']
        v_blocks.append(
            np.concatenate((v_ss, v_ss), axis=1)
        )
        Re_vals.append(md[1])
        u1_vals.append(0)
        u2_vals.append(0)

    if include_meta_data:
        Re_vals = np.array(Re_vals)
        u1_vals = np.array(u1_vals)
        u2_vals = np.array(u2_vals)
        return v_blocks, Re_vals, u1_vals, u2_vals

    return v_blocks


def normalize_v_data(v_blocks):
    v_matrix = np.concatenate(v_blocks, axis=1)
    v_mean = np.mean(v_matrix, axis=1)
    normalized_v_blocks = []
    for v_block in v_blocks:
        normalized_v_blocks.append(v_block - v_mean[:,None])

    return normalized_v_blocks, v_mean


def read_v_data(setup, data_str='data'):
    # read velocity data for a specific setup
    v_dir = os.path.join(data_str, setup, 'timesteps')

    v_files = [f for f in os.listdir(v_dir) if f.startswith('v')]
    v_x = []
    v_y = []
    t = []

    # Print the filenames
    for f in v_files:
        v_data = pv.read(os.path.join(v_dir, f))
        v_x.append(v_data.point_data['v'][:,0])
        v_y.append(v_data.point_data['v'][:,1])

        base = os.path.splitext(f)[0]
        t_str = base.split('_t')[-1]
        t.append(float(t_str))

    return np.array(t), np.array(v_x).T, np.array(v_y).T


def save_v_paraview(v=None, strtojson=None, visudict=None, vfile='vel.vtu'):
    if visudict is None:
        jsfile = open(strtojson)
        visudict = json.load(jsfile)
        vaux = np.zeros((visudict['vdim'], 1))
        for bcdict in visudict['bclist']:
            intbcidx = [int(bci) for bci in bcdict.keys()]
            vaux[intbcidx, 0] = list(bcdict.values())

    n_dof = len(visudict['vxvtxdofs'])

    with open(vfile, 'w') as velfile:
        velfile.write(visudict['vtuheader_v'])
        for j in range(n_dof):
            velfile.write(u'{0} {1} {2} '.format(v[j], v[n_dof+j], 0.))
        velfile.write(visudict['vtufooter_v'])