NexQuant/rdagent/core/experiment.py

from abc import ABC, abstractmethod
from pathlib import Path
from typing import Generic, Optional, Sequence, TypeVar

"""
This file contains the all the class about organizing the task in RD-Agent.
"""


class Task:
    # TODO: 把name放在这里作为主键
    # Please refer to rdagent/model_implementation/task.py for the implementation
    # I think the task version applies to the base class.
    pass


ASpecificTask = TypeVar("ASpecificTask", bound=Task)


class Implementation(ABC, Generic[ASpecificTask]):
    def __init__(self, target_task: ASpecificTask) -> None:
        self.target_task = target_task

    @abstractmethod
    def execute(self, data=None, config: dict = {}) -> object:
        """
        The execution of the implementation can be dynamic.

        So we may pass in the data and config dynamically.
        """
        raise NotImplementedError("execute method is not implemented.")

    @abstractmethod
    def execute_desc(self):
        """
        return the description how we will execute the code in the folder.
        """
        raise NotImplementedError(f"This type of input is not supported")

    # TODO:
    # After execution, it should return some results.
    # Some evaluators will input the results and output


ASpecificImp = TypeVar("ASpecificImp", bound=Implementation)


class ImpLoader(ABC, Generic[ASpecificTask, ASpecificImp]):
    @abstractmethod
    def load(self, task: ASpecificTask) -> ASpecificImp:
        raise NotImplementedError("load method is not implemented.")


class FBImplementation(Implementation):
    """
    File-based task implementation

    The implemented task will be a folder which contains related elements.
    - Data
    - Code Implementation
    - Output
        - After execution, it will generate the final output as file.

    A typical way to run the pipeline of FBImplementation will be
    (We didn't add it as a method due to that we may pass arguments into `prepare` or `execute` based on our requirements.)

    .. code-block:: python

        def run_pipeline(self, **files: str):
            self.prepare()
            self.inject_code(**files)
            self.execute()

    """

    # TODO:
    # FileBasedFactorImplementation should inherit from it.
    # Why not directly reuse FileBasedFactorImplementation.
    #   Because it has too much concrete dependencies.
    #   e.g.  dataframe, factors

    path: Optional[Path]

    @abstractmethod
    def prepare(self, *args, **kwargs):
        """
        Prepare all the files except the injected code
        - Data
        - Documentation
        - TODO: env?  Env is implicitly defined by the document?

            typical usage of `*args, **kwargs`:
                Different methods shares the same data. The data are passed by the arguments.
        """

    def inject_code(self, **files: str):
        """
        Inject the code into the folder.
        {
            "model.py": "<model code>"
        }
        """
        for k, v in files.items():
            with open(self.path / k, "w") as f:
                f.write(v)

    def get_files(self) -> list[Path]:
        """
        Get the environment description.

        To be general, we only return a list of filenames.
        How to summarize the environment is the responsibility of the TaskGenerator.
        """
        return list(self.path.iterdir())


class Experiment(ABC, Generic[ASpecificTask, ASpecificImp]):
    """
    The experiment is a sequence of tasks and the implementations of the tasks after generated by the TaskGenerator.
    """

    def __init__(self, sub_tasks: Sequence[ASpecificTask]) -> None:
        self.sub_tasks = sub_tasks
        self.sub_implementations: Sequence[ASpecificImp] = [None for _ in self.sub_tasks]
        self.based_experiments: Sequence[Experiment] = []
        self.result: object = None  # The result of the experiment, can be different types in different scenarios.


TaskOrExperiment = TypeVar("TaskOrExperiment", Task, Experiment)


class Loader(ABC, Generic[TaskOrExperiment]):
    @abstractmethod
    def load(self, *args, **kwargs) -> TaskOrExperiment:
        raise NotImplementedError("load method is not implemented.")