utils.py

# Copyright 2018 Google LLC
#
# 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
#
#     https://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.
# ==============================================================================

"""Helper functions for creating the training graph and plotting.

"""

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

import os
from matplotlib.backends.backend_pdf import PdfPages
import matplotlib.pyplot as plt
import numpy as np
import tensorflow as tf

import ensembles  # pylint: disable=g-bad-import-order


np.seterr(invalid="ignore")


def get_place_cell_ensembles(
    env_size, neurons_seed, targets_type, lstm_init_type, n_pc, pc_scale):
  """Create the ensembles for the Place cells."""
  place_cell_ensembles = [
      ensembles.PlaceCellEnsemble(
          n,
          stdev=s,
          pos_min=-env_size / 2.0,
          pos_max=env_size / 2.0,
          seed=neurons_seed,
          soft_targets=targets_type,
          soft_init=lstm_init_type)
      for n, s in zip(n_pc, pc_scale)
  ]
  return place_cell_ensembles


def get_head_direction_ensembles(
    neurons_seed, targets_type, lstm_init_type, n_hdc, hdc_concentration):
  """Create the ensembles for the Head direction cells."""
  head_direction_ensembles = [
      ensembles.HeadDirectionCellEnsemble(
          n,
          concentration=con,
          seed=neurons_seed,
          soft_targets=targets_type,
          soft_init=lstm_init_type)
      for n, con in zip(n_hdc, hdc_concentration)
  ]
  return head_direction_ensembles


def encode_initial_conditions(init_pos, init_hd, place_cell_ensembles,
                              head_direction_ensembles):
  initial_conds = []
  for ens in place_cell_ensembles:
    initial_conds.append(
        tf.squeeze(ens.get_init(init_pos[:, tf.newaxis, :]), axis=1))
  for ens in head_direction_ensembles:
    initial_conds.append(
        tf.squeeze(ens.get_init(init_hd[:, tf.newaxis, :]), axis=1))
  return initial_conds


def encode_targets(target_pos, target_hd, place_cell_ensembles,
                   head_direction_ensembles):
  ensembles_targets = []
  for ens in place_cell_ensembles:
    ensembles_targets.append(ens.get_targets(target_pos))
  for ens in head_direction_ensembles:
    ensembles_targets.append(ens.get_targets(target_hd))
  return ensembles_targets


def clip_all_gradients(g, var, limit):
  # print(var.name)
  return (tf.clip_by_value(g, -limit, limit), var)


def clip_bottleneck_gradient(g, var, limit):
  if ("bottleneck" in var.name or "pc_logits" in var.name):
    return (tf.clip_by_value(g, -limit, limit), var)
  else:
    return (g, var)


def no_clipping(g, var):
  return (g, var)


def concat_dict(acc, new_data):
  """Dictionary concatenation function."""

  def to_array(kk):
    if isinstance(kk, np.ndarray):
      return kk
    else:
      return np.asarray([kk])

  for k, v in new_data.iteritems():
    if isinstance(v, dict):
      if k in acc:
        acc[k] = concat_dict(acc[k], v)
      else:
        acc[k] = concat_dict(dict(), v)
    else:
      v = to_array(v)
      if k in acc:
        acc[k] = np.concatenate([acc[k], v])
      else:
        acc[k] = np.copy(v)
  return acc


def get_scores_and_plot(scorer,
                        data_abs_xy,
                        activations,
                        directory,
                        filename,
                        plot_graphs=True,  # pylint: disable=unused-argument
                        nbins=20,  # pylint: disable=unused-argument
                        cm="jet",
                        sort_by_score_60=True):
  """Plotting function."""

  # Concatenate all trajectories
  xy = data_abs_xy.reshape(-1, data_abs_xy.shape[-1])
  act = activations.reshape(-1, activations.shape[-1])
  n_units = act.shape[1]
  # Get the rate-map for each unit
  s = [
      scorer.calculate_ratemap(xy[:, 0], xy[:, 1], act[:, i])
      for i in xrange(n_units)
  ]
#   print('UNIT0')
#   print(s[0])
  # Get the scores
  score_60, score_90, max_60_mask, max_90_mask, sac = zip(
      *[scorer.get_scores(rate_map) for rate_map in s])
  # Separations
  # separations = map(np.mean, max_60_mask)
  # Sort by score if desired
  if sort_by_score_60:
    ordering = np.argsort(-np.array(score_60))
  else:
    ordering = range(n_units)
  # Plot
  cols = 16
  rows = int(np.ceil(n_units / cols))
  fig = plt.figure(figsize=(24, rows * 4))
  for i in xrange(n_units):
    rf = plt.subplot(rows * 2, cols, i + 1)
    acr = plt.subplot(rows * 2, cols, n_units + i + 1)
    if i < n_units:
      index = ordering[i]
      title = "%d (%.2f)" % (index, score_60[index])
      # Plot the activation maps
      scorer.plot_ratemap(s[index], ax=rf, title=title, cmap=cm)
      # Plot the autocorrelation of the activation maps
      scorer.plot_sac(
          sac[index],
          mask_params=max_60_mask[index],
          ax=acr,
          title=title,
          cmap=cm)
  # Save
  if not os.path.exists(directory):
    os.makedirs(directory)
  with PdfPages(os.path.join(directory, filename), "w") as f:
    plt.savefig(f, format="pdf")
  plt.close(fig)
  return (np.asarray(score_60), np.asarray(score_90),
          np.asarray(map(np.mean, max_60_mask)),
          np.asarray(map(np.mean, max_90_mask)))