-- SPDX-License-Identifier: MPL-2.0

{- | Structured-grid reference PDEs in one, two or three spatial dimensions.
Coordinates use consistent caller-selected length units; diffusivity is L²/T,
velocity L/T and reaction 1/T. Nodes are row-major, last axis fastest.
Nonperiodic axes include both endpoints; periodic axes exclude the upper
endpoint. Boundary derivatives are outward normal derivatives, not fluxes.

FD uses nodal kappa*(-Laplacian); FV uses conservative arithmetic face kappa
and half control volumes at physical faces. Q1 FEM uses two-point tensor
Gauss quadrature, a consistent mass matrix, and natural boundary integrals
(cf. <https://mfem.org/tutorial/fem/>). FEM supports nonperiodic grids.
Spectral paths require periodic axes. Sparse assembly is bounded to 32768
nodes (FEM 2048); dense reference solves/time evolution to 1024 nodes.
-}
module HMath.PDE (
    Grid,
    grid,
    gridShape,
    gridCoordinates,
    gridSpacing,
    flattenIndex,
    unflattenIndex,
    Boundary (..),
    Discretization (..),
    PDEOperator,
    assemblePDE,
    pdeStiffness,
    pdeMass,
    pdeBoundaryLoad,
    PDEResult (..),
    solveElliptic,
    poisson,
    laplace,
    helmholtz,
    spatialGradient,
    laplacianStencil,
    spectralDerivative,
    spectralLaplacian,
    spectralPoisson,
    pseudoSpectralDiffusion,
    methodOfLines,
    heat,
    diffusion,
    wave,
    advection,
    advectionDiffusion,
    stabilityLimit,
    cflNumber,
    manufacturedError,
) where

import Control.Monad (foldM)
import Data.Complex (Complex ((:+)), magnitude, realPart)
import qualified Data.Map.Strict as Map
import HMath.Error (MathError (..))
import qualified HMath.Fourier as F
import qualified HMath.ODE as O
import HMath.Signal.Internal
import qualified HMath.Sparse as S

-- | Opaque Cartesian grid with positive representable spacings and 1..3 axes.
data Grid = Grid [Int] [(Double, Double)] [Bool] [Double] deriving (Grid -> Grid -> Bool
(Grid -> Grid -> Bool) -> (Grid -> Grid -> Bool) -> Eq Grid
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Grid -> Grid -> Bool
== :: Grid -> Grid -> Bool
$c/= :: Grid -> Grid -> Bool
/= :: Grid -> Grid -> Bool
Eq, Int -> Grid -> ShowS
[Grid] -> ShowS
Grid -> String
(Int -> Grid -> ShowS)
-> (Grid -> String) -> ([Grid] -> ShowS) -> Show Grid
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Grid -> ShowS
showsPrec :: Int -> Grid -> ShowS
$cshow :: Grid -> String
show :: Grid -> String
$cshowList :: [Grid] -> ShowS
showList :: [Grid] -> ShowS
Show)

-- | Constant essential value, outward derivative, a*u+b*du/dn=c, or paired periodic faces.
data Boundary = Dirichlet Double | Neumann Double | Robin Double Double Double | Periodic
    deriving (Boundary -> Boundary -> Bool
(Boundary -> Boundary -> Bool)
-> (Boundary -> Boundary -> Bool) -> Eq Boundary
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Boundary -> Boundary -> Bool
== :: Boundary -> Boundary -> Bool
$c/= :: Boundary -> Boundary -> Bool
/= :: Boundary -> Boundary -> Bool
Eq, Int -> Boundary -> ShowS
[Boundary] -> ShowS
Boundary -> String
(Int -> Boundary -> ShowS)
-> (Boundary -> String) -> ([Boundary] -> ShowS) -> Show Boundary
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Boundary -> ShowS
showsPrec :: Int -> Boundary -> ShowS
$cshow :: Boundary -> String
show :: Boundary -> String
$cshowList :: [Boundary] -> ShowS
showList :: [Boundary] -> ShowS
Show)

-- | Nodal nondivergence FD, conservative face-flux FV, or tensor-product Q1 Galerkin FEM.
data Discretization = FiniteDifference | FiniteVolume | FiniteElement deriving (Discretization -> Discretization -> Bool
(Discretization -> Discretization -> Bool)
-> (Discretization -> Discretization -> Bool) -> Eq Discretization
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Discretization -> Discretization -> Bool
== :: Discretization -> Discretization -> Bool
$c/= :: Discretization -> Discretization -> Bool
/= :: Discretization -> Discretization -> Bool
Eq, Int -> Discretization -> ShowS
[Discretization] -> ShowS
Discretization -> String
(Int -> Discretization -> ShowS)
-> (Discretization -> String)
-> ([Discretization] -> ShowS)
-> Show Discretization
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Discretization -> ShowS
showsPrec :: Int -> Discretization -> ShowS
$cshow :: Discretization -> String
show :: Discretization -> String
$cshowList :: [Discretization] -> ShowS
showList :: [Discretization] -> ShowS
Show)

-- | Opaque assembled stiffness, source mass matrix, boundary load and essential constraints.
data PDEOperator = PDEOperator Grid (S.Sparse Double) (S.Sparse Double) [Double] [(Int, Double)]
    deriving (PDEOperator -> PDEOperator -> Bool
(PDEOperator -> PDEOperator -> Bool)
-> (PDEOperator -> PDEOperator -> Bool) -> Eq PDEOperator
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: PDEOperator -> PDEOperator -> Bool
== :: PDEOperator -> PDEOperator -> Bool
$c/= :: PDEOperator -> PDEOperator -> Bool
/= :: PDEOperator -> PDEOperator -> Bool
Eq, Int -> PDEOperator -> ShowS
[PDEOperator] -> ShowS
PDEOperator -> String
(Int -> PDEOperator -> ShowS)
-> (PDEOperator -> String)
-> ([PDEOperator] -> ShowS)
-> Show PDEOperator
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> PDEOperator -> ShowS
showsPrec :: Int -> PDEOperator -> ShowS
$cshow :: PDEOperator -> String
show :: PDEOperator -> String
$cshowList :: [PDEOperator] -> ShowS
showList :: [PDEOperator] -> ShowS
Show)

-- | Nodal solution with independently recomputed infinity residual; errorEstimate reports backward residual only.
data PDEResult = PDEResult
    { PDEResult -> [Double]
pdeValues :: [Double]
    -- ^ Computed grid solution values in the solver's documented spatial order.
    , PDEResult -> Double
pdeResidual :: Double
    -- ^ Residual measure at the returned grid solution.
    , PDEResult -> Double
pdeErrorEstimate :: Double
    -- ^ Solver-specific discretization/error estimate; not a universal certified bound.
    }
    deriving (PDEResult -> PDEResult -> Bool
(PDEResult -> PDEResult -> Bool)
-> (PDEResult -> PDEResult -> Bool) -> Eq PDEResult
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: PDEResult -> PDEResult -> Bool
== :: PDEResult -> PDEResult -> Bool
$c/= :: PDEResult -> PDEResult -> Bool
/= :: PDEResult -> PDEResult -> Bool
Eq, Int -> PDEResult -> ShowS
[PDEResult] -> ShowS
PDEResult -> String
(Int -> PDEResult -> ShowS)
-> (PDEResult -> String)
-> ([PDEResult] -> ShowS)
-> Show PDEResult
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> PDEResult -> ShowS
showsPrec :: Int -> PDEResult -> ShowS
$cshow :: PDEResult -> String
show :: PDEResult -> String
$cshowList :: [PDEResult] -> ShowS
showList :: [PDEResult] -> ShowS
Show)

-- | Construct a grid with >=3 nodes per axis; reject overflowing counts and rounded-away duplicate coordinates.
grid :: [Int] -> [(Double, Double)] -> [Bool] -> Either MathError Grid
grid :: [Int] -> [(Double, Double)] -> [Bool] -> Either MathError Grid
grid [Int]
dims [(Double, Double)]
bounds [Bool]
periodic = do
    Bool -> String -> Either MathError ()
require
        ( [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
dims Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
1
            Bool -> Bool -> Bool
&& [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
dims Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
3
            Bool -> Bool -> Bool
&& [(Double, Double)] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(Double, Double)]
bounds Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
dims
            Bool -> Bool -> Bool
&& [Bool] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Bool]
periodic Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
dims
            Bool -> Bool -> Bool
&& (Int -> Bool) -> [Int] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
3) [Int]
dims
            Bool -> Bool -> Bool
&& [Integer] -> Integer
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product ((Int -> Integer) -> [Int] -> [Integer]
forall a b. (a -> b) -> [a] -> [b]
map Int -> Integer
forall a. Integral a => a -> Integer
toInteger [Int]
dims) Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<= Integer
32768
        )
        String
"invalid PDE grid dimensions (1..3 axes, >=3 nodes/axis, <=32768 total)"
    let spacing :: [Double]
spacing = [(Double
b Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
a) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (if Bool
p then Int
n else Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) | (Int
n, (Double
a, Double
b), Bool
p) <- [Int]
-> [(Double, Double)] -> [Bool] -> [(Int, (Double, Double), Bool)]
forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3 [Int]
dims [(Double, Double)]
bounds [Bool]
periodic]
    Bool -> String -> Either MathError ()
require
        (((Double, Double) -> Bool) -> [(Double, Double)] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (\(Double
a, Double
b) -> Double -> Bool
finite Double
a Bool -> Bool -> Bool
&& Double -> Bool
finite Double
b Bool -> Bool -> Bool
&& Double
a Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
b) [(Double, Double)]
bounds Bool -> Bool -> Bool
&& (Double -> Bool) -> [Double] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (\Double
h -> Double -> Bool
finite Double
h Bool -> Bool -> Bool
&& Double
h Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
0) [Double]
spacing)
        String
"invalid PDE bounds/spacing"
    let axes :: [[Double]]
