Incorporate other simulators

docs/index.rst

@@ -6,7 +6,7 @@
 Welcome to fmu-sumo-sim2sumo's documentation!
 =============================================
 ``fmu.sumo.sim2sumo`` is a python package for uploading results to Sumo.
-The package facilitates upload of results from reservoir simulators such as **eclipse**, **IX**, and **OPM flow** as arrow files.
+The package facilitates upload of results from reservoir simulators such as **eclipse**, **IX**, **pflotran**, and **OPM flow** as arrow files.
 
 .. toctree::
    :maxdepth: 2

src/fmu/sumo/sim2sumo/_special_treatments.py

@@ -1,15 +1,18 @@
 """Special treatment of some options used in res2df"""
-from inspect import signature
+
 import importlib
 import logging
+from inspect import signature
 from pathlib import Path
-import res2df
-from res2df.common import convert_lyrlist_to_zonemap, parse_lyrfile
+
 import pandas as pd
 import pyarrow as pa
+import res2df
+from res2df.common import convert_lyrlist_to_zonemap, parse_lyrfile
 
 logging.getLogger(__name__).setLevel(logging.DEBUG)
 
+
 def convert_to_arrow(frame):
     """Convert pd.DataFrame to arrow
 
@@ -23,13 +26,17 @@ def convert_to_arrow(frame):
     logger.debug("!!!!Using convert to arrow!!!")
     standard = {"DATE": pa.timestamp("ms")}
     if "DATE" in frame.columns:
-        frame["DATE"] = pd.to_datetime(frame["DATE"], infer_datetime_format=True)
+        frame["DATE"] = pd.to_datetime(
+            frame["DATE"], infer_datetime_format=True
+        )
     scheme = []
     for column_name in frame.columns:
         if pd.api.types.is_string_dtype(frame[column_name]):
             scheme.append((column_name, pa.string()))
         else:
-            scheme.append((column_name, standard.get(column_name, pa.float32())))
+            scheme.append(
+                (column_name, standard.get(column_name, pa.float32()))
+            )
     logger.debug(scheme)
     table = pa.Table.from_pandas(frame, schema=pa.schema(scheme))
     return table
@@ -112,8 +119,12 @@ def _define_submodules():
         except AttributeError:
             logger.debug("No df function in %s", submod_path)
 
-    logger.debug("Returning the submodule names as a list: %s ", submodules.keys())
-    logger.debug("Returning the submodules extra args as a dictionary: %s ", submodules)
+    logger.debug(
+        "Returning the submodule names as a list: %s ", submodules.keys()
+    )
+    logger.debug(
+        "Returning the submodules extra args as a dictionary: %s ", submodules
+    )
 
     return tuple(submodules.keys()), submodules
 

src/fmu/sumo/sim2sumo/hook_implementations/jobs.py

@@ -1,13 +1,12 @@
 """Install in ert, including documentation to ert-docs"""
-import importlib
+
 import os
 import sys
 from pathlib import Path
 
 from ert.shared.plugins.plugin_manager import hook_implementation
 from ert.shared.plugins.plugin_response import plugin_response
 
-
 DESCRIPTION = """
 Makes result simulator (Eclipse, OPM, IX) available in sumo.
 This is done by SIM2SUMO in a three step process:

src/fmu/sumo/sim2sumo/sim2sumo.py

@@ -5,24 +5,26 @@
    3. Uploads to Sumo
 """
 
-import sys
+import argparse
+import logging
 import re
-from typing import Union
+import sys
 from pathlib import Path, PosixPath
-import logging
-import argparse
+from typing import Union
+
 import pandas as pd
-import res2df
 import pyarrow as pa
+import res2df
 import yaml
 from fmu.dataio import ExportData
 from fmu.sumo.uploader.scripts.sumo_upload import sumo_upload_main
-from ._special_treatments import SUBMODULES, SUBMOD_DICT, convert_options, tidy
+
+from ._special_treatments import SUBMOD_DICT, SUBMODULES, convert_options, tidy
 
 logging.getLogger(__name__).setLevel(logging.DEBUG)
 
 
-def yaml_load(file_name: str) -> dict:
+def yaml_load(file_name):
     """Load yaml config file into dict
 
     Args:
