Численно решаем волновое уравнение разностной схемой

Для меня уравнения в частных производных -- это очень красивая история из студенчества. Почему? Это невероятно красиво. Но что особенно стало для меня захватывающим, так это то, что дифуры в широком смысле прикладной математики -- это тот самый пример, когда математика и компьютер используются вместе, чтобы представить некоторую компьютерную модель вполне реальных процессов. Как вы уже, наверное, догадались, речь пойдёт про то, как вообще можно попробовать решать дифференциальные уравнения в частных производных на компьютере. Мы попробуем это сделать на примере волнового уравнения и с использованием уже ставших привычными python, scipy и numpy. Если вы примерно помните математику, но панически боялись дифуров или они просто как-то обошли вас стороной, то добро пожаловать.

Волновое уравнение

Если просто, то дифференциальное уравнение -- это некоторое уравнение на функцию, в которое входят также её производные. Много уже написано про обыкновенные дифференциальные уравнения, где неизвестная функция зависит от одного аргумента (например, см. здесь и здесь). Когда мы говорим об уравнениях в частных производных, то в них, соответственно, входят частные производные, также не обойдённые пользователями Хабра (зрелищно тут и тут). Вполне естественно, что решение таких уравнений не всегда существует, а если существует, то не обязательно единственное, а для единственности нам нужны некоторые дополнительные условия.

В данном случае мы рассматриваем волновое уравнение для функции

Почему именно такое уравнение?

Давайте представим, что мы имеем не струну, а набор из грузиков, соединённых пружинами в линию; при этом координата каждого грузика зафиксирована и горизонтальное расстояние между соседями равно . Пренебрегая силой тяжести (хотя её при желании легко внести позже в виде правой части), мы можем записать второй закон Ньютона для -координаты грузика , которую мы обозначим :

где в правой части мы записали силу упругости в соответствии с законом Гука. Теперь заметим, что коэффициент жёсткости всей линии грузиков (это же система пружин) равен , а масса грузика равна , за обозначена масса всей системы. Подставим и немного перепишем:

Что ещё постоянно? Длина между закреплёнными концами, так что горизонтальное расстояние между соседями равно . Давайте теперь умножим и разделим правую часть на и...

Знакомое выражение справа? Да, теперь можем перейти к пределу при и получим уже известное

Дополнительно нам понадобятся начальные условия

и в качестве краевых условий рассмотрим краевые условия Дирихле (функция как бы закреплена в нуле на краях)

Решением такой задачи будет динамика профиля натянутой cтруны, который изначально выглядел как , но под действием сил упругости изменяется и в момент выглядит как .

Конечно-разностная схема

Вообще, решение в данном случае даётся, например, формулой Д'Аламбера, но предположим, что мы не изучали уравнения в частных производных и сейчас этого не знаем. Какой разумный рецепт может прийти в голову? Вроде мы помним как примерно оценить производную со времён матана-1...

Почему бы и нет, но подойдём к вопросу строго. По сути мы пойдём в обратную сторону, пытаясь вообразить струну набором грузиков, соединённых пружиной. Давайте поделим струну на равных сегментов и представим её как набор точек.Аналогичным образом поступим со временем, поделив интересующий нас отрезок на равных кусков; итого имеем точек по пространству, точек по времени, всего различных точек, образующих сетку . Наша задача будет состоять в том, чтобы построить такую дискретную (как ещё говорят, сеточную) функцию , которая бы хорошо приближала настоящую . Конкретнее, нам бы хотелось, чтобы

Расстояние между узлами дискретной сетки по обозначим , а .Так, при более густой дискретизации приближённое решение будет ближе к истинному.

Теперь разберёмся с дифференциальным уравнением

Если пронумеруем точки по каждой координате от нуля, то у нас получатся индексы и . Заменим уравнение на похожее для нашей дискретной функции, оценив производную в точке . Вспомним, что у нас есть начальное и краевые условия, поэтому новое уравнение мы записываем не во всех точках:

