{-# LANGUAGE TupleSections #-}
{-| 
    Module      : Vocoder.Filter
    Description : Frequency-domain filters
    Copyright   : (c) Marek Materzok, 2021
    License     : BSD2

This module defines some useful frequency-domain filters for use in
the vocoder framework.
-}
module Vocoder.Filter (
    FreqStep,
    Filter,
    composeFilters,
    addFilters,
    idFilter,
    amplitudeFilter,
    linearAmplitudeFilter,
    amplify,
    lowpassBrickwall,
    highpassBrickwall,
    bandpassBrickwall,
    bandstopBrickwall,
    lowpassButterworth,
    highpassButterworth,
    bandpassButterworth,
    bandstopButterworth,
    pitchShiftInterpolate,
    convolution,
    convolutionFilter,
    envelope,
    envelopeFilter,
    randomPhaseFilter
    ) where

import Vocoder
import Vocoder.Window
import Control.Monad
import Control.Monad.IO.Class
import System.Random
import qualified Data.Vector.Storable as V

-- | A frequency step is a coefficient relating physical frequency (in Hz)
--   to FFT bin numbers. It is used to define filters independently of the
--   FFT window size.
type FreqStep = Double

-- | The type of frequency-domain filters. A frequency-domain filter is
--   a function transforming STFT frames which can depend on the
--   frequency step.
type Filter m = FreqStep -> STFTFrame -> m STFTFrame

-- | Sequential composition of filters.
composeFilters :: Monad m => Filter m -> Filter m -> Filter m
composeFilters :: Filter m -> Filter m -> Filter m
composeFilters Filter m
f1 Filter m
f2 FreqStep
step = Filter m
f1 FreqStep
step (STFTFrame -> m STFTFrame)
-> (STFTFrame -> m STFTFrame) -> STFTFrame -> m STFTFrame
forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> Filter m
f2 FreqStep
step

-- | Addition of filters.
addFilters :: Monad m => Filter m -> Filter m -> Filter m
addFilters :: Filter m -> Filter m -> Filter m
addFilters Filter m
f1 Filter m
f2 FreqStep
step STFTFrame
fr = STFTFrame -> STFTFrame -> STFTFrame
addFrames (STFTFrame -> STFTFrame -> STFTFrame)
-> m STFTFrame -> m (STFTFrame -> STFTFrame)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Filter m
f1 FreqStep
step STFTFrame
fr m (STFTFrame -> STFTFrame) -> m STFTFrame -> m STFTFrame
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Filter m
f2 FreqStep
step STFTFrame
fr

-- | Identity filter.
idFilter :: Monad m => Filter m
idFilter :: Filter m
idFilter FreqStep
_ = STFTFrame -> m STFTFrame
forall (m :: * -> *) a. Monad m => a -> m a
return

-- | Creates a filter which transforms only amplitudes, leaving phase
--   increments unchanged.
amplitudeFilter :: Monad m => (FreqStep -> Moduli -> Moduli) -> Filter m
amplitudeFilter :: (FreqStep -> Moduli -> Moduli) -> Filter m
amplitudeFilter FreqStep -> Moduli -> Moduli
f FreqStep
step (Moduli
mag, Moduli
ph_inc) = STFTFrame -> m STFTFrame
forall (m :: * -> *) a. Monad m => a -> m a
return (FreqStep -> Moduli -> Moduli
f FreqStep
step Moduli
mag, Moduli
ph_inc)

-- | Creates a filter which transforms amplitudes and zeroes the phase 
--   increments.
amplitudeFilter0 :: Monad m => (FreqStep -> Moduli -> Moduli) -> Filter m
amplitudeFilter0 :: (FreqStep -> Moduli -> Moduli) -> Filter m
amplitudeFilter0 FreqStep -> Moduli -> Moduli
f FreqStep
step (Moduli
mag, Moduli
ph_inc) = STFTFrame -> m STFTFrame
forall (m :: * -> *) a. Monad m => a -> m a
return (FreqStep -> Moduli -> Moduli
f FreqStep
step Moduli
mag, Int -> FreqStep -> Moduli
forall a. Storable a => Int -> a -> Vector a
V.replicate (Moduli -> Int
forall a. Storable a => Vector a -> Int
V.length Moduli
ph_inc) FreqStep
0)

-- | Creates a filter which scales amplitudes depending on frequency.
linearAmplitudeFilter :: Monad m => (Double -> Double) -> Filter m
linearAmplitudeFilter :: (FreqStep -> FreqStep) -> Filter m
linearAmplitudeFilter FreqStep -> FreqStep
f = (FreqStep -> Moduli -> Moduli) -> Filter m
forall (m :: * -> *).
Monad m =>
(FreqStep -> Moduli -> Moduli) -> Filter m
amplitudeFilter ((FreqStep -> Moduli -> Moduli) -> Filter m)
-> (FreqStep -> Moduli -> Moduli) -> Filter m
forall a b. (a -> b) -> a -> b
$ \FreqStep
step Moduli
mag -> (FreqStep -> FreqStep -> FreqStep) -> Moduli -> Moduli -> Moduli
forall a b c.
(Storable a, Storable b, Storable c) =>
(a -> b -> c) -> Vector a -> Vector b -> Vector c
V.zipWith FreqStep -> FreqStep -> FreqStep
forall a. Num a => a -> a -> a
(*) Moduli
mag (Moduli -> Moduli) -> Moduli -> Moduli
forall a b. (a -> b) -> a -> b
$ Int -> (Int -> FreqStep) -> Moduli
forall a. Storable a => Int -> (Int -> a) -> Vector a
V.generate (Moduli -> Int
forall a. Storable a => Vector a -> Int
V.length Moduli
mag) ((Int -> FreqStep) -> Moduli) -> (Int -> FreqStep) -> Moduli
forall a b. (a -> b) -> a -> b
$ \Int
k -> FreqStep -> FreqStep
f (FreqStep
step FreqStep -> FreqStep -> FreqStep
forall a. Num a => a -> a -> a
* Int -> FreqStep
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
k)

-- | Creates an "amplifier" which scales all frequencies.
amplify :: Monad m => Double -> Filter m
amplify :: FreqStep -> Filter m
amplify FreqStep
k = (FreqStep -> FreqStep) -> Filter m
forall (m :: * -> *). Monad m => (FreqStep -> FreqStep) -> Filter m
linearAmplitudeFilter (FreqStep -> FreqStep -> FreqStep
forall a b. a -> b -> a
const FreqStep
k)

-- | Creates a brickwall lowpass filter.
lowpassBrickwall :: Monad m => Double -> Filter m
lowpassBrickwall :: FreqStep -> Filter m
lowpassBrickwall FreqStep
t = (FreqStep -> FreqStep) -> Filter m
forall (m :: * -> *). Monad m => (FreqStep -> FreqStep) -> Filter m
linearAmplitudeFilter ((FreqStep -> FreqStep) -> Filter m)
-> (FreqStep -> FreqStep) -> Filter m
forall a b. (a -> b) -> a -> b
$ \FreqStep
x -> if FreqStep
x FreqStep -> FreqStep -> Bool
forall a. Ord a => a -> a -> Bool
<= FreqStep
t then FreqStep
1.0 else FreqStep
0.0 

-- | Creates a brickwall highpass filter.
highpassBrickwall :: Monad m => Double -> Filter m
highpassBrickwall :: FreqStep -> Filter m
highpassBrickwall FreqStep
t = (FreqStep -> FreqStep) -> Filter m
forall (m :: * -> *). Monad m => (FreqStep -> FreqStep) -> Filter m
linearAmplitudeFilter ((FreqStep -> FreqStep) -> Filter m)
-> (FreqStep -> FreqStep) -> Filter m
forall a b. (a -> b) -> a -> b
$ \FreqStep
x -> if FreqStep
x FreqStep -> FreqStep -> Bool
forall a. Ord a => a -> a -> Bool
>= FreqStep
t then FreqStep
1.0 else FreqStep
0.0 

-- | Creates a brickwall bandpass filter.
bandpassBrickwall :: Monad m => Double -> Double -> Filter m
bandpassBrickwall :: FreqStep -> FreqStep -> Filter m
bandpassBrickwall FreqStep
t FreqStep
u = (FreqStep -> FreqStep) -> Filter m
forall (m :: * -> *). Monad m => (FreqStep -> FreqStep) -> Filter m
linearAmplitudeFilter ((FreqStep -> FreqStep) -> Filter m)
-> (FreqStep -> FreqStep) -> Filter m
forall a b. (a -> b) -> a -> b
$ \FreqStep
x -> if FreqStep
x FreqStep -> FreqStep -> Bool
forall a. Ord a => a -> a -> Bool
>= FreqStep
t Bool -> Bool -> Bool
&& FreqStep
x FreqStep -> FreqStep -> Bool
forall a. Ord a => a -> a -> Bool
<= FreqStep
u then FreqStep
1.0 else FreqStep
0.0 

-- | Creates a brickwall bandstop filter.
bandstopBrickwall :: Monad m => Double -> Double -> Filter m
bandstopBrickwall :: FreqStep -> FreqStep -> Filter m
bandstopBrickwall FreqStep
t FreqStep
u = (FreqStep -> FreqStep) -> Filter m
forall (m :: * -> *). Monad m => (FreqStep -> FreqStep) -> Filter m
linearAmplitudeFilter ((FreqStep -> FreqStep) -> Filter m)
-> (FreqStep -> FreqStep) -> Filter m
forall a b. (a -> b) -> a -> b
$ \FreqStep
x -> if FreqStep
x FreqStep -> FreqStep -> Bool
forall a. Ord a => a -> a -> Bool
<= FreqStep
t Bool -> Bool -> Bool
|| FreqStep
x FreqStep -> FreqStep -> Bool
forall a. Ord a => a -> a -> Bool
>= FreqStep
u then FreqStep
1.0 else FreqStep
0.0 

butterworthGain :: Double -> Double -> Double -> Double
butterworthGain :: FreqStep -> FreqStep -> FreqStep -> FreqStep
butterworthGain FreqStep
n FreqStep
t FreqStep
x = FreqStep
1 FreqStep -> FreqStep -> FreqStep
forall a. Fractional a => a -> a -> a
/ FreqStep -> FreqStep
forall a. Floating a => a -> a
sqrt (FreqStep
1 FreqStep -> FreqStep -> FreqStep
forall a. Num a => a -> a -> a
+ (FreqStep
x FreqStep -> FreqStep -> FreqStep
forall a. Fractional a => a -> a -> a
/ FreqStep
t)FreqStep -> FreqStep -> FreqStep
forall a. Floating a => a -> a -> a
**(FreqStep
2 FreqStep -> FreqStep -> FreqStep
forall a. Num a => a -> a -> a
* FreqStep
n))

-- | Creates an n-th degree Butterworth-style lowpass filter.
lowpassButterworth :: Monad m => Double -> Double -> Filter m
lowpassButterworth :: FreqStep -> FreqStep -> Filter m
lowpassButterworth FreqStep
n FreqStep
t = (FreqStep -> FreqStep) -> Filter m
forall (m :: * -> *). Monad m => (FreqStep -> FreqStep) -> Filter m
linearAmplitudeFilter ((FreqStep -> FreqStep) -> Filter m)
-> (FreqStep -> FreqStep) -> Filter m
forall a b. (a -> b) -> a -> b
$ FreqStep -> FreqStep -> FreqStep -> FreqStep
butterworthGain FreqStep
n FreqStep
t

-- | Creates an n-th degree Butterworth-style highpass filter.
highpassButterworth :: Monad m => Double -> Double -> Filter m
highpassButterworth :: FreqStep -> FreqStep -> Filter m
highpassButterworth FreqStep
n FreqStep
t = (FreqStep -> FreqStep) -> Filter m
forall (m :: * -> *). Monad m => (FreqStep -> FreqStep) -> Filter m
linearAmplitudeFilter ((FreqStep -> FreqStep) -> Filter m)
-> (FreqStep -> FreqStep) -> Filter m
forall a b. (a -> b) -> a -> b
$ FreqStep -> FreqStep -> FreqStep -> FreqStep
butterworthGain (-FreqStep
n) FreqStep
t

-- | Creates an n-th degree Butterworth-style bandpass filter.
bandpassButterworth :: Monad m => Double -> Double -> Double -> Filter m
bandpassButterworth :: FreqStep -> FreqStep -> FreqStep -> Filter m
bandpassButterworth FreqStep
n FreqStep
t FreqStep
u = (FreqStep -> FreqStep) -> Filter m
forall (m :: * -> *). Monad m => (FreqStep -> FreqStep) -> Filter m
linearAmplitudeFilter ((FreqStep -> FreqStep) -> Filter m)
-> (FreqStep -> FreqStep) -> Filter m
forall a b. (a -> b) -> a -> b
$ \FreqStep
x -> FreqStep -> FreqStep -> FreqStep -> FreqStep
butterworthGain FreqStep
n FreqStep
u FreqStep
x FreqStep -> FreqStep -> FreqStep
forall a. Num a => a -> a -> a
* FreqStep -> FreqStep -> FreqStep -> FreqStep
butterworthGain (-FreqStep
n) FreqStep
t FreqStep
x

-- | Creates an n-th degree Butterworth-style bandstop filter.
bandstopButterworth :: Monad m => Double -> Double -> Double -> Filter m
bandstopButterworth :: FreqStep -> FreqStep -> FreqStep -> Filter m
bandstopButterworth FreqStep
n FreqStep
t FreqStep
u = (FreqStep -> FreqStep) -> Filter m
forall (m :: * -> *). Monad m => (FreqStep -> FreqStep) -> Filter m
linearAmplitudeFilter ((FreqStep -> FreqStep) -> Filter m)
-> (FreqStep -> FreqStep) -> Filter m
forall a b. (a -> b) -> a -> b
$ \FreqStep
x -> FreqStep -> FreqStep -> FreqStep -> FreqStep
butterworthGain (-FreqStep
n) FreqStep
t FreqStep
x FreqStep -> FreqStep -> FreqStep
forall a. Num a => a -> a -> a
+ FreqStep -> FreqStep -> FreqStep -> FreqStep
butterworthGain FreqStep
n FreqStep
u FreqStep
x

interpolate :: Double -> V.Vector Double -> V.Vector Double
interpolate :: FreqStep -> Moduli -> Moduli
interpolate FreqStep
n Moduli
v = Int -> (Int -> FreqStep) -> Moduli
forall a. Storable a => Int -> (Int -> a) -> Vector a
V.generate (Moduli -> Int
forall a. Storable a => Vector a -> Int
V.length Moduli
v) Int -> FreqStep
forall a. Integral a => a -> FreqStep
f 
    where
    f :: a -> FreqStep
f a
x | Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Moduli -> Int
forall a. Storable a => Vector a -> Int
V.length Moduli
v = FreqStep
0
        | Bool
otherwise = (FreqStep
1FreqStep -> FreqStep -> FreqStep
forall a. Num a => a -> a -> a
-FreqStep
k) FreqStep -> FreqStep -> FreqStep
forall a. Num a => a -> a -> a
* Moduli
v Moduli -> Int -> FreqStep
forall a. Storable a => Vector a -> Int -> a
V.! Int
i FreqStep -> FreqStep -> FreqStep
forall a. Num a => a -> a -> a
+ FreqStep
k FreqStep -> FreqStep -> FreqStep
forall a. Num a => a -> a -> a
* Moduli
v Moduli -> Int -> FreqStep
forall a. Storable a => Vector a -> Int -> a
V.! (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) where
            x' :: FreqStep
x' = FreqStep
n FreqStep -> FreqStep -> FreqStep
forall a. Num a => a -> a -> a
* a -> FreqStep
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
x
            i :: Int
i = FreqStep -> Int
forall a b. (RealFrac a, Integral b) => a -> b
floor FreqStep
x'
            k :: FreqStep
k = FreqStep
x' FreqStep -> FreqStep -> FreqStep
forall a. Num a => a -> a -> a
- Int -> FreqStep
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
i

-- | Creates an interpolative pitch-shifting filter.
pitchShiftInterpolate :: Monad m => Double -> Filter m
pitchShiftInterpolate :: FreqStep -> Filter m
pitchShiftInterpolate FreqStep
n FreqStep
_ (Moduli
mag, Moduli
ph_inc) = STFTFrame -> m STFTFrame
forall (m :: * -> *) a. Monad m => a -> m a
return (FreqStep -> Moduli -> Moduli
interpolate FreqStep
n Moduli
mag, (FreqStep -> FreqStep) -> Moduli -> Moduli
forall a b.
(Storable a, Storable b) =>
(a -> b) -> Vector a -> Vector b
V.map (FreqStep -> FreqStep -> FreqStep
forall a. Fractional a => a -> a -> a
/FreqStep
n) (Moduli -> Moduli) -> Moduli -> Moduli
forall a b. (a -> b) -> a -> b
$ FreqStep -> Moduli -> Moduli
interpolate FreqStep
n Moduli
ph_inc)

-- | Convolves the amplitude spectrum using a kernel.
convolution :: V.Vector Double -> Moduli -> Moduli
convolution :: Moduli -> Moduli -> Moduli
convolution Moduli
ker Moduli
mag = Int -> (Int -> FreqStep) -> Moduli
forall a. Storable a => Int -> (Int -> a) -> Vector a
V.generate (Moduli -> Int
forall a. Storable a => Vector a -> Int
V.length Moduli
mag) ((Int -> FreqStep) -> Moduli) -> (Int -> FreqStep) -> Moduli
forall a b. (a -> b) -> a -> b
$ \Int
k -> Moduli -> FreqStep
forall a. (Storable a, Num a) => Vector a -> a
V.sum (Moduli -> FreqStep) -> Moduli -> FreqStep
forall a b. (a -> b) -> a -> b
$ ((Int -> FreqStep -> FreqStep) -> Moduli -> Moduli)
-> Moduli -> (Int -> FreqStep -> FreqStep) -> Moduli
forall a b c. (a -> b -> c) -> b -> a -> c
flip (Int -> FreqStep -> FreqStep) -> Moduli -> Moduli
forall a b.
(Storable a, Storable b) =>
(Int -> a -> b) -> Vector a -> Vector b
V.imap Moduli
ker ((Int -> FreqStep -> FreqStep) -> Moduli)
-> (Int -> FreqStep -> FreqStep) -> Moduli
forall a b. (a -> b) -> a -> b
$ \Int
i FreqStep
v -> FreqStep
v FreqStep -> FreqStep -> FreqStep
forall a. Num a => a -> a -> a
* Moduli
gmag Moduli -> Int -> FreqStep
forall a. Storable a => Vector a -> Int -> a
V.! (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
k) FreqStep -> FreqStep -> FreqStep
forall a. Fractional a => a -> a -> a
/ FreqStep
s
    where
    h :: Int
h = Moduli -> Int
forall a. Storable a => Vector a -> Int
V.length Moduli
ker Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2
    gmag :: Moduli
gmag = Int -> FreqStep -> Moduli
forall a. Storable a => Int -> a -> Vector a
V.replicate Int
h FreqStep
0 Moduli -> Moduli -> Moduli
forall a. Storable a => Vector a -> Vector a -> Vector a
V.++ Moduli
mag Moduli -> Moduli -> Moduli
forall a. Storable a => Vector a -> Vector a -> Vector a
V.++ Int -> FreqStep -> Moduli
forall a. Storable a => Int -> a -> Vector a
V.replicate Int
h FreqStep
0
    s :: FreqStep
s = Moduli -> FreqStep
forall a. (Storable a, Num a) => Vector a -> a
V.sum Moduli
ker

-- | Creates a filter which convolves the spectrum using a kernel.
convolutionFilter :: Monad m => V.Vector Double -> Filter m
convolutionFilter :: Moduli -> Filter m
convolutionFilter Moduli
ker = (FreqStep -> Moduli -> Moduli) -> Filter m
forall (m :: * -> *).
Monad m =>
(FreqStep -> Moduli -> Moduli) -> Filter m
amplitudeFilter0 ((FreqStep -> Moduli -> Moduli) -> Filter m)
-> (FreqStep -> Moduli -> Moduli) -> Filter m
forall a b. (a -> b) -> a -> b
$ \FreqStep
_ -> Moduli -> Moduli -> Moduli
convolution Moduli
ker

-- | Calculates the envelope of an amplitude spectrum using convolution.
envelope :: Length -> Moduli -> Moduli
envelope :: Int -> Moduli -> Moduli
envelope Int
ksize = (FreqStep -> FreqStep) -> Moduli -> Moduli
forall a b.
(Storable a, Storable b) =>
(a -> b) -> Vector a -> Vector b
V.map ((FreqStep -> FreqStep -> FreqStep
forall a. Num a => a -> a -> a
+(-FreqStep
ee)) (FreqStep -> FreqStep)
-> (FreqStep -> FreqStep) -> FreqStep -> FreqStep
forall b c a. (b -> c) -> (a -> b) -> a -> c
. FreqStep -> FreqStep
forall a. Floating a => a -> a
exp) (Moduli -> Moduli) -> (Moduli -> Moduli) -> Moduli -> Moduli
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Moduli -> Moduli -> Moduli
convolution Moduli
ker (Moduli -> Moduli) -> (Moduli -> Moduli) -> Moduli -> Moduli
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (FreqStep -> FreqStep) -> Moduli -> Moduli
forall a b.
(Storable a, Storable b) =>
(a -> b) -> Vector a -> Vector b
V.map (FreqStep -> FreqStep
forall a. Floating a => a -> a
log (FreqStep -> FreqStep)
-> (FreqStep -> FreqStep) -> FreqStep -> FreqStep
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (FreqStep -> FreqStep -> FreqStep
forall a. Num a => a -> a -> a
+FreqStep
ee))
    where
    ee :: FreqStep
ee = FreqStep
2FreqStep -> FreqStep -> FreqStep
forall a. Floating a => a -> a -> a
**(-FreqStep
24)
    ker :: Moduli
ker = if Int
ksize Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
3 then Int -> Moduli
boxWindow Int
ksize else Int -> Moduli
blackmanWindow Int
ksize

-- | Creates a filter which replaces the amplitudes with their envelope.
envelopeFilter :: Monad m => Length -> Filter m
envelopeFilter :: Int -> Filter m
envelopeFilter Int
ksize = (FreqStep -> Moduli -> Moduli) -> Filter m
forall (m :: * -> *).
Monad m =>
(FreqStep -> Moduli -> Moduli) -> Filter m
amplitudeFilter0 ((FreqStep -> Moduli -> Moduli) -> Filter m)
-> (FreqStep -> Moduli -> Moduli) -> Filter m
forall a b. (a -> b) -> a -> b
$ \FreqStep
_ -> Int -> Moduli -> Moduli
envelope Int
ksize

-- | Sets the phase increments so that the bins have horizontal consistency.
--   This erases the phase information, introducing "phasiness".
randomPhaseFilter :: MonadIO m => Filter m
randomPhaseFilter :: Filter m
randomPhaseFilter FreqStep
_ (Moduli
mag, Moduli
ph_inc) = (Moduli
mag, ) (Moduli -> STFTFrame) -> m Moduli -> m STFTFrame
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> m FreqStep -> m Moduli
forall (m :: * -> *) a.
(Monad m, Storable a) =>
Int -> m a -> m (Vector a)
V.replicateM (Moduli -> Int
forall a. Storable a => Vector a -> Int
V.length Moduli
ph_inc) ((FreqStep, FreqStep) -> m FreqStep
forall a (m :: * -> *). (Random a, MonadIO m) => (a, a) -> m a
randomRIO (FreqStep
0, FreqStep
2FreqStep -> FreqStep -> FreqStep
forall a. Num a => a -> a -> a
*FreqStep
forall a. Floating a => a
pi))