Skip to content

Commit

Permalink
WIP: stochastic optimisation
Browse files Browse the repository at this point in the history
  • Loading branch information
@ruicoelhopedro
ruicoelhopedro committed Oct 7, 2024
1 parent 45694b0 commit 0374ff3
Show file tree
Hide file tree
Showing 18 changed files with 2,508 additions and 727 deletions.
2 changes: 1 addition & 1 deletion piglot/bin/piglot.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -71,7 +71,7 @@ def main(config_path: str = None):
stop_criteria=stop,
)
# Re-run the best case
if 'skip_last_run' not in config and best_params is not None:
if 'skip_last_run' not in config and best_params is not None and not objective.noisy:
objective(best_params)


Expand Down
227 changes: 223 additions & 4 deletions piglot/bin/piglot_plot.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,10 +12,17 @@
from scipy.spatial import ConvexHull
from PIL import Image
import torch
from botorch.sampling import SobolQMCNormalSampler
from botorch.utils.sampling import draw_sobol_normal_samples
from piglot.parameter import read_parameters
from piglot.objectives import read_objective
from piglot.utils.surrogate import get_model, optmise_posterior_mean
from piglot.utils.yaml_parser import parse_config_file
from piglot.optimisers.botorch.dataset import BayesDataset
from piglot.optimisers.botorch.model import build_gp_model
from piglot.optimisers.botorch.composition import BoTorchComposition, RiskComposition
from piglot.optimisers.botorch.acquisitions import get_acquisition, default_acquisition, build_and_optimise_acquisition, optimise_acquisition
from piglot.optimisers.botorch.risk_acquisitions import get_risk_measure


def cumulative_regret(values: np.ndarray, x_grid: np.ndarray) -> np.ndarray:
Expand Down Expand Up @@ -305,12 +312,14 @@ def plot_gp(args):
with torch.no_grad():
posterior = model.posterior(x.unsqueeze(1))
mean = posterior.mean.squeeze()
variance = posterior.variance.squeeze()
f_lb, f_ub = posterior.mvn.confidence_region()
noise_posterior = model.posterior(x.unsqueeze(1), observation_noise=True)
y_lb, y_ub = noise_posterior.mvn.confidence_region()
axis.plot(x, mean, label=f'{name}: Mean')
axis.fill_between(x, mean - variance.sqrt(), mean + variance.sqrt(), alpha=0.5,
label=f'{name}: Std. dev.')
axis.fill_between(x, f_lb, f_ub, alpha=0.5, label=f'{name}: f CI 95%')
axis.fill_between(x, y_lb, y_ub, alpha=0.5, label=f'{name}: y CI 95%')
if 'variances' in data_dict:
axis.errorbar(param_values, values, yerr=np.sqrt(variances), color='black', fmt='o')
axis.errorbar(param_values, values, yerr=2 * np.sqrt(variances), color='black', fmt='o')
else:
axis.scatter(param_values, values, color='black')
axis.set_xlim(x_min, x_max)
Expand All @@ -324,6 +333,184 @@ def plot_gp(args):
plt.close()


def plot_mcgp(args):
"""Driver for plotting a Gaussian process regression with Monte Carlo samplers.
Parameters
----------
args : dict
Passed arguments.
"""
# Build piglot problem
config = parse_config_file(args.config)
parameters = read_parameters(config)
if len(parameters) != 1:
raise ValueError("Can only plot a Gaussian process regression for a single parameter.")
objective = read_objective(config["objective"], parameters, config["output"])
data = objective.get_history()
fig, axis = plt.subplots()
axis: plt.Axes = axis
x_min = min(par.lbound for par in parameters)
x_max = max(par.ubound for par in parameters)

# Load objective data
data_dict = data[list(data.keys())[0]]
values = data_dict['values']
param_values = data_dict['params']
variances = data_dict['variances'] if 'variances' in data_dict else None
noisy = variances is not None

