docker-python/tests/test_rapids.py at main · hokbin/docker-python · GitHub

22 lines (16 loc) · 572 Bytes

import unittest

from common import gpu_test

class TestRapids(unittest.TestCase):

@gpu_test

def test_dbscan(self):

import cudf

from cuml.cluster import DBSCAN

# Create and populate a GPU DataFrame

gdf_float = cudf.DataFrame()

gdf_float['0'] = [1.0, 2.0, 5.0]

gdf_float['1'] = [4.0, 2.0, 1.0]

gdf_float['2'] = [4.0, 2.0, 1.0]

# Setup and fit clusters

dbscan_float = DBSCAN(eps=1.0, min_samples=1)

dbscan_float.fit(gdf_float)

self.assertEqual(3, dbscan_float.labels_.size)

Expand file tree