Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
# These names are gathered later in this function by inspecting the output from Stan.
self.sample_and_sampler_param_names: Sequence[str]
num_flat_params = sum(np.product(dims_ or 1) for dims_ in dims) # if dims == [] then it is a scalar
assert num_flat_params == len(constrained_param_names)
num_samples_saved = (self.num_samples + self.num_warmup * self.save_warmup) // self.num_thin
# self._draws holds all the draws. We cannot allocate it before looking at the draws
# because we do not know how many sampler-specific parameters are present. Later in this
# function we count them and only then allocate the array for `self._draws`.
self._draws: np.ndarray
for chain_index, stan_output in zip(range(self.num_chains), self.stan_outputs):
draw_index = 0
for msg in stan_output:
if msg.topic == callbacks_writer_pb2.WriterMessage.Topic.Value("SAMPLE"):
# Ignore sample message which is mixed together with proper draws.
if msg.feature and msg.feature[0].name == "":
continue
draw_row = [] # a "row" of values from a single draw from Stan C++
# for the first draw: collect sample and sampler parameter names.
if not hasattr(self, "_draws"):
feature_names = tuple(fea.name for fea in msg.feature)
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)