SOAP Kernel for IAP

experimental/shear_2x3_offset.ipynb

@@ -1005,7 +1005,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3 (ipykernel)",
+   "display_name": "gpshear_env",
    "language": "python",
    "name": "python3"
   },
@@ -1019,7 +1019,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.9"
+   "version": "3.9.18"
   }
  },
  "nbformat": 4,

experimental/shear_kernel.ipynb

@@ -1881,7 +1881,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3 (ipykernel)",
+   "display_name": "gp4iap",
    "language": "python",
    "name": "python3"
   },
@@ -1895,7 +1895,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.9"
+   "version": "3.11.7"
   }
  },
  "nbformat": 4,

src/MuyGPyS/_src/gp/kernels/soap/__init__.py

@@ -0,0 +1,8 @@
+# Copyright 2021-2024 Lawrence Livermore National Security, LLC and other
+# MuyGPyS Project Developers. See the top-level COPYRIGHT file for details.
+#
+# SPDX-License-Identifier: MIT
+
+from MuyGPyS._src.util import _collect_implementation
+
+(_soap_fn,) = _collect_implementation("MuyGPyS._src.gp.kernels.soap", "_soap_fn")

src/MuyGPyS/_src/gp/kernels/soap/jax.py

@@ -0,0 +1,28 @@
+# Copyright 2021-2024 Lawrence Livermore National Security, LLC and other
+# MuyGPyS Project Developers. See the top-level COPYRIGHT file for details.
+#
+# SPDX-License-Identifier: MIT
+
+
+def _omega(diffs, sensitivity):
+    raise NotImplementedError("JAX backend not yet supported for SOAPKernel")
+
+
+def _T1(diffs, sensitivity):
+    raise NotImplementedError("JAX backend not yet supported for SOAPKernel")
+
+
+def _T2(diffs, sensitivity):
+    raise NotImplementedError("JAX backend not yet supported for SOAPKernel")
+
+
+def _T3(diffs, sensitivity):
+    raise NotImplementedError("JAX backend not yet supported for SOAPKernel")
+
+
+def _Knm(omega, T1, T2, T3, sensitivity):
+    raise NotImplementedError("JAX backend not yet supported for SOAPKernel")
+
+
+def _soap_fn(diffs, sensitivity):
+    raise NotImplementedError("JAX backend not yet supported for SOAPKernel")

src/MuyGPyS/_src/gp/kernels/soap/mpi.py

@@ -0,0 +1,28 @@
+# Copyright 2021-2024 Lawrence Livermore National Security, LLC and other
+# MuyGPyS Project Developers. See the top-level COPYRIGHT file for details.
+#
+# SPDX-License-Identifier: MIT
+
+
+def _omega(diffs, sensitivity):
+    raise NotImplementedError("MPI backend not yet supported for SOAPKernel")
+
+
+def _T1(diffs, sensitivity):
+    raise NotImplementedError("MPI backend not yet supported for SOAPKernel")
+
+
+def _T2(diffs, sensitivity):
+    raise NotImplementedError("MPI backend not yet supported for SOAPKernel")
+
+
+def _T3(diffs, sensitivity):
+    raise NotImplementedError("MPI backend not yet supported for SOAPKernel")
+
+
+def _Knm(omega, T1, T2, T3, sensitivity):
+    raise NotImplementedError("MPI backend not yet supported for SOAPKernel")
+
+
+def _soap_fn(diffs, sensitivity):
+    raise NotImplementedError("MPI backend not yet supported for SOAPKernel")

src/MuyGPyS/_src/gp/kernels/soap/numpy.py

