mirror of
https://github.com/saymrwulf/prophet.git
synced 2026-09-15 22:10:22 +00:00
Allow to build models in-place. (#100)
This fixes the commands ``python setup.py develop`` and ``pip install -e``.
This commit is contained in:
parent
1a6b736cca
commit
1a57d19148
3 changed files with 43 additions and 15 deletions
3
.gitignore
vendored
3
.gitignore
vendored
|
|
@ -1,6 +1,9 @@
|
||||||
# Compiled python modules.
|
# Compiled python modules.
|
||||||
*.pyc
|
*.pyc
|
||||||
|
|
||||||
|
# Persisted models.
|
||||||
|
*.pkl
|
||||||
|
|
||||||
# Setuptools distribution folder.
|
# Setuptools distribution folder.
|
||||||
python/dist/
|
python/dist/
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -1,3 +1,6 @@
|
||||||
include stan/unix/*.stan
|
include stan/unix/*.stan
|
||||||
include stan/win/*.stan
|
include stan/win/*.stan
|
||||||
include LICENSE
|
include LICENSE
|
||||||
|
|
||||||
|
# Ensure in-place built models do not get included in the source dist.
|
||||||
|
prune fbprophet/stan_models
|
||||||
|
|
|
||||||
|
|
@ -11,33 +11,54 @@ from pkg_resources import (
|
||||||
)
|
)
|
||||||
from setuptools import setup
|
from setuptools import setup
|
||||||
from setuptools.command.build_py import build_py
|
from setuptools.command.build_py import build_py
|
||||||
|
from setuptools.command.develop import develop
|
||||||
from setuptools.command.test import test as test_command
|
from setuptools.command.test import test as test_command
|
||||||
|
|
||||||
|
|
||||||
|
PLATFORM = 'unix'
|
||||||
|
if platform.platform().startswith('Win'):
|
||||||
|
PLATFORM = 'win'
|
||||||
|
|
||||||
|
SETUP_DIR = os.path.dirname(os.path.abspath(__file__))
|
||||||
|
MODELS_DIR = os.path.join(SETUP_DIR, 'stan', PLATFORM)
|
||||||
|
MODELS_TARGET_DIR = os.path.join('fbprophet', 'stan_models')
|
||||||
|
|
||||||
|
|
||||||
|
def build_stan_models(target_dir, models_dir=MODELS_DIR):
|
||||||
|
from pystan import StanModel
|
||||||
|
for model_type in ['linear', 'logistic']:
|
||||||
|
model_name = 'prophet_{}_growth.stan'.format(model_type)
|
||||||
|
target_name = '{}_growth.pkl'.format(model_type)
|
||||||
|
with open(os.path.join(models_dir, model_name)) as f:
|
||||||
|
model_code = f.read()
|
||||||
|
sm = StanModel(model_code=model_code)
|
||||||
|
with open(os.path.join(target_dir, target_name), 'wb') as f:
|
||||||
|
pickle.dump(sm, f, protocol=pickle.HIGHEST_PROTOCOL)
|
||||||
|
|
||||||
|
|
||||||
class BuildPyCommand(build_py):
|
class BuildPyCommand(build_py):
|
||||||
"""Custom build command to pre-compile Stan models."""
|
"""Custom build command to pre-compile Stan models."""
|
||||||
|
|
||||||
def run(self):
|
def run(self):
|
||||||
if not self.dry_run:
|
if not self.dry_run:
|
||||||
self.build_stan_models()
|
target_dir = os.path.join(self.build_lib, MODELS_TARGET_DIR)
|
||||||
|
self.mkpath(target_dir)
|
||||||
|
build_stan_models(target_dir)
|
||||||
|
|
||||||
build_py.run(self)
|
build_py.run(self)
|
||||||
|
|
||||||
def build_stan_models(self):
|
|
||||||
from pystan import StanModel
|
|
||||||
target_dir = os.path.join(self.build_lib, 'fbprophet/stan_models')
|
|
||||||
self.mkpath(target_dir)
|
|
||||||
|
|
||||||
if platform.platform().startswith('Win'):
|
class DevelopCommand(develop):
|
||||||
plat = 'win'
|
"""Custom develop command to pre-compile Stan models in-place."""
|
||||||
else:
|
|
||||||
plat = 'unix'
|
def run(self):
|
||||||
|
if not self.dry_run:
|
||||||
|
target_dir = os.path.join(self.setup_path, MODELS_TARGET_DIR)
|
||||||
|
self.mkpath(target_dir)
|
||||||
|
build_stan_models(target_dir)
|
||||||
|
|
||||||
|
develop.run(self)
|
||||||
|
|
||||||
for model_type in ['linear', 'logistic']:
|
|
||||||
with open('stan/{}/prophet_{}_growth.stan'.format(plat, model_type)) as f:
|
|
||||||
model_code = f.read()
|
|
||||||
sm = StanModel(model_code=model_code)
|
|
||||||
with open(os.path.join(target_dir, '{}_growth.pkl'.format(model_type)), 'wb') as f:
|
|
||||||
pickle.dump(sm, f, protocol=pickle.HIGHEST_PROTOCOL)
|
|
||||||
|
|
||||||
class TestCommand(test_command):
|
class TestCommand(test_command):
|
||||||
"""We must run tests on the build directory, not source."""
|
"""We must run tests on the build directory, not source."""
|
||||||
|
|
@ -96,6 +117,7 @@ setup(
|
||||||
include_package_data=True,
|
include_package_data=True,
|
||||||
cmdclass={
|
cmdclass={
|
||||||
'build_py': BuildPyCommand,
|
'build_py': BuildPyCommand,
|
||||||
|
'develop': DevelopCommand,
|
||||||
'test': TestCommand,
|
'test': TestCommand,
|
||||||
},
|
},
|
||||||
test_suite='fbprophet.tests.test_prophet',
|
test_suite='fbprophet.tests.test_prophet',
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue