pro ts_stats

; Purpose
; -------
; Reads in annual-mean CSIRO OGCM output, and produces a single netCDF file
; containing the following temperature and salinity statistics:
;
; STFHT		mean surface heat flux [W/m^2]
; STFSAL        mean surface salinity tendency [s^-1]
;
; ISTFHT	integrated surface heat flux [J]
; ISTFSAL	integrated surface salinity tendency [kg/kg]
;
; TEMPYZ	zonal-mean potential temperature [degC]
; SALYZ		zonal-mean salinity [psu]
;
; TEMPZ		zonal- and meridional-mean potential temperature [degC]
; SALZ		zonal- and meridional-mean salinity [psu]
;
; TEMP		mean potential temperature [degC]
; SAL		mean salinity [psu]
;
; TEMP1		mean potential temperature of the upper ocean [0-800m, degC]
; TEMP2		mean potential temperature of the mid-ocean [800-2350m, degC]
; TEMP3		mean potential temperature of the deep ocean [2350-4600m, degC]
;
; SAL1		mean salinity of the upper ocean [0-800m, psu]
; SAL2		mean salinity of the mid-ocean [800-2350m, psu]
; SAL3		mean salinity of the deep ocean [2350-4600m, psu]
;
; DTE		equivalent change in mean potential temperature [K]
; DSE		equivalent change in mean salinity [psu]
;
; TERR		error in mean potential temperature [K]
; SERR		error in mean salinity [psu]
;
; History
; -------
; 2005 Mar 17	Steve Phipps	Original version
; 2005 Mar 25	Steve Phipps	Added calculation of DT and DS
; 2005 Mar 26	Steve Phipps	Added calculation of DTE and DSE; DT and DS
;				re-named TERR and SERR respectively
; 2006 May 16	Steve Phipps	Added calculation of TEMP1, TEMP2, TEMP3,
;				SAL1, SAL2 and SAL3

; Define parameters
nx = 64
ny = 56
nz = 21
missing_value = -1.0d34
secs_year = 3.1536d7
rhocw = 4.1d6

; Get details from user
run_name = ""
year1 = 0
year2 = 0
read, run_name, prompt="Please enter the run name [e.g. d30]    "
read, year1, prompt="Please enter the initial year [e.g. 1]  "
read, year2, prompt="Please enter the final year [e.g. 500]  "

; Get number of years of data to read
nyears = year2 - year1 + 1

; Declare output arrays
stfht = dblarr(50)
stfsal = dblarr(50)
istfht = dblarr(nyears)
istfsal = dblarr(nyears)
tempyz = dblarr(ny, nz, 50)
salyz = dblarr(ny, nz, 50)
tempz = dblarr(nz, 50)
salz = dblarr(nz, 50)
temp = dblarr(50)
sal = dblarr(50)
temp1 = dblarr(50)
temp2 = dblarr(50)
temp3 = dblarr(50)
sal1 = dblarr(50)
sal2 = dblarr(50)
sal3 = dblarr(50)
dte = dblarr(50)
dse = dblarr(50)
terr = dblarr(50)
serr = dblarr(50)

; Get areas, volumes and axes from csiro_ogcm_ts_area_volume.nc
ncid = ncdf_open("csiro_ogcm_ts_area_volume.nc")
lattsid = ncdf_varid(ncid, "latts")
ztsid = ncdf_varid(ncid, "zts")
daid = ncdf_varid(ncid, "da")
dvid = ncdf_varid(ncid, "dv")
areaid = ncdf_varid(ncid, "area")
volumeid = ncdf_varid(ncid, "volume")
ncdf_varget, ncid, lattsid, latts
ncdf_varget, ncid, ztsid, zts
ncdf_varget, ncid, daid, da
ncdf_varget, ncid, dvid, dv
ncdf_varget, ncid, areaid, area
ncdf_varget, ncid, volumeid, volume
ncdf_close, ncid

; Get land/sea mask from csiro_ogcm_ts_landsea3.nc
ncid = ncdf_open("csiro_ogcm_ts_landsea3.nc")
maskid = ncdf_varid(ncid, "mask")
ncdf_varget, ncid, maskid, mask
ncdf_close, ncid

