diff --git a/grudge/discretization.py b/grudge/discretization.py index e84016c61..3d37f8a56 100644 --- a/grudge/discretization.py +++ b/grudge/discretization.py @@ -29,6 +29,7 @@ THE SOFTWARE. """ +from typing import Any, Optional from pytools import memoize_method from grudge.dof_desc import ( @@ -698,7 +699,7 @@ def order(self): # {{{ Discretization-specific geometric properties - def nodes(self, dd=None): + def nodes(self, dd=None, dtype: Optional[np.dtype[Any]] = None): r"""Return the nodes of a discretization specified by *dd*. :arg dd: a :class:`~grudge.dof_desc.DOFDesc`, or a value convertible to one. @@ -707,7 +708,7 @@ def nodes(self, dd=None): """ if dd is None: dd = DD_VOLUME - return self.discr_from_dd(dd).nodes() + return self.discr_from_dd(dd).nodes(dtype) def normal(self, dd): r"""Get the unit normal to the specified surface discretization, *dd*. diff --git a/grudge/geometry/metrics.py b/grudge/geometry/metrics.py index 492d8b7f7..86a2ab725 100644 --- a/grudge/geometry/metrics.py +++ b/grudge/geometry/metrics.py @@ -58,11 +58,14 @@ """ +from typing import Optional, Any + import numpy as np from arraycontext import thaw, freeze, ArrayContext from meshmode.dof_array import DOFArray +from grudge.tools import to_real_dtype from grudge import DiscretizationCollection import grudge.dof_desc as dof_desc @@ -111,7 +114,8 @@ def to_quad(vec): def forward_metric_nth_derivative( actx: ArrayContext, dcoll: DiscretizationCollection, xyz_axis, ref_axes, dd=None, - *, _use_geoderiv_connection=False) -> DOFArray: + *, _use_geoderiv_connection=False, + dtype: Optional[np.dtype[Any]] = None) -> DOFArray: r"""Pointwise metric derivatives representing repeated derivatives of the physical coordinate enumerated by *xyz_axis*: :math:`x_{\mathrm{xyz\_axis}}` with respect to the coordiantes on the reference element :math:`\xi_i`: @@ -169,7 +173,8 @@ def forward_metric_nth_derivative( vec = num_reference_derivative( dcoll.discr_from_dd(inner_dd), flat_ref_axes, - thaw(dcoll.discr_from_dd(inner_dd).nodes(), actx)[xyz_axis] + thaw(dcoll.discr_from_dd(inner_dd).nodes( + dtype=to_real_dtype(dtype)), actx)[xyz_axis] ) return _geometry_to_quad_if_requested( @@ -178,7 +183,8 @@ def forward_metric_nth_derivative( def forward_metric_derivative_vector( actx: ArrayContext, dcoll: DiscretizationCollection, rst_axis, dd=None, - *, _use_geoderiv_connection=False) -> np.ndarray: + *, _use_geoderiv_connection=False, dtype: Optional[np.dtype[Any]] = None + ) -> np.ndarray: r"""Computes an object array containing the forward metric derivatives of each physical coordinate. @@ -195,7 +201,8 @@ def forward_metric_derivative_vector( return make_obj_array([ forward_metric_nth_derivative( actx, dcoll, i, rst_axis, dd=dd, - _use_geoderiv_connection=_use_geoderiv_connection) + _use_geoderiv_connection=_use_geoderiv_connection, + dtype=dtype) for i in range(dcoll.ambient_dim) ] ) @@ -203,7 +210,8 @@ def forward_metric_derivative_vector( def forward_metric_derivative_mv( actx: ArrayContext, dcoll: DiscretizationCollection, rst_axis, dd=None, - *, _use_geoderiv_connection=False) -> MultiVector: + *, _use_geoderiv_connection=False, dtype: Optional[np.dtype[Any]] = None + ) -> MultiVector: r"""Computes a :class:`pymbolic.geometric_algebra.MultiVector` containing the forward metric derivatives of each physical coordinate. @@ -220,13 +228,15 @@ def forward_metric_derivative_mv( return MultiVector( forward_metric_derivative_vector( actx, dcoll, rst_axis, dd=dd, - _use_geoderiv_connection=_use_geoderiv_connection) + _use_geoderiv_connection=_use_geoderiv_connection, + dtype=dtype) ) def forward_metric_derivative_mat( actx: ArrayContext, dcoll: DiscretizationCollection, dd=None, - *, _use_geoderiv_connection=False) -> np.ndarray: + *, _use_geoderiv_connection=False, + dtype: Optional[np.dtype[Any]] = None) -> np.ndarray: r"""Computes the forward metric derivative matrix, also commonly called the Jacobian matrix, with entries defined as the forward metric derivatives: @@ -260,13 +270,15 @@ def forward_metric_derivative_mat( for j in range(dim): result[:, j] = forward_metric_derivative_vector( actx, dcoll, j, dd=dd, - _use_geoderiv_connection=_use_geoderiv_connection) + _use_geoderiv_connection=_use_geoderiv_connection, + dtype=dtype) return result def first_fundamental_form(actx: ArrayContext, dcoll: DiscretizationCollection, - dd=None, *, _use_geoderiv_connection=False) -> np.ndarray: + dd=None, *, _use_geoderiv_connection=False, + dtype: Optional[np.dtype[Any]] = None) -> np.ndarray: r"""Computes the first fundamental form using the Jacobian matrix: .. math:: @@ -295,14 +307,16 @@ def first_fundamental_form(actx: ArrayContext, dcoll: DiscretizationCollection, dd = DD_VOLUME mder = forward_metric_derivative_mat( - actx, dcoll, dd=dd, _use_geoderiv_connection=_use_geoderiv_connection) + actx, dcoll, dd=dd, _use_geoderiv_connection=_use_geoderiv_connection, + dtype=dtype) return mder.T.dot(mder) def inverse_metric_derivative_mat( actx: ArrayContext, dcoll: DiscretizationCollection, dd=None, - *, _use_geoderiv_connection=False) -> np.ndarray: + *, _use_geoderiv_connection=False, dtype: Optional[np.dtype[Any]] = None + ) -> np.ndarray: r"""Computes the inverse metric derivative matrix, which is the inverse of the Jacobian (forward metric derivative) matrix. @@ -324,15 +338,16 @@ def inverse_metric_derivative_mat( for j in range(ambient_dim): result[i, j] = inverse_metric_derivative( actx, dcoll, i, j, dd=dd, - _use_geoderiv_connection=_use_geoderiv_connection - ) + _use_geoderiv_connection=_use_geoderiv_connection, + dtype=dtype) return result def inverse_first_fundamental_form( actx: ArrayContext, dcoll: DiscretizationCollection, dd=None, - *, _use_geoderiv_connection=False) -> np.ndarray: + *, _use_geoderiv_connection=False, dtype: Optional[np.dtype[Any]] + ) -> np.ndarray: r"""Computes the inverse of the first fundamental form: .. math:: @@ -361,11 +376,13 @@ def inverse_first_fundamental_form( if dcoll.ambient_dim == dim: inv_mder = inverse_metric_derivative_mat( - actx, dcoll, dd=dd, _use_geoderiv_connection=_use_geoderiv_connection) + actx, dcoll, dd=dd, _use_geoderiv_connection=_use_geoderiv_connection, + dtype=dtype) inv_form1 = inv_mder.dot(inv_mder.T) else: form1 = first_fundamental_form( - actx, dcoll, dd=dd, _use_geoderiv_connection=_use_geoderiv_connection) + actx, dcoll, dd=dd, _use_geoderiv_connection=_use_geoderiv_connection, + dtype=dtype) if dim == 1: inv_form1 = 1.0 / form1 @@ -383,7 +400,7 @@ def inverse_first_fundamental_form( def inverse_metric_derivative( actx: ArrayContext, dcoll: DiscretizationCollection, rst_axis, xyz_axis, dd, - *, _use_geoderiv_connection=False + *, _use_geoderiv_connection=False, dtype: Optional[np.dtype[Any]] = None ) -> DOFArray: r"""Computes the inverse metric derivative of the physical coordinate enumerated by *xyz_axis* with respect to the @@ -409,7 +426,8 @@ def inverse_metric_derivative( par_vecs = [ forward_metric_derivative_mv( actx, dcoll, rst, dd, - _use_geoderiv_connection=_use_geoderiv_connection) + _use_geoderiv_connection=_use_geoderiv_connection, + dtype=dtype) for rst in range(dim)] # Yay Cramer's rule! @@ -442,7 +460,8 @@ def outprod_with_unit(i, at): def inverse_surface_metric_derivative( actx: ArrayContext, dcoll: DiscretizationCollection, rst_axis, xyz_axis, dd=None, - *, _use_geoderiv_connection=False): + *, _use_geoderiv_connection=False, + dtype: Optional[np.dtype[Any]]): r"""Computes the inverse surface metric derivative of the physical coordinate enumerated by *xyz_axis* with respect to the reference axis *rst_axis*. These geometric terms are used in the @@ -467,24 +486,24 @@ def inverse_surface_metric_derivative( dd = dof_desc.as_dofdesc(dd) if ambient_dim == dim: - result = inverse_metric_derivative( + return inverse_metric_derivative( actx, dcoll, rst_axis, xyz_axis, dd=dd, - _use_geoderiv_connection=_use_geoderiv_connection - ) + _use_geoderiv_connection=_use_geoderiv_connection, + dtype=dtype) else: inv_form1 = inverse_first_fundamental_form(actx, dcoll, dd=dd) - result = sum( + return sum( inv_form1[rst_axis, d]*forward_metric_nth_derivative( actx, dcoll, xyz_axis, d, dd=dd, - _use_geoderiv_connection=_use_geoderiv_connection + _use_geoderiv_connection=_use_geoderiv_connection, + dtype=dtype, ) for d in range(dim)) - return result - def inverse_surface_metric_derivative_mat( actx: ArrayContext, dcoll: DiscretizationCollection, dd=None, - *, times_area_element=False, _use_geoderiv_connection=False): + *, times_area_element=False, _use_geoderiv_connection=False, + dtype: Optional[np.dtype[Any]] = None): r"""Computes the matrix of inverse surface metric derivatives, indexed by ``(xyz_axis, rst_axis)``. It returns all values of :func:`inverse_surface_metric_derivative_mat` in cached matrix form. @@ -505,7 +524,7 @@ def inverse_surface_metric_derivative_mat( @memoize_in(dcoll, (inverse_surface_metric_derivative_mat, dd, times_area_element, _use_geoderiv_connection)) - def _inv_surf_metric_deriv(): + def _inv_surf_metric_deriv(dtype): if times_area_element: multiplier = area_element(actx, dcoll, dd=dd, _use_geoderiv_connection=_use_geoderiv_connection) @@ -517,13 +536,17 @@ def _inv_surf_metric_deriv(): multiplier * inverse_surface_metric_derivative(actx, dcoll, rst_axis, xyz_axis, dd=dd, - _use_geoderiv_connection=_use_geoderiv_connection) + _use_geoderiv_connection=_use_geoderiv_connection, + dtype=dtype) for rst_axis in range(dcoll.dim)]) for xyz_axis in range(dcoll.ambient_dim)]) return freeze(mat, actx) - return thaw(_inv_surf_metric_deriv(), actx) + if dtype is not None: + dtype = to_real_dtype(dtype) + + return thaw(_inv_surf_metric_deriv(dtype), actx) def _signed_face_ones( @@ -557,7 +580,8 @@ def _signed_face_ones( def parametrization_derivative( actx: ArrayContext, dcoll: DiscretizationCollection, dd, - *, _use_geoderiv_connection=False) -> MultiVector: + *, _use_geoderiv_connection=False, + dtype: Optional[np.dtype[Any]]) -> MultiVector: r"""Computes the product of forward metric derivatives spanning the tangent space with topological dimension *dim*. @@ -585,13 +609,15 @@ def parametrization_derivative( return product( forward_metric_derivative_mv( actx, dcoll, rst_axis, dd, - _use_geoderiv_connection=_use_geoderiv_connection) + _use_geoderiv_connection=_use_geoderiv_connection, + dtype=dtype) for rst_axis in range(dim) ) def pseudoscalar(actx: ArrayContext, dcoll: DiscretizationCollection, - dd=None, *, _use_geoderiv_connection=False) -> MultiVector: + dd=None, *, _use_geoderiv_connection=False, + dtype: Optional[np.dtype[Any]]) -> MultiVector: r"""Computes the field of pseudoscalars for the domain/discretization identified by *dd*. @@ -607,12 +633,14 @@ def pseudoscalar(actx: ArrayContext, dcoll: DiscretizationCollection, return parametrization_derivative( actx, dcoll, dd, - _use_geoderiv_connection=_use_geoderiv_connection).project_max_grade() + _use_geoderiv_connection=_use_geoderiv_connection, + dtype=dtype).project_max_grade() def area_element( actx: ArrayContext, dcoll: DiscretizationCollection, dd=None, - *, _use_geoderiv_connection=False + *, _use_geoderiv_connection=False, + dtype: Optional[np.dtype[Any]] = None ) -> DOFArray: r"""Computes the scale factor used to transform integrals from reference to global space. @@ -623,6 +651,8 @@ def area_element( Defaults to the base volume discretization. :arg _use_geoderiv_connection: For internal use. See :func:`forward_metric_nth_derivative` for an explanation. + :arg dtype: the :class:`numpy.dtype` with which to return the area element + data. :returns: a :class:`~meshmode.dof_array.DOFArray` containing the transformed volumes for each element. """ @@ -630,15 +660,19 @@ def area_element( dd = DD_VOLUME @memoize_in(dcoll, (area_element, dd, _use_geoderiv_connection)) - def _area_elements(): + def _area_elements(dtype: np.dtype[Any]): result = actx.np.sqrt( pseudoscalar( actx, dcoll, dd=dd, - _use_geoderiv_connection=_use_geoderiv_connection).norm_squared()) + _use_geoderiv_connection=_use_geoderiv_connection, + dtype=dtype).norm_squared()) return freeze(result, actx) - return thaw(_area_elements(), actx) + if dtype is not None: + dtype = to_real_dtype(dtype) + + return thaw(_area_elements(dtype), actx) # }}} diff --git a/grudge/op.py b/grudge/op.py index 70bd1caf7..c4e6bc15d 100644 --- a/grudge/op.py +++ b/grudge/op.py @@ -49,6 +49,7 @@ THE SOFTWARE. """ +from typing import Any from arraycontext import ArrayContext, map_array_container from arraycontext.container import ArrayOrContainerT @@ -66,11 +67,12 @@ import numpy as np import grudge.dof_desc as dof_desc +from grudge.tools import to_real_dtype -from grudge.interpolation import interp # noqa: F401 -from grudge.projection import project # noqa: F401 +from grudge.interpolation import interp +from grudge.projection import project -from grudge.reductions import ( # noqa: F401 +from grudge.reductions import ( norm, nodal_sum, nodal_min, @@ -85,7 +87,7 @@ elementwise_integral, ) -from grudge.trace_pair import ( # noqa: F401 +from grudge.trace_pair import ( interior_trace_pair, interior_trace_pairs, connected_ranks, @@ -95,6 +97,19 @@ ) +__all__ = ( + "interp", "project", + + "norm", "nodal_sum", "nodal_min", "nodal_max", "nodal_sum_loc", + "nodal_min_loc", "nodal_max_loc", "integral", "elementwise_sum", + "elementwise_max", "elementwise_min", "elementwise_integral", + + "interior_trace_pair", "interior_trace_pairs", + "connected_ranks", "cross_rank_trace_pairs", + "bdry_trace_pair", "bv_trace_pair", + ) + + # {{{ common derivative "kernels" def _single_axis_derivative_kernel( @@ -140,7 +155,8 @@ def _gradient_kernel(actx, out_discr, in_discr, get_diff_mat, inv_jac_mat, vec, get_diff_mat( actx, out_element_group=out_grp, - in_element_group=in_grp + in_element_group=in_grp, + dtype=vec_i.dtype, ), vec_i, arg_names=("inv_jac_t", "ref_stiffT_mat", "vec"), @@ -244,7 +260,7 @@ def _grad_helper(dcoll, scalar_grad, *args, nested): # {{{ Derivative operators def _reference_derivative_matrices(actx: ArrayContext, - out_element_group, in_element_group): + out_element_group, in_element_group, dtype: np.dtype[Any]): # We're accepting in_element_group for interface consistency with # _reference_stiffness_transpose_matrix. assert out_element_group is in_element_group @@ -252,16 +268,17 @@ def _reference_derivative_matrices(actx: ArrayContext, @keyed_memoize_in( actx, _reference_derivative_matrices, lambda grp: grp.discretization_key()) - def get_ref_derivative_mats(grp): + def get_ref_derivative_mats(grp, dtype: np.dtype[Any]): from meshmode.discretization.poly_element import diff_matrices return actx.freeze( actx.from_numpy( np.asarray( [dfmat for dfmat in diff_matrices(grp)] - ) + ).astype(dtype) ) ) - return get_ref_derivative_mats(out_element_group) + + return get_ref_derivative_mats(out_element_group, to_real_dtype(dtype)) def _strong_scalar_grad(dcoll, dd_in, vec): @@ -273,7 +290,8 @@ def _strong_scalar_grad(dcoll, dd_in, vec): actx = vec.array_context inverse_jac_mat = inverse_surface_metric_derivative_mat(actx, dcoll, - _use_geoderiv_connection=actx.supports_nonscalar_broadcasting) + _use_geoderiv_connection=actx.supports_nonscalar_broadcasting, + dtype=vec.entry_dtype) return _gradient_kernel(actx, discr, discr, _reference_derivative_matrices, inverse_jac_mat, vec, metric_in_matvec=False) @@ -360,12 +378,13 @@ def local_div(dcoll: DiscretizationCollection, vecs) -> ArrayOrContainerT: # {{{ Weak derivative operators def _reference_stiffness_transpose_matrix( - actx: ArrayContext, out_element_group, in_element_group): + actx: ArrayContext, out_element_group, in_element_group, + dtype: np.dtype[Any]): @keyed_memoize_in( actx, _reference_stiffness_transpose_matrix, lambda out_grp, in_grp: (out_grp.discretization_key(), in_grp.discretization_key())) - def get_ref_stiffness_transpose_mat(out_grp, in_grp): + def get_ref_stiffness_transpose_mat(out_grp, in_grp, dtype: np.dtype[Any]): if in_grp == out_grp: from meshmode.discretization.poly_element import \ mass_matrix, diff_matrices @@ -397,11 +416,12 @@ def get_ref_stiffness_transpose_mat(out_grp, in_grp): weights, vand_inv_t, grad_vand - ).copy() # contigify the array + ).astype(dtype).copy() # contigify the array ) ) return get_ref_stiffness_transpose_mat(out_element_group, - in_element_group) + in_element_group, + to_real_dtype(dtype)) def _weak_scalar_grad(dcoll, dd_in, vec): @@ -546,12 +566,13 @@ def weak_local_div(dcoll: DiscretizationCollection, *args) -> ArrayOrContainerT: # {{{ Mass operator -def reference_mass_matrix(actx: ArrayContext, out_element_group, in_element_group): +def _reference_mass_matrix(actx: ArrayContext, out_element_group, in_element_group, + dtype: np.dtype[Any]): @keyed_memoize_in( - actx, reference_mass_matrix, + actx, _reference_mass_matrix, lambda out_grp, in_grp: (out_grp.discretization_key(), in_grp.discretization_key())) - def get_ref_mass_mat(out_grp, in_grp): + def get_ref_mass_mat(out_grp, in_grp, dtype: np.dtype[Any]): if out_grp == in_grp: from meshmode.discretization.poly_element import mass_matrix @@ -575,12 +596,15 @@ def get_ref_mass_mat(out_grp, in_grp): actx.from_numpy( np.asarray( np.einsum("j,ik,jk->ij", weights, vand_inv_t, o_vand), - order="C" + order="C", + dtype=dtype ) ) ) - return get_ref_mass_mat(out_element_group, in_element_group) + return get_ref_mass_mat( + out_element_group, in_element_group, + to_real_dtype(dtype)) def _apply_mass_operator( @@ -597,15 +621,17 @@ def _apply_mass_operator( actx = vec.array_context area_elements = area_element(actx, dcoll, dd=dd_in, - _use_geoderiv_connection=actx.supports_nonscalar_broadcasting) + _use_geoderiv_connection=actx.supports_nonscalar_broadcasting, + dtype=vec.entry_dtype) return DOFArray( actx, data=tuple( actx.einsum("ij,ej,ej->ei", - reference_mass_matrix( + _reference_mass_matrix( actx, out_element_group=out_grp, - in_element_group=in_grp + in_element_group=in_grp, + dtype=vec_i.dtype, ), ae_i, vec_i, @@ -659,11 +685,12 @@ def mass(dcoll: DiscretizationCollection, *args) -> ArrayOrContainerT: # {{{ Mass inverse operator -def reference_inverse_mass_matrix(actx: ArrayContext, element_group): +def _reference_inverse_mass_matrix(actx: ArrayContext, element_group, + dtype: np.dtype[Any]): @keyed_memoize_in( - actx, reference_inverse_mass_matrix, + actx, _reference_inverse_mass_matrix, lambda grp: grp.discretization_key()) - def get_ref_inv_mass_mat(grp): + def get_ref_inv_mass_mat(grp, dtype: np.dtype[Any]): from modepy import inverse_mass_matrix basis = grp.basis_obj() @@ -671,12 +698,13 @@ def get_ref_inv_mass_mat(grp): actx.from_numpy( np.asarray( inverse_mass_matrix(basis.functions, grp.unit_nodes), - order="C" + order="C", + dtype=dtype ) ) ) - return get_ref_inv_mass_mat(element_group) + return get_ref_inv_mass_mat(element_group, to_real_dtype(dtype)) def _apply_inverse_mass_operator( @@ -698,12 +726,14 @@ def _apply_inverse_mass_operator( actx = vec.array_context discr = dcoll.discr_from_dd(dd_in) inv_area_elements = 1./area_element(actx, dcoll, dd=dd_in, - _use_geoderiv_connection=actx.supports_nonscalar_broadcasting) + _use_geoderiv_connection=actx.supports_nonscalar_broadcasting, + dtype=vec.entry_dtype) group_data = [] for grp, jac_inv, vec_i in zip(discr.groups, inv_area_elements, vec): - ref_mass_inverse = reference_inverse_mass_matrix(actx, - element_group=grp) + ref_mass_inverse = _reference_inverse_mass_matrix(actx, + element_group=grp, + dtype=vec_i.entry_dtype) group_data.append( # Based on https://arxiv.org/pdf/1608.03836.pdf @@ -763,13 +793,14 @@ def inverse_mass(dcoll: DiscretizationCollection, vec) -> ArrayOrContainerT: # {{{ Face mass operator -def reference_face_mass_matrix( - actx: ArrayContext, face_element_group, vol_element_group, dtype): +def _reference_face_mass_matrix( + actx: ArrayContext, face_element_group, vol_element_group, + dtype: np.dtype[Any]): @keyed_memoize_in( - actx, reference_mass_matrix, + actx, _reference_mass_matrix, lambda face_grp, vol_grp: (face_grp.discretization_key(), vol_grp.discretization_key())) - def get_ref_face_mass_mat(face_grp, vol_grp): + def get_ref_face_mass_mat(face_grp, vol_grp, dtype: np.dtype[Any]): nfaces = vol_grp.mesh_el_group.nfaces assert face_grp.nelements == nfaces * vol_grp.nelements @@ -846,7 +877,8 @@ def get_ref_face_mass_mat(face_grp, vol_grp): return actx.freeze(actx.from_numpy(matrix)) - return get_ref_face_mass_mat(face_element_group, vol_element_group) + return get_ref_face_mass_mat(face_element_group, vol_element_group, + dtype=to_real_dtype(dtype)) def _apply_face_mass_operator(dcoll: DiscretizationCollection, dd, vec): @@ -864,13 +896,14 @@ def _apply_face_mass_operator(dcoll: DiscretizationCollection, dd, vec): assert len(face_discr.groups) == len(volm_discr.groups) surf_area_elements = area_element(actx, dcoll, dd=dd, - _use_geoderiv_connection=actx.supports_nonscalar_broadcasting) + _use_geoderiv_connection=actx.supports_nonscalar_broadcasting, + dtype=dtype) return DOFArray( actx, data=tuple( actx.einsum("ifj,fej,fej->ei", - reference_face_mass_matrix( + _reference_face_mass_matrix( actx, face_element_group=afgrp, vol_element_group=vgrp, diff --git a/grudge/tools.py b/grudge/tools.py index 25018b530..9a12b9ad1 100644 --- a/grudge/tools.py +++ b/grudge/tools.py @@ -22,6 +22,7 @@ THE SOFTWARE. """ +from typing import Any import numpy as np from pytools import levi_civita @@ -104,6 +105,10 @@ def __call__(self, x, y, three_mult=None): # }}} +def to_real_dtype(dtype: np.dtype[Any]) -> np.dtype[Any]: + return np.array(0, dtype=dtype).real.dtype + + def count_subset(subset): from pytools import len_iterable return len_iterable(uc for uc in subset if uc) @@ -199,3 +204,5 @@ def __eq__(self, other): return set(self) == set(other) # }}} + +# vim: foldmethod=marker