## 
# chain operators demo
##
from rsf.proj import *
import sys,wplot
sys.path.append('../../PYUTIL')
from difLOP import *
from chainLOP import *
from DPTEST import *

from LICGsolver import *
from LICGutils  import *

# ------------------------------------------------------------
do3D = 'y'
doINV = 'y'
# ------------------------------------------------------------

par = dict(
    nx=25,  ox=0, dx=1.0,  lx='x', ux='',
    ny=11,  oy=0, dy=1.0,  ly='y', uy='',
    nz=15,  oz=0, dz=1.0,  lz='z', uz='',
    otype = "c",
    N = 4
    )
wplot.param(par)

# ------------------------------------------------------------
Flow('s2D',None,
    '''
    spike nsp=1 mag=1 k1=8 k2=13
    n1=%(nz)d o1=%(oz)g d1=%(dz)g label1="" unit1=""
    n2=%(nx)d o2=%(ox)g d2=%(dx)g label2="" unit2=""
    '''%par)
Result('s2D',wplot.igrey2d('',par))

Flow('v2D',['s2D','s2D'],'cat axis=3 space=n ${SOURCES[1:2]}')

if(do3D=='y'):
    Flow('s3D',None,
    '''
    spike nsp=1 mag=1 k1=8 k2=13 k3=6
    n1=%(nz)d o1=%(oz)g d1=%(dz)g label1="" unit1=""
    n2=%(nx)d o2=%(ox)g d2=%(dx)g label2="" unit2=""
    n3=%(ny)d o3=%(oy)g d3=%(dy)g label3="" unit3=""
    '''%par)
    Result('s3D','byte gainpanel=a pclip=100|'+wplot.igrey3d('',par))

    Flow('v3D',['s3D','s3D','s3D'],'cat axis=4 space=n ${SOURCES[1:3]}')
    
# ------------------------------------------------------------
# ------------------------------------------------------------
# ------------------------------------------------------------
# ------------------------------------------------------------
# gradient
G2D = grad2d(par['otype'],par['N'])
G2D.FORW(  's2D',   'grd2Dfor')
G2D.ADJT('grd2Dadj','grd2Dfor')
Result('grd2Dfor',wplot.igrey2d('',par))
Result('grd2Dadj',wplot.igrey2d('',par))

if(do3D=='y'):
    G3D = grad3d(par['otype'],par['N'])
    G3D.FORW(  's3D',   'grd3Dfor')
    G3D.ADJT('grd3Dadj','grd3Dfor')
    Result('grd3Dfor','byte gainpanel=a pclip=100|'+wplot.igrey4d('',par))
    Result('grd3Dadj','byte gainpanel=a pclip=100|'+wplot.igrey4d('',par))

# ------------------------------------------------------------
# divergence
D2D = div2d(par['otype'],par['N'])
D2D.FORW(  'v2D',   'div2Dfor')
D2D.ADJT('div2Dadj','div2Dfor')
Result('div2Dfor',wplot.igrey2d('',par))
Result('div2Dadj',wplot.igrey2d('',par))

if(do3D=='y'):
    D3D = div3d(par['otype'],par['N'])
    D3D.FORW(  'v3D',   'div3Dfor')
    D3D.ADJT('div3Dadj','div3Dfor')
    Result('div3Dfor','byte gainpanel=a pclip=100|'+wplot.igrey4d('',par))
    Result('div3Dadj','byte gainpanel=a pclip=100|'+wplot.igrey4d('',par))

# ------------------------------------------------------------
# Laplacian
L2D = lapl2d(par['otype'],par['N'])
L2D.FORW(   's2D',  'lap2Dfor')
L2D.ADJT('lap2Dadj','lap2Dfor')
Result('lap2Dfor',wplot.igrey2d('',par))
Result('lap2Dadj',wplot.igrey2d('',par))

if(do3D=='y'):
    L3D = lapl3d(par['otype'],par['N'])
    L3D.FORW(  's3D',   'lap3Dfor')
    L3D.ADJT('lap3Dadj','lap3Dfor')
    Result('lap3Dfor','byte gainpanel=a pclip=100|'+wplot.igrey4d('',par))
    Result('lap3Dadj','byte gainpanel=a pclip=100|'+wplot.igrey4d('',par))

# ------------------------------------------------------------
# ------------------------------------------------------------
# ------------------------------------------------------------
# ------------------------------------------------------------
# chain: LAPL( DIV( GRAD() ) )
H2D = chainop(par,[G2D,D2D,L2D])
H2D.FORW(   's2D',  'chn2Dfor')
H2D.ADJT('chn2Dadj','chn2Dfor')
Result('chn2Dfor',wplot.igrey2d('',par))
Result('chn2Dadj',wplot.igrey2d('',par))

DH2D = DPTEST(H2D,['chn2Dadj'],['chn2Dfor'])           
DH2D.RUN()

if(do3D=='y'):
    H3D = chainop(par,[G3D,D3D,L3D])
    H3D.FORW(  's3D',   'chn3Dfor')
    H3D.ADJT('chn3Dadj','chn3Dfor')
    Result('chn3Dfor','byte gainpanel=a pclip=100|'+wplot.igrey4d('',par))
    Result('chn3Dadj','byte gainpanel=a pclip=100|'+wplot.igrey4d('',par))
    
    DH3D = DPTEST(H3D,['chn3Dadj'],['chn3Dfor'])           
    DH3D.RUN()

# ------------------------------------------------------------
if(doINV=='y'):
    NITER = 10

    Flow('z2D','s2D','math output=0')
    
    L2D = LICG('I2D',par, Lop=H2D, d=['chn2Dfor'], xo=['z2D'], xf=['inv2D'], niter=NITER)
    L2D.iterate()
    plotOFUN('I2D',NITER,par)
    Result('inv2D',wplot.igrey2d('',par))

    Result('mod2D','I2D/MOD.rsf',wplot.igrey2d('',par))
    Result('res2D','I2D/RES0.rsf',wplot.igrey2d('',par))
    
    if(do3D=='y'):
        
        Flow('z3D','s3D','math output=0')

        L3D = LICG('I3D',par, Lop=H3D, d=['chn3Dfor'], xo=['z3D'], xf=['inv3D'], niter=NITER)
        L3D.iterate()
        plotOFUN('I3D',NITER,par)
        Result('inv3D','byte gainpanel=a pclip=100|'+wplot.igrey4d('',par))

        Result('mod3D','I3D/MOD.rsf',
            'byte gainpanel=a pclip=100|'+wplot.igrey4d('',par))
        Result('res3D','I3D/RES0.rsf',
            'byte gainpanel=a pclip=100|'+wplot.igrey4d('',par))
    
End()
