Cudf

graphistry/feature_utils.py

@@ -3,6 +3,7 @@
 import os
 import pandas as pd
 from time import time
+from inspect import getmodule
 import warnings
 from functools import partial
 
@@ -167,7 +168,7 @@ def resolve_feature_engine(
 
 def resolve_y(df: Optional[pd.DataFrame], y: YSymbolic) -> pd.DataFrame:
 
-    if isinstance(y, pd.DataFrame):
+    if isinstance(y, pd.DataFrame) or 'cudf.core.dataframe' in str(getmodule(y)):
         return y
 
     if df is None:
@@ -188,7 +189,7 @@ def resolve_y(df: Optional[pd.DataFrame], y: YSymbolic) -> pd.DataFrame:
 
 def resolve_X(df: Optional[pd.DataFrame], X: XSymbolic) -> pd.DataFrame:
 
-    if isinstance(X, pd.DataFrame):
+    if isinstance(X, pd.DataFrame) or 'cudf.core.dataframe' in str(getmodule(X)):
         return X
 
     if df is None:

graphistry/umap_utils.py

@@ -1,6 +1,7 @@
 import copy
 from time import time
 from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple, Union
+from inspect import getmodule
 
 import pandas as pd
 
@@ -282,13 +283,20 @@ def transform_umap(  # noqa: E303
 
     def _bundle_embedding(self, emb, index):
         # Converts Embedding into dataframe and takes care if emb.dim > 2
-        if emb.shape[1] == 2:
+        if emb.shape[1] == 2 and 'cudf.core.dataframe' not in str(getmodule(emb)):
             emb = pd.DataFrame(emb, columns=[config.X, config.Y], index=index)
+        elif emb.shape[1] == 2 and 'cudf.core.dataframe' in str(getmodule(emb)):
+            import cudf
+            emb = cudf.DataFrame(emb.values, columns=[config.X, config.Y], index=index.values)
         else:
             columns = [config.X, config.Y] + [
                 f"umap_{k}" for k in range(2, emb.shape[1] - 2)
             ]
-            emb = pd.DataFrame(emb, columns=columns, index=index)
+            if 'cudf.core.dataframe' not in str(getmodule(emb)):
+                emb = pd.DataFrame(emb, columns=columns, index=index)
+            elif 'cudf.core.dataframe' in str(getmodule(emb)):
+                import cudf
+                emb = cudf.DataFrame(emb.values, columns=columns, index=index.values)
         return emb
 
     def _process_umap(
@@ -322,7 +330,6 @@ def _process_umap(
             # have to set _raw_data attribute on umap?
             fresh_res._umap = old_res._umap  # this saves the day!
             return fresh_res
-
         emb = res.umap_fit_transform(X_, y_)
         res._xy = emb
         return res
@@ -478,7 +485,6 @@ def umap(
                 )
 
             nodes = res._nodes[res._node].values
-            index_to_nodes_dict = dict(zip(range(len(nodes)), nodes))
 
             logger.debug("propagating with featurize_kwargs: %s", featurize_kwargs)
             (
@@ -492,6 +498,12 @@ def umap(
             logger.debug("umap X_: %s", X_)
             logger.debug("umap y_: %s", y_)
 
+            if isinstance(X_,pd.DataFrame):
+                index_to_nodes_dict = dict(zip(range(len(nodes)), nodes))
+            elif 'cudf.core.dataframe' in str(getmodule(X_)):
+                import cudf
+                index_to_nodes_dict = cudf.DataFrame(nodes).reset_index()
+
             res = res._process_umap(
                 res, X_, y_, kind, memoize, featurize_kwargs, **umap_kwargs
             )
@@ -500,6 +512,7 @@ def umap(
             if res._xy is None:
                 raise RuntimeError("This should not happen")
             res._node_embedding = res._xy
+
             # TODO add edge filter so graph doesn't have double edges
             # TODO user-guidable edge merge policies like upsert?
             res._weighted_edges_df_from_nodes = (
@@ -586,11 +599,10 @@ def _bind_xy_from_umap(
             emb = res._node_embedding
         else:
             emb = res._edge_embedding
-
-        df[x_name] = emb.values.T[0]  # if embedding is greater
-        # than two dimensions will only take first two coordinates
+
+        df[x_name] = emb.values.T[0]
         df[y_name] = emb.values.T[1]
-        #
+
         res = res.nodes(df) if kind == "nodes" else res.edges(df)
 
         if encode_weight and kind == "nodes":