Hasktorch.hs

{-# LANGUAGE BangPatterns #-}

module DataFrame.Hasktorch (
    toTensor,
    toIntTensor,
) where

import qualified Data.Vector as V
import qualified Data.Vector.Unboxed as VU
import qualified Data.Vector.Unboxed.Mutable as VUM
import qualified DataFrame as D

import Control.Exception (throw)
import DataFrame (DataFrame)
import Torch

{- | Converts a dataframe to a floating-point tensor.

This function converts all columns in the dataframe to floats and creates
a tensor suitable for machine learning operations. The tensor dimensions
are determined by the dataframe's shape.

==== __Dimensional behavior__

* Multi-column dataframe: Creates a 2D tensor with shape @[rows, columns]@
* Single-column dataframe: Creates a 1D tensor with shape @[rows]@

==== __Conversion process__

1. Converts the dataframe to a float matrix using 'D.toFloatMatrix'
2. Flattens the matrix features into a 1D representation
3. Reshapes into the appropriate tensor dimensions

==== __Throws__

* 'DataFrameException' - if any column cannot be converted to float

==== __Examples__

>>> toTensor df  -- where df has shape (100, 5)
Tensor with shape [100, 5]

>>> toTensor df  -- where df has shape (100, 1)
Tensor with shape [100]

==== __See also__

* 'toIntTensor' - for integer tensor conversion
-}
toTensor :: DataFrame -> Tensor
toTensor df = case D.toFloatMatrix df of
    Left e -> throw e
    Right m ->
        let
            (r, c) = D.dimensions df
            dims' = if c == 1 then [r] else [r, c]
         in
            reshape dims' (asTensor (flattenFeatures m))

{- | Converts a dataframe to an integer tensor.

This function converts all columns in the dataframe to integers and creates
a tensor suitable for machine learning operations (e.g., classification labels,
discrete features). The tensor dimensions are determined by the dataframe's shape.

==== __Dimensional behavior__

* Multi-column dataframe: Creates a 2D tensor with shape @[rows, columns]@
* Single-column dataframe: Creates a 1D tensor with shape @[rows]@

==== __Conversion process__

1. Converts the dataframe to an int matrix using 'D.toIntMatrix'
2. Flattens the matrix features into a 1D representation
3. Reshapes into the appropriate tensor dimensions

==== __Throws__

* 'DataFrameException' - if any column cannot be converted to int

==== __Examples__

>>> toIntTensor labelsDf  -- where labelsDf has shape (100, 1)
Tensor with shape [100]

>>> toIntTensor featuresDf  -- where featuresDf has shape (100, 3)
Tensor with shape [100, 3]

==== __Note__

Floating-point values in the dataframe will be rounded to the nearest integer.
See 'D.toIntMatrix' for details on the conversion behavior.

==== __See also__

* 'toTensor' - for floating-point tensor conversion
-}
toIntTensor :: DataFrame -> Tensor
toIntTensor df = case D.toIntMatrix df of
    Left e -> throw e
    Right m ->
        let
            (r, c) = D.dimensions df
            dims' = if c == 1 then [r] else [r, c]
         in
            reshape dims' (asTensor (flattenFeatures m))

flattenFeatures :: V.Vector (VU.Vector a) -> VU.Vector a
flattenFeatures rows =
    let
        total = V.foldl' (\s v -> s + VU.length v) 0 rows
     in
        VU.create $ do
            ret <- VUM.unsafeNew total
            let go !i !off
                    | i == V.length rows = pure ()
                    | otherwise = do
                        let v = rows V.! i
                            len = VU.length v
                        VU.unsafeCopy (VUM.unsafeSlice off len ret) v
                        go (i + 1) (off + len)
            go 0 0
            pure ret