def main():
    'Run the simulation that infers an embedding for three groups.'
    n_stimuli = 30
    n_dim = 4
    n_group = 3
    n_restart = 1
    epochs = 1000
    n_trial = 2000
    batch_size = 128
    model_true = ground_truth(n_stimuli, n_dim, n_group)
    generator = psiz.trials.RandomRank(n_stimuli, n_reference=8, n_select=2)
    docket = generator.generate(n_trial)
    agent_novice = psiz.agents.RankAgent(model_true, groups=[0])
    agent_interm = psiz.agents.RankAgent(model_true, groups=[1])
    agent_expert = psiz.agents.RankAgent(model_true, groups=[2])
    obs_novice = agent_novice.simulate(docket)
    obs_interm = agent_interm.simulate(docket)
    obs_expert = agent_expert.simulate(docket)
    obs = psiz.trials.stack((obs_novice, obs_interm, obs_expert))
    (obs_train, obs_val, obs_test) = psiz.utils.standard_split(obs)
    ds_obs_train = obs_train.as_dataset().shuffle(buffer_size=obs_train.n_trial, reshuffle_each_iteration=True).batch(batch_size, drop_remainder=False)
    ds_obs_val = obs_val.as_dataset().batch(batch_size, drop_remainder=False)
    ds_obs_test = obs_test.as_dataset().batch(batch_size, drop_remainder=False)
    early_stop = psiz.keras.callbacks.EarlyStoppingRe('val_cce', patience=15, mode='min', restore_best_weights=True)
    callbacks = [early_stop]
    compile_kwargs = {'loss': tf.keras.losses.CategoricalCrossentropy(), 'optimizer': tf.keras.optimizers.Adam(lr=0.001), 'weighted_metrics': [tf.keras.metrics.CategoricalCrossentropy(name='cce')]}
    model_inferred = build_model(n_stimuli, n_dim, n_group)
    restarter = psiz.keras.Restarter(model_inferred, compile_kwargs=compile_kwargs, monitor='val_loss', n_restart=n_restart)
    restart_record = restarter.fit(x=ds_obs_train, validation_data=ds_obs_val, epochs=epochs, callbacks=callbacks, verbose=0)
    model_inferred = restarter.model
    simmat_truth = (model_similarity(model_true, groups=[0]), model_similarity(model_true, groups=[1]), model_similarity(model_true, groups=[2]))
    simmat_inferred = (model_similarity(model_inferred, groups=[0]), model_similarity(model_inferred, groups=[1]), model_similarity(model_inferred, groups=[2]))
    r_squared = np.empty((n_group, n_group))
    for i_truth in range(n_group):
        for j_infer in range(n_group):
            (rho, _) = pearsonr(simmat_truth[i_truth], simmat_inferred[j_infer])
            r_squared[(i_truth, j_infer)] = (rho ** 2)
    attention_weight = tf.stack([model_inferred.kernel.subnets[0].distance.w, model_inferred.kernel.subnets[1].distance.w, model_inferred.kernel.subnets[2].distance.w], axis=0).numpy()
    idx_sorted = np.argsort((- attention_weight[0, :]))
    attention_weight = attention_weight[:, idx_sorted]
    group_labels = ['Novice', 'Intermediate', 'Expert']
    print('\n    Attention weights:')
    for i_group in range(attention_weight.shape[0]):
        print('    {0:>12} | {1}'.format(group_labels[i_group], np.array2string(attention_weight[i_group, :], formatter={'float_kind': (lambda x: ('%.2f' % x))})))
    print('\n    Model Comparison (R^2)')
    print('    ================================')
    print('      True  |        Inferred')
    print('            | Novice  Interm  Expert')
    print('    --------+-----------------------')
    print('     Novice | {0: >6.2f}  {1: >6.2f}  {2: >6.2f}'.format(r_squared[(0, 0)], r_squared[(0, 1)], r_squared[(0, 2)]))
    print('     Interm | {0: >6.2f}  {1: >6.2f}  {2: >6.2f}'.format(r_squared[(1, 0)], r_squared[(1, 1)], r_squared[(1, 2)]))
    print('     Expert | {0: >6.2f}  {1: >6.2f}  {2: >6.2f}'.format(r_squared[(2, 0)], r_squared[(2, 1)], r_squared[(2, 2)]))
    print('\n')
