algo.tex

\documentclass{article}

\usepackage{amsmath}
\usepackage{amssymb}
\usepackage{tikz-cd}

\usepackage{color}

\DeclareMathOperator{\id}{id}

\title{Rank Polymorphic MultiGrid}

\begin{document}

\textcolor{red}{Work in progress, not in readable state yet.}

\section{One-Dimensional Discretization}

We write $\{e_0, \cdots, e_{n - 1}\}$ for the standard basis of $\mathbb{R}^n$.
The one-dimensional discretization of $\nabla$ (periodic boundary conditions)
is $A / h^2$ for

\[
A: e_{i} \mapsto e_{(i - 1) \mod n} - 2 e_{i} + e_{(i + 1) \mod n}.
\]

We want to solve $(A / h^2)u = f$ or equivalently $Au = h^2 f$. For parallel
algorithms it is useful to solve a transformed system instead. To this end
we look at the vector space $\mathbb{R}^{n / 2} \oplus \mathbb{R}^{n / 2}$
instead, where we call the first summand the black part, with basis
$b_0, \cdots, b_{n / 2 - 1}$, and the second summand the red part, with
basis $r_0, \cdots, r_{n / 2 - 1}$. Let
$\sigma: \mathbb{R}^{n / 2} \oplus \mathbb{R}^{n / 2} \to \mathbb{R}$ be
the isomorphism that sends $b_i$ to $e_{2i}$ and $r_i$ to $e_{2i + 1}$.
Then the red-black form of $A$ is $A^{rb} := \sigma^{-1} A \sigma$.
We can equivalently solve the system $A^{rb}u^{rb} = \sigma^{-1} f$ and recover
$u$ as follows.

\begin{align*}
    A^{rb}u^{rb} = \sigma^{-1} f \\
    \sigma^{-1} A \sigma u^{rb} = \sigma^{-1} f \\
    A (\sigma u^{rb}) = f
\end{align*}

\newpage

From now on, we are not going to write $\mod$. The function $A^{rb}$ is
given more explicitly by

\begin{align*}
    A^{rb}: & \\
    b_{i} \mapsto & -2b_{i} + r_{i - 1} + r_{i} \\
    r_{i} \mapsto & -2r_{i} + b_{i} + b_{i + 1} \\ 
\end{align*}

We split $A^{rb} = M - N$, so we can write $A^{rb}u = h^2f \iff Mu = Nu + h^2f$,
so we can use update formula $u^{(k + 1)} = M^{-1}(Nu^{(k)} + h^2f)$.

We parameterize by $\omega$.

\begin{align*}
    M : & \\
    b_{i} \mapsto & -\frac{2}{\omega}\big(b_{i} - \frac{\omega}{2} (r_{i - 1} + r_{i}) \big) \\
    r_{i} \mapsto & -\frac{2}{\omega}\cdot r_{i} \\
\end{align*}

\begin{align*}
    N : & \\
    b_{i} \mapsto & -\frac{2}{\omega}(1 - \omega)b_{i} \\
    r_{i} \mapsto & -\frac{2}{\omega}\big((1 - \omega)r_{i} + \frac{\omega}{2} (b_{i} + b_{i + 1}) \big) \\
\end{align*}

\begin{align*}
    M^{-1} : & \\
    b_{i} \mapsto & -\frac{\omega}{2}\big(b_{i} + \frac{\omega}{2} (r_{i - 1} + r_{i}) \big) \\
    r_{i} \mapsto & -\frac{\omega}{2} r_{i} \\
\end{align*}

To check this is an inverse: trivial on the red part.

\begin{align*}
b_i \mapsto -\frac{2}{\omega}b_{i} + r_{i - 1} + r_{i} \\
    \mapsto  b_{i} +\frac{\omega}{2} (r_{i - 1} + r_{i}) - \frac{\omega}{2} r_{i - 1} - \frac{\omega}{2} r_i
\end{align*}

To get rid of the fractions, we rewrite $M^{-1}(Nu + h^2 f)$ to 

\[
    (-2 / \omega M^{-1})\left((-\omega / 2 N)(u) - \omega / 2 h^2 f\right).
\]

\textcolor{red}{It appears to need to be + omega / 2 h\^2 f}

We write $S[a, b, c]$ for the stencil function 

\[
    x[i] = a * x[i - 1] + b * x[i] + c * x[i + 1]
\]

