diff --git a/grudge/array_context.py b/grudge/array_context.py index 50fa9045..15e9f0f2 100644 --- a/grudge/array_context.py +++ b/grudge/array_context.py @@ -246,7 +246,7 @@ def _dag_to_compiled_func(self, dict_of_named_arrays, # type-ignore-reason: 'PytatoPyOpenCLArrayContext' has no 'mpi_communicator' # pylint: disable=no-member - distributed_partition, _new_mpi_base_tag = number_distributed_tags( + distributed_partition, new_mpi_base_tag = number_distributed_tags( self.actx.mpi_communicator, distributed_partition, base_tag=prev_mpi_base_tag) @@ -254,7 +254,7 @@ def _dag_to_compiled_func(self, dict_of_named_arrays, assert prev_mpi_base_tag == self.actx.mpi_base_tag # FIXME: Updating stuff inside the array context from here is *cough* # not super pretty. - self.actx.mpi_base_tag = _new_mpi_base_tag + self.actx.mpi_base_tag = new_mpi_base_tag self.actx._compile_trace_callback(self.f, "post_number_distributed_tags", distributed_partition) diff --git a/grudge/dof_desc.py b/grudge/dof_desc.py index 2fefbffa..91e7d7e1 100644 --- a/grudge/dof_desc.py +++ b/grudge/dof_desc.py @@ -389,8 +389,11 @@ def _normalize_domain_and_discr_tag( *, _contextual_volume_tag: VolumeTag | None = None ) -> tuple[DomainTag, DiscretizationTag]: - if _contextual_volume_tag is None: - _contextual_volume_tag = VTAG_ALL + contextual_volume_tag = _contextual_volume_tag + del _contextual_volume_tag + + if contextual_volume_tag is None: + contextual_volume_tag = VTAG_ALL if domain == "scalar": domain = DTAG_SCALAR @@ -399,13 +402,13 @@ def _normalize_domain_and_discr_tag( elif domain in [VTAG_ALL, "vol"]: domain = DTAG_VOLUME_ALL elif domain in [FACE_RESTR_ALL, "all_faces"]: - domain = BoundaryDomainTag(FACE_RESTR_ALL, _contextual_volume_tag) + domain = BoundaryDomainTag(FACE_RESTR_ALL, contextual_volume_tag) elif domain in [FACE_RESTR_INTERIOR, "int_faces"]: - domain = BoundaryDomainTag(FACE_RESTR_INTERIOR, _contextual_volume_tag) + domain = BoundaryDomainTag(FACE_RESTR_INTERIOR, contextual_volume_tag) elif isinstance(domain, BTAG_PARTITION): - domain = BoundaryDomainTag(domain, _contextual_volume_tag) + domain = BoundaryDomainTag(domain, contextual_volume_tag) elif domain in [BTAG_ALL, BTAG_REALLY_ALL, BTAG_NONE]: - domain = BoundaryDomainTag(domain, _contextual_volume_tag) + domain = BoundaryDomainTag(domain, contextual_volume_tag) else: raise ValueError(f"domain tag not understood: {domain}") diff --git a/grudge/geometry/metrics.py b/grudge/geometry/metrics.py index af8267ce..d80c7087 100644 --- a/grudge/geometry/metrics.py +++ b/grudge/geometry/metrics.py @@ -330,7 +330,8 @@ def inverse_metric_derivative_mat( result = np.zeros((ambient_dim, dim), dtype=object) for i in range(dim): for j in range(ambient_dim): - result[i, j] = inverse_metric_derivative( + # type-ignore because numpy typing doesn't like object arrays + result[i, j] = inverse_metric_derivative( # type: ignore[call-overload] actx, dcoll, i, j, dd=dd, _use_geoderiv_connection=_use_geoderiv_connection ) @@ -564,17 +565,17 @@ def _signed_face_ones( actx, dtype=dcoll.real_dtype ) + 1 - _signed_face_ones_numpy = actx.to_numpy(signed_ones) + signed_face_ones_numpy = actx.to_numpy(signed_ones) for igrp, grp in enumerate(all_faces_conn.groups): for batch in grp.batches: assert batch.to_element_face is not None i = actx.to_numpy(actx.thaw(batch.to_element_indices)) - grp_field = _signed_face_ones_numpy[igrp].reshape(-1) + grp_field = signed_face_ones_numpy[igrp].reshape(-1) grp_field[i] = \ (2.0 * (batch.to_element_face % 2) - 1.0) * grp_field[i] - return actx.from_numpy(_signed_face_ones_numpy) + return actx.from_numpy(signed_face_ones_numpy) def parametrization_derivative( @@ -724,10 +725,13 @@ def mv_normal( """ dd = dof_desc.as_dofdesc(dd) - if _use_geoderiv_connection is None: - _use_geoderiv_connection = actx.supports_nonscalar_broadcasting + use_geoderiv_connection = _use_geoderiv_connection + del _use_geoderiv_connection - @memoize_in(dcoll, (mv_normal, dd, _use_geoderiv_connection)) + if use_geoderiv_connection is None: + use_geoderiv_connection = actx.supports_nonscalar_broadcasting + + @memoize_in(dcoll, (mv_normal, dd, use_geoderiv_connection)) def _normal(): dim = dcoll.discr_from_dd(dd).dim ambient_dim = dcoll.ambient_dim @@ -740,7 +744,7 @@ def _normal(): if dim == ambient_dim - 1: result = rel_mv_normal( actx, dcoll, dd=dd, - _use_geoderiv_connection=_use_geoderiv_connection) + use_geoderiv_connection=use_geoderiv_connection) else: # NOTE: In the case of (d - 2)-dimensional curves, we don't really have # enough information on the face to decide what an "exterior face normal" @@ -759,12 +763,12 @@ def _normal(): rel_mv_normal( actx, dcoll, dd=dd.untrace(), - _use_geoderiv_connection=_use_geoderiv_connection + use_geoderiv_connection=use_geoderiv_connection ).as_vector(dtype=object)) ) pder = pseudoscalar( actx, dcoll, dd=dd, - _use_geoderiv_connection=_use_geoderiv_connection) + use_geoderiv_connection=use_geoderiv_connection) mv = -(volm_normal ^ pder) << volm_normal.I.inv()