add real_gaussian_integral function

mrmustard/physics/utils.py

@@ -0,0 +1,78 @@
+# Copyright 2024 Xanadu Quantum Technologies Inc.
+
+# 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.
+
+"""
+This module contains helper functions for physics.
+"""
+import numpy as np
+
+from mrmustard import math
+
+
+def real_gaussian_integral(
+    Abc: tuple,
+    idx: tuple[int, ...],
+):
+    r"""Computes the Gaussian integral of the exponential of a real quadratic form.
+    The integral is defined as (note that in general we integrate over a subset of m dimensions):
+
+    :math:`\int_{R^m} F(x) dx`
+
+    where
+
+    :math:`F(x) = \textrm{exp}(0.5 x^T A x + b^T x)`
+
+    Here, ``x`` is an ``n``-dim real vector, ``A`` is an ``n x n`` real matrix,
+    ``b`` is an ``n``-dim real vector, ``c`` is a real scalar. The integral indices
+    are specified by ``idx``.
+
+    Arguments:
+        A,b,c: the ``(A,b,c)`` triple
+        idx: the tuple of indices of the x variables
+
+    Returns:
+        The ``(A,b,c)`` triple of the result of the integral.
+    """
+    A, b, c = Abc
+
+    n = len(idx)
+
+    if not idx:
+        return A, b, c
+    not_idx = tuple(i for i in range(A.shape[-1]) if i not in idx)
+
+    I = math.eye(n, dtype=A.dtype)
+    Z = math.zeros((n, n), dtype=A.dtype)
+    X = math.block([[Z, I], [I, Z]])
+    M = math.gather(math.gather(A, idx, axis=-1), idx, axis=-2) + X
+    bM = math.gather(b, idx, axis=-1)
+
+    not_idx = tuple(i for i in range(A.shape[-1]) if i not in idx)
+    if math.asnumpy(not_idx).shape != (0,):
+        D = math.gather(math.gather(A, idx, axis=-1), not_idx, axis=-2)
+        R = math.gather(math.gather(A, not_idx, axis=-1), not_idx, axis=-2)
+        bR = math.gather(b, not_idx, axis=-1)
+        A_post = R - math.matmul(D, math.inv(M), math.transpose(D))
+        b_post = bR - math.matvec(D, math.solve(M, bM))
+    else:
+        A_post = math.astensor([])
+        b_post = math.astensor([])
+
+    c_post = (
+        c
+        * math.sqrt((2 * np.pi) ** n / math.det(M))
+        * math.exp(-0.5 * math.sum(bM * math.solve(M, bM)))
+    )
+
+    return A_post, b_post, c_post