```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, DecisionTreeRegressor
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import StratifiedKFold
from torch.utils.data import TensorDataset, DataLoader
import time
from sklearn.metrics import accuracy_score
from sympy import simplify_logic
import re

import lens
from lens.utils.base import validate_network, set_seed, tree_to_formula
from lens.utils.relu_nn import get_reduced_model, prune_features
from lens.utils.psi_nn import prune_equal_fanin
from lens.models.brl import XBRLClassifier
from lens import logic

set_seed(0)

#%%

data = pd.read_csv('data/vdem/V-Dem-CY-Core-v10.csv')
```

```python
data['country_name_year'] = data['country_name'] + '_' + data['year'].astype(str)
data['country_name_year']
```

    0                   Mexico_1789
    1                   Mexico_1790
    2                   Mexico_1791
    3                   Mexico_1792
    4                   Mexico_1793
                      ...
    27008    Piedmont-Sardinia_1857
    27009    Piedmont-Sardinia_1858
    27010    Piedmont-Sardinia_1859
    27011    Piedmont-Sardinia_1860
    27012    Piedmont-Sardinia_1861
    Name: country_name_year, Length: 27013, dtype: object

```python
data_2000 = data[data['year']>2000].iloc[:, 12:-1].dropna(axis=1)
data_2000
```

      v2x_polyarchy
      v2x_polyarchy_codelow
      v2x_polyarchy_codehigh
      v2x_polyarchy_sd
      v2x_delibdem
      v2x_delibdem_codelow
      v2x_delibdem_codehigh
      v2x_delibdem_sd
      v2x_egaldem
      v2x_egaldem_codelow
      ...
      v2xcl_slave_sd
      v2xel_elecparl
      v2xel_elecpres
      v2xex_elecleg
      v2xlg_leginter
      v2xme_altinf
      v2xme_altinf_codelow
      v2xme_altinf_codehigh
      v2xme_altinf_sd
      v2x_divparctrl

      212
      0.698
      0.632
      0.763
      0.067
      0.539
      0.467
      0.626
      0.080
      0.411
      0.349
      ...
      0.556
      0
      0
      1.00
      0.0
      0.767
      0.615
      0.855
      0.653
      1.616

      213
      0.711
      0.667
      0.778
      0.055
      0.543
      0.475
      0.608
      0.064
      0.418
      0.359
      ...
      0.556
      0
      0
      1.00
      0.0
      0.808
      0.683
      0.887
      0.649
      1.616

      214
      0.715
      0.668
      0.762
      0.047
      0.554
      0.497
      0.606
      0.053
      0.418
      0.378
      ...
      0.556
      1
      0
      1.00
      0.0
      0.808
      0.683
      0.887
      0.649
      1.616

      215
      0.720
      0.668
      0.746
      0.039
      0.566
      0.519
      0.605
      0.043
      0.423
      0.386
      ...
      0.556
      0
      0
      1.00
      0.0
      0.808
      0.683
      0.887
      0.649
      1.616

      216
      0.720
      0.668
      0.746
      0.039
      0.585
      0.522
      0.615
      0.047
      0.423
      0.386
      ...
      0.556
      0
      0
      1.00
      0.0
      0.808
      0.683
      0.887
      0.649
      1.616

      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...

      25618
      0.316
      0.283
      0.348
      0.035
      0.215
      0.173
      0.278
      0.052
      0.255
      0.214
      ...
      0.592
      1
      1
      0.85
      0.0
      0.516
      0.381
      0.627
      0.633
      1.602

      25619
      0.267
      0.236
      0.286
      0.028
      0.194
      0.152
      0.233
      0.041
      0.234
      0.207
      ...
      0.592
      1
      1
      0.85
      0.0
      0.516
      0.381
      0.627
      0.633
      -0.684

      25620
      0.260
      0.236
      0.282
      0.025
      0.168
      0.137
      0.208
      0.034
      0.213
      0.188
      ...
      0.599
      0
      0
      0.85
      0.0
      0.437
      0.321
      0.579
      0.628
      -0.448

      25621
      0.262
      0.245
      0.284
      0.020
      0.213
      0.181
      0.264
      0.041
      0.214
      0.192
      ...
      0.608
      0
      0
      0.85
      0.0
      0.502
      0.374
      0.620
      0.625
      -0.586

      25622
      0.245
      0.209
      0.265
      0.029
      0.183
      0.133
      0.224
      0.044
      0.187
      0.160
      ...
      0.609
      0
      0
      0.85
      0.0
      0.525
      0.378
      0.663
      0.640
      -0.586

3385 rows × 1052 columns

```python
high_level_indicators = [
    'v2x_polyarchy',
    #'v2x_libdem',
    #'v2x_partipdem',
    'v2x_delibdem',
    'v2x_egaldem'
]
mid_level_indicators = [
    'v2x_api',
    'v2x_mpi',
    'v2x_freexp_altinf',
    'v2x_frassoc_thick',
    'v2x_suffr',
    'v2xel_frefair',
    'v2x_elecoff',
    #'v2x_liberal',
    'v2xcl_rol',
    #'v2x_jucon',
    #'v2xlg_legcon',
    #'v2x_partip',
    'v2x_cspart',
    #'v2xdd_dd',
    #'v2xel_locelec',
    #'v2xel_regelec',
    'v2xdl_delib',
    'v2x_egal',
    'v2xeg_eqprotec',
    'v2xeg_eqaccess',
    'v2xeg_eqdr',
]

drop_list = ['codelow', 'codehigh', 'sd', 'osp', 'nr', 'mean']
low_level_indicators = []
for f in data_2000.columns:
    if f.endswith('_ord') and f not in high_level_indicators and f not in mid_level_indicators:
        low_level_indicators.append(f)

low_level_indicators_continuous = []
for f in data_2000.columns:
    if f.endswith('_codehigh') or f.endswith('_codelow') and f not in high_level_indicators and f not in mid_level_indicators:
        low_level_indicators_continuous.append(f)

print(f'Main {len(high_level_indicators)} - Area {len(mid_level_indicators)} - Raw {len(low_level_indicators)}')
```

    Main 3 - Area 14 - Raw 62

```python
data_low_continuous = data_2000[low_level_indicators_continuous]
data_low_continuous.shape
```

    (3385, 464)

```python
data_low_raw = data_2000[low_level_indicators]
one_hots = []
for indicator in low_level_indicators:
    c = data_low_raw[indicator].values
    n_bins = int(c.max())
    kbin = KBinsDiscretizer(n_bins=n_bins, encode='onehot-dense', strategy='uniform')
    c1h = kbin.fit_transform(c.reshape(-1, 1))
    one_hots.append(c1h)

new_indicator_names = []
for clist, cname in zip(one_hots, low_level_indicators):
    if clist.shape[1] > 1:
        for i in range(clist.shape[1]):
            new_indicator_names.append(f'{cname}_{i}')
    else:
        new_indicator_names.append(f'{cname}')

data_low = pd.DataFrame(np.hstack(one_hots), columns=new_indicator_names)
data_low.describe()
```

      v2psbars_ord_0
      v2psbars_ord_1
      v2psbars_ord_2
      v2psbars_ord_3
      v2psparban_ord_0
      v2psparban_ord_1
      v2psparban_ord_2
      v2psparban_ord_3
      v2psorgs_ord_0
      v2psorgs_ord_1
      ...
      v2pepwrgen_ord_1
      v2pepwrgen_ord_2
      v2peedueq_ord_0
      v2peedueq_ord_1
      v2peedueq_ord_2
      v2peedueq_ord_3
      v2pehealth_ord_0
      v2pehealth_ord_1
      v2pehealth_ord_2
      v2pehealth_ord_3

      count
      3385.000000
      3385.000000
      3385.000000
      3385.000000
      3385.000000
      3385.000000
      3385.000000
      3385.000000
      3385.000000
      3385.000000
      ...
      3385.000000
      3385.000000
      3385.000000
      3385.000000
      3385.000000
      3385.000000
      3385.000000
      3385.000000
      3385.000000
      3385.000000

      mean
      0.056425
      0.037518
      0.080059
      0.825997
      0.053471
      0.029838
      0.028951
      0.887740
      0.044018
      0.150960
      ...
      0.230133
      0.730871
      0.059970
      0.312851
      0.132939
      0.494239
      0.038109
      0.322304
      0.145643
      0.493944

      std
      0.230775
      0.190057
      0.271425
      0.379168
      0.225004
      0.170164
      0.167694
      0.315733
      0.205165
      0.358063
      ...
      0.420980
      0.443572
      0.237467
      0.463723
      0.339559
      0.500041
      0.191488
      0.467428
      0.352800
      0.500037

      min
      0.000000
      0.000000
      0.000000
      0.000000
      0.000000
      0.000000
      0.000000
      0.000000
      0.000000
      0.000000
      ...
      0.000000
      0.000000
      0.000000
      0.000000
      0.000000
      0.000000
      0.000000
      0.000000
      0.000000
      0.000000

      25%
      0.000000
      0.000000
      0.000000
      1.000000
      0.000000
      0.000000
      0.000000
      1.000000
      0.000000
      0.000000
      ...
      0.000000
      0.000000
      0.000000
      0.000000
      0.000000
      0.000000
      0.000000
      0.000000
      0.000000
      0.000000

      50%
      0.000000
      0.000000
      0.000000
      1.000000
      0.000000
      0.000000
      0.000000
      1.000000
      0.000000
      0.000000
      ...
      0.000000
      1.000000
      0.000000
      0.000000
      0.000000
      0.000000
      0.000000
      0.000000
      0.000000
      0.000000

      75%
      0.000000
      0.000000
      0.000000
      1.000000
      0.000000
      0.000000
      0.000000
      1.000000
      0.000000
      0.000000
      ...
      0.000000
      1.000000
      0.000000
      1.000000
      0.000000
      1.000000
      0.000000
      1.000000
      0.000000
      1.000000

      max
      1.000000
      1.000000
      1.000000
      1.000000
      1.000000
      1.000000
      1.000000
      1.000000
      1.000000
      1.000000
      ...
      1.000000
      1.000000
      1.000000
      1.000000
      1.000000
      1.000000
      1.000000
      1.000000
      1.000000
      1.000000

8 rows × 241 columns

```python
data_mid = data_2000[mid_level_indicators] > 0.5
data_mid.describe()
```

      v2x_api
      v2x_mpi
      v2x_freexp_altinf
      v2x_frassoc_thick
      v2x_suffr
      v2xel_frefair
      v2x_elecoff
      v2xcl_rol
      v2x_cspart
      v2xdl_delib
      v2x_egal
      v2xeg_eqprotec
      v2xeg_eqaccess
      v2xeg_eqdr

      count
      3385
      3385
      3385
      3385
      3385
      3385
      3385
      3385
      3385
      3385
      3385
      3385
      3385
      3385

      unique
      2
      2
      2
      2
      2
      2
      2
      2
      2
      2
      2
      2
      2
      2

      top
      True
      False
      True
      True
      True
      True
      True
      True
      True
      True
      True
      True
      True
      True

      freq
      2696
      2196
      2605
      2622
      3322
      1909
      2934
      2607
      2669
      2464
      2398
      2527
      2392
      2073

```python
data_high = data_2000[high_level_indicators].iloc[:, 0] > 0.5
data_high.describe()
```

    count     3385
    unique       2
    top       True
    freq      1799
    Name: v2x_polyarchy, dtype: object

```python
c0c = torch.FloatTensor(data_low_continuous.values)
c0 = torch.FloatTensor(data_low.values)
c1 = torch.FloatTensor(data_mid.values)
c2 = torch.FloatTensor(data_high.values)
print(c0c.shape)
print(c0.shape)
print(c1.shape)
print(c2.shape)
```

    torch.Size([3385, 464])
    torch.Size([3385, 241])
    torch.Size([3385, 14])
    torch.Size([3385])

