Skip to content

Commit

Permalink
stan and exe file loading refinement
Browse files Browse the repository at this point in the history
  • Loading branch information
edwinnglabs committed Mar 8, 2024
1 parent 1595e01 commit ce8b925
Show file tree
Hide file tree
Showing 3 changed files with 21 additions and 14 deletions.
2 changes: 1 addition & 1 deletion orbit/__version__.py
Original file line number Diff line number Diff line change
@@ -1 +1 @@
__version__ = "1.1.4.4.dev12"
__version__ = "1.1.4.4.dev13"
30 changes: 19 additions & 11 deletions orbit/utils/stan.py
Original file line number Diff line number Diff line change
@@ -1,12 +1,12 @@
import json
import os
import platform
from multiprocessing import cpu_count

import cmdstanpy
import importlib_resources
from cmdstanpy import CmdStanModel

from typing import Optional

from ..utils.logger import get_logger

logger = get_logger("orbit")
Expand All @@ -31,27 +31,33 @@


def get_compiled_stan_model(
stan_model_name: str = "", stan_file_path: str = ""
stan_model_name: str = "",
stan_file_path: Optional[str] = None,
exe_file_path: Optional[str] = None,
force_compile: bool = False,
) -> CmdStanModel:
"""Return a compiled Stan model using CmdStan.
This includes both prepackaged models as well as user provided models through stan_file_path.
Parameters
----------
stan_model_name : str
stan_model_name :
The name of the Stan model to use. Use this for the built in models (dlt, ets, ktrlite, lgt)
stan_file_path : str, optional
stan_file_path :
The path to the Stan file to use. If not provided, the default is to search for the file in the 'orbit' package.
If provided, function will ignore the stan_model_name parameter, and will compile the provide stan_file_path
into executable in place (same folder as stan_file_path)
exe_file_path :
The path to the Stan-exe file to use. If not provided, the default is to search for the file
in the 'orbit' package. If provided, function will ignore the stan_model_name parameter,
and will compile the provide stan_file_path into executable in place (same folder as stan_file_path)
Returns
-------
sm : CmdStanModel
A compiled Stan model.
"""
if stan_file_path != "":
stan_file = stan_file_path
sm = CmdStanModel(stan_file=stan_file)
if (stan_file_path is not None) or (exe_file_path is not None):
sm = CmdStanModel(stan_file=stan_file, exe_file=exe_file_path, force_compile=force_compile)
else:
# Load orbit included cmdstan models
# Some oddities here. if not providing exe_file, CmdStanModel would delete the actual executable file.
Expand All @@ -63,9 +69,11 @@ def get_compiled_stan_model(
)
# Check if exe is older than .stan file.
# This behavior is default on CmdStanModel if we don't have to specify the exe_file.
if not os.path.isfile(exe_file) or (
os.path.getmtime(exe_file) <= os.path.getmtime(stan_file)
):
# if not os.path.isfile(exe_file) or (
# os.path.getmtime(exe_file) <= os.path.getmtime(stan_file)
# ):

if not os.path.isfile(exe_file) or force_compile:
logger.info(f"Compiling stan model:{stan_file}. ETA 3 - 5 mins.")
sm = CmdStanModel(stan_file=stan_file)
else:
Expand Down
3 changes: 1 addition & 2 deletions setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,6 @@
import platform
import shutil
from pathlib import Path
from shutil import copy

from setuptools import find_packages, setup
from setuptools.command.build_py import build_py
Expand Down Expand Up @@ -90,7 +89,7 @@ def build_model(model: str, model_dir: str):
model_name = f"{model}.stan"
model_path = os.path.join(model_dir, model_name)
print(f"Compiling stan file in place: {model_path}")
sm = cmdstanpy.CmdStanModel(stan_file=model_path)
_ = cmdstanpy.CmdStanModel(stan_file=model_path)


def build_stan_models():
Expand Down

0 comments on commit ce8b925

Please sign in to comment.