```python
import numpy as np
from scipy.integrate import odeint

import numpy as np
from sympy import symbols,sqrt,sech,Rational,lambdify,Matrix,exp,cosh,cse,simplify,cos,sin
from sympy.vector import CoordSysCartesian

#from theano.scalar.basic_sympy import SymPyCCode
#from theano import function
#from theano.scalar import floats

from IRI import *
#from Symbolic import *
from scipy.integrate import simps
from ENUFrame import ENU

import astropy.coordinates as ac
import astropy.units as au
import astropy.time as at

from time import time as tictoc
from TricubicInterpolation import TriCubic

import math
import pp

from RadioArray import RadioArray

class Fermat(object):
    def __init__(self,neTCI=None,frequency = 120e6,type='s'):
        self.type = type
        self.frequency = frequency#Hz
        if neTCI is not None:
            self.ne2n(neTCI)
            return

    def loadFunc(self,file):
        '''Load the model given in `file`'''
        data = np.load(file)
        if 'ne' in data.keys():
            ne = data['ne']
            xvec = data['xvec']
            yvec = data['yvec']
            zvec = data['zvec']
            self.ne2n(TriCubic(xvec,yvec,zvec,ne))
            return
        if 'n' in data.keys():
            ne = data['n']
            xvec = data['xvec']
            yvec = data['yvec']
            zvec = data['zvec']
            self.n2ne(TriCubic(xvec,yvec,zvec,n))
            return

    def saveFunc(self,file):
        np.savez(file,xvec=self.nTCI.xvec,yvec=self.nTCI.yvec,zvec=self.nTCI.zvec,n=self.nTCI.m,ne=self.neTCI.m)

    def ne2n(self,neTCI):
        '''Analytically turn electron density to refractive index. Assume ne in m^-3'''
        self.neTCI = neTCI
        #copy object
        self.nTCI = neTCI.copy(default=1.)
        #inplace change to refractive index
        self.nTCI.m *= -8.980**2/self.frequency**2
        self.nTCI.m += 1.
        self.nTCI.m = np.sqrt(self.nTCI.m)
        #wp = 5.63e4*np.sqrt(ne/1e6)/2pi#Hz^2 m^3 lightman p 226
        return self.nTCI

    def n2ne(self,nTCI):
        """Get electron density in m^-3 from refractive index"""
        self.nTCI = nTCI
        #convert to
        self.neTCI = nTCI.copy()
        self.neTCI.m *= -self.neTCI.m
        self.neTCI.m += 1.
        self.neTCI.m *= self.frequency**2/8.980**2
        #wp = 5.63e4*np.sqrt(ne/1e6)/2pi#Hz^2 m^3 lightman p 226
        return self.neTCI

    def eulerODE(self,y,t,*args):
        '''return pxdot,pydot,pzdot,xdot,ydot,zdot,sdot'''
        #print(y)
        px,py,pz,x,y,z,s = y
        #n,nx,ny,nz,nxy,nxz,nyz,nxyz = self.nTCI.interp(x,y,z,doDiff=True)
        #ne,nex,ney,nez,nexy,nexz,neyz,nexyz = self.neTCI.interp(x,y,z,doDiff=True)
        #A = - 8.98**2/self.frequency**2
        #n = math.sqrt(1. + A*ne)
        #ndot = A/(2.*n)
        #nx = ndot * nex
        #ny = ndot * ney
        #nz = ndot * nez

        #print(n)
        n,nx,ny,nz = 1.,0,0,0
        #if (n>1):
        #    print(x,y,z,n)
        if self.type == 'z':
            sdot = n / pz
            pxdot = nx*n/pz
            pydot = ny*n/pz
            pzdot = nz*n/pz

            xdot = px / pz
            ydot = py / pz
            zdot = 1.

        if self.type == 's':
            sdot = 1.
            pxdot = nx
            pydot = ny
            pzdot = nz

            xdot = px / n
            ydot = py / n
            zdot = pz / n

        return [pxdot,pydot,pzdot,xdot,ydot,zdot,sdot]

    def jacODE(self,y,t,*args):
        '''return d ydot / d y, with derivatives down column for speed'''
        px,py,pz,x,y,z,s = y
        #n,nx,ny,nz,nxy,nxz,nyz,nxyz = self.nTCI.interp(x,y,z,doDiff=True)
        nxx,nyy,nzz = 0.,0.,0.
        n,nx,ny,nz,nxy,nxz,nyz = 1.,0,0,0,0,0,0

        #ne,nex,ney,nez,nexy,nexz,neyz,nexyz = self.neTCI.interp(x,y,z,doDiff=True)
        #A = - 8.98**2/self.frequency**2
        #n = math.sqrt(1. + A*ne)
        #ndot = A/(2.*n)
        #nx = ndot * nex
        #ny = ndot * ney
        #nz = ndot * nez

        #ndotdot = -(A * ndot)/(2. * n**2)

        #nxy = ndotdot * nex*ney + ndot * nexy
        #nxz = ndotdot * nex * nez + ndot * nexz
        #nyz = ndotdot * ney * nez + ndot * neyz

        #if (n>1):
        #    print(x,y,z,n)

        if self.type == 'z':
            x0 = n
            x1 = nx
            x2 = pz**(-2)
            x3 = x0*x2
            x4 = 1./pz
            x5 = ny
            x6 = x4*(x0*nxy + x1*x5)
            x7 = nz
            x8 = x4*(x0*nxz + x1*x7)
            x9 = x4*(x0*nyz + x5*x7)
            jac = np.array([[ 0,  0, -x1*x3, x4*(x0*nxx + x1**2),x6, x8, 0.],
                            [ 0,  0, -x3*x5,x6, x4*(x0*nyy + x5**2), x9, 0.],
                            [ 0,  0, -x3*x7,x8, x9, x4*(x0*nzz + x7**2), 0.],
                            [x4,  0, -px*x2, 0, 0,  0, 0.],
                            [ 0, x4, -py*x2, 0, 0, 0, 0.],
                            [ 0,  0, 0, 0, 0, 0, 0.],
                            [ 0,  0,-x3,x1*x4, x4*x5, x4*x7, 0.]])

        if self.type == 's':
            x0 = n
            x1 = nxy
            x2 = nxz
            x3 = nyz
            x4 = 1./x0
            x5 = nx
            x6 = x0**(-2)
            x7 = px*x6
            x8 = ny
            x9 = nz
            x10 = py*x6
            x11 = pz*x6
            jac = np.array([[ 0,  0,  0, nxx, x1, x2, 0.],
                            [ 0,  0,  0, x1, nyy, x3, 0.],
                            [ 0,  0,  0, x2, x3, nzz, 0.],
                            [x4,  0,  0, -x5*x7, -x7*x8, -x7*x9, 0.],
                            [ 0, x4,  0, -x10*x5, -x10*x8, -x10*x9, 0.],
                            [ 0,  0, x4, -x11*x5, -x11*x8, -x11*x9, 0.],
                            [ 0,  0,  0, 0, 0, 0, 0.]])
        return jac

    def integrateRay(self,X0,direction,tmax,time = 0,N=100):
        '''Integrate rays from x0 in initial direction where coordinates are (r,theta,phi)'''
        direction /= np.linalg.norm(direction)
        x0,y0,z0 = X0
        xdot0,ydot0,zdot0 = direction
        sdot = np.sqrt(xdot0**2 + ydot0**2 + zdot0**2)
        px0 = xdot0/sdot
        py0 = ydot0/sdot
        pz0 = zdot0/sdot
        init = [px0,py0,pz0,x0,y0,z0,0]
        if self.type == 'z':
            tarray = np.linspace(z0,tmax,N)
        if self.type == 's':
            tarray = np.linspace(0,tmax,N)
        #print("Integrating at {0} from {1} in direction {2} until {3}".format(time,X0,direction,tmax))
        #print(init)
        #print("Integrating from {0} in direction {1} until {2}".format(x0,directions,tmax))
        Y,info =  odeint(self.eulerODE, init, tarray, args=(time,),Dfun = self.jacODE, col_deriv = True, full_output=1)
        #print(info['hu'].shape,np.sum(info['hu']),info['hu'])
        #print(Y)
        x = Y[:,3]
        y = Y[:,4]
        z = Y[:,5]
        s = Y[:,6]
        return x,y,z,s

def plotWavefront(neTCI,rays,save=False,animate=False):
    xmin = neTCI.xvec[0]
    xmax = neTCI.xvec[-1]
    ymin = neTCI.yvec[0]
    ymax = neTCI.yvec[-1]
    zmin = neTCI.zvec[0]
    zmax = neTCI.zvec[-1]

    X,Y,Z = np.mgrid[xmin:xmax:len(neTCI.xvec)*1j,
                     ymin:ymax:len(neTCI.yvec)*1j,
                     zmin:zmax:len(neTCI.zvec)*1j]

    #reshape array
    data = neTCI.getShapedArray()
    print(np.mean(data),np.max(data),np.min(data))
    l = mlab.pipeline.volume(mlab.pipeline.scalar_field(X,Y,Z,data))#,vmin=min, vmax=min + .5*(max-min))
    l._volume_property.scalar_opacity_unit_distance = min((xmax-xmin)/4.,(ymax-ymin)/4.,(zmax-zmin)/4.)
    l._volume_property.shade = False
    mlab.contour3d(X,Y,Z,data,contours=5,opacity=0.2)
    mlab.colorbar()

    def getWave(rays,idx):
        xs = np.zeros(len(rays))
        ys = np.zeros(len(rays))
        zs = np.zeros(len(rays))
        ridx = 0
        while ridx < len(rays):
            xs[ridx] = rays[ridx]['x'][idx]
            ys[ridx] = rays[ridx]['y'][idx]
            zs[ridx] = rays[ridx]['z'][idx]
            ridx += 1
        return xs,ys,zs

    if rays is not None:
        for datumIdx in rays.keys():
            ray = rays[datumIdx]
            mlab.plot3d(ray["x"],ray["y"],ray["z"],tube_radius=0.25)
        if animate:
            plt = mlab.points3d(*getWave(rays,0),color=(1,0,0),scale_mode='vector', scale_factor=10.)
            #mlab.move(-200,0,0)
            view = mlab.view()
            @mlab.animate(delay=100)
            def anim():
                nt = len(rays[0]["s"])
                f = mlab.gcf()
                save = False
                while True:
                    i = 0
                    while i < nt:
                        #print("updating scene")
                        xs,ys,zs = getWave(rays,i)
                        plt.mlab_source.set(x=xs,y=ys,z=zs)
                        #mlab.view(*view)
                        if save:
                            #mlab.view(*view)
                            mlab.savefig('figs/wavefronts/wavefront_{0:04d}.png'.format(i))#,magnification = 2)#size=(1920,1080))
                        #f.scene.render()
                        i += 1
                        yield
                    save = False
            anim()
    mlab.show()
    if save and rays is not None:
        return
        import os
        os.system('ffmpeg -r 10 -f image2 -s 1900x1080 -i figs/wavefronts/wavefront_%04d.png -vcodec libx264 -crf 25  -pix_fmt yuv420p figs/wavefronts/wavefront.mp4')

def plotModel(neTCI,save=False):
    '''Plot the model contained in a tricubic interpolator (a convienient container for one)'''
    plotWavefront(neTCI,None,save=save)

def createPrioriModel(iri = None,L_ne=15.):
    if iri is None:
        iri = IriModel()
    xmin = -200.
    xmax = 200.
    ymin = -200.
    ymax = 200.
    zmin = -10.
    zmax = 3000.

    eastVec = np.linspace(xmin,xmax,int(np.ceil((xmax-xmin)/L_ne)))
    northVec = np.linspace(ymin,ymax,int(np.ceil((ymax-ymin)/L_ne)))
    upVec = np.linspace(zmin,zmax,int(np.ceil((zmax-zmin)/L_ne)))

    E,N,U = np.meshgrid(eastVec,northVec,upVec,indexing='ij')
    #get the points in ITRS frame
    points = ac.SkyCoord(E.flatten()*au.km,N.flatten()*au.km,U.flatten()*au.km,frame=iri.enu).transform_to('itrs').cartesian.xyz.to(au.km).value
    X = points[0,:].reshape(E.shape)
    Y = points[1,:].reshape(N.shape)
    Z = points[2,:].reshape(U.shape)
    #X,Y,Z = np.meshgrid(points[0,:],points[1,:],points[2,:],indexing='ij')

    ##generate cartesian grid in ITRS frame
    #Nx = int(np.ceil((np.max(points[0,:]) - np.min(points[0,:]))/30.))
    #Ny = int(np.ceil((np.max(points[1,:]) - np.min(points[1,:]))/30.))
    #Nz = int(np.ceil((np.max(points[2,:]) - np.min(points[2,:]))/30.))

    #xvec = np.linspace(np.min(points[0,:]),np.max(points[0,:]),Nx)
    #yvec = np.linspace(np.min(points[1,:]),np.max(points[1,:]),Ny)
    #zvec = np.linspace(np.min(points[2,:]),np.max(points[2,:]),Nz)
    #ij indexing - mgrid is that way by default
    #X,Y,Z = np.mgrid[np.min(points[0,:]):np.max(points[0,:]):1j*Nx,
    #                 np.min(points[1,:]):np.max(points[1,:]):1j*Ny,
    #                 np.min(points[2,:]):np.max(points[2,:]):1j*Nz]
    #X,Y,Z = np.meshgrid(xvec,yvec,zvec,indexing='ij')
    #Get values at points
    ne = iri.evaluate(X,Y,Z)
    print("created an a priori cube of shape: {0}".format(ne.shape))
    return eastVec,northVec,upVec,ne

def perturbModel(eastVec,northVec,upVec,ne,loc,width,amp):
    nePert = ne.copy()
    E,N,U = np.meshgrid(eastVec,northVec,upVec,indexing='ij')
    for l,w,a in zip(loc,width,amp):
        print("Adding amp:{0:1.2e} at: {1} scale:{2:0.2f}".format(a,l,w))
        nePert += a*np.exp(-((E-l[0])**2 + (N-l[1])**2 + (U-l[2])**2)/w**2)
    return nePert

def testSweep():
    '''Test the full system.'''
    # The priori ionosphere
    iri = IriModel()
    print("Creating priori model")
    eastVec,northVec,upVec,nePriori = createPrioriModel(iri)
    print("Creating perturbed model")
    nePert = perturbModel(eastVec,northVec,upVec,nePriori,([0,0,200.],),(40.,),(1e12,))
    print("creating TCI object")
    neTCI = TriCubic(eastVec,northVec,upVec,nePert)
    print("creating fermat object")
    f =  Fermat(neTCI = neTCI,type = 's')

    ### test interpolation in both
    for i in range(0):
        x = np.random.uniform(low=eastVec[0],high=eastVec[-1])
        y = np.random.uniform(low=northVec[0],high=northVec[-1])
        z = np.random.uniform(low=upVec[0],high=upVec[-1])
        n_,nx_,ny_,nz_,nxy_,nxz_,nyz_,nxyz_ = f.nTCI.interp(x,y,z,doDiff=True)

        ne,nex,ney,nez,nexy,nexz,neyz,nexyz = f.neTCI.interp(x,y,z,doDiff=True)
        A = - 8.98**2/f.frequency**2
        n = math.sqrt(1. + A*ne)
        ndot = A/(2.*n)
        nx = ndot * nex
        ny = ndot * ney
        nz = ndot * nez
        ndotdot = -(A * ndot)/(2. * n**2)
        nxy = ndotdot * nex*ney + ndot * nexy
        nxz = ndotdot * nex * nez + ndot * nexz
        nyz = ndotdot * ney * nez + ndot * neyz
        print(x,y,z)
        print(n,n_)
        print(nx,nx_)
        print(nxy,nxy)

    print("min and max n:",np.min(f.nTCI.m),np.max(f.nTCI.m))
    theta = np.linspace(-np.pi/15.,np.pi/15.,25)
    #phi = np.linspace(0,2*np.pi,6)
    rays = []
    origin = ac.ITRS(iri.enu.location).transform_to(iri.enu).cartesian.xyz.to(au.km).value
    print(origin)
    rayIdx = 0
    t1 = tictoc()
    for t in theta:
        for p in theta:
            #print("integrating ray: {0}".format(rayIdx))
            direction = ac.SkyCoord(np.sin(t),
                                    np.sin(p),
                                    1.,frame=iri.enu).cartesian.xyz.value#.transform_to('itrs').cartesian.xyz.value
            x,y,z,s = f.integrateRay(origin,direction,1000,time=0.)
            rayIdx += 1
            rays.append({'x':x,'y':y,'z':z,'s':s})
    print("time per ray:",(tictoc()-t1)/len(rays))
    #print(rays)
    plotWavefront(neTCI,rays,save=False)
    #plotWavefront(f.nFunc.subs({'t':0}),rays,*getSolitonCube(sol),save = False)
    #plotFuncCube(f.nFunc.subs({'t':0}), *getSolitonCube(sol),rays=rays)

def plot_dtec(Nant,directions,dtec,title='',subAnt=None):
    def getDatumIdx(antIdx,dirIdx,timeIdx,numDirections,numTimes):
        '''standarizes indexing'''
        idx = antIdx*numDirections*numTimes + dirIdx*numTimes + timeIdx
        return idx
    vmin = np.min(dtec)
    vmax = np.max(dtec)
    #data -= np.min(dtec)
    #data /= np.max(dtec)
    Nperaxis = int(np.ceil(np.sqrt(Nant)))
    import pylab as plt
    cm = plt.cm.get_cmap('RdYlBu')
    f = plt.figure(figsize=(22,17))
    #f,ax = plt.subplots(int(np.ceil(np.sqrt(numAntennas))),int(np.ceil(np.sqrt(numAntennas))))
    for antIdx in range(Nant):
        ax = plt.subplot(Nperaxis,Nperaxis,antIdx+1)
        ax.set_title("Antenna {}".format(antIdx))
        for dirIdx in range(len(directions)):
            datumIdx = getDatumIdx(antIdx,dirIdx,0,len(directions),1)
            if subAnt is not None:
                datumIdx0 = getDatumIdx(subAnt,dirIdx,0,len(directions),1)
                sc=ax.scatter(directions[dirIdx,0],directions[dirIdx,1],c=dtec[datumIdx]-dtec[datumIdx0],s=20**2,vmin=vmin,vmax=vmax,cmap=cm)
            else:
                sc=ax.scatter(directions[dirIdx,0],directions[dirIdx,1],c=dtec[datumIdx],s=20**2,vmin=vmin,vmax=vmax,cmap=cm)
        plt.colorbar(sc)
    if title is not "":
        f.savefig("figs/dtec/{}.png".format(title),format='png')
    #plt.show()

def SimulatedDataInversion(numThreads = 1,noise=None):
    '''Test the full system.'''

    def getDatumIdx(antIdx,dirIdx,timeIdx,numDirections,numTimes):
        '''standarizes indexing'''
        idx = antIdx*numDirections*numTimes + dirIdx*numTimes + timeIdx
        return idx

    def reverseDatumIdx(datumIdx,numTimes,numDirections):
        '''Reverse standardized indexing'''
        timeIdx = datumIdx % numTimes
        dirIdx = (datumIdx - timeIdx)/numTimes % numDirections
        antIdx = (datumIdx - timeIdx - dirIdx*numTimes)/numTimes/numDirections
        return antIdx, dirIdx, timeIdx

    def datumDicts2array(datumDicts):
        '''Given a tupel of dicts where each dict is of datumIdx:value
        convert into single array with index giving order'''
        N = 0
        for datumDict in datumDicts:
            N += len(datumDict)
        array = np.zeros(N,dtype=np.double)
        for datumDict in datumDicts:
            for datumIdx in datumDict.keys():#ordering set by datumIdx function 1-to-1
                array[datumIdx] = datumDict[datumIdx]
        return array

    raylength = 2000.
    print("Using lofar array")
    radioArray = RadioArray(arrayFile='arrays/lofar.hba.antenna.cfg')
    timestamp = '2017-02-7T15:37:00.000'
    timeIdx = 0#one time stamp for now
    numTimes = 1
    time = at.Time(timestamp,format='isot',scale='tai')
    enu = ENU(obstime=time,location=radioArray.getCenter().earth_location)
    phase = ac.SkyCoord(east=0,north=0,up=1,frame=enu).transform_to(ac.ITRS(obstime=time)).transform_to('icrs')#straight up for now
    dec = phase.dec.rad
    ra = phase.ra.rad
    print("Simulating observation on {0}: {1}".format(time.isot,phase))

    stations = radioArray.locs.transform_to(enu).cartesian.xyz.to(au.km).value.transpose()
    stations = stations[50:60,:]
    Nant = stations.shape[0]
    print("Using {0} stations".format(Nant))
    print(stations)
    #stations = np.random.multivariate_normal([0,0,0],[[20**2,0,0],[0,20**2,0],[0,0,0.01**2]],Nant)
    #stations = np.array([[0,0,0],[20,0,0]])

    Ndir = 5
    fov = radioArray.getFov()#radians
    print("Creating {0} directions in FOV of {1}".format(Ndir,fov))
    directions = np.random.multivariate_normal([ra,dec],[[(fov/2.)**2,0],[0,(fov/2.)**2]],Ndir)
    print(directions)

    directions = ac.SkyCoord(directions[:,0]*au.radian,directions[:,1]*au.radian,frame='icrs').transform_to(enu).cartesian.xyz.value.transpose()

    #print(directions)

    print("Setting up tri cubic interpolator")
    L_ne = 15.

    # The priori ionosphere
    iri = IriModel()
    print("Creating priori model")
    eastVec,northVec,upVec,nePriori = createPrioriModel(iri,L_ne)
    print("Creating perturbed model")
    nePert = perturbModel(eastVec,northVec,upVec,nePriori,([0,0,200.],),(40.,),(1e9,))
    print("Creating TCI object")
    neTCI = TriCubic(eastVec,northVec,upVec,nePert)
    neTCIModel = TriCubic(eastVec,northVec,upVec,nePriori)
    TCI = TriCubic(eastVec,northVec,upVec,np.zeros_like(nePert))

    print("Creating fermat object - based on a priori (second order corrections require iterating this)")
    f =  Fermat(neTCI = neTCIModel,type = 's')

    print("Integrating rays with fermats principle")
    t1 = tictoc()
    rays = {}
    for antIdx in range(Nant):
        for dirIdx in range(Ndir):
            datumIdx = getDatumIdx(antIdx,dirIdx,timeIdx,Ndir,numTimes)
            #print(antIdx,dirIdx,timeIdx,datumIdx)
            origin = stations[antIdx,:]#ENU frame, later use UVW frame
            direction = directions[dirIdx,:]
            x,y,z,s = f.integrateRay(origin,direction,raylength,time=0.)
            rays[datumIdx] = {'x':x,'y':y,'z':z,'s':s}
    Nd = len(rays)
    print("Time (total/per ray): {0:0.2f} / {1:0.2e} s".format(tictoc()-t1,(tictoc()-t1)/Nd))
    TCI.m = neTCI.m - neTCIModel.m
    TCI.clearCache()
    #plotWavefront(TCI,rays,save=False,animate=False)
    print("Setting up ray chunks for {0} threads".format(numThreads))
    #split up rays
    raypack = {i:{} for i in range(numThreads)}
    c = 0
    for datumIdx in rays.keys():
        raypack[c%numThreads][datumIdx] = rays[datumIdx]
        c += 1

    def ppForwardEquation(rays,TCI,mu,Kmu,rho,Krho,numTimes,numDirections):
        dtec, rho, Krho = ParallelInversionProducts.forwardEquations(rays,TCI,mu,Kmu,rho,Krho,numTimes,numDirections)
        return dtec, rho, Krho
    def ppPrimaryInversionSteps(dtec,rays,TCI,mu,Kmu,rho,Krho,muprior,rhoprior,sigma_ne,L_ne,sigma_rho,numTimes,numDirections,priorFlag=True):
        G, CmGt, ddGdmpm = ParallelInversionProducts.primaryInversionSteps(dtec,rays,TCI,mu,Kmu,rho,Krho,muprior,rhoprior,sigma_ne,L_ne,sigma_rho,numTimes,numDirections,priorFlag=True)
        return G, CmGt, ddGdmpm
    def ppSecondaryInversionSteps(rays, G, CmGt, TCI, sigma_rho, Cd,numTimes,numDirections):
        S = ParallelInversionProducts.secondaryInversionSteps(rays, G, CmGt, TCI, sigma_rho, Cd,numTimes,numDirections)
        return S

    jobs = {}
    job_server = pp.Server(numThreads, ppservers=())
    print("Creating dTec simulated data")
    job = job_server.submit(ppForwardEquation,
                   args=(rays,TCI,np.log(neTCI.m/np.mean(neTCI.m)),np.mean(neTCI.m),None,None,numTimes,Ndir),
                   depfuncs=(),
                   modules=('ParallelInversionProducts',))
    jobs['dtecSim'] = job

    job = job_server.submit(ppForwardEquation,
                   args=(rays,TCI,np.log(neTCIModel.m/np.mean(neTCIModel.m)),np.mean(neTCIModel.m),None,None,numTimes,Ndir),
                   depfuncs=(),
                   modules=('ParallelInversionProducts',))
    jobs['dtecModel'] = job

    dtecSim,rhoSim0, KrhoSim0 = jobs['dtecSim']()
    dobs = datumDicts2array((dtecSim,))
    #print("dobs: {0}".format(dobs))
    if noise is not None:
        print("Adding {0:0.2f}-sigma noise to simulated dtec".format(noise))
        dtecStd = np.std(dobs)
        dobs += np.random.normal(loc=0,scale=dtecStd*noise,size=np.size(dobs))
    #print("dobs: {0}".format(dobs))
    dtecModel,rhoModel0,KrhoModel0 = jobs['dtecModel']()
    g = datumDicts2array((dtecModel,))
    #print("g: {0}".format(g))
    job_server.print_stats()
    job_server.destroy()

    subAnt = None
    plot_dtec(Nant,directions,dobs,title='sim_dtec',subAnt=subAnt)
    plot_dtec(Nant,directions,g,title='model_dtec',subAnt=subAnt)
    plot_dtec(Nant,directions,dobs-g,title='sim-mod_dtec',subAnt=subAnt)

    print("Setting up inversion with parameters:")
    print("Number of rays: {0}".format(Nd))
    print("Forward equation: g(m) = int_R^i (K_mu * EXP[mu(x)] - K_rho * EXP[rho])/TECU ds")

    #gaussian process assumption, d = g + G.dm -> Cd = Gt.Cm.G (not sure)
    Cd = np.eye(Nd)*np.std(dobs)
    print("<Diag(Cd)> = {0:0.2e}".format(np.mean(np.diag(Cd))))

    print("a priori model is IRI")
    print("Define: mu(x) = LOG[ne(x) / K_mu]")
    Kmu = np.mean(neTCIModel.m)
    mu = np.log(neTCIModel.m/Kmu)
    muPrior = mu.copy()
    print("K_mu = {0:0.2e}".format(Kmu))

    #spatial-ergodic assumption
    sigma_ne = np.std(neTCIModel.m)
    print("Coherence scale: L_ne = {0:0.2e}".format(L_ne))
    print("C_ne = ({0:0.2e})**2 EXP[-|x1 - x2| / {1:0.1f}]".format(sigma_ne,L_ne))
    print("Define: rho = LOG[TEC_0 / K_rho / S]")
    Krho = KrhoModel0
    rho = rhoModel0
    rhoPrior = rho.copy()
    sigma_TEC = np.std(g*1e13)
    sigma_rho = np.sqrt(np.log(1+(sigma_TEC/Krho/raylength)**2))
    print("K_rho = {0:0.2e}".format(Krho))
    print("a priori rho (reference TEC): {0}".format(rho))
    print("sigma_rho = {0:0.2e}".format(sigma_rho))

    #inversion steps
    iter = 0
    residuals = np.inf
    while residuals > 1e-10:
        print("Performing iteration: {0}".format(iter))
        print("Performing primary inversion steps on {0}".format(numThreads))
        job_server = pp.Server(numThreads, ppservers=())

        for i in range(numThreads):
            job = job_server.submit(ppPrimaryInversionSteps,
                       args=(dtecSim,raypack[i],TCI,mu,Kmu,rho,Krho,muPrior,rhoPrior,sigma_ne,L_ne,sigma_rho,numTimes,Ndir,True),
                       depfuncs=(),
                       modules=('ParallelInversionProducts',))
            jobs['ppPrimaryInversionSteps_{0}'.format(i)] = job
        G,CmGt,ddGdmpm = {},{},{}
        for i in range(numThreads):
            G_, CmGt_, ddGdmpm_ = jobs['ppPrimaryInversionSteps_{0}'.format(i)]()
            #print(G_, CmGt_, ddGdmpm_)
            G.update(G_)
            CmGt.update(CmGt_)
            ddGdmpm.update(ddGdmpm_)
        job_server.print_stats()
        job_server.destroy()
        print("Performing secondary inversion steps")
        job_server = pp.Server(numThreads, ppservers=())
        for i in range(numThreads):
            job = job_server.submit(ppSecondaryInversionSteps,
                       args=(raypack[i], G, CmGt, TCI, sigma_rho, Cd,numTimes,Ndir),
                       depfuncs=(),
                       modules=('ParallelInversionProducts',))
            jobs['ppSecondaryInversionSteps_{0}'.format(i)] = job
        S = np.zeros([Nd,Nd],dtype=np.double)
        for i in range(numThreads):
            S_ = jobs['ppSecondaryInversionSteps_{0}'.format(i)]()
            S += S_
        print("Inverting S")
        T = np.linalg.pinv(S)
        if False:
            import pylab as plt
            ax = plt.subplot(121)
            p1 = ax.imshow(S)
            plt.colorbar(p1)
            ax = plt.subplot(122)
            p2 = ax.imshow(T)
            plt.colorbar(p2)
            print("S:",S)
            print("T:",T)
        #plt.show()
        job_server.print_stats()
        job_server.destroy()
        # dm = (mp-m) + CmGt.T.ddGdmpm
        ddGdmpmArray = datumDicts2array([ddGdmpm])
        TddGdmpmArray = T.dot(ddGdmpmArray)
        CmGtArray = np.zeros([np.size(mu)+np.size(rho),Nd])
        for i in range(Nd):
            CmGtArray[:np.size(mu),i] = CmGt[i][0]
            CmGtArray[np.size(mu):,i] = CmGt[i][1]
        dm = CmGtArray.dot(TddGdmpmArray)
        dmu = (muPrior - mu) + dm[:np.size(mu)]
        drho = (rhoPrior - rho) + dm[np.size(mu):]
        residuals = np.sum(dmu**2) / np.sum(mu**2) + np.sum(drho**2) / np.sum(rho**2)
        print("Residual:",residuals)
        print("Incrementing mu and rho")
        print("dmlog:",dmu)
        print("drho:",drho)
        mu += dmu
        rho += drho
        iter += 1
    print('Finished inversion with {0} iterations'.format(iter))
    #print(rays)
    TCI.m = Kmu*np.exp(mu) - neTCIModel.m
    TCI.clearCache()
    plotWavefront(TCI,rays,save=False)
    #plotWavefront(f.nFunc.subs({'t':0}),rays,*getSolitonCube(sol),save = False)
    #plotFuncCube(f.nFunc.subs({'t':0}), *getSolitonCube(sol),rays=rays)
def LMSol(G,mprior,Cd,Cm,dobs,mu=1.,octTree=None):
    """Assume the frechet derivative is,
    G(x) = exp"""
    import pylab as plt

    K = np.mean(mprior)
    mlog = np.log(mprior/K)

    Cm_log = transformCov2Log(Cm,K)#np.log(1. + Cm/K**2)#transformCov2Log(Cm,mprior)
    #Cdinv = np.linalg.pinv(Cd)
    if octTree is not None:
        voxels = getAllDecendants(octTree)
        scale = np.zeros(np.size(mprior))
        i = 0
        while i < np.size(mprior):
            scale[i] = voxels[i].volume**(1./3.)
            i+= 1
        C = np.sum(G,axis=0)/scale
        C = C/float(np.max(C))
        C[C==0] = np.min(C[C>0])/2.
    else:
        C = np.sum(G>0,axis=0)
        plt.hist(C)
        plt.show()
        C = C/float(np.max(C))
        C[C==0] = np.min(C[C>0])/2.
        #C = np.sum(G,axis=0)
        #C = C/np.max(C)
    res = 1
    iter = 0
    while res > 1e-6 and iter < 10000:
        #forward transform
        #print(mlog)
        mForward = K*np.exp(mlog)

        g = G.dot(mForward)
        J = G*mForward
        #residuals g - dobs -> -dm
        res = g - dobs
        #A1 = J.transpose().dot(Cdinv)
        #Cmlog_inv = A1.dot(J) + mu*Cm_log
        #dm,resi,rank,s = np.linalg.lstsq(Cmlog_inv,A1.dot(res))
        #S = mu Cd + J.Cm.J^t
        #S = int Ri Rj k^2 exp(m(x) + m(x')) sigma^2 exp(-|x-x'|/L) + Cd

        #K int dV Cm(x,x') J(x') del(i)
        P1 = Cm_log.dot(J.transpose())
        smooth = np.linalg.pinv(mu*Cd + J.dot(P1))
        dm = P1.dot(smooth).dot(res)
        res = np.sum(dm**2)/np.sum(mlog**2)
        print("Iter-{0} res: {1}".format(iter,res))
        #converage learn propto length of rays in cells
        #print(dm)
        mlog -= dm*C
        iter += 1
    CmlogPost = Cm_log - P1.dot(smooth).dot(P1.transpose())
    cmlin = transformCov2Linear(CmlogPost,K)
    #print(CmlogPost)
    #mMl,cmlin = metropolisPosteriorCovariance(G,dobs,Cd,CmlogPost,mlog,K)
    #print(mMl - K*np.exp(mlog))
    #print(transformCov2Linear(CmlogPost,K) - cmlin)
    return K*np.exp(mlog), cmlin

if __name__=='__main__':
    np.random.seed(1234)
    #testSquare()
    #testSweep()
    SimulatedDataInversion(4,noise=None)
    #testThreadedFermat()
    #testSmoothify()
    #testcseLam()

```

    Using lofar array

    WARNING: Tried to get polar motions for times after IERS data is valid. Defaulting to polar motion from the 50-yr mean for those.
    If you need enough precision such that this matters (~<10 arcsec), you can
    download the latest IERS predictions by running:

        >>> from astropy.utils.data import download_file
        >>> from astropy.utils import iers
        >>> iers.IERS.iers_table = iers.IERS_A.open(download_file(iers.IERS_A_URL, cache=True))

     [astropy.coordinates.builtin_frames.utils]

    Simulating observation on 2017-02-07T15:37:00.000: <SkyCoord (ICRS): (ra, dec) in deg
        (18.78277889, 52.82321826)>
    Using 10 stations
    [[  5.19440041e+00  -2.72439778e+01  -3.70815110e-02]
     [  2.18090006e+00  -6.49589411e+01  -2.99629857e-01]
     [ -4.72236748e+00  -1.69407482e+00   2.11281777e-02]
     [ -6.76306945e+00  -2.69925381e+00   1.98400188e-02]
     [ -1.08513739e+01  -1.24399002e+01   4.84456806e-03]
     [ -4.74886384e+01  -1.64321721e+01  -1.88068457e-01]
     [ -6.14736033e+00   1.15285624e+01   1.36414320e-02]
     [ -3.91072222e+00   1.97198299e+01  -2.15157932e-02]
     [ -3.26139435e+01   7.38267739e+00  -6.89811048e-02]
     [  5.75474767e-01   3.30397868e+00   1.41152971e-02]]
    Creating 5 directions in FOV of 0.0698131700798
    [[ 0.34427753  0.88036619]
     [ 0.37783224  0.91102547]
     [ 0.30266804  0.95290691]
     [ 0.35782663  0.89972022]
     [ 0.32836924  0.84365461]]
    Setting up tri cubic interpolator
    Generated IRI symbolic function with 9 params
    Creating priori model
    created an a priori cube of shape: (27L, 27L, 201L)
    Creating perturbed model
    Adding amp:1.00e+09 at: [0, 0, 200.0] scale:40.00
    Creating TCI object
    Creating fermat object - based on a priori (second order corrections require iterating this)
    Integrating rays with fermats principle
    Time (total/per ray): 0.01 / 1.20e-04 s
    Setting up ray chunks for 4 threads
    Creating dTec simulated data
    Job execution statistics:
     job count | % of all jobs | job time sum | time per job | job server

