import numpy as np
import matplotlib.pyplot as pl

pl.close('all')

scalar = np.random.random(49).reshape((7,7))   # j,i
u      = np.random.random(49).reshape((7,7))   # j,i
v      = np.random.random(49).reshape((7,7))   # j,i

i0 = 3  # central grid point
j0 = 3  # central grid point

dxy = 2.

n = 2       # Averaging domain (+/- n grid points)
s = 2*n+1   # Size of averaging domain

# Start/end indices of averaging domain
istart = i0-n
iend   = i0+n+1

jstart = j0-n
jend   = j0+n+1

# --------------
# Manual calculation
# --------------
tend_1 = np.zeros_like(scalar)

for i in range(istart, iend):
    for j in range(jstart, jend):
        tend_1[j,i] = -u[j,i] * (scalar[j,i+1] - scalar[j,i-1]) / (2*dxy) \
                      -v[j,i] * (scalar[j+1,i] - scalar[j-1,i]) / (2*dxy)

# --------------
# Manual numpy slices
# --------------
ic  = np.s_[istart:iend]
jc  = np.s_[jstart:jend]

ip2 = np.s_[istart+2:iend+2]    # [j,i+2]
ip1 = np.s_[istart+1:iend+1]    # [j,i+1]
im1 = np.s_[istart-1:iend-1]    # [j,i-1]
im2 = np.s_[istart-2:iend-2]    # [j,i-2]

jp2 = np.s_[jstart+2:jend+2]    # [j+2,i]
jp1 = np.s_[jstart+1:jend+1]    # [j+1,i]
jm1 = np.s_[jstart-1:jend-1]    # [j-1,i]
jm2 = np.s_[jstart-2:jend-2]    # [j-2,i]

tend_2 = np.zeros_like(scalar)
tend_2[jc,ic] = -u[jc,ic] * (scalar[jc,ip1] - scalar[jc,im1]) / (2*dxy) \
                -v[jc,ic] * (scalar[jp1,ic] - scalar[jm1,ic]) / (2*dxy)

# --------------
# Slice class
# --------------
class Slice:
    def __init__(self, istart, iend, jstart, jend):
        self.istart = istart
        self.iend   = iend
        self.jstart = jstart
        self.jend   = jend

    def __getitem__(self, pos):
        return np.s_[self.jstart+pos[0]:self.jend+pos[0],\
                     self.istart+pos[1]:self.iend+pos[1]]

    def __call__(self, j, i):
        return np.s_[self.jstart+j:self.jend+j,\
                     self.istart+i:self.iend+i]

offs = Slice(istart, iend, jstart, jend)

tend_3 = np.zeros_like(scalar)
tend_3[ offs(0,0) ] = -u[ offs(0,0) ] * (scalar[ offs(0,+1) ] - scalar[ offs(0,-1) ]) / (2*dxy) \
                      -v[ offs(0,0) ] * (scalar[ offs(+1,0) ] - scalar[ offs(-1,0) ]) / (2*dxy)
# OR:
tend_3[ offs(0,0) ] = -u[ offs[0,0] ] * (scalar[ offs[0,+1] ] - scalar[ offs[0,-1] ]) / (2*dxy) \
                      -v[ offs[0,0] ] * (scalar[ offs[+1,0] ] - scalar[ offs[-1,0] ]) / (2*dxy)


print(np.all(tend_2 == tend_1), np.all(tend_3 == tend_1))




pl.figure()

pl.subplot(131)
pl.imshow(tend_1)
pl.colorbar()

pl.subplot(132)
pl.imshow(tend_2)
pl.colorbar()

pl.subplot(133)
pl.imshow(tend_3)
pl.colorbar()