cPDS/train_cPDS.py at main · JSA-Insubria/cPDS · 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
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
import datetime
import numpy as np

import load_save_data
import cPDS
import key_gen_util


gp_param = 0
m = 0


def save_time(file, time_pre):
    time_post = datetime.datetime.now()
    load_save_data.writeIntoCSV(m, 'enc_' + str(gp_param), file, str((time_post - time_pre).total_seconds()))


def save_time_enc(file, time):
    load_save_data.writeIntoCSV(m, 'enc_' + str(gp_param), file, str(time))


def compute_Lx(L, x):
    res = np.zeros(shape=x.shape)
    for i in range(len(L)):
        if L[i] != 0:
            res[i] = L[i] * x[i]

    return res


def aggregator_sum(node, L, lambdaa_k, x, alone_list):
    time_pre = datetime.datetime.now()
    tmp_sum = np.zeros(shape=lambdaa_k.shape, dtype=object)
    for j in range(len(L)):
        if L[j] != 0:
            tmp_sum += x[j]

    if alone_list != -1:
        tmp_sum += x[alone_list]

    lambdaa_kplus1 = lambdaa_k + tmp_sum
    save_time('agent_sum_' + str(node), time_pre)

    return lambdaa_kplus1


def agent_encrypt(keys_dict, L, x, node, enc_time_nodes, alone_list):
    x_enc_node = np.zeros(shape=x.shape, dtype=object)
    key = 0
    for other_node in range(len(L)):
        if L[other_node] != 0:
            if other_node != node:
                time_pre = datetime.datetime.now()
                x_enc_node[other_node] = keys_dict['pk_list' + str(node)][key].encryptMatrix(x[other_node])
                enc_time_nodes[other_node] += (datetime.datetime.now() - time_pre).total_seconds()
                key += 1
            else:
                time_pre = datetime.datetime.now()
                x_enc_node[other_node] = keys_dict['pk_list' + str(node)][-1].encryptMatrix(x[other_node])
                enc_time_nodes[other_node] += (datetime.datetime.now() - time_pre).total_seconds()

    if alone_list != -1:
        x_enc_node[alone_list] = keys_dict['pk_list' + str(node)][-2].encryptMatrix(x[alone_list])

    return x_enc_node, enc_time_nodes


def main_decrypt(keys_dict, lambdaa_encrypted):
    lambdaa = np.zeros(lambdaa_encrypted.shape)
    for node in range(m):
        time_pre = datetime.datetime.now()
        lambdaa[node] = keys_dict['msk' + str(node)].decryptMatrix(lambdaa_encrypted[node])
        save_time('agent_dec_' + str(node), time_pre)

    return lambdaa


def train_cPDS(n_agent, graph_param, max_iters, L, tau, rho, n, gammas, data, labels, x_init, q, theta):

    global m, gp_param
    m = n_agent
    gp_param = graph_param

    alone_list = key_gen_util.get_alone_listL(L)

    keys_dict = key_gen_util.gen_keys(L)
    x = np.zeros(shape=x_init.shape, dtype=object)

    cPDSs = []
    for j in range(m):
        cPDSs.append(cPDS.cPDS(j, tau, rho, theta[j], gammas[j], data[j], labels[j], q[j], n[j], x_init[j], L[j]))

    lambdaa = L @ x_init

    for i in range(max_iters):
        iteration_time_pre = datetime.datetime.now()

        # compute x for all nodes
        x = np.asarray([cPDSs[node].compute(lambdaa[node]) for node in range(m)])

        # encrypt for node
        lambdaa_kplus1 = np.zeros(shape=lambdaa.shape, dtype=object)
        enc_time_nodes = np.zeros(shape=m)
        for node in range(m):
            res = compute_Lx(L[node], x)
            x_enc, enc_time_nodes = agent_encrypt(keys_dict, L[node], res, node, enc_time_nodes, alone_list[node])
            lambdaa_kplus1[node] = aggregator_sum(node, L[node], lambdaa[node], x_enc, alone_list[node])

        # save agent time
        [save_time_enc('agent_enc_' + str(node), enc_time_nodes[node]) for node in range(m)]

        lambdaa = main_decrypt(keys_dict, lambdaa_kplus1)
        save_time('iteration_time', iteration_time_pre)

    x_return = np.mean(x, axis=0)
    w_cPDS = x_return[:-1]
    b_cPDS = x_return[-1]
    return w_cPDS, b_cPDS