Skip to content

Commit

Permalink
misc
Browse files Browse the repository at this point in the history
  • Loading branch information
@pabloprf
pabloprf committed Nov 22, 2024
1 parent 8296ba8 commit e8ee02e
Show file tree
Hide file tree
Showing 3 changed files with 149 additions and 44 deletions.
139 changes: 117 additions & 22 deletions src/mitim_tools/opt_tools/BOTORCHtools.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -59,22 +59,58 @@ def __init__(
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,
)

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

self._validate_tensor_args(X=train_X, Y=train_Y, Yvar=train_Yvar)
if outcome_transform == DEFAULT:
outcome_transform = Standardize(
m=train_Y.shape[-1], batch_shape=train_X.shape[:-2]
)
with torch.no_grad():
transformed_X = self.transform_inputs(
X=train_X, input_transform=input_transform
)

self.ard_num_dims = transformed_X.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 = torch.cat((transformed_X, train_X_added), axis=0)
train_Y_usedToTrain = torch.cat((train_Y, train_Y_added), axis=0)
train_Yvar_usedToTrain = 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)

#self._aug_batch_shape = train_Y.shape[:-2] #<----- New



# Validate again after applying the transforms
self._validate_tensor_args(X=transformed_X, Y=train_Y, Yvar=train_Yvar)
ignore_X_dims = getattr(self, "_ignore_X_dims_scaling_check", None)
Expand All @@ -85,25 +121,65 @@ def __init__(
ignore_X_dims=ignore_X_dims,
)
self._set_dimensions(train_X=train_X, train_Y=train_Y)
self._aug_batch_shape = train_Y.shape[:-2] #<----- New


train_X, train_Y, train_Yvar = self._transform_tensor_args(
X=train_X, Y=train_Y, Yvar=train_Yvar
)

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

self._subset_batch_dict = {}

likelihood = (
gpytorch.likelihoods.gaussian_likelihood.FixedNoiseGaussianLikelihood(
noise=train_Yvar.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)
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,
learn_additional_noise=learn_additional_noise,
noise_constraint=noise_constraint,
)
)
self._is_custom_likelihood = True

ExactGP.__init__(
self, train_inputs=train_X, train_targets=train_Y, likelihood=likelihood
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,
)

"""
Expand Down Expand Up @@ -189,14 +265,33 @@ def __init__(
outputscale_prior=outputscale_prior,
)


# TODO: Allow subsetting of other covar modules
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,
Expand Down
50 changes: 32 additions & 18 deletions src/mitim_tools/opt_tools/SURROGATEtools.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -206,12 +206,12 @@ def __init__(
# Obtain normalization constants now (although during training this is messed up, so needed later too)
# -------------------------------------------------------------------------------------

# self.normalization_pass(
# input_transform_physics,
# input_transform_normalization,
# outcome_transform_physics,
# output_transformed_standardization,
# )
self.normalization_pass(
input_transform_physics,
input_transform_normalization,
outcome_transform_physics,
output_transformed_standardization,
)

# ------------------------------------------------------------------------------------
# Combine transformations in chain of PHYSICS + NORMALIZATION
Expand Down Expand Up @@ -395,8 +395,10 @@ def fit(self):
"""

# Train always in physics-transformed space, to enable mitim re-use training from file
#with fundamental_model_context(self):
track_fval = self.perform_model_fit(mll)
with fundamental_model_context(self):
track_fval = self.perform_model_fit(mll)

embed()

# ---------------------------------------------------------------------------------------------------
# Asses optimization
Expand All @@ -407,12 +409,12 @@ def fit(self):
# Go back to definining the right normalizations, because the optimizer has to work on training mode...
# ---------------------------------------------------------------------------------------------------

# self.normalization_pass(
# self.gpmodel.input_transform["tf1"],
# self.gpmodel.input_transform["tf2"],
# self.gpmodel.outcome_transform["tf1"],
# self.gpmodel.outcome_transform["tf2"],
# )
self.normalization_pass(
self.gpmodel.input_transform["tf1"],
self.gpmodel.input_transform["tf2"],
self.gpmodel.outcome_transform["tf1"],
self.gpmodel.outcome_transform["tf2"],
)

def perform_model_fit(self, mll):
self.gpmodel.train()
Expand Down Expand Up @@ -901,17 +903,29 @@ def __init__(self, surrogate_model):

def __enter__(self):
# Works for individual models, not ModelList
for i in range(len(self.surrogate_model.gpmodel.input_transform.tf1.transforms)):
self.surrogate_model.gpmodel.input_transform.tf1.transforms[i].flag_to_evaluate = False
self.surrogate_model.gpmodel.input_transform.tf1.flag_to_evaluate = False
self.surrogate_model.gpmodel.outcome_transform.tf1.flag_to_evaluate = False

return self.surrogate_model

def __exit__(self, *args):
for i in range(len(self.surrogate_model.gpmodel.input_transform.tf1.transforms)):
self.surrogate_model.gpmodel.input_transform.tf1.transforms[i].flag_to_evaluate = True
self.surrogate_model.gpmodel.input_transform.tf1.flag_to_evaluate = True
self.surrogate_model.gpmodel.outcome_transform.tf1.flag_to_evaluate = True

# def __enter__(self):
# # Works for individual models, not ModelList
# embed()
# for i in range(len(self.surrogate_model.gpmodel.input_transform.tf1.transforms)):
# self.surrogate_model.gpmodel.input_transform.tf1.transforms[i].flag_to_evaluate = False
# self.surrogate_model.gpmodel.outcome_transform.tf1.flag_to_evaluate = False

# return self.surrogate_model

# def __exit__(self, *args):
# for i in range(len(self.surrogate_model.gpmodel.input_transform.tf1.transforms)):
# self.surrogate_model.gpmodel.input_transform.tf1.transforms[i].flag_to_evaluate = True
# self.surrogate_model.gpmodel.outcome_transform.tf1.flag_to_evaluate = True

def create_df_portals(x, y, yvar, x_names, output, max_x = 20):

new_data = []
Expand Down
4 changes: 0 additions & 4 deletions src/mitim_tools/opt_tools/utils/TESTtools.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -47,10 +47,6 @@ def testBatchCapabilities(GPs, combinations=[2, 100, 1000]):
It stops running if the error gets larger than thrPercent in those cases
"""

from mitim_tools.misc_tools import IOtools
with IOtools.speeder("/Users/pablorf/PROJECTS/project_2024_PORTALSdevelopment/speed/profiler_evv.prof") as s:
GPs.predict(GPs.train_X[0:1, :].repeat(50, 1))

for i in combinations:
x = GPs.train_X[0:1, :].repeat(i, 1)

Expand Down

0 comments on commit e8ee02e

Please sign in to comment.