jexplore.sampling.space

This module provides the definitions for the base target space classes. The prototyping class jexplore.sampling.space.Space and the class defining a simple box space jexplore.sampling.space.Box

Classes

Space	Markov sampling target space prototype class.
Box	Simple rectangular box space.

Module Contents

class Space

Bases: Protocol

Markov sampling target space prototype class. Child space classes should implement an inspace method and a wrap method.

Parameters:

dim – dimension of the target space.

dim: int

Dimension of the target space

inspace(points)

Check which of a set of point lay in the defined space.

Parameters:

points (jax.Array) – set of points to be checked, it can have the state shape \((N_{chains}, D)\) or the samples shape \((N_{chains}, D, S)\).

Returns:

a boolean mask of shape \((N_{chains})\) or \((N_{chains}, S)\) selecting points that lay in the defined space.

Return type:

jax.Array

wrap(points)

folds the wrapped dimensions of the target space.

Parameters:

points (jax.Array) – set of points to be processed, it can have the state shape \((N_{chains}, D)\) or the samples shape \((N_{chains}, D, S)\).

Return type:

jax.Array

return: the same set of points with the wrapped dimensions folded.

class Box(dim=None, size=jnp.inf, box=None, wrapped=None)

Bases: Space

Simple rectangular box space.

Parameters:

• dim (int | None) – dimension of the box. Only used if the box parameter is not provided.

• size (float) – size of the box (assumed having equal size in all dimensions) only used if the box parameter is not provided. Default: infinity.

• box (list[list[float]] | None) – list of lists defining the segments boundaries for each box dimensions. If not provided a box with dim dimensions of equal size size is considered.

• wrapped (list[int] | None) – list of periodic dimensions indexes. These dimensions will be considererd unbound and the corresponding box intervals will be interpreted as principal domain intervals. Default: empty list.

bounds: jax.Array

Box bounds

wrap_dims: list[int]

Wrapped dimensions indexes

wrap_domain: jax.Array

wrapped dimensions principal domain

dim

Dimension of the target space

inspace(points)

Check which of a set of point lay in the defined space.

Parameters:

points (jax.Array) – set of points to be checked, it can have the state shape \((N_{chains}, D)\) or the samples shape \((N_{chains}, D, S)\).

Returns:

a boolean mask of shape \((N_{chains})\) or \((N_{chains}, S)\) selecting points that lay in the defined space.

Return type:

jax.Array

wrap(points)

folds the wrapped dimensions of the target space.

Parameters:

points (jax.Array) – set of points to be processed, it can have the state shape \((N_{chains}, D)\) or the samples shape \((N_{chains}, D, S)\).

Return type:

jax.Array

return: the same set of points with the wrapped dimensions folded.