schemist 0.0.1__py3-none-any.whl

schemist/__init__.py

@@ -0,0 +1,3 @@
+from importlib.metadata import version
+
+__version__ = version("schemist")

schemist/cleaning.py

@@ -0,0 +1,27 @@
+"""Chemical structure cleaning routines."""
+
+from carabiner.decorators import vectorize
+
+from datamol import sanitize_smiles
+import selfies as sf
+
+@vectorize
+def clean_smiles(smiles: str, 
+                 *args, **kwargs) -> str:
+
+    """Sanitize a SMILES string or list of SMILES strings.
+    
+    """
+
+    return sanitize_smiles(smiles, *args, **kwargs) 
+
+
+@vectorize
+def clean_selfies(selfies: str, 
+                  *args, **kwargs) -> str:
+
+    """Sanitize a SELFIES string or list of SELFIES strings.
+    
+    """
+
+    return sf.encode(sanitize_smiles(sf.decode(selfies), *args, **kwargs))

schemist/cli.py

@@ -0,0 +1,535 @@
+"""Command-line interface for schemist."""
+
+from typing import Any, Dict, List, Optional
+
+from argparse import FileType, Namespace
+from collections import Counter, defaultdict
+from functools import partial
+import os
+import sys
+from tempfile import NamedTemporaryFile, TemporaryDirectory
+
+from carabiner import pprint_dict, upper_and_lower
+from carabiner.cliutils import clicommand, CLIOption, CLICommand, CLIApp
+from carabiner.itertools import tenumerate
+from carabiner.pd import get_formats, write_stream
+
+from . import __version__
+from .collating import collate_inventory, deduplicate_file
+from .converting import _TO_FUNCTIONS, _FROM_FUNCTIONS
+from .generating import AA, REACTIONS
+from .io import _mutate_df_stream
+from .tables import (converter, cleaner, featurizer, assign_groups, 
+                     _assign_splits, splitter, _peptide_table, reactor)
+from .splitting import _SPLITTERS, _GROUPED_SPLITTERS
+
+def _option_parser(x: Optional[List[str]]) -> Dict[str, Any]:
+
+    options = {}
+
+    try:
+        for opt in x:
+
+            try:
+                key, value = opt.split('=')
+            except ValueError:
+                raise ValueError(f"Option {opt} is misformatted. It should be in the format keyword=value.")
+            
+            try:
+                value = int(value)
+            except ValueError:
+                try:
+                    value = float(value)
+                except ValueError:
+                    pass
+
+            options[key] = value
+
+    except TypeError:
+        
+        pass
+
+    return options
+
+
+def _sum_tally(tallies: Counter, 
+               message: str = "Error counts",
+               use_length: bool = False):
+
+    total_tally = Counter()
+    
+    for tally in tallies:
+
+        if use_length:
+            total_tally.update({key: len(value) for key, value in tally.items()})
+        else:
+            total_tally.update(tally)
+
+    if len(tallies) == 0:
+        raise ValueError(f"Nothing generated!")
+
+    pprint_dict(total_tally, message=message)
+
+    return total_tally
+
+
+@clicommand(message="Cleaning file with the following parameters")
+def _clean(args: Namespace) -> None:
+    
+    error_tallies = _mutate_df_stream(input_file=args.input, 
+                                      output_file=args.output, 
+                                      function=partial(cleaner, 
+                                                       column=args.column, 
+                                                       input_representation=args.representation, 
+                                                       prefix=args.prefix),
+                                      file_format=args.format)
+    
+    _sum_tally(error_tallies)
+    
+    return None
+
+
+@clicommand(message="Converting between string representations with the following parameters")
+def _convert(args: Namespace) -> None:
+    
+    options = _option_parser(args.options)
+    
+    error_tallies = _mutate_df_stream(input_file=args.input, 
+                                      output_file=args.output, 
+                                      function=partial(converter, 
+                                                       column=args.column, 
+                                                       input_representation=args.representation, 
+                                                       output_representation=args.to,
+                                                       prefix=args.prefix,
+                                                       options=options),
+                                      file_format=args.format)
+    
+    _sum_tally(error_tallies)
+
+    return None
+
+
+@clicommand(message="Adding features to files with the following parameters")
+def _featurize(args: Namespace) -> None:
+    
+    error_tallies = _mutate_df_stream(input_file=args.input, 
+                                      output_file=args.output, 
+                                      function=partial(featurizer,
+                                                       feature_type=args.feature,
+                                                       column=args.column, 
+                                                       ids=args.id,
+                                                       input_representation=args.representation, 
+                                                       prefix=args.prefix),
+                                      file_format=args.format)
+    
+    _sum_tally(error_tallies)
+
+    return None
+
+
+@clicommand(message="Splitting table with the following parameters")
+def _split(args: Namespace) -> None:
+    
+    split_type = args.type.casefold()
+
+    if split_type in _GROUPED_SPLITTERS:
+
+        chunk_processor, aggregator = _GROUPED_SPLITTERS[split_type]
+
+        with TemporaryDirectory() as dir:
+
+            with NamedTemporaryFile("w", dir=dir, delete=False) as f:
+
+                group_idxs = _mutate_df_stream(input_file=args.input, 
+                                                output_file=f, 
+                                                function=partial(assign_groups,
+                                                                grouper=chunk_processor,
+                                                                group_name=split_type,
+                                                                column=args.column, 
+                                                                input_representation=args.representation),
+                                                file_format=args.format)
+                f.close()
+                new_group_idx = defaultdict(list)
+
+                totals = 0
+                for group_idx in group_idxs:
+                    these_totals = 0
+                    for key, value in group_idx.items():
+                        these_totals += len(value)
+                        new_group_idx[key] += [idx + totals for idx in value]
+                    totals += these_totals
+
+                group_idx = aggregator(new_group_idx, 
+                                    train=args.train, 
+                                    test=args.test)
+                
+                split_tallies = _mutate_df_stream(input_file=f.name, 
+                                                  output_file=args.output, 
+                                                  function=partial(_assign_splits,
+                                                                  split_idx=group_idx,
+                                                                  use_df_index=True),
+                                                  file_format=args.format)
+                if os.path.exists(f.name):
+                    os.remove(f.name)
+        
+    else:
+
+        split_tallies = _mutate_df_stream(input_file=args.input, 
+                                          output_file=args.output, 
+                                          function=partial(splitter,
+                                                           split_type=args.type,
+                                                           column=args.column, 
+                                                           input_representation=args.representation, 
+                                                           train=args.train, 
+                                                           test=args.test,
+                                                           set_seed=args.seed),
+                                          file_format=args.format)
+    
+    _sum_tally(split_tallies, 
+               message="Split counts")
+
+    return None
+
+
+@clicommand(message="Collating files with the following parameters")
+def _collate(args: Namespace) -> None:
+    
+    root_dir = args.data_dir or '.'
+    
+    error_tallies = _mutate_df_stream(input_file=args.input, 
+                                      output_file=args.output, 
+                                      function=partial(collate_inventory, 
+                                                       root_dir=root_dir,
+                                                       drop_unmapped=not args.keep_extra_columns,
+                                                       catalog_smiles_column=args.column,
+                                                       id_column_name=args.id_column,
+                                                       id_n_digits=args.digits,
+                                                       id_prefix=args.prefix),
+                                      file_format=args.format)
+    
+    _sum_tally(error_tallies,
+               message="Collated chemicals:")
+
+    return None
+
+
+@clicommand(message="Deduplicating chemical structures with the following parameters")
+def _dedup(args: Namespace) -> None:
+
+    report, deduped_df = deduplicate_file(args.input,
+                                      format=args.format,
+                                      column=args.column, 
+                                      input_representation=args.representation,
+                                      index_columns=args.indexes)
+    
+    if args.prefix is not None and 'inchikey' in deduped_df:
+        deduped_df = deduped_df.rename(columns={'inchikey': f'{args.prefix}inchikey'})
+    
+    write_stream(deduped_df, 
+                 output=args.output,
+                 format=args.format)
+    
+    pprint_dict(report, message="Finished deduplicating:")
+
+    return None
+
+
+@clicommand(message="Enumerating peptides with the following parameters")
+def _enum(args: Namespace) -> None:
+    
+    tables = _peptide_table(max_length=args.max_length,
+                            min_length=args.min_length,
+                            n=args.number,
+                            indexes=args.slice,
+                            set_seed=args.seed,
+                            prefix=args.prefix,
+                            suffix=args.suffix,
+                            d_aa_only=args.d_aa_only,
+                            include_d_aa=args.include_d_aa, 
+                            generator=True)
+
+    dAA_use = any(aa.islower() for aa in args.prefix + args.suffix)
+    dAA_use = dAA_use or args.include_d_aa or args.d_aa_only
+    
+    tallies, error_tallies = [], []
+    options = _option_parser(args.options)
+    _converter = partial(converter, 
+                         column='peptide_sequence', 
+                         input_representation='minihelm' if dAA_use else 'aa_seq',  ## affects performance
+                         output_representation=args.to,
+                         options=options)
+    
+    for i, table in tenumerate(tables):
+
+        _err_tally, df = _converter(table)
+
+        tallies.append({"Number of peptides": df.shape[0]})
+        error_tallies.append(_err_tally)
+
+        write_stream(df, 
+                     output=args.output,
+                     format=args.format,
+                     mode='w' if i == 0  else 'a',
+                     header=i == 0)
+
+    _sum_tally(tallies,
+               message="Enumerated peptides")
+    _sum_tally(error_tallies,
+               message="Conversion errors")
+
+    return None
+
+
+@clicommand(message="Reacting peptides with the following parameters")
+def _react(args: Namespace) -> None:
+
+    error_tallies = _mutate_df_stream(input_file=args.input, 
+                                      output_file=args.output, 
+                                      function=partial(reactor, 
+                                                       column=args.column, 
+                                                       input_representation=args.representation,
+                                                       reaction=args.reaction,
+                                                       product_name=args.name),
+                                      file_format=args.format)
+    
+    _sum_tally(error_tallies)
+
+    return None
+
+
+def main() -> None:
+
+    inputs = CLIOption('input', 
+                       default=sys.stdin,
+                       type=FileType('r'), 
+                       nargs='?',
+                       help='Input columnar Excel, CSV or TSV file. Default: STDIN.')
+    representation = CLIOption('--representation', '-r', 
+                       type=str,
+                       default='SMILES',
+                       choices=upper_and_lower(_FROM_FUNCTIONS),
+                       help='Chemical representation to use for input. ')
+    column = CLIOption('--column', '-c', 
+                       default='smiles',
+                       type=str,
+                       help='Column to use as input string representation. ')
+    prefix = CLIOption('--prefix', '-p', 
+                       default=None,
+                       type=str,
+                       help='Prefix to add to new column name. Default: no prefix')
+    to = CLIOption('--to', '-2', 
+                       type=str,
+                       default='SMILES',
+                       nargs='*',
+                       choices=upper_and_lower(_TO_FUNCTIONS),
+                       help='Format to convert to.')
+    options = CLIOption('--options', '-x', 
+                       type=str,
+                       default=None,
+                       nargs='*',
+                       help='Options to pass to converter, in the format '
+                           '"keyword1=value1 keyword2=value2"')
+    output = CLIOption('--output', '-o', 
+                       type=FileType('w'),
+                       default=sys.stdout,
+                       help='Output file. Default: STDOUT')                
+    formatting = CLIOption('--format', '-f', 
+                           type=str,
+                           default=None,
+                           choices=upper_and_lower(get_formats()),
+                           help='Override file extensions for input and output. '
+                                'Default: infer from file extension.')
+
+    ## featurize
+    id_feat = CLIOption('--id', '-i', 
+                        type=str,
+                        default=None,
+                        nargs='*',
+                        help='Columns to retain in output table. Default: use all')
+    feature = CLIOption('--feature', '-t', 
+                           type=str,
+                           default='2d',
+                           choices=['2d', 'fp'],  ## TODO: implement 3d
+                           help='Which feature type to generate.')
+
+    ## split
+    type_ = CLIOption('--type', '-t', 
+                       type=str,
+                       default='random',
+                       choices=upper_and_lower(_SPLITTERS),
+                       help='Which split type to use.')
+    train = CLIOption('--train', '-a', 
+                       type=float,
+                       default=1.,
+                       help='Proportion of data to use for training. ')
+    test = CLIOption('--test', '-b', 
+                       type=float,
+                       default=0.,
+                       help='Proportion of data to use for testing. ')
+
+    ## collate
+    data_dir = CLIOption('--data-dir', '-d', 
+                         type=str,
+                         default=None,
+                         help='Directory containing data files. '
+                              'Default: current directory')
+    id_column = CLIOption('--id-column', '-s', 
+                         default=None,
+                         type=str,
+                         help='If provided, add a structure ID column with this name. '
+                              'Default: don\'t add structure IDs')
+    prefix_collate = CLIOption('--prefix', '-p', 
+                         default='ID-',
+                         type=str,
+                         help='Prefix to add to structure IDs. '
+                              'Default: no prefix')
+    digits = CLIOption('--digits', '-n', 
+                         default=8,
+                         type=int,
+                         help='Number of digits in structure IDs. ')
+    keep_extra_columns = CLIOption('--keep-extra-columns', '-x', 
+                         action='store_true',
+                         help='Whether to keep columns not mentioned in the catalog. '
+                              'Default: drop extra columns.')
+    keep_invalid_smiles = CLIOption('--keep-invalid-smiles', '-y', 
+                         action='store_true',
+                         help='Whether to keep rows with invalid SMILES. '
+                              'Default: drop invalid rows.')
+
+    ## dedup
+    indexes = CLIOption('--indexes', '-x', 
+                        type=str,
+                        default=None,
+                        nargs='*',
+                         help='Columns to retain and collapse (if multiple values per unique structure). '
+                              'Default: retain no other columns than structure and InchiKey.')
+    drop_inchikey = CLIOption('--drop-inchikey', '-d', 
+                         action='store_true',
+                         help='Whether to drop the calculated InchiKey column. '
+                              'Default: keep InchiKey.')
+    
+    ### enum 
+    max_length = CLIOption('--max-length', '-l', 
+                           type=int,
+                           help='Maximum length of enumerated peptide. '
+                                'Required.')
+    min_length = CLIOption('--min-length', '-m', 
+                      type=int,
+                      default=None,
+                      help='Minimum length of enumerated peptide. '
+                           'Default: same as maximum, i.e. all peptides same length.')
+    number_to_gen = CLIOption('--number', '-n', 
+                              type=float,
+                              default=None,
+                              help='Number of peptides to sample from all possible '
+                                   'within the constraints. If less than 1, sample '
+                                   'that fraction of all possible. If greater than 1, '
+                                   'sample that number. '
+                                   'Default: return all peptides.')
+    slicer = CLIOption('--slice', '-z', 
+                       type=str,
+                       default=None,
+                       nargs='*',
+                       help='Subset of (possibly sampled) population to return, in the format <stop> '
+                            'or <start> <stop> [<step>]. If "x" is used for <stop>, then it runs to the end. '
+                            'For example, 1000 gives the first 1000, 2 600 gives items 2-600, and '
+                            '3 500 2 gives every other from 3 to 500. Default: return all.')
+    alphabet = CLIOption('--alphabet', '-b', 
+                      type=str,
+                      default=''.join(AA),
+                      help='Alphabet to use in sampling.')
+    suffix = CLIOption('--suffix', '-s', 
+                      type=str,
+                      default='',
+                      help='Sequence to add to end. Lowercase for D-amino acids. '
+                           'Default: no suffix.')
+    set_seed = CLIOption('--seed', '-e', 
+                      type=int,
+                      default=None,
+                      help='Seed to use for reproducible randomness. '
+                           'Default: don\'t enable reproducibility.')
+    d_aa_only = CLIOption('--d-aa-only', '-a', 
+                      action='store_true',
+                      help='Whether to only use D-amino acids. '
+                           'Default: don\'t include.')
+    include_d_aa = CLIOption('--include-d-aa', '-y', 
+                      action='store_true',
+                      help='Whether to include D-amino acids in enumeration. '
+                           'Default: don\'t include.')
+
+    ## reaction
+    name = CLIOption('--name', '-n', 
+                     type=str,
+                     default=None,
+                     help='Name of column for product. '
+                          'Default: same as reaction name.')
+    reaction_opt = CLIOption('--reaction', '-x', 
+                             type=str,
+                             nargs='*',
+                             choices=list(REACTIONS),
+                             default='N_to_C_cyclization',
+                             help='Reaction(s) to apply.')
+
+    clean = CLICommand('clean', 
+                       description='Clean and normalize SMILES column of a table.',
+                       main=_clean,
+                       options=[output, formatting, inputs, representation, column, prefix])
+    convert = CLICommand('convert', 
+                         description='Convert between string representations of chemical structures.',
+                         main=_convert,
+                         options=[output, formatting, inputs, representation, column, prefix, to, options])
+    featurize = CLICommand('featurize', 
+                         description='Convert between string representations of chemical structures.',
+                         main=_featurize,
+                         options=[output, formatting, inputs, representation, column, prefix,
+                                  id_feat, feature])
+    collate = CLICommand('collate', 
+                         description='Collect disparate tables or SDF files of libraries into a single table.',
+                         main=_collate,
+                         options=[output, formatting, inputs, representation,
+                                  data_dir, column.replace(default='input_smiles'), id_column, prefix_collate,
+                                  digits, keep_extra_columns, keep_invalid_smiles])
+    dedup = CLICommand('dedup', 
+                         description='Deduplicate chemical structures and retain references.',
+                         main=_dedup,
+                         options=[output, formatting, inputs, representation, column, prefix,
+                                  indexes, drop_inchikey])
+    enum = CLICommand('enumerate', 
+                      description='Enumerate bio-chemical structures within length and sequence constraints.',
+                      main=_enum,
+                      options=[output, formatting, to, options,
+                               alphabet, max_length, min_length, number_to_gen,
+                               slicer, set_seed,
+                               prefix.replace(default='',  
+                                              help='Sequence to prepend. Lowercase for D-amino acids. '
+                                                   'Default: no prefix.'), 
+                               suffix, 
+                               type_.replace(default='aa',
+                                             choices=['aa'],
+                                             help='Type of bio sequence to enumerate. '
+                                                  'Default: %(default)s.'), 
+                               d_aa_only, include_d_aa])
+    reaction = CLICommand('react', 
+                         description='React compounds in silico in indicated columns using a named reaction.',
+                         main=_react,
+                         options=[output, formatting, inputs, representation, column, name,
+                                  reaction_opt])
+    split = CLICommand('split', 
+                         description='Split table based on chosen algorithm, optionally taking account of chemical structure during splits.',
+                         main=_split,
+                         options=[output, formatting, inputs, representation, column, prefix,
+                                  type_, train, test, set_seed])
+
+    app = CLIApp("schemist",
+                 version=__version__,
+                 description="Tools for cleaning, collating, and augmenting chemical datasets.",
+                 commands=[clean, convert, featurize, collate, dedup, enum, reaction, split])
+
+    app.run()
+
+    return None
+
+
+if __name__ == "__main__":
+
+    main()