Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 3 additions & 6 deletions .github/workflows/main.yml
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion .github/workflows/publish.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
6 changes: 3 additions & 3 deletions pyproject.toml
Original file line number Diff line number Diff line change
@@ -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 <riddella@indiana.edu>",
Expand All @@ -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 = "*"
Expand Down
5 changes: 2 additions & 3 deletions stan/common.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
import aiohttp
import aiohttp.web
import httpstan.app
import simdjson
import orjson


def unused_tcp_port():
Expand All @@ -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:
Expand Down
59 changes: 27 additions & 32 deletions stan/fit.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
from typing import Generator, Tuple, cast

import numpy as np
import simdjson
import orjson


class Fit(collections.abc.Mapping):
Expand Down Expand Up @@ -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
Expand Down
15 changes: 6 additions & 9 deletions stan/model.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand All @@ -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("<comment>Messages received during sampling:</comment>")
Expand Down
Loading