fix anchor device, fix spectral bugs (cvxgrp#59)

SSamDav · Mar 28, 2022 · 0e9ab9d · 0e9ab9d
1 parent e8b1b4b
commit 0e9ab9d
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 7 deletions.
diff --git a/pymde/quadratic.py b/pymde/quadratic.py
@@ -72,7 +72,7 @@ def _spectral(
     L,
     m,
     cg=False,
-    max_iter=40,
+    max_iter=400,
     edges=None,
     weights=None,
     warm_start=False,
@@ -91,6 +91,7 @@ def _spectral(
             v0=np.ones(L.shape[0]),
             maxiter=L.shape[0] * 5,
         )
+        order = np.argsort(eigenvalues)[1:k]
     else:
         if warm_start:
             mde = problem.MDE(
@@ -114,7 +115,7 @@ def _spectral(
             largest=False,
             niter=max_iter,
         )
-    order = np.argsort(eigenvalues)[1:k]
+        order = torch.argsort(eigenvalues)[1:k]
     return eigenvectors[:, order]
 
 

diff --git a/pymde/recipes.py b/pymde/recipes.py
@@ -63,8 +63,8 @@ def _remove_anchor_anchor_edges(edges, data, anchors, opt=None, blk=1e8):
         estep = int(blk / na)
         assert estep > 0
 
-        msk0 = torch.zeros(ne, dtype=torch.bool, device=edges.get_device())
-        msk1 = torch.zeros(ne, dtype=torch.bool, device=edges.get_device())
+        msk0 = torch.zeros(ne, dtype=torch.bool, device=edges.device)
+        msk1 = torch.zeros(ne, dtype=torch.bool, device=edges.device)
         e = 0
         while e < ne:
             ehi = e + estep if e + estep < ne else ne
@@ -79,8 +79,8 @@ def _remove_anchor_anchor_edges(edges, data, anchors, opt=None, blk=1e8):
         assert astep > 0
         assert estep > 0
 
-        msk0 = torch.zeros(ne, dtype=torch.bool, device=edges.get_device())
-        msk1 = torch.zeros(ne, dtype=torch.bool, device=edges.get_device())
+        msk0 = torch.zeros(ne, dtype=torch.bool, device=edges.device)
+        msk1 = torch.zeros(ne, dtype=torch.bool, device=edges.device)
         e = 0
         while e < ne:
             ehi = e + estep if e + estep < ne else ne
@@ -365,7 +365,13 @@ def preserve_neighbors(
         # use cg + torch when using GPU
         cg = device == "cuda"
         X_init = quadratic.spectral(
-            n, embedding_dim, edges, weights, device=device, cg=cg
+            n,
+            embedding_dim,
+            edges,
+            weights,
+            max_iter=1000,
+            device=device,
+            cg=cg,
         )
         if not isinstance(
             constraint, (constraints._Centered, constraints._Standardized)