# Finding an isomorphism between unlabeled graphs
## Preliminaries
### Problem definition
Given two anonymized versions of a graph $G=(V,E)$, an interesting problem is to find the true mapping between the nodes of these two graphs. Generally, these two anonymized graphs could also have a slightly different set of edges (assume for example that they are the connections of a set of people in two different online social networks)

This problem can be formulated as follows:
>Given two graphs, $G_1  = (V,E_1)$ and $G_2 = (V,E_2)$ of size $n=|V|$, a
>permutation $\pi$ on $V$ maps the vertices from $G_1$ onto those in $G_2$.
>
>Find the “best” such permutation, denoted by $\pi_0$.

### Measuring the error in a solution
The cost of a permutation can be expressed in many different ways. One such measure is the **edge-inconsistency** between two graphs, $G_1$ and $G_2$ _[Pedarsani2011]_. This represents the number of edges that exist in one graph and not the other, under a matching $\pi$. \begin{equation} \Delta_{\pi} = \sum_{e \in E_1} \mathbb{1}_{\{\pi(e) \not\in E_2\}} + \sum_{e \in E_2} \mathbb{1}_{\{\pi^{-1}(e) \not\in E_1\}}                \end{equation} Here $\mathbb{1}_{\{ A \}}$ denotes the indicator function, and  $\pi(e)$ denotes the mapped version of edge $e = (u,v)$, i.e.  $\pi(e) = (\pi(u), \pi(v))$.

### Problem instances
We will try two different ways of generating the input instances. First, we will just anonymize the graph $G$ and we will try to match it to itself. This is equivalent to generating $G_1$ and $G_2$ by starting from $G$ and randomly renaming the nodes in $V$.

Then, in order to simulate the generation of $G_1$ and $G_2$ based on two different social networks, we will start from $G$ and, with some probability $p$, we will remove edges from $E$.

We will examine these two scenarios separately, starting with the easy one.

```
%matplotlib inline
import networkx as nx
import matplotlib.pyplot as plt
import numpy as np
import graph_isomorphisms as giso

# This helper function will be used to generate the seed for the randomness, so that we can be consistent
import hashlib
def hash_string(text):
    return int(hashlib.md5(text).hexdigest()[0:7], 16)

giso.setup_plots()
```

## Algorithms
Our algorithms are all based on the same intuition. We will represent each node in the graph as a **multi-dimensional point**. Each dimension will represent the similarity of the node to some fixed **beacon node**. For these beacon nodes we assume that we know the correct mapping across the two graphs. This is a reasonable assumption to make, as in most social networks there are some nodes with a high and unique degree.

```
def find_beacons_max_unique(G, num_of_beacons=3):
    # Return the nodes with maximum unique degrees
    inverse_degree = {}
    for k, v in nx.degree(G).iteritems():
        inverse_degree[v] = inverse_degree.get(v, [])
        inverse_degree[v].append(k)
    degrees = sorted(inverse_degree.keys(),reverse=True)
    multiplicity=[len(inverse_degree[degree]) for degree in degrees]
    unique = [i for i,m in enumerate(multiplicity) if m == 1]
    return [inverse_degree[degrees[i]][0] for i in unique[0:num_of_beacons]]

def find_beacons_sample(G, num_of_beacons=3, seed=hash_string('find_beacons_sample')):
    # Sample a beacon based on its degree
    from numpy.random import RandomState
    prng = RandomState(seed)
    degrees = np.array(nx.degree(G).values())
    return prng.choice(np.arange(len(degrees)), num_of_beacons, p=degrees *1./sum(degrees), replace=False )

def find_beacons_sample_inverse(G, num_of_beacons=3, seed=hash_string('find_beacons_sample')):
    # Sample a beacon based on its degree
    from numpy.random import RandomState
    prng = RandomState(seed)
    degrees = np.array(nx.degree(G).values())
    return prng.choice(np.arange(len(degrees)), num_of_beacons, p=(1./degrees) *1./sum(1./degrees), replace=False )
```

### Generate the source graph

```
G = nx.generators.barabasi_albert_graph(50,2, seed=hash_string('example graphs'))
```

### Remove edges with some probability
For each edge in $G$, we remove it with probability $p=0.2$. We repeat this process two times and get back $G_1$ and $G_2$.

```
G1 = giso.noisy_remove(G,0.2,seed=giso.hash_string('noisy removal 2'))
G2 = giso.noisy_remove(G,0.2,seed=giso.hash_string('noisy removal 2'))
```

### Find the beacons
To pick the beacons, we use one of the techniques that we saw above. The one that seemed to be more practical is sampling according to the *inverse* degree.

```
beacons_G1 = find_beacons_sample_inverse(G1,4)
beacons_G2 = beacons_G1
```

### Anonymize the graphs by relabeling the non-beacon nodes
After having choosen the beacon nodes, we need to randomly label the rest of the nodes. We pick a random permutation and we label the nodes accordingly.

```
rest_of_nodes = [x for x in G1.nodes() if x not in beacons_G1]

anonymous_mapping = dict(zip(beacons_G1,beacons_G1))
anonymous_mapping.update(zip(rest_of_nodes,np.random.permutation(rest_of_nodes)))
G2 = nx.relabel_nodes(G2,anonymous_mapping, copy=True)
```

The end result can be seen below. The beacon nodes are the bigger ones, with the labels. For these nodes, we know their correspondence accross the two graphs.

```
giso.plot_graphs(G1, G2, beacons_G1, beacons_G2)
```

## Shortest-path distance projection
Now that we have picked the beacon sets $B_{G_1}$ and $B_{G_2}$, each of size $b$, and we know the true mapping $\pi_0(i), i \in B_{G_1}$, we will project each node in the graph to a multi-dimensional space.

The projection of node $v$ in $G_i$ is the vector $p(v) \in \mathcal{R}^{b}$ defined as
\begin{equation} p(v) = \left\lbrace d(v, B_{G_i}^{(1)}), \ldots, d(v, B_{G_i}^{(b)})  \right\rbrace \end{equation}
where $d(v, B_{G_i}^{(j)})$ is the shortest-path distance between node $v$ and the $j$-th beacon. Notice that when we do this projection, the dimensions are ordered according to the mapping between the beacons.

```
def shortest_path_project(G, beacons):
    projection = np.zeros((G.number_of_nodes() - len(beacons), len(beacons)))

    for i,beacon in enumerate(beacons):
        lengths = nx.shortest_path_length(G, source=beacon)
        node_index = 0
        for node in lengths:
            if node in beacons:
                continue
            projection[node_index][i] = lengths[node]
            node_index += 1
    return projection
```

```
p1 = shortest_path_project(G1,beacons_G1)
p2 = shortest_path_project(G2,beacons_G2)
```

### Projecting to a 2-dimensional space, for visualization purposes

```
from sklearn.manifold import Isomap
from sklearn.metrics.pairwise import euclidean_distances
```

Here, we get a sense of how these $b$-dimensional spaces look like. We project their coordinates to a $2$-dimensional space by using the **Isomap** algorithm. Keep in mind that the darker the marker is, the bigger the number of points that have the same projection.

```
coords1 = giso.isomap_project(p1)
coords2 = giso.isomap_project(p2)
ax = plt.subplot(121)
ax.plot(coords1[:, 0], coords1[:, 1], 'o', alpha = 0.4, markersize=8, c='#4b809c')
remove_border(ax)
ax = plt.subplot(122)
ax.plot(coords2[:, 0], coords2[:, 1], 'D', alpha = 0.3, markersize=8, c='#ba7c5a')
remove_border(ax)
```

### Run k-means on the $b$-dimensional projection space

A natural question is where the nodes that are close with each other in the high-dimensional space are in the graph. We answer this by running a k-medians algorithm in that space and then showing these clusters back on the graphs, side by side.

```
from sklearn.cluster import KMeans
def kmeans_plot(X, mds, ax, n_clusters=8, marker='o'):
    est = KMeans(n_clusters=n_clusters)
    est.fit(X-X.mean())
    kmeans_labels = est.labels_
    for k in range(n_clusters):
        my_members = kmeans_labels == k
        ax.plot(mds[my_members, 0], mds[my_members, 1], marker, alpha = 0.6, markersize=8)
    remove_border(ax)
    return kmeans_labels

def kmedians_plot(X, mds, ax, n_clusters=8, marker='o'):
    est = giso.KMedians(k=n_clusters)
    est.fit(X-X.mean())
    kmedians_labels = est.labels_
    for k in range(n_clusters):
        my_members = kmedians_labels == k
        ax.plot(mds[my_members, 0], mds[my_members, 1], marker, alpha = 0.6, markersize=8)
    remove_border(ax)
    return kmedians_labels
```

```
import brewer2mpl
from sklearn.cluster import KMeans
accent_colors = brewer2mpl.get_map('Accent', 'qualitative',8).mpl_colors
rcParams['axes.color_cycle'] = accent_colors
f, (ax1,ax2) = plt.subplots(1,2)
kmedians_labels1 = kmedians_plot(p1, coords1, ax1, n_clusters=4)
plt.figure()
kmedians_labels2 = kmedians_plot(p2, coords2, ax2, n_clusters=4, marker='D')
```

### Paint the graphs with the cluster colors
Make the clustering above viewable directly on the graph. The goal is to check if there is a certain structure in the coloring of the nodes that we can discern.

```
def cluster_colors(G, beacons, kmeans_labels, colors):
    j=0
    clustered_colors=[]
    for i,node in enumerate(G.nodes()):
        if node not in beacons:
            clustered_colors.append(colors[kmeans_labels[j]])
            j = j + 1
        else:
            clustered_colors.append((101./256,171./256,208./256))
    return clustered_colors

giso.paint_clusters(G1,G2,beacons_G1, beacons_G2, kmeans_labels1, kmeans_labels2)
```

### Match nodes from $G_1$ to nodes in $G_2$ by finding the nearest neigbor
To generate a matching, we assign each node from graph $G_1$ to its closest node in graph $G_2$

```
d= euclidean_distances(p1,p2)
best = {k:v for k,v in giso.find_best_match(d)}
```

## Visualize the matching

In the plot below, we visualize the result of the matching. The x and y axes are the node labels, ordered by the edge centrality of the node. If the cell $(i,j)$ is light green, it means that the node $i \in G_1$ and $j \in G_2$ were used as beacons. If the cell $(i,j)$ is dark green, it means that the method found a correct matching between nodes $i$ and $j$, in $G_1$ and $G_2$ respectively. Finally, a purple square represents a mismatch.

```
(correct, matching_matrix) = giso.count_matches(G1, G2, beacons_G1, beacons_G2, best, anonymous_mapping)
print "Found:", correct

giso.print_matching(G, matching_matrix)
```

### Number of beacons ###
The next thing we will check is how the number of beacons affects the recovery of nodes. For a fixed graph, and fixed noisy versions of the graph, we try different values of beacon percentage. Every time we repeat the experiment, we pick a different beacon set and measure how many of the node labels we can correctly recover.

```
G = nx.generators.barabasi_albert_graph(num_nodes,2)
G1 = giso.noisy_remove(G,0.05)
G2 = giso.noisy_remove(G,0.05)

repeats = 50
num_nodes = 200
beacon_percentages=[0.01,0.02,0.03,0.04,0.05,0.1, 0.15,0.2]
averages = []
for i,beacon_percentage in enumerate(beacon_percentages):
    total = 0
    for tries in xrange(repeats):
        (best, anonymous, beacons_G1, beacons_G2) = giso.match(G1, G2, beacon_percentage)
        (correct, matching_matrix) = giso.count_matches(G1, G2, beacons_G1, beacons_G2, best, anonymous)
        total += correct
    averages.append(total * 1. / repeats)
```

```
beacon_percentages=[0.01,0.02,0.03,0.04,0.05,0.1, 0.15,0.2]
fig, ax = plt.subplots()
ax.plot(averages)
ax.set_xticklabels(beacon_percentages)
ax.set_xlabel("Beacon set site (as a percentage)")
ax.set_ylabel("Average number of recovered labels")
```

We ran this experiment for a graph of 200 nodes. As expected, the number of recovered labels increases as we add more beacons to the graph. However, when the beacons exceed the 5% of the total graph, the number of recovered nodes drops rapidly. The reason might be that there is a predefined set of "possibly identifiable" nodes, that depends on the structure of the graph, and as we add more beacon nodes, we cover parts of this set as well.

###Size of $G$ ###

We tried this experiment for different sizes for $G$, and the trend seems to be similar in all the cases.

### Using Effective Resistance as the distance function ###

As the next step, instead of using the shortest path distance between each node and the beacons, we will use the effective resistance between them.

The **effective resistance** between two nodes is the resistance of the whole graph when voltage is applied to them. We imply here that we view the network as an electrical circuit, where the edges are unit resistors and the nodes are connection points. Consequently, we are allow to attach a battery at any two nodes in the network, and current will flow from one battery terminal to the other.

```
def effective_resistance_project(G, beacons):
    from numpy.linalg import pinv
    projection = np.zeros((G.number_of_nodes() - len(beacons), len(beacons)))
    L = nx.laplacian_matrix(G)
    B = nx.incidence_matrix(G).T
    B_e = B.copy()

    L_pseudo = pinv(L)
    for i in xrange(B.shape[0]):
        min_ace = np.min(np.where(B[i,:] ==1)[1])
        B_e[i, min_ace] = -1

    for i,beacon in enumerate(beacons):
        node_index = 0
        for j,node in enumerate(G.nodes()):
            if node in beacons:
                continue

            battery = np.zeros((B_e.shape[1],1))
            battery[i] = 1
            battery[node_index] = -1

            p = L_pseudo * battery
            projection[node_index][i] = abs(p[i] - p[j])
            node_index += 1
    return projection
```

