Move from pandas to polars for local_ensemble.py by akustn · Pull Request #13945 · equinor/ert

akustn · 2026-07-03T12:52:49Z

Issue
Resolves #13693

Approach
For moving from pandas to polars in local_ensemble.py there were multiple steps necessary:
First I introduced 2 tests to make sure the test coverage is good enough to move from pandas to polars without losing functionality.
In a next step I removed unused functionality from MeasuredData or functionality that was only used by tests and changed the tests accordingly.
Afterwards the logic from MeasuredData was moved to local_ensemble.py as MeasuredData was only called from local_ensemble.py and made calls back to local_ensemble.py.
In a final step all usage of pandas in local_ensemble.py was replaced by polars to improve consistency.

Remaining work to be found in #13954

(Screenshot of new behavior in GUI if applicable)

PR title captures the intent of the changes, and is fitting for release notes.
Added appropriate release note label
Commit history is consistent and clean, in line with the contribution guidelines.
Make sure unit tests pass locally after every commit (git rebase -i main --exec 'just rapid-tests')

When applicable

When there are user facing changes: Updated documentation
New behavior or changes to existing untested code: Ensured that unit tests are added (See Ground Rules).
Large PR: Prepare changes in small commits for more convenient review
Bug fix: Add regression test for the bug
Bug fix: Add backport label to latest release (format: 'backport release-branch-name')

codecov-commenter · 2026-07-06T11:14:53Z

Codecov Report

❌ Patch coverage is 97.67442% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 91.89%. Comparing base (31248f9) to head (5ef9532).
✅ All tests successful. No failed tests found.

Files with missing lines	Patch %	Lines
src/ert/storage/local_ensemble.py	97.29%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main   #13945      +/-   ##
==========================================
+ Coverage   91.85%   91.89%   +0.03%     
==========================================
  Files         484      482       -2     
  Lines       33581    33544      -37     
==========================================
- Hits        30847    30825      -22     
+ Misses       2734     2719      -15

Flag	Coverage Δ
cli-tests	`36.16% <18.60%> (-0.02%)`	⬇️
fuzz	`44.08% <18.60%> (-0.01%)`	⬇️
gui-tests	`58.39% <48.83%> (+<0.01%)`	⬆️
performance-and-unit-tests	`80.55% <97.67%> (+0.09%)`	⬆️
test	`46.32% <18.60%> (-0.01%)`	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

Files with missing lines	Coverage Δ
src/ert/__init__.py	`100.00% <ø> (ø)`
...rt/gui/tools/manage_experiments/ensemble_widget.py	`97.17% <100.00%> (+2.39%)`	⬆️
...ugins/hook_implementations/workflows/csv_export.py	`74.64% <100.00%> (+0.36%)`	⬆️
...ok_implementations/workflows/export_misfit_data.py	`92.85% <100.00%> (-0.90%)`	⬇️
src/ert/storage/local_ensemble.py	`96.79% <97.29%> (-0.04%)`	⬇️

codspeed-hq · 2026-07-06T11:55:41Z

Merging this PR will not alter performance

✅ 36 untouched benchmarks

_{Comparing change-pandas-to-polars (5ef9532) with main (31248f9)}

SAKavli

I have not finished reviewing, but will post the comments I have so far as I am joining a meeting now.
I can't put my finger on it yet, but I feel like there might be something to be done with the misfit changes, as we went from 12 lines to a lot more. We can look at this together later today!

SAKavli · 2026-07-07T12:01:56Z

Aha, when I come to the second commit, it makes sense why all the code in the first was added!
Perhaps we could review the rest together

SAKavli

We decided to redo the commit order and try to have a more clear separation of the removal of MeasuredData and the conversion from pandas to polars dataframes

Copilot

Pull request overview

This PR completes the migration of misfit collection/measured-data plumbing around LocalEnsemble from pandas to polars, and updates downstream GUI/workflows/tests accordingly, including removing the previously exposed MeasuredData API.

Changes:

Refactors LocalEnsemble.load_all_misfit_data() to return a polars.DataFrame and inlines measured-data loading/validation into local_ensemble.py.
Updates GUI/workflow export paths to work with polars misfit tables (converting to pandas only where required for HDF output / existing pandas joins).
Adjusts unit/UI tests and snapshots to reflect the new polars-based data structures and formatting.

