[refactor] Refactor _get_base_search_space in base_pipeline

Author: nabenabe0928

autoPyTorch/pipeline/base_pipeline.py

@@ -304,6 +304,101 @@ def __repr__(self) -> str:
         string += '_' * 40
         return string
 
+    def _get_search_space_modifications(
+        self,
+        include: Optional[Dict[str, Any]],
+        exclude: Optional[Dict[str, Any]],
+        pipeline: List[Tuple[str, autoPyTorchChoice]]
+    ) -> Tuple[Dict[str, Any], Dict[str, Any]]:
+        """
+        Get what pipeline or components to
+        include in or exclude from the searching space
+
+        Args:
+            include (Optional[Dict[str, Any]]): Components to include in the searching
+            exclude (Optional[Dict[str, Any]]): Components to exclude from the searching
+            pipeline (List[Tuple[str, autoPyTorchChoice]]):
+                Available components
+
+        Returns:
+            include, exclude (Tuple[Dict[str, Any]], Dict[str, Any]]]):
+                modified version of `include` and `exclude`
+        """
+
+        key_exist = {pair[0]: True for pair in pipeline}
+
+        if include is None:
+            include = {} if self.include is None else self.include
+
+        for key in include.keys():
+            if key_exist.get(key, False):
+                raise ValueError('Keys in `include` must be {}, but got {}'.format(
+                    list(key_exist.keys()), key)
+                )
+
+        if exclude is None:
+            exclude = {} if self.exclude is None else self.exclude
+
+        for key in exclude:
+            if key_exist.get(key, False):
+                raise ValueError('Keys in `exclude` must be {}, but got {}'.format(
+                    list(key_exist.keys()), key)
+                )
+
+        return include, exclude
+
+    @staticmethod
+    def _update_search_space_per_node(
+        cs: ConfigurationSpace,
+        dataset_properties: Dict[str, Any],
+        include: Optional[Dict[str, Any]],
+        exclude: Optional[Dict[str, Any]],
+        matches: np.ndarray,
+        node_name: str,
+        node_idx: int,
+        node: autoPyTorchChoice
+    ) -> ConfigurationSpace:
+        """
+        Args:
+            cs (ConfigurationSpace): Searching space information
+            dataset_properties (Dict[str, Any]): The properties of dataset
+            include (Optional[Dict[str, Any]]): Components to include in the searching
+            exclude (Optional[Dict[str, Any]]): Components to exclude from the searching
+            node_name (str): The name of the component choice
+            node_idx (int): The index of the component in a provided list of components
+            node (autoPyTorchChoice): The module of the component
+            matches (np.ndarray): ...
+
+        Returns:
+            modified cs (ConfigurationSpace):
+                modified searching space information based on the arguments.
+        """
+        is_choice = isinstance(node, autoPyTorchChoice)
+
+        if not is_choice:
+            # if the node isn't a choice we can add it immediately because
+            # it must be active (if it wasn't, np.sum(matches) would be zero
+            assert not isinstance(node, autoPyTorchChoice)
+            cs.add_configuration_space(
+                node_name,
+                node.get_hyperparameter_search_space(dataset_properties,  # type: ignore[arg-type]
+                                                     **node._get_search_space_updates()),
+            )
+        else:
+            # If the node is a choice, we have to figure out which of
+            # its choices are actually legal choices
+            choices_list = find_active_choices(
+                matches, node, node_idx,
+                dataset_properties,
+                include.get(node_name),
+                exclude.get(node_name)
+            )
+            sub_config_space = node.get_hyperparameter_search_space(
+                dataset_properties, include=choices_list)
+            cs.add_configuration_space(node_name, sub_config_space)
+
+        return cs
+
     def _get_base_search_space(
             self,
             cs: ConfigurationSpace,
@@ -312,75 +407,50 @@ def _get_base_search_space(
             exclude: Optional[Dict[str, Any]],
             pipeline: List[Tuple[str, autoPyTorchChoice]]
     ) -> ConfigurationSpace:
-        if include is None:
-            if self.include is None:
-                include = {}
-            else:
-                include = self.include
-
-        keys = [pair[0] for pair in pipeline]
-        for key in include:
-            if key not in keys:
-                raise ValueError('Invalid key in include: %s; should be one '
-                                 'of %s' % (key, keys))
-
-        if exclude is None:
-            if self.exclude is None:
-                exclude = {}
-            else:
-                exclude = self.exclude
+        """
+        Get the searching space
 
-        keys = [pair[0] for pair in pipeline]
-        for key in exclude:
-            if key not in keys:
-                raise ValueError('Invalid key in exclude: %s; should be one '
-                                 'of %s' % (key, keys))
+        Args:
+            cs (ConfigurationSpace): Searching space information
+            dataset_properties (Dict[str, Any]): The properties of dataset
+            include (Optional[Dict[str, Any]]): Components to include in the searching
+            exclude (Optional[Dict[str, Any]]): Components to exclude from the searching
+            pipeline (List[Tuple[str, autoPyTorchChoice]]):
+                Available components
 
+        Returns:
+            modified cs (ConfigurationSpace):
+                modified searching space information based on the arguments.
+        """
+        include, exclude = self._get_search_space_modifications(
+            include=include, exclude=exclude, pipeline=pipeline
+        )
         if self.search_space_updates is not None:
             self._check_search_space_updates(include=include,
                                              exclude=exclude)
             self.search_space_updates.apply(pipeline=pipeline)
 
+        # The size of this array exponentially grows, so it will be better to remove
         matches = get_match_array(
             pipeline, dataset_properties, include=include, exclude=exclude)
 
         # Now we have only legal combinations at this step of the pipeline
         # Simple sanity checks
-        assert np.sum(matches) != 0, "No valid pipeline found."
-
-        assert np.sum(matches) <= np.size(matches), \
-            "'matches' is not binary; %s <= %d, %s" % \
-            (str(np.sum(matches)), np.size(matches), str(matches.shape))
+        if np.sum(matches) == 0:
+            raise ValueError("No valid pipeline found.")
+        if np.sum(matches) > np.size(matches):
+            raise TypeError("'matches' is not binary; {} <= {}, {}".format(
+                np.sum(matches), np.size(matches), str(matches.shape)
+            ))
 
         # Iterate each dimension of the matches array (each step of the
         # pipeline) to see if we can add a hyperparameter for that step
-        for node_idx, n_ in enumerate(pipeline):
-            node_name, node = n_
-
-            is_choice = isinstance(node, autoPyTorchChoice)
-
-            # if the node isn't a choice we can add it immediately because it
-            #  must be active (if it wasn't, np.sum(matches) would be zero
-            if not is_choice:
-                # for mypy
-                assert not isinstance(node, autoPyTorchChoice)
-                cs.add_configuration_space(
-                    node_name,
-                    node.get_hyperparameter_search_space(dataset_properties,  # type: ignore[arg-type]
-                                                         **node._get_search_space_updates()),
-                )
-            # If the node is a choice, we have to figure out which of its
-            #  choices are actually legal choices
-            else:
-                choices_list = find_active_choices(
-                    matches, node, node_idx,
-                    dataset_properties,
-                    include.get(node_name),
-                    exclude.get(node_name)
-                )
-                sub_config_space = node.get_hyperparameter_search_space(
-                    dataset_properties, include=choices_list)
-                cs.add_configuration_space(node_name, sub_config_space)
+        for node_idx, (node_name, node) in enumerate(pipeline):
+            cs = self._update_search_space_per_node(
+                cs=ConfigurationSpace, dataset_properties=dataset_properties,
+                include=include, exclude=exclude, matches=matches,
+                node=node, node_idx=node_idx, node_name=node_name
+            )
 
         # And now add forbidden parameter configurations
         # According to matches