Repository: scipy/scipy

Issue:
Orthogonal Array for Latin hypercube in `scipy.stats.qmc`
## Motivations

As discussed in #13654 #13647, we need to implement a generator of orthogonal arrays (OAs) to support the orthogonal array based Latin hypercube sampling (LHS), which is a performant version of naive Latin hypercube (`scipy.stats.qmc.LatinHypercube`). However, there are so many algorithms to generate OAs that developers are not sure which one to adapt in Scipy at the moment. 

In this issue, I would like to have a discussion of the algorithms for generating OAs, which will be used to support OA-based Latin hypercube in Scipy. 


## Quick Summary of Discussions in #13654 #13647

### Orthogonal Arrays based Latin Hypercube  sampling(OA-LHS)
Good references for OA-LHS are 
- Owen, A. B. (1992). Orthogonal arrays for computer experiments, integration and visualization. Statistica Sinica, 439-452.
- Tang, B. (1993). Orthogonal array-based Latin hypercubes. Journal of the American statistical association, 88(424), 1392-1397.

In those articles, it is assumed that they already have OAs at their hands. By assuming that we have a function `get_orthogonal_array` that generates an orthogonal array, the OA-LHS algorithm can be written as follows.
```python
def SampleOALH(n, d):
    sample = self.rg_sample((n, d))
    # For example, 
    # symbols = [1, 2, 3, ..., n_grid_size] and 
    # oa is an array of size (n, d) whose all elements are one of [1, 2, 3, ..., n_grid_size] (=: symbols).
    # oa is designed so that its samples (rows) are more uniform (in a sense) than random samples.
    oa, symbols = get_orthogonal_array(n, d)  
    
    # Randomization of oa[:, :] to generate OA-LHSs
    for i in range(d):
        # permutate symbols, not the elements
        perm = self.rng.shuffle(len(symbols))
        oa[:, i] = perm[oa[:i]]

    oa_lhs = (oa + sample) / len(symbols)
    return oa_lhs
```

### Orthogonal Arrays (OAs)
OAs are briefly introduced in the article above. Also, it has a Wikipedia page on it
- https://en.wikipedia.org/wiki/Orthogonal_array

The generation of orthogonal arrays is a bit complicated business as there is not a single algorithm that can generate arbitrarily OAs, as discussed below  
- https://stackoverflow.com/questions/37851038/how-to-create-orthogonal-array

Additionally, to make matters worse, OAs do not exist for a certain combination of sample size, a dimension of sample space, and strength, which is a parameter of OAs. 

Thus, there are also approximation algorithms for generating nearly orthogonal arrays. One promising algorithm would be the following, as it is recommended by the original author in https://github.com/scipy/scipy/issues/13647#issuecomment-792378238
- Owen, A. B. (1994). Controlling correlations in Latin hypercube samples. Journal of the American Statistical Association, 89(428), 1517-1522.



