Discretizing Linear Operators.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Overview\n",
    "\n",
    "This notebook includes prototypical code for the operator overview write-up. It does not depend on any exisiting code in DiffEqOperators, but it should be easy to modify the current codebase to achieve the same functionality.\n",
    "\n",
    "To make things simpler, everything is assumed to use `Float64` datatype. I will also always use the out-of-place convention (i.e. `*` instead of `A_mul_B!`).\n",
    "\n",
    "The naming of different operators will use the following convention, as in the writeup:\n",
    "\n",
    "- `L` denotes a (quasi)linear operator.\n",
    "\n",
    "- `A` denotes an operator that can generally (but not necessarily) be affine.\n",
    "\n",
    "- `b` denotes the bias of an affine operator\n",
    "\n",
    "- (Not discussed but I guess we may as well make the call?) `M` for raw matrices (`AbstractMatrix`) and `x`, `y`, `u`, etc for vectors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import Base: +, -, *, \\\n",
    "\n",
    "abstract type DiffEqOperator end\n",
    "abstract type DiffEqLinearOperator <: DiffEqOperator end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 1. Constant case (i.e. no `update_coefficients!`)\n",
    "\n",
    "## 1.1 Abstract operator interface\n",
    "\n",
    "Below is a simplified operator interface used in my new interface draft. Basically we want to express lazy addition and multiplication of linear operators naturally using `+` and `*`. The `as_array` interface returns the most suitable (dense/sparse) representation of the underlying operator as an `AbstractMatrix` (or should we treat this as a type conversion and just use `AbstractMatrix`?).\n",
    "\n",
    "The affine operator is defined in such a way that arithmetic on them and linear operators always yield an `DiffEqAffineOperator`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We should improve type stability by including more inferrable fields in the type signature\n",
    "# But for the sake of this notebook I'll make things simple\n",
    "struct DiffEqArrayOperator <: DiffEqLinearOperator\n",
    "    M::AbstractMatrix{Float64}\n",
    "end\n",
    "\n",
    "*(L::DiffEqArrayOperator, x::Vector{Float64}) = L.M * x\n",
    "as_array(L::DiffEqArrayOperator) = L.M\n",
    "\n",
    "struct DiffEqOperatorCombination <: DiffEqLinearOperator\n",
    "    cs::Tuple{Vararg{Float64}} # Coefficients\n",
    "    Ls::Tuple{Vararg{DiffEqLinearOperator}}\n",
    "end\n",
    "\n",
    "*(L::DiffEqOperatorCombination, x::Vector{Float64}) = sum(ck * (Lk*x) for (ck,Lk) in zip(L.cs,L.Ls))\n",
    "as_array(L::DiffEqOperatorCombination) = sum(ck * as_array(Lk) for (ck,Lk) in zip(L.cs,L.Ls))\n",
    "\n",
    "struct DiffEqOperatorComposition <: DiffEqLinearOperator\n",
    "    Ls::Tuple{Vararg{DiffEqLinearOperator}}\n",
    "end\n",
    "\n",
    "*(L::DiffEqOperatorComposition, x::Vector{Float64}) = foldl((u, Lk) -> Lk*u, x, L.Ls)\n",
    "as_array(L::DiffEqOperatorComposition) = prod(as_array, reverse(L.Ls))\n",
    "\n",
    "struct DiffEqAffineOperator <: DiffEqOperator\n",
    "    L::DiffEqLinearOperator\n",
    "    b::Vector{Float64}\n",
    "end\n",
    "\n",
    "*(A::DiffEqAffineOperator, x::Vector{Float64}) = A.L * x + A.b;"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Operator arithmetic\n",
    "# Addition\n",
    "+(L1::DiffEqOperatorCombination, L2::DiffEqOperatorCombination) = DiffEqOperatorCombination(\n",
    "    (L1.cs...,L2.cs...), (L1.Ls...,L2.Ls...))\n",
    "+(L1::DiffEqOperatorCombination, L2::DiffEqLinearOperator) = DiffEqOperatorCombination((L1.cs...,1.0), (L1.Ls...,L2))\n",
    "+(L1::DiffEqLinearOperator, L2::DiffEqOperatorCombination) = L2 + L1\n",
    "+(L1::DiffEqLinearOperator, L2::DiffEqLinearOperator) = DiffEqOperatorCombination((1.0,1.0), (L1,L2))\n",
    "\n",
    "+(A1::DiffEqAffineOperator, L2::DiffEqLinearOperator) = DiffEqAffineOperator(A1.L + L2, A1.b)\n",
    "+(L1::DiffEqLinearOperator, A2::DiffEqAffineOperator) = A2 + L1\n",
    "+(A1::DiffEqAffineOperator, A2::DiffEqAffineOperator) = DiffEqAffineOperator(A1.L + A2.L, A1.b + A2.b)\n",
    "\n",
    "# Scalar multiplication\n",
    "*(α::Float64, L::DiffEqOperatorCombination) = DiffEqOperatorCombination(α.*L.cs, L.Ls)\n",
    "*(α::Float64, L::DiffEqLinearOperator) = DiffEqOperatorCombination((α,), (L,))\n",
    "*(α::Float64, A::DiffEqAffineOperator) = DiffEqAffineOperator(α * A.L, α * A.b)\n",
    "*(A::DiffEqOperator, α::Float64) = α * A\n",
    "\n",
    "# Subtraction/unary minus\n",
    "-(A::DiffEqOperator) = (-1.0) * A\n",
    "-(A1::DiffEqOperator, A2::DiffEqOperator) = A1 + (-A2)\n",
    "\n",
    "# Multiplication\n",
    "# Note the application order\n",
    "*(L1::DiffEqOperatorComposition, L2::DiffEqOperatorComposition) = DiffEqOperatorComposition((L2.Ls..., L1.Ls...))\n",
    "*(L1::DiffEqLinearOperator, L2::DiffEqOperatorComposition) = DiffEqOperatorComposition((L2.Ls..., L1))\n",
    "*(L1::DiffEqOperatorComposition, L2::DiffEqLinearOperator) = DiffEqOperatorComposition((L2, L1.Ls...))\n",
    "*(L1::DiffEqLinearOperator, L2::DiffEqLinearOperator) = DiffEqOperatorComposition((L2, L1))\n",
    "\n",
    "*(L1::DiffEqLinearOperator, A2::DiffEqAffineOperator) = DiffEqAffineOperator(L1 * A2.L, L1 * A2.b)\n",
    "*(A1::DiffEqAffineOperator, L2::DiffEqLinearOperator) = DiffEqAffineOperator(A1.L * L2, A1.b)\n",
    "*(A1::DiffEqAffineOperator, A2::DiffEqAffineOperator) = DiffEqAffineOperator(A1.L * A2.L, A1.L * A2.b + A1.b)\n",
    "\n",
    "# Right division (i.e. linear solve)\n",
    "# In the full version we should also include interface to lazy solvers\n",
    "\\(L::DiffEqLinearOperator, y::Vector{Float64}) = as_array(L) \\ y\n",
    "\\(A::DiffEqAffineOperator, y::Vector{Float64}) = A.L \\ (y - A.b);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1.2 Discretization of the differential operator (stencil convolution)\n",
    "\n",
    "The 2nd-order central difference approximation to $\\partial_x^2$ and the 1st-order upwind approximation to $\\partial_x$ are included. The general case can be modified from the stencil convolution code in DiffEqOperators.\n",
    "\n",
    "The upwind operator is implemented differently from DiffEqOperators, where the direction at each point is stored. Here only the left/right operator is constructed and we interpret\n",
    "\n",
    "$$ \\mu(x)\\partial_x = \\mu^+(x)\\partial_x^+ + \\mu^-(x)\\partial_x^- $$\n",
    "\n",
    "(Might not be a favorable approach, especially if we wish to extend to multidimensional upwind operators)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "struct DiffusionOperator <: DiffEqLinearOperator\n",
    "    dx::Float64\n",
    "    m::Int # number of interior points\n",
    "end\n",
    "\n",
    "*(L::DiffusionOperator, x::Vector{Float64}) = [x[i] + x[i+2] - 2*x[i+1] for i in 1:L.m] / L.dx^2\n",
    "as_array(L::DiffusionOperator) = spdiagm((ones(L.m), -2*ones(L.m), ones(L.m)), (0,1,2)) / L.dx^2;"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "struct DriftOperator <: DiffEqLinearOperator\n",
    "    dx::Float64\n",
    "    m::Int # number of interior points\n",
    "    direction::Bool # true = right, false = left\n",
    "end\n",
    "\n",
    "function *(L::DriftOperator, x::Vector{Float64})\n",
    "    if L.direction # right drift\n",
    "        [x[i+1] - x[i] for i in 1:L.m] / L.dx\n",
    "    else # left drift\n",
    "        [x[i+2] - x[i+1] for i in 1:L.m] / L.dx\n",
    "    end\n",
    "end\n",
    "function as_array(L::DriftOperator)\n",
    "    if L.direction # right drift\n",
    "        spdiagm((-ones(L.m), ones(L.m), zeros(L.m)), (0,1,2)) / L.dx\n",
    "    else # left drift\n",
    "        spdiagm((zeros(L.m), -ones(L.m), ones(L.m)), (0,1,2)) / L.dx\n",
    "    end\n",
    "end;"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 1.3 Boundary extrapolation operator $Q$\n",
    "\n",
    "Question: is it OK to shorthand \"boundary extrapolation operator\" as \"BEOperator\" or simply \"BE\"?\n",
    "\n",
    "A generic $Q$ from the generic boundary condition $Bu = b$ can be a bit difficult to implement, but the simple case of Dirichlet/Neumann BC is easy to handle.\n",
    "\n",
    "It should be easy to modify `AbsorbingBoundaryMap` and `ReflectingBoundaryMap` to incorporate boundaries that are absorbing at one end and reflecting at another.\n",
    "\n",
    "The non-zero Dirichlet/Neumann Qs do not have their own type. Instead we construct the operator as an affine map, with `AbsorbingBoundaryMap` and `ReflectingBoundaryMap` its linear part."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "struct AbsorbingBE <: DiffEqLinearOperator\n",
    "    m::Int # number of interior points\n",
    "end\n",
    "\n",
    "*(Q::AbsorbingBE, x::Vector{Float64}) = [0.0; x; 0.0]\n",
    "as_array(Q::AbsorbingBE) = sparse([zeros(Q.m)'; eye(Q.m); zeros(Q.m)'])\n",
    "\n",
    "struct ReflectingBE <: DiffEqLinearOperator\n",
    "    m::Int # number of interior points\n",
    "end\n",
    "\n",
    "*(Q::ReflectingBE, x::Vector{Float64}) = [x[1]; x; x[end]]\n",
    "as_array(Q::ReflectingBE) = sparse([1.0 zeros(Q.m-1)'; eye(Q.m); zeros(Q.m-1)' 1.0]);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "function Dirichlet_BE(m::Int, bl::Float64, br::Float64)\n",
    "    # y[1] = bl, y[end] = br\n",
    "    L = AbsorbingBE(m)\n",
    "    b = [bl; zeros(m); br]\n",
    "    DiffEqAffineOperator(L, b)\n",
    "end\n",
    "\n",
    "function Neumann_BE(m::Int, dx::Float64, bl::Float64, br::Float64)\n",
    "    # (y[2] - y[1])/dx = bl, (y[end] - y[end-1])/dx = br\n",
    "    L = ReflectingBE(m)\n",
    "    b = [-bl*dx; zeros(m); br*dx]\n",
    "    DiffEqAffineOperator(L, b)\n",
    "end;"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Examples of non-zero BC:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "dx = 1.0\n",
    "u = [1.,2.,3.,4.]\n",
    "QD = Dirichlet_BE(4, 10.0, 20.0)\n",
    "QN = Neumann_BE(4, dx, 1.0, 2.0);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "6-element Array{Float64,1}:\n",
       " 10.0\n",
       "  1.0\n",
       "  2.0\n",
       "  3.0\n",
       "  4.0\n",
       " 20.0"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "QD * u"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "6-element Array{Float64,1}:\n",
       " 0.0\n",
       " 1.0\n",
       " 2.0\n",
       " 3.0\n",
       " 4.0\n",
       " 6.0"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "QN * u"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1.4 Constructing operators for the Fokker-Planck equations\n",
    "\n",
    "The most general case is in section 3.5 of the write-up:\n",
    "\n",
    "$$ \\mathcal{L} = \\mu(x)\\partial_x + \\frac{\\sigma(x)^2}{2}\\partial_{xx} $$\n",
    "\n",
    "where the drift term is discretized using upwind operators as described in section 2:\n",
    "\n",
    "$$ \\mu(x)\\partial_x = \\mu^+(x)\\partial_x^+ + \\mu^-(x)\\partial_x^- $$\n",
    "\n",
    "The discretized operators are expressed as the composition of an interior stencil convolution operator `L` and boundary extrapolation operator `Q` (they are generally affine). The drift and diffusion coefficients can be expressed as diagonal matrices (wrapped in a `DiffEqArrayOperator`).\n",
    "\n",
    "The script below can be modified to represent each of the scenarios described in secition 3 of the write-up (except the mixed-BC case)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DiffEqOperatorComposition((AbsorbingBE(10), DiffEqOperatorCombination((1.0, 1.0, 1.0), (DiffEqOperatorComposition((DriftOperator(1.0, 10, true), DiffEqArrayOperator([0.887661 0.0 … 0.0 0.0; 0.0 0.213112 … 0.0 0.0; … ; 0.0 0.0 … 0.436606 0.0; 0.0 0.0 … 0.0 0.569388]))), DiffEqOperatorComposition((DriftOperator(1.0, 10, false), DiffEqArrayOperator([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]))), DiffEqOperatorComposition((DiffusionOperator(1.0, 10), DiffEqArrayOperator([1.89194 0.0 … 0.0 0.0; 0.0 0.597414 … 0.0 0.0; … ; 0.0 0.0 … 0.539805 0.0; 0.0 0.0 … 0.0 0.94818])))))))"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# The grid\n",
    "N = 10\n",
    "dx = 1.0\n",
    "xs = collect(1:N) * dx # interior nodes\n",
    "\n",
    "# Discretization of the differential operators\n",
    "L1p = DriftOperator(dx, N, true)\n",
    "L1m = DriftOperator(dx, N, false)\n",
    "L2 = DiffusionOperator(dx, N)\n",
    "\n",
    "# Boundary operators\n",
    "Q = AbsorbingBE(N)\n",
    "# Q = Neumann_BE(N, dx, 1.0, 2.0)\n",
    "\n",
    "# Coefficients\n",
    "mu = rand(N)\n",
    "mup = [mu[i] > 0.0 ? mu[i] : 0.0 for i in 1:N]\n",
    "mum = [mu[i] < 0.0 ? mu[i] : 0.0 for i in 1:N]\n",
    "sigma = rand(N) + 1.0\n",
    "\n",
    "# Construct the final product\n",
    "Ldrift = DiffEqArrayOperator(Diagonal(mup)) * L1p + DiffEqArrayOperator(Diagonal(mum)) * L1m\n",
    "Ldiffusion = DiffEqArrayOperator(Diagonal(sigma.^2 / 2)) * L2\n",
    "A = (Ldrift + Ldiffusion) * Q"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(The standard printout for the composed operator is a bit messy. Should probably implement `Base.show` for the composition types)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "10-element Array{Float64,1}:\n",
       " 24.4109\n",
       " 37.4853\n",
       " 45.6853\n",
       " 52.2336\n",
       " 51.9707\n",
       " 50.2344\n",
       " 46.5129\n",
       " 37.9023\n",
       " 30.9117\n",
       " 15.7658"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Solve the HJBE (rI - A)u = x\n",
    "r = 0.05\n",
    "LHS = DiffEqArrayOperator(r*speye(N)) - A\n",
    "u = LHS \\ xs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Julia 0.6.2",
   "language": "julia",
   "name": "julia-0.6"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "0.6.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}