Modularised calculations of dx, dy, dz, z into 2D and 3D functions ra…

…ther than copying and pasting in many files

avg_ismr.py

@@ -2,6 +2,7 @@
 from numpy import *
 from matplotlib.pyplot import *
 from os.path import *
+from cartesian_grid_2d import *
 
 # Calculate and plot timeseries of area-averaged ice shelf melt rates from 
 # major ice shelves during a ROMS simulation.
@@ -135,58 +136,27 @@ def avg_ismr (file_path, log_path):
 # i,j = i- and j- coordinates on the rho-grid, starting at 1
 def calc_grid (file_path):
 
-    # Radius of the Earth in m
-    r = 6.371e6
-    # Degrees to radians conversion factor
-    deg2rad = pi/180.0
-    # Northern boundary of ROMS grid
-    nbdry_val = -30
-
     # Read grid variables
     id = Dataset(file_path, 'r')
     lon = id.variables['lon_rho'][:,:]
     lat = id.variables['lat_rho'][:,:]
     zice = id.variables['zice'][:,:]
     id.close()
+
+    # Calculate dx and dy in another script
+    dx, dy = cartesian_grid_2d(lon, lat)
+
+    # Calculate dA and mask with zice
+    dA = ma.masked_where(zice==0, dx*dy)
+
     # Save dimensions
     num_lat = size(lon, 0)
     num_lon = size(lon, 1)
 
-    # Add or subtract 360 from longitude values which wrap around
-    # so that longitude increases monotonically from west to east
+    # Calculate i- and j-coordinates
     i = tile(arange(1, num_lon+1), (num_lat, 1))
-    index1 = nonzero((i > 1200)*(lon < 100))
-    lon[index1] = lon[index1] + 360
-    index2 = nonzero((i < 200)*(lon > 300))
-    lon[index2] = lon[index2] - 360
-
-    # Interpolate to get longitude at the edges of each cell
-    w_bdry = 0.5*(lon[:,0] + lon[:,-1] - 360)
-    middle_lon = 0.5*(lon[:,0:-1] + lon[:,1:])
-    e_bdry = 0.5*(lon[:,0] + 360 + lon[:,-1])
-    lon_edges = ma.concatenate((w_bdry[:,None], middle_lon, e_bdry[:,None]), axis=1)
-    # Subtract to get the change in longitude over each cell
-    dlon = abs(lon_edges[:,1:] - lon_edges[:,0:-1])
-
-    # Similarly for latitude
-    s_bdry = lat[0,:]
-    middle_lat = 0.5*(lat[0:-1,:] + lat[1:,:])
-    n_bdry = lat[-1,:]*0 + nbdry_val
-    lat_edges = ma.concatenate((s_bdry[None,:], middle_lat, n_bdry[None,:]))
-    dlat = lat_edges[1:,:] - lat_edges[0:-1,:]
-
-    # Also calculate j-coordinates
     j = transpose(tile(arange(1, num_lat+1), (num_lon, 1)))
 
-    # Convert from spherical to Cartesian coordinates
-    # dy = r*dlat where dlat is converted to radians
-    dy = r*dlat*deg2rad    
-    # dx = r*cos(lat)*dlon where lat and dlon are converted to radians
-    dx = r*cos(lat*deg2rad)*dlon*deg2rad
-
-    # Calculate dA and mask with zice
-    dA = ma.masked_where(zice==0, dx*dy)
-
     return dA, i, j
 
 

avg_zeta.py

@@ -1,6 +1,7 @@
 from netCDF4 import Dataset
 from numpy import *
-from matplotlib.pyplot import plot,xlabel,ylabel,clf,show
+from matplotlib.pyplot import *
+from cartesian_grid_2d import *
 
 # Calculate the average sea surface height (zeta) at each timestep of the given
 # ocean history file.
@@ -16,50 +17,18 @@ def avg_zeta (file_path):
     mask = file.variables['mask_rho'][:,:]
     avg_zeta = []
 
-    # Mask land points out of lat and lon
-    lon = ma.masked_where(mask==0, lon)
-    lat = ma.masked_where(mask==0, lat)
-    num_lat = size(lon, 0)
-    num_lon = size(lon, 1)
-    # Add or subtract 360 from longitude values which wrap around
-    # so that longitude increases monotonically from west to east
-    i = tile(arange(1, num_lon+1), (num_lat, 1))
-    index1 = nonzero((i > 1200)*(lon < 100))
-    lon[index1] = lon[index1] + 360
-    index2 = nonzero((i < 200)*(lon > 300))
-    lon[index2] = lon[index2] - 360
+    # Calculate dx and dy in another script
+    dx, dy = cartesian_grid_2d(lon, lat)
 
-    # Interpolate to get longitude at the edges of each cell
-    w_bdry = (lon[:,0] + lon[:,num_lon-1] - 360)/2
-    middle_lon = (lon[:,0:num_lon-1] + lon[:,1:num_lon])/2
-    e_bdry = (lon[:,0] + 360 + lon[:,num_lon-1])/2
-    lon_edges = ma.concatenate((w_bdry[:,None], middle_lon, e_bdry[:,None]), axis=1)
-    # Subtract to get the change in longitude over each cell
-    dlon = abs(lon_edges[:,1:num_lon+1] - lon_edges[:,0:num_lon])
-
-    # Similarly for latitude
-    s_bdry = lat[0,:]
-    middle_lat = (lat[0:num_lat-1,:] + lat[1:num_lat,:])/2
-    n_bdry = lat[num_lat-1,:]*0 - 50
-    lat_edges = ma.concatenate((s_bdry[None,:], middle_lat, n_bdry[None,:]))
-    dlat = lat_edges[1:num_lat+1,:] - lat_edges[0:num_lat,:]
-
-    # Convert from spherical to Cartesian coordinates
-    r = 6.371e6
-    # dy = r*dlat where dlat is converted to radians
-    dy = r*dlat*pi/180.0
-    # dx = r*cos(lat)*dlon where lat and dlon are converted to radians
-    dx = r*cos(pi*lat/180.0)*dlon*pi/180.0
-    # Multiply to get the area of each cell
-    area = nansum(dx*dy)
+    # Calculate dA and mask with land mask
+    dA = ma.masked_where(mask==0, dx*dy)
 
     for l in range(size(time)):
         print 'Processing timestep ' + str(l+1) + ' of ' + str(size(time))
         # Read zeta at this timestep
         zeta = file.variables['zeta'][l,:,:]
-        # Calculate Area-weighted average
-        zeta_int = nansum(zeta*dx*dy)
-        avg_zeta.append(zeta_int/area)
+        # Calculate area-weighted average
+        avg_zeta.append(sum(zeta*dA)/sum(dA))
 
     file.close()
 

cartesian_grid_2d.py

@@ -0,0 +1,50 @@
+from numpy import *
+
+# Given ROMS grid variables, calculate Cartesian integrands dx and dy.
+# Input:
+# lon, lat = 2D arrays containing values for latitude and longitude, with
+#            dimension latitude x longitude
+# Output:
+# dx, dy = 2D arrays (dimension latitude x longitude) containing Cartesian
+#          integrands
+def cartesian_grid_2d (lon, lat):
+
+    # Radius of the Earth in metres
+    r = 6.371e6
+    # Degrees to radians conversion factor
+    deg2rad = pi/180.0
+
+    # Save horizontal dimensions
+    num_lat = size(lon, 0)
+    num_lon = size(lon, 1)
+
+    # Add or subtract 360 from longitude values which wrap around so that
+    # longitude increases monotonically from west to east
+    i = tile(arange(1,num_lon+1), (num_lat,1))
+    index1 = nonzero((i > 1200)*(lon < 100))
+    lon[index1] = lon[index1] + 360
+    index2 = nonzero((i < 200)*(lon > 300))
+    lon[index2] = lon[index2] - 360
+
+    # Interpolate to get longitude at the edges of each cell
+    w_bdry = 0.5*(lon[:,0] + lon[:,-1] - 360)
+    middle_lon = 0.5*(lon[:,0:-1] + lon[:,1:])
+    e_bdry = 0.5*(lon[:,0] + 360 + lon[:,-1])
+    lon_edges = concatenate((w_bdry[:,None], middle_lon, e_bdry[:,None]), axis=1)
+    # Subtract to get the change in longitude over each cell
+    dlon = abs(lon_edges[:,1:] - lon_edges[:,0:-1])
+
+    # Similarly for latitude; linearly extrapolate for latitude at edges of
+    # N/S boundary cells
+    middle_lat = 0.5*(lat[0:-1,:] + lat[1:,:])
+    s_bdry = 2*lat[0,:] - middle_lat[0,:]
+    n_bdry = 2*lat[-1,:] - middle_lat[-1,:]
+    lat_edges = concatenate((s_bdry[None,:], middle_lat, n_bdry[None,:]), axis=0)
+    dlat = lat_edges[1:,:] - lat_edges[0:-1,:]
+
+    # dx = r*cos(lat)*dlon where lat and dlon are converted to radians
+    dx = r*cos(lat*deg2rad)*dlon*deg2rad
+    # dy = r*dlat where dlat is converted to radians
+    dy = r*dlat*deg2rad
+
+    return dx, dy

