plot sections of constant lon/lat from irregular ROMS grid

sections_roms.py

import pyroms
import matplotlib.pylab as plt
import numpy as np

def lon_section(grd,lon,debug=False):
	''' return i,j for the section of given lon '''
	# use a single contour of the desired longitude
	cont = plt.contour(grd.hgrid.lon_vert,[lon])
	plt.close() # contour triggers a figure
	isection, jsection = get_broken_line_from_contour(cont,debug=debug)
	return isection, jsection

def lat_section(grd,lat,debug=False):
	''' return i,j for the section of given lat '''
	# use a single contour of the desired latitude
	cont = plt.contour(grd.hgrid.lat_vert,[lat])
	plt.close() # contour triggers a figure
	isection, jsection = get_broken_line_from_contour(cont,debug=debug)
	return isection, jsection

def get_broken_line_from_contour(contour,debug=False):
	''' return a broken line from contour, suitable to integrate transport '''
	# first guess of our broken line is the contour position in (x,y) 
	# raw array has float grid indices
	xseg_raw = contour.allsegs[0][0][:,0]
	yseg_raw = contour.allsegs[0][0][:,1]

	# we round them up to the closest integer indices, this is our first guess
	nseg = len(xseg_raw)
	iseg_fg = np.zeros((nseg))
	jseg_fg = np.zeros((nseg))

	for kseg in np.arange(nseg):
		# NB: flooring works actually better than rounding as it reduces wiggles
		#     seen on arctic2 using constant longitude.
		#iseg_fg[kseg] = round(xseg_raw[kseg])
		#jseg_fg[kseg] = round(yseg_raw[kseg])
		iseg_fg[kseg] = int(xseg_raw[kseg])
		jseg_fg[kseg] = int(yseg_raw[kseg])

	# create a second guess by forcing broken line along cell edges
	iseg_sg = [iseg_fg[0]]
	jseg_sg = [jseg_fg[0]]

	for kseg in np.arange(1,nseg):
		if ( iseg_fg[kseg] - iseg_fg[kseg-1] == 0 ) or ( jseg_fg[kseg] - jseg_fg[kseg-1] == 0 ):
			# we are along one face of the cell so we're good
			iseg_sg.append(iseg_fg[kseg])
			jseg_sg.append(jseg_fg[kseg])
		else:
			# we need to create two segments stopping by an intermediate cell corner
			iseg_sg.append(iseg_fg[kseg])
			jseg_sg.append(jseg_fg[kseg-1])
			iseg_sg.append(iseg_fg[kseg])
			jseg_sg.append(jseg_fg[kseg])

	iseg_sg = np.array(iseg_sg)
	jseg_sg = np.array(jseg_sg)

	if debug:
		plt.figure()
		plt.plot(iseg_fg,jseg_fg,'r')
		plt.plot(iseg_sg,jseg_sg,'k',alpha=0.5)
		plt.show()

	return iseg_sg, jseg_sg


# testing the code with NWA
mylat_nwa = 20.
mylon_nwa = 290.
grd = pyroms.grid.get_ROMS_grid('NWA')

isec, jsec = lon_section(grd,mylon_nwa,debug=True)
isec, jsec = lat_section(grd,mylat_nwa,debug=True)

# testing the code with Arctic2
mylat_arctic2 = 72.
mylon_arctic2 = 180.
grd = pyroms.grid.get_ROMS_grid('ARCTIC2')

isec, jsec = lat_section(grd,mylat_arctic2,debug=True)
isec, jsec = lon_section(grd,mylon_arctic2,debug=True)