Merge pull request #134 from efiring/gibbs

Add gibbs and gibbs_ice via ufunc wrapping.

Author: efiring

.pre-commit-config.yaml

@@ -13,7 +13,8 @@ repos:
       exclude: >
           (?x)^(
               .*\.c|
-              tools/fix_wrapped_ufunc_typos\.py
+              tools/fix_wrapped_ufunc_typos\.py|
+              gsw/tests/test_gibbs\.py
           )$
       args:
         - --ignore-words-list=nin,preformed,wih,

gsw/_fixed_wrapped_ufuncs.py

@@ -3,6 +3,8 @@
 Users should import only from non-private modules, of course.
 """
 
+import numpy
+
 from ._wrapped_ufuncs import *
 
 _p_from_z = p_from_z
@@ -15,5 +17,25 @@ def z_from_p(p, lat, geo_strf_dyn_height=0, sea_surface_geopotential=0):
     return _z_from_p(p, lat, geo_strf_dyn_height, sea_surface_geopotential)
 z_from_p.__doc__ = _z_from_p.__doc__
 
+_gibbs = gibbs
+def gibbs(ns, nt, np, SA, t, p):
+    params = {"ns": ns, "nt": nt, "np": np}
+    for k, v in params.items():
+        u = numpy.unique(v)
+        if u.min() < 0 or u.max() > 2 or u.dtype.kind != "i":
+            raise ValueError("ns, nt, np must contain integers 0, 1, or 2;"
+                             f" found {k}={v}")
+    return _gibbs(ns, nt, np, SA, t, p)
+gibbs.__doc__ = _gibbs.__doc__
 
 
+_gibbs_ice = gibbs_ice
+def gibbs_ice(nt, np, t, p):
+    params = {"nt": nt, "np": np}
+    for k, v in params.items():
+        u = numpy.unique(v)
+        if u.min() < 0 or u.max() > 2 or u.dtype.kind != "i":
+            raise ValueError("nt, np must contain integers 0, 1, or 2;"
+                             f" found {k}={v}")
+    return _gibbs_ice(nt, np, t, p)
+gibbs_ice.__doc__ = _gibbs_ice.__doc__

gsw/_utilities.py

@@ -36,6 +36,16 @@ def wrapper(*args, **kw):
         hasmasked = np.any(ismasked)
         hasduck = np.any(isduck)
 
+        # Handle the leading integer arguments in gibbs and gibbs_ice.
+        # Wrapped ufuncs are constructed with the "types" attribute from the
+        # underlying ufunc.
+        if hasattr(f, "types"):
+            argtypes = f.types[0].split("->")[0]
+            first_double = argtypes.index("d")
+        else:
+            first_double = 0
+
+
         def fixup(ret):
             if hasduck:
                 return ret
@@ -50,7 +60,9 @@ def fixup(ret):
 
         newargs = []
         for i, arg in enumerate(args):
-            if ismasked[i]:
+            if i < first_double:
+                newargs.append(arg)  # for gibbs and gibbs_ice
+            elif ismasked[i]:
                 newargs.append(masked_to_nan(arg))
             elif isduck[i]:
                 newargs.append(arg)