text:
import numbers
import tempfile

import imageio
import matplotlib.pyplot as plt
import networkx as nx
import numpy as np
import numpy.linalg as linalg
import seaborn as sns
import pandas as pd

from matplotlib.colors import ListedColormap, to_hex
from matplotlib.patches import Ellipse, Rectangle, FancyArrowPatch
from scipy.stats import norm
from scipy.special import expit
from sklearn.preprocessing import LabelEncoder
from sklearn.utils import check_random_state
from dynetlsm.plots import get_colors

__all__ = ['plot_network', 'make_network_animation',
           'plot_sociability', 'plot_lambda', 'plot_node_trajectories',
           'plot_pairwise_distances', 'plot_pairwise_probabilities']

def normal_contour(mean, cov, n_std=2, ax=None, **kwargs):
    if cov.shape[0] != 2:
        raise ValueError('Only for bivariate normal densities.')

    eigenvalues, eigenvectors = linalg.eigh(cov)

    # sort the eigenvalues and eigenvectors in descending order
    order = eigenvalues.argsort()[::-1]
    eigenvalues = eigenvalues[order]
    eigenvectors = eigenvectors[:, order]

    # determine the angle of rotation
    angle = np.degrees(np.arctan2(*eigenvectors[:, 0][::-1]))

    if ax is None:
        ax = plt.gca()

    if isinstance(n_std, numbers.Integral):
        # the diameter of the ellipse is twice the square root of the evalues
        width, height = 2 * n_std * np.sqrt(eigenvalues)
        ellipse = Ellipse(xy=mean, width=width, height=height, angle=angle,
                          **kwargs)
        ax.add_artist(ellipse)

        return ellipse

    ellipses = []
    for std in n_std:
        width, height = 2 * std * np.sqrt(eigenvalues)
        ellipse = Ellipse(xy=mean, width=width, height=height, angle=angle,
                          **kwargs)

        ax.add_artist(ellipse)
        ellipses.append(ellipse)

    return ellipses

def plot_network(Y, X, X_sigma=None, delta=None,
                 z=None, tau_sq=None, normalize=False, figsize=(8, 6),
                 node_color='orangered', color_distance=False,
                 colors=None, alpha=1.0, contour_alpha=0.25,
                 size=300, edgecolors='w',
                 edge_width=0.25, node_labels=None,
                 font_size=12, legend_fontsize=12,
                 with_labels=False, ax=None):
    if ax is None:
        fig, ax = plt.subplots(figsize=figsize)
    else:
        fig = None

    r = np.sqrt((X ** 2).sum(axis=1)).reshape(-1, 1)
    if normalize:
        X = X / r

    cmap = None
    if not isinstance(node_color, np.ndarray):
        cmap = ListedColormap(
            sns.light_palette(node_color, n_colors=np.unique(r).shape[0]))

    G = nx.from_numpy_array(Y)
    if node_labels is not None:
        labels = {node_id : label for node_id, label in enumerate(node_labels)}
    else:
        labels = None

    if z is None:
        if not isinstance(node_color, np.ndarray):
            if color_distance:
                node_color = r.ravel() / r.min()
            else:
                node_color = np.asarray([node_color] * X.shape[0])
    else:
        encoder = LabelEncoder().fit(z)
        if colors is None:
            colors = get_colors(z.ravel())
        node_color = colors[encoder.transform(z)]

        # add a legend
        for i in range(encoder.classes_.shape[0]):
            ax.plot([0], [0], 'o', c=colors[i], label=encoder.classes_[i],
                    markeredgecolor='w', zorder=0)
        ax.plot([0], [0], 'o', markeredgecolor='w', c='w', zorder=0)

    # draw latent position credible interval ellipses
    if X_sigma is not None:
        for i in range(X.shape[0]):
            if isinstance(contour_alpha, np.ndarray):
                calpha = contour_alpha[i]
            else:
                calpha = contour_alpha

            normal_contour(X[i], X_sigma[i], edgecolor='gray',
                           facecolor=node_color[i] if isinstance(node_color, np.ndarray) else 'gray',
                           alpha=calpha, ax=ax, n_std=[2])

    nx.draw_networkx(G, X, edge_color='gray', width=edge_width,
                     node_color=node_color,
                     node_size=size if delta is None else 0,
                     alpha=alpha,
                     cmap=cmap,
                     labels=labels,
                     font_size=font_size,
                     with_labels=with_labels,
                     edgecolors=edgecolors,
                     ax=ax)

    if delta is not None:
        sizes = (delta - delta.min()) / (delta.max() - delta.min())
        ax.scatter(X[:, 0], X[:, 1], s=size * sizes,
                   c='gray')

    if X_sigma is not None:
        ax.collections[0].set_edgecolor(None)
    else:
        ax.collections[0].set_edgecolor('white')

    ax.axis('equal')
    ax.axis('off')

    # draw normal contour if available
    if tau_sq is not None:
        # draw center of latent space
        ax.scatter(0, 0, color='k', marker='+', s=200)

        # draw two standard deviation contour
        normal_contour([0, 0], tau_sq * np.eye(X.shape[1]), n_std=[1],
                       linestyle='--', edgecolor='k',
                       facecolor='none', zorder=1, ax=ax)

    if z is not None:
        ax.legend(loc='lower center', bbox_to_anchor=(0.5, -0.05), ncol=6,
                  fontsize=legend_fontsize)

    return fig, ax

def make_network_animation(filename, Y, X, X_sigma=None,
                           k=0, z=None, tau_sq=None, normalize=True,
                           figsize=(8, 6), node_color='orangered',
                           alpha=1.0, contour_alpha=0.25,
                           size=300, edge_width=0.25,
                           node_labels=None, font_size=12, with_labels=False,
                           layer_labels=None, time_labels=None,
                           title_fmt='{}, {}', border=0.5, duration=1):

    # XXX: hack to shut off plotting within a jupyter notebook...
    plt.ioff()

    n_layers, n_time_steps, _, _ = Y.shape

    if layer_labels is None:
        layer_labels =  ["k = {}".format(k) for k in range(n_layers)]

    if time_labels is None:
        time_labels = ["t = {}".format(t) for t in range(n_time_steps)]

    with tempfile.TemporaryDirectory() as tempdir:
        x_max, y_max = X.max(axis=(0, 1))
        x_min, y_min = X.min(axis=(0, 1))

        pngs = []
        for t in range(Y.shape[1]):
            fig, ax = plot_network(Y[k, t], X[t],
                X_sigma=X_sigma[t] if X_sigma is not None else None,
                z=z, tau_sq=tau_sq,
                normalize=normalize, figsize=figsize, node_color=node_color,
                alpha=alpha, contour_alpha=contour_alpha,
                size=size, edge_width=edge_width,
                node_labels=node_labels, font_size=font_size,
                with_labels=with_labels,)
            ax.set_title(title_fmt.format(layer_labels[k], time_labels[t]))
            ax.set_xlim(x_min - border, x_max + border)
            ax.set_ylim(y_min - border, y_max + border)

            fname = tempfile.TemporaryFile(dir=tempdir, suffix='.png')
            fig.savefig(fname, dpi=100)
            fname.seek(0)
            plt.close(fig)  # necessary to free memory
            pngs.append(fname)

        images = []
        for png in pngs:
            images.append(imageio.imread(png))
        imageio.mimsave(filename, images, duration=duration)

    plt.ion()

def plot_static_sociability(model, k=0, node_labels=None, layer_label=None,
                            ax=None, figsize=(10, 12), color_code=False):
    if ax is None:
        fig, ax = plt.subplots(figsize=figsize)
    else:
        fig = None

    if node_labels is None:
        node_labels = [str(i + 1) for i in range(model.X_.shape[1])]
    node_labels = np.asarray(node_labels)

    order = np.argsort(model.delta_[k])
    odds = np.exp(model.delta_[k][order])
    y_pos = np.arange(node_labels.shape[0])

    if color_code:
        colors = ['steelblue' if odds[i] >= 1. else 'gray' for i in
                  range(len(odds))]
    else:
        colors = 'gray'
    ax.barh(y_pos, odds, align='center', color=colors)
    ax.set_yticks(y_pos)
    ax.set_yticklabels(node_labels[order])
    ax.set_xlabel('odds [$\exp(\delta_k^i)]$')

    if layer_label is not None:
        ax.set_title(layer_label)
    else:
        ax.set_title('k = {}'.format(k))

    return fig, ax

def plot_social_trajectories(
                     model, k=0, q_alpha=0.05, node_list=None, node_colors=None,
                     node_labels=None, layer_label=None, ref_value=0,
                     ref_label=None,
                     plot_hline=True, xlabel='Time', alpha=0.15, fill_alpha=0.2,
                     line_width=3, ax=None, figsize=(10, 6), label_offset=1,
                     fontsize=12, color_code=False):

    n_layers, n_time_steps, n_nodes = model.delta_.shape

    if ax is None:
        fig, ax = plt.subplots(figsize=figsize)
    else:
        fig = None

    ax.set_clip_on(False)

    if node_labels is None:
        node_labels = [str(i + 1) for i in range(model.delta_.shape[2])]
    node_labels = np.asarray(node_labels)

    for i in range(n_nodes):
        if model.X_ is None:
            ax.plot(model.delta_[k, :, i].T, 'k-', alpha=alpha)
        else:
            ax.plot(model.gamma_[k, :, i].T, 'k-', alpha=alpha)

    if node_list is not None:
        node_list = np.asarray(node_list)
        if node_colors is None:
            node_colors = get_colors(np.arange(len(node_list)))

        for i, node_label in enumerate(node_list):
            node_id = np.where(node_labels == node_label)[0].item()
            if model.X_ is None:
                ax.plot(model.delta_[k, :, node_id].T, '--',
                        lw=line_width, c=node_colors[i])
            else:
                ax.plot(model.gamma_[k, :, node_id].T, '--',
                        lw=line_width, c=node_colors[i])
            ax.annotate(node_label,
                        xy=(n_time_steps + label_offset,
                            model.delta_[k, -1, node_id]),
                        color=node_colors[i], fontsize=fontsize,
                        annotation_clip=False)

            if q_alpha is not None and model.X_ is None:
                x_upp = np.zeros(n_time_steps)
                x_low = np.zeros(n_time_steps)
                z_alpha = norm.ppf(1 - q_alpha / 2.)
                ts = np.arange(n_time_steps)
                for t in range(n_time_steps):
                    se = z_alpha * np.sqrt(model.delta_sigma_[k, t, node_id])
                    x_upp[t] = model.delta_[k, t, node_id] + se
                    x_low[t] = model.delta_[k, t, node_id] - se
                ax.fill_between(
                    ts, x_low, x_upp, alpha=fill_alpha, color=node_colors[i])
            elif q_alpha is not None:
                gamma_ci = np.quantile(
                    model.gammas_, [q_alpha/2., 1 - q_alpha/2.], axis=0)
                ax.fill_between(
                    np.arange(n_time_steps), gamma_ci[0, k, :, node_id],
                    gamma_ci[1, k, :, node_id],
                    alpha=fill_alpha, color=node_colors[i])

    if plot_hline:
        ax.hlines(
            ref_value, 0, n_time_steps - 1, lw=2, linestyles='--', color='k')

        if ref_label:
            ax.annotate(ref_label,
                        xy=(n_time_steps + label_offset, ref_value),
                        color='k', fontsize=fontsize)

    # remove spines
    #ax.spines['right'].set_visible(False)
    #ax.spines['left'].set_visible(False)
    #ax.spines['top'].set_visible(False)
    #ax.spines['bottom'].set_visible(False)

    # axis-labels
    ax.set_ylabel('Sociality', fontsize=fontsize)
    ax.set_xlabel(xlabel, fontsize=fontsize)

    if layer_label is not None:
        ax.set_title(layer_label, fontsize=fontsize)
    else:
        ax.set_title('k = {}'.format(k), fontsize=fontsize)

    return fig, ax

