<filename>recsim_ng/lib/tensorflow/log_probability.py
# coding=utf-8
# Copyright 2021 The RecSim Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# python3
"""Compute the joint log-probability of a Network given an observation."""

import functools
from typing import Callable, Collection, Mapping, Sequence, Text

import edward2 as ed  # type: ignore
from recsim_ng.core import network as network_lib
from recsim_ng.core import value as value_lib
from recsim_ng.core import variable as variable_lib
from recsim_ng.lib import data
from recsim_ng.lib.tensorflow import field_spec
from recsim_ng.lib.tensorflow import runtime
import tensorflow as tf

FieldSpec = field_spec.FieldSpec
ValueSpec = value_lib.ValueSpec
FieldValue = value_lib.FieldValue
Value = value_lib.Value

Dependency = variable_lib.Dependency
ValueDef = variable_lib.ValueDef
Variable = variable_lib.Variable

Network = network_lib.Network
NetworkValue = network_lib.NetworkValue
NetworkValueTrajectory = runtime.NetworkValueTrajectory

_OBSERVATION_INDEX_FIELD = "__log_probability_observation_index"


def log_probability_from_value_trajectory(
    variables,
    value_trajectory,
    num_steps,
    graph_compile = True):
  """Log probability of a trajectory of network outputs.

  Provides a direct interface to evaluate the outputs of a network simulation,
  for example:
  ```
    variables = story()
    network = network_lib.Network(variables)
    tf_runtime = runtime.TFRuntime(network)
    trajectory = tf_runtime.trajectory(length=5)
    log_p = log_probability_from_value_trajectory(variables, trajectory, 4)
  ```

  Args:
    variables: A collection of `Variable`s defining a dynamic Bayesian network
      (DBN).
    value_trajectory: A trajectory generated from `TFRuntime.trajectory`.
    num_steps: The number of time steps over which to measure the probability.
    graph_compile: Boolean indicating whether the log prob computation is run in
      graph mode.

  Returns:
    A Tensor like that returned from `tfp.distributions.Distribution.log_prob`.
  """
  variables = tuple(variables)
  observations = replay_variables(variables, value_trajectory)
  return log_probability(
      variables, observations, num_steps, graph_compile=graph_compile)


def replay_variables(
    variables,
    value_trajectory):
  """Trajectory replay variables for log probability computation.

  Given a sequence of variables and a trajectory of observed values of these
  variables, this function constructs a sequence of observation variables with
  corresponding to the simulation variables replaying their logged values.

  Args:
    variables: A sequence of `Variable`s defining a dynamic Bayesian network
      (DBN).
    value_trajectory: A trajectory generated from `TFRuntime.trajectory`.

  Returns:
    A sequence of `Variable`.
  """
  observations = []
  for var in variables:
    obs = data.data_variable(
        name=var.name + " obs",
        spec=var.spec,
        data_sequence=data.SlicedValue(value=value_trajectory[var.name]),
        data_index_field=_OBSERVATION_INDEX_FIELD)
    observations.append(obs)
  return observations


def log_prob_variables_from_observation(
    variables,
    observation):
  """Log probability variables for a sequence of observational data.

  Given a sequence of simulation variables and a corresponding sequence
  observation variables (e.g. as generated by `replay_variables`), this function
  generates sequence of log probability variables corresponding to the
  simulation-observation variable pairs. A log probability variable has the
  same fields as its simulation variable with values corresponding to
  log p(observed field value | observations of dependencies), where p is
  the probability mass / density function extracted from the variable's `fn`.
  Deterministic field names are assigned a scalar value of 0.

  Args:
    variables: A sequence of `Variable`s defining a dynamic Bayesian network
      (DBN).
    observation: A sequence of `Variable`s that corresponds one-to-one with
      `variables` and which defines an observation of the DBN.

  Returns:
    A sequence of `Variable`.

  Throws:
    ValueError if the number of simulation variables does not correspond to the
    number of observation variables.
  """
  if len(variables) != len(observation):
    raise ValueError(
        "number of observation variabbles ({}) does not match number of"
        " variables ({})".format(len(observation), len(variables)))

  observation_name = {
      var.name: obs.name for var, obs in zip(variables, observation)
  }
  collision = set(observation_name).intersection(set(observation_name.values()))
  if collision:
    raise ValueError(
        "variables and observation variables share the same names: {}".format(
            collision))

  log_prob_transformation = functools.partial(
      _log_probability_from_observations, observation_name=observation_name)
  log_prob_vars = list(map(log_prob_transformation, variables))
  return log_prob_vars


def log_prob_variables_from_direct_output(
    variables):
  """Log probability variables for outputs at simulation time.

  Given a sequence of simulation variables, this function generates a sequence
  of log probability variables containing the log probabilities of the values
  of those fields of the variables which are stochastically generated. I.e.
  the log probability variable contains log(p(X)) where X is the corresponding
  field of the simulation variable. Deterministic field names are assigned a
  scalar value of 0.

  Args:
    variables: A sequence of `Variable`s defining a dynamic Bayesian network
      (DBN).

  Returns:
    A sequence of `Variable`.

  Throws:
    ValueError if the number of simulation variables does not correspond to the
    number of observation variables.
  """
  return [_log_probability_from_direct_output(var) for var in variables]


