Discrete Haar wavelet transform

To make this idea concrete, we will decompose the input vector \(y_1,\ldots,y_n\) according to a discrete Haar wavelet transform, which is a linear transformation that expands the \(n\)-dimensional input vector \(y_1,\ldots,y_n\) into an \((n-1)\)-dimensional vector of wavelet (or difference) coefficients \(d_1,\ldots,d_{n-1}\) and a single scaling (or average) coefficient \(c_0\). The discrete Haar wavelet transform is the most basic wavelet transform and is particularly well-suited to decompose piecewise constant signals, which generally produce very sparse Haar wavelet coefficient vectors with most wavelet coefficients equal to zero. For a description of the discrete Haar wavelet transform and a more general introduction to the use of wavelets in statistics, see e.g. (Nason 2008) or (Jansen 2012).

For simplicity, in the remainder of this post it is assumed that the number of observations is dyadic, i.e. \(n = 2^J\) for some integer \(J\), which is a common assumption in the context of wavelet regression³. Given the input signal \(f(x_1),\ldots,f(x_n)\), it is quite straightforward to encode the forward Haar transform in a recursive fashion:

## helper function to calculate scaling (average) or wavelet (difference) coefficients
filt <- function(C, fun) fun(C[c(T, F)], C[c(F, T)]) / sqrt(2)

## lists with scaling + wavelet coefficients from course to fine scales
C <- D <- vector(mode = "list", length = log2(N))

## recursively update course scale coefficients
for(l in log2(N):1) {
  C[[l]] <- filt(C = if(l < log2(N)) C[[l + 1]] else f, fun = `+`)
  D[[l]] <- filt(C = if(l < log2(N)) C[[l + 1]] else f, fun = `-`)
}

The list with scaling coefficients C consists of scaled local averages at increasingly coarse resolution scales, with the scaling coefficient at the coarsest scale (i.e. C[[1]]) being equivalent to the global mean scaled by a known factor:

all.equal(C[[1]], 2^(log2(N)/2) * mean(f))
#> [1] TRUE

Analogously, the list with wavelet coefficients D consists of scaled local differences of the scaling coefficients at increasingly coarse resolution scales. For the piecewise constant signal f, the list of wavelet coefficients is very sparse as most local differences are equal to zero and only a few non-zero wavelet coefficients are necessary to encode the jumps in the signal⁴:

D
#> [[1]]
#> [1] -11.31371
#> 
#> [[2]]
#> [1] -4 -4
#> 
#> [[3]]
#> [1] 0 0 0 0
#> 
#> [[4]]
#> [1] 0 0 0 0 0 0 0 0
#> 
#> [[5]]
#>  [1] 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
#> 
#> [[6]]
#>  [1] 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
#> 
#> [[7]]
#>  [1] 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
#> [39] 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0

Keeping track of all scaling and wavelet coefficients contained in C and D is redundant. To reconstruct the original input signal we need only the coarsest scaling coefficient C[[1]] and the \(n-1\) wavelet coefficients present in D. The inverse (or backward) Haar wavelet transform follows directly by applying the average and difference operations in the forward wavelet transform in the opposite sense:

## helper function to reconstruct scaling coefficients at finer scale
inv_filt <- function(C, D) c(t(cbind((C + D) / sqrt(2), (C - D) / sqrt(2))))

## recursively reconstruct fine scale coefficients
f1 <- C[[1]]
for(l in 1:log2(N)) {
  f1 <- inv_filt(C = f1, D = D[[l]])  
}

all.equal(f, f1)
#> [1] TRUE

The following functions encode the discrete (forward and backward) Haar wavelet transform in Stan and are saved in a file named haar.stan:

// haar.stan