def plot_node_trajectories(model, node_list, q_alpha=0.05, node_labels=None,
                           node_colors=None, nrows=None, ncols=1, alpha=0.2,
                           linestyle='o--', fontsize=12,
                           figsize=(10, 8)):

    if nrows is None:
        nrows = model.X_.shape[2]

    fig, axes = plt.subplots(nrows, ncols, figsize=figsize)
    ax = axes.flat

    if node_labels is None:
        node_labels = [i for i in range(model.X_.shape[1])]
    node_labels = np.asarray(node_labels)

    if node_colors is None:
        node_colors = get_colors(np.arange(len(node_list)))

    n_time_steps, n_nodes, n_features = model.Z_.shape
    z_alpha = norm.ppf(1 - q_alpha / 2.)
    ts = np.arange(n_time_steps)
    for i, node_label in enumerate(node_list):
        node_id = np.where(node_labels == node_label)[0].item()
        x_upp = np.zeros(n_time_steps)
        x_low = np.zeros(n_time_steps)
        for p in range(n_features):
            ax[p].plot(ts, model.Z_[:, node_id, p], linestyle,
                       label=node_labels[node_id], c=node_colors[i])
            for t in range(n_time_steps):
                se = z_alpha * np.sqrt(model.Z_sigma_[t, node_id, p, p])
                x_upp[t] = model.Z_[t, node_id, p] + se
                x_low[t] = model.Z_[t, node_id, p] - se
            ax[p].fill_between(
                ts, x_low, x_upp, alpha=alpha, color=node_colors[i])

    # accomodate legends and title
    ax[0].legend(bbox_to_anchor=(1.04, 1), loc='upper left', fontsize=fontsize)
    ax[-1].set_xlabel('t')
    for p in range(n_features):
        #ax[p].set_title('p = {}'.format(p + 1), fontsize=fontsize)
        ax[p].hlines(0, 1, n_time_steps, lw=2, linestyles='dotted', color='k', alpha=0)
        ax[p].set_ylabel('Latent Position [h = {}]'.format(p + 1),
                         fontsize=fontsize)
        ax[p].tick_params(axis='x', labelsize=fontsize)
        ax[p].tick_params(axis='y', labelsize=fontsize)

    plt.subplots_adjust(right=0.7)

    return fig, ax

def sample_distances(model, k, t, i, j, n_reps=1000, random_state=123):
    rng = check_random_state(random_state)

    Xi = rng.multivariate_normal(model.X_[t, i], model.X_sigma_[t, i],
                                 size=n_reps)
    Xj = rng.multivariate_normal(model.X_[t, j], model.X_sigma_[t, j],
                                 size=n_reps)

    if k == 0:
        lmbdak = np.zeros((n_reps, model.lambda_.shape[1]))
        for p in range(model.lambda_.shape[1]):
            lmbdak[:, p] = (
                2 * rng.binomial(1, model.lambda_proba_[p], size=n_reps) - 1)
    else:
        lmbdak = rng.multivariate_normal(
            model.lambda_[k], model.lambda_sigma_[k], size=n_reps)

    return np.sum(lmbdak * Xi * Xj, axis=1)

def plot_pairwise_distances(model, node_i, node_j,
                            node_labels=None,
                            layer_labels=None, q_alpha=0.05, n_reps=1000,
                            random_state=123, alpha=0.2, linestyle='--',
                            figsize=(10, 8), ax=None):
    if ax is not None:
        fig = None
    else:
        fig, ax = plt.subplots(figsize=figsize)

    if node_labels is None:
        node_labels = [i for i in range(model.X_.shape[1])]
    node_labels = np.asarray(node_labels)

    n_layers, n_time_steps, n_nodes, _ = model.dist_.shape
    ts = np.arange(n_time_steps)
    i = np.where(node_labels == node_i)[0].item()
    j = np.where(node_labels == node_j)[0].item()
    for k in range(n_layers):
        if layer_labels is None:
            label =  'k = {}'.format(k)
        else:
            label = layer_labels[k]

        if q_alpha is None:
            ax.plot(ts, model.dist_[k, :, i, j], linestyle,
                    label=label)
        else:
            dist_mean = np.zeros(n_time_steps)
            dist_low = np.zeros(n_time_steps)
            dist_upp = np.zeros(n_time_steps)
            for t in range(n_time_steps):
                dist = sample_distances(
                    model, k, t, i, j, n_reps=n_reps, random_state=random_state)
                dist_mean[t] = dist.mean()
                dist_low[t] = np.quantile(dist, q=q_alpha / 2.)
                dist_upp[t] = np.quantile(dist, q=1 - q_alpha / 2.)

            ax.plot(ts, dist_mean, linestyle, label=label)
            ax.fill_between(ts, dist_low, dist_upp, alpha=alpha)
        break

    # accomodate legends and title
    ax.legend(bbox_to_anchor=(1.04, 1), loc='upper left')
    ax.set_xlabel('t')
    ax.set_ylabel('Distances ({} - {})'.format(node_i, node_j))

    return fig, ax

def sample_link_probability(model, k, t, i, j, n_reps=1000, random_state=123):
    rng = check_random_state(random_state)

    deltai = rng.normal(
        loc=model.delta_[k, t, i], scale=np.sqrt(model.delta_sigma_[k, t, i]),
        size=n_reps)
    deltaj = rng.normal(
        loc=model.delta_[k, t, j], scale=np.sqrt(model.delta_sigma_[k, t, j]),
        size=n_reps)

    Xi = rng.multivariate_normal(model.X_[t, i], model.X_sigma_[t, i],
                                 size=n_reps)
    Xj = rng.multivariate_normal(model.X_[t, j], model.X_sigma_[t, j],
                                 size=n_reps)

    if k == 0:
        lmbdak = np.zeros((n_reps, model.lambda_.shape[1]))
        for p in range(model.lambda_.shape[1]):
            lmbdak[:, p] = (
                2 * rng.binomial(1, model.lambda_proba_[p], size=n_reps) - 1)
    else:
        lmbdak = rng.multivariate_normal(
            model.lambda_[k], model.lambda_sigma_[k], size=n_reps)

    return expit(deltai + deltaj + np.sum(lmbdak * Xi * Xj, axis=1))

def forecast_link_probability(model, k, i, j, horizon=1, n_reps=1000,
                              random_state=123):
    rng = check_random_state(random_state)
    n_features = model.X_.shape[-1]

    if k == 0:
        lmbdak = np.zeros((n_reps, model.lambda_.shape[1]))
        for p in range(model.lambda_.shape[1]):
            lmbdak[:, p] = (
                2 * rng.binomial(1, model.lambda_proba_[p], size=n_reps) - 1)
    else:
        lmbdak = rng.multivariate_normal(
            model.lambda_[k], model.lambda_sigma_[k], size=n_reps)

    deltai = rng.normal(
        loc=model.delta_[k, -1, i], scale=np.sqrt(model.delta_sigma_[k, -1, i]),
        size=n_reps)
    deltaj = rng.normal(
        loc=model.delta_[k, -1, j], scale=np.sqrt(model.delta_sigma_[k, -1, j]),
        size=n_reps)

    Xi = rng.multivariate_normal(model.X_[-1, i], model.X_sigma_[-1, i],
                                 size=n_reps)
    Xj = rng.multivariate_normal(model.X_[-1, j], model.X_sigma_[-1, j],
                                 size=n_reps)

    pis = np.zeros((horizon, n_reps))
    for h in range(horizon):
        deltai = deltai + rng.normal(
            loc=0, scale=np.sqrt(model.sigma_sq_delta_), size=n_reps)
        deltaj = deltaj + rng.normal(
            loc=0, scale=np.sqrt(model.sigma_sq_delta_), size=n_reps)

        Xi = Xi + rng.multivariate_normal(
            np.zeros(n_features), model.sigma_sq_ * np.eye(n_features),
            size=n_reps)
        Xj = Xj + rng.multivariate_normal(
            np.zeros(n_features), model.sigma_sq_ * np.eye(n_features),
            size=n_reps)

        pis[h] = expit(deltai + deltaj + np.sum(lmbdak * Xi * Xj, axis=1))

    return pis

