#!/usr/bin/python

import numpy as np
import matplotlib.pyplot as plt

# ========= functions

def blob(p,w,Xs,Ys):
    # helps generate nice test-function
    dXs = Xs-p[0]
    dYs = Ys-p[1]
    return 1/( 1 + (dXs**2+dYs**2)/(w*w) )

def matchPBC(f,Xs,Ys):
   '''
      match periodic boundary condions iteratevily (one axis at a tim)
   '''
    f=f.copy()
    ### X-boundary
    Dxs = f[:,0] - f[:,-1];    # difference for each pixel at boundary
    f += Dxs[:,None]*(Xs-0.5)  # linear interpolation for each x-line[!
    ### Y-boundary
    Dys = f[0,:] - f[-1,:];
    f += Dys[None,:]*(Ys-0.5)
    return f

# ========= main

### perpare support arryas
xs = np.linspace(0.0,1.0,100,endpoint=False)
ys = np.linspace(0.0,1.0,100,endpoint=False)
Xs,Ys = np.meshgrid(xs,ys)

#### test function
#f = np.sin(Xs*6) + np.sin(Xs*11) + np.cos(Xs*7-Ys*9) + np.sin(Xs*13+Ys*14) #+ np.sin(Xs*33+Ys*24)
#f = np.cos(Xs*3.5)  + np.sin(Ys*8)
#f = (Xs*5 + Ys*16)/21.0
f = blob([0.0,0.9],0.2,Xs,Ys) + blob([0.5,0.0],0.2,Xs,Ys)

#### do it !!!
fpbc = matchPBC(f,Xs,Ys)

#### plot results !!!
plt.subplot(1,2,1); plt.imshow( np.tile(f   ,(2,2)) )
plt.subplot(1,2,2); plt.imshow( np.tile(fpbc,(2,2)) )
plt.show()