A 2D benchmark of excavation modelling under two-phase flow condition

import os
from pathlib import Path

import gmsh
import matplotlib.pyplot as plt
import numpy as np
import ogstools as ot
from mesh import MeshGenerator

out_dir = Path(os.environ.get("OGS_TESTRUNNER_OUT_DIR", "_out"))
if not out_dir.exists():
    out_dir.mkdir(parents=True)
mesh_dir = Path(out_dir, "Mesh")

point_a = (6.5, -857.0, 0)
point_b = (70.0, -857.0, 0)

radius_tunnel = 6.5
points = np.asarray(
    [
        [x, -857, 0]
        for x in [
            radius_tunnel,
            radius_tunnel + 1,
            radius_tunnel + 2.5,
            radius_tunnel + 3.5,
            radius_tunnel + 8.5,
        ]
    ]
)
labels = [f"r = {point[0]:.1f} m" for i, point in enumerate(points)]

saturation = ot.variables.Scalar(
    data_name="saturation",
    output_name="Saturation",
    data_unit="",
    output_unit="",
)
pg_interpolated = ot.variables.Scalar(
    data_name="gas_pressure_interpolated",
    output_name="Gas pressure",
    data_unit="Pa",
    output_unit="MPa",
)
pc_interpolated = ot.variables.Scalar(
    data_name="capillary_pressure_interpolated",
    output_name="capillary pressure",
    data_unit="Pa",
    output_unit="MPa",
)
pl_interpolated = ot.variables.Scalar(
    data_name="liquid_pressure_interpolated",
    output_name="Liquid pressure",
    data_unit="Pa",
    output_unit="MPa",
)
pls_strain = ot.variables.Scalar(
    data_name="epsilon_pls",
    output_name="Equivalent plastic strain",
    data_unit="",
    output_unit="",
)

mohr_coulomb = {
    "Cohesion": 5e6,  # Large number for fast run
    "FrictionAngle": 25.0,
    "DilatancyAngle": 2.0,
    "TransitionAngle": 25.0,
    "TensionCutOffParameter": 0.5,
}

output_variable_names = [
    "displacement",
    "sigma",
    "gas_pressure",
    "capillary_pressure",
    "saturation",
]
comparison_atols = {
    "displacement": 1e-9,
    "sigma": 1e-6,
    "gas_pressure": 1e-9,
    "capillary_pressure": 1e-9,
    "saturation": 1e-6,
}

# Compare the obtained results with the reference results
def check_profile(profile, reference_file_name):
    reference_profile = np.load(reference_file_name)
    for variable_name in output_variable_names:
        atol = comparison_atols[variable_name]
        data = profile[variable_name]
        expected_data = reference_profile[variable_name]
        np.testing.assert_allclose(data, expected_data, atol=atol, rtol=1e-6)

    print("Test passed.")

if not gmsh.isInitialized():
    gmsh.initialize()

gmsh.model.add("Mesh")

mesh_generator = MeshGenerator(
    gmsh_model=gmsh.model,
    min_edge_length=1.0,
    max_edge_length=12.0,
    quad_zone_factor=0.25,
    n_side=15,
)
mesh_generator.generate_meshes(out_dir=mesh_dir, order=2)

gmsh.finalize()

def remove_plastic_model(prj):
    prj.remove_element("./processes/process/constitutive_relation[type='MFront']")
    prj.remove_element(
        "./processes/process/secondary_variables/secondary_variable[@internal_name='EquivalentPlasticStrain']"
    )
    prj.remove_element("./time_loop/output/variables/variable[7]")

def runBenchmark(
    project_file,
    out_dir,
    output_prefix,
    mesh_dir,
    use_taylor_hood=True,
    plastic_model=None,
    ogs_path=None,
):
    prj = ot.Project(
        input_file=project_file,
        output_file=Path(out_dir, f"{output_prefix}.prj"),
        OMP_NUM_THREADS=4,
        OGS_ASM_THREADS=4,
    )

    if use_taylor_hood:
        xpath_var = './process_variables/process_variable[name="displacement"]/'
        prj.replace_text(
            2,
            xpath=xpath_var + "order",
        )

    if plastic_model:
        print("Use plastic model")
        for key, value in plastic_model.items():
            prj.replace_text(value, f"./parameters/parameter[name='{key}']/value")
    else:
        print("Use elastic model")
        prj.add_block(
            blocktag="constitutive_relation",
            parent_xpath="./processes/process",
            taglist=["type", "youngs_modulus", "poissons_ratio"],
            textlist=["LinearElasticIsotropic", "E", "nu"],
        )
        remove_plastic_model(prj)

    prj.replace_text(
        output_prefix,
        xpath="./time_loop/output/prefix",
    )

    prj.write_input()

    prj.run_model(
        logfile=Path(out_dir, "log.txt"),
        path=ogs_path,
        args=f"-o {out_dir} -m {mesh_dir}",
    )

    print(f"Output prefix is {output_prefix}")

    return Path(out_dir, f"{output_prefix}.pvd")

def contour_plot(mesh_data, variables, cmap="jet", file_name=None):
    fig, axs = plt.subplots(2, 2, figsize=(8, 6), sharex=True, sharey=True)
    axs_flatten = axs.flatten()
    for ax, variable in zip(axs_flatten, variables, strict=True):
        ot.plot.contourf(
            mesh_data,
            variable,
            fig=fig,
            ax=ax,
            fontsize=8,
            cmap=cmap,
        )

    fig.tight_layout()

    if file_name:
        plt.savefig(file_name, dpi=320)

