source: trunk/SRC/Interpolation/extrapolate.pro

;+
;
; @file_comments
; extrapolate data (zinput) where maskinput equal 0 by filling step by
; step the coastline points with the mean value of the 8 neighbors
; (weighted by their mask values).
;
; @categories
; Interpolation
;
; @param zinput {in}{required}{type=2d array}
; data to be extrapolate
;
; @param maskinput {in}{required}{type=2d array or -1}
; a 2D array, the land-sea mask of the output data (1 on ocean, 0 on land)
; put -1 if input data are not masked
;
; @param nb_iteration {in}{optional}{type=integer}{default=large enough to fill everything}
; Maximum number of iterations done in the extrapolation process. If there
; is no more masked values we exit extrapolate before reaching nb_iteration
;
; @keyword FILLXDIR {type=scalar, 0 or 1}{default=0}
; put 1 to specify that filling of the data must be done only along x direction
;
; @keyword FILLYDIR {type=scalar, 0 or 1}{default=0}
; put 1 to specify that filling of the data must be done only along y direction
;
; @keyword X_PERIODIC {type=scalar, 0 or 1}{default=0}
; put 1 to specify that the data are periodic along x axis
;
; @keyword MINVAL {type=scalar}{default=not used}
; to specify a minimum value to the extrapolated values
;
; @keyword MAXVAL {type=scalar}{default=not used}
; to specify a maximum value to the extrapolated values
;
; @keyword GE0 {type=scalar 0 or 1}{default=0}
; put 1 to force the extrapolated values to be larger than 0, same as using minval=0.
;
; @keyword
; _EXTRA to be able to call extrapolate with _extra keyword
;
; @returns
; the extrapolated 2d array
;
; @examples
;
;   IDL> a=extrapolate(dist(jpi,jpj),tmask[*,*,0],/x_periodic)
;   IDL> tvplus, a
;   IDL> tvplus, a*(1-tmask[*,*,0])
;
; get the coastline:
;
;   IDL> a=extrapolate(tmask[*,*,0],tmask[*,*,0],1,/x_periodic)
;   IDL> tvplus, a-tmask[*,*,0]
;
; @history
;  Originally written by G. Roullet
;  Sebastien Masson (smasson\@lodyc.jussieu.fr)
;
; @version
; $Id$
;
;-
FUNCTION extrapolate, zinput, maskinput, nb_iteration, FILLXDIR = fillxdir, FILLYDIR = fillydir $
                      , X_PERIODIC = x_periodic, MINVAL = minval, MAXVAL = maxval, GE0 = ge0 $
                      , SMWIN = smwin, NSMOOTH = nsmooth, _EXTRA = ex
;
  compile_opt idl2, strictarrsubs
;
; check the number of iteration used in the extrapolation.
  szin = size(zinput)
  IF szin[0] EQ 1 THEN BEGIN
    zinput = reform(zinput, szin[1], 1)
    maskinput = reform(maskinput, szin[1], 1)
    szin = size(zinput)
    fillxdir = 1
  ENDIF
  
  IF szin[0] NE 2 THEN return, -1. ELSE szin = szin[1:2]
  nx = szin[0]
  ny = szin[1]
  IF n_elements(nb_iteration) EQ 0 THEN nb_iteration = max(szin)
  IF n_elements(smwin) EQ 0 THEN smwin = 0
  smwin = smwin < min(szin)
  IF n_elements(nsmooth) EQ 0 THEN nsmooth = 25
  IF smwin GT 1 AND nsmooth NE 0 THEN nb_iteration = max(szin)
  IF nb_iteration EQ 0 THEN return, zinput
; take care of the boundary conditions...
;
; for the x direction, we put 2 additional columns at the left and
; right side of the array.
; for the y direction, we put 2 additional lines at the bottom and
; top side of the array.
; These changes allow us to use shift function without taking care of
; the x and y periodicity.
;
  msk = bytarr(nx+2, ny+2)
  IF n_elements(maskinput) EQ 1 AND maskinput[0] EQ -1 THEN maskinput = replicate(1b, nx, ny)
  IF n_elements(maskinput) NE nx*ny THEN BEGIN
    ras = report('input grid mask do not have the good size')
    return, -1
  ENDIF
  msk[1:nx, 1:ny] = byte(maskinput)