# Low-to-mid level explanations

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

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
mid_concepts = list(data_mid.columns)
low_concepts = list(data_low.columns)
n_rep = 10
tot_epochs = 6001
prune_epochs = 3001

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

x = c0
xh = c1
y = c2

print(low_concepts[:5])
print(mid_concepts[:5])
```

    ['v2psbars_ord_0', 'v2psbars_ord_1', 'v2psbars_ord_2', 'v2psbars_ord_3', 'v2psparban_ord_0']
    ['v2x_api', 'v2x_mpi', 'v2x_freexp_altinf', 'v2x_frassoc_thick', 'v2x_suffr']

## Train loops

```python
def train_nn(x_train, y_train, need_pruning, seed, device, level, l1=0.001, lr=0.001, relu=False, verbose=False):
    set_seed(seed)
    x_train = x_train.to(device)
    y_train = y_train.to(device)

    if level == 'mid':
        layers = [
            torch.nn.Linear(x_train.size(1), 100),
            torch.nn.ReLU() if relu else torch.nn.LeakyReLU(),
            torch.nn.Linear(100, 50),
            torch.nn.ReLU() if relu else torch.nn.LeakyReLU(),
            torch.nn.Linear(50, 30),
            torch.nn.ReLU() if relu else torch.nn.LeakyReLU(),
            torch.nn.Linear(30, 10),
            torch.nn.ReLU() if relu else torch.nn.LeakyReLU(),
            torch.nn.Linear(10, y_train.size(1)),
            torch.nn.Sigmoid(),
        ]
        loss_form = torch.nn.BCELoss()
    else:
        layers = [
            torch.nn.Linear(x_train.size(1), 100),
            torch.nn.ReLU() if relu else torch.nn.LeakyReLU(),
            torch.nn.Linear(100, 50),
            torch.nn.ReLU() if relu else torch.nn.LeakyReLU(),
            torch.nn.Linear(50, 30),
            torch.nn.ReLU() if relu else torch.nn.LeakyReLU(),
            torch.nn.Linear(30, 10),
            torch.nn.ReLU() if relu else torch.nn.LeakyReLU(),
            torch.nn.Linear(10, 1),
            torch.nn.Sigmoid(),
        ]
        loss_form = torch.nn.BCELoss()

    model = torch.nn.Sequential(*layers).to(device)

    optimizer = torch.optim.AdamW(model.parameters(), lr=lr)
    model.train()
    for epoch in range(tot_epochs):
        # forward pass
        optimizer.zero_grad()
        y_pred = model(x_train).squeeze()
        # Compute Loss
        loss = loss_form(y_pred, y_train)

        for module in model.children():
            if isinstance(module, torch.nn.Linear):
#                 if level == 'mid':
#                     loss += 0.0001 * torch.norm(module.weight, 1)
#                     loss += 0.0001 * torch.norm(module.bias, 1)
                if level == 'high':
                    loss += l1 * torch.norm(module.weight, 1)
                    loss += l1 * torch.norm(module.bias, 1)
                break

        # backward pass
        loss.backward()
        optimizer.step()

        if epoch == prune_epochs and need_pruning and level == 'high':
            prune_features(model, n_classes=1, device=device)
            need_pruning = False

        # compute accuracy
        if epoch % 500 == 0 and verbose:
            if level == 'mid':
                y_pred_d = y_pred > 0.5
                accuracy = y_pred_d.eq(y_train).sum().item() / (y_train.size(0) * y_train.size(1))
                print(f'\t Epoch {epoch}: train accuracy: {accuracy:.4f}')
            else:
                y_pred_d = y_pred > 0.5
                accuracy = y_pred_d.eq(y_train).sum().item() / y_train.size(0)
                print(f'\t Epoch {epoch}: train accuracy: {accuracy:.4f}')
#             if level == 'mid':
#                 y_pred_d = y_pred > 0.5
#                 accuracy = y_pred_d.eq(y_train).sum().item() / (y_train.size(0) * y_train.size(1))
#                 print(f'\t Epoch {epoch}: train accuracy: {accuracy:.4f}')
#             else:
#                 y_pred_d = torch.argmax(y_pred, dim=1)
#                 accuracy = y_pred_d.eq(y_train).sum().item() / y_train.size(0)
#                 print(f'\t Epoch {epoch}: train accuracy: {accuracy:.4f}')

    return model
```

```python
def train_psi_nn(x_train, y_train, need_pruning, seed, device, level, verbose=False):
    set_seed(seed)
    x_train = x_train.to(device)
    y_train = y_train.to(device).to(torch.float)

    if level == 'mid':
        layers = [
            torch.nn.Linear(x_train.size(1), 50),
            torch.nn.Sigmoid(),
            torch.nn.Linear(50, 30),
            torch.nn.Sigmoid(),
            torch.nn.Linear(30, y_train.size(1)),
            torch.nn.Sigmoid(),
        ]
    else:
        layers = [
            torch.nn.Linear(x_train.size(1), 10),
            torch.nn.Sigmoid(),
            torch.nn.Linear(10, 4),
            torch.nn.Sigmoid(),
            torch.nn.Linear(4, 1),
            torch.nn.Sigmoid(),
        ]

    model = torch.nn.Sequential(*layers).to(device)

    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
    loss_form = torch.nn.BCELoss()
    model.train()
    for epoch in range(tot_epochs):
        # forward pass
        optimizer.zero_grad()
        y_pred = model(x_train).squeeze()
        # Compute Loss
        loss = loss_form(y_pred, y_train)

        for module in model.children():
            if isinstance(module, torch.nn.Linear):
                loss += 0.00001 * torch.norm(module.weight, 1)

        # backward pass
        loss.backward()
        optimizer.step()

        if epoch > prune_epochs and need_pruning:
            model = prune_equal_fanin(model, 2, validate=True, device=device)
            need_pruning = False

        # compute accuracy
        if epoch % 500 == 0 and verbose:
            if level == 'mid':
                y_pred_d = y_pred > 0.5
                accuracy = y_pred_d.eq(y_train).sum().item() / (y_train.size(0) * y_train.size(1))
                print(f'\t Epoch {epoch}: train accuracy: {accuracy:.4f}')
            else:
                y_pred_d = y_pred > 0.5
                accuracy = y_pred_d.eq(y_train).sum().item() / y_train.size(0)
                print(f'\t Epoch {epoch}: train accuracy: {accuracy:.4f}')

    return model
```

```python
def c_to_y(method, need_pruning, relu, l1=0.0001, lr=0.0001, verbose=False):
    methods = []
    splits = []
    explanations = []
    explanations_inv = []
    model_accuracies_mid = []
    model_accuracies_top = []
    explanation_accuracies = []
    explanation_accuracies_inv = []
    elapsed_times = []
    elapsed_times_inv = []
    explanation_fidelities = []
    explanation_complexities = []
    for split, (train_index, test_index) in enumerate(skf.split(x.cpu().detach().numpy(), y.cpu().detach().numpy())):
        print(f'Split [{split+1}/{n_splits}]')
        x_train, x_test = torch.FloatTensor(x[train_index]), torch.FloatTensor(x[test_index])
        xh_train, xh_test = torch.FloatTensor(xh[train_index]), torch.FloatTensor(xh[test_index])
        y_train, y_test = torch.FloatTensor(y[train_index]), torch.FloatTensor(y[test_index])

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

        start = time.time()

        if 'tree' in method:
            if method == 'tree-3':
                classifier = DecisionTreeRegressor(random_state=seed, max_depth=3)
            else:
                classifier = DecisionTreeRegressor(random_state=seed)

            classifier.fit(x_train.cpu().detach().numpy(), xh_train.cpu().detach().numpy())
            xh_train_preds = classifier.predict(x_train.cpu().detach().numpy())
            xh_test_preds = classifier.predict(x_test.cpu().detach().numpy())

            if method == 'tree-3':
                classifier_h = DecisionTreeClassifier(random_state=seed, max_depth=3)
            else:
                classifier_h = DecisionTreeClassifier(random_state=seed)
            classifier_h.fit(xh_train_preds, y_train)
            y_preds = classifier_h.predict(xh_test)

            model_accuracy_mid = accuracy_score(xh_test.cpu().detach().numpy(), xh_test_preds>0.5)
            model_accuracy_top = accuracy_score(y_test.cpu().detach().numpy(), y_preds)

            target_class = 1
            explanation = tree_to_formula(classifier_h, mid_concepts, target_class)

            explanation_accuracy = model_accuracy_top
            explanation_fidelity = accuracy_score(y_preds, y_preds)
            explanation_complexity = lens.logic.complexity(explanation)

        elif method == 'brl':
#             level = 'mid'
#             model = train_nn(x_train, xh_train, need_pruning=False, seed=seed, device=device,
#                              level=level, relu=False, verbose=verbose)
#             xh_train_preds = model(x_train.to(device)).cpu().detach().numpy()
#             xh_test_preds = model(x_test.to(device)).cpu().detach().numpy()

            classifier = DecisionTreeRegressor(random_state=seed)
            classifier.fit(x_train.cpu().detach().numpy(), xh_train.cpu().detach().numpy())
            xh_train_preds = classifier.predict(x_train.cpu().detach().numpy())
            xh_test_preds = classifier.predict(x_test.cpu().detach().numpy())

            xh_train_preds = torch.FloatTensor(xh_train_preds)
            xh_test_preds = torch.FloatTensor(xh_test_preds)

            train_dataset = TensorDataset(xh_train_preds, y_train)
            test_dataset = TensorDataset(xh_test_preds, y_test)
            dummy_concepts = [f'feature{i:010}' for i in range(len(mid_concepts))]
            model = XBRLClassifier(name=os.path.join(results_dir, f'{method}_{split}'),
                                   n_classes=1,
                                   n_features=x_train.shape[1], feature_names=dummy_concepts,
                                   class_names=['~democrat', 'democrat'], discretize=True)
            results = model.fit(train_dataset, metric=accuracy_score, save=True)
            model_accuracy_mid = 0
            global_explanation, elapsed_time = model.get_global_explanation(0, concepts=dummy_concepts, return_time=True)
            explanation = logic.base.replace_names(global_explanation, mid_concepts)

            # compute metrics
            y_preds = model.forward(xh_test_preds)
            model_accuracy_top = accuracy_score(y_test, y_preds)
            target_class = 1
            explanation_accuracy, y_formula = logic.base.test_explanation(global_explanation,
                                                                          target_class,
                                                                          x=xh_test, y=y_test,
                                                                          metric=accuracy_score)
            explanation_fidelity = accuracy_score(y_preds, y_formula)
            explanation_complexity = lens.logic.complexity(global_explanation)

        else:
            if method == 'psi':
                # positive class
                target_class = 1
                level = 'mid'
                model = train_psi_nn(x_train, xh_train, need_pruning, split, device, level, verbose)
                xh_train_preds = model(x_train.to(device)).cpu().detach().numpy()
                xh_test_preds = model(x_test.to(device)).cpu().detach().numpy()
                xh_train_preds = torch.FloatTensor(xh_train_preds)
                xh_test_preds = torch.FloatTensor(xh_test_preds)

                level = 'high'
                model_h = train_psi_nn(xh_train_preds, y_train, need_pruning, split, device, level, verbose)
                y_train_preds = model_h(xh_train_preds.to(device)).cpu().detach().numpy()
                y_test_preds = model_h(xh_test_preds.to(device)).cpu().detach().numpy()

                model_accuracy_mid = accuracy_score(xh_test.cpu().detach().numpy(), xh_test_preds>0.5)
                model_accuracy_top = accuracy_score(y_test.cpu().detach().numpy(), y_test_preds>0.5)

            else:
                level = 'mid'
                model = train_nn(x_train, xh_train, need_pruning, seed, device, level, l1, lr, relu, verbose)
                xh_train_preds = model(x_train.to(device)).cpu().detach().numpy()
                xh_test_preds = model(x_test.to(device)).cpu().detach().numpy()
                xh_train_preds = torch.FloatTensor(xh_train_preds)
                xh_test_preds = torch.FloatTensor(xh_test_preds)

                level = 'high'
                model_h = train_nn(xh_train_preds, y_train, need_pruning, seed, device, level, l1, lr, relu, verbose)
                y_train_preds = model_h(xh_train_preds.to(device)).cpu().detach().numpy()
                y_test_preds = model_h(xh_test_preds.to(device)).cpu().detach().numpy()

                model_accuracy_mid = accuracy_score(xh_test.cpu().detach().numpy(), xh_test_preds>0.5)
                model_accuracy_top = accuracy_score(y_test.cpu().detach().numpy(), y_test_preds>0.5)

            # positive class
            target_class = 1
            if method == 'psi':
                global_explanation = logic.generate_fol_explanations(model_h, device)[0]
            else:
                global_explanation, _, _ = logic.relu_nn.combine_local_explanations(model_h,
                                                                                xh_train.to(device), y_train.to(device),
                                                                                target_class=target_class,
                                                                                topk_explanations=2,
                                                                                method=method, device=device)
            explanation_fidelity = 0
            explanation_complexity = 0
            if global_explanation:
                explanation = logic.base.replace_names(global_explanation, mid_concepts)
                explanation_accuracy, y_formula = logic.base.test_explanation(global_explanation,
                                                                              target_class,
                                                                              x=xh_test, y=y_test,
                                                                              metric=accuracy_score)
                explanation_fidelity = accuracy_score(y_test_preds>0.5, y_formula)
                explanation_complexity = lens.logic.complexity(global_explanation)

        elapsed_time = time.time() - start

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

        methods.append(method)
        splits.append(split)
        explanations.append(explanation)
        model_accuracies_mid.append(model_accuracy_mid)
        model_accuracies_top.append(model_accuracy_top)
        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_mid': model_accuracies_mid,
        'model_accuracy_top': model_accuracies_top,
        '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
