# A Biological Example

Suppose we are observing exponential growth $\dot{y} = \theta y$ but we don't know $\theta$ and wish to estimate it. We could assume $\theta \sim {\cal{N}}(\mu, \sigma^2)$ and use least squares or better something like Markov Chain Monte Carlo or Hamiltonian Monte Carlo and any observations to infer $\mu$ and $\sigma$. However, we might want to model that the further we go into the future, the less we know about $\theta$. We can write our system as as

$$
\begin{aligned}
\mathrm{d}y & = \theta y\mathrm{d}t \\
\mathrm{d}\theta & = \sigma\mathrm{d}W_t
\end{aligned}
$$

where $W_t$ is Brownian Motion.

# Fokker-Planck

We can use the [Fokker-Planck equation](https://en.wikipedia.org/wiki/Fokker%E2%80%93Planck_equation) to convert a stochastic differential equation into a partial differential equation.

$$
d \mathbf{X}_{t}=\boldsymbol{\mu}\left(\mathbf{X}_{t}, t\right) d t+\boldsymbol{\sigma}\left(\mathbf{X}_{t}, t\right) d \mathbf{W}_{t}
$$

$$
\frac{\partial}{\partial t} p(t, \mathbf{x})+\sum_{k=1}^{n} \frac{\partial}{\partial x_{k}}\left({\mu}_{k}(t, \mathbf{x}) p(t, \mathbf{x})\right)=\frac{1}{2} \sum_{j=1, k=1}^{n} \frac{\partial^{2}}{\partial x_{j} \partial x_{k}}\left[\left(\sigma(t, \mathbf{x}) \sigma^{T}(t, \mathbf{x})\right)_{j k} p(t, \mathbf{x})\right]
$$

For our particular system we have

$$
\frac{\partial}{\partial t} p(t, y, \theta)+\frac{\partial}{\partial y}\left({\mu}_{1}(t, y, \theta) p(t, y, \theta)\right)+\frac{\partial}{\partial \theta}\left({\mu}_{2}(t, y, \theta) p(t, y, \theta)\right)=\frac{1}{2}\left[\sigma_{y}^{2} \frac{\partial^{2}}{\partial y^{2}} p(t, y, \theta)+\sigma_{\theta}^{2} \frac{\partial^{2}}{\partial \theta^{2}} p(t, y, \theta)\right]
$$

And since $\mu_1 = \theta y$, $\mu_2 = 0$ and $\sigma_y = 0$ this further simplifies to

$$
\frac{\partial}{\partial t} p(t, y, \theta)+\frac{\partial}{\partial y}(\theta y p(t, y, \theta))=\sigma_{\theta}^{2} \frac{\partial^{2}}{\partial \theta^{2}} p(t, y, \theta)
$$

We can note two things:

* This is an advection / diffusion equation with two spatial variables ($y$ and $\theta$).
* If $\sigma_\theta = 0$ then this is a transport (advection?) equation.

$$
\frac{\partial}{\partial t} p(t, y, \theta)+\frac{\partial}{\partial y}(\theta y p(t, y, \theta))=0
$$

Notice that there is nothing stochastic about the biology but we
express our uncertainty about the parameter by making it a
time-varying stochastic variable which says the further we go into the
future the less certain we are about it.

We are going to turn this into a Fokker-Planck equation which we can
then solve using e.g. the method of lines. But before turning to
Fokker-Planck, let's show that we can indeed solve a diffusion
equation using the method of lines.

# Warming Up

Let us solve the heat equation

$$
\frac{\partial u}{\partial t}=k_{x} \frac{\partial^{2} u}{\partial x^{2}}+k_{y} \frac{\partial^{2} u}{\partial y^{2}}+h
$$

with initial condition $u(0, x, y) = 0$ and stationary boundary conditions

$$
\frac{\partial u}{\partial t}(t, 0, y)=\frac{\partial u}{\partial t}(t, 1, y)=\frac{\partial u}{\partial t}(t, x, 0)=\frac{\partial u}{\partial t}(t, x, 1)=0
$$

and a periodic heat source

$$
h(x, y)=\sin (\pi x) \sin (2 \pi y)
$$

This has analytic solution

$$
u(t, x, y)=\frac{1-e^{-\left(k_{x}+4 k_{y}\right) \pi^{2} t}}{\left(k_{x}+4 k_{y}\right) \pi^{2}} \sin (\pi x) \sin (2 \pi y)
$$

The spatial derivatives are computed using second-order centered differences, with the data distributed over $n_x \times n_y$ points on a uniform spatial grid.

$$
u_{i\,j}(t) \triangleq u\left(t, x_{i}, y_{j}\right), \quad x_{i} \triangleq i \Delta x, \quad 0 \leq i \leq n_x-1, \quad  y_{j} \triangleq j \Delta y, \quad 0 \leq j \leq n_y-1
$$

$$
\begin{align}
u_{x x} &= \frac{u_{i+1\,j}-2 u_{i\,j}+u_{i-1\,j}}{\Delta x^{2}} \\
u_{y y} &= \frac{u_{i\,j+1}-2 u_{i\,j}+u_{i\,j-1}}{\Delta y^{2}}
\end{align}
$$

$$
\dot{u}_{i\, j} = \frac{k_x}{(\Delta x)^2}({u_{i+1\,j}-2 u_{i\,j}+u_{i-1\,j}})
                + \frac{k_y}{(\Delta y)^2}({u_{i\,j+1}-2 u_{i\,j}+u_{i\,j-1}})
                + h_{i\, j}
$$

We could try using [Naperian functors and APL-like programming in Haskell](https://www.cs.ox.ac.uk/people/jeremy.gibbons/publications/aplicative.pdf) via this [library](http://hackage.haskell.org/package/Naperian). But the performance is terrible (or it could be that the author's implementation was terrible). Moreover, applied mathematicans tend to think of everything as matrices and vectors. But flattening the above tensor operation into a matrix operation is not entirely trivial. Although the Haskell Ecosystem's support for symbolic mathematics is very rudimentary, we can use what there is to convince ourselves that we haven't made too many errors in the transcription.


```haskell
{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE OverloadedLists     #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE FlexibleInstances   #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE GADTs               #-}
{-# LANGUAGE TypeApplications    #-}
{-# LANGUAGE TypeOperators       #-}

import Data.Number.Symbolic
import qualified Data.Number.Symbolic as Sym
import Data.Proxy

import qualified Naperian as N
import qualified GHC.TypeLits as M
import           Data.Functor

import           Numeric.Sundials.ARKode.ODE
import           Numeric.LinearAlgebra

import           System.IO
```

We can re-write the semi-discretized equations in a tensor form from which we can derive an implementation.

$$
\dot{u}_{i\, j} = \frac{k_x}{(\Delta x)^2}\sum_{k=0}^{n_x-1}\sum_{l=0}^{n_y-1}A_{i\,j\,k\,l} u_{k\,l}
                + \frac{k_y}{(\Delta y)^2}\sum_{k=0}^{n_x-1}\sum_{l=0}^{n_y-1}B_{i\,j\,k\,l} u_{k\,l}
                + h_{i\, j}
$$


$$
A_{i\, j\, l\, k} =
\begin{cases}
0,& \text{if } i = 0 \\
0,& \text{if } j = 0 \\
1,& \text{if } k = i-1 \text{ and } l = j \\
-2,& \text{if } k = i \text{ and } l = j \\
1,& \text{if } k = i+1 \text{ and } l = j \\
0,& \text{if } i = n_x - 1 \\
0,& \text{if } j = n_y - 1 \\
0,& \text{otherwise}
\end{cases}
$$


$$
B_{i\, j\, l\, k} =
\begin{cases}
0,& \text{if } i = 0 \\
0,& \text{if } j = 0 \\
1,& \text{if } k = i \text{ and } l = j - 1 \\
-2,& \text{if } k = i \text{ and } l = j \\
1,& \text{if } k = i \text{ and } l = j + 1 \\
0,& \text{if } i = n_x - 1 \\
0,& \text{if } i = n_y - 1 \\
0,& \text{otherwise}
\end{cases}
$$

Let's write down the tensor $A$ in Haskell


```haskell
preA :: forall b m n . (M.KnownNat m, M.KnownNat n, Num b) =>
        N.Hyper '[N.Vector n, N.Vector m, N.Vector n, N.Vector m] b
preA = N.Prism $ N.Prism $ N.Prism $ N.Prism $ N.Scalar $
      N.viota @m <&> (\(N.Fin x) ->
      N.viota @n <&> (\(N.Fin w) ->
      N.viota @m <&> (\(N.Fin v) ->
      N.viota @n <&> (\(N.Fin u) ->
      (f m n x w v u)))))
        where
          m = fromIntegral $ M.natVal (undefined :: Proxy m)
          n = fromIntegral $ M.natVal (undefined :: Proxy n)
          f p q i j k l | i == 0               = 0
                        | j == 0               = 0
                        | i == p - 1           = 0
                        | j == q - 1           = 0
                        | k == i - 1 && l == j = 1
                        | k == i     && l == j = -2
                        | k == i + 1 && l == j = 1
                        | otherwise            = 0
```

We can concretize this to symbolic numbers


```haskell
a :: forall a m n . (M.KnownNat m, M.KnownNat n, Floating a, Eq a) =>
      N.Hyper '[N.Vector n, N.Vector m, N.Vector n, N.Vector m] (Sym a)
a = N.binary (*) (N.Scalar $ var "a") preA
```

And do the same for the tensor $B$


```haskell
preB :: forall b m n . (M.KnownNat m, M.KnownNat n, Num b) =>
        N.Hyper '[N.Vector n, N.Vector m, N.Vector n, N.Vector m] b
preB = N.Prism $ N.Prism $ N.Prism $ N.Prism $ N.Scalar $
      N.viota @m <&> (\(N.Fin x) ->
      N.viota @n <&> (\(N.Fin w) ->
      N.viota @m <&> (\(N.Fin v) ->
      N.viota @n <&> (\(N.Fin u) ->
      (f m n x w v u)))))
        where
          m = fromIntegral $ M.natVal (undefined :: Proxy m)
          n = fromIntegral $ M.natVal (undefined :: Proxy n)
          f :: Int -> Int -> Int -> Int -> Int -> Int -> b
          f p q i j k l | i == 0                   = 0
                        | j == 0                   = 0
                        | i == p - 1               = 0
                        | j == q - 1               = 0
                        | k == i     && l == j - 1 = 1
                        | k == i     && l == j     = -2
                        | k == i     && l == j + 1 = 1
                        | otherwise                = 0
```


```haskell
b :: forall a m n . (M.KnownNat m, M.KnownNat n, Floating a, Eq a) =>
           N.Hyper '[N.Vector n, N.Vector m, N.Vector n, N.Vector m] (Sym a)
b = N.binary (*) (N.Scalar $ var "b") preB
```

We can check that our implementation matches the mathematical formula
by rendering it as a $\LaTeX$.


```haskell
ps :: forall m n . (M.KnownNat m, M.KnownNat n) =>
       [N.Vector n (N.Vector m ((Int, Int), Sym Double))]
ps = N.elements $ N.crystal $ N.crystal $ N.hzipWith (,) ss rhs
  where
    h = N.Prism $ N.Prism $ N.Scalar $
      N.viota @n <&> (\(N.Fin x) ->
      N.viota @m <&> (\(N.Fin w) ->
      var ("u_{" ++ show x ++ "," ++ show w ++ "}")))
    rhs = N.foldrH (+) 0 $ N.foldrH (+) 0 $ N.binary (*) preFoo h
    preFoo = N.binary (+) (a @Double @n @m) (b @Double @n @m)
    ss = N.Prism $ N.Prism $ N.Scalar $
         N.viota @n <&> (\(N.Fin x) ->
         N.viota @m <&> (\(N.Fin w) -> (x,w)))
```


```haskell
eqns = mapM_ putStrLn $ zipWith (++) aaa bbb
  where
    aaa = concatMap (N.elements . N.Prism . N.Prism . N.Scalar) $
          fmap (fmap (fmap ((\(i, j)-> "u_{" ++ show i ++ show j ++ "} &= ") . fst))) x
    bbb = concatMap (N.elements . N.Prism . N.Prism . N.Scalar) $
          fmap (fmap (fmap ((++ " \\\\") . show . snd))) x
    x = ps @4 @3
```


```haskell
eqns
```


    u_{00} &= 0.0 \\
    u_{01} &= 0.0 \\
    u_{02} &= 0.0 \\
    u_{03} &= 0.0 \\
    u_{10} &= 0.0 \\
    u_{11} &= a*u_{0,1}+b*u_{1,0}+((-2.0)*a+(-2.0)*b)*u_{1,1}+b*u_{1,2}+a*u_{2,1} \\
    u_{12} &= a*u_{0,2}+b*u_{1,1}+((-2.0)*a+(-2.0)*b)*u_{1,2}+b*u_{1,3}+a*u_{2,2} \\
    u_{13} &= 0.0 \\
    u_{20} &= 0.0 \\
    u_{21} &= 0.0 \\
    u_{22} &= 0.0 \\
    u_{23} &= 0.0 \\


And then getting our notebook to render the $\LaTeX$.

$$
\begin{aligned}
u_{00} &= 0.0 \\
u_{01} &= 0.0 \\
u_{02} &= 0.0 \\
u_{03} &= 0.0 \\
u_{10} &= 0.0 \\
u_{11} &= a*u_{0,1}+b*u_{1,0}+((-2.0)*a+(-2.0)*b)*u_{1,1}+b*u_{1,2}+a*u_{2,1} \\
u_{12} &= a*u_{0,2}+b*u_{1,1}+((-2.0)*a+(-2.0)*b)*u_{1,2}+b*u_{1,3}+a*u_{2,2} \\
u_{13} &= 0.0 \\
u_{20} &= 0.0 \\
u_{21} &= 0.0 \\
u_{22} &= 0.0 \\
u_{23} &= 0.0 \\
\end{aligned}
$$

Now we have checked that our tensors look correct (at least for a particular and small tensor) we can try solving the system numerically

Spatial mesh size:


```haskell
nx, ny :: Int
nx = 3
ny = 4
```

Heat conductivity coefficients:


```haskell
kx, ky :: Floating a => a
kx = 0.5
ky = 0.75
```

x and y mesh spacing:


```haskell
dx :: Floating a => a
dx = 1 / (fromIntegral nx - 1)
dy :: Floating a => a
dy = 1 / (fromIntegral ny - 1)
```


```haskell
c1, c2 :: Floating a => a
c1 = kx/dx/dx
c2 = ky/dy/dy
```

Now we make the tensors more concrete by ensuring their elements come from `Floating`


```haskell
bNum :: forall a m n . (M.KnownNat m, M.KnownNat n, Floating a) =>
        N.Hyper '[N.Vector n, N.Vector m, N.Vector n, N.Vector m] a
bNum = N.binary (*) (N.Scalar c1) preB
```


```haskell
aNum :: forall a m n . (M.KnownNat m, M.KnownNat n, Floating a) =>
        N.Hyper '[N.Vector n, N.Vector m, N.Vector n, N.Vector m] a
aNum = N.binary (*) (N.Scalar c2) preA
```

Again we flatten the system into a matrix form so we can check everything looks as it should


```haskell
bigA :: Matrix Double
bigA = fromLists $
       fmap (N.elements . N.Prism . N.Prism . N.Scalar) $
       N.elements $ N.crystal $ N.crystal $ N.binary (+)
       (aNum @Double @4 @3) (bNum @Double @4 @3)
```


```haskell
bigA
```


    (12><12)
     [ 0.0,  0.0, 0.0, 0.0,   0.0, 0.0, 0.0,   0.0, 0.0, 0.0,  0.0, 0.0
     , 0.0,  0.0, 0.0, 0.0,   0.0, 0.0, 0.0,   0.0, 0.0, 0.0,  0.0, 0.0
     , 0.0,  0.0, 0.0, 0.0,   0.0, 0.0, 0.0,   0.0, 0.0, 0.0,  0.0, 0.0
     , 0.0,  0.0, 0.0, 0.0,   0.0, 0.0, 0.0,   0.0, 0.0, 0.0,  0.0, 0.0
     , 0.0, 6.75, 0.0, 2.0, -17.5, 2.0, 0.0,  6.75, 0.0, 0.0,  0.0, 0.0
     , 0.0,  0.0, 0.0, 0.0,   0.0, 0.0, 0.0,   0.0, 0.0, 0.0,  0.0, 0.0
     , 0.0,  0.0, 0.0, 0.0,   0.0, 0.0, 0.0,   0.0, 0.0, 0.0,  0.0, 0.0
     , 0.0,  0.0, 0.0, 0.0,  6.75, 0.0, 2.0, -17.5, 2.0, 0.0, 6.75, 0.0
     , 0.0,  0.0, 0.0, 0.0,   0.0, 0.0, 0.0,   0.0, 0.0, 0.0,  0.0, 0.0
     , 0.0,  0.0, 0.0, 0.0,   0.0, 0.0, 0.0,   0.0, 0.0, 0.0,  0.0, 0.0
     , 0.0,  0.0, 0.0, 0.0,   0.0, 0.0, 0.0,   0.0, 0.0, 0.0,  0.0, 0.0
     , 0.0,  0.0, 0.0, 0.0,   0.0, 0.0, 0.0,   0.0, 0.0, 0.0,  0.0, 0.0 ]



```haskell
h :: forall m n a . (M.KnownNat m, M.KnownNat n, Floating a) =>
                     N.Hyper '[N.Vector m, N.Vector n] a
h = N.Prism $ N.Prism $ N.Scalar $
     N.viota @n <&> (\(N.Fin x) ->
     N.viota @m <&> (\(N.Fin w) ->
     sin (pi * (fromIntegral w) * dx)
             * sin (2 * pi * (fromIntegral x) * dy)))
```


```haskell
c :: Vector Double
c = fromList $ N.elements (h @3 @4 @Double)
```


```haskell
t0, tf :: Double
t0 = 0.0
tf = 0.3

bigNt :: Int
bigNt = 20

dTout :: Double
dTout = (tf - t0) / (fromIntegral bigNt)

ts :: [Double]
ts = map (dTout *) $ map fromIntegral [1..bigNt]

sol :: Matrix Double
sol = odeSolveV SDIRK_5_3_4' Nothing 1.0e-5 1.0e-10 (const bigU') (assoc (nx * ny) 0.0 [] :: Vector Double) (fromList $ ts)
  where
    bigU' bigU = bigA #> bigU + c
```


```haskell
main :: IO ()
main = do
  h1 <- openFile "Haskell.txt" WriteMode
  mapM_ (hPutStrLn h1) $ map (concatMap (' ':)) $ map (map show) $ toLists sol
  hClose h1
  mapM_ (\i -> putStrLn $ show $ sqrt $ (sol!i) <.> (sol!i) / (fromIntegral nx) / (fromIntegral ny)) ([0 .. length ts - 1] :: [Int])
```


```haskell
main
```


    0.0
    4.445792552162428e-3
    7.536496589199581e-3
    9.685474590386838e-3
    1.1179180479250954e-2
    1.2215783380962438e-2
    1.2936469361667507e-2
    1.3437435527874617e-2
    1.3785614451622913e-2
    1.4027670958010267e-2
    1.4196099215739022e-2
    1.4312878038957881e-2
    1.4394259941155774e-2
    1.4450706139605892e-2
    1.4490277045502337e-2
    1.451741051111616e-2
    1.4536171194645382e-2
    1.4549482462569737e-2
    1.45586417237852e-2
    1.4564936158380733e-2


The grid is unrealistically coarse. Let's check the Haskell implementation against a [reference C implmentation](https://github.com/LLNL/sundials/blob/master/examples/arkode/CXX_parallel/README)  with a better grid size.

```
nix-shell -I nixpkgs=https://github.com/NixOS/nixpkgs/archive/19.09.tar.gz
mpicxx ark_heat2D.cpp -lm -lsundials_arkode -lsundials_nvecparallel -o ark_heat2D
mpiexec -n 1 ./ark_heat2D
python plot_heat2D.py
```

We could run with more processors but it's easier to modify the python to work on the Haskell output if we don't.

```
nix-shell -I nixpkgs=https://github.com/NixOS/nixpkgs/archive/19.09.tar.gz
ghc -O2 -fforce-recomp Heat2D.hs -main-is Heat2D -o Heat2D
./Heat2D
python plot_heat2D.py
```



903f76edcc3281a1ad272c2e992376d2e2266fa8
ipynb
Jupyter Notebook
Two Dimensional Heat Equation.ipynb
idontgetoutmuch/Diffusions
b07eb63aeb8d7b12e5a2ba18a01eca520af3e128
BSD-3-Clause
Two Dimensional Heat Equation.ipynb
idontgetoutmuch/Diffusions
b07eb63aeb8d7b12e5a2ba18a01eca520af3e128
BSD-3-Clause
Two Dimensional Heat Equation.ipynb
idontgetoutmuch/Diffusions
b07eb63aeb8d7b12e5a2ba18a01eca520af3e128
BSD-3-Clause
Qwen/Qwen-72B
1. YES
2. YES