Add script for merging UNRST files

docs/scripts/merge_unrst_files.rst

@@ -0,0 +1,8 @@
+
+MERGE_UNRST_FILES
+=================
+
+.. argparse::
+   :module: subscript.merge_unrst_files.merge_unrst_files
+   :func: get_parser
+   :prog: merge_unrst_files

pyproject.toml

@@ -101,6 +101,7 @@ fmu_copy_revision = "subscript.fmu_copy_revision.fmu_copy_revision:main"
 fmuobs = "subscript.fmuobs.fmuobs:main"
 interp_relperm = "subscript.interp_relperm.interp_relperm:main"
 merge_rft_ertobs = "subscript.merge_rft_ertobs.merge_rft_ertobs:main"
+merge_unrst_files = "subscript.merge_unrst_files.merge_unrst_files:main"
 ofmvol2csv = "subscript.ofmvol2csv.ofmvol2csv:main"
 pack_sim = "subscript.pack_sim.pack_sim:main"
 params2csv = "subscript.params2csv.params2csv:main"

src/subscript/merge_unrst_files/merge_unrst_files.py

@@ -0,0 +1,115 @@
+import argparse
+import logging
+
+import resfo
+
+from subscript import __version__, getLogger
+
+DESCRIPTION = """Read two ``UNRST`` files and export a merged version. This is useful in
+cases where history and prediction are run separately and one wants to calculate
+differences across dates in the the two files.
+"""
+
+CATEGORY = "utility.eclipse"
+
+EXAMPLES = """
+.. code-block:: console
+
+  DEFINE <RESTART_DIR>      iter-3
+  FORWARD_MODEL MERGE_UNRST_FILES(<UNRST1>=../<RESTART_DIR>/<ECLBASE>.UNRST, <UNRST2>=<ECLBASE>.UNRST, <MERGED_FILE>=eclipse/model/ECLIPSE_MERGED.UNRST)
+
+"""  # noqa
+
+logger = getLogger(__name__)
+logger.setLevel(logging.INFO)
+
+
+def get_parser() -> argparse.ArgumentParser:
+    """Function to create the argument parser that is going to be served to the user.
+
+    Returns:
+        argparse.ArgumentParser: The argument parser to be served
+
+    """
+    parser = argparse.ArgumentParser(
+        prog="merge_unrst_files.py", description=DESCRIPTION
+    )
+    parser.add_argument("UNRST_HIST_FILE", type=str, help="UNRST file 1, history part")
+    parser.add_argument(
+        "UNRST_PRED_FILE", type=str, help="UNRST file 2, prediction part"
+    )
+    parser.add_argument(
+        "-o",
+        "--output",
+        type=str,
+        help="Name of the merged UNRST file",
+        default="MERGED.UNRST",
+    )
+    parser.add_argument(
+        "-v",
+        "--version",
+        action="version",
+        version=f"%(prog)s (subscript version {__version__ })",
+    )
+    return parser
+
+
+def _check_report_number(
+    args: argparse.Namespace,
+    max_report_number_hist: int,
+    current_report_number: int,
+) -> None:
+    """Check that pred file report numbers are larger than in hist file.
+
+    Args:
+        args (argparse.Namespace): The Namespace object with the argument list.
+        max_report_number_hist (int): The largest report number in hist file.
+        current_report_number (int): The current restart report number in pred file.
+    """
+
+    if current_report_number <= max_report_number_hist:
+        logger.warning(
+            f"{args.UNRST_PRED_FILE} file has a restart report number"
+            + f" ({current_report_number}) which is smaller than largest report number"
+            + f" in {args.UNRST_HIST_FILE} ({max_report_number_hist})"
+        )
+        logger.warning(
+            "Check that you have entered arguments in correct order and/or"
+            + " that the unrst files are compatible."
+        )
+
+
+def main() -> None:
+    """Parse command line arguments and run"""
+
+    args: argparse.Namespace = get_parser().parse_args()
+
+    logger.info(f"Merge unrst files {args.UNRST_HIST_FILE} and {args.UNRST_PRED_FILE}.")
+    unrst_hist = resfo.read(args.UNRST_HIST_FILE)
+    unrst_pred = resfo.read(args.UNRST_PRED_FILE)
+
+    max_first_seqnum: int = 1
+    max_first_solver_step: int = 1
+    max_first_report_step: int = 1
+
+    for kw, val in unrst_hist:
+        if kw == "SEQNUM  ":  # restart report number
+            max_first_seqnum = max(max_first_seqnum, val[0])
+        if kw == "INTEHEAD":
+            max_first_solver_step = max(max_first_solver_step, val[67])
+            max_first_report_step = max(max_first_report_step, val[68])
+
+    for kw, val in unrst_pred:
+        if kw == "SEQNUM  ":
+            _check_report_number(args, max_first_seqnum, val[0])
+            val[0] += max_first_seqnum
+        if kw == "INTEHEAD":
+            val[67] += max_first_solver_step
+            val[68] += max_first_report_step
+
+    resfo.write(args.output, unrst_hist + unrst_pred)
+    logger.info(f"Done. Merged file is written to {args.output}")
+
+
+if __name__ == "__main__":
+    main()

