text:
# Load library dependencies
import numpy as np
import numpy.linalg

import scipy as sp
import scipy.optimize
import scipy.stats

import pandas as pd

# Import logit() command in ipt module since att() calls it
from .logit import logit

# Define att() function
#-----------------------------------------------------------------------------#

def att(D, Y, r_W, t_W, study_tilt=True, rlgrz=1, s_wgt=None, nocons=False, c_id=None, silent=False):

    """

    AUTHOR: Bryan S. Graham, UC - Berkeley, bgraham@econ.berkeley.edu
    DATE: Python 2.7 code on 26 May 2016, updated for Python 3.6 on 15 July 2018

    This function estimates the average treatment effect on the treated (ATT)
    using the "auxiliary-to-study tilting" (AST) method described by
    Graham, Pinto and Egel (2016, Journal of Business and Economic Statistics).
    The notation below mirrors that in the paper where possible. The Supplemental
    Web Appendix of the paper describes the estimation algorithm implemented here
    in detail. A copy of the paper and all supplemental appendices can be found
    online at http://bryangraham.github.io/econometrics/

    INPUTS
    ------
    D         : N x 1 pandas.Series with ith element equal to 1 if ith unit in the merged
                sample is from the study population and zero if from the auxiliary
                population (i.e., D is the "treatment" indicator)
    Y         : N x 1  pandas.Series of observed outcomes
    r_W       : r(W), N x 1+L pandas.DataFrame of functions of always observed covariates
                -- these are the propensity score basis functions
    t_W       : t(W), N x 1+M pandas.DataFrame of functions of always observed covariates
                -- these are the balancing functions
    study_tilt: If True compute the study sample tilt. This should be set to False
                if all the elements in t(W) are also contained in h(W). In that
                case the study_tilt coincides with its empirical measure.This
                measure is returned in the pi_s vector when  study_tilt = False.
    rlgrz     : Regularization parameter. Should positive and less than or equal
                to one. Smaller values correspond to less regularizations, but
                may cause underflow problems when overlap is poor. The default
                value will be adequate for most applications.
    s_wgt     : N x 1 pandas.Series of sampling weights variable (optional)
    nocons    : If True, then do NOT add constant to h_W and t_W matrix
                (only set to True if user passes in dataframes with constants included)
    c_id      : N X 1 pandas.Series of unique `cluster' id values (assumed to be integer valued) (optional)
                NOTE: Default is to assume independent observations and report heteroscedastic robust
                      standard errors
                NOTE: Data are assumed to be pre-sorted by groups.
    silent    : if silent = True display less optimization information and use
                lower tolerance levels (optional)

    OUTPUTS
    -------
    gamma_ast         : AST estimate of gamma (the ATT)
    vcov_gamma_ast    : estimated large sample variance of gamma
    pscore_tests      : list of [study_test, auxiliary_test] where
                        study_test     : ChiSq test statistic of H0 : lambda_s = 0; list with
                                         [statistic, dof, p-val]
                                         NOTE: returns [None, None, None] if study_tilt = False
                        auxiliary_test : ChiSq test statistic of H0 : lambda_a = 0; list with
                                         [statistic, dof, p-val]
    tilts             : numpy array with pi_eff, pi_s & pi_a as columns, sorted according
                        to the input data, and where
                        pi_eff : Semiparametrically efficient estimate of F_s(W)
                        pi_s   : Study sample tilt
                        pi_a   : Auxiliary sample tilt
    exitflag          : 1 = success, 2 = can't compute MLE of p-score, 3 = can't compute study/treated tilt,
                        4 = can't compute auxiliary/control tilt

    FUNCTIONS CALLED  : logit()                             (...logit_logl(), logit_score(), logit_hess()...)
    ----------------    ast_crit(), ast_foc(), ast_soc()    (...ast_phi()...)
    """

    def ast_phi(lmbda, t_W, p_W_index, NQ, rlgrz):

        """
        This function evaluates the regularized phi(v) function for
        the logit propensity score case (as well as its first and
        second derivatives) as described in the Supplemental
        Web Appendix of Graham, Pinto and Egel (2016, JBES).

        INPUTS
        ------
        lmbda         : vector of tilting parameters
        t_W           : vector of balancing moments
        p_W_index     : index of estimated logit propensity score
        NQ            : sample size times the marginal probability of missingness
        rlgrz         : Regularization parameter. See discussion in main header.

        OUTPUTS
        -------
        phi, phi1, phi2 : N x 1 vectors with elements phi(p_W_index + lmbda't_W)
                          and its first and second derivatives w.r.t to
                          v = p_W_index + lmbda't_W
        """

        # Adjust the NQ cut-off value used for quadratic extrapolation according
        # to the user-defined rlgrz parameter
        NQ =  NQ*rlgrz

        # Coefficients on quadratic extrapolation of phi(v) used to regularize
        # the problem
        c = -(NQ - 1)
        b = NQ + (NQ - 1)*np.log(1/(NQ - 1))
        a = -(NQ - 1)*(1 + np.log(1/(NQ - 1)) + 0.5*(np.log(1/(NQ - 1)))**2)

        v_star = np.log(1/(NQ - 1))

        # Evaluation of phi(v) and derivatives
        v          =  p_W_index + t_W @ lmbda
        phi        =  (v>v_star) * (v - np.exp(-v))   + (v<=v_star) * (a + b*v + 0.5*c*v**2)
        phi1       =  (v>v_star) * (1 + np.exp(-v))   + (v<=v_star) * (b + c*v)
        phi2       =  (v>v_star) * (  - np.exp(-v))   + (v<=v_star) * c

        return [phi, phi1, phi2]

    def ast_crit(lmbda, D, p_W, p_W_index, t_W, NQ, rlgrz, s_wgt):

        """
        This function constructs the AST criterion function
        as described in Graham, Pinto and Egel (2016, JBES).

        INPUTS
        ------
        lmbda         : vector of tilting parameters
        D             : N x 1 treatment indicator vector
        p_W           : N x 1 MLEs of the propensity score
        p_W_index     : index of estimated logit propensity score
        t_W           : vector of balancing moments
        NQ            : sample size times the marginal probability of missingness
        rlgrz         : Regularization parameter. See discussion in main header.
        s_wgt         : N x 1 vector of known sampling weights (optional)

        OUTPUTS
        -------
        crit          : AST criterion function at passed parameter values

        Functions called : ast_phi()
        """

        lmbda   = np.reshape(lmbda,(-1,1))                                # make lmda 2-dimensional object
        [phi, phi1, phi2] = ast_phi(lmbda, t_W, p_W_index, NQ, rlgrz)     # compute phi and 1st/2nd derivatives
        crit    = -np.sum(s_wgt * (D * phi - (t_W @ lmbda)) * (p_W / NQ)) # AST criterion (scalar)

        return crit

    def ast_foc(lmbda, D, p_W, p_W_index, t_W, NQ, rlgrz, s_wgt):

        """
        Returns first derivative vector of AST criterion function with respect
        to lmbda. See the header for ast_crit() for description of parameters.
        """

        lmbda   = np.reshape(lmbda,(-1,1))                              # make lmda 2-dimensional object
        [phi, phi1, phi2] = ast_phi(lmbda, t_W, p_W_index, NQ, rlgrz)   # compute phi and 1st/2nd derivatives
        foc     = -(t_W.T @ (s_wgt * (D * phi1 - 1) * (p_W / NQ)))      # AST gradient (1+M x 1 vector)
        foc     = np.ravel(foc)                                         # make foc 1-dimensional numpy array

        return foc

    def ast_soc(lmbda, D, p_W, p_W_index, t_W, NQ, rlgrz, s_wgt):

        """
        Returns hessian matrix of AST criterion function with respect
        to lmbda. See the header for ast_crit() for description of parameters.
        """

        lmbda   = np.reshape(lmbda,(-1,1))                            # make lmda 2-dimensional object
        [phi, phi1, phi2] = ast_phi(lmbda, t_W, p_W_index, NQ, rlgrz) # compute phi and 1st/2nd derivatives
        soc     = -(((s_wgt * D * phi2 * (p_W / NQ)) * t_W).T @ t_W)  # AST hessian (note use of numpy broadcasting rules)
                                                                      # (1 + M) x (1 + M) "matrix" (numpy array)
        return [soc]

    def ast_study_callback(lmbda):
        print("Value of ast_crit = "   + "%.6f" % ast_crit(lmbda, Ds, p_W, p_W_index, t_Ws, NQ, rlgrz, sw) + \
              ",  2-norm of ast_foc = "+ "%.6f" % numpy.linalg.norm(ast_foc(lmbda, Ds, p_W, p_W_index, t_Ws, \
                                                                            NQ, rlgrz, sw)))

    def ast_auxiliary_callback(lmbda):
        print("Value of ast_crit = "   + "%.6f" % ast_crit(lmbda, 1-Ds, p_W, -p_W_index, t_Ws, NQ, rlgrz, sw) + \
              ",  2-norm of ast_foc = "+ "%.6f" % numpy.linalg.norm(ast_foc(lmbda, 1-Ds, p_W, -p_W_index, t_Ws, \
                                                                            NQ, rlgrz, sw)))

    # ----------------------------------------------------------------------------------- #
    # - STEP 1 : ORGANIZE DATA                                                          - #
    # ----------------------------------------------------------------------------------- #

    # Extract variable names from pandas data objects
    dep_var     = Y.name                  # Get dependent variable names
    r_W_names   = list(r_W.columns)       # Get r_W variable names
    t_W_names   = list(t_W.columns)       # Get t_W variable names

    # Create pointers to pandas objects transformed into appropriately sized numpy arrays
    Ds         = D.values.reshape((-1,1)) # Turn pandas.Series into N x 1 numpy array
    Ys         = Y.values.reshape((-1,1)) # Turn pandas.Series into N x 1 numpy array
    r_Ws       = r_W.values               # Turn pandas.DataFrame into N x 1 + L numpy array
    t_Ws       = t_W.values               # Turn pandas.DataFrame into N x 1 + M numpy array

    # Extract basic information and set-up AST problem
    N         = len(D)                    # Number of units in sample
    Ns        = np.sum(D)                 # Number of study units in the sample (treated units)
    Na        = N-Ns                      # Number of auxiliary units in the sample (control units)
    M         = np.shape(t_Ws)[1]
    L         = np.shape(r_Ws)[1]

    if nocons:
        M = M - 1                         # Dimension of t_W (excluding constant)
        L = L - 1                         # Dimension of r_W (excluding constant)

    DY        = Ds * Ys                   # D*Y, N x 1  vector of observed outcomes for treated/study units
    mDX       = (1-Ds) * Ys               # (1-D)*X, N x 1  vector of observed outcomes for non-treated/auxiliary units

    # Add a constant to the regressor matrix (if needed)
    if not nocons:
        r_Ws      = np.concatenate((np.ones((N,1)), r_Ws), axis=1)
        r_W_names = ['constant'] + r_W_names
        t_Ws      = np.concatenate((np.ones((N,1)), t_Ws), axis=1)
        t_W_names = ['constant'] + t_W_names

    # Normalize weights to have mean one (if needed)
    if s_wgt is None:
        sw = 1
    else:
        s_wgt_var = s_wgt.name                 # Get sample weight variable name
        sw = np.asarray(s_wgt/s_wgt.mean()).reshape((-1,1))  # Normalized sampling weights with mean one

    # ----------------------------------------------------------------------------------- #
    # - STEP 2 : ESTIMATE PROPENSITY SCORE PARAMETER BY LOGIT ML                        - #
    # ----------------------------------------------------------------------------------- #
    try:
        if not silent:
            print("")
            print("--------------------------------------------------------------")
            print("- Computing propensity score by MLE                          -")
            print("--------------------------------------------------------------")

         # CMLE of p-score coefficients
        [delta_ml, vcov_delta_ml, hess_logl, score_i, p_W, _] = \
                                            logit(D, r_W, s_wgt=s_wgt, nocons=nocons, \
                                                  c_id=c_id, silent=silent, full=False)

        delta_ml                 = np.reshape(delta_ml,(-1,1))                 # Put delta_ml into 2-dimensional form
        p_W_index                = r_Ws @ delta_ml                             # Fitted p-score index
        NQ                       = np.sum(sw * p_W)                            # Sum of fitted p-scores
        pi_eff                   = (sw * p_W) / NQ                             # Efficient estimate of F(W)

    except:
        print("FATAL ERROR: exitflag = 2, unable to compute propensity score by maximum likelihood.")

        # Set all returnables to "None" and then exit function
        gamma_ast      = None
        vcov_gamma_ast = None
        pscore_tests   = None
        tilts          = None
        exitflag       = 2

        return [gamma_ast, vcov_gamma_ast, pscore_tests, tilts, exitflag]

    # ----------------------------------------------------------------------------------- #
    # - STEP 3 : SOLVE FOR AST TILTING PARAMETERS                                       - #
    # -----------------------------------------------------------------------------------

    # Set optimization parameters
    if silent:
        # Use Newton-CG solver with vector of zeros as starting values,
        # low tolerance levels, and smaller number of allowed iterations.
        # Hide iteration output.
        options_set = {'xtol': 1e-8, 'maxiter': 1000, 'disp': False}
    else:
        # Use Newton-CG solver with vector of zeros as starting values,
        # high tolerance levels, and larger number of allowed iterations.
        # Show iteration output.
        options_set = {'xtol': 1e-12, 'maxiter': 10000, 'disp': True}

    lambda_sv = np.zeros(1+M) # use vector of zeros as starting values

    #------------------------------#
    #- STUDY TILT                 -#
    #------------------------------#

    # NOTE: Only compute the study_tilt if directed to do so (this is the default). The study_tilt
    #       doesn't need to be computed if all the elements of t(W) are also included in h(W). It
    #       is the users responsibility to check this condition.

    if study_tilt:
        # -------------------------------------------------- #
        # - CASE 1: Non-trivial study sample tilt required - #
        # -------------------------------------------------- #

        # Compute lamba_s_hat (study or treated sample tilting parameters)
        try:
            if not silent:
                print("")
                print("--------------------------------------------------------------")
                print("- Computing study/treated sample tilt                        -")
                print("--------------------------------------------------------------")

                # Derivative check at starting values
                grad_norm = sp.optimize.check_grad(ast_crit, ast_foc, lambda_sv, Ds, p_W, \
                                                   p_W_index, t_Ws, NQ, rlgrz, \
                                                   sw, epsilon = 1e-12)
                print('Study sample tilt derivative check (2-norm): ' + "%.8f" % grad_norm)

                # Solve for tilting parameters
                lambda_s_res = sp.optimize.minimize(ast_crit, lambda_sv, args=(Ds, p_W, p_W_index, \
                                                                               t_Ws, NQ, rlgrz, sw), \
                                                    method='Newton-CG', jac=ast_foc, hess=ast_soc, \
                                                    callback = ast_study_callback, options=options_set)
            else:
                # Solve for tilting parameters
                lambda_s_res = sp.optimize.minimize(ast_crit, lambda_sv, args=(Ds, p_W, p_W_index, \
                                                                               t_Ws, NQ, rlgrz, sw), \
                                                    method='Newton-CG', jac=ast_foc, hess=ast_soc, \
                                                    options=options_set)
        except:
            print("FATAL ERROR: exitflag = 3, Unable to compute the study/treated vector of tilting parameters.")

            # Set all returnables to "None" and then exit function
            gamma_ast      = None
            vcov_gamma_ast = None
            pscore_tests   = None
            tilts          = None
            exitflag       = 3

            return [gamma_ast, vcov_gamma_ast, pscore_tests, tilts, exitflag]

        # Collect study tilt estimation results needed below
        lambda_s_hat = np.reshape(lambda_s_res.x,(-1,1))                           # study/treated sample tilting
                                                                                # parameter estimates
        p_W_s = (1+np.exp(-(p_W_index) - (t_Ws @ lambda_s_hat)))**-1               # study/treated sample tilted p-score
        pi_s  = Ds * pi_eff / p_W_s                                                # study/treated sample tilt

    else:
        # ------------------------------------------ #
        # - CASE 2: Study sample tilt NOT required - #
        # ------------------------------------------ #

        if not silent:
            print("")
            print("----------------------------------------------------------------------")
            print("- Tilt of study sample not requested by user (study_tilt = False).   -")
            print("- Validity of this requires all elements of t(W) to be elements of   -")
            print("- h(W) as well. User is advised to verify this condition.            -")
            print("----------------------------------------------------------------------")
            print("")

        # Collect study tilt objects needed below
        lambda_s_hat = np.reshape(lambda_sv ,(-1,1)) # study/treated sample tilting parameters set equal to zero
        p_W_s = p_W                                  # study/treated sample tilted p-score equals actual score
        pi_s  = Ds * pi_eff / p_W_s                  # set pi_s to "empirical measure" of study sub-sample
                                                     # (w/o sampling weights this puts mass 1/Ns on each study unit)

    #------------------------------#
    #- AUXILIARY TILT             -#
    #------------------------------#

    # Compute lamba_a_hat (auxiliary or control sample tilting parameters)
    try:
        if not silent:
            print("")
            print("--------------------------------------------------------------")
            print("- Computing auxiliary/control sample tilt                    -")
            print("--------------------------------------------------------------")

            # Derivative check at starting values
            grad_norm = sp.optimize.check_grad(ast_crit, ast_foc, lambda_sv, 1-Ds, p_W, \
                                               -p_W_index, t_Ws, NQ, rlgrz, \
                                               sw, epsilon = 1e-12)
            print('Auxiliary sample tilt derivative check (2-norm): ' + "%.8f" % grad_norm)

            # Solve for tilting parameters
            lambda_a_res = sp.optimize.minimize(ast_crit, lambda_sv, args=(1-Ds, p_W, -p_W_index, \
                                                                           t_Ws, NQ, rlgrz, sw), \
                                                method='Newton-CG', jac=ast_foc, hess=ast_soc, \
                                                callback = ast_auxiliary_callback, options=options_set)
        else:
            # Solve for tilting parameters
            lambda_a_res = sp.optimize.minimize(ast_crit, lambda_sv, args=(1-Ds, p_W, -p_W_index, \
                                                                           t_Ws, NQ, rlgrz, sw), \
                                                method='Newton-CG', jac=ast_foc, hess=ast_soc, \
                                                options=options_set)
    except:
        print("FATAL ERROR: exitflag = 4, Unable to compute the auxiliary/control vector of tilting parameters.")

        # Set returnables to "None" and then exit function
        gamma_ast      = None
        vcov_gamma_ast = None
        pscore_tests   = None
        tilts          = None
        exitflag       = 4

        return [gamma_ast, vcov_gamma_ast, pscore_tests, tilts, exitflag]

    # Collect auxiliary tilt estimation results needed below
    lambda_a_hat = -np.reshape(lambda_a_res.x,(-1,1))                          # auxiliary/control sample tilting
                                                                               # parameter estimates
    p_W_a = (1+np.exp(-(p_W_index) - (t_Ws @ lambda_a_hat)))**-1               # auxiliary sample tilted p-score
    pi_a  = (1-Ds) * (pi_eff / (1-p_W_a))                                      # auxiliary sample tilt

    # ----------------------------------------------------------------------------------- #
    # - STEP 4 : SOLVE FOR AST ESTIMATE OF GAMMA (i.e., ATT)                            - #
    # ----------------------------------------------------------------------------------- #

    # AST estimate of gamma -- the ATT %
    gamma_ast = np.sum(sw * p_W * ((Ds / p_W_s) * DY - (1-Ds) / (1-p_W_a) * mDX))/NQ;

    # ----------------------------------------------------------------------------------- #
    # - STEP 5 : FORM LARGE SAMPLE VARIANCE-COVARIANCE ESTIMATES                        - #
    # ----------------------------------------------------------------------------------- #

    # Form moment vector corresponding to full three step procedure
    m1 = (sw * (Ds - p_W) * r_Ws).T                          # 1+L x N matrix of m_1 moments (logit scores)
    m2 = (sw * ((1 - Ds) / (1 - p_W_a) - 1) * p_W * t_Ws).T  # 1+M x N matrix of m_2 moments
    m3 = (sw * (Ds / p_W_s - 1) * p_W * t_Ws).T              # 1+M x N matrix of m_3 moments
    m4 = (sw * p_W * ((Ds / p_W_s) * DY - ((1-Ds) / (1-p_W_a)) * (mDX+gamma_ast))).T  # 1 x N matrix of m_4 moments
    m  = np.concatenate((m1, m2, m3, m4), axis=0)            # 1 + L + 2(1 + M) + 1 x N matrix of all moments

    # Calculate covariance matrix of moment vector. Take into account any
    # within-group dependence/clustering as needed
    if c_id is None:

        # Case 1: No cluster dependence to account for when constructing covariance matrix
        C   = N                                         # Number of clusters equals number of observations
        fsc = N/(N - (1+L+2*(1+M)+1))                   # Finite-sample correction factor
        V_m = fsc*(m @ m.T)/N

    else:

        # Case 2: Need to correct for cluster dependence when constructing covariance matrix

        # Get number and unique list of clusters
        c_list  = np.unique(c_id)
        C       = len(c_list)

        # Calculate cluster-robust variance-covariance matrix of m
        # Sum moment vector within clusters
        sum_m   = np.empty((C,1+L+2*(1+M)+1))           # initiate vector of cluster-summed moments

        for c in range(0,C):

            # sum of moments for units in c-th cluster
            b_cluster    = np.nonzero((c_id == c_list[c]))[0]                             # Observations in c-th cluster
            sum_m[c,:]   = np.sum(m[np.ix_(range(0,1+L+2*(1+M)+1), b_cluster)], axis = 1) # Sum over "columns" within c-th cluster

        # Compute variance-covariance matrix of moment vector
        fsc = (N/(N - (1+L+2*(1+M)+1)))*(C/(C-1))     # Finite-sample correction factor
        V_m = fsc*(sum_m.T @ sum_m)/C                 # Variance-covariance of the summed moments

    # Form Jacobian matrix for entire parameter: theta = (rho, delta, lambda, gamma)
    e_V  = np.exp(np.dot(r_Ws, delta_ml))
    e_Va = np.exp(np.dot(r_Ws, delta_ml) + np.dot(t_Ws, lambda_a_hat))
    e_Vs = np.exp(np.dot(r_Ws, delta_ml) + np.dot(t_Ws, lambda_s_hat))

    M1_delta = np.dot((sw * (- e_V / (1 + e_V)**2) * r_Ws).T, r_Ws)/N                                                    # 1 + L x 1 + L
    M2_delta = np.dot((sw * ((1 - Ds) / (1 - p_W_a) - 1) * (e_V / (1 + e_V)**2) * t_Ws).T, r_Ws)/N                       # 1 + M x 1 + L
    M3_delta = np.dot((sw * (Ds / p_W_s - 1) * (e_V / (1 + e_V)**2) * t_Ws).T, r_Ws)/N                                   # 1 + M x 1 + L
    M4_delta = np.dot((sw * (e_V / (1 + e_V)**2) * \
                      ((Ds / p_W_s) * DY - ((1 - Ds) / (1 - p_W_a)) * (mDX + gamma_ast))).T, r_Ws)/N                     # 1     x 1 + L

    M2_lambda_a = np.dot(( sw * ((1 - Ds) / (1 - p_W_a)**2) * p_W * (e_Va / (1 + e_Va)**2) * t_Ws).T, t_Ws)/N            # 1 + M x 1 + M
    M4_lambda_a = np.dot((-sw * ((1 - Ds) / (1 - p_W_a)**2) * p_W * (mDX+gamma_ast) * (e_Va / (1 + e_Va)**2)).T, t_Ws)/N # 1     x 1 + M

    M3_lambda_s = np.dot((-sw * (Ds / p_W_s**2) * p_W * (e_Vs / (1 + e_Vs)**2) * t_Ws).T, t_Ws)/N                        # 1 + M x 1 + M
    M4_lambda_s = np.dot((-sw * (Ds / p_W_s**2) * p_W * DY * (e_Vs / (1 + e_Vs)**2)).T, t_Ws)/N                          # 1     x 1 + M

    M4_gamma    = -(NQ/N).reshape(1,1)                                                                                # 1     x 1

    M1 = np.hstack((M1_delta, np.zeros((1+L,1+M)), np.zeros((1+L,1+M)), np.zeros((1+L,1))))
    M2 = np.hstack((M2_delta, M2_lambda_a,         np.zeros((1+M,1+M)), np.zeros((1+M,1))))
    M3 = np.hstack((M3_delta, np.zeros((1+M,1+M)), M3_lambda_s,         np.zeros((1+M,1))))
    M4 = np.hstack((M4_delta, M4_lambda_a,         M4_lambda_s,         M4_gamma))

    # Concatenate Jacobian and compute inverse
    M_hat  = (N/C)*np.vstack((M1, M2, M3, M4))
    iM_hat = np.linalg.inv(M_hat)

    # Compute sandwich variance estimates
    vcov_theta_ast  = (iM_hat @ V_m @ iM_hat.T)/C
    vcov_gamma_ast  = vcov_theta_ast[-1,-1]

    exitflag = 1 # AST estimate of the ATT successfully computed!

    # ----------------------------------------------------------------------------------- #
    # - STEP 6 : COMPUTE TEST STATISTICS BASED ON TILTING PARAMETER                     - #
    # ----------------------------------------------------------------------------------- #

    # Compute propensity score specification test based on study tilt (if applicable)
    if study_tilt:
        iV_lambda_s = np.linalg.inv(vcov_theta_ast[1+L:1+L+1+M,1+L:1+L+1+M])
        ps_test_st  = np.dot(np.dot(lambda_s_hat.T, iV_lambda_s), lambda_s_hat)
        dof_st      = len(lambda_s_hat)
        pval_st     = 1 - sp.stats.chi2.cdf(ps_test_st[0,0], dof_st)
        study_test  = [ps_test_st[0,0], dof_st, pval_st]
    else:
        study_test  = [None, None, None]

    # Compute propensity score specification test based on auxiliary tilt (always done)
    iV_lambda_a    = np.linalg.inv(vcov_theta_ast[1+L+1+M:1+L+1+M+1+M,1+L+1+M:1+L+1+M+1+M])
    ps_test_at     = np.dot(np.dot(lambda_a_hat.T, iV_lambda_a), lambda_a_hat)
    dof_at         = len(lambda_a_hat)
    pval_at        = 1 - sp.stats.chi2.cdf(ps_test_at[0,0], dof_at)
    auxiliary_test = [ps_test_at[0,0], dof_at, pval_at]

    # ----------------------------------------------------------------------------------- #
    # - STEP 7 : DISPLAY RESULTS                                                        - #
    # ----------------------------------------------------------------------------------- #

    if not silent:
        print("")
        print("-------------------------------------------------------------------------------------------")
        print("- Auxiliary-to-Study (AST) estimates of the ATT                                           -")
        print("-------------------------------------------------------------------------------------------")
        print("ATT:  " + "%10.6f" % gamma_ast)
        print("     (" + "%10.6f" % np.sqrt(vcov_gamma_ast) + ")")
        print("")
        print("-------------------------------------------------------------------------------------------")
        if c_id is None:
            print("NOTE: Outcome variable   = " + dep_var)
            print("      Heteroscedastic-robust standard errors reported")
            print("      N1 = " "%0.0f" % Ns + ", N0 = " + "%0.0f" % Na)
        else:
            print("NOTE: Outcome variable   = " + dep_var)
            print("      Cluster-robust standard errors reported")
            print("      Cluster-variable   = " + c_id.name)
            print("      Number of clusters = " + "%0.0f" % C)
            print("      N1 = " "%0.0f" % Ns + ", N0 = " + "%0.0f" % Na)
        if s_wgt is not None:
            print("NOTE: (Sampling) Weighted AST estimates computed.")
            print("      Weight-variable   = " + s_wgt_var)

        print("")
        print("-------------------------------------------------------------------------------------------")
        print("- Maximum likelihood estimates of the p-score                                             -")
        print("-------------------------------------------------------------------------------------------")
        print("")
        print("Independent variable       Coef.    ( Std. Err.) ")
        print("-------------------------------------------------------------------------------------------")

        c = 0
        for names in r_W_names:
            print(names.ljust(25) + "%10.6f" % delta_ml[c] + \
                             " (" + "%10.6f" % np.sqrt(vcov_theta_ast[c,c]) + ")")
            c += 1

        print("")
        print("-------------------------------------------------------------------------------------------")
        print("- Tilting parameter estimates                                                             -")
        print("-------------------------------------------------------------------------------------------")

        if study_tilt:
            print("")
            print("TREATED (study) sample tilt")
            print("-------------------------------------------------------------------------------------------")
            print("")
            print("Independent variable       Coef.    ( Std. Err.) ")
            print("-------------------------------------------------------------------------------------------")

            c = 0
            for names in t_W_names:
                print(names.ljust(25) + "%10.6f" % lambda_s_hat[c] + \
                                 " (" + "%10.6f" % np.sqrt(vcov_theta_ast[1+L+c,1+L+c]) + ")")
                c += 1

            print("")
            print("Specification test for p-score (H_0 : lambda_s = 0)")
            print("-------------------------------------------------------------------------------------------")
            print("chi-square("+str(dof_st)+") = " + "%10.6f" % ps_test_st + "   p-value: " + "% .6f" % pval_st)

            print("")
            print("Summary statistics study/treated re-weighting")
            print("-------------------------------------------------------------------------------------------")

            j          = np.where(D)[0]        # find indices of treated units
            N_s_eff    = 1/np.sum(pi_s[j]**2)  # Kish's formula for effective sample size
            print("Kish's effective study/treated sample size = " "%0.0f" % N_s_eff)
            print("")

            print("Percentiles of N_s * pi_s distribution")
            quantiles  = [1, 5, 10, 25, 50, 75, 90, 95, 99]
            qnt_pi_s   = np.percentile(Ns*pi_s[j],quantiles)

            c = 0
            for q in quantiles:
                print("%2.0f" % quantiles[c] + " percentile = " "%2.4f" % qnt_pi_s[c])
                c += 1

        else:
            print("")
            print("--------------------------------------------------------")
            print("- NOTE: Study tilt not computed (study_tilt = False).  -")
            print("-       Components of t(W) assumed to be also in h(W). -")
            print("--------------------------------------------------------")
            print("")

        print("")
        print("CONTROL (auxiliary) sample tilt")
        print("-------------------------------------------------------------------------------------------")
        print("")
        print("Independent variable       Coef.    ( Std. Err.) ")
        print("-------------------------------------------------------------------------------------------")

        c = 0
        for names in t_W_names:
            print(names.ljust(25) + "%10.6f" % lambda_a_hat[c] + \
                             " (" + "%10.6f" % np.sqrt(vcov_theta_ast[1+L+1+M+c,1+L+1+M+c]) + ")")
            c += 1

        print("")
        print("Specification test for p-score (H_0 : lambda_a = 0)")
        print("-------------------------------------------------------------------------------------------")
        print("chi-square("+str(dof_at)+") = " + "%10.6f" % ps_test_at + "   p-value: " + "% .6f" % pval_at)

        print("")
        print("Summary statistics auxiliary/control re-weighting")
        print("-------------------------------------------------------------------------------------------")

        j          = np.where(1-D)[0]      # find indices of control units
        N_a_eff    = 1/np.sum(pi_a[j]**2)  # Kish's formula for effective sample size
        print("Kish's effective auxiliary/control sample size = " "%0.0f" % N_a_eff)
        print("")

        print("Percentiles of N_a * pi_a distribution")
        quantiles  = [1, 5, 10, 25, 50, 75, 90, 95, 99]
        qnt_pi_a   = np.percentile(Na*pi_a[j],quantiles)

        c = 0
        for q in quantiles:
            print("%2.0f" % quantiles[c] + " percentile = " "%2.4f" % qnt_pi_a[c])
            c += 1

        # ------------------------------------------- #
        # Construct "exact balancing" table         - #
        # ------------------------------------------- #

        Na_wgt = np.sum(sw * (1-Ds) , axis = 0)
        Ns_wgt = np.sum(sw * Ds , axis = 0)

        # Compute means of t_W across various distribution function estimates
        # Mean of t(W) across controls
        mu_t_D0      = np.sum(sw * (1-Ds) * t_Ws, axis = 0)/Na_wgt
        mu_t_D0_std  = np.sqrt(np.sum(sw * (1-Ds) * (t_Ws - mu_t_D0)**2, axis = 0)/Na_wgt)

        # Mean of t(W) across treated
        mu_t_D1      = np.sum(sw * Ds * t_Ws, axis = 0)/Ns_wgt
        mu_t_D1_std  = np.sqrt(np.sum(sw * Ds * (t_Ws - mu_t_D1)**2, axis = 0)/Ns_wgt)

        # Normalized mean differences across treatment and controls
        # (cf., Imbens, 2015, Journal of Human Resources)
        NormDif_t    = (mu_t_D1 - mu_t_D0)/np.sqrt((mu_t_D1_std**2 + mu_t_D0_std**2)/2)

        # Semiparametrically efficient estimate of mean of t(W) across treated
        mu_t_eff     = np.sum(pi_eff * t_Ws, axis = 0)
        mu_t_eff_std = np.sqrt(np.sum(pi_eff * (t_Ws - mu_t_eff)**2, axis = 0))

        # Mean of t(W) across controls after re-weighting
        mu_t_a     = np.sum(pi_a * t_Ws, axis = 0)
        mu_t_a_std = np.sqrt(np.sum(pi_a * (t_Ws - mu_t_a)**2, axis = 0))

        # Mean of t(W) across treated after re-weighting
        mu_t_s     = np.sum(pi_s * t_Ws, axis = 0)
        mu_t_s_std = np.sqrt(np.sum(pi_s * (t_Ws - mu_t_s)**2, axis = 0))

        # Pre-balance table
        print("")
        print("Means & standard deviations of t_W (pre-balance)                                           ")
        print("-------------------------------------------------------------------------------------------")
        print("                            Treated (D = 1)        Control (D = 0)        Norm. Diff.      ")
        print("-------------------------------------------------------------------------------------------")
        c = 0
        for names in t_W_names:
            print(names.ljust(25) + "%8.4f" % mu_t_D1[c]  + " (" + "%8.4f" % mu_t_D1_std[c] + ")    " \
                                  + "%8.4f" % mu_t_D0[c]  + " (" + "%8.4f" % mu_t_D0_std[c] + ")    " \
                                  + "%8.4f" % NormDif_t[c])
            c += 1

        # Post-balance table
        print("")
        print("Means and standard deviations of t_W (post-balance)                                        ")
        print("-------------------------------------------------------------------------------------------")
        print("                            Treated (D = 1)        Control (D = 0)        Efficient (D = 1)")
        print("-------------------------------------------------------------------------------------------")
        c = 0
        for names in t_W_names:
            print(names.ljust(25) + "%8.4f" % mu_t_s[c]    + " (" + "%8.4f" % mu_t_s_std[c]   + ")    " \
                                  + "%8.4f" % mu_t_a[c]    + " (" + "%8.4f" % mu_t_a_std[c]   + ")    " \
                                  + "%8.4f" % mu_t_eff[c]  + " (" + "%8.4f" % mu_t_eff_std[c] + ")    ")
            c += 1

    # Collect/format remaining returnables and exit function
    pscore_tests = [study_test, auxiliary_test]                 # Collect p-score test results
    tilts        = np.concatenate((pi_eff, pi_s, pi_a), axis=1) # Collect three sample tilts

    return [gamma_ast, vcov_gamma_ast, pscore_tests, tilts, exitflag]

max_stars_repo_name: ida-j/endogenous-peer-effects

max_issues_repo_name: ida-j/endogenous-peer-effects

max_forks_repo_name: ida-j/endogenous-peer-effects

reason: import numpy,import scipy


text:
import cv2
import os
import numpy as np

os.chdir(os.path.dirname(os.path.abspath(__file__)))
angle = lambda x,y: (np.arctan2(y,-1*x)*180/np.pi)%360
imgname = 'resized_game.png'

img = cv2.imread(imgname)
print('Shape', img.shape)
imgcopy = img
coords = []
end = []
RED  = (0, 0, 255)
BLUE = (255,0,0 )
"""These should be in inches"""
real_robot_side = 33.5
real_robot_len = 38.5

robot_side = real_robot_side*2
robot_len = real_robot_len*2

base_angle = np.arctan((robot_len/2)/(robot_side/2))*180/np.pi
target_radius = np.sqrt((robot_side/2)**2 + (robot_len/2)**2)

def lineDrawer(img, x):
    if(len(x) == 1):
        return
    start = (x[-2][0], x[-2][1])
    end = (x[-1][0], x[-1][1])
    cv2.line(img, start, end, RED, 1)
    x.pop(-1)
    lineDrawer(img, x)

def recurseRobotInstructions(coords, robotcoords = [], robot_scaledcoords = []):
    if(len(coords) == 1):
        return robotcoords, robot_scaledcoords
    start = coords[0]
    end = coords[1]
    dx = end[0] - start[0]
    dy = end[1] - start[1]
    pot_angle = angle(dy,dx)
    distance = np.sqrt((dx)**2 + (dy)**2)
    robotcoords.append('T%s' %(round(pot_angle,2)))
    robotcoords.append('F%s' %(round(distance,2)))

    sdx = dx/2
    sdy = dy/2
    sdistance = np.sqrt((sdx)**2 + (sdy)**2)
    robot_scaledcoords.append('T%s' %(round(pot_angle,2)))
    robot_scaledcoords.append('F%s' %(round(sdistance,2)))
    coords.pop(0)
    return recurseRobotInstructions(coords, robotcoords)

def eventHandler(event, x, y, k, s):
    global imgcopy, img, coords
    imgcopy = img.copy()
    if(event == cv2.EVENT_LBUTTONDOWN):
        if(len(coords) == 0):
            coords.append([x,y])
            coords.append([x,y])
        else:
            coords.append([x,y])
    elif(event == cv2.EVENT_MOUSEMOVE):
        if(len(coords) >= 2):
            coords[-1] = [x,y]
        else:
            pass
    else:
        pass

    x = coords.copy()
    while(len(x) > 1):
        start = (x[-2][0], x[-2][1])
        end = (x[-1][0], x[-1][1])
        cv2.line(imgcopy, start, start, BLUE, 5)
        cv2.line(imgcopy, end, end, BLUE, 5)
        dx = end[0] - start[0]
        dy = end[1] - start[1]
        pot_angle = angle(dy,dx)
        angle1 = base_angle - pot_angle
        angle2 = base_angle + pot_angle
        start_bottom = (start[0] - int(target_radius*np.cos(np.pi*angle1/180)), start[1] + int(target_radius*np.sin(np.pi*angle1/180)))
        start_top = (end[0]  - int(target_radius*np.cos(np.pi*angle2/180)), end[1] - int(target_radius*np.sin(np.pi*angle2/180)))

        end_bottom = (start[0]  + int(target_radius*np.cos(np.pi*angle2/180)), start[1] + int(target_radius*np.sin(np.pi*angle2/180)))
        end_top = (end[0] + int(target_radius*np.cos(np.pi*angle1/180)), end[1]- int(target_radius*np.sin(np.pi*angle1/180)))

        cv2.line(imgcopy, start_bottom, start_top ,BLUE, 1 )
        cv2.line(imgcopy, end_bottom, end_top ,BLUE, 1 )
        cv2.line(imgcopy, start_bottom, end_bottom ,BLUE, 1 )
        cv2.line(imgcopy, start_top, end_top ,BLUE, 1 )
        cv2.line(imgcopy, start, end, RED, 1)
        x.pop(-1)

cv2.namedWindow("img")
cv2.setMouseCallback("img", eventHandler)

while True:
    cv2.imshow("img", imgcopy)
    key = cv2.waitKey(1) & 0xFF
    # if the 'q' key is pressed, stop the loop
    if(key == ord("y")):
        print(coords)
        print(recurseRobotInstructions(coords[0:-1]))
        break
    if key == ord("q"):
        break
cv2.destroyAllWindows()

ext: py

reason: import numpy


text:
# -*- coding: utf-8 -*-
# emacs: -*- mode: python; py-indent-offset: 4; indent-tabs-mode: nil -*-
# vi: set ft=python sts=4 ts=4 sw=4 et:
"""
Miscellaneous algorithms for 2D contours and 3D triangularized meshes handling
"""
import os.path as op
import numpy as np
from numpy import linalg as nla

from .. import logging
from ..interfaces.base import (
    BaseInterface,
    traits,
    TraitedSpec,
    File,
    BaseInterfaceInputSpec,
)
from ..interfaces.vtkbase import tvtk
from ..interfaces import vtkbase as VTKInfo

IFLOGGER = logging.getLogger("nipype.interface")

class TVTKBaseInterface(BaseInterface):
    """A base class for interfaces using VTK"""

    _redirect_x = True

    def __init__(self, **inputs):
        if VTKInfo.no_tvtk():
            raise ImportError("This interface requires tvtk to run.")
        super(TVTKBaseInterface, self).__init__(**inputs)

class WarpPointsInputSpec(BaseInterfaceInputSpec):
    points = File(exists=True, mandatory=True, desc="file containing the point set")
    warp = File(
        exists=True, mandatory=True, desc="dense deformation field to be applied"
    )
    interp = traits.Enum(
        "cubic",
        "nearest",
        "linear",
        usedefault=True,
        mandatory=True,
        desc="interpolation",
    )
    out_points = File(
        name_source="points",
        name_template="%s_warped",
        output_name="out_points",
        keep_extension=True,
        desc="the warped point set",
    )

class WarpPointsOutputSpec(TraitedSpec):
    out_points = File(desc="the warped point set")

class WarpPoints(TVTKBaseInterface):
    """
    Applies a displacement field to a point set given in vtk format.
    Any discrete deformation field, given in physical coordinates and
    which volume covers the extent of the vtk point set, is a valid
    ``warp`` file. FSL interfaces are compatible, for instance any
    field computed with :class:`nipype.interfaces.fsl.utils.ConvertWarp`.

    Example::

        from nipype.algorithms.mesh import WarpPoints
        wp = WarpPoints()
        wp.inputs.points = 'surf1.vtk'
        wp.inputs.warp = 'warpfield.nii'
        res = wp.run()

    """

    input_spec = WarpPointsInputSpec
    output_spec = WarpPointsOutputSpec

    def _gen_fname(self, in_file, suffix="generated", ext=None):
        fname, fext = op.splitext(op.basename(in_file))

        if fext == ".gz":
            fname, fext2 = op.splitext(fname)
            fext = fext2 + fext

        if ext is None:
            ext = fext

        if ext[0] == ".":
            ext = ext[1:]
        return op.abspath("%s_%s.%s" % (fname, suffix, ext))

    def _run_interface(self, runtime):
        import nibabel as nb
        from scipy import ndimage

        r = tvtk.PolyDataReader(file_name=self.inputs.points)
        r.update()
        mesh = VTKInfo.vtk_output(r)
        points = np.array(mesh.points)
        warp_dims = nb.funcs.four_to_three(nb.load(self.inputs.warp))

        affine = warp_dims[0].affine
        # voxsize = warp_dims[0].header.get_zooms()
        vox2ras = affine[0:3, 0:3]
        ras2vox = np.linalg.inv(vox2ras)
        origin = affine[0:3, 3]
        voxpoints = np.array([np.dot(ras2vox, (p - origin)) for p in points])

        warps = []
        for axis in warp_dims:
            wdata = axis.dataobj  # four_to_three ensures this is an array
            if np.any(wdata != 0):
                warp = ndimage.map_coordinates(wdata, voxpoints.transpose())
            else:
                warp = np.zeros((points.shape[0],))

            warps.append(warp)

        disps = np.squeeze(np.dstack(warps))
        newpoints = [p + d for p, d in zip(points, disps)]
        mesh.points = newpoints
        w = tvtk.PolyDataWriter()
        VTKInfo.configure_input_data(w, mesh)
        w.file_name = self._gen_fname(self.inputs.points, suffix="warped", ext=".vtk")
        w.write()
        return runtime

    def _list_outputs(self):
        outputs = self._outputs().get()
        outputs["out_points"] = self._gen_fname(
            self.inputs.points, suffix="warped", ext=".vtk"
        )
        return outputs

class ComputeMeshWarpInputSpec(BaseInterfaceInputSpec):
    surface1 = File(
        exists=True,
        mandatory=True,
        desc=("Reference surface (vtk format) to which compute " "distance."),
    )
    surface2 = File(
        exists=True,
        mandatory=True,
        desc=("Test surface (vtk format) from which compute " "distance."),
    )
    metric = traits.Enum(
        "euclidean", "sqeuclidean", usedefault=True, desc="norm used to report distance"
    )
    weighting = traits.Enum(
        "none",
        "area",
        usedefault=True,
        desc=(
            '"none": no weighting is performed, surface": edge distance is '
            "weighted by the corresponding surface area"
        ),
    )
    out_warp = File(
        "surfwarp.vtk",
        usedefault=True,
        desc="vtk file based on surface1 and warpings mapping it " "to surface2",
    )
    out_file = File(
        "distance.npy",
        usedefault=True,
        desc="numpy file keeping computed distances and weights",
    )

class ComputeMeshWarpOutputSpec(TraitedSpec):
    distance = traits.Float(desc="computed distance")
    out_warp = File(
        exists=True,
        desc=("vtk file with the vertex-wise " "mapping of surface1 to surface2"),
    )
    out_file = File(
        exists=True, desc="numpy file keeping computed distances and weights"
    )

class ComputeMeshWarp(TVTKBaseInterface):
    """
    Calculates a the vertex-wise warping to get surface2 from surface1.
    It also reports the average distance of vertices, using the norm specified
    as input.

    .. warning:

      A point-to-point correspondence between surfaces is required

    Example::

        import nipype.algorithms.mesh as m
        dist = m.ComputeMeshWarp()
        dist.inputs.surface1 = 'surf1.vtk'
        dist.inputs.surface2 = 'surf2.vtk'
        res = dist.run()

    """

    input_spec = ComputeMeshWarpInputSpec
    output_spec = ComputeMeshWarpOutputSpec

    def _triangle_area(self, A, B, C):
        A = np.array(A)
        B = np.array(B)
        C = np.array(C)
        ABxAC = nla.norm(A - B) * nla.norm(A - C)
        prod = np.dot(B - A, C - A)
        angle = np.arccos(prod / ABxAC)
        area = 0.5 * ABxAC * np.sin(angle)
        return area

    def _run_interface(self, runtime):
        r1 = tvtk.PolyDataReader(file_name=self.inputs.surface1)
        r2 = tvtk.PolyDataReader(file_name=self.inputs.surface2)
        vtk1 = VTKInfo.vtk_output(r1)
        vtk2 = VTKInfo.vtk_output(r2)
        r1.update()
        r2.update()
        assert len(vtk1.points) == len(vtk2.points)

        points1 = np.array(vtk1.points)
        points2 = np.array(vtk2.points)

        diff = points2 - points1
        weights = np.ones(len(diff))

        try:
            errvector = nla.norm(diff, axis=1)
        except TypeError:  # numpy < 1.9
            errvector = np.apply_along_axis(nla.norm, 1, diff)

        if self.inputs.metric == "sqeuclidean":
            errvector **= 2

        if self.inputs.weighting == "area":
            faces = vtk1.polys.to_array().reshape(-1, 4).astype(int)[:, 1:]

            for i, p1 in enumerate(points2):
                # compute surfaces, set in weight
                w = 0.0
                point_faces = faces[(faces[:, :] == i).any(axis=1)]

                for idset in point_faces:
                    fp1 = points1[int(idset[0])]
                    fp2 = points1[int(idset[1])]
                    fp3 = points1[int(idset[2])]
                    w += self._triangle_area(fp1, fp2, fp3)
                weights[i] = w

        result = np.vstack([errvector, weights])
        np.save(op.abspath(self.inputs.out_file), result.transpose())

        out_mesh = tvtk.PolyData()
        out_mesh.points = vtk1.points
        out_mesh.polys = vtk1.polys
        out_mesh.point_data.vectors = diff
        out_mesh.point_data.vectors.name = "warpings"
        writer = tvtk.PolyDataWriter(file_name=op.abspath(self.inputs.out_warp))
        VTKInfo.configure_input_data(writer, out_mesh)
        writer.write()

        self._distance = np.average(errvector, weights=weights)
        return runtime

    def _list_outputs(self):
        outputs = self._outputs().get()
        outputs["out_file"] = op.abspath(self.inputs.out_file)
        outputs["out_warp"] = op.abspath(self.inputs.out_warp)
        outputs["distance"] = self._distance
        return outputs

class MeshWarpMathsInputSpec(BaseInterfaceInputSpec):
    in_surf = File(
        exists=True,
        mandatory=True,
        desc=(
            "Input surface in vtk format, with associated warp "
            "field as point data (ie. from ComputeMeshWarp"
        ),
    )
    float_trait = traits.Either(
        traits.Float(1.0),
        traits.Tuple(traits.Float(1.0), traits.Float(1.0), traits.Float(1.0)),
    )

    operator = traits.Either(
        float_trait,
        File(exists=True),
        default=1.0,
        usedefault=True,
        mandatory=True,
        desc="image, float or tuple of floats to act as operator",
    )

    operation = traits.Enum(
        "sum", "sub", "mul", "div", usedefault=True, desc="operation to be performed"
    )

    out_warp = File(
        "warp_maths.vtk",
        usedefault=True,
        desc="vtk file based on in_surf and warpings mapping it " "to out_file",
    )
    out_file = File("warped_surf.vtk", usedefault=True, desc="vtk with surface warped")

class MeshWarpMathsOutputSpec(TraitedSpec):
    out_warp = File(
        exists=True,
        desc=("vtk file with the vertex-wise " "mapping of surface1 to surface2"),
    )
    out_file = File(exists=True, desc="vtk with surface warped")

class MeshWarpMaths(TVTKBaseInterface):
    """
    Performs the most basic mathematical operations on the warping field
    defined at each vertex of the input surface. A surface with scalar
    or vector data can be used as operator for non-uniform operations.

    .. warning:

      A point-to-point correspondence between surfaces is required

    Example::

        import nipype.algorithms.mesh as m
        mmath = m.MeshWarpMaths()
        mmath.inputs.in_surf = 'surf1.vtk'
        mmath.inputs.operator = 'surf2.vtk'
        mmath.inputs.operation = 'mul'
        res = mmath.run()

    """

    input_spec = MeshWarpMathsInputSpec
    output_spec = MeshWarpMathsOutputSpec

    def _run_interface(self, runtime):
        r1 = tvtk.PolyDataReader(file_name=self.inputs.in_surf)
        vtk1 = VTKInfo.vtk_output(r1)
        r1.update()
        points1 = np.array(vtk1.points)

        if vtk1.point_data.vectors is None:
            raise RuntimeError("No warping field was found in in_surf")

        operator = self.inputs.operator
        opfield = np.ones_like(points1)

        if isinstance(operator, (str, bytes)):
            r2 = tvtk.PolyDataReader(file_name=self.inputs.surface2)
            vtk2 = VTKInfo.vtk_output(r2)
            r2.update()
            assert len(points1) == len(vtk2.points)

            opfield = vtk2.point_data.vectors

            if opfield is None:
                opfield = vtk2.point_data.scalars

            if opfield is None:
                raise RuntimeError("No operator values found in operator file")

            opfield = np.array(opfield)

            if opfield.shape[1] < points1.shape[1]:
                opfield = np.array([opfield.tolist()] * points1.shape[1]).T
        else:
            operator = np.atleast_1d(operator)
            opfield *= operator

        warping = np.array(vtk1.point_data.vectors)

        if self.inputs.operation == "sum":
            warping += opfield
        elif self.inputs.operation == "sub":
            warping -= opfield
        elif self.inputs.operation == "mul":
            warping *= opfield
        elif self.inputs.operation == "div":
            warping /= opfield

        vtk1.point_data.vectors = warping
        writer = tvtk.PolyDataWriter(file_name=op.abspath(self.inputs.out_warp))
        VTKInfo.configure_input_data(writer, vtk1)
        writer.write()

        vtk1.point_data.vectors = None
        vtk1.points = points1 + warping
        writer = tvtk.PolyDataWriter(file_name=op.abspath(self.inputs.out_file))
        VTKInfo.configure_input_data(writer, vtk1)
        writer.write()
        return runtime

    def _list_outputs(self):
        outputs = self._outputs().get()
        outputs["out_file"] = op.abspath(self.inputs.out_file)
        outputs["out_warp"] = op.abspath(self.inputs.out_warp)
        return outputs

class P2PDistance(ComputeMeshWarp):
    """
    Calculates a point-to-point (p2p) distance between two corresponding
    VTK-readable meshes or contours.

    A point-to-point correspondence between nodes is required

    .. deprecated:: 1.0-dev
       Use :py:class:`ComputeMeshWarp` instead.
    """

    def __init__(self, **inputs):
        super(P2PDistance, self).__init__(**inputs)
        IFLOGGER.warning(
            "This interface has been deprecated since 1.0, please "
            "use ComputeMeshWarp"
        )

max_issues_repo_name: tamires-consulting/nipype

max_forks_repo_name: tamires-consulting/nipype

reason: import numpy,from numpy,from scipy


text:
# -*- coding: utf-8 -*-
"""
@author: Adam
"""
import numpy as np
import pandas as pd
from tqdm import tqdm
from .trajectory import trajectory, final_position

def fly(fa, vol_t, initial, charge, mass, dt, **kwargs):
    """ Calculate the trajectories of charged particles in a
    time-varying electric field

        Parameters
        ----------
        fa :: FastAdjust
            fast-adjust potential arrays
        vol_t :: func(t),
            returns np.array([v0, v1, ... vn])              (V)
        initial :: pd.DataFrame
            initial positions and velocities
        charge :: float64
            particle charge                                 (C)
        mass :: float64
            particle mass                                   (kg)
        dt :: float64
            time step                                       (s)
        max_iterations :: int
            (default: 1 million)
        mode :: str
            'full' (default) or 'final'
        tqdm_kw :: dict
            keyword arguments for tqdm (progress bar)

        Returns
        -------
        mode='full': step-by-step trajectories for all particles ::
            pd.DataFrame(index=['particle', 'time'], columns=['x', 'y', 'z', 'KE', 'PE'])

        mode='final': final position for all particles ::
            pd.DataFrame(index=['particle'], columns=['t' 'x', 'y', 'z', 'vx', 'vy', 'vz'])
    """
    max_iterations = kwargs.get("max_iterations", int(1e6))
    mode = kwargs.get("mode", "full")
    tqdm_kw = kwargs.get("tqdm_kw", {})
    # fly ions
    num = len(initial.index)
    result = {}
    for i, row in tqdm(initial.iterrows(), total=len(initial.index), **tqdm_kw):
        t0 = row.time
        x0 = np.array([row.x, row.y, row.z])
        v0 = np.array([row.vx, row.vy, row.vz])
        if fa.electrode_r(x0):
            pass
        elif mode == 'full':
            result[i] = trajectory(fa, vol_t, t0, x0, v0,
                                   charge, mass, dt, max_iterations)
        elif mode == 'final':
            result[i] = final_position(fa, vol_t, t0, x0, v0,
                                       charge, mass, dt, max_iterations, to_series=False)
        else:
            raise ValueError("valid values for arg `mode` : 'full', 'final' ")
    # output
    if mode == "full":
        return pd.concat(result, names=["particle"])
    elif mode == "final":
        df = pd.DataFrame.from_dict(result,
                                    orient='index',
                                    columns=['time', 'x', 'y', 'z', 'vx', 'vy', 'vz'])
        df.index.rename('particle')
        return df
    else:
        raise ValueError("valid values for arg `mode` : 'full', 'final' ")


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.
#

"""Perform inference on a single video or all videos with a certain extension
(e.g., .mp4) in a folder.
"""

import detectron2
from detectron2.utils.logger import setup_logger
from detectron2.config import get_cfg
from detectron2 import model_zoo
from detectron2.engine import DefaultPredictor
from detectron2.modeling import build_model
from tqdm import tqdm

from chivideo.pipeline import *
from chivideo.cache_streams import *

import subprocess as sp
import numpy as np
import time
import argparse
import sys
import os
import glob

def parse_args():
    # right now we don't save intermediate results
    parser = argparse.ArgumentParser(description='End-to-end inference')
    parser.add_argument(
        '--cfg',
        dest='cfg',
        help='cfg model file (/path/to/model_config.yaml)',
        default=None,
        type=str
    )
    parser.add_argument(
        '--output-dir',
        dest='output_dir',
        help='directory for visualization pdfs (default: /tmp/infer_simple)',
        default='./resutls',
        type=str
    )
    parser.add_argument(
        '--image-ext',
        dest='image_ext',
        help='image file name extension (default: mp4)',
        default='mp4',
        type=str
    )
    parser.add_argument(
        'im_or_folder', help='image or folder of images', default=None
    )

    parser.add_argument('--dense', action='store_true', help='use dense convolutions instead of dilated convolutions')

    parser.add_argument('-ch', '--channels', default=1024, type=int, metavar='N', help='number of channels in convolution layers')
    parser.add_argument('-arc', '--architecture', default='3,3,3', type=str, metavar='LAYERS', help='filter widths separated by comma')

    if len(sys.argv) == 1:
        parser.print_help()
        sys.exit(1)
    return parser.parse_args()

class Skeleton:
    def __init__(self, parents, joints_left, joints_right):
        assert len(joints_left) == len(joints_right)

        self._parents = np.array(parents)
        self._joints_left = joints_left
        self._joints_right = joints_right
        self._compute_metadata()

    def num_joints(self):
        return len(self._parents)

    def parents(self):
        return self._parents

    def has_children(self):
        return self._has_children

    def children(self):
        return self._children

    def remove_joints(self, joints_to_remove):
        """
        Remove the joints specified in 'joints_to_remove'.
        """
        valid_joints = []
        for joint in range(len(self._parents)):
            if joint not in joints_to_remove:
                valid_joints.append(joint)

        for i in range(len(self._parents)):
            while self._parents[i] in joints_to_remove:
                self._parents[i] = self._parents[self._parents[i]]

        index_offsets = np.zeros(len(self._parents), dtype=int)
        new_parents = []
        for i, parent in enumerate(self._parents):
            if i not in joints_to_remove:
                new_parents.append(parent - index_offsets[parent])
            else:
                index_offsets[i:] += 1
        self._parents = np.array(new_parents)

        if self._joints_left is not None:
            new_joints_left = []
            for joint in self._joints_left:
                if joint in valid_joints:
                    new_joints_left.append(joint - index_offsets[joint])
            self._joints_left = new_joints_left
        if self._joints_right is not None:
            new_joints_right = []
            for joint in self._joints_right:
                if joint in valid_joints:
                    new_joints_right.append(joint - index_offsets[joint])
            self._joints_right = new_joints_right

        self._compute_metadata()

        return valid_joints

    def joints_left(self):
        return self._joints_left

    def joints_right(self):
        return self._joints_right

    def _compute_metadata(self):
        self._has_children = np.zeros(len(self._parents)).astype(bool)
        for i, parent in enumerate(self._parents):
            if parent != -1:
                self._has_children[parent] = True

        self._children = []
        for i, parent in enumerate(self._parents):
            self._children.append([])
        for i, parent in enumerate(self._parents):
            if parent != -1:
                self._children[parent].append(i)

class DetectronPoseOP(Operator):
    # need to initialize appropriate dstreams
    def __init__(self, name, vstream, predictor):
        input_names = [vstream]
        output_names = ['boxes', 'segments', 'keypoints']
        self.predictor = predictor
        self.vstream = vstream
        super().__init__(name, input_names, output_names)

    def next(self):
        im = self.istreams[self.vstream].next(self.name)
        if im == None:
            return False
        outputs = self.predictor(im)['instances'].to('cpu')
        has_bbox = False

        if outputs.has('pred_boxes'):
            bbox_tensor = outputs.pred_boxes.tensor.numpy()
            if len(bbox_tensor) > 0:
                has_bbox = True
                scores = outputs.scores.numpy()[:, None]
                bbox_tensor = np.concatenate((bbox_tensor, scores), axis=1)
        if has_bbox:
            kps = outputs.pred_keypoints.numpy()
            kps_xy = kps[:, :, :2]
            kps_prob = kps[:, :, 2:3]
            kps_logit = np.zeros_like(kps_prob) # Dummy
            kps = np.concatenate((kps_xy, kps_logit, kps_prob), axis=2)
            kps = kps.transpose(0, 2, 1)
        else:
            kps = []
            bbox_tensor = []

        # # Mimic Detectron1 format
        # cls_boxes = [[], bbox_tensor]
        # cls_keyps = [[], kps]

        self.ostreams['boxes'].append(bbox_tensor)
        self.ostreams['segments'].append(None)
        self.ostreams['keypoints'].append(kps)
        return True

    def init_outputs(self):
        for key in self.output_names:
            self.ostreams[key] = CacheStream(key)
        return super().init_outputs()

class FindBestOP(Operator):
    # need to initialize appropriate dstreams
    def __init__(self, name, boxes = 'boxes', keypoints = 'keypoints'):
        input_names = [boxes, keypoints]
        self.boxes = boxes
        self.keypoints = keypoints
        output_names = ['results_bb', 'results_kp']
        super().__init__(name, input_names, output_names)

    def next(self):
        bb = self.istreams[self.boxes].next()
        kp = self.istreams[self.keypoints].next()
        if bb == None and kp == None:
            return False
        if len(bb) == 0 or len(kp) == 0:
            # No bbox/keypoints detected for this frame -> will be interpolated
            self.ostreams['results_bb'].append(np.full(4, np.nan, dtype=np.float32)) # 4 bounding box coordinates
            self.ostreams['results_kp'].append(np.full((17, 4), np.nan, dtype=np.float32)) # 17 COCO keypoints
        else:
            best_match = np.argmax(bb[:, 4])
            best_bb = bb[best_match, :4]
            best_kp = kp[best_match].T.copy()
            self.ostreams['results_bb'].append(best_bb)
            self.ostreams['results_kp'].append(best_kp)
        return True

    def init_outputs(self):
        for key in self.output_names:
            self.ostreams[key] = CacheStream(key)
        return super().init_outputs()

class InterpolateOP(Operator):
# need to initialize appropriate dstreams
    def __init__(self, name, results_bb = 'results_bb', result_kp = 'results_kp'):
        input_names = [results_bb, result_kp]
        self.results_bb = results_bb
        self.result_kp = result_kp
        output_names = ['bbs_np', 'kps_np']
        super().__init__(name, input_names, output_names)

    def next(self):
        bb = []
        kp = []
        while True:
            bb_ = self.istreams[self.results_bb].next()
            kp_ = self.istreams[self.result_kp].next()
            if bb_ == None and kp_ == None:
                break
            bb.append(bb_)
            kp.append(kp_)

        bb = np.array(bb, dtype=np.float32)
        kp = np.array(kp, dtype=np.float32)
        kp = kp[:, :, :2] # Extract (x, y)

        # Fix missing bboxes/keypoints by linear interpolation
        mask = ~np.isnan(bb[:, 0])
        indices = np.arange(len(bb))
        for i in range(4):
            bb[:, i] = np.interp(indices, indices[mask], bb[mask, i])
        for i in range(17):
            for j in range(2):
                kp[:, i, j] = np.interp(indices, indices[mask], kp[mask, i, j])

        self.ostreams['bb_np'].insert(bb)
        self.ostreams['kp_np'].insert(kp)
        return True

    def init_outputs(self):
        for key in self.output_names:
            self.ostreams[key] = CacheStream(key)
        return super().init_outputs()

class ReFormat3DOP(Operator):
    def __init__(self, name, pad, causal_shift, meta = 'meta', kp = 'kp_np'):
        self.pad = pad
        self.causal_shift = causal_shift
        input_names = [meta, kp]
        self.meta = meta
        self.kp = kp
        output_names = ['out_kp']
        super().__init__(name, input_names, output_names)

    def next(self):
        X = self.istreams[self.kp].next()
        w, h = self.istreams[self.meta].next()
        X[..., :2] = X[..., :2]/w*2 - [1, h/w]
        X = np.expand_dims(np.pad(X, ((self.pad + self.causal_shift, self.pad - self.causal_shift), (0, 0), (0, 0)), 'edge'), axis=0)
        self.ostreams['out_kp'].insert(X)

    def init_outputs(self):
        for key in self.output_names:
            self.ostreams[key] = CacheStream(key)
        return super().init_outputs()

class Pose3DOP(Operator):
    def __init__(self, name, model, batch_2d = 'out_kp'):
        input_names = [batch_2d]
        self.batch_2d = batch_2d
        output_names = ['poses']
        self.model = model
        super.__init__(name, input_names, output_names)

    def next(self):
        inputs_2d = self.istreams[self.batch_2d].next()
        inputs_2d = torch.from_numpy(inputs_2d.astype('float32'))
        if torch.cuda.is_available():
                inputs_2d = inputs_2d.cuda()

        predicted_3d_pos = self.model(inputs_2d)
        predicted_3d_pos = predicted_3d_pos.squeeze(0).cpu().numpy()
        self.ostreams['poses'].insert(predicted_3d_pos)
        return True

    def init_outputs(self):
        for key in self.output_names:
            self.ostreams[key] = CacheStream(key)
        return super().init_outputs()

def initializePipeline(vstream, meta, x, y, predictor, skeleton, filter_widths, channels, dense):
    pipeline = Pipeline({'video': vstream, 'meta': meta})

    model_pos = TemporalModel(x, y, skeleton.num_joints(),
                            filter_widths=filter_widths, causal=False, dropout=0.0, channels=channels,
                            dense=dense)
    receptive_field = model_pos.receptive_field()
    pad = (receptive_field - 1) // 2
    causal_shift = 0

    pipeline.add_operator(DetectronPoseOP('detectron', 'video', predictor))
    pipeline.add_operator(FindBestOP('single'))
    pipeline.add_operatorInterpolateOP('interpolate'))
    pipeline.add_operator(ReFormat3DOP('reformat', pad, causal_shift))
    pipeline.add_operator(Pose3DOP('pose', model_pos))

    poses = pipeline.get_outputs(['poses'], keep_all = True)

    poses.add_iter('output')

    return pipeline, poses

def main(args):

    cfg = get_cfg()
    cfg.merge_from_file(model_zoo.get_config_file(args.cfg))
    cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.7
    cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(args.cfg)
    predictor = DefaultPredictor(cfg)

    if os.path.isdir(args.im_or_folder):
        im_list = glob.iglob(args.im_or_folder + '/*.' + args.image_ext)
    else:
        im_list = [args.im_or_folder]

    filter_widths = [int(x) for x in args.architecture.split(',')]

    skeleton = Skeleton(parents=[-1,  0,  1,  2,  3,  4,  0,  6,  7,  8,  9,  0, 11, 12, 13, 14, 12,
       16, 17, 18, 19, 20, 19, 22, 12, 24, 25, 26, 27, 28, 27, 30],
       joints_left=[6, 7, 8, 9, 10, 16, 17, 18, 19, 20, 21, 22, 23],
       joints_right=[1, 2, 3, 4, 5, 24, 25, 26, 27, 28, 29, 30, 31])

    remove_joints = True
    if remove_joints:
        skeleton.remove_joints([4, 5, 9, 10, 11, 16, 20, 21, 22, 23, 24, 28, 29, 30, 31])
        skeleton._parents[11] = 8
        skeleton._parents[14] = 8

    pipeline = None

    for video_name in tqdm(im_list):
        out_name = os.path.join(
                args.output_dir, os.path.basename(video_name)
            )

        print('Processing {}'.format(video_name))
        #predictor.filename = video_name
        #predictor.count = 0
        vstream = VideoStream('video', video_name)
        meta = JSONListStream([(vstream.width, vstream.hieght)])
        if pipeline == None:
            pipeline, poses = initializePipeline(vstream, meta, vstream.width, vstream.hieght, predictor, skeleton, filter_widths, args.channels, args.dense)

        else:
            pipeline.reset({'video':vstream, 'meta': meta})

        p = poses.next()

        # save poses in whatever format we want -> numpy array
        #snp.savez_compressed(out_name, boxes=boxes.all(), segments=segments.all(), keypoints=keypoints.all(), metadata=metadata)

if __name__ == '__main__':
    setup_logger()
    args = parse_args()
    main(args)


text:
import random
import torch
import collections
import numbers
import numpy as np

from torchvision.transforms import functional as F
from torchvision.transforms import InterpolationMode

def _flip_coco_person_keypoints(kps, width):
    flip_inds = [0, 2, 1, 4, 3, 6, 5, 8, 7, 10, 9, 12, 11, 14, 13, 16, 15]
    flipped_data = kps[:, flip_inds]
    flipped_data[..., 0] = width - flipped_data[..., 0]
    # Maintain COCO convention that if visibility == 0, then x, y = 0
    inds = flipped_data[..., 2] == 0
    flipped_data[inds] = 0
    return flipped_data

def _flip_coco_person_keypoints_v(kps, height):
    flip_inds = [0, 2, 1, 4, 3, 6, 5, 8, 7, 10, 9, 12, 11, 14, 13, 16, 15]
    flipped_data = kps[:, flip_inds]
    flipped_data[..., 1] = height - flipped_data[..., 1]
    # Maintain COCO convention that if visibility == 0, then x, y = 0
    inds = flipped_data[..., 2] == 0
    flipped_data[inds] = 0
    return flipped_data

class Compose(object):
    def __init__(self, transforms):
        self.transforms = transforms

    def __call__(self, image, target):
        for t in self.transforms:
            image, target = t(image, target)
        return image, target

class RandomHorizontalFlip(object):
    def __init__(self, prob):
        self.prob = prob

    def __call__(self, image, target):
        if random.random() < self.prob:
            height, width = image.shape[-2:]
            image = image.flip(-1)
            bbox = target["boxes"]
            bbox[:, [0, 2]] = width - bbox[:, [2, 0]]
            target["boxes"] = bbox
            if "masks" in target:
                target["masks"] = target["masks"].flip(-1)
            if "keypoints" in target:
                keypoints = target["keypoints"]
                keypoints = _flip_coco_person_keypoints(keypoints, width)
                target["keypoints"] = keypoints
        return image, target

class RandomVerticalFlip(object):
    def __init__(self, prob):
        self.prob = prob

    def __call__(self, image, target):
        if random.random() < self.prob:
            height, width = image.shape[-2:]
            image = image.flip(0)
            bbox = target["boxes"]
            bbox[:, [1, 3]] = height - bbox[:, [3, 1]]
            target["boxes"] = bbox
            if "masks" in target:
                target["masks"] = target["masks"].flip(0)
            if "keypoints" in target:
                keypoints = target["keypoints"]
                keypoints = _flip_coco_person_keypoints_v(keypoints, height)
                target["keypoints"] = keypoints
        return image, target

class Resize(object):

    def __init__(self, size=1500, interpolation=InterpolationMode.BILINEAR):
        assert isinstance(size, int) or (isinstance(size, collections.Iterable) and len(size) == 2)
        self.size = size
        self.interpolation = interpolation

    def __call__(self, image, target):

        image = F.resize(image, self.size, self.interpolation)
        bbox = (target["boxes"]).numpy()
        ratio = self.size/1024
        bbox = torch.from_numpy(np.trunc(bbox*ratio))
        target["boxes"] = bbox

        return image, target

class RandomCrop(object):
    """Crop the given PIL Image at a random location.

    Args:
        size (sequence or int): Desired output size of the crop. If size is an
            int instead of sequence like (h, w), a square crop (size, size) is
            made.
        padding (int or sequence, optional): Optional padding on each border
            of the image. Default is 0, i.e no padding. If a sequence of length
            4 is provided, it is used to pad left, top, right, bottom borders
            respectively.
    """

    def __init__(self, size=1024, padding=0):
        if isinstance(size, numbers.Number):
            self.size = (int(size), int(size))
        else:
            self.size = size
        self.padding = padding

    @staticmethod
    def get_params(image, output_size):

        w, h = image.size
        th, tw = output_size

        i = random.randint(0, h - th)
        j = random.randint(0, w - tw)
        return i, j, th, tw

    def __call__(self, image, target):

        if self.padding > 0:
            image = F.pad(image, self.padding)

        i, j, h, w = self.get_params(image, self.size)
        image = F.crop(image, i, j, h, w)

        th, tw = self.size
        bbox = (target["boxes"]).numpy()
        labels = (target['labels']).numpy()
        image_id = target["image_id"].numpy()
        area = target["area"].numpy()
        iscrowd = target["iscrowd"].numpy()
        delete_list = []

        for k in range(bbox.shape[0]):
            if bbox[k, 0] >= j + tw or bbox[k, 1] >= i + th or bbox[k, 2] <= j or bbox[k, 3] <= i:
                delete_list.append(k)
            else:
                if bbox[k, 0] < j:
                    bbox[k, 0] = 0
                else:
                    bbox[k, 0] = bbox[k, 0] - j

                if bbox[k, 1] < i:
                    bbox[k, 1] = 0
                else:
                    bbox[k, 1] = bbox[k, 1] - i

                if bbox[k, 2] > j + tw:
                    bbox[k, 2] = tw
                else:
                    bbox[k, 2] = bbox[k, 2] - j

                if bbox[k, 3] > i + th:
                    bbox[k, 3] = th
                else:
                    bbox[k, 3] = bbox[k, 3] - i

        #for m in range(len(delete_list)):
            #print('bbox.shape[0]: ', bbox.shape[0])
            #print('delete_list[m]: ', delete_list[m])
        bbox = np.delete(bbox, delete_list, axis=0)
        labels = np.delete(labels, delete_list)
        #image_id = np.delete(image_id, delete_list)
        area = np.delete(area, delete_list)
        iscrowd = np.delete(iscrowd, delete_list)

        target["boxes"] = torch.from_numpy(np.trunc(bbox))
        target["labels"] = torch.from_numpy(labels)
        target["image_id"] = torch.from_numpy(image_id)
        target["area"] = torch.from_numpy(area)
        target["iscrowd"] = torch.from_numpy(iscrowd)

        return image, target

class ToTensor(object):
    def __call__(self, image, target):
        image = F.to_tensor(image)
        return image, target


text:
# -*- coding: utf-8 -*-
"""
This pytorch custom dataset was modified from code in this repository:
https://github.com/jeromerony/survey_wsl_histology. Please cite their work:

@article{rony2019weak-loc-histo-survey,
  title={Deep weakly-supervised learning methods for classification and localization in histology images: a survey},
  author={Rony, J. and Belharbi, S. and Dolz, J. and Ben Ayed, I. and McCaffrey, L. and Granger, E.},
  journal={coRR},
  volume={abs/1909.03354},
  year={2019}
}
@author: jpeeples
"""

from PIL import Image
from torch.utils.data import Dataset
import matplotlib.pyplot as plt

from .utils import check_files
from torchvision.transforms import functional as F
from random import random, randint
from os.path import splitext
from os import listdir
import os
import pdb
from glob import glob
import itertools
import numpy as np
import torch

#Update load_data function
def load_data(samples, resize=None, min_resize=None):
    images = {}
    for image_path in samples:
        image = Image.open(image_path)
        if resize is not None:
            image = image.resize(resize, resample=Image.LANCZOS)
        elif min_resize:
            image = F.resize(image, min_resize, interpolation=Image.LANCZOS)
        images[image_path] = image.copy()
    return images

class PhotoDataset(Dataset):
    def __init__(self, data_path, files, transform, mask_transform, augment=False,
                 patch_size=None, rotate=False,
                 preload=False, resize=None, min_resize=None,class_label=True,
                 label_path=None,img_ext='.jpg',mask_ext='.png'):
        self.transform = transform
        self.mask_transform = mask_transform
        self.resize = resize
        self.min_resize = min_resize
        self.augment = augment
        self.patch_size = patch_size
        self.rotate = rotate
        self.masks_dir = label_path
        self.class_label = class_label
        self.img_ext = img_ext
        self.mask_ext = mask_ext

        #Get rgb mapping for CSAS data
        self.rgb = [(86, 115, 181), (132, 167, 77), (77, 77, 77), (141, 202, 207),
                    (211, 150, 202), (209, 205, 75), (255, 196, 52), (51, 69, 83),
                    (145, 58, 219), (58, 69, 219), (80, 139, 48), (164, 38, 41)]
        self.mapping = {tuple(c): t for c, t in zip(self.rgb, range(len(self.rgb)))}

       #Updated from previous loader
        if class_label:
            self.samples = [(os.path.join(data_path,file[0]), os.path.join(data_path,file[1]),
                         file[2]) for file in files]
        else:
            self.samples = [(os.path.join(data_path,file[0].replace(" ", "")), os.path.join(label_path,file[0].replace(" ", "")),
                          file[2]) for file in files]

            self.ids = [splitext(file)[0] for file in listdir(data_path)
                        if (not file.startswith('.') and os.path.isfile(os.path.join(data_path,file)))]

        self.n = len(self.samples)

        self.preloaded = False
        if preload:
            self.images = load_data([image_path for image_path, _, _ in self.samples],
                                    resize=resize, min_resize=min_resize)
            self.masks = load_data([mask_path for _, mask_path, _ in self.samples if mask_path != ''],
                                   resize=resize, min_resize=min_resize)
            self.preloaded = True
            print(self.__class__.__name__ + ' loaded with {} images'.format(len(self.images.keys())))

    def __len__(self):
        return len(self.samples)

    def __getitem__(self, index):
        image_path, mask_path, label = self.samples[index]
        img_name = image_path.rsplit('/',1)[-1].rsplit('.',1)[0]
        if self.preloaded:
                image = self.images[image_path].convert('RGB')
        else:
            try:
                image = Image.open(image_path).convert('RGB')
            except:
                image = Image.open(image_path+self.img_ext).convert('RGB')

            if self.resize is not None:
                image = image.resize(self.resize, resample=Image.LANCZOS)
            elif self.min_resize is not None:
                image = F.resize(image, self.min_resize, interpolation=Image.LANCZOS)
        image_size = image.size # to generate the mask if there is no file

        if mask_path == '':
            mask = Image.new('L', image_size)
        else:
            if self.preloaded:
                mask = self.masks[mask_path].convert('L')
            else:
                try:
                    mask = Image.open(mask_path).convert('L')
                except:
                    mask = Image.open(mask_path+self.mask_ext).convert('RGB')

                if self.resize is not None:
                    mask = mask.resize(self.resize, resample=Image.LANCZOS)
                elif self.min_resize is not None:
                    mask = F.resize(mask, self.min_resize, interpolation=Image.LANCZOS)

        if self.augment:

            # extract patch
            if self.patch_size is not None:
                left = randint(0, image_size[0] - self.patch_size)
                up = randint(0, image_size[1] - self.patch_size)
                image = image.crop(box=(left, up, left + self.patch_size, up + self.patch_size))
                mask = mask.crop(box=(left, up, left + self.patch_size, up + self.patch_size))

            # rotate
            if self.rotate:
                angle = randint(0, 3) * 90
                image = image.rotate(angle)
                mask = mask.rotate(angle)

            # flip
            if random() > 0.5:
                image = image.transpose(Image.FLIP_LEFT_RIGHT)
                mask = mask.transpose(Image.FLIP_LEFT_RIGHT)

        if self.transform is not None:
            image = self.transform(image)

            mask = self.mask_transform(mask)
            if self.class_label:
                #Preprocess to be binary
                mask = (mask != 0).long()
            else:
                #Clean up masks, software leaves two labels (1,2) for fat
                mask[mask>=.5] = 1
                mask[mask<.5] = 0

        return {'image':image,'mask': mask, 'index': img_name, 'label': label}


text:
# -*- coding: utf-8 -*-
"""
   File Name：     target.py
   Description :  rpn和rcnn 分类、回归、分割目标
   Author :       mick.yi
   Date：          2019/5/8
"""
import torch
from torch import nn
import numpy as np
from utils.np_utils import iou_3d, regress_target_3d
from utils import torch_utils
from cuda_functions.roi_align_3D.roi_align.crop_and_resize import CropAndResizeFunction as crop_and_resize

class RpnTarget(nn.Module):
    """
    rpn target 接收numpy类型的anchors和gt 的list列表
    """

    def __init__(self, train_anchors_per_image=800, positive_iou_threshold=0.5,
                 negative_iou_threshold=0.02, train_positive_anchors=2):
        self.train_anchors_per_image = train_anchors_per_image
        self.positive_iou_threshold = positive_iou_threshold
        self.negative_iou_threshold = negative_iou_threshold
        self.train_positive_anchors = train_positive_anchors
        super(RpnTarget, self).__init__()

    def forward(self, anchors, gt_boxes, gt_labels):
        """

        :param anchors: numpy 数组[anchors_num,(y1,x1,z1,y2,x2,z2)]
        :param gt_boxes: list of numpy [n,(y1,x1,z1,y2,x2,z2)]
        :param gt_labels: list of numpy [n]
        :return: anchors_tag: torch tensor [batch,anchors_num]  1-正样本，-1-负样本，0-不参与训练
        :return: deltas: torch tensor  [batch,anchors_num,(dy,dx,dz,dh,dw,dd)]
        """
        batch_anchors_tag = []
        batch_deltas = []
        # 逐个样本处理
        for boxes, labels in zip(gt_boxes, gt_labels):
            iou = iou_3d(boxes, anchors)  # [gt_num,anchors_num]
            anchors_iou_max = np.max(iou, axis=0)  # [anchors_num]
            anchors_tag = np.zeros_like(anchors[:, 0], np.int)  # [anchors_num]
            # 正负样本
            pos_indices = np.where(anchors_iou_max >= self.positive_iou_threshold)[0]
            neg_indices = np.where(anchors_iou_max <= self.negative_iou_threshold)[0]
            # 正样本采样
            pos_num = min(pos_indices.shape[0], self.train_positive_anchors)
            np.random.shuffle(pos_indices)
            pos_indices = pos_indices[:pos_num]
            anchors_tag[pos_indices] = 1
            # 负样本采样
            neg_num = min(neg_indices.shape[0], self.train_anchors_per_image - pos_num)
            np.random.shuffle(neg_indices)
            neg_indices = neg_indices[:neg_num]
            anchors_tag[neg_indices] = -1

            # 计算回归目标
            anchors_iou_argmax = np.argmax(iou, axis=0)  # [anchors_num]
            pos_gt_indices = anchors_iou_argmax[pos_indices]  # 正样本对应的gt索引号
            pos_gt_boxes = boxes[pos_gt_indices]
            pos_anchors = anchors[pos_indices]
            deltas = np.zeros_like(anchors)  # [anchors_num,(dy,dx,dz,dh,dw,dd)]
            deltas[pos_indices] = regress_target_3d(pos_anchors, pos_gt_boxes)

            # 转为tensor
            batch_anchors_tag.append(torch.from_numpy(anchors_tag).cuda())
            batch_deltas.append(torch.from_numpy(deltas).cuda())

        batch_anchors_tag = torch.stack(batch_anchors_tag, dim=0)
        batch_deltas = torch.stack(batch_deltas, dim=0)
        return batch_anchors_tag, batch_deltas

class MrcnnTarget(nn.Module):
    """
    计算mrcnn网络的分类、回归、分割 目标；还有一个作用就是过滤训练的proposals
    """

    def __init__(self, train_rois_per_image, mask_pool_size=28, positive_iou_threshold=0.5,
                 negative_iou_threshold=0.02, positive_ratio=0.1, ):
        """

        :param train_rois_per_image: 数值
        :param mask_pool_size: mask池化后的尺寸，数值，一般28
        :param positive_iou_threshold:
        :param negative_iou_threshold:
        :param positive_ratio:
        """
        super(MrcnnTarget, self).__init__()
        self.train_rois_per_image = train_rois_per_image
        self.mask_pool_size = mask_pool_size
        self.positive_iou_threshold = positive_iou_threshold
        self.negative_iou_threshold = negative_iou_threshold
        self.positive_ratio = positive_ratio

    def forward(self, proposals, batch_indices, gt_boxes, gt_labels, gt_masks):
        """

        :param proposals: tensor [proposals_num,(y1,x1,z1,y2,x2,z2)]
        :param batch_indices: proposal在原始的mini-batch中的索引号，tensor [proposals_num]
        :param gt_boxes: list of numpy [n,(y1,x1,z1,y2,x2,z2)]
        :param gt_labels: list of numpy [n]
        :param gt_masks: list of numpy [y,x,z]  0 1 值；每张只有一个mask

        :return: rois: tensor [rois_num,(y1,x1,z1,y2,x2,z2)]
        :return: deltas: tensor [rois_num,(dy,dx,dz,dh,dw,dd)]
        :return: labels: tensor [rois_num,(y1,x1,z1,y2,x2,z2)]
        :return: masks: tensor [rois_num,28,28,28] 0 1 值；28是mask_pool_size
        :return: rois_indices: tensor [rois_num] roi在原始的mini-batch中的索引号;roiAlign时用到
        """
        batch_rois = []
        batch_deltas = []
        batch_labels = []
        batch_masks = []
        batch_rois_indices = []
        # 逐个样本处理
        for i in range(len(gt_boxes)):
            # gt to gpu
            boxes = torch.from_numpy(gt_boxes[i]).float().cuda()
            labels = torch.from_numpy(gt_labels[i]).cuda()
            masks = torch.from_numpy(gt_masks[i]).int().cuda()
            # 属于第i个样本的proposals
            roi_indices = (batch_indices == i).nonzero()[:, 0]  # 索引
            rois = torch.index_select(proposals, 0, roi_indices)
            # 计算iou
            iou = torch_utils.iou_3d(boxes, rois)  # [gt_num,roi_num]
            # 正样本
            roi_max, _ = torch.max(iou, 0, keepdim=True)  # 每个roi最大的iou值
            pos_indices = (roi_max == iou) * (iou >= self.positive_iou_threshold)  # 正样本索引
            pos_indices = pos_indices.nonzero()  # [n,(gt_idx,roi_idx)]
            # 负样本
            neg_indices = (roi_max == iou) * (iou < self.negative_iou_threshold)
            neg_indices = neg_indices.nonzero()

            # 采样
            pos_num = int(self.positive_ratio * self.train_rois_per_image)
            pos_num = min(pos_num, pos_indices.shape[0])
            pos_indices = torch_utils.shuffle_and_select(pos_indices, pos_num)

            neg_num = min(self.train_rois_per_image - pos_num, neg_indices.shape[0])
            neg_indices = torch_utils.shuffle_and_select(neg_indices, neg_num)
            # 合并indices
            indices = torch.cat([pos_indices, neg_indices], dim=0)
            roi_indices = indices[:, 1]
            gt_indices = indices[:, 0]
            rois = torch.index_select(rois, 0, roi_indices)
            boxes = torch.index_select(boxes, 0, gt_indices)
            labels = torch.index_select(labels, 0, gt_indices)
            labels[-neg_num:] = 0  # 负样本标签为0
            # 计算回归目标
            deltas = torch_utils.regress_target_3d(rois, boxes)

            # 新的batch_indices
            batch_rois_indices.append(torch.Tensor([i] * (pos_num + neg_num)).cuda())

            # 添加rois,deltas,labels
            batch_rois.append(rois)
            batch_deltas.append(deltas)
            batch_labels.append(labels)
            batch_masks.append(masks)

        # 在拼接到一块
        batch_rois = torch.cat(batch_rois, dim=0)
        batch_deltas = torch.cat(batch_deltas, dim=0)
        batch_labels = torch.cat(batch_labels, dim=0)
        batch_rois_indices = torch.cat(batch_rois_indices, dim=0).int()
        batch_masks = torch.cat(batch_masks, dim=0)  # [n,y,x,z]
        # 最后处理masks
        boxes = batch_rois[:, [1, 0, 4, 3, 2, 5]]  # [n,(y1,x1,z1,y2,x2,z2)] => [n,(x1,y1,x2,y2,z1,z2)]
        batch_masks = crop_and_resize(self.mask_pool_size,
                                      self.mask_pool_size,
                                      self.mask_pool_size,
                                      0)(batch_masks, boxes, batch_indices)  # [rois_num,28,28,28]

        return batch_rois, batch_deltas, batch_labels, batch_masks, batch_rois_indices


text:
import numpy as np
import matplotlib.pyplot as plt
x = [1, 2, 3, 4, 5, 6, 7]
y = [17, 7, 14, 18, 12, 22, 13]
plt.figure()
plt.plot([0, 7], [15, 15], 'k-', lw=2)  # Population Mean = 15
plt.errorbar(x, y, yerr=4, fmt='.')
plt.xlabel("Confidence Intervals")
plt.ylabel("Sample Means")
plt.show()

max_stars_repo_name: yosais/Computer-Simulation-Book

max_issues_repo_name: yosais/Computer-Simulation-Book

max_forks_repo_name: yosais/Computer-Simulation-Book