bojaxns

bojaxns

Subpackages

• gaussian_process_formulation

Submodules

• base
• basic
• common
• experiment
• parameter_space
• service
• utils

Package Contents

class Trial

Bases: pydantic.BaseModel

trial_id: str

create_dt: datetime.datetime

param_values: bojaxns.common.ParamValues

U_value: bojaxns.common.UValue

trial_updates: Dict[str, TrialUpdate]

class TrialUpdate

Bases: pydantic.BaseModel

ref_id: str

measurement_dt: datetime.datetime

objective_measurement: float

class OptimisationExperiment

Bases: pydantic.BaseModel

experiment_id: str

parameter_space: bojaxns.parameter_space.ParameterSpace

trials: Dict[str, Trial]

ensure_parameters_match_space(value, values)

class NewExperimentRequest

Bases: pydantic.BaseModel

parameter_space: bojaxns.parameter_space.ParameterSpace

init_explore_size: conint(ge=1)

class AbstractAcquisition

A class that represents any acquisition function. All acquisition functions take a point in the U-domain and returns a metric that gives a proxy as to how valuable it would be to try that point. All acquisition values only make sense relatively.

abstract __call__(u_star)

Parameters:

u_star (jax.numpy.ndarray) –

class MarginalisedAcquisitionFunction(key, ns_results, acquisition_factory, S)

Bases: AbstractAcquisition

Class that represents a marginalisation of an acquisition function over samples.

Parameters:

• key (jax.random.PRNGKey) –

• ns_results (MarginalisationData) –

• acquisition_factory (AcquisitionFactory) –

• S (int) –

__call__(u_star)

Parameters:

u_star (jax.numpy.ndarray) –

class MarginalisationData

Bases: NamedTuple

samples: Dict[str, chex.Array]

log_dp_mean: chex.Array

class OptimisationExperiment

Bases: pydantic.BaseModel

experiment_id: str

parameter_space: bojaxns.parameter_space.ParameterSpace

trials: Dict[str, Trial]

ensure_parameters_match_space(value, values)

class GaussianProcessData

Bases: NamedTuple

U: jax.numpy.ndarray

Y: jax.numpy.ndarray

Y_var: jax.numpy.ndarray

sample_size: jax.numpy.ndarray

class GaussianProcessConditionalPredictiveFactory(data)

Bases: bojaxns.base.ConditionalPredictiveFactory

Parameters:

data (GaussianProcessData) –

ndims()

build_prior_model()

psd_kernels()

Return type:

List[Type[tensorflow_probability.substrates.jax.math.psd_kernels.PositiveSemidefiniteKernel]]

__call__(**samples)

Return type:

GaussianProcessConditionalPredictive

class ExpectedImprovementAcquisitionFactory(conditional_predictive_factory)

Bases: bojaxns.base.AcquisitionFactory

Parameters:

conditional_predictive_factory (GaussianProcessConditionalPredictiveFactory) –

__call__(**sample)

Return type:

bojaxns.base.AbstractAcquisition

class TopTwoAcquisitionFactory(conditional_predictive_factory, u1)

Bases: bojaxns.base.AcquisitionFactory

Parameters:

• conditional_predictive_factory (GaussianProcessConditionalPredictiveFactory) –

• u1 (jax.numpy.ndarray) –

__call__(**sample)

Return type:

bojaxns.base.AbstractAcquisition

run_multi_lookahead(rng_key, data, ns_results, batch_size, max_depth, num_actions, num_simulations, S)

Parameters:

• rng_key (chex.PRNGKey) –

• data (bojaxns.gaussian_process_formulation.distribution_math.GaussianProcessData) –

• ns_results (bojaxns.base.MarginalisationData) –

• batch_size (int) –

• max_depth (int) –

• num_actions (int) –

• num_simulations (int) –

• S (int) –

Return type:

Tuple[chex.Array, mctx.PolicyOutput[mctx.GumbelMuZeroExtraData]]

convert_tree_to_graph(tree, action_labels=None, batch_index=0)

Converts a search tree into a Graphviz graph.

Parameters:

• tree (mctx.Tree) – A Tree containing a batch of search data.

• action_labels (Optional[Sequence[str]]) – Optional labels for edges, defaults to the action index.

• batch_index (int) – Index of the batch element to plot.

Returns:

A Graphviz graph representation of tree.

tfpb

class BayesianOptimiser(experiment, num_parallel_solvers=1, beta=0.5, S=512)

Parameters:

