pro sigmat_annual

; Purpose
; -------
; Reads in annual-mean CSIRO ocean model output, and calculates the annual-mean
; potential density.
;
; Note that we make the important assumption that the equation of state is
; sufficiently linear that we can calculate the annual-mean density directly
; from the annual-mean temperature and salinity.
;
; History
; -------
; 2006 Jun 19	Steven Phipps	Original version, based on rho_annual.pro
; 2008 May 23	Steven Phipps	Modified to read grid and mask data from
;				working directory

; Define parameters
nx = 64
ny = 56
nz = 21
missing_value = -1.0e34

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

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

; Declare input and output arrays
temp = fltarr(nx, ny, nz, 50)
sal = fltarr(nx, ny, nz, 50)
sigmat_out = fltarr(nx, ny, nz, 50)
sigmat = dblarr(50)
sigmat1 = dblarr(50)
sigmat2 = dblarr(50)
sigmat3 = dblarr(50)

; Get latitudes and volumes from csiro_ogcm_ts_area_volume.nc
ncid = ncdf_open("csiro_ogcm_ts_area_volume.nc")
lattsid = ncdf_varid(ncid, "latts")
dvid = ncdf_varid(ncid, "dv")
ncdf_varget, ncid, lattsid, latts
ncdf_varget, ncid, dvid, dv
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

; GENERATE OUTPUT FILE
;
; (1) Create netCDF 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 = "sigma_t." + run_name + "." + year1s + "-" + year2s + ".nc"
outid = ncdf_create(outfile)

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

; (3) Define variables
lonvid = ncdf_vardef(outid, "longitude", londid, /float)
latvid = ncdf_vardef(outid, "latitude", latdid, /float)
depvid = ncdf_vardef(outid, "depth", depdid, /float)
timvid = ncdf_vardef(outid, "year", timdid, /short)
sigmatid = ncdf_vardef(outid, "sigma_t", [londid, latdid, depdid, timdid], $
                      /float)
sigmaid = ncdf_vardef(outid, "sigmat", timdid, /float)
sigma1id = ncdf_vardef(outid, "sigmat1", timdid, /float)
sigma2id = ncdf_vardef(outid, "sigmat2", timdid, /float)
sigma3id = ncdf_vardef(outid, "sigmat3", timdid, /float)

; (4) Put attributes
title = "CSIRO model run " + run_name + ", years " + year1s + " to " + year2s
ncdf_attput, outid, sigmatid, "long_name", "Potential density"
ncdf_attput, outid, sigmaid, "long_name", "Mean potential density"
ncdf_attput, outid, sigma1id, "long_name", $
  "Mean potential density of the upper ocean (0-800m)"
ncdf_attput, outid, sigma2id, "long_name", $
  "Mean potential density of the mid-ocean (800-2350m)"
ncdf_attput, outid, sigma3id, "long_name", $
  "Mean potential density of the deep ocean (2350-4600m)"
ncdf_attput, outid, lonvid, "units", "degrees_east"
ncdf_attput, outid, latvid, "units", "degrees_north"
ncdf_attput, outid, depvid, "units", "m"
ncdf_attput, outid, timvid, "units", "years"
ncdf_attput, outid, sigmatid, "units", "kg/m^3"
ncdf_attput, outid, sigmaid, "units", "kg/m^3"
ncdf_attput, outid, sigma1id, "units", "kg/m^3"
ncdf_attput, outid, sigma2id, "units", "kg/m^3"
ncdf_attput, outid, sigma3id, "units", "kg/m^3"
ncdf_attput, outid, sigmatid, "missing_value", missing_value
ncdf_attput, outid, sigmaid, "missing_value", missing_value
ncdf_attput, outid, sigma1id, "missing_value", missing_value
ncdf_attput, outid, sigma2id, "missing_value", missing_value
ncdf_attput, outid, sigma3id, "missing_value", missing_value
ncdf_attput, outid, lonvid, "modulo", " "
ncdf_attput, outid, lonvid, "point_spacing", "even"
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

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

  ; Initialise output array
  sigmat_out(*, *, *, *) = missing_value

  ; 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)

  ; If this is first input file, read axis data and write it to output file
  if (year eq year1) then begin

    ; Read axis data
    loninid = ncdf_varid(inid, "lonts")
    depinid = ncdf_varid(inid, "zts")
    ncdf_varget, inid, loninid, lon
    ncdf_varget, inid, depinid, z

    ; Write axis data
    ncdf_varput, outid, lonvid, lon
    ncdf_varput, outid, latvid, latts
    ncdf_varput, outid, depvid, z
    ncdf_varput, outid, timvid, indgen(nyears) + year1

  endif

  ; Read data from input file
  tempinid = ncdf_varid(inid, "temp")
  salinid = ncdf_varid(inid, "sal")
  ncdf_varget, inid, tempinid, temp
  ncdf_varget, inid, salinid, sal

  ; Calculate densities
  for l = 0, 49 do begin
    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
            sigmat_out(i, j, k, l) = sigma_t(temp(i, j, k, l), sal(i, j, k, l))
          endif
        endfor
      endfor
    endfor
  endfor

  ; Calculate the mean density of the upper ocean
  sigmat(*) = 0.0d
  vtotal = 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
          vtotal = vtotal + dv(j, k)
          sigmat(*) = sigmat(*) + dv(j, k) * sigmat_out(i, j, k, *)
        endif
      endfor
    endfor
  endfor
  sigmat(*) = sigmat(*) / vtotal

  ; Calculate the mean density of the upper ocean
  sigmat1(*) = 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)
          sigmat1(*) = sigmat1(*) + dv(j, k) * sigmat_out(i, j, k, *)
        endif
      endfor
    endfor
  endfor
  sigmat1(*) = sigmat1(*) / vtotal

  ; Calculate the mean density of the mid-ocean
  sigmat2(*) = 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)
          sigmat2(*) = sigmat2(*) + dv(j, k) * sigmat_out(i, j, k, *)
        endif
      endfor
    endfor
  endfor
  sigmat2(*) = sigmat2(*) / vtotal

  ; Calculate the mean density of the mid-ocean
  sigmat3(*) = 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)
          sigmat3(*) = sigmat3(*) + dv(j, k) * sigmat_out(i, j, k, *)
        endif
      endfor
    endfor
  endfor
  sigmat3(*) = sigmat3(*) / vtotal

  ; Write data to output file
  ncdf_varput, outid, sigmatid, sigmat_out, offset=[0, 0, 0, year-year1]
  ncdf_varput, outid, sigmaid, sigmat, offset=[year-year1]
  ncdf_varput, outid, sigma1id, sigmat1, offset=[year-year1]
  ncdf_varput, outid, sigma2id, sigmat2, offset=[year-year1]
  ncdf_varput, outid, sigma3id, sigmat3, offset=[year-year1]

endfor

; Close output file
ncdf_close, outid

end