; Derive name of output file
year1s = strtrim(string(year1), 2)
year2s = strtrim(string(year2), 2)
if (year1 lt 10000) then year1s = '0' + year1s
if (year1 lt 1000) then year1s = '0' + year1s
if (year1 lt 100) then year1s = '0' + year1s
if (year1 lt 10) then year1s = '0' + year1s
if (year2 lt 10000) then year2s = '0' + year2s
if (year2 lt 1000) then year2s = '0' + year2s
if (year2 lt 100) then year2s = '0' + year2s
if (year2 lt 10) then year2s = '0' + year2s
outfile = "ts." + run_name + "." + year1s + "-" + year2s + ".nc"

; Create output file
;
; (1) Create netCDF file
outid = ncdf_create(outfile)

; (2) Define dimensions
latdid = ncdf_dimdef(outid, "latitude", ny)
depdid = ncdf_dimdef(outid, "depth", nz)
timdid = ncdf_dimdef(outid, "year", nyears)

; (3) Define variables
latvid = ncdf_vardef(outid, "latitude", latdid, /float)
depvid = ncdf_vardef(outid, "depth", depdid, /float)
timvid = ncdf_vardef(outid, "year", timdid, /short)
stfhtid = ncdf_vardef(outid, "stfht", timdid, /float)
stfsalid = ncdf_vardef(outid, "stfsal", timdid, /float)
istfhtid = ncdf_vardef(outid, "istfht", timdid, /float)
istfsalid = ncdf_vardef(outid, "istfsal", timdid, /float)
tempyzid = ncdf_vardef(outid, "tempyz", [latdid, depdid, timdid], /float)
salyzid = ncdf_vardef(outid, "salyz", [latdid, depdid, timdid], /float)
tempzid = ncdf_vardef(outid, "tempz", [depdid, timdid], /float)
salzid = ncdf_vardef(outid, "salz", [depdid, timdid], /float)
tempid = ncdf_vardef(outid, "temp", timdid, /float)
salid = ncdf_vardef(outid, "sal", timdid, /float)
temp1id = ncdf_vardef(outid, "temp1", timdid, /float)
temp2id = ncdf_vardef(outid, "temp2", timdid, /float)
temp3id = ncdf_vardef(outid, "temp3", timdid, /float)
sal1id = ncdf_vardef(outid, "sal1", timdid, /float)
sal2id = ncdf_vardef(outid, "sal2", timdid, /float)
sal3id = ncdf_vardef(outid, "sal3", timdid, /float)
dteid = ncdf_vardef(outid, "dte", timdid, /float)
dseid = ncdf_vardef(outid, "dse", timdid, /float)
terrid = ncdf_vardef(outid, "terr", timdid, /float)
serrid = ncdf_vardef(outid, "serr", timdid, /float)

; (4) Put attributes
title = "CSIRO model run " + run_name + ", years " + year1s + " to " + year2s
ncdf_attput, outid, stfhtid, "long_name", "Mean surface heat flux"
ncdf_attput, outid, stfsalid, "long_name", "Mean surface salinity tendency"
ncdf_attput, outid, istfhtid, "long_name", "Integrated surface heat flux"
ncdf_attput, outid, istfsalid, "long_name",  $
  "Integrated surface salinity tendency"
ncdf_attput, outid, tempyzid, "long_name", "Zonal-mean potential temperature"
ncdf_attput, outid, salyzid, "long_name", "Zonal-mean salinity"
ncdf_attput, outid, tempzid, "long_name", $
  "Zonal- and meridional-mean potential temperature"
ncdf_attput, outid, salzid, "long_name", "Zonal- and meridional-mean salinity"
ncdf_attput, outid, tempid, "long_name", "Mean potential temperature"
ncdf_attput, outid, salid, "long_name", "Mean salinity"
ncdf_attput, outid, temp1id, "long_name", $
  "Mean potential temperature of the upper ocean (0-800m)"
ncdf_attput, outid, temp2id, "long_name", $
  "Mean potential temperature of the mid-ocean (800-2350m)"
ncdf_attput, outid, temp3id, "long_name", $
  "Mean potential temperature of the deep ocean (2350-4600m)"
ncdf_attput, outid, sal1id, "long_name", $
  "Mean salinity of the upper ocean (0-800m)"
ncdf_attput, outid, sal2id, "long_name", $
  "Mean salinity of the mid-ocean (800-2350m)"
ncdf_attput, outid, sal3id, "long_name", $
  "Mean salinity of the deep ocean (2350-4600m)"
ncdf_attput, outid, dteid, "long_name", $
  "Equivalent change in mean potential temperature"
