some new features

.venv/lib/python3.12/site-packages/stanio/csv.py

@@ -0,0 +1,50 @@
+"""
+Module to load the minimal information from a Stan CSV file.
+Only the header row and data are read, no metadata is parsed.
+"""
+from typing import List, Tuple, Union
+
+import numpy as np
+import numpy.typing as npt
+
+
+def read_csv(filenames: Union[str, List[str]]) -> Tuple[str, npt.NDArray[np.float64]]:
+    """
+    Reads CSV files like those produced by Stan, returning the header and data.
+
+    If multiple files are given, the data is stacked along the first axis,
+    so in typical usage, the shape of the returned data will be
+    ``(n_chains, n_samples, n_params)``.
+
+    Parameters
+    ----------
+    filenames : Union[str, List[str]]
+        Path to the CSV file(s) to read.
+
+    Returns
+    -------
+    Tuple[str, npt.NDArray[np.float64]]
+        The header row and data from the CSV file(s).
+
+    Raises
+    ------
+    ValueError
+        If multiple files are given and the headers do not match between them.
+    """
+
+    if not isinstance(filenames, list):
+        filenames = [filenames]
+
+    header = ""
+    data: List[npt.NDArray[np.float64]] = [None for _ in range(len(filenames))]  # type: ignore
+    for i, f in enumerate(filenames):
+        with open(f, "r") as fd:
+            while (file_header := fd.readline()).startswith("#"):
+                pass
+            if header == "":
+                header = file_header
+            elif header != file_header:
+                raise ValueError("Headers do not match")
+            data[i] = np.loadtxt(fd, delimiter=",", comments="#")
+
+    return header.strip(), np.stack(data, axis=0)