<commit_before>module Cloth where

import Data.Array.Accelerate  --         (Vector, Segments, Acc, use, (?))
import qualified Data.Array.Accelerate as Acc
import Data.Array.Accelerate.Math
import Data.Array.Accelerate.Math.SMVM
import Data.Array.Accelerate.Types


-- TODO Maybe this should return Acc (Vector Float)?
filterMPCG :: AccVector Float -> AccVector Float
filterMPCG v = v -- TODO Not implemented until now


mpcgSingleStep :: AccSparseMatrix Float -> AccVector Float -> AccVector Float -> AccVector Float -> AccScalar Float ->
  (AccVector Float, AccVector Float, AccVector Float, AccScalar Float)
mpcgSingleStep a dv r c delta = (dv', r', c', delta')
  where
    p_inv     = a -- TODO
    q         = filterMPCG (smvmAcc a c)
    p_cq      = dotpAcc c q
--    alpha   = rho / -- TODO
    alpha     = use $ Acc.fromList (Z) [1]   :: AccScalar Float
    alpha_c   = alpha *. c                   :: AccVector Float
    dv'       = Acc.zipWith (+) dv alpha_c   :: AccVector Float
    alpha_q   = alpha *. q
    r         = Acc.zipWith (-) r alpha_q
    s         = smvmAcc p_inv r 
    delta'    = dotpAcc r s                  :: AccScalar Float
    beta      = Acc.zipWith (/) delta' delta :: AccScalar Float
    beta_c    = beta *. c                    :: AccVector Float
    c'        = filterMPCG (Acc.zipWith (+) s beta_c)
    r' = r


-- Compute 'dv' in 'A * dv = B' with preconditioned conjugate gradient method.
--
-- * SparseMatrix in compressed row format
--
-- * b' given as 'normal' vector
--
-- * p is ...
--
-- * z is ...
--
-- * epsilon is the maximal error
--
-- * z is ...
--
-- mpcgAcc :: AccSparseMatrix Float -> AccVector Float -> AccSparseMatrix Float -> AccVector Float
--   -> AccScalar Float -> Int -> AccVector Float
-- mpcgAcc a b p z epsilon count = step dv r c rho_new count
--   where
--     p_inv   = p -- vectorInverseAcc p -- TODO This is a matrix, not a vector!
--     dv      = z
--     rho     = dotpAcc (filterMPCG b) (smvmAcc p b)
--     r       = filterMPCG (Acc.zipWith (-) b (smvmAcc a dv))
--     c       = smvmAcc p_inv r<commit_msg>Single step finished. Initialization still buggy.<commit_after>module Cloth where

import Data.Array.Accelerate  --         (Vector, Segments, Acc, use, (?))
import qualified Data.Array.Accelerate as Acc
import Data.Array.Accelerate.Math
import Data.Array.Accelerate.Math.SMVM
import Data.Array.Accelerate.Types


filterMPCG :: AccVector Float -> AccVector Float
filterMPCG v = v -- TODO Not implemented until now


-- TODO p_inv could be represented as a single vector!
mpcgSingleStep :: AccSparseMatrix Float -> AccSparseMatrix Float -> AccVector Float -> AccVector Float -> AccVector Float -> AccScalar Float -> Int ->
  AccVector Float
mpcgSingleStep p_inv a dv r c delta n
  | n == 0    = dv
  | otherwise = mpcgSingleStep p_inv a dv' r' c' delta' (pred n)
  where
    q         = filterMPCG (smvmAcc a c)
    p_cq      = dotpAcc c q
    alpha     = delta /. p_cq
    alpha_c   = alpha *. c                   :: AccVector Float
    dv'       = Acc.zipWith (+) dv alpha_c   :: AccVector Float
    alpha_q   = alpha *. q
    r         = Acc.zipWith (-) r alpha_q    :: AccVector Float
    s         = smvmAcc p_inv r              :: AccVector Float
    delta'    = dotpAcc r s                  
    beta      = Acc.zipWith (/) delta' delta :: AccScalar Float
    beta_c    = beta *. c                    :: AccVector Float
    c'        = filterMPCG (Acc.zipWith (+) s beta_c)
    r' = r


mpcgInitial a epsilon b z n = mpcgSingleStep p_inv a dv r c delta n
  where
    p_inv   = a -- smvmAcc (smunity $ smrows a) a -- This should be quadradic?!
    p       = p -- smScalarMul
    dv      = z
    rho     = dotpAcc (filterMPCG b) (smvmAcc p b)
    r       = filterMPCG (Acc.zipWith (-) b (smvmAcc a dv))
    c       = smvmAcc p_inv r
    delta   = dotpAcc r c

-- Compute 'dv' in 'A * dv = B' with preconditioned conjugate gradient method.
--
-- * SparseMatrix in compressed row format
--
-- * b' given as 'normal' vector
--
-- * p is ...
--
-- * z is ...
--
-- * epsilon is the maximal error
--
-- * z is ...
--
-- mpcgAcc :: AccSparseMatrix Float -> AccVector Float -> AccSparseMatrix Float -> AccVector Float
--   -> AccScalar Float -> Int -> AccVector Float
-- mpcgAcc a b p z epsilon count = step dv r c rho_new count
--   where
--     p_inv   = p -- vectorInverseAcc p -- TODO This is a matrix, not a vector!
--     dv      = z
--     rho     = dotpAcc (filterMPCG b) (smvmAcc p b)
--     r       = filterMPCG (Acc.zipWith (-) b (smvmAcc a dv))
--     c       = smvmAcc p_inv r
{'fraction_non_alphanumeric': 0.11222927149420259, 'fraction_numerical': 0.0004375410194705754, 'mean_word_length': 1.9901896664486594, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 29, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1010365', 'n_tokens_mistral': 1494, 'n_tokens_neox': 1460, 'n_words': 520}