def TimeseriesFromPSD_nd(param_noise):
"""
GPU only
"""
(asd_pos, asd_neg, low_f, high_f, high_f_, size, fs, fmin, fmax) = param_noise
(*D_, N) = size
D = reduce(lambda x, y: x * y, D_)
# Gauss noise and its one-sided PSD without window
gauss_noise = 1* nd.random_normal(loc=0,scale=64,shape=(D, N), ctx=ctx)
_, xf_noise, psd_gauss = oneSidedPeriodogram_nd(gauss_noise, fs=8192)
psd_gauss = nd.array(psd_gauss, ctx = ctx, dtype='float64')
psd_twosided = nd.concat( # low positive
nd.zeros((D, low_f), ctx = ctx, dtype='float64'),
# high positive
psd_gauss[:, low_f:high_f] * asd_pos,
nd.zeros((D, high_f_), ctx = ctx, dtype='float64'),
nd.zeros((D, high_f_), ctx = ctx, dtype='float64'),
# high negative
psd_gauss[:, low_f:high_f][::-1] * asd_neg,
# low negative
nd.zeros((D, low_f), ctx = ctx, dtype='float64'), dim=1)
amplitude = nd.sqrt(psd_twosided *2 *fs*N )
epsilon_imag = nd.random_uniform(low=0, high=1, shape=(D,N),ctx=ctx,dtype='float64')*2*np.pi
re = nd.cos(epsilon_imag)*amplitude
im = nd.sin(epsilon_imag)*amplitude
temp = nd.zeros((D, N*2) , ctx=ctx)
temp[:,::2] = re
temp[:,1::2] = im
timeseries = mx.contrib.ndarray.ifft(temp)/N
return timeseries.reshape(size), psd_twosided
def edge_func(self, edges):
real_head, img_head = nd.split(edges.src['emb'], num_outputs=2, axis=-1)
real_tail, img_tail = nd.split(edges.dst['emb'], num_outputs=2, axis=-1)
phase_rel = edges.data['emb'] / (self.emb_init / np.pi)
re_rel, im_rel = nd.cos(phase_rel), nd.sin(phase_rel)
real_score = real_head * re_rel - img_head * im_rel
img_score = real_head * im_rel + img_head * re_rel
real_score = real_score - real_tail
img_score = img_score - img_tail
#sqrt((x*x).sum() + eps)
score = mx.nd.sqrt(real_score * real_score + img_score * img_score + self.eps).sum(-1)
return {'score': self.gamma - score}
def infer(self, head_emb, rel_emb, tail_emb):
re_head, im_head = nd.split(head_emb, num_outputs=2, axis=-1)
re_tail, im_tail = nd.split(tail_emb, num_outputs=2, axis=-1)
phase_rel = rel_emb / (self.emb_init / np.pi)
re_rel, im_rel = nd.cos(phase_rel), nd.sin(phase_rel)
real_score = re_head.expand_dims(axis=1) * re_rel.expand_dims(axis=0) - im_head.expand_dims(axis=1) * im_rel.expand_dims(axis=0)
img_score = re_head.expand_dims(axis=1) * im_rel.expand_dims(axis=0) + im_head.expand_dims(axis=1) * re_rel.expand_dims(axis=0)
real_score = real_score.expand_dims(axis=2) - re_tail.expand_dims(axis=0).expand_dims(axis=0)
img_score = img_score.expand_dims(axis=2) - im_tail.expand_dims(axis=0).expand_dims(axis=0)
score = mx.nd.sqrt(real_score * real_score + img_score * img_score + self.eps).sum(-1)
return self.gamma - score
def fn(heads, relations, tails, num_chunks, chunk_size, neg_sample_size):
hidden_dim = heads.shape[1]
emb_real, emb_img = nd.split(heads, num_outputs=2, axis=-1)
phase_rel = relations / (emb_init / np.pi)
rel_real, rel_img = nd.cos(phase_rel), nd.sin(phase_rel)
real = emb_real * rel_real - emb_img * rel_img
img = emb_real * rel_img + emb_img * rel_real
emb_complex = nd.concat(real, img, dim=-1)
tmp = emb_complex.reshape(num_chunks, chunk_size, 1, hidden_dim)
tails = tails.reshape(num_chunks, 1, neg_sample_size, hidden_dim)
score = tmp - tails
score_real, score_img = nd.split(score, num_outputs=2, axis=-1)
score = mx.nd.sqrt(score_real * score_real + score_img * score_img + self.eps).sum(-1)
return gamma - score
