Added batch generation of multiple samples

Author: Rolf Johan Lorentzen

src/geostat/gaussian_sim.py

@@ -4,10 +4,9 @@
 from scipy.linalg import toeplitz
 
 
-def fast_gaussian(dimension, sdev, corr):
-
-    # print('Fast Gaussian!')
+def fast_gaussian(dimension, sdev, corr, num_samples=1):
     """
+
     Generates random vector from distribution satisfying Gaussian variogram in dimension up to 3-d.
 
     Parameters
@@ -23,10 +22,14 @@ def fast_gaussian(dimension, sdev, corr):
         If a single float is provided, it represents the correlation length in all directions.
         If an array-like object with length 3 is provided, it represents the correlation length in the x-, y-, and z-directions.
 
+    num_samples : int, optional
+        Number of samples to generate. Default is 1.
+        If greater than 1, the function will return an array of shape (dimension, num_samples).
+
     Returns
     -------
     x : array-like
-        The generated random vector.
+        The generated random vectors of shape (dimension, num_samples).
 
     Notes
     -----
@@ -39,15 +42,15 @@ def fast_gaussian(dimension, sdev, corr):
     Want to generate a field on a 3-d grid with dimension m x n x p, with correlation length a along first coordinate
     axis, b along second coordinate axis, c alone third coordinate axis, and standard deviation sigma:
 
-    >>> x=fast_gaussian(np.array([m, n, p]),np.array([sigma]),np.array([a b c]))
+    x=fast_gaussian(np.array([m, n, p]),np.array([sigma]),np.array([a b c]))
 
     If the dimension is n x 1 one can write
 
-    >>> x=fast_gaussian(np.array([n]),np.array([sigma]),np.array([a]))
+    x=fast_gaussian(np.array([n]),np.array([sigma]),np.array([a]))
 
     If the correlation length is the same in all directions:
 
-    >>> x=fast_gaussian(np.array([m n p]),np.array([sigma]),np.array([a]))
+    x=fast_gaussian(np.array([m n p]),np.array([sigma]),np.array([a]))
 
     The properties on the Kronecker product behind this algorithm can be found in
     Horn & Johnson: Topics in Matrix Analysis, Cambridge UP, 1991.
@@ -58,9 +61,9 @@ def fast_gaussian(dimension, sdev, corr):
     Also note that reshape with order='F' is used to keep the code identical to the Matlab code.
 
     The method was invented and implemented in Matlab by Geir Nævdal in 2011.
+    Memory-efficient implementation for batch generation of samples was added in 2025.
     """
 
-    # Initialize dimension:
     if len(dimension) == 0:
         sys.exit("fast_gaussian: Wrong input, dimension should have length at least 1")
     m = dimension[0]
@@ -75,13 +78,11 @@ def fast_gaussian(dimension, sdev, corr):
     if len(dimension) > 3:
         sys.exit("fast_gaussian: Wrong input, dimension should have length at most 3")
 
-    # Compute standard deviation
-    if len(sdev) > 1:  # check input
+    if len(sdev) > 1:
         std = 1
     else:
-        std = sdev  # the variance will come out through the kronecker product.
+        std = sdev
 
-    # Initialize correlation length
     if len(corr) == 0:
         sys.exit("fast_gaussian: Wrong input, corr should have length at least 1")
     if len(corr) == 1:
@@ -90,60 +91,52 @@ def fast_gaussian(dimension, sdev, corr):
         corr = np.append(corr, corr[1])
     corr = np.maximum(corr, 1)
 
-    # first generate the covariance matrix for first dimension
-    dist1 = np.arange(m)
-    dist1 = dist1 / corr[0]
+    dist1 = np.arange(m) / corr[0]
     t1 = toeplitz(dist1)
-    # to avoid problem with Cholesky factorization when the matrix is close to singular we add a small number on the
-    # diagonal entries
     t1 = std * np.exp(-t1 ** 2) + 1e-10 * np.eye(m)
-    # Cholesky decomposition
     cholt1 = np.linalg.cholesky(t1)
 
-    # generate the covariance matrix for the second dimension
-    # to save time - use a copy if possible
     if corr[0] == corr[1] and n == m:
         cholt2 = cholt1
     else:
-        dist2 = np.arange(n)
-        dist2 = dist2 / corr[1]
+        dist2 = np.arange(n) / corr[1]
         t2 = toeplitz(dist2)
         t2 = std * np.exp(-t2 ** 2) + 1e-10 * np.eye(n)
         cholt2 = np.linalg.cholesky(t2)
 
-    # generate the covariance matrix for the third dimension if required
     cholt3 = None
     if p is not None:
-        # use std = 1 to get the correct value
-        dist3 = np.arange(p)
-        dist3 = dist3 / corr[2]
+        dist3 = np.arange(p) / corr[2]
         t3 = toeplitz(dist3)
         t3 = np.exp(-t3 ** 2) + 1e-10 * np.eye(p)
         cholt3 = np.linalg.cholesky(t3)
 
-    # draw a random variable
-    x = np.random.randn(dim, 1)
+    z = np.random.randn(dim, num_samples)
 
-    # adjust to get the correct covariance matrix, applying Lemma 4.3.1. in Horn & Johnson:
-    # we need to adjust to get the correct covariance matrix
-    if p is None:  # 2-d
-        x = np.dot(np.dot(cholt1, np.reshape(x, (m, n))), cholt2.T)
-    else:  # 3-d
-        # either dimension 1 and 2 or 2 and 3 need to be grouped together.
+    # Memory-efficient multiplication without explicit large Kronecker product
+    if p is None:
+        x = z.reshape(m, n, num_samples, order='F')
+        x = np.tensordot(cholt1, x, axes=([1], [0]))
+        x = np.tensordot(cholt2, x, axes=([1], [1]))
+    else:
+        x = z.reshape(m, n, p, num_samples, order='F')
         if n <= p:
-            x = np.dot(np.dot(np.kron(cholt2, cholt1),
-                       np.reshape(x, (m * n, p))), cholt3.T)
+            x = np.tensordot(cholt1, x, axes=([1], [0]))
+            x = np.tensordot(cholt2, x, axes=([1], [1]))
+            x = np.tensordot(cholt3, x, axes=([1], [2]))
         else:
-            x = np.dot(np.dot(cholt1, np.reshape(x, (m, n * p))),
-                       np.kron(cholt3.T, cholt2.T))
+            x = np.tensordot(cholt1, x, axes=([1], [0]))
+            x = np.tensordot(cholt2, x, axes=([1], [1]))
+            x = np.tensordot(cholt3, x, axes=([1], [2]))
 
-    # reshape back
-    x = np.reshape(x, (dim,), order='F')
+    # Reshape back to (dim, num_samples) (order='C' is used here to match the original function's output)
+    x = x.reshape((dim, num_samples), order='C')
 
     if len(sdev) > 1:
-        if len(sdev) == len(x):
-            x = sdev * x
+        if len(sdev) == dim:
+            x = sdev[:, None] * x
         else:
             sys.exit('fast_gaussian: Inconsistent dimension of sdev')
 
     return x
+