jaxsr.sampling

jaxsr.sampling#

Adaptive Sampling Strategies for JAXSR.

Implements various strategies for suggesting new data points to query, inspired by ALAMO’s adaptive sampling approach.

class jaxsr.sampling.SamplingStrategy(value, names=None, *, module=None, qualname=None, type=None, start=1, boundary=None)#

Bases: Enum

Available sampling strategies.

UNCERTAINTY = 'uncertainty'#
ERROR = 'error'#
LEVERAGE = 'leverage'#
GRADIENT = 'gradient'#
SPACE_FILLING = 'space_filling'#
RANDOM = 'random'#
class jaxsr.sampling.SamplingResult(points: Array, scores: Array, strategy: str)#

Bases: object

Result of adaptive sampling suggestion.

Parameters:

  • points (jnp.ndarray) – Suggested points of shape (n_points, n_features).

  • scores (jnp.ndarray) – Acquisition scores for each point.

  • strategy (str) – Strategy used.

points: Array#
scores: Array#
strategy: str#
class jaxsr.sampling.AdaptiveSampler(model: SymbolicRegressor, bounds: list[tuple[float, float]], strategy: str = 'uncertainty', batch_size: int = 5, n_candidates: int = 1000, random_state: int | None = None, discrete_dims: dict[int, list] | None = None)#

Bases: object

Adaptive sampling for iterative model improvement.

Suggests new data points to query based on various strategies that identify regions of high uncertainty or potential improvement.

Parameters:

  • model (SymbolicRegressor) – Fitted model to improve.

  • bounds (list of tuple) – Bounds (lower, upper) for each feature.

  • strategy (str) – Sampling strategy: “uncertainty”, “error”, “leverage”, “gradient”, “space_filling”, “random”.

  • batch_size (int) – Number of points to suggest per call.

  • n_candidates (int) – Number of candidate points to evaluate.

  • random_state (int, optional) – Random seed.

Examples

>>> sampler = AdaptiveSampler(
...     model=model,
...     bounds=[(300, 500), (1, 10)],
...     strategy="uncertainty",
... )
>>> result = sampler.suggest(n_points=5)
>>> X_next = result.points    # shape (5, n_features)
>>> scores = result.scores    # acquisition function values
suggest(n_points: int | None = None, exclude_points: Array | None = None, min_distance: float = 0.01) → SamplingResult#

Suggest new points to query.

Parameters:

  • n_points (int, optional) – Number of points to suggest. Defaults to batch_size.

  • exclude_points (jnp.ndarray, optional) – Points to exclude (e.g., already queried).

  • min_distance (float) – Minimum normalized distance from existing points.

Returns:

result – Suggested points and their scores.

Return type:

SamplingResult

jaxsr.sampling.latin_hypercube_sample(n_samples: int, bounds: list[tuple[float, float]], random_state: int | None = None, discrete_dims: dict[int, list] | None = None) → Array#

Generate Latin Hypercube samples.

Parameters:

  • n_samples (int) – Number of samples.

  • bounds (list of tuple) – Bounds (lower, upper) for each dimension.

  • random_state (int, optional) – Random seed.

  • discrete_dims (dict, optional) – Mapping of dimension index to list of valid discrete values. Continuous dimensions are sampled normally; discrete dimensions are snapped to the nearest valid value.

Returns:

samples – Sample points of shape (n_samples, n_dims).

Return type:

jnp.ndarray

jaxsr.sampling.sobol_sample(n_samples: int, bounds: list[tuple[float, float]], random_state: int | None = None, discrete_dims: dict[int, list] | None = None) → Array#

Generate Sobol sequence samples.

Parameters:

  • n_samples (int) – Number of samples (rounded to power of 2).

  • bounds (list of tuple) – Bounds (lower, upper) for each dimension.

  • random_state (int, optional) – Random seed for scrambling.

  • discrete_dims (dict, optional) – Mapping of dimension index to list of valid discrete values.

Returns:

samples – Sample points.

Return type:

jnp.ndarray

jaxsr.sampling.halton_sample(n_samples: int, bounds: list[tuple[float, float]], random_state: int | None = None, discrete_dims: dict[int, list] | None = None) → Array#

Generate Halton sequence samples.

Parameters:

  • n_samples (int) – Number of samples.

  • bounds (list of tuple) – Bounds (lower, upper) for each dimension.

  • random_state (int, optional) – Random seed for scrambling.

  • discrete_dims (dict, optional) – Mapping of dimension index to list of valid discrete values.

Returns:

samples – Sample points.

Return type:

jnp.ndarray

jaxsr.sampling.grid_sample(n_per_dim: int, bounds: list[tuple[float, float]], discrete_dims: dict[int, list] | None = None) → Array#

Generate grid samples.

Parameters:

  • n_per_dim (int) – Number of samples per dimension (for continuous dims).

  • bounds (list of tuple) – Bounds (lower, upper) for each dimension.

  • discrete_dims (dict, optional) – Mapping of dimension index to list of valid discrete values. Discrete dimensions use their exact values instead of linspace.

Returns:

samples – Sample points.

Return type:

jnp.ndarray

jaxsr.sampling.d_optimal_select(candidates: Array, n_select: int, basis_library, selected_indices: Array, random_state: int | None = None) → Array#

Select D-optimal points from candidates.

Maximizes det(Phi.T @ Phi) for selected design matrix.

Parameters:

  • candidates (jnp.ndarray) – Candidate points of shape (n_candidates, n_features).

  • n_select (int) – Number of points to select.

  • basis_library (BasisLibrary) – Basis function library.

  • selected_indices (jnp.ndarray) – Indices of selected basis functions.

  • random_state (int, optional) – Random seed for initialization.

Returns:

selected – Indices of selected candidates.

Return type:

jnp.ndarray