<commit_before># -*- coding: utf-8 -*-
'''
Molecule Data
===============================
'''
import numpy as np
import pandas as pd
from networkx import Graph
from networkx.algorithms.components import connected_components
from itertools import combinations
from exa import DataFrame
from atomic import Isotope
from atomic.formula import string_to_dict


class Molecule(DataFrame):
    '''
    Description of molecules in the atomic universe.
    '''
    _index = ['molecule']
    _groupbys = ['frame']
    _categories = {'frame': np.int64, 'formula': str, 'classification': object}

    def classify(self, *classifiers, overwrite=False):
        '''
        Classify molecules into arbitrary categories.

        .. code-block:: Python

            u.molecule.classify(('Na', 'solute'), ('H(2)O(1)', 'solvent'))

        Args:
            classifiers: Any number of tuples of the form ("identifier", "label", exact) (see below)
            overwrite (bool): If true, overwrite existing ClassificationError

        Note:
            A classifier has 3 parts, "identifier", e.g. "H(2)O(1)", "label", e.g.
            "solvent", and exact (true or false). If exact is false (default),
            classification is greedy and (in this example) molecules with formulas
            "H(1)O(1)", "H(3)O(1)", etc. would get classified as "solvent". If,
            instead, exact were set to true, those molecules would remain
            unclassified.

        See Also:
            :func:`~atomic.algorithms.nearest.compute_nearest_molecules`
        '''
        for c in classifiers:
            n = len(c)
            if n != 3 and n != 2:
                raise ClassificationError()
        if 'classification' not in self:
            self['classification'] = None
        for classifier in classifiers:
            identifier = string_to_dict(classifier[0])
            classification = classifier[1]
            exact = classifier[2] if len(classifier) == 3 else False
            this = self if overwrite else self[self['classification'].isnull()]
            for symbol, count in identifier.items():
                this = this[this[symbol] == count] if exact else this[this[symbol] >= 1]
            if len(this) > 0:
                self.ix[self.index.isin(this.index), 'classification'] = classification
            else:
                raise KeyError('No records found for {}, with identifier {}.'.format(classification, identifier))
        self['classification'] = self['classification'].astype('category')


def compute_molecule(universe):
    '''
    Cluster atoms into molecules.

    The algorithm is to create a network graph containing every atom (in every
    frame as nodes and bonds as edges). Using this connectivity information,
    one can perform a (breadth first) traversal of the network graph to cluster
    all nodes (whose indices correspond to physical atoms).

    Args:
        universe (:class:`~atomic.universe.Universe`): Atomic universe<commit_msg>Added molecular atom count function<commit_after># -*- coding: utf-8 -*-
'''
Molecule Data
===============================
'''
import numpy as np
import pandas as pd
from networkx import Graph
from networkx.algorithms.components import connected_components
from itertools import combinations
from exa import DataFrame
from atomic import Isotope
from atomic.formula import string_to_dict


class Molecule(DataFrame):
    '''
    Description of molecules in the atomic universe.
    '''
    _index = ['molecule']
    _groupbys = ['frame']
    _categories = {'frame': np.int64, 'formula': str, 'classification': object}

    def classify(self, *classifiers, overwrite=False):
        '''
        Classify molecules into arbitrary categories.

        .. code-block:: Python

            u.molecule.classify(('Na', 'solute'), ('H(2)O(1)', 'solvent'))

        Args:
            classifiers: Any number of tuples of the form ("identifier", "label", exact) (see below)
            overwrite (bool): If true, overwrite existing ClassificationError

        Note:
            A classifier has 3 parts, "identifier", e.g. "H(2)O(1)", "label", e.g.
            "solvent", and exact (true or false). If exact is false (default),
            classification is greedy and (in this example) molecules with formulas
            "H(1)O(1)", "H(3)O(1)", etc. would get classified as "solvent". If,
            instead, exact were set to true, those molecules would remain
            unclassified.

        See Also:
            :func:`~atomic.algorithms.nearest.compute_nearest_molecules`
        '''
        for c in classifiers:
            n = len(c)
            if n != 3 and n != 2:
                raise ClassificationError()
        if 'classification' not in self:
            self['classification'] = None
        for classifier in classifiers:
            identifier = string_to_dict(classifier[0])
            classification = classifier[1]
            exact = classifier[2] if len(classifier) == 3 else False
            this = self if overwrite else self[self['classification'].isnull()]
            for symbol, count in identifier.items():
                this = this[this[symbol] == count] if exact else this[this[symbol] >= 1]
            if len(this) > 0:
                self.ix[self.index.isin(this.index), 'classification'] = classification
            else:
                raise KeyError('No records found for {}, with identifier {}.'.format(classification, identifier))
        self['classification'] = self['classification'].astype('category')

    def compute_atom_count(self):
        '''
        Compute the molecular atom count.
        '''
        symbols = [col for col in self.columns if len(col) < 3 and col[0].istitle()]
        self['atom_count'] = self[symbols].sum(axis=1)


def compute_molecule(universe):
    '''
    Cluster atoms into molecules.

    The algorithm is to create a network graph containing every atom (in every
    frame as nodes and bonds as edges). Using this connectivity information,
    one can perform a (breadth first) traversal of the network graph to cluster
    all nodes (whose indices correspond to physical atoms).

    Args:
        universe (:class:`~atomic.universe.Universe`): Atomic universe
{'fraction_non_alphanumeric': 0.11474883646284706, 'fraction_numerical': 0.006900978976087305, 'mean_word_length': 2.6465769455822117, 'pattern_counts': {'":': 0, '<': 4, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '762437', 'n_tokens_mistral': 1692, 'n_tokens_neox': 1563, 'n_words': 577}