problems: Allow correlated arguments in random problem generators

miplearn/problems/setcover.py

@@ -3,7 +3,7 @@
 #  Released under the modified BSD license. See COPYING.md for more details.
 
 from dataclasses import dataclass
-from typing import List, Union
+from typing import List, Union, Callable
 
 import gurobipy as gp
 import numpy as np
@@ -34,7 +34,7 @@ class SetCoverGenerator:
     def __init__(
         self,
         n_elements: rv_frozen = randint(low=50, high=51),
-        n_sets: rv_frozen = randint(low=100, high=101),
+        n_sets: Union[rv_frozen, Callable] = randint(low=100, high=101),
         costs: rv_frozen = uniform(loc=0.0, scale=100.0),
         K: rv_frozen = uniform(loc=25.0, scale=0.0),
         density: rv_frozen = uniform(loc=0.02, scale=0.00),
@@ -45,8 +45,9 @@ class SetCoverGenerator:
         ----------
         n_elements: rv_discrete
             Probability distribution for number of elements.
-        n_sets: rv_discrete
-            Probability distribution for number of sets.
+        n_sets: rv_discrete or callable
+            Probability distribution for number of sets, or a callable that takes
+            the number of elements and returns the number of sets.
         costs: rv_continuous
             Probability distribution for base set costs.
         K: rv_continuous
@@ -57,9 +58,9 @@ class SetCoverGenerator:
         assert isinstance(
             n_elements, rv_frozen
         ), "n_elements should be a SciPy probability distribution"
-        assert isinstance(
-            n_sets, rv_frozen
-        ), "n_sets should be a SciPy probability distribution"
+        assert isinstance(n_sets, rv_frozen) or callable(
+            n_sets
+        ), "n_sets should be a SciPy probability distribution or callable"
         assert isinstance(
             costs, rv_frozen
         ), "costs should be a SciPy probability distribution"
@@ -75,8 +76,11 @@ class SetCoverGenerator:
 
     def generate(self, n_samples: int) -> List[SetCoverData]:
         def _sample() -> SetCoverData:
-            n_sets = self.n_sets.rvs()
             n_elements = self.n_elements.rvs()
+            if callable(self.n_sets):
+                n_sets = self.n_sets(n_elements)
+            else:
+                n_sets = self.n_sets.rvs()
             density = self.density.rvs()
 
             incidence_matrix = np.random.rand(n_elements, n_sets) < density