Remove unused chain_id from MetricInfo

Author: amas0

cmdstanpy/stanfit/mcmc.py

@@ -400,12 +400,13 @@ def _validate_csv_files(self) -> dict[str, Any]:
     def _parse_metric_info(self) -> None:
         """Extracts metric type, inv_metric, and step size information from the
         parsed metric JSONs."""
-        self._chain_metric_info = [
-            MetricInfo.from_json(mf, chain_id)
-            for mf, chain_id in zip(
-                self.runset.metric_files, self.runset.chain_ids
-            )
-        ]
+        self._chain_metric_info = []
+        for mf in self.runset.metric_files:
+            with open(mf) as f:
+                self._chain_metric_info.append(
+                    MetricInfo.model_validate_json(f.read())
+                )
+
         metric_types = {cmi.metric_type for cmi in self._chain_metric_info}
         if len(metric_types) != 1:
             raise ValueError("Inconsistent metric types found across chains")

cmdstanpy/stanfit/metadata.py

@@ -3,14 +3,13 @@
 from __future__ import annotations
 
 import copy
-import json
 import math
 import os
 from typing import Any, Iterator, Literal
 
 import numpy as np
 import stanio
-from pydantic import BaseModel, Field, field_validator, model_validator
+from pydantic import BaseModel, field_validator, model_validator
 
 from cmdstanpy.utils import stancsv
 
@@ -91,7 +90,6 @@ class MetricInfo(BaseModel):
     """Structured representation of HMC-NUTS metric information,
     as output by CmdStan"""
 
-    chain_id: int = Field(gt=0)
     stepsize: float
     metric_type: Literal["diag_e", "dense_e", "unit_e"]
     inv_metric: np.ndarray
@@ -126,12 +124,3 @@ def validate_inv_metric_shape(self) -> MetricInfo:
                 raise ValueError("Dense inv_metric must be square")
 
         return self
-
-    @classmethod
-    def from_json(cls, file: str | os.PathLike, chain_id: int) -> MetricInfo:
-        """Parse and validate a metric json given a file path and chain_id"""
-        with open(file) as f:
-            info_dict = json.load(f)
-
-        info_dict['chain_id'] = chain_id
-        return cls.model_validate(info_dict)  # type: ignore

test/test_metadata.py

@@ -1,4 +1,4 @@
-"""Metadata tests"""
+"Metadata tests"
 
 import json
 import math
@@ -84,20 +84,17 @@ class TestMetricInfoValidators:
     def test_valid_diag_e_metric(self) -> None:
         """Test valid diag_e metric with 1D array"""
         metric = MetricInfo(
-            chain_id=1,
             stepsize=0.5,
             metric_type="diag_e",
             inv_metric=[1.0, 2.0, 3.0],  # type: ignore
         )
