diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 6917dd1..1e957f4 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -6,14 +6,11 @@ jobs: tests: name: pystan tests timeout-minutes: 30 - runs-on: ${{ matrix.runs-on }} + runs-on: ${{ matrix.os }} strategy: matrix: - include: - - {runs-on: ubuntu-24.04, python-version: "3.12"} - - {runs-on: ubuntu-24.04, python-version: "3.13"} - - {runs-on: ubuntu-24.04, python-version: "3.14"} - - {runs-on: macos-15-intel, python-version: "3.12"} + os: [ubuntu-24.04, macos-14] + python-version: ["3.12", "3.14"] steps: - name: Check out repository uses: actions/checkout@v4 diff --git a/.github/workflows/publish.yaml b/.github/workflows/publish.yaml index dfb9304..4eb098e 100644 --- a/.github/workflows/publish.yaml +++ b/.github/workflows/publish.yaml @@ -11,7 +11,7 @@ jobs: - name: Set up Python uses: actions/setup-python@v5 with: - python-version: "3.11" + python-version: "3.12" - name: Install dependencies run: | python -m pip install --upgrade pip diff --git a/pyproject.toml b/pyproject.toml index 28c9cba..6f4f190 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "pystan" -version = "3.10.1" +version = "3.14.0" description = "Python interface to Stan, a package for Bayesian inference" authors = [ "Allen Riddell ", @@ -24,8 +24,8 @@ classifiers = [ [tool.poetry.dependencies] python = "^3.12" aiohttp = "^3.6" -httpstan = "~4.13" -pysimdjson = ">=5.0.2" +httpstan = "~4.17" +orjson = ">=3.10" numpy = ">=1.19" clikit = "^0.6" setuptools = "*" diff --git a/stan/common.py b/stan/common.py index 60000ab..4d78852 100644 --- a/stan/common.py +++ b/stan/common.py @@ -6,7 +6,7 @@ import aiohttp import aiohttp.web import httpstan.app -import simdjson +import orjson def unused_tcp_port(): @@ -22,8 +22,7 @@ class HTTPResponse(typing.NamedTuple): content: bytes def json(self) -> dict: - # mypy 0.961 complains that simdjson lacks a `loads`. - return simdjson.loads(self.content) # type: ignore + return orjson.loads(self.content) class HttpstanClient: diff --git a/stan/fit.py b/stan/fit.py index 978a484..4e6bd66 100644 --- a/stan/fit.py +++ b/stan/fit.py @@ -4,7 +4,7 @@ from typing import Generator, Tuple, cast import numpy as np -import simdjson +import orjson class Fit(collections.abc.Mapping): @@ -65,43 +65,38 @@ def __init__( # _draws is an ndarray with shape (num_sample_and_sampler_params + num_flat_params, num_draws, num_chains) self._draws: np.ndarray - parser = simdjson.Parser() for chain_index, stan_output in zip(range(self.num_chains), self.stan_outputs): draw_index = 0 for line in stan_output.splitlines(): try: - msg = cast(simdjson.Object, parser.parse(line)) - except ValueError: - # Occurs when draws contain an nan or infinity. simdjson cannot parse such values. + msg = orjson.loads(line) + except orjson.JSONDecodeError: + # Occurs when draws contain a NaN or infinity. orjson cannot parse such values. msg = json.loads(line) - try: - if msg["topic"] == "sample": - # Ignore sample message which is mixed together with proper draws. - if not isinstance(msg["values"], (simdjson.Object, dict)): - continue - - # for the first draw: collect sample and sampler parameter names. - if not hasattr(self, "_draws"): - feature_names = cast(Tuple[str, ...], tuple(msg["values"].keys())) - self.sample_and_sampler_param_names = tuple( - name for name in feature_names if name.endswith("__") + if msg["topic"] == "sample": + # Ignore sample message which is mixed together with proper draws. + if not isinstance(msg["values"], dict): + continue + + # for the first draw: collect sample and sampler parameter names. + if not hasattr(self, "_draws"): + feature_names = cast(Tuple[str, ...], tuple(msg["values"].keys())) + self.sample_and_sampler_param_names = tuple( + name for name in feature_names if name.endswith("__") + ) + num_rows = len(self.sample_and_sampler_param_names) + num_flat_params + # column-major order ("F") aligns with how the draws are stored (in cols). + self._draws = np.empty((num_rows, num_samples_saved, num_chains), order="F") + # rudimentary check of parameter order (sample & sampler params must be first) + if num_flat_params and feature_names[-1].endswith("__"): + raise RuntimeError( + f"Expected last parameter name to be one declared in program code, found `{feature_names[-1]}`" ) - num_rows = len(self.sample_and_sampler_param_names) + num_flat_params - # column-major order ("F") aligns with how the draws are stored (in cols). - self._draws = np.empty((num_rows, num_samples_saved, num_chains), order="F") - # rudimentary check of parameter order (sample & sampler params must be first) - if num_flat_params and feature_names[-1].endswith("__"): - raise RuntimeError( - f"Expected last parameter name to be one declared in program code, found `{feature_names[-1]}`" - ) - - draw_row = tuple(msg["values"].values()) # a "row" of values from a single draw from Stan C++ - draw_row = cast(Tuple[float, ...], draw_row) - self._draws[:, draw_index, chain_index] = draw_row - draw_index += 1 - finally: - # clean up `Object`s produced by parser, required by simdjson - del msg + + draw_row = tuple(msg["values"].values()) # a "row" of values from a single draw from Stan C++ + draw_row = cast(Tuple[float, ...], draw_row) + self._draws[:, draw_index, chain_index] = draw_row + draw_index += 1 assert draw_index == num_samples_saved assert self.sample_and_sampler_param_names and self._draws.size self._draws.flags["WRITEABLE"] = False # type: ignore diff --git a/stan/model.py b/stan/model.py index e1db175..fb00468 100644 --- a/stan/model.py +++ b/stan/model.py @@ -10,7 +10,7 @@ import httpstan.services.arguments as arguments import httpstan.utils import numpy as np -import simdjson +import orjson from clikit.io import ConsoleIO import stan.common @@ -257,10 +257,10 @@ async def go(): stan_outputs = tuple(stan_outputs) # Fit constructor expects a tuple. - def is_nonempty_logger_message(msg: simdjson.Object): + def is_nonempty_logger_message(msg: dict): return msg["topic"] == "logger" and msg["values"][0] != "info:" # type: ignore - def is_iteration_or_elapsed_time_logger_message(msg: simdjson.Object): + def is_iteration_or_elapsed_time_logger_message(msg: dict): # Assumes `msg` is a message with topic `logger`. text = msg["values"][0] # type: ignore text = cast(str, text) @@ -271,19 +271,16 @@ def is_iteration_or_elapsed_time_logger_message(msg: simdjson.Object): or text.startswith("info:" + " " * 15) ) - parser = simdjson.Parser() nonstandard_logger_messages = [] for stan_output in stan_outputs: for line in stan_output.splitlines(): # Do not attempt to parse non-logger messages. Draws could contain nan or inf values. - # simdjson cannot parse lines containing such values. + # orjson cannot parse lines containing such values. if b'"logger"' not in line: continue - msg = parser.parse(line) + msg = orjson.loads(line) if is_nonempty_logger_message(msg) and not is_iteration_or_elapsed_time_logger_message(msg): - nonstandard_logger_messages.append(msg.as_dict()) - del msg - del parser # simdjson.Parser is no longer used at this point. + nonstandard_logger_messages.append(msg) if nonstandard_logger_messages: io.error_line("Messages received during sampling:")