Real.hs

{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE ForeignFunctionInterface #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE TypeFamilies #-}

{-# OPTIONS_GHC -fno-warn-orphans #-}

module Arpack.Foreign.Real () where

import qualified Control.Concurrent.Lock as Lock
import Control.Exception (bracket, throwIO)
import Data.Maybe (fromMaybe)
import qualified Data.Vector.Storable as VS
import Data.Vector.Storable.Mutable (IOVector)
import qualified Data.Vector.Storable.Mutable as VSM
import Foreign
import Foreign.C.String
import Foreign.C.Types (CChar)
import Foreign.Storable.Complex ()
import qualified Numeric.LinearAlgebra.Devel as Dense

import Arpack.Exceptions
import Arpack.Foreign.Class
import qualified Arpack.Lock as Arpack
import Arpack.Options

--
-- * Foreign functions
--

foreign import ccall unsafe "dnaupd_"
  dnaupd_ :: Ptr Int32  -- ido
          -> Ptr CChar  -- bmat
          -> Ptr Int32  -- n
          -> Ptr CChar  -- which
          -> Ptr Int32  -- nev
          -> Ptr Double -- tol
          -> Ptr Double -- resid
          -> Ptr Int32  -- ncv
          -> Ptr Double -- v
          -> Ptr Int32  -- ldv
          -> Ptr Int32  -- iparam
          -> Ptr Int32  -- ipntr
          -> Ptr Double -- workd
          -> Ptr Double -- workl
          -> Ptr Int32  -- lworkl
          -> Ptr Int32  -- info
          -> IO ()

foreign import ccall unsafe "dneupd_"
  dneupd_ :: Ptr Int32  -- rvec
          -> Ptr CChar  -- howmny
          -> Ptr Int32  -- select
          -> Ptr Double -- dr
          -> Ptr Double -- di
          -> Ptr Double -- z
          -> Ptr Int32  -- ldz
          -> Ptr Double -- sigmar
          -> Ptr Double -- sigmai
          -> Ptr Double -- workev
          -> CString    -- bmat
          -> Ptr Int32  -- n
          -> Ptr CChar  -- which
          -> Ptr Int32  -- nev
          -> Ptr Double -- tol
          -> Ptr Double -- resid
          -> Ptr Int32  -- ncv
          -> Ptr Double -- v
          -> Ptr Int32  -- ldv
          -> Ptr Int32  -- iparam
          -> Ptr Int32  -- ipntr
          -> Ptr Double -- workd
          -> Ptr Double -- workl
          -> Ptr Int32  -- lworkl
          -> Ptr Int32  -- info
          -> IO ()

--
-- * Types
--

data AUPD
  = AUPD
    { ido :: {-# UNPACK #-} !(Ptr Int32)
    , bmat :: {-# UNPACK #-} !(Ptr CChar)
    , n :: {-# UNPACK #-} !(Ptr Int32)
    , which :: {-# UNPACK #-} !(Ptr CChar)
    , nev :: {-# UNPACK #-} !(Ptr Int32)
    , tol :: {-# UNPACK #-} !(Ptr Double)
    , resid :: {-# UNPACK #-} !(IOVector Double)
    , ncv :: {-# UNPACK #-} !(Ptr Int32)
    , v :: {-# UNPACK #-} !(IOVector Double)
    , ldv :: {-# UNPACK #-} !(Ptr Int32)
    , iparam :: {-# UNPACK #-} !(IOVector Int32)
    , ipntr :: {-# UNPACK #-} !(IOVector Int32)
    , workd :: {-# UNPACK #-} !(IOVector Double)
    , workl :: {-# UNPACK #-} !(IOVector Double)
    , lworkl :: {-# UNPACK #-} !(Ptr Int32)
    , info :: {-# UNPACK #-} !(Ptr Int32)
    }

withAUPD :: Options Double -> Int -> (AUPD -> IO a) -> IO a
withAUPD options dim = bracket initAUPD freeAUPD where

  initAUPD = do
    let
      _nev = number options

      -- Largest number of basis vectors to use.
      -- Work per iteration is O(dim * ncv ^ 2).

      _ncv = min dim (4 * _nev)

      _lworkl = 3 * _ncv * _ncv + 6 * _ncv

    ido <- new 0
    bmat <- newCString "I"
    n <- new (fromIntegral dim)
    which <- newCString "SR"
    nev <- new (fromIntegral _nev)
    tol <- new 0
    resid <- VSM.new dim
    ncv <- new (fromIntegral _ncv)
    v <- VSM.new (dim * _ncv)
    ldv <- new (fromIntegral dim)

    iparam <- VSM.new 11
    -- shift strategy
    VSM.write iparam (1 - 1) 1
    -- maximum number of iterations
    VSM.write iparam (3 - 1)
      (fromIntegral (fromMaybe (3 * dim) (maxIterations options)))
    -- block size
    VSM.write iparam (4 - 1) 1
    -- eigenproblem type
    VSM.write iparam (7 - 1) 1

    ipntr <- VSM.new 14
    workd <- VSM.new (3 * dim)
    workl <- VSM.new _lworkl
    lworkl <- new (fromIntegral _lworkl)
    info <- new 0

    pure AUPD {..}

  freeAUPD AUPD {..} = do
    free ido
    free bmat
    free n
    free which
    free nev
    free tol
    free ncv
    free ldv
    free lworkl
    free info

data EUPD
  = EUPD
    { rvec :: {-# UNPACK #-} !(Ptr Int32)
    , howmny :: {-# UNPACK #-} !(Ptr CChar)
    , select :: {-# UNPACK #-} !(IOVector Int32)
    , dr :: {-# UNPACK #-} !(IOVector Double)
    , di :: {-# UNPACK #-} !(IOVector Double)
    , z :: {-# UNPACK #-} !(IOVector Double)
    , ldz :: {-# UNPACK #-} !(Ptr Int32)
    , sigmar :: {-# UNPACK #-} !(Ptr Double)
    , sigmai :: {-# UNPACK #-} !(Ptr Double)
    , workev :: {-# UNPACK #-} !(IOVector Double)
    }

withEUPD :: AUPD -> (EUPD -> IO a) -> IO a
withEUPD (AUPD {..}) = bracket initEUPD freeEUPD where

  initEUPD = do
    _nev <- fromIntegral <$> peek nev
    _ncv <- fromIntegral <$> peek ncv
    dim <- fromIntegral <$> peek n

    rvec <- new 1
    howmny <- newCString "A"
    select <- VSM.new _ncv
    dr <- VSM.new (_nev + 1)
    di <- VSM.new (_nev + 1)
    z <- VSM.new (dim * _nev)
    ldz <- new (fromIntegral dim)
    sigmar <- malloc
    sigmai <- malloc
    workev <- VSM.new (2 * _ncv)
    pure EUPD {..}

  freeEUPD (EUPD {..}) = do
    free rvec
    free howmny
    free ldz
    free sigmar
    free sigmai

--
-- * Wrappers
--

aupd :: AUPD -> IO ()
aupd (AUPD {..}) =
  VSM.unsafeWith resid $ \_resid ->
  VSM.unsafeWith v $ \_v ->
  VSM.unsafeWith iparam $ \_iparam ->
  VSM.unsafeWith ipntr $ \_ipntr ->
  VSM.unsafeWith workd $ \_workd ->
  VSM.unsafeWith workl $ \_workl ->

  dnaupd_ ido bmat n which nev tol _resid ncv _v ldv
          _iparam _ipntr _workd _workl lworkl info

eupd :: EUPD -> AUPD -> IO ()
eupd (EUPD {..}) (AUPD {..}) =
  VSM.unsafeWith select $ \_select ->
  VSM.unsafeWith dr $ \_dr ->
  VSM.unsafeWith di $ \_di ->
  VSM.unsafeWith z $ \_z ->
  VSM.unsafeWith workev $ \_workev ->
  VSM.unsafeWith resid $ \_resid ->
  VSM.unsafeWith v $ \_v ->
  VSM.unsafeWith iparam $ \_iparam ->
  VSM.unsafeWith ipntr $ \_ipntr ->
  VSM.unsafeWith workd $ \_workd ->
  VSM.unsafeWith workl $ \_workl ->

  dneupd_ rvec howmny _select _dr _di _z ldz sigmar sigmai _workev
          bmat n which nev tol _resid ncv _v ldv _iparam _ipntr
          _workd _workl lworkl info

instance Arpack Double where

  arpack !opts !dim !multiply
    -- These variables are all banged because we need to be strict
    -- in them _before_ we enter the locked segment of code! If we
    -- wait until we're inside the lock, and evaluating one of these
    -- variables invokes 'arpack' again, the program will deadlock!
    = withAUPD opts dim $ \stateA@(AUPD {..}) -> do

      let
        loop = do
          aupd stateA
          peek ido >>= \case
            99 -> do
              peek info >>= \case
                0 -> pure ()
                1 -> throwIO MaxIterations
                3 -> throwIO NoShifts
                i -> throwIO (XYAUPD i)

            i | abs i == 1 -> do
                  xi <- fromIntegral <$> VSM.read ipntr 0
                  let
                    x = VSM.slice (xi - 1) dim workd
                  yi <- fromIntegral <$> VSM.read ipntr 1
                  let
                    y = VSM.slice (yi - 1) dim workd
                  multiply y x
                  loop

              | otherwise -> throwIO (Unimplemented i)

        extract = withEUPD stateA $ \stateE@(EUPD {..}) -> do
          eupd stateE stateA

          peek info >>= \case
            0 -> pure ()
            1 -> throwIO Reallocate
            i -> throwIO (XYEUPD i)

          evals <- VS.unsafeFreeze (VSM.slice 0 (number opts) dr)
          evecs <- VS.unsafeFreeze (VSM.slice 0 (number opts * dim) z)
          let matrixFromVector = Dense.matrixFromVector Dense.ColumnMajor
          pure (evals, matrixFromVector dim (number opts) evecs)

      Lock.with Arpack.lock (loop >> extract)