diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 2b72dd7..fc5010f 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -28,7 +28,11 @@ jobs: - name: Install package run: pip install . - name: Unit tests - run: python -m unittest discover -v + # Specificity keeps its tests beside its source (see pyproject.toml), + # where discovery from the top level does not reach. + run: | + python -m unittest discover -v + python -m unittest discover -v -s src/immvoke/Sources/Specificity/tests -t src python-versions: runs-on: ubuntu-latest @@ -46,7 +50,11 @@ jobs: python -m pip install --upgrade pip pip install . - name: Unit tests - run: python -m unittest discover -v + # Specificity keeps its tests beside its source (see pyproject.toml), + # where discovery from the top level does not reach. + run: | + python -m unittest discover -v + python -m unittest discover -v -s src/immvoke/Sources/Specificity/tests -t src # Proves that the packaged schema snapshots and the runtime dependencies both # ship in the wheel. Running from /tmp guarantees the source tree cannot be diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index f10a6b5..af030c1 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -78,7 +78,11 @@ jobs: python -m pip install --upgrade pip build twine pip install . - name: Unit tests - run: python -m unittest discover -v + # Specificity keeps its tests beside its source (see pyproject.toml), + # where discovery from the top level does not reach. + run: | + python -m unittest discover -v + python -m unittest discover -v -s src/immvoke/Sources/Specificity/tests -t src - name: Build sdist and wheel run: python -m build - name: Check distribution metadata diff --git a/NEWS.rst b/NEWS.rst index 8916b58..549e99d 100644 --- a/NEWS.rst +++ b/NEWS.rst @@ -1,6 +1,88 @@ Release Notes ================================================================================ +Unreleased +------------------------------------------------------------------------------- + +Specificity: + ++ Added specificity databases IEDB_, CoV-AbDab_ and ASD_ under + ``immvoke specificity search|download``. IEDB's BCR and TCR tables + are downloaded as schema-validated AIRR rearrangement TSVs with one record + per chain, linked by ``cell_id``, taking IEDB's calculated (IMGT + standardized) V/D/J-gene calls over the curated ones, every receptor + species kept unless ``--air_species`` names one (e.g. ``--air_species + "Homo sapiens"``; a species includes its strains, and a chain with no + species recorded is left out), and every assay + outcome unless ``--outcome`` names one (``Positive`` or ``Negative``). + IEDB's B-cell assay tables (``bcell``, ``bcr_to_bcell``) are not + offered: each receptor record already carries its epitope, assay ids and + outcome. + ``--antigen_species`` keeps receptors whose epitope comes from an antigen + organism, by any name IEDB uses, ignoring case (``--antigen_species + "Influenza A virus"`` takes in every flu strain); it asks IEDB's API for + the organisms of the table's epitopes when the table is downloaded. ++ ``immvoke specificity search `` counts what the table holds + once its filters are applied (rows, receptor groups, epitopes and the most + common loci, receptor, antigen species and structure types and outcomes, + ``--top`` of each), instead of + printing the download URL; ``--out`` writes every count as a TSV. Its and + download's filters are grouped in ``--help``, each saying what it filters + on, and ``--limit`` is left only on ``download all``, the one place it + had an effect. ++ IEDB's bcr and tcr tables gain ``iedb_antigen``, + ``iedb_antigen_source_species`` and ``iedb_antigen_structure_type``, each + record's epitope's antigen, its source species and the epitope's + structure type, looked up on IEDB's API at download time. ++ ``--antigen_structure_type`` keeps receptors whose epitope has that IEDB + structure type: ``Linear peptide``, ``Discontinuous peptide`` (which + includes ``Discontinuous peptide on multi chain``) or ``Non-peptidic`` + (haptens, glycans, lipids...), IEDB's own Epitope Structure choice. ++ IEDB's bcr and tcr take ``--assay_type`` (e.g. ``neutralization``) and + ``--disease`` (e.g. ``COVID-19``; a disease includes those under it), + looked up on IEDB's assay API only when given, and ``--paired_only``, + keeping receptors with both chains of a pair. ++ IEDB's bulk ZIP is cached in ``~/.cache/immvoke/iedb/`` for a day, shared + by search and download, so ``bcr`` then ``tcr`` fetch it once. ++ Filter flags generated from a source's schema are spelled with + underscores, as the field is (``--air_species``), and so are + ``immvoke specificity ... download``'s own flags (``--dry_run``, + ``--no_resume``, ``--strict_airr``). ++ Each specificity download lands in its database's own folder, + ``//``: the normalized ``_
.tsv`` (e.g. + ``iedb/iedb_bcr.tsv``) and its own ``download_metadata.yml``; IEDB's + ``raw/`` (the bulk ZIP and what was taken from it) is deleted once the + tables are written. ++ Added ``immvoke annotate``, to annotate a query AIRR table against a + reference one (e.g. IEDB's BCR table) through symscan_: the heavy chain + CDR3, full VDJ or pooled CDR1-3 (``--sequence-region``), within the same + V/J-gene (``--partition-by``), exact or at an identity threshold + (``--identity-threshold``) under Hamming or Levenshtein distance + (``--distance-metric``), or heavy+light chain of the same antibody + (``--paired``, one chain of each kind per cell picked by + ``--count-field``). ++ ``immvoke annotate --outdir`` writes ``annotation.tsv`` and + ``annotation_metadata.yml`` (the options, inputs, reference database and + counts of the run), each named after the reference's database and table, + e.g. ``iedb_bcr_annotation.tsv``. The annotation keeps each query's best + hits, the highest identity (all of them on a tie): ``sequence_id``, + ``identity``, ``hit_id`` and that hit's ``antigen``, ``epitope`` and + ``antigen_source_species`` -- long rather than wide, one row per (query + row, matching reference row), since no separator is safe in IEDB's own + epitope names. A summary of the run in prose is written beside it as + ``annotation_summary.txt``. + ``--full-dist-matrix`` also writes ``annotation_full_matrix.txt``: + every pair the filters admit and how close it is (``sequence_id``, + ``hit_id``, ``identity``). ``--identity-threshold`` does not apply to it + -- it answers how close a query is to everything it was compared with, + not what it matched -- but ``--partition-by`` does, and under ``hamming`` + only equal-length sequences have a distance at all. + +.. _IEDB: https://www.iedb.org/ +.. _CoV-AbDab: https://opig.stats.ox.ac.uk/webapps/covabdab/ +.. _ASD: https://naturalantibody.com/agab/ +.. _symscan: https://github.com/yutanagano/symscan + Version 0.1.0: 2026.09.28 ------------------------------------------------------------------------------- diff --git a/pyproject.toml b/pyproject.toml index 83cda68..fcd6e40 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -37,6 +37,17 @@ dependencies = [ "pandas>=2.2.3", "airr>=2.0", "tqdm>=4.64", + # ASD's data unit is a Delta Lake of Parquet partitions, staged in a + # Google Drive folder: pyarrow reads the partitions, gdown fetches the + # folder, since it has no plain HTTP download. The floor is 16 for the + # same reason as the pandas floor above: pyarrow wheels before 16 were + # built against numpy 1.x and fail to import under numpy>=2. + "pyarrow>=16", + "gdown>=5.0", + # Symscan.py: annotate's matching engine for exact, Hamming and + # Levenshtein matching (a Rust symmetric-deletion index; numpy, which it + # returns its pairs in, comes with it and with pandas). + "symscan>=0.8", ] [project.optional-dependencies] @@ -56,6 +67,11 @@ path = "src/immvoke/Version.py" # automatically as part of the package. Do not add a force-include for them: it # duplicates every file and the build fails. packages = ["src/immvoke"] +# Specificity's own tests live alongside its source +# (Sources/Specificity/tests/) rather than under the top level tests/ +# directory; exclude them here so a built wheel still ships no test code, +# matching every other source under src/immvoke. +exclude = ["src/immvoke/Sources/Specificity/tests/*"] [tool.ruff] line-length = 90 diff --git a/requirements.txt b/requirements.txt index 4b841b6..96c770e 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,3 +4,6 @@ PyYAML>=6.0 pandas>=2.2.3 airr>=2.0 tqdm>=4.64 +pyarrow>=16 +gdown>=5.0 +symscan>=0.8 diff --git a/src/immvoke/Cli.py b/src/immvoke/Cli.py index 906ae04..bd90113 100644 --- a/src/immvoke/Cli.py +++ b/src/immvoke/Cli.py @@ -10,8 +10,10 @@ # Imports import logging +import shutil import sys -from argparse import ArgumentParser +import time +from argparse import ArgumentParser, Namespace from pathlib import Path # Immvoke imports @@ -22,6 +24,16 @@ from immvoke.Http import HttpClient from immvoke.Schema import loadSchema, saveSchema from immvoke.Sources import ALIASES, REGISTRY, canonicalName, getSource +from immvoke.Sources.Specificity import REGISTRY as SPECIFICITY_REGISTRY +from immvoke.Sources.Specificity import Report, getSpecificitySource +from immvoke.Sources.Specificity.AnnotateCli import addAnnotateParser, handleAnnotate +from immvoke.Sources.Specificity.Provenance import writeSpecificityMetadata +from immvoke.Sources.Specificity.Search import ( + DEFAULT_TOP, + SpecificityHelpFormatter, + addFilterGroup, + searchTables, +) from immvoke.Version import __date__, __version__ log = logging.getLogger('immvoke') @@ -40,6 +52,33 @@ #: exists so one field having many values can't blow up every --help screen. VALUE_LIST_CAP = 20 +#: Every source `immvoke schema` works on: both registries and the aliases. +#: Specificity databases keep a snapshot like any other source, so `schema +#: show` lists their filter values -- which is what a rejected filter value +#: tells the user to run (see Schema.validateFilters). +SCHEMA_SOURCES = sorted(set(REGISTRY) | set(SPECIFICITY_REGISTRY) | set(ALIASES)) + + +def getAnySource(name, client): + """ + Instantiate a source from whichever registry holds it. + + `immvoke schema` is the one command that spans both: every source keeps + a snapshot, whether it is mounted under `immvoke ` or + `immvoke specificity `. + + Arguments: + name (str): the source name or alias. + client (HttpClient): the shared HTTP client. + + Returns: + SourceBase: the source. + """ + if canonicalName(name) in REGISTRY: + return getSource(name, client) + + return getSpecificitySource(name, client) + def loadSchemaQuietly(name): """ @@ -81,6 +120,9 @@ def addFilterArgs(parser, schema, source, collection): for item in schema.getCollection(collection).fields: if item.pseudo_values: summary = 'filter on whether %s is recorded' % item.name + elif item.open_values: + summary = ('any name %s uses, checked when fetching, e.g. %s' + % (source, ', '.join(item.values))) elif len(item.values) <= VALUE_LIST_CAP: summary = '%d values: %s' % (len(item.values), ', '.join(item.values)) else: @@ -152,6 +194,8 @@ def getArgParser(): _addReferenceParser(commands) for name, source in sorted(REGISTRY.items()): _addSourceParser(commands, name, source) + _addSpecificityGroupParser(commands) + addAnnotateParser(commands) return parser @@ -245,8 +289,9 @@ def _addSchemaParser(commands): 'collection\'s fields and how many values each accepts, ' 'or every value a single field accepts.', formatter_class=CommonHelpFormatter) - show.add_argument('--source', required=True, choices=sorted(REGISTRY) + sorted(ALIASES), - help='which source (or alias) to read the snapshot of') + show.add_argument('--source', required=True, choices=SCHEMA_SOURCES, + help='which source (or alias) to read the snapshot of, ' + 'a repertoire source or a specificity database') show.add_argument('--collection', default=None, help='list this collection\'s fields; without it, print ' 'one summary line per collection') @@ -262,7 +307,7 @@ def _addSchemaParser(commands): 'fetched detail-page enrichment is carried forward ' 'unless --refresh-details asks to redo it.', formatter_class=CommonHelpFormatter) - refresh.add_argument('--source', required=True, choices=sorted(REGISTRY) + sorted(ALIASES), + refresh.add_argument('--source', required=True, choices=SCHEMA_SOURCES, help='which source to contact and re-harvest') refresh.add_argument('--out', default=None, type=Path, help='directory to write into; without it the packaged ' @@ -395,6 +440,157 @@ def _addSourceParser(commands, name, source): MIRROR_DELAY = 0.05 +def _sharedFilterTable(schema, source): + """ + A table whose filters every table of the database shares, so 'all' can + offer them too (IEDB's bcr and tcr filter on the same fields), or None + when the tables filter differently or not at all. + + Arguments: + schema (SourceSchema): the database's snapshot, or None. + source (type): the specificity source class. + + Returns: + str: the first table, or None. + """ + if schema is None or not all(t in schema.collections for t in source.collections): + return None + fields = {schema.getCollection(t).fields for t in source.collections} + if len(fields) != 1 or not next(iter(fields)): + return None + + return source.collections[0] + + +def _addSpecificityParser(commands, name, source): + """ + Add one specificity database's subcommand tree, with a level per table. + + A trimmed copy of `_addSourceParser`: no --format/samplesheet flags, since + those are repertoire specific, and one extra synthetic 'all' table that + loops every table the database offers. + """ + schema = loadSchemaQuietly(name) + + parser = commands.add_parser( + name, help=source.description, + description='%s\n\nHomepage: %s' % (source.description, source.homepage), + formatter_class=CommonHelpFormatter) + actions = parser.add_subparsers(dest='action', metavar='ACTION', + required=True) + + action_help = { + 'search': ('count what a table holds', + 'Count what a table holds, optionally narrowed by the ' + 'filters: its rows and the most common values of its key ' + 'columns. The table is fetched into a temporary folder and ' + 'nothing else is written, unless --out asks for the ' + 'counts.'), + 'download': ('download a table', + 'Fetch a table, optionally narrowed by the filters, and ' + 'write it as one normalized TSV: a real AIRR ' + 'rearrangement TSV for a table listed in the database\'s ' + 'airr_collections, a plain TSV of the source\'s own ' + 'columns otherwise.'), + } + for action, (helptext, description) in action_help.items(): + action_parser = actions.add_parser(action, help=helptext, + description=description, + formatter_class=CommonHelpFormatter) + tables = action_parser.add_subparsers(dest='table', metavar='TABLE', + required=True) + + entries = [(table, source.collection_help.get(table), table) + for table in source.collections] + entries.append(('all', 'every table', None)) + + for table_name, table_help, filter_table in entries: + if filter_table is None: + summary = 'every table this database offers' + else: + summary = 'the %s table: %s' % (table_name, table_help) + leaf = tables.add_parser( + table_name, help=table_help, + description='%s %s.' % (action.capitalize(), summary), + formatter_class=SpecificityHelpFormatter) + if filter_table is not None: + addFilterGroup(leaf, schema, source, filter_table) + elif _sharedFilterTable(schema, source) is not None: + addFilterGroup(leaf, schema, source, _sharedFilterTable(schema, source)) + if action == 'search': + leaf.add_argument('--top', type=int, default=DEFAULT_TOP, + help='show at most this many values per column') + leaf.add_argument('-o', '--out', type=Path, default=None, + help='also write every count to this TSV file') + else: + if filter_table is None: + leaf.add_argument('--limit', type=int, default=None, + help='stop after this many tables') + leaf.add_argument('--outdir', type=Path, required=True, + help='directory to write into, one ' + 'folder per database: ' + '/_
.tsv is the ' + 'normalized table') + leaf.add_argument('--dry_run', action='store_true', + help='report what would be fetched, then stop') + leaf.add_argument('--no_resume', action='store_true', + help='re-download in full rather than ' + 'continuing a partly fetched file') + leaf.add_argument('--strict_airr', action='store_true', + help='for a table written as AIRR, drop ' + 'columns the AIRR schema does not ' + 'define; no effect on other tables') + + +def _addSpecificityGroupParser(commands): + """ + Add the `specificity` subcommand: `list`, one per registered database, and + a synthetic `all` that fans out to every registered database. + """ + group = commands.add_parser( + 'specificity', help='search and download specificity databases', + description='Search and download specificity databases (reference ' + 'tables of epitopes, assays and receptor sequences), as ' + 'opposed to repertoire sequencing sources.', + formatter_class=CommonHelpFormatter) + dbs = group.add_subparsers(dest='db', metavar='DB', required=True) + + dbs.add_parser( + 'list', help='list the specificity databases immvoke knows about', + description='List every specificity database immvoke knows how to ' + 'fetch from, along with a one-line description and its ' + 'homepage.', + formatter_class=CommonHelpFormatter) + + for name, source in sorted(SPECIFICITY_REGISTRY.items()): + _addSpecificityParser(dbs, name, source) + + all_group = dbs.add_parser( + 'all', help='every registered specificity database', + description='Act on every registered specificity database at once.', + formatter_class=CommonHelpFormatter) + all_actions = all_group.add_subparsers(dest='action', metavar='ACTION', + required=True) + all_download = all_actions.add_parser( + 'download', help='download every table of every database', + description='Download every table of every registered specificity ' + 'database. Filters cannot be applied here, since ' + 'different databases have unrelated fields; target one ' + 'database directly to filter it.', + formatter_class=CommonHelpFormatter) + all_download.add_argument('--outdir', type=Path, required=True, + help='directory to write into; one ' + 'subdirectory per database') + all_download.add_argument('--limit', type=int, default=None, + help='stop after this many tables, per database') + all_download.add_argument('--dry_run', action='store_true', + help='report what would be fetched, then stop') + all_download.add_argument('--strict_airr', action='store_true', + help='for a table written as AIRR, drop columns ' + 'the AIRR schema does not define; no effect ' + 'on other tables') + + def makeClient(args): """Build the shared HTTP client.""" return HttpClient() @@ -432,15 +628,28 @@ def handleSchemaShow(args): collection = schema.getCollection(args.collection) if args.field is None: + width = max(len(x.name) for x in collection.fields) if collection.fields else 0 for item in collection.fields: - kind = 'presence only' if item.pseudo_values else '%d values' % len(item.values) - print('%-14s %s' % (item.name, kind)) + if item.pseudo_values: + kind = 'presence only' + elif item.open_values: + # The source checks the value itself, so what the snapshot + # holds are examples; calling them the vocabulary would read + # as though nothing else were accepted. + kind = 'any value (%d examples)' % len(item.values) + else: + kind = '%d values' % len(item.values) + print('%s %s' % (item.name.ljust(width), kind)) return 0 item = collection.getField(args.field) if item is None: raise ImmvokeError("no field '%s' in %s %s" % (args.field, args.source, args.collection)) + if item.open_values: + # On stderr, so that the values themselves still pipe cleanly. + log.info("%s accepts any value %s checks; these are examples", + args.field, args.source) for value in item.values: print(value) @@ -451,7 +660,7 @@ def handleSchemaRefresh(args): """Re-harvest a snapshot and its catalogs.""" args.source = canonicalName(args.source) client = makeClient(args) - source = getSource(args.source, client) + source = getAnySource(args.source, client) log.info('harvesting %s search schema', args.source) schema = source.harvestSchema() @@ -465,8 +674,14 @@ def handleSchemaRefresh(args): written, changed = saveSchema(schema, out) log.info('%s %s', 'wrote' if changed else 'unchanged, left alone:', written) - wanted = args.collection or list(source.collections) - for collection in wanted: + # Only a source whose units come from a catalog of its own has one to + # refresh. For a germline set or a specificity database the snapshot is + # the whole of it, so the harvest ends here rather than reaching for a + # harvestCatalog that source never defined. + if not hasattr(source, 'harvestCatalog'): + return 0 + + for collection in args.collection or list(source.collections): log.info('harvesting %s %s catalog', args.source, collection) rows = source.harvestCatalog(collection, schema=schema) @@ -789,6 +1004,172 @@ def handleDownload(args): return 0 +def _resolveSpecificityTables(source, args): + """ + Return the tables a specificity action should act on. + + Arguments: + source (SourceBase): the specificity source. + args (Namespace): parsed arguments, carrying `table` and, for + `download all`, `limit`. + + Returns: + list: table names. `--limit` caps how many are used when `table` is + 'all'; a single named table, and search, take no --limit. + """ + if args.table != 'all': + return [args.table] + + tables = list(source.collections) + limit = getattr(args, 'limit', None) + + return tables[:limit] if limit is not None else tables + + +def _writeSpecificityTsv(chunks, out): + """ + Write normalized specificity records as a TSV, streaming chunk by chunk. + + Unlike `Convert.writeAirr`, this performs no AIRR schema validation: + specificity columns are a source's own, not AIRR rearrangement fields. + + Arguments: + chunks (iterable): DataFrames of normalized records. + out (Path): output path. + + Returns: + int: rows written. + """ + out = Path(out) + out.parent.mkdir(parents=True, exist_ok=True) + + written = 0 + header_written = False + with open(out, 'w', newline='') as handle: + for frame in chunks: + frame.to_csv(handle, sep='\t', index=False, header=not header_written) + header_written = True + written += len(frame) + + if not header_written: + out.write_text('') + + return written + + +def handleSpecificityList(args): + """List registered specificity databases.""" + for name, source in sorted(SPECIFICITY_REGISTRY.items()): + print('%-10s %s' % (name, source.description)) + print('%-10s %s' % ('', source.homepage)) + if source.license: + print('%-10s license: %s' % ('', source.license)) + for paper in source.citation: + print('%-10s cite: %s' % ('', paper)) + + return 0 + + +def handleSpecificitySearch(args): + """Count what a table (or every table) holds; see Specificity.Search.""" + source = getSpecificitySource(args.db, makeClient(args)) + + return searchTables(source, _resolveSpecificityTables(source, args), + collectFilters(args), top=args.top, out=args.out) + + +def handleSpecificityDownload(args): + """ + Download one or more tables: a raw mirror plus one normalized TSV each, + all in the database's own folder, `//` -- so several + databases can share one outdir without mixing their files or their + download_metadata.yml. + """ + client = makeClient(args) + source = getSpecificitySource(args.db, client) + filters = collectFilters(args) + tables = _resolveSpecificityTables(source, args) + + outdir = Path(args.outdir) / args.db + raw_dir = outdir / 'raw' + + records = [] + formats_written = set() + for table in tables: + query = source.validateQuery(table, filters) + units = source.searchUnits(query) + + if args.dry_run: + for unit in units: + print('%-16s %s' % (table, unit.unit_id)) + continue + + for unit in units: + start = time.monotonic() + Report.started(source, table, filters) + result = source.fetchUnit(unit, raw_dir, resume=not args.no_resume) + _, chunks, report = source.convertUnit(result.path, unit) + tally = Report.Tally(source, table) + chunks = tally.wrap(chunks) + + # Named after the database as well as the table, so the file + # says where it came from once copied out of /. + dest = outdir / ('%s_%s.tsv' % (args.db, table)) + if table in source.airr_collections: + validation = Convert.writeAirr(chunks, dest, + strict=args.strict_airr) + Convert.writeValidationReport(validation, dest) + rows = validation['rows_checked'] + output_format = 'airr' + else: + validation = None + rows = _writeSpecificityTsv(chunks, dest) + output_format = 'tsv' + formats_written.add(output_format) + Report.completed(dest, time.monotonic() - start, report['rows_in'], rows, + tally, validation) + + records.append(Provenance.buildUnitRecord( + unit, result, outdir, {output_format: dest})) + + if args.dry_run: + log.info('dry run: nothing downloaded') + return 0 + + if not source.keep_raw: + shutil.rmtree(raw_dir, ignore_errors=True) + log.debug('removed %s', raw_dir) + + # A specificity download represents the database's current state, not one + # more addition to a growing set the way OAS downloads are, so this always + # replaces the metadata file rather than accumulating old runs into it -- + # see Specificity.Provenance for why that is a separate writer. + record = writeSpecificityMetadata( + outdir, args.db, args.table, filters, getattr(args, 'limit', None), records, + schema=source.schema, license=source.license, citation=source.citation, + formats=formats_written or ('tsv',)) + Report.finished(record) + + return 0 + + +def handleSpecificityDownloadAll(args): + """ + Download every table of every registered specificity database, each + into its own `//` (see handleSpecificityDownload). + """ + outdir = Path(args.outdir) + + for name in sorted(SPECIFICITY_REGISTRY): + log.info('downloading specificity source %s', name) + handleSpecificityDownload(Namespace( + db=name, table='all', outdir=outdir, limit=args.limit, + dry_run=args.dry_run, no_resume=False, + strict_airr=args.strict_airr)) + + return 0 + + def main(): """ Parse the commandline and dispatch to the selected subcommand. @@ -831,6 +1212,21 @@ def main(): if args.action == 'download': return handleDownload(args) + if args.command == 'specificity': + if args.db == 'list': + return handleSpecificityList(args) + if args.db == 'all': + return handleSpecificityDownloadAll(args) + # The action and table levels are required subparsers, so argparse + # has already rejected a commandline missing either. + if args.action == 'search': + return handleSpecificitySearch(args) + if args.action == 'download': + return handleSpecificityDownload(args) + + if args.command == 'annotate': + return handleAnnotate(args) + parser.print_help(sys.stderr) return 1 except ImmvokeError as error: diff --git a/src/immvoke/Commandline.py b/src/immvoke/Commandline.py index 32a0db5..d28b65a 100644 --- a/src/immvoke/Commandline.py +++ b/src/immvoke/Commandline.py @@ -43,7 +43,22 @@ def setupLogging(verbose=False, quiet=False): else: level = logging.INFO - logging.basicConfig(stream=sys.stderr, level=level, - format='%(levelname)s %(message)s') + handler = logging.StreamHandler(sys.stderr) + handler.setFormatter(StatusFormatter('%(message)s')) + logging.basicConfig(level=level, handlers=[handler]) return logging.getLogger('immvoke') + + +class StatusFormatter(logging.Formatter): + """ + Write status (INFO) lines as plain text, and name the level of every + other line: 'WARNING ...', 'ERROR ...', 'DEBUG ...'. + """ + + def format(self, record): + text = super().format(record) + if record.levelno == logging.INFO: + return text + + return '%s %s' % (record.levelname, text) diff --git a/src/immvoke/Convert.py b/src/immvoke/Convert.py index 783c3a7..74cdf04 100644 --- a/src/immvoke/Convert.py +++ b/src/immvoke/Convert.py @@ -13,6 +13,9 @@ import logging from pathlib import Path +# Immvoke imports +from immvoke import Provenance + log = logging.getLogger(__name__) #: Identifier written when a record has no sequence_id. Only the identifier is @@ -100,23 +103,51 @@ def newValidation(): 'rows_invalid': 0, 'errors': collections.Counter()} -def summarizeValidation(validation): +def summarizeValidation(validation, name='the file'): """ - Render a validation report as text. + Render a validation report as plain sentences. Arguments: validation (dict): a report from newValidation(). + name (str): the rearrangement file the report describes, as it should + be named in the text. Returns: - str: a human readable summary. + str: a human readable summary: whether the columns and rows are valid + AIRR, each distinct row error with how many rows it hit, and an + overall verdict. """ - lines = ['header_valid: %s' % validation['header_valid']] - if validation['header_error']: - lines.append('header_error: %s' % validation['header_error']) - lines.append('rows_checked: %d' % validation['rows_checked']) - lines.append('rows_invalid: %d' % validation['rows_invalid']) - for message, count in validation['errors'].most_common(): - lines.append(' %6d %s' % (count, message)) + checked, invalid = validation['rows_checked'], validation['rows_invalid'] + lines = ['AIRR validation of %s' % name, + 'Checked: %s' % Provenance.timestamp(), ''] + + if validation['header_valid'] is None: + lines.append('Columns: not checked, the table is empty.') + elif validation['header_valid']: + lines.append('Columns: OK, every required AIRR field is present.') + else: + lines.append('Columns: PROBLEM, %s' % validation['header_error']) + + if not checked: + lines.append('Rows: none to check.') + elif not invalid: + lines.append('Rows: %s checked, all valid.' % format(checked, ',')) + else: + lines.append('Rows: %s checked, %s invalid (%.1f%%).' + % (format(checked, ','), format(invalid, ','), + 100.0 * invalid / checked)) + lines += ['', 'Row errors, most common first (rows affected, error):'] + for message, count in validation['errors'].most_common(): + lines.append(' %8s %s' % (format(count, ','), message)) + + lines.append('') + if validation['header_valid'] is False or invalid: + lines.append('Result: FAIL. %s was still written in full; this log only ' + 'records what did not meet the AIRR schema.' % name) + elif checked: + lines.append('Result: PASS. %s is a valid AIRR rearrangement file.' % name) + else: + lines.append('Result: nothing to validate.') return '\n'.join(lines) + '\n' @@ -194,7 +225,8 @@ def writeAirr(chunks, out, strict=False, validation=None): def writeValidationReport(validation, out): """ - Write a validation summary next to its rearrangement file. + Write a validation summary next to its rearrangement file, as + .airr_validation.log (bcr.tsv -> bcr.airr_validation.log). Arguments: validation (dict): the report. @@ -203,8 +235,9 @@ def writeValidationReport(validation, out): Returns: Path: the report file. """ - report = Path(str(out) + '.validation.txt') - report.write_text(summarizeValidation(validation)) + out = Path(out) + report = out.with_suffix('.airr_validation.log') + report.write_text(summarizeValidation(validation, out.name)) return report diff --git a/src/immvoke/Exceptions.py b/src/immvoke/Exceptions.py index fd818fd..9079449 100644 --- a/src/immvoke/Exceptions.py +++ b/src/immvoke/Exceptions.py @@ -50,6 +50,26 @@ class OgrdbParseError(ParseError): pass +class IedbParseError(ParseError): + """IEDB content did not match the expected structure.""" + pass + + +class CovAbDabParseError(ParseError): + """CoV-AbDab content did not match the expected structure.""" + pass + + +class AsdParseError(ParseError): + """ASD (Delta Lake) content did not match the expected structure.""" + pass + + +class AsdFetchError(ImmvokeError): + """ASD's Google Drive folder could not be fetched.""" + pass + + class SchemaError(ImmvokeError): """A stored schema snapshot is missing, malformed or too new to understand.""" pass @@ -58,3 +78,8 @@ class SchemaError(ImmvokeError): class ConversionError(ImmvokeError): """A data unit could not be converted.""" pass + + +class AnnotationError(ImmvokeError): + """A specificity annotation request could not be run as given.""" + pass diff --git a/src/immvoke/Schema.py b/src/immvoke/Schema.py index 485ae69..ecaa48b 100644 --- a/src/immvoke/Schema.py +++ b/src/immvoke/Schema.py @@ -46,16 +46,21 @@ class Field: wildcard (str): the value meaning "all". pseudo_values (bool): True when values are presence flags rather than a real vocabulary, in which case offering them as choices would mislead. + open_values (bool): True when any value is accepted here and the source + checks it itself when fetching, because its vocabulary is too large + or changes too often to snapshot (IEDB's antigen organisms). values + then holds examples, not the full list. """ name: str values: tuple = () wildcard: str = '*' pseudo_values: bool = False + open_values: bool = False @property def flag(self): """str: the commandline flag generated for this field.""" - return '--%s' % self.name.lower().replace(' ', '-').replace('_', '-') + return '--%s' % self.name.lower().replace(' ', '_').replace('-', '_') def accepts(self, value): """ @@ -67,7 +72,7 @@ def accepts(self, value): Returns: bool: True if the value may be sent. """ - if value == self.wildcard: + if value == self.wildcard or self.open_values: return True # A presence-only field has no vocabulary of its own; what it accepts is @@ -279,7 +284,8 @@ def fromDict(payload): Field(name=x['name'], values=tuple(x.get('values') or ()), wildcard=x.get('wildcard', '*'), - pseudo_values=bool(x.get('pseudo_values', False))) + pseudo_values=bool(x.get('pseudo_values', False)), + open_values=bool(x.get('open_values', False))) for x in (body.get('fields') or [])) collections[name] = Collection( name=name, fields=fields, @@ -314,12 +320,15 @@ def toDict(schema): collections = {} for name in sorted(schema.collections): collection = schema.collections[name] + # open_values is written only where set, so snapshots of sources + # that never use it stay byte identical. collections[name] = { 'reported_totals': dict(collection.reported_totals), 'fields': [{'name': x.name, 'wildcard': x.wildcard, 'pseudo_values': x.pseudo_values, - 'values': list(x.values)} + 'values': list(x.values), + **({'open_values': True} if x.open_values else {})} for x in collection.fields]} return {'source': schema.source, diff --git a/src/immvoke/Sources/Base.py b/src/immvoke/Sources/Base.py index 8d1f641..e86deb4 100644 --- a/src/immvoke/Sources/Base.py +++ b/src/immvoke/Sources/Base.py @@ -117,6 +117,28 @@ class SourceBase(ABC): #: Shown in `immvoke sources list` and recorded in download provenance, #: since it is what governs how downloaded data may be reused. license = '' + #: Specificity collections whose normalizeChunk output is AIRR named + #: rearrangement records rather than the source's own column names. Empty + #: for every non-specificity source and for specificity collections that + #: are not sequences (assay results, join tables, ...), since those have + #: no natural AIRR shape. `immvoke specificity download` writes a + #: real, schema-validated rearrangement TSV for a listed collection and a + #: plain TSV for everything else. + airr_collections = frozenset() + #: Whether `immvoke specificity download` keeps the raw mirror, + #: /raw/, once every table has been written. False for a source whose + #: raw files are only an intermediate step, such as IEDB's bulk ZIP and + #: the tables extracted from it; the normalized TSVs are what remain. + keep_raw = True + #: Specificity filter field -> (placeholder, what it filters on), shown + #: for its flag in `immvoke specificity download|search --help`. A + #: field missing here gets a generic VALUE and only its values listed. + filter_help = {} + #: Specificity table -> the columns `immvoke specificity search` + #: reports on, as ((column, label, how), ...): how is 'distinct' (count + #: its distinct values) or 'top' (its most common values, with counts). + #: A table missing here is reported as a row count alone. + search_columns = {} #: How to cite this source, one string per paper, oldest first. Shown in #: `immvoke sources list` and recorded in download provenance, so that #: the record of what was downloaded travels with a reminder of how to @@ -212,6 +234,20 @@ def validateQuery(self, collection, filters): return Query(collection=collection, filters=resolved) + def labelValue(self, column, value): + """ + Name a column's value for `immvoke specificity search`, e.g. a + taxon code as its species. The value as is by default. + + Arguments: + column (str): the column, from search_columns. + value (str): one of its values. + + Returns: + str: the label to show. + """ + return value + def resolveOutputPath(self, unit, outdir): """ Build the mirror location for a unit. diff --git a/src/immvoke/Sources/Specificity/Annotate.py b/src/immvoke/Sources/Specificity/Annotate.py new file mode 100644 index 0000000..3b9e7a2 --- /dev/null +++ b/src/immvoke/Sources/Specificity/Annotate.py @@ -0,0 +1,1233 @@ +""" +Antigen specificity annotation: query-vs-reference receptor comparison + +For each row of a query AIRR table, find every row of a fixed reference AIRR +table it agrees with, and report its best hits -- one row per (query row, +matching reference row): its `sequence_id`, `identity`, `hit_id` and that +hit's `antigen`, `epitope` and `antigen_source_species`, long rather than +wide so no cell has to hold two values -- plus, on request, how close every pair the +filters admit is (full_dist_matrix). This is the last pipeline stage: download a +specificity database's table, which lands as an AIRR rearrangement TSV in +the database's own folder (`immvoke specificity iedb download bcr`), then +annotate a user's own AIRR table against it as the reference, into an output +folder of its own. + +annotateSpecificity() is the entry point `immvoke annotate` wraps: it picks +the comparison from a few knobs (sequence_region, partition_by, +paired, identity_threshold, distance_metric) -- see its section at the end of +this module -- and runs it through Symscan.runSymscanAnnotation, symscan's +symmetric-deletion index. Rows beyond symscan's limits are scored directly +instead (Distances.py, _scoreSkippedRows), so no row misses a hit. + +runPairedExactAnnotation builds on an exact match for the case where a +query row is one chain of an antibody rather than a self-contained unit: it +runs it once per chain (heavy, light) and reconciles the two per antibody, +classifying each as agreeing with the same reference antibody on both +chains, one chain only, or neither. + +V/J-gene agreement goes through Genes.py (`gene_cols`), which +parses gene calls rather than comparing them as strings; `exact_cols` +compares plain strings. Substitution-matrix-weighted (BLOSUM) annotation is +not included. +""" + +# Info +__author__ = 'Pramod Shinde' + +# Imports +import logging +import sys +import time +from pathlib import Path + +import numpy +import pandas +import yaml + +# Immvoke imports +from immvoke.Exceptions import AnnotationError +from immvoke.Provenance import ( + DOWNLOAD_METADATA, + METADATA_VERSION, + commandLine, + timestamp, +) +from immvoke.Sources.Specificity import Report +from immvoke.Version import __version__ + +log = logging.getLogger(__name__) + +#: Reference column read as each hit's identifier when the caller does not +#: say otherwise -- the AIRR rearrangement primary key. +DEFAULT_ID_COL = 'sequence_id' +#: An unresolved residue: under exact and Hamming matching, a position where +#: either side has it never counts as a difference (see Symscan.py). +WILDCARD = 'X' + +#: runPairedExactAnnotation's fixed column names -- one antibody's heavy and +#: light chain are two rows of the same AIRR table, linked by `cell_id` and +#: told apart by `locus`; there is nothing to make configurable here beyond +#: what the AIRR schema already fixes these fields to mean. +CELL_ID_COL = 'cell_id' +LOCUS_COL = 'locus' +HEAVY_LOCUS = 'IGH' +LIGHT_LOCI = ('IGK', 'IGL') +#: The AIRR count column that picks one heavy and one light chain per query +#: cell in paired annotation, unless the caller names another. +DEFAULT_COUNT_FIELD = 'duplicate_count' + +#: What annotateSpecificity writes into its outdir: the result, the record +#: of the run and, with full_dist_matrix, how close every compared pair is -- +#: each name prefixed with the reference's database and table +#: (_outputPrefix), e.g. iedb_bcr_annotation.tsv. +ANNOTATION_FILE = 'annotation.tsv' +ANNOTATION_METADATA = 'annotation_metadata.yml' +FULL_MATRIX_FILE = 'annotation_full_matrix.txt' +ANNOTATION_SUMMARY = 'annotation_summary.txt' + +#: What the result reports about each best hit -> the reference column it +#: comes from (IEDB's, see Iedb.py); a reference without one leaves it blank. +HIT_ANNOTATIONS = ( + ('antigen', 'iedb_antigen'), + ('epitope', 'iedb_epitope_name'), + ('antigen_source_species', 'iedb_antigen_source_species'), +) + +#: Slack added before truncating a fraction-of-length threshold to a count, +#: so float error can't drop it by one (e.g. (1 - 0.9) * 10 is 0.99999...). +_FRACTION_EPSILON = 1e-9 + + +def _allowedDifferences(fraction, length): + """ + Resolve a fraction-of-length tolerance to a whole number of differences. + + Arguments: + fraction (float): the tolerated fraction, e.g. 0.1. + length (int): the query's length the fraction applies to. + + Returns: + int: floor(fraction * length), robust to float representation error. + """ + return int(fraction * length + _FRACTION_EPSILON) + + +def _mismatchPositions(a, b): + """ + Every substitution mismatch between `a` and `b`, a WILDCARD on either + side of a position never counting as one. + + Arguments: + a (tuple): one row's values, one per comparison column. + b (tuple): the row it is being compared to, same shape. + + Returns: + list: (column index, 0-based position) pairs. + """ + return [(col, pos) + for col, (x, y) in enumerate(zip(a, b)) + for pos, (cx, cy) in enumerate(zip(x, y)) + if cx != cy and cx != WILDCARD and cy != WILDCARD] + + +def _formatPositions(positions, n_cols): + """ + Render mismatch positions as a comma-separated string. + + Arguments: + positions (list): (column index, position) pairs, from + _mismatchPositions. + n_cols (int): how many comparison columns are in play. + + Returns: + str: bare positions ("3,7") for the common single-column case, + column-prefixed ("0:3,1:2") when there is more than one -- a bare + position is ambiguous once more than one column is being compared. + """ + if n_cols == 1: + return ','.join(str(pos) for _, pos in positions) + + return ','.join('%d:%d' % (col, pos) for col, pos in positions) + + +def _rowValues(frame, columns): + """ + Read one or more columns as a list of per-row tuples. + + Arguments: + frame (pandas.DataFrame): the table to read from. + columns (list): column names, in the order values should be compared. + + Returns: + list: one tuple per row, values in `columns` order, blank standing in + for a missing value. + """ + # astype(object): an all-blank column read as float takes a '' fill too. + return list(frame[columns].astype(object).fillna('').itertuples(index=False, + name=None)) + + +def _exactValues(frame, exact_cols): + """ + _rowValues for the exact-agreement gate, or one empty tuple per row when + there is no gate. + + Arguments: + frame (pandas.DataFrame): the table to read from. + exact_cols (list): column names, or None. + + Returns: + list: one tuple per row. + """ + return _rowValues(frame, exact_cols) if exact_cols else [()] * len(frame) + + +# ── paired heavy+light exact annotation ───────────────────────────────────── + +def _hitsByCellId(query_chain, chain_result, cell_id_col): + """ + Fold one chain's exact match result into a hit set per antibody. + + Arguments: + query_chain (pandas.DataFrame): the query rows matched for this + chain, in the same order (its index reset to 0..n-1, matching + `chain_result`'s `query_row`). + chain_result (pandas.DataFrame): the exact match result for + `query_chain`. + cell_id_col (str): column in `query_chain` identifying the antibody a + chain row belongs to. + + Returns: + dict: query antibody id to the set of reference antibody ids (read off + the reference's `cell_id_col`) whose corresponding chain agreed with + it. + """ + cell_ids = query_chain[cell_id_col].reset_index(drop=True) + hits = {} + for row in chain_result.itertuples(index=False): + if not row.hit_ids: + continue + cell_id = cell_ids.iloc[row.query_row] + hits.setdefault(cell_id, set()).update(row.hit_ids.split(';')) + + return hits + + +def runPairedExactAnnotation(queries, reference, cdr3_col='cdr3_aa', + vgene_col='v_call', jgene_col='j_call', + threads=1, return_pairs=False): + """ + Classify each query antibody by whether the *same* reference antibody + agrees with it on the heavy chain, the light chain, both, or neither. + + An antibody is two rows of an AIRR table -- a heavy chain and, where + present, a light chain -- sharing one `cell_id` and told apart by + `locus` (`IGH` for heavy, `IGK`/`IGL` for light). Agreement on a chain + is the same V-gene and J-gene (any allele, see Genes.py) and + exact CDR3, run once for the heavy rows and once for the light rows; a + reference antibody counts toward "both" only if its `cell_id` appears in + *both* chains' hit sets for the query antibody, not merely in each + independently -- a query whose heavy chain happens to match one + reference antibody and whose light chain happens to match an unrelated + one is `conflicting`, not `both`. + + Arguments: + queries (pandas.DataFrame): query AIRR rows, both chains, with `locus` + and `cell_id` columns. + reference (pandas.DataFrame): reference AIRR rows, same shape. + cdr3_col (str): the CDR3 amino acid column. + vgene_col (str): the V-gene call column. + jgene_col (str): the J-gene call column. + threads (int): symscan's thread pool; see Symscan.runSymscanAnnotation. + return_pairs (bool): also return the (query, reference) antibody pairs + behind the hit ids. + + Returns: + pandas.DataFrame: one row per query `cell_id`: `classification` (one + of `both`, `heavy_only`, `light_only`, `conflicting`, `neither`), + `heavy_hit_ids`, `light_hit_ids` (every reference antibody agreeing on + that chain alone, semicolon joined, sorted) and `both_hit_ids` (their + intersection -- non-empty only when `classification` is `both`); its + `attrs['skipped_query_rows']` counts the query chain rows symscan + could not compare (see Symscan.py). With `return_pairs`, a tuple: that table, and one row per (query antibody, + reference antibody) agreeing on either chain -- `cell_id`, `hit_id` + and `matched_chains` (`both`, `heavy` or `light`). + """ + query_heavy = queries[queries[LOCUS_COL] == HEAVY_LOCUS] + query_light = queries[queries[LOCUS_COL].isin(LIGHT_LOCI)] + reference_heavy = reference[reference[LOCUS_COL] == HEAVY_LOCUS] + reference_light = reference[reference[LOCUS_COL].isin(LIGHT_LOCI)] + + query_heavy = query_heavy.reset_index(drop=True) + query_light = query_light.reset_index(drop=True) + reference_heavy = reference_heavy.reset_index(drop=True) + reference_light = reference_light.reset_index(drop=True) + + gene_cols = [c for c in (vgene_col, jgene_col) if c is not None] or None + # Deferred: Symscan builds on this module's helpers. + from immvoke.Sources.Specificity.Symscan import runSymscanAnnotation + + def _exact(query_chain, reference_chain): + return runSymscanAnnotation(query_chain, reference_chain, [cdr3_col], 0, + id_col=CELL_ID_COL, gene_cols=gene_cols, + threads=threads) + + heavy_result = _exact(query_heavy, reference_heavy) + light_result = _exact(query_light, reference_light) + + heavy_hits = _hitsByCellId(query_heavy, heavy_result, CELL_ID_COL) + light_hits = _hitsByCellId(query_light, light_result, CELL_ID_COL) + + rows, pairs = [], [] + for cell_id in sorted(set(queries[CELL_ID_COL])): + heavy = heavy_hits.get(cell_id, set()) + light = light_hits.get(cell_id, set()) + both = heavy & light + for hit in sorted(heavy | light): + matched = 'both' if hit in both else 'heavy' if hit in heavy else 'light' + pairs.append({'cell_id': cell_id, 'hit_id': hit, 'matched_chains': matched}) + + if both: + classification = 'both' + elif heavy and light: + classification = 'conflicting' + elif heavy: + classification = 'heavy_only' + elif light: + classification = 'light_only' + else: + classification = 'neither' + + rows.append({ + 'cell_id': cell_id, + 'classification': classification, + 'heavy_hit_ids': ';'.join(sorted(heavy)), + 'light_hit_ids': ';'.join(sorted(light)), + 'both_hit_ids': ';'.join(sorted(both)), + }) + + result = pandas.DataFrame(rows) + result.attrs['skipped_query_rows'] = (heavy_result.attrs['skipped_query_rows'] + + light_result.attrs['skipped_query_rows']) + if not return_pairs: + return result + + return result, pandas.DataFrame(pairs, columns=['cell_id', 'hit_id', + 'matched_chains']) + + +# ── high level API: annotateSpecificity() ─────────────────────────────────── +# +# A file-or-DataFrame-in, DataFrame-and-files-out entry point, for a caller +# who wants "match this receptor table against that reference table". It +# picks the comparison from a small set of knobs (paired, count_field, +# sequence_region, partition_by, identity_threshold, distance_metric) -- the +# flow `immvoke annotate` exposes on the commandline: +# +# paired -> one heavy and one light chain per query cell (count_field) +# -> runPairedExactAnnotation (V, J, exact CDR3 on both chains) +# else -> heavy chain only -> sequence_region columns +# -> partition_by gate (none, V, VJ, VJ + CDR3 length) +# -> identity_threshold: 1.0 exact, below it Hamming/Levenshtein +# -> symscan; rows beyond its limits are scored directly instead +# then -> each query row's best hits (_bestHits) +# -> outdir: ANNOTATION_FILE, ANNOTATION_METADATA and, with +# full_dist_matrix, FULL_MATRIX_FILE, each named after the +# reference (e.g. iedb_bcr_annotation.tsv) + +#: sequence_region -> AIRR column(s) to compare. 'CDRs' pools CDR1-3, so it +#: only runs under Hamming distance (or an exact match), never Levenshtein: +#: edit distance summed over independently aligned columns is not the edit +#: distance of any one alignment, the way Hamming distance pools cleanly. +SEQUENCE_REGION_COLUMNS = { + 'cdr3h': ('cdr3_aa',), + 'vdj': ('sequence_aa',), + 'cdrs': ('cdr1_aa', 'cdr2_aa', 'cdr3_aa'), +} + +#: partition_by -> gene call columns two rows must agree on (same gene, any +#: allele; see Genes.py) before their sequence counts as a hit. 'VJL' also +#: gates on exact CDR3 length; see annotateSpecificity. +PARTITION_GENE_COLUMNS = { + 'v': ('v_call',), + 'vj': ('v_call', 'j_call'), + 'vjl': ('v_call', 'j_call'), +} + +#: distance_metric values accepted ('levenstein' is a common misspelling). +DISTANCE_METRICS = {'hamming': 'hamming', 'levenshtein': 'levenshtein', + 'levenstein': 'levenshtein'} + +#: Synthetic gate column added to both tables under partition_by='VJL'; it +#: only feeds the exact-column gate and is never written to the result. +_CDR3_LENGTH_COL = '__cdr3_length__' + + +def _resolveSequenceRegion(sequence_region): + """ + Arguments: + sequence_region (str): 'CDR3H', 'VDJ' or 'CDRs', case insensitive. + + Returns: + list: the AIRR column(s) to compare. + + Raises: + AnnotationError: if `sequence_region` is not recognized. + """ + columns = SEQUENCE_REGION_COLUMNS.get(str(sequence_region).strip().lower()) + if columns is None: + raise AnnotationError("unknown sequence_region %r; choose one of: CDR3H, " + "VDJ, CDRs" % (sequence_region,)) + + return list(columns) + + +def _resolvePartitionBy(partition_by): + """ + Arguments: + partition_by (str): None, 'V', 'VJ' or 'VJL', case insensitive. + + Returns: + tuple: (gate columns (list), whether to also gate on CDR3 length). + + Raises: + AnnotationError: if `partition_by` is not recognized. + """ + if partition_by is None: + return [], False + + key = str(partition_by).strip().lower() + columns = PARTITION_GENE_COLUMNS.get(key) + if columns is None: + raise AnnotationError("unknown partition_by %r; choose one of: None, V, " + "VJ, VJL" % (partition_by,)) + + return list(columns), key == 'vjl' + + +def _resolveDistanceMetric(distance_metric): + """ + Arguments: + distance_metric (str): 'hamming' or 'levenshtein', case insensitive. + + Returns: + str: 'hamming' or 'levenshtein'. + + Raises: + AnnotationError: if `distance_metric` is not recognized. + """ + metric = DISTANCE_METRICS.get(str(distance_metric).strip().lower()) + if metric is None: + raise AnnotationError("unknown distance_metric %r; choose 'hamming' or " + "'levenshtein'" % (distance_metric,)) + + return metric + + +def _readAirrTsv(table): + """ + Read an AIRR rearrangement TSV as plain text columns. + + Arguments: + table (str or pandas.DataFrame): the file to read, or an already + loaded table (returned as is). + + Returns: + pandas.DataFrame: every column as a string, blank cells as '' rather + than NaN -- the row comparisons here expect a value to always be a + string. + """ + if isinstance(table, pandas.DataFrame): + return table + + return pandas.read_csv(table, sep='\t', dtype=str, na_filter=False) + + +def _requireColumns(frame, columns, label, why): + """ + Raise a readable error, instead of a KeyError, for a missing column. + + Arguments: + frame (pandas.DataFrame): the table to check. + columns (list): the column names it needs. + label (str): 'query' or 'reference', for the message. + why (str): what needs them, for the message. + + Raises: + AnnotationError: if any of `columns` is missing from `frame`. + """ + missing = [c for c in columns if c not in frame.columns] + if missing: + raise AnnotationError('%s has no %s column, needed by %s' + % (label, ', '.join(repr(c) for c in missing), why)) + + +def _restrictToHeavyChain(frame, label): + """ + Keep only the heavy chain (locus IGH), if `frame` has a locus column; a + table without one is kept whole. + + Arguments: + frame (pandas.DataFrame): the table to filter. + label (str): 'query' or 'reference', for messages. + + Returns: + pandas.DataFrame: the filtered table, index reset to 0..n-1. + + Raises: + AnnotationError: if no row is left -- rather than silently annotating + nothing. + """ + if LOCUS_COL in frame.columns: + filtered = frame[frame[LOCUS_COL] == HEAVY_LOCUS] + log.info(' %s: %s of %s rows are heavy chains (locus %s)', label, + format(len(filtered), ','), format(len(frame), ','), HEAVY_LOCUS) + if not len(filtered): + raise AnnotationError( + '%s has no heavy chain rows (locus %s) to annotate; its loci ' + 'are: %s' % (label, HEAVY_LOCUS, + ', '.join(sorted(set(frame[LOCUS_COL]))) or 'none')) + frame = filtered + + return frame.reset_index(drop=True) + + +def _withCellId(frame, label): + """ + Drop the rows with a blank cell_id: nothing links such a chain to a + partner, so paired annotation cannot pair it -- and must not lump every + such row together as one antibody. + + Arguments: + frame (pandas.DataFrame): the table, with a cell_id column. + label (str): 'query' or 'reference', for the log line only. + + Returns: + pandas.DataFrame: `frame` without those rows. + """ + blank = frame[CELL_ID_COL].astype(object).fillna('').astype(str).str.strip() == '' + if blank.any(): + log.warning('%s: %d rows have no cell_id to pair them by; left out of ' + 'paired annotation', label, int(blank.sum())) + + return frame[~blank] + + +def _oneChainPerCell(frame, count_field, label): + """ + Keep one heavy chain (IGH) and one light chain (IGK or IGL) per cell_id: + the one with the highest `count_field`. On a tie, a blank count, or no + such column at all, the row that comes first in the file is kept. Rows + of other loci are kept as they are. + + Arguments: + frame (pandas.DataFrame): the table, with cell_id and locus columns. + count_field (str): the AIRR count column to rank chains by, e.g. + 'duplicate_count' or 'umi_count'. + label (str): 'query' or 'reference', for the log line only. + + Returns: + pandas.DataFrame: `frame` without the extra chains, rows in their + original order. + """ + frame = frame.reset_index(drop=True) + if count_field in frame.columns: + counts = pandas.to_numeric(frame[count_field], errors='coerce') + # A blank count ranks below any number. + counts = counts.fillna(float('-inf')) + else: + counts = pandas.Series(0, index=frame.index) + + drop = [] + report = {} + for name, chain in (('heavy', frame[LOCUS_COL] == HEAVY_LOCUS), + ('light', frame[LOCUS_COL].isin(LIGHT_LOCI))): + rows = frame.index[chain] + cells = frame.loc[rows, CELL_ID_COL] + several = cells[cells.duplicated(keep=False)] + if several.empty: + continue + # Highest count first; a stable sort keeps file order among ties. + ranked = counts[several.index].sort_values(ascending=False, kind='stable') + best = frame.loc[ranked.index, CELL_ID_COL].drop_duplicates() + drop.extend(several.index.difference(best.index)) + report[name] = several.nunique() + + if not drop: + return frame + if count_field in frame.columns: + rule = 'the one with the highest %s (the first row on a tie)' % count_field + else: + rule = 'the first row (no %s column to rank them by)' % count_field + log.info(' %s: %s; kept %s', label, + ', '.join('%s cells had several %s chains' % (format(n, ','), name) + for name, n in report.items()), rule) + + return frame.drop(index=drop) + + +def _requireAirr(frame, label): + """ + Require `frame` to be an AIRR rearrangement table: every field the AIRR + schema marks required present as a column (its values may be blank). + + Arguments: + frame (pandas.DataFrame): the table to check. + label (str): 'query' or 'reference', for the message. + + Raises: + AnnotationError: naming the missing required fields. + """ + from airr.schema import RearrangementSchema, ValidationError + + try: + RearrangementSchema.validate_header(list(frame.columns)) + except ValidationError as error: + raise AnnotationError('%s is not an AIRR rearrangement table: %s' + % (label, error)) from None + + +def _runSingleChain(queries, reference, compare_cols, gene_cols, exact_cols, + threshold, metric, threads): + """ + Run one single-chain comparison through symscan, keeping every hit. + + Arguments: + queries (pandas.DataFrame): the query rows, already filtered. + reference (pandas.DataFrame): the reference rows, already filtered. + compare_cols (list): the column(s) to compare. + gene_cols (list): gene call columns that must name the same gene, + possibly empty. + exact_cols (list): columns that must agree exactly, possibly empty. + threshold (float): identity in (0.0, 1.0]; 1.0 is an exact match. + metric (str): 'hamming' or 'levenshtein'. + threads (int): symscan's thread pool; see annotateSpecificity. + + Returns: + tuple: (every hit within the threshold, one row per (query row, hit + id), closest first within a query row -- see + Symscan.runSymscanAnnotation's `return_pairs`; the query rows beyond + symscan's limits, scored directly instead). + """ + tolerance, options = _compareOptions(gene_cols, exact_cols, threshold, metric) + + # Deferred: Symscan builds on this module's helpers. + from immvoke.Sources.Specificity.Symscan import runSymscanAnnotation + + result, pairs = runSymscanAnnotation(queries, reference, compare_cols, tolerance, + threads=threads, return_pairs=True, + warn_skipped=False, **options) + skipped = result.attrs['skipped_rows'] + found = _scoreSkippedRows(skipped, queries, reference, compare_cols, tolerance, + options) + if len(found): + pairs = (pandas.concat([pairs, found[pairs.columns]]) + .sort_values(['query_row', 'mismatches'], kind='stable') + .drop_duplicates(['query_row', 'reference_id']) + .reset_index(drop=True)) + + return pairs, skipped + + +def _compareOptions(gene_cols, exact_cols, threshold, metric): + """ + Turn the run's options into the arguments the comparison takes, so the + thresholded search and the full matrix are gated identically. + + Arguments: + gene_cols (list): gene call columns that must agree, possibly empty. + exact_cols (list): columns that must agree exactly, possibly empty. + threshold (float): identity in (0.0, 1.0]; 1.0 is an exact match. + metric (str): 'hamming' or 'levenshtein'. + + Returns: + tuple: (tolerance, keyword arguments). + """ + exact = threshold >= 1.0 + + # An exact match is 0 substitutions under either metric -- the same + # thing, and the only way to run exact agreement on pooled 'CDRs'. + return (0 if exact else 1 - threshold, + dict(distance_type='hamming' if exact else metric, + threshold_is_fraction=not exact, id_col=DEFAULT_ID_COL, + exact_cols=exact_cols or None, gene_cols=gene_cols or None)) + + +def _fullMatrixBlocks(queries, reference, compare_cols, tolerance, options): + """ + Every (query, reference) pair the gates admit, with its identity, + whatever its distance -- the --full-dist-matrix table. + + The annotation table answers "what did this query match"; the matrix + answers "how close is this query to everything it was compared with", + so --identity-threshold does not apply to it. The gates still do: a + pair of different V or J gene was never a candidate (--partition-by), + and under 'hamming' only equal-length sequences have a distance at all, + so unequal ones are left out rather than given a blank. + + Yielded in blocks rather than one table: without --partition-by this is + every pair of both tables, which for a repertoire against IEDB is + hundreds of millions of rows -- far more than fits in memory, though + each block streams to disk fine. + + Arguments: + queries (pandas.DataFrame): the query rows that were compared. + reference (pandas.DataFrame): the reference rows. + compare_cols (list): the column(s) compared. + tolerance (float): unused for selection; `allPairDistances` needs it + to mark `within_threshold`, which this drops. + options (dict): the rest of the comparison's keyword arguments. + + Yields: + pandas.DataFrame: `query_row`, `hit_id`, `identity`, in query order, + closest first within a query. + """ + from immvoke.Sources.Specificity.Distances import allPairDistances + + for block in allPairDistances(queries, reference, compare_cols, tolerance, + **options): + yield (block[['query_row', 'reference_id', 'identity']] + .rename(columns={'reference_id': 'hit_id'})) + + +def _scoreSkippedRows(rows, queries, reference, compare_cols, tolerance, options): + """ + Find the hits of the query rows beyond symscan's limits (see Symscan.py) + by scoring their pairs directly (Distances.allPairDistances), so no row + misses a hit -- slower, but only for those rows. + + Arguments: + rows (list): those query rows, runSymscanAnnotation's + `attrs['skipped_rows']`. + queries, reference, compare_cols, tolerance: as given to symscan. + options (dict): the rest of symscan's keyword arguments. + + Returns: + pandas.DataFrame: their hits as pairs, see + Symscan.runSymscanAnnotation's `return_pairs`; empty when there are + none. + """ + rows = numpy.array(rows, dtype=numpy.int64) + if not len(rows): + return pandas.DataFrame() + log.info(" %d query rows are beyond symscan's limits (see Symscan.py); scoring " + 'their pairs directly', len(rows)) + + # Deferred: Distances builds on this module's helpers. + from immvoke.Sources.Specificity.Distances import allPairDistances + + subset = queries.iloc[rows].reset_index(drop=True) + pairs = pandas.concat(list(allPairDistances(subset, reference, compare_cols, + tolerance, **options))) + found = pairs[pairs['within_threshold']].drop(columns='within_threshold') + found = found.assign(query_row=rows[found['query_row'].to_numpy(dtype=int)]) + + return found.reset_index(drop=True) + + +def _hitRows(reference): + """ + What each reference id says for HIT_ANNOTATIONS, one entry per + reference row with that id -- IEDB lists an antibody once per epitope + it was tested against, and each of those rows becomes a row of the + annotation table. + + Arguments: + reference (pandas.DataFrame): the reference rows that were compared. + + Returns: + dict: reference id -> list of value tuples, one per row with that id, + in reference order, duplicates dropped; a column the reference lacks + reads as ''. An id absent from the reference is not a key. + """ + ids = reference[DEFAULT_ID_COL].astype(str).tolist() + columns = [reference[column].astype(object).fillna('').astype(str).tolist() + if column in reference.columns else [''] * len(reference) + for _, column in HIT_ANNOTATIONS] + + rows = {} + for position, hit in enumerate(ids): + # dict.fromkeys: an id whose rows repeat the same values (one + # epitope recorded twice) should not repeat them in the output. + rows.setdefault(hit, {})[tuple(c[position] for c in columns)] = None + + return {hit: list(values) for hit, values in rows.items()} + + +def _bestHits(pairs, reference): + """ + Each query row's best hits -- the ones with the highest identity, all + of them on a tie -- one row per matching reference row. + + Long rather than wide: a query with three hits gets three rows, not one + row holding three joined values. Nothing has to be separated, so + nothing can be mis-separated, which matters because IEDB's own epitope + names contain both the ';' and the '|' a joined cell would need. + + Arguments: + pairs (pandas.DataFrame): every hit, from _runSingleChain. + reference (pandas.DataFrame): the reference rows that were compared, + for HIT_ANNOTATIONS. + + Returns: + pandas.DataFrame: `query_row`, `identity` (to 4 decimals), `hit_id` + and HIT_ANNOTATIONS' columns, one row per (query row, reference row + of one of its best hits), by query row then identity. A query row + without a hit has no row at all. + """ + names = [name for name, _ in HIT_ANNOTATIONS] + columns = ['query_row', 'identity', 'hit_id'] + names + if not len(pairs): + return pandas.DataFrame({column: [] for column in columns}) + + # One query row's identities share its length, so the fewest mismatches + # is the highest identity, without comparing floats. + fewest = pairs.groupby('query_row')['mismatches'].transform('min') + best = pairs[pairs['mismatches'] == fewest].sort_values( + ['query_row', 'reference_id'], kind='stable') + rows = _hitRows(reference) + + blank = ('',) * len(names) + records = [(query, identity, hit) + values + for query, identity, hit in zip(best['query_row'], + best['identity'].round(4), + best['reference_id']) + for values in rows.get(hit, [blank])] + + return pandas.DataFrame(records, columns=columns) + + +def _withQueryIds(table, queries): + """ + Name each query by its own `sequence_id` instead of its position, so + the output joins straight back to the query file. A query without a + `sequence_id` column keeps `query_row`. + + Arguments: + table (pandas.DataFrame): a result or pairs table, with `query_row` + indexing `queries`. + queries (pandas.DataFrame): the query rows that were compared, in + `query_row` order. + + Returns: + pandas.DataFrame: `table` with `sequence_id` in place of `query_row`. + """ + if DEFAULT_ID_COL not in queries.columns: + return table + ids = queries[DEFAULT_ID_COL].astype(object).fillna('').astype(str).to_numpy() + named = table.drop(columns='query_row') + named.insert(0, DEFAULT_ID_COL, ids[table['query_row'].to_numpy(dtype=int)]) + + return named + + +def _tableRecord(table, frame): + """ + Describe an input table for the annotation metadata. + + Arguments: + table (str or pandas.DataFrame): the table as the caller passed it. + frame (pandas.DataFrame): the table as read. + + Returns: + dict: its `path` (None for a DataFrame) and `rows`. + """ + path = None if isinstance(table, pandas.DataFrame) else str(table) + + return {'path': path, 'rows': len(frame)} + + +def _referenceSource(reference): + """ + Read which database a reference table came from, off the + download_metadata.yml `immvoke specificity` writes in the database's + folder, beside the table. + + Arguments: + reference (str or pandas.DataFrame): the reference as the caller + passed it. + + Returns: + dict: `source` (e.g. 'iedb') and `downloaded` (when), or empty when + the reference is a DataFrame, or no readable record sits beside it. + """ + if isinstance(reference, pandas.DataFrame): + return {} + path = Path(reference).parent / DOWNLOAD_METADATA + try: + record = yaml.safe_load(path.read_text()) + except (OSError, yaml.YAMLError): + return {} + if not isinstance(record, dict) or 'source' not in record: + return {} + + runs = record.get('runs') or [{}] + return {'source': record['source'], 'downloaded': runs[-1].get('finished')} + + +def _outputPrefix(reference, source): + """ + Name the reference for the output files, so runs against different + databases stay apart even in one outdir. + + Arguments: + reference (str or pandas.DataFrame): the reference as the caller + passed it. + source (str): its database, from _referenceSource, or None. + + Returns: + str: '_
' (e.g. 'iedb_bcr') for a table `immvoke + specificity` downloaded, read off its file name (iedb_bcr.tsv), with + the database added when the name lacks it (an older bcr.tsv); the file + name alone for any other reference file; 'reference' for a DataFrame. + """ + if isinstance(reference, pandas.DataFrame): + return 'reference' + table = Path(reference).name.split('.')[0] + if not source or table.startswith(source + '_'): + return table + + return '%s_%s' % (source, table) + + +def _summaryText(prefix, metadata, files): + """ + The run in prose: what was compared, what came of it, and what each + file beside this one holds -- so the folder explains itself without + reading the YAML or this module. + + Arguments: + prefix (str): the reference's name, from _outputPrefix. + metadata (dict): the run record, see annotateSpecificity. + files (list): the file names written, in order. + + Returns: + str: the summary, ending in a newline. + """ + options, results = metadata['options'], metadata['results'] + paired = options['paired'] + threshold = options['identity_threshold'] + + def named(record, fallback): + return record['path'] or fallback + + lines = ['Annotation of %s' % named(metadata['query'], 'a table passed from Python'), + ' against %s' % named(metadata['reference'], + 'a reference passed from Python'), + 'Finished: %s' % metadata['finished'], ''] + + lines.append('What was compared') + if paired: + lines.append(' heavy and light chain of the same antibody, matching on ' + 'V gene, J gene') + lines.append(' and an exact CDR3.') + else: + lines.append(' %s of each query row, against the same of every reference row.' + % {'CDR3H': 'the heavy chain CDR3 (cdr3_aa)', + 'VDJ': 'the V(D)J amino acid sequence (sequence_aa)', + 'CDRs': 'CDR1, CDR2 and CDR3 pooled'} + .get(options['sequence_region'], options['sequence_region'])) + lines.append(' A hit is %s.' + % ('an exact match' if threshold is None else + '%g%% identity or better, by %s distance' + % (threshold * 100, options['distance_metric']))) + lines.append(' %s' % {None: 'Any pair of rows was compared.', + 'V': 'Only rows of the same V gene were compared.', + 'VJ': 'Only rows of the same V and J gene were compared.', + 'VJL': 'Only rows of the same V gene, J gene and CDR3 ' + 'length were compared.'} + .get(options['partition_by'], + 'Partitioned by %s.' % options['partition_by'])) + lines.append('') + + lines.append('What was found') + figures = ([('query antibodies', results['query_antibodies'])] + + [(label, count) for label, count + in results.get('classification', {}).items()] + if paired else + [('query rows compared', results['query_rows_compared']), + ('query rows with a hit', results['query_rows_with_hits']), + ('rows in the annotation', results['annotation_rows'])]) + if 'full_matrix_rows' in results: + figures.append(('pairs in the full matrix', results['full_matrix_rows'])) + pad = max(len(label) for label, _ in figures) + 2 + lines += [' %s%s' % (label.ljust(pad), format(count, ',')) + for label, count in figures] + lines.append('') + + lines.append('The files') + what = { + ANNOTATION_FILE: ('one row per (query antibody)' if paired else + 'one row per (query row, matching reference row): a query ' + 'row\n with several hits has a row for each, so no cell ' + 'holds two values.\n Only the best hits are here -- the ' + 'highest identity, all of them on a\n tie -- and a query ' + 'row that matched nothing has no row at all'), + FULL_MATRIX_FILE: ('every pair the filters admit and how close it is, whether ' + 'or not\n it is a hit: the identity threshold does not ' + 'apply here, the\n partition does'), + ANNOTATION_METADATA: ('the whole record of the run: every option, the input ' + 'tables and\n the reference database, the counts ' + 'above, and the immvoke version,\n time and command'), + } + for name in files: + lines.append(' %s' % name) + lines.append(' %s.' % what[name[len(prefix) + 1:]]) + + return '\n'.join(lines) + '\n' + + +def _writeOutputs(outdir, prefix, result, matrix, metadata): + """ + Write the annotation bundle into `outdir`: ANNOTATION_FILE, + ANNOTATION_METADATA and, when there is a matrix, FULL_MATRIX_FILE, each + named `_`. + + A run writes the whole bundle, so a matrix an earlier run left in the + same folder is removed when this one was not asked for a matrix. + Otherwise it would stay beside a table and a metadata record that have + been overwritten, describing a comparison that no longer happened -- + a stale 0.7 matrix next to an exact-match table, say. + + Arguments: + outdir (str): the folder to write into, or None to skip writing. + prefix (str): the reference's name, from _outputPrefix. + result (pandas.DataFrame): the annotation result. + matrix (iterable): the full matrix as DataFrames to write one after + another -- it can outgrow memory, so it is streamed rather than + built whole (see _fullMatrixBlocks) -- or None when + full_dist_matrix was not asked for. + metadata (dict): the run record, see annotateSpecificity. + """ + if outdir is None: + return + + # Rendered once up front, so a value YAML cannot represent fails before + # any file is written rather than leaving a partial bundle behind. + yaml.safe_dump(metadata) + out = Path(outdir) + out.mkdir(parents=True, exist_ok=True) + names = ['%s_%s' % (prefix, name) for name in (ANNOTATION_FILE, ANNOTATION_METADATA)] + result.to_csv(out / names[0], sep='\t', index=False) + stale = out / ('%s_%s' % (prefix, FULL_MATRIX_FILE)) + if matrix is not None: + names.append(stale.name) + rows = 0 + with open(stale, 'w', newline='') as handle: + for block in matrix: + # Identity to 4 decimals: 0.4231, not 0.42307692307692313. + block.round({'identity': 4}).to_csv(handle, sep='\t', index=False, + header=not rows) + rows += len(block) + metadata['results']['full_matrix_rows'] = rows + elif stale.is_file(): + stale.unlink() + log.info(' removed %s, left by an earlier run that asked for it', + stale.name) + metadata['finished'] = timestamp() + (out / names[1]).write_text(yaml.safe_dump(metadata, sort_keys=False)) + summary = '%s_%s' % (prefix, ANNOTATION_SUMMARY) + (out / summary).write_text(_summaryText(prefix, metadata, names)) + names.append(summary) + log.info(' wrote %s', ', '.join(names)) + log.info(' into %s', out) + + +def annotateSpecificity(query, reference, sequence_region='CDR3H', + partition_by=None, paired=False, identity_threshold=None, + distance_metric='hamming', threads=1, + outdir=None, count_field=DEFAULT_COUNT_FIELD, + full_dist_matrix=False): + """ + Annotate a query receptor table against a reference specificity table. + + Arguments: + query (str or pandas.DataFrame): the query AIRR rearrangement TSV, or + the table itself. + reference (str or pandas.DataFrame): the reference AIRR rearrangement + TSV -- e.g. a specificity database's downloaded table -- or the + table itself. It must be an AIRR rearrangement table: every field + the AIRR schema requires present as a column (values may be blank), + as immvoke's own AIRR writer always produces. + sequence_region (str): which sequence to compare: 'CDR3H' (the heavy + chain CDR3, cdr3_aa), + 'VDJ' (sequence_aa, the full V(D)J amino acid sequence) or 'CDRs' + (cdr1_aa, cdr2_aa and cdr3_aa pooled -- Hamming only). Ignored when + `paired` is True. + partition_by (str): only compare row pairs in the same partition: None + (no restriction), 'V' (same V-gene), 'VJ' (same V- and J-gene) or + 'VJL' (same V-gene, J-gene and CDR3 length -- needs `cdr3_aa` in both + tables whatever `sequence_region` is). Ignored when `paired` is True. + paired (bool): classify each query antibody (`cell_id`) by V-gene, + J-gene and exact CDR3 agreement on both its heavy and light chain + with the same reference antibody (runPairedExactAnnotation); rows + with a blank `cell_id` are left out, with a warning. When False, + only heavy chain (locus IGH) rows are compared; a table without a + locus column is compared whole. + count_field (str): under `paired`, a query cell with several heavy + (or light) chains keeps only the one with the highest value in this + AIRR count column -- 'duplicate_count' by default, or e.g. + 'umi_count' -- or the first in the file on a tie, a blank count, or + no such column. + identity_threshold (float): minimum sequence identity in (0.0, 1.0]; + None or 1.0 requires an exact match, a lower value allows up to + floor((1 - identity_threshold) * query length) differences. Ignored + when `paired` is True. + distance_metric (str): 'hamming' (substitutions only) or 'levenshtein' + (substitutions, insertions and deletions; a single-column + `sequence_region` only, i.e. not 'CDRs'). Only matters when + `identity_threshold` is below 1.0. + threads (int): symscan's thread pool size (-1 for every core); only + the first symscan run in a process sets it. + outdir (str): if given, also write into this folder, each file named + after the reference's database and table (e.g. + iedb_bcr_annotation.tsv; see _outputPrefix): ANNOTATION_FILE (the + result) and ANNOTATION_METADATA (the options, the query and + reference with the reference's source database, the counts, and + the immvoke version, time and command). + full_dist_matrix (bool): also write FULL_MATRIX_FILE into `outdir`: + every pair the gates admit and how close it is, one row per + (query, hit) -- `sequence_id`, `hit_id`, `identity`. + `identity_threshold` does not apply to it (see _fullMatrixBlocks), + `partition_by` does; under `paired`, `cell_id`, `hit_id` and + `matched_chains` (`both`, `heavy` or `light`). + + Returns: + pandas.DataFrame: when `paired`, one row per query `cell_id` (see + runPairedExactAnnotation); otherwise one row per (query heavy chain + row, matching reference row), a query row without a hit left out, + with `sequence_id` (the query's own; `query_row`, its position among + the heavy chain rows, for a query without that column), then the hit, + one of the best, the ones + with the highest identity (all of them on a tie): `identity`, + `hit_id`, and that hit's `antigen`, `epitope` and + `antigen_source_species` from the reference (HIT_ANNOTATIONS). A hit + standing for several reference rows -- IEDB lists a receptor once per + epitope -- gets a row for each, so no cell holds two values. + + Raises: + AnnotationError: if the reference is not an AIRR rearrangement table, + an argument is not recognized, a needed column is missing, no heavy + chain row is left to annotate, an argument combination cannot run + (e.g. 'CDRs' under 'levenshtein'). Single-chain rows beyond + symscan's limits are scored directly instead (_scoreSkippedRows); + under `paired`, such rows (over 14 X's) are skipped with a warning. + """ + queries = _readAirrTsv(query) + ref = _readAirrTsv(reference) + _requireAirr(ref, 'reference') + source = _referenceSource(reference) + prefix = _outputPrefix(reference, source.get('source')) + + metadata = { + 'immvoke_metadata_version': METADATA_VERSION, + 'generated_by': 'immvoke %s' % __version__, + 'finished': None, + 'query': _tableRecord(query, queries), + 'reference': dict(_tableRecord(reference, ref), **source), + # .item(): a numpy scalar (e.g. a threshold from numpy.arange) as the + # plain Python value YAML can write. + 'options': {name: value.item() if hasattr(value, 'item') else value + for name, value in (('sequence_region', sequence_region), + ('partition_by', partition_by), + ('paired', paired), + ('identity_threshold', identity_threshold), + ('distance_metric', distance_metric), + ('count_field', count_field), + ('full_dist_matrix', full_dist_matrix))}, + } + # Run through the commandline rather than called from Python, where + # sys.argv is some other program's. + if Path(sys.argv[0]).stem == 'immvoke': + metadata['command'] = commandLine() + + started = time.monotonic() + Report.annotationStarted(metadata['query'], metadata['reference'], + metadata['options']) + + if paired: + if partition_by is not None or identity_threshold is not None: + log.warning('paired annotation always compares V-gene, J-gene and ' + 'exact CDR3; partition_by and identity_threshold are ' + 'ignored') + paired_cols = [CELL_ID_COL, LOCUS_COL, 'cdr3_aa', 'v_call', 'j_call'] + _requireColumns(queries, paired_cols, 'query', 'paired annotation') + _requireColumns(ref, paired_cols, 'reference', 'paired annotation') + queries = _withCellId(queries, 'query') + ref = _withCellId(ref, 'reference') + chain_loci = (HEAVY_LOCUS,) + LIGHT_LOCI + for frame, label in ((queries, 'query'), (ref, 'reference')): + if not frame[LOCUS_COL].isin(chain_loci).any(): + raise AnnotationError('%s has no antibody chain rows (locus %s) ' + 'to pair' % (label, ', '.join(chain_loci))) + queries = _oneChainPerCell(queries, count_field, 'query') + result, pairs = runPairedExactAnnotation(queries, ref, threads=threads, + return_pairs=True) + classes = {label: int(count) for label, count + in result['classification'].value_counts().items()} + metadata['results'] = { + 'query_antibodies': len(result), 'classification': classes, + 'skipped_query_rows': result.attrs['skipped_query_rows'], + 'pairs': len(pairs)} + # One block: a paired run's matrix is one row per matched antibody + # pair, which cannot outgrow memory the way a distance matrix can. + _writeOutputs(outdir, prefix, result, [pairs] if full_dist_matrix else None, + metadata) + Report.annotationCompleted(time.monotonic() - started, + metadata['results'], paired=True) + + return result + + threshold = 1.0 if identity_threshold is None else float(identity_threshold) + if not 0 < threshold <= 1.0: + raise AnnotationError('identity_threshold must be in (0.0, 1.0]; got %s' + % (identity_threshold,)) + compare_cols = _resolveSequenceRegion(sequence_region) + gene_cols, gate_on_cdr3_length = _resolvePartitionBy(partition_by) + metric = _resolveDistanceMetric(distance_metric) + if threshold < 1.0 and metric == 'levenshtein' and len(compare_cols) != 1: + raise AnnotationError("distance_metric='levenshtein' compares a single " + "column; sequence_region=%r pools %d, which only " + "'hamming' supports" + % (sequence_region, len(compare_cols))) + + needed = compare_cols + gene_cols + (['cdr3_aa'] if gate_on_cdr3_length else []) + why = 'sequence_region=%r, partition_by=%r' % (sequence_region, partition_by) + _requireColumns(queries, needed, 'query', why) + _requireColumns(ref, needed + [DEFAULT_ID_COL], 'reference', why) + + queries = _restrictToHeavyChain(queries, 'query') + ref = _restrictToHeavyChain(ref, 'reference') + exact_cols = [] + if gate_on_cdr3_length: + queries = queries.assign( + **{_CDR3_LENGTH_COL: queries['cdr3_aa'].str.len().astype(str)}) + ref = ref.assign(**{_CDR3_LENGTH_COL: ref['cdr3_aa'].str.len().astype(str)}) + exact_cols = [_CDR3_LENGTH_COL] + + pairs, skipped = _runSingleChain(queries, ref, compare_cols, gene_cols, + exact_cols, threshold, metric, threads) + # Long: one row per (query row, matching reference row), so a query + # row with no hit simply has none, as on a real repertoire almost every + # row would otherwise be an empty one. + result = _bestHits(pairs, ref) + metadata['results'] = { + 'query_rows_compared': len(queries), + 'query_rows_with_hits': int(result['query_row'].nunique()), + 'annotation_rows': len(result), + 'query_rows_scored_directly': len(skipped), 'hits': len(pairs)} + result = _withQueryIds(result, queries) + matrix = None + if full_dist_matrix: + if not gene_cols: + log.warning('--full-dist-matrix without --partition-by scores every ' + 'pair of both tables, which on a repertoire is tens of ' + 'millions of rows; --partition-by V, VJ or VJL narrows it') + tolerance, options = _compareOptions(gene_cols, exact_cols, threshold, metric) + matrix = (_withQueryIds(block, queries) for block in + _fullMatrixBlocks(queries, ref, compare_cols, tolerance, options)) + _writeOutputs(outdir, prefix, result, matrix, metadata) + Report.annotationCompleted(time.monotonic() - started, metadata['results'], + paired=False) + + return result diff --git a/src/immvoke/Sources/Specificity/AnnotateCli.py b/src/immvoke/Sources/Specificity/AnnotateCli.py new file mode 100644 index 0000000..cd12ffe --- /dev/null +++ b/src/immvoke/Sources/Specificity/AnnotateCli.py @@ -0,0 +1,158 @@ +""" +`immvoke annotate` -- the commandline front end for +Annotate.annotateSpecificity() + +Kept apart from immvoke.Cli (which only imports addAnnotateParser/ +handleAnnotate and wires them into the top level parser and dispatcher) so +that the annotation step -- argument definitions, help text, and the handler +itself -- lives alongside the rest of the specificity annotation code it +operates on. Every flag maps to one annotateSpecificity() parameter, and +handleAnnotate only forwards them. + +Reads two AIRR rearrangement TSVs already on disk -- typically a user's own +receptor table as `--query` and a specificity database's downloaded table +(e.g. `immvoke specificity iedb download bcr`'s output) as `--reference` -- +and writes the annotation (each query row's best hits) and a record of the +run into `--outdir`, plus every pair the filters admit and how close it +is with `--full-dist-matrix`. It does not itself download or convert anything; that +is handled by `immvoke specificity iedb download bcr`. +""" + +# Info +__author__ = 'Pramod Shinde' + +# Imports +from pathlib import Path + +# Immvoke imports +from immvoke.Commandline import CommonHelpFormatter +from immvoke.Sources.Specificity.Annotate import ( + ANNOTATION_FILE, + ANNOTATION_METADATA, + DEFAULT_COUNT_FIELD, + FULL_MATRIX_FILE, + annotateSpecificity, +) + + +def addAnnotateParser(commands): + """ + Add the `annotate` subcommand to the top level parser. + + Arguments: + commands (argparse._SubParsersAction): the top level subparsers + object, as built by `immvoke.Cli.getArgParser`. + """ + parser = commands.add_parser( + 'annotate', help='annotate a receptor table against a reference table', + description='Annotate each row of a query AIRR rearrangement table ' + 'with antigen specificity information from a fixed ' + 'reference AIRR rearrangement table, e.g. `immvoke ' + 'specificity iedb download bcr`\'s output.', + formatter_class=CommonHelpFormatter) + + parser.add_argument('--query', type=Path, required=True, + help='AIRR rearrangement TSV to annotate, one row ' + 'per chain/sequence') + parser.add_argument('--reference', type=Path, required=True, + help='AIRR rearrangement TSV to annotate against, ' + 'e.g. a specificity database\'s downloaded ' + 'table; its sequence_id is reported as each ' + 'hit\'s identifier (cell_id under --paired)') + parser.add_argument('--sequence-region', dest='sequence_region', + choices=['CDR3H', 'VDJ', 'CDRs'], default='CDR3H', + help='which sequence to compare: CDR3H (the heavy ' + 'chain CDR3, cdr3_aa), VDJ ' + '(sequence_aa, the full V(D)J amino acid ' + 'sequence) or CDRs (cdr1_aa, cdr2_aa and cdr3_aa ' + 'pooled -- --distance-metric hamming only)') + parser.add_argument('--partition-by', dest='partition_by', + choices=['V', 'VJ', 'VJL'], default=None, + help='only compare row pairs with the same V-gene ' + '(V), the same V- and J-gene (VJ), or the same ' + 'V-gene, J-gene and CDR3 length (VJL); unset, ' + 'every row pair is compared. Genes match by ' + 'name, any allele (IGHV1-2*02 = IGHV1-2*04); a ' + 'cell with several calls agrees if any one does, ' + 'a blank or unrecognized call never does') + parser.add_argument('--paired', action='store_true', + help='require V-gene and J-gene agreement (same ' + 'gene, any allele) and exact CDR3 agreement on both ' + 'the heavy and the light chain ' + 'of the same reference antibody, and classify ' + 'each query cell_id as both/heavy_only/' + 'light_only/conflicting/neither. Without it, ' + 'only heavy chain (locus IGH) rows are compared. ' + '--sequence-region, --partition-by, ' + '--identity-threshold and --distance-metric ' + 'have no effect here') + parser.add_argument('--count-field', dest='count_field', + default=DEFAULT_COUNT_FIELD, metavar='COLUMN', + help='--paired only: when a query cell has several ' + 'heavy (or light) chains, keep the one with the ' + 'highest value in this AIRR count column (e.g. ' + 'umi_count); on a tie, a blank count, or no such ' + 'column, the first in the file') + parser.add_argument('--identity-threshold', dest='identity_threshold', + type=float, default=None, metavar='FRACTION', + help='minimum sequence identity in (0.0, 1.0], e.g. ' + '0.9 allows up to 10%% of the query\'s positions ' + 'to differ; unset, an exact match is required. ' + 'Under hamming (and an exact match) an X ' + 'residue matches any residue. Rows a threshold ' + 'makes too large for symscan (over 14 ' + 'differences, or a very long sequence) are ' + 'scored directly instead, more slowly') + parser.add_argument('--distance-metric', dest='distance_metric', + choices=['hamming', 'levenshtein'], default='hamming', + help='hamming (substitutions only) or levenshtein ' + '(substitutions, insertions and deletions -- ' + 'not with --sequence-region CDRs). Only matters ' + 'when --identity-threshold is below 1.0') + parser.add_argument('--threads', type=int, default=1, + help='threads for symscan\'s thread pool (-1 for ' + 'every core; only the first symscan run in a ' + 'process sets it)') + parser.add_argument('--outdir', type=Path, required=True, + help='folder to write into, created if needed, each ' + 'file named after the reference\'s database and ' + 'table (e.g. iedb_bcr_%s): _
_%s (each ' + 'query row\'s best hits -- the highest identity, ' + 'all of them on a tie -- with their antigen, ' + 'epitope and antigen source species) and ' + '_
_%s (the options, inputs and ' + 'counts of this run). A reference not downloaded ' + 'by `immvoke specificity` is named by its file ' + 'name alone' + % (ANNOTATION_FILE, ANNOTATION_FILE, ANNOTATION_METADATA)) + parser.add_argument('--full-dist-matrix', dest='full_dist_matrix', + action='store_true', + help='also write _
_%s: every pair the ' + 'filters admit and how close it is, one row per ' + '(query sequence_id, hit_id, identity). ' + '--identity-threshold does not apply to it, ' + '--partition-by does; under --paired, each ' + "antibody's hits and the chains they matched on" + % FULL_MATRIX_FILE) + + +def handleAnnotate(args): + """ + Forward parsed arguments to annotateSpecificity(). + + Arguments: + args (argparse.Namespace): parsed arguments, as built by + addAnnotateParser. + + Returns: + int: process exit status. + """ + annotateSpecificity( + query=args.query, reference=args.reference, + sequence_region=args.sequence_region, partition_by=args.partition_by, + paired=args.paired, identity_threshold=args.identity_threshold, + distance_metric=args.distance_metric, threads=args.threads, + outdir=args.outdir, count_field=args.count_field, + full_dist_matrix=args.full_dist_matrix) + + return 0 diff --git a/src/immvoke/Sources/Specificity/Asd.py b/src/immvoke/Sources/Specificity/Asd.py new file mode 100644 index 0000000..e7cd9e2 --- /dev/null +++ b/src/immvoke/Sources/Specificity/Asd.py @@ -0,0 +1,317 @@ +""" +ASD (Antigen-Specific antibody Database) - Ab-Ag pairs as a Delta Lake + +ASD (https://naturalantibody.com/agab/) publishes its whole collection as a +Delta Lake (Parquet partitions plus a _delta_log transaction log) staged in a +Google Drive folder behind a Colab notebook; there is no plain HTTP file to +GET. fetchUnit therefore reaches for gdown rather than the shared +HttpClient, and what it downloads is one folder rather than one file. + +A table is not an index into many separately downloadable files: the Delta +Lake itself *is* the data unit, so there is exactly one. ASD has no HTML +search form, so harvestSchema returns a hand curated snapshot rather than +one built from a live page; the one filterable field (dataset) is the +controlled vocabulary observed directly in the Delta Lake, since the Drive +folder offers no per-dataset export to filter server side. +""" + +# Info +__author__ = 'Pramod Shinde' + +# Imports +import hashlib +import json +import logging +from datetime import UTC +from pathlib import Path + +import pandas + +# Immvoke imports +from immvoke.Exceptions import AsdFetchError, AsdParseError +from immvoke.Http import hashFile +from immvoke.Sources.Base import DataUnit, DownloadResult, SourceBase +from immvoke.Sources.Specificity.RowHash import rowHashes + +log = logging.getLogger(__name__) + +#: Endpoints. There is no per-file download URL; the whole Delta Lake is one +#: Google Drive folder, reached through the Colab notebook linked from HOME_URL. +HOME_URL = 'https://naturalantibody.com/agab/' +DRIVE_FOLDER_URL = ('https://drive.google.com/drive/folders/' + '1HY2GOVj-HR8t6Jmhe5MIJ6xmBCOxC45q') + +#: The one table this source offers; the whole Delta Lake is the data unit. +COLLECTION = 'asd' + +#: dataset has no server-side filter on the Drive folder download, so it is +#: applied client-side after fetching; this is its known controlled +#: vocabulary, observed directly in the Delta Lake. +_DATASETS = ( + 'aae', 'aatp', 'ab-bind', 'abbd', 'abdesign', 'alphaseq', 'biomap', + 'buzz', 'covid-19', 'dlgo', 'flab_hie2022', 'flab_koenig2017', + 'flab_rosace2023', 'flab_shanehsazzadeh2023', 'flab_warszawski2019', + 'genbank', 'hiv', 'literature', 'met', 'osh', 'patents', 'rmna', + 'skempiv2', 'structures-antibodies', 'structures-nanobodies', +) + +#: Columns whose values are nested (Arrow struct/list), which pandas reads +#: back as Python dicts/lists. A TSV cell that is str()'d from a dict is not +#: valid JSON and cannot be round tripped, so these are serialized instead. +_NESTED_COLUMNS = ('metadata',) + + +def _hashDirectory(path): + """ + Build a content digest for a downloaded Delta Lake directory. + + HttpClient.fetch hashes a single completed file; a Delta Lake is a + directory of them, so this combines a hash of every file's relative path + and content into one digest, computed in a fixed (sorted) order so the + result is stable regardless of the order gdown wrote files in. + + Arguments: + path (Path): the Delta Lake directory. + + Returns: + str: hexadecimal SHA-256 digest. + """ + digest = hashlib.sha256() + for item in sorted(Path(path).rglob('*')): + if item.is_file(): + digest.update(str(item.relative_to(path)).encode('utf-8')) + digest.update(hashFile(item).encode('utf-8')) + + return digest.hexdigest() + + +def _directorySize(path): + """ + Total size in bytes of every file under a Delta Lake directory. + + Arguments: + path (Path): the Delta Lake directory. + + Returns: + int: total size in bytes. + """ + return sum(item.stat().st_size for item in Path(path).rglob('*') + if item.is_file()) + + + +class AsdSource(SourceBase): + """ + The ASD specificity source. + """ + + name = 'asd' + description = 'ASD: antigen-specific antibody and nanobody Ab-Ag pairs' + homepage = HOME_URL + collections = (COLLECTION,) + collection_help = { + COLLECTION: ('the full Ab-Ag pair Delta Lake (structures, affinity ' + 'assays, literature and patent mined pairs, …), ' + 'optionally narrowed to one dataset'), + } + + #: TODO: confirm license wording against https://naturalantibody.com/agab/ + #: before this is relied on for redistribution terms. + license = ('CC BY-NC 4.0; see the LICENSE.txt distributed alongside the ' + 'Delta Lake in the Drive folder') + citation = ( + 'naturalantibody.com. ASD: Antigen-Specific antibody Database. ' + 'https://naturalantibody.com/agab/', + ) + + def harvestSchema(self): + """ + Build the hand curated snapshot. + + There is no live form to scrape, so this returns the same field + definitions packaged in schema.yaml, timestamped now. `immvoke + schema refresh --source asd` exists mainly so the timestamp and any + future edits to this method stay reproducible the same way a real + harvest would. + + Returns: + SourceSchema: the snapshot. + """ + from datetime import datetime + + from immvoke.Schema import Collection, Field, SourceSchema + from immvoke.Version import __version__ + + collections = { + COLLECTION: Collection(name=COLLECTION, fields=( + Field(name='dataset', values=_DATASETS), + )), + } + + return SourceSchema( + source=self.name, + harvested=datetime.now(UTC).strftime('%Y-%m-%dT%H:%M:%SZ'), + harvested_by=('immvoke %s (hand curated, not scraped: ASD ' + 'exposes no HTML form)' % __version__), + source_urls={'home': HOME_URL, 'drive_folder': DRIVE_FOLDER_URL}, + collections=collections) + + def searchUnits(self, query): + """ + Resolve a query to the one DataUnit the Delta Lake represents. + + Arguments: + query (Query): the validated request. + + Returns: + list: a single DataUnit. + """ + return [DataUnit(unit_id=COLLECTION, collection=query.collection, + url=DRIVE_FOLDER_URL, metadata=dict(query.filters))] + + def fetchUnit(self, unit, outdir, resume=True, progress=True): + """ + Download the Delta Lake folder from Google Drive via gdown. + + Unlike HttpClient.fetch, gdown has no resumable-range support, so + `resume` only controls whether an already complete looking directory + (one that already has *.parquet files) is trusted and left alone, + not whether a partial transfer is continued. + + Arguments: + unit (DataUnit): the asd table. + outdir (Path): the mirror root. + resume (bool): reuse a directory that already looks complete. + progress (bool): show gdown's own progress output. + + Returns: + DownloadResult: the outcome. + + Raises: + AsdFetchError: if gdown is not installed, the folder could not be + fetched, or the result contains no *.parquet files. + """ + dest = self.resolveOutputPath(unit, outdir) + + if resume and dest.exists() and any(dest.rglob('*.parquet')): + log.info('%s already present; skipping', dest) + else: + self._downloadDeltaLake(unit.url, dest, progress=progress) + + return DownloadResult(unit=unit, path=dest, + sha256=_hashDirectory(dest), + size_bytes=_directorySize(dest)) + + @staticmethod + def _downloadDeltaLake(url, dest, progress=True): + """ + Download one Google Drive folder into dest with gdown. + + Arguments: + url (str): the Drive folder URL. + dest (Path): where to write it. + progress (bool): show gdown's own progress output. + + Raises: + AsdFetchError: if gdown is not installed, the download failed, or + the folder contained no *.parquet files. + """ + try: + import gdown + except ImportError as error: + raise AsdFetchError( + "gdown is required to download ASD's Delta Lake from Google " + "Drive; install it with 'pip install gdown'") from error + + dest.mkdir(parents=True, exist_ok=True) + try: + gdown.download_folder(url, output=str(dest), quiet=not progress, + use_cookies=False) + except Exception as error: + raise AsdFetchError( + 'failed to download the ASD Delta Lake folder from %s: %s' + % (url, error)) from error + + if not any(dest.rglob('*.parquet')): + raise AsdFetchError( + 'gdown reported success but no *.parquet files were found ' + 'under %s; the ASD Drive folder layout has changed' % dest) + + def readUnit(self, path, unit, chunksize=None): + """ + Open the downloaded Delta Lake, one partition file per chunk. + + Arguments: + path (Path): the downloaded Delta Lake directory. + unit (DataUnit): what it is. + chunksize (int): unused; each Parquet partition is already a + manageably sized chunk. Kept for interface parity with SourceBase. + + Returns: + tuple: ({}, iterator of raw record chunks). ASD carries no + per-unit metadata line. + + Raises: + AsdParseError: if the directory contains no *.parquet partitions. + """ + partitions = sorted(Path(path).glob('*.parquet')) + if not partitions: + raise AsdParseError( + 'no *.parquet files under %s; the ASD Delta Lake layout has ' + 'changed' % path) + + return {}, (pandas.read_parquet(x) for x in partitions) + + def normalizeChunk(self, metadata, chunk, unit, offset, report): + """ + Pass a chunk's columns through unchanged, applying provenance and any + dataset filter the fetch step could not apply server-side. + + Arguments: + metadata (dict): unused; ASD carries no per-unit metadata. + chunk (pandas.DataFrame): raw records, ASD's own columns. + unit (DataUnit): what they came from. + offset (int): unused; kept for interface parity with SourceBase. + report (dict): counters to accumulate into. + + Returns: + pandas.DataFrame: the chunk with provenance columns added. + """ + report['rows_in'] += len(chunk) + + frame = self._applyClientFilter(chunk, unit) + frame = frame.copy() + + for column in _NESTED_COLUMNS: + if column in frame.columns: + frame[column] = frame[column].apply( + lambda x: json.dumps(x, default=str) if x is not None else '') + + row_hash = rowHashes(frame) + + frame['immvoke_source'] = self.name + frame['immvoke_collection'] = unit.collection + frame['immvoke_unit_id'] = unit.unit_id + frame['immvoke_row_hash'] = row_hash + + report['rows_out'] += len(frame) + + return frame + + def _applyClientFilter(self, frame, unit): + """ + Apply dataset filtering that the Drive folder download cannot do itself. + + Arguments: + frame (pandas.DataFrame): raw records for one chunk. + unit (DataUnit): the table, carrying the resolved filters. + + Returns: + pandas.DataFrame: the filtered chunk, or the chunk unchanged when no + filter applies. + """ + wanted = (unit.metadata or {}).get('dataset') + if not wanted or wanted == '*' or 'dataset' not in frame.columns: + return frame + + return frame[frame['dataset'] == wanted] diff --git a/src/immvoke/Sources/Specificity/CovAbDab.py b/src/immvoke/Sources/Specificity/CovAbDab.py new file mode 100644 index 0000000..f9292eb --- /dev/null +++ b/src/immvoke/Sources/Specificity/CovAbDab.py @@ -0,0 +1,183 @@ +""" +CoV-AbDab - coronavirus-binding antibody and nanobody entries + +CoV-AbDab (https://opig.stats.ox.ac.uk/webapps/covabdab/) publishes its whole +collection as a single bulk CSV, so the one table this source offers *is* +that CSV rather than an index into many separately downloadable files. The +CSV is not served from a fixed endpoint, though: its filename is stamped +with the release date (e.g. CoV-AbDab_080224.csv), so the current download +URL has to be read off the homepage, and that costs an HTTP request every +time the table is resolved. + +There is no search form to build a controlled vocabulary from, so +harvestSchema returns a hand curated snapshot rather than one scraped from a +live page. +""" + +# Info +__author__ = 'Pramod Shinde' + +# Imports +import logging +from datetime import UTC + +import pandas +from bs4 import BeautifulSoup + +# Immvoke imports +from immvoke.Exceptions import CovAbDabParseError +from immvoke.Sources.Base import DataUnit, SourceBase +from immvoke.Sources.Specificity.RowHash import rowHashes + +log = logging.getLogger(__name__) + +#: Endpoints. +HOST = 'https://opig.stats.ox.ac.uk' +HOME_URL = HOST + '/webapps/covabdab/' + +#: The one table this source offers; the whole CSV is the data unit. +COLLECTION = 'antibodies' + + +def findCsvUrl(client): + """ + Read the current bulk CSV download link off the CoV-AbDab homepage. + + Arguments: + client (HttpClient): the shared HTTP client. + + Returns: + str: the absolute CSV URL. + + Raises: + CovAbDabParseError: if no download link can be found, meaning the page + layout has changed. + """ + html = client.get(HOME_URL).text + soup = BeautifulSoup(html, 'html.parser') + + for link in soup.find_all('a', href=True): + href = link['href'] + if 'download' in href.lower() and href.lower().endswith('.csv'): + return href if href.startswith('http') else HOST + href + + raise CovAbDabParseError( + 'no CSV download link found on %s; the homepage layout has changed' + % HOME_URL) + + + +class CovAbDabSource(SourceBase): + """ + The CoV-AbDab specificity source. + """ + + name = 'covabdab' + description = ('CoV-AbDab: coronavirus-binding antibody and nanobody ' + 'sequences') + homepage = HOME_URL + collections = (COLLECTION,) + collection_help = { + COLLECTION: ('antibody and nanobody entries against SARS-CoV-2 and ' + 'related coronaviruses, from the bulk CSV export'), + } + + license = 'CC BY-NC 4.0 (https://creativecommons.org/licenses/by-nc/4.0/)' + citation = ( + 'Raybould MIJ, Kovaltsuk A, Marks C, Deane CM. CoV-AbDab: the ' + 'Coronavirus Antibody Database. Bioinformatics. 2021;37(5):734-735. ' + 'doi:10.1093/bioinformatics/btaa739', + ) + + def harvestSchema(self): + """ + Build the hand curated snapshot. + + There is no live form to scrape for a controlled vocabulary, so this + returns the same field definitions packaged in schema.yaml, timestamped + now. `immvoke schema refresh --source covabdab` exists mainly so the + timestamp and any future edits to this method stay reproducible the + same way a real harvest would. + + Returns: + SourceSchema: the snapshot. + """ + from datetime import datetime + + from immvoke.Schema import Collection, SourceSchema + from immvoke.Version import __version__ + + collections = {COLLECTION: Collection(name=COLLECTION)} + + return SourceSchema( + source=self.name, + harvested=datetime.now(UTC).strftime('%Y-%m-%dT%H:%M:%SZ'), + harvested_by=('immvoke %s (hand curated, not scraped: ' + 'CoV-AbDab exposes no HTML form)' % __version__), + source_urls={'home': HOME_URL}, + collections=collections) + + def searchUnits(self, query): + """ + Resolve a query to the one DataUnit the current CSV represents. + + The CSV's filename and URL change with every release, so this reads + the current download link off the homepage rather than caching one. + + Arguments: + query (Query): the validated request. + + Returns: + list: a single DataUnit. + """ + url = findCsvUrl(self.client) + unit_id = url.rsplit('/', 1)[-1] + + return [DataUnit(unit_id=unit_id, collection=query.collection, url=url)] + + def readUnit(self, path, unit, chunksize=50000): + """ + Open the downloaded CSV. + + Arguments: + path (Path): the downloaded file. + unit (DataUnit): what it is. + chunksize (int): rows per chunk. + + Returns: + tuple: ({}, iterator of raw record chunks). CoV-AbDab carries no + per-unit metadata line. + """ + # utf-8-sig: the export is saved with a leading BOM. + chunks = pandas.read_csv(path, chunksize=chunksize, dtype=str, + na_filter=False, encoding='utf-8-sig') + + return {}, chunks + + def normalizeChunk(self, metadata, chunk, unit, offset, report): + """ + Pass a chunk's columns through unchanged, applying provenance. + + Arguments: + metadata (dict): unused; CoV-AbDab carries no per-unit metadata. + chunk (pandas.DataFrame): raw records, CoV-AbDab's own columns. + unit (DataUnit): what they came from. + offset (int): unused; kept for interface parity with SourceBase. + report (dict): counters to accumulate into. + + Returns: + pandas.DataFrame: the chunk with provenance columns added. + """ + report['rows_in'] += len(chunk) + + row_hash = rowHashes(chunk) + + frame = chunk.copy() + frame['immvoke_source'] = self.name + frame['immvoke_collection'] = unit.collection + frame['immvoke_unit_id'] = unit.unit_id + frame['immvoke_row_hash'] = row_hash + + report['rows_out'] += len(frame) + + return frame diff --git a/src/immvoke/Sources/Specificity/Distances.py b/src/immvoke/Sources/Specificity/Distances.py new file mode 100644 index 0000000..dec925b --- /dev/null +++ b/src/immvoke/Sources/Specificity/Distances.py @@ -0,0 +1,291 @@ +""" +Every gated pair's distance, whatever the threshold + +symscan only ever returns the pairs within a threshold. The pairs file of a +Hamming or Levenshtein run lists every (query, reference) pair that passes +the run's gates instead -- the gene gate (`gene_cols`), the exact columns +(`exact_cols`) and, under Hamming, each compared column's length -- with +its distance, within the threshold or not. allPairDistances scores those +pairs directly with numpy: + + Hamming a block of queries against the same-length references, + compared position by position; an X on either side never + counts (as in Symscan.py) + Levenshtein the edit distance table, filled one query residue at a + time for a whole block of pairs at once; an X is an + ordinary letter (as in Symscan.py) + +Candidates come from indexes rather than from testing every pair: +references are indexed by their exact-column key (with the per-column +lengths under Hamming) and by every gene they name, so a query only meets +the references that pass its gates. Queries go in blocks of consecutive +rows, so the pairs come out in `query_row` order and never all sit in +memory at once -- without a gene gate, a repertoire against IEDB is on the +order of 10^8 pairs. +""" + +# Info +__author__ = 'Pramod Shinde' + +# Imports +import logging + +import numpy +import pandas + +# Immvoke imports +from immvoke.Sources.Specificity.Annotate import ( + DEFAULT_ID_COL, + WILDCARD, + _exactValues, + _formatPositions, + _rowValues, +) +from immvoke.Sources.Specificity.Genes import parseCall +from immvoke.Sources.Specificity.Symscan import _resolveLimits + +log = logging.getLogger(__name__) + +#: Consecutive query rows scored, and yielded, together. +QUERY_BLOCK = 2000 + +#: Most residues (pairs x sequence length) one vectorized step holds, which +#: bounds its memory to some tens of MB. +RESIDUE_BUDGET = 5_000_000 + +_WILDCARD_CODE = ord(WILDCARD) +_NO_ROWS = numpy.array([], dtype=numpy.int64) + + +def _codes(seqs): + """ + Arguments: + seqs (list): equal-length strings. + + Returns: + numpy.ndarray: one row of byte codes per string. + """ + joined = ''.join(seqs).encode('latin-1', 'replace') + + return numpy.frombuffer(joined, dtype=numpy.uint8).reshape(len(seqs), -1) + + +def _byLength(rows, seqs): + """ + Returns: + dict: sequence length to the `rows` (a numpy array) with that length. + """ + lengths = numpy.array([len(seqs[i]) for i in rows], dtype=numpy.int64) + + return {int(length): rows[lengths == length] for length in numpy.unique(lengths)} + + +def _levenshtein(a, b): + """ + The edit distance of every pair at once. + + Arguments: + a (numpy.ndarray): query codes, one row per pair. + b (numpy.ndarray): reference codes, one row per pair. + + Returns: + numpy.ndarray: one distance per pair. + """ + pairs, n = b.shape + steps = numpy.arange(1, n + 1, dtype=numpy.int32) + prev = numpy.tile(numpy.arange(n + 1, dtype=numpy.int32), (pairs, 1)) + for i in range(1, a.shape[1] + 1): + # Substitution (or match) and deletion first; insertion runs along + # the row: cur[j] = min(best[j], cur[j-1] + 1) with cur[0] = i, + # which unrolls to min(i + j, j + min over k <= j of best[k] - k). + best = numpy.minimum(prev[:, :-1] + (b != a[:, i - 1:i]), prev[:, 1:] + 1) + cur = numpy.empty_like(prev) + cur[:, 0] = i + cur[:, 1:] = numpy.minimum( + i + steps, numpy.minimum.accumulate(best - steps, axis=1) + steps) + prev = cur + + return prev[:, n] + + +def _levenshteinPairs(query_seqs, ref_seqs, query_rows, ref_rows): + """ + Score every query of `query_rows` against every reference of `ref_rows` + under Levenshtein, a block of same-length queries against same-length + references at a time. + + Yields: + tuple: (query positions, reference positions, distances), numpy + arrays. + """ + query_groups = _byLength(query_rows, query_seqs) + for ref_length, refs in _byLength(ref_rows, ref_seqs).items(): + ref_codes = _codes([ref_seqs[i] for i in refs]) + for query_length, queries in query_groups.items(): + query_codes = _codes([query_seqs[i] for i in queries]) + step = max(1, RESIDUE_BUDGET // (len(refs) * (query_length + ref_length))) + for start in range(0, len(queries), step): + block = query_codes[start:start + step] + distances = _levenshtein(numpy.repeat(block, len(refs), axis=0), + numpy.tile(ref_codes, (len(block), 1))) + yield (numpy.repeat(queries[start:start + step], len(refs)), + numpy.tile(refs, len(block)), distances) + + +def _hammingPairs(query_seqs, ref_seqs, query_rows, ref_rows, col_lengths): + """ + Score every query of `query_rows` against every reference of `ref_rows` + under Hamming -- all one length, the gate saw to that. + + Arguments: + col_lengths (tuple): the length of each compared column, pooled into + the sequences in this order, for the mismatch positions. + + Yields: + tuple: (query positions, reference positions, distances, mismatch + positions as strings, see Annotate._formatPositions). + """ + starts = numpy.cumsum((0,) + tuple(col_lengths))[:-1] + ref_codes = _codes([ref_seqs[i] for i in ref_rows]) + query_codes = _codes([query_seqs[i] for i in query_rows]) + step = max(1, RESIDUE_BUDGET // (len(ref_rows) * ref_codes.shape[1])) + for start in range(0, len(query_rows), step): + a = query_codes[start:start + step, None, :] + b = ref_codes[None, :, :] + differ = (a != b) & (a != _WILDCARD_CODE) & (b != _WILDCARD_CODE) + differ = differ.reshape(-1, ref_codes.shape[1]) + positions = [] + for row in differ: + where = numpy.flatnonzero(row) + cols = numpy.searchsorted(starts, where, side='right') - 1 + positions.append(_formatPositions( + list(zip(cols.tolist(), (where - starts[cols]).tolist())), + len(col_lengths))) + count = len(a) + yield (numpy.repeat(query_rows[start:start + step], len(ref_rows)), + numpy.tile(ref_rows, count), differ.sum(axis=1), positions) + + +def allPairDistances(queries, reference, compare_cols, max_distance, + distance_type='hamming', id_col=DEFAULT_ID_COL, exact_cols=None, + threshold_is_fraction=False, gene_cols=None): + """ + Score every (query, reference) pair that passes the gates, whatever its + distance. + + Arguments: as Symscan.runSymscanAnnotation -- `max_distance` and + `threshold_is_fraction` only decide `within_threshold` here. + + Yields: + pandas.DataFrame: the pairs of one block of consecutive query rows, + in `query_row` order, closest first: `query_row`, `reference_id` (the + reference's `id_col`, once per query row -- its closest row), + `mismatches` (the distance), `identity` (1 - mismatches / the query's + compared length), `within_threshold` (whether symscan's threshold + admits it), plus `mismatch_positions` under 'hamming'. At least one + table, empty when no pair passes. + """ + hamming = distance_type == 'hamming' + query_values = _rowValues(queries, compare_cols) + ref_values = _rowValues(reference, compare_cols) + query_seqs = [''.join(v) for v in query_values] + ref_seqs = [''.join(v) for v in ref_values] + limits = _resolveLimits(query_values, max_distance, threshold_is_fraction) + + # Gate keys, as Symscan applies them: exact columns, and under Hamming + # each compared column's length (only same-length strings compare). + query_keys = _exactValues(queries, exact_cols) + ref_keys = _exactValues(reference, exact_cols) + if hamming: + query_keys = [k + tuple(map(len, v)) for k, v in zip(query_keys, query_values)] + ref_keys = [k + tuple(map(len, v)) for k, v in zip(ref_keys, ref_values)] + query_blank = [any(v == '' for v in values) for values in query_values] + by_key = {} + for row, (key, values) in enumerate(zip(ref_keys, ref_values)): + if not any(v == '' for v in values): + by_key.setdefault(key, []).append(row) + by_key = {key: numpy.array(rows, dtype=numpy.int64) for key, rows in by_key.items()} + + # Gene gate: each reference row under every gene it names, per column. + def parsed(frame, col): + values = frame[col].astype(object).fillna('').astype(str).tolist() + calls = {value: parseCall(value) for value in set(values)} + return [calls[value] for value in values] + + gene_cols = gene_cols or [] + query_genes = list(zip(*(parsed(queries, col) for col in gene_cols))) or \ + [()] * len(queries) + gene_index = [] + for col in gene_cols: + index = {} + for row, call in enumerate(parsed(reference, col)): + for gene in call: + index.setdefault(gene, []).append(row) + gene_index.append({gene: numpy.array(rows, dtype=numpy.int64) + for gene, rows in index.items()}) + + cache = {} + + def candidates(key, genes): + if (key, genes) not in cache: + rows = by_key.get(key, _NO_ROWS) + for index, call in zip(gene_index, genes): + # A blank or unrecognized call agrees with nothing. + agreeing = [index[g] for g in call if g in index] + rows = (numpy.intersect1d(rows, numpy.concatenate(agreeing)) + if agreeing else _NO_ROWS) + cache[key, genes] = rows + return cache[key, genes] + + ref_ids = reference[id_col].astype(str).to_numpy() + lengths = numpy.array([len(seq) for seq in query_seqs], dtype=numpy.int64) + columns = ['query_row', 'reference_id', 'mismatches', 'identity', + 'within_threshold'] + (['mismatch_positions'] if hamming else []) + yielded = False + for start in range(0, len(queries), QUERY_BLOCK): + groups = {} + for row in range(start, min(start + QUERY_BLOCK, len(queries))): + if not query_blank[row]: + groups.setdefault((query_keys[row], query_genes[row]), []).append(row) + + parts = [] + for (key, genes), rows in groups.items(): + refs = candidates(key, genes) + if not len(refs): + continue + rows = numpy.array(rows, dtype=numpy.int64) + if hamming: + col_lengths = key[len(key) - len(compare_cols):] + parts.extend(_hammingPairs(query_seqs, ref_seqs, rows, refs, col_lengths)) + else: + parts.extend(part + (None,) for part in + _levenshteinPairs(query_seqs, ref_seqs, rows, refs)) + + if parts: + pairs = pandas.DataFrame({ + 'query_row': numpy.concatenate([p[0] for p in parts]), + 'col': numpy.concatenate([p[1] for p in parts]), + 'dist': numpy.concatenate([p[2] for p in parts])}) + pairs['id'] = ref_ids[pairs['col'].to_numpy()] + if hamming: + pairs['positions'] = [x for p in parts for x in p[3]] + # One reference id can sit on several rows; keep its closest one + # (lowest reference position on a tie), as Symscan does. + pairs = (pairs.sort_values(['query_row', 'dist', 'id', 'col']) + .drop_duplicates(['query_row', 'id'])) + else: + pairs = pandas.DataFrame({'query_row': _NO_ROWS, 'col': _NO_ROWS, + 'dist': _NO_ROWS, 'id': [], 'positions': []}) + rows, dist = pairs['query_row'].to_numpy(), pairs['dist'].to_numpy() + table = pandas.DataFrame({ + 'query_row': rows, 'reference_id': pairs['id'].to_numpy(), + 'mismatches': dist, 'identity': 1 - dist / lengths[rows], + 'within_threshold': dist <= limits[rows]}) + if hamming: + table['mismatch_positions'] = pairs['positions'].to_numpy() + if len(table) or not yielded: + yielded = True + yield table[columns] + + if not yielded: + yield pandas.DataFrame({col: [] for col in columns}) diff --git a/src/immvoke/Sources/Specificity/Genes.py b/src/immvoke/Sources/Specificity/Genes.py new file mode 100644 index 0000000..ac467e3 --- /dev/null +++ b/src/immvoke/Sources/Specificity/Genes.py @@ -0,0 +1,139 @@ +""" +V/J-gene call parsing and comparison for the annotation gene gate + +An AIRR `v_call`/`j_call` cell is not always one clean name: IEDB alone mixes +allele level calls (`IGHV1-2*02`), gene level ones (`IGHV1-2`), several calls +in one cell (`IGHV1-2*02, IGHV1-2*04`) and blanks. Comparing those as plain +strings makes `IGHV1-2*02` and `IGHV1-2` disagree, so the gene gate parses +each cell into the genes it names and compares those instead: two calls +agree on the same gene, whatever the allele (IGHV1-2*02 == IGHV1-2*04 == +IGHV1-2, but IGHV1-2 != IGHV1-69). + +Two cells agree when they name a gene in common. A blank cell, or one naming +no recognizable gene, has nothing to compare and agrees with nothing -- not +even another blank. + +The gene pattern is Change-O's (changeo/Gene.py, Immcantation), so a name +parses here the way Change-O's getGene would parse it. The allele is read +as an optional `*` suffix of the gene instead of with Change-O's allele +pattern, which also accepts `-` as the allele separator and so would read a +gene level `IGHV1-2` as gene `IGHV1` allele `2`. +""" + +# Info +__author__ = 'Pramod Shinde' + +# Imports +import logging +import re + +import numpy + +log = logging.getLogger(__name__) + +#: Change-O's gene pattern, plus an optional `*allele` suffix. +_CALL_REGEX = re.compile(r'(?P(IG[HLK]|TR[ABGD])([VDJ][A-R0-9]+[-/\w]*))' + r'(\*(?P[.\w]+))?') + + +def parseCall(call): + """ + Parse one gene call cell into the genes it names. + + Arguments: + call (str): the cell, e.g. 'IGHV1-2*02, IGHV1-2*04'. + + Returns: + frozenset: the gene names, alleles dropped (e.g. {'IGHV1-2'}); empty + when the cell is blank or names no recognizable gene. + """ + return frozenset(match.group('gene') for match in _CALL_REGEX.finditer(call or '')) + + +def callsAgree(a, b): + """ + Whether two parsed calls (see parseCall) agree. + + Arguments: + a (frozenset): one parsed call. + b (frozenset): the other. + + Returns: + bool: True if they name a gene in common. A blank or unrecognized call + (empty) agrees with nothing. + """ + return bool(a & b) + + +class GeneGate: + """ + The gene agreement test between every query row and every reference + row, over one or more gene call columns -- a pair passes only if it + agrees (callsAgree) on each of them. + """ + + def __init__(self, queries, reference, gene_cols): + """ + Arguments: + queries (pandas.DataFrame): the query rows, in result order. + reference (pandas.DataFrame): the reference rows, by position. + gene_cols (list): gene call columns to gate on, e.g. + ['v_call', 'j_call']. + """ + self.query_calls = [self._parseColumn(queries, col, 'query') + for col in gene_cols] + self.ref_calls = [self._parseColumn(reference, col, 'reference') + for col in gene_cols] + + @staticmethod + def _parseColumn(frame, col, label): + """ + Returns: + list: parseCall of every value of `frame[col]`, in row order -- + each distinct value parsed once. + """ + values = frame[col].astype(object).fillna('').astype(str).tolist() + parsed = {value: parseCall(value) for value in set(values)} + unparsed = sum(1 for value in values if value.strip() and not parsed[value]) + if unparsed: + log.warning("%d %s %s values name no recognizable gene; those rows " + 'agree with no gene', unparsed, label, col) + + return [parsed[value] for value in values] + + def agrees(self, query_row, ref_row): + """ + Arguments: + query_row (int): a query position. + ref_row (int): a reference position. + + Returns: + bool: whether the pair agrees on every gene column. + """ + return all(callsAgree(query_calls[query_row], ref_calls[ref_row]) + for query_calls, ref_calls in zip(self.query_calls, self.ref_calls)) + + def keep(self, query_rows, ref_rows): + """ + Arguments: + query_rows (numpy.ndarray): query positions, one per pair. + ref_rows (numpy.ndarray): reference positions, one per pair. + + Returns: + numpy.ndarray: one bool per pair, True where it agrees. + """ + return numpy.fromiter((self.agrees(q, r) for q, r in zip(query_rows.tolist(), + ref_rows.tolist())), + dtype=bool, count=len(query_rows)) + + +def buildGeneGate(queries, reference, gene_cols): + """ + Returns: + GeneGate: for `gene_cols`, or None when `gene_cols` is empty or None + -- no gene gate at all. + """ + if not gene_cols: + return None + + return GeneGate(queries, reference, gene_cols) diff --git a/src/immvoke/Sources/Specificity/Iedb.py b/src/immvoke/Sources/Specificity/Iedb.py new file mode 100644 index 0000000..deb073f --- /dev/null +++ b/src/immvoke/Sources/Specificity/Iedb.py @@ -0,0 +1,961 @@ +""" +IEDB - B-cell and T-cell receptor sequences with their epitopes and assays + +IEDB (https://www.iedb.org/) publishes receptor sequences as a single bulk +ZIP export; its PostgREST API (https://query-api.iedb.org/api/v1) is used +only to look up what the export leaves out: the antigens of those +receptors' epitopes, and, when filtered on, their assays. A table +is not an index into many separately downloadable files: the table itself +*is* the data unit, so each table produces exactly one DataUnit. + +IEDB has no HTML search form to scrape, so harvestSchema returns a hand +curated snapshot rather than one built from a live page. + +The receptor tables (bcr, tcr) come from IEDB's own AIRR export, +airr_full.zip: one AIRR rearrangement TSV per table, a record per chain +linked by cell_id, built from IEDB's calculated (IMGT standardized) fields +and carrying the epitope and assay context alongside. They are already AIRR, +so normalizeChunk passes them through, adding only provenance columns. Every +receptor species is kept unless --air_species names one (see +RECEPTOR_SPECIES and _speciesCells), every antigen unless +--antigen_species names its organism (see _antigenEpitopes), every +epitope structure unless --antigen_structure_type names one (see +_structureEpitopes), every assay type and disease unless --assay_type or +--disease names one (see _typeAssays and _diseaseAssays), and single +chains unless --paired_only asks for both (see _pairedCells). Each record +already carries its epitope, assay ids and outcome, so IEDB's separate +B-cell assay tables (bcell_search, bcr_to_bcell) are not offered. +""" + +# Info +__author__ = 'Pramod Shinde' + +# Imports +import difflib +import logging +import os +import time +import zipfile +from datetime import UTC +from pathlib import Path + +import pandas + +# Immvoke imports +from immvoke.Convert import coerceAirrTypes +from immvoke.Exceptions import IedbParseError, SchemaError +from immvoke.Http import hashFile +from immvoke.Sources.Base import DataUnit, DownloadResult, Query, SourceBase +from immvoke.Sources.Specificity.Paginate import pageByRange +from immvoke.Sources.Specificity.Report import humanDuration, humanPath +from immvoke.Sources.Specificity.RowHash import rowHashes + +log = logging.getLogger(__name__) + +#: Endpoints. +API_BASE = 'https://query-api.iedb.org/api/v1' +BULK_URL = 'https://www.iedb.org/downloader.php?file_name=doc/airr_full.zip' +#: The bulk ZIP's file name in the cache (see cacheDir). +BULK_ZIP = 'airr_full.zip' +#: Seconds a cached bulk ZIP is reused before it is fetched again. IEDB +#: refreshes the export weekly and its download sends no Last-Modified or +#: ETag to check against, so freshness goes by age. +CACHE_MAX_AGE = 24 * 60 * 60 + +#: Rows requested per page from the PostgREST API. +API_PAGE_SIZE = 2000 + +#: table -> its member of the AIRR bulk ZIP. +_TABLE_MEMBER = { + 'bcr': 'bcr_full_airr.tsv', + 'tcr': 'tcr_full_airr.tsv', +} + +#: The receptor species IEDB's AIRR export records in `locus_species`, as +#: NCBI taxon code -> (name, parent code). Names are IEDB's own +#: (receptor_species_names on its bcr_search/tcr_search API) without the +#: common name in brackets, so 'Homo sapiens (human)' is 'Homo sapiens'. +#: Codes from 10000000 up are IEDB's own strain and line codes (its ONTIE +#: terms), not real NCBI taxa; each names its parent species, so asking for +#: 'Mus musculus' also keeps BALB/c, C57BL/6, Xenomouse and the rest. +#: Built by joining the export's codes to those names on +#: iedb_receptor_group (export of 2026-10-20); a code IEDB adds later is +#: kept when no species is asked for, but cannot be asked for by name until +#: it is added here. +RECEPTOR_SPECIES = { + 'NCBITAXON:7797': ('Squalus acanthias', None), + 'NCBITAXON:7801': ('Ginglymostoma cirratum', None), + 'NCBITAXON:9031': ('Gallus gallus', None), + 'NCBITAXON:9534': ('Chlorocebus aethiops', None), + 'NCBITAXON:9541': ('Macaca fascicularis', None), + 'NCBITAXON:9544': ('Macaca mulatta', None), + 'NCBITAXON:9598': ('Pan troglodytes', None), + 'NCBITAXON:9606': ('Homo sapiens', None), + 'NCBITAXON:9823': ('Sus scrofa', None), + 'NCBITAXON:9835': ('Camelid', None), + 'NCBITAXON:9837': ('Camelus bactrianus', None), + 'NCBITAXON:9838': ('Camelus dromedarius', None), + 'NCBITAXON:9844': ('Lama glama', None), + 'NCBITAXON:9913': ('Bos taurus', None), + 'NCBITAXON:9986': ('Oryctolagus cuniculus', None), + 'NCBITAXON:10090': ('Mus musculus', None), + 'NCBITAXON:10114': ('Rattus', None), + 'NCBITAXON:10116': ('Rattus norvegicus', None), + 'NCBITAXON:10117': ('Rattus rattus', None), + 'NCBITAXON:30538': ('Vicugna pacos', None), + 'NCBITAXON:32630': ('synthetic construct', None), + 'NCBITAXON:36176': ('Chiloscyllium plagiosum', None), + 'NCBITAXON:168098': ('Orectolobus maculatus', None), + 'NCBITAXON:419612': ('Camelus ferus', None), + 'NCBITAXON:447135': ('Myodes glareolus', None), + 'NCBITAXON:3122392': ('Nothocricetulus migratorius', None), + 'NCBITAXON:10000000': ('Mus musculus BALB/c', 'NCBITAXON:10090'), + 'NCBITAXON:10000038': ('Mus musculus C3H/He', 'NCBITAXON:10090'), + 'NCBITAXON:10000061': ('Mus musculus FVB', 'NCBITAXON:10090'), + 'NCBITAXON:10000067': ('Mus musculus C57BL/6', 'NCBITAXON:10090'), + 'NCBITAXON:10000206': ('Oryctolagus cuniculus New Zealand White', + 'NCBITAXON:9986'), + 'NCBITAXON:10000228': ('Mus musculus CBA', 'NCBITAXON:10090'), + 'NCBITAXON:10000236': ('Mus musculus CBA/J', 'NCBITAXON:10090'), + 'NCBITAXON:10000248': ('Mus musculus 129 GIX+', 'NCBITAXON:10090'), + 'NCBITAXON:10000275': ('Mus musculus B10.RIII', 'NCBITAXON:10090'), + 'NCBITAXON:10000617': ('Mus musculus HuMAb (Medarex)', 'NCBITAXON:10090'), + 'NCBITAXON:10000645': ('Mus musculus PrP null', 'NCBITAXON:10090'), + 'NCBITAXON:10000655': ('Mus musculus Swiss Webster', 'NCBITAXON:10090'), + 'NCBITAXON:10000657': ('Mus musculus Xenomouse', 'NCBITAXON:10090'), + 'NCBITAXON:10000667': ('Rattus norvegicus Sprague-Dawley', + 'NCBITAXON:10116'), + 'NCBITAXON:10000915': ('Mus musculus Kunming', 'NCBITAXON:10090'), + 'NCBITAXON:10001602': ('Mus musculus RBF', 'NCBITAXON:10090'), + 'NCBITAXON:10002257': ('Mus musculus Kymouse', 'NCBITAXON:10090'), +} + +#: --air_species's vocabulary: every name in RECEPTOR_SPECIES, sorted. +_AIR_SPECIES = tuple(sorted(name for name, _ in RECEPTOR_SPECIES.values())) + +#: --outcome's vocabulary. The AIRR export writes iedb_assay_outcome as +#: positive or negative, in lower case, so --outcome matches it ignoring +#: case. +_RECEPTOR_OUTCOMES = ('Positive', 'Negative') + +#: --antigen_species takes any antigen organism name IEDB uses, so the schema +#: carries only these as examples (see Field.open_values). +_ANTIGEN_SPECIES_EXAMPLES = ('Influenza A virus', 'SARS-CoV2', + 'Human immunodeficiency virus 1', 'Homo sapiens') + +#: --antigen_structure_type's vocabulary: IEDB's epitope structure_type, the +#: same choice as its home page search's Epitope Structure. 'Discontinuous +#: peptide' also keeps 'Discontinuous peptide on multi chain', the way a +#: species keeps its strains. +_ANTIGEN_STRUCTURE_TYPES = ('Linear peptide', 'Discontinuous peptide', + 'Discontinuous peptide on multi chain', 'Non-peptidic') + +#: --assay_type and --disease take any assay type or disease IEDB uses for +#: the table's assays, so the schema carries only these as examples (see +#: Field.open_values). An assay type is the first part of IEDB's +#: assay_names: 'neutralization' of 'neutralization |biological activity'. +_ASSAY_TYPE_EXAMPLES = ('neutralization', '3D structure', 'qualitative binding', + 'dissociation constant KD') +_DISEASE_EXAMPLES = ('COVID-19', 'influenza', 'rheumatoid arthritis', + 'viral infectious disease') + +#: Epitope ids asked for per epitope_search request (see _epitopeRecords); +#: keeps each request URL a few kilobytes long. Assay ids likewise (see +#: _assayRecords). +EPITOPE_BATCH = 300 + +#: table -> IEDB's assay search endpoint and its id field. +_ASSAY_SEARCH = { + 'bcr': ('bcell_search', 'bcell_id'), + 'tcr': ('tcell_search', 'tcell_id'), +} + +#: --paired_only: the chains a receptor needs, one of each side of a pair: +#: a heavy chain and a light chain, a beta and an alpha, or a delta and a +#: gamma. +CHAIN_PAIRS = ( + ('IGH', frozenset(['IGK', 'IGL'])), + ('TRB', frozenset(['TRA'])), + ('TRD', frozenset(['TRG'])), +) +_LOCI = frozenset(['IGH', 'IGK', 'IGL', 'TRA', 'TRB', 'TRD', 'TRG']) + +#: The columns added to bcr and tcr from each record's epitope: its +#: antigen, the organism that antigen comes from, and the epitope's +#: structure type (see _epitopeAntigens), placed after iedb_epitope_name. +ANTIGEN_COL = 'iedb_antigen' +ANTIGEN_SPECIES_COL = 'iedb_antigen_source_species' +ANTIGEN_STRUCTURE_COL = 'iedb_antigen_structure_type' +#: Joins an epitope's antigens when IEDB curates several. +ANTIGEN_SEP = ' | ' + + +def _baseName(name): + """ + An IEDB organism name without its bracketed part, lower case: + 'Homo sapiens (human)' -> 'homo sapiens', and a strain such as + 'Influenza A virus (A/Puerto Rico/8/1934(H1N1))' -> 'influenza a virus'. + """ + return name.split(' (', 1)[0].strip().lower() + + +def _resolveAntigen(name, records): + """ + Find the IEDB organism an --antigen_species name stands for. + + IEDB's organism names and ids come as two lists per epitope, not always + in the same order, so a name is paired with its id only where an epitope + has a single source organism. The name matches ignoring case, and + without IEDB's bracketed common name ('Homo sapiens' for 'Homo sapiens + (human)') as long as that leaves a single organism; the full name always + wins, so 'Influenza A virus' is the species, not one of its strains. + + Arguments: + name (str): the --antigen_species value. + records (list): epitope_search records, with source_organism_names + and source_organism_iris. + + Returns: + tuple: (IEDB's own name, its organism id). + + Raises: + SchemaError: if no organism of the table's epitopes has that name, + naming the closest ones. + """ + exact, base = {}, {} + for record in records: + names = record.get('source_organism_names') or [] + iris = record.get('source_organism_iris') or [] + if len(names) == 1 and len(iris) == 1: + exact.setdefault(names[0].lower(), (names[0], iris[0])) + base.setdefault(_baseName(names[0]), set()).add((names[0], iris[0])) + + wanted = name.strip().lower() + if wanted in exact: + return exact[wanted] + if len({iri for _, iri in base.get(wanted, ())}) == 1: + return sorted(base[wanted])[0] + + close = difflib.get_close_matches(wanted, list(exact), n=5, cutoff=0.6) + raise SchemaError( + "no antigen organism named '%s' among this table's epitopes%s" + % (name, '; closest IEDB names: %s' % ', '.join(exact[x][0] for x in close) + if close else '')) + + +def _epitopeRecords(client, path, chunksize): + """ + Look up the epitopes of a receptor AIRR table on IEDB's epitope_search. + + The AIRR export records each record's epitope (reactivity_ref) but not + the epitope's antigen, that antigen's organism or the epitope's + structure type, so this asks the API about the table's own epitopes, in + batches of EPITOPE_BATCH. + + Arguments: + client (HttpClient): the shared HTTP client. + path (Path): the downloaded AIRR TSV. + chunksize (int): rows per chunk while scanning it. + + Returns: + list: epitope_search records: structure_id, the source organism names, + ids and ids with every organism above them (source_organism_iri_search), + curated_source_antigens and structure_type. + """ + ids = set() + for chunk in pandas.read_csv(path, sep='\t', dtype=str, na_filter=False, + usecols=['reactivity_ref'], chunksize=chunksize): + ids.update(chunk['reactivity_ref']) + numbers = sorted(x.rsplit(':', 1)[-1] for x in ids if x) + log.info(' antigens: looking up %s epitopes on IEDB', format(len(numbers), ',')) + + records = [] + for start in range(0, len(numbers), EPITOPE_BATCH): + params = {'select': 'structure_id,source_organism_names,' + 'source_organism_iris,source_organism_iri_search,' + 'curated_source_antigens,structure_type', + 'structure_id': 'in.(%s)' + % ','.join(numbers[start:start + EPITOPE_BATCH])} + for batch in pageByRange(client, API_BASE + '/epitope_search', + page_size=API_PAGE_SIZE, params=params): + records.extend(batch) + + return records + + +def _antigenEpitopes(records, name): + """ + The epitopes (reactivity_ref) whose antigen comes from the organism + `name` or one under it. + + IEDB's source_organism_iri_search lists each epitope's organism and + every organism above it, so 'Influenza A virus' takes in every flu + strain without immvoke knowing the taxonomy. + + Arguments: + records (list): epitope_search records, from _epitopeRecords. + name (str): the --antigen_species value. + + Returns: + frozenset: the reactivity_ref values to keep. + """ + own, iri = _resolveAntigen(name, records) + kept = frozenset('IEDB_EPITOPE:%s' % record['structure_id'] for record in records + if iri in (record.get('source_organism_iri_search') or [])) + log.info(" antigen species '%s' (%s): %s of %s epitopes kept", own, iri, + format(len(kept), ','), format(len(records), ',')) + + return kept + + +def _structureEpitopes(records, name): + """ + The epitopes (reactivity_ref) of the structure type `name`, or of one + under it ('Discontinuous peptide' takes in 'Discontinuous peptide on + multi chain'). The schema has already checked the name against + _ANTIGEN_STRUCTURE_TYPES; IEDB's own value is compared ignoring case, in + case its capitalization drifts. + + Arguments: + records (list): epitope_search records, from _epitopeRecords. + name (str): the --antigen_structure_type value. + + Returns: + frozenset: the reactivity_ref values to keep. + """ + wanted = name.strip().lower() + kept = frozenset( + 'IEDB_EPITOPE:%s' % record['structure_id'] for record in records + if (record.get('structure_type') or '').lower() in (wanted, wanted + ' on multi chain')) + log.info(" antigen structure type '%s': %s of %s epitopes kept", name, + format(len(kept), ','), format(len(records), ',')) + + return kept + + +def _epitopeAntigens(records): + """ + Each epitope's antigen and that antigen's source species, from IEDB's + curated_source_antigens, and its structure_type, for the iedb_antigen, + iedb_antigen_source_species and iedb_antigen_structure_type columns. + + Arguments: + records (list): epitope_search records, from _epitopeRecords. + + Returns: + dict: reactivity_ref -> (antigen names, their source species, + structure type), the first two ' | ' joined in IEDB's order when an + epitope has several antigens, each blank when IEDB names none. + """ + antigens = {} + for record in records: + curated = record.get('curated_source_antigens') or [] + names = list(dict.fromkeys(x.get('name') or '' for x in curated)) + species = list(dict.fromkeys(x.get('source_organism_name') or '' + for x in curated)) + antigens['IEDB_EPITOPE:%s' % record['structure_id']] = ( + ANTIGEN_SEP.join(x for x in names if x), + ANTIGEN_SEP.join(x for x in species if x), + record.get('structure_type') or '') + + return antigens + + +def _intersect(*sets): + """ + What every given set keeps: their intersection, ignoring the None of a + filter not given, and None when none was given. + """ + given = [x for x in sets if x is not None] + + return frozenset.intersection(*given) if given else None + + +def _assayRecords(client, path, chunksize, table): + """ + Look up the assays of a receptor AIRR table on IEDB's bcell_search (bcr) + or tcell_search (tcr). + + The AIRR export lists each record's assays (iedb_assay_ids) but not what + they were or the disease studied, so this asks the API about the table's + own assays, in batches of EPITOPE_BATCH. Only --assay_type and --disease + need it, so it runs only when one of them is given; given both, a record + is kept when one of its assays is of that type and studied that disease. + + Arguments: + client (HttpClient): the shared HTTP client. + path (Path): the downloaded AIRR TSV. + chunksize (int): rows per chunk while scanning it. + table (str): bcr or tcr. + + Returns: + list: records with `id` (the assay id, as text), assay_names, + disease_names, disease_iris and disease_iri_search (each disease and + every disease above it). + """ + endpoint, id_field = _ASSAY_SEARCH[table] + ids = set() + for chunk in pandas.read_csv(path, sep='\t', dtype=str, na_filter=False, + usecols=['iedb_assay_ids'], chunksize=chunksize): + for value in chunk['iedb_assay_ids']: + ids.update(x.strip() for x in value.split(',') if x.strip()) + numbers = sorted(ids, key=int) + log.info(' assays: looking up %s assays on IEDB', format(len(numbers), ',')) + + records = [] + for start in range(0, len(numbers), EPITOPE_BATCH): + params = {'select': '%s,assay_names,disease_names,disease_iris,' + 'disease_iri_search' % id_field, + id_field: 'in.(%s)' % ','.join(numbers[start:start + EPITOPE_BATCH])} + for batch in pageByRange(client, '%s/%s' % (API_BASE, endpoint), + page_size=API_PAGE_SIZE, params=params): + records.extend(dict(x, id=str(x[id_field])) for x in batch) + + return records + + +def _assayType(record): + """ + An assay's type: the first part of IEDB's assay_names, so + 'neutralization |biological activity' -> 'neutralization'. + """ + return (record.get('assay_names') or '').split('|', 1)[0].strip() + + +def _typeAssays(records, name): + """ + The assays (iedb_assay_ids) of the type `name`, matched ignoring case. + + Arguments: + records (list): assay records, from _assayRecords. + name (str): the --assay_type value. + + Returns: + frozenset: the assay ids to keep. + + Raises: + SchemaError: if none of the table's assays is of that type, naming + the closest types. + """ + types = {} + for record in records: + types.setdefault(_assayType(record).lower(), _assayType(record)) + wanted = name.strip().lower() + if wanted not in types: + close = difflib.get_close_matches(wanted, list(types), n=5, cutoff=0.5) + raise SchemaError( + "no assay type '%s' among this table's assays%s" + % (name, '; closest IEDB types: %s' % ', '.join(types[x] for x in close) + if close else '')) + + kept = frozenset(record['id'] for record in records + if _assayType(record).lower() == wanted) + log.info(" assay type '%s': %s of %s assays kept", types[wanted], + format(len(kept), ','), format(len(records), ',')) + + return kept + + +def _diseaseAssays(records, name): + """ + The assays (iedb_assay_ids) studying the disease `name` or one under it. + + Like an antigen organism, a disease is matched ignoring case and paired + with its id only where an assay names a single disease; IEDB's + disease_iri_search lists each disease and every disease above it, so + 'viral infectious disease' takes in COVID-19, influenza and the rest. + + Arguments: + records (list): assay records, from _assayRecords. + name (str): the --disease value. + + Returns: + frozenset: the assay ids to keep. + + Raises: + SchemaError: if no assay of the table names that disease, naming the + closest ones. + """ + diseases = {} + for record in records: + names = record.get('disease_names') or [] + iris = record.get('disease_iris') or [] + if len(names) == 1 and len(iris) == 1: + diseases.setdefault(names[0].lower(), (names[0], iris[0])) + + wanted = name.strip().lower() + if wanted not in diseases: + close = difflib.get_close_matches(wanted, list(diseases), n=5, cutoff=0.6) + raise SchemaError( + "no disease named '%s' among this table's assays%s" + % (name, '; closest IEDB names: %s' % ', '.join(diseases[x][0] for x in close) + if close else '')) + + own, iri = diseases[wanted] + kept = frozenset(record['id'] for record in records + if iri in (record.get('disease_iri_search') or [])) + log.info(" disease '%s' (%s): %s of %s assays kept", own, iri, + format(len(kept), ','), format(len(records), ',')) + + return kept + + +def _chainLoci(chunk): + """ + Each chain's locus: its locus, or, where IEDB left that blank, the locus + its V or J gene names (IGKV1-39*01 -> IGK). + + Arguments: + chunk (pandas.DataFrame): records with locus, v_call and j_call. + + Returns: + pandas.Series: the loci, blank where none is known. + """ + loci = chunk['locus'].str.upper() + for column in ('v_call', 'j_call'): + guess = chunk[column].str[:3].str.upper() + loci = loci.where(loci.isin(_LOCI), guess) + + return loci.where(loci.isin(_LOCI), '') + + +def _pairedCells(path, chunksize, codes=None): + """ + The receptors (cell_id) with both chains of a pair (see CHAIN_PAIRS), + for --paired_only. + + Arguments: + path (Path): the downloaded AIRR TSV. + chunksize (int): rows per chunk while scanning it. + codes (frozenset): with --air_species, the wanted locus_species codes: + only chains of them count, since those are the only ones kept (see + _speciesCells). None counts every chain. + + Returns: + frozenset: the cell_id values to keep. + """ + loci = {} + for chunk in pandas.read_csv(path, sep='\t', dtype=str, na_filter=False, + usecols=['cell_id', 'locus', 'locus_species', + 'v_call', 'j_call'], + chunksize=chunksize): + if codes is not None: + chunk = chunk[chunk['locus_species'].isin(codes)] + for cell, locus in zip(chunk['cell_id'], _chainLoci(chunk)): + loci.setdefault(cell, set()).add(locus) + + kept = frozenset(cell for cell, found in loci.items() + if cell and any(one in found and found & other + for one, other in CHAIN_PAIRS)) + log.info(' paired only: %s receptors with both chains kept, %s with one dropped', + format(len(kept), ','), format(len(loci) - len(kept), ',')) + + return kept + + +def _speciesCodes(name): + """ + The locus_species codes an --air_species name stands for: the species' + own code and those of its strains. + + Arguments: + name (str): a name from RECEPTOR_SPECIES. + + Returns: + frozenset: the matching codes. + """ + return frozenset(code for code, (own, parent) in RECEPTOR_SPECIES.items() + if own == name + or (parent is not None and RECEPTOR_SPECIES[parent][0] == name)) + + +def _speciesCells(path, chunksize, codes): + """ + The receptors (cell_id) of a receptor AIRR table to keep: those with a + chain of the wanted species and no chain of another, decided per + receptor, so a pair with a chain of another species is dropped whole. A + chain with no species recorded does not count against its receptor, but + is not kept either: normalizeChunk drops it, so asking for a species + gives only chains IEDB says are of it, and a heavy chain whose light + chain has no species recorded is kept alone. + + Arguments: + path (Path): the downloaded AIRR TSV. + chunksize (int): rows per chunk while scanning it. + codes (frozenset): the wanted locus_species codes (see _speciesCodes). + + Returns: + frozenset: the cell_id values to keep. + """ + wanted, other, blank = set(), set(), [] + for chunk in pandas.read_csv(path, sep='\t', dtype=str, na_filter=False, + usecols=['cell_id', 'locus_species'], + chunksize=chunksize): + species = chunk['locus_species'] + matched = species.isin(codes) + wanted.update(chunk.loc[matched, 'cell_id']) + other.update(chunk.loc[~matched & (species != ''), 'cell_id']) + blank.extend(chunk.loc[species == '', 'cell_id']) + kept = frozenset(wanted - other) + log.info(' receptor species: %s receptors kept, %s with a chain of another ' + 'species dropped, %s chains with no species recorded dropped', + format(len(kept), ','), format(len(wanted & other), ','), + format(sum(cell in kept for cell in blank), ',')) + + return kept + + + +def _findZipMember(archive, filename): + """ + Locate a member of the AIRR bulk ZIP by its filename. + + Arguments: + archive (zipfile.ZipFile): the opened archive. + filename (str): the filename to find, ignoring any directory prefix. + + Returns: + str: the matching member name. + + Raises: + IedbParseError: if no member has that filename. + """ + for member in archive.namelist(): + if Path(member).name == filename: + return member + + raise IedbParseError( + "'%s' not found in %s; the AIRR bulk export layout has changed" + % (filename, BULK_URL)) + + +def cacheDir(): + """ + Where IEDB's bulk ZIP is cached between runs: $XDG_CACHE_HOME/immvoke/iedb, + ~/.cache/immvoke/iedb when that is not set. + + Returns: + Path: the cache folder (not created here). + """ + root = os.environ.get('XDG_CACHE_HOME') or Path.home() / '.cache' + + return Path(root) / 'immvoke' / 'iedb' + + +class IedbSource(SourceBase): + """ + The IEDB specificity source. + """ + + name = 'iedb' + description = ('IEDB: curated B-cell and T-cell receptor sequences with ' + 'their epitopes and assays') + homepage = 'https://www.iedb.org/' + collections = ('bcr', 'tcr') + collection_help = { + 'bcr': 'BCR (antibody) receptor sequences, from the AIRR bulk export', + 'tcr': 'TCR receptor sequences, from the same AIRR bulk export', + } + airr_collections = frozenset(collections) + #: The tables extracted from the bulk ZIP are deleted once iedb_bcr.tsv + #: and iedb_tcr.tsv are written; the ZIP itself stays in the cache. + keep_raw = False + filter_help = { + 'air_species': ('SPECIES', "the receptor's species, from each chain's " + 'locus_species; a species includes its ' + 'strains'), + 'antigen_species': ('ORGANISM', 'the organism its epitope comes from: ' + 'any name IEDB uses, ignoring case; an ' + 'organism includes those under it'), + 'antigen_structure_type': ('TYPE', "its epitope's structure, IEDB's " + 'structure_type; Discontinuous ' + 'peptide includes its multi chain ' + 'kind'), + 'assay_type': ('ASSAY', 'the type of one of its assays: any type IEDB ' + 'uses, ignoring case'), + 'disease': ('DISEASE', 'the disease one of its assays studied: any name ' + 'IEDB uses, ignoring case; a disease includes ' + 'those under it'), + 'outcome': ('OUTCOME', 'the assay outcome, iedb_assay_outcome'), + } + #: On/off filters, outside the schema since they take no value; see + #: Search.addFilterGroup and validateQuery. + filter_switches = { + 'paired_only': 'only receptors with both chains of a pair: a heavy and ' + 'a light chain (IGH with IGK or IGL), a beta and an ' + 'alpha, or a delta and a gamma', + } + search_columns = dict.fromkeys(collections, ( + ('iedb_receptor_group', 'receptor groups', 'distinct'), + ('reactivity_ref', 'epitopes', 'distinct'), + ('locus', 'locus', 'top'), + ('locus_species', 'receptor species', 'top'), + (ANTIGEN_SPECIES_COL, 'antigen species', 'top'), + (ANTIGEN_STRUCTURE_COL, 'antigen structure type', 'top'), + ('iedb_assay_outcome', 'outcome', 'top'), + )) + + #: TODO: confirm exact license wording against https://www.iedb.org/about + #: before this is relied on for redistribution terms. + license = 'Freely available; see https://www.iedb.org/about' + citation = ( + 'Vita R, Mahajan S, Overton JA, Dhanda SK, Martini S, Cantrell JR, ' + 'Wheeler DK, Sette A, Peters B. The Immune Epitope Database (IEDB): ' + '2018 update. Nucleic Acids Res. 2019;47(D1):D339-D343. ' + 'doi:10.1093/nar/gky1006', + ) + + def harvestSchema(self): + """ + Build the hand curated snapshot. + + There is no live form to scrape, so this returns the same field + definitions packaged in schema.yaml, timestamped now. `immvoke + schema refresh --source iedb` exists mainly so the timestamp and any + future edits to this method stay reproducible the same way a real + harvest would. + + Returns: + SourceSchema: the snapshot. + """ + from datetime import datetime + + from immvoke.Schema import Collection, Field, SourceSchema + from immvoke.Version import __version__ + + air_species = Field(name='air_species', values=_AIR_SPECIES) + outcome = Field(name='outcome', values=_RECEPTOR_OUTCOMES) + antigen_species = Field(name='antigen_species', + values=_ANTIGEN_SPECIES_EXAMPLES, open_values=True) + antigen_structure_type = Field(name='antigen_structure_type', + values=_ANTIGEN_STRUCTURE_TYPES) + assay_type = Field(name='assay_type', values=_ASSAY_TYPE_EXAMPLES, + open_values=True) + disease = Field(name='disease', values=_DISEASE_EXAMPLES, open_values=True) + receptor_fields = (air_species, antigen_species, antigen_structure_type, + assay_type, disease, outcome) + collections = { + 'bcr': Collection(name='bcr', fields=receptor_fields), + 'tcr': Collection(name='tcr', fields=receptor_fields), + } + + return SourceSchema( + source=self.name, + harvested=datetime.now(UTC).strftime('%Y-%m-%dT%H:%M:%SZ'), + harvested_by=('immvoke %s (hand curated, not scraped: IEDB ' + 'exposes no HTML form)' % __version__), + source_urls={'airr_bulk': BULK_URL, 'api': API_BASE}, + collections=collections) + + def validateQuery(self, collection, filters): + """ + Validate filters against the stored snapshot, letting the on/off + filters of filter_switches through: they take no value for the + schema to check. + + Arguments: + collection (str): the table being searched. + filters (dict): user supplied filters. + + Returns: + Query: the resolved query, a switch given as True. + """ + switches = {k: True for k in self.filter_switches if filters.get(k)} + query = super().validateQuery( + collection, {k: v for k, v in filters.items() if k not in self.filter_switches}) + + return Query(collection=collection, filters={**query.filters, **switches}) + + def searchUnits(self, query): + """ + Resolve a query to the one DataUnit its table represents. + + A table is not an index of many files; the table itself is the unit, + so there is always exactly one. + + Arguments: + query (Query): the validated request. + + Returns: + list: a single DataUnit. + """ + return [DataUnit(unit_id=_TABLE_MEMBER[query.collection], + collection=query.collection, url=BULK_URL, + metadata=dict(query.filters))] + + def labelValue(self, column, value): + """ + Name a locus_species code as its species (see RECEPTOR_SPECIES), and + a blank cell as not recorded. + + Arguments: + column (str): the column. + value (str): one of its values. + + Returns: + str: the label to show. + """ + if value == '': + return '(not recorded)' + if column == 'locus_species': + return RECEPTOR_SPECIES.get(value, (value, None))[0] + + return value + + def fetchUnit(self, unit, outdir, resume=True, progress=True): + """ + Fetch the shared AIRR bulk ZIP and extract one table from it. + + The ZIP is kept in the user's cache (see cacheDir), not under outdir, + so download and search share it, and fetching bcr and then tcr + downloads it only once. It is reused for CACHE_MAX_AGE, then fetched + again. + + Arguments: + unit (DataUnit): the bcr or tcr table. + outdir (Path): the mirror root, where the table is extracted to. + resume (bool): continue an interrupted ZIP transfer if possible. + progress (bool): show a progress bar for the ZIP transfer. + + Returns: + DownloadResult: the outcome, describing the extracted table file. + """ + zip_dest = cacheDir() / BULK_ZIP + age = time.time() - zip_dest.stat().st_mtime if zip_dest.exists() else None + if age is not None and age > CACHE_MAX_AGE: + log.info(" IEDB export: the cached copy is over a day old; " + "downloading IEDB's current one") + zip_dest.unlink() + age = None + elif age is None: + log.info(' IEDB export: downloading %s', BULK_ZIP) + if age is None: + self.client.fetch(unit.url, zip_dest, resume=resume, progress=progress) + else: + log.info(' IEDB export: using the copy downloaded %s ago (%s)', + humanDuration(age), humanPath(zip_dest)) + + dest = self.resolveOutputPath(unit, outdir) + dest.parent.mkdir(parents=True, exist_ok=True) + with zipfile.ZipFile(zip_dest) as archive: + member = _findZipMember(archive, unit.unit_id) + with archive.open(member) as source, open(dest, 'wb') as target: + target.write(source.read()) + + return DownloadResult(unit=unit, path=dest, sha256=hashFile(dest), + size_bytes=dest.stat().st_size) + + def readUnit(self, path, unit, chunksize=50000): + """ + Open a downloaded table. + + Arguments: + path (Path): the downloaded AIRR TSV. + unit (DataUnit): what it is. + chunksize (int): rows per chunk. + + Returns: + tuple: (metadata, iterator of raw record chunks). metadata holds + what normalizeChunk keeps: `kept_cells`, the receptors (see + _speciesCells, _pairedCells), `kept_epitopes`, the epitopes (see + _antigenEpitopes, _structureEpitopes) and `kept_assays`, the + assays (see _typeAssays, _diseaseAssays), each None to keep + everything when none of its filters was given and the records all + given filters agree on otherwise; `species_only`, whether + --air_species was given, so chains with no species recorded are + dropped; and `antigens`, each epitope's antigen, source species + and structure type (see _epitopeAntigens). The table is read as + text, blank cells as '' rather than NaN. + """ + filters = {k: v for k, v in (unit.metadata or {}).items() if v not in (None, '*')} + species, antigen = filters.get('air_species'), filters.get('antigen_species') + structure = filters.get('antigen_structure_type') + assay_type, disease = filters.get('assay_type'), filters.get('disease') + + records = _epitopeRecords(self.client, path, chunksize) + kept_epitopes = _intersect( + _antigenEpitopes(records, antigen) if antigen else None, + _structureEpitopes(records, structure) if structure else None) + + kept_assays = None + if assay_type or disease: + assays = _assayRecords(self.client, path, chunksize, unit.collection) + kept_assays = _intersect( + _typeAssays(assays, assay_type) if assay_type else None, + _diseaseAssays(assays, disease) if disease else None) + + codes = _speciesCodes(species) if species else None + kept_cells = _intersect( + _speciesCells(path, chunksize, codes) if species else None, + _pairedCells(path, chunksize, codes) if filters.get('paired_only') else None) + + metadata = { + 'species_only': species is not None, + 'kept_cells': kept_cells, + 'kept_epitopes': kept_epitopes, + 'kept_assays': kept_assays, + 'antigens': _epitopeAntigens(records), + } + + return metadata, pandas.read_csv(path, sep='\t', dtype=str, + na_filter=False, chunksize=chunksize) + + def normalizeChunk(self, metadata, chunk, unit, offset, report): + """ + Normalize a chunk, applying provenance and the filters. + + bcr and tcr are already AIRR rearrangement records, one per chain, + and pass through unchanged, less the rows the filters drop, plus + ANTIGEN_COL, ANTIGEN_SPECIES_COL and ANTIGEN_STRUCTURE_COL after + iedb_epitope_name. Those + come from IEDB's API rather than its export, so the row hash leaves + them out. + + Arguments: + metadata (dict): from readUnit: `species_only`, `kept_cells`, + `kept_epitopes`, `kept_assays` and `antigens`. + chunk (pandas.DataFrame): raw records, IEDB's own columns. + unit (DataUnit): what they came from. + offset (int): unused; kept for interface parity with SourceBase. + report (dict): counters to accumulate into. + + Returns: + pandas.DataFrame: the normalized chunk, with provenance columns + added. + """ + report['rows_in'] += len(chunk) + frame = chunk + if metadata['kept_cells'] is not None: + frame = frame[frame['cell_id'].isin(metadata['kept_cells'])] + if metadata['species_only']: + frame = frame[frame['locus_species'] != ''] + if metadata['kept_epitopes'] is not None: + frame = frame[frame['reactivity_ref'].isin(metadata['kept_epitopes'])] + if metadata['kept_assays'] is not None: + kept = metadata['kept_assays'] + frame = frame[pandas.Series([any(x.strip() in kept for x in ids.split(',')) + for ids in frame['iedb_assay_ids']], + index=frame.index, dtype=bool)] + outcome = (unit.metadata or {}).get('outcome') + if outcome and outcome != '*': + frame = frame[frame['iedb_assay_outcome'].str.lower() == outcome.lower()] + report['rows_out'] += len(frame) + frame = coerceAirrTypes(frame.copy()) + + row_hash = rowHashes(frame) + + frame = frame.copy() + found = frame['reactivity_ref'].map(metadata['antigens']) + at = frame.columns.get_loc('iedb_epitope_name') + 1 + frame.insert(at, ANTIGEN_COL, [x[0] if isinstance(x, tuple) else '' + for x in found]) + frame.insert(at + 1, ANTIGEN_SPECIES_COL, [x[1] if isinstance(x, tuple) else '' + for x in found]) + frame.insert(at + 2, ANTIGEN_STRUCTURE_COL, [x[2] if isinstance(x, tuple) else '' + for x in found]) + frame['immvoke_source'] = self.name + frame['immvoke_collection'] = unit.collection + frame['immvoke_unit_id'] = unit.unit_id + frame['immvoke_row_hash'] = row_hash + + return frame diff --git a/src/immvoke/Sources/Specificity/Paginate.py b/src/immvoke/Sources/Specificity/Paginate.py new file mode 100644 index 0000000..cc4fabb --- /dev/null +++ b/src/immvoke/Sources/Specificity/Paginate.py @@ -0,0 +1,74 @@ +""" +Range-header pagination, for IEDB + +IEDB's query API (https://query-api.iedb.org/api/v1) is PostgREST, which pages +results through the HTTP `Range` request header rather than a query parameter, +and signals the end of data with a 416 status rather than an empty page. Kept +here rather than under Sources/ directly since only specificity sources page +this way so far. +""" + +# Info +__author__ = 'Pramod Shinde' + +# Imports +import logging + +# Immvoke imports +from immvoke.Exceptions import HttpError + +log = logging.getLogger(__name__) + + +def pageByRange(client, url, page_size=2000, params=None, headers=None): + """ + Page a Range-header paginated endpoint. + + Arguments: + client (HttpClient): the shared HTTP client. + url (str): the endpoint's absolute URL. + page_size (int): rows requested per page. + params (dict): query parameters sent with every request, e.g. PostgREST + column filters. + headers (dict): additional headers sent with every request. + + Yields: + list: one page of parsed JSON records. A page shorter than page_size, + or a Content-Range total that has been reached, ends iteration. + + Raises: + HttpError: if a page comes back with a status other than 200, 206 or 416. + """ + base_headers = dict(headers or {}) + start = 0 + fetched = 0 + + while True: + end = start + page_size - 1 + page_headers = {**base_headers, 'Range-Unit': 'items', + 'Range': '%d-%d' % (start, end)} + response = client.get(url, params=params, headers=page_headers) + + # 416 past the last row is how this style of API says "no more data". + if response.status_code == 416: + return + if response.status_code not in (200, 206): + raise HttpError('GET %s returned HTTP %d' + % (url, response.status_code)) + + batch = response.json() + if batch: + yield batch + fetched += len(batch) + + # Content-Range: items 0-1999/42000 reports the total row count, when + # the endpoint sends it, so pagination can stop as soon as every row + # has arrived rather than always waiting for a short final page. + content_range = response.headers.get('Content-Range', '') + total_text = content_range.rsplit('/', 1)[-1] if '/' in content_range else '' + if total_text.isdigit() and fetched >= int(total_text): + return + + if len(batch) < page_size: + return + start += page_size diff --git a/src/immvoke/Sources/Specificity/Provenance.py b/src/immvoke/Sources/Specificity/Provenance.py new file mode 100644 index 0000000..b7637d9 --- /dev/null +++ b/src/immvoke/Sources/Specificity/Provenance.py @@ -0,0 +1,86 @@ +""" +Download provenance for specificity sources + +A separate writer from immvoke.Provenance.writeDownloadMetadata: a +specificity download represents a database's current state rather than one +more addition to a growing repertoire dataset, so each run replaces the +record instead of merging into it. Kept apart from Provenance.py itself so +this does not change what repertoire sources rely on. +""" + +# Info +__author__ = 'Pramod Shinde' + +# Imports +from pathlib import Path + +import yaml + +# Immvoke imports +from immvoke.Provenance import DOWNLOAD_METADATA, METADATA_VERSION, commandLine, timestamp +from immvoke.Version import __version__ + + +def writeSpecificityMetadata(out, source, table, filters, limit, units, + schema=None, license=None, citation=None, + formats=('tsv',)): + """ + Write the provenance record for a specificity download directory. + + Unlike Provenance.writeDownloadMetadata, this always replaces the file + rather than merging into whatever was already there. + + Arguments: + out (Path): the download root. The file is written at its top level. + source (str): the specificity database name. + table (str): the table (or 'all') downloaded. + filters (dict): the filters the user supplied. + limit (int): the table cap, or None. + units (list): unit records from Provenance.buildUnitRecord. + schema (SourceSchema): the snapshot the query resolved against, if known. + license (str): the source's data license, if known. Recorded so a reader + of this directory alone, with no access to immvoke's own docs, still + knows the terms the data was obtained under. + citation (tuple): the source's requested citation(s), if known, for the + same reason. + formats (tuple): the normalized shape(s) written this run: 'tsv' for a + source-column dump, 'airr' where the source's collection is AIRR + shaped (see SourceBase.airr_collections). A run spanning several + tables can produce both. + + Returns: + Path: the file written. + """ + out = Path(out) + out.mkdir(parents=True, exist_ok=True) + path = out / DOWNLOAD_METADATA + + run = { + 'finished': timestamp(), + 'command': commandLine(), + 'collection': table, + 'filters': dict(filters or {}), + 'limit': limit, + 'formats': sorted(set(formats)), + 'units': len(units), + } + if schema is not None: + run['schema_harvested'] = schema.harvested + run['schema_harvested_by'] = schema.harvested_by + + record = { + 'immvoke_metadata_version': METADATA_VERSION, + 'source': source, + 'generated_by': 'immvoke %s' % __version__, + } + if license is not None: + record['data_license'] = license + if citation: + record['data_citation'] = list(citation) + record['runs'] = [run] + record['units'] = list(units) + + with open(path, 'w') as handle: + yaml.safe_dump(record, handle, sort_keys=False, default_flow_style=False) + + return path diff --git a/src/immvoke/Sources/Specificity/README.md b/src/immvoke/Sources/Specificity/README.md new file mode 100644 index 0000000..13d8c57 --- /dev/null +++ b/src/immvoke/Sources/Specificity/README.md @@ -0,0 +1,255 @@ +# Specificity annotations + +Annotate your own receptors with antigen specificity from public databases: +download a reference table in AIRR format, then match your receptors against +it. How the matching works, step by step, is in [idea.md](idea.md). + +## IEDB: download + +[IEDB](https://www.iedb.org/) (`Iedb.py`) publishes its BCR (antibody) +receptors in AIRR format, as `bcr_full_airr.tsv` inside its bulk +`airr_full.zip` export, refreshed weekly. `immvoke specificity iedb download +bcr` downloads that table, keeping **every receptor species** unless +`--air_species` names one. With `--air_species "Homo sapiens"`, a receptor +is kept when it has a human chain (`locus_species` `NCBITAXON:9606`) and no +chain of another species, and only its human chains are written: a chain +with no species recorded is left out (`search` would show it as +`(not recorded)`), so a heavy chain whose light chain IEDB gave no species +is kept alone. A species includes its strains: +`"Mus musculus"` also keeps BALB/c, C57BL/6, Xenomouse and the rest, while +`"Mus musculus BALB/c"` keeps that strain alone. The names are IEDB's; +`immvoke specificity iedb download bcr --help` lists them all. Likewise +every assay outcome is kept unless `--outcome` names one (`Positive` or +`Negative`, matched against `iedb_assay_outcome`), and every antigen unless +`--antigen_species` names its organism: any name IEDB uses, ignoring case, +with or without its bracketed common name (`"Homo sapiens"` for +`Homo sapiens (human)`). An organism includes those under it, so +`"Influenza A virus"` keeps every flu strain. The AIRR export does not say +which organism an epitope comes from, so this asks IEDB's API about the +table's epitopes at download time (about 15 seconds for bcr); an unknown name fails, +naming the closest IEDB names. Every epitope structure is kept unless +`--antigen_structure_type` names one: `Linear peptide`, `Discontinuous +peptide` (which includes `Discontinuous peptide on multi chain`) or +`Non-peptidic` (haptens, glycans, lipids...) -- IEDB's `structure_type`, +the same choice as its own search's Epitope Structure, looked up in the +same API call. `--assay_type` keeps records with an assay of that type +(`neutralization`, `3D structure`, `qualitative binding`, ... -- the first +part of IEDB's assay name), and `--disease` those with an assay studying +that disease or one under it (`"viral infectious disease"` takes in +COVID-19, influenza and the rest), both any name IEDB uses, ignoring case; +given both, one assay must satisfy both. They ask IEDB's API about the +table's assays, adding about 40 seconds, and only when given. IEDB records +no disease for most assays, so for a pathogen `--antigen_species` finds +more (SARS-CoV2 neutralizing human antibodies: 63 receptor groups by +`--antigen_species SARS-CoV2`, 39 by `--disease COVID-19`). `--paired_only` +keeps only receptors with both chains of a pair -- heavy with kappa or +lambda, beta with alpha, delta with gamma -- naming a chain IEDB gave no +locus by its V or J gene. The rest is +kept as is -- one +rearrangement record per chain, linked by `cell_id`, +built from IEDB's calculated (IMGT standardized) gene calls and CDRs, with +the epitope and assay context (`iedb_epitope_name`, `iedb_assay_ids`, ...) +alongside -- and writes it as `iedb_bcr.tsv` with an `iedb_bcr.airr_validation.log` report next to it. + +```bash +immvoke specificity iedb download bcr --air_species "Homo sapiens" \ + --antigen_species "Influenza A virus" --antigen_structure_type "Discontinuous peptide" \ + --outcome Positive --outdir databases + +# neutralizing, paired human antibodies +immvoke specificity iedb download bcr --air_species "Homo sapiens" \ + --assay_type neutralization --paired_only --outdir databases +``` + +`search` takes the same filters and counts what a download would give, +writing nothing (unless `--out` asks for the counts as a TSV): rows, +receptor groups, epitopes, and the most common loci, receptor species, +antigen species, antigen structure types and outcomes (`--top` of each, 10 +by default). + +```bash +immvoke specificity iedb search bcr --air_species "Homo sapiens" +``` + +``` +iedb bcr: BCR (antibody) receptor sequences, from the AIRR bulk export +filters: air_species = Homo sapiens + + rows 10,861 of 18,650 + receptor groups 3,417 + epitopes 3,128 + + locus + 5,917 IGH + 2,920 IGK + 1,747 IGL + 277 (not recorded) +... +``` + +IEDB's bulk ZIP is kept in `~/.cache/immvoke/iedb/` (`$XDG_CACHE_HOME` if +set) for a day, so a search, a download, and `bcr` then `tcr`, share one +copy; after a day the current export is fetched again. + +Every download lands in its database's own folder under `--outdir`, so +more databases can share one outdir without mixing their files: + +``` +databases/ +└── iedb/ + ├── iedb_bcr.tsv # the reference table for annotate + ├── iedb_bcr.airr_validation.log # AIRR validation report + └── download_metadata.yml # source, license, citation, when, command +``` + +Each record also gets its epitope's antigen, that antigen's source +species and the epitope's structure type, `iedb_antigen`, +`iedb_antigen_source_species` and `iedb_antigen_structure_type` (after +`iedb_epitope_name`), which IEDB's AIRR export lacks: the download looks the +table's epitopes up on IEDB's API (`curated_source_antigens`, several +` | ` joined, and `structure_type`), adding about 15 seconds. + +The written `databases/iedb/iedb_bcr.tsv` keeps every column IEDB provides and +always carries every field the AIRR schema requires; one IEDB lacks is +written blank. + +## Annotate + +`immvoke annotate` matches a query AIRR table -- your own receptors -- +against a reference AIRR table, such as the `databases/iedb/iedb_bcr.tsv` +downloaded above, and reports which reference rows each query row matches. +The reference must be an AIRR rearrangement table; see +[idea.md](idea.md#input-contract) for exactly what that requires. From +Python, `immvoke.Sources.Specificity.Annotate.annotateSpecificity()` takes +the same options, as file paths or DataFrames, returns the result table, +and, given `outdir`, writes the same files the command does. + +### Options + +| `immvoke annotate` | `annotateSpecificity()` | Values | Meaning | +| --- | --- | --- | --- | +| `--query`, `--reference` | `query`, `reference` | AIRR TSV (or DataFrame in Python) | the tables to match | +| `--sequence-region` | `sequence_region` | `CDR3H` (default), `VDJ`, `CDRs` | compare the heavy chain CDR3 (`cdr3_aa`); `sequence_aa`; or CDR1-3 pooled | +| `--partition-by` | `partition_by` | unset (default), `V`, `VJ`, `VJL` | only compare pairs with the same V-gene; V- and J-gene; or V-gene, J-gene and CDR3 length (same gene, any allele) | +| `--identity-threshold` | `identity_threshold` | unset (exact), or a fraction in (0, 1] | e.g. `0.9` allows up to 10% of the query's positions to differ | +| `--distance-metric` | `distance_metric` | `hamming` (default), `levenshtein` | substitutions only, or substitutions, insertions and deletions | +| `--paired` | `paired` | flag | match heavy and light chain of the same reference antibody | +| `--count-field` | `count_field` | an AIRR count column; `duplicate_count` (default) | `--paired`: which of a query cell's several heavy (or light) chains to keep -- the highest count | +| `--threads` | `threads` | `1` (default); `-1` for every core | symscan's threads | +| `--outdir` | `outdir` | folder, created if needed | where to write this run's files (see Output) | +| `--full-dist-matrix` | `full_dist_matrix` | flag | also write how close every pair the filters admit is, threshold or not (see Output) | + +`--partition-by` is off by default, but an unpartitioned CDR3 comparison at +a loose threshold is prone to coincidental hits, especially on short CDR3s, +so `VJ` is usually what you want there. + +### Examples + +```bash +# exact CDR3 match +immvoke annotate --query mine.tsv --reference databases/iedb/iedb_bcr.tsv \ + --outdir annotation/exact + +# exact full VDJ match within the same V- and J-gene +immvoke annotate --query mine.tsv --reference databases/iedb/iedb_bcr.tsv \ + --sequence-region VDJ --partition-by VJ --outdir annotation/vdj_exact + +# same V+J gene, CDR3 at 85% amino acid identity (Hamming distance) +immvoke annotate --query mine.tsv --reference databases/iedb/iedb_bcr.tsv \ + --partition-by VJ --identity-threshold 0.85 --outdir annotation/hamming_0.85 + +# CDR1-3 pooled at 90% identity, same V-gene, J-gene and CDR3 length +immvoke annotate --query mine.tsv --reference databases/iedb/iedb_bcr.tsv \ + --sequence-region CDRs --partition-by VJL --identity-threshold 0.9 \ + --outdir annotation/cdrs_0.9 + +# CDR3 at 90% identity allowing insertions/deletions +immvoke annotate --query mine.tsv --reference databases/iedb/iedb_bcr.tsv \ + --partition-by VJ --identity-threshold 0.9 --distance-metric levenshtein \ + --outdir annotation/levenshtein_0.9 + +# paired heavy+light match +immvoke annotate --query mine.tsv --reference databases/iedb/iedb_bcr.tsv \ + --paired --outdir annotation/paired +``` + +### Output + +Each run writes three files into its `--outdir`, kept apart from the +downloaded databases, each named after the reference's database and table, +`_
_`, and a fourth with `--full-dist-matrix`: + +``` +annotation/exact/ +├── iedb_bcr_annotation.tsv # one row per (query row, matching +│ # reference row): its best hits +├── iedb_bcr_annotation_summary.txt # the run and its files in prose +├── iedb_bcr_annotation_metadata.yml # the options, query, reference and its +│ # database, counts, version, time and command +└── iedb_bcr_annotation_full_matrix.txt # --full-dist-matrix: every compared pair +``` + +`_
_annotation_summary.txt` says in plain words what was +compared, what came of it and what each of the other files holds, so the +folder explains itself without reading the YAML. + +The database is read off the `download_metadata.yml` that `immvoke +specificity` writes beside the reference table, and the table off its file +name, so runs against different databases can share one `--outdir` without +overwriting each other. A reference without that record is named by its +file name alone (`myref.tsv` gives `myref_annotation.tsv`), and a DataFrame +passed from Python by `reference_`. + +`_
_annotation.tsv`, without `--paired`, is **long**: one row +per (query heavy chain row, matching reference row). A query row with +three hits has three rows; a query row with no hit has none, its count +kept in the metadata. Of a query's hits it keeps the **best**: the ones +with the highest identity, all of them when several tie. + +| Column | Meaning | +| --- | --- | +| `sequence_id` | the query's own `sequence_id`, so the table joins straight back to the query file (`query_row`, its position among the heavy chain rows, for a query without that column) | +| `identity` | the hit's identity, `1 - mismatches / query length` (an `X` counts as identical under `hamming`) | +| `hit_id` | the hit's reference `sequence_id` | +| `antigen` | its antigen (`iedb_antigen`) | +| `epitope` | its epitope (`iedb_epitope_name`) | +| `antigen_source_species` | the organism that antigen comes from (`iedb_antigen_source_species`) | + +One row per hit rather than a joined cell, because there is no separator +IEDB's own names leave free: a discontinuous epitope is named by its +residues (`G58, L59, S61`), and 1,576 of its epitope names contain a `;`. +IEDB also lists a receptor once per epitope it was tested against, all +under one `sequence_id`, so a hit standing for several epitopes gets a row +for each. A reference without IEDB's columns leaves those three blank. + +With `--paired`, one row per query antibody (`cell_id`): `classification` +(`both`, `heavy_only`, `light_only`, `conflicting` or `neither`), +`heavy_hit_ids`, `light_hit_ids` and `both_hit_ids`. + +`_
_annotation_full_matrix.txt`, with `--full-dist-matrix`, lists +**every pair the filters admit** and how close it is: one row per (query, +hit), by query and closest first -- `sequence_id`, `hit_id` and +`identity`. The annotation table answers what a query matched, so +`--identity-threshold` decides it; the matrix answers how close a query is +to everything it was compared with, so the threshold does not apply -- +a pair far below it still gets a row. `--partition-by` does apply: a +reference of another V or J gene was never compared, so it has no distance +to report, and under `hamming` neither has one of a different CDR3 length +(a Hamming distance needs equal lengths). Without `--partition-by` this is +every same-length pair of both tables -- 14.6 million rows for a +repertoire against IEDB's BCR table, against 28,671 under `VJ` -- so it is +written a block at a time rather than built in memory, and a run without +`--partition-by` warns. With `--paired`, each antibody's hits: `cell_id`, `hit_id` +(a reference `cell_id`) and `matched_chains` (`both`, `heavy` or `light`). +Nothing is reduced here: where the annotation table keeps a query's best +hits only, the matrix keeps every hit the threshold let through, weaker +ones included. + +A run writes the whole bundle, so running again into the same `--outdir` +without `--full-dist-matrix` removes the matrix the earlier run left, +rather than leaving it beside a table and a metadata record that have been +overwritten -- a 0.7 matrix next to an exact-match table describes a +comparison that no longer happened. + +Reference command: `immvoke specificity iedb download bcr --help`, +`immvoke annotate --help`. diff --git a/src/immvoke/Sources/Specificity/Report.py b/src/immvoke/Sources/Specificity/Report.py new file mode 100644 index 0000000..0737548 --- /dev/null +++ b/src/immvoke/Sources/Specificity/Report.py @@ -0,0 +1,238 @@ +""" +What `immvoke specificity download` and `immvoke annotate` tell the +user as they run + +Both say when they start and what they are about to do, report what they +are doing in between, and end with how long it took and what came of it: +a download with the rows kept and the counts the source lists in +search_columns, an annotation with how many query rows matched something. +It is all status, so it goes through logging like the rest of immvoke's +status output and --quiet silences it. +""" + +# Info +__author__ = 'Pramod Shinde' + +# Imports +import logging +from pathlib import Path + +log = logging.getLogger(__name__) + + +def humanDuration(seconds): + """ + A duration for people: '0.4 s', '8 s', '2 min 5 s', '3 h 10 min'. + + Arguments: + seconds (float): the duration. + + Returns: + str: the duration, to a tenth under a second (where rounding to the + second would say '0 s' for work that did happen), to the second under + an hour, to the minute above. + """ + if seconds < 1: + return '%.1f s' % seconds + seconds = int(round(seconds)) + if seconds < 60: + return '%d s' % seconds + if seconds < 3600: + return '%d min %d s' % divmod(seconds, 60) + + return '%d h %d min' % (seconds // 3600, seconds % 3600 // 60) + + +def humanPath(path): + """ + A path for people: under the home folder as '~/...', else as given. + + Arguments: + path (Path): the path. + + Returns: + str: the path to show. + """ + path = Path(path) + try: + return '~/%s' % path.resolve().relative_to(Path.home()) + except ValueError: + return str(path) + + +class Tally: + """ + Count a table's distinct values as its chunks stream past on their way + to the file, for the columns the source lists as 'distinct' in + search_columns, so the summary costs no second pass over the table. + + Arguments: + source (SourceBase): the specificity source. + table (str): the table being written. + """ + + def __init__(self, source, table): + self.columns = [(column, label) for column, label, how + in source.search_columns.get(table, ()) if how == 'distinct'] + self.seen = {column: set() for column, _ in self.columns} + + def wrap(self, chunks): + """ + Pass the chunks through, counting as they go. + + Arguments: + chunks (iterable): DataFrames of normalized records. + + Yields: + pandas.DataFrame: each chunk, unchanged. + """ + for chunk in chunks: + for column, _ in self.columns: + if column in chunk.columns: + self.seen[column].update(chunk[column].tolist()) + yield chunk + + def distinct(self): + """ + Returns: + list: (label, count) per column, blanks not counted. + """ + return [(label, len(self.seen[column] - {'', None})) + for column, label in self.columns] + + +def started(source, table, filters): + """ + Say a table's download has started, and with which filters. + + Arguments: + source (SourceBase): the specificity source. + table (str): the table. + filters (dict): the filter flags given, field to value. + """ + given = {k: v for k, v in filters.items() if v not in (None, '*')} + log.info('Download started: %s %s (%s)', source.name, table, + source.collection_help.get(table, '')) + log.info(' filters: %s', ', '.join('%s = %s' % x for x in given.items()) + if given else 'none, every record kept') + + +def completed(dest, seconds, rows_in, rows, tally, validation=None): + """ + Say a table's download has finished, with what it holds. + + Arguments: + dest (Path): the table written. + seconds (float): how long the table took. + rows_in (int): rows in the source's table, before filtering. + rows (int): rows written. + tally (Tally): the distinct counts, from the rows written. + validation (dict): from Convert.writeAirr, or None for a table that is + not AIRR. + """ + figures = [('rows', '%s of %s kept' % (format(rows, ','), format(rows_in, ',')))] + figures += [(label, format(count, ',')) for label, count in tally.distinct()] + if validation is not None: + invalid = validation['rows_invalid'] + verdict = ('every row valid' + if not invalid and validation['header_valid'] is not False + else '%s of %s rows invalid' % (format(invalid, ','), + format(validation['rows_checked'], ','))) + figures.append(('AIRR validation', '%s (%s)' % ( + verdict, dest.with_name(dest.stem + '.airr_validation.log').name))) + pad = max(len(label) for label, _ in figures) + 2 + + log.info('Download completed in %s: %s', humanDuration(seconds), dest) + for label, value in figures: + log.info(' %s%s', label.ljust(pad), value) + + +#: sequence_region -> what it compares, for annotationStarted. +_REGIONS = { + 'CDR3H': 'the heavy chain CDR3 (cdr3_aa)', + 'VDJ': 'the whole V(D)J amino acid sequence (sequence_aa)', + 'CDRs': 'CDR1, CDR2 and CDR3 pooled', +} + +#: partition_by -> what a pair must share before it is compared at all. +_PARTITIONS = { + None: 'any pair of rows is compared', + 'V': 'only rows with the same V gene are compared', + 'VJ': 'only rows with the same V and J gene are compared', + 'VJL': 'only rows with the same V gene, J gene and CDR3 length are compared', +} + + +def annotationStarted(query, reference, options): + """ + Say an annotation has started, and what it is about to compare. + + Arguments: + query (dict): the query's metadata record: `path` and `rows`. + reference (dict): the reference's, which may also name its `source` + database. + options (dict): the run's options, as annotateSpecificity records + them. + """ + def named(record, fallback): + # As the caller wrote it: resolving a relative path to an absolute + # one only makes a familiar argument harder to recognize. + return record['path'] or fallback + + log.info('Annotation started: %s rows of %s', format(query['rows'], ','), + named(query, 'a table passed from Python')) + log.info(' against %s rows of %s', format(reference['rows'], ','), + named(reference, 'a reference passed from Python')) + + if options['paired']: + log.info(' comparing: heavy and light chain of the same antibody, ' + 'same V gene, J gene and exact CDR3') + return + + threshold = options['identity_threshold'] + log.info(' comparing: %s', _REGIONS.get(options['sequence_region'], + options['sequence_region'])) + log.info(' a hit is: %s', 'an exact match' if threshold is None else + '%g%% identity or better, %s' % (threshold * 100, + options['distance_metric'])) + log.info(' %s', _PARTITIONS.get(options['partition_by'], + options['partition_by'])) + + +def annotationCompleted(seconds, results, paired): + """ + Say an annotation has finished, and what it found. + + Arguments: + seconds (float): how long it took. + results (dict): the run's counts, as annotateSpecificity records them. + paired (bool): whether this was a paired run. + """ + log.info('Annotation completed in %s', humanDuration(seconds)) + if paired: + figures = [('query antibodies', format(results['query_antibodies'], ','))] + figures += [(label, format(count, ',')) + for label, count in results['classification'].items()] + else: + figures = [ + ('query rows matched', + '%s of %s' % (format(results['query_rows_with_hits'], ','), + format(results['query_rows_compared'], ','))), + ('hits within the threshold', format(results['hits'], ',')), + ] + if 'full_matrix_rows' in results: + figures.append(('pairs in the full matrix', + format(results['full_matrix_rows'], ','))) + pad = max(len(label) for label, _ in figures) + 2 + for label, value in figures: + log.info(' %s%s', label.ljust(pad), value) + + +def finished(record): + """ + Say where the download's provenance was written. + + Arguments: + record (Path): the download_metadata.yml written. + """ + log.info('Provenance (source, license, citation, command): %s', record) diff --git a/src/immvoke/Sources/Specificity/RowHash.py b/src/immvoke/Sources/Specificity/RowHash.py new file mode 100644 index 0000000..625ef2a --- /dev/null +++ b/src/immvoke/Sources/Specificity/RowHash.py @@ -0,0 +1,45 @@ +""" +Row hashes, for specificity sources + +Every specificity source tags each normalized record with +immvoke_row_hash: a short content hash of its upstream columns, so a row can +be recognized across downloads. Kept here rather than in each source, since +they all hash the same way. +""" + +# Info +__author__ = 'Pramod Shinde' + +# Imports +import hashlib + +import pandas + + +def rowHashes(frame): + """ + Build a short content hash for each row of a frame. + + A row's key is its values as text joined by '|', the way + '|'.join(str(x) for x in row) spells it, so hashes do not change with + how they are computed. The values are turned to text a column at a time + rather than with apply(..., axis=1), which builds a Series per row and + took 45 of the 46 seconds of normalizing IEDB's 338,000 row tcr table. + + Arguments: + frame (pandas.DataFrame): upstream columns, before provenance columns + are added. + + Returns: + pandas.Series: the first 12 hex characters of each row's SHA-256 + digest, on the frame's index. + """ + if not len(frame) or not len(frame.columns): + return pandas.Series([], dtype=str, index=frame.index) + + # str() on each value, as the old row-wise key did; pandas' own + # astype(str) would leave missing values as NaN rather than 'nan'. + columns = [map(str, frame.iloc[:, i].tolist()) for i in range(len(frame.columns))] + + return pandas.Series([hashlib.sha256('|'.join(row).encode('utf-8')).hexdigest()[:12] + for row in zip(*columns)], index=frame.index, dtype=str) diff --git a/src/immvoke/Sources/Specificity/Search.py b/src/immvoke/Sources/Specificity/Search.py new file mode 100644 index 0000000..e7c2b30 --- /dev/null +++ b/src/immvoke/Sources/Specificity/Search.py @@ -0,0 +1,279 @@ +""" +`immvoke specificity search`, and the filter flags it shares with download + +search counts what a table holds once the filters are applied: its rows, +the distinct values of key columns (receptor groups, epitopes, ...) and the +most common values of others (locus, species, outcome, ...), as each source +lists them in search_columns. It fetches the table the way download does, +but into a temporary folder that is removed afterwards (IEDB's bulk ZIP +lands in its cache, so a download straight after reuses it), and writes +nothing else unless --out asks for the counts as a TSV. + +The filter flags of both actions come from the source's schema snapshot, +grouped under 'filters', each saying what it filters on (the source's +filter_help) and laying its values out in columns. +""" + +# Info +__author__ = 'Pramod Shinde' + +# Imports +import collections +import csv +import logging +import tempfile +import textwrap +from pathlib import Path + +# Immvoke imports +from immvoke.Commandline import CommonHelpFormatter + +log = logging.getLogger(__name__) + +#: Most common values shown per column unless --top says otherwise. +DEFAULT_TOP = 10 + +#: Marks a help line holding a value list, for SpecificityHelpFormatter to +#: lay out once it knows the width; VALUE_SEP separates values. +VALUE_LIST = '\x00' +VALUE_SEP = '\x1f' + +#: A short vocabulary is written inline rather than as an indented list. +INLINE_VALUES = 4 + + +class SpecificityHelpFormatter(CommonHelpFormatter): + """ + CommonHelpFormatter for the specificity table commands: it keeps a help + string's own line breaks, lays a filter's values out so that no value is + split across lines, and leaves '(default: None)' and '(default: False)' + off a flag that is simply off unless given. + """ + + def _split_lines(self, text, width): + lines = [] + for line in text.splitlines(): + if line.startswith(VALUE_LIST): + lines.extend(_flow(line[len(VALUE_LIST):].split(VALUE_SEP), width)) + else: + lines.extend(textwrap.wrap(line, width) or ['']) + + return lines + + def _get_help_string(self, action): + if action.default is None or action.default is False: + return action.help + + return super()._get_help_string(action) + + +def _flow(values, width): + """ + Pack values onto lines, comma separated, breaking only between values: + many species names have spaces, so ordinary wrapping would split them. + + Arguments: + values (list): the values, in order. + width (int): characters available. + + Returns: + list: the lines, each indented by two spaces. + """ + lines, line = [], '' + for i, value in enumerate(values): + piece = value + (',' if i < len(values) - 1 else '') + if line and len(line) + 1 + len(piece) > width - 2: + lines.append(' ' + line) + line = piece + else: + line = '%s %s' % (line, piece) if line else piece + if line: + lines.append(' ' + line) + + return lines + + +def addFilterGroup(parser, schema, source, table): + """ + Add a table's filter flags to its subcommand, as a 'filters' group. + + Each flag's placeholder and description come from the source's + filter_help, and its values from the schema snapshot: listed inline when + there are a few, in columns when there are more, as examples when the + source checks the value itself (Field.open_values). The source's + filter_switches, if any, follow as flags that take no value. + + Arguments: + parser (ArgumentParser): the table's subparser. + schema (SourceSchema): the source's snapshot, or None if none is + installed. + source (type): the source class, for its filter_help. + table (str): the table whose fields to add. + """ + if schema is None or table not in schema.collections: + return + fields = schema.getCollection(table).fields + if not fields: + return + + group = parser.add_argument_group( + 'filters', 'Each narrows the table down; leave one out to keep every value.') + for item in fields: + metavar, what = source.filter_help.get(item.name, ('VALUE', '')) + lines = [what] if what else [] + if item.pseudo_values: + lines.append('defined or undefined') + elif item.open_values: + lines.append('e.g. %s' % ', '.join('"%s"' % x for x in item.values)) + elif len(item.values) <= INLINE_VALUES: + lines.append('one of: %s' % ', '.join(item.values)) + else: + lines.append('%d values:' % len(item.values)) + lines.append(VALUE_LIST + VALUE_SEP.join(item.values)) + + group.add_argument(item.flag, dest='filter_%s' % item.name, + metavar=metavar, default=None, help='\n'.join(lines)) + for name, what in getattr(source, 'filter_switches', {}).items(): + group.add_argument('--%s' % name, dest='filter_%s' % name, action='store_true', + default=None, help=what) + + +def searchTable(source, table, filters, workdir): + """ + Count what one table holds once the filters are applied. + + Arguments: + source (SourceBase): the specificity source. + table (str): the table. + filters (dict): the filter flags given, field to value. + workdir (Path): where the table is fetched to. + + Returns: + dict: `table`, `filters` (those given), `rows_in` (before filtering), + `rows` (after), `distinct` (label -> count) and `top` (label -> + Counter of labelled values), from the source's search_columns. + """ + query = source.validateQuery(table, filters) + columns = source.search_columns.get(table, ()) + seen = {column: set() for column, _, how in columns if how == 'distinct'} + counts = {column: collections.Counter() for column, _, how in columns + if how == 'top'} + + rows_in = rows = 0 + for unit in source.searchUnits(query): + result = source.fetchUnit(unit, workdir) + _, chunks, report = source.convertUnit(result.path, unit) + for chunk in chunks: + rows += len(chunk) + for column in seen.keys() & set(chunk.columns): + seen[column].update(chunk[column].tolist()) + for column in counts.keys() & set(chunk.columns): + counts[column].update(chunk[column].tolist()) + rows_in += report['rows_in'] + + top = {} + for column, label, how in columns: + if how == 'top': + labelled = collections.Counter() + for value, count in counts[column].items(): + text = '' if value is None else str(value) + labelled[source.labelValue(column, text)] += count + top[label] = labelled + + return {'table': table, + 'filters': {k: v for k, v in query.filters.items() if v != '*'}, + 'rows_in': rows_in, 'rows': rows, + 'distinct': {label: len(seen[column] - {''}) + for column, label, how in columns if how == 'distinct'}, + 'top': top} + + +def formatSummary(source, summary, top=DEFAULT_TOP): + """ + Render a searchTable summary for the terminal. + + Arguments: + source (SourceBase): the source, for its name and table description. + summary (dict): from searchTable. + top (int): most common values shown per column. + + Returns: + str: the summary, ending in a newline. + """ + table = summary['table'] + lines = ['%s %s: %s' % (source.name, table, source.collection_help.get(table, ''))] + if summary['filters']: + lines.append('filters: %s' % ', '.join('%s = %s' % item + for item in summary['filters'].items())) + lines.append('') + + figures = [('rows', '%s of %s' % (format(summary['rows'], ','), + format(summary['rows_in'], ',')))] + figures += [(label, format(count, ',')) + for label, count in summary['distinct'].items()] + pad = max(len(label) for label, _ in figures) + 2 + lines += [' %s%s' % (label.ljust(pad), value) for label, value in figures] + + for label, counter in summary['top'].items(): + if not counter: + continue + shown = counter.most_common(top) + more = len(counter) - len(shown) + lines += ['', ' %s%s' % (label, ' (top %d of %d)' % (len(shown), len(counter)) + if more else '')] + digits = len(format(shown[0][1], ',')) + lines += [' %s %s' % (format(count, ',').rjust(digits), value) + for value, count in shown] + + return '\n'.join(lines) + '\n' + + +def writeSummaries(summaries, out): + """ + Write every count of the summaries as a TSV: table, field, value, rows. + + Arguments: + summaries (list): from searchTable. + out (Path): the TSV to write. + """ + out = Path(out) + out.parent.mkdir(parents=True, exist_ok=True) + with open(out, 'w', newline='') as handle: + writer = csv.writer(handle, delimiter='\t', lineterminator='\n') + writer.writerow(['table', 'field', 'value', 'rows']) + for summary in summaries: + table = summary['table'] + writer.writerow([table, 'rows', '', summary['rows']]) + for label, count in summary['distinct'].items(): + writer.writerow([table, label, '', count]) + for label, counter in summary['top'].items(): + for value, count in counter.most_common(): + writer.writerow([table, label, value, count]) + + +def searchTables(source, tables, filters, top=DEFAULT_TOP, out=None): + """ + Search one or more tables and print what each holds. + + Arguments: + source (SourceBase): the specificity source. + tables (list): the tables to search. + filters (dict): the filter flags given, field to value. + top (int): most common values shown per column. + out (Path): also write every count here as a TSV, or None. + + Returns: + int: process exit status. + """ + summaries = [] + with tempfile.TemporaryDirectory(prefix='immvoke-search-') as workdir: + for table in tables: + summary = searchTable(source, table, filters, Path(workdir)) + summaries.append(summary) + print(formatSummary(source, summary, top=top)) + + if out is not None: + writeSummaries(summaries, out) + log.info('wrote %s', out) + + return 0 diff --git a/src/immvoke/Sources/Specificity/Symscan.py b/src/immvoke/Sources/Specificity/Symscan.py new file mode 100644 index 0000000..a8a6501 --- /dev/null +++ b/src/immvoke/Sources/Specificity/Symscan.py @@ -0,0 +1,653 @@ +""" +symscan annotation: exact, Hamming and Levenshtein + + runSymscanAnnotation up to `max_distance` substitutions (distance_type + 'hamming', one or more pooled columns) or + substitutions/insertions/deletions combined + ('levenshtein', a single column); an exact match is + `max_distance` 0 + +symscan (https://github.com/yutanagano/symscan) builds a Rust +symmetric-deletion index and only ever looks at likely matches, rather than +scoring every query x reference pair -- so it stays fast as the reference +grows. It is the only engine that finds hits: no hit is searched for in +Python. (Distances.py scores every gated pair for a threshold run's pairs +file; the hits still come from here.) + +Blanks. A row with a blank value in any compared column has nothing to +compare, so it never matches: a blank query gets no hit, and a blank +reference is never hit. + +Gates. symscan has no gate of its own: `exact_cols` (exact string +agreement), each column's length for a pooled multi-column Hamming +comparison, and `gene_cols` (V/J-gene agreement, same gene and any allele, +see Genes.py) are applied to the pairs it returns. + +`gene_cols` is also applied before the search, as a partition: the +comparison runs once per gene key rather than once over everything (see +_geneGroups). The index cost grows combinatorially with the tolerance, and +in one index every query collides with references of every gene +combination -- collisions whose distance is computed only for the gate to +discard them. Grouping removes that work instead of paying for it: on a +21,412 row repertoire against IEDB at 0.7 identity, 57 seconds became 3. +The gate still runs afterwards and still decides what is kept, so the +answer does not depend on the partition being clever. + +The "X" wildcard. An X is an unresolved residue: under 'hamming' (and so +under exact matching) a position where either side has one never counts as a +difference. symscan has no wildcard and compares X as an ordinary letter, +so the X's are taken out of the comparison instead: references are grouped +by where their X's sit, and those positions are cut out of the group's +references and of every query of the same length before symscan sees them +-- a reference X then never counts, and the search runs at the real limit. +A query's own X's, rarer, are covered by slack (its limit raised by its X +count), and every returned pair involving an X is re-scored under the +wildcard rule, keeping only those within the real limit. Levenshtein has no +wildcard rule: an X there is an ordinary letter. + +Limits. symscan's index holds every deletion variant of every string on both +sides -- sum(C(length, k) for k <= distance) per string -- and it aborts the +whole process (an uncatchable Rust panic or failed allocation) once that +outgrows memory, or once the distance exceeds MAX_DISTANCE. Each call is +therefore kept within MAX_VARIANTS by grouping queries by their limit (and, +under 'hamming', by length, since only same-length references can match) +and chunking both the queries and the references of a group to half the +budget each. What cannot fit at all -- a limit above MAX_DISTANCE, X slack +included, or one sequence whose index alone exceeds half the budget (a very +long sequence at a loose threshold) -- is skipped, with a warning naming the +query rows that may therefore miss hits. + +symscan parallelizes via rayon's global pool, which reads RAYON_NUM_THREADS +once, on first use -- so `threads` only takes effect on the first call in a +process; later calls reuse that pool. +""" + +# Info +__author__ = 'Pramod Shinde' + +# Imports +import itertools +import logging +import math +import os +from functools import cache + +import numpy +import pandas +import symscan + +# Immvoke imports +from immvoke.Exceptions import AnnotationError +from immvoke.Sources.Specificity.Annotate import ( + DEFAULT_ID_COL, + WILDCARD, + _allowedDifferences, + _exactValues, + _formatPositions, + _mismatchPositions, + _rowValues, +) +from immvoke.Sources.Specificity.Genes import buildGeneGate, parseCall + +log = logging.getLogger(__name__) + +#: Largest max_distance symscan 0.8 accepts without panicking. +MAX_DISTANCE = 14 + +#: Most deletion variants, both sides together, one symscan call may index. +#: symscan's peak memory measured at roughly 13-40 bytes per variant, so this +#: keeps a call to around 1-2 GB. +MAX_VARIANTS = 40_000_000 + +#: RAYON_NUM_THREADS the rayon pool was started with, once one has been. +_pool_threads = None + + +@cache +def _deletionVariants(length, distance): + """ + Arguments: + length (int): a string's length. + distance (int): the distance it is indexed at. + + Returns: + int: how many deletion variants symscan's index holds for it. + """ + return sum(math.comb(length, k) for k in range(min(distance, length) + 1)) + + +def _resolveLimits(query_values, max_distance, threshold_is_fraction): + """ + Resolve the tolerance to an absolute distance per query row. + + Arguments: + query_values (list): per-query comparison column tuples. + max_distance (float): an absolute count, or a fraction of the query's + summed comparison column length when `threshold_is_fraction` is True. + threshold_is_fraction (bool): see runSymscanAnnotation. + + Returns: + numpy.ndarray: one int per query row. + """ + if threshold_is_fraction: + return numpy.array([_allowedDifferences(max_distance, + sum(len(v) for v in values)) + for values in query_values], dtype=numpy.int64) + + return numpy.full(len(query_values), int(max_distance), dtype=numpy.int64) + + +def _startPool(threads): + """ + Size rayon's global pool before symscan's first call in this process. + + Arguments: + threads (int): worker threads, -1 for every available core. + """ + global _pool_threads + wanted = os.cpu_count() if threads == -1 else threads + if _pool_threads is None: + os.environ['RAYON_NUM_THREADS'] = str(wanted) + _pool_threads = wanted + elif wanted != _pool_threads: + log.debug('symscan thread pool already started with %d threads; ' + 'threads=%d ignored', _pool_threads, threads) + + +def _keyCodes(query_keys, reference_keys): + """ + Encode each row's gate key as an integer, consistently across both + sides, so candidate pairs can be filtered with one vectorized compare. + + Arguments: + query_keys (list): one tuple per query row. + reference_keys (list): one tuple per reference row. + + Returns: + tuple: (query codes, reference codes), numpy int arrays. + """ + joined = ['\x1f'.join(map(str, key)) for key in query_keys + reference_keys] + codes, _ = pandas.factorize(pandas.Series(joined, dtype=object)) + + return codes[:len(query_keys)], codes[len(query_keys):] + + +def _without(seq, positions): + """ + Returns: + str: `seq` with the characters at `positions` (a set) cut out. + """ + return ''.join(c for i, c in enumerate(seq) if i not in positions) + + +def _describeRows(rows, queries, shown=5): + """ + Returns: + str: the first few query rows, for a log line -- by their own + `sequence_id` where the query has one, else by position. + """ + more = ', ...' if len(rows) > shown else '' + if DEFAULT_ID_COL in queries.columns: + ids = queries[DEFAULT_ID_COL].iloc[rows[:shown]].astype(str) + return '%s %s%s' % (DEFAULT_ID_COL, ', '.join(ids), more) + + return 'query_row %s%s' % (', '.join(str(int(r)) for r in rows[:shown]), more) + + +def _concat(parts): + """ + Arguments: + parts (list): (query positions, reference positions, distances) tuples. + + Returns: + tuple: the three arrays, each concatenated across `parts`. + """ + if not parts: + return tuple(numpy.array([], dtype=numpy.int64) for _ in range(3)) + + return tuple(numpy.concatenate(x).astype(numpy.int64) for x in zip(*parts)) + + +def _chunks(rows, costs, budget): + """ + Split rows into consecutive chunks whose summed cost stays within budget. + + Arguments: + rows (numpy.ndarray): positions, each costing no more than `budget`. + costs (list): each row's deletion-variant count, same order. + budget (int): the most one chunk may cost. + + Returns: + list: numpy int arrays of positions. + """ + chunks, chunk, total = [], [], 0 + for row, cost in zip(rows.tolist(), costs): + if chunk and total + cost > budget: + chunks.append(chunk) + chunk, total = [], 0 + chunk.append(row) + total += cost + if chunk: + chunks.append(chunk) + + return [numpy.array(c, dtype=numpy.int64) for c in chunks] + + +def _geneKeys(frame, gene_cols): + """ + Each row's gene keys: one per combination of the genes it calls, so a + row calling two V genes and one J gene gets two keys. + + Two rows pass the gene gate when they name a gene in common in every + gene column (Genes.callsAgree is set intersection), which is exactly + when their key sets intersect -- so grouping by key never separates a + pair the gate would keep. A row with a blank or unrecognized call in + any column gets no key at all, matching the gate's rule that such a + call agrees with nothing. + + Arguments: + frame (pandas.DataFrame): the rows. + gene_cols (list): gene call columns, e.g. ['v_call', 'j_call']. + + Returns: + list: one tuple of keys per row, each key a tuple of one gene per + column. + """ + parsed = [] + for col in gene_cols: + values = frame[col].astype(object).fillna('').astype(str).tolist() + # Each distinct call parsed once: a repertoire repeats its calls. + cache = {value: sorted(parseCall(value)) for value in set(values)} + parsed.append([cache[value] for value in values]) + + return [tuple(itertools.product(*genes)) for genes in zip(*parsed)] + + +def _geneGroups(queries, reference, gene_cols, query_rows, ref_rows): + """ + Split the comparison into one group per gene key, so symscan indexes + only the rows that could pass the gene gate. + + symscan's cost grows combinatorially with the tolerance, and in one + index every query collides with references of every gene combination -- + collisions the gate then throws away, after their distance has been + computed. Grouping first removes that work rather than paying for it. + The gate still runs afterwards: grouping decides what is compared, the + gate decides what is kept. + + Arguments: + queries (pandas.DataFrame): the query rows. + reference (pandas.DataFrame): the reference rows. + gene_cols (list): gene call columns to group on, or None for no + grouping -- then the whole comparison is one group. + query_rows (numpy.ndarray): the query positions to compare. + ref_rows (numpy.ndarray): the reference positions to compare. + + Yields: + tuple: (query positions, reference positions) per group, only where + both sides have rows. + """ + if not gene_cols: + yield query_rows, ref_rows + return + + ref_keys = _geneKeys(reference, gene_cols) + by_key = {} + for ref in ref_rows.tolist(): + for key in ref_keys[ref]: + by_key.setdefault(key, []).append(ref) + + query_keys = _geneKeys(queries, gene_cols) + wanted = {} + for query in query_rows.tolist(): + for key in query_keys[query]: + if key in by_key: + wanted.setdefault(key, []).append(query) + + for key, members in wanted.items(): + yield (numpy.array(members, dtype=numpy.int64), + numpy.array(by_key[key], dtype=numpy.int64)) + + +def _searchGroup(query_seqs, ref_seqs, query_lengths, limits, query_x, ref_x, + distance_type, query_rows, ref_rows): + """ + Run one group's comparison through symscan, handling the X wildcard. + + References are grouped by where their X's sit, and those positions are + cut out of the group's references and of every same-length query, so a + reference X never counts; a query's own X's are covered by slack + instead (see the module note). + + Arguments: + query_seqs (list): every query's compared sequence. + ref_seqs (list): every reference's. + query_lengths (numpy.ndarray): len() of each query sequence. + limits (numpy.ndarray): each query's allowed differences. + query_x (numpy.ndarray): each query's X count. + ref_x (numpy.ndarray): each reference's X count. + distance_type (str): 'hamming' or 'levenshtein'. + query_rows (numpy.ndarray): the query positions to compare. + ref_rows (numpy.ndarray): the reference positions to compare. + + Returns: + tuple: (list of (rows, cols, dists) pair blocks, list of skipped + (query positions, reference positions) blocks), all in the positions + of the whole table. + """ + found, skipped = [], [] + patterns = {} + for ref in ref_rows[ref_x[ref_rows] > 0].tolist(): + seq = ref_seqs[ref] + key = (len(seq), tuple(i for i, c in enumerate(seq) if c == WILDCARD)) + patterns.setdefault(key, []).append(ref) + + clean_refs = ref_rows[ref_x[ref_rows] == 0] + if len(clean_refs) and len(query_rows): + pairs, blocks = _symscanPairs(query_seqs, ref_seqs, limits + query_x, + distance_type, query_rows, clean_refs) + found.append(pairs) + skipped.extend(blocks) + for (length, wild), refs in patterns.items(): + members = query_rows[query_lengths[query_rows] == length] + if not len(members): + continue + cut = set(wild) + masked_queries = {q: _without(query_seqs[q], cut) for q in members.tolist()} + masked_refs = {r: _without(ref_seqs[r], cut) for r in refs} + pairs, blocks = _symscanPairs(masked_queries, masked_refs, limits + query_x, + distance_type, members, + numpy.array(refs, dtype=numpy.int64)) + found.append(pairs) + skipped.extend(blocks) + + return found, skipped + + +def _symscanPairs(query_seqs, ref_seqs, limits, distance_type, query_rows, + ref_rows): + """ + Every (query, reference) pair within its query's limit, via symscan, + with each call kept within MAX_VARIANTS and MAX_DISTANCE. + + Queries are grouped by limit (and, under Hamming, length), and both + sides of a group are chunked to half the budget each, so every call + fits however large the tables. What cannot fit at all is skipped: a + group whose limit exceeds MAX_DISTANCE, or a single sequence whose own + index exceeds half the budget. + + Arguments: + query_seqs (list): every query row's comparison string. + ref_seqs (list): every reference row's comparison string. + limits (numpy.ndarray): every query row's resolved distance limit. + distance_type (str): 'hamming' or 'levenshtein'. + query_rows (numpy.ndarray): the query positions to compare. + ref_rows (numpy.ndarray): the reference positions to compare them to. + + Returns: + tuple: ((query positions, reference positions, distances), skipped) + -- numpy int arrays indexing the full query and reference tables, and + a list of (query positions, reference positions) blocks that were not + compared. + """ + hamming = distance_type == 'hamming' + half = MAX_VARIANTS // 2 + query_lengths = numpy.array([len(query_seqs[i]) for i in query_rows], dtype=int) + ref_lengths = numpy.array([len(ref_seqs[i]) for i in ref_rows], dtype=int) + + # Under Hamming only same-length strings can match, so a group is one + # (limit, length) and indexes only the references of that length. + if hamming: + keys = numpy.stack([limits[query_rows], query_lengths], axis=1) + else: + keys = limits[query_rows][:, None] + groups = {} + for position, key in enumerate(map(tuple, keys)): + groups.setdefault(key, []).append(position) + + found, skipped = [], [] + for key, positions in groups.items(): + limit = int(key[0]) + members = query_rows[positions] + refs = ref_rows[ref_lengths == key[1]] if hamming else ref_rows + if len(refs) == 0: + continue + if limit > MAX_DISTANCE: + skipped.append((members, refs)) + continue + + # Python ints: a long sequence at a high distance overflows int64. + query_costs = [_deletionVariants(len(query_seqs[i]), limit) for i in members] + ref_costs = [_deletionVariants(len(ref_seqs[i]), limit) for i in refs] + query_fits = numpy.array([cost <= half for cost in query_costs], dtype=bool) + ref_fits = numpy.array([cost <= half for cost in ref_costs], dtype=bool) + fit_queries, fit_refs = members[query_fits], refs[ref_fits] + if not ref_fits.all(): + skipped.append((members, refs[~ref_fits])) + if not query_fits.all(): + skipped.append((members[~query_fits], fit_refs)) + + query_chunks = _chunks(fit_queries, [c for c, ok in zip(query_costs, query_fits) + if ok], half) + ref_chunks = _chunks(fit_refs, [c for c, ok in zip(ref_costs, ref_fits) if ok], + half) + for ref_chunk in ref_chunks: + subset_refs = [ref_seqs[i] for i in ref_chunk] + for chunk in query_chunks: + r, c, d = symscan.get_neighbors_across( + [query_seqs[i] for i in chunk], subset_refs, + max_distance=limit, distance_type=distance_type) + found.append((chunk[r.astype(numpy.int64)], + ref_chunk[c.astype(numpy.int64)], d.astype(numpy.int64))) + + return _concat(found), skipped + + +def runSymscanAnnotation(queries, reference, compare_cols, max_distance, + distance_type='hamming', id_col=DEFAULT_ID_COL, + exact_cols=None, threshold_is_fraction=False, threads=1, + gene_cols=None, return_pairs=False, warn_skipped=True): + """ + Annotate each query row against `reference` allowing up to + `max_distance` differences in `compare_cols`, via + symscan.get_neighbors_across. + + symscan has no gate of its own, so `exact_cols` (and, for a pooled + multi-column Hamming comparison, each column's length, so the columns + line up one by one) and `gene_cols` are applied to the pairs symscan + returns. `gene_cols` additionally partitions the search, which is only + an optimization: the gate decides what is kept either way (see the + module note). An 'X' matches + any residue under 'hamming'; see the module note. + + Arguments: + queries (pandas.DataFrame): the input rows to annotate. + reference (pandas.DataFrame): the fixed table being annotated against. + compare_cols (list): column names to compare -- pooled (concatenated) + when more than one, which only 'hamming' supports. + max_distance (float): the tolerance -- an absolute count, or a + fraction of the query's summed comparison column length when + `threshold_is_fraction` is True. + distance_type (str): 'hamming' (substitutions only) or 'levenshtein' + (substitutions, insertions and deletions combined). + id_col (str): column in `reference` holding each row's identifier. + exact_cols (list): column names that must agree exactly, e.g. + `['v_call']`. None for no such gate. + threshold_is_fraction (bool): treat `max_distance` as a fraction of + the query's length instead of a fixed count. + threads (int): rayon worker threads (-1 for every available core); + only the first call in a process sets this -- see the module note. + gene_cols (list): gene call columns that must name the same gene + (any allele), e.g. `['v_call', 'j_call']`; a blank or unrecognized + call agrees with nothing (see Genes.py). None for no such gate. + return_pairs (bool): also return the pairs behind `hit_ids`. + warn_skipped (bool): warn about the query rows beyond symscan's + limits; False for a caller that scores them some other way (see + Annotate._scoreSkippedRows). + + Returns: + pandas.DataFrame: one row per query row, columns `query_row`, + `n_hits_total`, `min_mismatches`, `is_unique_hit`, `hit_ids`, plus + `mismatch_positions` under 'hamming'. `min_mismatches` is the + closest hit's distance (edit distance under 'levenshtein'), and + `hit_ids` the hits' `id_col` values, closest first. Pairs beyond + symscan's limits are skipped with a warning (see the module note); + `attrs['skipped_query_rows']` counts the query rows that may + therefore miss hits, and `attrs['skipped_rows']` lists them. + With `return_pairs`, a tuple: that table, and one row per (query row, + hit id) in the same order -- `query_row`, `reference_id`, `mismatches`, + `identity` (1 - mismatches / the query's compared length, an X + counting as identical under 'hamming'), plus `mismatch_positions` + under 'hamming'. + + Raises: + AnnotationError: if `distance_type` is not recognized, or + 'levenshtein' is given several columns. + """ + if distance_type not in ('hamming', 'levenshtein'): + raise AnnotationError("unknown distance_type %r; choose 'hamming' or " + "'levenshtein'" % (distance_type,)) + hamming = distance_type == 'hamming' + if not hamming and len(compare_cols) != 1: + raise AnnotationError("distance_type='levenshtein' compares a single column; " + 'got %d' % len(compare_cols)) + + query_values = _rowValues(queries, compare_cols) + ref_values = _rowValues(reference, compare_cols) + query_seqs = [''.join(v) for v in query_values] + ref_seqs = [''.join(v) for v in ref_values] + limits = _resolveLimits(query_values, max_distance, threshold_is_fraction) + query_exact = _exactValues(queries, exact_cols) + ref_exact = _exactValues(reference, exact_cols) + + _startPool(threads) + # A blank value has nothing to compare: a row with one in any compared + # column never matches, on either side. + query_blank = numpy.array([any(v == '' for v in values) for values in query_values], + dtype=bool) + ref_blank = numpy.array([any(v == '' for v in values) for values in ref_values], + dtype=bool) + if query_blank.any() or ref_blank.any(): + # Only the side that has them, so a clean query says nothing. + blank = ', '.join('%s %s rows' % (format(int(n.sum()), ','), label) + for label, n in (('query', query_blank), + ('reference', ref_blank)) if n.any()) + log.info(' %s have a blank %s; they are not compared', blank, + '/'.join(compare_cols)) + all_queries = numpy.flatnonzero(~query_blank).astype(numpy.int64) + all_refs = numpy.flatnonzero(~ref_blank).astype(numpy.int64) + + # The X wildcard (Hamming only); see the module note. References are + # grouped by where their X's sit, and those positions are cut out of the + # group's references and of every same-length query, so a reference X + # never counts; a query's own X's are covered by slack instead. + if hamming: + query_x = numpy.array([seq.count(WILDCARD) for seq in query_seqs], + dtype=numpy.int64) + ref_x = numpy.array([seq.count(WILDCARD) for seq in ref_seqs], dtype=numpy.int64) + else: + query_x = numpy.zeros(len(query_seqs), dtype=numpy.int64) + ref_x = numpy.zeros(len(ref_seqs), dtype=numpy.int64) + # One search per gene key rather than one over everything (see + # _geneGroups): the gene gate below still decides what is kept, but + # pairs of different genes are no longer compared only to be discarded. + query_lengths = numpy.array([len(seq) for seq in query_seqs], dtype=int) + found, skipped = [], [] + for group_queries, group_refs in _geneGroups(queries, reference, gene_cols, + all_queries, all_refs): + pairs, blocks = _searchGroup(query_seqs, ref_seqs, query_lengths, limits, + query_x, ref_x, distance_type, + group_queries, group_refs) + found.extend(pairs) + skipped.extend(blocks) + rows, cols, dists = _concat(found) + skipped_queries = [] + if skipped: + skipped_queries = numpy.unique(numpy.concatenate([q for q, _ in skipped])) + skipped_refs = numpy.unique(numpy.concatenate([r for _, r in skipped])) + if skipped and warn_skipped: + log.warning( + '%d query rows could not be compared with %d reference rows, beyond ' + "symscan's limits (over %d differences, counting one per %r, or a " + 'sequence too long for its index at this threshold); those pairs ' + 'were skipped, so these query rows may miss hits: %s', + len(skipped_queries), len(skipped_refs), MAX_DISTANCE, WILDCARD, + _describeRows(skipped_queries, queries)) + + # Gate symscan's pairs: exact columns always, per-column lengths when + # several Hamming columns were pooled (the concatenated length alone + # could match while the individual columns don't line up). + query_keys, ref_keys = query_exact, ref_exact + if hamming and len(compare_cols) > 1: + query_keys = [k + tuple(map(len, v)) for k, v in zip(query_keys, query_values)] + ref_keys = [k + tuple(map(len, v)) for k, v in zip(ref_keys, ref_values)] + if exact_cols or (hamming and len(compare_cols) > 1): + query_codes, ref_codes = _keyCodes(query_keys, ref_keys) + keep = query_codes[rows] == ref_codes[cols] + rows, cols, dists = rows[keep], cols[keep], dists[keep] + + # Re-score every pair involving an X under the wildcard rule, and drop + # those the slack let in beyond the real limit. + wild = numpy.flatnonzero((query_x[rows] > 0) | (ref_x[cols] > 0)) + if len(wild): + dists = dists.copy() + dists[wild] = [len(_mismatchPositions(query_values[rows[i]], ref_values[cols[i]])) + for i in wild] + keep = dists <= limits[rows] + rows, cols, dists = rows[keep], cols[keep], dists[keep] + + gate = buildGeneGate(queries, reference, gene_cols) + if gate is not None: + keep = gate.keep(rows, cols) + rows, cols, dists = rows[keep], cols[keep], dists[keep] + + ref_ids = reference[id_col].to_numpy() + pairs = pandas.DataFrame({'query_row': rows, 'col': cols, 'dist': dists, + 'id': ref_ids[cols].astype(str)}) + # One reference id can appear on several rows; keep its closest one + # (lowest reference position on a tie). + pairs = (pairs.sort_values(['query_row', 'dist', 'id', 'col']) + .drop_duplicates(['query_row', 'id'])) + + grouped = pairs.groupby('query_row', sort=True) + hits = pandas.DataFrame({'n_hits_total': grouped.size(), + 'min_mismatches': grouped['dist'].min().astype(float), + 'hit_ids': grouped['id'].agg(';'.join)}) + columns = ['query_row', 'n_hits_total', 'min_mismatches', 'is_unique_hit', + 'hit_ids'] + if hamming: + n_cols = len(compare_cols) + pairs['positions'] = [ + _formatPositions(_mismatchPositions(query_values[q], ref_values[c]), n_cols) + for q, c in zip(pairs['query_row'].tolist(), pairs['col'].tolist())] + hits['mismatch_positions'] = pairs.groupby('query_row', sort=True)[ + 'positions'].agg(';'.join) + columns.append('mismatch_positions') + + result = pandas.DataFrame({'query_row': numpy.arange(len(queries))}).join( + hits, on='query_row') + result['n_hits_total'] = result['n_hits_total'].fillna(0).astype(int) + # An all-missing column comes out float: cast it to object before filling + # with '' (pandas will refuse a str fill on a float column), then to str. + text_cols = ['hit_ids'] + (['mismatch_positions'] if hamming else []) + for col in text_cols: + result[col] = result[col].astype(object).fillna('').astype(str) + result['is_unique_hit'] = result['n_hits_total'] <= 1 + result = result[columns] + result.attrs['skipped_query_rows'] = len(skipped_queries) + result.attrs['skipped_rows'] = [int(row) for row in skipped_queries] + if not return_pairs: + return result + + # The pairs behind hit_ids, in the same order. Identity is against the + # query's compared length, the length identity_threshold scales with. + query_rows = pairs['query_row'].to_numpy() + lengths = numpy.array([len(seq) for seq in query_seqs], dtype=numpy.int64) + pair_table = pandas.DataFrame({ + 'query_row': query_rows, 'reference_id': pairs['id'].to_numpy(), + 'mismatches': pairs['dist'].to_numpy(), + 'identity': 1 - pairs['dist'].to_numpy() / lengths[query_rows]}) + if hamming: + pair_table['mismatch_positions'] = pairs['positions'].to_numpy() + + return result, pair_table diff --git a/src/immvoke/Sources/Specificity/__init__.py b/src/immvoke/Sources/Specificity/__init__.py new file mode 100644 index 0000000..4f84dbd --- /dev/null +++ b/src/immvoke/Sources/Specificity/__init__.py @@ -0,0 +1,43 @@ +""" +Specificity source registry + +A second, separate registry from immvoke.Sources: specificity databases +(IEDB, CoV-AbDab, ASD, and others to follow) are SourceBase subclasses too, +but are kept apart from repertoire sources so the commandline can mount them +under their own `immvoke specificity ` tree instead of `immvoke +`. +""" + +# Info +__author__ = 'Pramod Shinde' + +# Immvoke imports +from immvoke.Sources.Specificity.Asd import AsdSource +from immvoke.Sources.Specificity.CovAbDab import CovAbDabSource +from immvoke.Sources.Specificity.Iedb import IedbSource + +#: Every specificity source immvoke knows about, by commandline name. +REGISTRY = {IedbSource.name: IedbSource, CovAbDabSource.name: CovAbDabSource, + AsdSource.name: AsdSource} + + +def getSpecificitySource(name, client, schema=None): + """ + Instantiate a specificity source by name. + + Arguments: + name (str): the source name. + client (HttpClient): the shared HTTP client. + schema (SourceSchema): a preloaded snapshot, or None to load on demand. + + Returns: + SourceBase: the source. + + Raises: + KeyError: if the name is not registered. + """ + if name not in REGISTRY: + raise KeyError("unknown specificity source '%s'; known: %s" + % (name, ', '.join(sorted(REGISTRY)))) + + return REGISTRY[name](client, schema=schema) diff --git a/src/immvoke/Sources/Specificity/idea.md b/src/immvoke/Sources/Specificity/idea.md new file mode 100644 index 0000000..aa5a0e8 --- /dev/null +++ b/src/immvoke/Sources/Specificity/idea.md @@ -0,0 +1,189 @@ +# Specificity annotation: how it works + +`immvoke annotate` and `annotateSpecificity()` match a query AIRR table +against a reference AIRR table and report, for every query row, which +reference rows it matches. This page explains the flow and the rules behind +each step; [README.md](README.md) covers how to run it. + +## Workflow + +Each query row goes through these steps in order. Steps marked *automatic* +always run; the others are set with the option shown (the +`annotateSpecificity()` argument has the same name, with `_` for `-`). + +| # | Step | Option | Choices | Default | +| --- | --- | --- | --- | --- | +| 1 | Check the reference is AIRR | *automatic* | -- | -- | +| 2 | Paired or single chain? | `--paired` | on, off | off | +| 2a | Paired: one heavy and one light chain per query cell | `--count-field` | an AIRR count column | `duplicate_count` | +| 3 | Keep heavy chain rows | *automatic* | `IGH` only | -- | +| 4 | Sequence to compare | `--sequence-region` | `CDR3H`, `VDJ`, `CDRs` | `CDR3H` | +| 5 | Gene gate | `--partition-by` | none, `V`, `VJ`, `VJL` | none | +| 6 | How close is a hit | `--identity-threshold` | a fraction in (0, 1] | exact | +| 7 | Distance | `--distance-metric` | `hamming`, `levenshtein` | `hamming` | +| 8 | Search | *automatic* | symscan; rows too large for it are scored directly instead, and counted in `annotation_metadata.yml` | -- | +| 9 | Output | `--outdir` (`outdir` in Python), `--full-dist-matrix` | `annotation.tsv` (one row per query row and matching reference row, its best hits, also returned as a DataFrame), `annotation_summary.txt` (the run in prose) and `annotation_metadata.yml`, plus `annotation_full_matrix.txt` (every pair the filters admit, whatever its distance) with `--full-dist-matrix`, each prefixed with the reference's database and table (`iedb_bcr_annotation.tsv`) | -- | + +With `--paired` on, steps 3-7 are replaced by one fixed comparison: an +antibody (`cell_id`) matches a reference antibody only if both its heavy and +its light chain have the same V gene, the same J gene and an exact CDR3 -- +classified `both`, `heavy_only`, `light_only`, `conflicting` or `neither`. +Without it, every query heavy chain row that matched something gets one +result row -- an unmatched row is left out, counted in the metadata +instead: its best hits -- the highest identity, all of them on a tie -- +with `identity`, `hit_id` and that hit's `antigen`, `epitope` and +`antigen_source_species` from the reference -- a row per (query row, +matching reference row), so nothing is joined into a cell. +`annotation_full_matrix.txt` (`--full-dist-matrix`) lists every pair the +filters admit and how close it is, one (query, hit, identity) per row -- +`--identity-threshold` does not apply to it, `--partition-by` does. + +The output files are named after the reference, `_
_` -- the +database read off the `download_metadata.yml` `immvoke specificity` writes +beside the reference table, the table off its file name -- so runs against +different databases stay apart, even in one `--outdir`. A reference without +that record is named by its file name alone. + +`annotation_metadata.yml` records the run: the options, the query and +reference (with the reference's source database, read from the +`download_metadata.yml` beside it), the counts -- rows scored directly +included -- and the immvoke version, time and command. + +## Input contract + +The reference must always be an **AIRR rearrangement table**, whichever +database it comes from: every field the AIRR schema requires (`sequence_id`, +`sequence`, `rev_comp`, `productive`, `v_call`, `d_call`, `j_call`, +`sequence_alignment`, `germline_alignment`, `junction`, `junction_aa`, +`v_cigar`, `d_cigar`, `j_cigar`) present as a column, though values may be +blank. Any table immvoke writes as AIRR already is one; anything else is +refused with an error naming the missing fields. + +On top of that, each step reads its own AIRR fields: `cdr3_aa`, +`sequence_aa` or `cdr1_aa`/`cdr2_aa` for the sequence region, `v_call` and +`j_call` for partitions and paired matching, `locus` for the chain filter, +and `cell_id` to link the two chains of a paired antibody. The query only +needs the fields the chosen options read. A table without a `locus` column +is compared whole; a table with one but no `IGH` row left is an error, never +an empty result. + +`cdr3_aa` must follow the AIRR definition in both tables: the CDR3 without +its two flanking conserved residues, i.e. `junction_aa` minus its first and +last residue (`ARGLDS`, not `CARGLDSW`). IEDB's does; annotate does not check +or convert a query's, so a query whose `cdr3_aa` holds junctions will find +no exact match. + +A blank `locus` is kept as the reference has it -- nothing is inferred +from the gene calls -- so such a row is in neither the heavy nor the light +chain pool: single-chain matching never compares it, and paired matching +cannot pair it. (Of IEDB's human chains, 277 have a blank locus, nearly all +of them light chains.) + +## Matching rules + +**Sequence region.** `CDR3H` compares the heavy chain CDR3 (`cdr3_aa`), `VDJ` compares +`sequence_aa`, and `CDRs` pools `cdr1_aa`, `cdr2_aa` and `cdr3_aa` under one +shared tolerance. Pooling only works with Hamming distance, because edit +distance does not add up across columns. + +**Blanks.** A row with a blank value in any compared column (say, no +`cdr1_aa` under `CDRs`) has nothing to compare and never matches -- not +even another blank. A blank query row gets no hit; a blank reference row is +never hit. + +**Partition.** `V` and `VJ` only compare row pairs whose V-gene (and +J-gene) agree; `VJL` also requires the same CDR3 length, and so needs +`cdr3_aa` in both tables whatever the sequence region is. + +**Gene names.** `v_call`/`j_call` cells are not always one clean name: IEDB +mixes allele calls (`IGHV1-2*02`), gene calls (`IGHV1-2`), several calls in +one cell (`IGHV1-2*02, IGHV1-2*04`) and blanks. They are parsed with +Change-O's gene pattern (`Genes.py`) and compared by exact gene name, +whatever the allele: `IGHV1-2*02`, `IGHV1-2*04` and `IGHV1-2` all agree, +while `IGHV1-2` and `IGHV1-69` do not. + +A cell with several calls agrees if any one of them does. A blank or +unrecognized call **agrees with nothing**, not even another blank; the +number of unrecognized values is logged. So under `VJ`, a row with no J +call can never match -- use `V`, or no partition, to include such rows. + +**Threshold and metric.** An unset (or 1.0) `identity_threshold` is an +exact match. Below that, a query may differ from a hit by up to +floor((1 - threshold) x query length) positions: substitutions only under +`hamming`, or substitutions, insertions and deletions under `levenshtein`, +which also lets the two differ in length. Hamming reports which positions +differ; Levenshtein cannot, since an indel shifts every position after it. + +**X residues.** An `X` is an unresolved residue. Under exact and Hamming +matching it matches any residue, so `CARXYW` equals `CARDYW`; under +Levenshtein it is an ordinary letter. + +**Paired.** An antibody is its heavy (`IGH`) and light (`IGK`/`IGL`) chain +rows sharing one `cell_id`. It matches a reference antibody only if that +*same* antibody agrees on V-gene and J-gene (same gene, any allele) and +exact CDR3 on both chains; a heavy hit on one reference antibody and a light +hit on another is `conflicting`, not `both`. Region, partition, threshold +and metric do not apply. A row with a blank `cell_id` has no partner chain, +so it is left out, on either side, with a warning giving the count. + +A query cell can carry more than one heavy chain, or more than one light +chain (a doublet, ambient RNA, or a second light chain). Before matching, +each query cell keeps one heavy (`IGH`) and one light (`IGK` or `IGL`) +chain: the one with the highest `--count-field` value (`duplicate_count` by +default, or e.g. `umi_count`). On a tie, a blank count, or a file without +that column, the chain that comes first in the file is kept. The reference +is not reduced this way: a reference antibody with several chains of one +kind (4 heavy and 26 light cases among IEDB's human antibodies, which carry +no count column) keeps them all. + +## Engine + +- **symscan finds the hits.** It is a required dependency, a + symmetric-deletion index that only looks at likely matches instead of + scoring every query x reference pair. There is no install check and no + pure-Python matcher. +- **Every gated pair through numpy.** symscan never scores a pair beyond + the threshold, so the pairs file of a Hamming or Levenshtein run is + scored separately (`Distances.py`): references are indexed by gene and + gate key so each query only meets the pairs its gates allow, then a + block of pairs is compared position by position (Hamming, `X` never + counting) or through an edit distance table filled for the whole block + at once (Levenshtein). Queries go in blocks, so the file is written as it + goes, and its `within_threshold` pairs are exactly symscan's hits. +- **`X` through symscan.** symscan has no wildcard and compares `X` as a + letter, so the `X`s are taken out of the comparison instead: reference + rows are grouped by where their `X`s sit, and those positions are cut out + of them and of every same-length query before symscan sees them. A + query's own `X`s (rarer) are covered by slack -- one extra difference + each -- and every pair involving an `X` is re-scored with the wildcard + rule, keeping only those within the real threshold. +- **Rows beyond symscan's limits are scored directly.** symscan aborts the + process above 14 differences, or once its index -- every deletion variant + of every sequence, which grows steeply with length and tolerance -- + outgrows memory. Both the queries and the references are chunked to a + fixed budget, so table size alone never matters; what cannot fit at all + (a threshold allowing over 14 differences, `X` slack included, or one + sequence too long to index at that threshold) is kept away from symscan, + and those query rows' pairs are scored directly instead (`Distances.py`), + so no row misses a hit; their count is in `annotation_metadata.yml`. On + today's IEDB table that is a handful of outliers, such as reference + "CDR3s" of 60+ residues at a loose threshold. A VDJ comparison (~110-130 + residues) fits symscan up to 4 differences -- exact matching and + thresholds of about 0.97 and up; below that most VDJ rows are scored + directly, which is correct but slower, so `--partition-by` matters there. + (Under `--paired`, which is always exact, only a query with over 14 `X`s + goes beyond symscan; such a row is skipped, with a warning.) +- **Order of steps.** symscan measures distance first; the exact-column and + gene checks are applied to the pairs it returns. + +## Current limits + +- Only IEDB's `bcr` table comes in AIRR so far (IEDB publishes it that + way); CoV-AbDab and ASD are still downloaded in their own columns, so they + cannot be used as a reference until they are converted. +- Antibodies only: TCR tables are not supported yet (no `IGH` rows, so they + are refused). +- No substitution-matrix (BLOSUM) scoring. +- Fuzzy VDJ matching below a threshold of about 0.97 is out of symscan's + reach; those rows are scored directly, which is slow without + `--partition-by`. diff --git a/src/immvoke/Sources/Specificity/tests/BruteForce.py b/src/immvoke/Sources/Specificity/tests/BruteForce.py new file mode 100644 index 0000000..272301f --- /dev/null +++ b/src/immvoke/Sources/Specificity/tests/BruteForce.py @@ -0,0 +1,163 @@ +""" +A brute-force reference matcher, for tests only + +Compares every query row with every reference row directly -- no index, no +grouping, no slack -- and builds the same result and pairs tables as +Symscan.runSymscanAnnotation, so symscan's answers can be checked against an +implementation simple enough to trust by reading it. Far too slow for real +tables. +""" + +# Info +__author__ = 'Pramod Shinde' + +# Imports +import pandas + +# Immvoke imports +from immvoke.Sources.Specificity.Annotate import WILDCARD, _allowedDifferences +from immvoke.Sources.Specificity.Genes import GeneGate + + +def _hamming(a, b): + """ + Returns: + list: (column, position) of every difference between two rows' column + tuples, an X on either side never counting -- or None if any column's + lengths differ. + """ + if any(len(x) != len(y) for x, y in zip(a, b)): + return None + + return [(col, pos) for col, (x, y) in enumerate(zip(a, b)) + for pos, (cx, cy) in enumerate(zip(x, y)) + if cx != cy and WILDCARD not in (cx, cy)] + + +def _levenshtein(a, b): + """ + Returns: + int: the plain edit distance between two strings. + """ + previous = list(range(len(b) + 1)) + for i, ca in enumerate(a, 1): + current = [i] + for j, cb in enumerate(b, 1): + current.append(min(previous[j] + 1, current[j - 1] + 1, + previous[j - 1] + (ca != cb))) + previous = current + + return previous[-1] + + +def bruteForceAnnotation(queries, reference, compare_cols, max_distance, + distance_type='hamming', id_col='id', exact_cols=None, + threshold_is_fraction=False, gene_cols=None, + return_pairs=False): + """ + The result table runSymscanAnnotation must produce, by brute force. + + Arguments: as Symscan.runSymscanAnnotation. + + Returns: + pandas.DataFrame: as Symscan.runSymscanAnnotation. + """ + result, pairs, _ = _bruteForce(queries, reference, compare_cols, max_distance, + distance_type, id_col, exact_cols, + threshold_is_fraction, gene_cols) + + return (result, pairs) if return_pairs else result + + +def bruteForceAllPairs(queries, reference, compare_cols, max_distance, + distance_type='hamming', id_col='id', exact_cols=None, + threshold_is_fraction=False, gene_cols=None): + """ + The pairs Distances.allPairDistances must yield, concatenated, by brute + force: every pair through the gates, whatever its distance. + + Arguments: as Symscan.runSymscanAnnotation. + + Returns: + pandas.DataFrame: as Distances.allPairDistances's tables. + """ + return _bruteForce(queries, reference, compare_cols, max_distance, distance_type, + id_col, exact_cols, threshold_is_fraction, gene_cols)[2] + + +def _bruteForce(queries, reference, compare_cols, max_distance, distance_type, + id_col, exact_cols, threshold_is_fraction, gene_cols): + """ + Returns: + tuple: (the result table, its pairs, every gated pair). + """ + hamming = distance_type == 'hamming' + query_values = list(queries[compare_cols].fillna('').itertuples(index=False, + name=None)) + ref_values = list(reference[compare_cols].fillna('').itertuples(index=False, + name=None)) + exact = exact_cols or [] + gate = GeneGate(queries, reference, gene_cols) if gene_cols else None + + rows, pairs, all_pairs = [], [], [] + for q, values in enumerate(query_values): + blank = '' in values + length = sum(len(v) for v in values) + limit = (_allowedDifferences(max_distance, length) if threshold_is_fraction + else int(max_distance)) + best = {} # id -> (distance, reference position, mismatch positions) + scored = {} # the same, whatever the distance + for r, candidate in enumerate(ref_values): + if blank or '' in candidate: + continue + if any(queries[c].iloc[q] != reference[c].iloc[r] for c in exact): + continue + if gate is not None and not gate.agrees(q, r): + continue + if hamming: + positions = _hamming(values, candidate) + if positions is None: + continue + distance = len(positions) + else: + positions, distance = None, _levenshtein(values[0], candidate[0]) + hit = str(reference[id_col].iloc[r]) + if hit not in scored or (distance, r) < scored[hit][:2]: + scored[hit] = (distance, r, positions) + if distance > limit: + continue + if hit not in best or (distance, r) < best[hit][:2]: + best[hit] = (distance, r, positions) + + single = len(compare_cols) == 1 + + def formatted(found, x): + return ','.join(str(pos) if single else '%d:%d' % (col, pos) + for col, pos in found[x][2]) + + order = sorted(best, key=lambda x: (best[x][0], x)) + row = {'query_row': q, 'n_hits_total': len(order), + 'min_mismatches': float(best[order[0]][0]) if order else float('nan'), + 'is_unique_hit': len(order) <= 1, 'hit_ids': ';'.join(order)} + if hamming: + row['mismatch_positions'] = ';'.join(formatted(best, x) for x in order) + rows.append(row) + for x in order: + pair = {'query_row': q, 'reference_id': x, 'mismatches': best[x][0], + 'identity': 1 - best[x][0] / length} + if hamming: + pair['mismatch_positions'] = formatted(best, x) + pairs.append(pair) + for x in sorted(scored, key=lambda x: (scored[x][0], x)): + pair = {'query_row': q, 'reference_id': x, 'mismatches': scored[x][0], + 'identity': 1 - scored[x][0] / length, + 'within_threshold': scored[x][0] <= limit} + if hamming: + pair['mismatch_positions'] = formatted(scored, x) + all_pairs.append(pair) + + extra = ['mismatch_positions'] if hamming else [] + columns = ['query_row', 'reference_id', 'mismatches', 'identity'] + + return (pandas.DataFrame(rows), pandas.DataFrame(pairs, columns=columns + extra), + pandas.DataFrame(all_pairs, columns=columns + ['within_threshold'] + extra)) diff --git a/src/immvoke/Sources/Specificity/tests/__init__.py b/src/immvoke/Sources/Specificity/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/immvoke/Sources/Specificity/tests/test_Annotate.py b/src/immvoke/Sources/Specificity/tests/test_Annotate.py new file mode 100644 index 0000000..8a88ce5 --- /dev/null +++ b/src/immvoke/Sources/Specificity/tests/test_Annotate.py @@ -0,0 +1,1001 @@ +""" +Unit tests for antigen specificity annotation (Annotate.py, AnnotateCli.py) +""" + +# Info +__author__ = 'Pramod Shinde' + +# Imports +import io +import tempfile +import unittest +import warnings +from pathlib import Path +from unittest import mock + +import numpy +import pandas +import yaml +from airr.schema import RearrangementSchema + +# Immvoke imports +from immvoke.Cli import getArgParser +from immvoke.Convert import writeAirr +from immvoke.Exceptions import AnnotationError +from immvoke.Sources.Specificity.Annotate import ( + ANNOTATION_FILE, + ANNOTATION_METADATA, + ANNOTATION_SUMMARY, + FULL_MATRIX_FILE, + HIT_ANNOTATIONS, + annotateSpecificity, + runPairedExactAnnotation, +) +from immvoke.Sources.Specificity.AnnotateCli import handleAnnotate +from immvoke.Sources.Specificity.Symscan import runSymscanAnnotation +from immvoke.Sources.Specificity.tests.BruteForce import bruteForceAnnotation + +ANNOTATE = 'immvoke.Sources.Specificity.Annotate' + + +def frame(rows): + return pandas.DataFrame(rows) + + +class TestRunPairedExactAnnotation(unittest.TestCase): + """ + Tests for the heavy+light reconciliation (exact V-gene, J-gene and CDR3 + agreement, both chains against the same reference antibody) + """ + + def _antibodies(self): + """ + 5 query antibodies, one per classification, against 2 reference + antibodies -- see the inline comments for which reference row each + is built to agree or disagree with. + """ + queries = frame([ + # q1: heavy and light both agree with refA -> both + {'cell_id': 'q1', 'locus': 'IGH', 'cdr3_aa': 'CARDYW', + 'v_call': 'IGHV1', 'j_call': 'IGHJ1'}, + {'cell_id': 'q1', 'locus': 'IGL', 'cdr3_aa': 'CQQYNS', + 'v_call': 'IGKV1', 'j_call': 'IGKJ1'}, + # q2: heavy agrees with refA, light agrees with nothing -> heavy_only + {'cell_id': 'q2', 'locus': 'IGH', 'cdr3_aa': 'CARDYW', + 'v_call': 'IGHV1', 'j_call': 'IGHJ1'}, + {'cell_id': 'q2', 'locus': 'IGL', 'cdr3_aa': 'CDIFFER', + 'v_call': 'IGKV9', 'j_call': 'IGKJ9'}, + # q3: heavy agrees with nothing, light agrees with refA -> light_only + {'cell_id': 'q3', 'locus': 'IGH', 'cdr3_aa': 'CDIFF2', + 'v_call': 'IGHV9', 'j_call': 'IGHJ9'}, + {'cell_id': 'q3', 'locus': 'IGL', 'cdr3_aa': 'CQQYNS', + 'v_call': 'IGKV1', 'j_call': 'IGKJ1'}, + # q4: neither chain agrees with anything -> neither + {'cell_id': 'q4', 'locus': 'IGH', 'cdr3_aa': 'CDIFF3', + 'v_call': 'IGHV8', 'j_call': 'IGHJ8'}, + {'cell_id': 'q4', 'locus': 'IGL', 'cdr3_aa': 'CDIFF4', + 'v_call': 'IGKV8', 'j_call': 'IGKJ8'}, + # q5: heavy agrees with refA, light agrees with refB -> conflicting + {'cell_id': 'q5', 'locus': 'IGH', 'cdr3_aa': 'CARDYW', + 'v_call': 'IGHV1', 'j_call': 'IGHJ1'}, + {'cell_id': 'q5', 'locus': 'IGL', 'cdr3_aa': 'CZZZZZZ', + 'v_call': 'IGKV7', 'j_call': 'IGKJ7'}, + ]) + reference = frame([ + {'cell_id': 'refA', 'locus': 'IGH', 'cdr3_aa': 'CARDYW', + 'v_call': 'IGHV1', 'j_call': 'IGHJ1'}, + {'cell_id': 'refA', 'locus': 'IGL', 'cdr3_aa': 'CQQYNS', + 'v_call': 'IGKV1', 'j_call': 'IGKJ1'}, + {'cell_id': 'refB', 'locus': 'IGH', 'cdr3_aa': 'CWWWWWW', + 'v_call': 'IGHV6', 'j_call': 'IGHJ6'}, + {'cell_id': 'refB', 'locus': 'IGL', 'cdr3_aa': 'CZZZZZZ', + 'v_call': 'IGKV7', 'j_call': 'IGKJ7'}, + ]) + return queries, reference + + def test_classifications(self): + queries, reference = self._antibodies() + + result = runPairedExactAnnotation(queries, reference) + + by_cell_id = result.set_index('cell_id')['classification'].to_dict() + self.assertEqual(by_cell_id, { + 'q1': 'both', 'q2': 'heavy_only', 'q3': 'light_only', + 'q4': 'neither', 'q5': 'conflicting', + }) + + def test_both_hit_ids_is_the_intersection_not_the_union(self): + queries, reference = self._antibodies() + + result = runPairedExactAnnotation(queries, reference) + + row = result.set_index('cell_id').loc['q1'] + self.assertEqual(row['heavy_hit_ids'], 'refA') + self.assertEqual(row['light_hit_ids'], 'refA') + self.assertEqual(row['both_hit_ids'], 'refA') + + row5 = result.set_index('cell_id').loc['q5'] + self.assertEqual(row5['heavy_hit_ids'], 'refA') + self.assertEqual(row5['light_hit_ids'], 'refB') + self.assertEqual(row5['both_hit_ids'], '') + +def airr(rows): + """ + An AIRR rearrangement table: `rows`, plus every field the AIRR schema + requires that they lack, blank -- what immvoke's AIRR writer produces. + """ + table = pandas.DataFrame(rows) + for field in RearrangementSchema.required: + if field not in table.columns: + table[field] = '' + + return table + + +#: sequence_region -> the columns it compares, as annotateSpecificity does. +REGION_COLUMNS = {'CDR3H': ['cdr3_aa'], 'VDJ': ['sequence_aa'], + 'CDRs': ['cdr1_aa', 'cdr2_aa', 'cdr3_aa']} + + +def bruteForceSpecificity(queries, reference, sequence_region='CDR3H', + partition_by=None, identity_threshold=None, + distance_metric='hamming'): + """ + annotateSpecificity's single-chain result table, by brute force: the + same heavy chain filter, columns, gates and threshold, every pair + compared directly (see BruteForce.py). + """ + def heavyChains(rows): + table = pandas.DataFrame(rows) + if 'locus' in table.columns: + table = table[table['locus'] == 'IGH'] + return table.reset_index(drop=True) + + queries, reference = heavyChains(queries), heavyChains(reference) + gene_cols = {'V': ['v_call'], 'VJ': ['v_call', 'j_call'], + 'VJL': ['v_call', 'j_call']}.get(partition_by) + exact_cols = None + if partition_by == 'VJL': + queries = queries.assign(cdr3_length=queries['cdr3_aa'].str.len()) + reference = reference.assign(cdr3_length=reference['cdr3_aa'].str.len()) + exact_cols = ['cdr3_length'] + exact = identity_threshold is None or identity_threshold >= 1.0 + + _, pairs = bruteForceAnnotation( + queries, reference, REGION_COLUMNS[sequence_region], + 0 if exact else 1 - identity_threshold, + distance_type='hamming' if exact else distance_metric, id_col='sequence_id', + exact_cols=exact_cols, threshold_is_fraction=not exact, gene_cols=gene_cols, + return_pairs=True) + + # Each query's best hits: the fewest mismatches, all of them on a tie. + # A query row that matched nothing is left out, as the real one does. + # Each reference id's annotation values, one entry per row with that id. + annotations = {} + for _, row in reference.iterrows(): + values = tuple(str(row.get(column, '') or '') + for _, column in HIT_ANNOTATIONS) + annotations.setdefault(row['sequence_id'], {})[values] = None + + rows = [] + for row, sequence_id in enumerate(queries['sequence_id']): + hits = pairs[pairs['query_row'] == row] + if not len(hits): + continue + best = hits[hits['mismatches'] == hits['mismatches'].min()] + for hit_id, identity in zip(best['reference_id'], best['identity']): + for values in annotations.get(hit_id, {(('',) * len(HIT_ANNOTATIONS))}): + rows.append(dict(zip(['sequence_id', 'identity', 'hit_id'], + (sequence_id, round(identity, 4), hit_id)), + **dict(zip([n for n, _ in HIT_ANNOTATIONS], values)))) + + columns = ['sequence_id', 'identity', 'hit_id'] + [n for n, _ in HIT_ANNOTATIONS] + + return pandas.DataFrame(rows, columns=columns).sort_values( + ['sequence_id', 'hit_id'], kind='stable').reset_index(drop=True) + + +def heavy(sequence_id, cdr3_aa, v_call='IGHV1', j_call='IGHJ1', **extra): + return dict(sequence_id=sequence_id, locus='IGH', cdr3_aa=cdr3_aa, + v_call=v_call, j_call=j_call, **extra) + + +class TestAnnotateSpecificity(unittest.TestCase): + """ + Tests for the high level entry point, step by step through its flow: + paired?, heavy chain, sequence_region, partition_by, + identity_threshold/distance_metric, symscan, the outdir bundle + """ + + def _hits(self, queries, reference, **kwargs): + """Each heavy chain query's hit_ids, in query order, '' where it + matched nothing. The real table leaves an unmatched query out + entirely; this puts it back so that a test can say which query + matched what by position.""" + result = annotateSpecificity(frame(queries), airr(reference), **kwargs) + found = {} + for sequence_id, hit in zip(result['sequence_id'], result['hit_id']): + found.setdefault(sequence_id, dict())[hit] = None + + return [';'.join(found.get(x['sequence_id'], ())) + for x in queries if x.get('locus', 'IGH') == 'IGH'] + + def test_default_is_an_exact_cdr3_match_reporting_only_matched_rows(self): + """A query row that matched nothing is left out, so a repertoire + where almost nothing matches gives a short table, not one row per + query; the counts stay in the metadata.""" + with tempfile.TemporaryDirectory() as tmp: + result = annotateSpecificity( + frame([heavy('q1', 'CARDYW'), heavy('q2', 'CARDYY')]), + airr([heavy('ref1', 'CARDYW', v_call='IGHV3')]), outdir=tmp) + metadata = yaml.safe_load(Path(tmp, 'reference_' + ANNOTATION_METADATA) + .read_text()) + + self.assertEqual(result['sequence_id'].tolist(), ['q1']) + self.assertEqual(result['hit_id'].tolist(), ['ref1']) + self.assertEqual(result['identity'].iloc[0], 1.0) + self.assertEqual(result.columns.tolist(), + ['sequence_id', 'identity', 'hit_id', 'antigen', + 'epitope', 'antigen_source_species']) + self.assertEqual(metadata['results']['query_rows_compared'], 2) + self.assertEqual(metadata['results']['query_rows_with_hits'], 1) + + def test_queries_are_named_by_their_sequence_id(self): + """By the query's own sequence_id -- not by position, which counts + heavy chain rows only -- or by position without that column.""" + light = dict(sequence_id='q1-light', locus='IGK', cdr3_aa='CQQYNS', + v_call='IGKV1', j_call='IGKJ1') + queries = [light, heavy('q2', 'CARDYW'), heavy('q3', 'CQQQQQ')] + with tempfile.TemporaryDirectory() as tmp: + result = annotateSpecificity(frame(queries), airr([heavy('ref1', 'CARDYW')]), + identity_threshold=0.8, outdir=tmp, + full_dist_matrix=True) + + matrix = pandas.read_csv(Path(tmp) / ('reference_' + FULL_MATRIX_FILE), + sep='\t') + self.assertEqual(result[['sequence_id', 'hit_id']].values.tolist(), + [['q2', 'ref1']]) + self.assertEqual(matrix[['sequence_id', 'hit_id']].values.tolist(), + [['q2', 'ref1'], ['q3', 'ref1']]) + + anonymous = frame([{'cdr3_aa': 'CARDYW'}, {'cdr3_aa': 'CQQQQQ'}]) + result = annotateSpecificity(anonymous, airr([heavy('ref1', 'CARDYW')])) + self.assertEqual(result[['query_row', 'hit_id']].values.tolist(), + [[0, 'ref1']]) + + def test_only_heavy_chain_rows_are_compared(self): + light = dict(sequence_id='q2', locus='IGK', cdr3_aa='CQQYNS', + v_call='IGKV1', j_call='IGKJ1') + reference = [heavy('ref1', 'CARDYW'), dict(light, sequence_id='ref2')] + + result = annotateSpecificity(frame([heavy('q1', 'CARDYW'), light]), + airr(reference)) + + self.assertEqual(result['hit_id'].tolist(), ['ref1']) + + def test_sequence_region_picks_the_compared_columns(self): + queries = [heavy('q1', 'CARDYW', sequence_aa='EVQLCARDYW', + cdr1_aa='GFT', cdr2_aa='ISY')] + reference = [heavy('ref1', 'CARDYW', sequence_aa='QVQLCARDYW', + cdr1_aa='GFT', cdr2_aa='ISS')] + + self.assertEqual(self._hits(queries, reference, sequence_region='CDR3H'), + ['ref1']) + self.assertEqual(self._hits(queries, reference, sequence_region='VDJ'), ['']) + self.assertEqual(self._hits(queries, reference, sequence_region='CDRs'), ['']) + + def test_cdrs_pools_the_threshold_across_cdr1_to_cdr3(self): + """12 pooled residues at 0.85 identity -> floor(1.8) == 1 substitution, + here in CDR2 -- no single column is compared on its own.""" + queries = [heavy('q1', 'CARDYW', cdr1_aa='GFT', cdr2_aa='ISY')] + reference = [heavy('ref1', 'CARDYW', cdr1_aa='GFT', cdr2_aa='ISS'), + heavy('ref2', 'CARDYQ', cdr1_aa='GFA', cdr2_aa='ISY')] + + self.assertEqual(self._hits(queries, reference, sequence_region='CDRs', + identity_threshold=0.85), ['ref1']) + + def test_partition_by_gates_on_genes_and_cdr3_length(self): + queries = [heavy('q1', 'CARDYW', sequence_aa='EVQLCARDYW')] + # Each 1 insertion away, so all three tie as best hits when let in. + reference = [ + heavy('sameVJ', 'CARDDYW', sequence_aa='EVQLCARDDYW'), + heavy('otherJ', 'CARDDYW', j_call='IGHJ4', sequence_aa='EVQLCARDDYW'), + heavy('otherV', 'CARDDYW', v_call='IGHV3', sequence_aa='EVQLCARDDYW'), + ] + loose = dict(sequence_region='VDJ', identity_threshold=0.9, + distance_metric='levenshtein') + + def hits(partition_by): + ids = self._hits(queries, reference, partition_by=partition_by, **loose) + return set(ids[0].split(';')) - {''} + + self.assertEqual(hits(None), {'sameVJ', 'otherJ', 'otherV'}) + self.assertEqual(hits('V'), {'sameVJ', 'otherJ'}) + self.assertEqual(hits('VJ'), {'sameVJ'}) + self.assertEqual(hits('VJL'), set()) + + def test_identity_threshold_scales_with_query_length(self): + """length 10 at 0.85 identity -> floor(1.5) == 1 substitution.""" + queries = [heavy('q1', 'AAAAAAAAAA')] + reference = [heavy('close', 'AAAAAAAAAB'), heavy('far', 'AABBAAAAAB')] + + self.assertEqual(self._hits(queries, reference, identity_threshold=0.85), + ['close']) + + def test_levenshtein_tolerates_an_indel_hamming_does_not(self): + queries = [heavy('q1', 'CARDYWAAAA')] + reference = [heavy('ref1', 'CARDXYWAAAA')] # 1 insertion + + self.assertEqual(self._hits(queries, reference, identity_threshold=0.9, + distance_metric='levenshtein'), ['ref1']) + self.assertEqual(self._hits(queries, reference, identity_threshold=0.9), ['']) + + def test_matches_brute_force(self): + queries = [heavy('q%d' % i, cdr3, cdr1_aa='GFT', cdr2_aa='ISY') + for i, cdr3 in enumerate(['CARDYW', 'CARXYW', 'CARDDYW', 'CQ'])] + reference = [heavy('r%d' % i, cdr3, cdr1_aa='GFT', cdr2_aa='ISS') + for i, cdr3 in enumerate(['CARDYW', 'CARQYW', 'CXRDYW', 'CQ'])] + modes = [{}, {'identity_threshold': 0.8}, + {'identity_threshold': 0.8, 'distance_metric': 'levenshtein'}, + {'sequence_region': 'CDRs', 'identity_threshold': 0.9}, + {'partition_by': 'VJL', 'identity_threshold': 0.8}] + for mode in modes: + with self.subTest(**mode): + result = annotateSpecificity(frame(queries), airr(reference), **mode) + pandas.testing.assert_frame_equal( + result, bruteForceSpecificity(queries, reference, **mode), + check_dtype=False) + + def test_x_matches_any_residue(self): + self.assertEqual(self._hits([heavy('q1', 'CARXYW'), heavy('q2', 'CAKDYW')], + [heavy('ref1', 'CARDYW')]), ['ref1', '']) + + def test_a_row_beyond_symscan_is_scored_directly(self): + """20 residues at 0.25 identity -> 15 differences, above symscan's 14: + symscan is never called with it, and the row is scored directly + instead -- so its hit, 15 differences away, is still found.""" + with mock.patch('symscan.get_neighbors_across', + side_effect=AssertionError), \ + self.assertLogs(ANNOTATE, 'INFO') as logs: + hits = self._hits([heavy('q1', 'A' * 20)], + [heavy('ref1', 'B' * 15 + 'A' * 5)], + identity_threshold=0.25) + + self.assertEqual(hits, ['ref1']) + self.assertTrue(any('scoring their pairs directly' in line + for line in logs.output)) + + def test_an_exact_row_beyond_symscan_keeps_its_hit_and_pair(self): + """15 X's are 15 differences of slack, beyond symscan; scored + directly, an X still matches any residue.""" + query = frame([heavy('q1', 'X' * 15 + 'CA'), heavy('q2', 'CARDYW')]) + reference = airr([heavy('ref1', 'ACDEFGHIKLMNPQRCA'), heavy('ref2', 'CARDYW')]) + with tempfile.TemporaryDirectory() as tmp: + result = annotateSpecificity(query, reference, outdir=tmp, + full_dist_matrix=True) + + matrix = pandas.read_csv(Path(tmp) / ('reference_' + FULL_MATRIX_FILE), + sep='\t') + metadata = yaml.safe_load(Path(tmp, 'reference_' + ANNOTATION_METADATA) + .read_text()) + self.assertEqual(result['hit_id'].tolist(), ['ref1', 'ref2']) + self.assertEqual(matrix[['sequence_id', 'hit_id']].values.tolist(), + [['q1', 'ref1'], ['q2', 'ref2']]) + self.assertEqual(metadata['results']['query_rows_scored_directly'], 1) + + def test_paired_classifies_each_antibody(self): + rows = [dict(cell_id='ab1', locus='IGH', cdr3_aa='CARDYW', v_call='IGHV1', + j_call='IGHJ1'), + dict(cell_id='ab1', locus='IGK', cdr3_aa='CQQYNS', v_call='IGKV1', + j_call='IGKJ1')] + reference = [dict(row, cell_id='refA') for row in rows] + + result = annotateSpecificity(frame(rows), airr(reference), paired=True) + self.assertEqual(result['classification'].tolist(), ['both']) + self.assertEqual(result['both_hit_ids'].tolist(), ['refA']) + + with self.assertLogs(ANNOTATE, 'WARNING'): + annotateSpecificity(frame(rows), airr(reference), paired=True, + identity_threshold=0.9) + + def test_reads_paths_and_writes_the_bundle(self): + with tempfile.TemporaryDirectory() as tmp: + tmp = Path(tmp) + frame([heavy('q1', 'CARDYW'), heavy('q2', 'CQQQQQ')]).to_csv( + tmp / 'q.tsv', sep='\t', index=False) + airr([heavy('ref1', 'CARDYW')]).to_csv(tmp / 'r.tsv', sep='\t', + index=False) + + result = annotateSpecificity(str(tmp / 'q.tsv'), tmp / 'r.tsv', + partition_by='VJ', outdir=tmp / 'sub' / 'run') + + # Not downloaded by `immvoke specificity`: named by its file name. + run = tmp / 'sub' / 'run' + self.assertEqual(sorted(p.name for p in run.iterdir()), + ['r_' + ANNOTATION_FILE, 'r_' + ANNOTATION_METADATA, + 'r_' + ANNOTATION_SUMMARY]) + written = pandas.read_csv(run / ('r_' + ANNOTATION_FILE), sep='\t', + keep_default_na=False) + self.assertEqual(written['hit_id'].tolist(), result['hit_id'].tolist()) + metadata = yaml.safe_load((run / ('r_' + ANNOTATION_METADATA)).read_text()) + self.assertEqual(metadata['query'], {'path': str(tmp / 'q.tsv'), 'rows': 2}) + self.assertEqual(metadata['reference'], {'path': str(tmp / 'r.tsv'), 'rows': 1}) + self.assertEqual(metadata['options']['partition_by'], 'VJ') + self.assertEqual(metadata['results'], { + 'query_rows_compared': 2, 'query_rows_with_hits': 1, + 'annotation_rows': 1, 'query_rows_scored_directly': 0, 'hits': 1}) + self.assertFalse(metadata['options']['full_dist_matrix']) + self.assertNotIn('command', metadata) # called from Python, not the commandline + + def test_a_downloaded_reference_names_its_database_and_table(self): + with tempfile.TemporaryDirectory() as tmp: + database = Path(tmp) / 'databases' / 'iedb' + database.mkdir(parents=True) + airr([heavy('ref1', 'CARDYW')]).to_csv(database / 'iedb_bcr.tsv', + sep='\t', index=False) + (database / 'download_metadata.yml').write_text(yaml.safe_dump( + {'source': 'iedb', 'runs': [{'finished': '2026-10-14T17:52:12Z'}]})) + + annotateSpecificity(frame([heavy('q1', 'CARDYW')]), + database / 'iedb_bcr.tsv', outdir=Path(tmp) / 'run') + + run = Path(tmp) / 'run' + self.assertEqual(sorted(p.name for p in run.iterdir()), + ['iedb_bcr_' + ANNOTATION_FILE, + 'iedb_bcr_' + ANNOTATION_METADATA, + 'iedb_bcr_' + ANNOTATION_SUMMARY]) + metadata = yaml.safe_load(Path(run, 'iedb_bcr_' + ANNOTATION_METADATA).read_text()) + self.assertEqual(metadata['query']['path'], None) + self.assertEqual(metadata['reference']['source'], 'iedb') + self.assertEqual(metadata['reference']['downloaded'], '2026-10-14T17:52:12Z') + + def test_metadata_counts_the_rows_scored_directly(self): + with tempfile.TemporaryDirectory() as tmp: + result = annotateSpecificity( + frame([heavy('q1', 'A' * 20), heavy('q2', 'CARDYW')]), + airr([heavy('ref1', 'CARDYW'), heavy('ref2', 'B' * 15 + 'A' * 5)]), + identity_threshold=0.25, outdir=tmp) + + metadata = yaml.safe_load(Path(tmp, 'reference_' + ANNOTATION_METADATA) + .read_text()) + self.assertEqual(metadata['results']['query_rows_scored_directly'], 1) + self.assertEqual(result[['sequence_id', 'hit_id']].values.tolist(), + [['q1', 'ref2'], ['q2', 'ref1']]) + self.assertEqual(metadata['results']['hits'], 2) + + def test_symscan_names_skipped_rows_by_sequence_id(self): + with self.assertLogs('immvoke.Sources.Specificity.Symscan', 'WARNING') as logs: + runSymscanAnnotation(frame([heavy('q1', 'A' * 20)]), + airr([heavy('ref1', 'B' * 15 + 'A' * 5)]), + ['cdr3_aa'], 0.75, threshold_is_fraction=True) + + self.assertIn('may miss hits: sequence_id q1', logs.output[0]) + + def test_metadata_takes_numpy_options(self): + with tempfile.TemporaryDirectory() as tmp: + annotateSpecificity(frame([heavy('q1', 'CARDYW')]), + airr([heavy('ref1', 'CARDYW')]), + identity_threshold=numpy.float64(0.9), outdir=tmp) + + metadata = yaml.safe_load(Path(tmp, 'reference_' + ANNOTATION_METADATA).read_text()) + self.assertEqual(metadata['options']['identity_threshold'], 0.9) + + def test_a_run_without_hits_fills_text_columns_cleanly(self): + """No hit at all leaves the text columns all missing, which pandas + reads as float: filling those with '' must not warn (a future pandas + raises instead). Every row is then dropped, so the table is empty + but still has its columns.""" + with warnings.catch_warnings(): + warnings.simplefilter('error') + result = annotateSpecificity(frame([heavy('q1', 'CQQQQQ')]), + airr([heavy('ref1', 'CARDYW')])) + + self.assertEqual(len(result), 0) + self.assertEqual(result.columns.tolist(), + ['sequence_id', 'identity', 'hit_id', 'antigen', + 'epitope', 'antigen_source_species']) + + def test_only_the_best_hits_are_reported(self): + """A query keeps the hits with the highest identity only; the full + matrix keeps every hit within the threshold.""" + queries = [heavy('q1', 'CARDYW'), heavy('q2', 'CQQQQQ'), heavy('q3', 'CARDYY'), + heavy('q4', 'CARDYW', v_call='IGHV3')] + reference = [heavy('ref1', 'CARDYW'), heavy('ref2', 'CARDYY')] + with tempfile.TemporaryDirectory() as tmp: + result = annotateSpecificity(frame(queries), airr(reference), + identity_threshold=0.8, partition_by='V', + outdir=tmp, full_dist_matrix=True) + + matrix = pandas.read_csv(Path(tmp) / ('reference_' + FULL_MATRIX_FILE), + sep='\t') + metadata = yaml.safe_load(Path(tmp, 'reference_' + ANNOTATION_METADATA) + .read_text()) + self.assertEqual(result[['sequence_id', 'hit_id']].values.tolist(), + [['q1', 'ref1'], ['q3', 'ref2']]) + self.assertEqual(result['identity'].tolist()[0], 1.0) + self.assertEqual(matrix.columns.tolist(), ['sequence_id', 'hit_id', 'identity']) + # The matrix is not thresholded: q2 matches nothing, yet its + # distance to each reference is still reported. + self.assertEqual(matrix.values.tolist(), + [['q1', 'ref1', 1.0], ['q1', 'ref2', 0.8333], + ['q2', 'ref1', 0.1667], ['q2', 'ref2', 0.1667], + ['q3', 'ref2', 1.0], ['q3', 'ref1', 0.8333]]) + self.assertEqual(metadata['results']['hits'], 4) + + def test_tied_best_hits_are_all_reported_with_their_antigens(self): + """Hits at the same identity are all kept, one row each; a hit + standing for several reference rows (IEDB lists a receptor once per + epitope) gets a row per reference row, so no cell ever holds two + values -- which matters because IEDB's own epitope names contain + both the ';' and the '|' a joined cell would need.""" + queries = [heavy('q1', 'CARDYW')] + reference = [ + heavy('ref1', 'CARDYA', iedb_antigen='Spike', iedb_epitope_name='G58, L59', + iedb_antigen_source_species='SARS-CoV2'), + heavy('ref1', 'CARDYA', iedb_antigen='Spike', iedb_epitope_name='NWFDIT', + iedb_antigen_source_species='SARS-CoV2'), + heavy('ref2', 'CARDYQ', iedb_antigen='hemagglutinin', + iedb_epitope_name='K180', iedb_antigen_source_species='Influenza A virus'), + heavy('ref3', 'CARDQQ', iedb_antigen='far', iedb_epitope_name='far', + iedb_antigen_source_species='far')] + + result = annotateSpecificity(frame(queries), airr(reference), + identity_threshold=0.6) + + # ref1 twice (its two epitopes) and ref2 once; ref3 is too far. + self.assertEqual(result[['hit_id', 'identity', 'antigen', 'epitope', + 'antigen_source_species']].values.tolist(), + [['ref1', 0.8333, 'Spike', 'G58, L59', 'SARS-CoV2'], + ['ref1', 0.8333, 'Spike', 'NWFDIT', 'SARS-CoV2'], + ['ref2', 0.8333, 'hemagglutinin', 'K180', + 'Influenza A virus']]) + self.assertEqual(result['sequence_id'].tolist(), ['q1'] * 3) + + def test_the_summary_explains_the_run_and_its_files(self): + """A reader who opens the folder and nothing else should be able to + tell what was run and what each file holds.""" + with tempfile.TemporaryDirectory() as tmp: + annotateSpecificity(frame([heavy('q1', 'CARDYW'), heavy('q2', 'CQQQQQ')]), + airr([heavy('ref1', 'CARDYA')]), outdir=tmp, + identity_threshold=0.8, partition_by='VJ', + full_dist_matrix=True) + text = Path(tmp, 'reference_' + ANNOTATION_SUMMARY).read_text() + + self.assertIn('A hit is 80% identity or better, by hamming distance.', text) + self.assertIn('Only rows of the same V and J gene were compared.', text) + self.assertIn('query rows compared', text) + self.assertIn('query rows with a hit', text) + # every file it wrote is named and described + for name in ('reference_' + ANNOTATION_FILE, + 'reference_' + ANNOTATION_METADATA, + 'reference_' + FULL_MATRIX_FILE): + self.assertIn(name, text) + self.assertIn('the identity threshold does not apply here', text) + + def test_the_summary_says_what_a_paired_run_compared(self): + def antibody(cell_id, heavy_cdr3, light_cdr3): + return [dict(cell_id=cell_id, locus='IGH', cdr3_aa=heavy_cdr3, + v_call='IGHV1', j_call='IGHJ1'), + dict(cell_id=cell_id, locus='IGK', cdr3_aa=light_cdr3, + v_call='IGKV1', j_call='IGKJ1')] + with tempfile.TemporaryDirectory() as tmp: + annotateSpecificity(frame(antibody('ab1', 'CARDYW', 'CQQYNS')), + airr(antibody('refA', 'CARDYW', 'CQQYNS')), + paired=True, outdir=tmp) + text = Path(tmp, 'reference_' + ANNOTATION_SUMMARY).read_text() + + self.assertIn('heavy and light chain of the same antibody', text) + self.assertIn('query antibodies', text) + + def test_a_rerun_without_the_matrix_removes_the_one_left_behind(self): + """A run writes the whole bundle: a matrix from an earlier run would + otherwise sit beside an overwritten table describing a comparison + that no longer happened.""" + queries, reference = [heavy('q1', 'CARDYW')], [heavy('ref1', 'CARDYY')] + with tempfile.TemporaryDirectory() as tmp: + matrix = Path(tmp) / ('reference_' + FULL_MATRIX_FILE) + annotateSpecificity(frame(queries), airr(reference), outdir=tmp, + identity_threshold=0.8, full_dist_matrix=True) + self.assertTrue(matrix.is_file()) + + # The same folder, this time an exact match and no matrix asked for. + annotateSpecificity(frame(queries), airr(reference), outdir=tmp) + + self.assertFalse(matrix.is_file()) + self.assertEqual(sorted(x.name for x in Path(tmp).iterdir()), + ['reference_' + ANNOTATION_FILE, + 'reference_' + ANNOTATION_METADATA, + 'reference_' + ANNOTATION_SUMMARY]) + + def test_the_full_matrix_ignores_the_identity_threshold(self): + """The table answers what a query matched, so --identity-threshold + decides it; the matrix answers how close the query is to everything + it was compared with, so the threshold does not apply -- 'far' is + well below it and still reported.""" + queries = [heavy('q1', 'CARDYWSSTA')] + reference = [heavy('exact', 'CARDYWSSTA'), heavy('one_off', 'CARDYWSSTG'), + heavy('three_off', 'CARDYWSGGG'), heavy('far', 'CGGGGGGGGG')] + with tempfile.TemporaryDirectory() as tmp: + result = annotateSpecificity(frame(queries), airr(reference), outdir=tmp, + identity_threshold=0.7, full_dist_matrix=True) + matrix = pandas.read_csv(Path(tmp) / ('reference_' + FULL_MATRIX_FILE), + sep='\t') + metadata = yaml.safe_load(Path(tmp, 'reference_' + ANNOTATION_METADATA) + .read_text()) + + self.assertEqual(result['hit_id'].tolist(), ['exact']) + self.assertEqual(matrix[['hit_id', 'identity']].values.tolist(), + [['exact', 1.0], ['one_off', 0.9], ['three_off', 0.7], + ['far', 0.1]]) + self.assertEqual(metadata['results']['full_matrix_rows'], 4) + + def test_the_full_matrix_holds_only_pairs_the_filters_admit(self): + """--partition-by still applies: a reference of another V or J gene + was never compared, so it has no distance to report. Under + 'hamming' a different CDR3 length has no distance either.""" + queries = [heavy('q1', 'CARDYW')] + reference = [heavy('sameVJ', 'CAKKKK'), + heavy('otherV', 'CAKKKK', v_call='IGHV9-9'), + heavy('otherJ', 'CAKKKK', j_call='IGHJ9'), + heavy('longer', 'CARDYWAAA')] + + def matrix(**kwargs): + with tempfile.TemporaryDirectory() as tmp: + annotateSpecificity(frame(queries), airr(reference), outdir=tmp, + full_dist_matrix=True, **kwargs) + return pandas.read_csv(Path(tmp) / ('reference_' + FULL_MATRIX_FILE), + sep='\t')['hit_id'].tolist() + + self.assertEqual(matrix(partition_by='VJ'), ['sameVJ']) + self.assertEqual(sorted(matrix(partition_by='V')), ['otherJ', 'sameVJ']) + self.assertEqual(sorted(matrix()), ['otherJ', 'otherV', 'sameVJ']) + + def test_no_pairs_file_and_no_matrix_unless_asked(self): + with tempfile.TemporaryDirectory() as tmp: + annotateSpecificity(frame([heavy('q1', 'CARDYW')]), + airr([heavy('ref1', 'CARDYW')]), outdir=tmp) + + self.assertEqual(sorted(p.name for p in Path(tmp).iterdir()), + ['reference_' + ANNOTATION_FILE, + 'reference_' + ANNOTATION_METADATA, + 'reference_' + ANNOTATION_SUMMARY]) + + def test_paired_full_matrix_names_the_matched_chains(self): + def antibody(cell_id, heavy_cdr3, light_cdr3): + return [dict(cell_id=cell_id, locus='IGH', cdr3_aa=heavy_cdr3, + v_call='IGHV1', j_call='IGHJ1'), + dict(cell_id=cell_id, locus='IGK', cdr3_aa=light_cdr3, + v_call='IGKV1', j_call='IGKJ1')] + queries = antibody('ab1', 'CARDYW', 'CQQYNS') + antibody('ab2', 'CAKKKK', 'CQQAAA') + reference = antibody('refA', 'CARDYW', 'CQQYNS') + antibody('refB', 'CARDYW', + 'CQQTTT') + with tempfile.TemporaryDirectory() as tmp: + annotateSpecificity(frame(queries), airr(reference), paired=True, + outdir=tmp, full_dist_matrix=True) + + matrix = pandas.read_csv(Path(tmp) / ('reference_' + FULL_MATRIX_FILE), + sep='\t') + metadata = yaml.safe_load(Path(tmp, 'reference_' + ANNOTATION_METADATA).read_text()) + self.assertEqual(matrix.values.tolist(), [['ab1', 'refA', 'both'], + ['ab1', 'refB', 'heavy']]) + self.assertEqual(metadata['results'], { + 'query_antibodies': 2, 'classification': {'both': 1, 'neither': 1}, + 'pairs': 2, 'skipped_query_rows': 0, 'full_matrix_rows': 2}) + + def test_genes_match_by_name_any_allele(self): + queries = [heavy('q1', 'CARDYW', v_call='IGHV1-2*02', j_call='IGHJ4*02')] + reference = [ + heavy('sameAllele', 'CARDYW', v_call='IGHV1-2*02', j_call='IGHJ4*02'), + heavy('otherAllele', 'CARDYW', v_call='IGHV1-2*04', j_call='IGHJ4*02'), + heavy('geneOnly', 'CARDYW', v_call='IGHV1-2', j_call='IGHJ4'), + heavy('several', 'CARDYW', v_call='IGHV3-23*01, IGHV1-2*04', + j_call='IGHJ4*02'), + heavy('blankJ', 'CARDYW', v_call='IGHV1-2*04', j_call=''), + heavy('sameFamily', 'CARDYW', v_call='IGHV1-69*01', j_call='IGHJ4*02'), + heavy('otherFamily', 'CARDYW', v_call='IGHV3-23*01', j_call='IGHJ4*02'), + ] + + hits = self._hits(queries, reference, partition_by='VJ') + + # blankJ: a blank call agrees with nothing; sameFamily: another gene + self.assertEqual(set(hits[0].split(';')), + {'sameAllele', 'otherAllele', 'geneOnly', 'several'}) + + def test_messy_gene_calls_match_brute_force(self): + calls = ['IGHV1-2*02', 'IGHV1-2*04', 'IGHV1-2', '', 'IGVL3-10', + 'IGHV3-23*01, IGHV1-2*02', 'IGHV1-69*01'] + queries = [heavy('q%d' % i, cdr3, v_call=calls[i % len(calls)], + j_call=('IGHJ4*02', '', 'IGHJ6')[i % 3]) + for i, cdr3 in enumerate(['CARDYW', 'CARQYW', 'CARDDYW', 'CAKDYW'] * 3)] + reference = [heavy('r%d' % i, cdr3, v_call=calls[(i + 2) % len(calls)], + j_call=('IGHJ4', 'IGHJ4*02', '')[i % 3]) + for i, cdr3 in enumerate(['CARDYW', 'CARQYW', 'CARDYY', 'CQ'] * 3)] + for mode in ({}, {'identity_threshold': 0.8}, + {'identity_threshold': 0.8, 'distance_metric': 'levenshtein'}): + with self.subTest(**mode): + kwargs = dict(partition_by='VJ', **mode) + result = annotateSpecificity(frame(queries), airr(reference), **kwargs) + pandas.testing.assert_frame_equal( + result, bruteForceSpecificity(queries, reference, **kwargs), + check_dtype=False) + + def _pairedCell(self, heavies, lights=None, **kwargs): + """ + Classify one query cell `ab1` against reference antibody `refA` + (heavy CDR3 CARDYW, light CDR3 CQQYNS). `heavies`/`lights` are + (cdr3_aa, extra fields) per query chain. + """ + def chain(locus, cdr3, v, j, extra): + return dict(cell_id='ab1', locus=locus, cdr3_aa=cdr3, v_call=v, + j_call=j, **extra) + lights = lights if lights is not None else [('CQQYNS', {})] + queries = ([chain('IGH', c, 'IGHV1-2*02', 'IGHJ4*02', e) for c, e in heavies] + + [chain('IGK', c, 'IGKV1-39*01', 'IGKJ1*01', e) for c, e in lights]) + reference = [chain('IGH', 'CARDYW', 'IGHV1-2*02', 'IGHJ4*02', {}), + chain('IGK', 'CQQYNS', 'IGKV1-39*01', 'IGKJ1*01', {})] + reference = [dict(row, cell_id='refA') for row in reference] + + result = annotateSpecificity(frame(queries), airr(reference), paired=True, + **kwargs) + return result['classification'].tolist() + + def test_several_heavy_chains_keep_the_most_abundant(self): + self.assertEqual(self._pairedCell([('CAKQQQ', {'duplicate_count': '5'}), + ('CARDYW', {'duplicate_count': '50'})]), + ['both']) + self.assertEqual(self._pairedCell([('CAKQQQ', {'duplicate_count': '50'}), + ('CARDYW', {'duplicate_count': '5'})]), + ['light_only']) + + def test_several_light_chains_keep_the_most_abundant(self): + heavy = [('CARDYW', {'duplicate_count': '9'})] + self.assertEqual(self._pairedCell(heavy, [('CQQAAA', {'duplicate_count': '3'}), + ('CQQYNS', {'duplicate_count': '7'})]), + ['both']) + self.assertEqual(self._pairedCell(heavy, [('CQQAAA', {'duplicate_count': '7'}), + ('CQQYNS', {'duplicate_count': '3'})]), + ['heavy_only']) + + def test_ties_and_missing_counts_keep_the_first_row(self): + tie = [('CAKQQQ', {'duplicate_count': '5'}), ('CARDYW', {'duplicate_count': '5'})] + no_column = [('CAKQQQ', {}), ('CARDYW', {})] + blank = [('CARDYW', {'duplicate_count': ''}), ('CAKQQQ', {'duplicate_count': '1'})] + self.assertEqual(self._pairedCell(tie), ['light_only']) + self.assertEqual(self._pairedCell(no_column), ['light_only']) + # A blank count ranks below any number. + self.assertEqual(self._pairedCell(blank), ['light_only']) + + def test_the_count_field_can_be_named(self): + heavies = [('CAKQQQ', {'duplicate_count': '50', 'umi_count': '1'}), + ('CARDYW', {'duplicate_count': '5', 'umi_count': '9'})] + self.assertEqual(self._pairedCell(heavies), ['light_only']) + self.assertEqual(self._pairedCell(heavies, count_field='umi_count'), ['both']) + + def test_paired_leaves_out_rows_without_a_cell_id(self): + """Blank cell_ids are not one antibody: each such chain has no + partner, so it is left out rather than paired with the others.""" + def chain(cell_id, locus, cdr3): + v, j = ('IGHV1', 'IGHJ1') if locus == 'IGH' else ('IGKV1', 'IGKJ1') + return dict(cell_id=cell_id, locus=locus, cdr3_aa=cdr3, v_call=v, j_call=j) + queries = [chain('', 'IGH', 'CARDYW'), chain('', 'IGH', 'CAKKKK'), + chain('', 'IGK', 'CQQYNS'), chain('ab1', 'IGH', 'CAKKKK'), + chain('ab1', 'IGK', 'CQQYNS')] + reference = [chain('refA', 'IGH', 'CAKKKK'), chain('refA', 'IGK', 'CQQYNS')] + + with self.assertLogs(ANNOTATE, 'WARNING') as logs: + result = annotateSpecificity(frame(queries), airr(reference), paired=True) + + self.assertIn('query: 3 rows have no cell_id', logs.output[0]) + self.assertEqual(result[['cell_id', 'classification']].values.tolist(), + [['ab1', 'both']]) + + def test_paired_compares_genes_by_name(self): + rows = [dict(cell_id='ab1', locus='IGH', cdr3_aa='CARDYW', v_call='IGHV1-2*02', + j_call='IGHJ4*02'), + dict(cell_id='ab1', locus='IGK', cdr3_aa='CQQYNS', v_call='IGKV1-39*01', + j_call='IGKJ1*01')] + same_gene = [dict(rows[0], cell_id='refA', v_call='IGHV1-2*04'), + dict(rows[1], cell_id='refA')] + other_gene = [dict(rows[0], cell_id='refA', v_call='IGHV1-69*01'), + dict(rows[1], cell_id='refA')] + + for reference, classification in ((same_gene, 'both'), + (other_gene, 'light_only')): + with self.subTest(classification): + result = annotateSpecificity(frame(rows), airr(reference), paired=True) + self.assertEqual(result['classification'].tolist(), [classification]) + reference = same_gene + + # A blank J-gene call agrees with nothing, so the light chain cannot match. + blank_j = [rows[0], dict(rows[1], j_call='')] + result = annotateSpecificity(frame(blank_j), airr(reference), paired=True) + self.assertEqual(result['classification'].tolist(), ['heavy_only']) + + def test_invalid_arguments_raise(self): + queries = [heavy('q1', 'CARDYW', cdr1_aa='GFT', cdr2_aa='ISY')] + cases = [{'sequence_region': 'CDR4'}, {'partition_by': 'J'}, + {'distance_metric': 'blosum'}, + {'identity_threshold': 0}, {'identity_threshold': 1.5}, + {'sequence_region': 'CDRs', 'identity_threshold': 0.9, + 'distance_metric': 'levenshtein'}, + {'sequence_region': 'VDJ'}] # no sequence_aa column + for kwargs in cases: + with self.subTest(**kwargs), self.assertRaises(AnnotationError): + self._hits(queries, queries, **kwargs) + + +class TestReferenceContract(unittest.TestCase): + """ + The reference is always an AIRR rearrangement table, from whichever + database -- annotate must accept any such table and refuse anything else + """ + + def test_a_non_airr_reference_is_refused(self): + query = frame([heavy('q1', 'CARDYW', cell_id='ab1')]) + cov_abdab = frame([{'Name': 'ab1', 'VHorVHH': 'EVQL', 'CDRH3': 'CARDYW', + 'Heavy V Gene': 'IGHV1-2 (Human)'}]) + for paired in (False, True): + with self.subTest(paired=paired), \ + self.assertRaisesRegex(AnnotationError, + 'reference is not an AIRR rearrangement'): + annotateSpecificity(query, cov_abdab, paired=paired) + + def test_a_missing_required_field_is_named(self): + reference = airr([heavy('ref1', 'CARDYW')]).drop(columns='productive') + with self.assertRaisesRegex(AnnotationError, 'missing required fields ' + r'\(productive\)'): + annotateSpecificity(frame([heavy('q1', 'CARDYW')]), reference) + + def test_the_query_needs_only_the_columns_it_uses(self): + query = frame([{'locus': 'IGH', 'cdr3_aa': 'CARDYW'}]) + result = annotateSpecificity(query, airr([heavy('ref1', 'CARDYW')])) + self.assertEqual(result['hit_id'].tolist(), ['ref1']) + + def test_a_reference_written_by_immvoke_is_accepted(self): + """Every immvoke source that writes AIRR goes through writeAirr.""" + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / 'reference.tsv' + writeAirr([frame([heavy('ref1', 'CARDYW', sequence_aa='EVQLCARDYW')])], + path) + + result = annotateSpecificity(frame([heavy('q1', 'CARDYW')]), path) + + self.assertEqual(result['hit_id'].tolist(), ['ref1']) + + def test_any_airr_reference_finds_its_own_rows(self): + """ + References shaped like different databases: IEDB (allele calls, blank + J-genes, provenance columns), a user's own IgBLAST/Change-O table + (several calls in one cell, extra AIRR fields), and the bare minimum. + Every query row is in its reference, so must find itself, and the + whole table must match brute force. + """ + cdr3s = ['CARDYW', 'CARQYW', 'CAKDGYSSGWYFDYW', 'CTTDLW'] + references = { + 'iedb': [heavy('%d_1' % i, c, v_call='IGHV1-2*0%d' % (i % 2 + 1), + j_call='' if i % 2 else 'IGHJ4*02', + epitope_name='spike', immvoke_source='iedb') + for i, c in enumerate(cdr3s)], + 'own': [heavy('seq%d' % i, c, v_call='IGHV3-23*01, IGHV3-23D*01', + j_call='IGHJ4*02', productive='T', clone_id=str(i)) + for i, c in enumerate(cdr3s)], + 'minimal': [{'sequence_id': 'r%d' % i, 'cdr3_aa': c} + for i, c in enumerate(cdr3s)], + } + modes = [{}, {'identity_threshold': 0.8}, + {'identity_threshold': 0.8, 'distance_metric': 'levenshtein'}] + for name, rows in references.items(): + reference = airr(rows) + queries = frame(rows) + gated = [{'partition_by': 'VJ'}] if 'v_call' in rows[0] else [] + for mode in modes + gated: + with self.subTest(reference=name, **mode): + result = annotateSpecificity(queries, reference, **mode) + pandas.testing.assert_frame_equal( + result.sort_values(['sequence_id', 'hit_id'], + kind='stable').reset_index(drop=True), + bruteForceSpecificity(rows, rows, **mode), + check_dtype=False) + found = {sid: set() for sid in result['sequence_id']} + for sid, hit in zip(result['sequence_id'], result['hit_id']): + found[sid].add(hit) + for row in rows: + hits = found.get(row['sequence_id'], set()) + # Under a partition, a blank J-gene call agrees with + # nothing -- not even the row's own. + if mode.get('partition_by') and not row['j_call']: + self.assertNotIn(row['sequence_id'], hits) + else: + self.assertIn(row['sequence_id'], hits) + + def test_no_heavy_chain_rows_left_raises(self): + light = dict(sequence_id='k1', locus='IGK', cdr3_aa='CQQYNS', + v_call='IGKV1-39*01', j_call='IGKJ1*01') + tcr = dict(light, sequence_id='b1', locus='TRB', v_call='TRBV20-1*01', + j_call='TRBJ2-7*01') + for query, reference in (([light], [heavy('ref1', 'CARDYW')]), + ([heavy('q1', 'CARDYW')], [tcr])): + with self.subTest(query=query[0]['locus'], + reference=reference[0]['locus']), \ + self.assertRaisesRegex(AnnotationError, 'no heavy chain rows'): + annotateSpecificity(frame(query), airr(reference)) + + +class TestAnnotateCli(unittest.TestCase): + """ + Tests for `immvoke annotate`, a thin wrapper over annotateSpecificity + """ + + BASE = ['annotate', '--query', 'q.tsv', '--reference', 'r.tsv', '--outdir', 'run'] + + def test_annotate_is_a_top_level_subcommand(self): + with mock.patch('sys.stdout', new_callable=io.StringIO) as out: + with self.assertRaises(SystemExit): + getArgParser().parse_args(['annotate', '--help']) + + text = out.getvalue() + for flag in ('--query', '--reference', '--sequence-region', '--partition-by', + '--paired', '--identity-threshold', '--distance-metric', + '--count-field', '--threads', '--outdir', '--full-dist-matrix'): + self.assertIn(flag, text) + for gone in ('--backend', '--gene-match', 'compairr', 'python', '-o,', '--out '): + self.assertNotIn(gone, text.lower()) + + def test_defaults(self): + args = getArgParser().parse_args(self.BASE) + + self.assertEqual((args.sequence_region, args.partition_by, args.paired, + args.identity_threshold, args.distance_metric), + ('CDR3H', None, False, None, 'hamming')) + + def test_invalid_choices_are_rejected(self): + for extra in (['--backend', 'symscan'], ['--partition-by', 'J'], + ['--sequence-region', 'CDR4'], ['--distance-metric', 'x'], + ['--gene-match', 'gene']): + with self.subTest(extra=extra), \ + mock.patch('sys.stderr', new_callable=io.StringIO), \ + self.assertRaises(SystemExit): + getArgParser().parse_args(self.BASE + extra) + + def test_required_flags_are_enforced(self): + with mock.patch('sys.stderr', new_callable=io.StringIO): + with self.assertRaises(SystemExit): + getArgParser().parse_args(['annotate']) + + def test_handle_annotate_forwards_every_flag(self): + args = getArgParser().parse_args(self.BASE + [ + '--sequence-region', 'CDRs', '--partition-by', 'VJL', + '--identity-threshold', '0.9', '--threads', '4', + '--count-field', 'umi_count', '--full-dist-matrix']) + + with mock.patch('immvoke.Sources.Specificity.AnnotateCli.annotateSpecificity') \ + as annotate: + self.assertEqual(handleAnnotate(args), 0) + + annotate.assert_called_once_with( + query=Path('q.tsv'), reference=Path('r.tsv'), sequence_region='CDRs', + partition_by='VJL', paired=False, identity_threshold=0.9, + distance_metric='hamming', threads=4, + outdir=Path('run'), count_field='umi_count', full_dist_matrix=True) + + def test_end_to_end(self): + with tempfile.TemporaryDirectory() as tmp: + tmp = Path(tmp) + frame([heavy('q1', 'CARDYW'), heavy('q2', 'CQQQQQ')]).to_csv( + tmp / 'q.tsv', sep='\t', index=False) + airr([heavy('ref1', 'CARQYW')]).to_csv(tmp / 'r.tsv', sep='\t', + index=False) + argv = ['annotate', '--query', str(tmp / 'q.tsv'), '--reference', + str(tmp / 'r.tsv'), '--outdir', str(tmp / 'run'), + '--partition-by', 'VJ', '--identity-threshold', '0.8'] + args = getArgParser().parse_args(argv) + + with mock.patch('sys.argv', ['/venv/bin/immvoke'] + argv): + self.assertEqual(handleAnnotate(args), 0) + + written = pandas.read_csv(tmp / 'run' / ('r_' + ANNOTATION_FILE), sep='\t', + keep_default_na=False) + metadata = yaml.safe_load(Path(tmp, 'run', 'r_' + ANNOTATION_METADATA).read_text()) + self.assertEqual(written['hit_id'].tolist(), ['ref1']) + self.assertEqual(written['sequence_id'].tolist(), ['q1']) + self.assertEqual(metadata['command'], ' '.join(['immvoke'] + argv)) diff --git a/src/immvoke/Sources/Specificity/tests/test_Asd.py b/src/immvoke/Sources/Specificity/tests/test_Asd.py new file mode 100644 index 0000000..d5fe2d8 --- /dev/null +++ b/src/immvoke/Sources/Specificity/tests/test_Asd.py @@ -0,0 +1,228 @@ +""" +Unit tests for the ASD specificity source +""" + +# Info +__author__ = 'Pramod Shinde' + +# Imports +import sys +import tempfile +import types +import unittest +from pathlib import Path +from unittest import mock + +import pandas + +# Immvoke imports +from immvoke.Exceptions import AsdFetchError, AsdParseError +from immvoke.Http import HttpClient +from immvoke.Schema import loadSchema +from immvoke.Sources.Base import DataUnit, Query +from immvoke.Sources.Specificity import REGISTRY, getSpecificitySource +from immvoke.Sources.Specificity.Asd import COLLECTION, DRIVE_FOLDER_URL, AsdSource +from tests.FakeHttp import FakeResponse, FakeSession + + +def makeClient(handler): + """Build an HttpClient with no politeness delay over a scripted session.""" + return HttpClient(delay=0, backoff=0, session=FakeSession(handler)) + + +class FakeGdown(types.ModuleType): + """ + A stand-in for the gdown module, driven by a callback. + """ + + def __init__(self, make_files): + """ + Arguments: + make_files (callable): called as make_files(Path) once + download_folder is invoked, to populate the destination directory. + """ + super().__init__('gdown') + self._make_files = make_files + self.calls = [] + + def download_folder(self, url, output=None, quiet=True, use_cookies=False): + """Record the call and populate output via the configured callback.""" + self.calls.append({'url': url, 'output': output}) + self._make_files(Path(output)) + + +def writeParquet(path, records): + """Write a small parquet file at path from a list of dict records.""" + pandas.DataFrame.from_records(records).to_parquet(path) + + +class TestAsdSource(unittest.TestCase): + """ + Tests for the ASD specificity source + """ + + def test_registered(self): + self.assertIs(REGISTRY['asd'], AsdSource) + client = makeClient(lambda *a: FakeResponse(200)) + self.assertIsInstance(getSpecificitySource('asd', client), AsdSource) + + def test_packaged_schema_offers_a_dataset_field(self): + schema = loadSchema('asd') + collection = schema.getCollection(COLLECTION) + self.assertIn('dataset', collection.field_names) + self.assertIn('hiv', collection.getField('dataset').values) + + def test_search_units_returns_one_unit_carrying_the_filters(self): + source = AsdSource(makeClient(lambda *a: FakeResponse(200))) + + units = source.searchUnits( + Query(collection=COLLECTION, filters={'dataset': 'hiv'})) + + self.assertEqual(len(units), 1) + self.assertEqual(units[0].unit_id, COLLECTION) + self.assertEqual(units[0].url, DRIVE_FOLDER_URL) + self.assertEqual(units[0].metadata, {'dataset': 'hiv'}) + + def test_fetch_unit_downloads_the_folder_with_gdown(self): + def make_files(path): + path.mkdir(parents=True, exist_ok=True) + (path / 'part-0.snappy.parquet').write_bytes(b'PAR1fake-parquet-bytes') + + fake_gdown = FakeGdown(make_files) + source = AsdSource(makeClient(lambda *a: FakeResponse(200))) + unit = DataUnit(unit_id=COLLECTION, collection=COLLECTION, + url=DRIVE_FOLDER_URL) + + with tempfile.TemporaryDirectory() as outdir: + with mock.patch.dict(sys.modules, {'gdown': fake_gdown}): + result = source.fetchUnit(unit, Path(outdir)) + + self.assertEqual(len(fake_gdown.calls), 1) + self.assertEqual(fake_gdown.calls[0]['url'], DRIVE_FOLDER_URL) + self.assertTrue(result.path.is_dir()) + self.assertTrue(any(result.path.glob('*.parquet'))) + self.assertEqual(len(result.sha256), 64) + self.assertGreater(result.size_bytes, 0) + + def test_fetch_unit_skips_a_directory_that_already_has_parquet_files(self): + source = AsdSource(makeClient(lambda *a: FakeResponse(200))) + unit = DataUnit(unit_id=COLLECTION, collection=COLLECTION, + url=DRIVE_FOLDER_URL) + + def unexpected_call(*args, **kwargs): + raise AssertionError('gdown should not be called when resuming') + + with tempfile.TemporaryDirectory() as outdir: + dest = Path(outdir) / COLLECTION / COLLECTION + dest.mkdir(parents=True) + (dest / 'part-0.snappy.parquet').write_bytes(b'PAR1already-here') + + with mock.patch.dict(sys.modules, + {'gdown': types.SimpleNamespace( + download_folder=unexpected_call)}): + result = source.fetchUnit(unit, Path(outdir), resume=True) + + self.assertEqual(result.path, dest) + + def test_fetch_unit_raises_when_gdown_is_not_installed(self): + source = AsdSource(makeClient(lambda *a: FakeResponse(200))) + unit = DataUnit(unit_id=COLLECTION, collection=COLLECTION, + url=DRIVE_FOLDER_URL) + + with tempfile.TemporaryDirectory() as outdir: + # Setting a module to None in sys.modules makes the next `import + # gdown` raise ImportError, without requiring gdown to actually be + # absent from the environment running the tests. + with mock.patch.dict(sys.modules, {'gdown': None}): + with self.assertRaises(AsdFetchError): + source.fetchUnit(unit, Path(outdir)) + + def test_fetch_unit_raises_when_the_folder_has_no_parquet_files(self): + fake_gdown = FakeGdown(lambda path: path.mkdir(parents=True, exist_ok=True)) + source = AsdSource(makeClient(lambda *a: FakeResponse(200))) + unit = DataUnit(unit_id=COLLECTION, collection=COLLECTION, + url=DRIVE_FOLDER_URL) + + with tempfile.TemporaryDirectory() as outdir: + with mock.patch.dict(sys.modules, {'gdown': fake_gdown}): + with self.assertRaises(AsdFetchError): + source.fetchUnit(unit, Path(outdir)) + + def test_read_unit_reads_every_partition_as_a_chunk(self): + source = AsdSource(makeClient(lambda *a: FakeResponse(200))) + unit = DataUnit(unit_id=COLLECTION, collection=COLLECTION, + url=DRIVE_FOLDER_URL) + + with tempfile.TemporaryDirectory() as outdir: + path = Path(outdir) + writeParquet(path / 'part-0.parquet', + [{'dataset': 'hiv', 'heavy_sequence': 'QVQ'}]) + writeParquet(path / 'part-1.parquet', + [{'dataset': 'covid-19', 'heavy_sequence': 'EVQ'}]) + + metadata, chunks = source.readUnit(path, unit) + frame = pandas.concat(list(chunks), ignore_index=True) + + self.assertEqual(metadata, {}) + self.assertEqual(sorted(frame['dataset'].tolist()), ['covid-19', 'hiv']) + + def test_read_unit_raises_when_no_parquet_files_are_present(self): + source = AsdSource(makeClient(lambda *a: FakeResponse(200))) + unit = DataUnit(unit_id=COLLECTION, collection=COLLECTION, + url=DRIVE_FOLDER_URL) + + with tempfile.TemporaryDirectory() as outdir: + with self.assertRaises(AsdParseError): + source.readUnit(Path(outdir), unit) + + def test_normalize_chunk_serializes_nested_metadata_to_json(self): + source = AsdSource(makeClient(lambda *a: FakeResponse(200))) + unit = DataUnit(unit_id=COLLECTION, collection=COLLECTION, + url=DRIVE_FOLDER_URL, metadata={}) + chunk = pandas.DataFrame([ + {'dataset': 'hiv', 'metadata': {'target_pdb': '5ovw'}}, + ]) + report = {'rows_in': 0, 'rows_out': 0} + + frame = source.normalizeChunk({}, chunk, unit, 0, report) + + self.assertEqual(frame['metadata'].iloc[0], '{"target_pdb": "5ovw"}') + self.assertEqual(frame['immvoke_source'].tolist(), ['asd']) + self.assertEqual(frame['immvoke_collection'].tolist(), [COLLECTION]) + self.assertEqual(len(frame['immvoke_row_hash'].iloc[0]), 12) + self.assertEqual(report['rows_in'], 1) + self.assertEqual(report['rows_out'], 1) + + def test_normalize_chunk_applies_the_dataset_filter(self): + source = AsdSource(makeClient(lambda *a: FakeResponse(200))) + unit = DataUnit(unit_id=COLLECTION, collection=COLLECTION, + url=DRIVE_FOLDER_URL, metadata={'dataset': 'hiv'}) + chunk = pandas.DataFrame([ + {'dataset': 'hiv', 'metadata': None}, + {'dataset': 'covid-19', 'metadata': None}, + ]) + report = {'rows_in': 0, 'rows_out': 0} + + frame = source.normalizeChunk({}, chunk, unit, 0, report) + + self.assertEqual(frame['dataset'].tolist(), ['hiv']) + self.assertEqual(report['rows_in'], 2) + self.assertEqual(report['rows_out'], 1) + + def test_normalize_chunk_keeps_everything_for_the_wildcard(self): + source = AsdSource(makeClient(lambda *a: FakeResponse(200))) + unit = DataUnit(unit_id=COLLECTION, collection=COLLECTION, + url=DRIVE_FOLDER_URL, metadata={'dataset': '*'}) + chunk = pandas.DataFrame([ + {'dataset': 'hiv', 'metadata': None}, + {'dataset': 'covid-19', 'metadata': None}, + ]) + + frame = source.normalizeChunk({}, chunk, unit, 0, + {'rows_in': 0, 'rows_out': 0}) + + self.assertEqual(len(frame), 2) + + +if __name__ == '__main__': + unittest.main() diff --git a/src/immvoke/Sources/Specificity/tests/test_CovAbDab.py b/src/immvoke/Sources/Specificity/tests/test_CovAbDab.py new file mode 100644 index 0000000..73db3d5 --- /dev/null +++ b/src/immvoke/Sources/Specificity/tests/test_CovAbDab.py @@ -0,0 +1,145 @@ +""" +Unit tests for the CoV-AbDab specificity source +""" + +# Info +__author__ = 'Pramod Shinde' + +# Imports +import tempfile +import unittest +from pathlib import Path + +import pandas + +# Immvoke imports +from immvoke.Exceptions import CovAbDabParseError +from immvoke.Http import HttpClient +from immvoke.Schema import loadSchema +from immvoke.Sources.Base import DataUnit, Query +from immvoke.Sources.Specificity import REGISTRY, getSpecificitySource +from immvoke.Sources.Specificity.CovAbDab import ( + COLLECTION, + HOME_URL, + CovAbDabSource, + findCsvUrl, +) +from tests.FakeHttp import FakeResponse, FakeSession, rangeHandler + + +def makeClient(handler): + """Build an HttpClient with no politeness delay over a scripted session.""" + return HttpClient(delay=0, backoff=0, session=FakeSession(handler)) + + +HOME_HTML = """ + +Download + +""" + + +class TestFindCsvUrl(unittest.TestCase): + """ + Tests for scraping the current bulk CSV link off the homepage + """ + + def test_resolves_a_relative_download_link_to_an_absolute_url(self): + client = makeClient(lambda *a: FakeResponse(200, HOME_HTML.encode())) + + url = findCsvUrl(client) + + self.assertEqual( + url, 'https://opig.stats.ox.ac.uk/webapps/covabdab/static/' + 'downloads/CoV-AbDab_080224.csv') + + def test_raises_when_no_csv_link_is_present(self): + client = makeClient( + lambda *a: FakeResponse(200, b'no links here')) + + with self.assertRaises(CovAbDabParseError): + findCsvUrl(client) + + +class TestCovAbDabSource(unittest.TestCase): + """ + Tests for the CoV-AbDab specificity source + """ + + def test_registered(self): + self.assertIs(REGISTRY['covabdab'], CovAbDabSource) + client = makeClient(lambda *a: FakeResponse(200)) + self.assertIsInstance(getSpecificitySource('covabdab', client), + CovAbDabSource) + + def test_packaged_schema_loads_and_offers_the_antibodies_collection(self): + schema = loadSchema('covabdab') + self.assertEqual(schema.collection_names, (COLLECTION,)) + + def test_search_units_scrapes_the_current_csv_url(self): + client = makeClient(lambda *a: FakeResponse(200, HOME_HTML.encode())) + source = CovAbDabSource(client) + + units = source.searchUnits(Query(collection=COLLECTION, filters={})) + + self.assertEqual(len(units), 1) + self.assertEqual(units[0].collection, COLLECTION) + self.assertEqual(units[0].unit_id, 'CoV-AbDab_080224.csv') + self.assertTrue(units[0].url.startswith(HOME_URL.rstrip('/')) or + units[0].url.startswith('https://opig.stats.ox.ac.uk')) + + def test_fetch_unit_downloads_the_resolved_csv(self): + """fetchUnit falls back to SourceBase's plain URL download.""" + body = b'a,b\n1,2\n' + client = makeClient(rangeHandler(body)) + source = CovAbDabSource(client) + unit = DataUnit(unit_id='CoV-AbDab_080224.csv', collection=COLLECTION, + url='https://opig.stats.ox.ac.uk/x.csv') + + with tempfile.TemporaryDirectory() as outdir: + result = source.fetchUnit(unit, Path(outdir)) + self.assertEqual(result.path.read_bytes(), body) + + def test_read_unit_strips_the_bom_and_parses_the_csv(self): + source = CovAbDabSource(makeClient(lambda *a: FakeResponse(200))) + unit = DataUnit(unit_id='x.csv', collection=COLLECTION, url='https://x/x.csv') + + with tempfile.TemporaryDirectory() as outdir: + path = Path(outdir) / 'x.csv' + path.write_bytes('Name,Ab or Nb\nCurtis_3548,Ab\n'.encode('utf-8-sig')) + metadata, chunks = source.readUnit(path, unit) + frame = pandas.concat(list(chunks), ignore_index=True) + + self.assertEqual(metadata, {}) + self.assertEqual(list(frame.columns), ['Name', 'Ab or Nb']) + self.assertEqual(frame.to_dict('records'), + [{'Name': 'Curtis_3548', 'Ab or Nb': 'Ab'}]) + + def test_normalize_chunk_adds_provenance_columns(self): + source = CovAbDabSource(makeClient(lambda *a: FakeResponse(200))) + unit = DataUnit(unit_id='x.csv', collection=COLLECTION, url='https://x/x.csv') + chunk = pandas.DataFrame([{'Name': 'Curtis_3548', 'Ab or Nb': 'Ab'}]) + report = {'rows_in': 0, 'rows_out': 0} + + frame = source.normalizeChunk({}, chunk, unit, 0, report) + + self.assertEqual(frame['immvoke_source'].tolist(), ['covabdab']) + self.assertEqual(frame['immvoke_collection'].tolist(), [COLLECTION]) + self.assertEqual(frame['immvoke_unit_id'].tolist(), ['x.csv']) + self.assertEqual(len(frame['immvoke_row_hash'].iloc[0]), 12) + self.assertEqual(report['rows_in'], 1) + self.assertEqual(report['rows_out'], 1) + + def test_normalize_chunk_handles_an_empty_frame(self): + source = CovAbDabSource(makeClient(lambda *a: FakeResponse(200))) + unit = DataUnit(unit_id='x.csv', collection=COLLECTION, url='https://x/x.csv') + chunk = pandas.DataFrame(columns=['Name', 'Ab or Nb']) + + frame = source.normalizeChunk({}, chunk, unit, 0, + {'rows_in': 0, 'rows_out': 0}) + + self.assertEqual(len(frame), 0) + + +if __name__ == '__main__': + unittest.main() diff --git a/src/immvoke/Sources/Specificity/tests/test_Distances.py b/src/immvoke/Sources/Specificity/tests/test_Distances.py new file mode 100644 index 0000000..bffdf83 --- /dev/null +++ b/src/immvoke/Sources/Specificity/tests/test_Distances.py @@ -0,0 +1,133 @@ +""" +Unit tests for every gated pair's distance (Distances.py) +""" + +# Info +__author__ = 'Pramod Shinde' + +# Imports +import unittest +from unittest import mock + +import pandas + +# Immvoke imports +from immvoke.Sources.Specificity import Distances +from immvoke.Sources.Specificity.Distances import allPairDistances +from immvoke.Sources.Specificity.Symscan import runSymscanAnnotation +from immvoke.Sources.Specificity.tests.BruteForce import bruteForceAllPairs +from immvoke.Sources.Specificity.tests.test_Symscan import randomTables + + +def frame(rows): + return pandas.DataFrame(rows) + + +def allPairs(*args, **kwargs): + """allPairDistances' tables, as one.""" + return pandas.concat(list(allPairDistances(*args, **kwargs)), ignore_index=True) + + +class TestAllPairDistances(unittest.TestCase): + """ + Tests for scoring every gated pair, whatever the threshold + """ + + def test_every_pair_is_scored_beyond_the_threshold(self): + queries = frame([{'cdr3': 'CARDYW'}]) + reference = frame([{'id': 'r1', 'cdr3': 'CARDYW'}, {'id': 'r2', 'cdr3': 'CQQQYW'}, + {'id': 'r3', 'cdr3': 'CARD'}]) + + pairs = allPairs(queries, reference, ['cdr3'], 1, distance_type='levenshtein', + id_col='id') + + self.assertEqual(pairs[['reference_id', 'mismatches', 'within_threshold']] + .values.tolist(), + [['r1', 0, True], ['r3', 2, False], ['r2', 3, False]]) + self.assertEqual(pairs['identity'].round(4).tolist(), [1.0, 0.6667, 0.5]) + + def test_hamming_scores_same_length_pairs_only(self): + queries = frame([{'cdr3': 'CARDYW'}]) + reference = frame([{'id': 'r1', 'cdr3': 'CQQQYW'}, {'id': 'r2', 'cdr3': 'CARD'}, + {'id': 'r3', 'cdr3': 'XARDYF'}]) + + pairs = allPairs(queries, reference, ['cdr3'], 0, id_col='id') + + # An X never counts under Hamming; CARD is another length. + self.assertEqual(pairs[['reference_id', 'mismatches', 'mismatch_positions']] + .values.tolist(), [['r3', 1, '5'], ['r1', 3, '1,2,3']]) + + def test_gates_still_apply(self): + queries = frame([{'cdr3': 'CARDYW', 'v_call': 'IGHV1-2*02'}, + {'cdr3': '', 'v_call': 'IGHV1-2*02'}]) + reference = frame([{'id': 'r1', 'cdr3': 'CQQQQQ', 'v_call': 'IGHV1-2*04'}, + {'id': 'r2', 'cdr3': 'CARDYW', 'v_call': 'IGHV3-23*01'}, + {'id': 'r3', 'cdr3': 'CARDYW', 'v_call': ''}]) + + pairs = allPairs(queries, reference, ['cdr3'], 0, id_col='id', + gene_cols=['v_call']) + + # Only the same gene, and nothing for the blank query. + self.assertEqual(pairs[['query_row', 'reference_id', 'mismatches']] + .values.tolist(), [[0, 'r1', 5]]) + + def test_no_pair_yields_one_empty_table(self): + tables = list(allPairDistances(frame([{'cdr3': 'CARDYW'}]), + frame([{'id': 'r1', 'cdr3': 'CA'}]), ['cdr3'], 0, + id_col='id')) + + self.assertEqual(len(tables), 1) + self.assertEqual(len(tables[0]), 0) + self.assertEqual(tables[0].columns.tolist(), + ['query_row', 'reference_id', 'mismatches', 'identity', + 'within_threshold', 'mismatch_positions']) + + +class TestAllPairDistancesMatchBruteForce(unittest.TestCase): + """ + allPairDistances must return exactly the pairs a brute-force comparison + of every pair returns, and its within_threshold pairs exactly symscan's + -- on the seeded random tables test_Symscan checks symscan with. + """ + + CASES = [(['cdr3'], 1, {}), + (['cdr3'], 2, {'exact_cols': ['j_call']}), + (['cdr1', 'cdr3'], 2, {'gene_cols': ['v_call']}), + (['cdr3'], 0.34, {'threshold_is_fraction': True, + 'gene_cols': ['v_call', 'j_call']}), + (['cdr3'], 1, {'distance_type': 'levenshtein'}), + (['cdr3'], 2, {'distance_type': 'levenshtein', + 'gene_cols': ['v_call', 'j_call']})] + + def _check(self, queries, reference, cols, tolerance, **kwargs): + pairs = allPairs(queries, reference, cols, tolerance, id_col='id', **kwargs) + expected = bruteForceAllPairs(queries, reference, cols, tolerance, id_col='id', + **kwargs) + pandas.testing.assert_frame_equal(pairs, expected, check_dtype=False) + + _, hits = runSymscanAnnotation(queries, reference, cols, tolerance, id_col='id', + return_pairs=True, **kwargs) + within = pairs[pairs['within_threshold']].drop(columns='within_threshold') + pandas.testing.assert_frame_equal(within.reset_index(drop=True), hits, + check_dtype=False) + + def test_random_tables(self): + for seed in range(5): + queries, reference = randomTables(seed) + for cols, tolerance, kwargs in self.CASES: + with self.subTest(seed=seed, cols=cols, tolerance=tolerance, **kwargs): + self._check(queries, reference, cols, tolerance, **kwargs) + + def test_small_blocks_give_the_same_pairs(self): + """Query blocks and vectorized steps small enough to split every + group many times over.""" + queries, reference = randomTables(0) + with mock.patch.object(Distances, 'QUERY_BLOCK', 7), \ + mock.patch.object(Distances, 'RESIDUE_BUDGET', 50): + for cols, tolerance, kwargs in self.CASES: + with self.subTest(cols=cols, tolerance=tolerance, **kwargs): + self._check(queries, reference, cols, tolerance, **kwargs) + + +if __name__ == '__main__': + unittest.main() diff --git a/src/immvoke/Sources/Specificity/tests/test_Genes.py b/src/immvoke/Sources/Specificity/tests/test_Genes.py new file mode 100644 index 0000000..cfdd7bd --- /dev/null +++ b/src/immvoke/Sources/Specificity/tests/test_Genes.py @@ -0,0 +1,105 @@ +""" +Unit tests for V/J-gene call parsing and comparison (Genes.py) +""" + +# Info +__author__ = 'Pramod Shinde' + +# Imports +import unittest + +import numpy +import pandas + +# Immvoke imports +from immvoke.Sources.Specificity.Genes import ( + GeneGate, + callsAgree, + parseCall, +) + + +def agree(a, b): + return callsAgree(parseCall(a), parseCall(b)) + + +class TestParseCall(unittest.TestCase): + """ + Tests for reading the gene names out of one gene call cell + """ + + def test_the_allele_is_dropped(self): + self.assertEqual(parseCall('IGHV1-2*02'), {'IGHV1-2'}) + self.assertEqual(parseCall('IGHV1-2'), {'IGHV1-2'}) + + def test_several_calls_in_one_cell(self): + self.assertEqual(parseCall('IGHV5-9*01, IGHV5-9*03, IGHV5-9-5*01'), + {'IGHV5-9', 'IGHV5-9-5'}) + + def test_a_gene_level_name_is_not_read_as_an_allele(self): + """Not gene 'IGHV1' allele '2', as Change-O's allele pattern reads it.""" + self.assertEqual(parseCall('IGHV1-2'), {'IGHV1-2'}) + + def test_names_change_o_parses(self): + for call, gene in [('IGHV3/OR16-10*02', 'IGHV3/OR16-10'), + ('IGHV1-69D*01', 'IGHV1-69D'), + ('IGKV1D-39*01', 'IGKV1D-39'), + ('IGHV3-30-3*01', 'IGHV3-30-3'), + ('TRBV20-1*01', 'TRBV20-1')]: + with self.subTest(call): + self.assertEqual(parseCall(call), {gene}) + + def test_blank_or_unrecognized_names_nothing(self): + for call in ('', None, 'IGVL3-10', 'unknown'): + with self.subTest(call): + self.assertEqual(parseCall(call), frozenset()) + + +class TestCallsAgree(unittest.TestCase): + """ + Tests for whether two cells name the same gene, any allele + """ + + def test_same_gene_any_allele(self): + self.assertTrue(agree('IGHV1-2*02', 'IGHV1-2*04')) + self.assertTrue(agree('IGHV1-2*02', 'IGHV1-2')) + self.assertFalse(agree('IGHV1-2*02', 'IGHV1-69*01')) + self.assertFalse(agree('IGHV1-69*01', 'IGHV1-69D*01')) + + def test_any_shared_call_agrees(self): + self.assertTrue(agree('IGHV3-23*01, IGHV1-2*02', 'IGHV1-2*04')) + self.assertFalse(agree('IGHV3-23*01, IGHV1-2*02', 'IGHV1-69*01')) + + def test_blank_or_unrecognized_agrees_with_nothing(self): + for other in ('IGHV1-2*02', '', 'IGVL3-10'): + with self.subTest(other=other): + self.assertFalse(agree('', other)) + self.assertFalse(agree('IGVL3-10', other)) + + +class TestGeneGate(unittest.TestCase): + """ + Tests for the pairwise gate over several gene columns + """ + + def test_every_column_must_agree(self): + queries = pandas.DataFrame([{'v_call': 'IGHV1-2*01', 'j_call': 'IGHJ4*02'}]) + reference = pandas.DataFrame([ + {'v_call': 'IGHV1-2*02', 'j_call': 'IGHJ4*02'}, # both agree + {'v_call': 'IGHV1-2*02', 'j_call': 'IGHJ6*02'}, # J differs + {'v_call': 'IGHV1-2', 'j_call': ''}, # J blank + ]) + + gate = GeneGate(queries, reference, ['v_call', 'j_call']) + + keep = gate.keep(numpy.array([0, 0, 0]), numpy.array([0, 1, 2])) + self.assertEqual(keep.tolist(), [True, False, False]) + + def test_unrecognized_values_are_logged(self): + frame = pandas.DataFrame([{'v_call': 'IGVL3-10'}]) + with self.assertLogs('immvoke.Sources.Specificity.Genes', 'WARNING'): + GeneGate(frame, frame, ['v_call']) + + +if __name__ == '__main__': + unittest.main() diff --git a/src/immvoke/Sources/Specificity/tests/test_Report.py b/src/immvoke/Sources/Specificity/tests/test_Report.py new file mode 100644 index 0000000..9b89cbc --- /dev/null +++ b/src/immvoke/Sources/Specificity/tests/test_Report.py @@ -0,0 +1,120 @@ +""" +Unit tests for what `immvoke specificity download` reports as it runs +""" + +# Info +__author__ = 'Pramod Shinde' + +# Imports +import logging +import unittest +from pathlib import Path + +import pandas + +# Immvoke imports +from immvoke.Commandline import StatusFormatter +from immvoke.Sources.Specificity import Report +from immvoke.Sources.Specificity.Iedb import IedbSource + + +class TestReport(unittest.TestCase): + """ + Tests for the download's status lines + """ + + def test_durations_read_well(self): + self.assertEqual(Report.humanDuration(0.42), '0.4 s') + self.assertEqual(Report.humanDuration(8.4), '8 s') + self.assertEqual(Report.humanDuration(125), '2 min 5 s') + self.assertEqual(Report.humanDuration(3 * 3600 + 610), '3 h 10 min') + + def test_tally_counts_distinct_values_as_chunks_pass(self): + tally = Report.Tally(IedbSource, 'bcr') + chunks = [pandas.DataFrame({'iedb_receptor_group': ['g1', 'g1'], + 'reactivity_ref': ['e1', '']}), + pandas.DataFrame({'iedb_receptor_group': ['g2'], + 'reactivity_ref': ['e1']})] + + passed = list(tally.wrap(chunks)) + + self.assertEqual(len(passed), 2) + self.assertEqual(tally.distinct(), [('receptor groups', 2), ('epitopes', 1)]) + + def test_completed_says_what_was_written(self): + tally = Report.Tally(IedbSource, 'bcr') + list(tally.wrap([pandas.DataFrame({'iedb_receptor_group': ['g1'], + 'reactivity_ref': ['e1']})])) + validation = {'header_valid': True, 'rows_checked': 1, 'rows_invalid': 0} + + with self.assertLogs(Report.log, 'INFO') as logs: + Report.completed(Path('databases/iedb/iedb_bcr.tsv'), 15.2, 3, 1, tally, + validation) + text = '\n'.join(record.getMessage() for record in logs.records) + + self.assertIn('Download completed in 15 s: databases/iedb/iedb_bcr.tsv', text) + self.assertIn('rows 1 of 3 kept', text) + self.assertIn('every row valid (iedb_bcr.airr_validation.log)', text) + + def test_an_annotation_says_what_it_compares_and_what_it_found(self): + """The run's own options, in words, rather than the raw values.""" + query = {'path': 'mine.tsv', 'rows': 43225} + reference = {'path': 'databases/iedb/iedb_bcr.tsv', 'rows': 18650} + options = {'sequence_region': 'CDR3H', 'partition_by': 'VJ', + 'paired': False, 'identity_threshold': 0.7, + 'distance_metric': 'hamming'} + + with self.assertLogs(Report.log, 'INFO') as logs: + Report.annotationStarted(query, reference, options) + Report.annotationCompleted( + 56.2, {'query_rows_compared': 21412, 'query_rows_with_hits': 17, + 'hits': 22}, paired=False) + text = '\n'.join(record.getMessage() for record in logs.records) + + self.assertIn('Annotation started: 43,225 rows of mine.tsv', text) + self.assertIn('against 18,650 rows of databases/iedb/iedb_bcr.tsv', text) + self.assertIn('comparing: the heavy chain CDR3 (cdr3_aa)', text) + self.assertIn('a hit is: 70% identity or better, hamming', text) + self.assertIn('only rows with the same V and J gene are compared', text) + self.assertIn('Annotation completed in 56 s', text) + self.assertIn('query rows matched 17 of 21,412', text) + + def test_an_exact_paired_annotation_says_so(self): + options = {'sequence_region': 'CDR3H', 'partition_by': None, + 'paired': False, 'identity_threshold': None, + 'distance_metric': 'hamming'} + with self.assertLogs(Report.log, 'INFO') as logs: + Report.annotationStarted({'path': None, 'rows': 2}, + {'path': None, 'rows': 3}, options) + text = '\n'.join(x.getMessage() for x in logs.records) + + self.assertIn('a table passed from Python', text) + self.assertIn('a hit is: an exact match', text) + self.assertIn('any pair of rows is compared', text) + + with self.assertLogs(Report.log, 'INFO') as logs: + Report.annotationStarted({'path': None, 'rows': 2}, + {'path': None, 'rows': 3}, + dict(options, paired=True)) + Report.annotationCompleted( + 1.5, {'query_antibodies': 2, 'classification': {'both': 1, 'neither': 1}}, + paired=True) + text = '\n'.join(x.getMessage() for x in logs.records) + + self.assertIn('heavy and light chain of the same antibody', text) + self.assertIn('query antibodies', text) + self.assertIn('both', text) + + def test_status_lines_carry_no_level_but_warnings_do(self): + formatter = StatusFormatter('%(message)s') + + def line(level): + return formatter.format(logging.LogRecord('immvoke', level, __file__, 1, + 'hello', None, None)) + + self.assertEqual(line(logging.INFO), 'hello') + self.assertEqual(line(logging.WARNING), 'WARNING hello') + + +if __name__ == '__main__': + unittest.main() diff --git a/src/immvoke/Sources/Specificity/tests/test_Search.py b/src/immvoke/Sources/Specificity/tests/test_Search.py new file mode 100644 index 0000000..3a0c003 --- /dev/null +++ b/src/immvoke/Sources/Specificity/tests/test_Search.py @@ -0,0 +1,211 @@ +""" +Unit tests for `immvoke specificity search` and the filter flags' help +""" + +# Info +__author__ = 'Pramod Shinde' + +# Imports +import json +import os +import tempfile +import time +import unittest +from argparse import ArgumentParser +from pathlib import Path +from unittest import mock + +# Immvoke imports +from immvoke.Schema import loadSchema +from immvoke.Sources.Base import DataUnit +from immvoke.Sources.Specificity.Iedb import BULK_URL, CACHE_MAX_AGE, IedbSource +from immvoke.Sources.Specificity.Search import ( + SpecificityHelpFormatter, + _flow, + addFilterGroup, + formatSummary, + searchTable, + searchTables, +) +from immvoke.Sources.Specificity.tests.test_Specificity import ( + iedbAirrRows, + isolateCache, + makeClient, + makeZip, +) +from tests.FakeHttp import FakeResponse, rangeHandler + + +def iedbZip(): + """A bulk ZIP holding a small bcr table: a human and a mouse antibody.""" + bcr = iedbAirrRows( + {'cell_id': 'h', 'sequence_id': 'h_1', 'locus': 'IGH', + 'locus_species': 'NCBITAXON:9606', 'iedb_receptor_group': 'IEDB_RECEPTOR:1', + 'reactivity_ref': 'IEDB_EPITOPE:7', 'iedb_assay_outcome': 'positive'}, + {'cell_id': 'h', 'sequence_id': 'h_2', 'locus': 'IGK', + 'locus_species': 'NCBITAXON:9606', 'iedb_receptor_group': 'IEDB_RECEPTOR:1', + 'reactivity_ref': 'IEDB_EPITOPE:7', 'iedb_assay_outcome': 'positive'}, + {'cell_id': 'm', 'sequence_id': 'm_1', 'locus': 'IGH', + 'locus_species': 'NCBITAXON:10000000', 'iedb_receptor_group': 'IEDB_RECEPTOR:2', + 'reactivity_ref': 'IEDB_EPITOPE:8', 'iedb_assay_outcome': 'negative'}) + + return makeZip({'bcr_full_airr.tsv': bcr.to_csv(sep='\t', index=False), + 'tcr_full_airr.tsv': 'cell_id\n'}) + + +#: epitope_search's answer for iedbZip's two epitopes. +EPITOPES = [ + {'structure_id': 7, 'source_organism_iri_search': ['NCBITaxon:11320'], + 'structure_type': 'Linear peptide', + 'curated_source_antigens': [{'name': 'hemagglutinin', + 'source_organism_name': 'Influenza A virus'}]}, + {'structure_id': 8, 'source_organism_iri_search': ['NCBITaxon:9606'], + 'structure_type': 'Non-peptidic', + 'curated_source_antigens': [{'name': 'Angiogenin', + 'source_organism_name': 'Homo sapiens (human)'}]}, +] + + +def iedbHandler(zip_bytes): + """Serve the bulk ZIP, and EPITOPES to any epitope_search request.""" + serve_zip = rangeHandler(zip_bytes) + + def handler(method, url, headers, index): + if 'epitope_search' in url: + return FakeResponse(200, json.dumps(EPITOPES).encode('utf-8')) + return serve_zip(method, url, headers, index) + + return handler + + +class CountingHandler: + """Serve a body like rangeHandler, counting the requests made.""" + + def __init__(self, body): + self.serve = rangeHandler(body) + self.calls = 0 + + def __call__(self, *args): + self.calls += 1 + return self.serve(*args) + + +class TestZipCache(unittest.TestCase): + """ + Tests for the bulk ZIP cache download and search share + """ + + def setUp(self): + self.cache = isolateCache(self) + + def test_the_zip_is_reused_then_refetched_once_stale(self): + handler = CountingHandler(iedbZip()) + source = IedbSource(makeClient(handler)) + unit = DataUnit(unit_id='bcr_full_airr.tsv', collection='bcr', url=BULK_URL) + + with tempfile.TemporaryDirectory() as outdir: + source.fetchUnit(unit, Path(outdir), progress=False) + fetched = handler.calls + source.fetchUnit(unit, Path(outdir), progress=False) + self.assertEqual(handler.calls, fetched) + + stale = time.time() - CACHE_MAX_AGE - 60 + os.utime(self.cache / 'airr_full.zip', (stale, stale)) + source.fetchUnit(unit, Path(outdir), progress=False) + self.assertGreater(handler.calls, fetched) + + +class TestFilterHelp(unittest.TestCase): + """ + Tests for the filter flags' help + """ + + def test_flow_never_splits_a_value(self): + values = ['Mus musculus', 'Mus musculus BALB/c', 'Homo sapiens', 'Rattus'] + lines = _flow(values, 30) + + self.assertTrue(all(len(line) <= 30 for line in lines)) + self.assertEqual(' '.join(x.strip() for x in lines), + 'Mus musculus, Mus musculus BALB/c, Homo sapiens, Rattus') + + def test_filters_are_grouped_and_described(self): + parser = ArgumentParser(formatter_class=SpecificityHelpFormatter) + addFilterGroup(parser, loadSchema('iedb'), IedbSource, 'bcr') + with mock.patch.dict(os.environ, {'COLUMNS': '100'}): + text = parser.format_help() + + self.assertIn('filters:', text) + self.assertIn('--air_species SPECIES', text) + self.assertIn('--antigen_species ORGANISM', text) + self.assertIn('--antigen_structure_type TYPE', text) + self.assertIn('--assay_type ASSAY', text) + self.assertIn('--disease DISEASE', text) + self.assertIn('--paired_only ', text) + self.assertIn('one of: Positive, Negative', text) + self.assertIn('Mus musculus BALB/c', text) + self.assertNotIn('default: None', text) + args = parser.parse_args(['--air_species', 'Homo sapiens', '--paired_only']) + self.assertEqual(args.filter_air_species, 'Homo sapiens') + self.assertIs(args.filter_paired_only, True) + self.assertIsNone(parser.parse_args([]).filter_paired_only) + + +class TestSearch(unittest.TestCase): + """ + Tests for counting what a table holds + """ + + def setUp(self): + isolateCache(self) + self.source = IedbSource(makeClient(iedbHandler(iedbZip()))) + + def test_counts_rows_distinct_values_and_top_values(self): + with tempfile.TemporaryDirectory() as workdir: + summary = searchTable(self.source, 'bcr', {}, Path(workdir)) + + self.assertEqual((summary['rows_in'], summary['rows']), (3, 3)) + self.assertEqual(summary['distinct'], {'receptor groups': 2, 'epitopes': 2}) + self.assertEqual(dict(summary['top']['locus']), {'IGH': 2, 'IGK': 1}) + # taxon codes are named, strains by their own name + self.assertEqual(dict(summary['top']['receptor species']), + {'Homo sapiens': 2, 'Mus musculus BALB/c': 1}) + self.assertEqual(dict(summary['top']['antigen species']), + {'Influenza A virus': 2, 'Homo sapiens (human)': 1}) + self.assertEqual(dict(summary['top']['antigen structure type']), + {'Linear peptide': 2, 'Non-peptidic': 1}) + + def test_filters_narrow_the_counts(self): + with tempfile.TemporaryDirectory() as workdir: + summary = searchTable(self.source, 'bcr', {'air_species': 'Mus musculus'}, + Path(workdir)) + + self.assertEqual(summary['filters'], {'air_species': 'Mus musculus'}) + self.assertEqual((summary['rows_in'], summary['rows']), (3, 1)) + self.assertEqual(dict(summary['top']['outcome']), {'negative': 1}) + + def test_summary_reads_well(self): + with tempfile.TemporaryDirectory() as workdir: + summary = searchTable(self.source, 'bcr', {'outcome': 'Positive'}, + Path(workdir)) + text = formatSummary(self.source, summary, top=1) + + self.assertIn('filters: outcome = Positive', text) + self.assertIn('rows 2 of 3', text) + self.assertIn('locus (top 1 of 2)', text) + + def test_search_writes_nothing_but_the_counts_asked_for(self): + with tempfile.TemporaryDirectory() as tmp: + out = Path(tmp) / 'counts.tsv' + with mock.patch('builtins.print'): + searchTables(self.source, ['bcr'], {}, out=out) + + self.assertEqual(sorted(p.name for p in Path(tmp).iterdir()), ['counts.tsv']) + lines = out.read_text().splitlines() + + self.assertEqual(lines[0], 'table\tfield\tvalue\trows') + self.assertIn('bcr\trows\t\t3', lines) + self.assertIn('bcr\treceptor species\tHomo sapiens\t2', lines) + + +if __name__ == '__main__': + unittest.main() diff --git a/src/immvoke/Sources/Specificity/tests/test_Specificity.py b/src/immvoke/Sources/Specificity/tests/test_Specificity.py new file mode 100644 index 0000000..71f6905 --- /dev/null +++ b/src/immvoke/Sources/Specificity/tests/test_Specificity.py @@ -0,0 +1,734 @@ +""" +Unit tests for specificity sources (IEDB) and the Range-header paginator +""" + +# Info +__author__ = 'Pramod Shinde' + +# Imports +import json +import os +import tempfile +import unittest +import zipfile +from io import BytesIO +from pathlib import Path +from unittest import mock + +import pandas + +# Immvoke imports +from immvoke.Exceptions import HttpError, IedbParseError, SchemaError +from immvoke.Http import HttpClient +from immvoke.Sources.Base import DataUnit, Query +from immvoke.Sources.Specificity import REGISTRY, getSpecificitySource +from immvoke.Sources.Specificity.Iedb import BULK_URL, IedbSource +from immvoke.Sources.Specificity.Paginate import pageByRange +from tests.FakeHttp import FakeResponse, FakeSession, rangeHandler + +#: The columns of IEDB's AIRR export, bcr_full_airr.tsv, in file order. +IEDB_AIRR_COLUMNS = ( + 'cell_id', 'sequence_id', 'iedb_receptor_group', 'reactivity_ref', + 'iedb_epitope_name', 'iedb_assay_outcome', 'iedb_assay_ids', + 'iedb_mhc_alleles', 'productive', 'locus', 'locus_species', 'v_call', + 'd_call', 'j_call', 'sequence', 'sequence_aa', 'cdr1_aa', 'cdr2_aa', + 'cdr3_aa', 'junction_aa', 'rev_comp', 'sequence_alignment', + 'germline_alignment', 'junction', 'cdr3', 'v_cigar', 'd_cigar', 'j_cigar', +) + + +def iedbAirrRows(*rows): + """ + Rows shaped like IEDB's AIRR export: every column of IEDB_AIRR_COLUMNS, + blank unless a row gives it. + """ + return pandas.DataFrame([{c: row.get(c, '') for c in IEDB_AIRR_COLUMNS} + for row in rows]) + + +def makeClient(handler): + """Build an HttpClient with no politeness delay over a scripted session.""" + return HttpClient(delay=0, backoff=0, session=FakeSession(handler)) + + +def itemRangeHandler(pages): + """ + Build a handler serving JSON pages via item Range headers, PostgREST style. + + Arguments: + pages (list): one list of records per page, in the order they are served. + A request past the last page gets a 416, the way IEDB's API ends + pagination. + + Returns: + callable: a handler for FakeSession. + """ + total = sum(len(page) for page in pages) + + def handler(method, url, headers, index): + if index >= len(pages): + return FakeResponse(416, b'') + + page = pages[index] + start = sum(len(p) for p in pages[:index]) + end = start + len(page) - 1 + body = json.dumps(page).encode('utf-8') + + return FakeResponse(206, body, { + 'Content-Range': 'items %d-%d/%d' % (start, end, total)}) + + return handler + + +class TestPageByRange(unittest.TestCase): + """ + Tests for the PostgREST Range-header paginator + """ + + def test_pages_until_content_range_total_is_reached(self): + client = makeClient(itemRangeHandler([[{'id': 1}, {'id': 2}], [{'id': 3}]])) + records = [r for batch in pageByRange(client, 'https://x/y', page_size=2) + for r in batch] + + self.assertEqual(records, [{'id': 1}, {'id': 2}, {'id': 3}]) + + def test_a_short_final_page_also_ends_pagination(self): + """A page shorter than page_size ends iteration even with no total.""" + def handler(method, url, headers, index): + if index == 0: + return FakeResponse(206, json.dumps([{'id': 1}]).encode()) + raise AssertionError('paginator should have stopped after one page') + + client = makeClient(handler) + records = [r for batch in pageByRange(client, 'https://x/y', page_size=5) + for r in batch] + + self.assertEqual(records, [{'id': 1}]) + + def test_416_past_the_last_row_ends_pagination(self): + client = makeClient(itemRangeHandler([[{'id': 1}]])) + records = [r for batch in pageByRange(client, 'https://x/y', page_size=1) + for r in batch] + + self.assertEqual(records, [{'id': 1}]) + + def test_unexpected_status_raises(self): + client = makeClient(lambda method, url, headers, index: FakeResponse(404)) + + with self.assertRaises(HttpError): + list(pageByRange(client, 'https://x/y')) + + +def makeZip(members): + """ + Build ZIP bytes containing the given filename to text-content mapping. + """ + buffer = BytesIO() + with zipfile.ZipFile(buffer, 'w') as archive: + for name, content in members.items(): + archive.writestr(name, content) + + return buffer.getvalue() + + +def isolateCache(test): + """ + Point IEDB's ZIP cache (see Iedb.cacheDir) at a temporary folder for the + length of a test, so no test reads or writes the user's own cache. + """ + cache = tempfile.TemporaryDirectory() + test.addCleanup(cache.cleanup) + patcher = mock.patch.dict(os.environ, {'XDG_CACHE_HOME': cache.name}) + patcher.start() + test.addCleanup(patcher.stop) + + return Path(cache.name) / 'immvoke' / 'iedb' + + +class TestIedbSource(unittest.TestCase): + """ + Tests for the IEDB specificity source + """ + + def setUp(self): + self.cache = isolateCache(self) + + def test_registered(self): + self.assertIs(REGISTRY['iedb'], IedbSource) + client = makeClient(lambda *a: FakeResponse(200)) + self.assertIsInstance(getSpecificitySource('iedb', client), IedbSource) + + def test_unknown_source_raises(self): + client = makeClient(lambda *a: FakeResponse(200)) + with self.assertRaises(KeyError): + getSpecificitySource('not-a-real-db', client) + + def test_search_units_returns_one_unit_per_table(self): + source = IedbSource(makeClient(lambda *a: FakeResponse(200))) + + for table in source.collections: + units = source.searchUnits(Query(collection=table, filters={})) + self.assertEqual(len(units), 1) + self.assertEqual(units[0].collection, table) + + bcr = source.searchUnits( + Query(collection='bcr', filters={'outcome': 'Positive'}))[0] + self.assertEqual(bcr.url, BULK_URL) + self.assertEqual(bcr.metadata, {'outcome': 'Positive'}) + + def test_only_the_receptor_tables_are_offered(self): + """IEDB's B-cell assay tables are not offered: each receptor record + already carries its epitope, assay ids and outcome.""" + self.assertEqual(IedbSource.collections, ('bcr', 'tcr')) + + def test_fetch_unit_extracts_the_right_member_from_the_shared_zip(self): + """bcr and tcr are both members of one AIRR ZIP, extracted by name.""" + zip_bytes = makeZip({ + 'tcr_full_airr.tsv': 'sequence_id\tlocus\nt1\tTRB\n', + 'bcr_full_airr.tsv': 'sequence_id\tlocus\nb1\tIGH\n', + }) + client = makeClient(rangeHandler(zip_bytes)) + source = IedbSource(client) + unit = DataUnit(unit_id='bcr_full_airr.tsv', collection='bcr', url=BULK_URL) + + with tempfile.TemporaryDirectory() as outdir: + result = source.fetchUnit(unit, Path(outdir)) + self.assertEqual(result.path.read_text(), 'sequence_id\tlocus\nb1\tIGH\n') + # The ZIP is kept in the cache, not beside the extracted table, so + # a second table or a later search does not re-download it. + self.assertTrue((self.cache / 'airr_full.zip').exists()) + self.assertFalse((Path(outdir) / 'airr_full.zip').exists()) + + def test_fetch_unit_raises_when_the_zip_member_is_missing(self): + zip_bytes = makeZip({'tcr_full_airr.tsv': 'sequence_id\nt1\n'}) + client = makeClient(rangeHandler(zip_bytes)) + source = IedbSource(client) + unit = DataUnit(unit_id='bcr_full_airr.tsv', collection='bcr', url=BULK_URL) + + with tempfile.TemporaryDirectory() as outdir: + with self.assertRaises(IedbParseError): + source.fetchUnit(unit, Path(outdir)) + + def test_read_unit_bulk_reads_the_airr_tsv_as_text(self): + """Blank cells stay '' rather than NaN, and nothing turns numeric.""" + source = IedbSource(makeClient(lambda *a: FakeResponse(200))) + unit = DataUnit(unit_id='bcr_full_airr.tsv', collection='bcr', url=BULK_URL) + frame_in = iedbAirrRows({'cell_id': '1', 'sequence_id': '1_1', 'locus': 'IGH', + 'locus_species': 'NCBITAXON:9606', + 'cdr3_aa': 'ARDYW', 'iedb_assay_ids': '32'}) + + with tempfile.TemporaryDirectory() as outdir: + path = Path(outdir) / 'bcr_full_airr.tsv' + frame_in.to_csv(path, sep='\t', index=False) + metadata, chunks = source.readUnit(path, unit) + frame = pandas.concat(list(chunks), ignore_index=True) + + self.assertEqual(metadata, {'species_only': False, 'kept_cells': None, + 'kept_epitopes': None, 'kept_assays': None, + 'antigens': {}}) + self.assertEqual(list(frame.columns), list(IEDB_AIRR_COLUMNS)) + self.assertEqual(frame['cdr3_aa'].tolist(), ['ARDYW']) + self.assertEqual(frame['iedb_assay_ids'].tolist(), ['32']) + self.assertEqual(frame['j_call'].tolist(), ['']) + + +class TestIedbAirrExport(unittest.TestCase): + """ + Tests for the receptor tables, read from IEDB's own AIRR export + """ + + def setUp(self): + self.cache = isolateCache(self) + + def test_bcr_and_tcr_are_airr_collections(self): + self.assertEqual(IedbSource.airr_collections, frozenset(['bcr', 'tcr'])) + self.assertEqual(IedbSource.airr_collections & set(IedbSource.collections), + frozenset(['bcr', 'tcr'])) + + def test_bcr_comes_from_the_airr_export(self): + source = IedbSource(makeClient(lambda *a: FakeResponse(200))) + unit = source.searchUnits(Query(collection='bcr', filters={}))[0] + + self.assertEqual(unit.unit_id, 'bcr_full_airr.tsv') + self.assertTrue(unit.url.endswith('airr_full.zip')) + + def _normalize(self, frame, air_species=None, outcome=None, + antigen_species=None, antigen_structure_type=None, handler=None, + **more): + """Read and normalize `frame` the way `immvoke specificity iedb + download bcr` does, from a downloaded TSV, with --air_species, + --outcome, --antigen_species, --antigen_structure_type and any + `more` filters if given; `handler` serves IEDB's API.""" + source = IedbSource(makeClient(handler or (lambda *a: FakeResponse(200)))) + filters = {k: v for k, v in (('air_species', air_species), + ('outcome', outcome), + ('antigen_species', antigen_species), + ('antigen_structure_type', antigen_structure_type), + *more.items()) + if v is not None} + unit = DataUnit(unit_id='bcr_full_airr.tsv', collection='bcr', url=BULK_URL, + metadata=filters) + report = {'rows_in': 0, 'rows_out': 0} + + with tempfile.TemporaryDirectory() as outdir: + path = Path(outdir) / 'bcr_full_airr.tsv' + frame.to_csv(path, sep='\t', index=False) + metadata, chunks = source.readUnit(path, unit, chunksize=1) + normalized = [source.normalizeChunk(metadata, chunk, unit, 0, report) + for chunk in chunks] + + return pandas.concat(normalized, ignore_index=True), report + + def _antibody(self): + """One antibody, both chains, as IEDB's AIRR export writes it.""" + return iedbAirrRows( + {'cell_id': '1284_2134780p', 'sequence_id': '1284_1', 'locus': 'IGH', + 'locus_species': 'NCBITAXON:9606', + 'productive': 'T', 'v_call': 'IGHV5-9*01, IGHV5-9*03', + 'j_call': 'IGHJ4*02', 'cdr3_aa': 'TRLGDYGYAYTMDY', 'iedb_epitope_name': 'G58, L59', + 'iedb_assay_outcome': 'positive'}, + {'cell_id': '1284_2134780p', 'sequence_id': '1284_2', 'locus': 'IGK', + 'productive': 'T', 'v_call': 'IGKV3-2*01', 'j_call': 'IGKJ2*01', + 'cdr3_aa': 'QQSKEVPLT', + 'iedb_epitope_name': 'G58, L59', 'iedb_assay_outcome': 'positive'}) + + def _mixedSpecies(self): + """Receptors of several species, one of a BALB/c mouse.""" + human, mouse, construct = 'NCBITAXON:9606', 'NCBITAXON:10090', 'NCBITAXON:32630' + return iedbAirrRows( + {'cell_id': 'human', 'sequence_id': '1_1', 'locus_species': human}, + {'cell_id': 'human', 'sequence_id': '1_2', 'locus_species': human}, + {'cell_id': 'humanBlank', 'sequence_id': '2_1', 'locus_species': human}, + {'cell_id': 'humanBlank', 'sequence_id': '2_2', 'locus_species': ''}, + {'cell_id': 'mouse', 'sequence_id': '3_1', 'locus_species': mouse}, + {'cell_id': 'mixed', 'sequence_id': '4_1', 'locus_species': human}, + {'cell_id': 'mixed', 'sequence_id': '4_2', 'locus_species': construct}, + {'cell_id': 'blank', 'sequence_id': '5_1', 'locus_species': ''}, + {'cell_id': 'balbc', 'sequence_id': '6_1', + 'locus_species': 'NCBITAXON:10000000'}) + + def test_every_species_is_kept_without_air_species(self): + kept, report = self._normalize(self._mixedSpecies()) + + self.assertEqual(len(kept), 9) + self.assertEqual((report['rows_in'], report['rows_out']), (9, 9)) + + def test_air_species_keeps_only_that_species(self): + """Only chains of the asked species: a chain with no species + recorded is dropped, leaving its partner alone; any chain of another + species drops the whole receptor.""" + kept, report = self._normalize(self._mixedSpecies(), 'Homo sapiens') + + self.assertEqual(kept['sequence_id'].tolist(), ['1_1', '1_2', '2_1']) + self.assertEqual((report['rows_in'], report['rows_out']), (9, 3)) + + def test_air_species_includes_its_strains(self): + kept, _ = self._normalize(self._mixedSpecies(), 'Mus musculus') + self.assertEqual(kept['sequence_id'].tolist(), ['3_1', '6_1']) + + kept, _ = self._normalize(self._mixedSpecies(), 'Mus musculus BALB/c') + self.assertEqual(kept['sequence_id'].tolist(), ['6_1']) + + def test_outcome_keeps_only_that_outcome(self): + """IEDB's AIRR export writes the outcome in lower case; --outcome + matches it either way, and combines with --air_species.""" + frame = iedbAirrRows( + {'cell_id': 'a', 'sequence_id': '1_1', 'iedb_assay_outcome': 'positive', + 'locus_species': 'NCBITAXON:9606'}, + {'cell_id': 'b', 'sequence_id': '2_1', 'iedb_assay_outcome': 'negative', + 'locus_species': 'NCBITAXON:9606'}, + {'cell_id': 'c', 'sequence_id': '3_1', 'iedb_assay_outcome': 'positive', + 'locus_species': 'NCBITAXON:10090'}) + + kept, _ = self._normalize(frame, outcome='Negative') + self.assertEqual(kept['sequence_id'].tolist(), ['2_1']) + + kept, _ = self._normalize(frame, outcome='*') + self.assertEqual(len(kept), 3) + + kept, report = self._normalize(frame, air_species='Homo sapiens', + outcome='Positive') + self.assertEqual(kept['sequence_id'].tolist(), ['1_1']) + self.assertEqual((report['rows_in'], report['rows_out']), (3, 1)) + + def test_air_species_is_a_bcr_and_tcr_filter(self): + source = IedbSource(makeClient(lambda *a: FakeResponse(200))) + + for table in ('bcr', 'tcr'): + query = source.validateQuery(table, {'air_species': 'Homo sapiens'}) + self.assertEqual(query.filters, {'air_species': 'Homo sapiens', + 'antigen_species': '*', + 'antigen_structure_type': '*', + 'assay_type': '*', 'disease': '*', + 'outcome': '*'}) + with self.assertRaises(Exception): + source.validateQuery('bcr', {'air_species': 'human'}) + + #: epitope_search records: a flu strain, the flu species itself, a human + #: self antigen, and one epitope with two organisms whose name and id + #: lists are not in the same order. + EPITOPES = [ + {'structure_id': 1, + 'source_organism_names': ['Influenza A virus (A/Puerto Rico/8/1934(H1N1))'], + 'source_organism_iris': ['NCBITaxon:211044'], + 'source_organism_iri_search': ['NCBITaxon:1', 'NCBITaxon:11320', + 'NCBITaxon:211044']}, + {'structure_id': 2, + 'source_organism_names': ['Influenza A virus'], + 'source_organism_iris': ['NCBITaxon:11320'], + 'source_organism_iri_search': ['NCBITaxon:1', 'NCBITaxon:11320']}, + {'structure_id': 3, + 'source_organism_names': ['Homo sapiens (human)'], + 'source_organism_iris': ['NCBITaxon:9606'], + 'source_organism_iri_search': ['NCBITaxon:1', 'NCBITaxon:9606']}, + {'structure_id': 4, + 'source_organism_names': ['Borreliella burgdorferi B31', + 'Borreliella burgdorferi'], + 'source_organism_iris': ['NCBITaxon:139', 'NCBITaxon:224326'], + 'source_organism_iri_search': ['NCBITaxon:1', 'NCBITaxon:139', + 'NCBITaxon:224326']}, + ] + + def test_antigen_name_resolves_to_iedb_organism(self): + from immvoke.Exceptions import SchemaError + from immvoke.Sources.Specificity.Iedb import _resolveAntigen + + # The full name wins over a strain sharing it before the brackets. + self.assertEqual(_resolveAntigen('influenza A VIRUS', self.EPITOPES), + ('Influenza A virus', 'NCBITaxon:11320')) + # Without the bracketed common name, when that leaves one organism. + self.assertEqual(_resolveAntigen('Homo sapiens', self.EPITOPES), + ('Homo sapiens (human)', 'NCBITaxon:9606')) + # Names only from an epitope with two organisms cannot be paired. + with self.assertRaises(SchemaError): + _resolveAntigen('Borreliella burgdorferi', self.EPITOPES) + with self.assertRaisesRegex(SchemaError, 'closest IEDB names: Influenza A virus'): + _resolveAntigen('Influenza A', self.EPITOPES) + + def test_antigen_species_keeps_that_organism_and_those_under_it(self): + """'Influenza A virus' keeps the species' own epitope and its + strain's, through IEDB's source_organism_iri_search.""" + frame = iedbAirrRows(*[ + {'cell_id': str(n), 'sequence_id': '%d_1' % n, + 'reactivity_ref': 'IEDB_EPITOPE:%d' % n} for n in (1, 2, 3, 4)]) + + kept, report = self._normalize(frame, antigen_species='Influenza A virus', + handler=itemRangeHandler([self.EPITOPES])) + + self.assertEqual(kept['sequence_id'].tolist(), ['1_1', '2_1']) + self.assertEqual((report['rows_in'], report['rows_out']), (4, 2)) + + def test_each_record_gets_its_epitopes_antigen_and_source_species(self): + """From IEDB's curated_source_antigens, several ' | ' joined, and + its structure_type; an epitope IEDB names no antigen for, or a blank + one, is left blank. The row hash covers IEDB's export only, so it + does not change.""" + epitopes = [ + {'structure_id': 7, 'structure_type': 'Linear peptide', + 'curated_source_antigens': [ + {'name': 'hemagglutinin', + 'source_organism_name': 'Influenza A virus (A/Korea/01/2009(H1N1))'}]}, + {'structure_id': 8, 'structure_type': 'Discontinuous peptide', + 'curated_source_antigens': [ + {'name': 'Spike glycoprotein', 'source_organism_name': 'SARS-CoV2'}, + {'name': 'Spike protein S2', 'source_organism_name': 'SARS-CoV2'}]}, + {'structure_id': 9, 'curated_source_antigens': None}] + frame = iedbAirrRows(*[ + {'cell_id': str(n), 'sequence_id': '%d_1' % n, + 'reactivity_ref': 'IEDB_EPITOPE:%d' % n if n else ''} for n in (7, 8, 9, 0)]) + + kept, _ = self._normalize(frame, handler=itemRangeHandler([epitopes])) + plain, _ = self._normalize(frame.assign(reactivity_ref='')) + + self.assertEqual(kept['iedb_antigen'].tolist(), + ['hemagglutinin', 'Spike glycoprotein | Spike protein S2', '', '']) + self.assertEqual(kept['iedb_antigen_source_species'].tolist(), + ['Influenza A virus (A/Korea/01/2009(H1N1))', 'SARS-CoV2', '', '']) + self.assertEqual(kept['iedb_antigen_structure_type'].tolist(), + ['Linear peptide', 'Discontinuous peptide', '', '']) + self.assertNotEqual(kept['immvoke_row_hash'].tolist(), + plain['immvoke_row_hash'].tolist()) + + def test_antigen_structure_type_keeps_that_structure(self): + """IEDB's value is matched ignoring case; 'Discontinuous peptide' + also keeps its multi chain kind, and combines with + --antigen_species.""" + epitopes = [dict(record, structure_type=kind) for record, kind in zip( + self.EPITOPES, ('Linear peptide', 'Discontinuous peptide', + 'Discontinuous peptide on multi chain', 'non-peptidic'))] + frame = iedbAirrRows(*[ + {'cell_id': str(n), 'sequence_id': '%d_1' % n, + 'reactivity_ref': 'IEDB_EPITOPE:%d' % n} for n in (1, 2, 3, 4)]) + + def kept(**filters): + frame_out, _ = self._normalize(frame, handler=itemRangeHandler([epitopes]), + **filters) + return frame_out['sequence_id'].tolist() + + self.assertEqual(kept(antigen_structure_type='Non-peptidic'), ['4_1']) + self.assertEqual(kept(antigen_structure_type='Discontinuous peptide'), + ['2_1', '3_1']) + self.assertEqual(kept(antigen_structure_type='Discontinuous peptide on multi chain'), + ['3_1']) + self.assertEqual(kept(antigen_structure_type='Discontinuous peptide', + antigen_species='Influenza A virus'), ['2_1']) + + def test_antigen_structure_type_is_a_closed_vocabulary(self): + source = IedbSource(makeClient(lambda *a: FakeResponse(200))) + query = source.validateQuery('bcr', {'antigen_structure_type': 'Linear peptide'}) + self.assertEqual(query.filters['antigen_structure_type'], 'Linear peptide') + with self.assertRaises(Exception): + source.validateQuery('bcr', {'antigen_structure_type': 'conformational'}) + + def _receptors(self): + """Receptors for --paired_only: heavy+kappa, heavy alone, heavy with + a light chain IEDB gave no locus (its V gene names it), alpha+beta, + and heavy+light whose light chain has no species.""" + human = 'NCBITAXON:9606' + return iedbAirrRows( + {'cell_id': 'pair', 'sequence_id': '1_1', 'locus': 'IGH', 'locus_species': human}, + {'cell_id': 'pair', 'sequence_id': '1_2', 'locus': 'IGK', 'locus_species': human}, + {'cell_id': 'heavy', 'sequence_id': '2_1', 'locus': 'IGH', 'locus_species': human}, + {'cell_id': 'noLocus', 'sequence_id': '3_1', 'locus': 'IGH', + 'locus_species': human}, + {'cell_id': 'noLocus', 'sequence_id': '3_2', 'v_call': 'IGLV1-44*01', + 'locus_species': human}, + {'cell_id': 'tcr', 'sequence_id': '4_1', 'locus': 'TRB', 'locus_species': human}, + {'cell_id': 'tcr', 'sequence_id': '4_2', 'locus': 'TRA', 'locus_species': human}, + {'cell_id': 'noSpecies', 'sequence_id': '5_1', 'locus': 'IGH', + 'locus_species': human}, + {'cell_id': 'noSpecies', 'sequence_id': '5_2', 'locus': 'IGK'}) + + def test_paired_only_keeps_receptors_with_both_chains(self): + """A chain with no locus counts by its gene; with --air_species, a + chain with no species does not count, as it is not kept.""" + kept, _ = self._normalize(self._receptors(), paired_only=True) + self.assertEqual(kept['sequence_id'].tolist(), + ['1_1', '1_2', '3_1', '3_2', '4_1', '4_2', '5_1', '5_2']) + + kept, _ = self._normalize(self._receptors(), paired_only=True, + air_species='Homo sapiens') + self.assertEqual(kept['sequence_id'].tolist(), + ['1_1', '1_2', '3_1', '3_2', '4_1', '4_2']) + + def test_paired_only_is_a_switch(self): + source = IedbSource(makeClient(lambda *a: FakeResponse(200))) + query = source.validateQuery('bcr', {'paired_only': True}) + + self.assertIs(query.filters['paired_only'], True) + self.assertEqual(query.filters['air_species'], '*') + self.assertNotIn('paired_only', + source.validateQuery('bcr', {}).filters) + + #: bcell_search records: a neutralization assay on COVID-19, an ELISA + #: on influenza, and a crystal structure with no disease. + ASSAYS = [ + {'bcell_id': 10, 'assay_names': 'neutralization |biological activity', + 'disease_names': ['COVID-19'], 'disease_iris': ['DOID:0080600'], + 'disease_iri_search': ['DOID:4', 'DOID:934', 'DOID:0080600']}, + {'bcell_id': 11, 'assay_names': 'qualitative binding|ELISA', + 'disease_names': ['influenza'], 'disease_iris': ['DOID:8469'], + 'disease_iri_search': ['DOID:4', 'DOID:934', 'DOID:8469']}, + {'bcell_id': 12, 'assay_names': '3D structure|x-ray crystallography|angstroms', + 'disease_names': None, 'disease_iris': None, 'disease_iri_search': None}, + {'bcell_id': 13, 'assay_names': 'qualitative binding|ELISA', + 'disease_names': ['viral infectious disease'], 'disease_iris': ['DOID:934'], + 'disease_iri_search': ['DOID:4', 'DOID:934']}, + ] + + def _assayHandler(self): + """Serve ASSAYS to bcell_search, no epitopes to epitope_search.""" + def handler(method, url, headers, index): + body = self.ASSAYS if 'bcell_search' in url else [] + return FakeResponse(200, json.dumps(body).encode('utf-8')) + return handler + + def _assayed(self): + """A record per assay mix: neutralization+ELISA, ELISA alone, the + structure alone.""" + return iedbAirrRows( + {'cell_id': 'a', 'sequence_id': '1_1', 'iedb_assay_ids': '11, 10'}, + {'cell_id': 'b', 'sequence_id': '2_1', 'iedb_assay_ids': '11'}, + {'cell_id': 'c', 'sequence_id': '3_1', 'iedb_assay_ids': '12'}) + + def test_assay_type_keeps_records_with_an_assay_of_that_type(self): + kept, _ = self._normalize(self._assayed(), handler=self._assayHandler(), + assay_type='Neutralization') + self.assertEqual(kept['sequence_id'].tolist(), ['1_1']) + + kept, _ = self._normalize(self._assayed(), handler=self._assayHandler(), + assay_type='qualitative binding') + self.assertEqual(kept['sequence_id'].tolist(), ['1_1', '2_1']) + + with self.assertRaisesRegex(SchemaError, 'closest IEDB types: neutralization'): + self._normalize(self._assayed(), handler=self._assayHandler(), + assay_type='neutralisation') + + def test_disease_keeps_that_disease_and_those_under_it(self): + kept, _ = self._normalize(self._assayed(), handler=self._assayHandler(), + disease='covid-19') + self.assertEqual(kept['sequence_id'].tolist(), ['1_1']) + + kept, _ = self._normalize(self._assayed(), handler=self._assayHandler(), + disease='viral infectious disease') + self.assertEqual(kept['sequence_id'].tolist(), ['1_1', '2_1']) + + # Both must hold for the same assay: 1_1's neutralization was on + # COVID-19, its influenza assay an ELISA. + kept, _ = self._normalize(self._assayed(), handler=self._assayHandler(), + disease='COVID-19', assay_type='neutralization') + self.assertEqual(kept['sequence_id'].tolist(), ['1_1']) + kept, _ = self._normalize(self._assayed(), handler=self._assayHandler(), + disease='influenza', assay_type='neutralization') + self.assertEqual(kept['sequence_id'].tolist(), []) + + with self.assertRaisesRegex(SchemaError, "no disease named 'measles'"): + self._normalize(self._assayed(), handler=self._assayHandler(), + disease='measles') + + def test_antigen_species_accepts_any_name(self): + """The schema lets any name through; IEDB is asked when fetching.""" + source = IedbSource(makeClient(lambda *a: FakeResponse(200))) + query = source.validateQuery('tcr', {'antigen_species': 'Yellow fever virus'}) + self.assertEqual(query.filters['antigen_species'], 'Yellow fever virus') + + def test_every_species_name_has_codes(self): + from immvoke.Sources.Specificity.Iedb import ( + RECEPTOR_SPECIES, + _speciesCodes, + ) + + for name, parent in RECEPTOR_SPECIES.values(): + self.assertTrue(_speciesCodes(name), name) + self.assertTrue(parent is None or parent in RECEPTOR_SPECIES, name) + + def test_records_pass_through_with_provenance(self): + frame, report = self._normalize(self._antibody()) + + at = IEDB_AIRR_COLUMNS.index('iedb_epitope_name') + 1 + self.assertEqual(list(frame.columns), + list(IEDB_AIRR_COLUMNS[:at]) + + ['iedb_antigen', 'iedb_antigen_source_species', + 'iedb_antigen_structure_type'] + + list(IEDB_AIRR_COLUMNS[at:]) + + ['immvoke_source', 'immvoke_collection', + 'immvoke_unit_id', 'immvoke_row_hash']) + self.assertEqual(frame['sequence_id'].tolist(), ['1284_1', '1284_2']) + self.assertEqual(frame['cell_id'].tolist(), ['1284_2134780p'] * 2) + self.assertEqual(frame['v_call'].tolist(), + ['IGHV5-9*01, IGHV5-9*03', 'IGKV3-2*01']) + self.assertEqual(frame['immvoke_source'].tolist(), ['iedb'] * 2) + self.assertEqual((report['rows_in'], report['rows_out']), (2, 2)) + + def test_writes_as_a_header_and_row_valid_rearrangement(self): + """What `immvoke specificity iedb download bcr` writes is valid AIRR, as is.""" + from immvoke import Convert + + frame, _ = self._normalize(self._antibody()) + + with tempfile.TemporaryDirectory() as outdir: + validation = Convert.writeAirr(iter([frame]), Path(outdir) / 'bcr.tsv') + + self.assertTrue(validation['header_valid']) + self.assertEqual(validation['rows_invalid'], 0) + + def test_written_file_keeps_iedb_columns_and_every_required_field(self): + """ + The way `immvoke specificity iedb download bcr` writes it (strict_airr + off): IEDB's own columns survive, and a required field the export + lacks is still written, blank -- so the file is always a full AIRR + table. + """ + from airr.schema import RearrangementSchema + + from immvoke import Convert + + frame, _ = self._normalize(self._antibody().drop(columns=['d_cigar', + 'rev_comp'])) + + with tempfile.TemporaryDirectory() as outdir: + path = Path(outdir) / 'bcr.tsv' + validation = Convert.writeAirr(iter([frame]), path, strict=False) + written = pandas.read_csv(path, sep='\t', dtype=str, na_filter=False) + + self.assertTrue(validation['header_valid']) + self.assertEqual([f for f in RearrangementSchema.required + if f not in written.columns], []) + self.assertEqual(written['d_cigar'].tolist(), ['', '']) + for column in ('iedb_epitope_name', 'iedb_assay_outcome', 'locus_species', + 'immvoke_source'): + self.assertIn(column, written.columns) + self.assertEqual(written['iedb_epitope_name'].tolist(), ['G58, L59'] * 2) + + def test_annotate_accepts_it_as_the_reference(self): + from immvoke import Convert + from immvoke.Sources.Specificity.Annotate import annotateSpecificity + + frame, _ = self._normalize(self._antibody()) + + with tempfile.TemporaryDirectory() as outdir: + path = Path(outdir) / 'bcr.tsv' + Convert.writeAirr(iter([frame]), path) + single = annotateSpecificity(frame, path) + paired = annotateSpecificity(frame, path, paired=True) + + self.assertEqual(single['hit_id'].tolist(), ['1284_1']) + self.assertEqual(paired['classification'].tolist(), ['both']) + + def test_download_removes_the_raw_mirror(self): + """ + `immvoke specificity iedb download` keeps only the normalized tables: + the bulk ZIP and the tables extracted from it are deleted afterwards. + """ + from argparse import Namespace + from unittest import mock + + from immvoke import Cli + + zip_bytes = makeZip({ + 'bcr_full_airr.tsv': self._antibody().to_csv(sep='\t', index=False), + 'tcr_full_airr.tsv': 'sequence_id\tlocus\nt1\tTRB\n', + }) + args = Namespace(db='iedb', table='bcr', limit=None, dry_run=False, + no_resume=False, strict_airr=False) + + with tempfile.TemporaryDirectory() as outdir: + args.outdir = Path(outdir) + with mock.patch.object(Cli, 'makeClient', + return_value=makeClient(rangeHandler(zip_bytes))): + Cli.handleSpecificityDownload(args) + + self.assertFalse((Path(outdir) / 'iedb' / 'raw').exists()) + self.assertTrue((Path(outdir) / 'iedb' / 'iedb_bcr.tsv').exists()) + self.assertTrue((Path(outdir) / 'iedb' / 'download_metadata.yml').exists()) + + +class TestRowHashes(unittest.TestCase): + """ + Tests for the shared row hash + """ + + def test_matches_hashing_each_row_as_text(self): + """Same hashes as '|'.join(str(x) for x in row), whatever the dtype, + so a download's row hashes do not change with how they are built.""" + import hashlib + + import numpy + + from immvoke.Sources.Specificity.RowHash import rowHashes + + frame = pandas.DataFrame({ + 'a': [1, None, 3], 'b': [1.0, numpy.nan, 2.5], 'c': [True, False, None], + 'd': pandas.array([1, None, 3], dtype='Int64'), + 'e': pandas.array([True, None, False], dtype='boolean'), + 'f': ['x', '', None]}) + expected = [hashlib.sha256('|'.join(str(x) for x in row).encode('utf-8')) + .hexdigest()[:12] for row in frame.itertuples(index=False)] + + self.assertEqual(rowHashes(frame).tolist(), expected) + + def test_empty_frame(self): + from immvoke.Sources.Specificity.RowHash import rowHashes + + frame = pandas.DataFrame({'a': []}, index=pandas.Index([], dtype=int)) + self.assertEqual(len(rowHashes(frame)), 0) diff --git a/src/immvoke/Sources/Specificity/tests/test_Symscan.py b/src/immvoke/Sources/Specificity/tests/test_Symscan.py new file mode 100644 index 0000000..192942e --- /dev/null +++ b/src/immvoke/Sources/Specificity/tests/test_Symscan.py @@ -0,0 +1,492 @@ +""" +Unit tests for symscan-accelerated annotation (Symscan.py) +""" + +# Info +__author__ = 'Pramod Shinde' + +# Imports +import random +import unittest +from unittest import mock + +import pandas + +# Immvoke imports +from immvoke.Exceptions import AnnotationError +from immvoke.Sources.Specificity import Symscan +from immvoke.Sources.Specificity.Genes import buildGeneGate +from immvoke.Sources.Specificity.Symscan import ( + MAX_DISTANCE, + runSymscanAnnotation, +) +from immvoke.Sources.Specificity.tests.BruteForce import bruteForceAnnotation + + +def frame(rows): + return pandas.DataFrame(rows) + + +class TestRunSymscanAnnotation(unittest.TestCase): + """ + Tests for runSymscanAnnotation's criteria and result table + """ + + def test_hamming_within_budget_hits(self): + queries = frame([{'cdr3': 'CARDYW'}]) + reference = frame([{'id': 'r1', 'cdr3': 'CARQYW'}]) # 1 substitution + + result = runSymscanAnnotation(queries, reference, ['cdr3'], 1, id_col='id') + + self.assertEqual(result['hit_ids'].tolist(), ['r1']) + self.assertEqual(result['min_mismatches'].tolist(), [1.0]) + self.assertEqual(result['mismatch_positions'].tolist(), ['3']) + + def test_pairs_score_each_hit(self): + queries = frame([{'cdr3': 'CARDYW'}, {'cdr3': 'CQQQQQ'}, {'cdr3': 'CARDYW'}]) + reference = frame([{'id': 'r1', 'cdr3': 'CARQYW'}, # 1 substitution + {'id': 'r2', 'cdr3': 'CARDYW'}, # identical + {'id': 'r2', 'cdr3': 'CAKDYW'}]) # same id, farther + + result, pairs = runSymscanAnnotation(queries, reference, ['cdr3'], 1, + id_col='id', return_pairs=True) + + self.assertEqual(result['hit_ids'].tolist(), ['r2;r1', '', 'r2;r1']) + self.assertEqual(pairs.columns.tolist(), ['query_row', 'reference_id', 'mismatches', + 'identity', 'mismatch_positions']) + self.assertEqual(pairs[['query_row', 'reference_id', 'mismatches']].values.tolist(), + [[0, 'r2', 0], [0, 'r1', 1], [2, 'r2', 0], [2, 'r1', 1]]) + self.assertEqual(pairs['identity'].tolist(), [1.0, 5 / 6, 1.0, 5 / 6]) + self.assertEqual(pairs['mismatch_positions'].tolist(), ['', '3', '', '3']) + + def test_pairs_without_a_hit_are_empty(self): + _, pairs = runSymscanAnnotation(frame([{'cdr3': 'CARDYW'}]), + frame([{'id': 'r1', 'cdr3': 'CQQQQQ'}]), + ['cdr3'], 0, id_col='id', + distance_type='levenshtein', return_pairs=True) + + self.assertEqual(len(pairs), 0) + self.assertEqual(pairs.columns.tolist(), + ['query_row', 'reference_id', 'mismatches', 'identity']) + + def test_hamming_over_budget_does_not_hit(self): + queries = frame([{'cdr3': 'CARDYW'}]) + reference = frame([{'id': 'r1', 'cdr3': 'CQRQYW'}]) # 2 substitutions + + result = runSymscanAnnotation(queries, reference, ['cdr3'], 1, id_col='id') + + self.assertEqual(result['n_hits_total'].tolist(), [0]) + self.assertTrue(pandas.isna(result['min_mismatches'].iloc[0])) + self.assertEqual(result['hit_ids'].tolist(), ['']) + + def test_exact_cols_gate_the_pairs(self): + queries = frame([{'cdr3': 'CARDYW', 'v_call': 'IGHV1-2'}]) + reference = frame([{'id': 'r1', 'cdr3': 'CARDYW', 'v_call': 'IGHV3-9'}]) + + result = runSymscanAnnotation(queries, reference, ['cdr3'], 1, id_col='id', + exact_cols=['v_call']) + + self.assertEqual(result['n_hits_total'].tolist(), [0]) + + def test_pooled_columns_must_line_up_column_by_column(self): + """'AB'+'CDE' and 'ABC'+'DE' concatenate identically but are not a + column-by-column match.""" + queries = frame([{'cdr1': 'AB', 'cdr3': 'CDE'}]) + reference = frame([{'id': 'r1', 'cdr1': 'ABC', 'cdr3': 'DE'}, + {'id': 'r2', 'cdr1': 'AB', 'cdr3': 'CDF'}]) + + result = runSymscanAnnotation(queries, reference, ['cdr1', 'cdr3'], 1, + id_col='id') + + self.assertEqual(result['hit_ids'].tolist(), ['r2']) + self.assertEqual(result['mismatch_positions'].tolist(), ['1:2']) + + def test_threshold_is_fraction_scales_with_query_length(self): + queries = frame([{'cdr3': 'AAAAAAAAAA'}]) # length 10 -> 1 allowed at 0.1 + reference = frame([{'id': 'one', 'cdr3': 'AAAAAAAAAB'}, + {'id': 'two', 'cdr3': 'AAAAAAAABB'}]) + + result = runSymscanAnnotation(queries, reference, ['cdr3'], 0.1, id_col='id', + threshold_is_fraction=True) + + self.assertEqual(result['hit_ids'].tolist(), ['one']) + + def test_levenshtein_tolerates_an_indel(self): + queries = frame([{'cdr3': 'CARDYW'}]) + reference = frame([{'id': 'r1', 'cdr3': 'CARDGYW'}]) # 1 insertion + + result = runSymscanAnnotation(queries, reference, ['cdr3'], 1, + distance_type='levenshtein', id_col='id') + + self.assertEqual(result['hit_ids'].tolist(), ['r1']) + self.assertNotIn('mismatch_positions', result.columns) + + def test_best_hit_per_id_is_kept(self): + queries = frame([{'cdr3': 'AAAA'}]) + reference = frame([{'id': 'r1', 'cdr3': 'AABB'}, # distance 2 + {'id': 'r1', 'cdr3': 'AAAB'}]) # distance 1, same id + + result = runSymscanAnnotation(queries, reference, ['cdr3'], 2, id_col='id') + + self.assertEqual(result['n_hits_total'].tolist(), [1]) + self.assertEqual(result['min_mismatches'].tolist(), [1.0]) + self.assertEqual(result['mismatch_positions'].tolist(), ['3']) + + def test_x_matches_any_residue_through_symscan_alone(self): + """No pair is compared in Python: X rows go to symscan too.""" + queries = frame([{'cdr3': 'CARXYW'}, {'cdr3': 'CARDYW'}]) + reference = frame([{'id': 'r1', 'cdr3': 'CARDYW'}, + {'id': 'r2', 'cdr3': 'CXRDYW'}]) + + with self._spy() as spy: + result = runSymscanAnnotation(queries, reference, ['cdr3'], 0, + id_col='id') + + self.assertTrue(spy.called) + self.assertEqual(result['hit_ids'].tolist(), ['r1;r2', 'r1;r2']) + self.assertEqual(result['min_mismatches'].tolist(), [0.0, 0.0]) + + def test_x_slack_is_rescored_away(self): + """CARXYW vs CAQDYW is 2 letters apart but 1 real difference: within + the slack symscan is searched with, beyond an exact match.""" + queries = frame([{'cdr3': 'CARXYW'}]) + reference = frame([{'id': 'r1', 'cdr3': 'CAQDYW'}]) + + exact = runSymscanAnnotation(queries, reference, ['cdr3'], 0, id_col='id') + one = runSymscanAnnotation(queries, reference, ['cdr3'], 1, id_col='id') + + self.assertEqual(exact['hit_ids'].tolist(), ['']) + self.assertEqual(one['hit_ids'].tolist(), ['r1']) + self.assertEqual(one['min_mismatches'].tolist(), [1.0]) + self.assertEqual(one['mismatch_positions'].tolist(), ['2']) + + def test_x_is_an_ordinary_letter_under_levenshtein(self): + queries = frame([{'cdr3': 'CARXYW'}]) + reference = frame([{'id': 'r1', 'cdr3': 'CARDYW'}]) + + result = runSymscanAnnotation(queries, reference, ['cdr3'], 0, + distance_type='levenshtein', id_col='id') + + self.assertEqual(result['hit_ids'].tolist(), ['']) + + def test_levenshtein_compares_a_single_column(self): + queries = frame([{'a': 'AA', 'b': 'BB'}]) + reference = frame([{'id': 'r1', 'a': 'AA', 'b': 'BB'}]) + + with self.assertRaises(AnnotationError): + runSymscanAnnotation(queries, reference, ['a', 'b'], 1, + distance_type='levenshtein', id_col='id') + + def test_unknown_distance_type_raises(self): + queries = frame([{'cdr3': 'AA'}]) + reference = frame([{'id': 'r1', 'cdr3': 'AA'}]) + + with self.assertRaises(AnnotationError): + runSymscanAnnotation(queries, reference, ['cdr3'], 1, + distance_type='blosum', id_col='id') + + def _spy(self): + return mock.patch.object(Symscan.symscan, 'get_neighbors_across', + side_effect=Symscan.symscan.get_neighbors_across) + + def _assertSkipped(self, queries, reference, *args, **kwargs): + """symscan is never called; the query row gets no hit, with a + warning naming it.""" + with self._spy() as spy, \ + self.assertLogs('immvoke.Sources.Specificity.Symscan', 'WARNING') as logs: + result = runSymscanAnnotation(queries, reference, *args, id_col='id', + **kwargs) + + spy.assert_not_called() + self.assertEqual(result['hit_ids'].tolist(), ['']) + self.assertIn('query_row 0', logs.output[0]) + + def test_a_distance_symscan_would_crash_on_is_skipped(self): + """symscan panics (killing the process) above MAX_DISTANCE, so it + must never be called with one.""" + self._assertSkipped(frame([{'cdr3': 'A' * 40}]), + frame([{'id': 'r1', 'cdr3': 'B' * 15 + 'A' * 25}]), + ['cdr3'], MAX_DISTANCE + 1) + + def test_reference_xs_are_cut_out_not_searched_around(self): + """A reference's X positions are removed before symscan sees them: + 15 X's cost nothing, and the search stays exact.""" + queries = frame([{'cdr3': 'A' * 20}, {'cdr3': 'B' + 'A' * 19}]) + reference = frame([{'id': 'r1', 'cdr3': 'X' * 15 + 'A' * 5}]) + + with self._spy() as spy: + result = runSymscanAnnotation(queries, reference, ['cdr3'], 0, + id_col='id') + + self.assertEqual(spy.call_args.args, (['AAAAA'] * 2, ['AAAAA'])) + self.assertEqual(result['hit_ids'].tolist(), ['r1', 'r1']) + self.assertEqual(result['mismatch_positions'].tolist(), ['', '']) + + def test_query_x_slack_beyond_symscan_is_skipped(self): + """A query's own X's are covered by slack: 15 of them need 15.""" + self._assertSkipped(frame([{'cdr3': 'X' * 15 + 'A' * 5}]), + frame([{'id': 'r1', 'cdr3': 'A' * 20}]), + ['cdr3'], 0) + + def test_a_sequence_over_the_variant_budget_is_skipped(self): + for distance_type in ('hamming', 'levenshtein'): + with self.subTest(distance_type), \ + mock.patch.object(Symscan, 'MAX_VARIANTS', 1): + self._assertSkipped(frame([{'cdr3': 'CARDYW'}]), + frame([{'id': 'r1', 'cdr3': 'CARQYW'}]), + ['cdr3'], 1, distance_type=distance_type) + + def test_only_the_rows_that_do_not_fit_are_skipped(self): + """At distance 1 a length-4 string has 5 variants and a length-12 + one 13: with half a budget of 20 the long reference cannot be + indexed, but the short pair is still compared.""" + queries = frame([{'cdr3': 'AAAA'}]) + reference = frame([{'id': 'short', 'cdr3': 'AAAC'}, + {'id': 'long', 'cdr3': 'AAAAAAAAAAAA'}]) + + with mock.patch.object(Symscan, 'MAX_VARIANTS', 20), \ + self.assertLogs('immvoke.Sources.Specificity.Symscan', 'WARNING'): + result = runSymscanAnnotation(queries, reference, ['cdr3'], 1, + distance_type='levenshtein', id_col='id') + + self.assertEqual(result['hit_ids'].tolist(), ['short']) + + def test_queries_are_chunked_to_fit_the_variant_budget(self): + """At distance 1 a length-4 string has 5 variants: each side gets + half a budget of 20, so two queries per call -- 4 queries take two + calls.""" + queries = frame([{'cdr3': 'AAAA'}, {'cdr3': 'AAAC'}, {'cdr3': 'AACC'}, + {'cdr3': 'CCCC'}]) + reference = frame([{'id': 'r1', 'cdr3': 'AAAA'}]) + + with self._spy() as spy, mock.patch.object(Symscan, 'MAX_VARIANTS', 20): + result = runSymscanAnnotation(queries, reference, ['cdr3'], 1, id_col='id') + + self.assertEqual(spy.call_count, 2) + self.assertEqual(result['hit_ids'].tolist(), ['r1', 'r1', '', '']) + + def test_references_are_chunked_too(self): + """A reference side larger than half the budget is split across + calls rather than skipped.""" + queries = frame([{'cdr3': 'AAAA'}]) + reference = frame([{'id': 'r%d' % i, 'cdr3': cdr3} + for i, cdr3 in enumerate(['AAAA', 'AAAC', 'AACA', 'CCCC'])]) + + with self._spy() as spy, mock.patch.object(Symscan, 'MAX_VARIANTS', 20): + result = runSymscanAnnotation(queries, reference, ['cdr3'], 1, id_col='id') + + self.assertEqual(spy.call_count, 2) + self.assertEqual(result['hit_ids'].tolist(), ['r0;r1;r2']) + + def test_hamming_only_indexes_same_length_references(self): + queries = frame([{'cdr3': 'CARDYW'}]) + reference = frame([{'id': 'r1', 'cdr3': 'CARQYW'}, + {'id': 'r2', 'cdr3': 'CARDYWW'}]) + + with self._spy() as spy: + runSymscanAnnotation(queries, reference, ['cdr3'], 1, id_col='id') + + self.assertEqual(spy.call_args.args, (['CARDYW'], ['CARQYW'])) + + +class TestSymscanExactMatch(unittest.TestCase): + """ + Tests for an exact match: runSymscanAnnotation at 0 substitutions + """ + + def test_exact_cols_gate_an_exact_match(self): + queries = frame([{'cdr3': 'CARDYW', 'v_call': 'V1', 'j_call': 'J1'}, + {'cdr3': 'CARDYW', 'v_call': 'V2', 'j_call': 'J1'}, + {'cdr3': 'CARDY', 'v_call': 'V1', 'j_call': 'J1'}]) + reference = frame([{'id': 'r2', 'cdr3': 'CARDYW', 'v_call': 'V1', 'j_call': 'J1'}, + {'id': 'r1', 'cdr3': 'CARDYW', 'v_call': 'V1', 'j_call': 'J1'}]) + + result = runSymscanAnnotation(queries, reference, ['cdr3'], 0, id_col='id', + exact_cols=['v_call', 'j_call']) + + self.assertEqual(result['hit_ids'].tolist(), ['r1;r2', '', '']) + self.assertEqual(result['is_unique_hit'].tolist(), [False, True, True]) + + def test_a_blank_value_never_matches(self): + """Not even another blank: there is nothing to compare.""" + queries = frame([{'cdr3': ''}, {'cdr3': 'A'}]) + reference = frame([{'id': 'r1', 'cdr3': ''}, {'id': 'r2', 'cdr3': 'A'}]) + + result = runSymscanAnnotation(queries, reference, ['cdr3'], 0, id_col='id') + + self.assertEqual(result['hit_ids'].tolist(), ['', 'r2']) + + def test_a_blank_in_any_pooled_column_never_matches(self): + queries = frame([{'cdr1': '', 'cdr3': 'CAR'}]) + reference = frame([{'id': 'r1', 'cdr1': '', 'cdr3': 'CAR'}]) + + result = runSymscanAnnotation(queries, reference, ['cdr1', 'cdr3'], 1, + id_col='id') + + self.assertEqual(result['hit_ids'].tolist(), ['']) + + +def randomTables(seed): + """ + Seeded random query and reference tables built to hit the edge cases (X + residues, blanks, duplicate ids, gene calls that agree or not); also + used by test_Distances. + """ + rng = random.Random(seed) + + def residues(n): + value = ''.join(rng.choice('ACD') for _ in range(n)) + if value and rng.random() < 0.2: # at most one X per column + pos = rng.randrange(len(value)) + value = value[:pos] + 'X' + value[pos + 1:] + return value + + def row(i=None): + values = {'cdr1': residues(rng.randint(0, 3)), + 'cdr3': residues(rng.randint(3, 6)), + 'v_call': rng.choice(['IGHV1-2*01', 'IGHV1-2*02', 'IGHV3-23*01', + '']), + 'j_call': rng.choice(['IGHJ4*02', 'IGHJ6*01'])} + if i is not None: + values['id'] = 'r%d' % (i % 50) # some ids span several rows + return values + + return (frame([row() for _ in range(40)]), + frame([row(i) for i in range(70)])) + + +class TestGeneGrouping(unittest.TestCase): + """ + Tests for searching one gene group at a time (_geneGroups) + """ + + @staticmethod + def _rows(*specs): + return frame([{'sequence_id': sid, 'cdr3_aa': cdr3, 'v_call': v, 'j_call': j} + for sid, cdr3, v, j in specs]) + + def _pairs(self, queries, reference, gene_cols, **kwargs): + """The (query_row, reference_id) pairs found, as a sorted set.""" + _, pairs = runSymscanAnnotation( + queries, reference, ['cdr3_aa'], kwargs.pop('tolerance', 2), + id_col='sequence_id', gene_cols=gene_cols, return_pairs=True, + warn_skipped=False, **kwargs) + + return sorted(zip(pairs['query_row'].tolist(), + pairs['reference_id'].tolist())) + + def test_grouping_finds_what_the_gate_alone_would_keep(self): + """Grouping decides what is compared, the gate what is kept: the + two must agree, so searching per gene group loses no pair an + ungrouped search would have kept.""" + queries = self._rows( + ('q1', 'CARDYW', 'IGHV1-2*01', 'IGHJ4*02'), + ('q2', 'CARDYW', 'IGHV3-7*01', 'IGHJ4*02'), + ('q3', 'CARDYY', 'IGHV1-2*03', 'IGHJ6*01')) + reference = self._rows( + ('r_same', 'CARDYW', 'IGHV1-2*02', 'IGHJ4*01'), + ('r_otherV', 'CARDYW', 'IGHV5-9*01', 'IGHJ4*02'), + ('r_otherJ', 'CARDYW', 'IGHV1-2*01', 'IGHJ6*02'), + ('r_q3', 'CARDYY', 'IGHV1-2*01', 'IGHJ6*01')) + + grouped = self._pairs(queries, reference, ['v_call', 'j_call']) + # Every pair an ungrouped search finds, kept by the same gate. + gate = buildGeneGate(queries, reference, ['v_call', 'j_call']) + ids = reference['sequence_id'].tolist() + ungrouped = [(q, rid) for q, rid in self._pairs(queries, reference, None) + if gate.agrees(q, ids.index(rid))] + + self.assertEqual(grouped, sorted(ungrouped)) + # r_otherJ has q1's V but not its J, yet shares both with q3. + self.assertEqual(grouped, [(0, 'r_same'), (2, 'r_otherJ'), (2, 'r_q3')]) + + def test_a_row_calling_several_genes_is_searched_under_each(self): + """A call naming two genes agrees with either, so its row belongs + to both groups rather than only the first.""" + queries = self._rows(('q1', 'CARDYW', 'IGHV1-2*01, IGHV3-7*02', 'IGHJ4*02')) + reference = self._rows(('r_first', 'CARDYW', 'IGHV1-2*05', 'IGHJ4*01'), + ('r_second', 'CARDYW', 'IGHV3-7*01', 'IGHJ4*03'), + ('r_neither', 'CARDYW', 'IGHV5-9*01', 'IGHJ4*02')) + + self.assertEqual(self._pairs(queries, reference, ['v_call', 'j_call']), + [(0, 'r_first'), (0, 'r_second')]) + + def test_an_unrecognized_call_matches_nothing(self): + """As the gate has it: a blank or unparseable call agrees with no + gene, so its row gets no group and no hit.""" + queries = self._rows(('q1', 'CARDYW', '', 'IGHJ4*02'), + ('q2', 'CARDYW', 'not a gene', 'IGHJ4*02'), + ('q3', 'CARDYW', 'IGHV1-2*01', 'IGHJ4*02')) + reference = self._rows(('r1', 'CARDYW', 'IGHV1-2*01', 'IGHJ4*02')) + + self.assertEqual(self._pairs(queries, reference, ['v_call', 'j_call']), + [(2, 'r1')]) + + def test_grouping_narrows_what_symscan_is_asked_to_compare(self): + """The point of grouping: references of another gene are never + indexed, rather than compared and then discarded.""" + queries = self._rows(('q1', 'CARDYW', 'IGHV1-2*01', 'IGHJ4*02')) + # One reference of the query's gene, nineteen of other genes. + reference = self._rows( + ('r_same', 'CARDYW', 'IGHV1-2*04', 'IGHJ4*02'), + *[('r%d' % i, 'CARDYW', 'IGHV%d-9*01' % i, 'IGHJ4*02') for i in range(19)]) + + with mock.patch.object(Symscan.symscan, 'get_neighbors_across', + side_effect=Symscan.symscan.get_neighbors_across) as call: + self._pairs(queries, reference, ['v_call', 'j_call']) + grouped = max(len(c.args[1]) for c in call.call_args_list) + call.reset_mock() + self._pairs(queries, reference, None) + ungrouped = max(len(c.args[1]) for c in call.call_args_list) + + self.assertEqual(ungrouped, 20) + self.assertEqual(grouped, 1) + + +class TestSymscanMatchesBruteForce(unittest.TestCase): + """ + runSymscanAnnotation must return exactly the result and pairs tables a + brute-force comparison of every pair returns -- checked on seeded + random tables built to hit the edge cases (X residues, blanks, duplicate + ids, exact and gene gates, pooled columns, fraction thresholds). + """ + + def _check(self, queries, reference, *args, **kwargs): + expected, expected_pairs = bruteForceAnnotation( + queries, reference, *args, id_col='id', return_pairs=True, **kwargs) + result, pairs = runSymscanAnnotation(queries, reference, *args, id_col='id', + return_pairs=True, **kwargs) + pandas.testing.assert_frame_equal(result, expected, check_dtype=False) + pandas.testing.assert_frame_equal(pairs, expected_pairs, check_dtype=False) + + def test_hamming(self): + cases = [(['cdr3'], 0, {}), (['cdr3'], 1, {}), + (['cdr3'], 2, {'exact_cols': ['j_call']}), + (['cdr1', 'cdr3'], 2, {'gene_cols': ['v_call']}), + (['cdr3'], 0.34, {'threshold_is_fraction': True, + 'gene_cols': ['v_call', 'j_call']}), + (['cdr1', 'cdr3'], 0.25, {'threshold_is_fraction': True})] + for seed in range(5): + queries, reference = randomTables(seed) + for cols, tolerance, kwargs in cases: + with self.subTest(seed=seed, cols=cols, tolerance=tolerance, **kwargs): + self._check(queries, reference, cols, tolerance, **kwargs) + + def test_levenshtein(self): + for seed in range(5): + queries, reference = randomTables(seed) + for max_edits, kwargs in ((1, {}), (2, {'gene_cols': ['v_call', 'j_call']})): + with self.subTest(seed=seed, max_edits=max_edits): + self._check(queries, reference, ['cdr3'], max_edits, + distance_type='levenshtein', **kwargs) + + def test_chunked_queries_give_the_same_result(self): + """A budget small enough to split every group's queries over many + symscan calls, yet still fitting each single query.""" + queries, reference = randomTables(0) + with mock.patch.object(Symscan, 'MAX_VARIANTS', 2000): + self._check(queries, reference, ['cdr3'], 1) + + +if __name__ == '__main__': + unittest.main() diff --git a/src/immvoke/data/schemas/asd/schema.yaml b/src/immvoke/data/schemas/asd/schema.yaml new file mode 100644 index 0000000..bb0360c --- /dev/null +++ b/src/immvoke/data/schemas/asd/schema.yaml @@ -0,0 +1,43 @@ +collections: + asd: + fields: + - name: dataset + pseudo_values: false + values: + - aae + - aatp + - ab-bind + - abbd + - abdesign + - alphaseq + - biomap + - buzz + - covid-19 + - dlgo + - flab_hie2022 + - flab_koenig2017 + - flab_rosace2023 + - flab_shanehsazzadeh2023 + - flab_warszawski2019 + - genbank + - hiv + - literature + - met + - osh + - patents + - rmna + - skempiv2 + - structures-antibodies + - structures-nanobodies + wildcard: '*' + reported_totals: {} +field_aliases: {} +harvested: '2026-08-19T00:00:00Z' +harvested_by: 'immvoke 0.1.0 (hand curated, not scraped: ASD exposes no HTML form)' +parse_contracts: {} +schema_version: 1 +source: asd +source_urls: + drive_folder: https://drive.google.com/drive/folders/1HY2GOVj-HR8t6Jmhe5MIJ6xmBCOxC45q + home: https://naturalantibody.com/agab/ +url_rules: {} diff --git a/src/immvoke/data/schemas/covabdab/schema.yaml b/src/immvoke/data/schemas/covabdab/schema.yaml new file mode 100644 index 0000000..6e87e8b --- /dev/null +++ b/src/immvoke/data/schemas/covabdab/schema.yaml @@ -0,0 +1,13 @@ +collections: + antibodies: + fields: [] + reported_totals: {} +field_aliases: {} +harvested: '2026-08-19T00:00:00Z' +harvested_by: 'immvoke 0.1.0 (hand curated, not scraped: CoV-AbDab exposes no HTML form)' +parse_contracts: {} +schema_version: 1 +source: covabdab +source_urls: + home: https://opig.stats.ox.ac.uk/webapps/covabdab/ +url_rules: {} diff --git a/src/immvoke/data/schemas/iedb/schema.yaml b/src/immvoke/data/schemas/iedb/schema.yaml new file mode 100644 index 0000000..fe4733e --- /dev/null +++ b/src/immvoke/data/schemas/iedb/schema.yaml @@ -0,0 +1,193 @@ +source: iedb +schema_version: 1 +harvested: '2026-08-14T00:00:00Z' +harvested_by: 'immvoke 0.1.0 (hand curated, not scraped: IEDB exposes no HTML form)' +source_urls: + airr_bulk: https://www.iedb.org/downloader.php?file_name=doc/airr_full.zip + api: https://query-api.iedb.org/api/v1 +field_aliases: {} +url_rules: {} +parse_contracts: {} +collections: + bcr: + fields: + - name: air_species + wildcard: '*' + pseudo_values: false + values: + - Bos taurus + - Camelid + - Camelus bactrianus + - Camelus dromedarius + - Camelus ferus + - Chiloscyllium plagiosum + - Chlorocebus aethiops + - Gallus gallus + - Ginglymostoma cirratum + - Homo sapiens + - Lama glama + - Macaca fascicularis + - Macaca mulatta + - Mus musculus + - 'Mus musculus 129 GIX+' + - Mus musculus B10.RIII + - 'Mus musculus BALB/c' + - 'Mus musculus C3H/He' + - 'Mus musculus C57BL/6' + - Mus musculus CBA + - 'Mus musculus CBA/J' + - Mus musculus FVB + - 'Mus musculus HuMAb (Medarex)' + - Mus musculus Kunming + - Mus musculus Kymouse + - Mus musculus PrP null + - Mus musculus RBF + - Mus musculus Swiss Webster + - Mus musculus Xenomouse + - Myodes glareolus + - Nothocricetulus migratorius + - Orectolobus maculatus + - Oryctolagus cuniculus + - Oryctolagus cuniculus New Zealand White + - Pan troglodytes + - Rattus + - Rattus norvegicus + - Rattus norvegicus Sprague-Dawley + - Rattus rattus + - Squalus acanthias + - Sus scrofa + - Vicugna pacos + - synthetic construct + - name: antigen_species + wildcard: '*' + pseudo_values: false + open_values: true + values: + - Influenza A virus + - SARS-CoV2 + - Human immunodeficiency virus 1 + - Homo sapiens + - name: antigen_structure_type + wildcard: '*' + pseudo_values: false + values: + - Linear peptide + - Discontinuous peptide + - Discontinuous peptide on multi chain + - Non-peptidic + - name: assay_type + wildcard: '*' + pseudo_values: false + open_values: true + values: + - neutralization + - 3D structure + - qualitative binding + - dissociation constant KD + - name: disease + wildcard: '*' + pseudo_values: false + open_values: true + values: + - COVID-19 + - influenza + - rheumatoid arthritis + - viral infectious disease + - name: outcome + wildcard: '*' + pseudo_values: false + values: + - Positive + - Negative + reported_totals: {} + tcr: + fields: + - name: air_species + wildcard: '*' + pseudo_values: false + values: + - Bos taurus + - Camelid + - Camelus bactrianus + - Camelus dromedarius + - Camelus ferus + - Chiloscyllium plagiosum + - Chlorocebus aethiops + - Gallus gallus + - Ginglymostoma cirratum + - Homo sapiens + - Lama glama + - Macaca fascicularis + - Macaca mulatta + - Mus musculus + - 'Mus musculus 129 GIX+' + - Mus musculus B10.RIII + - 'Mus musculus BALB/c' + - 'Mus musculus C3H/He' + - 'Mus musculus C57BL/6' + - Mus musculus CBA + - 'Mus musculus CBA/J' + - Mus musculus FVB + - 'Mus musculus HuMAb (Medarex)' + - Mus musculus Kunming + - Mus musculus Kymouse + - Mus musculus PrP null + - Mus musculus RBF + - Mus musculus Swiss Webster + - Mus musculus Xenomouse + - Myodes glareolus + - Nothocricetulus migratorius + - Orectolobus maculatus + - Oryctolagus cuniculus + - Oryctolagus cuniculus New Zealand White + - Pan troglodytes + - Rattus + - Rattus norvegicus + - Rattus norvegicus Sprague-Dawley + - Rattus rattus + - Squalus acanthias + - Sus scrofa + - Vicugna pacos + - synthetic construct + - name: antigen_species + wildcard: '*' + pseudo_values: false + open_values: true + values: + - Influenza A virus + - SARS-CoV2 + - Human immunodeficiency virus 1 + - Homo sapiens + - name: antigen_structure_type + wildcard: '*' + pseudo_values: false + values: + - Linear peptide + - Discontinuous peptide + - Discontinuous peptide on multi chain + - Non-peptidic + - name: assay_type + wildcard: '*' + pseudo_values: false + open_values: true + values: + - neutralization + - 3D structure + - qualitative binding + - dissociation constant KD + - name: disease + wildcard: '*' + pseudo_values: false + open_values: true + values: + - COVID-19 + - influenza + - rheumatoid arthritis + - viral infectious disease + - name: outcome + wildcard: '*' + pseudo_values: false + values: + - Positive + - Negative + reported_totals: {} diff --git a/tests/FakeHttp.py b/tests/FakeHttp.py index 7286a11..813db0d 100644 --- a/tests/FakeHttp.py +++ b/tests/FakeHttp.py @@ -11,6 +11,8 @@ __author__ = 'Susanna Marquez' # Imports +import json + import requests from requests.structures import CaseInsensitiveDict @@ -32,11 +34,21 @@ def __init__(self, status_code=200, body=b'', headers=None): self.headers = CaseInsensitiveDict(headers or {}) self.closed = False + @property + def text(self): + """str: the body decoded as UTF-8, like requests.Response.text.""" + return self.content.decode('utf-8') if isinstance(self.content, bytes) \ + else self.content + def iter_content(self, chunk_size=1024): """Yield the body in chunk_size pieces.""" for start in range(0, len(self.content), chunk_size): yield self.content[start:start + chunk_size] + def json(self): + """Parse the body as JSON, like requests.Response.json().""" + return json.loads(self.content) + def close(self): """Mark the response closed.""" self.closed = True diff --git a/tests/test_Cli.py b/tests/test_Cli.py index 7b4b088..afa47c7 100644 --- a/tests/test_Cli.py +++ b/tests/test_Cli.py @@ -26,6 +26,7 @@ handleReferenceDiff, handleReferenceDownload, handleReferenceShow, + handleSpecificityDownload, loadMap, ) from immvoke.Exceptions import ImmvokeError @@ -91,6 +92,61 @@ def test_action_help_lists_the_collections(self): self.assertIn(collection, text) self.assertIn(OasSource.collection_help[collection], text) + def test_specificity_help_lists_the_tables(self): + """`immvoke specificity iedb download --help` names the tables it accepts.""" + from immvoke.Sources.Specificity.Iedb import IedbSource + + parser = getArgParser() + with mock.patch('sys.stdout', new_callable=io.StringIO) as out: + with self.assertRaises(SystemExit): + parser.parse_args(['specificity', 'iedb', 'download', '--help']) + + text = out.getvalue() + for table in IedbSource.collections: + self.assertIn(table, text) + + def test_specificity_all_download_takes_no_table(self): + """ + `specificity all` fans out to every database's every table, so it has + no table level of its own the way `specificity ` does. + """ + parser = getArgParser() + args = parser.parse_args(['specificity', 'all', 'download', + '--outdir', '/tmp/does-not-need-to-exist']) + + self.assertEqual(args.db, 'all') + self.assertEqual(args.action, 'download') + self.assertFalse(hasattr(args, 'table')) + + def test_specificity_all_tables_take_the_filters_they_share(self): + """ + IEDB's bcr and tcr filter on the same fields, so `iedb + all` takes them; CoV-AbDab has no filters to share. + """ + parser = getArgParser() + for action in ('search', 'download'): + outdir = ['--outdir', '/tmp/does-not-need-to-exist'] if action == 'download' else [] + args = parser.parse_args(['specificity', 'iedb', action, 'all', + '--air_species', 'Homo sapiens'] + outdir) + self.assertEqual(args.filter_air_species, 'Homo sapiens') + + with self.assertRaises(SystemExit), mock.patch('sys.stderr'): + parser.parse_args(['specificity', 'covabdab', 'search', 'all', + '--air_species', 'Homo sapiens']) + + def test_schema_commands_accept_specificity_databases(self): + """ + A rejected filter value tells the user to run `immvoke schema show` + for that source (see Schema.validateFilters), so the schema commands + have to accept a specificity database, not only a repertoire source. + """ + parser = getArgParser() + for action in ('show', 'refresh'): + for source in ('iedb', 'covabdab', 'asd', 'oas'): + with self.subTest(action=action, source=source): + args = parser.parse_args(['schema', action, '--source', source]) + self.assertEqual(args.source, source) + def test_collection_is_required(self): """ Omitting the collection is a parse error, not a runtime one. @@ -262,6 +318,27 @@ def test_both_formats_write_one_samplesheet_each(self): self.assertTrue((self.outdir / 'samplesheet_airrflow_fasta.tsv').exists()) +class TestHandleSpecificityDownload(unittest.TestCase): + """ + Tests for the specificity download command's output layout + """ + + def test_each_database_gets_its_own_folder(self): + with tempfile.TemporaryDirectory() as tmp: + args = Namespace(db='fakedb', table='paired', outdir=Path(tmp), limit=None, + dry_run=False, no_resume=False, strict_airr=False) + + with mock.patch('immvoke.Cli.getSpecificitySource', + return_value=StubSource(None)): + self.assertEqual(handleSpecificityDownload(args), 0) + + written = sorted(str(p.relative_to(tmp)) for p in Path(tmp).rglob('*') + if p.is_file()) + self.assertEqual(written, ['fakedb/download_metadata.yml', + 'fakedb/fakedb_paired.tsv', + 'fakedb/raw/' + StubSource.unit.unit_id]) + + def makeReference(root, chain='IGHV', records=(('IGHV1-2*02', 'ACGT'),)): """Write a minimal reference_base and return its root.""" root = Path(root) diff --git a/tests/test_Convert.py b/tests/test_Convert.py index ad0111c..ee391ef 100644 --- a/tests/test_Convert.py +++ b/tests/test_Convert.py @@ -408,6 +408,27 @@ def test_extras_are_kept_by_default(self): self.assertIn('immvoke_unit_id', frame.columns) self.assertIn('immvoke_row_hash', frame.columns) + def test_validation_log_is_named_after_the_table(self): + out, validation = self.writeCase(UNPAIRED) + report = Convert.writeValidationReport(validation, out) + + self.assertEqual(report.name, 'out.airr_validation.log') + text = report.read_text() + self.assertIn('AIRR validation of out.tsv', text) + self.assertIn('all valid', text) + self.assertIn('Result: PASS', text) + + def test_validation_log_lists_row_errors(self): + validation = Convert.newValidation() + validation.update(header_valid=True, rows_checked=200, rows_invalid=3) + validation['errors']['productive is not a boolean'] = 3 + + text = Convert.summarizeValidation(validation, 'bcr.tsv') + + self.assertIn('200 checked, 3 invalid (1.5%)', text) + self.assertIn(' 3 productive is not a boolean', text) + self.assertIn('Result: FAIL', text) + class TestFastaOutput(unittest.TestCase): """ diff --git a/tests/test_Schema.py b/tests/test_Schema.py index 6219e14..993e17f 100644 --- a/tests/test_Schema.py +++ b/tests/test_Schema.py @@ -104,6 +104,13 @@ def test_unknown_collection_raises(self): with self.assertRaises(SchemaError): self.schema.validateFilters('nonesuch', {}) + def test_open_values_accept_anything(self): + """An open field's values are examples; the source checks the rest.""" + field = Field(name='antigen', values=('Influenza A virus',), open_values=True) + + self.assertTrue(field.accepts('Yellow fever virus')) + self.assertFalse(Field(name='antigen', values=('x',)).accepts('y')) + class TestRoundTrip(unittest.TestCase): """ @@ -118,6 +125,19 @@ def test_round_trip(self): self.assertEqual(again.getCollection('unpaired').field_names, ('Species', 'Primer')) + def test_open_values_round_trip_and_stay_out_of_other_fields(self): + schema = SourceSchema(source='demo', collections={ + 'c': Collection(name='c', fields=( + Field(name='antigen', values=('a',), open_values=True), + Field(name='Species', values=('human',))))}) + + payload = toDict(schema) + fields = payload['collections']['c']['fields'] + + self.assertTrue(fields[0]['open_values']) + self.assertNotIn('open_values', fields[1]) + self.assertTrue(fromDict(payload).getCollection('c').fields[0].open_values) + def test_a_newer_snapshot_is_refused(self): """ A snapshot from a future version is refused rather than partly read. @@ -197,6 +217,28 @@ def test_the_two_oas_collections_differ_as_the_live_forms_do(self): self.assertIn('Primer', unpaired) self.assertTrue({'BSource', 'BType', 'Subject'} <= paired) + def test_every_packaged_snapshot_round_trips(self): + """ + A snapshot must serialize back to what is checked in, so that + `immvoke schema refresh` on an unchanged source leaves the file + alone (see saveSchema) rather than reformatting a tracked file and + opening a pull request with no real change in it. + """ + from importlib import resources + + from immvoke.Schema import sameContent, serializeSchema + + root = resources.files('immvoke').joinpath('data/schemas') + packaged = sorted(x.name for x in root.iterdir() + if x.joinpath('schema.yaml').is_file()) + self.assertTrue(packaged, 'no packaged snapshots found') + + for name in packaged: + with self.subTest(source=name): + stored = root.joinpath(name, 'schema.yaml').read_text() + self.assertTrue(sameContent(stored, serializeSchema(loadSchema(name))), + '%s would be rewritten by schema refresh' % name) + if __name__ == '__main__': unittest.main()