lu-group / mionet Goto Github PK

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import itertools

import numpy as np
from deepxde.backend import tf
import os
from scipy.integrate import solve_ivp
import deepxde as dde
from spaces import FinitePowerSeries, FiniteChebyshev, GRF
from utils import merge_values, trim_to_65535, mean_squared_error_outlier, safe_test

from deepxde.nn.tensorflow_compat_v1.mionet import MIONetCartesianProd
from deepxde.data.quadruple import QuadrupleCartesianProd

def int_index(x, t, T):
    mat = np.linspace(0, T, x)
    return int(t / mat[1])

def ode(m, T, sensor_values1, sensor_values2):
    """ODE system"""
    s0 = [0, 0]  # initial condition

    def model(t, s):
        k = 1
        u1 = lambda t: sensor_values1[t]
        u2 = lambda t: sensor_values2[t]
        return [
            s[1] + u1(int_index(m, t, T)),
            -k * np.sin(s[0]) + u2(int_index(m, t, T)),
        ]

    res = solve_ivp(model, [0, T], s0, method="RK45", t_eval=np.linspace(0, T, m), vectorized = True)
    return res.y[0, :], res.y[1, :]

def network(problem, m):
    if problem == "ODE":
        branch = [m, 200, 200]
        trunk = [1, 200, 200]
    elif problem == "DR":
        branch = [m, 200, 200]
        trunk = [2, 200, 200]
    elif problem == "ADVD":
        branch = [m, 300, 300, 300]
        trunk = [2, 300, 300, 300]
    return branch, trunk

problem = "ODE"
T = 1
m = 100
lr = 0.0002 if problem in ["ADVD"] else 0.001
epochs = 100000
activation = (
    ["relu", None, "relu"] if problem in ["ADVD"] else ["relu", "relu", "relu"]
)
initializer = "Glorot normal"

training_data = np.load("../data/" + problem + "_train_1.npz", allow_pickle=True)
testing_data = np.load("../data/" + problem + "_test_1.npz", allow_pickle=True)

X_train = training_data["X_train"]
y_train = training_data["y_train"]
X_test = testing_data["X_test"]
y_test = testing_data["y_test"]


branch_net, trunk_net = network(problem, m)

data = QuadrupleCartesianProd(X_train, y_train, X_test, y_test)

net = MIONetCartesianProd(
    branch_net,
    branch_net,
    trunk_net,
    {"branch1": activation[0], "branch2": activation[1], "trunk": activation[2]},
    initializer,
    regularization=None,
)

model = dde.Model(data, net)
model.compile("adam", lr=lr)
checker = dde.callbacks.ModelCheckpoint(
    "model/mionet_model", save_better_only=True, period=1000
)

model.restore(os.path.normpath("model/mionet_model-96000.ckpt"), verbose=1)

model.predict(X_test)