Мы знаем т.к. в у нас есть начальное условие, таким же образом нам известны и из краевых условий Дирихле. Есть, правда, небольшая проблема: нам неизвестно , -- но её мы можем решить используя второе начальное условие. Помните, была первая производная в момент ? Вспомним старый добрый ряд Тэйлора и аппроксимируем :

Теперь остаётся только выразить из уравнения выше и voilà, мы получаем явную схему (в том смысле, что есть простая итерационная формула при известном ). Наш алгоритм заключается в последовательном вычислении до самого последнего слоя по времени .

Подытожим, обозначив вектор и введя матрицу размера как матрицу с элементами

на диагонали и над диагональю и под ней. Такая матрица называется трёхдиагональной. Наш алгоритм выглядит как-то так:

Вычислений потребуется всего лишь порядка . Напишем небольшой класс, имплементирующий явную схему с таким решением. Для имплементации мы будем использовать scipy.sparse, который нам позволит задействовать разреженные матрицы, сэкономив много памяти и времени на умножения. Это будет особенно важно, если мы потом захотим решать волновое уравнение в размерности .

import numpy as np
import scipy.sparse as sps
import scipy.sparse.linalg as splin


class ExplicitEulerSolver:

    def __init__(self, hs, cCoef):
        '''
        Input
        float[] hs -- the steps of the discretization, [d+1], the first is the time
        '''
        super(ExplicitEulerSolver,self).__init__()

        self.hs = hs
        self.cCoef = cCoef
        self.dim = hs.shape[0]-1
        print(self.dim)


    def solveBVProblem1D(self, uInits, numSteps, rightPart=None, spaceGrid=None):
        '''
            Solves Dirichlet Boundary-Value Problem
            Input
            float[] uInits -- two initial layers of shape [2,...dims...]
            int numSteps -- number of time steps
            funcHandler rightPart -- right part depending on time and space (t,x)
            float[][] spaceGrid -- space grid of points
        '''

        #trim inits
        uSol = [ (uInits[0,...])[None,...], (uInits[1,...])[None,...] ]
        shape= np.array(list(uSol[0].shape)[1:])-2#-2 for dirichlet, trimmed shape for later restoration
        uSol[0] = np.take(uSol[0],np.arange(1,uSol[0].shape[1]-1),axis=1)
        uSol[1] = np.take(uSol[1],np.arange(1,uSol[1].shape[1]-1),axis=1)
        uSol[0] = np.reshape(uSol[0], newshape=(-1,), order='C')[None,:]
        uSol[1] = np.reshape(uSol[1], newshape=(-1,), order='C')[None,:]

        #assembling iteration matrix
        diags = self.hs[0]**2 *np.concatenate([self.cCoef**2 /(self.hs[1]**2) * np.ones([1,shape[0]]),
                    (2/(self.hs[0]**2) - 2*self.cCoef**2 /(self.hs[1]**2) ) * np.ones([1,shape[0]]),
                     self.cCoef**2  /(self.hs[1]**2) * np.ones([1,shape[0]])],axis=0)
        offsets = np.array([-1,0,1])
            
        C = sps.spdiags(diags,offsets, m=len(diags[1]),n=len(diags[1]), format="csr")

        for i in np.arange(numSteps-1):
            if(rightPart is None):
                toAdd = C@uSol[-1][0,:] - uSol[-2][0,:]
            else:
                toAdd = C@uSol[-1][0,:] - uSol[-2][0,:] + np.reshape(rightPart(self.hs[0]*(i+2),spaceGrid),newshape=(-1,))
            uSol.append(toAdd[None,:])

        uSol=np.concatenate(uSol, axis=0)
        toShape= tuple([uSol.shape[0]] + [shape[j] for j in np.arange(len(shape))])
        return np.pad(np.reshape(uSol, newshape=toShape), pad_width=[(0,0)]+[(1,1) for _ in np.arange(self.dim)])

  

Давайте проверим! Пусть , зададим начальные условия

Возьмём, к примеру, и нарисуем результат, используя следующий скрипт:

import numpy as np
import pickle

import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation 

import LinearPDESolvers as linPDEs


import time
#solve wave eq