def plot_pairwise_probabilities(model, node_i, node_j, horizon=0,
                                node_labels=None,
                                layer_labels=None, q_alpha=0.05, n_reps=1000,
                                random_state=123, alpha=0.2, linestyle='--',
                                fontsize=16, figsize=(10, 8), ax=None):

    if ax is None:
        fig, ax = plt.subplots(figsize=figsize)
    else:
        fig = None

    if node_labels is None:
        node_labels = [i for i in range(model.X_.shape[1])]
    node_labels = np.asarray(node_labels)

    n_layers, n_time_steps, n_nodes, _ = model.probas_.shape
    ts = np.arange(n_time_steps + horizon)
    i = np.where(node_labels == node_i)[0].item()
    j = np.where(node_labels == node_j)[0].item()
    for k in range(n_layers):
        if layer_labels is None:
            label =  'k = {}'.format(k)
        else:
            label = layer_labels[k]

        if q_alpha is None:
            ax.plot(ts, model.probas_[k, :, i, j], linestyle,
                    label=label)
        else:
            pi_mean = np.zeros(n_time_steps + horizon)
            pi_low = np.zeros(n_time_steps + horizon)
            pi_upp = np.zeros(n_time_steps + horizon)
            for t in range(n_time_steps):
                pis = sample_link_probability(
                    model, k, t, i, j, n_reps=n_reps, random_state=random_state)
                pi_mean[t] = pis.mean()
                pi_low[t] = np.quantile(pis, q=q_alpha / 2.)
                pi_upp[t] = np.quantile(pis, q=1 - q_alpha / 2.)

            if horizon > 0:
                pis = forecast_link_probability(
                    model, k, i, j, horizon=horizon, n_reps=n_reps,
                    random_state=random_state)

                for h in range(horizon):
                    pi_mean[n_time_steps + h] = pis[h].mean()
                    pi_low[n_time_steps + h] = (
                        np.quantile(pis[h], q=q_alpha / 2.))
                    pi_upp[n_time_steps + h] = (
                        np.quantile(pis[h], q=1 - q_alpha / 2.))

            ax.plot(ts, pi_mean, linestyle, label=label)
            ax.fill_between(ts, pi_low, pi_upp, alpha=alpha)

    # accomodate legends and title
    ax.legend(bbox_to_anchor=(1.04, 1), loc='upper left', fontsize=fontsize)
    ax.set_xlabel('t')
    ax.set_ylabel('Link Probability ({} - {})'.format(node_i, node_j), fontsize=fontsize)

    return fig, ax

def plot_homophily_matrix(model, q_alpha=0.05,
                           layer_labels=None, height=0.5, hspace=1.,
                           fontsize=12, figsize=(12, 6)):
    n_layers, n_features = model.lambda_.shape

    if layer_labels is None:
        layer_labels = ['k = {}'.format(k + 1) for k in range(n_layers)]

    fig, axes = plt.subplots(n_layers, 1, figsize=figsize, sharex=True)

    z_alpha = norm.ppf(1 - q_alpha / 2.)
    for p in range(n_features):
        xerr = z_alpha * np.sqrt(model.lambda_sigma_[:, p, p])

        for k in range(n_layers):
            colors = 'red' if model.lambda_[k, p] > 0 else 'blue'
            axes[k].hlines(k + hspace * p, 0, model.lambda_[k, p], lw=1,
                           color=colors, linestyles='--')
            axes[k].errorbar(model.lambda_[k, p], k + hspace * p,
                             fmt='o',
                             xerr=xerr[k], ecolor='k', capsize=5,
                             color='k', markersize=9, markeredgecolor='w')

        # add text
        for k in range(n_layers):
            align = 'right' if model.lambda_[k, p]  >= 0 else 'left'

            lmbda = model.lambda_[k, p]
            if k == 0:
                txt = '{}    (d = {})'.format(lmbda, p+1)
            else:
                txt = '{:.3f} ({:.3f}, {:.3f})'.format(
                    lmbda, lmbda - xerr[k], lmbda + xerr[k])
            axes[k].text(lmbda, k + hspace * p - 0.1,
                         txt, horizontalalignment=align)

    for k in range(n_layers):
        axes[k].set_yticks([k + hspace / n_features])
        axes[k].set_yticklabels([layer_labels[k]], fontsize=fontsize)
        axes[k].invert_yaxis()

        if k != (n_layers - 1):
            axes[k].spines['bottom'].set_visible(False)
            axes[k].tick_params(bottom=False)

    axes[-1].set_xlabel('Homophily Parameter ($\lambda_{kd}$)',
                        fontsize=fontsize)

    x_max = max([ax.get_xlim()[1] for ax in axes.flat])
    for k in range(n_layers):
        if np.all(model.lambda_ >= 0):
            axes[k].set_xlim(0, axes[k].get_xlim()[1])
        else:
            axes[k].vlines(0, k, k + hspace * (n_features - 1),
                           linestyles='dotted', color='k')
        sns.despine(ax=axes[k], bottom=True)

    sns.set_style('white')
    plt.subplots_adjust(hspace=0.5)

    return fig, axes

#def plot_homophily_matrix(model, q_alpha=0.05, colors=None,
#                          layer_labels=None, height=0.5,
#                          fontsize=12, figsize=(12, 6)):
#    n_layers, n_features = model.lambda_.shape
#
#    if layer_labels is None:
#        layer_labels = ['k = {}'.format(k + 1) for k in range(n_layers)]
#
#    fig, ax = plt.subplots(figsize=figsize)
#
#    if colors is None:
#        colors = get_colors(np.arange(n_layers))
#
#    z_alpha = norm.ppf(1 - q_alpha / 2.)
#    for p in range(n_features):
#        xerr = z_alpha * np.sqrt(model.lambda_sigma_[:, p, p])
#
#        #colors = ['red' if model.lambda_[k, p] > 0 else 'blue' for
#        #          k in range(n_layers)]
#        ax.hlines(np.arange(n_layers) + 0.5 * p, 0, model.lambda_[:, p], lw=1,
#                  color=colors, linestyles='--')
#        ax.errorbar(model.lambda_[:, p], np.arange(n_layers) + 0.5 * p,
#                    fmt='o',
#                    xerr=xerr, ecolor='k', capsize=5,
#                    color='k', markersize=9, markeredgecolor='w')
#
#        # add text
#        for k in range(n_layers):
#            align = 'right' if model.lambda_[k, p]  >= 0 else 'left'
#
#            lmbda = model.lambda_[k, p]
#            if k == 0:
#                txt = '{}'.format(lmbda)
#            else:
#                txt = '{:.3f} ({:.3f}, {:.3f})'.format(
#                    lmbda, lmbda - xerr[k], lmbda + xerr[k])
#            ax.text(lmbda, k + 0.5 * p - 0.1, txt, horizontalalignment=align)
#
#    ax.set_yticks(np.arange(n_layers) + 0.25)
#    ax.set_yticklabels(layer_labels, fontsize=fontsize)
#    ax.invert_yaxis()
#    #ax.set_title('p = {}'.format(p + 1), fontsize=fontsize)
#
#    ax.set_xlabel('Homophily Parameter ($\lambda_{kp}$)',
#                  fontsize=fontsize)
#
#    #x_max = max([ax.get_xlim()[1] for ax in axes.flat])
#    if np.all(model.lambda_ >= 0):
#        ax.set_xlim(0, ax.get_xlim()[1])
#    else:
#        ax.vlines(0, 0, n_layers - 0.5 * (n_features - 1), linestyles='dotted', color='k')
#    sns.despine(ax=ax, bottom=True)
#
#    sns.set_style('white')
#
#    return fig, ax

def plot_lambda(model, q_alpha=0.05, layer_labels=None, height=0.5,
                fontsize=12,
                figsize=(12, 6), include_gridlines=False):
    n_layers, n_features = model.lambda_.shape

    if layer_labels is None:
        layer_labels = ['k = {}'.format(k + 1) for k in range(n_layers)]

    if include_gridlines:
        sns.set_style('whitegrid')

    fig, axes = plt.subplots(n_features, 1, figsize=figsize, sharex=True)
    colors = [to_hex(c) for c in sns.color_palette(
              'muted', n_colors=n_layers, desat=0.75)]

    z_alpha = norm.ppf(1 - q_alpha / 2.)
    for p, ax in enumerate(axes.flat):
        xerr = z_alpha * np.sqrt(model.lambda_sigma_[:, p, p])

        colors = ['red' if model.lambda_[k, p] > 0 else 'blue' for
                  k in range(n_layers)]
        ax.hlines(np.arange(n_layers), 0, model.lambda_[:, p], lw=1,
                  color=colors, linestyles='--')
        ax.errorbar(model.lambda_[:, p], np.arange(n_layers), fmt='o',
                    xerr=xerr, ecolor='k', capsize=5,
                    color='k', markersize=9, markeredgecolor='w')

        # add text
        for k in range(n_layers):
            align = 'right' if model.lambda_[k, p]  >= 0 else 'left'

            lmbda = model.lambda_[k, p]
            if k == 0:
                txt = '{}'.format(lmbda)
            else:
                txt = '{:.3f} ({:.3f}, {:.3f})'.format(
                    lmbda, lmbda - xerr[k], lmbda + xerr[k])
            ax.text(lmbda, k - 0.1, txt, horizontalalignment=align)

        ax.set_yticks(np.arange(n_layers))
        ax.set_yticklabels(layer_labels, fontsize=fontsize)
        ax.invert_yaxis()
        ax.set_title('h = {}'.format(p + 1), fontsize=fontsize)

        axes.flat[-1].set_xlabel('Homophily Parameter ($\lambda_{kh}$)',
                                 fontsize=fontsize)

    x_max = max([ax.get_xlim()[1] for ax in axes.flat])
    for ax in axes.flat:
        if np.all(model.lambda_ >= 0):
            ax.set_xlim(0, x_max)
        else:
            ax.vlines(0, 0, n_layers - 1, linestyles='dotted', color='k')
        sns.despine(ax=ax, bottom=True)

    sns.set_style('white')

    return fig, axes

def plot_network_statistics(stat_sim, stat_obs=None, nrow=1, ncol=None,
                            time_labels=None, stat_label='Statistic',
                            time_step=1,
                            layer_labels=None, figsize=(16, 10),
                            xlabel='Time'):
    n_layers, n_time_steps, _ = stat_sim.shape

    if ncol is None:
        ncol = n_layers
    fig, axes = plt.subplots(nrow, ncol, sharey=True, figsize=figsize)

    if time_labels is None:
        time_labels = np.arange(n_time_steps) + 1

    if layer_labels is None:
        layer_labels = np.arange(n_layers) + 1

    for k, ax in enumerate(axes.flat):
        data = pd.DataFrame()
        for t in range(n_time_steps):
            data[time_labels[t]] = stat_sim[k, t]

        if stat_obs is not None:
            ax.plot(np.arange(n_time_steps), stat_obs[k], 'o--', c='k')
        sns.boxplot(x='variable', y='value', data=pd.melt(data),
                    ax=ax, color='white')

        ax.set_xticklabels(time_labels[::time_step], rotation=45, fontsize=12)
        plt.setp(ax.artists, edgecolor='black')
        plt.setp(ax.lines, color='black')

        ax.set_xticks([i for i in range(0, n_time_steps, time_step)])
        ax.tick_params(axis='y', labelsize=16)
        ax.tick_params(axis='x', labelsize=16)

        ax.grid(axis='x')
        ax.set_title(layer_labels[k], fontsize=24)
        ax.set_xlabel(xlabel, fontsize=24)

        if k == 0:
            ax.set_ylabel(stat_label, fontsize=24)
        else:
            ax.set_ylabel('')

    return fig, axes