// filter C coefficient vector
vector filtC(vector C, int N) {
    vector[N] C1;
    for (n in 1 : N) {
        C1[n] = (C[2 * n - 1] + C[2 * n]) / sqrt2();
    }
    return C1;
}
// filter D coefficient vector
vector filtD(vector D, int N) {
    vector[N] D1;
    for (n in 1 : N) {
        D1[n] = (D[2 * n - 1] - D[2 * n]) / sqrt2();
    }
    return D1;
}
// reconstruct C coefficient vector
vector inv_filt(vector C, vector D, int N) {
    vector[2 * N] C1;
    for (n in 1 : N) {
        C1[2 * n - 1] = (C[n] + D[n]) / sqrt2();
        C1[2 * n] = (C[n] - D[n]) / sqrt2();
    }
    return C1;
}
// forward Haar wavelet transform
vector fwt(vector y) {
    int N = rows(y);
    int Ni = 0;
    vector[N] ywd;
    vector[N] C = y;
    while (N > 1) {
        N /= 2;
        ywd[(Ni + 1):(Ni + N)] = filtD(C[1 : (2 * N)], N);
        C[1:N] = filtC(C[1 : (2 * N)], N);
        Ni += N;
    }
    ywd[Ni + 1] = C[1];
    return ywd;
}
// inverse Haar wavelet transform
vector iwt(vector ywd) {
    int N = rows(ywd);
    vector[N] y;
    int Nj = 1;
    y[1] = ywd[N];
    while (Nj < N) {
        y [1 : (2 * Nj)] = inv_filt(y[1 : Nj], ywd[(N - 2 * Nj + 1) : (N - Nj)], Nj);
        Nj *= 2;
    }
    return y;
}

The above code can then easily be included in the functions block of another Stan file with #include haar.stan. For instance, we can verify that the forward and backward wavelet transforms produce the expected outputs:

// haar_test.stan
functions{
  #include haar.stan
}
data {
  int<lower=1> N;
  vector[N] y;
}
parameters {
}
generated quantities {
   vector[N] ywd = fwt(y);
   vector[N] y1 = iwt(ywd);
}

## compile model
dwt_model <- cmdstan_model("haar_test.stan", include_paths = ".")

## draw single sample with no parameters
dwt_fit <- dwt_model$sample(
  data = list(N = N, y = f),
  chains = 1,
  iter_sampling = 1,
  iter_warmup = 0,
  sig_figs = 18,
  fixed_param = TRUE
)
#> Running MCMC with 1 chain...
#> 
#> Chain 1 Iteration: 1 / 1 [100%]  (Sampling) 
#> Chain 1 finished in 0.0 seconds.

## check forward wavelet transform
all.equal(c(dwt_fit$draws(variables = "ywd")), c(unlist(rev(D)), C[[1]]))
#> [1] TRUE

## check inverse wavelet transform 
all.equal(c(dwt_fit$draws(variables = "y1")), f)
#> [1] TRUE

Wavelet domain regression

