Add opt_pressured_kCO2 to results dict and fix derivatives

PyCO2SYS/engine/nd.py

@@ -297,6 +297,8 @@ def _get_results_dict(
             # Added in v1.8.0:
             "pressure_atmosphere": args["pressure_atmosphere"],
             "pressure_atmosphere_out": args["pressure_atmosphere_out"],
+            # Added in v1.8.3 (but should have been in v1.8.2):
+            "opt_pressured_kCO2": args["opt_pressured_kCO2"],
         }
     )
     results.update(_get_in_out(core_in, others_in, k_constants_in, suffix=""))

PyCO2SYS/uncertainty.py

@@ -413,13 +413,14 @@ def forward_nd(
             if gradable.startswith("k_") and not gradable.endswith("_out")
         ]
     )
-    assert np.all(
-        np.isin(grads_of, gradables)
-    ), "PyCO2SYS error: all grads_of must be in the list at PyCO2SYS.engine.nd.gradables."
+    assert np.all(np.isin(grads_of, gradables)), (
+        "PyCO2SYS error: all grads_of must be in the list at "
+        + "PyCO2SYS.engine.nd.gradables."
+    )
     if np.any([of.endswith("_out") for of in grads_of]):
         assert "temperature_out" in CO2SYS_nd_results, (
-            "PyCO2SYS error: you can only get gradients at output conditions if you calculated"
-            + " results at output conditions!"
+            "PyCO2SYS error: you can only get gradients at output conditions if you "
+            + "calculated results at output conditions!"
         )
     # Extract CO2SYS_nd fixed args from CO2SYS_nd_results and CO2SYS_nd_kwargs.
     # These are arguments that always get a specific value, rather than being calculated
@@ -444,6 +445,7 @@ def forward_nd(
             "opt_pH_scale",
             "opt_total_borate",
             "opt_buffers_mode",
+            "opt_pressured_kCO2",
         ]
         + list(CO2SYS_nd_kwargs.keys())
     )

docs/versions.md

@@ -43,10 +43,14 @@ Adds atmospheric pressure input for *p*CO<sub>2</sub>-*f*CO<sub>2</sub>-*x*CO<su
 
     * Reverted default `opt_k_carbonic` to `10` (i.e., [LDK00](../refs/#l)) for consistency with best practice guide.
 
+    ***Bug fixes***
+
+    * Updated `pyco2.equilibria.p1atm.kH2CO3_NBS_MCHP73` (used for `opt_k_carbonic` options `6` and `7`) to update any salinity values less than 10<sup>–16</sup> to be 10<sup>–16</sup>, because zero salinities give a NaN for <i>K</i><sub>2</sub>, which causes autograd problems.  This should not make any practical difference, because the parameterisation is only valid for salinities above 19.
+    * Added `opt_pressured_kCO2` to results dict and incorporated it properly into the uncertainty propagation functions.
+
     ***Technical***
 
     * Updated from building with setup.py to pyproject.toml.
-    * Updated `pyco2.equilibria.p1atm.kH2CO3_NBS_MCHP73` (used for `opt_k_carbonic` options `6` and `7`) to update any salinity values less than 10<sup>–16</sup> to be 10<sup>–16</sup>, because zero salinities give a NaN for <i>K</i><sub>2</sub>, which causes autograd problems.  This should not make any practical difference, because the parameterisation is only valid for salinities above 19.
 
 ### 1.8.2 (19 January 2023)
 

tests/test_dlnpCO2_dT.py

@@ -3,7 +3,7 @@
 
 # Simulate arguments
 rng = np.random.default_rng(1)
-npts = 1000
+npts = 10_000
 alkalinity = rng.normal(loc=2300, scale=50, size=npts)
 dic = rng.normal(loc=2150, scale=50, size=npts)
 salinity = rng.uniform(low=0, high=40, size=npts)
@@ -16,6 +16,7 @@
 total_ammonia = rng.uniform(low=0, high=100, size=npts)
 total_sulfide = rng.uniform(low=0, high=100, size=npts)
 opt_k_carbonic = rng.integers(low=1, high=18, size=npts)
+opt_pressured_kCO2 = rng.integers(low=0, high=2, size=npts)
 
 # Solve
 results = pyco2.sys(
@@ -33,7 +34,7 @@
     total_ammonia=total_ammonia,
     total_sulfide=total_sulfide,
     opt_k_carbonic=opt_k_carbonic,
-    opt_pressured_kCO2=0,
+    opt_pressured_kCO2=opt_pressured_kCO2,
     grads_of=["pCO2"],
     grads_wrt=["temperature"],
 )
@@ -51,9 +52,23 @@
 
 def test_dlnpCO2_dT():
     """Are all of the autograd values the same as the forward finite difference values
-    to within 0.02%?
+    to within:
+        - 0.03% where opt_pressured_kCO2 is 0?
+        - 2.00% where opt_pressured_kCO2 is 1?
+        - 1e-3 everywhere?
+        - 1e-6 in the mean across all inputs?
+        - 0.01% in the mean across all inputs?
     """
-    assert np.allclose(dlnpCO2_dT_autograd, dlnpCO2_dT_ffd, rtol=0.02 / 100, atol=0)
+    L = opt_pressured_kCO2 == 0
+    assert np.allclose(
+        dlnpCO2_dT_autograd[L], dlnpCO2_dT_ffd[L], rtol=0.03 / 100, atol=0
+    )
+    assert np.allclose(
+        dlnpCO2_dT_autograd[~L], dlnpCO2_dT_ffd[~L], rtol=2 / 100, atol=0
+    )
+    assert np.allclose(dlnpCO2_dT_autograd, dlnpCO2_dT_ffd, rtol=0, atol=1e-3)
+    assert np.mean(abs_diff) < 1e-6
+    assert np.mean(abs_diff_pct) < 0.01
 
 
 test_dlnpCO2_dT()