• experiment (bojaxns.experiment.OptimisationExperiment) –

• num_parallel_solvers (int) –

• beta (float) –

• S (int) –

posterior_solve(key)

Parameters:

key (chex.PRNGKey) –

Return type:

jaxns.internals.types.NestedSamplerResults

search_U_top1(key, ns_results, batch_size, num_search)

Parameters:

• key (chex.PRNGKey) –

• ns_results (bojaxns.base.MarginalisationData) –

• batch_size (int) –

• num_search (int) –

search_U_top2(key, ns_results, u1, batch_size, num_search)

Parameters:

• key (chex.PRNGKey) –

• ns_results (bojaxns.base.MarginalisationData) –

• u1 (jax.numpy.ndarray) –

• batch_size (int) –

• num_search (int) –

choose_next_U_toptwo(key, batch_size, num_search)

Parameters:

• key (chex.PRNGKey) –

• batch_size (int) –

• num_search (int) –

choose_next_U_multistep(key, batch_size, max_depth, num_simulations, branch_factor)

Parameters:

• key (chex.PRNGKey) –

• batch_size (int) –

• max_depth (int) –

• num_simulations (int) –

• branch_factor (int) –

class ContinuousPrior

Bases: pydantic.BaseModel

type: Literal[continuous_prior] = 'continuous_prior'

lower: float

upper: float

mode: float

uncert: confloat(gt=0.0)

class IntegerPrior

Bases: pydantic.BaseModel

type: Literal[integer_prior] = 'integer_prior'

lower: int

upper: int

mode: float

uncert: confloat(gt=0.0)

class CategoricalPrior

Bases: pydantic.BaseModel

type: Literal[categorical_prior] = 'categorical_prior'

probs: List[confloat(ge=0.0)]

class Parameter

Bases: pydantic.BaseModel

name: str

prior: ParamPrior

class ParameterSpace

Bases: pydantic.BaseModel

parameters: List[Parameter]

unique_parameters(value)

build_prior_model(parameter_space)

Constructs a prior model given the parameter space.

Parameters:

parameter_space (ParameterSpace) –

Returns:

exception InvalidTrial

Bases: Exception

Common base class for all non-exit exceptions.

Initialize self. See help(type(self)) for accurate signature.

class BayesianOptimisation(experiment)

Parameters:

experiment (bojaxns.experiment.OptimisationExperiment) –

property experiment

classmethod create_new_experiment(new_experiment)

Parameters:

new_experiment (bojaxns.experiment.NewExperimentRequest) –

Return type:

BayesianOptimisation

add_trial_from_data(key, param_values)

Parameters:

• key (jax._src.random.PRNGKey) –

• param_values (bojaxns.common.ParamValues) –

Return type:

str

create_new_trial(key, random_explore=False, beta=0.5)

Parameters:

• key (jax._src.random.PRNGKey) –

• random_explore (bool) –

• beta (float) –

Return type:

str

get_trial(trial_id)

Parameters:

trial_id (str) –

Return type:

bojaxns.experiment.Trial

delete_trial(trial_id)

Parameters:

trial_id (str) –

post_measurement(trial_id, trial_update)

Parameters:

• trial_id (str) –

• trial_update (bojaxns.experiment.TrialUpdate) –

trial_size(trial_id)

Parameters:

trial_id (str) –

visualise(main_color='#7e97bf', grid_color='#969396')

Constructs a visual breakdown of condition

Parameters:

• main_color – color of main axes

• grid_color – color of grid

Returns:

a pylab Figure

Raises:

NotEnoughData if not enough to compute a breakdown –

Return type:

pylab.Figure

latin_hypercube(seed, num_samples, num_dim)

Sample from the latin-hypercube defined as the continuous analog of the discrete latin-hypercube. That is, if you partition each dimension into num_samples equal volume intervals then there is (conditionally) exactly one point in each interval. We guarantee that uniformity by randomly assigning the permutation of each dimension. The degree of randomness is controlled by cube_scale. A value of 0 places the sample at the center of the grid point, and a value of 1 places the value randomly inside the grid-cell.

Parameters:

• key – PRNG key

• num_samples (int) – number of samples in total to draw

• num_dim (int) – number of dimensions in each sample

• cube_scale – The scale of randomness, in (0,1).

• seed (int) –

Returns:

latin-hypercube samples of shape [num_samples, num_dim]

build_example(model)

Parameters:

model (Type[_T]) –

Return type:

_T

current_utc()

Return type:

datetime.datetime