diff --git a/cmdstanpy/model.py b/cmdstanpy/model.py index d1d96f48..c4fad9da 100644 --- a/cmdstanpy/model.py +++ b/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 + '-_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 '-_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( diff --git a/cmdstanpy/stanfit/laplace.py b/cmdstanpy/stanfit/laplace.py index 5906d0e7..991c65b7 100644 --- a/cmdstanpy/stanfit/laplace.py +++ b/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) diff --git a/cmdstanpy/stanfit/runset.py b/cmdstanpy/stanfit/runset.py index 665f8174..aecacc34 100644 --- a/cmdstanpy/stanfit/runset.py +++ b/cmdstanpy/stanfit/runset.py @@ -76,11 +76,13 @@ 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 = [ @@ -88,7 +90,9 @@ def __init__( ] 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 ] diff --git a/test/test_laplace.py b/test/test_laplace.py index 0e8e5571..71de4455 100644 --- a/test/test_laplace.py +++ b/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')