📦 Fix packaging, pre-compile BimodalNonNegative

0001-Remove-dynamic-cmdstan-version-selection.patch

@@ -0,0 +1,37 @@
+From cbb7b3d1797066477aa6c1c8c694cd873bd1f8fc Mon Sep 17 00:00:00 2001
+From: =?UTF-8?q?David=20P=C3=A4tzel?= <david.paetzel@posteo.de>
+Date: Tue, 18 Oct 2022 13:33:29 +0200
+Subject: [PATCH] Remove dynamic cmdstan version selection
+
+---
+ cmdstanpy/utils/cmdstan.py | 14 +-------------
+ 1 file changed, 1 insertion(+), 13 deletions(-)
+
+diff --git a/cmdstanpy/utils/cmdstan.py b/cmdstanpy/utils/cmdstan.py
+index 273e0d2..46563e6 100644
+--- a/cmdstanpy/utils/cmdstan.py
++++ b/cmdstanpy/utils/cmdstan.py
+@@ -163,19 +163,7 @@ def cmdstan_path() -> str:
+     if 'CMDSTAN' in os.environ and len(os.environ['CMDSTAN']) > 0:
+         cmdstan = os.environ['CMDSTAN']
+     else:
+-        cmdstan_dir = os.path.expanduser(os.path.join('~', _DOT_CMDSTAN))
+-        if not os.path.exists(cmdstan_dir):
+-            raise ValueError(
+-                'No CmdStan installation found, run command "install_cmdstan"'
+-                'or (re)activate your conda environment!'
+-            )
+-        latest_cmdstan = get_latest_cmdstan(cmdstan_dir)
+-        if latest_cmdstan is None:
+-            raise ValueError(
+-                'No CmdStan installation found, run command "install_cmdstan"'
+-                'or (re)activate your conda environment!'
+-            )
+-        cmdstan = os.path.join(cmdstan_dir, latest_cmdstan)
++        cmdstan = ...
+         os.environ['CMDSTAN'] = cmdstan
+     validate_cmdstan_path(cmdstan)
+     return os.path.normpath(cmdstan)
+-- 
+2.36.1
+

README.md