T=4
xLeft=0
xRight=2
c=2

#Nx = 510
#Nt = 1000
Nx=20
Nt=200
hx = (xRight-xLeft)/Nx
ht = T/Nt


solver = linPDEs.ExplicitEulerSolver(np.array([ht,hx]), cCoef=c)


xs = np.arange(xLeft,xRight+hx/2,hx)
ts = np.arange(0,T+ht/2,ht)

#assume u'(0,x)=0
def initFunc(xx):
    return ((2*xx)*(xx<0.5) + (xx>=0.5)*(-2/3*xx + 4/3))
def dtFunc(xx):
    return 0
uInits = np.zeros([2,len(xs)])
uInits[0,:] = initFunc(xs)
duts = dtFunc(xs)
uInits[1,:] = uInits[0,:] + duts*ht

#SOLVE!
res=solver.solveBVProblem1D(uInits, Nt)


# initializing a figure in 
# which the graph will be plotted
fig = plt.figure() 
   
# marking the x-axis and y-axis
axis = plt.axes(xlim =(xLeft, xRight), 
                ylim =(-1.5, 1.5)) 
  
# initializing a line variable
line, = axis.plot([], [], lw = 3) 

cross = axis.scatter([xLeft,xRight],[0,0], marker="x", s=90, color="red")

txt = axis.text(xLeft+1e-1,-1,"t="+str(ht*0))
# data which the line will 
# contain (x, y)
def init(): 
    line.set_data(xs, res[0,:])
    txt.set_text("t="+str(ht*0))
    return line,txt,
   
def animate(i):
    line.set_data(xs, res[i,:])
    txt.set_text("t={:.2f}".format(ht*i) )
    return line,txt,
   
anim = FuncAnimation(fig, animate, init_func = init,
                     frames = np.arange(0,res.shape[0],1), interval = ht*1000, blit = False, repeat=False)#ms scale

anim.save('explEuler.mp4',fps=30)

with open("./solveTempOut.pkl", "wb") as f:
    pickle.dump({"res": res},f)

Так стоп. Может, нужно попробовать более жесткую дискретизацию, сделаем, например, .

В реальности гитарные струны не разлетаются по Вселенной. В чём же дело? Может, слишком, грубая сетка, возьмём и...

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

Устойчивость разностной схемы

Давайте вспомним наш итеративный метод; при известном

Наш взгляд неизбежно останавливается на том, что вообще-то линейно зависит от . Действительно, в блочном виде

где внизу мы просто записали очевидное равенство . Мы наблюдаем, что решение очень быстро растёт по модулю, значит, определённо, у этой матрицы посередине есть собственные вектора с собственным значением по модулю больше . Которые со временем просто взрываются. Это в чистом виде артефакт алгоритма, с физикой всё в порядке.

Но как проверить? Есть, например, теорема Гершгорина, однако она даёт обычно очень грубую оценку собственных чисел. Считать собственные числа? Но это дорого и не будем же мы для каждого так проверять адекватность схемы. Хотелось бы получить простой критерий, в который мы могли бы подставить и сказать сразу, будет ли алгоритм численно устойчивым.

Именно так и называется это свойство, которое упрощённо заключается в том, что при малых возмущениях полученный также мало отличается от , полученного без возмущений. Возмущения же раскладываются на собственные вектора? Если там будет даже с коэффициентом собственный вектор с собственным значением, скажем, , то уже где-то после шага гитарная струна улетит в небо.

Знаменитая теорема Лакса (оригинальная статья реально хорошая математически, посмотрите, если не видели) утверждает, что если численная схема устойчива, то есть, и аппроксимирует дифференциальное уравнение (в том смысле, что для полиномов от степени не выше приближённая формула для уравнения точна), то численное решение будет равномерно по сетке сходиться к настоящему и ошибка будет порядка . Можно по-разному строить критерий устойчивости разностной схемы, но в любом случае нам надо как-то исследовать спектр матрицы

Матрица симметричная и трёхдиагональная с повторяющейся тройкой чисел. Поэтому, введя множество , векторы рассмотреть как сигналы (добавив нули из условия Дирихле по краям), а их умножение на матрицу можно рассмотреть как свёртку c фильтром , который определяется как