```python
import numpy as np
1./np.tan(0.5*np.pi/180.) * 1
```

    114.58865012930961

```python

```


```python
import sys
sys.path.append('..')
import torch
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.preprocessing import LabelEncoder, MinMaxScaler
from sklearn.model_selection import StratifiedKFold
from sklearn.tree import DecisionTreeClassifier, plot_tree
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score
from sympy import simplify_logic

from lens.utils.relu_nn import get_reduced_model, prune_features
from lens import logic
from lens.utils.base import collect_parameters, set_seed

set_seed(0)
```

```python
gene_expression_matrix = pd.read_csv('w_1/data_0.csv', index_col=None, header=None)
labels = pd.read_csv('w_1/tempLabels_W-1.csv', index_col=None, header=None)
genes = pd.read_csv('w_1/features_0.csv', index_col=None, header=None)
```

```python
gene_expression_matrix
```

      0
      1
      2
      3
      4
      5
      6
      7
      8
      9
      ...
      28392
      28393
      28394
      28395
      28396
      28397
      28398
      28399
      28400
      28401

      0
      14.622486
      11.162004
      3.320000
      3.320000
      3.320000
      12.788433
      6.143456
      3.320000
      4.876620
      3.320000
      ...
      3.32
      3.320000
      3.885589
      3.914260
      3.320000
      3.32
      3.320000
      4.465420
      3.320000
      4.973620

      1
      14.398743
      11.000080
      3.320000
      3.320000
      3.320000
      12.845914
      6.147482
      3.320000
      4.484223
      3.320000
      ...
      3.32
      3.575025
      4.236519
      4.047825
      3.320000
      3.32
      4.176269
      3.320000
      4.553796
      4.967418

      2
      14.692079
      11.100175
      3.320000
      4.171535
      3.320000
      12.712544
      5.583210
      3.320000
      3.478171
      3.320000
      ...
      3.32
      3.320000
      3.992331
      4.865538
      3.320000
      3.32
      3.488281
      3.406285
      3.320000
      6.676063

      3
      14.613382
      11.023209
      3.320000
      3.320000
      3.320000
      12.750496
      5.688023
      3.320000
      4.464426
      3.320000
      ...
      3.32
      3.855643
      3.320000
      4.905350
      3.320000
      3.32
      4.158393
      4.433457
      3.874214
      5.981160

      4
      14.482065
      10.989851
      3.320000
      3.992726
      4.574745
      12.878702
      6.195418
      4.177962
      3.872567
      3.320000
      ...
      3.32
      3.320000
      3.320000
      4.879493
      3.320000
      3.32
      4.571869
      3.320000
      4.982136
      6.145585

      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...

      56
      14.565031
      11.699843
      3.320000
      3.320000
      3.320000
      12.789212
      6.504027
      3.320000
      3.320000
      6.182912
      ...
      3.32
      3.320000
      3.320000
      4.338456
      3.771718
      3.32
      3.320000
      3.320000
      3.320000
      6.100644

      57
      14.624502
      11.918757
      3.320000
      4.292406
      3.430485
      10.728709
      6.197159
      3.320000
      4.089918
      5.201608
      ...
      3.32
      3.320000
      3.320000
      4.790134
      3.320000
      3.32
      4.985474
      4.444057
      3.580523
      6.301926

      58
      14.585190
      11.090112
      3.320000
      3.674768
      3.320000
      12.877485
      6.326960
      3.320000
      3.320000
      3.320000
      ...
      3.32
      3.320000
      3.320000
      4.547342
      3.320000
      3.32
      4.064473
      3.320000
      4.254152
      5.964505

      59
      14.449554
      10.805855
      3.320000
      3.320000
      3.320000
      12.660038
      6.261395
      3.320000
      4.125096
      3.320000
      ...
      3.32
      3.784260
      3.644823
      4.546974
      3.427441
      3.32
      4.666265
      3.888525
      3.765754
      5.452018

      60
      14.439020
      11.080826
      4.152651
      3.320000
      3.320000
      12.187504
      6.154358
      3.320000
      4.247009
      3.320000
      ...
      3.32
      3.671102
      3.519157
      4.771458
      3.320000
      3.32
      5.115644
      3.320000
      3.410361
      6.449961

