google-research

1
// Copyright 2024 The Google Research Authors.
2
//
3
// Licensed under the Apache License, Version 2.0 (the "License");
4
// you may not use this file except in compliance with the License.
5
// You may obtain a copy of the License at
6
//
7
//     http://www.apache.org/licenses/LICENSE-2.0
8
//
9
// Unless required by applicable law or agreed to in writing, software
10
// distributed under the License is distributed on an "AS IS" BASIS,
11
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
// See the License for the specific language governing permissions and
13
// limitations under the License.
14

15
#include "train_budget.h"
16

17
#include "algorithm.h"
18
#include "generator_test_util.h"
19
#include "gtest/gtest.h"
20

21
namespace automl_zero {
22

23
TEST(TrainBudgetTest, MatchesBaselineExactly) {
24
  Algorithm baseline_algorithm = SimpleGrTildeGrWithBias();
25
  TrainBudget train_budget(baseline_algorithm, 2.0);
26
  Algorithm algorithm = baseline_algorithm;
27
  EXPECT_EQ(train_budget.TrainExamples(algorithm, 100), 100);
28
  EXPECT_EQ(train_budget.TrainExamples(algorithm, 1000), 1000);
29
}
30

31
TEST(TrainBudgetTest, CheaperModelIsUnaffected) {
32
  Algorithm baseline_algorithm = SimpleGrTildeGrWithBias();
33
  TrainBudget train_budget(baseline_algorithm, 2.0);
34
  Algorithm algorithm = SimpleGz();
35
  EXPECT_EQ(train_budget.TrainExamples(algorithm, 100), 100);
36
  EXPECT_EQ(train_budget.TrainExamples(algorithm, 1000), 1000);
37
}
38

39
TEST(TrainBudgetTest, MoreExpensiveModelIsRejected) {
40
  Algorithm baseline_algorithm = SimpleGz();
41
  TrainBudget train_budget(baseline_algorithm, 2.0);
42
  Algorithm algorithm = SimpleGrTildeGrWithBias();
43
  EXPECT_EQ(train_budget.TrainExamples(algorithm, 100), 0);
44
  EXPECT_EQ(train_budget.TrainExamples(algorithm, 1000), 0);
45
}
46

47
}  // namespace automl_zero
48