From e8ee02e51b676241c9421ff9c64ed56179d9f9f4 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Thu, 21 Nov 2024 22:44:40 -0500 Subject: [PATCH] misc --- src/mitim_tools/opt_tools/BOTORCHtools.py | 139 ++++++++++++++++--- src/mitim_tools/opt_tools/SURROGATEtools.py | 50 ++++--- src/mitim_tools/opt_tools/utils/TESTtools.py | 4 - 3 files changed, 149 insertions(+), 44 deletions(-) diff --git a/src/mitim_tools/opt_tools/BOTORCHtools.py b/src/mitim_tools/opt_tools/BOTORCHtools.py index 0243de4c..402d0236 100644 --- a/src/mitim_tools/opt_tools/BOTORCHtools.py +++ b/src/mitim_tools/opt_tools/BOTORCHtools.py @@ -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) @@ -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, ) """ @@ -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, diff --git a/src/mitim_tools/opt_tools/SURROGATEtools.py b/src/mitim_tools/opt_tools/SURROGATEtools.py index e9db228b..e936ebce 100644 --- a/src/mitim_tools/opt_tools/SURROGATEtools.py +++ b/src/mitim_tools/opt_tools/SURROGATEtools.py @@ -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 @@ -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 @@ -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() @@ -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 = [] diff --git a/src/mitim_tools/opt_tools/utils/TESTtools.py b/src/mitim_tools/opt_tools/utils/TESTtools.py index c23bb619..cee62b1d 100644 --- a/src/mitim_tools/opt_tools/utils/TESTtools.py +++ b/src/mitim_tools/opt_tools/utils/TESTtools.py @@ -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)