max_stars_repo_name: joshloyal/multidynet

max_issues_repo_name: joshloyal/multidynet

max_forks_repo_name: joshloyal/multidynet

reason: import numpy,from scipy,import networkx


text:
# Copyright 2019 ducandu GmbH. All Rights Reserved.
# (this is a modified version of the Apache 2.0 licensed RLgraph file of the same name).
#
# 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.
# ==============================================================================

import numpy as np

from surreal.components.preprocessors.preprocessor import Preprocessor

class ImageCrop(Preprocessor):
    """
    Crops one or more images to a new size without touching the color channel.
    """
    def __init__(self, x=0, y=0, width=0, height=0):
        """
        Args:
            x (int): Start x coordinate.
            y (int): Start y coordinate.
            width (int): Width of resulting image.
            height (int): Height of resulting image.
        """
        super().__init__()
        self.x = x
        self.y = y
        self.width = width
        self.height = height

        assert self.x >= 0
        assert self.y >= 0
        assert self.width > 0
        assert self.height > 0

    def call(self, inputs):
        """
        Images come in with either a batch dimension or not.
        """
        #if isinstance(inputs, list):
        #    inputs = np.asarray(inputs)
        # Preserve batch dimension.
        if len(inputs.shape) >= 4:
            return inputs[:, self.y:self.y + self.height, self.x:self.x + self.width]
        else:
            return inputs[self.y:self.y + self.height, self.x:self.x + self.width]

ext: py

max_stars_repo_name: rosea-tf/surreal

max_issues_repo_name: rosea-tf/surreal

max_forks_repo_name: rosea-tf/surreal

reason: import numpy


text:
"""A proper search function for threshold estimation.

Author: Yuhuang Hu
Email : yuhuang.hu@ini.uzh.ch
"""
import os

import subprocess
import numpy as np
import h5py

import matplotlib
import matplotlib.pyplot as plt

from v2ecore.v2e_utils import select_events_in_roi

def evaluate_threshold(
        cfg,
        idx, threshold, event_count_diffs,
        ref_event_count=0):

    # v2e commands
    v2e_command = [
        "v2e.py",
        "-i", cfg.input,
        "-o", cfg.output_folder,
        "--dvs_emulator_seed", "42",
        "--start_time", "{}".format(cfg.start),
        "--stop_time", "{}".format(cfg.stop),
        "--overwrite",
        "--unique_output_folder", "false",
        "--no_preview",
        "--skip_video_output",
        "--pos_thres", "{}".format(threshold),
        "--neg_thres", "{}".format(threshold),
        "--sigma_thres", "0",
        "--cutoff_hz", "0",
        "--leak_rate_hz", "0",
        "--shot_noise_rate_hz", "0",
        "--input_frame_rate", "{}".format(cfg.input_frame_rate),
        #  "--input_slowmotion_factor", "17.866666708",
        "--dvs_h5", "dvs_events.h5",
        "--dvs_aedat2", "None",
        "--dvs_text", "None",
        "--dvs_exposure", "duration", "0.001",
        #  "--disable_slomo",
        "--auto_timestamp_resolution", "True",
        "--slomo_model", cfg.slomo_model
        ]

    # TODO: a proper way to avoid repeatly running slomo model
    #  slomo_vid_path = os.path.join(cfg.output_folder, "video_slomo.avi")
    #  if os.path.isfile(slomo_vid_path):
    #      # has slow motion file
    #      v2e_command += [
    #          "--disable_slomo",
    #          "--auto_timestamp_resolution", "False"]
    #  else:
    #      v2e_command += [
    #          "--auto_timestamp_resolution", "True",
    #          "--slomo_model", cfg.slomo_model]

    if event_count_diffs[idx] is None:
        # evaluate and assign
        subprocess.run(v2e_command)

        # load events
        data = h5py.File(os.path.join(
            cfg.output_folder, "dvs_events.h5"), "r")
        curr_events = data["events"][()]
        curr_events = select_events_in_roi(curr_events, cfg.x, cfg.y)
        curr_event_count = curr_events.shape[0]

        # assign event counts
        event_count_diffs[idx] = abs(
            curr_event_count-ref_event_count)

    return event_count_diffs[idx], event_count_diffs

def threshold_estimator(
        cfg,
        low_thres_idx, high_thres_idx, threshold_range,
        event_count_diffs, ref_event_count=0):
    """Binary search for estimating threshold.

    Assume only single minimum.

    # Arguments
    cfg: a result of parsing arguments, NameSpace that contains
        all necessary arguments. such as IO, V2E settings, etc
    low_thres: float
        threshold boundary (low)
    high_thres: float
        threshold boundary (high)
    threshold_range: list
        a list of thresholds, all values larger than 0.
        assume small to large
    event_count_diffs: list
        a list of event count absolute differences
    v2e_commands: for subprocess to call
        should be able to modify threshold parameter.
    ref_event_count: int
        should be the reference event count.
    """

    # boundary
    start_idx, end_idx = 0, len(threshold_range)-1

    curr_mid_idx = (high_thres_idx+low_thres_idx)//2
    while (low_thres_idx < high_thres_idx):
        print("current range: {}-{}".format(
            threshold_range[low_thres_idx],
            threshold_range[high_thres_idx]))
        # current middle point
        curr_mid_idx = (high_thres_idx+low_thres_idx)//2
        curr_mid = threshold_range[curr_mid_idx]

        curr_mid_count, event_count_diffs = evaluate_threshold(
            cfg, curr_mid_idx, curr_mid, event_count_diffs, ref_event_count)

        # terminate condition, if neighbors are all higher than current,
        # then terminate
        # left neighbor
        if curr_mid_idx > 5:
            left_count, event_count_diffs = evaluate_threshold(
                cfg, curr_mid_idx-5, threshold_range[curr_mid_idx-5],
                event_count_diffs, ref_event_count)
        else:
            left_count = np.inf

        # right neighbor
        if curr_mid_idx < end_idx-5:
            right_count, event_count_diffs = evaluate_threshold(
                cfg, curr_mid_idx+5, threshold_range[curr_mid_idx+5],
                event_count_diffs, ref_event_count)
        else:
            right_count = np.inf

        if (curr_mid_idx == start_idx+5 or left_count >= curr_mid_count) and \
                (curr_mid_idx == end_idx-5 or right_count >= curr_mid_count):
            return curr_mid, event_count_diffs
        elif (curr_mid_idx > start_idx and left_count < curr_mid_count):
            high_thres_idx = curr_mid_idx
        elif (curr_mid_idx < end_idx and right_count < curr_mid_count):
            low_thres_idx = curr_mid_idx+1

    return curr_mid, event_count_diffs


text:
# -*- coding: utf-8 -*-
"""
@author: Xinrui Zhou
"""

import numpy as np

from numpy import inf, tan, arctan2, arctan

