Support copy and __add__ for StructuralTimeSeries models.

This allows model construction using the sugar:

model = tfp.sts.LocalLevel(observed_time_series=series)
model += tfp.sts.Seasonal(num_seasons=7, observed_time_series=series)
model += tfp.sts.LinearRegression(design_matrix)
# etc.

This change also adds an `init_parameters` property analogous to the `parameters` property of Distributions and Bijectors. (bare `parameters` was already taken for STS components). This was required in order to support `__add__`. After adding it, it was trivial to also support `copy`.

PiperOrigin-RevId: 384768503

Author: davmre

tensorflow_probability/python/sts/components/autoregressive.py

@@ -324,6 +324,7 @@ def __init__(self,
       name: the name of this model component.
         Default value: 'Autoregressive'.
     """
+    init_parameters = dict(locals())
     with tf.name_scope(name or 'Autoregressive') as name:
       masked_time_series = None
       if observed_time_series is not None:
@@ -386,6 +387,7 @@ def __init__(self,
                                    tfb.Softplus()]))
           ],
           latent_size=order,
+          init_parameters=init_parameters,
           name=name)
 
   @property

tensorflow_probability/python/sts/components/dynamic_regression.py

@@ -269,9 +269,8 @@ def __init__(self,
         Default value: 'DynamicLinearRegression'.
 
     """
-
+    init_parameters = dict(locals())
     with tf.name_scope(name or 'DynamicLinearRegression') as name:
-
       dtype = dtype_util.common_dtype(
           [design_matrix, drift_scale_prior, initial_weights_prior])
 
@@ -306,6 +305,7 @@ def __init__(self,
                                    tfb.Softplus()]))
           ],
           latent_size=num_features,
+          init_parameters=init_parameters,
           name=name)
 
   @property

tensorflow_probability/python/sts/components/local_level.py

