<commit_before>{-# LANGUAGE MultiParamTypeClasses, FlexibleInstances #-}
-----------------------------------------------------------------------------
-- |
-- Module     : Data.Matrix.Herm.Banded
-- Copyright  : Copyright (c) , <NAME> <<EMAIL>>
-- License    : BSD3
-- Maintainer : <NAME> <<EMAIL>>
-- Stability  : experimental
--

module Data.Matrix.Herm.Banded (
    module Data.Matrix.Herm,
    module BLAS.Matrix.Immutable,
    module BLAS.Matrix.ReadOnly,

    hbmv,
    hbmm,
    ) where

import Control.Monad ( zipWithM_ )
import System.IO.Unsafe
import Unsafe.Coerce

import BLAS.Access
import BLAS.Elem ( BLAS2 )
import BLAS.C ( colMajor, cblasUpLo )
import BLAS.Types ( flipUpLo )
import qualified BLAS.Elem as E
import qualified BLAS.C as BLAS

import Data.Matrix.Banded.Internal
import Data.Matrix.Dense.Internal ( DMatrix, IOMatrix )
import Data.Vector.Dense.Internal
import qualified Data.Matrix.Banded.Internal as B
import qualified Data.Matrix.Dense.Internal as M
import qualified Data.Vector.Dense.Internal as V
import qualified Data.Vector.Dense.Operations as V

import Data.Matrix.Herm
import BLAS.Matrix.Immutable
import BLAS.Matrix.ReadOnly


instance (BLAS2 e) => IMatrix (Herm (BMatrix Imm)) e where
    (<*>) h x = unsafePerformIO $ getApply h x
    {-# NOINLINE (<*>) #-}

    (<**>) h a = unsafePerformIO $ getApplyMat h a
    {-# NOINLINE (<**>) #-}


instance (BLAS2 e) => RMatrix (Herm (BMatrix s)) e where
    getApply h x = do
        y <- newZero (dim x)
        hbmv 1 (unsafeCoerce h) x 1 y
        return (unsafeCoerce y)
    
    getApplyMat h a = do
        b <- newZero (shape a)
        hbmm 1 (unsafeCoerce h) a 1 b
        return (unsafeCoerce b)


hbmv :: (BLAS2 e) => e -> Herm (BMatrix t) (n,n) e -> DVector s n e -> e -> IOVector n e -> IO ()
hbmv alpha h x beta y
    | numRows h == 0 =
        return ()
    | isConj y = do
        V.doConj y
        hbmv alpha h x beta (V.conj y)
        V.doConj y
    | isConj x = do
        x' <- newCopy x
        V.doConj (V.unsafeThaw x')
        hbmv alpha h (conj x') beta y
    | otherwise =
        let order   = colMajor
            (u,e,a) = toBase h
            alpha'  = alpha * e
            n       = numCols a
            k       = case u of 
                          Upper -> numUpper a
                          Lower -> numLower a      
            (u',alpha'') 
                    = case (isHerm a) of
                          True  -> (flipUpLo u, E.conj alpha')
                          False -> (u, alpha')
            uploA   = cblasUpLo u'
            ldA     = ldaOf a
            incX    = strideOf x
            incY    = strideOf y
        in B.unsafeWithBasePtr a  $ \pA ->
               V.unsafeWithElemPtr x 0 $ \pX ->
                    V.unsafeWithElemPtr y 0 $ \pY -> do
                        BLAS.hbmv order uploA n k alpha'' pA ldA pX incX beta pY incY

hbmm :: (BLAS2 e) => e -> Herm (BMatrix t) (m,m) e -> DMatrix s (m,n) e -> e -> IOMatrix (m,n) e -> IO ()
hbmm alpha h b beta c =
    sequence_ $
        zipWith (\x y -> hbmv alpha h x beta y) (M.cols b) (M.cols c)<commit_msg>Data/Matrix/Herm/Bandes.hs: use zipWithM_ instead of sequence_/zipWith<commit_after>{-# LANGUAGE MultiParamTypeClasses, FlexibleInstances #-}
-----------------------------------------------------------------------------
-- |
-- Module     : Data.Matrix.Herm.Banded
-- Copyright  : Copyright (c) , <NAME> <<EMAIL>>
-- License    : BSD3
-- Maintainer : <NAME> <<EMAIL>>
-- Stability  : experimental
--

module Data.Matrix.Herm.Banded (
    module Data.Matrix.Herm,
    module BLAS.Matrix.Immutable,
    module BLAS.Matrix.ReadOnly,

    hbmv,
    hbmm,
    ) where

import Control.Monad ( zipWithM_ )
import System.IO.Unsafe
import Unsafe.Coerce

import BLAS.Access
import BLAS.Elem ( BLAS2 )
import BLAS.C ( colMajor, cblasUpLo )
import BLAS.Types ( flipUpLo )
import qualified BLAS.Elem as E
import qualified BLAS.C as BLAS

import Data.Matrix.Banded.Internal
import Data.Matrix.Dense.Internal ( DMatrix, IOMatrix )
import Data.Vector.Dense.Internal
import qualified Data.Matrix.Banded.Internal as B
import qualified Data.Matrix.Dense.Internal as M
import qualified Data.Vector.Dense.Internal as V
import qualified Data.Vector.Dense.Operations as V

import Data.Matrix.Herm
import BLAS.Matrix.Immutable
import BLAS.Matrix.ReadOnly


instance (BLAS2 e) => IMatrix (Herm (BMatrix Imm)) e where
    (<*>) h x = unsafePerformIO $ getApply h x
    {-# NOINLINE (<*>) #-}

    (<**>) h a = unsafePerformIO $ getApplyMat h a
    {-# NOINLINE (<**>) #-}


instance (BLAS2 e) => RMatrix (Herm (BMatrix s)) e where
    getApply h x = do
        y <- newZero (dim x)
        hbmv 1 (unsafeCoerce h) x 1 y
        return (unsafeCoerce y)
    
    getApplyMat h a = do
        b <- newZero (shape a)
        hbmm 1 (unsafeCoerce h) a 1 b
        return (unsafeCoerce b)


hbmv :: (BLAS2 e) => e -> Herm (BMatrix t) (n,n) e -> DVector s n e -> e -> IOVector n e -> IO ()
hbmv alpha h x beta y
    | numRows h == 0 =
        return ()
    | isConj y = do
        V.doConj y
        hbmv alpha h x beta (V.conj y)
        V.doConj y
    | isConj x = do
        x' <- newCopy x
        V.doConj (V.unsafeThaw x')
        hbmv alpha h (conj x') beta y
    | otherwise =
        let order   = colMajor
            (u,e,a) = toBase h
            alpha'  = alpha * e
            n       = numCols a
            k       = case u of 
                          Upper -> numUpper a
                          Lower -> numLower a      
            (u',alpha'') 
                    = case (isHerm a) of
                          True  -> (flipUpLo u, E.conj alpha')
                          False -> (u, alpha')
            uploA   = cblasUpLo u'
            ldA     = ldaOf a
            incX    = strideOf x
            incY    = strideOf y
        in B.unsafeWithBasePtr a  $ \pA ->
               V.unsafeWithElemPtr x 0 $ \pX ->
                    V.unsafeWithElemPtr y 0 $ \pY -> do
                        BLAS.hbmv order uploA n k alpha'' pA ldA pX incX beta pY incY

hbmm :: (BLAS2 e) => e -> Herm (BMatrix t) (m,m) e -> DMatrix s (m,n) e -> e -> IOMatrix (m,n) e -> IO ()
hbmm alpha h b beta c =
    zipWithM_ (\x y -> hbmv alpha h x beta y) (M.cols b) (M.cols c)
{'fraction_non_alphanumeric': 0.1334505354003516, 'fraction_numerical': 0.004155346012466038, 'mean_word_length': 1.6939302625914765, 'pattern_counts': {'":': 0, '<': 29, '<?xml version=': 0, '>': 67, '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': '5364355', 'n_tokens_mistral': 2183, 'n_tokens_neox': 2085, 'n_words': 663}