def total_log_prob_accumulator_variable(
    log_prob_vars):
  """Accumulated joint log probability variable."""
  log_prob_vars = list(log_prob_vars)  # make sure ordering is stable.
  log_prob_accum = Variable(
      name="total_log_prob_accum", spec=ValueSpec(accum=FieldSpec()))

  def summation(*log_probs):
    reduced_log_probs = []
    for value in log_probs:
      reduced_log_probs.extend(
          [tf.reduce_sum(v) for v in value.as_dict.values()])
    return Value(accum=sum(reduced_log_probs))

  def accumulate(prev, *log_probs):
    return Value(accum=prev.get("accum") + summation(*log_probs).get("accum"))

  log_prob_accum.initial_value = variable_lib.value(
      fn=summation, dependencies=log_prob_vars)
  log_prob_accum.value = variable_lib.value(
      fn=accumulate, dependencies=[log_prob_accum.previous] + log_prob_vars)
  return log_prob_accum


def log_prob_accumulator_variable(log_prob_var):
  """Temporal accumulation of log probability variables.

  Given a log probability variable, outputs temporal per-field accumulator
  of the log probability values of the variable up to the current time instance.

  Args:
    log_prob_var: An instance of `Variable` computing the per-time-step log
      probability of an simulation-observation variable pair (e.g. as generated
      by `log_prob_variables`.

  Returns:
    A `Variable` outputting the per-field sum of all values of the input
    variable up to the current time-step.
  """
  log_prob_accum = Variable(
      name=log_prob_var.name + "_accum", spec=log_prob_var.spec)

  def accumulate_fields(prev, log_prob):
    accumulated_values = {}
    for field_name, value in log_prob.as_dict.items():
      prev_value = prev.get(field_name)
      if not isinstance(prev_value, tf.Tensor):
        prev_value = tf.constant(prev_value)
      if not isinstance(value, tf.Tensor):
        value = tf.constant(value)
      shape_mismatch_msg = (
          f"The log probability of field {field_name} of {log_prob_var.name}"
          f" changes shape during iteration from {prev_value.shape} to"
          f" {value.shape}.")
      tf.debugging.assert_equal(
          tf.shape(prev_value), tf.shape(value), message=shape_mismatch_msg)
      accumulated_values[field_name] = prev_value + value
    return Value(
        **{
            field_name: prev.get(field_name) + value
            for field_name, value in log_prob.as_dict.items()
        })

  log_prob_accum.initial_value = variable_lib.value(
      fn=lambda log_prob: log_prob, dependencies=[log_prob_var])
  log_prob_accum.value = variable_lib.value(
      fn=accumulate_fields,
      dependencies=[log_prob_accum.previous, log_prob_var])
  return log_prob_accum


def log_prob_accumulator_variables(
    log_prob_vars):
  """List version of `log_prob_accumulator_variable`."""
  return [log_prob_accumulator_variable(lpvar) for lpvar in log_prob_vars]


def log_probability(variables,
                    observation,
                    num_steps,
                    graph_compile = True):
  """Returns the joint log probability of an observation given a network.

  Please note that the correctness of the result requires that all of the value
  functions of all the `Variable`s create `ed.RandomVariable` objects in a
  stable order. In other words, if a value function is invoked twice, it will
  create logically corresponding `ed.RandomVariable` objects in the same order.

  Args:
    variables: A sequence of `Variable`s defining a dynamic Bayesian network
      (DBN).
    observation: A sequence of `Variable`s that corresponds one-to-one with
      `variables` and which defines an observation of the DBN.
    num_steps: The number of time steps over which to measure the probability.
    graph_compile: Boolean indicating whether the computation should be run in
      graph mode.

  Returns:
    A Tensor like that returned from `tfp.distributions.Distribution.log_prob`.
  """

  log_prob_vars = log_prob_variables_from_observation(variables, observation)
  accumulator = total_log_prob_accumulator_variable(log_prob_vars)
  tf_runtime = runtime.TFRuntime(
      network=Network(
          variables=list(observation) + list(log_prob_vars) + [accumulator]),
      graph_compile=graph_compile)

  return tf_runtime.execute(num_steps)["total_log_prob_accum"].get("accum")


def _log_probability_from_direct_output(variable):
  """Generates log probability vars from the outputs of a given variable."""
  log_prob_spec = {
      field_name: FieldSpec() for field_name in variable.spec.as_dict.keys()
  }
  log_prob_variable = Variable(
      name=variable.name + "_log_prob", spec=ValueSpec(**log_prob_spec))

  def get_log_probs(value):

    def lp_fn(field):
      if not isinstance(field, ed.RandomVariable):
        return tf.constant(0.0)
      return field.distribution.log_prob(field.value)

    return value.map(lp_fn)

  log_prob_variable.value = variable_lib.value(
      fn=get_log_probs, dependencies=(variable,))
  return log_prob_variable