61 rows × 28402 columns

```python
labels
```

      0

      0
      diagnosis: healthy control

      1
      diagnosis: healthy control

      2
      diagnosis: healthy control

      3
      diagnosis: healthy control

      4
      diagnosis: healthy control

      ...
      ...

      56
      omalizumab responder status: Responder

      57
      omalizumab responder status: Responder

      58
      omalizumab responder status: Responder

      59
      omalizumab responder status: Responder

      60
      omalizumab responder status: Responder

61 rows × 1 columns

```python
encoder = LabelEncoder()
labels_encoded = encoder.fit_transform(labels.values)
labels_encoded_noncontrols = labels_encoded[labels_encoded!=0] - 1

data_controls = gene_expression_matrix[labels_encoded==0]
data = gene_expression_matrix[labels_encoded!=0]

gene_signature = data_controls.mean(axis=0)
data_scaled = data - gene_signature

scaler = MinMaxScaler((0, 1))
scaler.fit(data_scaled)
data_normalized = scaler.transform(data_scaled)

x = torch.FloatTensor(data_normalized)
y = torch.FloatTensor(labels_encoded_noncontrols).unsqueeze(1)
print(x.shape)
print(y.shape)
```

    c:\users\pietr\anaconda3\envs\deep-logic\lib\site-packages\sklearn\utils\validation.py:63: DataConversionWarning: A column-vector y was passed when a 1d array was expected. Please change the shape of y to (n_samples, ), for example using ravel().
      return f(*args, **kwargs)

    torch.Size([40, 28402])
    torch.Size([40, 1])