Reviewed changes

Copilot reviewed 14 out of 14 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
tests/ert/unit_tests/test_misfit_collector.py	Updates snapshot generation and row-access expectations for polars misfit output.
tests/ert/unit_tests/snapshots/test_misfit_collector/test_misfit_collector/0/misfit_collector.csv	Refreshes snapshot values/formatting emitted by polars CSV writer.
tests/ert/unit_tests/scenarios/test_summary_response.py	Switches measured-data assertions from pandas/MeasuredData to polars `_load_measured_data()`.
tests/ert/unit_tests/gui/ertwidgets/test_ensemble_widget.py	Adds test ensuring GUI misfit retrieval returns a polars DataFrame.
tests/ert/unit_tests/data/test_integration_data.py	Reworks measured-data integration tests to use `_load_measured_data()` and polars columns.
tests/ert/unit_tests/data/snapshots/test_integration_data/test_all_measured_snapshot/0/snake_oil_measured_output.csv	Removes the old MeasuredData snapshot artifact.
tests/ert/ui_tests/gui/test_full_manual_update_workflow.py	Updates UI workflow test to use `_load_measured_data()` instead of MeasuredData.
src/ert/storage/local_ensemble.py	Implements polars misfit computation and moves measured-data loading/validation into `LocalEnsemble`.
src/ert/plugins/hook_implementations/workflows/export_misfit_data.py	Adapts misfit export workflow to polars misfit output, converting to pandas for HDF writing.
src/ert/plugins/hook_implementations/workflows/csv_export.py	Adapts CSV export workflow to polars misfit output, converting to pandas for joins.
src/ert/gui/tools/manage_experiments/ensemble_widget.py	Simplifies GUI misfit retrieval now that `load_all_misfit_data()` returns polars directly.
src/ert/data/_measured_data.py	Removes the pandas-based `MeasuredData` implementation.
src/ert/data/init.py	Removes `ert.data` export surface for `MeasuredData`.
src/ert/init.py	Removes top-level `MeasuredData` export from the public `ert` API.

SAKavli · 2026-07-10T12:51:52Z

 import pytest


 def test_misfit_collector(snake_oil_case_storage, snake_oil_default_storage, snapshot):


Out of scope of this PR

SAKavli · 2026-07-10T12:53:31Z

