Solving an ODE on a complex domain

If you are not yet familiar with solving ODEs on real domains, it is recommended you learn more about it via the ODE TFC tutorial.

Consider the nonlinear ODE,

\[y_{xx}+y_x y = e^{-2x}\sin(x)\Big(\cos(x)-\sin(x)\Big)-2e^{-x}\cos(x),\]

where a subscript with respect to \(x\) denotes a derivative with respect to \(x\), subject to the boundary constraints

\[y(1) = e^{-1}\sin(1) \quad \text{and} \quad y(1.2 + 1.2i) = e^{-1.2 -1.2i} \sin(1.2 + 1.2i).\]

The domain for this ODE is a circle of radius 1 in the complex plane centered around \(1+1i\). The analytical solution to this differential equation is:

\[y(x) = e^{-x} \sin(x).\]

To begin, let’s create the univariate TFC class, the analytical solution, and points in the domain that will be used to solve the problem.

# Setup random number generation with the same seed for repeatability
from numpy import random
random.seed(0)

import jax.numpy as np
from tfc import utfc

# Analytical solution
f2 = lambda x: np.exp(-2.*x)*np.sin(x)*(np.cos(x)-np.sin(x))-2.*np.exp(-x)*np.cos(x)
realSoln = lambda x: np.exp(-x)*np.sin(x)

# Create the univariate TFC class
tfc = utfc(100,0,95,basis="ELMTanh",x0=0.0+0.0j,xf=1.0+0.0j, backend="Python")

# Create the analytical solution
realSoln = lambda x: np.exp(-x)*np.sin(x)

# Create points in the domain
r = np.linspace(0.,1.,10).reshape((1,10))
th = np.linspace(0.,2*np.pi,10).reshape((10,1))
real = r*np.sin(th)
imag = r*np.cos(th)*1.j
x = (real+imag).flatten() + 1. + 1.j

Notice that \(x_0 = 0\) and \(x_f = 1.0\) even though these are not in the domain. The reason is the TFC basis function will automatically map the problem domain to the basis function domain. For ELMs, the default is \([0.0, 1.0]\). Therefore, setting \(x_0 = 0\) and \(x_f = 1.0\) means that no mapping will be done, since for this problem we will do it ourselves; the mapping will be to the unit circle centered around 0.0, which can be done by simply subtracting \(1+1i\) from \(x\).

For this problem, we will set the random weights and biases of the ELM to be uniformly distributed in the unit circle in the complex plane centered around 0.0.

# Set weigths and biases
size = tfc.basisClass.b.size

r = random.uniform(low=0.0, high=1.0, size=size)
th = random.uniform(low=0.0, high=2.0 * np.pi, size=size)
tfc.basisClass.w = r*(np.cos(th)+np.sin(th)*1.j)

r = random.uniform(low=0.0, high=1.0, size=size)
th = random.uniform(low=0.0, high=2.0 * np.pi, size=size)
tfc.basisClass.b = r*(np.cos(th)+np.sin(th)*1.j)

The next step is to develop the constrained expression,

\[y(x,g(x)) = g(x) + \frac{x-1}{0.2+1.2i}\Big(e^{-1.2-1.2i}\sin(1.2+1.2i) - g(1.2+1.2i)\Big) + \frac{x-1.2-1.2i}{-0.2-1.2j}\Big(e^{-1}\sin(1) - g(1)\Big).\]

# Get the basis functions from the TFC class
H = tfc.H

# Create the constrained expression
g = lambda x,xi: np.dot(H(x-1.0-1.0j),xi)
u = lambda x,xi: g(x,xi) + (x-1.0)/(0.2 + 1.2j)*(np.exp(-1.2 -1.2j)*np.sin(1.2+1.2j)-g(np.ones_like(x)*(1.2 + 1.2j), xi)) + (x-1.2 - 1.2j)/(-0.2 -1.2j)*(np.exp(-1.0)*np.sin(1.0) - g(np.ones_like(x),xi))

where \(g\) subtracts \(1+1i\) from x to normalize it as discussed earlier.

Finally, form the residual of the differential equation and minimize it using nonlinear least-squares.

from tfc.utils import egrad, NLLS

# Create the residual
du = egrad(u)
d2u = egrad(du)
L = lambda xi,x: d2u(x,xi)+du(x,xi)*u(x,xi)-np.exp(-2.*x)*np.sin(x)*(np.cos(x)-np.sin(x))+2.*np.exp(-x)*np.cos(x)

# Minimize the residual using least-squares
xi0 = np.zeros(H(x).shape[1], dtype=x.dtype)
xi, _, time = NLLS(xi0,L,x,constant_arg_nums=[1], method="lstsq", timer=True, holomorphic=True, tol=1e-9)

In the above code, constant_arg_nums must be used to make all functions that depend solely on x compile time constants. This is necessary when using the Python backend, since no JIT transform is defined for the basis function primitives, see Basis function backends for more details.

lstsq is used since the ELM basis classes will likely produce near-singular matrices in the nonlinear least squares if enough basis functions are used. In addition, holomorphic=True must be set so that JAX knows our complex functions are holomorphic.

Finally, lets compare the results to the true solution on a test set and show some statistics about the TFC solution.

# Create test points
numTest = 100
r = np.linspace(0.,1.,numTest).reshape((1,numTest))
th = np.linspace(0.,2*np.pi,numTest).reshape((numTest,1))
real = r*np.sin(th)
imag = r*np.cos(th)*1.j
test = (real+imag).flatten() + 1. + 1.j

# Calculate the error
U = u(test,xi)
err = U - realSoln(test)
maxErr = np.max(np.abs(err))

# Display error statistics
print(f"Time: {time}")
print(f"Max error: {maxErr}")

# Plot the results
from tfc.utils.PlotlyMakePlot import MakePlot
test = test.reshape((numTest, numTest))
U = U.reshape((numTest, numTest))
p = MakePlot([["x<sub>real</sub>"],["x<sub>real</sub>"]],[["x<sub>imag</sub>"],["x<sub>imag</sub>"]],zlabs=[["u<sub>real</sub>"],["u<sub>imag</sub>"]])

p.Surface(x=np.real(test), y=np.imag(test), z=np.real(U), row=1, col=1, showscale=False)
p.Surface(x=np.real(test), y=np.imag(test), z=np.imag(U), row=2, col=1, showscale=False)

p.Scatter3d(x=[1., 1.2],
            y=[0., 1.2],
            z=[np.exp(-1.0)*np.sin(1.0), np.real(np.exp(-1.2 - 1.2j)*np.sin(1.2j + 1.2))],
            mode="markers",
            marker=dict(color="red",size=5),
            row = 1, col = 1
           )
p.Scatter3d(x=[1., 1.2],
            y=[0., 1.2],
            z=[0.0, np.imag(np.exp(-1.2 -1.2j)*np.sin(1.2 + 1.2j))],
            mode="markers",
            marker=dict(color="red",size=5),
            row = 2, col = 1
           )

p.fig.update_layout(showlegend=False,scene_aspectmode='cube')
p.view(azimuth=225,elevation=25, row=1, col=1)
p.view(azimuth=225,elevation=25, row=2, col=1)
p.fig.layout.margin.t = 0.0; p.fig.layout.margin.b = 0.0
p.PartScreen(750,1500,units='px')
p.show()

Time: 0.05192915000000031
Max error: 3.838183681079111e-09

The TFC estimated solution has a maximum error less than \(5.0\times10^{-9}\) and was obtained in a few hunderedths of a second.