Add support for exporting EDITNNC

docs/usage/nnc.rst

@@ -35,4 +35,37 @@ extra columns.
 If you only want vertical connections, add the option ``--pillars`` or ``-vertical``,
 or set ``pillars=True`` if using the Python API (:func:`ecl2df.nnc.df`)
 
+``EDITNNC`` export
+^^^^^^^^^^^^^^^^^^
+
+Data for the ``EDITNNC`` keyword can be dumped, in order to scale the NNC connections
+using Pandas operations. Select the connections you want to scale by slicing
+the nnc dataframe (either from the nnc module, or from the trans module), and fill
+transmissibility multipliers in a new column ``TRANM``, then this can be exported
+to an Eclipse include file:
+
+.. code-block:: python
+
+   from ecl2f import nnc, EclFiles
+
+   eclfiles = EclFiles("MYDATADECK.DATA")
+   nnc_df = nnc.df(eclfiles)
+   nnc_df["TRANM"] = 0.1  # Reduce all NNC transmissibilities
+
+   nnc.df2ecl_editnnc(nnc_df, filename="editnnc.inc")
+
+and the contents of the exported file can be:
+
+..
+   print(nnc.df2ecl_editnnc(nnc.df(eclfiles).head(4).assign(TRANM=0.1)))
+
+.. code-block:: console
+
+    EDITNNC
+    --  I1  J1  K1  I2  J2  K2  TRANM
+        30   4   2  31   4   1    0.1 /
+        30   4   3  31   4   1    0.1 /
+        30   4   3  31   4   2    0.1 /
+        30   4   4  31   4   1    0.1 /
+    / -- 4 nnc connections, avg multiplier 0.1
 

ecl2df/nnc.py

@@ -6,12 +6,13 @@
 from __future__ import absolute_import
 from __future__ import division
 
+import os
 import logging
+import datetime
 import argparse
 import pandas as pd
 
-from .eclfiles import EclFiles
-from .grid import gridgeometry2df
+from ecl2df import common, EclFiles, grid, __version__
 
 logging.basicConfig()
 logger = logging.getLogger(__name__)
@@ -113,7 +114,7 @@ def add_nnc_coords(nncdf, eclfiles):
     Returns:
         DataFrame: Incoming dataframe augmented with the columns X, Y and Z.
     """
-    gridgeometry = gridgeometry2df(eclfiles)
+    gridgeometry = grid.gridgeometry2df(eclfiles)
     gnncdf = pd.merge(
         nncdf,
         gridgeometry,
@@ -195,6 +196,87 @@ def fill_parser(parser):
     return parser
 
 
+def df2ecl_editnnc(nnc_df, filename=None, nocomments=False):
+    """Write an EDITNNC keyword
+
+    This will write
+        EDITNNC
+           IX IY IZ JX JY JZ TRANM /
+           ...
+        /
+
+    and return as string and/or dump to filename.
+
+    The column TRANM must be in the incoming dataframe
+    and it should be the multiplier to be written for
+    each connection.
+
+    If you want to edit only a subset of the non-neighbour
+    connections, filter the dataframe upfront.
+
+    Only rows where the column "DIR" is "NNC" will be considered.
+
+    Args:
+        nnc_df (pd.DataFrame): Dataframe with I1, J1, K1, I2, J2, K2 and a TRANM
+            column, where the multiplier to be written is in TRANM. If the DIR
+            column is present, only the rows with 'NNC' in the DIR column
+            are included.
+        filename (str): Filename to write to
+        nocomments (bool): Set to True if you don't want any comments
+            in the produced string/file
+
+    Returns:
+        string with the EDITNNC keyword.
+    """
+
+    string = ""
+    ecl2df_header = (
+        "Output file printed by ecl2df.nnc"
+        + " "
+        + __version__
+        + "\n"
+        + " at "
+        + str(datetime.datetime.now())
+    )
+    if not nocomments:
+        string += common.comment_formatter(ecl2df_header)
+    string += "\n"
+
+    if "DIR" in nnc_df:
+        nnc_df = nnc_df[nnc_df["DIR"] == "NNC"]
+
+    if "TRANM" not in nnc_df:
+        raise ValueError("TRANM not supplied in nnc_df")
+    string += "EDITNNC" + os.linesep
+    table_str = nnc_df[["I1", "J1", "K1", "I2", "J2", "K2", "TRANM"]].to_string(
+        header=True, index=False
+    )
+    lines = table_str.rstrip().split(os.linesep)
+    indent = "   "
+    string += "-- " + lines[0] + os.linesep
+    string += os.linesep.join([indent + line + " /" for line in lines[1:]])
+    string += os.linesep
+    string += "/"
+    if not nocomments:
+        string += " "
+        string += common.comment_formatter(
+            "  {} nnc connections, avg multiplier {}".format(
+                len(nnc_df), nnc_df["TRANM"].mean()
+            )
+        )
+    string += "\n\n"
+
+    if filename is not None:
+        # Make directory if not present:
+        filenamedir = os.path.dirname(filename)
+        if filenamedir and not os.path.exists(filenamedir):
+            os.makedirs(filenamedir)
+        with open(filename, "w") as file_handle:
+            file_handle.write(string)
+
+    return string
+
+
 def main():
     """Entry-point for module, for command line utility
 

tests/test_nnc.py

@@ -8,8 +8,9 @@
 import sys
 
 import pandas as pd
+import pytest
 
-from ecl2df import nnc, faults, ecl2csv
+from ecl2df import nnc, faults, ecl2csv, trans
 from ecl2df.eclfiles import EclFiles
 
 TESTDIR = os.path.dirname(os.path.abspath(__file__))
@@ -71,6 +72,35 @@ def test_nnc2df_faultnames():
     # Remove I_x, J_x, K_x (and _y) which is not needed
 
 
+def test_df2ecl_editnnc(tmpdir):
+    """Test generation of EDITNNC keyword"""
+    eclfiles = EclFiles(DATAFILE)
+    nncdf = nnc.df(eclfiles)
+    tmpdir.chdir()
+
+    nncdf["TRANM"] = 2
+    editnnc = nnc.df2ecl_editnnc(nncdf, filename="editnnc.inc")
+    editnnc_fromfile = "".join(open("editnnc.inc").readlines())
+    assert editnnc == editnnc_fromfile
+    assert "EDITNNC" in editnnc
+    assert editnnc.count("/") == len(nncdf) + 1
+    assert "avg multiplier" in editnnc
+
+    # Fails when columns are missing
+    with pytest.raises((KeyError, ValueError)):
+        nnc.df2ecl_editnnc(nncdf[["I1", "I2"]])
+
+    editnnc = nnc.df2ecl_editnnc(nncdf, nocomments=True)
+    assert "avg multiplier" not in editnnc
+
+    # Test compatibility with trans module:
+    trans_df = trans.df(eclfiles, addnnc=True)
+    editnnc = nnc.df2ecl_editnnc(trans_df.assign(TRANM=0.3))
+    assert "avg multiplier 0.3" in editnnc or "avg multiplier 0.29999" in editnnc
+
+    print(nnc.df2ecl_editnnc(nnc.df(eclfiles).head(4).assign(TRANM=0.1)))
+
+
 def test_main(tmpdir):
     """Test command line interface"""
     tmpcsvfile = tmpdir.join(".TMP-nnc.csv")