method = 'pruning'
need_pruning = True
relu = False
results_pruning = c_to_y(method, need_pruning, relu, l1=0.01, lr=0.0001, verbose=True)
results_pruning
```

```python
results_pruning.mean(axis=0)
```

    split                       4.500000
    model_accuracy_mid          0.728509
    model_accuracy_top          0.923477
    explanation_accuracy        0.946819
    explanation_fidelity        0.951254
    explanation_complexity      5.000000
    explanation_consistency     1.000000
    elapsed_time               31.864813
    dtype: float64

## ReLU

```python
method = 'weights'
need_pruning = False
relu = True
results_weights = c_to_y(method, need_pruning, relu, l1=0.00001, lr=0.0001, True)
results_weights
```

    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.9800

      method
      split
      explanation
      model_accuracy_mid
      model_accuracy_top
      explanation_accuracy
      explanation_fidelity
      explanation_complexity
      explanation_consistency
      elapsed_time

      0
      weights
      0
      v2x_freexp_altinf & v2xel_frefair & v2x_elecof...
      0.752212
      0.941003
      0.943953
      0.955752
      4
      0.98
      36.240852

      1
      weights
      1
      v2x_freexp_altinf & v2x_frassoc_thick & v2xel_...
      0.731563
      0.929204
      0.955752
      0.967552
      5
      0.98
      35.648936

      2
      weights
      2
      v2x_freexp_altinf & v2x_frassoc_thick & v2xel_...
      0.710914
      0.923304
      0.938053
      0.938053
      5
      0.98
      34.755997

      3
      weights
      3
      v2x_freexp_altinf & v2x_frassoc_thick & v2xel_...
      0.696165
      0.911504
      0.949853
      0.943953
      5
      0.98
      33.704165

      4
      weights
      4
      v2x_freexp_altinf & v2x_frassoc_thick & v2xel_...
      0.731563
      0.917404
      0.961652
      0.938053
      5
      0.98
      33.458029

      5
      weights
      5
      v2x_freexp_altinf & v2x_frassoc_thick & v2xel_...
      0.718935
      0.899408
      0.934911
      0.934911
      5
      0.98
      33.571011

      6
      weights
      6
      v2x_freexp_altinf & v2x_frassoc_thick & v2xel_...
      0.733728
      0.920118
      0.943787
      0.946746
      5
      0.98
      33.346951

      7
      weights
      7
      v2x_freexp_altinf & v2x_frassoc_thick & v2xel_...
      0.742604
      0.923077
      0.961538
      0.961538
      5
      0.98
      33.487034

      8
      weights
      8
      v2x_freexp_altinf & v2x_frassoc_thick & v2xel_...
      0.713018
      0.917160
      0.940828
      0.934911
      5
      0.98
      33.645938

      9
      weights
      9
      v2x_freexp_altinf & v2x_frassoc_thick & v2xel_...
      0.754438
      0.905325
      0.934911
      0.940828
      5
      0.98
      34.565918

## Psi network

```python
method = 'psi'
need_pruning = True
relu = False
results_psi = c_to_y(method, need_pruning, relu, verbose=True)
results_psi
```

    Split [1/10]
    	 Epoch 0: train accuracy: 0.4649
    	 Epoch 500: train accuracy: 0.9104
    	 Epoch 1000: train accuracy: 0.9367
    	 Epoch 1500: train accuracy: 0.9503
    	 Epoch 2000: train accuracy: 0.9622
    	 Epoch 2500: train accuracy: 0.9742
    	 Epoch 3000: train accuracy: 0.9847
    	 Epoch 3500: train accuracy: 0.8533
    	 Epoch 4000: train accuracy: 0.8856
    	 Epoch 4500: train accuracy: 0.8930
    	 Epoch 5000: train accuracy: 0.8943
    	 Epoch 5500: train accuracy: 0.8964
    	 Epoch 6000: train accuracy: 0.8992
    	 Epoch 0: train accuracy: 0.4685
    	 Epoch 500: train accuracy: 0.7173
    	 Epoch 1000: train accuracy: 0.8546
    	 Epoch 1500: train accuracy: 0.8651
    	 Epoch 2000: train accuracy: 0.8766
    	 Epoch 2500: train accuracy: 0.8775
    	 Epoch 3000: train accuracy: 0.8805
    	 Epoch 3500: train accuracy: 0.8513
    	 Epoch 4000: train accuracy: 0.8480
    	 Epoch 4500: train accuracy: 0.8510
    	 Epoch 5000: train accuracy: 0.8595
    	 Epoch 5500: train accuracy: 0.8618
    	 Epoch 6000: train accuracy: 0.8638
    	 Model's accuracy (mid): 0.3451 - Model's accuracy (top): 0.8673
    	 Class 1 - Global explanation: "(v2x_mpi)" - Accuracy: 0.8142
    	 Elapsed time 20.810575246810913
    	 Fidelity: "0.8230" - Complexity: "1"
    Split [2/10]
    	 Epoch 0: train accuracy: 0.4695
    	 Epoch 500: train accuracy: 0.9072
    	 Epoch 1000: train accuracy: 0.9349
    	 Epoch 1500: train accuracy: 0.9480
    	 Epoch 2000: train accuracy: 0.9592
    	 Epoch 2500: train accuracy: 0.9750
    	 Epoch 3000: train accuracy: 0.9841
    	 Epoch 3500: train accuracy: 0.8153
    	 Epoch 4000: train accuracy: 0.8607
    	 Epoch 4500: train accuracy: 0.8692
    	 Epoch 5000: train accuracy: 0.8751
    	 Epoch 5500: train accuracy: 0.8810
    	 Epoch 6000: train accuracy: 0.8829
    	 Epoch 0: train accuracy: 0.5315
    	 Epoch 500: train accuracy: 0.7754
    	 Epoch 1000: train accuracy: 0.8450
    	 Epoch 1500: train accuracy: 0.8572
    	 Epoch 2000: train accuracy: 0.8615
    	 Epoch 2500: train accuracy: 0.8697
    	 Epoch 3000: train accuracy: 0.8710
    	 Epoch 3500: train accuracy: 0.8194
    	 Epoch 4000: train accuracy: 0.8276
    	 Epoch 4500: train accuracy: 0.8244
    	 Epoch 5000: train accuracy: 0.8273
    	 Epoch 5500: train accuracy: 0.8296
    	 Epoch 6000: train accuracy: 0.8336
    	 Model's accuracy (mid): 0.3363 - Model's accuracy (top): 0.8496
    	 Class 1 - Global explanation: "(v2x_freexp_altinf & v2xcl_rol)" - Accuracy: 0.8289
    	 Elapsed time 21.03696036338806
    	 Fidelity: "0.8319" - Complexity: "2"
    Split [3/10]
    	 Epoch 0: train accuracy: 0.5062
    	 Epoch 500: train accuracy: 0.9127
    	 Epoch 1000: train accuracy: 0.9378
    	 Epoch 1500: train accuracy: 0.9504
    	 Epoch 2000: train accuracy: 0.9652
    	 Epoch 2500: train accuracy: 0.9788
    	 Epoch 3000: train accuracy: 0.9865
    	 Epoch 3500: train accuracy: 0.8423
    	 Epoch 4000: train accuracy: 0.8808
    	 Epoch 4500: train accuracy: 0.8839
    	 Epoch 5000: train accuracy: 0.8922
    	 Epoch 5500: train accuracy: 0.8999
    	 Epoch 6000: train accuracy: 0.9007
    	 Epoch 0: train accuracy: 0.4685
    	 Epoch 500: train accuracy: 0.8362
    	 Epoch 1000: train accuracy: 0.8723
    	 Epoch 1500: train accuracy: 0.8802
    	 Epoch 2000: train accuracy: 0.8808
    	 Epoch 2500: train accuracy: 0.8795
    	 Epoch 3000: train accuracy: 0.8743
    	 Epoch 3500: train accuracy: 0.7991
    	 Epoch 4000: train accuracy: 0.8060
    	 Epoch 4500: train accuracy: 0.8102
    	 Epoch 5000: train accuracy: 0.8125
    	 Epoch 5500: train accuracy: 0.8178
    	 Epoch 6000: train accuracy: 0.8240
    	 Model's accuracy (mid): 0.3540 - Model's accuracy (top): 0.8348
    	 Class 1 - Global explanation: "(v2xeg_eqaccess & (v2xdl_delib | ~v2x_suffr))" - Accuracy: 0.8466
    	 Elapsed time 20.94809055328369
    	 Fidelity: "0.8584" - Complexity: "3"
    Split [4/10]
    	 Epoch 0: train accuracy: 0.3644
    	 Epoch 500: train accuracy: 0.9121
    	 Epoch 1000: train accuracy: 0.9362
    	 Epoch 1500: train accuracy: 0.9495
    	 Epoch 2000: train accuracy: 0.9636
    	 Epoch 2500: train accuracy: 0.9739
    	 Epoch 3000: train accuracy: 0.9837
    	 Epoch 3500: train accuracy: 0.8111
    	 Epoch 4000: train accuracy: 0.8292
    	 Epoch 4500: train accuracy: 0.8642
    	 Epoch 5000: train accuracy: 0.8708
    	 Epoch 5500: train accuracy: 0.8749
    	 Epoch 6000: train accuracy: 0.8750
    	 Epoch 0: train accuracy: 0.5315
    	 Epoch 500: train accuracy: 0.8342
    	 Epoch 1000: train accuracy: 0.8631
    	 Epoch 1500: train accuracy: 0.8707
    	 Epoch 2000: train accuracy: 0.8762
    	 Epoch 2500: train accuracy: 0.8795
    	 Epoch 3000: train accuracy: 0.8841
    	 Epoch 3500: train accuracy: 0.7180
    	 Epoch 4000: train accuracy: 0.7124
    	 Epoch 4500: train accuracy: 0.7183
    	 Epoch 5000: train accuracy: 0.7410
    	 Epoch 5500: train accuracy: 0.7387
    	 Epoch 6000: train accuracy: 0.7702
    	 Model's accuracy (mid): 0.2655 - Model's accuracy (top): 0.7640
    	 Class 1 - Global explanation: "(v2x_elecoff)" - Accuracy: 0.6814
    	 Elapsed time 20.736979484558105
    	 Fidelity: "0.7050" - Complexity: "1"
    Split [5/10]
    	 Epoch 0: train accuracy: 0.5734
    	 Epoch 500: train accuracy: 0.9120
    	 Epoch 1000: train accuracy: 0.9376
    	 Epoch 1500: train accuracy: 0.9500
    	 Epoch 2000: train accuracy: 0.9641
    	 Epoch 2500: train accuracy: 0.9770
    	 Epoch 3000: train accuracy: 0.9847
    	 Epoch 3500: train accuracy: 0.7999
    	 Epoch 4000: train accuracy: 0.8295
    	 Epoch 4500: train accuracy: 0.8427
    	 Epoch 5000: train accuracy: 0.8587
    	 Epoch 5500: train accuracy: 0.8615
    	 Epoch 6000: train accuracy: 0.8672
    	 Epoch 0: train accuracy: 0.5315
    	 Epoch 500: train accuracy: 0.8414
    	 Epoch 1000: train accuracy: 0.8601
    	 Epoch 1500: train accuracy: 0.8644
    	 Epoch 2000: train accuracy: 0.8664
    	 Epoch 2500: train accuracy: 0.8697
    	 Epoch 3000: train accuracy: 0.8785
    	 Epoch 3500: train accuracy: 0.6202
    	 Epoch 4000: train accuracy: 0.7981
    	 Epoch 4500: train accuracy: 0.8142
    	 Epoch 5000: train accuracy: 0.8230
    	 Epoch 5500: train accuracy: 0.8148
    	 Epoch 6000: train accuracy: 0.8135
    	 Model's accuracy (mid): 0.1681 - Model's accuracy (top): 0.8201
    	 Class 1 - Global explanation: "(v2x_frassoc_thick & v2xeg_eqaccess)" - Accuracy: 0.8761
    	 Elapsed time 20.86394429206848
    	 Fidelity: "0.8260" - Complexity: "2"
    Split [6/10]
    	 Epoch 0: train accuracy: 0.4149
    	 Epoch 500: train accuracy: 0.9151
    	 Epoch 1000: train accuracy: 0.9371
    	 Epoch 1500: train accuracy: 0.9476
    	 Epoch 2000: train accuracy: 0.9619
    	 Epoch 2500: train accuracy: 0.9750
    	 Epoch 3000: train accuracy: 0.9843
    	 Epoch 3500: train accuracy: 0.8164
    	 Epoch 4000: train accuracy: 0.8335
    	 Epoch 4500: train accuracy: 0.8559
    	 Epoch 5000: train accuracy: 0.8668
    	 Epoch 5500: train accuracy: 0.8738
    	 Epoch 6000: train accuracy: 0.8736
    	 Epoch 0: train accuracy: 0.4687
    	 Epoch 500: train accuracy: 0.8326
    	 Epoch 1000: train accuracy: 0.8697
    	 Epoch 1500: train accuracy: 0.8792
    	 Epoch 2000: train accuracy: 0.8825
    	 Epoch 2500: train accuracy: 0.8858
    	 Epoch 3000: train accuracy: 0.8904
    	 Epoch 3500: train accuracy: 0.7345
    	 Epoch 4000: train accuracy: 0.7565
    	 Epoch 4500: train accuracy: 0.7683
    	 Epoch 5000: train accuracy: 0.7758
    	 Epoch 5500: train accuracy: 0.7801
    	 Epoch 6000: train accuracy: 0.7896
    	 Model's accuracy (mid): 0.3787 - Model's accuracy (top): 0.7870
    	 Class 1 - Global explanation: "(v2x_frassoc_thick & v2xeg_eqaccess)" - Accuracy: 0.7959
    	 Elapsed time 20.888019323349
    	 Fidelity: "0.8077" - Complexity: "2"
    Split [7/10]
    	 Epoch 0: train accuracy: 0.5186
    	 Epoch 500: train accuracy: 0.9088
    	 Epoch 1000: train accuracy: 0.9361
    	 Epoch 1500: train accuracy: 0.9501
    	 Epoch 2000: train accuracy: 0.9611
    	 Epoch 2500: train accuracy: 0.9769
    	 Epoch 3000: train accuracy: 0.9853
    	 Epoch 3500: train accuracy: 0.8163
    	 Epoch 4000: train accuracy: 0.8711
    	 Epoch 4500: train accuracy: 0.8824
    	 Epoch 5000: train accuracy: 0.9013
    	 Epoch 5500: train accuracy: 0.9029
    	 Epoch 6000: train accuracy: 0.9034
    	 Epoch 0: train accuracy: 0.5313
    	 Epoch 500: train accuracy: 0.8392
    	 Epoch 1000: train accuracy: 0.8654
    	 Epoch 1500: train accuracy: 0.8694
    	 Epoch 2000: train accuracy: 0.8730
    	 Epoch 2500: train accuracy: 0.8779
    	 Epoch 3000: train accuracy: 0.8871
    	 Epoch 3500: train accuracy: 0.7384
    	 Epoch 4000: train accuracy: 0.7535
    	 Epoch 4500: train accuracy: 0.7529
    	 Epoch 5000: train accuracy: 0.7466
    	 Epoch 5500: train accuracy: 0.7496
    	 Epoch 6000: train accuracy: 0.7473
    	 Model's accuracy (mid): 0.3462 - Model's accuracy (top): 0.7337
    	 Class 1 - Global explanation: "(v2x_egal)" - Accuracy: 0.7101
    	 Elapsed time 20.66694474220276
    	 Fidelity: "0.8402" - Complexity: "1"
    Split [8/10]
    	 Epoch 0: train accuracy: 0.5142
    	 Epoch 500: train accuracy: 0.9026
    	 Epoch 1000: train accuracy: 0.9372
    	 Epoch 1500: train accuracy: 0.9485
    	 Epoch 2000: train accuracy: 0.9651
    	 Epoch 2500: train accuracy: 0.9782
    	 Epoch 3000: train accuracy: 0.9863
    	 Epoch 3500: train accuracy: 0.8006
    	 Epoch 4000: train accuracy: 0.8554
    	 Epoch 4500: train accuracy: 0.8759
    	 Epoch 5000: train accuracy: 0.8800
    	 Epoch 5500: train accuracy: 0.8835
    	 Epoch 6000: train accuracy: 0.8845
    	 Epoch 0: train accuracy: 0.5313
    	 Epoch 500: train accuracy: 0.8461
    	 Epoch 1000: train accuracy: 0.8661
    	 Epoch 1500: train accuracy: 0.8677
    	 Epoch 2000: train accuracy: 0.8691
    	 Epoch 2500: train accuracy: 0.8710
    	 Epoch 3000: train accuracy: 0.8809
    	 Epoch 3500: train accuracy: 0.7460
    	 Epoch 4000: train accuracy: 0.7699
    	 Epoch 4500: train accuracy: 0.7781
    	 Epoch 5000: train accuracy: 0.7749
    	 Epoch 5500: train accuracy: 0.7870
    	 Epoch 6000: train accuracy: 0.7968
    	 Model's accuracy (mid): 0.3166 - Model's accuracy (top): 0.7811
    	 Class 1 - Global explanation: "(v2xel_frefair)" - Accuracy: 0.9201
    	 Elapsed time 20.79501700401306
    	 Fidelity: "0.7959" - Complexity: "1"
    Split [9/10]
    	 Epoch 0: train accuracy: 0.5690
    	 Epoch 500: train accuracy: 0.9119
    	 Epoch 1000: train accuracy: 0.9372
    	 Epoch 1500: train accuracy: 0.9493
    	 Epoch 2000: train accuracy: 0.9618
    	 Epoch 2500: train accuracy: 0.9753
    	 Epoch 3000: train accuracy: 0.9836
    	 Epoch 3500: train accuracy: 0.7855
    	 Epoch 4000: train accuracy: 0.8039
    	 Epoch 4500: train accuracy: 0.8427
    	 Epoch 5000: train accuracy: 0.8544
    	 Epoch 5500: train accuracy: 0.8617
    	 Epoch 6000: train accuracy: 0.8674
    	 Epoch 0: train accuracy: 0.5313
    	 Epoch 500: train accuracy: 0.8014
    	 Epoch 1000: train accuracy: 0.8707
    	 Epoch 1500: train accuracy: 0.8756
    	 Epoch 2000: train accuracy: 0.8773
    	 Epoch 2500: train accuracy: 0.8750
    	 Epoch 3000: train accuracy: 0.8841
    	 Epoch 3500: train accuracy: 0.8014
    	 Epoch 4000: train accuracy: 0.8303
    	 Epoch 4500: train accuracy: 0.8513
    	 Epoch 5000: train accuracy: 0.8487
    	 Epoch 5500: train accuracy: 0.8503
    	 Epoch 6000: train accuracy: 0.8563
    	 Model's accuracy (mid): 0.3047 - Model's accuracy (top): 0.8491
    	 Class 1 - Global explanation: "(v2x_freexp_altinf & v2x_cspart & (v2xel_frefair | ~v2x_suffr) & (v2xeg_eqaccess | ~v2x_suffr))" - Accuracy: 0.9083
    	 Elapsed time 20.954056978225708
    	 Fidelity: "0.8225" - Complexity: "6"
    Split [10/10]
    	 Epoch 0: train accuracy: 0.5914
    	 Epoch 500: train accuracy: 0.9018
    	 Epoch 1000: train accuracy: 0.9371
    	 Epoch 1500: train accuracy: 0.9494
    	 Epoch 2000: train accuracy: 0.9649
    	 Epoch 2500: train accuracy: 0.9785
    	 Epoch 3000: train accuracy: 0.9861
    	 Epoch 3500: train accuracy: 0.7995
    	 Epoch 4000: train accuracy: 0.8485
    	 Epoch 4500: train accuracy: 0.8572
    	 Epoch 5000: train accuracy: 0.8672
    	 Epoch 5500: train accuracy: 0.8754
    	 Epoch 6000: train accuracy: 0.8801
    	 Epoch 0: train accuracy: 0.4683
    	 Epoch 500: train accuracy: 0.8011
    	 Epoch 1000: train accuracy: 0.8540
    	 Epoch 1500: train accuracy: 0.8664
    	 Epoch 2000: train accuracy: 0.8664
    	 Epoch 2500: train accuracy: 0.8740
    	 Epoch 3000: train accuracy: 0.8796
    	 Epoch 3500: train accuracy: 0.6387
    	 Epoch 4000: train accuracy: 0.6387
    	 Epoch 4500: train accuracy: 0.6387
    	 Epoch 5000: train accuracy: 0.6387
    	 Epoch 5500: train accuracy: 0.6387
    	 Epoch 6000: train accuracy: 0.6387
    	 Model's accuracy (mid): 0.3314 - Model's accuracy (top): 0.6361
    	 Class 1 - Global explanation: "(v2xeg_eqdr)" - Accuracy: 0.6479
    	 Elapsed time 20.614274740219116
    	 Fidelity: "0.9349" - Complexity: "1"
    Consistency of explanations: 0.1583

      method
      split
      explanation
      model_accuracy_mid
      model_accuracy_top
      explanation_accuracy
      explanation_fidelity
      explanation_complexity
      explanation_consistency
      elapsed_time

      0
      psi
      0
      (v2x_mpi)
      0.345133
      0.867257
      0.814159
      0.823009
      1
      0.158333
      20.810575

      1
      psi
      1
      (v2x_freexp_altinf & v2xcl_rol)
      0.336283
      0.849558
      0.828909
      0.831858
      2
      0.158333
      21.036960

      2
      psi
      2
      (v2xeg_eqaccess & (v2xdl_delib | ~v2x_suffr))
      0.353982
      0.834808
      0.846608
      0.858407
      3
      0.158333
      20.948091

      3
      psi
      3
      (v2x_elecoff)
      0.265487
      0.764012
      0.681416
      0.705015
      1
      0.158333
      20.736979

      4
      psi
      4
      (v2x_frassoc_thick & v2xeg_eqaccess)
      0.168142
      0.820059
      0.876106
      0.825959
      2
      0.158333
      20.863944

      5
      psi
      5
      (v2x_frassoc_thick & v2xeg_eqaccess)
      0.378698
      0.786982
      0.795858
      0.807692
      2
      0.158333
      20.888019

      6
      psi
      6
      (v2x_egal)
      0.346154
      0.733728
      0.710059
      0.840237
      1
      0.158333
      20.666945

      7
      psi
      7
      (v2xel_frefair)
      0.316568
      0.781065
      0.920118
      0.795858
      1
      0.158333
      20.795017

      8
      psi
      8
      (v2x_freexp_altinf & v2x_cspart & (v2xel_frefa...
      0.304734
      0.849112
      0.908284
      0.822485
      6
      0.158333
      20.954057

      9
      psi
      9
      (v2xeg_eqdr)
      0.331361
      0.636095
      0.647929
      0.934911
      1
      0.158333
      20.614275

## Decision tree

```python
method = 'tree'
need_pruning = False
relu = False
results_tree = c_to_y(method, need_pruning, relu, 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.3686

      method
      split
      explanation
      model_accuracy_mid
      model_accuracy_top
      explanation_accuracy
      explanation_fidelity
      explanation_complexity
      explanation_consistency
      elapsed_time

      0
      tree
      0
      (v2xel_frefair <= 0.33 & v2xcl_rol > 0.90 & v2...
      0.749263
      0.946903
      0.946903
      1.0
      305
      0.368571
      0.090787

      1
      tree
      1
      (v2xel_frefair <= 0.33 & v2xdl_delib > 0.58 & ...
      0.784661
      0.952802
      0.952802
      1.0
      294
      0.368571
      0.072805

      2
      tree
      2
      (v2xel_frefair <= 0.58 & v2xcl_rol > 0.88 & v2...
      0.746313
      0.938053
      0.938053
      1.0
      236
      0.368571
      0.078789

      3
      tree
      3
      (v2xel_frefair <= 0.88 & v2xcl_rol > 0.90 & v2...
      0.734513
      0.949853
      0.949853
      1.0
      289
      0.368571
      0.069811

      4
      tree
      4
      (v2xel_frefair <= 0.44 & v2xcl_rol > 0.90 & v2...
      0.722714
      0.958702
      0.958702
      1.0
      289
      0.368571
      0.070817

      5
      tree
      5
      (v2xel_frefair <= 0.46 & v2xdl_delib > 0.61 & ...
      0.745562
      0.931953
      0.931953
      1.0
      319
      0.368571
      0.082806

      6
      tree
      6
      (v2xel_frefair <= 0.47 & v2xcl_rol > 0.90 & v2...
      0.760355
      0.943787
      0.943787
      1.0
      313
      0.368571
      0.073838

      7
      tree
      7
      (v2xel_frefair <= 0.33 & v2xcl_rol > 0.88 & v2...
      0.769231
      0.961538
      0.961538
      1.0
      309
      0.368571
      0.075795

      8
      tree
      8
      (v2xel_frefair <= 0.88 & v2xcl_rol > 0.90 & v2...
      0.754438
      0.940828
      0.940828
      1.0
      273
      0.368571
      0.073803

      9
      tree
      9
      (v2xel_frefair <= 0.58 & v2xcl_rol > 0.88 & v2...
      0.789941
      0.934911
      0.934911
      1.0
      314
      0.368571
      0.071769

```python
method = 'tree-3'
need_pruning = False
relu = False
results_tree = c_to_y(method, need_pruning, relu, verbose=True)
results_tree
```

    Split [1/10]
    	 Model's accuracy (mid): 0.3540 - Model's accuracy (top): 0.7021
    	 Class 1 - Global explanation: "(v2x_api > 0.94 & v2x_cspart <= 0.96 & v2xeg_eqaccess <= 0.88) | (v2x_api > 0.94 & v2x_cspart <= 0.96 & v2xeg_eqaccess > 0.88) | (v2x_api > 0.94 & v2x_cspart > 0.96)" - Accuracy: 0.7021
    	 Elapsed time 0.05385541915893555
    	 Fidelity: "1.0000" - Complexity: "8"
    Split [2/10]
    	 Model's accuracy (mid): 0.3687 - Model's accuracy (top): 0.6519
    	 Class 1 - Global explanation: "(v2x_elecoff > 0.91 & v2x_cspart <= 0.96 & v2xeg_eqaccess <= 0.89) | (v2x_elecoff > 0.91 & v2x_cspart <= 0.96 & v2xeg_eqaccess > 0.89) | (v2x_elecoff > 0.91 & v2x_cspart > 0.96)" - Accuracy: 0.6519
    	 Elapsed time 0.05037879943847656
    	 Fidelity: "1.0000" - Complexity: "8"
    Split [3/10]
    	 Model's accuracy (mid): 0.3776 - Model's accuracy (top): 0.6637
    	 Class 1 - Global explanation: "(v2x_elecoff > 0.92 & v2x_cspart <= 0.96 & v2xeg_eqaccess <= 0.87) | (v2x_elecoff > 0.92 & v2x_cspart <= 0.96 & v2xeg_eqaccess > 0.87) | (v2x_elecoff > 0.92 & v2x_cspart > 0.96)" - Accuracy: 0.6637
    	 Elapsed time 0.04688262939453125
    	 Fidelity: "1.0000" - Complexity: "8"
    Split [4/10]
    	 Model's accuracy (mid): 0.3658 - Model's accuracy (top): 0.6814
    	 Class 1 - Global explanation: "(v2x_elecoff > 0.92 & v2x_cspart <= 0.95 & v2xeg_eqaccess <= 0.88) | (v2x_elecoff > 0.92 & v2x_cspart <= 0.95 & v2xeg_eqaccess > 0.88) | (v2x_elecoff > 0.92 & v2x_cspart > 0.95)" - Accuracy: 0.6814
    	 Elapsed time 0.04882049560546875
    	 Fidelity: "1.0000" - Complexity: "8"
    Split [5/10]
    	 Model's accuracy (mid): 0.3746 - Model's accuracy (top): 0.7788
    	 Class 1 - Global explanation: "(v2x_api > 0.94 & v2x_cspart <= 0.96 & v2xeg_eqaccess <= 0.88) | (v2x_api > 0.94 & v2x_cspart <= 0.96 & v2xeg_eqaccess > 0.88) | (v2x_api > 0.94 & v2x_cspart > 0.96)" - Accuracy: 0.7788
    	 Elapsed time 0.04886960983276367
    	 Fidelity: "1.0000" - Complexity: "8"
    Split [6/10]
    	 Model's accuracy (mid): 0.4083 - Model's accuracy (top): 0.6183
    	 Class 1 - Global explanation: "(v2x_elecoff > 0.91 & v2x_cspart <= 0.96 & v2xeg_eqaccess <= 0.87) | (v2x_elecoff > 0.91 & v2x_cspart <= 0.96 & v2xeg_eqaccess > 0.87) | (v2x_elecoff > 0.91 & v2x_cspart > 0.96)" - Accuracy: 0.6183
    	 Elapsed time 0.050864458084106445
    	 Fidelity: "1.0000" - Complexity: "8"
    Split [7/10]
    	 Model's accuracy (mid): 0.3432 - Model's accuracy (top): 0.6657
    	 Class 1 - Global explanation: "(v2x_elecoff > 0.91 & v2x_cspart <= 0.96 & v2xeg_eqaccess <= 0.87) | (v2x_elecoff > 0.91 & v2x_cspart <= 0.96 & v2xeg_eqaccess > 0.87) | (v2x_elecoff > 0.91 & v2x_cspart > 0.96)" - Accuracy: 0.6657
    	 Elapsed time 0.043816328048706055
    	 Fidelity: "1.0000" - Complexity: "8"
    Split [8/10]
    	 Model's accuracy (mid): 0.3462 - Model's accuracy (top): 0.7456
    	 Class 1 - Global explanation: "(v2x_api > 0.94 & v2x_cspart <= 0.96 & v2xeg_eqaccess <= 0.87) | (v2x_api > 0.94 & v2x_cspart <= 0.96 & v2xeg_eqaccess > 0.87) | (v2x_api > 0.94 & v2x_cspart > 0.96)" - Accuracy: 0.7456
    	 Elapsed time 0.03590130805969238
    	 Fidelity: "1.0000" - Complexity: "8"
    Split [9/10]
    	 Model's accuracy (mid): 0.3817 - Model's accuracy (top): 0.6953
    	 Class 1 - Global explanation: "(v2x_elecoff > 0.92 & v2x_cspart <= 0.96 & v2xeg_eqaccess <= 0.89) | (v2x_elecoff > 0.92 & v2x_cspart <= 0.96 & v2xeg_eqaccess > 0.89) | (v2x_elecoff > 0.92 & v2x_cspart > 0.96)" - Accuracy: 0.6953
    	 Elapsed time 0.05502510070800781
    	 Fidelity: "1.0000" - Complexity: "8"
    Split [10/10]
    	 Model's accuracy (mid): 0.3609 - Model's accuracy (top): 0.6479
    	 Class 1 - Global explanation: "(v2x_elecoff > 0.91 & v2x_cspart <= 0.96 & v2xeg_eqaccess <= 0.87) | (v2x_elecoff > 0.91 & v2x_cspart <= 0.96 & v2xeg_eqaccess > 0.87) | (v2x_elecoff > 0.91 & v2x_cspart > 0.96)" - Accuracy: 0.6479
    	 Elapsed time 0.05886125564575195
    	 Fidelity: "1.0000" - Complexity: "8"
    Consistency of explanations: 0.3846

      method
      split
      explanation
      model_accuracy_mid
      model_accuracy_top
      explanation_accuracy
      explanation_fidelity
      explanation_complexity
      explanation_consistency
      elapsed_time

      0
      tree-3
      0
      (v2x_api > 0.94 & v2x_cspart <= 0.96 & v2xeg_e...
      0.353982
      0.702065
      0.702065
      1.0
      8
      0.384615
      0.053855

      1
      tree-3
      1
      (v2x_elecoff > 0.91 & v2x_cspart <= 0.96 & v2x...
      0.368732
      0.651917
      0.651917
      1.0
      8
      0.384615
      0.050379

      2
      tree-3
      2
      (v2x_elecoff > 0.92 & v2x_cspart <= 0.96 & v2x...
      0.377581
      0.663717
      0.663717
      1.0
      8
      0.384615
      0.046883

      3
      tree-3
      3
      (v2x_elecoff > 0.92 & v2x_cspart <= 0.95 & v2x...
      0.365782
      0.681416
      0.681416
      1.0
      8
      0.384615
      0.048820

      4
      tree-3
      4
      (v2x_api > 0.94 & v2x_cspart <= 0.96 & v2xeg_e...
      0.374631
      0.778761
      0.778761
      1.0
      8
      0.384615
      0.048870

      5
      tree-3
      5
      (v2x_elecoff > 0.91 & v2x_cspart <= 0.96 & v2x...
      0.408284
      0.618343
      0.618343
      1.0
      8
      0.384615
      0.050864

      6
      tree-3
      6
      (v2x_elecoff > 0.91 & v2x_cspart <= 0.96 & v2x...
      0.343195
      0.665680
      0.665680
      1.0
      8
      0.384615
      0.043816

      7
      tree-3
      7
      (v2x_api > 0.94 & v2x_cspart <= 0.96 & v2xeg_e...
      0.346154
      0.745562
      0.745562
      1.0
      8
      0.384615
      0.035901

      8
      tree-3
      8
      (v2x_elecoff > 0.92 & v2x_cspart <= 0.96 & v2x...
      0.381657
      0.695266
      0.695266
      1.0
      8
      0.384615
      0.055025

      9
      tree-3
      9
      (v2x_elecoff > 0.91 & v2x_cspart <= 0.96 & v2x...
      0.360947
      0.647929
      0.647929
      1.0
      8
      0.384615
      0.058861

## BRL

```python
method = 'brl'
need_pruning = False
relu = False
results_brl = c_to_y(method, need_pruning, relu, verbose=True)
results_brl
```

# Summary

```python
import glob

cols = ['model_accuracy_top', '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]

df_list = []
files = glob.glob(f'{results_dir}/**/results**.csv', recursive = True)
for file in files:
    df = pd.read_csv(file, index_col=0)
    method = df['method'][0]
    df_mean = df[cols].mean()
    df_sem = df[cols].sem()
    df_mean.columns = mean_cols
    df_sem.columns = sem_cols
    summary = pd.concat([df_mean, df_sem])
    summary.name = method
    df_list.append(summary)

summary = pd.concat(df_list, axis=1).T
summary.columns = mean_cols + sem_cols
summary
```

      model_accuracy_top_mean
      explanation_accuracy_mean
      explanation_fidelity_mean
      explanation_complexity_mean
      elapsed_time_mean
      explanation_consistency_mean
      model_accuracy_top_sem
      explanation_accuracy_sem
      explanation_fidelity_sem
      explanation_complexity_sem
      elapsed_time_sem
      explanation_consistency_sem

      brl
      0.922293
      0.946819
      1.000000
      8.7
      53.689056
      0.762500
      0.003764
      0.003217
      0.000000
      0.966667
      0.406669
      0.000000e+00

      pruning
      0.923477
      0.946819
      0.951254
      5.0
      31.864813
      1.000000
      0.003091
      0.003217
      0.003742
      0.000000
      0.265803
      0.000000e+00

      psi
      0.792268
      0.802945
      0.824543
      2.0
      20.831486
      0.158333
      0.021982
      0.029882
      0.017960
      0.494413
      0.042169
      9.251859e-18

      tree-3
      0.685066
      0.685066
      1.000000
      8.0
      0.049328
      0.384615
      0.015131
      0.015131
      0.000000
      0.000000
      0.002012
      0.000000e+00

      tree
      0.945933
      0.945933
      1.000000
      294.1
      0.076102
      0.368571
      0.003127
      0.003127
      0.000000
      7.878875
      0.002043
      1.850372e-17

      weights
      0.918751
      0.946524
      0.946230
      4.9
      34.242483
      0.980000
      0.003737
      0.003229
      0.003658
      0.100000
      0.323599
      3.700743e-17

```python
summary.to_csv(os.path.join(results_dir, 'summary.csv'))
```


# Learning Disentangled Representations using sequential images of a teapot

```python
from mpl_toolkits import mplot3d
import matplotlib.pyplot as plt
import random
import numpy as np
from PIL import Image
from tqdm import tqdm
import os
```

### Create dataset

Code to generate this dataset borrows from
https://medium.com/@yzhong.cs/beyond-data-scientist-3d-plots-in-python-with-examples-2a8bd7aa654b

The following cells generate sequential images of a teapot, where at each step an action, corresponding to either a rotation in viewpoint or a change in colour, was performed to generate the next image.

```python
CREATE_DATASET = True
```

First load teapot object, which is borrowed from the Stanford Computer Graphics Lab at https://graphics.stanford.edu/courses/cs148-10-summer/as3/code/as3/teapot.obj

```python
# Load teapot.obj

if CREATE_DATASET:

    import numpy as np
    def read_obj(filename):
        triangles = []
        vertices = []
        with open(filename) as file:
            for line in file:
                components = line.strip(' \n').split(' ')
                if components[0] == "f": # face data
                    # e.g. "f 1/1/1/ 2/2/2 3/3/3 4/4/4 ..."
                    indices = list(map(lambda c: int(c.split('/')[0]) - 1, components[1:]))
                    for i in range(0, len(indices) - 2):
                        triangles.append(indices[i: i+3])
                elif components[0] == "v": # vertex data
                    # e.g. "v  30.2180 89.5757 -76.8089"
                    vertex = list(map(lambda c: float(c), components[1:]))
                    vertices.append(vertex)
        return np.array(vertices), np.array(triangles)
```

Now generate a dataset, consisting of 1000 sequential images

```python
if CREATE_DATASET:

    N_DATA = 1000

    folder = 'teapot/'

    vertices, triangles = read_obj(folder+'teapot.obj')
    angle = 2*np.pi / 5
    colors = [[0,0,0],[255,0,0],[255,255,255],[0,255,0],[0,0,255]]
    color_index = 0
    actions = []

    for i in tqdm(range(N_DATA)):

        # First, plot 3D image of a teapot and save as image

        x = np.asarray(vertices[:,0]).squeeze()
        y = np.asarray(vertices[:,1]).squeeze()
        z = np.asarray(vertices[:,2]).squeeze()

        fig = plt.figure()
        ax = fig.add_subplot(111, projection='3d')
        ax.grid(None)
        ax.axis('off')
        ax.set_xlim([-3, 3])
        ax.set_ylim([-3, 3])
        ax.set_zlim([0, 3])
        ax.plot_trisurf(x, z, triangles, y, shade=True, color='white')
        ax.view_init(100, angle)
        plt.savefig(folder+'teapot'+str(i)+'.png')
        plt.close()

        # Then load the image, crop, resize it, and change background color

        img = Image.open('teapot/teapot'+str(i)+'.png').convert('RGB')
        img = img.crop((100,0,350,258))
        img = img.resize((84,84))
        arr = np.array(img)
        arr = np.where(arr == [255,255,255], colors[color_index], arr)
        np.save(folder+'small_teapot/teapot'+str(i),arr)

        # Now select an action to perform that changes the scene.

        action = random.randrange(8)

        if action == 0: # y rotation, positive
            m = np.matrix([[np.cos(angle), 0, np.sin(angle)],[0,1,0],[-np.sin(angle), 0, np.cos(angle)]])
        elif action == 1: # y rotation, negative
            m = np.matrix([[np.cos(angle), 0, -np.sin(angle)],[0,1,0],[np.sin(angle), 0, np.cos(angle)]])
        elif action == 2: # z rotation, positive
            m = np.matrix([[1,0,0],[0, np.cos(angle), np.sin(angle)],[0, -np.sin(angle), np.cos(angle)]])
        elif action == 3: # z rotation, positive
            m = np.matrix([[1,0,0],[0, np.cos(angle), -np.sin(angle)],[0, np.sin(angle), np.cos(angle)]])
        elif action == 4: # x rotation, positive
            m = np.matrix([[np.cos(angle), np.sin(angle), 0],[-np.sin(angle), np.cos(angle), 0],[0,0,1]])
        elif action == 5: # x rotation, positive
            m = np.matrix([[np.cos(angle), -np.sin(angle), 0],[np.sin(angle), np.cos(angle), 0],[0,0,1]])
        elif action ==6: # Change color by +1 increment
            m = np.matrix([[1, 0, 0],[0,1,0],[0,0,1]])
            color_index = (color_index + 1) % 5
        elif action ==7: # Change color by -1 increment
            m = np.matrix([[1, 0, 0],[0,1,0],[0,0,1]])
            color_index = (color_index - 1) % 5

        actions.append(action)

        # Change viewpoint of teapot
        vertices = vertices*m

    # Save action sequence
    np.save(folder+'actions',actions)
```

### Show a sample from this dataset

```python
plt.imshow(np.load('teapot/small_teapot/teapot290.npy'))
```

# Learning Disentangled Representations

Note: This works faster using a GPU

```python
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
```

```python
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
```

### First, define a TeapotWorld class that can load sequences of states and actions from our dataset

```python
class TeapotWorld():

    class action_space():
        def __init__(self,n_actions):
            self.n = n_actions

    class observation_space():
        def __init__(self,shape):
            self.shape = shape

    def __init__(self, action_file, data_dir):

        self.action_space = self.action_space(8)
        self.observation_space = self.observation_space([84,84,3])

        self.actions = np.load(action_file)
        self.data_dir = data_dir
        self.n_data = len(self.actions)
        self.dataset = self.load_dataset()

        self.current_idx = 0

        self.reset()

    def load_dataset(self):
        data = []
        for idx in tqdm(range(self.n_data)):
            data.append(self.load_image(idx))
        return torch.stack(data)

    def reset(self):
        self.current_idx = random.randrange(self.n_data-10)
        return self.get_observation()

    def load_image(self, idx):
        data_file = 'teapot'+str(idx)+'.npy'
        obs = np.load(self.data_dir+data_file)
        return torch.FloatTensor(obs/255)

    def get_observation(self, idx=None):
        if idx == None:
            idx = self.current_idx
        obs = self.dataset[idx]
        return obs.to(device)

    def get_action(self):
        return self.actions[self.current_idx]

    def step(self):
        self.current_idx += 1
        return self.get_observation()

    def get_batch(self, batch_size):
        idx = random.sample(range(self.n_data),batch_size)
        batch = []
        for i in idx:
            batch.append(self.get_observation(idx=i))
        batch = torch.stack(batch)
        return batch.to(device)
```

**Check that TeapotWorld loads correctly**

```python
env = TeapotWorld('teapot/actions.npy','teapot/small_teapot/')
```

    100%|██████████| 1000/1000 [00:01<00:00, 676.83it/s]

### Define encoder and decoder.

These are based on the neural networks used in the "Human-level control through deep reinforcement learning" DQN paper, since the observations have pretty much the same dimensions

```python
def init_weights(m, gain):
    if (type(m) == nn.Linear) | (type(m) == nn.Conv2d) | (type(m) == nn.ConvTranspose2d):
        nn.init.orthogonal_(m.weight, gain)
        nn.init.zeros_(m.bias)

class Encoder(nn.Module):

    def __init__(self, n_out=5, n_hid = 128, weight_scale=5):

        super().__init__()

        self.conv = nn.Sequential(nn.Conv2d(3, 16, 8, stride=4),
                                  nn.ReLU(),
                                  nn.Conv2d(16, 32, 4, stride=2),
                                  nn.ReLU(),
                                  nn.Conv2d(32, 32, 3, stride=1),
                                  nn.ReLU())

        self.output = nn.Sequential(nn.Linear(32 * 7 * 7, n_hid),
                                  nn.ReLU(),
                                  nn.Linear(n_hid, n_out))

        self.conv.apply(lambda x: init_weights(x, weight_scale))
        self.output.apply(lambda x: init_weights(x, weight_scale))

    def forward(self, obs):
        if len(obs.shape) != 4:
            obs = obs.unsqueeze(0)
        obs = obs.permute(0, 3, 1, 2)
        obs = obs/255
        obs = self.conv(obs)
        obs = obs.contiguous().view(obs.size(0), -1)
        return F.normalize(self.output(obs)).squeeze()

class Decoder(nn.Module):

    def __init__(self, n_in=5, n_hid = 128, weight_scale=5):

        super().__init__()

        self.fc1 = nn.Linear(n_in, n_hid)
        self.fc2 = nn.Linear(n_hid, 32 * 7 * 7)

        self.conv = nn.Sequential(nn.ConvTranspose2d(32, 32, 3, stride=1),
                                  nn.ReLU(),
                                  nn.ConvTranspose2d(32, 16, 4, stride=2),
                                  nn.ReLU(),
                                  nn.ConvTranspose2d(16, 3, 8, stride=4),
                                  )

        self.conv.apply(lambda x: init_weights(x, weight_scale))
        init_weights(self.fc1, weight_scale)
        init_weights(self.fc2, weight_scale)

    def forward(self, x):
        if len(x.shape) == 1:
            x = x.unsqueeze(0)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        batch_size = x.shape[0]
        x = x.reshape(batch_size,32,7,7)
        x = self.conv(x).permute(0,2,3,1)
        return torch.sigmoid(x).squeeze()
```

**Check dimensions**

```python
encoder = Encoder(n_out=5).to(device)
decoder = Decoder(n_in=5).to(device)
print(encoder)
print(decoder)
```

    Encoder(
      (conv): Sequential(
        (0): Conv2d(3, 16, kernel_size=(8, 8), stride=(4, 4))
        (1): ReLU()
        (2): Conv2d(16, 32, kernel_size=(4, 4), stride=(2, 2))
        (3): ReLU()
        (4): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1))
        (5): ReLU()
      )
      (output): Sequential(
        (0): Linear(in_features=1568, out_features=128, bias=True)
        (1): ReLU()
        (2): Linear(in_features=128, out_features=5, bias=True)
      )
    )
    Decoder(
      (fc1): Linear(in_features=5, out_features=128, bias=True)
      (fc2): Linear(in_features=128, out_features=1568, bias=True)
      (conv): Sequential(
        (0): ConvTranspose2d(32, 32, kernel_size=(3, 3), stride=(1, 1))
        (1): ReLU()
        (2): ConvTranspose2d(32, 16, kernel_size=(4, 4), stride=(2, 2))
        (3): ReLU()
        (4): ConvTranspose2d(16, 3, kernel_size=(8, 8), stride=(4, 4))
      )
    )

```python
obs = env.reset()
print(obs.shape)
latent = encoder(obs)
print(latent.shape)
reconstructed = decoder(latent)
print(reconstructed.shape)
```

    torch.Size([84, 84, 3])
    torch.Size([5])
    torch.Size([84, 84, 3])

**Representation**

The crux of the matter is learning to 'represent' actions in the observation space with actions in latent space.  Here, we will do this by assuming every action is a generalized rotation in latent space, which we denote with a series of 2-dimensional rotations.

A 2-d rotation is given by:

\begin{pmatrix}
\cos(\theta) & \sin(\theta) \\
-\sin(\theta) & \cos(\theta)
\end{pmatrix}

and we denote a rotation in dimensions $i$ and $j$ of a higher dimensional space as $R_{i,j}(\theta)$.  For $i=1$, $j=4$, in a 4-dimensional space:

\begin{equation}
R_{1,4}(\theta) =
\begin{pmatrix}
\cos(\theta) & 0 & 0 & \sin(\theta) \\
0 & 1 & 0 & 0 \\
0 & 0 & 1 & 0 \\
-\sin(\theta) & 0 & 0 & \cos(\theta)
\end{pmatrix}
\end{equation}

An arbitrary rotation, denoted $g$ as I am subtly moving towards this being a group action, can then be written as:

\begin{equation}
    g(\theta_{1,2},\theta_{1,3},\dots,\theta_{n-1,n}) = \prod_{i=1}^{n-1} \prod_{j=1+1}^{n} R_{i,j}(\theta_{i,j})
\end{equation}

which has $n(n-1)/2$ free parameters (i.e. $\theta_{i,j}$'s).

```python
class Representation():

    def __init__(self, dim=5):
        self.dim = dim
        self.params = dim*(dim-1)//2
        self.thetas = torch.FloatTensor(2*torch.rand(self.params)-1).to(device).requires_grad_()

        self.__matrix = None

    def set_thetas(self, thetas):
        self.thetas = thetas
        self.thetas.requires_grad = True
        self.clear_matrix()

    def clear_matrix(self):
        self.__matrix = None

    def get_matrix(self):
        if self.__matrix is None:
            k = 0
            mats = []
            for i in range(self.dim-1):
                for j in range(self.dim-1-i):
                    theta_ij = self.thetas[k]
                    k+=1
                    c, s = torch.cos(theta_ij), torch.sin(theta_ij)

                    rotation_i = torch.eye(self.dim, self.dim).to(device)
                    rotation_i[i, i] = c
                    rotation_i[i, i+j+1] = s
                    rotation_i[j+i+1, i] = -s
                    rotation_i[j+i+1, j+i+1] = c

                    mats.append(rotation_i)

            def chain_mult(l):
                if len(l)>=3:
                    return l[0]@l[1]@chain_mult(l[2:])
                elif len(l)==2:
                    return l[0]@l[1]
                else:
                    return l[0]

            self.__matrix = chain_mult(mats)

        return self.__matrix
```

**LatentWorld**

Now, for symmetry's sake, we'll also have a `LatentWorld` which acts as the environment in the latent space.

```python
class LatentWorld():

    class action_space():
        def __init__(self,n_actions):
            self.n = n_actions

        def sample(self, k=1):
            return torch.randint(0,self.n,(k,))

    class observation_space():
        def __init__(self,n_features):
            self.shape = [n_features]

    def __init__(self,
                 dim=5,
                 n_actions=8,
                 action_reps=None):

        self.dim = dim

        self.action_space = self.action_space(n_actions)
        self.observation_space = self.observation_space(dim)

        if action_reps is None:
            self.action_reps = [Representation(dim=self.dim) for _ in range(n_actions)]
        else:
            if len(action_reps)!=n_actions:
                raise Exception("Must pass an action representation for every action.")
            if not all([rep.dim==self.dim]):
                raise Exception("Action representations do not act on the dimension of the latent space.")
            self.action_reps = action_reps

    def reset(self, state):
        self.state = state
        return self.get_observation()

    def clear_representations(self):
        for rep in self.action_reps:
            rep.clear_matrix()

    def get_representation_params(self):
        params = []
        for rep in self.action_reps:
            params.append(rep.thetas)
        return params

    def save_representations(self, path):
        if os.path.splitext(path)[-1] != '.pth':
            path += '.pth'
        rep_thetas = [rep.thetas for rep in self.action_reps]
        return torch.save(rep_thetas, path)

    def load_reprentations(self, path):
        rep_thetas = torch.load(path)
        for rep in self.action_reps:
            rep.set_thetas(rep_thetas.pop(0))

    def get_observation(self):
        return self.state

    def step(self,action):
        self.state = torch.mv(self.action_reps[action].get_matrix(), self.state.squeeze())
        obs = self.get_observation()
        return obs
```

**Entanglement regularisation**

So for $m$ parameters, ${\theta_1, \dots, \theta_m}$, we want to regularise with
\begin{equation}
    \sum_{i \neq j} \vert\theta_i\vert^2, \mathrm{where\ } \theta_j {=} \mathrm{max_k}({\vert\theta_k\vert}).
\end{equation}
We will also use this term as our metric of 'entanglement'.

```python
def calc_entanglement(params):
    params = params.abs().pow(2)
    return params.sum() - params.max()

params = torch.FloatTensor([1,1,0.5,0,0])
calc_entanglement(params)
```

    tensor(1.2500)

### Training with regularization

We find it helpful to increase the regularization strength halfway during training

```python
dim=5

obs_env = TeapotWorld('teapot/actions.npy','teapot/small_teapot/')
lat_env = LatentWorld(dim = dim,
                      n_actions = obs_env.action_space.n
                     )
decoder = Decoder(n_in = dim, n_hid = 128).to(device)
encoder = Encoder(n_out = dim, n_hid = 128).to(device)

optimizer_dec = optim.Adam(decoder.parameters(),
                           lr=1e-3,
#                            betas=(0.9, 0.99),
                           weight_decay=0)

optimizer_enc = optim.Adam(encoder.parameters(),
                           lr=1e-3,
#                            betas=(0.9, 0.99),
                           weight_decay=0)

optimizer_rep = optim.Adam(lat_env.get_representation_params(),
                           lr=1e-2,
#                            betas=(0.9, 0.99),
                           weight_decay=0)

losses = []
entanglement = []
orthogonality = []
```

    100%|██████████| 1000/1000 [00:00<00:00, 1754.53it/s]

```python
import time

n_sgd_steps = 10000
ep_steps = 5
batch_eps = 16

i = 1

t_start = time.time()

temp = 0

while i < n_sgd_steps:
    loss = torch.zeros(1).to(device)

    for _ in range(batch_eps):
        t_ep = -1
        while t_ep < ep_steps:
            if t_ep == -1:
                obs_x = obs_env.reset()
                obs_z = lat_env.reset(encoder(obs_x))
            else:
                action = obs_env.get_action()
                obs_x = obs_env.step()
                obs_z = lat_env.step(action)

            t_ep += 1

            obs_x_recon = decoder(obs_z)

            loss += F.binary_cross_entropy(obs_x_recon, obs_x)

    loss /= (ep_steps * batch_eps)
    loss_raw = loss

    reg_loss = sum([calc_entanglement(r.thetas) for r in lat_env.action_reps])/8
    if i < 5000:
        loss += reg_loss*1e-3
    else:
        loss += reg_loss*3e-2

    losses.append(loss_raw.item())
    entanglement.append(reg_loss.item())

    optimizer_dec.zero_grad()
    optimizer_rep.zero_grad()
    optimizer_enc.zero_grad()
    loss.backward()
    optimizer_dec.step()
    optimizer_rep.step()
    optimizer_enc.step()

    # Rember to clear the cached action representations after we update the parameters!
    lat_env.clear_representations()

    i+=1

    if i%10==0:
        print("iter {} : loss={:.3e} : entanglement={:.2e} : last 10 iters in {:.3f}s".format(
            i, loss.item(), reg_loss.item(), time.time() - t_start
            ), end="\r" if i%100 else "\n")
        t_start = time.time()
```

    iter 100 : loss=8.301e-01 : entanglement=1.50e+00 : last 10 iters in 3.398s
    iter 200 : loss=8.234e-01 : entanglement=1.10e+00 : last 10 iters in 3.412s
    iter 300 : loss=7.464e-01 : entanglement=7.17e-01 : last 10 iters in 3.326s
    iter 400 : loss=6.625e-01 : entanglement=5.25e-01 : last 10 iters in 3.296s
    iter 500 : loss=6.165e-01 : entanglement=4.72e-01 : last 10 iters in 3.372s
    iter 600 : loss=4.312e-01 : entanglement=4.46e-01 : last 10 iters in 3.389s
    iter 700 : loss=3.550e-01 : entanglement=4.39e-01 : last 10 iters in 3.328s
    iter 800 : loss=2.399e-01 : entanglement=4.42e-01 : last 10 iters in 3.372s
    iter 900 : loss=1.676e-01 : entanglement=4.29e-01 : last 10 iters in 3.408s
    iter 1000 : loss=1.614e-01 : entanglement=5.21e-01 : last 10 iters in 3.788s
    iter 1100 : loss=1.492e-01 : entanglement=7.08e-01 : last 10 iters in 3.340s
    iter 1200 : loss=1.425e-01 : entanglement=7.89e-01 : last 10 iters in 3.372s
    iter 1300 : loss=1.357e-01 : entanglement=8.48e-01 : last 10 iters in 3.384s
    iter 1400 : loss=1.248e-01 : entanglement=8.88e-01 : last 10 iters in 3.416s
    iter 1500 : loss=1.295e-01 : entanglement=8.86e-01 : last 10 iters in 3.374s
    iter 1600 : loss=1.296e-01 : entanglement=8.93e-01 : last 10 iters in 3.390s
    iter 1700 : loss=1.212e-01 : entanglement=8.90e-01 : last 10 iters in 3.384s
    iter 1800 : loss=1.178e-01 : entanglement=8.89e-01 : last 10 iters in 3.327s
    iter 1900 : loss=1.226e-01 : entanglement=8.66e-01 : last 10 iters in 3.345s
    iter 2000 : loss=1.226e-01 : entanglement=8.80e-01 : last 10 iters in 3.365s
    iter 2100 : loss=1.156e-01 : entanglement=8.64e-01 : last 10 iters in 3.392s
    iter 2200 : loss=1.142e-01 : entanglement=8.59e-01 : last 10 iters in 3.336s
    iter 2300 : loss=1.164e-01 : entanglement=8.46e-01 : last 10 iters in 3.398s
    iter 2400 : loss=1.144e-01 : entanglement=8.43e-01 : last 10 iters in 3.356s
    iter 2500 : loss=1.121e-01 : entanglement=8.32e-01 : last 10 iters in 3.356s
    iter 2600 : loss=1.132e-01 : entanglement=8.24e-01 : last 10 iters in 3.344s
    iter 2700 : loss=1.144e-01 : entanglement=8.18e-01 : last 10 iters in 3.440s
    iter 2800 : loss=1.130e-01 : entanglement=8.05e-01 : last 10 iters in 3.405s
    iter 2900 : loss=1.130e-01 : entanglement=8.01e-01 : last 10 iters in 3.392s
    iter 3000 : loss=1.140e-01 : entanglement=7.92e-01 : last 10 iters in 3.428s
    iter 3100 : loss=1.165e-01 : entanglement=7.86e-01 : last 10 iters in 3.373s
    iter 3200 : loss=1.119e-01 : entanglement=7.82e-01 : last 10 iters in 3.740s
    iter 3300 : loss=1.125e-01 : entanglement=7.73e-01 : last 10 iters in 3.768s
    iter 3400 : loss=1.111e-01 : entanglement=7.62e-01 : last 10 iters in 3.769s
    iter 3500 : loss=1.106e-01 : entanglement=7.54e-01 : last 10 iters in 3.817s
    iter 3600 : loss=1.067e-01 : entanglement=7.44e-01 : last 10 iters in 3.804s
    iter 3700 : loss=1.155e-01 : entanglement=7.35e-01 : last 10 iters in 3.772s
    iter 3800 : loss=1.100e-01 : entanglement=7.31e-01 : last 10 iters in 3.788s
    iter 3900 : loss=1.091e-01 : entanglement=7.24e-01 : last 10 iters in 3.884s
    iter 4000 : loss=1.132e-01 : entanglement=7.20e-01 : last 10 iters in 3.892s
    iter 4100 : loss=1.092e-01 : entanglement=7.18e-01 : last 10 iters in 3.855s
    iter 4200 : loss=1.080e-01 : entanglement=7.05e-01 : last 10 iters in 3.933s
    iter 4300 : loss=1.078e-01 : entanglement=7.01e-01 : last 10 iters in 3.920s
    iter 4400 : loss=1.101e-01 : entanglement=7.02e-01 : last 10 iters in 3.965s
    iter 4500 : loss=1.109e-01 : entanglement=6.86e-01 : last 10 iters in 3.944s
    iter 4600 : loss=1.099e-01 : entanglement=6.84e-01 : last 10 iters in 3.988s
    iter 4700 : loss=1.076e-01 : entanglement=6.73e-01 : last 10 iters in 3.995s
    iter 4800 : loss=1.073e-01 : entanglement=6.74e-01 : last 10 iters in 4.090s
    iter 4900 : loss=1.042e-01 : entanglement=6.74e-01 : last 10 iters in 4.064s
    iter 5000 : loss=1.099e-01 : entanglement=6.69e-01 : last 10 iters in 4.104s
    iter 5100 : loss=1.503e-01 : entanglement=3.34e-01 : last 10 iters in 3.988s
    iter 5200 : loss=1.384e-01 : entanglement=2.47e-01 : last 10 iters in 3.882s
    iter 5300 : loss=1.256e-01 : entanglement=1.54e-01 : last 10 iters in 3.868s
    iter 5400 : loss=1.206e-01 : entanglement=8.89e-02 : last 10 iters in 4.004s
    iter 5500 : loss=1.173e-01 : entanglement=4.64e-02 : last 10 iters in 3.896s
    iter 5600 : loss=1.177e-01 : entanglement=2.57e-02 : last 10 iters in 3.940s
    iter 5700 : loss=1.125e-01 : entanglement=1.45e-02 : last 10 iters in 3.963s
    iter 5800 : loss=1.116e-01 : entanglement=8.64e-03 : last 10 iters in 3.824s
    iter 5900 : loss=1.104e-01 : entanglement=4.99e-03 : last 10 iters in 3.836s
    iter 6000 : loss=1.097e-01 : entanglement=3.90e-03 : last 10 iters in 3.912s
    iter 6100 : loss=1.103e-01 : entanglement=2.56e-03 : last 10 iters in 3.904s
    iter 6200 : loss=1.065e-01 : entanglement=1.83e-03 : last 10 iters in 3.884s
    iter 6300 : loss=1.079e-01 : entanglement=1.38e-03 : last 10 iters in 3.880s
    iter 6400 : loss=1.089e-01 : entanglement=1.41e-03 : last 10 iters in 3.888s
    iter 6500 : loss=1.066e-01 : entanglement=9.59e-04 : last 10 iters in 3.941s
    iter 6600 : loss=1.051e-01 : entanglement=9.38e-04 : last 10 iters in 3.915s
    iter 6700 : loss=1.085e-01 : entanglement=7.41e-04 : last 10 iters in 3.960s
    iter 6800 : loss=1.065e-01 : entanglement=6.89e-04 : last 10 iters in 3.920s
    iter 6900 : loss=1.074e-01 : entanglement=8.87e-04 : last 10 iters in 4.000s
    iter 7000 : loss=1.072e-01 : entanglement=6.43e-04 : last 10 iters in 4.032s
    iter 7100 : loss=1.077e-01 : entanglement=6.19e-04 : last 10 iters in 4.036s
    iter 7200 : loss=1.054e-01 : entanglement=5.41e-04 : last 10 iters in 4.092s
    iter 7300 : loss=1.081e-01 : entanglement=5.11e-04 : last 10 iters in 4.076s
    iter 7400 : loss=1.095e-01 : entanglement=5.76e-04 : last 10 iters in 4.067s
    iter 7500 : loss=1.025e-01 : entanglement=5.35e-04 : last 10 iters in 4.107s
    iter 7600 : loss=1.071e-01 : entanglement=4.75e-04 : last 10 iters in 4.160s
    iter 7700 : loss=1.068e-01 : entanglement=4.39e-04 : last 10 iters in 4.124s
    iter 7800 : loss=1.032e-01 : entanglement=4.82e-04 : last 10 iters in 4.155s
    iter 7900 : loss=1.051e-01 : entanglement=5.41e-04 : last 10 iters in 4.195s
    iter 8000 : loss=1.053e-01 : entanglement=5.58e-04 : last 10 iters in 4.228s
    iter 8100 : loss=1.063e-01 : entanglement=5.11e-04 : last 10 iters in 4.224s
    iter 8200 : loss=1.064e-01 : entanglement=5.91e-04 : last 10 iters in 4.285s
    iter 8300 : loss=1.036e-01 : entanglement=5.89e-04 : last 10 iters in 4.201s
    iter 8400 : loss=1.065e-01 : entanglement=5.15e-04 : last 10 iters in 4.307s
    iter 8500 : loss=1.041e-01 : entanglement=6.81e-04 : last 10 iters in 4.344s
    iter 8600 : loss=1.039e-01 : entanglement=4.47e-04 : last 10 iters in 4.293s
    iter 8700 : loss=1.035e-01 : entanglement=7.87e-04 : last 10 iters in 4.317s
    iter 8800 : loss=1.047e-01 : entanglement=8.44e-04 : last 10 iters in 4.324s
    iter 8900 : loss=1.053e-01 : entanglement=7.56e-04 : last 10 iters in 4.372s
    iter 9000 : loss=1.030e-01 : entanglement=6.96e-04 : last 10 iters in 4.400s
    iter 9100 : loss=1.038e-01 : entanglement=1.03e-03 : last 10 iters in 4.399s
    iter 9200 : loss=1.041e-01 : entanglement=7.93e-04 : last 10 iters in 4.428s
    iter 9300 : loss=9.924e-02 : entanglement=1.13e-03 : last 10 iters in 4.428s
    iter 9400 : loss=1.006e-01 : entanglement=6.01e-04 : last 10 iters in 4.451s
    iter 9500 : loss=1.023e-01 : entanglement=7.11e-04 : last 10 iters in 4.428s
    iter 9600 : loss=1.001e-01 : entanglement=5.78e-04 : last 10 iters in 4.449s
    iter 9700 : loss=1.041e-01 : entanglement=7.96e-04 : last 10 iters in 4.540s
    iter 9800 : loss=1.036e-01 : entanglement=9.06e-04 : last 10 iters in 4.548s
    iter 9900 : loss=1.031e-01 : entanglement=7.36e-04 : last 10 iters in 4.528s
    iter 10000 : loss=1.048e-01 : entanglement=5.82e-04 : last 10 iters in 4.440s

### Show reconstructed states

```python
obs = env.reset()
latent = encoder(obs)
reconstructed = decoder(latent)

fig, (ax1,ax2) = plt.subplots(1, 2)
ax1.imshow(obs.to('cpu'))
ax2.imshow(reconstructed.detach().to('cpu'))
```

### Show action representations

```python
width=0.5

rep_thetas = [rep.thetas.detach().to('cpu').numpy() for rep in lat_env.action_reps]
#print(rep_thetas)
plt_lim = max( 0.22, max([max(t) for t in rep_thetas])/(2*np.pi) )
titles = ["transformation +", "transformation -"]
cols=["r","b","g","black"]
labels=["x","y","z","color"]

with plt.style.context('seaborn-paper', after_reset=True):

    fig, axs = plt.subplots(1, 2, figsize=(12, 4), gridspec_kw={"wspace":0.5, "hspace":0.5})

    for i, thetas in enumerate(rep_thetas):
        x = np.arange(len(thetas))
        axs[i%2].bar(x - width/2, thetas/(2*np.pi), width, label=labels[i//2], color=cols[i//2])
    for i in range(2):
        axs[i].hlines([0.2,-0.2], xmin=-10, xmax=10, linestyles="dashed")
        axs[i].hlines(0., xmin=-10, xmax=10)
        axs[i].set_yticks([-0.2, 0., 0.2])
        axs[i].set_xticks(x-0.25)
        axs[i].set_xticklabels(["12","13","14","15","23","24","25","34","35","45"], fontsize=10)
        axs[i].set_xlabel("$ij$", fontsize=20)

        axs[i].set_ylim(-plt_lim,plt_lim)
        axs[i].set_xlim(-0.5,9)
        axs[i].set_title(titles[i], fontsize=20)
        axs[i].tick_params(labelsize=15)

        axs[i].set_ylabel(r"$\theta / 2\pi$", fontsize=20)

    axs[1].legend(loc="center right", bbox_to_anchor=(1.5,0.5), fontsize=15)

plt.savefig("teapot/representations.png", bbox_inches='tight')
```

### For the paper, generate sequence of images from the dataset

```python
import matplotlib.gridspec as gridspec
from IPython import display

plt.figure(figsize = (6,6))
gs1 = gridspec.GridSpec(3, 3)
gs1.update(wspace=0.02, hspace=0.02)
plt.grid(None)
state = env.reset()
for i in range(9):
    ax = plt.subplot(gs1[i])
    ax.axis('off')
    ax.set_aspect('equal')
    ax.imshow(state.to('cpu'))
    display.display(plt.gcf())
    time.sleep(0.2)
    display.clear_output(wait=True)
    state = env.step()

plt.savefig("teapot/env.png", bbox_inches='tight')
```

```python

```