@@ -0,0 +1,66 @@
+# Copyright 2021-2024 Lawrence Livermore National Security, LLC and other
+# MuyGPyS Project Developers. See the top-level COPYRIGHT file for details.
+#
+# SPDX-License-Identifier: MIT
+
+import MuyGPyS._src.math.numpy as mm
+
+
+def _omega(diffs) -> mm.ndarray:
+    ndim = diffs.ndim
+    slicer = [slice(None)] * ndim
+    slicer[-3] = 0
+
+    qq_slice = diffs[tuple(slicer)]
+
+    return qq_slice
+
+
+def _T1(diffs) -> mm.ndarray:
+    ndim = diffs.ndim
+    slicer = [slice(None)] * ndim
+    slicer[-3] = 3
+
+    dd_slice = diffs[tuple(slicer)]
+
+    return dd_slice
+
+
+def _T2(diffs) -> mm.ndarray:
+    ndim = diffs.ndim
+    slicer = [slice(None)] * ndim
+    slicer[-3] = 1
+
+    diq_slice = diffs[tuple(slicer)]
+
+    return diq_slice
+
+
+def _T3(diffs) -> mm.ndarray:
+    ndim = diffs.ndim
+    slicer = [slice(None)] * ndim
+    slicer[-3] = 2
+
+    djq_slice = diffs[tuple(slicer)]
+
+    return djq_slice
+
+
+def _Knm(omega, T1, T2, T3, sensitivity) -> mm.ndarray:
+    Knm = (sensitivity - 1.0) * (omega ** (sensitivity - 2.0)) * (T2 * T3) + (
+        omega ** (sensitivity - 1.0)
+    ) * T1
+    return Knm
+
+
+def _soap_fn(diffs: mm.ndarray, sensitivity: float, **kwargs) -> mm.ndarray:
+    omega = _omega(diffs)
+    T1 = _T1(diffs)
+    T2 = _T2(diffs)
+    T3 = _T3(diffs)
+
+    Knm = _Knm(omega, T1, T2, T3, sensitivity)
+
+    Kij = sensitivity * mm.sum(Knm, axis=(-2, -1))
+
+    return Kij

src/MuyGPyS/_src/gp/kernels/soap/torch.py

@@ -0,0 +1,28 @@
+# Copyright 2021-2024 Lawrence Livermore National Security, LLC and other
+# MuyGPyS Project Developers. See the top-level COPYRIGHT file for details.
+#
+# SPDX-License-Identifier: MIT
+
+
+def _omega(diffs, sensitivity):
+    raise NotImplementedError("MPI backend not yet supported for SOAPKernel")
+
+
+def _T1(diffs, sensitivity):
+    raise NotImplementedError("MPI backend not yet supported for SOAPKernel")
+
+
+def _T2(diffs, sensitivity):
+    raise NotImplementedError("MPI backend not yet supported for SOAPKernel")
+
+
+def _T3(diffs, sensitivity):
+    raise NotImplementedError("MPI backend not yet supported for SOAPKernel")
+
+
+def _Knm(omega, T1, T2, T3, sensitivity):
+    raise NotImplementedError("MPI backend not yet supported for SOAPKernel")
+
+
+def _soap_fn(diffs, sensitivity):
+    raise NotImplementedError("MPI backend not yet supported for SOAPKernel")

src/MuyGPyS/_src/gp/tensors/__init__.py

@@ -9,8 +9,11 @@
     _make_fast_predict_tensors,
     _batch_features_tensor,
     _crosswise_differences,
+    _crosswise_similarity,
+    _out_similarity,
     _crosswise_tensor,
     _pairwise_differences,
+    _pairwise_similarity,
     _pairwise_tensor,
     _fast_nn_update,
     _make_heteroscedastic_tensor,
@@ -21,8 +24,11 @@
     "_make_fast_predict_tensors",
     "_batch_features_tensor",
     "_crosswise_differences",
+    "_crosswise_similarity",
+    "_out_similarity",
     "_crosswise_tensor",
     "_pairwise_differences",
+    "_pairwise_similarity",
     "_pairwise_tensor",
     "_fast_nn_update",
     "_make_heteroscedastic_tensor",

src/MuyGPyS/_src/gp/tensors/jax.py

@@ -87,6 +87,32 @@ def _pairwise_tensor(
     return _pairwise_differences(points)
 
 
+@jit
+def _crosswise_similarity(
+    data: jnp.ndarray,
+    nn_data: jnp.ndarray,
+    data_indices: jnp.ndarray,
+    nn_indices: jnp.ndarray,
+):
+    raise NotImplementedError("JAX backend not yet supported for similarity tensors")
+
+
+@jit
+def _pairwise_similarity(
+    data: jnp.ndarray,
+    nn_indices: jnp.ndarray,
+):
+    raise NotImplementedError("JAX backend not yet supported for similarity tensors")
+
+
+@jit
+def _out_similarity(
+    data: jnp.ndarray,
+    data_indices: jnp.ndarray
+):
+    raise NotImplementedError("JAX backend not yet supported for similarity tensors")
+
+
 @jit
 def _F2(diffs: jnp.ndarray) -> jnp.ndarray:
     return jnp.sum(diffs**2, axis=-1)

src/MuyGPyS/_src/gp/tensors/mpi.py

@@ -93,6 +93,29 @@ def _pairwise_differences(points: np.ndarray) -> np.ndarray:
     )
 
 
