Merge pull request #840 from amas0/add-compilation-kwargs

Add multithreading and stanc_optimization arguments to CmdStanModel

Author: amas0

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

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
 

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"
    ]
   },
   {

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,28 +76,30 @@ 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
 
     import os
     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

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"}