{-# OPTIONS_GHC -Wall #-}

{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE TypeApplications #-}

-- | Module for running training
module Run where

import           System.ProgressBar
import           Lib
import           Net
import           HyperParameters
import           Data.Frame                as DF
import qualified Torch                     as T
import qualified Torch.Extensions          as T
import qualified Torch.Functional.Internal as T         (negative)
import qualified Torch.NN           as NN

import           Prelude                         hiding (exp)

------------------------------------------------------------------------------
-- Training
------------------------------------------------------------------------------

-- | Run one Update Step
trainStep :: T.Tensor -> T.Tensor 
          -> OpNet -> T.Adam -> IO (OpNet, T.Adam, T.Tensor)
trainStep :: Tensor -> Tensor -> OpNet -> Adam -> IO (OpNet, Adam, Tensor)
trainStep Tensor
trueX Tensor
trueY OpNet
net Adam
opt = do
    (OpNet
net', Adam
opt') <- OpNet -> Adam -> Tensor -> Tensor -> IO (OpNet, Adam)
forall optimizer model.
(Optimizer optimizer, Parameterized model) =>
model -> optimizer -> Tensor -> Tensor -> IO (model, optimizer)
T.runStep OpNet
net Adam
opt Tensor
loss Tensor
α
    (OpNet, Adam, Tensor) -> IO (OpNet, Adam, Tensor)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (OpNet
net', Adam
opt', Tensor
loss)
  where
    predY :: Tensor
predY = OpNet -> Tensor -> Tensor
forward OpNet
net Tensor
trueX
    loss :: Tensor
loss  = Reduction -> Tensor -> Tensor -> Tensor
T.l1Loss Reduction
T.ReduceSum Tensor
trueY Tensor
predY

-- | Run through all Batches performing an update for each
trainingEpoch :: ProgressBar s -> [T.Tensor]  -> [T.Tensor] -> [T.Tensor] 
              -> OpNet -> T.Adam -> IO (OpNet, T.Adam, T.Tensor)
trainingEpoch :: ProgressBar s
-> [Tensor]
-> [Tensor]
-> [Tensor]
-> OpNet
-> Adam
-> IO (OpNet, Adam, Tensor)
trainingEpoch ProgressBar s
_     [Tensor]
_      []   [Tensor]
losses OpNet
net Adam
opt = (OpNet, Adam, Tensor) -> IO (OpNet, Adam, Tensor)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (OpNet
net, Adam
opt, Tensor
losses')
  where
    losses' :: Tensor
losses' = Dim -> [Tensor] -> Tensor
T.cat (Int -> Dim
T.Dim Int
0) ([Tensor] -> Tensor)
-> ([Tensor] -> [Tensor]) -> [Tensor] -> Tensor
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Tensor -> Tensor) -> [Tensor] -> [Tensor]
forall a b. (a -> b) -> [a] -> [b]
map ([Int] -> Tensor -> Tensor
T.reshape [-Int
1]) ([Tensor] -> Tensor) -> [Tensor] -> Tensor
forall a b. (a -> b) -> a -> b
$ [Tensor]
losses
trainingEpoch ProgressBar s
_     []     [Tensor]
_    [Tensor]
losses OpNet
net Adam
opt = (OpNet, Adam, Tensor) -> IO (OpNet, Adam, Tensor)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (OpNet
net, Adam
opt, Tensor
losses')
  where
    losses' :: Tensor
losses' = Dim -> [Tensor] -> Tensor
T.cat (Int -> Dim
T.Dim Int
0) ([Tensor] -> Tensor)
-> ([Tensor] -> [Tensor]) -> [Tensor] -> Tensor
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Tensor -> Tensor) -> [Tensor] -> [Tensor]
forall a b. (a -> b) -> [a] -> [b]
map ([Int] -> Tensor -> Tensor
T.reshape [-Int
1]) ([Tensor] -> Tensor) -> [Tensor] -> Tensor
forall a b. (a -> b) -> a -> b
$ [Tensor]
losses
trainingEpoch ProgressBar s
bar (Tensor
x:[Tensor]
xs) (Tensor
y:[Tensor]
ys) [Tensor]
losses OpNet
net Adam
opt = do
        ProgressBar s -> Int -> IO ()
forall s. ProgressBar s -> Int -> IO ()
incProgress ProgressBar s
bar Int
1
        Tensor -> Tensor -> OpNet -> Adam -> IO (OpNet, Adam, Tensor)
trainStep Tensor
x Tensor
y OpNet
net Adam
opt IO (OpNet, Adam, Tensor)
-> ((OpNet, Adam, Tensor) -> IO (OpNet, Adam, Tensor))
-> IO (OpNet, Adam, Tensor)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (OpNet, Adam, Tensor) -> IO (OpNet, Adam, Tensor)
trainingEpoch' 
  where
    trainingEpoch' :: (OpNet, Adam, Tensor) -> IO (OpNet, Adam, Tensor)
trainingEpoch' = \ (OpNet
n, Adam
o, Tensor
l) -> ProgressBar s
-> [Tensor]
-> [Tensor]
-> [Tensor]
-> OpNet
-> Adam
-> IO (OpNet, Adam, Tensor)
forall s.
ProgressBar s
-> [Tensor]
-> [Tensor]
-> [Tensor]
-> OpNet
-> Adam
-> IO (OpNet, Adam, Tensor)
trainingEpoch ProgressBar s
bar [Tensor]
xs [Tensor]
ys (Tensor
lTensor -> [Tensor] -> [Tensor]
forall a. a -> [a] -> [a]
:[Tensor]
losses) OpNet
n Adam
o

------------------------------------------------------------------------------
-- Validation
------------------------------------------------------------------------------

-- | Run one Update Step
validStep :: T.Tensor -> T.Tensor -> OpNet -> IO T.Tensor
validStep :: Tensor -> Tensor -> OpNet -> IO Tensor
validStep Tensor
trueX Tensor
trueY OpNet
net = Tensor -> IO Tensor
T.detach Tensor
loss
  where
    predY :: Tensor
predY = OpNet -> Tensor -> Tensor
forward OpNet
net Tensor
trueX
    loss :: Tensor
loss  = Reduction -> Tensor -> Tensor -> Tensor
T.l1Loss Reduction
T.ReduceMean Tensor
trueY Tensor
predY

-- | Run through all Batches performing an update for each
validationEpoch :: ProgressBar s -> [T.Tensor] -> [T.Tensor] -> OpNet 
                -> [T.Tensor] -> IO T.Tensor
validationEpoch :: ProgressBar s
-> [Tensor] -> [Tensor] -> OpNet -> [Tensor] -> IO Tensor
validationEpoch ProgressBar s
_     [Tensor]
_      []   OpNet
_   [Tensor]
losses = Tensor -> IO Tensor
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Tensor -> IO Tensor)
-> ([Tensor] -> Tensor) -> [Tensor] -> IO Tensor
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Dim -> [Tensor] -> Tensor
T.cat (Int -> Dim
T.Dim Int
0) 
                                             ([Tensor] -> Tensor)
-> ([Tensor] -> [Tensor]) -> [Tensor] -> Tensor
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Tensor -> Tensor) -> [Tensor] -> [Tensor]
forall a b. (a -> b) -> [a] -> [b]
map ([Int] -> Tensor -> Tensor
T.reshape [-Int
1]) ([Tensor] -> IO Tensor) -> [Tensor] -> IO Tensor
forall a b. (a -> b) -> a -> b
$ [Tensor]
losses
validationEpoch ProgressBar s
_     []     [Tensor]
_    OpNet
_   [Tensor]
losses = Tensor -> IO Tensor
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Tensor -> IO Tensor)
-> ([Tensor] -> Tensor) -> [Tensor] -> IO Tensor
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Dim -> [Tensor] -> Tensor
T.cat (Int -> Dim
T.Dim Int
0) 
                                             ([Tensor] -> Tensor)
-> ([Tensor] -> [Tensor]) -> [Tensor] -> Tensor
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Tensor -> Tensor) -> [Tensor] -> [Tensor]
forall a b. (a -> b) -> [a] -> [b]
map ([Int] -> Tensor -> Tensor
T.reshape [-Int
1]) ([Tensor] -> IO Tensor) -> [Tensor] -> IO Tensor
forall a b. (a -> b) -> a -> b
$ [Tensor]
losses
validationEpoch ProgressBar s
bar (Tensor
x:[Tensor]
xs) (Tensor
y:[Tensor]
ys) OpNet
net [Tensor]
losses = do
        ProgressBar s -> Int -> IO ()