def _log_probability_from_observations(
    variable, observation_name):
  """Returns a `Variable` that computes log-probability.

  Given a `Variable` with inputs `[a_1, ..., a_k]` and output `b`, this returns
  a `Variable` with inputs `[c_1, ..., c_k, d]` and output `{"log_prob": p}`
  where `c_i` is the observed value of `a_i`, `d` is the observed value of `b`,
  and `p` is the joint log-probability of `d` conditioned on `[c_1, ..., c_k]`.

  Args:
    variable: A `Variable`.
    observation_name: A mapping that, given the `Variable` name of an input or
      output of `variable`, gives the name of the corresponding observation
      `Variable`.

  Returns:
    A `Variable` that yields the log-probability as described above.
  """
  # TODO(ccolby): can we have remove_data_index for variables?
  field_names_no_data = [
      key for key in variable.spec.as_dict.keys() if not key.startswith("__")
  ]
  log_prob_spec = {
      field_name: FieldSpec() for field_name in field_names_no_data
  }
  transformed_variable = Variable(
      name=variable.name + "_log_prob", spec=ValueSpec(**log_prob_spec))

  def transform_value_def(value_def):

    def rewire_dependency_to_observation(dep):
      return Dependency(
          variable_name=observation_name[dep.variable_name],
          on_current_value=dep.on_current_value)

    transformed_dependencies = ([
        Dependency(
            variable_name=observation_name[variable.name],
            on_current_value=True)
    ] + list(map(rewire_dependency_to_observation, value_def.dependencies)))

    def log_prob_fn(observed_output, *observed_inputs):
      unfiltered_log_prob = _transform_fn(value_def.fn)(observed_output,
                                                        *observed_inputs)
      filtered_dict = {
          key: (unfiltered_log_prob.get(key)
                if key in unfiltered_log_prob.as_dict else 0.0)
          for key in field_names_no_data
      }
      return Value(**filtered_dict)

    return ValueDef(fn=log_prob_fn, dependencies=transformed_dependencies)

  if variable.has_explicit_initial_value:
    transformed_variable.initial_value = transform_value_def(
        variable.initial_value)
  if variable.has_explicit_value:
    transformed_variable.value = transform_value_def(variable.value)

  return transformed_variable


def _transform_fn(fn):
  """Returns a transformed value function that computes log-probability.

  For example, suppose `fn(x, y) = {"a": rv1, "b": rv2}` where
  ```
    rv1 = F(x, y)
    rv2 = G(rv1)
  ```
  and suppsoed the observed output is `{"a": ov1, "b": ov2}`. This will return a
  transformed function
  ```
      _transform_fn(fn)({"a": ov1, "b": ov2}, x, y) = {"log_prob": p}
  ```
  that has four random variables, two for each of `rv1` and `rv2`:
  ```
    rv1_a = F(x, y)
    rv1_b = ov1
    rv2_a = G(rv1_b)
    rv2_b = ov2
  ```
  The output, `p`, is the sum of:
  ```
    rv1_a.distribution.log_prob(ov1)
    rv2_a.distribution.log_prob(ov2)
  ```
  Note that in this example, `rv2_b` is created but not used.

  Args:
    fn: A `Variable` value function (appearing in `ValueDef.fn`).

  Returns:
    A transformed function as described above.

  Raises:
    RuntimeError: If `fn` creates any `ed.RandomVariable` objects that are not
      exposed as fields of its output `Value`.
  """

  def transformed_fn(observed_output, *observed_inputs):
    cleaned_inputs = list(
        map(
            functools.partial(
                data.remove_data_index,
                data_index_field=_OBSERVATION_INDEX_FIELD), observed_inputs))
    observed_output_values_by_rv_order = _observed_output_values_by_rv_order(
        fn, observed_output, *cleaned_inputs)
    num_rvs = len(observed_output_values_by_rv_order)

    rv_index = 0
    log_probs = {}

    def log_prob_tracer(rv_constructor, *args, **kwargs):
      nonlocal rv_index
      if rv_index >= num_rvs:
        raise RuntimeError(
            "function created {} random variables the first time it was called,"
            " but created more the second time".format(num_rvs))
      field_name, observed_value = observed_output_values_by_rv_order[rv_index]
      rv_index += 1
      rv = rv_constructor(*args, **kwargs)
      logp = rv.distribution.log_prob(observed_value)
      log_probs[field_name] = logp
      kwargs["value"] = observed_value
      # be nice to higher tracers
      return ed.traceable(rv_constructor)(*args, **kwargs)

    with ed.trace(log_prob_tracer):
      _ = fn(*cleaned_inputs)

    return Value(**log_probs)

  return transformed_fn