Discussion:
Thank you for this great description!
Art Owen told me that there are actually explicit implementations. See _A. S. HedayatN. J. A. SloaneJohn Stufken 1999. "Orthogonal Arrays Theory and Applications. DOI: 10.1007/978-1-4612-1478-6"_ for many examples.
I tried to implement the RGS algorithm that prof Owen mentioned ([my implementation is available here](https://gist.github.com/kstoneriv3/3daddb7b05942a1b8282e556ddf03fec)), but it did not work at all. It is supposed to produce samples in a unit cube that is as uncorrelated as possible but the magnitude of correlation was almost the same as naive Latin hypercube...

RGS takes samples generated as LHS and applies a kind of Gram-Schmidt orthogonalization to reduce correlations among dimensions of samples, but the orthogonalization step did not improve the quality of sequences as far as I could see (though I am not sure if my implementation is correct).

FYI @alegresor this could interest you (or someone of your group) as I see you are maintaining QMCPy.
@kstoneriv3 I had some exchanges with Art and I wrote the following which is successfully producing an orthogonal LHS:

```python
import numpy as np
from scipy.stats.qmc import LatinHypercube

p = 5
d = 2  # note that this cannot be changed and is not the dimension of the final sample.

arrays = np.tile(np.arange(p), (2, 1))
oa_sample = np.zeros(shape=(p**2, p+1))
oa_sample[:, :2] = np.stack(np.meshgrid(*arrays), axis=-1).reshape(-1, d)
for p_ in range(1, p):
    oa_sample[:, 2+p_-1] = np.mod(oa_sample[:, 0] + p_*oa_sample[:, 1], p)

for j in range(p+1):
    np.random.shuffle(oa_sample[:, j])

# sample is a randomized OA from now
# and the following is making it an OA-LHS

oa_lhs_sample = np.zeros(shape=(p**2, p+1))
for j in range(p+1):
    for k in range(p):
        idx = np.where(oa_sample[:, j] == k)[0]
        lhs = LatinHypercube(d=1).random(p).flatten()
        oa_lhs_sample[:, j][idx] = lhs + oa_sample[:, j][idx]

oa_lhs_sample /= p
```

I will clean this snippet and streamline it as some operations can be made more efficient, but the idea is here. It creates an OA of strength 2. With the constraint that `p` must be a prime number. So you cannot freely use any number of samples. It produces a sample of shape (p**2, p+1). What can be done though is to have the dimension you want from 1 to p+1 selecting any subset. 
That's great news! I read part of your code (I did not follow the construction of the OA (first half of the code)) and have a few comments.
- Shouldn't columns be shuffled using the same random permutation? If you permutate each column independently, the resulting array might not be orthogonal with strength 2 anymore by a chance (e.g. `[[0, 0], [0, 1], [1, 0], [1, 1]]` is an 4x2 OA but a there are permutations of columns that change this OA to `[[0, 0], [0, 0], [1, 1], [1, 1]]`, which is no longer an OA. Owen, A. B. (1992) uses independent random permutation for each column so probably the independent permutation is correct but I am not sure about this.... My intuition is that if you want to preserve the orthogonal property of the array after permutation, you should keep the content of each row.
- Is `sample` and `oa_sample` the same thing in the code above? It seems that `sample` is not defined anywhere.
> * Shouldn't columns be shuffled using the same random permutation? If you permutate each column independently, the resulting array might not be orthogonal with strength 2 anymore by a chance (e.g. `[[0, 0], [0, 1], [1, 0], [1, 1]]` is an 4x2 OA but a there are permutations of columns that change this OA to `[[0, 0], [0, 0], [1, 1], [1, 1]]`, which is no longer an OA. Owen, A. B. (1992) uses independent random permutation for each column so probably the independent permutation is correct but I am not sure about this.... My intuition is that if you want to preserve the orthogonal property of the array after permutation, you should keep the content of each row.

I find this strange as well, but keeping the same permutations for each column would just mean changing the order between rows which has no impact. I can forward your question if you want. Art told me that not scrambling would produce bad samplings (which I could see afterwards).

> * Is `sample` and `oa_sample` the same thing in the code above? It seems that `sample` is not defined anywhere.

Indeed, it was a typo, I fixed it.

I just checked and the functions even works with the `centered` parameter. So I might just add a parameter to use an OA. It could also be automatic, if we can use an OA (Ns compatible), otherwise we use a normal LHS. Soon there will be an optimized version of LHS. If that's of interest for you, check #13471
> > * Shouldn't columns be shuffled using the same random permutation? If you permutate each column independently, the resulting array might not be orthogonal with strength 2 anymore by a chance (e.g. `[[0, 0], [0, 1], [1, 0], [1, 1]]` is an 4x2 OA but a there are permutations of columns that change this OA to `[[0, 0], [0, 0], [1, 1], [1, 1]]`, which is no longer an OA. Owen, A. B. (1992) uses independent random permutation for each column so probably the independent permutation is correct but I am not sure about this.... My intuition is that if you want to preserve the orthogonal property of the array after permutation, you should keep the content of each row.
> 
> I find this strange as well, but keeping the same permutations for each column would just mean changing the order between rows which has no impact. I can forward your question if you want. Art told me that not scrambling would produce bad samplings (which I could see afterwards).
> 

Then, could you ask Prof Owen about this? Independent random permutation of each column should only produce LHS of strength 1, but not 2. I feel that there are better ways to randomize an OA, such as permutation of columns, shifting points from the origin by a random constant vector, and flipping the sign of randomly picked columns (and taking mod 1.0 afterward). Maybe the scrambling is the smartest of operations of the kind.

> > * Is `sample` and `oa_sample` the same thing in the code above? It seems that `sample` is not defined anywhere.
> 
> Indeed, it was a typo, I fixed it.
> 
> I just checked and the functions even works with the `centered` parameter. So I might just add a parameter to use an OA. It could also be automatic, if we can use an OA (Ns compatible), otherwise we use a normal LHS. Soon there will be an optimized version of LHS. If that's of interest for you, check #13471

Thanks! That makes sense. #13471 seems interesting but I guess I cannot catch up with the thread very soon. I might pop into the thread if I can.

> Then, could you ask Prof Owen about this? Independent random permutation of each column should only produce LHS of strength 1, but not 2. I feel that there are better ways to randomize an OA, such as permutation of columns, shifting points from the origin by a random constant vector, and flipping the sign of randomly picked columns (and taking mod 1.0 afterward). Maybe the scrambling is the smartest of operations of the kind.

Sure, I will let you know. I should also be able to do a PR soon for this OALHS and will ping you. We could add more randomization method if that's helpful.

> Thanks! That makes sense. #13471 seems interesting but I guess I cannot catch up with the thread very soon. I might pop into the thread if I can.

No problem. As I said for now it will be very "simple" and we will soon merge. I could not make the famous ESE work, it was not performing any better than doing n-random permutations for me. If I missed something and we can make it work, it can be added later as I left the structure to add other methods.
@kstoneriv3 I was doing it wrong! I should permute the levels in a column, not the elements.

From Art:
> If you scramble the levels within each column of an OA it keeps it's strength.  If you scramble the n values, then you break the strength.  E.g., if you have 121 numbers in 0,1,2,...,10 then scramble those 11 levels among themselves.  Don't scramble the 121 values in the column.

Resolving patch:
diff --git a/scipy/stats/_qmc.py b/scipy/stats/_qmc.py
--- a/scipy/stats/_qmc.py
+++ b/scipy/stats/_qmc.py
@@ -885,16 +885,12 @@ def random(self, n: IntNumber = 1) -> np.ndarray:
 
 
 class LatinHypercube(QMCEngine):
-    """Latin hypercube sampling (LHS).
+    r"""Latin hypercube sampling (LHS).
 
     A Latin hypercube sample [1]_ generates :math:`n` points in
     :math:`[0,1)^{d}`. Each univariate marginal distribution is stratified,
     placing exactly one point in :math:`[j/n, (j+1)/n)` for
     :math:`j=0,1,...,n-1`. They are still applicable when :math:`n << d`.
-    LHS is extremely effective on integrands that are nearly additive [2]_.
-    LHS on :math:`n` points never has more variance than plain MC on
-    :math:`n-1` points [3]_. There is a central limit theorem for LHS [4]_,
-    but not necessarily for optimized LHS.
 
     Parameters
     ----------
@@ -914,6 +910,15 @@ class LatinHypercube(QMCEngine):
 
         .. versionadded:: 1.8.0
 
+    strength : {1, 2}, optional
+        Strength of the LHS. ``strength=1`` produces a plain LHS while
+        ``strength=2`` produces an orthogonal array based LHS of strength 2
+        [7]_, [8]_. In that case, only ``n=p**2`` points can be sampled,
+        with ``p`` a prime number. It also constrains ``d <= p + 1``.
+        Default is 1.
+
+        .. versionadded:: 1.8.0
+
     seed : {None, int, `numpy.random.Generator`}, optional
         If `seed` is None the `numpy.random.Generator` singleton is used.
         If `seed` is an int, a new ``Generator`` instance is used,
@@ -921,6 +926,36 @@ class LatinHypercube(QMCEngine):
         If `seed` is already a ``Generator`` instance then that instance is
         used.
 
+    Notes
+    -----
+
+    When LHS is used for integrating a function :math:`f` over :math:`n`,
+    LHS is extremely effective on integrands that are nearly additive [2]_.
+    With a LHS of :math:`n` points, the variance of the integral is always
+    lower than plain MC on :math:`n-1` points [3]_. There is a central limit
+    theorem for LHS on the mean and variance of the integral [4]_, but not
+    necessarily for optimized LHS due to the randomization.
+
+    :math:`A` is called an orthogonal array of strength :math:`t` if in each
+    n-row-by-t-column submatrix of :math:`A`: all :math:`p^t` possible
+    distinct rows occur the same number of times. The elements of :math:`A`
+    are in the set :math:`\{0, 1, ..., p-1\}`, also called symbols.
+    The constraint that :math:`p` must be a prime number is to allow modular
+    arithmetic.
+
+    Strength 1 (plain LHS) brings an advantage over strength 0 (MC) and
+    strength 2 is a useful increment over strength 1. Going to strength 3 is
+    a smaller increment and scrambled QMC like Sobol', Halton are more
+    performant [7]_.
+
+    To create a LHS of strength 2, the orthogonal array :math:`A` is
+    randomized by applying a random, bijective map of the set of symbols onto
+    itself. For example, in column 0, all 0s might become 2; in column 1,
+    all 0s might become 1, etc.
+    Then, for each column :math:`i` and symbol :math:`j`, we add a plain,
+    one-dimensional LHS of size :math:`p` to the subarray where
+    :math:`A^i = j`. The resulting matrix is finally divided by :math:`p`.
+
     References
     ----------
     .. [1] Mckay et al., "A Comparison of Three Methods for Selecting Values
@@ -937,6 +972,10 @@ class LatinHypercube(QMCEngine):
     .. [6] Damblin et al., "Numerical studies of space filling designs:
        optimization of Latin Hypercube Samples and subprojection properties."
        Journal of Simulation, 2013.
+    .. [7] A. B. Owen , "Orthogonal arrays for computer experiments,
+       integration and visualization." Statistica Sinica, 1992.
+    .. [8] B. Tang, "Orthogonal Array-Based Latin Hypercubes."
+       Journal of the American Statistical Association, 1993.
 
     Examples
     --------
@@ -976,32 +1015,61 @@ class LatinHypercube(QMCEngine):
     >>> qmc.discrepancy(sample)
     0.0176...  # random
 
+    Use the `strength` keyword argument to produce an orthogonal array based
+    LHS of strength 2. In this case, the number of sample points must be the
+    square of a prime number.
+
+    >>> sampler = qmc.LatinHypercube(d=2, strength=2)
+    >>> sample = sampler.random(n=9)
+    >>> qmc.discrepancy(sample)
+    0.00526...  # random
+
+    Options could be combined to produce an optimized centered
+    orthogonal array based LHS. After optimization, the result would not
+    be guaranteed to be of strength 2.
+
     """
 
     def __init__(
         self, d: IntNumber, *, centered: bool = False,
+        strength: int = 1,
         optimization: Optional[Literal["random-cd"]] = None,
         seed: SeedType = None
     ) -> None:
         super().__init__(d=d, seed=seed)
         self.centered = centered
 
-        lhs_methods: Dict[Optional[Literal["random-cd"]], Callable] = {
-            None: self._random,
-            "random-cd": self._random_cd,
+        lhs_method_strength = {
+            1: self._random,
+            2: self._random_oa_lhs
         }
 
         try:
-            if optimization is not None:
-                optimization = optimization.lower()  # type: ignore[assignment]
-            self.lhs_method = lhs_methods[optimization]
-        except KeyError:
-            raise ValueError(f"{optimization!r} is not a valid optimization"
-                             " method. It must be one of"
-                             f" {set(lhs_methods)!r}")
+            self.lhs_method = lhs_method_strength[strength]
+        except KeyError as exc:
+            message = (f"{strength!r} is not a valid strength. It must be one"
+                       f" of {set(lhs_method_strength)!r}")
+            raise ValueError(message) from exc
 
-        self._n_nochange = 100
-        self._n_iters = 10_000
+        optimization_method: Dict[Literal["random-cd"], Callable] = {
+            "random-cd": self._random_cd,
+        }
+
+        self.optimization_method: Optional[Callable]
+        if optimization is not None:
+            try:
+                optimization = optimization.lower()  # type: ignore[assignment]
+                self.optimization_method = optimization_method[optimization]
+            except KeyError as exc:
+                message = (f"{optimization!r} is not a valid optimization"
+                           f" method. It must be one of"
+                           f" {set(optimization_method)!r}")
+                raise ValueError(message) from exc
+
+            self._n_nochange = 100
+            self._n_iters = 10_000
+        else:
+            self.optimization_method = None
 
     def random(self, n: IntNumber = 1) -> np.ndarray:
         """Draw `n` in the half-open interval ``[0, 1)``.
@@ -1017,7 +1085,12 @@ def random(self, n: IntNumber = 1) -> np.ndarray:
             LHS sample.
 
         """
-        return self.lhs_method(n)
+        lhs = self.lhs_method(n)
+        if self.optimization_method is not None:
+            lhs = self.optimization_method(lhs)
+
+        self.num_generated += n
+        return lhs
 
     def _random(self, n: IntNumber = 1) -> np.ndarray:
         """Base LHS algorithm."""
@@ -1033,10 +1106,56 @@ def _random(self, n: IntNumber = 1) -> np.ndarray:
         perms = perms.T
 
         samples = (perms - samples) / n
-        self.num_generated += n
         return samples
 
-    def _random_cd(self, n: IntNumber = 1) -> np.ndarray:
+    def _random_oa_lhs(self, n: IntNumber = 4) -> np.ndarray:
+        """Orthogonal array based LHS of strength 2."""
+        p = np.sqrt(n).astype(int)
+        n_row = p**2
+        n_col = p + 1
+
+        primes = primes_from_2_to(p + 1)
+        if p not in primes or n != n_row:
+            raise ValueError(
+                "n is not the square of a prime number. Close"
+                f" values are {primes[-2:]**2}"
+            )
+        if self.d > p + 1:
+            raise ValueError("n is too small for d. Must be n > (d-1)**2")
+
+        oa_sample = np.zeros(shape=(n_row, n_col), dtype=int)
+
+        # OA of strength 2
+        arrays = np.tile(np.arange(p), (2, 1))
+        oa_sample[:, :2] = np.stack(np.meshgrid(*arrays),
+                                    axis=-1).reshape(-1, 2)
+        for p_ in range(1, p):
+            oa_sample[:, 2+p_-1] = np.mod(oa_sample[:, 0]
+                                          + p_*oa_sample[:, 1], p)
+
+        # scramble the OA
+        oa_sample_ = np.empty(shape=(n_row, n_col), dtype=int)
+        for j in range(n_col):
+            perms = self.rng.permutation(p)
+            oa_sample_[:, j] = perms[oa_sample[:, j]]
+
+        # following is making a scrambled OA into an OA-LHS
+        oa_lhs_sample = np.zeros(shape=(n_row, n_col))
+        lhs_engine = LatinHypercube(d=1, centered=self.centered, strength=1,
+                                    seed=self.rng)  # type: QMCEngine
+        for j in range(n_col):
+            for k in range(p):
+                idx = oa_sample[:, j] == k
+                lhs = lhs_engine.random(p).flatten()
+                oa_lhs_sample[:, j][idx] = lhs + oa_sample[:, j][idx]
+
+                lhs_engine = lhs_engine.reset()
+
+        oa_lhs_sample /= p
+
+        return oa_lhs_sample[:, :self.d]  # type: ignore[misc]
+
+    def _random_cd(self, best_sample: np.ndarray) -> np.ndarray:
         """Optimal LHS on CD.
 
         Create a base LHS and do random permutations of coordinates to
@@ -1048,10 +1167,11 @@ def _random_cd(self, n: IntNumber = 1) -> np.ndarray:
         `_n_iters` iterations are performed; or if there is no improvement
         for `_n_nochange` consecutive iterations.
         """
+        n = len(best_sample)
+
         if self.d == 0 or n == 0:
             return np.empty((n, self.d))
 
-        best_sample = self._random(n=n)
         best_disc = discrepancy(best_sample)
 
         if n == 1:
@@ -1081,7 +1201,6 @@ def _random_cd(self, n: IntNumber = 1) -> np.ndarray:
             else:
                 n_nochange += 1
 
-        self.num_generated += n
         return best_sample
 
 
@@ -1331,7 +1450,7 @@ def fast_forward(self, n: IntNumber) -> Sobol:
 
         Returns
         -------
-        engine: Sobol
+        engine : Sobol
             The fast-forwarded engine.
 
         """


Test patch:
diff --git a/scipy/stats/tests/test_qmc.py b/scipy/stats/tests/test_qmc.py
--- a/scipy/stats/tests/test_qmc.py
+++ b/scipy/stats/tests/test_qmc.py
@@ -1,5 +1,6 @@
 import os
 from collections import Counter
+from itertools import combinations, product
 
 import pytest
 import numpy as np
@@ -534,23 +535,41 @@ def test_sample(self, *args):
         pytest.skip("Not applicable: the value of reference sample is"
                     " implementation dependent.")
 
+    @pytest.mark.parametrize("strength", [1, 2])
+    @pytest.mark.parametrize("centered", [False, True])
     @pytest.mark.parametrize("optimization", [None, "random-CD"])
-    def test_sample_stratified(self, optimization):
-        d, n = 4, 20
+    def test_sample_stratified(self, optimization, centered, strength):
+        seed = np.random.RandomState(123456)
+        p = 5
+        n = p**2
+        d = 6
         expected1d = (np.arange(n) + 0.5) / n
         expected = np.broadcast_to(expected1d, (d, n)).T
 
-        engine = self.engine(d=d, scramble=False, centered=True,
-                             optimization=optimization)
+        engine = qmc.LatinHypercube(d=d, centered=centered,
+                                    strength=strength,
+                                    optimization=optimization,
+                                    seed=seed)
         sample = engine.random(n=n)
+        assert sample.shape == (n, d)
+        assert engine.num_generated == n
+
         sorted_sample = np.sort(sample, axis=0)
 
-        assert_equal(sorted_sample, expected)
         assert np.any(sample != expected)
-
         assert_allclose(sorted_sample, expected, atol=0.5 / n)
         assert np.any(sample - expected > 0.5 / n)
 
+        if strength == 2 and optimization is None:
+            unique_elements = np.arange(p)
+            desired = set(product(unique_elements, unique_elements))
+
+            for i, j in combinations(range(engine.d), 2):
+                samples_2d = sample[:, [i, j]]
+                res = (samples_2d * p).astype(int)
+                res_set = set((tuple(row) for row in res))
+                assert_equal(res_set, desired)
+
     def test_discrepancy_hierarchy(self):
         seed = np.random.RandomState(123456)
         lhs = qmc.LatinHypercube(d=2, seed=seed)
@@ -564,12 +583,32 @@ def test_discrepancy_hierarchy(self):
 
         assert disc_ < disc_ref
 
-    def test_warning(self):
-        with pytest.raises(ValueError, match=r"'toto' is not a valid"
-                                             r" optimization method"):
-            seed = np.random.RandomState(12345)
+    def test_raises(self):
+        seed = np.random.RandomState(12345)
+
+        message = r"'toto' is not a valid optimization method"
+        with pytest.raises(ValueError, match=message):
             qmc.LatinHypercube(1, seed=seed, optimization="toto")
 
+        message = r"not a valid strength"
+        with pytest.raises(ValueError, match=message):
+            qmc.LatinHypercube(1, strength=3, seed=seed)
+
+        message = r"n is not the square of a prime number"
+        with pytest.raises(ValueError, match=message):
+            engine = qmc.LatinHypercube(d=2, strength=2, seed=seed)
+            engine.random(16)
+
+        message = r"n is not the square of a prime number"
+        with pytest.raises(ValueError, match=message):
+            engine = qmc.LatinHypercube(d=2, strength=2, seed=seed)
+            engine.random(5)  # because int(sqrt(5)) would result in 2
+
+        message = r"n is too small for d"
+        with pytest.raises(ValueError, match=message):
+            engine = qmc.LatinHypercube(d=5, strength=2, seed=seed)
+            engine.random(9)
+
 
 class TestSobol(QMCEngineTests):
     qmce = qmc.Sobol