<commit_before>#!/usr/bin/python
import numpy as np
import pandas as pd
from mdtraj import Trajectory
from mdtraj import Topology
from mdtraj import element
from mdtraj.geometry import distance

mapping_options = {}

def map_forces(traj,atom_indices=None,use_pbc=True):
    """Compute the center of mass for each frame.
    Parameters
    ----------
    traj : Trajectory
        Trajectory to sum forces on
    atom_indices : array-like, dtype=int, shape=(n_atoms)
            List of indices of atoms to use in computing com
    Returns
    -------
    forces : np.ndarray, shape=(n_frames, 3)
         Summed up forces on atom indices
    """

    if atom_indices is not None and len(atom_indices)>0:
        forces = traj.forces[:,atom_indices,:]
    else:
        forces = traj.forces

    mapped_forces = forces.sum(axis=1)

    return mapped_forces

def map_molecules(trj,selection_list,bead_label_list,molecule_types=None,*args,**kwargs):
    """ This performs the mapping where each molecule has been assigned a type"""
    index_list = []
    resSeq_list = []
    label_list = []

    n_molecule_types = len(selection_list)
    if molecule_types is None:
        molecule_types = [0 for idx in range(trj.top.n_residues)]

    if not sorted(set(molecule_types))==range(n_molecule_types):
        raise ValueError("Error in map molecules, molecule types list must contain only and all numbers from 0 to n_molecule_types-1")
    if len(molecule_types) != trj.top.n_residues:
        raise ValueError("Error in map molecules, molecule types list must have the same length as number of residues")

    # get the first molecule with molecule type i
    first_molecules = [molecule_types.index(i) for i in range(n_molecule_types)]

    if len(selection_list) != len(bead_label_list):
        raise ValueError("Error in map molecules, must submit selection list and bead label list of same length")
    for i in range(n_molecule_types):
        if len(selection_list[i]) != len(bead_label_list[i]):
            raise ValueError("Error in map molecules, selection list %i and bead label list %i must be of same length"%(i,i))

    first_indices = []
    internal_indices_list = [] 
    for i in range(n_molecule_types):
        internal_indices_list.append([])
        first_index = trj.top.select("(resid == %i)"%(first_molecules[i])).min()
        for sel in selection_list[i]:
            if sel.find("index")>-1 and sel.find("name")>-1:
                raise ValueError("Error in map molecules, do not specify selection by index and by type")
            elif sel.find("index")>-1:
                # use atom selection language to parse selection string containing only indices on whole system, then offset later
                internal_indices = trj.top.select("%s"%(sel))
            elif sel.find("name")>-1:
                # have to un-shift list because this will be added to current id later
                internal_indices = trj.top.select("(resid == %i) and (%s)"%(first_molecules[i],sel)) - first_index

            if len(internal_indices)==0:
                raise ValueError("Error in map_molecules, selection string '%s' produced an empty list of atom indices"%sel)
            internal_indices_list[i].append(internal_indices)

    start_index = 0
    resSeq = 1
    for r in trj.top.residues:
        molecule_type = molecule_types[resSeq-1]
        for bead_idx, internal_indices in enumerate(internal_indices_list[molecule_type]):
            system_indices = internal_indices + start_index
            index_list.append(system_indices) 
            resSeq_list.append(resSeq)
            label_list.append(bead_label_list[molecule_type][bead_idx])
        resSeq = resSeq+1<commit_msg>added option to map_molecules so that it can sort topology into blocks of different atom types<commit_after>#!/usr/bin/python
import numpy as np
import pandas as pd
from mdtraj import Trajectory
from mdtraj import Topology
from mdtraj import element
from mdtraj.geometry import distance

mapping_options = {}

def map_forces(traj,atom_indices=None,use_pbc=True):
    """Compute the center of mass for each frame.
    Parameters
    ----------
    traj : Trajectory
        Trajectory to sum forces on
    atom_indices : array-like, dtype=int, shape=(n_atoms)
            List of indices of atoms to use in computing com
    Returns
    -------
    forces : np.ndarray, shape=(n_frames, 3)
         Summed up forces on atom indices
    """

    if atom_indices is not None and len(atom_indices)>0:
        forces = traj.forces[:,atom_indices,:]
    else:
        forces = traj.forces

    mapped_forces = forces.sum(axis=1)

    return mapped_forces

def map_molecules(trj,selection_list,bead_label_list,molecule_types=None,molecule_type_order=False,*args,**kwargs):
    """ This performs the mapping where each molecule has been assigned a type. 

        Specifying molecule_type_order means that the map will be reordered so that all molecules of type 0 come first, then 1, etc
    """
    index_list = []
    resSeq_list = []
    label_list = []

    n_molecule_types = len(selection_list)
    if molecule_types is None:
        molecule_types = [0 for idx in range(trj.top.n_residues)]

    if not sorted(set(molecule_types))==range(n_molecule_types):
        raise ValueError("Error in map molecules, molecule types list must contain only and all numbers from 0 to n_molecule_types-1")
    if len(molecule_types) != trj.top.n_residues:
        raise ValueError("Error in map molecules, molecule types list must have the same length as number of residues")

    # get the first molecule with molecule type i
    first_molecules = [molecule_types.index(i) for i in range(n_molecule_types)]

    if len(selection_list) != len(bead_label_list):
        raise ValueError("Error in map molecules, must submit selection list and bead label list of same length")
    for i in range(n_molecule_types):
        if len(selection_list[i]) != len(bead_label_list[i]):
            raise ValueError("Error in map molecules, selection list %i and bead label list %i must be of same length"%(i,i))

    first_indices = []
    internal_indices_list = [] 
    for i in range(n_molecule_types):
        internal_indices_list.append([])
        first_index = trj.top.select("(resid == %i)"%(first_molecules[i])).min()
        for sel in selection_list[i]:
            if sel.find("index")>-1 and sel.find("name")>-1:
                raise ValueError("Error in map molecules, do not specify selection by index and by type")
            elif sel.find("index")>-1:
                # use atom selection language to parse selection string containing only indices on whole system, then offset later
                internal_indices = trj.top.select("%s"%(sel))
            elif sel.find("name")>-1:
                # have to un-shift list because this will be added to current id later
                internal_indices = trj.top.select("(resid == %i) and (%s)"%(first_molecules[i],sel)) - first_index

            if len(internal_indices)==0:
                raise ValueError("Error in map_molecules, selection string '%s' produced an empty list of atom indices"%sel)
            internal_indices_list[i].append(internal_indices)

    start_index = 0

    # get list of type [ (0,r0), (1,r1) etc ]
    if molecule_type_order is True:
        residue_list = sorted( enumerate(trj.top.residues),key=lambda x: molecule_types[x[0]])
    else:
        residue_list = enumerate(trj.top.residues)

    resSeq = 1
    for ridx,r in residue_list:
        molecule_type = molecule_types[ridx]
        for bead_idx, internal_indices in enumerate(internal_indices_list[molecule_type]):
            system_indices = internal_indices + start_index
            index_list.append(system_indices) 
            resSeq_list.append(resSeq)
            label_list.append(bead_label_list[molecule_type][bead_idx])
        resSeq = resSeq+1
{'fraction_non_alphanumeric': 0.09309385277004807, 'fraction_numerical': 0.004553503668100177, 'mean_word_length': 3.079979360165119, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 13, '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': '2024565', 'n_tokens_mistral': 2426, 'n_tokens_neox': 2262, 'n_words': 708}