packages feed

sound-collage-0.0: src/SoundCollage.hs

{-# LANGUAGE RebindableSyntax #-}
{- |
Very much inspired by NoiseReduction.
-}
module SoundCollage (
   testChopCompose, runDecompose, runAssociate, runCompose,
   Parameters(..), defltParams, chunkSizeFromPool,
   ) where

import qualified Sound.SoxLib as SoxLib
import qualified Sound.Frame as Sample

import qualified Math.FFT as FFT

import qualified Synthesizer.Storable.Signal as SigSt
import qualified Synthesizer.Generic.Filter.NonRecursive as FiltNRG
import qualified Synthesizer.Basic.Binary as Bin

import qualified Data.StorableVector as SV
import qualified Data.StorableVector.Lazy as SVL
import qualified Data.StorableVector.CArray as SVCArr

import qualified System.Directory as Dir
import qualified System.FilePath as FilePath
import qualified System.IO as IO
import System.FilePath ((</>), )
import Foreign.Storable (Storable, peek, )
import Control.Monad (forM, forM_, zipWithM_, )
import Data.Maybe (fromMaybe, )

import qualified Data.List.Key as Key
import qualified Data.List.HT as ListHT
import qualified Data.List as List
import Data.Tuple.HT (mapPair, )
import Data.List (isSuffixOf, )

import qualified Data.Complex as Complex98
import Data.Int (Int16, )

import Text.Printf (printf, )

import qualified Algebra.Transcendental as Trans
import qualified Algebra.RealRing as Real
import qualified Algebra.Additive as Additive

import NumericPrelude.Numeric
import NumericPrelude.Base
import qualified Prelude as P



pillow ::
   (Trans.C a, Storable a) =>
   Int -> SV.Vector a
pillow n =
   SV.sample n (\i -> sin (pi * fromIntegral i / fromIntegral n))

chop ::
   (Additive.C a, Storable a) =>
   Int -> Int -> SVL.Vector a -> [SVL.Vector a]
chop numOverlap shift =
   takeWhile (not . SVL.null) .
   iterate (SVL.drop shift) .
   SVL.append (SVL.fromChunks [SV.replicate ((numOverlap-1) * shift) zero])

{- |
The chunks in the chunks list must have length numOverlap*shift.
This is not checked.
If you use the result of 'chop'  then you are safe.
-}
compose ::
   (Additive.C a, Storable a) =>
   Int -> Int -> [SVL.Vector a] -> SVL.Vector a
compose numOverlap shift =
   foldl1 SigSt.mix .
   zipWith SVL.drop [0, shift ..] .
   reverse . map SVL.concat .
   ListHT.sliceHorizontal numOverlap


{- |
test reconstruction
-}
testChopCompose :: IO ()
testChopCompose =
   let shift = 42
       overlap = 5
       chunkSize = shift*overlap
       pilPre = SVL.fromChunks [pillow chunkSize]
       pilPost = SVL.map (2 / fromIntegral overlap *) pilPre
       xs :: SVL.Vector Float
       xs =
          SVL.fromChunks
             [SV.sample 10000 (\i -> cos (pi * fromIntegral i / 10000))]
   in  SVL.writeFile "/tmp/test.f32" $
          SVL.zipWith (\_ y -> y) xs $
          compose overlap shift $
          map (FiltNRG.envelope pilPost) $
          map (FiltNRG.envelope pilPre) $
          chop overlap shift $
          SVL.append xs $ SVL.fromChunks [SV.replicate chunkSize zero]


data Parameters =
   Parameters {
      paramShift, paramOverlap :: Int
   }

defltParams :: Parameters
defltParams =
   Parameters {
      paramShift = 1024,
      paramOverlap = 2
   }


maxCoeff :: SV.Vector Float -> Int
maxCoeff ys =
   Key.maximum (SV.index ys) $ take (SV.length ys) [0..]

featuresFromChunk ::
   Parameters -> SV.Vector Float -> SV.Vector Float
featuresFromChunk _params =
   SV.map Complex98.magnitude .
   SVCArr.from . FFT.dftRC . SVCArr.to

chopChannel ::
   Parameters -> SVL.Vector Float -> [SV.Vector Float]
chopChannel params input =
   let shift = paramShift params
       overlap = paramOverlap params
       chunkSize = shift*overlap

       pilPre = pillow chunkSize

   in  takeWhile ((chunkSize==) . SV.length) $
       map (FiltNRG.envelope pilPre . SV.concat . SVL.chunks) $
       chop overlap shift $
       SVL.append input $
       SVL.fromChunks [SV.replicate chunkSize zero]




spectrumSuffix, chunkSuffix :: String
spectrumSuffix = "-spec.f32"
chunkSuffix = "-chunk.s16"

runDecompose :: Parameters -> FilePath -> FilePath -> IO ()
runDecompose params src dst =
   SoxLib.withRead SoxLib.defaultReaderInfo src $ \fmtInPtr -> do
      fmtIn <- peek fmtInPtr
      let numChan =
             fromMaybe 1 $ SoxLib.channels $ SoxLib.signalInfo fmtIn

      inputs <-
         fmap (SVL.deinterleave numChan . SVL.map Bin.toCanonical) $
         SoxLib.readStorableVectorLazy fmtInPtr (SVL.ChunkSize 16384)

      let write n (spec,chunk) = do
            let path = printf dst (maxCoeff spec) (n::Int)
            Dir.createDirectoryIfMissing True $ FilePath.takeDirectory path
            SV.writeFile (path ++ spectrumSuffix) spec
            SV.writeFile (path ++ chunkSuffix) $
               (SV.map (Bin.fromCanonicalWith Real.roundSimple) chunk
                  :: SV.Vector Int16)

      zipWithM_ write [0..] $ uncurry zip $
         mapPair
            (foldl1 (zipWith (SV.zipWith (+))),
             map SV.concat . List.transpose) $
         unzip $
         map ((\chunks -> (map (featuresFromChunk params) chunks, chunks)) .
              chopChannel params)
            inputs

getDirectoryContents :: FilePath -> IO [FilePath]
getDirectoryContents =
   fmap (filter (not . flip elem [".", ".."])) .
   Dir.getDirectoryContents

replaceSuffix :: FilePath -> FilePath
replaceSuffix name =
   take (length name - length spectrumSuffix) name ++ chunkSuffix

fileSize :: FilePath -> IO Integer
fileSize path = IO.withFile path IO.ReadMode IO.hFileSize

divByteSize :: (Sample.C a) => a -> Integer -> Integer
divByteSize x n =
   div n (fromIntegral (Sample.sizeOfElement x))

chunkSizeFromPool :: FilePath -> IO Int
chunkSizeFromPool dir = do
   dirs <- getDirectoryContents dir
   firstBucket <-
      case dirs of
         [] -> ioError $ userError "chunk size determination: no bucket"
         bucket:_ -> return bucket
   files <-
      fmap (filter (isSuffixOf spectrumSuffix)) $
      getDirectoryContents (dir </> firstBucket)
   case files of
      [] -> ioError $ userError "chunk size determination: empty pool"
      file:_ -> do
         let path = dir </> firstBucket </> file
         specSize <- fmap (divByteSize (0::Float)) $ fileSize path
         chunkSize <-
            fmap (divByteSize (0::Int16)) $ fileSize $ replaceSuffix path
         return $ fromInteger $
            if mod chunkSize (specSize-1) == 0
              then (specSize-1)*2
              else specSize*2-1

loadSpectra :: FilePath -> IO [(SV.Vector Float, FilePath)]
loadSpectra dir = do
   files <-
      fmap (filter (isSuffixOf spectrumSuffix)) $
      getDirectoryContents dir
   forM files $ \file -> do
      spec <- SV.readFile $ dir </> file
      return (spec, file)

norm2 :: SV.Vector Float -> Float
norm2 = sqrt . SV.foldl' (+) 0 . SV.map (\x -> x*x)

createSpectrumMap ::
   Parameters -> FilePath -> IO [(SV.Vector Float, (Float, FilePath))]
createSpectrumMap _params poolDir = do
   files <- loadSpectra poolDir
   return $ flip map files $ \(spec, file) ->
      let norm = norm2 spec
      in  (SV.map (/norm) spec, (norm, file))

matchSpectrum :: SV.Vector Float -> SV.Vector Float -> Float
matchSpectrum spec dict =
   SV.foldl' (+) 0 $ SV.zipWith (*) spec dict

asVector16 :: SV.Vector Int16 -> SV.Vector Int16
asVector16 = id

clip16 :: Float -> Float
clip16 =
   min (fromIntegral (maxBound::Int16)) .
   max (fromIntegral (minBound::Int16))

associateBucket :: Parameters -> FilePath -> FilePath -> FilePath -> IO ()
associateBucket params poolDir src dst = do
   dict <- createSpectrumMap params poolDir
   files <- loadSpectra src
   forM_ files $ \(spec, file) -> do
      let srcNorm = norm2 spec
      let (poolNorm, matching) =
            snd $ Key.maximum (matchSpectrum spec . fst) dict
      SV.writeFile (dst </> replaceSuffix file)
         . asVector16
         . SV.map
              (Real.roundSimple . clip16 .
               (*(srcNorm/poolNorm)) . fromIntegral)
         . asVector16
         =<< SV.readFile (poolDir </> replaceSuffix matching)

merge :: (Ord a) => [a] -> [a] -> [(a,a)]
merge =
   let go (x:xs) yt@(y0:ys0) =
          if x <= y0
             then (x,y0) : go xs yt
             else
                case ys0 of
                   [] -> (x,y0) : []
                   y1:_ys1 ->
                      if x < y1
                         then (x,y0) : go xs ys0
                         else go (x:xs) ys0
       go [] _ = []
       go (_:_) [] = error "merge: second list empty"
   in  go

runAssociate :: Parameters -> FilePath -> FilePath -> FilePath -> IO ()
runAssociate params poolDir src dst = do
   Dir.createDirectoryIfMissing True dst
   srcDirs <- fmap List.sort $ getDirectoryContents src
   poolDirs <- fmap List.sort $ getDirectoryContents poolDir
   forM_ (merge srcDirs poolDirs) $ \(sdir,pdir) ->
      associateBucket params (poolDir </> pdir) (src </> sdir) dst


sliceVertical :: Int -> SV.Vector Float -> [SV.Vector Float]
sliceVertical n xs =
   map (SV.take n . flip SV.drop xs) $
   takeWhile (< SV.length xs) $ iterate (n+) 0

runCompose :: Parameters -> FilePath -> FilePath -> IO ()
runCompose params src dst = do
   let shift = paramShift params
       overlap = paramOverlap params
       chunkSize = shift*overlap

       pilPost = pillow chunkSize

   chunks <-
      mapM (SV.readFile . (src </>)) . List.sort =<< getDirectoryContents src
   SVL.writeFile dst $ SVL.interleaveFirstPattern $
      map (compose overlap shift) $ List.transpose $
      map
         (map (SVL.fromChunks . (:[]) . FiltNRG.envelope pilPost) .
          sliceVertical chunkSize .
          SV.map (Bin.toCanonical :: Int16 -> Float)) $
      chunks