def _observed_output_values_by_rv_order(
    fn, observed_output,
    *observed_inputs):
  """Returns a sequence of observed values in order of RandomVariable creation.

  For example, suppose `fn(x, y) = {"a": rv1, "b": rv2}` where
  ```
    rv1 = F(x, y)
    rv2 = G(rv1)
  ```
  Then,
  ```
    _observed_values_by_rv_order(fn, {"a": ov1, "b": ov2}, x, y) = [ov1, ov2]
  ```
  Here, `fn` constructs two `ed.RandomVariable` objects, first `rv1` and then
  `rv2`. Because `rv1` was constructed first, `ov1` appears first in the
  returned sequence.

  Args:
    fn: A `Variable` value function; see `ValueDef.fn`.
    observed_output: A `Value` containing field names matching the value
      returned by `fn`.
    *observed_inputs: The input arguments of `fn`.

  Returns:
    A sequence of (field_name, `observed_output`) pairs of field values
    corresponding to the `ed.RandomVariable` field values output by `fn`,
    arranged in the order in which `fn` constructs those `ed.RandomVariable`
    objects together with the names of the fields they get assigned to.

  Raises:
    RuntimeError: If `fn` creates any `ed.RandomVariable` objects that are not
      exposed as fields of its output `Value`.
  """

  rvs_in_order_of_construction = []

  def index_random_variables(rv_constructor, *args, **kwargs):
    rv = rv_constructor(*args, **kwargs)
    rvs_in_order_of_construction.append(rv)
    return rv

  with ed.trace(index_random_variables):
    temporary_output = fn(*observed_inputs)

  rv_to_output_value = {
      value: (field_name, observed_output.get(field_name))
      for field_name, value in temporary_output.as_dict.items()
      if isinstance(value, ed.RandomVariable)
  }
  unobserved = [
      rv for rv in rvs_in_order_of_construction if rv not in rv_to_output_value
  ]
  if unobserved:
    raise RuntimeError(
        "unobserved random variables; log-probability cannot be computed: {}"
        .format(unobserved))

  return [rv_to_output_value[rv] for rv in rvs_in_order_of_construction]

<filename>recsim_ng/lib/tensorflow/log_probability.py
# coding=utf-8
# Copyright 2021 The RecSim Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# python3
"""Compute the joint log-probability of a Network given an observation."""

import functools
from typing import Callable, Collection, Mapping, Sequence, Text

import edward2 as ed  # type: ignore
from recsim_ng.core import network as network_lib
from recsim_ng.core import value as value_lib
from recsim_ng.core import variable as variable_lib
from recsim_ng.lib import data
from recsim_ng.lib.tensorflow import field_spec
from recsim_ng.lib.tensorflow import runtime
import tensorflow as tf

FieldSpec = field_spec.FieldSpec
ValueSpec = value_lib.ValueSpec
FieldValue = value_lib.FieldValue
Value = value_lib.Value

Dependency = variable_lib.Dependency
ValueDef = variable_lib.ValueDef
Variable = variable_lib.Variable

Network = network_lib.Network
NetworkValue = network_lib.NetworkValue
NetworkValueTrajectory = runtime.NetworkValueTrajectory

_OBSERVATION_INDEX_FIELD = "__log_probability_observation_index"


def log_probability_from_value_trajectory(
    variables,
    value_trajectory,
    num_steps,
    graph_compile = True):
  """Log probability of a trajectory of network outputs.

  Provides a direct interface to evaluate the outputs of a network simulation,
  for example:
  ```
    variables = story()
    network = network_lib.Network(variables)
    tf_runtime = runtime.TFRuntime(network)
    trajectory = tf_runtime.trajectory(length=5)
    log_p = log_probability_from_value_trajectory(variables, trajectory, 4)
  ```

  Args:
    variables: A collection of `Variable`s defining a dynamic Bayesian network
      (DBN).
    value_trajectory: A trajectory generated from `TFRuntime.trajectory`.
    num_steps: The number of time steps over which to measure the probability.
    graph_compile: Boolean indicating whether the log prob computation is run in
      graph mode.

  Returns:
    A Tensor like that returned from `tfp.distributions.Distribution.log_prob`.
  """
  variables = tuple(variables)
  observations = replay_variables(variables, value_trajectory)
  return log_probability(
      variables, observations, num_steps, graph_compile=graph_compile)


def replay_variables(
    variables,
    value_trajectory):
  """Trajectory replay variables for log probability computation.

  Given a sequence of variables and a trajectory of observed values of these
  variables, this function constructs a sequence of observation variables with
  corresponding to the simulation variables replaying their logged values.

  Args:
    variables: A sequence of `Variable`s defining a dynamic Bayesian network
      (DBN).
    value_trajectory: A trajectory generated from `TFRuntime.trajectory`.

  Returns:
    A sequence of `Variable`.
  """
  observations = []
  for var in variables:
    obs = data.data_variable(
        name=var.name + " obs",
        spec=var.spec,
        data_sequence=data.SlicedValue(value=value_trajectory[var.name]),
        data_index_field=_OBSERVATION_INDEX_FIELD)
    observations.append(obs)
  return observations


