<commit_before>#! /usr/bin/env python

########################################################################
#                                                                      #
# This script was written by <NAME> in 2017.                    #
#        <EMAIL>     <EMAIL>                    #
#                                                                      #
# Copyright 2017 <NAME>                                   #
#                                                                      #
# Licensed under the Apache License, Version 2.0 (the "License");      #
# you may not use this file except in compliance with the License.     #
# You may obtain a copy of the License at                              #
#                                                                      #
#    http://www.apache.org/licenses/LICENSE-2.0                        #
#                                                                      #
# Unless required by applicable law or agreed to in writing, software  #
# distributed under the License is distributed on an "AS IS" BASIS,    #
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or      #
# implied.                                                             #
# See the License for the specific language governing permissions and  #
# limitations under the License.                                       #
#                                                                      #
########################################################################

from __future__ import absolute_import, division, print_function
from six.moves import range

import os.path
from warnings import warn

import MDAnalysis as MDa
import matplotlib as mpl
import matplotlib.pyplot as plt
# import mdtraj as md  # Think I'm going with MDAnalysis instead
import numpy as np
import pandas as pd
from typing import Iterable, Tuple

from . import exceptions
from .exceptions import InputError


# TODO move all import statements to the beginning (out of functions)


class Universe(MDa.Universe):

    def __init__(self, *args, **kwargs):
        """

        :type verbosity: int or bool
        :param verbosity: Default: 1. Setting whether to print details. If in
            the future more levels of verbosity are desired, this may be
            changed to an int.

        :param float temp: Default: None. Temperature of this simulation to be
            used for calculating free energy surfaces with weighted histograms.
        :param args:
        :param kwargs:
        """
        # self.univ = (line below): I'm not sure if this is needed or if this
        # just automatically inherits everything
        # Maybe use the super() command? need to learn more about this
        self._verbosity = kwargs.pop('verbosity', 1)
        self._temperature = kwargs.pop('temp', None)
        super(Universe, self).__init__(*args, **kwargs)
        self._num_frames = self.trajectory.n_frames
        self._last_time = self.trajectory.totaltime
        self._data = self._init_dataframe()
        # dict of distance definitions
        self._dict_dist_defs = {}
        self._dict_dihed_defs = {}

    def _init_dataframe(self):
        """
        Initialize a pandas.DataFrame with Times column.

        This uses self._last_time as the final time and self._num_frames as
        the number of rows to put into the DataFrame to be returned.
        This uses np.linspace to make the evenly spaced times that should
        match with the times in the trajectory file.

        :return: a DataFrame with one column of Times
        :rtype: pd.DataFrame
        """
        return pd.DataFrame(np.linspace(0, self._last_time,
                                        num=self._num_frames),
                            columns=['Time'])

    def save_data(self, filename=None, overwrite=False):
        """
        Save calculated data to disk

        :param str filename: Filename to save the data as. Defaults to the
            name of the trajectory with a '.h5' extension.

        :param bool overwrite: Whether to overwrite existing data on disk.
            If it's True, it will completely overwrite the existing data store.
            If it's False, but a store for this time already exists, only new
            columns in self.data will be added to the store, and no data will be
            overwritten.

        :return: None
        """
        if filename is None:
            filename = os.path.splitext(self.trajectory.filename)[0] + '.h5'
        with pd.HDFStore(filename) as store:
            time = 'time_' + str(int(self._last_time/1000)) + 'ns'
            # TODO speed this up for no changes
            # can quickly(?) get the columns with
            # store.get_node("df").axis0.read() to see if I have any new
            # columns to add in the first place.
            if overwrite or ('/'+time not in store.keys()):
                store[time] = self._data
            else:
                store_cols = store.get_node(time).axis0.read()
                set_diff_cols = set(self._data.columns).difference(store_cols)
                if not set_diff_cols:
                    if self._verbosity:
                        print('No data added to {}[{}]'.format(filename, time))
                    return
                store_df = store[time]  # seems this has to be done to add cols
                # see https://stackoverflow.com/questions/15939603/
                # append-new-columns-to-hdfstore-with-pandas
                for col in set_diff_cols:
                    store_df[col] = self._data[col]
                store[time] = store_df
        if self._verbosity:
            print('Saved data to {}[{}]'.format(filename, time))

    def read_data(self, filename=None, ignore_no_data=False):
        """
        Read calculated data from disk

        This will read the data from disk and add it to self.data. Any
        existing data will not be overwritten.

        :param str filename: Filename from which to read the data.
            Defaults to the name of the trajectory with a '.h5' extension.
        :param bool ignore_no_data: Default: False. If True, not having data
            in the file will not raise an error.
        :return: None
        :raises: IOError
        """
        if filename is None:
            filename = os.path.splitext(self.trajectory.filename)[0] + '.h5'
        with pd.HDFStore(filename) as store:
            time = 'time_' + str(int(self._last_time/1000)) + 'ns'
            try:
                read_df = store[time]
                keys_to_read = set(read_df.columns).difference(
                    self._data.columns)
            except KeyError:
                keys_to_read = []
                read_df = pd.DataFrame()  # not necessary; stops IDE complaint
                if not ignore_no_data:
                    raise IOError('This data does not exist!\n{}[{}]'.format(
                        filename, time))
                else:
                    if self._verbosity:
                        print('No data to read in '
                              '{}[{}]'.format(filename, time))
        for key in keys_to_read:
            self._data[key] = read_df[key]

    def calculate_distances(self, recalculate=False, ignore_file_change=False,
                            *args, **kwargs):
        """
        Calculate distances by iterating through the trajectory

        :param recalculate: Default: False. If True, all requested
            distances will be calculated.
            If False, the intersection of the set of requested distance names
            and the existing column titles in self.data will be removed from the
            information to be calculated.<commit_msg>remove completed TODO; added two to use final_time_str internally<commit_after>#! /usr/bin/env python

########################################################################
#                                                                      #
# This script was written by <NAME> in 2017.                    #
#        <EMAIL>     <EMAIL>                    #
#                                                                      #
# Copyright 2017 <NAME>                                   #
#                                                                      #
# Licensed under the Apache License, Version 2.0 (the "License");      #
# you may not use this file except in compliance with the License.     #
# You may obtain a copy of the License at                              #
#                                                                      #
#    http://www.apache.org/licenses/LICENSE-2.0                        #
#                                                                      #
# Unless required by applicable law or agreed to in writing, software  #
# distributed under the License is distributed on an "AS IS" BASIS,    #
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or      #
# implied.                                                             #
# See the License for the specific language governing permissions and  #
# limitations under the License.                                       #
#                                                                      #
########################################################################

from __future__ import absolute_import, division, print_function
from six.moves import range

import os.path
from warnings import warn

import MDAnalysis as MDa
import matplotlib as mpl
import matplotlib.pyplot as plt
# import mdtraj as md  # Think I'm going with MDAnalysis instead
import numpy as np
import pandas as pd
from typing import Iterable, Tuple

from . import exceptions
from .exceptions import InputError


# TODO move all import statements to the beginning (out of functions)


class Universe(MDa.Universe):

    def __init__(self, *args, **kwargs):
        """

        :type verbosity: int or bool
        :param verbosity: Default: 1. Setting whether to print details. If in
            the future more levels of verbosity are desired, this may be
            changed to an int.

        :param float temp: Default: None. Temperature of this simulation to be
            used for calculating free energy surfaces with weighted histograms.
        :param args:
        :param kwargs:
        """
        # self.univ = (line below): I'm not sure if this is needed or if this
        # just automatically inherits everything
        # Maybe use the super() command? need to learn more about this
        self._verbosity = kwargs.pop('verbosity', 1)
        self._temperature = kwargs.pop('temp', None)
        super(Universe, self).__init__(*args, **kwargs)
        self._num_frames = self.trajectory.n_frames
        self._last_time = self.trajectory.totaltime
        self._data = self._init_dataframe()
        # dict of distance definitions
        self._dict_dist_defs = {}
        self._dict_dihed_defs = {}

    def _init_dataframe(self):
        """
        Initialize a pandas.DataFrame with Times column.

        This uses self._last_time as the final time and self._num_frames as
        the number of rows to put into the DataFrame to be returned.
        This uses np.linspace to make the evenly spaced times that should
        match with the times in the trajectory file.

        :return: a DataFrame with one column of Times
        :rtype: pd.DataFrame
        """
        return pd.DataFrame(np.linspace(0, self._last_time,
                                        num=self._num_frames),
                            columns=['Time'])

    def save_data(self, filename=None, overwrite=False):
        """
        Save calculated data to disk

        :param str filename: Filename to save the data as. Defaults to the
            name of the trajectory with a '.h5' extension.

        :param bool overwrite: Whether to overwrite existing data on disk.
            If it's True, it will completely overwrite the existing data store.
            If it's False, but a store for this time already exists, only new
            columns in self.data will be added to the store, and no data will be
            overwritten.

        :return: None
        """
        if filename is None:
            filename = os.path.splitext(self.trajectory.filename)[0] + '.h5'
        with pd.HDFStore(filename) as store:
            time = 'time_' + str(int(self._last_time/1000)) + 'ns'
            # TODO use self.final_time_str
            if overwrite or ('/'+time not in store.keys()):
                store[time] = self._data
            else:
                store_cols = store.get_node(time).axis0.read()
                set_diff_cols = set(self._data.columns).difference(store_cols)
                if not set_diff_cols:
                    if self._verbosity:
                        print('No data added to {}[{}]'.format(filename, time))
                    return
                store_df = store[time]  # seems this has to be done to add cols
                # see https://stackoverflow.com/questions/15939603/
                # append-new-columns-to-hdfstore-with-pandas
                for col in set_diff_cols:
                    store_df[col] = self._data[col]
                store[time] = store_df
        if self._verbosity:
            print('Saved data to {}[{}]'.format(filename, time))

    def read_data(self, filename=None, ignore_no_data=False):
        """
        Read calculated data from disk

        This will read the data from disk and add it to self.data. Any
        existing data will not be overwritten.

        :param str filename: Filename from which to read the data.
            Defaults to the name of the trajectory with a '.h5' extension.
        :param bool ignore_no_data: Default: False. If True, not having data
            in the file will not raise an error.
        :return: None
        :raises: IOError
        """
        if filename is None:
            filename = os.path.splitext(self.trajectory.filename)[0] + '.h5'
        with pd.HDFStore(filename) as store:
            time = 'time_' + str(int(self._last_time/1000)) + 'ns'
            # TODO use self.final_time_str
            try:
                read_df = store[time]
                keys_to_read = set(read_df.columns).difference(
                    self._data.columns)
            except KeyError:
                keys_to_read = []
                read_df = pd.DataFrame()  # not necessary; stops IDE complaint
                if not ignore_no_data:
                    raise IOError('This data does not exist!\n{}[{}]'.format(
                        filename, time))
                else:
                    if self._verbosity:
                        print('No data to read in '
                              '{}[{}]'.format(filename, time))
        for key in keys_to_read:
            self._data[key] = read_df[key]

    def calculate_distances(self, recalculate=False, ignore_file_change=False,
                            *args, **kwargs):
        """
        Calculate distances by iterating through the trajectory

        :param recalculate: Default: False. If True, all requested
            distances will be calculated.
            If False, the intersection of the set of requested distance names
            and the existing column titles in self.data will be removed from the
            information to be calculated.
{'fraction_non_alphanumeric': 0.09004800830413909, 'fraction_numerical': 0.004995458673932788, 'mean_word_length': 1.624723310063, 'pattern_counts': {'":': 0, '<': 11, '<?xml version=': 0, '>': 11, 'https://': 2, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 34, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1984202', 'n_tokens_mistral': 3803, 'n_tokens_neox': 3603, 'n_words': 1377}