# Load dataset
dataset = BayesDataset.load(os.path.join(config["output"], 'dataset.pt'))
infer_noise = objective.noisy and not objective.stochastic
noise_model = config['optimiser'].get('noise_model', 'homoscedastic')
risk_measure = (
config['optimiser'].get('risk_measure', None) or
config['optimiser'].get('composite_risk_measure', None)
) or 'var'
model = build_gp_model(dataset, infer_noise, noise_model)
composition = BoTorchComposition(model, dataset, objective)
risk_composition = RiskComposition(
model,
dataset,
objective,
args.num_samples,
config['optimiser'].get('seed', None),
get_risk_measure(
risk_measure,
alpha=config['optimiser'].get('alpha', 0.1),
),
)
mc_samples = config['optimiser'].get('mc_samples', 256)
mc_sampler = SobolQMCNormalSampler(torch.Size([mc_samples]), seed=0)
sampler = SobolQMCNormalSampler(torch.Size([args.num_samples]), seed=0)

# Build the acquisition function
acq_name = config['optimiser'].get(
'acquisition',
default_acquisition(
objective.composition is not None,
objective.multi_objective,
variances is not None,
config['optimiser'].get('q', 1),
),
)
options = config['optimiser']
options.pop('name')
_, best = build_and_optimise_acquisition(
'qsr',
model,
dataset,
parameters,
composition,
noisy,
**options,
)
acq = get_acquisition(
acq_name,
model,
dataset,
risk_composition if risk_measure else composition,
best=best,
sampler=mc_sampler,
**options,
)
risk_acq = get_acquisition(
'qsr', model, dataset, composition, best=None, sampler=mc_sampler, risk_measure=risk_measure, alpha=config.get('alpha', 0.1),
)
comp_risk_acq = get_acquisition(
'qsr', model, dataset, risk_composition, best=None, sampler=mc_sampler, alpha=config.get('alpha', 0.1),
)
# candidates, _ = optimise_acquisition(
# acq,
# parameters,
# dataset,
# q=config['optimiser'].get('q', 1),
# sequential=config['optimiser'].get('sequential', False),
# batch_size=config['optimiser'].get('batch_size', 32),
# raw_samples=4096,
# )
# print(candidates)

# Evaluate the model
x = torch.linspace(x_min, x_max, 256, dtype=dataset.dtype, device=dataset.device)
with torch.no_grad():
# Posterior for f
f_posterior = model.posterior(x.unsqueeze(1))
f_samples = -composition.from_model_samples(sampler(f_posterior), x.unsqueeze(1))
f_mean = f_samples.mean(dim=0)
f_lb, f_ub = f_samples.quantile(0.025, dim=0), f_samples.quantile(0.975, dim=0)

# Posterior for y
y_posterior = model.posterior(x.unsqueeze(1), observation_noise=True)
y_samples = -composition.from_model_samples(sampler(y_posterior), x.unsqueeze(1))
y_lb, y_ub = y_samples.quantile(0.025, dim=0), y_samples.quantile(0.975, dim=0)

# Observations
std_values, std_vars = dataset.transform_outcomes()
z_samples = draw_sobol_normal_samples(std_values.shape[-1], args.num_samples, seed=0)
phat_samples = std_values + z_samples.unsqueeze(1) * std_vars.sqrt()
obs_samples = -composition.from_model_samples(phat_samples, dataset.params)
obs_mean = obs_samples.mean(dim=0)
obs_lb, obs_ub = obs_samples.quantile(0.025, dim=0), obs_samples.quantile(0.975, dim=0)
obs_bounds = torch.stack([obs_mean - obs_lb, obs_ub - obs_mean], dim=0).clamp_min(0)

# Noise model
noise_level = model.noise_prediction(x.unsqueeze(1)).sqrt()

# # Risk measure
# if risk_measure:
# risk_samples = -risk_composition.from_model_samples(sampler(f_posterior), x.unsqueeze(1))
# risk_mean = risk_samples.mean(dim=0)
# risk_lb, risk_ub = risk_samples.quantile(0.025, dim=0), risk_samples.quantile(0.975, dim=0)

x_acq = torch.linspace(x_min, x_max, 128, dtype=dataset.dtype, device=dataset.device)
risk_acq = -risk_acq(x_acq.reshape(-1, 1, 1)).squeeze().detach()
comp_risk_acq = -comp_risk_acq(x_acq.reshape(-1, 1, 1)).squeeze().detach()

