Inference Gym: Add a smoothing option to PlasmaSpectroscopy problem.

I didn't add a new dataset, as I wasn't able to produce a good example of
multimodality with this setting on.

PiperOrigin-RevId: 380846216

Author: SiegeLordEx

spinoffs/inference_gym/inference_gym/targets/BUILD

@@ -393,6 +393,7 @@ py_test(
     srcs_version = "PY3",
     deps = [
         ":plasma_spectroscopy",
+        # absl/testing:parameterized dep,
         # numpy dep,
         # tensorflow dep,
         # tensorflow_probability/python/internal:test_util dep,

spinoffs/inference_gym/inference_gym/targets/plasma_spectroscopy.py

@@ -92,6 +92,7 @@ def __init__(
       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',
   ):
@@ -124,6 +125,8 @@ def __init__(
         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.
 
@@ -151,6 +154,7 @@ def __init__(
       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():
@@ -262,6 +266,10 @@ def absolute_noise_scale(self):
   def relative_noise_scale(self):
     return self._relative_noise_scale
 
+  @property
+  def use_bump_function(self):
+    return self._use_bump_function
+
   def forward_model(self, sample):
     """The forward model.
 
@@ -314,6 +322,10 @@ def forward_model(self, sample):
             -(wavelengths[:, tf.newaxis] - doppler_shifted_center_wavelength)**2
             / (2 * bandwidth**2))
 
+    if self.use_bump_function:
+      emissivity *= tfp.math.round_exponential_bump_function(
+          tf.linspace(-1., 1., self.num_shells))
+
     x = tf.linspace(-self.outer_shell_radius, self.outer_shell_radius,
                     self.num_integration_points)
     y = tf.linspace(-self.sensor_span, self.sensor_span, self.num_sensors)

spinoffs/inference_gym/inference_gym/targets/plasma_spectroscopy_test.py

@@ -15,6 +15,7 @@
 # ============================================================================
 """Tests for inference_gym.targets.plasma_spectroscopy."""
 
+from absl.testing import parameterized
 import numpy as np
 import tensorflow.compat.v2 as tf
 
@@ -36,13 +37,21 @@ def _test_dataset():
 @test_util.multi_backend_test(globals(), 'targets.plasma_spectroscopy_test')
 class PlasmaSpectroscopyTest(test_util.InferenceGymTestCase):
 
-  def testBasic(self):
+  @parameterized.named_parameters(
+      ('Smooth', True),
+      ('NotSmooth', False),
+  )
+  def testBasic(self, use_bump_function):
     """Checks that you get finite values given unconstrained samples.
 
     We check `unnormalized_log_prob` as well as the values of the sample
     transformations.
+
+    Args:
+      use_bump_function: Whether to use the bump function.
     """
-    model = plasma_spectroscopy.PlasmaSpectroscopy(**_test_dataset())
+    model = plasma_spectroscopy.PlasmaSpectroscopy(
+        **_test_dataset(), use_bump_function=use_bump_function)
     self.validate_log_prob_and_transforms(
         model,
         sample_transformation_shapes=dict(