class PlotProfile:
    def __init__(self, profile_e, profile_pls):
        self.xs = profile_e["Distance"] + 6.5
        self.profile_e = profile_e
        self.profile_pls = profile_pls

    def get_data(self, variable_name):
        if variable_name == "sigma_xx":
            return (
                self.profile_e["sigma"][:, 0] * 1e-6,
                self.profile_pls["sigma"][:, 0] * 1e-6,
            )
        if variable_name == "sigma_yy":
            return (
                self.profile_e["sigma"][:, 1] * 1e-6,
                self.profile_pls["sigma"][:, 1] * 1e-6,
            )
        scaling = 1e-6 if "pressure" in variable_name else 1.0
        return (
            self.profile_e[variable_name] * scaling,
            self.profile_pls[variable_name] * scaling,
        )

    def plot(self, variable_names, y_labels):
        fig, axs = plt.subplots(1, 2, figsize=(10, 4), sharey=False)

        for ax, variable_name, y_label in zip(
            axs, variable_names, y_labels, strict=True
        ):
            data_e, data_pls = self.get_data(variable_name)
            ax.plot(self.xs, data_e, color="C0", label="elastic")
            ax.plot(self.xs, data_pls, color="C7", label="plastic")

            ax.set_xlabel("x / m")
            ax.set_ylabel(y_label)
            ax.grid()
            ax.legend()

        fig.tight_layout()

        plt.show()

class VariableVariations:
    def __init__(self, point_data_e, point_data_pls):
        self.point_data_e = point_data_e
        self.point_data_pls = point_data_pls
        self.labels_e = (np.array(labels) + " (elastic)").tolist()
        self.labels_pls = (np.array(labels) + " (plastic)").tolist()

    def plot(self, variables):
        fig, axs = plt.subplots(1, 2, figsize=(12, 4))

        for ax, variable in zip(axs, variables, strict=True):
            ot.plot.line(
                self.point_data_e,
                "time",
                variable,
                labels=self.labels_e,
                linestyle="dashed",
                ax=ax,
                fontsize=11,
                linewidth=0.5,
            )
            ot.plot.line(
                self.point_data_pls,
                "time",
                variable,
                labels=self.labels_pls,
                ax=ax,
                fontsize=11,
                linewidth=0.5,
            )
            ax.grid()
            ax.legend(loc="center left", bbox_to_anchor=(1, 0.5))

        fig.tight_layout()
        plt.show()

pvd = runBenchmark(
    project_file="excavation_th2m.prj",
    out_dir=out_dir,
    output_prefix="excavation_th2m_e",
    mesh_dir=mesh_dir,
)

contour_plot(
    ms_e_last,
    [
        ot.variables.displacement["x"],
        ot.variables.displacement["y"],
        ot.variables.stress["xx"],
        ot.variables.stress["yy"],
    ],
)

contour_plot(
    ms_e_last,
    [saturation, pc_interpolated, pg_interpolated, pl_interpolated],
    cmap="coolwarm",
)

x_profile_e = ms_e_last.sample_over_line(point_a, point_b, resolution=55)

## Write the reference results.
# fields = {name: x_profile_e[name] for name in output_variable_names}
# np.savez("reference_profile_e.npz", **fields)

check_profile(x_profile_e, "reference_profile_e.npz")

pvd_pls = runBenchmark(
    project_file="excavation_th2m.prj",
    out_dir=out_dir,
    output_prefix="excavation_th2m_pls",
    mesh_dir=mesh_dir,
    plastic_model=mohr_coulomb,
)

contour_plot(
    ms_pls_last,
    [
        ot.variables.displacement["x"],
        ot.variables.displacement["y"],
        ot.variables.stress["xx"],
        pls_strain,
    ],
)

contour_plot(
    ms_pls_last,
    [saturation, pc_interpolated, pg_interpolated, pl_interpolated],
    cmap="coolwarm",
)

x_profile_pls = ms_pls_last.sample_over_line(point_a, point_b, resolution=55)

## Write the reference results.
# fields = {name: x_profile_pls[name] for name in output_variable_names}
# np.savez("reference_profile_pls.npz", **fields)

check_profile(x_profile_pls, "reference_profile_pls.npz")

fig, ax = plt.subplots(1, 2, figsize=(10, 4))

ax[0].plot(x_profile_pls["Distance"] + 6.5, x_profile_pls["epsilon_pls"], color="C0")
ax[0].set_xlabel("x / m")
ax[0].set_ylabel("Equivalent plastic strain")
ax[0].set_title("Profile at the bottom")
ax[0].grid()

ms_pts_pls = ot.MeshSeries.extract_probe(ms_pls, points)
ot.plot.line(
    ms_pts_pls,
    "time",
    pls_strain,
    labels=labels,
    # linestyle="dashed",
    ax=ax[1],
    fontsize=11,
    linewidth=0.5,
)
ax[1].set_title("Variations at the selected bottom points")

fig.tight_layout()
plt.show()

plot_profile.plot(
    ["sigma_xx", "sigma_yy"],
    [r"Stress $\sigma_{xx}$ /MPa", r"Stress $\sigma_{yy}$ /MPa"],
)

plot_profile.plot(
    ["gas_pressure_interpolated", "liquid_pressure_interpolated"],
    ["Gas pressure /MPa", "Liquid pressure /MPa"],
)

plot_profile.plot(
    ["capillary_pressure_interpolated", "saturation"],
    ["Capillary pressure /MPa", "Saturation"],
)