text:
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon Jul 27 20:00:38 2020

@author: lukepinkel
"""
import tqdm
import arviz as az
import numpy as np
import scipy as sp
import scipy.stats
import matplotlib.pyplot as plt
from .lmm import LMM, make_theta
from ..utilities.data_utils import _check_shape, sign_change
from ..utilities.linalg_operations import vech
from sksparse.cholmod import cholesky
from ..utilities.random import trnorm, r_invwishart, r_invgamma
from ..utilities.func_utils import poisson_logp, log1p, norm_cdf

def get_u_indices(dims):
    u_indices = {}
    start=0
    for key, val in dims.items():
        q = val['n_groups']*val['n_vars']
        u_indices[key] = np.arange(start, start+q)
        start+=q
    return u_indices

def wishart_info(dims):
    ws = {}
    for key in dims.keys():
        ws[key] = {}
        q = dims[key]['n_groups']
        k = dims[key]['n_vars']
        nu = q-(k+1)
        ws[key]['q'] = q
        ws[key]['k'] = k
        ws[key]['nu'] = nu
    return ws

def sample_gcov(theta, u, wsinfo, indices, key, priors):
    u_i = u[indices['u'][key]]
    U_i = u_i.reshape(-1, wsinfo[key]['k'], order='C')
    Sg_i =  U_i.T.dot(U_i)
    Gs = r_invwishart(wsinfo[key]['nu']+priors[key]['n'],
                             Sg_i+priors[key]['V'])
    theta[indices['theta'][key]] = vech(Gs)
    return theta

def sample_rcov(theta, y, yhat, wsinfo, priors):
    resid = y - yhat
    sse = resid.T.dot(resid)
    nu = wsinfo['r']['nu']
    ss = r_invgamma((nu+priors['R']['n'])/2, scale=(sse+priors['R']['V'])/2)
    theta[-1] = ss
    return theta

def to_arviz_dict(samples, var_dict, burnin=0):
    az_dict = {}
    for key, val in var_dict.items():
        az_dict[key] = samples[:, burnin:, val]
    return az_dict

def _plot_trace(idata, plot_kws=None):
    default_plot_kws = dict(combined=False, compact=False, figsize=(12, 8))
    plot_kws = {} if plot_kws is None else plot_kws
    plot_kws = {**default_plot_kws, **plot_kws}
    axes = az.plot_trace(idata, **plot_kws)
    n = axes.shape[0]
    for i in range(n):
        label = axes[i, 0].get_title().replace("\n", "[")+"]"
        axes[i, 0].set_title("")
        axes[i, 1].set_title("")
        axes[i, 0].set_ylabel(label, rotation=0, labelpad=15)
    plt.subplots_adjust(left=0.04, right=0.97, top=0.97, bottom=0.04, wspace=0.1)
    return axes

def _plot_ess(idata, plot_kws=None):
    default_plot_kws = dict(figsize=(16, 8))
    plot_kws = {} if plot_kws is None else plot_kws
    plot_kws = {**default_plot_kws, **plot_kws}

    axes = az.plot_ess(idata, **plot_kws)
    if axes.ndim==1:
        axes = axes.reshape(1, -1)
    n, m = axes.shape
    for i in range(n):
        for j in range(m):
            label = axes[i, j].get_title().replace("\n", "[")+"]"
            axes[i, j].set_title(label)
            axes[i, j].xaxis.set_tick_params(labelsize=8)
            axes[i, j].yaxis.set_tick_params(labelsize=8)
            if j>0:
                axes[i, j].set_ylabel("")
            else:
                label = axes[i, j].get_ylabel()
                axes[i, j].set_ylabel(label, size=12)

            if (i+1)<n:
                axes[i, j].set_xlabel("")
            else:
                label = axes[i, j].get_xlabel()
                axes[i, j].set_xlabel(label, size=12)

    plt.subplots_adjust(left=0.08, right=0.97, top=0.97, bottom=0.1, wspace=0.15,
                        hspace=0.15)
    return axes

def _plot_rank(idata, plot_kws=None):
    default_plot_kws = dict(figsize=(12, 10))
    plot_kws = {} if plot_kws is None else plot_kws
    plot_kws = {**default_plot_kws, **plot_kws}

    axes = az.plot_rank(idata, **plot_kws)
    if axes.ndim==1:
        axes = axes.reshape(1, -1)
    n, m = axes.shape
    for i in range(n):
        for j in range(m):
            label = axes[i, j].get_title().replace("\n", "[")+"]"
            axes[i, j].set_title(label)
            axes[i, j].xaxis.set_tick_params(labelsize=8)
            axes[i, j].yaxis.set_tick_params(labelsize=8)
            if j>0:
                axes[i, j].set_ylabel("")
            else:
                label = axes[i, j].get_ylabel()
                axes[i, j].set_ylabel(label, size=12)

            if (i+1)<n:
                axes[i, j].set_xlabel("")
            else:
                label = axes[i, j].get_xlabel()
                axes[i, j].set_xlabel(label, size=12)

    plt.subplots_adjust(left=0.08, right=0.97, top=0.97, bottom=0.08, wspace=0.15,
                        hspace=0.15)
    return axes

class MixedMCMC(LMM):

    def __init__(self, formula, data, response_dist, priors=None, weights=None,
                 rng=None, vnames=None, freeR=None):
        """

        Parameters
        ----------
        formula : string
            lme4 style formula specifying random effects with parentheses and
            a vertical bar.
        data : dataframe
            Dataframe containing data.  Missing values should be dropped
            manually before passing the dataframe.
        response_dist : string
            String specifying the type of model. One of `ordinal_probit`,
            `bernoulli`, `binomial`, or `normal`.
        priors : dict, optional
            A dictionary with keys consisting of random effect factors or
            `R` for the residual covariance. The corresponding value should be
            a dictionary with keys `V` and `n`, specifying the prior covariance
            matrix and number of observations. The default is None.
        weights : ndarray, optional
            Weight matrix. The default is None.
        rng : numpy.random._generator.Generator, optional
            Numpy random generator. The default is None.
        vnames : list, optional
            List of variable names. The default is None.
        freeR : bool, optional
            If true, allows the residual covariance to vary in `binomial` and
            `bernoulli` models. The default is None.

        Returns
        -------
        None.

        """
        super().__init__(formula, data)

        #Initialize parameters, get misc information for sampling
        self.t_init, _ = make_theta(self.dims)
        self.W = sp.sparse.csc_matrix(self.XZ)
        self.WtW = self.W.T.dot(self.W)
        self.wsinfo = wishart_info(self.dims)
        self.y = _check_shape(self.y, 1)
        self.indices['u'] = get_u_indices(self.dims)

        if weights is None:
            self.weights = np.ones_like(self.y)
        else:
            self.weights=weights

        #Get various model constants
        self.n_re = self.G.shape[0]
        self.n_fe = self.X.shape[1]
        self.n_lc = self.W.shape[1]
        self.n_ob = self.W.shape[0]
        self.n_params = len(self.t_init)+self.n_fe

        #Handle priors
        if priors is None:
            if response_dist == 'normal':
                priors = dict(R=dict(V=0, n=-2))
            else:
                priors = dict(R=dict(V=0.500*self.n_ob, n=self.n_ob))
            for level in self.levels:
                Vi = np.eye(self.dims[level]['n_vars'])*0.001
                priors[level] = dict(V=Vi, n=4)
        self.priors = priors
        self.wsinfo['r'] = dict(nu=self.n_ob-2)

        #Initialize containers requiring constants above
        self.offset = np.zeros(self.n_lc)
        self.location = np.zeros(self.n_re+self.n_fe)
        self.zero_mat = sp.sparse.eye(self.n_fe)*0.0
        self.re_mu = np.zeros(self.n_re)

        #Get model specific indices and constants
        if response_dist == 'bernoulli':
            self.ix0, self.ix1 = self.y==0, self.y==1
            self.jv0, self.jv1 = np.ones(len(self.ix0)), np.ones(len(self.ix1))
            freeR = False if freeR is None else freeR
        elif response_dist == 'ordinal_probit':
            self.jv = np.ones(self.n_ob)
            self.n_cats = np.unique(self.y).shape[0]
            self.n_thresh = self.n_cats - 1
            self.n_tau = self.n_cats + 1
            self.y_cat = np.zeros((self.n_ob, self.n_cats))
            self.y_ix = {}
            for i in range(self.n_cats):
                self.y_ix[i] = self.y==i
                self.y_cat[self.y==i, i] = 1
            freeR = False if freeR is None else freeR
        elif response_dist == 'binomial':
            freeR = False if freeR is None else freeR
        elif response_dist == 'normal':
            freeR = True if freeR is None else freeR

        self.response_dist = response_dist
        self.freeR = freeR
        self.rng = np.random.default_rng() if rng is None else rng

        #Handle variable names for arviz and variable names for results
        if vnames is None:
            vnames = {"$\\beta$":np.arange(self.n_fe),
                       "$\\theta$":np.arange(self.n_fe, self.n_params)}
            if response_dist == 'ordinal_probit':
                    vnames["$\\tau$"]=np.arange(self.n_params+1, self.n_params+self.n_thresh)
            if not self.freeR:
                vnames["$\\theta$"] = np.arange(self.n_fe, self.n_params-1)
        self.vnames = vnames

        param_names = list(self.fe_vars)
        for level in self.levels:
            for i, j in list(zip(*np.triu_indices(self.dims[level]['n_vars']))):
                param_names.append(f"{level}:G[{i}][{j}]")

        if self.freeR:
            param_names.append("resid_cov")
        if response_dist=='ordinal_probit':
            param_names = param_names + [f"t{i}" for i in range(1, self.n_thresh+1)]
        self.param_names = param_names

    def sample_location(self, theta, x1, x2, y):
        """

        Parameters
        ----------
        theta: array_like
            Vector of covariance parameters

        x1: array_like
            Sample from a standard normal distribution

        x2: array_like
            Sample from a standard normal distribution

        y: array_like
            Dependent variable

        Returns
        -------
        location: array_like
            Sample from P(beta, u|y, G, R)

        """
        s, s2 =  np.sqrt(theta[-1]), theta[-1]
        M = self.WtW.copy() / s2
        Ginv = self.update_gmat(theta, inverse=True).copy()
        Omega = sp.sparse.block_diag([self.zero_mat, Ginv])
        M+=Omega
        chol_fac = cholesky(Ginv, ordering_method='natural')
        a_star = chol_fac.solve_Lt(x1, use_LDLt_decomposition=False)
        y_z = y - (self.Zs.dot(a_star) + x2 * s)
        ofs = self.offset.copy()
        ofs[-self.n_re:] = a_star
        m_chol = cholesky(M.tocsc())
        u = sp.sparse.csc_matrix.dot(y_z, self.W) / s2
        location = ofs + m_chol.solve_A(u)
        return location

    def mh_lvar_binomial(self, pred, s, z, x_step, u_accept, propC):
        z_prop = x_step * propC + z

        mndenom1, mndenom2 = np.exp(z)+1, np.exp(z_prop)+1
        densityl1 = ((self.y*z) - np.log(mndenom1))*self.weights
        densityl2 = ((self.y*z_prop) - np.log(mndenom2))*self.weights
        densityl1 += sp.stats.norm(pred, s).logpdf(z)
        densityl2 += sp.stats.norm(pred, s).logpdf(z_prop)
        density_diff = densityl2 - densityl1
        accept = (density_diff>u_accept)
        z[accept] = z_prop[accept]
        return z, accept

    def mh_lvar_poisson(self, pred, s, z, x_step, u_accept, propC):
        z_prop = x_step * propC + z

        densityl1 = poisson_logp(self.y, np.exp(z))
        densityl2 = poisson_logp(self.y, np.exp(z_prop))
        densityl1 += sp.stats.norm(pred, s).logpdf(z)
        densityl2 += sp.stats.norm(pred, s).logpdf(z_prop)
        density_diff = densityl2 - densityl1
        accept = (density_diff>u_accept)
        z[accept] = z_prop[accept]
        return z, accept

    def mh_lvar_ordinal_probit(self, theta, t, pred, v):
        tau = np.pad(t, ((1, 1)), mode='constant', constant_values=[-1e17, 1e17])
        mu, sd = np.zeros_like(pred), np.ones_like(pred)
        s = np.sqrt(theta[-1])
        for i in range(1, self.y_cat.shape[1]+1):
            ix = self.y_ix[i-1]
            j = self.jv[ix]
            mu = pred[ix]
            sd = j*s
            lb = j*tau[i-1]
            ub = j*tau[i]
            v[ix] = trnorm(mu, sd, lb, ub)
        return v

    def sample_theta(self, theta, u, z, pred, freeR=True):
        for key in self.levels:
            theta = sample_gcov(theta.copy(), u, self.wsinfo, self.indices,
                                key, self.priors)
        if freeR:
            theta = sample_rcov(theta, z, pred, self.wsinfo, self.priors)
        return theta

    def slice_sample_lvar(self, rexpon, v, z, theta, pred):
        v[self.ix1] = z[self.ix1] - log1p(np.exp(z[self.ix1]))
        v[self.ix1]-= rexpon[self.ix1]
        v[self.ix1] = v[self.ix1] - log1p(-np.exp(v[self.ix1]))

        v[self.ix0] = -log1p(np.exp(z[self.ix0]))
        v[self.ix0]-= rexpon[self.ix0]
        v[self.ix0] = log1p(-np.exp(v[self.ix0])) - v[self.ix0]
        s = np.sqrt(theta[-1])
        z[self.ix1] = trnorm(mu=pred[self.ix1], sd=s*self.jv1,
                             lb=v[self.ix1], ub=200*self.jv1)
        z[self.ix0] = trnorm(mu=pred[self.ix0], sd=s*self.jv0,
                             lb=-200*self.jv0, ub=v[self.ix0])
        return z

    def sample_tau(self, theta, t, pred, v, propC):
        alpha_prev = np.pad(np.log(np.diff(t)), (1, 0), mode='constant', constant_values=[np.log(t[0])])
        alpha_prop = self.rng.normal(alpha_prev, propC)
        ll = 0.0
        t_prop = np.cumsum(np.exp(alpha_prop))
        s = np.sqrt(theta[-1])
        for i in range(1, self.n_thresh):
            m = pred[self.y_ix[i]]
            ll += (np.log(norm_cdf((t_prop[i]-m) / s)-norm_cdf((t_prop[i-1]-m) / s))).sum()
            ll -= (np.log(norm_cdf((t[i]-m)/s)-norm_cdf((t[i-1]-m)/s))).sum()
        m = pred[self.y_ix[i+1]]
        ll += np.sum(np.log(1.0 - norm_cdf((t_prop[i]-m)/s)))
        ll -= np.sum(np.log(1.0 - norm_cdf((t[i]-m)/s)))
        if ll>np.log(np.random.uniform(0, 1)):
            t_accept = True
            t = t_prop
        else:
            t_accept = False
            t = t
        return t, t_accept

    def sample_ordinal_probit(self, n_samples, chain=0, save_pred=False,
                              save_u=False, save_lvar=False, propC=0.04, damping=0.99,
                              adaption_rate=1.02,  target_accept=0.44, n_adapt=None):

        if n_adapt is None:
            n_adapt = np.minimum(int(n_samples/2), 1000)

        freeR = self.freeR
        param_samples = np.zeros((n_samples, self.n_params+self.n_thresh))
        t_acceptances = np.zeros((n_samples))
        rng = self.rng
        location = self.location.copy()
        pred =  self.W.dot(location)
        theta = self.t_init.copy()
        t = sp.stats.norm(0, 1).ppf((self.y_cat.sum(axis=0).cumsum()/np.sum(self.y_cat))[:-1])
        t = t - t[0]
        t = t+1.0
        z = np.zeros_like(self.y).astype(float)
        wtrace, waccept = 1.0, 1.0
        pbar = tqdm.tqdm(range(n_samples), smoothing=0.01)
        for i in range(n_samples):
            t, t_accept = self.sample_tau(theta, t, pred, z, propC)
            wtrace = wtrace * damping + 1.0
            waccept *= damping
            if t_accept:
                waccept +=1
            if i<n_adapt:
                propC = propC * np.sqrt(adaption_rate**((waccept/wtrace)-target_accept))
            if t_accept:
                z = self.mh_lvar_ordinal_probit(theta, t, pred, z)
            location = self.sample_location(theta, rng.normal(0, 1, size=self.n_re),
                                            rng.normal(0, 1, size=self.n_ob), z)
            pred = self.W.dot(location)
            u = location[-self.n_re:]
            theta  = self.sample_theta(theta, u, z, pred, freeR)
            param_samples[i, self.n_fe:-self.n_thresh] = theta.copy()
            param_samples[i, :self.n_fe] = location[:self.n_fe]
            param_samples[i, -self.n_thresh:] = t
            t_acceptances[i] = t_accept
            self._secondary_samples(chain, i, pred, u, z)
            if i>1:
                pbar.set_description(f"Chain {chain+1} Tau Acceptance Prob: {t_acceptances[:i].mean():.3f} C: {propC:.5f}")
            pbar.update(1)
        pbar.close()
        return param_samples

    def sample_binomial(self, n_samples, chain=0, save_pred=False, save_u=False,
                        save_lvar=False, propC=1.0, damping=0.99,
                        adaption_rate=1.01,  target_accept=0.44,
                        n_adapt=None):
        freeR = self.freeR
        if n_adapt is None:
            n_adapt = np.minimum(int(n_samples/2), 1000)

        param_samples = np.zeros((n_samples, self.n_params))
        rng = self.rng
        acceptances = np.zeros((n_samples))
        location = self.location.copy()
        pred =  self.W.dot(location)
        z = sp.stats.norm(self.y, self.y.var()).rvs()
        theta = self.t_init.copy()
        progress_bar = tqdm.tqdm(range(n_samples), smoothing=0.01)
        wtrace, waccept = 1.0, 1.0

        for i in progress_bar:
            s2 = theta[-1]
            s = np.sqrt(s2)
            z, accept = self.mh_lvar_binomial(pred, s, z, rng.normal(0, 1, size=self.n_ob),
                                              np.log(rng.uniform(0, 1, size=self.n_ob)), propC)

            mean_accept = accept.mean()
            wtrace = wtrace * damping + 1.0
            waccept = waccept * damping + mean_accept

            if i<n_adapt:
                propC = propC * np.sqrt(adaption_rate**((waccept/wtrace)-target_accept))

            location = self.sample_location(theta, rng.normal(0, 1, size=self.n_re),
                                            rng.normal(0, 1, size=self.n_ob), z)
            pred = self.W.dot(location)
            u = location[-self.n_re:]
            theta = self.sample_theta(theta, u, z, pred, freeR)
            acceptances[i] = mean_accept
            param_samples[i, self.n_fe:] = theta.copy()
            param_samples[i, :self.n_fe] = location[:self.n_fe]
            self._secondary_samples(chain, i, pred, u, z)
            if i>1:
                progress_bar.set_description(f"Chain {chain+1} Acceptance Prob: {acceptances[:i].mean():.4f} C: {propC:.5f}")
        progress_bar.close()
        return param_samples

    def sample_bernoulli(self, n_samples, chain=0, save_pred=False, save_u=False,
                         save_lvar=False):
        freeR = self.freeR
        rng = self.rng
        n_pr, n_ob = self.n_params, self.n_ob
        n_smp = n_samples
        samples = np.zeros((n_smp, n_pr))
        location, pred = self.location.copy(), self.W.dot(self.location)
        theta, z = self.t_init.copy(), sp.stats.norm(0, 1).rvs(n_ob)
        v = np.zeros_like(z).astype(float)
        progress_bar = tqdm.tqdm(range(n_smp), smoothing=0.01)
        progress_bar.set_description(f"Chain {chain+1}")
        for i in progress_bar:
            #P(z|location, theta)
            z = self.slice_sample_lvar(rng.exponential(scale=1.0, size=self.n_ob),
                                       v, z, theta, pred)
            #P(location|z, theta)
            location = self.sample_location(theta, rng.normal(0, 1, size=self.n_re),
                                             rng.normal(0, 1, size=self.n_ob), z)
            pred, u = self.W.dot(location), location[-self.n_re:]
            #P(theta|z, location)
            theta  = self.sample_theta(theta, u, z, pred, freeR)
            samples[i, self.n_fe:] = theta.copy()
            samples[i, :self.n_fe] = location[:self.n_fe]
            self._secondary_samples(chain, i, pred, u, z)
        progress_bar.close()
        return samples

    def sample_normal(self, n_samples, chain=0, save_pred=False, save_u=False,
                      save_lvar=False):
        freeR = self.freeR
        rng = self.rng
        n_pr= self.n_params
        n_smp = n_samples
        samples = np.zeros((n_smp, n_pr))

        y = self.y
        location, pred = self.location.copy(), self.W.dot(self.location)
        theta = self.t_init.copy()
        progress_bar = tqdm.tqdm(range(n_smp), smoothing=0.01)
        progress_bar.set_description(f"Chain {chain+1}")
        for i in progress_bar:
            #P(location|y, theta)
            location = self.sample_location(theta, rng.normal(0, 1, size=self.n_re),
                                             rng.normal(0, 1, size=self.n_ob), y)
            pred, u = self.W.dot(location), location[-self.n_re:]
            #P(theta|z, location)
            theta  = self.sample_theta(theta, u, y, pred, freeR)
            samples[i, self.n_fe:] = theta.copy()
            samples[i, :self.n_fe] = location[:self.n_fe]
            self._secondary_samples(chain, i, pred, u, None)
        progress_bar.close()
        return samples

    def _secondary_samples(self, chain, i, pred=None, u=None, z=None):
        if self.save_pred:
            self.secondary_samples["pred"][chain, i] = pred
        if self.save_u:
            self.secondary_samples["u"][chain, i] = u
        if self.save_lvar:
            self.secondary_samples["lvar"][chain, i] = z

    def sample(self, n_samples=5000, n_chains=8, burnin=1000, save_pred=False,
               save_u=False, save_lvar=False, sampling_kws={},
               summary_kws={}):
        n_params = self.n_params
        if self.response_dist=="ordinal_probit":
            n_params = n_params+np.unique(self.y).shape[0]-1
        samples = np.zeros((n_chains, n_samples, n_params))

        self.secondary_samples = {}
        if save_u:
            self.secondary_samples['u'] = np.zeros((n_chains, n_samples, self.n_re))
        if save_pred:
            self.secondary_samples['pred'] = np.zeros((n_chains, n_samples, self.n_ob))
        if save_lvar:
            self.secondary_samples['lvar'] = np.zeros((n_chains, n_samples, self.n_ob))
        self.save_pred, self.save_u, self.save_lvar = save_pred, save_u, save_lvar

        if self.response_dist=='binomial':
            func = self.sample_binomial
        elif self.response_dist=='bernoulli':
            func = self.sample_bernoulli
        elif self.response_dist=='ordinal_probit':
            func = self.sample_ordinal_probit
        elif self.response_dist=='normal':
            func = self.sample_normal

        for i in range(n_chains):
            samples[i] = func(n_samples, chain=i, **sampling_kws)

        #Heuristic for a heuristic...
        if burnin=="heuristic":
            tmp = []
            for i in np.concatenate(list(self.vnames.values())):
                xv = samples[:, :, i].T
                xm = np.mean(xv, axis=0)
                tmp.append(sign_change(xv, nth_change=50, offset=xm))
            burnin = np.max(tmp)

        if self.response_dist == 'bernoulli':
            c = np.sqrt(1.0+(16*np.sqrt(3))/(15*np.pi))
            if save_u:
                self.secondary_samples['u'] /= c
            if save_pred:
                self.secondary_samples['pred'] /=c
            samples[:, :, self.vnames["$\\beta$"]] /= c
            samples[:, :, self.vnames["$\\theta$"]] /= c

        self.burnin = burnin
        self.samples = samples
        self.az_dict = to_arviz_dict(samples, self.vnames, burnin=burnin)
        self.az_data = az.from_dict(self.az_dict)
        self.summary = az.summary(self.az_data, round_to=6, **summary_kws)
        self.res = self.summary.copy()
        self.res.index = self.param_names
        self.res.insert(2, "z-value", self.res["mean"]/self.res["sd"])
        self.res.insert(3, "p-value", sp.stats.norm(0, 1).sf(np.abs(self.res["z-value"])))
        self.beta = np.mean(samples[:, burnin:, self.vnames["$\\beta$"]], axis=(0, 1))
        self.theta = np.mean(samples[:, burnin:, self.vnames["$\\theta$"]], axis=(0, 1))
        if self.response_dist == 'ordinal_probit':
            self.tau = np.mean(samples[:, burnin:, self.vnames["$\\tau$"]], axis=(0, 1))

    def aggregate_samples(self, burnin=None):
        burnin = self.burnin if burnin is None else burnin
        agg_samples = {}
        for name, ix in self.vnames.items():
            x = self.samples[:, burnin, ix]
            x = x.reshape(np.product(x.shape[:-1]), x.shape[-1])
            agg_samples[name] = x

        if 'u' in self.secondary_samples.keys():
            u = self.secondary_samples["u"][:, burnin:]
            u = u.reshape(np.product(u.shape[:-1]), u.shape[-1])
        agg_samples['u'] = u
        return agg_samples

    def plot_trace(self, plot_kws=None):
        axes = _plot_trace(self.az_data, plot_kws)
        return axes

    def plot_ess(self, plot_kws=None):
        axes = _plot_ess(self.az_data, plot_kws)
        return axes

    def plot_rank(self, plot_kws=None):
        axes = _plot_rank(self.az_data, plot_kws)
        return axes

max_stars_repo_name: lukepinkel/pystats

max_issues_repo_name: lukepinkel/pystats

max_forks_repo_name: lukepinkel/pystats

reason: import numpy,import scipy


text:
# encoding: utf-8
'''
@author: Minghao Guo
@contact: mh.guo0111@gmail.com
@software: nef
@file: sl_phantom.py
@date: 8/21/2019
@desc:
'''
from nefct import nef_class
import numpy as np

def phantom(n = 256, p_type = 'Modified Shepp-Logan', ellipses = None):
    """
     phantom (n = 256, p_type = 'Modified Shepp-Logan', ellipses = None)

    Create a Shepp-Logan or modified Shepp-Logan phantom.

    A phantom is a known object (either real or purely mathematical)
    that is used for testing image reconstruction algorithms.  The
    Shepp-Logan phantom is a popular mathematical model of a cranial
    slice, made up of a set of ellipses.  This allows rigorous
    testing of computed tomography (CT) algorithms as it can be
    analytically transformed with the radon transform (see the
    function `radon').

    Inputs
    ------
    n : The edge length of the square image to be produced.

    p_type : The type of phantom to produce. Either
      "Modified Shepp-Logan" or "Shepp-Logan".  This is overridden
      if `ellipses' is also specified.

    ellipses : Custom set of ellipses to use.  These should be in
      the form
          [[I, a, b, x0, y0, phi],
           [I, a, b, x0, y0, phi],
           ...]
      where each row defines an ellipse.
      I : Additive intensity of the ellipse.
      a : Length of the major axis.
      b : Length of the minor axis.
      x0 : Horizontal offset of the centre of the ellipse.
      y0 : Vertical offset of the centre of the ellipse.
      phi : Counterclockwise rotation of the ellipse in degrees,
            measured as the angles between the horizontal axis and
            the ellipse major axis.
      The image bounding box in the algorithm is [-1, -1], [1, 1],
      so the values of a, b, x0, y0 should all be specified with
      respect to this box.

    Output
    ------
    P : A phantom image.

    Usage example
    -------------
      import matplotlib.pyplot as pl
      P = phantom ()
      pl.imshow (P)

    References
    ----------
    Shepp, L. A.; Logan, B. F.; Reconstructing Interior Head Tissue
    from X-Ray Transmissions, IEEE Transactions on Nuclear Science,
    Feb. 1974, p. 232.

    Toft, P.; "The Radon Transform - Theory and Implementation",
    Ph.D. thesis, Department of Mathematical Modelling, Technical
    University of Denmark, June 1996.

    """

    if (ellipses is None):
        ellipses = _select_phantom(p_type)
    elif (np.size(ellipses, 1) != 6):
        raise AssertionError("Wrong number of columns in user phantom")

    if '3d' not in p_type.lower():
        # Blank image
        p = np.zeros((n, n))

        # Create the pixel grid
        ygrid, xgrid = np.mgrid[-1:1:(1j * n), -1:1:(1j * n)]

        for ellip in ellipses:
            I = ellip[0]
            a2 = ellip[1] ** 2
            b2 = ellip[2] ** 2
            x0 = ellip[3]
            y0 = ellip[4]
            phi = ellip[5] * np.pi / 180  # Rotation angles in radians

            # Create the offset x and y values for the grid
            x = xgrid - x0
            y = ygrid - y0

            cos_p = np.cos(phi)
            sin_p = np.sin(phi)

            # Find the pixels within the ellipse
            locs = (((x * cos_p + y * sin_p) ** 2) / a2
                    + ((y * cos_p - x * sin_p) ** 2) / b2) <= 1

            # Add the ellipse intensity to those pixels
            p[locs] += I
    else:
        p = np.zeros((n, n, n))

        # rng = (np.arange(n) - (n - 1) / 2) / ((n - 1) / 2)
        rng = np.linspace(-1 + 2 / n, 1 - 2 / n, n)
        # Create the pixel grid
        xgrid, ygrid, zgrid = np.meshgrid(rng, rng, rng, indexing = 'ij')
        coord = np.vstack((xgrid.ravel(), ygrid.ravel(), zgrid.ravel()))
        for ellip in ellipses:
            A = ellip[0]
            asq = ellip[1] ** 2
            bsq = ellip[2] ** 2
            csq = ellip[3] ** 2
            x0 = ellip[4]
            y0 = ellip[5]
            z0 = ellip[6]
            phi = ellip[7] * np.pi / 180
            theta = ellip[8] * np.pi / 180
            psi = ellip[9] * np.pi / 180

            cphi = np.cos(phi)
            sphi = np.sin(phi)
            ctheta = np.cos(theta)
            stheta = np.sin(theta)
            cpsi = np.cos(psi)
            spsi = np.sin(psi)

            alpha = np.array([[cpsi * cphi - ctheta * sphi * spsi, cpsi * sphi + ctheta * cphi *
                               spsi, spsi * stheta],
                              [-spsi * cphi - ctheta * sphi * cpsi,
                               -spsi * sphi + ctheta * cphi * cpsi,
                               cpsi * stheta],
                              [stheta * sphi, -stheta * cphi, ctheta]])

            x = xgrid * alpha[0, 0] + ygrid * alpha[1, 0] + zgrid * alpha[2, 0]
            y = xgrid * alpha[0, 1] + ygrid * alpha[1, 1] + zgrid * alpha[2, 1]
            z = xgrid * alpha[0, 2] + ygrid * alpha[1, 2] + zgrid * alpha[2, 2]

            idx = (x - x0) ** 2 / asq + (y - y0) ** 2 / bsq + (z - z0) ** 2 / csq <= 1
            p[idx] += A
    return p.astype(np.float32)

def _select_phantom(name):
    if (name.lower() == 'shepp-logan'):
        e = _shepp_logan()
    elif (name.lower() == 'modified shepp-logan'):
        e = _mod_shepp_logan()
    elif (name.lower() == 'shepp-logan-3d'):
        e = _shepp_logan_3d()
    else:
        raise ValueError("Unknown phantom type: %s" % name)

    return e

def _shepp_logan():
    #  Standard head phantom, taken from Shepp & Logan
    return [[2, .69, .92, 0, 0, 0],
            [-.98, .6624, .8740, 0, -.0184, 0],
            [-.02, .1100, .3100, .22, 0, -18],
            [-.02, .1600, .4100, -.22, 0, 18],
            [.01, .2100, .2500, 0, .35, 0],
            [.01, .0460, .0460, 0, .1, 0],
            [.02, .0460, .0460, 0, -.1, 0],
            [.01, .0460, .0230, -.08, -.605, 0],
            [.01, .0230, .0230, 0, -.606, 0],
            [.01, .0230, .0460, .06, -.605, 0]]

def _shepp_logan_3d():
    #  Standard head phantom, taken from Shepp & Logan
    # return [[2, .69, .92, .81, 0, 0, 0, 0, 0, 0],
    #         [-.98, .6624, .8740, .78, 0, -.0184, 0, 0, 0, 0],
    #         [-.02, .1100, .3100, .22, .22, 0, 0, -18, 0, 10],
    #         [-.02, .1600, .4100, .28, -.22, 0, 0, 18, 0, 10],
    #         [.01, .2100, .2500, .41, 0, .35, -.15, 0, 0, 0],
    #         [.01, .0460, .0460, .05, 0, .1, .25, 0, 0, 0],
    #         [.02, .0460, .0460, .05, 0, -.1, .25, 0, 0, 0],
    #         [.01, .0460, .0230, .05, -.08, -.605, 0, 0, 0, 0],
    #         [.01, .0230, .0230, .02, 0, -.606, 0, 0, 0, 0],
    #         [.01, .0230, .0460, .02, .06, -.605, 0, 0, 0, 0]]

    return [[1, .6900, .920, .810, 0, 0, 0, 0, 0, 0],
            [- .8, .6624, .874, .780, 0, - .0184, 0, 0, 0, 0],
            [- .2, .1100, .310, .220, .22, 0, 0, - 18, 0, 10],
            [- .2, .1600, .410, .280, - .22, 0, 0, 18, 0, 10],
            [.1, .2100, .250, .410, 0, .35, - .15, 0, 0, 0],
            [.1, .0460, .046, .050, 0, .1, .25, 0, 0, 0],
            [.1, .0460, .046, .050, 0, - .1, .25, 0, 0, 0],
            [.1, .0460, .023, .050, - .08, - .605, 0, 0, 0, 0],
            [.1, .0230, .023, .020, 0, - .606, 0, 0, 0, 0],
            [.1, .0230, .046, .020, .06, - .605, 0, 0, 0, 0]]

def _mod_shepp_logan():
    #  Modified version of Shepp & Logan's head phantom,
    #  adjusted to improve contrast.  Taken from Toft.
    return [[1, .69, .92, 0, 0, 0],
            [-.80, .6624, .8740, 0, -.0184, 0],
            [-.20, .1100, .3100, .22, 0, -18],
            [-.20, .1600, .4100, -.22, 0, 18],
            [.10, .2100, .2500, 0, .35, 0],
            [.10, .0460, .0460, 0, .1, 0],
            [.10, .0460, .0460, 0, -.1, 0],
            [.10, .0460, .0230, -.08, -.605, 0],
            [.10, .0230, .0230, 0, -.606, 0],
            [.10, .0230, .0460, .06, -.605, 0]]

ext: py

reason: import numpy


text:
from __future__ import division
import numpy as np
import matplotlib.pyplot as plt
from astropy.io import fits
from astropy.table import Table
import scipy.special as sp

#- redshift errors and zwarn fractions from DESI-1657
#- sigmav = c sigmaz / (1+z)
_sigma_v = {
    'ELG': 19.,
    'LRG': 40.,
    'BGS': 13.,
    'QSO': 423.,
    'STAR': 18.,
    'SKY': 9999,      #- meaningless
    'UNKNOWN': 9999,  #- meaningless
}

_cata_fail_fraction = {
   # Catastrophic error fractions from redmonster on oak (ELG, LRG, QSO)
    'ELG': 0.08,
    'LRG': 0.013,
    'QSO': 0.20,
    'BGS': 0.,
    'STAR': 0.,
    'SKY': 0.,
    'UNKNOWN': 0.,
}

def get_zeff_obs_newmodel(simtype, obsconditions):
    #- only model for 'ELG'
    if(simtype=='ELG'):
        p_v = [0.86, 0.11, -0.02]
        p_w = [0.86, 0.92, 0.62]
        p_x = [0.86, -0.10, -0.01]
        p_y = [0.86, -0.15, -0.20]
        p_z = [0.86, -10.0, 300.0]
        sigma_r = 0.071 #- adding all the delta eff residuals in quad.

    else:
        log.warning('No model for how observing conditions impact {} redshift efficiency'.format(simtype))
        return np.ones(len(obsconditions))
    # airmass
    v = obsconditions['AIRMASS'] - 1.0
    pv  = p_v[0] + p_v[1] * v + p_v[2] * (v**2. - np.mean(v**2))

    # ebmv
    w = obsconditions['EBMV'] - 0.0
    pw = p_w[0] + p_w[1] * w + p_w[2] * (w**2 - np.mean(w**2))

    # seeing
    x = obsconditions['SEEING'] - 1.1
    px = p_x[0] + p_x[1]*x + p_x[2] * (x**2 - np.mean(x**2))

    # transparency
    y = obsconditions['LINTRANS'] - 1.0
    py = p_y[0] + p_y[1]*y + p_y[2] * (y**2 - np.mean(y**2))

    pr = 1.0 + np.random.normal(size=len(y), scale=sigma_r)

    pobs = (pv * pw * px * py * pr).clip(min=0.0)
    return pobs

def get_zeff_obs(simtype, obsconditions):
    '''
    '''
    if(simtype=='LRG'):
        p_v = [1.0, 0.15, -0.5]
        p_w = [1.0, 0.4, 0.0]
        p_x = [1.0, 0.06, 0.05]
        p_y = [1.0, 0.0, 0.08]
        p_z = [1.0, 0.0, 0.0]
        sigma_r = 0.02
    elif(simtype=='QSO'):
        p_v = [1.0, -0.2, 0.3]
        p_w = [1.0, -0.5, 0.6]
        p_x = [1.0, -0.1, -0.075]
        p_y = [1.0, -0.08, -0.04]
        p_z = [1.0, 0.0, 0.0]
        sigma_r = 0.05
    elif(simtype=='ELG'):
        p_v = [1.0, -0.1, -0.2]
        p_w = [1.0, 0.25, -0.75]
        p_x = [1.0, 0.0, 0.05]
        p_y = [1.0, 0.2, 0.1]
        p_z = [1.0, -10.0, 300.0]
        sigma_r = 0.075
    else:
        log.warning('No model for how observing conditions impact {} redshift efficiency'.format(simtype))
        return np.ones(len(obsconditions))

    # airmass
    v = obsconditions['AIRMASS'] - np.mean(obsconditions['AIRMASS'])
    pv  = p_v[0] + p_v[1] * v + p_v[2] * (v**2. - np.mean(v**2))

    # ebmv
    w = obsconditions['EBMV'] - np.mean(obsconditions['EBMV'])
    pw = p_w[0] + p_w[1] * w + p_w[2] * (w**2 - np.mean(w**2))

    # seeing
    x = obsconditions['SEEING'] - np.mean(obsconditions['SEEING'])
    px = p_x[0] + p_x[1]*x + p_x[2] * (x**2 - np.mean(x**2))

    # transparency
    y = obsconditions['LINTRANS'] - np.mean(obsconditions['LINTRANS'])
    py = p_y[0] + p_y[1]*y + p_y[2] * (y**2 - np.mean(y**2))

    # moon illumination fraction
    #z = obsconditions['MOONFRAC'] - np.mean(obsconditions['MOONFRAC'])
    #pz = p_z[0] + p_z[1]*z + p_z[2] * (z**2 - np.mean(z**2))

    #- if moon is down phase doesn't matter
    #pz[obsconditions['MOONALT'] < 0] = 1.0

    pr = 1.0 + np.random.normal(size=len(y), scale=sigma_r)

    #- this correction factor can be greater than 1, but not less than 0
    pobs = (pv * pw * px * py * pr).clip(min=0.0)

    return pobs

def get_redshift_efficiency(simtype, catfile,plot=False, newmodel=False):
    """
    Simple model to get the redshift effiency from the observational conditions or observed magnitudes+redshuft
    Args:
        simtype: ELG, LRG, QSO, MWS, BGS
        targets: target catalog table; currently used only for TARGETID
        truth: truth table with OIIFLUX, TRUEZ
        targets_in_tile: dictionary. Keys correspond to tileids, its values are the
            arrays of targetids observed in that tile.
        obsconditions: table observing conditions with columns
           'TILEID': array of tile IDs
           'AIRMASS': array of airmass values on a tile
           'EBMV': array of E(B-V) values on a tile
           'LINTRANS': array of atmospheric transparency during spectro obs; floats [0-1]
           'MOONFRAC': array of moonfraction values on a tile.
           'SEEING': array of FWHM seeing during spectroscopic observation on a tile.
    Returns:
        tuple of arrays (observed, p) both with same length as targets
        observed: boolean array of whether the target was observed in these tiles
        p: probability to get this redshift right
    """

    catalog=fits.open(catfile)
    datacat=catalog[1].data
    n=len(datacat)

    obsconditions={'AIRMASS': datacat['AIRMASS'], 'SEEING': datacat['SEEING'], 'LINTRANS': datacat['TRANSPARENCY'], 'EBMV': np.zeros(n)}#datacat['EBMV']}

    if (simtype == 'ELG'):
        # Read the model OII flux threshold (FDR fig 7.12 modified to fit redmonster efficiency on OAK)
        filename = ('./quickcat_elg_oii_flux_threshold.txt')
        fdr_z, modified_fdr_oii_flux_threshold = np.loadtxt(filename, unpack=True)

        # Get OIIflux from truth
        true_oii_flux = datacat['OIIFLUX']

        # Compute OII flux thresholds for truez
        oii_flux_threshold = np.interp(datacat['Z'],fdr_z,modified_fdr_oii_flux_threshold)
        assert (oii_flux_threshold.size == true_oii_flux.size),"oii_flux_threshold and true_oii_flux should have the same size"
        if plot:
            fig=plt.figure()
            ax=fig.add_subplot(111)
            selz=np.where((fdr_z>= 0.6) & (fdr_z<=1.6))[0]
            ax.plot(fdr_z[selz],modified_fdr_oii_flux_threshold[selz],'r-',label='[OII] flux threshold',lw=2)
            ax.legend(fontsize=22,numpoints=1)
            ax.set_xlabel(r'$z$',fontsize=20)
            ax.set_xlim(0.6,1.6)
            ax.set_ylabel(r'${\rm [OII]\ flux\ [ergs/s/cm^2]}$',fontsize=20)
            ax.tick_params('both',labelsize=16)
            plt.show()

        # efficiency is modeled as a function of flux_OII/f_OII_threshold(z) and an arbitrary sigma_fudge
        sigma_fudge = 1.0
        max_efficiency = 1.0
        simulated_eff = eff_model(true_oii_flux/oii_flux_threshold,sigma_fudge,max_efficiency)

    else:
        default_zeff = 0.98
        log.warning('using default redshift efficiency of {} for {}'.format(default_zeff, simtype))
        simulated_eff = default_zeff * np.ones(n)

    if newmodel:
        print("using new efficiency model instead of original")
        zeff_obs = get_zeff_obs_newmodel(simtype,obsconditions)
    else:
        print("Using original quickcat efficiency model")
        zeff_obs = get_zeff_obs(simtype, obsconditions)
    pfail = np.ones(n)
    #observed = np.zeros(n, dtype=bool)

    tmp=(simulated_eff*zeff_obs).clip(0,1)
    pfail*=(1-tmp)
    simulated_eff=(1-pfail)
    return simulated_eff

# Efficiency model
def eff_model(x, sigma, max_efficiency):
    return 0.5*max_efficiency*(1.+sp.erf((x-1)/(np.sqrt(2.)*sigma)))

def get_observed_redshifts(catfile,newmodel=False):
    """
    Returns observed z, zerr, zwarn arrays given true object types and redshifts
    Args:
        targets: target catalog table; currently used only for target mask bits
        truth: truth table with OIIFLUX, TRUEZ
        targets_in_tile: dictionary. Keys correspond to tileids, its values are the
            arrays of targetids observed in that tile.
        obsconditions: table observing conditions with columns
           'TILEID': array of tile IDs
           'AIRMASS': array of airmass values on a tile
           'EBMV': array of E(B-V) values on a tile
           'LINTRANS': array of atmospheric transparency during spectro obs; floats [0-1]
           'MOONFRAC': array of moonfraction values on a tile.
           'SEEING': array of FWHM seeing during spectroscopic observation on a tile.
    Returns:
        tuple of (zout, zerr, zwarn)
    """

    catalog=fits.open(catfile)
    datacat=catalog[1].data
    truez = datacat['Z']

    zout = truez.copy()
    zerr = np.zeros(len(truez), dtype=np.float32)
    zwarn = np.zeros(len(truez), dtype=np.int32)
    simtype = np.array(['ELG']*len(truez))
    objtypes = ['ELG']
    print len(objtypes)

    for objtype in objtypes:
        if objtype in _sigma_v.keys():
            ii = np.where(simtype==objtype)[0]

            n = len(ii)#np.count_nonzero(ii)
            print n
            # Error model for ELGs
            if (objtype =='ELG'):
                filename = './quickcat_elg_oii_errz.txt'
                oii, errz_oii = np.loadtxt(filename, unpack=True)
                try:
                    true_oii_flux = datacat['OIIFLUX'][ii]
                except:
                    raise Exception('Missing OII flux information to estimate redshift error for ELGs')

                mean_err_oii = np.interp(true_oii_flux,oii,errz_oii)
                zerr[ii] = mean_err_oii*(1.+truez[ii])
                zout[ii] += np.random.normal(scale=zerr[ii])

            else:
                zerr[ii] = _sigma_v[objtype] * (1+truez[ii]) / c
                zout[ii] += np.random.normal(scale=zerr[ii])

            # Set ZWARN flags for some targets
            # the redshift efficiency only sets warning, but does not impact
            # the redshift value and its error.
            goodz_prob = get_redshift_efficiency(objtype, catfile, newmodel=newmodel)

            assert len(goodz_prob) == n
            r = np.random.random(n)
            zwarn[ii] = 4 * (r > goodz_prob)

            # Add fraction of catastrophic failures (zwarn=0 but wrong z)
            nzwarnzero = np.count_nonzero(zwarn[ii] == 0)
            num_cata = np.random.poisson(_cata_fail_fraction[objtype] * nzwarnzero)
            if (objtype == 'ELG'): zlim=[0.6,0.7]
            if num_cata > 0:
                #- tmp = boolean array for all targets, flagging only those
                #- that are of this simtype and were observed this epoch
                kk, = np.where((zwarn==0))
                index = np.random.choice(kk, size=num_cata, replace=False)
                assert np.all(np.in1d(index, np.where(ii)[0]))
                assert np.all(zwarn[index] == 0)

                zout[index] = np.random.uniform(zlim[0],zlim[1],len(index))

        else:
            msg = 'No redshift efficiency model for {}; using true z\n'.format(objtype) + \
                  'Known types are {}'.format(list(_sigma_v.keys()))
            print(msg)

    return zout, zerr, zwarn

def quickcat(catfile,outfile,newmodel=False):

    cat=fits.open(catfile)
    datacat=cat[1].data

    ra=datacat['RA']
    dec=datacat['DEC']

    zout,zerr,zwarn=get_observed_redshifts(catfile,newmodel=newmodel)

    kk=np.where(zwarn==0)[0]
    ra=ra[kk]
    dec=dec[kk]
    z=zout[kk]
    zwarn=zwarn[kk]
    zerr=zerr[kk]

    zcatalog=Table([ra,dec,z,zerr,zwarn],names=('RA','DEC','Z','ZERR','ZWARN'))
    zcatalog.write(outfile,format='fits',overwrite=True)

reason: import numpy,import scipy,from astropy


text:
import numpy as np

class Matrix:
    def __init__(self, s):
        self.matrix = np.array([col for col in [list(map(int, row.split())) for row in s.split("\n")]])

    def row(self, index):
        return list(self.matrix[index-1, :])

    def column(self, index):
        return list(self.matrix[:, index-1])


text:
"""
Load datasets for Low Data processing.
"""
from __future__ import print_function
from __future__ import division
from __future__ import unicode_literals

import os
import sys
import shutil
import tempfile
import numpy as np
import deepchem as dc

def to_numpy_dataset(dataset):
  """Converts dataset to numpy dataset."""
  return dc.data.NumpyDataset(dataset.X, dataset.y, dataset.w, dataset.ids)

def load_tox21_ecfp(num_train=7200):
  """Load Tox21 datasets. Does not do train/test split"""
  # Set some global variables up top
  current_dir = os.path.dirname(os.path.realpath(__file__))
  dataset_file = os.path.join(current_dir, "../../datasets/tox21.csv.gz")
  # Featurize Tox21 dataset
  print("About to featurize Tox21 dataset.")
  featurizer = dc.feat.CircularFingerprint(size=1024)
  tox21_tasks = [
      'NR-AR', 'NR-AR-LBD', 'NR-AhR', 'NR-Aromatase', 'NR-ER', 'NR-ER-LBD',
      'NR-PPAR-gamma', 'SR-ARE', 'SR-ATAD5', 'SR-HSE', 'SR-MMP', 'SR-p53'
  ]

  loader = dc.data.CSVLoader(
      tasks=tox21_tasks, smiles_field="smiles", featurizer=featurizer)
  dataset = loader.featurize(dataset_file, shard_size=8192)

  # Initialize transformers
  transformers = [dc.trans.BalancingTransformer(dataset=dataset)]

  print("About to transform data")
  for transformer in transformers:
    dataset = transformer.transform(dataset)

  return tox21_tasks, dataset, transformers

def load_tox21_convmol(base_dir=None, num_train=7200):
  """Load Tox21 datasets. Does not do train/test split"""
  # Set some global variables up top
  current_dir = os.path.dirname(os.path.realpath(__file__))
  dataset_file = os.path.join(current_dir, "../../datasets/tox21.csv.gz")

  # Featurize Tox21 dataset
  print("About to featurize Tox21 dataset.")
  featurizer = dc.feat.ConvMolFeaturizer()
  tox21_tasks = [
      'NR-AR', 'NR-AR-LBD', 'NR-AhR', 'NR-Aromatase', 'NR-ER', 'NR-ER-LBD',
      'NR-PPAR-gamma', 'SR-ARE', 'SR-ATAD5', 'SR-HSE', 'SR-MMP', 'SR-p53'
  ]

  loader = dc.data.CSVLoader(
      tasks=tox21_tasks, smiles_field="smiles", featurizer=featurizer)
  dataset = loader.featurize(dataset_file, shard_size=8192)

  # Initialize transformers
  transformers = [dc.trans.BalancingTransformer(dataset=dataset)]

  print("About to transform data")
  for transformer in transformers:
    dataset = transformer.transform(dataset)

  return tox21_tasks, dataset, transformers

def load_muv_ecfp():
  """Load MUV datasets. Does not do train/test split"""
  # Load MUV dataset
  print("About to load MUV dataset.")
  current_dir = os.path.dirname(os.path.realpath(__file__))
  dataset_file = os.path.join(current_dir, "../../datasets/muv.csv.gz")
  # Featurize MUV dataset
  print("About to featurize MUV dataset.")
  featurizer = dc.feat.CircularFingerprint(size=1024)
  MUV_tasks = sorted([
      'MUV-692', 'MUV-689', 'MUV-846', 'MUV-859', 'MUV-644', 'MUV-548',
      'MUV-852', 'MUV-600', 'MUV-810', 'MUV-712', 'MUV-737', 'MUV-858',
      'MUV-713', 'MUV-733', 'MUV-652', 'MUV-466', 'MUV-832'
  ])

  loader = dc.data.CSVLoader(
      tasks=MUV_tasks, smiles_field="smiles", featurizer=featurizer)
  dataset = loader.featurize(dataset_file)

  # Initialize transformers
  transformers = [dc.trans.BalancingTransformer(dataset=dataset)]
  print("About to transform data")
  for transformer in transformers:
    dataset = transformer.transform(dataset)

  return MUV_tasks, dataset, transformers

def load_muv_convmol():
  """Load MUV datasets. Does not do train/test split"""
  # Load MUV dataset
  print("About to load MUV dataset.")
  current_dir = os.path.dirname(os.path.realpath(__file__))
  dataset_file = os.path.join(current_dir, "../../datasets/muv.csv.gz")
  # Featurize MUV dataset
  print("About to featurize MUV dataset.")
  featurizer = dc.feat.ConvMolFeaturizer()
  MUV_tasks = sorted([
      'MUV-692', 'MUV-689', 'MUV-846', 'MUV-859', 'MUV-644', 'MUV-548',
      'MUV-852', 'MUV-600', 'MUV-810', 'MUV-712', 'MUV-737', 'MUV-858',
      'MUV-713', 'MUV-733', 'MUV-652', 'MUV-466', 'MUV-832'
  ])

  loader = dc.data.CSVLoader(
      tasks=MUV_tasks, smiles_field="smiles", featurizer=featurizer)
  dataset = loader.featurize(dataset_file)

  # Initialize transformers
  transformers = [dc.trans.BalancingTransformer(dataset=dataset)]
  print("About to transform data")
  for transformer in transformers:
    dataset = transformer.transform(dataset)

  return MUV_tasks, dataset, transformers

def load_sider_ecfp():
  """Load SIDER datasets. Does not do train/test split"""
  # Featurize SIDER dataset
  print("About to featurize SIDER dataset.")
  current_dir = os.path.dirname(os.path.realpath(__file__))
  dataset_file = os.path.join(current_dir, "../sider/sider.csv.gz")
  featurizer = dc.feat.CircularFingerprint(size=1024)

  dataset = dc.utils.save.load_from_disk(dataset_file)
  SIDER_tasks = dataset.columns.values[1:].tolist()
  print("SIDER tasks: %s" % str(SIDER_tasks))
  print("%d tasks in total" % len(SIDER_tasks))

  loader = dc.data.CSVLoader(
      tasks=SIDER_tasks, smiles_field="smiles", featurizer=featurizer)
  dataset = loader.featurize(dataset_file)
  print("%d datapoints in SIDER dataset" % len(dataset))

  # Initialize transformers
  transformers = [dc.trans.BalancingTransformer(dataset=dataset)]
  print("About to transform data")
  for transformer in transformers:
    dataset = transformer.transform(dataset)

  return SIDER_tasks, dataset, transformers

def load_sider_convmol():
  """Load SIDER datasets. Does not do train/test split"""
  # Featurize SIDER dataset
  print("About to featurize SIDER dataset.")
  current_dir = os.path.dirname(os.path.realpath(__file__))
  dataset_file = os.path.join(current_dir, "../sider/sider.csv.gz")
  featurizer = dc.feat.ConvMolFeaturizer()

  dataset = dc.utils.save.load_from_disk(dataset_file)
  SIDER_tasks = dataset.columns.values[1:].tolist()
  print("SIDER tasks: %s" % str(SIDER_tasks))
  print("%d tasks in total" % len(SIDER_tasks))

  loader = dc.data.CSVLoader(
      tasks=SIDER_tasks, smiles_field="smiles", featurizer=featurizer)
  dataset = loader.featurize(dataset_file)
  print("%d datapoints in SIDER dataset" % len(dataset))

  # Initialize transformers
  transformers = [dc.trans.BalancingTransformer(dataset=dataset)]
  print("About to transform data")
  for transformer in transformers:
    dataset = transformer.transform(dataset)

  return SIDER_tasks, dataset, transformers

max_stars_repo_name: cjgalvin/deepchem

max_issues_repo_name: cjgalvin/deepchem

max_forks_repo_name: cjgalvin/deepchem


text:
# Collection of chatbot models, each held within its own class
# The imports needed for each model are presented before that class

# Vanilla three lawyer artificial neural network

from keras.models import Sequential # basic Keras model
from keras.layers import Dense, Activation # Necessary functions for NN layers and computation

class NLPv1(object):
	def __init__(self,name):
		self.name = name

	def create(self, name):
		self.name = str(name)
		model = Sequential([

			# creates a hiidden layer with 12 nodes and an input layer
			# with 20 nodes
		   Dense(12, input_shape=(20,)),
		  #Dense(12, input_dim=20), # same as above

		  # Set softmax as the normalizing function for hidden layer outputs
		   Activation('softmax'),

		   Dense(5), # create a 5 node output layer
		   Activation('softmax'), # Give the output layer a softmax function too
		])

		# This is for reference if you wanted the NN to be deeper
		#model.add to add layers
		#model.add(Dense(32, activation='relu'))

		return model

	def train(self,model,data,labels,batch,epoch):

		# Set the model to train with the following optimization function
		#to find the best 'metric' listed
		model.compile(optimizer='rmsprop',
			loss='categorical_crossentropy',
			metrics=['accuracy'])

		# Train the model to the data
		#train with 32 samples/batches (batch: number of samples before backproping)
		model.fit(data, labels, epochs=epoch, batch_size=batch)

		return model

	def save(self,name,model):
		# save the weights of the model once trained
		name = str(name)
		try:
			model.save_weights(name+".h5")
			print("Model weights saved successfully.")
		except:
			raise('Unknown error risen while saving weights. Perhaps try again.')
		return model

	def load(self,name):
		name= str(name)
		model = Sequential([
		   Dense(12, input_shape=(20,)),
		  #Dense(12, input_dim=20), # same as above
		   Activation('softmax'),
		   Dense(5),
		   Activation('softmax'),
		])

		#To reload the weights later, use:
		model.load_weights(name+".h5")
		return model

	def query(self,model,input,batch_number=1):
		return model.predict(input, batch_size=batch_number)

		#test the model
		#model.evaluate(input_data,output_data,batch_size=sample_number)

		#query the model
		#model.predict(input_data, batch_size=number_of_samples_for_prediction)

from __future__ import print_function

from keras.models import Model
from keras.layers import Input, LSTM, Dense
import numpy as np

class SeqtoSeq(object):
	""" NLP model taken from Keras' example of Seq2Seq for language translation. Should also work
	adequately on text generation and as a starting point for more advanced models. Keras' model
	is built to translate based on character level, but this model uses a word level text
	generation. This model uses Keras' built in LSTM function. It would also be good to build
	this model from stratch, but then again, that would just be like rewritting keras, which isn't
	useful."""

	def __init__(self,batch_size=64,epochs=100,latent_dim=256,num_samples=10000): # All of this inf
		self.batch_size = batch_size
		self.epochs = epochs
		self.latent_dim = latent_dim
		self.num_samples = num_samples

	def create(self):

		enoder_input_data = np.zeros()

		return True

	def train(self,sample,target):

		# Vectorize the data
		sample = sample.split(' ')
		target = target.split(' ')
		sample = sorted(list(sample))
		target = sorted(list(target))
		num_encoder_tokens = len(sample)
		num_decoder_tokens = len(target)

		# The max encoder length should be set based on multiple samples,
		#however this function is designed to take in a single input. Thus
		#the following lines are commented out.
		#max_encoder_seq_length = max([len(txt) for txt in sample_texts])
		#max_decoder_seq_length = max([len(txt) for txt in target_texts])

		sample_token_index = dict([(char, i) for i, char in enumerate(sample)])
		target_token_index = dict([(char, i) for i, char in enumerate(target)])

		#TODO - Put lines 103 to 111 in the create() function
		#TODO - continue implementing Seq2Seq example, have read to line 113

		return True

	def save(self):
		return True

	def load(self):
		return True

	def query(self):
		return True

max_stars_repo_name: francoisrice/conversation-bot

max_issues_repo_name: francoisrice/conversation-bot

max_forks_repo_name: francoisrice/conversation-bot


text:
import os
import pathlib
import shutil
import tempfile
import unittest
from unittest.mock import patch

import anndata
import boto3
import numpy
import pandas

from moto import mock_s3

from backend.corpora.common.corpora_orm import (
    CollectionVisibility,
    DatasetArtifactType,
    DatasetArtifactFileType,
    UploadStatus,
    ValidationStatus,
)
from backend.corpora.common.entities.collection import Collection
from backend.corpora.common.entities.dataset import Dataset

from backend.corpora.dataset_processing import process

class TestDatasetProcessing(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.tmp_dir = tempfile.mkdtemp()
        cls.h5ad_filename = pathlib.Path(cls.tmp_dir, "test.h5ad")
        cls.seurat_filename = pathlib.Path(cls.tmp_dir, "test.rds")
        cls.loom_filename = pathlib.Path(cls.tmp_dir, "test.loom")

        cls.h5ad_filename.touch()
        cls.seurat_filename.touch()
        cls.loom_filename.touch()

    @classmethod
    def tearDownClass(cls):
        shutil.rmtree(cls.tmp_dir)

    @patch.dict(
        os.environ,
        {
            "DROPBOX_URL": "xxx",
            "ARTIFACT_BUCKET": "yyy",
            "CELLXGENE_BUCKET": "zzz",
            "DATASET_ID": "aaa",
            "DEPLOYMENT_STAGE": "test",
        },
    )
    def test_check_env_good(self):
        process.check_env()

    @patch.dict(os.environ, {"DROPBOX_URL": "xxx", "ARTIFACT_BUCKET": "yyy", "CELLXGENE_BUCKET": "zzz"})
    def test_check_env_bad(self):
        with self.assertRaises(EnvironmentError):
            process.check_env()

    @patch("scanpy.read_h5ad")
    def test_extract_metadata(self, mock_read_h5ad):

        df = pandas.DataFrame(
            numpy.random.randint(10, size=(50001, 5)) * 50, columns=list("ABCDE"), index=(str(i) for i in range(50001))
        )

        tissue = numpy.random.choice([0, 1], size=(50001))
        assay = numpy.random.choice([0, 1, 2], size=(50001))
        eth = numpy.random.choice([0, 1], size=(50001))
        dev = numpy.random.choice([0, 1, 2], size=(50001))

        obs = pandas.DataFrame(
            numpy.hstack(
                [
                    numpy.array([["lung", "liver"][i] for i in tissue]).reshape(50001, 1),
                    numpy.array([["UBERON:01", "UBERON:10"][i] for i in tissue]).reshape(50001, 1),
                    numpy.array([["10x", "smartseq", "cite-seq"][i] for i in assay]).reshape(50001, 1),
                    numpy.array([["EFO:001", "EFO:010", "EFO:011"][i] for i in assay]).reshape(50001, 1),
                    numpy.random.choice(["healthy"], size=(50001, 1)),
                    numpy.random.choice(["MONDO:123"], size=(50001, 1)),
                    numpy.random.choice(["male", "female"], size=(50001, 1)),
                    numpy.array([["solomon islander", "orcadian"][i] for i in eth]).reshape(50001, 1),
                    numpy.array([["HANCESTRO:321", "HANCESTRO:456"][i] for i in eth]).reshape(50001, 1),
                    numpy.array([["adult", "baby", "tween"][i] for i in dev]).reshape(50001, 1),
                    numpy.array([["HsapDv:0", "HsapDv:1", "HsapDv:2"][i] for i in dev]).reshape(50001, 1),
                ]
            ),
            columns=[
                "tissue",
                "tissue_ontology_term_id",
                "assay",
                "assay_ontology_term_id",
                "disease",
                "disease_ontology_term_id",
                "sex",
                "ethnicity",
                "ethnicity_ontology_term_id",
                "development_stage",
                "development_stage_ontology_term_id",
            ],
            index=(str(i) for i in range(50001)),
        )
        uns = {
            "organism": "Homo sapiens",
            "organism_ontology_term_id": "NCBITaxon:8505",
            "layer_descriptions": {"X": "raw"},
        }

        adata = anndata.AnnData(X=df, obs=obs, uns=uns)
        mock_read_h5ad.return_value = adata

        extracted_metadata = process.extract_metadata("dummy")
        lab, ont = "label", "ontology_term_id"

        self.assertDictEqual(extracted_metadata["organism"], {lab: "Homo sapiens", ont: "NCBITaxon:8505"})

        def list_equal(list1, list2, cmp_func):
            self.assertEqual(len(list1), len(list2))
            for el1 in list1:
                self.assertIn(el1, list2)
                el2 = list2[list2.index(el1)]
                if cmp_func:
                    cmp_func(el1, el2)

        list_equal(
            extracted_metadata["tissue"],
            [{lab: "lung", ont: "UBERON:01"}, {lab: "liver", ont: "UBERON:10"}],
            self.assertDictEqual,
        )

        list_equal(
            extracted_metadata["assay"],
            [{lab: "10x", ont: "EFO:001"}, {lab: "smartseq", ont: "EFO:010"}, {lab: "cite-seq", ont: "EFO:011"}],
            self.assertDictEqual,
        )

        list_equal(extracted_metadata["disease"], [{lab: "healthy", ont: "MONDO:123"}], self.assertDictEqual)

        self.assertListEqual(sorted(extracted_metadata["sex"]), sorted(["male", "female"]))

        list_equal(
            extracted_metadata["ethnicity"],
            [{lab: "solomon islander", ont: "HANCESTRO:321"}, {lab: "orcadian", ont: "HANCESTRO:456"}],
            self.assertDictEqual,
        )

        list_equal(
            extracted_metadata["development_stage"],
            [{lab: "adult", ont: "HsapDv:0"}, {lab: "baby", ont: "HsapDv:1"}, {lab: "tween", ont: "HsapDv:2"}],
            self.assertDictEqual,
        )
        self.assertEqual(extracted_metadata["cell_count"], 50001)
        self.assertAlmostEqual(extracted_metadata["mean_genes_per_cell"], numpy.count_nonzero(df) / 50001)

    def test_update_db(self):

        collection = Collection.create(visibility=CollectionVisibility.PRIVATE)
        dataset = Dataset.create(collection_id=collection.id, collection_visibility=CollectionVisibility.PRIVATE)
        dataset_id = dataset.id

        fake_env = patch.dict(os.environ, {"DATASET_ID": dataset_id, "DEPLOYMENT_STAGE": "test"})
        fake_env.start()

        process.update_db(metadata={"sex": ["male", "female"]})

        self.assertListEqual(Dataset.get(dataset_id).sex, ["male", "female"])

        artifact = {
            "filename": "test_filename",
            "filetype": DatasetArtifactFileType.H5AD,
            "type": DatasetArtifactType.REMIX,
            "user_submitted": True,
            "s3_uri": "s3://test_uri",
        }
        dep_dir = {"url": "https://cellxgene.com/data"}
        process.update_db(metadata={"artifacts": [artifact], "deployment_directories": [dep_dir]})

        self.assertEqual(len(Dataset.get(dataset_id).artifacts), 1)
        self.assertEqual(Dataset.get(dataset_id).artifacts[0].filename, "test_filename")
        self.assertEqual(Dataset.get(dataset_id).deployment_directories[0].url, "https://cellxgene.com/data")

        process.update_db(processing_status={"upload_status": UploadStatus.UPLOADING, "upload_progress": 0.5})
        self.assertEqual(Dataset.get(dataset_id).processing_status.upload_status, UploadStatus.UPLOADING)
        self.assertEqual(Dataset.get(dataset_id).processing_status.upload_progress, 0.5)
        self.assertIsNone(Dataset.get(dataset_id).processing_status.validation_status)

        process.update_db(
            processing_status={
                "upload_status": UploadStatus.UPLOADED,
                "upload_progress": 1,
                "validation_status": ValidationStatus.VALIDATING,
            }
        )
        self.assertEqual(Dataset.get(dataset_id).processing_status.upload_status, UploadStatus.UPLOADED)
        self.assertEqual(Dataset.get(dataset_id).processing_status.upload_progress, 1)
        self.assertEqual(Dataset.get(dataset_id).processing_status.validation_status, ValidationStatus.VALIDATING)

        fake_env.stop()

    def test_create_artifacts(self):
        s3_mock = mock_s3()
        s3_mock.start()
        s3 = boto3.client("s3")
        s3.create_bucket(Bucket="test_bucket")

        fake_env = patch.dict(os.environ, {"DATASET_ID": "aaaa-bbbb-cccc", "ARTIFACT_BUCKET": "test_bucket"})
        fake_env.start()

        artifacts = process.create_artifacts(
            str(self.h5ad_filename), str(self.seurat_filename), str(self.loom_filename)
        )

        self.assertEqual(len(artifacts), 3)

        self.assertTrue(all(a["user_submitted"] for a in artifacts))
        self.assertTrue(all(a["s3_uri"].startswith("s3://test_bucket/aaaa-bbbb-cccc/") for a in artifacts))
        self.assertEqual(len(set(a["filetype"] for a in artifacts)), 3)
        self.assertTrue(all(a["type"] == DatasetArtifactType.REMIX for a in artifacts))

        resp = s3.list_objects_v2(Bucket="test_bucket", Prefix="aaaa-bbbb-cccc")
        s3_filenames = [os.path.basename(c["Key"]) for c in resp["Contents"]]
        self.assertEqual(len(s3_filenames), 3)
        self.assertIn(str(self.h5ad_filename.parts[-1]), s3_filenames)
        self.assertIn(str(self.seurat_filename.parts[-1]), s3_filenames)
        self.assertIn(str(self.loom_filename.parts[-1]), s3_filenames)

        fake_env.stop()
        s3_mock.stop()


text:
import pandas as pd
import numpy as np
from pathlib import Path

if __name__ == "__main__":
      csv_path = "geninus_prediction_scnym_ssl.csv"

      df = pd.read_csv(csv_path)
      gen_pred = df["pred_celltype_l2"].values
      sc_pred = df["celltype_pred_scnym"].values

      print(f'Number of same predictions are {np.sum(gen_pred==sc_pred)} among total {len(gen_pred)} cells, which is ratio of {np.sum(gen_pred==sc_pred)/len(gen_pred)}.')

      confusion_matrix = pd.crosstab(df['pred_celltype_l2'], df['celltype_pred_scnym'], rownames=['Geninus'], colnames=['scNym'])
      confusion_matrix['Total'] = confusion_matrix.sum(axis=1)
      confusion_matrix.loc["Total"] = confusion_matrix.sum()

      for key in confusion_matrix.keys():
            if key not in ["Total", "Recall"]:
                  confusion_matrix.at["Precision", key]= confusion_matrix[key][key] / confusion_matrix[key]["Total"]
      for key in confusion_matrix.keys():
            if key not in ["Total", "Recall", "Precision"]:
                  confusion_matrix.at[key, "Recall"]= confusion_matrix[key][key] / confusion_matrix["Total"][key]

      confusion_matrix.to_csv("confusion_matrix_geninus_scnym_ssl.csv")

      seurat_dir = Path("/home/svcapp/tbrain_x/SKT_data_prediction_by_Seurat")
      csv_files = seurat_dir.glob("*.csv")
      seurat_df = pd.concat([pd.read_csv(csv) for csv in csv_files])

      df = df.merge(seurat_df)

      sr_pred = df["predicted.celltype.l2"]
      print(f'Geninus-Seurat: Number of same predictions are {np.sum(gen_pred==sr_pred)} among total {len(gen_pred)} cells, which is ratio of {np.sum(gen_pred==sr_pred)/len(gen_pred)}.')
      print(f'scNym-Seurat: Number of same predictions are {np.sum(sc_pred==sr_pred)} among total {len(sc_pred)} cells, which is ratio of {np.sum(sc_pred==sr_pred)/len(sc_pred)}.')

      confusion_matrix = pd.crosstab(df['pred_celltype_l2'], df['predicted.celltype.l2'], rownames=['Geninus'], colnames=['Seurat'])
      # confusion_matrix.to_csv("confusion_matrix_geninus_seurat.csv")
      # confusion_matrix = pd.crosstab(df['celltype_pred_scnym'], df['predicted.celltype.l2'], rownames=['scNym'], colnames=['Seurat'])
      # confusion_matrix.to_csv("confusion_matrix_scnym_ssl_seurat.csv")

max_stars_repo_name: jdasam/scnym

max_issues_repo_name: jdasam/scnym

max_forks_repo_name: jdasam/scnym


text:
from tensorflow import keras
from tensorflow.keras import backend as K
from keras_contrib.losses.jaccard import jaccard_distance
import numpy as np

# Test and plot
y_pred = np.array([np.arange(-10, 10 + 0.1, 0.1)]).T
y_true = np.zeros(y_pred.shape)
name = 'jaccard_distance_loss'
try:
    loss = jaccard_distance_loss(
        K.variable(y_true), K.variable(y_pred)
    ).eval(session=K.get_session())
except Exception as e:
    print("error plotting", name, e)
else:
    plt.title(name)
    plt.plot(y_pred, loss)
    plt.show()

print("TYPE                 |Almost_right |half right |all_wrong")
y_true = np.array([[0, 0, 1, 0], [0, 0, 1, 0], [0, 0, 1., 0.]])
y_pred = np.array([[0, 0, 0.9, 0], [0, 0, 0.1, 0], [1, 1, 0.1, 1.]])

r = jaccard_distance(
    K.variable(y_true),
    K.variable(y_pred),
).eval(session=K.get_session())
print('jaccard_distance_loss', r)
assert r[0] < r[1]
assert r[1] < r[2]

r = keras.losses.binary_crossentropy(
    K.variable(y_true),
    K.variable(y_pred),
).eval(session=K.get_session())
print('binary_crossentropy', r)
print('binary_crossentropy_scaled', r / r.max())
assert r[0] < r[1]
assert r[1] < r[2]

"""
TYPE                 |Almost_right |half right |all_wrong
jaccard_distance_loss [ 0.09900928  0.89108944  3.75000238]
binary_crossentropy [  0.02634021   0.57564634  12.53243446]
binary_crossentropy_scaled [ 0.00210176  0.04593252  1.        ]
"""

max_stars_repo_name: HaemanthSP/keras-contrib

max_issues_repo_name: HaemanthSP/keras-contrib

max_forks_repo_name: HaemanthSP/keras-contrib


text:
import logging
from typing import List
import numpy as np
import scipy
from deepchem.feat.base_classes import Featurizer
from typing import Any, Iterable

logger = logging.getLogger(__name__)
CHARSET = [
    'A', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'K', 'L', 'M', 'N', 'P', 'Q', 'R',
    'S', 'T', 'V', 'W', 'Y', 'X', 'Z', 'B', 'U', 'O'
]

class SparseMatrixOneHotFeaturizer(Featurizer):
  """Encodes any arbitrary string as a one-hot array.

  This featurizer uses the sklearn OneHotEncoder to create
  sparse matrix representation of a one-hot array of any string.
  It is expected to be used in large datasets that produces memory overload
  using standard featurizer such as OneHotFeaturizer. For example: SwissprotDataset

  Examples
  --------
  >>> import deepchem as dc
  >>> featurizer = dc.feat.SparseMatrixOneHotFeaturizer()
  >>> sequence = "MMMQLA"
  >>> encodings = featurizer.featurize([sequence])
  >>> encodings[0].shape
  (6, 25)

  """

  def __init__(self, charset: List[str] = CHARSET):
    """Initialize featurizer.

    Parameters
    ----------
    charset: List[str] (default code)
      A list of strings, where each string is length 1 and unique.
    """
    if len(charset) != len(set(charset)):
      raise ValueError("All values in charset must be unique.")
    self.charset = charset
    from sklearn.preprocessing import OneHotEncoder
    cat = np.array(self.charset).reshape(1, len(self.charset))
    self.ohe = OneHotEncoder(categories=list(cat), handle_unknown='ignore')

  def featurize(self,
                datapoints: Iterable[Any],
                log_every_n: int = 1000,
                **kwargs) -> np.ndarray:
    """Featurize strings.

    Parameters
    ----------
    datapoints: list
      A list of either strings (str or numpy.str_)
    log_every_n: int, optional (default 1000)
      How many elements are featurized every time a featurization is logged.
    """
    # Featurize data using featurize() in parent class
    return Featurizer.featurize(self, datapoints, log_every_n)

  def _featurize(self, datapoint: Any, **kwargs):
    """ Use parent method of base clase Featurizer.

    Parameters
    ----------
    datapoint : list of string
        string to be converted to a sparse one hot matrix.

    Returns
    -------
    scipy sparse matrix
      A scipy sparse matrix of the one hot representation of the given string.
    """
    # Featurize str data
    if isinstance(datapoint, (str, np.str_)):
      sequence = np.array(list(datapoint)).reshape(-1, 1)
      sparse_mat = self.ohe.fit_transform(sequence)
      return sparse_mat
    else:
      raise ValueError("Datapoint is not a string")

  def untransform(self, one_hot_vectors: scipy.sparse) -> str:
    """Convert from one hot representation back to original string

    Parameters
    ----------
    one_hot_vectors: np.ndarray
      An array of one hot encoded features.

    Returns
    -------
    str
      Original string for an one hot encoded array.
    """
    string = ""
    invers_trans = self.ohe.inverse_transform(one_hot_vectors)
    for one_hot in invers_trans:
      string += one_hot[0]
    return string

max_stars_repo_name: Chahalprincy/deepchem

max_issues_repo_name: Chahalprincy/deepchem

max_forks_repo_name: Chahalprincy/deepchem


text:
# Copyright (c) 2018-present, Facebook, Inc.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
#

import numpy as np
import copy,json
from common.skeleton import Skeleton
from common.mocap_dataset import MocapDataset
from common.camera import normalize_screen_coordinates, image_coordinates, world_to_camera
import time,os
import pickle

h36m_skeleton = Skeleton(parents=[-1,  0,  0,  1,  2,  6,  5,  5,  6,  7,  8,  5, 6, 11, 12, 13, 14, 15, 16, 15, 16, 0, 9, 10],
       joints_left=[1,3,5,7,9,11,13,15,17,19,22],
       joints_right=[2,4,6,8,10,12,14,16,18,20,23])

class Human36mDataset(MocapDataset):
    def __init__(self, path3d, path2d, remove_static_joints=True, chunk_length=1, augment=False, pad=121, causal_shift=0, is_train=False):
        super().__init__(fps=50, skeleton=copy.deepcopy(h36m_skeleton))
        self.chunk_length = chunk_length
        self.augment = augment
        self.pad = pad
        self.causal_shift = causal_shift
        self.is_train = is_train
        self.train_subjects = ["S1", "S2", "S5", "S6", "S7", "S8", "setting1", "setting101", "setting102", "setting2", "setting4","setting6", "setting71", "setting72", "setting81", "setting82", "setting91", "setting92"]
        self.test_subjects = ["S9", "S11", "setting3", "setting5"]
        cam_params = json.load(open("data/cam_all_z+.json"))
        self._cameras = copy.deepcopy(cam_params)
        for sub,cameras in self._cameras.items():
            for i, cam in enumerate(cameras):
                for k, v in cam.items():
                    if k not in ['id', 'res_w', 'res_h']:
                        cam[k] = np.array(v, dtype='float32')

                # Normalize camera frame
                # cam['center'] = normalize_screen_coordinates(cam['center'], w=cam['res_w'], h=cam['res_h']).astype('float32')
                # cam['focal_length'] = cam['focal_length'] / cam['res_w'] * 2
                if 'translation' in cam:
                    cam['translation'] = cam['translation'] / 1000  # mm to meters

                # Add intrinsic parameters vector
                cam['intrinsic'] = np.concatenate((cam['focal_length'], cam['center'], cam['radial_distortion'], cam['tangential_distortion']))

        # Load serialized dataset
        self._data_3d = np.load(path3d, allow_pickle=True)['positions_3d'].item()

        # if remove_static_joints:
        #     # Bring the skeleton to 17 joints instead of the original 32
        #     self.remove_joints([4, 5, 9, 10, 11, 16, 20, 21, 22, 23, 24, 28, 29, 30, 31])

        #     # Rewire shoulders to the correct parents
        #     self._skeleton._parents[11] = 8
        #     self._skeleton._parents[14] = 8

        data2d_file = np.load(path2d, allow_pickle=True)
        self._data_2d = data2d_file['positions_2d'].item()
        keypoints_symmetry = data2d_file['metadata'].item()['keypoints_symmetry']
        self.kps_left, self.kps_right = list(keypoints_symmetry[0]), list(keypoints_symmetry[1])

        if os.path.exists(path3d.replace(".npz",".pkl")):
            with open(path3d.replace(".npz",".pkl"), 'rb') as handle:
                self._data_3d = pickle.load(handle)
        else:
            self.process3DPose()
            with open(path3d.replace(".npz",".pkl"), 'wb') as handle:
                pickle.dump(self._data_3d,handle)

        self.normalize2DPose()
        self.getPairs()

        self.joints_left = self.skeleton().joints_left()
        self.joints_right = self.skeleton().joints_right()

    def supports_semi_supervised(self):
        return True

    def process3DPose(self):
        data_3d = self._data_3d
        self._data_3d = {}
        for subject, actions in data_3d.items():
            self._data_3d[subject] = {}
            for action_name, positions in actions.items():
                positions = positions[:, self.kept_joints]
                positions_3d = []
                for cam in self._cameras[subject]:
                    pos_3d = world_to_camera(positions, R=cam['orientation'], t=cam['translation'])
                    # pos_3d[:, 1:] -= pos_3d[:, :1]  # Remove global offset, but keep trajectory in first position
                    if len(pos_3d.shape) == 4:
                        pos_3d = pos_3d.reshape(pos_3d.shape[0] * pos_3d.shape[1],pos_3d.shape[2],pos_3d.shape[3])
                    pos_3d[:, :] -= np.mean(pos_3d[:, 11:13],axis = 1,keepdims=True)  # Remove global offset, but keep trajectory in first position
                    positions_3d.append(pos_3d)
                self._data_3d[subject][action_name] = {'positions': positions, 'cameras': self._cameras[subject], 'positions_3d': positions_3d}

    def normalize2DPose(self):
        for subject in self._data_3d.keys():
            assert subject in self._data_2d, 'Subject {} is missing from the 2D detections dataset'.format(subject)
            for action in self._data_3d[subject].keys():
                assert action in self._data_2d[subject], 'Action {} of subject {} is missing from the 2D detections dataset'.format(action, subject)
                if 'positions_3d' not in self._data_3d[subject][action]:
                    continue
                for cam_idx in range(len(self._data_2d[subject][action])):
                    # We check for >= instead of == because some videos in H3.6M contain extra frames
                    mocap_length = self._data_3d[subject][action]['positions_3d'][cam_idx].shape[0]
                    assert self._data_2d[subject][action][cam_idx].shape[0] >= mocap_length

                    if self._data_2d[subject][action][cam_idx].shape[0] > mocap_length:
                        # Shorten sequence
                        self._data_2d[subject][action][cam_idx] = self._data_2d[subject][action][cam_idx][:mocap_length]

                    cam = self._cameras[subject][cam_idx]
                    kps = self._data_2d[subject][action][cam_idx]
                    kps[..., :2] = normalize_screen_coordinates(kps[..., :2], cx=cam['center'][0], cy=cam['center'][1],fx = cam["focal_length"][0],fy = cam["focal_length"][1])
                    self._data_2d[subject][action][cam_idx] = kps
                assert len(self._data_2d[subject][action]) == len(self._data_3d[subject][action]['positions_3d'])

    def getPairs(self):
        if self.is_train:
            subjects = self.train_subjects
        else:
            subjects = self.test_subjects

        for subject in subjects:
            assert subject in self._data_2d, 'Subject {} is missing from the 2D detections dataset'.format(subject)
            for action in self._data_3d[subject].keys():
                pose3dlist = self._data_3d[subject][action]["positions_3d"]
                pose2dlist = self._data_2d[subject][action]
                assert len(pose3dlist) == len(pose2dlist)
                for i in range(len(pose3dlist)):
                    assert pose3dlist[i].shape[0] == pose2dlist[i].shape[0]
                    n_chunks = (pose2dlist[i].shape[0] + self.chunk_length - 1) // self.chunk_length
                    offset = (n_chunks * self.chunk_length - pose2dlist[i].shape[0]) // 2
                    bounds = np.arange(n_chunks + 1) * self.chunk_length - offset
                    augment_vector = np.full(len(bounds - 1), False, dtype=bool)
                    self._pairs += zip(["{:s}_{:s}_{:d}".format(subject, action, i)] * len(bounds - 1), bounds[:-1], bounds[1:], augment_vector)
                    if self.augment:
                        self._pairs += zip(["{:s}_{:s}_{:d}".format(subject, action, i)] * len(bounds - 1), bounds[:-1], bounds[1:], ~augment_vector)

    def getDataFromPair(self, index):

        subject_action_camIndex, start_3d, end_3d, flip = self._pairs[index]
        str_split_list = subject_action_camIndex.split("_")
        if len(str_split_list) == 3:
            subject, action, camIndex = str_split_list
        elif len(str_split_list) >= 4:
            subject = str_split_list[0]
            camIndex = str_split_list[-1]
            action = "_".join(str_split_list[1:-1])
        # print("subject:",subject," action:",action," camIndex:",camIndex)

        seq_2d = self._data_2d[subject][action][int(camIndex)]

        start_2d = start_3d - self.pad - self.causal_shift
        end_2d = end_3d + self.pad - self.causal_shift
        low_2d = max(start_2d, 0)
        high_2d = min(end_2d, seq_2d.shape[0])
        pad_left_2d = low_2d - start_2d
        pad_right_2d = end_2d - high_2d
        if pad_left_2d != 0 or pad_right_2d != 0:
            group_2d = np.pad(copy.deepcopy(seq_2d[low_2d:high_2d]), ((pad_left_2d, pad_right_2d), (0, 0), (0, 0)), 'edge')
        else:
            group_2d = copy.deepcopy(seq_2d[low_2d:high_2d])

        if flip:
            # Flip 2D keypoints
            group_2d[:, :, 0] *= -1
            group_2d[:, self.kps_left + self.kps_right] = group_2d[:, self.kps_right + self.kps_left]

        seq_3d = self._data_3d[subject][action]["positions_3d"][int(camIndex)]

        low_3d = max(start_3d, 0)
        high_3d = min(end_3d, seq_3d.shape[0])
        pad_left_3d = low_3d - start_3d
        pad_right_3d = end_3d - high_3d
        if pad_left_3d != 0 or pad_right_3d != 0:
            group_3d = np.pad(copy.deepcopy(seq_3d[low_3d:high_3d]), ((pad_left_3d, pad_right_3d), (0, 0), (0, 0)), 'edge')
        else:
            group_3d = copy.deepcopy(seq_3d[low_3d:high_3d])

        if flip:
            # Flip 3D joints
            group_3d[:, :, 0] *= -1
            group_3d[:, self.joints_left + self.joints_right] = group_3d[:, self.joints_right + self.joints_left]

        group_cam = self._cameras[subject][int(camIndex)]['intrinsic']

        if flip:
            # Flip horizontal distortion coefficients
            group_cam[2] *= -1
            group_cam[7] *= -1

        return subject_action_camIndex, group_cam, group_3d, group_2d


text:
import os
import pandas as pd
import numpy as np
import preprocessing as proc
import randomforest as rf
from pandas.api.types import CategoricalDtype
import seaborn as sns
from matplotlib import pyplot as plt
from scipy import stats
import joblib

from imblearn.over_sampling import SMOTE
from sklearn.model_selection import train_test_split, RandomizedSearchCV
from sklearn.preprocessing import StandardScaler, MaxAbsScaler, MinMaxScaler, Normalizer
from sklearn.decomposition import PCA
from sklearn.ensemble import RandomForestClassifier,RandomForestRegressor
from sklearn.cluster import KMeans, DBSCAN, AgglomerativeClustering
from sklearn.metrics import mean_squared_error, accuracy_score, f1_score, confusion_matrix

plt.style.use('seaborn-ticks')
sns.set_style('ticks')
plt.rcParams['figure.figsize'] = (8, 8)
plt.rcParams['axes.titlesize'] = 22
plt.rcParams['axes.labelsize'] = 22
plt.rcParams['xtick.labelsize'] = 20
plt.rcParams['ytick.labelsize'] = 20

pd.options.display.max_columns = 1000

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'

def plot_corr(correlations, mask=False, threshold=0.9, annot=False):
    """
    Plot a correlation matrix as a heatmap
    masking coefficients below the threshold.
    """
    if mask:
        mask = np.zeros_like(correlations)
        mask[abs(correlations) < threshold] = 1

    plt.rcParams['figure.figsize'] = (10, 10)

    sns.heatmap(correlations,
                mask=mask, annot=annot,
                vmin=-1, vmax=1,
                cmap=sns.color_palette("RdBu_r", 100));

def create_palette(df, by='stiffness', show=False):
    """
    Create a palette to make pretty plots.
    """
    if by == 'stiffness':
        n_levels = df.stiffness.unique().size
        palette = dict(zip(df.stiffness.unique(), sns.color_palette("Set3", n_levels)))
        row_colors = df.stiffness.map(palette)
    elif by == 'cluster':
        n_levels = df.cluster.unique().size
        if n_levels == 2:
            palette = {0 : sns.color_palette("PRGn", 20)[15], 1 : sns.color_palette("PRGn", 20)[4]}
        else:
            palette = dict(zip(df.cluster.unique(), sns.color_palette("Set2", n_levels)))
        row_colors = df.cluster.map(palette)
    elif by == 'biom':
        palette = {df.biom.unique()[0] : sns.color_palette("RdBu", 10)[1],
                   df.biom.unique()[1] : sns.color_palette("RdBu", 10)[8]}
        row_colors = df.biom.map(palette)
    elif by == 'comb':
        palette = {df.combination.unique()[0] : sns.color_palette("RdBu", 10)[1],
                   df.combination.unique()[1] : sns.color_palette("RdBu", 10)[8]}
        row_colors = df.combination.map(palette)
    elif by == 'isclumped':
        palette = {0 : sns.color_palette("Set3", 6)[4], 1 : sns.color_palette("Set3", 6)[5]}
        row_colors = df.isclumped.map(palette)

    if show:
#         print(list(palette.keys()))
        sns.palplot(palette.values());

    return row_colors, palette
# plt.savefig('../results/My palette.png', bbox_inches='tight', dpi=300);

def pca(df, cols):
    """
    Scale data and run PCA.
    Returns the first N components
    explaining 90% of variance.
    """
    # Feature map
    X = df[cols]

    # Standartise the data
    scaler = StandardScaler()
    X_scaled = scaler.fit_transform(X)

    # Number of components requred to describe 90% variance
    pca = PCA(n_components=0.9)

    # Apply PCA to the scaled feature matrix
    X_reduced = pca.fit_transform(X_scaled)

    pc_cols = [('pc_' + str(i)) for i in range(1, X_reduced.shape[1] + 1)]
    pc_df = pd.concat([df.loc[:, ['label', 'combination', 'stiffness', 'cluster']],
                      pd.DataFrame(data = X_reduced, columns = pc_cols)],
                      axis=1)

    print("The first {} components explain 90% of variance.\n".format(pca.n_components_ ))

    print("Explained variance: {}\n".format(pca.explained_variance_ratio_))

    print(X.shape, pca.components_.shape, X_reduced.shape, pc_df.shape)

    return pca, pc_cols, pc_df

def hc(df, cols, n_clusters=2, save=True):
    X = df[cols]
    scaler = StandardScaler()
    X_scaled = scaler.fit_transform(X)

    agg = AgglomerativeClustering(n_clusters=n_clusters)
    agg.fit(X_scaled)

#     1 is epit, 0 => 2 is mesench
#     if n_clusters == 2:
#         df['cluster'] = np.where(agg.labels_ == 1, 1, 2)
#     else:
#         df['cluster'] = agg.labels_
    df['cluster'] = agg.labels_ + 1

    cluster_type = CategoricalDtype(categories=list(range(1, n_clusters+1)), ordered=True)
    df.cluster = df.cluster.astype(cluster_type)

    row_colors, palette = create_palette(df, by='cluster')

    # Cluster map
    sns.clustermap(pd.DataFrame(X_scaled, columns=cols),
                   metric='euclidean', method='ward',
                   col_cluster=False,
                   cmap=sns.color_palette('RdBu_r', 100),
                   vmin=-2.5, vmax=2.5,
                   robust=True, row_colors=row_colors,
    #                figsize=(6.5, 10),
                   xticklabels=False, yticklabels=False);
    if save:
        plt.savefig('../results/Clustering geom_cols.png', bbox_inches='tight', dpi=300);

def ttest(df, by, col, equal_var, transform=None, verbose=True):
    """
    Run a t-testWelch t-test or a Welch t-test for unequal variances.
    """
    x_start = []
    x_end = []
    signif = []

    group = df[by].unique().sort_values()
    for i in range(len(group) - 1):
        x1 = df.loc[df[by] == group[i], col]
        x2 = df.loc[df[by] == group[i + 1], col]

        if transform == 'log':
            x1 = np.log(x1)
            x2 = np.log(x2)
        elif transform == 'boxcox':
            x1,_ = stats.boxcox(x1)
            x2,_ = stats.boxcox(x2)

        _, p = stats.ttest_ind(x1, x2, equal_var=equal_var, nan_policy='omit')
        if p < 0.05:
            x_start.append(i)
            x_end.append(i + 1)
            if p < 0.001:
                sign = '***'
            elif p < 0.01:
                sign = '**'
            elif p < 0.05:
                sign = '*'
            signif.append(sign)
        else:
            sign = ''

        if verbose:
            if equal_var:
                print("A two-tailed t-test on samples {} vs {}. {} \t p-value = {:.2}."
                      .format(group[i], group[i + 1],  sign, p))
            else:
                print("Welch's unequal variances t-test on samples {} vs {}. {} \t p-value = {:.2}."
                      .format(group[i], group[i + 1],  sign, p))

    return x_start, x_end, signif

def mannwhitneytest(df, by, col, alternative, transform=None):
    """
    Run a Mann-whitney U test.s
    """
    group = df[by].unique()
    for i in range(len(group) - 1):
        x1 = df.loc[df[by] == group[i], col]
        x2 = df.loc[df[by] == group[i + 1], col]

        if transform == 'log':
            x1 = np.log(x1)
            x2 = np.log(x2)
        elif transform == 'boxcox':
            x1,_ = stats.boxcox(x1)
            x2,_ = stats.boxcox(x2)

        _, p = stats.mannwhitneyu(x1, x2, alternative=alternative, nan_policy='omit')
        if p < 0.001:
            sign = '***'
        elif p < 0.01:
            sign = '**'
        elif p < 0.05:
            sign = '*'
        else:
            sign = ''

        print("A {} Mann Whitney test on samples {} vs {}. {} \t p-value = {:.2}."
              .format(alternative, group[i], group[i + 1],  sign, p))

def stat_annot(df, by, col, x_start, x_end, signif, ylim, kind='barplot'):
    """
    Add annotation to show the results of statistical testing.
    """
    s = df[by].unique().sort_values()
    h = (ylim[1] - ylim[0])/50

    for x1, x2, label in zip(x_start, x_end, signif):

        if kind == 'barplot':
            y = max(df.loc[df[by] == s[x1], col].mean() + df.loc[df[by] == s[x1], col].std(),
                    df.loc[df[by] == s[x2], col].mean() + df.loc[df[by] == s[x2], col].std()) + h
        elif kind == 'boxplot':
            y = max(upper_whisker(df.loc[df[by] == s[x1], col]),
                    upper_whisker(df.loc[df[by] == s[x2], col])) + h

        plt.plot([x1+0.05, x1+0.05, x2-0.05, x2-0.05], [y, y + h, y + h, y],
                 lw=1.5, color='k');

        plt.text((x1 + x2) * 0.5, y + h, s=label,
                 ha='center', va='bottom',
                 color='k', fontsize=20);
        plt.ylim(ylim)

def upper_whisker(x):
    iqr = x.quantile(0.75) - x.quantile(0.25)
    return x.quantile(0.75) + 1.5 * iqr

reason: import numpy,from scipy


text:
"""This module provides functions for editing traces as if they were tapes you can trim, reverse, etc."""
import numpy as np
from public import public
from typing import Union, Tuple, Any

from .trace import Trace

@public
def trim(trace: Trace, start: int = None, end: int = None) -> Trace:
    """
    Trim the `trace` samples, output contains samples between the `start` and `end` indices.

    :param trace:
    :param start:
    :param end:
    :return:
    """
    if start is None:
        start = 0
    if end is None:
        end = len(trace.samples)
    if start > end:
        raise ValueError("Invalid trim arguments.")
    return trace.with_samples(trace.samples[start:end].copy())

@public
def reverse(trace: Trace) -> Trace:
    """
    Reverse the samples of the `trace`.

    :param trace:
    :return:
    """
    return trace.with_samples(np.flipud(trace.samples))

@public
def pad(
    trace: Trace,
    lengths: Union[Tuple[int, int], int],
    values: Union[Tuple[Any, Any], Any] = (0, 0),
) -> Trace:
    """
    Pad the samples of the `trace` by `values` at the beginning and end.

    :param trace:
    :param lengths: How much to pad at the beginning and end, either symmetric (if integer) or asymmetric (if tuple).
    :param values: What value to pad with,  either symmetric or asymmetric (if tuple).
    :return:
    """
    if not isinstance(lengths, tuple):
        lengths = (lengths, lengths)
    if not isinstance(values, tuple):
        values = (values, values)
    return trace.with_samples(
        np.pad(trace.samples, lengths, "constant", constant_values=values)
    )


text:
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# 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 onnx_graphsurgeon as gs
import onnx
import numpy as np
import json

import sys, os
sys.path.insert(1, os.path.join(sys.path[0], ".."))
from downloader import getFilePath

def drop_category_mapper_nodes(graph):
    new_inputs = []
    for org_input in graph.inputs:
        # head node, simply disconnect it with others
        assert len(org_input.outputs) == 1
        category_mapper_node = org_input.outputs[0]
        assert category_mapper_node.op == 'CategoryMapper'
        assert len(category_mapper_node.outputs) == 1
        new_inputs.append(category_mapper_node.outputs[0])
        category_mapper_node.inputs.clear()
        category_mapper_node.outputs.clear()

        # Save mapping info to preprocess inputs.
        with open(category_mapper_node.name + '.json', 'w') as fp:
            json.dump(category_mapper_node.attrs, fp)

    graph.inputs = new_inputs

def replace_unsupported_ops(graph):
    # replace hardmax with ArgMax
    hardmaxes = [node for node in graph.nodes if node.op == "Hardmax"]
    assert len(hardmaxes) == 1
    hardmax = hardmaxes[0]
    hardmax.op = "ArgMax"
    hardmax.name = "ArgMax(org:" + hardmax.name + ")"
    hardmax.attrs["axis"] = 1
    hardmax.attrs["keepdims"] = 0

    cast = hardmax.o()
    reshape = cast.o()

    hardmax.outputs = reshape.outputs
    assert len(hardmax.outputs) == 1
    hardmax.outputs[0].dtype = np.int64
    hardmax.outputs[0].shape = [1]

    compress = reshape.o()
    compress.op = "Gather"
    compress.name = "Gather(org:" + compress.name + ")"
    compress.attrs["axis"] = 1

    cast.outputs.clear()
    reshape.outputs.clear()
    # Remove the node from the graph completely
    graph.cleanup().toposort()

def save_weights_for_refitting(graph):
    # Save weights for refitting
    tmap = graph.tensors()
    np.save("Parameter576_B_0.npy", tmap["Parameter576_B_0"].values)
    np.save("W_0.npy", tmap["W_0"].values)

def main():
    org_model_file_path = getFilePath('samples/python/engine_refit_onnx_bidaf/bidaf-original.onnx')

    print("Modifying the ONNX model ...")
    original_model = onnx.load(org_model_file_path)
    graph = gs.import_onnx(original_model)

    drop_category_mapper_nodes(graph)
    replace_unsupported_ops(graph)
    save_weights_for_refitting(graph)

    new_model = gs.export_onnx(graph)

    modified_model_name = "bidaf-modified.onnx"
    onnx.checker.check_model(new_model)
    onnx.save(new_model, modified_model_name)
    print("Modified ONNX model saved as {}".format(modified_model_name))
    print("Done.")

if __name__ == '__main__':
    main()