ncdf_attput, outid, dseid, "long_name", "Equivalent change in mean salinity"
ncdf_attput, outid, terrid, "long_name", "Error in mean potential temperature"
ncdf_attput, outid, serrid, "long_name", "Error in mean salinity"
ncdf_attput, outid, latvid, "units", "degrees_north"
ncdf_attput, outid, depvid, "units", "m"
ncdf_attput, outid, timvid, "units", "years"
ncdf_attput, outid, stfhtid, "units", "W/m^2"
ncdf_attput, outid, stfsalid, "units", "s^-1"
ncdf_attput, outid, istfhtid, "units", "J"
ncdf_attput, outid, istfsalid, "units", "kg/kg"
ncdf_attput, outid, tempyzid, "units", "degC"
ncdf_attput, outid, salyzid, "units", "psu"
ncdf_attput, outid, tempzid, "units", "degC"
ncdf_attput, outid, salzid, "units", "psu"
ncdf_attput, outid, tempid, "units", "degC"
ncdf_attput, outid, salid, "units", "psu"
ncdf_attput, outid, temp1id, "units", "degC"
ncdf_attput, outid, temp2id, "units", "degC"
ncdf_attput, outid, temp3id, "units", "degC"
ncdf_attput, outid, sal1id, "units", "psu"
ncdf_attput, outid, sal2id, "units", "psu"
ncdf_attput, outid, sal3id, "units", "psu"
ncdf_attput, outid, dteid, "units", "K"
ncdf_attput, outid, dseid, "units", "psu"
ncdf_attput, outid, terrid, "units", "K"
ncdf_attput, outid, serrid, "units", "psu"
ncdf_attput, outid, tempyzid, "missing_value", float(missing_value)
ncdf_attput, outid, salyzid, "missing_value", float(missing_value)
ncdf_attput, outid, latvid, "point_spacing", "uneven"
ncdf_attput, outid, depvid, "point_spacing", "uneven"
ncdf_attput, outid, timvid, "point_spacing", "even"
ncdf_attput, outid, depvid, "positive", "down"
ncdf_attput, outid, /global, "title", title

; (5) Exit define mode
ncdf_control, outid, /endef

; (6) Write axis data
ncdf_varput, outid, latvid, latts
ncdf_varput, outid, depvid, zts
ncdf_varput, outid, timvid, indgen(nyears) + year1

