Source code for mdsuite.experiment.run

"""
This program and the accompanying materials are made available under the terms of the
Eclipse Public License v2.0 which accompanies this distribution, and is available at
https://www.eclipse.org/legal/epl-v20.html
SPDX-License-Identifier: EPL-2.0

Copyright Contributors to the Zincware Project.

Description: Collection of calculators / transformations for exp.run
"""
from __future__ import annotations

import functools
from typing import TYPE_CHECKING, Any, List, Type, Union

from mdsuite.calculators import (
    AngularDistributionFunction,  # SpatialDistributionFunction,
)
from mdsuite.calculators import (  # StructureFactor,
    CoordinationNumbers,
    EinsteinDiffusionCoefficients,
    EinsteinDistinctDiffusionCoefficients,
    EinsteinHelfandIonicConductivity,
    EinsteinHelfandThermalConductivity,
    EinsteinHelfandThermalKinaci,
    GreenKuboDiffusionCoefficients,
    GreenKuboDistinctDiffusionCoefficients,
    GreenKuboIonicConductivity,
    GreenKuboThermalConductivity,
    GreenKuboViscosity,
    GreenKuboViscosityFlux,
    KirkwoodBuffIntegral,
    NernstEinsteinIonicConductivity,
    PotentialOfMeanForce,
    RadialDistributionFunction,
)
from mdsuite.transformations import (
    CoordinateUnwrapper,
    CoordinateWrapper,
    IntegratedHeatCurrent,
    IonicCurrent,
    KinaciIntegratedHeatCurrent,
    MolecularMap,
    MomentumFlux,
    ScaleCoordinates,
    ThermalFlux,
    Transformations,
    TranslationalDipoleMoment,
    UnwrapViaIndices,
    VelocityFromPositions,
)

if TYPE_CHECKING:
    from mdsuite.experiment import Experiment


