Return 1D array when parsing diagnoal hmc mass matrix

Author: amas0

cmdstanpy/stanfit/mcmc.py

@@ -18,7 +18,6 @@
 )
 
 import numpy as np
-import numpy.typing as npt
 import pandas as pd
 
 try:
@@ -442,7 +441,7 @@ def _assemble_draws(self) -> None:
         )
         self._step_size = np.empty(self.chains, dtype=np.float64)
 
-        mass_matrix_per_chain: List[Optional[npt.NDArray[np.float64]]] = []
+        mass_matrix_per_chain = []
         for chain in range(self.chains):
             try:
                 comments, draws = stancsv.parse_stan_csv_comments_and_draws(
@@ -465,13 +464,7 @@ def _assemble_draws(self) -> None:
                 ) from exc
 
         if all(mm is not None for mm in mass_matrix_per_chain):
-            if self.metric_type == "diag_e":
-                # Mass matrix will have shape (1, num_params)
-                self._metric = np.array(
-                    [mm[0] for mm in mass_matrix_per_chain]  # type: ignore
-                )
-            else:
-                self._metric = np.array(mass_matrix_per_chain)
+            self._metric = np.array(mass_matrix_per_chain)
 
         assert self._draws is not None
 

cmdstanpy/utils/stancsv.py

@@ -98,14 +98,17 @@ def parse_hmc_adaptation_lines(
     comment_lines: List[bytes],
 ) -> Tuple[float, Optional[npt.NDArray[np.float64]]]:
     """Extracts step size/mass matrix information from the Stan CSV comment
-    lines by parsing the adaptation section. If unit metric is used, the mass
-    matrix field will be None, otherwise an appropriate numpy array.
+    lines by parsing the adaptation section. If the diag_e metric is used,
+    the returned mass matrix will be a 1D array of the diagnoal elements,
+    if the dense_e metric is used, it will be a 2D array representing the
+    entire matrix, and if unit_e is used then None will be returned.
 
     Returns a (step_size, mass_matrix) tuple"""
     step_size, mass_matrix = None, None
 
     cleaned_lines = (ln.lstrip(b"# ") for ln in comment_lines)
     in_matrix_block = False
+    diag_e_metric = False
     matrix_lines = []
     for line in cleaned_lines:
         if in_matrix_block and line.strip():
@@ -120,12 +123,16 @@ def parse_hmc_adaptation_lines(
             in_matrix_block = True
         elif line.startswith(b"No free"):
             break
+        elif b"diag_e" in line:
+            diag_e_metric = True
     if step_size is None:
         raise ValueError("Unable to parse adapated step size")
     if matrix_lines:
         mass_matrix = csv_bytes_list_to_numpy(
             matrix_lines, includes_header=False
         )
+        if diag_e_metric and mass_matrix.shape[0] == 1:
+            mass_matrix = mass_matrix[0]
     return step_size, mass_matrix
 
 

test/test_stancsv.py

@@ -212,6 +212,7 @@ def test_parsing_adaptation_lines():
 
 def test_parsing_adaptation_lines_diagonal():
     lines = [
+        b"diag_e",  # Will be present in the Stan CSV config
         b"# Adaptation terminated\n",
         b"# Step size = 0.787025\n",
         b"# Diagonal elements of inverse mass matrix:\n",
@@ -220,7 +221,7 @@ def test_parsing_adaptation_lines_diagonal():
     step_size, mass_matrix = stancsv.parse_hmc_adaptation_lines(lines)
     assert step_size == 0.787025
     assert mass_matrix is not None
-    assert np.array_equal(mass_matrix, np.array([[1, 2, 3]]))
+    assert np.array_equal(mass_matrix, np.array([1, 2, 3]))
 
 
 def test_parsing_adaptation_lines_dense():