forall s. ProgressBar s -> Int -> IO ()
incProgress ProgressBar s
bar Int
1
        Tensor -> Tensor -> OpNet -> IO Tensor
validStep Tensor
x Tensor
y OpNet
net IO Tensor -> (Tensor -> IO Tensor) -> IO Tensor
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= 
                ProgressBar s
-> [Tensor] -> [Tensor] -> OpNet -> [Tensor] -> IO Tensor
forall s.
ProgressBar s
-> [Tensor] -> [Tensor] -> OpNet -> [Tensor] -> IO Tensor
validationEpoch ProgressBar s
bar [Tensor]
xs [Tensor]
ys OpNet
net ([Tensor] -> IO Tensor)
-> (Tensor -> [Tensor]) -> Tensor -> IO Tensor
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Tensor -> [Tensor] -> [Tensor]
forall a. a -> [a] -> [a]
:[Tensor]
losses) 

------------------------------------------------------------------------------
-- Running
------------------------------------------------------------------------------

-- | Run Training and Validation for a given number of Epochs
runEpochs :: FilePath -> Int -> Int -> T.Tensor -> T.Tensor -> T.Tensor -> T.Tensor
          -> OpNet -> T.Adam -> IO (OpNet, T.Adam)
runEpochs :: FilePath
-> Int
-> Int
-> Tensor
-> Tensor
-> Tensor
-> Tensor
-> OpNet
-> Adam
-> IO (OpNet, Adam)
runEpochs FilePath
path Int
0     Int
_  Tensor
_      Tensor
_      Tensor
_      Tensor
_      OpNet
net Adam
opt = do
    FilePath -> OpNet -> Adam -> IO ()
saveCheckPoint FilePath
path OpNet
net Adam
opt
    (OpNet, Adam) -> IO (OpNet, Adam)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (OpNet
net, Adam
opt)
runEpochs FilePath
path Int
epoch Int
bs Tensor
trainX Tensor
validX Tensor
trainY Tensor
validY OpNet
net Adam
opt = do
    ([Tensor]
trainXs, [Tensor]
trainYs) <- Int -> Tensor -> Tensor -> IO ([Tensor], [Tensor])
shuffleBatches Int
bs Tensor
trainX Tensor
trainY
    ([Tensor]
validXs, [Tensor]
validYs) <- Int -> Tensor -> Tensor -> IO ([Tensor], [Tensor])
shuffleBatches Int
bs Tensor
trainX Tensor
trainY

    ProgressBar ()
tBar <- Style () -> Double -> Progress () -> IO (ProgressBar ())
forall s. Style s -> Double -> Progress s -> IO (ProgressBar s)
newProgressBar (Int -> Style ()
forall s. Int -> Style s
trainStyle Int
epoch) Double
10 (Int -> Int -> () -> Progress ()
forall s. Int -> Int -> s -> Progress s
Progress Int
0 ([Tensor] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Tensor]
trainXs) ())
    (OpNet
net', Adam
opt', Tensor
lss) <- ProgressBar ()
-> [Tensor]
-> [Tensor]
-> [Tensor]
-> OpNet
-> Adam
-> IO (OpNet, Adam, Tensor)
forall s.
ProgressBar s
-> [Tensor]
-> [Tensor]
-> [Tensor]
-> OpNet
-> Adam
-> IO (OpNet, Adam, Tensor)
trainingEpoch ProgressBar ()
tBar [Tensor]
trainXs [Tensor]
trainYs [] OpNet
net Adam
opt

    FilePath -> IO ()
putStrLn (FilePath -> IO ()) -> FilePath -> IO ()
forall a b. (a -> b) -> a -> b
$ FilePath
"\tTraining Loss: " FilePath -> FilePath -> FilePath
forall a. [a] -> [a] -> [a]
++ Float -> FilePath
forall a. Show a => a -> FilePath
show (Tensor -> Float
forall a. TensorLike a => Tensor -> a
T.asValue (Tensor -> Float) -> Tensor -> Float
forall a b. (a -> b) -> a -> b
$ Tensor -> Tensor
T.mean Tensor
lss :: Float)

    ProgressBar ()
vBar <- Style () -> Double -> Progress () -> IO (ProgressBar ())
forall s. Style s -> Double -> Progress s -> IO (ProgressBar s)
newProgressBar (Int -> Style ()
forall s. Int -> Style s
validStyle Int
epoch) Double
10 (Int -> Int -> () -> Progress ()
forall s. Int -> Int -> s -> Progress s
Progress Int
0 ([Tensor] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Tensor]
validXs) ())
    Tensor
mae  <- ProgressBar ()
-> [Tensor] -> [Tensor] -> OpNet -> [Tensor] -> IO Tensor
forall s.
ProgressBar s
-> [Tensor] -> [Tensor] -> OpNet -> [Tensor] -> IO Tensor
validationEpoch ProgressBar ()
vBar [Tensor]
validXs [Tensor]
validYs OpNet
net' []
    
    FilePath -> IO ()
putStrLn (FilePath -> IO ()) -> FilePath -> IO ()
forall a b. (a -> b) -> a -> b
$ FilePath
"\tValidataion Loss: " FilePath -> FilePath -> FilePath
forall a. [a] -> [a] -> [a]
++ Float -> FilePath
forall a. Show a => a -> FilePath
show (Tensor -> Float
forall a. TensorLike a => Tensor -> a
T.asValue (Tensor -> Float) -> Tensor -> Float
forall a b. (a -> b) -> a -> b
$ Tensor -> Tensor
T.mean Tensor
mae :: Float)

    FilePath -> OpNet -> Adam -> IO ()
saveCheckPoint FilePath
path OpNet
net' Adam
opt'

    FilePath
-> Int
-> Int
-> Tensor
-> Tensor
-> Tensor
-> Tensor
-> OpNet
-> Adam
-> IO (OpNet, Adam)
runEpochs FilePath
path Int
epoch' Int
bs Tensor
trainX Tensor
validX Tensor
trainY Tensor
validY OpNet
net' Adam
opt'
  where
    epoch' :: Int
epoch' = Int
epoch Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1

-- | Initiate Training Run for given Args
run :: Args -> IO ()
run :: Args -> IO ()
run args :: Args
args@Args{Bool
Int
FilePath
Device
PDK
exp :: Args -> Bool
size :: Args -> Int
reg :: Args -> Int
num :: Args -> Int
dir :: Args -> FilePath
dev :: Args -> Device
pdk :: Args -> PDK
exp :: Bool
size :: Int
reg :: Int
num :: Int
dir :: FilePath
dev :: Device
pdk :: PDK
..} | Bool
exp       = Args -> IO ()
runExp Args
args
                  | Bool
otherwise = Args -> IO ()
runBase Args
args

-- | Experimental Training Run
runExp :: Args -> IO ()
runExp :: Args -> IO ()
runExp Args{Bool
Int
FilePath
Device
PDK
exp :: Bool
size :: Int
reg :: Int
num :: Int
dir :: FilePath
dev :: Device
pdk :: PDK
exp :: Args -> Bool
size :: Args -> Int
reg :: Args -> Int
num :: Args -> Int
dir :: Args -> FilePath
dev :: Args -> Device
pdk :: Args -> PDK
..} = do
    FilePath -> IO ()
putStrLn (FilePath -> IO ()) -> FilePath -> IO ()
forall a b. (a -> b) -> a -> b
$ FilePath
"Training " FilePath -> FilePath -> FilePath
forall a. [a] -> [a] -> [a]
++ Device -> FilePath
forall a. Show a => a -> FilePath
show Device
dev FilePath -> FilePath -> FilePath
forall a. [a] -> [a] -> [a]
++ FilePath
" Model in " FilePath -> FilePath -> FilePath
forall a. [a] -> [a] -> [a]
++ PDK -> FilePath
forall a. Show a => a -> FilePath
show PDK
pdk FilePath -> FilePath -> FilePath
forall a. [a] -> [a] -> [a]
++ FilePath
"."

    DataFrame Tensor
df' <- FilePath -> IO (DataFrame Tensor)
DF.fromFile FilePath
dir

    let cls' :: [FilePath]