class RunComputation:
    """Collection of all calculators that can be used by an experiment"""

    def __init__(
        self, experiment: Experiment = None, experiments: List[Experiment] = None
    ):
        """Collection of all calculators

        Parameters
        ----------
        experiment: Experiment
            Experiment to run the computations for
        experiments: List[Experiment]
            A list of experiments passed by running the computation from the project
            class
        """
        self.experiment = experiment
        self.experiments = experiments

        self.kwargs = {"experiment": experiment, "experiments": experiments}

    def exp_wrapper(self, func):
        # preparation for https://github.com/zincware/MDSuite/issues/404
        # currently this does nothing
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            func_instance = func(*args, **kwargs)
            # self.experiment._run(func_instance)
            return func_instance

        return wrapper

    def transformation_wrapper(self, func: Union[Type[Transformations], Any]):
        """Run the transformation for every selected experiment

        Parameters
        ----------
        func: a transformation to be attached to the experiment/s
        """

        @functools.wraps(func.run_transformation)
        def wrapper(*args, **kwargs):
            if self.experiments is None:
                self.experiments = [self.experiment]
            for experiment in self.experiments:
                func_instance = func()
                # attach the transformation to the experiment
                experiment.cls_transformation_run(func_instance, *args, **kwargs)

        return wrapper

    #######################
    ### Transformations ###
    #######################

    @property
    def CoordinateWrapper(self) -> Type[CoordinateWrapper]:
        return self.transformation_wrapper(CoordinateWrapper)

    @property
    def CoordinateUnwrapper(self) -> Type[CoordinateUnwrapper]:
        return self.transformation_wrapper(CoordinateUnwrapper)

    @property
    def IntegratedHeatCurrent(self) -> Type[IntegratedHeatCurrent]:
        return self.transformation_wrapper(IntegratedHeatCurrent)

    @property
    def IonicCurrent(self) -> Type[IonicCurrent]:
        return self.transformation_wrapper(IonicCurrent)

    @property
    def KinaciIntegratedHeatCurrent(self) -> Type[KinaciIntegratedHeatCurrent]:
        return self.transformation_wrapper(KinaciIntegratedHeatCurrent)

    @property
    def MolecularMap(self) -> Type[MolecularMap]:
        return self.transformation_wrapper(MolecularMap)

    @property
    def MomentumFlux(self) -> Type[MomentumFlux]:
        return self.transformation_wrapper(MomentumFlux)

    @property
    def ScaleCoordinates(self) -> Type[ScaleCoordinates]:
        return self.transformation_wrapper(ScaleCoordinates)

    @property
    def ThermalFlux(self) -> Type[ThermalFlux]:
        return self.transformation_wrapper(ThermalFlux)

    @property
    def TranslationalDipoleMoment(self) -> Type[TranslationalDipoleMoment]:
        return self.transformation_wrapper(TranslationalDipoleMoment)

    @property
    def UnwrapViaIndices(self) -> Type[UnwrapViaIndices]:
        return self.transformation_wrapper(UnwrapViaIndices)

    @property
    def VelocityFromPositions(self) -> Type[VelocityFromPositions]:
        return self.transformation_wrapper(VelocityFromPositions)

    #####################
    #### Calculators ####
    #####################
    @property
    def AngularDistributionFunction(self) -> AngularDistributionFunction:
        return self.exp_wrapper(AngularDistributionFunction)(**self.kwargs)

    @property
    def CoordinationNumbers(self) -> CoordinationNumbers:
        return self.exp_wrapper(CoordinationNumbers)(**self.kwargs)

    @property
    def EinsteinDiffusionCoefficients(self) -> EinsteinDiffusionCoefficients:
        return self.exp_wrapper(EinsteinDiffusionCoefficients)(**self.kwargs)

    @property
    def EinsteinDistinctDiffusionCoefficients(
        self,
    ) -> EinsteinDistinctDiffusionCoefficients:
        return self.exp_wrapper(EinsteinDistinctDiffusionCoefficients)(**self.kwargs)

    @property
    def EinsteinHelfandIonicConductivity(self) -> EinsteinHelfandIonicConductivity:
        return self.exp_wrapper(EinsteinHelfandIonicConductivity)(**self.kwargs)

    @property
    def EinsteinHelfandThermalKinaci(self) -> EinsteinHelfandThermalKinaci:
        return self.exp_wrapper(EinsteinHelfandThermalKinaci)(**self.kwargs)

    @property
    def GreenKuboViscosityFlux(self) -> GreenKuboViscosityFlux:
        return self.exp_wrapper(GreenKuboViscosityFlux)(**self.kwargs)

    @property
    def GreenKuboDistinctDiffusionCoefficients(
        self,
    ) -> GreenKuboDistinctDiffusionCoefficients:
        return self.exp_wrapper(GreenKuboDistinctDiffusionCoefficients)(**self.kwargs)

    @property
    def GreenKuboIonicConductivity(self) -> GreenKuboIonicConductivity:
        return self.exp_wrapper(GreenKuboIonicConductivity)(**self.kwargs)

    @property
    def GreenKuboDiffusionCoefficients(self) -> GreenKuboDiffusionCoefficients:
        return self.exp_wrapper(GreenKuboDiffusionCoefficients)(**self.kwargs)

    @property
    def GreenKuboThermalConductivity(self) -> GreenKuboThermalConductivity:
        return self.exp_wrapper(GreenKuboThermalConductivity)(**self.kwargs)

    @property
    def GreenKuboViscosity(self) -> GreenKuboViscosity:
        return self.exp_wrapper(GreenKuboViscosity)(**self.kwargs)

    @property
    def KirkwoodBuffIntegral(self) -> KirkwoodBuffIntegral:
        return self.exp_wrapper(KirkwoodBuffIntegral)(**self.kwargs)

    @property
    def NernstEinsteinIonicConductivity(self) -> NernstEinsteinIonicConductivity:
        return self.exp_wrapper(NernstEinsteinIonicConductivity)(**self.kwargs)

    @property
    def PotentialOfMeanForce(self) -> PotentialOfMeanForce:
        return self.exp_wrapper(PotentialOfMeanForce)(**self.kwargs)

    @property
    def RadialDistributionFunction(self) -> RadialDistributionFunction:
        return self.exp_wrapper(RadialDistributionFunction)(**self.kwargs)

    # @property
    # def StructureFactor(self) -> StructureFactor:
    #     return self.exp_wrapper(StructureFactor)(**self.kwargs)

    @property
    def EinsteinHelfandThermalConductivity(self) -> EinsteinHelfandThermalConductivity:
        return self.exp_wrapper(EinsteinHelfandThermalConductivity)(**self.kwargs)

    # @property
    # def SpatialDistributionFunction(self) -> SpatialDistributionFunction:
    #     return self.exp_wrapper(SpatialDistributionFunction)(**self.kwargs)