def Lanczos(Sigma_func,b):
""" Lanczos method to approximate Sigma^1/2 * b, with b random vec
Note: this only gives you a single draw though, which may not be ideal.
Inputs:
Sigma_func: function to premultiply a vector by Sigma, which you
might not want to explicitly construct if it's huge.
b: random vector of N(0,1)'
Returns:
random vector approximately equal to Sigma^1/2 * b
"""
n = tf.shape(b)[0]
k = tf.div(n,500) + 3 #this many Lanczos iterations
betas = tf.zeros(1)
alphas = tf.zeros(0)
D = tf.zeros((n,1))
b_norm = tf.norm(b)
D = tf.concat([D,tf.reshape(b/b_norm,[-1,1])],1)
def cond(j,alphas,betas,D):
return j < k+1
def body(j,alphas,betas,D):
d_j = tf.slice(D,[0,j],[-1,1])
d = Sigma_func(d_j) - tf.slice(betas,[j-1],[1])*tf.slice(D,[0,j-1],[-1,1])
alphas = tf.concat([alphas,[dot(d_j,d)]],0)
d = d - tf.slice(alphas,[j-1],[1])*d_j
betas = tf.concat([betas,[tf.norm(d)]],0)
D = tf.concat([D,d/tf.slice(betas,[j],[1])],1)
return j+1,alphas,betas,D
j = tf.constant(1)
j,alphas,betas,D = tf.while_loop(cond,body,loop_vars=[j,alphas,betas,D],
shape_invariants=[j.get_shape(),tf.TensorShape([None]),
tf.TensorShape([None]),tf.TensorShape([None,None])])
betas_ = tf.diag(tf.slice(betas,[1],[k-1]))
D_ = tf.slice(D,[0,1],[-1,k])
#build out tridiagonal H: alphas_1:k on main, betas_2:k on off
H = tf.diag(alphas) + tf.pad(betas_,[[1,0],[0,1]]) + tf.pad(betas_,[[0,1],[1,0]])
e,v = tf.self_adjoint_eig(H)
e_pos = tf.maximum(0.0,e)+1e-6 #make sure positive definite
e_sqrt = tf.diag(tf.sqrt(e_pos))
sq_H = tf.matmul(v,tf.matmul(e_sqrt,tf.transpose(v)))
out = b_norm*tf.matmul(D_,sq_H)
return tf.slice(out,[0,0],[-1,1]) #grab last column = *e_1
评论列表
文章目录