tests/test_merge_unrst_files.py

@@ -0,0 +1,92 @@
+import filecmp
+import os
+import shutil
+import subprocess
+from pathlib import Path
+
+import pytest
+import resfo
+
+from subscript.merge_unrst_files import merge_unrst_files
+
+UNRST_HIST = (
+    Path(__file__).absolute().parent / "testdata_merge_unrst_files" / "HIST.UNRST"
+)
+UNRST_PRED = (
+    Path(__file__).absolute().parent / "testdata_merge_unrst_files" / "PRED.UNRST"
+)
+
+# pylint: disable=protected-access
+
+
+@pytest.mark.integration
+def test_integration():
+    """Test that endpoint is installed"""
+    assert subprocess.check_output(["merge_unrst_files", "-h"])
+
+
+@pytest.mark.integration
+def test_main_default_output(tmp_path, mocker):
+    """Test invocation from command line"""
+    os.chdir(tmp_path)
+
+    mocker.patch("sys.argv", ["merge_unrst_files", str(UNRST_HIST), str(UNRST_PRED)])
+    merge_unrst_files.main()
+
+    assert Path("MERGED.UNRST").exists()
+
+
+@pytest.mark.integration
+def test_main_with_output(tmp_path, mocker):
+    """Test invocation from command line"""
+    os.chdir(tmp_path)
+
+    mocker.patch(
+        "sys.argv",
+        [
+            "merge_unrst_files",
+            str(UNRST_HIST),
+            str(UNRST_PRED),
+            "-o",
+            "MY_MERGED.UNRST",
+        ],
+    )
+    merge_unrst_files.main()
+
+    assert Path("MY_MERGED.UNRST").exists()
+
+
+def get_restart_report_numbers(unrst_merged):
+    """Get restart report numbers from merged unrst file. It should be [0, 82, 206]"""
+    report_numbers = []
+    for kw, val in unrst_merged:
+        if kw == "SEQNUM  ":
+            report_numbers.append(val[0])
+    return report_numbers
+
+
+@pytest.mark.integration
+def test_check_report_numbers(tmp_path, mocker):
+    """Verify that merged restart has the expected restart report numbers."""
+    os.chdir(tmp_path)
+
+    mocker.patch(
+        "sys.argv",
+        [
+            "merge_unrst_files",
+            str(UNRST_HIST),
+            str(UNRST_PRED),
+            "-o",
+            "MY_MERGED.UNRST",
+        ],
+    )
+    merge_unrst_files.main()
+
+    expected_report_numbers = [0, 82, 206]
+    report_numbers = get_restart_report_numbers(resfo.read("MY_MERGED.UNRST"))
+
+    print(
+        f"expected restart report numbers: {expected_report_numbers}, "
+        + f"actual restart report_numbers: {report_numbers}"
+    )
+    assert report_numbers == expected_report_numbers