gprMax/reframe_tests/tests/mixins.py

from pathlib import Path
from shutil import copyfile
from typing import Optional

import reframe.utility.sanity as sn
import reframe.utility.typecheck as typ
from numpy import prod
from reframe import RegressionMixin
from reframe.core.builtins import parameter, run_after, variable
from typing_extensions import TYPE_CHECKING

from reframe_tests.tests.base_tests import GprMaxBaseTest
from reframe_tests.tests.regression_checks import (
    GeometryObjectMaterialsRegressionCheck,
    GeometryObjectRegressionCheck,
    GeometryViewRegressionCheck,
    H5RegressionCheck,
    ReceiverRegressionCheck,
    SnapshotRegressionCheck,
)

# If using a static type checker, inherit from GprMaxBaseTest as the
# Mixin classes should always have access to resources from that class.
# However, during execution inherit from RegressionMixin.
if TYPE_CHECKING:
    GprMaxMixin = GprMaxBaseTest
else:
    GprMaxMixin = RegressionMixin


class ReceiverMixin(GprMaxMixin):
    number_of_receivers = variable(int, value=-1)

    @run_after("setup", always_last=True)
    def add_receiver_regression_checks(self):
        test_dependency = self.get_test_dependency()
        if test_dependency is not None:
            reference_file = self.build_reference_filepath(test_dependency.output_file)
        else:
            reference_file = self.build_reference_filepath(self.output_file)

        if self.number_of_receivers > 0:
            for i in range(self.number_of_receivers):
                regression_check = ReceiverRegressionCheck(
                    self.output_file, reference_file, f"r{i}"
                )
                self.regression_checks.append(regression_check)
        else:
            regression_check = H5RegressionCheck(self.output_file, reference_file)
            self.regression_checks.append(regression_check)


class SnapshotMixin(GprMaxMixin):
    """Add regression tests for snapshots.

    Attributes:
        snapshots (list[str]): List of snapshots to run regression
            checks on.
    """

    snapshots = variable(typ.List[str], value=[])

    def build_snapshot_filepath(self, snapshot: str) -> Path:
        """Build filepath to the specified snapshot.

        Args:
            snapshot: Name of the snapshot.
        """
        return Path(f"{self.model}_snaps", snapshot).with_suffix(".h5")

    @run_after("setup")
    def add_snapshot_regression_checks(self):
        """Add a regression check for each snapshot.

        The test will be skipped if no snapshots have been specified.
        """
        self.skip_if(
            len(self.snapshots) < 0,
            f"Must provide a list of snapshots.",
        )

        for snapshot in self.snapshots:
            snapshot_file = self.build_snapshot_filepath(snapshot)
            reference_file = self.build_reference_filepath(snapshot)
            regression_check = SnapshotRegressionCheck(snapshot_file, reference_file)
            self.regression_checks.append(regression_check)


class GeometryOnlyMixin(GprMaxMixin):
    """Run test with geometry only flag"""

    @run_after("setup")
    def add_geometry_only_flag(self):
        self.executable_opts += ["--geometry-only"]


class GeometryObjectMixinBase(GprMaxMixin):
    def build_geometry_object_filepath(self, geometry_object: str) -> Path:
        """Build filepath to the specified geometry object.

        Args:
            geometry_object: Name of the geometry object.
        """
        return Path(geometry_object).with_suffix(".h5")

    def build_materials_filepath(self, geometry_object: str) -> Path:
        """Build filepath to the materials output by the geometry object.

        Args:
            geometry_object: Name of the geometry object.
        """
        return Path(f"{geometry_object}_materials").with_suffix(".txt")


class GeometryObjectsReadMixin(GeometryObjectMixinBase):
    geometry_objects_read = variable(typ.Dict[str, str], value={})

    @run_after("setup")
    def copy_geometry_objects_from_test_dependency(self):
        self.skip_if(
            len(self.geometry_objects_read) < 0,
            f"Must provide a list of geometry objects being read by the test.",
        )

        test_dependency = self.get_test_dependency()

        self.skip_if(
            test_dependency is None,
            f"GeometryObjectsReadMixin must be used with a test dependency.",
        )

        for geometry_object_input, geometry_object_output in self.geometry_objects_read.items():
            geometry_object_input_file = self.build_geometry_object_filepath(geometry_object_input)
            geometry_object_input_file = Path(test_dependency.stagedir, geometry_object_input_file)

            materials_input_file = self.build_materials_filepath(geometry_object_input)
            materials_input_file = Path(test_dependency.stagedir, materials_input_file)

            self.skip_if(
                not sn.path_isfile(geometry_object_input_file),
                f"Test dependency did not create geometry object file.",
            )

            self.skip_if(
                not sn.path_isfile(materials_input_file),
                f"Test dependency did not create geometry object materials file.",
            )

            geometry_object_output_file = self.build_geometry_object_filepath(
                geometry_object_output
            )
            geometry_object_output_file = Path(self.stagedir, geometry_object_output_file)

            materials_output_file = self.build_materials_filepath(geometry_object_output)
            materials_output_file = Path(self.stagedir, materials_output_file)

            copyfile(geometry_object_input_file, geometry_object_output_file)
            copyfile(materials_input_file, materials_output_file)


class GeometryObjectsWriteMixin(GeometryObjectMixinBase):
    """Add regression tests for geometry objects.

    Attributes:
        geometry_objects_write (list[str]): List of geometry objects to
            run regression checks on.
    """

    geometry_objects_write = variable(typ.List[str], value=[])

    @run_after("setup")
    def add_geometry_object_regression_checks(self):
        """Add a regression check for each geometry object.

        The test will be skipped if no geometry objects have been specified.
        """
        self.skip_if(
            len(self.geometry_objects_write) < 0,
            f"Must provide a list of geometry objects being written by the test.",
        )

        for geometry_object in self.geometry_objects_write:
            # Add materials regression check first as if this fails,
            # checking the .h5 file will almost definitely fail.
            materials_file = self.build_materials_filepath(geometry_object)
            materials_reference_file = self.build_reference_filepath(
                materials_file.name, suffix=materials_file.suffix
            )
            materials_regression_check = GeometryObjectMaterialsRegressionCheck(
                materials_file, materials_reference_file
            )
            self.regression_checks.append(materials_regression_check)

            geometry_object_file = self.build_geometry_object_filepath(geometry_object)
            reference_file = self.build_reference_filepath(
                geometry_object, suffix=geometry_object_file.suffix
            )
            regression_check = GeometryObjectRegressionCheck(geometry_object_file, reference_file)
            self.regression_checks.append(regression_check)


class GeometryViewMixin(GprMaxMixin):
    """Add regression tests for geometry views.

    Attributes:
        geometry_views (list[str]): List of geometry views to run
            regression checks on.
    """

    geometry_views = variable(typ.List[str], value=[])

    def build_geometry_view_filepath(self, geometry_view: str) -> Path:
        """Build filepath to the specified geometry view.

        Args:
            geometry_view: Name of the geometry view.
        """
        return Path(geometry_view).with_suffix(".vtkhdf")

    @run_after("setup")
    def add_geometry_view_regression_checks(self):
        """Add a regression check for each geometry view.

        The test will be skipped if no geometry views have been specified.
        """
        self.skip_if(
            len(self.geometry_views) < 0,
            f"Must provide a list of geometry views.",
        )

        for geometry_view in self.geometry_views:
            geometry_view_file = self.build_geometry_view_filepath(geometry_view)
            reference_file = self.build_reference_filepath(geometry_view, ".vtkhdf")
            regression_check = GeometryViewRegressionCheck(geometry_view_file, reference_file)
            self.regression_checks.append(regression_check)


class PythonApiMixin(GprMaxMixin):
    """Use the GprMax Python API rather than a standard input file."""

    @run_after("setup")
    def use_python_input_file(self):
        """Input files for API tests will be python files."""
        self.executable = "time -p python"
        self.input_file = self.input_file.with_suffix(".py")


class MpiMixin(GprMaxMixin):
    """Run test using GprMax MPI functionality.

    Attributes:
        mpi_layout (parameter[list[int]]): ReFrame parameter to specify
            how MPI tasks should be arranged.
    """

    mpi_layout = parameter()

    @run_after("setup")
    def configure_mpi_tasks(self):
        """Set num_tasks and add MPI specific commandline arguments."""
        self.num_tasks = int(prod(self.mpi_layout))
        self.executable_opts += ["--mpi", *map(str, self.mpi_layout)]


class BScanMixin(GprMaxMixin):
    """Test a B-scan model - a model with a moving source and receiver.

    Attributes:
        num_models (parameter[int]): Number of models to run.
    """

    num_models = parameter()

    @run_after("setup")
    def setup_bscan_test(self):
        """Add B-scan specific commandline arguments and postrun cmds.

        Set the number of models to run, and merge the output files.
        """
        self.executable_opts += ["-n", str(self.num_models)]

        self.postrun_cmds += [
            f"python -m toolboxes.Utilities.outputfiles_merge {self.model}",
            f"mv {self.model}_merged.h5 {self.output_file}",
        ]

    def get_test_dependency_variant_name(self, **kwargs) -> Optional[str]:
        """Get unique ReFrame name of the test dependency variant.

        By default, filter test dependencies by the model name and the
        number of models.

        Args:
            **kwargs: Additional key-value pairs to filter the parameter
                space of the test dependency. The key is the test
                parameter name and the value is either a single value or
                a unary function that evaluates to True if the parameter
                point must be kept, False otherwise.

        Returns:
            variant_name: Unique name of the test dependency variant.
        """

        kwargs.setdefault("num_models", self.num_models)
        return super().get_test_dependency_variant_name(**kwargs)


class TaskfarmMixin(GprMaxMixin):
    """Run test using GprMax taskfarm functionality."""

    # TODO: Make this a required variabe, or create a new variable to
    # proxy it.
    # num_tasks = required

    @run_after("setup")
    def add_taskfarm_flag(self):
        """Add taskfarm specific commandline arguments."""
        self.executable_opts += ["--taskfarm"]


class AntennaModelMixin(GprMaxMixin):
    """Test an antenna model."""

    pass