# # Acquisition values
# x_acq = torch.clone(x)
# x_acq = torch.linspace(x_min, x_max, 64, dtype=dataset.dtype, device=dataset.device)
# if acq_name.startswith('q') and acq_name.endswith('kg'):
# bounds = torch.tensor([[par.lbound, par.ubound] for par in parameters], dtype=dataset.dtype, device=dataset.device)
# acq_values = torch.tensor([
# acq.evaluate(torch.tensor([[[val]]], dtype=dataset.dtype), bounds=bounds.T, num_restarts=32, raw_samples=4096)
# for val in x_acq
# ]).detach()
# # acq_values = acq.evaluate(x_acq.reshape(-1, 1, 1), bounds=bounds.T).squeeze().detach()
# else:
# acq_values = acq(x_acq.reshape(-1, 1, 1)).squeeze().detach()

p, = axis.plot(x, f_mean, label='Mean')
axis.fill_between(x, f_lb, f_ub, alpha=0.6, color=p.get_color(), label='f CI 95%')
axis.fill_between(x, y_lb, y_ub, alpha=0.3, color=p.get_color(), label='y CI 95%')
axis.axhline(-best.item(), color='black', linestyle='--', label='Best')
# axis.twinx().plot(x_acq, acq_values, label='Acquisition', color='red')
# ax2: plt.Axes = axis.twinx()
axis.plot(x_acq, risk_acq, label='Acquisition', color='red')
axis.plot(x_acq, comp_risk_acq, label='Composite Acquisition', color='green')
axis.plot(x, noise_level.squeeze(), label='Noise Level', color='purple')
# ax2.legend()
# if risk_measure:
# p, = axis.plot(x, risk_mean, label='Risk')
# axis.fill_between(x, risk_lb, risk_ub, alpha=0.5, color=p.get_color(), label='Risk CI 95%')
if variances is not None:
axis.errorbar(dataset.params[:, 0], obs_mean, yerr=obs_bounds, fmt='o', capsize=6.0, elinewidth=2.0, label='MC Posterior')
axis.errorbar(param_values, values, yerr=2 * np.sqrt(variances), fmt='+', capsize=6.0, elinewidth=1.0, label='Observations')
else:
axis.scatter(dataset.params[:, 0], obs_mean, marker='o', label='MC Posterior')
axis.scatter(param_values, values, marker='+', label='Observations')
axis.set_xlim(x_min, x_max)
axis.legend()#loc='lower left')
axis.grid()
fig.tight_layout()
if args.save_fig:
fig.savefig(args.save_fig)
else:
plt.show()
plt.close()


def plot_pareto(args):
"""Driver for plotting the Pareto front for a given multi-objective optimisation problem.
Expand Down Expand Up @@ -657,6 +844,38 @@ def main(passed_args: List[str] = None):
)
sp_gp.set_defaults(func=plot_gp)

# Gaussian process with Monte Carlo plotting
sp_mcgp = subparsers.add_parser(
'mcgp',
help='plot a Gaussian process regression with Monte Carlo sampling',
description=("Plot a Gaussian process regression with Monte Carlo sampling. This must be "
"executed in the same path as the running piglot instance."),
)
sp_mcgp.add_argument(
'config',
type=str,
help="Path for the used or generated configuration file.",
)
sp_mcgp.add_argument(
'--save_fig',
default=None,
type=str,
help=("Path to save the generated figure. If used, graphical output is skipped."),
)
sp_mcgp.add_argument(
'--max_calls',
default=None,
type=int,
help=("Max number of calls to plot."),
)
sp_mcgp.add_argument(
'--num_samples',
default=512,
type=int,
help=("Number of samples for the Monte Carlo integration."),
)
sp_mcgp.set_defaults(func=plot_mcgp)

# Pareto front plotting
sp_pareto = subparsers.add_parser(
'pareto',
Expand Down
Loading

0 comments on commit 0374ff3

Please sign in to comment.