Make Sample abstract; create MemorySample

miplearn/features/sample.py

@@ -2,10 +2,25 @@
 #  Copyright (C) 2020-2021, UChicago Argonne, LLC. All rights reserved.
 #  Released under the modified BSD license. See COPYING.md for more details.
 
+from abc import ABC, abstractmethod
 from typing import Dict, Optional, Any
 
 
-class Sample:
+class Sample(ABC):
+    """Abstract dictionary-like class that stores training data."""
+
+    @abstractmethod
+    def get(self, key: str) -> Optional[Any]:
+        pass
+
+    @abstractmethod
+    def put(self, key: str, value: Any) -> None:
+        pass
+
+
+class MemorySample(Sample):
+    """Dictionary-like class that stores training data in-memory."""
+
     def __init__(
         self,
         data: Optional[Dict[str, Any]] = None,

miplearn/solvers/learning.py

@@ -15,7 +15,7 @@ from miplearn.components.dynamic_user_cuts import UserCutsComponent
 from miplearn.components.objective import ObjectiveValueComponent
 from miplearn.components.primal import PrimalSolutionComponent
 from miplearn.features.extractor import FeaturesExtractor
-from miplearn.features.sample import Sample
+from miplearn.features.sample import Sample, MemorySample
 from miplearn.instance.base import Instance
 from miplearn.instance.picklegz import PickleGzInstance
 from miplearn.solvers import _RedirectOutput
@@ -150,7 +150,7 @@ class LearningSolver:
 
         # Initialize training sample
         # -------------------------------------------------------
-        sample = Sample()
+        sample = MemorySample()
         instance.push_sample(sample)
 
         # Initialize stats