Релятивистская трассировка лучей

В этой статье я покажу как можно самому, бесплатно и без смс, нарисовать черную дыру при помощи OpenGL, в полном соответствии с ОТО.

Для этого, мы сначала выведем уравнения движения лучей света, напишем интегратор Рунге-Кутты на GLSL ^[1], и наконец, объединив одно с другим, получим фрагментный шейдер ^[2], который вычисляет путь лучей, отправленных из камеры назад во времени.

Подготовка: метрика, связность и геодезические

Как известно, свет распространяется по прямой линии. Это означает, что вектор скорости не меняется при параллельном переносе вдоль самого себя.

Однако «прямая» — понятие относительное. Так, уравнение прямой в полярных координатах выглядит уже совсем не линейным. В случае же, когда и само пространство не плоское (к примеру, на поверхности тора), выбрать систему координат, в которой «прямая» определяется линейным уравнением, попросту невозможно.

Как тогда можно вообще считать некую кривую «прямой»? Самое логичное — взять вектор в любой её точке, направленный по касательной, и начать его переносить вдоль этой кривой. Если он при любом таком переносе останется направленным по касательной — то кривая называется геодезической, и это самое близкое что есть к понятию прямой на произвольном многообразии.

Выводить формулы для геодезических на четырехмерном многообразии — довольно трудоемкое занятие. К счастью, это как раз тот случай, когда труд можно переложить на машину. Воспользуемся библиотекой символьной алгебры sympy ^[3].

from sympy import *
from sympy.diffgeom import *
import numpy as np

Модуль sympy.diffgeom ^[4] довольно нетривиален в использовании и, я бы даже сказал, уродлив, но он неплохо справляется с задачей вывода всяких уравнений на многообразиях, описанных атласом ^[5]. Наш атлас будет иметь всего одну карту, но с двумя разными системами координат - сферической (т.к. именно в ней удобно работать с метрикой Шварцшильда), и декартовой - для рисования на экране.

Объявим символы для координат:

t_, x_, y_, z_ = symbols("t x y z")
rho_, theta_, phi_ = symbols("\rho \theta \phi")

Я использую постфикс _ для терминальных символов, чтобы их было сразу видно в коде.

Составим уравнения перехода из сферических координат в декартовы:

x = rho_ * sin(theta_) * sin(phi_)
y = rho_ * sin(theta_) * cos(phi_)
z = rho_ * cos(theta_)
print(latex(x_) + " = " + latex(x))
print(latex(y_) + " = " + latex(y))
print(latex(z_) + " = " + latex(z))

Создадим 4х-мерное многообразие spacetime, добавим в его атлас карту patch, зададим на ней две наши координатные системы, и добавим возможность переходить из сферической системы в декартову:

spacetime = Manifold("spacetime", 4)
patch = Patch("patch", spacetime)

relation_dict = {
    ("spherical", "cart"): [(x_, y_, z_, t_), (x, y, z, t_)]
}

coord_sh = CoordSystem("spherical", patch, (t_, rho_, theta_, phi_), relation_dict)
coord_cart = CoordSystem("cart", patch, (x_, y_, z_, t_), relation_dict)

Одно из неудобств sympy.diffgeom заключается в том, что уже заданные символы координат нельзя использовать для работы с многообразием. Нужно взять оттуда новые:

c_t, c_rho, c_theta, c_phi = coord_sh.base_scalars()
print(",".join(map(latex, (c_t, c_rho, c_theta, c_phi))))

Как видно, это те же самые символы, которые мы задавали в начале, но пригодные для использования в sympy.giffgeom.

Зададим на нашем многообразии метрику Шварцшильда:

Rs = symbols("r_s")

d_t, d_rho, d_theta, d_phi = coord_sh.base_oneforms()

TP = TensorProduct
metric = (
    + (1 - Rs/c_rho) * TP(d_t, d_t)
    - 1 / (1 - Rs / c_rho) * TP(d_rho, d_rho)
    - (c_rho**2) * TP(d_theta, d_theta)
    - (c_rho**2) * (sin(c_theta)**2) * TP(d_phi, d_phi))

