df_host_b.py

class B:
    def __init__(self, x, pub_key=None):
        self.x = x # Host B's X.
        self.features = x.shape[1]
        self.pub_key = pub_key

    # Called by Host (A) with current model Theta and A's
    # (encypted) part of the gradient calculation.
    def gradients(self, theta, u):

        # B's Theta
        b_theta = theta[-self.features:]

        # B's part of the gradient (result is a 1d vector of length = A's y)
        v = 1/4 * np.dot(self.x, b_theta)

        # Add A's u to B's v.
        # Since u is encrypted, the resulting w is also encrypted.
        w = u + v

        # Calculate B's gradient (result is encrypted since dot-product with w)
        gradient_b = np.dot(w, self.x)

        # Send back (encypted) w and B gradient to A.
        return w, gradient_b