diff --git a/pyscal/plotting.py b/pyscal/plotting.py new file mode 100644 index 00000000..f29020e5 --- /dev/null +++ b/pyscal/plotting.py @@ -0,0 +1,541 @@ +"""Module for plotting relative permeability and capillary pressure curves. + +Potential improvements: + Plotting low/base/high curves before interpolation + + Plotting curve sets on the same plot, e.g. curves from different SATNUMs or + low/base/high curves from a SCAL recommendation + + Option to plot GasWater curves with Sg axis (instead of Sw; the Sg column + is already present in the dataframe) + +""" + +from pathlib import Path + +import matplotlib.pyplot as plt +import pandas as pd + +from pyscal import GasOil, GasWater, PyscalList, WaterOil, WaterOilGas, getLogger_pyscal + +logger = getLogger_pyscal(__name__) + +# Data for configuring plot based on pyscal model type +PLOT_CONFIG_OPTIONS = { + "WaterOil": { + "axis": "SW", + "kra_name": "KRW", + "krb_name": "KROW", + "kra_colour": "blue", + "krb_colour": "green", + "pc_name": "PCOW", + "xlabel": "Sw", + "ylabel": "krw, krow", + "curves": "krw_krow", + }, + "GasOil": { + "axis": "SG", + "kra_name": "KRG", + "krb_name": "KROG", + "kra_colour": "red", + "krb_colour": "green", + "xlabel": "Sg", + "ylabel": "krg, krog", + "pc_name": "PCOG", + "curves": "krg_krog", + }, + "GasWater": { + "axis": "SW", + "kra_name": "KRW", + "krb_name": "KRG", + "kra_colour": "blue", + "krb_colour": "red", + "pc_name": "PCGW", + "xlabel": "Sw", + "ylabel": "krg, krw", + "curves": "krg_krw", + }, +} + + +def format_gaswater_table(model: GasWater) -> pd.DataFrame: + """ + + Format the tables held by the GasWater object (through the GasOil and + WaterOil objects). + + The GasWater object is a bit trickier to work with. Like WaterOilGas, + GasWater takes from both WaterOil and GasOil objects, but unlike + WaterOilGas, it is a two-phase model. It is based on one curve from GasOil, + krg, and one curve from WaterOil, krw. This is a different format to the + other models, where the relperm curves to be plotted in the same figure are + found in the same table and are accessed easily using the "table" instance + variable, e.g. WaterOil.table. To plot the correct curves for the GasWater + model, additional formatting is required. That is handled by this function. + + Process: sort both the GasOil and WaterOil tables by increasing SW and then + merge on index. Can't join on "SW" due to potential floating-point number + errors. + + Args: + model (GasWater): GasWater model + + Returns: + pd.DataFrame: saturation table with Sw, Sg, krg from the GasOil + instance, and krw and Pc from the WaterOil instance + """ + + gasoil = model.gasoil.table[["SL", "KRG"]].copy() + gasoil.rename(columns={"SL": "SW"}, inplace=True) + gasoil["SG"] = 1 - gasoil["SW"] + wateroil = model.wateroil.table[["SW", "KRW", "PC"]].copy() + + # Sort by SW to be able to join on index + gasoil.sort_values("SW", ascending=True, inplace=True) + wateroil.sort_values("SW", ascending=True, inplace=True) + gasoil.reset_index(inplace=True, drop=True) + wateroil.reset_index(inplace=True, drop=True) + + # Check if SW in the two models differs using an epsilon of 1E-6 If any + # absolute differences are greater than this threshold, an assertion error + # will be raised. + assert ( + abs(gasoil["SW"] - wateroil["SW"]) < 1e-6 + ).all(), "SW for the GasOil model does not match SW for the WaterOil model" + + # Merge dataframes and format + gaswater = gasoil.merge(wateroil, left_index=True, right_index=True) + gaswater.rename(columns={"SW_x": "SW"}, inplace=True) + gaswater.drop("SW_y", axis=1, inplace=True) + + return gaswater + + +def get_satnum_from_tag(string: str) -> int: + """ + Get SATNUM from the model tag. Used in the naming of figures. + + Args: + string (str): String from the model .tag instance variable + + Returns: + int: SATNUM number + """ + satnum = int(string.split("SATNUM")[1].strip()) + return satnum + + +def get_plot_config_options(curve_type: str, **kwargs) -> dict: + """ + Get config data from plot config dictionary based on the curve (model) type. + + Args: + curve_type (str): Name of the curve type. Allowed types are given in + the PLOT_CONFIG_OPTIONS dictionary + + Returns: + dict: Config parameters for the chosen model type + """ + + config = PLOT_CONFIG_OPTIONS[curve_type].copy() + + # If semilog plot, add suffix to the name of the saved relperm figure + if kwargs["semilog"]: + suffix = "_semilog" + else: + suffix = "" + + config["suffix"] = suffix + + return config + + +def format_relperm_plot(fig: plt.Figure, **kwargs) -> plt.Figure: + """ + Formatting options for individual relative permeability plots. + + Args: + fig (plt.Figure): Relative permeability figure to be formatted + + Returns: + plt.Figure: Formatted relative permeability figure + """ + + ax = fig.gca() + + # Set axis lables + ax.set_xlabel(kwargs["xlabel"]) + ax.set_ylabel(kwargs["ylabel"]) + + # Set axis limits + ax.set_xlim((0.0, 1.0)) + + # Log y-axis fro semilog plot + # Some cases may have kr > 1 depending on base/absolute + # permeability used to calculate kr + if kwargs["semilog"]: + ax.set_yscale("log") + ax.set_ylim((1e-6, 1.0)) + else: + ax.set_ylim((0.0, 1.0)) + + # Add legend + plt.legend() + ax.legend( + loc="upper center", + bbox_to_anchor=(0.5, -0.12), + ncol=5, + ) + + return fig + + +def format_cap_pressure_plot( + fig: plt.Figure, neg_pc: bool = False, **kwargs +) -> plt.Figure: + """ + Formatting options for individual capillary pressure plots. + + Args: + fig (plt.Figure): Capillary pressure figure to be formatted + neg_pc (bool, optional): Negative Pc flag. True if the Pc curve has a + negative part. Defaults to False. + + Returns: + plt.Figure: Formatted capillary pressure figure + """ + + ax = fig.gca() + + # Set axis lables + ax.set_xlabel(kwargs["xlabel"]) + ax.set_ylabel(kwargs["pc_name"].lower().capitalize()) + + # Set axis limits + ax.set_xlim((0.0, 1.0)) + + # Set lower y-axis limit to 0 if Pc >= 0 + if not neg_pc: + ax.set_ylim(bottom=0.0) + + return fig + + +def plot_pc(table: pd.DataFrame, satnum: int, **kwargs) -> plt.Figure: + """ + Plot capillary pressure curves. + + Called if the Pc plot is requested, regardless of if Pc is non-zero. + + Args: + table (pd.DataFrame): Saturation table with Pc curves to be plotted + satnum (int): SATNUM number + + Returns: + plt.Figure: Pc figure with a single Pc curve for a given model and + SATNUM + """ + + fig = plt.figure(1, figsize=(5, 5), dpi=300) + plt.title(f"SATNUM {satnum}") + plt.plot(table[kwargs["axis"]], table["PC"]) + + # Flag for negative Pc + neg_pc = table["PC"].min() < 0 + fig = format_cap_pressure_plot(fig, neg_pc, **kwargs) + + # Log warning if Pc plot is requested but is zero or practically zero. + # There is no checking of units of Pc in pyscal, but this should work as + # intended for Pc in bar, Pa, MPa, atm and psi, i.e. the most common units + # of pressure. In any case, the figures will still be made. + if not (abs(table["PC"]) > 1e-6).any(): + logger.warning("Pc plots were requested, but Pc is zero.") + + return fig + + +def plot_relperm( + table: pd.DataFrame, satnum: int, config: dict, **kwargs +) -> plt.Figure: + """ + + Function for plotting one relperm curve set for each SATNUM. + + Takes kwargs from the plotter() function, which in turn come from the + pyscal CLI, and are passed on to the plot formatting functions. + + Args: + table (pd.DataFrame): Saturation table with curves to be plotted + satnum (int): SATNUM number + config (dict): Plot config + + Returns: + plt.Figure: Relative permeability figure with two relative permeability + curves for a given model and SATNUM + """ + + # Plotting relative permeability curves + fig = plt.figure(1, figsize=(5, 5), dpi=300) + + plt.title(f"SATNUM {satnum}") + + # Plot first relperm curve + plt.plot( + table[config["axis"]], + table[config["kra_name"]], + label=config["kra_name"].lower(), + color=config["kra_colour"], + ) + + # Plot second relperm curve + plt.plot( + table[config["axis"]], + table[config["krb_name"]], + label=config["krb_name"].lower(), + color=config["krb_colour"], + ) + + fig = format_relperm_plot(fig, **kwargs, **config) + + return fig + + +def save_figure( + fig: plt.Figure, + satnum: int, + config: dict, + plot_type: str, + outdir: str, +) -> None: + """ + + Save the provided figure. + + Args: + fig (plt.Figure): Figure to be saved + satnum (int): SATNUM number + config (dict): Plot config + plot_type (str): Figure type. Allowed types are 'relperm' and 'pc' + outdir (str): Directory where the figure will be saved + """ + + # Get curve name + if plot_type == "relperm": + curve_names = config["curves"] + suffix = config["suffix"] + elif plot_type == "pc": + curve_names = config["pc_name"].lower() + suffix = "" + else: + raise ValueError(f"'type' given ({plot_type}) must be one of 'relperm' or 'pc'") + + fname = f"{curve_names}_SATNUM_{satnum}{suffix}".replace(" ", "") + fout = Path(outdir).joinpath(fname) + + fig.savefig( + fout, + bbox_inches="tight", + ) + + print(f"Figure saved to {fout}.png") + + # Clear figure so that it is empty for the next SATNUM's plot + fig.clear() + + +def wog_plotter(model: WaterOilGas, **kwargs) -> None: + """ + + Plot a WaterOilGas (WaterOil and GasOil) model. + + + For a WaterOilGas instance, the WaterOil and GasOil instances can be + accessed, then the "table" instance variable. + + Args: + model (WaterOilGas): WaterOilGas instance + """ + + outdir = kwargs["outdir"] + + # the wateroil and gasoil instance variables are optional for the + # WaterOilGas class. If statements used to check if they are provided + if model.wateroil: + config_wo = get_plot_config_options("WaterOil", **kwargs) + satnum_wo = get_satnum_from_tag(model.wateroil.tag) + + fig_wo = plot_relperm( + model.wateroil.table, + satnum_wo, + config_wo, + **kwargs, + ) + + save_figure(fig_wo, satnum_wo, config_wo, "relperm", outdir) + + if kwargs["pc"]: + fig_pcwo = plot_pc( + model.wateroil.table, + get_satnum_from_tag(model.wateroil.tag), + **config_wo, + ) + + save_figure(fig_pcwo, satnum_wo, config_wo, "pc", outdir) + + if model.gasoil: + config_go = get_plot_config_options("GasOil", **kwargs) + satnum_go = get_satnum_from_tag(model.gasoil.tag) + + fig_go = plot_relperm( + model.gasoil.table, + satnum_go, + config_go, + **kwargs, + ) + + save_figure(fig_go, satnum_go, config_go, "relperm", outdir) + + +def wo_plotter(model: WaterOil, **kwargs) -> None: + """ + + Plot a WaterOil model. + + For a WaterOil instance, the saturation table can be accessed using the + "table" instance variable. + + Args: + model (WaterOil): WaterOil instance + """ + config = get_plot_config_options("WaterOil", **kwargs) + satnum = get_satnum_from_tag(model.tag) + + outdir = kwargs["outdir"] + + fig = plot_relperm( + model.table, + satnum, + config, + **kwargs, + ) + + save_figure(fig, satnum, config, "relperm", outdir) + + if kwargs["pc"]: + fig_pc = plot_pc(model.table, get_satnum_from_tag(model.tag), **config) + + save_figure(fig_pc, satnum, config, "pc", outdir) + + +def go_plotter(model: GasOil, **kwargs) -> None: + """ + + Plot a GasOil model. + + For a GasOil instance, the saturation table can be accessed using the + "table" instance variable. + + Args: + model (GasOil): GasOil instance + """ + + config = get_plot_config_options("GasOil", **kwargs) + satnum = get_satnum_from_tag(model.tag) + + outdir = kwargs["outdir"] + + fig = plot_relperm( + model.table, + satnum, + config, + **kwargs, + ) + + save_figure(fig, satnum, config, "relperm", outdir) + + # Note that there are no supporting functions for adding Pc to GasOil + # instances. This can only be done by modifying the "table" instance + # variable for a GasOil object + if kwargs["pc"]: + fig_pc = plot_pc(model.table, get_satnum_from_tag(model.tag), **config) + + save_figure(fig_pc, satnum, config, "pc", outdir) + + +def gw_plotter(model: GasWater, **kwargs) -> None: + """ + + For GasWater, the format is different, and an additional formatting step is + required. Use the formatted table as an argument to the plotter function, + instead of the "table" instance variable + + Args: + model (GasWater): GasWater instance + """ + + table = format_gaswater_table(model) + config = get_plot_config_options("GasWater", **kwargs) + satnum = get_satnum_from_tag(model.tag) + outdir = kwargs["outdir"] + + fig = plot_relperm( + table, + satnum, + config, + **kwargs, + ) + + save_figure(fig, satnum, config, "relperm", outdir) + + if kwargs["pc"]: + fig_pc = plot_pc(table, get_satnum_from_tag(model.tag), **config) + + save_figure(fig_pc, satnum, config, "pc", outdir) + + +def plotter( + models: PyscalList, pc: bool = False, semilog: bool = False, outdir: str = "./" +) -> None: + """ + + Runner function for creating plots. + + Iterate over PyscalList and plot curves based on type of pyscal objects + encountered. + + PyscalList is a list of WaterOilGas, WaterOil, GasOil or GasWater objects. + + For WaterOil and GasOil, the saturation table can be accessed using the + "table" instance variable. + + For WaterOilGas, the WaterOil and GasOil instances can be accessed, then + the "table" instance variable. + + For GasWater, the format is different, and an additional formatting step is + required. + + Args: + models (PyscalList): List of models + pc (bool, optional): Plot Pc flag. Defaults to False. + semilog (bool, optional): Plot relperm with log y-axis. Defaults to + False. + + """ + + # kwargs to be passed on to other functions + kwargs = {"pc": pc, "semilog": semilog, "outdir": outdir} + + for model in models.pyscal_list: + if isinstance(model, WaterOilGas): + wog_plotter(model, **kwargs) + elif isinstance(model, WaterOil): + wo_plotter(model, **kwargs) + elif isinstance(model, GasOil): + go_plotter(model, **kwargs) + elif isinstance(model, GasWater): + gw_plotter(model, **kwargs) + else: + raise Exception( + f"Model type received was {type(model)} but\ + must be one of: {WaterOil, WaterOilGas, GasOil, GasWater}" + ) diff --git a/pyscal/pyscalcli.py b/pyscal/pyscalcli.py index af70ef5a..50abc275 100644 --- a/pyscal/pyscalcli.py +++ b/pyscal/pyscalcli.py @@ -15,6 +15,7 @@ WaterOilGas, __version__, getLogger_pyscal, + plotting, ) from .factory import PyscalFactory @@ -149,6 +150,36 @@ def get_parser() -> argparse.ArgumentParser: "Implicit for gas-water input." ), ) + + # Plotting arguments + parser.add_argument( + "--plot", + action="store_true", + default=False, + help=("Make and save relative permeability figures."), + ) + parser.add_argument( + "--plot_pc", + action="store_true", + default=False, + help=("Make and save capillary pressure figures."), + ) + parser.add_argument( + "--plot_semilog", + action="store_true", + default=False, + help=( + "Plot relative permeability figures with log y-axis." + "Run both with and without this flag to plot" + "both linear and semi-log relperm plots." + ), + ) + parser.add_argument( + "--plot_outdir", + default="./", + help="Directory where the plot output figures will be saved.", + ) + return parser @@ -176,6 +207,10 @@ def main() -> None: sheet_name=args.sheet_name, slgof=args.slgof, family2=args.family2, + plot=args.plot, + plot_pc=args.plot_pc, + plot_semilog=args.plot_semilog, + plot_outdir=args.plot_outdir, ) except (OSError, ValueError) as err: print("".join(traceback.format_tb(err.__traceback__))) @@ -193,6 +228,10 @@ def pyscal_main( sheet_name: Optional[str] = None, slgof: bool = False, family2: bool = False, + plot: bool = False, + plot_pc: bool = False, + plot_semilog: bool = False, + plot_outdir: str = "./", ) -> None: """A "main()" method not relying on argparse. This can be used for testing, and also by an ERT forward model, e.g. @@ -209,6 +248,10 @@ def pyscal_main( sheet_name: Which sheet in XLSX file slgof: Use SLGOF family2: Dump family 2 keywords + plot: Plot relative permeability figures and save + plot_pc: Plot capillary pressure curves in addition to relperm curves + plot_semilog: Plot relative permeability figures with log y-axis + """ logger = getLogger_pyscal( @@ -274,3 +317,6 @@ def pyscal_main( with open(output, "w", newline="\n", encoding="utf-8") as fh: fh.write(wog_list.build_eclipse_data(family=family, slgof=slgof)) print("Written to " + output) + + if plot: + plotting.plotter(wog_list, plot_pc, plot_semilog, plot_outdir) diff --git a/tests/test_plotting.py b/tests/test_plotting.py new file mode 100644 index 00000000..2500b5da --- /dev/null +++ b/tests/test_plotting.py @@ -0,0 +1,239 @@ +"""Test the plotting module""" + +from pathlib import Path + +import matplotlib.pyplot as plt +import pytest + +from pyscal import GasOil, GasWater, PyscalList, WaterOil, WaterOilGas, plotting + + +def test_get_satnum_from_tag(): + """Check that the SATNUM number can be retrieved from the model tag""" + # Several PyscalLists of different model types to be checked + pyscal_lists = [ + PyscalList( + [ + WaterOil(tag="SATNUM 1"), + WaterOil(tag="SATNUM 2"), + WaterOil(tag="SATNUM 3"), + ] + ), + PyscalList( + [ + GasWater(tag="SATNUM 1"), + GasWater(tag="SATNUM 2"), + GasWater(tag="SATNUM 3"), + ] + ), + PyscalList( + [GasOil(tag="SATNUM 1"), GasOil(tag="SATNUM 2"), GasOil(tag="SATNUM 3")] + ), + PyscalList( + [ + WaterOilGas(tag="SATNUM 1"), + WaterOilGas(tag="SATNUM 2"), + WaterOilGas(tag="SATNUM 3"), + ] + ), + ] + + # Test that an integer is returned from this function + for pyscal_list in pyscal_lists: + for model in pyscal_list.pyscal_list: + satnum = plotting.get_satnum_from_tag(model.tag) + assert isinstance(satnum, int) + + # Test that this function doesn't work if SATNUM is represented as a float + # in the model.tag string + with pytest.raises(ValueError): + satnum = plotting.get_satnum_from_tag(WaterOil(tag="SATNUM 3.0").tag) + + +def test_plotter(): + """Check if an Exception is raised if a model type is not included. This is done + to check that all models have been implemented in the plotting module.""" + + class DummyPyscalList: + """ + Can't use the actual PyscalList, as this will raise its own exception + (DummyModel is not a pyscal object), so a dummy PyscalList is used + + #If the PyscalList.pyscal_list instance variable name changes, this + will still pass...""" + + def __init__(self, models: list) -> None: + self.pyscal_list = models + + class DummyModel: + """Dummy model""" + + def __init__(self, tag: str) -> None: + self.tag = tag + + dummy_model_list = [ + DummyModel("SATNUM 1"), + DummyModel("SATNUM 2"), + DummyModel("SATNUM 3"), + ] + + dummy_pyscal_list = DummyPyscalList(dummy_model_list) + + with pytest.raises(Exception): + plotting.plotter(dummy_pyscal_list) + + +def test_pyscal_list_attr(): + """Check that the PyscalList class has an pyscal_list instance variable. + This is accessed by the plotting module to loop through models to plot.""" + assert ( + hasattr(PyscalList(), "pyscal_list") is True + ), "The PyscalList object should have a pyscal_list instance variable.\ + This is accessed by the plotting module." + + +def test_plot_relperm(): + """Test that a matplotlib.pyplot Figure instance is returned""" + wateroil = WaterOil(swl=0.1, h=0.1) + wateroil.add_corey_water() + wateroil.add_corey_oil() + + kwargs = {"semilog": False, "pc": False} + config = plotting.get_plot_config_options("WaterOil", **kwargs) + fig = plotting.plot_relperm(wateroil.table, 1, config, **kwargs) + + assert isinstance( + fig, plt.Figure + ), "Type of returned object is not a matplotlib.pyplot Figure" + + +def test_plot_pc(): + """Test that a matplotlib.pyplot Figure instance is returned""" + wateroil = WaterOil(swl=0.1, h=0.1) + wateroil.add_corey_water() + wateroil.add_corey_oil() + wateroil.add_simple_J() + + kwargs = {"semilog": False, "pc": True} + config = plotting.get_plot_config_options("WaterOil", **kwargs) + fig = plotting.plot_pc(wateroil.table, 1, **config) + + assert isinstance( + fig, plt.Figure + ), "Type of returned object is not a matplotlib.pyplot Figure" + + +def test_wog_plotter(tmpdir): + """Test that relative permeability figures are created by the plotter function""" + wateroil = WaterOil(swl=0.1, h=0.1, tag="SATNUM 1") + wateroil.add_corey_water() + wateroil.add_corey_oil() + wateroil.add_simple_J() + + gasoil = GasOil(swl=0.1, h=0.1, tag="SATNUM 1") + gasoil.add_corey_gas() + gasoil.add_corey_oil() + + wateroilgas = WaterOilGas() + wateroilgas.wateroil = wateroil + wateroilgas.gasoil = gasoil + + # Testing both rel perm and Pc plots are created + kwargs = {"semilog": False, "pc": True, "outdir": tmpdir} + kr_ow_plot_name = "krw_krow_SATNUM_1.png" + kr_og_plot_name = "krw_krow_SATNUM_1.png" + pc_ow_plot_name = "pcow_SATNUM_1.png" + + plotting.wog_plotter(wateroilgas, **kwargs) + + assert Path.exists(Path(tmpdir).joinpath(kr_ow_plot_name)), "Plot not created" + assert Path.exists(Path(tmpdir).joinpath(kr_og_plot_name)), "Plot not created" + assert Path.exists(Path(tmpdir).joinpath(pc_ow_plot_name)), "Plot not created" + + +def test_wo_plotter(tmpdir): + """Test that relative permeability figures are created by the plotter function""" + wateroil = WaterOil(swl=0.1, h=0.1, tag="SATNUM 1") + wateroil.add_corey_water() + wateroil.add_corey_oil() + wateroil.add_simple_J() + + # Testing both rel perm and Pc plots are created + kwargs = {"semilog": False, "pc": True, "outdir": tmpdir} + kr_plot_name = "krw_krow_SATNUM_1.png" + pc_plot_name = "pcow_SATNUM_1.png" + + plotting.wo_plotter(wateroil, **kwargs) + + assert Path.exists(Path(tmpdir).joinpath(kr_plot_name)), "Plot not created" + assert Path.exists(Path(tmpdir).joinpath(pc_plot_name)), "Plot not created" + + +def test_wo_plotter_relperm_only(tmpdir): + """Test that relative permeability figures are created by the plotter function""" + wateroil = WaterOil(swl=0.1, h=0.1, tag="SATNUM 1") + wateroil.add_corey_water() + wateroil.add_corey_oil() + wateroil.add_simple_J() + + # Testing only relperm plots are created + kwargs = {"semilog": False, "pc": False, "outdir": tmpdir} + kr_plot_name = "krw_krow_SATNUM_1.png" + pc_plot_name = "pcow_SATNUM_1.png" + + plotting.wo_plotter(wateroil, **kwargs) + + assert Path.exists(Path(tmpdir).joinpath(kr_plot_name)), "Plot not created" + assert not Path.exists( + Path(tmpdir).joinpath(pc_plot_name) + ), "Plot created when it shouldn't have been" + + +def test_go_plotter(tmpdir): + """Test that relative permeability figures are created by the plotter function""" + gasoil = GasOil(swl=0.1, h=0.1, tag="SATNUM 1") + gasoil.add_corey_gas() + gasoil.add_corey_oil() + + # Add Pc manually as there are currently no supporting functions for adding + # Pc to GasOil instances + gasoil.table["PC"] = 1 * len(gasoil.table) + + # Testing both rel perm and Pc plots are created + kwargs = {"semilog": False, "pc": True, "outdir": tmpdir} + kr_plot_name = "krg_krog_SATNUM_1.png" + pc_plot_name = "pcog_SATNUM_1.png" + + plotting.go_plotter(gasoil, **kwargs) + + assert Path.exists(Path(tmpdir).joinpath(kr_plot_name)), "Plot not created" + assert Path.exists(Path(tmpdir).joinpath(pc_plot_name)), "Plot not created" + + +def test_gw_plotter(tmpdir): + """Test that relative permeability figures are created by the plotter function""" + gaswater = GasWater(swl=0.1, h=0.1, tag="SATNUM 1") + gaswater.add_corey_water() + gaswater.add_corey_gas() + gaswater.add_simple_J() + + # Testing both rel perm and Pc plots are created + kwargs = {"semilog": False, "pc": True, "outdir": tmpdir} + kr_plot_name = "krg_krw_SATNUM_1.png" + pc_plot_name = "pcgw_SATNUM_1.png" + + plotting.gw_plotter(gaswater, **kwargs) + + assert Path.exists(Path(tmpdir).joinpath(kr_plot_name)), "Plot not created" + assert Path.exists(Path(tmpdir).joinpath(pc_plot_name)), "Plot not created" + + +def test_save_figure(tmpdir): + """Test that figure is saved""" + fig = plt.Figure() + + config = {"curves": "dummy", "suffix": ""} + fig_name = "dummy_SATNUM_1.png" + plotting.save_figure(fig, 1, config, plot_type="relperm", outdir=tmpdir) + + assert Path.exists(Path(tmpdir).joinpath(fig_name)), "Figure not saved" diff --git a/tests/test_pyscalcli.py b/tests/test_pyscalcli.py index aabb459d..543aea78 100644 --- a/tests/test_pyscalcli.py +++ b/tests/test_pyscalcli.py @@ -586,3 +586,44 @@ def test_pyscal_main(): with pytest.raises(ValueError, match="Interpolation parameter provided"): pyscalcli.pyscal_main(relperm_file, int_param_wo=-1, output=os.devnull) + + +def test_pyscalcli_plot(capsys, mocker, tmpdir): + """Test that plots are created through the CLI. This is done by testing to + see if the print statements in the save_figure function are present in stdout""" + scalrec_file = Path(__file__).absolute().parent / "data/scal-pc-input-example.xlsx" + + mocker.patch( + "sys.argv", + [ + "pyscal", + str(scalrec_file), + "--int_param_wo", + "0", + "--output", + "-", + "--plot", + "--plot_pc", + "--plot_outdir", + str(tmpdir), + ], + ) + + pyscalcli.main() + + expected_plots = [ + "krw_krow_SATNUM_1.png", + "krg_krog_SATNUM_1.png", + "krw_krow_SATNUM_2.png", + "krg_krog_SATNUM_2.png", + "krw_krow_SATNUM_3.png", + "krg_krog_SATNUM_3.png", + "pcow_SATNUM_1.png", + "pcow_SATNUM_2.png", + "pcow_SATNUM_3.png", + ] + + captured = capsys.readouterr() + + for plot in expected_plots: + assert plot in captured.out