diff --git a/requirements.txt b/requirements.txt index 6da6a87..31ef45c 100644 Binary files a/requirements.txt and b/requirements.txt differ diff --git a/src/deeponet_pde.py b/src/deeponet_pde.py index 61259bc..09205f2 100644 --- a/src/deeponet_pde.py +++ b/src/deeponet_pde.py @@ -1,16 +1,13 @@ -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - +import argparse import itertools - +import deepxde as dde import numpy as np import tensorflow as tf -import deepxde as dde -from spaces import FinitePowerSeries, FiniteChebyshev, GRF -from system import LTSystem, ODESystem, DRSystem, CVCSystem, ADVDSystem -from utils import merge_values, trim_to_65535, mean_squared_error_outlier, safe_test +from spaces import GRF, FiniteChebyshev, FinitePowerSeries +from system import ADVDSystem, CVCSystem, DRSystem, LTSystem, ODESystem +from utils import (mean_squared_error_outlier, merge_values, safe_test, + trim_to_65535) def test_u_lt(nn, system, T, m, model, data, u, fname): @@ -21,9 +18,10 @@ def test_u_lt(nn, system, T, m, model, data, u, fname): ns = np.arange(system.npoints_output)[:, None] X_test = [np.tile(sensor_value, (system.npoints_output, 1)), ns] y_test = s - if nn != "opnn": + if nn != "deeponet": X_test = merge_values(X_test) - y_pred = model.predict(data.transform_inputs(X_test)) + X_test = tuple([arr.astype(dtype=np.float32) for arr in X_test]) + y_pred = model.predict(X_test) np.savetxt("test/u_" + fname, sensor_value) np.savetxt("test/s_" + fname, np.hstack((ns, y_test, y_pred))) @@ -35,9 +33,11 @@ def test_u_ode(nn, system, T, m, model, data, u, fname, num=100): x = np.linspace(0, T, num=num)[:, None] X_test = [np.tile(sensor_values.T, (num, 1)), x] y_test = system.eval_s_func(u, x) - if nn != "opnn": + if nn != "deeponet": X_test = merge_values(X_test) - y_pred = model.predict(data.transform_inputs(X_test)) + + X_test = tuple([arr.astype(dtype=np.float32) for arr in X_test]) + y_pred = model.predict(X_test) np.savetxt(fname, np.hstack((x, y_test, y_pred))) print("L2relative error:", dde.metrics.l2_relative_error(y_test, y_pred)) @@ -51,9 +51,11 @@ def test_u_dr(nn, system, T, m, model, data, u, fname): xt = xt * [1 / (m - 1), T / (system.Nt - 1)] X_test = [np.tile(sensor_value, (m * system.Nt, 1)), xt] y_test = s.reshape([m * system.Nt, 1]) - if nn != "opnn": + if nn != "deeponet": X_test = merge_values(X_test) - y_pred = model.predict(data.transform_inputs(X_test)) + + X_test = tuple([arr.astype(dtype=np.float32) for arr in X_test]) + y_pred = model.predict(X_test) np.savetxt(fname, np.hstack((xt, y_test, y_pred))) @@ -66,9 +68,11 @@ def test_u_cvc(nn, system, T, m, model, data, u, fname): xt = xt * [1 / (m - 1), T / (system.Nt - 1)] X_test = [np.tile(sensor_value, (m * system.Nt, 1)), xt] y_test = s.reshape([m * system.Nt, 1]) - if nn != "opnn": + if nn != "deeponet": X_test = merge_values(X_test) - y_pred = model.predict(data.transform_inputs(X_test)) + + X_test = tuple([arr.astype(dtype=np.float32) for arr in X_test]) + y_pred = model.predict(X_test) np.savetxt("test/u_" + fname, sensor_value) np.savetxt("test/s_" + fname, np.hstack((xt, y_test, y_pred))) @@ -82,16 +86,18 @@ def test_u_advd(nn, system, T, m, model, data, u, fname): xt = xt * [1 / (m - 1), T / (system.Nt - 1)] X_test = [np.tile(sensor_value, (m * system.Nt, 1)), xt] y_test = s.reshape([m * system.Nt, 1]) - if nn != "opnn": + if nn != "deeponet": X_test = merge_values(X_test) - y_pred = model.predict(data.transform_inputs(X_test)) + + X_test = tuple([arr.astype(dtype=np.float32) for arr in X_test]) + y_pred = model.predict(X_test) np.savetxt("test/u_" + fname, sensor_value) np.savetxt("test/s_" + fname, np.hstack((xt, y_test, y_pred))) -def lt_system(npoints_output): +def lt_system(npoints_output, T): """Legendre transform""" - return LTSystem(npoints_output) + return LTSystem(npoints_output, T) def ode_system(T): @@ -103,11 +109,11 @@ def g(s, u, x): # Nonlinear ODE # return -s**2 + u # Gravity pendulum - # k = 1 - # return [s[1], - k * np.sin(s[0]) + u] + k = 1 + return [s[1], - k * np.sin(s[0]) + u] - s0 = [0] - # s0 = [0, 0] # Gravity pendulum + #s0 = [0] + s0 = [0, 0] # Gravity pendulum return ODESystem(g, s0, T) @@ -129,8 +135,10 @@ def cvc_system(T, npoints_output): def advd_system(T, npoints_output): """Advection-diffusion""" - f = None - g = None + # source term + f = 10 + # boundary condition + g = [3 , 0] Nt = 100 return ADVDSystem(f, g, T, Nt, npoints_output) @@ -140,7 +148,7 @@ def run(problem, system, space, T, m, nn, net, lr, epochs, num_train, num_test): X_train, y_train = system.gen_operator_data(space, m, num_train) X_test, y_test = system.gen_operator_data(space, m, num_test) - if nn != "opnn": + if nn != "deeponet": X_train = merge_values(X_train) X_test = merge_values(X_test) @@ -155,8 +163,8 @@ def run(problem, system, space, T, m, nn, net, lr, epochs, num_train, num_test): X_test_trim = trim_to_65535(X_test)[0] y_test_trim = trim_to_65535(y_test)[0] - if nn == "opnn": - data = dde.data.OpDataSet( + if nn == "deeponet": + data = dde.data.Triple( X_train=X_train, y_train=y_train, X_test=X_test_trim, y_test=y_test_trim ) else: @@ -167,13 +175,12 @@ def run(problem, system, space, T, m, nn, net, lr, epochs, num_train, num_test): model = dde.Model(data, net) model.compile("adam", lr=lr, metrics=[mean_squared_error_outlier]) checker = dde.callbacks.ModelCheckpoint( - "model/model.ckpt", save_better_only=True, period=1000 + "model/model", save_better_only=True, period=1000 ) losshistory, train_state = model.train(epochs=epochs, callbacks=[checker]) - print("# Parameters:", np.sum([np.prod(v.get_shape().as_list()) for v in tf.trainable_variables()])) + print("# Parameters:", np.sum([np.prod(v.get_shape().as_list()) for v in tf.compat.v1.trainable_variables()])) dde.saveplot(losshistory, train_state, issave=True, isplot=True) - - model.restore("model/model.ckpt-" + str(train_state.best_step), verbose=1) + model.restore("model/model-" + str(train_state.best_step) + ".ckpt", verbose=1) safe_test(model, data, X_test, y_test) tests = [ @@ -196,7 +203,7 @@ def run(problem, system, space, T, m, nn, net, lr, epochs, num_train, num_test): if problem == "lt": features = space.random(10) - sensors = np.linspace(0, 2, num=m)[:, None] + sensors = np.linspace(0, T, num=m)[:, None] u = space.eval_u(features, sensors) for i in range(u.shape[0]): test_u_lt(nn, system, T, m, model, data, lambda x: u[i], str(i) + ".dat") @@ -221,18 +228,18 @@ def run(problem, system, space, T, m, nn, net, lr, epochs, num_train, num_test): test_u_advd(nn, system, T, m, model, data, lambda x: u[i], str(i) + ".dat") -def main(): +def main(args): # Problems: # - "lt": Legendre transform # - "ode": Antiderivative, Nonlinear ODE, Gravity pendulum # - "dr": Diffusion-reaction # - "cvc": Advection # - "advd": Advection-diffusion - problem = "ode" - T = 1 + problem = args.problem + T = args.t if problem == "lt": npoints_output = 20 - system = lt_system(npoints_output) + system = lt_system(npoints_output, T) elif problem == "ode": system = ode_system(T) elif problem == "dr": @@ -249,29 +256,29 @@ def main(): # space = FinitePowerSeries(N=100, M=1) # space = FiniteChebyshev(N=20, M=1) # space = GRF(2, length_scale=0.2, N=2000, interp="cubic") # "lt" - space = GRF(1, length_scale=0.2, N=1000, interp="cubic") - # space = GRF(T, length_scale=0.2, N=1000 * T, interp="cubic") + # space = GRF(1, length_scale=0.2, N=1000, interp="cubic") + space = GRF(T, length_scale=0.2, N=1000 * T, interp="cubic") # Hyperparameters - m = 100 - num_train = 10000 - num_test = 100000 - lr = 0.001 - epochs = 50000 + m = args.m + num_train = args.num_train + num_test = args.num_test + lr = args.lr + epochs = args.epochs # Network - nn = "opnn" - activation = "relu" - initializer = "Glorot normal" # "He normal" or "Glorot normal" + nn = args.nn + activation = args.activation + initializer = args.init # "He normal" or "Glorot normal" dim_x = 1 if problem in ["ode", "lt"] else 2 - if nn == "opnn": - net = dde.maps.OpNN( + if nn == "deeponet": + net = dde.maps.DeepONet( [m, 40, 40], [dim_x, 40, 40], activation, initializer, use_bias=True, - stacked=False, + stacked=args.stacked, ) elif nn == "fnn": net = dde.maps.FNN([m + dim_x] + [100] * 2 + [1], activation, initializer) @@ -282,4 +289,33 @@ def main(): if __name__ == "__main__": - main() + + def process_initializer(value): + if value == 'Glorot': + return 'Glorot normal' + else: + return 'He normal' + + parser = argparse.ArgumentParser() + parser.add_argument('-p','--problem', type=str, choices=['lt','ode', 'dr', 'cvc', 'advd'], help="Type of differential equation", required=True) + parser.add_argument('-t', type=int, default=1, help="Final time in domain (defualt=1)") + parser.add_argument('-m', type=int, help="number of sensors", required=True) + parser.add_argument('--num-train', type=int, help="number of train data", required=True) + parser.add_argument('--num-test', type=int, help="number of test data", required=True) + parser.add_argument('--lr', type=float, default=1e-3, help="learning rate (default=1e-3)") + parser.add_argument('--epochs', type=int, help="number of epochs", required=True) + parser.add_argument('--nn', type=str, choices=['fnn', 'deeponet', 'resnet'], help="Type of nueral network", default='deeponet') + parser.add_argument('--activation', type=str, choices=['elu', 'gelu', 'relu', 'selu', 'sigmoid', 'silu', 'sin', 'swish', 'tanh'], help="Activation function", default='relu') + parser.add_argument( + '--init', + type=str, + choices=['He', 'Glorot'], + default='Glorot', # Default initializer + help="Specify the initializer (choose from 'He [normal]', 'Glorot [normal]')" + ) + parser.add_argument('--stacked',type=str,default='False',help="Specify whether to use stacked architecture (usage for --nn is 'deeponet')") + args = parser.parse_args() + args.init = process_initializer(args.init) + if args.nn == 'deeponet': + args.stacked = args.stacked.lower() == 'true' + main(args) diff --git a/src/spaces.py b/src/spaces.py index 2734a13..03f99e5 100644 --- a/src/spaces.py +++ b/src/spaces.py @@ -1,15 +1,11 @@ -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - import matplotlib.pyplot as plt import numpy as np from pathos.pools import ProcessPool -from scipy import linalg, interpolate +from scipy import interpolate, linalg from sklearn import gaussian_process as gp -import config from utils import eig +import config class FinitePowerSeries: diff --git a/src/system.py b/src/system.py index eef6c69..7ac6ab2 100644 --- a/src/system.py +++ b/src/system.py @@ -1,28 +1,25 @@ -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - import numpy as np from pathos.pools import ProcessPool from scipy import interpolate from scipy.integrate import solve_ivp from scipy.special import legendre -import config from ADR_solver import solve_ADR from ADVD_solver import solve_ADVD from CVC_solver import solve_CVC from utils import timing +import config class LTSystem(object): - def __init__(self, npoints_output): + def __init__(self, npoints_output, T): """Legendre transform J_n{f(x)}. Args: npoints_output: For a input function, choose n=0,1,2,...,`npoints_output`-1 as data. """ self.npoints_output = npoints_output + self.T = T @timing def gen_operator_data(self, space, m, num): @@ -30,7 +27,7 @@ def gen_operator_data(self, space, m, num): """ print("Generating operator data...", flush=True) features = space.random(num) - sensors = np.linspace(0, 2, num=m)[:, None] + sensors = np.linspace(0, self.T, num=m)[:, None] sensor_values = space.eval_u(features, sensors) sensor_values_tile = np.tile(sensor_values, (1, self.npoints_output)).reshape( diff --git a/src/utils.py b/src/utils.py index 4f9af3a..cd5478a 100644 --- a/src/utils.py +++ b/src/utils.py @@ -1,7 +1,3 @@ -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - import sys import time from functools import wraps @@ -53,7 +49,7 @@ def is_nonempty(X): X = X_test while is_nonempty(X): X_add, X = trim_to_65535(X) - y_pred.append(model.predict(data.transform_inputs(X_add))) + y_pred.append(model.predict(X_add)) y_pred = np.vstack(y_pred) error = np.mean((y_test - y_pred) ** 2) print("Test MSE: {}".format(error))