Interpret NaNs as missing values in STS observed_time_series arguments.

davmre · tensorflower-gardener · commit d4609c2dbb4f · 2021-07-08T11:41:49.000-07:00
This supports Pandas Series and DataFrame objects returned by the recently-added `regularize_series`. Note that these objects were *already* supported, insofar as they convert automatically to numpy arrays, but now they can use NaNs to represent missing values. We also now raise an error if the input is a Pandas instance
indexed by a DatetimeIndex with no set frequency.

PiperOrigin-RevId: 383679050
diff --git a/tensorflow_probability/python/sts/components/autoregressive.py b/tensorflow_probability/python/sts/components/autoregressive.py
@@ -315,11 +315,11 @@ def __init__(self,
         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.
+        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). May
-        optionally be an instance of `tfp.sts.MaskedTimeSeries`, which includes
-        a mask `Tensor` to specify timesteps with missing observations.
+        the scale of the observed time series (or batch of time series).
         Default value: `None`.
       name: the name of this model component.
         Default value: 'Autoregressive'.
diff --git a/tensorflow_probability/python/sts/components/dynamic_regression.py b/tensorflow_probability/python/sts/components/dynamic_regression.py
@@ -257,13 +257,13 @@ def __init__(self,
         Must have event shape `[num_features]`. If `None`, a weakly-informative
         Normal(0., 10.) prior is used.
         Default value: `None`.
-      observed_time_series: `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 priors not explicitly set will
-        be given default values according to the scale of the observed time
-        series (or batch of time series). May optionally be an instance of
-        `tfp.sts.MaskedTimeSeries`, which includes a mask `Tensor` to specify
-        timesteps with missing observations.
+      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` for the name of this component.
         Default value: 'DynamicLinearRegression'.
diff --git a/tensorflow_probability/python/sts/components/local_level.py b/tensorflow_probability/python/sts/components/local_level.py
@@ -272,11 +272,11 @@ def __init__(self,
         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.
+        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). May
-        optionally be an instance of `tfp.sts.MaskedTimeSeries`, which includes
-        a mask `Tensor` to specify timesteps with missing observations.
+        the scale of the observed time series (or batch of time series).
         Default value: `None`.
       name: the name of this model component.
         Default value: 'LocalLevel'.
diff --git a/tensorflow_probability/python/sts/components/local_linear_trend.py b/tensorflow_probability/python/sts/components/local_linear_trend.py
@@ -341,11 +341,11 @@ def __init__(self,
         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.
+        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). May
-        optionally be an instance of `tfp.sts.MaskedTimeSeries`, which includes
-        a mask `Tensor` to specify timesteps with missing observations.
+        the scale of the observed time series (or batch of time series).
         Default value: `None`.
       name: the name of this model component.
         Default value: 'LocalLinearTrend'.
diff --git a/tensorflow_probability/python/sts/components/seasonal.py b/tensorflow_probability/python/sts/components/seasonal.py
@@ -803,11 +803,11 @@ def __init__(self,
         Default value: `True`.
       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.
+        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). May
-        optionally be an instance of `tfp.sts.MaskedTimeSeries`, which includes
-        a mask `Tensor` to specify timesteps with missing observations.
+        the scale of the observed time series (or batch of time series).
         Default value: `None`.
       name: the name of this model component.
         Default value: 'Seasonal'.
diff --git a/tensorflow_probability/python/sts/components/semilocal_linear_trend.py b/tensorflow_probability/python/sts/components/semilocal_linear_trend.py
@@ -348,11 +348,11 @@ def __init__(self,
         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.
+        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). May
-        optionally be an instance of `tfp.sts.MaskedTimeSeries`, which includes
-        a mask `Tensor` to specify timesteps with missing observations.
+        the scale of the observed time series (or batch of time series).
         Default value: `None`.
       constrain_ar_coef_stationary: if `True`, perform inference using a
         parameterization that restricts `autoregressive_coef` to the interval
diff --git a/tensorflow_probability/python/sts/components/smooth_seasonal.py b/tensorflow_probability/python/sts/components/smooth_seasonal.py
@@ -389,11 +389,11 @@ def __init__(self,
         is constructed based on the provided `observed_time_series`.
       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.
+        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). May
-        optionally be an instance of `tfp.sts.MaskedTimeSeries`, which includes
-        a mask `Tensor` to specify timesteps with missing observations.
+        the scale of the observed time series (or batch of time series).
         Default value: `None`.
       name: the name of this model component.
         Default value: 'SmoothSeasonal'.
diff --git a/tensorflow_probability/python/sts/components/sum.py b/tensorflow_probability/python/sts/components/sum.py
@@ -408,11 +408,11 @@ def __init__(self,
         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. This is
-        used to set the constant offset, if not provided, and to construct a
-        default heuristic `observation_noise_scale_prior` if not provided. May
-        optionally be an instance of `tfp.sts.MaskedTimeSeries`, which includes
-        a mask `Tensor` to specify timesteps with missing observations.
+        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.
diff --git a/tensorflow_probability/python/sts/decomposition.py b/tensorflow_probability/python/sts/decomposition.py
@@ -124,11 +124,11 @@ def decompose_by_component(model, observed_time_series, parameter_samples):
   Args:
     model: An instance of `tfp.sts.Sum` representing a structural time series
       model.
-    observed_time_series: `float` `Tensor` of shape
-      `batch_shape + [num_timesteps, 1]` (omitting the trailing unit dimension
-      is also supported when `num_timesteps > 1`), specifying an observed time
-      series. May optionally be an instance of `tfp.sts.MaskedTimeSeries`, which
-      includes a mask `Tensor` to specify timesteps with missing observations.
+    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.
     parameter_samples: Python `list` of `Tensors` representing posterior
       samples of model parameters, with shapes `[concat([
       [num_posterior_draws], param.prior.batch_shape,
diff --git a/tensorflow_probability/python/sts/fitting.py b/tensorflow_probability/python/sts/fitting.py
@@ -244,9 +244,9 @@ def fit_with_hmc(model,
     observed_time_series: `float` `Tensor` of shape
       `concat([sample_shape, model.batch_shape, [num_timesteps, 1]])` where
       `sample_shape` corresponds to i.i.d. observations, and the trailing `[1]`
-      dimension may (optionally) be omitted if `num_timesteps > 1`. May
-      optionally be an instance of `tfp.sts.MaskedTimeSeries`, which includes
-      a mask `Tensor` to specify timesteps with missing observations.
+      dimension may (optionally) be omitted if `num_timesteps > 1`. Any `NaN`s
+        are interpreted as missing observations; missingness may be also be
+        explicitly specified by passing a `tfp.sts.MaskedTimeSeries` instance.
     num_results: Integer number of Markov chain draws.
       Default value: `100`.
     num_warmup_steps: Integer number of steps to take before starting to
diff --git a/tensorflow_probability/python/sts/forecast.py b/tensorflow_probability/python/sts/forecast.py
@@ -45,11 +45,11 @@ def one_step_predictive(model, observed_time_series, parameter_samples):
       time-series model. This represents a joint distribution over
       time-series and their parameters with batch shape `[b1, ..., bN]`.
     observed_time_series: `float` `Tensor` of shape
-      `concat([sample_shape, model.batch_shape, [num_timesteps, 1]]) where
+      `concat([sample_shape, model.batch_shape, [num_timesteps, 1]])` where
       `sample_shape` corresponds to i.i.d. observations, and the trailing `[1]`
-      dimension may (optionally) be omitted if `num_timesteps > 1`. May
-      optionally be an instance of `tfp.sts.MaskedTimeSeries` including a
-      mask `Tensor` to encode the locations of missing observations.
+      dimension may (optionally) be omitted if `num_timesteps > 1`. Any `NaN`s
+        are interpreted as missing observations; missingness may be also be
+        explicitly specified by passing a `tfp.sts.MaskedTimeSeries` instance.
     parameter_samples: Python `list` of `Tensors` representing posterior samples
       of model parameters, with shapes `[concat([[num_posterior_draws],
       param.prior.batch_shape, param.prior.event_shape]) for param in
@@ -190,9 +190,9 @@ def forecast(model,
     observed_time_series: `float` `Tensor` of shape
       `concat([sample_shape, model.batch_shape, [num_timesteps, 1]])` where
       `sample_shape` corresponds to i.i.d. observations, and the trailing `[1]`
-      dimension may (optionally) be omitted if `num_timesteps > 1`. May
-      optionally be an instance of `tfp.sts.MaskedTimeSeries` including a
-      mask `Tensor` to encode the locations of missing observations.
+      dimension may (optionally) be omitted if `num_timesteps > 1`. Any `NaN`s
+        are interpreted as missing observations; missingness may be also be
+        explicitly specified by passing a `tfp.sts.MaskedTimeSeries` instance.
     parameter_samples: Python `list` of `Tensors` representing posterior samples
       of model parameters, with shapes `[concat([[num_posterior_draws],
       param.prior.batch_shape, param.prior.event_shape]) for param in
@@ -402,9 +402,9 @@ def impute_missing_values(model,
     observed_time_series: `float` `Tensor` of shape
       `concat([sample_shape, model.batch_shape, [num_timesteps, 1]])` where
       `sample_shape` corresponds to i.i.d. observations, and the trailing `[1]`
-      dimension may (optionally) be omitted if `num_timesteps > 1`. May
-      optionally be an instance of `tfp.sts.MaskedTimeSeries` including a
-      mask `Tensor` to encode the locations of missing observations.
+      dimension may (optionally) be omitted if `num_timesteps > 1`. Any `NaN`s
+        are interpreted as missing observations; missingness may be also be
+        explicitly specified by passing a `tfp.sts.MaskedTimeSeries` instance.
     parameter_samples: Python `list` of `Tensors` representing posterior
       samples of model parameters, with shapes `[concat([
       [num_posterior_draws], param.prior.batch_shape,
diff --git a/tensorflow_probability/python/sts/internal/BUILD b/tensorflow_probability/python/sts/internal/BUILD
@@ -106,8 +106,10 @@ py_test(
     srcs = ["util_test.py"],
     python_version = "PY3",
     srcs_version = "PY3",
+    tags = ["tf1-broken"],
     deps = [
         # numpy dep,
+        # pandas dep,
         # tensorflow dep,
         "//tensorflow_probability",
         "//tensorflow_probability/python/internal:test_util",
diff --git a/tensorflow_probability/python/sts/internal/util.py b/tensorflow_probability/python/sts/internal/util.py
@@ -18,6 +18,8 @@
 from __future__ import print_function
 
 # Dependency imports
+import numpy as np
+
 import tensorflow.compat.v1 as tf1
 import tensorflow.compat.v2 as tf
 
@@ -103,16 +105,13 @@ def do_padding(observed_time_series_tensor):
                                  model_batch_ndims + event_ndims,
                                  lambda: do_padding(observed_time_series),
                                  lambda: observed_time_series)
-
   if is_missing is not None:
     is_missing = ps.cond(ps.rank(is_missing) >
                          model_batch_ndims + event_ndims,
                          lambda: do_padding(is_missing),
                          lambda: is_missing)
-    return missing_values_util.MaskedTimeSeries(observed_time_series,
-                                                is_missing=is_missing)
-
-  return observed_time_series
+  return missing_values_util.MaskedTimeSeries(observed_time_series,
+                                              is_missing=is_missing)
 
 
 def factored_joint_mvn(distributions):
@@ -310,7 +309,9 @@ def canonicalize_observed_time_series_with_mask(
   Args:
     maybe_masked_observed_time_series: a `Tensor`-like object with shape
       `[..., num_timesteps]` or `[..., num_timesteps, 1]`, or a
-      `tfp.sts.MaskedTimeSeries` containing such an object.
+      `tfp.sts.MaskedTimeSeries` containing such an object, or a Pandas
+      Series or DataFrame instance with set frequency
+      (i.e., `.index.freq is not None`).
   Returns:
     masked_time_series: a `tfp.sts.MaskedTimeSeries` namedtuple, in which
       the `observed_time_series` is converted to `Tensor` with canonical shape
@@ -319,18 +320,42 @@ def canonicalize_observed_time_series_with_mask(
   """
 
   with tf.name_scope('canonicalize_observed_time_series_with_mask'):
-    if hasattr(maybe_masked_observed_time_series, 'is_missing'):
+
+    is_missing_is_specified = hasattr(maybe_masked_observed_time_series,
+                                      'is_missing')
+    if is_missing_is_specified:
+      # Input is a MaskedTimeSeries.
       observed_time_series = (
           maybe_masked_observed_time_series.time_series)
       is_missing = maybe_masked_observed_time_series.is_missing
+    elif (hasattr(maybe_masked_observed_time_series, 'index') and
+          hasattr(maybe_masked_observed_time_series, 'to_numpy')):
+      # Input is a Pandas Series or DataFrame.
+      index = maybe_masked_observed_time_series.index
+      if hasattr(index, 'freq') and index.freq is None:
+        raise ValueError('Pandas DataFrame or Series has a DatetimeIndex with '
+                         'no set frequency, but STS requires regularly spaced '
+                         'observations. Consider using '
+                         '`tfp.sts.regularize_series` to infer a frequency and '
+                         'build a regularly spaced series (by marking '
+                         'unobserved steps as missing observations).')
+      # When a DataFrame has multiple columns representing a batch of series,
+      # we want shape `[batch_size, num_steps]` rather than vice versa.
+      observed_time_series = np.squeeze(np.transpose(
+          maybe_masked_observed_time_series.to_numpy()))
     else:
       observed_time_series = maybe_masked_observed_time_series
-      is_missing = None
 
     observed_time_series = tf.convert_to_tensor(value=observed_time_series,
                                                 name='observed_time_series')
     observed_time_series = _maybe_expand_trailing_dim(observed_time_series)
 
+    # Treat `NaN` values as missing.
+    if not is_missing_is_specified:
+      is_missing = tf.math.is_nan(observed_time_series[..., 0])
+    is_missing_static = tf.get_static_value(is_missing)
+    if is_missing_static is not None and not np.any(is_missing_static):
+      is_missing = None
     if is_missing is not None:
       is_missing = tf.convert_to_tensor(
           value=is_missing, name='is_missing', dtype_hint=tf.bool)
diff --git a/tensorflow_probability/python/sts/internal/util_test.py b/tensorflow_probability/python/sts/internal/util_test.py
diff --git a/tensorflow_probability/python/sts/structural_time_series.py b/tensorflow_probability/python/sts/structural_time_series.py