class CorotationalBeamElement2D():
    """It is a class formulating two-dimensional corotational beam elements."""

    def __init__(self, beamtype, youngs_modulus, poisson_ratio, width, height, initial_coordinate_node_1, initial_coordinate_node_2, element_id):
        """ Initialize a 2D co-rotational beam element.

        Parameters
        ----------
        beamtype: str
            the type of beam element. Here only Bernoulli element implemented.
        youngs_modulus: int/float
            Young's modulus.
        poisson_ratio: float
            Poisson ratio, not interesting for Bernoulli elements.
        width: int/float
            width of the beam, perpendicular to local z-axis.
        height: int/float
            height of the beam, perpendicular to local y-axis.
        initial_coordinate_node_1: numpy.ndarray
            initial coordinate of the 1st node: [X1, Z1].T.
        initial_coordinate_node_2: numpy.ndarray
            initial coordinate of the 2nd node: [X2, Z2].T.
        element_id: int
            The number of the element in the system.

        Returns
        -------
        None.

        """
        self.beamtype = beamtype
        self.youngs_modulus = youngs_modulus
        self.poisson_ratio = poisson_ratio
        self.width = width
        self.height = height
        self.initial_coordinate_node_1 = initial_coordinate_node_1
        self.initial_coordinate_node_2 = initial_coordinate_node_2

        self.initial_length = self.calculate_initial_length()
        self.initial_angle = self.calculate_initial_angle()
        self.area = self.calculate_area()
        self.moment_of_inertia = self.calculate_moment_of_inertia()

        self.analysis = "elastic"  # default value
        self.global_displacement = np.zeros((6, 1), dtype=float)
        self.element_freedom_table = np.linspace(
            3 * element_id, 3 * element_id + 5, num=6, dtype=int)

    def calculate_area(self):
        """Calculate the cross-sectional area of the beam element.

        Parameters
        ----------
        None.

        Returns
        -------
        cross-sectional area: A = b * h

        """
        return self.width * self.height

    def calculate_moment_of_inertia(self):
        """Calculate the moment of inertia of the beam element.

        Parameters
        ----------
        None.

        Returns
        -------
        moment of inertia: I = 1/12 * b * h^3

        """
        return 1/12 * self.width * (self.height) ** 3

    def current_coordinate_node_1(self):
        """Calculate the coodinate of node 1 in current configuration.

        Parameters
        ----------
        None.

        Returns
        -------
        current coordinate of the 1st node: x1 = X1 + u1.

        """
        return self.initial_coordinate_node_1 + self.global_displacement[0: 2]

    def current_coordinate_node_2(self):
        """Calculate the coodinate of node 2 in current configuration.

        Parameters
        ----------
        None.

        Returns
        -------
        current coordinate of the 2nd node: x2 = X2 + u2.

        """
        return self.initial_coordinate_node_2 + self.global_displacement[3: 5]

    def apply_hardening_model(self, hardening_model, gauss_number, yield_stress, kinematic_hardening_modulus, plastic_modulus, saturation_stress,
                              modified_modulus, exponent):
        """Apply hardening model on the beam element.

        Parameters
        ----------
        hardening_model: str
            hardening model. linear hardening, exponential hardening and ramberg-osgood hardening are implemented.
        gauss_number: tuple
            Number of Gauss points on xi-axis and eta-axis. Set as a tuple with 2 numbers, such as (2, 2)
        yield_stress: int/float
            Value of yield stress.
        kinematic_hardening_modulus: int/float
            Value of kinematic hardening modulus.
        plastic_modulus: int/float
            Value of plastic modulus.
        saturation_stress: int/float
            Value of saturation stress.
        modified_modulus: int/float
            Value of modified modulus.
        exponent: int/float
            Value of exponent.

        Returns
        -------
        None.

        """
        self.num_of_gauss_locations_xi = gauss_number[0]
        self.num_of_gauss_locations_eta = gauss_number[1]

        self.num_of_gauss_locations = self.num_of_gauss_locations_xi * \
            self.num_of_gauss_locations_eta

        [gauss_locations_xi, weights_xi] = np.polynomial.legendre.leggauss(
            self.num_of_gauss_locations_xi)
        [gauss_locations_eta, weights_eta] = np.polynomial.legendre.leggauss(
            self.num_of_gauss_locations_eta)

        global_gauss_locations_mat = np.zeros(
            (2, self.num_of_gauss_locations), dtype=float)

        weights_mat = np.zeros(self.num_of_gauss_locations, dtype=float)

        for ixi in range(self.num_of_gauss_locations_xi):
            for jeta in range(self.num_of_gauss_locations_eta):
                x, z = self.map_local_coordinate_to_global_coordinate(
                    gauss_locations_xi[ixi], gauss_locations_eta[jeta])

                global_gauss_locations_mat[0, ixi *
                                           self.num_of_gauss_locations_eta + jeta] = x
                global_gauss_locations_mat[1, ixi *
                                           self.num_of_gauss_locations_eta + jeta] = z
                weights_mat[ixi * self.num_of_gauss_locations_eta +
                            jeta] = weights_xi[ixi] * weights_eta[jeta]

        self.global_gauss_locations_mat = global_gauss_locations_mat
        self.weights_mat = weights_mat

        if hardening_model == "linear hardening":

            self.analysis = hardening_model
            self.yield_stress = yield_stress
            self.kinematic_hardening_modulus = kinematic_hardening_modulus
            self.plastic_modulus = plastic_modulus

            self.stress = np.zeros(self.num_of_gauss_locations, dtype=float)
            self.plastic_strain = np.zeros(
                self.num_of_gauss_locations, dtype=float)
            self.back_stress = np.zeros(
                self.num_of_gauss_locations, dtype=float)
            self.tangent_modulus = self.youngs_modulus * \
                np.ones(self.num_of_gauss_locations, dtype=float)
            self.internal_hardening_variable = np.zeros(
                self.num_of_gauss_locations, dtype=float)

        elif hardening_model == "exponential hardening":

            self.analysis = hardening_model
            self.yield_stress = yield_stress
            self.kinematic_hardening_modulus = kinematic_hardening_modulus
            self.saturation_stress = saturation_stress
            self.exponent = exponent

            self.stress = np.zeros(self.num_of_gauss_locations, dtype=float)
            self.plastic_strain = np.zeros(
                self.num_of_gauss_locations, dtype=float)
            self.back_stress = np.zeros(
                self.num_of_gauss_locations, dtype=float)
            self.tangent_modulus = self.youngs_modulus * \
                np.ones(self.num_of_gauss_locations, dtype=float)
            self.internal_hardening_variable = np.zeros(
                self.num_of_gauss_locations, dtype=float)

        elif hardening_model == "ramberg-osgood hardening":

            self.analysis = hardening_model
            self.yield_stress = yield_stress
            self.kinematic_hardening_modulus = kinematic_hardening_modulus
            self.modified_modulus = modified_modulus
            self.exponent = exponent

            self.stress = np.zeros(self.num_of_gauss_locations, dtype=float)
            self.plastic_strain = np.zeros(
                self.num_of_gauss_locations, dtype=float)
            self.back_stress = np.zeros(
                self.num_of_gauss_locations, dtype=float)
            self.tangent_modulus = self.youngs_modulus * \
                np.ones(self.num_of_gauss_locations, dtype=float)
            self.internal_hardening_variable = np.zeros(
                self.num_of_gauss_locations, dtype=float)

        else:
            print("Three hardening models are implemented here: 1. linear hardening; 2. exponential hardening; 3. ramberg-osgood hardening.")
            raise ValueError("Wrong hardening model!")

    def local_displacement(self):
        """Calculate local displacement vector at current state.

        Parameters
        ----------
        None.

        Returns
        -------
        p_l: numpy.ndarray
            local displacement p_l = [u_bar, theta1_bar, theta2_bar].T

        """
        global_rotation_node_1 = float(self.global_displacement[2])
        global_rotation_node_2 = float(self.global_displacement[5])

        u_bar = self.current_length() - self.initial_length
        tanbeta = self.current_angle_trigonometric("local")

        if tanbeta == inf:
            theta1_bar = arctan(-1 /
                                tan(global_rotation_node_1 + self.initial_angle))
            theta2_bar = arctan(-1 /
                                tan(global_rotation_node_2 + self.initial_angle))
        else:
            theta1_bar = arctan((tan(global_rotation_node_1 + self.initial_angle) - tanbeta) / (
                1 + tan(global_rotation_node_1 + self.initial_angle) * tanbeta))
            theta2_bar = arctan((tan(global_rotation_node_2 + self.initial_angle) - tanbeta) / (
                1 + tan(global_rotation_node_2 + self.initial_angle) * tanbeta))

        return np.array([[u_bar, theta1_bar, theta2_bar]]).T

    def calculate_initial_length(self):
        """Calculate the initial length of the beam element.

        Parameters
        ----------
        None.

        Returns
        -------
        L: float

        """
        return np.linalg.norm(self.initial_coordinate_node_2 -
                              self.initial_coordinate_node_1)

    def calculate_initial_angle(self):
        """Calculate the initial angle of the beam element.

        Parameters
        ----------
        None.

        Returns
        -------
        beta_0: float

        """
        return float(arctan2(self.initial_coordinate_node_2[1] -
                     self.initial_coordinate_node_1[1],
                     self.initial_coordinate_node_2[0] -
                     self.initial_coordinate_node_1[0]))

    def current_length(self):
        """Calculate the current length of the beam element.

        Parameters
        ----------
        None.

        Returns
        -------
        l: float

        """
        return np.linalg.norm(self.current_coordinate_node_2() -
                              self.current_coordinate_node_1())

    def current_angle_trigonometric(self, str):
        """Calculate the sine and cosine value of current angle of the beam element.

        Parameters
        ----------
        str: str
            can be 'local' or 'global'.

        Returns
        -------
        cos(beta): float
        sin(beta): float
        tan(beta): float

            If calculate local displacement, the function returns only tangent;
            If calculate global matrix or force, return cosine and sine.
        """
        x1 = self.current_coordinate_node_1()
        x2 = self.current_coordinate_node_2()
        l = self.current_length()
        if str == "global":
            sinbeta = float((x2[1] - x1[1]) / l)
            cosbeta = float((x2[0] - x1[0]) / l)
            return cosbeta, sinbeta
        elif str == "local":
            tanbeta = float((x2[1] - x1[1]) / (x2[0] - x1[0]))
            return tanbeta
        else:
            raise ValueError("Wrong input!")

    def local_stiffness_force(self):
        """Calculate the local stiffness matrix kl and local force vector fl of the beam element.

        Parameters
        ----------
        None.

        Returns
        -------
        kl: np.ndarray
            local stiffness matrix.
        fl: np.ndarray
            local force vector

        """
        E, A, I, L = self.youngs_modulus, self.area, self.moment_of_inertia, self.initial_length

        if self.analysis == "elastic":
            k_l = np.array([[A,  0.,    0.],
                           [0., 4 * I, 2 * I],
                           [0., 2 * I, 4 * I]])
            k_l *= E / L
            f_l = k_l @ self.local_displacement()

        else:
            kl_11, kl_22, kl_33, kl_12, kl_13, kl_23, N, M1, M2 = 0., 0., 0., 0., 0., 0., 0., 0., 0.
            for n in range(self.num_of_gauss_locations):
                x, z = self.global_gauss_locations_mat[0: 2, n]
                fac = self.width * self.height * L / 4

                kl_11 += fac * self.weights_mat[n] * \
                    self.tangent_modulus[n] / L ** 2
                kl_22 += fac * self.weights_mat[n] * self.tangent_modulus[n] * z ** 2 * (4 / L - 6 *
                                                                                x / L ** 2) ** 2
                kl_33 += fac * self.weights_mat[n] * self.tangent_modulus[n] * z ** 2 * (2 / L - 6 *
                                                                                x / L ** 2) ** 2
                kl_12 += fac * self.weights_mat[n] * self.tangent_modulus[n] * z * (
                    4 / L ** 2 - 6 * x / L ** 3)
                kl_13 += fac * self.weights_mat[n] * self.tangent_modulus[n] * z * (
                    2 / L ** 2 - 6 * x / L ** 3)
                kl_23 += fac * self.weights_mat[n] * self.tangent_modulus[n] * z ** 2 * (4 / L - 6 * x / (
                    L) ** 2) * (2 / L - 6 * x / L ** 2)

                N += fac * self.weights_mat[n] * \
                    self.stress[n] / L
                M1 += fac * self.weights_mat[n] * self.stress[n] * z * (
                    4 / L - 6 * x / L ** 2)
                M2 += fac * self.weights_mat[n] * self.stress[n] * z * (
                    2 / L - 6 * x / L ** 2)

            k_l = np.array([[kl_11, kl_12, kl_13],
                            [kl_12, kl_22, kl_23],
                            [kl_13, kl_23, kl_33]])

            f_l = np.array([[N, M1, M2]]).T

        return k_l, f_l

    def map_local_coordinate_to_global_coordinate(self, xi, eta):
        """Define the linear map between Gaussian space and local space.

        Parameters
        ----------
        xi: float
            the 1st Gaussian coordinate.
        eta: float
            the 2nd Gaussian coordinate.

        Returns
        -------
        x: float
            the 1st local coordinate.
        z: float
            the 2nd local coordinate.

        """
        x = self.initial_length / 2 * (xi + 1)
        z = self.height / 2 * eta

        return x, z

    def strain(self, x, z):
        """Calculate the strain at a point expressed in local coordinate.

        Parameters
        ----------
        x: float
            the 1st local coordinate.
        z: float
            the 2nd local coordinate.

        Returns
        -------
        strain: float/numpy.ndarray
            If it is Timoshenko beam, shear strain will occur so that there are two components
            in the strain.

        """
        pl = self.local_displacement()
        if self.beamtype == "Bernoulli":
            return float(pl[0] / self.initial_length + z * ((4 / self.initial_length - 6 *
                                                            x / (self.initial_length) ** 2) * pl[1] +
                                                            (2 / self.initial_length - 6 *
                                                            x / (self.initial_length) ** 2) * pl[2]))
        else:
            raise ValueError("Wrong element type, please check!")

    def perform_linear_hardening(self):
        """Update state variables using return-mapping algorithm for linear hardening model.

        Parameters
        ----------
        None.

        Returns
        -------
        None.

        """
        def yield_stress_function(
            alpha): return self.yield_stress + self.plastic_modulus * alpha

        for n in range(self.num_of_gauss_locations):
            x, z = self.global_gauss_locations_mat[0: 2, n]

            stress_trial = self.youngs_modulus * \
                (self.strain(x, z) - self.plastic_strain[n])
            relative_stress_trial = stress_trial - \
                self.back_stress[n]
            yield_condition_trial = np.abs(
                relative_stress_trial) - yield_stress_function(self.internal_hardening_variable[n])

            if yield_condition_trial <= 0:
                self.stress[n] = stress_trial
                self.tangent_modulus[n] = self.youngs_modulus
            else:
                deltagamma = yield_condition_trial / \
                    (self.youngs_modulus + self.plastic_modulus +
                        self.kinematic_hardening_modulus)
                self.stress[n] = stress_trial - deltagamma * \
                    self.youngs_modulus * np.sign(relative_stress_trial)
                self.plastic_strain[n] += deltagamma * \
                    np.sign(relative_stress_trial)
                self.back_stress[n] += deltagamma * \
                    self.kinematic_hardening_modulus * \
                    np.sign(relative_stress_trial)
                self.internal_hardening_variable[n] += deltagamma
                self.tangent_modulus[n] = self.youngs_modulus * (self.plastic_modulus + self.kinematic_hardening_modulus) / (
                    self.youngs_modulus + self.plastic_modulus + self.kinematic_hardening_modulus)

    def perform_exponential_hardening(self):
        """Update state variables using return-mapping algorithm for exponential hardening model.

        Parameters
        ----------
        None.

        Returns
        -------
        None.

        """

        def yield_stress_function(alpha): return self.yield_stress + (
            self.saturation_stress - self.yield_stress) * (1 - np.exp(-self.exponent * alpha))

        hardening_limit_stress = self.saturation_stress - self.yield_stress

        for n in range(self.num_of_gauss_locations):
            x, z = self.global_gauss_locations_mat[0: 2, n]

            stress_trial = self.youngs_modulus * \
                (float(self.strain(x, z)) -
                    self.plastic_strain[n])
            relative_stress_trial = stress_trial - \
                self.back_stress[n]
            yield_condition_trial = np.abs(
                relative_stress_trial) - yield_stress_function(self.internal_hardening_variable[n])

            if yield_condition_trial <= 0:
                self.stress[n] = stress_trial
                self.tangent_modulus[n] = self.youngs_modulus
            else:
                deltagamma = 0
                residual = yield_condition_trial - deltagamma * (self.youngs_modulus + self.kinematic_hardening_modulus) - yield_stress_function(
                    self.internal_hardening_variable[n] + deltagamma) + yield_stress_function(self.internal_hardening_variable[n])

                maxiter = 10
                iteration_counter = 0
                tolerance = 1.0e-5

                while residual > tolerance and iteration_counter < maxiter:
                    dR_ddeltagamma = -(self.youngs_modulus + self.kinematic_hardening_modulus) - hardening_limit_stress * self.exponent * \
                        np.exp(-self.exponent *
                               (self.internal_hardening_variable[n] + deltagamma))
                    d_g = - residual / dR_ddeltagamma
                    deltagamma += d_g
                    residual = yield_condition_trial - deltagamma * (self.youngs_modulus + self.kinematic_hardening_modulus) - yield_stress_function(
                        self.internal_hardening_variable[n] + deltagamma) + yield_stress_function(self.internal_hardening_variable[n])
                    iteration_counter += 1

                self.stress[n] = stress_trial - deltagamma * \
                    self.youngs_modulus * np.sign(relative_stress_trial)
                self.plastic_strain[n] += deltagamma * \
                    np.sign(relative_stress_trial)
                self.back_stress[n] += deltagamma * \
                    self.kinematic_hardening_modulus * \
                    np.sign(relative_stress_trial)
                self.internal_hardening_variable[n] += deltagamma
                self.tangent_modulus[n] = self.youngs_modulus * (self.kinematic_hardening_modulus + hardening_limit_stress * self.exponent * np.exp(-self.exponent * self.internal_hardening_variable[n])) / (
                    self.youngs_modulus + self.kinematic_hardening_modulus + hardening_limit_stress * self.exponent * np.exp(-self.exponent * self.internal_hardening_variable[n]))

    def perform_ramberg_osgood_hardening(self):
        """Update state variables using return-mapping algorithm for Ramberg-Osgood hardening model.

        Parameters
        ----------
        None.

        Returns
        -------
        None.

        """
        def yield_stress_function(alpha): return self.yield_stress + \
            self.modified_modulus * np.power(alpha, self.exponent)

        for n in range(self.num_of_gauss_locations):
            x, z = self.global_gauss_locations_mat[0: 2, n]
            stress_trial = self.youngs_modulus * \
                (float(self.strain(x, z)) -
                    self.plastic_strain[n])
            relative_stress_trial = stress_trial - \
                self.back_stress[n]
            yield_condition_trial = np.abs(
                relative_stress_trial) - yield_stress_function(self.internal_hardening_variable[n])

            if yield_condition_trial <= 0:
                self.stress[n] = stress_trial
                self.tangent_modulus[n] = self.youngs_modulus
            else:
                deltagamma = 0
                residual = yield_condition_trial - deltagamma * (self.youngs_modulus + self.kinematic_hardening_modulus) - yield_stress_function(
                    self.internal_hardening_variable[n] + deltagamma) + yield_stress_function(self.internal_hardening_variable[n])

                maxiter = 10
                iteration_counter = 0
                tolerance = 1.0e-5

                while residual > tolerance and iteration_counter < maxiter:
                    dR_ddeltagamma = -(self.youngs_modulus + self.kinematic_hardening_modulus) - self.exponent * self.modified_modulus * \
                        np.power(
                            self.internal_hardening_variable[n] + deltagamma, self.exponent - 1)
                    d_g = - residual / dR_ddeltagamma
                    deltagamma += d_g
                    residual = yield_condition_trial - deltagamma * (self.youngs_modulus + self.kinematic_hardening_modulus) - yield_stress_function(
                        self.internal_hardening_variable[n] + deltagamma) + yield_stress_function(self.internal_hardening_variable[n])
                    iteration_counter += 1

                self.stress[n] = stress_trial - deltagamma * \
                    self.youngs_modulus * np.sign(relative_stress_trial)
                self.plastic_strain[n] += deltagamma * \
                    np.sign(relative_stress_trial)
                self.back_stress[n] += deltagamma * \
                    self.kinematic_hardening_modulus * \
                    np.sign(relative_stress_trial)
                self.internal_hardening_variable[n] += deltagamma
                self.tangent_modulus[n] = self.youngs_modulus * (self.kinematic_hardening_modulus + self.exponent * self.modified_modulus * np.power(self.internal_hardening_variable[n], self.exponent - 1)) / (
                    self.youngs_modulus + self.kinematic_hardening_modulus + self.exponent * self.modified_modulus * np.power(self.internal_hardening_variable[n], self.exponent - 1))

    def global_stiffness_force(self):
        """Calculate the global stiffness matrix kg and global force vector fg of the beam element.

        Parameters
        ----------
        None.

        Returns
        -------
        kg: np.ndarray
            global stiffness matrix.
        fg: np.ndarray
            global force vector

        """
        cosbeta, sinbeta = self.current_angle_trigonometric("global")
        l = self.current_length()

        B = np.array([[-cosbeta, -sinbeta, 0., cosbeta, sinbeta, 0.],
                      [-sinbeta / l, cosbeta/l, 1., sinbeta/l, -cosbeta/l, 0.],
                      [-sinbeta / l, cosbeta/l, 0., sinbeta/l, -cosbeta/l, 1.]])
        r = np.array([[-cosbeta, -sinbeta, 0.,  cosbeta, sinbeta, 0.]]).T
        z = np.array([[sinbeta, -cosbeta, 0., -sinbeta, cosbeta, 0.]]).T

        K_l, f_l = self.local_stiffness_force()

        K_m = f_l[0] / l * z @ z.T
        + (f_l[1] + f_l[2]) / (l ** 2) * (r @ z.T + z @ r.T)

        K_g = B.T @ K_l @ B + K_m
        f_g = B.T @ f_l

        return K_g, f_g