For a random graph $G$, we plot the nodes by using the effective resistance projection above. We notice that, even though the two graphs are noisy versions of each other, the geometry of the points is very similar in both plots.

```
G = nx.generators.barabasi_albert_graph(50,2)
G1 = giso.noisy_remove(G,0.05)
G2 = giso.noisy_remove(G,0.05)
beacons = giso.find_beacons_sample_inverse(G, num_of_beacons=int(0.05 * G.number_of_nodes()))
```

```
ef_project1  = effective_resistance_project(G1,beacons)
ef_project2 = effective_resistance_project(G2,beacons)
```

    (48L, 2L)
    (48L, 2L)

```
coords1 = ef_project1
coords2 = ef_project2
ax = plt.subplot(121)
ax.plot(coords1[:, 0], coords1[:, 1], 'o', alpha = 0.4, markersize=8, c='#4b809c')
giso.remove_border(ax)
ax = plt.subplot(122)
ax.plot(coords2[:, 0], coords2[:, 1], 'D', alpha = 0.3, markersize=8, c='#ba7c5a')
giso.remove_border(ax)
```

### The effect of the number of beacons when using the effective resistance projection ###

Here, we will repeat the experiment we had above, but we will now use the effective resistance as the similarity function between two nodes.

```
repeats = 50
num_nodes = 200
G = nx.generators.barabasi_albert_graph(num_nodes,2)
G1 = giso.noisy_remove(G,0.05)
G2 = giso.noisy_remove(G,0.05)

beacon_percentages=[0.01,0.02,0.03,0.04,0.05,0.1, 0.15,0.2]
averages = []
for i,beacon_percentage in enumerate(beacon_percentages):
    total = 0
    for tries in xrange(repeats):
        (best, anonymous, beacons_G1, beacons_G2) = giso.match(G1, G2, beacon_percentage, electrical=True)
        (correct, matching_matrix) = giso.count_matches(G1, G2, beacons_G1, beacons_G2, best, anonymous)
        total += correct
    averages.append(total * 1. / repeats)
```

```
beacon_percentages=[0.01,0.02,0.03,0.04,0.05,0.1, 0.15,0.2]
fig, ax = plt.subplots()
ax.plot(averages)
ax.set_xticklabels(beacon_percentages)
ax.set_xlabel("Beacon set site (as a percentage)")
ax.set_ylabel("Average number of recovered labels")
```

As it seems, although the geometry of the projection looks very promising, the nearest neighbor distance fails to identify the correct isomorphism in the presence of noise. The reason for this might be that the projection is very dense, i.e. there are many neighbors to each node even when the neighborhood is small.

## References
Pedarsani, Pedram, and Matthias Grossglauser. "On the privacy of anonymized networks." Proceedings of the 17th ACM SIGKDD international conference on Knowledge discovery and data mining. ACM, 2011.


This tutorial describes how to use the SMT toolbox to do some Bayesian Optimization (EGO method) to solve unconstrained optimization problem

Rémy Priem and Nathalie BARTOLI ONERA/DTIS/M2CI - April 2020

To use SMT models, please follow this link : https://github.com/SMTorg/SMT/blob/master/README.md. The documentation is available here: http://smt.readthedocs.io/en/latest/

The reference paper is available
here https://www.sciencedirect.com/science/article/pii/S0965997818309360?via%3Dihub

or as a preprint: http://mdolab.engin.umich.edu/content/python-surrogate-modeling-framework-derivatives

In this notebook, two examples are presented to illustrate Bayesian Optimization

 - a 1D-example (xsinx function) where the algorithm is explicitely given and the use of different criteria is presented

 - a 2D-exemple (Rosenbrock function) where the EGO algorithm from SMT is used

# Bayesian Optimization

```python

import numpy as np
%matplotlib notebook
import matplotlib.pyplot as plt

plt.ion()

def fun(point):
    return np.atleast_2d((point-3.5)*np.sin((point-3.5)/(np.pi)))

X_plot = np.atleast_2d(np.linspace(0, 25, 10000)).T
Y_plot = fun(X_plot)

lines = []
fig = plt.figure(figsize=[5,5])
ax = fig.add_subplot(111)
true_fun, = ax.plot(X_plot,Y_plot)
lines.append(true_fun)
ax.set_title('$x \sin{x}$ function')
ax.set_xlabel('x')
ax.set_ylabel('y')
plt.show()

#dimension of the problem
ndim = 1
```

Here, the training data are the points xdata=[0,7,25].

```python
x_data = np.atleast_2d([0,7,25]).T
y_data = fun(x_data)
```

Build the GP model with a square exponential kernel with SMT toolbox knowing $(x_{data}, y_{data})$.

```python
from smt.surrogate_models import KPLS, KRG, KPLSK

########### The Kriging model

# The variable 'theta0' is a list of length ndim.
t = KRG(theta0=[1e-2]*ndim,print_prediction = False, corr='squar_exp')

#Training
t.set_training_values(x_data,y_data)
t.train()

# Prediction of the  points for the plot
Y_GP_plot = t.predict_values(X_plot)
Y_GP_plot_var = t.predict_variances(X_plot)
fig = plt.figure(figsize=[5,5])
ax = fig.add_subplot(111)
true_fun, = ax.plot(X_plot,Y_plot)
data, = ax.plot(x_data,y_data,linestyle='',marker='o')
gp, = ax.plot(X_plot,Y_GP_plot,linestyle='--',color='g')
sig_plus = Y_GP_plot+3*np.sqrt(Y_GP_plot_var)
sig_moins = Y_GP_plot-3*np.sqrt(Y_GP_plot_var)
un_gp = ax.fill_between(X_plot.T[0],sig_plus.T[0],sig_moins.T[0],alpha=0.3,color='g')
lines = [true_fun,data,gp,un_gp]
ax.set_title('$x \sin{x}$ function')
ax.set_xlabel('x')
ax.set_ylabel('y')
ax.legend(lines,['True function','Data','GPR prediction','99 % confidence'])
plt.show()
```

Bayesian optimization is defined by Jonas Mockus in (Mockus, 1975) as an optimization technique based upon the minimization of the expected deviation from the extremum of the studied function.

The objective function is treated as a black-box function. A Bayesian strategy sees the objective as a random function and places a prior over it. The prior captures our beliefs about the behavior of the function. After gathering the function evaluations, which are treated as data, the prior is updated to form the posterior distribution over the objective function. The posterior distribution, in turn, is used to construct an acquisition function (often also referred to as infill sampling criterion) that determines what the next query point should be.

One of the earliest bodies of work on Bayesian optimisation that we are aware of is (Kushner, 1962 ; Kushner, 1964). Kushner used Wiener processes for one-dimensional problems. Kushner’s decision model was based on maximizing the probability of improvement, and included a parameter that controlled the trade-off between ‘more global’ and ‘more local’ optimization, in the same spirit as the Exploration/Exploitation trade-off.

Meanwhile, in the former Soviet Union, Mockus and colleagues developed a multidimensional Bayesian optimization method using linear combinations of Wiener fields, some of which was published in English in (Mockus, 1975). This paper also describes an acquisition function that is based on myopic expected improvement of the posterior, which has been widely adopted in Bayesian optimization as the Expected Improvement function.

In 1998, Jones used Gaussian processes together with the expected improvement function to successfully perform derivative-free optimization and experimental design through an algorithm called  Efficient  Global  Optimization, or EGO (Jones, 1998).

## Efficient Global Optimization

In what follows, we describe the Efficient Global Optimization (EGO) algorithm, as published in (Jones, 1998).

Let $F$ be an expensive black-box function to be minimized. We sample $F$ at the different locations  $X = \{x_1, x_2,\ldots,x_n\}$ yielding the responses $Y = \{y_1, y_2,\ldots,y_n\}$. We build a Kriging model (also called Gaussian process) with a mean function $\mu$ and a variance function $\sigma^{2}$.

The next step is to compute the criterion EI. To do this, let us denote:
$$f_{min} = \min \{y_1, y_2,\ldots,y_n\}.$$
The Expected Improvement funtion (EI) can be expressed:
$$E[I(x)] = E[\max(f_{min}-Y, 0)],$$
where $Y$ is the random variable following the distribution $\mathcal{N}(\mu(x), \sigma^{2}(x))$.
By expressing the right-hand side of EI expression as an integral, and applying some tedious integration by parts, one can express the expected improvement in closed form:
$$
E[I(x)] = (f_{min} - \mu(x))\Phi\left(\frac{f_{min} - \mu(x)}{\sigma(x)}\right) + \sigma(x) \phi\left(\frac{f_{min} - \mu(x)}{\sigma(x)}\right)
$$
where $\Phi(\cdot)$ and $\phi(\cdot)$ are respectively the cumulative and probability density functions of $\mathcal{N}(0,1)$.

Next, we determine our next sampling point as :
\begin{align}
x_{n+1} = \arg \max_{x} \left(E[I(x)]\right)
\end{align}

We then test the response $y_{n+1}$ of our black-box function $F$ at $x_{n+1}$, rebuild the model taking into account the new information gained, and research the point of maximum expected improvement again.

We summarize here the EGO algorithm:

EGO(F, $n_{iter}$) \# Find the best minimum of $\operatorname{F}$ in $n_{iter}$ iterations
For ($i=0:n_{iter}$)

* $mod = {model}(X, Y)$  \# surrogate model based on sample vectors $X$ and $Y$
* $f_{min} = \min Y$
* $x_{i+1} = \arg \max {EI}(mod, f_{min})$ \# choose $x$ that maximizes EI
* $y_{i+1} = {F}(x_{i+1})$ \# Probe the function at most promising point $x_{i+1}$
* $X = [X,x_{i+1}]$
* $Y = [Y,y_{i+1}]$
* $i = i+1$

$f_{min} = \min Y$
Return : $f_{min}$ \# This is the best known solution after $n_{iter}$ iterations

Now we want to optimize this function by using Bayesian Optimization and comparing
- Surrogate Based optimization (SBO)
- Expected Improvement criterion (EI)

In a first step we compute the EI criterion

```python
from scipy.stats import norm
from scipy.optimize import minimize

def EI(GP,points,f_min):
    pred = GP.predict_values(points)
    var = GP.predict_variances(points)
    args0 = (f_min - pred)/np.sqrt(var)
    args1 = (f_min - pred)*norm.cdf(args0)
    args2 = np.sqrt(var)*norm.pdf(args0)

    if var.size == 1 and var == 0.0:  # can be use only if one point is computed
        return 0.0

    ei = args1 + args2
    return ei

Y_GP_plot = t.predict_values(X_plot)
Y_GP_plot_var  =  t.predict_variances(X_plot)
Y_EI_plot = EI(t,X_plot,np.min(y_data))

fig = plt.figure(figsize=[10,10])
ax = fig.add_subplot(111)
true_fun, = ax.plot(X_plot,Y_plot)
data, = ax.plot(x_data,y_data,linestyle='',marker='o')
gp, = ax.plot(X_plot,Y_GP_plot,linestyle='--',color='g')
sig_plus = Y_GP_plot+3*np.sqrt(Y_GP_plot_var)
sig_moins = Y_GP_plot-3*np.sqrt(Y_GP_plot_var)
un_gp = ax.fill_between(X_plot.T[0],sig_plus.T[0],sig_moins.T[0],alpha=0.3,color='g')
ax1 = ax.twinx()
ei, = ax1.plot(X_plot,Y_EI_plot,color='red')
lines = [true_fun,data,gp,un_gp,ei]
ax.set_title('$x \sin{x}$ function')
ax.set_xlabel('x')
ax.set_ylabel('y')
ax1.set_ylabel('ei')
fig.legend(lines,['True function','Data','GPR prediction','99 % confidence','Expected Improvement'],loc=[0.13,0.64])
plt.show()

```

Now we compute the EGO method and compare it to other infill criteria
- SBO (surrogate based optimization): directly using the prediction of the surrogate model ($\mu$)
- LCB (Lower Confidence bound): using the confidence interval : $\mu -3 \times \sigma$
- EI for expected Improvement (EGO)

```python
#surrogate Based optimization: min the Surrogate model by using the mean mu
def SBO(GP,point):
    res = GP.predict_values(point)
    return res

#lower confidence bound optimization: minimize by using mu - 3*sigma
def LCB(GP,point):
    pred = GP.predict_values(point)
    var = GP.predict_variances(point)
    res = pred-3.*np.sqrt(var)
    return res

```

```python
IC = 'EI'
```

```python
import matplotlib.image as mpimg
import matplotlib.animation as animation
from IPython.display import HTML

plt.ioff()

x_data = np.atleast_2d([0,7,25]).T
y_data = fun(x_data)

n_iter = 15

gpr = KRG(theta0=[1e-2]*ndim,print_global = False)

for k in range(n_iter):
    x_start = np.atleast_2d(np.random.rand(20)*25).T
    f_min_k = np.min(y_data)
    gpr.set_training_values(x_data,y_data)
    gpr.train()
    if IC == 'EI':
        obj_k = lambda x: -EI(gpr,np.atleast_2d(x),f_min_k)[:,0]
    elif IC =='SBO':
        obj_k = lambda x: SBO(gpr,np.atleast_2d(x))
    elif IC == 'LCB':
        obj_k = lambda x: LCB(gpr,np.atleast_2d(x))

    opt_all = np.array([minimize(lambda x: float(obj_k(x)), x_st, method='SLSQP', bounds=[(0,25)]) for x_st in x_start])
    opt_success = opt_all[[opt_i['success'] for opt_i in opt_all]]
    obj_success = np.array([opt_i['fun'] for opt_i in opt_success])
    ind_min = np.argmin(obj_success)
    opt = opt_success[ind_min]
    x_et_k = opt['x']

    y_et_k = fun(x_et_k)

    y_data = np.atleast_2d(np.append(y_data,y_et_k)).T
    x_data = np.atleast_2d(np.append(x_data,x_et_k)).T

    Y_GP_plot = gpr.predict_values(X_plot)
    Y_GP_plot_var  =  gpr.predict_variances(X_plot)
    Y_EI_plot = -EI(gpr,X_plot,f_min_k)

    fig = plt.figure(figsize=[10,10])
    ax = fig.add_subplot(111)
    if IC == 'LCB' or IC == 'SBO':
        ei, = ax.plot(X_plot,Y_EI_plot,color='red')
    else:
        ax1 = ax.twinx()
        ei, = ax1.plot(X_plot,Y_EI_plot,color='red')
    true_fun, = ax.plot(X_plot,Y_plot)
    data, = ax.plot(x_data[0:k+3],y_data[0:k+3],linestyle='',marker='o',color='orange')
    opt, = ax.plot(x_data[k+3],y_data[k+3],linestyle='',marker='*',color='r')
    gp, = ax.plot(X_plot,Y_GP_plot,linestyle='--',color='g')
    sig_plus = Y_GP_plot+3*np.sqrt(Y_GP_plot_var)
    sig_moins = Y_GP_plot-3*np.sqrt(Y_GP_plot_var)
    un_gp = ax.fill_between(X_plot.T[0],sig_plus.T[0],sig_moins.T[0],alpha=0.3,color='g')
    lines = [true_fun,data,gp,un_gp,opt,ei]
    ax.set_title('$x \sin{x}$ function')
    ax.set_xlabel('x')
    ax.set_ylabel('y')
    ax.legend(lines,['True function','Data','GPR prediction','99 % confidence','Next point to Evaluate','Infill Criteria'])
    plt.savefig('Optimisation %d' %k)
    plt.close(fig)

ind_best = np.argmin(y_data)
x_opt = x_data[ind_best]
y_opt = y_data[ind_best]

print('Results : X = %s, Y = %s' %(x_opt,y_opt))

fig = plt.figure(figsize=[10,10])

ax = plt.gca()
ax.axes.get_xaxis().set_visible(False)
ax.axes.get_yaxis().set_visible(False)

ims = []
for k in range(n_iter):
    image_pt = mpimg.imread('Optimisation %d.png' %k)
    im = plt.imshow(image_pt)
    ims.append([im])

ani = animation.ArtistAnimation(fig, ims,interval=500)
HTML(ani.to_jshtml())

```

 ## Use the EGO from SMT

```python
from smt.applications.ego import EGO
from smt.sampling_methods import LHS

```

* Choose your criterion to perform the optimization: EI, SBO or LCB
* Choose the size of the initial DOE
* Choose the number of EGO iterations

## Try with a 2D function : 2D Rosenbrock function

Rosenbrock Function  in dimension N

$$
f(\mathbf{x}) = \sum_{i=1}^{N-1} 100 (x_{i+1} - x_i^2 )^2 + (1-x_i)^2 \quad \mbox{where} \quad \mathbf{x} = [x_1, \ldots, x_N] \in \mathbb{R}^N.
$$

$$x_i \in [-2,2]$$

```python
#define the rosenbrock function
def rosenbrock(x):
    """
    Evaluate objective and constraints for the Rosenbrock test case:
    """
    n,dim = x.shape

    #parameters:
    Opt =[]
    Opt_point_scalar = 1
    #construction of O vector
    for i in range(0, dim):
        Opt.append(Opt_point_scalar)

    #Construction of Z vector
    Z= np.zeros((n,dim))
    for i in range(0,dim):
        Z[:,i] = (x[:,i]-Opt[i]+1)

    #Sum
    sum1 = np.zeros((n,1))
    for i in range(0,dim-1):
        sum1[:,0] += 100*(((Z[:,i]**2)-Z[:,i+1])**2)+((Z[:,i]-1)**2)

    return sum1

xlimits=np.array([[-2,2], [-2,2]])
```

```python
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
from matplotlib import cm

#To plot the Rosenbrock function
num_plot = 50 #to plot rosenbrock
x = np.linspace(xlimits[0][0],xlimits[0][1],num_plot)
res = []
for x0 in x:
    for x1 in x:
        res.append(rosenbrock(np.array([[x0,x1]])))
res = np.array(res)
res = res.reshape((50,50)).T
X,Y = np.meshgrid(x,x)
fig = plt.figure(figsize=[10,10])
ax = fig.gca(projection='3d')
surf = ax.plot_surface(X, Y, res, cmap=cm.coolwarm,
                       linewidth=0, antialiased=False,alpha=0.5)
plt.title(' Rosenbrock function')
plt.xlabel('x1')
plt.ylabel('x2')
plt.show()
```

```python
criterion='EI' #'EI' or 'SBO' or 'LCB'

#number of points in the initial DOE
ndoe = 10 #(at least ndim+1)

#number of iterations with EGO
n_iter = 50

#Build the initial DOE, add the random_state option to have the reproducibility of the LHS points
sampling = LHS(xlimits=xlimits, random_state=1)
xdoe = sampling(ndoe)

#EGO call
ego = EGO(n_iter=n_iter, criterion=criterion, xdoe=xdoe, xlimits=xlimits)

x_opt, y_opt, ind_best, x_data, y_data = ego.optimize(fun=rosenbrock)

print('Xopt for Rosenbrock ', x_opt,y_opt, ' obtained using EGO criterion = ', criterion )
print('Check if the optimal point is Xopt= (1,1) with the Y value=0')
print('if not you can increase the number of iterations with n_iter but the CPU will increase also.')
print('---------------------------')
```

```python
#To plot the Rosenbrock function
#3D plot
x = np.linspace(xlimits[0][0],xlimits[0][1],num_plot)
res = []
for x0 in x:
    for x1 in x:
        res.append(rosenbrock(np.array([[x0,x1]])))
res = np.array(res)
res = res.reshape((50,50)).T
X,Y = np.meshgrid(x,x)
fig = plt.figure(figsize=(10, 10))
ax = fig.gca(projection='3d')
surf = ax.plot_surface(X, Y, res, cmap=cm.coolwarm,
                       linewidth=0, antialiased=False,alpha=0.5)
#to add the points provided by EGO
ax.scatter(x_data[:ndoe,0],x_data[:ndoe,1],y_data[:ndoe],zdir='z',marker = '.',c='k',s=100, label='Initial DOE')
ax.scatter(x_data[ndoe:,0],x_data[ndoe:,1],y_data[ndoe:],zdir='z',marker = 'x',c='r', s=100, label= 'Added point')
ax.scatter(x_opt[0],x_opt[1],y_opt,zdir='z',marker = '*',c='g', s=100, label= 'EGO optimal point')

plt.title(' Rosenbrock function during EGO algorithm')
plt.xlabel('x1')
plt.ylabel('x2')
plt.legend()
plt.show()

#2D plot
#to add the points provided by EGO
plt.plot(x_data[:ndoe,0],x_data[:ndoe,1],'.', label='Initial DOE')
plt.plot(x_data[ndoe:,0],x_data[ndoe:,1],'x', c='r', label='Added point')
plt.plot(x_opt[:1],x_opt[1:],'*',c='g', label= 'EGO optimal point')
plt.plot([1], [1],'*',c='m', label= 'Optimal point')

plt.title(' Rosenbrock function during EGO algorithm')
plt.xlabel('x1')
plt.ylabel('x2')
plt.legend()
plt.show()

```

We can now compare the results by using only the mean information provided by surrogate model approximation

```python
criterion='SBO' #'EI' or 'SBO' or 'LCB'

#number of points in the initial DOE
ndoe = 10 #(at least ndim+1)

#number of iterations with EGO
n_iter = 50

#Build the initial DOE
sampling = LHS(xlimits=xlimits, random_state=1)
xdoe = sampling(ndoe)

#EGO call
ego = EGO(n_iter=n_iter, criterion=criterion, xdoe=xdoe, xlimits=xlimits)

x_opt, y_opt, ind_best, x_data, y_data = ego.optimize(fun=rosenbrock)

print('Xopt for Rosenbrock ', x_opt, y_opt, ' obtained using EGO criterion = ', criterion)
print('Check if the optimal point is Xopt=(1,1) with the Y value=0')
print('---------------------------')
```

```python

```


# Project 5, deadline  December 16, 2020

**Partial Differential equations**

Date: **Fall semester 2020**

## Theoretical background and description of the system

This project is more numerically directed, with tests
of algorithms for solving partial differential equations using finite difference schemes.
Chapter 10 of the lecture notes provides the necessary theoretical background.

**For this project you can collaborate with fellow students and you can  hand in a common report.**
This project (together with projects 3 and 4) counts 1/3 of the final mark.

The project has a strong mathematical axis. However, for those interested, it should be straightforward to replace the dimensionless analysis in  parts 5a-5f) with specific boundary and initial conditions. This is done in part 5g), with a focus on the heat equation and geophysical processes. That part is optional, but gives an additional 30 points to a score of 100.

The physical problem can be that of the temperature gradient in a rod of length $L=1$ or that of channel flow
between two flat and infinite plates at $x=0$ and $x=1$, where the fluid is initially at rest
and the plate $x=1$ is given a sudden initial movement.  We are looking at a one-dimensional
problem

$$
\frac{\partial^2 u(x,t)}{\partial x^2} =\frac{\partial u(x,t)}{\partial t}, t> 0, x\in [0,L]
$$

or

$$
u_{xx} = u_t,
$$

with initial conditions, i.e., the conditions at $t=0$,

$$
u(x,0)= 0 \hspace{0.5cm} 0 < x < L
$$

with $L=1$ the length of the $x$-region of interest. The
boundary conditions are

$$
u(0,t)= 0 \hspace{0.5cm} t \ge 0,
$$

and

$$
u(L,t)= 1 \hspace{0.5cm} t \ge 0.
$$

The function $u(x,t)$  can be the temperature gradient of a  rod or represent the fluid velocity
in a direction parallel to the plates, that is normal to the $x$-axis. In the latter case,
for small $t$, only the part of the fluid
close to the moving plate is set in significant  motion, resulting in a thin boundary layer at $x=1$.
As time increases, the velocity approaches a linear variation with $x$. In this case, which can be derived
from the incompressible Navier-Stokes, the above equations constitute a model for
studying friction between moving surfaces separated by a thin fluid film.

In this project we want to study the numerical stability of three methods for partial differential equations
(PDEs).
These methods we will discuss are
the explicit forward Euler algorithm with discretized versions of time given by a forward formula and a centered difference in space resulting in

$$
u_t\approx \frac{u(x,t+\Delta t)-u(x,t)}{\Delta t}=\frac{u(x_i,t_j+\Delta t)-u(x_i,t_j)}{\Delta t}
$$

and

$$
u_{xx}\approx \frac{u(x+\Delta x,t)-2u(x,t)+u(x-\Delta x,t)}{\Delta x^2},
$$

or

$$
u_{xx}\approx \frac{u(x_i+\Delta x,t_j)-2u(x_i,t_j)+u(x_i-\Delta x,t_j)}{\Delta x^2}.
$$

We will also discuss the implicit Backward Euler with

$$
u_t\approx \frac{u(x,t)-u(x,t-\Delta t)}{\Delta t}=\frac{u(x_i,t_j)-u(x_i,t_j-\Delta t)}{\Delta t}
$$

and

$$
u_{xx}\approx \frac{u(x+\Delta x,t)-2u(x,t)+u(x-\Delta x,t)}{\Delta x^2},
$$

or

$$
u_{xx}\approx \frac{u(x_i+\Delta x,t_j)-2u(x_i,t_j)+u(x_i-\Delta x,t_j)}{\Delta x^2},
$$

Finally we use the implicit Crank-Nicolson scheme with  a time-centered scheme at $(x,t+\Delta t/2)$

$$
u_t\approx \frac{u(x,t+\Delta t)-u(x,t)}{\Delta t}=\frac{u(x_i,t_j+\Delta t)-u(x_i,t_j)}{\Delta t}.
$$

The corresponding spatial second-order derivative reads

$$
\begin{align*}
u_{xx}\approx &\frac{1}{2}\left(\frac{u(x_i+\Delta x,t_j)-2u(x_i,t_j)+u(x_i-\Delta x,t_j)}{\Delta x^2}\right. +\\
&\left. \frac{u(x_i+\Delta x,t_j+\Delta t)-2u(x_i,t_j+\Delta t)+u(x_i-\Delta x,t_j+\Delta t)}{\Delta x^2}\right).
\end{align*}
$$

Note well that we are using a time-centered scheme wih $t+\Delta t/2$ as center.

### Project 5a): Setting up the algorithms

Write down the algorithms for these three methods and the equations you need to implement.
For the implicit schemes show that the equations lead to a tridiagonal matrix system for the new values.

### Project 5b): Truncation errors and analytic solutions

Find the truncation errors of these three schemes and investigate their stability properties.
Find also the analytic solution to the continuous problem.

### Project 5c): Implementation

Implement the three algorithms in the same code and perform tests of the solution
for these three approaches
for $\Delta x=1/10$, $\Delta x=1/100$ using  $\Delta t$ as dictated by the stability limit of the explicit scheme.
Study the solutions at two time points $t_1$ and $t_2$ where $u(x,t_1)$ is smooth but still significantly curved
and $u(x,t_2)$ is almost linear, close to the stationary state.
Remember that forsolving the tridiagonal equations you can use your code from project 1.
### Project 5d):

Compare the solutions at $t_1$ and $t_2$ with the analytic result for the continuous problem.
Which of the schemes would you classify as the best?

### Project 5e): Moving to two dimensions

Extend the code you have developed here to two
  dimensions.
It means that we deal with a $2+1$ dimensional problem. Our differential equation becomes

$$
\frac{\partial^2 u(x,y,t)}{\partial x^2}+\frac{\partial^2 u(x,y,t)}{\partial y^2} =\frac{\partial u(x,y,t)}{\partial t}, t> 0, x,y\in [0,1],
$$

where we now have made a model with a square lattice for $x$ and $y$.
How would you extend the boundary conditions from one dimension to two dimensions? And can you
  find a closed form solution here as well?  It is left to you to decide upon what kind of boundary conditions
you deem appropriate.
### Project 5f): Solving the two-dimensional equations numerically

Here you can choose between an explicit or an implicit scheme.
For the implicit scheme discussed in chapter 10 of the lecture notes, you need to use for example an iterative method like Jacobi's.

Outline the algorithm for solving the two-dimensional diffusion equation and
implement the explicit or the implicit scheme as function of $\Delta x$ (assuming
$\Delta x = \Delta y$) and $\Delta t$. Solve the equations numerically and give a critical discussion of your results.
Compare your results with the closed-form answer. Discuss the stability
of the solution as function of different values of $\Delta x$  and $\Delta t$.

You should, if you have time, try to  parallelize the equations using MPI or OpenMP.

### Project 5g) Temperature distribution in the lithosphere (this part is optional and gives an additional score of 30 points)

The Scientific background is  based on geological evidence at the surface. Geologists have proposed
that there was an active subduction zone on the western coast of
Norway about 1Gy ago. When subducting, the oceanic lithosphere
releases water and other chemical components that go mostly to the
surface but may also be trapped in the mantle above the subducting
slab. This process is called the refertilization of the mantle wedge
(see Figure). These chemical components contain more radioactive
elements than the normal mantle. We can therefore expect that ancient
mantle wedges are enriched in radioactive elements and therefore
warmer than normal mantle. The question is how much warmer they are
and if we can find evidence for this enrichment in geophysical
studies.

Cartoon of the enrichment of the mantle by a subduction slab.

The purpose of this part is to
calculate the thermal evolution of the lithosphere up to present following the emplacement of radioactive elements in the mantle wedge 1 Gy ago.

The equation to compute is the time evolution of the temperature distribution is the heat equation, which will be used only in 2D here, namely

$$
\begin{equation}
\vec \nabla(k \vec \nabla T) + Q = \rho c_p \frac{\partial T}{\partial t}
\label{eq:heat2D} \tag{1}
\end{equation}
$$

where $T$ is the temperature, $\rho$ is the density, $k$ is the thermal conductivity, $c_p$ is the specific heat capacity and $Q$ is the heat production.

With the boundary conditions and initial conditions, you should rescale your equations in order to use the results from parts 5a-5f). The paramaters of the model are described here.

### Parameters of the model

The boundary conditions are $8^{\circ}\,C$ at the surface and $1300^{\circ}C$ at the bottom at a depth of 120 km.

In this study, we assume a constant density for the lithosphere of $3.5 10^3\,Kg/m^3$,
a constant thermal conductivity of $2.5\,W/m/^{\circ}C$
and a constant specific heat capacity of $1000\,J/Kg/^{\circ}C^{-1}$.

The heat production, caused by radioactive decay, cannot be taken as uniform in the whole lithosphere as radioactive elements are much more abundant in the crust than in the mantle.
We need therefore to separate the lithosphere in 3 units: the upper crust from 0 to 20 km depth, the lower crust from 20 to 40 km depth and the mantle from 40 to 120 km depth, where the heat production is set respectively to $1.4\,\mu W/m^3$, $0.35\,\mu W/m^3$ and $0.05\,\mu W/m^3$.

### Temperature distribution before radioactive enrichment

Before radioactive enrichment, the temperature is steady-state and depends only on depth. In order to test your results you should
try to obtain analytical results for  the temperature profile as function of depth neglecting the presence of radioactice sources and. Compare thereafter these results with your numerical results.  Include then the radioactive sources and compare eventual analytical results with numerical calculations.

### Temperature distribution after radioactive enrichment

We assume that the lithosphere above the slab gets enriched in
radioactive elements U, Th and K 1 Gy ago and that this results in an
additional heat production of $0.5\mu W/m^3$ over the whole depth of
the mantle (not the crust) in a 150 km wide area above the
slab. Neglecting all other thermal perturbations that this slab could
produce, and assuming that the heat production will remain constant
over geological time, compute the thermal evolution of the lithosphere
from 1 Gy to present.

The radioactivity will decrease with time. Assuming that the
additional heat is produced at 40% by U, 40% by Th and 20% by K,
which have halflives of 4.47 Gy, 14.0 Gy and 1.25 Gy respectively,
compute the thermal evolution of the lithosphere and compare with the
result obtained above.

### References

