Merge remote-tracking branch 'origin/main' into main

Author: levimcclenny

examples/steady-state-poisson.py

@@ -0,0 +1,101 @@
+import math
+import matplotlib.pyplot as plt
+import tensorflow as tf
+import tensordiffeq as tdq
+from tensordiffeq.boundaries import *
+from tensordiffeq.models import CollocationSolverND
+from tensorflow.math import sin
+from tensordiffeq.utils import constant
+
+Domain = DomainND(["x", "y"])
+
+Domain.add("x", [0, 1.0], 11)
+Domain.add("y", [0, 1.0], 11)
+
+N_f = 100
+Domain.generate_collocation_points(N_f)
+
+
+def f_model(u_model, x, y):
+    u = u_model(tf.concat([x, y], 1))
+    u_x = tf.gradients(u, x)[0]
+    u_y = tf.gradients(u, y)[0]
+    u_xx = tf.gradients(u_x, x)[0]
+    u_yy = tf.gradients(u_y, y)[0]
+
+    a1 = constant(1.0)
+    a2 = constant(1.0)
+    pi = constant(math.pi)
+
+    # we use this specific forcing term because we have an exact analytical solution for this case
+    # to compare the results of the PINN solution
+    # note that we must use tensorflow math primitives such as sin, cos, etc!
+    forcing = - sin(a1 * pi * x) * sin(a2 * pi * y)
+
+    f_u = u_xx + u_yy - forcing  # = 0
+
+    return f_u
+
+
+def func_upper_x(y):
+    return -sin(constant(math.pi) * y) * sin(constant(math.pi))
+
+
+def func_upper_y(x):
+    return -sin(constant(math.pi) * x) * sin(constant(math.pi))
+
+
+lower_x = dirichletBC(Domain, val=0.0, var='x', target="upper")
+upper_x = FunctionDirichletBC(Domain, fun=[func_upper_x], var='x', target="upper", func_inputs=["y"], n_values=10)
+upper_y = FunctionDirichletBC(Domain, fun=[func_upper_y], var='y', target="upper", func_inputs=["x"], n_values=10)
+lower_y = dirichletBC(Domain, val=0.0, var='y', target="lower")
+
+BCs = [upper_x, lower_x, upper_y, lower_y]
+
+layer_sizes = [2, 16, 16, 1]
+
+model = CollocationSolverND()
+model.compile(layer_sizes, f_model, Domain, BCs)
+model.tf_optimizer = tf.keras.optimizers.Adam(lr=.005)
+model.fit(tf_iter=4000)
+
+# get exact solution
+nx, ny = (11, 11)
+x = np.linspace(0, 1, nx)
+y = np.linspace(0, 1, ny)
+
+xv, yv = np.meshgrid(x, y)
+
+x = np.reshape(x, (-1, 1))
+y = np.reshape(y, (-1, 1))
+
+# Exact analytical soln is available:
+Exact_u = (np.sin(math.pi * xv) * np.sin(math.pi * yv)) / (2*math.pi**2)
+
+# Flatten for use
+u_star = Exact_u.flatten()[:, None]
+
+# Plotting
+x = Domain.domaindict[0]['xlinspace']
+y = Domain.domaindict[1]["ylinspace"]
+
+X, Y = np.meshgrid(x, y)
+
+# print(np.shape((X,Y))) # 2, 256, 256
+X_star = np.hstack((X.flatten()[:, None], Y.flatten()[:, None]))
+
+lb = np.array([0.0, 0.0])
+ub = np.array([1.0, 1])
+
+u_pred, f_u_pred = model.predict(X_star)
+
+#error_u = tdq.helpers.find_L2_error(u_pred, u_star)
+#print('Error u: %e' % (error_u))
+
+U_pred = tdq.plotting.get_griddata(X_star, u_pred.flatten(), (X, Y))
+FU_pred = tdq.plotting.get_griddata(X_star, f_u_pred.flatten(), (X, Y))
+
+lb = np.array([0.0, 0.0])
+ub = np.array([1.0, 1.0])
+
+tdq.plotting.plot_solution_domain1D(model, [x, y], ub=ub, lb=lb, Exact_u=Exact_u)

examples/transfer-learn.py