max_stars_repo_name: Xinrui/CorotationalBeamSolver

max_issues_repo_name: Xinrui/CorotationalBeamSolver

max_forks_repo_name: Xinrui/CorotationalBeamSolver

reason: import numpy,from numpy


text:
"""
When calling Main.py, there are imports that lead tensorflow to be initialized. This file must always be completely
imported before utilizing other parts of our code pipeline as tensorflow will automatically allocate all available
VRAM. This will prevent new runs from accessing the GPUs and may result in a crash.

Additionally, this file contains a variety of monitoring flags and functions for tracking progress of agents.
"""
from abc import ABC, abstractmethod
import logging
import typing
import os

# Dynamic VRAM growth: https://github.com/tensorflow/tensorflow/issues/24496#issuecomment-464909727
from tensorflow.compat.v1 import ConfigProto
from tensorflow.compat.v1 import InteractiveSession

import tensorflow as tf
import numpy as np

from utils.loss_utils import support_to_scalar, safe_l2norm

# Dynamic VRAM growth: https://github.com/tensorflow/tensorflow/issues/24496#issuecomment-464909727
config = ConfigProto()
config.gpu_options.allow_growth = True
session = InteractiveSession(config=config)  # TODO: Set main session on the CPU. Activate Pipeline with GPU session.
os.environ['TF_FORCE_GPU_ALLOW_GROWTH'] = 'true'

# Suppress warnings from TENSORFLOW's side
logger = tf.get_logger()
logger.setLevel(logging.ERROR)
tf.autograph.set_verbosity(1)

# Debugging flags. SET AS GLOBAL VARIABLES BEFORE RUNNING COACH.
DEBUG_MODE = False
LOG_RATE = 1
RENDER = False

class Monitor(ABC):
    """ Main interface for monitoring progress of agent """

    def __init__(self, reference):
        self.reference = reference  # Instance of a Neural Network framework to track statistics on.

    def log(self, tensor: typing.Union[tf.Tensor, float], name: str) -> None:
        """ Log a scalar annotated by the number of backpropagation steps """
        if self.reference.steps % LOG_RATE == 0:
            tf.summary.scalar(name, data=tensor, step=self.reference.steps)

    def log_distribution(self, tensor: typing.Union[tf.Tensor, np.ndarray], name: str) -> None:
        """ Log an array of scalars annotated by the number of backpropagation steps """
        if self.reference.steps % LOG_RATE == 0:
            tf.summary.histogram(name, tensor, step=self.reference.steps)

    @abstractmethod
    def log_batch(self, data_batch: typing.List) -> None:
        """ Ad-Hoc function to log specific statistic based on a batch of training-data. """

class MuZeroMonitor(Monitor):
    """
    Implementation of the Monitor class to add batch logging functionality for MuZero agents.
    Also adds functionality to log the losses of each recurrent head during unrolling/ training of the RNN.
    """

    def __init__(self, reference):
        super().__init__(reference)

    def log_recurrent_losses(self, t: int, v_loss: tf.Tensor, r_loss: tf.Tensor, pi_loss: tf.Tensor,
                             absorb: tf.Tensor, o_loss: tf.Tensor = None) -> None:
        """ Log each prediction head loss from the MuZero RNN as a scalar (optionally includes decoder loss). """
        step = self.reference.steps
        if self.reference.steps % LOG_RATE == 0:
            tf.summary.scalar(f"r_loss_{t}", data=tf.reduce_mean(r_loss), step=step)
            tf.summary.scalar(f"v_loss_{t}", data=tf.reduce_mean(v_loss), step=step)
            tf.summary.scalar(f"pi_loss_{t}", data=tf.reduce_sum(pi_loss) / tf.reduce_sum(1 - absorb), step=step)

            if o_loss is not None:  # Decoder option.
                tf.summary.scalar(f"decode_loss_{t}", data=tf.reduce_mean(o_loss), step=step)

    def log_batch(self, data_batch: typing.List) -> None:
        """
        Log a large amount of neural network statistics based on the given batch.
        Functionality can be toggled on by specifying '--debug' as a console argument to Main.py.
        Note: toggling this functionality on will produce significantly larger tensorboard event files!

        Statistics include:
         - Priority sampling sample probabilities.
         - Loss of each recurrent head per sample as a distribution.
         - Loss discrepancy between cross-entropy and MSE for the reward/ value predictions.
         - Norm of the neural network's weights.
         - Divergence between the dynamics and encoder functions.
         - Squared error of the decoding function.
        """
        if DEBUG_MODE and self.reference.steps % LOG_RATE == 0:
            observations, actions, targets, forward_observations, sample_weight = list(zip(*data_batch))
            actions, sample_weight = np.asarray(actions), np.asarray(sample_weight)
            target_vs, target_rs, target_pis = list(map(np.asarray, zip(*targets)))

            priority = sample_weight * len(data_batch)  # Undo 1/n scaling to get priority
            tf.summary.histogram(f"sample probability", data=priority, step=self.reference.steps)

            s, pi, v = self.reference.neural_net.forward.predict_on_batch(np.asarray(observations))

            v_real = support_to_scalar(v, self.reference.net_args.support_size).ravel()

            tf.summary.histogram(f"v_predict_{0}", data=v_real, step=self.reference.steps)
            tf.summary.histogram(f"v_target_{0}", data=target_vs[:, 0], step=self.reference.steps)
            tf.summary.scalar(f"v_mse_{0}", data=np.mean((v_real - target_vs[:, 0]) ** 2), step=self.reference.steps)

            # Sum over one-hot-encoded actions. If this sum is zero, then there is no action --> leaf node.
            absorb_k = 1.0 - tf.reduce_sum(target_pis, axis=-1)

            collect = list()
            for k in range(actions.shape[1]):
                r, s, pi, v = self.reference.neural_net.recurrent.predict_on_batch([s, actions[:, k, :]])

                collect.append((s, v, r, pi, absorb_k[k, :]))

            for t, (s, v, r, pi, absorb) in enumerate(collect):
                k = t + 1

                pi_loss = -np.sum(target_pis[:, k] * np.log(pi + 1e-8), axis=-1)
                self.log_distribution(pi_loss, f"pi_dist_{k}")

                v_real = support_to_scalar(v, self.reference.net_args.support_size).ravel()
                r_real = support_to_scalar(r, self.reference.net_args.support_size).ravel()

                self.log_distribution(r_real, f"r_predict_{k}")
                self.log_distribution(v_real, f"v_predict_{k}")

                self.log_distribution(target_rs[:, k], f"r_target_{k}")
                self.log_distribution(target_vs[:, k], f"v_target_{k}")

                self.log(np.mean((r_real - target_rs[:, k]) ** 2), f"r_mse_{k}")
                self.log(np.mean((v_real - target_vs[:, k]) ** 2), f"v_mse_{k}")

            l2_norm = tf.reduce_sum([safe_l2norm(x) for x in self.reference.get_variables()])
            self.log(l2_norm, "l2 norm")

            # Option to track statistical properties of the dynamics model.
            if self.reference.net_args.dynamics_penalty > 0:
                forward_observations = np.asarray(forward_observations)
                # Compute statistics related to auto-encoding state dynamics:
                for t, (s, v, r, pi, absorb) in enumerate(collect):
                    k = t + 1
                    stacked_obs = forward_observations[:, t, ...]

                    s_enc = self.reference.neural_net.encoder.predict_on_batch(stacked_obs)
                    kl_divergence = tf.keras.losses.kullback_leibler_divergence(s_enc, s)

                    # Relative entropy of dynamics model and encoder.
                    # Lower values indicate that the prediction model receives more stable input.
                    self.log_distribution(kl_divergence, f"KL_Divergence_{k}")
                    self.log(np.mean(kl_divergence), f"Mean_KLDivergence_{k}")

                    # Internal entropy of the dynamics model
                    s_entropy = tf.keras.losses.categorical_crossentropy(s, s)
                    self.log(np.mean(s_entropy), f"mean_dynamics_entropy_{k}")

                    if hasattr(self.reference.neural_net, "decoder"):
                        # If available, track the performance of a neural decoder from latent to real state.
                        stacked_obs_predict = self.reference.neural_net.decoder.predict_on_batch(s)
                        se = (stacked_obs - stacked_obs_predict) ** 2

                        self.log(np.mean(se), f"decoder_error_{k}")

class AlphaZeroMonitor(Monitor):
    """ Implementation of the Monitor class to add batch logging functionality for AlphaZero agents. """

    def __init__(self, reference):
        super().__init__(reference)

    def log_batch(self, data_batch: typing.List) -> None:
        """
        Log a large amount of neural network statistics based on the given batch.
        Functionality can be toggled on by specifying '--debug' as a console argument to Main.py.
        Note: toggling this functionality on will produce significantly larger tensorboard event files!

        Statistics include:
         - Priority sampling sample probabilities.
         - Values of each target/ prediction for the data batch.
         - Loss discrepancy between cross-entropy and MSE for the reward/ value predictions.
        """
        if DEBUG_MODE and self.reference.steps % LOG_RATE == 0:
            observations, targets, sample_weight = list(zip(*data_batch))
            target_pis, target_vs = list(map(np.asarray, zip(*targets)))
            observations = np.asarray(observations)

            priority = sample_weight * len(data_batch)  # Undo 1/n scaling to get priority
            tf.summary.histogram(f"sample probability", data=priority, step=self.reference.steps)

            pis, vs = self.reference.neural_net.model.predict_on_batch(observations)
            v_reals = support_to_scalar(vs, self.reference.net_args.support_size).ravel()  # as scalars

            tf.summary.histogram(f"v_targets", data=target_vs, step=self.reference.steps)
            tf.summary.histogram(f"v_predict", data=v_reals, step=self.reference.steps)

            mse = np.mean((v_reals - target_vs) ** 2)
            tf.summary.scalar("v_mse", data=mse, step=self.reference.steps)

max_stars_repo_name: morozig/muzero

max_issues_repo_name: frankbryce/muzero

max_forks_repo_name: frankbryce/muzero


text:
import os
from glob import glob
from shutil import rmtree
from tempfile import mkdtemp

import numpy as np
import pytest

from graph.dim import Conv2DFilterDim, Dim, FcFilterDim, PadDim, StrideDim
from graph.matches.matches import get_std_match_group
from graph.nngraph import NNGraph
from graph.types import Conv2DParameters, FcParameters, NNEdge
from importer.importer import create_graph
from quantization.simple_auto_quantify import SimpleQuantizer
from stats.activation_stats_collector import ActivationStatsCollector
from stats.filter_stats_collector import FilterStatsCollector
from utils.data_importer import import_data
from utils.intermediate_cache import IntermediateCache
from utils.new_param_state import dump_state
from utils.sparse_list import SparseList

MNIST_GRAPH = 'tests/graph/mnist_model.tflite'
IR_GRAPH = 'tests/graph/ir_model.tflite'
SSD_GRAPH = 'tests/graph/ssd_model.tflite'
CIFAR10_GRAPH = 'tests/graph/cifar10_model.tflite'
KWS_GRAPH = 'tests/graph/kws.tflite'
VISUAL_GRAPH = 'tests/graph/visual_wake.tflite'
CONCAT_TEST_GRAPH = 'tests/graph/concat_test.tflite'

MNIST_IMAGES = glob('tests/images/*.pgm')
VWW_IMAGES = glob('tests/vwwimages/*.png')

def save_state(temp_dir, width, fusions=False, adjust=False):
    file_name = os.path.join(temp_dir, "state_file")
    G = create_graph(MNIST_GRAPH, opts={"load_tensors":True})
    G.add_dimensions()
    if adjust:
        G.adjust_order()
    if fusions:
        get_std_match_group().match(G)
        G.add_dimensions()
    stats_collector = ActivationStatsCollector()
    for input_file in MNIST_IMAGES:
        data = import_data(input_file, offset=0, divisor=255)
        if not adjust:
            data = data.reshape((28, 28, 1))
        stats_collector.collect_stats(G, [data])
    astats = stats_collector.reduce_stats()
    stats_collector = FilterStatsCollector()
    fstats = stats_collector.collect_stats(G)
    quantizer = SimpleQuantizer(astats, fstats, force_width=width)
    qrecs = quantizer.quantize(G)
    G.quantization = qrecs
    dump_state(G, include_parameters=True, state_path=file_name)
    return file_name

@pytest.fixture()
def two_conv_graph():
    G = NNGraph(name='two_conv_graph')
    ti = G.add_input(Dim.unnamed([10, 10, 2]))
    c1filt = Conv2DFilterDim(3, 3, 2, in_c=2)
    c1filt.impose_order(['out_c', 'h', 'w', 'in_c'])
    n1 = Conv2DParameters("node1",
                          filt=c1filt,
                          stride=StrideDim(1, 1),
                          padding=PadDim(0),
                          in_dims_hint=SparseList([['h', 'w', 'c']]),
                          out_dims_hint=SparseList([['h', 'w', 'c']]))
    G.add_node(n1)
    w1 = [[0.25, 0.25], [0.25, 0.25], [0.25, 0.25]]
    w1 = [w1, w1, w1]
    w2 = [[0.75, 0.75], [0.75, 0.75], [0.75, 0.75]]
    w2 = [w2, w2, w2]
    n1.weights = np.array([w1, w2])
    c2filt = Conv2DFilterDim(3, 3, 2, in_c=2)
    c2filt.impose_order(['out_c', 'h', 'w', 'in_c'])
    n2 = Conv2DParameters("node2",
                          filt=c2filt,
                          stride=StrideDim(1, 1),
                          padding=PadDim(0),
                          in_dims_hint=SparseList([['h', 'w', 'c']]),
                          out_dims_hint=SparseList([['h', 'w', 'c']]))
    G.add_node(n2)
    w3 = [[0.75, 0.25], [0.75, 0.25], [0.75, 0.25]]
    w3 = [w3, w3, w3]
    n2.weights = np.array([w3, w3])
    to = G.add_output()
    G.add_edge(NNEdge(ti, n1))
    G.add_edge(NNEdge(n1, n2))
    G.add_edge(NNEdge(n2, to))
    G.add_dimensions()
    yield G

@pytest.fixture(scope="session")
def mnist_unfused_16bit_state():
    temp_dir = mkdtemp()
    yield save_state(temp_dir, 16)
    rmtree(temp_dir)

@pytest.fixture(scope="session")
def mnist_unfused_8bit_state():
    temp_dir = mkdtemp()
    yield save_state(temp_dir, 8)
    rmtree(temp_dir)

@pytest.fixture(scope="session")
def mnist_fused_8bit_state():
    temp_dir = mkdtemp()
    yield save_state(temp_dir, 8, True, True)
    rmtree(temp_dir)

@pytest.fixture(scope="session")
def value_cache():
    temp_dir = mkdtemp()
    yield IntermediateCache(temp_dir)
    rmtree(temp_dir)

@pytest.fixture(scope="session")
def vww_graph():
    yield VISUAL_GRAPH

@pytest.fixture(scope="session")
def mnist_graph():
    yield MNIST_GRAPH

@pytest.fixture(scope="session")
def ir_graph():
    yield IR_GRAPH

@pytest.fixture(scope="session")
def ssd_graph():
    yield SSD_GRAPH

@pytest.fixture(scope="session")
def cifar10_graph():
    yield CIFAR10_GRAPH

@pytest.fixture(scope="session")
def kws_graph():
    yield KWS_GRAPH

@pytest.fixture(scope="session")
def concat_test_graph():
    yield CONCAT_TEST_GRAPH

@pytest.fixture(scope="session")
def ir_images():
    yield [
        "tests/images/im1.png",
        "tests/images/im2.png",
        "tests/images/im3.png",
        "tests/images/im4.png",
    ]

@pytest.fixture(scope="session")
def kws_sounds():
    yield [
        "tests/sounds/feature.txt_475.pgm",
        "tests/sounds/feature.txt_737.pgm",
        "tests/sounds/feature.txt_99.pgm",
        "tests/sounds/feature.txt_476.pgm",
        "tests/sounds/feature.txt_738.pgm",
        "tests/sounds/feature.txt_9.pgm"
    ]

@pytest.fixture(scope="session")
def mnist_images():
    yield MNIST_IMAGES

@pytest.fixture(scope="session")
def vww_images():
    yield VWW_IMAGES


text:
#   Copyright 2020 The PyMC Developers
#
#   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.

import aesara
import numpy as np
import pytest

from pymc.variational.updates import (
    adadelta,
    adagrad,
    adagrad_window,
    adam,
    adamax,
    momentum,
    nesterov_momentum,
    rmsprop,
    sgd,
)

_a = aesara.shared(1.0)
_b = _a * 2

_m = aesara.shared(np.empty((10,), aesara.config.floatX))
_n = _m.sum()
_m2 = aesara.shared(np.empty((10, 10, 10), aesara.config.floatX))
_n2 = _b + _n + _m2.sum()

@pytest.mark.parametrize(
    "opt",
    [sgd, momentum, nesterov_momentum, adagrad, rmsprop, adadelta, adam, adamax, adagrad_window],
    ids=[
        "sgd",
        "momentum",
        "nesterov_momentum",
        "adagrad",
        "rmsprop",
        "adadelta",
        "adam",
        "adamax",
        "adagrad_window",
    ],
)
@pytest.mark.parametrize(
    "getter",
    [
        lambda t: t,  # all params -> ok
        lambda t: (None, t[1]),  # missing loss -> fail
        lambda t: (t[0], None),  # missing params -> fail
        lambda t: (None, None),
    ],  # all missing -> partial
    ids=["all_params", "missing_loss", "missing_params", "all_missing"],
)
@pytest.mark.parametrize(
    "kwargs", [dict(), dict(learning_rate=1e-2)], ids=["without_args", "with_args"]
)
@pytest.mark.parametrize(
    "loss_and_params",
    [(_b, [_a]), (_n, [_m]), (_n2, [_a, _m, _m2])],
    ids=["scalar", "matrix", "mixed"],
)
def test_updates_fast(opt, loss_and_params, kwargs, getter):
    with aesara.config.change_flags(compute_test_value="ignore"):
        loss, param = getter(loss_and_params)
        args = dict()
        args.update(**kwargs)
        args.update(dict(loss_or_grads=loss, params=param))
        if loss is None and param is None:
            updates = opt(**args)
            # Here we should get new callable
            assert callable(updates)
            # And be able to get updates
            updates = opt(_b, [_a])
            assert isinstance(updates, dict)
        # case when both are None is above
        elif loss is None or param is None:
            # Here something goes wrong and user provides not full set of [params + loss_or_grads]
            # We raise Value error
            with pytest.raises(ValueError):
                opt(**args)
        else:
            # Usual call to optimizer, old behaviour
            updates = opt(**args)
            assert isinstance(updates, dict)


text:
from fastapi import APIRouter, Query, HTTPException
import json
import numpy as np
import pandas as pd
from typing import List
import functools
from txmlcommonlibrary.utilities.jsonencoder import NumpyEncoder
from fastapi.encoders import jsonable_encoder

from ..params import models_dict, configs, Params
from ..log import logger

router = APIRouter()

def call_function(model_name, obj, data, call_name, call_params):

    assert call_name in ["transform", "predict", "predict_proba",
                         "mean", "predict_alert", "predict_trip",
                         "column_list"]

    # assert call_params
    if call_name == "predict" or call_name == "predict_proba":
        return predict_with_features(model_name, data, obj, call_name, **call_params)
    elif call_name == "mean":
        return np.mean(data)
    elif call_name == "column_list":
        return obj.column_list

def predict_with_features(model_name, features, model, predict_call, call_params=None, **other_params):
    """
    given a dict of features turn it into a df and use it to predict.
    This assumes the column order is already correct from the scaling.
    This method is relevant in most traditional techniques,
    using a flat featrue based input.
    """
    try:
        df = pd.DataFrame(features)
    except ValueError:
        df = pd.DataFrame(features, index=[0])
    df = df.apply(functools.partial(pd.to_numeric, errors='coerce'))

    # Checks columns match
    # emergency cleaning
    data = df

    if call_params:
        assert isinstance(call_params, dict)
        _call_params = call_params
    else:
        _call_params = {}

    if predict_call == "predict":
        results = model.predict(data, **_call_params)
    elif predict_call == "predict_proba":
        results = model.predict_proba(data, **_call_params).flatten()
    else:
        raise Exception("Model is to be called with predict ot predict_proba")

    if type(results) == dict:
        results = json.loads(json.dumps(results, cls=NumpyEncoder))
        results['model_name'] = configs.get('MODELS', {}).get(model_name, {}).get('INFO', {}).get('MODEL_NAME')
        results['version'] = configs.get('MODELS', {}).get(model_name, {}).get('INFO', {}).get('VERSION')
        return results
    else:
        res = {}
        res['model_name'] = configs.get('MODELS', {}).get(model_name, {}).get('INFO', {}).get('MODEL_NAME')
        res['version'] = configs.get('MODELS', {}).get(model_name, {}).get('INFO', {}).get('VERSION')
        if isinstance(results, list):
            res['data'] = results
            return res
        elif isinstance(results, np.ndarray):
            res['data'] = results.tolist()
            return res
        else:
            raise ValueError("Unexpected result type. Model results must be dict, list or numpy arrays")

@router.post("/model/{model_name}/{method}")
async def predict_single(model_name: str,
                         method: str,
                         data: dict):
    """
    For a model name, return the prediction given the data.
    The method can be found using an endpoint or prior user knowledge.
    """
    obj = models_dict.get(model_name, None)
    calls = configs.get("MODELS", {}).get(model_name, {}).get("CALLS", {}).get("CALL_PARAMS", {})
    logger.info(f"{obj}")
    # TODO All excpetions here ar e500 internals
    if method in calls.keys():
        logger.info(f"{calls}")
        call_params = calls[method]
    else:
        raise HTTPException(status_code=501, detail=f"{method} not a registered call")
    if obj:
        try:
            results = call_function(model_name, obj, data, method, call_params)
            if results is None:
                raise HTTPException(status_code=501, detail=f"null results")
            return results
        except Exception as err:
            logger.exception(f"{err}")
            raise HTTPException(status_code=500, detail=f"{err}")
    else:
        raise HTTPException(status_code=501, detail=f"Obj is Null")

# @router.post("/model/{model_name}/{method}/object")
# async def predict_single_from_object(model_name: str,
#                                      method: str,
#                                      object: str):
#     pass

# @router.post("/model/{model_name}/{method}/{id}")
# async def predict_single_from_id(model_name: str,
#                                  method: str,
#                                  id: str):
#     pass

max_stars_repo_name: Tim-ty-tang/mlflow-fastapi-deploy

max_issues_repo_name: Tim-ty-tang/mlflow-fastapi-deploy

max_forks_repo_name: Tim-ty-tang/mlflow-fastapi-deploy