@@ -128,6 +128,11 @@ pip install git+https://github.com/dpaetzel/cmpbayes
 ```
 
 
+However, this is not tested so far and since we're precompiling Stan models,
+this may not work out of the box. Please open an issue if this is the case for
+you.
+
+
 ## Other libraries to check out
 
 

flake.nix

@@ -13,6 +13,31 @@
       pkgs = import nixpkgs { inherit system; };
       python = pkgs.python39;
     in rec {
+
+      cmdstanpy = python.pkgs.buildPythonPackage rec {
+        pname = "cmdstanpy";
+        version = "1.0.7";
+
+        propagatedBuildInputs = with python.pkgs; [ numpy pandas tqdm ujson ];
+
+        patches = [
+          "${self}/0001-Remove-dynamic-cmdstan-version-selection.patch"
+        ];
+
+        postPatch = ''
+          sed -i \
+            "s|\(cmdstan = \)\.\.\.|\1\"${pkgs.cmdstan}/opt/cmdstan\"|" \
+            cmdstanpy/utils/cmdstan.py
+        '';
+
+        doCheck = false;
+
+        src = python.pkgs.fetchPypi {
+          inherit pname version;
+          sha256 = "sha256-AyzbqfVKup4pLl/JgDcoNKFi5te4QfO7KKt3pCNe4N8=";
+        };
+      };
+
       defaultPackage.${system} = python.pkgs.buildPythonPackage rec {
         pname = "cmpbayes";
         version = "1.0.0-beta";
@@ -22,17 +47,12 @@
         # We use pyproject.toml.
         format = "pyproject";
 
-        # TODO In order to provide a proper default flake here we need to
-        # package pystan/httpstan properly. For now, we assume that pystan is
-        # already there.
-        postPatch = ''
-          sed -i "s/^.*pystan.*$//" setup.cfg
-        '';
+        buildInputs = [ pkgs.cmdstan ];
 
         propagatedBuildInputs = with python.pkgs; [
-          pkgs.cmdstan
-
           arviz
+          click
+          matplotlib
           numpy
           pandas
           scipy
@@ -73,7 +93,7 @@
         # missing symbols error on NixOS. 4.7.1 works, however, so we use that.
         postVenvCreation = ''
           unset SOURCE_DATE_EPOCH
-          pip install httpstan==4.7.1 pystan==3.4.0 cmdstanpy==1.0.7
+          pip install httpstan==4.7.1 pystan==3.4.0
         '';
 
       };

pyproject.toml

@@ -2,5 +2,6 @@
 requires = [
     "setuptools>=30.3.0",
     "wheel",
+    "cmdstanpy>=1.0.7",
 ]
 build-backend = "setuptools.build_meta"

setup.cfg

@@ -37,7 +37,5 @@ extra_requires =
 where = src
 
 [options.package_data]
-cmpbayes =
-  bimodal_nonnegative.stan
-  kruschke.stan
-  pl_model.stan
+* =
+  *.stan

setup.py

@@ -0,0 +1,144 @@
+"""
+Large parts of this code was taken from [this
+setup.py](https://github.com/WardBrian/stan_py_example/blob/main/setup.py) by
+Brian Ward.
+
+This setup.py is used only for the building of the Stan models; the rest of the
+metadata is in `setup.cfg`.
+"""
+
+import os
+import platform
+from pathlib import Path
+from shutil import copy, copytree, rmtree
+from typing import Tuple
+
+import cmdstanpy
+from setuptools import Extension, setup
+from setuptools.command.build_ext import build_ext
+from wheel.bdist_wheel import bdist_wheel
+
+MODEL_DIR = "cmpbayes/stan"
+
+MODELS = [
+    "bimodal_nonnegative",
+]
+
+CMDSTAN_VERSION = "2.29.2"
+BINARIES_DIR = "bin"
+BINARIES = ["diagnose", "print", "stanc", "stansummary"]
+MATH_LIB = "stan/lib/stan_math/lib"
+TBB_DIRS = ["tbb", "tbb_2020.3"]
+
+
+def prune_cmdstan(cmdstan_dir: os.PathLike) -> None:
+    """
+    Keep only the cmdstan executables and tbb files (minimum required to run
+    cmdstanpy commands on a pre-compiled model).
+    """
+    original_dir = Path(cmdstan_dir).resolve()
+    parent_dir = original_dir.parent
+    temp_dir = parent_dir / "temp"
+    if temp_dir.is_dir():
+        rmtree(temp_dir)
+    temp_dir.mkdir()
+
+    print("Copying ", original_dir, " to ", temp_dir, " for pruning")
+    copytree(original_dir / BINARIES_DIR, temp_dir / BINARIES_DIR)
+    copy(original_dir / "makefile", temp_dir / "makefile")
+    for f in (temp_dir / BINARIES_DIR).iterdir():
+        if f.is_dir():
+            rmtree(f)
+        elif f.is_file() and f.stem not in BINARIES:
+            os.remove(f)
+    for tbb_dir in TBB_DIRS:
+        copytree(original_dir / MATH_LIB / tbb_dir,
+                 temp_dir / MATH_LIB / tbb_dir)
+
+    rmtree(original_dir)
+    temp_dir.rename(original_dir)
+
+
+def repackage_cmdstan() -> bool:
+    return os.environ.get("STAN_PY_EXAMPLE_REPACKAGE_CMDSTAN", "").lower() in [
+        "true",
+        "1",
+    ]
+
+
+def maybe_install_cmdstan_toolchain() -> None:
+    """
+    Install C++ compilers required to build stan models on Windows machines.
+    """
+
+    try:
+        cmdstanpy.utils.cxx_toolchain_path()
+    except Exception:
+        from cmdstanpy.install_cxx_toolchain import run_rtools_install
+
+        run_rtools_install({"version": None, "dir": None, "verbose": True})
+        cmdstanpy.utils.cxx_toolchain_path()
+
+
+def install_cmdstan_deps(cmdstan_dir: Path) -> None:
+    from multiprocessing import cpu_count
+
+    if repackage_cmdstan():
+        if platform.platform().startswith("Win"):
+            maybe_install_cmdstan_toolchain()
+        print("Installing cmdstan to", cmdstan_dir)
+        if os.path.isdir(cmdstan_dir):
+            print("Removing existing dir", cmdstan_dir)
+            rmtree(cmdstan_dir)
+
+        if not cmdstanpy.install_cmdstan(
+                version=CMDSTAN_VERSION,
+                dir=cmdstan_dir.parent,
+                overwrite=True,
+                verbose=True,
+                cores=cpu_count(),
+        ):
+            raise RuntimeError(
+                "CmdStan failed to install in repackaged directory")
+    else:
+        try:
+            cmdstanpy.cmdstan_path()
+        except ValueError as e:
+            raise SystemExit(
+                "CmdStan not installed, but the package is building from source"
+            ) from e
+
+
+def build_models(target_dir: str) -> None:
+    cmdstan_dir = (Path(target_dir) / f"cmdstan-{CMDSTAN_VERSION}").resolve()
+    install_cmdstan_deps(cmdstan_dir)
+    for model_name in MODELS:
+        model = cmdstanpy.CmdStanModel(
+            stan_file=os.path.join("src", MODEL_DIR, f"{model_name}.stan"),
+            stanc_options={"O1": True},
+        )
+        copy(model.exe_file, os.path.join(target_dir, f"{model_name}.exe"))
+
+    if repackage_cmdstan():
+        prune_cmdstan(cmdstan_dir)
+
+
+class BuildModels(build_ext):
+    """
+    Custom build command to pre-compile Stan models.
+    """
+
+    def run(self) -> None:
+        if not self.dry_run:
+            target_dir = os.path.join(self.build_lib, MODEL_DIR)
+            self.mkpath(target_dir)
+            build_models(target_dir)
+
+
+setup(
+    # Mark this as platform-specific.
+    ext_modules=[Extension("stan_py_example.stan", [])],
+    # Override the build command.
+    cmdclass={
+        "build_ext": BuildModels,
+    })

src/cmpbayes/__init__.py

@@ -3,40 +3,41 @@
 
 __version__ = '1.0.0-beta'
 
-import subprocess
+from pathlib import Path
 
 import arviz as az  # type: ignore
 import cmdstanpy  # type: ignore
 import matplotlib.pyplot as plt  # type: ignore
 import numpy as np  # type: ignore
 import scipy.stats as st  # type: ignore
-import stan  # type: ignore
 from pkg_resources import resource_filename  # type: ignore
 
+import stan  # type: ignore
+
 pl_stan_filename = resource_filename(__name__, "pl_model.stan")
 kruschke_filename = resource_filename(__name__, "kruschke.stan")
-bimodal_nonnegative_filename = resource_filename(__name__,
-                                                 "bimodal_nonnegative.stan")
-
-try:
-    # Note that we need what's in the opt directory under NixOS since only there
-    # is the makefile. Otherwise we get a "ValueError: Unable to compile Stan
-    # model file".
-    nixos_version = subprocess.run(
-        ['nixos-version'],
-        capture_output=True).stdout.decode().removesuffix('\n')
-    print(f"Running NixOS {nixos_version}, attempting to fix cmdstan_path.")
-    _stan_path = subprocess.run(["which", "stan"],
-                                capture_output=True).stdout.decode()
-    _cmdstan_path = _stan_path.replace("bin/stan\n", "opt/cmdstan")
-    cmdstanpy.set_cmdstan_path(_cmdstan_path)
-    print(f"Successfully set cmdstan_path to {_cmdstan_path}.")
-except FileNotFoundError:
-    print("Not running on NixOS, not applying automated NixOS fix. "
-          "If you see “ValueError: Unable to compile Stan model file”,"
-          "consider setting cmdstan path to the directory containing "
-          "cmdstan's makefile and bin directory using "
-          "cmdstanpy.set_cmdstan_path(…).")
+
+STAN_FILES_FOLDER = Path(__file__).parent / "stan"
+
+
+def _load_stan_model(name: str) -> cmdstanpy.CmdStanModel:
+    """
+    Load the precompiled Stan model via `cmdstanpy`.
+    """
+    try:
+        model = cmdstanpy.CmdStanModel(
+            exe_file=STAN_FILES_FOLDER / f"{name}.exe",
+            stan_file=STAN_FILES_FOLDER / f"{name}.stan",
+            compile=False,
+        )
+    except ValueError:
+        raise ValueError(
+            "Precompiled model could not be found at "
+            f"{STAN_FILES_FOLDER}/{name}.exe. "
+            "This is typically due to the cmpbayes library not having been "
+            "installed correctly.")
+
+    return model
 
 
 def _generate_random_seed():
@@ -419,8 +420,7 @@ def fit(self, **kwargs):
             var_upper=1.0,
         )
 
-        self.model_ = cmdstanpy.CmdStanModel(
-            stan_file=bimodal_nonnegative_filename)
+        self.model_ = _load_stan_model("bimodal_nonnegative")
 
         self.fit_ = self.model_.sample(data=data, **kwargs)