Commit c3fba6d2 authored by tomsko's avatar tomsko
Browse files

save all provided and loaded files

parent ebc2e63e

src/executor/task_creator.py

+13 −0
Original line number Diff line number Diff line
@@ -78,6 +78,17 @@ def unpack_archives_in_directory(dir_path: str, verbose: bool) -> None: 
            zip_ref.extractall(dir_path)





def save_files_in_directory(path: str) -> None:

    """

    Saves list of all files in directory into the file "all_files.txt"

    :param path: to the directory

    :return: None (creates the file)

    """

    files = [f for f in listdir(path) if isfile(join(path, f))]

    with open(join(path, "all_files.txt"), "w") as f:

        json.dump(files, f)





def create_tasks_for_simulations(simulation_folders: list[str],

                                 task_types: set[TaskType],

                                 verbose: bool = True) -> list[Task]:
@@ -98,6 +109,8 @@ def create_tasks_for_simulations(simulation_folders: list[str], 
        # begin with extracting any zip archives

        unpack_archives_in_directory(folder, verbose)



        save_files_in_directory(folder)



        tasks: list[Task] = []

        def mk_task_list() -> None:

            """

src/loader/universe_loader.py

+33 −10
Original line number Diff line number Diff line 
import logging

import json

import os

from enum import Enum

from .results_file import ResultsFile

from MDAnalysis import Universe # type: ignore
@@ -9,6 +11,17 @@ LOAD_STRATEGY_PATH = "simulation/load_strategy" 


logger = logging.getLogger("base")



def save_used_files(dir_path: str, files: list[str]) -> None:

    """

    Saves used files used for loading the simulation

    :param dir_path: to save used files to (used_files.txt)

    :param files: used to load the simulation

    :return: None (creates the file)

    """



    with open(os.path.join(dir_path, "used_files.txt"), "w") as fd:

        json.dump(files, fd)



class LoadStrategy(Enum):

    ITP_GRO = "ITP_GRO"  # ITP-based topology and gro coords

    TOP_GRO = "TOP_GRO"  # TOP-based topology and gro coords
@@ -57,18 +70,22 @@ def _load_using(simulation_directory: str, results_file: ResultsFile, strategy: 


    match strategy:

        case LoadStrategy.ITP_GRO:

            return _load_itp_gro(simulation_files, results_file)

            return _load_itp_gro(simulation_files, results_file, simulation_directory)

        case LoadStrategy.TOP_GRO:

            return _load_top_gro(simulation_files, results_file)

            return _load_top_gro(simulation_files, results_file, simulation_directory)

        case LoadStrategy.TPR_GRO:

            return _load_tpr_gro(simulation_files, results_file)

            return _load_tpr_gro(simulation_files, results_file, simulation_directory)

        case LoadStrategy.TPR:

            return _load_tpr(simulation_files, results_file)

            return _load_tpr(simulation_files, results_file, simulation_directory)

        case LoadStrategy.TOP:

            return _load_top(simulation_files, results_file, simulation_directory)

        case LoadStrategy.ITP:

            return _load_itp(simulation_files, results_file, simulation_directory)

        case _:

            raise RuntimeError("Unknown load strategy")





def _load_itp_gro(simulation_files: list[str], results_file: ResultsFile) -> Universe:

def _load_itp_gro(simulation_files: list[str], results_file: ResultsFile, simulation_directory: str) -> Universe:

    """

    Attempts to load simulation into MDAnalysis' universe using ITP topology and GRO coordinates

    :param simulation_files: files representing the simulation
@@ -92,10 +109,11 @@ def _load_itp_gro(simulation_files: list[str], results_file: ResultsFile) -> Uni 
                                  infer_system=True)



    results_file.set_item(LOAD_STRATEGY_PATH, LoadStrategy.ITP_GRO.value)

    save_used_files(simulation_directory, [top_files[0], gro_files[0]])

    return universe





