class_docstr: Seasonal state space model with effects constrained to sum to zero.
  
  See `SeasonalStateSpaceModel` for background.
  
  #### Mathematical details
  
  The constrained model implements a reparameterization of the
  naive `SeasonalStateSpaceModel`. Instead of directly representing the
  seasonal effects in the latent space, the latent space of the constrained
  model represents the difference between each effect and the mean effect.
  The following discussion assumes familiarity with the mathematical details
  of `SeasonalStateSpaceModel`.
  
  *Reparameterization and constraints*: let the seasonal effects at a given
  timestep be `E = [e_1, ..., e_N]`. The difference between each effect `e_i`
  and the mean effect is `z_i = e_i - sum_i(e_i)/N`. By itself, this
  transformation is not invertible because recovering the absolute effects
  requires that we know the mean as well. To fix this, we'll define
  `z_N = sum_i(e_i)/N` as the mean effect. It's easy to see that this is
  invertible: given the mean effect and the differences of the first `N - 1`
  effects from the mean, it's easy to solve for all `N` effects. Formally,
  we've defined the invertible linear reparameterization `Z = R E`, where
  
  ```
  R = [1 - 1/N, -1/N,    ..., -1/N,      -1/N
       -1/N,    1 - 1/N, ..., -1/N,      -1/N
       ...
       -1/N,    -1/N,    ...,  1/ - 1/N, -1/N
       1/N,     1/N,     ...,  1/N,       1/N]
  E = [e_1, ..., e_(N-1), z_N]
  ```
  
  represents the change of basis from 'effect coordinates' E to
  'residual coordinates' Z. The `Z`s form the latent space of the
  `ConstrainedSeasonalStateSpaceModel`.
  
  To constrain the mean effect `z_N` to zero, we fix the prior to zero,
  `p(z_N) ~ N(0., 0)`, and after the transition at each timestep we project
  `z_N` back to zero. Note that this projection is linear: to set the Nth
  dimension to zero, we simply multiply by the identity matrix with a missing
  element in the bottom right, i.e., `Z_constrained = P Z`,
  where `P = eye(N) - scatter((N-1, N-1), 1)`.
  
  *Model*: concretely, suppose a naive seasonal effect model has initial state
  prior `N(m, S)`, transition matrix `F` and noise covariance
  `Q`, and observation matrix `H`. Then the corresponding constrained seasonal
  effect model has initial state prior `N(P R m, P R S R' P')`,
  transition matrix `P R F R^-1` and noise covariance `F R Q R' F'`, and
  observation matrix `H R^-1`, where the change-of-basis matrix `R` and
  constraint projection matrix `P` are as defined above. This follows
  directly from applying the reparameterization `Z = R E`, and then enforcing
  the zero-sum constraint on the prior and transition noise covariances.
  
  In practice, because the sum of effects `z_N` is constrained to be zero, it
  will never contribute a term to any linear operation on the latent space,
  so we can drop that dimension from the model entirely.
  `ConstrainedSeasonalStateSpaceModel` does this, so that it implements the
  `N - 1` dimension latent space `z_1, ..., z_[N-1]`.
  
  Note that since we constrained the mean effect to be zero, the latent
  `z_i`'s now recover their interpretation as the *actual* effects,
  `z_i = e_i` for `i = `1, ..., N - 1`, even though they were originally
  defined as residuals. The `N`th effect is represented only implicitly, as
  the nonzero mean of the first `N - 1` effects. Although the computational
  represention is not symmetric across all `N` effects, we derived the
  `ConstrainedSeasonalStateSpaceModel` by starting with a symmetric
  representation and imposing only a symmetric constraint (the zero-sum
  constraint), so the probability model remains symmetric over all `N`
  seasonal effects.
  
  #### Examples
  
  A constrained state-space model with day-of-week seasonality on hourly data:
  
  ```python
  day_of_week = ConstrainedSeasonalStateSpaceModel(
    num_timesteps=30,
    num_seasons=7,
    drift_scale=0.1,
    initial_state_prior=tfd.MultivariateNormalDiag(
      scale_diag=tf.ones([7-1], dtype=tf.float32)),
    num_steps_per_season=24)
  ```
  
  A model with basic month-of-year seasonality on daily data, demonstrating
  seasons of varying length:
  
  ```python
  month_of_year = ConstrainedSeasonalStateSpaceModel(
    num_timesteps=2 * 365,  # 2 years
    num_seasons=12,
    drift_scale=0.1,
    initial_state_prior=tfd.MultivariateNormalDiag(
      scale_diag=tf.ones([12-1], dtype=tf.float32)),
    num_steps_per_season=[31, 28, 31, 30, 30, 31, 31, 31, 30, 31, 30, 31],
    initial_step=22)
  ```
  
  Note that we've used `initial_step=22` to denote that the model begins
  on January 23 (steps are zero-indexed). This version works over time periods
  not involving a leap year. A general implementation of month-of-year
  seasonality would require additional logic:
  
  ```python
  num_days_per_month = np.array(
    [[31, 28, 31, 30, 30, 31, 31, 31, 30, 31, 30, 31],
     [31, 29, 31, 30, 30, 31, 31, 31, 30, 31, 30, 31],  # year with leap day
     [31, 28, 31, 30, 30, 31, 31, 31, 30, 31, 30, 31],
     [31, 28, 31, 30, 30, 31, 31, 31, 30, 31, 30, 31]])
  
  month_of_year = ConstrainedSeasonalStateSpaceModel(
    num_timesteps=4 * 365 + 2,  # 8 years with leap days
    num_seasons=12,
    drift_scale=0.1,
    initial_state_prior=tfd.MultivariateNormalDiag(
      scale_diag=tf.ones([12-1], dtype=tf.float32)),
    num_steps_per_season=num_days_per_month,
    initial_step=22)
  ```
class_docstr_tok: ["Seasonal","state","space","model","with","effects","constrained","to","sum","to","zero",".","See","`","SeasonalStateSpaceModel","`","for","background",".","#","#","#","#","Mathematical","details","The","constrained","model","implements","a","reparameterization","of","the","naive","`","SeasonalStateSpaceModel","`",".","Instead","of","directly","representing","the","seasonal","effects","in","the","latent","space",",","the","latent","space","of","the","constrained","model","represents","the","difference","between","each","effect","and","the","mean","effect",".","The","following","discussion","assumes","familiarity","with","the","mathematical","details","of","`","SeasonalStateSpaceModel","`",".","*","Reparameterization","and","constraints","*",":","let","the","seasonal","effects","at","a","given","timestep","be","`","E","=","[","e_1",",","...",",","e_N","]","`",".","The","difference","between","each","effect","`","e_i","`","and","the","mean","effect","is","`","z_i","=","e_i","-","sum_i","(","e_i",")","\/N","`",".","By","itself",",","this","transformation","is","not","invertible","because","recovering","the","absolute","effects","requires","that","we","know","the","mean","as","well",".","To","fix","this",",","we","'ll","define","`","z_N","=","sum_i","(","e_i",")","\/N","`","as","the","mean","effect",".","It","'s","easy","to","see","that","this","is","invertible",":","given","the","mean","effect","and","the","differences","of","the","first","`","N","-","1","`","effects","from","the","mean",",","it","'s","easy","to","solve","for","all","`","N","`","effects",".","Formally",",","we","'ve","defined","the","invertible","linear","reparameterization","`","Z","=","R","E","`",",","where","``","`","R","=","[","1","-","1\/N",",","-1\/N",",","...",",","-1\/N",",","-1\/N","-1\/N",",","1","-","1\/N",",","...",",","-1\/N",",","-1\/N","...","-1\/N",",","-1\/N",",","...",",","1\/","-","1\/N",",","-1\/N","1\/N",",","1\/N",",","...",",","1\/N",",","1\/N","]","E","=","[","e_1",",","...",",","e_","(","N-1",")",",","z_N","]","``","`","represents","the","change","of","basis","from","'effect","coordinates","'","E","to","'residual","coordinates","'","Z",".","The","`","Z","`","s","form","the","latent","space","of","the","`","ConstrainedSeasonalStateSpaceModel","`",".","To","constrain","the","mean","effect","`","z_N","`","to","zero",",","we","fix","the","prior","to","zero",",","`","p","(","z_N",")","~","N","(","0.",",","0",")","`",",","and","after","the","transition","at","each","timestep","we","project","`","z_N","`","back","to","zero",".","Note","that","this","projection","is","linear",":","to","set","the","Nth","dimension","to","zero",",","we","simply","multiply","by","the","identity","matrix","with","a","missing","element","in","the","bottom","right",",","i.e.",",","`","Z_constrained","=","P","Z","`",",","where","`","P","=","eye","(","N",")","-","scatter","(","(","N-1",",","N-1",")",",","1",")","`",".","*","Model","*",":","concretely",",","suppose","a","naive","seasonal","effect","model","has","initial","state","prior","`","N","(","m",",","S",")","`",",","transition","matrix","`","F","`","and","noise","covariance","`","Q","`",",","and","observation","matrix","`","H","`",".","Then","the","corresponding","constrained","seasonal","effect","model","has","initial","state","prior","`","N","(","P","R","m",",","P","R","S","R","'","P","'",")","`",",","transition","matrix","`","P","R","F","R^-1","`","and","noise","covariance","`","F","R","Q","R","'","F","'","`",",","and","observation","matrix","`","H","R^-1","`",",","where","the","change-of-basis","matrix","`","R","`","and","constraint","projection","matrix","`","P","`","are","as","defined","above",".","This","follows","directly","from","applying","the","reparameterization","`","Z","=","R","E","`",",","and","then","enforcing","the","zero-sum","constraint","on","the","prior","and","transition","noise","covariances",".","In","practice",",","because","the","sum","of","effects","`","z_N","`","is","constrained","to","be","zero",",","it","will","never","contribute","a","term","to","any","linear","operation","on","the","latent","space",",","so","we","can","drop","that","dimension","from","the","model","entirely",".","`","ConstrainedSeasonalStateSpaceModel","`","does","this",",","so","that","it","implements","the","`","N","-","1","`","dimension","latent","space","`","z_1",",","...",",","z_","[","N-1","]","`",".","Note","that","since","we","constrained","the","mean","effect","to","be","zero",",","the","latent","`","z_i","`","'s","now","recover","their","interpretation","as","the","*","actual","*","effects",",","`","z_i","=","e_i","`","for","`","i","=","`","1",",","...",",","N","-","1","`",",","even","though","they","were","originally","defined","as","residuals",".","The","`","N","`","th","effect","is","represented","only","implicitly",",","as","the","nonzero","mean","of","the","first","`","N","-","1","`","effects",".","Although","the","computational","represention","is","not","symmetric","across","all","`","N","`","effects",",","we","derived","the","`","ConstrainedSeasonalStateSpaceModel","`","by","starting","with","a","symmetric","representation","and","imposing","only","a","symmetric","constraint","(","the","zero-sum","constraint",")",",","so","the","probability","model","remains","symmetric","over","all","`","N","`","seasonal","effects",".","#","#","#","#","Examples","A","constrained","state-space","model","with","day-of-week","seasonality","on","hourly","data",":","``","`","python","day_of_week","=","ConstrainedSeasonalStateSpaceModel","(","num_timesteps=30",",","num_seasons=7",",","drift_scale=0.1",",","initial_state_prior=tfd.MultivariateNormalDiag","(","scale_diag=tf.ones","(","[","7-1","]",",","dtype=tf.float32",")",")",",","num_steps_per_season=24",")","``","`","A","model","with","basic","month-of-year","seasonality","on","daily","data",",","demonstrating","seasons","of","varying","length",":","``","`","python","month_of_year","=","ConstrainedSeasonalStateSpaceModel","(","num_timesteps=2","*","365",",","#","2","years","num_seasons=12",",","drift_scale=0.1",",","initial_state_prior=tfd.MultivariateNormalDiag","(","scale_diag=tf.ones","(","[","12-1","]",",","dtype=tf.float32",")",")",",","num_steps_per_season=","[","31",",","28",",","31",",","30",",","30",",","31",",","31",",","31",",","30",",","31",",","30",",","31","]",",","initial_step=22",")","``","`","Note","that","we","'ve","used","`","initial_step=22","`","to","denote","that","the","model","begins","on","January","23","(","steps","are","zero-indexed",")",".","This","version","works","over","time","periods","not","involving","a","leap","year",".","A","general","implementation","of","month-of-year","seasonality","would","require","additional","logic",":","``","`","python","num_days_per_month","=","np.array","(","[","[","31",",","28",",","31",",","30",",","30",",","31",",","31",",","31",",","30",",","31",",","30",",","31","]",",","[","31",",","29",",","31",",","30",",","30",",","31",",","31",",","31",",","30",",","31",",","30",",","31","]",",","#","year","with","leap","day","[","31",",","28",",","31",",","30",",","30",",","31",",","31",",","31",",","30",",","31",",","30",",","31","]",",","[","31",",","28",",","31",",","30",",","30",",","31",",","31",",","31",",","30",",","31",",","30",",","31","]","]",")","month_of_year","=","ConstrainedSeasonalStateSpaceModel","(","num_timesteps=4","*","365","+","2",",","#","8","years","with","leap","days","num_seasons=12",",","drift_scale=0.1",",","initial_state_prior=tfd.MultivariateNormalDiag","(","scale_diag=tf.ones","(","[","12-1","]",",","dtype=tf.float32",")",")",",","num_steps_per_season=num_days_per_month",",","initial_step=22",")","``","`"]
docstr: Build a seasonal effect state space model with a zero-sum constraint.
  
  {seasonal_init_args}
