`Arr` is a type alias for convenience, `CGState` is the state of a CG iteration.
CG takes a function implementing a linear operator (type `Arr sh U -> Arr sh D`), a right hand side and an initial guess and returns the (lazy, infinte) list of the `CGState`s of all iterations.
Computations use Repa's parallel computation in a straight-forward way and needs to be performed in a monad. We use the `Identity` monad which does essentially nothing except for providing the sequencing. All operations are performed strictly using `BangPatterns`.
%% Cell type:code id: tags:
``` haskell
cg :: Shape sh => (Arr sh U -> Arr sh D) -> Arr sh U -> Arr sh U -> [CGState sh]
-- the main iteration; maps a CGState to another CGState
cgStep (CGState x p r r2) =
runIdentity $ do
-- need inline type annotation here because computeP is polymorphic in its
-- output type (syntax like this requires ScopedTypeVariables)
!(q :: Arr sh U) <- computeP $ op p
!qp <- sumAllP $ q *^ p
let alpha = r2 / qp
!x' <- computeP $ x +^ scale alpha p
!r' <- computeP $ r -^ scale alpha q
!r2' <- normSquaredP r'
let beta = r2' / r2
!p' <- computeP $ r' +^ scale beta p
return $ CGState x' p' r' r2'
```
%% Cell type:markdown id: tags:
Using lazyness, the iterations can be decoupled from stopping rule and output of information.
%% Cell type:code id: tags:
``` haskell
takeUntil :: (a -> Bool) -> [a] -> [a]
takeUntil _ [] = []
takeUntil predicate (x:xs)
| predicate x = [x]
| otherwise = x : takeUntil predicate xs
process :: Monad m => [a] -> (a -> m ()) -> m a
process xs f = foldM (\_ x -> f x >> return x) undefined xs
runCG :: Shape sh => Double -> (Arr sh U -> Arr sh D) -> Arr sh U -> IO (Arr sh U)
runCG tol op rhs = do
let initial = computeS $ fromFunction (extent rhs) (const 0)
let steps' = cg op rhs initial
let r20 = cgr2 $ head steps'
let steps = takeUntil (\x -> sqrt (cgr2 x / r20) < tol) steps'
result <-process(zip[(1::Int)..]steps)$\(n,cgs)->
putStrLn $ show n ++ " " ++ show (sqrt $ cgr2 cgs / r20)
return $ cgx . snd $ result
```
%% Cell type:markdown id: tags:
- `takeUntil` is like `takeWhile`, but also returns the final iterate (why waste it?).
- `process` runs over a list, performs a monadic action on each element and returns the last element. It does not retain the start of the list while iterating, so it gets garbage collected properly.
- `runCG` executes `cg` with initial guess 0, a stopping rule based on residuals relative to the first initial guess, and outputs the residual while iterating.