@@ -59,6 +61,24 @@ def give_name(datafile_path: str) -> str:
     return base_name
 
 
+def fix_suffix(datafile_path: str):
+    """Check if suffix is .DATA, if not change to
+
+    Args:
+        datafile_path (PosixPath): path to check
+
+    Returns:
+        str: the corrected path
+    """
+    logger = logging.getLogger(__file__ + ".fix_suffix")
+    string_datafile_path = str(datafile_path)
+    if not string_datafile_path.endswith(".DATA"):
+        corrected_path = re.sub(r"\..*", ".DATA", string_datafile_path)
+        logger.debug("Changing %s to %s", string_datafile_path, corrected_path)
+        datafile_path = corrected_path
+    return datafile_path
+
+
 def get_results(
     datafile_path: str, submod: str, print_help=False, **kwargs
 ) -> Union[pa.Table, pd.DataFrame]:
@@ -75,6 +95,7 @@ def get_results(
     logger = logging.getLogger(__file__ + ".get_dataframe")
     extract_df = SUBMOD_DICT[submod]["extract"]
     arrow = kwargs.get("arrow", True)
+    datafile_path = fix_suffix(datafile_path)
     output = None
     trace = None
     if print_help:
@@ -88,6 +109,11 @@ def get_results(
         }
         logger.debug("Exporting with arguments %s", right_kwargs)
         try:
+            logger.info(
+                "Extracting data from %s with %s",
+                datafile_path,
+                extract_df.__name__,
+            )
             output = extract_df(
                 res2df.ResdataFiles(datafile_path),
                 **convert_options(right_kwargs),
@@ -161,71 +187,29 @@ def export_results(
     return exp_path
 
 
-def read_config(config, datafile=None, datatype=None) -> tuple:
+def read_config(config, datafile=None, datatype=None):
     """Read config settings
 
     Args:
         config (dict): the settings for export of simulator results
         kwargs (dict): overiding settings
 
     Returns:
-        tuple: datafiles as list, submodules to use as list, and options as dict
+        tuple: datafiles as list, submodules to use as list, and options as kwargs
     """
     # datafile can be read as list, or string which can be either folder or filepath
     logger = logging.getLogger(__file__ + ".read_config")
     logger.debug("Input config keys are %s", config.keys())
 
-    defaults = {
-        "datafile": "eclipse/model/",
-        "datatypes": ["summary", "rft", "satfunc"],
-        "options": {"arrow": True},
-    }
-    try:
-        simconfig = config["sim2sumo"]
-    except KeyError:
-        logger.warning(
-            "No specification in config, will use defaults %s", defaults
-        )
-        simconfig = defaults
+    simconfig = config.get("sim2sumo", {})
     if isinstance(simconfig, bool):
-        simconfig = defaults
-    if datafile is None:
-        datafile = simconfig.get("datafile", "eclipse/model/")
+        simconfig = {}
+    datafiles = find_datafiles(datafile, simconfig)
 
-    if isinstance(datafile, str):
-        logger.debug("Using this string %s to find datafile(s)", datafile)
-        datafile_posix = Path(datafile)
-        if datafile_posix.is_dir():
-            logger.debug("Directory, globbing for datafiles")
-            datafiles = list(datafile_posix.glob("*.DATA"))
+    submods = find_datatypes(datatype, simconfig)
 
-        else:
-            logger.debug("File path, will just use this one")
-            datafiles = [datafile]
-    else:
-        logger.debug("String is list")
-        datafiles = datafile
-    logger.debug("Datafile(s) to use %s", datafiles)
-    if datatype is None:
-        try:
-            submods = simconfig["datatypes"]
-            if isinstance(submods, str):
-                if submods == "all":
-                    submods = SUBMODULES
-                else:
-                    submods = [submods]
-        except KeyError:
-            submods = defaults["datatypes"]
-    else:
-        submods = [datatype]
+    options = simconfig.get("options", {"arrow": True})
 
-    try:
-        options = simconfig["options"]
-        logger.info("Will use these options %s", options)
-    except KeyError:
-        logger.info("No special options selected")
-        options = {}
-    options["arrow"] = options.get("arrow", True)
     logger.info(
         "Running with: datafile(s): \n%s \n Types: \n %s \noptions:\n %s",
         datafiles,
@@ -235,7 +219,79 @@ def read_config(config, datafile=None, datatype=None) -> tuple:
     return datafiles, submods, options
 
 
-def export_with_config(config_path, datafile=None, datatype=None) -> tuple:
+def find_datatypes(datatype, simconfig):
+    """Find datatypes to extract
+
+    Args:
+        datatype (str or None): datatype to extract
+        simconfig (dict): the config file settings
+
+    Returns:
+        list or dict: data types to extract
+    """
+
+    if datatype is None:
+        submods = simconfig.get("datatypes", ["summary", "rft", "satfunc"])
+
+        if submods == "all":
+            submods = SUBMODULES
+    else:
+        submods = [datatype]
+    return submods
+
+
+def find_datafiles(seedpoint, simconfig):
+    """Find all relevant paths that can be datafiles
+
+    Args:
+        seedpoint (str, list): path of datafile, or list of folders where one can find one
+        simconfig (dict): the sim2sumo config settings
+
+    Returns:
+        list: list of datafiles to interrogate
+    """
+
+    logger = logging.getLogger(__file__ + ".find_datafiles")
+    datafiles = []
+
+    if seedpoint is None:
+        seedpoint = simconfig.get(
+            "datafile",
+            [
+                "eclipse/model/",
+                "ix/model/",
+                "opm/model/",
+                "pflotran/model",
+            ],
+        )
+
+    if isinstance(seedpoint, (str, PosixPath)):
+        logger.debug("Using this string %s to find datafile(s)", seedpoint)
+        seedpoint_posix = Path(seedpoint)
+
+        if seedpoint_posix.is_dir():
+            logger.debug("%s is directory, globbing for datafiles", seedpoint)
+            glob_list = (
+                list(seedpoint_posix.glob("*.DATA"))
+                + list(seedpoint_posix.glob("*.afi"))
+                + list(seedpoint_posix.glob("*.in"))
+            )
+
+            logger.debug("Results are %s", glob_list)
+            datafiles.extend(find_datafiles(glob_list, simconfig))
+
+        else:
+            logger.debug("%s is file path, will just use this one", seedpoint)
+            datafiles.append(seedpoint)
+    else:
+        logger.debug("%s is list", seedpoint)
+        for item in seedpoint:
+            datafiles.extend(find_datafiles(item, simconfig))
+    logger.debug("Datafile(s) to use %s", datafiles)
+    return datafiles
+
+
+def export_with_config(config_path, datafile=None, datatype=None):
     """Export several datatypes with yaml config file
 
     Args:
@@ -277,10 +333,10 @@ def export_with_config(config_path, datafile=None, datatype=None) -> tuple:
 
 
 def upload(
-    upload_folder: str,
-    suffixes: list,
-    env: str = "prod",
-    threads: int = 5,
+    upload_folder,
+    suffixes,
+    env="prod",
+    threads=5,
     start_del="real",
     config_path="fmuconfig/output/global_variables.yml",
 ):
@@ -315,7 +371,7 @@ def upload(
         logger.warning("Nothing to export..")
 
 
-def parse_args() -> argparse.Namespace:
+def parse_args():
     """Parse arguments for command line tool
 
     Returns:
@@ -373,7 +429,7 @@ def parse_args() -> argparse.Namespace:
     return args
 
 
-def give_help(submod: Union[None, str], only_general: bool = False) -> str:
+def give_help(submod, only_general=False):
     """Give descriptions of variables available for submodule
 
     Args:
@@ -408,9 +464,7 @@ def give_help(submod: Union[None, str], only_general: bool = False) -> str:
     return text_to_return
 
 
-def upload_with_config(
-    config_path: str, datafile: Union[str, PosixPath], datatype: str, env: str
-):
+def upload_with_config(config_path, datafile, datatype, env):
     """Upload simulator results to sumo
 
     Args:
@@ -433,9 +487,11 @@ def main():
     """Main function to be called"""
     logger = logging.getLogger(__file__ + ".main")
     args = parse_args()
+    logger.debug("Running with arguments %s", args)
     try:
         print(give_help(args.help_on))
     except AttributeError:
+        logger.info("Will be extracting results")
         upload_with_config(
             args.config_path, args.datafile, args.datatype, args.env
         )