<jupyter_start><jupyter_text>Target function<jupyter_code>class HyperbolicTarget(sbo.TargetFunction):
    
    def eval(self, x):
        
        x1 = x[...,0]
        x2 = x[...,1]

        return 0.3*x1**2 + 2*x2**2 - 20.0

target = HyperbolicTarget([[-10, 10], [-10, 10]])<jupyter_output><empty_output><jupyter_text>Training data<jupyter_code>utilities.set_random_seed(112)
N_train_points = 3
X_train = torch.rand(N_train_points, 2)*20 + torch.FloatTensor([-10, -10])
y_train = target.eval(X_train)
utilities.plot2d_func(target.eval, target.ranges)
plt.scatter(X_train[:,0].detach().numpy(), X_train[:,1].detach().numpy(), 
            marker="x", s=200, c='orange', zorder=2, linewidth=4);<jupyter_output><empty_output><jupyter_text>Parametric model<jupyter_code>class ParametricMeanFn(PyroModule):
    
    def __init__(self):
        super().__init__()
        self.alpha = PyroSample(dist.Uniform(0, 10))
        
    def forward(self, X):

        x1 = X[...,0]
        x2 = X[...,1]
        
        return self.alpha * x2**2<jupyter_output><empty_output><jupyter_text>SBO<jupyter_code>pyro.clear_param_store()

# SBO parameters
sbo_steps = 10

opti_num_steps = 1000
opti_params = {"lr": 0.1}
optimizer = pyro.optim.Adam(opti_params)
loss = Trace_ELBO()

return_site = "EI"

# Acquisition function optimizer
acqf_optimizer = torch.optim.Adam
acqf_opti_num_steps = 100
acqf_opti_lr = 0.5

num_candidates = 10

# Sampling
num_samples = 5
# GP kernel
kernel = gp.kernels.Matern52(
    input_dim = X_train.shape[1], lengthscale = 10*torch.ones(X_train.shape[1]))

# Semi paramteric GP model
model = sbo.SemiParametricModel(X_train, y_train, ParametricMeanFn(), kernel)

# Defining GP's parameters
model.gp.kernel.lengthscale = PyroSample(dist.Uniform(0, 15).expand([X_train.shape[1]]).to_event())
model.gp.kernel.variance = PyroSample(dist.Uniform(0, 50))
model.gp.noise = PyroSample(dist.Uniform(0, 1))

guide = None
for i in range(sbo_steps):
    
    guide, predict, losses = sbo.step(model, guide, optimizer, loss, target, acqf_optimizer, 
                                      opti_num_steps=opti_num_steps, acqf_opti_num_steps=acqf_opti_num_steps,
                                      acqf_opti_lr=acqf_opti_lr, num_samples=num_samples, 
                                      num_candidates=num_candidates,
                                      return_site=return_site)
            
    ######################
    # Visualising the step
    ######################
    if i == 0:
        plt.figure(figsize=(5, 3))
        
        plt.title("Training")
        plt.semilogy(losses)
        
    else:
        plt.figure(figsize=(15, 3))
        
        plt.subplot(1,3,1)
        plt.title("mean (samples=%d)" % (num_samples))
        
        with torch.no_grad():
            utilities.plot2d_func(lambda x: predict(x)["y"].mean(0), target.ranges)
        
        plt.subplot(1,3,2)
        plt.title("Find x (samples=%d)" % (num_samples))
        
        with torch.no_grad():
            utilities.plot2d_func(lambda x: predict(x)[return_site].mean(0), target.ranges)
        
        plt.scatter(model.X[-1, 0].detach().numpy(), model.X[-1, 1].detach().numpy(), 
            marker="x", s=200, c='orange', zorder=2, linewidth=4);
        
        plt.subplot(1,3,3)
        plt.title("Training")
        plt.semilogy(losses)
    
    plt.show()
    
    print("SBO step: %d" % (i), model.X[-1])<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.15471698113207547, 'fraction_numerical': 0.03193033381712627, 'mean_word_length': 2.8077348066298344, 'pattern_counts': {'":': 1, '<': 17, '<?xml version=': 0, '>': 17, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '4', 'dir': 'jupyter-structured-clean-dedup', 'id': '324164', 'license': 'MIT', 'n_tokens_mistral': 1329, 'n_tokens_neox': 1247, 'n_words': 187, 'path': 'examples/sbo_hyperbolic.ipynb', 'repo_name': 'tbrx/SBO_Pyro', 'size': '3445'}