Run the simulation that infers an embedding for three groups.
def main():
    
    n_stimuli = 30
    n_dim = 4
    n_group = 3
    n_restart = 1
    epochs = 1000
    n_trial = 2000
    batch_size = 128
    model_true = ground_truth(n_stimuli, n_dim, n_group)
    generator = psiz.trials.RandomRank(n_stimuli, n_reference=8, n_select=2)
    docket = generator.generate(n_trial)
    agent_novice = psiz.agents.RankAgent(model_true, groups=[0])
    agent_interm = psiz.agents.RankAgent(model_true, groups=[1])
    agent_expert = psiz.agents.RankAgent(model_true, groups=[2])
    obs_novice = agent_novice.simulate(docket)
    obs_interm = agent_interm.simulate(docket)
    obs_expert = agent_expert.simulate(docket)
    obs = psiz.trials.stack((obs_novice, obs_interm, obs_expert))
    (obs_train, obs_val, obs_test) = psiz.utils.standard_split(obs)
    ds_obs_train = obs_train.as_dataset().shuffle(buffer_size=obs_train.n_trial, reshuffle_each_iteration=True).batch(batch_size, drop_remainder=False)
    ds_obs_val = obs_val.as_dataset().batch(batch_size, drop_remainder=False)
    ds_obs_test = obs_test.as_dataset().batch(batch_size, drop_remainder=False)
    early_stop = psiz.keras.callbacks.EarlyStoppingRe('val_cce', patience=15, mode='min', restore_best_weights=True)
    callbacks = [early_stop]
    compile_kwargs = {'loss': tf.keras.losses.CategoricalCrossentropy(), 'optimizer': tf.keras.optimizers.Adam(lr=0.001), 'weighted_metrics': [tf.keras.metrics.CategoricalCrossentropy(name='cce')]}
    model_inferred = build_model(n_stimuli, n_dim, n_group)
    restarter = psiz.keras.Restarter(model_inferred, compile_kwargs=compile_kwargs, monitor='val_loss', n_restart=n_restart)
    restart_record = restarter.fit(x=ds_obs_train, validation_data=ds_obs_val, epochs=epochs, callbacks=callbacks, verbose=0)
    model_inferred = restarter.model
    simmat_truth = (model_similarity(model_true, groups=[0]), model_similarity(model_true, groups=[1]), model_similarity(model_true, groups=[2]))
    simmat_inferred = (model_similarity(model_inferred, groups=[0]), model_similarity(model_inferred, groups=[1]), model_similarity(model_inferred, groups=[2]))
    r_squared = np.empty((n_group, n_group))
    for i_truth in range(n_group):
        for j_infer in range(n_group):
            (rho, _) = pearsonr(simmat_truth[i_truth], simmat_inferred[j_infer])
            r_squared[(i_truth, j_infer)] = (rho ** 2)
    attention_weight = tf.stack([model_inferred.kernel.subnets[0].distance.w, model_inferred.kernel.subnets[1].distance.w, model_inferred.kernel.subnets[2].distance.w], axis=0).numpy()
    idx_sorted = np.argsort((- attention_weight[0, :]))
    attention_weight = attention_weight[:, idx_sorted]
    group_labels = ['Novice', 'Intermediate', 'Expert']
    print('\n    Attention weights:')
    for i_group in range(attention_weight.shape[0]):
        print('    {0:>12} | {1}'.format(group_labels[i_group], np.array2string(attention_weight[i_group, :], formatter={'float_kind': (lambda x: ('%.2f' % x))})))
    print('\n    Model Comparison (R^2)')
    print('    ================================')
    print('      True  |        Inferred')
    print('            | Novice  Interm  Expert')
    print('    --------+-----------------------')
    print('     Novice | {0: >6.2f}  {1: >6.2f}  {2: >6.2f}'.format(r_squared[(0, 0)], r_squared[(0, 1)], r_squared[(0, 2)]))
    print('     Interm | {0: >6.2f}  {1: >6.2f}  {2: >6.2f}'.format(r_squared[(1, 0)], r_squared[(1, 1)], r_squared[(1, 2)]))
    print('     Expert | {0: >6.2f}  {1: >6.2f}  {2: >6.2f}'.format(r_squared[(2, 0)], r_squared[(2, 1)], r_squared[(2, 2)]))
    print('\n')