Add save_diagnostics to laplace_sample and expose the Hessian

cmdstanpy/model.py

@@ -1689,6 +1689,7 @@ def laplace_sample(
         seed: int | None = None,
         output_dir: OptionalPath = None,
         sig_figs: int | None = None,
+        save_diagnostics: bool = False,
         save_profile: bool = False,
         show_console: bool = False,
         refresh: int | None = None,
@@ -1728,6 +1729,13 @@ def laplace_sample(
             precision for the system file I/O is used; the usual value is 6.
             Introduced in CmdStan-2.25.
 
+        :param save_diagnostics: Whether or not to save CmdStan's diagnostic
+            file.  If ``True``, a JSON file is written to
+            '<model_name>-<YYYYMMDDHHMM>_diagnostic.json'.  It holds the
+            Hessian of the log density at the mode, on the unconstrained
+            scale, which is available as :attr:`CmdStanLaplace.hessian`.
+            Default is ``False``.
+
         :param save_profile: Whether or not to profile auto-diff operations in
             labelled blocks of code.  If ``True``, CSV outputs are written to
             file '<model_name>-<YYYYMMDDHHMM>_profile.csv'.
@@ -1812,6 +1820,7 @@ def laplace_sample(
                 seed=seed,
                 output_dir=output_dir,
                 sig_figs=sig_figs,
+                save_latent_dynamics=save_diagnostics,
                 save_profile=save_profile,
                 method_args=laplace_args,
                 refresh=refresh,
@@ -1830,6 +1839,9 @@ def laplace_sample(
             config_file=runset.config_files[0],
             stdout_file=runset.stdout_files[0],
             mode=cmdstan_mode,
+            diagnostic_file=(
+                runset.diagnostic_files[0] if save_diagnostics else None
+            ),
         )
 
     def _run_cmdstan(

cmdstanpy/stanfit/laplace.py

@@ -4,11 +4,12 @@
 
 from __future__ import annotations
 
+import json
 import os
 from collections.abc import Hashable, Mapping, Sequence
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Any, MutableMapping
+from typing import Any, ClassVar, MutableMapping
 
 import numpy as np
 import pandas as pd
@@ -195,6 +196,16 @@ class CmdStanLaplace(SingleFileFit[LaplaceConfig]):
     """
 
     mode: CmdStanMLE
+    diagnostic_file: str | None = None
+
+    _FILE_ATTRS: ClassVar[tuple[str, ...]] = (
+        'csv_file',
+        'config_file',
+        'stdout_file',
+        'diagnostic_file',
+    )
+
+    _hessian: np.ndarray | None = field(default=None, init=False)
 
     @classmethod
     def from_files(
@@ -203,6 +214,7 @@ def from_files(
         config_file: str | os.PathLike,
         stdout_file: str | os.PathLike | None = None,
         mode: CmdStanMLE | None = None,
+        diagnostic_file: str | os.PathLike | None = None,
     ) -> CmdStanLaplace:
         kwargs = cls._from_files_kwargs(
             csv_file, config_file, stdout_file, LaplaceRunConfig
@@ -211,7 +223,31 @@ def from_files(
             mode = _mode_from_files(
                 kwargs['config'].method_config.mode, csv_file
             )
-        return cls(mode=mode, **kwargs)
+        return cls(
+            mode=mode,
+            diagnostic_file=(
+                os.fspath(diagnostic_file)
+                if diagnostic_file is not None
+                else None
+            ),
+            **kwargs,
+        )
+
+    @property
+    def hessian(self) -> np.ndarray | None:
+        """
+        The Hessian of the log density at the mode, on the unconstrained
+        scale, read from CmdStan's diagnostic file.  This is the matrix from
+        which the Laplace approximation is built.  ``None`` if the diagnostic
+        file was not saved; see the ``save_diagnostics`` argument of
+        :meth:`CmdStanModel.laplace_sample`.
+        """
+        if self.diagnostic_file is None:
+            return None
+        if self._hessian is None:
+            with open(self.diagnostic_file) as fd:
+                self._hessian = np.array(json.load(fd)['Hessian'], dtype=float)
+        return self._hessian
 
     def save_output_files(self, dir: str | None = None) -> None:
         """Move the Laplace outputs and optimization mode outputs together."""
@@ -321,4 +357,6 @@ def __repr__(self) -> str:
             lines.append(f' config_file:\n\t{self.config_file}')
         if self.stdout_file is not None:
             lines.append(f' output_file:\n\t{self.stdout_file}')
+        if self.diagnostic_file is not None:
+            lines.append(f' diagnostic_file:\n\t{self.diagnostic_file}')
         return '\n'.join(lines)

cmdstanpy/stanfit/runset.py

@@ -76,19 +76,23 @@ def __init__(
                 ]
 
         # per-chain output files
+        # Laplace writes its diagnostic file (the Hessian) as JSON
+        diagnostic_ext = ".json" if args.method == Method.LAPLACE else ".csv"
         if chains == 1:
             self._csv_files = [self.gen_file_name(".csv")]
             if args.save_latent_dynamics:
                 self._diagnostic_files = [
-                    self.gen_file_name(".csv", extra="diagnostic")
+                    self.gen_file_name(diagnostic_ext, extra="diagnostic")
                 ]
         else:
             self._csv_files = [
                 self.gen_file_name(".csv", id=id) for id in self._chain_ids
             ]
             if args.save_latent_dynamics:
                 self._diagnostic_files = [
-                    self.gen_file_name(".csv", extra="diagnostic", id=id)
+                    self.gen_file_name(
+                        diagnostic_ext, extra="diagnostic", id=id
+                    )
                     for id in self._chain_ids
                 ]
 

test/test_laplace.py

@@ -1,5 +1,6 @@
 """Tests for the Laplace sampling method."""
 
+import json
 import os
 from pathlib import Path
 from tempfile import TemporaryDirectory
@@ -137,6 +138,62 @@ def test_laplace_outputs() -> None:
     assert 'y' in fit_pd.columns
     assert fit_pd['x'].shape == (123,)
 
+    assert fit.diagnostic_file is None
+    assert fit.hessian is None
+
+
+def test_laplace_hessian(tmp_path: Path) -> None:
+    model_file = os.path.join(DATAFILES_PATH, 'linear_regression.stan')
+    data_file = os.path.join(DATAFILES_PATH, 'linear_regression.data.json')
+    model = cmdstanpy.CmdStanModel(stan_file=model_file)
+    fit = model.laplace_sample(
+        data=data_file,
+        jacobian=False,
+        seed=1234,
+        sig_figs=18,
+        save_diagnostics=True,
+    )
+    assert fit.diagnostic_file is not None
+    assert fit.diagnostic_file.endswith('_diagnostic.json')
+
+    # unconstrained parameters: alpha, beta, log(sigma)
+    hessian = fit.hessian
+    assert hessian is not None
+    assert hessian.shape == (3, 3)
+    np.testing.assert_array_equal(hessian, hessian.T)
+    assert np.all(np.linalg.eigvalsh(hessian) < 0)
+
+    # Analytic Hessian of the log likelihood of y ~ normal(alpha + beta * x,
+    # sigma) in (alpha, beta, log(sigma)), evaluated at the mode CmdStan used.
+    with open(data_file) as fd:
+        data = json.load(fd)
+    x, y = np.array(data['x']), np.array(data['y'])
+    mode = fit.mode.stan_variables()
+    resid = y - mode['alpha'] - mode['beta'] * x
+    expected = (
+        -np.array(
+            [
+                [len(x), x.sum(), 2 * resid.sum()],
+                [x.sum(), (x**2).sum(), 2 * (x * resid).sum()],
+                [2 * resid.sum(), 2 * (x * resid).sum(), 2 * (resid**2).sum()],
+            ]
+        )
+        / mode['sigma'] ** 2
+    )
+    np.testing.assert_allclose(hessian, expected, rtol=1e-6, atol=1e-6)
+
+    # the diagnostic file moves with the other output files
+    fit.save_output_files(os.fspath(tmp_path))
+    assert Path(fit.diagnostic_file).parent == tmp_path
+    assert fit.config_file is not None
+    rebuilt = cmdstanpy.CmdStanLaplace.from_files(
+        csv_file=fit.csv_file,
+        config_file=fit.config_file,
+        mode=fit.mode,
+        diagnostic_file=fit.diagnostic_file,
+    )
+    np.testing.assert_array_equal(rebuilt.hessian, hessian)
+
 
 def test_laplace_create_inits() -> None:
     stan = os.path.join(DATAFILES_PATH, 'bernoulli.stan')