/
DLDMD.py
119 lines (103 loc) · 5.16 KB
/
DLDMD.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
"""
Author:
Jay Lago, NIWC/SDSU, 2021
"""
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras.layers import *
class DLDMD(keras.Model):
def __init__(self, hyp_params, **kwargs):
super(DLDMD, self).__init__(**kwargs)
# Parameters
self.batch_size = hyp_params['batch_size']
self.phys_dim = hyp_params['phys_dim']
self.latent_dim = hyp_params['latent_dim']
self.num_time_steps = int(hyp_params['num_time_steps'])
self.num_pred_steps = int(hyp_params['num_pred_steps'])
self.time_final = hyp_params['time_final']
self.num_en_layers = hyp_params['num_en_layers']
self.num_neurons = hyp_params['num_en_neurons']
self.delta_t = hyp_params['delta_t']
self.enc_input = (self.num_time_steps, self.phys_dim)
self.dec_input = (self.num_time_steps, self.latent_dim)
self.precision = hyp_params['precision']
if self.precision == 'float32':
self.precision_complex = tf.complex64
else:
self.precision_complex = tf.complex128
# Construct the ENCODER network
self.encoder = keras.Sequential(name="encoder")
self.encoder.add(Dense(self.num_neurons,
input_shape=self.enc_input,
activation=hyp_params['hidden_activation'],
kernel_initializer=hyp_params['kernel_init_enc'],
bias_initializer=hyp_params['bias_initializer'],
trainable=True, name='enc_in'))
for ii in range(self.num_en_layers):
self.encoder.add(Dense(self.num_neurons,
activation=hyp_params['hidden_activation'],
kernel_initializer=hyp_params['kernel_init_enc'],
bias_initializer=hyp_params['bias_initializer'],
trainable=True, name='enc_' + str(ii)))
self.encoder.add(Dense(self.latent_dim,
activation=hyp_params['ae_output_activation'],
kernel_initializer=hyp_params['kernel_init_enc'],
bias_initializer=hyp_params['bias_initializer'],
trainable=True, name='enc_out'))
# Construct the DECODER network
self.decoder = keras.Sequential(name="decoder")
self.decoder.add(Dense(self.num_neurons,
input_shape=self.dec_input,
activation=hyp_params['hidden_activation'],
kernel_initializer=hyp_params['kernel_init_enc'],
bias_initializer=hyp_params['bias_initializer'],
trainable=True, name='dec_in'))
for ii in range(self.num_en_layers):
self.decoder.add(Dense(self.num_neurons,
activation=hyp_params['hidden_activation'],
kernel_initializer=hyp_params['kernel_init_dec'],
bias_initializer=hyp_params['bias_initializer'],
trainable=True, name='dec_' + str(ii)))
self.decoder.add(Dense(self.phys_dim,
activation=hyp_params['ae_output_activation'],
kernel_initializer=hyp_params['kernel_init_dec'],
bias_initializer=hyp_params['bias_initializer'],
trainable=True, name='dec_out'))
def call(self, x):
# Encode the entire time series
y = self.encoder(x)
x_ae = self.decoder(y)
# Reshape for DMD step
yt = tf.transpose(y, [0, 2, 1])
# Generate latent time series using DMD prediction
y_adv, evals, evecs, phi = self.edmd(yt)
# Decode the latent trajectories
x_adv = self.decoder(y_adv)
# Model weights
weights = self.trainable_weights
return [y, x_ae, x_adv, y_adv, weights, evals, evecs, phi]
def edmd(self, Y):
Y_m = Y[:, :, :-1]
Y_p = Y[:, :, 1:]
sig, U, V = tf.linalg.svd(Y_m, compute_uv=True, full_matrices=False)
sigr_inv = tf.linalg.diag(1.0 / sig)
Uh = tf.linalg.adjoint(U)
A = Y_p @ V @ sigr_inv @ Uh
evals, evecs = tf.linalg.eig(A)
phi = tf.linalg.solve(evecs, tf.cast(Y_m, dtype=self.precision_complex))
y0 = phi[:, :, 0]
y0 = y0[:, :, tf.newaxis]
recon = tf.TensorArray(self.precision_complex, size=self.num_pred_steps)
recon = recon.write(0, evecs @ y0)
evals_k = tf.identity(evals)
for ii in tf.range(1, self.num_pred_steps):
tmp = evecs @ (tf.linalg.diag(evals_k) @ y0)
recon = recon.write(ii, tmp)
evals_k = evals_k * evals
recon = tf.math.real(tf.transpose(tf.squeeze(recon.stack()), perm=[1, 0, 2]))
return recon, evals, evecs, phi
def get_config(self):
base_config = super().get_config()
return {**base_config,
'encoder': self.encoder,
'decoder': self.decoder}