-        assert metric.chain_id == 1
         assert metric.stepsize == 0.5
         assert isinstance(metric.inv_metric, np.ndarray)
         assert metric.inv_metric.ndim == 1
 
     def test_valid_unit_e_metric(self) -> None:
         """Test valid unit_e metric with 1D array"""
         metric = MetricInfo(
-            chain_id=2,
             stepsize=0.1,
             metric_type="unit_e",
             inv_metric=[1.0, 1.0, 1.0],  # type: ignore
@@ -108,7 +105,6 @@ def test_valid_unit_e_metric(self) -> None:
     def test_valid_dense_e_metric(self) -> None:
         """Test valid dense_e metric with 2D square array"""
         metric = MetricInfo(
-            chain_id=1,
             stepsize=0.3,
             metric_type="dense_e",
             inv_metric=[[1.0, 0.5], [0.5, 1.0]],  # type: ignore
@@ -120,7 +116,6 @@ def test_valid_dense_e_metric(self) -> None:
     def test_convert_inv_metric_from_list(self) -> None:
         """Test that inv_metric is converted to numpy array from list"""
         metric = MetricInfo(
-            chain_id=1,
             stepsize=0.5,
             metric_type="diag_e",
             inv_metric=[1.0, 2.0, 3.0],  # type: ignore
@@ -130,7 +125,6 @@ def test_convert_inv_metric_from_list(self) -> None:
     def test_convert_inv_metric_from_nested_list(self) -> None:
         """Test that inv_metric is converted to numpy array from nested list"""
         metric = MetricInfo(
-            chain_id=1,
             stepsize=0.5,
             metric_type="dense_e",
             inv_metric=[[1.0, 0.0], [0.0, 1.0]],  # type: ignore
@@ -140,7 +134,6 @@ def test_convert_inv_metric_from_nested_list(self) -> None:
     def test_stepsize_positive(self) -> None:
         """Test valid positive stepsize"""
         metric = MetricInfo(
-            chain_id=1,
             stepsize=0.5,
             metric_type="diag_e",
             inv_metric=[1.0],  # type: ignore
@@ -150,7 +143,6 @@ def test_stepsize_positive(self) -> None:
     def test_stepsize_nan_allowed(self) -> None:
         """Test that NaN stepsize is allowed"""
         metric = MetricInfo(
-            chain_id=1,
             stepsize=math.nan,
             metric_type="diag_e",
             inv_metric=[1.0],  # type: ignore
@@ -161,7 +153,6 @@ def test_stepsize_zero_raises_error(self) -> None:
         """Test that zero stepsize raises ValueError"""
         with pytest.raises(ValidationError) as exc_info:
             MetricInfo(
-                chain_id=1,
                 stepsize=0.0,
                 metric_type="diag_e",
                 inv_metric=[1.0],  # type: ignore
@@ -172,7 +163,6 @@ def test_stepsize_negative_raises_error(self) -> None:
         """Test that negative stepsize raises ValueError"""
         with pytest.raises(ValidationError) as exc_info:
             MetricInfo(
-                chain_id=1,
                 stepsize=-0.5,
                 metric_type="diag_e",
                 inv_metric=[1.0],  # type: ignore
@@ -183,7 +173,6 @@ def test_diag_e_with_2d_array_raises_error(self) -> None:
         """Test that diag_e with 2D array raises ValueError"""
         with pytest.raises(ValidationError) as exc_info:
             MetricInfo(
-                chain_id=1,
                 stepsize=0.5,
                 metric_type="diag_e",
                 inv_metric=[[1.0, 2.0]],  # type: ignore
@@ -196,7 +185,6 @@ def test_unit_e_with_2d_array_raises_error(self) -> None:
         """Test that unit_e with 2D array raises ValueError"""
         with pytest.raises(ValidationError) as exc_info:
             MetricInfo(
-                chain_id=1,
                 stepsize=0.5,
                 metric_type="unit_e",
                 inv_metric=[[1.0], [1.0]],  # type: ignore
@@ -209,7 +197,6 @@ def test_dense_e_with_1d_array_raises_error(self) -> None:
         """Test that dense_e with 1D array raises ValueError"""
         with pytest.raises(ValidationError) as exc_info:
             MetricInfo(
-                chain_id=1,
                 stepsize=0.5,
                 metric_type="dense_e",
                 inv_metric=[1.0, 2.0],  # type: ignore
@@ -220,27 +207,15 @@ def test_dense_e_non_square_raises_error(self) -> None:
         """Test that dense_e with non-square array raises ValueError"""
         with pytest.raises(ValidationError) as exc_info:
             MetricInfo(
-                chain_id=1,
                 stepsize=0.5,
                 metric_type="dense_e",
                 inv_metric=[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]],  # type: ignore
             )
         assert "Dense inv_metric must be square" in str(exc_info.value)
 
-    def test_chain_id_must_be_positive(self) -> None:
-        """Test that chain_id must be greater than 0"""
-        with pytest.raises(ValidationError) as exc_info:
-            MetricInfo(
-                chain_id=0,
-                stepsize=0.5,
-                metric_type="diag_e",
-                inv_metric=[1.0],  # type: ignore
-            )
-        assert "greater than 0" in str(exc_info.value)
-
 
-class TestMetricInfoFromJson:
-    """Test from_json class method"""
+class TestMetricInfoModelValidateJson:
+    """Test model_validate_json class method"""
 
     def test_from_json_diag_e(self) -> None:
         """Test loading diag_e metric from JSON file"""
@@ -258,8 +233,8 @@ def test_from_json_diag_e(self) -> None:
             temp_path = f.name
 
         try:
-            metric = MetricInfo.from_json(temp_path, chain_id=1)
-            assert metric.chain_id == 1
+            with open(temp_path) as f:
+                metric = MetricInfo.model_validate_json(f.read())
             assert metric.stepsize == 0.5
             assert metric.metric_type == "diag_e"
             assert np.array_equal(metric.inv_metric, np.array([1.0, 2.0, 3.0]))
@@ -282,8 +257,8 @@ def test_from_json_dense_e(self) -> None:
             temp_path = f.name
 
         try:
-            metric = MetricInfo.from_json(temp_path, chain_id=2)
-            assert metric.chain_id == 2
+            with open(temp_path) as f:
+                metric = MetricInfo.model_validate_json(f.read())
             assert metric.stepsize == 0.3
             assert metric.metric_type == "dense_e"
             assert metric.inv_metric.shape == (2, 2)
@@ -307,7 +282,8 @@ def test_from_json_invalid_data_raises_error(self) -> None:
 
         try:
             with pytest.raises(ValidationError):
-                MetricInfo.from_json(temp_path, chain_id=1)
+                with open(temp_path) as f:
+                    MetricInfo.model_validate_json(f.read())
         finally:
             Path(temp_path).unlink()
 
@@ -327,8 +303,8 @@ def test_from_json_pathlike(self) -> None:
             temp_path = Path(f.name)
 
         try:
-            metric = MetricInfo.from_json(temp_path, chain_id=3)
-            assert metric.chain_id == 3
+            with open(temp_path) as f:
+                metric = MetricInfo.model_validate_json(f.read())
             assert metric.metric_type == "unit_e"
         finally:
             temp_path.unlink()