Merged in support-array-types (pull request #175)

Support any type of array in equations/integrators

pysph/base/cython_generator.py

@@ -79,6 +79,9 @@ class KnownType(object):
     def __init__(self, type_str):
         self.type = type_str
 
+    def __repr__(self):
+        return 'KnownType("%s")'%self.type
+
 
 class CythonGenerator(object):
     def __init__(self, known_types=None, python_methods=False):
@@ -112,12 +115,12 @@ def ctype_to_python(self, type_str):
     def detect_type(self, name, value):
         """Given the variable name and value, detect its type.
         """
+        if isinstance(value, KnownType):
+            return value.type
         if name.startswith(('s_', 'd_')) and name not in ['s_idx', 'd_idx']:
             return 'double*'
         if name in ['s_idx', 'd_idx']:
             return 'long'
-        if isinstance(value, KnownType):
-            return value.type
         if value is Undefined or isinstance(value, Undefined):
             raise CodeGenerationError('Unknown type, for %s'%name)
 

pysph/base/tests/test_cython_generator.py

@@ -27,7 +27,7 @@ def func(self, d_idx=0, d_x=[0.0, 0.0]):
         return d_x[d_idx]
 
 class EqWithKnownTypes:
-    def some_func(self, d_idx, d_p, WIJ, DWIJ, user):
+    def some_func(self, d_idx, d_p, WIJ, DWIJ, user, d_user, s_user):
         d_p[d_idx] = WIJ*DWIJ[0]
 
 class EqWithMatrix:
@@ -250,7 +250,10 @@ def __init__(self, **kwargs):
     def test_method_with_known_types(self):
         cg = CythonGenerator(
             known_types={'WIJ':0.0, 'DWIJ':[0.0, 0.0, 0.0],
-                         'user': KnownType('ndarray')}
+                         'user': KnownType('ndarray'),
+                         'd_user': KnownType('long*'),
+                         's_user': KnownType('int*'),
+                         }
         )
         cg.parse(EqWithKnownTypes())
         expect = dedent("""
@@ -259,7 +262,7 @@ def __init__(self, **kwargs):
                 for key, value in kwargs.iteritems():
                     setattr(self, key, value)
 
-            cdef inline void some_func(self, long d_idx, double* d_p, double WIJ, double* DWIJ, ndarray user):
+            cdef inline void some_func(self, long d_idx, double* d_p, double WIJ, double* DWIJ, ndarray user, long* d_user, int* s_user):
                 d_p[d_idx] = WIJ*DWIJ[0]
         """)
         self.assert_code_equal(cg.get_code().strip(), expect.strip())

pysph/sph/acceleration_eval_cython.mako

@@ -136,15 +136,13 @@ from pysph.base.nnps import get_number_of_threads
 from pyzoltan.core.carray cimport (DoubleArray, IntArray, UIntArray,
     aligned, aligned_free, aligned_malloc)
 
-${header}
+${helper.get_header()}
 
 # #############################################################################
 cdef class ParticleArrayWrapper:
     cdef public int index
     cdef public ParticleArray array
-    cdef public IntArray tag, pid
-    cdef public UIntArray gid
-    cdef public DoubleArray ${array_names}
+    ${indent(helper.get_array_decl_for_wrapper(), 1)}
     cdef public str name
 
     def __init__(self, pa, index):
@@ -169,7 +167,7 @@ cdef class ParticleArrayWrapper:
 # #############################################################################
 cdef class AccelerationEval:
     cdef public tuple particle_arrays
-    cdef public ParticleArrayWrapper ${pa_names}
+    cdef public ParticleArrayWrapper ${helper.get_particle_array_names()}
     cdef public NNPS nnps
     cdef public int n_threads
     cdef public list _nbr_refs

pysph/sph/acceleration_eval_cython_helper.py

@@ -1,8 +1,10 @@
+from collections import defaultdict
 from mako.template import Template
 from os.path import dirname, join
 
+from pyzoltan.core import carray
 from pysph.base.config import get_config
-from pysph.base.cython_generator import CythonGenerator
+from pysph.base.cython_generator import CythonGenerator, KnownType
 
 
 ###############################################################################
@@ -25,37 +27,95 @@ def get_code(obj):
     return result
 
 ###############################################################################
