Grid evaluation

Utilities for evaluating a per-point function over a cartesian product of grid points.

Conventions

• The function being evaluated must have the signature: fn(x1, x2, ..., aux) -> scalar, where coordinates x1, x2, ... are scalar entries and aux is a tuple of auxiliary components (PDE parameters, discretized source functions, initial conditions, etc). Components of aux may be scalars or arrays of any (possibly different) shapes.

• The function is evaluated over a cartesian product of the 1D coordinate arrays and the leading axis of each aux element. The trailing axes of each aux element are treated as the per-call input for that component:

  # scalar input
  a = jnp.linspace(0,1,10)
  
  # function input
  x = jnp.linspace(0,10,20) # sensor coordinate
  b = jnp.linspace(0,100,15) # function parameter
  func = lambda x, b: b * x**2
  f = jax.vmap(func, in_axes = (None, 0))(x, b) # shape = (15, 20)
  
  #the aux tuple
  aux = (a, f) # per-call shape: a = (1,), f = (20,)
  

• An element of aux can be a tuple of arrays that share their leading axis length. This is useful when the function requires a parameter that a functional auxiliary input was calculated from. In the above example, this would be the b array of parameter values that were used to calculate f:

  aux = (a, (f, b)) # per-call shape: a = (1,), f = (20,), b = (1,)
  
  residual_fn(model):
      r(x, aux):
          a, (f, b) = aux
          # du/dx = b*u
          res = model.D(0)(x, (a, f)) - b * model(x, (a, f)) # b isn't a model input, but it is required to calculate the PDE
          return res
      return r
  

  Tuple elements of aux are zipped across their shared axis to preserve their correlated structure.

• Each function returns the same output layout (*reversed(aux axis lengths), *reversed(coord lengths)). For the example above, aux = (a, (f, b)) and a single collocation coordinate x = jnp.linspace(0,10,100), the output has a shape (15, 10, 100).

---

Memory dials

By default, the cartesian product evaluation is performed by nested jax.vmap calls, which materializes the full grid of evaluation points in memory at once. Each utility has the ability to trade jax.vmap for jax.lax.map on a single axis in order to execute the computation in batches to be more memory efficient:

• outer_batch_size in eval_grid: batches the outermost grid axis only -- axis 0 of aux[-1] when aux is non-empty, otherwise the last coordinate axis (so models with no auxiliary inputs can batch over a coordinate).

• eval_grid_flat_aux: enumerate the aux combinations and stack them into an array with shape (N_combinations, N_aux_elements), then batch over the first axis of this array.

---

popinn.eval_grid(fn, coords_1d, aux=(), outer_batch_size=None)

Evaluate fn on the cartesian product of coordinate and auxiliary axes.

Parameters:

Name	Type	Description	Default
fn	collections.abc.Callable	Per-point function with signature fn(*coords, aux) -> scalar, with aux a tuple of auxiliary inputs.	required
coords_1d	tuple[jax.Array, ...]	Tuple of 1-D coordinate arrays.	required
aux	tuple	Tuple of auxiliary components. Each TOP-LEVEL element is one grid axis (its leading axis); trailing dims are the per-call input for that component. Independent top-level elements are crossed: aux = (a, b) # a.shape=(Na,), b.shape=(Nb,) -> Na x Nb grid A top-level element that is itself a tuple is zipped, not crossed: all its leaves share their leading axis. Use this for correlated quantities, e.g. a DeepONet whose initial conditions were generated from PDE parameters (IC[k] from param[k]): # IC.shape = (N_IC, N_IC_PTS), param.shape = (N_IC,) aux = ((IC, param),) # ONE axis of length N_IC, zipped so per call the function receives the pair (IC[k], param[k]).	()
outer_batch_size	int | None	If not None, the OUTERMOST grid axis is batched with jax.lax.map, which dispatches jax.vmap on each batch. The outermost axis is axis 0 of aux[-1] when aux is non-empty and the last coordinate axis (axis 0 of coords_1d[-1]) when aux is empty. Note that batching an axis shrinks the inner vectorized batch to the product of the remaining axes, so keep the largest coordinate axis as the first element and the largest aux axis as the last element.	None

Returns:

Type	Description
jax.Array	Array of shape (*reversed(aux lens), *reversed(coord lens)).

