Colocated cokriging in python with r

colocated_cokriging_r.py

    def cokriging_in_r(self, x, y, z):
        '''
        Cokriging (and ordinary kriging) is quite fast in R.
        This would anyway be more pragmatic than rewriting/porting it to Python.
        For the moment this will be the 'best' way as R makes it very easy to
        use kriging without fitting a variogram model, but using a standard 
        variogram.
        '''
        import rpy2
        import rpy2.robjects as robj
        import time
        time0 = time.time()
        robj.r.library('gstat')
        self.create_interpolation_grid()
        xi, yi = robj.FloatVector(self.xi.tolist()), robj.FloatVector(self.yi.tolist())
        dataset = self.dataloader.dataset.ReadAsArray().flatten()
        mask = numpy.equal(dataset, -9999)
        rxi = robj.FloatVector(dataset.tolist())
        radar = self.get_radar_for_locations()
        radar = robj.FloatVector(radar)
        x,y,z = robj.FloatVector(x), robj.FloatVector(y), robj.FloatVector(z)
        rain_frame = robj.DataFrame({'x': x, 'y': y, 'z':z})
        radar_frame = robj.DataFrame({'x': xi, 'y': yi, 'radar': rxi})
        target_frame = robj.DataFrame({'x':xi, 'y':yi})
        vgm_args = {'model_type': 'Sph', 'range1': 20000, 'range2':1040000}
        # following command gives a python output to r functions
        vgm_args['nugget'] = 0.035
        vgm_args['sill1'] = 0.473
        vgm_args['sill2'] = 8.994
        radar_variance = dataset[~mask].var()
        rain_variance = robj.r.var(z)[0]
        correlation_radar_rain = numpy.abs(robj.r.cor(z, radar))
        variance_correction = numpy.sqrt(rain_variance * radar_variance)
        # The cross coefficient is used in the cross variogram (crossgram)
        # 
        cross_coef = (correlation_radar_rain * variance_correction)[0]
        # change it back to rpy strict

        # The variogram is combined. This is a bit awkward in Rpy.
        # So one way is change the args manually (see below)
        # or load variables in R before hand.
        variogram_radar = robj.r.vgm(vgm_args['sill1'] * radar_variance, 
            vgm_args['model_type'], vgm_args['range1'], vgm_args['nugget'], 
            robj.r("add.to=vgm("+ str(vgm_args['sill2'] * radar_variance)
                + ", 'Sph', 1040000, " + str(vgm_args['nugget'] * 
                    radar_variance)+")"))
        variogram_rain = robj.r.vgm(vgm_args['sill1'] * rain_variance,
            vgm_args['model_type'], vgm_args['range1'], vgm_args['nugget'], 
            robj.r("add.to=vgm("+ str(vgm_args['sill2'] * rain_variance)
                + ", 'Sph', 1040000, " + str(vgm_args['nugget'] * 
                    rain_variance)+")"))
        crossgram = robj.r.vgm(vgm_args['sill1'] * cross_coef, 
            vgm_args['model_type'], vgm_args['range1'], vgm_args['nugget'], 
            robj.r("add.to=vgm("+ str(vgm_args['sill2'] * cross_coef)
                + ", 'Sph', 1040000, " + str(vgm_args['nugget'] * cross_coef) + ")"))

        # #waarde van de radar op de modelnodes (=predictielocatie), nodig voor colocated cokriging (betekent sec. variabele op elke predictie locatie bekend)
        #         mn.radar=krige(rain~1,rads,mn,nmax=1)
        #         names(mn.radar)[1]='radar.rain'
        
        # # gstat object voor colocated cokriging                            
        # g.cck=NULL
        #         g.cck=gstat(g.cck,id='gauge',formula=sqrt(rain)~1,data=raingauge,model=vgm.raing,nmax = 40)
        #         g.cck=gstat(g.cck,id='radar',sqrt(radar.rain)~1,data=mn.radar,model=vgm.radar,merge=c('gauge','radar'),nmax=1)
        #         g.cck=gstat(g.cck,id=c('gauge','radar'),model=vgm.cross,nmax = 40)
                
        # #cokriging naar modelknooppunten (mn)
        #         cck.out=predict(g.cck,mn,nsim=0) 

        # radar_krige = (robj.r.krige(robj.r('radar ~ 1'), robj.r('~ x + y'),
        #     radar_frame, target_frame, nmax=1))
        cck = robj.r('NULL')
        cck = robj.r.gstat(cck, "rain", robj.r('z ~ 1'), robj.r('~ x + y'), 
            data=rain_frame, model=variogram_rain, nmax=40)
        cck = robj.r.gstat(cck, "radar", robj.r('radar~ 1'), robj.r('~ x + y'),
            data=radar_frame, model=variogram_radar, 
            merge=robj.r("c('rain','radar')"), nmax=1)
        cck = robj.r.gstat(cck, robj.r('c("rain", "radar")'), model=crossgram, nmax=40)        
        # pdb.set_trace()

        result = robj.r.predict(cck, target_frame)
        time1 = time.time()
        deltatime = time1-time0
        print('Kriging took '+ str(deltatime) + ' seconds')
        radar_est = numpy.array(result[4])
        radar_est = radar_est.reshape((self.ny, self.nx))
        rain_est = numpy.array(result[2])
        rain_est = radar_est.reshape((self.ny, self.nx))
        return radar_est, rain_est