Package: boax
Boax: A Bayesian Optimization library for JAX.Overview|Installation|Getting Started|DocumentationBoax is currently in early alpha and under active development!OverviewBoax is a composable library of core components for Bayesian Optimization
that isdesigned for flexibility. It comes with a low-level interfaces for:Core capabilities(boax.core):Common DistributionsMonte-Carlo SamplersFitting a surrogate model to data(boax.prediction):Kernels FunctionsLikelihood FunctionsMean FunctionsModel FunctionsObjective FunctionsConstructing and optimizing acquisition functions(boax.optimization):Acquisition FunctionsConstraint FunctionsOptimizer FunctionsInstallationYou can install the latest released version of Boax from PyPI via:pipinstallboaxor you can install the latest development version from GitHub:pipinstallgit+https://github.com/Lando-L/boax.gitGetting StartedHere is a quick start example of the two main compoments that form the Bayesian optimization loop.
For more details check out thedocs.Create a synthetic dataset.fromjaximportconfig# Double precision is highly recommended.config.update("jax_enable_x64",True)fromjaximportjitfromjaximportlaxfromjaximportnnfromjaximportnumpyasjnpfromjaximportrandomfromjaximportvalue_and_gradimportoptaxfromboaximportprediction,optimizationfromboax.coreimportdistributions,samplersfromboax.predictionimportkernels,likelihoods,means,models,objectivesfromboax.optimizationimportacquisitions,optimizersbounds=jnp.array([[0.0,1.0]])defobjective(x):return1-jnp.linalg.norm(x-0.5)data_key,sampler_key,optimizer_key=random.split(random.key(0),3)x_train=random.uniform(random.fold_in(data_key,0),minval=bounds[:,0],maxval=bounds[:,1],shape=(10,1))y_train=objective(x_train)+0.1*random.normal(random.fold_in(data_key,1),shape=(10,))Fit a Gaussian Process surrogate model to the training dataset.params={'amplitude':jnp.zeros(()),'length_scale':jnp.zeros(()),'noise':jnp.zeros(()),}adam=optax.adam(0.01)deffit(x_train,y_train):defmodel(params):returnmodels.outcome_transformed(models.gaussian_process(means.zero(),kernels.scaled(kernels.rbf(params['amplitude']),params['length_scale'],),),likelihoods.gaussian(params['noise']),)defobjective(params):returnobjectives.negative_log_likelihood(distributions.multivariate_normal.logpdf)defprojection(params):return{'amplitude':nn.softplus(params['amplitude']),'length_scale':nn.softplus(params['length_scale']),'noise':nn.softplus(params['noise'])+1e-4,}defstep(state,iteration):loss_fn=prediction.construct(model,objective,projection)loss,grads=value_and_grad(loss_fn)(state[0],x_train,y_train)updates,opt_state=adam.update(grads,state[1])params=optax.apply_updates(state[0],updates)return(params,opt_state),loss(next_params,_),_=lax.scan(jit(step),(params,adam.init(params)),jnp.arange(500))returnprojection(next_params)Construct and optimize an UCB acquisition function.x0=jnp.reshape(samplers.halton_uniform(distributions.uniform.uniform(bounds[:,0],bounds[:,1]))(sampler_key,100,),(100,1,-1))defoptimize(x_train,y_train):defmodel(params):returnmodels.outcome_transformed(models.gaussian_process_regression(means.zero(),kernels.scaled(kernels.rbf(params['amplitude']),params['length_scale']))(x_train,y_train,),likelihoods.gaussian(params['noise']),distributions.multivariate_normal.as_normal,)foriinrange(10):params=fit(x_train,y_train)acqf=optimization.construct(model(params),acquisitions.upper_confidence_bound(2.0),)bfgs=optimizers.bfgs(acqf,bounds,x0,10)candidates=bfgs.init(random.fold_in(optimizer_key,i))next_candidates,values=bfgs.update(candidates)next_x=next_candidates[jnp.argmax(values)]next_y=objective(next_x)x_train=jnp.vstack([x_train,next_x])y_train=jnp.hstack([y_train,next_y])returnx_train,y_trainnext_x_train,next_y_train=optimize(x_train,y_train)Citing BoaxTo cite Boax please use the citation:@software{boax2023github,author={Lando L{\"o}per},title={{B}oax: A Bayesian Optimization library for {JAX}},url={https://github.com/Lando-L/boax},version={0.0.4},year={2023},}In the above bibtex entry, the version number
is intended to be that fromboax/version.py, and the year corresponds to the project's open-source release.