УРЧП

Общая информация

Нетрудно обобщить описанный выше метод для решения УРЧП. Рассмотрим двумерное уравнение Пуассона (как один из самых простых примеров УРЧП):

\[ \begin{equation*} \frac{\partial^{2}}{\partial x^{2}} f(x, y)+\frac{\partial^{2}}{\partial y^{2}} f(x, y)=G(x, y) \end{equation*} \]

где \(x \in[a, b], y \in[c, d]\)

Можно задавать разные граничные условия, в том числе Неймана или смешанные, однако для простоты ограничимся условиями Дирихле - для других случаев всё идейно аналогично, но формулы более громоздкие.

Таким образом \(f(a, y)=A(y), \; f(b, y)=B(y), \; f(x, c)=C(x)\) и \(f(x, d)=D(x)\).

\[ \begin{equation*} \hat{f}(x, y ; \theta)=E(x, y)+(x-a)(b-x)(y-c)(d-y) N(x, y ; \theta) \end{equation*} \]

Функция перед \(N(x, y ; \theta)\) подбирается легко, ведь она должна просто занулиться на границах прямоугольника. А вот с \(E(x, y)\) всё сложнее - придумать функцию, принимающую заданные значения на границе быстро не получается.

К счастью, всё уже придумали за нас, можно просто подставить граничные значения и убедиться, что всё работает:

\[ \begin{align*} E(x, y)= & \frac{b-x}{b-a} A(y)+\frac{x-a}{b-a} B(y)+\frac{d-y}{d-c}\left\{C(x)-\left[\frac{b-x}{b-a} C(a)+\frac{x-a}{b-a} C(b)\right]\right\} +\frac{y-c}{d-c}\left\{D(x)-\left[\frac{b-x}{b-a} D(a)+\frac{x-a}{b-a} D(b)\right]\right\} \end{align*} \]

Стандартным образом также пишется квадрат невязки:

\[ \begin{equation*} J(\theta)=\sum_{j=1}^{M}\left(\frac{\partial^{2}}{\partial x^{2}} \hat{f}\left(x_{j}, y_{j} ; \theta\right)+\frac{\partial^{2}}{\partial y^{2}} f\left(x_{j}, y_{j} ; \theta\right)-G\left(x_{j}, y_{j}\right)\right)^{2} \end{equation*} \]

Пример

Рассмотрим следующее уравнение на квадрате \(\in[0,1] \times[0,1]\)

\[ \begin{equation*} \frac{\partial^{2}}{\partial x^{2}} f(x, y)+\frac{\partial^{2}}{\partial y^{2}} f(x, y)=e^{-x}\left(x-2+y^{3}+6 y\right) \end{equation*} \]

с граничными условиями

\[f(0, y)=y^{3}, \; f(1, y)=\left(1+y^{3}\right) e^{-1}, \; f(x, 0)=x e^{-x}, \; f(x, 1)=e^{-x}(x+1)\]

Как описано в общей информации к этому разделу, выбираем

\[ \begin{equation*} \hat{f}(x, y ; \theta)=E(x, y)+x(1-x) y(1-y) N(x, y ; \theta) \end{equation*} \]

\[\begin{split} \begin{align*} E(x, y)= & (1-x) y^{3}+x\left(1+y^{3}\right) e^{-1}+(1-y) x\left(e^{-x}-e^{-1}\right) \\ & +y\left[e^{-x}(x+1)-\left(1-x+2 e^{-1} x\right)\right] \end{align*} \end{split}\]

Аналитическое решение при этом выглядит очень просто:

\[f(x, y)=e^{-x}\left(x+y^{3}\right)\]

import torch
import torch.utils.data
import numpy as np
import matplotlib.pyplot as plt
from torch.autograd import grad


class DataSet(torch.utils.data.Dataset):

    def __init__(self, xRange, yRange, numSamples):
        X  = torch.linspace(xRange[0],xRange[1],numSamples, requires_grad=True)
        Y  = torch.linspace(yRange[0],yRange[1],numSamples, requires_grad=True)
        grid = torch.meshgrid(X,Y)
        self.data_in = torch.cat((grid[0].reshape(-1,1),grid[1].reshape(-1,1)),1)

    def __len__(self):
        return self.data_in.shape[0]

    def __getitem__(self, i):
        return self.data_in[i]

class PDESolver(torch.nn.Module):
    def __init__(self, numHiddenNodes):
        super(PDESolver, self).__init__()
        self.fc1 = torch.nn.Linear(in_features = 2, out_features = numHiddenNodes)
        self.fc2 = torch.nn.Linear(in_features = numHiddenNodes, out_features = 1)

    def forward(self, input):
        h = torch.tanh(self.fc1(input))
        z = self.fc2(h)
        return z

def solution(x,y):
    return torch.exp(-x) * (x + y**3)

def trial_term(x,y):
    e_inv = np.exp(-1)
    return ((1-x)*(y**3) + x*(1+(y**3))*e_inv + (1-y)*x*(torch.exp(-x)-e_inv) +
            y * ((1+x)*torch.exp(-x) - (1-x+(2*x*e_inv))))

def trial(x,y,n_out):
    return trial_term(x,y) + x*(1-x)*y*(1-y)*n_out

def dx_trial(x,y,n_out,n_x):
    return (-torch.exp(-x) * (x+y-1) - y**3 + (y**2 +3) * y * np.exp(-1) + y
            + y*(1-y)* ((1-2*x) * n_out + x*(1-x)*n_x))

def dx2_trial(x,y,n_out,n_x,n_xx):
    return (torch.exp(-x) * (x+y-2)
            + y*(1-y)* ((-2*n_out) + 2*(1-2*x)*n_x) + x*(1-x)*n_xx)

def dy_trial(x,y,n_out, n_y):
    return (3*x*(y**2 +1) *np.exp(-1) - (x-1)*(3*(y**2)-1) + torch.exp(-x)
            + x*(1-x)* ((1-2*y) * n_out + y*(1-y)*n_y) )

def dy2_trial(x,y,n_out,n_y,n_yy):
    return (np.exp(-1) * 6 * y * (-np.exp(1)*x + x + np.exp(1))
            + x*(1-x)* ((-2*n_out) + 2*(1-2*y)*n_y) + y*(1-y)*n_yy)

def diffEq(x,y,trial_dx2,trial_dy2):
    RHS = torch.exp(-x) * (x - 2 + y**3 + 6*y)
    return trial_dx2 + trial_dy2 - RHS


def train(network, loader, lossFn, optimiser, numEpochs):

    cost_list=[]
    network.train(True)
    for epoch in range(numEpochs):
        for batch in loader:
            n_out = network(batch)

            dn = grad(n_out, batch, torch.ones_like(n_out), retain_graph=True, create_graph=True)[0]
            dn2 = grad(dn, batch, torch.ones_like(dn), retain_graph=True, create_graph=True)[0]

            n_x, n_y = torch.split(dn, split_size_or_sections=1, dim=1)
            n_xx, n_yy = torch.split(dn2, split_size_or_sections=1, dim=1)

            x, y = torch.split(batch, 1, dim=1)
            trial_dx2 = dx2_trial(x,y,n_out,n_x,n_xx)
            trial_dy2 = dy2_trial(x,y,n_out,n_y,n_yy)
            D = diffEq(x,y, trial_dx2, trial_dy2)
            cost = lossFn(D, torch.zeros_like(D))
            cost.backward()
            optimiser.step()
            optimiser.zero_grad()

        cost_list.append(cost.item())
    network.train(False)
    return cost_list

