3 parameter ARD kernel Tensorflow for Gaussian Process marginal likelihood optimization

save_the_kernel__save_the_world.py

import numpy as np
import tensorflow as tf
import tensorflow_probability as tfp
tfd = tfp.distributions
tfk = tfp.positive_semidefinite_kernels

from Starfish.emulator import PCAGrid

class InputTransformedKernel(tfk.PositiveSemidefiniteKernel):

    def __init__(self, kernel, transformation, name='InputTransformedKernel'):
        self._kernel = kernel
        self._transformation = transformation
        super().__init__(
            feature_ndims=kernel.feature_ndims,
            dtype=kernel.dtype,
            name=name)

    def apply(self, x1, x2):
        return self._kernel.apply(
            self._transformation(x1),
            self._transformation(x2))

    def matrix(self, x1, x2):
        return self._kernel.matrix(
            self._transformation(x1),
            self._transformation(x2))

    @property
    def batch_shape(self):
        return self._kernel.batch_shape

    def batch_shape_tensor(self):
        return self._kernel.batch_shape_tensor

class InputScaledKernel(InputTransformedKernel):

    def __init__(self, kernel, length_scales):
        super().__init__(
            kernel,
            lambda x: x / tf.expand_dims(length_scales,
                                     -(kernel.feature_ndims + 1)))
        
pca = PCAGrid.open('testPCA.hdf5')

X = tf.Variable(pca.gparams)
Y = tf.Variable(pca.w)

conf_priors = np.array(Starfish.config.PCA["priors"]).T
amplitude = tf.Variable(10. * tf.ones(4, dtype=tf.float64),  name='amplitude')

rv_scale = tfd.Gamma(
    concentration=tf.Variable(conf_priors[0], name='s'), 
    rate=tf.Variable(conf_priors[1], name='r'),
    name='length_scale_prior')

lscales = np.array([[300., 30., 30.]], dtype='f8')
length_scale = tf.Variable(tf.tile(lscales, (4, 1)), name='length_scale')
se_kernel = tfk.ExponentiatedQuadratic(amplitude, name='ARD_kernel')
kernel = InputScaledKernel(se_kernel, length_scale)

gp = tfd.GaussianProcess(
    kernel=kernel, 
    index_points=X,
)

log_likelihood = tf.reduce_sum(gp.log_prob(Y)) + tf.reduce_sum(rv_scale.log_prob(length_scale))
opt = tf.train.AdamOptimizer(1e-6).minimize(-log_likelihood)

mat = kernel.matrix(X, X)
eigs = tf.linalg.eigvalsh(mat)
with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    X_, mat_, eigs_ = sess.run([X, mat, eigs])
print(np.all(eigs_ > 0))

get the PCA file here and place it in the root of wherever you plan to run the code in.
Get the starfish package like this

pip install git+https://github.com/iancze/Starfish.git@develop#egg=astrostarfish