;
  ztmp = replicate(1.e20, nx+2, ny+2)
  ztmp[1:nx, 1:ny] = zinput
  if keyword_set(x_periodic) then begin
    ztmp[0, 1:ny] = zinput[nx-1, *]
    ztmp[nx+1, 1:ny] = zinput[0, *]
  ENDIF
; remove NaN points if there is some...
  finztmp = finite(ztmp)
  nan = where(finztmp EQ 0, cnt_nan)
  IF cnt_nan NE 0 THEN BEGIN
    ztmp[temporary(nan)] = 1.e20
    msk = temporary(msk) * temporary(finztmp)
  ENDIF ELSE finztmp = -1       ; free memory
  z = temporary(ztmp)
  nx2 = nx+2
  ny2 = ny+2
;---------------------------------------------------------------
;---------------------------------------------------------------
; extrapolation
;---------------------------------------------------------------
  sqrtinv = 1./sqrt(2)
;
  cnt = 1
; When we look for the coastline, we don't want to select the
; borderlines of the array. -> we force the value of the mask for
; those lines.
  msk[0, *] = 1b
  msk[nx+1, *] = 1b
  msk[*, 0] = 1b
  msk[*, ny+1] = 1b
; find the land points
  land = where(msk EQ 0, cnt_land)
  IF keyword_set(fillxdir) THEN BEGIN
    tst = total(msk[1:nx, 1:ny], 1)
    cnt_land = total(tst NE 0 AND tst NE nx)
  ENDIF
  IF keyword_set(fillydir) THEN BEGIN
    tst = total(msk[1:nx, 1:ny], 2)
    cnt_land = total(tst NE 0 AND tst NE ny)
  ENDIF
;---------------------------------------------------------------
  WHILE cnt LE nb_iteration AND cnt_land NE 0 DO BEGIN
;---------------------------------------------------------------
; find the coastline points...
;---------------------------------------------------------------
; Once the land points list has been found, we change back the
; mask values for the boundary conditions.
    msk[0, *] = 0b
    msk[nx+1, *] = 0b
    msk[*, 0] = 0b
    msk[*, ny+1] = 0b
    if keyword_set(x_periodic) then begin
      msk[0, *] = msk[nx, *]
      msk[nx+1, *] = msk[1, *]
    endif
;
; we compute the weighted number of sea neighbors.
; those 4 neighbors have a weight of 1:
;    *
;   *+*
;    *
;
; those 4 neighbors have a weight of 1/sqrt(2):
;
;    * *
;     +
;    * *
;
; As we make sure that none of the land points are located on the
; border of the array, we can compute the weight without shift
; (faster).
;
    CASE 1 OF
      keyword_set(fillxdir):weight = msk[land+1]+msk[land-1]
      keyword_set(fillydir):weight = msk[land+nx2]+msk[land-nx2]
      ELSE:weight = msk[land+1]+msk[land-1]+msk[land+nx2]+msk[land-nx2] $
                    +sqrtinv*(msk[land+nx2+1]+msk[land+nx2-1] $
                              +msk[land-nx2+1]+msk[land-nx2-1])
    ENDCASE
; list all the points that have sea neighbors
    ok = where(weight GT 0)
; the coastline points
    coast = land[ok]
; their weighted number of sea neighbors.
    weight = weight[temporary(ok)]
;---------------------------------------------------------------
; fill the coastline points
;---------------------------------------------------------------
    z = temporary(z)*msk
;
    CASE 1 OF
      keyword_set(fillxdir):zcoast = z[1+coast]+z[-1+coast]
      keyword_set(fillydir):zcoast = z[nx2+coast]+z[-nx2+coast]
      ELSE:zcoast = z[1+coast]+z[-1+coast]+z[nx2+coast]+z[-nx2+coast] $
                    +1./sqrt(2)*(z[nx2+1+coast]+z[nx2-1+coast] $
                                 +z[-nx2+1+coast]+z[-nx2-1+coast])
    ENDCASE
