util.py

# coding=utf-8
# Copyright 2018 The Google Research Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Tools."""
# pylint: disable=missing-docstring,g-doc-args,g-doc-return-or-yield
# pylint: disable=g-short-docstring-punctuation,g-no-space-after-docstring-summary
# pylint: disable=invalid-name,broad-except

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections
import time

from absl import logging


import numpy as np
import tensorflow.compat.v1 as tf

from tensorflow_probability import bijectors as tfb
from tensorflow_probability import distributions as tfd
from tensorflow_probability import edward2 as ed

from tensorflow_probability.python.experimental.edward2.generated_random_variables import Normal
from tensorflow_probability.python.experimental.edward2.interceptor import tape
from tensorflow_probability.python.experimental.edward2.program_transformations import make_log_joint_fn

from tensorflow.python.ops.parallel_for import pfor

import program_transformations as program_transformations

FLAGS = tf.app.flags.FLAGS

# pylint: disable=g-import-not-at-top
try:
  import __builtin__
except ImportError:
  # Python 3
  import builtins as __builtin__
# pylint: enable=g-import-not-at-top

__all__ = [
    'condition_number_cp',
    'condition_number_ncp',
    'compute_V_cp',
    'compute_V_ncp',
    'mean_field_variational_inference',
    'approximate_mcmc_step_size',
]


def compute_V_cp(q, v):
  r = (v * q + q + 1.)
  return np.array([[1. + v, 1.], [1., q*v + 1.]]) / r


def compute_V_ncp(q, v):
  r = 1 / (v * q + q + 1)
  return r * np.array([[q + 1, -np.sqrt(v)*q], [-np.sqrt(v)*q, v*q + 1]])


def condition_number_cp(q, v):
  sqrt_det = 2 * np.sqrt((v * q + 1) * (v * q + 1) - v * (v * q + q + 1) *
                         (v * q + 1) / (v + 1))
  lambda1 = 2*(v*q + 1) - sqrt_det
  lambda2 = 2*(v*q + 1) + sqrt_det
  return lambda2 / lambda1


def condition_number_ncp(q, v):
  sqrt_det = 2 * np.sqrt((v * q + 1) * (v * q + 1) - (v * q + q + 1) *
                         (v * q + 1) / (q + 1))
  lambda1 = 2*(v*q + 1) - sqrt_det
  lambda2 = 2*(v*q + 1) + sqrt_det
  return lambda2 / lambda1


def mean_field_variational_inference(model, *args, **kwargs):
  num_optimization_steps = kwargs.get('num_optimization_steps', 2000)
  del kwargs['num_optimization_steps']

  (variational_model,
   variational_parameters) = program_transformations.make_variational_model(
       model, *args, **kwargs)

  log_joint = make_log_joint_fn(model)
  def target(**parameters):
    full_kwargs = dict(parameters, **kwargs)
    return log_joint(*args, **full_kwargs)

  log_joint_q = make_log_joint_fn(variational_model)
  def target_q(**parameters):
    return log_joint_q(*args, **parameters)

  elbo_sum = 0.
  for _ in range(16):
    with tape() as variational_tape:
      _ = variational_model(*args)

    params = variational_tape
    elbo_sum = elbo_sum + target(**params) - target_q(**params)

  elbo = elbo_sum / 16.
  best_elbo = None

  learning_rate_ph = tf.compat.v1.placeholder(shape=[], dtype=tf.float32)
  learning_rate = tf.Variable(learning_rate_ph, trainable=False)
  optimizer = tf.compat.v1.train.AdamOptimizer(learning_rate=learning_rate)
  train = optimizer.minimize(-elbo)
  init = tf.compat.v1.global_variables_initializer()

  start_time = time.time()
  for learning_rate_val in [0.01, 0.1, 0.01, 0.1, 0.01, 0.1]:
    feed_dict = {learning_rate_ph: learning_rate_val}
    with tf.compat.v1.Session() as sess:
      sess.run(init, feed_dict=feed_dict)

      this_timeline = []
      print('VI with {} optimization steps'.format(num_optimization_steps))
      for _ in range(num_optimization_steps):
        _, e = sess.run([train, elbo], feed_dict=feed_dict)
        this_timeline.append(e)

      this_elbo = np.mean(this_timeline[-100:])
      if best_elbo is None or best_elbo < this_elbo:
        timeline = this_timeline
        best_elbo = this_elbo

        vals = sess.run(
            list(variational_parameters.values()), feed_dict=feed_dict)
        learned_variational_params = collections.OrderedDict(
            zip(variational_parameters.keys(), vals))

      vi_time = time.time() - start_time

  results = collections.OrderedDict()
  results['vp'] = learned_variational_params
  print('ELBO: {}'.format(best_elbo))

  return results, best_elbo, timeline, vi_time


def _marshal(*rvs):
  """Args: a list of ed.RandomVariables each with vector or scalar event shape
  (which must be staticly known), and all having the same batch shape.

  Returns: a Tensor from concatenating their values along a single vector
  dimension.
  """
  vector_rvs = []
  for rv in rvs:
    v = rv.value
    if v.shape.ndims == 0:
      vector_rvs.append([v])
    else:
      vector_rvs.append(v)
  print(vector_rvs)
  return tf.concat(vector_rvs, axis=-1)


