Calculate Log Likelihood with particular mean and std
        x must be 2D. [batch_size * eqsamples* iwsamples, num_latent]
556022
def neg_log_likelihood(self, x, eps=1e-10):
        
        x_reshape = tf.reshape(x, [self.batch_size, self.eq_samples, self.iw_samples, self.num_latent])

        # mean over the softmax outputs inside the log domain.
        pi = tf.reduce_mean(self.pi, axis=(1, 2))

        density = tf.reduce_mean(x * tf.log(pi + eps))
        return tf.reduce_sum(-density, axis=2)
54615
def _fit_diagV_noGP(
            self, XTY, XTDY, XTFY, YTY_diag, YTDY_diag, YTFY_diag,
            XTX, XTDX, XTFX, X, Y, X_base, X_res, D, F, run_TRs,
            current_vec_U_chlsk_l,
            current_a1, current_logSNR2,
            idx_param_fitU, idx_param_fitV,
            l_idx, n_C, n_T, n_V, n_l, n_run, n_X0, rank):
        
        init_iter = self.init_iter
        logger.info('second fitting without GP prior'
                    ' for {} times'.format(init_iter))

        # Initial parameters
        param0_fitU = np.empty(
            np.sum(np.size(v) for v in idx_param_fitU.values()))
        param0_fitV = np.empty(np.size(idx_param_fitV['log_SNR2']))
        # We cannot use the same logic as the line above because
        # idx_param_fitV also includes entries for GP parameters.
        param0_fitU[idx_param_fitU['Cholesky']] = \
            current_vec_U_chlsk_l.copy()
        param0_fitU[idx_param_fitU['a1']] = current_a1.copy()
        param0_fitV[idx_param_fitV['log_SNR2']] = \
            current_logSNR2[:-1].copy()

        L = np.zeros((n_C, rank))
        tol = self.tol * 5
        for it in range(0, init_iter):
            X0TX0, X0TDX0, X0TFX0, XTX0, XTDX0, XTFX0, \
                X0TY, X0TDY, X0TFY, X0, X_base, n_X0, _ \
                = self._prepare_data_XYX0(
                    X, Y, X_base, X_res, D, F, run_TRs, no_DC=True)

            # fit U, the covariance matrix, together with AR(1) param
            param0_fitU[idx_param_fitU['Cholesky']] = \
                current_vec_U_chlsk_l \
                + self.random_state_.randn(n_l) \
                * np.linalg.norm(current_vec_U_chlsk_l) \
                / n_l**0.5 * np.exp(-it / init_iter * self.anneal_speed - 1)
            param0_fitU[idx_param_fitU['a1']] = current_a1
            res_fitU = scipy.optimize.minimize(
                self._loglike_AR1_diagV_fitU, param0_fitU,
                args=(XTX, XTDX, XTFX, YTY_diag, YTDY_diag, YTFY_diag,
                      XTY, XTDY, XTFY, X0TX0, X0TDX0, X0TFX0,
                      XTX0, XTDX0, XTFX0, X0TY, X0TDY, X0TFY,
                      current_logSNR2, l_idx, n_C,
                      n_T, n_V, n_run, n_X0, idx_param_fitU, rank),
                method=self.optimizer, jac=True, tol=tol,
                options=self.minimize_options)
            current_vec_U_chlsk_l = \
                res_fitU.x[idx_param_fitU['Cholesky']]
            current_a1 = res_fitU.x[idx_param_fitU['a1']]
            norm_fitUchange = np.linalg.norm(res_fitU.x - param0_fitU)
            logger.debug('norm of parameter change after fitting U: '
                         '{}'.format(norm_fitUchange))
            param0_fitU = res_fitU.x.copy()

            # fit V, reflected in the log(SNR^2) of each voxel
            rho1 = np.arctan(current_a1) * 2 / np.pi
            L[l_idx] = current_vec_U_chlsk_l
            X0TAX0, XTAX0, X0TAY, X0TAX0_i, \
                XTAcorrX, XTAcorrY, YTAcorrY, \
                LTXTAcorrY, XTAcorrXL, LTXTAcorrXL = \
                self._precompute_ar1_quad_forms(XTY, XTDY, XTFY,
                                                YTY_diag, YTDY_diag, YTFY_diag,
                                                XTX, XTDX, XTFX,
                                                X0TX0, X0TDX0, X0TFX0,
                                                XTX0, XTDX0, XTFX0,
                                                X0TY, X0TDY, X0TFY,
                                                L, rho1, n_V, n_X0)
            res_fitV = scipy.optimize.minimize(
                self._loglike_AR1_diagV_fitV, param0_fitV,
                args=(X0TAX0, XTAX0, X0TAY,
                      X0TAX0_i, XTAcorrX, XTAcorrY, YTAcorrY,
                      LTXTAcorrY, XTAcorrXL, LTXTAcorrXL,
                      current_vec_U_chlsk_l,
                      current_a1, l_idx, n_C, n_T, n_V, n_run,
                      n_X0, idx_param_fitV, rank,
                      False, False),
                method=self.optimizer, jac=True, tol=tol,
                options=self.minimize_options)

            current_logSNR2[0:n_V - 1] = res_fitV.x
            current_logSNR2[-1] = - np.sum(current_logSNR2[0:n_V - 1])

            norm_fitVchange = np.linalg.norm(res_fitV.x - param0_fitV)
            logger.debug('norm of parameter change after fitting V: '
                         '{}'.format(norm_fitVchange))
            logger.debug('E[log(SNR2)^2]: {}'.format(
                np.mean(current_logSNR2**2)))

            # The lines below are for debugging purpose.
            # If any voxel's log(SNR^2) gets to non-finite number,
            # something might be wrong -- could be that the data has
            # nothing to do with the design matrix.
            if np.any(np.logical_not(np.isfinite(current_logSNR2))):
                logger.warning('Initial fitting: iteration {}'.format(it))
                logger.warning('current log(SNR^2): '
                               '{}'.format(current_logSNR2))
                logger.warning('log(sigma^2) has non-finite number')

            param0_fitV = res_fitV.x.copy()

            # Re-estimating X_res from residuals
            current_SNR2 = np.exp(current_logSNR2)
            if self.auto_nuisance:
                LL, LAMBDA_i, LAMBDA, YTAcorrXL_LAMBDA, current_sigma2 \
                    = self._calc_LL(rho1, LTXTAcorrXL, LTXTAcorrY, YTAcorrY,
                                    X0TAX0, current_SNR2,
                                    n_V, n_T, n_run, rank, n_X0)
                betas = current_SNR2 * np.dot(L, YTAcorrXL_LAMBDA.T)
                beta0s = np.einsum(
                    'ijk,ki->ji', X0TAX0_i,
                    (X0TAY - np.einsum('ikj,ki->ji', XTAX0, betas)))
                residuals = Y - np.dot(X, betas) - np.dot(
                    X_base, beta0s[:np.shape(X_base)[1], :])
                X_res = self.nureg_method(
                    self.n_nureg_).fit_transform(
                    self.preprocess_residual(residuals))
            if norm_fitVchange / np.sqrt(param0_fitV.size) < tol \
                    and norm_fitUchange / np.sqrt(param0_fitU.size) \
                    < tol:
                break
        return current_vec_U_chlsk_l, current_a1, current_logSNR2, X_res
658908
def _build_index_param(self, n_l, n_V, n_smooth):
        
        idx_param_sing = {'Cholesky': np.arange(n_l), 'a1': n_l}
        # for simplified fitting
        idx_param_fitU = {'Cholesky': np.arange(n_l),
                          'a1': np.arange(n_l, n_l + n_V)}
        # for the likelihood function when we fit U (the shared covariance).
        idx_param_fitV = {'log_SNR2': np.arange(n_V - 1),
                          'c_space': n_V - 1, 'c_inten': n_V,
                          'c_both': np.arange(n_V - 1, n_V - 1 + n_smooth)}
        # for the likelihood function when we fit V (reflected by SNR of
        # each voxel)
        return idx_param_sing, idx_param_fitU, idx_param_fitV
904929
def sample_stats_to_xarray(self):
        """"""
        dtypes = {"divergent__": bool, "n_leapfrog__": np.int64, "treedepth__": np.int64}

        # copy dims and coords
        dims = deepcopy(self.dims) if self.dims is not None else {}
        coords = deepcopy(self.coords) if self.coords is not None else {}

        sampler_params = self.sample_stats
        log_likelihood = self.log_likelihood
        if isinstance(log_likelihood, str):
            log_likelihood_cols = [
                col for col in self.posterior[0].columns if log_likelihood == col.split(".")[0]
            ]
            log_likelihood_vals = [item[log_likelihood_cols] for item in self.posterior]

            # Add log_likelihood to sampler_params
            for i, _ in enumerate(sampler_params):
                # slice log_likelihood to keep dimensions
                for col in log_likelihood_cols:
                    col_ll = col.replace(log_likelihood, "log_likelihood")
                    sampler_params[i][col_ll] = log_likelihood_vals[i][col]

            # change dims and coords for log_likelihood if defined
            if log_likelihood in dims:
                dims["log_likelihood"] = dims.pop(log_likelihood)

            log_likelihood_dims = np.array(
                [list(map(int, col.split(".")[1:])) for col in log_likelihood_cols]
            )
            max_dims = log_likelihood_dims.max(0)
            max_dims = max_dims if hasattr(max_dims, "__iter__") else (max_dims,)
            default_dim_names, _ = generate_dims_coords(shape=max_dims, var_name=log_likelihood)
            log_likelihood_dim_names, _ = generate_dims_coords(
                shape=max_dims, var_name="log_likelihood"
            )
            for default_dim_name, log_likelihood_dim_name in zip(
                default_dim_names, log_likelihood_dim_names
            ):
                if default_dim_name in coords:
                    coords[log_likelihood_dim_name] = coords.pop(default_dim_name)

        for j, s_params in enumerate(sampler_params):
            rename_dict = {}
            for key in s_params:
                key_, *end = key.split(".")
                name = re.sub("__$", "", key_)
                name = "diverging" if name == "divergent" else name
                rename_dict[key] = ".".join((name, *end))
                sampler_params[j][key] = s_params[key].astype(dtypes.get(key))
            sampler_params[j] = sampler_params[j].rename(columns=rename_dict)
        data = _unpack_dataframes(sampler_params)
        return dict_to_dataset(data, coords=coords, dims=dims)
427270
def _joint_likelihood(self,
                          logits: torch.Tensor,
                          tags: torch.Tensor,
                          mask: torch.LongTensor) -> torch.Tensor:
        
        batch_size, sequence_length, _ = logits.data.shape

        # Transpose batch size and sequence dimensions:
        logits = logits.transpose(0, 1).contiguous()
        mask = mask.float().transpose(0, 1).contiguous()
        tags = tags.transpose(0, 1).contiguous()

        # Start with the transition scores from start_tag to the first tag in each input
        if self.include_start_end_transitions:
            score = self.start_transitions.index_select(0, tags[0])
        else:
            score = 0.0

        # Add up the scores for the observed transitions and all the inputs but the last
        for i in range(sequence_length - 1):
            # Each is shape (batch_size,)
            current_tag, next_tag = tags[i], tags[i+1]

            # The scores for transitioning from current_tag to next_tag
            transition_score = self.transitions[current_tag.view(-1), next_tag.view(-1)]

            # The score for using current_tag
            emit_score = logits[i].gather(1, current_tag.view(batch_size, 1)).squeeze(1)

            # Include transition score if next element is unmasked,
            # input_score if this element is unmasked.
            score = score + transition_score * mask[i + 1] + emit_score * mask[i]

        # Transition from last state to "stop" state. To start with, we need to find the last tag
        # for each instance.
        last_tag_index = mask.sum(0).long() - 1
        last_tags = tags.gather(0, last_tag_index.view(1, batch_size)).squeeze(0)

        # Compute score of transitioning to `stop_tag` from each "last tag".
        if self.include_start_end_transitions:
            last_transition_score = self.end_transitions.index_select(0, last_tags)
        else:
            last_transition_score = 0.0

        # Add the last input if it's not masked.
        last_inputs = logits[-1]                                         # (batch_size, num_tags)
        last_input_score = last_inputs.gather(1, last_tags.view(-1, 1))  # (batch_size, 1)
        last_input_score = last_input_score.squeeze()                    # (batch_size,)

        score = score + last_transition_score + last_input_score * mask[-1]

        return score
851071
def data_log_likelihood(self, successes, trials, beta):
        ''''''
        return binom.logpmf(successes, trials, 1.0 / (1 + np.exp(-beta))).sum()
804654
def predict(self, Xnew, kern=None, pred_init=None, Y_metadata=None,
                median=False, deg_gauss_hermite=20, likelihood=None):
        
        #mu, var = GP._raw_predict(self, Xnew)
        # now push through likelihood
        #mean, var = self.likelihood.predictive_values(mu, var)
        
        mean, var = super(WarpedGP, self).predict(Xnew, kern=kern, full_cov=False, likelihood=likelihood)


        if self.predict_in_warped_space:
            std = np.sqrt(var)
            if median:
                wmean = self.warping_function.f_inv(mean, y=pred_init)
            else:
                wmean = self._get_warped_mean(mean, std, pred_init=pred_init,
                                              deg_gauss_hermite=deg_gauss_hermite).T
            wvar = self._get_warped_variance(mean, std, pred_init=pred_init,
                                             deg_gauss_hermite=deg_gauss_hermite).T
        else:
            wmean = mean
            wvar = var
        return wmean, wvar
917828
def inference_minibatch(self, kern, X, Z, likelihood, Y, Zp):
        

        num_data, output_dim = Y.shape

        #see whether we've got a different noise variance for each datum
        beta = 1./np.fmax(likelihood.variance, 1e-6)
        het_noise = beta.size > 1
        # VVT_factor is a matrix such that tdot(VVT_factor) = VVT...this is for efficiency!
        #self.YYTfactor = beta*self.get_YYTfactor(Y)
        if self.Y_speedup and not het_noise:
            YYT_factor = self.get_YYTfactor(Y)
        else:
            YYT_factor = Y

        n_start = self.batch_pos
        batchsize = num_data if self.batchsize is None else self.batchsize
        n_end = min(batchsize+n_start, num_data)
        if n_end == num_data:
            isEnd = True
            self.batch_pos = 0
        else:
            isEnd = False
            self.batch_pos = n_end

        if batchsize == num_data:
            Y_slice = YYT_factor
            X_slice = X
        else:
            Y_slice = YYT_factor[n_start:n_end]
            X_slice = X[n_start:n_end]

        psi0 = kern._Kdiag(X_slice) #Kffdiag
        psi1 = kern.K(X_slice, Z) #Kfu
        betapsi1 = np.einsum('n,nm->nm', beta, psi1)

        X_slice = X_slice.values
        Z = Z.values

        Zp = Zp.gamma.values
        indX = np.int_(X_slice[:, -1])
        indZ = np.int_(Z[:, -1]) - Zp.shape[0]

        betaY = beta*Y_slice

        #======================================================================
        # Load Intermediate Results
        #======================================================================

        dL_dpsi2R = self.midRes['dL_dpsi2R']
        v = self.midRes['v']

        #======================================================================
        # Compute dL_dpsi
        #======================================================================

        dL_dpsi0 = -.5*output_dim*(beta * Zp[indX, :]) #XxQ #TODO: Check this gradient

        dL_dpsi1 = np.dot(betaY, v.T)
        dL_dEZp  = psi1*dL_dpsi1
        dL_dpsi1 = Zp[np.ix_(indX, indZ)]*dL_dpsi1
        dL_dgamma = np.zeros(Zp.shape)
        for d in np.unique(indX):
            indd = indX == d
            betapsi1d = betapsi1[indd, :]
            psi1d = psi1[indd, :]
            Zpd = Zp[d, :]
            Zp2 = Zpd[:, None]*Zpd[None, :] - np.diag(np.power(Zpd, 2)) + np.diag(Zpd)
            dL_dpsi1[indd, :] += np.dot(betapsi1d, Zp2[np.ix_(indZ, indZ)] * dL_dpsi2R)*2.

            dL_EZp2 = dL_dpsi2R * (np.dot(psi1d.T, psi1d) * beta)*2.  # Zpd*Kufd*Kfud*beta
            #Gradient of Likelihood wrt gamma is calculated here
            EZ = Zp[d, indZ]
            for q in range(Zp.shape[1]):
                EZt = EZ.copy()
                indq = indZ == q
                EZt[indq] = .5
                dL_dgamma[d, q] = np.sum(dL_dEZp[np.ix_(indd, indq)]) + np.sum(dL_EZp2[:, indq]*EZt[:, None]) -\
                    .5*beta*(np.sum(psi0[indd, q]))

        #======================================================================
        # Compute dL_dthetaL
        #======================================================================
        if isEnd:
            dL_dthetaL = self.midRes['dL_dthetaL']
        else:
            dL_dthetaL = 0.

        grad_dict = {'dL_dKdiag': dL_dpsi0,
                     'dL_dKnm': dL_dpsi1,
                     'dL_dthetaL': dL_dthetaL,
                     'dL_dgamma': dL_dgamma}

        return isEnd, (n_start, n_end), grad_dict
655214
def estimate_column_means(
            self,
            X,
            observed,
            row_means,
            row_scales):
        
        n_rows, n_cols = X.shape
        row_means = np.asarray(row_means)

        if len(row_means) != n_rows:
            raise ValueError("Expected length %d but got shape %s" % (
                n_rows, row_means.shape))
        column_means = np.zeros(n_cols, dtype=X.dtype)

        X = X - row_means.reshape((n_rows, 1))
        row_weights = 1.0 / row_scales
        X *= row_weights.reshape((n_rows, 1))
        col_residual_sums = np.nansum(X, axis=0)
        for j in range(n_cols):
            col_mask = observed[:, j]
            sum_weights = row_weights[col_mask].sum()
            column_means[j] = col_residual_sums[j] / sum_weights
        return column_means
1161820
def convenience_calc_log_likelihood(self, params):
        
        orig_nest_coefs, betas = self.convenience_split_params(params)
        natural_nest_coefs = nc.naturalize_nest_coefs(orig_nest_coefs)

        args = [natural_nest_coefs,
                betas,
                self.design,
                self.rows_to_obs,
                self.rows_to_nests,
                self.choice_vector]
        kwargs = {"ridge": self.ridge, "weights": self.weights}

        log_likelihood = general_log_likelihood(*args, **kwargs)

        return log_likelihood
93237
def batchnorm_2d(nf:int, norm_type:NormType=NormType.Batch):
    ""
    bn = nn.BatchNorm2d(nf)
    with torch.no_grad():
        bn.bias.fill_(1e-3)
        bn.weight.fill_(0. if norm_type==NormType.BatchZero else 1.)
    return bn
562419
def posterior_to_xarray(self):
        """"""
        data = self.posterior
        if not isinstance(data, dict):
            raise TypeError("DictConverter.posterior is not a dictionary")

        if "log_likelihood" in data:
            warnings.warn(
                "log_likelihood found in posterior."
                " For stats functions log_likelihood needs to be in sample_stats.",
                SyntaxWarning,
            )

        return dict_to_dataset(data, library=None, coords=self.coords, dims=self.dims)
1197401
def neg_log_perplexity(batch, model_predictions):
  """"""
  _, targets = batch
  model_predictions, targets = _make_list(model_predictions, targets)
  xent = []
  for (prediction, target) in zip(model_predictions, targets):
    hot_target = layers.one_hot(target, prediction.shape[-1])
    xent.append(np.sum(prediction * hot_target, axis=-1))
  return masked_mean(xent, targets)
657029
def sample_stats_to_xarray(self):
        """"""
        posterior = self.posterior

        # copy dims and coords
        dims = deepcopy(self.dims) if self.dims is not None else {}
        coords = deepcopy(self.coords) if self.coords is not None else {}

        # log_likelihood
        log_likelihood = self.log_likelihood
        if log_likelihood is not None:
            if isinstance(log_likelihood, str) and log_likelihood in dims:
                dims["log_likelihood"] = dims.pop(log_likelihood)

        data = get_sample_stats(posterior, log_likelihood)

        return dict_to_dataset(data, library=self.pystan, coords=coords, dims=dims)
697893
def sample_W(self, nSamples, raw_samples=False):
        
        assert isinstance(self.kern, kern.Linear)
        from ..util.linalg import pdinv
        N, D = self.Y.shape
        Q = self.X.shape[1]
        noise_var = self.likelihood.variance.values
        
        # Draw samples for X
        Xs = np.random.randn(*((nSamples,)+self.X.shape))*np.sqrt(self.X.variance.values)+self.X.mean.values
        b = np.random.rand(*((nSamples,)+self.X.shape))
        Xs[b>self.X.gamma.values] = 0
        
        invcov = (Xs[:,:,:,None]*Xs[:,:,None,:]).sum(1)/noise_var+np.eye(Q)
        cov = np.array([pdinv(invcov[s_idx])[0] for s_idx in range(invcov.shape[0])])
        Ws = np.empty((nSamples, Q, D))
        tmp = (np.transpose(Xs, (0,2,1)).reshape(nSamples*Q,N).dot(self.Y)).reshape(nSamples,Q,D)
        mean = (cov[:,:,:,None]*tmp[:,None,:,:]).sum(2)/noise_var
        zeros = np.zeros((Q,))
        for s_idx in range(Xs.shape[0]):
            Ws[s_idx] = (np.random.multivariate_normal(mean=zeros,cov=cov[s_idx],size=(D,))).T+mean[s_idx]
        
        if raw_samples:
            return Ws
        else:
            return Ws.mean(0), Ws.std(0)
190907
def expected_log_likelihood(self, mus, sigmas, y):
        

        # Flatten the covariance
        T = mus.shape[0]
        D = self.D_in
        sigs_vec = sigmas.reshape((T, D ** 2))

        # Compute the log likelihood of each column
        ll = np.zeros((T, self.D_out))
        for n in range(self.D_out):

            an = self.A[n]

            E_loglmbda = np.dot(mus, an)
            ll[:,n] += y[:,n] * E_loglmbda

            # Vectorized log likelihood calculation
            aa_vec = np.outer(an, an).reshape((D ** 2,))
            ll[:,n] = -np.exp(E_loglmbda + 0.5 * np.dot(sigs_vec, aa_vec))

        return ll
1190905
def lnprior(pars):
    
    # Limit norm and B to be positive
    logprob = (
        naima.uniform_prior(pars[0], 0.0, np.inf)
        + naima.uniform_prior(pars[1], -1, 5)
        + naima.uniform_prior(pars[3], 0, np.inf)
    )

    return logprob