+        Returns:
+            DataFrame: A DataFrame containing the measured data for all
+                realizations in the ensemble. Each column corresponds to a key
+                in the measured data, and each row corresponds to a realization.
+        """


Not changed in this PR, just moved

SAKavli · 2026-07-10T12:54:59Z

+    def get_simulated_data(self, data: pd.DataFrame) -> pd.DataFrame:
+        """Dimension of data is (number of responses x number of realizations)."""
+        return data[~data.index.isin(["OBS", "STD"])]


Perhaps we can delete this, @akustn ?

SAKavli · 2026-07-10T13:02:56Z

 from .config import (
    ForwardModelStepJSON,
    ForwardModelStepPlugin,
    ForwardModelStepValidationError,
    ForwardModelStepWarning,
    ForwardModelStepDocumentation,
    ErtScript,
    ErtScriptWorkflow,
    WorkflowConfigs,
 )
-from .data import MeasuredData
 from .workflow_runner import WorkflowRunner


We don't believe this is used elsewhere

SAKavli

Seems like copilot made one fair point, we should consider removing that.
Apart from that, I think this can be merged!
Good work @akustn 🎉

MeasuredData is only called from local_ensemble.py. In get_data we have calls back to local_ensemble.py. Moving the logic to local_ensemble.py prevents circular calls.

akustn self-assigned this Jul 3, 2026

akustn added the fin_ny_byrjing label Jul 3, 2026

akustn added this to SCOUT Jul 3, 2026

akustn force-pushed the change-pandas-to-polars branch 3 times, most recently from 6a36ef0 to c2bad13 Compare July 6, 2026 08:57

akustn moved this to In Progress in SCOUT Jul 6, 2026

akustn force-pushed the change-pandas-to-polars branch 3 times, most recently from 9743fd0 to 2091eae Compare July 6, 2026 11:02

akustn force-pushed the change-pandas-to-polars branch 2 times, most recently from 340b2f0 to 4cf340e Compare July 7, 2026 07:00

akustn mentioned this pull request Jul 7, 2026

Make csv_export.py only use polars Dataframe #13954

Open

akustn changed the title ~~Make load_all_misfit_data use polars~~ Move from pandas to polars for local_ensemble.py Jul 7, 2026

akustn marked this pull request as ready for review July 7, 2026 07:29

akustn moved this from In Progress to Ready for Review in SCOUT Jul 7, 2026

SAKavli reviewed Jul 7, 2026

View reviewed changes

Comment thread src/ert/plugins/hook_implementations/workflows/export_misfit_data.py Outdated

Comment thread src/ert/plugins/hook_implementations/workflows/export_misfit_data.py

Comment thread src/ert/storage/local_ensemble.py

SAKavli reviewed Jul 7, 2026

View reviewed changes

Comment thread src/ert/storage/local_ensemble.py

Comment thread src/ert/storage/local_ensemble.py

SAKavli moved this from Ready for Review to In Progress in SCOUT Jul 8, 2026

SAKavli removed this from SCOUT Jul 8, 2026

akustn force-pushed the change-pandas-to-polars branch 2 times, most recently from 213f9d6 to fc47809 Compare July 9, 2026 08:43

SAKavli requested changes Jul 9, 2026

View reviewed changes

akustn force-pushed the change-pandas-to-polars branch 2 times, most recently from 8ae1522 to 1151623 Compare July 9, 2026 11:15

akustn added improvement Something nice to have, that will make life easier for developers or users or both. release-notes:improvement Automatically categorise as improvement in release notes labels Jul 9, 2026

akustn requested a review from SAKavli July 9, 2026 11:55

akustn mentioned this pull request Jul 10, 2026

Cast dataframe to pandas equinor/semeio#923

Merged

SAKavli requested a review from Copilot July 10, 2026 12:45

Copilot started reviewing on behalf of SAKavli July 10, 2026 12:45 View session

Copilot AI reviewed Jul 10, 2026

View reviewed changes

SAKavli approved these changes Jul 10, 2026

View reviewed changes

akustn force-pushed the change-pandas-to-polars branch from 4106e50 to fcc166d Compare August 5, 2026 06:51

Add test for missing keys in MeasuredData

2bc5acd

akustn force-pushed the change-pandas-to-polars branch 2 times, most recently from a81fd0c to dbca7ca Compare August 5, 2026 08:33

akustn added 4 commits August 5, 2026 10:42

Add test for get_misfit_df

6293a0e

Remove unused functions from MeasuredData

a614fc6

Move MeasuredData logic to local_ensemble.py

ed408d5

MeasuredData is only called from local_ensemble.py. In get_data we have calls back to local_ensemble.py. Moving the logic to local_ensemble.py prevents circular calls.

Make load_all_misfit_data use polars

5ef9532

akustn force-pushed the change-pandas-to-polars branch from dbca7ca to 5ef9532 Compare August 5, 2026 08:42

akustn merged commit 62af0fd into main Aug 5, 2026
38 checks passed

akustn deleted the change-pandas-to-polars branch August 5, 2026 09:48

		import pytest


		def test_misfit_collector(snake_oil_case_storage, snake_oil_default_storage, snapshot):

Uh oh!

Conversation

akustn commented Jul 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

When applicable

Uh oh!

codecov-commenter commented Jul 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

codspeed-hq Bot commented Jul 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Merging this PR will not alter performance

Uh oh!

SAKavli left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

SAKavli commented Jul 7, 2026

Uh oh!

SAKavli left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

SAKavli Jul 10, 2026

Choose a reason for hiding this comment

Uh oh!

SAKavli Jul 10, 2026

Choose a reason for hiding this comment

Uh oh!

SAKavli Jul 10, 2026

Choose a reason for hiding this comment

Uh oh!

SAKavli Jul 10, 2026

Choose a reason for hiding this comment

Uh oh!

SAKavli left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

akustn commented Jul 3, 2026 •

edited

Loading

codecov-commenter commented Jul 6, 2026 •

edited

Loading

codspeed-hq Bot commented Jul 6, 2026 •

edited

Loading