Sampling

popinn.UniformCollocationSampler

Resample one residual term's collocation points, typically the interior coordinates where the PDE is enforced.

Source code in src/popinn/sampling.py

class UniformCollocationSampler(AbstractSampler):
    """Resample one residual term's collocation points, typically the interior coordinates where the PDE is enforced."""

    data: eqx.Module
    bounds: tuple = eqx.field(static=True)
    counts: tuple = eqx.field(static=True)
    field: str = eqx.field(static=True)

    def __init__(
        self,
        data: eqx.Module,
        bounds: Sequence[tuple[float, float]],
        num_samples: int | Sequence[int] | None = None,
        field: str = "pde_coords",
    ):
        """Configure the sampler.

        Args:
            data (eqx.Module): The training data container to resample from.
            bounds (Sequence[tuple[float, float]]): One (min, max) pair per
                coordinate axis of `data.<field>`. Its length must equal the
                number of axes in that field. Example: for `data.pde_coords`
                with two coordinate axes `x` and `t`, then
                `bounds = ((xmin,xmax), (tmin,tmax))`.
            num_samples (int | Sequence[int] | None): Number of collocation
                points to sample per axis. An int applies to every axis; a sequence
                gives a per-axis count. None (default) reuses the current
                per-axis sizes of `data.<field>`. NOTE: if the resolved counts
                differ from the sizes in the batch first passed to training,
                the jitted step recompiles on the first resample.
            field (str): Name of the coordinate field to resample. Defaults to
                'pde_coords', i.e. the interior collocation points.

        Raises:
            ValueError: If `bounds` or a sequence `num_samples` does not
                have one entry per coordinate axis of `data.<field>`.
        """
        coords = getattr(data, field)
        n_axes = len(coords)

        if len(bounds) != n_axes:
            raise ValueError(f"bounds has {len(bounds)} entries but data.{field} has {n_axes} coordinate axes; expected one `(min, max)` per axis.")

        if num_samples is None:
            counts = tuple(c.shape[0] for c in coords)
        elif isinstance(num_samples, int):
            counts = (num_samples,) * n_axes
        else:
            if len(num_samples) != n_axes:
                raise ValueError(f"num_samples has {len(num_samples)} entries but data.{field} has {n_axes} coordinate axes.")
            counts = tuple(num_samples)

        self.data = data
        self.bounds = tuple(bounds)
        self.counts = counts
        self.field = field

    def __call__(self, key: PRNGKeyArray) -> eqx.Module:
        """Draw a new batch with `field` resampled uniformly.

        Args:
            key (PRNGKeyArray): PRNG key.

        Returns:
            (eqx.Module): A new batch with `field` uniformly re-sampled within the specified `bounds` and
                all other fields (including `aux`) unchanged.
        """
        keys = jr.split(key, len(self.counts))
        new_coords = tuple(
            jr.uniform(keys[i], (self.counts[i],), minval=self.bounds[i][0], maxval=self.bounds[i][1]) for i in range(len(self.counts))
        )
        # Replace only `field`; tree_at returns a new module, leaving the
        # reference `data` (and its aux / other fields) untouched.
        return eqx.tree_at(lambda b: getattr(b, self.field), self.data, new_coords)

__init__(data, bounds, num_samples=None, field='pde_coords')

Configure the sampler.

Parameters:

Name	Type	Description	Default
data	equinox.Module	The training data container to resample from.	required
bounds	collections.abc.Sequence[tuple[float, float]]	One (min, max) pair per coordinate axis of data.<field>. Its length must equal the number of axes in that field. Example: for data.pde_coords with two coordinate axes x and t, then bounds = ((xmin,xmax), (tmin,tmax)).	required
num_samples	int | collections.abc.Sequence[int] | None	Number of collocation points to sample per axis. An int applies to every axis; a sequence gives a per-axis count. None (default) reuses the current per-axis sizes of data.<field>. NOTE: if the resolved counts differ from the sizes in the batch first passed to training, the jitted step recompiles on the first resample.	None
field	str	Name of the coordinate field to resample. Defaults to 'pde_coords', i.e. the interior collocation points.	'pde_coords'

Raises:

Type	Description
ValueError	If bounds or a sequence num_samples does not have one entry per coordinate axis of data.<field>.

Source code in src/popinn/sampling.py

def __init__(
    self,
    data: eqx.Module,
    bounds: Sequence[tuple[float, float]],
    num_samples: int | Sequence[int] | None = None,
    field: str = "pde_coords",
):
    """Configure the sampler.

    Args:
        data (eqx.Module): The training data container to resample from.
        bounds (Sequence[tuple[float, float]]): One (min, max) pair per
            coordinate axis of `data.<field>`. Its length must equal the
            number of axes in that field. Example: for `data.pde_coords`
            with two coordinate axes `x` and `t`, then
            `bounds = ((xmin,xmax), (tmin,tmax))`.
        num_samples (int | Sequence[int] | None): Number of collocation
            points to sample per axis. An int applies to every axis; a sequence
            gives a per-axis count. None (default) reuses the current
            per-axis sizes of `data.<field>`. NOTE: if the resolved counts
            differ from the sizes in the batch first passed to training,
            the jitted step recompiles on the first resample.
        field (str): Name of the coordinate field to resample. Defaults to
            'pde_coords', i.e. the interior collocation points.

    Raises:
        ValueError: If `bounds` or a sequence `num_samples` does not
            have one entry per coordinate axis of `data.<field>`.
    """
    coords = getattr(data, field)
    n_axes = len(coords)

    if len(bounds) != n_axes:
        raise ValueError(f"bounds has {len(bounds)} entries but data.{field} has {n_axes} coordinate axes; expected one `(min, max)` per axis.")

    if num_samples is None:
        counts = tuple(c.shape[0] for c in coords)
    elif isinstance(num_samples, int):
        counts = (num_samples,) * n_axes
    else:
        if len(num_samples) != n_axes:
            raise ValueError(f"num_samples has {len(num_samples)} entries but data.{field} has {n_axes} coordinate axes.")
        counts = tuple(num_samples)

    self.data = data
    self.bounds = tuple(bounds)
    self.counts = counts
    self.field = field

__call__(key)

Draw a new batch with field resampled uniformly.

Parameters:

Name	Type	Description	Default
key	typing.Union	PRNG key.	required

Returns:

Type	Description
equinox.Module	A new batch with field uniformly re-sampled within the specified bounds and all other fields (including aux) unchanged.

Source code in src/popinn/sampling.py

def __call__(self, key: PRNGKeyArray) -> eqx.Module:
    """Draw a new batch with `field` resampled uniformly.

    Args:
        key (PRNGKeyArray): PRNG key.

    Returns:
        (eqx.Module): A new batch with `field` uniformly re-sampled within the specified `bounds` and
            all other fields (including `aux`) unchanged.
    """
    keys = jr.split(key, len(self.counts))
    new_coords = tuple(
        jr.uniform(keys[i], (self.counts[i],), minval=self.bounds[i][0], maxval=self.bounds[i][1]) for i in range(len(self.counts))
    )
    # Replace only `field`; tree_at returns a new module, leaving the
    # reference `data` (and its aux / other fields) untouched.
    return eqx.tree_at(lambda b: getattr(b, self.field), self.data, new_coords)