Solving Partial Differential Equations (PDEs)#
If you are not yet familiar with the multivarite TFC class, mtfc, it is recommended you learn more about it via the multivariate TFC notebook.
Consider the nonlinear PDE,
where a subscript with respect to \(x\) denotes a derivative with respect to \(x\), subject to the boundary constraints
on the domain \((x,y)\in[0,1]\times[0,2\pi]\). The analytical solution to this differential equation is:
To begin, let’s create the multivariate TFC class and create the analytical solution so we can compare against it later.
[1]:
import jax.numpy as np
from tfc import mtfc
# Create the multivariate TFC class
N = [25,25] # Number of points in the domain
m = 25 # Degree of basis function expansion
nC = [2,[1,-1]] # Indicates which basis functions need to be removed from the expansion
x0 = [0.0,0.0] # Start of the domain for each dimension
xf = [1,2.*np.pi] # End of the domain in each dimension
myTfc = mtfc(N,nC,m,x0=x0,xf=xf)
# Create the analytical solution
realSoln = lambda x,y: x**2*np.cos(y)
WARNING:jax._src.lib.xla_bridge:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
Notice that \(nC\) has been chosen such that the first two basis functions will be removed in from the first dimension, and the second basis function will be removed from the second dimension. That is because these Chebyshev orthogonal polynomials are linearly dependent to the support functions used to create the constrained expression. The remaining variables, \(N\) and \(m\), were chosen based on the author’s previous experience solving PDEs using TFC.
The next step is to develop the constrained expression,
If you are unfamiliar with what a constrained expression is or how to derive one, this article is a good starting point.
[2]:
x = myTfc.x # Collocation points from the TFC class
# Get the basis functions from the TFC class
H = myTfc.H
# Create the constrained expression
g = lambda xi,*x: np.dot(H(*x),xi)
u1 = lambda xi,*x: g(xi,*x)-(1.-x[0])*g(xi,np.zeros_like(x[0]),x[1])+x[0]*(np.cos(x[1])-g(xi,np.ones_like(x[0]),x[1]))
u = lambda xi,*x: u1(xi,*x)+x[1]/(2.*np.pi)*(u1(xi,x[0],2.*np.pi*np.ones_like(x[1]))-u1(xi,x[0],np.zeros_like(x[1])))
Finally, form the residual of the differential equation and minimize it using nonlinear least-squares.
[3]:
from tfc.utils import egrad, NLLS
# Create the residual
ux = egrad(u,1)
uxx = egrad(ux,1)
uy = egrad(u,2)
L = lambda xi: uxx(xi,*x)+ux(xi,*x)*uy(xi,*x)-2.*np.cos(x[1])+2.*x[0]**3*np.sin(x[1])*np.cos(x[1])
# Minimize the residual using least-squares
xi0 = np.zeros(H(*x).shape[1])
xi,_,time = NLLS(xi0,L,timer=True,timerType="perf_counter")
Note that the last line in the above code block calls the JIT. Therefore, it may take a few seconds to run, because code is being compiled. However, once the code is finished compiling it runs very fast. The time returned by the NLLS function is the time it takes the compiled code to run, not the time it takes to compile the code itself. For more information on this function (and an associated class form) see the tutorial on NLLS.
Finally, lets compare the results to the true solution on a test set, and show some statistics about the TFC solution.
[4]:
# Calculate the error on the test set
testSetMat = np.meshgrid(np.linspace(0,1,100),np.linspace(0,2.*np.pi,100))
testSet = [testSetMat[k].flatten() for k in range(2)]
error = np.abs(u(xi,*testSet)-realSoln(*testSet))
# Print out the results
print("Maximum error on test set: " +str(np.max(error)))
print("Mean error on the test set: "+str(np.mean(error)))
print("TFC run time: "+str(time)+" seconds")
Maximum error on test set: 2.9909408283401717e-13
Mean error on the test set: 4.101125863749403e-15
TFC run time: 0.5350202089975937 seconds
The TFC estimated solution is has a mean error on the order of \(10^{-15}\) and was obtained in less than one second.