add geostrophic adjustment exmple

Author: miniufo

README.md

@@ -70,6 +70,7 @@ This is a list of the problems that can be solved by `xinvert`:
 | <img src="./pics/Gallery_SWMReference.png" width="380"><br/>[invert PV balance equation for<br/> steady reference state](./docs/source/notebooks/05_reference_SWM.ipynb)| <img src="./pics/Gallery_GillMatsuno.png" width="380"><br/>[invert Gill-Matsuno model for<br/> wind and mass fields](./docs/source/notebooks/07_Gill_Matsuno_model.ipynb) |
 | <img src="./pics/Gallery_StommelMunk.png" width="380"><br/>[invert Stommel-Munk model for<br/> wind-driven ocean circulation](./docs/source/notebooks/08_Stommel_Munk_model.ipynb) | <img src="./pics/Gallery_Fofonoff.png" width="380"><br/>[invert Fofonoff model for<br/> inviscid/adiabatic steady state](./docs/source/notebooks/09_Fofonoff_flow.ipynb) |
 | <img src="./pics/Gallery_Bretherton.png" width="380"><br/>[invert Bretherton model for<br/> steady flow over topography](./docs/source/notebooks/10_Bretherton_flow_over_topography.ipynb) | <img src="./pics/Gallery_Omega.png" width="380"><br/>[invert Omega equation for<br/> QG vertical velocity](./docs/source/notebooks/11_Omega_equation.ipynb) |
+| <img src="./pics/Gallery_GeoAdjust.png" width="380"><br/>[invert geostrophic-adjustment problem](./docs/source/notebooks/12_GeostrophicAdjustment.ipynb) | ... add more examples ... |
 
 
 

docs/source/notebooks/00_Introduction.ipynb

@@ -132,7 +132,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.9"
+   "version": "3.13.2"
   }
  },
  "nbformat": 4,

docs/source/notebooks/12_GeostrophicAdjustment.ipynb

@@ -8,34 +8,37 @@
 #%% classical cases
 import numpy as np
 import xarray as xr
-from xinvert.xinvert import invert_GeoAdjustment
-
-# def test_adjustment():
-yc   = 501
-ydef = np.linspace(-75, -25, yc)
-ydef = xr.DataArray(ydef, dims='lat', coords={'lat':ydef})
-
-R  = 6371200
-O  = 7.292e-5
-h0 = ydef - ydef + 1500
-h0[int(yc/2):] = 1520
-g  = 9.81
-
-f   = 2 * O * np.sin(np.deg2rad(ydef))
-PV0 = f / h0
-
-# invert
-iParams = {
-    'BCs'      : ['extend'],
-    'mxLoop'   : 100000,
-    'tolerance': -1e-11,
-    'optArg'   : 1.8,
-    'undef'    : -9999,
-}
-
-deg2m = R/180*np.pi
-
-h  = invert_GeoAdjustment(h0, dims=['lat'], coords='lat', iParams=iParams)
-u  = - h.differentiate('lat') / deg2m * g / f
-PV = (f - u.differentiate('lat') / deg2m) / h
+from xinvert import invert_GeoAdjustment
+
+def test_adjustment():
+	yc   = 501
+	ydef = np.linspace(-75, -25, yc)
+	ydef = xr.DataArray(ydef, dims='lat', coords={'lat':ydef})
+
+	R  = 6371200
+	O  = 7.292e-5
+	h0 = ydef - ydef + 1500
+	h0[int(yc/2):] = 1520
+	g  = 9.81
+
+	f   = 2 * O * np.sin(np.deg2rad(ydef))
+	PV0 = f / h0
+
+	# invert
+	iParams = {
+	    'BCs'      : ['extend'],
+	    'mxLoop'   : 100000,
+	    'tolerance': -1e-11,
+	    'optArg'   : 1.8,
+	    'undef'    : -9999,
+	}
+
+	deg2m = R/180*np.pi
+
+	h  = invert_GeoAdjustment(h0, dims=['lat'], coords='lat', iParams=iParams)
+	u  = - h.differentiate('lat') / deg2m * g / f
+	PV = (f - u.differentiate('lat') / deg2m) / h
+    
+    assert h.dims  == h0.dims
+    assert h.shape == h0.shape
 

pics/Gallery_GeoAdjust.png

@@ -145,7 +145,7 @@ def invert_RefState(PV, dims, coords='z-lat', icbc=None,
                       icbc, ['Ang0', 'Gamma', 'g', 'Omega', 'Rearth'], mParams, iParams)
 
 
-def invert_GeoAdjustment(h0, dims, coords='lat', icbc=None,
+def invert_GeoAdjustment(PV0, dims, coords='lat', icbc=None,
                          mParams=default_mParams, iParams=default_iParams):
     r"""(PV) inversion for a geostrophic adjustment model.
 
@@ -161,18 +161,18 @@ def invert_GeoAdjustment(h0, dims, coords='lat', icbc=None,
     .. math::
          
          A = \cos\phi / f
-         B = - f / g / h0
+         B = - q_0 \ cos\phi / g
          F = - f \cos\phi / g
     
     Invert this equation for geostrophically adjusted free surface :math:`h` given
-    the initial free surface distribution.
+    the initial PV distribution :math:`q_0`.
     
     Parameters
     ----------
-    h0: xarray.DataArray
-        Initial free surface.
+    PV0: xarray.DataArray
+        Initial PV distribution.
     dims: list
-        Dimension combination for the inversion e.g., ['lat', 'lon'].
+        Dimension name for the inversion e.g., ['lat'].
     coords: {'lat-lon', 'cartesian'}, optional
         Coordinate combinations in which inversion is performed.
     icbc: xarray.DataArray, optional
@@ -185,9 +185,9 @@ def invert_GeoAdjustment(h0, dims, coords='lat', icbc=None,
     Returns
     -------
     xarray.DataArray
-        Results (angular momentum Λ) of the SOR inversion.
+        Results (free surface after adjustment) of the SOR inversion.
     """
-    return __template(__coeffs_GeoAdjustment, inv_standard1D, 1, h0, dims, coords,
+    return __template(__coeffs_GeoAdjustment, inv_standard1D, 1, PV0, dims, coords,
                       icbc, ['g', 'Rearth', 'Omega'], mParams, iParams)
 
 
@@ -1524,12 +1524,12 @@ def diff_2nd(M, cosH, delY):
     return F, initS, (A, B)
 
 
-def __coeffs_GeoAdjustment(h0, dims, coords, mParams, iParams, icbc):
+def __coeffs_GeoAdjustment(PV0, dims, coords, mParams, iParams, icbc):
     """Calculating coefficients for geostrophic adjustment model."""
     g     = mParams['g']
     Omega = mParams['Omega']
 
-    maskF, initS, zero = __mask_FS(h0, dims, iParams, icbc)
+    maskF, initS, zero = __mask_FS(PV0, dims, iParams, icbc)
 
     if coords.lower() == 'lat': # dims[0] is θ, dims[1] is lat
         lats = np.deg2rad(maskF[dims[0]])
@@ -1539,7 +1539,7 @@ def __coeffs_GeoAdjustment(h0, dims, coords, mParams, iParams, icbc):
         fH   = 2 * Omega * np.sin((lats+lats.shift({dims[0]:1}))/2.0)
 
         A = zero + cosH / fH
-        B = zero - f * cosG / g / h0
+        B = zero - PV0 * cosG / g
         F = zero - f * cosG / g
 
     elif coords.lower() == 'cartesian': # dims[0] is θ, dims[1] is r