Так мы получим в другом виде

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

Поскольку мы можем взять обратное преобразование Фурье, а ещё прямое и обратное дискретные преобразования Фурье сохраняют расстояния, так что спектр матрицы по модулю меньше тогда и только тогда, когда решение полученного рекуррентного уравнения не будет расти по модулю для всех . Оно не такое сложное, как кажется, например, можем вычислить явно

Теперь вернёмся к записи итерации. Давайте для краткости положим

тогда получим

Нам, таким образом, нужно, чтобы у характеристического уравнения

не было корней, по модулю больших единицы. Давайте не будем торопиться решать его; посмотрев внимательно, мы можем сказать, что

1. По формулам Виета произведение корней равно единице;

2. Уравнение имеет вещественные коэффициенты, следовательно, у него может быть либо пара вещественных решений, либо пара комплексных, но сопряжённых.

Пункт 1 говорит, что если корни вещественные , то один из них точно больше единицы, а кратная единица для всех точно решением быть не может. С другой стороны, если имеем два комплексно сопряжённых корня, то они будут лежать ровно на единичной окружности. Наш случай! Это происходит тогда и только тогда, когда дискриминант отрицателен, то есть

Дальше дело техники, которая нас приводит к

и далее для выполнения неравенства для всех можем потребовать

Число слева иногда называют числом Куранта; в результате мы получили достаточное условие устойчивости. Таким образом, если мы задаём в алгоритме такие шаги , что

полученная схема будет устойчивой и мы не увидим аномалий. Интересный вопрос состоит в том, что будет, если достигается равенство? В этом случае схема тоже будет устойчива, так как подставив

в дискриминант и посчитав корни, мы увидим, что они по модулю меньше единицы.

Наконец, является ли это условие необходимым? Вдруг есть устойчивые схемы, для которых этот критерий не выполняется... Нет, такого не бывает. Если мы посмотрим, что происходит с дискриминантом при равенстве 1 , то мы заметим, что у нас возникает точка , где дискриминант равен нулю и там кратный корень, по модулю меньше 1, который проблемой не является. Однако при числе Куранта больше 1 для некоторых возникнут некратные вещественные корни и это уже однозначно плохой случай. В терминах преобразования Фурье это будет означать, что какие-то частоты в ошибке будут экспоненциально усиливаться с каждой итерацией алгоритма.

Попробуйте схему с числом Куранта равным 1 и немного больше. Результат налицо.

А что насчёт уравнения в размерности d>1?

Физически не всё так очевидно, но мы получим аналогичную краевую задачу

где -- оператор Лапласа, а -- интересующий нас регион пространства и его граница (например, куб и его граница). Построенная нами схема Эйлера будет строиться точно так же, но интересно, что и анализ устойчивости мы можем провести аналогично, заменив одномерную свёртку на d-мерную, а одномерное преобразование Фурье на d-мерное. В полученной итеративной схеме на слое центральная точка будет иметь вес

а точки, отходящие по пространству на 1 клетку, получат вес . И даже что-то типа числа Куранта мы тоже получим, выведя критерий устойчивости, так же записав условие для дискриминанта (не буду утомлять формулами, да и мне кажется, можно получше результат получить):

что в случае d=1 даёт то, что мы получили ранее.

А что дальше?

Если пойти немного дальше и подумать про конструкцию разностной схемы, то можно построить такие, которые будут устойчивы при любом выборе шага. Вместе с теоремой Лакса это дало бы нам возможность не быть такими параноидальными в плане выборе шага по времени. С другой стороны, построенная явная схема Эйлера -- это всего лишь схема первого порядка и для того, чтобы увидеть (более красивую картинку, см. ниже) много деталей, шаг по времени нужно задавать достаточно большим, что будет приводить к большим вычислительным затратам. Возможно, в следующий раз построим что-нибудь покруче. Да и уравнение в размерности d>1 тоже интересно решить.

Думаете, слишком регулярно? Посмотрите на реальное видео: