text:
"""TRIVARIATE_UNQ.py -- Python Module

Creates the optimization problems needed to compute unique information

The optimization problems are:

min -H(S|XY), min -H(S|XZ) and min -H(S|YZ)

(c) Abdullah Makkeh, Dirk Oliver Theis

Permission to use and modify under Apache License version 2.0
"""
import ecos
from scipy import sparse
import numpy as np
from numpy import linalg as LA
from collections import defaultdict
import math
import time
from collections import defaultdict
ln  = math.log
log = math.log2

def core_initialization(_S, X_1, X_2, b_sx1, b_sx2, idx_of, of_idx):
    # (c) Abdullah Makkeh, Dirk Oliver Theis
    # Permission to use and modify under Apache License version 2.0

    for s in _S:
        for x1 in X_1:
            if (s,x1) in b_sx1.keys():
                for x2 in X_2:
                    if (s,x2) in b_sx2.keys():
                        idx_of[ (s,x1,x2) ] = len(of_idx)
                        of_idx.append( (s,x1,x2) )
                    #^ if sx2
                #^ for x2
            #^if sx1
        #^ for x1
    #^ for s

    return 0
#^ core_initialization()

def initialization(self, which_sources):
    """Initialize the data for the triplets (T,U,V) where U,V in {X,Y,Z}

        Args:
             which_sources:
                           [1,2] if sources are X and Y

                           [1,3] if sources are X and Z

                           [2,3] if sources are Y and Z

        Returns:
            (if [1,2] u,v=x,y|if [1,3] u,v=x,z|if [2,3] u,v=y,z|)

            dictionary

              keys: (t,u,v)
              values: their indices

            list of (t,u,v)
    """

    idx_of_trip  = dict()
    trip_of_idx  = []

    if which_sources == [1,2]:
        core_initialization(self.T, self.X, self.Y, self.b_tx, self.b_ty, idx_of_trip, trip_of_idx)
    elif which_sources == [1,3]:
        core_initialization(self.T, self.X, self.Z, self.b_tx, self.b_tz, idx_of_trip, trip_of_idx)
    elif which_sources == [2,3]:
        core_initialization(self.T, self.Y, self.Z, self.b_ty, self.b_tz, idx_of_trip, trip_of_idx)
    else:
        print("TRIVARIATE_UNQ.initialization(): which_sources takes the values [1,2], [1,3], or [2,3]")
        exit(1)

    return idx_of_trip, trip_of_idx
#^ initialization()

# ECOS's exp cone: (r,p,w)     w/   w>0  &  exp(r/w) ≤ p/w
# Variables here:  (r,p,w)U(q)
# Translation:     (0,1,2,3)   w/   2>0  &  0/2      ≤ ln(1/2)
def sr_vidx(i):
    """Computes the index of the optimal r (r_vars) in the optimal solution of the Exponential Cone Programming

        Args:
             i: int

        Returns:
            int
    """

    return 3*i
def sp_vidx(i):
    """Computes the index of the optimal (p_vars) in the optimal solution of the Exponential Cone Programming

        Args:
             i: int

        Returns:
            int
    """

    return 3*i+1
def sw_vidx(i):
    """Computes the index of the optimal (w_vars) in the optimal solution of the Exponential Cone Programming

        Args:
             i: int

        Returns:
             int
    """

    return 3*i+2
def sq_vidx(self, i, ltrip_of_idx):
    """Computes the index of the optimal distribution (q_vars) in the optimal solution of the Exponential Cone Programming

        Args:
             i: int
             ltrip_of_idx: int - length of triplet t,u,v
        Returns:
             int
    """

    return 3*ltrip_of_idx + i

def create_model(self, which_sources, output):
    """Creates the exponential Cone Program min_{q in Delta_d}H(T|U,V) of the form

              min. c'x

              s.t.

              Ax = b

              Gx <=_K h

           where

           x = (r,p,q)

           K represents a vector representing cones (K_1, K_2) such that K_1 is a vector repesenting exponential cones and K_2 is a vector repesenting nonnegative cones

        Args:
             which_sources:
                            [1,2] if sources are X and Y

                            [1,3] if sources are X and Z

                            [2,3] if sources are Y and Z

        Returns:

            numpy.array - objective function weights

            scipy.sparse.csc_matrix - matrix of exponential and nonnegative inequalities

            numpy.array - L.H.S. of inequalities

            dictionary -  cones to be used

                keys: string - cone type (exponential or nonegative)
                values: int - number of cones

            scipy.sparse.csc_matrix - Matrix of marginal, q-w coupling, and q-p coupling equations

            numpy.array - L.H.S. of equalities

    """

    # Initialize which sources for the model
    tic_all = time.process_time()
    idx_of_trip,trip_of_idx = self.initialization(which_sources)
    m = len(self.b_tx) + len(self.b_ty) + len(self.b_tz)
    n = len(trip_of_idx)
    ltrip_of_idx = n
    n_vars = 3*n + len(self.quad_of_idx)
    n_cons = 2*n + m

    # Create the equations: Ax = b
    self.b = np.zeros((n_cons,),dtype=np.double)

    Eqn   = []
    Var   = []
    Coeff = []

    # The q-w coupling eqautions:
    #         if Sources = X,Y  q_{stv*} - w_{stv} = 0
    #         if Sources = X,Z  q_{st*v} - w_{stv} = 0
    #         if Sources = Y,Z  q_{s*tv} - w_{stv} = 0
    tic_w = time.process_time()
    for i,stv in enumerate(trip_of_idx):
        eqn   = i
        w_var = sw_vidx(i)
        Eqn.append( eqn )
        Var.append( w_var )
        Coeff.append( -1. )
        (s,t,v) = stv
        if which_sources == [1,2]:
            for u in self.Z:
                if (s,t,v,u) in self.idx_of_quad.keys():
                    q_var = self.sq_vidx(self.idx_of_quad[ (s,t,v,u) ], ltrip_of_idx)
                    Eqn.append( eqn )
                    Var.append( q_var )
                    Coeff.append( +1. )
                #^ if q_{stv*}
            #^ loop *xy*
        #^ if SXY
        elif which_sources == [1,3]:
            for u in self.Y:
                if (s,t,u,v) in self.idx_of_quad.keys():
                    q_var = self.sq_vidx(self.idx_of_quad[ (s,t,u,v) ], ltrip_of_idx)
                    Eqn.append( eqn )
                    Var.append( q_var )
                    Coeff.append( +1. )
                #^ if q_{st*v}
            #^ loop *x*z
        #^ if SXZ
        elif which_sources == [2,3]:
            for u in self.X:
                if (s,u,t,v) in self.idx_of_quad.keys():
                    q_var = self.sq_vidx(self.idx_of_quad[ (s,u,t,v) ], ltrip_of_idx)
                    Eqn.append( eqn )
                    Var.append( q_var )
                    Coeff.append( +1. )
                #^ if q_{s*tv}
            #^ loop **yz
        #^if SYZ
    #^ for stv
    toc_w = time.process_time()

    if output == 2:
        if which_sources == [1,2]: print("TRIVARIATE_UNQ.create_model(): Time to create q-w coupling equations [min -H(S|X,Y)]:", toc_w - tic_w, "secs")
        if which_sources == [1,3]: print("TRIVARIATE_UNQ.create_model(): Time to create q-w coupling equations [min -H(S|X,Z)]:", toc_w - tic_w, "secs")
        if which_sources == [2,3]: print("TRIVARIATE_UNQ.create_model(): Time to create q-w coupling equations [min -H(S|Y,Z)]:", toc_w - tic_w, "secs")

    # running number
    eqn = -1 + len(trip_of_idx)

    # The q-p coupling equations:
    #         if Sources = X,Y  q_{*tv*} - p_{stv} = 0
    #         if Sources = X,Z  q_{*t*v} - p_{stv} = 0
    #         if Sources = Y,Z  q_{**tv} - p_{stv} = 0
    # ( Expensive step )
    tic_p = time.process_time()

    Eqn_dict = defaultdict(lambda: 0.)
    Eqn_dict_num = defaultdict(int)
    Eqn_dict_acc = defaultdict(list)
    Var_dict = defaultdict(list)
    Coeff_dict = defaultdict(list)

    if which_sources == [1,2]:
        for i,stv in enumerate(trip_of_idx):
            s,t,v = stv
            eqn     += 1
            Eqn_dict[(s,t,v)] = eqn
        #^ for

        for i,utvw in enumerate(self.quad_of_idx):
            u,t,v,w = utvw
            if (t,v) in self.b_xy.keys():
                Eqn_dict_num[(t,v)] += 1
                Coeff_dict[(t,v)].append(+1.)
                q_var = self.sq_vidx(self.idx_of_quad[ (u,t,v,w) ], ltrip_of_idx)
                Var_dict[(t,v)].append(q_var)
            #^ if *xy* exists
        #^ for utvw
        for i,stv in enumerate(trip_of_idx):
            s,t,v = stv
            temp = [Eqn_dict[(s,t,v)]]*Eqn_dict_num[(t,v)]
            Eqn_dict_acc[(s,t,v)] += temp
        # for
    elif which_sources == [1,3]:
        for i,stv in enumerate(trip_of_idx):
            s,t,v = stv
            eqn     += 1
            Eqn_dict[(s,t,v)] = eqn
        #^ for
        for i,utwv in enumerate(self.quad_of_idx):
            u,t,w,v = utwv
            if (t,v) in self.b_xz.keys():
                Eqn_dict_num[(t,v)] += 1
                Coeff_dict[(t,v)].append(+1.)
                q_var = self.sq_vidx(self.idx_of_quad[ (u,t,w,v) ], ltrip_of_idx)
                Var_dict[(t,v)].append(q_var)
            #^ if *x*z exists
        #^ for utwv

        for i,stv in enumerate(trip_of_idx):
            s,t,v = stv
            temp = [Eqn_dict[(s,t,v)]]*Eqn_dict_num[(t,v)]
            Eqn_dict_acc[(s,t,v)] += temp
        # for
    elif which_sources == [2,3]:
        for i,stv in enumerate(trip_of_idx):
            s,t,v = stv
            eqn     += 1
            Eqn_dict[(s,t,v)] = eqn
        #^ for
        for i,uwtv in enumerate(self.quad_of_idx):
            u,w,t,v = uwtv
            if (t,v) in self.b_yz.keys():
                Eqn_dict_num[(t,v)] += 1
                Coeff_dict[(t,v)].append(+1.)
                q_var = self.sq_vidx(self.idx_of_quad[ (u,w,t,v) ], ltrip_of_idx)
                Var_dict[(t,v)].append(q_var)
            #^ if **yz exists
        #^ for utwv

        for i,stv in enumerate(trip_of_idx):
            s,t,v = stv
            temp = [Eqn_dict[(s,t,v)]]*Eqn_dict_num[(t,v)]
            Eqn_dict_acc[(s,t,v)] += temp
        # for
    #^ if which sources

    for i,stv in enumerate(trip_of_idx):
            s,t,v = stv
            p_var   = sp_vidx(i)
            Eqn.append( Eqn_dict[(s,t,v)] )
            Eqn += Eqn_dict_acc[(s,t,v)]
            Var.append( p_var )
            Var += Var_dict[(t,v)]
            Coeff.append( -1. )
            Coeff += Coeff_dict[(t,v)]
    #^ for stv
    toc_p = time.process_time()

    if output == 2:
        if which_sources == [1,2]: print("TRIVARIATE_UNQ.create_model(): Time to create q-t coupling equations [min -H(S|X,Y)]:", toc_p - tic_p, "secs")
        if which_sources == [1,3]: print("TRIVARIATE_UNQ.create_model(): Time to create q-t coupling equations [min -H(S|X,Z)]:", toc_p - tic_p, "secs")
        if which_sources == [2,3]: print("TRIVARIATE_UNQ.create_model(): Time to create q-t coupling equations [min -H(S|Y,Z)]:", toc_p - tic_p, "secs")

    # Create the marginal constraints

    # The sx marginals q_{sx**} = b^x_{sx}
    tic_m = time.process_time()
    Eqn_marg = defaultdict(lambda: 0.)
    Eqn_marg_num = defaultdict(int)
    Eqn_marg_acc = defaultdict(list)
    Var_marg = defaultdict(list)
    Coeff_marg = defaultdict(list)

    for sx,i in self.b_tx.items():
        (s,x) = sx
        eqn += 1
        Eqn_marg[(s,x)] = eqn
    #^ for sx exists

    for i,sxyz in enumerate(self.quad_of_idx):
        s,x,y,z = sxyz
        if (s,x) in self.b_tx.keys():
            q_var = self.sq_vidx(self.idx_of_quad[ (s,x,y,z) ], ltrip_of_idx)
            Var_marg[(s,x)].append(q_var)
            Coeff_marg[(s,x)].append(+1.)
            Eqn_marg_num[(s,x)] += 1
            self.b[ Eqn_marg[(s,x)] ] = self.b_tx[(s,x)]
        #^ if sx exits
    # for sxyz

    for sx,i in self.b_tx.items():
        (s,x) = sx
        temp = [ Eqn_marg[(s,x)] ]*Eqn_marg_num[(s,x)]
        Eqn_marg_acc[(s,x)] += temp
    #^ for  sx exists

    for sx,i in self.b_tx.items():
        (s,x) = sx
        Eqn += Eqn_marg_acc[(s,x)]
        Var += Var_marg[(s,x)]
        Coeff += Coeff_marg[(s,x)]
    #^ for sx exists

    # The sy marginals q_{s*y*} = b^y_{sy}

    Eqn_marg = defaultdict(lambda: 0.)
    Eqn_marg_num = defaultdict(int)
    Eqn_marg_acc = defaultdict(list)
    Var_marg = defaultdict(list)
    Coeff_marg = defaultdict(list)
    for sy,i in self.b_ty.items():
        (s,y) = sy
        eqn += 1
        Eqn_marg[(s,y)] = eqn
    #^ for sy exists

    for i,sxyz in enumerate(self.quad_of_idx):
        s,x,y,z = sxyz
        if (s,y) in self.b_ty.keys():
            q_var = self.sq_vidx(self.idx_of_quad[ (s,x,y,z) ], ltrip_of_idx)
            Var_marg[(s,y)].append(q_var)
            Coeff_marg[(s,y)].append(+1.)
            Eqn_marg_num[(s,y)] += 1
            self.b[ Eqn_marg[(s,y)] ] = self.b_ty[(s,y)]
        #^ if sy exits
    # for sxyz

    for sy,i in self.b_ty.items():
        (s,y) = sy
        temp = [ Eqn_marg[(s,y)] ]*Eqn_marg_num[(s,y)]
        Eqn_marg_acc[(s,y)] += temp
    #^ for  sy exists

    for sy,i in self.b_ty.items():
        (s,y) = sy
        Eqn += Eqn_marg_acc[(s,y)]
        Var += Var_marg[(s,y)]
        Coeff += Coeff_marg[(s,y)]
    #^ for sy exists

    # The sz marginals q_{s**z} = b^z_{sz}
    Eqn_marg = defaultdict(lambda: 0.)
    Eqn_marg_num = defaultdict(int)
    Eqn_marg_acc = defaultdict(list)
    Var_marg = defaultdict(list)
    Coeff_marg = defaultdict(list)

    for sz,i in self.b_tz.items():
        (s,z) = sz
        eqn += 1
        Eqn_marg[(s,z)] = eqn
    #^ for sz exists

    for i,sxyz in enumerate(self.quad_of_idx):
        s,x,y,z = sxyz
        if (s,z) in self.b_tz.keys():
            q_var = self.sq_vidx(self.idx_of_quad[ (s,x,y,z) ], ltrip_of_idx)
            Var_marg[(s,z)].append(q_var)
            Coeff_marg[(s,z)].append(+1.)
            Eqn_marg_num[(s,z)] += 1
            self.b[ Eqn_marg[(s,z)] ] = self.b_tz[(s,z)]
        #^ if sz exits
    # for sxyz

    for sz,i in self.b_tz.items():
        (s,z) = sz
        temp = [ Eqn_marg[(s,z)] ]*Eqn_marg_num[(s,z)]
        Eqn_marg_acc[(s,z)] += temp
    #^ for  sz exists

    for sz,i in self.b_tz.items():
        (s,z) = sz
        Eqn += Eqn_marg_acc[(s,z)]
        Var += Var_marg[(s,z)]
        Coeff += Coeff_marg[(s,z)]
    #^ for sz exists
    toc_m = time.process_time()

    if output == 2:
        if which_sources == [1,2]: print("TRIVARIATE_UNQ.create_model(): Time to create marginal equations [min -H(S|X,Y)]:", toc_m - tic_m, "secs")
        if which_sources == [1,3]: print("TRIVARIATE_UNQ.create_model(): Time to create marginal equations [min -H(S|X,Z)]:", toc_m - tic_m, "secs")
        if which_sources == [2,3]: print("TRIVARIATE_UNQ.create_model(): Time to create marginal equations [min -H(S|Y,Z)]:", toc_m - tic_m, "secs")

    # Store the constraints in A
    tic_rest = time.process_time()
    self.A = sparse.csc_matrix( (Coeff, (Eqn,Var)), shape=(n_cons,n_vars), dtype=np.double)

    # Generalized ieqs: gen.nneg of the variable triple (r_i,w_i,p_i), i=0,dots,n-1:
    Ieq   = []
    Var   = []
    Coeff = []

    # Adding q_{s,x,y,z} >= 0 or q_{s,x,y,z} is free variable
    for i,sxyz in enumerate(self.quad_of_idx):
        q_var = self.sq_vidx(i, ltrip_of_idx)
        Ieq.append( len(Ieq) )
        Var.append( q_var )
        Coeff.append( -1. )
    #^ for sxyz

    for i,stv in enumerate(trip_of_idx):
        r_var = sr_vidx(i)
        w_var = sw_vidx(i)
        p_var = sp_vidx(i)

        Ieq.append( len(Ieq) )
        Var.append( r_var )
        Coeff.append( -1. )

        Ieq.append( len(Ieq) )
        Var.append( p_var )
        Coeff.append( -1. )

        Ieq.append( len(Ieq) )
        Var.append( w_var )
        Coeff.append( -1. )
    #^ for stv

    self.G         = sparse.csc_matrix( (Coeff, (Ieq,Var)), shape=(n_vars,n_vars), dtype=np.double)
    self.h         = np.zeros( (n_vars,),dtype=np.double )
    self.dims = dict()
    self.dims['e'] = n
    self.dims['l'] = len(self.quad_of_idx)

    # Objective function:
    self.c = np.zeros( (n_vars,),dtype=np.double )
    for i,stv in enumerate(trip_of_idx):
        self.c[ sr_vidx(i) ] = -1.
    #^ for stv

    toc_rest = time.process_time()
    if output == 2:
        if which_sources == [1,2]: print("TRIVARIATE_UNQ.create_model(): Time to create the matrices [min -H(S|X,Y)]:", toc_rest - tic_rest, "secs")
        if which_sources == [1,3]: print("TRIVARIATE_UNQ.create_model(): Time to create the matrices [min -H(S|X,Z)]:", toc_rest - tic_rest, "secs")
        if which_sources == [2,3]: print("TRIVARIATE_UNQ.create_model(): Time to create the matrices [min -H(S|Y,Z)]:", toc_rest - tic_rest, "secs")
    toc_all = time.process_time()
    if output > 0:
        if which_sources == [1,2]: print("TRIVARIATE_UNQ.create_model(): Time to create model [min -H(S|X,Y)]:", toc_all - tic_all, "secs")
        if which_sources == [1,3]: print("TRIVARIATE_UNQ.create_model(): Time to create model [min -H(S|X,Z)]:", toc_all - tic_all, "secs")
        if which_sources == [2,3]: print("TRIVARIATE_UNQ.create_model(): Time to create model [min -H(S|Y,Z)]:", toc_all - tic_all, "secs")

    return self.c, self.G, self.h, self.dims, self.A, self.b

#^ create_model()

def solve(self, c, G, h, dims, A, b, output):
    """Solves the exponential Cone Program min_{Delta_p}H(T|U,V) where U,V in {X,Y,Z}

        Args:

            c: numpy.array - objective function weights

            G: scipy.sparse.csc_matrix - matrix of exponential and nonnegative inequalities

            h: numpy.array - L.H.S. of inequalities

            dims: dictionary -  cones to be used

                    keys: string - cone type (exponential or nonegative)
                    values: int - number of cones

            A: scipy.sparse.csc_matrix - Matrix of marginal, q-w coupling, and q-p coupling equations

            b: numpy.array - L.H.S. of equalities

            output: int - print different outputs based on (int) to console

       Returns:
            sol_rpq:    numpy.array - primal optimal solution

            sol_slack:  numpy.array - slack of primal optimal solution (G*sol_rpq - h)

            sol_lambda: numpy.array - equalities dual optimal solution

            sol_mu:     numpy.array - inequalities dual  optimal solution

            sol_info:   dictionary - Brief stats of the optimization from ECOS

    """

    itic = time.process_time()
    self.marg_xyz = None # for cond[]mutinf computation below

    if self.verbose != None:
        # print(self.verbose)
        self.ecos_kwargs["verbose"] = self.verbose
    #^ if

    solution = ecos.solve(c, G, h, dims, A, b, **self.ecos_kwargs)

    if 'x' in solution.keys():
        self.sol_rpq    = solution['x']
        self.sol_slack  = solution['s']
        self.sol_lambda = solution['y']
        self.sol_mu     = solution['z']
        self.sol_info   = solution['info']
        itoc = time.process_time()
        if output == 2: print("TRIVARIATE_UNQ.solve(): Time to solve the Exponential Program of H(S|W,T)", itoc - itic, "secs")
        return "success", self.sol_rpq, self.sol_slack, self.sol_lambda, self.sol_mu, self.sol_info
    else: # "x" not in dict solution
        return "TRIVARIATE_UNQ.solve(): x not in dict solution -- No Solution Found!!!"
    #^ if/esle
#^ solve()