; Loop over years in 50-year increments
for year = year1, year2, 50 do begin
  iyear = year - year1

  ; Derive name of input file
  yr1 = year
  yr2 = year + 49
  yr1s = strtrim(string(yr1), 2)
  yr2s = strtrim(string(yr2), 2)
  if (yr1 lt 10000) then yr1s = '0' + yr1s
  if (yr1 lt 1000) then yr1s = '0' + yr1s
  if (yr1 lt 100) then yr1s = '0' + yr1s
  if (yr1 lt 10) then yr1s = '0' + yr1s
  if (yr2 lt 10000) then yr2s = '0' + yr2s
  if (yr2 lt 1000) then yr2s = '0' + yr2s
  if (yr2 lt 100) then yr2s = '0' + yr2s
  if (yr2 lt 10) then yr2s = '0' + yr2s
  infile = "com.ann." + run_name + "." + yr1s + "-" + yr2s + ".nc"
  print, "Reading data from file ", infile, " ..."

  ; Open input file
  inid = ncdf_open(infile)

  ; Read data
  stfhtinid = ncdf_varid(inid, "stfht")
  stfsalinid = ncdf_varid(inid, "stfsal")
  tempinid = ncdf_varid(inid, "temp")
  salinid = ncdf_varid(inid, "sal")
  ncdf_varget, inid, stfhtinid, stfht_in
  ncdf_varget, inid, stfsalinid, stfsal_in
  ncdf_varget, inid, tempinid, temp_in
  ncdf_varget, inid, salinid, sal_in

  ; Calculate the mean surface fluxes
  stfht(*) = 0.0d
  stfsal(*) = 0.0d
  for j = 0, ny-1 do begin
    for i = 0, nx-1 do begin
      if (mask(i, j, 0) eq 0) then begin
        stfht = stfht + da(j) * stfht_in(i, j, *)
        stfsal = stfsal + da(j) * stfsal_in(i, j, *)
      endif
    endfor
  endfor
  stfht = stfht / area
  stfsal = stfsal / area

  ; Integrate the surface fluxes with respect to time
  for l = 0, 49 do begin
    if ((iyear eq 0) and (l eq 0)) then begin
      istfht(0) = secs_year * area * stfht(0)
      istfsal(0) = secs_year * stfsal(0)
    endif else begin
      istfht(iyear+l) = istfht(iyear+l-1) + secs_year * area * stfht(l)
      istfsal(iyear+l) = istfsal(iyear+l-1) + secs_year * stfsal(l)
    endelse
  endfor

  ; Calculate the zonal-mean temperature and salinity
  for k = 0, nz-1 do begin
    for j = 0, ny-1 do begin
      nvalues = 0
      tempyz(j, k, *) = 0.0d
      salyz(j, k, *) = 0.0d
      for i = 0, nx-1 do begin
        if (mask(i, j, k) eq 0) then begin
          nvalues = nvalues + 1
          tempyz(j, k, *) = tempyz(j, k, *) + temp_in(i, j, k, *)
          salyz(j, k, *) = salyz(j, k, *) + sal_in(i, j, k, *)
        endif
      endfor
      if (nvalues eq 0) then begin
        tempyz(j, k, *) = missing_value
        salyz(j, k, *) = missing_value
      endif else begin
        tempyz(j, k, *) = tempyz(j, k, *) / double(nvalues)
        salyz(j, k, *) = salyz(j, k, *) / double(nvalues)
      endelse
    endfor
  endfor

  ; Calculate the zonal- and meridional-mean temperature and salinity
  for k = 0, nz-1 do begin
    vtotal = 0.0d
    tempz(k, *) = 0.0d
    salz(k, *) = 0.0d
    for j = 0, ny-1 do begin
      for i = 0, nx-1 do begin
        if (mask(i, j, k) eq 0) then begin
          vtotal = vtotal + dv(j, k)
          tempz(k, *) = tempz(k, *) + dv(j, k) * temp_in(i, j, k, *)
          salz(k, *) = salz(k, *) + dv(j, k) * sal_in(i, j, k, *)
        endif
      endfor
    endfor
    tempz(k, *) = tempz(k, *) / vtotal
    salz(k, *) = salz(k, *) / vtotal
  endfor

  ; Calculate the mean temperature and salinity
  temp(*) = 0.0d
  sal(*) = 0.0d
  for k = 0, nz-1 do begin
    for j = 0, ny-1 do begin
      for i = 0, nx-1 do begin
        if (mask(i, j, k) eq 0) then begin
          temp(*) = temp(*) + dv(j, k) * temp_in(i, j, k, *)
          sal(*) = sal(*) + dv(j, k) * sal_in(i, j, k, *)
        endif
      endfor
    endfor
  endfor
  temp(*) = temp(*) / volume
  sal(*) = sal(*) / volume

  ; Calculate the mean temperature and salinity of the upper ocean
  temp1(*) = 0.0d
  sal1(*) = 0.0d
  vtotal = 0.0d
  for k = 0, 10 do begin
    for j = 0, ny-1 do begin
      for i = 0, nx-1 do begin
        if (mask(i, j, k) eq 0) then begin
          vtotal = vtotal + dv(j, k)
          temp1(*) = temp1(*) + dv(j, k) * temp_in(i, j, k, *)
          sal1(*) = sal1(*) + dv(j, k) * sal_in(i, j, k, *)
        endif
      endfor
    endfor
  endfor
  temp1(*) = temp1(*) / vtotal
  sal1(*) = sal1(*) / vtotal

  ; Calculate the mean temperature and salinity of the mid-ocean
  temp2(*) = 0.0d
  sal2(*) = 0.0d
  vtotal = 0.0d
  for k = 11, 15 do begin
    for j = 0, ny-1 do begin
      for i = 0, nx-1 do begin
        if (mask(i, j, k) eq 0) then begin
          vtotal = vtotal + dv(j, k)
          temp2(*) = temp2(*) + dv(j, k) * temp_in(i, j, k, *)
          sal2(*) = sal2(*) + dv(j, k) * sal_in(i, j, k, *)
        endif
      endfor
    endfor
  endfor
  temp2(*) = temp2(*) / vtotal
  sal2(*) = sal2(*) / vtotal

  ; Calculate the mean temperature and salinity of the deep ocean
  temp3(*) = 0.0d
  sal3(*) = 0.0d
  vtotal = 0.0d
  for k = 16, 20 do begin
    for j = 0, ny-1 do begin
      for i = 0, nx-1 do begin
        if (mask(i, j, k) eq 0) then begin
          vtotal = vtotal + dv(j, k)
          temp3(*) = temp3(*) + dv(j, k) * temp_in(i, j, k, *)
          sal3(*) = sal3(*) + dv(j, k) * sal_in(i, j, k, *)
        endif
      endfor
    endfor
  endfor
  temp3(*) = temp3(*) / vtotal
  sal3(*) = sal3(*) / vtotal

  ; Calculate the errors in the mean temperature and salinity
  ;
  ; (1) The equivalent change in mean temperature between one year and the next
  ;     is given by [K]:
  ;
  ;     mean surface heat flux * secs_year * area / (volume * rhocw)
  ;
  ;     where rhocw = 4.1e6 J/m^3/K is the volumetric heat capacity of seawater
  ;
  ; (2) The equivalent change in mean salinity between one year and the next is
  ;     given by [psu]:
  ;
  ;     1000.0 * mean surface salinity tendency * secs_year * 25.0 * area /
  ;                                                                  volume
  ;
  ;     where 1000.0 converts from kg/kg to psu, and 25.0m is the thickness of
  ;     the upper layer of the ocean model
  for l = 0, 49 do begin
    if ((iyear eq 0) and (l eq 0)) then begin
      dte(0) = 0.0d
      dse(0) = 0.0d
      terr(0) = 0.0d
      serr(0) = 0.0d
      dte_old = 0.0d
      dse_old = 0.0d
      terr_old = 0.0d
      serr_old = 0.0d
    endif else begin
      dte(l) = dte_old + 0.5d * secs_year * area * (stfht_old + stfht(l)) / $
                                                   (volume * rhocw)
      dse(l) = dse_old + 1.25d4 * secs_year * area * (stfsal_old + stfsal(l)) $
                                                     / volume
      terr(l) = terr_old + temp(l) - temp_old - dte(l) + dte_old
      serr(l) = serr_old + sal(l) - sal_old - dse(l) + dse_old
      dte_old = dte(l)
      dse_old = dse(l)
      terr_old = terr(l)
      serr_old = serr(l)
    endelse
    stfht_old = stfht(l)
    stfsal_old = stfsal(l)
    temp_old = temp(l)
    sal_old = sal(l)
  endfor

  ; Write data to output file
  ncdf_varput, outid, stfhtid, stfht, count=[50], offset=[iyear]
  ncdf_varput, outid, stfsalid, stfsal, count=[50], offset=[iyear]
  ncdf_varput, outid, istfhtid, istfht(iyear:iyear+49), count=[50], $
               offset=[iyear]
  ncdf_varput, outid, istfsalid, istfsal(iyear:iyear+49), count=[50], $
               offset=[iyear]
  ncdf_varput, outid, tempyzid, tempyz, count=[ny, nz, 50], $
               offset=[0, 0, iyear]
  ncdf_varput, outid, salyzid, salyz, count=[ny, nz, 50], offset=[0, 0, iyear]
  ncdf_varput, outid, tempzid, tempz, count=[nz, 50], offset=[0, iyear]
  ncdf_varput, outid, salzid, salz, count=[nz, 50], offset=[0, iyear]
  ncdf_varput, outid, tempid, temp, count=[50], offset=[iyear]
  ncdf_varput, outid, salid, sal, count=[50], offset=[iyear]
  ncdf_varput, outid, temp1id, temp1, count=[50], offset=[iyear]
  ncdf_varput, outid, temp2id, temp2, count=[50], offset=[iyear]
  ncdf_varput, outid, temp3id, temp3, count=[50], offset=[iyear]
  ncdf_varput, outid, sal1id, sal1, count=[50], offset=[iyear]
  ncdf_varput, outid, sal2id, sal2, count=[50], offset=[iyear]
  ncdf_varput, outid, sal3id, sal3, count=[50], offset=[iyear]
  ncdf_varput, outid, dteid, dte, count=[50], offset=[iyear]
  ncdf_varput, outid, dseid, dse, count=[50], offset=[iyear]
  ncdf_varput, outid, terrid, terr, count=[50], offset=[iyear]
  ncdf_varput, outid, serrid, serr, count=[50], offset=[iyear]

endfor

; Close output file
ncdf_close, outid

end