parameters: (
          coord : np.ndarray, atom_type : np.ndarray, sel_atoms : List[int] = None
      )
docstring: Sort atoms in the system according their types.
          
          Parameters
          ----------
          coord
                  The coordinates of atoms.
                  Should be of shape [nframes, natoms, 3]
          atom_type
                  The type of atoms
                  Should be of shape [natoms]
          sel_atom
                  The selected atoms by type
          
          Returns
          -------
          coord_out
                  The coordinates after sorting
          atom_type_out
                  The atom types after sorting
          idx_map
                  The index mapping from the input to the output. 
                  For example coord_out = coord[:,idx_map,:]
          sel_atom_type
                  Only output if sel_atoms is not None
                  The sorted selected atom types
          sel_idx_map
                  Only output if sel_atoms is not None
                  The index mapping from the selected atoms to sorted selected atoms.
docstring_summary: Sort atoms in the system according their types.
          
          Parameters
          ----------
          coord
                  The coordinates of atoms.
                  Should be of shape [nframes, natoms, 3]
          atom_type
                  The type of atoms
                  Should be of shape [natoms]
          sel_atom
                  The selected atoms by type
          
          Returns
          -------
          coord_out
                  The coordinates after sorting
          atom_type_out
                  The atom types after sorting
          idx_map
                  The index mapping from the input to the output. 
                  For example coord_out = coord[:,idx_map,:]
          sel_atom_type
                  Only output if sel_atoms is not None
                  The sorted selected atom types
          sel_idx_map
                  Only output if sel_atoms is not None
                  The index mapping from the selected atoms to sorted selected atoms.
docstring_tokens: Sort atoms in the system according their types . Parameters ---------- coord The coordinates of atoms . Should be of shape [ nframes natoms 3 ] atom_type The type of atoms Should be of shape [ natoms ] sel_atom The selected atoms by type Returns ------- coord_out The coordinates after sorting atom_type_out The atom types after sorting idx_map The index mapping from the input to the output . For example coord_out = coord [ : idx_map : ] sel_atom_type Only output if sel_atoms is not None The sorted selected atom types sel_idx_map Only output if sel_atoms is not None The index mapping from the selected atoms to sorted selected atoms .
function: def sort_input(
          coord : np.ndarray, atom_type : np.ndarray, sel_atoms : List[int] = None
      ):
          """
          Sort atoms in the system according their types.
          
          Parameters
          ----------
          coord
                  The coordinates of atoms.
                  Should be of shape [nframes, natoms, 3]
          atom_type
                  The type of atoms
                  Should be of shape [natoms]
          sel_atom
                  The selected atoms by type
          
          Returns
          -------
          coord_out
                  The coordinates after sorting
          atom_type_out
                  The atom types after sorting
          idx_map
                  The index mapping from the input to the output. 
                  For example coord_out = coord[:,idx_map,:]
          sel_atom_type
                  Only output if sel_atoms is not None
                  The sorted selected atom types
          sel_idx_map
                  Only output if sel_atoms is not None
                  The index mapping from the selected atoms to sorted selected atoms.
          """
          if sel_atoms is not None:
              selection = [False] * np.size(atom_type)
              for ii in sel_atoms:
                  selection += (atom_type == ii)
              sel_atom_type = atom_type[selection]
          natoms = atom_type.size
          idx = np.arange (natoms)
          idx_map = np.lexsort ((idx, atom_type))
          nframes = coord.shape[0]
          coord = coord.reshape([nframes, -1, 3])
          coord = np.reshape(coord[:,idx_map,:], [nframes, -1])
          atom_type = atom_type[idx_map]
          if sel_atoms is not None:
              sel_natoms = np.size(sel_atom_type)
              sel_idx = np.arange(sel_natoms)
              sel_idx_map = np.lexsort((sel_idx, sel_atom_type))
              sel_atom_type = sel_atom_type[sel_idx_map]
              return coord, atom_type, idx_map, sel_atom_type, sel_idx_map
          else:
              return coord, atom_type, idx_map
function_tokens: def sort_input ( coord : np . ndarray , atom_type : np . ndarray , sel_atoms : List [ int ] = None ) : if sel_atoms is not None : selection = [ False ] * np . size ( atom_type ) for ii in sel_atoms : selection += ( atom_type == ii ) sel_atom_type = atom_type [ selection ] natoms = atom_type . size idx = np . arange ( natoms ) idx_map = np . lexsort ( ( idx , atom_type ) ) nframes = coord . shape [ 0 ] coord = coord . reshape ( [ nframes , - 1 , 3 ] ) coord = np . reshape ( coord [ : , idx_map , : ] , [ nframes , - 1 ] ) atom_type = atom_type [ idx_map ] if sel_atoms is not None : sel_natoms = np . size ( sel_atom_type ) sel_idx = np . arange ( sel_natoms ) sel_idx_map = np . lexsort ( ( sel_idx , sel_atom_type ) ) sel_atom_type = sel_atom_type [ sel_idx_map ] return coord , atom_type , idx_map , sel_atom_type , sel_idx_map else : return coord , atom_type , idx_map