```python
set_seed(0)
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
x = x.to(device)
y = y.to(device)

layers = [
    torch.nn.Linear(x.size(1), 10),
    torch.nn.ReLU(),
    torch.nn.Linear(10, 5),
    torch.nn.ReLU(),
    torch.nn.Linear(5, 1),
    torch.nn.Sigmoid(),
]
model = torch.nn.Sequential(*layers).to(device)

optimizer = torch.optim.Adam(model.parameters(), lr=0.0001)
model.train()
need_pruning = True
for epoch in range(15000):
    # forward pass
    optimizer.zero_grad()
    y_pred = model(x)
    # Compute Loss
    loss = torch.nn.functional.mse_loss(y_pred, y)

    for module in model.children():
        if isinstance(module, torch.nn.Linear):
            loss += 0.01 * torch.norm(module.weight, 1)
            loss += 0.01 * torch.norm(module.bias, 1)

    # backward pass
    loss.backward()
    optimizer.step()

    # compute accuracy
    if epoch % 1000 == 0:
        y_pred_d = (y_pred > 0.5)
        accuracy = (y_pred_d.eq(y).sum(dim=1) == y.size(1)).sum().item() / y.size(0)
        print(f'Epoch {epoch}: train accuracy: {accuracy:.4f}')

    if epoch > 8000 and need_pruning and epoch % 1000 == 0:
        prune_features(model, 1, device)
        need_pruning = True
```

    Epoch 0: train accuracy: 0.2250
    Epoch 1000: train accuracy: 0.7500
    Epoch 2000: train accuracy: 0.7500
    Epoch 3000: train accuracy: 0.7500
    Epoch 4000: train accuracy: 0.9500
    Epoch 5000: train accuracy: 0.9750
    Epoch 6000: train accuracy: 1.0000
    Epoch 7000: train accuracy: 1.0000
    Epoch 8000: train accuracy: 1.0000
    Epoch 9000: train accuracy: 1.0000
    Epoch 10000: train accuracy: 1.0000
    Epoch 11000: train accuracy: 1.0000
    Epoch 12000: train accuracy: 1.0000
    Epoch 13000: train accuracy: 1.0000
    Epoch 14000: train accuracy: 1.0000

## Local explanations

```python
np.set_printoptions(precision=2, suppress=True)
outputs = []
for i, (xin, yin) in enumerate(zip(x, y)):
    model_reduced = get_reduced_model(model, xin.to(device)).to(device)
    for module in model_reduced.children():
        if isinstance(module, torch.nn.Linear):
            wa = module.weight.cpu().detach().numpy()
            break
    output = model_reduced(xin)

    pred_class = torch.argmax(output)
    true_class = torch.argmax(y[i])

    # generate local explanation only if the prediction is correct
    if pred_class.eq(true_class):
        local_explanation = logic.relu_nn.explain_local(model.to(device), x, y, xin, yin, device=device)
        print(f'Input {(i+1)}')
        print(f'\tx={xin.cpu().detach().numpy()}')
        print(f'\ty={y[i].cpu().detach().numpy()}')
        print(f'\ty={output.cpu().detach().numpy()}')
        #print(f'\tw={wa}')
        print(f'\tExplanation: {local_explanation}')
        print()
    outputs.append(output)
```

    Input 1
    	x=[0.94 0.31 0.   ... 0.26 0.8  0.68]
    	y=[0.]
    	y=[0.99]
    	Explanation: ~feature0000006749 & ~feature0000010382 & ~feature0000013464 & feature0000013710 & feature0000022010 & ~feature0000025379 & ~feature0000026329

    Input 2
    	x=[0.84 0.71 0.   ... 0.26 0.   0.62]
    	y=[0.]
    	y=[0.99]
    	Explanation: ~feature0000026329

    Input 3
    	x=[0.62 0.36 0.   ... 0.   0.   0.71]
    	y=[0.]
    	y=[0.99]
    	Explanation: ~feature0000006749 & ~feature0000010382 & ~feature0000013464 & ~feature0000013710 & ~feature0000022010 & ~feature0000025379 & feature0000026329

    Input 4
    	x=[0.81 0.95 0.   ... 0.83 0.   0.95]
    	y=[0.]
    	y=[0.99]
    	Explanation: ~feature0000006749 & feature0000013710 & ~feature0000026329

    Input 5
    	x=[0.   0.68 0.   ... 0.43 0.   0.89]
    	y=[0.]
    	y=[0.99]
    	Explanation: ~feature0000006749 & ~feature0000010382 & ~feature0000013464 & ~feature0000013710 & ~feature0000022010 & ~feature0000025379 & ~feature0000026329

    Input 6
    	x=[0.79 0.65 0.   ... 0.3  0.   0.69]
    	y=[0.]
    	y=[0.99]
    	Explanation:

    Input 7
    	x=[0.7  0.75 0.   ... 0.   0.52 0.42]
    	y=[0.]
    	y=[0.99]
    	Explanation: ~feature0000006749 & ~feature0000010382 & feature0000013464 & ~feature0000013710 & ~feature0000022010 & ~feature0000025379 & ~feature0000026329

    Input 8
    	x=[0.88 0.52 0.   ... 0.   0.33 0.74]
    	y=[0.]
    	y=[0.99]
    	Explanation: ~feature0000013710

    Input 9
    	x=[0.61 0.6  0.   ... 1.   0.   0.69]
    	y=[0.]
    	y=[0.99]
    	Explanation: ~feature0000006749 & ~feature0000010382 & ~feature0000013464 & ~feature0000013710 & feature0000022010 & ~feature0000025379 & feature0000026329

    Input 10
    	x=[0.76 0.07 0.   ... 0.41 0.88 0.1 ]
    	y=[0.]
    	y=[0.99]
    	Explanation: ~feature0000006749 & feature0000013710

    Input 11
    	x=[0.3  0.71 0.   ... 0.51 0.   0.97]
    	y=[1.]
    	y=[1.]
    	Explanation: ~feature0000006749 & ~feature0000010382 & ~feature0000013464 & feature0000013710 & ~feature0000022010 & ~feature0000025379 & feature0000026329

    Input 12
    	x=[0.6  0.51 0.   ... 0.23 0.   0.15]
    	y=[1.]
    	y=[1.]
    	Explanation: ~feature0000006749 & feature0000010382 & feature0000013710 & ~feature0000022010 & feature0000025379

    Input 13
    	x=[0.72 0.56 0.   ... 0.22 0.   0.82]
    	y=[1.]
    	y=[1.]
    	Explanation: feature0000006749 & feature0000026329

    Input 14
    	x=[0.31 0.53 0.   ... 0.12 0.   0.  ]
    	y=[1.]
    	y=[1.]
    	Explanation: ~feature0000006749 & feature0000010382 & ~feature0000013464 & feature0000013710 & ~feature0000022010 & ~feature0000025379 & ~feature0000026329

    Input 15
    	x=[0.96 0.57 0.   ... 0.44 0.54 0.59]
    	y=[1.]
    	y=[1.]
    	Explanation: feature0000006749 & feature0000010382 & ~feature0000013464

    Input 16
    	x=[0.86 0.53 0.   ... 0.69 0.08 0.73]
    	y=[1.]
    	y=[1.]
    	Explanation: feature0000006749 & ~feature0000010382 & feature0000013464 & feature0000013710 & ~feature0000022010 & ~feature0000025379 & ~feature0000026329

    Input 17
    	x=[0.83 0.67 0.   ... 0.   1.   0.5 ]
    	y=[1.]
    	y=[1.]
    	Explanation: ~feature0000006749 & feature0000010382 & ~feature0000013464 & ~feature0000022010 & feature0000025379

    Input 18
    	x=[0.95 0.67 0.   ... 0.68 0.   0.94]
    	y=[1.]
    	y=[1.]
    	Explanation: feature0000006749 & ~feature0000010382 & ~feature0000013464 & ~feature0000013710 & ~feature0000022010 & feature0000025379 & feature0000026329

    Input 19
    	x=[0.94 0.63 0.   ... 0.   0.   0.77]
    	y=[1.]
    	y=[1.]
    	Explanation: feature0000006749 & feature0000010382 & ~feature0000013464 & feature0000013710 & ~feature0000022010 & ~feature0000025379 & ~feature0000026329

    Input 20
    	x=[0.7  0.14 0.   ... 0.38 0.75 0.77]
    	y=[1.]
    	y=[1.]
    	Explanation: feature0000006749 & ~feature0000010382 & ~feature0000013464 & feature0000013710 & feature0000022010 & ~feature0000025379 & ~feature0000026329

    Input 21
    	x=[0.58 1.   0.   ... 0.   0.   0.56]
    	y=[1.]
    	y=[1.]
    	Explanation: ~feature0000006749 & feature0000010382 & ~feature0000013464 & ~feature0000013710 & ~feature0000022010 & ~feature0000025379 & feature0000026329

    Input 22
    	x=[0.7  0.63 0.   ... 0.86 0.   0.57]
    	y=[1.]
    	y=[1.]
    	Explanation: ~feature0000006749 & feature0000010382 & ~feature0000013464 & feature0000013710 & ~feature0000022010 & ~feature0000025379 & ~feature0000026329

    Input 23
    	x=[0.68 0.3  0.   ... 0.26 0.   0.87]
    	y=[1.]
    	y=[1.]
    	Explanation: feature0000006749 & feature0000010382 & ~feature0000013464 & ~feature0000013710 & ~feature0000022010 & ~feature0000025379 & ~feature0000026329

    Input 24
    	x=[0.8  0.41 0.   ... 0.54 0.   0.84]
    	y=[1.]
    	y=[1.]
    	Explanation: feature0000013710 & ~feature0000022010 & feature0000025379

    Input 25
    	x=[0.69 0.12 0.   ... 0.72 0.   0.34]
    	y=[1.]
    	y=[1.]
    	Explanation: ~feature0000006749 & feature0000010382 & ~feature0000013464 & ~feature0000013710 & ~feature0000022010 & feature0000025379 & feature0000026329

    Input 26
    	x=[0.61 0.82 0.   ... 0.85 0.   1.  ]
    	y=[1.]
    	y=[1.]
    	Explanation: ~feature0000006749 & ~feature0000010382 & ~feature0000013464 & feature0000013710 & ~feature0000022010 & feature0000025379 & feature0000026329

    Input 27
    	x=[1.   0.6  0.   ... 0.14 0.   0.64]
    	y=[1.]
    	y=[1.]
    	Explanation: ~feature0000006749 & feature0000010382 & ~feature0000013464 & feature0000013710 & ~feature0000022010 & ~feature0000025379 & feature0000026329

    Input 28
    	x=[0.37 0.75 0.   ... 0.07 0.   0.5 ]
    	y=[1.]
    	y=[1.]
    	Explanation: feature0000013710 & ~feature0000022010 & feature0000025379

    Input 29
    	x=[0.72 0.4  0.   ... 0.21 0.56 0.46]
    	y=[1.]
    	y=[1.]
    	Explanation: feature0000006749 & ~feature0000010382 & ~feature0000013464 & ~feature0000013710 & ~feature0000022010 & ~feature0000025379 & feature0000026329

    Input 30
    	x=[0.49 0.49 0.   ... 0.   0.   0.13]
    	y=[1.]
    	y=[1.]
    	Explanation: ~feature0000006749 & ~feature0000010382 & ~feature0000013464 & feature0000013710 & ~feature0000022010 & feature0000025379 & feature0000026329

    Input 31
    	x=[0.97 0.35 0.   ... 0.31 0.   0.41]
    	y=[1.]
    	y=[1.]
    	Explanation: feature0000013710 & ~feature0000022010 & feature0000025379

    Input 32
    	x=[0.86 0.57 0.   ... 0.38 0.01 0.92]
    	y=[1.]
    	y=[1.]
    	Explanation: feature0000006749 & feature0000010382 & ~feature0000013464

    Input 33
    	x=[0.56 0.72 0.   ... 0.29 1.   0.35]
    	y=[1.]
    	y=[1.]
    	Explanation: ~feature0000006749 & feature0000010382 & ~feature0000013464 & feature0000013710 & ~feature0000022010 & ~feature0000025379 & feature0000026329

    Input 34
    	x=[0.51 0.45 0.   ... 0.3  0.   0.3 ]
    	y=[1.]
    	y=[1.]
    	Explanation: ~feature0000006749 & ~feature0000010382 & ~feature0000013464 & feature0000013710 & ~feature0000022010 & ~feature0000025379 & feature0000026329

    Input 35
    	x=[0.3  0.28 0.   ... 0.   0.   0.42]
    	y=[1.]
    	y=[1.]
    	Explanation: feature0000013710 & ~feature0000022010 & feature0000025379

    Input 36
    	x=[0.71 0.78 0.   ... 0.   0.   0.54]
    	y=[1.]
    	y=[1.]
    	Explanation: feature0000006749 & ~feature0000010382 & ~feature0000013464 & ~feature0000013710 & ~feature0000022010 & ~feature0000025379 & feature0000026329

    Input 37
    	x=[0.79 0.97 0.   ... 0.57 0.23 0.62]
    	y=[1.]
    	y=[1.]
    	Explanation: ~feature0000006749 & feature0000010382 & ~feature0000013464 & feature0000013710 & ~feature0000022010 & feature0000025379 & ~feature0000026329

    Input 38
    	x=[0.74 0.25 0.   ... 0.   0.83 0.48]
    	y=[1.]
    	y=[1.]
    	Explanation: ~feature0000006749 & feature0000010382 & ~feature0000013464 & feature0000013710 & ~feature0000022010 & ~feature0000025379 & feature0000026329

    Input 39
    	x=[0.57 0.   0.   ... 0.29 0.39 0.26]
    	y=[1.]
    	y=[1.]
    	Explanation: ~feature0000006749 & feature0000010382 & ~feature0000013464 & feature0000013710 & ~feature0000022010 & ~feature0000025379 & feature0000026329

    Input 40
    	x=[0.56 0.24 1.   ... 0.   0.08 0.68]
    	y=[1.]
    	y=[1.]
    	Explanation: ~feature0000006749 & feature0000010382 & ~feature0000013464 & ~feature0000013710 & ~feature0000022010 & ~feature0000025379 & feature0000026329

