# -*- coding: utf-8 -*-
# Copyright 2020 The PsiZ Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""Example that infers a shared embedding for three groups.

Fake data is generated from a ground truth model for three different
groups. In this example, these groups represent groups of agents with
varying levels of skill: novices, intermediates, and experts. Each group
has a different set of attention weights. An embedding model is
inferred from the simulated data and compared to the ground truth
model.

Results are saved in the directory specified by `fp_example`. By
default, a `psiz_examples` directory is created in your home directory.

Example output:

    Restart Summary
    n_valid_restart 1 | total_duration: 2104 s
    best | n_epoch: 999 | val_loss: 3.0700
    mean ±stddev | n_epoch: 999 ±0 | val_loss: 3.0700 ±0.0000 |
        2088 ±0 s | 2090 ±0 ms/epoch

    Model Comparison (R^2)
    ================================
      True  |        Inferred
            | Novice  Interm  Expert
    --------+-----------------------
     Novice |   0.94    0.67    0.06
     Interm |   0.55    0.87    0.35
     Expert |   0.09    0.42    0.85

"""

import copy
import os
from pathlib import Path
import shutil

import imageio
import matplotlib
import matplotlib.pyplot as plt
import numpy as np
import tensorflow as tf
import tensorflow_probability as tfp

import psiz

# Uncomment the following line to force eager execution.
# tf.config.experimental_run_functions_eagerly(True)

# Modify the following to control GPU visibility.
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "0"


def main():
    """Run script."""
    # Settings.
    fp_example = Path.home() / Path('psiz_examples', 'rank', 'vi_3ge')
    fp_board = fp_example / Path('logs', 'fit')
    n_stimuli = 30
    n_dim = 3
    n_group = 3
    n_trial = 2000
    epochs = 1000
    n_restart = 1
    batch_size = 128
    n_frame = 1

    # Directory preparation.
    fp_example.mkdir(parents=True, exist_ok=True)
    # Remove existing TensorBoard logs.
    if fp_board.exists():
        shutil.rmtree(fp_board)

    # Plot settings.
    small_size = 6
    medium_size = 8
    large_size = 10
    plt.rc('font', size=small_size)  # controls default text sizes
    plt.rc('axes', titlesize=medium_size)
    plt.rc('axes', labelsize=small_size)
    plt.rc('xtick', labelsize=small_size)
    plt.rc('ytick', labelsize=small_size)
    plt.rc('legend', fontsize=small_size)
    plt.rc('figure', titlesize=large_size)

    model_true = ground_truth(n_stimuli, n_group)
    proxy_true = psiz.models.Proxy(model=model_true)

    # Generate a random docket of trials to show each group.
    generator = psiz.generators.RandomRank(
        n_stimuli, n_reference=8, n_select=2
    )
    docket = generator.generate(n_trial)

    # Create virtual agents for each group.
    agent_novice = psiz.agents.RankAgent(proxy_true.model, group_id=0)
    agent_interm = psiz.agents.RankAgent(proxy_true.model, group_id=1)
    agent_expert = psiz.agents.RankAgent(proxy_true.model, group_id=2)

    # Simulate similarity judgments for each group.
    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))

    # Compute ground truth similarity matrices.
    def truth_sim_func0(z_q, z_ref):
        return proxy_true.similarity(z_q, z_ref, group_id=0)

    def truth_sim_func1(z_q, z_ref):
        return proxy_true.similarity(z_q, z_ref, group_id=1)

    def truth_sim_func2(z_q, z_ref):
        return proxy_true.similarity(z_q, z_ref, group_id=2)

    simmat_truth = (
        psiz.utils.pairwise_matrix(truth_sim_func0, proxy_true.z[0]),
        psiz.utils.pairwise_matrix(truth_sim_func1, proxy_true.z[0]),
        psiz.utils.pairwise_matrix(truth_sim_func2, proxy_true.z[0])
    )

    # Partition observations into 80% train, 10% validation and 10% test set.
    obs_train, obs_val, obs_test = psiz.utils.standard_split(obs)

    compile_kwargs = {
        'loss': tf.keras.losses.CategoricalCrossentropy(),
        'optimizer': tf.keras.optimizers.Adam(lr=.001),
        'weighted_metrics': [
            tf.keras.metrics.CategoricalCrossentropy(name='cce')
        ]
    }

    # Infer independent models with increasing amounts of data.
    if n_frame == 1:
        n_obs = np.array([obs_train.n_trial], dtype=int)
    else:
        n_obs = np.round(
            np.linspace(15, obs_train.n_trial, n_frame)
        ).astype(np.int64)
    r2 = np.empty([n_frame, n_group, n_group]) * np.nan
    train_loss = np.empty((n_frame)) * np.nan
    val_loss = np.empty((n_frame)) * np.nan
    test_loss = np.empty((n_frame)) * np.nan
    for i_frame in range(n_frame):
        include_idx = np.arange(0, n_obs[i_frame])
        obs_round_train = obs_train.subset(include_idx)
        print(
            '\n  Frame {0} ({1} obs)'.format(i_frame, obs_round_train.n_trial)
        )

        # Define model.
        model = build_model(n_stimuli, n_dim, n_group, obs_round_train.n_trial)
        proxy_inferred = psiz.models.Proxy(model=model)

        # Define callbacks.
        fp_board_frame = fp_board / Path('frame_{0}'.format(i_frame))
        cb_board = psiz.keras.callbacks.TensorBoardRe(
            log_dir=fp_board_frame, histogram_freq=0,
            write_graph=False, write_images=False, update_freq='epoch',
            profile_batch=0, embeddings_freq=0, embeddings_metadata=None
        )
        cb_early = psiz.keras.callbacks.EarlyStoppingRe(
            'loss', patience=100, mode='min', restore_best_weights=False,
            verbose=1
        )
        callbacks = [cb_board, cb_early]

        # Infer model.
        restart_record = proxy_inferred.fit(
            obs_round_train, validation_data=obs_val, epochs=epochs,
            batch_size=batch_size, callbacks=callbacks, n_restart=n_restart,
            monitor='val_loss', verbose=1, compile_kwargs=compile_kwargs
        )

        train_loss[i_frame] = restart_record.record['loss'][0]
        val_loss[i_frame] = restart_record.record['val_loss'][0]

        tf.keras.backend.clear_session()
        proxy_inferred.model.n_sample = 100
        proxy_inferred.compile(**compile_kwargs)
        test_metrics = proxy_inferred.evaluate(
            obs_test, verbose=0, return_dict=True
        )
        test_loss[i_frame] = test_metrics['loss']

        z = proxy_inferred.model.stimuli.embeddings.mode().numpy()
        if proxy_inferred.model.stimuli.mask_zero:
            z = z[:, 1:, :]

        # Compare the inferred model with ground truth by comparing the
        # similarity matrices implied by each model.
        def infer_sim_func(z_q, z_ref):
            return proxy_inferred.similarity(z_q, z_ref)

        simmat_infer = (
            psiz.utils.pairwise_matrix(infer_sim_func, z[0]),
            psiz.utils.pairwise_matrix(infer_sim_func, z[1]),
            psiz.utils.pairwise_matrix(infer_sim_func, z[2])
        )
        for i_truth in range(n_group):
            for j_infer in range(n_group):
                r2[i_frame, i_truth, j_infer] = psiz.utils.matrix_comparison(
                    simmat_truth[i_truth], simmat_infer[j_infer],
                    score='r2'
                )

        # Display comparison results. A good inferred model will have a high
        # R^2 value on the diagonal elements (max is 1) and relatively low R^2
        # values on the off-diagonal elements.
        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(
            r2[i_frame, 0, 0], r2[i_frame, 0, 1], r2[i_frame, 0, 2]))
        print('     Interm | {0: >6.2f}  {1: >6.2f}  {2: >6.2f}'.format(
            r2[i_frame, 1, 0], r2[i_frame, 1, 1], r2[i_frame, 1, 2]))
        print('     Expert | {0: >6.2f}  {1: >6.2f}  {2: >6.2f}'.format(
            r2[i_frame, 2, 0], r2[i_frame, 2, 1], r2[i_frame, 2, 2]))
        print('\n')

        # Create and save visual frame.
        fig0 = plt.figure(figsize=(12, 5), dpi=200)
        plot_frame(
            fig0, n_obs, train_loss, val_loss, test_loss, r2, proxy_true,
            proxy_inferred, i_frame
        )
        fname = fp_example / Path('frame_{0}.tiff'.format(i_frame))
        plt.savefig(
            os.fspath(fname), format='tiff', bbox_inches="tight", dpi=300
        )

    # Create animation.
    if n_frame > 1:
        frames = []
        for i_frame in range(n_frame):
            fname = fp_example / Path('frame_{0}.tiff'.format(i_frame))
            frames.append(imageio.imread(fname))
        imageio.mimwrite(fp_example / Path('evolution.gif'), frames, fps=1)


def ground_truth(n_stimuli, n_group):
    """Return a ground truth embedding."""
    n_dim = 4
    embedding = tf.keras.layers.Embedding(
        n_stimuli+1, n_dim, mask_zero=True,
        embeddings_initializer=tf.keras.initializers.RandomNormal(
            stddev=.17, seed=58
        )
    )
    stimuli = psiz.keras.layers.Stimuli(embedding=embedding)
    kernel = psiz.keras.layers.AttentionKernel(
        group_level=1,
        distance=psiz.keras.layers.WeightedMinkowski(
            rho_initializer=tf.keras.initializers.Constant(2.),
            trainable=False,
        ),
        attention=psiz.keras.layers.GroupAttention(
            n_dim=n_dim, n_group=n_group
        ),
        similarity=psiz.keras.layers.ExponentialSimilarity(
            tau_initializer=tf.keras.initializers.Constant(1.),
            gamma_initializer=tf.keras.initializers.Constant(0.),
            trainable=False,
        )
    )
    kernel.attention.embeddings.assign(
        np.array((
            (1.8, 1.8, .2, .2),
            (1., 1., 1., 1.),
            (.2, .2, 1.8, 1.8)
        ))
    )
    model = psiz.models.Rank(stimuli=stimuli, kernel=kernel)
    return model


def build_model(n_stimuli, n_dim, n_group, n_obs_train):
    """Build model.

    Arguments:
        n_stimuli: Integer indicating the number of stimuli in the
            embedding.
        n_dim: Integer indicating the dimensionality of the embedding.
        n_group: Integer indicating the number of groups.
        n_obs_train: Integer indicating the number of training
            observations. Used to determine KL weight for variational
            inference.

    Returns:
        model: A TensorFlow Keras model.

    """
    kl_weight = 1. / n_obs_train
    prior_scale = .2

    n_source_embeddings = n_group * (n_stimuli + 1)
    embedding_posterior = psiz.keras.layers.EmbeddingNormalDiag(
        n_source_embeddings, n_dim, mask_zero=True,
        scale_initializer=tf.keras.initializers.Constant(
            tfp.math.softplus_inverse(prior_scale).numpy()
        )
    )
    embedding_prior = psiz.keras.layers.EmbeddingShared(
        n_source_embeddings, n_dim, mask_zero=True,
        embedding=psiz.keras.layers.EmbeddingNormalDiag(
            1, 1,
            loc_initializer=tf.keras.initializers.Constant(0.),
            scale_initializer=tf.keras.initializers.Constant(
                tfp.math.softplus_inverse(prior_scale).numpy()
            ),
            loc_trainable=False
        )
    )
    embedding_variational = psiz.keras.layers.EmbeddingVariational(
        posterior=embedding_posterior, prior=embedding_prior,
        kl_weight=kl_weight, kl_n_sample=30
    )
    stimuli = psiz.keras.layers.Stimuli(
        embedding=embedding_variational, group_level=1, n_group=n_group
    )

    kernel = psiz.keras.layers.Kernel(
        distance=psiz.keras.layers.WeightedMinkowski(
            rho_initializer=tf.keras.initializers.Constant(2.),
            trainable=False,
        ),
        similarity=psiz.keras.layers.ExponentialSimilarity(
            beta_initializer=tf.keras.initializers.Constant(10.),
            tau_initializer=tf.keras.initializers.Constant(1.),
            gamma_initializer=tf.keras.initializers.Constant(0.),
            trainable=False
        )
    )
    model = psiz.models.Rank(
        stimuli=stimuli, kernel=kernel, n_sample=1
    )
    return model


def plot_frame(
        fig0, n_obs, train_loss, val_loss, test_loss, r2, proxy_true,
        proxy_inferred, i_frame):
    """Plot posteriors."""
    # Settings.
    group_labels = ['Novice', 'Intermediate', 'Expert']

    n_group = proxy_inferred.model.stimuli.n_group
    n_dim = proxy_inferred.model.n_dim

    gs = fig0.add_gridspec(n_group + 1, 4)

    f0_ax0 = fig0.add_subplot(gs[0, 0:2])
    plot_loss(f0_ax0, n_obs, train_loss, val_loss, test_loss)

    f0_ax1 = fig0.add_subplot(gs[0, 2])
    plot_convergence(fig0, f0_ax1, n_obs, r2[i_frame])

    gs.tight_layout(fig0)


def plot_logitnormal(ax, dist, name=None, c=None):
    """Plot univariate distribution.

    Arguments:
        ax:
        dist:
        name:

    """
    if name is None:
        name = 'x'

    x = np.linspace(.001, .999, 1000)
    y = dist.prob(x)

    # Determine mode from samples.
    idx = np.argmax(y)
    x_mode = x[idx]

    ax.plot(x, y, c=c)
    ax.text(x_mode, .75 * np.max(y), '{0:.2f}'.format(x_mode))
    ax.set_xlabel(r'${0}$'.format(name))
    ax.set_ylabel(r'$p({0})$'.format(name))
    ax.set_xlim([0, 1])
    ax.set_xticks([0, 1])
    ax.set_xticklabels([0, 1])


def plot_loss(ax, n_obs, train_loss, val_loss, test_loss):
    """Plot loss."""
    # Settings
    ms = 2

    ax.plot(n_obs, train_loss, 'bo-', ms=ms, label='Train Loss')
    ax.plot(n_obs, val_loss, 'go-', ms=ms, label='Val. Loss')
    ax.plot(n_obs, test_loss, 'ro-', ms=ms, label='Test Loss')
    ax.set_title('Optimization Objective')

    ax.set_xlabel('Trials')
    limits = [0, np.max(n_obs) + 10]
    ax.set_xlim(limits)
    ticks = [np.min(n_obs), np.max(n_obs)]
    ax.set_xticks(ticks)

    ax.set_ylabel('Loss')
    ax.legend()


def plot_convergence(fig, ax, n_obs, r2):
    """Plot convergence."""
    # Settings.
    cmap = matplotlib.cm.get_cmap('Greys')
    labels = ['Nov', 'Int', 'Exp']

    im = ax.imshow(r2, cmap=cmap, vmin=0., vmax=1.)
    fig.colorbar(im, ax=ax)
    ax.set_xticks([0, 1, 2])
    ax.set_xticklabels(labels)
    ax.set_yticks([0, 1, 2])
    ax.set_yticklabels(labels)
    ax.set_ylabel('True')
    ax.set_xlabel('Inferred')
    ax.set_title(r'$R^2$ Convergence')


if __name__ == "__main__":
    main()

tensorflow.keras.initializers.Constant tensorflow.keras.losses.CategoricalCrossentropy numpy.linspace numpy.min numpy.arange tensorflow.keras.initializers.RandomNormal matplotlib.pyplot.rc numpy.max tensorflow.keras.metrics.CategoricalCrossentropy numpy.argmax tensorflow.keras.optimizers.Adam tensorflow.keras.backend.clear_session matplotlib.cm.get_cmap numpy.array numpy.empty matplotlib.pyplot.figure
[(93, 'matplotlib.pyplot.rc', 'plt.rc', (['"""font"""'], {'size': 'small_size'}), True, 'import matplotlib.pyplot as plt\n'), (94, 'matplotlib.pyplot.rc', 'plt.rc', (['"""axes"""'], {'titlesize': 'medium_size'}), True, 'import matplotlib.pyplot as plt\n'), (95, 'matplotlib.pyplot.rc', 'plt.rc', (['"""axes"""'], {'labelsize': 'small_size'}), True, 'import matplotlib.pyplot as plt\n'), (96, 'matplotlib.pyplot.rc', 'plt.rc', (['"""xtick"""'], {'labelsize': 'small_size'}), True, 'import matplotlib.pyplot as plt\n'), (97, 'matplotlib.pyplot.rc', 'plt.rc', (['"""ytick"""'], {'labelsize': 'small_size'}), True, 'import matplotlib.pyplot as plt\n'), (98, 'matplotlib.pyplot.rc', 'plt.rc', (['"""legend"""'], {'fontsize': 'small_size'}), True, 'import matplotlib.pyplot as plt\n'), (99, 'matplotlib.pyplot.rc', 'plt.rc', (['"""figure"""'], {'titlesize': 'large_size'}), True, 'import matplotlib.pyplot as plt\n'), (102, 'psiz.models.Proxy', 'psiz.models.Proxy', ([], {'model': 'model_true'}), False, 'import psiz\n'), (105, 'psiz.generators.RandomRank', 'psiz.generators.RandomRank', (['n_stimuli'], {'n_reference': '(8)', 'n_select': '(2)'}), False, 'import psiz\n'), (111, 'psiz.agents.RankAgent', 'psiz.agents.RankAgent', (['proxy_true.model'], {'group_id': '(0)'}), False, 'import psiz\n'), (112, 'psiz.agents.RankAgent', 'psiz.agents.RankAgent', (['proxy_true.model'], {'group_id': '(1)'}), False, 'import psiz\n'), (113, 'psiz.agents.RankAgent', 'psiz.agents.RankAgent', (['proxy_true.model'], {'group_id': '(2)'}), False, 'import psiz\n'), (119, 'psiz.trials.stack', 'psiz.trials.stack', (['(obs_novice, obs_interm, obs_expert)'], {}), False, 'import psiz\n'), (138, 'psiz.utils.standard_split', 'psiz.utils.standard_split', (['obs'], {}), False, 'import psiz\n'), (267, 'psiz.keras.layers.Stimuli', 'psiz.keras.layers.Stimuli', ([], {'embedding': 'embedding'}), False, 'import psiz\n'), (290, 'psiz.models.Rank', 'psiz.models.Rank', ([], {'stimuli': 'stimuli', 'kernel': 'kernel'}), False, 'import psiz\n'), (331, 'psiz.keras.layers.EmbeddingVariational', 'psiz.keras.layers.EmbeddingVariational', ([], {'posterior': 'embedding_posterior', 'prior': 'embedding_prior', 'kl_weight': 'kl_weight', 'kl_n_sample': '(30)'}), False, 'import psiz\n'), (335, 'psiz.keras.layers.Stimuli', 'psiz.keras.layers.Stimuli', ([], {'embedding': 'embedding_variational', 'group_level': '(1)', 'n_group': 'n_group'}), False, 'import psiz\n'), (351, 'psiz.models.Rank', 'psiz.models.Rank', ([], {'stimuli': 'stimuli', 'kernel': 'kernel', 'n_sample': '(1)'}), False, 'import psiz\n'), (390, 'numpy.linspace', 'np.linspace', (['(0.001)', '(0.999)', '(1000)'], {}), True, 'import numpy as np\n'), (394, 'numpy.argmax', 'np.argmax', (['y'], {}), True, 'import numpy as np\n'), (429, 'matplotlib.cm.get_cmap', 'matplotlib.cm.get_cmap', (['"""Greys"""'], {}), False, 'import matplotlib\n'), (72, 'pathlib.Path.home', 'Path.home', ([], {}), False, 'from pathlib import Path\n'), (72, 'pathlib.Path', 'Path', (['"""psiz_examples"""', '"""rank"""', '"""vi_3ge"""'], {}), False, 'from pathlib import Path\n'), (73, 'pathlib.Path', 'Path', (['"""logs"""', '"""fit"""'], {}), False, 'from pathlib import Path\n'), (87, 'shutil.rmtree', 'shutil.rmtree', (['fp_board'], {}), False, 'import shutil\n'), (132, 'psiz.utils.pairwise_matrix', 'psiz.utils.pairwise_matrix', (['truth_sim_func0', 'proxy_true.z[0]'], {}), False, 'import psiz\n'), (133, 'psiz.utils.pairwise_matrix', 'psiz.utils.pairwise_matrix', (['truth_sim_func1', 'proxy_true.z[0]'], {}), False, 'import psiz\n'), (134, 'psiz.utils.pairwise_matrix', 'psiz.utils.pairwise_matrix', (['truth_sim_func2', 'proxy_true.z[0]'], {}), False, 'import psiz\n'), (141, 'tensorflow.keras.losses.CategoricalCrossentropy', 'tf.keras.losses.CategoricalCrossentropy', ([], {}), True, 'import tensorflow as tf\n'), (142, 'tensorflow.keras.optimizers.Adam', 'tf.keras.optimizers.Adam', ([], {'lr': '(0.001)'}), True, 'import tensorflow as tf\n'), (150, 'numpy.array', 'np.array', (['[obs_train.n_trial]'], {'dtype': 'int'}), True, 'import numpy as np\n'), (155, 'numpy.empty', 'np.empty', (['[n_frame, n_group, n_group]'], {}), True, 'import numpy as np\n'), (156, 'numpy.empty', 'np.empty', (['n_frame'], {}), True, 'import numpy as np\n'), (157, 'numpy.empty', 'np.empty', (['n_frame'], {}), True, 'import numpy as np\n'), (158, 'numpy.empty', 'np.empty', (['n_frame'], {}), True, 'import numpy as np\n'), (160, 'numpy.arange', 'np.arange', (['(0)', 'n_obs[i_frame]'], {}), True, 'import numpy as np\n'), (168, 'psiz.models.Proxy', 'psiz.models.Proxy', ([], {'model': 'model'}), False, 'import psiz\n'), (172, 'psiz.keras.callbacks.TensorBoardRe', 'psiz.keras.callbacks.TensorBoardRe', ([], {'log_dir': 'fp_board_frame', 'histogram_freq': '(0)', 'write_graph': '(False)', 'write_images': '(False)', 'update_freq': '"""epoch"""', 'profile_batch': '(0)', 'embeddings_freq': '(0)', 'embeddings_metadata': 'None'}), False, 'import psiz\n'), (177, 'psiz.keras.callbacks.EarlyStoppingRe', 'psiz.keras.callbacks.EarlyStoppingRe', (['"""loss"""'], {'patience': '(100)', 'mode': '"""min"""', 'restore_best_weights': '(False)', 'verbose': '(1)'}), False, 'import psiz\n'), (193, 'tensorflow.keras.backend.clear_session', 'tf.keras.backend.clear_session', ([], {}), True, 'import tensorflow as tf\n'), (239, 'matplotlib.pyplot.figure', 'plt.figure', ([], {'figsize': '(12, 5)', 'dpi': '(200)'}), True, 'import matplotlib.pyplot as plt\n'), (284, 'numpy.array', 'np.array', (['((1.8, 1.8, 0.2, 0.2), (1.0, 1.0, 1.0, 1.0), (0.2, 0.2, 1.8, 1.8))'], {}), True, 'import numpy as np\n'), (419, 'numpy.min', 'np.min', (['n_obs'], {}), True, 'import numpy as np\n'), (419, 'numpy.max', 'np.max', (['n_obs'], {}), True, 'import numpy as np\n'), (144, 'tensorflow.keras.metrics.CategoricalCrossentropy', 'tf.keras.metrics.CategoricalCrossentropy', ([], {'name': '"""cce"""'}), True, 'import tensorflow as tf\n'), (211, 'psiz.utils.pairwise_matrix', 'psiz.utils.pairwise_matrix', (['infer_sim_func', 'z[0]'], {}), False, 'import psiz\n'), (212, 'psiz.utils.pairwise_matrix', 'psiz.utils.pairwise_matrix', (['infer_sim_func', 'z[1]'], {}), False, 'import psiz\n'), (213, 'psiz.utils.pairwise_matrix', 'psiz.utils.pairwise_matrix', (['infer_sim_func', 'z[2]'], {}), False, 'import psiz\n'), (246, 'os.fspath', 'os.fspath', (['fname'], {}), False, 'import os\n'), (263, 'tensorflow.keras.initializers.RandomNormal', 'tf.keras.initializers.RandomNormal', ([], {'stddev': '(0.17)', 'seed': '(58)'}), True, 'import tensorflow as tf\n'), (274, 'psiz.keras.layers.GroupAttention', 'psiz.keras.layers.GroupAttention', ([], {'n_dim': 'n_dim', 'n_group': 'n_group'}), False, 'import psiz\n'), (398, 'numpy.max', 'np.max', (['y'], {}), True, 'import numpy as np\n'), (417, 'numpy.max', 'np.max', (['n_obs'], {}), True, 'import numpy as np\n'), (217, 'psiz.utils.matrix_comparison', 'psiz.utils.matrix_comparison', (['simmat_truth[i_truth]', 'simmat_infer[j_infer]'], {'score': '"""r2"""'}), False, 'import psiz\n'), (254, 'imageio.imread', 'imageio.imread', (['fname'], {}), False, 'import imageio\n'), (255, 'pathlib.Path', 'Path', (['"""evolution.gif"""'], {}), False, 'from pathlib import Path\n'), (153, 'numpy.linspace', 'np.linspace', (['(15)', 'obs_train.n_trial', 'n_frame'], {}), True, 'import numpy as np\n'), (271, 'tensorflow.keras.initializers.Constant', 'tf.keras.initializers.Constant', (['(2.0)'], {}), True, 'import tensorflow as tf\n'), (278, 'tensorflow.keras.initializers.Constant', 'tf.keras.initializers.Constant', (['(1.0)'], {}), True, 'import tensorflow as tf\n'), (279, 'tensorflow.keras.initializers.Constant', 'tf.keras.initializers.Constant', (['(0.0)'], {}), True, 'import tensorflow as tf\n'), (324, 'tensorflow.keras.initializers.Constant', 'tf.keras.initializers.Constant', (['(0.0)'], {}), True, 'import tensorflow as tf\n'), (341, 'tensorflow.keras.initializers.Constant', 'tf.keras.initializers.Constant', (['(2.0)'], {}), True, 'import tensorflow as tf\n'), (345, 'tensorflow.keras.initializers.Constant', 'tf.keras.initializers.Constant', (['(10.0)'], {}), True, 'import tensorflow as tf\n'), (346, 'tensorflow.keras.initializers.Constant', 'tf.keras.initializers.Constant', (['(1.0)'], {}), True, 'import tensorflow as tf\n'), (347, 'tensorflow.keras.initializers.Constant', 'tf.keras.initializers.Constant', (['(0.0)'], {}), True, 'import tensorflow as tf\n'), (317, 'tensorflow_probability.math.softplus_inverse', 'tfp.math.softplus_inverse', (['prior_scale'], {}), True, 'import tensorflow_probability as tfp\n'), (326, 'tensorflow_probability.math.softplus_inverse', 'tfp.math.softplus_inverse', (['prior_scale'], {}), True, 'import tensorflow_probability as tfp\n')]