diff --git a/pystencils/data_types.py b/pystencils/data_types.py index 3b2fe35dadbfb2a0577973ed2e32d0280d637aa6..3880edbcfc95f7f011d469a4754f35f2537ea4a5 100644 --- a/pystencils/data_types.py +++ b/pystencils/data_types.py @@ -84,6 +84,37 @@ class cast_func(sp.Function): def dtype(self): return self.args[1] + @property + def is_integer(self): + if hasattr(self.dtype, 'numpy_dtype'): + return np.issubdtype(self.dtype.numpy_dtype, np.integer) or super().is_integer + else: + return super().is_integer + + @property + def is_negative(self): + if hasattr(self.dtype, 'numpy_dtype'): + if np.issubdtype(self.dtype.numpy_dtype, np.unsignedinteger): + return False + + return super().is_negative + + @property + def is_nonnegative(self): + if self.is_negative is False: + return True + else: + return super().is_nonnegative + + @property + def is_real(self): + if hasattr(self.dtype, 'numpy_dtype'): + return np.issubdtype(self.dtype.numpy_dtype, np.integer) or \ + np.issubdtype(self.dtype.numpy_dtype, np.floating) or \ + super().is_real + else: + return super().is_real + # noinspection PyPep8Naming class boolean_cast_func(cast_func, Boolean): diff --git a/pystencils/math_optimizations.py b/pystencils/math_optimizations.py index aa866570af196cbb0c32c25920cbfe062d065b21..ad0114782e15977fb329baacfb6ad4b6bac1e33b 100644 --- a/pystencils/math_optimizations.py +++ b/pystencils/math_optimizations.py @@ -9,7 +9,6 @@ import itertools from pystencils import Assignment from pystencils.astnodes import SympyAssignment -from pystencils.integer_functions import IntegerFunctionTwoArgsMixIn try: from sympy.codegen.rewriting import optims_c99, optimize @@ -38,7 +37,7 @@ def optimize_assignments(assignments, optimizations): if HAS_REWRITING: assignments = [Assignment(a.lhs, optimize(a.rhs, optimizations)) - if hasattr(a, 'lhs') and not a.rhs.atoms(IntegerFunctionTwoArgsMixIn) + if hasattr(a, 'lhs') else a for a in assignments] assignments_nodes = [a.atoms(SympyAssignment) for a in assignments] for a in itertools.chain.from_iterable(assignments_nodes):