Moved corner interpolation methods into `cso_tools`. (357b802d) · Commits · CAMS / CSO

src/cso/cso_plot.py

+6 −218

Original line number	Diff line number	Diff line
		@@ -1484,226 +1484,10 @@ def GetColorMap(colors=None, color_under=None, color_over=None, color_bad=None,

		# enddef # GetColorMap

		# ***


		def mid2bounds(x):
		"""
		Return 1D arrays with boundary values.
		"""

		# modules:
		import numpy

		# size:
		nx = len(x)

		# result:
		xb = numpy.zeros((nx + 1), float)
		# edge:
		xb[0] = x[0] - 0.5 * (x[1] - x[0])
		# mid:
		xb[1:nx] = 0.5 * ( numpy.array(x[:nx-1]) + numpy.array(x[1:nx]) )
		# end:
		xb[nx] = x[nx - 1] + 0.5 * (x[nx - 1] - x[nx - 2])

		# ok
		return xb


		# enddef mid2bounds

		# ***


		def mid2corners(xx, dim=None):
		"""
		Return 2D fields with corner values.
		Eventually in single dim "x" or "y" only.
		"""

		# modules:
		import numpy

		# size:
		ny, nx = xx.shape

		# extend in x-direction?
		if (dim is None) or (dim == "x") :
		# intermediate:
		bxx = numpy.zeros((ny, nx + 1), float)
		# fill:
		for iy in range(ny):
		bxx[iy, :] = mid2bounds(xx[iy, :])
		# endfor
		else :
		# create copy:
		bxx = xx * 1.0
		#endif

		# extend in x-direction?
		if (dim is None) or (dim == "y") :
		# target:
		cxx = numpy.zeros((ny + 1, nx + 1), float)
		# fill:
		for ix in range(nx + 1):
		cxx[:, ix] = mid2bounds(bxx[:, ix])
		# endfor
		else :
		# create copy:
		cxx = bxx * 1.0
		#endif

		# ok
		return cxx


		# enddef mid2corners

		# *


		def GetGrid( shp, x=None, y=None, domain=None ):
		"""
		Return 2D grid arrays with corner points.

		Arguments:

		* ``shp`` : shape (ny,nx) of values

		Optional grid definitions:

		* ``x`` and ``y`` are either 1D or 2D and define the coordinates;
		supported shapes:

		* 2D corner field with shape ``(ny+1,nx+1)`` ; the result will have this shape too;
		* 2D mid of edges with shape ``(ny,nx+1)`` or ``(ny+1,nx)`` ;
		* 2D mid of cell with shape ``(ny,nx)`` ;
		* 1D corner points with shape ``(nx+1)`` or ``(ny+1)``;
		* 1D mid of cell with shape ``(nx)`` and ``(ny)``.

		* ``domain=(left,right,bottom,top)`` defines the boundaries of the domain;
		default = ``(0,nx,0,ny)``.

		If one or both of ``x`` and ``y`` are not defined, the ``domain`` or its default
		is used to define a regular grid in one or both of the directions.
		"""

		# modules:
		import numpy

		# extract:
		ny, nx = shp

		# domain definition:
		if domain is None:
		left, right, bottom, top = 0, nx, 0, ny
		else:
		left, right, bottom, top = domain
		# endif

		# form xx:
		if x is None :
		# regular boundaries:
		xb = left + numpy.arange(nx + 1) * (right - left) / float(nx)
		# extend to 2D:
		xx = numpy.zeros((ny+1,nx+1))
		for j in range(ny+1) :
		xx[j,:] = xb
		#endfor
		#
		elif x.shape == (ny+1,nx+1) :
		# copy:
		xx = x
		#
		elif x.shape == (ny,nx) :
		# extrapolate both directions:
		xx = mid2corners(x)
		#
		elif x.shape == (ny,nx+1) :
		# only in one direction:
		xx = mid2corners(x,dim="y")
		#
		elif x.shape == (ny+1,nx) :
		# only in one direction:
		xx = mid2corners(x,dim="x")
		#
		elif x.shape == (nx+1,):
		# extend to 2D:
		xx = numpy.zeros((ny+1,nx+1))
		for j in range(ny+1) :
		xx[j,:] = x
		#endfor
		#
		elif x.shape == (nx,):
		# form bounds:
		xb = mid2bounds(x)
		# extend to 2D:
		xx = numpy.zeros((ny+1,nx+1))
		for j in range(ny+1) :
		xx[j,:] = xb
		#endfor
		#
		else :
		print( f"ERROR - unsupported x shape {x.shape} for field shape {shp}" )
		raise Exception
		#endif

		# form yy:
		if y is None :
		# regular boundaries:
		yb = bottom + numpy.arange(ny + 1) * (top - bottom) / float(ny)
		# extend to 2D:
		yy = numpy.zeros((ny+1,nx+1))
		for i in range(nx+1) :
		yy[:,i] = yb
		#endfor
		#
		elif y.shape == (ny+1,nx+1) :
		# copy:
		yy = y
		#
		elif y.shape == (ny,nx) :
		# extrapolate both directions:
		yy = mid2corners(y)
		#
		elif y.shape == (ny,nx+1) :
		# only in one direction:
		yy = mid2corners(y,dim="y")
		#
		elif y.shape == (ny+1,nx) :
		# only in one direction:
		yy = mid2corners(y,dim="x")
		#
		elif y.shape == (ny+1,):
		# extend to 2D:
		yy = numpy.zeros((ny+1,nx+1))
		for i in range(nx+1) :
		yy[:,i] = y
		#endfor
		#
		elif y.shape == (ny,):
		# form bounds:
		yb = mid2bounds(y)
		# extend to 2D:
		yy = numpy.zeros((ny+1,nx+1))
		for i in range(nx+1) :
		yy[:,i] = yb
		#endfor
		#
		else :
		print( f"ERROR - unsupported y shape {y.shape} for field shape {shp}" )
		raise Eyception
		#endif

		# ok
		return xx, yy

		# enddef GetGrid


		# *

		def MapFigure( domain=None ):

		"""
		@@ -1749,7 +1533,7 @@ def QuickPat(

		Plot 2D data as colored pattern.

		Arguments passed to :py:meth:`GetGrid` :
		Arguments passed to :py:meth:`.cso_tools.GetCornerGridX` and :py:meth:`.cso_tools.GetCornerGridY` :

		* ``x``, ``y`` : 1D or 2D coordinates; used to define 2D corner fields;
		* ``domain=(left,right,bottom,top)`` : coordinate bounds; if ``domain`` is ``None`` or
		@@ -1788,6 +1572,9 @@ def QuickPat(
		# modules:
		import numpy

		# tools:
		from . import cso_tools

		# extract known coordinates:
		if hasattr(cc,'coords') :
		# longitudes:
		@@ -1820,7 +1607,8 @@ def QuickPat(
		# endif

		# corner points:
		xx, yy = GetGrid( cc.shape, x=x, y=y, domain=domain )
		xx = cso_tools.GetCornerGridX( cc.shape, x=x, domain=domain )
		yy = cso_tools.GetCornerGridY( cc.shape, y=y, domain=domain )

		# set domain if not present yet:
		if domain is None:

src/cso/cso_tools.py

+290 −0

Original line number	Diff line number	Diff line
		@@ -4,6 +4,9 @@
		# 2024-03, Arjo Segers
		# Tool module, started with copy of "linearize_avg_kernel" method from Enrico Dammers.
		#
		# 2024-03, Arjo Segers
		# Added `mid2bounds`, `mid2corners`, `GetCornerGridX`, and `GetCornerGridY` methods.
		#

		########################################################################
		###
		@@ -25,6 +28,293 @@ Methods
		"""


		########################################################################
		###
		### grid tools
		###
		########################################################################


		def mid2bounds(x):
		"""
		Given input ``x`` with shape ``(nx,)``,
		return array ``xb`` with shape ``(nx+1,)`` with boundary values
		between and outside the input locations.
		"""

		# modules:
		import numpy

		# size:
		nx = len(x)

		# result:
		xb = numpy.zeros((nx + 1), float)
		# edge:
		xb[0] = x[0] - 0.5 * (x[1] - x[0])
		# mid:
		xb[1:nx] = 0.5 * ( numpy.array(x[:nx-1]) + numpy.array(x[1:nx]) )
		# end:
		xb[nx] = x[nx - 1] + 0.5 * (x[nx - 1] - x[nx - 2])

		# ok
		return xb


		# enddef mid2bounds

		# ***


		def mid2corners(xx, dim=None):
		"""
		Return 2D fields with corner values.
		Eventually in single dim "x" or "y" only.
		"""

		# modules:
		import numpy

		# size:
		ny, nx = xx.shape

		# extend in x-direction?
		if (dim is None) or (dim == "x") :
		# intermediate:
		bxx = numpy.zeros((ny, nx + 1), float)
		# fill:
		for iy in range(ny):
		bxx[iy, :] = mid2bounds(xx[iy, :])
		# endfor
		else :
		# create copy:
		bxx = xx * 1.0
		#endif

		# extend in x-direction?
		if (dim is None) or (dim == "y") :
		# target:
		cxx = numpy.zeros((ny + 1, nx + 1), float)
		# fill:
		for ix in range(nx + 1):
		cxx[:, ix] = mid2bounds(bxx[:, ix])
		# endfor
		else :
		# create copy:
		cxx = bxx * 1.0
		#endif

		# ok
		return cxx


		# enddef mid2corners

		# *


		def GetCornerGridX( shp, x=None, domain=None, bounds=False ):
		"""
		Return grid array of shape ``(ny+1,nx+1)`` with 'x' locations of corner points.
		If ``bounds=True``, the result is a ``(ny,nx,4)`` array.
		or

		Arguments:

		* ``shp`` : shape ``(ny,nx)`` of values

		Optional grid definitions:

		* ``x`` is either 1D or 2D and define the coordinates; supported shapes:

		* 2D corner field with shape ``(ny+1,nx+1)`` ; the result will have this shape too;
		* 2D mid of edges with shape ``(ny,nx+1)`` ;
		* 2D mid of cell with shape ``(ny,nx)`` ;
		* 1D corner points with shape ``(nx+1)``;
		* 1D mid of cell with shape ``(nx)``.

		* ``domain=(left,right,bottom,top)`` defines the boundaries of the domain;
		default = ``(0,nx,0,ny)``.

		If ``x`` is not defined, the ``domain`` or its default is used to define a regular grid.
		"""

		# modules:
		import numpy

		# extract:
		ny, nx = shp

		# domain definition:
		if domain is None:
		left, right, bottom, top = 0, nx, 0, ny
		else:
		left, right, bottom, top = domain
		# endif

		# form xx:
		if x is None :
		# regular boundaries:
		xb = left + numpy.arange(nx + 1) * (right - left) / float(nx)
		# extend to 2D:
		xx = numpy.zeros((ny+1,nx+1))
		for j in range(ny+1) :
		xx[j,:] = xb
		#endfor
		#
		elif x.shape == (ny+1,nx+1) :
		# copy:
		xx = x
		#
		elif x.shape == (ny,nx) :
		# extrapolate both directions:
		xx = mid2corners(x)
		#
		elif x.shape == (ny,nx+1) :
		# only in one direction:
		xx = mid2corners(x,dim="y")
		#
		elif x.shape == (ny+1,nx) :
		# only in one direction:
		xx = mid2corners(x,dim="x")
		#
		elif x.shape == (nx+1,):
		# extend to 2D:
		xx = numpy.zeros((ny+1,nx+1))
		for j in range(ny+1) :
		xx[j,:] = x
		#endfor
		#
		elif x.shape == (nx,):
		# form bounds:
		xb = mid2bounds(x)
		# extend to 2D:
		xx = numpy.zeros((ny+1,nx+1))
		for j in range(ny+1) :
		xx[j,:] = xb
		#endfor
		#
		else :
		print( f"ERROR - unsupported x shape {x.shape} for field shape {shp}" )
		raise Exception
		#endif

		# ok
		if bounds :
		xxb = numpy.zeros((ny,nx,4))
		xxb[:,:,0] = xx[0:ny ,0:nx ]
		xxb[:,:,1] = xx[0:ny ,1:nx+1]
		xxb[:,:,2] = xx[1:ny+1,1:nx+1]
		xxb[:,:,3] = xx[1:ny+1,0:nx ]
		return xxb
		else :
		return xx
		#endif

		# enddef GetCornerGridX

		# *


		def GetCornerGridY( shp, y=None, domain=None, bounds=False ):
		"""
		Return grid array of shape ``(ny+1,nx+1)`` with 'x' locations of corner points.
		If ``bounds=True``, the result is a ``(ny,nx,4)`` array.

		Arguments:

		* ``shp`` : shape (ny,nx) of values

		Optional grid definitions:

		* ``y`` is either 1D or 2D and define the coordinates; supported shapes:

		* 2D corner field with shape ``(ny+1,nx+1)`` ; the result will have this shape too;
		* 2D mid of edges with shape ``(ny+1,nx)`` ;
		* 2D mid of cell with shape ``(ny,nx)`` ;
		* 1D corner points with shape ``(ny+1)``;
		* 1D mid of cell with shape ``(ny)``.

		* ``domain=(left,right,bottom,top)`` defines the boundaries of the domain; default = ``(0,nx,0,ny)``.

		If ``y`` is not defined, the ``domain`` or its default is used define a regular grid.
		"""

		# modules:
		import numpy

		# extract:
		ny, nx = shp

		# domain definition:
		if domain is None:
		left, right, bottom, top = 0, nx, 0, ny
		else:
		left, right, bottom, top = domain
		# endif

		# form yy:
		if y is None :
		# regular boundaries:
		yb = bottom + numpy.arange(ny + 1) * (top - bottom) / float(ny)
		# extend to 2D:
		yy = numpy.zeros((ny+1,nx+1))
		for i in range(nx+1) :
		yy[:,i] = yb
		#endfor
		#
		elif y.shape == (ny+1,nx+1) :
		# copy:
		yy = y
		#
		elif y.shape == (ny,nx) :
		# extrapolate both directions:
		yy = mid2corners(y)
		#
		elif y.shape == (ny,nx+1) :
		# only in one direction:
		yy = mid2corners(y,dim="y")
		#
		elif y.shape == (ny+1,nx) :
		# only in one direction:
		yy = mid2corners(y,dim="x")
		#
		elif y.shape == (ny+1,):
		# extend to 2D:
		yy = numpy.zeros((ny+1,nx+1))
		for i in range(nx+1) :
		yy[:,i] = y
		#endfor
		#
		elif y.shape == (ny,):
		# form bounds:
		yb = mid2bounds(y)
		# extend to 2D:
		yy = numpy.zeros((ny+1,nx+1))
		for i in range(nx+1) :
		yy[:,i] = yb
		#endfor
		#
		else :
		print( f"ERROR - unsupported y shape {y.shape} for field shape {shp}" )
		raise Eyception
		#endif

		# ok
		if bounds :
		yyb = numpy.zeros((ny,nx,4))
		yyb[:,:,0] = yy[0:ny ,0:nx ]
		yyb[:,:,1] = yy[0:ny ,1:nx+1]
		yyb[:,:,2] = yy[1:ny+1,1:nx+1]
		yyb[:,:,3] = yy[1:ny+1,0:nx ]
		return yyb
		else :
		return yy
		#endif

		# enddef GetCornerGridY


		########################################################################
		###
		### averaging kernel tool