def get_NEI( model: Model, objective_weights: Tensor, outcome_constraints: Optional[Tuple[Tensor, Tensor]] = None, X_observed: Optional[Tensor] = None, X_pending: Optional[Tensor] = None, **kwargs: Any, ) -> AcquisitionFunction: r"""Instantiates a qNoisyExpectedImprovement acquisition function. Args: objective_weights: The objective is to maximize a weighted sum of the columns of f(x). These are the weights. outcome_constraints: A tuple of (A, b). For k outcome constraints and m outputs at f(x), A is (k x m) and b is (k x 1) such that A f(x) <= b. (Not used by single task models) X_observed: A tensor containing points observed for all objective outcomes and outcomes that appear in the outcome constraints (if there are any). X_pending: A tensor containing points whose evaluation is pending (i.e. that have been submitted for evaluation) present for all objective outcomes and outcomes that appear in the outcome constraints (if there are any). mc_samples: The number of MC samples to use (default: 512). qmc: If True, use qMC instead of MC (default: True). prune_baseline: If True, prune the baseline points for NEI (default: True). Returns: qNoisyExpectedImprovement: The instantiated acquisition function. """ if X_observed is None: raise ValueError("There are no feasible observed points.") # Parse random_scalarization params objective_weights = _extract_random_scalarization_settings( objective_weights, outcome_constraints, **kwargs) # construct Objective module if outcome_constraints is None: objective = LinearMCObjective(weights=objective_weights) else: obj_tf = get_objective_weights_transform(objective_weights) con_tfs = get_outcome_constraint_transforms(outcome_constraints) X_observed = torch.as_tensor(X_observed) inf_cost = get_infeasible_cost(X=X_observed, model=model, objective=obj_tf) objective = ConstrainedMCObjective(objective=obj_tf, constraints=con_tfs or [], infeasible_cost=inf_cost) return get_acquisition_function( acquisition_function_name="qNEI", model=model, objective=objective, X_observed=X_observed, X_pending=X_pending, prune_baseline=kwargs.get("prune_baseline", True), mc_samples=kwargs.get("mc_samples", 512), qmc=kwargs.get("qmc", True), seed=torch.randint(1, 10000, (1, )).item(), )
def get_PosteriorMean( model: Model, objective_weights: Tensor, outcome_constraints: Optional[Tuple[Tensor, Tensor]] = None, X_observed: Optional[Tensor] = None, X_pending: Optional[Tensor] = None, **kwargs: Any, ) -> AcquisitionFunction: r"""Instantiates a PosteriorMean acquisition function. Note: If no OutcomeConstraints given, return an analytic acquisition function. This requires {optimizer_kwargs: {joint_optimization: True}} or an optimizer that does not assume pending point support. Args: objective_weights: The objective is to maximize a weighted sum of the columns of f(x). These are the weights. outcome_constraints: A tuple of (A, b). For k outcome constraints and m outputs at f(x), A is (k x m) and b is (k x 1) such that A f(x) <= b. (Not used by single task models) X_observed: A tensor containing points observed for all objective outcomes and outcomes that appear in the outcome constraints (if there are any). X_pending: A tensor containing points whose evaluation is pending (i.e. that have been submitted for evaluation) present for all objective outcomes and outcomes that appear in the outcome constraints (if there are any). Returns: PosteriorMean: The instantiated acquisition function. """ if X_observed is None: raise ValueError("There are no feasible observed points.") # construct Objective module if kwargs.get("chebyshev_scalarization", False): obj_tf = get_chebyshev_scalarization( weights=objective_weights, Y=torch.stack(kwargs.get("Ys")).transpose(0, 1).squeeze(-1), ) else: obj_tf = get_objective_weights_transform(objective_weights) def obj_fn(samples: Tensor, X: Optional[Tensor] = None) -> Tensor: return obj_tf(samples) if outcome_constraints is None: objective = GenericMCObjective(objective=obj_fn) else: con_tfs = get_outcome_constraint_transforms(outcome_constraints) inf_cost = get_infeasible_cost(X=X_observed, model=model, objective=obj_fn) objective = ConstrainedMCObjective(objective=obj_fn, constraints=con_tfs or [], infeasible_cost=inf_cost) # Use qSimpleRegret, not analytic posterior, to handle arbitrary objective fns. acq_func = qSimpleRegret(model, objective=objective) return acq_func
def test_BotorchModel(self, dtype=torch.float, cuda=False): Xs1, Ys1, Yvars1, bounds, task_features, feature_names = _get_torch_test_data( dtype=dtype, cuda=cuda, constant_noise=True) Xs2, Ys2, Yvars2, _, _, _ = _get_torch_test_data(dtype=dtype, cuda=cuda, constant_noise=True) model = BotorchModel() # Test ModelListGP # make training data different for each output Xs2_diff = [Xs2[0] + 0.1] with mock.patch(FIT_MODEL_MO_PATH) as _mock_fit_model: model.fit( Xs=Xs1 + Xs2_diff, Ys=Ys1 + Ys2, Yvars=Yvars1 + Yvars2, bounds=bounds, task_features=task_features, feature_names=feature_names, fidelity_features=[], ) _mock_fit_model.assert_called_once() # Check attributes self.assertTrue(torch.equal(model.Xs[0], Xs1[0])) self.assertTrue(torch.equal(model.Xs[1], Xs2_diff[0])) self.assertEqual(model.dtype, Xs1[0].dtype) self.assertEqual(model.device, Xs1[0].device) self.assertIsInstance(model.model, ModelListGP) # Check fitting model_list = model.model.models self.assertTrue(torch.equal(model_list[0].train_inputs[0], Xs1[0])) self.assertTrue(torch.equal(model_list[1].train_inputs[0], Xs2_diff[0])) self.assertTrue( torch.equal(model_list[0].train_targets, Ys1[0].view(-1))) self.assertTrue( torch.equal(model_list[1].train_targets, Ys2[0].view(-1))) self.assertIsInstance(model_list[0].likelihood, _GaussianLikelihoodBase) self.assertIsInstance(model_list[1].likelihood, _GaussianLikelihoodBase) # Test batched multi-output FixedNoiseGP with mock.patch(FIT_MODEL_MO_PATH) as _mock_fit_model: model.fit( Xs=Xs1 + Xs2, Ys=Ys1 + Ys2, Yvars=Yvars1 + Yvars2, bounds=bounds, task_features=task_features, feature_names=feature_names, fidelity_features=[], ) _mock_fit_model.assert_called_once() # Check attributes self.assertTrue(torch.equal(model.Xs[0], Xs1[0])) self.assertTrue(torch.equal(model.Xs[1], Xs2[0])) self.assertEqual(model.dtype, Xs1[0].dtype) self.assertEqual(model.device, Xs1[0].device) self.assertIsInstance(model.model, FixedNoiseGP) # Check fitting # train inputs should be `o x n x 1` self.assertTrue( torch.equal( model.model.train_inputs[0], Xs1[0].unsqueeze(0).expand(torch.Size([2]) + Xs1[0].shape), )) # train targets should be `o x n` self.assertTrue( torch.equal(model.model.train_targets, torch.cat(Ys1 + Ys2, dim=-1).permute(1, 0))) self.assertIsInstance(model.model.likelihood, _GaussianLikelihoodBase) # Check infeasible cost can be computed on the model device = torch.device("cuda") if cuda else torch.device("cpu") objective_weights = torch.tensor([1.0, 0.0], dtype=dtype, device=device) objective_transform = get_objective_weights_transform( objective_weights) infeasible_cost = torch.tensor( get_infeasible_cost(X=Xs1[0], model=model.model, objective=objective_transform)) expected_infeasible_cost = -1 * torch.min( objective_transform( model.model.posterior(Xs1[0]).mean - 6 * model.model.posterior(Xs1[0]).variance.sqrt()).min(), torch.tensor(0.0, dtype=dtype, device=device), ) self.assertTrue( torch.abs(infeasible_cost - expected_infeasible_cost) < 1e-5) # Check prediction X = torch.tensor([[6.0, 7.0, 8.0]], dtype=dtype, device=device) f_mean, f_cov = model.predict(X) self.assertTrue(f_mean.shape == torch.Size([1, 2])) self.assertTrue(f_cov.shape == torch.Size([1, 2, 2])) # Check generation objective_weights = torch.tensor([1.0, 0.0], dtype=dtype, device=device) outcome_constraints = ( torch.tensor([[0.0, 1.0]], dtype=dtype, device=device), torch.tensor([[5.0]], dtype=dtype, device=device), ) linear_constraints = (torch.tensor([[0.0, 1.0, 1.0]]), torch.tensor([[100.0]])) fixed_features = None pending_observations = [ torch.tensor([[1.0, 3.0, 4.0]], dtype=dtype, device=device), torch.tensor([[2.0, 6.0, 8.0]], dtype=dtype, device=device), ] n = 3 X_dummy = torch.tensor([[[1.0, 2.0, 3.0]]], dtype=dtype, device=device) model_gen_options = {} # test sequential optimize with mock.patch("ax.models.torch.botorch_defaults.sequential_optimize", return_value=X_dummy) as mock_optimize_acqf: Xgen, wgen = model.gen( n=n, bounds=bounds, objective_weights=objective_weights, outcome_constraints=outcome_constraints, linear_constraints=linear_constraints, fixed_features=fixed_features, pending_observations=pending_observations, model_gen_options=model_gen_options, rounding_func=dummy_func, ) # note: gen() always returns CPU tensors self.assertTrue(torch.equal(Xgen, X_dummy.cpu())) self.assertTrue(torch.equal(wgen, torch.ones(n, dtype=dtype))) # test joint optimize with mock.patch("ax.models.torch.botorch_defaults.joint_optimize", return_value=X_dummy) as mock_optimize_acqf: Xgen, wgen = model.gen( n=n, bounds=bounds, objective_weights=objective_weights, outcome_constraints=None, linear_constraints=None, fixed_features=fixed_features, pending_observations=pending_observations, model_gen_options={ "optimizer_kwargs": { "joint_optimization": True } }, ) # note: gen() always returns CPU tensors self.assertTrue(torch.equal(Xgen, X_dummy.cpu())) self.assertTrue(torch.equal(wgen, torch.ones(n, dtype=dtype))) mock_optimize_acqf.assert_called_once() # test get_rounding_func dummy_rounding = get_rounding_func(rounding_func=dummy_func) X_temp = torch.rand(1, 2, 3, 4) self.assertTrue(torch.equal(X_temp, dummy_rounding(X_temp))) # Check best point selection xbest = model.best_point(bounds=bounds, objective_weights=objective_weights) xbest = model.best_point( bounds=bounds, objective_weights=objective_weights, fixed_features={0: 100.0}, ) self.assertIsNone(xbest) # Test cross-validation mean, variance = model.cross_validate( Xs_train=Xs1 + Xs2, Ys_train=Ys1 + Ys2, Yvars_train=Yvars1 + Yvars2, X_test=torch.tensor([[1.2, 3.2, 4.2], [2.4, 5.2, 3.2]], dtype=dtype, device=device), ) self.assertTrue(mean.shape == torch.Size([2, 2])) self.assertTrue(variance.shape == torch.Size([2, 2, 2])) # Test cross-validation with refit_on_cv model.refit_on_cv = True mean, variance = model.cross_validate( Xs_train=Xs1 + Xs2, Ys_train=Ys1 + Ys2, Yvars_train=Yvars1 + Yvars2, X_test=torch.tensor([[1.2, 3.2, 4.2], [2.4, 5.2, 3.2]], dtype=dtype, device=device), ) self.assertTrue(mean.shape == torch.Size([2, 2])) self.assertTrue(variance.shape == torch.Size([2, 2, 2])) # Test update model.refit_on_update = False model.update(Xs=Xs2 + Xs2, Ys=Ys2 + Ys2, Yvars=Yvars2 + Yvars2) # Test feature_importances importances = model.feature_importances() self.assertEqual(importances.shape, torch.Size([2, 1, 3])) # When calling update directly, the data is completely overwritten. self.assertTrue(torch.equal(model.Xs[0], Xs2[0])) self.assertTrue(torch.equal(model.Xs[1], Xs2[0])) self.assertTrue(torch.equal(model.Ys[0], Ys2[0])) self.assertTrue(torch.equal(model.Yvars[0], Yvars2[0])) model.refit_on_update = True with mock.patch(FIT_MODEL_MO_PATH) as _mock_fit_model: model.update(Xs=Xs2 + Xs2, Ys=Ys2 + Ys2, Yvars=Yvars2 + Yvars2) # test unfit model CV, update, and feature_importances unfit_model = BotorchModel() with self.assertRaises(RuntimeError): unfit_model.cross_validate( Xs_train=Xs1 + Xs2, Ys_train=Ys1 + Ys2, Yvars_train=Yvars1 + Yvars2, X_test=Xs1[0], ) with self.assertRaises(RuntimeError): unfit_model.update(Xs=Xs1 + Xs2, Ys=Ys1 + Ys2, Yvars=Yvars1 + Yvars2) with self.assertRaises(RuntimeError): unfit_model.feature_importances() # Test loading state dict tkwargs = {"device": device, "dtype": dtype} true_state_dict = { "mean_module.constant": [3.5004], "covar_module.raw_outputscale": 2.2438, "covar_module.base_kernel.raw_lengthscale": [[-0.9274, -0.9274, -0.9274]], "covar_module.base_kernel.lengthscale_prior.concentration": 3.0, "covar_module.base_kernel.lengthscale_prior.rate": 6.0, "covar_module.outputscale_prior.concentration": 2.0, "covar_module.outputscale_prior.rate": 0.15, } true_state_dict = { key: torch.tensor(val, **tkwargs) for key, val in true_state_dict.items() } model = get_and_fit_model( Xs=Xs1, Ys=Ys1, Yvars=Yvars1, task_features=[], fidelity_features=[], state_dict=true_state_dict, refit_model=False, ) for k, v in chain(model.named_parameters(), model.named_buffers()): self.assertTrue(torch.equal(true_state_dict[k], v)) # Test for some change in model parameters & buffer for refit_model=True true_state_dict["mean_module.constant"] += 0.1 true_state_dict["covar_module.raw_outputscale"] += 0.1 true_state_dict["covar_module.base_kernel.raw_lengthscale"] += 0.1 true_state_dict = { key: torch.tensor(val, **tkwargs) for key, val in true_state_dict.items() } model = get_and_fit_model( Xs=Xs1, Ys=Ys1, Yvars=Yvars1, task_features=[], fidelity_features=[], state_dict=true_state_dict, refit_model=True, ) self.assertTrue( any(not torch.equal(true_state_dict[k], v) for k, v in chain( model.named_parameters(), model.named_buffers())))
def get_NEI( model: Model, objective_weights: Tensor, outcome_constraints: Optional[Tuple[Tensor, Tensor]] = None, X_observed: Optional[Tensor] = None, X_pending: Optional[Tensor] = None, **kwargs: Any, ) -> AcquisitionFunction: r"""Instantiates a qNoisyExpectedImprovement acquisition function. Args: model: The underlying model which the acqusition function uses to estimate acquisition values of candidates. objective_weights: The objective is to maximize a weighted sum of the columns of f(x). These are the weights. outcome_constraints: A tuple of (A, b). For k outcome constraints and m outputs at f(x), A is (k x m) and b is (k x 1) such that A f(x) <= b. (Not used by single task models) X_observed: A tensor containing points observed for all objective outcomes and outcomes that appear in the outcome constraints (if there are any). X_pending: A tensor containing points whose evaluation is pending (i.e. that have been submitted for evaluation) present for all objective outcomes and outcomes that appear in the outcome constraints (if there are any). mc_samples: The number of MC samples to use (default: 512). qmc: If True, use qMC instead of MC (default: True). prune_baseline: If True, prune the baseline points for NEI (default: True). chebyshev_scalarization: Use augmented Chebyshev scalarization. Returns: qNoisyExpectedImprovement: The instantiated acquisition function. """ if X_observed is None: raise ValueError("There are no feasible observed points.") # construct Objective module if kwargs.get("chebyshev_scalarization", False): if "Ys" not in kwargs: raise ValueError("Chebyshev Scalarization requires Ys argument") Y_tensor = torch.cat(kwargs.get("Ys"), dim=-1) obj_tf = get_chebyshev_scalarization(weights=objective_weights, Y=Y_tensor) else: obj_tf = get_objective_weights_transform(objective_weights) if outcome_constraints is None: objective = GenericMCObjective(objective=obj_tf) else: con_tfs = get_outcome_constraint_transforms(outcome_constraints) inf_cost = get_infeasible_cost(X=X_observed, model=model, objective=obj_tf) objective = ConstrainedMCObjective(objective=obj_tf, constraints=con_tfs or [], infeasible_cost=inf_cost) return get_acquisition_function( acquisition_function_name="qNEI", model=model, objective=objective, X_observed=X_observed, X_pending=X_pending, prune_baseline=kwargs.get("prune_baseline", True), mc_samples=kwargs.get("mc_samples", 512), qmc=kwargs.get("qmc", True), # pyre-fixme[6]: Expected `Optional[int]` for 9th param but got # `Union[float, int]`. seed=torch.randint(1, 10000, (1, )).item(), )
def test_NoMCSamples(self): Y = torch.ones(2, 4, 2) objective_transform = get_objective_weights_transform( torch.tensor([1.0, 1.0])) Y_transformed = objective_transform(Y) self.assertTrue(torch.equal(torch.sum(Y, dim=-1), Y_transformed))
def test_IncompatibleNumberOfWeights(self): Y = torch.ones(5, 2, 4, 1) objective_transform = get_objective_weights_transform( torch.tensor([1.0, 2.0])) with self.assertRaises(RuntimeError): objective_transform(Y)
def test_OneWeightBroadcasting(self): Y = torch.ones(5, 2, 4, 1) objective_transform = get_objective_weights_transform( torch.tensor([0.5])) Y_transformed = objective_transform(Y) self.assertTrue(torch.equal(0.5 * Y.sum(dim=-1), Y_transformed))
def test_NoWeights(self): Y = torch.ones(5, 2, 4, 1) objective_transform = get_objective_weights_transform(None) Y_transformed = objective_transform(Y) self.assertTrue(torch.equal(Y.squeeze(-1), Y_transformed))
def _get_acquisition_func( model: Model, acquisition_function_name: str, objective_weights: Tensor, outcome_constraints: Optional[Tuple[Tensor, Tensor]] = None, X_observed: Optional[Tensor] = None, X_pending: Optional[Tensor] = None, mc_objective: Type[GenericMCObjective] = GenericMCObjective, constrained_mc_objective: Optional[ Type[ConstrainedMCObjective] ] = ConstrainedMCObjective, mc_objective_kwargs: Optional[Dict] = None, **kwargs: Any, ) -> AcquisitionFunction: r"""Instantiates a acquisition function. Args: model: The underlying model which the acqusition function uses to estimate acquisition values of candidates. acquisition_function_name: Name of the acquisition function. objective_weights: The objective is to maximize a weighted sum of the columns of f(x). These are the weights. outcome_constraints: A tuple of (A, b). For k outcome constraints and m outputs at f(x), A is (k x m) and b is (k x 1) such that A f(x) <= b. (Not used by single task models) X_observed: A tensor containing points observed for all objective outcomes and outcomes that appear in the outcome constraints (if there are any). X_pending: A tensor containing points whose evaluation is pending (i.e. that have been submitted for evaluation) present for all objective outcomes and outcomes that appear in the outcome constraints (if there are any). mc_objective: GenericMCObjective class, used for constructing a MC-objective. If constructing a penalized MC-objective, pass in PenalizedMCObjective together with mc_objective_kwargs . constrained_mc_objective: ConstrainedMCObjective class, used when applying constraints on the outcomes. mc_objective_kwargs: kwargs for constructing MC-objective. For GenericMCObjective, leave it as None. For PenalizedMCObjective, it needs to be specified in the format of kwargs. mc_samples: The number of MC samples to use (default: 512). qmc: If True, use qMC instead of MC (default: True). prune_baseline: If True, prune the baseline points for NEI (default: True). chebyshev_scalarization: Use augmented Chebyshev scalarization. Returns: The instantiated acquisition function. """ if X_observed is None: raise ValueError("There are no feasible observed points.") # construct Objective module if kwargs.get("chebyshev_scalarization", False): with torch.no_grad(): Y = model.posterior(X_observed).mean obj_tf = get_chebyshev_scalarization(weights=objective_weights, Y=Y) else: obj_tf = get_objective_weights_transform(objective_weights) def objective(samples: Tensor, X: Optional[Tensor] = None) -> Tensor: return obj_tf(samples) if outcome_constraints is None: mc_objective_kwargs = {} if mc_objective_kwargs is None else mc_objective_kwargs objective = mc_objective(objective=objective, **mc_objective_kwargs) else: if constrained_mc_objective is None: raise ValueError( "constrained_mc_objective cannot be set to None " "when applying outcome constraints." ) if issubclass(mc_objective, PenalizedMCObjective): raise RuntimeError( "Outcome constraints are not supported for PenalizedMCObjective." ) con_tfs = get_outcome_constraint_transforms(outcome_constraints) inf_cost = get_infeasible_cost(X=X_observed, model=model, objective=objective) objective = constrained_mc_objective( objective=objective, constraints=con_tfs or [], infeasible_cost=inf_cost ) return get_acquisition_function( acquisition_function_name=acquisition_function_name, model=model, objective=objective, X_observed=X_observed, X_pending=X_pending, prune_baseline=kwargs.get("prune_baseline", True), mc_samples=kwargs.get("mc_samples", 512), qmc=kwargs.get("qmc", True), # pyre-fixme[6]: Expected `Optional[int]` for 9th param but got # `Union[float, int]`. seed=torch.randint(1, 10000, (1,)).item(), marginalize_dim=kwargs.get("marginalize_dim"), )
def fit_albo_objective(self) -> AlboMCObjective: r"""Inner loop of Augmented Lagrangian algorithm Args: model: A BoTorch model, fitted to observed data objective_callable: A callable transformation from model outputs to objective constraints_callable_list: A callable transformation from model outputs to constraints, with negative values imply feasibility sampler: An MCSampler instance for monte-carlo acquisition Returns: albo_objective: augmented objective with fitted Lagrangian multipliers trace: optimization trace """ objective_callable = get_objective_weights_transform( self.objective_weights) constraints_callable_list = get_outcome_constraint_transforms( self.outcome_constraints) penalty_rate = self.init_penalty_rate num_mults = self.outcome_constraints[0].shape[0] if self.init_mults is not None: assert num_mults == self.init_mults.shape[-1] mults = self.init_mults else: mults = torch.Tensor( [self._default_init_mult for _ in range(num_mults)]) x_trace = torch.zeros_like(self.bounds[0].unsqueeze(0)) mults_trace = mults.unsqueeze(0) output_means = torch.zeros((1, self.model.num_outputs), dtype=float) output_variances = torch.zeros((1, self.model.num_outputs), dtype=float) for i in range(self.num_iter): self._execute_callbacks("on_iter_start", locals()) # 1. Optimize the augmented objective with fixed multipliers to find the next point for multipliers update albo_objective = self.albo_objective_constructor( objective=objective_callable, constraints=constraints_callable_list, penalty_rate=penalty_rate, lagrange_mults=mults) # Using predictive mean for inner loop optimization acq_function = qSimpleRegret(model=self.model, objective=albo_objective, sampler=self.sampler) x, val = optimize_acqf(acq_function=acq_function, bounds=self.bounds, q=1, num_restarts=self.num_restarts, raw_samples=self.raw_samples) # 2. Compute update of lagrange multipliers at the optimal point of augmented objective posterior = self.model.posterior(x.unsqueeze(0)) samples = self.sampler(posterior) mults_next, mults_stds_next = albo_objective.get_mults_update( samples) # 3. Possibly apply heuristics here, i.e clamp mults before update and increase penalty rate mults = mults_next # currently not using any heuristics penalty_rate = self.init_penalty_rate # seems to work just fine with a constant penalty rate # 4. Write trace of inner-loop optimization for debugging x_trace = torch.cat([x_trace, x], dim=0) mults_trace = torch.cat([mults_trace, mults.unsqueeze(0)], dim=0) output_means = torch.cat( [output_means, posterior.mean.detach().squeeze(dim=0)], dim=0) output_variances = torch.cat( [output_variances, posterior.variance.detach().squeeze(dim=0)], dim=0) # 5. Check stopping condition for inner loop (not implemented) self._execute_callbacks("on_iter_end", locals()) continue # Construct final objective albo_objective = self.albo_objective_constructor( objective=objective_callable, constraints=constraints_callable_list, penalty_rate=penalty_rate, lagrange_mults=mults) trace = { 'x': x_trace, 'mults': mults_trace, 'output': { 'mean': output_means, 'variance': output_variances } } return albo_objective, trace
def test_FullyBayesianBotorchModel(self, dtype=torch.float, cuda=False): Xs1, Ys1, Yvars1, bounds, tfs, fns, mns = get_torch_test_data( dtype=dtype, cuda=cuda, constant_noise=True) Xs2, Ys2, Yvars2, _, _, _, _ = get_torch_test_data( dtype=dtype, cuda=cuda, constant_noise=True) Yvars_inferred_noise = [ torch.full_like(Yvars1[0], float("nan")), torch.full_like(Yvars2[0], float("nan")), ] # make input different for each output Xs2_diff = [Xs2[0] + 0.1] Xs = Xs1 + Xs2_diff Ys = Ys1 + Ys2 for inferred_noise, use_input_warping, use_saas in product( (True, False), repeat=3): Yvars = Yvars_inferred_noise if inferred_noise else Yvars1 + Yvars2 model = self.model_cls( use_input_warping=use_input_warping, thinning=1, num_samples=4, use_saas=use_saas, disable_progbar=True, max_tree_depth=1, ) if use_input_warping: self.assertTrue(model.use_input_warping) # Test ModelListGP # make training data different for each output tkwargs = {"dtype": dtype, "device": Xs1[0].device} dummy_samples_list = _get_dummy_mcmc_samples(num_samples=4, num_outputs=2, **tkwargs) for dummy_samples in dummy_samples_list: if use_input_warping: dummy_samples["c0"] = ( torch.rand(4, 1, Xs1[0].shape[-1], **tkwargs) * 0.5 + 0.1) dummy_samples["c1"] = ( torch.rand(4, 1, Xs1[0].shape[-1], **tkwargs) * 0.5 + 0.1) if inferred_noise: dummy_samples["noise"] = torch.rand( 4, 1, **tkwargs).clamp_min(MIN_INFERRED_NOISE_LEVEL) with mock.patch( RUN_INFERENCE_PATH, side_effect=dummy_samples_list, ) as _mock_fit_model: model.fit( Xs=Xs, Ys=Ys, Yvars=Yvars, search_space_digest=SearchSpaceDigest( feature_names=fns, bounds=bounds, task_features=tfs, ), metric_names=mns, ) self.assertEqual(_mock_fit_model.call_count, 2) for i, call in enumerate(_mock_fit_model.call_args_list): _, ckwargs = call X = Xs[i] Y = Ys[i] Yvar = Yvars[i] self.assertIs(ckwargs["pyro_model"], pyro_model) self.assertTrue(torch.equal(ckwargs["X"], X)) self.assertTrue(torch.equal(ckwargs["Y"], Y)) if inferred_noise: self.assertTrue(torch.isnan(ckwargs["Yvar"]).all()) else: self.assertTrue(torch.equal(ckwargs["Yvar"], Yvar)) self.assertEqual(ckwargs["num_samples"], 4) self.assertEqual(ckwargs["warmup_steps"], 1024) self.assertEqual(ckwargs["max_tree_depth"], 1) self.assertTrue(ckwargs["disable_progbar"]) self.assertEqual(ckwargs["use_input_warping"], use_input_warping) self.assertEqual(ckwargs["use_saas"], use_saas) # Check attributes self.assertTrue(torch.equal(model.Xs[i], Xs[i])) self.assertEqual(model.dtype, Xs[i].dtype) self.assertEqual(model.device, Xs[i].device) self.assertIsInstance(model.model, ModelListGP) # Check fitting # Note each model in the model list is a batched model, where # the batch dim corresponds to the MCMC samples model_list = model.model.models # Put model in `eval` mode to transform the train inputs. m = model_list[i].eval() # check mcmc samples dummy_samples = dummy_samples_list[i] expected_train_inputs = Xs[i].expand(4, *Xs[i].shape) if use_input_warping: # train inputs should be warped inputs expected_train_inputs = m.input_transform( expected_train_inputs) self.assertTrue( torch.equal( m.train_inputs[0], expected_train_inputs, )) self.assertTrue( torch.equal( m.train_targets, Ys[i].view(1, -1).expand(4, Ys[i].numel()), )) expected_noise = (dummy_samples["noise"].view( m.likelihood.noise.shape) if inferred_noise else Yvars[i].view(1, -1).expand( 4, Yvars[i].numel())) self.assertTrue( torch.allclose( m.likelihood.noise.detach(), expected_noise, )) self.assertIsInstance(m.likelihood, _GaussianLikelihoodBase) self.assertTrue( torch.allclose( m.covar_module.base_kernel.lengthscale.detach( ), dummy_samples["lengthscale"].view( m.covar_module.base_kernel.lengthscale. shape), )) self.assertTrue( torch.allclose( m.covar_module.outputscale.detach(), dummy_samples["outputscale"].view( m.covar_module.outputscale.shape), )) self.assertTrue( torch.allclose( m.mean_module.constant.detach(), dummy_samples["mean"].view( m.mean_module.constant.shape), )) if use_input_warping: self.assertTrue(hasattr(m, "input_transform")) self.assertIsInstance(m.input_transform, Warp) self.assertTrue( torch.equal( m.input_transform.concentration0, dummy_samples_list[i]["c0"], )) self.assertTrue( torch.equal( m.input_transform.concentration1, dummy_samples_list[i]["c1"], )) else: self.assertFalse(hasattr(m, "input_transform")) # test that multi-task is not implemented ( Xs_mt, Ys_mt, Yvars_mt, bounds_mt, tfs_mt, fns_mt, mns_mt, ) = get_torch_test_data(dtype=dtype, cuda=cuda, constant_noise=True, task_features=[2]) with mock.patch( RUN_INFERENCE_PATH, side_effect=dummy_samples_list, ) as _mock_fit_model, self.assertRaises(NotImplementedError): model.fit( Xs=Xs_mt, Ys=Ys_mt, Yvars=Yvars_mt, search_space_digest=SearchSpaceDigest( feature_names=fns_mt, bounds=bounds_mt, task_features=tfs_mt, ), metric_names=mns_mt, ) with mock.patch( RUN_INFERENCE_PATH, side_effect=dummy_samples_list, ) as _mock_fit_model, self.assertRaises(NotImplementedError): model.fit( Xs=Xs1 + Xs2, Ys=Ys1 + Ys2, Yvars=Yvars1 + Yvars2, search_space_digest=SearchSpaceDigest( feature_names=fns, bounds=bounds, fidelity_features=[0], ), metric_names=mns, ) # fit model with same inputs (otherwise X_observed will be None) model = self.model_cls( use_input_warping=use_input_warping, thinning=1, num_samples=4, use_saas=use_saas, disable_progbar=True, max_tree_depth=1, ) Yvars = Yvars1 + Yvars2 dummy_samples_list = _get_dummy_mcmc_samples(num_samples=4, num_outputs=2, **tkwargs) with mock.patch( RUN_INFERENCE_PATH, side_effect=dummy_samples_list, ) as _mock_fit_model: model.fit( Xs=Xs1 + Xs2, Ys=Ys1 + Ys2, Yvars=Yvars, search_space_digest=SearchSpaceDigest( feature_names=fns, bounds=bounds, task_features=tfs, ), metric_names=mns, ) # Check infeasible cost can be computed on the model device = torch.device("cuda") if cuda else torch.device("cpu") objective_weights = torch.tensor([1.0, 0.0], dtype=dtype, device=device) objective_transform = get_objective_weights_transform( objective_weights) infeasible_cost = torch.tensor( get_infeasible_cost(X=Xs1[0], model=model.model, objective=objective_transform)) expected_infeasible_cost = -1 * torch.min( objective_transform( model.model.posterior(Xs1[0]).mean - 6 * model.model.posterior(Xs1[0]).variance.sqrt()).min(), torch.tensor(0.0, dtype=dtype, device=device), ) self.assertTrue( torch.abs(infeasible_cost - expected_infeasible_cost) < 1e-5) # Check prediction X = torch.tensor([[6.0, 7.0, 8.0]], **tkwargs) f_mean, f_cov = model.predict(X) self.assertTrue(f_mean.shape == torch.Size([1, 2])) self.assertTrue(f_cov.shape == torch.Size([1, 2, 2])) # Check generation objective_weights = torch.tensor( [1.0, 0.0] if self.model_cls is FullyBayesianBotorchModel else [1.0, 1.0], **tkwargs, ) outcome_constraints = ( torch.tensor([[0.0, 1.0]], **tkwargs), torch.tensor([[5.0]], **tkwargs), ) gen_kwargs = ({ "objective_thresholds": torch.zeros(2, **tkwargs) } if self.model_cls is FullyBayesianMOOBotorchModel else {}) linear_constraints = ( torch.tensor([[0.0, 1.0, 1.0]]), torch.tensor([[100.0]]), ) fixed_features = None pending_observations = [ torch.tensor([[1.0, 3.0, 4.0]], **tkwargs), torch.tensor([[2.0, 6.0, 8.0]], **tkwargs), ] n = 3 X_dummy = torch.tensor([[[1.0, 2.0, 3.0]]], **tkwargs) acqfv_dummy = torch.tensor([[[1.0, 2.0, 3.0]]], **tkwargs) model_gen_options = { Keys.OPTIMIZER_KWARGS: { "maxiter": 1 }, Keys.ACQF_KWARGS: { "mc_samples": 3 }, } # test sequential optimize with constraints with mock.patch( "ax.models.torch.botorch_defaults.optimize_acqf", return_value=(X_dummy, acqfv_dummy), ) as _: Xgen, wgen, gen_metadata, cand_metadata = model.gen( n=n, bounds=bounds, objective_weights=objective_weights, outcome_constraints=outcome_constraints, linear_constraints=linear_constraints, fixed_features=fixed_features, pending_observations=pending_observations, model_gen_options=model_gen_options, rounding_func=dummy_func, **gen_kwargs, ) # note: gen() always returns CPU tensors self.assertTrue(torch.equal(Xgen, X_dummy.cpu())) self.assertTrue( torch.equal(wgen, torch.ones(n, dtype=dtype))) # actually test optimization for 1 step without constraints with mock.patch( "ax.models.torch.botorch_defaults.optimize_acqf", wraps=optimize_acqf, return_value=(X_dummy, acqfv_dummy), ) as _: Xgen, wgen, gen_metadata, cand_metadata = model.gen( n=n, bounds=bounds, objective_weights=objective_weights, outcome_constraints=outcome_constraints, fixed_features=fixed_features, pending_observations=pending_observations, model_gen_options=model_gen_options, **gen_kwargs, ) # note: gen() always returns CPU tensors self.assertTrue(torch.equal(Xgen, X_dummy.cpu())) self.assertTrue( torch.equal(wgen, torch.ones(n, dtype=dtype))) # Check best point selection xbest = model.best_point(bounds=bounds, objective_weights=objective_weights) xbest = model.best_point( bounds=bounds, objective_weights=objective_weights, fixed_features={0: 100.0}, ) self.assertIsNone(xbest) # Test cross-validation mean, variance = model.cross_validate( Xs_train=Xs1 + Xs2, Ys_train=Ys, Yvars_train=Yvars, X_test=torch.tensor([[1.2, 3.2, 4.2], [2.4, 5.2, 3.2]], dtype=dtype, device=device), ) self.assertTrue(mean.shape == torch.Size([2, 2])) self.assertTrue(variance.shape == torch.Size([2, 2, 2])) # Test cross-validation with refit_on_cv model.refit_on_cv = True with mock.patch( RUN_INFERENCE_PATH, side_effect=dummy_samples_list, ) as _mock_fit_model: mean, variance = model.cross_validate( Xs_train=Xs1 + Xs2, Ys_train=Ys, Yvars_train=Yvars, X_test=torch.tensor( [[1.2, 3.2, 4.2], [2.4, 5.2, 3.2]], dtype=dtype, device=device, ), ) self.assertTrue(mean.shape == torch.Size([2, 2])) self.assertTrue(variance.shape == torch.Size([2, 2, 2])) # Test update model.refit_on_update = False model.update(Xs=Xs2 + Xs2, Ys=Ys2 + Ys2, Yvars=Yvars2 + Yvars2) # Test feature_importances importances = model.feature_importances() self.assertEqual(importances.shape, torch.Size([2, 1, 3])) # When calling update directly, the data is completely overwritten. self.assertTrue(torch.equal(model.Xs[0], Xs2[0])) self.assertTrue(torch.equal(model.Xs[1], Xs2[0])) self.assertTrue(torch.equal(model.Ys[0], Ys2[0])) self.assertTrue(torch.equal(model.Yvars[0], Yvars2[0])) model.refit_on_update = True with mock.patch( RUN_INFERENCE_PATH, side_effect=dummy_samples_list) as _mock_fit_model: model.update(Xs=Xs2 + Xs2, Ys=Ys2 + Ys2, Yvars=Yvars2 + Yvars2) # test unfit model CV, update, and feature_importances unfit_model = self.model_cls() with self.assertRaises(RuntimeError): unfit_model.cross_validate( Xs_train=Xs1 + Xs2, Ys_train=Ys1 + Ys2, Yvars_train=Yvars1 + Yvars2, X_test=Xs1[0], ) with self.assertRaises(RuntimeError): unfit_model.update(Xs=Xs1 + Xs2, Ys=Ys1 + Ys2, Yvars=Yvars1 + Yvars2) with self.assertRaises(RuntimeError): unfit_model.feature_importances()