From e5ff16703b04a0ca386c67ef95669f64397742e0 Mon Sep 17 00:00:00 2001 From: xjules Date: Tue, 10 Dec 2024 15:02:05 +0100 Subject: [PATCH] WIP: writting tests for categorical data in design sheet --- tests/ert/ui_tests/cli/analysis/test_design_matrix.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/tests/ert/ui_tests/cli/analysis/test_design_matrix.py b/tests/ert/ui_tests/cli/analysis/test_design_matrix.py index 738c92a8963..d93e376731c 100644 --- a/tests/ert/ui_tests/cli/analysis/test_design_matrix.py +++ b/tests/ert/ui_tests/cli/analysis/test_design_matrix.py @@ -22,6 +22,7 @@ def test_run_poly_example_with_design_matrix(): { "REAL": list(range(num_realizations)), "a": a_values, + "category": 5 * ["cat1"] + 5 * ["cat2"], } ) default_sheet_df = pd.DataFrame([["b", 1], ["c", 2]]) @@ -59,6 +60,7 @@ def _load_coeffs(filename): return json.load(f)["DESIGN_MATRIX"] def _evaluate(coeffs, x): + assert coeffs["category"] in ["cat1", "cat2"] return coeffs["a"] * x**2 + coeffs["b"] * x + coeffs["c"] if __name__ == "__main__": @@ -88,8 +90,9 @@ def _evaluate(coeffs, x): "DESIGN_MATRIX" )["values"] np.testing.assert_array_equal(params[:, 0], a_values) - np.testing.assert_array_equal(params[:, 1], 10 * [1]) - np.testing.assert_array_equal(params[:, 2], 10 * [2]) + np.testing.assert_array_equal(params[:, 0], 5 * ["cat1"] + 5 * ["cat2"]) + np.testing.assert_array_equal(params[:, 2], 10 * [1]) + np.testing.assert_array_equal(params[:, 3], 10 * [2]) @pytest.mark.usefixtures("copy_poly_case")