import asyncio
import dataclasses
import json
import re
import time
from typing import Dict, List, Optional, Sequence, Tuple, Union, cast
import httpstan.models
import httpstan.schemas
import httpstan.services.arguments as arguments
import httpstan.utils
import numpy as np
import simdjson
from clikit.io import ConsoleIO
import stan.common
import stan.fit
import stan.plugins
Data = Dict[str, Union[int, float, Sequence[Union[int, float]]]]
class DataJSONEncoder(json.JSONEncoder):
def default(self, obj):
# numpy.ndarray is *unofficially* supported as there is no easy way to
# construct tabular data using the Python standard library.
if isinstance(obj, np.ndarray):
return obj.tolist()
# unofficially support np.int64, np.int32, etc.
if hasattr(obj, "dtype") and np.issubdtype(obj.dtype, np.integer):
return int(obj)
return json.JSONEncoder.default(self, obj)
# superficial frozendict implementation. Only used for function signatures
class frozendict(dict):
def __setitem__(self, key, value):
raise TypeError("'frozendict' object is immutable.")
[docs]
@dataclasses.dataclass(frozen=True)
class Model:
"""Stores data associated with and proxies calls to a Stan model.
Returned by ``build``. Users will not instantiate this class directly.
"""
model_name: str
program_code: str
data: Data
param_names: Tuple[str, ...]
constrained_param_names: Tuple[str, ...]
dims: Tuple[Tuple[int, ...]]
random_seed: Optional[int]
def __post_init__(self):
if self.model_name != httpstan.models.calculate_model_name(self.program_code):
raise ValueError("`model_name` does not match `program_code`.")
[docs]
def sample(self, *, num_chains=4, **kwargs) -> stan.fit.Fit:
"""Draw samples from the model.
Parameters in ``kwargs`` will be passed to the default sample function.
The default sample function is currently
``stan::services::sample::hmc_nuts_diag_e_adapt``. Parameter names are
identical to those used in CmdStan. See the CmdStan documentation for
parameter descriptions and default values.
There is one exception: `num_chains`. `num_chains` is a
PyStan-specific keyword argument. It indicates the number of
independent processes to use when drawing samples.
Returns:
Fit: instance of Fit allowing access to draws.
Examples:
User-defined initial values for parameters must be provided
for each chain. Typically they will be the same for each chain.
The following example shows how user-defined initial parameters
are provided:
>>> program_code = "parameters {real y;} model {y ~ normal(0,1);}"
>>> posterior = stan.build(program_code)
>>> fit = posterior.sample(num_chains=2, init=[{"y": 3}, {"y": 3}])
"""
return self.hmc_nuts_diag_e_adapt(num_chains=num_chains, **kwargs)
[docs]
def hmc_nuts_diag_e_adapt(self, *, num_chains=4, **kwargs) -> stan.fit.Fit:
"""Draw samples from the model using ``stan::services::sample::hmc_nuts_diag_e_adapt``.
Parameters in ``kwargs`` will be passed to the (Python wrapper of)
``stan::services::sample::hmc_nuts_diag_e_adapt``. Parameter names are
identical to those used in CmdStan. See the CmdStan documentation for
parameter descriptions and default values.
There is one exception: `num_chains`. `num_chains` is a
PyStan-specific keyword argument. It indicates the number of
independent processes to use when drawing samples.
Returns:
Fit: instance of Fit allowing access to draws.
"""
function = "stan::services::sample::hmc_nuts_diag_e_adapt"
return self._create_fit(function=function, num_chains=num_chains, **kwargs)
[docs]
def fixed_param(self, *, num_chains=4, **kwargs) -> stan.fit.Fit:
"""Draw samples from the model using ``stan::services::sample::fixed_param``.
Parameters in ``kwargs`` will be passed to the (Python wrapper of)
``stan::services::sample::fixed_param``. Parameter names are
identical to those used in CmdStan. See the CmdStan documentation for
parameter descriptions and default values.
There is one exception: `num_chains`. `num_chains` is a
PyStan-specific keyword argument. It indicates the number of
independent processes to use when drawing samples.
Returns:
Fit: instance of Fit allowing access to draws.
"""
function = "stan::services::sample::fixed_param"
return self._create_fit(function=function, num_chains=num_chains, **kwargs)
def _create_fit(self, *, function, num_chains, **kwargs) -> stan.fit.Fit:
"""Make a request to httpstan's ``create_fit`` endpoint and process results.
Users should not use this function.
Parameters in ``kwargs`` will be passed to the (Python wrapper of)
`function`. Parameter names are identical to those used in CmdStan.
See the CmdStan documentation for parameter descriptions and default
values.
Returns:
Fit: instance of Fit allowing access to draws.
"""
assert "chain" not in kwargs, "`chain` id is set automatically."
assert "data" not in kwargs, "`data` is set in `build`."
assert "random_seed" not in kwargs, "`random_seed` is set in `build`."
# copy kwargs and verify everything is JSON-encodable
kwargs = json.loads(DataJSONEncoder().encode(kwargs))
# FIXME: special handling here for `init`, consistent with PyStan 2 but needs docs
init: List[Data] = kwargs.pop("init", [dict() for _ in range(num_chains)])
if len(init) != num_chains:
raise ValueError("Initial values must be provided for each chain.")
payloads = []
for chain in range(1, num_chains + 1):
payload = kwargs.copy()
payload["function"] = function
payload["chain"] = chain # type: ignore
payload["data"] = self.data # type: ignore
payload["init"] = init.pop(0)
if self.random_seed is not None:
payload["random_seed"] = self.random_seed # type: ignore
# fit needs to know num_samples, num_warmup, num_thin, save_warmup
# progress reporting needs to know some of these
num_warmup = payload.get("num_warmup", arguments.lookup_default(arguments.Method["SAMPLE"], "num_warmup"))
num_samples = payload.get(
"num_samples",
arguments.lookup_default(arguments.Method["SAMPLE"], "num_samples"),
)
num_thin = payload.get("num_thin", arguments.lookup_default(arguments.Method["SAMPLE"], "num_thin"))
save_warmup = payload.get(
"save_warmup",
arguments.lookup_default(arguments.Method["SAMPLE"], "save_warmup"),
)
payloads.append(payload)
async def go():
io = ConsoleIO()
sampling_output = io.section().error_output
percent_complete = 0
sampling_output.write_line(f"<comment>Sampling:</comment> {percent_complete:3.0f}%")
current_and_max_iterations_re = re.compile(r"Iteration:\s+(\d+)\s+/\s+(\d+)")
async with stan.common.HttpstanClient() as client:
operations = []
for payload in payloads:
resp = await client.post(f"/{self.model_name}/fits", json=payload)
if resp.status == 422:
raise ValueError(str(resp.json()))
elif resp.status != 201:
raise RuntimeError(resp.json()["message"])
assert resp.status == 201
operations.append(resp.json())
# poll to get progress for each chain until all chains finished
current_iterations = {}
while not all(operation["done"] for operation in operations):
for operation in operations:
if operation["done"]:
continue
resp = await client.get(f"/{operation['name']}")
assert resp.status != 404
operation.update(resp.json())
progress_message = operation["metadata"].get("progress")
if not progress_message:
continue
iteration, iteration_max = map(
int, current_and_max_iterations_re.findall(progress_message).pop(0)
)
if current_iterations.get(operation["name"]) == iteration:
continue
current_iterations[operation["name"]] = iteration
iterations_count = sum(current_iterations.values())
total_iterations = iteration_max * num_chains
percent_complete = 100 * iterations_count / total_iterations
sampling_output.clear() if io.supports_ansi() else sampling_output.write("\n")
sampling_output.write_line(
f"<comment>Sampling:</comment> {round(percent_complete):3.0f}% ({iterations_count}/{total_iterations})"
)
await asyncio.sleep(0.05)
fit_in_cache = len(current_iterations) < num_chains
stan_outputs = []
for operation in operations:
fit_name = operation["result"].get("name")
if fit_name is None: # operation["result"] is an error
assert not str(operation["result"]["code"]).startswith("2"), operation
message = operation["result"]["message"]
if """ValueError('Initialization failed.')""" in message:
sampling_output.clear()
sampling_output.write_line("<info>Sampling:</info> <error>Initialization failed.</error>")
raise RuntimeError("Initialization failed.")
raise RuntimeError(message)
resp = await client.get(f"/{fit_name}")
if resp.status != 200:
raise RuntimeError((resp.json())["message"])
stan_outputs.append(resp.content)
# clean up after ourselves when fit is uncacheable (no random seed)
if self.random_seed is None:
resp = await client.delete(f"/{fit_name}")
if resp.status not in {200, 202, 204}:
raise RuntimeError((resp.json())["message"])
sampling_output.clear() if io.supports_ansi() else sampling_output.write("\n")
sampling_output.write_line(
"<info>Sampling:</info> 100%, done."
if fit_in_cache
else f"<info>Sampling:</info> {percent_complete:3.0f}% ({iterations_count}/{total_iterations}), done."
)
if not io.supports_ansi():
sampling_output.write("\n")
stan_outputs = tuple(stan_outputs) # Fit constructor expects a tuple.
def is_nonempty_logger_message(msg: simdjson.Object):
return msg["topic"] == "logger" and msg["values"][0] != "info:" # type: ignore
def is_iteration_or_elapsed_time_logger_message(msg: simdjson.Object):
# Assumes `msg` is a message with topic `logger`.
text = msg["values"][0] # type: ignore
text = cast(str, text)
return (
text.startswith("info:Iteration:")
or text.startswith("info: Elapsed Time:")
# this detects lines following "Elapsed Time:", part of a multi-line Stan message
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.
if b'"logger"' not in line:
continue
msg = parser.parse(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.
if nonstandard_logger_messages:
io.error_line("<comment>Messages received during sampling:</comment>")
for msg in nonstandard_logger_messages:
text = msg["values"][0].replace("info:", " ").replace("error:", " ")
if text.strip():
io.error_line(f"{text}")
fit = stan.fit.Fit(
stan_outputs,
num_chains,
self.param_names,
self.constrained_param_names,
self.dims,
num_warmup,
num_samples,
num_thin,
save_warmup,
)
for entry_point in stan.plugins.get_plugins():
Plugin = entry_point.load()
fit = Plugin().on_post_sample(fit)
return fit
try:
return asyncio.run(go())
except KeyboardInterrupt:
return # type: ignore
[docs]
def constrain_pars(
self, unconstrained_parameters: Sequence[float], include_tparams: bool = True, include_gqs: bool = True
) -> Sequence[float]:
"""Transform a sequence of unconstrained parameters to their defined support,
optionally including transformed parameters and generated quantities.
Arguments:
unconstrained_parameters: A sequence of unconstrained parameters.
include_tparams: Boolean to control whether we include transformed parameters.
include_gqs: Boolean to control whether we include generated quantities.
Returns:
A sequence of constrained parameters, optionally including transformed parameters.
Note:
The unconstrained parameters are passed to the `write_array` method of the `model_base`
instance. See `model_base.hpp` in the Stan C++ library for details.
"""
payload = {
"data": self.data,
"unconstrained_parameters": unconstrained_parameters,
"include_tparams": include_tparams,
"include_gqs": include_gqs,
}
async def go():
async with stan.common.HttpstanClient() as client:
resp = await client.post(f"/{self.model_name}/write_array", json=payload)
if resp.status != 200:
raise RuntimeError(resp.json())
return resp.json()["params_r_constrained"]
return asyncio.run(go())
[docs]
def unconstrain_pars(self, constrained_parameters: Sequence[float]) -> Sequence[float]:
"""Reads constrained parameter values from their specified context and returns a
sequence of unconstrained parameter values.
Arguments:
constrained_parameters: Constrained parameter values and their specified context
Returns:
A sequence of unconstrained parameters.
Note:
The unconstrained parameters are passed to the `transform_inits` method of the
`model_base` instance. See `model_base.hpp` in the Stan C++ library for details.
"""
payload = {"data": self.data, "constrained_parameters": constrained_parameters}
async def go():
async with stan.common.HttpstanClient() as client:
resp = await client.post(f"/{self.model_name}/transform_inits", json=payload)
if resp.status != 200:
raise RuntimeError(resp.json())
return resp.json()["params_r_unconstrained"]
return asyncio.run(go())
[docs]
def log_prob(self, unconstrained_parameters: Sequence[float], adjust_transform: bool = True) -> float:
"""Calculate the log probability of a set of unconstrained parameters.
Arguments:
unconstrained_parameters: A sequence of unconstrained parameters.
adjust_transform: Apply jacobian adjust transform.
Returns:
The log probability of the unconstrained parameters.
Notes:
The unconstrained parameters are passed to the log_prob
function in stan::model.
"""
payload = {
"data": self.data,
"unconstrained_parameters": unconstrained_parameters,
"adjust_transform": adjust_transform,
}
async def go():
async with stan.common.HttpstanClient() as client:
resp = await client.post(f"/{self.model_name}/log_prob", json=payload)
if resp.status != 200:
raise RuntimeError(resp.json())
return resp.json()["log_prob"]
return asyncio.run(go())
[docs]
def grad_log_prob(self, unconstrained_parameters: Sequence[float]) -> float:
"""Calculate the gradient of the log posterior evaluated at
the unconstrained parameters.
Arguments:
unconstrained_parameters: A sequence of unconstrained parameters.
Returns:
The gradient of the log posterior evaluated at the
unconstrained parameters.
Notes:
The unconstrained parameters are passed to the log_prob_grad
function in stan::model.
"""
payload = {
"data": self.data,
"unconstrained_parameters": unconstrained_parameters,
}
async def go():
async with stan.common.HttpstanClient() as client:
resp = await client.post(f"/{self.model_name}/log_prob_grad", json=payload)
if resp.status != 200:
raise RuntimeError(resp.json())
return resp.json()["log_prob_grad"]
return asyncio.run(go())
[docs]
def build(program_code: str, data: Data = frozendict(), random_seed: Optional[int] = None) -> Model:
"""Build (compile) a Stan program.
Arguments:
program_code: Stan program code describing a Stan model.
data: A Python dictionary or mapping providing the data for the
model. Variable names are the keys and the values are their
associated values. Default is an empty dictionary, suitable
for Stan programs with no `data` block.
random_seed: Random seed, a positive integer for random number
generation. Used to make sure that results can be reproduced.
Returns:
Model: an instance of Model
Notes:
C++ reserved words and Stan reserved words may not be used for
variable names; see the Stan User's Guide for a complete list.
"""
# `data` must be JSON-serializable in order to send to httpstan
data = json.loads(DataJSONEncoder().encode(data))
async def go():
io = ConsoleIO()
# hack: use stdout instead of stderr because httpstan silences stderr during compilation
building_output = io.section().output
if not io.supports_ansi():
building_output.write("<comment>Building...</comment>")
async with stan.common.HttpstanClient() as client:
# Check to see if model is in cache.
model_name = httpstan.models.calculate_model_name(program_code)
resp = await client.post(f"/{model_name}/params", json={"data": data})
model_in_cache = resp.status != 404
task = asyncio.create_task(client.post("/models", json={"program_code": program_code}))
start = time.time()
while True:
done, pending = await asyncio.wait({task}, timeout=0.1)
if done:
break
if io.supports_ansi():
building_output.clear()
building_output.write(f"<comment>Building:</comment> {time.time() - start:0.1f}s")
building_output.clear() if io.supports_ansi() else building_output.write("\n")
# now that httpstan has released stderr, we can use error_output again
building_output = io.section().error_output
resp = task.result()
if resp.status != 201:
match = re.search(r"""ValueError\(['"](.*)['"]\)""", resp.json()["message"])
if not match: # unknown error, should not happen
raise RuntimeError(resp.json()["message"])
exception_body = match.group(1).encode().decode("unicode_escape")
error_type_match = re.match(r"(Semantic|Syntax) error", exception_body)
if error_type_match:
error_type = error_type_match.group(0)
exception_body_without_first_line = exception_body.split("\n", 1)[1]
building_output.write_line(f"<info>Building:</info> <error>{error_type}:</error>")
building_output.write_line(f"<error>{exception_body_without_first_line}</error>")
raise ValueError(error_type)
else:
raise RuntimeError(exception_body)
building_output.clear() if io.supports_ansi() else building_output.write("\n")
if model_in_cache:
building_output.write("<info>Building:</info> found in cache, done.")
else:
building_output.write(f"<info>Building:</info> {time.time() - start:0.1f}s, done.")
assert model_name == resp.json()["name"]
if resp.json().get("stanc_warnings"):
io.error_line("<comment>Messages from <fg=cyan;options=bold>stanc</>:</comment>")
io.error_line(resp.json()["stanc_warnings"])
resp = await client.post(f"/{model_name}/params", json={"data": data})
if resp.status != 200:
raise RuntimeError(resp.json()["message"])
params_list = resp.json()["params"]
assert len({param["name"] for param in params_list}) == len(params_list)
param_names, dims = zip(*((param["name"], param["dims"]) for param in params_list))
constrained_param_names = sum((tuple(param["constrained_names"]) for param in params_list), ())
return Model(model_name, program_code, data, param_names, constrained_param_names, dims, random_seed)
try:
return asyncio.run(go())
except KeyboardInterrupt:
return # type: ignore