and $a$ for $x \mapsto ax$. Note that linear map $e_i \mapsto e_{i - 1} + e_i$
corresponds with $S[0, 1, 1]$ and not $S[1, 1, 0]$!

With this notation, we can draw 
$-2 / \omega M$ and $-\omega / 2 N$ in a commutative diagram.

\begin{tikzcd}[column sep = huge, row sep = huge]
    B \arrow[d, "(1 - \omega)"']  & 
    R \oplus B \arrow[d, dotted, pos = 0.4, "-\omega / 2 N"']& 
    R \arrow[d, "(1 - \omega)"] 
      \arrow[dll, "{\frac{\omega}{2} S([1, 1, 0])}"]\\
    B \arrow[d, "\id"'] 
      \arrow[drr, pos=0.3, "{\frac{\omega}{2} S[0, 1, 1]}"'] & 
    R \oplus B \arrow[d, pos = 0.3, dotted, "-2 / \omega M^{-1}"] & 
    R \arrow[d, "\id"] \\
    B & R \oplus B & R \\
\end{tikzcd}

the -1 is reversed, if $e_{i} \mapsto e_{i - 1} + e_i$, then
$e_{j} = e_{j} + e_{j + 1}$. 

We can move the R $(1 - \omega)$ to the second function, and then both 
$R \otimes I$ and $I \otimes R$ are the identity good.

\section{Rank polymorphic}

We have

\[
\nabla f: \mathbb{R^{n_1}} \times \cdots \times \mathbb{R^{n_d}} \to \mathbb{R} 
= \sum_{i = 1}^{d} \frac{\partial^2}{\partial x_i} f,
\]

so discretizing gives

\[
\nabla f: \mathbb{R^{n_1}} \times \cdots \times \mathbb{R^{n_d}} \to \mathbb{R} 
= \sum_{i = 1}^d A_{n_i} \otimes \bigotimes_{j \neq i} I_{n_j}
\]

In 2D

\begin{align*}
    A \otimes I + I \otimes A = (M - N) \otimes I + I \otimes (M - N) = \\
    M \otimes I + I \otimes M - (N \otimes I + I \otimes N) \\
\end{align*}

this gives rise to the similar update

\begin{align*}
    u^{(k + 1)} = (M \otimes I + I \otimes M)^{-1} 
            \left((N \otimes I + I \otimes N)u^{(k)} + f \right)
\end{align*}

\begin{align*}
    N \otimes I: & \\
    b_i \otimes b_j = & -\frac{2}{\omega}(1 - \omega)b_i \otimes b_j            \\
                      & -\frac{2}{\omega}\frac{\omega}{2} r_{i - 1} \otimes b_j \\
                      & -\frac{2}{\omega}\frac{\omega}{2} r_i       \otimes b_j \\
\end{align*}

\begin{align*}
    I \otimes N: & \\
    b_i \otimes b_j = & -\frac{2}{\omega}(1 - \omega)b_i \otimes b_j            \\
                      & -\frac{2}{\omega}\frac{\omega}{2} b_j \otimes r_{i - 1} \\
                      & -\frac{2}{\omega}\frac{\omega}{2} b_j \otimes r         \\
\end{align*}

So

\begin{align*}
    -\frac{2}{\omega} (I \otimes N + N \otimes I): & \\
    b_i \otimes b_j = & 2(1 - \omega)b_i \otimes b_j           \\
                      & \frac{\omega}{2} r_{i - 1} \otimes b_j \\
                      & \frac{\omega}{2} r_i       \otimes b_j \\
                      & \frac{\omega}{2} b_j \otimes r_{i - 1} \\
                      & \frac{\omega}{2} b_j \otimes r         \\
\end{align*}

\begin{align*}
    N \otimes I + I \otimes N: & \\
    r_i \otimes r_j = & -\frac{4}{\omega}(1 - \omega)r_i \otimes r_j \\
    r_i \otimes b_j = & -\frac{2}{\omega}(1 - \omega)r_i \otimes b_j \\
                      & +\frac{2}{\omega}\frac{\omega}{2} r_i \otimes r_j \\
                      & +\frac{2}{\omega}\frac{\omega}{2} r_i \otimes r_{j - 1} \\
\end{align*}

\newpage

\section{Stencil RB}

What if $b_{i} = e_{2i}$ and $r_i = e_{n - (2i + 1)}$?

\begin{align*}
    b_i = e_{2i} \mapsto & e_{2i - 1} -2 e_{2i} + e_{2i + 1} \\
                    & = e_{n - (n - 2i + 1)} -2 e_{2i} + e_{n - (n - 2i - 1)} \\
                    & = e_{n - (2(n/2 - i) + 1)} - 2e_{2i}
                        + e_{n - (2(n/2 - i - 1) + 1)} \\
                    & = r_{n / 2 - i} -2 b_{i} + r_{n / 2 - i - 1}
\end{align*}

\begin{align*}
    r_i = e_{n - (2i + 1)} \mapsto & e_{n - 2i - 2} -2 e_{n - (2i + 1)} + 
                                     e_{n - 2i} \\
                    & = e_{2(n / 2 - i - 1)} -2 e_{n - (2i + 1)} + 
                                     e_{2(n / 2 - i)} \\
                    & = b_{n / 2 - i - 1} -2 r_{i} + b_{n / 2 - i}
\end{align*}

\newpage

Let $A$ be a finite element discretization of $f \mapsto \frac{d^2}{d x^2} f$.
The Laplacian operator in higher dimensions consists of tensor product of $A$ 
and identities. In 2D the Laplacian $\nabla$ is $A \otimes I + I \otimes A$. 
I feel like a rank-polymorphic SaC program should exist for solving 
$\nabla u = f$ for $u$.

Methods for solving this equation are based on splitting $A = M - N$ for 
suitable maps $M$ and $N$ and then iteratively updating 
$u^{(k + 1)} = M^{-1}(Nu^{(k)} + f)$. In categorical terms, they use (co)products
in the following way. Let $V$ be a $n$-dimensional vector space, 
and let $B$(lack), $R$(ed) be $n / 2$-dimensional vector spaces. If we can find
an isomorphism $\sigma$ such that the following diagram commutes,
we can compose them so we don't have to explicitly compute $\sigma$, except for
$\sigma f$ and $\sigma^{-1}u$ on the final solution. The structure on the
right-hand side allows for an easily invertible $M$.

\begin{tikzcd}
    V \arrow[d, "A"'] \arrow[r, "\sigma"] & B \oplus R \arrow[d, "something\ nice"] \\
    V \arrow[r, "\sigma"'] & B \oplus R \\
\end{tikzcd}

For $e$ basis $V$, $r$ basis $R$ and $b$ basis $B$, the standard choice is
$b_i \mapsto e_{2i}$, $r_i \mapsto e_{2i + 1}$. 

Tensor products distribute over direct sums, which gives a rank-polymorphic
formulation for FFTs, so I hope that we can do something similar here. I am
however less and less convinced that this works for the standard choice of 
$\sigma$. Maybe we need $r_i \mapsto e_{n - (2i + 1)}$ instead of $e_{2i + 1}$.

\newpage

For $e$ basis $V$, $r$ basis $R$ and $b$ basis $B$, we define $\sigma$ by
$b_i \mapsto e_{2i}$, $r_i \mapsto e_{n - (2i + 1)}$. 

\begin{tikzcd}
    V \arrow[d, "A"'] \arrow[r, "\sigma"] & B \oplus R \arrow[d, "A^{rb}"] \\
    V \arrow[r, "\sigma"'] & B \oplus R \\
\end{tikzcd}

Here $A^{rb}$ is the unique morphism making the following diagram commute
(the direct sum $\oplus$ is both product and coproduct in Vect).
The function $S$ is a reversed stencil $x_{n - i} = w_1 (x_i + x_{i + 1})$.

\begin{tikzcd}[column sep = huge, row sep = huge]
    B \arrow[d, "w_0"'] \arrow[r, "\iota"] \arrow[drr, pos = 0.3, "S"] & 
        B \oplus R \arrow[d, pos = 0.7, dotted, "A^{rb}"] &  
        R \arrow[dll, pos = 0.3, "S"'] \arrow[d, "w_0"] \arrow[l, "\iota"'] \\
    B & B \oplus R \arrow[r, "\pi"'] \arrow[l, "\pi"] & R
\end{tikzcd}

$A^{rb} \otimes I + I \otimes A^{rb}$

$S[a, b, c] \otimes I + I \otimes S[d, e, f]$ becomes

\begin{align*}
    S[&[0,& d    , & 0], \\
      &[a,& b + e, & c], \\
      &[0,& f    , & 0]]
\end{align*}

\end{document}