;
    IF keyword_set(ge0) THEN zcoast = 0. > temporary(zcoast)
    IF n_elements(minval) NE 0 THEN zcoast = minval > temporary(zcoast)
    IF n_elements(maxval) NE 0 THEN zcoast = temporary(zcoast) < maxval
    z[coast] = temporary(zcoast)/temporary(weight)
; we update the boundary conditions of z
    if keyword_set(x_periodic) then begin
      z[0, *] = z[nx, *]
      z[nx+1, *] = z[1, *]
    endif
;---------------------------------------------------------------
; we update the mask
;---------------------------------------------------------------
    msk[temporary(coast)] = 1
;
    cnt = cnt + 1
; When we look for the coast line, we don't want to select the
; borderlines of the array. -> we force the value of the mask for
; those lines.
    msk[0, *] = 1b
    msk[nx+1, *] = 1b
    msk[*, 0] = 1b
    msk[*, ny+1] = 1b
; find the land points
    land = where(msk EQ 0, cnt_land)
    IF keyword_set(fillxdir) THEN BEGIN
      tst = total(msk[1:nx, 1:ny], 1)
      cnt_land = total(tst NE 0 AND tst NE nx)
    ENDIF
    IF keyword_set(fillydir) THEN BEGIN
      tst = total(msk[1:nx, 1:ny], 2)
      cnt_land = total(tst NE 0 AND tst NE ny)
    ENDIF
;
  ENDWHILE
;
  z = z[1:nx, 1:ny]
;
;---------------------------------------------------------------
; smooth the filled values
;---------------------------------------------------------------
  IF smwin LE 1 OR nsmooth EQ 0 THEN return, z

; add extra bands to avoid edge errors...

  one = replicate(1, smwin)
  IF keyword_set(x_periodic) THEN BEGIN
;   nx-smwin   nx-1    0       smwin-1                      nx-smwin      nx-1    0       smwin-1
;     |----------|     |----------|------------------------------|----------|     |----------|
;     0      smwin-1  smwin   2*smwin-1                          nx  nx+smwin-1  nx+smwin nx+2*smwin-1
    z   = [         z[nx-smwin:nx-1, *],         z,         z[0:smwin-1, *] ]
    msk = [ maskinput[nx-smwin:nx-1, *], maskinput, maskinput[0:smwin-1, *] ]
;; IF array_equal(z[0:smwin-1, *], z[nx:nx+smwin-1, *]) NE 1 THEN stop
;; IF array_equal(z[nx+smwin:nx+2*smwin-1, *], z[smwin:2*smwin-1, *]) NE 1 THEN stop
  ENDIF ELSE BEGIN
    z   = [         one#(z[0, *])[*],         z,         one#(z[nx-1, *])[*] ]
    msk = [ one#(maskinput[0, *])[*], maskinput, one#(maskinput[nx-1, *])[*] ]
  ENDELSE
  z   = [ [  z[*, 0]#one], [  z], [  z[*, ny-1]#one] ]
  msk = [ [msk[*, 0]#one], [msk], [msk[*, ny-1]#one] ]

  zorg = z
  mskm1 = 1-msk
  FOR i = 1, nsmooth DO BEGIN
    z = smooth( temporary(z)*mskm1 + zorg*msk, smwin )
    IF keyword_set(x_periodic) THEN BEGIN
      z[0:smwin-1, *] = z[nx:nx+smwin-1, *]
      z[nx+smwin:nx+2*smwin-1, *] =  z[smwin:2*smwin-1, *]
    ENDIF ELSE BEGIN
      z[0:smwin-1, *] = one#(z[smwin, *])[*]
      z[nx+smwin:nx+2*smwin-1, *] = one#(z[nx+smwin-1, *])[*]
    ENDELSE
    z[*, 0:smwin-1] = z[*, smwin]#one
    z[*, ny+smwin:ny+2*smwin-1] = z[*, ny+smwin-1]#one
  ENDFOR
  z = temporary(z)*mskm1 + zorg*msk
  z = (temporary(z))[smwin:nx+smwin-1, smwin:ny+smwin-1]
;
;---------------------------------------------------------------
; we return the original size of the array
;---------------------------------------------------------------
;
  return, z
END