Include localization columns when joining obs dataframes

Author: SAKavli

src/ert/storage/local_ensemble.py

@@ -1152,6 +1152,9 @@ def get_observations_and_responses(
                                 "observation_key",
                                 "observations",
                                 "std",
+                                "east",
+                                "north",
+                                "radius",
                             ]
                         )
 

tests/ert/performance_tests/test_analysis.py

@@ -83,6 +83,9 @@ def g(X):
             "index": np.arange(len(observations)),
             "observations": observations,
             "std": observation_noise,
+            "east": pl.Series([None] * len(observations), dtype=pl.Float32),
+            "north": pl.Series([None] * len(observations), dtype=pl.Float32),
+            "radius": pl.Series([None] * len(observations), dtype=pl.Float32),
         }
     )
 

tests/ert/performance_tests/test_obs_and_responses_performance.py

@@ -142,6 +142,9 @@ def create_experiment_args(
                 rng.normal(loc=0.2, scale=0.1, size=num_gen_data_obs),
                 dtype=pl.Float32,
             ),
+            "east": pl.Series([None] * num_gen_data_obs, dtype=pl.Float32),
+            "north": pl.Series([None] * num_gen_data_obs, dtype=pl.Float32),
+            "radius": pl.Series([None] * num_gen_data_obs, dtype=pl.Float32),
         }
     )
 
@@ -195,6 +198,9 @@ def create_experiment_args(
                 rng.normal(loc=0.2, scale=0.1, size=num_summary_obs),
                 dtype=pl.Float32,
             ),
+            "east": pl.Series([None] * num_summary_obs, dtype=pl.Float32),
+            "north": pl.Series([None] * num_summary_obs, dtype=pl.Float32),
+            "radius": pl.Series([None] * num_summary_obs, dtype=pl.Float32),
         }
     )
 

tests/ert/unit_tests/analysis/test_es_update.py

@@ -42,6 +42,9 @@ def obs() -> pl.DataFrame:
             "index": pl.Series([0, 1, 2], dtype=pl.UInt16),
             "observations": pl.Series([1.0, 1.0, 1.0], dtype=pl.Float32),
             "std": pl.Series([0.1, 1.0, 10.0], dtype=pl.Float32),
+            "east": pl.Series([None, None, None], dtype=pl.Float32),
+            "north": pl.Series([None, None, None], dtype=pl.Float32),
+            "radius": pl.Series([None, None, None], dtype=pl.Float32),
         }
     )
 
@@ -250,6 +253,9 @@ def test_update_handles_precision_loss_in_std_dev(tmp_path):
                             [559437122.6211826, 999999999.9999999, 1.9],
                             dtype=pl.Float32,
                         ),
+                        "east": pl.Series([None, None, None], dtype=pl.Float32),
+                        "north": pl.Series([None, None, None], dtype=pl.Float32),
+                        "radius": pl.Series([None, None, None], dtype=pl.Float32),
                     }
                 )
             },
@@ -364,6 +370,9 @@ def test_update_raises_on_singular_matrix(tmp_path):
                             [0.33333334, 0.14142136, 0.0],
                             dtype=pl.Float32,
                         ),
+                        "east": pl.Series([None, None, None], dtype=pl.Float32),
+                        "north": pl.Series([None, None, None], dtype=pl.Float32),
+                        "radius": pl.Series([None, None, None], dtype=pl.Float32),
                     }
                 )
             },
@@ -940,6 +949,9 @@ def test_gen_data_obs_data_mismatch(storage, uniform_parameter):
             "index": pl.Series([1000], dtype=pl.UInt16),
             "observations": pl.Series([1.0], dtype=pl.Float32),
             "std": pl.Series([0.1], dtype=pl.Float32),
+            "east": pl.Series([None], dtype=pl.Float32),
+            "north": pl.Series([None], dtype=pl.Float32),
+            "radius": pl.Series([None], dtype=pl.Float32),
         }
     )
 

tests/ert/unit_tests/storage/test_local_storage.py

@@ -787,6 +787,9 @@ def test_asof_joining_summary(tmp_path, perturb_observations, perturb_responses)
                     [0.1] * len(response_keys),
                     dtype=pl.Float32,
                 ),
+                "north": pl.Series([None] * len(response_keys), dtype=pl.Float32),
+                "east": pl.Series([None] * len(response_keys), dtype=pl.Float32),
+                "radius": pl.Series([None] * len(response_keys), dtype=pl.Float32),
             }
         )