try to fix

Author: voorhs

autointent/nodes/_node_optimizer.py

@@ -3,11 +3,13 @@
 import gc
 import itertools as it
 import logging
+from abc import ABC, abstractmethod
 from copy import deepcopy
 from functools import partial
 from pathlib import Path
-from typing import Any
+from typing import Any, TypeVar
 
+import numpy as np
 import optuna
 import torch
 from optuna.trial import Trial
@@ -20,27 +22,65 @@
 from autointent.nodes.info import NODES_INFO
 
 
-class ParamSpaceInt(BaseModel):
+class ParamSpace(BaseModel, ABC):
+    """Base class for parameter search space configuration."""
+
+    @abstractmethod
+    def n_possible_values(self) -> int | None:
+        """Calculate the number of possible values in the search space.
+
+        Returns:
+            The number of possible values or None if search space is continuous.
+        """
+
+
+class ParamSpaceInt(ParamSpace):
     """Integer parameter search space configuration."""
 
     low: int = Field(..., description="Lower boundary of the search space.")
     high: int = Field(..., description="Upper boundary of the search space.")
     step: int = Field(1, description="Step size for the search space.")
     log: bool = Field(False, description="Indicates whether to use a logarithmic scale.")
 
+    def n_possible_values(self) -> int:
+        """Calculate the number of possible values in the search space.
 
-class ParamSpaceFloat(BaseModel):
+        Returns:
+            The number of possible values.
+        """
+        if self.log:
+            return int(np.logspace(np.log10(self.low), np.log10(self.high), num=self.step))
+
+        return (self.high - self.low) // self.step + 1
+
+
+class ParamSpaceFloat(ParamSpace):
     """Float parameter search space configuration."""
 
     low: float = Field(..., description="Lower boundary of the search space.")
     high: float = Field(..., description="Upper boundary of the search space.")
     step: float | None = Field(None, description="Step size for the search space (if applicable).")
     log: bool = Field(False, description="Indicates whether to use a logarithmic scale.")
 
+    def n_possible_values(self) -> int | None:
+        """Calculate the number of possible values in the search space.
+
+        Returns:
+            The number of possible values or None if search space is continuous.
+        """
+        if self.step is None:
+            return None
+        if self.log:
+            return int(np.logspace(np.log10(self.low), np.log10(self.high), num=self.step))
+        return (self.high - self.low) // self.step + 1
+
 
 logger = logging.getLogger(__name__)
 
 
+ParamSpaceT = TypeVar("ParamSpaceT", bound=ParamSpace)
+
+
 class NodeOptimizer:
     """Class for optimizing nodes in a computational pipeline.
 
@@ -104,6 +144,8 @@ def fit(self, context: Context, sampler: SamplerType = "brute", n_jobs: int = 1)
             else:
                 assert_never(sampler)
 
+            n_trials = None if n_trials is None else min(self._n_possible_combinations(search_space), n_trials)
+
             study, finished_trials, n_trials = load_or_create_study(
                 study_name=f"{self.node_info.node_type}_{module_name}",
                 storage_dir=context.get_dump_dir(),
@@ -205,23 +247,42 @@ def suggest(self, trial: Trial, search_space: dict[str, Any | list[Any]]) -> dic
         for param_name, param_space in search_space.items():
             if isinstance(param_space, list):
                 res[param_name] = trial.suggest_categorical(param_name, choices=param_space)
-            elif self._is_valid_param_space(param_space, ParamSpaceInt):
+            elif self._parse_param_space(param_space, ParamSpaceInt):
                 res[param_name] = trial.suggest_int(param_name, **param_space)
-            elif self._is_valid_param_space(param_space, ParamSpaceFloat):
+            elif self._parse_param_space(param_space, ParamSpaceFloat):
                 res[param_name] = trial.suggest_float(param_name, **param_space)
             else:
                 msg = f"Unsupported type of param search space: {param_space}"
                 raise TypeError(msg)
         return res
 
-    def _is_valid_param_space(
-        self, param_space: dict[str, Any], space_type: type[ParamSpaceInt | ParamSpaceFloat]
-    ) -> bool:
+    def _n_possible_combinations(self, search_space: dict[str, Any]) -> int:
+        """Calculate the number of possible combinations in the search space.
+
+        Args:
+            search_space: The parameter search space.
+        """
+        n_combinations = 1
+        for param_space in search_space.values():
+            if isinstance(param_space, list):
+                n_combinations *= len(param_space)
+            else:
+                param_space_int = self._parse_param_space(param_space, ParamSpaceInt)
+                if param_space_int is not None:
+                    n_combinations *= param_space_int.n_possible_values()
+                    continue
+                param_space_float = self._parse_param_space(param_space, ParamSpaceFloat)
+                if param_space_float is not None:
+                    n_combinations *= param_space_float.n_possible_values()
+                    continue
+                assert_never(param_space)
+        return n_combinations
+
+    def _parse_param_space(self, param_space: dict[str, Any], space_type: type[ParamSpaceT]) -> ParamSpaceT | None:
         try:
-            space_type(**param_space)
-            return True  # noqa: TRY300
+            return space_type(**param_space)
         except ValueError:
-            return False
+            return None
 
     def get_module_dump_dir(self, dump_dir: Path, module_name: str, j_combination: int) -> str:
         """Creates and returns the path to the module dump directory.