print(latex(metric))

Необходимость выписывать метрику в таком виде - это еще одно неудобство simpy.diffgeom. Казалось бы, дайте возможность скормить туда компоненты тензора , но нет. Только дифференциальные формы, только хардкор!

Как бы то ни было, вооружившись метрикой в таком виде, можно рассчитать коэффициенты метрически‑совместимой связности ^[6] (также известные как символы Кристоффеля второго рода). Метрически‑совместимая связность — это такой способ движения из касательного пространства в одной точке к касательному пространству в другой точке, при котором метрический тензор инвариантен. То есть его ковариантная производная ^[7] равняется нулю (но компоненты тензора в криволинейной системе координат, разумеется, могут меняться).

Модуль simpy.diffgeom услужливо предоставляет функцию для вычисления символов Кристоффеля metric_to_Christoffel_2nd:

christoffel = simplify(metric_to_Christoffel_2nd(metric))

print("\begin{cases}")
for (i,j,k) in np.ndindex(*np.shape(christoffel)):
    if not christoffel[i,j,k].is_zero:
        print(f"\Gamma_{{{i}{j}{k}}} =" + latex(christoffel[i,j,k]) + " \\")
print("\end{cases}")

А вот вывод уравнения геодезической придется написать самим. Он, к счастью, относительно несложный.

Пусть кривая задана параметрически, как . Тогда касательный к ней вектор имеет координаты . По определению геодезической, ковариантная производная касательного вектора вдоль самого себя равняется нулю:

Раскроем ее через символы Кристоффеля:

В первом слагаемом поменяeм порядок множителей и воспользуемся цепным правилом:

Тогда всё уравнение сводится к

Или, проще говоря

На питоне:

def geo_second_derivatives(christoffel, derivatives):
    '''
    Parameters:
        christoffel: tensor of rank 3 - analytical expression of christoffel symbol in given coordinates
        derivatives: iterable consisting of sympy symbols corresponding to each first derivative of a coordinate variable
    '''
    d2u = np.zeros(len(derivatives), dtype='object')
    for n,k,m in np.ndindex(*np.shape(christoffel)):
        d2u[n] -= christoffel[n,k,m] * derivatives[k] * derivatives[m]
    return d2u
