diff --git a/pymbolic/mapper/c_code.py b/pymbolic/mapper/c_code.py index f002d8e1e51ed4cbe8efd25b76e310899b77f975..26fa4e3b79944253ebf2c366d9e1da45da554343 100644 --- a/pymbolic/mapper/c_code.py +++ b/pymbolic/mapper/c_code.py @@ -4,13 +4,27 @@ from pymbolic.mapper.stringifier import SimplifyingSortingStringifyMapper class CCodeMapper(SimplifyingSortingStringifyMapper): - def __init__(self, constant_mapper=repr, reverse=True, cse_prefix="_cse"): + def __init__(self, constant_mapper=repr, reverse=True, + cse_prefix="_cse", complex_constant_base_type="double"): SimplifyingSortingStringifyMapper.__init__(self, constant_mapper, reverse) self.cse_prefix = cse_prefix self.cses = [] self.cse_to_index = {} + self.complex_constant_base_type = complex_constant_base_type + # mappings ---------------------------------------------------------------- + def map_constant(self, x, enclosing_prec): + import numpy + if isinstance(x, complex): + return "std::complex<%s>(%s, %s)" % ( + complex_constant_base_type, + self.constant_mapper(x.real), + self.constant_mapper(x.imag)) + else: + return SimplifyingSortingStringifyMapper.map_constant( + self, x, enclosing_prec) + def map_call(self, expr, enclosing_prec): from pymbolic.primitives import Variable from pymbolic.mapper.stringifier import PREC_NONE, PREC_CALL