FileHandler.py

"""FileHandler class provides convenience functions to handle collections of trajectory files generated by
AdaptiveSampling runs.
"""
import os
import mdtraj as md


class FileHandler:
    """FileHandler object for single agent runs.
    """
    def __init__(self, root, basename="", save_format=""):
        """Constructor for FileHandler object.

        :param root: str.
            Root directory.
        :param basename: str.
            Basename to append to trajectory files.
        :param save_format: str.
            Format used to save trajectories (only .dcd is used so far).
        """
        self.root = root if root else os.getcwd()  # Root directory
        self.basename = basename  # Basename for files
        self.subdirs = []  # List of subdirectories where trajectories will be stored
        self.files = {}  # List of files inside subdirectory
        self.n_frames = {}  # Number of frames in a trajectory
        self.init_states = {}  # The initial state file for each trajectory file
        self.save_format = save_format

    def generate_new_round(self, init_states, n_repeats, n_steps, n_round, save_rate, add_basename=""):
        """Determine paths and filenames for new trajectory files.

        :param init_states: dict.
            Information required to determine the initial state. This is used for record-keeping purposes only.
        :param n_repeats: int.
            Number of clones per trajectory.
        :param n_steps: int.
            Trajectory length in number of frames.
        :param n_round: int.
            Round number to generate.
        :param save_rate: int.
            Frame save rate for trajectory files. This is used for record-keeping purposes only.
        :param add_basename: str.
            Include additional string after self.basename.
        :return: list[str].
            Filenames for new round of trajectories.
        """
        path = os.path.join(self.root, "round_" + str(n_round))
        os.makedirs(path, exist_ok=True)

        self.subdirs.append(path)
        self.files[path] = []

        for i, state in enumerate(init_states):
            for j in range(n_repeats):
                fname = self.basename + \
                        add_basename + \
                        "_traj_{}".format(i) + \
                        "_repeat_{}".format(j) + \
                        self.save_format
                fname = os.path.join(path, fname)
                self.files[path].append(fname)
                self.n_frames[fname] = n_steps // save_rate
                self.init_states[fname] = state

        return self.files[path]

    def list_all_files(self):
        """List all trajectory files handled by this object.

        :return: list.
            List of trajectory filenames.
        """
        filenames = []
        for subdir in self.subdirs:
            for fname in self.files[subdir]:
                filenames.append(fname)

        return filenames

    def check_files(self, top_file=None):
        """Check that all files that are handled by this object exist and contain the correct number of frames.

        :param top_file: str.
            Top file to load trajectories. Specifying the top_file will trigger the check for the number of frames.
        :return: None.
        """
        all_files = self.list_all_files()

        # Check if file exists
        for fname in all_files:
            try:
                assert (os.path.exists(fname))
            except AssertionError as err:
                print("File {} should exist but was not found".format(fname))
                raise err

        # Check if number of frames is correct
        if top_file:
            for fname in all_files:
                traj = md.load(fname, top=top_file)
                try:
                    assert (traj.n_frames == self.n_frames[fname])
                except AssertionError as err:
                    print("File {} contains {} frames but {} were expected".format(fname,
                                                                                   traj.n_frames,
                                                                                   self.n_frames[fname]))
                    raise err

    def find_state_info(self, indices, traj_list, top_file):
        """Find the relevant information for each starting state.
        For each frame index, returns the dict {fname, frame_idx, top_file}, which allows to map the index to a frame
        in a trajectory file.

        :param indices: iterable(int).
            Frame indices.
        :param traj_list: list[str].
            List of trajectories to search.
        :param top_file: str.
            Path to the corresponding topology file for the trajectories.
        :return: list[dict].
            List with information from each state (fname, frame_idx, top_file).
        """
        all_files = traj_list
        states = []
        for idx in indices:
            counter = 0
            for fname in all_files:
                counter += self.n_frames[fname]
                if (counter - 1) >= idx:
                    frame_idx = idx - (counter - self.n_frames[fname])
                    state = {'fname': fname, 'frame_idx': frame_idx, 'top_file': top_file}
                    states.append(state)
                    break
        # Make sure all indices correspond to a trajectory frame
        assert (len(indices) == len(states))

        return states