def _to_vector_shape(tensor_shape):
  if tensor_shape.ndims > 1:
    raise Exception('cannot convert {} to vector shape!'.format(tensor_shape))
  elif tensor_shape.ndims == 0:
    return tf.TensorShape([1])
  return tensor_shape


def _tensorshape_size(tensor_shape):
  if tensor_shape.ndims > 1:
    raise Exception(
        'shapes of ndims >1 are bad! (saw: {})!'.format(tensor_shape))
  elif tensor_shape.ndims == 0:
    return 1
  return tensor_shape[0].value


def get_iaf_elbo(target, num_mc_samples, param_shapes):
  shape_sizes = [_tensorshape_size(pshape) for pshape in param_shapes.values()]
  overall_shape = [sum(shape_sizes)]

  def unmarshal(variational_sample):
    results = []
    n_dimensions_used = 0
    for (n_to_add, result_shape) in zip(shape_sizes, param_shapes.values()):
      result = variational_sample[Ellipsis, n_dimensions_used:
                                  n_dimensions_used + n_to_add]
      results.append(tf.reshape(result, result_shape))
      n_dimensions_used += n_to_add
    return tuple(results)

  variational_dist = tfd.TransformedDistribution(
      distribution=tfd.Normal(loc=0., scale=1.),
      bijector=tfb.Invert(
          tfb.MaskedAutoregressiveFlow(
              shift_and_log_scale_fn=tfb.masked_autoregressive_default_template(
                  hidden_layers=[256, 256]))),
      event_shape=overall_shape,
      name='q_iaf')

  variational_samples = variational_dist.sample(num_mc_samples)
  target_q_sum = tf.reduce_sum(input_tensor=variational_dist.log_prob(variational_samples))
  target_sum = 0.
  for s in range(num_mc_samples):
    params = unmarshal(variational_samples[s, Ellipsis])
    target_sum = target_sum + target(*params)

  energy = target_sum / float(num_mc_samples)
  entropy = -target_q_sum / float(num_mc_samples)
  elbo = energy + entropy

  tf.compat.v1.summary.scalar('energy', energy)
  tf.compat.v1.summary.scalar('entropy', entropy)
  tf.compat.v1.summary.scalar('elbo', elbo)

  return elbo


def get_mean_field_elbo(model, target, num_mc_samples, model_args,
                        model_obs_kwargs, vi_kwargs):
  if FLAGS.reparameterise_variational and 'cVIP' in FLAGS.method:
    combined_kwargs = model_obs_kwargs.copy()
    combined_kwargs.update(vi_kwargs)
    variational_model, variational_parameters = make_variational_model_special(
        model, *model_args, **combined_kwargs)
  else:
    variational_model, variational_parameters = program_transformations.make_variational_model(
        model, *model_args, **model_obs_kwargs)

  log_joint_q = make_log_joint_fn(variational_model)

  def target_q(**parameters):
    return log_joint_q(*model_args, **parameters)

  #beta = tf.get_variable("beta", trainable=False, initializer=0.)
  #beta_incr = tf.assign(beta, tf.clip_by_value(beta + 0.1*beta + 0.0000001, 0., 1.))

  #with tf.control_dependencies([beta_incr]):

  def loop_body(mc_sample):
    with tape() as variational_tape:
      _ = variational_model(*model_args)

      params = variational_tape.values()

      energy = target(*params)
      entropy = tf.negative(target_q(**variational_tape))
      return energy + entropy

  if num_mc_samples == 1:
    elbo = tf.reduce_sum(loop_body(0))
  else:
    elbo = tf.reduce_sum(input_tensor=pfor(loop_body, num_mc_samples)) / num_mc_samples
  tf.compat.v1.summary.scalar('elbo', elbo)
  return elbo, variational_parameters


def get_approximate_step_size(variational_parameters, num_leapfrog_steps):
  return [
      variational_parameters[key] / num_leapfrog_steps**2
      for key in variational_parameters.keys()
      if key.endswith('_scale')
  ]


# FIXME: need to make this nicer than with all these weird kwargs
def approximate_mcmc_step_size(model, *args, **kwargs):

  with tf.compat.v1.variable_scope('approx_step_size_{}'.format(model.__name__)):
    if 'diagnostics' in kwargs.keys():
      diagnostics = kwargs.pop('diagnostics')
    else:
      diagnostics = False

    if 'num_leapfrog_steps' in kwargs.keys():
      num_leapfrog_steps = kwargs.pop('num_leapfrog_steps')
    else:
      num_leapfrog_steps = 4

    results, final_elbo_val, _, vi_time = mean_field_variational_inference(
        model, *args, **kwargs)
    stepsize = [(np.array(np.array(results['vp'][key], dtype=np.float32)) /
                 (float(num_leapfrog_steps)**2))
                for key in results['vp'].keys()
                if key.endswith('_scale')]

    if diagnostics:
      print('Estimated goodness of {}: {}'.format(model.__name__,
                                                  final_elbo_val))
      print('Estimated stepsize of {}: {}'.format(model.__name__, stepsize))

  return stepsize, final_elbo_val, vi_time