```

d2_t, d2_rho, d2_theta, d2_phi = geo_second_derivatives(christoffel, symbols("t' \rho' \theta' \phi'"))

print("\begin{array}{l}")
print(f"{{{latex(c_t)}}}'' = {latex(d2_t)} \\")
print(f"{{{latex(c_rho)}}}'' = {latex(d2_rho)} \\")
print(f"{{{latex(c_theta)}}}'' = {latex(d2_theta)} \\")
print(f"{{{latex(c_phi)}}}'' = {latex(d2_phi)} \\")
print("\end{array}")

Так мы получили систему дифференциальных уравнений, описывающих координаты геодезических линий.

Трассировка лучей

Итак, лучи света (а также и любые другие объекты в гравитационном поле) следуют по "прямым" линиям, заданным дифференциальным уравнением второго порядка. Рассчитать путь одного такого лучика можно легко и непринужденно на том же питоне, воспользовавшись scipy.integrate ^[8].

Но что если мы хотим увидеть черную дыру? Если отправить по лучику света из одной точки через каждый пиксель экрана, то они обернутся вокруг черной дыры, и упадут на то что расположено вокруг нее. А если они при этом будут обращены во времени, то получится что мы как бы приняли лучи, испущенные объектами в прошлом!

Но отправить миллион параллельных лучей — небыстрое занятие. Если только мы не можем это сделать параллельно на какой‑нибудь крутой вычислительной машине с кучей ядер. Именно для таких задач существуют GPU.

Для исполнения одной и той же задачи на каждый пиксель окна есть специальный тип программ для GPU — фрагментные шейдеры. Для написания своего шейдера возьмем OpenGL, и его язык GLSL.

План предельно простой — разместить цветную сферическую оболочку вокруг черной дыры, разместить внутри нее камеру, и отправить из камеры лучи, которые либо упадут на эту оболочку, окрасив пиксель экрана в соответствующий цвет, либо потеряются где‑то на горизонте событий.

Уравнение геодезической в GLSL...

Перепечатывать уравнения, которые были получены в sympy, в другой язык, конечно, мало кому хочется, но к счастью, там есть функция ccode ^[9], а в любом текстовом редакторе есть автозамена наших latex-фредли переменных на текстовые. После небольшого причесывания руками, уравнение геодезической на GLSL выглядит так:

vec4 d2p(vec4 p, vec4 dp)
{
    float t = p[0];
    float rho = p[1];
    float theta = p[2];
    float phi = p[3];

    float D_t = dp[0];
    float D_rho = dp[1];
    float D_theta = dp[2];
    float D_phi = dp[3];

    // Geodesics equation:
    float D2_t = - (D_rho * rs * D_t) / (rho * (rho - rs));
    float D2_rho = (rho - rs) * (
        + pow(D_phi * sin(theta), 2)
        + pow(D_theta, 2)
        - (rs * pow(D_t, 2)) / (2 * pow(rho, 3))
    )
    + (rs * pow(D_rho, 2)) / (2 * rho * (rho - rs));
    float D2_theta = pow(D_phi, 2) * sin(theta) * cos(theta) - (2 * D_rho * D_theta) / rho;
    float D2_phi = -(2 * D_phi * D_theta) / tan(theta) - (2 * D_rho * D_phi) / rho;

    vec4 d2p = {D2_t, D2_rho, D2_theta, D2_phi};
    return d2p;
}

Эта функция принимает на вход координаты точки луча и 4-вектор скорости движения вдоль него. Возвращаемое значение — вторая производная координат при движении вдоль луча.

Здесь сразу можно заметить проблему — видите деление на ? Когда этот угол равняется нулю, значение второй производной будет неопределено. Также, когда угол околонулевой — точность интегрирования просядет. К счастью, это не такая большая проблема. Мы ее решим при помощи переменного шага интегрирования. Альтернативно, можно было бы вспомнить что для невращающейся черной дыры движение луча никогда не покидает одну плоскость, и оптимизировать вычисления таким образом. Но мне лень это я оставлю в качестве упражнения читателю.

Итак, теперь мы можем вычислить вторую производную любой геодезической, проходящей через точку p c 4-скоростью dp. Эта вторая производная в свою очередь показывает как изменяется 4-скорость при проходе через эту точку, а 4-скорость определяет в какую следующую точку мы попадем. Классическая задача для численного интегрирования.

... и его численное решение

В поиске хрупкого баланса между точностью и производительностью, я перепробовал несколько разных методов интегрирования и остановился на адаптивном методе Рунге-Кутты ^[10] 3-4 порядков.

Суть адаптивных методов Рунге-Кутты заключается в том, что параллельно работают два метода разных порядков (в данном случае 3-го и 4-го). Разница между полученными значениями позволяет оценить величину погрешности, и скорректировать шаг интегрирования: если она слишком большая - уменьшить, если достаточно маленькая - увеличить.

Вдвойне удобно то что промежуточные результаты метода одного порядка могут быть переиспользованы с другими коэффициентами внутри метода другого порядка, таким образом увеличивая производительность. Подбор этих коэффициентов - целая наука и на эту тему написана не одна сотня научных статей.

Я эти сотни статей читать, конечно же, не стал. По рабоче-крестьянски я взял классический RK4 как основной метод, и эмпирически подогнал метод третьего порядка чтобы по-максимуму переиспользовать уже посчитанное.

Но погодите, ведь RK4 решает уравнения 1-го порядка, а у нас-то - второго! К счастью, это не проблема. Любое уравнение вида

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

тогда

Это уравнение — уже первого порядка, а значит мы умеем его решать.

Обернем функцию d2p для использования в таком виде:

void rk4_diff(vec4 y[2], out vec4 dy[2])
{
    dy[0] = y[1];
    dy[1] = d2p(y[0], y[1]);
}

Эта обертка на вход принимает 8-мерный вектор, включающий в себя как положение точки, так и ее скорость, а на выходе отдает его производную: скорость точки как производную от положения, и вычисленное функцией d2p «ускорение».

Итоговый код интегратора получился довольно-таки страшным:

/**
 * Adaptive Runge-Kutta of orders 3 and 4, tailored for taking 2nd order derivative for a moving 4D point.
 *  [in]    epsilon - desired value for error estimate. Step is updated based on that value
 *  [in]    max_step - hard limit to how big a step can get
 *  [in]    dtau - integration step
 *  [inout] p0 - point position
 *  [inout] p1 - point velocity
 *  [out]   error_estimate - estimated error value (difference between RK3 and RK4 results)
 * 
 * Return value:
 *  true - error is not too big, can continue processing
 *  false - error is too big, need to rerun with the updated dtau
 */

bool rk34_step(float epsilon, float max_step, inout float dtau, inout vec4 p0, inout vec4 p1, out float error_estimate)
{
    vec4 a[2], b[2], c[2], d[2];

    vec4 y0[2] = {p0, p1};
    rk4_diff(y0, a);

    vec4 y1[2] = {p0 + a[0] * dtau/2, p1 + a[1] * dtau/2};
    rk4_diff(y1, b);

    vec4 y2[2] = {p0 + b[0] * dtau/2, p1 + b[1] * dtau/2};
    rk4_diff(y2, c);

    vec4 y3[2] = {p0 + c[0] * dtau, p1 + c[1] * dtau};
    rk4_diff(y3, d);

    vec4 rk3_0 = (a[0] + 2 * b[0] + 3 * c[0]) * dtau / 6;
    vec4 rk3_1 = (a[1] + 2 * b[1] + 3 * c[1]) * dtau / 6;

    vec4 rk4_0 = (a[0] + 2 * b[0] + 2 * c[0] + d[0]) * dtau / 6;
    vec4 rk4_1 = (a[1] + 2 * b[1] + 2 * c[1] + d[1]) * dtau / 6;

    vec4 diff0 = rk3_0 - rk4_0;
    vec4 diff1 = rk3_1 - rk4_1;
    error_estimate = sqrt(abs(dot(diff0, diff0) + dot(diff1, diff1)));

    float new_dtau = dtau*sqrt(epsilon/error_estimate);
    if (new_dtau < dtau / 2) {
        dtau = new_dtau;
        return false;
    }
    dtau = new_dtau;
    if (dtau > max_step) {
        dtau = max_step;
    }
    p0 += rk4_0;
    p1 += rk4_1;

    return true;
}

Координаты

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

Теперь у нас есть все инструменты чтобы отправить лучи света куда угодно! Но пока нет инструментов, чтобы понять куда же их отправлять.

А для этого, так же как и раньше, нам понадобится две системы координат: сферическая для вычислений, и декартова — для работы с камерой. И возможность переходить между ними.

Преобразование координат точек

Мы уже умеем преобразовывать координаты точек из сферической системы координат в декартову. Просто перенесем это преобразование на GLSL:

vec4 sh_to_cart(vec4 sh)
{
    float t = sh[0];
    float r = sh[1];
    float th = sh[2];
    float phi = sh[3];

    float x = r*sin(th)*sin(phi);
    float y = r*sin(th)*cos(phi);
    float z = r*cos(th);

    vec4 ret = {x, y, z, t};
    return ret;
}

У декартовых координат, в отличие от сферических, время выбрано последней - для того чтобы к можно было обращаться через точку как обычным компонентам GLSL-ного 4-вектора: point.x, point.y, point.z. К можно будет обращаться через point.w.

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

Такое обратное преобразование — это задача, которая для sympy уже не под силу. А вот человеку ее решить довольно просто, из чисто геометрических соображений:

vec4 cart_to_sh(vec4 pos)
{
    float r = sqrt(pos.x*pos.x + pos.y*pos.y + pos.z*pos.z);
    float th = PI / 2. - atan(pos.z/sqrt(pos.x*pos.x + pos.y*pos.y));
    float phi;
    if (pos.y == 0) {
        phi = sign(pos.x) * PI/2;
    } else if (pos.y >= 0) {
        phi = atan(pos.x/pos.y);
    } else {
        phi = PI + atan(pos.x/pos.y);
    }
    float t = pos.w;

    vec4 ret = {t, r, th, phi};
    return ret;
}

Обратите внимание на вычисление угла . Дело в том, что при делении теряется информация об исходных знаках переменных (++ неотличим от --, а +- от -+). Без условного перехода пары точек слились бы в одну. Именно это мешает sympy вычислить обратное преобразование.

Преобразование координат векторов

Напомню, что геодезическая линия (т. е. наш луч света) — это путь перемещения вектора вдоль самого себя. А в криволинейных координатах один и тот же вектор, исходящий из двух разных точек, будет иметь разные компоненты. Для того чтобы преобразовывать координаты векторов из одной системы в другую, существует математический инструмент под названием матрица Якоби ^[12].

И тут на помощь снова приходит sympy:

J = coord_sh.jacobian(coord_cart)
IJ = simplify(J.inv())
print(f"J_{{spherical \rightarrow cart}} = {latex(J)}")
print(f"J_{{cart \rightarrow spherical}} = {latex(IJ)}")

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

Кроме точек, у которых . В них происходит деление на ноль. Но худшее что может случиться - артефакт в виде полоски, бьющей с полюсов черной дыры. Поэтому не будем слишком заморачиваться на этот счет.

При переносе на GLSL эти матрицы нужно транспонировать, так как они там хранятся в column-major формате:

// Value of Jacobi matrix at point sh_pos (point is given in Schwartschield coordinates)
// cartesian_vector = J * sh_vector
mat4 jacobian_inner_to_cart_at(vec4 sh_pos)
{
    float rho = sh_pos[1];
    float theta = sh_pos[2];
    float phi = sh_pos[3];

    // OpenGL stores matrices in column-major format. So this is a TRANSPOSED jacobian matrix
    mat4 ret = {
        {0, 0, 0, 1},
        {sin(phi)*sin(theta),      sin(theta)*cos(phi),      cos(theta),     0},
        {rho*sin(phi)*cos(theta),  rho*cos(phi)*cos(theta), -rho*sin(theta), 0},
        {rho*sin(theta)*cos(phi), -rho*sin(phi)*sin(theta),  0,              0},
    };
    return ret;
}

// Value of inverse Jacobi matrix at point sh_pos (point is given in Schwartschield coordinates)
// sh_vector = IJ * cartesian_vector
mat4 jacobian_cart_to_inner_at(vec4 sh_pos)
{
    float rho = sh_pos[1];
    float theta = sh_pos[2];
    float phi = sh_pos[3];

    // OpenGL stores matrices in column-major format. So this is a TRANSPOSED jacobian matrix
    mat4 ret = {
        {0, sin(theta)*sin(phi), sin(phi)*cos(theta)/rho,  cos(phi)/(rho*sin(theta))},
        {0, sin(theta)*cos(phi), cos(phi)*cos(theta)/rho, -sin(phi)/(rho*sin(theta))},
        {0, cos(theta),         -sin(theta)/rho,           0},
        {1, 0, 0, 0},
    };
    return ret;
}

Рисуем остаток совы

Итак, наш фрагментный шейдер будет вызван для каждого пикселя в области отрисовки. Результат его работы — цвет этого пикселя.

layout(location = 0) out vec4 diffuseColor;

#define PI 3.1415926535897932384626433832795

uniform vec2 window_size;
uniform vec4 camera_pos;
uniform mat3 camera_rotation;
uniform float viewport_dist = 1.0;

uniform float rs = 0.25;
uniform float shell_radius = 2.0;

На вход он принимает следующие uniform ^[13] (т.е. общие для всех пикселей) переменные:

• размер области отрисовки

• параметры камеры (позицию в сферических координатах, матрицу поворота в декартовых, расстояние до точки схождения лучей)

• радиус Шварцшильда

• радиус оболочки

Координаты самого пикселя доступны во встроенной переменной gl_FragCoord.

Осталось самое простое:

• Отправляем луч света через пиксель назад во времени

• В цикле двигаем его по геодезической с помощью rk34_step

• При контакте с оболочкой, красим пиксель в её цвет. Если он при этом еще и покидает оболочку - выходим.

Поехали!

void main()
{
    // Position of the current pixel in the (-1, 1) range
    vec2 window_center = {window_size.x / 2, window_size.y / 2};
    vec2 pixel_pos = (vec2(gl_FragCoord) - window_center) / max(window_center.x, window_center.y);

    // Light ray.
    // Direction from the viewport into the screen pixel:
    vec3 direction3 = {viewport_dist, pixel_pos.x, pixel_pos.y};
    direction3 = normalize(direction3);
    // Adjusted for the camera pointing direction:
    direction3 = camera_rotation * direction3;
    vec4 direction4 = {direction3.x, direction3.y, direction3.z, -length(direction3)};

    // Current position and 4-velocity in Schwartschield coordinates:
    vec4 pos_s = camera_pos;
    vec4 dp = jacobian_cart_to_inner_at(pos_s) * direction4;

    // Initial step
    float dtau = 0.01;
    
    vec3 color = {0,0,0};
    float prev_pho = pos_s[1];
    float max_step = 0.25;
    int i;
    for (i = 0; i < 2000; i++) {
        float r = pos_s[1];

        if (abs(r - shell_radius) < 1) {
            // increase precision near the shell
            max_step = 0.01 + abs(r - shell_radius)/20.;
        } else {
            max_step = 0.5;
        }

        float error_estimate;
        if (rk34_step(0.01, max_step, dtau, pos_s, dp, error_estimate) == false) {
            // Spend this iteration on adjusting the step size.
            continue;
        }
        if ((r > shell_radius * 2) && (r - prev_pho > 0)) {
            // Optimization for looking from outside.
            // This ray has missed the shell completely and keeps moving away from it.
            break;
        }

        // Optimization. If we are this close, not going away anyway
        if (r < rs*1.01) {
            break;
        }

        bool entering_shell = (prev_pho >= shell_radius) && (r < shell_radius);
        bool exiting_shell = (prev_pho < shell_radius) && (r >= shell_radius);
        if ( entering_shell || exiting_shell) {
            float lat = pos_s[2]/PI*180; // in degrees
            float lon = pos_s[3]/PI*180; // in degrees
            float line_freq = 15; // degrees between grid lines

            // Grid lines
            float shine = 1.0;
            float thicc = 0.5;

            float prev_lat_line = floor(lat/line_freq)*line_freq;
            float next_lat_line = prev_lat_line + line_freq;
            color.b += shine * exp(-pow((lat - prev_lat_line)/thicc, 2));
            color.b += shine * exp(-pow((lat - next_lat_line)/thicc, 2));

            float prev_lon_line = floor(lon/line_freq)*line_freq;
            float next_lon_line = prev_lon_line + line_freq;
            color.b += shine * exp(-pow((lon - prev_lon_line)/thicc, 2));
            color.b += shine * exp(-pow((lon - next_lon_line)/thicc, 2));

            color.g = color.b/2;

            if (exiting_shell)
                break;
        }
        prev_pho = r;
    }

    diffuseColor = vec4(color, 1.0);
}

Результат:

Получилось на самом деле довольно-таки неоптимально. Деление на почти-ноль около полюсов сильно портит производительность, заставляя интегратор выбирать очень маленькие шаги для сохранения точности. Только посмотрите на это зловещее красное свечение: оно пропорционально количеству итераций, потраченному на каждый из пикселей!

А это - максимальная оценка ошибки интегратором:

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

Думаю, можно оптимизировать это всё еще в пару десятков раз. Но это я оставлю в качестве упражнения читателю.