pytorch-submodular/test.py at master · Pratik08/pytorch-submodular · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
from nesterov import FacilityLocationSelection
from nesterov import MaxMarginRelSelection
from nesterov import DisparityMin
import numpy
import time
import torch
import matplotlib.pyplot as plt

numpy.random.seed(8)
for n in [25]:
    X = numpy.concatenate([numpy.random.normal(7.5, 1, size=(n, 2)),
                           numpy.random.normal(2, 1, size=(n, 2)),
                           numpy.random.normal(15, 1, size=(n, 2))])

    X_1 = numpy.concatenate([numpy.random.normal(7.5, 1, size=(n, 2)),
                             numpy.random.rand(n, 2),
                             numpy.random.normal(15, 1, size=(n, 2))])

    X_2 = numpy.concatenate([numpy.random.normal(7.5, 1, size=(n, 2)),
                             numpy.array([[5, 16]]),
                             numpy.random.normal(2, 1, size=(n, 2)),
                             numpy.array([[16, 2], [12, 4]]),
                             numpy.random.normal(15, 1, size=(n, 2)),
                             numpy.array([[1, 12], [8, 8]])])

    start_time = time.time()
    #func = FacilityLocationSelection(2, 'euclidean', n_greedy_samples=1)
    #func = MaxMarginRelSelection(5, 'euclidean')
    func = DisparityMin(5, 'euclidean', n_greedy_samples=5)
    func.fit(X)
    end_time = time.time()
    print("n = ", str(n), " time: " + str(end_time - start_time))
    print(func.ranking)

    plt_x = []
    plt_y = []
    out_x = []
    out_y = []
    # out_x = [5, 16, 12, 1, 8]
    # out_y = [16, 2, 4, 12, 8]

    for idx in range(X.shape[0]):
        if idx not in func.ranking:
            plt_x.append(X[idx][0])
            plt_y.append(X[idx][1])
        else:
            out_x.append(X[idx][0])
            out_y.append(X[idx][1])

    plt.scatter(plt_x, plt_y, color='blue')
    plt.scatter(out_x, out_y, color='red')
    plt.savefig("plot.png")