grids.mixins.mathops.DenseMathOpsMixin.compute_function_on_grid#

DenseMathOpsMixin.compute_function_on_grid(func: Callable, /, result_shape: Sequence[int] | None = None, out: _Buffer | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | complex | bytes | str | _NestedSequence[complex | bytes | str] | None = None, *, in_chunks: bool = False, output_axes: Sequence[str] | None = None, func_type: Literal['all', 'axes'] = 'axes', pbar: bool = True, pbar_kwargs: dict | None = None, **kwargs)[source]#

Evaluate a coordinate-based function over the grid domain.

This method computes the values of a user-supplied callable func at each point on the grid spanned by output_axes, using the physical (coordinate system-defined) coordinates at those points. It supports both global (in-memory) and streaming (chunked) evaluation strategies.

This is useful for constructing scalar, vector, or tensor fields defined analytically in terms of coordinates, and can be applied efficiently even for large grids.

Parameters:

  • func (callable) –

    A Python function that returns values defined over physical coordinates. Its signature must match the coordinate mode:

    • If func_type="axes" (default), the function takes only the coordinates of output_axes (in canonical coordinate order).

    • If func_type="all", the function must take all coordinate axes of the system in canonical order.

    The return value should be an array with shape result_shape per grid point.

  • result_shape (sequence of int, optional) – The shape of the result at each grid point (e.g., () for scalar fields, (3,) for 3-vectors). Defaults to () if unspecified.

  • out (array-like, optional) – An optional buffer in which to store the result. This can be used to reduce memory usage when performing computations. The shape of out must be compliant with broadcasting rules (see Notes). out may be a buffer or any other array-like object.

  • in_chunks (bool, optional) – Whether to perform the computation in chunks. This can help reduce memory usage during the operation but will increase runtime due to increased computational load. If input buffers are all fully-loaded into memory, chunked performance will only marginally improve; however, if buffers are lazy loaded, then chunked operations will significantly improve efficiency. Defaults to False.

  • output_axes (list of str, optional) –

    The axes of the coordinate system over which the result should span. By default, output_axes is the same as field_axes and the output field matches the span of the input field.

    This argument may be specified to expand the number of axes onto which the output field is computed. output_axes must be a superset of the field_axes.

  • pbar (bool, optional) – Whether to show a progress bar during chunked evaluation.

  • pbar_kwargs (dict, optional) – Extra keyword arguments passed to tqdm when pbar=True.

  • func_type ({"axes", "all"}, optional) – Specifies whether func accepts all coordinate axes or only those listed in output_axes. By default, "axes".

  • **kwargs – Additional keyword arguments forwarded to func.

Returns:

A NumPy array with shape (grid_shape over output_axes) + result_shape, containing the evaluated function values at each grid point.

Return type:

np.ndarray

Notes

This method evaluates the function:

\[f(x^1, x^2, ..., x^n)\]

where \(x^i\) are the physical coordinates (e.g., r, theta, z, etc.) defined by the coordinate system.

If in_chunks=True, the grid is traversed in blocks (including 1-cell ghost zones), and the function is evaluated independently on each block. This enables out-of-core and lazy buffer evaluations.

This method is commonly used to:

  • Construct analytic scalar/vector/tensor fields on structured grids

  • Convert closed-form physics models into discrete field data

  • Generate test fields for numerical method validation