```python
import sys
sys.path.append('..')
import os
import torch
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.preprocessing import LabelEncoder, MinMaxScaler, KBinsDiscretizer
from sklearn.impute import SimpleImputer
from sklearn.model_selection import cross_val_score
from sklearn.tree import DecisionTreeClassifier
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import StratifiedKFold
import time
from sklearn.metrics import accuracy_score
from sympy import simplify_logic
from sklearn.model_selection import train_test_split

import lens
from lens.utils.base import validate_network, set_seed, tree_to_formula
from lens.utils.layer import prune_logic_layers
from lens import logic

set_seed(0)

#%%

data = pd.read_csv('data/mimic-ii/full_cohort_data.csv')
# data.drop('hgb_first')
fs = [
    'aline_flg',
    'gender_num',
    # 'hosp_exp_flg',
    # 'icu_exp_flg',
    # 'day_28_flg',
    # 'censor_flg',
    'sepsis_flg', 'chf_flg', 'afib_flg',
    'renal_flg', 'liver_flg', 'copd_flg', 'cad_flg', 'stroke_flg',
    'mal_flg', 'resp_flg',
]
features = fs
data1 = data[fs].values
imp_mean = SimpleImputer(missing_values=np.nan, strategy='mean')
data1 = imp_mean.fit_transform(data1)

f2 = fs.copy()
f2.append('day_icu_intime')
f2.append('service_unit')
f2.append('day_28_flg')
f2.append('hospital_los_day')
f2.append('icu_exp_flg')
f2.append('hosp_exp_flg')
f2.append('censor_flg')
f2.append('mort_day_censored')
f2 = data.columns.difference(f2)
data2 = data[f2].values
imp_mean = SimpleImputer(missing_values=np.nan, strategy='mean')
data2 = imp_mean.fit_transform(data2)
scaler = MinMaxScaler((0, 1))
data2 = scaler.fit_transform(data2)
features = features + list(f2)
est = KBinsDiscretizer(n_bins=3, encode='onehot-dense', strategy='uniform')
data2d = est.fit_transform(data2)
f2d = []
for feature in f2:
    #f2d.append(feature + '_VLOW')
    f2d.append(feature + '_LOW')
    f2d.append(feature + '_NORMAL')
    f2d.append(feature + '_HIGH')
    #f2d.append(feature + '_VHIGH')
features = fs + f2d

datax = np.hstack((data1, data2d))
datay = data['day_28_flg'].values

x = torch.FloatTensor(datax)
y = torch.FloatTensor(datay)
print(x.shape)
print(y.shape)
```

    torch.Size([1776, 90])
    torch.Size([1776])

```python
results_dir = 'results_ll/mimic'
if not os.path.isdir(results_dir):
    os.makedirs(results_dir)

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
# concepts = [f'c{i:03}' for i in range(x.shape[1])]
concepts =  features
n_rep = 10
tot_epochs = 5001
prune_epochs = 2001
seed = 42

n_splits = 10
skf = StratifiedKFold(n_splits=n_splits, shuffle=True, random_state=seed)
```

```python
def train_nn(x_train, y_train, seed, device, l1=0.001, lr=0.001, verbose=False):
    set_seed(seed)
    x_train = x_train.to(device)
    y_train = y_train.to(device)
    layers = [
        lens.nn.XLogic(x_train.size(1), 50, first=True),
        torch.nn.LeakyReLU(),
        torch.nn.Linear(50, 30),
        torch.nn.LeakyReLU(),
        torch.nn.Linear(30, 1),
        torch.nn.LeakyReLU(),
        lens.nn.XLogic(1, 1, top=True),
    ]
    model = torch.nn.Sequential(*layers).to(device)

    optimizer = torch.optim.AdamW(model.parameters(), lr=lr)
    loss_form = torch.nn.BCELoss()
    model.train()
    need_pruning = True
    for epoch in range(tot_epochs):
        optimizer.zero_grad()
        y_pred = model(x_train).squeeze()
        loss = loss_form(y_pred, y_train)

        for module in model.children():
            if isinstance(module, lens.nn.XLogic):
                loss += l1 * torch.norm(module.weight, 1)
                loss += l1 * torch.norm(module.bias, 1)

        loss.backward()
        optimizer.step()

        if epoch > prune_epochs and need_pruning:
            lens.utils.layer.prune_logic_layers(model.to(device), fan_in=5, device=device)
            need_pruning = False

        # compute accuracy
        if epoch % 100 == 0 and verbose:
            y_pred_d = y_pred > 0.5
            accuracy = y_pred_d.eq(y_train).sum().item() / y_train.size(0)
            print(f'Epoch {epoch}: train accuracy: {accuracy:.4f}')

    return model
```

```python
def c_to_y(method, verbose=False):
    methods = []
    splits = []
    explanations = []
    model_accuracies = []
    explanation_accuracies = []
    explanation_fidelities = []
    explanation_complexities = []
    elapsed_times = []
    for split, (trainval_index, test_index) in enumerate(skf.split(x.cpu().detach().numpy(), y.cpu().detach().numpy())):
        print(f'Split [{split+1}/{n_splits}]')
        x_trainval, x_test = torch.FloatTensor(x[trainval_index]), torch.FloatTensor(x[test_index])
        y_trainval, y_test = torch.FloatTensor(y[trainval_index]), torch.FloatTensor(y[test_index])
        x_train, x_val, y_train, y_val = train_test_split(x_trainval, y_trainval, test_size=0.3, random_state=42)

        explanation, explanation_inv = '', ''
        explanation_accuracy, explanation_accuracy_inv = 0, 0

        if method == 'tree':
            classifier = DecisionTreeClassifier(random_state=seed)
            classifier.fit(x_train.detach().numpy(), y_train.detach().numpy())
            y_preds = classifier.predict(x_test.detach().numpy())
            model_accuracy = accuracy_score(y_test.detach().numpy(), y_preds)
            explanation_accuracy = model_accuracy

            target_class = 1
            start = time.time()
            explanation = tree_to_formula(classifier, concepts, target_class)
            elapsed_time = time.time() - start
            explanation_fidelity = 1.
            explanation_complexity = lens.logic.complexity(explanation)

            target_class_inv = 0
            start = time.time()
            explanation_inv = tree_to_formula(classifier, concepts, target_class_inv)
            elapsed_time = time.time() - start

        else:
            model = train_nn(x_trainval, y_trainval, seed, device, verbose=False)
            y_preds = model(x_test.to(device)).cpu().detach().numpy() > 0.5
            model_accuracy = accuracy_score(y_test.cpu().detach().numpy(), y_preds)

            # positive class
            start = time.time()
            class_explanation, class_explanations = lens.logic.explain_class(model.cpu(),
                                                                           x_trainval.cpu(),
                                                                           y_trainval.cpu(),
                                                                           binary=True, target_class=1,
                                                                           topk_explanations=3)
            elapsed_time = time.time() - start

            if class_explanation:
                explanation = logic.base.replace_names(class_explanation, concepts)
                explanation_accuracy, y_formula = logic.base.test_explanation(class_explanation,
                                                                              target_class=1,
                                                                              x=x_test, y=y_test,
                                                                              metric=accuracy_score)
                explanation_fidelity = lens.logic.fidelity(y_formula, y_preds)
                explanation_complexity = lens.logic.complexity(class_explanation)

        if verbose:
            print(f'\t Model\'s accuracy: {model_accuracy:.4f}')
            print(f'\t Class 1 - Global explanation: "{explanation}" - Accuracy: {explanation_accuracy:.4f}')
            print(f'\t Fidelity: "{explanation_fidelity:.4f}" - Complexity: "{explanation_complexity}"')
            print(f'\t Elapsed time {elapsed_time}')

        methods.append(method)
        splits.append(seed)
        explanations.append(explanation)
        model_accuracies.append(model_accuracy)
        explanation_accuracies.append(explanation_accuracy)
        explanation_fidelities.append(explanation_fidelity)
        explanation_complexities.append(explanation_complexity)
        elapsed_times.append(elapsed_time)

    explanation_consistency = lens.logic.formula_consistency(explanations)
    print(f'Consistency of explanations: {explanation_consistency:.4f}')

    results = pd.DataFrame({
        'method': methods,
        'split': splits,
        'explanation': explanations,
        'model_accuracy': model_accuracies,
        'explanation_accuracy': explanation_accuracies,
        'explanation_fidelity': explanation_fidelities,
        'explanation_complexity': explanation_complexities,
        'explanation_consistency': explanation_consistency,
        'elapsed_time': elapsed_times,
    })
    results.to_csv(os.path.join(results_dir, f'results_{method}.csv'))

    return results
```

# General pruning

```python
results_pruning = c_to_y(method='logic_layer', verbose=True)
results_pruning
```

    Split [1/10]
    	 Model's accuracy: 0.8596
    	 Class 1 - Global explanation: "stroke_flg & age_HIGH & ~sapsi_first_LOW" - Accuracy: 0.8596
    	 Fidelity: "1.0000" - Complexity: "3"
    	 Elapsed time 0.3629953861236572
    Split [2/10]
    	 Model's accuracy: 0.8820
    	 Class 1 - Global explanation: "(stroke_flg & abg_count_LOW & age_HIGH & ~age_LOW & ~sapsi_first_LOW) | (age_HIGH & ~stroke_flg & ~abg_count_LOW & ~age_LOW & ~sapsi_first_LOW)" - Accuracy: 0.8820
    	 Fidelity: "1.0000" - Complexity: "10"
    	 Elapsed time 0.3061814308166504
    Split [3/10]
    	 Model's accuracy: 0.8764
    	 Class 1 - Global explanation: "stroke_flg & age_HIGH & ~sapsi_first_LOW" - Accuracy: 0.8764
    	 Fidelity: "1.0000" - Complexity: "3"
    	 Elapsed time 0.26928281784057617
    Split [4/10]
    	 Model's accuracy: 0.8652
    	 Class 1 - Global explanation: "(stroke_flg & abg_count_LOW & age_HIGH & ~age_LOW & ~sapsi_first_LOW) | (age_HIGH & ~stroke_flg & ~abg_count_LOW & ~age_LOW & ~sapsi_first_LOW)" - Accuracy: 0.8652
    	 Fidelity: "1.0000" - Complexity: "10"
    	 Elapsed time 0.31247735023498535
    Split [5/10]
    	 Model's accuracy: 0.8596
    	 Class 1 - Global explanation: "stroke_flg & age_HIGH & ~sapsi_first_LOW" - Accuracy: 0.8596
    	 Fidelity: "1.0000" - Complexity: "3"
    	 Elapsed time 0.2738156318664551
    Split [6/10]
    	 Model's accuracy: 0.8708
    	 Class 1 - Global explanation: "stroke_flg & age_HIGH & ~sapsi_first_LOW" - Accuracy: 0.8708
    	 Fidelity: "1.0000" - Complexity: "3"
    	 Elapsed time 0.370011568069458
    Split [7/10]
    	 Model's accuracy: 0.8418
    	 Class 1 - Global explanation: "stroke_flg & age_HIGH & ~sapsi_first_LOW" - Accuracy: 0.8418
    	 Fidelity: "1.0000" - Complexity: "3"
    	 Elapsed time 0.2564055919647217
    Split [8/10]
    	 Model's accuracy: 0.8701
    	 Class 1 - Global explanation: "(stroke_flg & age_HIGH & ~age_LOW & ~sapsi_first_LOW) | (age_HIGH & ~age_LOW & ~bun_first_LOW & ~sapsi_first_LOW)" - Accuracy: 0.8701
    	 Fidelity: "1.0000" - Complexity: "8"
    	 Elapsed time 0.3371310234069824
    Split [9/10]
    	 Model's accuracy: 0.8870
    	 Class 1 - Global explanation: "stroke_flg & age_HIGH & ~sapsi_first_LOW" - Accuracy: 0.8870
    	 Fidelity: "1.0000" - Complexity: "3"
    	 Elapsed time 0.30785655975341797
    Split [10/10]
    	 Model's accuracy: 0.8701
    	 Class 1 - Global explanation: "stroke_flg & age_HIGH & ~sapsi_first_LOW" - Accuracy: 0.8701
    	 Fidelity: "1.0000" - Complexity: "3"
    	 Elapsed time 0.28518128395080566
    Consistency of explanations: 0.6000

      method
      split
      explanation
      model_accuracy
      explanation_accuracy
      explanation_fidelity
      explanation_complexity
      explanation_consistency
      elapsed_time

      0
      logic_layer
      42
      stroke_flg & age_HIGH & ~sapsi_first_LOW
      0.859551
      0.859551
      1.0
      3
      0.6
      0.362995

      1
      logic_layer
      42
      (stroke_flg & abg_count_LOW & age_HIGH & ~age_...
      0.882022
      0.882022
      1.0
      10
      0.6
      0.306181

      2
      logic_layer
      42
      stroke_flg & age_HIGH & ~sapsi_first_LOW
      0.876404
      0.876404
      1.0
      3
      0.6
      0.269283

      3
      logic_layer
      42
      (stroke_flg & abg_count_LOW & age_HIGH & ~age_...
      0.865169
      0.865169
      1.0
      10
      0.6
      0.312477

      4
      logic_layer
      42
      stroke_flg & age_HIGH & ~sapsi_first_LOW
      0.859551
      0.859551
      1.0
      3
      0.6
      0.273816

      5
      logic_layer
      42
      stroke_flg & age_HIGH & ~sapsi_first_LOW
      0.870787
      0.870787
      1.0
      3
      0.6
      0.370012

      6
      logic_layer
      42
      stroke_flg & age_HIGH & ~sapsi_first_LOW
      0.841808
      0.841808
      1.0
      3
      0.6
      0.256406

      7
      logic_layer
      42
      (stroke_flg & age_HIGH & ~age_LOW & ~sapsi_fir...
      0.870056
      0.870056
      1.0
      8
      0.6
      0.337131

      8
      logic_layer
      42
      stroke_flg & age_HIGH & ~sapsi_first_LOW
      0.887006
      0.887006
      1.0
      3
      0.6
      0.307857

      9
      logic_layer
      42
      stroke_flg & age_HIGH & ~sapsi_first_LOW
      0.870056
      0.870056
      1.0
      3
      0.6
      0.285181

# Decision tree

```python
results_tree = c_to_y(method='tree', verbose=False)
results_tree
```

    Split [1/10]
    Split [2/10]
    Split [3/10]
    Split [4/10]
    Split [5/10]
    Split [6/10]
    Split [7/10]
    Split [8/10]
    Split [9/10]
    Split [10/10]
    Consistency of explanations: 0.7510

      method
      split
      explanation
      model_accuracy
      explanation_accuracy
      explanation_fidelity
      explanation_complexity
      explanation_consistency
      elapsed_time

      0
      tree
      42
      (age_HIGH <= 0.50 & resp_flg <= 0.50 & stroke_...
      0.752809
      0.752809
      1.0
      838
      0.750955
      0.002992

      1
      tree
      42
      (stroke_flg <= 0.50 & age_HIGH <= 0.50 & copd_...
      0.775281
      0.775281
      1.0
      994
      0.750955
      0.004987

      2
      tree
      42
      (stroke_flg <= 0.50 & age_HIGH <= 0.50 & resp_...
      0.808989
      0.808989
      1.0
      993
      0.750955
      0.002990

      3
      tree
      42
      (age_HIGH <= 0.50 & stroke_flg <= 0.50 & sapsi...
      0.758427
      0.758427
      1.0
      1002
      0.750955
      0.002991

      4
      tree
      42
      (age_HIGH <= 0.50 & sapsi_first_HIGH <= 0.50 &...
      0.735955
      0.735955
      1.0
      927
      0.750955
      0.004987

      5
      tree
      42
      (stroke_flg <= 0.50 & age_HIGH <= 0.50 & mal_f...
      0.792135
      0.792135
      1.0
      903
      0.750955
      0.004003

      6
      tree
      42
      (stroke_flg <= 0.50 & age_HIGH <= 0.50 & mal_f...
      0.774011
      0.774011
      1.0
      864
      0.750955
      0.002991

      7
      tree
      42
      (stroke_flg <= 0.50 & age_HIGH <= 0.50 & copd_...
      0.790960
      0.790960
      1.0
      987
      0.750955
      0.004986

      8
      tree
      42
      (age_HIGH <= 0.50 & stroke_flg <= 0.50 & resp_...
      0.779661
      0.779661
      1.0
      811
      0.750955
      0.002992

      9
      tree
      42
      (age_HIGH <= 0.50 & age_NORMAL <= 0.50 & chlor...
      0.785311
      0.785311
      1.0
      898
      0.750955
      0.003988

# Summary

```python
cols = ['model_accuracy', 'explanation_accuracy', 'explanation_fidelity',
        'explanation_complexity', 'elapsed_time', 'explanation_consistency']
mean_cols = [f'{c}_mean' for c in cols]
sem_cols = [f'{c}_sem' for c in cols]

# pruning
df_mean = results_pruning[cols].mean()
df_sem = results_pruning[cols].sem()
df_mean.columns = mean_cols
df_sem.columns = sem_cols
summary_pruning = pd.concat([df_mean, df_sem])
summary_pruning.name = 'pruning'

# tree
df_mean = results_tree[cols].mean()
df_sem = results_tree[cols].sem()
df_mean.columns = mean_cols
df_sem.columns = sem_cols
summary_tree = pd.concat([df_mean, df_sem])
summary_tree.name = 'tree'

summary = pd.concat([summary_pruning,
                     summary_tree], axis=1).T
summary.columns = mean_cols + sem_cols
summary.to_csv(os.path.join(results_dir, 'summary.csv'))

summary
```

      model_accuracy_mean
      explanation_accuracy_mean
      explanation_fidelity_mean
      explanation_complexity_mean
      elapsed_time_mean
      explanation_consistency_mean
      model_accuracy_sem
      explanation_accuracy_sem
      explanation_fidelity_sem
      explanation_complexity_sem
      elapsed_time_sem
      explanation_consistency_sem

      pruning
      0.868241
      0.868241
      1.0
      4.9
      0.308134
      0.600000
      0.004056
      0.004056
      0.0
      0.982627
      0.012289
      3.700743e-17

      tree
      0.775354
      0.775354
      1.0
      921.7
      0.003791
      0.750955
      0.006768
      0.006768
      0.0
      22.249120
      0.000290
      0.000000e+00


# Adding the $8\pi T^{\mu\nu}$ Matter Source Terms to Einstein's equations of general relativity in the [BSSN](http://www2.yukawa.kyoto-u.ac.jp/~yuichiro.sekiguchi/3+1.pdf) formalism, in ***curvilinear*** coordinates, using a covariant reference metric approach: C code generation of the evolution equations' right-hand sides

## Python module containing these expressions: [BSSN/BSSN_RHS.py](../edit/BSSN/BSSN_RHSs.py)

### ***Citations***: Generic curvilinear coordinate reference metric approach matches that of [Ruchlin, Etienne, and Baumgarte (2018)](https://arxiv.org/abs/1712.07658), which is an extension of the spherical coordinate reference metric approach of [Baumgarte, Montero, Cordero-Carrión, and Müller (2012)](https://arxiv.org/abs/1211.6632), which builds upon the covariant "Lagrangian" BSSN formalism of [Brown (2009)](https://arxiv.org/abs/0902.3652). *See also citations within each article.*

In the [previous module](Tutorial-BSSNCurvilinear.ipynb), we generated SymPy expressions for all of the covariant BSSN equations in curvilinear coordinates, rescaling tensors to factor out coordinate singularities (and render them harmless). To wit:

\begin{align}
  \partial_t \bar{\gamma}_{i j} {} = {} & \left[\beta^k \partial_k \bar{\gamma}_{ij} + \partial_i \beta^k \bar{\gamma}_{kj} + \partial_j \beta^k \bar{\gamma}_{ik} \right] + \frac{2}{3} \bar{\gamma}_{i j} \left (\alpha \bar{A}_{k}^{k} - \bar{D}_{k} \beta^{k}\right ) - 2 \alpha \bar{A}_{i j} \; , \\
  \partial_t \bar{A}_{i j} {} = {} & \left[\beta^k \partial_k \bar{A}_{ij} + \partial_i \beta^k \bar{A}_{kj} + \partial_j \beta^k \bar{A}_{ik} \right] - \frac{2}{3} \bar{A}_{i j} \bar{D}_{k} \beta^{k} - 2 \alpha \bar{A}_{i k} {\bar{A}^{k}}_{j} + \alpha \bar{A}_{i j} K \nonumber \\
  & + e^{-4 \phi} \left \{-2 \alpha \bar{D}_{i} \bar{D}_{j} \phi + 4 \alpha \bar{D}_{i} \phi \bar{D}_{j} \phi  + 4 \bar{D}_{(i} \alpha \bar{D}_{j)} \phi - \bar{D}_{i} \bar{D}_{j} \alpha + \alpha \bar{R}_{i j} \right \}^{\text{TF}} \; , \\
  \partial_t \phi {} = {} & \left[\beta^k \partial_k \phi \right] + \frac{1}{6} \left (\bar{D}_{k} \beta^{k} - \alpha K \right ) \; , \\
  \partial_{t} K {} = {} & \left[\beta^k \partial_k K \right] + \frac{1}{3} \alpha K^{2} + \alpha \bar{A}_{i j} \bar{A}^{i j} - e^{-4 \phi} \left (\bar{D}_{i} \bar{D}^{i} \alpha + 2 \bar{D}^{i} \alpha \bar{D}_{i} \phi \right ) \; , \\
  \partial_t \bar{\Lambda}^{i} {} = {} & \left[\beta^k \partial_k \bar{\Lambda}^i - \partial_k \beta^i \bar{\Lambda}^k \right] + \bar{\gamma}^{j k} \hat{D}_{j} \hat{D}_{k} \beta^{i} + \frac{2}{3} \Delta^{i} \bar{D}_{j} \beta^{j} + \frac{1}{3} \bar{D}^{i} \bar{D}_{j} \beta^{j} \nonumber \\
  & - 2 \bar{A}^{i j} \left (\partial_{j} \alpha - 6 \partial_{j} \phi \right ) + 2 \alpha \bar{A}^{j k} \Delta_{j k}^{i}  -\frac{4}{3} \alpha \bar{\gamma}^{i j} \partial_{j} K \\
\partial_t \alpha &= \left[\beta^i \partial_i \alpha\right] - 2 \alpha K \\
  \partial_{t} \beta^{i} &= \left[\beta^j \partial_j \beta^i\right] + B^{i} \\
  \partial_{t} B^{i} &= \left[\beta^j \partial_j B^i\right] + \frac{3}{4} \partial_{0} \bar{\Lambda}^{i} - \eta B^{i}, \; .
\end{align}

where the terms moved from the left-hand side are enclosed in square braces. Note that $\partial_{0} \bar{\Lambda}^{i}$ in the right-hand side of the $\partial_{t} B^{i}$ equation is computed by adding $\beta^j \partial_j \bar{\Lambda}^i$ to the right-hand side expression given for $\partial_t \bar{\Lambda}^i$, so no explicit time dependence occurs in the right-hand sides of the BSSN evolution equations and the Method of Lines can be applied directly.

Of course the above BSSN evolution equations cannot appear as written above in our actual code, because (as discussed in the [previous module](Tutorial-BSSNCurvilinear.ipynb)) tensorial expressions can diverge at coordinate singularities. So the equations above will simply need to be rescaled according to the prescription (again, described in the [previous module](Tutorial-BSSNCurvilinear.ipynb)).

## Adding the matter source terms

Analyzing Eqs. 9 of [Baumgarte, Montero, Cordero-Carrión, and Müller](https://arxiv.org/pdf/1211.6632.pdf), we see that matter source terms $T_{\mu\nu}$ to Einstein's equations only adjust the right-hand sides of the $\partial_t \bar{A}_{ij}$, $\partial_t K$, and $\partial_t \bar{\Lambda}^i$ equations, as follows:

\begin{array}
\ \partial_t \bar{A}_{ij} &= [\text{as above}]\ {\color{blue}{-\ 8\pi \alpha e^{4\phi} \left(S_{ij}\right)^{\rm TF}}} \\
\partial_t K &= [\text{as above}]\ {\color{blue}{+\ 4\pi \alpha (\rho + S)}} \\
\partial_t \bar{\Lambda}^i &= [\text{as above}]\ {\color{blue}{-\ 16\pi \alpha \gamma^{ij} S_j}}, \\
\end{array}

where $\rho$, $S$, $S_i$, and $S_{ij}$ are related to the stress-energy tensor $T^{\mu\nu}$ as follows (note that Greek indices are 4-dimensional and Latin indices are 3-dimensional):

\begin{array}
\ S_{ij} &= \gamma_{i \mu} \gamma_{j \nu} T^{\mu \nu} \\
S_{i} &= -\gamma_{i\mu} n_\nu T^{\mu\nu} \\
S &= \gamma^{ij} S_{ij} \\
\rho &= n_\mu n_\nu T^{\mu\nu},
\end{array}

where the definition of 3-metric (Eq. 2.27 in Baumgarte & Shapiro) implies

$$\gamma_{\mu \nu} = g_{\mu \nu} + n_\mu n_\nu.$$

Here, $n_\mu = (-\alpha,0,0,0)$ is the normal one-form on each spatial slice, and

$$
g_{\mu\nu} = \begin{pmatrix}
-\alpha^2 + \beta^k \beta_k & \beta_i \\
\beta_j & \bar{\gamma}_{ij}
\end{pmatrix}.
$$

## Numerical Implementation

The above expressions are written in backwards order from what we need in our numerical implementation, as for example $S_{ij}$ depends on $\gamma_{i\mu}$, which depends on $n_{\mu}$ and $g_{\mu\nu}$; $g_{\mu\nu}$ depends on known quantities $\bar{\gamma}_{ij}$, $\alpha$, and $\beta^i$; and $\beta_k = \bar{\gamma}_{ik} \beta^i$.

So let's set about doing this in the proper order.

### Step 1: Load needed modules and set NRPy+ parameters
First we load the needed NRPy+ modules and set parameters for this module.

```python
import sympy as sp
import NRPy_param_funcs as par
import indexedexp as ixp
import reference_metric as rfm
import BSSN.BSSN_RHSs as bssn
bssn.BSSN_RHSs()

thismodule = "BSSN_RHSs_matter_source_terms"
PI = par.Cparameters("BUILTIN",thismodule,["M_PI"])

# We choose dimension DIM=3 here, as our base variables use a 3+1 decomposition.
DIM = 3
```

### Step 2: Set the 4-metric $g_{\mu \nu}$

$$
g_{\mu\nu} = \begin{pmatrix}
-\alpha^2 + \beta^k \beta_k & \beta_i \\
\beta_j & \bar{\gamma}_{ij}
\end{pmatrix},
$$
which are all provided by the BSSN_RHSs module, except we need $\beta_k = \bar{\gamma}_{ik} \beta^i$, which is provided by B&S (Baumgarte & Shapiro, *Numerical Relativity*) Eq. 2.121:

```python
# To get \gamma_{\mu \nu} = gammabar4DD[mu][nu], we'll need to construct the 4-metric, using Eq. 2.122 in B&S:
g4DD = ixp.zerorank2(DIM=4)

# Eq. 2.121 in B&S
betaD = ixp.zerorank1()
for i in range(DIM):
    for j in range(DIM):
        betaD[i] += bssn.gammabarDD[i][j]*bssn.betaU[j]

# Now compute the beta contraction.
beta2 = sp.sympify(0)
for i in range(DIM):
    beta2 += bssn.betaU[i]*betaD[i]

# Eq. 2.122 in B&S
g4DD[0][0] = -bssn.alpha**2 + beta2
for mu in range(1,4):
    g4DD[mu][0] = g4DD[0][mu] = betaD[mu-1]
for mu in range(1,4):
    for nu in range(1,4):
        g4DD[mu][nu] = bssn.gammabarDD[mu-1][nu-1]
```

### Step 3: Set $\gamma_{\mu\nu} = g_{\mu\nu} + n_\mu n_\nu$ (B&S Eq. 2.27)

Eq. 2.117 in B&S defines $n_\mu = \{-\alpha,0,0,0\}$, and we have already declared $g_{\mu\nu} = \text{g4DD[mu][nu]}$ above, so we have all we need to compute $\gamma_{\mu\nu} = g_{\mu\nu} + n_\mu n_\nu$:

```python
# n_{\mu}: Eq. 2.117 in B&S
nD  = [-bssn.alpha,sp.sympify(0),sp.sympify(0),sp.sympify(0)]

# Eq. 2.27 in B&S
gammabar4DD = ixp.zerorank2(DIM=4)
for mu in range(4):
    for nu in range(4):
        gammabar4DD[mu][nu] = g4DD[mu][nu] + nD[mu]*nD[nu]
```

### Step 4: Compute $S_{ij}$, $S_{i}$, $S$, and $\rho$, via Eqs. 10 of [Baumgarte, Montero, Cordero-Carrión, and Müller](https://arxiv.org/pdf/1211.6632.pdf):

\begin{array}
\ S_{ij} &= \gamma_{i \mu} \gamma_{j \nu} T^{\mu \nu} \\
S_{i} &= -\gamma_{i\mu} n_\nu T^{\mu\nu} \\
S &= \gamma^{ij} S_{ij} \\
\rho &= n_\mu n_\nu T^{\mu\nu},
\end{array}

```python
# T4UU[mu][nu] = T^{\mu\nu} (the *un-rescaled* stress-energy tensor) is assumed to be provided as input.
T4UU = ixp.register_gridfunctions_for_single_rank2("AUX","T4UU", "sym01",DIM=4)

SDD = ixp.zerorank2()
SD  = ixp.zerorank1()
S   = sp.sympify(0)
rho = sp.sympify(0)

# Compute S_{ij}:
for i in range(DIM):
    for j in range(DIM):
        for mu in range(4):
            for nu in range(4):
                SDD[i][j] += gammabar4DD[i][mu]*gammabar4DD[j][nu]*T4UU[mu][nu]

# Compute S_{i}
for i in range(DIM):
    for mu in range(4):
        for nu in range(4):
            SD[i] += -gammabar4DD[i][mu]*nD[nu]*T4UU[mu][nu]

# Compute S
for i in range(DIM):
    for j in range(DIM):
        S += bssn.gammabarUU[i][j]*SDD[i][j]

# Compute rho
for mu in range(4):
    for nu in range(4):
        rho += nD[mu]*nD[nu]*T4UU[mu][nu]
```

### Step 5: Add matter source terms to BSSN right-hand sides *a la* Eqs. 9 of [Baumgarte, Montero, Cordero-Carrión, and Müller](https://arxiv.org/pdf/1211.6632.pdf):

\begin{array}
\ \partial_t \bar{A}_{ij} &= [\text{vacuum terms}]\ {\color{blue}{-\ 8\pi \alpha e^{4\phi} \left(S_{ij}\right)^{\rm TF}}} \\
\partial_t K &= [\text{vacuum terms}]\ {\color{blue}{+\ 4\pi \alpha (\rho + S)}} \\
\partial_t \bar{\Lambda}^i &= [\text{vacuum terms}]\ {\color{blue}{-\ 16\pi \alpha \gamma^{ij} S_j}}, \\
\end{array}

which require that we first compute $S_{ij}^{\rm TF} = \gamma^{ij} S_{ij}$.

After setting these expressions, we apply the usual rescalings:

```python
# Step 1: Add matter source term to RHS of \partial_t \bar{A}_{ij} equation:
# Step 1a: Compute trace of S_{ij}:
S_trace = sp.sympify(0)
for i in range(DIM):
    for j in range(DIM):
        S_trace += bssn.gammabarUU[i][j]*SDD[i][j]

# Step 1b: find the trace-free S_{ij}, and add *rescaled* matter source term to Abar_RHS:
S_TFDD = ixp.zerorank2()
for i in range(DIM):
    for j in range(DIM):
        S_TFDD[i][j] = SDD[i][j] - sp.Rational(1,3)*bssn.gammabarDD[i][j]*S_trace
        bssn.a_rhsDD[i][j] += -8*PI*bssn.alpha/bssn.exp_m4phi * S_TFDD[i][j] / rfm.ReDD[i][j]

# Step 2: Add matter source term to RHS of \partial_t K equation:
bssn.trK_rhs += 4*PI*bssn.alpha*(rho + S)

# Step 3: Next add the *rescaled* matter source term to Lambar_RHS.
# Step 3a: Define the matter source term
Lambar_sourcetermU = ixp.zerorank1()
for i in range(DIM):
    for j in range(DIM):
        Lambar_sourcetermU[i] += -16*PI*bssn.alpha*bssn.gammabarUU[i][j]*SD[j]

# Step 3b: Add the *rescaled* matter source term just computed to lambda_RHSU[i]:
for i in range(DIM):
    bssn.lambda_rhsU[i] += Lambar_sourcetermU[i] / rfm.ReU[i]
```

```python
!jupyter nbconvert --to latex --template latex_nrpy_style.tplx Tutorial-BSSNCurvilinear_Plus_Matter_Sources.ipynb
!pdflatex -interaction=batchmode Tutorial-BSSNCurvilinear_Plus_Matter_Sources.tex
!pdflatex -interaction=batchmode Tutorial-BSSNCurvilinear_Plus_Matter_Sources.tex
!pdflatex -interaction=batchmode Tutorial-BSSNCurvilinear_Plus_Matter_Sources.tex
!rm -f Tut*.out Tut*.aux Tut*.log
```

    [NbConvertApp] Converting notebook Tutorial-BSSNCurvilinear_Plus_Matter_Sources.ipynb to latex
    [NbConvertApp] Writing 38355 bytes to Tutorial-BSSNCurvilinear_Plus_Matter_Sources.tex
    This is pdfTeX, Version 3.14159265-2.6-1.40.18 (TeX Live 2017/Debian) (preloaded format=pdflatex)
     restricted \write18 enabled.
    entering extended mode
    This is pdfTeX, Version 3.14159265-2.6-1.40.18 (TeX Live 2017/Debian) (preloaded format=pdflatex)
     restricted \write18 enabled.
    entering extended mode
    This is pdfTeX, Version 3.14159265-2.6-1.40.18 (TeX Live 2017/Debian) (preloaded format=pdflatex)
     restricted \write18 enabled.
    entering extended mode


**Objective:**

In this tutorial we will create a simple magnetic problem from scratch using the SimPEG framework.

We are using the integral form of the magnetostatic problem. In the absence of free-currents or changing magnetic field, magnetic material can give rise to a secondary magnetic field according to:

$$\vec b = \frac{\mu_0}{4\pi}  \int_{V}   \vec M \cdot \nabla \nabla \left(\frac{1}{r}\right) \; dV $$

Where $\mu_0$ is the magnetic permealitity of free-space, $\vec M$ is the magnetization per unit volume and $r$  defines the distance between the observed field $\vec b$ and the magnetized object. Assuming a purely induced response, the strenght of magnetization can be written as:

$$ \vec M = \mu_0 \kappa \vec H_0 $$

where $\vec H$ is an external inducing magnetic field, and $\kappa$ the magnetic susceptibility of matter.
As derived by Sharma 1966, the integral can be evaluated for rectangular prisms such that:

$$ \vec b(P) =  \mathbf{T} \cdot \vec H_0 \; \kappa $$

Where the tensor matrix $\bf{T}$ relates the three components of magnetization $\vec M$ to the components of the field $\vec b$:

$$\mathbf{T} =
	 \begin{pmatrix}
       		T_{xx} & T_{xy} & T_{xz}    \\
		T_{yx} & T_{yy} & T_{yz}    \\
		T_{zx} & T_{zy} & T_{zz}
	\end{pmatrix} $$

In general, we discretize the earth into a collection of cells, each contributing to the magnetic data such that:

$$\vec b(P) = \sum_{j=1}^{nc} \mathbf{T}_j \cdot \vec H_0 \; \kappa_j$$

giving rise to a linear problem.

```python
## Need to be on \pf\dev branch !!!
%pylab inline
import SimPEG.PF as PF
from SimPEG import *
from SimPEG.Utils import io_utils
```

    Populating the interactive namespace from numpy and matplotlib
    Efficiency Warning: Interpolation will be slow, use setup.py!

                python setup.py build_ext --inplace

```python
import matplotlib
matplotlib.rcParams['font.size'] = 14
```

```python
# We first need to create a susceptibility model.
# Based on a set of parametric surfaces representing TKC,
# we use VTK to discretize the 3-D space.

mshfile = 'MEsh_TEst.msh'

model_dir = '../../Geological_model/'

# Load mesh file
mesh = Mesh.TensorMesh.readUBC(model_dir+mshfile)

# Create our own mesh!
csx, csy, csz = 10., 10., 10.
ncx, ncy, ncz = 55, 55, 30
npad = 10
hx = [(csx,npad, -1.3),(csx,ncx),(csx,npad, 1.3)]
hy = [(csy,npad, -1.3),(csy,ncy),(csy,npad, 1.3)]
hz = [(csz,npad, -1.3),(csz,ncz), (10,20)]
mesh = Mesh.TensorMesh([hx, hy, hz],x0="CCN")
xc = 300+5.57e5
yc = 600+7.133e6
zc = 450.
x0_new = np.r_[mesh.x0[0]+xc, mesh.x0[1]+yc, mesh.x0[2]+zc]
mesh._x0 = x0_new

Mesh.TensorMesh.writeUBC(mesh,model_dir+"PF_mesh_UTM.msh")
# Define no-data-value
ndv = -100

# Define survey flight height
Z_bird = 20.

# Read in topo surface
topofile = model_dir+'TKCtopo.dat'
geosurf = [
    [model_dir+'Till.ts',True,True,0],
    [model_dir+'XVK.ts',True,True,1],
    [model_dir+'PK1.ts',True,True,2],
    [model_dir+'PK2.ts',True,True,3],
    [model_dir+'PK3.ts',True,True,4],
    [model_dir+'HK1.ts',True,True,5],
    [model_dir+'VK.ts',True,True,6]
]

```

```python
print model_dir+"PF_mesh_UTM.msh"
```

    ../../Geological_model/PF_mesh_UTM.msh

```python
import time as tm
import mpl_toolkits.mplot3d as a3
import matplotlib.colors as colors
import scipy as sp
# from mayavi import mlab
# mlab.close(all=True)
# fig = plt.figure(figsize=(8,6))

# mlab.figure(bgcolor=(1.,1.,1.))
modelInd = np.ones(mesh.nC)*ndv
for ii in range(len(geosurf)):
    tin = tm.time()
    print "Computing indices with VTK: " + geosurf[ii][0]
    T, S = io_utils.read_GOCAD_ts(geosurf[ii][0])
    indx = io_utils.surface2inds(T,S,mesh, boundaries=geosurf[ii][1], internal=geosurf[ii][2])
    print "VTK operation completed in " + str(tm.time() - tin) + " sec"
    modelInd[indx] = geosurf[ii][3]

#     S = S-1
#     mlab.triangular_mesh(T[:,0], T[:,1], T[:,2], S)
#     mlab.view(azimuth=-45.,distance=1000,elevation=75.)

# arr = mlab.screenshot()

# plt.imshow(arr)
```

    Computing indices with VTK: ../../Geological_model/Till.ts
    Extracting indices from grid...
    VTK operation completed in 28.1850001812 sec
    Computing indices with VTK: ../../Geological_model/XVK.ts
    Extracting indices from grid...
    VTK operation completed in 5.34999990463 sec
    Computing indices with VTK: ../../Geological_model/PK1.ts
    Extracting indices from grid...
    VTK operation completed in 13.1600000858 sec
    Computing indices with VTK: ../../Geological_model/PK2.ts
    Extracting indices from grid...
    VTK operation completed in 5.3599998951 sec
    Computing indices with VTK: ../../Geological_model/PK3.ts
    Extracting indices from grid...
    VTK operation completed in 7.12800002098 sec
    Computing indices with VTK: ../../Geological_model/HK1.ts
    Extracting indices from grid...
    VTK operation completed in 12.5909998417 sec
    Computing indices with VTK: ../../Geological_model/VK.ts
    Extracting indices from grid...
    VTK operation completed in 13.7820000648 sec

```python
# Load topography file in UBC format and find the active cells
#T, S = io_utils.read_GOCAD_ts(topsurf)
#indx = io_utils.surface2inds(T,S, mesh, boundaries=True, internal=True)
#actv = np.zeros(mesh.nC)
#actv[indx] = 1

topo = np.genfromtxt(topofile,skip_header=1)
# Find the active cells
actv = Utils.surface2ind_topo(mesh, topo, gridLoc='N')
# actv = PF.Magnetics.getActiveTopo(mesh,topo,'N')

# Create active map to go from reduce set to full
actvMap = Maps.InjectActiveCells(mesh, actv, ndv)

print "Active cells created from topography!"
```

    Active cells created from topography!

```python
# Build model

def getModel(Till=0.0, XVK=2e-3, PK1=5e-3, PK2=2e-3, PK3=1e-3, HK1=2e-2, VK=1e-2, bkgr=0.):
    vals = [Till, XVK, PK1, PK2, PK3, HK1, VK]
    model= np.ones(mesh.nC) * bkgr

    for ii, sus in zip(range(7),vals):
        model[modelInd == ii] = sus
    return model
model = getModel()
sus = getModel()
model = model[actv]
```

```python
from ipywidgets.widgets import interact, IntSlider
```

```python
# Here you can visualize the current model
m_true = actvMap*model
Mesh.TensorMesh.writeModelUBC(mesh,model_dir+"Synthetic_mag.sus",m_true)
airc = m_true==ndv
m_true[airc] = np.nan
def slide(s,normal):
    colorbar(mesh.plotSlice(m_true, normal=normal, ind=s, clim=np.r_[model.min(), model.max()])[0])
    plt.gca().set_aspect('equal')
interact(slide, s=(0,60), normal=['X','Y','Z'])
```

**Forward system:**

Now that we have all our spatial components, we can create our linear system. For a single location and single component of the data, the system would looks like this:

$$ b_x =
	\begin{bmatrix}
	T_{xx}^1 &... &T_{xx}^{nc} & T_{xy}^1 & ... & T_{xy}^{nc} & T_{xz}^1 & ... & T_{xz}^{nc}\\
	 \end{bmatrix}
	 \begin{bmatrix}
		\mathbf{M}_x \\ \mathbf{M}_y \\ \mathbf{M}_z
	\end{bmatrix} \\ $$

where each of $T_{xx},\;T_{xy},\;T_{xz}$ are [nc x 1] long. For the $y$ and $z$ component, we need the two other rows of the tensor $\mathbf{T}$.
In our simple induced case, the magnetization direction $\mathbf{M_x,\;M_y\;,Mz}$ are known and assumed to be constant everywhere, so we can reduce the size of the system such that:

$$ \vec{\mathbf{d}}_{\text{pred}} = (\mathbf{T\cdot M})\; \kappa$$

In most geophysical surveys, we are not collecting all three components, but rather the magnitude of the field, or $Total\;Magnetic\;Intensity$ (TMI) data.
Because the inducing field is really large, we will assume that the anomalous fields are parallel to $H_0$:

$$ d^{TMI}  = \hat H_0 \cdot \vec d$$

We then end up with a much smaller system:

$$ d^{TMI} = \mathbf{F\; \kappa}$$

where $\mathbf{F} \in \mathbb{R}^{nd \times nc}$ is our $forward$ operator.

```python
from scipy.interpolate import NearestNDInterpolator
# We need to define the direction of the inducing field
# From old convention, field orientation is given as an azimuth from North
# (positive clockwise) and dip from the horizontal (positive downward).
# The field parameters at TKC are [H:60,308 nT, I:83.8 d D:25.4 d ]
H0 = (60308.,83.8,25.4)

# We create a synthetic survey with observations in cell center.
X, Y = np.meshgrid(mesh.vectorCCx[npad:-npad:2], mesh.vectorCCy[npad:-npad:2])

# Using our topography, we trape the survey and shift it up by the flight height
Ftopo = NearestNDInterpolator(topo[:,:2], topo[:,2])
Z = Ftopo(Utils.mkvc(X.T),Utils.mkvc(Y.T)) + Z_bird

rxLoc = np.c_[Utils.mkvc(X.T), Utils.mkvc(Y.T), Utils.mkvc(Z.T)]
rxLoc = PF.BaseMag.RxObs(rxLoc)

srcField = PF.BaseMag.SrcField([rxLoc])
srcField.param = H0
survey = PF.BaseMag.LinearSurvey(srcField)
```

```python
# Now that we have a model and a survey we can build the linear system ...
nactv = np.int(np.sum(actv))

# Creat reduced identity map
idenMap = Maps.IdentityMap(nP=nactv)

# Create the forward model operator
prob = PF.Magnetics.Problem3D_Integral(mesh, mapping=idenMap, actInd=actv, forwardOnly=False, rtype = 'tmi')

# Pair the survey and problem
survey.pair(prob)

```

```python
# Fist time that we ask for predicted data,
# the dense matrix T is calculated.
# This is generally the bottleneck of the integral formulation,
# in terms of time and memory.
d = prob.fields(model)

# Add noise to the data and assign uncertainties
data = d + randn(len(d)) # We add some random Gaussian noise (1nT)
wd = np.ones(len(data))*1. # Assign flat uncertainties

survey.dobs = data
survey.std = wd

PF.Magnetics.writeUBCobs('MAG_Synthetic_data.obs',survey,data)

d2D = data.reshape(X.shape)
dat = plt.contourf(X-xc,Y-yc, d2D,40)
plt.gca().set_aspect('equal')
plt.plot(X.flatten()-xc,Y.flatten()-yc,'k.', ms=2)
plt.colorbar(dat)
plt.xlabel("Easting (m)")
plt.ylabel("Northing (m)")
plt.title("Total Magnetic Intensity (nT)")
xlim(-500, 500)
ylim(-500, 500)
```

```python
# d2Dx, d2Dy, d2Dz = d[0:survey.nRx].reshape(X.shape), d[survey.nRx:2*survey.nRx].reshape(X.shape),  d[2*survey.nRx:].reshape(X.shape)

# plt.figure(figsize=[5,5])
# xc, yc = (X.min()+X.max())*0.5, (Y.min()+Y.max())*0.5
# #     colorbar(imshow(d2D,extent=[X.min()-xc, X.max()-xc, Y.min()-yc, Y.max()-yc],origin = 'lower'))
# #     plt.contour(X-xc,Y-yc, d2D,20,colors='k')
# # dat = plt.contourf(X-xc,Y-yc, d2Dz,40)

# # plt.gca().set_aspect('equal')

# mlab.figure(bgcolor=(1.,1.,1.))
# mlab.triangular_mesh(T[:,0], T[:,1], T[:,2], S)
# mlab.quiver3d(X,Y,Z,d2Dx, d2Dy, d2Dz, color)
# mlab.view(azimuth=-45.,distance=1500,elevation=75.)
# arr = mlab.screenshot()

# plt.imshow(arr)
```

**Inverse problem**

We have generated synthetic data, we now what to see if we can solve the inverse. Using the usual formulation, we seek a model that can reproduce the data, let’s say a least-squares measure of the form:

\begin{equation}
\phi_d =   \|\mathbf{W}_d \left( \mathbb{F}[\mathbf{m}] - \mathbf{d}^{obs} \right)\|_2^2
\end{equation}

The inverse problem is hard because we don’t have great data coverage, and the Earth is big, and there is usually noise in the data. So we need to add something to regularize it.
The simplest way to do it is to penalize solutions that won’t make sense geologically, for example to assume that the model is small.
The usual smooth inversion function use an l2-norm measure:

\begin{equation}
\phi_d =   \|\mathbf{W}_d \left( \mathbb{F}[\mathbf{m}] - \mathbf{d}^{obs} \right)\|_2^2 \\
\phi_m = \beta \Big [ {\| \mathbf{W}_s \;( \mathbf{m - m^{ref}})\|}^2_2  + \sum_{i = x,y,z}  {\|   \mathbf{W}_i  \; \mathbf{G}_i \; \mathbf{m}\|}^2_2  \Big ]\;,
\end{equation}

The full objective function to be minimized can be written as:
\begin{equation}
\phi(m) =  \phi_d + \beta \phi_m\;,
\end{equation}
which will yield our usual *small* and *smooth* models.

We propose a fancier regularization function that can allow to recover *sparse* and *blocky* solutions.
Starting with the well known Ekblom norm:
\begin{equation}
\phi_m =  \sum_{i=1}^{nc} {(x_i^2 + \epsilon^2)}^{p/2} \;,
\end{equation}
where $x_i$ denotes some function of the model parameter, and $\epsilon$ is a small value to avoid singularity as $m\rightarrow0$.
For p=2, we get the usual least-squares measure and we recover the regularization presented above. For $p \leq 1$, the function becomes non-linear which requires some tweaking.

We can linearize the function by updating the penality function iteratively, commonly known as an Iterative Re-weighted Least-Squares (IRLS) method:
\begin{equation}
\phi_m^{(k)} =  \frac{1}{2}\sum_{i=1}^{nc} r_i \; x_i^2
\end{equation}
where we added the superscript $\square^{(k)}$ to denote the IRLS iterations. The weights $r(x)$ are computed from model values obtained at a previous iteration such that:
\begin{equation}
	{r}_i  ={\Big( {({x_i}^{(k-1)})}^{2} + \epsilon^2 \Big)}^{p/2 - 1} \;,
\end{equation}
where ${r}(x) \in \mathbb{R}^{nc}$.

In matrix form, our objective function simply becomes:
\begin{equation}
\phi(m) =   \|\mathbf{W}_d \left( \mathbb{F}[\mathbf{m}] - \mathbf{d}^{obs} \right)\|_2^2 + \beta \Big [ {\| \mathbf{W}_s \;\mathbf{R}_s\;( \mathbf{m - m^{ref}})\|}^2_2  + \sum_{i = x,y,z}  {\|   \mathbf{W}_i\; \mathbf{R}_i  \; \mathbf{G}_i \; \mathbf{m}\|}^2_2  \Big ]\;,
\end{equation}
where the IRLS weights $\mathbf{R}_s$ and $\mathbf{R}_i$ are diagonal matrices defined as:
\begin{equation}
\begin{split}
	{R}_{s_{jj}}  &=  \sqrt{\eta_p}{\Big[ {({m_j}^{(k-1)})}^{2} + \epsilon_p^2 \Big]}^{(p/2 - 1)/2} \\
	{R}_{i_{jj}}  &=  \sqrt{\eta_q}{\Big[ {\left ({{(G_i\;m^{(k-1)})}_j }\right)}^{2} + \epsilon_q^2 \Big]}^{(q/2 - 1)/2} \\
\eta_p &=  {\epsilon_p}^{(1-p/2)} \\
\eta_q &=   {\epsilon_q}^{(1-q/2)}  \;,
\end{split}
\end{equation}

we added two scaling parameters $\eta_p$ and $\eta_q$ for reasons that we won't dicuss here, but turn out to be important to get stable solves.

In order to initialize the IRLS and get an estimate for the stabilizing parameters $\epsilon_p$ and $\epsilon_q$, we first invert with the smooth $l_2$-norm.
The whole IRLS process is implemented with a directive added to the inversion workflow (see below).

```python
# It is potential fields, so we will need to push the inverison down
# Create distance weights from our linera forward operator
wr = np.sum(prob.G**2.,axis=0)**0.5
wr = ( wr/np.max(wr) )
#wr = PF.Magnetics.get_dist_wgt(mesh, survey.srcField.rxList[0].locs, actv, 3., np.min(mesh.hx)/4.)
#wr = wr**2. # Need to square to be included as weights

reg = Regularization.Sparse(mesh, indActive = actv, mapping = idenMap)
reg.cell_weights = wr

dmis = DataMisfit.l2_DataMisfit(survey)
dmis.Wd = 1/wd

# Add directives to the inversion
opt = Optimization.ProjectedGNCG(maxIter=100 ,lower=0.,upper=1., maxIterLS = 20, maxIterCG= 10, tolCG = 1e-3)
invProb = InvProblem.BaseInvProblem(dmis, reg, opt)
betaest = Directives.BetaEstimate_ByEig()

# Here is where the norms are applied
eps = [3e-4,3e-4] # Threshold parameters to penalize different model parameters
norms = [0,1,1,1] # Norms applies on the model and 3 gradients [p, qx, qy, qz]

IRLS = Directives.Update_IRLS( norms=norms,  eps=eps, f_min_change = 1e-1, minGNiter=6)
update_Jacobi = Directives.Update_lin_PreCond()
inv = Inversion.BaseInversion(invProb, directiveList=[IRLS,betaest,update_Jacobi])

m0 = np.ones(idenMap.nP)*1e-4

```

```python
# Run inversion...
mrec = inv.run(m0)
```

    SimPEG.InvProblem will set Regularization.mref to m0.
    SimPEG.InvProblem is setting bfgsH0 to the inverse of the eval2Deriv.
                        ***Done using same Solver and solverOpts as the problem***
    =============================== Projected GNCG ===============================
      #     beta     phi_d     phi_m       f      |proj(x-g)-x|  LS    Comment
    -----------------------------------------------------------------------------
       0  1.18e+09  6.86e+05  0.00e+00  6.86e+05    4.76e+02      0
       1  5.89e+08  5.68e+05  4.45e-05  5.94e+05    4.73e+02      0
       2  2.95e+08  4.63e+05  1.61e-04  5.11e+05    4.68e+02      0   Skip BFGS
       3  1.47e+08  3.49e+05  4.31e-04  4.12e+05    4.62e+02      0   Skip BFGS
       4  7.36e+07  2.35e+05  9.74e-04  3.06e+05    4.50e+02      0   Skip BFGS
       5  3.68e+07  1.41e+05  1.85e-03  2.09e+05    4.32e+02      0   Skip BFGS
       6  1.84e+07  7.56e+04  3.06e-03  1.32e+05    4.10e+02      0   Skip BFGS
       7  9.21e+06  3.42e+04  4.55e-03  7.61e+04    4.16e+02      0   Skip BFGS
       8  4.60e+06  1.31e+04  6.02e-03  4.08e+04    4.15e+02      0   Skip BFGS
       9  2.30e+06  4.81e+03  7.16e-03  2.13e+04    3.61e+02      0   Skip BFGS
      10  1.15e+06  1.79e+03  7.95e-03  1.09e+04    3.87e+02      0   Skip BFGS
      11  5.75e+05  7.55e+02  8.47e-03  5.63e+03    4.14e+02      0   Skip BFGS
    Convergence with smooth l2-norm regularization: Start IRLS steps...
    L[p qx qy qz]-norm : [0, 1, 1, 1]
    eps_p: 0.0003 eps_q: 0.0003
    Regularization decrease: 0.000e+00
      12  5.75e+05  3.76e+02  8.47e-03  5.25e+03    3.97e+02      0   Skip BFGS
      13  5.75e+05  5.50e+02  4.95e-03  3.40e+03    3.21e+02      0
      14  5.75e+05  3.97e+02  4.44e-03  2.95e+03    3.16e+02      0
      15  5.75e+05  4.09e+02  4.10e-03  2.77e+03    3.11e+02      0
      16  5.75e+05  3.77e+02  3.97e-03  2.66e+03    2.50e+02      0
      17  5.75e+05  3.84e+02  3.83e-03  2.59e+03    2.98e+02      0
    Regularization decrease: 5.747e-01
      18  6.13e+05  3.68e+02  3.83e-03  2.71e+03    2.29e+02      0
      19  6.13e+05  4.08e+02  2.70e-03  2.06e+03    3.33e+02      0
      20  6.13e+05  3.79e+02  2.48e-03  1.90e+03    3.31e+02      0
      21  6.13e+05  3.86e+02  2.35e-03  1.82e+03    2.72e+02      0
      22  6.13e+05  3.74e+02  2.29e-03  1.78e+03    2.85e+02      0
      23  6.13e+05  3.78e+02  2.23e-03  1.75e+03    2.49e+02      0
    Regularization decrease: 4.142e-01
      24  6.48e+05  3.71e+02  2.23e-03  1.82e+03    2.50e+02      0
      25  6.48e+05  3.92e+02  1.78e-03  1.54e+03    3.19e+02      0
      26  6.48e+05  3.74e+02  1.69e-03  1.47e+03    2.81e+02      0
      27  6.48e+05  3.78e+02  1.63e-03  1.44e+03    2.67e+02      0
      28  6.48e+05  3.71e+02  1.61e-03  1.41e+03    2.42e+02      0
      29  6.48e+05  3.73e+02  1.58e-03  1.39e+03    2.47e+02      0
    Regularization decrease: 2.917e-01
      30  6.88e+05  3.69e+02  1.58e-03  1.45e+03    2.30e+02      0
      31  6.88e+05  3.80e+02  1.36e-03  1.32e+03    3.06e+02      0
      32  6.88e+05  3.70e+02  1.32e-03  1.27e+03    2.57e+02      0
      33  6.88e+05  3.72e+02  1.28e-03  1.25e+03    2.67e+02      0
      34  6.88e+05  3.67e+02  1.27e-03  1.24e+03    2.26e+02      0
      35  6.88e+05  3.68e+02  1.25e-03  1.23e+03    2.45e+02      0
    Regularization decrease: 2.059e-01
      36  7.37e+05  3.66e+02  1.25e-03  1.29e+03    2.30e+02      0
      37  7.37e+05  3.73e+02  1.13e-03  1.20e+03    2.85e+02      0
      38  7.37e+05  3.65e+02  1.10e-03  1.18e+03    2.53e+02      0
      39  7.37e+05  3.67e+02  1.08e-03  1.16e+03    2.53e+02      0
      40  7.37e+05  3.63e+02  1.07e-03  1.15e+03    2.16e+02      0
      41  7.37e+05  3.64e+02  1.06e-03  1.14e+03    2.35e+02      0
    Regularization decrease: 1.502e-01
      42  7.99e+05  3.62e+02  1.06e-03  1.21e+03    2.31e+02      0
      43  7.99e+05  3.69e+02  9.74e-04  1.15e+03    2.64e+02      0
      44  7.99e+05  3.62e+02  9.60e-04  1.13e+03    2.31e+02      0
      45  7.99e+05  3.64e+02  9.43e-04  1.12e+03    2.25e+02      0
      46  7.99e+05  3.60e+02  9.39e-04  1.11e+03    2.10e+02      0
      47  7.99e+05  3.62e+02  9.29e-04  1.10e+03    2.13e+02      0
    Regularization decrease: 1.204e-01
      48  8.72e+05  3.59e+02  9.29e-04  1.17e+03    2.38e+02      0
      49  8.72e+05  3.67e+02  8.70e-04  1.13e+03    2.61e+02      0
      50  8.72e+05  3.60e+02  8.60e-04  1.11e+03    2.28e+02      0
      51  8.72e+05  3.63e+02  8.47e-04  1.10e+03    2.20e+02      0
      52  8.72e+05  3.59e+02  8.45e-04  1.10e+03    2.09e+02      0
      53  8.72e+05  3.61e+02  8.37e-04  1.09e+03    2.05e+02      0
    Regularization decrease: 9.805e-02
    Minimum decrease in regularization. End of IRLS
    ------------------------- STOP! -------------------------
    1 : |fc-fOld| = 0.0000e+00 <= tolF*(1+|f0|) = 6.8632e+04
    1 : |xc-x_last| = 8.3941e-03 <= tolX*(1+|x0|) = 1.0567e-01
    0 : |proj(x-g)-x|    = 2.0483e+02 <= tolG          = 1.0000e-01
    0 : |proj(x-g)-x|    = 2.0483e+02 <= 1e3*eps       = 1.0000e-02
    0 : maxIter   =     100    <= iter          =     54
    ------------------------- DONE! -------------------------

```python
# Get the final model back to full space
m_lp = actvMap*mrec
m_lp[airc] = np.nan

# Get the smooth model aslo
m_l2 = actvMap*reg.l2model
m_l2[airc] = np.nan

# Save both models to file
Mesh.TensorMesh.writeModelUBC(mesh,'SimPEG_MAG_l2l2.sus',m_l2)
Mesh.TensorMesh.writeModelUBC(mesh,'SimPEG_MAG_lplq.sus',m_lp)
```

```python
# Plot the recoverd models
vmin, vmax = 0., 0.015

mesh = Mesh.TensorMesh([mesh.hx, mesh.hy, mesh.hz],x0="CCN")

def slide(s,normal):

    if normal == "Z":
        fig = plt.figure(figsize(10*1.2, 8))
    else:
        fig = plt.figure(figsize(10*1.2, 4))

    ax1 = plt.subplot(2,2,3)
    dat = mesh.plotSlice(m_lp, ax = ax1, normal=normal, ind=s, clim=np.r_[vmin, vmax],pcolorOpts={'cmap':'viridis'})
#     plt.colorbar(dat[0])
    plt.gca().set_aspect('equal')
    plt.title('Compact model')

    if normal == "Z":
        xlim(-600, 600)
        ylim(-600, 600.)
    else:
        xlim(-600, 600)
        ylim(-500, 0.)

    ax2 = plt.subplot(2,2,1)
    dat = mesh.plotSlice(m_l2, ax = ax2, normal=normal, ind=s, clim=np.r_[vmin, vmax],pcolorOpts={'cmap':'viridis'})
#     plt.colorbar(dat[0])
    plt.gca().set_aspect('equal')
    plt.title('Smooth model')

    if normal == "Z":
        xlim(-600, 600)
        ylim(-600, 600.)
    else:
        xlim(-600, 600)
        ylim(-500, 0.)

    ax2.set_xticklabels([])

    ax2 = plt.subplot(1,2,2)
    dat = mesh.plotSlice(m_true, ax = ax2, normal=normal, ind=s, clim=np.r_[vmin, vmax],pcolorOpts={'cmap':'viridis'})
#     plt.colorbar(dat[0])
    plt.gca().set_aspect('equal')
    plt.title('True model')

    pos =  ax2.get_position()

    ax2.yaxis.set_visible(False)
    if normal == "Z":
        xlim(-600, 600)
        ylim(-600, 600.)
        ax2.set_position([pos.x0 -0.04 , pos.y0,  pos.width, pos.height])
    else:
        xlim(-600, 600)
        ylim(-500, 0.)

    pos =  ax2.get_position()
    cbarax = fig.add_axes([pos.x0 + 0.375 , pos.y0 + 0.05,  pos.width*0.1, pos.height*0.75])  ## the parameters are the specified position you set
    cb = fig.colorbar(dat[0],cax=cbarax, orientation="vertical", ax = ax2, ticks=np.linspace(vmin,vmax, 4))
    cb.set_label("Susceptibility (SI)",size=12)
    fig.savefig('PF_Compact.png',dpi = 150)

interact(slide, s=(0,60), normal=['X','Y','Z'])

# interact(lambda ind: viz(m_l2, ind, normal="Z"), ind=IntSlider(min=0, max=32,step=1, value=28))

```

```python
# Lets compare the distribution of model parameters and model gradients
plt.figure(figsize=[15,5])
ax = plt.subplot(121)
plt.hist(reg.l2model,100)
plt.plot((eps[0],eps[0]),(0,mesh.nC),'r--')
plt.yscale('log', nonposy='clip')
plt.title('Hist model - l2-norm')
ax = plt.subplot(122)
plt.hist(reg.regmesh.cellDiffxStencil*reg.l2model,100)
plt.plot((eps[1],eps[1]),(0,mesh.nC),'r--')
plt.yscale('log', nonposy='clip')
plt.title('Hist model gradient - l2-norm')

# Lets look at the distribution of model parameters and model gradients
plt.figure(figsize=[15,5])
ax = plt.subplot(121)
plt.hist(mrec,100)
plt.plot((eps[0],eps[0]),(0,mesh.nC),'r--')
plt.yscale('log', nonposy='clip')
plt.title('Hist model - lp-norm')
ax = plt.subplot(122)
plt.hist(reg.regmesh.cellDiffxStencil*mrec,100)
plt.plot((eps[1],eps[1]),(0,mesh.nC),'r--')
plt.yscale('log', nonposy='clip')
plt.title('Hist model gradient - lp-norm')
```

```python
Mesh.MeshIO.TensorMeshIO.writeUBC(mesh,'MAG_mesh.msh')

# import pickle
# Results = {"mesh":mesh, "model_true":sus, "model_pred":m_IRLS, "Obs":data, "XYZ": xyz}
# outputs = open("Magresults", 'wb')
# pickle.dump(Results, outputs)
# outputs.close()
```

```python

```


# [NTDS'18] milestone 4: graph signal processing
[ntds'18]: https://github.com/mdeff/ntds_2018

[Rodrigo Pena](https://people.epfl.ch/254838), [EPFL LTS2](http://lts2.epfl.ch)

## Students

* Team: `Team 2`
* Students: `Michael Allemann, Roman Bachmann, Liamarcia Bifano, Virginia Bordignon`
* Dataset: `US Senators`

## Rules

* Milestones have to be completed by teams. No collaboration between teams is allowed.
* Textual answers shall be short. Typically one to two sentences.
* Code has to be clean.
* You cannot import any other library than we imported.
* When submitting, the notebook is executed and the results are stored. I.e., if you open the notebook again it should show numerical results and plots. We won't be able to execute your notebooks.
* The notebook is re-executed from a blank state before submission. That is to be sure it is reproducible. You can click "Kernel" then "Restart & Run All" in Jupyter.

## Objective

The goal of this milestone is to do some Graph Signal Processing (GSP) on the data of your project.

### A note about plotting

There are several questions in this milestone that ask you to plot a signal on your network.
There are several ways from which you could approach it.
In all cases, compute the position of the nodes a single time at the beginning, as this is likely to be a costly operation.
Using a single layout for all the graph plots will also make it easier to compare the plots.
Indeed, the only thing changing between plots is the signal displayed.
You can represent the features/labels lying on the graph via node **colors**.
To do so, make sure to have a consistent color map throughout and remember to display a colorbar and scale in all plots, so that we can tell what numbers the colors represent.

* An option is to use the **Laplacian eigenmaps** that you have seen in the previous milestone to embed your graph on the plane. For example:
  ```
  from matplotlib import pyplot as plt
  plt.scatter(eigenvectors[:, 1], eigenvectors[:, 2], c=signal, alpha=0.5)
  plt.colorbar()
  ```
* Another option is to use the plotting capabilities of **[NetworkX](https://networkx.github.io)**.
  See the documentation of its [drawing methods](https://networkx.github.io/documentation/stable/reference/drawing.html).
  For example:
  ```
  import networkx as nx
  graph = nx.from_scipy_sparse_matrix(adjacency)
  coords = nx.spring_layout(graph)  # Force-directed layout.
  coords = eigenvectors[:, 1:3]  # Laplacian eigenmaps.
  nx.draw_networkx_nodes(graph, coords, node_size=60, node_color=signal)
  nx.draw_networkx_edges(graph, coords, alpha=0.3)
  ```
* Another option is to use the plotting capabilities of the **[PyGSP](https://github.com/epfl-lts2/pygsp)**, a Python package for Graph Signal Processing.
  **Note that your are forbidden to use the PyGSP for anything else than plotting.**
  See the documentation of its [plotting utilities](https://pygsp.readthedocs.io/en/stable/reference/plotting.html).
  For example:
  ```
  import pygsp as pg
  graph = pg.graphs.Graph(adjacency)
  graph.set_coordinates('spring')  # Force-directed layout.
  graph.set_coordinates(eigenvectors[:, 1:3])  # Laplacian eigenmaps.
  graph.plot_signal(signal)
  ```
* Yet another option is to save your graph on disk, use **[Gephi](https://gephi.org)** externally, to visualize the graph, save the graph with the Gephi coordinates and finally load the nodes coordinates back into the notebook.

We encourage you to try all the above methods before making your choice. Then be consistent and use only one throughout the milestone.
NetworkX and PyGSP should already be installed in your environement. If that's not the case, install with `conda install networkx pygsp` (after activating the `ntds_2018` environment).

## 0 - Load your network

```python
%matplotlib inline
```

If you get a `No module named 'pyunlocbox'` error when running the below cell, install the [pyunlocbox](https://github.com/epfl-lts2/pyunlocbox) with `conda install pyunlocbox` (after activating the `ntds_2018` environment).

```python
import numpy as np
import pandas as pd
from scipy import sparse
import scipy.sparse.linalg
from matplotlib import pyplot as plt
from pyunlocbox import functions, solvers
import networkx as nx
```

For this milestone, all we will need is a set of features/labels for each of the nodes on the network, as well as the Laplacian, $L,$ and Gradient, $\nabla_G,$ matrices that you have computed for your network while working on milestone 3.

Import those objects in the cell below (or recompute the Laplacian and Gradient from your stored adjacency matrix, if you wish).

_Note_: If your features/labels are not floating-point numbers, please convert them. For example, if your data has labels "cat" and "dog" for nodes that represent cats or dogs, respectively, you may assign the number `1.0` for the label "cat" and the number `-1.0` for the label "dog".

```python
adjacency = np.loadtxt("../data/adjacency.csv", delimiter=",")
n_nodes =  len(adjacency)
n_edges = int(adjacency.sum() // 2)
```

```python
# Recomputing the Laplacian
D = np.diag(np.sum(adjacency, 1)) # Degree matrix
D_norm = np.diag(np.sum(adjacency, 1)**(-1/2))
laplacian_combinatorial =  D - adjacency
laplacian_normalized =  D_norm @ laplacian_combinatorial @ D_norm

laplacian = laplacian_normalized
```

```python
# Computing the gradient (incident matrix S)
S = np.zeros((n_nodes, n_edges))
edge_idx = 0
for i in range(n_nodes):
    for k in range(i):
        if adjacency[i,k] == 1.0:
            S[i,edge_idx] = 1
            S[k,edge_idx] = -1
            edge_idx += 1

assert np.allclose(S @ S.T, laplacian_combinatorial)
```

Here we map the three parties to unique integers to represent our label signal.
We chose to give the Democrats a lower number and the Republicans a higher one, such that when plotting using the Matplotlib color map *"bwr"*, the Democrats are represented in blue, while the Republicans are red.

Furthermore, we chose to label the independent party members between the two main parties.

```python
string_labels = pd.read_csv('../data/all_votes.csv').sort_values(by=['party']).party.values

labels = np.zeros(string_labels.shape)
labels[np.where(string_labels == 'D')] = -1 # Democrats
labels[np.where(string_labels == 'I')] = 0  # Independent
labels[np.where(string_labels == 'R')] = 1  # Republicans
labels
```

## 1 - Graph Fourier Transform

In this section we will observe how your feature/label vector looks like in the "Graph Fourier" domain.

### Question 1

Compute the Fourier basis vectors and the Laplacian eigenvalues. Make sure to order those from smaller to larger, $\lambda_0 \leq \lambda_1 \leq \dots \leq \lambda_{N-1},$ and use the same ordering for the Fourier basis vectors.

```python
# e: Ordered Laplacian eigenvalues
# U: Ordered graph Fourier basis
e, U = np.linalg.eigh(laplacian)
```

The Eigenvalues are already sorted in ascending order:

```python
plt.figure(figsize=(20,5))

(markerline, stemlines, baseline)=plt.stem(np.arange(len(e)),e);
plt.setp(baseline, visible=False)
plt.ylim((0,1.2))
plt.xlim((-1,98))
plt.ylabel('$\lambda$');
plt.xlabel('Eigenvalue index');
```

Plot the first 3 and the last Fourier basis vectors as signals on your graph. Clearly indicate which plot belongs to which basis vector.

```python
# Using the generalized second and fourth Eigenvectors as a two-dimensional embedding
D_norm = np.diag(np.sum(adjacency, 1)**(-1/2))
network_emb = D_norm @ U[:,[1,3]]
emb_x = network_emb[:,0]
emb_y = network_emb[:,1]
```

We plot the first four and two last Fourier basis vectors instead to see some more variation:

```python
fig, ax = plt.subplots(3, 2, figsize=(15,17))

fourier_bases = [0,1,2,3,96,97]

vmax = max(-U[:,fourier_bases].min(), U[:,fourier_bases].max())
vmin = -vmax

for ax_idx, fourier_basis in enumerate(fourier_bases):
    ax_x, ax_y = ax_idx//2, ax_idx%2
    im = ax[ax_x, ax_y].scatter(emb_x, emb_y, c=U[:,fourier_basis], cmap='bwr', s=70,
                                edgecolors='black', vmin=vmin, vmax=vmax)
    ax[ax_x, ax_y].set_title('Signal = Fourier basis {}'.format(fourier_basis))
    ax[ax_x, ax_y].set_xlabel('Generalized eigenvector embedding $U_1$')
    ax[ax_x, ax_y].set_ylabel('Generalized eigenvector embedding $U_3$')

fig.subplots_adjust(right=0.9)
cbar_ax = fig.add_axes([0.925, 0.15, 0.025, 0.7])
fig.colorbar(im, cax=cbar_ax)
```

In addition to our own embedding using two chosen Eigenvectors, we would also like to display all graphs using *NetworkX*'s Force-directed layout.
This finds the party separation between the two main clusters and arranges them nicely (with one exception).

Also, in this representation we can plot all the edges without distracting too much from the nodes.

```python
graph = nx.from_numpy_matrix(adjacency)
coords = nx.spring_layout(graph) # Force-directed layout.
```

```python
fig, ax = plt.subplots(3, 2, figsize=(15,15))

fourier_bases = [0,1,2,3,96,97]

vmax = max(-U[:,fourier_bases].min(), U[:,fourier_bases].max())
vmin = -vmax

for ax_idx, fourier_basis in enumerate(fourier_bases):
    ax_x, ax_y = ax_idx//2, ax_idx%2
    im = nx.draw_networkx_nodes(graph, coords, node_size=60, node_color=U[:,fourier_basis],
                                cmap='bwr', edgecolors='black', ax=ax[ax_x, ax_y], vmin=vmin, vmax=vmax)
    nx.draw_networkx_edges(graph, coords, alpha=0.2, ax=ax[ax_x, ax_y])
    ax[ax_x, ax_y].set_title('Signal = Fourier basis {}'.format(fourier_basis))

fig.subplots_adjust(right=0.9)
cbar_ax = fig.add_axes([0.925, 0.15, 0.025, 0.7])
fig.colorbar(im, cax=cbar_ax)
```

### Question 2

What can you observe in terms of local variations when comparing the basis vectors corresponding to the smallest eigenvalues to those corresponding to the largest eigenvalue? How would this justify the interpretation of the eigenvalues as "graph frequencies"?

**Your answer here.**

For the signals corresponding to the smallest eigenvalues, we can observe them having very low frequencies over the network while still having very large weights.
The basis vectors of larger eigenvalues contain mostly lower amplitudes and slightly higher frequencies. The one outlier node has a large influence over some of the Fourier bases, like basis 2.

### Question 3

Implement a function that returns the Graph Fourier Transform (GFT) of a given vector $x \in \mathbb{R}^{N},$ with respect to your graph, and a function that computes the corresponding inverse GFT (iGFT).

```python
def GFT(x):
    return U.T @ x

def iGFT(x):
    return U @ x
```

### Question 4

Plot your feature/label vector as a signal on your graph

```python
def coplot_network_signal(signal, title='Signal = ...'):
    '''
    Plots a signal on a graph using both a Laplacian embedding and the NetworkX force-directed layout.

    Args:
        signal: The signal of each node to plot on the graph
        title: Plot title
    '''
    fig, ax = plt.subplots(1, 2, figsize=(16,7))

    vmax = max(-np.nanmin(signal), np.nanmax(signal))
    vmin = -vmax

    im = ax[0].scatter(emb_x, emb_y, c=signal, cmap='bwr', s=70, edgecolors='black',
                       vmin=vmin, vmax=vmax)
    ax[0].set_title('Laplacian Embedding')
    ax[0].set_xlabel('Generalized eigenvector embedding $U_1$')
    ax[0].set_ylabel('Generalized eigenvector embedding $U_3$')

    nx.draw_networkx_nodes(graph, coords, node_size=60, node_color=signal, cmap='bwr',
                           edgecolors='black', ax=ax[1], vmin=vmin, vmax=vmax)
    nx.draw_networkx_edges(graph, coords, alpha=0.2, ax=ax[1])
    ax[1].set_title('NetworkX Force-directed layout')

    fig.suptitle(title, fontsize=16)

    fig.subplots_adjust(right=0.9)
    cbar_ax = fig.add_axes([0.925, 0.15, 0.025, 0.7])
    fig.colorbar(im, cax=cbar_ax)
```

```python
coplot_network_signal(labels, title='Signal = Ground truth labels')
```

Plot the absolute values of the GFT of your feature/label signal as a function of the graph eigenvalues. Make sure to add a marker indicating the position of each graph eigenvalue, and remember to properly name the axes.

*We plot the signal with respect to the (normalized) Eigenvalue frequency as well as the respective index.*

```python
plt.figure(figsize=(20,8))

(markerline, stemlines, baseline)=plt.stem(np.arange(len(e)),abs(GFT(labels)));
plt.setp(baseline, visible=False)
plt.ylim((0,10))
plt.xlim((-1,98))
plt.title('Graph Fourier Transform of the real label signal')
plt.ylabel('$|U^T x|$');
plt.xlabel('Eigenvalue index');
```

```python
plt.figure(figsize=(20,8))

(markerline, stemlines, baseline)=plt.stem(e,abs(GFT(labels)));
plt.setp(baseline, visible=False)
plt.ylim((0,10))
plt.title('Graph Fourier Transform of the real label signal')
plt.ylabel('$|U^T x|$');
plt.xlabel('Eigenvalue frequency');
```

### Question 5

Discuss the behavior of the GFT that you plotted in the last question via comparing the plot of your label signal and those of the Fourier basis of Question 1. Would you consider your labels a "low-pass" or "high-pass" signal, or yet something else entirely?

**Your answer here.**

We observe one very strong value corresponding to eigenvalue 1 with several smaller peaks located in the lower and upper 20 eigenvalues.
This main peak represents a single, very smooth transition over the whole network between our labels of -1 and 1.
This is coherent with the underlying assumption that Democrats and Republicans form two internally similar but opposing voting patterns.
Indeed, comparing the above plot of our ground truth labels with the plot of Fourier basis 1 shows them being very similar.

We would consider our labels as a "low-pass" signal, as most mass is concentrated in the very low frequencies, with some noise present in the upper end of the spectrum.

## 2 - Filtering on graphs

In this section we will check how filtered Dirac impulses diffuse on your graph.

### Question 6

Implement the following three filter kernels and the graph filtering operation.

- The **heat kernel** is supposed to take in a vector of eigenvalues `e` and a parameter `t` and output a vector of evaluations of the heat kernel at those eigenvalues (see the course slides for help).
- The **inverse filter** kernel is supposed to take in a vector of eigenvalues `e` and a parameter `t` and implement spectrally the  filter defined in the node domain by $f_{out}  = (I + t L)^{-1} f_{in},$ where $f_{in}, f_{out} \in \mathbb{R}^{N}$ are, repectively, the input and output signals to the filter.
- The **rectangle kernel** takes in a vector of eigenvalues `e` and parameters `l_min` and `l_max` and returns `1.0` at coordinates satisfying $(e[l] \geq l_{min}) \wedge (e[l] \leq l_{max}),$ and `0.0` otherwise.
- The **graph filtering** operation takes a graph signal $x \in \mathbb{R}^{N}$, a spectral graph `kernel` and a set of keyworded variables, and returns the corresponding filtered signal.
    - _Hint:_ Remember that you have implemented the `GFT` and `iGFT` operations in Question 3.
    - The `**kwargs` is a placeholder to collect supplementary pairs of keyword-values that are not known by the implementation before execution time.
      The `kwargs` variable is a dictionary whose keyes and values are the parameter names and values.
      This is useful to allow both `graph_filter(x, heat_kernel, tau=1.0)` and `graph_filter(x, rectangle_kernel, lambda_min=0.0, lambda_max=1.0)` to be valid calls from the same implementation.
      One can then defer the keyword-value assignment to the `kernel` call: `foo = kernel(bar, **kwargs)`.

```python
def heat_kernel(e, t):
    return np.exp(-t * e)

def inverse_kernel(e, t):
    return 1/(1 + t*e)

def rectangle_kernel(e, l_min, l_max):
    return ((e >= l_min) & (e <= l_max)).astype(float)

def graph_filter(x, kernel, **kwargs):
    return iGFT(kernel(e, **kwargs) * GFT(x))
```

### Question 7

Plot all three filter kernels in the spectral domain. Remember to properly name the axes and title the plots. Choose filter parameters that best approximate the behavior of the GFT of your feature/label signal (as seen in Question 4).

*We plot the signal with respect to the (normalized) Eigenvalue frequency as well as the respective index.*

```python
fig, ax = plt.subplots(3, 1, figsize=(20,16))

kernels = {
    'Heat kernel': heat_kernel(e,5),
    'Inverse kernel': inverse_kernel(e,10),
    'Rectangular kernel': rectangle_kernel(e,0.1,0.8)
}

for idx, (kernel_name, kernel) in enumerate(kernels.items()):
    (markerline, stemlines, baseline)=ax[idx].stem(np.arange(len(e)), kernel);
    plt.setp(baseline, visible=False)
    ax[idx].set_ylim((0,1.1))
    ax[idx].set_xlim((-1,98))
    ax[idx].set_title(kernel_name)
    ax[idx].set_ylabel('Kernel strength');
    ax[idx].set_xlabel('Eigenvalue index');
```

```python
fig, ax = plt.subplots(3, 1, figsize=(20,16))

kernels = {
    'Heat kernel': heat_kernel(e,5),
    'Inverse kernel': inverse_kernel(e,10),
    'Rectangular kernel': rectangle_kernel(e,0.1,0.8)
}

for idx, (kernel_name, kernel) in enumerate(kernels.items()):
    (markerline, stemlines, baseline)=ax[idx].stem(e, kernel);
    plt.setp(baseline, visible=False)
    ax[idx].set_ylim((0,1.1))
    ax[idx].set_title(kernel_name)
    ax[idx].set_ylabel('Kernel strength');
    ax[idx].set_xlabel('Eigenvalue frequency');
```

We chose the filter parameters such that the second peak gets preserved well, while larger frequencies are attenuated. With the rectangle kernel, we can very easily select the desired frequencies, while with the other two we have to find some equilibrium between the strength of the desired signal and the rest. In case of the heat kernel, we could nicely isolate the second peak while not having to worry about the first one, as it models a constant signal over the network. In the inverse kernel, we still have a tail with a lot of weight.

When filtering, we multiply this filter kernel by the Fourier spectrum of a signal, whereby the low frequencies are kept and the larger frequencies are discarded.

### Question 8

Consider two Dirac impulses arbitrarily placed on your graph. Plot their filtered versions by the three filter kernels implemented in Question 6.

```python
def plot_filtered_diracs(dirac_dict, kernel_name, kernel, **kwargs):
    '''
    Plots the filtered signal of two randomly chosen dirac spikes on a graph.
    Plots using both Laplacian embedding and NetworkX Force-directed layout.

    Args:
        dirac_dict: Dictionary specifying index and value of diracs. Eg: {30: 1, 60: -1}
        kernel_name: Used in the title to indicate used kernel type
        kernel: Kernel function to apply to diracs
        kwargs: Additional arguments for specific kernel
    '''
    diracs = np.zeros(n_nodes)
    dirac_idxs = list(dirac_dict.keys())
    for idx, dirac_value in dirac_dict.items():
        diracs[idx] = dirac_value

    filtered = graph_filter(diracs, kernel, **kwargs)

    vmax = max(-filtered.min(), filtered.max())
    vmin = -vmax

    # Plot using Laplacian embedding
    fig, ax = plt.subplots(1, 2, figsize=(16,7))
    # Plotting all nodes
    ax[0].scatter(emb_x, emb_y, c=filtered, cmap='PiYG', s=70,
                  edgecolors='black', vmin=vmin, vmax=vmax)
    # Plotting the dirac locations
    im = ax[0].scatter(emb_x[dirac_idxs], emb_y[dirac_idxs], c=diracs[dirac_idxs],
                       cmap='PiYG', s=300, edgecolors='black', vmin=vmin, vmax=vmax, marker='x')
    ax[0].set_title('Laplacian Embedding')
    ax[0].set_xlabel('Generalized eigenvector embedding $U_1$')
    ax[0].set_ylabel('Generalized eigenvector embedding $U_3$')

    # Plot using NetworkX
    nx.draw_networkx_nodes(graph, coords, node_size=60, node_color=filtered,
                           cmap='PiYG', edgecolors='black', ax=ax[1], vmin=vmin, vmax=vmax)
    nx.draw_networkx_edges(graph, coords, alpha=0.2, ax=ax[1])
    # Plotting the dirac locations
    d1_coords = coords[dirac_idxs[0]]
    d2_coords = coords[dirac_idxs[1]]
    im = ax[1].scatter(d1_coords[0], d1_coords[1], c=diracs[dirac_idxs[0]],
                       cmap='PiYG', s=300, edgecolors='black', vmin=vmin, vmax=vmax, marker='x')
    im = ax[1].scatter(d2_coords[0], d2_coords[1], c=diracs[dirac_idxs[1]],
                       cmap='PiYG', s=300, edgecolors='black', vmin=vmin, vmax=vmax, marker='x')
    ax[1].set_title('NetworkX Force-directed layout')

    fig.suptitle('Signal = Two diracs filtered using {}'.format(kernel_name), fontsize=16)

    fig.subplots_adjust(right=0.9)
    cbar_ax = fig.add_axes([0.925, 0.15, 0.025, 0.7])
    fig.colorbar(im, cax=cbar_ax)
```

In the following plots, as our graph consists of two main clusters, we will plot the following possible combinations of placing two dirac spikes:

- 8.1: Diracs of opposite signs in opposite clusters
- 8.2: Diracs of opposite signs in same cluster
- 8.3: Diracs of equal signs in opposite clusters
- 8.4: Diracs of equal signs in same cluster

For reproducibility reasons we will hereby not initialize the nodes randomly, but choose them arbitrarily ourself. We chose nodes 15, 30 and 60 for that purpose.

*We plot the bellow graphs using a different color palette, as not to conflict the colours with any party associations in this exercise.*

#### 8.1 Diracs of opposite signs in opposite clusters

```python
plot_filtered_diracs({30: -1, 60: 1}, 'Heat Kernel', heat_kernel, t=5)
```

```python
plot_filtered_diracs({30: -1, 60: 1}, 'Inverse Kernel', inverse_kernel, t=10)
```

```python
plot_filtered_diracs({30: -1, 60: 1}, 'Rectangle Kernel', rectangle_kernel, l_min=0.1, l_max=0.8)
```

#### 8.2 Diracs of opposite signs in same clusters

```python
plot_filtered_diracs({30: -1, 15: 1}, 'Heat Kernel', heat_kernel, t=5)
```

```python
plot_filtered_diracs({30: -1, 15: 1}, 'Inverse Kernel', inverse_kernel, t=10)
```

```python
plot_filtered_diracs({30: -1, 15: 1}, 'Rectangle Kernel', rectangle_kernel, l_min=0.1, l_max=0.8)
```

#### 8.3 Diracs of equal signs in opposite clusters

```python
plot_filtered_diracs({30: 1, 60: 1}, 'Heat Kernel', heat_kernel, t=5)
```

```python
plot_filtered_diracs({30: 1, 60: 1}, 'Inverse Kernel', inverse_kernel, t=10)
```

```python
plot_filtered_diracs({30: 1, 60: 1}, 'Rectangle Kernel', rectangle_kernel, l_min=0.1, l_max=0.8)
```

#### 8.4 Diracs of equal signs in same cluster

```python
plot_filtered_diracs({30: 1, 15: 1}, 'Heat Kernel', heat_kernel, t=5)
```

```python
plot_filtered_diracs({30: 1, 15: 1}, 'Inverse Kernel', inverse_kernel, t=10)
```

```python
plot_filtered_diracs({30: 1, 15: 1}, 'Rectangle Kernel', rectangle_kernel, l_min=0.1, l_max=0.8)
```

Comment on the "diffusion" of the Diracs induced by the filters. What does it say about the "communication" of information across your network? Relate that to the network connectivity measures that you analyzed during the previous milestones.

**Your answer here.**

*8.1: Diracs of opposite signs in opposite clusters*

Both communities get nicely coloured in their respective sign. The inverse kernel diffusion results in a very weak signal across the network, but separation can be faintly visible.

*8.2: Diracs of opposite signs in same cluster*

The opposite spikes difused equally over the clusters, cancelling each other out in the process. The rectangular kernel on the other hand inhibits a particularly strong separation. This is the case because we are only preserving the signal belonging to a divide with the lowest frequency over the cluster.

*8.3: Diracs of equal signs in opposite clusters*

In this case we see an equal diffusion over the whole network. An exception is again made by the rectangular kernel, as seen in 8.2.

*8.4: Diracs of equal signs in same cluster*

The spikes are mainly difused into the cluster where they originated. The resulting signal is stronger in those clusters and weaker in the others.

The above results are in line with our findings that communication within clusters is very high (i.e. distances are small), while distances between the two clusters are higher and thus communication lower.

## 3 - De-noising

In this section we will add some centered Gaussian noise to your feature/label signal and attempt to recover it.

### Question 9

In the cell below, set the noise variance $\sigma^2$ by making sure that the signal-to-noise ratio $SNR = \frac{\operatorname{Var}(\text{labels})}{\sigma^2}$ is about  $1.5$.

_Note:_ Actually, you might want to play with the noise variance here and set it to different values and see how the denoising filters behave.

```python
noise_variance = labels.std()**2 / 1.5
noisy_measurements = labels + noise_variance * np.random.randn(n_nodes)
```

### Question 10

In the denoising setting, a common graph signal processing assumption is that the signal $z$ that we want to recover is "smooth", in the sense that $\|\nabla_G z\|_2 = \sqrt{z^{\top} L z}$ is small, while remaining "close" to the measurements that we start with. This leads to denoising by solving the following optimization problem:

$$
z^\star = \text{arg} \, \underset{z \in \mathbb{R}^{N}}{\min} \, \|z - y\|_2^2 + \gamma z^{\top} L z,
$$

where $y \in \mathbb{R}^{N}$ is the vector of noisy measurements.

Derive the close form solution to this problem giving $z^\star$ as a function of $y$, $\gamma$ and $L$. Does this solution correspond to any graph filtering operation that you know?

**Your answer here.**

To find the optimum, we derive the above term (let's denote it as $g$) with respect to the vector $z$:

$$ \nabla_z g(z,y,\gamma) = \nabla_z \|z - y\|_2^2 + \gamma z^{\top} L z = 2(z-y) + 2 \gamma L $$

Setting $ \nabla_z g(z,y,\gamma) = 0 $ to find our minimum:

\begin{align}
0 &= 2(z-y) + 2 \gamma L \\
  &= 2 (z(1+\gamma L) - y)
\end{align}

Solving for $z$ we get:

$$ z^* = \frac{y}{1 + \gamma L} $$

We have to check that $z^*$ is indeed our global minima. The function is convex, since it is a linear combination of other convex functions.
Furthermore, $z^*$ is a minima, as the second derivative is positive semi-definite:

$$ \nabla_z^2 g(z,y,\gamma) = \nabla_z^2 \|z - y\|_2^2 + \gamma z^{\top} L z = 2 + 2 \gamma L $$

The solution for $z^*$ corresponds to the inverse kernel filter from above, that models $f_{out} = (I + tL)^{-1} f_{in}$ in the node domain.

### Question 11

Now, denoise the noisy measurements by passing them through the filters that you implemented in Question 6. Choose the filter parameters based on the behavior of the GFT of your original label signal (this is the prior knowledge that you input to the problem).

Plot, on your graph, the original label signal, the noisy measurements, and the three denoised version obtained above. Report on each plot the value of the corresponding relative error
$$
\text{rel-err} = \frac{\|\text{labels} - z \|_2}{\|\text{labels}\|_2},
$$
where $z$ is the plotted signal.

```python
def rel_err(labels, z):
    '''
    Calculates the relative error between the true labels and an estimate z

    Args:
        labels: Ground truth signal
        z: Estimated signal
    '''
    return np.linalg.norm(labels - z, 2) / np.linalg.norm(labels, 2)
```

```python
coplot_network_signal(labels, title='Signal = Ground truth labels')
print('Relative Error: {:.2f}'.format(rel_err(labels, labels)))
```

```python
coplot_network_signal(noisy_measurements, title='Signal = Noisy measurements')
print('Relative Error: {:.2f}'.format(rel_err(labels, noisy_measurements)))
```

```python
z_heat_denoised = graph_filter(noisy_measurements, heat_kernel, t=5)
coplot_network_signal(z_heat_denoised, title='Signal = Heat denoised measurements')
print('Relative Error: {:.2f}'.format(rel_err(labels, z_heat_denoised)))
```

```python
z_inv_denoised = graph_filter(noisy_measurements, inverse_kernel, t=5)
coplot_network_signal(z_inv_denoised, title='Signal = Inverse denoised measurements')
print('Relative Error: {:.2f}'.format(rel_err(labels, z_inv_denoised)))
```

```python
z_rect_denoised = graph_filter(noisy_measurements, rectangle_kernel, l_min=0.1, l_max=0.8)
coplot_network_signal(z_rect_denoised, title='Signal = Rectangle denoised measurements')
print('Relative Error: {:.2f}'.format(rel_err(labels, z_rect_denoised)))
```

Finally, overlay on the same plot the GFT of all five signals above.

```python
signals = {
    'Ground truth': labels,
    'Noisy measurements': noisy_measurements,
    'Heat filtered': z_heat_denoised,
    'Inverse filtered': z_inv_denoised,
    'Rectangle filtered': z_rect_denoised
}

plt.figure(figsize=(15,10))

for signal_name, signal in signals.items():
    plt.plot(signal, label=signal_name)

plt.xlabel('Node index')
plt.ylabel('Signal strength')
plt.title('Signals on sorted nodes')
plt.legend()
```

### Question 12

Comment on which denoised version seems to best match the original label signal. What is the underlying assumption behind the three filtering approaches? Do you think it holds for your label signal? Why?

*We plot the signal with respect to the (normalized) Eigenvalue frequency as well as the respective index.*

```python
plt.figure(figsize=(20,8))

(markerline, stemlines, baseline)=plt.stem(np.arange(len(e)),abs(GFT(noisy_measurements)));
plt.setp(baseline, visible=False)
plt.ylim((0,11))
plt.xlim((-1,98))
plt.title('Graph Fourier Transform of the noisy measurements')
plt.ylabel('$|U^T x|$');
plt.xlabel('Eigenvalue index');
```

```python
plt.figure(figsize=(20,8))

(markerline, stemlines, baseline)=plt.stem(e,abs(GFT(noisy_measurements)));
plt.setp(baseline, visible=False)
plt.ylim((0,11))
plt.title('Graph Fourier Transform of the noisy measurements')
plt.ylabel('$|U^T x|$');
plt.xlabel('Eigenvalue frequency');
```

**Your answer here.**

From the graph plots and the plot with all the signals overlayed, the rectangle filter seems to recover the underlying labels best, without losing much signal strength.
It also has the best relative error between the ground truth and denoised signal.

Since our graph has only two big clusters, all other above filters did not have much of a problem denoising the noisy measurements either. We could threshold by taking the sign function of each denoised signal and retrieve a very accurate reconstruction of our labels.

We tried to match the strength of our filters in the spectral domain to the strength of our Laplacian's eigenvalues, whereby the aim was to preserve as much as possible the frequencies of the very strong eigenvalue 1 and attenuate all others.
As seen in our plots from Question 1, the eigenvector associated with eigenvalue 1 presents a smooth transition between the two main parties, accurately modeling their respective clusters.

The underlying assumption is then, that if we design a filter that only keeps the preselected frequencies (and weakens or completely removes all others) over the network that we selected based on the ground truth, that noise is filtered out.
This again stands on the assumption, that the introduced noise has frequencies that can be well separated from the signal in the frequency domain.

For our signal, these assumptions hold, as we can see in the above plot of the GFT of our noisy signal. The main peak from eigenvalue 1 is still clearly visible and easy to filter out.

## 4 - Transductive learning

It is often the case in large networks that we can only afford to query properties/labels on a small subset of nodes. Nonetheless, if the underlying labels signal is "regular" enough, we might still be able to recover a good approximation of it by solving an offline variational problem, with constraints on the values of the measured nodes.

In this section, we will be interested in solving such transductive learning problems by minimizing a (semi-) p-norm of the graph gradient applied to the signal of interest:

$$
\text{arg} \, \underset{z|_S = y}{\min} \|\nabla_G z\|_p^p,
$$

where $S$ is the set of measured nodes.

In English, we can say that we are looking for solutions with small "aggregated local variations", as measured by $\|\nabla_G z\|_p^p = \sum_{i=1}^{n} \sum_{j=1}^{n} \left( \sqrt{W_{ij}} |z[i] - z[j]| \right)^p,$ while satisfying the measurement constraints $z[i] = y[i]$ for $i \in S.$

We will work with two cases, according to the choices $p=1$ or $p=2.$ For $p=1,$ the problem is known as "interpolation by graph total-variation minimization," whereas for $p=2$ it is sometimes called "interpolation by Tikhonov regularization".

In order to solve these variational problems with the black-box solver provided to you, you will use the [pyunlocbox](https://pyunlocbox.readthedocs.io). This toolbox implements iterative solvers based on so-called ["proximal-splitting"](https://en.wikipedia.org/wiki/Proximal_gradient_method) methods.

### Question 13

Throughout this section, we will consider only a binarized version of your label signal. If your variable `labels` currently has values other than $\{-1, 1\},$ threshold them so that those are the only values taken in this vector. This can be done for example by choosing a number $t \in \mathbb{R}$ and then setting $\text{labels_bin}[i] = 1$ if $\text{labels}[i] \geq t$ and $\text{labels_bin}[i] = 0$ otherwise.

```python
# Setting the independent party members to be Democrats
labels_bin = labels.copy()
labels_bin[labels_bin==0] = -1
```

Now, subsample this binarized label signal by $70\%$ by choosing, uniformly at random, $30\%$ of the nodes whose labels we will keep.

You will do this by computing a "measurement mask" vector `w` with `1.0`'s at the measured coordinates, and $0.0$'s otherwise.

```python
mn_ratio = 0.3
w = np.random.binomial(n=1, p=mn_ratio, size=n_nodes)
m = sum(w) # Number of measurements
print('Sampled {} out of {} nodes'.format(m, n_nodes))
```

Plot the subsampled signal on the graph. _Hint:_ you might want to set to `numpy.nan` the values of the un-measured nodes for a cleaner plot.

```python
subsampled = labels_bin*w

subsampled_nan = subsampled.copy()
subsampled_nan[subsampled_nan==0] = np.nan
subsampled_nan
```

```python
coplot_network_signal(subsampled_nan, title='Signal = Subsampled labels')
```

### Interlude

For the solution of the variational problems you can use the following function as a "black-box".

You will just need to provide a `gradient` matrix (which you should already have from Section 0), and an orthogonal projection operator `P` onto the span of the measured coordinates (made precise in the next question).

```python
def graph_pnorm_interpolation(gradient, P, w, x0=None, p=1., **kwargs):
    r"""
    Solve an interpolation problem via gradient p-norm minimization.

    A signal :math:`x` is estimated from its measurements :math:`y = A(x)` by solving
    :math:`\text{arg}\underset{z \in \mathbb{R}^n}{\min}
    \| \nabla_G z \|_p^p \text{ subject to } Az = y`
    via a primal-dual, forward-backward-forward algorithm.

    Parameters
    ----------
    gradient : array_like
        A matrix representing the graph gradient operator
    P : callable
        Orthogonal projection operator mapping points in :math:`z \in \mathbb{R}^n`
        onto the set satisfying :math:`A P(z) = A z`.
    x0 : array_like, optional
        Initial point of the iteration. Must be of dimension n.
        (Default is `numpy.random.randn(n)`)
    p : {1., 2.}
    kwargs :
        Additional solver parameters, such as maximum number of iterations
        (maxit), relative tolerance on the objective (rtol), and verbosity
        level (verbosity). See :func:`pyunlocbox.solvers.solve` for the full
        list of options.

    Returns
    -------
    x : array_like
        The solution to the optimization problem.

    """

    grad = lambda z: gradient.dot(z)
    div = lambda z: gradient.transpose().dot(z)

    # Indicator function of the set satisfying :math:`y = A(z)`
    f = functions.func()
    f._eval = lambda z: 0
    f._prox = lambda z, gamma: P(z, w)

    # :math:`\ell_1` norm of the dual variable :math:`d = \nabla_G z`
    g = functions.func()
    g._eval = lambda z: np.sum(np.abs(grad(z)))
    g._prox = lambda d, gamma: functions._soft_threshold(d, gamma)

    # :math:`\ell_2` norm of the gradient (for the smooth case)
    h = functions.norm_l2(A=grad, At=div)

    stepsize = (0.9 / (1. + scipy.sparse.linalg.norm(gradient, ord='fro'))) ** p

    solver = solvers.mlfbf(L=grad, Lt=div, step=stepsize)

    if p == 1.:
        problem = solvers.solve([f, g, functions.dummy()], x0=x0, solver=solver, **kwargs)
        return problem['sol']
    if p == 2.:
        problem = solvers.solve([f, functions.dummy(), h], x0=x0, solver=solver, **kwargs)
        return problem['sol']
    else:
        return x0
```

### Question 14

During the iterations of the algorithm used for solving the variational problem, we have to make sure that the labels at the measured nodes stay the same. We will do this by means of an operator `P` which, given a vector $a \in \mathbb{R}^{N},$ returns another vector $b \in \mathbb{R}^{N}$ satisfying $b[i] = \text{labels_bin}[i]$ for every node $i$ in the set $S$ of known labels, and $b[i] = a[i]$ otherwise. Write in the cell below the function for this orthogonal projection operator `P`.

_Hint:_ remember you have already computed the mask `w`.

```python
def P(a, w):
    mask_pos = np.where(w==1)
    b = a.copy()
    b[mask_pos] = labels_bin[mask_pos]
    return b
```

### Question 15

Solve the variational problems for $p = 1$ and $p = 2$. Record the solution for the $1-$norm minimization under `sol_1norm_min` and the one for $2-$norm minimization under `sol_2norm_min`.

Compute also binarized versions of these solutions by thresholding the values with respect to $0$, that is, non-negative values become `1.0`, while negative values become `-1.0`. Store those binarized versions under `sol_1norm_bin` and `sol_2norm_bin`, respectively.

```python
sol_1norm_min = graph_pnorm_interpolation(sparse.csr_matrix(S).T, P, w, x0=subsampled.copy(), p=1.)

sol_2norm_min = graph_pnorm_interpolation(sparse.csr_matrix(S).T, P, w, x0=subsampled.copy(), p=2.)

threshold = 0

sol_1norm_bin = (sol_1norm_min > threshold).astype(int)
sol_1norm_bin[sol_1norm_bin==0] = -1

sol_2norm_bin = (sol_2norm_min > threshold).astype(int)
sol_2norm_bin[sol_2norm_bin==0] = -1
```

Plot, on your graph, the original `labels_bin` signal, as well as the solutions to the variational problems (both binarized and otherwise). Indicate on each plot the value of the relative error $\text{rel-err} = \frac{\|\text{labels_bin} - z\|_2}{\|\text{labels_bin}\|_2}$, where $z$ is the signal in the corresponding plot.

```python
coplot_network_signal(labels_bin, title='Signal = Ground truth labels')
print('Relative Error: {:.2f}'.format(rel_err(labels_bin, labels_bin)))
```

```python
coplot_network_signal(sol_1norm_min, title='Signal = Variational problem solution for p=1')
print('Relative Error: {:.2f}'.format(rel_err(labels_bin, sol_1norm_min)))
```

```python
coplot_network_signal(sol_1norm_bin, title='Signal = Binarized variational problem solution for p=1')
print('Relative Error: {:.2f}'.format(rel_err(labels_bin, sol_1norm_bin)))
```

```python
coplot_network_signal(sol_2norm_min, title='Signal = Variational problem solution for p=2')
print('Relative Error: {:.2f}'.format(rel_err(labels_bin, sol_2norm_min)))
```

```python
coplot_network_signal(sol_2norm_bin, title='Signal = Binarized variational problem solution for p=2')
print('Relative Error: {:.2f}'.format(rel_err(labels_bin, sol_2norm_bin)))
```

### Question 16

Now that you have got a feeling for the sort of solutions that the transductive learning problems studied can give, we will see what is the effect of the number of measurements on the accuracy of both $p-$norm minimization problems.

Towards this goal, you will write a `phase_transition()` function. This function will basically go over all the procedures that you have implemented in this section, but for varying numbers of measurements and thresholding values. It will also compute the relative error, $\text{rel-err},$ of the solutions and average them over a number of trials.

The output of the `phase_transition()` function has to be a matrix with `len(mn_ratios)` columns and `len(thresholds)` rows. Each pixel $(i,j)$ in the output matrix has to contain the average, over `n_trials` trials, of the relative error $\text{rel-err}$ in the binarized (with threshold `thresholds[i]`) solution given by `graph_pnorm_interpolation()` from observing an `mn_ratios[j]` fraction of nodes. The randomness comes from a different choice of mask `w` at each trial, hence the averaging.

The interest of this phase transition matrix is to assess what level of recovery error one could expect for a certain fraction of measurements and a certain threshold level.

```python
def phase_transition(mn_ratios, thresholds, n_trials, labels_bin, p, **kwargs):
    pt_matrix = np.zeros((len(thresholds), len(mn_ratios)))

    for i, threshold in enumerate(thresholds):
        for j, ratio in enumerate(mn_ratios):
            # Simulate n_trials times
            for trial in range(n_trials):
                # Subsample randomly
                w = np.random.binomial(n=1, p=ratio, size=n_nodes)
                subsampled = labels_bin * w
                # Solve p-norm interpolation
                sol = graph_pnorm_interpolation(sparse.csr_matrix(S).T, P, w,
                                                x0=subsampled.copy(), p=p, **kwargs)
                # Threshold to -1 and 1
                sol_bin = (sol > threshold).astype(int)
                sol_bin[sol_bin==0] = -1
                # Compute and store the error
                pt_matrix[i,j] += rel_err(labels_bin, sol_bin)

    # Computing the mean of all trials
    pt_matrix /= n_trials

    return pt_matrix
```

### Question 17

Pick 5 "m/n" ratios in $(0, 1)$ and 5 threshold levels in $(-1, 1)$ and run the `phase_transition()` function with `n_trials` = 20, for both $p = 1$ and $p = 2$.

```python
N = 11 # For more detail in the plots below, but longer computation time!
n_trials = 20
#mn_ratios = np.linspace(0.05, 0.9, N) # To space linearly
mn_ratios = np.geomspace(0.01, 0.9, N) # Small ratios are more interesting
thresholds = np.linspace(-0.9, 0.9, N)

pt_matrix_1norm = phase_transition(mn_ratios, thresholds, n_trials, labels_bin, 1, verbosity='NONE')
pt_matrix_2norm = phase_transition(mn_ratios, thresholds, n_trials, labels_bin, 2, verbosity='NONE')
```

Plot both phase transition matrices as images with a colorbar. Make sure to properly name the axes and title the images.

```python
ratio_strings = ['{:.2f}'.format(ratio) for ratio in mn_ratios]
thresh_strings = ['{:.2f}'.format(thresh) for thresh in thresholds]

vmax = max(pt_matrix_1norm.max(), pt_matrix_2norm.max())
vmin = min(pt_matrix_1norm.min(), pt_matrix_2norm.min())

fig, ax = plt.subplots(1,2,figsize=(16,8))

im = ax[0].imshow(pt_matrix_1norm, cmap='inferno', vmin=vmin, vmax=vmax)
ax[0].set_title('1-Norm')
ax[0].set_xlabel('m/n ratio')
ax[0].set_ylabel('Threshold')
ax[0].set_xticks(np.arange(N))
ax[0].set_yticks(np.arange(N))
ax[0].set_xticklabels(ratio_strings)
ax[0].set_yticklabels(thresh_strings)

im = ax[1].imshow(pt_matrix_2norm, cmap='inferno', vmin=vmin, vmax=vmax)
ax[1].set_title('2-Norm')
ax[1].set_xlabel('m/n ratio')
ax[1].set_ylabel('Threshold')
ax[1].set_xticks(np.arange(N))
ax[1].set_yticks(np.arange(N))
ax[1].set_xticklabels(ratio_strings)
ax[1].set_yticklabels(thresh_strings)

fig.suptitle('Phase transition matrices (Darker is better)', fontsize=16)

fig.subplots_adjust(right=0.9)
cbar_ax = fig.add_axes([0.925, 0.15, 0.025, 0.7])
fig.colorbar(im, cax=cbar_ax)

plt.show()
```

### Question 18

Do the phase transition plots above provide any justification for choosing one $p-$norm interpolation over the other? Why?

**Your answer here.**

When picking a threshold of 0, the 2-Norm yields very good results up to very small ratios. If we choose different thresholds, the results quickly degrade. This is because the 2-Norm tends to produce estimates between the smallest and largest inputs -1 and 1, which makes chosing a good threshold important. As we have two balanced clusters, choosing the threshold to be zero is the most reasonable choice.

For the 1-Norm, results are almost entirely independent of the threshold we choose. Indeed, the estimated values tend to be bellow or above the smallest and largest inputs -1 and 1, making the choice of a threshold between -1 and 1 futile. Performance is only slightly worse than using the 2-Norm with a carefully selected threshold.

If there is an easy and justified way to pick a good threshold, taking the 2-Norm might produce better results. If this is not the case, one can get good results across the board with the 1-Norm using an almost arbitrary threshold.

```python

```