From 03e5acb11a2560ba00707d53fc766c6f8e97b931 Mon Sep 17 00:00:00 2001
From: "Alinson S. Xavier" <git@axavier.org>
Date: Tue, 22 Feb 2022 09:20:17 -0600
Subject: [PATCH] Make MultiKnapsackGenerator return data class

---
 miplearn/problems/knapsack.py   | 37 ++++++---------------------------
 tests/problems/test_knapsack.py | 17 +++++++++------
 2 files changed, 17 insertions(+), 37 deletions(-)

diff --git a/miplearn/problems/knapsack.py b/miplearn/problems/knapsack.py
index 1dd06ef..3ca9096 100644
--- a/miplearn/problems/knapsack.py
+++ b/miplearn/problems/knapsack.py
@@ -1,7 +1,7 @@
 #  MIPLearn: Extensible Framework for Learning-Enhanced Mixed-Integer Optimization
 #  Copyright (C) 2020-2021, UChicago Argonne, LLC. All rights reserved.
 #  Released under the modified BSD license. See COPYING.md for more details.
-
+from dataclasses import dataclass
 from typing import List, Dict, Optional
 
 import numpy as np
@@ -13,36 +13,11 @@ from scipy.stats.distributions import rv_frozen
 from miplearn.instance.base import Instance
 
 
-class ChallengeA:
-    """
-    - 250 variables, 10 constraints, fixed weights
-    - w ~ U(0, 1000), jitter ~ U(0.95, 1.05)
-    - K = 500, u ~ U(0., 1.)
-    - alpha = 0.25
-    """
-
-    def __init__(
-        self,
-        seed: int = 42,
-        n_training_instances: int = 500,
-        n_test_instances: int = 50,
-    ) -> None:
-        np.random.seed(seed)
-        self.gen = MultiKnapsackGenerator(
-            n=randint(low=250, high=251),
-            m=randint(low=10, high=11),
-            w=uniform(loc=0.0, scale=1000.0),
-            K=uniform(loc=500.0, scale=0.0),
-            u=uniform(loc=0.0, scale=1.0),
-            alpha=uniform(loc=0.25, scale=0.0),
-            fix_w=True,
-            w_jitter=uniform(loc=0.95, scale=0.1),
-        )
-        np.random.seed(seed + 1)
-        self.training_instances = self.gen.generate(n_training_instances)
-
-        np.random.seed(seed + 2)
-        self.test_instances = self.gen.generate(n_test_instances)
+@dataclass
+class MultiKnapsackData:
+    prices: np.ndarray
+    capacities: np.ndarray
+    weights: np.ndarray
 
 
 class MultiKnapsackInstance(Instance):
diff --git a/tests/problems/test_knapsack.py b/tests/problems/test_knapsack.py
index 06a1ee8..59561d0 100644
--- a/tests/problems/test_knapsack.py
+++ b/tests/problems/test_knapsack.py
@@ -6,7 +6,7 @@ import numpy as np
 from scipy.stats import uniform, randint
 
 from miplearn import LearningSolver
-from miplearn.problems.knapsack import MultiKnapsackGenerator
+from miplearn.problems.knapsack import MultiKnapsackGenerator, MultiKnapsackInstance
 
 
 def test_knapsack_generator() -> None:
@@ -18,17 +18,22 @@ def test_knapsack_generator() -> None:
         u=uniform(loc=1.0, scale=1.0),
         alpha=uniform(loc=0.50, scale=0.0),
     )
-    instances = gen.generate(100)
-    w_sum = sum(instance.weights for instance in instances) / len(instances)
-    b_sum = sum(instance.capacities for instance in instances) / len(instances)
+    data = gen.generate(100)
+    w_sum = sum(d.weights for d in data) / len(data)
+    b_sum = sum(d.capacities for d in data) / len(data)
     assert round(float(np.mean(w_sum)), -1) == 500.0
     assert round(float(np.mean(b_sum)), -3) == 25000.0
 
 
 def test_knapsack() -> None:
-    instance = MultiKnapsackGenerator(
+    data = MultiKnapsackGenerator(
         n=randint(low=5, high=6),
         m=randint(low=5, high=6),
-    ).generate(1)[0]
+    ).generate(1)
+    instance = MultiKnapsackInstance(
+        prices=data[0].prices,
+        capacities=data[0].capacities,
+        weights=data[0].weights,
+    )
     solver = LearningSolver()
     solver.solve(instance)