Extractzm

pro extractzm, all=all, head=head

;
; Process zonal-mean variables from a set of ncdf files
;  if keyword 'all' is set, all variables are chosen for a specified year
;

  if(keyword_set(all)) then begin

; process all variables for one year
    yr = ' '
    read, yr, prompt='process all variables for what year? '

  endif else begin

; variables to process (varnames) for all years
    varnames = 
    read, varnames, prompt='enter variables to process separated by commas (var1,var2,var3, ...) '

; T and some other fields need to be added to all files
    res = strmatch(varnames,'*T*')
    if(res eq 0) then varadd ='date,hyam,hybm,hyai,hybi,P0,PS,PSL,T' $
               else varadd ='date,hyam,hybm,hyai,hybi,P0,PS,PSL' 

    varlist = varnames + ',' + varadd

  endelse

; will look for head*.nc files in current directory, where 'head' is file header
  if(not(keyword_set(head))) then head='wa319_4x5*'

;
; Produce a concatenated file containing variables in 'varlist'
;  or for all variables in year 'yr'
;

  spawn, 'mkdir tmpdir'

; select all files for the current ensemble
  if(keyword_set(all)) then pattern = head + '*' + yr + '-*.nc' $
                       else pattern = head + '*.nc'
  select, pattern, filenames, n1

; zonally-average all files -- output to tmpdir
  print, 'zonally-averaging all files'

  for j = 0, n1-1 do begin
    print, 'averaging file ' + strtrim(string(j+1,'(i3)'), 2) + ' of ' $
           + strtrim(string(n1,'(i3)'), 2)
    inputfile = ' ' + filenames(j)
    outputfile = ' tmpdir/' + filenames(j)
    if(keyword_set(all)) then command = 'ncwa -O -a lon ' + inputfile + outputfile $
                         else command = 'ncwa -O -a lon -v ' + varlist + inputfile + outputfile
    spawn, command ;& print, command
  endfor

; concantenate variables across all files 
  print, 'concatenating results'

  inputfiles = ' tmpdir/' + head + '*.nc'

  if(keyword_set(all)) then begin
    outfile = ' concat_zm_all_' + yr + '.nc' 
    command = 'ncrcat -O ' + inputfiles + outfile 
  endif else begin  
    outfile = ' concat_zm_' + varnames + '.nc'
    command = 'ncrcat -O -v ' + varlist + inputfiles + outfile
  endelse

  spawn, command ;& print, command

;
; Deseasonalize if requested (only the variables in 'varnames'),
; or all the variables for the single year 'yr'
;

  cont = ' '
  if(not(keyword_set(all))) then read, cont, prompt='create anomaly file? (y or n) '
  if(cont ne 'y') then goto, finish

  for i = 1, 12 do begin
    mo = strtrim(string(i,'(i2)'), 2)

; ensemble-average over the zonal-mean files for each month -- output to tmpdir
    print, 'ensemble-averaging files for month ' + mo

    if(i lt 10) then inputfiles = ' tmpdir/' + head + '*-0' + mo + '.nc' $
                else inputfiles = ' tmpdir/' + head + '*-' + mo + '.nc'
    avefile = ' tmpdir/ave' + mo + '.nc'

    command = 'ncea -O -v ' + varnames + inputfiles + avefile
    spawn, command ;& print, command

; subtract the ensemble-mean zonal-mean for each month from each zonal-mean monthly file,
;   and output anomaly file to tmpdir 
    print, 'computing anomalies for month ' + mo

    if(i lt 10) then pattern = 'tmpdir/' + head + '*-0' + mo + '.nc' $
                else pattern = 'tmpdir/' + head + '*-' + mo + '.nc'
    select, pattern, filenames, n1

    for j = 0, n1-1 do begin
      inputfile = filenames(j)
      outputfile = ' ' + filenames(j) + '_anom'

      command = 'ncbo -O --op_typ=sub -v ' + varnames + ' ' + inputfile + avefile + outputfile
      spawn, command ;& print, command
    endfor

  endfor

; concatenate the zonal-mean anomaly files -- output to 'outanom'
  print, 'concatenating zonal-mean anomaly files' 

  inputfiles = ' tmpdir/' + head + '*.nc_anom'
  outanom = ' anom_zm_' + varnames + '.nc'
  command = 'ncrcat -O -v ' + varnames + inputfiles + outanom
  spawn, command ;& print, command

; now add zonal-mean T and other fields to the anomaly file (recall that 'outfile' contains
;  raw concatenated results)

  print, 'adding default fields to anomaly file'
  command = 'ncks -A -v ' + varadd + ' ' + outfile + outanom
  spawn, command ;& print, command
  print, 'wrote ' + outanom

; clean up
  finish:
  print, 'cleaning up'
  spawn, '/bin/rm -rf tmpdir'


end

pro select, pattern, filenames, n1

  command = 'ls -1 ' + pattern + ' > names'
  spawn, command ;& print, command
  command = 'wc names > number'
  spawn, command ;& print, command

; n1 is the number of files in the ensemble
  openr, 10, 'number'
  readf, 10, n1, dummy, dummy
  close, 10
  print, 'selected ' + strtrim( string(n1,'(i4)'), 2) + ' files'

; filenames contains the names of the n1 files 
  filenames = strarr(n1)
  openr, 10, 'names'
  readf, 10, filenames
  close, 10

  spawn, '/bin/rm names number'

return
end