diff --git a/examples/dagrt-fusion.py b/examples/dagrt-fusion.py index 2a285ad533b0d1ffb545da65f69efc8665072a10..3cd6b71faf176db5a067b06d914ecd690909d10c 100755 --- a/examples/dagrt-fusion.py +++ b/examples/dagrt-fusion.py @@ -107,6 +107,8 @@ def dof_array_nbytes(ary: np.ndarray): return sum( dof_array_nbytes(ary[idx]) for idx in np.ndindex(ary.shape)) + elif isinstance(ary, DOFArray): + return sum(dof_array_nbytes(ary_i) for ary_i in ary) else: return ary.nbytes diff --git a/grudge/eager.py b/grudge/eager.py index 60fc7bd3e19e18e2190a7aa647b41bfc5cec9a6e..482a37ca8afe841b218789e4b5d556b2cecb3cb4 100644 --- a/grudge/eager.py +++ b/grudge/eager.py @@ -28,7 +28,7 @@ import pyopencl.array as cla # noqa from grudge import sym, bind from meshmode.mesh import BTAG_ALL, BTAG_NONE, BTAG_PARTITION # noqa -from meshmode.dof_array import freeze, DOFArray, flatten, unflatten +from meshmode.dof_array import freeze, flatten, unflatten from grudge.discretization import DGDiscretizationWithBoundaries from grudge.symbolic.primitives import TracePair @@ -83,9 +83,7 @@ class EagerDGDiscretization(DGDiscretizationWithBoundaries): :arg tgt: a :class:`~grudge.sym.DOFDesc`, or a value convertible to one :arg vec: a :class:`~meshmode.dof_array.DOFArray` """ - if (isinstance(vec, np.ndarray) - and vec.dtype.char == "O" - and not isinstance(vec, DOFArray)): + if isinstance(vec, np.ndarray): return obj_array_vectorize( lambda el: self.project(src, tgt, el), vec) @@ -270,9 +268,7 @@ class EagerDGDiscretization(DGDiscretizationWithBoundaries): local_only=True) def inverse_mass(self, vec): - if (isinstance(vec, np.ndarray) - and vec.dtype.char == "O" - and not isinstance(vec, DOFArray)): + if isinstance(vec, np.ndarray): return obj_array_vectorize( lambda el: self.inverse_mass(el), vec) @@ -292,9 +288,7 @@ class EagerDGDiscretization(DGDiscretizationWithBoundaries): else: raise TypeError("invalid number of arguments") - if (isinstance(vec, np.ndarray) - and vec.dtype.char == "O" - and not isinstance(vec, DOFArray)): + if isinstance(vec, np.ndarray): return obj_array_vectorize( lambda el: self.face_mass(dd, el), vec) @@ -312,9 +306,7 @@ class EagerDGDiscretization(DGDiscretizationWithBoundaries): dd = sym.as_dofdesc(dd) - if (isinstance(vec, np.ndarray) - and vec.dtype.char == "O" - and not isinstance(vec, DOFArray)): + if isinstance(vec, np.ndarray): if p == 2: return sum( self.norm(vec[idx], dd=dd)**2 @@ -350,9 +342,7 @@ class EagerDGDiscretization(DGDiscretizationWithBoundaries): def interior_trace_pair(discrwb, vec): i = discrwb.project("vol", "int_faces", vec) - if (isinstance(vec, np.ndarray) - and vec.dtype.char == "O" - and not isinstance(vec, DOFArray)): + if isinstance(vec, np.ndarray): e = obj_array_vectorize( lambda el: discrwb.opposite_face_connection()(el), i) @@ -412,9 +402,7 @@ def _cross_rank_trace_pairs_scalar_field(discrwb, vec, tag=None): def cross_rank_trace_pairs(discrwb, vec, tag=None): - if (isinstance(vec, np.ndarray) - and vec.dtype.char == "O" - and not isinstance(vec, DOFArray)): + if isinstance(vec, np.ndarray): n, = vec.shape result = {} diff --git a/grudge/execution.py b/grudge/execution.py index 59af300fbdf5aa25ba4ad3c0517daf05be2dec97..c80be0419b47801f2ec4139f21f88a72f1f87c9e 100644 --- a/grudge/execution.py +++ b/grudge/execution.py @@ -281,7 +281,7 @@ class ExecutionMapper(mappers.Evaluator, for in_grp, out_grp in zip(in_discr.groups, out_discr.groups): - cache_key = "elwise_linear", in_grp, out_grp, op, field.dtype + cache_key = "elwise_linear", in_grp, out_grp, op, field.entry_dtype try: matrix = self.bound_op.operator_data_cache[cache_key] except KeyError: @@ -347,7 +347,7 @@ class ExecutionMapper(mappers.Evaluator, assert len(all_faces_discr.groups) == len(vol_discr.groups) for afgrp, volgrp in zip(all_faces_discr.groups, vol_discr.groups): - cache_key = "face_mass", afgrp, op, field.dtype + cache_key = "face_mass", afgrp, op, field.entry_dtype nfaces = volgrp.mesh_el_group.nfaces