@@ -281,7 +281,7 @@ def __init__(self,
       name: the name of this model component.
         Default value: 'LocalLevel'.
     """
-
+    init_parameters = dict(locals())
     with tf.name_scope(name or 'LocalLevel') as name:
 
       dtype = dtype_util.common_dtype([level_scale_prior, initial_level_prior])
@@ -319,6 +319,7 @@ def __init__(self,
                                    tfb.Softplus()])),
           ],
           latent_size=1,
+          init_parameters=init_parameters,
           name=name)
 
   @property

tensorflow_probability/python/sts/components/local_linear_trend.py

@@ -350,7 +350,7 @@ def __init__(self,
       name: the name of this model component.
         Default value: 'LocalLinearTrend'.
     """
-
+    init_parameters = dict(locals())
     with tf.name_scope(name or 'LocalLinearTrend') as name:
       _, observed_stddev, observed_initial = (
           sts_util.empirical_statistics(observed_time_series)
@@ -400,6 +400,7 @@ def __init__(self,
               Parameter('slope_scale', slope_scale_prior, scaled_softplus)
           ],
           latent_size=2,
+          init_parameters=init_parameters,
           name=name)
 
   @property

tensorflow_probability/python/sts/components/regression.py

@@ -176,6 +176,7 @@ def __init__(self,
       name: the name of this model component.
         Default value: 'LinearRegression'.
     """
+    init_parameters = dict(locals())
     with tf.name_scope(name or 'LinearRegression') as name:
 
       if not isinstance(design_matrix, tfl.LinearOperator):
@@ -225,6 +226,7 @@ def __init__(self,
                   ),
           ],
           latent_size=0,
+          init_parameters=init_parameters,
           name=name)
 
   @property
@@ -412,8 +414,8 @@ def __init__(self,
       name: the name of this model component.
         Default value: 'SparseLinearRegression'.
     """
+    init_parameters = dict(locals())
     with tf.name_scope(name or 'SparseLinearRegression') as name:
-
       if not isinstance(design_matrix, tfl.LinearOperator):
         design_matrix = tfl.LinearOperatorFullMatrix(
             tf.convert_to_tensor(value=design_matrix, name='design_matrix'),
@@ -468,6 +470,7 @@ def __init__(self,
                         bijector=tfb.Identity())
           ],
           latent_size=0,
+          init_parameters=init_parameters,
           name=name)
 
   @property

tensorflow_probability/python/sts/components/seasonal.py

@@ -812,7 +812,7 @@ def __init__(self,
       name: the name of this model component.
         Default value: 'Seasonal'.
     """
-
+    init_parameters = dict(locals())
     with tf.name_scope(name or 'Seasonal') as name:
 
       _, observed_stddev, observed_initial = (
@@ -875,6 +875,7 @@ def __init__(self,
           parameters,
           latent_size=(num_seasons - 1
                        if self.constrain_mean_effect_to_zero else num_seasons),
+          init_parameters=init_parameters,
           name=name)
 
   @property

tensorflow_probability/python/sts/components/semilocal_linear_trend.py

@@ -368,7 +368,7 @@ def __init__(self,
       name: the name of this model component.
         Default value: 'SemiLocalLinearTrend'.
     """
-
+    init_parameters = dict(locals())
     with tf.name_scope(name or 'SemiLocalLinearTrend') as name:
       if observed_time_series is not None:
         _, observed_stddev, observed_initial = sts_util.empirical_statistics(
@@ -429,6 +429,7 @@ def __init__(self,
                         autoregressive_coef_bijector),
           ],
           latent_size=2,
+          init_parameters=init_parameters,
           name=name)
 
   @property

tensorflow_probability/python/sts/components/smooth_seasonal.py

@@ -399,7 +399,7 @@ def __init__(self,
         Default value: 'SmoothSeasonal'.
 
     """
-
+    init_parameters = dict(locals())
     with tf.name_scope(name or 'SmoothSeasonal') as name:
 
       _, observed_stddev, observed_initial = (
@@ -436,6 +436,7 @@ def __init__(self,
       super(SmoothSeasonal, self).__init__(
           parameters=parameters,
           latent_size=latent_size,
+          init_parameters=init_parameters,
           name=name)
 
   @property

tensorflow_probability/python/sts/components/sum.py

@@ -421,7 +421,7 @@ def __init__(self,
     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, _ = (
@@ -474,6 +474,7 @@ def __init__(self,
           parameters=parameters,
           latent_size=sum(
               [component.latent_size for component in components]),
+          init_parameters=init_parameters,
           name=name)
 
   @property

tensorflow_probability/python/sts/structural_time_series.py

@@ -43,7 +43,8 @@ class StructuralTimeSeries(object):
   structure determined by the child class.
   """
 
-  def __init__(self, parameters, latent_size, name='StructuralTimeSeries'):
+  def __init__(self, parameters, latent_size, init_parameters=None,
+               name='StructuralTimeSeries'):
     """Construct a specification for a structural time series model.
 
     Args:
@@ -54,13 +55,27 @@ def __init__(self, parameters, latent_size, name='StructuralTimeSeries'):
         parameter ordering used by fitting and inference algorithms.
       latent_size: Python `int` specifying the dimensionality of the latent
         state space for this model.
+      init_parameters: Python `dict` of parameters used to instantiate this
+        model component.
       name: Python `str` name for this model component.
     """
 
+    self._init_parameters = init_parameters
     self._parameters = parameters
     self._latent_size = latent_size
     self._name = name
 
+  @property
+  def init_parameters(self):
+    """Parameters used to instantiate this `StructuralTimeSeries`."""
+    if self._init_parameters is None:
+      raise ValueError(
+          'Component has `init_parameters` of `None`. See the built-in '
+          'components (e.g., `LocalLevel`) for examples of how to properly '
+          'store parameters to `__init__`.')
+    return {k: v for k, v in self._init_parameters.items()
+            if not k.startswith('__') and v is not self}
+
   @property
   def parameters(self):
     """List of Parameter(name, prior, bijector) namedtuples for this model."""
@@ -108,6 +123,66 @@ def batch_shape_tensor(self):
           batch_shape, param.prior.batch_shape_tensor())
     return batch_shape
 
+  def __add__(self, other):
+    """Models the sum of the series from the two components."""
+    # Local import to avoid circular dependency.
+    from tensorflow_probability.python.sts.components import sum as sts_sum  # pylint: disable=g-import-not-at-top
+
+    sum_kwargs = {}
+    if isinstance(self, sts_sum.Sum) and isinstance(other, sts_sum.Sum):
+      # Debatably, a Sum + Sum should sum the `constant_offset` parameters, and
+      # should model the sum of the component observation noises (e.g., if
+      # `noise1 ~ N(0, scale=1)` and `noise2 ~ N(0, scale=1)`, then
+      # `noise1 + noise2 ~ N(0, scale=sqrt(2))`). But this would be surprising
+      # in the likely-common case where both sums are modeling the same
+      # observed_time_series. So we instead treat Sum + Sum as a logical
+      # 'merge' operation on the components, requiring that all other parameters
+      # match.
+      _assert_dict_contents_are_equal(
+          self.init_parameters, other.init_parameters,
+          ignore=['components'],
+          message='Cannot add Sum components with different parameter values.')
+      sum_kwargs = self.init_parameters
+    elif isinstance(self, sts_sum.Sum):
+      sum_kwargs = self.init_parameters
+    elif isinstance(other, sts_sum.Sum):
+      sum_kwargs = other.init_parameters
+    else:
+      # If creating a Sum from scratch, try to infer a heuristic noise prior.
+      observed_time_series = self.init_parameters.get(
+          'observed_time_series', None)
+      if observed_time_series is None:
+        observed_time_series = other.init_parameters.get(
+            'observed_time_series', None)
+      if observed_time_series is None:
+        raise ValueError('Could not automatically create a `Sum` component '
+                         'because neither summand was initialized with an '
+                         '`observed_time_series`. You may still instantiate '
+                         'the component manually as `tfp.sts.Sum(...).')
+      sum_kwargs['observed_time_series'] = observed_time_series
+
+    my_components = getattr(self, 'components', [self])
+    other_components = getattr(other, 'components', [other])
+    sum_kwargs['components'] = list(my_components) + list(other_components)
+    return sts_sum.Sum(**sum_kwargs)
+
+  def copy(self, **override_parameters_kwargs):
+    """Creates a deep copy.
+
+    Note: the copy distribution may continue to depend on the original
+    initialization arguments.
+
+    Args:
+      **override_parameters_kwargs: String/value dictionary of initialization
+        arguments to override with new values.
+    Returns:
+      copy: A new instance of `type(self)` initialized from the union
+        of self.init_parameters and override_parameters_kwargs, i.e.,
+        `dict(self.init_parameters, **override_parameters_kwargs)`.
+    """
+    parameters = dict(self.init_parameters, **override_parameters_kwargs)
+    return type(self)(**parameters)
+
   def _canonicalize_param_vals_as_map(self, param_vals):
     """If given an ordered list of parameter values, build a name:value map.
 
@@ -278,3 +353,22 @@ def log_joint_fn(*param_vals, **param_kwargs):
         return param_lp + observation_lp
 
     return log_joint_fn
+
+
+def _strict_equals(a, b):
+  try:
+    return bool(a == b)
+  except ValueError:
+    # `a == b` is an array, Tensor, or similar.
+    return id(a) == id(b)
+
+
+def _assert_dict_contents_are_equal(
+    a, b, message, ignore=(), equals_fn=_strict_equals):
+  combined_keys = set(a.keys()) | set(b.keys())
+  for k in combined_keys - set(ignore):
+    a_val = a.get(k, None)
+    b_val = b.get(k, None)
+    if not _strict_equals(a_val, b_val):
+      raise ValueError(message +
+                       ' `{}`: {} vs {}.'.format(k, a_val, b_val))