def _load_top_gro(simulation_files: list[str], results_file: ResultsFile) -> Universe:

def _load_top_gro(simulation_files: list[str], results_file: ResultsFile, simulation_directory: str) -> Universe:

    """

    Attempts to load simulation into MDAnalysis' universe using TOP topology and GRO coordinates

    :param simulation_files: files representing the simulation
@@ -112,10 +130,11 @@ def _load_top_gro(simulation_files: list[str], results_file: ResultsFile) -> Uni 
        raise RuntimeError(ERROR_STRATEGY_NOT_APPLICABLE)

    universe: Universe = Universe(top_files[0],gro_files[0], tpr_resid_from_one=True, in_memory=True)

    results_file.set_item(LOAD_STRATEGY_PATH, LoadStrategy.TOP_GRO.value)

    save_used_files(simulation_directory, [top_files[0], gro_files[0]])

    return universe





def _load_tpr_gro(simulation_files: list[str], results_file: ResultsFile) -> Universe:

def _load_tpr_gro(simulation_files: list[str], results_file: ResultsFile, simulation_directory: str) -> Universe:

    """

    Attempts to load simulation into MDAnalysis' universe using TPR topology and GRO coordinates

    :param simulation_files: files representing the simulation
@@ -133,10 +152,11 @@ def _load_tpr_gro(simulation_files: list[str], results_file: ResultsFile) -> Uni 


    universe: Universe = Universe(top_files[0],gro_files[0], tpr_resid_from_one=True, in_memory=True)

    results_file.set_item(LOAD_STRATEGY_PATH, LoadStrategy.TPR_GRO.value)

    save_used_files(simulation_directory, [top_files[0], gro_files[0]])

    return universe





def _load_tpr(simulation_files: list[str], results_file: ResultsFile) -> Universe:

def _load_tpr(simulation_files: list[str], results_file: ResultsFile, simulation_directory: str) -> Universe:

    """

    Attempts to load simulation into MDAnalysis' universe using TPR topology w/o coordinates

    :param simulation_files: files representing the simulation
@@ -153,10 +173,11 @@ def _load_tpr(simulation_files: list[str], results_file: ResultsFile) -> Univers 


    universe: Universe = Universe(top_files[0], tpr_resid_from_one=True, in_memory=True)

    results_file.set_item(LOAD_STRATEGY_PATH, LoadStrategy.TPR.value)

    save_used_files(simulation_directory, [top_files[0]])

    return universe





def _load_top(simulation_files: list[str], results_file: ResultsFile) -> Universe:

def _load_top(simulation_files: list[str], results_file: ResultsFile, simulation_directory: str) -> Universe:

    """

    Attempts to load simulation into MDAnalysis' universe using TOP topology w/o coordinates

    :param simulation_files: files representing the simulation
@@ -173,10 +194,11 @@ def _load_top(simulation_files: list[str], results_file: ResultsFile) -> Univers 


    universe: Universe = Universe(top_files[0], tpr_resid_from_one=True, in_memory=True)

    results_file.set_item(LOAD_STRATEGY_PATH, LoadStrategy.TOP.value)

    save_used_files(simulation_directory, [top_files[0]])

    return universe





def _load_itp(simulation_files: list[str], results_file: ResultsFile) -> Universe:

def _load_itp(simulation_files: list[str], results_file: ResultsFile, simulation_directory: str) -> Universe:

    """

    Attempts to load simulation into MDAnalysis' universe using TOP topology w/o coordinates

    :param simulation_files: files representing the simulation
@@ -198,4 +220,5 @@ def _load_itp(simulation_files: list[str], results_file: ResultsFile) -> Univers 
                                  infer_system=True)



    results_file.set_item(LOAD_STRATEGY_PATH, LoadStrategy.ITP.value)

    save_used_files(simulation_directory, [top_files[0]])

    return universe
 
 No newline at end of file