-def get_array_names(particle_arrays):
-    """Get the union of the names of all particle array properties.
+def get_all_array_names(particle_arrays):
+    """For each type of carray, find the union of the names of all particle
+    array properties/constants along with their array type.  Returns a
+    dictionary keyed on the name of the Array class with values being a set of
+    property names for each.
+
+    Parameters
+    ----------
+
+    particle_arrays : list
+        A list of particle arrays.
+
+    Examples
+    --------
+
+    A simple example would be::
+
+        >>> x = np.linspace(0, 1, 10)
+        >>> pa = ParticleArray(name='f', x=x)
+        >>> get_all_array_names([pa])
+        {'DoubleArray': {'x'}, 'IntArray': {'pid', 'tag'}, 'UIntArray': {'gid'}}
     """
-    props = set()
+    props = defaultdict(set)
     for array in particle_arrays:
-        for name in array.properties.keys() + array.constants.keys():
-            props.add(name)
-    props.difference_update(set(('tag', 'pid', 'gid')))
-    array_names = ', '.join(sorted(props))
-    return array_names
+        for properties in (array.properties, array.constants):
+            for name, arr in properties.iteritems():
+                a_type = arr.__class__.__name__
+                props[a_type].add(name)
+    return dict(props)
+
+def get_known_types_for_arrays(array_names):
+    """Given all the array names from `get_all_array_names` this creates known
+    types for each of them so that the code generators can use this type
+    information when needed.  Note that known type info is generated for both
+    source and destination style arrays.
+
+    Parameters
+    ----------
+
+    array_names: dict
+        A dictionary produced by `get_all_array_names`.
+
+    Examples
+    --------
+
+    A simple example would be::
+
+        >>> x = np.linspace(0, 1, 10)
+        >>> pa = ParticleArray(name='f', x=x)
+        >>> pa.remove_property('pid')
+        >>> info = get_all_array_names([pa])
+        >>> get_known_types_for_arrays(info)
+        {'d_gid': KnownType("unsigned int*"),
+         'd_tag': KnownType("int*"),
+         'd_x': KnownType("double*"),
+         's_gid': KnownType("unsigned int*"),
+         's_tag': KnownType("int*"),
+         's_x': KnownType("double*")}
+
+    """
+    result = {}
+    for arr_type, arrays in array_names.iteritems():
+        c_type = getattr(carray, arr_type)().get_c_type()
+        for arr in arrays:
+            known_type = KnownType(c_type + '*')
+            result['s_' + arr] = known_type
+            result['d_' + arr] = known_type
+    return result
 
 
 ###############################################################################
 class AccelerationEvalCythonHelper(object):
     def __init__(self, acceleration_eval):
         self.object = acceleration_eval
         self.config = get_config()
+        self.all_array_names = get_all_array_names(
+            self.object.particle_arrays
+        )
+        self.known_types = get_known_types_for_arrays(
+            self.all_array_names
+        )
 
     ##########################################################################
     # Public interface.
     ##########################################################################
     def get_code(self):
-        object = self.object
-        header = self.get_header()
-        array_names =  get_array_names(object.particle_arrays)
-        parrays = [pa.name for pa in object.particle_arrays]
-        pa_names = ', '.join(parrays)
         path = join(dirname(__file__), 'acceleration_eval_cython.mako')
         template = Template(filename=path)
-        main = template.render(header=header, array_names=array_names,
-                               pa_names=pa_names, helper=self)
+        main = template.render(helper=self)
         return main
 
     def setup_compiled_module(self, module):
@@ -70,6 +130,18 @@ def setup_compiled_module(self, module):
     ##########################################################################
     # Mako interface.
     ##########################################################################
+    def get_array_decl_for_wrapper(self):
+        array_names = self.all_array_names
+        decl = []
+        for a_type in sorted(array_names.keys()):
+            props = array_names[a_type]
+            decl.append(
+                'cdef public {a_type} {attrs}'.format(
+                    a_type=a_type, attrs=', '.join(sorted(props))
+                )
+            )
+        return '\n'.join(decl)
+
     def get_header(self):
         object = self.object
         headers = []
@@ -80,12 +152,14 @@ def get_header(self):
             headers.extend(get_code(equation))
 
         # Kernel wrappers.
-        cg = CythonGenerator()
+        cg = CythonGenerator(known_types=self.known_types)
         cg.parse(object.kernel)
         headers.append(cg.get_code())
 
         # Equation wrappers.
-        headers.append(object.all_group.get_equation_wrappers())
+        headers.append(object.all_group.get_equation_wrappers(
+            self.known_types
+        ))
 
         return '\n'.join(headers)
 
@@ -111,7 +185,7 @@ def get_array_declarations(self):
         group = self.object.all_group
         src, dest = group.get_array_names()
         src.update(dest)
-        return group.get_array_declarations(src)
+        return group.get_array_declarations(src, self.known_types)
 
     def get_dest_array_setup(self, dest_name, eqs_with_no_source, sources, real):
         src, dest_arrays = eqs_with_no_source.get_array_names()
@@ -136,3 +210,6 @@ def get_parallel_block(self):
         else:
             return "if True: # Placeholder used for OpenMP."
 
+    def get_particle_array_names(self):
+        parrays = [pa.name for pa in self.object.particle_arrays]
+        return ', '.join(parrays)

pysph/sph/equation.py

@@ -580,10 +580,15 @@ def get_variable_names(self):
 
         return filtered_vars
 
-    def get_array_declarations(self, names):
+    def get_array_declarations(self, names, known_types={}):
         decl = []
         for arr in sorted(names):
-            decl.append('cdef double* %s'%arr)
+            if arr in known_types:
+                decl.append('cdef {type} {arr}'.format(
+                    type=known_types[arr].type, arr=arr
+                ))
+            else:
+                decl.append('cdef double* %s'%arr)
         return '\n'.join(decl)
 
     def get_variable_declarations(self, context):
@@ -641,7 +646,7 @@ def get_converged_condition(self):
             # and not be short-circuited by the first one that returns False.
             return ' & '.join(code)
 
-    def get_equation_wrappers(self):
+    def get_equation_wrappers(self, known_types={}):
         classes = defaultdict(lambda: 0)
         eqs = {}
         for equation in self.equations:
@@ -651,7 +656,8 @@ def get_equation_wrappers(self):
             classes[cls] += 1
             eqs[cls] = equation
         wrappers = []
-        predefined = get_predefined_types(self.pre_comp)
+        predefined = dict(get_predefined_types(self.pre_comp))
+        predefined.update(known_types)
         code_gen = CythonGenerator(known_types=predefined)
         for cls in sorted(classes.keys()):
             code_gen.parse(eqs[cls])

pysph/sph/integrator_cython_helper.py

@@ -19,8 +19,9 @@
 class IntegratorCythonHelper(object):
     """A helper that generates Cython code for the Integrator class.
     """
-    def __init__(self, integrator):
+    def __init__(self, integrator, acceleration_eval_helper):
         self.object = integrator
+        self.acceleration_eval_helper = acceleration_eval_helper
 
     def get_code(self):
         if self.object is not None:
@@ -51,7 +52,9 @@ def get_stepper_code(self):
             classes[cls] = stepper
 
         wrappers = []
-        code_gen = CythonGenerator()
+        code_gen = CythonGenerator(
+            known_types=self.acceleration_eval_helper.known_types
+        )
         for cls in sorted(classes.keys()):
             code_gen.parse(classes[cls])
             wrappers.append(code_gen.get_code())
@@ -87,9 +90,15 @@ def get_array_declarations(self, method):
             s, d = get_array_names(self.get_args(dest, method))
             arrays.update(s | d)
 
+        known_types = self.acceleration_eval_helper.known_types
         decl = []
         for arr in sorted(arrays):
-            decl.append('cdef double* %s'%arr)
+            if arr in known_types:
+                decl.append('cdef {type} {arr}'.format(
+                    type=known_types[arr].type, arr=arr
+                ))
+            else:
+                decl.append('cdef double* %s'%arr)
         return '\n'.join(decl)
 
     def get_array_setup(self, dest, method):
@@ -120,4 +129,4 @@ def get_timestep_code(self):
         method = self.object.one_timestep
         sourcelines = inspect.getsourcelines(method)[0]
         defn, lines = get_func_definition(sourcelines)
-	return dedent(''.join(lines))
+        return dedent(''.join(lines))

pysph/sph/sph_compiler.py

@@ -11,7 +11,9 @@ def __init__(self, acceleration_eval, integrator):
             self.acceleration_eval
         )
         self.integrator = integrator
-        self.integrator_helper = IntegratorCythonHelper(integrator)
+        self.integrator_helper = IntegratorCythonHelper(
+            integrator, self.acceleration_eval_helper
+        )
         self.ext_mod = None
         self.module = None