<jupyter_start><jupyter_text>Reading trial data<jupyter_code>def get_all_epochs(trial):
    try:
        df = get_df_from_logdir(trial, do_tqdm=False)
    except FileNotFoundError:
        return None
    except json.JSONDecodeError:
        return None
    
    return df
dfs = {trial: get_all_epochs(trial) for trial in tqdm(trials)}
df = list(dfs.values())[0]
data = [df[~pd.isna(df.epochs)].iloc[-1] for df in dfs.values() if hasattr(df, 'epochs')]
df = pd.DataFrame(data)
df.to_csv('sm5_linear_lagrange_no_component.csv', index=False)
df = pd.read_csv('sm5_linear_lagrange_no_component.csv')
list(df.columns)
# plot_bar('epochs')
plt.rcParams.update({
    #'font.size': 8,
    'text.usetex': False,
#     'text.latex.preamble': r'\usepackage{amsfonts}',
    'font.family' : 'normal',
        'font.weight' : 'normal',
        'font.size'   : 20
})

plt.hist(df.epochs)
plt.title("Sparsity w.r.t. $c_{eye}$")
plt.scatter(df['config/_gin__eyecoeff__loguniform'], df['metrics/nnz'])
plt.xscale('log')
plt.xlabel("$c_{eye}$")
plt.ylabel("Number of edges")
plt.axhline(15, ls='--')
df_fine = df[df['metrics/nnz'] <= 15]
def find_trial(row, path=path):
    """Get folder name for a trial row."""
    trials = os.listdir(path)
    is_match = [row.trial_id in trial and row.experiment_tag in trial for trial in trials]
    assert sum(is_match) == 1
    idx = np.where(is_match)[0][0]
    trial = trials[idx]
    return trial
from path import Path
import gin
import sparse_causal_model_learner_rl.config
from causal_util import load_env, get_true_graph
gin.bind_parameter('Config._unpickle_skip_init', True)
gin.enter_interactive_mode()
trial_names = [find_trial(row) for _, row in df_fine.iterrows()]
class model_():
    def __init__(self, Mf, Ma):
        self.Mf = Mf
        self.Ma = Ma
        
def adhoc_model(l):
    """Get an instance with Mf, Ma attributes."""
    n_f = l._unpickled_state['config'].get('feature_shape')[0]
    if gin.query_parameter('ModelDecoder.add_batch_number'):
        n_f += 1
    fadd = l._unpickled_state['config'].get('additional_feature_keys')
    print(n_f, fadd)
    
    gin.bind_parameter('graph_for_matrices.additional_features', fadd)
    gin.bind_parameter('plot_model.vmin', 0.0)
    gin.bind_parameter('plot_model.singlecolor_palette', True)
    gin.bind_parameter('plot_model.additional_features', fadd)
    P = l._unpickled_state['trainables_weights']['model']['model.switch.probas'].T
    
    print(P.shape)
    
    Mf = P[:, :n_f].detach().cpu().numpy()
    Ma = P[:, n_f:].detach().cpu().numpy()
    model = model_(Mf, Ma)
    return model
def model_from_trial(trial):
    trial_path = os.path.join(path, trial)
    checkpoints = [x for x in os.listdir(trial_path) if x.startswith('checkpoint')]
    checkpoint_epochs = {x: int(x.split('_')[1]) for x in checkpoints}
    checkpoints_rev = sorted(checkpoints, key=lambda x: checkpoint_epochs[x], reverse=True)
    
    l = None
    for checkpoint in checkpoints_rev:
        print("Trying checkpoint", checkpoint)
        try:
            ckpt_path = os.path.join(trial_path, checkpoint, 'checkpoint')
            l = unpickle(ckpt_path)
#             print(gin.config_str())
            break
        except Exception as e:
            print("Can't read", checkpoint)
            
    if l is None:
        raise ValueError("No checkpoints for trial", trial)
        
    gin.parse_config(l._unpickled_state['gin_config'])
    model = adhoc_model(l)
    
    return l, model
from itertools import permutations
from scipy.spatial.distance import cosine
def permute_model(m, perm, thr=0.9):
    permuted_Mf = m.Mf[perm, :][:, perm] > thr
    permuted_Ma = m.Ma[perm, :] > thr
    return permuted_Mf, permuted_Ma

def distance(m, m_true, perm, thr=0.9):
    permuted_Mf, permuted_Ma = permute_model(m, perm, thr=thr)
    Mf_true, Ma_true = m_true.Mf > thr, m_true.Ma > thr
    
    cosF = cosine(permuted_Mf.flatten(), Mf_true.flatten())
    cosA = cosine(permuted_Ma.flatten(), Ma_true.flatten())

    cos = cosF + cosA
    return cos
def check_last_feature(m, thr=0.9):
    Mf = m.Mf > 0.9
    Ma = m.Ma > 0.9
    assert np.sum(Mf[-1, :-1]) == 0
    assert np.sum(Mf[:-1, -1]) == 0
    assert Mf[-1, -1] == True, Mf[-1, -1]
    assert np.sum(Ma[-1, :]) == 0
check_last_feature(m)
def visualize_trial(trial):

    l, m = model_from_trial(trial)
    check_last_feature(m)

    env = load_env()
    G = get_true_graph(env)
    m_true = model_(G.As, G.Aa)

    fig = l.visualize_model(m)
    display(fig)
    gvz_m = l.visualize_graph(model=m)[1]
    gvz_m_true = l.visualize_graph(model=m_true)[1]
    
    display(gvz_m)
    
    n_f = m.Mf.shape[0]
    perms = list(permutations(range(n_f - 1)))

    distances = [distance(m, m_true, perm) for perm in perms]
    idx = np.argmin(distances)
    print("Best cosine", distances[idx])
    print("Best permutation", perms[idx])

    best_f, best_a = permute_model(m, perms[idx])
    display(l.visualize_graph(model=model_(1. * best_f, 1. * best_a))[1])
    return distances[idx]
cosines = []
for trial in trial_names:
    print(f"===== {trial} ======")
    dst = visualize_trial(trial)
    cosines.append(dst)
plt.title("Cosine distance to best matching\ngraph with nnz=15")

plt.hist(cosines)
plt.scatter(df_fine['config/_gin__eyecoeff__loguniform'], cosines)#df['metrics/nnz']))
plt.xscale('log')
plt.scatter(df_fine['metrics/nnz'], cosines)
nx.draw(nx_m_true, pos = nx.spring_layout(nx_m_true))
from networkx.algorithms import isomorphism
match = isomorphism.GraphMatcher(nx_m, nx_m_true)
dir(match)
match.subgraph_is_isomorphic()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1462457642232923, 'fraction_numerical': 0.008917424647761726, 'mean_word_length': 3.7086481947942906, 'pattern_counts': {'":': 0, '<': 6, '<?xml version=': 0, '>': 11, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '334855', 'license': 'CC0-1.0', 'n_tokens_mistral': 2203, 'n_tokens_neox': 2103, 'n_words': 409, 'path': 'causal_analysis/sm5_linear_lagrange_no_component.ipynb', 'repo_name': 'sergeivolodin/causality-disentanglement-rl', 'size': '5607'}