docstr_tok: ["Build","a","seasonal","effect","state","space","model","with","a","zero-sum","constraint",".","{","seasonal_init_args","}"]
code: def __init__(
      self,
      num_timesteps,
      num_seasons,
      drift_scale,
      initial_state_prior,
      observation_noise_scale=1e-4,  # Avoid degeneracy.
      num_steps_per_season=1,
      name=None,
      **linear_gaussian_ssm_kwargs,
  ):  # pylint: disable=g-doc-args
      """Build a seasonal effect state space model with a zero-sum constraint.
  
      {seasonal_init_args}
      """
      parameters = dict(locals())
      parameters.update(linear_gaussian_ssm_kwargs)
      del parameters["linear_gaussian_ssm_kwargs"]
      with tf.name_scope(
          name or "ConstrainedSeasonalStateSpaceModel"
      ) as name:
          # The initial state prior determines the dtype of sampled values.
          # Other model parameters must have the same dtype.
          dtype = initial_state_prior.dtype
          drift_scale = tf.convert_to_tensor(
              value=drift_scale, name="drift_scale", dtype=dtype
          )
          observation_noise_scale = tf.convert_to_tensor(
              value=observation_noise_scale,
              name="observation_noise_scale",
              dtype=dtype,
          )
  
          # Coerce `num_steps_per_season` to a canonical form, an array of
          # `num_seasons` integers.
          num_steps_per_season = np.squeeze(
              np.asarray(num_steps_per_season)
          )
          if num_steps_per_season.ndim == 0:  # scalar case
              num_steps_per_season = np.tile(
                  num_steps_per_season, num_seasons
              )
          elif (num_steps_per_season.ndim <= 2) and (  # 1D and 2D case
              num_steps_per_season.shape[-1] != num_seasons
          ):
              raise ValueError(
                  "num_steps_per_season must either be scalar (shape [])"
                  " or have the last dimension equal to [num_seasons] = "
                  "[{}] (saw: shape {})".format(
                      num_seasons, num_steps_per_season.shape
                  )
              )
  
          is_last_day_of_season = build_is_last_day_of_season(
              num_steps_per_season
          )
  
          [
              effects_to_residuals,
              residuals_to_effects,
          ] = build_effects_to_residuals_matrix(
              num_seasons, dtype=dtype
          )
  
          seasonal_transition_matrix = build_seasonal_transition_matrix(
              num_seasons=num_seasons,
              is_last_day_of_season=is_last_day_of_season,
              dtype=dtype,
              basis_change_matrix=effects_to_residuals,
              basis_change_matrix_inv=residuals_to_effects,
          )
  
          seasonal_transition_noise = (
              build_constrained_seasonal_transition_noise(
                  drift_scale, num_seasons, is_last_day_of_season
              )
          )
  
          observation_matrix = tf.concat(
              [
                  tf.ones([1, 1], dtype=dtype),
                  tf.zeros([1, num_seasons - 1], dtype=dtype),
              ],
              axis=-1,
          )
          observation_matrix = tf.matmul(
              observation_matrix, residuals_to_effects
          )
  
          self._drift_scale = drift_scale
          self._observation_noise_scale = observation_noise_scale
          self._num_seasons = num_seasons
          self._num_steps_per_season = num_steps_per_season
  
          super(ConstrainedSeasonalStateSpaceModel, self).__init__(
              num_timesteps=num_timesteps,
              transition_matrix=seasonal_transition_matrix,
              transition_noise=seasonal_transition_noise,
              observation_matrix=observation_matrix,
              observation_noise=mvn_diag.MultivariateNormalDiag(
                  scale_diag=observation_noise_scale[..., tf.newaxis]
              ),
              initial_state_prior=initial_state_prior,
              name=name,
              **linear_gaussian_ssm_kwargs,
          )
          self._parameters = parameters
  
imports: import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import chain
  from tensorflow_probability.python.bijectors import invert
  from tensorflow_probability.python.bijectors import scale
  from tensorflow_probability.python.bijectors import softplus
  from tensorflow_probability.python.bijectors import square
  from tensorflow_probability.python.distributions import inverse_gamma
  from tensorflow_probability.python.distributions import linear_gaussian_ssm
  from tensorflow_probability.python.distributions import lognormal
  from tensorflow_probability.python.distributions import mvn_diag
  from tensorflow_probability.python.distributions import mvn_tril
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.distributions import transformed_distribution
  from tensorflow_probability.python.experimental.sts_gibbs import sample_parameters
  from tensorflow_probability.python.internal import distribution_util
  from tensorflow_probability.python.internal import docstring_util
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.sts.internal import util
  from tensorflow_probability.python.sts.structural_time_series import Parameter
  from tensorflow_probability.python.sts.structural_time_series import StructuralTimeSeries
label_desc: Use image node_id 1 to create a new ConstrainedSeasonalStateSpaceModel object from inherited base classes: linear_gaussian_ssm with example: obj = ConstrainedSeasonalStateSpaceModel(num_timesteps, num_seasons, drift_scale, initial_state_prior, observation_noise_scale, num_steps_per_season, name)