text: # Copyright 2018 The TensorFlow Probability 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.
  # ============================================================================
  """Sums of time-series models."""
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import collections
  
  # Dependency imports
  import tensorflow.compat.v2 as tf
  
  from tensorflow_probability.python import bijectors as tfb
  from tensorflow_probability.python import distributions as tfd
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import prefer_static
  
  from tensorflow_probability.python.sts.internal import util as sts_util
  from tensorflow_probability.python.sts.structural_time_series import Parameter
  from tensorflow_probability.python.sts.structural_time_series import StructuralTimeSeries
  
  tfl = tf.linalg
  
  
  class AdditiveStateSpaceModel(tfd.LinearGaussianStateSpaceModel):
    """A state space model representing a sum of component state space models.
  
    A state space model (SSM) posits a set of latent (unobserved) variables that
    evolve over time with dynamics specified by a probabilistic transition model
    `p(z[t+1] | z[t])`. At each timestep, we observe a value sampled from an
    observation model conditioned on the current state, `p(x[t] | z[t])`. The
    special case where both the transition and observation models are Gaussians
    with mean specified as a linear function of the inputs, is known as a linear
    Gaussian state space model and supports tractable exact probabilistic
    calculations; see `tfp.distributions.LinearGaussianStateSpaceModel` for
    details.
  
    The `AdditiveStateSpaceModel` represents a sum of component state space
    models. Each of the `N` components describes a random process
    generating a distribution on observed time series `x1[t], x2[t], ..., xN[t]`.
    The additive model represents the sum of these
    processes, `y[t] = x1[t] + x2[t] + ... + xN[t] + eps[t]`, where
    `eps[t] ~ N(0, observation_noise_scale)` is an observation noise term.
  
    #### Mathematical Details
  
    The additive model concatenates the latent states of its component models.
    The generative process runs each component's dynamics in its own subspace of
    latent space, and then observes the sum of the observation models from the
    components.
  
    Formally, the transition model is linear Gaussian:
  
    ```
    p(z[t+1] | z[t]) ~ Normal(loc = transition_matrix.matmul(z[t]),
                              cov = transition_cov)
    ```
  
    where each `z[t]` is a latent state vector concatenating the component
    state vectors, `z[t] = [z1[t], z2[t], ..., zN[t]]`, so it has size
    `latent_size = sum([c.latent_size for c in components])`.
  
    The transition matrix is the block-diagonal composition of transition
    matrices from the component processes:
  
    ```
    transition_matrix =
      [[ c0.transition_matrix,  0.,                   ..., 0.                   ],
       [ 0.,                    c1.transition_matrix, ..., 0.                   ],
       [ ...                    ...                   ...                       ],
       [ 0.,                    0.,                   ..., cN.transition_matrix ]]
    ```
  
    and the noise covariance is similarly the block-diagonal composition of
    component noise covariances:
  
    ```
    transition_cov =
      [[ c0.transition_cov, 0.,                ..., 0.                ],
       [ 0.,                c1.transition_cov, ..., 0.                ],
       [ ...                ...                     ...               ],
       [ 0.,                0.,                ..., cN.transition_cov ]]
    ```
  
    The observation model is also linear Gaussian,
  
    ```
    p(y[t] | z[t]) ~ Normal(loc = observation_matrix.matmul(z[t]),
                            stddev = observation_noise_scale)
    ```
  
    This implementation assumes scalar observations, so
    `observation_matrix` has shape `[1, latent_size]`. The additive
    observation matrix simply concatenates the observation matrices from each
    component:
  
    ```
    observation_matrix =
      concat([c0.obs_matrix, c1.obs_matrix, ..., cN.obs_matrix], axis=-1)
    ```
  
    The effect is that each component observation matrix acts on the dimensions
    of latent state corresponding to that component, and the overall expected
    observation is the sum of the expected observations from each component.
  
    If `observation_noise_scale` is not explicitly specified, it is also computed
    by summing the noise variances of the component processes:
  
    ```
    observation_noise_scale = sqrt(sum([
      c.observation_noise_scale**2 for c in components]))
    ```
  
    #### Examples
  
    To construct an additive state space model combining a local linear trend
    and day-of-week seasonality component (note, the `StructuralTimeSeries`
    classes, e.g., `Sum`, provide a higher-level interface for this
    construction, which will likely be preferred by most users):
  
    ```
      num_timesteps = 30
      local_ssm = tfp.sts.LocalLinearTrendStateSpaceModel(
          num_timesteps=num_timesteps,
          level_scale=0.5,
          slope_scale=0.1,
          initial_state_prior=tfd.MultivariateNormalDiag(
              loc=[0., 0.], scale_diag=[1., 1.]))
      day_of_week_ssm = tfp.sts.SeasonalStateSpaceModel(
          num_timesteps=num_timesteps,
          num_seasons=7,
          initial_state_prior=tfd.MultivariateNormalDiag(
              loc=tf.zeros([7]), scale_diag=tf.ones([7])))
      additive_ssm = tfp.sts.AdditiveStateSpaceModel(
          component_ssms=[local_ssm, day_of_week_ssm],
          observation_noise_scale=0.1)
  
      y = additive_ssm.sample()
      print(y.shape)
      # => []
    ```
  
    """
  
    # TODO(b/115656646): test the docstring example using
    # `tfp.sts.SeasonalStateSpaceModel` once that component is checked in.
  
    def __init__(self,
                 component_ssms,
                 constant_offset=0.,
                 observation_noise_scale=None,
                 initial_state_prior=None,
                 initial_step=0,
                 validate_args=False,
                 name=None,
                 **linear_gaussian_ssm_kwargs):
      """Build a state space model representing the sum of component models.
  
      Args:
        component_ssms: Python `list` containing one or more
          `tfd.LinearGaussianStateSpaceModel` instances. The components
          will in general implement different time-series models, with possibly
          different `latent_size`, but they must have the same `dtype`, event
          shape (`num_timesteps` and `observation_size`), and their batch shapes
          must broadcast to a compatible batch shape.
        constant_offset: `float` `Tensor` of shape broadcasting to
          `concat([batch_shape, [num_timesteps]]`) specifying a constant value
          added to the sum of outputs from the component models. This allows the
          components to model the shifted series
          `observed_time_series - constant_offset`.
          Default value: `0.`
        observation_noise_scale: Optional scalar `float` `Tensor` indicating the
          standard deviation of the observation noise. May contain additional
          batch dimensions, which must broadcast with the batch shape of elements
          in `component_ssms`. If `observation_noise_scale` is specified for the
          `AdditiveStateSpaceModel`, the observation noise scales of component
          models are ignored. If `None`, the observation noise scale is derived
          by summing the noise variances of the component models, i.e.,
          `observation_noise_scale = sqrt(sum(
          [ssm.observation_noise_scale**2 for ssm in component_ssms]))`.
        initial_state_prior: Optional instance of `tfd.MultivariateNormal`
          representing a prior distribution on the latent state at time
          `initial_step`. If `None`, defaults to the independent priors from
          component models, i.e.,
          `[component.initial_state_prior for component in component_ssms]`.
          Default value: `None`.
        initial_step: Optional scalar `int` `Tensor` specifying the starting
          timestep.
          Default value: 0.
        validate_args: Python `bool`. Whether to validate input
          with asserts. If `validate_args` is `False`, and the inputs are
          invalid, correct behavior is not guaranteed.
          Default value: `False`.
        name: Python `str` name prefixed to ops created by this class.
          Default value: "AdditiveStateSpaceModel".
        **linear_gaussian_ssm_kwargs: Optional additional keyword arguments to
          to the base `tfd.LinearGaussianStateSpaceModel` constructor.
      Raises:
        ValueError: if components have different `num_timesteps`.
      """
      parameters = dict(locals())
      parameters.update(linear_gaussian_ssm_kwargs)
      del parameters['linear_gaussian_ssm_kwargs']
      with tf.name_scope(name or 'AdditiveStateSpaceModel') as name:
        # Check that all components have the same dtype
        dtype = tf.debugging.assert_same_float_dtype(component_ssms)
  
        # Convert scalar offsets to canonical shape `[..., num_timesteps]`.
        constant_offset = (tf.convert_to_tensor(value=constant_offset,
                                                name='constant_offset',
                                                dtype=dtype) *
                           tf.ones([1], dtype=dtype))
        offset_length = prefer_static.shape(constant_offset)[-1]
        assertions = []
  
        # Construct an initial state prior as a block-diagonal combination
        # of the component state priors.
        if initial_state_prior is None:
          initial_state_prior = sts_util.factored_joint_mvn(
              [ssm.initial_state_prior for ssm in component_ssms])
        dtype = initial_state_prior.dtype
  
        static_num_timesteps = [
            tf.get_static_value(ssm.num_timesteps)
            for ssm in component_ssms
            if tf.get_static_value(ssm.num_timesteps) is not None
        ]
  
        # If any components have a static value for `num_timesteps`, use that
        # value for the additive model. (and check that all other static values
        # match it).
        if static_num_timesteps:
          num_timesteps = static_num_timesteps[0]
          if not all([component_timesteps == num_timesteps
                      for component_timesteps in static_num_timesteps]):
            raise ValueError('Additive model components must all have the same '
                             'number of timesteps '
                             '(saw: {})'.format(static_num_timesteps))
        else:
          num_timesteps = component_ssms[0].num_timesteps
        if validate_args and len(static_num_timesteps) != len(component_ssms):
          assertions += [
              tf.debugging.assert_equal(  # pylint: disable=g-complex-comprehension
                  num_timesteps,
                  ssm.num_timesteps,
                  message='Additive model components must all have '
                  'the same number of timesteps.') for ssm in component_ssms
          ]
  
        # Define the transition and observation models for the additive SSM.
        # See the "mathematical details" section of the class docstring for
        # further information. Note that we define these as callables to
        # handle the fully general case in which some components have time-
        # varying dynamics.
        def transition_matrix_fn(t):
          return tfl.LinearOperatorBlockDiag(
              [ssm.get_transition_matrix_for_timestep(t)
               for ssm in component_ssms])
  
        def transition_noise_fn(t):
          return sts_util.factored_joint_mvn(
              [ssm.get_transition_noise_for_timestep(t)
               for ssm in component_ssms])
  
        # Build the observation matrix, concatenating (broadcast) observation
        # matrices from components. We also take this as an opportunity to enforce
        # any dynamic assertions we may have generated above.
        broadcast_batch_shape = tf.convert_to_tensor(
            value=sts_util.broadcast_batch_shape(
                [ssm.get_observation_matrix_for_timestep(initial_step)
                 for ssm in component_ssms]), dtype=tf.int32)
        broadcast_obs_matrix = tf.ones(
            tf.concat([broadcast_batch_shape, [1, 1]], axis=0), dtype=dtype)
        if assertions:
          with tf.control_dependencies(assertions):
            broadcast_obs_matrix = tf.identity(broadcast_obs_matrix)
  
        def observation_matrix_fn(t):
          return tfl.LinearOperatorFullMatrix(
              tf.concat([ssm.get_observation_matrix_for_timestep(t).to_dense() *
                         broadcast_obs_matrix for ssm in component_ssms],
                        axis=-1))
  
        # Broadcast the constant offset across timesteps.
        offset_at_step = lambda t: (  # pylint: disable=g-long-lambda
            constant_offset if offset_length == 1
            else tf.gather(constant_offset, tf.minimum(t, offset_length - 1),
                           axis=-1)[..., tf.newaxis])
  
        if observation_noise_scale is not None:
          observation_noise_scale = tf.convert_to_tensor(
              value=observation_noise_scale,
              name='observation_noise_scale',
              dtype=dtype)
          def observation_noise_fn(t):
            return tfd.MultivariateNormalDiag(
                loc=(sum([ssm.get_observation_noise_for_timestep(t).mean()
                          for ssm in component_ssms]) + offset_at_step(t)),
                scale_diag=observation_noise_scale[..., tf.newaxis])
        else:
          def observation_noise_fn(t):
            offset = offset_at_step(t)
            return sts_util.sum_mvns(
                [tfd.MultivariateNormalDiag(
                    loc=offset,
                    scale_diag=tf.zeros_like(offset))] +
                [ssm.get_observation_noise_for_timestep(t)
                 for ssm in component_ssms])
  
        super(AdditiveStateSpaceModel, self).__init__(
            num_timesteps=num_timesteps,
            transition_matrix=transition_matrix_fn,
            transition_noise=transition_noise_fn,
            observation_matrix=observation_matrix_fn,
            observation_noise=observation_noise_fn,
            initial_state_prior=initial_state_prior,
            initial_step=initial_step,
            validate_args=validate_args,
            name=name,
            **linear_gaussian_ssm_kwargs)
        self._parameters = parameters
  
  
  class Sum(StructuralTimeSeries):
    """Sum of structural time series components.
  
    This class enables compositional specification of a structural time series
    model from basic components. Given a list of component models, it represents
    an additive model, i.e., a model of time series that may be decomposed into a
    sum of terms corresponding to the component models.
  
    Formally, the additive model represents a random process
    `g[t] = f1[t] + f2[t] + ... + fN[t] + eps[t]`, where the `f`'s are the
    random processes represented by the components, and
    `eps[t] ~ Normal(loc=0, scale=observation_noise_scale)` is an observation
    noise term. See the `AdditiveStateSpaceModel` documentation for mathematical
    details.
  
    This model inherits the parameters (with priors) of its components, and
    adds an `observation_noise_scale` parameter governing the level of noise in
    the observed time series.
  
    #### Examples
  
    To construct a model combining a local linear trend with a day-of-week effect:
  
    ```
      local_trend = tfp.sts.LocalLinearTrend(
          observed_time_series=observed_time_series,
          name='local_trend')
      day_of_week_effect = tfp.sts.Seasonal(
          num_seasons=7,
          observed_time_series=observed_time_series,
          name='day_of_week_effect')
      additive_model = tfp.sts.Sum(
          components=[local_trend, day_of_week_effect],
          observed_time_series=observed_time_series)
  
      print([p.name for p in additive_model.parameters])
      # => `[observation_noise_scale,
      #      local_trend_level_scale,
      #      local_trend_slope_scale,
      #      day_of_week_effect_drift_scale`]
  
      print(local_trend.latent_size,
            seasonal.latent_size,
            additive_model.latent_size)
      # => `2`, `7`, `9`
    ```
  
    """
  
    def __init__(self,
                 components,
                 constant_offset=None,
                 observation_noise_scale_prior=None,
                 observed_time_series=None,
                 name=None):
      """Specify a structural time series model representing a sum of components.
  
      Args:
        components: Python `list` of one or more StructuralTimeSeries instances.
          These must have unique names.
        constant_offset: optional `float` `Tensor` of shape broadcasting to
          `concat([batch_shape, [num_timesteps]]`) specifying a constant value
          added to the sum of outputs from the component models.
          This allows the components to model the shifted series
          `observed_time_series - constant_offset`. If `None`, this is set to the
          mean of the provided `observed_time_series`.
          Default value: `None`.
        observation_noise_scale_prior: optional `tfd.Distribution` instance
          specifying a prior on `observation_noise_scale`. If `None`, a heuristic
          default prior is constructed based on the provided
          `observed_time_series`.
          Default value: `None`.
        observed_time_series: optional `float` `Tensor` of shape
          `batch_shape + [T, 1]` (omitting the trailing unit dimension is also
          supported when `T > 1`), specifying an observed time series. Any `NaN`s
          are interpreted as missing observations; missingness may be also be
          explicitly specified by passing a `tfp.sts.MaskedTimeSeries` instance.
          Any priors not explicitly set will be given default values according to
          the scale of the observed time series (or batch of time series).
          Default value: `None`.
        name: Python `str` name of this model component; used as `name_scope`
          for ops created by this class.
          Default value: 'Sum'.
  
      Raises:
        ValueError: if components do not have unique names.
      """
      init_parameters = dict(locals())
      with tf.name_scope(name or 'Sum') as name:
        if observed_time_series is not None:
          observed_mean, observed_stddev, _ = (
              sts_util.empirical_statistics(observed_time_series))
        else:
          observed_mean, observed_stddev = 0., 1.
  
        if observation_noise_scale_prior is None:
          observation_noise_scale_prior = tfd.LogNormal(
              loc=tf.math.log(.01 * observed_stddev), scale=2.)
  
        dtype = dtype_util.common_dtype([constant_offset,
                                         observation_noise_scale_prior,
                                         observed_mean,
                                         observed_stddev])
  
        # Ensure that offsets have canonical shape `[..., num_timesteps]`.
        if constant_offset is None:
          constant_offset = tf.convert_to_tensor(
              observed_mean, dtype=dtype)[..., tf.newaxis]
        constant_offset *= tf.ones([1], dtype=dtype)
  
        # Check that components have unique names, to ensure that inherited
        # parameters will be assigned unique names.
        component_names = [c.name for c in components]
        if len(component_names) != len(set(component_names)):
          raise ValueError(
              'Components must have unique names: {}'.format(component_names))
        components_by_name = collections.OrderedDict(
            [(c.name, c) for c in components])
  
        # Build parameters list for the combined model, by inheriting parameters
        # from the component models in canonical order.
        parameters = [Parameter('observation_noise_scale',
                                observation_noise_scale_prior,
                                tfb.Chain([
                                    tfb.Scale(scale=observed_stddev),
                                    tfb.Softplus()]))]
        for component in components:
          for parameter in component.parameters:
            parameters.append(Parameter(
                name='{}_{}'.format(component.name, parameter.name),
                prior=parameter.prior, bijector=parameter.bijector))
  
        self._components = components
        self._components_by_name = components_by_name
        self._constant_offset = constant_offset
  
        super(Sum, self).__init__(
            parameters=parameters,
            latent_size=sum(
                [component.latent_size for component in components]),
            init_parameters=init_parameters,
            name=name)
  
    @property
    def components(self):
      """List of component `StructuralTimeSeries` models."""
      return self._components
  
    @property
    def components_by_name(self):
      """OrderedDict mapping component names to components."""
      return self._components_by_name
  
    @property
    def constant_offset(self):
      """Constant value subtracted from observed data."""
      return self._constant_offset
  
    def make_component_state_space_models(self,
                                          num_timesteps,
                                          param_vals,
                                          **linear_gaussian_ssm_kwargs):
      """Build an ordered list of Distribution instances for component models.
  
      Args:
        num_timesteps: Python `int` number of timesteps to model.
        param_vals: a list of `Tensor` parameter values in order corresponding to
          `self.parameters`, or a dict mapping from parameter names to values.
        **linear_gaussian_ssm_kwargs: Optional additional keyword arguments to
          to the base `tfd.LinearGaussianStateSpaceModel` constructors.
  
      Returns:
        component_ssms: a Python list of `LinearGaussianStateSpaceModel`
          Distribution objects, in order corresponding to `self.components`.
      """
  
      with tf.name_scope('make_component_state_space_models'):
  
        # List the model parameters in canonical order
        param_map = self._canonicalize_param_vals_as_map(param_vals)
        param_vals_list = [param_map[p.name] for p in self.parameters]
  
        # Build SSMs for each component model. We process the components in
        # canonical order, extracting the parameters for each component from the
        # (ordered) list of parameters.
        remaining_param_vals = param_vals_list[1:]
        component_ssms = []
        for component in self.components:
          num_parameters = len(component.parameters)
          component_param_vals = remaining_param_vals[:num_parameters]
          remaining_param_vals = remaining_param_vals[num_parameters:]
  
          component_ssms.append(
              component.make_state_space_model(
                  num_timesteps,
                  param_vals=component_param_vals,
                  **linear_gaussian_ssm_kwargs))
  
      return component_ssms
  
    def _make_state_space_model(self,
                                num_timesteps,
                                param_map,
                                initial_state_prior=None,
                                constant_offset=None,
                                **linear_gaussian_ssm_kwargs):
  
      # List the model parameters in canonical order
      param_vals_list = [param_map[p.name] for p in self.parameters]
      observation_noise_scale = param_vals_list[0]
  
      component_ssms = self.make_component_state_space_models(
          num_timesteps=num_timesteps,
          param_vals=param_map,
          **linear_gaussian_ssm_kwargs)
  
      if constant_offset is None:
        constant_offset = self.constant_offset
  
      return AdditiveStateSpaceModel(
          component_ssms=component_ssms,
          constant_offset=constant_offset,
          observation_noise_scale=observation_noise_scale,
          initial_state_prior=initial_state_prior,
          **linear_gaussian_ssm_kwargs)
  