In [1]:

    

import pystan
import numpy as np
from scipy.io.matlab import loadmat
from scipy.sparse import csr_matrix
import matplotlib
import matplotlib.pyplot as plt
%matplotlib inline


    

In [2]:

    

mat = loadmat('../data/alyawarradata.mat')
T = np.array(mat['Rs'], np.float32)

D = 5
T = np.swapaxes(T, 1, 2)
T = np.swapaxes(T, 0, 1)  # [relation, entity, entity]
num_relation, num_entity, _ = T.shape
print(T.shape)


    

    




(26, 104, 104)

In [ ]:

    

code = """
data {
    int<lower=0> n_dim;
    int<lower=0> n_entity;
    int<lower=0> n_relation;
    real<lower=0> var_x;
    real<lower=0> var_e;
    real<lower=0> var_r;    
    real x[n_relation, n_entity, n_entity];
}
parameters {
    matrix[n_entity, n_dim] E;
    real R[n_relation, n_dim, n_dim];
}
transformed parameters {
    real mu[n_relation, n_entity, n_entity];
    for (k in 1:n_relation)
        mu[k] <- to_array_2d((to_matrix(E) * to_matrix(R[k])) * to_matrix(E)');
}
model {
    vector[n_dim] mean_e;
    vector[n_dim] diag_e;
    
    mean_e <- rep_vector(0, n_dim);
    diag_e <- rep_vector(var_e, n_dim);
    
    for (i in 1:n_entity)
        E[i] ~ multi_normal(mean_e, diag_matrix(diag_e));
    for (k in 1:n_relation)
        for (i in 1:n_dim)
            for (j in 1:n_dim)
                R[k,i,j] ~ normal(0, var_r);
    for (k in 1:n_relation)
        for (i in 1:n_entity)
            for (j in 1:n_entity)
                x[k, i, j] ~ normal(mu[k, i, j], var_x);
}
"""
dat = {'var_x':1.,
       'var_r':1.,
       'var_e':1.,
       'n_dim':5,
       'n_entity':num_entity,
       'n_relation':num_relation,
       'x':T
      }

fit = pystan.stan(model_code=code,data=dat,iter=1,chains=1)


    

In [ ]:

    

samples = fit.extract()


    

In [ ]:

    

E = samples['E']
R = samples['R']

X = np.zeros_like(trainX)

for k in range(num_relation)
    X = np.dot(np.dot(E,R[k]), E.T)
    
A = np.sum((trainX-_X)**2)
B = np.sum(trainX**2)
fit = 1.-A/B

print(fit)