<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 Bayesian optimization with context variables
In this notebook we are going to see how to use Emukit to solve optimization problems in which certain variables are fixed during the optimization phase. These are called context variables [1]. This is useful when some of the variables in the optimization are controllable/known factors. And example is the optimization of a the movement of a robot under  conditions of the environment change (but the change is known).
Step1:  Loading the problem and the loop
Step2:  Now we define the domain of the function to optimize.
We build the model
Step3:  And prepare the optimization object to run the loop.
Step4:  Now, we set the number of iterations to run to 10.
Step5:  Running the optimization by setting a context variable
To set a context, we just need to create a dictionary with the variables to fix and pass it to the Bayesian optimization object when running the optimization. Note that, every time we run new iterations we can set other variables to be the context. We run 3 sequences of 10 iterations each with different values of the context.
Step6:  We can now inspect the collected points.
<ASSISTANT_TASK:>
Python Code:
from emukit.test_functions import branin_function
from emukit.core import ParameterSpace, ContinuousParameter, DiscreteParameter
from emukit.core.initial_designs import RandomDesign
from GPy.models import GPRegression
from emukit.model_wrappers import GPyModelWrapper
from emukit.bayesian_optimization.acquisitions import ExpectedImprovement
from emukit.bayesian_optimization.loops import BayesianOptimizationLoop
from emukit.core.loop import FixedIterationsStoppingCondition
Explanation: Bayesian optimization with context variables
In this notebook we are going to see how to use Emukit to solve optimization problems in which certain variables are fixed during the optimization phase. These are called context variables [1]. This is useful when some of the variables in the optimization are controllable/known factors. And example is the optimization of a the movement of a robot under  conditions of the environment change (but the change is known).
End of explanation
f, parameter_space = branin_function()
Explanation: Loading the problem and the loop
End of explanation
design = RandomDesign(parameter_space) # Collect random points
X = design.get_samples(10)
Y = f(X)
model_gpy = GPRegression(X,Y) # Train and wrap the model in Emukit
model_emukit = GPyModelWrapper(model_gpy)
Explanation: Now we define the domain of the function to optimize.
We build the model:
End of explanation
expected_improvement = ExpectedImprovement(model = model_emukit)
bayesopt_loop = BayesianOptimizationLoop(model = model_emukit,
                                         space = parameter_space,
                                         acquisition = expected_improvement,
                                         batch_size = 1)
Explanation: And prepare the optimization object to run the loop.
End of explanation
max_iter = 10
Explanation: Now, we set the number of iterations to run to 10.
End of explanation
bayesopt_loop.run_loop(f, max_iter, context={'x1':0.3}) # we set x1 as the context variable
bayesopt_loop.run_loop(f, max_iter, context={'x2':0.1}) # we set x2 as the context variable
bayesopt_loop.run_loop(f, max_iter) # no context
Explanation: Running the optimization by setting a context variable
To set a context, we just need to create a dictionary with the variables to fix and pass it to the Bayesian optimization object when running the optimization. Note that, every time we run new iterations we can set other variables to be the context. We run 3 sequences of 10 iterations each with different values of the context.
End of explanation
bayesopt_loop.loop_state.X
Explanation: We can now inspect the collected points.
End of explanation
<END_TASK>