class FileHandlerMultiagent(FileHandler):
    """FileHandler object for multiple agent runs.
    """
    def __init__(self, root, n_agents, basename="", save_format=""):
        """Constructor for FileHandlerMultiagent object.

        :param root: str.
            Root directory.
        :param n_agents: int.
            Number of agents.
        :param basename: str.
            Basename to append to trajectory files.
        :param save_format: str.
            Format used to save trajectories (only .dcd is used so far).
        """
        FileHandler.__init__(self, root, basename=basename, save_format=save_format)
        self.n_agents = n_agents
        self.agent_trajs = [[] for _ in range(n_agents)]  # Save the paths to the trajectories belonging to each agent

    def generate_new_round(self, init_states, n_repeats, n_steps, n_round, save_rate, add_basename=""):
        """Determine paths and filenames for new trajectory files.

        :param init_states: dict.
            Information required to determine the initial state. This is used for record-keeping purposes only.
        :param n_repeats: int.
            Number of clones per trajectory.
        :param n_steps: int.
            Trajectory length in number of frames.
        :param n_round: int.
            Round number to generate.
        :param save_rate: int.
            Frame save rate for trajectory files. This is used for record-keeping purposes only.
        :param add_basename: str.
            Include additional string after self.basename.
        :return: list[str].
            Filenames for new round of trajectories.
        """
        path = os.path.join(self.root, "round_" + str(n_round))
        os.makedirs(path, exist_ok=True)

        self.subdirs.append(path)
        self.files[path] = []

        for i, state in enumerate(init_states):
            agent_idx = state["agent_idx"]
            for j in range(n_repeats):
                fname = self.basename + \
                        add_basename + \
                        "_agent_{}".format(agent_idx) + \
                        "_traj_{}".format(i) + \
                        "_repeat_{}".format(j) + \
                        self.save_format
                fname = os.path.join(path, fname)
                self.files[path].append(fname)
                self.n_frames[fname] = n_steps // save_rate
                self.init_states[fname] = state
                self.agent_trajs[agent_idx].append(fname)

        return self.files[path]

    def list_agent_files(self, agent_idx):
        """List the trajectory files that belong to a specific agent.

        :param agent_idx: int.
            Agent index.
        :return: list[str].
            List of trajectory files for the corresponding agent.
        """
        return self.agent_trajs[agent_idx]

    def list_all_files(self):
        """List all trajectory files handled by this object.

        :return: list.
            List of trajectory filenames.
        """
        filenames = []
        for agent_idx in range(self.n_agents):
            for fname in self.agent_trajs[agent_idx]:
                filenames.append(fname)

        return filenames

    def find_state_info(self, indices, traj_list, top_file, executors):
        """Find the relevant information for each starting state.
        For each frame index, returns the dict {fname, frame_idx, top_file, agent_idx}, which allows to map the index to
        a frame in a trajectory file.

        :param indices: list[int].
            Frame indices.
        :param traj_list: list[str].
            List of trajectories to search.
        :param top_file: str.
            Path to the corresponding topology file for the trajectories.
        :param executors: list[int].
            List of agents that will execute each trajectory in the next round starting from the specified state.
        :return: list[dict].
            List with information from each state (fname, frame_idx, top_file, agent_idx).
        """
        all_files = traj_list
        states = []
        for i, idx in enumerate(indices):
            counter = 0
            executor = executors[i]
            for fname in all_files:
                counter += self.n_frames[fname]
                if (counter - 1) >= idx:
                    frame_idx = idx - (counter - self.n_frames[fname])
                    state = {'fname': fname, 'frame_idx': frame_idx, 'top_file': top_file, 'agent_idx': executor}
                    states.append(state)
                    break
        # Make sure all indices correspond to a trajectory frame
        assert (len(indices) == len(states))

        return states