Add sig_figs option for log_prob.

Author: tillahoffmann

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,11 @@ 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.
+            Introduced in CmdStan-2.25.
+
         :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 +1695,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 Optional
 
 import numpy as np
 import pytest
@@ -20,17 +21,36 @@
 BERN_BASENAME = 'bernoulli'
 
 
-def test_lp_good() -> None:
+@pytest.mark.parametrize("sig_figs", [15, 3, None])
+def test_lp_good(sig_figs: Optional[int]) -> 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.
+    expected_values = {
+        None: ["-7.02147", "-1.18847"],
+        3: ["-7.02", "-1.19"],
+        15: ["-7.02146677130525","-1.18847260704286"],
+    }[sig_figs]
+    for actual, expected in zip(out.values[0], expected_values):
+        assert str(actual) == expected
+
     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())
 
+    expected_values = {
+        None: ["-5.53959", "-1.49039"],
+        3: ["-5.54", "-1.49"],
+        15: ["-5.53959011989073", "-1.49039383920238"],
+    }[sig_figs]
+    for actual, expected in zip(out_unadjusted.values[0], expected_values):
+        assert str(actual) == expected
+
 
 def test_lp_bad(
     caplog: pytest.LogCaptureFixture,