Table of Contents

VarianceComponentModels.jl

Package Features

Heritability Analysis

Data files

Read in binary SNP data

EUR_subset

Empirical kinship matrix

Simulating phenotypes

Phenotypes

Pre-processing data for heritability analysis

Heritability of single trait

Multivariate trait analysis

Exercise

Testing SNP association using maximum likelihoods of variance component models

Fit the null model

Fit the alternative model

Likelihood ratio test

Exercise

# Heritability analysis and testing SNP association using maximum likelihoods of variance component models

**Lange Symposium**

**Juhyun Kim, juhkim111@ucla.edu**

**Department of Biostatistics, UCLA**

**Feb 22, 2020**

Machine information:

```julia
versioninfo()
```

## VarianceComponentModels.jl

[`VarianceComponentModels.jl`](https://github.com/OpenMendel/VarianceComponentModels.jl/) is a package that resides in [OpenMendel](https://github.com/OpenMendel) ecosystem. It implements computation routines for fitting and testing variance component model of form

$$\text{vec}(Y) \sim \text{Normal}(XB, \Sigma_1 \otimes V_1 + \cdots + \Sigma_m \otimes V_m)$$

where $\otimes$ is the [Kronecker product](https://en.wikipedia.org/wiki/Kronecker_product).

In this model, **data** is represented by

* `Y`: `n x d` response matrix
* `X`: `n x p` covariate matrix
* `V=(V1,...,Vm)`: a tuple of `m` `n x n` covariance matrices

and **parameters** are

* `B`: `p x d` mean parameter matrix
* `Σ=(Σ1,...,Σm)`: a tuple of `m` `d x d` variance components.

### Package Features
* Maximum likelihood estimation (MLE) and restricted maximum likelihood estimation (REML) of mean parameters $B$ and variance component parameters $Σ$
* Allow constraints in the mean parameters $B$
* Choice of optimization algorithms: [Fisher scoring](https://books.google.com/books?id=QYqeYTftPNwC&lpg=PP1&pg=PA142#v=onepage&q&f=false) and [minorization-maximization algorithm](http://hua-zhou.github.io/media/pdf/ZhouHuZhouLange19VCMM.pdf)
* [Heritability analysis](https://openmendel.github.io/VarianceComponentModels.jl/latest/man/heritability/#Heritability-Analysis-1) in genetics

## Heritability Analysis

Variance component estimation can be used to estimate heritability of a quantitative trait.

### Data files

For this analysis, we use a sample data set [`EUR_subset`](https://openmendel.github.io/SnpArrays.jl/latest/#Example-data-1) from `SnpArrays.jl`. This data set is available in the `data` folder of the package.

```julia
using SnpArrays
```

```julia
datapath = normpath(SnpArrays.datadir())
```

`EUR_subset.bed`, `EUR_subset.bim`, and `EUR_subset.fam` is a set of Plink files in binary format.

```julia
using Glob
readdir(glob"EUR_subset.*", datapath)
```

### Read in binary SNP data

We use the [`SnpArrays.jl`](https://openmendel.github.io/SnpArrays.jl/latest) package to read in binary SNP data and compute the empirical kinship matrix.

```julia
# read in genotype data from Plink binary file
const EUR_subset = SnpArray(SnpArrays.datadir("EUR_subset.bed"))
```

### `EUR_subset`

`EUR_subset` contains **379** individuals and **54,051** SNPs. There is no missing genotype in `EUR_subset`.

Minor allele frequencies (MAF) for each SNP:

```julia
maf_EUR = maf(EUR_subset)
```

Histogram of minor allele frequency:

```julia
## generate a histogram of MAF
# using Plots, PyPlot
# gr(size=(600, 500), html_output_format=:png)
# hist_maf = histogram(maf_EUR, xlab = "Minor Allele Frequency (MAF)",
#                    ylab = "Number of SNPs", label="")
# png(hist_maf, "hist_MAF.png")
```

Note that about 29% of SNPs are rare variants (MAF < 0.05).

```julia
count(!iszero, maf_EUR .< 0.05) / length(maf_EUR)
```

### Empirical kinship matrix

For a measure of relatedness, we compute empirical kinship matrix based on all SNPs by the genetic relation matrix (GRM). If there are missing genotypes, they are imputed on the fly by drawing according to the minor allele frequencies.

Kinship coefficients summarize genetic similarity between pairs of individuals. To estimate kinship coefficient $\Phi_{ij}$ between individuals $i$ and $j$ using GRM:

$$\widehat{\Phi}_{GRMij} = \frac{1}{2S} \sum_{k=1}^S \frac{(G_{ik}-2p_k)(G_{jk}-2p_k)}{2p_k(1-p_k)},$$

where

* $S$: number of SNPs in this set
* $p_k$: minor allele frequency of SNP $k$
* $G_{ik} \in \{0,1,2\}$: number of copies of minor alleles at the $k$-th SNP of the $i$-th individual

```julia
## GRM using SNPs with maf > 0.01 (default)
Φgrm = grm(EUR_subset; method = :GRM) # classical genetic relationship matrix
# Φgrm = grm(EUR_subset; method = :MoM) # method of moment method
# Φgrm = grm(EUR_subset; method = :Robust) # robust method
```

### Simulating phenotypes

We simulate phenotype vector from

$$\mathbf{y} \sim \text{Normal}(\mathbf{1}, 0.1 \widehat{\Phi}_{GRM} + 0.9 \mathbf{I})$$

where $\widehat{\Phi}_{GRM}$ is the estimated empirical kinship matrix `Φgrm`.

The data should be available in `pheno.txt`.

```julia
## simulate `pheno.txt`
# using LinearAlgebra, DelimitedFiles
# Random.seed!(1234)
# Ω = 0.1 * Φgrm + 0.9 * Matrix(1.0*I, nobs, nobs)
# Ωchol = cholesky(Symmetric(Ω))
# y = ones(nobs) + Ωchol.L * randn(nobs)
# writedlm("pheno.txt", y)
```

### Phenotypes

Read in the phenotype data and plot a histogram.

```julia
using DelimitedFiles
pheno = readdlm("pheno.txt")
```

Histogram of phenotype values:

```julia
## generate histogram of phenotype values
#hist_pheno = histogram(pheno, xlab="Phenotype", ylab="Frequency", label="")
#png(hist_pheno, "hist_pheno.png")
```

### Pre-processing data for heritability analysis

To prepare variance component model fitting, we form an instance of VarianceComponentVariate. The two covariance matrices are $(2\Phi, I)$.

```julia
using VarianceComponentModels, LinearAlgebra
# no. of observations
nobs = size(pheno, 1)

# form data as VarianceComponentVariate
X = ones(nobs)
EURdata = VarianceComponentVariate(pheno, X, (2Φgrm, Matrix(1.0I, nobs, nobs)))
fieldnames(typeof(EURdata))
```

```julia
EURdata
```

### Heritability of single trait

Before fitting the variance component model, we pre-compute the eigen-decomposition of $2\Phi_{\text{GRM}}$, the rotated responses, and the constant part in log-likelihood, and store them as a TwoVarCompVariateRotate instance, which is re-used in various variane component estimation procedures.

We use Fisher scoring algorithm to fit variance component model for our trait.

```julia
# pre-compute eigen-decomposition
@time EURdata_rotated = TwoVarCompVariateRotate(EURdata)
fieldnames(typeof(EURdata_rotated))

# form data set for trait
trait_data = TwoVarCompVariateRotate(EURdata_rotated.Yrot,
    EURdata_rotated.Xrot, EURdata_rotated.eigval, EURdata_rotated.eigvec,
    EURdata_rotated.logdetV2)

# initialize model parameters
trait_model = VarianceComponentModel(trait_data)

# estimate variance components
_, _, _, Σcov, = mle_fs!(trait_model, trait_data; solver=:Ipopt, verbose=false)
σ2a = trait_model.Σ[1][1] # additive genetic variance
σ2e = trait_model.Σ[2][1] # environmental variance
@show σ2a, σ2e
```

Additive genetic variance:

```julia
σ2a
```

Environmental/non-genetic variance:

```julia
σ2e
```

```julia
# heritability and its standard error from single trait analysis
h, hse = heritability(trait_model.Σ, Σcov)

[h[1], hse[1]]
```

We can also run MM algorithm.

```julia
trait_model = VarianceComponentModel(trait_data)
@time _, _, _, Σcov, = mle_mm!(trait_model, trait_data; verbose=false)
σ2a = trait_model.Σ[1][1]
σ2e = trait_model.Σ[2][1]
@show σ2a, σ2e
```

Heritability and its standard error.

```julia
h, hse = heritability(trait_model.Σ, Σcov)
[h[1], hse[1]]
```

### Multivariate trait analysis

For the joint analysis of multiple traits, go to [`VarianceComponentModels` documentation](https://openmendel.github.io/VarianceComponentModels.jl/latest/man/heritability/).

### Exercise

Repeat the above analysis computing the empirical kinship matrix using the method of moment method or the robust method. See [SnpArrays.jl documentation](https://openmendel.github.io/SnpArrays.jl/latest/#Genetic-relationship-matrix-1).

```julia
# Φgrm = grm(EUR_subset; method = )
```

## Testing SNP association using maximum likelihoods of variance component models
credit: [Heritability tutorial by Sarah Ji, Janet Sinsheimer and Hua Zhou](https://github.com/OpenMendel/Tutorials/blob/master/Heritability/HERITABILITY-VCexample.ipynb)

Suppose we want to see a particular SNP has an effect on a given phenotype after accounting for relatedness among individuals. Here we fit variance component model with a single SNP *s* as fixed effect.

$$\hspace{5em}  \mathbf{y} = \mathbf{X}\mathbf{\beta} + \mathbf{G}_s \gamma + \mathbf{g} + \mathbf{\epsilon} \hspace{5em} (1)$$

\begin{equation}
\begin{array}{ll}
\mathbf{g} \sim N(\mathbf{0}, \sigma_g^2\mathbf{\Phi}) \\
\mathbf{\epsilon} \sim N(\mathbf{0}, \sigma_e^2\mathbf{I})
\end{array}
\end{equation}

where

* $\mathbf{y}$: phenotype

and

* Fixed effects:
    * $\mathbf{X}$: matrix of covariates including intercept
    * $\beta$: vector of covariate effects, including intercept
    * $\mathbf{G}_s$: genotype of SNP *s*
    * $\gamma$: (scalar) association parameter of interest, measuring the effect of genotype on phenotype
* Random effects:
    * $\mathbf{g}$: random vector of polygenic effects with $\mathbf{g} \sim N(\mathbf{0}, \sigma_g^2 \mathbf{\Phi})$
        * $\sigma_g^2$: additive genetic variance
        * $\mathbf{\Phi}$: matrix of pairwise measures of genetic relatedness
    * $\epsilon$: random vector with $\epsilon \sim N(\mathbf{0}, \sigma_e^2\mathbf{I})$
        * $\sigma_e^2$: non-genetic variance due to non-genetic effects assumed to be acting independently on individuals

To test whether SNP *s* is associated with phenotype, we fit two models. First consider the model without SNP *s* as fixed effects (aka null model):

$$\hspace{5em}  \mathbf{y} = \mathbf{X}\mathbf{\beta} + \mathbf{g} + \mathbf{\epsilon} \hspace{5em} (2)$$

and the model with SNP *s* as fixed effects (1). Then we can compare the log likelihood to see if there is improvement in the model fit with inclusion of the SNP of interest.

### Fit the null model

```julia
using VarianceComponentModels
# null data model has two variance components but no SNP fixed effects
# form data as VarianceComponentVariate matrix
X = ones(nobs)
nulldata = VarianceComponentVariate(pheno, X, (2Φgrm, Matrix(1.0I, nobs, nobs)))
```

```julia
nullmodel = VarianceComponentModel(nulldata)
```

```julia
@time nulllogl, nullmodel, = fit_mle!(nullmodel, nulldata; algo=:FS)
```

The null model log-likelihood (no SNP effects)

```julia
nulllogl
```

The null model mean effects (a grand mean)

```julia
nullmodel.B
```

The null model additive genetic variance

```julia
nullmodel.Σ[1]
```

The null model environmental variance

```julia
nullmodel.Σ[2]
```

### Fit the alternative model

```julia
snp_vec = convert(Vector{Float64}, EUR_subset[:, 10])
Xalt = [ones(nobs) snp_vec]
altdata = VarianceComponentVariate(pheno, Xalt, (2Φgrm, Matrix(1.0I, nobs, nobs)))
```

```julia
altmodel = VarianceComponentModel(altdata)
```

```julia
@time altlogl, altmodel, = fit_mle!(altmodel, altdata; algo=:FS)
```

The alternative model log-likelihood:

```julia
altlogl
```

The alternative model mean effects (a grand mean)

```julia
altmodel.B
```

The alternative model additive genetic variance

```julia
altmodel.Σ[1]
```

The alternative model environmental variance

```julia
altmodel.Σ[2]
```

### Likelihood ratio test

We use likelihood ratio test (LRT) to test the goodness-of-fit between two models.

Our likelihood ratio test statistic is 1.786 (distributed chi-squared), with one degree of freedom.

```julia
using Distributions
LRT = 2(altlogl - nulllogl)
```

The associated p-value:

```julia
pval = ccdf(Chisq(1), LRT)
```

We see that adding this SNP as a covariate to the model does not fit significantly better than the null model. In other words, the SNP does not explain more of the variation in our trait.

### Exercise

Use minorization-maximization algorithm (`algo=:MM`) to find MLEs of both null model and alternative model. Then conduct the likelihood ratio test.


Massimo Nocentini

June 27, 2016: add external references

May 31, 2016: primality testing by sets (broken)

May 22, 2016: subsets sums with Gray codes

May 21, 2016: basics, subsets sums

Abstract

This document collect some examples about *bit manipulation techniques* and some application of them to contest problems

# Intro

The current notebook born with the aim to collect introductory material, references, tutorial, code snippet and challenging problems around the topic of *bit manipulation techniques*. It follows to fullfil our need to get a deep understanding of this topic to be competitive in programming (contests).

First of all a collection of general references follows:

   - [__The Aggregate Magic Algorithms__][aggregate], by _Department of Electrical and Computer Engineering at the University of Kentucky_
   - [__Bit Twiddling Hacks__][seander], by _Sean Eron Anderson_
   - [__Hacker's delight__][hd], by _Warren_

On the other hand, the following pages have a spirit toward programming contests:

   - [__CPSC 490 202 - Problem Solving in Computer Science__][CPSC 490 202 current] ([old][CPSC 490 202 old] version)

A focused application about compressed representation can be found here:
   - https://code.google.com/archive/p/compressedbitset/
   - https://sdm.lbl.gov/fastbit/
   - http://crd-legacy.lbl.gov/~kewu/fastbit/compression.html

Blog posts that inspired this document are the following:
   - http://redrating.blogspot.it/2013/08/algo-and-techniques-bitmask.html
   - https://www.quora.com/What-is-bitmasking-What-kind-of-problems-can-be-solved-using-it

[aggregate]:http://aggregate.org/MAGIC/
[seander]:http://graphics.stanford.edu/~seander/bithacks.html
[hd]:http://www.hackersdelight.org/

[CPSC 490 202 old]:http://www.ugrad.cs.ubc.ca/~cs490/sec202/
[CPSC 490 202 current]:http://www.ugrad.cs.ubc.ca/~cs490/2015W2/

# *bit masking* core idea

In this section we study [this][basic] introductory document, where the first paragraph reads:

>Suppose you have a set of objects and you want some way to represent which objects to pick and
which ones not to pick. How do you represent that in in a program? More generally, *how do you
represent a subest of a set*? We can represent whether an object *is picked or not by a single bit*! Using a boolean to represent this is an overkill in terms of memory usage. However, neither C++ nor Java has any data type representing a single bit, so how can we cut down the memory usage?
**The answer is to use an integer!** We know an integer is just a bunch of bits stringed together, so
why don’t we use the integer to represent the entire set?

> For example, suppose in a set of 5 objects, we have picked the 1st, 3rd, and 4th object. The bitmask to represent this in binary is 01101 or 13 in decimal (in the notes, the 1st bit will always be the least significant bit and
will always appear at the very right). *We have just cut down the memory usage from five booleans to
a single integer*!

Then follows a list of basic manipulations, which we coded in [this][our:bits] Python module, loaded in the next cell; two functions deserve attention, namely `subsets` and `subsets_of`, read the docstring for info.

[basic]:http://www.ugrad.cs.ubc.ca/~cs490/sec202/notes/intro/bitmask.pdf
[our:bits]:https://github.com/massimo-nocentini/competitive-programming/blob/master/python-libs/bits.py

```python
%matplotlib inline
%run ../python-libs/bits.py
%run ../python-libs/timing.py
%run ../python-libs/graycodes.py
%run ../python-libs/symbols.py
```

# Sum of subsets: a first example

## With "usual" subsets generation

Previous doc finishes with an first application of bitmasking, here the request:

>Bitmask is an efficient and convenient way to represent subsets. For example, **given a set of numbers,
we want to find the sum of all subsets**. This is easy to code using bitmasks. Furthermore, we can use
an array to store all the results (imagine storing the results when you are using an array to represent
the subset)!

The following cell contains a pythonic implementation of their code, which is given in `C`: generator `subsets` yields subsets one by one, such that a successor subset `s` is computed by a previous subset `p` by *increment*, namely `s = p + 1`.

```python
def subsets_sums(S):
    """
    Returns a map of `(s, v)`, where `s` is a subset of `S`, and `v` is `sum(i in S & s)`.
    """
    n = len(S)
    sums = [0 for _ in range(1 << n)]
    for s in subsets(n):
        for j, v in enumerate(S):
            sums[s] += is_on(s, j, return_int=True) * v

    return sums
```

As a first experiment, let `S` be a list of 4 random integers, each one less than $2^{8}$.

```python
from random import randrange
S = [randrange(1 << 8) for _ in range(1 << 2)]
S
```

    [184, 150, 30, 153]

We time the execution of the function `subsets_sums` applied to `S` and bind `sums` to the resulting mapping:

```python
with timing(lambda: subsets_sums(S)) as (sums, start, end):
    print("Elapsed time: {:.5} secs".format(end-start))
```

    Elapsed time: 0.00037456 secs

Have a look at results, after the definition of the pretty printing utility `pretty_dict`:

```python
def pretty(sums):
    """
    Utility pretty printer.
    """
    return {pretty_mask(j, coding='little', width=len(S)): v for j, v in enumerate(sums)}
```

```python
pretty(sums)
```

    {'0000': 0,
     '0001': 153,
     '0010': 30,
     '0011': 183,
     '0100': 150,
     '0101': 303,
     '0110': 180,
     '0111': 333,
     '1000': 184,
     '1001': 337,
     '1010': 214,
     '1011': 367,
     '1100': 334,
     '1101': 487,
     '1110': 364,
     '1111': 517}

### Repeating experiments

```python
def plotter(do, n):

    import matplotlib.pyplot as plt

    def doer(i):
        """
        Repeat the previous experiment with lists of integers lengths `i`.
        """
        S = [randrange(1 << 8) for _ in range(i)]
        with timing(lambda: do(S)) as (sums, start, end):
            return end - start

    l = range(n)
    line, = plt.plot(l, list(map(doer, l)), '--', linewidth=2)

    plt.xlabel("lengths")
    plt.ylabel("secs")
    plt.show()
```

Moreover, we repeat the above experiment for lists of integers with different lengths to understand the *exponential* complexity of *subsets* enumeration: the longer list will have **24 integers** and it took aproximately **5 minutes** to return an answer.

```python
plotter(subsets_sums, n=25)
```

## Subsets enumeration as Gray codes

Another order of generation can be used for subsets, namely we see each subset as a [Gray code][gray]. In this way we know that consecutive subsets differs *exactly* by one bit, which is the same to say that they have the same objects but one, so *one* object goes away or comes in.

Here is the a new function `subsets_sums_gray` that uses two hooks; given two consecutive subsets and `p` is the position of the bit `b` that toggles, if `b` switches *off* then `off` hook removes the corresponding value `S[p]` from the sum cumulated up to now, otherwise `add` hook adds `S[p]` to the sum.

[gray]:http://nbviewer.jupyter.org/github/massimo-nocentini/competitive-programming/blob/master/tutorials/graycodes.ipynb?flush_cache=true

```python
def subsets_sums_gray(S):
    """
    Returns a map of `(s, v)`, where `s` is a subset of `S`, and `v` is `sum(i in S & s)`.
    """
    return list(high(gray_codes(length=len(S)),
                on= lambda p, r: r + S[p],
                off=lambda p, r: r - S[p],
                redux=0))
```

We time the new experiment applying it to the same list of integers `S`:

```python
with timing(lambda: subsets_sums_gray(S)) as (res, start, end):

    print("Elapsed time: {:.5} secs".format(end-start))

    sums_gray = [0 for _ in range(1 << len(S))]
    for o, s, p, r in res:
        sums_gray[s] = r
```

    Elapsed time: 0.00019288 secs

and we pretty print the subsets sums

```python
pretty(sums_gray)
```

    {'0000': 0,
     '0001': 153,
     '0010': 30,
     '0011': 183,
     '0100': 150,
     '0101': 303,
     '0110': 180,
     '0111': 333,
     '1000': 184,
     '1001': 337,
     '1010': 214,
     '1011': 367,
     '1100': 334,
     '1101': 487,
     '1110': 364,
     '1111': 517}

checking that the answers are the same

```python
assert sums == sums_gray
```

We repeat the experiment for lists of integers of different lengths as done before:

```python
plotter(subsets_sums_gray, n=25)
```

using this order of subsets generations we gain approximately a **10-factor** respect the "usual" generation order.

# Primality testing - **not working**

The following is a very **unstable** attempt to test primality of an integer $n$, working on subsets instead of integer values.

```python
from sympy import IndexedBase, Eq, solve, init_printing, Xor
from math import sqrt, ceil
from itertools import product
from collections import defaultdict

init_printing()

def prefix_dictionaries(n, indexes): # be the integer under study, we would like to decide if is a *prime* or not

    x, y = indexes

    ones = ones_of(n) # is the list of positions in the binary representation of `n` where bits are 1
    bits = ones[-1] + 1  # bits is the number of bits necessary to represent `n` in binary

    for prefix in range(ceil(sqrt(bits)), bits):
        compl = bits - prefix
        pairs = defaultdict(lambda: 0)
        for i, j in product(range(0,prefix+1), range(0,compl)): # compl
            pairs[i+j] += x[i]*y[j]
        yield compl, pairs

def prefix_systems(n, dicts, init={}):
    for compl, dictionary in dicts:
        yield { (k, compl) : Eq(is_on(n, k, return_int=True), v).subs(init)
               for k, v in dictionary.items() if k}

def rec(eqs, subscripts_sum, substitutions):

    if not eqs: yield substitutions

    eq = eqs.pop(subscripts_sum)

    substitutions[y[subscripts_sum]] = 0
    subst_eq = eq.subs(substitutions)
    if not subst_eq:
        yield False
    elif subst_eq == True:
        yield from rec(dict(eqs), subscripts_sum + 1, substitutions)
    else:
        sol = solve(subst_eq, x[subscripts_sum]).pop()
        new_substitutions = dict(substitutions)
        new_substitutions[x[subscripts_sum]] = sol.rhs % 2

def diofantine(systems):
    for s in systems:
        for sol in solve(s):
            for (k, compl), v in sol.items():
                for i in range(1, compl):
                    eq = v[i, compl]
                    substitutions = {y[i]:1}

                if not v.is_integer: break
            else:
                yield sol

def factorize(n, sols, indexes, init={}):
    x, y = indexes
    for sol in sols:
        working_sol = defaultdict(lambda: 0)
        working_sol.update(init)
        working_sol.update(sol)
        p, q = 0, 0
        for i in range(n.bit_length()):
            p = set_bit(p, i) if working_sol[x[i]] else clear_bit(p, i)
            q = set_bit(q, i) if working_sol[y[i]] else clear_bit(q, i)
        yield p, q

```

```python
x, y = IndexedBase('x'), IndexedBase('y')

n = 35

g = prefix_dictionaries(n, (x,y))
```

```python
next(g)
```

```python
init = {x[0]:1, y[0]:1}
eqs = prefix_systems(n, g, init)
```

```python
next(eqs)
```

```python
sols = diofantine(eqs)
```

```python
next(sols)
```

```python
factors = factorize(n, sols, (x,y), init)
```

```python
next(factors)
```

---

Bitmasking tutorial by Massimo Nocentini is licensed under a Creative Commons Attribution-NonCommercial 4.0 International License.
Based on a work at https://github.com/massimo-nocentini/competitive-programming/blob/master/tutorials/bitmasking.ipynb.


# R: Cluster Robust Double Machine Learning

## Motivation

In many empirical applications, errors exhibit a clustered structure such that the usual i.i.d. assumption does not hold anymore. In order to perform valid statistical inference, researchers have to account for clustering. In this notebook, we will shortly emphasize the consequences of clustered data on inference based on the double machine learning (DML) approach as has been considered in [Chiang et al. (2021)](https://doi.org/10.1080/07350015.2021.1895815). We will demonstrate how users of the [DoubleML](https://docs.doubleml.org/stable/index.html) package can account for one- and two-way clustering in their analysis.

Clustered errors in terms of one or multiple dimensions might arise in many empirical applications. For example, in a cross-sectional study, errors might be correlated (i) within regions (one-way clustering) or (ii) within regions and industries at the same time (two-way clustering). Another example for two-way clustering, discussed in [Chiang et al. (2021)](https://doi.org/10.1080/07350015.2021.1895815), refers to market share data with market shares being subject to shocks on the market and product level at the same time. We refer to [Cameron et al. (2011)](https://doi.org/10.1198/jbes.2010.07136) for an introduction to multiway clustering and a illustrative list of empirical examples.

## Clustering and double machine learning

Clustering creates a challenge to the double machine learning (DML) approach in terms of

1. a necessary adjustment of the formulae used for estimation of the variance covariance matrix, standard errors, p-values etc., and,
2. an adjusted resampling scheme for the cross-fitting algorithm.

The first point equally applies to classical statistical models, for example a linear regression model (see, for example [Cameron et al. 2011](https://doi.org/10.1198/jbes.2010.07136)). The second point arises because the clustering implies a correlation of errors from train and test samples if the standard cross-fitting procedure suggested in [Chernozhukov et al. (2018)](https://doi.org/10.1111/ectj.12097) was employed. The DML approach builds on independent sample splits into partitions that are used for training of the machine learning (ML) model learners and generation of predictions that are eventually used for solving the score function. For a motivation of the necessity of sample splitting, we refer to the illustration example in the [user guide](
https://docs.doubleml.org/stable/guide/basics.html#sample-splitting-to-remove-bias-induced-by-overfitting) as well as to the explanation in [Chernozhukov et al. (2018)](https://doi.org/10.1111/ectj.12097) .

In order to achieve independent data splits in a setting with one-way or multi-way clustering, [Chiang et al. (2021)](https://doi.org/10.1080/07350015.2021.1895815) develop an updated  $K$-fold sample splitting procedure that ensures independent sample splits: The data set is split into disjoint partitions in terms of all clustering dimensions. For example, in a situation with two-way clustering, the data is split into $K^2$ folds. The machine learning models are then trained on a specific fold and used for generation of predictions in hold-out samples. Thereby, the sample splitting procedure ensures that the hold-out samples do not contain observations of the same clusters as used for training.

```R
library('hdm')
library('DoubleML')
library('mlr3')
library('mlr3learners')
# surpress messages from mlr3 package during fitting
lgr::get_logger("mlr3")$set_threshold("warn")

library('ggplot2')
library('reshape2')
library('gridExtra')
```

## A Motivating Example: Two-Way Cluster Robust DML

In a first part, we show how the two-way cluster robust double machine learning (DML) ([Chiang et al. 2021](https://doi.org/10.1080/07350015.2021.1895815)) can be implemented with the [DoubleML](https://docs.doubleml.org/stable/index.html) package.
[Chiang et al. (2021)](https://doi.org/10.1080/07350015.2021.1895815) consider double-indexed data

\begin{equation}
\lbrace W_{ij}: i \in \lbrace 1, \ldots, N \rbrace, j \in \lbrace 1, \ldots, M \rbrace \rbrace
\end{equation}

and the partially linear IV regression model (PLIV)

$$\begin{aligned}
Y_{ij} = D_{ij} \theta_0 +  g_0(X_{ij}) + \epsilon_{ij}, & &\mathbb{E}(\epsilon_{ij} | X_{ij}, Z_{ij}) = 0, \\
Z_{ij} = m_0(X_{ij}) + v_{ij}, & &\mathbb{E}(v_{ij} | X_{ij}) = 0.
\end{aligned}$$

### Simulate two-way cluster data

We use the PLIV data generating process described in Section 4.1 of [Chiang et al. (2021)](https://doi.org/10.1080/07350015.2021.1895815).
The DGP is defined as
$$\begin{aligned}
Z_{ij} &= X_{ij}' \xi_0 + V_{ij}, \\
D_{ij} &= Z_{ij}' \pi_{10} + X_{ij}' \pi_{20} + v_{ij}, \\
Y_{ij} &= D_{ij} \theta + X_{ij}' \zeta_0 + \varepsilon_{ij},
\end{aligned}$$
with
$$\begin{aligned}
X_{ij} &= (1 - \omega_1^X - \omega_2^X) \alpha_{ij}^X
+ \omega_1^X \alpha_{i}^X + \omega_2^X \alpha_{j}^X, \\
\varepsilon_{ij} &= (1 - \omega_1^\varepsilon - \omega_2^\varepsilon) \alpha_{ij}^\varepsilon
+ \omega_1^\varepsilon \alpha_{i}^\varepsilon + \omega_2^\varepsilon \alpha_{j}^\varepsilon, \\
v_{ij} &= (1 - \omega_1^v - \omega_2^v) \alpha_{ij}^v
+ \omega_1^v \alpha_{i}^v + \omega_2^v \alpha_{j}^v, \\
V_{ij} &= (1 - \omega_1^V - \omega_2^V) \alpha_{ij}^V
+ \omega_1^V \alpha_{i}^V + \omega_2^V \alpha_{j}^V,
\end{aligned}$$
and $\alpha_{ij}^X, \alpha_{i}^X, \alpha_{j}^X \sim \mathcal{N}(0, \Sigma)$
where  $\Sigma$ is a $p_x \times p_x$ matrix with entries
$\Sigma_{kj} = s_X^{|j-k|}$.
Further
$$\begin{aligned}
\left(\begin{matrix} \alpha_{ij}^\varepsilon \\ \alpha_{ij}^v \end{matrix}\right),
\left(\begin{matrix} \alpha_{i}^\varepsilon \\ \alpha_{i}^v \end{matrix}\right),
\left(\begin{matrix} \alpha_{j}^\varepsilon \\ \alpha_{j}^v \end{matrix}\right)
\sim \mathcal{N}\left(0, \left(\begin{matrix} 1 & s_{\varepsilon v} \\
s_{\varepsilon v} & 1 \end{matrix} \right) \right)
\end{aligned}$$
and $\alpha_{ij}^V, \alpha_{i}^V, \alpha_{j}^V \sim \mathcal{N}(0, 1)$.

Data from this DGP can be generated with the [make_pliv_multiway_cluster_CKMS2021()](https://docs.doubleml.org/r/stable/reference/make_pliv_multiway_cluster_CKMS2021.html) function from [DoubleML](https://docs.doubleml.org/stable/index.html).
Analogously to [Chiang et al. (2021, Section 5)](https://doi.org/10.1080/07350015.2021.1895815)
we use the following parameter setting:
$\theta=1.0$, $N=M=25$, $p_x=100$, $\pi_{10}=1.0$, $\omega_X = \omega_{\varepsilon} = \omega_V = \omega_v = (0.25, 0.25)$, $s_X = s_{\varepsilon v} = 0.25$ and the $j$-th entries of the $p_x$-vectors $\zeta_0 = \pi_{20} = \xi_0$ are $(\zeta_{0})_j = 0.5^j$.
This are also the default values of [make_pliv_multiway_cluster_CKMS2021()](https://docs.doubleml.org/r/stable/reference/make_pliv_multiway_cluster_CKMS2021.html).

```R
# Set the simulation parameters
N = 25  # number of observations (first dimension)
M = 25  # number of observations (second dimension)
dim_X = 100  # dimension of X
set.seed(3141) # set seed

obj_dml_data = make_pliv_multiway_cluster_CKMS2021(N, M, dim_X)
```

### Data-Backend for Cluster Data
The implementation of cluster robust double machine learning is based on a special data-backend called [DoubleMLClusterData](https://docs.doubleml.org/r/stable/reference/DoubleMLClusterData.html). As compared to the standard data-backend [DoubleMLData](https://docs.doubleml.org/r/stable/reference/DoubleMLData.html), users can specify the clustering variables during instantiation of a  [DoubleMLClusterData](https://docs.doubleml.org/r/stable/reference/DoubleMLClusterData.html) object. The estimation framework will subsequently account for the provided clustering options.

```R
# The simulated data is of type DoubleMLClusterData
print(obj_dml_data)
```

```R
# The cluster variables are part of the DataFrame
head(obj_dml_data$data)
```

### Initialize the objects of class `DoubleMLPLIV`

```R
# Set machine learning methods for m, g & r
ml_g = lrn("regr.cv_glmnet", nfolds = 10, s = "lambda.min")
ml_m = lrn("regr.cv_glmnet", nfolds = 10, s = "lambda.min")
ml_r = lrn("regr.cv_glmnet", nfolds = 10, s = "lambda.min")

# initialize the DoubleMLPLIV object
dml_pliv_obj = DoubleMLPLIV$new(obj_dml_data,
                                ml_g, ml_m, ml_r,
                                n_folds=3)
```

```R
print(dml_pliv_obj)
```

```R
library(RColorBrewer)
coul <- rev(colorRampPalette(brewer.pal(8, "RdBu"))(3))
options(repr.plot.width = 10, repr.plot.height = 10)

plt_smpls = function(smpls, n_folds) {
    df = matrix(0, nrow = N*M, ncol = n_folds)
    for (i_fold in 1:n_folds){
        df[smpls$train_ids[[i_fold]], i_fold] = -1
        df[smpls$test_ids[[i_fold]], i_fold] = 1
    }

    heatmap(df, Rowv=NA, Colv=NA, col=coul, cexRow=1.5, cexCol=1.5, scale='none')
}

plt_smpls_cluster = function(smpls_cluster, n_folds, n_folds_per_cluster) {
    #options(repr.plot.width = 6, repr.plot.height = 6)
    plots = list()
    for (i_fold in 1:n_folds){
        mat = matrix(0, nrow = M, ncol = N)
        for (k in smpls_cluster$train_ids[[i_fold]][[1]]) {
            for (l in smpls_cluster$train_ids[[i_fold]][[2]]) {
                mat[k, l] = -1
            }
        }
        for (k in smpls_cluster$test_ids[[i_fold]][[1]]) {
            for (l in smpls_cluster$test_ids[[i_fold]][[2]]) {
                mat[k, l] = 1
            }
        }
        l = (i_fold-1) %% n_folds_per_cluster + 1
        k = ((i_fold-1) %/% n_folds_per_cluster)+1
        df = data.frame(mat)
        cols = names(df)
        names(df) = 1:N
        df$id = 1:N
        df_plot = melt(df,  id.var = 'id')
        df_plot$value = factor(df_plot$value)
        plots[[i_fold]] = ggplot(data = df_plot, aes(x=id, y=variable)) +
          geom_tile(aes(fill=value), colour = "grey50") +
        scale_fill_manual(values = c("darkblue", "white", "darkred")) +
        theme(text = element_text(size=15))
        # ToDo: Add Subplot titles
        if (k == 3) {
            plots[[i_fold]] = plots[[i_fold]] + xlab(expression(paste('Second Cluster Variable ', l)))
        } else {
            plots[[i_fold]] = plots[[i_fold]] + xlab('')
        }
        if (l == 1) {
            plots[[i_fold]] = plots[[i_fold]] + ylab(expression(paste('First Cluster Variable ', k)))
        } else {
            plots[[i_fold]] = plots[[i_fold]] + ylab('')
        }
    }
    return(plots)
}
```

### Cluster Robust Cross Fitting
A key element of cluster robust DML ([Chiang et al. 2021](https://doi.org/10.1080/07350015.2021.1895815)) is a special sample splitting used for the cross-fitting.
In case of two-way clustering, we assume $N$ clusters in the first dimension and $M$ clusters in the second dimension.

For $K$-fold cross-fitting, [Chiang et al. (2021)](https://doi.org/10.1080/07350015.2021.1895815) proposed to randomly partition $[N]:=\{1,\ldots,N\}$ into $K$ subsets $\{I_1, \ldots, I_K\}$ and $[M]:=\{1,\ldots,N\}$ into $K$ subsets $\{J_1, \ldots, J_K\}$.
Effectively, one then considers $K^2$ folds.
Basically for each $(k, \ell) \in \{1, \ldots, K\} \times \{1, \ldots, K\}$, the nuisance functions are estimated for all double-indexed observations in $([N]\setminus I_K) \times ([M]\setminus J_\ell)$, i.e.,
$$
\hat{\eta}_{k\ell} = \hat{\eta}\left((W_{ij})_{(i,j)\in ([N]\setminus I_K) \times ([M]\setminus J_\ell)}\right)
$$
The causal parameter is then estimated as usual by solving a moment condition with a Neyman orthogonal score function.
For two-way cluster robust double machine learning with algorithm [DML2](https://docs.doubleml.org/stable/guide/algorithms.html#algorithm-dml2) this results in solving
$$
\frac{1}{K^2} \sum_{k=1}^{K} \sum_{\ell=1}^{K} \frac{1}{|I_k| |J_\ell|} \sum_{(i,j) \in I_K \times J_\ell}
\psi(W_{ij}, \tilde{\theta}_0, \hat{\eta}_{k\ell}) = 0
$$
for $\tilde{\theta}_0$.
Here $|I_k|$ denotes the cardinality, i.e., the number of clusters in the $k$-th fold for the first cluster variable.

We can visualize the sample splitting of the $N \cdot M = 625$ observations into $K \cdot K = 9$ folds. The following heat map illustrates the partitioned data set that is split into $K=9$ folds. The horizontal axis corresponds to the fold indices and the vertical axis to the indices of the observations. A blue field indicates that the observation $i$ is used for fitting the nuisance part, red indicates that the fold is used for prediction generation and white means that an observation is left out from the sample splitting.

For example, the first observation as displayed on the very bottom of the figure is used for training of the nuisance parts in the first, second, fourth and fifth fold and used for generation of the predictions in fold nine. At the same time the observation is left out from the sample splitting procedure in folds three, six, seven and eight.

```R
# The function plt_smpls is defined at the end of the Notebook
plt_smpls(dml_pliv_obj$smpls[[1]], dml_pliv_obj$n_folds)
```

If we visualize the sample splitting in terms of the cluster variables, the partitioning of the data into $9$ folds $I_k \times J_\ell$ becomes clear.
The identifiers for the first cluster variable $[N]:=\{1,\ldots,N\}$ have been randomly partioned into $K=3$ folds denoted by $\{I_1, I_2, I_3\}$ and the identifiers for the second cluster variable $[M]:=\{1,\ldots,M\}$ have also been randomly partioned into $K=3$ folds denoted by $\{J_1, J_2, J_3\}$.
By considering every combination $I_k \times J_\ell$ for $1 \leq k, \ell \leq K = 3$ we effectively base the cross-fitting on $9$ folds.

We now want to focus on the top-left sub-plot showing the partitioning of the cluster data for the first fold.
The $x$-axis corresponds to the first cluster variable and the $y$-axis to the second cluster variable.
Observations with cluster variables $(i,j) \in I_K \times J_\ell$ are used for estimation of the target parameter $\tilde{\theta}_0$ by solving a Neyman orthogonal score function.
For estimation of the nuisance function, we only use observation where neither the first cluster variable is in $I_K$ nor the second cluster variable is in $J_\ell$, i.e., we use observations indexed by $(i,j)\in ([N]\setminus I_K) \times ([M]\setminus J_\ell)$ to estimate the nuisance functions
$$
\hat{\eta}_{k\ell} = \hat{\eta}\left((W_{ij})_{(i,j)\in ([N]\setminus I_K) \times ([M]\setminus J_\ell)}\right).
$$
This way we guarantee that there are never observations from the same cluster (first and/or second cluster dimension) in the sample for the nuisance function estimation (blue) and at the same time in the sample for solving the score function (red). As a result of this special sample splitting proposed by [Chiang et al. (2021)](https://doi.org/10.1080/07350015.2021.1895815), the observations in the score (red) and nuisance (blue) sample can be considered independent and the standard cross-fitting approach for double machine learning can be applied.

```R
# The function plt_smpls_cluster is defined at the end of the Notebook
options(repr.plot.width = 12, repr.plot.height = 10)
plots = plt_smpls_cluster(dml_pliv_obj$smpls_cluster[[1]],
                         dml_pliv_obj$n_folds,
                         sqrt(dml_pliv_obj$n_folds))
grid.arrange(grobs=plots, ncol = 3, nrow = 3)
```

### Cluster Robust Standard Errors
In the abstract base class `DoubleML` the estimation of cluster robust standard errors is implemented for all supported double machine learning models.
It is based on the assumption of a linear Neyman orthogonal score function.
We use the notation $n \wedge m := \min\{n,m\}$.
For the the asymptotic variance of
$\sqrt{\underline{C}}(\tilde{\theta_0} - \theta_0)$ with
$\underline{C} := N \wedge M$
[Chiang et al. (2021)](https://doi.org/10.1080/07350015.2021.1895815) then propose the following estimator
$$
\hat{\sigma}^2 = \hat{J}^{-1} \hat{\Gamma} \hat{J}^{-1}
$$
where
$$
\begin{aligned}
\hat{\Gamma} = \frac{1}{K^2} \sum_{(k, \ell) \in[K]^2}
\Bigg[ \frac{|I_k| \wedge |J_\ell|}{(|I_k||J_\ell|)^2}
\bigg(&\sum_{i \in I_k} \sum_{j \in J_\ell} \sum_{j' \in J_\ell}
\psi(W_{ij}; \tilde{\theta}, \hat{\eta}_{k \ell}) \psi(W_{ij'}; \tilde{\theta}_0, \hat{\eta}_{k \ell}) \\
&+ \sum_{i \in I_k} \sum_{i' \in I_k} \sum_{j \in J_\ell}
\psi(W_{ij}; \tilde{\theta}, \hat{\eta}_{k \ell}) \psi(W_{i'j}; \tilde{\theta}_0, \hat{\eta}_{k \ell})
\bigg)
\Bigg]
\end{aligned}$$
and
$$
\begin{aligned}
\hat{J} = \frac{1}{K^2} \sum_{(k, \ell) \in[K]^2} \frac{1}{|I_k||J_\ell|}
\sum_{i \in I_k} \sum_{j \in J_\ell}
\psi_a(W_{ij}; \tilde{\theta}_0, \hat{\eta}_{k \ell}).
\end{aligned}
$$
A $(1-\alpha)$ confidence interval is then given by ([Chiang et al. 2021](https://doi.org/10.1080/07350015.2021.1895815))
$$\begin{aligned}
\left[
\tilde{\theta} \pm \Phi^{-1}(1-\alpha/2) \sqrt{\hat{\sigma}^2 / \underline{C}}
\right]
\end{aligned}
$$
with $\underline{C} = N \wedge M$.

```R
# Estimate the PLIV model with cluster robust double machine learning
dml_pliv_obj$fit()
dml_pliv_obj$summary()
```

## (One-Way) Cluster Robust Double Machine Learing

We again use the PLIV data generating process described in Section 4.1 of [Chiang et al. (2021)](https://doi.org/10.1080/07350015.2021.1895815).
To obtain one-way clustered data, we set the following weights to zero
$$
\omega_2^X = \omega_2^\varepsilon = \omega_2^v = \omega_2^V = 0.
$$
Again we can simulate this data with [make_pliv_multiway_cluster_CKMS2021()](https://docs.doubleml.org/r/stable/reference/make_pliv_multiway_cluster_CKMS2021.html). To prepare the data-backend for one-way clustering, we only have to alter the `cluster_cols` to be `'cluster_var_i'`.

```R
obj_dml_data = make_pliv_multiway_cluster_CKMS2021(N, M, dim_X,
                                                   omega_X = c(0.25, 0),
                                                   omega_epsilon = c(0.25, 0),
                                                   omega_v = c(0.25, 0),
                                                   omega_V = c(0.25, 0))
```

```R
obj_dml_data$cluster_cols = 'cluster_var_i'
print(obj_dml_data)
```

```R
# Set machine learning methods for m, g & r
ml_g = lrn("regr.cv_glmnet", nfolds = 10, s = "lambda.min")
ml_m = lrn("regr.cv_glmnet", nfolds = 10, s = "lambda.min")
ml_r = lrn("regr.cv_glmnet", nfolds = 10, s = "lambda.min")

# initialize the DoubleMLPLIV object
dml_pliv_obj = DoubleMLPLIV$new(obj_dml_data,
                                ml_g, ml_m, ml_r,
                                n_folds=3)
```

```R
dml_pliv_obj$fit()
dml_pliv_obj$summary()
```

## Real-Data Application
As a real-data application we revist the consumer demand example from [Chiang et al. (2021)](https://doi.org/10.1080/07350015.2021.1895815).
The U.S. automobile data of [Berry, Levinsohn, and Pakes (1995)](https://doi.org/10.2307/2171802) is obtained from the `R` package [hdm](https://cran.r-project.org/web/packages/hdm/index.html). In this example, we consider different specifications for the cluster dimensions.

### Load and Process Data

```R
## Prepare the BLP data
data(BLP);
blp_data <- BLP$BLP;
blp_data$price <- blp_data$price + 11.761
blp_data$log_p = log(blp_data$price)
```

```R
x_cols = c('hpwt', 'air', 'mpd', 'space')
head(blp_data[x_cols])
```

```R
iv_vars = as.data.frame(hdm:::constructIV(blp_data$firm.id,
                            blp_data$cdid,
                            blp_data$id,
                            blp_data[x_cols]))
```

```R
formula = formula(paste0(" ~ -1 + (hpwt + air + mpd + space)^2",
                         "+ I(hpwt^2)*(air + mpd + space)",
                         "+ I(air^2)*(hpwt + mpd + space)",
                         "+ I(mpd^2)*(hpwt + air + space)",
                         "+ I(space^2)*(hpwt + air + mpd)",
                         "+ I(space^2) + I(hpwt^3) + I(air^3) + I(mpd^3) + I(space^3)"))
data_transf = data.frame(model.matrix(formula, blp_data))
names(data_transf)
```

```R
y_col = 'y'
d_col = 'log_p'
cluster_cols = c('model.id', 'cdid')
all_z_cols = c('sum.other.hpwt', 'sum.other.mpd', 'sum.other.space')
z_col = all_z_cols[1]
```

```R
dml_df = cbind(blp_data[c(y_col, d_col, cluster_cols)],
               data_transf,
               iv_vars[all_z_cols])
```

### Initialize `DoubleMLClusterData` object

```R
dml_data = DoubleMLClusterData$new(dml_df,
                                   y_col=y_col,
                                   d_cols=d_col,
                                   z_cols=z_col,
                                   cluster_cols=cluster_cols,
                                   x_cols=names(data_transf))
```

```R
print(dml_data)
```

```R
lasso = lrn("regr.cv_glmnet", nfolds = 10, s = "lambda.min")
```

```R
coef_df = data.frame(matrix(NA_real_, ncol = 4, nrow = 1))
colnames(coef_df) = c('zero-way', 'one-way-product', 'one-way-market', 'two-way')
rownames(coef_df) = all_z_cols[1]
se_df = coef_df
n_rep = 10
```

### Two-Way Clustering with Respect to Product and Market

```R
set.seed(1111)
dml_data$z_cols = z_col
dml_data$cluster_cols = c('model.id', 'cdid')
dml_pliv = DoubleMLPLIV$new(dml_data,
                            lasso, lasso, lasso,
                            n_folds=2, n_rep=n_rep)
dml_pliv$fit()
coef_df[1, 4] = dml_pliv$coef
se_df[1, 4] = dml_pliv$se
```

### One-Way Clustering with Respect to the Product

```R
set.seed(2222)
dml_data$z_cols = z_col
dml_data$cluster_cols = 'model.id'
dml_pliv = DoubleMLPLIV$new(dml_data,
                            lasso, lasso, lasso,
                            n_folds=4, n_rep=n_rep)
dml_pliv$fit()
coef_df[1, 2] = dml_pliv$coef
se_df[1, 2] = dml_pliv$se
```

### One-Way Clustering with Respect to the Market

```R
set.seed(3333)
dml_data$z_cols = z_col
dml_data$cluster_cols = 'cdid'
dml_pliv = DoubleMLPLIV$new(dml_data,
                            lasso, lasso, lasso,
                            n_folds=4, n_rep=n_rep)
dml_pliv$fit()
coef_df[1, 3] = dml_pliv$coef
se_df[1, 3] = dml_pliv$se
```

### No Clustering / Zero-Way Clustering

```R
dml_data = DoubleMLData$new(dml_df,
                            y_col=y_col,
                            d_cols=d_col,
                            z_cols=z_col,
                            x_cols=names(data_transf))
```

```R
print(dml_data)
```

```R
set.seed(4444)
dml_data$z_cols = z_col
dml_pliv = DoubleMLPLIV$new(dml_data,
                            lasso, lasso, lasso,
                            n_folds=4, n_rep=n_rep)
dml_pliv$fit()
coef_df[1, 1] = dml_pliv$coef
se_df[1, 1] = dml_pliv$se
```

### Application Results

```R
coef_df
```

```R
se_df
```

## References
Berry, S., Levinsohn, J., and Pakes, A. (1995), Automobile Prices in Market
Equilibrium, Econometrica: Journal of the Econometric Society, 63, 841-890, doi: [10.2307/2171802](https://doi.org/10.2307/2171802).

Cameron, A. C., Gelbach, J. B. and Miller, D. L. (2011), Robust Inference with Multiway Clustering, Journal of Business & Economic Statistics, 29:2, 238-249, doi: [10.1198/jbes.2010.07136](https://doi.org/10.1198/jbes.2010.07136).

Chernozhukov, V., Chetverikov, D., Demirer, M., Duflo, E., Hansen, C., Newey, W. and Robins, J. (2018), Double/debiased machine learning for treatment and structural parameters. The Econometrics Journal, 21: C1-C68, doi: [10.1111/ectj.12097](https://doi.org/10.1111/ectj.12097).

Chiang, H. D., Kato K., Ma, Y. and Sasaki, Y. (2021), Multiway Cluster Robust Double/Debiased Machine Learning, Journal of Business & Economic Statistics, doi: [10.1080/07350015.2021.1895815](https://doi.org/10.1080/07350015.2021.1895815), arXiv: [1909.03489](https://arxiv.org/abs/1909.03489).

## Define Helper Functions for Plotting

```R
library(RColorBrewer)
coul <- rev(colorRampPalette(brewer.pal(8, "RdBu"))(3))
options(repr.plot.width = 10, repr.plot.height = 10)

plt_smpls = function(smpls, n_folds) {
    df = matrix(0, nrow = N*M, ncol = n_folds)
    for (i_fold in 1:n_folds){
        df[smpls$train_ids[[i_fold]], i_fold] = -1
        df[smpls$test_ids[[i_fold]], i_fold] = 1
    }

    heatmap(df, Rowv=NA, Colv=NA, col=coul, cexRow=1.5, cexCol=1.5, scale='none')
}

plt_smpls_cluster = function(smpls_cluster, n_folds, n_folds_per_cluster) {
    #options(repr.plot.width = 6, repr.plot.height = 6)
    plots = list()
    for (i_fold in 1:n_folds){
        mat = matrix(0, nrow = M, ncol = N)
        for (k in smpls_cluster$train_ids[[i_fold]][[1]]) {
            for (l in smpls_cluster$train_ids[[i_fold]][[2]]) {
                mat[k, l] = -1
            }
        }
        for (k in smpls_cluster$test_ids[[i_fold]][[1]]) {
            for (l in smpls_cluster$test_ids[[i_fold]][[2]]) {
                mat[k, l] = 1
            }
        }
        l = (i_fold-1) %% n_folds_per_cluster + 1
        k = ((i_fold-1) %/% n_folds_per_cluster)+1
        df = data.frame(mat)
        cols = names(df)
        names(df) = 1:N
        df$id = 1:N
        df_plot = melt(df,  id.var = 'id')
        df_plot$value = factor(df_plot$value)
        plots[[i_fold]] = ggplot(data = df_plot, aes(x=id, y=variable)) +
          geom_tile(aes(fill=value), colour = "grey50") +
        scale_fill_manual(values = c("darkblue", "white", "darkred")) +
        theme(text = element_text(size=15))
        # ToDo: Add Subplot titles
        if (k == 3) {
            plots[[i_fold]] = plots[[i_fold]] + xlab(expression(paste('Second Cluster Variable ', l)))
        } else {
            plots[[i_fold]] = plots[[i_fold]] + xlab('')
        }
        if (l == 1) {
            plots[[i_fold]] = plots[[i_fold]] + ylab(expression(paste('First Cluster Variable ', k)))
        } else {
            plots[[i_fold]] = plots[[i_fold]] + ylab('')
        }
    }
    return(plots)
}
```