Merge pull request #692 from tillahoffmann/log-prob-sig-figs

Add `sig_figs` option for `log_prob`.

Author: WardBrian

cmdstanpy/model.py

@@ -1640,6 +1640,7 @@ def log_prob(
         data: Union[Mapping[str, Any], str, os.PathLike, None] = None,
         *,
         jacobian: bool = True,
+        sig_figs: Optional[int] = None,
     ) -> pd.DataFrame:
         """
         Calculate the log probability and gradient at the given parameter
@@ -1663,6 +1664,10 @@ def log_prob(
         :param jacobian: Whether or not to enable the Jacobian adjustment
             for constrained parameters. Defaults to ``True``.
 
+        :param sig_figs: Numerical precision used for output CSV and text files.
+            Must be an integer between 1 and 18.  If unspecified, the default
+            precision for the system file I/O is used; the usual value is 6.
+
         :return: A pandas.DataFrame containing columns "lp__" and additional
             columns for the gradient values. These gradients will be for the
             unconstrained parameters of the model.
@@ -1689,6 +1694,8 @@ def log_prob(
 
             output = os.path.join(output_dir, "output.csv")
             cmd += ["output", f"file={output}"]
+            if sig_figs is not None:
+                cmd.append(f"sig_figs={sig_figs}")
 
             get_logger().debug("Cmd: %s", str(cmd))
 

test/test_log_prob.py

@@ -4,6 +4,7 @@
 import os
 import re
 from test import check_present
+from typing import List, Optional
 
 import numpy as np
 import pytest
@@ -20,17 +21,31 @@
 BERN_BASENAME = 'bernoulli'
 
 
-def test_lp_good() -> None:
+@pytest.mark.parametrize("sig_figs, expected, expected_unadjusted", [
+    (11, ["-7.0214667713","-1.188472607"], ["-5.5395901199", "-1.4903938392"]),
+    (3, ["-7.02", "-1.19"], ["-5.54", "-1.49"]),
+    (None, ["-7.02147", "-1.18847"], ["-5.53959", "-1.49039"])
+])
+def test_lp_good(sig_figs: Optional[int], expected: List[str],
+                 expected_unadjusted: List[str]) -> None:
     model = CmdStanModel(stan_file=BERN_STAN)
-    out = model.log_prob({"theta": 0.1}, data=BERN_DATA)
+    params = {"theta": 0.34903938392023830482}
+    out = model.log_prob(params, data=BERN_DATA, sig_figs=sig_figs)
     assert "lp_" in out.columns[0]
 
+    # Check the number of digits.
+    for actual, value in zip(out.values[0], expected):
+        assert str(actual) == value
+
     out_unadjusted = model.log_prob(
-        {"theta": 0.1}, data=BERN_DATA, jacobian=False
+        params, data=BERN_DATA, jacobian=False, sig_figs=sig_figs
     )
     assert "lp_" in out_unadjusted.columns[0]
     assert not np.allclose(out.to_numpy(), out_unadjusted.to_numpy())
 
+    for actual, value in zip(out_unadjusted.values[0], expected_unadjusted):
+        assert str(actual) == value
+
 
 def test_lp_bad(
     caplog: pytest.LogCaptureFixture,