From 2c9cc44b4c112e07f219e013a8984fcc6fe697b8 Mon Sep 17 00:00:00 2001
From: Ema Puljak <ema.puljak@cern.ch>
Date: Tue, 14 Jun 2022 16:05:14 +0200
Subject: [PATCH] Files updated

---
 qkmedians/distance_calc.py |  2 +-
 qkmedians/qkmedians.py     | 12 ++++++------
 2 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/qkmedians/distance_calc.py b/qkmedians/distance_calc.py
index 6c56140..8323e43 100644
--- a/qkmedians/distance_calc.py
+++ b/qkmedians/distance_calc.py
@@ -10,7 +10,7 @@ def pad_input(X):
         X = np.pad(X, (0, size_needed-num_features), "constant")
     return X
 
-def DistCalc_DI(a, b, device_name, shots_n=10000):
+def DistCalc_DI(a, b, device_name='/GPU:0', shots_n=10000):
     """ Distance calculation using destructive interference """
     num_features = len(a)
     norm = calc_norm(a, b)
diff --git a/qkmedians/qkmedians.py b/qkmedians/qkmedians.py
index 510be0b..b10c6a4 100644
--- a/qkmedians/qkmedians.py
+++ b/qkmedians/qkmedians.py
@@ -14,7 +14,7 @@ def initialize_centroids(points, k):
     indexes = np.random.randint(points.shape[0], size=k)
     return points[indexes]
 
-def find_distance_matrix_quantum(points, centroid):
+def find_distance_matrix_quantum(points, centroid, device_name):
     """ 
     Modified version of scipy.spatial.distance.cdist() function.
     Args:
@@ -35,11 +35,11 @@ def find_distance_matrix_quantum(points, centroid):
 
     dist_matrix = np.zeros((n_events, centroid.shape[0]))
     for i in range(n_events):
-        distance, _ = distc.DistCalc_DI(points[i,:], centroid[0])
+        distance, _ = distc.DistCalc_DI(points[i,:], centroid[0], device_name)
         dist_matrix[i,:] = distance
     return dist_matrix
 
-def geometric_median(points, eps=1e-6):
+def geometric_median(points, eps=1e-6, device_name='/GPU:0'):
     """
     Implementation from Reference - DOI: 10.1007/s00180-011-0262-4
     Args:
@@ -55,7 +55,7 @@ def geometric_median(points, eps=1e-6):
     median = np.mean(points, 0) # starting median
 
     while True:
-        D = find_distance_matrix_quantum(points, [y])
+        D = find_distance_matrix_quantum(points, [median], device_name)
         nonzeros = (D != 0)[:, 0]
         Dinv = 1 / D[nonzeros]
         Dinv_sum = np.sum(Dinv)
@@ -75,7 +75,7 @@ def geometric_median(points, eps=1e-6):
             new_median = (1-gamma)*T1 + gamma*median # Eq. (11)
         
         # converge condition    
-        dist_med_newmed,_ = distc.DistCalc_DI(median, new_median)
+        dist_med_newmed,_ = distc.DistCalc_DI(median, new_median, device_name=device_name)
         if dist_med_newmed < eps:
             return new_median
         median = new_median # next median
@@ -100,7 +100,7 @@ def find_centroids_GM(points, cluster_labels, clusters=2):
         print(f'Found for cluster {j}')
     return np.array(centroids)
 
-def find_nearest_neighbour_DI(points, centroids, device_name):
+def find_nearest_neighbour_DI(points, centroids, device_name='/GPU:0'):
     """
     Args:
         points: array of shape (N, X)