contrib/haskell/src/Hkl/Flat.hs

   1 {-# LANGUAGE FlexibleInstances #-}
   2 {-# LANGUAGE GADTs #-}
   3 {-# LANGUAGE StandaloneDeriving #-}
   4 {-# LANGUAGE UnicodeSyntax #-}
   5
   6 module Hkl.Flat
   7        ( Flat(..)
   8        , Npy
   9        , computeFlat
  10        )
  11        where
  12
  13 import Data.Text ( unlines, pack )
  14 import System.Exit ( ExitCode( ExitSuccess ) )
  15 import System.FilePath.Posix ( replaceExtension )
  16
  17 import Hkl.DataSource ( DataItem ( DataItemH5 ) )
  18 import Hkl.Nxs ( Nxs ( Nxs )
  19                , XrdFlat
  20                , DataFrameH5Path ( XrdFlatH5Path )
  21                )
  22 import Hkl.Python ( PyVal
  23                   , toPyVal
  24                   )
  25 import Hkl.Script ( Py2
  26                   , Script ( Py2Script )
  27                   , run
  28                   )
  29
  30 data Npy
  31
  32 data Flat a where
  33   FlatNpy ∷ FilePath → Flat Npy
  34 deriving instance (Show) (Flat a)
  35
  36 scriptPy2Flat ∷ [Nxs XrdFlat] → FilePath → Script Py2
  37 scriptPy2Flat ns output = Py2Script (script, scriptName)
  38     where
  39       script = Data.Text.unlines $
  40                map pack ["#!/bin/env python"
  41                         , ""
  42                         , "import numpy"
  43                         , "from h5py import File"
  44                         , ""
  45                         , "NEXUSFILES = " ++ toPyVal nxs'
  46                         , "IMAGEPATHS = " ++ toPyVal hpaths
  47                         , "OUTPUT = " ++ toPyVal output
  48                         , ""
  49                         , "flat = None"
  50                         , "n = None"
  51                         , "with File(NEXUSFILES[0], mode='r') as f:"
  52                         , "    imgs = f[IMAGEPATHS[0]]"
  53                         , "    flat = numpy.sum(imgs[:], axis=0)"
  54                         , "    n = imgs.shape[0]"
  55                         , "for idx, (nxs, h5path) in enumerate(zip(NEXUSFILES[1:], IMAGEPATHS[1:])):"
  56                         , "    with File(nxs, mode='r') as f:"
  57                         , "        imgs = f[h5path]"
  58                         , "        flat += numpy.sum(imgs[:], axis=0)"
  59                         , "        n += imgs.shape[0]"
  60                         , "numpy.save(OUTPUT, flat.astype('f') / n)"
  61                         ]
  62       nxs' ∷ [String]
  63       nxs' = [f | (Nxs f _) ← ns]
  64
  65       hpaths ∷ [String]
  66       hpaths = [h | (Nxs _ (XrdFlatH5Path (DataItemH5 h _))) ← ns]
  67
  68       scriptName ∷ FilePath
  69       scriptName = output `replaceExtension` "py"
  70
  71 computeFlat ∷ [Nxs XrdFlat] → FilePath → IO (Flat Npy)
  72 computeFlat ns o = do
  73   -- create the python script.
  74   let script = scriptPy2Flat ns o
  75   -- execute this script.
  76   ExitSuccess ← run script False
  77   -- return the filepath of the generated file.
  78   return (FlatNpy o)
  79
  80 instance PyVal (Flat a) where
  81   toPyVal (FlatNpy v) = "numpy.load(" ++ show v ++ ")"