-
Notifications
You must be signed in to change notification settings - Fork 8
/
Copy pathsetup.py
47 lines (40 loc) · 1.36 KB
/
setup.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
from distutils.core import setup
from distutils.extension import Extension
from Cython.Build import cythonize
import numpy
import os
USE_FLOAT = 0 # use float (1) or double (0)
eigen_dir = os.path.join('include', 'Eigen')
if not os.path.exists(eigen_dir):
print('Downloading Eigen...')
from glob import glob
from urllib.request import urlretrieve
import shutil
import tarfile
if not os.path.exists('include'):
os.mkdir('include')
eigen_url = 'http://bitbucket.org/eigen/eigen/get/3.3.3.tar.gz'
tar_path = os.path.join('include', 'Eigen.tar.gz')
urlretrieve(eigen_url, tar_path)
with tarfile.open(tar_path, 'r') as tar:
tar.extractall('include')
thedir = glob(os.path.join('include', 'eigen-eigen-*'))[0]
shutil.move(os.path.join(thedir, 'Eigen'), eigen_dir)
print('done!')
setup(
name = 'stochs',
ext_modules = cythonize([Extension(
'stochs',
['stochs.pyx',
'solvers/Loss.cpp',
'solvers/SGD.cpp',
'solvers/MISO.cpp',
'solvers/SAGA.cpp'],
language='c++',
include_dirs=[numpy.get_include(), 'include', 'solvers'],
extra_compile_args=['-std=c++11', '-fopenmp'],
extra_link_args=['-std=c++11', '-fopenmp', '-lglog'],
define_macros=[('USE_FLOAT', USE_FLOAT)],
)],
compile_time_env={'USE_FLOAT': USE_FLOAT})
)