<jupyter_start><jupyter_text>Load SUPPORT dataset- fill values are found in the table belowBaseline Variable | Column Name | Normal Fill-in Value--- | --- | ---Serum albumin | alb | 3.5PaO2/FiO2 ratio (pafi) | pafi | 333.3Bilirubin | bili | 1.01Creatinine | crea | 1.01BUN | bun | 6.51White blood count | wblc | 9 (thousands)Urine output | urine | 2502<jupyter_code>import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

import sys
sys.path.append('../src')
from models import DNMC, NMC, NSurv, MLP, discrete_ci

import tensorflow as tf
from tensorflow import keras

tf.random.set_seed(2021)
FILL_VALUES = {
    'alb': 3.5,
    'pafi': 333.3,
    'bili': 1.01,
    'crea': 1.01,
    'bun': 6.51,
    'wblc': 9.,
    'urine': 2502.
}

TO_DROP = ['aps', 'sps', 'surv2m', 'surv6m', 'prg2m', 'prg6m', 'dnr', 'dnrday']
TO_DROP = TO_DROP + ['sfdm2', 'hospdead']

# load, drop columns, fill using specified fill values
df = pd.read_csv('../datasets/support2.csv').drop(TO_DROP,axis=1).fillna(value=FILL_VALUES)

# get dummies for categorical vars
df = pd.get_dummies(df, dummy_na=True)

# fill remaining values to the median

df = df.fillna(df.median())

# standardize numeric columns

numrc_cols = df.dtypes == 'float64'
df.loc[:, numrc_cols] = (df.loc[:, numrc_cols] - df.loc[:, numrc_cols].mean()) / df.loc[:, numrc_cols].std()

df
df['death'].value_counts()
OUTCOMES = ['death', 'd.time']
X = df.drop(OUTCOMES, axis=1).sample(frac=1, random_state=2021)
X = X.values<jupyter_output><empty_output><jupyter_text>We create the semi-synthetic dataset as follows- real X- make up new failure times- make up new failure probabilities- make up new censoring times<jupyter_code># GLOBAL PARAMETERS

N_FEATURES = 10
N_BINS = 10
LEARNING_RATE = 6e-4
N_OVERLAPPING = 5

rs = np.random.RandomState(8)<jupyter_output><empty_output><jupyter_text>Separate Factors- randomly select separate factors for failure times vs probabilities (no overlap)- censoring times are random uniform<jupyter_code>def discretize(*arrs, nbins=10, pad=1e-5):
    min_val = np.amin(np.stack(arrs))
    max_val = np.amax(np.stack(arrs))
    binsize = (max_val - min_val + pad) / nbins
    return ((arr - arr.min()) // binsize for arr in arrs)
    
def onehot(arr, ncategories=None):
    if ncategories is None:
        ncategories = len(np.unique(arr))
    return np.eye(ncategories)[arr.astype(int)]
feature_order = rs.permutation(len(X.T))
T_features = feature_order[:N_FEATURES]
E_features = feature_order[N_FEATURES:(2 * N_FEATURES)]

T_weights = rs.randn(len(T_features))
E_weights = rs.randn(len(E_features))

T = X[:, T_features] @ T_weights
E = X[:, E_features] @ E_weights

t_cont = np.exp((T - np.mean(T)) / (2 * np.std(T)))
c_cont = rs.rand(len(t_cont)) * t_cont.max()

e_prob = 1 - 1 / (1 + np.exp(4. * (E - E.mean()) / E.std()))
e = (rs.rand(len(e_prob)) < e_prob).astype(int)

ct_cont = np.stack([c_cont, t_cont]).T
y_cont = e * np.amin(ct_cont, axis=1) + (1 - e) * c_cont

s = e * np.argmin(ct_cont, axis=1)
y_disc, t_disc, c_disc = discretize(y_cont, t_cont, c_cont, nbins=N_BINS)

y = onehot(y_disc, ncategories=N_BINS)
print('How many are susceptible?')
print(pd.value_counts(e).sort_index())

print('How many are observed?')
print(pd.value_counts(s).sort_index())

fig, ax = plt.subplots(nrows=1, ncols=4, figsize=(16, 4))

ax[0].hist(e_prob)
ax[1].hist(t_disc)
ax[2].hist(c_disc)
ax[3].hist(y_disc)

plt.show()
x_train, y_train, s_train = X[:6000], y[:6000], s[:6000]
x_val, y_val, s_val = X[6000:7500], y[6000:7500], s[6000:7500]
x_test, y_test, s_test, e_test = X[7500:], y[7500:], s[7500:], e[7500:]
t_disc_test, c_disc_test, y_disc_test = t_disc[7500:], c_disc[7500:], y_disc[7500:]
def get_batches(*arrs, batch_size=1):
    l = len(arrs[0])
    for ndx in range(0, l, batch_size):
        yield (arr[ndx:min(ndx + batch_size, l)] for arr in arrs)
        
import time

def train_model(
    model, train_data, val_data, n_epochs,
    batch_size=50, learning_rate=1e-3, early_stopping_criterion=2):

    optimizer = keras.optimizers.Adam(learning_rate=learning_rate)

    #@tf.function
    def train_step(x, y, s):
        with tf.GradientTape() as tape:
            train_loss, train_nll = model.loss(x, y, s)
            #print(train_loss, train_nll)
        grads = tape.gradient(train_loss, model.trainable_weights)
        optimizer.apply_gradients(zip(grads, model.trainable_weights))
        return train_loss, train_nll

    #@tf.function
    def test_step(x, y, s):
        val_loss, val_nll = model.loss(x, y, s)
        return val_loss, val_nll
    
    best_val_loss = np.inf
    no_decrease = 0

    for epoch_idx in range(n_epochs):

        #print("\nStart of epoch %d" % (epoch_idx,))
        start_time = time.time()

        train_losses = []
        train_nlls = []

        for batch_idx, (xt, yt, st) in enumerate(get_batches(*train_data, batch_size=batch_size)):

            train_loss, train_nll = train_step(xt, yt, st)

            train_losses.append(train_loss)
            train_nlls.append(train_nll)

        # Display metrics at the end of each epoch.
        #print('Epoch training loss: %.4f, NLL = %.4f' % (np.mean(batch_losses), np.mean(batch_nll)))

        val_losses = []
        val_nlls = []

        # Run a validation loop at the end of each epoch.
        for batch_idx, (xv, yv, sv) in enumerate(get_batches(*val_data, batch_size=batch_size)):

            val_loss, val_nll = test_step(xv, yv, sv)

            val_losses.append(val_loss)
            val_nlls.append(val_nll)
            
        new_val_loss = np.mean(val_losses)

        print(
            'Epoch %2i | Train Loss: %.4f | Train NLL: %.4f | Val Loss: %.4f | Val NLL: %.4f'
            % (epoch_idx, np.mean(train_losses), np.mean(train_nlls), np.mean(val_losses), np.mean(val_nlls))
        )
        #print('Time taken: %.2fs' % (time.time() - start_time))
                
        if new_val_loss > best_val_loss:
            no_decrease += 1
        else:
            no_decrease = 0
            best_val_loss = new_val_loss
            
        if no_decrease == early_stopping_criterion:
            break
from sklearn.metrics import roc_auc_score


def nll(t_true, t_pred, tol=1e-8):
    ncat = np.shape(t_pred)[1]
    nll_ = -1 * np.log(np.sum(onehot(t_true, ncategories=ncat) * t_pred, axis=1) + tol)
    return np.mean(nll_)


def evaluate_model(
    model, test_data, e_test, 
    batch_size=50, dataset='support', factors='unknown'):
    
    modelname = type(model).__name__
    
    test_losses = []
    test_nlls = []
    
    test_e_pred = []
    test_t_pred = []
    test_c_pred = []
    
    for batch_idx, (xt, yt, st) in enumerate(get_batches(*test_data, batch_size=batch_size)):
        
        test_loss, test_nll = model.loss(xt, yt, st)
        
        if modelname == 'NSurv':
            t_pred, c_pred = model(xt)
            test_t_pred.append(t_pred)
            test_c_pred.append(c_pred)
        elif modelname == 'MLP':
            e_pred = model(xt)
            test_e_pred.append(e_pred)
        else:
            e_pred, t_pred, c_pred = model(xt)
            test_e_pred.append(e_pred)
            test_t_pred.append(t_pred)
            test_c_pred.append(c_pred)
        
        test_losses.append(test_loss)
        test_nlls.append(test_nll)
    
    if modelname == 'NSurv':
        e_auc = None
    else:
        e_auc = roc_auc_score(e_test, np.concatenate(test_e_pred, axis=0))
        
    if modelname == 'MLP':
        ci = None
    else:
        ci = discrete_ci(test_data[2], test_data[1], np.concatenate(test_t_pred, axis=0))
    
    results = {
        'dataset': dataset,
        'factors': factors,
        'model': modelname,
        'ld': model.ld,
        'lr': model.lr,
        'avg_test_loss': np.mean(test_losses),
        'avg_test_nll': np.mean(test_nlls),
        'e_auc': e_auc,
        'y_ci': ci
    }
    
    return results
### BEGIN COLLECTING RESULTS HERE ###
all_results = []

pstrue = s_train.mean()
importance_weights = [1 / (1 - pstrue), 1 / pstrue]
model = DNMC(n_bins=N_BINS, importance_weights=importance_weights, ld=1e-2, lr=1e-2)
train_model(model, (x_train, y_train, s_train), (x_val, y_val, s_val), 100, learning_rate=LEARNING_RATE)
all_results.append(evaluate_model(model, (x_test, y_test, s_test), e_test, factors='separate'))
model = DNMC(n_bins=N_BINS, importance_weights=importance_weights, ld=3e-3, lr=1e-2)
train_model(model, (x_train, y_train, s_train), (x_val, y_val, s_val), 100, learning_rate=LEARNING_RATE)
all_results.append(evaluate_model(model, (x_test, y_test, s_test), e_test, factors='separate'))
model = DNMC(n_bins=N_BINS, importance_weights=importance_weights, ld=3e-2, lr=1e-2)
train_model(model, (x_train, y_train, s_train), (x_val, y_val, s_val), 100, learning_rate=LEARNING_RATE)
all_results.append(evaluate_model(model, (x_test, y_test, s_test), e_test, factors='separate'))
model = NMC(n_bins=N_BINS, importance_weights=importance_weights, lr=1e-2)
train_model(model, (x_train, y_train, s_train), (x_val, y_val, s_val), 100, learning_rate=LEARNING_RATE)
all_results.append(evaluate_model(model, (x_test, y_test, s_test), e_test, factors='separate'))
model = NSurv(n_bins=N_BINS, importance_weights=importance_weights, lr=1e-2)
train_model(model, (x_train, y_train, s_train), (x_val, y_val, s_val), 100, learning_rate=LEARNING_RATE)
all_results.append(evaluate_model(model, (x_test, y_test, s_test), e_test, factors='separate'))
model = MLP(importance_weights=importance_weights, lr=1e-2)
train_model(model, (x_train, y_train, s_train), (x_val, y_val, s_val), 100, learning_rate=LEARNING_RATE)
all_results.append(evaluate_model(model, (x_test, y_test, s_test), e_test, factors='separate'))<jupyter_output>Epoch  0 | Train Loss: 3.5809 | Train NLL: 1.2390 | Val Loss: 2.6526 | Val NLL: 1.0909
Epoch  1 | Train Loss: 2.3427 | Train NLL: 1.1014 | Val Loss: 2.0575 | Val NLL: 1.0434
Epoch  2 | Train Loss: 1.9639 | Train NLL: 1.0739 | Val Loss: 1.8125 | Val NLL: 1.0253
Epoch  3 | Train Loss: 1.7687 | Train NLL: 1.0528 | Val Loss: 1.6645 | Val NLL: 1.0119
Epoch  4 | Train Loss: 1.6443 | Train NLL: 1.0410 | Val Loss: 1.5665 | Val NLL: 1.0082
Epoch  5 | Train Loss: 1.5485 | Train NLL: 1.0267 | Val Loss: 1.4888 | Val NLL: 1.0010
Epoch  6 | Train Loss: 1.4698 | Train NLL: 1.0110 | Val Loss: 1.4278 | Val NLL: 0.9958
Epoch  7 | Train Loss: 1.4088 | Train NLL: 1.0002 | Val Loss: 1.3757 | Val NLL: 0.9893
Epoch  8 | Train Loss: 1.3615 | Train NLL: 0.9944 | Val Loss: 1.3345 | Val NLL: 0.9856
Epoch  9 | Train Loss: 1.3317 | Train NLL: 0.9991 | Val Loss: 1.3196 | Val NLL: 1.0021
Epoch 10 | Train Loss: 1.2856 | Train NLL: 0.9822 | Val Loss: 1.2736 | Val NLL: 0.9842
Epoch 11 | Train Loss: 1.2546 | Train NLL: [...]<jupyter_text>Overlapping Factors<jupyter_code>feature_order = rs.permutation(len(X.T))
T_features = feature_order[N_OVERLAPPING:(N_FEATURES + N_OVERLAPPING)]
E_features = feature_order[N_FEATURES:(2 * N_FEATURES)]

T_weights = rs.randn(len(T_features))
E_weights = rs.randn(len(E_features))

T = X[:, T_features] @ T_weights
E = X[:, E_features] @ E_weights

t_cont = np.exp((T - np.mean(T)) / (2 * np.std(T)))
c_cont = rs.rand(len(t_cont)) * t_cont.max()

e_prob = 1 - 1 / (1 + np.exp(4. * (E - E.mean()) / E.std()))
e = (rs.rand(len(e_prob)) < e_prob).astype(int)

ct_cont = np.stack([c_cont, t_cont]).T
y_cont = e * np.amin(ct_cont, axis=1) + (1 - e) * c_cont

s = e * np.argmin(ct_cont, axis=1)
y_disc, t_disc, c_disc = discretize(y_cont, t_cont, c_cont, nbins=N_BINS)

y = onehot(y_disc, ncategories=N_BINS)
pstrue = s_train.mean()
importance_weights = [1 / (1 - pstrue), 1 / pstrue]
print('How many are susceptible?')
print(pd.value_counts(e).sort_index())

print('How many are observed?')
print(pd.value_counts(s).sort_index())

fig, ax = plt.subplots(nrows=1, ncols=4, figsize=(16, 4))

ax[0].hist(e_prob)
ax[1].hist(t_disc)
ax[2].hist(c_disc)
ax[3].hist(y_disc)

plt.show()
x_train, y_train, s_train = X[:6000], y[:6000], s[:6000]
x_val, y_val, s_val = X[6000:7500], y[6000:7500], s[6000:7500]
x_test, y_test, s_test, e_test = X[7500:], y[7500:], s[7500:], e[7500:]
t_disc_test, c_disc_test, y_disc_test = t_disc[7500:], c_disc[7500:], y_disc[7500:]
x_train.shape, y_train.shape, s_train.shape, x_val.shape, y_val.shape, s_val.shape
model = DNMC(n_bins=N_BINS, importance_weights=importance_weights, ld=1e-2, lr=1e-2)
train_model(model, (x_train, y_train, s_train), (x_val, y_val, s_val), 100, learning_rate=LEARNING_RATE)
all_results.append(evaluate_model(model, (x_test, y_test, s_test), e_test, factors='overlapping'))
model = DNMC(n_bins=N_BINS, importance_weights=importance_weights, ld=3e-3, lr=1e-2)
train_model(model, (x_train, y_train, s_train), (x_val, y_val, s_val), 100, learning_rate=LEARNING_RATE)
all_results.append(evaluate_model(model, (x_test, y_test, s_test), e_test, factors='overlapping'))
model = DNMC(n_bins=N_BINS, importance_weights=importance_weights, ld=3e-2, lr=1e-2)
train_model(model, (x_train, y_train, s_train), (x_val, y_val, s_val), 100, learning_rate=LEARNING_RATE)
all_results.append(evaluate_model(model, (x_test, y_test, s_test), e_test, factors='overlapping'))
model = NMC(n_bins=N_BINS, importance_weights=importance_weights, lr=1e-2)
train_model(model, (x_train, y_train, s_train), (x_val, y_val, s_val), 100, learning_rate=LEARNING_RATE)
all_results.append(evaluate_model(model, (x_test, y_test, s_test), e_test, factors='overlapping'))
model = NSurv(n_bins=N_BINS, importance_weights=importance_weights, lr=1e-2)
train_model(model, (x_train, y_train, s_train), (x_val, y_val, s_val), 100, learning_rate=LEARNING_RATE)
all_results.append(evaluate_model(model, (x_test, y_test, s_test), e_test, factors='overlapping'))
model = MLP(importance_weights=importance_weights, lr=1e-2)
train_model(model, (x_train, y_train, s_train), (x_val, y_val, s_val), 100, learning_rate=LEARNING_RATE)
all_results.append(evaluate_model(model, (x_test, y_test, s_test), e_test, factors='overlapping'))<jupyter_output>Epoch  0 | Train Loss: 3.6507 | Train NLL: 1.3185 | Val Loss: 2.6634 | Val NLL: 1.1265
Epoch  1 | Train Loss: 2.3147 | Train NLL: 1.1033 | Val Loss: 2.0202 | Val NLL: 1.0370
Epoch  2 | Train Loss: 1.9211 | Train NLL: 1.0630 | Val Loss: 1.7862 | Val NLL: 1.0300
Epoch  3 | Train Loss: 1.7282 | Train NLL: 1.0430 | Val Loss: 1.6437 | Val NLL: 1.0208
Epoch  4 | Train Loss: 1.6012 | Train NLL: 1.0263 | Val Loss: 1.5365 | Val NLL: 1.0056
Epoch  5 | Train Loss: 1.5149 | Train NLL: 1.0192 | Val Loss: 1.4582 | Val NLL: 0.9955
Epoch  6 | Train Loss: 1.4493 | Train NLL: 1.0135 | Val Loss: 1.3977 | Val NLL: 0.9877
Epoch  7 | Train Loss: 1.3971 | Train NLL: 1.0084 | Val Loss: 1.3497 | Val NLL: 0.9816
Epoch  8 | Train Loss: 1.3550 | Train NLL: 1.0046 | Val Loss: 1.3109 | Val NLL: 0.9777
Epoch  9 | Train Loss: 1.3202 | Train NLL: 1.0014 | Val Loss: 1.2793 | Val NLL: 0.9746
Epoch 10 | Train Loss: 1.2895 | Train NLL: 0.9972 | Val Loss: 1.2516 | Val NLL: 0.9716
Epoch 11 | Train Loss: 1.2647 | Train NLL: [...]<jupyter_text>Dependent Censoring<jupyter_code>rs = np.random.RandomState(11)

feature_order = rs.permutation(len(X.T))
T_features = feature_order[N_OVERLAPPING:(N_FEATURES + N_OVERLAPPING)]
E_features = feature_order[N_FEATURES:(2 * N_FEATURES)]
C_features = feature_order[(2 * N_FEATURES - N_OVERLAPPING):(3 * N_FEATURES - N_OVERLAPPING)]

T_weights = rs.randn(len(T_features))
E_weights = rs.randn(len(E_features))
C_weights = rs.randn(len(C_features))

T = X[:, T_features] @ T_weights
E = X[:, E_features] @ E_weights
C = X[:, C_features] @ C_weights

t_cont = np.exp((T - np.mean(T)) / (2 * np.std(T)))
c_cont = np.exp((C - np.mean(C)) / (2 * np.std(C)))

e_prob = 1 - 1 / (1 + np.exp(4. * (E - E.mean()) / E.std()))
e = (rs.rand(len(e_prob)) < e_prob).astype(int)

ct_cont = np.stack([c_cont, t_cont]).T
y_cont = e * np.amin(ct_cont, axis=1) + (1 - e) * c_cont

s = e * np.argmin(ct_cont, axis=1)
y_disc, t_disc, c_disc = discretize(y_cont, t_cont, c_cont, nbins=N_BINS)

y = onehot(y_disc, ncategories=N_BINS)
pstrue = s_train.mean()
importance_weights = [1 / (1 - pstrue), 1 / pstrue]
print('How many are susceptible?')
print(pd.value_counts(e).sort_index())

print('How many are observed?')
print(pd.value_counts(s).sort_index())

fig, ax = plt.subplots(nrows=1, ncols=4, figsize=(16, 4))

ax[0].hist(e_prob)
ax[1].hist(t_disc)
ax[2].hist(c_disc)
ax[3].hist(y_disc)

plt.show()
x_train, y_train, s_train = X[:6000], y[:6000], s[:6000]
x_val, y_val, s_val = X[6000:7500], y[6000:7500], s[6000:7500]
x_test, y_test, s_test, e_test = X[7500:], y[7500:], s[7500:], e[7500:]
t_disc_test, c_disc_test, y_disc_test = t_disc[7500:], c_disc[7500:], y_disc[7500:]
model = DNMC(n_bins=N_BINS, importance_weights=importance_weights, ld=1e-2, lr=1e-2)
train_model(model, (x_train, y_train, s_train), (x_val, y_val, s_val), 100, learning_rate=LEARNING_RATE)
all_results.append(evaluate_model(model, (x_test, y_test, s_test), e_test, factors='dependent'))
model = DNMC(n_bins=N_BINS, importance_weights=importance_weights, ld=3e-3, lr=1e-2)
train_model(model, (x_train, y_train, s_train), (x_val, y_val, s_val), 100, learning_rate=LEARNING_RATE)
all_results.append(evaluate_model(model, (x_test, y_test, s_test), e_test, factors='dependent'))
model = DNMC(n_bins=N_BINS, importance_weights=importance_weights, ld=3e-2, lr=1e-2)
train_model(model, (x_train, y_train, s_train), (x_val, y_val, s_val), 100, learning_rate=LEARNING_RATE)
all_results.append(evaluate_model(model, (x_test, y_test, s_test), e_test, factors='dependent'))
model = NMC(n_bins=N_BINS, importance_weights=importance_weights, lr=1e-2)
train_model(model, (x_train, y_train, s_train), (x_val, y_val, s_val), 100, learning_rate=LEARNING_RATE)
all_results.append(evaluate_model(model, (x_test, y_test, s_test), e_test, factors='dependent'))
model = NSurv(n_bins=N_BINS, importance_weights=importance_weights, lr=1e-2)
train_model(model, (x_train, y_train, s_train), (x_val, y_val, s_val), 100, learning_rate=LEARNING_RATE)
all_results.append(evaluate_model(model, (x_test, y_test, s_test), e_test, factors='dependent'))
model = MLP(importance_weights=importance_weights, lr=1e-2)
train_model(model, (x_train, y_train, s_train), (x_val, y_val, s_val), 100, learning_rate=LEARNING_RATE)
all_results.append(evaluate_model(model, (x_test, y_test, s_test), e_test, factors='dependent'))<jupyter_output>Epoch  0 | Train Loss: 3.6624 | Train NLL: 1.2143 | Val Loss: 2.7460 | Val NLL: 0.9991
Epoch  1 | Train Loss: 2.3054 | Train NLL: 0.8541 | Val Loss: 2.1627 | Val NLL: 0.9298
Epoch  2 | Train Loss: 1.9031 | Train NLL: 0.7994 | Val Loss: 1.8536 | Val NLL: 0.8588
Epoch  3 | Train Loss: 1.6766 | Train NLL: 0.7613 | Val Loss: 1.6466 | Val NLL: 0.8027
Epoch  4 | Train Loss: 1.5205 | Train NLL: 0.7325 | Val Loss: 1.5165 | Val NLL: 0.7792
Epoch  5 | Train Loss: 1.4037 | Train NLL: 0.7077 | Val Loss: 1.4077 | Val NLL: 0.7500
Epoch  6 | Train Loss: 1.3114 | Train NLL: 0.6856 | Val Loss: 1.3199 | Val NLL: 0.7234
Epoch  7 | Train Loss: 1.2370 | Train NLL: 0.6663 | Val Loss: 1.2453 | Val NLL: 0.6983
Epoch  8 | Train Loss: 1.1740 | Train NLL: 0.6480 | Val Loss: 1.1946 | Val NLL: 0.6884
Epoch  9 | Train Loss: 1.1203 | Train NLL: 0.6323 | Val Loss: 1.1557 | Val NLL: 0.6847
Epoch 10 | Train Loss: 1.0808 | Train NLL: 0.6255 | Val Loss: 1.1310 | Val NLL: 0.6902
Epoch 11 | Train Loss: 1.0421 | Train NLL: [...]<jupyter_text>Summary of results<jupyter_code>summary = pd.DataFrame(all_results)
summary.to_csv('~/Downloads/support_results.csv', index=False)
summary
# separate = non-overlapping subsets of predictors that predict occurrence vs time
# another way to think of it is that there are just two factors in the model, vs three<jupyter_output><empty_output><jupyter_text>Identification of Factors?<jupyter_code># # Trying to work toward evidence of identification of factors
# # Phi -> E, Omega -> T

# E_mask = np.array([False] * X.shape[1])
# E_mask[E_features] = True

# T_mask = np.array([False] * X.shape[1])
# T_mask[T_features] = True

# phi_on_e = np.abs(model.phi_layers[0].kernel.numpy()).max(axis=1)[E_mask].mean()
# phi_off_e = np.abs(model.phi_layers[0].kernel.numpy()).max(axis=1)[~E_mask].mean()

# omega_on_t = np.abs(model.omega_layers[0].kernel.numpy()).max(axis=1)[T_mask].mean()
# omega_off_t = np.abs(model.omega_layers[0].kernel.numpy()).max(axis=1)[~T_mask].mean()

# print(phi_on_e, phi_off_e, omega_on_t, omega_off_t)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.14983346394984326, 'fraction_numerical': 0.058679467084639496, 'mean_word_length': 4.3447643979057595, 'pattern_counts': {'":': 0, '<': 29, '<?xml version=': 0, '>': 29, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '7', 'dir': 'jupyter-structured-clean-dedup', 'id': '397341', 'license': 'MIT', 'n_tokens_mistral': 9232, 'n_tokens_neox': 8416, 'n_words': 1812, 'path': 'notebooks/support_semi_synthetic.ipynb', 'repo_name': 'mengelhard/dnmc', 'size': '20416'}