cartesian_grid_3d.py

@@ -0,0 +1,40 @@
+from numpy import *
+from cartesian_grid_2d import *
+from calc_z import *
+
+# Given ROMS grid variables, calculate Cartesian integrands (dx, dy, dz) as
+# well as z-coordinates.
+# Input:
+# lon, lat, h, zice = 2D arrays containing values for latitude, longitude,
+#                     bathymetry, and ice shelf draft. All have dimension
+#                     latitude x longitude.
+# theta_s, theta_b, hc, N = scalar parameters (check your grid file and roms.in)
+# zeta = optional 2D array containing values for sea surface height at the
+#        desired timestep
+# Output:
+# dx, dy, dz, z = 3D arrays (dimension depth x latitude x longitude) containing
+#                 Cartesian integrands (dx, dy, dz) and z-coordinates (z).
+def cartesian_grid_3d (lon, lat, h, zice, theta_s, theta_b, hc, N, zeta=None):
+
+    # Calculate 2D dx and dy in another script
+    dx, dy = cartesian_grid_2d(lon, lat)
+    # Copy into 3D arrays, same at each depth level
+    dx = tile(dx, (N,1,1))
+    dy = tile(dy, (N,1,1))
+    # Save horizontal dimensions
+    num_lat = size(lon, 0)
+    num_lon = size(lon, 1)
+
+    # Get a 3D array of z-coordinates; sc_r and Cs_r are unused
+    z, sc_r, Cs_r = calc_z(h, zice, theta_s, theta_b, hc, N, zeta)
+    # We have z at the midpoint of each cell, now find it on the top and
+    # bottom edges of each cell
+    z_edges = zeros((N+1, num_lat, num_lon))
+    z_edges[1:-1,:,:] = 0.5*(z[0:-1,:,:] + z[1:,:,:])
+    # At surface, z=zice; at bottom, extrapolate
+    z_edges[-1,:,:] = zice[:,:]
+    z_edges[0,:,:] = 2*z[0,:,:] - z_edges[1,:,:]
+    # Now find dz
+    dz = z_edges[1:,:,:] - z_edges[0:-1,:,:]    
+
+    return dx, dy, dz, z

circumpolar_plot.py

@@ -1,7 +1,7 @@
 from netCDF4 import Dataset
 from numpy import *
 from matplotlib.pyplot import *
-from calc_z import *
+from cartesian_grid_3d import *
 
 # Make a circumpolar Antarctic plot of the given (horizontal) ROMS variable.
 # Input:
@@ -111,28 +111,17 @@ def circumpolar_plot (file_path, var_name, tstep, depth_key, depth, depth_bounds
             # Bottom layer
             data = data_full[0,:,:]
         else:
-            # We will need to calculate z-coordinates
-            z, sc_r, Cs_r = calc_z(h, zice, theta_s, theta_b, hc, N)
+            # We will need z-coordinates and possibly dz
+            dx, dy, dz, z = cartesian_grid_3d(lon, lat, h, zice, theta_s, theta_b, hc, N)
             if depth_key == 2:
                 # Interpolate to given depth
                 data = interp_depth(data_full, z, depth)
-            elif depth_key in [3, 4]:
-                # We need dz for averaging
-                # We have z on the midpoint of each cell, now find it on the top and bottom edges
-                # of each cell
-                z_edges = zeros((size(z,0)+1, size(z,1), size(z,2)))
-                z_edges[1:-1,:,:] = 0.5*(z[0:-1,:,:] + z[1:,:,:])
-                # At the surface, z=zice; at bottom, extrapolate
-                z_edges[-1,:,:] = zice[:,:]
-                z_edges[0,:,:] = 2*z[0,:,:] - z_edges[1,:,:]
-                # Now find dz
-                dz = z_edges[1:,:,:] - z_edges[0:-1,:,:]
-                if depth_key == 3:
-                    # Vertically average entire water column
-                    data = sum(data_full*dz, axis=0)/sum(dz, axis=0)
-                elif depth_key == 4:                    
-                    # Vertically average between given depths
-                    data = average_btw_depths (data_full, z, dz, depth_bounds)
+            elif depth_key == 3:
+                # Vertically average entire water column
+                data = sum(data_full*dz, axis=0)/sum(dz, axis=0)
+            elif depth_key == 4:
+                # Vertically average between given depths
+                data = average_btw_depths(data_full, z, dz, depth_bounds)
 
     # Center levels on 0 for certain variables, with a blue-to-red colourmap
     if var_name in ['u', 'v', 'ubar', 'vbar', 'm']:

file_guide.txt

@@ -16,9 +16,9 @@ romscice_ini_woa.py: Builds a ROMS initialisation file using World Ocean Atlas
                      data for temperature and salinity; starts with a motionless
 		     ocean and zero sea surface height. Under ice shelves, sets
 		     the salinity to a constant value (34.5 psu) and the
-		     temperature to the local freezing point. Depends on
-		     calc_z.py, and uses WOA data assumed to be converted from
-		     FESOM input using woa_netcdf.py
+		     temperature to the local freezing point. Uses WOA data
+		     assumed to be converted from FESOM input using
+		     woa_netcdf.py
 		     To run: First edit user parameters (file paths and grid
 		             parameters) near the bottom of the file (after the
 			     Then make sure that you have scipy version 0.14 or
@@ -30,19 +30,14 @@ romscice_ini_woa.py: Builds a ROMS initialisation file using World Ocean Atlas
 romscice_nbc.py: Builds a ROMS northern lateral boundary condition file from 1
                  year of monthly ECCO2 data for temperature, salinity, and
 		 velocity (set sea surface height to zero). Also contains useful
-		 code for calculating volume averages. Depends on calc_z.py.
+		 code for calculating volume averages.
 		 To run: Edit user parameters near the top of the script
 		         (mainly just file paths). Then make sure that you have
 			 scipy version 0.14 or higher (on raijin, this means
 			 switching to python/2.7.6; instructions at the top
 			 of the file). Then open python or ipython and type
 			 "run romscice_nbc.py".
 
-calc_z.py: Given ROMS grid variables, calculate the s-coordinates, stretching
-           curves, and z-coordinates. Assumes Vtransform=2 and Vstretching=2.
-	   To run: This is a function designed to be called from other
-	           scripts. 
-
 romscice_atm_monthly.py: Convert ERA-Interim files of monthly averaged
                          atmospheric forcing to ROMS-CICE input forcing files
 			 on the correct grid and with the correct units.
@@ -472,6 +467,24 @@ overturning_plot.py: Calculate the meridional overturning streamfunction at the
 			     and the filename with which to save the plot.
 
 
+***UTILITY FUNCTIONS***
+
+calc_z.py: Given ROMS grid variables, calculate the s-coordinates, stretching
+           curves, and z-coordinates. Assumes Vtransform=2 and Vstretching=2.
+	   To run: This is a function designed to be called from other
+	           scripts. See for example romscice_nbc.py.
+
+cartesian_grid_2d.py: Given ROMS grid variables, calculate 2D Cartesian
+                      integrands dx and dy.
+		      To run: This is a function designed to be called from
+		              other scripts. See for example massloss.py
+
+cartesian_grid_3d.py: Given ROMS grid variables, calculate 3D Cartesian
+                      integrands dx, dy, and dz, as well as 3D z-coordinates.
+		      To run: This is a function designed to be called from
+		              other scripts. See for example spinup_plots.py.
+
+
 ***OBSOLETE***
 
 roms_grid_rtopo2.py: Interpolate RTopo-2 data to an existing ROMS grid, and do
@@ -528,7 +541,7 @@ romscice_atm_subdaily.py: Convert 100 6-hour timesteps of an ERA-Interim
 
 romscice_ini_ecco.py: Builds a ROMS initialisation file using ECCO2 data for
                       temperature and salinity; starts with a motionless ocean
-		      and zero sea surface height. Depends on calc_z.py.
+		      and zero sea surface height.
 		      To run: First edit user parameters (file paths and grid
 		              parameters) near the bottom of the file. Then make
 			      sure that you have scipy version 0.14 or higher
@@ -537,8 +550,7 @@ romscice_ini_ecco.py: Builds a ROMS initialisation file using ECCO2 data for
 			      python/ipython and type "run romscice_ini.py".
 
 convert_ecco.job: A batch job which converts 1 year of ECCO2 data into
-                  northern lateral bounday conditions for ROMS. Depends on
-		  romscice_nbc.py, which depends on calc_z.py.
+                  northern lateral bounday conditions for ROMS.
 		  To run: First edit user parameters in romscice_nbc.py
 		          (see below). Then edit the PBS job settings at the top
 			  of this file, and the module commands if necessary to