diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,5 @@ +# Revision history for trexio-hs + +## 0.1.0 -- YYYY-mm-dd + +* First version. Released on an unsuspecting world. diff --git a/src-int/TREXIO/Internal/Base.hsc b/src-int/TREXIO/Internal/Base.hsc new file mode 100644 --- /dev/null +++ b/src-int/TREXIO/Internal/Base.hsc @@ -0,0 +1,255 @@ +module TREXIO.Internal.Base where + +import Control.Exception.Safe +import Control.Monad.IO.Class +import Foreign +import Foreign.C.String +import Foreign.C.Types +import GHC.Generics +import System.IO.Unsafe (unsafePerformIO) +import Foreign.C.ConstPtr + +#include + +type ExitCodeC = #type trexio_exit_code + +-- | TREXIO Exit Codes. Can be thrown as Haskell exceptions in IO. +data ExitCode + = Failure + | Success + | InvalidArg1 + | InvalidArg2 + | InvalidArg3 + | InvalidArg4 + | InvalidArg5 + | End + | ReadOnly + | Errno + | InvalidID + | AllocationFailed + | HasNot + | InvalidNum + | AttrAlreadyExists + | DSetAlreadyExists + | OpenError + | LockError + | UnlockError + | FileError + | GroupReadError + | GroupWriteError + | ElemReadError + | ElemWriteError + | UnsafeArrayDim + | AttrMissing + | DSetMissing + | BackEndMissing + | InvalidArg6 + | InvalidArg7 + | InvalidArg8 + | InvalidStrLen + | IntSizeOverflow + | SafeMode + | InvalidElectronNum + | InvalidDeterminantNum + | InvalidState + | VersionParsingIssue + | PhaseChange + deriving (Show, Eq, Ord, Generic) + +instance Enum ExitCode where + fromEnum Failure = #const TREXIO_FAILURE + fromEnum Success = #const TREXIO_SUCCESS + fromEnum InvalidArg1 = #const TREXIO_INVALID_ARG_1 + fromEnum InvalidArg2 = #const TREXIO_INVALID_ARG_2 + fromEnum InvalidArg3 = #const TREXIO_INVALID_ARG_3 + fromEnum InvalidArg4 = #const TREXIO_INVALID_ARG_4 + fromEnum InvalidArg5 = #const TREXIO_INVALID_ARG_5 + fromEnum End = #const TREXIO_END + fromEnum ReadOnly = #const TREXIO_READONLY + fromEnum Errno = #const TREXIO_ERRNO + fromEnum InvalidID = #const TREXIO_INVALID_ID + fromEnum AllocationFailed = #const TREXIO_ALLOCATION_FAILED + fromEnum HasNot = #const TREXIO_HAS_NOT + fromEnum InvalidNum = #const TREXIO_INVALID_NUM + fromEnum AttrAlreadyExists = #const TREXIO_ATTR_ALREADY_EXISTS + fromEnum DSetAlreadyExists = #const TREXIO_DSET_ALREADY_EXISTS + fromEnum OpenError = #const TREXIO_OPEN_ERROR + fromEnum LockError = #const TREXIO_LOCK_ERROR + fromEnum UnlockError = #const TREXIO_UNLOCK_ERROR + fromEnum FileError = #const TREXIO_FILE_ERROR + fromEnum GroupReadError = #const TREXIO_GROUP_READ_ERROR + fromEnum GroupWriteError = #const TREXIO_GROUP_WRITE_ERROR + fromEnum ElemReadError = #const TREXIO_ELEM_READ_ERROR + fromEnum ElemWriteError = #const TREXIO_ELEM_WRITE_ERROR + fromEnum UnsafeArrayDim = #const TREXIO_UNSAFE_ARRAY_DIM + fromEnum AttrMissing = #const TREXIO_ATTR_MISSING + fromEnum DSetMissing = #const TREXIO_DSET_MISSING + fromEnum BackEndMissing = #const TREXIO_BACK_END_MISSING + fromEnum InvalidArg6 = #const TREXIO_INVALID_ARG_6 + fromEnum InvalidArg7 = #const TREXIO_INVALID_ARG_7 + fromEnum InvalidArg8 = #const TREXIO_INVALID_ARG_8 + fromEnum InvalidStrLen = #const TREXIO_INVALID_STR_LEN + fromEnum IntSizeOverflow = #const TREXIO_INT_SIZE_OVERFLOW + fromEnum SafeMode = #const TREXIO_SAFE_MODE + fromEnum InvalidElectronNum = #const TREXIO_INVALID_ELECTRON_NUM + fromEnum InvalidDeterminantNum = #const TREXIO_INVALID_DETERMINANT_NUM + fromEnum InvalidState = #const TREXIO_INVALID_STATE + fromEnum VersionParsingIssue = #const TREXIO_VERSION_PARSING_ISSUE + fromEnum PhaseChange = #const TREXIO_PHASE_CHANGE + + toEnum (#const TREXIO_FAILURE) = Failure + toEnum (#const TREXIO_SUCCESS) = Success + toEnum (#const TREXIO_INVALID_ARG_1) = InvalidArg1 + toEnum (#const TREXIO_INVALID_ARG_2) = InvalidArg2 + toEnum (#const TREXIO_INVALID_ARG_3) = InvalidArg3 + toEnum (#const TREXIO_INVALID_ARG_4) = InvalidArg4 + toEnum (#const TREXIO_INVALID_ARG_5) = InvalidArg5 + toEnum (#const TREXIO_END) = End + toEnum (#const TREXIO_READONLY) = ReadOnly + toEnum (#const TREXIO_ERRNO) = Errno + toEnum (#const TREXIO_INVALID_ID) = InvalidID + toEnum (#const TREXIO_ALLOCATION_FAILED) = AllocationFailed + toEnum (#const TREXIO_HAS_NOT) = HasNot + toEnum (#const TREXIO_INVALID_NUM) = InvalidNum + toEnum (#const TREXIO_ATTR_ALREADY_EXISTS) = AttrAlreadyExists + toEnum (#const TREXIO_DSET_ALREADY_EXISTS) = DSetAlreadyExists + toEnum (#const TREXIO_OPEN_ERROR) = OpenError + toEnum (#const TREXIO_LOCK_ERROR) = LockError + toEnum (#const TREXIO_UNLOCK_ERROR) = UnlockError + toEnum (#const TREXIO_FILE_ERROR) = FileError + toEnum (#const TREXIO_GROUP_READ_ERROR) = GroupReadError + toEnum (#const TREXIO_GROUP_WRITE_ERROR) = GroupWriteError + toEnum (#const TREXIO_ELEM_READ_ERROR) = ElemReadError + toEnum (#const TREXIO_ELEM_WRITE_ERROR) = ElemWriteError + toEnum (#const TREXIO_UNSAFE_ARRAY_DIM) = UnsafeArrayDim + toEnum (#const TREXIO_ATTR_MISSING) = AttrMissing + toEnum (#const TREXIO_DSET_MISSING) = DSetMissing + toEnum (#const TREXIO_BACK_END_MISSING) = BackEndMissing + toEnum (#const TREXIO_INVALID_ARG_6) = InvalidArg6 + toEnum (#const TREXIO_INVALID_ARG_7) = InvalidArg7 + toEnum (#const TREXIO_INVALID_ARG_8) = InvalidArg8 + toEnum (#const TREXIO_INVALID_STR_LEN) = InvalidStrLen + toEnum (#const TREXIO_INT_SIZE_OVERFLOW) = IntSizeOverflow + toEnum (#const TREXIO_SAFE_MODE) = SafeMode + toEnum (#const TREXIO_INVALID_ELECTRON_NUM) = InvalidElectronNum + toEnum (#const TREXIO_INVALID_DETERMINANT_NUM) = InvalidDeterminantNum + toEnum (#const TREXIO_INVALID_STATE) = InvalidState + toEnum (#const TREXIO_VERSION_PARSING_ISSUE) = VersionParsingIssue + toEnum (#const TREXIO_PHASE_CHANGE) = PhaseChange + toEnum _ = error "toEnum(ExitCode): invalid argument" + +instance Exception ExitCode where + displayException = stringOfError + +exitCodeH :: ExitCodeC -> ExitCode +exitCodeH = toEnum . fromIntegral + +-- | Helper to make a C function returning an exit code to an easier to handle +-- exception +checkEC :: (MonadIO m, MonadThrow m) => IO ExitCodeC -> m () +checkEC f = do + ec <- exitCodeH <$> liftIO f + if ec == Success + then return () + else throwIO ec + +foreign import capi "trexio.h trexio_string_of_error" stringOfError_ :: (#type trexio_exit_code) -> IO (ConstPtr CChar) +stringOfError :: ExitCode -> String +stringOfError exitCode = unsafePerformIO $ do + stringOfError_ exitCodeInt >>= peekCString . unConstPtr + where + exitCodeInt = fromIntegral . fromEnum $ exitCode + +-------------------------------------------------------------------------------- + +-- | TREXIO version: Major, Minor, Patch +version :: (Int, Int, Int) +version = + ( #const TREXIO_VERSION_MAJOR + , #const TREXIO_VERSION_MINOR + , #const TREXIO_VERSION_PATCH + ) + +-------------------------------------------------------------------------------- + +data Backend + = Hdf5 + | Text + | Invalid + | Auto + deriving (Eq, Ord, Show) + +instance Enum Backend where + fromEnum Hdf5 = #const TREXIO_HDF5 + fromEnum Text = #const TREXIO_TEXT + fromEnum Invalid = #const TREXIO_INVALID_BACK_END + fromEnum Auto = #const TREXIO_AUTO + + toEnum (#const TREXIO_HDF5) = Hdf5 + toEnum (#const TREXIO_TEXT) = Text + toEnum (#const TREXIO_INVALID_BACK_END) = Invalid + toEnum (#const TREXIO_AUTO) = Auto + toEnum _ = error "toEnum(Backend): invalid argument" + +backendC :: Backend -> (#type back_end_t) +backendC = fromIntegral . fromEnum + +foreign import capi "trexio.h trexio_has_backend" hasBackend_ :: (#type back_end_t) -> CBool +hasBackend :: Backend -> Bool +hasBackend backend = toBool . hasBackend_ . fromIntegral . fromEnum $ backend + +-------------------------------------------------------------------------------- + +-- | TREXIO file handle +newtype Trexio = Trexio (Ptr Trexio) + +-- | File access mode for 'Trexio'. +data FileMode + = FileRead + | FileWrite + | FileUnsafe -- ^ Allows to delete and override data + deriving (Eq, Show, Ord, Generic) + +modeC :: FileMode -> CChar +modeC FileRead = castCharToCChar 'r' +modeC FileWrite = castCharToCChar 'w' +modeC FileUnsafe = castCharToCChar 'u' + +foreign import capi "trexio.h trexio_open" open_ :: + CString -> + CChar -> + (#type back_end_t) -> + Ptr (#type trexio_exit_code) -> + IO Trexio +open :: (MonadIO m, MonadThrow m) => FilePath -> FileMode -> Backend -> m Trexio +open filepath mode backend = liftIO . withCString filepath $ \filepathC -> + alloca $ \ecPtr -> do + trexio <- open_ filepathC (modeC mode) (backendC backend) ecPtr + ec <- exitCodeH <$> peek ecPtr + if ec == Success + then return trexio + else throwM ec + +foreign import capi "trexio.h trexio_close" close_ :: Trexio -> IO ExitCodeC +close :: (MonadIO m, MonadThrow m) => Trexio -> m () +close trexio = checkEC $ close_ trexio + +foreign import capi "trexio.h trexio_get_int64_num" intsPerDet_ :: + Trexio -> + Ptr CInt -> + IO ExitCodeC + +-- | Get the number of 'Int64's required to store a single determinant. This is a +-- convenience function that avoids manual calculation via the size if Int64 and +-- the number of MOs. +intsPerDet :: (MonadIO m, MonadThrow m) => Trexio -> m Int +intsPerDet trexio = liftIO . alloca $ \nPtr -> do + checkEC $ intsPerDet_ trexio nPtr + fromIntegral <$> peek nPtr + +foreign import capi "trexio.h trexio_mark_safety" markSafety_ :: Trexio -> Int32 -> IO ExitCodeC + +-- | After a file has been opened unsafely, it can be marked as safe by this function +-- forcefully. You are responsible for internal incosistencies that may arise from this. +markSafety :: (MonadIO m, MonadThrow m) => Trexio -> m () +markSafety trexio = checkEC $ markSafety_ trexio 0 diff --git a/src/TREXIO.hs b/src/TREXIO.hs new file mode 100644 --- /dev/null +++ b/src/TREXIO.hs @@ -0,0 +1,102 @@ +{-# LANGUAGE TemplateHaskell #-} + +{- | +Module: TREXIO +Description: High-Level bindings to TREXIO library +Copyright: Phillip Seeber 2024 +License: BSD-3-Clause +Maintainer: phillip.seeber@uni-jena.de +Stability: experimental +Portability: POSIX + +This module provides the high-level bindings to the TREXIO library for wave function data. +The 'TREXIO.HighLevel' modules provides complete high-level bindings to the TREXIO and the function names are generated by stripping the @trexio_@ prefix from the C function names and converting to camel case. +E.g. the C function 'TREXIO.LowLevel.trexio_read_rdm_2e_updn_cholesky' is available as 'readRdm2eUpdnCholesky'. + +The high-level bindings abstract mainly over three aspects of the C- (and also Python-API): + +Memory management is done by the Haskell garbage collector and no pointers need to be moved around. +All multidimensional data is safely handled by the 'Data.Massiv.Array.Array' type or the 'TREXIO.CooArray.CooArray' type. + +Error handling is done by throwing 'TREXIO.ExitCode' exceptions in 'IO', i.e. you don't need to check for error codes manually. +You may 'catch' 'ExitCode' exceptions via the usual mechanisms, however. + +In the Python- and C-APIs, Mutlidimensional quantities require writing their size to another field, referenced by the TREXIO specification first. +For example, see this example for Python in TREXIO: + +@ +import trexio +coord = [ # xyz coordinates in atomic units + [0. , 0., -0.24962655], + [0. , 2.70519714, 1.85136466], + [0. , -2.70519714, 1.85136466] +] +with trexio.File("water.trexio", 'w', + back_end=trexio.TREXIO_HDF5) as f: + trexio.write_nucleus_num(f, len(coord)) + trexio.write_nucleus_coord(f, coord) +@ + +This high-level API abstracts over this and automatically writes the size of the array to the corresponding field. +Safety checks are employed to ensure, should the size already exist, that it is consistent with the array size and other arrays utilising the same size field. +Should this safety check be violated, an the 'AttrAlreadyExists' exception will be thrown, as the corresponding size field already exists and is inconsistent with the new size. +Thus, the Haskell equivalent to this is: + +@ +import TREXIO +import Data.Massiv.Array as Massiv + +coord <- Massiv.fromListsM Par + [ [0. , 0., -0.24962655] + , [0. , 2.70519714, 1.85136466] + , [0. , -2.70519714, 1.85136466] + ] +withTrexio "water.trexio" FileWrite Hdf5 $ \\trexio -> + writeNucleusCoord trexio coord +@ + +-} +module TREXIO ( + -- * Basic Operations + ExitCode (..), + ExitCodeC, + version, + Backend (..), + Trexio, + FileMode (..), + hasBackend, + open, + close, + markSafety, + TrexioScheme (..), + GroupName, + Group (..), + Typ (..), + Length (..), + DimLength (..), + DataName, + + -- * High Level Interface + scheme, + intsPerDet, + withTrexio, + module TREXIO.HighLevel, +) where + +import Control.Exception.Safe +import Control.Monad.IO.Class +import Data.Aeson +import Language.Haskell.TH +import Language.Haskell.TH.Syntax (lift) +import TREXIO.HighLevel +import TREXIO.Internal.Base +import TREXIO.Internal.TH (DataName, Group (..), GroupName, TrexioScheme (..), Typ (..), Length (..), DimLength (..)) +import TREXIO.LowLevel.Scheme (scheme) + +-- | Work safely with a TREXIO file handle. Prefer over 'open' and 'close'. +withTrexio :: (MonadMask m, MonadIO m) => FilePath -> FileMode -> Backend -> (Trexio -> m a) -> m a +withTrexio path mode backend f = + bracket + (liftIO $ open path mode backend) + (\trexio -> liftIO . close $ trexio) + f diff --git a/src/TREXIO/CooArray.hs b/src/TREXIO/CooArray.hs new file mode 100644 --- /dev/null +++ b/src/TREXIO/CooArray.hs @@ -0,0 +1,90 @@ +{-# LANGUAGE UndecidableInstances #-} + +{- | +Module: TREXIO.CooArray +Description: Coordinate list sparse array representation for TREXIO +Copyright: Phillip Seeber 2024 +License: BSD-3-Clause +Maintainer: phillip.seeber@uni-jena.de +Stability: experimental +Portability: POSIX +-} +module TREXIO.CooArray ( + CooArray, + values, + coords, + cooSize, + mkCooArrayF, + mkCooArray, +) +where + +import Data.Foldable +import Data.Massiv.Array as Massiv hiding (all, toList) +import GHC.Generics (Generic) + +-- | A coordinate list array representation. +data CooArray r ix a = CooArray + { values_ :: Vector r a + , coords_ :: Vector r ix + , cooSize_ :: Sz ix + -- ^ Size of the COO array + } + deriving (Generic) + +instance (Eq (Vector r a), Eq (Vector r ix), Eq ix) => Eq (CooArray r ix a) where + CooArray v1 c1 s1 == CooArray v2 c2 s2 = v1 == v2 && c1 == c2 && s1 == s2 + +instance (Ord (Vector r a), Ord (Vector r ix), Ord ix) => Ord (CooArray r ix a) where + compare (CooArray v1 c1 s1) (CooArray v2 c2 s2) = + compare v1 v2 <> compare c1 c2 <> compare s1 s2 + +instance + (Show (Vector r a), Show (Vector r ix), Index ix, Show ix) => + Show (CooArray r ix a) + where + show CooArray{..} = + "CooArray " + <> show values_ + <> " " + <> show coords_ + <> " " + <> show cooSize_ + +values :: CooArray r ix a -> Vector r a +values CooArray{values_} = values_ + +coords :: CooArray r ix a -> Vector r ix +coords CooArray{coords_} = coords_ + +cooSize :: CooArray r ix a -> Sz ix +cooSize CooArray{cooSize_} = cooSize_ + +-- | Make a 'CooArray' from a list of coordinate-value pairs. +mkCooArrayF :: + (Foldable f, Index ix, Manifest r a, Manifest r ix, MonadThrow m, Stream r Ix1 ix) => + Sz ix -> + f (ix, a) -> + m (CooArray r ix a) +mkCooArrayF cooSize_ coo + | isNull unsafeInds = return CooArray{..} + | otherwise = throwM $ IndexOutOfBoundsException cooSize_ (shead' unsafeInds) + where + arr = Massiv.fromList @B Par . toList $ coo + values_ = Massiv.compute . Massiv.map snd $ arr + coords_ = Massiv.compute . Massiv.map fst $ arr + unsafeInds = Massiv.sfilter (not . isSafeIndex cooSize_) coords_ + +-- | Make a 'CooArray' from a indices and values. +mkCooArray :: + (MonadThrow m, Index ix, Size r, Stream r Ix1 ix) => + Sz ix -> + Vector r ix -> + Vector r a -> + m (CooArray r ix a) +mkCooArray cooSize_ coords_ values_ + | Massiv.size coords_ /= Massiv.size values_ = throwM $ SizeMismatchException (Massiv.size coords_) (Massiv.size values_) + | not . isNull $ unsafeInds = throwM $ IndexOutOfBoundsException cooSize_ (shead' unsafeInds) + | otherwise = return CooArray{..} + where + unsafeInds = Massiv.sfilter (not . isSafeIndex cooSize_) coords_ diff --git a/src/TREXIO/HighLevel.hs b/src/TREXIO/HighLevel.hs new file mode 100644 --- /dev/null +++ b/src/TREXIO/HighLevel.hs @@ -0,0 +1,40 @@ +{-# LANGUAGE QuasiQuotes #-} +{-# LANGUAGE TemplateHaskell #-} + +{- | +Module: TREXIO.HighLevel +Description: Abstracted Haskell bindings to the TREXIO library +Copyright: Phillip Seeber 2024 +License: BSD-3-Clause +Maintainer: phillip.seeber@uni-jena.de +Stability: experimental +Portability: POSIX +-} +module TREXIO.HighLevel where + +import Control.Monad +import Data.Aeson +import Data.Map qualified as Map +import Data.Massiv.Array as Massiv hiding (Dim, forM) +import Foreign.C.ConstPtr +import Foreign.C.Types +import Language.Haskell.TH +import TREXIO.Internal.Base +import TREXIO.Internal.TH +import TREXIO.LowLevel +import TREXIO.LowLevel.Scheme + +$( do + let trexio@(TrexioScheme trexioScheme) = scheme + groups = Map.toList trexioScheme + + -- Generate abstracted Haskell wrappers around all C functions + fmap concat . forM groups $ \(groupName, Group group) -> do + deleteBind <- mkHsDeleteFn groupName + fieldBinds <- fmap concat . forM (Map.toList group) $ \(dataName, fieldType) -> do + hasBind <- mkHsHasFn groupName dataName fieldType + readBind <- mkHsReadFn groupName dataName fieldType + writeBind <- mkHsWriteFn trexio groupName dataName fieldType + return $ hasBind <> readBind <> writeBind + return $ deleteBind <> fieldBinds + ) diff --git a/src/TREXIO/HighLevel/Records.hs b/src/TREXIO/HighLevel/Records.hs new file mode 100644 --- /dev/null +++ b/src/TREXIO/HighLevel/Records.hs @@ -0,0 +1,31 @@ +{-# LANGUAGE QuasiQuotes #-} +{-# LANGUAGE TemplateHaskell #-} + +{- | +Module: TREXIO.HighLevel.Records +Description: Records representing the TREXIO scheme +Copyright: Phillip Seeber 2024 +License: BSD-3-Clause +Maintainer: phillip.seeber@uni-jena.de +Stability: experimental +Portability: POSIX +-} +module TREXIO.HighLevel.Records where + +import Data.Map qualified as Map +import Data.Massiv.Array as Massiv hiding (Dim, forM) +import TREXIO.Internal.TH +import TREXIO.LowLevel.Scheme + +$( do + let trexio@(TrexioScheme trexioScheme) = scheme + groups = Map.toList trexioScheme + + -- Generate records for all data groups + dataGroupRecords <- traverse (uncurry mkRecord) groups + + -- Generate the TREXIO super record with fields for all data groups + trexioSuperRecord <- mkTrexioScheme trexio + + return $ trexioSuperRecord : dataGroupRecords + ) diff --git a/src/TREXIO/Internal/Marshaller.hs b/src/TREXIO/Internal/Marshaller.hs new file mode 100644 --- /dev/null +++ b/src/TREXIO/Internal/Marshaller.hs @@ -0,0 +1,176 @@ +module TREXIO.Internal.Marshaller where + +import Data.Massiv.Array as Massiv hiding (withMArray) +import Data.Massiv.Array.Unsafe +import Data.Maybe (fromJust) +import Foreign + +{- | Passes a mutable 'MArray' to a C function. The array may be deallocated +on C side or modified in place. This funciton is super dangerous. You array may +suddenly be gone! +-} +withMArray :: (Index ix) => MArray s S ix a -> (Ptr b -> IO c) -> IO c +withMArray v f = do + let (arrFPtr, _arrL) = unsafeMArrayToForeignPtr v + withForeignPtr arrFPtr $ \arrPtr -> f (castPtr arrPtr) + +{- | Pass a 'Massiv.Array' to a C function. This function is safe, but the array +is copied to a new memory location. +-} +withArray :: (Storable a, Index ix) => Array S ix a -> (Ptr b -> IO c) -> IO c +withArray v f = do + mArr <- thaw v + withMArray mArr f + +-- | Get an 'MArray' from C memory. Haskell and C side use the same memory reference. +-- Be careful with this function. When the undelying pointer vanishes, the array is nonsense +unsafeToMArray :: (Index ix, Storable a) => Sz ix -> Ptr a -> IO (MArray s S ix a) +unsafeToMArray sz ptr = do + fPtr <- newForeignPtr_ ptr + let mArr = unsafeMArrayFromForeignPtr0 fPtr (toLinearSz sz) + resizeMArrayM sz mArr + +{- | Get an 'Array' from C memory. The underlying C array is copied and the +result is safe to use. +-} +peekArray :: (Index ix, Storable a) => Sz ix -> Ptr a -> IO (Array S ix a) +peekArray sz ptr = do + mArr <- unsafeToMArray sz ptr + freeze Par mArr + +-- | Peek an integer array from C memory and cast it into standard int types. +peekIntArray :: (Index ix, Manifest r Int) => Sz ix -> Ptr Int32 -> IO (Array r ix Int) +peekIntArray sz ptr = do + mArr <- unsafeToMArray sz (castPtr ptr) + compute . Massiv.map (fromIntegral :: Int32 -> Int) <$> freeze Par mArr + +-- | Peek 2D coordinates from C memory. +peek2DCoords :: (Manifest r Ix2) => Sz1 -> Ptr Int32 -> IO (Vector r Ix2) +peek2DCoords sz ptr = do + coordsSimple :: Matrix S Int <- peekIntArray (consSz sz 2) ptr + return . compute . Massiv.map mkIx2 . outerSlices $ coordsSimple + where + mkIx2 v = Ix2 (v ! 0) (v ! 1) + +-- | Peek 3D coordinates from C memory. +peek3DCoords :: (Manifest r Ix3) => Sz1 -> Ptr Int32 -> IO (Vector r Ix3) +peek3DCoords sz ptr = do + coordsSimple :: Matrix S Int <- peekIntArray (consSz sz 3) ptr + return . compute . Massiv.map mkIx3 . outerSlices $ coordsSimple + where + mkIx3 v = Ix3 (v ! 0) (v ! 1) (v ! 2) + +-- | Peek 4D coordinates from C memory. +peek4DCoords :: (Manifest r Ix4) => Sz1 -> Ptr Int32 -> IO (Vector r Ix4) +peek4DCoords sz ptr = do + coordsSimple :: Matrix S Int <- peekIntArray (consSz sz 4) ptr + return . compute . Massiv.map mkIx4 . outerSlices $ coordsSimple + where + mkIx4 v = Ix4 (v ! 0) (v ! 1) (v ! 2) (v ! 3) + +-- | Peek 5D coordinates from C memory. +peek6DCoords :: (Manifest r (IxN 6)) => Sz1 -> Ptr Int32 -> IO (Vector r (IxN 6)) +peek6DCoords sz ptr = do + coordsSimple :: Matrix S Int <- peekIntArray (consSz sz 6) ptr + return . compute . Massiv.map mkIx6 . outerSlices $ coordsSimple + where + mkIx6 v = (v ! 0) :> (v ! 1) :> (v ! 2) :> (v ! 3) :> (v ! 4) :. (v ! 5) + +peek8DCoords :: (Manifest r (IxN 8)) => Sz1 -> Ptr Int32 -> IO (Vector r (IxN 8)) +peek8DCoords sz ptr = do + coordsSimple :: Matrix S Int <- peekIntArray (consSz sz 8) ptr + return . compute . Massiv.map mkIx8 . outerSlices $ coordsSimple + where + mkIx8 v = (v ! 0) :> (v ! 1) :> (v ! 2) :> (v ! 3) :> (v ! 4) :> (v ! 5) :> (v ! 6) :. (v ! 7) + +-- | Convert a vector of 'Ix2' to a \( m \times 2 \) matrix. +castCoords2D :: (Manifest r1 Ix2, Manifest r2 Int32) => Vector r1 Ix2 -> Matrix r2 Int32 +castCoords2D ixVec = + compute + . fromJust + . stackOuterSlicesM + . Massiv.map ix2ToRow + $ ixVec + where + ix2ToRow :: Ix2 -> Vector S Int32 + ix2ToRow (Ix2 x y) = fromList Seq [fromIntegral x, fromIntegral y] + +-- | Convert a vector of 'Ix3' to a \( m \times 3 \) matrix. +castCoords3D :: (Manifest r1 Ix3, Manifest r2 Int32) => Vector r1 Ix3 -> Matrix r2 Int32 +castCoords3D ixVec = + compute + . fromJust + . stackOuterSlicesM + . Massiv.map ix3ToRow + $ ixVec + where + ix3ToRow :: Ix3 -> Vector S Int32 + ix3ToRow (Ix3 x y z) = + fromList + Seq + [ fromIntegral x + , fromIntegral y + , fromIntegral z + ] + +-- | Convert a vector of 'Ix4' to a \( m \times 4 \) matrix. +castCoords4D :: (Manifest r1 Ix4, Manifest r2 Int32) => Vector r1 Ix4 -> Matrix r2 Int32 +castCoords4D ixVec = + compute + . fromJust + . stackOuterSlicesM + . Massiv.map ix4ToRow + $ ixVec + where + ix4ToRow :: Ix4 -> Vector S Int32 + ix4ToRow (Ix4 x y z w) = + fromList + Seq + [ fromIntegral x + , fromIntegral y + , fromIntegral z + , fromIntegral w + ] + +-- | Convert a vector of 'Ix6' to a \( m \times 6 \) matrix. +castCoords6D :: (Manifest r1 (IxN 6), Manifest r2 Int32) => Vector r1 (IxN 6) -> Matrix r2 Int32 +castCoords6D ixVec = + compute + . fromJust + . stackOuterSlicesM + . Massiv.map ix6ToRow + $ ixVec + where + ix6ToRow :: IxN 6 -> Vector S Int32 + ix6ToRow (x :> y :> z :> w :> u :. v) = + fromList + Seq + [ fromIntegral x + , fromIntegral y + , fromIntegral z + , fromIntegral w + , fromIntegral u + , fromIntegral v + ] + +castCoords8D :: (Manifest r1 (IxN 8), Manifest r2 Int32) => Vector r1 (IxN 8) -> Matrix r2 Int32 +castCoords8D ixVec = + compute + . fromJust + . stackOuterSlicesM + . Massiv.map ix8ToRow + $ ixVec + where + ix8ToRow :: IxN 8 -> Vector S Int32 + ix8ToRow (x :> y :> z :> w :> u :> v :> t :. s) = + fromList + Seq + [ fromIntegral x + , fromIntegral y + , fromIntegral z + , fromIntegral w + , fromIntegral u + , fromIntegral v + , fromIntegral t + , fromIntegral s + ] diff --git a/src/TREXIO/Internal/TH.hs b/src/TREXIO/Internal/TH.hs new file mode 100644 --- /dev/null +++ b/src/TREXIO/Internal/TH.hs @@ -0,0 +1,1173 @@ +{-# LANGUAGE OverloadedLists #-} +{-# LANGUAGE QuasiQuotes #-} +{-# LANGUAGE TemplateHaskell #-} + +module TREXIO.Internal.TH where + +import Control.Exception.Safe +import Control.Monad +import Data.Aeson hiding (Success, withArray) +import Data.Bit.ThreadSafe (Bit) +import Data.Bit.ThreadSafe qualified as BV +import Data.ByteString qualified as BS +import Data.ByteString.Lazy qualified as BL +import Data.ByteString.Lazy.Char8 qualified as BLC +import Data.ByteString.Unsafe qualified as BS +import Data.Char +import Data.Coerce +import Data.List qualified as L +import Data.Map (Map) +import Data.Map qualified as Map +import Data.Massiv.Array as Massiv hiding (Dim, dropWhile, forM, forM_, mapM, product, replicate, takeWhile, throwM, toList, zip) +import Data.Massiv.Array qualified as Massiv +import Data.Massiv.Array.Manifest.Vector qualified as Massiv +import Data.Massiv.Array.Unsafe (unsafeWithPtr) +import Data.Maybe +import Data.Text (Text) +import Data.Text qualified as T +import Data.Text.IO qualified as T +import Data.Vector qualified as V +import Foreign hiding (peekArray, void, withArray) +import Foreign.C.ConstPtr +import Foreign.C.String +import Foreign.C.Types +import GHC.Generics (Generic) +import Language.Haskell.TH +import Language.Haskell.TH.Syntax (Lift (..)) +import System.IO +import System.IO.Temp +import System.Process.Typed +import TREXIO.CooArray +import TREXIO.Internal.Base +import TREXIO.Internal.Marshaller +import Text.Casing +import Text.Read (readMaybe) + +tshow :: (Show a) => a -> Text +tshow = T.pack . show + +-------------------------------------------------------------------------------- + +{- | Attempts to obtain the JSON specification from the trexio.h header. This is +a little bit arcane process: + +1. Write a temporary file @trexio.c@ that merely includes the header @#include @ +2. Run the C preprocessor on it using @gcc -E trexio.c@. Comments will include + the included header paths +3. Parse the output to find the header paths +4. From the extracted header path, get the JSON specification +-} +getJsonSpec :: (MonadIO m, MonadMask m) => m TrexioScheme +getJsonSpec = withSystemTempFile "trexio.c" $ \tmpPath tmpHandle -> do + -- Write the temporary file + liftIO $ do + T.hPutStrLn tmpHandle "#include " + hFlush tmpHandle + + -- Run the C preprocessor + (stdo, _) <- readProcess_ . shell $ "gcc -E " <> tmpPath + + -- Filter for trexio.h header paths + let trexioLines = + filter ("/trexio.h" `BL.isSuffixOf`) + . fmap (BLC.filter (/= '"') . BLC.dropWhileEnd (/= '"') . BLC.dropWhile (/= '"')) + . filter ("#" `BL.isPrefixOf`) + . BLC.lines + $ stdo + trexioPath <- case trexioLines of + t : _ -> liftIO . BS.toFilePath . BS.toStrict $ t + _ -> throwString "Could not find trexio.h header path" + + -- Get the JSON specification from the header + trexioHeader <- liftIO $ BL.readFile trexioPath + let jsonString = + BLC.unlines + . L.drop 1 + . takeWhile (/= "*/") + . dropWhile (/= "/* JSON configuration") + . BLC.lines + $ trexioHeader + case eitherDecode jsonString of + Right trexio -> return trexio + Left err -> throwString $ "Could not parse JSON specification: " <> err + +-------------------------------------------------------------------------------- + +{- | The overall data structure TREXIO uses to represent a wave function as a +JSON specification. A TREXIO scheme consists of multiple data groups and each +data group has multiple fields. A field may require knowledge of other fields. +-} +newtype TrexioScheme = TrexioScheme (Map GroupName Group) + deriving (Generic, Show, Eq, Ord, Lift) + deriving (ToJSON, FromJSON) via Map GroupName Group + +{- | The name of a data group, e.g. @ao@ for atomic orbitals, @basis@ for basis +functions, etc. +-} +newtype GroupName = GroupName Text + deriving (Generic, Show, Eq, Ord, Lift) + deriving (ToJSONKey, FromJSONKey) via Text + +{- | A data group is a record like data structure with named fields of different +types. Each field may or may not be set, thus the 'Maybe' type. +-} +newtype Group = Group (Map DataName Typ) + deriving (Generic, Show, Eq, Ord, Lift) + deriving (ToJSON, FromJSON) via Map DataName Typ + +{- | The name of a data field, as specified by the TREXIO scheme. There is no +guarantee that the name is a valid Haskell identifier. To ensure that, use the +'sanId' function. +-} +newtype DataName = DataName Text + deriving (Generic, Show, Eq, Ord, Lift) + deriving (ToJSONKey, FromJSONKey) via Text + +instance ToJSON DataName where + toJSON (DataName name) = String name + +instance FromJSON DataName where + parseJSON (String name) = return . DataName $ name + parseJSON _ = fail "parseJSON(DataName): could not parse" + +{- | The TREXIO type of a data field including sparsity, buffering, dimensionality +etc. +-} +data Typ + = -- | A 32 integer but meant to represent the size in a given dimension. The + -- Bool indicates if field can also be written + Dim Bool Length + | -- | A 32 bit integer + Int Length + | -- | A double precision float. The Bool indicates whether this field is + -- buffered + Float Bool Length + | -- | A string with a given length + Str Length + | -- | An index type + Idx Length + | -- | Sparse array of floats + SparseFloat Length + | -- | A bit field + BitField Length + deriving (Generic, Show, Eq, Ord, Lift) + +instance ToJSON Typ where + toJSON (Dim False len) = Array ["dim", toJSON len] + toJSON (Dim True len) = Array ["dim readonly", toJSON len] + toJSON (Int len) = Array ["int", toJSON len] + toJSON (Float False len) = Array ["float", toJSON len] + toJSON (Float True len) = Array ["float buffered", toJSON len] + toJSON (Str len) = Array ["str", toJSON len] + toJSON (Idx len) = Array ["index", toJSON len] + toJSON (SparseFloat len) = Array ["float sparse", toJSON len] + toJSON (BitField len) = Array ["int special", toJSON len] + +instance FromJSON Typ where + parseJSON (Array ["dim", len]) = Dim True <$> parseJSON len + parseJSON (Array ["dim readonly", len]) = Dim False <$> parseJSON len + parseJSON (Array ["int", len]) = Int <$> parseJSON len + parseJSON (Array ["float", len]) = Float False <$> parseJSON len + parseJSON (Array ["float buffered", len]) = Float True <$> parseJSON len + parseJSON (Array ["str", len]) = Str <$> parseJSON len + parseJSON (Array ["index", len]) = Idx <$> parseJSON len + parseJSON (Array ["float sparse", len]) = SparseFloat <$> parseJSON len + parseJSON (Array ["int special", len]) = BitField <$> parseJSON len + parseJSON _ = fail "parseJSON(Typ): could not parse" + +{- | TREXIO data fields are annotated with a length specification. This +specification is a list of sizes along the dimensions of an $n$D array. +An empty length specification refers to a scalar. A dimension may have a +constant size or refer to another field that stores its size, see 'DimLength'. +-} +newtype Length = Length [DimLength] deriving (Generic, Show, Eq, Ord, Lift) + +instance ToJSON Length where + toJSON (Length dim) = Array . V.fromList . fmap toJSON $ dim + +instance FromJSON Length where + parseJSON (Array arr) = + Length . V.toList + <$> traverse (parseJSON @DimLength) arr + parseJSON _ = fail "parseJSON(Length): could not parse" + +{- | The size along a dimension of a field. It can be a constant or refer to +a field that stores a scalar describing a length. +-} +data DimLength + = Const Int + | Field GroupName DataName + deriving (Generic, Show, Eq, Ord, Lift) + +instance ToJSON DimLength where + toJSON (Const int) = String . tshow $ int + toJSON (Field (GroupName groupName) (DataName dataName)) = String $ groupName <> "." <> dataName + +instance FromJSON DimLength where + parseJSON (String s) = case readMaybe . T.unpack $ s of + Just i -> return . Const $ i + Nothing -> case T.splitOn "." s of + [groupName, dataName] -> return $ Field (GroupName groupName) (DataName dataName) + _ -> fail "parseJSON(DimLength): could not parse" + parseJSON _ = fail "parseJSON(DimLength): could not parse" + +-------------------------------------------------------------------------------- +-- Helper functions + +{- | Sanitise an identifier, e.g. a field name or function name. I.e. we ensure +it starts with a valid lower case letter or symbol. +-} +sanId :: String -> String +sanId "" = error "sanId: empty string" +sanId ind@(c : cs) + | isUpperCase c = sanId $ toLower c : cs + | isDigit c = sanId $ '_' : c : cs + | ind == "type" = "type'" + | ind == "class" = "class'" + | otherwise = c : cs + +-------------------------------------------------------------------------------- +-- Template Haskell binding generator + +-- | The standard operations on data fields. +data FieldOps + = -- | Check if a field is set + Has + | -- | Read data from a field + Read + | -- | Write data to a field + Write + deriving (Generic, Eq, Show, Ord) + +opsFnName :: FieldOps -> String +opsFnName Has = "has" +opsFnName Read = "read" +opsFnName Write = "write" + +-- | Associate a TREXIO 'Typ' with a Haskell 'Type'. +typToType :: (Quote m) => Typ -> m Type +typToType (Dim _ (Length [])) = [t|Int|] +typToType (Dim _ (Length [_])) = [t|Vector S Int|] +typToType (Int (Length [])) = [t|Int|] +typToType (Int (Length [_])) = [t|Vector S Int|] +typToType (Float False (Length [])) = [t|Double|] +typToType (Float False (Length [_])) = [t|Vector S Double|] +typToType (Float False (Length [_, _])) = [t|Matrix S Double|] +typToType (Float False (Length [_, _, _])) = [t|Massiv.Array S Ix3 Double|] +typToType (Float False (Length [_, _, _, _])) = [t|Massiv.Array S Ix4 Double|] +typToType (Float True (Length [_])) = [t|Vector S Double|] +typToType (Str (Length [])) = [t|Text|] +typToType (Str (Length [_])) = [t|Vector B Text|] +typToType (Idx (Length [])) = [t|Int|] +typToType (Idx (Length [_])) = [t|Vector S Int|] +typToType (SparseFloat (Length [_, _])) = [t|CooArray U Ix2 Double|] +typToType (SparseFloat (Length [_, _, _])) = [t|CooArray U Ix3 Double|] +typToType (SparseFloat (Length [_, _, _, _])) = [t|CooArray U Ix4 Double|] +typToType (SparseFloat (Length [_, _, _, _, _, _])) = [t|CooArray U (IxN 6) Double|] +typToType (SparseFloat (Length [_, _, _, _, _, _, _, _])) = [t|CooArray U (IxN 8) Double|] +typToType (BitField (Length [_])) = [t|BV.Vector Word8|] +typToType t = error $ "Can not associate " <> show t <> " with a Type" + +-- | Associate a 'FieldOps' and a TREXIO field 'Typ' with a Haskell function type. +mkCFnSig :: FieldOps -> Typ -> Q Type +mkCFnSig Has _ = [t|Trexio -> IO ExitCodeC|] +mkCFnSig Read (Dim _ _) = [t|Trexio -> Ptr Int32 -> IO ExitCodeC|] +mkCFnSig Read (Int _) = [t|Trexio -> Ptr Int32 -> IO ExitCodeC|] +mkCFnSig Read (Float False _) = [t|Trexio -> Ptr CDouble -> IO ExitCodeC|] +mkCFnSig Read (Float True _) = [t|Trexio -> Int64 -> Ptr Int64 -> Ptr CDouble -> IO ExitCodeC|] +mkCFnSig Read (Str (Length [])) = [t|Trexio -> Ptr CChar -> Int32 -> IO ExitCodeC|] +mkCFnSig Read (Str (Length [_])) = [t|Trexio -> Ptr (Ptr CChar) -> Int32 -> IO ExitCodeC|] +mkCFnSig Read (Idx _) = [t|Trexio -> Ptr Int32 -> IO ExitCodeC|] +mkCFnSig Read (SparseFloat _) = [t|Trexio -> Int64 -> Ptr Int64 -> Ptr Int32 -> Ptr CDouble -> IO ExitCodeC|] +mkCFnSig Read (BitField _) = [t|Trexio -> Int64 -> Ptr Int64 -> Ptr Int64 -> IO ExitCodeC|] +mkCFnSig Write (Dim _ (Length [])) = [t|Trexio -> Int32 -> IO ExitCodeC|] +mkCFnSig Write (Dim _ (Length [_])) = [t|Trexio -> Ptr Int32 -> IO ExitCodeC|] +mkCFnSig Write (Int (Length [])) = [t|Trexio -> Int32 -> IO ExitCodeC|] +mkCFnSig Write (Int (Length [_])) = [t|Trexio -> Ptr Int32 -> IO ExitCodeC|] +mkCFnSig Write (Float False (Length [])) = [t|Trexio -> CDouble -> IO ExitCodeC|] +mkCFnSig Write (Float False (Length _)) = [t|Trexio -> Ptr CDouble -> IO ExitCodeC|] +mkCFnSig Write (Float True (Length _)) = [t|Trexio -> Int64 -> Int64 -> Ptr CDouble -> IO ExitCodeC|] +mkCFnSig Write (Str (Length [])) = [t|Trexio -> ConstPtr CChar -> Int32 -> IO ExitCodeC|] +mkCFnSig Write (Str (Length [_])) = [t|Trexio -> ConstPtr (ConstPtr CChar) -> Int32 -> IO ExitCodeC|] +mkCFnSig Write (Idx (Length [])) = [t|Trexio -> Int32 -> IO ExitCodeC|] +mkCFnSig Write (Idx (Length [_])) = [t|Trexio -> Ptr Int32 -> IO ExitCodeC|] +mkCFnSig Write (SparseFloat _) = [t|Trexio -> Int64 -> Int64 -> Ptr Int32 -> Ptr CDouble -> IO ExitCodeC|] +mkCFnSig Write (BitField _) = [t|Trexio -> Int64 -> Int64 -> Ptr Int64 -> IO ExitCodeC|] +mkCFnSig op t = error $ "Can not associate " <> show op <> " and " <> show t <> " with a Type" + +{- | Associate a 'FieldOps' and a field 'Typ' with the type of a Haskell function. +The Haskell function is already abstracted and expected to perform other queries +such as vector sizes as necessary. +-} +mkHsFnSig :: FieldOps -> Typ -> Q Type +mkHsFnSig Has _ = [t|forall m. (MonadIO m) => Trexio -> m Bool|] +mkHsFnSig Read (Dim _ (Length [])) = [t|forall m. (MonadIO m) => Trexio -> m Int|] +mkHsFnSig Read (Dim _ (Length [_])) = [t|forall m. (MonadIO m) => Trexio -> m (Vector S Int)|] +mkHsFnSig Read (Int (Length [])) = [t|forall m. (MonadIO m) => Trexio -> m Int|] +mkHsFnSig Read (Int (Length [_])) = [t|forall m. (MonadIO m) => Trexio -> m (Vector S Int)|] +mkHsFnSig Read (Float False (Length [])) = [t|forall m. (MonadIO m) => Trexio -> m Double|] +mkHsFnSig Read (Float False (Length [_])) = [t|forall m. (MonadIO m) => Trexio -> m (Vector S Double)|] +mkHsFnSig Read (Float False (Length [_, _])) = [t|forall m. (MonadIO m) => Trexio -> m (Matrix S Double)|] +mkHsFnSig Read (Float False (Length [_, _, _])) = [t|forall m. (MonadIO m) => Trexio -> m (Massiv.Array S Ix3 Double)|] +mkHsFnSig Read (Float False (Length [_, _, _, _])) = [t|forall m. (MonadIO m) => Trexio -> m (Massiv.Array S Ix4 Double)|] +mkHsFnSig Read (Float True (Length [_])) = [t|forall m. (MonadIO m) => Trexio -> m (Vector S Double)|] +mkHsFnSig Read (Str (Length [])) = [t|forall m. (MonadIO m) => Trexio -> m Text|] +mkHsFnSig Read (Str (Length [_])) = [t|forall m. (MonadIO m) => Trexio -> m (Vector B Text)|] +mkHsFnSig Read (Idx (Length [])) = [t|forall m. (MonadIO m) => Trexio -> m Int|] +mkHsFnSig Read (Idx (Length [_])) = [t|forall m. (MonadIO m) => Trexio -> m (Vector S Int)|] +mkHsFnSig Read (SparseFloat (Length [_, _])) = [t|forall m. (MonadIO m) => Trexio -> m (CooArray U Ix2 Double)|] +mkHsFnSig Read (SparseFloat (Length [_, _, _])) = [t|forall m. (MonadIO m) => Trexio -> m (CooArray U Ix3 Double)|] +mkHsFnSig Read (SparseFloat (Length [_, _, _, _])) = [t|forall m. (MonadIO m) => Trexio -> m (CooArray U Ix4 Double)|] +mkHsFnSig Read (SparseFloat (Length [_, _, _, _, _, _])) = [t|forall m. (MonadIO m) => Trexio -> m (CooArray U (IxN 6) Double)|] +mkHsFnSig Read (SparseFloat (Length [_, _, _, _, _, _, _, _])) = [t|forall m. (MonadIO m) => Trexio -> m (CooArray U (IxN 8) Double)|] +mkHsFnSig Read (BitField (Length [_])) = [t|forall m. (MonadIO m) => Trexio -> m (Matrix U (Bit, Bit))|] +mkHsFnSig Write (Dim _ (Length [])) = [t|forall m. (MonadIO m) => Trexio -> Int -> m ()|] +mkHsFnSig Write (Dim _ (Length [_])) = [t|forall m. (MonadIO m) => Trexio -> Vector S Int -> m ()|] +mkHsFnSig Write (Int (Length [])) = [t|forall m. (MonadIO m) => Trexio -> Int -> m ()|] +mkHsFnSig Write (Int (Length [_])) = [t|forall m. (MonadIO m) => Trexio -> Vector S Int -> m ()|] +mkHsFnSig Write (Float False (Length [])) = [t|forall m. (MonadIO m) => Trexio -> Double -> m ()|] +mkHsFnSig Write (Float False (Length [_])) = [t|forall m. (MonadIO m) => Trexio -> Vector S Double -> m ()|] +mkHsFnSig Write (Float False (Length [_, _])) = [t|forall m. (MonadIO m) => Trexio -> Matrix S Double -> m ()|] +mkHsFnSig Write (Float False (Length [_, _, _])) = [t|forall m. (MonadIO m) => Trexio -> Massiv.Array S Ix3 Double -> m ()|] +mkHsFnSig Write (Float False (Length [_, _, _, _])) = [t|forall m. (MonadIO m) => Trexio -> Massiv.Array S Ix4 Double -> m ()|] +mkHsFnSig Write (Float True (Length [_])) = [t|forall m. (MonadIO m) => Trexio -> Vector S Double -> m ()|] +mkHsFnSig Write (Str (Length [])) = [t|forall m. (MonadIO m) => Trexio -> Text -> m ()|] +mkHsFnSig Write (Str (Length [_])) = [t|forall m. (MonadIO m) => Trexio -> Vector B Text -> m ()|] +mkHsFnSig Write (Idx (Length [])) = [t|forall m. (MonadIO m) => Trexio -> Int -> m ()|] +mkHsFnSig Write (Idx (Length [_])) = [t|forall m. (MonadIO m) => Trexio -> Vector S Int -> m ()|] +mkHsFnSig Write (SparseFloat (Length [_, _])) = [t|forall m. (MonadIO m) => Trexio -> CooArray U Ix2 Double -> m ()|] +mkHsFnSig Write (SparseFloat (Length [_, _, _])) = [t|forall m. (MonadIO m) => Trexio -> CooArray U Ix3 Double -> m ()|] +mkHsFnSig Write (SparseFloat (Length [_, _, _, _])) = [t|forall m. (MonadIO m) => Trexio -> CooArray U Ix4 Double -> m ()|] +mkHsFnSig Write (SparseFloat (Length [_, _, _, _, _, _])) = [t|forall m. (MonadIO m) => Trexio -> CooArray U (IxN 6) Double -> m ()|] +mkHsFnSig Write (SparseFloat (Length [_, _, _, _, _, _, _, _])) = [t|forall m. (MonadIO m) => Trexio -> CooArray U (IxN 8) Double -> m ()|] +mkHsFnSig Write (BitField (Length [_])) = [t|forall m. (MonadIO m) => Trexio -> Matrix U (Bit, Bit) -> m ()|] +mkHsFnSig op t = error $ "Can not associate " <> show op <> " and " <> show t <> " with a Type" + +-- | Generate a Haskell function name for a given operation, group and field of that group. +mkHsFnName :: FieldOps -> GroupName -> DataName -> String +mkHsFnName op (GroupName groupName) (DataName dataName) = + sanId . camel $ opsFnName op <> "_" <> T.unpack groupName <> "_" <> T.unpack dataName + +-- | Generate a C function name for a given operation, group and field of that group. +mkCFnName :: FieldOps -> GroupName -> DataName -> String +mkCFnName op (GroupName groupName) (DataName dataName) = + "trexio_" <> opsFnName op <> "_" <> (T.unpack . T.toLower $ groupName) <> "_" <> (T.unpack . T.toLower $ dataName) + +-- | Convert a field to a type +fieldToType :: (Quote m) => DataName -> Typ -> m VarBangType +fieldToType (DataName dataName) typ = do + let fieldName = mkName . sanId . camel . T.unpack $ dataName + fieldType <- typToType typ + maybeFieldType <- [t|Maybe $(return fieldType)|] + return (fieldName, Bang NoSourceUnpackedness NoSourceStrictness, maybeFieldType) + +stdDerivs :: [DerivClause] +stdDerivs = [DerivClause Nothing [ConT ''Generic, ConT ''Show, ConT ''Ord, ConT ''Eq]] + +-- | Create a record from a given data group +mkRecord :: GroupName -> Group -> Q Dec +mkRecord (GroupName groupName) (Group fields) = do + groupNameTD <- newName . pascal . T.unpack $ groupName + groupNameTC <- newName . pascal . T.unpack $ groupName + fieldsT <- traverse (uncurry fieldToType) . Map.toList $ fields + return $ DataD [] groupNameTD [] Nothing [RecC groupNameTC fieldsT] stdDerivs + +-- | Create the TREXIO scheme type with subrecords for each data group. +mkTrexioScheme :: TrexioScheme -> Q Dec +mkTrexioScheme (TrexioScheme groups) = do + dataName <- newName "TREXIO" + constructorName <- newName "TREXIO" + fieldsT <- forM (Map.toList groups) $ \(GroupName groupName, _) -> do + groupFieldName <- newName . camel . T.unpack $ groupName + groupFieldType <- [t|$(conT . mkName . pascal . T.unpack $ groupName)|] + return (groupFieldName, Bang NoSourceUnpackedness NoSourceStrictness, groupFieldType) + return $ DataD [] dataName [] Nothing [RecC constructorName fieldsT] stdDerivs + +-- | Create all C function bindings for a given group +mkCBindings :: GroupName -> Group -> Q [Dec] +mkCBindings groupName (Group fields) = do + -- Group bindings for delete + groupDelBind <- mkCDeleteFn groupName + fieldBinds <- fmap (catMaybes . concat) . forM (Map.toList fields) $ \(fieldName, fieldTyp) -> do + -- Standard bindings + stdBindings <- forM [Has, Read, Write] $ \op -> do + let cFnName = mkCFnName op groupName fieldName + cFnNameT <- newName cFnName + cFnSig <- mkCFnSig op fieldTyp + + -- Dim fields, that are read only, do not have a write function + if fieldTyp == Dim False (Length []) && op == Write + then return Nothing + else return . Just . ForeignD $ ImportF CApi Unsafe ("trexio.h " <> cFnName) cFnNameT cFnSig + + -- "size" bindings: bitfields, sparse arrays and buffered arrays have an + -- additional function "_size", that tells how many COO elements there are. + let cSizeFnString = mkCSizeFnName groupName fieldName + cSizeFnName <- newName cSizeFnString + cFnSig <- [t|Trexio -> Ptr Int64 -> IO Int32|] + let imprt = + ForeignD $ + ImportF CApi Unsafe ("trexio.h " <> cSizeFnString) cSizeFnName cFnSig + let sizeBinding = case fieldTyp of + SparseFloat _ -> Just imprt + Float True _ -> Just imprt + _ -> Nothing + + -- Return all bindings + return $ stdBindings <> [sizeBinding] + return $ groupDelBind : fieldBinds + +-- Make a Has function for a given field +mkHsHasFn :: GroupName -> DataName -> Typ -> Q [Dec] +mkHsHasFn groupName dataName fieldTyp = do + let hsFnName = mkHsFnName Has groupName dataName + cFnName = mkCFnName Has groupName dataName + hsFnSig <- mkHsFnSig Has fieldTyp + hsExp <- + [e| + \trexio -> liftIO $ do + cRes <- $(varE . mkName $ cFnName) trexio + if exitCodeH cRes == Success + then return True + else return False + |] + return + [ SigD (mkName hsFnName) hsFnSig + , FunD (mkName hsFnName) [Clause [] (NormalB hsExp) []] + ] + +{- | Generate an expression to obtain the size of an array field along a given +dimension. +-} +mkSizeFn :: DimLength -> Q Exp +mkSizeFn (Const i) = [e|\_ -> return i|] +mkSizeFn (Field groupName dataName) = do + let cFnName = mkCFnName Read groupName dataName + [e| + ( \trexio -> alloca $ \(dimPtr :: Ptr Int32) -> do + ec <- exitCodeH <$> $(varE . mkName $ cFnName) trexio dimPtr + case ec of + Success -> fromIntegral <$> peek dimPtr + _ -> throwM ec + ) + |] + +isIntField :: Typ -> Bool +isIntField (Dim _ _) = True +isIntField (Int _) = True +isIntField (Idx _) = True +isIntField _ = False + +isWritableIntField :: Typ -> Bool +isWritableIntField (Dim True _) = True +isWritableIntField (Int _) = True +isWritableIntField (Idx _) = True +isWritableIntField _ = False + +isProtectedIntField :: Typ -> Bool +isProtectedIntField (Dim False _) = True +isProtectedIntField _ = False + +isFloatField :: Typ -> Bool +isFloatField (Float False _) = True +isFloatField _ = False + +isBufferedFloat :: Typ -> Bool +isBufferedFloat (Float True _) = True +isBufferedFloat _ = False + +isSparseFloat :: Typ -> Bool +isSparseFloat (SparseFloat _) = True +isSparseFloat _ = False + +isStringField :: Typ -> Bool +isStringField (Str _) = True +isStringField _ = False + +isBitField :: Typ -> Bool +isBitField (BitField _) = True +isBitField _ = False + +{- | Sparse fields have an associated @_size@ function, that returns the number +of COO elements. +-} +mkCSizeFnName :: GroupName -> DataName -> String +mkCSizeFnName (GroupName groupName) (DataName dataName) = + T.unpack $ + "trexio_read_" <> groupName <> "_" <> dataName <> "_size" + +{- | Create abstracted read functions, that automatically obtain sizes for arrays +as required from other fields. 'CooArray's are read in a single, big chunk and +need to fit in memory. +-} +mkReadFns :: GroupName -> DataName -> Typ -> Q Exp +mkReadFns groupName dataName fieldType = case dims of + [] + | isIntField fieldType -> + [e| + \trexio -> liftIO . alloca $ \buf -> do + ec <- exitCodeH <$> $(varE . mkName $ mkCFnName Read groupName dataName) trexio buf + case ec of + Success -> fromIntegral <$> peek buf + _ -> throwM ec + |] + | isFloatField fieldType -> + [e| + \trexio -> liftIO . alloca $ \buf -> do + ec <- exitCodeH <$> $(varE . mkName $ mkCFnName Read groupName dataName) trexio buf + case ec of + Success -> peek (castPtr buf) + _ -> throwM ec + |] + | isStringField fieldType -> + [e| + \trexio -> liftIO . allocaBytes 256 $ \strPtr -> do + ec <- exitCodeH <$> $(varE . mkName $ mkCFnName Read groupName dataName) trexio strPtr 256 + case ec of + Success -> T.pack <$> peekCString strPtr + _ -> throwM ec + |] + | otherwise -> error $ "mkReadFns: unsupported field type for 0D data: " <> show fieldType + [d1] + | isIntField fieldType -> + [e| + \trexio -> liftIO $ do + sz1 <- $(mkSizeFn d1) trexio + allocaArray sz1 $ \buf -> do + ec <- exitCodeH <$> $(varE . mkName $ mkCFnName Read groupName dataName) trexio buf + case ec of + Success -> peekIntArray (Sz1 sz1) buf + _ -> throwM ec + |] + | isFloatField fieldType -> + [e| + \trexio -> liftIO $ do + sz1 <- $(mkSizeFn d1) trexio + allocaArray sz1 $ \buf -> do + ec <- exitCodeH <$> $(varE . mkName $ mkCFnName Read groupName dataName) trexio buf + case ec of + Success -> peekArray (Sz1 sz1) (castPtr buf) + _ -> throwM ec + |] + | isStringField fieldType -> + [e| + \trexio -> liftIO $ do + nStrings <- $(mkSizeFn d1) trexio + let maxStrLen = 256 + + allocaArray nStrings $ \(superPtr :: Ptr (Ptr CChar)) -> + -- Allocate the buffers for the strings + bracket + (replicateM nStrings $ callocArray0 maxStrLen) + (traverse free) + $ \(strPtrs :: [Ptr CChar]) -> do + -- Write the individual buffers to the super buffer + forM_ (zip [0 ..] strPtrs) $ \(i, strPtr) -> + pokeElemOff superPtr i strPtr + + -- Call the C function + ec <- exitCodeH <$> $(varE . mkName $ mkCFnName Read groupName dataName) trexio superPtr (fromIntegral maxStrLen) + case ec of + Success -> Massiv.fromList Seq . fmap T.pack <$> traverse peekCString strPtrs + _ -> throwM ec + |] + | isBitField fieldType -> + [e| + \trexio -> liftIO $ do + moNum <- $(mkSizeFn $ Field (GroupName "mo") (DataName "num")) trexio + nDets <- $(mkSizeFn d1) trexio + nInt64PerDet <- intsPerDet trexio + + -- Allocate a buffer + allocaArray (nDets * nInt64PerDet * 2) $ \detBuf -> do + let readDets :: IO (Matrix U (Bit, Bit)) + readDets = do + -- Read each determinant individually + dets <- forM [0 .. nDets - 1] $ \i -> do + let upPtr = detBuf `plusPtr` (i * nInt64PerDet * 2 * sizeOf (undefined :: Int64)) + downPtr = upPtr `plusPtr` (nInt64PerDet * sizeOf (undefined :: Int64)) + nBytes = nInt64PerDet * sizeOf (undefined :: Int64) + + upBS <- BS.unsafePackCStringLen (castPtr upPtr, nBytes) + downBS <- BS.unsafePackCStringLen (castPtr downPtr, nBytes) + + let toDet = + compute @U + . Massiv.take moNum + . (Massiv.fromVector' Par (Sz $ nBytes * 8) :: BV.Vector Bit -> Vector U Bit) + . BV.cloneFromByteString + upDet = toDet upBS + downDet = toDet downBS + + return $ Massiv.zip upDet downDet + + compute <$> stackOuterSlicesM dets + + -- Call the C function and populate the buffer + with (fromIntegral nDets) $ \bufSz -> do + ec <- + exitCodeH + <$> $(varE . mkName $ mkCFnName Read groupName dataName) + trexio + 0 + bufSz + detBuf + case ec of + Success -> readDets + End -> readDets + _ -> throwM ec + |] + | isBufferedFloat fieldType -> + [e| + \trexio -> liftIO $ do + sz1 <- $(mkSizeFn d1) trexio + with (fromIntegral sz1) $ \bufSz -> + allocaArray sz1 $ \buf -> do + ec <- exitCodeH <$> $(varE . mkName $ mkCFnName Read groupName dataName) trexio 0 bufSz buf + case ec of + Success -> peekArray (Sz1 sz1) (castPtr buf) + End -> peekArray (Sz1 sz1) (castPtr buf) + _ -> throwM ec + |] + | otherwise -> error $ "mkReadFns: unsupported field type for 1D data: " <> show fieldType + [d1, d2] + | isFloatField fieldType -> + [e| + \trexio -> liftIO $ do + sz1 <- $(mkSizeFn d1) trexio + sz2 <- $(mkSizeFn d2) trexio + allocaArray (sz1 * sz2) $ \buf -> do + ec <- exitCodeH <$> $(varE . mkName $ mkCFnName Read groupName dataName) trexio buf + case ec of + Success -> peekArray (Sz2 sz1 sz2) (castPtr buf) + _ -> throwM ec + |] + | isSparseFloat fieldType -> + [e| + \trexio -> liftIO $ do + -- Size of the array + sz1 <- $(mkSizeFn d1) trexio + sz2 <- $(mkSizeFn d2) trexio + + -- Number of COO elements in the sparse array + nCoo <- alloca $ \buf -> do + ec <- exitCodeH <$> $(varE . mkName $ mkCSizeFnName groupName dataName) trexio buf + case ec of + Success -> fromIntegral <$> peek buf + _ -> throwM ec + + -- Read the COO array in a single chunk + with (fromIntegral nCoo) $ \bufSz -> + allocaArray (nCoo * 2) $ \ixBuf -> + allocaArray nCoo $ \valBuf -> do + ec <- exitCodeH <$> $(varE . mkName $ mkCFnName Read groupName dataName) trexio 0 bufSz ixBuf valBuf + case ec of + Success -> do + ixs <- peek2DCoords (Sz1 nCoo) ixBuf + vals <- peekArray (Sz1 nCoo) . castPtr $ valBuf + mkCooArray (Sz2 sz1 sz2) ixs . compute @U $ vals + _ -> throwM ec + |] + | otherwise -> error $ "mkReadFns: unsupported field type for 2D data: " <> show fieldType + [d1, d2, d3] + | isFloatField fieldType -> + [e| + \trexio -> liftIO $ do + sz1 <- $(mkSizeFn d1) trexio + sz2 <- $(mkSizeFn d2) trexio + sz3 <- $(mkSizeFn d3) trexio + allocaArray (sz1 * sz2 * sz3) $ \buf -> do + ec <- exitCodeH <$> $(varE . mkName $ mkCFnName Read groupName dataName) trexio buf + case ec of + Success -> peekArray (Sz3 sz1 sz2 sz3) (castPtr buf) + _ -> throwM ec + |] + | isSparseFloat fieldType -> + [e| + \trexio -> liftIO $ do + -- Size of the array + sz1 <- $(mkSizeFn d1) trexio + sz2 <- $(mkSizeFn d2) trexio + sz3 <- $(mkSizeFn d3) trexio + + -- Number of COO elements in the sparse array + nCoo <- alloca $ \buf -> do + ec <- exitCodeH <$> $(varE . mkName $ mkCSizeFnName groupName dataName) trexio buf + case ec of + Success -> fromIntegral <$> peek buf + _ -> throwM ec + + -- Read the COO array in a single chunk + with (fromIntegral nCoo) $ \bufSz -> + allocaArray (nCoo * 3) $ \ixBuf -> + allocaArray nCoo $ \valBuf -> do + ec <- exitCodeH <$> $(varE . mkName $ mkCFnName Read groupName dataName) trexio 0 bufSz ixBuf valBuf + case ec of + Success -> do + ixs <- peek3DCoords (Sz1 nCoo) ixBuf + vals <- peekArray (Sz1 nCoo) . castPtr $ valBuf + mkCooArray (Sz3 sz1 sz2 sz3) ixs . compute @U $ vals + _ -> throwM ec + |] + | otherwise -> error $ "mkReadFns: unsupported field type for 3D data: " <> show fieldType + [d1, d2, d3, d4] + | isFloatField fieldType -> + [e| + \trexio -> + liftIO $ do + sz1 <- $(mkSizeFn d1) trexio + sz2 <- $(mkSizeFn d2) trexio + sz3 <- $(mkSizeFn d3) trexio + sz4 <- $(mkSizeFn d4) trexio + allocaArray (sz1 * sz2 * sz3 * sz4) $ \buf -> do + ec <- exitCodeH <$> $(varE . mkName $ mkCFnName Read groupName dataName) trexio buf + case ec of + Success -> peekArray (Sz4 sz1 sz2 sz3 sz4) (castPtr buf) + _ -> throwM ec + |] + | isSparseFloat fieldType -> + [e| + \trexio -> liftIO $ do + -- Size of the array + sz1 <- $(mkSizeFn d1) trexio + sz2 <- $(mkSizeFn d2) trexio + sz3 <- $(mkSizeFn d3) trexio + sz4 <- $(mkSizeFn d4) trexio + + -- Number of COO elements in the sparse array + nCoo <- alloca $ \buf -> do + ec <- exitCodeH <$> $(varE . mkName $ mkCSizeFnName groupName dataName) trexio buf + case ec of + Success -> fromIntegral <$> peek buf + _ -> throwM ec + + -- Read the COO array in a single chunk + with (fromIntegral nCoo) $ \bufSz -> + allocaArray (nCoo * 4) $ \ixBuf -> + allocaArray nCoo $ \valBuf -> do + ec <- exitCodeH <$> $(varE . mkName $ mkCFnName Read groupName dataName) trexio 0 bufSz ixBuf valBuf + case ec of + Success -> do + ixs <- peek4DCoords (Sz1 nCoo) ixBuf + vals <- peekArray (Sz1 nCoo) . castPtr $ valBuf + mkCooArray (Sz4 sz1 sz2 sz3 sz4) ixs . compute @U $ vals + _ -> throwM ec + |] + | otherwise -> error $ "mkReadFns: unsupported field type for 4D data: " <> show fieldType + [d1, d2, d3, d4, d5, d6] + | isSparseFloat fieldType -> + [e| + \trexio -> liftIO $ do + -- Size of the array + sz1 <- $(mkSizeFn d1) trexio + sz2 <- $(mkSizeFn d2) trexio + sz3 <- $(mkSizeFn d3) trexio + sz4 <- $(mkSizeFn d4) trexio + sz5 <- $(mkSizeFn d5) trexio + sz6 <- $(mkSizeFn d6) trexio + + -- Number of COO elements in the sparse array + nCoo <- alloca $ \buf -> do + ec <- exitCodeH <$> $(varE . mkName $ mkCSizeFnName groupName dataName) trexio buf + case ec of + Success -> fromIntegral <$> peek buf + _ -> throwM ec + + -- Read the COO array in a single chunk + with (fromIntegral nCoo) $ \bufSz -> + allocaArray (nCoo * 6) $ \ixBuf -> + allocaArray nCoo $ \valBuf -> do + ec <- exitCodeH <$> $(varE . mkName $ mkCFnName Read groupName dataName) trexio 0 bufSz ixBuf valBuf + case ec of + Success -> do + ixs <- peek6DCoords (Sz1 nCoo) ixBuf + vals <- peekArray (Sz1 nCoo) . castPtr $ valBuf + mkCooArray (Sz $ sz1 :> sz2 :> sz3 :> sz4 :> sz5 :. sz6) ixs . compute @U $ vals + _ -> throwM ec + |] + | otherwise -> error $ "mkReadFns: unsupported field type for 6D data: " <> show fieldType + [d1, d2, d3, d4, d5, d6, d7, d8] + | isSparseFloat fieldType -> + [e| + \trexio -> liftIO $ do + -- Size of the array + sz1 <- $(mkSizeFn d1) trexio + sz2 <- $(mkSizeFn d2) trexio + sz3 <- $(mkSizeFn d3) trexio + sz4 <- $(mkSizeFn d4) trexio + sz5 <- $(mkSizeFn d5) trexio + sz6 <- $(mkSizeFn d6) trexio + sz7 <- $(mkSizeFn d7) trexio + sz8 <- $(mkSizeFn d8) trexio + + -- Number of COO elements in the sparse array + nCoo <- alloca $ \buf -> do + ec <- exitCodeH <$> $(varE . mkName $ mkCSizeFnName groupName dataName) trexio buf + case ec of + Success -> fromIntegral <$> peek buf + _ -> throwM ec + + -- Read the COO array in a single chunk + with (fromIntegral nCoo) $ \bufSz -> + allocaArray (nCoo * 8) $ \ixBuf -> + allocaArray nCoo $ \valBuf -> do + ec <- exitCodeH <$> $(varE . mkName $ mkCFnName Read groupName dataName) trexio 0 bufSz ixBuf valBuf + case ec of + Success -> do + ixs <- peek8DCoords (Sz1 nCoo) ixBuf + vals <- peekArray (Sz1 nCoo) . castPtr $ valBuf + mkCooArray (Sz $ sz1 :> sz2 :> sz3 :> sz4 :> sz5 :> sz6 :> sz7 :. sz8) ixs . compute @U $ vals + _ -> throwM ec + |] + | otherwise -> error $ "mkReadFns: unsupported field type for 8D data: " <> show fieldType + dl -> error $ "mkReadFns: unsupported number of dimensions: " <> show dl + where + dims = getCrossRefs fieldType + +-- | Get the Length specifications of a field +getCrossRefs :: Typ -> [DimLength] +getCrossRefs (Dim _ (Length lspec)) = lspec +getCrossRefs (Int (Length lspec)) = lspec +getCrossRefs (Float _ (Length lspec)) = lspec +getCrossRefs (Str (Length lspec)) = lspec +getCrossRefs (Idx (Length lspec)) = lspec +getCrossRefs (SparseFloat (Length lspec)) = lspec +getCrossRefs (BitField (Length lspec)) = lspec + +{- | Make a Read function for a given field. This generator takes care to query +referenced 'Dim' fields to obtain the correct size of the result. If any +of this 'Dim' fields is not set, the function will fail. +-} +mkHsReadFn :: GroupName -> DataName -> Typ -> Q [Dec] +mkHsReadFn groupName dataName fieldTyp = do + -- Generate the function name for Haskell + let hsFnName = mkHsFnName Read groupName dataName + + -- Generate the Haskell function + hsFnSig <- mkHsFnSig Read fieldTyp + hsExp <- mkReadFns groupName dataName fieldTyp + return + [ SigD (mkName hsFnName) hsFnSig + , FunD (mkName hsFnName) [Clause [] (NormalB hsExp) []] + ] + +-- | Make a writer function for a given 'DimLength'. +mkWriteSzFn :: TrexioScheme -> DimLength -> Q Exp +mkWriteSzFn _ (Const i) = [e|\_ _ -> return i|] +mkWriteSzFn (TrexioScheme scheme) dimLength@(Field groupName dataName) + | isReadOnly = [e|\_ _ -> return ()|] + | otherwise = do + let cFnName = mkCFnName Write groupName dataName + [e| + \trexio sz -> liftIO $ do + ec <- exitCodeH <$> $(varE . mkName $ cFnName) trexio (fromIntegral sz) + case ec of + Success -> return () + ReadOnly -> return () + -- If the attribute already exists, read it and check if it is the + -- same value we want to write + AttrAlreadyExists -> do + currentSz <- $(mkSizeFn dimLength) trexio + if currentSz == sz + then return () + else throwM AttrAlreadyExists + _ -> throwM ec + |] + where + Group grp = scheme Map.! groupName + fieldTyp = grp Map.! dataName + isReadOnly = case fieldTyp of + Dim False _ -> True + _ -> False + +-- | Make a writer function for a given field +mkWriteFns :: TrexioScheme -> GroupName -> DataName -> Typ -> Q Exp +mkWriteFns scheme groupName dataName fieldType = case dims of + [] + | isWritableIntField fieldType -> + [e| + \trexio int -> liftIO $ do + ec <- exitCodeH <$> $(varE . mkName $ mkCFnName Write groupName dataName) trexio (fromIntegral int) + case ec of + Success -> return () + _ -> throwM ec + |] + | isFloatField fieldType -> + [e| + \trexio float -> liftIO $ do + ec <- exitCodeH <$> $(varE . mkName $ mkCFnName Write groupName dataName) trexio (coerce float) + case ec of + Success -> return () + _ -> throwM ec + |] + | isStringField fieldType -> + [e| + \trexio str -> liftIO . withCStringLen (T.unpack str) $ \(strPtr, len) -> do + ec <- exitCodeH <$> $(varE . mkName $ mkCFnName Write groupName dataName) trexio (ConstPtr strPtr) (fromIntegral len) + case ec of + Success -> return () + _ -> throwM ec + |] + | isProtectedIntField fieldType -> [e|\_ _ -> return ()|] + | otherwise -> error $ "mkWriteFns: unsupported field type for 0D data: " <> show fieldType + [d1] + | isIntField fieldType -> + [e| + \trexio arr -> liftIO . unsafeWithPtr (compute . Massiv.map fromIntegral $ arr) $ \arrPtr -> do + let Sz1 sz1 = size arr + $(mkWriteSzFn scheme d1) trexio sz1 + checkEC $ $(varE . mkName $ mkCFnName Write groupName dataName) trexio arrPtr + |] + | isFloatField fieldType -> + [e| + \trexio arr -> liftIO . unsafeWithPtr arr $ \arrPtr -> do + let Sz1 sz1 = size arr + $(mkWriteSzFn scheme d1) trexio sz1 + checkEC $ $(varE . mkName $ mkCFnName Write groupName dataName) trexio (castPtr arrPtr) + |] + | isStringField fieldType -> + [e| + \trexio arr -> liftIO $ do + let Sz1 nStrings = size arr + maxStrLen = 255 + $(mkWriteSzFn scheme d1) trexio nStrings + ptrArr <- compute <$> mapM (fmap ConstPtr . newCString . T.unpack) arr + unsafeWithPtr ptrArr $ \arrPtr -> + checkEC $ + $(varE . mkName $ mkCFnName Write groupName dataName) + trexio + (ConstPtr arrPtr) + maxStrLen + |] + | isBitField fieldType -> + [e| + \trexio dets -> liftIO $ do + nInt64PerDet <- intsPerDet trexio + let Sz2 nDets _nMos = size dets + $(mkWriteSzFn scheme d1) trexio nDets + + allocaArray (nDets * nInt64PerDet * 2) $ \detBuf -> do + -- Write each determinant to the buffer + forM_ [0 .. nDets - 1] $ \i -> do + let det = dets !> i + detToByteString bv accFn = + BV.cloneToByteString + . Massiv.toVector + . compute @U + . Massiv.map accFn + $ bv + up = detToByteString det fst + down = detToByteString det snd + nBytes = BS.length up + upPtr = detBuf `plusPtr` (i * nInt64PerDet * 2 * sizeOf (undefined :: Int64)) + downPtr = upPtr `plusPtr` (nInt64PerDet * sizeOf (undefined :: Int64)) + + -- Up spin + BS.unsafeUseAsCString up $ \charPtr -> do + copyBytes (castPtr upPtr) charPtr nBytes + + -- Down spin + BS.unsafeUseAsCString down $ \charPtr -> do + copyBytes (castPtr downPtr) charPtr nBytes + + -- Call the C funciton with the buffer + checkEC $ + $(varE . mkName $ mkCFnName Write groupName dataName) + trexio + 0 + (fromIntegral nDets) + detBuf + |] + | isBufferedFloat fieldType -> + [e| + \trexio vec -> liftIO $ do + let Sz1 sz1 = size vec + $(mkWriteSzFn scheme d1) trexio sz1 + unsafeWithPtr vec $ \arrPtr -> + checkEC $ $(varE . mkName $ mkCFnName Write groupName dataName) trexio 0 (fromIntegral sz1) (castPtr arrPtr) + |] + | otherwise -> error $ "mkWriteFns: unsupported field type for 1D data: " <> show fieldType + [d1, d2] + | isFloatField fieldType -> + [e| + \trexio arr -> liftIO . unsafeWithPtr arr $ \arrPtr -> do + let Sz2 sz1 sz2 = size arr + $(mkWriteSzFn scheme d1) trexio sz1 + $(mkWriteSzFn scheme d2) trexio sz2 + checkEC $ $(varE . mkName $ mkCFnName Write groupName dataName) trexio (castPtr arrPtr) + |] + | isSparseFloat fieldType -> + [e| + \trexio cooArr -> liftIO $ do + let Sz2 sz1 sz2 = cooSize cooArr + $(mkWriteSzFn scheme d1) trexio sz1 + $(mkWriteSzFn scheme d2) trexio sz2 + let cooVals = convert . values $ cooArr :: Vector S Double + cooIxs = castCoords2D . coords $ cooArr :: Matrix S Int32 + Sz1 nCoo = size cooVals + unsafeWithPtr cooVals $ \valPtr -> + unsafeWithPtr cooIxs $ \ixPtr -> do + checkEC $ + $(varE . mkName $ mkCFnName Write groupName dataName) + trexio + 0 + (fromIntegral nCoo :: Int64) + ixPtr + (castPtr valPtr) + |] + | otherwise -> error $ "mkWriteFns: unsupported field type for 2D data: " <> show fieldType + [d1, d2, d3] + | isFloatField fieldType -> + [e| + \trexio arr -> liftIO . unsafeWithPtr arr $ \arrPtr -> do + let Sz3 sz1 sz2 sz3 = size arr + $(mkWriteSzFn scheme d1) trexio sz1 + $(mkWriteSzFn scheme d2) trexio sz2 + $(mkWriteSzFn scheme d3) trexio sz3 + checkEC $ $(varE . mkName $ mkCFnName Write groupName dataName) trexio (castPtr arrPtr) + |] + | isSparseFloat fieldType -> + [e| + \trexio cooArr -> liftIO $ do + let Sz3 sz1 sz2 sz3 = cooSize cooArr + $(mkWriteSzFn scheme d1) trexio sz1 + $(mkWriteSzFn scheme d2) trexio sz2 + $(mkWriteSzFn scheme d3) trexio sz3 + let cooVals = convert . values $ cooArr + cooIxs = castCoords3D . coords $ cooArr + Sz1 nCoo = size cooVals + unsafeWithPtr cooVals $ \valPtr -> + unsafeWithPtr cooIxs $ \ixPtr -> + checkEC $ + $(varE . mkName $ mkCFnName Write groupName dataName) + trexio + 0 + (fromIntegral nCoo) + ixPtr + (castPtr valPtr) + |] + | otherwise -> error $ "mkWriteFns: unsupported field type for 3D data: " <> show fieldType + [d1, d2, d3, d4] + | isFloatField fieldType -> + [e| + \trexio arr -> liftIO . unsafeWithPtr arr $ \arrPtr -> do + let Sz4 sz1 sz2 sz3 sz4 = size arr + $(mkWriteSzFn scheme d1) trexio sz1 + $(mkWriteSzFn scheme d2) trexio sz2 + $(mkWriteSzFn scheme d3) trexio sz3 + $(mkWriteSzFn scheme d4) trexio sz4 + checkEC $ $(varE . mkName $ mkCFnName Write groupName dataName) trexio (castPtr arrPtr) + |] + | isSparseFloat fieldType -> + [e| + \trexio cooArr -> liftIO $ do + let Sz4 sz1 sz2 sz3 sz4 = cooSize cooArr + $(mkWriteSzFn scheme d1) trexio sz1 + $(mkWriteSzFn scheme d2) trexio sz2 + $(mkWriteSzFn scheme d3) trexio sz3 + $(mkWriteSzFn scheme d4) trexio sz4 + let cooVals = convert . values $ cooArr + cooIxs = castCoords4D . coords $ cooArr + Sz1 nCoo = size cooVals + unsafeWithPtr cooVals $ \valPtr -> + unsafeWithPtr cooIxs $ \ixPtr -> + checkEC $ + $(varE . mkName $ mkCFnName Write groupName dataName) + trexio + 0 + (fromIntegral nCoo) + ixPtr + (castPtr valPtr) + |] + | otherwise -> error $ "mkWriteFns: unsupported field type for 4D data: " <> show fieldType + [d1, d2, d3, d4, d5, d6] + | isSparseFloat fieldType -> + [e| + \trexio cooArr -> liftIO $ do + let Sz (sz1 :> sz2 :> sz3 :> sz4 :> sz5 :. sz6) = cooSize cooArr + $(mkWriteSzFn scheme d1) trexio sz1 + $(mkWriteSzFn scheme d2) trexio sz2 + $(mkWriteSzFn scheme d3) trexio sz3 + $(mkWriteSzFn scheme d4) trexio sz4 + $(mkWriteSzFn scheme d5) trexio sz5 + $(mkWriteSzFn scheme d6) trexio sz6 + let cooVals = convert . values $ cooArr + cooIxs = castCoords6D . coords $ cooArr + Sz1 nCoo = size cooVals + unsafeWithPtr cooVals $ \valPtr -> + unsafeWithPtr cooIxs $ \ixPtr -> + checkEC $ + $(varE . mkName $ mkCFnName Write groupName dataName) + trexio + 0 + (fromIntegral nCoo) + ixPtr + (castPtr valPtr) + |] + | otherwise -> error $ "mkWriteFns: unsupported field type for 6D data: " <> show fieldType + [d1, d2, d3, d4, d5, d6, d7, d8] + | isSparseFloat fieldType -> + [e| + \trexio cooArr -> liftIO $ do + let Sz (sz1 :> sz2 :> sz3 :> sz4 :> sz5 :> sz6 :> sz7 :. sz8) = cooSize cooArr + $(mkWriteSzFn scheme d1) trexio sz1 + $(mkWriteSzFn scheme d2) trexio sz2 + $(mkWriteSzFn scheme d3) trexio sz3 + $(mkWriteSzFn scheme d4) trexio sz4 + $(mkWriteSzFn scheme d5) trexio sz5 + $(mkWriteSzFn scheme d6) trexio sz6 + $(mkWriteSzFn scheme d7) trexio sz7 + $(mkWriteSzFn scheme d8) trexio sz8 + let cooVals = convert . values $ cooArr + cooIxs = castCoords8D . coords $ cooArr + Sz1 nCoo = size cooVals + unsafeWithPtr cooVals $ \valPtr -> + unsafeWithPtr cooIxs $ \ixPtr -> + checkEC $ + $(varE . mkName $ mkCFnName Write groupName dataName) + trexio + 0 + (fromIntegral nCoo) + ixPtr + (castPtr valPtr) + |] + | otherwise -> error $ "mkWriteFns: unsupported field type for 8D data: " <> show fieldType + dl -> error $ "mkWriteFns: unsupported number of dimensions: " <> show dl + where + dims = getCrossRefs fieldType + +mkHsWriteFn :: TrexioScheme -> GroupName -> DataName -> Typ -> Q [Dec] +mkHsWriteFn scheme groupName dataName fieldTyp = do + -- Generate the function names in C and Haskell + let hsFnName = mkHsFnName Write groupName dataName + + -- Generate the Haskell function + hsFnSig <- mkHsFnSig Write fieldTyp + hsExp <- mkWriteFns scheme groupName dataName fieldTyp + return + [ SigD (mkName hsFnName) hsFnSig + , FunD (mkName hsFnName) [Clause [] (NormalB hsExp) []] + ] + +mkCDeleteName :: GroupName -> String +mkCDeleteName (GroupName groupName) = "trexio_delete_" <> T.unpack groupName + +mkHsDeleteName :: GroupName -> String +mkHsDeleteName (GroupName groupName) = sanId . camel $ "delete_" <> T.unpack groupName + +mkCDeleteFn :: GroupName -> Q Dec +mkCDeleteFn groupName = do + let cFnName = mkCDeleteName groupName + cTyp <- [t|Trexio -> IO ExitCodeC|] + return . ForeignD $ ImportF CApi Unsafe ("trexio.h " <> cFnName) (mkName cFnName) cTyp + +mkHsDeleteFn :: GroupName -> Q [Dec] +mkHsDeleteFn groupName = do + let cFnName = mkName . mkCDeleteName $ groupName + hsFnName = mkName . mkHsDeleteName $ groupName + hsTyp <- [t|forall m. (MonadIO m) => Trexio -> m ()|] + hsFn <- [e|\trexio -> liftIO . checkEC $ $(varE cFnName) trexio|] + return + [ SigD hsFnName hsTyp + , FunD hsFnName [Clause [] (NormalB hsFn) []] + ] diff --git a/src/TREXIO/LowLevel.hs b/src/TREXIO/LowLevel.hs new file mode 100644 --- /dev/null +++ b/src/TREXIO/LowLevel.hs @@ -0,0 +1,32 @@ +{-# LANGUAGE QuasiQuotes #-} +{-# LANGUAGE TemplateHaskell #-} + +{- | +Module: TREXIO.LowLevel +Description: Generated direct low-level bindings to the TREXIO library +Copyright: Phillip Seeber 2024 +License: BSD-3-Clause +Maintainer: phillip.seeber@uni-jena.de +Stability: experimental +Portability: POSIX + +These are the low-level bindings to the TREXIO library. +They are generated by the Template Haskell and provide direct access to the C functions in the TREXIO library. +Consequently, they are unsafe and require manual memory management. +-} +module TREXIO.LowLevel where + +import Data.Map qualified as Map +import Foreign.C.ConstPtr +import Foreign.C.Types +import TREXIO.Internal.Base +import TREXIO.Internal.TH +import TREXIO.LowLevel.Scheme + +$( do + let TrexioScheme trexioScheme = scheme + groups = Map.toList trexioScheme + + -- Import all C functions for all fields and operations + concat <$> traverse (uncurry mkCBindings) groups + ) diff --git a/src/TREXIO/LowLevel/Scheme.hs b/src/TREXIO/LowLevel/Scheme.hs new file mode 100644 --- /dev/null +++ b/src/TREXIO/LowLevel/Scheme.hs @@ -0,0 +1,25 @@ +{-# LANGUAGE QuasiQuotes #-} +{-# LANGUAGE TemplateHaskell #-} + +{- | +Module: TREXIO.LowLevel.Scheme +Description: The TREXIO scheme +Copyright: Phillip Seeber 2024 +License: BSD-3-Clause +Maintainer: phillip.seeber@uni-jena.de +Stability: experimental +Portability: POSIX +-} +module TREXIO.LowLevel.Scheme where + +import Language.Haskell.TH +import TREXIO.Internal.TH +import Language.Haskell.TH.Syntax (lift) + +-- | The JSON specification of the code generator, that constructs the C-API and +-- that this package binds to. +scheme :: TrexioScheme +scheme = $(do + trexio <- runIO getJsonSpec + lift trexio + ) diff --git a/test/trexio-test.hs b/test/trexio-test.hs new file mode 100644 --- /dev/null +++ b/test/trexio-test.hs @@ -0,0 +1,446 @@ +import Control.Exception.Safe +import Control.Monad +import Data.Massiv.Array as Massiv hiding (Size, elem, forM, forM_, mapM, mapM_, take, zip, zipWith) +import Data.Maybe (catMaybes, fromJust) +import Data.Set qualified as Set +import Data.Text (Text) +import Hedgehog (Gen, MonadGen, Size, forAll, property, (===)) +import Hedgehog.Gen qualified as Gen +import Hedgehog.Range qualified as Range +import System.Directory +import System.IO.Temp +import TREXIO +import TREXIO.CooArray +import Test.Tasty +import Test.Tasty.HUnit +import Test.Tasty.Hedgehog + +main :: IO () +main = defaultMain tests + +tests :: TestTree +tests = + testGroup + "TREXIO" + [ testGroup + "0D" + [ testGroup "Integers" . appFn $ + [ ("nucleus.num", genDim, deleteNucleus, hasNucleusNum, readNucleusNum, writeNucleusNum) + , ("grid.max_ang_num", genInt, deleteGrid, hasGridMaxAngNum, readGridMaxAngNum, writeGridMaxAngNum) + , ("state.id", genIndex, deleteState, hasStateId, readStateId, writeStateId) + ] + , testGroup "Floats" . appFn $ + [ ("nucleus.repulsion", genFloat, deleteNucleus, hasNucleusRepulsion, readNucleusRepulsion, writeNucleusRepulsion) + ] + , testGroup "Strings" . appFn $ + [ ("metadata.description", genIdentifier, deleteMetadata, hasMetadataDescription, readMetadataDescription, writeMetadataDescription) + ] + ] + , testGroup + "1D" + [ testGroup "Integers" . appFn $ + [ ("ecp.ang_num", genVector genInt, deleteEcp, hasEcpAngMom, readEcpAngMom, writeEcpAngMom) + ] + , testGroup "Floats" . appFn $ + [ ("basis.shell_factor", genVector genFloat, deleteBasis, hasBasisShellFactor, readBasisShellFactor, writeBasisShellFactor) + ] + , testGroup "Strings" . appFn $ + [ ("metadata.author", genVector genIdentifier, deleteMetadata, hasMetadataAuthor, readMetadataAuthor, writeMetadataAuthor) + ] + , testCase "Determinant IO" . withSystemTempFile "trexio.dat" $ \fp _ -> + withTrexio fp FileWrite Hdf5 $ \trexio -> do + writeMoNum trexio 3 + + has1 <- hasDeterminantList trexio + has1 @?= False + + ingoreExcp [AttrMissing] (readDeterminantList trexio) + + let detList = + Massiv.fromLists' + Seq + [ [(1, 1), (1, 0), (0, 0)] + , [(1, 0), (1, 1), (0, 0)] + ] + detCoeffs = Massiv.fromList Seq [1.0, 2.0] + + writeDeterminantList trexio detList + + has2 <- hasDeterminantList trexio + has2 @?= True + + detList' <- readDeterminantList trexio + detList' @?= detList + + writeDeterminantCoefficient trexio detCoeffs + + detCoeffs' <- readDeterminantCoefficient trexio + detCoeffs' @?= detCoeffs + ] + , testGroup + "2D" + [ testGroup "Dense" . appFn $ + [ ("ao_1e_int.overlap", genMatrix SameAs1 genFloat, deleteAo1eInt, hasAo1eIntOverlap, readAo1eIntOverlap, writeAo1eIntOverlap) + ] + , testGroup "Sparse" . appFn $ + [ ("amplitude.single", genSparseArr2 SameAs1 genFloat, deleteAmplitude, hasAmplitudeSingle, readAmplitudeSingle, writeAmplitudeSingle) + ] + ] + , testGroup + "3D" + [ testGroup "Sparse" . appFn $ + [ ("mo_2e_int.eri_lr_cholesky", genSparseArr3 SameAs1 SameAs1 genFloat, deleteMo2eInt, hasMo2eIntEriLrCholesky, readMo2eIntEriLrCholesky, writeMo2eIntEriLrCholesky) + ] + ] + , testGroup "4D" . appFn $ + [ ("mo_2e_int.eri", genSparseArr4 SameAs1 SameAs1 SameAs1 genFloat, deleteMo2eInt, hasMo2eIntEri, readMo2eIntEri, writeMo2eIntEri) + ] + , testGroup "6D" . appFn $ + [ ("amplitude.triple", genSparseArr6 SameAs1 SameAs1 SameAs1 SameAs1 SameAs1 genFloat, deleteAmplitude, hasAmplitudeTriple, readAmplitudeTriple, writeAmplitudeTriple) + ] + , testGroup "8D" . appFn $ + [ ("amplitude.quadruple", genSparseArr8 SameAs1 SameAs1 SameAs1 SameAs1 SameAs1 SameAs1 SameAs1 genFloat, deleteAmplitude, hasAmplitudeQuadruple, readAmplitudeQuadruple, writeAmplitudeQuadruple) + ] + ] + where + appFn :: (Eq a, Show a) => [(TestName, Gen a, Trexio -> IO (), Trexio -> IO Bool, Trexio -> IO a, Trexio -> a -> IO ())] -> [TestTree] + appFn = fmap (\(name, val, delFn, hasFn, readFn, writeFn) -> testField name val delFn hasFn readFn writeFn) + +ingoreExcp :: (MonadCatch m) => [ExitCode] -> m a -> m () +ingoreExcp excps action = catch (void action) $ \e -> + if e `elem` excps + then return () + else throw e + +testField :: + (Eq a, Show a) => + -- | Name of the test + TestName -> + -- | Value to write + Gen a -> + -- | The group delete function + (Trexio -> IO ()) -> + -- | The "Has" function + (Trexio -> IO Bool) -> + -- | The "Read" function + (Trexio -> IO a) -> + -- | The "Write" function + (Trexio -> a -> IO ()) -> + TestTree +testField name gen delFn hasFn readFn writeFn = testProperty name . property $ do + val <- forAll gen + fp <- liftIO $ emptySystemTempFile "trexio.h5" + + trexio <- open fp FileUnsafe Hdf5 + + -- Mark operations as unsafe + safetyFlagU <- readMetadataUnsafe trexio + safetyFlagU === 1 + markSafety trexio + safetyFlagS <- readMetadataUnsafe trexio + safetyFlagS === 0 + + -- Nothing should be there yet + has1 <- liftIO $ hasFn trexio + has1 === False + + -- Reading should return a missing attribute exception + ingoreExcp [AttrMissing] (liftIO $ readFn trexio) + + -- Write the value to the file + liftIO $ writeFn trexio val + + -- Writing again should return an attribute already exists exception + -- ingoreExcp [AttrAlreadyExists, DSetAlreadyExists] (write trexio val) + + -- Check if it is there + has2 <- liftIO $ hasFn trexio + has2 === True + + -- Read it back + val' <- liftIO $ readFn trexio + val' === val + + -- Delete the entire group + liftIO $ delFn trexio + + -- Nothing should be there anymore + has3 <- liftIO $ hasFn trexio + has3 === False + + close trexio + liftIO $ removeFile fp + +data DimDep + = Independent + | SameAs1 + | SameAs2 + deriving (Show, Eq, Ord) + +-- | Generate a random integer +genInt :: Gen Int +genInt = Gen.integral (Range.linearFrom 0 (-1_000_000) 1_000_000) + +-- | Generate a @dim@ value, which is a positive integer +genDim :: Gen Int +genDim = Gen.integral (Range.linear 1 100) + +-- | Generate a random index, which is a non-negative integer +genIndex :: Gen Int +genIndex = Gen.integral (Range.linear 0 1000) + +genFloat :: Gen Double +genFloat = Gen.realFloat (Range.linearFrac (-1_000_000) 1_000_000) + +-- | Generate a identifier, that is a single word without spaces or stuff +genIdentifier :: Gen Text +genIdentifier = Gen.text (Range.linear 1 10) Gen.alphaNum + +-- | Generate a Massiv vector from elements from another generator +genVector :: (Manifest r a) => Gen a -> Gen (Vector r a) +genVector gen = do + dim <- genDim + Massiv.fromList Seq <$> Gen.list (Range.singleton dim) gen + +-- | Generate a Massiv matrix from elements from another generator +genMatrix :: (Manifest r a, Ord a) => DimDep -> Gen a -> Gen (Matrix r a) +genMatrix dimDep2 gen = do + rows <- genDim + cols <- case dimDep2 of + Independent -> genDim + SameAs1 -> return rows + SameAs2 -> error "genMatrix: SameAs2 not supported for Dim2" + Massiv.fromLists' Seq <$> Gen.list (Range.singleton rows) (Gen.list (Range.singleton cols) gen) + +genArr3 :: (Manifest r a, Ord a) => DimDep -> DimDep -> Gen a -> Gen (Array r Ix3 a) +genArr3 dimDep2 dimDep3 gen = do + d1 <- genDim + d2 <- case dimDep2 of + Independent -> genDim + SameAs1 -> return d1 + SameAs2 -> error "genArr3: SameAs2 not supported for Dim2" + d3 <- case dimDep3 of + Independent -> genDim + SameAs1 -> return d1 + SameAs2 -> return d2 + vals <- Gen.list (Range.singleton $ d1 * d2 * d3) gen + pure . Massiv.resize' (Sz3 d1 d2 d3) $ Massiv.fromList Par vals + +genSparseArr2 :: + forall r a. + (Manifest r a, Ord a, Manifest r Ix2, Stream r Ix1 Ix2) => + DimDep -> + Gen a -> + Gen (CooArray r Ix2 a) +genSparseArr2 dimDep2 gen = scaleSparse $ do + d1 <- genDim + d2 <- case dimDep2 of + Independent -> genDim + SameAs1 -> return d1 + SameAs2 -> error "genSparseArr2: SameAs2 not supported for Dim2" + let sz = Sz2 d1 d2 + cooVals <- + catMaybes . Set.toAscList + <$> Gen.set + (Range.singleton $ d1 * d2) + ( do + c1 <- Gen.int $ Range.linear 0 (d1 - 1) + c2 <- Gen.int $ Range.linear 0 (d2 - 1) + v <- gen + Gen.maybe . pure $ (c1 :. c2, v) + ) + if null cooVals + then Gen.discard + else pure . fromJust $ mkCooArrayF sz cooVals + +genSparseArr3 :: + forall r a. + (Manifest r a, Manifest r Ix3, Ord a, Stream r Ix1 Ix3) => + DimDep -> + DimDep -> + Gen a -> + Gen (CooArray r Ix3 a) +genSparseArr3 dimDep2 dimDep3 gen = scaleSparse $ do + d1 <- genDim + d2 <- case dimDep2 of + Independent -> genDim + SameAs1 -> return d1 + SameAs2 -> error "genSparseArr3: SameAs2 not supported for Dim2" + d3 <- case dimDep3 of + Independent -> genDim + SameAs1 -> return d1 + SameAs2 -> return d2 + let sz = Sz3 d1 d2 d3 + cooVals <- + catMaybes . Set.toAscList + <$> Gen.set + (Range.singleton $ d1 * d2 * d3) + ( do + c1 <- Gen.int $ Range.linear 0 (d1 - 1) + c2 <- Gen.int $ Range.linear 0 (d2 - 1) + c3 <- Gen.int $ Range.linear 0 (d3 - 1) + v <- gen + Gen.maybe . pure $ (c1 :> c2 :. c3, v) + ) + if null cooVals + then Gen.discard + else pure . fromJust $ mkCooArrayF sz cooVals + +genSparseArr4 :: + forall r a. + (Manifest r a, Ord a, Manifest r Ix4, Stream r Ix1 Ix4) => + DimDep -> + DimDep -> + DimDep -> + Gen a -> + Gen (CooArray r Ix4 a) +genSparseArr4 dimDep2 dimDep3 dimDep4 gen = scaleSparse $ do + d1 <- genDim + d2 <- case dimDep2 of + Independent -> genDim + SameAs1 -> return d1 + SameAs2 -> error "genSparseArr4: SameAs2 not implemented" + d3 <- case dimDep3 of + Independent -> genDim + SameAs1 -> return d1 + SameAs2 -> return d2 + d4 <- case dimDep4 of + Independent -> genDim + SameAs1 -> return d1 + SameAs2 -> return d2 + let sz = Sz4 d1 d2 d3 d4 + cooVals <- + catMaybes . Set.toAscList + <$> Gen.set + (Range.singleton $ d1 * d2 * d3) + ( do + c1 <- Gen.int $ Range.linear 0 (d1 - 1) + c2 <- Gen.int $ Range.linear 0 (d2 - 1) + c3 <- Gen.int $ Range.linear 0 (d3 - 1) + c4 <- Gen.int $ Range.linear 0 (d4 - 1) + v <- gen + Gen.maybe . pure $ (c1 :> c2 :> c3 :. c4, v) + ) + if null cooVals + then Gen.discard + else pure . fromJust $ mkCooArrayF sz cooVals + +genSparseArr6 :: + forall r a. + (Manifest r a, Ord a, Manifest r (IxN 6), Stream r Ix1 (IxN 6)) => + DimDep -> + DimDep -> + DimDep -> + DimDep -> + DimDep -> + Gen a -> + Gen (CooArray r (IxN 6) a) +genSparseArr6 dimDep2 dimDep3 dimDep4 dimDep5 dimDep6 gen = scaleSparse $ do + d1 <- genDim + d2 <- case dimDep2 of + Independent -> genDim + SameAs1 -> return d1 + SameAs2 -> error "genSparseArr4: SameAs2 not implemented" + d3 <- case dimDep3 of + Independent -> genDim + SameAs1 -> return d1 + SameAs2 -> return d2 + d4 <- case dimDep4 of + Independent -> genDim + SameAs1 -> return d1 + SameAs2 -> return d2 + d5 <- case dimDep5 of + Independent -> genDim + SameAs1 -> return d1 + SameAs2 -> return d2 + d6 <- case dimDep6 of + Independent -> genDim + SameAs1 -> return d1 + SameAs2 -> return d2 + let sz = Sz $ d1 :> d2 :> d3 :> d4 :> d5 :. d6 + cooVals <- + catMaybes . Set.toAscList + <$> Gen.set + (Range.singleton $ d1 * d2 * d3) + ( do + c1 <- Gen.int $ Range.linear 0 (d1 - 1) + c2 <- Gen.int $ Range.linear 0 (d2 - 1) + c3 <- Gen.int $ Range.linear 0 (d3 - 1) + c4 <- Gen.int $ Range.linear 0 (d4 - 1) + c5 <- Gen.int $ Range.linear 0 (d5 - 1) + c6 <- Gen.int $ Range.linear 0 (d6 - 1) + v <- gen + Gen.maybe . pure $ (c1 :> c2 :> c3 :> c4 :> c5 :. c6, v) + ) + if null cooVals + then Gen.discard + else pure . fromJust $ mkCooArrayF sz cooVals + +genSparseArr8 :: + forall r a. + (Manifest r a, Ord a, Manifest r (IxN 8), Stream r Ix1 (IxN 8)) => + DimDep -> + DimDep -> + DimDep -> + DimDep -> + DimDep -> + DimDep -> + DimDep -> + Gen a -> + Gen (CooArray r (IxN 8) a) +genSparseArr8 dimDep2 dimDep3 dimDep4 dimDep5 dimDep6 dimDep7 dimDep8 gen = scaleSparse $ do + d1 <- genDim + d2 <- case dimDep2 of + Independent -> genDim + SameAs1 -> return d1 + SameAs2 -> error "genSparseArr4: SameAs2 not implemented" + d3 <- case dimDep3 of + Independent -> genDim + SameAs1 -> return d1 + SameAs2 -> return d2 + d4 <- case dimDep4 of + Independent -> genDim + SameAs1 -> return d1 + SameAs2 -> return d2 + d5 <- case dimDep5 of + Independent -> genDim + SameAs1 -> return d1 + SameAs2 -> return d2 + d6 <- case dimDep6 of + Independent -> genDim + SameAs1 -> return d1 + SameAs2 -> return d2 + d7 <- case dimDep7 of + Independent -> genDim + SameAs1 -> return d1 + SameAs2 -> return d2 + d8 <- case dimDep8 of + Independent -> genDim + SameAs1 -> return d1 + SameAs2 -> return d2 + let sz = Sz $ d1 :> d2 :> d3 :> d4 :> d5 :> d6 :> d7 :. d8 + cooVals <- + catMaybes . Set.toAscList + <$> Gen.set + (Range.singleton $ d1 * d2 * d3) + ( do + c1 <- Gen.int $ Range.linear 0 (d1 - 1) + c2 <- Gen.int $ Range.linear 0 (d2 - 1) + c3 <- Gen.int $ Range.linear 0 (d3 - 1) + c4 <- Gen.int $ Range.linear 0 (d4 - 1) + c5 <- Gen.int $ Range.linear 0 (d5 - 1) + c6 <- Gen.int $ Range.linear 0 (d6 - 1) + c7 <- Gen.int $ Range.linear 0 (d7 - 1) + c8 <- Gen.int $ Range.linear 0 (d8 - 1) + v <- gen + Gen.maybe . pure $ (c1 :> c2 :> c3 :> c4 :> c5 :> c6 :> c7 :. c8, v) + ) + if null cooVals + then Gen.discard + else pure . fromJust $ mkCooArrayF sz cooVals + +scaleSparse :: (MonadGen m) => m a -> m a +scaleSparse = Gen.scale sz2zs + where + sz2zs :: Size -> Size + sz2zs x = round $ fromIntegral x * (0.25 :: Double) diff --git a/trexio-hs.cabal b/trexio-hs.cabal new file mode 100644 --- /dev/null +++ b/trexio-hs.cabal @@ -0,0 +1,151 @@ +cabal-version: 3.0 +name: trexio-hs +-- The package version. +-- See the Haskell package versioning policy (PVP) for standards +-- guiding when and how versions should be incremented. +-- https://pvp.haskell.org +-- PVP summary: +-+------- breaking API changes +-- | | +----- non-breaking API additions +-- | | | +--- code changes with no API change +version: 0.1.0 +synopsis: Bindings to the TREXIO library for wave function data +homepage: https://github.com/TREX-CoE/trexio-hs +license: BSD-3-Clause +author: Phillip Seeber +maintainer: phillip.seeber@uni-jena.de +category: Data +build-type: Simple +extra-doc-files: CHANGELOG.md +tested-with: GHC == {9.6, 9.8, 9.10} +description: + This package provides low- and high-level Haskell bindings for [TREXIO, a portable file format for storing wave function data](https://trex-coe.github.io/trexio/). + The vast majority of the bindings in this package is generated via TemplateHaskell from the TREXIO JSON specification, that then defines the C-API. + For more details see the [TREXIO documentation](https://trex-coe.github.io/trexio/lib.html). + Consequently, this package is able to adapt to changes in the TREXIO specification, but the Hackage version reflects the canonical upstream TREXIO specification. + + The low-level API is a direct mapping of the C-API, with minimal abstraction, i.e. passing raw pointers around and responsibility for memory management is on the user. + See the @TREXIO.LowLevel@ module for more details. + The @TREXIO.LowLevel.Scheme@ module is the Aeson representation of the TREXIO JSON specification and reflects the structure that was used by TemplateHaskell for code generation. + However, the low-level bindings are very complete and allow for buffered IO of large quantities such as the Electron Repulsion Integrals or Configuration State Functions. + + The high-level API is more haskellish and provides a more type-safe and user-friendly interface to the TREXIO library. + The @TREXIO@ module provides your entry point to the high-level API. + Function names are automatically generated by stripping the @trexio_@ prefix from the C-API and converting to camel case. + Furthermore, @TREXIO.HighLevel.Records@ provides a direct translation of the TREXIO JSON specification to Haskell records (without buffered IO, however). + The @TREXIO.HighLevel.Records@ is not used for any other purpose within this library, however. + + The high-level API utilises @TREXIO.CooArray@ to provide a more idiomatic interface to sparse arrays used for many high-dimensional quantities within TREXIO. + + These bindings heavily rely on [Massiv](https://hackage.haskell.org/package/massiv) for handling arrays. + + Example usage: + + @ + import TREXIO + import TREXIO.CooArray + import Data.Massiv.Array as Massiv + + main :: IO () + main = do + \-- Get the version of TREXIO we have linked against + ver <- version + putStrLn $ "TREXIO version: " <> ver + + withTrexio "example.h5" FileWrite Hdf5 $ \\trexio -> do + \-- Write nuclear coordinates to a TREXIO file + coords <- Massiv.fromListsM Par + [ [0. , 0., -0.24962655] + , [0. , 2.70519714, 1.85136466] + , [0. , -2.70519714, 1.85136466] + ] + writeNucleusCoord trexio coords + + \-- Read reduced 2-electron density matrix (assuming it exists in the example.h5) + rdm2e <- readRdm2e trexio + print rdm2e + @ + +common warnings + ghc-options: -Wall + +common deps + build-depends: + aeson >= 2.1 && < 2.3, + base >= 4.18 && < 4.22, + bitvec >= 1.1.5.0 && < 1.2, + bytestring >= 0.10 && < 0.13, + casing >= 0.1.4 && < 0.2, + containers >= 0.6 && < 0.8, + filepath >= 1.4 && < 1.6, + massiv >= 1.0.0.0 && < 1.1, + safe-exceptions >= 0.1.7 && < 0.2, + template-haskell >= 2.19 && < 2.24, + temporary >= 1.3 && < 1.4, + typed-process >= 0.2.12 && < 0.3, + text >= 2.0 && < 2.2, + vector >= 0.13 && < 0.14 + +common extensions + default-extensions: + CApiFFI, + DataKinds, + DeriveAnyClass, + DerivingVia, + DerivingStrategies, + OverloadedStrings, + OverloadedRecordDot, + RecordWildCards, + DuplicateRecordFields, + LambdaCase + +library trexio-internal + import: + deps, + extensions, + warnings + exposed-modules: + TREXIO.Internal.Base + hs-source-dirs: src-int + default-language: GHC2021 + extra-libraries: + trexio + +library + import: + deps, + extensions, + warnings + exposed-modules: + TREXIO, + TREXIO.LowLevel, + TREXIO.LowLevel.Scheme, + TREXIO.HighLevel, + TREXIO.HighLevel.Records, + TREXIO.CooArray + other-modules: + TREXIO.Internal.Marshaller, + TREXIO.Internal.TH + hs-source-dirs: src + default-language: GHC2021 + build-depends: + trexio-internal + extra-libraries: + trexio + +test-suite trexio-test + import: + deps, + extensions, + warnings + default-language: GHC2021 + type: exitcode-stdio-1.0 + hs-source-dirs: test + main-is: trexio-test.hs + build-depends: + trexio-hs, + directory >= 1.3 && < 1.4, + tasty >= 1.4 && < 1.6, + tasty-hunit >= 0.10 && < 0.11, + tasty-hedgehog >= 1.4 && < 1.5, + hedgehog >= 1.4 && < 1.6 +