axes = [[Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
i Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
h | Int
i <- [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]] | (Int
n, (Double
a, Double
_), Double
h) <- [Int]
-> [(Double, Double)]
-> [Double]
-> [(Int, (Double, Double), Double)]
forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3 [Int]
dims [(Double, Double)]
bounds [Double]
spacing]
    Bool -> String -> Either MathError ()
require
        (([Double] -> Bool) -> [[Double]] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (\[Double]
xs -> (Double -> Bool) -> [Double] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Double -> Bool
finite [Double]
xs Bool -> Bool -> Bool
&& [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
and ((Double -> Double -> Bool) -> [Double] -> [Double] -> [Bool]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
(<) [Double]
xs (Int -> [Double] -> [Double]
forall a. Int -> [a] -> [a]
drop Int
1 [Double]
xs))) [[Double]]
axes)
        String
"rounded-away/nonfinite PDE coordinates"
    Grid -> Either MathError Grid
forall a b. b -> Either a b
Right ([Int] -> [(Double, Double)] -> [Bool] -> [Double] -> Grid
Grid [Int]
dims [(Double, Double)]
bounds [Bool]
periodic [Double]
spacing)

-- | Number of nodes along each axis, preserving last-axis-fastest storage order.
gridShape :: Grid -> [Int]
gridShape :: Grid -> [Int]
gridShape (Grid [Int]
dims [(Double, Double)]
_ [Bool]
_ [Double]
_) = [Int]
dims

-- | Positive uniform distance between adjacent nodes on each axis.
gridSpacing :: Grid -> [Double]
gridSpacing :: Grid -> [Double]
gridSpacing (Grid [Int]
_ [(Double, Double)]
_ [Bool]
_ [Double]
hs) = [Double]
hs

strides :: [Int] -> [Int]
strides :: [Int] -> [Int]
strides [Int]
dims = Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
drop Int
1 ((Int -> Int -> Int) -> Int -> [Int] -> [Int]
forall a b. (a -> b -> b) -> b -> [a] -> [b]
scanr Int -> Int -> Int
forall a. Num a => a -> a -> a
(*) Int
1 [Int]
dims)

-- | Checked conversion of zero-based multidimensional node coordinates to a row-major flat index.
flattenIndex :: Grid -> [Int] -> Either MathError Int
flattenIndex :: Grid -> [Int] -> Either MathError Int
flattenIndex (Grid [Int]
dims [(Double, Double)]
_ [Bool]
_ [Double]
_) [Int]
indices = do
    Bool -> String -> Either MathError ()
require
        ([Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
indices Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
dims Bool -> Bool -> Bool
&& [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
and ((Int -> Int -> Bool) -> [Int] -> [Int] -> [Bool]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\Int
i Int
n -> Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n) [Int]
indices [Int]
dims))
        String
"PDE index outside grid"
    Int -> Either MathError Int
forall a b. b -> Either a b
Right ([Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((Int -> Int -> Int) -> [Int] -> [Int] -> [Int]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Int -> Int -> Int
forall a. Num a => a -> a -> a
(*) [Int]
indices ([Int] -> [Int]
strides [Int]
dims)))

-- | Inverse checked row-major index conversion.
unflattenIndex :: Grid -> Int -> Either MathError [Int]
unflattenIndex :: Grid -> Int -> Either MathError [Int]
unflattenIndex (Grid [Int]
dims [(Double, Double)]
_ [Bool]
_ [Double]
_) Int
i = do
    Bool -> String -> Either MathError ()
require (Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product [Int]
dims) String
"PDE flat index outside grid"
    [Int] -> Either MathError [Int]
forall a b. b -> Either a b
Right ((Int -> Int -> Int) -> [Int] -> [Int] -> [Int]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\Int
s Int
n -> Int
i Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
s Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
n) ([Int] -> [Int]
strides [Int]
dims) [Int]
dims)

-- | Physical coordinates in flat storage order; periodic upper endpoints are omitted.
gridCoordinates :: Grid -> [[Double]]
gridCoordinates :: Grid -> [[Double]]
gridCoordinates (Grid [Int]
dims [(Double, Double)]
bounds [Bool]
_ [Double]
hs) =
    [ [Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
i Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
s Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
n) Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
h | ((Double
a, Double
_), Double
h, Int
s, Int
n) <- [(Double, Double)]
-> [Double]
-> [Int]
-> [Int]
-> [((Double, Double), Double, Int, Int)]
forall {a} {b} {c} {d}. [a] -> [b] -> [c] -> [d] -> [(a, b, c, d)]
zip4 [(Double, Double)]
bounds [Double]
hs ([Int] -> [Int]
strides [Int]
dims) [Int]
dims]
    | Int
i <- [Int
0 .. [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product [Int]
dims Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
    ]
  where
    zip4 :: [a] -> [b] -> [c] -> [d] -> [(a, b, c, d)]
zip4 [a]
as [b]
bs [c]
cs [d]
ds = ((a, b) -> (c, d) -> (a, b, c, d))
-> [(a, b)] -> [(c, d)] -> [(a, b, c, d)]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\(a
a, b
b) (c
c, d
d) -> (a
a, b
b, c
c, d
d)) ([a] -> [b] -> [(a, b)]
forall a b. [a] -> [b] -> [(a, b)]
zip [a]
as [b]
bs) ([c] -> [d] -> [(c, d)]
forall a b. [a] -> [b] -> [(a, b)]
zip [c]
cs [d]
ds)

validateBoundary :: Grid -> [(Boundary, Boundary)] -> Either MathError ()
validateBoundary :: Grid -> [(Boundary, Boundary)] -> Either MathError ()
validateBoundary (Grid [Int]
dims [(Double, Double)]
_ [Bool]
periodic [Double]
_) [(Boundary, Boundary)]
bc = do
    Bool -> String -> Either MathError ()
require ([(Boundary, Boundary)] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(Boundary, Boundary)]
bc Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
dims) String
"boundary count must equal grid rank"
    ((Bool, (Boundary, Boundary)) -> Either MathError ())
-> [(Bool, (Boundary, Boundary))] -> Either MathError ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_
        ( \(Bool
p, (Boundary
a, Boundary
b)) ->
            if Bool
p
                then Bool -> String -> Either MathError ()
require (Boundary
a Boundary -> Boundary -> Bool
forall a. Eq a => a -> a -> Bool
== Boundary
Periodic Bool -> Bool -> Bool
&& Boundary
b Boundary -> Boundary -> Bool
forall a. Eq a => a -> a -> Bool
== Boundary
Periodic) String
"periodic axis needs both periodic faces"
                else Bool -> String -> Either MathError ()
require (Boundary
a Boundary -> Boundary -> Bool
forall a. Eq a => a -> a -> Bool
/= Boundary
Periodic Bool -> Bool -> Bool
&& Boundary
b Boundary -> Boundary -> Bool
forall a. Eq a => a -> a -> Bool
/= Boundary
Periodic) String
"nonperiodic axis cannot have periodic face"
        )
        ([Bool] -> [(Boundary, Boundary)] -> [(Bool, (Boundary, Boundary))]
forall a b. [a] -> [b] -> [(a, b)]
zip [Bool]
periodic [(Boundary, Boundary)]
bc)
    (Boundary -> Either MathError ())
-> [Boundary] -> Either MathError ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ Boundary -> Either MathError ()
valid (((Boundary, Boundary) -> [Boundary])
-> [(Boundary, Boundary)] -> [Boundary]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (\(Boundary
a, Boundary
b) -> [Boundary
a, Boundary
b]) [(Boundary, Boundary)]
bc)
  where
    valid :: Boundary -> Either MathError ()
valid Boundary
Periodic = () -> Either MathError ()
forall a b. b -> Either a b
Right ()
    valid (Dirichlet Double
x) = Double -> Either MathError Double
checked Double
x Either MathError Double
-> Either MathError () -> Either MathError ()
forall a b.
Either MathError a -> Either MathError b -> Either MathError b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> () -> Either MathError ()
forall a b. b -> Either a b
Right ()
    valid (Neumann Double
x) = Double -> Either MathError Double
checked Double
x Either MathError Double
-> Either MathError () -> Either MathError ()
forall a b.
Either MathError a -> Either MathError b -> Either MathError b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> () -> Either MathError ()
forall a b. b -> Either a b
Right ()
    valid (Robin Double
a Double
b Double
c) = [Double] -> Either MathError [Double]
vector [Double
a, Double
b, Double
c] Either MathError [Double]
-> Either MathError () -> Either MathError ()
forall a b.
Either MathError a -> Either MathError b -> Either MathError b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Bool -> String -> Either MathError ()
require (Double
a Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
/= Double
0 Bool -> Bool -> Bool
|| Double
b Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
/= Double
0) String
"Robin coefficients both zero"

-- | Resolve essential values on corners; contradictory Dirichlet faces fail.
fixedNodes :: Grid -> [(Boundary, Boundary)] -> Either MathError [(Int, Double)]
fixedNodes :: Grid -> [(Boundary, Boundary)] -> Either MathError [(Int, Double)]
fixedNodes g :: Grid
g@(Grid [Int]
dims [(Double, Double)]
_ [Bool]
periodic [Double]
_) [(Boundary, Boundary)]
bc = ([[(Int, Double)]] -> [(Int, Double)])
-> Either MathError [[(Int, Double)]]
-> Either MathError [(Int, Double)]
forall a b. (a -> b) -> Either MathError a -> Either MathError b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [[(Int, Double)]] -> [(Int, Double)]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ((Int -> Either MathError [(Int, Double)])
-> [Int] -> Either MathError [[(Int, Double)]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM Int -> Either MathError [(Int, Double)]
fixed [Int
0 .. [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product [Int]
dims Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1])
  where
    fixed :: Int -> Either MathError [(Int, Double)]
fixed Int
i = do
        ix <- Grid -> Int -> Either MathError [Int]
unflattenIndex Grid
g Int
i
        let faces =
                [ if Int
side Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 then (Boundary, Boundary) -> Boundary
forall a b. (a, b) -> a
fst (Boundary, Boundary)
pair else (Boundary, Boundary) -> Boundary
forall a b. (a, b) -> b
snd (Boundary, Boundary)
pair
                | (Int
k, (Int
n, Bool
p, (Boundary, Boundary)
pair)) <- [Int]
-> [(Int, Bool, (Boundary, Boundary))]
-> [(Int, (Int, Bool, (Boundary, Boundary)))]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 ..] ([Int]
-> [Bool]
-> [(Boundary, Boundary)]
-> [(Int, Bool, (Boundary, Boundary))]
forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3 [Int]
dims [Bool]
periodic [(Boundary, Boundary)]
bc)
                , Bool -> Bool
not Bool
p
                , Int
side <- [Int
0, Int
1 :: Int]
                , [Int]
ix [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
k Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== if Int
side Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 then Int
0 else Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
                ]
            value (Dirichlet Double
x) = Double -> Maybe Double
forall a. a -> Maybe a
Just Double
x; value (Robin Double
a Double
0 Double
c) = Double -> Maybe Double
forall a. a -> Maybe a
Just (Double
c Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
a); value Boundary
_ = Maybe Double
forall a. Maybe a
Nothing
            values = [Double
v | Just Double
v <- (Boundary -> Maybe Double) -> [Boundary] -> [Maybe Double]
forall a b. (a -> b) -> [a] -> [b]
map Boundary -> Maybe Double
value [Boundary]
faces]
        _ <- vector values
        case values of
            [] -> [(Int, Double)] -> Either MathError [(Int, Double)]
forall a b. b -> Either a b
Right []
            Double
v : [Double]
_ ->
                Bool -> String -> Either MathError ()
require
                    ((Double -> Bool) -> [Double] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (\Double
w -> Double -> Double
forall a. Num a => a -> a
abs (Double
w Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
v) Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
1e-12 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double -> Double
forall a. Ord a => a -> a -> a
max Double
1 (Double -> Double
forall a. Num a => a -> a
abs Double
v)) [Double]
values)
                    String
"inconsistent essential boundary values at grid corner"
                    Either MathError ()
-> Either MathError [(Int, Double)]
-> Either MathError [(Int, Double)]
forall a b.
Either MathError a -> Either MathError b -> Either MathError b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> [(Int, Double)] -> Either MathError [(Int, Double)]
forall a b. b -> Either a b
Right [(Int
i, Double
v)]

-- | Assemble -kappa*Laplacian (FD) or -div(kappa grad) (FV/FEM) plus constant reaction; coefficient field is positive at every node.
assemblePDE ::
    Discretization ->
    Grid ->
    [(Boundary, Boundary)] ->
    [Double] ->
    Double ->
    Either MathError PDEOperator
assemblePDE :: Discretization
-> Grid
-> [(Boundary, Boundary)]
-> [Double]
-> Double
-> Either MathError PDEOperator
assemblePDE Discretization
method g :: Grid
g@(Grid [Int]
dims [(Double, Double)]
_ [Bool]
periodic [Double]
_) [(Boundary, Boundary)]
bc [Double]
diffusivity Double
reaction = do
    Grid -> [(Boundary, Boundary)] -> Either MathError ()
validateBoundary Grid
g [(Boundary, Boundary)]
bc
    let n :: Int
n = [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product [Int]
dims
    Bool -> String -> Either MathError ()
require
        ([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
diffusivity Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
n Bool -> Bool -> Bool
&& (Double -> Bool) -> [Double] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
0) [Double]
diffusivity Bool -> Bool -> Bool
&& Double -> Bool
finite Double
reaction)
        String
"invalid positive diffusivity field/reaction"
    _ <- [Double] -> Either MathError [Double]
vector [Double]
diffusivity
    fixed <- fixedNodes g bc
    (entries, massEntries, load) <-
        if method == FiniteElement
            then do
                require (n <= 2048 && not (or periodic)) "Q1 FEM reference supports <=2048 nonperiodic nodes"
                femAssembly g bc diffusivity reaction
            else do
                rows <- mapM (stencilRow method g bc diffusivity reaction) [0 .. n - 1]
                Right (concatMap fst rows, [(i, i, 1) | i <- [0 .. n - 1]], map snd rows)
    let essential = ((Int, Double) -> Int) -> [(Int, Double)] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Int, Double) -> Int
forall a b. (a, b) -> a
fst [(Int, Double)]
fixed
        filtered = [(Int
i, Int
j, Double
v) | (Int
i, Int
j, Double
v) <- [(Int, Int, Double)]
entries, Int
i Int -> [Int] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` [Int]
essential] [(Int, Int, Double)]
-> [(Int, Int, Double)] -> [(Int, Int, Double)]
forall a. [a] -> [a] -> [a]
++ [(Int
i, Int
i, Double
1) | Int
i <- [Int]
essential]
        massFiltered = [(Int
i, Int
j, Double
v) | (Int
i, Int
j, Double
v) <- [(Int, Int, Double)]
massEntries, Int
i Int -> [Int] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` [Int]
essential]
        finalLoad = [Double -> (Double -> Double) -> Maybe Double -> Double
forall b a. b -> (a -> b) -> Maybe a -> b
maybe ([Double]
load [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
i) Double -> Double
forall a. a -> a
id (Int -> [(Int, Double)] -> Maybe Double
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup Int
i [(Int, Double)]
fixed) | Int
i <- [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]
    _ <- vector ([v | (_, _, v) <- filtered] ++ [v | (_, _, v) <- massFiltered] ++ finalLoad)
    stiffness <- S.sparse n n filtered
    mass <- S.sparse n n massFiltered
    Right (PDEOperator g stiffness mass finalLoad fixed)

stencilRow ::
    Discretization ->
    Grid ->
    [(Boundary, Boundary)] ->
    [Double] ->
    Double ->
    Int ->
    Either MathError ([(Int, Int, Double)], Double)
stencilRow :: Discretization
-> Grid
-> [(Boundary, Boundary)]
-> [Double]
-> Double
-> Int
-> Either MathError ([(Int, Int, Double)], Double)
stencilRow Discretization
method g :: Grid
g@(Grid [Int]
dims [(Double, Double)]
_ [Bool]
periodic [Double]
hs) [(Boundary, Boundary)]
bc [Double]
kappas Double
reaction Int
i = do
    ix <- Grid -> Int -> Either MathError [Int]
unflattenIndex Grid
g Int
i
    let parts Int
axis =
            let h :: Double
h = [Double]
hs [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
axis
                k :: Double
k = [Double]
kappas [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
i
                n :: Int
n = [Int]
dims [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
axis
                q :: Int
q = [Int]
ix [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
axis
                period :: Bool
period = [Bool]
periodic [Bool] -> Int -> Bool
forall a. HasCallStack => [a] -> Int -> a
!! Int
axis
                neighbor :: Int -> (Int, Double)
neighbor Int
side =
                    let qq :: Int
qq = (Int
q Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
side) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
n; j :: Int
j = [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((Int -> Int -> Int) -> [Int] -> [Int] -> [Int]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Int -> Int -> Int
forall a. Num a => a -> a -> a
(*) (Int -> Int -> [Int] -> [Int]
forall a. Int -> a -> [a] -> [a]
replace Int
axis Int
qq [Int]
ix) ([Int] -> [Int]
strides [Int]
dims))
                     in (Int
j, if Discretization
method Discretization -> Discretization -> Bool
forall a. Eq a => a -> a -> Bool
== Discretization
FiniteVolume then (Double
k Double -> Double -> Double
forall a. Num a => a -> a -> a
+ [Double]
kappas [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
j) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
2 else Double
k)
                inner :: Int -> ([(Int, Int, Double)], b)
inner Int
side = let (Int
j, Double
faceK) = Int -> (Int, Double)
neighbor Int
side in ([(Int
i, Int
i, Double
faceK Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
h Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
h)), (Int
i, Int
j, -Double
faceK Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
h Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
h))], b
0)
                face :: Int -> Boundary -> ([(Int, Int, Double)], Double)
face Int
side Boundary
boundary =
                    let (Int
j, Double
kface) = Int -> (Int, Double)
neighbor (-Int
side)
                        scale :: Double
scale = Double
2 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
kface Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
h Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
h)
                        (Double
a, Double
b, Double
c) = case Boundary
boundary of Neumann Double
d -> (Double
0, Double
1, Double
d); Robin Double
aa Double
bb Double
cc -> (Double
aa, Double
bb, Double
cc); Boundary
_ -> (Double
0, Double
1, Double
0)
                     in if Double
b Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
0
                            then ([], Double
0)
                            else ([(Int
i, Int
i, Double
scale Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
2 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
k Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
a Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
h Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
b)), (Int
i, Int
j, -Double
scale)], Double
2 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
k Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
c Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
h Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
b))
             in if Bool
period Bool -> Bool -> Bool
|| Int
q Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 Bool -> Bool -> Bool
&& Int
q Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
                    then ([(Int, Int, Double)], Double)
-> ([(Int, Int, Double)], Double) -> ([(Int, Int, Double)], Double)
forall {b} {a}. Num b => ([a], b) -> ([a], b) -> ([a], b)
combine (Int -> ([(Int, Int, Double)], Double)
forall {b}. Num b => Int -> ([(Int, Int, Double)], b)
inner (-Int
1)) (Int -> ([(Int, Int, Double)], Double)
forall {b}. Num b => Int -> ([(Int, Int, Double)], b)
inner Int
1)
                    else if Int
q Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 then Int -> Boundary -> ([(Int, Int, Double)], Double)
face (-Int
1) ((Boundary, Boundary) -> Boundary
forall a b. (a, b) -> a
fst ([(Boundary, Boundary)]
bc [(Boundary, Boundary)] -> Int -> (Boundary, Boundary)
forall a. HasCallStack => [a] -> Int -> a
!! Int
axis)) else Int -> Boundary -> ([(Int, Int, Double)], Double)
face Int
1 ((Boundary, Boundary) -> Boundary
forall a b. (a, b) -> b
snd ([(Boundary, Boundary)]
bc [(Boundary, Boundary)] -> Int -> (Boundary, Boundary)
forall a. HasCallStack => [a] -> Int -> a
!! Int
axis))
        combine ([a]
a, b
x) ([a]
b, b
y) = ([a]
a [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ [a]
b, b
x b -> b -> b
forall a. Num a => a -> a -> a
+ b
y)
        (entries, load) = foldl combine ([(i, i, reaction)], 0) [parts a | a <- [0 .. length dims - 1]]
    Right (entries, load)

-- Tensor Q1 element and boundary quadrature. Integration uses cell-specific
-- interpolated positive diffusivity and consistent source/mass interpolation.
femAssembly ::
    Grid ->
    [(Boundary, Boundary)] ->
    [Double] ->
    Double ->
    Either MathError ([(Int, Int, Double)], [(Int, Int, Double)], [Double])
femAssembly :: Grid
-> [(Boundary, Boundary)]
-> [Double]
-> Double
-> Either
     MathError ([(Int, Int, Double)], [(Int, Int, Double)], [Double])
femAssembly (Grid [Int]
dims [(Double, Double)]
_ [Bool]
_ [Double]
hs) [(Boundary, Boundary)]
bc [Double]
kappas Double
reaction = do
    let d :: Int
d = [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
dims
        n :: Int
n = [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product [Int]
dims
        bits :: [[Int]]
bits = [[Int]] -> [[Int]]
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
forall (m :: * -> *) a. Monad m => [m a] -> m [a]
sequence (Int -> [Int] -> [[Int]]
forall a. Int -> a -> [a]
replicate Int
d [Int
0, Int
1 :: Int])
        cells :: [[Int]]
cells = [[Int]] -> [[Int]]
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
forall (m :: * -> *) a. Monad m => [m a] -> m [a]
sequence [[Int
0 .. Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2] | Int
k <- [Int]
dims]
        gauss :: [Double]
gauss = [(Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double -> Double
forall a. Floating a => a -> a
sqrt Double
3) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
2, (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double -> Double
forall a. Floating a => a -> a
sqrt Double
3) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
2]
        phi :: [a] -> [a] -> a
phi [a]
corner [a]
z = [a] -> a
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product ((a -> a -> a) -> [a] -> [a] -> [a]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\a
v a
t -> if a
v a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
0 then a
1 a -> a -> a
forall a. Num a => a -> a -> a
- a
t else a
t) [a]
corner [a]
z)
        grad :: [a] -> [Double] -> Int -> Double
grad [a]
corner [Double]
z Int
axis =
            (if [a]
corner [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! Int
axis a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
0 then -Double
1 else Double
1)
                Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ ([Double]
hs [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
axis)
                Double -> Double -> Double
forall a. Num a => a -> a -> a
* [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product [if [a]
corner [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! Int
k a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
0 then Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- [Double]
z [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
k else [Double]
z [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
k | Int
k <- [Int
0 .. Int
d Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1], Int
k Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
axis]
        cellData :: [Int] -> ([(Int, Int, Double)], [(Int, Int, Double)], [Int])
cellData [Int]
base =
            let nodes :: [Int]
nodes = [[Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((Int -> Int -> Int) -> [Int] -> [Int] -> [Int]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Int -> Int -> Int
forall a. Num a => a -> a -> a
(*) ((Int -> Int -> Int) -> [Int] -> [Int] -> [Int]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) [Int]
base [Int]
corner) ([Int] -> [Int]
strides [Int]
dims)) | [Int]
corner <- [[Int]]
bits]
                quadrature :: [[Double]]
quadrature = [[Double]] -> [[Double]]
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
forall (m :: * -> *) a. Monad m => [m a] -> m [a]
sequence (Int -> [Double] -> [[Double]]
forall a. Int -> a -> [a]
replicate Int
d [Double]
gauss)
                weight :: Double
weight = [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product [Double]
hs Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
2 Int -> Int -> Int
forall a b. (Num a, Integral b) => a -> b -> a
^ Int
d :: Int)
                kval :: [Double] -> Double
kval [Double]
z = [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [[Double]
kappas [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
j Double -> Double -> Double
forall a. Num a => a -> a -> a
* [Int] -> [Double] -> Double
forall {a} {a}. (Eq a, Num a, Num a) => [a] -> [a] -> a
phi [Int]
corner [Double]
z | (Int
j, [Int]
corner) <- [Int] -> [[Int]] -> [(Int, [Int])]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int]
nodes [[Int]]
bits]
                localMass :: [a] -> [a] -> Double
localMass [a]
a [a]
b = Double
weight Double -> Double -> Double
forall a. Num a => a -> a -> a
* [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [[a] -> [Double] -> Double
forall {a} {a}. (Eq a, Num a, Num a) => [a] -> [a] -> a
phi [a]
a [Double]
z Double -> Double -> Double
forall a. Num a => a -> a -> a
* [a] -> [Double] -> Double
forall {a} {a}. (Eq a, Num a, Num a) => [a] -> [a] -> a
phi [a]
b [Double]
z | [Double]
z <- [[Double]]
quadrature]
                stiff :: [a] -> [a] -> Double
stiff [a]
a [a]
b =
                    Double
weight
                        Double -> Double -> Double
forall a. Num a => a -> a -> a
* [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum
                            [ [Double] -> Double
kval [Double]
z Double -> Double -> Double
forall a. Num a => a -> a -> a
* [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [[a] -> [Double] -> Int -> Double
forall {a}. (Eq a, Num a) => [a] -> [Double] -> Int -> Double
grad [a]
a [Double]
z Int
k Double -> Double -> Double
forall a. Num a => a -> a -> a
* [a] -> [Double] -> Int -> Double
forall {a}. (Eq a, Num a) => [a] -> [Double] -> Int -> Double
grad [a]
b [Double]
z Int
k | Int
k <- [Int
0 .. Int
d Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]] Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
reaction Double -> Double -> Double
forall a. Num a => a -> a -> a
* [a] -> [Double] -> Double
forall {a} {a}. (Eq a, Num a, Num a) => [a] -> [a] -> a
phi [a]
a [Double]
z Double -> Double -> Double
forall a. Num a => a -> a -> a
* [a] -> [Double] -> Double
forall {a} {a}. (Eq a, Num a, Num a) => [a] -> [a] -> a
phi [a]
b [Double]
z
                            | [Double]
z <- [[Double]]
quadrature
                            ]
             in ( [(Int
i, Int
j, [Int] -> [Int] -> Double
forall {a} {a}. (Eq a, Eq a, Num a, Num a) => [a] -> [a] -> Double
stiff [Int]
a [Int]
b) | (Int
i, [Int]
a) <- [Int] -> [[Int]] -> [(Int, [Int])]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int]
nodes [[Int]]
bits, (Int
j, [Int]
b) <- [Int] -> [[Int]] -> [(Int, [Int])]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int]
nodes [[Int]]
bits]
                , [(Int
i, Int
j, [Int] -> [Int] -> Double
forall {a} {a}. (Eq a, Eq a, Num a, Num a) => [a] -> [a] -> Double
localMass [Int]
a [Int]
b) | (Int
i, [Int]
a) <- [Int] -> [[Int]] -> [(Int, [Int])]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int]
nodes [[Int]]
bits, (Int
j, [Int]
b) <- [Int] -> [[Int]] -> [(Int, [Int])]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int]
nodes [[Int]]
bits]
                , [Int]
nodes
                )
        volume :: [([(Int, Int, Double)], [(Int, Int, Double)], [Int])]
volume = ([Int] -> ([(Int, Int, Double)], [(Int, Int, Double)], [Int]))
-> [[Int]] -> [([(Int, Int, Double)], [(Int, Int, Double)], [Int])]
forall a b. (a -> b) -> [a] -> [b]
map [Int] -> ([(Int, Int, Double)], [(Int, Int, Double)], [Int])
cellData [[Int]]
cells
        faceData :: [Int] -> Int -> p -> ([(Int, Int, Double)], [(Int, Double)])
faceData [Int]
base Int
axis p
side =
            let ([(Int, Int, Double)]
_, [(Int, Int, Double)]
_, [Int]
nodes) = [Int] -> ([(Int, Int, Double)], [(Int, Int, Double)], [Int])
cellData [Int]
base
                boundary :: Boundary
boundary = if p
side p -> p -> Bool
forall a. Eq a => a -> a -> Bool
== p
0 then (Boundary, Boundary) -> Boundary
forall a b. (a, b) -> a
fst ([(Boundary, Boundary)]
bc [(Boundary, Boundary)] -> Int -> (Boundary, Boundary)
forall a. HasCallStack => [a] -> Int -> a
!! Int
axis) else (Boundary, Boundary) -> Boundary
forall a b. (a, b) -> b
snd ([(Boundary, Boundary)]
bc [(Boundary, Boundary)] -> Int -> (Boundary, Boundary)
forall a. HasCallStack => [a] -> Int -> a
!! Int
axis)
                (Double
a, Double
b, Double
c) = case Boundary
boundary of Neumann Double
v -> (Double
0, Double
1, Double
v); Robin Double
aa Double
bb Double
cc -> (Double
aa, Double
bb, Double
cc); Boundary
_ -> (Double
0, Double
0, Double
0)
                zs :: [[Double]]
zs =
                    ([Double] -> [Double]) -> [[Double]] -> [[Double]]
forall a b. (a -> b) -> [a] -> [b]
map (\[Double]
z -> Int -> [Double] -> [Double]
forall a. Int -> [a] -> [a]
take Int
axis [Double]
z [Double] -> [Double] -> [Double]
forall a. [a] -> [a] -> [a]
++ [p -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral p
side] [Double] -> [Double] -> [Double]
forall a. [a] -> [a] -> [a]
++ Int -> [Double] -> [Double]
forall a. Int -> [a] -> [a]
drop Int
axis [Double]
z) ([[Double]] -> [[Double]]
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
forall (m :: * -> *) a. Monad m => [m a] -> m [a]
sequence (Int -> [Double] -> [[Double]]
forall a. Int -> a -> [a]
replicate (Int
d Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) [Double]
gauss))
                weight :: Double
weight = [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product [[Double]
hs [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
k | Int
k <- [Int
0 .. Int
d Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1], Int
k Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
axis] Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
2 Int -> Int -> Int
forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
d Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) :: Int)
                kval :: [Double] -> Double
kval [Double]
z = [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [[Double]
kappas [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
j Double -> Double -> Double
forall a. Num a => a -> a -> a
* [Int] -> [Double] -> Double
forall {a} {a}. (Eq a, Num a, Num a) => [a] -> [a] -> a
phi [Int]
corner [Double]
z | (Int
j, [Int]
corner) <- [Int] -> [[Int]] -> [(Int, [Int])]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int]
nodes [[Int]]
bits]
             in if Double
b Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
0
                    then ([], [])
                    else
                        ( [ (Int
i, Int
j, Double
weight Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
a Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
b Double -> Double -> Double
forall a. Num a => a -> a -> a
* [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [[Double] -> Double
kval [Double]
z Double -> Double -> Double
forall a. Num a => a -> a -> a
* [Int] -> [Double] -> Double
forall {a} {a}. (Eq a, Num a, Num a) => [a] -> [a] -> a
phi [Int]
ca [Double]
z Double -> Double -> Double
forall a. Num a => a -> a -> a
* [Int] -> [Double] -> Double
forall {a} {a}. (Eq a, Num a, Num a) => [a] -> [a] -> a
phi [Int]
cb [Double]
z | [Double]
z <- [[Double]]
zs])
                          | (Int
i, [Int]
ca) <- [Int] -> [[Int]] -> [(Int, [Int])]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int]
nodes [[Int]]
bits
                          , (Int
j, [Int]
cb) <- [Int] -> [[Int]] -> [(Int, [Int])]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int]
nodes [[Int]]
bits
                          ]
                        , [(Int
i, Double
weight Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
c Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
b Double -> Double -> Double
forall a. Num a => a -> a -> a
* [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [[Double] -> Double
kval [Double]
z Double -> Double -> Double
forall a. Num a => a -> a -> a
* [Int] -> [Double] -> Double
forall {a} {a}. (Eq a, Num a, Num a) => [a] -> [a] -> a
phi [Int]
ca [Double]
z | [Double]
z <- [[Double]]
zs]) | (Int
i, [Int]
ca) <- [Int] -> [[Int]] -> [(Int, [Int])]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int]
nodes [[Int]]
bits]
                        )
        boundaryParts :: [([(Int, Int, Double)], [(Int, Double)])]
boundaryParts =
            [ [Int] -> Int -> Int -> ([(Int, Int, Double)], [(Int, Double)])
forall {p}.
Integral p =>
[Int] -> Int -> p -> ([(Int, Int, Double)], [(Int, Double)])
faceData [Int]
base Int
axis Int
side
            | [Int]
base <- [[Int]]
cells
            , Int
axis <- [Int
0 .. Int
d Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
            , Int
side <- [Int
0, Int
1 :: Int]
            , [Int]
base [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
axis Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== if Int
side Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 then Int
0 else [Int]
dims [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
axis Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2
            ]
        entries :: [(Int, Int, Double)]
entries = [[(Int, Int, Double)]] -> [(Int, Int, Double)]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[(Int, Int, Double)]
a | ([(Int, Int, Double)]
a, [(Int, Int, Double)]
_, [Int]
_) <- [([(Int, Int, Double)], [(Int, Int, Double)], [Int])]
volume] [(Int, Int, Double)]
-> [(Int, Int, Double)] -> [(Int, Int, Double)]
forall a. [a] -> [a] -> [a]
++ (([(Int, Int, Double)], [(Int, Double)]) -> [(Int, Int, Double)])
-> [([(Int, Int, Double)], [(Int, Double)])]
-> [(Int, Int, Double)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap ([(Int, Int, Double)], [(Int, Double)]) -> [(Int, Int, Double)]
forall a b. (a, b) -> a
fst [([(Int, Int, Double)], [(Int, Double)])]
boundaryParts
        mass :: [(Int, Int, Double)]
mass = [[(Int, Int, Double)]] -> [(Int, Int, Double)]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[(Int, Int, Double)]
a | ([(Int, Int, Double)]
_, [(Int, Int, Double)]
a, [Int]
_) <- [([(Int, Int, Double)], [(Int, Int, Double)], [Int])]
volume]
        load :: [Double]
load = [[Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Double
v | (Int
j, Double
v) <- (([(Int, Int, Double)], [(Int, Double)]) -> [(Int, Double)])
-> [([(Int, Int, Double)], [(Int, Double)])] -> [(Int, Double)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap ([(Int, Int, Double)], [(Int, Double)]) -> [(Int, Double)]
forall a b. (a, b) -> b
snd [([(Int, Int, Double)], [(Int, Double)])]
boundaryParts, Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
j] | Int
i <- [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]
    _ <- [Double] -> Either MathError [Double]
vector ([Double
v | (Int
_, Int
_, Double
v) <- [(Int, Int, Double)]
entries] [Double] -> [Double] -> [Double]
forall a. [a] -> [a] -> [a]
++ [Double
v | (Int
_, Int
_, Double
v) <- [(Int, Int, Double)]
mass] [Double] -> [Double] -> [Double]
forall a. [a] -> [a] -> [a]
++ [Double]
load)
    Right (entries, mass, load)

-- | Read the canonical sparse stiffness and source mass operators; essential rows are identity/zero respectively.
pdeStiffness, pdeMass :: PDEOperator -> S.Sparse Double
pdeStiffness :: PDEOperator -> Sparse Double
pdeStiffness (PDEOperator Grid
_ Sparse Double
a Sparse Double
_ [Double]
_ [(Int, Double)]
_) = Sparse Double
a
pdeMass :: PDEOperator -> Sparse Double
pdeMass (PDEOperator Grid
_ Sparse Double
_ Sparse Double
m [Double]
_ [(Int, Double)]
_) = Sparse Double
m

-- | Read boundary contributions, including prescribed values at essential rows.
pdeBoundaryLoad :: PDEOperator -> [Double]
pdeBoundaryLoad :: PDEOperator -> [Double]
pdeBoundaryLoad (PDEOperator Grid
_ Sparse Double
_ Sparse Double
_ [Double]
b [(Int, Double)]
_) = [Double]
b

{- | Dense pivoted reference solve and independently recomputed infinity
residual. The error estimate is backward residual, not continuum error.
-}
solveElliptic :: PDEOperator -> [Double] -> Either MathError PDEResult
solveElliptic :: PDEOperator -> [Double] -> Either MathError PDEResult
solveElliptic (PDEOperator Grid
g Sparse Double
a Sparse Double
mass [Double]
boundary [(Int, Double)]
fixed) [Double]
forcing = do
    let n :: Int
n = [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product (Grid -> [Int]
gridShape Grid
g)
    Bool -> String -> Either MathError ()
require
        (Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
1024 Bool -> Bool -> Bool
&& [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
forcing Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
n)
        String
"elliptic reference solve requires <=1024 nodes and matching source"
    _ <- [Double] -> Either MathError [Double]
vector [Double]
forcing
    source <- S.spmv mass forcing
    let rhs = (Double -> Double -> Double) -> [Double] -> [Double] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Double -> Double -> Double
forall a. Num a => a -> a -> a
(+) [Double]
source [Double]
boundary
    -- A pure Neumann/periodic nullspace is explicitly rejected before floating
    -- elimination can manufacture a pivot from roundoff.
    rows <- denseRows a
    require
        (not (null fixed) || any (\[Double]
r -> Double -> Double
forall a. Num a => a -> a
abs ([Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Double]
r) Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
1e-10 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double -> Double
forall a. Ord a => a -> a -> a
max Double
1 ([Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((Double -> Double) -> [Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map Double -> Double
forall a. Num a => a -> a
abs [Double]
r))) rows)
        "underdetermined zero-reaction Neumann/periodic problem; fix a gauge or use spectralPoisson"
    values <- solve rows rhs
    residual <- checked (maximum (0 : map abs (zipWith (-) (mv rows values) rhs)))
    Right (PDEResult values residual residual)

denseRows :: S.Sparse Double -> Either MathError [[Double]]
denseRows :: Sparse Double -> Either MathError [[Double]]
denseRows Sparse Double
a = do
    let (Int
n, Int
m) = Sparse Double -> (Int, Int)
forall a. Sparse a -> (Int, Int)
S.sparseShape Sparse Double
a
    Bool -> String -> Either MathError ()
require (Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
1024 Bool -> Bool -> Bool
&& Int
m Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
1024) String
"dense PDE reference budget exceeded"
    let entries :: Map (Int, Int) Double
entries = [((Int, Int), Double)] -> Map (Int, Int) Double
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList [((Int
i, Int
j), Double
v) | (Int
i, Int
j, Double
v) <- Sparse Double -> [(Int, Int, Double)]
forall a. Sparse a -> [(Int, Int, a)]
S.sparseEntries Sparse Double
a]
    [[Double]] -> Either MathError [[Double]]
forall a b. b -> Either a b
Right [[Double -> (Int, Int) -> Map (Int, Int) Double -> Double
forall k a. Ord k => a -> k -> Map k a -> a
Map.findWithDefault Double
0 (Int
i, Int
j) Map (Int, Int) Double
entries | Int
j <- [Int
0 .. Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]] | Int
i <- [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]

-- | Solve a unit-diffusivity zero-reaction Poisson problem with the selected spatial discretization.
poisson ::
    Discretization -> Grid -> [(Boundary, Boundary)] -> [Double] -> Either MathError PDEResult
poisson :: Discretization
-> Grid
-> [(Boundary, Boundary)]
-> [Double]
-> Either MathError PDEResult
poisson Discretization
method Grid
g [(Boundary, Boundary)]
bc [Double]
forcing = Discretization
-> Grid
-> [(Boundary, Boundary)]
-> [Double]
-> Double
-> Either MathError PDEOperator
assemblePDE Discretization
method Grid
g [(Boundary, Boundary)]
bc (Int -> Double -> [Double]
forall a. Int -> a -> [a]
replicate ([Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product (Grid -> [Int]
gridShape Grid
g)) Double
1) Double
0 Either MathError PDEOperator
-> (PDEOperator -> Either MathError PDEResult)
-> Either MathError PDEResult
forall a b.
Either MathError a
-> (a -> Either MathError b) -> Either MathError b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (PDEOperator -> [Double] -> Either MathError PDEResult
`solveElliptic` [Double]
forcing)

-- | Solve the corresponding homogeneous-source boundary-value problem.
laplace :: Discretization -> Grid -> [(Boundary, Boundary)] -> Either MathError PDEResult
laplace :: Discretization
-> Grid -> [(Boundary, Boundary)] -> Either MathError PDEResult
laplace Discretization
method Grid
g [(Boundary, Boundary)]
bc = Discretization
-> Grid
-> [(Boundary, Boundary)]
-> [Double]
-> Either MathError PDEResult
poisson Discretization
method Grid
g [(Boundary, Boundary)]
bc (Int -> Double -> [Double]
forall a. Int -> a -> [a]
replicate ([Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product (Grid -> [Int]
gridShape Grid
g)) Double
0)

-- | Solve -Laplacian(u)-k²u=f; resonances/singular systems fail.
helmholtz ::
    Discretization -> Grid -> [(Boundary, Boundary)] -> Double -> [Double] -> Either MathError PDEResult
helmholtz :: Discretization
-> Grid
-> [(Boundary, Boundary)]
-> Double
-> [Double]
-> Either MathError PDEResult
helmholtz Discretization
method Grid
g [(Boundary, Boundary)]
bc Double
k [Double]
forcing = do
    Bool -> String -> Either MathError ()
require (Double -> Bool
finite Double
k Bool -> Bool -> Bool
&& Double
k Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Double
0 Bool -> Bool -> Bool
&& Double -> Bool
finite (Double
k Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
k)) String
"invalid Helmholtz wave number"
    op <- Discretization
-> Grid
-> [(Boundary, Boundary)]
-> [Double]
-> Double
-> Either MathError PDEOperator
assemblePDE Discretization
method Grid
g [(Boundary, Boundary)]
bc (Int -> Double -> [Double]
forall a. Int -> a -> [a]
replicate ([Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product (Grid -> [Int]
gridShape Grid
g)) Double
1) (-Double
k Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
k)
    solveElliptic op forcing

-- | Second-order central spatial gradients; physical endpoints use second-order one-sided stencils.
spatialGradient :: Grid -> [Double] -> Either MathError [[Double]]
spatialGradient :: Grid -> [Double] -> Either MathError [[Double]]
spatialGradient g :: Grid
g@(Grid [Int]
dims [(Double, Double)]
_ [Bool]
periodic [Double]
hs) [Double]
values = do
    Bool -> String -> Either MathError ()
require ([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
values Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product [Int]
dims) String
"gradient grid/value mismatch"
    _ <- [Double] -> Either MathError [Double]
vector [Double]
values
    mapM
        ( \Int
axis ->
            (Int -> Either MathError Double)
-> [Int] -> Either MathError [Double]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM
                ( \Int
i -> do
                    ix <- Grid -> Int -> Either MathError [Int]
unflattenIndex Grid
g Int
i
                    let q = [Int]
ix [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
axis
                        n = [Int]
dims [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
axis
                        h = [Double]
hs [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
axis
                        at Int
j = [Double]
values [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((Int -> Int -> Int) -> [Int] -> [Int] -> [Int]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Int -> Int -> Int
forall a. Num a => a -> a -> a
(*) (Int -> Int -> [Int] -> [Int]
forall a. Int -> a -> [a] -> [a]
replace Int
axis (Int
j Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
n) [Int]
ix) ([Int] -> [Int]
strides [Int]
dims))
                        result
                            | [Bool]
periodic [Bool] -> Int -> Bool
forall a. HasCallStack => [a] -> Int -> a
!! Int
axis Bool -> Bool -> Bool
|| Int
q Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 Bool -> Bool -> Bool
&& Int
q Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1 = (Int -> Double
at (Int
q Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Double -> Double -> Double
forall a. Num a => a -> a -> a
- Int -> Double
at (Int
q Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
2 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
h)
                            | Int
q Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = (-Double
3 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Int -> Double
at Int
0 Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
4 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Int -> Double
at Int
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Int -> Double
at Int
2) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
2 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
h)
                            | Bool
otherwise = (Double
3 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Int -> Double
at (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
4 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Int -> Double
at (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2) Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Int -> Double
at (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
3)) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
2 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
h)
                    checked result
                )
                [Int
0 .. [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product [Int]
dims Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
        )
        [0 .. length dims - 1]

-- | Apply the FD Laplacian with ghost normal/Robin boundary data. Essential rows return prescribed-value minus current-value residuals.
laplacianStencil :: Grid -> [(Boundary, Boundary)] -> [Double] -> Either MathError [Double]
laplacianStencil :: Grid
-> [(Boundary, Boundary)] -> [Double] -> Either MathError [Double]
laplacianStencil Grid
g [(Boundary, Boundary)]
bc [Double]
values = do
    op <- Discretization
-> Grid
-> [(Boundary, Boundary)]
-> [Double]
-> Double
-> Either MathError PDEOperator
assemblePDE Discretization
FiniteDifference Grid
g [(Boundary, Boundary)]
bc (Int -> Double -> [Double]
forall a. Int -> a -> [a]
replicate ([Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product (Grid -> [Int]
gridShape Grid
g)) Double
1) Double
0
    applied <- S.spmv (pdeStiffness op) values
    vector (zipWith (-) (pdeBoundaryLoad op) applied)

waveNumbers :: Grid -> Either MathError [[Double]]
waveNumbers :: Grid -> Either MathError [[Double]]
waveNumbers (Grid [Int]
dims [(Double, Double)]
_ [Bool]
periodic [Double]
hs) = do
    Bool -> String -> Either MathError ()
require ([Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
and [Bool]
periodic) String
"spectral differentiation requires all-periodic grid"
    [Either MathError [Double]] -> Either MathError [[Double]]
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
forall (m :: * -> *) a. Monad m => [m a] -> m [a]
sequence [(Double -> Double) -> [Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (Double
2 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
forall a. Floating a => a
pi Double -> Double -> Double
forall a. Num a => a -> a -> a
*) ([Double] -> [Double])
-> Either MathError [Double] -> Either MathError [Double]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Double -> Either MathError [Double]
F.frequencies Int
n Double
h | (Int
n, Double
h) <- [Int] -> [Double] -> [(Int, Double)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int]
dims [Double]
hs]

-- | Periodic Fourier collocation derivative of order 0..8 along one axis; odd real Nyquist derivatives are zero.
spectralDerivative :: Grid -> Int -> Int -> [Double] -> Either MathError [Double]
spectralDerivative :: Grid -> Int -> Int -> [Double] -> Either MathError [Double]
spectralDerivative g :: Grid
g@(Grid [Int]
dims [(Double, Double)]
_ [Bool]
_ [Double]
_) Int
axis Int
order [Double]
values = do
    Bool -> String -> Either MathError ()
require
        (Int
axis Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
axis Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
dims Bool -> Bool -> Bool
&& Int
order Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
order Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
8 Bool -> Bool -> Bool
&& [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
values Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product [Int]
dims)
        String
"invalid spectral derivative axis/order/shape"
    ks <- Grid -> Either MathError [[Double]]
waveNumbers Grid
g
    spectrum <- F.fftND F.Forward dims (map (:+ 0) values)
    let multiplier Int
i =
            let q :: Int
q = Int
i Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` ([Int] -> [Int]
strides [Int]
dims [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
axis) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` ([Int]
dims [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
axis); k :: Double
k = [[Double]]
ks [[Double]] -> Int -> [Double]
forall a. HasCallStack => [a] -> Int -> a
!! Int
axis [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
q
             in -- Odd derivatives of a real Nyquist mode are zero.
                if Int -> Bool
forall a. Integral a => a -> Bool
odd Int
order Bool -> Bool -> Bool
&& Int -> Bool
forall a. Integral a => a -> Bool
even ([Int]
dims [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
axis) Bool -> Bool -> Bool
&& Int
q Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [Int]
dims [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
axis Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2 then Complex Double
0 else (Double
0 Double -> Double -> Complex Double
forall a. a -> a -> Complex a
:+ Double
k) Complex Double -> Int -> Complex Double
forall a b. (Num a, Integral b) => a -> b -> a
^ Int
order
    output <- F.fftND F.Inverse dims [z * multiplier i | (i, z) <- zip [0 ..] spectrum]
    vector (map realPart output)

-- | Sum periodic second Fourier derivatives over all spatial axes.
spectralLaplacian :: Grid -> [Double] -> Either MathError [Double]
spectralLaplacian :: Grid -> [Double] -> Either MathError [Double]
spectralLaplacian Grid
g [Double]
values = do
    terms <- (Int -> Either MathError [Double])
-> [Int] -> Either MathError [[Double]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (\Int
axis -> Grid -> Int -> Int -> [Double] -> Either MathError [Double]
spectralDerivative Grid
g Int
axis Int
2 [Double]
values) [Int
0 .. [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (Grid -> [Int]
gridShape Grid
g) Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
    vector (foldl1 (zipWith (+)) terms)

-- | Solve (-Laplacian+reaction)u=f spectrally; reaction>=0, and zero reaction requires a mean-free source with zero-mean solution gauge.
spectralPoisson :: Grid -> Double -> [Double] -> Either MathError PDEResult
spectralPoisson :: Grid -> Double -> [Double] -> Either MathError PDEResult
spectralPoisson g :: Grid
g@(Grid [Int]
dims [(Double, Double)]
_ [Bool]
_ [Double]
_) Double
reaction [Double]
forcing = do
    Bool -> String -> Either MathError ()
require
        (Double -> Bool
finite Double
reaction Bool -> Bool -> Bool
&& Double
reaction Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Double
0 Bool -> Bool -> Bool
&& [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
forcing Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product [Int]
dims)
        String
"invalid spectral Poisson reaction/source"
    ks <- Grid -> Either MathError [[Double]]
waveNumbers Grid
g
    spectrum <- F.fftND F.Forward dims (map (:+ 0) forcing)
    let eigenvalue Int
i =
            Double
reaction
                Double -> Double -> Double
forall a. Num a => a -> a -> a
+ [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum
                    [ ([[Double]]
ks [[Double]] -> Int -> [Double]
forall a. HasCallStack => [a] -> Int -> a
!! Int
axis [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! (Int
i Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` ([Int] -> [Int]
strides [Int]
dims [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
axis) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` ([Int]
dims [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
axis))) Double -> Int -> Double
forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
2 :: Int)
                    | Int
axis <- [Int
0 .. [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
dims Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
                    ]
    require
        (reaction > 0 || magnitude (spectrum !! 0) <= 1e-10 * max 1 (sum (map abs forcing)))
        "periodic Poisson source must have zero mean"
    values <-
        map realPart
            <$> F.fftND
                F.Inverse
                dims
                [if eigenvalue i == 0 then 0 else z / (eigenvalue i :+ 0) | (i, z) <- zip [0 ..] spectrum]
    applied <- spectralLaplacian g values
    residual <-
        checked (maximum (0 : map abs (zipWith3 (\Double
u Double
l Double
f -> Double
reaction Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
u Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
l Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
f) values applied forcing)))
    Right (PDEResult values residual residual)

{- | Pseudospectral divergence(kappa*gradient u). Products are evaluated at
collocation nodes; no implicit de-aliasing, so caller must resolve coefficients.
-}
pseudoSpectralDiffusion :: Grid -> [Double] -> [Double] -> Either MathError [Double]
pseudoSpectralDiffusion :: Grid -> [Double] -> [Double] -> Either MathError [Double]
pseudoSpectralDiffusion Grid
g [Double]
kappas [Double]
values = do
    Bool -> String -> Either MathError ()
require
        ([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
kappas Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
values Bool -> Bool -> Bool
&& (Double -> Bool) -> [Double] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
0) [Double]
kappas)
        String
"invalid pseudospectral coefficient field"
    _ <- [Double] -> Either MathError [Double]
vector [Double]
kappas
    terms <-
        mapM
            (\Int
axis -> Grid -> Int -> Int -> [Double] -> Either MathError [Double]
spectralDerivative Grid
g Int
axis Int
1 [Double]
values Either MathError [Double]
-> ([Double] -> Either MathError [Double])
-> Either MathError [Double]
forall a b.
Either MathError a
-> (a -> Either MathError b) -> Either MathError b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Grid -> Int -> Int -> [Double] -> Either MathError [Double]
spectralDerivative Grid
g Int
axis Int
1 ([Double] -> Either MathError [Double])
-> ([Double] -> [Double]) -> [Double] -> Either MathError [Double]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Double -> Double -> Double) -> [Double] -> [Double] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Double -> Double -> Double
forall a. Num a => a -> a -> a
(*) [Double]
kappas)
            [0 .. length (gridShape g) - 1]
    vector (foldl1 (zipWith (+)) terms)

{- | Method of lines with a constant assembled operator and checked time-source
callback. Essential Dirichlet values remain constant in time; initial state
must satisfy them. Consistent FEM mass is inverted on unconstrained rows.
-}
methodOfLines ::
    O.Method ->
    PDEOperator ->
    (Double -> [Double]) ->
    [Double] ->
    [Double] ->
    Either MathError O.Trajectory
methodOfLines :: Method
-> PDEOperator
-> (Double -> [Double])
-> [Double]
-> [Double]
-> Either MathError Trajectory
methodOfLines Method
method (PDEOperator Grid
g Sparse Double
stiffness Sparse Double
mass [Double]
boundary [(Int, Double)]
fixed) Double -> [Double]
source [Double]
times [Double]
initial = do
    let n :: Int
n = [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product (Grid -> [Int]
gridShape Grid
g)
    Bool -> String -> Either MathError ()
require
        (Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
1024 Bool -> Bool -> Bool
&& [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
initial Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
n Bool -> Bool -> Bool
&& ((Int, Double) -> Bool) -> [(Int, Double)] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (\(Int
i, Double
v) -> Double -> Double
forall a. Num a => a -> a
abs ([Double]
initial [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
i Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
v) Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
1e-10) [(Int, Double)]
fixed)
        String
"invalid PDE initial state/Dirichlet values or reference budget"
    a <- Sparse Double -> Either MathError [[Double]]
denseRows Sparse Double
stiffness
    m0 <- denseRows mass
    let m =
            [ if Int
i Int -> [Int] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` ((Int, Double) -> Int) -> [(Int, Double)] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Int, Double) -> Int
forall a b. (a, b) -> a
fst [(Int, Double)]
fixed then [Double]
row else [Double]
r
            | (Int
i, [Double]
r) <- [Int] -> [[Double]] -> [(Int, [Double])]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 ..] [[Double]]
m0
            , let row :: [Double]
row = [if Int
j Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
i then Double
1 else Double
0 | Int
j <- [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]
            ]
        rhs Double
t [Double]
u =
            let f :: [Double]
f = Double -> [Double]
source Double
t
                load :: [Double]
load = (Double -> Double -> Double -> Double)
-> [Double] -> [Double] -> [Double] -> [Double]
forall a b c d. (a -> b -> c -> d) -> [a] -> [b] -> [c] -> [d]
zipWith3 (\Double
v Double
b Double
au -> Double
v Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
b Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
au) ([[Double]] -> [Double] -> [Double]
mv [[Double]]
m0 [Double]
f) [Double]
boundary ([[Double]] -> [Double] -> [Double]
mv [[Double]]
a [Double]
u)
                constrained :: [Double]
constrained = [if Int
i Int -> [Int] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` ((Int, Double) -> Int) -> [(Int, Double)] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Int, Double) -> Int
forall a b. (a, b) -> a
fst [(Int, Double)]
fixed then Double
0 else Double
v | (Int
i, Double
v) <- [Int] -> [Double] -> [(Int, Double)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 ..] [Double]
load]
             in if [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
f Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
n Bool -> Bool -> Bool
|| Bool -> Bool
not ((Double -> Bool) -> [Double] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Double -> Bool
finite [Double]
f)
                    then Int -> Double -> [Double]
forall a. Int -> a -> [a]
replicate Int
n (Double
0 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
0)
                    else (MathError -> [Double])
-> ([Double] -> [Double]) -> Either MathError [Double] -> [Double]
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either ([Double] -> MathError -> [Double]
forall a b. a -> b -> a
const (Int -> Double -> [Double]
forall a. Int -> a -> [a]
replicate Int
n (Double
0 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
0))) [Double] -> [Double]
forall a. a -> a
id ([[Double]] -> [Double] -> Either MathError [Double]
solve [[Double]]
m [Double]
constrained)
    require
        (not (null times) && all finite times && and (zipWith (<) times (drop 1 times)))
        "method of lines needs finite increasing time samples"
    (_, states) <-
        foldM
            ( \([Double]
u, [[Double]]
acc) (Double
aTime, Double
bTime) -> do
                trajectory <- Method
-> ODEOptions
-> Int
-> (Double -> [Double] -> [Double])
-> Double
-> Double
-> [Double]
-> Either MathError Trajectory
O.solveIVP Method
method ODEOptions
O.defaultODEOptions Int
1 Double -> [Double] -> [Double]
rhs Double
aTime Double
bTime [Double]
u
                let next = [[Double]] -> [Double]
forall a. HasCallStack => [a] -> a
last (Trajectory -> [[Double]]
O.trajectoryStates Trajectory
trajectory)
                Right (next, acc ++ [next])
            )
            (initial, [initial])
            (zip times (drop 1 times))
    Right (O.Trajectory times states (zipWith rhs times states))

{- | Explicit heat/diffusion uses RK4, with a conservative forward-Euler CFL
check (therefore valid for RK4 too). Positive diffusivity and ordered times.
-}
heat ::
    Grid -> [(Boundary, Boundary)] -> Double -> [Double] -> [Double] -> Either MathError O.Trajectory
heat :: Grid
-> [(Boundary, Boundary)]
-> Double
-> [Double]
-> [Double]
-> Either MathError Trajectory
heat Grid
g [(Boundary, Boundary)]
bc Double
k = Discretization
-> Grid
-> [(Boundary, Boundary)]
-> [Double]
-> [Double]
-> [Double]
-> Either MathError Trajectory
diffusion Discretization
FiniteVolume Grid
g [(Boundary, Boundary)]
bc (Int -> Double -> [Double]
forall a. Int -> a -> [a]
replicate ([Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product (Grid -> [Int]
gridShape Grid
g)) Double
k)

-- | Variable positive nodal diffusivity, chosen spatial method and increasing output times; FEM uses a stricter consistent-mass CFL bound.
diffusion ::
    Discretization ->
    Grid ->
    [(Boundary, Boundary)] ->
    [Double] ->
    [Double] ->
    [Double] ->
    Either MathError O.Trajectory
diffusion :: Discretization
-> Grid
-> [(Boundary, Boundary)]
-> [Double]
-> [Double]
-> [Double]
-> Either MathError Trajectory
diffusion Discretization
method Grid
g [(Boundary, Boundary)]
bc [Double]
kappas [Double]
times [Double]
initial = do
    Bool -> String -> Either MathError ()
require (Bool -> Bool
not ([Double] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Double]
times)) String
"diffusion needs times"
    limit <- Grid -> Double -> [Double] -> Either MathError Double
stabilityLimit Grid
g ([Double] -> Double
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum (Double
0 Double -> [Double] -> [Double]
forall a. a -> [a] -> [a]
: [Double]
kappas)) (Int -> Double -> [Double]
forall a. Int -> a -> [a]
replicate ([Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (Grid -> [Int]
gridShape Grid
g)) Double
0)
    require
        ( and
            ( zipWith
                (\Double
a Double
b -> Double
b Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
a Bool -> Bool -> Bool
&& Double
b Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
a Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
limit Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (if Discretization
method Discretization -> Discretization -> Bool
forall a. Eq a => a -> a -> Bool
== Discretization
FiniteElement then Double
3 else Double
1))
                times
                (drop 1 times)
            )
        )
        "diffusion time step violates explicit CFL"
    op <- assemblePDE method g bc kappas 0
    methodOfLines O.RK4 op (const (replicate (length initial) 0)) times initial

{- | Periodic velocity-Verlet wave equation u_tt=c² Laplacian(u); returns u at
requested equally or unequally spaced increasing times. Spectral Laplacian.
-}
wave :: Grid -> Double -> [Double] -> [Double] -> [Double] -> Either MathError [[Double]]
wave :: Grid
-> Double
-> [Double]
-> [Double]
-> [Double]
-> Either MathError [[Double]]
wave Grid
g Double
speed [Double]
times [Double]
initial [Double]
velocity = do
    _ <- Grid -> Either MathError [[Double]]
waveNumbers Grid
g
    let n = [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product (Grid -> [Int]
gridShape Grid
g)
    require
        (finite speed && speed > 0 && not (null times) && length initial == n && length velocity == n)
        "invalid wave inputs"
    _ <- vector (times ++ initial ++ velocity)
    let limit = Double
2 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
speed Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double
forall a. Floating a => a -> a
sqrt ([Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [(Double
forall a. Floating a => a
pi Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
h) Double -> Int -> Double
forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
2 :: Int) | Double
h <- Grid -> [Double]
gridSpacing Grid
g]))
    require
        (and (zipWith (\Double
a Double
b -> Double
b Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
a Bool -> Bool -> Bool
&& Double
b Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
a Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
limit) times (drop 1 times)))
        "wave step violates spectral Verlet stability"
    let go [Double]
_ [Double]
_ [] = [[Double]] -> Either MathError [[Double]]
forall a b. b -> Either a b
Right []
        go [Double]
u [Double]
v ((Double
a, Double
b) : [(Double, Double)]
steps) = do
            acc <- (Double -> Double) -> [Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (Double
speed Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
speed Double -> Double -> Double
forall a. Num a => a -> a -> a
*) ([Double] -> [Double])
-> Either MathError [Double] -> Either MathError [Double]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Grid -> [Double] -> Either MathError [Double]
spectralLaplacian Grid
g [Double]
u
            let dt = Double
b Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
a; half = (Double -> Double -> Double) -> [Double] -> [Double] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\Double
vv Double
aa -> Double
vv Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
dt Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
aa Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
2) [Double]
v [Double]
acc
            next <- vector (zipWith (\Double
uu Double
vv -> Double
uu Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
dt Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
vv) u half)
            acc2 <- map (speed * speed *) <$> spectralLaplacian g next
            vv <- vector (zipWith (\Double
hh Double
aa -> Double
hh Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
dt Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
aa Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
2) half acc2)
            rest <- go next vv steps
            Right (next : rest)
    rest <- go initial velocity (zip times (drop 1 times))
    Right (initial : rest)

{- | Conservative first-order upwind periodic advection, central conservative
variable-coefficient diffusion. Explicit Euler preserves constants and total
mass for periodic data when the combined positivity CFL is satisfied.
-}
advection :: Grid -> [Double] -> [Double] -> [Double] -> Either MathError [[Double]]
advection :: Grid
-> [Double] -> [Double] -> [Double] -> Either MathError [[Double]]
advection Grid
g [Double]
velocity = Grid
-> [Double]
-> Double
-> [Double]
-> [Double]
-> Either MathError [[Double]]
advectionDiffusion Grid
g [Double]
velocity Double
0

-- | Periodic upwind advection plus central diffusion using explicit Euler and the combined positivity CFL bound.
advectionDiffusion ::
    Grid -> [Double] -> Double -> [Double] -> [Double] -> Either MathError [[Double]]
advectionDiffusion :: Grid
-> [Double]
-> Double
-> [Double]
-> [Double]
-> Either MathError [[Double]]
advectionDiffusion g :: Grid
g@(Grid [Int]
dims [(Double, Double)]
_ [Bool]
periodic [Double]
hs) [Double]
velocity Double
diffusivity [Double]
times [Double]
initial = do
    Bool -> String -> Either MathError ()
require
        ([Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
and [Bool]
periodic Bool -> Bool -> Bool
&& [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
velocity Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
dims Bool -> Bool -> Bool
&& [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
initial Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product [Int]
dims Bool -> Bool -> Bool
&& Bool -> Bool
not ([Double] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Double]
times))
        String
"advection reference requires periodic grid and matching data"
    limit <- Grid -> Double -> [Double] -> Either MathError Double
stabilityLimit Grid
g Double
diffusivity [Double]
velocity
    _ <- vector (times ++ initial)
    require
        (and (zipWith (\Double
a Double
b -> Double
b Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
a Bool -> Bool -> Bool
&& Double
b Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
a Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
limit) times (drop 1 times)))
        "advection-diffusion time step violates positivity CFL"
    let derivative [Double]
u Int
i =
            let ix :: [Int]
ix = [Int
i Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
s Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
n | (Int
s, Int
n) <- [Int] -> [Int] -> [(Int, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip ([Int] -> [Int]
strides [Int]
dims) [Int]
dims]
                at :: Int -> Int -> Double
at Int
axis Int
step =
                    [Double]
u [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((Int -> Int -> Int) -> [Int] -> [Int] -> [Int]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Int -> Int -> Int
forall a. Num a => a -> a -> a
(*) (Int -> Int -> [Int] -> [Int]
forall a. Int -> a -> [a] -> [a]
replace Int
axis (([Int]
ix [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
axis Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
step) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` ([Int]
dims [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
axis)) [Int]
ix) ([Int] -> [Int]
strides [Int]
dims))
             in [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum
                    [ let h :: Double
h = [Double]
hs [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
axis; v :: Double
v = [Double]
velocity [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
axis; left :: Double
left = Int -> Int -> Double
at Int
axis (-Int
1); right :: Double
right = Int -> Int -> Double
at Int
axis Int
1; center :: Double
center = [Double]
u [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
i
                       in -Double
v Double -> Double -> Double
forall a. Num a => a -> a -> a
* (if Double
v Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Double
0 then Double
center Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
left else Double
right Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
center) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
h
                            Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
diffusivity Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
right Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
2 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
center Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
left) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
h Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
h)
                    | Int
axis <- [Int
0 .. [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
dims Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
                    ]
    (_, states) <-
        foldM
            ( \([Double]
u, [[Double]]
acc) (Double
a, Double
b) -> do
                next <- [Double] -> Either MathError [Double]
vector [Double
uu Double -> Double -> Double
forall a. Num a => a -> a -> a
+ (Double
b Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
a) Double -> Double -> Double
forall a. Num a => a -> a -> a
* [Double] -> Int -> Double
derivative [Double]
u Int
i | (Int
i, Double
uu) <- [Int] -> [Double] -> [(Int, Double)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 ..] [Double]
u]
                Right (next, acc ++ [next])
            )
            (initial, [initial])
            (zip times (drop 1 times))
    Right states

-- | Conservative explicit positivity bound 1/sum(|v|/h+2*kappa/h^2); infinity denotes a zero spatial operator.
stabilityLimit :: Grid -> Double -> [Double] -> Either MathError Double
stabilityLimit :: Grid -> Double -> [Double] -> Either MathError Double
stabilityLimit Grid
g Double
diffusivity [Double]
velocity = do
    Bool -> String -> Either MathError ()
require
        ( Double -> Bool
finite Double
diffusivity
            Bool -> Bool -> Bool
&& Double
diffusivity Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Double
0
            Bool -> Bool -> Bool
&& [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
velocity Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (Grid -> [Int]
gridShape Grid
g)
            Bool -> Bool -> Bool
&& (Double -> Bool) -> [Double] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Double -> Bool
finite [Double]
velocity
        )
        String
"invalid CFL coefficients"
    let inverseStep :: Double
inverseStep = [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Double -> Double
forall a. Num a => a -> a
abs Double
v Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
h Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
2 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
diffusivity Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
h Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
h) | (Double
v, Double
h) <- [Double] -> [Double] -> [(Double, Double)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Double]
velocity (Grid -> [Double]
gridSpacing Grid
g)]
    Bool -> String -> Either MathError ()
require (Double -> Bool
finite Double
inverseStep) String
"CFL coefficient overflow"
    Double -> Either MathError Double
forall a b. b -> Either a b
Right (if Double
inverseStep Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
0 then Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
0 else Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
inverseStep)

-- | Requested nonnegative timestep divided by the conservative combined stability limit.
cflNumber :: Grid -> Double -> [Double] -> Double -> Either MathError Double
cflNumber :: Grid -> Double -> [Double] -> Double -> Either MathError Double
cflNumber Grid
g Double
diffusivity [Double]
velocity Double
dt = do
    Bool -> String -> Either MathError ()
require (Double -> Bool
finite Double
dt Bool -> Bool -> Bool
&& Double
dt Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Double
0) String
"invalid CFL timestep"
    limit <- Grid -> Double -> [Double] -> Either MathError Double
stabilityLimit Grid
g Double
diffusivity [Double]
velocity
    checked (dt / limit)

-- | Maximum absolute nodal difference from a finite exact-solution callback; this measures continuum error when an exact solution is known.
manufacturedError :: Grid -> ([Double] -> Double) -> [Double] -> Either MathError Double
manufacturedError :: Grid -> ([Double] -> Double) -> [Double] -> Either MathError Double
manufacturedError Grid
g [Double] -> Double
exact [Double]
values = do
    Bool -> String -> Either MathError ()
require ([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
values Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product (Grid -> [Int]
gridShape Grid
g)) String
"manufactured solution shape mismatch"
    errors <- [Double] -> Either MathError [Double]
vector (([Double] -> Double -> Double)
-> [[Double]] -> [Double] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\[Double]
x Double
u -> Double -> Double
forall a. Num a => a -> a
abs ([Double] -> Double
exact [Double]
x Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
u)) (Grid -> [[Double]]
gridCoordinates Grid
g) [Double]
values)
    Right (maximum (0 : errors))