diff --git a/bin/jet b/bin/jet index e7b54ba..08c441d 100644 --- a/bin/jet +++ b/bin/jet @@ -1,10 +1,10 @@ #!/usr/bin/python import argparse import sys, os, shutil -import pkgutil +import imp from jet.version import __version__ -jet_path = pkgutil.get_loader("jet").filename +jet_path = imp.find_module("jet")[1] description = \ """ diff --git a/jet/__init__.py b/jet/__init__.py index ddc1b62..e432796 100644 --- a/jet/__init__.py +++ b/jet/__init__.py @@ -1,5 +1,6 @@ from sys import modules as _modules import types as _types +import inspect as _inspect import numpy as _numpy from jet import config jet_mode = config.jet_mode @@ -13,19 +14,21 @@ module = _modules[__name__] # decorate numpy attributes -for name, attr in _numpy.__dict__.iteritems(): +for name, attr in _numpy.__dict__.items(): if not name.startswith("_") and not isinstance(attr, _types.BuiltinFunctionType): if callable(attr): - module.__dict__[name] = _helpers.numpy_mode(attr) + module.__dict__[name] = _helpers.numpy_method(attr) else: module.__dict__[name] = attr # decorate jet-intake attributes -for name, attr in _intake.__dict__.iteritems(): - if not name.startswith("_") and \ - (callable(attr) or hasattr(attr, '__class__')) and \ - not isinstance(attr, _types.BuiltinFunctionType): - module.__dict__[name] = _helpers.jet_mode(attr) +for name, attr in _intake.__dict__.items(): + if not name.startswith("_"): + if _inspect.isclass(attr): + setattr(module, name, attr) + elif callable(attr) and \ + not isinstance(attr, _types.BuiltinFunctionType): + module.__dict__[name] = _helpers.jet_method(attr) def set_options(jet_mode=True, debug=False, @@ -36,8 +39,8 @@ def set_options(jet_mode=True, group_class=False, group_func=False, DTYPE=_numpy.float64): - - jet_mode = jet_mode + + setattr(module, 'jet_mode', jet_mode) config.jet_mode = jet_mode config.debug = debug config.draw_graph = draw_graph diff --git a/jet/expander.py b/jet/expander.py index 3b9c234..2c5365c 100644 --- a/jet/expander.py +++ b/jet/expander.py @@ -563,7 +563,7 @@ def find_node(n): def check_type(*args): arg_list = list(args) num_types = [float, int, bool, np.ndarray, np.float64, np.float32] - jet_types = [intake.variable, intake.constant, intake.placeholder, intake.array] + jet_types = [intake.variable, intake.constant, intake.placeholder, intake.ndarray] for i, arg in enumerate(args): if arg is not None and type(arg) not in jet_types: diff --git a/jet/helpers.py b/jet/helpers.py index 233fcd9..37a36b7 100644 --- a/jet/helpers.py +++ b/jet/helpers.py @@ -1,28 +1,48 @@ import numpy -from jet import jet_mode -import jet.intake as intake +import jet def jet_error(attr_name): raise NotImplementedError('\'{}\' is a special member of ' 'JET which is not part of Numpy\'s API.'.format(attr_name)) - def numpy_error(attr_name): raise NotImplementedError('Attribute \'{}\' is available in Numpy but not ' 'in JET yet.'.format(attr_name)) -def jet_mode(attr): - if jet_mode: - return attr - attr_name = attr.__name__ - np_obj = getattr(numpy, attr_name, None) - if np_obj is not None: - return np_obj - return lambda *args, **kwargs: jet_error(attr_name) +def jet_method(attr): + def wrapper(*args, **kwargs): + if jet.jet_mode: + return attr(*args, **kwargs) + attr_name = attr.__name__ + np_obj = getattr(numpy, attr_name, None) + if np_obj is not None: + return np_obj(*args, **kwargs) + jet_error(attr_name) + return wrapper + +def numpy_method(np_obj): + def wrapper(*args, **kwargs): + if not jet.jet_mode: + return np_obj(*args, **kwargs) + attr_name = np_obj.__name__ + numpy_error(attr_name) + return wrapper + +def jet_class_method(cls): + def decorator(attr): + def wrapper(*args, **kwargs): + if jet.jet_mode: + return attr(*args, **kwargs) + attr_name = attr.__name__ + np_obj = getattr(getattr(numpy, cls, None), attr_name, None) + if np_obj is not None: + return np_obj(*args, **kwargs) + jet_error(attr_name) + return wrapper + return decorator -def numpy_mode(np_obj): - if not jet_mode: - return np_obj - attr_name = np_obj.__name__ - attr = getattr(intake, attr_name, None) - return lambda *args, **kwargs: numpy_error(attr_name) +def jet_static_class(cls): + for name, attr in cls.__dict__.items(): + if callable(attr): + setattr(cls, name, staticmethod(jet_class_method(cls.__name__)(attr))) + return cls diff --git a/jet/include/supportlib.h b/jet/include/supportlib.h index abdca1c..b00de0b 100644 --- a/jet/include/supportlib.h +++ b/jet/include/supportlib.h @@ -68,7 +68,7 @@ void set_items(T& lhs, const Q& rhs, const std::array& start, } } -template -T mod(const T& x, const float m) { +template +T mod(const T& x, const S& m) { return std::fmod(x, m); -} \ No newline at end of file +} diff --git a/jet/intake.py b/jet/intake.py index 306f095..82c8a4d 100644 --- a/jet/intake.py +++ b/jet/intake.py @@ -2,14 +2,15 @@ from jet import config from jet import utils from jet import expander +from jet import helpers ########################################################################## #### General Functions #### ########################################################################## -def clip(a, a_min, a_max): - op = expander.ClipOp([a, a_min, a_max]) +def clip(array, a_min, a_max): + op = expander.ClipOp([array, a_min, a_max]) return op.get_output() def where(condition, x, y): @@ -174,7 +175,6 @@ def transpose(array): def ravel(array): return array.ravel() - def reshape(array, shape): return array.reshape(shape) @@ -214,27 +214,28 @@ def logical_not(x): #### Other Functions #### ########################################################################## +@helpers.jet_static_class class linalg(object): - def solve(self, lhs, rhs): + def solve( lhs, rhs): op = expander.SolveOp([lhs, rhs]) return op.get_output() - def norm(self, x, order=2): + def norm(x, order=2): op = expander.NormOp([x], order) return op.get_output() -linalg = linalg() +@helpers.jet_static_class class random(object): - def normal(self, mean=0, sd=1): + def normal(mean=0, sd=1): op = expander.RandomNormalOp([mean, sd]) return op.get_output() -random = random() ########################################################################## #### Array Objects #### ########################################################################## -class array(object): # TODO dominique: PEP8 requires CamelCase for class names +# not compatible with numpy's ndarray +class ndarray(object): # TODO dominique: PEP8 requires CamelCase for class names def __init__(self, value=None, name='array', shape=(), dtype=config.DTYPE, producer=None, **kwargs): if value is None: @@ -440,7 +441,10 @@ def __abs__(self): def __array_wrap__(self, result): return constant(result) -class variable(array): # TODO dominique: PEP8 requires CamelCase for class names +def array(*args, **kwargs): + return ndarray(*args, **kwargs) + +class variable(ndarray): # TODO dominique: PEP8 requires CamelCase for class names # variable results in member variable in generated class def __init__(self, value=None, name='variable', shape=(), dtype=config.DTYPE): _check_not_jet_type(value, name) @@ -452,7 +456,7 @@ def __init__(self, value=None, name='variable', shape=(), dtype=config.DTYPE): op = expander.VariableOp([self]) self.producer = op -class placeholder(array): # TODO dominique: PEP8 requires CamelCase for class names +class placeholder(ndarray): # TODO dominique: PEP8 requires CamelCase for class names # placeholder results in argument to be passed in generated function def __init__(self, name='placeholder', shape=(), dtype=config.DTYPE): super(placeholder, self).__init__(name=name, shape=shape, dtype=dtype) @@ -466,7 +470,7 @@ def __repr__(self): return ''.format(self.dtype.name, self.shape, self.name) -class constant(array): # TODO dominique: PEP8 requires CamelCase for class names +class constant(ndarray): # TODO dominique: PEP8 requires CamelCase for class names # constant results in constant expression in generated class def __init__(self, value, name='constant', dtype=config.DTYPE): _check_not_jet_type(value, name) @@ -487,7 +491,7 @@ def __repr__(self): ########################################################################## def _check_not_jet_type(obj, name): - if isinstance(obj, array) or (isinstance(obj, numpy.ndarray) and + if isinstance(obj, ndarray) or (isinstance(obj, numpy.ndarray) and obj.dtype == object): raise ValueError('Can\'t add this object to jet {}.'.format(name)) diff --git a/jet/jit.py b/jet/jit.py index d76f8b2..1028838 100644 --- a/jet/jit.py +++ b/jet/jit.py @@ -1,3 +1,4 @@ +import inspect from jet.compressor import JetBuilder from jet.utils import sanitize_name, get_caller_info from jet.intake import placeholder @@ -22,7 +23,14 @@ def wrapper(*args): shapes = _func_cached_dict[func_id]['shapes'] - arg_names = func.__code__.co_varnames[0:func.__code__.co_argcount] + if inspect.ismethod(func): + arg_names = func.__code__.co_varnames[1:func.__code__.co_argcount] + else: + arg_names = func.__code__.co_varnames[:func.__code__.co_argcount] + + if len(arg_names) != len(args): + assert(len(arg_names) == 0) + arg_names = [get_unique_name('ph') for each in args] if len(shapes) != len(arg_names) and shapes: raise ValueError('Shapes length does not match the arguments length.') @@ -40,7 +48,7 @@ def wrapper(*args): file_name=get_unique_name(sanitize_name('{}_{}_{func_name}'.format( *get_caller_info('jit.py')[1:-1], func_name=fun_name))), - fun_name=fun_name) + fun_name=get_unique_name(fun_name)) jet_class = getattr(jb.build(), jb.class_name) jet_func = getattr(jet_class(), jb.fun_name) diff --git a/setup.py b/setup.py index 79302cf..6b69e63 100644 --- a/setup.py +++ b/setup.py @@ -57,7 +57,7 @@ def package_files(directory): ], install_requires=[ - 'numpy', + 'numpy; python_version < "3"', 'networkx'], packages=['jet'], scripts=['bin/jet'], diff --git a/tests/0d b/tests/0d new file mode 100644 index 0000000..93fb9ea --- /dev/null +++ b/tests/0d @@ -0,0 +1,19 @@ +maximum +minimum +logical_not +logical_or +logical_and +logical_xor +sin +cos +tan +arcsin +arccos +arctan +arctan2 +mod +power +sqrt +sign +log +exp diff --git a/tests/1d b/tests/1d new file mode 100644 index 0000000..7df9347 --- /dev/null +++ b/tests/1d @@ -0,0 +1,13 @@ +cross +dot +multiply +matmul +divide +true_divide +reciprocal +clip +ravel +transpose +concatenate +vstack +hstack diff --git a/tests/2d b/tests/2d new file mode 100644 index 0000000..5f010a1 --- /dev/null +++ b/tests/2d @@ -0,0 +1,25 @@ +where +max +min +amax +amin +fabs +negative +add +subtract +multiply +matmul +dot +divide +true_divide +reciprocal +clip +eye +zeros +ones +reshape +ravel +transpose +concatenate +vstack +hstack diff --git a/tests/helpers.py b/tests/helpers.py new file mode 100644 index 0000000..43a86b0 --- /dev/null +++ b/tests/helpers.py @@ -0,0 +1,15 @@ +class bcolors: + HEADER = '\033[95m' + OKBLUE = '\033[94m' + OKGREEN = '\033[92m' + WARNING = '\033[93m' + FAIL = '\033[91m' + ENDC = '\033[0m' + BOLD = '\033[1m' + UNDERLINE = '\033[4m' + +def cprint(text, success, metadata=''): + if success: + print(bcolors.OKGREEN + text + bcolors.ENDC) + else: + print(bcolors.FAIL + text + bcolors.ENDC + (metadata and " " + str(metadata))) diff --git a/tests/mix0d b/tests/mix0d new file mode 100644 index 0000000..754adf5 --- /dev/null +++ b/tests/mix0d @@ -0,0 +1,8 @@ +add +subtract +power +divide +true_divide +multiply +dot +mod diff --git a/tests/mix1d b/tests/mix1d new file mode 100644 index 0000000..1a5a824 --- /dev/null +++ b/tests/mix1d @@ -0,0 +1,9 @@ +matmul +add +subtract +power +divide +true_divide +multiply +dot +mod diff --git a/tests/run_test.py b/tests/run_test.py new file mode 100644 index 0000000..0c6b954 --- /dev/null +++ b/tests/run_test.py @@ -0,0 +1,180 @@ +import jet +import sys +from random import randint +from inspect import getargspec +from helpers import cprint +from jet.jit import jit +import numpy as np + + +ERROR_MAX = 1.e-12 + +mode = 0 if 'continue' in sys.argv else 1 + +dimensions = [] +for f_name in ['0d', '1d', '2d']: + with open(f_name) as file: + dimensions_raw = file.readlines() + dimensions.append([line.strip() for line in dimensions_raw if line.strip()]) + +mix = [] +for f_name in ['mix0d', 'mix1d']: + with open(f_name) as file: + mix_raw = file.readlines() + mix.append([line.strip() for line in mix_raw if line.strip()]) + +def equal(x, y, member, tol=ERROR_MAX): + if np.any(np.isnan(x)): + x = np.nan_to_num(x) + y = np.nan_to_num(y) + if mode: + assert(np.all(abs(x - y) < tol)) + else: + try: + cprint(member, np.all(abs(x - y) < tol), (x, y)) + except: + cprint(member, False, (x, y)) + +def rand_mat(shape): + return 2.0 * np.random.rand(*shape) - 1.0 + +def compare(dim, member): + jet_func = jet.__dict__[member] + numpy_func = np.__dict__[member] + args = getargspec(jet.intake.__dict__[member]).args + if dim == 0: + shape = () + elif dim == 1: + shape = (3,) + elif dim == 2: + shape = (4, 4) + + mat_1 = rand_mat(shape) + mat_2 = rand_mat(shape) + + try: + if args in [['array'], ['x']]: + equal(numpy_func(mat_1), + jit(jet_func)(mat_1), member) + elif args == ['x', 'y']: + equal(numpy_func(mat_1, mat_2), + jit(jet_func)(mat_1, mat_2), member) + elif 'input_tuple' in args: + equal(numpy_func((mat_1, mat_2)), + jit(lambda x, y: jet_func((x, y)))(mat_1, mat_2), member) + elif 'size' in args: + equal(numpy_func(shape), + jit(lambda: jet_func(shape))(), member) + elif member == 'eye': + equal(numpy_func(shape[0]), + jit(lambda: jet_func(shape[0]))(), member) + elif member == 'where': + equal(numpy_func(False, mat_1, mat_2), + jit(lambda x, y: jet_func(False, x, y))(mat_1, mat_2), member + ', False') + equal(numpy_func(True, mat_1, mat_2), + jit(lambda x, y: jet_func(True, x, y))(mat_1, mat_2), member + ', True') + elif member == 'clip': + equal(numpy_func(mat_1, -0.5, 0.5), + jit(lambda x: jet_func(x, -0.5, 0.5))(mat_1), member) + elif member == 'reshape': + equal(numpy_func(mat_1, (shape[0] / 2, shape[1] * 2)), + jit(lambda x: jet_func(x, (shape[0] / 2, shape[1] * 2)))(mat_1), member + ', 2D') + equal(numpy_func(mat_1, (shape[0] ** 2,)), + jit(lambda x: jet_func(x, (shape[0] ** 2,)))(mat_1), member + ', 1D') + else: + assert(False) + except KeyboardInterrupt: + sys.exit(0) + except: + e = sys.exc_info()[0] + if mode: + raise(e) + else: + cprint(member, False, e) + +def test_0D(): + print("-------") + print("0D") + print("-------") + + for member in dimensions[0]: + compare(0, member) + +def test_1D(): + print("-------") + print("1D") + print("-------") + + for member in dimensions[1]: + compare(1, member) + +def test_2D(): + print("-------") + print("2D") + print("-------") + + for member in dimensions[2]: + compare(2, member) + +def test_mix(): + for dim, members in enumerate(mix): + print("-------") + print("2D x {dim}D".format(dim=dim)) + print("-------") + + mat = rand_mat((4, 4)) + vec = rand_mat((4,)) + scalar = rand_mat(()) + + for member in members: + jet_func = jet.__dict__[member] + numpy_func = np.__dict__[member] + try: + if dim == 0: + equal(numpy_func(mat, scalar), + jit(jet_func)(mat, scalar), member) + else: + equal(numpy_func(mat, vec), + jit(jet_func)(mat, vec), member) + except SystemExit: + cprint(member, False) + +def test_linalg(): + mat = rand_mat((3, 3)) + vec = rand_mat((3,)) + + print("-------") + print("linalg") + print("-------") + + equal(np.linalg.solve(mat, vec), + jit(jet.linalg.solve)(mat, vec), 'solve') + equal(np.linalg.norm(vec), + jit(lambda x: jet.linalg.norm(x))(vec), 'norm') + +def test_random(): + print("-------") + print("random") + print("-------") + + scalar_1 = rand_mat(()) + scalar_2 = rand_mat(()) + 1.0 + + n = int(1e5) + random = np.zeros((n,)) + jet_normal = jit(jet.random.normal) + for i in xrange(n): + random[i] = jet_normal(scalar_1, scalar_2) + equal(np.sum(random) / n, + scalar_1, 'normal, mean', tol=1.e-1) + equal(np.std(random), + scalar_2, 'normal, sd', tol=1.e-2) + + +if __name__ == '__main__': + test_0D() + test_1D() + test_2D() + test_mix() + test_random() + test_linalg()