Move feature classes to features.py

miplearn/components/component.py

@@ -7,7 +7,8 @@ from typing import Any, List, TYPE_CHECKING, Tuple, Dict, Hashable
 import numpy as np
 
 from miplearn.instance import Instance
-from miplearn.types import LearningSolveStats, TrainingSample, Features
+from miplearn.types import LearningSolveStats
+from miplearn.features import TrainingSample, Features
 
 if TYPE_CHECKING:
     from miplearn.solvers.learning import LearningSolver
@@ -49,7 +50,7 @@ class Component:
             their own statistics here. For example, PrimalSolutionComponent adds
             statistics regarding the number of predicted variables. All statistics in
             this dictionary are exported to the benchmark CSV file.
-        features: Features
+        features: miplearn.features.Features
             Features describing the model.
         training_data: TrainingSample
             A dictionary containing data that may be useful for training machine

miplearn/components/lazy_static.py

@@ -11,7 +11,8 @@ from miplearn import Classifier
 from miplearn.classifiers.counting import CountingClassifier
 from miplearn.classifiers.threshold import MinProbabilityThreshold, Threshold
 from miplearn.components.component import Component
-from miplearn.types import TrainingSample, Features, LearningSolveStats
+from miplearn.types import LearningSolveStats
+from miplearn.features import TrainingSample, Features
 
 logger = logging.getLogger(__name__)
 

miplearn/components/objective.py

@@ -12,7 +12,8 @@ from miplearn.classifiers import Regressor
 from miplearn.classifiers.sklearn import ScikitLearnRegressor
 from miplearn.components.component import Component
 from miplearn.instance import Instance
-from miplearn.types import TrainingSample, LearningSolveStats, Features
+from miplearn.types import LearningSolveStats
+from miplearn.features import TrainingSample, Features
 
 if TYPE_CHECKING:
     from miplearn.solvers.learning import LearningSolver

miplearn/components/primal.py

@@ -22,11 +22,10 @@ from miplearn.components import classifier_evaluation_dict
 from miplearn.components.component import Component
 from miplearn.instance import Instance
 from miplearn.types import (
-    TrainingSample,
     Solution,
     LearningSolveStats,
-    Features,
 )
+from miplearn.features import TrainingSample, Features
 
 logger = logging.getLogger(__name__)
 

miplearn/features.py

@@ -4,20 +4,57 @@
 
 import collections
 import numbers
-from typing import TYPE_CHECKING, Dict
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Dict, Optional, Set, List, Hashable
 
-from miplearn.types import (
-    Features,
-    ConstraintFeatures,
-    InstanceFeatures,
-    VariableFeatures,
-    VarIndex,
-)
+from miplearn.types import VarIndex, Solution
 
 if TYPE_CHECKING:
     from miplearn import InternalSolver, Instance
 
 
+@dataclass
+class TrainingSample:
+    lp_log: Optional[str] = None
+    lp_solution: Optional[Solution] = None
+    lp_value: Optional[float] = None
+    lazy_enforced: Optional[Set[str]] = None
+    lower_bound: Optional[float] = None
+    mip_log: Optional[str] = None
+    solution: Optional[Solution] = None
+    upper_bound: Optional[float] = None
+    slacks: Optional[Dict[str, float]] = None
+
+
+@dataclass
+class InstanceFeatures:
+    user_features: Optional[List[float]] = None
+    lazy_constraint_count: int = 0
+
+
+@dataclass
+class VariableFeatures:
+    category: Optional[Hashable] = None
+    user_features: Optional[List[float]] = None
+
+
+@dataclass
+class ConstraintFeatures:
+    rhs: Optional[float] = None
+    lhs: Optional[Dict[str, float]] = None
+    sense: Optional[str] = None
+    category: Optional[Hashable] = None
+    user_features: Optional[List[float]] = None
+    lazy: bool = False
+
+
+@dataclass
+class Features:
+    instance: Optional[InstanceFeatures] = None
+    variables: Optional[Dict[str, Dict[VarIndex, VariableFeatures]]] = None
+    constraints: Optional[Dict[str, ConstraintFeatures]] = None
+
+
 class FeaturesExtractor:
     def __init__(
         self,

miplearn/instance.py

@@ -9,7 +9,8 @@ import pickle
 from abc import ABC, abstractmethod
 from typing import Any, List, Optional, Hashable, IO, cast
 
-from miplearn.types import TrainingSample, VarIndex, Features
+from miplearn.types import VarIndex
+from miplearn.features import TrainingSample, Features
 
 logger = logging.getLogger(__name__)
 

miplearn/solvers/learning.py

@@ -13,12 +13,12 @@ from miplearn.components.cuts import UserCutsComponent
 from miplearn.components.lazy_dynamic import DynamicLazyConstraintsComponent
 from miplearn.components.objective import ObjectiveValueComponent
 from miplearn.components.primal import PrimalSolutionComponent
-from miplearn.features import FeaturesExtractor
+from miplearn.features import FeaturesExtractor, TrainingSample
 from miplearn.instance import Instance, PickleGzInstance
 from miplearn.solvers import _RedirectOutput
 from miplearn.solvers.internal import InternalSolver
 from miplearn.solvers.pyomo.gurobi import GurobiPyomoSolver
-from miplearn.types import TrainingSample, LearningSolveStats
+from miplearn.types import LearningSolveStats
 
 logger = logging.getLogger(__name__)
 

miplearn/types.py

@@ -2,8 +2,7 @@
 #  Copyright (C) 2020, UChicago Argonne, LLC. All rights reserved.
 #  Released under the modified BSD license. See COPYING.md for more details.
 
-from typing import Optional, Dict, Callable, Any, Union, Tuple, List, Set, Hashable
-from dataclasses import dataclass
+from typing import Optional, Dict, Callable, Any, Union, Tuple
 
 from mypy_extensions import TypedDict
 
@@ -11,20 +10,6 @@ VarIndex = Union[str, int, Tuple[Union[str, int]]]
 
 Solution = Dict[str, Dict[VarIndex, Optional[float]]]
 
-
-@dataclass
-class TrainingSample:
-    lp_log: Optional[str] = None
-    lp_solution: Optional[Solution] = None
-    lp_value: Optional[float] = None
-    lazy_enforced: Optional[Set[str]] = None
-    lower_bound: Optional[float] = None
-    mip_log: Optional[str] = None
-    solution: Optional[Solution] = None
-    upper_bound: Optional[float] = None
-    slacks: Optional[Dict[str, float]] = None
-
-
 LPSolveStats = TypedDict(
     "LPSolveStats",
     {
@@ -75,36 +60,6 @@ LearningSolveStats = TypedDict(
     total=False,
 )
 
-
-@dataclass
-class InstanceFeatures:
-    user_features: Optional[List[float]] = None
-    lazy_constraint_count: int = 0
-
-
-@dataclass
-class VariableFeatures:
-    category: Optional[Hashable] = None
-    user_features: Optional[List[float]] = None
-
-
-@dataclass
-class ConstraintFeatures:
-    rhs: Optional[float] = None
-    lhs: Optional[Dict[str, float]] = None
-    sense: Optional[str] = None
-    category: Optional[Hashable] = None
-    user_features: Optional[List[float]] = None
-    lazy: bool = False
-
-
-@dataclass
-class Features:
-    instance: Optional[InstanceFeatures] = None
-    variables: Optional[Dict[str, Dict[VarIndex, VariableFeatures]]] = None
-    constraints: Optional[Dict[str, ConstraintFeatures]] = None
-
-
 IterationCallback = Callable[[], bool]
 
 LazyCallback = Callable[[Any, Any], None]