<commit_before>import asyncio
import dataclasses
import json
import re
from typing import Dict, List, Optional, Sequence, Tuple, Union

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
from clikit.ui.components import ProgressBar

import stan.common
import stan.fit
import stan.plugins

Data = Dict[str, Union[int, float, Sequence[Union[int, float]], np.ndarray]]


class DataJSONEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, np.ndarray):
            return obj.tolist()
        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.")


@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`.")

    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)

    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.

        """
        kwargs["num_chains"] = num_chains
        kwargs["function"] = "stan::services::sample::hmc_nuts_diag_e_adapt"
        return self._create_fit(kwargs)

    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.

        """
        kwargs["num_chains"] = num_chains
        kwargs["function"] = "stan::services::sample::fixed_param"
        return self._create_fit(kwargs)

    def _create_fit(self, payload: dict) -> stan.fit.Fit:
        """Make a request to httpstan's ``create_fit`` endpoint and process results.

        Users should not use this function.

        Arguments:
            payload: dict whose JSON-encoded contents will be sent as the request body.

        Returns:
            Fit: instance of Fit allowing access to draws.

        """
        assert "chain" not in payload, "`chain` id is set automatically."
        assert "data" not in payload, "`data` is set in `build`."
        assert "random_seed" not in payload, "`random_seed` is set in `build`."
        assert "function" in payload

        # copy kwargs and verify everything is JSON-encodable
        payload = json.loads(DataJSONEncoder().encode(payload))
        num_chains = payload.pop("num_chains")

        init: List[Data]
        init = payload.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["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 bar 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.copy())

        async def go():
            io = ConsoleIO()
            io.error_line("<info>Sampling...</info>")
            progress_bar = ProgressBar(io)
            progress_bar.set_format("very_verbose")

            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:<commit_msg>refactor: Improve clarity of `_create_fit` code<commit_after>import asyncio
import dataclasses
import json
import re
from typing import Dict, List, Optional, Sequence, Tuple, Union

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
from clikit.ui.components import ProgressBar

import stan.common
import stan.fit
import stan.plugins

Data = Dict[str, Union[int, float, Sequence[Union[int, float]], np.ndarray]]


class DataJSONEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, np.ndarray):
            return obj.tolist()
        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.")


@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`.")

    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)

    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)

    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 bar 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()
            io.error_line("<info>Sampling...</info>")
            progress_bar = ProgressBar(io)
            progress_bar.set_format("very_verbose")

            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:
{'fraction_non_alphanumeric': 0.09873400763614441, 'fraction_numerical': 0.002344430303436265, 'mean_word_length': 2.625546381738708, 'pattern_counts': {'":': 4, '<': 7, '<?xml version=': 0, '>': 33, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1593097', 'n_tokens_mistral': 4198, 'n_tokens_neox': 3946, 'n_words': 1287}