def log_prob_variables_from_observation(
    variables,
    observation):
  """Log probability variables for a sequence of observational data.

  Given a sequence of simulation variables and a corresponding sequence
  observation variables (e.g. as generated by `replay_variables`), this function
  generates sequence of log probability variables corresponding to the
  simulation-observation variable pairs. A log probability variable has the
  same fields as its simulation variable with values corresponding to
  log p(observed field value | observations of dependencies), where p is
  the probability mass / density function extracted from the variable's `fn`.
  Deterministic field names are assigned a scalar value of 0.

  Args:
    variables: A sequence of `Variable`s defining a dynamic Bayesian network
      (DBN).
    observation: A sequence of `Variable`s that corresponds one-to-one with
      `variables` and which defines an observation of the DBN.

  Returns:
    A sequence of `Variable`.

  Throws:
    ValueError if the number of simulation variables does not correspond to the
    number of observation variables.
  """
  if len(variables) != len(observation):
    raise ValueError(
        "number of observation variabbles ({}) does not match number of"
        " variables ({})".format(len(observation), len(variables)))

  observation_name = {
      var.name: obs.name for var, obs in zip(variables, observation)
  }
  collision = set(observation_name).intersection(set(observation_name.values()))
  if collision:
    raise ValueError(
        "variables and observation variables share the same names: {}".format(
            collision))

  log_prob_transformation = functools.partial(
      _log_probability_from_observations, observation_name=observation_name)
  log_prob_vars = list(map(log_prob_transformation, variables))
  return log_prob_vars


def log_prob_variables_from_direct_output(
    variables):
  """Log probability variables for outputs at simulation time.

  Given a sequence of simulation variables, this function generates a sequence
  of log probability variables containing the log probabilities of the values
  of those fields of the variables which are stochastically generated. I.e.
  the log probability variable contains log(p(X)) where X is the corresponding
  field of the simulation variable. Deterministic field names are assigned a
  scalar value of 0.

  Args:
    variables: A sequence of `Variable`s defining a dynamic Bayesian network
      (DBN).

  Returns:
    A sequence of `Variable`.

  Throws:
    ValueError if the number of simulation variables does not correspond to the
    number of observation variables.
  """
  return [_log_probability_from_direct_output(var) for var in variables]


def total_log_prob_accumulator_variable(
    log_prob_vars):
  """Accumulated joint log probability variable."""
  log_prob_vars = list(log_prob_vars)  # make sure ordering is stable.
  log_prob_accum = Variable(
      name="total_log_prob_accum", spec=ValueSpec(accum=FieldSpec()))

  def summation(*log_probs):
    reduced_log_probs = []
    for value in log_probs:
      reduced_log_probs.extend(
          [tf.reduce_sum(v) for v in value.as_dict.values()])
    return Value(accum=sum(reduced_log_probs))

  def accumulate(prev, *log_probs):
    return Value(accum=prev.get("accum") + summation(*log_probs).get("accum"))

  log_prob_accum.initial_value = variable_lib.value(
      fn=summation, dependencies=log_prob_vars)
  log_prob_accum.value = variable_lib.value(
      fn=accumulate, dependencies=[log_prob_accum.previous] + log_prob_vars)
  return log_prob_accum


def log_prob_accumulator_variable(log_prob_var):
  """Temporal accumulation of log probability variables.

  Given a log probability variable, outputs temporal per-field accumulator
  of the log probability values of the variable up to the current time instance.

  Args:
    log_prob_var: An instance of `Variable` computing the per-time-step log
      probability of an simulation-observation variable pair (e.g. as generated
      by `log_prob_variables`.

  Returns:
    A `Variable` outputting the per-field sum of all values of the input
    variable up to the current time-step.
  """
  log_prob_accum = Variable(
      name=log_prob_var.name + "_accum", spec=log_prob_var.spec)

  def accumulate_fields(prev, log_prob):
    accumulated_values = {}
    for field_name, value in log_prob.as_dict.items():
      prev_value = prev.get(field_name)
      if not isinstance(prev_value, tf.Tensor):
        prev_value = tf.constant(prev_value)
      if not isinstance(value, tf.Tensor):
        value = tf.constant(value)
      shape_mismatch_msg = (
          f"The log probability of field {field_name} of {log_prob_var.name}"
          f" changes shape during iteration from {prev_value.shape} to"
          f" {value.shape}.")
      tf.debugging.assert_equal(
          tf.shape(prev_value), tf.shape(value), message=shape_mismatch_msg)
      accumulated_values[field_name] = prev_value + value
    return Value(
        **{
            field_name: prev.get(field_name) + value
            for field_name, value in log_prob.as_dict.items()
        })

  log_prob_accum.initial_value = variable_lib.value(
      fn=lambda log_prob: log_prob, dependencies=[log_prob_var])
  log_prob_accum.value = variable_lib.value(
      fn=accumulate_fields,
      dependencies=[log_prob_accum.previous, log_prob_var])
  return log_prob_accum


def log_prob_accumulator_variables(
    log_prob_vars):
  """List version of `log_prob_accumulator_variable`."""
  return [log_prob_accumulator_variable(lpvar) for lpvar in log_prob_vars]