def stddvs_to_mcmc_step_sizes(results, num_leapfrog_steps):
  stepsize = [(np.sqrt(2 * np.mean(results[key])) / float(num_leapfrog_steps))
              for key in results.keys()
              if key.endswith('_scale')]

  return stepsize


def estimate_true_mean(sample_groups, esss):

  true_mean = [0 for group in range(len(sample_groups))]

  r = float(sum(esss))

  for group in range(len(sample_groups)):

    samples = sample_groups[group]
    mean = [np.mean(v) for v in samples]

    true_mean[group] = [
        (true_mean[group] + esss[group] * var_mean / r) for var_mean in mean
    ]

  return true_mean


def make_variational_model_special(model, *args, **kwargs):

  variational_parameters = collections.OrderedDict()
  param_params = kwargs['parameterisation']

  def get_or_init(name, a, b, L=None, std_mean=None,
                  prior_mean=None, prior_scale=None, shape=None):

    loc_name = name + '_loc'
    scale_name = name + '_scale'

    if loc_name in variational_parameters.keys() and \
        scale_name in variational_parameters.keys():
      return (variational_parameters[loc_name],
              variational_parameters[scale_name])
    else:
      # shape must not be None
      pre_loc = tf.compat.v1.get_variable(
          name=loc_name, initializer=1e-2 * tf.random.normal(shape, dtype=tf.float32))
      pre_scale = tf.nn.softplus(
          tf.compat.v1.get_variable(
              name=scale_name,
              initializer=-2 * tf.ones(shape, dtype=tf.float32)))
      variational_parameters[loc_name] = (a + 0.1) * pre_loc
      variational_parameters[scale_name] = pre_scale**(b + 0.1)

      return (variational_parameters[loc_name],
              variational_parameters[scale_name])

  def mean_field(rv_constructor, *rv_args, **rv_kwargs):

    name = rv_kwargs['name']
    if name not in kwargs.keys():
      rv = rv_constructor(*rv_args, **rv_kwargs)

      try:
        a, b = param_params[name + '_a'], param_params[name + '_b']
      except Exception as err:
        print('couldn\'t get centering params for variable {}: {}'.format(
            name, err))
        a, b = 1., 1.
      loc, scale = get_or_init(name, a=a, b=b, shape=rv.shape)

      # NB: name must be the same as original variable,
      # in order to be able to do black-box VI (setting
      # parameters to variational values obtained via trace).
      return Normal(loc=loc, scale=scale, name=name)
    else:
      rv_kwargs['value'] = kwargs[name]
      return rv_constructor(*rv_args, **rv_kwargs)

  def variational_model(*args):
    with ed.interception(mean_field):
      return model(*args)

  _ = variational_model(*args)

  return variational_model, variational_parameters


def variational_inits_from_params(learned_variational_params, param_names,
                                  num_inits):
  """Sample from a normal variational dist, given saved parameters."""
  locs = collections.OrderedDict()
  stddevs = collections.OrderedDict()
  samples = collections.OrderedDict()
  for k, v in learned_variational_params.items():
    if k.endswith('_loc'):
      locs[k[:-4]] = v
    elif k.endswith('_scale'):
      stddevs[k[:-6]] = v

  for k in param_names:
    shape = (num_inits,) + np.asarray(locs[k]).shape
    samples[k] = (np.random.randn(*shape) * stddevs[k] + locs[k]).astype(
        np.float32)
  return samples


def print(*args):  # pylint: disable=redefined-builtin
  __builtin__.print(*args)
  logging.info(' '.join(args))


def reject_outliers(data, m=1.5):
  ret = data[abs(data - np.mean(data)) < m * np.std(data)]
  if len(ret) > 0:
    return ret
  else:
    return data


def get_min_ess_other(ess):

  ess = [[np.nan_to_num(e) for e in es] for es in ess]

  min_ess = []
  for c in range(FLAGS.num_chains):
    min_ess_c = min([np.array(e).min() for e in ess[c]])
    #print("   Min ess of chain {} is {}.".format(c, min_ess_c))
    min_ess.append(min_ess_c)

  min_ess = reject_outliers(np.array(min_ess))
  print('   Filtred {} outliers.'.format(FLAGS.num_chains - len(min_ess)))

  mean_ess = np.mean(min_ess)
  sem_ess = np.std(min_ess) / np.sqrt(len(min_ess))

  return mean_ess, sem_ess


def get_min_ess(ess):

  ess = [np.nan_to_num(e) for e in ess]

  min_ess = []
  for c in range(FLAGS.num_chains):
    min_ess_c = min([np.array(e[c]).min() for e in ess])
    min_ess.append(min_ess_c)

  # min_ess = reject_outliers(np.array(min_ess))
  # print("   Filtred {} outliers.".format(FLAGS.num_chains - len(min_ess)))

  mean_ess = np.mean(min_ess)
  sem_ess = np.std(min_ess) / np.sqrt(len(min_ess))

  return mean_ess, sem_ess