From 7235302d5ef85a63a18d28644608bc3e19cab8bb Mon Sep 17 00:00:00 2001
From: John Platt <john@carbonrelay.com>
Date: Thu, 21 Feb 2019 12:17:04 -0500
Subject: [PATCH] covered all parameter types

---
 .../src/algorithm_manager.py                  |  4 ++--
 pkg/suggestion/tests/parameter_values.yaml    |  8 +++++--
 pkg/suggestion/tests/study_config.yaml        | 12 ++++++++---
 .../tests/test_algorithm_manager.py           | 21 +++++++++++--------
 4 files changed, 29 insertions(+), 16 deletions(-)

diff --git a/pkg/suggestion/bayesianoptimization/src/algorithm_manager.py b/pkg/suggestion/bayesianoptimization/src/algorithm_manager.py
index 48084ac7c9b..3ff5115b7b6 100644
--- a/pkg/suggestion/bayesianoptimization/src/algorithm_manager.py
+++ b/pkg/suggestion/bayesianoptimization/src/algorithm_manager.py
@@ -15,8 +15,8 @@ def deal_with_discrete(feasible_values, current_value):
 
 def deal_with_categorical(feasible_values, one_hot_values):
     """ function to do the one hot encoding of the categorical values """
-    #index = np.argmax(one_hot_values)
-    index = one_hot_values.argmax()
+    index = np.argmax(one_hot_values)
+    #index = one_hot_values.argmax()
     return feasible_values[int(index)]
 
 
diff --git a/pkg/suggestion/tests/parameter_values.yaml b/pkg/suggestion/tests/parameter_values.yaml
index ec2c3fc8070..d9cef801e80 100644
--- a/pkg/suggestion/tests/parameter_values.yaml
+++ b/pkg/suggestion/tests/parameter_values.yaml
@@ -2,15 +2,19 @@ parameters:
   - - name: "x"
       value: 1.0
     - name: "y"
-      value: 1.0 
+      value: 1
     - name: "fake_discrete"
       value: 2
+    - name: "fake_categorical"
+      value: "true"
   - - name: "x"
       value: 1.0
     - name: "y"
-      value: 1.0
+      value: 1
     - name: "fake_discrete"
       value: 3
+    - name: "fake_categorical"
+      value: "false"
 metrics:
   - 1.0
   - 1.0
diff --git a/pkg/suggestion/tests/study_config.yaml b/pkg/suggestion/tests/study_config.yaml
index 8f8f3899839..7917899df03 100644
--- a/pkg/suggestion/tests/study_config.yaml
+++ b/pkg/suggestion/tests/study_config.yaml
@@ -6,11 +6,11 @@ parameter_configs:
       feasible:
         min: -5.0
         max: 5.0
-    - parameter_type: 1 # Code for DOUBLE
+    - parameter_type: 2 # Code for INT
       name: "y"
       feasible:
-        min: -5.0
-        max: 5.0
+        min: -5
+        max: 5
     - parameter_type: 3 # Code for DISCRETE
       name: "fake_discrete"
       feasible:
@@ -18,3 +18,9 @@ parameter_configs:
           - 2
           - 3
           - 5
+    - parameter_type: 4 # Code for DISCRETE
+      name: "fake_categorical"
+      feasible:
+        list:
+          - "true"
+          - "false"
diff --git a/pkg/suggestion/tests/test_algorithm_manager.py b/pkg/suggestion/tests/test_algorithm_manager.py
index a5cea0026d2..c7ea6c3df94 100644
--- a/pkg/suggestion/tests/test_algorithm_manager.py
+++ b/pkg/suggestion/tests/test_algorithm_manager.py
@@ -24,23 +24,26 @@ def observations():
 
 def test_algorithm_manager(study_config, observations):
     study_id = "test_id"
-    x_next = [1.0, 1.0, 5]
+    x_next = [1.0, 1, 5, "true"]
     manager = AlgorithmManager(study_id, study_config,
                                observations.parameters, observations.metrics)
     assert manager.study_id == study_id
     assert manager.study_config == study_config
     assert manager.goal == api_pb2.MAXIMIZE
-    assert manager.types == [api_pb2.DOUBLE, api_pb2.DOUBLE, api_pb2.DISCRETE]
-    assert manager.names == ["x", "y", "fake_discrete"]
-    assert manager.dim == 3
-    assert manager.lower_bound == [-5.0, -5.0, 2]
-    assert manager.upper_bound == [5.0, 5.0, 5]
+    assert manager.types == [api_pb2.DOUBLE, api_pb2.INT, api_pb2.DISCRETE, api_pb2.CATEGORICAL]
+    assert manager.names == ["x", "y", "fake_discrete", "fake_categorical"]
+    assert manager.dim == 5
+    assert manager.lower_bound == [-5.0, -5, 2, 0, 0]
+    assert manager.upper_bound == [5.0, 5, 5, 1, 1]
     assert manager.discrete_info == [{"name": "fake_discrete", "values": [2, 3, 5]}]
-    assert np.allclose(manager.X_train, np.array([[1.0, 1.0, 2], [1.0, 1.0, 3]]))
+    assert manager.categorical_info == \
+           [{"name": "fake_categorical", "values": ["true", "false"], "number": 2}]
+    assert np.allclose(manager.X_train, np.array([[1.0, 1, 2, 1, 0], [1.0, 1, 3, 0, 1]]))
     assert np.allclose(manager.y_train, np.array([1.0, 1.0]))
     parsed_x_next = manager.parse_x_next(x_next)
     x_next_dict = manager.convert_to_dict(parsed_x_next)
     assert x_next_dict == \
            [{"name": "x", "value": 1.0, "type": api_pb2.DOUBLE},
-            {"name": "y", "value": 1.0, "type": api_pb2.DOUBLE},
-            {"name": "fake_discrete", "value": 5, "type": api_pb2.DISCRETE}]
+            {"name": "y", "value": 1, "type": api_pb2.INT},
+            {"name": "fake_discrete", "value": 5, "type": api_pb2.DISCRETE},
+            {"name": "fake_categorical", "value": "true", "type": api_pb2.CATEGORICAL}]