From a domain to collocation points
A sampler only reads three things from a domain: its bounds, its level set and its mapping. Interior points and boundary points are drawn in two different ways, and both come back in a dict keyed by domain label, the format the residuals read.
Top panel: VolumetricSampler draws points
uniformly in the domain's bounding box, keeps only those with
is_inside(x) (the level set test), then maps the
remaining points through the domain's mapping.
Bottom panel: SurfacicSampler samples the
boundary's parametric domain directly (no rejection needed
there) and pushes the points through the surface
mapping, computing normals from its Jacobian.
1. Rejection sampling against the level set
VolumetricSampler is the building block: it
estimates what fraction of the bounding box the domain
occupies, and then loops to generate the collocation points.
The loop draws a batch uniformly in
domain.bounds, keeps the ones with
is_inside(x), and repeats until enough points
remain. After the loop, they are mapped to physical space.
import jax
from scimba_jax.nonlinear_approximation.integration.monte_carlo import (
VolumetricSampler,
)
# sample() draws points in domain.bounds, keeps those with
# domain.is_inside(x),
# repeats until n points are accepted, then maps them through
# domain.mapping
vsampler = VolumetricSampler(outer)
key, pts = vsampler.sample(jax.random.PRNGKey(0), 2000)
2. Sampling the interior: sample()
DomainSampler wraps one
VolumetricSampler per labelled piece: main
domain and each subdomain (excluding holes).
Holes have no interior sampler, as they are
not part of the collocation set.
Each dict key is exactly the piece's
get_label(), which is the
label_str given when the domain was built. A
suffix is added if label_idx != 0. A domain
built without an explicit label defaults to
"interior", as it corresponds to the case where
only one domain is present. In practice the user will rarely
call DomainSampler.sample() on its own, as
explained in the next section.
from scimba_jax.nonlinear_approximation.integration.monte_carlo import (
DomainSampler,
)
sampler = DomainSampler(outer)
key, samples = sampler.sample(jax.random.PRNGKey(0), 2000)
list(samples.keys())
# ['outer', 'inner']: one label for the main domain and one
# per subdomain
# (holes are not sampled)
samples["outer"]
# (Array of shape (n_outer, 2),), always a 1-tuple
3. TensorizedSampler is actually used in scripts
Real training code doesn't call sample() or
bc_sample() directly. It wraps the
DomainSampler (and any other axis like time or
parameters) in a TensorizedSampler. Sampling
from a TensorizedSampler then involves only
one sample(key, n, n_bc) call, which
merges interior and boundary points into a single flat dict.
The result is the dict that a Projector or a
loss function actually uses.
from scimba_jax.nonlinear_approximation.integration.monte_carlo import (
DomainSampler,
TensorizedSampler,
)
sampler = TensorizedSampler([DomainSampler(outer)], bc=True)
# one call samples the interior and the boundary, merged into
# a single dict
key, sample_dict = sampler.sample(
jax.random.PRNGKey(0), 2000, 2000
)
list(sample_dict.keys())
# ['outer', 'inner', 'boundary']: (points,) for the interior
# labels,
# (points, normals) for the boundary label
4. Boundary groups
By default, every boundary piece belongs to one group,
"boundary" (domain.boundaries = {"boundary": [""]}). It can be split with
set_boundaries_dict(...) to get one key per
edge or per boundary condition. For instance, a domain with
Dirichlet and Neumann boundary conditions groups its
boundaries into "bc N" and "bc D".
Those keys then directly drive both the residual and the
loss weights.
Adapted from view source ↗.
from scimba_jax.domains.meshless_domains.domains_2d import Square2D
from scimba_jax.nonlinear_approximation.integration.monte_carlo import (
DomainSampler,
TensorizedSampler,
)
domain_x = Square2D([[0.0, 1.0], [0.0, 1.0]], is_main_domain=True)
domain_x.set_boundaries_dict(
{
"bc N": ["bc north"], # Neumann on the north edge
"bc D": ["bc east", "bc south", "bc west"], # Dirichlet elsewhere
}
)
sampler = TensorizedSampler([DomainSampler(domain_x)], bc=True)
key, sample_dict = sampler.sample(key, N_COLLOC, N_BC_COLLOC)
# sample_dict.keys() -> ['interior', 'bc N', 'bc D']
weights = {"interior": [1.0], "bc N": [30.0], "bc D": [30.0]}
5. Output shape and type
DomainSampler.sample() and
.bc_sample() each return
(key, dict) with the same convention
(domain/group label, tuple). However, each may have a
different tuple width as the number of collocation points on
boundaries or inner domains differ. Then,
TensorizedSampler merges both into the one flat
dict that the loss function actually uses Stacking more axes
(time, parameters...) extends every per-label tuple in
place, ordered by model_type, and prefixes
initial-condition entries with "ic ".
SAMPLE_TYPE = dict[str, tuple[jnp.ndarray, ...]]
# DomainSampler.sample(key, n) -> {label: (points,)}
# DomainSampler.bc_sample(key, n) -> {group_label: (points,
# normals)}
# TensorizedSampler(...).sample(key, n, n_bc) merges both, for
# the Projector:
# {label: (points,), ..., group_label: (points, normals), ...}
6. Writing a new sampler
For space domains ("x"), the best way to
proceed is to just create a new domain (a level set
+ mapping, see
Domains & meshes). Then, the existing DomainSampler wraps it
for free. This is because
TensorizedSampler requires that "x" sampler to
be a real DomainSampler, which already contains
everything (rejection, etc).
For other domains ("t", "mu", etc.), the
user can just implement
sample(key, n) -> (key, array)
and a JIT-friendly
build_sample_func(n) -> Callable. One
remark: the time ("t") slot must specifically be a
UniformTimeSampler.
from scimba_jax.nonlinear_approximation.integration.monte_carlo_box import (
BoxSampler,
)
class UniformTimeSampler:
def __init__(self, bounds: tuple[float, float]):
self.bounds = bounds
self.base_sampler = BoxSampler([self.bounds])
def sample(self, key, n):
return self.base_sampler.sample(key, n)
def build_sample_func(self, n):
return self.base_sampler.build_sample_func(n)