Skip to content

Commit

Permalink
Recovering previous combined_model
Browse files Browse the repository at this point in the history
  • Loading branch information
@pabloprf
pabloprf committed Nov 22, 2024
1 parent e035379 commit 8296ba8
Show file tree
Hide file tree
Showing 6 changed files with 113 additions and 391 deletions.
2 changes: 1 addition & 1 deletion src/mitim_modules/maestro/utils/PORTALSbeat.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -253,7 +253,7 @@ def _inform(self, use_previous_residual = True, use_previous_surrogate_data = Tr
if use_previous_surrogate_data and ('portals_surrogate_data_file' in self.maestro_instance.parameters_trans_beat):
if 'surrogateOptions' not in self.optimization_options:
self.optimization_options['surrogateOptions'] = {}
self.optimization_options['surrogateOptions']["extrapointsFile"] = self.maestro_instance.parameters_trans_beat['portals_surrogate_data_file']
self.optimization_options['surrogateOptions']["add_data_from_file"] = self.maestro_instance.parameters_trans_beat['portals_surrogate_data_file']

self.folder_starting_point = self.maestro_instance.parameters_trans_beat['portals_last_run_folder']

Expand Down
299 changes: 4 additions & 295 deletions src/mitim_tools/opt_tools/BOTORCHtools.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -197,7 +197,7 @@ def __init__(
self.input_transform = input_transform
self.to(train_X)

def posterior_full(
def posterior(
self,
X,
output_indices=None,
Expand Down Expand Up @@ -249,303 +249,11 @@ def _Xs_and_transforms(self, X):
Xs = (X,) * len(self.transforms)
return zip(Xs, self.transforms)

class SingleTaskGP_MITIM2(botorch.models.gp_regression.SingleTaskGP):
def __init__(
self,
train_X,
train_Y,
train_Yvar,
input_transform=None,
outcome_transform=None,
surrogateOptions={},
variables=None,
train_X_added=torch.Tensor([]),
train_Y_added=torch.Tensor([]),
train_Yvar_added=torch.Tensor([]),
):
"""
_added refers to already-transformed variables that are added from table
"""

TypeMean = surrogateOptions.get("TypeMean", 0)
TypeKernel = surrogateOptions.get("TypeKernel", 0)
FixedNoise = surrogateOptions.get("FixedNoise", False)
ConstrainNoise = surrogateOptions.get("ConstrainNoise", -1e-4)
learn_additional_noise = surrogateOptions.get("ExtraNoise", False)
print("\t\t* Surrogate model options:")
print(
f"\t\t\t- FixedNoise: {FixedNoise} (extra noise: {learn_additional_noise}), TypeMean: {TypeMean}, TypeKernel: {TypeKernel}, ConstrainNoise: {ConstrainNoise:.1e}"
)

# self.store_training(
# train_X,
# train_X_added,
# train_Y,
# train_Y_added,
# train_Yvar,
# train_Yvar_added,
# input_transform,
# outcome_transform,
# )

"""
----------------------------------------------------------------------------------------
What set_dimensions did, and select things to train (already transformed and normalized)
----------------------------------------------------------------------------------------
"""

# Grab num_outputs
self._num_outputs = train_Y.shape[-1]

# Grab ard_num_dims
if train_X.shape[0] > 0:
with torch.no_grad():
transformed_X = self.transform_inputs(
X=train_X, input_transform=input_transform
)
self.ard_num_dims = transformed_X.shape[-1]
else:
self.ard_num_dims = train_X_added.shape[-1]
transformed_X = torch.empty((0, self.ard_num_dims)).to(train_X)

# Transform outcomes
if outcome_transform is not None:
train_Y, train_Yvar = outcome_transform(train_X, train_Y, train_Yvar)

# # Added points are raw transformed, so I need to normalize them
# if train_X_added.shape[0] > 0:
# train_X_added = input_transform["tf2"](train_X_added)
# train_Y_added, train_Yvar_added = outcome_transform["tf2"](
# train_Y_added, train_Yvar_added
# )
# -----

train_X_usedToTrain = transformed_X #torch.cat((transformed_X, train_X_added), axis=0)
train_Y_usedToTrain = train_Y #torch.cat((train_Y, train_Y_added), axis=0)
train_Yvar_usedToTrain = train_Yvar #torch.cat((train_Yvar, train_Yvar_added), axis=0)

self._input_batch_shape, self._aug_batch_shape = self.get_batch_dimensions(
train_X=train_X_usedToTrain, train_Y=train_Y_usedToTrain
)

# train_Y_usedToTrain = train_Y_usedToTrain.squeeze(-1)
# train_Yvar_usedToTrain = train_Yvar_usedToTrain.squeeze(-1)

"""
-----------------------------------------------------------------------
Likelihood and Noise
-----------------------------------------------------------------------
"""

self._subset_batch_dict = {}

if FixedNoise:
# Noise not inferred, given by data
likelihood = (
gpytorch.likelihoods.gaussian_likelihood.FixedNoiseGaussianLikelihood(
noise=train_Yvar_usedToTrain.clip(1e-6), # I clip the noise to avoid numerical issues (gpytorch would do it anyway, but this way it doesn't throw a warning)
batch_shape=self._aug_batch_shape,
learn_additional_noise=learn_additional_noise,
)
)

else:
# Infer Noise

noise_prior = gpytorch.priors.torch_priors.GammaPrior(1.1, 0.05)
noise_prior_mode = (noise_prior.concentration - 1) / noise_prior.rate

if ConstrainNoise < 0:
noise_constraint = gpytorch.constraints.constraints.GreaterThan(
-ConstrainNoise, transform=None, initial_value=noise_prior_mode
)
else:
noise_constraint = gpytorch.constraints.constraints.Interval(
1e-6, ConstrainNoise, transform=None, initial_value=noise_prior_mode
)

likelihood = gpytorch.likelihoods.gaussian_likelihood.GaussianLikelihood(
noise_prior=noise_prior,
batch_shape=self._aug_batch_shape,
noise_constraint=noise_constraint,
)

self._subset_batch_dict["likelihood.noise_covar.raw_noise"] = -2

"""
-----------------------------------------------------------------------
Initialize ExactGP
-----------------------------------------------------------------------
"""

gpytorch.models.exact_gp.ExactGP.__init__(
self,
train_inputs=train_X_usedToTrain,
train_targets=train_Y_usedToTrain,
likelihood=likelihood,
)

"""
-----------------------------------------------------------------------
GP Mean
-----------------------------------------------------------------------
"""

if TypeMean == 0:
self.mean_module = gpytorch.means.constant_mean.ConstantMean(
batch_shape=self._aug_batch_shape
)
elif TypeMean == 1:
self.mean_module = gpytorch.means.linear_mean.LinearMean(
self.ard_num_dims, batch_shape=self._aug_batch_shape, bias=True
)
elif TypeMean == 2:
self.mean_module = MITIM_LinearMeanGradients(
batch_shape=self._aug_batch_shape, variables=variables
)
elif TypeMean == 3:
self.mean_module = MITIM_CriticalGradient(
batch_shape=self._aug_batch_shape, variables=variables
)

"""
-----------------------------------------------------------------------
GP Kernel - Covariance
-----------------------------------------------------------------------
"""

# Priors
lengthscale_prior = gpytorch.priors.torch_priors.GammaPrior(3.0, 6.0)
outputscale_prior = gpytorch.priors.torch_priors.GammaPrior(2.0, 0.15)

# Do not allow too small lengthscales?
lengthscale_constraint = (
None # gpytorch.constraints.constraints.GreaterThan(0.05)
)

self._subset_batch_dict["covar_module.raw_outputscale"] = -1
self._subset_batch_dict["covar_module.base_kernel.raw_lengthscale"] = -3

if TypeKernel == 0:
self.covar_module = gpytorch.kernels.scale_kernel.ScaleKernel(
base_kernel=gpytorch.kernels.matern_kernel.MaternKernel(
nu=2.5,
ard_num_dims=self.ard_num_dims,
batch_shape=self._aug_batch_shape,
lengthscale_prior=lengthscale_prior,
lengthscale_constraint=lengthscale_constraint,
),
batch_shape=self._aug_batch_shape,
outputscale_prior=outputscale_prior,
)
elif TypeKernel == 1:
self.covar_module = gpytorch.kernels.scale_kernel.ScaleKernel(
base_kernel=gpytorch.kernels.rbf_kernel.RBFKernel(
ard_num_dims=self.ard_num_dims,
batch_shape=self._aug_batch_shape,
lengthscale_prior=lengthscale_prior,
lengthscale_constraint=lengthscale_constraint,
),
batch_shape=self._aug_batch_shape,
outputscale_prior=outputscale_prior,
)
elif TypeKernel == 2:
self.covar_module = MITIM_ConstantKernel(
ard_num_dims=self.ard_num_dims,
batch_shape=self._aug_batch_shape,
lengthscale_prior=lengthscale_prior,
lengthscale_constraint=lengthscale_constraint,
)
elif TypeKernel == 3:
self.covar_module = gpytorch.kernels.scale_kernel.ScaleKernel(
base_kernel=MITIM_NNKernel(
ard_num_dims=self.ard_num_dims,
batch_shape=self._aug_batch_shape,
lengthscale_prior=lengthscale_prior,
lengthscale_constraint=lengthscale_constraint,
),
batch_shape=self._aug_batch_shape,
outputscale_prior=outputscale_prior,
)

if outcome_transform is not None:
self.outcome_transform = outcome_transform
if input_transform is not None:
self.input_transform = input_transform

self.to(train_X)

def store_training(self, x, xa, y, ya, yv, yva, input_transform, outcome_transform):

# x, y are raw untransformed, and I want raw transformed
if input_transform is not None:
x_tr = input_transform["tf1"](x)
else:
x_tr = x
if outcome_transform is not None:
y_tr, yv_tr = outcome_transform["tf1"](x, y, yv)
else:
y_tr, yv_tr = y, yv

# xa, ya are raw transformed
xa_tr = xa
ya_tr, yva_tr = ya, yva

self.train_X_usedToTrain = torch.cat((xa_tr, x_tr), axis=0)
self.train_Y_usedToTrain = torch.cat((ya_tr, y_tr), axis=0)
self.train_Yvar_usedToTrain = torch.cat((yva_tr, yv_tr), axis=0)

# Modify posterior call from BatchedMultiOutputGPyTorchModel to call posterior untransform with "X"

def posterior(
self,
X,
output_indices=None,
observation_noise=False,
posterior_transform=None,
**kwargs,
):

self.eval() # make sure model is in eval mode
# input transforms are applied at `posterior` in `eval` mode, and at
# `model.forward()` at the training time
Xtr = self.transform_inputs(X)
with botorch.models.utils.gpt_posterior_settings():
# insert a dimension for the output dimension
if self._num_outputs > 1:
Xtr, output_dim_idx = botorch.models.utils.add_output_dim(
X=Xtr, original_batch_shape=self._input_batch_shape
)
# NOTE: BoTorch's GPyTorchModels also inherit from GPyTorch's ExactGP, thus
# self(X) calls GPyTorch's ExactGP's __call__, which computes the posterior,
# rather than e.g. SingleTaskGP's forward, which computes the prior.
mvn = self(Xtr)
mvn = self._apply_noise(X=Xtr, mvn=mvn, observation_noise=observation_noise)
if self._num_outputs > 1:
mean_x = mvn.mean
covar_x = mvn.lazy_covariance_matrix
output_indices = output_indices or range(self._num_outputs)
mvns = [
gpytorch.distributions.MultivariateNormal(
mean_x.select(dim=output_dim_idx, index=t),
covar_x[(slice(None),) * output_dim_idx + (t,)],
)
for t in output_indices
]
mvn = gpytorch.distributions.MultitaskMultivariateNormal.from_independent_mvns(mvns=mvns)

posterior = botorch.posteriors.gpytorch.GPyTorchPosterior(distribution=mvn)
if hasattr(self, "outcome_transform"):
posterior = self.outcome_transform.untransform_posterior(X, posterior)
if posterior_transform is not None:
return posterior_transform(posterior)
return posterior

# ----------------------------------------------------------------------------------------------------------------------------
# ModelListGP needs to be modified to allow me to have "common" parameters to models, to not run at every transformation again
# ----------------------------------------------------------------------------------------------------------------------------

class ModifiedModelListGP(botorch.models.model_list_gp_regression.ModelListGP):
class ModelListGP_MITIM(botorch.models.model_list_gp_regression.ModelListGP):
def __init__(self, *gp_models):
super().__init__(*gp_models)

Expand Down Expand Up @@ -1237,6 +945,7 @@ def untransform_posterior(self, posterior: Posterior) -> Posterior:
# could potentially use from_independent_mvns
# print(f"{mvn._covar.shape = }")
# print(f"{covar.shape=}")
dis = MultitaskMultivariateNormal(mean=mean, covariance_matrix=covar)
from gpytorch.distributions import MultivariateNormal
dis = MultivariateNormal(mean=mean, covariance_matrix=covar)
return GPyTorchPosterior(distribution=dis)

4 changes: 2 additions & 2 deletions src/mitim_tools/opt_tools/STEPtools.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -262,7 +262,7 @@ def _fit_individual_models(self, fit_output_contains=None):
models = ()
for GP in self.GP["individual_models"]:
models += (GP.gpmodel,)
self.GP["combined_model"].gpmodel = BOTORCHtools.ModifiedModelListGP(*models)
self.GP["combined_model"].gpmodel = BOTORCHtools.ModelListGP_MITIM(*models)

# ------------------------------------------------------------------------------------------------------
# Make sure each model has the right surrogate_transformation_variables inside the combined model
Expand Down Expand Up @@ -298,7 +298,7 @@ def defineFunctions(self, scalarized_objective):
I create this so that, upon reading a pickle, I re-call it. Otherwise, it is very heavy to store lambdas
"""

self.evaluators = {"GP": self.GP["mo_model"]}
self.evaluators = {"GP": self.GP["combined_model"]}#mo_model"]}

# **************************************************************************************************
# Objective (multi-objective model -> single objective residual)
Expand Down
Loading

0 comments on commit 8296ba8

Please sign in to comment.