original_docstring: Make a unscented Kalman predict.
  
      Args:
          states (jax.numpy.array): Array of shape (n_obs, n_mixtures, n_states) with
              pre-update states estimates.
          upper_chols (jax.numpy.array): Array of shape (n_obs, n_mixtures, n_states,
              n_states) with the transpose of the lower triangular cholesky factor
              of the pre-update covariance matrix of the state estimates.
          sigma_scaling_factor (float): A scaling factor that controls the spread of the
              sigma points. Bigger means that sigma points are further apart. Depends on
              the sigma_point algorithm chosen.
          sigma_weights (jax.numpy.array): 1d array of length n_sigma with non-negative
              sigma weights.
          transition_functions (tuple): tuple of tuples where the first element is the
              name of the transition function and the second the actual transition
              function. Order is important and corresponds to the latent
              factors in alphabetical order.
          trans_coeffs (tuple): Tuple of 1d jax.numpy.arrays with transition parameters.
          anchoring_scaling_factors (jax.numpy.array): Array of shape (2, n_fac) with
              the scaling factors for anchoring. The first row corresponds to the input
              period, the second to the output period (i.e. input period + 1).
          anchoring_constants (jax.numpy.array): Array of shape (2, n_states) with the
              constants for anchoring. The first row corresponds to the input
              period, the second to the output period (i.e. input period + 1).
  
      Returns:
          jax.numpy.array: Predicted states, same shape as states.
          jax.numpy.array: Predicted upper_chols, same shape as upper_chols.
  
      
docstring: Make a unscented Kalman predict.
docstring_tokens: Make a unscented Kalman predict .
code: def kalman_predict(
      states,
      upper_chols,
      sigma_scaling_factor,
      sigma_weights,
      transition_functions,
      trans_coeffs,
      shock_sds,
      anchoring_scaling_factors,
      anchoring_constants,
  ):
      """Make a unscented Kalman predict.
  
      Args:
          states (jax.numpy.array): Array of shape (n_obs, n_mixtures, n_states) with
              pre-update states estimates.
          upper_chols (jax.numpy.array): Array of shape (n_obs, n_mixtures, n_states,
              n_states) with the transpose of the lower triangular cholesky factor
              of the pre-update covariance matrix of the state estimates.
          sigma_scaling_factor (float): A scaling factor that controls the spread of the
              sigma points. Bigger means that sigma points are further apart. Depends on
              the sigma_point algorithm chosen.
          sigma_weights (jax.numpy.array): 1d array of length n_sigma with non-negative
              sigma weights.
          transition_functions (tuple): tuple of tuples where the first element is the
              name of the transition function and the second the actual transition
              function. Order is important and corresponds to the latent
              factors in alphabetical order.
          trans_coeffs (tuple): Tuple of 1d jax.numpy.arrays with transition parameters.
          anchoring_scaling_factors (jax.numpy.array): Array of shape (2, n_fac) with
              the scaling factors for anchoring. The first row corresponds to the input
              period, the second to the output period (i.e. input period + 1).
          anchoring_constants (jax.numpy.array): Array of shape (2, n_states) with the
              constants for anchoring. The first row corresponds to the input
              period, the second to the output period (i.e. input period + 1).
  
      Returns:
          jax.numpy.array: Predicted states, same shape as states.
          jax.numpy.array: Predicted upper_chols, same shape as upper_chols.
  
      """
      sigma_points = _calculate_sigma_points(states, upper_chols, sigma_scaling_factor)
      transformed = _transform_sigma_points(
          sigma_points,
          transition_functions,
          trans_coeffs,
          anchoring_scaling_factors,
          anchoring_constants,
      )
  
      n_obs, n_mixtures, n_sigma, n_fac = sigma_points.shape
  
      predicted_states = jnp.dot(sigma_weights, transformed)
  
      devs = transformed - predicted_states.reshape(n_obs, n_mixtures, 1, n_fac)
  
      qr_weights = jnp.sqrt(sigma_weights).reshape(n_sigma, 1)
      qr_points = jnp.zeros((n_obs, n_mixtures, n_sigma + n_fac, n_fac))
      qr_points = index_update(qr_points, index[:, :, 0:n_sigma], devs * qr_weights)
      qr_points = index_update(qr_points, index[:, :, n_sigma:], jnp.diag(shock_sds))
      predicted_covs = array_qr_jax(qr_points)[1][:, :, :n_fac]
  
      return predicted_states, predicted_covs
code_tokens: def kalman_predict ( states , upper_chols , sigma_scaling_factor , sigma_weights , transition_functions , trans_coeffs , shock_sds , anchoring_scaling_factors , anchoring_constants , ) : sigma_points = _calculate_sigma_points ( states , upper_chols , sigma_scaling_factor ) transformed = _transform_sigma_points ( sigma_points , transition_functions , trans_coeffs , anchoring_scaling_factors , anchoring_constants , ) n_obs , n_mixtures , n_sigma , n_fac = sigma_points . shape predicted_states = jnp . dot ( sigma_weights , transformed ) devs = transformed - predicted_states . reshape ( n_obs , n_mixtures , 1 , n_fac ) qr_weights = jnp . sqrt ( sigma_weights ) . reshape ( n_sigma , 1 ) qr_points = jnp . zeros ( ( n_obs , n_mixtures , n_sigma + n_fac , n_fac ) ) qr_points = index_update ( qr_points , index [ : , : , 0 : n_sigma ] , devs * qr_weights ) qr_points = index_update ( qr_points , index [ : , : , n_sigma : ] , jnp . diag ( shock_sds ) ) predicted_covs = array_qr_jax ( qr_points ) [ 1 ] [ : , : , : n_fac ] return predicted_states , predicted_covs
short_docstring: Make a unscented Kalman predict.
short_docstring_tokens: Make a unscented Kalman predict .
comment: """Make a unscented Kalman predict.
  
      Args:
          states (jax.numpy.array): Array of shape (n_obs, n_mixtures, n_states) with
              pre-update states estimates.
          upper_chols (jax.numpy.array): Array of shape (n_obs, n_mixtures, n_states,
              n_states) with the transpose of the lower triangular cholesky factor
              of the pre-update covariance matrix of the state estimates.
          sigma_scaling_factor (float): A scaling factor that controls the spread of the
              sigma points. Bigger means that sigma points are further apart. Depends on
              the sigma_point algorithm chosen.
          sigma_weights (jax.numpy.array): 1d array of length n_sigma with non-negative
              sigma weights.
          transition_functions (tuple): tuple of tuples where the first element is the
              name of the transition function and the second the actual transition
              function. Order is important and corresponds to the latent
              factors in alphabetical order.
          trans_coeffs (tuple): Tuple of 1d jax.numpy.arrays with transition parameters.
          anchoring_scaling_factors (jax.numpy.array): Array of shape (2, n_fac) with
              the scaling factors for anchoring. The first row corresponds to the input
              period, the second to the output period (i.e. input period + 1).
          anchoring_constants (jax.numpy.array): Array of shape (2, n_states) with the
              constants for anchoring. The first row corresponds to the input
              period, the second to the output period (i.e. input period + 1).
  
      Returns:
          jax.numpy.array: Predicted states, same shape as states.
          jax.numpy.array: Predicted upper_chols, same shape as upper_chols.
  
      """