+def _crosswise_similarity(
+    data: np.ndarray,
+    nn_data: np.ndarray,
+    data_indices: np.ndarray,
+    nn_indices: np.ndarray,
+):
+    raise NotImplementedError("MPI backend not yet supported for similarity tensors")
+
+
+def _pairwise_similarity(
+    data: np.ndarray,
+    nn_indices: np.ndarray,
+):
+    raise NotImplementedError("MPI backend not yet supported for similarity tensors")
+
+
+def _out_similarity(
+    data: np.ndarray,
+    data_indices: np.ndarray
+):
+    raise NotImplementedError("MPI backend not yet supported for similarity tensors")
+
+
 def _fast_nn_update(
     train_nn_indices: np.ndarray,
 ) -> np.ndarray:

src/MuyGPyS/_src/gp/tensors/numpy.py

@@ -86,6 +86,66 @@ def _pairwise_differences(points: np.ndarray) -> np.ndarray:
         raise ValueError(f"points shape {points.shape} is not supported.")
 
 
+def _crosswise_similarity(
+    data: np.ndarray,
+    nn_data: np.ndarray,
+    data_indices: np.ndarray,
+    nn_indices: np.ndarray,
+) -> np.ndarray:
+    locations = data[data_indices]
+    points = nn_data[nn_indices].swapaxes(2, 1)
+
+    # working implementation without einsum
+    # dot = np.sum(
+    #     locations[:, None, :, None, :, None, :, None, :]
+    #     * points[:, :, None, :, None, :, None, :, :],
+    #     axis=-1,
+    # ).swapaxes(1, 2)
+    # shape = dot.shape
+
+    # locations.shape = (i, x, d, a, q)
+    # points.shape = (i, y, k, e, b, q)
+    dot = np.einsum("ixdaq, iykebq -> iykxdeab", locations, points)
+
+    crosswise_similarity = dot.reshape(*dot.shape[:-4], -1, *dot.shape[-2:])
+
+    return crosswise_similarity
+
+
+def _pairwise_similarity(
+    data: np.ndarray,
+    nn_indices: np.ndarray,
+) -> np.ndarray:
+    points = data[nn_indices].swapaxes(2, 1)
+
+    # working implementation without einsum
+    # dot = np.sum(
+    #     points[:, :, :, None, None, :, None, :, None, :]
+    #     * points[:, None, None, :, :, None, :, None, :, :],
+    #     axis=-1,
+    # )
+
+    # points.shape=(i, x, k, d, a, q) / (i, y, l, e, b, q)
+    dot = np.einsum("ixkdaq,iylebq->ixkyldeab", points, points)
+
+    pairwise_similarity = dot.reshape(*dot.shape[:5], -1, *dot.shape[-2:])
+
+    return pairwise_similarity
+
+
+def _out_similarity(
+    data: np.ndarray,
+    data_indices: np.ndarray
+) -> np.ndarray:
+    points = data[data_indices]
+
+    dot = np.einsum("ixdaq, iyebq -> ixydeab", points, points)
+
+    out_similarity = dot.reshape(*dot.shape[:3], -1, *dot.shape[-2:])
+
+    return out_similarity
+
+
 def _F2(diffs: np.ndarray) -> np.ndarray:
     return np.sum(diffs**2, axis=-1)
 

src/MuyGPyS/_src/gp/tensors/torch.py

@@ -78,6 +78,29 @@ def _pairwise_differences(points: torch.ndarray) -> torch.ndarray:
         raise ValueError(f"points shape {points.shape} is not supported.")
 
 
+def _crosswise_similarity(
+    data: torch.ndarray,
+    nn_data: torch.ndarray,
+    data_indices: torch.ndarray,
+    nn_indices: torch.ndarray,
+):
+    raise NotImplementedError("Torch backend not yet supported for similarity tensors")
+
+
+def _pairwise_similarity(
+    data: torch.ndarray,
+    nn_indices: torch.ndarray,
+):
+    raise NotImplementedError("Torch backend not yet supported for similarity tensors")
+
+
+def _out_similarity(
+    data: torch.ndarray,
+    data_indices: torch.ndarray
+):
+    raise NotImplementedError("Torch backend not yet supported for similarity tensors")
+
+
 def _F2(diffs: torch.ndarray) -> torch.ndarray:
     return torch.sum(diffs**2, axis=-1)