def log_probability(variables,
                    observation,
                    num_steps,
                    graph_compile = True):
  """Returns the joint log probability of an observation given a network.

  Please note that the correctness of the result requires that all of the value
  functions of all the `Variable`s create `ed.RandomVariable` objects in a
  stable order. In other words, if a value function is invoked twice, it will
  create logically corresponding `ed.RandomVariable` objects in the same order.

  Args:
    variables: A sequence of `Variable`s defining a dynamic Bayesian network
      (DBN).
    observation: A sequence of `Variable`s that corresponds one-to-one with
      `variables` and which defines an observation of the DBN.
    num_steps: The number of time steps over which to measure the probability.
    graph_compile: Boolean indicating whether the computation should be run in
      graph mode.

  Returns:
    A Tensor like that returned from `tfp.distributions.Distribution.log_prob`.
  """

  log_prob_vars = log_prob_variables_from_observation(variables, observation)
  accumulator = total_log_prob_accumulator_variable(log_prob_vars)
  tf_runtime = runtime.TFRuntime(
      network=Network(
          variables=list(observation) + list(log_prob_vars) + [accumulator]),
      graph_compile=graph_compile)

  return tf_runtime.execute(num_steps)["total_log_prob_accum"].get("accum")


def _log_probability_from_direct_output(variable):
  """Generates log probability vars from the outputs of a given variable."""
  log_prob_spec = {
      field_name: FieldSpec() for field_name in variable.spec.as_dict.keys()
  }
  log_prob_variable = Variable(
      name=variable.name + "_log_prob", spec=ValueSpec(**log_prob_spec))

  def get_log_probs(value):

    def lp_fn(field):
      if not isinstance(field, ed.RandomVariable):
        return tf.constant(0.0)
      return field.distribution.log_prob(field.value)

    return value.map(lp_fn)

  log_prob_variable.value = variable_lib.value(
      fn=get_log_probs, dependencies=(variable,))
  return log_prob_variable


def _log_probability_from_observations(
    variable, observation_name):
  """Returns a `Variable` that computes log-probability.

  Given a `Variable` with inputs `[a_1, ..., a_k]` and output `b`, this returns
  a `Variable` with inputs `[c_1, ..., c_k, d]` and output `{"log_prob": p}`
  where `c_i` is the observed value of `a_i`, `d` is the observed value of `b`,
  and `p` is the joint log-probability of `d` conditioned on `[c_1, ..., c_k]`.

  Args:
    variable: A `Variable`.
    observation_name: A mapping that, given the `Variable` name of an input or
      output of `variable`, gives the name of the corresponding observation
      `Variable`.

  Returns:
    A `Variable` that yields the log-probability as described above.
  """
  # TODO(ccolby): can we have remove_data_index for variables?
  field_names_no_data = [
      key for key in variable.spec.as_dict.keys() if not key.startswith("__")
  ]
  log_prob_spec = {
      field_name: FieldSpec() for field_name in field_names_no_data
  }
  transformed_variable = Variable(
      name=variable.name + "_log_prob", spec=ValueSpec(**log_prob_spec))

  def transform_value_def(value_def):

    def rewire_dependency_to_observation(dep):
      return Dependency(
          variable_name=observation_name[dep.variable_name],
          on_current_value=dep.on_current_value)

    transformed_dependencies = ([
        Dependency(
            variable_name=observation_name[variable.name],
            on_current_value=True)
    ] + list(map(rewire_dependency_to_observation, value_def.dependencies)))

    def log_prob_fn(observed_output, *observed_inputs):
      unfiltered_log_prob = _transform_fn(value_def.fn)(observed_output,
                                                        *observed_inputs)
      filtered_dict = {
          key: (unfiltered_log_prob.get(key)
                if key in unfiltered_log_prob.as_dict else 0.0)
          for key in field_names_no_data
      }
      return Value(**filtered_dict)

    return ValueDef(fn=log_prob_fn, dependencies=transformed_dependencies)

  if variable.has_explicit_initial_value:
    transformed_variable.initial_value = transform_value_def(
        variable.initial_value)
  if variable.has_explicit_value:
    transformed_variable.value = transform_value_def(variable.value)

  return transformed_variable


def _transform_fn(fn):
  """Returns a transformed value function that computes log-probability.

  For example, suppose `fn(x, y) = {"a": rv1, "b": rv2}` where
  ```
    rv1 = F(x, y)
    rv2 = G(rv1)
  ```
  and suppsoed the observed output is `{"a": ov1, "b": ov2}`. This will return a
  transformed function
  ```
      _transform_fn(fn)({"a": ov1, "b": ov2}, x, y) = {"log_prob": p}
  ```
  that has four random variables, two for each of `rv1` and `rv2`:
  ```
    rv1_a = F(x, y)
    rv1_b = ov1
    rv2_a = G(rv1_b)
    rv2_b = ov2
  ```
  The output, `p`, is the sum of:
  ```
    rv1_a.distribution.log_prob(ov1)
    rv2_a.distribution.log_prob(ov2)
  ```
  Note that in this example, `rv2_b` is created but not used.

  Args:
    fn: A `Variable` value function (appearing in `ValueDef.fn`).

  Returns:
    A transformed function as described above.

  Raises:
    RuntimeError: If `fn` creates any `ed.RandomVariable` objects that are not
      exposed as fields of its output `Value`.
  """

  def transformed_fn(observed_output, *observed_inputs):
    cleaned_inputs = list(
        map(
            functools.partial(
                data.remove_data_index,
                data_index_field=_OBSERVATION_INDEX_FIELD), observed_inputs))
    observed_output_values_by_rv_order = _observed_output_values_by_rv_order(
        fn, observed_output, *cleaned_inputs)
    num_rvs = len(observed_output_values_by_rv_order)

    rv_index = 0
    log_probs = {}

    def log_prob_tracer(rv_constructor, *args, **kwargs):
      nonlocal rv_index
      if rv_index >= num_rvs:
        raise RuntimeError(
            "function created {} random variables the first time it was called,"
            " but created more the second time".format(num_rvs))
      field_name, observed_value = observed_output_values_by_rv_order[rv_index]
      rv_index += 1
      rv = rv_constructor(*args, **kwargs)
      logp = rv.distribution.log_prob(observed_value)
      log_probs[field_name] = logp
      kwargs["value"] = observed_value
      # be nice to higher tracers
      return ed.traceable(rv_constructor)(*args, **kwargs)

    with ed.trace(log_prob_tracer):
      _ = fn(*cleaned_inputs)

    return Value(**log_probs)

  return transformed_fn


def _observed_output_values_by_rv_order(
    fn, observed_output,
    *observed_inputs):
  """Returns a sequence of observed values in order of RandomVariable creation.

  For example, suppose `fn(x, y) = {"a": rv1, "b": rv2}` where
  ```
    rv1 = F(x, y)
    rv2 = G(rv1)
  ```
  Then,
  ```
    _observed_values_by_rv_order(fn, {"a": ov1, "b": ov2}, x, y) = [ov1, ov2]
  ```
  Here, `fn` constructs two `ed.RandomVariable` objects, first `rv1` and then
  `rv2`. Because `rv1` was constructed first, `ov1` appears first in the
  returned sequence.

  Args:
    fn: A `Variable` value function; see `ValueDef.fn`.
    observed_output: A `Value` containing field names matching the value
      returned by `fn`.
    *observed_inputs: The input arguments of `fn`.

  Returns:
    A sequence of (field_name, `observed_output`) pairs of field values
    corresponding to the `ed.RandomVariable` field values output by `fn`,
    arranged in the order in which `fn` constructs those `ed.RandomVariable`
    objects together with the names of the fields they get assigned to.

  Raises:
    RuntimeError: If `fn` creates any `ed.RandomVariable` objects that are not
      exposed as fields of its output `Value`.
  """

  rvs_in_order_of_construction = []

  def index_random_variables(rv_constructor, *args, **kwargs):
    rv = rv_constructor(*args, **kwargs)
    rvs_in_order_of_construction.append(rv)
    return rv

  with ed.trace(index_random_variables):
    temporary_output = fn(*observed_inputs)

  rv_to_output_value = {
      value: (field_name, observed_output.get(field_name))
      for field_name, value in temporary_output.as_dict.items()
      if isinstance(value, ed.RandomVariable)
  }
  unobserved = [
      rv for rv in rvs_in_order_of_construction if rv not in rv_to_output_value
  ]
  if unobserved:
    raise RuntimeError(
        "unobserved random variables; log-probability cannot be computed: {}"
        .format(unobserved))

  return [rv_to_output_value[rv] for rv in rvs_in_order_of_construction]

# coding=utf-8 # Copyright 2021 The RecSim Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # #     http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # python3 Compute the joint log-probability of a Network given an observation. # type: ignore Log probability of a trajectory of network outputs.    Provides a direct interface to evaluate the outputs of a network simulation,   for example:   ```     variables = story()     network = network_lib.Network(variables)     tf_runtime = runtime.TFRuntime(network)     trajectory = tf_runtime.trajectory(length=5)     log_p = log_probability_from_value_trajectory(variables, trajectory, 4)   ```    Args:     variables: A collection of `Variable`s defining a dynamic Bayesian network       (DBN).     value_trajectory: A trajectory generated from `TFRuntime.trajectory`.     num_steps: The number of time steps over which to measure the probability.     graph_compile: Boolean indicating whether the log prob computation is run in       graph mode.    Returns:     A Tensor like that returned from `tfp.distributions.Distribution.log_prob`. Trajectory replay variables for log probability computation.    Given a sequence of variables and a trajectory of observed values of these   variables, this function constructs a sequence of observation variables with   corresponding to the simulation variables replaying their logged values.    Args:     variables: A sequence of `Variable`s defining a dynamic Bayesian network       (DBN).     value_trajectory: A trajectory generated from `TFRuntime.trajectory`.    Returns:     A sequence of `Variable`. Log probability variables for a sequence of observational data.    Given a sequence of simulation variables and a corresponding sequence   observation variables (e.g. as generated by `replay_variables`), this function   generates sequence of log probability variables corresponding to the   simulation-observation variable pairs. A log probability variable has the   same fields as its simulation variable with values corresponding to   log p(observed field value | observations of dependencies), where p is   the probability mass / density function extracted from the variable's `fn`.   Deterministic field names are assigned a scalar value of 0.    Args:     variables: A sequence of `Variable`s defining a dynamic Bayesian network       (DBN).     observation: A sequence of `Variable`s that corresponds one-to-one with       `variables` and which defines an observation of the DBN.    Returns:     A sequence of `Variable`.    Throws:     ValueError if the number of simulation variables does not correspond to the     number of observation variables. Log probability variables for outputs at simulation time.    Given a sequence of simulation variables, this function generates a sequence   of log probability variables containing the log probabilities of the values   of those fields of the variables which are stochastically generated. I.e.   the log probability variable contains log(p(X)) where X is the corresponding   field of the simulation variable. Deterministic field names are assigned a   scalar value of 0.    Args:     variables: A sequence of `Variable`s defining a dynamic Bayesian network       (DBN).    Returns:     A sequence of `Variable`.    Throws:     ValueError if the number of simulation variables does not correspond to the     number of observation variables. Accumulated joint log probability variable. # make sure ordering is stable. Temporal accumulation of log probability variables.    Given a log probability variable, outputs temporal per-field accumulator   of the log probability values of the variable up to the current time instance.    Args:     log_prob_var: An instance of `Variable` computing the per-time-step log       probability of an simulation-observation variable pair (e.g. as generated       by `log_prob_variables`.    Returns:     A `Variable` outputting the per-field sum of all values of the input     variable up to the current time-step. List version of `log_prob_accumulator_variable`. Returns the joint log probability of an observation given a network.    Please note that the correctness of the result requires that all of the value   functions of all the `Variable`s create `ed.RandomVariable` objects in a   stable order. In other words, if a value function is invoked twice, it will   create logically corresponding `ed.RandomVariable` objects in the same order.    Args:     variables: A sequence of `Variable`s defining a dynamic Bayesian network       (DBN).     observation: A sequence of `Variable`s that corresponds one-to-one with       `variables` and which defines an observation of the DBN.     num_steps: The number of time steps over which to measure the probability.     graph_compile: Boolean indicating whether the computation should be run in       graph mode.    Returns:     A Tensor like that returned from `tfp.distributions.Distribution.log_prob`. Generates log probability vars from the outputs of a given variable. Returns a `Variable` that computes log-probability.    Given a `Variable` with inputs `[a_1, ..., a_k]` and output `b`, this returns   a `Variable` with inputs `[c_1, ..., c_k, d]` and output `{"log_prob": p}`   where `c_i` is the observed value of `a_i`, `d` is the observed value of `b`,   and `p` is the joint log-probability of `d` conditioned on `[c_1, ..., c_k]`.    Args:     variable: A `Variable`.     observation_name: A mapping that, given the `Variable` name of an input or       output of `variable`, gives the name of the corresponding observation       `Variable`.    Returns:     A `Variable` that yields the log-probability as described above. # TODO(ccolby): can we have remove_data_index for variables? Returns a transformed value function that computes log-probability.    For example, suppose `fn(x, y) = {"a": rv1, "b": rv2}` where   ```     rv1 = F(x, y)     rv2 = G(rv1)   ```   and suppsoed the observed output is `{"a": ov1, "b": ov2}`. This will return a   transformed function   ```       _transform_fn(fn)({"a": ov1, "b": ov2}, x, y) = {"log_prob": p}   ```   that has four random variables, two for each of `rv1` and `rv2`:   ```     rv1_a = F(x, y)     rv1_b = ov1     rv2_a = G(rv1_b)     rv2_b = ov2   ```   The output, `p`, is the sum of:   ```     rv1_a.distribution.log_prob(ov1)     rv2_a.distribution.log_prob(ov2)   ```   Note that in this example, `rv2_b` is created but not used.    Args:     fn: A `Variable` value function (appearing in `ValueDef.fn`).    Returns:     A transformed function as described above.    Raises:     RuntimeError: If `fn` creates any `ed.RandomVariable` objects that are not       exposed as fields of its output `Value`. # be nice to higher tracers Returns a sequence of observed values in order of RandomVariable creation.    For example, suppose `fn(x, y) = {"a": rv1, "b": rv2}` where   ```     rv1 = F(x, y)     rv2 = G(rv1)   ```   Then,   ```     _observed_values_by_rv_order(fn, {"a": ov1, "b": ov2}, x, y) = [ov1, ov2]   ```   Here, `fn` constructs two `ed.RandomVariable` objects, first `rv1` and then   `rv2`. Because `rv1` was constructed first, `ov1` appears first in the   returned sequence.    Args:     fn: A `Variable` value function; see `ValueDef.fn`.     observed_output: A `Value` containing field names matching the value       returned by `fn`.     *observed_inputs: The input arguments of `fn`.    Returns:     A sequence of (field_name, `observed_output`) pairs of field values     corresponding to the `ed.RandomVariable` field values output by `fn`,     arranged in the order in which `fn` constructs those `ed.RandomVariable`     objects together with the names of the fields they get assigned to.    Raises:     RuntimeError: If `fn` creates any `ed.RandomVariable` objects that are not       exposed as fields of its output `Value`.