diff --git a/boxtree/tools.py b/boxtree/tools.py index 5dc58319420b02fee0d39927bc089de51fc44f7a..9c74d960aac4e3c696a9d2e8ade6d3ed8aa13e3f 100644 --- a/boxtree/tools.py +++ b/boxtree/tools.py @@ -46,9 +46,9 @@ VectorArg = partial(_VectorArg, with_offset=True) AXIS_NAMES = ("x", "y", "z", "w") -def padded_bin(i, l): - """Format *i* as binary number, pad it to length *l*.""" - return bin(i)[2:].rjust(l, "0") +def padded_bin(i, nbits): + """Format *i* as binary number, pad it to length *nbits*.""" + return bin(i)[2:].rjust(nbits, "0") # NOTE: Order of positional args should match GappyCopyAndMapKernel.__call__() @@ -270,7 +270,7 @@ class DeviceDataRecord(Record): instances on the host. """ - def _transform_arrays(self, f): + def _transform_arrays(self, f, exclude_fields=frozenset()): result = {} def transform_val(val): @@ -290,6 +290,9 @@ class DeviceDataRecord(Record): return f(val) for field_name in self.__class__.fields: + if field_name in exclude_fields: + continue + try: attr = getattr(self, field_name) except AttributeError: @@ -335,9 +338,12 @@ class DeviceDataRecord(Record): return self._transform_arrays(try_with_queue) - def to_device(self, queue): + def to_device(self, queue, exclude_fields=frozenset()): """ Return a copy of `self` in all :class:`numpy.ndarray` arrays are transferred to device memory as :class:`pyopencl.array.Array` objects. + + :arg exclude_fields: a :class:`frozenset` containing fields excluding from + transferring to the device memory. """ def _to_device(attr): @@ -346,7 +352,7 @@ class DeviceDataRecord(Record): else: return attr - return self._transform_arrays(_to_device) + return self._transform_arrays(_to_device, exclude_fields) # }}} diff --git a/boxtree/tree.py b/boxtree/tree.py index c09ee65a24960b1eccb993b717e0c3291f5dcf8e..a67a7c2db637610ed08ae04894e8adaeca05d717 100644 --- a/boxtree/tree.py +++ b/boxtree/tree.py @@ -389,6 +389,13 @@ class Tree(DeviceDataRecord): # }}} + def to_device(self, queue, exclude_fields=frozenset()): + # level_start_box_nrs should remain in host memory + exclude_fields = set(exclude_fields) + exclude_fields.add("level_start_box_nrs") + + return super(Tree, self).to_device(queue, frozenset(exclude_fields)) + # }}} diff --git a/boxtree/tree_build.py b/boxtree/tree_build.py index ff944c6ae783e377eeb1771772ea80beec239165..8f652aa8397787b35199e74d18bb2ef041697eb9 100644 --- a/boxtree/tree_build.py +++ b/boxtree/tree_build.py @@ -718,8 +718,8 @@ class TreeBuilder(object): # Currently undocumented. lr_lookbehind_levels = kwargs.get("lr_lookbehind", 1) minimal_new_level_length += sum( - 2**(l*dimensions) * new_level_leaf_counts[level - l] - for l in range(1, 1 + min(level, lr_lookbehind_levels))) + 2**(lev*dimensions) * new_level_leaf_counts[level - lev] + for lev in range(1, 1 + min(level, lr_lookbehind_levels))) nboxes_minimal = \ sum(minimal_upper_level_lengths) + minimal_new_level_length @@ -740,8 +740,8 @@ class TreeBuilder(object): # Recompute the level padding. for ulevel in range(level): upper_level_padding[ulevel] = sum( - 2**(l*dimensions) * new_level_leaf_counts[ulevel - l] - for l in range( + 2**(lev*dimensions) * new_level_leaf_counts[ulevel - lev] + for lev in range( 1, 1 + min(ulevel, lr_lookbehind_levels))) new_upper_level_unused_box_counts = np.max(