cls'   = [FilePath
"id", FilePath
"gm", FilePath
"gds", FilePath
"gmbs", FilePath
"cgs", FilePath
"cgd", FilePath
"cgb", FilePath
"csd", FilePath
"cbs", FilePath
"cbd"]
        cls'' :: [FilePath]
cls''  = [FilePath
"idoverW", FilePath
"gmoverW", FilePath
"gdsoverW", FilePath
"gmbsoverW", FilePath
"cgsoverW", FilePath
"cgdoverW", FilePath
"cgboverW", FilePath
"csdoverW", FilePath
"cbsoverW", FilePath
"cbdoverW"]
        dfNorm :: DataFrame Tensor
dfNorm = [FilePath] -> Tensor -> DataFrame Tensor
forall a. [FilePath] -> a -> DataFrame a
DF.DataFrame [FilePath]
cls'' (Tensor -> DataFrame Tensor) -> Tensor -> DataFrame Tensor
forall a b. (a -> b) -> a -> b
$ Tensor -> Tensor -> Tensor
T.div (DataFrame Tensor -> Tensor
forall a. DataFrame a -> a
DF.values (DataFrame Tensor -> Tensor) -> DataFrame Tensor -> Tensor
forall a b. (a -> b) -> a -> b
$ [FilePath] -> DataFrame Tensor -> DataFrame Tensor
DF.lookup [FilePath]
cls' DataFrame Tensor
df') (DataFrame Tensor
df' DataFrame Tensor -> FilePath -> Tensor
?? FilePath
"W")
        dfRaw :: DataFrame Tensor
dfRaw  = DataFrame Tensor -> DataFrame Tensor -> DataFrame Tensor
DF.union DataFrame Tensor
df' DataFrame Tensor
dfNorm

    FilePath
modelPath <- FilePath -> FilePath -> IO FilePath
createModelDir FilePath
pdk' FilePath
dev'
    let tracePath :: FilePath
tracePath = FilePath
modelPath FilePath -> FilePath -> FilePath
forall a. [a] -> [a] -> [a]
++ FilePath
"/" FilePath -> FilePath -> FilePath
forall a. [a] -> [a] -> [a]
++ FilePath
dev' FilePath -> FilePath -> FilePath
forall a. [a] -> [a] -> [a]
++ FilePath
"-exp.pt"

    let dfRaw' :: DataFrame Tensor
dfRaw' = DataFrame Tensor -> DataFrame Tensor
DF.dropNan (DataFrame Tensor -> DataFrame Tensor)
-> DataFrame Tensor -> DataFrame Tensor
forall a b. (a -> b) -> a -> b
$ [FilePath] -> DataFrame Tensor -> DataFrame Tensor
DF.lookup [FilePath]
cols DataFrame Tensor
dfRaw
        dfReg :: DataFrame Tensor
dfReg  = Tensor -> DataFrame Tensor -> DataFrame Tensor
DF.rowFilter (Int -> DataFrame Tensor -> Tensor
region (-Int
1) DataFrame Tensor
dfRaw') DataFrame Tensor
dfRaw'

    DataFrame Tensor
dfShuff    <- DataFrame Tensor -> IO (DataFrame Tensor)
DF.shuffleIO DataFrame Tensor
dfReg

    let dfT :: DataFrame Tensor
dfT    = DataFrame Tensor -> DataFrame Tensor
DF.dropNan 
               (DataFrame Tensor -> DataFrame Tensor)
-> DataFrame Tensor -> DataFrame Tensor
forall a b. (a -> b) -> a -> b
$ DataFrame Tensor -> DataFrame Tensor -> DataFrame Tensor
DF.union (Tensor -> Tensor -> Tensor
trafo Tensor
maskX (Tensor -> Tensor) -> DataFrame Tensor -> DataFrame Tensor
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [FilePath] -> DataFrame Tensor -> DataFrame Tensor
DF.lookup [FilePath]
paramsX DataFrame Tensor
dfShuff)
                          (Tensor -> Tensor -> Tensor
trafo Tensor
maskY (Tensor -> Tensor) -> DataFrame Tensor -> DataFrame Tensor
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [FilePath] -> DataFrame Tensor -> DataFrame Tensor
DF.lookup [FilePath]
paramsY DataFrame Tensor
dfShuff)
        dfX' :: DataFrame Tensor
dfX'   = [FilePath] -> DataFrame Tensor -> DataFrame Tensor
DF.lookup [FilePath]
paramsX DataFrame Tensor
dfT
        dfY' :: DataFrame Tensor
dfY'   = [FilePath] -> DataFrame Tensor -> DataFrame Tensor
DF.lookup [FilePath]
paramsY DataFrame Tensor
dfT

    let minX :: Tensor
minX   = (Tensor, Tensor) -> Tensor
forall a b. (a, b) -> a
fst ((Tensor, Tensor) -> Tensor)
-> (DataFrame Tensor -> (Tensor, Tensor))
-> DataFrame Tensor
-> Tensor
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Dim -> KeepDim -> Tensor -> (Tensor, Tensor)
T.minDim (Int -> Dim
T.Dim Int
0) KeepDim
T.RemoveDim (Tensor -> (Tensor, Tensor))
-> (DataFrame Tensor -> Tensor)
-> DataFrame Tensor
-> (Tensor, Tensor)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DataFrame Tensor -> Tensor
forall a. DataFrame a -> a
values (DataFrame Tensor -> Tensor) -> DataFrame Tensor -> Tensor
forall a b. (a -> b) -> a -> b
$ DataFrame Tensor
dfX'
        maxX :: Tensor
maxX   = (Tensor, Tensor) -> Tensor
forall a b. (a, b) -> a
fst ((Tensor, Tensor) -> Tensor)
-> (DataFrame Tensor -> (Tensor, Tensor))
-> DataFrame Tensor
-> Tensor
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Dim -> KeepDim -> Tensor -> (Tensor, Tensor)
T.maxDim (Int -> Dim
T.Dim Int
0) KeepDim
T.RemoveDim (Tensor -> (Tensor, Tensor))
-> (DataFrame Tensor -> Tensor)
-> DataFrame Tensor
-> (Tensor, Tensor)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DataFrame Tensor -> Tensor
forall a. DataFrame a -> a
values (DataFrame Tensor -> Tensor) -> DataFrame Tensor -> Tensor
forall a b. (a -> b) -> a -> b
$ DataFrame Tensor
dfX'
        minY :: Tensor
minY   = (Tensor, Tensor) -> Tensor
forall a b. (a, b) -> a
fst ((Tensor, Tensor) -> Tensor)
-> (DataFrame Tensor -> (Tensor, Tensor))
-> DataFrame Tensor
-> Tensor
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Dim -> KeepDim -> Tensor -> (Tensor, Tensor)
T.minDim (Int -> Dim
T.Dim Int
0) KeepDim
T.RemoveDim (Tensor -> (Tensor, Tensor))
-> (DataFrame Tensor -> Tensor)
-> DataFrame Tensor
-> (Tensor, Tensor)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DataFrame Tensor -> Tensor
forall a. DataFrame a -> a
values (DataFrame Tensor -> Tensor) -> DataFrame Tensor -> Tensor
forall a b. (a -> b) -> a -> b
$ DataFrame Tensor
dfY'
        maxY :: Tensor
maxY   = (Tensor, Tensor) -> Tensor
forall a b. (a, b) -> a
fst ((Tensor, Tensor) -> Tensor)
-> (DataFrame Tensor -> (Tensor, Tensor))
-> DataFrame Tensor
-> Tensor
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Dim -> KeepDim -> Tensor -> (Tensor, Tensor)
T.maxDim (Int -> Dim
T.Dim Int
0) KeepDim
T.RemoveDim (Tensor -> (Tensor, Tensor))
-> (DataFrame Tensor -> Tensor)
-> DataFrame Tensor
-> (Tensor, Tensor)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DataFrame Tensor -> Tensor
forall a. DataFrame a -> a
values (DataFrame Tensor -> Tensor) -> DataFrame Tensor -> Tensor
forall a b. (a -> b) -> a -> b
$ DataFrame Tensor
dfY'

    let !dfX :: DataFrame Tensor
dfX   = Tensor -> Tensor -> Tensor -> Tensor
scale Tensor
minX Tensor
maxX (Tensor -> Tensor) -> DataFrame Tensor -> DataFrame Tensor
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> DataFrame Tensor
dfX'
        !dfY :: DataFrame Tensor
dfY   = Tensor -> Tensor -> Tensor -> Tensor
scale Tensor
minY Tensor
maxY (Tensor -> Tensor) -> DataFrame Tensor -> DataFrame Tensor
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> DataFrame Tensor
dfY'
        !df :: DataFrame Tensor
df    = DataFrame Tensor -> DataFrame Tensor
DF.dropNan (DataFrame Tensor -> DataFrame Tensor)
-> DataFrame Tensor -> DataFrame Tensor
forall a b. (a -> b) -> a -> b
$ DataFrame Tensor -> DataFrame Tensor -> DataFrame Tensor
DF.union DataFrame Tensor
dfX DataFrame Tensor
dfY

    OpNet
net        <- Device -> OpNet -> OpNet
forall a. HasTypes a Tensor => Device -> a -> a
T.toDevice Device
T.gpu (OpNet -> OpNet) -> IO OpNet -> IO OpNet
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> OpNetSpec -> IO OpNet
forall spec f. Randomizable spec f => spec -> IO f
T.sample (Int -> Int -> OpNetSpec
OpNetSpec Int
numInputs Int
numOutputs)
    let opt :: Adam
opt    =  Int -> Float -> Float -> [Parameter] -> Adam
T.mkAdam Int
0 Float
β1 Float
β2 ([Parameter] -> Adam) -> [Parameter] -> Adam
forall a b. (a -> b) -> a -> b
$ OpNet -> [Parameter]
forall f. Parameterized f => f -> [Parameter]
NN.flattenParameters OpNet
net

    let (!Tensor
trainX', !Tensor
validX', !Tensor
trainY', !Tensor
validY') = 
                [FilePath]
-> [FilePath]
-> Float
-> DataFrame Tensor
-> (Tensor, Tensor, Tensor, Tensor)
trainTestSplit [FilePath]
paramsX [FilePath]
paramsY Float
testSplit DataFrame Tensor
df

    (OpNet
net', Adam
opt') <- FilePath
-> Int
-> Int
-> Tensor
-> Tensor
-> Tensor
-> Tensor
-> OpNet
-> Adam
-> IO (OpNet, Adam)
runEpochs FilePath
modelPath Int
num Int
size Tensor
trainX' Tensor
validX' Tensor
trainY' Tensor
validY' OpNet
net Adam
opt

    FilePath -> OpNet -> Adam -> IO ()
saveCheckPoint FilePath
modelPath OpNet
net' Adam
opt'

    -- !net'' <- T.toDevice T.cpu <$> noGrad net'
    !OpNet
net'' <- FilePath -> OpNetSpec -> Int -> IO (OpNet, Adam)
loadCheckPoint FilePath
modelPath (Int -> Int -> OpNetSpec
OpNetSpec Int
numInputs Int
numOutputs) Int
num 
                    IO (OpNet, Adam) -> ((OpNet, Adam) -> IO OpNet) -> IO OpNet
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= OpNet -> IO OpNet
forall f. Parameterized f => f -> IO f
noGrad (OpNet -> IO OpNet)
-> ((OpNet, Adam) -> OpNet) -> (OpNet, Adam) -> IO OpNet
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (OpNet, Adam) -> OpNet
forall a b. (a, b) -> a
fst

    let predict :: Tensor -> Tensor
predict = Tensor -> Tensor -> Tensor
neg Tensor
maskNeg (Tensor -> Tensor) -> (Tensor -> Tensor) -> Tensor -> Tensor
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Tensor -> Tensor -> Tensor
trafo' Tensor
maskY (Tensor -> Tensor) -> (Tensor -> Tensor) -> Tensor -> Tensor
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Tensor -> Tensor -> Tensor -> Tensor
scale' Tensor
minY Tensor
maxY
                (Tensor -> Tensor) -> (Tensor -> Tensor) -> Tensor -> Tensor
forall b c a. (b -> c) -> (a -> b) -> a -> c
. OpNet -> Tensor -> Tensor
forward OpNet
net''
                (Tensor -> Tensor) -> (Tensor -> Tensor) -> Tensor -> Tensor
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Tensor -> Tensor -> Tensor -> Tensor
scale Tensor
minX Tensor
maxX (Tensor -> Tensor) -> (Tensor -> Tensor) -> Tensor -> Tensor
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Tensor -> Tensor -> Tensor
trafo Tensor
maskX

    Device
-> PDK
-> Int
-> [FilePath]
-> [FilePath]
-> (Tensor -> Tensor)
-> IO ScriptModule
traceModel Device
dev PDK
pdk Int
numInputs [FilePath]
paramsX [FilePath]
paramsY Tensor -> Tensor
predict IO ScriptModule -> (ScriptModule -> IO ()) -> IO ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= FilePath -> ScriptModule -> IO ()
saveInferenceModel FilePath
tracePath

    Tensor -> Tensor
mdl'' <- ScriptModule -> Tensor -> Tensor
unTraceModel (ScriptModule -> Tensor -> Tensor)
-> IO ScriptModule -> IO (Tensor -> Tensor)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> LoadMode -> FilePath -> IO ScriptModule
T.loadScript LoadMode
T.WithoutRequiredGrad FilePath
tracePath

    FilePath -> IO ()
putStrLn (FilePath -> IO ()) -> FilePath -> IO ()
forall a b. (a -> b) -> a -> b
$ FilePath
"Final Checkpoint saved at: " FilePath -> FilePath -> FilePath
forall a. [a] -> [a] -> [a]
++ FilePath
modelPath
    FilePath -> IO ()
putStrLn (FilePath -> IO ()) -> FilePath -> IO ()
forall a b. (a -> b) -> a -> b
$ FilePath
"Traced Model saved at: " FilePath -> FilePath -> FilePath
forall a. [a] -> [a] -> [a]
++ FilePath
tracePath
    
    DataFrame Tensor
-> [FilePath] -> [FilePath] -> (Tensor -> Tensor) -> IO ()
testModelExp DataFrame Tensor
dfRaw' [FilePath]
paramsX [FilePath]
paramsY Tensor -> Tensor
mdl''
  where
    pdk' :: FilePath
pdk'       = PDK -> FilePath
forall a. Show a => a -> FilePath
show PDK
pdk
    dev' :: FilePath
dev'       = Device -> FilePath
forall a. Show a => a -> FilePath
show Device
dev
    testSplit :: Float
testSplit  = Float
0.75
    paramsX :: [FilePath]
paramsX    = [FilePath
"L", FilePath
"vds", FilePath
"vgs", FilePath
"vbs"]
    paramsY :: [FilePath]
paramsY    = [ FilePath
"fug", FilePath
"gmoverid"
                 , FilePath
"idoverW", FilePath
"gmoverW", FilePath
"gdsoverW", FilePath
"gmbsoverW"
                 , FilePath
"vdsat", FilePath
"vth"
                 , FilePath
"cgsoverW", FilePath
"cgdoverW", FilePath
"cgboverW", FilePath
"csdoverW", FilePath
"cbsoverW", FilePath
"cbdoverW" ]
    cols :: [FilePath]
cols       = [FilePath]
paramsX [FilePath] -> [FilePath] -> [FilePath]
forall a. [a] -> [a] -> [a]
++ [FilePath]
paramsY [FilePath] -> [FilePath] -> [FilePath]
forall a. [a] -> [a] -> [a]
++ [FilePath
"region"]
    maskX :: Tensor
maskX      = Device -> Tensor -> Tensor
forall a. HasTypes a Tensor => Device -> a -> a
T.toDevice Device
T.cpu (Tensor -> Tensor) -> Tensor -> Tensor
forall a b. (a -> b) -> a -> b
$ [FilePath] -> [FilePath] -> Tensor
boolMask' [FilePath
"L"] [FilePath]
paramsX
    maskY :: Tensor
maskY      = Device -> Tensor -> Tensor
forall a. HasTypes a Tensor => Device -> a -> a
T.toDevice Device
T.cpu (Tensor -> Tensor) -> Tensor -> Tensor
forall a b. (a -> b) -> a -> b
$ [FilePath] -> [FilePath] -> Tensor
boolMask' [ FilePath
"idoverW", FilePath
"fug", FilePath
"gmoverW", FilePath
"gdsoverW", FilePath
"gmbsoverW"
                                              , FilePath
"cgsoverW", FilePath
"cgdoverW", FilePath
"cgboverW", FilePath
"csdoverW", FilePath
"cbsoverW", FilePath
"cbdoverW"
                                              ] [FilePath]
paramsY
    neg :: Tensor -> Tensor -> Tensor
neg Tensor
m Tensor
t    = Tensor -> Tensor -> Tensor
T.add (Tensor -> Tensor
T.negative (Tensor -> Tensor -> Tensor
T.mul Tensor
m Tensor
t)) (Tensor -> Tensor -> Tensor
T.mul (Tensor
1.0 Tensor -> Tensor -> Tensor
forall a. Num a => a -> a -> a
- Tensor
m) Tensor
t)
    !maskNeg' :: [FilePath]
maskNeg'   = if Device
dev Device -> Device -> Bool
forall a. Eq a => a -> a -> Bool
== Device
NMOS then [FilePath
"cgsoverW", FilePath
"cgdoverW", FilePath
"cgboverW", FilePath
"csdoverW", FilePath
"cbsoverW", FilePath
"cbdoverW"]
                                 else [FilePath
"id",FilePath
"cgsoverW", FilePath
"cgdoverW", FilePath
"cgboverW", FilePath
"csdoverW", FilePath
"cbsoverW", FilePath
"cbdoverW"]
    !maskNeg :: Tensor
maskNeg    = DType -> Tensor -> Tensor
T.toDType DType
T.Float (Tensor -> Tensor) -> (Tensor -> Tensor) -> Tensor -> Tensor
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Device -> Tensor -> Tensor
forall a. HasTypes a Tensor => Device -> a -> a
T.toDevice Device
T.cpu (Tensor -> Tensor) -> Tensor -> Tensor
forall a b. (a -> b) -> a -> b
$ [FilePath] -> [FilePath] -> Tensor
boolMask' [FilePath]
maskNeg' [FilePath]
paramsY
    numInputs :: Int
numInputs  = [FilePath] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [FilePath]
paramsX
    numOutputs :: Int
numOutputs = [FilePath] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [FilePath]
paramsY

-- | Basic Training Run
runBase :: Args -> IO ()
runBase :: Args -> IO ()
runBase Args{Bool
Int
FilePath
Device
PDK
exp :: Bool
size :: Int
reg :: Int
num :: Int
dir :: FilePath
dev :: Device
pdk :: PDK
exp :: Args -> Bool
size :: Args -> Int
reg :: Args -> Int
num :: Args -> Int
dir :: Args -> FilePath
dev :: Args -> Device
pdk :: Args -> PDK
..} = do
    FilePath -> IO ()
putStrLn (FilePath -> IO ()) -> FilePath -> IO ()
forall a b. (a -> b) -> a -> b
$ FilePath
"Training " FilePath -> FilePath -> FilePath
forall a. [a] -> [a] -> [a]
++ Device -> FilePath
forall a. Show a => a -> FilePath
show Device
dev FilePath -> FilePath -> FilePath
forall a. [a] -> [a] -> [a]
++ FilePath
" Model in " FilePath -> FilePath -> FilePath
forall a. [a] -> [a] -> [a]
++ PDK -> FilePath
forall a. Show a => a -> FilePath
show PDK
pdk FilePath -> FilePath -> FilePath
forall a. [a] -> [a] -> [a]
++ FilePath
"."

    DataFrame Tensor
df' <- FilePath -> IO (DataFrame Tensor)
DF.fromFile FilePath
dir

    let idoverW :: Tensor
idoverW = Tensor -> Tensor -> Tensor
T.div (DataFrame Tensor
df' DataFrame Tensor -> FilePath -> Tensor
?? FilePath
"id") (DataFrame Tensor
df' DataFrame Tensor -> FilePath -> Tensor
?? FilePath
"W")
        dfRaw :: DataFrame Tensor
dfRaw   = [FilePath] -> Tensor -> DataFrame Tensor -> DataFrame Tensor
DF.insert [FilePath
"idoverW"] Tensor
idoverW DataFrame Tensor
df'

    FilePath
modelPath <- FilePath -> FilePath -> IO FilePath
createModelDir FilePath
pdk' FilePath
dev'
    let tracePath :: FilePath
tracePath = FilePath
modelPath FilePath -> FilePath -> FilePath
forall a. [a] -> [a] -> [a]
++ FilePath
"/" FilePath -> FilePath -> FilePath
forall a. [a] -> [a] -> [a]
++ FilePath
dev' FilePath -> FilePath -> FilePath
forall a. [a] -> [a] -> [a]
++ FilePath
"-" FilePath -> FilePath -> FilePath
forall a. [a] -> [a] -> [a]
++ FilePath
reg' FilePath -> FilePath -> FilePath
forall a. [a] -> [a] -> [a]
++ FilePath
".pt"

    let dfRaw' :: DataFrame Tensor
dfRaw' = DataFrame Tensor -> DataFrame Tensor
DF.dropNan (DataFrame Tensor -> DataFrame Tensor)
-> DataFrame Tensor -> DataFrame Tensor
forall a b. (a -> b) -> a -> b
$ [FilePath] -> DataFrame Tensor -> DataFrame Tensor
DF.lookup [FilePath]
cols DataFrame Tensor
dfRaw
        dfReg :: DataFrame Tensor
dfReg  = Tensor -> DataFrame Tensor -> DataFrame Tensor
DF.rowFilter (Int -> DataFrame Tensor -> Tensor
region Int
reg DataFrame Tensor
dfRaw') DataFrame Tensor
dfRaw'

    DataFrame Tensor
dfShuff    <- DataFrame Tensor -> IO (DataFrame Tensor)
DF.shuffleIO DataFrame Tensor
dfReg

    let dfT :: DataFrame Tensor
dfT    = DataFrame Tensor -> DataFrame Tensor
DF.dropNan 
               (DataFrame Tensor -> DataFrame Tensor)
-> DataFrame Tensor -> DataFrame Tensor
forall a b. (a -> b) -> a -> b
$ DataFrame Tensor -> DataFrame Tensor -> DataFrame Tensor
DF.union (Tensor -> Tensor -> Tensor
trafo Tensor
maskX (Tensor -> Tensor) -> DataFrame Tensor -> DataFrame Tensor
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [FilePath] -> DataFrame Tensor -> DataFrame Tensor
DF.lookup [FilePath]
paramsX DataFrame Tensor
dfShuff)
                          (Tensor -> Tensor -> Tensor
trafo Tensor
maskY (Tensor -> Tensor) -> DataFrame Tensor -> DataFrame Tensor
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [FilePath] -> DataFrame Tensor -> DataFrame Tensor
DF.lookup [FilePath]
paramsY DataFrame Tensor
dfShuff)
        dfX' :: DataFrame Tensor
dfX'   = [FilePath] -> DataFrame Tensor -> DataFrame Tensor
DF.lookup [FilePath]
paramsX DataFrame Tensor
dfT
        dfY' :: DataFrame Tensor
dfY'   = [FilePath] -> DataFrame Tensor -> DataFrame Tensor
DF.lookup [FilePath]
paramsY DataFrame Tensor
dfT

    let minX :: Tensor
minX   = (Tensor, Tensor) -> Tensor
forall a b. (a, b) -> a
fst ((Tensor, Tensor) -> Tensor)
-> (DataFrame Tensor -> (Tensor, Tensor))
-> DataFrame Tensor
-> Tensor
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Dim -> KeepDim -> Tensor -> (Tensor, Tensor)
T.minDim (Int -> Dim
T.Dim Int
0) KeepDim
T.RemoveDim (Tensor -> (Tensor, Tensor))
-> (DataFrame Tensor -> Tensor)
-> DataFrame Tensor
-> (Tensor, Tensor)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DataFrame Tensor -> Tensor
forall a. DataFrame a -> a
values (DataFrame Tensor -> Tensor) -> DataFrame Tensor -> Tensor
forall a b. (a -> b) -> a -> b
$ DataFrame Tensor
dfX'
        maxX :: Tensor
maxX   = (Tensor, Tensor) -> Tensor
forall a b. (a, b) -> a
fst ((Tensor, Tensor) -> Tensor)
-> (DataFrame Tensor -> (Tensor, Tensor))
-> DataFrame Tensor
-> Tensor
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Dim -> KeepDim -> Tensor -> (Tensor, Tensor)
T.maxDim (Int -> Dim
T.Dim Int
0) KeepDim
T.RemoveDim (Tensor -> (Tensor, Tensor))
-> (DataFrame Tensor -> Tensor)
-> DataFrame Tensor
-> (Tensor, Tensor)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DataFrame Tensor -> Tensor
forall a. DataFrame a -> a
values (DataFrame Tensor -> Tensor) -> DataFrame Tensor -> Tensor
forall a b. (a -> b) -> a -> b
$ DataFrame Tensor
dfX'
        minY :: Tensor
minY   = (Tensor, Tensor) -> Tensor
forall a b. (a, b) -> a
fst ((Tensor, Tensor) -> Tensor)
-> (DataFrame Tensor -> (Tensor, Tensor))
-> DataFrame Tensor
-> Tensor
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Dim -> KeepDim -> Tensor -> (Tensor, Tensor)
T.minDim (Int -> Dim
T.Dim Int
0) KeepDim
T.RemoveDim (Tensor -> (Tensor, Tensor))
-> (DataFrame Tensor -> Tensor)
-> DataFrame Tensor
-> (Tensor, Tensor)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DataFrame Tensor -> Tensor
forall a. DataFrame a -> a
values (DataFrame Tensor -> Tensor) -> DataFrame Tensor -> Tensor
forall a b. (a -> b) -> a -> b
$ DataFrame Tensor
dfY'
        maxY :: Tensor
maxY   = (Tensor, Tensor) -> Tensor
forall a b. (a, b) -> a
fst ((Tensor, Tensor) -> Tensor)
-> (DataFrame Tensor -> (Tensor, Tensor))
-> DataFrame Tensor
-> Tensor
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Dim -> KeepDim -> Tensor -> (Tensor, Tensor)
T.maxDim (Int -> Dim
T.Dim Int
0) KeepDim
T.RemoveDim (Tensor -> (Tensor, Tensor))
-> (DataFrame Tensor -> Tensor)
-> DataFrame Tensor
-> (Tensor, Tensor)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DataFrame Tensor -> Tensor
forall a. DataFrame a -> a
values (DataFrame Tensor -> Tensor) -> DataFrame Tensor -> Tensor
forall a b. (a -> b) -> a -> b
$ DataFrame Tensor
dfY'

    let !dfX :: DataFrame Tensor
dfX   = Tensor -> Tensor -> Tensor -> Tensor
scale Tensor
minX Tensor
maxX (Tensor -> Tensor) -> DataFrame Tensor -> DataFrame Tensor
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> DataFrame Tensor
dfX'
        !dfY :: DataFrame Tensor
dfY   = Tensor -> Tensor -> Tensor -> Tensor
scale Tensor
minY Tensor
maxY (Tensor -> Tensor) -> DataFrame Tensor -> DataFrame Tensor
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> DataFrame Tensor
dfY'
        !df :: DataFrame Tensor
df    = DataFrame Tensor -> DataFrame Tensor
DF.dropNan (DataFrame Tensor -> DataFrame Tensor)
-> DataFrame Tensor -> DataFrame Tensor
forall a b. (a -> b) -> a -> b
$ DataFrame Tensor -> DataFrame Tensor -> DataFrame Tensor
DF.union DataFrame Tensor
dfX DataFrame Tensor
dfY

    OpNet
net        <- Device -> OpNet -> OpNet
forall a. HasTypes a Tensor => Device -> a -> a
T.toDevice Device
T.gpu (OpNet -> OpNet) -> IO OpNet -> IO OpNet
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> OpNetSpec -> IO OpNet
forall spec f. Randomizable spec f => spec -> IO f
T.sample (Int -> Int -> OpNetSpec
OpNetSpec Int
numInputs Int
numOutputs)
    let opt :: Adam
opt    =  Int -> Float -> Float -> [Parameter] -> Adam
T.mkAdam Int
0 Float
β1 Float
β2 ([Parameter] -> Adam) -> [Parameter] -> Adam
forall a b. (a -> b) -> a -> b
$ OpNet -> [Parameter]
forall f. Parameterized f => f -> [Parameter]
NN.flattenParameters OpNet
net

    let (!Tensor
trainX', !Tensor
validX', !Tensor
trainY', !Tensor
validY') = 
                [FilePath]
-> [FilePath]
-> Float
-> DataFrame Tensor
-> (Tensor, Tensor, Tensor, Tensor)
trainTestSplit [FilePath]
paramsX [FilePath]
paramsY Float
testSplit DataFrame Tensor
df

    (OpNet
net', Adam
opt') <- FilePath
-> Int
-> Int
-> Tensor
-> Tensor
-> Tensor
-> Tensor
-> OpNet
-> Adam
-> IO (OpNet, Adam)
runEpochs FilePath
modelPath Int
num Int
size Tensor
trainX' Tensor
validX' Tensor
trainY' Tensor
validY' OpNet
net Adam
opt

    FilePath -> OpNet -> Adam -> IO ()
saveCheckPoint FilePath
modelPath OpNet
net' Adam
opt'

    -- !net'' <- T.toDevice T.cpu <$> noGrad net'
    !OpNet
net'' <- FilePath -> OpNetSpec -> Int -> IO (OpNet, Adam)
loadCheckPoint FilePath
modelPath (Int -> Int -> OpNetSpec
OpNetSpec Int
numInputs Int
numOutputs) Int
num 
                    IO (OpNet, Adam) -> ((OpNet, Adam) -> IO OpNet) -> IO OpNet
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= OpNet -> IO OpNet
forall f. Parameterized f => f -> IO f
noGrad (OpNet -> IO OpNet)
-> ((OpNet, Adam) -> OpNet) -> (OpNet, Adam) -> IO OpNet
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (OpNet, Adam) -> OpNet
forall a b. (a, b) -> a
fst

    let predict :: Tensor -> Tensor
predict Tensor
x = Tensor
y
          where
            !i' :: Tensor
i'  = Int -> Int -> Int -> TensorOptions -> Tensor
T.arange Int
1 Int
num' Int
1 (TensorOptions -> Tensor)
-> (TensorOptions -> TensorOptions) -> TensorOptions -> Tensor
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DType -> TensorOptions -> TensorOptions
T.withDType DType
T.Int64 (TensorOptions -> Tensor) -> TensorOptions -> Tensor
forall a b. (a -> b) -> a -> b
$ TensorOptions
T.defaultOpts 
            !x' :: Tensor
x'  = Int -> Tensor -> Tensor -> Tensor
T.indexSelect Int
1 Tensor
i' Tensor
x
            !y' :: Tensor
y'  = Tensor -> Tensor -> Tensor
trafo' Tensor
maskY (Tensor -> Tensor) -> (Tensor -> Tensor) -> Tensor -> Tensor
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Tensor -> Tensor -> Tensor -> Tensor
scale' Tensor
minY Tensor
maxY (Tensor -> Tensor) -> (Tensor -> Tensor) -> Tensor -> Tensor
forall b c a. (b -> c) -> (a -> b) -> a -> c
. OpNet -> Tensor -> Tensor
forward OpNet
net''
                 (Tensor -> Tensor) -> (Tensor -> Tensor) -> Tensor -> Tensor
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Tensor -> Tensor -> Tensor -> Tensor
scale Tensor
minX Tensor
maxX (Tensor -> Tensor) -> (Tensor -> Tensor) -> Tensor -> Tensor
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Tensor -> Tensor -> Tensor
trafo Tensor
maskX (Tensor -> Tensor) -> Tensor -> Tensor
forall a b. (a -> b) -> a -> b
$ Tensor
x'
            !w :: Tensor
w   = [Int] -> Tensor -> Tensor
T.reshape [-Int
1,Int
1] (Tensor -> Tensor) -> Tensor -> Tensor
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Tensor -> Tensor
T.select Int
1 Int
0 Tensor
x Tensor -> Tensor -> Tensor
forall a. Fractional a => a -> a -> a
/ Int -> Int -> Tensor -> Tensor
T.select Int
1 Int
0 Tensor
y'
            !l :: Tensor
l   = [Int] -> Tensor -> Tensor
T.reshape [-Int
1,Int
1] (Tensor -> Tensor) -> Tensor -> Tensor
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Tensor -> Tensor
T.select Int
1 Int
1 Tensor
y'
            !v :: Tensor
v   = [Int] -> Tensor -> Tensor
T.reshape [-Int
1,Int
1] (Tensor -> Tensor) -> Tensor -> Tensor
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Tensor -> Tensor
T.select Int
1 Int
2 Tensor
y'
            !y :: Tensor
y   = Tensor -> Tensor
T.abs (Tensor -> Tensor) -> Tensor -> Tensor
forall a b. (a -> b) -> a -> b
$ Dim -> [Tensor] -> Tensor
T.cat (Int -> Dim
T.Dim Int
1) [ Tensor
w, Tensor
l, Tensor
v]

    Device
-> PDK
-> Int
-> [FilePath]
-> [FilePath]
-> (Tensor -> Tensor)
-> IO ScriptModule
traceModel Device
dev PDK
pdk Int
num' [FilePath]
paramsX [FilePath]
paramsY Tensor -> Tensor
predict IO ScriptModule -> (ScriptModule -> IO ()) -> IO ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= FilePath -> ScriptModule -> IO ()
saveInferenceModel FilePath
tracePath

    FilePath -> IO ()
putStrLn (FilePath -> IO ()) -> FilePath -> IO ()
forall a b. (a -> b) -> a -> b
$ FilePath
"Final Checkpoint saved at: " FilePath -> FilePath -> FilePath
forall a. [a] -> [a] -> [a]
++ FilePath
modelPath
    FilePath -> IO ()
putStrLn (FilePath -> IO ()) -> FilePath -> IO ()
forall a b. (a -> b) -> a -> b
$ FilePath
"Traced Model saved at: " FilePath -> FilePath -> FilePath
forall a. [a] -> [a] -> [a]
++ FilePath
tracePath

    Tensor -> Tensor
mdl'' <- ScriptModule -> Tensor -> Tensor
unTraceModel (ScriptModule -> Tensor -> Tensor)
-> IO ScriptModule -> IO (Tensor -> Tensor)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> LoadMode -> FilePath -> IO ScriptModule
T.loadScript LoadMode
T.WithoutRequiredGrad FilePath
tracePath

    Int
-> DataFrame Tensor
-> [FilePath]
-> [FilePath]
-> (Tensor -> Tensor)
-> IO ()
testModel Int
reg DataFrame Tensor
dfRaw' [FilePath]
paramsX' [FilePath]
paramsY' Tensor -> Tensor
mdl''
  where
    reg' :: FilePath
reg'       = if Int
reg Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
2 then FilePath
"sat" else FilePath
"ulp"
    pdk' :: FilePath
pdk'       = PDK -> FilePath
forall a. Show a => a -> FilePath
show PDK
pdk
    dev' :: FilePath
dev'       = Device -> FilePath
forall a. Show a => a -> FilePath
show Device
dev
    testSplit :: Float
testSplit  = Float
0.75
    paramsX :: [FilePath]
paramsX    = [FilePath
"gmoverid", FilePath
"fug", FilePath
"vds", FilePath
"vbs"]
    paramsY :: [FilePath]
paramsY    = [FilePath
"idoverW", FilePath
"L", FilePath
"vgs"]
    paramsX' :: [FilePath]
paramsX'   = FilePath
"id" FilePath -> [FilePath] -> [FilePath]
forall a. a -> [a] -> [a]
: [FilePath]
paramsX
    paramsY' :: [FilePath]
paramsY'   = [FilePath
"W", FilePath
"L", FilePath
"vgs"]
    cols :: [FilePath]
cols       = [FilePath]
paramsX [FilePath] -> [FilePath] -> [FilePath]
forall a. [a] -> [a] -> [a]
++ [FilePath]
paramsY [FilePath] -> [FilePath] -> [FilePath]
forall a. [a] -> [a] -> [a]
++ [FilePath
"id", FilePath
"region", FilePath
"vth", FilePath
"W"]
    maskX :: Tensor
maskX      = Device -> Tensor -> Tensor
forall a. HasTypes a Tensor => Device -> a -> a
T.toDevice Device
T.cpu (Tensor -> Tensor) -> Tensor -> Tensor
forall a b. (a -> b) -> a -> b
$ [FilePath] -> [FilePath] -> Tensor
boolMask' [FilePath
"fug"] [FilePath]
paramsX
    maskY :: Tensor
maskY      = Device -> Tensor -> Tensor
forall a. HasTypes a Tensor => Device -> a -> a
T.toDevice Device
T.cpu (Tensor -> Tensor) -> Tensor -> Tensor
forall a b. (a -> b) -> a -> b
$ [FilePath] -> [FilePath] -> Tensor
boolMask' [FilePath
"idoverW", FilePath
"L"] [FilePath]
paramsY
    numInputs :: Int
numInputs  = [FilePath] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [FilePath]
paramsX
    numOutputs :: Int
numOutputs = [FilePath] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [FilePath]
paramsY
    num' :: Int
num'       = Int
numInputs Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1

-- | Model Test
testModel :: Int -> DF.DataFrame T.Tensor -> [String] -> [String]
          -> (T.Tensor -> T.Tensor) -> IO ()
testModel :: Int
-> DataFrame Tensor
-> [FilePath]
-> [FilePath]
-> (Tensor -> Tensor)
-> IO ()
testModel Int
reg DataFrame Tensor
df [FilePath]
paramsX [FilePath]
paramsY Tensor -> Tensor
mdl = do
    DataFrame Tensor
df' <- Int -> Bool -> DataFrame Tensor -> IO (DataFrame Tensor)
DF.sampleIO Int
10 Bool
False (DataFrame Tensor -> IO (DataFrame Tensor))
-> DataFrame Tensor -> IO (DataFrame Tensor)
forall a b. (a -> b) -> a -> b
$  Tensor -> DataFrame Tensor -> DataFrame Tensor
DF.rowFilter (Int -> DataFrame Tensor -> Tensor
region Int
reg DataFrame Tensor
df) DataFrame Tensor
df
    let x :: DataFrame Tensor
x  = [FilePath] -> DataFrame Tensor -> DataFrame Tensor
DF.lookup [FilePath]
paramsX DataFrame Tensor
df'
        y :: DataFrame Tensor
y  = [FilePath] -> DataFrame Tensor -> DataFrame Tensor
DF.lookup [FilePath]
paramsY DataFrame Tensor
df'
           --   DF.DataFrame paramsY 
           --  $ T.cat (T.Dim 1) [ T.abs $ (df' ?? "id") / (df' ?? "idoverW")
           --                    , T.abs $ df' ?? "L" ]
        y' :: DataFrame Tensor
y' = [FilePath] -> Tensor -> DataFrame Tensor
forall a. [FilePath] -> a -> DataFrame a
DF.DataFrame [FilePath]
paramsY (Tensor -> DataFrame Tensor)
-> (Tensor -> Tensor) -> Tensor -> DataFrame Tensor
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Tensor -> Tensor
mdl (Tensor -> DataFrame Tensor) -> Tensor -> DataFrame Tensor
forall a b. (a -> b) -> a -> b
$ DataFrame Tensor -> Tensor
forall a. DataFrame a -> a
DF.values DataFrame Tensor
x
    DataFrame Tensor -> IO ()
forall a. Show a => a -> IO ()
print DataFrame Tensor
x
    DataFrame Tensor -> IO ()
forall a. Show a => a -> IO ()
print DataFrame Tensor
y
    DataFrame Tensor -> IO ()
forall a. Show a => a -> IO ()
print DataFrame Tensor
y'
    () -> IO ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

-- | Experimental Model Test
testModelExp :: DF.DataFrame T.Tensor -> [String] -> [String]
          -> (T.Tensor -> T.Tensor) -> IO ()
testModelExp :: DataFrame Tensor
-> [FilePath] -> [FilePath] -> (Tensor -> Tensor) -> IO ()
testModelExp DataFrame Tensor
df [FilePath]
paramsX [FilePath]
paramsY Tensor -> Tensor
mdl = do
    DataFrame Tensor
df' <- Int -> Bool -> DataFrame Tensor -> IO (DataFrame Tensor)
DF.sampleIO Int
10 Bool
False (DataFrame Tensor -> IO (DataFrame Tensor))
-> DataFrame Tensor -> IO (DataFrame Tensor)
forall a b. (a -> b) -> a -> b
$ Tensor -> DataFrame Tensor -> DataFrame Tensor
DF.rowFilter (Int -> DataFrame Tensor -> Tensor
region (-Int
1) DataFrame Tensor
df) DataFrame Tensor
df
    let x :: DataFrame Tensor
x  = [FilePath] -> DataFrame Tensor -> DataFrame Tensor
DF.lookup [FilePath]
paramsX DataFrame Tensor
df'
        y :: DataFrame Tensor
y  = [FilePath] -> DataFrame Tensor -> DataFrame Tensor
DF.lookup [FilePath]
paramsY DataFrame Tensor
df'
        y' :: DataFrame Tensor
y' = [FilePath] -> Tensor -> DataFrame Tensor
forall a. [FilePath] -> a -> DataFrame a
DF.DataFrame [FilePath]
paramsY (Tensor -> DataFrame Tensor)
-> (Tensor -> Tensor) -> Tensor -> DataFrame Tensor
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Tensor -> Tensor
mdl (Tensor -> DataFrame Tensor) -> Tensor -> DataFrame Tensor
forall a b. (a -> b) -> a -> b
$ DataFrame Tensor -> Tensor
forall a. DataFrame a -> a
DF.values DataFrame Tensor
x
    DataFrame Tensor -> IO ()
forall a. Show a => a -> IO ()
print DataFrame Tensor
x
    DataFrame Tensor -> IO ()
forall a. Show a => a -> IO ()
print DataFrame Tensor
y
    DataFrame Tensor -> IO ()
forall a. Show a => a -> IO ()
print DataFrame Tensor
y'
    () -> IO ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

-- | Filter region
-- Experimental: -1
-- Saturation: 2
-- Sub-Threshold: 3
region :: Int -> DF.DataFrame T.Tensor -> T.Tensor
region :: Int -> DataFrame Tensor -> Tensor
region (-1) DataFrame Tensor
df = Tensor -> Tensor -> Tensor
T.logicalAnd Tensor
m' Tensor
m''
  where
    m' :: Tensor
m' = Tensor -> Tensor -> Tensor
T.logicalOr (Tensor -> Tensor -> Tensor
T.eq (DataFrame Tensor
df DataFrame Tensor -> FilePath -> Tensor
?? FilePath
"region") Tensor
1.0)
       (Tensor -> Tensor) -> Tensor -> Tensor
forall a b. (a -> b) -> a -> b
$ Tensor -> Tensor -> Tensor
T.logicalOr (Tensor -> Tensor -> Tensor
T.eq (DataFrame Tensor
df DataFrame Tensor -> FilePath -> Tensor
?? FilePath
"region") Tensor
2.0) 
                     (Tensor -> Tensor -> Tensor
T.eq (DataFrame Tensor
df DataFrame Tensor -> FilePath -> Tensor
?? FilePath
"region") Tensor
3.0)
    -- m'' = T.gt (df ?? "gmoverid") 0.0
    m'' :: Tensor
m'' = Tensor -> Tensor -> Tensor
T.logicalAnd (Tensor -> Tensor -> Tensor
T.lt (DataFrame Tensor
df DataFrame Tensor -> FilePath -> Tensor
?? FilePath
"cgsoverW") Tensor
0.0)
        (Tensor -> Tensor) -> (Tensor -> Tensor) -> Tensor -> Tensor
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Tensor -> Tensor -> Tensor
T.logicalAnd (Tensor -> Tensor -> Tensor
T.lt (DataFrame Tensor
df DataFrame Tensor -> FilePath -> Tensor
?? FilePath
"cgdoverW") Tensor
0.0)
        (Tensor -> Tensor) -> (Tensor -> Tensor) -> Tensor -> Tensor
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Tensor -> Tensor -> Tensor
T.logicalAnd (Tensor -> Tensor -> Tensor
T.lt (DataFrame Tensor
df DataFrame Tensor -> FilePath -> Tensor
?? FilePath
"cgboverW") Tensor
0.0)
        (Tensor -> Tensor) -> (Tensor -> Tensor) -> Tensor -> Tensor
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Tensor -> Tensor -> Tensor
T.logicalAnd (Tensor -> Tensor -> Tensor
T.lt (DataFrame Tensor
df DataFrame Tensor -> FilePath -> Tensor
?? FilePath
"csdoverW") Tensor
0.0)
        (Tensor -> Tensor) -> (Tensor -> Tensor) -> Tensor -> Tensor
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Tensor -> Tensor -> Tensor
T.logicalAnd (Tensor -> Tensor -> Tensor
T.lt (DataFrame Tensor
df DataFrame Tensor -> FilePath -> Tensor
?? FilePath
"cbsoverW") Tensor
0.0)
        (Tensor -> Tensor) -> Tensor -> Tensor
forall a b. (a -> b) -> a -> b
$ Tensor -> Tensor -> Tensor
T.logicalAnd (Tensor -> Tensor -> Tensor
T.lt (DataFrame Tensor
df DataFrame Tensor -> FilePath -> Tensor
?? FilePath
"cbdoverW") Tensor
0.0)
                       (Tensor -> Tensor -> Tensor
T.gt (DataFrame Tensor
df DataFrame Tensor -> FilePath -> Tensor
?? FilePath
"gmoverid") Tensor
0.0)
region Int
3 DataFrame Tensor
df = Tensor -> Tensor -> Tensor
T.logicalOr  (Tensor -> Tensor -> Tensor
T.eq (DataFrame Tensor
df DataFrame Tensor -> FilePath -> Tensor
?? FilePath
"region") Tensor
3) 
            (Tensor -> Tensor) -> (Tensor -> Tensor) -> Tensor -> Tensor
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Tensor -> Tensor -> Tensor
T.logicalAnd (Tensor -> Tensor -> Tensor
T.gt Tensor
vds Tensor
10.0e-3)
            (Tensor -> Tensor) -> Tensor -> Tensor
forall a b. (a -> b) -> a -> b
$ Tensor -> Tensor -> Tensor
T.logicalAnd (Tensor -> Tensor -> Tensor
T.gt Tensor
vgs (Tensor
vth Tensor -> Tensor -> Tensor
forall a. Num a => a -> a -> a
- Tensor
vt')) (Tensor -> Tensor -> Tensor
T.lt Tensor
vgs (Tensor
vth Tensor -> Tensor -> Tensor
forall a. Num a => a -> a -> a
+ Tensor
vt''))
    where
      vt :: Tensor
vt    = Tensor
25.0e-3
      vt' :: Tensor
vt'   = Tensor
20 Tensor -> Tensor -> Tensor
forall a. Num a => a -> a -> a
* Tensor
vt
      vt'' :: Tensor
vt''  = Tensor
10 Tensor -> Tensor -> Tensor
forall a. Num a => a -> a -> a
* Tensor
vt
      vgs :: Tensor
vgs   = Tensor -> Tensor
T.abs (Tensor -> Tensor) -> Tensor -> Tensor
forall a b. (a -> b) -> a -> b
$ DataFrame Tensor
df DataFrame Tensor -> FilePath -> Tensor
?? FilePath
"vgs"
      vds :: Tensor
vds   = Tensor -> Tensor
T.abs (Tensor -> Tensor) -> Tensor -> Tensor
forall a b. (a -> b) -> a -> b
$ DataFrame Tensor
df DataFrame Tensor -> FilePath -> Tensor
?? FilePath
"vds"
      vth :: Tensor
vth   = Tensor -> Tensor
T.abs (Tensor -> Tensor) -> Tensor -> Tensor
forall a b. (a -> b) -> a -> b
$ DataFrame Tensor
df DataFrame Tensor -> FilePath -> Tensor
?? FilePath
"vth"
region Int
r DataFrame Tensor
df = Tensor -> Tensor -> Tensor
T.eq (DataFrame Tensor
df DataFrame Tensor -> FilePath -> Tensor
?? FilePath
"region") (Tensor -> Tensor) -> (Float -> Tensor) -> Float -> Tensor
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TensorLike Float => Float -> Tensor
forall a. TensorLike a => a -> Tensor
T.asTensor @Float (Float -> Tensor) -> Float -> Tensor
forall a b. (a -> b) -> a -> b
$ Int -> Float
forall a b. (Real a, Fractional b) => a -> b
realToFrac Int
r
-- region r _  = error $ "Region " ++ show r ++ " not defined"