Provided that the target signal \(f(x)\) has a sparse representation in the wavelet domain, a sensible estimation approach is to: (1) transform the noisy observation vector \(\boldsymbol{y} = (y_1,\ldots,y_n)'\) to the wavelet domain; (2) perform a sparse regression on the wavelet coefficients; (3) transform the result back to the functional domain to obtain an estimate \(\hat{f}(x)\).

As mentioned previously, the discrete Haar wavelet transform is a linear transformation \(\boldsymbol{d}^y = \boldsymbol{W} \boldsymbol{y}\), with \((n \times n)\)-dimensional wavelet transformation matrix \(\boldsymbol{W}\). Given the signal plus i.i.d. Gaussian noise model for the observations, \[ \boldsymbol{y} \ \sim N(\boldsymbol{f}, \sigma^2\boldsymbol{I}_{n \times n}) \] this implies that the transformed observations in the wavelet domain also follow a Gaussian signal plus noise model: \[ \boldsymbol{d}^y \ \sim N(\boldsymbol{d}^f, \sigma^2 \boldsymbol{W}\boldsymbol{W}') \] where \(\boldsymbol{d}^f = \boldsymbol{W}\boldsymbol{f}\) is the wavelet transformation of the target signal \(\boldsymbol{f} = (f(x_1),\ldots,f(x_n))'\).

Moreover, the linear transformation matrix \(\boldsymbol{W}\) is a unitary matrix, i.e. the transpose \(\boldsymbol{W}'\) and the inverse \(\boldsymbol{W}^{-1}\) coincide, (see e.g. (Nason 2008, Ch. 2)). This is particularly useful as it means that we also have a signal plus i.i.d. Gaussian noise model in the wavelet domain with the same noise variance as in the functional domain:

\[ d^y_{i} \ \overset{\text{iid}}{\sim} N(d^f_{i}, \sigma^2), \quad i = 1,\ldots, n \]

That is, the regression problem in the wavelet domain comes down to sparse linear regression in an i.i.d. Gaussian noise model.

Sparse linear regression in Stan

To induce sparsity in the estimated wavelet coefficient vector, we use a simplified version of the Finnish horseshoe prior as described in (Betancourt 2018), which is summarized as:

\[ \begin{aligned} d^f_i & \sim \ N(0, \tilde{\lambda}_i^2 ) \\ \tilde{\lambda}_i & \ = \ \frac{\tau \lambda_i}{\sqrt{1 + \tau^2 \lambda_i^2}} \\ \lambda_i & \sim \ \text{Half-}\mathcal{C}(0, 1) \\ \tau & \sim \ \text{Half-}N(0, \tau_0) \end{aligned} \]

The difference with respect to the specification in (Betancourt 2018) is that the additional scale parameter \(c\) is set to 1 and the Cauchy prior for \(\tau\) is replaced by a light-tailed normal prior. The \(\tau_0\) parameter is calculated as:

\[ \tau_0 \ = \ \frac{m_0}{1 - m_0} \frac{\sigma_0}{\sqrt{N}} \] where \(m_0\) is the expected fraction of non-zero wavelet coefficients (provided as input data), and \(\sigma_0\) is an initial estimate of the noise variance. The value for \(\sigma_0\) is calculated automatically by taking the standard deviation of the finest-scale coefficients in the noisy vector \(\boldsymbol{d}^y\), which are expected to contain primarily noise and (almost) no signal. As in (Betancourt 2018), the sampling results are not very sensitive to the value of \(m_0\) (and consequently \(\tau_0\)), but it does provide a convenient global tuning parameter for the amount of regularization applied to the estimated coefficients.

The complete model is encoded in a new Stan file step2.stan. First, the input vector \(y\) is transformed to the wavelet domain (ywd) in the transformed data block. The wavelet coefficient vector of the target signal \(f\) is constructed in the transformed parameters block based on the considerations above. To reduce the number of model parameters, the wavelet coefficients at the finest resolution scale are directly set to zero, (as the finest-scale wavelet coefficients are expected to contain only noise), leaving a set of \(N/2 - 1\) (sparse) wavelet coefficients to be estimated. The model block specifies the simplified Finnish horseshoe priors as well as a naive normal prior for the noise standard deviation \(\sigma\). In addition, the likelihood contributions based on the wavelet domain Gaussian linear model are specified. Finally, the regularized wavelet coefficient vector is back-transformed to the functional domain in the generated quantities block.

// step2.stan
functions{
  #include haar.stan 
}
data {
  int<lower=1> N;
  vector[N] y;        // input vector
  real<lower = 0> m0; // expected fraction of non-zero coefficients
}
transformed data{
  int M = (N / 2) - 1;                               // # estimated coefficients
  vector[N] ywd = fwt(y);                            // wavelet coefficients input
  real sigma0 = sd(ywd[1 : (N / 2)]);                // initial estimate sigma
  real tau0 = m0 / (1 - m0) * sigma0 / sqrt(N - 1);  // irrelevance scale
}
parameters {
  real<lower=0> sigma;        // noise standard deviation
  real<lower=0> tau;          // global scale horseshoe
  vector[M] z;                // unscaled estimated coefficients
  vector<lower=0>[M] lambda;  // local scales horseshoe
}
transformed parameters {
  // regularized (sparse) wavelet coefficients
  vector[N] fwd = rep_vector(0.0, N);
  fwd[(N - M) : (N - 1)] = (tau * lambda ./ sqrt(1 + square(tau * lambda))) .* z;
  fwd[N] = ywd[N];
}
model {  
  // (sparse) priors
  lambda ~ cauchy(0, 1);
  sigma ~ normal(sigma0, 5 * sigma0);
  tau ~ normal(0, tau0);
  z ~ std_normal();
  // likelhood contributions
  ywd[1 : (N - 1)] ~ normal(fwd[1 : (N - 1)], sigma);
}
generated quantities {
  // back-transformed coefficients
  vector[N] f = iwt(fwd);
}

We compile the model with cmdstanr and draw 1000 (after 1000 warm-up samples) per chain from 4 individual chains as before. For the expected fraction of non-zero wavelet coefficients, we use \(m_0 = 0.05\), which is quite conservative given our prior knowledge on the sparseness of the signal.

## compile model
step2_model <- cmdstan_model("step2.stan", include_paths = ".")

## draw samples
step2_fit <- step2_model$sample(
  data = list(N = N, y = y, m0 = 0.05),
  chains = 4,
  iter_sampling = 1000,
  iter_warmup = 1000
)
#> Running MCMC with 4 sequential chains...
#> 
#> Chain 1 Iteration:    1 / 2000 [  0%]  (Warmup) 
#> Chain 1 Iteration:  100 / 2000 [  5%]  (Warmup) 
#> Chain 1 Iteration:  200 / 2000 [ 10%]  (Warmup) 
#> Chain 1 Iteration:  300 / 2000 [ 15%]  (Warmup) 
#> Chain 1 Iteration:  400 / 2000 [ 20%]  (Warmup) 
#> Chain 1 Iteration:  500 / 2000 [ 25%]  (Warmup) 
#> Chain 1 Iteration:  600 / 2000 [ 30%]  (Warmup) 
#> Chain 1 Iteration:  700 / 2000 [ 35%]  (Warmup) 
#> Chain 1 Iteration:  800 / 2000 [ 40%]  (Warmup) 
....

## sampling results
step2_fit
#>  variable   mean median   sd  mad     q5    q95 rhat ess_bulk ess_tail
#>     lp__  -69.75 -69.58 9.20 9.44 -85.25 -55.05 1.00     1084     1954
#>     sigma   0.18   0.18 0.01 0.01   0.16   0.20 1.00     4411     3292
#>     tau     0.00   0.00 0.00 0.00   0.00   0.00 1.00     4376     2659
#>     z[1]   -0.04  -0.05 1.01 1.01  -1.63   1.61 1.00     6083     2605
#>     z[2]   -0.04  -0.05 1.00 0.97  -1.65   1.59 1.00     6483     2697
#>     z[3]    0.00  -0.02 1.04 1.05  -1.69   1.70 1.00     5897     2130
#>     z[4]   -0.06  -0.08 1.04 1.07  -1.74   1.65 1.00     5885     2644
#>     z[5]   -0.02  -0.02 1.01 1.00  -1.70   1.61 1.00     6283     2566
#>     z[6]    0.07   0.06 0.99 1.00  -1.53   1.68 1.00     5923     2462
#>     z[7]    0.02   0.01 1.01 1.01  -1.62   1.67 1.00     6808     2943
#> 
#>  # showing 10 of 385 rows (change via 'max_rows' argument)

## sampling diagnostics
step2_fit$cmdstan_diagnose()
#> Processing csv files: /tmp/Rtmp4e8Bg8/step2-202106161838-1-186086.csv, /tmp/Rtmp4e8Bg8/step2-202106161838-2-186086.csv, /tmp/Rtmp4e8Bg8/step2-202106161838-3-186086.csv, /tmp/Rtmp4e8Bg8/step2-202106161838-4-186086.csv
#> 
#> Checking sampler transitions treedepth.
#> Treedepth satisfactory for all transitions.
#> 
#> Checking sampler transitions for divergences.
#> No divergent transitions found.
#> 
#> Checking E-BFMI - sampler transitions HMC potential energy.
#> E-BFMI satisfactory for all transitions.
#> 
#> Effective sample size satisfactory.
#> 
#> Split R-hat values satisfactory all parameters.
#> 
#> Processing complete, no problems detected.

In contrast to the first estimation attempt in step1.stan, the sampling diagnostics now produce satisfying results as we have reparametrized the model to avoid the problematic log-likelihood gradient.

Below, we plot the posterior median of \(f\) as well as 90%- and 99%-credible bands (pointwise in time):

Blocks test function

To conclude, we apply the same sampling procedure to a more challenging example using the blocks test function available through DJ.EX() in the wavethresh-package, see also (Nason 2008, Ch. 3). The observations \(y_1, \ldots, y_n\) with \(n = 256\) are sampled from a signal plus i.i.d. Gaussian noise model as before:

library(wavethresh)

## data
set.seed(1)
N <- 256
x <- (1:N) / N
f <- DJ.EX(n = N, signal = 1)$blocks
y <- DJ.EX(n = N, signal = 1, rsnr = 5, noisy = TRUE)$blocks

ggplot(data = data.frame(x = x, y = y, f = f), aes(x = x)) + 
  geom_line(aes(y = f), lty = 2, color = "grey50") + 
  geom_point(aes(y = y)) +
  theme_light() +
  labs(x = "Time", y = "Response", title = "Blocks test function with i.i.d. Gaussian noise (N = 256)")

We draw 1000 posterior samples (after 1000 warm-up samples) per chain from 4 individual chains, with the expected fraction of non-zero wavelet coefficients set to \(m_0 = 0.05\) as before. Note that the number of breakpoints present in the signal does not need to be known prior to fitting the model.

## draw samples
blocks_fit <- step2_model$sample(
  data = list(N = N, y = y, m0 = 0.05),
  chains = 4,
  iter_sampling = 1000,
  iter_warmup = 1000
)
#> Running MCMC with 4 sequential chains...
#> 
#> Chain 1 Iteration:    1 / 2000 [  0%]  (Warmup) 
#> Chain 1 Iteration:  100 / 2000 [  5%]  (Warmup) 
#> Chain 1 Iteration:  200 / 2000 [ 10%]  (Warmup) 
#> Chain 1 Iteration:  300 / 2000 [ 15%]  (Warmup) 
#> Chain 1 Iteration:  400 / 2000 [ 20%]  (Warmup) 
#> Chain 1 Iteration:  500 / 2000 [ 25%]  (Warmup) 
#> Chain 1 Iteration:  600 / 2000 [ 30%]  (Warmup) 
#> Chain 1 Iteration:  700 / 2000 [ 35%]  (Warmup) 
#> Chain 1 Iteration:  800 / 2000 [ 40%]  (Warmup) 
....

The sampling results and diagnostics all look satisfactory:

## sampling results
blocks_fit
#>  variable    mean  median    sd   mad      q5     q95 rhat ess_bulk ess_tail
#>     lp__  -330.02 -329.62 13.39 13.15 -352.96 -308.54 1.00      436     1860
#>     sigma    0.38    0.38  0.04  0.03    0.33    0.45 1.01      126      239
#>     tau      0.00    0.00  0.00  0.00    0.00    0.01 1.00     1181     1866
#>     z[1]    -0.02    0.00  1.02  1.00   -1.71    1.66 1.00     6564     2765
#>     z[2]    -0.02   -0.03  1.00  1.01   -1.65    1.63 1.00     6298     2965
#>     z[3]    -0.01   -0.02  0.96  0.96   -1.58    1.57 1.00     5677     2899
#>     z[4]    -0.04   -0.06  1.02  1.03   -1.68    1.62 1.00     6102     2553
#>     z[5]     0.00    0.01  1.03  1.06   -1.74    1.72 1.00     6355     3067
#>     z[6]     0.02    0.03  1.02  1.05   -1.63    1.68 1.00     6465     2925
#>     z[7]    -0.11   -0.11  1.01  1.00   -1.80    1.54 1.00     5285     2612
#> 
#>  # showing 10 of 769 rows (change via 'max_rows' argument)

## sampling diagnostics
blocks_fit$cmdstan_diagnose()
#> Processing csv files: /tmp/Rtmp4e8Bg8/step2-202106161838-1-0ec2d6.csv, /tmp/Rtmp4e8Bg8/step2-202106161838-2-0ec2d6.csv, /tmp/Rtmp4e8Bg8/step2-202106161838-3-0ec2d6.csv, /tmp/Rtmp4e8Bg8/step2-202106161838-4-0ec2d6.csv
#> 
#> Checking sampler transitions treedepth.
#> Treedepth satisfactory for all transitions.
#> 
#> Checking sampler transitions for divergences.
#> No divergent transitions found.
#> 
#> Checking E-BFMI - sampler transitions HMC potential energy.
#> E-BFMI satisfactory for all transitions.
#> 
#> Effective sample size satisfactory.
#> 
#> Split R-hat values satisfactory all parameters.
#> 
#> Processing complete, no problems detected.

And finally we evaluate several posterior (pointwise) quantiles of \(f(x)\) analogous to the previous example: