# Dominant frequency versus depth for scalar waves in von Karman 3-D medium
#
# December 2007
#
# Thomas Jules Browaeys 
# Bureau of Economic Geology
# University of Texas at Austin
# mailto:jules.browaeys@beg.utexas.edu


from rsf.proj import *
from math import pi


# ----------------------------------------------
# DEPTH AND FREQUENCY GRIDS
# ----------------------------------------------

pgrid = {'nx':150,  # Frequency points
         'ox':1.,
         'dx':1.,   # Frequency steps
         'nz':150,  # Depth points
         'oz':1.,
         'dz':10.   # Depth steps
         }

# Create frequency and depth 2-D grid
Flow('dfgrd',None,
     '''
     spike nsp=1 mag=1 n1=%(nx)d d1=%(dx)g o1=%(ox)g n2=%(nz)d d2=%(dz)g o2=%(oz)g |
     put label1=f unit1=Hz label2=z unit2=m
     ''' % pgrid)


# ----------------------------------------------
# STOCHASTIC MEDIUM
# ----------------------------------------------
# hu = H exponent > -0.5 
# cu = H spectrum constant

# Exponents and constants
pgrids = {'M025':{'hu':-0.25, 'cu':1.311,  'shu':'\F15 8\F2 0.25'},
           '025':{'hu':0.25,  'cu':0.5991, 'shu':'0.25'},
            '05':{'hu':0.5,   'cu':1.0,    'shu':'0.5'},
           '075':{'hu':0.75,  'cu':1.3317, 'shu':'0.75'},
             '1':{'hu':1.0,   'cu':1.57,   'shu':'1.0'}
          }


# Correlation lengths (m)
bvalues = {'b25':2.5, 'b05':5.0, 'b10':10.0}
bvalstr = {'b25':'2.5', 'b05':'5', 'b10':'10'}

# VP VS variations and frequency
parq = {
    'f0':60.,    # Ricker wavelet frequency
    'sd':0.3,    # standard deviation
    'cp':2700.,  # P wave velocity
    'cs':1230.   # S wave velocity
    }

# Variance
parq['s2'] = parq['sd']*parq['sd']


# Ricker wavelet
Flow('srick','dfgrd','math output="(x1/(%g))^2" | math output="input*exp(-input)"' % (parq['f0']))
Flow('psrick','srick','window n2=1 f2=1 | window | put label2=S unit2=')


for lb in bvalues.keys():

    parq['b'] = bvalues[lb]
    bstring = bvalstr[lb]

    ksb  = 'ksb'   + lb
    ksb2 = 'ksb2'  + lb
    kpb  = 'kpb'   + lb
    kpb2 = 'kpb2'  + lb

    # Wave vector k*b for VS
    Flow(ksb,'dfgrd','math output="(%g)*x1/(%g)"' % (2.*pi*parq['b'],parq['cs']))
    # Squared (k*b) for VS
    Flow(ksb2,ksb,'math output="input*input"')

    # Wave vector k*b for VP
    Flow(kpb,'dfgrd','math output="(%g)*x1/(%g)"' % (2.*pi*parq['b'],parq['cp']))
    # Squared (k*b) for VP
    Flow(kpb2,kpb,'math output="input*input"')


    # Exponential correlation
    # -----------------------

    depthse = 'depthse'  + lb
    depthpe = 'depthpe'  + lb

    # Depth penetration VS
    Flow(depthse,ksb2,'math k2=${SOURCES[0]} output="(%g)/(%g)*(1.+4.*k2)/(k2*k2)"' % (parq['b'],4.*parq['s2']))
    # Depth penetration VP
    Flow(depthpe,kpb2,'math k2=${SOURCES[0]} output="(%g)/(%g)*(1.+4.*k2)/(k2*k2)"' % (parq['b'],4.*parq['s2']))

    specse = 'specse'  + lb
    sumse  = 'sumse'   + lb
    fumse  = 'fumse'   + lb
    fdomse = 'fdomse'  + lb

    specpe = 'specpe'  + lb
    sumpe  = 'sumpe'   + lb
    fumpe  = 'fumpe'   + lb
    fdompe = 'fdompe'  + lb

    fdome = 'fdome' + lb

    # Peak frequency versus depth VS
    Flow(specse,[depthse, 'srick'],'math lc=${SOURCES[0]} s=${SOURCES[1]} output="s*exp(-x2/lc)"')
    Flow(sumse,specse,'stack axis=1 norm=n | window')
    Flow(fumse,specse,'math output="x1*input" | stack axis=1 norm=n | window')
    Flow(fdomse,[fumse, sumse],'math f=${SOURCES[0]} s=${SOURCES[1]} output="f/s" | put unit1=m unit2=Hz label2=f')

    # Peak frequency versus depth VP
    Flow(specpe,[depthpe, 'srick'],'math lc=${SOURCES[0]} s=${SOURCES[1]} output="s*exp(-x2/lc)"')
    Flow(sumpe,specpe,'stack axis=1 norm=n | window')
    Flow(fumpe,specpe,'math output="x1*input" | stack axis=1 norm=n | window')
    Flow(fdompe,[fumpe ,sumpe],'math f=${SOURCES[0]} s=${SOURCES[1]} output="f/s" | put unit1=m unit2=Hz label2=f')

    # Plots
    Result(fdome,[fdompe, fdomse],
           '''
           cat ${SOURCES[0:2]} axis=2 |
           put label2='f' unit2='Hz' |
           graph dash=0,3 title="Dominant frequency H=0.5" min1=0. max1=1500. min2=0. max2=68.
           parallel2=n plotfat=4,4
           ''',stdin=0)

    # Fractal correlation
    # -------------------

    for fexp in ('M025','025','05','075'):

        parq.update(pgrids[fexp])

         # S waves
        hksb4   = 'hksb4'   + lb + fexp
        depthsf = 'depthsf' + lb + fexp

        Flow(hksb4,[ksb2, ksb],'math k2=${SOURCES[0]} k=${SOURCES[1]} output="k-k/(1.+4.*k2)^(%g)"' % (0.5+parq['hu']))
        # Depth penetration 
        Flow(depthsf,[hksb4, ksb],'math hk4=${SOURCES[0]} k=${SOURCES[1]} output="(%g)/(%g)/(k*hk4)"' % (parq['b'],parq['s2']*parq['cu']))

        # P waves
        hkpb4   = 'hkpb4'   + lb + fexp
        depthpf = 'depthpf' + lb + fexp

        Flow(hkpb4,[kpb2, kpb],'math k2=${SOURCES[0]} k=${SOURCES[1]} output="k-k/(1.+4.*k2)^(%g)"' % (0.5+parq['hu']))
        # Depth penetration
        Flow(depthpf,[hkpb4, kpb],'math hk4=${SOURCES[0]} k=${SOURCES[1]} output="(%g)/(%g)/(k*hk4)"' % (parq['b'],parq['s2']*parq['cu']))

        specsf = 'specsf'  + lb + fexp
        sumsf  = 'sumsf'   + lb + fexp
        fumsf  = 'fumsf'   + lb + fexp
        fdomsf = 'fdomsf'  + lb + fexp

        specpf = 'specpf'  + lb + fexp
        sumpf  = 'sumpf'   + lb + fexp
        fumpf  = 'fumpf'   + lb + fexp
        fdompf = 'fdompf'  + lb + fexp

        fdomf = 'fdomf' + lb + fexp

        # Peak frequency versus depth VS
        Flow(specsf,[depthsf, 'srick'],'math lc=${SOURCES[0]} s=${SOURCES[1]} output="s*exp(-x2/lc)"')
        Flow(sumsf,specsf,'stack axis=1 norm=n | window')
        Flow(fumsf,specsf,'math output="x1*input" | stack axis=1 norm=n | window')
        Flow(fdomsf,[fumsf, sumsf],'math f=${SOURCES[0]} s=${SOURCES[1]} output="f/s" | put unit1=m unit2=Hz label2=f')

        # Peak frequency versus depth VP
        Flow(specpf,[depthpf, 'srick'],'math lc=${SOURCES[0]} s=${SOURCES[1]} output="s*exp(-x2/lc)"')
        Flow(sumpf,specpf,'stack axis=1 norm=n | window')
        Flow(fumpf,specpf,'math output="x1*input" | stack axis=1 norm=n | window')
        Flow(fdompf,[fumpf, sumpf],'math f=${SOURCES[0]} s=${SOURCES[1]} output="f/s" | put unit1=m unit2=Hz label2=f')

        # Plots
        titledfq = 'title="Dominant frequency \k100 \F5 b\F2 \k70 =\k70 ' + bstring + '\k10 m\g \G \F5 H\F2 \k60 =\k60 ' + parq['shu'] + '"'
        Result(fdomf,[fdompf, fdomsf],
               '''
               cat ${SOURCES[0:2]} axis=2 |
               put label2='f' unit2='Hz' |
               graph dash=0,3 min1=0. max1=1500. min2=0. max2=68. font=2
               parallel2=n plotfat=4,4 n2tic=7 o2num=0.0 d2num=10.0 %s
               ''' % (titledfq),stdin=0)

End()