diff --git a/cmdstanpy/compilation.py b/cmdstanpy/compilation.py index 54e099fe..afc75812 100644 --- a/cmdstanpy/compilation.py +++ b/cmdstanpy/compilation.py @@ -486,3 +486,23 @@ def format_stan_file( except (ValueError, RuntimeError) as e: raise RuntimeError("Stanc formatting failed") from e + + +def resolve_cpp_options( + cpp_options: dict[str, Any] | None, multithreading: bool +) -> dict[str, Any]: + out = cpp_options or {} + out = out.copy() + if multithreading and "STAN_THREADS" not in out: + out["STAN_THREADS"] = "TRUE" + return out + + +def resolve_stanc_options( + stanc_options: dict[str, Any] | None, stanc_optimizations: bool +) -> dict[str, Any]: + out = stanc_options or {} + out = out.copy() + if stanc_optimizations and "O" not in out: + out["O"] = 1 + return out diff --git a/cmdstanpy/model.py b/cmdstanpy/model.py index 5cea19f2..ece9fd2e 100644 --- a/cmdstanpy/model.py +++ b/cmdstanpy/model.py @@ -93,6 +93,9 @@ def __init__( stanc_options: dict[str, Any] | None = None, cpp_options: dict[str, Any] | None = None, user_header: OptionalPath = None, + *, + multithreading: bool = False, + stanc_optimizations: bool = False, ) -> None: """ Initialize object given constructor args. @@ -101,14 +104,27 @@ def __init__( :param exe_file: Path to compiled executable file. :param force_compile: Whether or not to force recompilation if executable file already exists. - :param stanc_options: Options for stanc compiler. - :param cpp_options: Options for C++ compiler. + :param multithreading: Enables multithreading in a Stan model. + Equivalent to `cpp_options = {"STAN_THREADS": "TRUE"}`. + Defaults to False. + :param stanc_optimizations: Enables O1 optimizations in the + stanc compiler. Equivalent to `stanc_options = {"O": 1}`. + Defaults to False. + :param stanc_options: Options for stanc compiler. Note, this + will override the `stanc_optimizations` if in conflict. + :param cpp_options: Options for C++ compiler. Note, this will + override the `multithreading` option if in conflict. :param user_header: A path to a header file to include during C++ compilation. """ self._name = '' self._stan_file = None - self._stanc_options: dict[str, Any] = stanc_options or {} + self._stanc_options = compilation.resolve_stanc_options( + stanc_options, stanc_optimizations + ) + cpp_options = compilation.resolve_cpp_options( + cpp_options, multithreading + ) self._fixed_param = False diff --git a/docsrc/users-guide/examples/Pathfinder.ipynb b/docsrc/users-guide/examples/Pathfinder.ipynb index e5f3c564..4a0ed002 100644 --- a/docsrc/users-guide/examples/Pathfinder.ipynb +++ b/docsrc/users-guide/examples/Pathfinder.ipynb @@ -44,7 +44,7 @@ "outputs": [], "source": [ "import os\n", - "from cmdstanpy.model import CmdStanModel, cmdstan_path" + "from cmdstanpy import CmdStanModel, cmdstan_path" ] }, { diff --git a/docsrc/users-guide/workflow.rst b/docsrc/users-guide/workflow.rst index 82bafc46..453301a1 100644 --- a/docsrc/users-guide/workflow.rst +++ b/docsrc/users-guide/workflow.rst @@ -39,21 +39,14 @@ managing the resulting inference for a single model and set of inputs. Compile the Stan model ^^^^^^^^^^^^^^^^^^^^^^ -The: :class:`CmdStanModel` class provides methods -to compile and run the Stan program. -A CmdStanModel object can be instantiated by specifying -either a Stan file or the executable file, or both. -If only the Stan file path is specified, the constructor will -check for the existence of a correspondingly named exe file in -the same directory. If found, it will use this as the exe file path. - -By default, when a CmdStanModel object is instantiated from a Stan file, -the constructor will compile the model as needed. -The constructor argument `compile` controls this behavior. - -* ``compile=False``: never compile the Stan file. -* ``compile="Force"``: always compile the Stan file. -* ``compile=True``: (default) compile the Stan file as needed, i.e., if no exe file exists or if the Stan file is newer than the exe file. +The :class:`CmdStanModel` class provides methods to compile and run the Stan +program. A CmdStanModel object can be instantiated by specifying a Stan file, +the executable file, or both. If only the Stan file path is specified, the +constructor will check for the existence of a correspondingly named executable in +the same directory. If found, it will use this as the exe file path. + +When a CmdStanModel object is instantiated from a Stan file, the constructor +will compile the model if the executable is non-existent or out-of-date. .. code-block:: python @@ -67,8 +60,8 @@ The constructor argument `compile` controls this behavior. my_model.exe_file my_model.code() -The CmdStanModel class also provides the :meth:`~CmdStanModel.compile` method, -which can be called at any point to (re)compile the model as needed. +The ``force_compile=True`` argument can be passed to the CmdStanModel +constructor, which will force (re)compilation of the model. Model compilation is carried out via the GNU Make build tool. The CmdStan ``makefile`` contains a set of general rules which @@ -83,20 +76,22 @@ Model compilation is done in two steps: * The C++ compiler compiles the generated code and links in the necessary supporting libraries. -Therefore, both the constructor and the ``compile`` method -allow optional arguments ``stanc_options`` and ``cpp_options`` which -specify options for each compilation step. -Options are specified as a Python dictionary mapping -compiler option names to appropriate values. +The constructor accepts arguments to specify both ``stanc`` and C++ compilation +options, if desired. Passing `multithreading=True` enables the **STAN_THREADS** +C++ flag, which is needed to parallelize within-chain computations, such as +with ``reduce_sum``, or to parallelize the NUTS-HMC sampler across chains. +Passing ``stanc_optimizations=True`` will enable ``O1`` optimizations in the +``stanc`` compiler. -In order parallelize within-chain computations using the -Stan language ``reduce_sum`` function, or to parallelize -running the NUTS-HMC sampler across chains, -the Stan model must be compiled with -C++ compiler flag **STAN_THREADS**. -While any value can be used, -we recommend the value ``True``, e.g.: +Outside of these common options, the constructor accepts the optional arguments +``stanc_options`` and ``cpp_options``, which allow specifying arbitrary +compilation options. Some more advanced Stan features, like MPI or OpenCL +support, require using these. Note that if the lower-level compilation options +conflict with an argument like ``multithreading=True``, the option in +``stanc_options`` or ``cpp_options`` takes precedence. +An example model compilation that enables multithreading and +basic optimization can be done like so: .. code-block:: python @@ -104,7 +99,7 @@ we recommend the value ``True``, e.g.: from cmdstanpy import CmdStanModel my_stanfile = os.path.join('.', 'my_model.stan') - my_model = CmdStanModel(stan_file=my_stanfile, cpp_options={'STAN_THREADS':'true'}) + my_model = CmdStanModel(stan_file=my_stanfile, multithreading=True, stanc_optimizations=True) Assemble input and initialization data diff --git a/test/test_compilation.py b/test/test_compilation.py index 731a5c8b..addaea55 100644 --- a/test/test_compilation.py +++ b/test/test_compilation.py @@ -9,7 +9,12 @@ import pytest -from cmdstanpy.compilation import CompilerOptions, format_stan_file +from cmdstanpy.compilation import ( + CompilerOptions, + format_stan_file, + resolve_cpp_options, + resolve_stanc_options, +) HERE = os.path.dirname(os.path.abspath(__file__)) DATAFILES_PATH = os.path.join(HERE, 'data') @@ -225,3 +230,23 @@ def test_model_format_options() -> None: formatted = sys_stdout.getvalue() assert formatted.count('{') == 3 assert formatted.count('(') == 1 + + +def test_compilation_options_resolution() -> None: + out = resolve_cpp_options(None, multithreading=False) + assert not out + out = resolve_cpp_options(None, multithreading=True) + assert out == {"STAN_THREADS": "TRUE"} + out = resolve_cpp_options({"STAN_THREADS": ""}, multithreading=True) + assert out == {"STAN_THREADS": ""} + out = resolve_cpp_options({"STAN_OPENCL": "TRUE"}, multithreading=True) + assert out == {"STAN_THREADS": "TRUE", "STAN_OPENCL": "TRUE"} + + out = resolve_stanc_options(None, stanc_optimizations=False) + assert not out + out = resolve_stanc_options(None, stanc_optimizations=True) + assert out == {"O": 1} + out = resolve_stanc_options({"O": 0}, stanc_optimizations=True) + assert out == {"O": 0} + out = resolve_stanc_options({"O": "experimental"}, stanc_optimizations=True) + assert out == {"O": "experimental"}