AssertionError: expected 'TITLE="my title"\nimport argparse\nim…' to deeply equal 'TITLE="my title"\nimport argparse\nim…' - Expected + Received TITLE="my title" import argparse import json import os import cmdstanpy HERE = os.path.dirname(os.path.abspath(__file__)) argparser = argparse.ArgumentParser(prog=f"Stan-Playground: {TITLE}") argparser.add_argument("--install-cmdstan", action="store_true", help="Install cmdstan if it is missing") args, _ = argparser.parse_known_args() + data = "" _option_names_map = { "init_radius": "inits", "num_warmup": "iter_warmup", "num_samples": "iter_sampling", "num_chains": "chains", } if os.path.isfile(os.path.join(HERE, "sampling_opts.json")): print("loading sampling_opts.json") with open(os.path.join(HERE, "sampling_opts.json")) as f: s = json.load(f) sampling_opts = {_option_names_map.get(k, k): v for k, v in s.items()} else: sampling_opts = {} try: cmdstanpy.cmdstan_path() except Exception: if args.install_cmdstan: cmdstanpy.install_cmdstan() else: raise ValueError("cmdstan not found, use --install-cmdstan to install") print("compiling model") model = cmdstanpy.CmdStanModel(stan_file=os.path.join(HERE, "main.stan")) print("sampling") fit = model.sample(data=data, **sampling_opts) print(fit.summary()) # Used in pyodideWorker for running analysis.py from typing import TYPE_CHECKING, List, TypedDict import numpy as np import pandas as pd import stanio # We don't import this unconditionaly because # we only install it when the user's script needs it if TYPE_CHECKING: from arviz import InferenceData class SpData(TypedDict): draws: List[List[float]] paramNames: List[str] numChains: int class DrawsObject: def __init__(self, sp_data: SpData): self._all_parameter_names: List[str] = sp_data["paramNames"] self._params = stanio.parse_header(",".join(self._all_parameter_names)) self._num_chains: int = sp_data["numChains"] # draws come in as num_params by (num_chains * num_draws) self._draws = ( np.array(sp_data["draws"]) .transpose() .reshape(self._num_chains, -1, len(self._all_parameter_names)) ) def __repr__(self) -> str: return f"""SpDraws with {self._num_chains} chains, {self._draws.shape[1]} draws, and {self._draws.shape[2]} parameters. Methods: - as_dataframe(): return a pandas DataFrame of the draws. - as_numpy(): return a numpy array indexed by (chain, draw, parameter) - as_arviz(): return an arviz InferenceData object - get(pname: str): return a numpy array of the parameter values for the given parameter name""" def as_dataframe(self) -> pd.DataFrame: # The first column is the chain id # The second column is the draw number # The remaining columns are the parameter values (num_chains, num_draws, num_params) = self._draws.shape flattened = self._draws.reshape(-1, num_params) chain_ids = np.repeat(np.arange(1, num_chains + 1), num_draws) draw_numbers = np.tile(np.arange(1, num_draws + 1), num_chains) data = np.column_stack((chain_ids, draw_numbers, flattened)) df = pd.DataFrame(data, columns=["chain", "draw"] + self._all_parameter_names) return df def as_numpy(self) -> np.ndarray: return np.array(self._draws) def get(self, pname: str) -> np.ndarray: if pname not in self._params: raise ValueError(f"Parameter {pname} not found") return self._params[pname].extract_reshape(self._draws) def to_arviz(self) -> "InferenceData": import arviz as az return az.from_dict( posterior={pname: self.get(pname) for pname in self.parameter_names}, ) @Property def parameter_names(self) -> List[str]: