class_docstr: Plasma spectroscopy model.
  
  This implements the fusion plasma spectroscopy problem described in [1]. The
  goal of the model is to reconstruct the plasma configuration in a linear
  fusion reactor via spectroscopy. The plasma blob is instrumented by an array
  of regularly spaced spectrometers arranged in the transverse plane to the
  reactor. Each spectrometer can measure the emission spectrum integrated along
  the line of sight of the sensor. The sensor array can resolve shifts
  perpendicular to the their lines of sight, but cannot measure shifts along
  them. The plasma is parameterized via amplitude, temperature and velocity ,
  which are assumed to be radially symmetric about the plasma blob center and
  are spatially correlated.
  
  The overall model is as follows:
  
  ```none
  for r in range(num_shells):
    amplitude[r] ~ Normal(0, 1)
    temperature[r] ~ Normal(0, 1)
    velocity[r] ~ Normal(0, 1)
  shift ~ Normal(0, 1)
  
  mean_emission = forward_model(amplitude, temperature, velocity, shift)
  
  measurements[wavelength, sensor] ~ Normal(mean_emission,
    (absolute_noise_scale**2 + relative_noise_scale**2 * mean_emission)**0.5)
  ```
  
  where the `forward_model` describes the mechanism of the spectrometers. Note
  that the latent parameters are unconstrained and the correlations between them
  are baked into the forward model.
  
  This model is notable for poor conditioning and presence of multiple low
  probability modes, requiring careful initialization for MCMC to progress.
  
  #### References
  
  [1]: Langmore I, Dikovsky M, Geraedts S, Norgaard P, von Behren R. Hamiltonian
       Monte Carlo in Inverse Problems; Ill-Conditioning and Multi-Modality.
       2021. http://arxiv.org/abs/2103.07515
class_docstr_tok: ["Plasma","spectroscopy","model",".","This","implements","the","fusion","plasma","spectroscopy","problem","described","in","[","1","]",".","The","goal","of","the","model","is","to","reconstruct","the","plasma","configuration","in","a","linear","fusion","reactor","via","spectroscopy",".","The","plasma","blob","is","instrumented","by","an","array","of","regularly","spaced","spectrometers","arranged","in","the","transverse","plane","to","the","reactor",".","Each","spectrometer","can","measure","the","emission","spectrum","integrated","along","the","line","of","sight","of","the","sensor",".","The","sensor","array","can","resolve","shifts","perpendicular","to","the","their","lines","of","sight",",","but","can","not","measure","shifts","along","them",".","The","plasma","is","parameterized","via","amplitude",",","temperature","and","velocity",",","which","are","assumed","to","be","radially","symmetric","about","the","plasma","blob","center","and","are","spatially","correlated",".","The","overall","model","is","as","follows",":","``","`","none","for","r","in","range","(","num_shells",")",":","amplitude","[","r","]","~","Normal","(","0",",","1",")","temperature","[","r","]","~","Normal","(","0",",","1",")","velocity","[","r","]","~","Normal","(","0",",","1",")","shift","~","Normal","(","0",",","1",")","mean_emission","=","forward_model","(","amplitude",",","temperature",",","velocity",",","shift",")","measurements","[","wavelength",",","sensor","]","~","Normal","(","mean_emission",",","(","absolute_noise_scale","*","*","2","+","relative_noise_scale","*","*","2","*","mean_emission",")","*","*","0.5",")","``","`","where","the","`","forward_model","`","describes","the","mechanism","of","the","spectrometers",".","Note","that","the","latent","parameters","are","unconstrained","and","the","correlations","between","them","are","baked","into","the","forward","model",".","This","model","is","notable","for","poor","conditioning","and","presence","of","multiple","low","probability","modes",",","requiring","careful","initialization","for","MCMC","to","progress",".","#","#","#","#","References","[","1","]",":","Langmore","I",",","Dikovsky","M",",","Geraedts","S",",","Norgaard","P",",","von","Behren","R.","Hamiltonian","Monte","Carlo","in","Inverse","Problems",";","Ill-Conditioning","and","Multi-Modality",".","2021.","http",":","\/\/arxiv.org\/abs\/2103.07515"]
docstr: Construct the PlasmaSpectroscopy model.
  
  Args:
    measurements: Float `Tensor` with shape [num_wavelengths, num_sensors].
      The spectrometer measurements.
    wavelengths: Float `Tensor` with shape [num_wavelengths]. Wavelengths
      measured by the spectrometers.
    center_wavelength: Float `Tensor` scalar. The center wavelength of the
      target emission line.
    num_shells: Python integer. Number of radial shells to model.
    outer_shell_radius: Python float. Outermost shell radius.
    sensor_span: Python float. Half-span of the sensors (i.e. the sensors are
      regularly placed on the interval `[-sensor_span, sensor_span]`.
    num_integration_points: Python integer. How many points to use to
      discretize the integral used to compute the integrated emission spectra.
    prior_diag_noise_variance: Float `Tensor` scalar. Diagonal variance of the
      latent variables.
    prior_length_scale: Float `Tensor` scalar. Length scale of the radial
      correlations between the latent variables.
    amplitude_scale: Float `Tensor` scalar. Scale factor of the constrained
      amplitude.
    temperature_scale: Float `Tensor` scalar. Scale factor of the constrained
      temperature.
    velocity_scale: Float `Tensor` scalar. Scale factor of the constrained
      velocity.
    absolute_noise_scale: Float `Tensor` scalar. Absolute noise scale of the
      observation noise.
    relative_noise_scale: Float `Tensor` scalar. Absolute noise scale of the
      observation noise.
    use_bump_function: Python `bool`. If True, use a bump function to smoothly
      decay the plasma blob to 0 at the edges.
    name: Python `str` name prefixed to Ops created by this class.
    pretty_name: A Python `str`. The pretty name of this model.
  
  Raises:
    ValueError: If `measurements` and `wavelengths` have inconsistent shape.
docstr_tok: ["Construct","the","PlasmaSpectroscopy","model",".","Args",":","measurements",":","Float","`","Tensor","`","with","shape","[","num_wavelengths",",","num_sensors","]",".","The","spectrometer","measurements",".","wavelengths",":","Float","`","Tensor","`","with","shape","[","num_wavelengths","]",".","Wavelengths","measured","by","the","spectrometers",".","center_wavelength",":","Float","`","Tensor","`","scalar",".","The","center","wavelength","of","the","target","emission","line",".","num_shells",":","Python","integer",".","Number","of","radial","shells","to","model",".","outer_shell_radius",":","Python","float",".","Outermost","shell","radius",".","sensor_span",":","Python","float",".","Half-span","of","the","sensors","(","i.e",".","the","sensors","are","regularly","placed","on","the","interval","`","[","-sensor_span",",","sensor_span","]","`",".","num_integration_points",":","Python","integer",".","How","many","points","to","use","to","discretize","the","integral","used","to","compute","the","integrated","emission","spectra",".","prior_diag_noise_variance",":","Float","`","Tensor","`","scalar",".","Diagonal","variance","of","the","latent","variables",".","prior_length_scale",":","Float","`","Tensor","`","scalar",".","Length","scale","of","the","radial","correlations","between","the","latent","variables",".","amplitude_scale",":","Float","`","Tensor","`","scalar",".","Scale","factor","of","the","constrained","amplitude",".","temperature_scale",":","Float","`","Tensor","`","scalar",".","Scale","factor","of","the","constrained","temperature",".","velocity_scale",":","Float","`","Tensor","`","scalar",".","Scale","factor","of","the","constrained","velocity",".","absolute_noise_scale",":","Float","`","Tensor","`","scalar",".","Absolute","noise","scale","of","the","observation","noise",".","relative_noise_scale",":","Float","`","Tensor","`","scalar",".","Absolute","noise","scale","of","the","observation","noise",".","use_bump_function",":","Python","`","bool","`",".","If","True",",","use","a","bump","function","to","smoothly","decay","the","plasma","blob","to","0","at","the","edges",".","name",":","Python","`","str","`","name","prefixed","to","Ops","created","by","this","class",".","pretty_name",":","A","Python","`","str","`",".","The","pretty","name","of","this","model",".","Raises",":","ValueError",":","If","`","measurements","`","and","`","wavelengths","`","have","inconsistent","shape","."]
code: def __init__(
      self,
      measurements,
      wavelengths,
      center_wavelength,
      num_shells=16,
      outer_shell_radius=1.0,
      sensor_span=1.5,
      num_integration_points=64,
      prior_diag_noise_variance=0.1,
      prior_length_scale=0.05,
      amplitude_scale=2.0,
      temperature_scale=0.2,
      velocity_scale=0.5,
      absolute_noise_scale=0.5,
      relative_noise_scale=0.05,
      use_bump_function=False,
      name="plasma_spectroscopy",
      pretty_name="Plasma Spectroscopy",
  ):
      """Construct the PlasmaSpectroscopy model.
  
      Args:
        measurements: Float `Tensor` with shape [num_wavelengths, num_sensors].
          The spectrometer measurements.
        wavelengths: Float `Tensor` with shape [num_wavelengths]. Wavelengths
          measured by the spectrometers.
        center_wavelength: Float `Tensor` scalar. The center wavelength of the
          target emission line.
        num_shells: Python integer. Number of radial shells to model.
        outer_shell_radius: Python float. Outermost shell radius.
        sensor_span: Python float. Half-span of the sensors (i.e. the sensors are
          regularly placed on the interval `[-sensor_span, sensor_span]`.
        num_integration_points: Python integer. How many points to use to
          discretize the integral used to compute the integrated emission spectra.
        prior_diag_noise_variance: Float `Tensor` scalar. Diagonal variance of the
          latent variables.
        prior_length_scale: Float `Tensor` scalar. Length scale of the radial
          correlations between the latent variables.
        amplitude_scale: Float `Tensor` scalar. Scale factor of the constrained
          amplitude.
        temperature_scale: Float `Tensor` scalar. Scale factor of the constrained
          temperature.
        velocity_scale: Float `Tensor` scalar. Scale factor of the constrained
          velocity.
        absolute_noise_scale: Float `Tensor` scalar. Absolute noise scale of the
          observation noise.
        relative_noise_scale: Float `Tensor` scalar. Absolute noise scale of the
          observation noise.
        use_bump_function: Python `bool`. If True, use a bump function to smoothly
          decay the plasma blob to 0 at the edges.
        name: Python `str` name prefixed to Ops created by this class.
        pretty_name: A Python `str`. The pretty name of this model.
  
      Raises:
        ValueError: If `measurements` and `wavelengths` have inconsistent shape.
      """
      if measurements.shape[0] != wavelengths.shape[0]:
          raise ValueError(
              "Number of measurement rows does not match the number "
              f"wavelengths {measurements.shape[0]} != "
              f"{len(wavelengths)}"
          )
  
      with tf.name_scope(name):
          self._measurements = measurements
          self._wavelengths = wavelengths
          self._center_wavelength = center_wavelength
          self._num_shells = num_shells
          self._outer_shell_radius = outer_shell_radius
          self._sensor_span = sensor_span
          self._num_integration_points = num_integration_points
          self._prior_diag_noise_variance = prior_diag_noise_variance
          self._prior_length_scale = prior_length_scale
          self._amplitude_scale = amplitude_scale
          self._temperature_scale = temperature_scale
          self._velocity_scale = velocity_scale
          self._absolute_noise_scale = absolute_noise_scale
          self._relative_noise_scale = relative_noise_scale
          self._use_bump_function = use_bump_function
  
          @tfd.JointDistributionCoroutine
          def prior():
              yield Root(
                  tfd.Sample(
                      tfd.Normal(0.0, 1.0), num_shells, name="amplitude"
                  )
              )
              yield Root(
                  tfd.Sample(
                      tfd.Normal(0.0, 1.0),
                      num_shells,
                      name="temperature",
                  )
              )
              yield Root(
                  tfd.Sample(
                      tfd.Normal(0.0, 1.0), num_shells, name="velocity"
                  )
              )
              yield Root(tfd.Normal(0.0, 1.0, name="shift"))
  
          self._prior_dist = prior
  
          def observation_noise_fn(mean_measurement):
              """Creates the observation noise distribution."""
              return tfd.Independent(
                  tfd.Normal(
                      mean_measurement,
                      tf.sqrt(
                          absolute_noise_scale**2
                          + relative_noise_scale**2
                          * mean_measurement**2
                      ),
                  ),
                  2,
              )
  
          self._observation_noise_fn = observation_noise_fn
  
          def log_likelihood_fn(sample):
              """The log_likelihood function."""
              _, mean_measurement = self.forward_model(sample)
              return observation_noise_fn(mean_measurement).log_prob(
                  measurements
              )
  
          self._log_likelihood_fn = log_likelihood_fn
  
          sample_transformations = {
              "identity": model.Model.SampleTransformation(
                  fn=lambda params: params,
                  pretty_name="Identity",
                  dtype=self._prior_dist.dtype,
              )
          }
  
      super(PlasmaSpectroscopy, self).__init__(
          default_event_space_bijector=type(self._prior_dist.dtype)(
              tfb.Identity(),
              tfb.Identity(),
              tfb.Identity(),
              tfb.Identity(),
          ),
          event_shape=self._prior_dist.event_shape,
          dtype=self._prior_dist.dtype,
          name=name,
          pretty_name=pretty_name,
          sample_transformations=sample_transformations,
      )
  
imports: import numpy
  import tensorflow.compat.v2
  import tensorflow_probability
  from inference_gym.internal import data
  from inference_gym.targets import bayesian_model
  from inference_gym.targets import model
label_desc: Use image node_id 1 to create a new PlasmaSpectroscopy object from inherited base classes: bayesian_model with example: obj = PlasmaSpectroscopy(measurements, wavelengths, center_wavelength, num_shells, outer_shell_radius, sensor_span, num_integration_points, prior_diag_noise_variance, prior_length_scale, amplitude_scale, temperature_scale, velocity_scale, absolute_noise_scale, relative_noise_scale, use_bump_function, name, pretty_name)