text: # Copyright 2017 The TensorFlow Authors. All Rights Reserved.
  #
  # 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.
  # ==============================================================================
  """Abstract base for state space models."""
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import abc
  import collections
  
  import numpy
  
  from tensorflow.contrib.layers.python.layers import layers
  
  from tensorflow.contrib.timeseries.python.timeseries import math_utils
  from tensorflow.contrib.timeseries.python.timeseries import model
  from tensorflow.contrib.timeseries.python.timeseries import model_utils
  from tensorflow.contrib.timeseries.python.timeseries.feature_keys import TrainEvalFeatures
  from tensorflow.contrib.timeseries.python.timeseries.state_space_models import kalman_filter
  
  from tensorflow.python.estimator import estimator_lib
  from tensorflow.python.framework import constant_op
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import ops
  from tensorflow.python.framework import tensor_shape
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import control_flow_ops
  from tensorflow.python.ops import gen_math_ops
  from tensorflow.python.ops import linalg_ops
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops import variable_scope
  
  
  class StateSpaceModelConfiguration(
      collections.namedtuple(
          typename="StateSpaceModelConfiguration",
          field_names=[
              "num_features", "use_observation_noise", "dtype",
              "covariance_prior_fn", "bayesian_prior_weighting",
              "filtering_postprocessor", "trainable_start_state",
              "exogenous_noise_increases", "exogenous_noise_decreases",
              "exogenous_feature_columns", "exogenous_update_condition",
              "filtering_maximum_posterior_variance_ratio",
              "filtering_minimum_posterior_variance",
              "transition_covariance_initial_log_scale_bias",
              "static_unrolling_window_size_threshold"])):
    """Configuration options for StateSpaceModels."""
  
    def __new__(
        cls,
        num_features=1,
        use_observation_noise=True,
        dtype=dtypes.float32,
        covariance_prior_fn=math_utils.log_noninformative_covariance_prior,
        bayesian_prior_weighting=True,
        filtering_postprocessor=None,
        trainable_start_state=False,
        exogenous_noise_increases=True,
        exogenous_noise_decreases=False,
        exogenous_feature_columns=None,
        exogenous_update_condition=None,
        filtering_maximum_posterior_variance_ratio=1e6,
        filtering_minimum_posterior_variance=0.,
        transition_covariance_initial_log_scale_bias=-5.,
        static_unrolling_window_size_threshold=None):
      """Configuration options for StateSpaceModels.
  
      Args:
        num_features: Output dimension for model
        use_observation_noise: If true, observations are modeled as noisy
          functions of the current state. If false, observations are a
          deterministic function of the current state. Only applicable to the
          top-level model in an ensemble. Consider also changing the
          transition_covariance_initial_log_scale_bias when disabling observation
          noise, as its default setting assumes that observation noise is part of
          the model.
        dtype: The float dtype to use when defining the model.
        covariance_prior_fn: A function mapping from a covariance matrix to a
            scalar value (e.g. log likelihood) which can be summed across
            matrices. Defaults to an independent Jeffreys prior on the diagonal
            elements (regularizing as log(1. / variance)). To use a flat prior
            (i.e. no regularization), set to `lambda _: 0.`.  Defaults to
            relatively uninformative priors on state transition and observation
            noise, which have the effect of encouraging low-noise solutions which
            provide confident predictions when possible. Without regularization,
            transition noise tends to remain high, and multi-step predictions are
            under-confident.
        bayesian_prior_weighting: If True, weights the prior (covariance_prior_fn)
            based on an estimate of the full dataset size. If False, weights it
            based on the mini-batch window size, which (while statistically
            improper) can lead to more desirable low-noise solutions in cases
            where the full dataset is large enough to overwhelm the prior.
        filtering_postprocessor: A FilteringStepPostprocessor object to use,
            useful for ignoring anomalies in training data.
        trainable_start_state: If True, start state may depend on trainable
            Variables. If False, it will not.
        exogenous_noise_increases: If True, exogenous regressors can add to model
            state, increasing uncertainty. If both this parameter and
            exogenous_noise_decreases are False, exogenous regressors are ignored.
        exogenous_noise_decreases: If True, exogenous regressors can "set" model
            state, decreasing uncertainty. If both this parameter and
            exogenous_noise_increases are False, exogenous regressors are ignored.
        exogenous_feature_columns: A list of `tf.feature_column`s (for example
            `tf.feature_column.embedding_column`) corresponding to exogenous
            features which provide extra information to the model but are not part
            of the series to be predicted. Passed to
            `tf.feature_column.input_layer`.
        exogenous_update_condition: A function taking two Tensor arguments `times`
            (shape [batch size]) and `features` (a dictionary mapping exogenous
            feature keys to Tensors with shapes [batch size, ...]) and returning a
            boolean Tensor with shape [batch size] indicating whether state should
            be updated using exogenous features for each part of the batch. Where
            it is False, no exogenous update is performed. If None (default),
            exogenous updates are always performed. Useful for avoiding "leaky"
            frequent exogenous updates when sparse updates are desired. Called
            only during graph construction.
        filtering_maximum_posterior_variance_ratio: The maximum allowed ratio of
            two diagonal entries in a state covariance matrix just prior to
            filtering. Lower values mean that filtering will be more numerically
            stable, at the cost of artificially increasing estimated uncertainty
            in some cases. This parameter can be important when learning a
            transition matrix.
        filtering_minimum_posterior_variance: The minimum diagonal value in a
            state covariance matrix just prior to filtering, preventing numerical
            instability due to deterministic beliefs (sometimes an issue when
            learning transition matrices). This value should be set several orders
            of magnitude below any expected minimum state uncertainty.
        transition_covariance_initial_log_scale_bias: Controls the initial
            tradeoff between the transition noise covariance matrix and the
            observation noise covariance matrix, on a log scale (the elements of
            the transition noise covariance matrix are proportional to `e^{X +
            transition_covariance_initial_log_scale_bias}` where `X` is learned
            and may depend on input statistics, observation noise covariance is
            proportional to `e^{Y -
            transition_covariance_initial_log_scale_bias}`). For models *with*
            observation noise, -5 is a reasonable value. Models which do not use
            observation noise, and are not part of an ensemble which does use
            observation noise, should have this set to 0 or more to avoid
            numerical issues due to filtering with too little noise.
        static_unrolling_window_size_threshold: Only relevant for the top-level
            StateSpaceModel in an ensemble; enables switching between static and
            dynamic looping (if not None, default, meaning that no static
            unrolling is performed) based on the window size (windows with this
            size and smaller will have their graphs unrolled statically). See the
            SequentialTimeSeriesModel constructor for details.
      Returns:
        A StateSpaceModelConfiguration object.
      """
      if exogenous_feature_columns is None:
        exogenous_feature_columns = []
      return super(StateSpaceModelConfiguration, cls).__new__(
          cls, num_features, use_observation_noise, dtype,
          covariance_prior_fn, bayesian_prior_weighting,
          filtering_postprocessor, trainable_start_state,
          exogenous_noise_increases, exogenous_noise_decreases,
          exogenous_feature_columns, exogenous_update_condition,
          filtering_maximum_posterior_variance_ratio,
          filtering_minimum_posterior_variance,
          transition_covariance_initial_log_scale_bias,
          static_unrolling_window_size_threshold)
  
  
  class StateSpaceModel(model.SequentialTimeSeriesModel):
    """Base class for linear state space models.
  
    Sub-classes can specify the model to be learned by overriding
    get_state_transition, get_noise_transform, and get_observation_model.
  
    See kalman_filter.py for a detailed description of the class of models covered
    by StateSpaceModel.
  
    Briefly, state space models are defined by a state transition equation:
  
    state[t] = StateTransition * state[t-1] + NoiseTransform * StateNoise[t]
               + ExogenousNoiseIncreasing[t]
    StateNoise[t] ~ Gaussian(0, StateNoiseCovariance)
    ExogenousNoiseIncreasing[t] ~ Gaussian(ExogenousNoiseIncreasingMean[t],
                                           ExogenousNoiseIncreasingCovariance[t])
  
    And an observation model:
  
    observation[t] = ObservationModel * state[t] + ObservationNoise[t]
    ObservationNoise[t] ~ Gaussian(0, ObservationNoiseCovariance)
  
    Additionally, exogenous regressors can act as observations, decreasing
    uncertainty:
  
    ExogenousNoiseDecreasingObservation[t] ~ Gaussian(
        ExogenousNoiseDecreasingMean[t], ExogenousNoiseDecreasingCovariance[t])
  
    Attributes:
      kalman_filter: If initialize_graph has been called, the initialized
          KalmanFilter to use for inference. None otherwise.
      prior_state_mean: If initialize_graph has been called, a
          Variable-parameterized Tensor with shape [state dimension];
          the initial prior mean for one or more time series. None otherwise.
      prior_state_var: If initialize_graph has been called, a
          Variable-parameterized Tensor with shape [state dimension x state
          dimension]; the initial prior covariance. None otherwise.
      state_transition_noise_covariance: If initialize_graph has been called, a
          Variable-parameterized Tensor with shape [state noise dimension x state
          noise dimension] indicating the amount of noise added at each
          transition.
    """
  
    def __init__(self, configuration):
      """Initialize a state space model.
  
      Args:
        configuration: A StateSpaceModelConfiguration object.
      """
      self._configuration = configuration
      if configuration.filtering_postprocessor is not None:
        filtering_postprocessor_names = (
            configuration.filtering_postprocessor.output_names)
      else:
        filtering_postprocessor_names = []
      super(StateSpaceModel, self).__init__(
          train_output_names=(["mean", "covariance", "log_likelihood"]
                              + filtering_postprocessor_names),
          predict_output_names=["mean", "covariance"],
          num_features=configuration.num_features,
          normalize_features=True,
          dtype=configuration.dtype,
          exogenous_feature_columns=configuration.exogenous_feature_columns,
          exogenous_update_condition=configuration.exogenous_update_condition,
          static_unrolling_window_size_threshold=
          configuration.static_unrolling_window_size_threshold)
      self._kalman_filter = None
      self.prior_state_mean = None
      self.prior_state_var = None
      self.state_transition_noise_covariance = None
      self._total_observation_count = None
      self._observation_noise_covariance = None
      # Capture the current variable scope and use it to define all model
      # variables. Especially useful for ensembles, where variables may be defined
      # for every component model in one function call, which would otherwise
      # prevent the user from separating variables from different models into
      # different scopes.
      self._variable_scope = variable_scope.get_variable_scope()
  
    def transition_power_noise_accumulator(self, num_steps):
      r"""Sum a transitioned covariance matrix over a number of steps.
  
      Computes
  
        \sum_{i=0}^{num_steps - 1} (
          state_transition^i
          * state_transition_noise_covariance
          * (state_transition^i)^T)
  
      If special cases are available, overriding this function can lead to more
      efficient inferences.
  
      Args:
        num_steps: A [...] shape integer Tensor with numbers of steps to compute
          power sums for.
      Returns:
        The computed power sum, with shape [..., state dimension, state
        dimension].
      """
      # TODO(allenl): This general case should use cumsum if transition_to_powers
      # can be computed in constant time (important for correlated ensembles,
      # where transition_power_noise_accumulator special cases cannot be
      # aggregated from member models).
      noise_transform = ops.convert_to_tensor(self.get_noise_transform(),
                                              self.dtype)
      noise_transformed = math_ops.matmul(
          math_ops.matmul(noise_transform,
                          self.state_transition_noise_covariance),
          noise_transform,
          transpose_b=True)
      noise_additions = math_utils.power_sums_tensor(
          math_ops.reduce_max(num_steps) + 1,
          ops.convert_to_tensor(self.get_state_transition(), dtype=self.dtype),
          noise_transformed)
      return array_ops.gather(noise_additions, indices=num_steps)
  
    def transition_to_powers(self, powers):
      """Raise the transition matrix to a batch of powers.
  
      Computes state_transition^powers. If special cases are available, overriding
      this function can lead to more efficient inferences.
  
      Args:
        powers: A [...] shape integer Tensor with powers to raise the transition
          matrix to.
      Returns:
        The computed matrix powers, with shape [..., state dimension, state
        dimension].
      """
      return math_utils.matrix_to_powers(
          ops.convert_to_tensor(self.get_state_transition(), dtype=self.dtype),
          powers)
  
    def _window_initializer(self, times, state):
      """Prepare to impute across the gaps in a window."""
      _, _, priors_from_time = state
      times = ops.convert_to_tensor(times)
      priors_from_time = ops.convert_to_tensor(priors_from_time)
      intra_batch_gaps = array_ops.reshape(times[:, 1:] - times[:, :-1], [-1])
      # Ignore negative starting gaps, since there will be transient start times
      # as inputs statistics are computed.
      starting_gaps = math_ops.maximum(times[:, 0] - priors_from_time, 0)
      # Pre-define transition matrices raised to powers (and their sums) for every
      # gap in this window. This avoids duplicate computation (for example many
      # steps will use the transition matrix raised to the first power) and
      # batches the computation rather than doing it inside the per-step loop.
      unique_gaps, _ = array_ops.unique(
          array_ops.concat([intra_batch_gaps, starting_gaps], axis=0))
      self._window_power_sums = self.transition_power_noise_accumulator(
          unique_gaps)
      self._window_transition_powers = self.transition_to_powers(unique_gaps)
      self._window_gap_sizes = unique_gaps
  
    def _lookup_window_caches(self, caches, indices):
      _, window_power_ids = array_ops.unique(
          array_ops.concat(
              [
                  self._window_gap_sizes, math_ops.cast(
                      indices, self._window_gap_sizes.dtype)
              ],
              axis=0))
      all_gathered_indices = []
      for cache in caches:
        gathered_indices = array_ops.gather(
            cache, window_power_ids[-array_ops.shape(indices)[0]:])
        gathered_indices.set_shape(indices.get_shape().concatenate(
            gathered_indices.get_shape()[-2:]))
        all_gathered_indices.append(gathered_indices)
      return all_gathered_indices
  
    def _cached_transition_powers_and_sums(self, num_steps):
      return self._lookup_window_caches(
          caches=[self._window_transition_powers, self._window_power_sums],
          indices=num_steps)
  
    def _imputation_step(self, current_times, state):
      """Add state transition noise to catch `state` up to `current_times`.
  
      State space models are inherently sequential, so we need to "predict
      through" any missing time steps to catch up each element of the batch to its
      next observation/prediction time.
  
      Args:
        current_times: A [batch size] Tensor of times to impute up to, not
            inclusive.
        state: A tuple of (mean, covariance, previous_times) having shapes
            mean; [batch size x state dimension]
            covariance; [batch size x state dimension x state dimension]
            previous_times; [batch size]
      Returns:
        Imputed model state corresponding to the `state` argument.
      """
      estimated_state, estimated_state_var, previous_times = state
      # Ignore negative imputation intervals due to transient start time
      # estimates.
      catchup_times = math_ops.maximum(current_times - previous_times, 0)
      transition_matrices, transition_noise_sums = (  # pylint: disable=unbalanced-tuple-unpacking
          self._cached_transition_powers_and_sums(catchup_times))
      estimated_state = self._kalman_filter.predict_state_mean(
          estimated_state, transition_matrices)
      estimated_state_var = self._kalman_filter.predict_state_var(
          estimated_state_var, transition_matrices, transition_noise_sums)
      return (estimated_state, estimated_state_var,
              previous_times + catchup_times)
  
    def _filtering_step(self, current_times, current_values, state, predictions):
      """Compute posteriors and accumulate one-step-ahead predictions.
  
      Args:
        current_times: A [batch size] Tensor for times for each observation.
        current_values: A [batch size] Tensor of values for each observation.
        state: A tuple of (mean, covariance, previous_times) having shapes
            mean; [batch size x state dimension]
            covariance; [batch size x state dimension x state dimension]
            previous_times; [batch size]
        predictions: A dictionary containing mean and covariance Tensors, the
            output of _prediction_step.
      Returns:
        A tuple of (posteriors, outputs):
          posteriors: Model state updated to take `current_values` into account.
          outputs: The `predictions` dictionary updated to include "loss" and
              "log_likelihood" entries (loss simply being negative log
              likelihood).
      """
      estimated_state, estimated_state_covariance, previous_times = state
      observation_model = self.get_broadcasted_observation_model(current_times)
      imputed_to_current_step_assert = control_flow_ops.Assert(
          math_ops.reduce_all(math_ops.equal(current_times, previous_times)),
          ["Attempted to perform filtering without imputation/prediction"])
      with ops.control_dependencies([imputed_to_current_step_assert]):
        estimated_state_covariance = math_utils.clip_covariance(
            estimated_state_covariance,
            self._configuration.filtering_maximum_posterior_variance_ratio,
            self._configuration.filtering_minimum_posterior_variance)
        (filtered_state, filtered_state_covariance,
         log_prob) = self._kalman_filter.do_filter(
             estimated_state=estimated_state,
             estimated_state_covariance=estimated_state_covariance,
             predicted_observation=predictions["mean"],
             predicted_observation_covariance=predictions["covariance"],
             observation=current_values,
             observation_model=observation_model,
             observation_noise=self._observation_noise_covariance)
      filtered_state = (filtered_state, filtered_state_covariance, current_times)
      log_prob.set_shape(current_times.get_shape())
      predictions["loss"] = -log_prob
      predictions["log_likelihood"] = log_prob
      if self._configuration.filtering_postprocessor is not None:
        return self._configuration.filtering_postprocessor.process_filtering_step(
            current_times=current_times,
            current_values=current_values,
            predicted_state=state,
            filtered_state=filtered_state,
            outputs=predictions)
      return (filtered_state, predictions)
  
    def _scale_back_predictions(self, predictions):
      """Return a window of predictions to input scale."""
      predictions["mean"] = self._scale_back_data(predictions["mean"])
      predictions["covariance"] = self._scale_back_variance(
          predictions["covariance"])
      return predictions
  
    def _prediction_step(self, current_times, state):
      """Make a prediction based on `state`.
  
      Computes predictions based on the current `state`, checking that it has
      already been updated (in `_imputation_step`) to `current_times`.
  
      Args:
        current_times: A [batch size] Tensor for times to make predictions for.
        state: A tuple of (mean, covariance, previous_times) having shapes
            mean; [batch size x state dimension]
            covariance; [batch size x state dimension x state dimension]
            previous_times; [batch size]
      Returns:
        A tuple of (updated state, predictions):
          updated state: Model state with added transition noise.
          predictions: A dictionary with "mean" and "covariance", having shapes
              "mean": [batch size x num features]
              "covariance: [batch size x num features x num features]
      """
      estimated_state, estimated_state_var, previous_times = state
      advanced_to_current_assert = control_flow_ops.Assert(
          math_ops.reduce_all(math_ops.less_equal(current_times, previous_times)),
          ["Attempted to predict without imputation"])
      with ops.control_dependencies([advanced_to_current_assert]):
        observation_model = self.get_broadcasted_observation_model(current_times)
        predicted_obs, predicted_obs_var = (
            self._kalman_filter.observed_from_state(
                state_mean=estimated_state,
                state_var=estimated_state_var,
                observation_model=observation_model,
                observation_noise=self._observation_noise_covariance))
        predicted_obs_var.set_shape(
            ops.convert_to_tensor(current_times).get_shape()
            .concatenate([self.num_features, self.num_features]))
      predicted_obs.set_shape(current_times.get_shape().concatenate(
          (self.num_features,)))
      predicted_obs_var.set_shape(current_times.get_shape().concatenate(
          (self.num_features, self.num_features)))
      # Not scaled back to input-scale, since this also feeds into the
      # loss. Instead, predictions are scaled back before being returned to the
      # user in _scale_back_predictions.
      predictions = {
          "mean": predicted_obs,
          "covariance": predicted_obs_var}
      state = (estimated_state, estimated_state_var, current_times)
      return (state, predictions)
  
    def _exogenous_noise_decreasing(self, current_times, exogenous_values, state):
      """Update state with exogenous regressors, decreasing uncertainty.
  
      Constructs a mean and covariance based on transformations of
      `exogenous_values`, then performs Bayesian inference on the constructed
      observation. This has the effect of lowering uncertainty.
  
      This update refines or overrides previous inferences, useful for modeling
      exogenous inputs which "set" state, e.g. we dumped boiling water on the
      thermometer so we're pretty sure it's 100 degrees C.
  
      Args:
        current_times: A [batch size] Tensor of times for the exogenous values
            being input.
        exogenous_values: A [batch size x exogenous input dimension] Tensor of
            exogenous values for each part of the batch.
        state: A tuple of (mean, covariance, previous_times) having shapes
            mean; [batch size x state dimension]
            covariance; [batch size x state dimension x state dimension]
            previous_times; [batch size]
      Returns:
        Updated state taking the exogenous regressors into account (with lower
        uncertainty than the input state).
  
      """
      estimated_state, estimated_state_covariance, previous_times = state
      state_transition = ops.convert_to_tensor(
          self.get_state_transition(), dtype=self.dtype)
      state_dimension = tensor_shape.dimension_value(state_transition.shape[0])
      # Learning the observation model would be redundant since we transform
      # `exogenous_values` to the state space via a linear transformation anyway.
      observation_model = linalg_ops.eye(
          state_dimension,
          batch_shape=array_ops.shape(exogenous_values)[:-1],
          dtype=self.dtype)
      with variable_scope.variable_scope("exogenous_noise_decreasing_covariance"):
        observation_noise = math_utils.transform_to_covariance_matrices(
            exogenous_values, state_dimension)
      with variable_scope.variable_scope(
          "exogenous_noise_decreasing_observation"):
        observation = layers.fully_connected(
            exogenous_values, state_dimension, activation_fn=None)
      # Pretend that we are making an observation with an observation model equal
      # to the identity matrix (i.e. a direct observation of the latent state),
      # with learned observation noise.
      posterior_state, posterior_state_var = (
          self._kalman_filter.posterior_from_prior_state(
              prior_state=estimated_state,
              prior_state_var=estimated_state_covariance,
              observation=observation,
              observation_model=observation_model,
              predicted_observations=(
                  estimated_state,
                  # The predicted noise covariance is noise due to current state
                  # uncertainty plus noise learned based on the exogenous
                  # observation (a somewhat trivial call to
                  # self._kalman_filter.observed_from_state has been omitted).
                  observation_noise + estimated_state_covariance),
              observation_noise=observation_noise))
      return (posterior_state, posterior_state_var, previous_times)
  
    def _exogenous_noise_increasing(self, current_times, exogenous_values, state):
      """Update state with exogenous regressors, increasing uncertainty.
  
      Adds to the state mean a linear transformation of `exogenous_values`, and
      increases uncertainty by constructing a covariance matrix based on
      `exogenous_values` and adding it to the state covariance.
  
      This update is useful for modeling changes relative to current state,
      e.g. the furnace turned on so the temperature will be increasing at an
      additional 1 degree per minute with some uncertainty, this uncertainty being
      added to our current uncertainty in the per-minute change in temperature.
  
      Args:
        current_times: A [batch size] Tensor of times for the exogenous values
            being input.
        exogenous_values: A [batch size x exogenous input dimension] Tensor of
            exogenous values for each part of the batch.
        state: A tuple of (mean, covariance, previous_times) having shapes
            mean; [batch size x state dimension]
            covariance; [batch size x state dimension x state dimension]
            previous_times; [batch size]
      Returns:
        Updated state taking the exogenous regressors into account (with higher
        uncertainty than the input state).
  
      """
      start_mean, start_covariance, previous_times = state
      with variable_scope.variable_scope("exogenous_noise_increasing_mean"):
        mean_addition = layers.fully_connected(
            exogenous_values,
            tensor_shape.dimension_value(start_mean.shape[1]), activation_fn=None)
      state_dimension = tensor_shape.dimension_value(start_covariance.shape[1])
      with variable_scope.variable_scope("exogenous_noise_increasing_covariance"):
        covariance_addition = (
            math_utils.transform_to_covariance_matrices(
                exogenous_values, state_dimension))
      return (start_mean + mean_addition,
              start_covariance + covariance_addition,
              previous_times)
  
    def _exogenous_input_step(
        self, current_times, current_exogenous_regressors, state):
      """Update state with exogenous regressors.
  
      Allows both increases and decreases in uncertainty.
  
      Args:
        current_times: A [batch size] Tensor of times for the exogenous values
            being input.
        current_exogenous_regressors: A [batch size x exogenous input dimension]
            Tensor of exogenous values for each part of the batch.
        state: A tuple of (mean, covariance, previous_times) having shapes
            mean; [batch size x state dimension]
            covariance; [batch size x state dimension x state dimension]
            previous_times; [batch size]
      Returns:
        Updated state taking the exogenous regressors into account.
      """
      if self._configuration.exogenous_noise_decreases:
        state = self._exogenous_noise_decreasing(
            current_times, current_exogenous_regressors, state)
      if self._configuration.exogenous_noise_increases:
        state = self._exogenous_noise_increasing(
            current_times, current_exogenous_regressors, state)
      return state
  
    def _loss_additions(self, times, values, mode):
      """Add regularization during training."""
      if mode == estimator_lib.ModeKeys.TRAIN:
        if (self._input_statistics is not None
            and self._configuration.bayesian_prior_weighting):
          normalization = 1. / math_ops.cast(
              self._input_statistics.total_observation_count, self.dtype)
        else:
          # If there is no total observation count recorded, or if we are not
          # doing a Bayesian prior weighting, assumes/pretends that the full
          # dataset size is the window size.
          normalization = 1. / math_ops.cast(
              array_ops.shape(times)[1], self.dtype)
        transition_contribution = ops.convert_to_tensor(
            self._configuration.covariance_prior_fn(
                self.state_transition_noise_covariance),
            dtype=self.dtype)
        if (self._configuration.use_observation_noise
            and self._observation_noise_covariance is not None):
          observation_contribution = ops.convert_to_tensor(
              self._configuration.covariance_prior_fn(
                  self._observation_noise_covariance),
              dtype=self.dtype)
          regularization_sum = transition_contribution + observation_contribution
        else:
          regularization_sum = transition_contribution
        return -normalization * regularization_sum
      else:
        return array_ops.zeros([], dtype=self.dtype)
  
    def _variable_observation_transition_tradeoff_log(self):
      """Define a variable to trade off observation and transition noise."""
      return variable_scope.get_variable(
          name="observation_transition_tradeoff_log_scale",
          initializer=constant_op.constant(
              -self._configuration.transition_covariance_initial_log_scale_bias,
              dtype=self.dtype),
          dtype=self.dtype)
  
    def _define_parameters(self, observation_transition_tradeoff_log=None):
      """Define extra model-specific parameters.
  
      Models should wrap any variables defined here in the model's variable scope.
  
      Args:
        observation_transition_tradeoff_log: An ensemble-global parameter
          controlling the tradeoff between observation noise and transition
          noise. If its value is not None, component transition noise should scale
          with e^-observation_transition_tradeoff_log.
      """
      with variable_scope.variable_scope(self._variable_scope):
        # A scalar which allows the optimizer to quickly shift from observation
        # noise to transition noise (this value is subtracted from log transition
        # noise and added to log observation noise).
        if observation_transition_tradeoff_log is None:
          self._observation_transition_tradeoff_log_scale = (
              self._variable_observation_transition_tradeoff_log())
        else:
          self._observation_transition_tradeoff_log_scale = (
              observation_transition_tradeoff_log)
        self.state_transition_noise_covariance = (
            self.get_state_transition_noise_covariance())
  
    def _set_input_statistics(self, input_statistics=None):
      super(StateSpaceModel, self).initialize_graph(
          input_statistics=input_statistics)
  
    def initialize_graph(self, input_statistics=None):
      """Define variables and ops relevant to the top-level model in an ensemble.
  
      For generic model parameters, _define_parameters() is called recursively on
      all members of an ensemble.
  
      Args:
        input_statistics: A math_utils.InputStatistics object containing input
            statistics. If None, data-independent defaults are used, which may
            result in longer or unstable training.
      """
      self._set_input_statistics(input_statistics=input_statistics)
      self._define_parameters()
      with variable_scope.variable_scope(self._variable_scope):
        self._observation_noise_covariance = ops.convert_to_tensor(
            self.get_observation_noise_covariance(), dtype=self.dtype)
      self._kalman_filter = kalman_filter.KalmanFilter(dtype=self.dtype)
      (self.prior_state_mean,
       self.prior_state_var) = self._make_priors()
  
    def _make_priors(self):
      """Creates and returns model priors."""
      prior_state_covariance = self.get_prior_covariance()
      prior_state_mean = self.get_prior_mean()
      return (prior_state_mean, prior_state_covariance)
  
    def get_prior_covariance(self):
      """Constructs a variable prior covariance with data-based initialization.
  
      Models should wrap any variables defined here in the model's variable scope.
  
      Returns:
        A two-dimensional [state dimension, state dimension] floating point Tensor
        with a (positive definite) prior state covariance matrix.
      """
      with variable_scope.variable_scope(self._variable_scope):
        state_dimension = ops.convert_to_tensor(
            self.get_state_transition()).get_shape().dims[0].value
        if self._configuration.trainable_start_state:
          base_covariance = math_utils.variable_covariance_matrix(
              state_dimension, "prior_state_var",
              dtype=self.dtype)
        else:
          return linalg_ops.eye(state_dimension, dtype=self.dtype)
        if self._input_statistics is not None:
          # Make sure initial latent value uncertainty is at least on the same
          # scale as noise in the data.
          covariance_multiplier = math_ops.reduce_max(
              self._scale_variance(
                  self._input_statistics.series_start_moments.variance))
          return base_covariance * gen_math_ops.maximum(
              covariance_multiplier, 1.0)
        else:
          return base_covariance
  
    def get_prior_mean(self):
      """Constructs a Variable-parameterized prior mean.
  
      Models should wrap any variables defined here in the model's variable scope.
  
      Returns:
        A one-dimensional floating point Tensor with shape [state dimension]
        indicating the prior mean.
      """
      with variable_scope.variable_scope(self._variable_scope):
        state_transition = ops.convert_to_tensor(
            self.get_state_transition(), dtype=self.dtype)
        state_dimension = state_transition.get_shape().dims[0].value
        return variable_scope.get_variable(
            name="prior_state_mean",
            shape=[state_dimension],
            dtype=self.dtype,
            trainable=self._configuration.trainable_start_state)
  
    # TODO(allenl): It would be nice if the generation were done with TensorFlow
    # ops, and if the model parameters were somehow set instead of being passed
    # around in a dictionary. Maybe unconditional generation should be through a
    # special set of initializers?
    def random_model_parameters(self, seed=None):
      if self.num_features != 1:
        raise NotImplementedError("Generation for multivariate state space models"
                                  " is not currently implemented.")
      if seed:
        numpy.random.seed(seed)
      state_dimension, noise_dimension = ops.convert_to_tensor(
          self.get_noise_transform()).get_shape().as_list()
      transition_var = 1.0 / numpy.random.gamma(shape=10., scale=10.,
                                                size=[noise_dimension])
      initial_state = numpy.random.normal(size=[state_dimension])
      params_dict = {}
      if self.prior_state_mean is not None:
        params_dict[self.prior_state_mean] = initial_state
      if self.state_transition_noise_covariance is not None:
        params_dict[self.state_transition_noise_covariance] = numpy.diag(
            transition_var)
      if self.prior_state_var is not None:
        params_dict[self.prior_state_var] = numpy.zeros(
            [state_dimension, state_dimension])
      if self._configuration.use_observation_noise:
        observation_var = 1.0 / numpy.random.gamma(shape=4, scale=4)
        params_dict[self._observation_noise_covariance] = [[observation_var]]
      return params_dict
  
    def generate(self, number_of_series, series_length,
                 model_parameters=None, seed=None, add_observation_noise=None):
      if seed is not None:
        numpy.random.seed(seed)
      if self.num_features != 1:
        raise NotImplementedError("Generation for multivariate state space models"
                                  " is not currently implemented.")
      if add_observation_noise is None:
        add_observation_noise = self._configuration.use_observation_noise
      if model_parameters is None:
        model_parameters = {}
      transitions = ops.convert_to_tensor(
          self.get_state_transition(), dtype=self.dtype).eval(
              feed_dict=model_parameters)
      noise_transform = ops.convert_to_tensor(self.get_noise_transform()).eval(
          feed_dict=model_parameters)
  
      noise_dimension = noise_transform.shape[1]
      get_passed_or_trained_value = model_utils.parameter_switch(model_parameters)
      transition_var = numpy.diag(get_passed_or_trained_value(
          self.state_transition_noise_covariance))
      transition_std = numpy.sqrt(transition_var)
      if add_observation_noise:
        observation_var = get_passed_or_trained_value(
            self._observation_noise_covariance)[0][0]
        observation_std = numpy.sqrt(observation_var)
      initial_state = get_passed_or_trained_value(self.prior_state_mean)
      current_state = numpy.tile(numpy.expand_dims(initial_state, 0),
                                 [number_of_series, 1])
      observations = numpy.zeros([number_of_series, series_length])
      observation_models = self.get_broadcasted_observation_model(
          times=math_ops.range(series_length)).eval(feed_dict=model_parameters)
      for timestep, observation_model in enumerate(observation_models):
        current_state = numpy.dot(current_state, transitions.T)
        current_state += numpy.dot(
            numpy.random.normal(
                loc=numpy.zeros([number_of_series, noise_dimension]),
                scale=numpy.tile(numpy.expand_dims(transition_std, 0),
                                 [number_of_series, 1])),
            noise_transform.T)
        observation_mean = numpy.dot(current_state, observation_model[0].T)
        if add_observation_noise:
          observations[:, timestep] = numpy.random.normal(loc=observation_mean,
                                                          scale=observation_std)
        else:
          observations[:, timestep] = observation_mean
      observations = numpy.expand_dims(observations, -1)
      times = numpy.tile(
          numpy.expand_dims(numpy.arange(observations.shape[1]), 0),
          [observations.shape[0], 1])
      return {TrainEvalFeatures.TIMES: times,
              TrainEvalFeatures.VALUES: observations}
  
    @abc.abstractmethod
    def get_state_transition(self):
      """Specifies the state transition model to use.
  
      Returns:
        A [state dimension x state dimension] Tensor specifying how states
        transition from one timestep to the next.
      """
      pass
  
    @abc.abstractmethod
    def get_noise_transform(self):
      """Specifies the noise transition model to use.
  
      Returns:
        A [state dimension x state noise dimension] Tensor specifying how noise
        (generated with shape [state noise dimension]) affects the model's state.
      """
      pass
  
    @abc.abstractmethod
    def get_observation_model(self, times):
      """Specifies the observation model to use.
  
      Args:
        times: A [batch dimension] int32 Tensor with times for each part of the
            batch, on which the observation model can depend.
      Returns:
        This function, when overridden, has three possible return values:
          - A [state dimension] Tensor with a static, univariate observation
            model.
          - A [self.num_features x state dimension] static, multivariate model.
          - A [batch dimension x self.num_features x state dimension] observation
            model, which may depend on `times`.
        See get_broadcasted_observation_model for details of the broadcasting.
      """
      pass
  
    def get_broadcasted_observation_model(self, times):
      """Broadcast this model's observation model if necessary.
  
      The model can define a univariate observation model which will be broadcast
      over both self.num_features and the batch dimension of `times`.
  
      The model can define a multi-variate observation model which does not depend
      on `times`, and it will be broadcast over the batch dimension of `times`.
  
      Finally, the model can define a multi-variate observation model with a batch
      dimension, which will not be broadcast.
  
      Args:
        times: A [batch dimension] int32 Tensor with times for each part of the
            batch, on which the observation model can depend.
      Returns:
        A [batch dimension x self.num_features x state dimension] Tensor
        specifying the observation model to use for each time in `times` and each
        feature.
      """
      unbroadcasted_model = ops.convert_to_tensor(
          self.get_observation_model(times), dtype=self.dtype)
      unbroadcasted_shape = (unbroadcasted_model.get_shape()
                             .with_rank_at_least(1).with_rank_at_most(3))
      if unbroadcasted_shape.ndims is None:
        # Pass through fully undefined shapes, but make sure they're rank 3 at
        # graph eval time
        assert_op = control_flow_ops.Assert(
            math_ops.equal(array_ops.rank(unbroadcasted_model), 3),
            [array_ops.shape(unbroadcasted_model)])
        with ops.control_dependencies([assert_op]):
          return array_ops.identity(unbroadcasted_model)
      if unbroadcasted_shape.ndims == 1:
        # Unbroadcasted shape [state dimension]
        broadcasted_model = array_ops.tile(
            array_ops.reshape(tensor=unbroadcasted_model, shape=[1, 1, -1]),
            [array_ops.shape(times)[0], self.num_features, 1])
      elif unbroadcasted_shape.ndims == 2:
        # Unbroadcasted shape [num features x state dimension]
        broadcasted_model = array_ops.tile(
            array_ops.expand_dims(unbroadcasted_model, axis=0),
            [array_ops.shape(times)[0], 1, 1])
      elif unbroadcasted_shape.ndims == 3:
        broadcasted_model = unbroadcasted_model
      broadcasted_model.get_shape().assert_has_rank(3)
      return broadcasted_model
  
    def get_state_transition_noise_covariance(
        self, minimum_initial_variance=1e-5):
      state_noise_transform = ops.convert_to_tensor(
          self.get_noise_transform(), dtype=self.dtype)
      state_noise_dimension = state_noise_transform.get_shape().dims[1].value
      if self._input_statistics is not None:
        feature_variance = self._scale_variance(
            self._input_statistics.series_start_moments.variance)
        initial_transition_noise_scale = math_ops.log(
            gen_math_ops.maximum(
                math_ops.reduce_mean(feature_variance) / math_ops.cast(
                    self._input_statistics.total_observation_count, self.dtype),
                minimum_initial_variance))
      else:
        initial_transition_noise_scale = 0.
      # Generally high transition noise is undesirable; we want to set it quite
      # low to start so that we don't need too much training to get to good
      # solutions (i.e. with confident predictions into the future if possible),
      # but not so low that training can't yield a high transition noise if the
      # data demands it.
      initial_transition_noise_scale -= (
          self._observation_transition_tradeoff_log_scale)
      return math_utils.variable_covariance_matrix(
          state_noise_dimension, "state_transition_noise",
          dtype=self.dtype,
          initial_overall_scale_log=initial_transition_noise_scale)
  
    def get_observation_noise_covariance(self, minimum_initial_variance=1e-5):
      if self._configuration.use_observation_noise:
        if self._input_statistics is not None:
          # Get variance across the first few values in each batch for each
          # feature, for an initial observation noise (over-)estimate.
          feature_variance = self._scale_variance(
              self._input_statistics.series_start_moments.variance)
        else:
          feature_variance = None
        if feature_variance is not None:
          feature_variance = gen_math_ops.maximum(feature_variance,
                                                  minimum_initial_variance)
        return math_utils.variable_covariance_matrix(
            size=self.num_features,
            dtype=self.dtype,
            name="observation_noise_covariance",
            initial_diagonal_values=feature_variance,
            initial_overall_scale_log=(
                self._observation_transition_tradeoff_log_scale))
      else:
        return array_ops.zeros(
            shape=[self.num_features, self.num_features],
            name="observation_noise_covariance",
            dtype=self.dtype)
  
    def get_start_state(self):
      """Defines and returns a non-batched prior state and covariance."""
      # TODO(allenl,vitalyk): Add an option for non-Gaussian priors once extended
      # Kalman filtering is implemented (ideally any Distribution object).
      if self._input_statistics is not None:
        start_time = self._input_statistics.start_time
      else:
        start_time = array_ops.zeros([], dtype=dtypes.int64)
      return (self.prior_state_mean,
              self.prior_state_var,
              start_time - 1)
  
    def get_features_for_timesteps(self, timesteps):
      """Get features for a batch of timesteps. Default to no features."""
      return array_ops.zeros([array_ops.shape(timesteps)[0], 0], dtype=self.dtype)
  
  
  class StateSpaceEnsemble(StateSpaceModel):
    """Base class for combinations of state space models."""
  
    def __init__(self, ensemble_members, configuration):
      """Initialize the ensemble by specifying its members.
  
      Args:
        ensemble_members: A list of StateSpaceModel objects which will be included
            in this ensemble.
        configuration: A StateSpaceModelConfiguration object.
      """
      self._ensemble_members = ensemble_members
      super(StateSpaceEnsemble, self).__init__(configuration=configuration)
  
    def _set_input_statistics(self, input_statistics):
      super(StateSpaceEnsemble, self)._set_input_statistics(input_statistics)
      for member in self._ensemble_members:
        member._set_input_statistics(input_statistics)  # pylint: disable=protected-access
  
    def _loss_additions(self, times, values, mode):
      # Allow sub-models to regularize
      return (super(StateSpaceEnsemble, self)._loss_additions(
          times, values, mode) + math_ops.add_n([
              member._loss_additions(times, values, mode)  # pylint: disable=protected-access
              for member in self._ensemble_members
          ]))
  
    def _compute_blocked(self, member_fn, name):
      with variable_scope.variable_scope(self._variable_scope):
        return math_utils.block_diagonal(
            [member_fn(member)
             for member in self._ensemble_members],
            dtype=self.dtype,
            name=name)
  
    def transition_to_powers(self, powers):
      return self._compute_blocked(
          member_fn=lambda member: member.transition_to_powers(powers),
          name="ensemble_transition_to_powers")
  
    def _define_parameters(self, observation_transition_tradeoff_log=None):
      with variable_scope.variable_scope(self._variable_scope):
        if observation_transition_tradeoff_log is None:
          # Define the tradeoff parameter between observation and transition noise
          # once for the whole ensemble, and pass it down to members.
          observation_transition_tradeoff_log = (
              self._variable_observation_transition_tradeoff_log())
        for member in self._ensemble_members:
          member._define_parameters(observation_transition_tradeoff_log=(  # pylint: disable=protected-access
              observation_transition_tradeoff_log))
        super(StateSpaceEnsemble, self)._define_parameters(
            observation_transition_tradeoff_log
            =observation_transition_tradeoff_log)
  
    def random_model_parameters(self, seed=None):
      param_union = {}
      for i, member in enumerate(self._ensemble_members):
        member_params = member.random_model_parameters(
            seed=seed + i if seed else None)
        param_union.update(member_params)
      param_union.update(
          super(StateSpaceEnsemble, self).random_model_parameters(seed=seed))
      return param_union
  
    def get_prior_mean(self):
      return array_ops.concat(
          values=[member.get_prior_mean() for member in self._ensemble_members],
          axis=0,
          name="ensemble_prior_state_mean")
  
    def get_state_transition(self):
      return self._compute_blocked(
          member_fn=
          lambda member: member.get_state_transition(),
          name="ensemble_state_transition")
  
    def get_noise_transform(self):
      return self._compute_blocked(
          member_fn=
          lambda member: member.get_noise_transform(),
          name="ensemble_noise_transform")
  
    def get_observation_model(self, times):
      raise NotImplementedError("No un-broadcasted observation model defined for"
                                " ensembles.")
  
    def get_broadcasted_observation_model(self, times):
      """Computes a combined observation model based on member models.
  
      The effect is that predicted observations from each model are summed.
  
      Args:
        times: A [batch dimension] int32 Tensor with times for each part of the
            batch, on which member observation models can depend.
      Returns:
        A [batch dimension x num features x combined state dimension] Tensor with
        the combined observation model.
      """
      member_observation_models = [
          ops.convert_to_tensor(
              member.get_broadcasted_observation_model(times), dtype=self.dtype)
          for member in self._ensemble_members
      ]
      return array_ops.concat(values=member_observation_models, axis=2)
  
  
  class StateSpaceIndependentEnsemble(StateSpaceEnsemble):
    """Implements ensembles of independent state space models.
  
    Useful for fitting multiple independent state space models together while
    keeping their specifications decoupled. The "ensemble" is simply a state space
    model with the observation models of its members concatenated, and the
    transition matrices and noise transforms stacked in block-diagonal
    matrices. This means that the dimensionality of the ensemble's state is the
    sum of those of its components, which can lead to slow and memory-intensive
    training and inference as the posterior (shape [state dimension x state
    dimension]) gets large.
  
    Each individual model j's state at time t is defined by:
  
    state[t, j] = StateTransition[j] * state[t-1, j]
        + NoiseTransform[j] * StateNoise[t, j]
    StateNoise[t, j] ~ Gaussian(0, StateNoiseCovariance[j])
  
    and the ensemble observation model is:
  
    observation[t] = Sum { ObservationModel[j] * state[t, j] }
        + ObservationNoise[t]
    ObservationNoise[t] ~ Gaussian(0, ObservationNoiseCovariance)
    """
  
    def transition_power_noise_accumulator(self, num_steps):
      return self._compute_blocked(
          member_fn=lambda m: m.transition_power_noise_accumulator(num_steps),
          name="ensemble_power_noise_accumulator")
  
    def get_prior_covariance(self):
      """Construct the ensemble prior covariance based on component models."""
      return self._compute_blocked(
          member_fn=
          lambda member: member.get_prior_covariance(),
          name="ensemble_prior_state_covariance")
  
    def get_state_transition_noise_covariance(self):
      """Construct the ensemble transition noise covariance from components."""
      return self._compute_blocked(
          member_fn=
          lambda member: member.state_transition_noise_covariance,
          name="ensemble_state_transition_noise")
  
  
  # TODO(allenl): It would be nice to have replicated feature models which are
  # identical batched together to reduce the graph size.
  # TODO(allenl): Support for sharing M independent models across N features, with
  # N > M.
  # TODO(allenl): Stack component prior covariances while allowing cross-model
  # correlations to be learned (currently a full covariance prior is learned, but
  # custom component model covariances are not used).
  class StateSpaceCorrelatedFeaturesEnsemble(StateSpaceEnsemble):
    """An correlated ensemble where each model represents a feature.
  
    Unlike `StateSpaceIndependentEnsemble`, a full state transition noise
    covariance matrix is learned for this ensemble; the models are not assumed to
    be independent. Rather than concatenating observation models (i.e. summing the
    contributions of each model to each feature),
    StateSpaceCorrelatedFeaturesEnsemble stacks observation models diagonally,
    meaning that each model corresponds to one feature of the series.
  
    Behaves like (and is) a single state space model where:
  
    StateTransition = Diag(StateTransition[j] for models j)
    ObservationModel = Diag(ObservationModel[j] for models j)
  
    Note that each ObservationModel[j] is a [1 x S_j] matrix (S_j being the state
    dimension of model j), i.e. a univariate model. The combined model is
    multivariate, the number of features of the series being equal to the number
    of component models in the ensemble.
    """
  
    def __init__(self, ensemble_members, configuration):
      """Specify the ensemble's configuration and component models.
  
      Args:
        ensemble_members: A list of `StateSpaceModel` objects, with length equal
          to `configuration.num_features`. Each of these models, which must be
          univariate, corresponds to a single feature of the time series.
        configuration: A StateSpaceModelConfiguration object.
      Raises:
        ValueError: If the length of `ensemble_members` does not equal the number
          of features in the series, or any component is not univariate.
      """
      if len(ensemble_members) != configuration.num_features:
        raise ValueError(
            "The number of members in a StateSpaceCorrelatedFeaturesEnsemble "
            "must equal the number of features in the time series.")
      for member in ensemble_members:
        if member.num_features != 1:
          raise ValueError(
              "StateSpaceCorrelatedFeaturesEnsemble components must be "
              "univariate.")
      super(StateSpaceCorrelatedFeaturesEnsemble, self).__init__(
          ensemble_members=ensemble_members, configuration=configuration)
  
    def transition_power_noise_accumulator(self, num_steps):
      """Use a noise accumulator special case when possible."""
      if len(self._ensemble_members) == 1:
        # If this is a univariate series, we should use the special casing built
        # into the single component model.
        return self._ensemble_members[0].transition_power_noise_accumulator(
            num_steps=num_steps)
      # If we have multiple features, and therefore multiple models, we have
      # introduced correlations which make noise accumulation more
      # complicated. Here we fall back to the general case, since we can't just
      # aggregate member special cases.
      return super(StateSpaceCorrelatedFeaturesEnsemble,
                   self).transition_power_noise_accumulator(num_steps=num_steps)
  
    def get_broadcasted_observation_model(self, times):
      """Stack observation models diagonally."""
      def _member_observation_model(member):
        return ops.convert_to_tensor(
            member.get_broadcasted_observation_model(times), dtype=self.dtype)
      return self._compute_blocked(member_fn=_member_observation_model,
                                   name="feature_ensemble_observation_model")
  