@@ -0,0 +1,104 @@
+import scipy.io
+import math
+import tensorflow as tf
+import tensordiffeq as tdq
+from tensordiffeq.models import CollocationSolverND
+from tensordiffeq.boundaries import *
+
+Domain = DomainND(["x", "t"], time_var='t')
+
+Domain.add("x", [-1.0, 1.0], 512)
+Domain.add("t", [0.0, 1.0], 201)
+
+N_f = 50000
+Domain.generate_collocation_points(N_f)
+
+
+def func_ic(x):
+    return x ** 2 * np.cos(math.pi * x)
+
+
+# Conditions to be considered at the boundaries for the periodic BC
+def deriv_model(u_model, x, t):
+    u = u_model(tf.concat([x, t], 1))
+    u_x = tf.gradients(u, x)[0]
+    # u_xx = tf.gradients(u_x, x)[0]
+    # u_xxx = tf.gradients(u_xx, x)[0]
+    # u_xxxx = tf.gradients(u_xxx, x)[0]
+    return u, u_x
+
+
+init = IC(Domain, [func_ic], var=[['x']])
+x_periodic = periodicBC(Domain, ['x'], [deriv_model])
+
+BCs = [init, x_periodic]
+
+
+def f_model(u_model, x, t):
+    u = u_model(tf.concat([x, t], 1))
+    u_x = tf.gradients(u, x)
+    u_xx = tf.gradients(u_x, x)
+    u_t = tf.gradients(u, t)
+    c1 = tdq.utils.constant(.0001)
+    c2 = tdq.utils.constant(5.0)
+    f_u = u_t - c1 * u_xx + c2 * u * u * u - c2 * u
+    return f_u
+
+
+col_weights = tf.Variable(tf.random.uniform([N_f, 1]), trainable=True, dtype=tf.float32)
+u_weights = tf.Variable(100 * tf.random.uniform([512, 1]), trainable=True, dtype=tf.float32)
+
+layer_sizes = [2, 128, 128, 128, 128, 1]
+
+model = CollocationSolverND()
+model.compile(layer_sizes, f_model, Domain, BCs, isAdaptive=True, col_weights=col_weights, u_weights=u_weights)
+model.fit(tf_iter=5000)
+model.save("test_model")
+
+# Must re-initialize the model class in order to effectively transfer learn or resume training
+model = CollocationSolverND()
+model.compile(layer_sizes, f_model, Domain, BCs, isAdaptive=True, col_weights=col_weights, u_weights=u_weights)
+model.tf_optimizer = tf.keras.optimizers.Adam(.0001)
+model.tf_optimizer_weights= tf.keras.optimizers.Adam(.0001)
+model.load_model("test_model")
+model.fit(tf_iter=5000)
+
+# Must re-initialize the model class in order to effectively transfer learn or resume training
+model = CollocationSolverND()
+model.compile(layer_sizes, f_model, Domain, BCs, isAdaptive=True, col_weights=col_weights, u_weights=u_weights)
+model.tf_optimizer = tf.keras.optimizers.Adam(.00001)
+model.tf_optimizer_weights= tf.keras.optimizers.Adam(.00001)
+model.load_model("test_model")
+model.fit(tf_iter=5000)
+
+# Load high-fidelity data for error calculation
+data = scipy.io.loadmat('AC.mat')
+
+Exact = data['uu']
+Exact_u = np.real(Exact)
+
+
+
+x = Domain.domaindict[0]['xlinspace']
+t = Domain.domaindict[1]["tlinspace"]
+
+# create mesh for plotting
+
+X, T = np.meshgrid(x, t)
+
+X_star = np.hstack((X.flatten()[:, None], T.flatten()[:, None]))
+u_star = Exact_u.T.flatten()[:, None]
+
+# forward pass through model
+u_pred, f_u_pred = model.predict(X_star)
+
+error_u = tdq.helpers.find_L2_error(u_pred, u_star)
+print('Error u: %e' % (error_u))
+
+U_pred = tdq.plotting.get_griddata(X_star, u_pred.flatten(), (X, T))
+FU_pred = tdq.plotting.get_griddata(X_star, f_u_pred.flatten(), (X, T))
+
+lb = np.array([-1.0, 0.0])
+ub = np.array([1.0, 1])
+
+tdq.plotting.plot_solution_domain1D(model, [x, t], ub=ub, lb=lb, Exact_u=Exact_u)