class_docstr: A model augmented with parameter perturbations for iterated filtering.
class_docstr_tok: ["A","model","augmented","with","parameter","perturbations","for","iterated","filtering","."]
docstr: Builds an iterated filter for parameter estimation in sequential models.
  
  Iterated filtering is a parameter estimation method in which parameters
  are included in an augmented state space, with dynamics that introduce
  parameter perturbations, and a filtering
  algorithm such as particle filtering is run several times with perturbations
  of decreasing size. This class implements the IF2 algorithm of
  [Ionides et al., 2015][1], for which, under appropriate conditions
  (including a uniform prior) the final parameter distribution approaches a
  point mass at the maximum likelihood estimate. If a non-uniform prior is
  provided, the final parameter distribution will (under appropriate
  conditions) approach a point mass at the maximum a posteriori (MAP) value.
  
  This class augments the state space of a sequential model to include
  parameter perturbations, and provides utilities to run particle filtering
  on that augmented model. Alternately, the augmented components may be passed
  directly into a filtering algorithm of the user's choice.
  
  Args:
    parameter_prior: prior `tfd.Distribution` over parameters (may be a joint
      distribution).
    parameterized_initial_state_prior_fn: `callable` with signature
      `initial_state_prior = parameterized_initial_state_prior_fn(parameters)`
      where `parameters` has the form of a sample from `parameter_prior`,
      and `initial_state_prior` is a distribution over the initial state.
    parameterized_transition_fn: `callable` with signature
      `next_state_dist = parameterized_transition_fn(
      step, state, parameters, **kwargs)`.
    parameterized_observation_fn: `callable` with signature
      `observation_dist = parameterized_observation_fn(
      step, state, parameters, **kwargs)`.
    parameterized_initial_state_proposal_fn: optional `callable` with
      signature `initial_state_proposal =
      parameterized_initial_state_proposal_fn(parameters)` where `parameters`
      has the form of a sample from `parameter_prior`, and
      `initial_state_proposal` is a distribution over the initial state.
    parameterized_proposal_fn: optional `callable` with signature
      `next_state_dist = parameterized_transition_fn(
      step, state, parameters, **kwargs)`.
      Default value: `None`.
    parameter_constraining_bijector: optional `tfb.Bijector` instance
      such that `parameter_constraining_bijector.forward(x)` returns valid
      parameters for any real-valued `x` of the same structure and shape
      as `parameters`. If `None`, the default bijector of the provided
      `parameter_prior` will be used.
      Default value: `None`.
    name: `str` name for ops constructed by this object.
      Default value: `iterated_filter`.
  
  #### Example
  
  We'll walk through applying iterated filtering to a toy
  Susceptible-Infected-Recovered (SIR) model, a [compartmental model](
  https://en.wikipedia.org/wiki/Compartmental_models_in_epidemiology#The_SIR_model)
  of infectious disease. Note that the model we use here is extremely
  simplified and is intended as a pedagogical example; it should not be
  interpreted to describe disease spread in the real world.
  
  We begin by specifying a prior distribution over the parameters to be
  inferred, thus defining the structure of the parameter space and the support
  of the parameters (which will imply a default constraining bijector). Here
  we'll use uniform priors over ranges that we expect to contain the
  parameters:
  
  ```python
  parameter_prior = tfd.JointDistributionNamed({
      'infection_rate': tfd.Uniform(low=0., high=3.),
      'recovery_rate': tfd.Uniform(low=0., high=3.),
  })
  ```
  
  The model specification itself is identical to that used by
  `tfp.experimental.mcmc.infer_trajectories`, except that each component
  accepts an additional `parameters` keyword argument. We start by specifying
  a parameterized prior on initial states. In this case, our state
  includes the current number of susceptible and infected individuals
  (the third compartment, recovered individuals, is implicitly defined
  to include the remaining population). We'll also include, as auxiliary
  variables, the daily counts of new infections and new recoveries; these
  will help ensure that people shift consistently across compartments.
  
  ```python
  population_size = 1000
  initial_state_prior_fn = lambda parameters: tfd.JointDistributionNamed({
      'new_infections': tfd.Poisson(parameters['infection_rate']),
      'new_recoveries': tfd.Deterministic(
          tf.broadcast_to(0., tf.shape(parameters['recovery_rate']))),
      'susceptible': (lambda new_infections:
                      tfd.Deterministic(population_size - new_infections)),
      'infected': (lambda new_infections:
                   tfd.Deterministic(new_infections))})
  ```
  
  **Note**: the state prior must have the same batch shape as the
  passed-in parameters; equivalently, it must sample a full state for each
  parameter particle. If any part of the state prior does not depend
  on the parameters, you must manually ensure that it has the appropriate
  batch shape. For example, in the definition of `new_recoveries` above,
  applying `broadcast_to` with the shape of a parameter ensures that
  the batch shape is maintained.
  
  Next, we specify a transition model. This takes the state at the
  previous day, along with parameters, and returns a distribution
  over the state for the current day.
  
  ```python
  def parameterized_infection_dynamics(_, previous_state, parameters):
    new_infections = tfd.Poisson(
        parameters['infection_rate'] * previous_state['infected'] *
        previous_state['susceptible'] / population_size)
    new_recoveries = tfd.Poisson(
        previous_state['infected'] * parameters['recovery_rate'])
    return tfd.JointDistributionNamed({
        'new_infections': new_infections,
        'new_recoveries': new_recoveries,
        'susceptible': lambda new_infections: tfd.Deterministic(
          tf.maximum(0., previous_state['susceptible'] - new_infections)),
        'infected': lambda new_infections, new_recoveries: tfd.Deterministic(
          tf.maximum(0.,
                     (previous_state['infected'] +
                      new_infections - new_recoveries)))})
  ```
  
  Finally, assume that every day we get to observe noisy counts of new
  infections and recoveries.
  
  ```python
  def parameterized_infection_observations(_, state, parameters):
    del parameters  # Not used.
    return tfd.JointDistributionNamed({
        'new_infections': tfd.Poisson(state['new_infections'] + 0.1),
        'new_recoveries': tfd.Poisson(state['new_recoveries'] + 0.1)})
  ```
  
  Combining these components, an `IteratedFilter` augments
  the state space to include parameters that may change over time.
  
  ```python
  iterated_filter = tfp.experimental.sequential.IteratedFilter(
    parameter_prior=parameter_prior,
    parameterized_initial_state_prior_fn=initial_state_prior_fn,
    parameterized_transition_fn=parameterized_infection_dynamics,
    parameterized_observation_fn=parameterized_infection_observations)
  ```
  
  We may then run the filter to estimate parameters from a series
  of observations:
  
  ```python
   # Simulated with `infection_rate=1.2` and `recovery_rate=0.1`.
   observed_values = {
     'new_infections': tf.convert_to_tensor([
        2., 7., 14., 24., 45., 93., 160., 228., 252., 158.,  17.,
        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]),
     'new_recoveries': tf.convert_to_tensor([
        0., 0., 3., 4., 3., 8., 12., 31., 49., 73., 85., 65., 71.,
        58., 42., 65., 36., 31., 32., 27., 31., 20., 19., 19., 14., 27.])
   }
   parameter_particles = iterated_filter.estimate_parameters(
       observations=observed_values,
       num_iterations=20,
       num_particles=4096,
       initial_perturbation_scale=1.0,
       cooling_schedule=(
           tfp.experimental.sequential.geometric_cooling_schedule(
               0.001, k=20)),
       seed=test_util.test_seed())
   print('Mean of parameter particles from final iteration: {}'.format(
     tf.nest.map_structure(lambda x: tf.reduce_mean(x[-1], axis=0),
                           parameter_particles)))
   print('Standard deviation of parameter particles from '
         'final iteration: {}'.format(
         tf.nest.map_structure(lambda x: tf.math.reduce_std(x[-1], axis=0),
                               parameter_particles)))
  ```
  
  For more control, we could alternately choose to run filtering iterations
  on the augmented model manually, using the filter of our choice.
  For example, manually invoking `infer_trajectories` would allow us
  to inspect the parameter and state values at all timesteps, and their
  corresponding log-probabilities:
  
  ```python
  trajectories, lps = tfp.experimental.mcmc.infer_trajectories(
    observations=observations,
    initial_state_prior=iterated_filter.joint_initial_state_prior,
    transition_fn=functools.partial(
        iterated_filter.joint_transition_fn,
        perturbation_scale=perturbation_scale),
    observation_fn=iterated_filter.joint_observation_fn,
    proposal_fn=iterated_filter.joint_proposal_fn,
    initial_state_proposal=iterated_filter.joint_initial_state_proposal(
        initial_unconstrained_parameters),
    num_particles=4096)
  ```
  
  #### References:
  
  [1] Edward L. Ionides, Dao Nguyen, Yves Atchade, Stilian Stoev, and Aaron A.
  King. Inference for dynamic and latent variable models via iterated,
  perturbed Bayes maps. _Proceedings of the National Academy of Sciences_
  112, no. 3: 719-724, 2015.
  https://www.pnas.org/content/pnas/112/3/719.full.pdf
docstr_tok: ["Builds","an","iterated","filter","for","parameter","estimation","in","sequential","models",".","Iterated","filtering","is","a","parameter","estimation","method","in","which","parameters","are","included","in","an","augmented","state","space",",","with","dynamics","that","introduce","parameter","perturbations",",","and","a","filtering","algorithm","such","as","particle","filtering","is","run","several","times","with","perturbations","of","decreasing","size",".","This","class","implements","the","IF2","algorithm","of","[","Ionides","et","al.",",","2015","]","[","1","]",",","for","which",",","under","appropriate","conditions","(","including","a","uniform","prior",")","the","final","parameter","distribution","approaches","a","point","mass","at","the","maximum","likelihood","estimate",".","If","a","non-uniform","prior","is","provided",",","the","final","parameter","distribution","will","(","under","appropriate","conditions",")","approach","a","point","mass","at","the","maximum","a","posteriori","(","MAP",")","value",".","This","class","augments","the","state","space","of","a","sequential","model","to","include","parameter","perturbations",",","and","provides","utilities","to","run","particle","filtering","on","that","augmented","model",".","Alternately",",","the","augmented","components","may","be","passed","directly","into","a","filtering","algorithm","of","the","user","'s","choice",".","Args",":","parameter_prior",":","prior","`","tfd.Distribution","`","over","parameters","(","may","be","a","joint","distribution",")",".","parameterized_initial_state_prior_fn",":","`","callable","`","with","signature","`","initial_state_prior","=","parameterized_initial_state_prior_fn","(","parameters",")","`","where","`","parameters","`","has","the","form","of","a","sample","from","`","parameter_prior","`",",","and","`","initial_state_prior","`","is","a","distribution","over","the","initial","state",".","parameterized_transition_fn",":","`","callable","`","with","signature","`","next_state_dist","=","parameterized_transition_fn","(","step",",","state",",","parameters",",","*","*","kwargs",")","`",".","parameterized_observation_fn",":","`","callable","`","with","signature","`","observation_dist","=","parameterized_observation_fn","(","step",",","state",",","parameters",",","*","*","kwargs",")","`",".","parameterized_initial_state_proposal_fn",":","optional","`","callable","`","with","signature","`","initial_state_proposal","=","parameterized_initial_state_proposal_fn","(","parameters",")","`","where","`","parameters","`","has","the","form","of","a","sample","from","`","parameter_prior","`",",","and","`","initial_state_proposal","`","is","a","distribution","over","the","initial","state",".","parameterized_proposal_fn",":","optional","`","callable","`","with","signature","`","next_state_dist","=","parameterized_transition_fn","(","step",",","state",",","parameters",",","*","*","kwargs",")","`",".","Default","value",":","`","None","`",".","parameter_constraining_bijector",":","optional","`","tfb.Bijector","`","instance","such","that","`","parameter_constraining_bijector.forward","(","x",")","`","returns","valid","parameters","for","any","real-valued","`","x","`","of","the","same","structure","and","shape","as","`","parameters","`",".","If","`","None","`",",","the","default","bijector","of","the","provided","`","parameter_prior","`","will","be","used",".","Default","value",":","`","None","`",".","name",":","`","str","`","name","for","ops","constructed","by","this","object",".","Default","value",":","`","iterated_filter","`",".","#","#","#","#","Example","We","'ll","walk","through","applying","iterated","filtering","to","a","toy","Susceptible-Infected-Recovered","(","SIR",")","model",",","a","[","compartmental","model","]","(","https",":","\/\/en.wikipedia.org\/wiki\/Compartmental_models_in_epidemiology","#","The_SIR_model",")","of","infectious","disease",".","Note","that","the","model","we","use","here","is","extremely","simplified","and","is","intended","as","a","pedagogical","example",";","it","should","not","be","interpreted","to","describe","disease","spread","in","the","real","world",".","We","begin","by","specifying","a","prior","distribution","over","the","parameters","to","be","inferred",",","thus","defining","the","structure","of","the","parameter","space","and","the","support","of","the","parameters","(","which","will","imply","a","default","constraining","bijector",")",".","Here","we","'ll","use","uniform","priors","over","ranges","that","we","expect","to","contain","the","parameters",":","``","`","python","parameter_prior","=","tfd.JointDistributionNamed","(","{","'infection_rate","'",":","tfd.Uniform","(","low=0.",",","high=3",".",")",",","'recovery_rate","'",":","tfd.Uniform","(","low=0.",",","high=3",".",")",",","}",")","``","`","The","model","specification","itself","is","identical","to","that","used","by","`","tfp.experimental.mcmc.infer_trajectories","`",",","except","that","each","component","accepts","an","additional","`","parameters","`","keyword","argument",".","We","start","by","specifying","a","parameterized","prior","on","initial","states",".","In","this","case",",","our","state","includes","the","current","number","of","susceptible","and","infected","individuals","(","the","third","compartment",",","recovered","individuals",",","is","implicitly","defined","to","include","the","remaining","population",")",".","We","'ll","also","include",",","as","auxiliary","variables",",","the","daily","counts","of","new","infections","and","new","recoveries",";","these","will","help","ensure","that","people","shift","consistently","across","compartments",".","``","`","python","population_size","=","1000","initial_state_prior_fn","=","lambda","parameters",":","tfd.JointDistributionNamed","(","{","'new_infections","'",":","tfd.Poisson","(","parameters","[","'infection_rate","'","]",")",",","'new_recoveries","'",":","tfd.Deterministic","(","tf.broadcast_to","(","0.",",","tf.shape","(","parameters","[","'recovery_rate","'","]",")",")",")",",","'susceptible","'",":","(","lambda","new_infections",":","tfd.Deterministic","(","population_size","-","new_infections",")",")",",","'infected","'",":","(","lambda","new_infections",":","tfd.Deterministic","(","new_infections",")",")","}",")","``","`","*","*","Note","*","*",":","the","state","prior","must","have","the","same","batch","shape","as","the","passed-in","parameters",";","equivalently",",","it","must","sample","a","full","state","for","each","parameter","particle",".","If","any","part","of","the","state","prior","does","not","depend","on","the","parameters",",","you","must","manually","ensure","that","it","has","the","appropriate","batch","shape",".","For","example",",","in","the","definition","of","`","new_recoveries","`","above",",","applying","`","broadcast_to","`","with","the","shape","of","a","parameter","ensures","that","the","batch","shape","is","maintained",".","Next",",","we","specify","a","transition","model",".","This","takes","the","state","at","the","previous","day",",","along","with","parameters",",","and","returns","a","distribution","over","the","state","for","the","current","day",".","``","`","python","def","parameterized_infection_dynamics","(","_",",","previous_state",",","parameters",")",":","new_infections","=","tfd.Poisson","(","parameters","[","'infection_rate","'","]","*","previous_state","[","'infected","'","]","*","previous_state","[","'susceptible","'","]","\/","population_size",")","new_recoveries","=","tfd.Poisson","(","previous_state","[","'infected","'","]","*","parameters","[","'recovery_rate","'","]",")","return","tfd.JointDistributionNamed","(","{","'new_infections","'",":","new_infections",",","'new_recoveries","'",":","new_recoveries",",","'susceptible","'",":","lambda","new_infections",":","tfd.Deterministic","(","tf.maximum","(","0.",",","previous_state","[","'susceptible","'","]","-","new_infections",")",")",",","'infected","'",":","lambda","new_infections",",","new_recoveries",":","tfd.Deterministic","(","tf.maximum","(","0.",",","(","previous_state","[","'infected","'","]","+","new_infections","-","new_recoveries",")",")",")","}",")","``","`","Finally",",","assume","that","every","day","we","get","to","observe","noisy","counts","of","new","infections","and","recoveries",".","``","`","python","def","parameterized_infection_observations","(","_",",","state",",","parameters",")",":","del","parameters","#","Not","used",".","return","tfd.JointDistributionNamed","(","{","'new_infections","'",":","tfd.Poisson","(","state","[","'new_infections","'","]","+","0.1",")",",","'new_recoveries","'",":","tfd.Poisson","(","state","[","'new_recoveries","'","]","+","0.1",")","}",")","``","`","Combining","these","components",",","an","`","IteratedFilter","`","augments","the","state","space","to","include","parameters","that","may","change","over","time",".","``","`","python","iterated_filter","=","tfp.experimental.sequential.IteratedFilter","(","parameter_prior=parameter_prior",",","parameterized_initial_state_prior_fn=initial_state_prior_fn",",","parameterized_transition_fn=parameterized_infection_dynamics",",","parameterized_observation_fn=parameterized_infection_observations",")","``","`","We","may","then","run","the","filter","to","estimate","parameters","from","a","series","of","observations",":","``","`","python","#","Simulated","with","`","infection_rate=1.2","`","and","`","recovery_rate=0.1","`",".","observed_values","=","{","'new_infections","'",":","tf.convert_to_tensor","(","[","2.",",","7.",",","14.",",","24.",",","45.",",","93.",",","160.",",","228.",",","252.",",","158.",",","17.",",","0.",",","0.",",","0.",",","0.",",","0.",",","0.",",","0.",",","0.",",","0.",",","0.",",","0.",",","0.",",","0.",",","0.",",","0",".","]",")",",","'new_recoveries","'",":","tf.convert_to_tensor","(","[","0.",",","0.",",","3.",",","4.",",","3.",",","8.",",","12.",",","31.",",","49.",",","73.",",","85.",",","65.",",","71.",",","58.",",","42.",",","65.",",","36.",",","31.",",","32.",",","27.",",","31.",",","20.",",","19.",",","19.",",","14.",",","27",".","]",")","}","parameter_particles","=","iterated_filter.estimate_parameters","(","observations=observed_values",",","num_iterations=20",",","num_particles=4096",",","initial_perturbation_scale=1.0",",","cooling_schedule=","(","tfp.experimental.sequential.geometric_cooling_schedule","(","0.001",",","k=20",")",")",",","seed=test_util.test_seed","(",")",")","print","(","'Mean","of","parameter","particles","from","final","iteration",":","{","}","'.format","(","tf.nest.map_structure","(","lambda","x",":","tf.reduce_mean","(","x","[","-1","]",",","axis=0",")",",","parameter_particles",")",")",")","print","(","'Standard","deviation","of","parameter","particles","from","'","'final","iteration",":","{","}","'.format","(","tf.nest.map_structure","(","lambda","x",":","tf.math.reduce_std","(","x","[","-1","]",",","axis=0",")",",","parameter_particles",")",")",")","``","`","For","more","control",",","we","could","alternately","choose","to","run","filtering","iterations","on","the","augmented","model","manually",",","using","the","filter","of","our","choice",".","For","example",",","manually","invoking","`","infer_trajectories","`","would","allow","us","to","inspect","the","parameter","and","state","values","at","all","timesteps",",","and","their","corresponding","log-probabilities",":","``","`","python","trajectories",",","lps","=","tfp.experimental.mcmc.infer_trajectories","(","observations=observations",",","initial_state_prior=iterated_filter.joint_initial_state_prior",",","transition_fn=functools.partial","(","iterated_filter.joint_transition_fn",",","perturbation_scale=perturbation_scale",")",",","observation_fn=iterated_filter.joint_observation_fn",",","proposal_fn=iterated_filter.joint_proposal_fn",",","initial_state_proposal=iterated_filter.joint_initial_state_proposal","(","initial_unconstrained_parameters",")",",","num_particles=4096",")","``","`","#","#","#","#","References",":","[","1","]","Edward","L.","Ionides",",","Dao","Nguyen",",","Yves","Atchade",",","Stilian","Stoev",",","and","Aaron","A",".","King",".","Inference","for","dynamic","and","latent","variable","models","via","iterated",",","perturbed","Bayes","maps",".","_Proceedings","of","the","National","Academy","of","Sciences_","112",",","no",".","3",":","719-724",",","2015.","https",":","\/\/www.pnas.org\/content\/pnas\/112\/3\/719.full.pdf"]
code: def __init__(
      self,
      parameter_prior,
      parameterized_initial_state_prior_fn,
      parameterized_transition_fn,
      parameterized_observation_fn,
      parameterized_initial_state_proposal_fn=None,
      parameterized_proposal_fn=None,
      parameter_constraining_bijector=None,
      name=None,
  ):
      """Builds an iterated filter for parameter estimation in sequential models.
  
      Iterated filtering is a parameter estimation method in which parameters
      are included in an augmented state space, with dynamics that introduce
      parameter perturbations, and a filtering
      algorithm such as particle filtering is run several times with perturbations
      of decreasing size. This class implements the IF2 algorithm of
      [Ionides et al., 2015][1], for which, under appropriate conditions
      (including a uniform prior) the final parameter distribution approaches a
      point mass at the maximum likelihood estimate. If a non-uniform prior is
      provided, the final parameter distribution will (under appropriate
      conditions) approach a point mass at the maximum a posteriori (MAP) value.
  
      This class augments the state space of a sequential model to include
      parameter perturbations, and provides utilities to run particle filtering
      on that augmented model. Alternately, the augmented components may be passed
      directly into a filtering algorithm of the user's choice.
  
      Args:
        parameter_prior: prior `tfd.Distribution` over parameters (may be a joint
          distribution).
        parameterized_initial_state_prior_fn: `callable` with signature
          `initial_state_prior = parameterized_initial_state_prior_fn(parameters)`
          where `parameters` has the form of a sample from `parameter_prior`,
          and `initial_state_prior` is a distribution over the initial state.
        parameterized_transition_fn: `callable` with signature
          `next_state_dist = parameterized_transition_fn(
          step, state, parameters, **kwargs)`.
        parameterized_observation_fn: `callable` with signature
          `observation_dist = parameterized_observation_fn(
          step, state, parameters, **kwargs)`.
        parameterized_initial_state_proposal_fn: optional `callable` with
          signature `initial_state_proposal =
          parameterized_initial_state_proposal_fn(parameters)` where `parameters`
          has the form of a sample from `parameter_prior`, and
          `initial_state_proposal` is a distribution over the initial state.
        parameterized_proposal_fn: optional `callable` with signature
          `next_state_dist = parameterized_transition_fn(
          step, state, parameters, **kwargs)`.
          Default value: `None`.
        parameter_constraining_bijector: optional `tfb.Bijector` instance
          such that `parameter_constraining_bijector.forward(x)` returns valid
          parameters for any real-valued `x` of the same structure and shape
          as `parameters`. If `None`, the default bijector of the provided
          `parameter_prior` will be used.
          Default value: `None`.
        name: `str` name for ops constructed by this object.
          Default value: `iterated_filter`.
  
      #### Example
  
      We'll walk through applying iterated filtering to a toy
      Susceptible-Infected-Recovered (SIR) model, a [compartmental model](
      https://en.wikipedia.org/wiki/Compartmental_models_in_epidemiology#The_SIR_model)
      of infectious disease. Note that the model we use here is extremely
      simplified and is intended as a pedagogical example; it should not be
      interpreted to describe disease spread in the real world.
  
      We begin by specifying a prior distribution over the parameters to be
      inferred, thus defining the structure of the parameter space and the support
      of the parameters (which will imply a default constraining bijector). Here
      we'll use uniform priors over ranges that we expect to contain the
      parameters:
  
      ```python
      parameter_prior = tfd.JointDistributionNamed({
          'infection_rate': tfd.Uniform(low=0., high=3.),
          'recovery_rate': tfd.Uniform(low=0., high=3.),
      })
      ```
  
      The model specification itself is identical to that used by
      `tfp.experimental.mcmc.infer_trajectories`, except that each component
      accepts an additional `parameters` keyword argument. We start by specifying
      a parameterized prior on initial states. In this case, our state
      includes the current number of susceptible and infected individuals
      (the third compartment, recovered individuals, is implicitly defined
      to include the remaining population). We'll also include, as auxiliary
      variables, the daily counts of new infections and new recoveries; these
      will help ensure that people shift consistently across compartments.
  
      ```python
      population_size = 1000
      initial_state_prior_fn = lambda parameters: tfd.JointDistributionNamed({
          'new_infections': tfd.Poisson(parameters['infection_rate']),
          'new_recoveries': tfd.Deterministic(
              tf.broadcast_to(0., tf.shape(parameters['recovery_rate']))),
          'susceptible': (lambda new_infections:
                          tfd.Deterministic(population_size - new_infections)),
          'infected': (lambda new_infections:
                       tfd.Deterministic(new_infections))})
      ```
  
      **Note**: the state prior must have the same batch shape as the
      passed-in parameters; equivalently, it must sample a full state for each
      parameter particle. If any part of the state prior does not depend
      on the parameters, you must manually ensure that it has the appropriate
      batch shape. For example, in the definition of `new_recoveries` above,
      applying `broadcast_to` with the shape of a parameter ensures that
      the batch shape is maintained.
  
      Next, we specify a transition model. This takes the state at the
      previous day, along with parameters, and returns a distribution
      over the state for the current day.
  
      ```python
      def parameterized_infection_dynamics(_, previous_state, parameters):
        new_infections = tfd.Poisson(
            parameters['infection_rate'] * previous_state['infected'] *
            previous_state['susceptible'] / population_size)
        new_recoveries = tfd.Poisson(
            previous_state['infected'] * parameters['recovery_rate'])
        return tfd.JointDistributionNamed({
            'new_infections': new_infections,
            'new_recoveries': new_recoveries,
            'susceptible': lambda new_infections: tfd.Deterministic(
              tf.maximum(0., previous_state['susceptible'] - new_infections)),
            'infected': lambda new_infections, new_recoveries: tfd.Deterministic(
              tf.maximum(0.,
                         (previous_state['infected'] +
                          new_infections - new_recoveries)))})
      ```
  
      Finally, assume that every day we get to observe noisy counts of new
      infections and recoveries.
  
      ```python
      def parameterized_infection_observations(_, state, parameters):
        del parameters  # Not used.
        return tfd.JointDistributionNamed({
            'new_infections': tfd.Poisson(state['new_infections'] + 0.1),
            'new_recoveries': tfd.Poisson(state['new_recoveries'] + 0.1)})
      ```
  
      Combining these components, an `IteratedFilter` augments
      the state space to include parameters that may change over time.
  
      ```python
      iterated_filter = tfp.experimental.sequential.IteratedFilter(
        parameter_prior=parameter_prior,
        parameterized_initial_state_prior_fn=initial_state_prior_fn,
        parameterized_transition_fn=parameterized_infection_dynamics,
        parameterized_observation_fn=parameterized_infection_observations)
      ```
  
      We may then run the filter to estimate parameters from a series
      of observations:
  
      ```python
       # Simulated with `infection_rate=1.2` and `recovery_rate=0.1`.
       observed_values = {
         'new_infections': tf.convert_to_tensor([
            2., 7., 14., 24., 45., 93., 160., 228., 252., 158.,  17.,
            0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]),
         'new_recoveries': tf.convert_to_tensor([
            0., 0., 3., 4., 3., 8., 12., 31., 49., 73., 85., 65., 71.,
            58., 42., 65., 36., 31., 32., 27., 31., 20., 19., 19., 14., 27.])
       }
       parameter_particles = iterated_filter.estimate_parameters(
           observations=observed_values,
           num_iterations=20,
           num_particles=4096,
           initial_perturbation_scale=1.0,
           cooling_schedule=(
               tfp.experimental.sequential.geometric_cooling_schedule(
                   0.001, k=20)),
           seed=test_util.test_seed())
       print('Mean of parameter particles from final iteration: {}'.format(
         tf.nest.map_structure(lambda x: tf.reduce_mean(x[-1], axis=0),
                               parameter_particles)))
       print('Standard deviation of parameter particles from '
             'final iteration: {}'.format(
             tf.nest.map_structure(lambda x: tf.math.reduce_std(x[-1], axis=0),
                                   parameter_particles)))
      ```
  
      For more control, we could alternately choose to run filtering iterations
      on the augmented model manually, using the filter of our choice.
      For example, manually invoking `infer_trajectories` would allow us
      to inspect the parameter and state values at all timesteps, and their
      corresponding log-probabilities:
  
      ```python
      trajectories, lps = tfp.experimental.mcmc.infer_trajectories(
        observations=observations,
        initial_state_prior=iterated_filter.joint_initial_state_prior,
        transition_fn=functools.partial(
            iterated_filter.joint_transition_fn,
            perturbation_scale=perturbation_scale),
        observation_fn=iterated_filter.joint_observation_fn,
        proposal_fn=iterated_filter.joint_proposal_fn,
        initial_state_proposal=iterated_filter.joint_initial_state_proposal(
            initial_unconstrained_parameters),
        num_particles=4096)
      ```
  
      #### References:
  
      [1] Edward L. Ionides, Dao Nguyen, Yves Atchade, Stilian Stoev, and Aaron A.
      King. Inference for dynamic and latent variable models via iterated,
      perturbed Bayes maps. _Proceedings of the National Academy of Sciences_
      112, no. 3: 719-724, 2015.
      https://www.pnas.org/content/pnas/112/3/719.full.pdf
      """
      name = name or "IteratedFilter"
      with tf.name_scope(name):
          self._parameter_prior = parameter_prior
          self._parameterized_initial_state_prior_fn = (
              parameterized_initial_state_prior_fn
          )
  
          if parameter_constraining_bijector is None:
              parameter_constraining_bijector = (
                  parameter_prior.experimental_default_event_space_bijector()
              )
          self._parameter_constraining_bijector = (
              parameter_constraining_bijector
          )
  
          # Augment the prior to include both parameters and states.
          self._joint_initial_state_prior = (
              joint_prior_on_parameters_and_state(
                  parameter_prior,
                  parameterized_initial_state_prior_fn,
                  parameter_constraining_bijector,
                  prior_is_constrained=True,
              )
          )
  
          # Check that prior samples have a consistent number of particles.
          # TODO(davmre): remove the need for dummy shape dependencies,
          # and this check, by using `JointDistributionNamedAutoBatched` with
          # auto-vectorization enabled in `joint_prior_on_parameters_and_state`.
  
          num_particles_canary = 13
          canary_seed = samplers.zeros_seed()
  
          def _get_shape_1(x):
              if hasattr(x, "state"):
                  x = x.state
              return tf.TensorShape(x.shape[1:2])
  
          prior_static_sample_shapes = tf.nest.map_structure(
              # Sample shape [0, num_particles_canary] particles (size will be zero)
              # then trim off the leading 0 and (possibly) any event shape.
              # We expect shape [num_particles_canary] to remain.
              _get_shape_1,
              self._joint_initial_state_prior.sample(
                  [0, num_particles_canary], seed=canary_seed
              ),
          )
          if not all(
              [
                  tensorshape_util.is_compatible_with(
                      s[:1], [num_particles_canary]
                  )
                  for s in tf.nest.flatten(prior_static_sample_shapes)
              ]
          ):
              raise ValueError(
                  "The specified prior does not generate consistent "
                  "shapes when sampled. Please verify that all parts of "
                  "`initial_state_prior_fn` have batch shape matching "
                  "that of the parameters. This may require creating "
                  '"dummy" dependencies on parameters; for example: '
                  "`tf.broadcast_to(value, tf.shape(parameter))`. (in a "
                  f"test sample with {num_particles_canary} particles, we expected "
                  "all) values to have shape compatible with "
                  f"[{num_particles_canary}, ...]; "
                  f"saw shapes {prior_static_sample_shapes})"
              )
  
          # Augment the transition and observation fns to cover both
          # parameters and states.
          self._joint_transition_fn = (
              augment_transition_fn_with_parameters(
                  parameter_prior,
                  parameterized_transition_fn,
                  parameter_constraining_bijector,
              )
          )
          self._joint_observation_fn = (
              augment_observation_fn_with_parameters(
                  parameterized_observation_fn,
                  parameter_constraining_bijector,
              )
          )
  
          # If given a proposal for the initial state, augment it into a joint
          # proposal over parameters and states.
          joint_initial_state_proposal = None
          if parameterized_initial_state_proposal_fn:
              joint_initial_state_proposal = (
                  joint_prior_on_parameters_and_state(
                      parameter_prior,
                      parameterized_initial_state_proposal_fn,
                      parameter_constraining_bijector,
                  )
              )
          else:
              parameterized_initial_state_proposal_fn = (
                  parameterized_initial_state_prior_fn
              )
          self._joint_initial_state_proposal = (
              joint_initial_state_proposal
          )
          self._parameterized_initial_state_proposal_fn = (
              parameterized_initial_state_proposal_fn
          )
  
          # If given a conditional proposal fn (for non-initial states), augment
          # it to be joint over states and parameters.
          self._joint_proposal_fn = None
          if parameterized_proposal_fn:
              self._joint_proposal_fn = (
                  augment_transition_fn_with_parameters(
                      parameter_prior,
                      parameterized_proposal_fn,
                      parameter_constraining_bijector,
                  )
              )
  
          self._batch_ndims = tf.nest.map_structure(
              ps.rank_from_shape, parameter_prior.batch_shape_tensor()
          )
          self._name = name
  
imports: import collections
  import contextlib
  import functools
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import invert
  from tensorflow_probability.python.distributions import distribution
  from tensorflow_probability.python.distributions import independent
  from tensorflow_probability.python.distributions import joint_distribution_named
  from tensorflow_probability.python.distributions import joint_distribution_sequential
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.distributions import transformed_distribution
  from tensorflow_probability.python.experimental.mcmc.particle_filter import infer_trajectories
  from tensorflow_probability.python.internal import distribution_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import reparameterization
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import tensorshape_util
label_desc: Use image node_id 1 to create a new IteratedFilter object from inherited base classes: object with example: obj = IteratedFilter(parameter_prior, parameterized_initial_state_prior_fn, parameterized_transition_fn, parameterized_observation_fn, parameterized_initial_state_proposal_fn, parameterized_proposal_fn, parameter_constraining_bijector, name)