Clarify amplitude and length_scale (#32)

Author: a-ws-m

tests/test_kernels.py

@@ -41,14 +41,22 @@ def test_custom_kernel_reloading(tmp_path: Path):
 @pytest.mark.parametrize("kernel_type", [RBFKernelFn, MaternOneHalfFn])
 def test_reload(tmp_path: Path, kernel_type: Type[AmpAndLengthScaleFn]):
     """Test saving and reloading a builtin kernel to/from disk."""
-    # Example layer weights: [_amplitude, _length_scale]
+    # Example layer weights: [_amplitude_basis, _length_scale_basis]
     example_weights = [np.array(2.0), np.array(3.0)]
     save_dir = tmp_path / "kernel"
     orig_kernel = kernel_type()
     orig_kernel.set_weights(example_weights)
     orig_kernel.save(save_dir)
 
-    loaded_kernel = load_kernel(save_dir)
+    orig_amplitude = orig_kernel.amplitude
+    orig_length_scale = orig_kernel.length_scale
+
+    loaded_kernel: AmpAndLengthScaleFn = load_kernel(save_dir)
     assert isinstance(loaded_kernel, kernel_type)
     loaded_weights = loaded_kernel.get_weights()
     assert loaded_weights == example_weights
+
+    loaded_amp = loaded_kernel.amplitude
+    loaded_ls = loaded_kernel.length_scale
+    assert loaded_amp == pytest.approx(orig_amplitude)
+    assert loaded_ls == pytest.approx(orig_length_scale)

unlocknn/kernel_layers.py

@@ -89,8 +89,10 @@ class AmpAndLengthScaleFn(KernelLayer, ABC):
     """An ABC for kernels with amplitude and length scale parameters.
 
     Attributes:
-        _amplitude (tf.Tensor): The amplitude of the kernel.
-        _length_scale (tf.Tensor): The length scale of the kernel.
+        _amplitude_basis (tf.Tensor): The basis for the kernel amplitude,
+            which is passed through a softplus to calculate the actual amplitude.
+        _length_scale_basis (tf.Tensor): The basis for the length scale of the kernel.
+            which is passed through a softplus to calculate the actual amplitude.
 
     """
 
@@ -99,30 +101,42 @@ def __init__(self, **kwargs):
         super().__init__(**kwargs)
         dtype = kwargs.get("dtype", tf.float64)
 
-        self._amplitude = self.add_weight(
+        self._amplitude_basis = self.add_weight(
             initializer=tf.constant_initializer(0), dtype=dtype, name="amplitude"
         )
 
-        self._length_scale = self.add_weight(
+        self._length_scale_basis = self.add_weight(
             initializer=tf.constant_initializer(0), dtype=dtype, name="length_scale"
         )
+    
+    @property
+    def amplitude(self) -> float:
+        """Get the current kernel amplitude."""
+        return tf.nn.softplus(0.1 * self._amplitude_basis).numpy().item()
+    
+    @property
+    def length_scale(self) -> float:
+        """Get the current kernel length scale."""
+        return tf.nn.softplus(5.0 * self._length_scale_basis).numpy().item()
 
 
 class RBFKernelFn(AmpAndLengthScaleFn):
     """A radial basis function implementation that works with keras.
 
     Attributes:
-        _amplitude (tf.Tensor): The amplitude of the kernel.
-        _length_scale (tf.Tensor): The length scale of the kernel.
+        _amplitude_basis (tf.Tensor): The basis for the kernel amplitude,
+            which is passed through a softplus to calculate the actual amplitude.
+        _length_scale_basis (tf.Tensor): The basis for the length scale of the kernel.
+            which is passed through a softplus to calculate the actual amplitude.
 
     """
 
     @property
     def kernel(self) -> tfp.math.psd_kernels.PositiveSemidefiniteKernel:
         """Get a callable kernel."""
         return tfp.math.psd_kernels.ExponentiatedQuadratic(
-            amplitude=tf.nn.softplus(0.1 * self._amplitude),
-            length_scale=tf.nn.softplus(5.0 * self._length_scale),
+            amplitude=self.amplitude,
+            length_scale=self.length_scale,
         )
 
     @property
@@ -134,17 +148,19 @@ class MaternOneHalfFn(AmpAndLengthScaleFn):
     """A Matern kernel with parameter 1/2 implementation that works with keras.
 
     Attributes:
-        _amplitude (tf.Tensor): The amplitude of the kernel.
-        _length_scale (tf.Tensor): The length scale of the kernel.
+        _amplitude_basis (tf.Tensor): The basis for the kernel amplitude,
+            which is passed through a softplus to calculate the actual amplitude.
+        _length_scale_basis (tf.Tensor): The basis for the length scale of the kernel.
+            which is passed through a softplus to calculate the actual amplitude.
 
     """
 
     @property
     def kernel(self) -> tfp.math.psd_kernels.PositiveSemidefiniteKernel:
         """Get a callable kernel."""
         return tfp.math.psd_kernels.MaternOneHalf(
-            amplitude=tf.nn.softplus(0.1 * self._amplitude),
-            length_scale=tf.nn.softplus(5.0 * self._length_scale),
+            amplitude=self.amplitude,
+            length_scale=self.length_scale,
         )
 
     @property