Visualize the example of # Laplace equation on a disk#

[Steph-Yhf]

Hi, doctor Lu! I want to visualize your example # Laplace equation on a disk# like this :

.
And this is the code:
`
"""Backend supported: tensorflow.compat.v1, tensorflow, pytorch, paddle"""
import deepxde as dde
import matplotlib.pyplot as plt
import numpy as np
from deepxde.backend import torch
def pde(x, y):
dy_r = dde.grad.jacobian(y, x, i=0, j=0)
dy_rr = dde.grad.hessian(y, x, i=0, j=0)
dy_thetatheta = dde.grad.hessian(y, x, i=1, j=1)
return x[:, 0:1] * dy_r + x[:, 0:1] ** 2 * dy_rr + dy_thetatheta
def solution(x):
r, theta = x[:, 0:1], x[:, 1:]
return r * np.cos(theta)
geom = dde.geometry.Rectangle(xmin=[0, 0], xmax=[1, 2 * np.pi])
bc_rad = dde.icbc.DirichletBC(
geom,
lambda x: np.cos(x[:, 1:2]),
lambda x, on_boundary: on_boundary and dde.utils.isclose(x[0], 1),
)
data = dde.data.PDE(
geom, pde, bc_rad, num_domain=2540, num_boundary=80, solution=solution
)
net = dde.nn.FNN([2] + [20] * 3 + [1], "tanh", "Glorot normal")
def feature_transform(x):
return torch.cat(
[x[:, 0:1] * torch.sin(x[:, 1:2]), x[:, 0:1] * torch.cos(x[:, 1:2])], dim=1
)
net.apply_feature_transform(feature_transform)
model = dde.Model(data, net)
model.compile("adam", lr=1e-3, metrics=["l2 relative error"])
losshistory, train_state = model.train(iterations=5000)
dde.saveplot(losshistory, train_state, issave=True, isplot=True)

#draw prediction graph
Nx=500
Ny=500
xmin, xmax, ymin, ymax = [-1, 1, -1, 1]
plot_grid = np.mgrid[xmin : xmax : Nx * 1j, ymin : ymax : Ny * 1j]

#generates a set of three-dimensional arrays
points = np.vstack(
(plot_grid[0].ravel(), plot_grid[1].ravel(), np.zeros(plot_grid[0].size))
)

#slice to be a set of two-dimendsional arrays
points_2d = points[:2, :]

x_tran=points_2d[0, :]
y_tran=points_2d[1, :]
N=points_2d[0, :].size

R=np.zeros_like(x_tran)
theta=np.zeros_like(x_tran)
p=np.zeros_like(x_tran)

#calculate the polar diameter and angle
for i in range(0, N, 1):
R[i]=np.sqrt(x_tran[i]**2+y_tran[i]**2)
theta[i]=np.arctan2(y_tran[i], x_tran[i])

X_star = np.column_stack((R, theta))

for i in range(0, N, 1):
if R[i]<=1:
p[i] = model.predict(X_star[i, :])

else:
    p[i]=0

p=p.reshape((Nx, Ny))

plt.rc("font",family="serif",size=22)

fig,ax=plt.subplots(1,1, figsize=(24, 12))
matrix=np.fliplr(p).T

pcm=ax.imshow(
matrix,
extent=[-1, 1, -1, 1],
cmap=plt.cm.get_cmap("seismic"),
interpolation="spline16",
label="PINN",
)
plt.gca().set_aspect('equal')

fig.colorbar(pcm, ax=ax)

ax.set_title("PINN")
plt.savefig("plot_Reynolds.pdf")
The following error occurs:

`
How should I modify it ?