Source code in src/popinn/eval.py

def eval_grid(
    fn: Callable,
    coords_1d: tuple[Array, ...],
    aux: tuple = (),
    outer_batch_size: int | None = None,
) -> Array:
    """Evaluate `fn` on the cartesian product of coordinate and auxiliary axes.

    Args:
        fn (Callable): Per-point function with signature `fn(*coords, aux) -> scalar`, with `aux` a
            tuple of auxiliary inputs.
        coords_1d (tuple[Array, ...]): Tuple of 1-D coordinate arrays.
        aux (tuple): Tuple of auxiliary components. Each TOP-LEVEL
            element is one grid axis (its leading axis); trailing dims are the
            per-call input for that component. Independent top-level elements are
            crossed:

            ```python
                aux = (a, b)        # a.shape=(Na,), b.shape=(Nb,) -> Na x Nb grid
            ```

            A top-level element that is itself a tuple is zipped, not
            crossed: all its leaves share their leading axis. Use this for
            correlated quantities, e.g. a DeepONet whose initial conditions
            were generated from PDE parameters (IC[k] from param[k]):

            ```python
                # IC.shape = (N_IC, N_IC_PTS), param.shape = (N_IC,)
                aux = ((IC, param),)   # ONE axis of length N_IC, zipped
            ```

            so per call the function receives the pair `(IC[k], param[k])`.
        outer_batch_size (int | None): If not `None`, the OUTERMOST
            grid axis is batched with `jax.lax.map`, which dispatches `jax.vmap`
            on each batch. The outermost axis is axis 0 of `aux[-1]` when `aux` is non-empty
            and the last coordinate axis (axis 0 of `coords_1d[-1]`) when
            `aux` is empty. Note that batching an axis shrinks the inner
            vectorized batch to the product of the remaining
            axes, so keep the largest coordinate axis as the first element and the largest
             `aux` axis as the last element.

    Returns:
        (Array): Array of shape `(*reversed(aux lens), *reversed(coord lens))`.
    """
    n_c, n_a = len(coords_1d), len(aux)
    n_total = n_c + n_a

    g = fn
    n_vmapped = n_total if outer_batch_size is None else n_total - 1
    for i in range(n_vmapped):
        g = _vmap_axis(g, n_c, n_a, i)

    if outer_batch_size is None:
        return g(*coords_1d, aux)

    # Chunked outermost axis: sequentially map over its slices,
    # `outer_batch_size` at a time. lax.map(batch_size=k) vmaps within
    # each chunk and scans across chunks, so only one chunk's activations
    # exist at any moment.
    if n_a > 0:
        # Outermost axis is the last aux component.
        body = lambda last: g(*coords_1d, (*aux[:-1], last))
        return jax.lax.map(body, aux[-1], batch_size=outer_batch_size)

    # No aux components: the outermost grid axis is the last coordinate.
    body = lambda last_c: g(*coords_1d[:-1], last_c, aux)
    return jax.lax.map(body, coords_1d[-1], batch_size=outer_batch_size)

popinn.eval_grid_flat_aux(fn, coords_1d, aux=(), batch_size=None)

Like eval_grid, but enumerates the aux combinations by integer index and stacks them into an array with shape (N_combinations, N_aux_elements). The computation is then performed on batch_size chunks of the combinations.

If you don't intend to utilize the batching scheme, use eval_grid instead.

Use this when even one full outermost-axis batch in eval_grid is too big, or when the aux axes are very uneven (e.g. Na=2, Nb=10_000).

Args: fn (Callable): Per-point function with signature fn(*coords, aux) -> scalar, where aux is a tuple of auxiliary inputs. coords_1d (tuple[Array, ...]): Tuple of 1-D coordinate arrays. aux (tuple): Tuple of auxiliary components. Each TOP-LEVEL element is one grid axis (its leading axis); trailing dims are the per-call input for that component. Independent top-level elements are crossed:

     ```python
         aux = (a, b)        # a.shape=(Na,), b.shape=(Nb,) -> Na x Nb grid
     ```

     A top-level element that is itself a tuple is zipped, not
     crossed: all its leaves share their leading axis. Use this for
     correlated quantities, e.g. a DeepONet whose initial conditions
     were generated from PDE parameters (IC[k] from param[k]):

     ```python
         # IC.shape = (N_IC, N_IC_PTS), param.shape = (N_IC,)
         aux = ((IC, param),)   # ONE axis of length N_IC, zipped
     ```

     so per call the function receives the pair `(IC[k], param[k])`.
 batch_size (int | None): The batch size for `jax.lax.map`.