def plotNetwork(network, epoch):

    numTestSamples = 12
    X  = torch.linspace(xRange[0],xRange[1],numTestSamples, requires_grad=True)
    Y  = torch.linspace(yRange[0],yRange[1],numTestSamples, requires_grad=True)
    x_mesh,y_mesh = torch.meshgrid(X,Y)

    input = torch.cat((x_mesh.reshape(-1,1),y_mesh.reshape(-1,1)),1)
    N = network.forward(input)

    dn = grad(N, input, torch.ones_like(N), retain_graph=True, create_graph=True)[0]
    dn2 = grad(dn, input, torch.ones_like(dn), retain_graph=True, create_graph=True)[0]
    n_x, n_y = torch.split(dn, split_size_or_sections=1, dim=1)
    n_xx, n_yy = torch.split(dn2, split_size_or_sections=1, dim=1)

    x, y = torch.split(input, split_size_or_sections=1, dim=1)

    trial_dx2 = dx2_trial(x,y,N,n_x,n_xx)
    trial_dy2 = dy2_trial(x,y,N,n_y,n_yy)

    D = diffEq(x,y, trial_dx2, trial_dy2)

    cost = lossFn(D, torch.zeros_like(D))
    print("test cost = ", cost.item())

    output = trial(x_mesh.reshape(-1,1),y_mesh.reshape(-1,1),N)
    output = output.reshape(numTestSamples,numTestSamples).detach().numpy()
    exact = solution(x_mesh,y_mesh).detach().numpy()

    surfaceError = ((output-exact)**2).mean()
    print("MSE полученным и аналитическим решением = ", surfaceError )

    x_mesh = x_mesh.detach().numpy()
    y_mesh = y_mesh.detach().numpy()
    ax = plt.axes(projection='3d')
    ax.plot_surface(x_mesh,y_mesh,output,rstride=1, cstride=1,
                cmap='plasma', edgecolor='none')
    ax.scatter(x_mesh,y_mesh,exact, label = 'Exact Solution')
    ax.set_xlabel('x', fontsize = 16)
    ax.set_ylabel('y', fontsize = 16)
    ax.set_zlabel('z', fontsize = 16)
    ax.legend()
    ax.set_title("Learning Rate = " + str(lr) + ": " + str(epoch) + " Epochs", fontsize = 16)
    plt.show()
    return surfaceError, cost.item()


lr = 1e-3

network    = PDESolver(numHiddenNodes=16)

xRange = [0,1]
yRange = [0,1]
numSamples = 10
numEpochs = 1000
costListDict = {}

lossFn      = torch.nn.MSELoss()
trainSet    = DataSet(xRange,yRange,numSamples)
trainLoader = torch.utils.data.DataLoader(dataset=trainSet, batch_size=int(numSamples**2), shuffle=True)

optimiser  = torch.optim.Adam(network.parameters(), lr = lr)
costList = []
epoch = 0
totalEpochs = 30000
while epoch < totalEpochs:
    costList.extend(train(network, trainLoader, lossFn, optimiser, numEpochs))
    epoch += numEpochs

print(f"{epoch} epochs total, final cost = {costList[-1]}")
plt.semilogy(costList)
plt.xlabel("Epochs", fontsize = 16)
plt.ylabel("Cost", fontsize = 16)
plt.title(f"Training Cost, Learning Rate = {lr}", fontsize = 16)
plt.show()

a, b = plotNetwork(network, epoch)

/usr/local/lib/python3.10/dist-packages/torch/functional.py:507: UserWarning: torch.meshgrid: in an upcoming release, it will be required to pass the indexing argument. (Triggered internally at ../aten/src/ATen/native/TensorShape.cpp:3549.)
  return _VF.meshgrid(tensors, **kwargs)  # type: ignore[attr-defined]

30000 epochs total, final cost = 1.3263698406262847e-07

test cost =  1.2643322122585232e-07
MSE полученным и аналитическим решением =  1.1546887e-06

Получили очень маленькую ошибку, взяв всего 10*10 = 100 точек, один скрытый слой и обучившись за 3 минуты.