In [74]:
import numpy as np
import tensorflow as tf
import logging
logging.basicConfig(format='%(asctime)s %(message)s')
import pylab as plt
import cmocean
from scipy.spatial import cKDTree
from ionotomo.tomography.pipeline import Pipeline
from ionotomo.settings import TFSettings
from timeit import default_timer
from ionotomo import *
import astropy.coordinates as ac
import astropy.units as au
import gpflow as gp
import sys
import h5py
import threading
from timeit import default_timer
#%matplotlib notebook
from concurrent import futures
from functools import partial
from threading import Lock
import astropy.units as au
import astropy.time as at
from collections import deque
from doubly_stochastic_dgp.dgp import DGP
from ionotomo.bayes.gpflow_contrib import GPR_v2,Gaussian_v2
from scipy.cluster.vq import kmeans2
from scipy.spatial.distance import pdist,squareform
import os
In [75]:
class NNComposedKernel(gp.kernels.Kernel):
"""
This kernel class allows for easily adding a NN (or other function) to a GP model.
The kernel does not actually do anything with the NN.
"""
def __init__(self, kern, f, f_scope):
"""
kern.input_dim needs to be consistent with the output dimension of f
"""
super().__init__(kern.input_dim,active_dims=kern.active_dims)
self.kern = kern
self._f = lambda x: tf.cast(f(x), gp.settings.float_type) #function to call on input
self._f_scope = f_scope #learnable variables that f depends on
def f(self, X):
if X is not None:
return self._f(X)
def _get_f_vars(self):
return tf.trainable_variables(scope=self._f_scope)
@gp.autoflow([gp.settings.float_type, [None,None]])
def compute_f(self, X):
return self.f(X)
def K(self, X, X2=None, presliced=False):
if not presliced:
X, X2 = self._slice(X, X2)
return self.kern.K(self.f(X), self.f(X2), presliced=True)
def Kdiag(self, X, presliced=False):
if not presliced:
X,_ = self._slice(X, None)
return self.kern.Kdiag(self.f(X))
# we need to add these extra functions to the model so the tensorflow variables get picked up
class NN_SVGP(gp.models.svgp.SVGP):
def get_NNKernels(self,kern=None):
if kern is None:
kern = self.kern
out = []
for c in kern.children.values():
if isinstance(c,gp.kernels.Kernel):
if isinstance(c,NNComposedKernel):
out.append(c)
else:
out = out + self.get_NNKernels(c)
return out
@property
def all_f_vars(self):
NN_kerns = self.get_NNKernels()
f_vars = []
for k in NN_kerns:
f_vars = f_vars + k._get_f_vars()
return f_vars
@property
def trainable_tensors(self):
f_vars = self.all_f_vars
try:
return super().trainable_tensors + f_vars
except:
return super().trainable_tensors
@property
def initializables(self):
f_vars = self.all_f_vars
try:
return super().initializables + f_vars
except:
return super().initializables
def scaled_square_dist(X,X2,lengthscales):
'''
r_ij = sum_k (x_ik - x_jk)*(x_ik - x_jk)
= sum_k x_ik*x_ik - x_ik*x_jk - x_jk*x_ik + x_jk*x_jk
= Xs - 2 X.X^t + Xs.T
r_ij = sum_k (x_ik - y_jk)*(x_ik - y_jk)
= sum_k x_ik*x_ik - x_ik*y_jk - y_jk*x_ik + y_jk*y_jk
= Xs - 2 X.X^t + Ys.T
'''
X = X / lengthscales
Xs = tf.reduce_sum(tf.square(X), axis=1)
if X2 is None:
dist = -2 * tf.matmul(X, X, transpose_b=True)
dist += tf.reshape(Xs, (-1, 1)) + tf.reshape(Xs, (1, -1))
return dist
X2 = X2 / lengthscales
X2s = tf.reduce_sum(tf.square(X2), axis=1)
dist = -2 * tf.matmul(X, X2, transpose_b=True)
dist += tf.reshape(Xs, (-1, 1)) + tf.reshape(X2s, (1, -1))
return dist
def anisotropic_modulation(ndim, M, scope):
def modulation(X,ndim=ndim,M=M,scope=scope):
with tf.variable_scope(scope,reuse=tf.AUTO_REUSE) as scope:
factor = tf.get_variable("factor",shape=(M,1),dtype=gp.settings.float_type,
initializer=\
tf.zeros_initializer(dtype=gp.settings.float_type))
factor = 1.5*tf.nn.sigmoid(factor) + 0.25 # between 0.25 and 1.75 modulations starting at 1.
points = tf.get_variable("points",shape=(M,ndim),dtype=gp.settings.float_type,
initializer=\
tf.random_uniform_initializer(minval=-2,maxval=2,dtype=gp.settings.float_type))
scale = tf.nn.softplus(tf.get_variable("scale",shape=(),dtype=gp.settings.float_type,
initializer=tf.ones_initializer(dtype=gp.settings.float_type))) + 1e-6
dist = scaled_square_dist(X,points,scale)
weights = tf.exp(-dist/2.) #N, M
weights /= tf.reduce_sum(weights,axis=1,keepdims=True,name='weights')# N,1
factor = tf.matmul(weights, factor,name='factor')#N, 1
res = X/factor
return res
return modulation
def _synced_minibatch(*X,minibatch_size=100,seed=0, sess=None, shuffle = True):
init_placeholders = tuple([tf.placeholder(gp.settings.tf_float,shape=x.shape) for x in X])
data = tf.data.Dataset.from_tensor_slices(init_placeholders)
data = data.repeat()
if shuffle:
data = data.shuffle(buffer_size=X[0].shape[0], seed=seed)
data = data.batch(batch_size=tf.constant(minibatch_size,dtype=tf.int64))
iterator_tensor = data.make_initializable_iterator()
if sess is not None:
sess.run(iterator_tensor.initializer, feed_dict={p:x for p,x in zip(init_placeholders,X)})
return init_placeholders, iterator_tensor.initializer, iterator_tensor.get_next()
class WeightedSVGP(NN_SVGP):
def __init__(self, obs_weight, X, Y, kern, likelihood, feat=None,
mean_function=None,
num_latent=None,
q_diag=False,
whiten=True,
minibatch_size=None,
Z=None,
num_data=None,
**kwargs):
super(WeightedSVGP,self).__init__(X, Y, kern, likelihood, feat=feat,
mean_function=mean_function,
num_latent=num_latent,
q_diag=q_diag,
whiten=whiten,
minibatch_size=None,
Z=Z,
num_data=num_data,
**kwargs)
self.obs_weight = gp.DataHolder(obs_weight) \
if minibatch_size is None else gp.Minibatch(obs_weight,batch_size=minibatch_size, seed=0)
@gp.params_as_tensors
def _build_likelihood(self):
"""
This gives a variational bound on the model likelihood.
"""
# Get prior KL.
KL = self.build_prior_KL()
# Get conditionals
fmean, fvar = self._build_predict(self.X, full_cov=False)
# Get variational expectations.
var_exp = self.likelihood.variational_expectations(fmean, fvar, self.Y) * self.obs_weight
# re-scale for minibatch size
scale = tf.cast(self.num_data, gp.settings.float_type) / tf.cast(tf.shape(self.X)[0], gp.settings.float_type)
scale = scale / tf.reduce_mean(self.obs_weight)
return tf.reduce_sum(var_exp) * scale - KL
def get_only_vars_in_model(variables, model):
reader = tf.train.NewCheckpointReader(model)
var_to_shape_map = reader.get_variable_to_shape_map()
vars_in_model = [k for k in sorted(var_to_shape_map)]
out_vars = []
for var in variables:
v = var.name.split(":")[0]
if v in vars_in_model:
if tuple(var.shape.as_list()) != reader.get_tensor(v).shape:
logging.warning("{} has shape mis-match: {} {}".format(v,
tuple(var.shape.as_list()), reader.get_tensor(v).shape))
continue
out_vars.append(var)
return out_vars
def rename(model,prefix='WeightedSVGP',index=""):
with tf.Session(graph=tf.Graph()) as sess:
reader = tf.train.NewCheckpointReader(model)
var_to_shape_map = reader.get_variable_to_shape_map()
vars_in_model = [k for k in sorted(var_to_shape_map)]
for v in vars_in_model:
t = reader.get_tensor(v)
if 'WeightedSVGP' in v:
new_name = "/".join(['WeightedSVGP-{}'.format(index)] + v.split('/')[1:])
var = tf.Variable(t,name=new_name)
else:
var = tf.Variable(t,name=v)
saver = tf.train.Saver()
sess.run(tf.global_variables_initializer())
saver.save(sess, model)
In [73]:
class WeightedDGP(DGP):
def __init__(self,obs_weight, X, Y, Z, kernels, likelihood,
num_outputs=None,num_data=None,
mean_function=gp.mean_functions.Zero(), # the final layer mean function
**kwargs):
pass
class Smoothing(object):
"""
Class for all types of GP smoothing/conditioned prediction
"""
def __init__(self,datapack, proj_dir):
if isinstance(datapack, str):
datapack = DataPack(filename=datapack)
self.datapack = datapack
self.proj_dir = os.path.abspath(proj_dir)
try:
os.makedirs(self.proj_dir)
except:
pass
def _make_coord_array(t,d):
"""Static method to pack coordinates
"""
Nt,Nd = t.shape[0],d.shape[0]
X = np.zeros([Nt,Nd,3],dtype=np.float64)
for j in range(Nt):
for k in range(Nd):
X[j,k,0] = t[j]
X[j,k,1:3] = d[k,:]
X = np.reshape(X,(Nt*Nd,3))
return X
def _make_coord_array_full(a,t,d,f):
"""Static method to pack coordinates
"""
Na,Nt,Nd,Nf = a.shape[0],t.shape[0],d.shape[0],f.shape[0]
X = np.zeros([Na,Nt,Nd,Nf,6],dtype=np.float64)
for i in range(Na):
for j in range(Nt):
for k in range(Nd):
for l in range(Nf):
X[i,j,k,l,0:2] = a[i,:]
X[i,j,k,l,2] = t[j]
X[i,j,k,l,3:5] = d[k,:]
X[i,j,k,l,5] = f[l]
X = np.reshape(X,(Na*Nt*Nd*Nf,6))
return X
def _build_sgp_model(self, sess, weights, X, Y, ls_scale, y_scale,
minibatch_size=500, M=1000,Z=None, feature_trainable=False,ls_init=(200,1.),
ls_trainable=(True,True), likelihood_var_trainable=True, verbose=False):
"""
Build svgp model
"""
N, num_latent = Y.shape
Z = kmeans2(X, M, minit='points')[0] if Z is None else Z
num_data = X.shape[0]
_,_, data = _synced_minibatch(weights, X, Y,minibatch_size=minibatch_size, sess = sess,shuffle=True)
weights,X,Y = data
with gp.defer_build():
k_time = gp.kernels.RBF(1,active_dims = [0],
lengthscales=[0.5 if ls_init[0] is None else ls_init[0]/ls_scale[0]])
k_space = gp.kernels.RBF(2,active_dims = [1,2],
lengthscales=[1.0 if ls_init[1] is None else ls_init[1]/ls_scale[1]])
for k,f in zip([k_time,k_space],ls_trainable):
k.lengthscales.set_trainable(f)
if not f:
logging.warning("Setting {} non-trainable".format(k))
k_time.lengthscales.prior = gp.priors.Gaussian(0.,200./ls_scale[0])
k_space.lengthscales.prior = gp.priors.Gaussian(1./ls_scale[1],1./ls_scale[1])
kern = k_time*k_space
mean = gp.mean_functions.Zero()
m = WeightedSVGP(weights, X, Y, kern, mean_function = mean,
likelihood=gp.likelihoods.Gaussian(),
Z=Z, num_latent=num_latent,num_data=num_data,
minibatch_size=None, whiten=True)
m.likelihood.variance.set_trainable(likelihood_var_trainable)
m.q_sqrt = m.q_sqrt.value * 0.4
m.feature.set_trainable(feature_trainable)
m.compile()
if verbose:
logging.warning(m)
return m
def _build_sgp_model_full(self, sess, weights, X, Y, ls_scale, y_scale,
minibatch_size=500, M=1000,Z=None, feature_trainable=False,ls_init=(5.,200,1.),
ls_trainable=(True, True,True), likelihood_var_trainable=True, verbose=False):
"""
Build svgp model
"""
N, num_latent = Y.shape
Z = kmeans2(X, M, minit='points')[0] if Z is None else Z
num_data = X.shape[0]
_,_, data = _synced_minibatch(weights, X, Y,minibatch_size=minibatch_size, sess = sess,shuffle=True)
weights,X,Y = data
with gp.defer_build():
k_space = gp.kernels.RBF(2,active_dims = [0,1],
lengthscales=[1.0 if ls_init[0] is None else ls_init[0]/ls_scale[0]])
k_time = gp.kernels.RBF(1,active_dims = [2],
lengthscales=[0.5 if ls_init[1] is None else ls_init[1]/ls_scale[1]])
k_dir = gp.kernels.RBF(2,active_dims = [3,4],
lengthscales=[1.0 if ls_init[2] is None else ls_init[2]/ls_scale[2]])
k_freq = gp.kernels.Polynomial(1,active_dims = [5], degree=1.,variance=1.)
for k,f in zip([k_space,k_time,k_dir],ls_trainable):
k.lengthscales.set_trainable(f)
if not f:
logging.warning("Setting {} non-trainable".format(k))
k_space.lengthscales.prior = gp.priors.Gaussian(0.,10./ls_scale[0])
k_time.lengthscales.prior = gp.priors.Gaussian(0.,200./ls_scale[1])
k_dir.lengthscales.prior = gp.priors.Gaussian(1./ls_scale[2],1./ls_scale[2])
# allow length scale to change depending on loc |X - X'|^2/ls^2 -> |X/f(X) - X'/f(X')|^2/ls^2
k_space = NNComposedKernel(k_space, anisotropic_modulation(ndim=2, M=7, scope='f_space_ls'), 'f_space_ls')
k_time = NNComposedKernel(k_time, anisotropic_modulation(ndim=1, M=4, scope='f_time_ls'), 'f_time_ls')
k_variance = NNComposedKernel(
gp.kernels.Polynomial(2,active_dims = [0,1], degree=1.,variance=1.),
anisotropic_modulation(ndim=2, M=7, scope='f_space_var'), 'f_space_var')
kern = k_variance*k_space*k_time*k_dir*k_freq
mean = gp.mean_functions.Zero()
m = WeightedSVGP(weights, X, Y, kern, mean_function = mean,
likelihood=gp.likelihoods.Gaussian(),
Z=Z, num_latent=num_latent,num_data=num_data,
minibatch_size=None, whiten=True)
m.likelihood.variance.set_trainable(likelihood_var_trainable)
m.q_sqrt = m.q_sqrt.value * 0.4
m.feature.set_trainable(feature_trainable)
m.compile()
if verbose:
logging.warning(m)
return m
def _build_dgp_model(self, depth, sess, weight, X, Y, ls_scale, y_scale,
minibatch_size=500, Z=None,M=100,feature_trainable=False,ls_init=(None,None,None),
ls_trainable=(True,True,True),likelihood_var_trainable=True, verbose=False):
"""
Build svgp model
"""
N, num_latent = Y.shape
Z = kmeans2(X, M, minit='points')[0]
with gp.defer_build():
k_time = gp.kernels.RBF(1,active_dims = [0],
lengthscales=[0.3 if ls_init[0] is None else ls_init[0]/ls_scale[0]])
k_space = gp.kernels.RBF(2,active_dims = [1,2],
lengthscales=[0.3 if ls_init[1] is None else ls_init[1]/ls_scale[1]])
k_freq = gp.kernels.RBF(1,active_dims = [3],
lengthscales=[10. if ls_init[2] is None else ls_init[2]/ls_scale[2]])
for k,f in zip([k_time,k_space,k_freq],ls_trainable):
k.lengthscales.set_trainable(f)
if not f:
logging.warning("Setting {} non-trainable".format(k))
k_time.lengthscales.prior = gp.priors.Gaussian(0,1./3.)
k_space.lengthscales.prior = gp.priors.Gaussian(1./ls_scale[1],0.5/ls_scale[1])
kern = k_time*k_space*k_freq
mean = gp.mean_functions.Zero()
kernels = [kern]
for l in range(1,depth):
kernels.append(RBF(4-l, lengthscales=2., variance=2.,ARD=True))
#kernels[-1].lengthscales.prior = gp.priors.Gaussian(0,1./3.)
m = DGP(X, Y, Z, kernels, gp.likelihoods.Gaussian(),
minibatch_size=minibatch_size,
num_outputs=num_latent,num_samples=1)
# start things deterministic
for layer in m.layers[:-1]:
layer.q_sqrt = layer.q_sqrt.value * 1e-5
for layer in m.layers:
layer.feature.Z.set_trainable(feature_trainable)
m.compile()
if verbose:
logging.warning(m)
return m
def _build_model(self,m_type, sess, weight, X, Y, ls_scale, y_scale, **kwargs):
"""
Build a GP model depending on m_type
m_type: str, one of 'sgp', 'dgp2', 'dgp3'
**kwargs are passes to the constructor of the model type.
"""
if m_type == 'sgp':
return self._build_sgp_model(sess, weight, X, Y, ls_scale, y_scale,**kwargs)
elif m_type == 'sgp_full':
return self._build_sgp_model_full(sess, weight, X, Y, ls_scale, y_scale,**kwargs)
# elif m_type == 'dgp2':
# return self._build_dgp_model(2,sess,weight, X, Y, ls_scale, y_scale,**kwargs)
# elif m_type == 'dgp3':
# return self._build_dgp_model(3,sess, weight, X, Y, ls_scale, y_scale,**kwargs)
raise ValueError("{} is invalid model type".format(m_type))
def _solve_interval(self,phase, error, coords, lock, error_sigma_clip=None, m_type='sgp',
iterations=1000, pargs=None,verbose=False,model_kwargs={}):
"""
Solve the block of data independently over antennas assuming homogeneity.
phase: array of shape (Na, Nt, Nd, Nf)
errors: array of shape (Na, Nt, Nd, Nf), -1 in an element means to mask
coords: tuple of arrays of shape (Nt,) (Nd,2) (Nf,)
lock: a mutable lock or None
m_type: str the model type to use, see build_model
pargs: str or None, thing to print on start of block
"""
try:
if pargs is not None:
logging.warning("{}".format(pargs))
Na,Nt,Nd,Nf = phase.shape
y = phase.transpose((1,2,0,3)).reshape((Nt*Nd,Na*Nf))#Nt*Nd,Na*Nf
sigma_y = error.transpose((1,2,0,3)).reshape((Nt*Nd,Na*Nf))#Nt*Nd,Na*Nf
mask = sigma_y < 0.#Nt*Nd,Na*Nf
y_mean = (y*np.bitwise_not(mask)).sum(axis=0) / (np.bitwise_not(mask).sum(axis=0))#Na*Nf
y_scale = np.sqrt((y**2*np.bitwise_not(mask)).sum(axis=0) \
/ (np.bitwise_not(mask).sum(axis=0)) - y_mean**2) + 1e-6#Na*Nf
y = (y - y_mean)/y_scale#Nt*Nd,Na*Nf
var_y = (sigma_y/y_scale)**2#Nt*Nd,Na*Nf
if error_sigma_clip is not None:
log_var_y = np.log(var_y)#Nt*Nd,Na*Nf
log_var_y[mask] = np.nan
E_log_var_y = np.nanmean(log_var_y,axis=0)#Na*Nf
std_log_var_y = np.nanstd(log_var_y,axis=0)#Na*Nf
clip_mask = (log_var_y - E_log_var_y) > error_sigma_clip*std_log_var_y#Nt*Nd,Na*Nf
ignore_mask = np.bitwise_or(mask,clip_mask)#Nt*Nd,Na*Nf
else:
ignore_mask = mask
keep_mask = np.bitwise_not(ignore_mask)#Nt*Nd,Na*Nf
weight = 1./(var_y+1e-6)#Nt*Nd,Na*Nf
weight_norm = np.stack([np.percentile(weight[keep_mask[:,i],i],50) for i in range(Na*Nf)],axis=-1)
weight /= weight_norm + 1e-6
# plt.hist(weight.flatten(),bins=20)
# plt.show()
weight = np.ones(y.shape)
weight[ignore_mask] = 0.
t,d = coords
t_scale = t.std() + 1e-6
d_scale = np.sqrt((d.std(axis=0)**2).mean()) + 1e-6
ls_scale = (t_scale,d_scale)
t = (t - t.mean()) / t_scale
d = (d - d.mean(axis=0)) / d_scale
X = Smoothing._make_coord_array(t,d)#Nt*Nd,3
###
# set Z explicitly to spacing of ::3 in time
model_kwargs['Z'] = Smoothing._make_coord_array(t[::3],d)
with tf.Session(graph=tf.Graph()) as sess:
lock.acquire() if lock is not None else None
try:
model = self._build_model(m_type, sess, weight, X, y, ls_scale, y_scale, **model_kwargs)
finally:
lock.release() if lock is not None else None
logging.warning("Initial log-likelihood {}".format(model.compute_log_likelihood()))
opt = gp.train.AdamOptimizer(1e-2)
opt.minimize(model, maxiter=iterations)
logging.warning("Final log-likelihood {}".format(model.compute_log_likelihood()))
# smooth
kern_lengthscales = (
model.kern.rbf_1.lengthscales.value[0]*ls_scale[0],
model.kern.rbf_2.lengthscales.value[0]*ls_scale[1]
)
kern_variance = model.kern.rbf_1.variance.value*model.kern.rbf_2.variance.value*y_scale**2
if verbose:
logging.warning(model)
logging.warning(kern_lengthscales)
logging.warning(kern_variance)
predict_minibatch = 1000
for start in range(0,X.shape[0],predict_minibatch):
stop = min(start+predict_minibatch,X.shape[0])
Xs = X[start:stop,:]
ystar,varstar = model.predict_y(Xs)#batch,Na
ystar = ystar * y_scale + y_mean
varstar = varstar * y_scale**2
y[start:stop,:] = ystar
var_y[start:stop,:] = varstar
phase = y.reshape([Nt,Nd,Na,Nf]).transpose((2,0,1,3))
variance = var_y.reshape([Nt,Nd,Na,Nf]).transpose((2,0,1,3))
return phase,variance,kern_lengthscales, kern_variance
except Exception as e:
logging.warning(e)
def _solve_interval_full(self,model_file, phase, error, coords, lock, load_model = None, error_sigma_clip=None, m_type='sgp_full',
iterations=1000, pargs=None,verbose=False,model_kwargs={}):
"""
Solve the block of data independently over antennas assuming homogeneity.
phase: array of shape (Na, Nt, Nd, Nf)
errors: array of shape (Na, Nt, Nd, Nf), -1 in an element means to mask
coords: tuple of arrays of shape (Nt,) (Nd,2) (Nf,)
lock: a mutable lock or None
m_type: str the model type to use, see build_model
pargs: str or None, thing to print on start of block
"""
assert "_full" in m_type
# try:
if pargs is not None:
logging.warning("{}".format(pargs))
Na,Nt,Nd,Nf = phase.shape
freqs = coords[3]
y = (phase*(freqs/-8.4480e9)).reshape((Na*Nt*Nd*Nf,1)) #Na*Nt*Nd*Nf, 1
sigma_y = (error*(freqs/8.4480e9)).reshape((Na*Nt*Nd*Nf,1)) #Na*Nt*Nd*Nf, 1
mask = sigma_y < 0. #Na*Nt*Nd*Nf, 1
y_mean = np.average(y,weights = np.bitwise_not(mask))# scalar
y_scale = np.sqrt(np.average(y**2,weights = np.bitwise_not(mask))\
- y_mean**2) + 1e-6 #scalar
y = (y - y_mean)/y_scale#Na*Nt*Nd*Nf, 1
var_y = (sigma_y/y_scale)**2#Na*Nt*Nd*Nf, 1
if error_sigma_clip is not None:
log_var_y = np.log(var_y)#Na*Nt*Nd*Nf, 1
log_var_y[mask] = np.nan
E_log_var_y = np.nanmean(log_var_y,axis=0)#1
std_log_var_y = np.nanstd(log_var_y,axis=0)#1
clip_mask = (log_var_y - E_log_var_y) > error_sigma_clip*std_log_var_y##Na*Nt*Nd*Nf, 1
ignore_mask = np.bitwise_or(mask,clip_mask)#Na*Nt*Nd*Nf, 1
else:
ignore_mask = mask
keep_mask = np.bitwise_not(ignore_mask)#Na*Nt*Nd*Nf, 1
# weight = 1./(var_y+1e-6)#Na*Nt*Nd*Nf, 1
# weight_norm = np.stack([np.percentile(weight[keep_mask[:,i],i],50) for i in range(Na*Nf)],axis=-1)
# weight /= weight_norm + 1e-6
# # plt.hist(weight.flatten(),bins=20)
# # plt.show()
weight = np.ones(y.shape)
weight[ignore_mask] = 0.
x,t,d,f = coords
x_scale = np.sqrt((x.std(axis=0)**2).mean()) + 1e-6
t_scale = t.std() + 1e-6
d_scale = np.sqrt((d.std(axis=0)**2).mean()) + 1e-6
f_scale = f.std() + 1e-6
ls_scale = (x_scale,t_scale,d_scale,f_scale)
x = (x - x.mean(axis=0)) / x_scale
t = (t - t.mean()) / t_scale
d = (d - d.mean(axis=0)) / d_scale
f = (f - f.mean()) / f_scale
X = Smoothing._make_coord_array_full(x,t,d,f)#Na*Nt*Nd*Nf,6
###
# set Z explicitly to spacing of ::3 in time
model_kwargs['Z'] = None#Smoothing._make_coord_array(t[::3],d)
with tf.Session(graph=tf.Graph()) as sess:
lock.acquire() if lock is not None else None
try:
model = self._build_model(m_type, sess, weight, X, y, ls_scale, y_scale, **model_kwargs)
finally:
lock.release() if lock is not None else None
if load_model is not None:
try:
all_vars = model.trainable_tensors
rename(load_model,prefix='WeightedSVGP',index=model.index)
all_vars = get_only_vars_in_model(all_vars,load_model)
saver = tf.train.Saver(all_vars)
saver.restore(sess, load_model)
model.compile()
logging.warning("Loaded model {}".format(load_model))
logging.warning(model)
except Exception as e:
logging.warning(e)
logging.warning("Unable to load {}".format(load_model))
logging.warning("Initial log-likelihood {}".format(model.compute_log_likelihood()))
opt = gp.train.AdamOptimizer(1e-3)
opt.minimize(model, maxiter=iterations)
logging.warning("Final log-likelihood {}".format(model.compute_log_likelihood()))
f_vars = model.all_f_vars
for var,val in zip(f_vars,sess.run(f_vars)):
logging.warning("{} {}".format(var.name,val))
all_vars = model.trainable_tensors
saver = tf.train.Saver(all_vars)
save_path = saver.save(sess, model_file)
logging.warning("Saved model to {}".format(save_path))
if verbose:
logging.warning(model)
predict_minibatch = 1000
for start in range(0,X.shape[0],predict_minibatch):
stop = min(start+predict_minibatch,X.shape[0])
Xs = X[start:stop,:]
ystar,varstar = model.predict_y(Xs)#minibatch,1
ystar = ystar * y_scale + y_mean
varstar = varstar * y_scale**2#minibaatch,1
y[start:stop,:] = ystar
var_y[start:stop,:] = varstar
phase = y.reshape([Na,Nt,Nd,Nf])*(-8.4480e9/freqs)**2
variance = var_y.reshape([Na,Nt,Nd,Nf])*(-8.4480e9/freqs)**2
return phase,variance
# except Exception as e:
# logging.warning(e)
# logging.warning("Failed interval solve {}".format(model_file))
def solve_and_apply_ensemble(self, save_datapack, ant_idx, time_idx, dir_idx, freq_idx, iterations,
interval, shift, init_solutions = None, num_threads=1,verbose=False,
model_kwargs = {}):
"""
Solve the problem using model_kwargs and then take an ensemble average over interval
and shift.
"""
if init_solutions is not None:
data = np.load(init_solutions)
kern_ls = data['kern_ls']
kern_var = data['kern_var']
kern_times = data['time']
kern_antenna_labels = data['antenna']
datapack = self.datapack
directions, patch_names = datapack.get_directions(dir_idx)
times,timestamps = datapack.get_times(time_idx)
antennas,antenna_labels = datapack.get_antennas(ant_idx)
freqs = datapack.get_freqs(freq_idx)
if ant_idx is -1:
ant_idx = range(len(antennas))
if time_idx is -1:
time_idx = range(len(times))
if freq_idx is -1:
freq_idx = range(len(freqs))
if dir_idx is -1:
dir_idx = range(len(directions))
phase = datapack.get_phase(ant_idx,time_idx,dir_idx,freq_idx)
Na,Nt,Nd,Nf = phase.shape
logging.warning("Working on shapes {}".format(phase.shape))
if interval is None:
interval = Nt
assert interval <= Nt
variance = datapack.get_variance(ant_idx,time_idx,dir_idx,freq_idx)
error = np.sqrt(variance)
data_mask = variance < 0
error[data_mask] = -1
logging.warning("Total masked phases: {}".format(np.sum(data_mask)))
enu = ENU(obstime=times[0],location = self.datapack.radio_array.get_center())
ant_enu = antennas.transform_to(enu)
x = np.array([ant_enu.east.to(au.km).value, ant_enu.north.to(au.km).value]).T
t = times.mjd*86400.#mjs
d = np.array([directions.ra.deg, directions.dec.deg]).T
f = freqs
lock = Lock()
with futures.ThreadPoolExecutor(max_workers=num_threads) as executor:
jobs = []
for j,aj in enumerate(time_idx[::shift]):
start = j*shift
stop = min(start+interval,Nt)
time_slice = slice(start,stop,1)
# if init_solutions is not None:
# ###
# # interpolate kern_params with this interval/shift
# mean_time = np.mean(times.gps[time_slice])
# model_kwargs['ls_init'] = (
# np.interp(mean_time, kern_times, kern_ls[ant_idx,:,1].mean(0)),
# np.interp(mean_time, kern_times, kern_ls[ant_idx,:,0].mean(0))
# )
#logging.warning(model_kwargs['ls_init'])
# initial ls_scale (if they exist)
# (phase, error, coords, lock, error_sigma_clip=4., m_type='sgp',
# iterations=1000, pargs=None,verbose=False,model_kwargs={}):
model_file = os.path.join(self.proj_dir,"model_{}_{}".format(start,stop))
jobs.append(executor.submit(
self._solve_interval_full,
model_file,
phase[:,time_slice,:,:],
error[:,time_slice,:,:],
(x, t[time_slice],d,f),
lock,
load_model = model_file,
error_sigma_clip = None,
m_type='sgp_full',
iterations=iterations,
pargs="Working on time chunk ({}) {} to ({}) {}".format(
start,timestamps[start],stop-1,timestamps[stop-1]),
verbose=verbose,
model_kwargs = model_kwargs
)
)
results = futures.wait(jobs)
if verbose:
logging.warning(results)
results = [j.result() for j in jobs]
phase_mean = np.zeros(phase.shape)
phase_weights = np.zeros(phase.shape)
variance_mean = np.zeros(variance.shape)
res_idx = 0
for j,aj in enumerate(time_idx[::shift]):
start = j*interval
stop = min((j+1)*interval,Nt)
time_slice = slice(start,stop,1)
res = results[res_idx]
p,v = res
phase_mean[:,time_slice,:,:] += p/(v+1e-6)
variance_mean[:,time_slice,:,:] += 1.
phase_weights[:,time_slice,:,:] += 1./(v+1e-6)
res_idx += 1
phase_mean /= (phase_weights+1e-6)
variance_mean /= (phase_weights+1e-6)
datapack.set_phase(phase_mean, ant_idx=ant_idx,time_idx=time_idx,dir_idx=dir_idx,freq_idx=freq_idx)
datapack.set_variance(variance_mean, ant_idx=ant_idx,time_idx=time_idx,dir_idx=dir_idx,freq_idx=freq_idx)
datapack.save(save_datapack)
if __name__=='__main__':
import os
if len(sys.argv) == 2:
starting_datapack = sys.argv[1]
else:
starting_datapack = "../../data/rvw_datapack_full_phase_dec27_unwrap.hdf5"
smoothing = Smoothing(starting_datapack,'projects')
model_kwargs = {'minibatch_size':500, 'M':1000,'feature_trainable':False,'ls_init':(5.,70,0.3),
'ls_trainable':(True,True,True), 'verbose':False,'likelihood_var_trainable':True}
smoothing.solve_and_apply_ensemble(starting_datapack.replace('.hdf5','_smoothed_ensemble.hdf5'),
-1, -1, -1, -1, iterations=1000,
interval = None, shift = 6, init_solutions='../../bayes/gp_params_fixed_scales.npz',
num_threads=1,verbose=True,model_kwargs = model_kwargs)
# # smoothing.solve_time_intervals("gp_params.npz",range(1,62),-1,-1,range(0,20),32,32,num_threads=16,verbose=True)
# refined_params = smoothing.refine_statistics_timeonly('gp_params.npz')
# print(refined_params.shape)
# smoothing.solve_time_intervals("gp_params_fixed_scales.npz",range(1,62),-1,-1,range(0,20),32,32,num_threads=16,verbose=True,refined_params=refined_params)
# plt.ion()
# smoothing.refine_statistics_timeonly('gp_params.npz')
# smoothing.refine_statistics('gp_params.npz')
# smoothing.refine_statistics_timeonly('gp_params_fixed_scales.npz')
# smoothing.refine_statistics('gp_params_fixed_scales.npz')
# plt.ioff()
# smoothing.apply_solutions(starting_datapack.replace('.hdf5','_refined_smoothed.hdf5'),
# "gp_params_fixed_scales.npz",range(1,62), -1, -1, range(0,20), 32, 32, num_threads=1,verbose=True)
May 4: Pick-up Truck (Early as possible) Drive to Antsirabe (170km 4h) City/Lakes/Fill Gas/Fill travel food May 5: Drive to Ranomafana National Park (230km 5h) May 6: Hike Ranomafana May 7: Ranomafana to Isalo National Park (350km 6h) May 8: Hike Isalo May 9: Isalo to Antsirabe (530km 10h) May 10: Antsirabe to Mahambo (590km 12h) -- OR -- Antsirabe to Toamasina (500km 10h) May 11: Mahambo to Ile Sainte Marie (Boat) -- OR -- Toamasina to Mahambo (90km 3h) Mahambo to Ile Sainte Marie (Boat) May 12: Ile Sainte Marie / Ile Aux Nattes May 13: Ile Sainte Marie / Ile Aux Nattes May 14: Ile Sainte Marie / Ile Aux Nattes May 15: Ile Sainte Marie to Mahambo (early Boat) Mahambo to Antanarivo (450km 10h) May 16: Antananrivo to Andasibe (150km 4-6h) May 17: Hike Andasibe May 18: Andasibe to Antanarivo (150km 4-6h) May 19: 2 am flight