# Combine local explanations

```python
dummy_names = [f'g{i}' for i in range(x.shape[1])]
```

```python
concept_names = genes.values.squeeze()
concept_names
```

    array(['ILMN_1343291', 'ILMN_1343295', 'ILMN_1651199', ...,
           'ILMN_3311165', 'ILMN_3311180', 'ILMN_3311190'], dtype=object)

```python
from sklearn.preprocessing import OneHotEncoder
y1h = OneHotEncoder(sparse=False).fit_transform(y.cpu().detach().numpy())
y2 = torch.FloatTensor(y1h).to(device)
```

```python
global_explanation, predictions, counter = logic.combine_local_explanations(model,
                                                                          x=x, y=y, method='pruning',
                                                                          target_class=0,
                                                                          topk_explanations=8,
                                                                          device=device)

accuracy, preds = logic.base.test_explanation(global_explanation, 0, x, y)
final_formula = logic.base.replace_names(global_explanation, concept_names)
final_formula_with_dummies = logic.base.replace_names(global_explanation, dummy_names)
print(f'Accuracy of when using the formula "{final_formula}": {accuracy:.4f}\n')
print(f'Accuracy of when using the formula "{final_formula_with_dummies}": {accuracy:.4f}')
```

    Accuracy of when using the formula "(~ILMN_1708983 & ~ILMN_1745049 & ~ILMN_1777811 & ~ILMN_2241168 & ~ILMN_3228700) | (ILMN_1777811 & ILMN_2241168 & ~ILMN_1708983 & ~ILMN_1745049 & ~ILMN_3228700 & ~ILMN_3243714) | (ILMN_3243714 & ~ILMN_1708983 & ~ILMN_1745049 & ~ILMN_1775520 & ~ILMN_1777811 & ~ILMN_3228700)": 0.9750

    Accuracy of when using the formula "(~g6749 & ~g10382 & ~g13710 & ~g22010 & ~g25379) | (g13710 & g22010 & ~g6749 & ~g10382 & ~g25379 & ~g26329) | (g26329 & ~g6749 & ~g10382 & ~g13464 & ~g13710 & ~g25379)": 0.9750

```python
global_explanation, predictions, counter = logic.combine_local_explanations(model,
                                                                          x=x, y=y, method='pruning',
                                                                          target_class=1,
                                                                          topk_explanations=10,
                                                                          device=device)

accuracy, preds = logic.base.test_explanation(global_explanation, 1, x, y)
final_formula = logic.base.replace_names(global_explanation, concept_names)
final_formula_with_dummies = logic.base.replace_names(global_explanation, dummy_names)
print(f'Accuracy of when using the formula "{final_formula}": {accuracy:.4f}\n')
print(f'Accuracy of when using the formula "{final_formula_with_dummies}": {accuracy:.4f}')
```

    Accuracy of when using the formula "(ILMN_1708983 & ILMN_1745049 & ~ILMN_1775520) | (ILMN_1777811 & ILMN_3228700 & ~ILMN_2241168) | (ILMN_1745049 & ILMN_1777811 & ~ILMN_1775520 & ~ILMN_2241168) | (ILMN_1745049 & ILMN_3228700 & ~ILMN_1775520 & ~ILMN_2241168) | (ILMN_1745049 & ILMN_3243714 & ~ILMN_1775520 & ~ILMN_2241168) | (ILMN_1777811 & ILMN_3243714 & ~ILMN_1708983 & ~ILMN_1775520 & ~ILMN_2241168) | (ILMN_1708983 & ILMN_3243714 & ~ILMN_1775520 & ~ILMN_1777811 & ~ILMN_2241168 & ~ILMN_3228700)": 0.9000

    Accuracy of when using the formula "(g6749 & g10382 & ~g13464) | (g13710 & g25379 & ~g22010) | (g10382 & g13710 & ~g13464 & ~g22010) | (g10382 & g25379 & ~g13464 & ~g22010) | (g10382 & g26329 & ~g13464 & ~g22010) | (g13710 & g26329 & ~g6749 & ~g13464 & ~g22010) | (g6749 & g26329 & ~g13464 & ~g13710 & ~g22010 & ~g25379)": 0.9000

```python
w, b = collect_parameters(model, device)
feature_weights = w[0]
feature_used_bool = np.sum(np.abs(feature_weights), axis=0) > 0
feature_used = np.nonzero(feature_used_bool)[0]
genes.iloc[feature_used]
```

      0

      6749
      ILMN_1708983

      10382
      ILMN_1745049

      13464
      ILMN_1775520

      13710
      ILMN_1777811

      22010
      ILMN_2241168

      25379
      ILMN_3228700

      26329
      ILMN_3243714

```python
global_explanation2 = logic.relu_nn.explain_global(model, n_classes=1, target_class=0, device=device)
explanation = logic.relu_nn.explain_global(model, n_classes=1, target_class=0, device=device)
if explanation not in ['False', 'True', 'The formula is too complex!']:
    accuracy, _ = logic.relu_nn.test_explanation(explanation, target_class=0, x=x.cpu(), y=y.cpu())
    print(f'Class {0} - Global explanation: "{global_explanation2}" - Accuracy: {accuracy:.4f}')
```

    Class 0 - Global explanation: "(feature0000006749 & feature0000010382 & feature0000013710) | (feature0000006749 & feature0000010382 & feature0000025379) | (feature0000006749 & feature0000010382 & feature0000026329) | (feature0000006749 & feature0000013710 & feature0000025379) | (feature0000006749 & feature0000013710 & feature0000026329) | (feature0000006749 & feature0000025379 & feature0000026329) | (feature0000010382 & feature0000013710 & feature0000025379) | (feature0000010382 & feature0000013710 & feature0000026329) | (feature0000010382 & feature0000025379 & feature0000026329) | (feature0000013710 & feature0000025379 & feature0000026329) | (feature0000006749 & feature0000010382 & ~feature0000013464) | (feature0000006749 & feature0000010382 & ~feature0000022010) | (feature0000006749 & feature0000013710 & ~feature0000013464) | (feature0000006749 & feature0000013710 & ~feature0000022010) | (feature0000006749 & feature0000025379 & ~feature0000022010) | (feature0000006749 & feature0000026329 & ~feature0000022010) | (feature0000010382 & feature0000013710 & ~feature0000013464) | (feature0000010382 & feature0000013710 & ~feature0000022010) | (feature0000010382 & feature0000025379 & ~feature0000022010) | (feature0000010382 & feature0000026329 & ~feature0000022010) | (feature0000013710 & feature0000025379 & ~feature0000022010) | (feature0000013710 & feature0000026329 & ~feature0000022010) | (feature0000025379 & feature0000026329 & ~feature0000022010) | (feature0000006749 & ~feature0000013464 & ~feature0000022010) | (feature0000010382 & ~feature0000013464 & ~feature0000022010) | (feature0000013710 & ~feature0000013464 & ~feature0000022010)" - Accuracy: 0.0250

## From previous works

ILMN_3286286, **ILMN_1775520**, ILMN_1656849, ILMN_1781198, ILMN_1665457

# Cross validation

```python
import copy
def train_nn(x_train, y_train, layers, device):
    set_seed(0)
    layers = [
        torch.nn.Linear(x_train.size(1), 200),
        torch.nn.LeakyReLU(),
        torch.nn.Linear(200, 100),
        torch.nn.LeakyReLU(),
        torch.nn.Linear(100, 50),
        torch.nn.LeakyReLU(),
        torch.nn.Linear(50, 10),
        torch.nn.LeakyReLU(),
        torch.nn.Linear(10, 1),
        torch.nn.Sigmoid(),
    ]
    model = torch.nn.Sequential(*layers).to(device)

    optimizer = torch.optim.Adam(model.parameters(), lr=0.0001)
    model.train()
    need_pruning = True
    for epoch in range(20000):
        # forward pass
        optimizer.zero_grad()
        y_pred = model(x_train)
        # Compute Loss
        loss = torch.nn.functional.binary_cross_entropy(y_pred, y_train)

        for module in model.children():
            if isinstance(module, torch.nn.Linear):
                loss += 0.05 * torch.norm(module.weight, 1)
                loss += 0.05 * torch.norm(module.bias, 1)
            break

        # backward pass
        loss.backward()
        optimizer.step()

        if epoch > 10000 and need_pruning and epoch % 1000 == 0:
            prune_features(model, 1, device)
            need_pruning = True

        # compute accuracy
        if epoch % 1000 == 0:
            y_pred_d = (y_pred > 0.5)
            accuracy = (y_pred_d.eq(y_train).sum(dim=1) == y_train.size(1)).sum().item() / y_train.size(0)
            print(f'Epoch {epoch}: train accuracy: {accuracy:.4f}')

    global_explanation, predictions, counter = logic.combine_local_explanations(model,
                                                                              x=x_train, y=y_train,
                                                                              method='pruning',
                                                                              target_class=0,
                                                                              topk_explanations=30,
                                                                              device=device)
#     global_explanation2, predictions, counter = logic.combine_local_explanations(model,
#                                                                               x=x_train, y=y_train,
#                                                                               method='pruning',
#                                                                               target_class=1,
#                                                                               topk_explanations=30,
#                                                                               device=device)

    return model, global_explanation#, global_explanation2
```

```python
n_splits = 5
skf = StratifiedKFold(n_splits=n_splits, shuffle=True, random_state=0)

explanations = []
model_accuracies = []
formula_accuracies = []
for i, (train_index, test_index) in enumerate(skf.split(x.cpu().detach().numpy(), y.cpu().detach().numpy())):
    print(f'Split [{i+1}/{n_splits}]')
    x_train, x_test = x[train_index], x[test_index]
    y_train, y_test = y[train_index], y[test_index]

    model, explanation = train_nn(x_train, y_train, layers, device)
    explanations.append(explanation)

    y_preds = (model(x_test) > 0.5).cpu().detach().numpy()
    accuracy = accuracy_score(y_test.cpu().detach().numpy(), y_preds)
    print(f'\t Accuracy: {accuracy:.4f}')
    model_accuracies.append(accuracy)

    accuracy, preds = logic.base.test_explanation(explanation, 0, x_test, y_test)
    final_formula = logic.base.replace_names(explanation, concept_names)
    print(f'\t Formula: "{final_formula}" \n\t Test accuracy (formula): {accuracy*100:.2f}%\n')
    formula_accuracies.append(accuracy)

#     accuracy, preds = logic.base.test_explanation(explanation, 1, x_test, y_test)
#     final_formula = logic.base.replace_names(explanation, concept_names)
#     print(f'\t Formula 2: "{final_formula}" \n\t Test accuracy (formula): {accuracy*100:.2f}%\n')
```

    Split [1/5]
    Epoch 0: train accuracy: 0.2500
    Epoch 1000: train accuracy: 0.7500
    Epoch 2000: train accuracy: 0.9062
    Epoch 3000: train accuracy: 1.0000
    Epoch 4000: train accuracy: 1.0000
    Epoch 5000: train accuracy: 1.0000
    Epoch 6000: train accuracy: 1.0000
    Epoch 7000: train accuracy: 1.0000
    Epoch 8000: train accuracy: 1.0000
    Epoch 9000: train accuracy: 1.0000
    Epoch 10000: train accuracy: 1.0000
    Epoch 11000: train accuracy: 1.0000
    Epoch 12000: train accuracy: 1.0000
    Epoch 13000: train accuracy: 1.0000
    Epoch 14000: train accuracy: 1.0000
    Epoch 15000: train accuracy: 1.0000
    Epoch 16000: train accuracy: 1.0000
    Epoch 17000: train accuracy: 1.0000
    Epoch 18000: train accuracy: 1.0000
    Epoch 19000: train accuracy: 1.0000
    	 Accuracy: 0.7500
    	 Formula: "(ILMN_1713875 & ILMN_1775520) | (ILMN_1713875 & ILMN_1823468) | (ILMN_1775520 & ILMN_1823468)"
    	 Test accuracy (formula): 75.00%

    Split [2/5]
    Epoch 0: train accuracy: 0.2500
    Epoch 1000: train accuracy: 0.7500
    Epoch 2000: train accuracy: 0.7500
    Epoch 3000: train accuracy: 0.7500
    Epoch 4000: train accuracy: 0.7500
    Epoch 5000: train accuracy: 0.7500
    Epoch 6000: train accuracy: 0.7500
    Epoch 7000: train accuracy: 0.7500
    Epoch 8000: train accuracy: 0.7500
    Epoch 9000: train accuracy: 0.7500
    Epoch 10000: train accuracy: 0.7500
    Epoch 11000: train accuracy: 0.7500
    Epoch 12000: train accuracy: 1.0000
    Epoch 13000: train accuracy: 1.0000
    Epoch 14000: train accuracy: 1.0000
    Epoch 15000: train accuracy: 1.0000
    Epoch 16000: train accuracy: 1.0000
    Epoch 17000: train accuracy: 1.0000
    Epoch 18000: train accuracy: 1.0000
    Epoch 19000: train accuracy: 0.9375
    	 Accuracy: 0.5000
    	 Formula: "~ILMN_1745049 & ~ILMN_1777650"
    	 Test accuracy (formula): 75.00%

    Split [3/5]
    Epoch 0: train accuracy: 0.2500
    Epoch 1000: train accuracy: 0.7500
    Epoch 2000: train accuracy: 0.7500
    Epoch 3000: train accuracy: 0.7500
    Epoch 4000: train accuracy: 0.7500
    Epoch 5000: train accuracy: 0.7500
    Epoch 6000: train accuracy: 1.0000
    Epoch 7000: train accuracy: 1.0000
    Epoch 8000: train accuracy: 1.0000
    Epoch 9000: train accuracy: 0.7500
    Epoch 10000: train accuracy: 1.0000
    Epoch 11000: train accuracy: 1.0000
    Epoch 12000: train accuracy: 0.9375
    Epoch 13000: train accuracy: 0.9375
    Epoch 14000: train accuracy: 0.9375
    Epoch 15000: train accuracy: 0.9375
    Epoch 16000: train accuracy: 0.9688
    Epoch 17000: train accuracy: 0.9688
    Epoch 18000: train accuracy: 0.9688
    Epoch 19000: train accuracy: 0.8438
    	 Accuracy: 0.6250
    	 Formula: "ILMN_3310514"
    	 Test accuracy (formula): 62.50%

    Split [4/5]
    Epoch 0: train accuracy: 0.2500
    Epoch 1000: train accuracy: 0.7500
    Epoch 2000: train accuracy: 0.7500
    Epoch 3000: train accuracy: 0.7500
    Epoch 4000: train accuracy: 0.7500
    Epoch 5000: train accuracy: 0.7500
    Epoch 6000: train accuracy: 0.7500
    Epoch 7000: train accuracy: 0.7500
    Epoch 8000: train accuracy: 0.7500
    Epoch 9000: train accuracy: 0.7500
    Epoch 10000: train accuracy: 0.8750
    Epoch 11000: train accuracy: 1.0000
    Epoch 12000: train accuracy: 1.0000
    Epoch 13000: train accuracy: 1.0000
    Epoch 14000: train accuracy: 1.0000
    Epoch 15000: train accuracy: 1.0000
    Epoch 16000: train accuracy: 1.0000
    Epoch 17000: train accuracy: 1.0000
    Epoch 18000: train accuracy: 1.0000
    Epoch 19000: train accuracy: 1.0000
    	 Accuracy: 0.7500
    	 Formula: "ILMN_2241168 | (ILMN_1893728 & ~ILMN_3310514) | (ILMN_3310514 & ~ILMN_1893728)"
    	 Test accuracy (formula): 87.50%

    Split [5/5]
    Epoch 0: train accuracy: 0.2500
    Epoch 1000: train accuracy: 0.7500
    Epoch 2000: train accuracy: 0.7500
    Epoch 3000: train accuracy: 0.7500
    Epoch 4000: train accuracy: 0.7500
    Epoch 5000: train accuracy: 1.0000
    Epoch 6000: train accuracy: 1.0000
    Epoch 7000: train accuracy: 1.0000
    Epoch 8000: train accuracy: 1.0000
    Epoch 9000: train accuracy: 1.0000
    Epoch 10000: train accuracy: 1.0000
    Epoch 11000: train accuracy: 0.9062
    Epoch 12000: train accuracy: 0.9062
    Epoch 13000: train accuracy: 0.9062
    Epoch 14000: train accuracy: 0.9062
    Epoch 15000: train accuracy: 0.9062
    Epoch 16000: train accuracy: 0.9062
    Epoch 17000: train accuracy: 0.9062
    Epoch 18000: train accuracy: 0.9062
    Epoch 19000: train accuracy: 0.9062
    	 Accuracy: 0.8750
    	 Formula: "ILMN_1893728"
    	 Test accuracy (formula): 75.00%

```python
model_accuracies = np.array(model_accuracies)
formula_accuracies = np.array(formula_accuracies)
print(f'Model accuracy: {model_accuracies.mean():.4f} +- {model_accuracies.std():.4f}')
print(f'Formula accuracy: {formula_accuracies.mean():.4f} +- {formula_accuracies.std():.4f}')
```

    Model accuracy: 0.7000 +- 0.1275
    Formula accuracy: 0.7500 +- 0.0791

```python
n_splits = 5
skf = StratifiedKFold(n_splits=n_splits, shuffle=True, random_state=0)
classifier = DecisionTreeClassifier(random_state=0)

classifier_accuracies = []
for i, (train_index, test_index) in enumerate(skf.split(x.cpu().detach().numpy(), y.cpu().detach().numpy())):
    print(f'Split [{i+1}/{n_splits}]')
    x_train, x_test = x[train_index].cpu().detach().numpy(), x[test_index].cpu().detach().numpy()
    y_train, y_test = y[train_index].cpu().detach().numpy().squeeze(), y[test_index].cpu().detach().numpy().squeeze()

    classifier.fit(x_train, y_train)
    y_preds = classifier.predict(x_test)
    accuracy = accuracy_score(y_test, y_preds)
    print(f'\tAccuracy: {accuracy:.4f}')
    classifier_accuracies.append(accuracy)

classifier_accuracies = np.array(classifier_accuracies)
print(f'Model accuracy: {classifier_accuracies.mean():.4f} +- {classifier_accuracies.std():.4f}')
```

    Split [1/5]
    	Accuracy: 0.6250
    Split [2/5]
    	Accuracy: 0.5000
    Split [3/5]
    	Accuracy: 0.7500
    Split [4/5]
    	Accuracy: 0.8750
    Split [5/5]
    	Accuracy: 0.5000
    Model accuracy: 0.6500 +- 0.1458

```python
plot_tree(classifier, feature_names=concept_names, filled=True)
plt.show()
```

```python
n_splits = 5
skf = StratifiedKFold(n_splits=n_splits, shuffle=True, random_state=0)
classifier = RandomForestClassifier(random_state=0)

classifier_accuracies = []
for i, (train_index, test_index) in enumerate(skf.split(x.cpu().detach().numpy(), y.cpu().detach().numpy())):
    print(f'Split [{i+1}/{n_splits}]')
    x_train, x_test = x[train_index].cpu().detach().numpy(), x[test_index].cpu().detach().numpy()
    y_train, y_test = y[train_index].cpu().detach().numpy().squeeze(), y[test_index].cpu().detach().numpy().squeeze()

    classifier.fit(x_train, y_train)
    y_preds = classifier.predict(x_test)
    accuracy = accuracy_score(y_test, y_preds)
    print(f'\tAccuracy: {accuracy:.4f}')
    classifier_accuracies.append(accuracy)

classifier_accuracies = np.array(classifier_accuracies)
print(f'Model accuracy: {classifier_accuracies.mean():.4f} +- {classifier_accuracies.std():.4f}')
```

    Split [1/5]
    	Accuracy: 0.8750
    Split [2/5]
    	Accuracy: 0.7500
    Split [3/5]
    	Accuracy: 0.7500
    Split [4/5]
    	Accuracy: 0.7500
    Split [5/5]
    	Accuracy: 0.7500
    Model accuracy: 0.7750 +- 0.0500


```python
import numpy as np
import itertools
```

## Compressed sensing

http://statweb.stanford.edu/~candes/papers/StableRecovery.pdf

$$
\begin{align}
&\text{Assume $\forall x_i \in X$;} \tag{the true signals}\\
&\quad\text{(1) $x_i \in \mathbb R^m$} \\
&\quad\text{(2) $x_i$ is sufficiently sparse} \\
&\text{Given $f(x_i) = Ax_i+e$ such that;} \tag{the function we know}\\
&\quad\text{(2) $A$ is approx orthonormal}\\
&\quad\text{(3) and $e$ is bounded} \\
&\text{Given $y_i = f(x_i)$ such that;} \tag{the observations}\\
&\quad\text{(2) where $n << m$} \\
x_i^* &= \text{min} \parallel \tilde x_i\parallel_1 \text{subject to} \parallel f(\tilde x_i) - y_i \parallel \le \epsilon \tag{?}\\
\end{align}
$$

Requirements.
- $x_i$ is sufficiently sparse
- $A$ is approximatley orthonormal
- $e$ is bounded (this was just needed to make the proofs easier?!? otherwise would need measures?)
- iid samples!?

#### Misunderstanding.

I thought we were matching the signals in a sparse basis.
We take our signal, $x_i$ and project it into a new basis where the signal has a sparse representation. We then match signals in this basis.

Rather the above requires $x_i$ itself to be sparse.

Oh, if $A$ is approximately orthonormal then $y_i$ will also be sparse!? No.

#### Search

Want a simple 2D case. Show how it is underdetermined.
> convex program searching, among all signals consistent with the data y

#### Structured spaces and regularisers and priors

The trick is using structure in the spaces, X/Y, to help us ...?
What about other structure in X or Y? How can that be translated into a regulariser?

Want to learn the distinct strutures of a space (e.g. smooth, sparse, ...) and then apply them when solving for x.

## Sparse representations

Ok, so can we just learn a sparse representation?
Take an AE, use a large hidden space, and regularise it to be sparse!?

# Isometry

Let $X$ and $Y$ be metric spaces with metrics $d_X$ and $d_Y$. A map $f : X \to Y$ is called an isometry or distance preserving if for any $a,b \in X$ one has

$$d_Y( f(a), f(b)) = d_X (a, b) $$

## Restricted isometry

> This property essentially requires that every set of columns with cardinality less than
$S$ approximately behaves like an orthonormal system. [source](http://statweb.stanford.edu/~candes/papers/StableRecovery.pdf)

$$
\begin{align}
(1-\delta_S)\parallel x \parallel_2^2 \le \parallel A_Sx \parallel \le (1+\delta_S)\parallel x \parallel_2^2 \\
\end{align}
$$

I am confused. If $A_{S_N}$ (all the columns) satisfies RIP then shouldnt the other $A_S$?

$$
\begin{align}
\parallel A_s^*A_s-I\parallel_{2\to 2}\leq \delta_s \\
\end{align}
$$

> A challenging aspect of RIP is its computational cost. Indeed, RIP is a property of the submatrices of a specific size. At present, no subexponential-time algorithm is known for testing RIP. [source](https://arxiv.org/pdf/0812.3137.pdf)

```python
def combinations(n, S):
    if S == 1:
        return list(itertools.combinations(range(n), 1))
    else:
        return list(itertools.combinations(range(n), S)) + combinations(n, S-1)
```

```python
def approx_orth(A, t):
    A_s = A[:, t]
    u, s, v = np.linalg.svd(np.dot(A_s.conj().T, A_s))
    e_S = np.abs(s - np.ones(len(t)))
    return np.max(e_S)
```

```python
def ft_matrix(N):
    i, j = np.meshgrid(np.arange(N), np.arange(N))
    omega = np.exp( - 2 * np.pi * 1J / N )
    W = np.power( omega, i * j ) / np.sqrt(N)
    return W
```

```python
def rim(A, S):
    T = combinations(A.shape[-1], S)
    return np.max([approx_orth(A, t) for t in T])
```

```python
[rim(ft_matrix(15), i) for i in range(1, 15)]
```

    [1.0658141036401503e-14,
     1.099120794378905e-14,
     1.2545520178264269e-14,
     1.3433698597964394e-14,
     1.3766765505351941e-14,
     1.4210854715202004e-14,
     1.4432899320127035e-14,
     1.4432899320127035e-14,
     1.4432899320127035e-14,
     1.4432899320127035e-14,
     1.4432899320127035e-14,
     1.4432899320127035e-14,
     1.4432899320127035e-14,
     1.4432899320127035e-14]

```python
A = np.random.standard_normal((10, 10))
A /= np.linalg.norm(A, axis=1, keepdims=True)
# axis = 0 or 1!?

[rim(A, i) for i in range(1, A.shape[-1])]
```

    [0.7082281239704467,
     1.1847212362098167,
     1.41341056257362,
     1.5968500301894166,
     1.6782821675460822,
     1.7341228419466281,
     1.7639518661661806,
     1.7694464508204697,
     1.7699056758830496]

## Non-linear RIP

Relationship to the gradient!? (and hessian?! where was that reference!?!?)

```python

```

## Modelling the data.

Can use signals in a sparse basis.
Thus the model is many causes/bases -> image.

What about other types of model?


```javascript
%%javascript
MathJax.Hub.Config({
    TeX: { equationNumbers: { autoNumber: "AMS" } }
});
MathJax.Hub.Queue(
  ["resetEquationNumbers", MathJax.InputJax.TeX],
  ["PreProcess", MathJax.Hub],
  ["Reprocess", MathJax.Hub]
);
```

    <IPython.core.display.Javascript object>

# Greek Letters

| Name    | lower case | lower LaTeX | upper case | upper LaTeX |
|:--------|:----------:|:------------|:----------:|:------------|
| alpha   | $\alpha$   | `\alpha`    | $A$        | `A`         |
| beta    | $\beta$    | `\beta`     | $B$        | `B`         |
| gamma   | $\gamma$   | `\gamma`    | $\Gamma$   | `\Gamma`    |
| delta   | $\delta$   | `\delta`    | $\Delta$   | `\Delta`    |
| epsilon | $\epsilon$ | `\epsilon`  | $E$        | `E`         |
| zeta    | $\zeta$    | `\zeta`     | $Z$        | `Z`         |
| eta     | $\eta$     | `\eta`      | $H$        | `H`         |
| theta   | $\theta$   | `\theta`    | $\Theta$   | `\Theta`    |
| iota    | $\iota$    | `\iota`     | $I$        | `I`         |
| kappa   | $\kappa$   | `\kappa`    | $K$        | `K`         |
| lambda  | $\lambda$  | `\lambda`   | $\Lambda$  | `\Lambda`   |
| mu      | $\mu$      | `\mu`       | $M$        | `M`         |
| nu      | $\nu$      | `\nu`       | $N$        | `N`         |
| xi      | $\xi$      | `\xi`       | $\Xi$      | `\Xi`       |
| omicron | $o$        | `o`         | $O$        | `O`         |
| pi      | $\pi$      | `\pi`       | $\Pi$      | `\pi`       |
| rho     | $\rho$     | `\rho`      | $P$        | `P`         |
| sigma   | $\sigma$   | `\sigma`    | $\Sigma$   | `\Sigma`    |
| tau     | $\tau$     | `\tau`      | $T$        | `T`         |
| upsilon | $\upsilon$ | `\upsilon`  | $\Upsilon$ | `\Upsilon`  |
| phi     | $\phi$     | `\phi`      | $\Phi$     | `\Phi`      |
| chi     | $\chi$     | `\chi`      | $X$        | `X`         |
| psi     | $\psi$     | `\psi`      | $\Psi$     | `\Psi`      |
| omega   | $\omega$   | `\omega`    | $\Omega$   | `\Omega`    |

# Limit

$$
f'(x) = \lim_{h \rightarrow 0} \frac{f(x+h) - f(x)}{h}
$$

$$
\lim_{x \rightarrow \infty} (1 + \frac{1}{x})^x = e
$$

$$
\lim_{n=1, 2, \ldots} a_n
$$

# Delta and Differentiation

$$
\lim_{\Delta x \rightarrow 0} \frac{\Delta y}{\Delta x} = \frac{dy}{dx}
$$

# Partial Derivatives

```
$$
\frac{\partial z}{\partial x} = \lim_{\Delta x \rightarrow 0} \frac{f(x+\Delta x, y) - f(x,y)}{\Delta x}
$$
```

$$
\frac{\partial z}{\partial x} = \lim_{\Delta x \rightarrow 0} \frac{f(x+\Delta x, y) - f(x,y)}{\Delta x}
$$

```latex
%%latex

\begin{align}
\label{partial}
\frac{\partial z}{\partial x} = \lim_{\Delta x \rightarrow 0} \frac{f(x+\Delta x, y) - f(x,y)}{\Delta x}
\end{align}
```

\begin{align}
\label{partial}
\frac{\partial z}{\partial x} = \lim_{\Delta x \rightarrow 0} \frac{f(x+\Delta x, y) - f(x,y)}{\Delta x}
\end{align}

Equation ($\ref{partial}$) is partial derivative

```python

```


# Model Selection for Causal Effect Model with the RScorer

```python
%load_ext autoreload
%autoreload 2
```

```python
import econml
```

```python
## Ignore warnings
import warnings
warnings.filterwarnings('ignore')
```

```python
# Main imports
from econml.dml import DML, LinearDML, SparseLinearDML, NonParamDML
from econml.metalearners import XLearner, TLearner, SLearner, DomainAdaptationLearner
from econml.dr import DRLearner

import numpy as np
from sklearn.ensemble import RandomForestRegressor, RandomForestClassifier
from sklearn.preprocessing import PolynomialFeatures
from sklearn.linear_model import LassoCV
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split

%matplotlib inline
```

## 2. Example Usage with Single Binary Treatment Synthetic Data

### 2.1. DGP
We use the following DGP:

\begin{align}
T \sim & \text{Bernoulli}\left(f(W)\right), &\; f(W)=\sigma(\langle W, \beta\rangle + \eta), \;\eta \sim \text{Uniform}(-1, 1)\\
Y = & T\cdot \theta(X) + \langle W, \gamma\rangle + \epsilon, & \; \epsilon \sim \text{Uniform}(-1, 1)\\
W \sim & \text{Normal}(0,\, I_{n_w}) & \\
X \sim & \text{Uniform}(0,\, 1)^{n_x}
\end{align}

where $W$ is a matrix of high-dimensional confounders, $\beta, \gamma$ have high sparsity and $\sigma$ is the sigmoid function.

For this DGP,
\begin{align}
\theta(x) = 1\{x_0 > .5\}
\end{align}

```python
# Treatment effect function
def exp_te(x):
    return x[:, 0] > 0.5

np.random.seed(123)
n = 5000
support_size = 5
n_x = 10
# Outcome support
support_Y = np.random.choice(range(n_x), size=support_size, replace=False)
coefs_Y = np.random.uniform(0, 1, size=support_size)
epsilon_sample = lambda n:np.random.uniform(-1, 1, size=n)
# Treatment support
support_T = support_Y
coefs_T = np.random.uniform(0, 1, size=support_size)
eta_sample = lambda n: np.random.uniform(-1, 1, size=n)

# Generate controls, covariates, treatments and outcomes
X = np.random.uniform(0, 1, size=(n, n_x))
# Heterogeneous treatment effects
TE = exp_te(X)
# Define treatment
log_odds = np.dot(X[:, support_T], coefs_T) + eta_sample(n)
T_sigmoid = 1/(1 + np.exp(-log_odds))
T = np.array([np.random.binomial(1, p) for p in T_sigmoid])
# Define the outcome
Y = TE * T + np.dot(X[:, support_Y], coefs_Y) + epsilon_sample(n)

# get testing data
X_test = np.random.uniform(0, 1, size=(n, n_x))
X_test[:, 0] = np.linspace(0, 1, n)
expected_te_test = exp_te(X_test)
```

### 2.2. Train Estimator

```python
reg = lambda: RandomForestRegressor(min_samples_leaf=10)
clf = lambda: RandomForestClassifier(min_samples_leaf=10)
```

```python
X_train, X_val, T_train, T_val, Y_train, Y_val = train_test_split(X, T, Y, test_size=.4)
```

```python
models = [('ldml', LinearDML(model_y=reg(), model_t=clf(), discrete_treatment=True,
                             linear_first_stages=False, cv=3)),
          ('sldml', SparseLinearDML(model_y=reg(), model_t=clf(), discrete_treatment=True,
                                    featurizer=PolynomialFeatures(degree=2, include_bias=False),
                                    linear_first_stages=False, cv=3)),
          ('xlearner', XLearner(models=reg(), cate_models=reg(), propensity_model=clf())),
          ('dalearner', DomainAdaptationLearner(models=reg(), final_models=reg(), propensity_model=clf())),
          ('slearner', SLearner(overall_model=reg())),
          ('tlearner', TLearner(models=reg())),
          ('drlearner', DRLearner(model_propensity=clf(), model_regression=reg(),
                                  model_final=reg(), cv=3)),
          ('rlearner', NonParamDML(model_y=reg(), model_t=clf(), model_final=reg(),
                                   discrete_treatment=True, cv=3)),
          ('dml3dlasso', DML(model_y=reg(), model_t=clf(), model_final=LassoCV(), discrete_treatment=True,
                             featurizer=PolynomialFeatures(degree=3),
                             linear_first_stages=False, cv=3))
]
```

```python
from joblib import Parallel, delayed

def fit_model(name, model):
    return name, model.fit(Y_train, T_train, X=X_train)

models = Parallel(n_jobs=-1, verbose=1)(delayed(fit_model)(name, mdl) for name, mdl in models)
```

    [Parallel(n_jobs=-1)]: Using backend LokyBackend with 8 concurrent workers.
    [Parallel(n_jobs=-1)]: Done   4 out of   9 | elapsed:   35.0s remaining:   43.8s
    [Parallel(n_jobs=-1)]: Done   9 out of   9 | elapsed:   54.0s finished

```python
from econml.score import RScorer

scorer = RScorer(model_y=reg(), model_t=clf(),
                 discrete_treatment=True, cv=3,
                 mc_iters=3, mc_agg='median')
scorer.fit(Y_val, T_val, X=X_val)
```

    <econml.score.rscorer.RScorer at 0x23a3270da90>

```python
rscore = [scorer.score(mdl) for _, mdl in models]
```

```python
expected_te_val = exp_te(X_val)
```

```python
rootpehe = [np.sqrt(np.mean((expected_te_val.flatten() - mdl.effect(X_val).flatten())**2)) for _, mdl in models]
```

```python
plt.scatter(rootpehe, rscore)
plt.xlabel('rpehe')
plt.ylabel('rscore')
plt.show()
```

### 2.3. Performance Visualization

```python
plt.figure(figsize=(16, 16))
rows = int(np.ceil(len(models) / 3))
for it, (name, mdl) in enumerate(models):
    plt.subplot(rows, 3, it + 1)
    plt.title('{}. RScore: {:.3f}, Root-PEHE: {:.3f}'.format(name, rscore[it], rootpehe[it]))
    plt.plot(X_test[:, 0], mdl.effect(X_test), label='{}'.format(name))
    plt.plot(X_test[:, 0], expected_te_test, 'b--', label='True effect')
    plt.ylabel('Treatment Effect')
    plt.xlabel('x')
    plt.legend()
plt.show()
```

# Getting the Best Model

```python
mdl, score  = scorer.best_model([mdl for _, mdl in models])
rootpehe_best = np.sqrt(np.mean((expected_te_val.flatten() - mdl.effect(X_val).flatten())**2))
plt.figure()
plt.title('RScore: {:.3f}, Root-PEHE: {:.3f}'.format(score, rootpehe_best))
plt.plot(X_test[:, 0], mdl.effect(X_test), label='best')
plt.plot(X_test[:, 0], expected_te_test, 'b--', label='True effect')
plt.ylabel('Treatment Effect')
plt.xlabel('x')
plt.legend()
plt.show()
```

# Getting an Ensemble based on Scores

```python
mdl, score  = scorer.ensemble([mdl for _, mdl in models])
rootpehe_ensemble = np.sqrt(np.mean((expected_te_val.flatten() - mdl.effect(X_val).flatten())**2))
plt.figure()
plt.title('RScore: {:.3f}, Root-PEHE: {:.3f}'.format(score, rootpehe_ensemble))
plt.plot(X_test[:, 0], mdl.effect(X_test), label='ensemble')
plt.plot(X_test[:, 0], expected_te_test, 'b--', label='True effect')
plt.ylabel('Treatment Effect')
plt.xlabel('x')
plt.legend()
plt.show()
```

# Semi-Synthetic Data

```python
reg = lambda: RandomForestRegressor(min_samples_leaf=10, random_state=123)
clf = lambda: RandomForestClassifier(min_samples_leaf=10, random_state=123)
```

```python
from econml.data.dgps import ihdp_surface_B, ihdp_surface_A
Y, T, X, expected_te = ihdp_surface_B(random_state=123)
```

```python
X_train, X_val, T_train, T_val,\
Y_train, Y_val, expected_te_train, expected_te_val = train_test_split(X, T, Y, expected_te, test_size=.3, random_state=123)
```

```python
models = [('ldml', LinearDML(model_y=reg(), model_t=clf(), discrete_treatment=True,
                             linear_first_stages=False, cv=3)),
          ('xlearner', XLearner(models=reg(), cate_models=reg(), propensity_model=clf())),
          ('dalearner', DomainAdaptationLearner(models=reg(), final_models=reg(), propensity_model=clf())),
          ('slearner', SLearner(overall_model=reg())),
          ('tlearner', TLearner(models=reg())),
          ('drlearner', DRLearner(model_propensity=clf(), model_regression=reg(),
                                  model_final=reg(), cv=3)),
          ('rlearner', NonParamDML(model_y=reg(), model_t=clf(), model_final=reg(),
                                   discrete_treatment=True, cv=3)),
          ('dml3dlasso', DML(model_y=reg(), model_t=clf(), model_final=LassoCV(), discrete_treatment=True,
                             featurizer=PolynomialFeatures(degree=2, interaction_only=True, include_bias=False),
                             linear_first_stages=False, cv=3))
]
```

```python
from joblib import Parallel, delayed

def fit_model(name, model):
    print("Training: ", name)
    model.fit(Y_train, T_train, X=X_train)
    print("Done training: ", name)
    return name, model

models = Parallel(n_jobs=-1, verbose=1)(delayed(fit_model)(name, mdl) for name, mdl in models)
```

    [Parallel(n_jobs=-1)]: Using backend LokyBackend with 8 concurrent workers.
    [Parallel(n_jobs=-1)]: Done   2 out of   8 | elapsed:    1.4s remaining:    4.3s
    [Parallel(n_jobs=-1)]: Done   8 out of   8 | elapsed:   10.2s finished

```python
from econml.score import RScorer

scorer = RScorer(model_y=reg(), model_t=clf(),
                 discrete_treatment=True, cv=3,
                 mc_iters=3, mc_agg='median')
scorer.fit(Y_val, T_val, X=X_val)
```

    <econml.score.rscorer.RScorer at 0x23a332a05f8>

```python
rscore = [scorer.score(mdl) for _, mdl in models]
```

```python
rootpehe = [np.sqrt(np.mean((expected_te.flatten() - mdl.effect(X).flatten())**2)) for _, mdl in models]
```

```python
plt.scatter(rootpehe, rscore)
plt.xlabel('rpehe')
plt.ylabel('rscore')
plt.show()
```

```python
best, score  = scorer.best_model([mdl for _, mdl in models])
rootpehe_best = np.sqrt(np.nanmean((expected_te_val.flatten() - best.effect(X_val).flatten())**2))
rootpehe_best
```

    0.7393917952588073

```python
ensemble, score  = scorer.ensemble([mdl for _, mdl in models])
rootpehe_ensemble = np.sqrt(np.nanmean((expected_te_val.flatten() - ensemble.effect(X_val).flatten())**2))
rootpehe_ensemble
```

    0.7346069084804471

```python
# Visualization of bias distribution
plt.figure(figsize=(15, 5))
plt.violinplot([np.abs(mdl.effect(X).flatten() - expected_te) for _, mdl in models] +
               [np.abs(best.effect(X).flatten() - expected_te)] +
               [np.abs(ensemble.effect(X).flatten() - expected_te)], showmeans=True)
plt.ylabel("Bias distribution")
plt.xticks(np.arange(1, len(models) + 3), [name for name, _ in models] + ['best', 'ensemble'])
plt.show()
```

notebooks/Causal Model Selection with the RScorer.ipynb


## Selection of secondary sampling units (SSUs)

To select the second stage sample, we need the second stage frame which is the list of all the households in the 10 selected clusters (psus). DHS, PHIA, MICS and other large scale surveys visit the selected clusters and construct the list of all households in the selected clusters.

Before starting the second stage selection, let us import the data from the first stage sampling information that is the first stage sample (psu_sample). For clarity, we explicitly import the packages and modules needed for this notebook.

```python
%%capture
%run psu_selection.ipynb

import numpy as np
import pandas as pd
from samplics.sampling import SampleSelection
```

In this tutorial, we will simulate the second stage frame. For the simulation, assume that the psu frame was obtained from a previous census conducted several years before. We also assume that, the change in the number of households since the previous census follows a normal distribution with a mean equal to 5% higher than the census value and a variance of 0.15 times the number of households from the census. Under these assumptions, we generate the following second stage frame of households. Note that the frame is created only for the selected PSUs.

```python
# Create a synthetic second stage frame
census_size = psu_frame.loc[
    psu_frame["psu_sample"] == 1, "number_households_census"
].values
stratum_names = psu_frame.loc[psu_frame["psu_sample"] == 1, "region"].values
cluster = psu_frame.loc[psu_frame["psu_sample"] == 1, "cluster"].values

np.random.seed(15)

listing_size = np.zeros(census_size.size)
for k in range(census_size.size):
    listing_size[k] = np.random.normal(1.05 * census_size[k], 0.15 * census_size[k])

listing_size = listing_size.astype(int)
hh_id = rr_id = cl_id = []
for k, s in enumerate(listing_size):
    hh_k1 = np.char.array(np.repeat(stratum_names[k], s)).astype(str)
    hh_k2 = np.char.array(np.arange(1, s + 1)).astype(str)
    cl_k = np.repeat(cluster[k], s)
    hh_k = np.char.add(np.char.array(cl_k).astype(str), hh_k2)
    hh_id = np.append(hh_id, hh_k)
    rr_id = np.append(rr_id, hh_k1)
    cl_id = np.append(cl_id, cl_k)

ssu_frame = pd.DataFrame(cl_id.astype(int))
ssu_frame.rename(columns={0: "cluster"}, inplace=True)
ssu_frame["region"] = rr_id
ssu_frame["household"] = hh_id

nb_obs = 15
print(f"\nFirst {nb_obs} observations of the SSU frame\n")
ssu_frame.head(nb_obs)
```

    First 15 observations of the SSU frame

      cluster
      region
      household

      0
      7
      North
      71

      1
      7
      North
      72

      2
      7
      North
      73

      3
      7
      North
      74

      4
      7
      North
      75

      5
      7
      North
      76

      6
      7
      North
      77

      7
      7
      North
      78

      8
      7
      North
      79

      9
      7
      North
      710

      10
      7
      North
      711

      11
      7
      North
      712

      12
      7
      North
      713

      13
      7
      North
      714

      14
      7
      North
      715

```python
psu_sample = psu_frame.loc[psu_frame["psu_sample"] == 1]
ssu_counts = ssu_frame.groupby("cluster").count()
ssu_counts.drop(columns="region", inplace=True)
ssu_counts.reset_index(inplace=True)
ssu_counts.rename(columns={"household": "number_households_listed"}, inplace=True)

pd.merge(
    psu_sample[["cluster", "region", "number_households_census"]],
    ssu_counts[["cluster", "number_households_listed"]],
    on=["cluster"],
)
```

      cluster
      region
      number_households_census
      number_households_listed

      0
      7
      North
      130
      130

      1
      10
      North
      600
      660

      2
      16
      South
      190
      195

      3
      24
      South
      75
      73

      4
      29
      South
      200
      217

      5
      34
      East
      305
      239

      6
      45
      East
      450
      398

      7
      52
      East
      700
      620

      8
      64
      West
      300
      301

      9
      86
      West
      280
      274

According to the simulated second stage frame, we get the same number of households in cluster 7 as the census. However, in strata 10, 16, 29, and 64, we listed more households than during than the census. And finally, we found less households in the remaining clusters than the census.

Now that we have a second stage frame, let's use *samplics* to calculate the probabilities of selection and to select a sample. The second stage sample size is 150 households and the strategy is to select 15 households per cluster.

### SSU (household) Probability of Selection

The second stage probabilities of selection are conditional on the first stage realization. For this stage, simple random selection (srs) and systematic selection(sys) are common methods used to select households. For this example, we use srs to select 15 households from each cluster. Conditionally to teh first stage, the second stage selection is a stratified srs where the clusters are the strata. More generally, we have that \begin{equation} p_{hij} = \frac{m_{hi}}{M_{hi}^{'}} \end{equation} where $p_{hij}$ is the conditional probability of selection for unit $j$ from stratum $h$ and cluster $j$, $m_{hi}$ and $M_{hi}^{'}$ are the sample size and the number of secondary sampling units listed for stratum $h$ and cluster $j$, respectively.

In this scenario, sample size is the same in each stratum. Hence, the parameter *sample_size* does not need to be a Python dictionary; we will only provide 15 in the function call.

```python
stage2_design = SampleSelection(
    method="srs", stratification=True, with_replacement=False
)

ssu_frame["ssu_prob"] = stage2_design.inclusion_probs(
    ssu_frame["household"], 15, ssu_frame["cluster"]
)

ssu_frame.sample(20)
```

      cluster
      region
      household
      ssu_prob

      1438
      34
      East
      34164
      0.062762

      2517
      52
      East
      52606
      0.024194

      2943
      86
      West
      86111
      0.054745

      3002
      86
      West
      86170
      0.054745

      559
      10
      North
      10430
      0.022727

      1216
      29
      South
      29159
      0.069124

      2751
      64
      West
      64220
      0.049834

      412
      10
      North
      10283
      0.022727

      2549
      64
      West
      6418
      0.049834

      1508
      34
      East
      34234
      0.062762

      1517
      45
      East
      454
      0.037688

      2658
      64
      West
      64127
      0.049834

      1072
      29
      South
      2915
      0.069124

      1074
      29
      South
      2917
      0.069124

      1669
      45
      East
      45156
      0.037688

      1286
      34
      East
      3412
      0.062762

      1162
      29
      South
      29105
      0.069124

      1051
      24
      South
      2467
      0.205479

      3085
      86
      West
      86253
      0.054745

      2148
      52
      East
      52237
      0.024194

### SSU (household) Selection

The second stage sample is selected from the SSU frame (*ssu_frame*) using the variable *cluster* as the stratification variable. The sample is selected without replacement according to the specification of the second stage design. Hence, both *ssu_sample* and *ssu_hits* sum to 150 and each selected household was hit only ounce (i.e. *ssu_hits* = 1).
```python
>>> ssu_frame["ssu_sample"].sum()
>>> 150
>>> ssu_frame["ssu_hits"].sum()
>>> 150
```

```python
np.random.seed(11)
ssu_sample, ssu_hits, ssu_probs = stage2_design.select(
    ssu_frame["household"], 15, ssu_frame["cluster"]
)

ssu_frame["ssu_sample"] = ssu_sample
ssu_frame["ssu_hits"] = ssu_hits
ssu_frame["ssu_probs"] = ssu_probs

ssu_frame[ssu_frame["ssu_sample"] == 1].sample(15)
```

      cluster
      region
      household
      ssu_prob
      ssu_sample
      ssu_hits
      ssu_probs

      2319
      52
      East
      52408
      0.024194
      1
      1
      0.024194

      2931
      86
      West
      8699
      0.054745
      1
      1
      0.054745

      2642
      64
      West
      64111
      0.049834
      1
      1
      0.049834

      122
      7
      North
      7123
      0.115385
      1
      1
      0.115385

      60
      7
      North
      761
      0.115385
      1
      1
      0.115385

      945
      16
      South
      16156
      0.076923
      1
      1
      0.076923

      338
      10
      North
      10209
      0.022727
      1
      1
      0.022727

      2218
      52
      East
      52307
      0.024194
      1
      1
      0.024194

      2870
      86
      West
      8638
      0.054745
      1
      1
      0.054745

      1764
      45
      East
      45251
      0.037688
      1
      1
      0.037688

      630
      10
      North
      10501
      0.022727
      1
      1
      0.022727

      1441
      34
      East
      34167
      0.062762
      1
      1
      0.062762

      986
      24
      South
      242
      0.205479
      1
      1
      0.205479

      1796
      45
      East
      45283
      0.037688
      1
      1
      0.037688

      1264
      29
      South
      29207
      0.069124
      1
      1
      0.069124

To use systematic selection, we just need to replace ```method="srs"``` by ```method="sys"```.

Another common approach is to use a rate for selecting the sample. Instead of selecting 15 households from 130 in the first cluster, we may want to select with a rate of 15/130, and similarly for the other clusters.

```python
rates = np.repeat(15, 10) / ssu_counts["number_households_listed"].values

ssu_rates = dict(zip(np.unique(ssu_frame["cluster"]), rates))

ssu_rates
```

    {7: 0.11538461538461539,
     10: 0.022727272727272728,
     16: 0.07692307692307693,
     24: 0.2054794520547945,
     29: 0.06912442396313365,
     34: 0.06276150627615062,
     45: 0.03768844221105527,
     52: 0.024193548387096774,
     64: 0.04983388704318937,
     86: 0.05474452554744526}

A sample is selected using the rates as follows:

```python
np.random.seed(22)

stage2_design2 = SampleSelection(
    method="sys", stratification=True, with_replacement=False
)

ssu_sample_r, ssu_hits_r, _ = stage2_design2.select(
    ssu_frame["household"], stratum=ssu_frame["cluster"], samp_rate=ssu_rates
)

ssu_sample2 = pd.DataFrame(
    data={
        "household": ssu_frame["household"],
        "ssu_sample_r": ssu_sample_r,
        "ssu_hits_r": ssu_hits_r,
    }
)

ssu_sample2.head(25)
```

      household
      ssu_sample_r
      ssu_hits_r

      0
      71
      0
      0

      1
      72
      0
      0

      2
      73
      0
      0

      3
      74
      0
      0

      4
      75
      0
      0

      5
      76
      0
      0

      6
      77
      0
      0

      7
      78
      0
      0

      8
      79
      0
      0

      9
      710
      0
      0

      10
      711
      0
      0

      11
      712
      0
      0

      12
      713
      0
      0

      13
      714
      0
      0

      14
      715
      0
      0

      15
      716
      0
      0

      16
      717
      0
      0

      17
      718
      0
      0

      18
      719
      0
      0

      19
      720
      0
      0

      20
      721
      0
      0

      21
      722
      0
      0

      22
      723
      0
      0

      23
      724
      0
      0

      24
      725
      0
      0

Let's store the first and second stages samples.

```python
psu_sample[["cluster", "region", "psu_prob"]].to_csv("psu_sample.csv")

ssu_sample = ssu_frame.loc[ssu_frame["ssu_sample"] == 1]
ssu_sample[["cluster", "household", "ssu_prob"]].to_csv("ssu_sample.csv")
```