Returns: (Array): Array of shape (*reversed(aux lens), *reversed(coord lens)).

Source code in src/popinn/eval.py

def eval_grid_flat_aux(
    fn: Callable,
    coords_1d: tuple[Array, ...],
    aux: tuple = (),
    batch_size: int | None = None,
) -> Array:
    """Like `eval_grid`, but enumerates the `aux` combinations by integer index and stacks them into an array with
     shape `(N_combinations, N_aux_elements)`. The computation is then performed on `batch_size` chunks of the combinations.

    If you don't intend to utilize the batching scheme, use `eval_grid` instead.

     Use this when even one full outermost-axis batch in `eval_grid` is too big, or when
     the aux axes are very uneven (e.g. `Na=2, Nb=10_000`).

     Args:
         fn (Callable): Per-point function with signature `fn(*coords, aux) -> scalar`, where `aux` is a
             tuple of auxiliary inputs.
         coords_1d (tuple[Array, ...]): Tuple of 1-D coordinate arrays.
         aux (tuple): Tuple of auxiliary components. Each TOP-LEVEL
             element is one grid axis (its leading axis); trailing dims are the
             per-call input for that component. Independent top-level elements are
             crossed:

             ```python
                 aux = (a, b)        # a.shape=(Na,), b.shape=(Nb,) -> Na x Nb grid
             ```

             A top-level element that is itself a tuple is zipped, not
             crossed: all its leaves share their leading axis. Use this for
             correlated quantities, e.g. a DeepONet whose initial conditions
             were generated from PDE parameters (IC[k] from param[k]):

             ```python
                 # IC.shape = (N_IC, N_IC_PTS), param.shape = (N_IC,)
                 aux = ((IC, param),)   # ONE axis of length N_IC, zipped
             ```

             so per call the function receives the pair `(IC[k], param[k])`.
         batch_size (int | None): The batch size for `jax.lax.map`.
     Returns:
         (Array): Array of shape `(*reversed(aux lens), *reversed(coord lens))`.
    """
    n_c, n_a = len(coords_1d), len(aux)

    # Zero aux axes -> exactly one (empty) combination: chunking over
    # combinations is a no-op, so this is just the coord-grid evaluation.
    if n_a == 0:
        return eval_grid(fn, coords_1d, aux)

    # Axis length of a component = leading-axis length of its first leaf
    # (for grouped components, all leaves share it by convention).
    aux_lens = tuple(jax.tree_util.tree_leaves(comp)[0].shape[0] for comp in aux)

    # Vectorize the coordinate axes only; the aux tuple stays per-call.
    g = fn
    for i in range(n_c):
        g = _vmap_axis(g, n_c, n_a, i)
    # g now maps (coord_vecs..., aux_slices_tuple) -> coord grid

    # Enumerate combinations in 'ij' (first-component-major) index order.
    idx_mesh = jnp.meshgrid(*(jnp.arange(n) for n in aux_lens), indexing="ij")
    idx = jnp.stack([m.ravel() for m in idx_mesh], axis=-1)  # (N_combos, n_a)

    # Gather one slice per component inside the loop body (dynamic gather
    # is fine under lax.map). tree_map handles plain arrays and grouped
    # (sub-tuple) components uniformly: every leaf is indexed at axis 0
    # by the same index, preserving zip pairing within a group.
    def body(ix):
        slices = tuple(jax.tree_util.tree_map(lambda leaf, k=k: leaf[ix[k]], aux[k]) for k in range(n_a))
        return g(*coords_1d, slices)

    out = jax.lax.map(body, idx, batch_size=batch_size)
    # out: (N_combos, *reversed(coord lens)), combos first-component-major.

    # Match eval_grid's layout: unflatten, then reverse the aux axes so the
    # LAST component leads, as in the nested-vmap convention.
    out = out.reshape(*aux_lens, *out.shape[1:])
    perm = tuple(reversed(range(n_a))) + tuple(range(n_a, out.ndim))
    return jnp.transpose(out, perm)