def check_feasibility(self, which_sources, sol_rpq, sol_slack, sol_lambda, sol_mu, output= 0):
    """Checks the KKT conditions of the exponential Cone Program min_{Delta_p}H(T|U,V) where U,V in {X,Y,Z}

        Args:
             which_sources:
                            [1,2] if sources are X and Y

                            [1,3] if sources are X and Z

                            [2,3] if sources are Y and Z

             sol_rpq:    numpy.array - primal optimal solution

             sol_slack:  numpy.array - slack of primal optimal solution (G*sol_rpq - h)

             sol_lambda: numpy.array - equalities dual optimal solution

             sol_mu:     numpy.array - inequalities dual  optimal solution

             output: int - print different outputs based on (int) to console

        Returns:
             primal_infeasability: float - maximum violation of the optimal primal solution for primal equalities and inequalities

             dual_infeasability:   float - maximum violation of the optimal dual solution for dual equalities and inequalities

    """

    # returns pair (p,d) of primal/dual infeasibility (maxima)
    idx_of_trip,trip_of_idx = self.initialization(which_sources)

    n = len(trip_of_idx)
    ltrip_of_idx = n
    # Primal infeasiblility

    # non-negative ineqaulity
    itic_neg = time.process_time()
    max_q_negativity = 0.
    for i in range(len(self.quad_of_idx)):
        max_q_negativity = max(max_q_negativity, -sol_rpq[self.sq_vidx(i, ltrip_of_idx)])
    #^ for
    toc_neg = time.time()
    itoc_neg = time.process_time()
    if output == 2:
        if which_sources == [1,2]: print("TRIVARIATE_UNQ.check_feasibility(): Time to compute primal negativity violations [min -H(S|XY)]:", itoc_neg - itic_neg, "secs")
        if which_sources == [1,3]: print("TRIVARIATE_UNQ.check_feasibility(): Time to compute primal negativity violations [min -H(S|XZ)]:", itoc_neg - itic_neg, "secs")
        if which_sources == [2,3]: print("TRIVARIATE_UNQ.check_feasibility(): Time to compute primal negativity violations [min -H(S|YZ)]:", itoc_neg - itic_neg, "secs")
    #^ if printing
    max_violation_of_eqn = 0.

    itic_marg = time.process_time()
    # sx** - marginals:
    sol_b_sx = defaultdict(lambda: 0.)
    for i,sxyz in enumerate(self.quad_of_idx):
        s,x,y,z = sxyz
        sol_b_sx[s,x] += max(0., sol_rpq[self.sq_vidx(i, ltrip_of_idx)])
    #^ for sxyz exists
    for sx,i in self.b_tx.items():
        s,x  = sx
        mysum = i - sol_b_sx[s,x]
        max_violation_of_eqn = max( max_violation_of_eqn, abs(mysum) )
    #^ for sx exists

    # s*y* - marginals:
    sol_b_sy = defaultdict(lambda: 0.)
    for i,sxyz in enumerate(self.quad_of_idx):
        s,x,y,z = sxyz
        sol_b_sy[s,y] += max(0., sol_rpq[self.sq_vidx(i, ltrip_of_idx)])
    #^ for sxyz exists
    for sy,i in self.b_ty.items():
        s,y  = sy
        mysum = i - sol_b_sy[s,y]
        max_violation_of_eqn = max( max_violation_of_eqn, abs(mysum) )
    #^ for sy exists

    # s**z - marginals:
    sol_b_sz = defaultdict(lambda: 0.)
    for i,sxyz in enumerate(self.quad_of_idx):
        s,x,y,z = sxyz
        sol_b_sz[s,z] += max(0., sol_rpq[self.sq_vidx(i, ltrip_of_idx)])
    #^ for sxyz exists
    for sz,i in self.b_tz.items():
        s,z  = sz
        mysum = i - sol_b_sz[s,z]
        max_violation_of_eqn = max( max_violation_of_eqn, abs(mysum) )
    #^ for sz exists
    itoc_marg = time.process_time()

    if output == 2:
        if which_sources == [1,2]: print("TRIVARIATE_UNQ.check_feasibility(): Time to compute violation of marginal eqautions [min -H(S|XY)]:", itoc_marg - itic_marg, "secs")
        if which_sources == [1,3]: print("TRIVARIATE_UNQ.check_feasibility(): Time to compute violation of marginal eqautions [min -H(S|XZ)]:", itoc_marg - itic_marg, "secs")
        if which_sources == [2,3]: print("TRIVARIATE_UNQ.check_feasibility(): Time to compute violation of marginal eqautions [min -H(S|YZ)]:", itoc_marg - itic_marg, "secs")
    #^ if printing

    primal_infeasability = max(max_violation_of_eqn, max_q_negativity)

    # Dual infeasiblility

    dual_infeasability = 0.
    tic_idx = time.time()
    itic_idx = time.process_time()
    idx_of_sx = dict()
    i = 0
    for s in self.T:
        for x in self.X:
            if (s,x) in self.b_tx.keys():
                idx_of_sx[(s,x)] = i
                i += 1
            #^ if sx exists
        #^ for x
    #^ for s

    idx_of_sy = dict()
    i = 0
    for s in self.T:
        for y in self.Y:
            if (s,y) in self.b_ty.keys():
                idx_of_sy[(s,y)] = i
                i += 1
            #^ if sy exists
        #^ for y
    #^ for s

    idx_of_sz = dict()
    i = 0
    for s in self.T:
        for z in self.Z:
            if (s,z) in self.b_tz.keys():
                idx_of_sz[(s,z)] = i
                i += 1
            #^ if sz exists
        #^ for z
    #^ for s
    itoc_idx = time.process_time()

    if output == 2:
        if which_sources == [1,2]: print("TRIVARIATE_UNQ.check_feasibility(): Time to find correct dual idx [min -H(S|XY)]:", itoc_idx - itic_idx, "secs")
        if which_sources == [1,3]: print("TRIVARIATE_UNQ.check_feasibility(): Time to find correct dual idx [min -H(S|XZ)]:", itoc_idx - itic_idx, "secs")
        if which_sources == [2,3]: print("TRIVARIATE_UNQ.check_feasibility(): Time to find correct dual idx [min -H(S|YZ)]:", itoc_idx - itic_idx, "secs")

    # non-negativity dual ineqaulity

    itic_negD12 = time.process_time()
    if which_sources == [1,2]:

        # Get indices of dual variables of the marginal constriants
        sz_idx = defaultdict(lambda: 0.)
        for i,sxyz in enumerate(self.quad_of_idx):
            s,x,y,z = sxyz
            sz_idx[(s,x,y)] = 2*n + len(self.b_tx) + len(self.b_ty) + idx_of_sz[(s,z)]
        #^ for

        # Compute mu_xy: dual varaible of the q-t coupling contsraints
        mu_xy = defaultdict(lambda: 0.)
        for k, sxy in enumerate(trip_of_idx):
            s,x,y = sxy
            mu_xy[(x,y)] += sol_lambda[ n + idx_of_trip[(s,x,y)] ]
        #^ for sxy

        for k, sxy in enumerate(trip_of_idx):
            s,x,y = sxy
            sx_idx = 2*n + idx_of_sx[(s,x)]
            sy_idx = 2*n + len(self.b_tx) + idx_of_sy[(s,y)]
            # nu_sxy: dual variable of the q-w coupling constraints
            nu_sxy = sol_lambda[ idx_of_trip[(s,x,y)] ]
            assert nu_sxy == sol_lambda[k], "problem in nu"

            # mu_sxy: dual varaible of the q-t coupling contsraints
            mu_sxy = sol_lambda[ n + idx_of_trip[(s,x,y)] ]
            assert mu_sxy == sol_lambda[ n + k ], "problem in mu"

            # Find the most violated nonnegative dual ieq
            #     a      >= 0
            dual_infeasability = max(dual_infeasability, -sol_lambda[sx_idx]
                                     - sol_lambda[sy_idx]
                                     - sol_lambda[ int(sz_idx[(s,x,y)]) ]
                                     - mu_xy[(x,y)]
                                     - nu_sxy
            )

            # Find the most violated K_exp dual ieq
            # print("dual of kexp: ", sol_lambda[sx_idx]
            #   + sol_lambda[sy_idx]
            #   + mu_xy
            #   - nu_sxy
            #   +ln(-mu_sxy)
            #   +1)
        #^ for
        itoc_negD12 = time.process_time()
        if output == 2: print("TRIVARIATE_UNQ.check_feasibility(): Time to compute neagtive dual violations [min -H(S|XY)]:", itoc_negD12 - itic_negD12, "secs")
    #^ if sources

    itic_negD13 = time.process_time()
    if which_sources == [1,3]:

        # Get indices of dual variables of the marginal constriants
        sy_idx = defaultdict(lambda: 0.)
        for i,sxyz in enumerate(self.quad_of_idx):
            s,x,y,z = sxyz
            sy_idx[(s,x,z)] = 2*n + len(self.b_tx) + idx_of_sy[(s,y)]
        #^ for

        # mu_xz: dual varaible of the q-t coupling contsraints
        mu_xz = defaultdict(lambda: 0.)
        for k, sxz in enumerate(trip_of_idx):
            s,x,z = sxz
            mu_xz[(x,z)] += sol_lambda[ n + idx_of_trip[(s,x,z)] ]
        #^ for sxz

        for k, sxz in enumerate(trip_of_idx):
            s,x,z = sxz
            # Get indices of dual variables of the marginal constriants
            sx_idx = 2*n + idx_of_sx[(s,x)]
            sz_idx = 2*n + len(self.b_tx) + len(self.b_ty) + idx_of_sz[(s,z)]

            # nu_sxz: dual variable of the q-w coupling constraints

            nu_sxz = sol_lambda[ idx_of_trip[(s,x,z)] ]
            assert nu_sxz == sol_lambda[k], "problem in nu"

            # mu_sxz: dual varaible of the q-t coupling contsraints
            mu_sxz = sol_lambda[ n + idx_of_trip[(s,x,z)] ]
            assert mu_sxz == sol_lambda[ n + k ], "problem in mu"

            # Find the most violated nonnegative dual ieq
            #     a      >= 0
            dual_infeasability = max(dual_infeasability, -sol_lambda[sx_idx]
                                         - sol_lambda[sy_idx[(s,x,z)]]
                                         - sol_lambda[sz_idx]
                                         - mu_xz[(x,z)]
                                         - nu_sxz
            )

            # Find the most violated K_exp dual ieq
            # print("dual of kexp: ", sol_lambda[sx_idx]
            #   + sol_lambda[sz_idx]
            #   + mu_xz
            #   - nu_sxz
            #   +ln(-mu_sxz)
            #   +1)
        #^ for
        itoc_negD13 = time.process_time()
        if output == 2 : print("TRIVARIATE_UNQ.check_feasibility(): Time to compute neagtive dual violations [min -H(S|XZ)]:", itoc_negD13 - itic_negD13, "secs")
    #^ if sources

    itic_negD23 = time.process_time()
    if which_sources == [2,3]:

        # Get indices of dual variables of the marginal constriants
        sx_idx = defaultdict(lambda: 0.)
        for i,sxyz in enumerate(self.quad_of_idx):
            s,x,y,z = sxyz
            sx_idx[(s,y,z)] = 2*n + idx_of_sx[(s,x)]
        #^ for

        # mu_yz: dual varaible of the q-t coupling contsraints
        mu_yz = defaultdict(lambda: 0.)
        for k, syz in enumerate(trip_of_idx):
            s,y,z = syz
            mu_yz[(y,z)] += sol_lambda[ n + idx_of_trip[(s,y,z)] ]
        #^ for syz

        for k, syz in enumerate(trip_of_idx):
            s,y,z = syz
            # Get indices of dual variables of the marginal constriants
            sy_idx = 2*n + len(self.b_tx) + idx_of_sy[(s,y)]
            sz_idx = 2*n + len(self.b_tx) + len(self.b_ty) + idx_of_sz[(s,z)]

            # nu_syz: dual variable of the q-w coupling constraints

            nu_syz = sol_lambda[idx_of_trip[(s,y,z)]]
            assert nu_syz == sol_lambda[k], "problem in nu"

            # mu_sxy: dual varaible of the q-t coupling contsraints
            mu_syz = sol_lambda[ n + idx_of_trip[(s,y,z)] ]
            assert mu_syz == sol_lambda[ n + k ], "problem in mu"

            # Find the most violated nonnegative dual ieq
            #     a      >= 0
            dual_infeasability = max(dual_infeasability, -sol_lambda[sx_idx[(s,y,z)]]
                                         - sol_lambda[sy_idx]
                                         - sol_lambda[sz_idx]
                                         - mu_yz[(y,z)]
                                         - nu_syz
            )

            # Find the most violated K_exp dual ieq
            # print("dual of kexp: ", sol_lambda[sy_idx]
            #   + sol_lambda[sz_idx]
            #   + mu_yz
            #   - nu_syz
            #   +ln(-mu_syz)
            #   +1)
        #^ for
        itoc_negD23 = time.process_time()
        if output == 2: print("TRIVARIATE_UNQ.check_feasibility(): Time to compute neagtive dual violations [min -H(S|YZ)]:", itoc_negD23 - itic_negD23, "secs")
    #^ if sources

    return primal_infeasability, dual_infeasability
#^ check_feasibility()

def dual_value(self, sol_lambda, b):
    """Evaluates the dual value of H(T|U,V) where U,V in {X,Y,Z}

        Args:
             sol_lambda: numpy.array - equalities dual optimal solution

             b: numpy.array - L.H.S. of equalities

        Returns:
            float

    """

    return -np.dot(sol_lambda, b)

def marginals(self, which_sources, sol_rpq, output):
    """Computes all the marginal distributions of the optimal distribution

        Args:
             which_sources:
                            [1,2] if sources are X and Y and Q is the optimal distribution of min_{Delta_P} H(T|X,Y)

                            [1,3] if sources are X and Z and Q is the optimal distribution of min_{Delta_P} H(T|X,Z)

                            [2,3] if sources are Y and Z and Q is the optimal distribution of min_{Delta_P} H(T|Y,Z)

             sol_rpq: numpy.array - array of triplets (r,p,q) of Exponential cone where q is the optimal distribution

             output: int - print different outputs based on (int) to console

        Returns:
            dictionary - optimal marginal distribution of T
              keys: t
              values: Q(t)
            dictionary - optimal marginal distribution of X
              keys: x
              values: Q(x)
            dictionary - optimal marginal distribution of Y
              keys: y
              values: Q(y)
            dictionary - optimal marginal distribution of Z
              keys: z
              values: Q(z)
            dictionary - optimal marginal distribution of (T,X)
              keys: t,x
              values: Q(t,x)
            dictionary - optimal marginal distribution of (T,Y)
              keys: t,y
              values: Q(t,y)
            dictionary - optimal marginal distribution of (T,Z)
              keys: t,z
              values: Q(t,z)
            dictionary - optimal marginal distribution of (X,Y)
              keys: x,y
              values: Q(x,y)
            dictionary - optimal marginal distribution of (X,Z)
              keys: x,z
              values: Q(x,z)
            dictionary - optimal marginal distribution of (Y,Z)
              keys: y,z
              values: Q(y,z)
            dictionary - optimal marginal distribution of (T,X,Y)
              keys: t,x,y
              values: Q(t,x,y)
            dictionary - optimal marginal distribution of (T,X,Z)
              keys: t,x,z
              values: Q(t,x,z)
            dictionary - optimal marginal distribution of (T,Y,Z)
              keys: t,y,z
              values: Q(t,y,z)

    """

    # provide the positive marginals all of a pdf for random varibles (A,B,C,D)

    itic = time.process_time()

    # First order marginals
    marg_S = defaultdict(float)
    marg_X = defaultdict(float)
    marg_Y = defaultdict(float)
    marg_Z = defaultdict(float)

    # Second order marginals
    marg_SX = defaultdict(float)
    marg_SY = defaultdict(float)
    marg_SZ = defaultdict(float)
    marg_XY = defaultdict(float)
    marg_XZ = defaultdict(float)
    marg_YZ = defaultdict(float)

    # Third order marginals
    marg_SXY = defaultdict(float)
    marg_SXZ = defaultdict(float)
    marg_SYZ = defaultdict(float)

    # Initialize the triplet
    idx_of_trip,trip_of_idx = self.initialization(which_sources)
    ltrip_of_idx = len(trip_of_idx)

    # Compute the marginals
    for sxyz,i in self.idx_of_quad.items():
        s,x,y,z = sxyz
        marg_S[s] += max(0, sol_rpq[self.sq_vidx(i, ltrip_of_idx)])
        marg_X[x] += max(0, sol_rpq[self.sq_vidx(i, ltrip_of_idx)])
        marg_Y[y] += max(0, sol_rpq[self.sq_vidx(i, ltrip_of_idx)])
        marg_Z[z] += max(0, sol_rpq[self.sq_vidx(i, ltrip_of_idx)])
        marg_SX[(s,x)] += max(0, sol_rpq[self.sq_vidx(i, ltrip_of_idx)])
        marg_SY[(s,y)] += max(0, sol_rpq[self.sq_vidx(i, ltrip_of_idx)])
        marg_SZ[(s,z)] += max(0, sol_rpq[self.sq_vidx(i, ltrip_of_idx)])
        marg_XY[(x,y)] += max(0, sol_rpq[self.sq_vidx(i, ltrip_of_idx)])
        marg_XZ[(x,z)] += max(0, sol_rpq[self.sq_vidx(i, ltrip_of_idx)])
        marg_YZ[(y,z)] += max(0, sol_rpq[self.sq_vidx(i, ltrip_of_idx)])
        marg_SXY[(s,x,y)] += max(0,sol_rpq[self.sq_vidx(i, ltrip_of_idx)])
        marg_SXZ[(s,x,z)] += max(0,sol_rpq[self.sq_vidx(i, ltrip_of_idx)])
        marg_SYZ[(s,y,z)] += max(0,sol_rpq[self.sq_vidx(i, ltrip_of_idx)])
    #^ for
    itoc = time.process_time()
    if output == 2:
        if which_sources == [1,2]:
            print("TRIVARIATE_UNQ.marginals(): Time to compute marginals() of H(S|XY):", itoc - itic, "secs")
        elif which_sources == [1,3]:
            print("TRIVARIATE_UNQ.marginals(): Time to compute marginals() of H(S|XZ):", itoc - itic, "secs")
        elif which_sources == [2,3]:
            print("TRIVARIATE_UNQ.marginals(): Time to compute marginals() of H(S|YZ):", itoc - itic, "secs")
        #^ if sources
    #^ if printing

    return marg_S, marg_X, marg_Y, marg_Z, marg_SX, marg_SY, marg_SZ, marg_XY, marg_XZ, marg_YZ, marg_SXY, marg_SXZ, marg_SYZ

#^ marginals()

def condentropy_2vars(self, which_sources, sol_rpq, output, marg_XY, marg_XZ, marg_YZ, marg_SXY, marg_SXZ, marg_SYZ):
    """Evalutes the value of H(T|U,V) w.r.t. the optimal distribution where U,V in {X,Y,Z}

        Args:
             which_sources:
                            [1,2] if sources are X and Y

                            [1,3] if sources are X and Z

                            [2,3] if sources are Y and Z

             sol_rpq:  numpy.array - primal optimal solution

             output:   int - print different outputs based on (int) to console

             marg_XY:  dictionary - optimal marginal distribution of (X,Y)
                       keys: x,y
                       values: Q(x,y)
             marg_XZ:  dictionary - optimal marginal distribution of (X,Z)
                       keys: x,z
                       values: Q(x,z)
             marg_YZ:  dictionary - optimal marginal distribution of (Y,Z)
                       keys: y,z
                       values: Q(y,z)
             marg_TXY: dictionary - optimal marginal distribution of (T,X,Y)
                       keys: t,x,y
                       values: Q(t,x,y)
             marg_TXZ: dictionary - optimal marginal distribution of (T,X,Z)
                        keys: t,x,z
                        values: Q(t,x,z)
             marg_TYZ: dictionary - optimal marginal distribution of (T,Y,Z)
                       keys: t,y,z
                       values: Q(t,y,z)

        Returns:
            (if [1,2]: U,V=X,Y | if [1,3]: U,V=X,Z | if [2,3]: U,V=Y,Z)

            mysum: float - H(T|U,V)

    """

    # compute cond entropy of the distribution in self.sol_rpq

    mysum = 0.
    idx_of_trip,trip_of_idx = self.initialization(which_sources)
    if which_sources == [1,2]:
        itic = time.process_time()
        # H( S | X, Y )
        for sxy in idx_of_trip.keys():
            s,x,y = sxy
            if marg_XY[(x,y)] > 0 and  marg_SXY[(s,x,y)] > 0:
                # subtract q_{sxy}*log( q_{sxy}/q_{xy} )
                mysum -= marg_SXY[(s,x,y)]*log(marg_SXY[(s,x,y)]/marg_XY[(x,y)])
            #^ if
        #^ for
        itoc = time.process_time()

        if output == 2: print("TRIVARIATE_UNQ.condentropy_2vars(): Time to compute H(S|XY):", itoc - itic,"secs")
        # print("H(T|XY)", mysum)
        return mysum
    #^ if sources
    elif which_sources == [1,3]:
        # H( S | X, Z )
        itic = time.process_time()

        for sxz in idx_of_trip.keys():
            s,x,z = sxz
            if marg_XZ[(x,z)] > 0 and  marg_SXZ[(s,x,z)] > 0:
                # subtract q_{sxz}*log( q_{sxz}/q_{xz} )
                mysum -= marg_SXZ[(s,x,z)]*log(marg_SXZ[(s,x,z)]/marg_XZ[(x,z)])
            #^ if
        #^ for
        itoc = time.process_time()

        if output == 2: print("TRIVARIATE_UNQ.condentropy_2vars(): Time to compute H(S|XZ):", itoc - itic,"secs")
        # print("H(T|XZ)", mysum)
        return mysum
    #^ if sources
    elif which_sources == [2,3]:
        # H( S | Y, Z )
        itic = time.process_time()

        for syz in idx_of_trip.keys():
            s,y,z = syz
            if marg_YZ[(y,z)] > 0 and  marg_SYZ[(s,y,z)] > 0:
                # subtract q_{syz}*log( q_{syz}/q_{yz} )
                mysum -= marg_SYZ[(s,y,z)]*log(marg_SYZ[(s,y,z)]/marg_YZ[(y,z)])
            #^ if
        #^ for
        itoc = time.process_time()

        if output == 2: print("TRIVARIATE_UNQ.condentropy_2vars(): Time to compute H(S|YZ):", itoc - itic,"secs")
        # print("H(T|YZ)", mysum)
        return mysum
    #^ if sources
    else:
        print("TRIVARIATE_UNQ.condentropy_2vars(): which_sources takes the values [1,2], [1,3], [2,3]")
        exit(1)
    #^ if sources

#^ condentropy_2vars()

def condentropy_1var(self,which_sources,sol_rpq, marg_SX, marg_SY, marg_SZ, marg_X, marg_Y, marg_Z):
    # (c) Abdullah Makkeh, Dirk Oliver Theis
    # Permission to use and modify under Apache License version 2.0

    mysum = 0.
    if which_sources == [1,2]:
        # H( S | Z )
        itic = time.process_time()
        for s in self.T:
            for z in self.Z:
                if (s,z) in self.b_tz.keys() and marg_SZ[(s,z)] > 0 and marg_Z[z] > 0:
                    # Subtract q_{s,z}*log( q_{s,z}/ q_{z} )
                    mysum -= marg_SZ[(s,z)]*log(marg_SZ[(s,z)]/marg_Z[z])
                #^ if sz exists
            #^ for z
        #^ for z
        itoc = time.process_time()
        print("Time to compute H(S|Z):", itoc - itic,"secs")
        return mysum
    #^ if sources
    elif which_sources == [1,3]:
        # H( S | Y )
        itic = time.process_time()
        for s in self.T:
            for y in self.Y:
                if (s,y) in self.b_ty.keys()and marg_SY[(s,y)] > 0 and marg_Y[y] > 0:
                    # Subtract q_{s,y}*log( q_{s,y}/ q_{y} )
                    mysum -= marg_SY[(s,y)]*log(marg_SY[(s,y)]/marg_Y[y])
                #^ if sy exists
            #^ for y
        #^ for s
        itoc = time.process_time()
        print("Time to compute H(S|Y):", itoc - itic,"secs")
        return mysum
    #^ if sources

    elif which_sources == [2,3]:
        # H ( S | X )
        itic = time.process_time()
        for s in self.T:
            for x in self.X:
                if (s,x) in self.b_tx.keys() and marg_SX[(s,x)] > 0 and marg_X[x] > 0:
                    # Subtract q_{s,x}*log( q_{s,x}/ q_{x} )
                    mysum -= marg_SX[(s,x)]*log(marg_SX[(s,x)]/marg_X[x])
                #^ if sx exists
            #^ for y
        #^ for s
        itoc = time.process_time()
        print("Time to compute H(S|X):", itoc - itic,"secs")
        return mysum
    #^ if sources
    else:
        print(" which sources takes the values: [1,2], [1,3], [2,3]")
        exit(1)
    #^ if sources

#^ condentropy_1var()

# def entropy_S(self,pdf, sol_rpq, which_sources, output):
#     # (c) Abdullah Makkeh, Dirk Oliver Theis
#     # Permission to use and modify under Apache License version 2.0
#     mysum = 0.
#     marg_s = defaultdict(lambda: 0.)
#     idx_of_trip,trip_of_idx = self.initialization(which_sources)
#     ltrip_of_idx = len(trip_of_idx)
#     for sxyz, i in self.idx_of_quad.items():
#         s,x,y,z = sxyz
#         marg_s[s] += max(0, sol_rpq[self.sq_vidx(i, ltrip_of_idx)])
#     #^ for
#     for s in self.S:
#         if marg_s[s] > 0: mysum -= marg_s[s]*log(marg_s[s])

#     print("H(S) new:", mysum)
#     itic = time.process_time()
#     mysum = 0.
#     for s in self.S:
#         psum = 0.
#         for x in self.X:
#             if not (s,x) in self.b_sx: continue
#             for y in self.Y:
#                 if not (s,y) in self.b_sy:  continue
#                 for z in self.Z:
#                     if (s,x,y,z) in pdf.keys():
#                         psum += pdf[(s,x,y,z)]
#                     #^ if
#                 #^ for z
#             #^ for y
#         #^ for x
#         mysum -= psum * log(psum)
#     #^ for x
#     print("H(S) old:", mysum)
#     itoc = time.process_time()
#     if output == 2: print("Time to compute H(S):", itoc - itic,"secs")
#     return mysum

# #^ entropy_S()

#EOF


text:
from sklearn.model_selection import train_test_split
import numpy as np
import torch
from torch.utils.data import Dataset
from torchvision import transforms
import h5py
from skimage import io, transform
import os
import matplotlib.pyplot as plt

def dataset_factory(use_images=True,
height_folder="/hpctmp/e0310631/DSA4299/PathLossPredictionSatelliteImages/Data_Folder/Height_Images_2_resized",
Occupancy_Folder = "/hpctmp/e0310631/DSA4299/PathLossPredictionSatelliteImages/Data_Folder/Occupancy_Images_2_resized",
Tx_Position_Folder = "/hpctmp/e0310631/DSA4299/PathLossPredictionSatelliteImages/Data_Folder/Tx_Azimuth_2_resized",
Rx_Position_Folder = "/hpctmp/e0310631/DSA4299/PathLossPredictionSatelliteImages/Data_Folder/Rx_Azimuth_2_resized",
transform=True, data_augment_angle=10):
    #Longitude,Latitude,Speed,Distance,Distance_x,Distance_y,PCI_64,PCI_65,PCI_302
    #selected_features = [0, 1, 3, 4, 5, 6, 7, 8]
    #['Tx_Lon', 'Tx_Lat', 'Rx_Lon', 'Rx_Lat', 'Tx_Height', 'Rx_Height','Tx_Rx_Distance']
    selected_features = [0, 1, 2, 3, 4, 5, 6]
    # ['SINR', 'RSRP', 'RSRQ', 'Power']
    #selected_targets = [1]
    # ["RSS"]
    selected_targets = [0]
    print(selected_features)
    dataset_path='/hpctmp/e0310631/DSA4299/PathLossPredictionSatelliteImages/Data_Folder'
    features = np.load("{}/training_features.npy".format(dataset_path))
    targets = np.load("{}/training_targets.npy".format(dataset_path))
    test_features =  np.load("{}/test_features.npy".format(dataset_path))
    test_targets = np.load("{}/test_targets.npy".format(dataset_path))
    target_mu = np.load("{}/targets_mu.npy".format(dataset_path))
    target_std = np.load("{}/targets_std.npy".format(dataset_path))
    features_mu = np.load("{}/features_mu.npy".format(dataset_path))
    features_std = np.load("{}/features_std.npy".format(dataset_path))
    images = np.load("{}/train_image_idx.npy".format(dataset_path))
    test_images = np.load("{}/test_image_idx.npy".format(dataset_path))

    features = features[:, selected_features]
    test_features = test_features[:, selected_features]
    features_mu = features_mu[selected_features]
    features_std = features_std[selected_features]

    targets = targets[:, selected_targets]
    test_targets = test_targets[:, selected_targets]
    target_mu = target_mu[selected_targets]
    target_std = target_std[selected_targets]

    # Data augmentation
    if transform:
        #composed = transforms.Compose([transforms.ToPILImage(), transforms.Grayscale(), transforms.RandomAffine(data_augment_angle, shear=10), transforms.ToTensor()])
        composed = transforms.Compose([transforms.ToPILImage(), transforms.Grayscale(), transforms.ToTensor()])
        composed_2 = transforms.Compose([transforms.RandomAffine(data_augment_angle, shear=10)])
    else:
        composed = None
        composed_2 = None

    # Dataset
    train_dataset = DrivetestDataset(features, targets, images, target_mu, target_std, features_mu, features_std, use_images, height_folder, Occupancy_Folder, Tx_Position_Folder, Rx_Position_Folder, transform=composed, transform_2 = composed_2, dataset_type = "Train")
    #valid_dataset = DrivetestDataset(images, features, targets, valid_idx, target_mu, target_std, features_mean, features_std, use_images, image_folder)
    test_dataset = DrivetestDataset(test_features, test_targets, test_images,  target_mu, target_std, features_mu, features_std, use_images, height_folder, Occupancy_Folder,Tx_Position_Folder, Rx_Position_Folder, transform=transforms.Compose([transforms.ToPILImage(), transforms.Grayscale(), transforms.ToTensor()]), dataset_type = "Test")
    return train_dataset, test_dataset

class DrivetestDataset(Dataset):
    def __init__(self, features, targets, images, target_mu, target_std, feature_mu, feature_std, use_images, height_folder, Occupancy_Folder, Tx_Position_Folder, Rx_Position_Folder, transform=None, transform_2 = None, dataset_type = None):
        self.features = features
        self.targets = targets
        self.image_idx = images
        self.target_mu = target_mu
        self.target_std = target_std
        self.feature_mu = feature_mu
        self.feature_std = feature_std
        self.distances = (self.features[:,2] * self.feature_std[2])+self.feature_mu[2]
        self.targets_unnorm = (self.targets * self.target_std)+self.target_mu
        self.use_images = use_images
        self.height_folder = height_folder
        self.Occupancy_Folder = Occupancy_Folder
        self.Tx_Position_Folder = Tx_Position_Folder
        self.Rx_Position_Folder = Rx_Position_Folder
        self.transform = transform
        self.transform_2 = transform_2
        self.dataset_type = dataset_type

    def get_811Mhz_idx(self):
        return np.argwhere(np.asarray(self.features[:,7] != 1))

    def get_2630Mhz_idx(self):
        return np.argwhere(np.asarray(self.features[:,7] == 1))

    def __getitem__(self, index):
        idx = self.image_idx[index]
        X = torch.from_numpy(self.features[index]).float() # Features (normalized)
        if self.use_images:
            if self.height_folder == None: #images are then pointer to hdf5
                image = self.image_idx[index]
            else:
                height_name = os.path.join(self.height_folder, "{}.jpg".format(idx))
                height_image = io.imread(height_name)
                height_image = height_image / 255

                occupancy_name = os.path.join(self.Occupancy_Folder, "{}.jpg".format(idx))
                Occupancy_image = io.imread(occupancy_name)
                Occupancy_image = Occupancy_image / 255

                Tx_P_name = os.path.join(self.Tx_Position_Folder, "{}.jpg".format(idx))
                Tx_P_image = io.imread(Tx_P_name)
                Tx_P_image = Tx_P_image / 255

                Rx_P_name = os.path.join(self.Rx_Position_Folder, "{}.jpg".format(idx))
                Rx_P_image = io.imread(Rx_P_name)
                Rx_P_image = Rx_P_image / 255

            A = torch.from_numpy(height_image).float().permute(2,0,1)
            B = torch.from_numpy(Occupancy_image).float().permute(2,0,1)
            C = torch.from_numpy(Tx_P_image).float().permute(2,0,1)
            D = torch.from_numpy(Rx_P_image).float().permute(2,0,1)

        else:
            A = torch.tensor(0)
            B = torch.tensor(0)
            C = torch.tensor(0)
            D = torch.tensor(0)

        y = torch.from_numpy(self.targets[index]).float() # Target
        dist = torch.abs(torch.tensor(self.distances[index])).float().view(1) # Unormalized distance
        dist = dist * 1000 # to meters

        if self.use_images:
            if self.transform:
                A = self.transform(A)
                B = self.transform(B)
                C = self.transform(C)
                D = self.transform(D)

                Composed_Image = torch.cat((A, B, C, D), 0)

                if self.dataset_type == "Train":
                    Composed_Image = self.transform_2(Composed_Image)
        else:
            Composed_Image = A

        return X, Composed_Image, y, dist

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

if __name__ == '__main__':
    train, test = dataset_factory()
    data = train.__getitem__(2)
    print(data[1].shape)

ext: py

reason: import numpy


text:
# coding=utf-8
"""The proposed model training code."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from absl import flags

from ieg import utils
from ieg.dataset_utils.utils import autoaug_batch_process_map_fn
from ieg.models import networks
from ieg.models.basemodel import BaseModel
from ieg.models.custom_ops import logit_norm
from ieg.models.custom_ops import MixMode

import numpy as np
import tensorflow.compat.v1 as tf
from tqdm import tqdm

FLAGS = flags.FLAGS
logging = tf.logging

class IEG(BaseModel):
  """Model training class."""

  def __init__(self, sess, strategy, dataset):
    super(IEG, self).__init__(sess, strategy, dataset)
    logging.info('Init IEG model')

    self.augment = MixMode()
    self.beta = 0.5  # MixUp hyperparam
    self.nu = 2      # K value for label guessing

  def set_input(self):
    if len(self.dataset.train_dataflow.output_shapes[0]) == 3:
      # Use for cifar
      train_ds = self.dataset.train_dataflow.shuffle(
          buffer_size=self.batch_size * 10).repeat().batch(
              self.batch_size, drop_remainder=True
          ).map(
              # strong augment each batch data and expand to 5D [Bx2xHxWx3]
              autoaug_batch_process_map_fn,
              num_parallel_calls=tf.data.experimental.AUTOTUNE).prefetch(
                  buffer_size=tf.data.experimental.AUTOTUNE)
    else:
      train_ds = self.dataset.train_dataflow.shuffle(
          buffer_size=self.batch_size * 10).repeat().batch(
              self.batch_size, drop_remainder=True).prefetch(
                  buffer_size=tf.data.experimental.AUTOTUNE)
    # no shuffle for probe, so a batch is class balanced.
    probe_ds = self.dataset.probe_dataflow.repeat().batch(
        self.batch_size, drop_remainder=True).prefetch(
            buffer_size=tf.data.experimental.AUTOTUNE)

    val_ds = self.dataset.val_dataflow.batch(
        FLAGS.val_batch_size, drop_remainder=False).prefetch(
            buffer_size=tf.data.experimental.AUTOTUNE)

    self.train_input_iterator = (
        self.strategy.experimental_distribute_dataset(
            train_ds).make_initializable_iterator())
    self.probe_input_iterator = (
        self.strategy.experimental_distribute_dataset(
            probe_ds).make_initializable_iterator())

    self.eval_input_iterator = (
        self.strategy.experimental_distribute_dataset(
            val_ds).make_initializable_iterator())

  def meta_momentum_update(self, grad, var_name, optimizer):
    # Finds corresponding momentum of a var name
    accumulation = utils.get_var(optimizer.variables(), var_name.split(':')[0])
    if len(accumulation) != 1:
      raise ValueError('length of accumulation {}'.format(len(accumulation)))
    new_grad = tf.math.add(
        tf.stop_gradient(accumulation[0]) * FLAGS.meta_momentum, grad)
    return new_grad

  def guess_label(self, logit, temp=0.5):
    logit = tf.reshape(logit, [-1, self.dataset.num_classes])
    logit = tf.split(logit, self.nu, axis=0)
    logit = [logit_norm(x) for x in logit]
    logit = tf.concat(logit, 0)
    ## Done with logit norm
    p_model_y = tf.reshape(
        tf.nn.softmax(logit), [self.nu, -1, self.dataset.num_classes])
    p_model_y = tf.reduce_mean(p_model_y, axis=0)

    p_target = tf.pow(p_model_y, 1.0 / temp)
    p_target /= tf.reduce_sum(p_target, axis=1, keepdims=True)

    return p_target

  def crossentropy_minimize(self,
                            u_logits,
                            u_images,
                            l_images,
                            l_labels,
                            u_labels=None):
    """Cross-entropy optimization step implementation for TPU."""
    batch_size = self.batch_size // self.strategy.num_replicas_in_sync
    guessed_label = self.guess_label(u_logits)
    self.guessed_label = guessed_label

    guessed_label = tf.reshape(
        tf.stop_gradient(guessed_label), shape=(-1, self.dataset.num_classes))

    l_labels = tf.reshape(
        tf.one_hot(l_labels, self.dataset.num_classes),
        shape=(-1, self.dataset.num_classes))
    augment_images, augment_labels = self.augment(
        [l_images, u_images], [l_labels] + [guessed_label] * self.nu,
        [self.beta, self.beta])
    logit = self.net(augment_images, name='model', training=True)

    zbs = batch_size * 2
    halfzbs = batch_size

    split_pos = [tf.shape(l_images)[0], halfzbs, halfzbs]

    logit = [logit_norm(lgt) for lgt in tf.split(logit, split_pos, axis=0)]
    u_logit = tf.concat(logit[1:], axis=0)

    split_pos = [tf.shape(l_images)[0], zbs]
    l_augment_labels, u_augment_labels = tf.split(
        augment_labels, split_pos, axis=0)

    u_loss = tf.losses.softmax_cross_entropy(u_augment_labels, u_logit)
    l_loss = tf.losses.softmax_cross_entropy(l_augment_labels, logit[0])

    loss = tf.math.add(
        l_loss, u_loss * FLAGS.ce_factor, name='crossentropy_minimization_loss')

    return loss

  def consistency_loss(self, logit, aug_logit):

    def kl_divergence(q_logits, p_logits):
      q = tf.nn.softmax(q_logits)
      per_example_kl_loss = q * (
          tf.nn.log_softmax(q_logits) - tf.nn.log_softmax(p_logits))
      return tf.reduce_mean(per_example_kl_loss) * self.dataset.num_classes

    return tf.math.multiply(
        kl_divergence(tf.stop_gradient(logit), aug_logit),
        FLAGS.consistency_factor,
        name='consistency_loss')

  def unsupervised_loss(self):
    """Creates unsupervised losses.

    Here we create two cross-entropy losses and a KL-loss defined in the paper.

    Returns:
      A list of losses.
    """

    if FLAGS.ce_factor == 0 and FLAGS.consistency_factor == 0:
      return [tf.constant(0, tf.float32), tf.constant(0, tf.float32)]
    logits = self.logits
    images = self.images
    aug_images = self.aug_images
    probe_images, probe_labels = self.probe_images, self.probe_labels
    im_shape = (-1, int(probe_images.shape[1]), int(probe_images.shape[2]),
                int(probe_images.shape[3]))

    aug_logits = self.net(aug_images, name='model', training=True)

    n_probe_to_mix = tf.shape(aug_images)[0]
    probe = tf.tile(tf.constant([[10.]]), [1, tf.shape(probe_images)[0]])
    idx = tf.squeeze(tf.random.categorical(probe, n_probe_to_mix))

    l_images = tf.reshape(tf.gather(probe_images, idx), im_shape)
    l_labels = tf.reshape(tf.gather(probe_labels, idx), (-1,))

    u_logits = tf.concat([logits, aug_logits], axis=0)
    u_images = tf.concat([images, aug_images], axis=0)

    losses = []
    if FLAGS.ce_factor > 0:
      logging.info('Use crossentropy minimization loss {}'.format(
          FLAGS.ce_factor))
      ce_min_loss = self.crossentropy_minimize(u_logits, u_images, l_images,
                                               l_labels)
      losses.append(ce_min_loss)
    else:
      losses.append(tf.constant(0, tf.float32))

    if FLAGS.consistency_factor > 0:
      logging.info('Use consistency loss {}'.format(
          FLAGS.consistency_factor))
      consis_loss = self.consistency_loss(logits, aug_logits)
      losses.append(consis_loss)

    else:
      losses.append(tf.constant(0, tf.float32))

    return losses

  def meta_optimize(self):
    """Meta optimization step."""

    probe_images, probe_labels = self.probe_images, self.probe_labels
    labels = self.labels
    net = self.net
    logits = self.logits
    gate_gradients = 1

    batch_size = int(self.batch_size / self.strategy.num_replicas_in_sync)
    init_eps_val = float(1) / batch_size

    meta_net = networks.MetaImage(self.net, name='meta_model')

    if FLAGS.meta_momentum and not self.optimizer.variables():
      # Initializing momentum state of optimizer for meta momentum update.
      # It is a hacky implementation
      logging.info('Pre-initialize optimizer momentum states.')
      idle_net_cost = tf.losses.sparse_softmax_cross_entropy(
          self.labels, logits)
      tmp_var_grads = self.optimizer.compute_gradients(
          tf.reduce_mean(idle_net_cost), net.trainable_variables)
      self.optimizer.apply_gradients(tmp_var_grads)

    with tf.name_scope('coefficient'):
      # Data weight coefficient
      target = tf.constant(
          [init_eps_val] * batch_size,
          shape=(batch_size,),
          dtype=np.float32,
          name='weight')
      # Data re-labeling coefficient
      eps = tf.constant(
          [FLAGS.grad_eps_init] * batch_size,
          shape=(batch_size,),
          dtype=tf.float32,
          name='eps')

    onehot_labels = tf.one_hot(labels, self.dataset.num_classes)
    onehot_labels = tf.cast(onehot_labels, tf.float32)
    eps_k = tf.reshape(eps, [batch_size, 1])

    mixed_labels = eps_k * onehot_labels + (1 - eps_k) * self.guessed_label
    # raw softmax loss
    log_softmax = tf.nn.log_softmax(logits)
    net_cost = -tf.reduce_sum(mixed_labels * log_softmax, 1)

    lookahead_loss = tf.reduce_sum(tf.multiply(target, net_cost))
    lookahead_loss = lookahead_loss + net.regularization_loss

    with tf.control_dependencies([lookahead_loss]):
      train_vars = net.trainable_variables
      var_grads = tf.gradients(
          lookahead_loss, train_vars, gate_gradients=gate_gradients)

      static_vars = []
      for i in range(len(train_vars)):
        if FLAGS.meta_momentum > 0:
          actual_grad = self.meta_momentum_update(var_grads[i],
                                                  train_vars[i].name,
                                                  self.optimizer)
          static_vars.append(
              tf.math.subtract(train_vars[i],
                               FLAGS.meta_stepsize * actual_grad))
        else:
          static_vars.append(
              tf.math.subtract(train_vars[i],
                               FLAGS.meta_stepsize * var_grads[i]))
        # new style
        meta_net.add_variable_alias(
            static_vars[-1], var_name=train_vars[i].name)

      for uv in net.updates_variables:
        meta_net.add_variable_alias(
            uv, var_name=uv.name, var_type='updates_variables')
      meta_net.verbose()

    with tf.control_dependencies(static_vars):
      g_logits = meta_net(
          probe_images, name='meta_model', reuse=True, training=True)

      desired_y = tf.one_hot(probe_labels, self.dataset.num_classes)
      meta_loss = tf.nn.softmax_cross_entropy_with_logits_v2(
          desired_y, g_logits)
      meta_loss = tf.reduce_mean(meta_loss, name='meta_loss')
      meta_loss = meta_loss + meta_net.get_regularization_loss(net.wd)
      meta_acc, meta_acc_op = tf.metrics.accuracy(probe_labels,
                                                  tf.argmax(g_logits, axis=1))

    with tf.control_dependencies([meta_loss] + [meta_acc_op]):
      meta_train_vars = meta_net.trainable_variables
      grad_meta_vars = tf.gradients(
          meta_loss, meta_train_vars, gate_gradients=gate_gradients)
      grad_target, grad_eps = tf.gradients(
          static_vars, [target, eps],
          grad_ys=grad_meta_vars,
          gate_gradients=gate_gradients)
    # updates weight
    raw_weight = target - grad_target
    raw_weight = raw_weight - init_eps_val
    unorm_weight = tf.clip_by_value(
        raw_weight, clip_value_min=0, clip_value_max=float('inf'))
    norm_c = tf.reduce_sum(unorm_weight)
    weight = tf.divide(unorm_weight, norm_c + 0.00001)

    # gets new lambda by the sign of gradient
    new_eps = tf.where(grad_eps < 0, x=tf.ones_like(eps), y=tf.zeros_like(eps))

    return tf.stop_gradient(weight), tf.stop_gradient(
        new_eps), meta_loss, meta_acc

  def train_step(self):

    def step_fn(inputs):
      """Step functon.

      Args:
        inputs: inputs from data iterator

      Returns:
        a set of variables want to observe in Tensorboard
      """
      net = self.net
      (all_images, labels), (self.probe_images, self.probe_labels) = inputs
      assert len(all_images.shape) == 5
      images, self.aug_images = all_images[:, 0], all_images[:, 1]

      self.images, self.labels = images, labels
      batch_size = int(self.batch_size / self.strategy.num_replicas_in_sync)

      logits = net(images, name='model', reuse=tf.AUTO_REUSE, training=True)
      self.logits = logits

      # other losses
      # initialized first to use self.guessed_label for meta step
      xe_loss, cs_loss = self.unsupervised_loss()

      # meta optimization
      weight, eps, meta_loss, meta_acc = self.meta_optimize()

      ## losses w.r.t new weight and loss
      onehot_labels = tf.one_hot(labels, self.dataset.num_classes)
      onehot_labels = tf.cast(onehot_labels, tf.float32)
      eps_k = tf.reshape(eps, [batch_size, 1])

      mixed_labels = tf.math.add(
          eps_k * onehot_labels, (1 - eps_k) * self.guessed_label,
          name='mixed_labels')
      net_cost = tf.losses.softmax_cross_entropy(
          mixed_labels, logits, reduction=tf.losses.Reduction.NONE)
      # loss with initial weight
      net_loss1 = tf.reduce_mean(net_cost)

      # loss with initial eps
      init_eps = tf.constant(
          [FLAGS.grad_eps_init] * batch_size, dtype=tf.float32)
      init_eps = tf.reshape(init_eps, (-1, 1))
      init_mixed_labels = tf.math.add(
          init_eps * onehot_labels, (1 - init_eps) * self.guessed_label,
          name='init_mixed_labels')

      net_cost2 = tf.losses.softmax_cross_entropy(
          init_mixed_labels, logits, reduction=tf.losses.Reduction.NONE)
      net_loss2 = tf.reduce_sum(tf.math.multiply(net_cost2, weight))

      net_loss = (net_loss1 + net_loss2) / 2

      net_loss = net_loss + tf.add_n([xe_loss, cs_loss])
      net_loss += net.regularization_loss
      net_loss /= self.strategy.num_replicas_in_sync

      # rescale by gpus
      with tf.control_dependencies(net.updates):
        net_grads = tf.gradients(net_loss, net.trainable_variables)
        minimizer_op = self.optimizer.apply_gradients(
            zip(net_grads, net.trainable_variables),
            global_step=self.global_step)

      with tf.control_dependencies([minimizer_op]):
        train_op = self.ema.apply(net.trainable_variables)

      acc_op, acc_update_op = self.acc_func(labels, tf.argmax(logits, axis=1))

      with tf.control_dependencies([train_op, acc_update_op]):
        return (tf.identity(net_loss), tf.identity(xe_loss),
                tf.identity(cs_loss), tf.identity(meta_loss),
                tf.identity(meta_acc), tf.identity(acc_op), tf.identity(weight),
                tf.identity(labels))

    # end of parallel
    (pr_net_loss, pr_xe_loss, pr_cs_loss, pr_metaloss, pr_metaacc, pr_acc,
     pr_weight, pr_labels) = self.strategy.run(
         step_fn,
         args=((next(self.train_input_iterator),
                next(self.probe_input_iterator)),))
    # collect device variables
    weights = self.strategy.unwrap(pr_weight)
    weights = tf.concat(weights, axis=0)
    labels = self.strategy.unwrap(pr_labels)
    labels = tf.concat(labels, axis=0)

    mean_acc = self.strategy.reduce(tf.distribute.ReduceOp.MEAN, pr_acc)
    mean_metaacc = self.strategy.reduce(tf.distribute.ReduceOp.MEAN, pr_metaacc)
    net_loss = self.strategy.reduce(tf.distribute.ReduceOp.MEAN, pr_net_loss)
    xe_loss = self.strategy.reduce(tf.distribute.ReduceOp.MEAN, pr_xe_loss)
    cs_loss = self.strategy.reduce(tf.distribute.ReduceOp.MEAN, pr_cs_loss)
    meta_loss = self.strategy.reduce(tf.distribute.ReduceOp.MEAN, pr_metaloss)

    # The following add variables for tensorboard visualization
    merges = []
    merges.append(tf.summary.scalar('acc/train', mean_acc))
    merges.append(tf.summary.scalar('loss/xemin', xe_loss))
    merges.append(tf.summary.scalar('loss/consistency', cs_loss))
    merges.append(tf.summary.scalar('loss/net', net_loss))
    merges.append(tf.summary.scalar('loss/meta', meta_loss))
    merges.append(tf.summary.scalar('acc/meta', mean_metaacc))
    merges.append(
        tf.summary.scalar('acc/eval_on_train', self.eval_acc_on_train[0]))
    merges.append(
        tf.summary.scalar('acc/eval_on_train_top5', self.eval_acc_on_train[1]))
    merges.append(tf.summary.scalar('acc/num_eval', self.eval_acc_on_train[2]))

    zw_inds = tf.squeeze(
        tf.where(tf.less_equal(weights, 0), name='zero_weight_index'))
    merges.append(
        tf.summary.scalar(
            'weights/zeroratio',
            tf.math.divide(
                tf.cast(tf.size(zw_inds), tf.float32),
                tf.cast(tf.size(weights), tf.float32))))

    self.epoch_var = tf.cast(
        self.global_step / self.iter_epoch, tf.float32, name='epoch')
    merges.append(tf.summary.scalar('epoch', self.epoch_var))
    merges.append(tf.summary.scalar('learningrate', self.learning_rate))
    summary = tf.summary.merge(merges)

    return [
        net_loss, meta_loss, xe_loss, cs_loss, mean_acc, mean_metaacc, summary,
        weights
    ]

  def train(self):
    self.set_input()
    self.build_graph()

    with self.strategy.scope():
      self.initialize_variables()

      self.sess.run([
          self.train_input_iterator.initializer,
          self.probe_input_iterator.initializer
      ])
      self.sess.run([self.eval_input_iterator.initializer])

      logging.info('Finish variable initialization')
      iter_epoch = self.iter_epoch

      self.saver = tf.train.Saver(max_to_keep=4)
      self.load_model()
      FLAGS.restore_step = self.global_step.eval()

      pbar = tqdm(total=(FLAGS.max_iteration - FLAGS.restore_step))
      for iteration in range(FLAGS.restore_step + 1, FLAGS.max_iteration + 1):
        self.update_learning_rate(iteration)
        (lr, net_loss, meta_loss, xe_loss, cs_loss, acc, meta_acc,
         merged_summary, weights) = (
             self.sess.run([self.learning_rate] + self.train_op))
        pbar.update(1)
        message = ('Epoch {}[{}/{}] lr{:.3f} meta_loss:{:.2f} loss:{:.2f} '
                   'mc_loss:{:.2f} uc_loss:{:.2f} weight{:.2f}({:.2f}) '
                   'acc:{:.2f} mata_acc{:.2f}').format(iteration // iter_epoch,
                                                       iteration % iter_epoch,
                                                       iter_epoch, lr,
                                                       float(meta_loss),
                                                       float(net_loss),
                                                       float(xe_loss),
                                                       float(cs_loss),
                                                       float(np.mean(weights)),
                                                       float(np.std(weights)),
                                                       float(acc),
                                                       float(meta_acc))
        pbar.set_description(message)
        self.summary_writer.add_summary(merged_summary, iteration)

        # checkpoint
        if self.time_for_evaluation(iteration, lr):
          logging.info(message)
          self.evaluate(iteration, lr)
          self.save_model(iteration)
          self.summary_writer.flush()
      # end of iterations
      pbar.close()

max_stars_repo_name: deepneuralmachine/google-research

max_issues_repo_name: deepneuralmachine/google-research

max_forks_repo_name: deepneuralmachine/google-research


text:
#author:vatsal verma
import sys
import sorting as s
import visualizer as vs
import numpy as np
import random

def get_random_array(length):
    n = list(range(length))
    random.shuffle(n)
    return n

array_size = 50

algorithms = [s.bubble_sort, s.heap_sort, s.selection_sort,
              s.insertion_sort, s.quick_sort, s.insertion_sort, s.merge_sort]

if len(sys.argv) < 2:
    print('Usage:', 'python', sys.argv[0], 'function_name', '\n')
    print('Choose from:', '\n')
    for algorithm in algorithms:
        print(algorithm.__name__)
    print('\n')
    sys.exit(0)

algorithms = list(filter(lambda a: a.__name__ == sys.argv[1], algorithms))
if (len(algorithms) == 0):
    print('Method does not exist.')
else:
    algorithm = algorithms[0]
    arr = s.Array(get_random_array(array_size))
    algorithm(arr)
    vs.show()

max_stars_repo_name: vatsalcode/Sorting-Algo-visulizer

max_issues_repo_name: vatsalcode/Sorting-Algo-visulizer

max_forks_repo_name: vatsalcode/Sorting-Algo-visulizer


text:
import joblib

try:
    import Tkinter as tk
except:
    import tkinter as tk
    from tkinter import *
    from tkinter.ttk import *
    from tkinter import messagebox
    from PIL import Image, ImageTk, ImageGrab
    import tkinter.font as font
    import pandas as pd
    import PIL
    from PIL import Image, ImageOps
    import numpy as np
    import sys
    import os, cv2
    import csv
    from tkinter import filedialog

    myDir = "..\GujOCR\Output"

    class SampleApp(tk.Tk):
        def __init__(self):
            tk.Tk.__init__(self)
            self._frame = None
            self.switch_frame(StartPage)

        def switch_frame(self, frame_class):
            new_frame = frame_class(self)
            if self._frame is not None:
                self._frame.destroy()
            self._frame = new_frame
            self._frame.pack()

    class StartPage(tk.Frame):
        def __init__(self, master):
            tk.Frame.__init__(self, master)

            def prep():
                # Useful function
                def createFileList(myDir, format='.png'):
                    fileList = []
                    print(myDir)
                    for root, dirs, files in os.walk(myDir, topdown=False):
                        for name in files:
                            if name.endswith(format):
                                fullName = os.path.join(root, name)
                                fileList.append(fullName)
                    return fileList

                columnNames = list()
                for i in range(784):
                    pixel = 'p'
                    pixel += str(i)
                    columnNames.append(pixel)
                l = os.listdir("..\GujOCR\Output")
                print(l)

                dic = {val: idx for idx, val in enumerate(l)}
                print(dic)

                train_data = pd.DataFrame(columns=columnNames)
                train_data.to_csv("train.csv", index=False)
                label_count = list()

                print(len(l))

                for i in range(len(l)):
                    mydir = 'OUTPUT/' + l[i]
                    fileList = createFileList(mydir)
                    for file in fileList:
                        img_file = Image.open(file)  # imgfile.show()
                        width, height = img_file.size
                        format = img_file.format
                        mode = img_file.mode

                        label_count.append(dic[l[i]])
                        inverted_image = img_file.convert('RGB')
                        im_invert = ImageOps.invert(inverted_image)
                        size = (28, 28)
                        new_image = img_file.resize(size)

                        img_grey = new_image.convert('L')
                        value = np.asarray(img_grey.getdata(), dtype=np.int).reshape(
                            (img_grey.size[1], img_grey.size[0]))
                        value = value.flatten()
                        with open("train.csv", 'a', newline='') as f:
                            writer = csv.writer(f)
                            writer.writerow(value)

                read_data = pd.read_csv('train.csv')
                read_data['Label'] = label_count
                print(read_data)
                messagebox.showinfo("completed successfully.", "Training data saved in train.csv")

            master.title("ગુજરાતી Handwritten Character recognition")
            master.resizable(0, 0)
            master.geometry("846x400")
            master.configure(bg='black')
            tk.Label(text="Gujarati Handwritten Character Recognition", fg="#e8edf3", bg="#22264b", width=44, height=2,font=("Courier", 23), anchor="center",relief="ridge").place_configure(x=2, y=2)
            canvas1 = Canvas(master, width=530, height=210, bg='#5A79A5')
            tk.Label(text="Main Menu", fg="#e8edf3", bg="#22264b", width=44, height=2,font=("Courier", 15), anchor="center").place_configure(x=145, y=85)
            canvas = Canvas(master, width=100, height=100, bg="white")
            load = Image.open("ka2.png")
            load = load.resize((100, 100), Image.ANTIALIAS)
            ren = ImageTk.PhotoImage(load)
            img = tk.Label(self, image=ren)
            img.image = ren
            l1 = tk.Label(canvas,image=ren,bg="black")
            b1 = tk.Button(canvas1, text="1. Dataset Collection", bg="#22264b", fg="#e6cf8b",font=("Courier", 15),command=lambda: master.switch_frame(Collect1))
            b2 = tk.Button(canvas1, text="2. Preprocessing the dataset", font=('Courier', 15), bg="#22264b", fg="#e6cf8b",command=prep)
            b3 = tk.Button(canvas1, text="3. Train the model and calculate accuracy", font=('Courier', 15), bg="#22264b",fg="#e6cf8b",command=lambda: master.switch_frame(Train))
            b4 = tk.Button(canvas1, text="4. Prediction", font=('Courier', 15), bg="#22264b", fg="#e6cf8b",command=lambda: master.switch_frame(Predict))
            canvas1.place(x=145, y=150)
            canvas.place(x=10,y=80)
            b1.place(x=15, y=20)
            b2.place(x=15, y=65)
            b3.place(x=15, y=110)
            b4.place(x=15, y=155)
            l1.place(x=0,y=0)

    class Collect1(tk.Frame):
        def __init__(self, master):
            tk.Frame.__init__(self, master)
            tk.Frame.configure(self, bg="gray")
            tk.Label(self, text="", font=('Helvetica', 18, "bold")).pack(side="top", padx=750, pady=500)
            tk.Label(self,text="Gujarati Handwritten Character Recognition", fg="#e8edf3", bg="#22264b", width=44, height=2,font=("Courier", 23), anchor="center",relief="ridge").place_configure(x=2, y=2)
            tk.Label(self, text="Collection of dataset", fg="#e8edf3", bg="#22264b", width=44, height=2,font=("Courier", 15), anchor="center").place_configure(x=145, y=85)
            canvas1 = Canvas(master, width=530, height=210, bg='black')
            b1 = tk.Button(canvas1, text="1. Runtime Dataset upgradation with paint", bg="#22264b", fg="#e6cf8b", font=("Courier", 15),command=lambda: master.switch_frame(Collect2))
            b2 = tk.Button(canvas1, text="2. Upload file", font=('Courier', 15), bg="#22264b", fg="#e6cf8b",command=lambda: master.switch_frame(Collect3))
            canvas2 = Canvas(master, width=50, height=50, bg='black')

            load = Image.open("back3.png")
            load = load.resize((40, 40), Image.ANTIALIAS)
            render = ImageTk.PhotoImage(load)

            img = tk.Label(self, image=render)
            img.image = render
            b3 = tk.Button(canvas2, text="BACK", font=('Courier', 10), image=render, bg="#22264b", fg="#e6cf8b",
                           height=40,
                           width=40, command=lambda: master.switch_frame(StartPage))
            canvas2.place(x=20, y=310)

            canvas1.place(x=145, y=150)
            b1.place(x=15, y=20)
            b2.place(x=15, y=65)
            b3.place(x=4, y=4)
    class Collect2(tk.Frame):
        def __init__(self, master):
            tk.Frame.__init__(self, master)
            tk.Frame.configure(self, bg="gray")
            tk.Label(self, text="", font=('Helvetica', 18, "bold")).pack(side="top", padx=750, pady=500)
            tk.Label(self,text="Gujarati Handwritten Character Recognition", fg="#e8edf3", bg="#22264b", width=44, height=2,font=("Courier", 23), anchor="center",relief="ridge").place_configure(x=2, y=2)
            tk.Label(self, text="Collection of dataset", fg="#e8edf3", bg="#22264b", width=44, height=2,font=("Courier", 15), anchor="center").place_configure(x=145, y=85)

            def screen_capture():
                import pyscreenshot as ImageGrab
                import time
                import os
                os.startfile("C:/ProgramData/Microsoft/Windows/Start Menu/Programs/Accessories/Paint")
                s1 = t1.get()
                os.chdir("C:/Users/helid/PycharmProjects/GujOCR/folder")
                images_folder = "C:/Users/helid/PycharmProjects/GujOCR/folder/" + s1 + "/"
                time.sleep(5)
                for i in range(0, 1):
                    time.sleep(8)
                    im = ImageGrab.grab(bbox=(60, 170, 400, 550))  # x1,y1,x2,y2
                    print("saved......", i)
                    im.save(images_folder + str(i) + '.png')
                    print("clear screen now and redraw now........")
                messagebox.showinfo("Result", "Capturing screen is completed!!")

            canvas1 = Canvas(master, width=530, height=210, bg='black')
            l1 = tk.Label(canvas1, text="Enter Character: ", font=('Courier', 15), bg='black', fg="#e6cf8b")
            t1 = tk.Entry(canvas1, width=20, border=5)
            b1 = tk.Button(canvas1, text="Submit and Collect with Paint", bg="#22264b", fg="#e6cf8b", font=("Courier", 15), command=screen_capture)
            canvas2 = Canvas(master, width=50,height=50,bg='black')
            load = Image.open("back3.png")
            load = load.resize((40, 40), Image.ANTIALIAS)
            render = ImageTk.PhotoImage(load)

            img = tk.Label(self, image=render)
            img.image = render
            b3 = tk.Button(canvas2, text="BACK", font=('Courier', 10),image=render, bg="#22264b", fg="#e6cf8b",height=40,
                      width=40,command=lambda: master.switch_frame(Collect1))
            canvas2.place(x=20,y=310)

            canvas1.place(x=145, y=150)
            l1.place(x=75, y=5)
            t1.place(x=290, y=5)
            b1.place(x=75,y=45)
            b3.place(x=4,y=4)

    class Collect3(tk.Frame):
        def __init__(self, master):
            tk.Frame.__init__(self, master)
            tk.Frame.configure(self, bg="gray")
            tk.Label(self, text="", font=('Helvetica', 18, "bold")).pack(side="top", padx=750, pady=500)
            tk.Label(self,text="Gujarati Handwritten Character Recognition", fg="#e8edf3", bg="#22264b", width=44, height=2,font=("Courier", 23), anchor="center",relief="ridge").place_configure(x=2, y=2)
            tk.Label(self, text="Collection of dataset", fg="#e8edf3", bg="#22264b", width=44, height=2,font=("Courier", 15), anchor="center").place_configure(x=145, y=85)
            canvas1 = Canvas(master, width=530, height=210, bg='black')
            l3 = tk.Label(canvas1, text="Upload your boxed image:", font=('Courier', 15), bg='black', fg="#e6cf8b")

            def UploadAction(event=None):
                filename = filedialog.askopenfilename()
                print('Selected:', filename)

                def sort_contours(cnts, method="left-to-right"):
                    reverse = False
                    i = 0
                    if method == "right-to-left" or method == "bottom-to-top":
                        reverse = True
                    if method == "top-to-bottom" or method == "bottom-to-top":
                        i = 1
                    boundingBoxes = [cv2.boundingRect(c) for c in cnts]
                    (cnts, boundingBoxes) = zip(
                        *sorted(zip(cnts, boundingBoxes), key=lambda b: b[1][i], reverse=reverse))
                    return (cnts, boundingBoxes)
                def box_extraction(img_for_box_extraction_path, cropped_dir_path):
                    print("Reading image..")
                    img = cv2.imread(img_for_box_extraction_path, 0)
                    (thresh, img_bin) = cv2.threshold(img, 128, 255,cv2.THRESH_BINARY | cv2.THRESH_OTSU)
                    img_bin = 255 - img_bin
                    print("Storing binary image to Images/Image_bin.jpg..")
                    cv2.imwrite("Images/Image_bin.jpg", img_bin)
                    print("Applying Morphological Operations..")
                    kernel_length = np.array(img).shape[1] // 40
                    verticle_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (1, kernel_length))
                    hori_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (kernel_length, 1))
                    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
                    img_temp1 = cv2.erode(img_bin, verticle_kernel, iterations=3)
                    verticle_lines_img = cv2.dilate(img_temp1, verticle_kernel, iterations=3)
                    cv2.imwrite("Images/verticle_lines.jpg", verticle_lines_img)
                    img_temp2 = cv2.erode(img_bin, hori_kernel, iterations=3)
                    horizontal_lines_img = cv2.dilate(img_temp2, hori_kernel, iterations=3)
                    cv2.imwrite("Images/horizontal_lines.jpg", horizontal_lines_img)
                    alpha = 0.5
                    beta = 1.0 - alpha
                    img_final_bin = cv2.addWeighted(verticle_lines_img, alpha, horizontal_lines_img, beta, 0.0)
                    img_final_bin = cv2.erode(~img_final_bin, kernel, iterations=2)
                    (thresh, img_final_bin) = cv2.threshold(img_final_bin, 128, 255,cv2.THRESH_BINARY | cv2.THRESH_OTSU)
                    print("Binary image which only contains boxes: Images/img_final_bin.jpg")
                    cv2.imwrite("Images/img_final_bin.jpg", img_final_bin)
                    contours, hierarchy = cv2.findContours(
                        img_final_bin, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
                    (contours, boundingBoxes) = sort_contours(contours, method="top-to-bottom")
                    print("Output stored in Output directiory!")
                    idx = 0
                    for c in contours:
                        x, y, w, h = cv2.boundingRect(c)
                        if (w > 22 and h > 22) and w >= h:
                            idx += 1
                            new_img = img[y:y + h, x:x + w]
                            cv2.imwrite("./Output/" + str(idx) + '.png', new_img)
                    cv2.drawContours(img, contours, -1, (0, 0, 255), 3)
                    cv2.imwrite("./Temp/img_contour.jpg", img)
                    messagebox.showinfo("Result", "Data segregated from the box image and stored in OUTPUT directory")
                box_extraction(filename, "./Output/")

            b2 = tk.Button(canvas1, text='Open..',bg="#22264b", fg="#e6cf8b", font=("Courier", 12), command=UploadAction)
            canvas2 = Canvas(master, width=50, height=50, bg='black')
            load = Image.open("back3.png")
            load = load.resize((40, 40), Image.ANTIALIAS)
            render = ImageTk.PhotoImage(load)
            img = tk.Label(self, image=render)
            img.image = render
            b3 = tk.Button(canvas2, text="BACK", font=('Courier', 10), image=render, bg="#22264b", fg="#e6cf8b",height=40, width=40, command=lambda: master.switch_frame(StartPage))

            b3.place(x=4, y=4)
            canvas2.place(x=20, y=310)
            canvas1.place(x=145, y=150)
            l3.place(x=75,y=5)
            b2.place(x=290,y=5)
            b3.place(x=4, y=4)

    class Train(tk.Frame):
        def __init__(self, master):
            tk.Frame.__init__(self, master)
            tk.Frame.configure(self, bg="gray")
            tk.Label(self, text="", font=('Helvetica', 18, "bold")).pack(side="top", padx=750, pady=500)
            tk.Label(self,text="Gujarati Handwritten Character Recognition", fg="#e8edf3", bg="#22264b", width=44, height=2,font=("Courier", 23), anchor="center",relief="ridge").place_configure(x=2, y=2)
            tk.Label(self, text="Train And Model Data", fg="#e8edf3", bg="#22264b", width=44, height=2,font=("Courier", 15), anchor="center").place_configure(x=145, y=85)
            canvas1 = Canvas(master, width=530, height=210, bg='black')
            canvas2 = Canvas(master, width=50, height=50, bg='black')
            load = Image.open("back3.png")
            load = load.resize((40, 40), Image.ANTIALIAS)
            render = ImageTk.PhotoImage(load)

            img = tk.Label(self, image=render)
            img.image = render
            b3 = tk.Button(canvas2, text="BACK", font=('Courier', 10), image=render, bg="#22264b", fg="#e6cf8b",height=40,width=40, command=lambda: master.switch_frame(StartPage))
            b3.place(x=4, y=4)
            canvas2.place(x=20, y=310)
            canvas1.place(x=145, y=150)

    class Predict(tk.Frame):
        def __init__(self, master):
            tk.Frame.__init__(self, master)
            tk.Frame.configure(self, bg="gray")
            tk.Label(self, text="", font=('Helvetica', 18, "bold")).pack(side="top", padx=750, pady=500)
            tk.Label(self, text="Gujarati Handwritten Character Recognition", fg="#e8edf3", bg="#22264b",
                     width=44, height=2, font=("Courier", 23), anchor="center", relief="ridge").place_configure(
                x=2, y=2)
            tk.Label(self, text="Prediction of your data", fg="#e8edf3", bg="#22264b", width=44, height=2,
                     font=("Courier", 15), anchor="center").place_configure(x=145, y=85)
            canvas1 = Canvas(master, width=530, height=210, bg='black')
            canvas2 = Canvas(master, width=50, height=50, bg='black')
            b1 = tk.Button(canvas1, text="1. Live Prediction using Blackboard", bg="#22264b", fg="#e6cf8b", font=("Courier", 15), command=lambda: master.switch_frame(Predict2))
            b2 = tk.Button(canvas1, text="2. Upload file and Predict", font=('Courier', 15), bg="#22264b", fg="#e6cf8b", command=lambda: master.switch_frame(Predict3))
            load = Image.open("back3.png")
            load = load.resize((40, 40), Image.ANTIALIAS)
            render = ImageTk.PhotoImage(load)

            img = tk.Label(self, image=render)
            img.image = render
            b3 = tk.Button(canvas2, text="BACK", font=('Courier', 10), image=render, bg="#22264b", fg="#e6cf8b",height=40, width=40, command=lambda: master.switch_frame(StartPage))
            b3.place(x=4, y=4)
            canvas2.place(x=20, y=310)
            canvas1.place(x=145, y=150)
            b1.place(x=15, y=20)
            b2.place(x=15, y=65)
            b3.place(x=4, y=4)

    class Predict2(tk.Frame):
        def __init__(self, master):
            tk.Frame.__init__(self, master)
            tk.Frame.configure(self, bg="gray")
            tk.Label(self, text="", font=('Helvetica', 18, "bold")).pack(side="top", padx=750, pady=500)
            tk.Label(self,text="Gujarati Handwritten Character Recognition", fg="#e8edf3", bg="#22264b", width=44, height=2,font=("Courier", 23), anchor="center",relief="ridge").place_configure(x=2, y=2)
            tk.Label(self, text="Live Prediction using Blackboard", fg="#e8edf3", bg="#22264b", width=44, height=2,font=("Courier", 15), anchor="center").place_configure(x=145, y=85)
            canvas1 = Canvas(master, width=530, height=210, bg='black')
            canvas2 = Canvas(master, width=50, height=50, bg='black')
            load = Image.open("back3.png")
            load = load.resize((40, 40), Image.ANTIALIAS)
            render = ImageTk.PhotoImage(load)

            def prediction():

                model = joblib.load("model/digit_recognizer")

                img = ImageGrab.grab(bbox=(130, 500, 500, 700))
                img.save("paint.png")

                im = cv2.imread("paint.png")
                load = Image.open("paint.png")
                load = load.resize((280, 280))
                photo = ImageTk.PhotoImage(load)

                # Labels can be text or images
                img = Label(canvas3, image=photo, width=280, height=280)
                img.image = photo
                img.place(x=0, y=0)

                im_gray = cv2.cvtColor(im, cv2.COLOR_BGR2GRAY)
                im_gray = cv2.GaussianBlur(im_gray, (15, 15), 0)

                # Threshold the image
                ret, im_th = cv2.threshold(im_gray, 100, 255, cv2.THRESH_BINARY)
                roi = cv2.resize(im_th, (28, 28), interpolation=cv2.INTER_AREA)

                rows, cols = roi.shape

                X = []

                ## Add pixel one by one into data array
                for i in range(rows):
                    for j in range(cols):
                        k = roi[i, j]
                        if k > 100:
                            k = 1
                        else:
                            k = 0
                        X.append(k)

                predictions = model.predict([X])

                a1 = tk.Label(canvas3, text="Prediction= ", font=("verdana", 20))
                a1.place(x=5, y=350)

                b1 = tk.Label(canvas3, text=predictions[0], font=("verdana", 20))
                b1.place(x=200, y=350)

            canvas4 = Canvas(master, width=530, height=210, bg='black')
            canvas4.place(x=145, y=150)

            def activate_paint(e):
                global lastx, lasty
                canvas4.bind('<B1-Motion>', paint)
                lastx, lasty = e.x, e.y

            def paint(e):
                global lastx, lasty
                x, y = e.x, e.y
                canvas4.create_line((lastx, lasty, x, y), width=40, fill="white")
                lastx, lasty = x, y

            canvas4.bind('<1>', activate_paint)

            def clear():
                canvas4.delete("all")

            tk.Label(canvas4, text="Draw in the below blackboard", fg="#e8edf3", bg="#22264b", width=40, height=1,font=("Courier", 15), anchor="center").place_configure(x=50, y=0)
            btn = tk.Button(canvas4, text="clear", fg="white", bg="green", command=clear)
            btn.place(x=0, y=0)

            canvas3 = Canvas(master, width=280, height=530, bg="green")
            # canvas3.place(x=515, y=120)
            img = tk.Label(self, image=render)
            img.image = render

            b3 = tk.Button(canvas2, text="BACK", font=('Courier', 10), image=render, bg="#22264b", fg="#e6cf8b",height=40,width=40, command=lambda: master.switch_frame(Predict))
            b3.place(x=4, y=4)
            canvas2.place(x=20, y=310)
            canvas1.place(x=145, y=150)

    class Predict3(tk.Frame):
        def __init__(self, master):
            tk.Frame.__init__(self, master)
            tk.Frame.configure(self, bg="gray")
            tk.Label(self, text="", font=('Helvetica', 18, "bold")).pack(side="top", padx=750, pady=500)
            tk.Label(self,text="Gujarati Handwritten Character Recognition", fg="#e8edf3", bg="#22264b", width=44, height=2,font=("Courier", 23), anchor="center",relief="ridge").place_configure(x=2, y=2)
            tk.Label(self, text="Upload file and Predict", fg="#e8edf3", bg="#22264b", width=44, height=2,font=("Courier", 15), anchor="center").place_configure(x=145, y=85)

            canvas1 = Canvas(master, width=530, height=210, bg='black')
            l3 = tk.Label(canvas1, text="Upload your image to be predicted:", font=('Courier', 15), bg='black', fg="#e6cf8b")
            def UploadAction(event=None):
                filename = filedialog.askopenfilename()
                print('Selected:', filename)
            b2 = tk.Button(canvas1, text='Open..', bg="#22264b", fg="#e6cf8b", font=("Courier", 12),command=UploadAction)
            canvas2 = Canvas(master, width=50, height=50, bg='black')
            load = Image.open("back3.png")
            load = load.resize((40, 40), Image.ANTIALIAS)
            render = ImageTk.PhotoImage(load)
            img = tk.Label(self, image=render)
            img.image = render
            b3 = tk.Button(canvas2, text="BACK", font=('Courier', 10), image=render, bg="#22264b", fg="#e6cf8b",height=40, width=40, command=lambda: master.switch_frame(Predict))
            b3.place(x=4, y=4)
            canvas2.place(x=20, y=310)
            canvas1.place(x=145, y=150)
            l3.place(x=75, y=5)
            b2.place(x=220, y=40)
            b3.place(x=4, y=4)

if __name__ == "__main__":
    app = SampleApp()
    app.mainloop()


text:
from .clustering_model import ClusteringModel

from scipy.cluster.hierarchy import linkage, fclusterdata
@ClusteringModel.register('hac')
class HAC(ClusteringModel):
    def __init__(self, **parameters):
        super().__init__(False)
        self.params = parameters

    def fit(self, x):
        self.data = x

    def get_labels(self):
        return fclusterdata(self.data, **self.params)

reason: from scipy


text:
import math
import torch
import numpy as np
import operators as op
import stability
import nn
from PIL import Image

with Image.open("test_image.png") as im:
    im.show()
    X = np.asarray(im).astype(np.float64) / 255

# parameters
eta = 1e-1
N, _ = X.shape
inner_iters = 5
outer_iters = 8
mu0 = 1e-1
delta = 1e-12
r = 0.5

# generate sampling mask
mask = torch.zeros((N,N), dtype=bool)
for i in range(-N//2+1,N//2+1):
    for j in range(-N//2+1,N//2+1):
        if (max(abs(i),abs(j)) <= 0.25*(N/2)):
            mask[i+N//2,j+N//2] = True

mask = mask | (torch.rand((N,N)) < 0.1)

n_mask = mask.numpy()

(Image.fromarray(n_mask)).show()

m = torch.sum(mask)

print('Image size (number of pixels):', N*N)
print('Number of measurements:', m)
print('Sample rate:', m/(N*N))

# generate functions for measurement and weight operators
subsampled_ft = lambda x, mode: op.fourier_2d(x,mode,N,mask)*(N/math.sqrt(m))
discrete_grad = lambda x, mode: op.discrete_gradient_2d(x,mode,N,N)
L_grad = 2.0

# define the inverse problem
noise = (eta/math.sqrt(m))*torch.randn(m)
T = torch.from_numpy(np.reshape(X.transpose(),N*N))
y = subsampled_ft(T,1) + noise

# compute worst-case perturbation
# (see Sec. 19.4 of Adcock & Hansen)

lam = 2
x = torch.from_numpy(np.reshape(np.transpose(X),(N*N)))
z0 = torch.zeros(N*N,dtype=y.dtype)

def recon_fn(y):
    rec, _ = nn.restarted_nesta_wqcbp(
                y, z0,
                subsampled_ft, discrete_grad, L_grad,
                inner_iters, outer_iters,
                eta, mu0, delta, r, False)
    return rec

adv_noise, adv_rec = stability.adversarial_perturbation(
    x, x, subsampled_ft, recon_fn, lam, 25, 1e-1, 0.9)

assert len(adv_noise.shape) == 1
print('l2 norm of perturbation:', torch.linalg.norm(adv_noise,2))

print(adv_noise.dtype)
print(adv_noise)

X_pert = np.transpose(np.reshape(np.real(adv_rec.numpy()),(N,N)))
X_pert = np.clip((X_pert*255),0,255).astype('uint8')

pert_im = np.transpose(np.reshape(np.abs(adv_noise.numpy()),(N,N)))
pert_im = np.clip(255*(pert_im-np.min(pert_im))/(np.max(pert_im)-np.min(pert_im)),0,255).astype('uint8')

Image.fromarray(X_pert).show()
Image.fromarray(pert_im).show()

max_stars_repo_name: mneyrane/AS-NESTA-net

max_issues_repo_name: mneyrane/AS-NESTA-net

max_forks_repo_name: mneyrane/AS-NESTA-net