From 12ac2fd98b4671eda288c92b34e51befd1785ba4 Mon Sep 17 00:00:00 2001
From: Feda Curic <feda.curic@gmail.com>
Date: Mon, 28 Nov 2022 10:38:17 +0100
Subject: [PATCH] Property test parameter-reponse corr function

---
 src/ert/analysis/_es_update.py              |  4 +-
 tests/unit_tests/analysis/test_es_update.py | 45 ++++++++++++++++++++-
 2 files changed, 46 insertions(+), 3 deletions(-)

diff --git a/src/ert/analysis/_es_update.py b/src/ert/analysis/_es_update.py
index 2546991c5d0..9ab3bf83fe2 100644
--- a/src/ert/analysis/_es_update.py
+++ b/src/ert/analysis/_es_update.py
@@ -132,7 +132,7 @@ def _create_temporary_parameter_storage(
     return temporary_storage
 
 
-def correlated_parameter_response_pairs(
+def _correlated_parameter_response_pairs(
     A: npt.NDArray[np.float_], Y: npt.NDArray[np.float_], correlation_threshold: float
 ) -> npt.NDArray[np.int_]:
     N = A.shape[1]
@@ -210,7 +210,7 @@ def analysis_ES(
                 p = A.shape[0]
                 for i in range(A.shape[0]):
                     A_chunk = A[i, :].reshape(-1, N)
-                    corr_idx_Y = correlated_parameter_response_pairs(
+                    corr_idx_Y = _correlated_parameter_response_pairs(
                         A_chunk,
                         Y,
                         module.localization_correlation_threshold(),
diff --git a/tests/unit_tests/analysis/test_es_update.py b/tests/unit_tests/analysis/test_es_update.py
index 0f39363f14f..feb9eb2b626 100644
--- a/tests/unit_tests/analysis/test_es_update.py
+++ b/tests/unit_tests/analysis/test_es_update.py
@@ -5,6 +5,8 @@
 import numpy as np
 
 rng = np.random.default_rng()
+from scipy.linalg import toeplitz
+
 import pandas as pd
 import pytest
 from iterative_ensemble_smoother import IterativeEnsembleSmoother
@@ -13,7 +15,10 @@
 from ert.__main__ import ert_parser
 from ert._c_wrappers.enkf import EnKFMain, EnkfNode, NodeId, ResConfig, RunContext
 from ert.analysis import ErtAnalysisError, ESUpdate
-from ert.analysis._es_update import _create_temporary_parameter_storage
+from ert.analysis._es_update import (
+    _create_temporary_parameter_storage,
+    _correlated_parameter_response_pairs,
+)
 from ert.cli import ENSEMBLE_EXPERIMENT_MODE, ENSEMBLE_SMOOTHER_MODE
 from ert.cli.main import run_cli
 
@@ -492,3 +497,41 @@ def test_update_multiple_param(copy_case):
     # https://en.wikipedia.org/wiki/Variance#For_vector-valued_random_variables
     for prior_name, prior_data in prior.items():
         assert np.trace(np.cov(posterior[prior_name])) < np.trace(np.cov(prior_data))
+
+
+def test_correlated_parameter_response_pair():
+    p = 4
+    m = 2
+    N = 1000
+    rho = 0.9
+    # Correlation matrix or AR(1) model
+    R = toeplitz([rho**i for i in range(p + m)])
+
+    # Get correlated samples from multivariate normal distribution
+    Z = rng.standard_normal(size=(p + m, N))
+    X = np.linalg.cholesky(R) @ Z
+
+    assert (
+        len(
+            _correlated_parameter_response_pairs(
+                X[0, :].reshape(-1, N), X[1, :].reshape(-1, N), rho - 0.05
+            )
+        )
+        == 1
+    )
+    assert (
+        len(
+            _correlated_parameter_response_pairs(
+                X[0, :].reshape(-1, N), X[1:4, :], rho**2 - 0.05
+            )
+        )
+        == 2
+    )
+    assert (
+        len(
+            _correlated_parameter_response_pairs(
+                X[0, :].reshape(-1, N), X[1:6, :], rho**3 - 0.05
+            )
+        )
+        == 3
+    )