Including tests for make_classification and make_regression

Signed-off-by: Julio Faracco <[email protected]>

tests/datasets/test_datasets.py

@@ -17,6 +17,8 @@
 from dasf.utils.types import is_dask_gpu_array
 
 from dasf.datasets import make_blobs
+from dasf.datasets import make_classification
+from dasf.datasets import make_regression
 
 
 class TestDatasets(unittest.TestCase):
@@ -81,3 +83,120 @@ def test_make_blobs_mgpu(self):
 
             # Compute everything to gracefully shutdown
             client.close()
+
+    @patch('dasf.datasets.datasets.is_gpu_supported', Mock(return_value=False))
+    @patch('dasf.datasets.datasets.is_dask_supported', Mock(return_value=False))
+    @patch('dasf.datasets.datasets.is_dask_gpu_supported', Mock(return_value=False))
+    def test_make_classification_cpu(self):
+        n_samples = 500000
+
+        X, y = make_classification(n_samples=n_samples, n_classes=2,
+                                   shuffle=False, random_state=42)
+
+        self.assertTrue(is_cpu_array(X))
+        self.assertTrue(is_cpu_array(y))
+
+    @patch('dasf.datasets.datasets.is_gpu_supported', Mock(return_value=False))
+    @patch('dasf.datasets.datasets.is_dask_supported', Mock(return_value=True))
+    @patch('dasf.datasets.datasets.is_dask_gpu_supported', Mock(return_value=False))
+    def test_make_classification_mcpu(self):
+        n_samples = 500000
+
+        X, y = make_classification(n_samples=n_samples, n_classes=2, shuffle=False,
+                                   random_state=42, chunks=(5000))
+
+        self.assertTrue(is_dask_array(X))
+        self.assertTrue(is_dask_array(y))
+
+    @unittest.skipIf(not is_gpu_supported(),
+                     "not supported CUDA in this platform")
+    @patch('dasf.datasets.datasets.is_gpu_supported', Mock(return_value=True))
+    @patch('dasf.datasets.datasets.is_dask_supported', Mock(return_value=False))
+    @patch('dasf.datasets.datasets.is_dask_gpu_supported', Mock(return_value=False))
+    def test_make_classification_gpu(self):
+        n_samples = 500000
+
+        X, y = make_classification(n_samples=n_samples, n_classes=2,
+                                   shuffle=False, random_state=42)
+
+        self.assertTrue(is_gpu_array(X))
+        self.assertTrue(is_gpu_array(y))
+
+    @unittest.skipIf(not is_gpu_supported(),
+                     "not supported CUDA in this platform")
+    @patch('dasf.datasets.datasets.is_gpu_supported', Mock(return_value=False))
+    @patch('dasf.datasets.datasets.is_dask_supported', Mock(return_value=False))
+    @patch('dasf.datasets.datasets.is_dask_gpu_supported', Mock(return_value=True))
+    def test_make_classification_mgpu(self):
+        with LocalCUDACluster() as cluster:
+            client = Client(cluster)
+
+            n_samples = 500000
+
+            X, y = make_classification(n_samples=n_samples, n_classes=2,
+                                       shuffle=False, random_state=42)
+
+            self.assertTrue(is_dask_gpu_array(X))
+            self.assertTrue(is_dask_gpu_array(y))
+
+            # Compute everything to gracefully shutdown
+            client.close()
+
+    @patch('dasf.datasets.datasets.is_gpu_supported', Mock(return_value=False))
+    @patch('dasf.datasets.datasets.is_dask_supported', Mock(return_value=False))
+    @patch('dasf.datasets.datasets.is_dask_gpu_supported', Mock(return_value=False))
+    def test_make_regression_cpu(self):
+        n_samples = 500000
+
+        X, y = make_regression(n_samples=n_samples, n_features=2,
+                               noise=1, random_state=42)
+
+        self.assertTrue(is_cpu_array(X))
+        self.assertTrue(is_cpu_array(y))
+
+    @patch('dasf.datasets.datasets.is_gpu_supported', Mock(return_value=False))
+    @patch('dasf.datasets.datasets.is_dask_supported', Mock(return_value=True))
+    @patch('dasf.datasets.datasets.is_dask_gpu_supported', Mock(return_value=False))
+    def test_make_regression_mcpu(self):
+        n_samples = 500000
+
+        X, y = make_regression(n_samples=n_samples, n_features=2,
+                               noise=1, random_state=42,
+                               chunks=(5000, 2))
+
+        self.assertTrue(is_dask_array(X))
+        self.assertTrue(is_dask_array(y))
+
+    @unittest.skipIf(not is_gpu_supported(),
+                     "not supported CUDA in this platform")
+    @patch('dasf.datasets.datasets.is_gpu_supported', Mock(return_value=True))
+    @patch('dasf.datasets.datasets.is_dask_supported', Mock(return_value=False))
+    @patch('dasf.datasets.datasets.is_dask_gpu_supported', Mock(return_value=False))
+    def test_make_regression_gpu(self):
+        n_samples = 500000
+
+        X, y = make_regression(n_samples=n_samples, n_features=2,
+                               noise=1, random_state=42)
+
+        self.assertTrue(is_gpu_array(X))
+        self.assertTrue(is_gpu_array(y))
+
+    @unittest.skipIf(not is_gpu_supported(),
+                     "not supported CUDA in this platform")
+    @patch('dasf.datasets.datasets.is_gpu_supported', Mock(return_value=False))
+    @patch('dasf.datasets.datasets.is_dask_supported', Mock(return_value=False))
+    @patch('dasf.datasets.datasets.is_dask_gpu_supported', Mock(return_value=True))
+    def test_make_regression_mgpu(self):
+        with LocalCUDACluster() as cluster:
+            client = Client(cluster)
+
+            n_samples = 500000
+
+            X, y = make_regression(n_samples=n_samples, n_features=2,
+                                   noise=1, random_state=42)
+
+            self.assertTrue(is_dask_gpu_array(X))
+            self.assertTrue(is_dask_gpu_array(y))
+
+            # Compute everything to gracefully shutdown
+            client.close()