A very good reference is the textbook by [Winther and Tveito on partial differential equations](http://www.springer.com/us/book/9783540225515). It is available online [from the University library](https://vpn2.uio.no/+CSCO+0h756767633A2F2F797661782E66636576617472652E70627A++/book/10.1007/b138016/page/1).

## Introduction to numerical projects

Here follows a brief recipe and recommendation on how to write a report for each
project.

  * Give a short description of the nature of the problem and the eventual  numerical methods you have used.

  * Describe the algorithm you have used and/or developed. Here you may find it convenient to use pseudocoding. In many cases you can describe the algorithm in the program itself.

  * Include the source code of your program. Comment your program properly.

  * If possible, try to find analytic solutions, or known limits in order to test your program when developing the code.

  * Include your results either in figure form or in a table. Remember to        label your results. All tables and figures should have relevant captions        and labels on the axes.

  * Try to evaluate the reliabilty and numerical stability/precision of your results. If possible, include a qualitative and/or quantitative discussion of the numerical stability, eventual loss of precision etc.

  * Try to give an interpretation of you results in your answers to  the problems.

  * Critique: if possible include your comments and reflections about the  exercise, whether you felt you learnt something, ideas for improvements and  other thoughts you've made when solving the exercise. We wish to keep this course at the interactive level and your comments can help us improve it.

  * Try to establish a practice where you log your work at the  computerlab. You may find such a logbook very handy at later stages in your work, especially when you don't properly remember  what a previous test version  of your program did. Here you could also record  the time spent on solving the exercise, various algorithms you may have tested or other topics which you feel worthy of mentioning.

## Format for electronic delivery of report and programs

The preferred format for the report is a PDF file. You can also use DOC or postscript formats or as an ipython notebook file.  As programming language we prefer that you choose between C/C++, Fortran2008 or Python. The following prescription should be followed when preparing the report:

  * Use **Canvas** to hand in your projects, log in  at  <https://www.uio.no/english/services/it/education/canvas/> with your normal UiO username and password.

  * Upload **only** the report file!  For the source code file(s) you have developed please provide us with your link to your github domain.  The report file should include all of your discussions and a list of the codes you have developed.  Do not include library files which are available at the course homepage, unless you have made specific changes to them. Alternatively, you can just upload the address to your GitHub or GitLab repository.

  * In your git repository, please include a folder which contains selected results. These can be in the form of output from your code for a selected set of runs and input parameters.

  * In this and all later projects, you should include tests (for example unit tests) of your code(s).

  * Comments  from us on your projects, approval or not, corrections to be made  etc can be found under your **Canvas** domain and are only visible to you and the teachers of the course.

Finally,
we encourage you to work two and two together. Optimal working groups consist of
2-3 students. You can then hand in a common report.


```python
%matplotlib inline
import matplotlib.pyplot as plt
import seaborn as sns
sns.set_theme(context='notebook', style='white', font_scale=1.6, color_codes=True)
plt.style.use('dark_background')
```

# Sparse Regression

Here we compare models of fully Bayesian sparse linear regression with pairwais interaction terms. We start by introducing hierarchical shrinkage priors in the form of regularised horseshoe distribution [1], next we introduce QR decomposition and illustrate its utility in the presence of highly corelated predictors (elements of the design matrix) [2], and finally introduce the sparse kernel interaction model (SKIM) proposed in [3]. SKIM is particularly suitable for situations with many feature dimensions (large D) but not too many datapoints (small N).

We will consider a quadratic regressor of the form:

\begin{equation}
    \begin{split}
        f\left( X, \pmb{\beta} \right) &= \alpha + \sum_i \beta_i X_i + \sum_{i<j} \beta_{ij} X_i X_j \\
        y &\sim p \left(y | f\left( X, \pmb{\beta} \right), \theta  \right)
    \end{split}
\end{equation}

In the case of normal likelihood, that is:

$$ p \left(y | f\left( X, \pmb{\beta} \right), \theta  \right) = \mathcal{N} \left(y; f(X, \pmb{\beta}), \sigma^2 \right) $$

we are dealing with a linear regression problem. In the case of Bernoulli likelihood, that is:

$$ p \left(y | f\left( X, \pmb{\beta} \right), \theta  \right) = s \left( f(X, \pmb{\beta}) \right)^y \left(1 - s \left( f(X, \pmb{\beta}) \right) \right)^{1-y}; \:\: \text{for } y \in \{ 0, 1 \} $$

we are dealing with a logistic regression problem.

**References:**

    1. Piironen, Juho, and Aki Vehtari (2017),
        "Sparsity information and regularization in the horseshoe and other shrinkage priors."
        Electronic Journal of Statistics 11.2 : 5018-5051.

    2. https://mc-stan.org/users/documentation/case-studies/qr_regression.html

    3. Raj Agrawal, Jonathan H. Huggins, Brian Trippe, Tamara Broderick (2019),
       "The kernel interaction trick: Fast bayesian discovery of pairwise interactions in high dimensions."
       International Conference on Machine Learning. PMLR.

```python
import itertools
import numpyro
numpyro.set_platform("cpu")
num_chains=1
numpyro.set_host_device_count(num_chains)
numpyro.enable_validation(True)

import numpy as np

import jax
from jax import vmap, nn
import jax.numpy as jnp
import jax.random as random
# from jax.scipy.linalg import cho_factor, cho_solve, solve_triangular
# from scipy.cluster.vq import kmeans2

import numpyro.distributions as dist
from numpyro import handlers
from numpyro.infer import MCMC, NUTS, init_to_feasible
from numpyro.infer.reparam import TransformReparam
from numpyro.distributions.transforms import Transform, AffineTransform, ComposeTransform
```

```python
# Create artificial regression dataset where only S out of P feature
# dimensions contain signal. If pairwise=True then there is a single
# pairwise interaction between the first and second dimensions.
def get_data(X, S=2, sigma=1., pairwise=False, logistic=False):
    N, D = X.shape
    assert S < D and D > 1 and S > 0

    np.random.seed(101)

    # generate S coefficients with non-negligible magnitude
    W = 1.0 + 2.0 * np.random.rand(S)
    if pairwise:
        # generate data using the S coefficients and a single pairwise interaction
        f = np.sum(X[:, 0:S] * W, axis=-1) + X[:, 0] * X[:, 1]
    else:
        f = np.sum(X[:, 0:S] * W, axis=-1)

    if logistic:
        Y = random.bernoulli(random.PRNGKey(0), p=nn.sigmoid(f))
        return X, Y, W, 1.0
    else:
        Y = f + sigma * np.random.randn(N)
        Y -= jnp.mean(Y)
        Y_std = jnp.std(Y)
        return X, Y/Y_std, W / Y_std, 1.0 / Y_std
```

```python
import numpyro.distributions as dist

S = 5
D = 50
N = 200

X = np.random.randn(N, D)
X, Y1, expected_thetas1, expected_pairwise = get_data(X, S=S, pairwise=False, logistic=False)

print(expected_thetas1)
```

    [0.5007176  0.5274526  0.26034698 0.33081773 0.5839136 ]

# Bayesian linear regression

Given a normal likelihood
$$ p \left(y | f\left( X, \pmb{\beta} \right), \theta  \right) = \mathcal{N} \left(y; f(X, \pmb{\beta}), \sigma^2 \right) $$

and the following prior over regression coefficients
$$ p(\pmb{\beta}) = \mathcal{N}_D\left(\pmb{\beta}; 0, \tau^2 \pmb{\Lambda} \right); \:\: \text{for } \pmb{\Lambda} = diag(\lambda_1^2, \ldots, \lambda_D^2)$$

The posterior for the coefficients $\pmb{\beta}$ given the (hyper)paramters and data $\mathcal{D} = \left( \pmb{X}, \pmb{y} \right)$ is
\begin{equation}
    \begin{split}
        p\left(\pmb{\beta}|\pmb{\Lambda}, \tau, \sigma, \mathcal{D} \right) &= \mathcal{N}\left(\pmb{\beta}; \bar{\pmb{\beta}}, \pmb{\Sigma} \right) \\
        \hat{\pmb{\beta}} &= \left( \pmb{X}^T \pmb{X} \right)^{-1} \pmb{X}^T \pmb{y} \\
        \bar{\pmb{\beta}} &= \tau^2 \pmb{\Lambda} \left( \tau^2 \pmb{\Lambda} + \sigma^2 \left(\pmb{X}^T \pmb{X} \right)^{-1} \right)^{-1} \hat{\pmb{\beta}}\\
        \pmb{\Sigma} &= \left( \frac{1}{\tau^2} \pmb{\Lambda}^{-1} + \frac{1}{\sigma^2} \pmb{X}^T \pmb{X} \right)^{-1}
    \end{split}
\end{equation}

```python
def blr(X, y, tau=1, sigma=1):
    # bayesian linear regression
    N, D = X.shape
    Lambda = np.eye(D)
    S = X.T @ X
    hat_beta = np.linalg.inv(S) @ X.T @ y
    bar_beta = tau**2 * Lambda @ np.linalg.inv(tau**2*Lambda + sigma**2 * np.linalg.inv(S)) @ hat_beta
    Sigma = np.linalg.inv( np.linalg.inv(Lambda)/tau**2 + S/sigma**2)
    return bar_beta, hat_beta, Sigma

fig, axes = plt.subplots(3, 1, figsize=(16, 10), sharex=True, sharey=True)
tau = [1., .1, .01]
for i in range(3):
    bar_beta, hat_beta, Sigma = blr(X, Y1, tau=tau[i])
    axes[i].errorbar(np.arange(1, D + 1), bar_beta, yerr=3*np.sqrt(np.diag(Sigma)), color='r', marker='o', ls='', label=r'$\bar{\beta}$')
    axes[i].scatter(np.arange(1, S+1), expected_thetas1, marker='*', s=100, label=r'true $\beta$')
    axes[i].scatter(np.arange(1, D + 1) + .25, hat_beta, marker='d', label=r'$\hat{\beta}$')
    axes[i].hlines(0, 1, D + 1, 'w', ls='--')
    axes[i].set_ylabel(r'$\beta$')
    axes[i].set_title(r'$\tau={}$'.format(tau[i]))

axes[-1].set_xticks([])
axes[-1].set_xlabel(r'$d$')
axes[0].legend(bbox_to_anchor=(1,1), loc="upper left", fontsize=12)
sns.despine()

# fig.savefig('blr_example.png', transparent=True, bbox_inches='tight', dpi=300)
```

```python
from jax import nn, lax
from numpyro.distributions import constraints

class QRTransform(Transform):
    domain = constraints.real_vector
    codomain = constraints.real_vector

    def __init__(self, R, R_inv):
        if jnp.ndim(R) != 2:
            raise ValueError(
                "Only support 2-dimensional R matrix. "
            )
        self.R = R
        self.R_inv = R_inv

    def __call__(self, x):
        return jnp.squeeze(
            jnp.matmul(self.R, x[..., jnp.newaxis]), axis=-1
        )

    def _inverse(self, y):
        return jnp.squeeze(
            jnp.matmul(self.R_inv, y[..., jnp.newaxis]), axis=-1
        )

    def log_abs_det_jacobian(self, x, y, intermediates=None):
        return jnp.broadcast_to(
            jnp.log(jnp.diagonal(self.R, axis1=-2, axis2=-1)).sum(-1),
            jnp.shape(x)[:-1],
        )

    def forward_shape(self, shape):
        if len(shape) < 1:
            raise ValueError("Too few dimensions on input")
        return lax.broadcast_shapes(shape, self.R.shape[:-1])

    def inverse_shape(self, shape):
        if len(shape) < 1:
            raise ValueError("Too few dimensions on input")
        return lax.broadcast_shapes(shape, self.R.shape[:-1])

class BayesRegression(object):
    def __init__(self, rng_key, X, tau=1., p0=None, fixed_lam=True, with_qr=False, reg_type='linear'):
        self.N, self.D = X.shape
        self.X = X
        self.rng_key = rng_key
        self.tau = tau
        self.fixed_lam = fixed_lam
        self.with_qr = with_qr
        self.p0 = p0
        self.type = reg_type # type of the rergression problem

        if self.with_qr:
            # use QR decomposition
            self.Q, self.R = jnp.linalg.qr(X)
            self.R_inv = jnp.linalg.inv(self.R)

    def model(self, obs=None):
        sigma_sqr = numpyro.sample('sigma^2', dist.InverseGamma(2., 2.))
        sigma = numpyro.deterministic('sigma', jnp.sqrt(sigma_sqr))

        if self.p0 is None:
            tau = self.tau * sigma
        else:
            tau0 = self.p0 * sigma / ((self.D - self.p0) * jnp.sqrt(self.N))
            tau = numpyro.sample('tau', dist.HalfCauchy(tau0))

        if self.fixed_lam:
            lam = numpyro.deterministic('lam', jnp.ones(self.D))
        else:
            lam = numpyro.sample('lam', dist.HalfCauchy(1.).expand([self.D]).to_event(1))

        if self.with_qr:
            rt = QRTransform(self.R, self.R_inv)
            aff = AffineTransform(0., tau * lam)
            ct = ComposeTransform([aff, rt])
            with handlers.reparam(config={"theta": TransformReparam()}):
                theta = numpyro.sample(
                    'theta',
                    dist.TransformedDistribution(dist.Normal(0., 1.).expand([D]).to_event(1), ct)
                )

            numpyro.deterministic('beta', rt.inv(theta))
            tmp = self.Q.dot(theta)

        else:
            aff = AffineTransform(0., tau * lam)
            with handlers.reparam(config={"beta": TransformReparam()}):
                beta = numpyro.sample(
                    'beta',
                    dist.TransformedDistribution(dist.Normal(0., 1.).expand([self.D]).to_event(1), aff)
                )
            tmp = self.X.dot(beta)

        alpha = numpyro.sample('alpha', dist.Normal(0., 10.))
        mu = numpyro.deterministic('mu', alpha + tmp)

        with numpyro.plate('data', self.N):
            if self.type == 'linear':
                numpyro.sample('obs', dist.Normal(mu, sigma), obs=obs)
            elif self.type == 'logistic':
                numpyro.sample('obs', dist.Bernoulli(logits=mu), obs=obs)
            elif self.type == 'possion':
                numpyro.sampel('obs', dist.Poisson(jnp.exp(mu)), obs=obs)

    def fit(self, data, num_samples=1000, warmup_steps=1000, num_chains=1, summary=True, progress_bar=True):
        self.rng_key, _rng_key = random.split(self.rng_key)

        nuts_kernel = NUTS(self.model, init_strategy=init_to_feasible)
        mcmc = MCMC(nuts_kernel,
                    num_warmup=warmup_steps,
                    num_samples=num_samples,
                    num_chains=num_chains,
                    chain_method='vectorized',
                    progress_bar=progress_bar)
        mcmc.run(_rng_key, obs=data)

        if summary:
            mcmc.print_summary()

        samples = mcmc.get_samples(group_by_chain=False)
        self.mcmc = mcmc
        self.samples = samples

        return samples
```

```python
rng_key = random.PRNGKey(101)
fig, axes = plt.subplots(3, 1, figsize=(16, 10), sharex=True, sharey=True)
tau = [1., .1, .01]
for i in range(3):
    rng_key, _rng_key = random.split(rng_key)
    BLR = BayesRegression(rng_key, jnp.array(X), tau=tau[i], fixed_lam=True)
    samples = BLR.fit(jnp.array(Y1), summary=False, num_chains=num_chains)

    sns.boxenplot(data=np.array(samples['beta']), color='c', ax=axes[i], showfliers=False)
    axes[i].scatter(np.arange(S), expected_thetas1, color='w', marker='*', s=100, label=r'true $\beta$')
    axes[i].scatter(np.arange(D) + .25, hat_beta, color='r', marker='d', label=r'$\hat{\beta}$')
    axes[i].hlines(0, 0, D, 'w', ls='--')
    axes[i].set_ylabel(r'$\beta$')
    axes[i].set_title(r'$\tau={}$'.format(tau[i]))

axes[-1].set_xticks([])
axes[-1].set_xlabel(r'$d$')
axes[0].legend(bbox_to_anchor=(1,1), loc="upper left")
sns.despine()

# fig.savefig('example_sampling.png', transparent=True, bbox_inches='tight', dpi=300)
```

# Horseshoe prior

The horseshoe prior consist of a half-cauchy prior distribution for individual $\lambda_d$ parameters

$$ \lambda_d \sim C^+(0, 1) = \begin{cases} \frac{2}{\pi} \frac{1}{1 + \lambda_d^2}, &\text{for } \lambda_d \geq 0 \\ 0 & \text{otherwise}\end{cases}$$

The origin of the name is related to the corresponding shape of the shrinkage factor $\kappa_d$ under the given prior over $\lambda_d$.
The shrinkage factor depends on $\lambda_d$ as follows:

$$ \bar{\beta}_d = (1 - \kappa_d) \hat{\beta}_d$$
$$ \kappa_d = \frac{1}{1 + N \tau^2 \lambda_d^2 s_d^2}$$
$$ Var(\pmb{x}_d) = s_d^2$$

Then the corresponding prior over shrinkage factor has the following functional form

$$ p(\kappa_d) = \frac{1}{\pi} \frac{a_d}{(a_d^2-1)\kappa_d + 1} \frac{1}{\sqrt{\kappa_d}\sqrt{1 - \kappa_d}}$$
$$ a_d = N \sigma^{-1} \tau s_d$$
When $a_d=1$ we have a prior over $\kappa_d$ in the shape of a horseshoe.

```python
kappa = np.arange(0.01, 1, 0.01)
plt.plot(kappa, 1/(np.sqrt(kappa)*np.sqrt(1 - kappa)))
plt.ylabel(r'$p(\kappa_d)$');
plt.xlabel(r'$\kappa_d$');
plt.title(r'$\tau \sigma^{-1} \sqrt{N} s_d = 1$');
# plt.savefig('horseshoe.png', transparent=True, bbox_inches='tight', dpi=300)
```

In what follows we illustrate the effect of horseshoe prior on the results for different values of global shrinkage parameter $\tau$

```python
rng_key = random.PRNGKey(101)
fig, axes = plt.subplots(3, 1, figsize=(16, 10), sharex=True, sharey=True)
tau = [1., .1, .01]
for i in range(3):
    rng_key, _rng_key = random.split(rng_key)
    BLR = BayesRegression(rng_key, jnp.array(X), tau=tau[i], fixed_lam=False)
    samples = BLR.fit(jnp.array(Y1), summary=False, num_chains=num_chains)

    sns.boxenplot(data=np.array(samples['beta']), color='c', ax=axes[i], showfliers=False)
    axes[i].scatter(np.arange(S), expected_thetas1, color='w', marker='*', s=100, label=r'true $\beta$')
    axes[i].scatter(np.arange(D) + .25, hat_beta, color='r', marker='d', label=r'$\hat{\beta}$')
    axes[i].hlines(0, 0, D, 'w', ls='--')
    axes[i].set_ylabel(r'$\beta$')
    axes[i].set_title(r'$\tau={}$'.format(tau[i]))

axes[-1].set_xticks([])
axes[0].legend(loc="upper right", fontsize=12)
sns.despine()
# fig.savefig('example_sampling2.png', transparent=True, bbox_inches='tight', dpi=300)
```

Finally we can also introduce a prior over $\tau$ in the form of half-cauchy distribution
$$ \tau \sim C^{+}(0, \tau_0^2) $$
This type of hierachical prior is known as the horseshoe+ (horseshoe +) prior. Importantly, $\tau_0$ can be selected in the way to constrain the global shrinkage to an initial guess about the number of non-zero coefficients. We will denote this quantity as $p_0$ and the corresponding value of $\tau_0$ is then obtained as
$$ \tau_0 = \frac{p_0}{D - p_0} \frac{\sigma}{\sqrt{N}}$$
where $N$ denotes number of data points and $D$ number of coefficeints. The derivation and justification for the above relation is described in detail in Piironen et al. 2017.

```python
def m_eff(D, N, p0, s_sqr, sigma, tau, lam):
    kappa = 1/(1 + N * jnp.expand_dims(tau, -1)**2 * lam**2 * s_sqr / jnp.expand_dims(sigma, -1) ** 2)
    return jnp.sum( 1 - kappa, -1)

rng_key = random.PRNGKey(101)
fig, axes = plt.subplots(1, 2, figsize=(16, 5), gridspec_kw={'width_ratios': [2, 1], 'wspace': 0.2})
p0 = 5
rng_key, _rng_key = random.split(rng_key)
BLR = BayesRegression(rng_key, jnp.array(X), fixed_lam=False, p0=p0)
samples = BLR.fit(jnp.array(Y1), summary=False, num_chains=num_chains)

sns.boxenplot(data=np.array(samples['beta']), color='c', ax=axes[0], showfliers=False)
axes[0].scatter(np.arange(S), expected_thetas1, color='w', marker='*', s=100, label=r'true $\beta$')
axes[0].scatter(np.arange(D) + .25, hat_beta, color='r', marker='d', label=r'$\hat{\beta}$')

axes[0].hlines(0, 0, D, 'w', ls='--')
axes[0].set_ylabel(r'$\beta$')
axes[0].set_title(r'$p_0={}$'.format(p0))
axes[0].set_xticks([])
axes[0].legend(loc="upper right")

m_e = m_eff(D, N, p0, np.var(X, 0, ddof=1), samples['sigma'], samples['tau'], samples['lam'])
axes[1].hist(np.array(m_e), bins=50)
axes[1].set_xlabel(r'$m_{eff}$')

sns.despine()
# fig.savefig('full_model.png', transparent=True, bbox_inches='tight', dpi=300)
```

```python
# posterior distribution for other model parameters
fig, axes = plt.subplots(1, 3, figsize=(16, 5))
axes[0].hist(np.array(samples['sigma']), bins=50);
axes[0].set_xlabel(r'$\sigma$')
axes[1].hist(np.array(samples['alpha']), bins=50);
axes[1].set_xlabel(r'$\alpha$')
axes[2].hist(np.array(samples['tau']), bins=50);
axes[2].set_xlabel(r'$\tau$');
# fig.savefig('other_parameters.png', transparent=True, bbox_inches='tight', dpi=300)
```

```python
N = 400
X = np.random.randn(N, D)
X, Y1, expected_thetas1, expected_pairwise = get_data(X, S=S, pairwise=False, logistic=False)

_, hat_beta, _ = blr(X, Y1)

rng_key = random.PRNGKey(101)
fig, axes = plt.subplots(1, 2, figsize=(16, 5), gridspec_kw={'width_ratios': [2, 1], 'wspace': 0.2})
p0 = 5
rng_key, _rng_key = random.split(rng_key)
BLR = BayesRegression(rng_key, jnp.array(X), fixed_lam=False, p0=p0)
samples = BLR.fit(jnp.array(Y1), summary=False, num_chains=num_chains)

sns.boxenplot(data=np.array(samples['beta']), color='c', ax=axes[0], showfliers=False)
axes[0].scatter(np.arange(S), expected_thetas1, color='w', marker='*', s=100, label=r'true $\beta$')
axes[0].scatter(np.arange(D) + .25, hat_beta, color='r', marker='d', label=r'$\hat{\beta}$')

axes[0].hlines(0, 0, D, 'w', ls='--')
axes[0].set_ylabel(r'$\beta$')
axes[0].set_title(r'$p_0={}$'.format(p0))
axes[0].set_xticks([])
axes[0].legend(loc="upper right")

m_e = m_eff(D, N, p0, np.var(X, 0, ddof=1), samples['sigma'], samples['tau'], samples['lam'])
axes[1].hist(np.array(m_e), bins=50)
axes[1].set_xlabel(r'$m_{eff}$')

sns.despine()
```

```python
N = 800
X = np.random.randn(N, D)
X, Y1, expected_thetas1, expected_pairwise = get_data(X, S=S, pairwise=False, logistic=False)
_, hat_beta, _ = blr(X, Y1)

rng_key = random.PRNGKey(101)
fig, axes = plt.subplots(1, 2, figsize=(16, 5), gridspec_kw={'width_ratios': [2, 1], 'wspace': 0.2})
p0 = 5
rng_key, _rng_key = random.split(rng_key)
BLR = BayesRegression(rng_key, jnp.array(X), fixed_lam=False, p0=p0)
samples = BLR.fit(jnp.array(Y1), summary=False, num_chains=num_chains)

sns.boxenplot(data=np.array(samples['beta']), color='c', ax=axes[0], showfliers=False)
axes[0].scatter(np.arange(S), expected_thetas1, color='w', marker='*', s=100, label=r'true $\beta$')
axes[0].scatter(np.arange(D) + .25, hat_beta, color='r', marker='d', label=r'$\hat{\beta}$')

axes[0].hlines(0, 0, D, 'w', ls='--')
axes[0].set_ylabel(r'$\beta$')
axes[0].set_title(r'$p_0={}$'.format(p0))
axes[0].set_xticks([])
axes[0].legend(loc="upper right")

m_e = m_eff(D, N, p0, np.var(X, 0, ddof=1), samples['sigma'], samples['tau'], samples['lam'])
axes[1].hist(np.array(m_e), bins=50)
axes[1].set_xlabel(r'$m_{eff}$')

sns.despine()
```

# Correlated predictors
In practice, the predictors (the elements of a desing matrix $\pmb{X}$) will often be correlated. Here we illustrate how the results change under these conditions.

```python
L = dist.LKJCholesky(D, 1.).sample(random.PRNGKey(1))

N = 200
X = np.random.randn(N, D)
X_corr = X @ L.T
X_corr = (X_corr - X_corr.mean(0))/X_corr.std(0)

fig, axes = plt.subplots(1, 3, figsize=(20, 5))

sns.heatmap(np.corrcoef(X.T), ax=axes[0], cmap='magma')
axes[0].set_title(r'$X_k \sim \mathcal{N}_D(0, I_D)$')
sns.heatmap(np.corrcoef(X_corr.T), ax=axes[1])
axes[1].set_title(r'$X_k^c \sim \mathcal{N}_D(0, \Sigma)$')
Q, R = np.linalg.qr(X_corr)
sns.heatmap(np.corrcoef(Q.T), ax=axes[2])
axes[2].set_title(r'$Q_k = R \cdot X_k^c$');
# fig.savefig('cm.png', transparent=True, bbox_inches='tight', dpi=300)
```

```python
rng_key = random.PRNGKey(101)

N = 200
X = np.random.randn(N, D)
X_corr = X @ L.T
X_corr = (X_corr - X_corr.mean(0))/X_corr.std(0)
X_corr, Y2, expected_thetas2, expected_pairwise = get_data(X_corr, S=S, pairwise=False, logistic=False)

fig, axes = plt.subplots(3, 1, figsize=(16, 10), sharex=True)
p0 = 5
rng_key, _rng_key = random.split(rng_key)
BLR = BayesRegression(rng_key, jnp.array(X_corr), fixed_lam=False, p0=5, with_qr=True)
samples = BLR.fit(jnp.array(Y2), summary=False, num_chains=num_chains)

rng_key, _rng_key = random.split(rng_key)
BLR = BayesRegression(rng_key, jnp.array(X_corr), fixed_lam=True, p0=None)
samples2 = BLR.fit(jnp.array(Y2), summary=False, num_chains=num_chains)

rng_key, _rng_key = random.split(rng_key)
BLR = BayesRegression(rng_key, jnp.array(X_corr), tau=1e3, fixed_lam=True, p0=None)
samples3 = BLR.fit(jnp.array(Y2), summary=False, num_chains=num_chains)

for i, beta in enumerate([samples3['beta'], samples2['beta'], samples['beta']]):
    sns.boxenplot(data=np.array(beta), color='c', ax=axes[i], showfliers=False)
    axes[i].scatter(np.arange(S), expected_thetas2, color='w', marker='*', s=100, label=r'non-zero $\beta$')
    axes[i].hlines(0, 0, D, 'w', ls='--', lw=3)
    axes[i].set_ylabel(r'$\beta$')

axes[-1].set_xticks([])
axes[-1].set_xlabel(r'$d$')
axes[0].legend(loc="upper right", fontsize=10)
fig.suptitle(r'$N={}, D={}$'.format(N, D))
axes[0].set_title('Maximum likelihood estimate')
axes[1].set_title('Bayesian regression')
axes[2].set_title('Bayesian sparse regression')
sns.despine()
# fig.savefig('comparison.png', transparent=True, bbox_inches='tight', dpi=300)
```

# Logistic regression
The same structure of sparse hierarchical priors can be applied to other regression problems, like logistic regression or poisson regression.

```python
S = 5
D = 50
N = 1000

X = np.random.randn(N, D)
X, Y1_log, expected_thetas1_log, expected_pairwise = get_data(X, S=S, pairwise=False, logistic=True)

L = dist.LKJCholesky(D, 1.).sample(random.PRNGKey(1))
X_corr = X @ L
X_corr = (X_corr - X_corr.mean(0))/X_corr.std(0)
X_corr, Y2_log, expected_thetas2_log, expected_pairwise = get_data(X_corr, S=S, pairwise=False, logistic=True)
```

```python
# independent predictors
fig, axes = plt.subplots(1, 2, figsize=(16, 5), gridspec_kw={'width_ratios': [2, 1], 'wspace': 0.2})
p0 = 5
rng_key, _rng_key = random.split(rng_key)
LQR_reg = BayesRegression(rng_key, jnp.array(X), fixed_lam=False, p0=5, with_qr=True, reg_type='logistic')
samples = LQR_reg.fit(jnp.array(Y1_log), summary=False, num_chains=num_chains)

sns.boxenplot(data=np.array(samples['beta']), color='c', ax=axes[0], showfliers=False)
axes[0].scatter(np.arange(S), expected_thetas1_log, color='w', marker='*', s=100, label=r'true $\beta$')
axes[0].hlines(0, 0, D, 'w', ls='--')

axes[0].set_ylabel(r'$\beta$')
axes[0].set_title(r'$p_0={}$'.format(p0))
axes[0].set_xticks([])
axes[0].legend(loc="upper right", fontsize=10)

m_e = m_eff(D, N, p0, np.var(X, 0, ddof=1), samples['sigma'], samples['tau'], samples['lam'])
axes[1].hist(np.array(m_e), bins=50)
axes[1].set_xlabel(r'$m_{eff}$')

sns.despine()
```

```python
# correlated predictors
fig, axes = plt.subplots(1, 2, figsize=(16, 5), gridspec_kw={'width_ratios': [2, 1], 'wspace': 0.2})
p0 = 5
rng_key, _rng_key = random.split(rng_key)
LQR_reg = BayesRegression(rng_key, jnp.array(X_corr), fixed_lam=False, p0=5, with_qr=True, reg_type='logistic')
samples = LQR_reg.fit(jnp.array(Y2_log), summary=False, num_chains=num_chains)

sns.boxenplot(data=np.array(samples['beta']), color='c', ax=axes[0], showfliers=False)
axes[0].scatter(np.arange(S), expected_thetas2_log, color='w', marker='*', s=100, label=r'true $\beta$')
axes[0].hlines(0, 0, D, 'w', ls='--')

axes[0].set_ylabel(r'$\beta$')
axes[0].set_title(r'$p_0={}$'.format(p0))
axes[0].set_xticks([])
axes[0].legend(loc="upper right", fontsize=10)

m_e = m_eff(D, N, p0, np.var(X, 0, ddof=1), samples['sigma'], samples['tau'], samples['lam'])
axes[1].hist(np.array(m_e), bins=50)
axes[1].set_xlabel(r'$m_{eff}$')

sns.despine()
```

# The kernel interaction trick:
### Fast bayesian discovery of pairwise interactions in high dimensions.

Often encountered (linear) regression problem is contribution of pairwise interactions of predictors $x_i x_j$ to the measured responses $y$. The structural hierachical prior for the generative model with the interaction terms and corresponding coefficients $\beta_{ij}$ can be expressed as follows
\begin{equation}
    \begin{split}
        \lambda_i & \sim C^+(0, 1) \\
        \tau_1 & \sim C^+(0, \tau_0^2) \\
        \tau_2 &\sim C^+(0, \tau_1^2) \\
        \beta_i & \sim \mathcal{N}\left(0, \tau_1 \lambda_i \right) \\
        \beta_{ij} &\sim \mathcal{N}\left(0, \tau_2 \lambda_i \lambda_j \right) \\
        \alpha &\sim \mathcal{N}(0, c^2)
    \end{split}
\end{equation}

Importantly, direct esimation of pairwise interactions is computationally chalanging for large number of predictors $D$. One can achieve substantial runtime and memory improvments by marginalising out the regresssion coefficients from the generative model and expressing the regression problem as a gaussian process model. This approach is called Sparse Kernel Interaction Model (SKIM) introduced by Agrawal et al. 2019. The modifiied generative model is expressed as

\begin{equation}
    \begin{split}
        \lambda_i & \sim C^+(0, 1) \\
        \tau_1 & \sim C^+(0, \tau_0^2) \\
        \tau_2 &\sim C^+(0, \tau_1^2) \\
         y &\sim \mathcal{N}\left( 0, \pmb{K} \right)
    \end{split}
\end{equation}
where the kernel matrix is obtained as a sum of polynomial kernels
$$ k(x, \tilde{x}) = \frac{\tau_2^2}{2} \left(1 + x\cdot\tilde{x}\right)^2 - \frac{\tau_2^2}{2} \left( 1 + x^2 \cdot \tilde{x}^2 \right) + (\tau_1^2 - \tau_2^2) x \cdot \tilde{x} + c^2$$
$$ k(x, \tilde{x}) = \frac{\tau_2^2}{2} \left( (x \cdot \tilde{x})^2 - x^2 \cdot \tilde{x}^2 \right) + \tau_1^2 x \cdot \tilde{x} + c^2$$
where
$$ x_i, \tilde{x}_j = \lambda_i X_i, \lambda_j X_j $$

In what follows we will test the SKIM on data generated with independent and correlated predictors. The code we use here is a copy of examples provided by the Numpyro's documentation [page](https://num.pyro.ai/en/stable/examples/sparse_regression.html).

```python
S = 5
D = 50
N = 200

X = np.random.randn(N, D)
X, Y1_pair, expected_thetas1_pair, expected_pairwise1 = get_data(X, S=S, pairwise=True, logistic=False)

L = dist.LKJCholesky(D, 1.).sample(random.PRNGKey(1))
X_corr = X @ L
X_corr = (X_corr - X_corr.mean(0))/X_corr.std(0)
X_corr, Y2_pair, expected_thetas2_pair, expected_pairwise2 = get_data(X_corr, S=S, pairwise=True, logistic=False)
```

```python
def dot(X, Z):
    return jnp.dot(X, Z[..., None])[..., 0]

# The kernel that corresponds to a quadratic regressor.
def kernel(X, Z, tau1, tau2, c, jitter=1e-4):
    tau1sq = jnp.square(tau1)
    tau2sq = jnp.square(tau2)
    K1 = dot(X, Z)
    K2 = dot(jnp.square(X), jnp.square(Z))
    k1 = 0.5 * tau2sq * jnp.square(1.0 + K1)
    k2 = - 0.5 * tau2sq * K2
    k3 = (tau1sq - tau2sq) * K1
    k4 = jnp.square(c) - 0.5 * tau2sq
    if X.shape == Z.shape:
        k4 += jitter * jnp.eye(X.shape[0])
    return k1 + k2 + k3 + k4

def skim(X, Y=None, p0=5, c=10.):
    N, D = X.shape

    sigma_sqr = numpyro.sample('sigma^2', dist.InverseGamma(2., 2.))
    sigma = numpyro.deterministic('sigma', jnp.sqrt(sigma_sqr))

    tau0 = sigma * (p0/jnp.sqrt(N)) / (D - p0)
    tau1 = numpyro.sample('tau1', dist.HalfCauchy(tau0))
    tau2 = numpyro.sample('tau2', dist.HalfCauchy(tau1))

    lam = numpyro.sample('lam', dist.HalfCauchy(1.).expand([D]).to_event(1))

    # compute kernel
    lX = lam * X
    k = kernel(lX, lX, tau1, tau2, c) + sigma_sqr * jnp.eye(N)
    assert k.shape == (N, N)

    # sample Y according to the standard gaussian process formula
    numpyro.sample(
        "Y",
        dist.MultivariateNormal(loc=jnp.zeros(X.shape[0]), scale_tril=jnp.linalg.cholesky(k)),
        obs=Y,
    )

# Compute the mean and variance of coefficient theta_i (where i = dimension) for a
# MCMC sample of the kernel hyperparameters (eta1, xisq, ...).
# Compare to theorem 5.1 in reference [1].
def compute_singleton_mean_variance(X, Y, dimension, lam, tau1, tau2, sigma, c):
    P, N = X.shape[1], X.shape[0]

    probe = jnp.zeros((2, P))
    probe = jax.ops.index_update(
        probe, jax.ops.index[:, dimension], jnp.array([1.0, -1.0])
    )

    kX = lam * X
    kprobe = lam * probe

    k_xx = kernel(kX, kX, tau1, tau2, c) + sigma ** 2 * jnp.eye(N)
    k_xx_inv = jnp.linalg.inv(k_xx)
    k_probeX = kernel(kprobe, kX, tau1, tau2, c)
    k_prbprb = kernel(kprobe, kprobe, tau1, tau2, c)

    vec = jnp.array([0.50, -0.50])
    mu = jnp.matmul(k_probeX, jnp.matmul(k_xx_inv, Y))
    mu = jnp.dot(mu, vec)

    var = k_prbprb - jnp.matmul(k_probeX, jnp.matmul(k_xx_inv, jnp.transpose(k_probeX)))
    var = jnp.matmul(var, vec)
    var = jnp.dot(var, vec)

    return mu, var

# Compute the mean and variance of coefficient theta_ij for a MCMC sample of the
# kernel hyperparameters (eta1, xisq, ...). Compare to theorem 5.1 in reference [1].
def compute_pairwise_mean_variance(X, Y, dim1, dim2, lam, tau1, tau2, sigma, c):
    P, N = X.shape[1], X.shape[0]

    probe = jnp.zeros((4, P))
    probe = jax.ops.index_update(
        probe, jax.ops.index[:, dim1], jnp.array([1.0, 1.0, -1.0, -1.0])
    )
    probe = jax.ops.index_update(
        probe, jax.ops.index[:, dim2], jnp.array([1.0, -1.0, 1.0, -1.0])
    )

    kX = lam * X
    kprobe = lam * probe

    k_xx = kernel(kX, kX, tau1, tau2, c) + sigma ** 2 * jnp.eye(N)
    k_xx_inv = jnp.linalg.inv(k_xx)
    k_probeX = kernel(kprobe, kX, tau1, tau2, c)
    k_prbprb = kernel(kprobe, kprobe, tau1, tau2, c)

    vec = jnp.array([0.25, -0.25, -0.25, 0.25])
    mu = jnp.matmul(k_probeX, jnp.matmul(k_xx_inv, Y))
    mu = jnp.dot(mu, vec)

    var = k_prbprb - jnp.matmul(k_probeX, jnp.matmul(k_xx_inv, jnp.transpose(k_probeX)))
    var = jnp.matmul(var, vec)
    var = jnp.dot(var, vec)

    return mu, var

# Sample coefficients theta from the posterior for a given MCMC sample.
# The first P returned values are {theta_1, theta_2, ...., theta_P}, while
# the remaining values are {theta_ij} for i,j in the list `active_dims`,
# sorted so that i < j.
from jax.scipy.linalg import cho_factor, cho_solve, solve_triangular
def sample_theta_space(X, Y, active_dims, lam, tau1, tau2, sigma, c):
    P, N, M = X.shape[1], X.shape[0], len(active_dims)
    # the total number of coefficients we return
    num_coefficients = P + M * (M - 1) // 2

    probe = jnp.zeros((2 * P + 2 * M * (M - 1), P))
    vec = jnp.zeros((num_coefficients, 2 * P + 2 * M * (M - 1)))
    start1 = 0
    start2 = 0

    for dim in range(P):
        probe = jax.ops.index_update(
            probe, jax.ops.index[start1 : start1 + 2, dim], jnp.array([1.0, -1.0])
        )
        vec = jax.ops.index_update(
            vec, jax.ops.index[start2, start1 : start1 + 2], jnp.array([0.5, -0.5])
        )
        start1 += 2
        start2 += 1

    for dim1 in active_dims:
        for dim2 in active_dims:
            if dim1 >= dim2:
                continue
            probe = jax.ops.index_update(
                probe,
                jax.ops.index[start1 : start1 + 4, dim1],
                jnp.array([1.0, 1.0, -1.0, -1.0]),
            )
            probe = jax.ops.index_update(
                probe,
                jax.ops.index[start1 : start1 + 4, dim2],
                jnp.array([1.0, -1.0, 1.0, -1.0]),
            )
            vec = jax.ops.index_update(
                vec,
                jax.ops.index[start2, start1 : start1 + 4],
                jnp.array([0.25, -0.25, -0.25, 0.25]),
            )
            start1 += 4
            start2 += 1

    kX = lam * X
    kprobe = lam * probe

    k_xx = kernel(kX, kX, tau1, tau2, c) + sigma ** 2 * jnp.eye(N)
    L = cho_factor(k_xx, lower=True)[0]
    k_probeX = kernel(kprobe, kX, tau1, tau2, c)
    k_prbprb = kernel(kprobe, kprobe, tau1, tau2, c)

    mu = jnp.matmul(k_probeX, cho_solve((L, True), Y))
    mu = jnp.sum(mu * vec, axis=-1)

    Linv_k_probeX = solve_triangular(L, jnp.transpose(k_probeX), lower=True)
    covar = k_prbprb - jnp.matmul(jnp.transpose(Linv_k_probeX), Linv_k_probeX)
    covar = jnp.matmul(vec, jnp.matmul(covar, jnp.transpose(vec)))

    # sample from N(mu, covar)
    L = jnp.linalg.cholesky(covar)
    sample = mu + jnp.matmul(L, np.random.randn(num_coefficients))

    return sample

# Get the mean and variance of a gaussian mixture
def gaussian_mixture_stats(mus, variances):
    mean_mu = jnp.mean(mus)
    mean_var = jnp.mean(variances) + jnp.mean(jnp.square(mus)) - jnp.square(mean_mu)
    return mean_mu, mean_var

# Helper function for analyzing the posterior statistics for coefficient theta_i
def analyze_dimension(samples, X, Y, dimension, c=10.):
    vmap_args = (
        samples["lam"],
        samples["tau1"],
        samples["tau2"],
        samples["sigma"],
    )
    mus, variances = vmap(
        lambda lam, tau1, tau2, sigma: compute_singleton_mean_variance(
            X, Y, dimension, lam, tau1, tau2, sigma, c
        )
    )(*vmap_args)
    mean, variance = gaussian_mixture_stats(mus, variances)
    std = jnp.sqrt(variance)
    return mean, std

# Helper function for analyzing the posterior statistics for coefficient theta_ij
def analyze_pair_of_dimensions(samples, X, Y, dim1, dim2, c=10.):
    vmap_args = (
        samples["lam"],
        samples["tau1"],
        samples["tau2"],
        samples["sigma"],
    )
    mus, variances = vmap(
        lambda lam, tau1, tau2, sigma: compute_pairwise_mean_variance(
            X, Y, dim1, dim2, lam, tau1, tau2, sigma, c
        )
    )(*vmap_args)
    mean, variance = gaussian_mixture_stats(mus, variances)
    std = jnp.sqrt(variance)
    return mean, std
```

```python
%%time
c=1.
nuts = NUTS(skim)
mcmc = MCMC(nuts, num_warmup=1000, num_samples=1000)
mcmc.run(random.PRNGKey(101010), X, Y=Y1_pair, p0=S, c=c)

samples = mcmc.get_samples()
```

    sample: 100%|█| 2000/2000 [02:01<00:00, 16.52it/s, 15 steps of size 2.04e-01. ac

    CPU times: user 16min 42s, sys: 45min 20s, total: 1h 2min 2s
    Wall time: 2min 3s

```python
# posterior distribution for other model parameters
fig, axes = plt.subplots(1, 3, figsize=(16, 5))
axes[0].hist(np.array(samples['sigma']), bins=50);
axes[0].set_xlabel(r'$\sigma$')
axes[1].hist(np.array(samples['tau1']), bins=50);
axes[1].set_xlabel(r'$\tau_1$')
axes[2].hist(np.array(samples['tau2']), bins=50);
axes[2].set_xlabel(r'$\tau_2$');
```

```python
 # compute the mean and square root variance of each coefficient theta_i
means, stds = vmap(lambda dim: analyze_dimension(samples, X, Y1_pair, dim, c=c))(
    jnp.arange(D)
)

print(
    "Coefficients theta_1 to theta_%d used to generate the data:"
    % S,
    expected_thetas1_pair
)
print(
    "The single quadratic coefficient theta_{1,2} used to generate the data:",
    expected_pairwise1,
)
```

    Coefficients theta_1 to theta_5 used to generate the data: [0.4465585  0.47040176 0.23218709 0.2950355  0.5207558 ]
    The single quadratic coefficient theta_{1,2} used to generate the data: 0.21967684

```python
act_dimensions = []

for dim, (mean, std) in enumerate(zip(means, stds)):
    # we mark the dimension as inactive if the interval [mean - 3 * std, mean + 3 * std] contains zero
    lower, upper = mean - 3.0 * std, mean + 3.0 * std
    inactive = "inactive" if lower < 0.0 and upper > 0.0 else "active"
    if inactive == "active":
        act_dimensions.append(dim)
    print(
        "[dimension %02d/%02d]  %s:\t%.2e +- %.2e"
        % (dim + 1, D, inactive, mean, std)
    )

print(
    "Identified a total of %d active dimensions; expected %d."
    % (len(act_dimensions), S)
)
```

    [dimension 01/50]  active:	4.62e-01 +- nan
    [dimension 02/50]  active:	4.62e-01 +- 1.10e-01
    [dimension 03/50]  active:	1.96e-01 +- 2.20e-02
    [dimension 04/50]  active:	2.99e-01 +- 2.21e-02
    [dimension 05/50]  active:	5.09e-01 +- 2.01e-02
    [dimension 06/50]  inactive:	7.29e-04 +- 9.96e-03
    [dimension 07/50]  inactive:	1.11e-03 +- 9.80e-03
    [dimension 08/50]  inactive:	7.79e-03 +- 1.55e-02
    [dimension 09/50]  inactive:	-3.62e-03 +- 1.15e-02
    [dimension 10/50]  inactive:	-4.93e-04 +- 1.02e-02
    [dimension 11/50]  inactive:	-3.42e-03 +- 1.16e-02
    [dimension 12/50]  inactive:	-6.69e-04 +- 1.02e-02
    [dimension 13/50]  inactive:	1.72e-03 +- 1.04e-02
    [dimension 14/50]  inactive:	1.35e-03 +- 1.02e-02
    [dimension 15/50]  inactive:	-7.63e-04 +- 9.73e-03
    [dimension 16/50]  inactive:	-2.22e-03 +- 1.08e-02
    [dimension 17/50]  inactive:	-2.89e-03 +- 1.11e-02
    [dimension 18/50]  inactive:	3.22e-03 +- 1.18e-02
    [dimension 19/50]  inactive:	-7.57e-04 +- 1.03e-02
    [dimension 20/50]  inactive:	-2.74e-03 +- 1.14e-02
    [dimension 21/50]  inactive:	1.49e-03 +- 1.02e-02
    [dimension 22/50]  inactive:	-3.58e-03 +- 1.13e-02
    [dimension 23/50]  inactive:	2.16e-03 +- 1.07e-02
    [dimension 24/50]  inactive:	6.97e-04 +- 1.03e-02
    [dimension 25/50]  inactive:	7.32e-04 +- 1.01e-02
    [dimension 26/50]  inactive:	-3.93e-04 +- 9.83e-03
    [dimension 27/50]  inactive:	4.60e-04 +- 1.04e-02
    [dimension 28/50]  inactive:	4.65e-03 +- 1.26e-02
    [dimension 29/50]  inactive:	5.72e-04 +- 1.01e-02
    [dimension 30/50]  inactive:	-3.29e-03 +- 1.17e-02
    [dimension 31/50]  inactive:	7.35e-04 +- 1.02e-02
    [dimension 32/50]  inactive:	-6.01e-04 +- 1.05e-02
    [dimension 33/50]  inactive:	8.03e-04 +- 9.74e-03
    [dimension 34/50]  inactive:	-2.96e-04 +- 1.01e-02
    [dimension 35/50]  inactive:	-1.38e-03 +- 9.94e-03
    [dimension 36/50]  inactive:	-1.43e-03 +- 1.04e-02
    [dimension 37/50]  inactive:	-5.04e-03 +- 1.28e-02
    [dimension 38/50]  inactive:	-8.05e-04 +- 1.05e-02
    [dimension 39/50]  inactive:	-2.55e-04 +- 9.87e-03
    [dimension 40/50]  inactive:	-4.39e-04 +- 1.01e-02
    [dimension 41/50]  inactive:	7.44e-04 +- 1.01e-02
    [dimension 42/50]  inactive:	2.13e-04 +- 1.02e-02
    [dimension 43/50]  inactive:	1.05e-03 +- 9.92e-03
    [dimension 44/50]  inactive:	5.90e-04 +- 9.80e-03
    [dimension 45/50]  inactive:	-5.03e-04 +- 1.01e-02
    [dimension 46/50]  inactive:	-3.43e-04 +- 1.01e-02
    [dimension 47/50]  inactive:	5.62e-04 +- 9.76e-03
    [dimension 48/50]  inactive:	-1.51e-03 +- 1.03e-02
    [dimension 49/50]  inactive:	-1.20e-05 +- 9.57e-03
    [dimension 50/50]  inactive:	-1.04e-03 +- 1.03e-02
    Identified a total of 5 active dimensions; expected 5.

```python
# Compute the mean and square root variance of coefficients theta_ij for i,j active dimensions.
# Note that the resulting numbers are only meaningful for i != j.
if len(act_dimensions) > 0:
    dim_pairs = jnp.array(
        list(itertools.product(act_dimensions, act_dimensions))
    )
    means, stds = vmap(
        lambda dim_pair: analyze_pair_of_dimensions(
            samples, X, Y1_pair, dim_pair[0], dim_pair[1], c=c
        )
    )(dim_pairs)
    for dim_pair, mean, std in zip(dim_pairs, means, stds):
        dim1, dim2 = dim_pair
        if dim1 >= dim2:
            continue
        lower, upper = mean - 3.0 * std, mean + 3.0 * std
        if not (lower < 0.0 and upper > 0.0):
            format_str = "Identified pairwise interaction between dimensions %d and %d: %.2e +- %.2e"
            print(format_str % (dim1 + 1, dim2 + 1, mean, std))
    # Draw a single sample of coefficients theta from the posterior, where we return all singleton
    # coefficients theta_i and pairwise coefficients theta_ij for i, j active dimensions. We use the
    # final MCMC sample obtained from the HMC sampler.
    thetas = sample_theta_space(
        X,
        Y1_pair,
        act_dimensions,
        samples["lam"][-1],
        samples["tau1"][-1],
        samples["tau2"][-1],
        samples["sigma"][-1],
        c
    )
    print("Single posterior sample theta:\n", thetas)
```

    Identified pairwise interaction between dimensions 1 and 2: 2.11e-01 +- 1.83e-02
    Single posterior sample theta:
     [ 0.4306354   0.485133    0.19135524  0.2955897   0.52012146 -0.00326966
      0.0108312  -0.01407628 -0.00548292  0.00316599  0.00326666  0.00479578
     -0.0071921  -0.00854616  0.00240428  0.01050893 -0.00742017 -0.01194203
     -0.00177869  0.00064334  0.00713305  0.00833935  0.00336977 -0.00722155
      0.00241847 -0.01148906 -0.00296297 -0.00630507 -0.0099332  -0.00883699
     -0.00103072  0.00351536  0.00765699  0.00569884 -0.01701669 -0.00840545
     -0.01197023  0.00811246  0.01220147 -0.0068486  -0.00806344 -0.0027991
     -0.01159613  0.00250444  0.00508768 -0.00217367  0.01692843 -0.00787118
      0.00555956 -0.02159943  0.2233778   0.02470442 -0.01145326 -0.01218557
     -0.02235619 -0.0190891   0.02199975  0.00086919  0.01243104  0.00085664]

Finally we can do the same analysis in the case of correlated predictors.

```python
%%time
nuts = NUTS(skim)
mcmc = MCMC(nuts, num_warmup=1000, num_samples=1000)
mcmc.run(random.PRNGKey(101010), X_corr, Y=Y2_pair, p0=S, c=c)

samples = mcmc.get_samples()
```

    sample: 100%|█| 2000/2000 [02:19<00:00, 14.37it/s, 15 steps of size 1.83e-01. ac

    CPU times: user 19min 13s, sys: 51min 50s, total: 1h 11min 4s
    Wall time: 2min 19s

```python
# posterior distribution for other model parameters
fig, axes = plt.subplots(1, 3, figsize=(16, 5))
axes[0].hist(np.array(samples['sigma']), bins=50);
axes[0].set_xlabel(r'$\sigma$')
axes[1].hist(np.array(samples['tau1']), bins=50);
axes[1].set_xlabel(r'$\tau_1$')
axes[2].hist(np.array(samples['tau2']), bins=50);
axes[2].set_xlabel(r'$\tau_2$');
```

```python
 # compute the mean and square root variance of each coefficient theta_i
means, stds = vmap(lambda dim: analyze_dimension(samples, X_corr, Y2_pair, dim, c=c))(
    jnp.arange(D)
)

print(
    "Coefficients theta_1 to theta_%d used to generate the data:"
    % S,
    expected_thetas2_pair
)
print(
    "The single quadratic coefficient theta_{1,2} used to generate the data:",
    expected_pairwise2,
)
```

    Coefficients theta_1 to theta_5 used to generate the data: [0.4325807  0.45567766 0.22491936 0.28580058 0.50445557]
    The single quadratic coefficient theta_{1,2} used to generate the data: 0.2128007

```python
act_dimensions = []

for dim, (mean, std) in enumerate(zip(means, stds)):
    # we mark the dimension as inactive if the interval [mean - 3 * std, mean + 3 * std] contains zero
    lower, upper = mean - 3.0 * std, mean + 3.0 * std
    inactive = "inactive" if lower < 0.0 and upper > 0.0 else "active"
    if inactive == "active":
        act_dimensions.append(dim)
    print(
        "[dimension %02d/%02d]  %s:\t%.2e +- %.2e"
        % (dim + 1, D, inactive, mean, std)
    )

print(
    "Identified a total of %d active dimensions; expected %d."
    % (len(act_dimensions), S)
)
```

    [dimension 01/50]  active:	4.50e-01 +- 2.31e-02
    [dimension 02/50]  active:	4.58e-01 +- 8.85e-03
    [dimension 03/50]  active:	2.13e-01 +- 2.17e-02
    [dimension 04/50]  active:	2.82e-01 +- 2.22e-02
    [dimension 05/50]  active:	4.77e-01 +- 2.14e-02
    [dimension 06/50]  inactive:	2.62e-03 +- 1.11e-02
    [dimension 07/50]  inactive:	-1.89e-03 +- 1.05e-02
    [dimension 08/50]  inactive:	4.67e-03 +- 1.28e-02
    [dimension 09/50]  inactive:	-3.20e-03 +- 1.15e-02
    [dimension 10/50]  inactive:	-8.33e-05 +- 9.96e-03
    [dimension 11/50]  inactive:	-6.10e-04 +- 1.03e-02
    [dimension 12/50]  inactive:	1.52e-03 +- 1.04e-02
    [dimension 13/50]  inactive:	3.22e-03 +- 1.16e-02
    [dimension 14/50]  inactive:	-1.79e-04 +- 9.86e-03
    [dimension 15/50]  inactive:	-3.04e-04 +- 9.88e-03
    [dimension 16/50]  inactive:	-1.45e-03 +- 1.02e-02
    [dimension 17/50]  inactive:	-8.26e-04 +- 1.02e-02
    [dimension 18/50]  inactive:	1.28e-03 +- 1.02e-02
    [dimension 19/50]  inactive:	-1.23e-03 +- 1.04e-02
    [dimension 20/50]  inactive:	-3.29e-04 +- 9.87e-03
    [dimension 21/50]  inactive:	1.21e-03 +- 1.04e-02
    [dimension 22/50]  inactive:	-3.63e-03 +- 1.21e-02
    [dimension 23/50]  inactive:	2.95e-03 +- 1.15e-02
    [dimension 24/50]  inactive:	-4.98e-04 +- 1.01e-02
    [dimension 25/50]  inactive:	5.43e-04 +- 1.03e-02
    [dimension 26/50]  inactive:	2.13e-03 +- 1.12e-02
    [dimension 27/50]  inactive:	9.41e-04 +- 1.02e-02
    [dimension 28/50]  inactive:	5.83e-04 +- 1.06e-02
    [dimension 29/50]  inactive:	1.02e-03 +- 1.04e-02
    [dimension 30/50]  inactive:	-1.42e-03 +- 1.05e-02
    [dimension 31/50]  inactive:	1.33e-03 +- 1.06e-02
    [dimension 32/50]  inactive:	5.93e-04 +- 1.05e-02
    [dimension 33/50]  inactive:	7.75e-04 +- 1.01e-02
    [dimension 34/50]  inactive:	1.59e-03 +- 1.05e-02
    [dimension 35/50]  inactive:	-1.80e-03 +- 1.07e-02
    [dimension 36/50]  inactive:	-1.61e-03 +- 1.08e-02
    [dimension 37/50]  inactive:	-1.01e-03 +- 1.05e-02
    [dimension 38/50]  inactive:	-4.04e-04 +- 1.04e-02
    [dimension 39/50]  inactive:	1.85e-04 +- 1.01e-02
    [dimension 40/50]  inactive:	-1.78e-03 +- 1.07e-02
    [dimension 41/50]  inactive:	1.00e-03 +- 1.05e-02
    [dimension 42/50]  inactive:	-1.58e-03 +- 1.15e-02
    [dimension 43/50]  inactive:	2.85e-03 +- 1.18e-02
    [dimension 44/50]  inactive:	4.51e-04 +- 1.00e-02
    [dimension 45/50]  inactive:	-1.72e-04 +- 1.04e-02
    [dimension 46/50]  inactive:	-1.34e-03 +- 1.11e-02
    [dimension 47/50]  inactive:	-1.67e-04 +- 1.02e-02
    [dimension 48/50]  inactive:	-5.83e-04 +- 1.03e-02
    [dimension 49/50]  inactive:	9.99e-04 +- 1.11e-02
    [dimension 50/50]  inactive:	-1.99e-03 +- 1.13e-02
    Identified a total of 5 active dimensions; expected 5.

```python
# Compute the mean and square root variance of coefficients theta_ij for i,j active dimensions.
# Note that the resulting numbers are only meaningful for i != j.
if len(act_dimensions) > 0:
    dim_pairs = jnp.array(
        list(itertools.product(act_dimensions, act_dimensions))
    )
    means, stds = vmap(
        lambda dim_pair: analyze_pair_of_dimensions(
            samples, X_corr, Y2_pair, dim_pair[0], dim_pair[1], c=c
        )
    )(dim_pairs)
    for dim_pair, mean, std in zip(dim_pairs, means, stds):
        dim1, dim2 = dim_pair
        if dim1 >= dim2:
            continue
        lower, upper = mean - 3.0 * std, mean + 3.0 * std
        if not (lower < 0.0 and upper > 0.0):
            format_str = "Identified pairwise interaction between dimensions %d and %d: %.2e +- %.2e"
            print(format_str % (dim1 + 1, dim2 + 1, mean, std))
    # Draw a single sample of coefficients theta from the posterior, where we return all singleton
    # coefficients theta_i and pairwise coefficients theta_ij for i, j active dimensions. We use the
    # final MCMC sample obtained from the HMC sampler.
    thetas = sample_theta_space(
        X_corr,
        Y2_pair,
        act_dimensions,
        samples["lam"][-1],
        samples["tau1"][-1],
        samples["tau2"][-1],
        samples["sigma"][-1],
        c
    )
    print("Single posterior sample theta:\n", thetas)
```

    Identified pairwise interaction between dimensions 1 and 2: 2.06e-01 +- 2.22e-02
    Single posterior sample theta:
     [ 4.8349717e-01  4.8069978e-01  2.1730670e-01  2.9150867e-01
      4.6403784e-01 -6.6083786e-04  5.0153071e-03 -8.2247304e-03
      2.3886025e-02 -2.1383921e-03 -4.8805494e-05  1.9332971e-02
     -2.3297151e-03 -1.1333087e-02  8.2426826e-03  1.3720994e-03
     -1.5972648e-03  1.5419638e-02 -1.7237801e-02 -1.0476515e-03
     -1.9991331e-02 -1.3375759e-02 -1.3664368e-02 -2.8204645e-03
      2.0179967e-03 -4.6986872e-03  8.7081976e-03 -1.8950242e-03
      6.0884496e-03 -1.0451022e-02 -1.2721230e-03  1.9996975e-02
      4.9252054e-03 -1.7527594e-03  5.2826968e-04 -1.2368436e-02
      9.9975662e-03  9.2484232e-04  1.0955558e-03  5.6325411e-04
     -1.0480522e-04 -5.3823125e-03 -7.6159891e-03  3.3704909e-03
      4.2474084e-03  2.0400750e-02  4.0809754e-03 -1.1416207e-02
      3.1625610e-03  2.2048359e-03  2.3782669e-01  2.9079003e-02
     -2.6484149e-02 -1.0421222e-02  2.0873636e-02 -2.3493353e-02
      2.1627620e-02 -2.0127231e-02 -5.6583737e-04  6.4932869e-04]

Importantly, SKIM performs equally well on data sets generated both with independent and correlated predictors.