diff --git a/genome_kit/variant_genome.py b/genome_kit/variant_genome.py index 50c619eb..f1837be4 100644 --- a/genome_kit/variant_genome.py +++ b/genome_kit/variant_genome.py @@ -1,13 +1,21 @@ # Copyright (C) 2016-2023 Deep Genomics Inc. All Rights Reserved. from __future__ import absolute_import +import gzip +import hashlib +import os import pickle +import re +import tempfile from functools import partial from operator import attrgetter from . import _util +from . import gk_data from . import interval as _interval from ._apply_variants import apply_variants, check_variants_list +from .genome import Genome +from .variant import Variant class VariantGenome(object): @@ -21,10 +29,22 @@ class VariantGenome(object): will be aligned to the variant genome's coordinate system. TODO: see explanation of anchors. + + A variant genome can be persisted and recreated later, in another process + or on another machine, from its content-derived :py:attr:`id`:: + + >>> variant_genome = VariantGenome(genome, variants) + >>> variant_genome_id = variant_genome.save() # keep this id, e.g. in a database + >>> VariantGenome.load(variant_genome_id) == variant_genome + True + + Saving requires a writable data manager; see :py:meth:`save`. """ __slots__ = ('genome', 'variants') + _FILENAME_FORMAT = "variant_genome.{}.txt.gz" + def __init__(self, reference_genome, variants): """Initialize a variant genome. @@ -204,20 +224,83 @@ def __repr__(self): def reference_genome(self): return self.genome.reference_genome + def _canonical_variants(self) -> list[Variant]: + # Normalized variants in a deterministic order, shared by `__eq__` and `id`. + return sorted((vv for v in self.variants for vv in v._normalized_variant), + key=attrgetter("chromosome", "start", "ref", "alt")) + + @property + def id(self) -> str: + """A content-derived identifier for this variant genome (hex SHA-256 digest of the genome config and the normalized variants). + + Two variant genomes have the same id if and only if they compare equal, + i.e. they have the same genome config and the same variants after normalization. + """ + hasher = hashlib.sha256(self.genome.config.encode()) + for variant in self._canonical_variants(): + hasher.update(b"\n" + variant.as_variant_string().encode()) + return hasher.hexdigest() + + def save(self) -> str: + """Save this variant genome so that it can be recreated with :py:meth:`load`. + Returns the :py:attr:`id` to pass to :py:meth:`load`. + + In the current implementation, a file is uploaded per-variant genome with ``genome_kit.gk_data.data_manager``. + Equal variant genomes share a single file. + + The default open source data manager points at GenomeKit's public bucket, which is + read-only. To save, bring your own data manager (see :ref:`sharing-data`). + """ + variant_genome_id = self.id + filename = self._FILENAME_FORMAT.format(variant_genome_id) + with tempfile.TemporaryDirectory() as tmpdir: + path = os.path.join(tmpdir, filename) + lines = [self.genome.config] + [v.as_variant_string() for v in self._canonical_variants()] + # Fixed gzip header so equal genomes produce identical bytes + with open(path, "wb") as raw, gzip.GzipFile(fileobj=raw, mode="wb", filename="", mtime=0) as f: + f.write("".join(line + "\n" for line in lines).encode()) + gk_data.data_manager.upload_file(path, filename) + return variant_genome_id + + @classmethod + def load(cls, variant_genome_id: str) -> "VariantGenome": + """Recreate a variant genome that was stored with :py:meth:`save`. + Returns a variant genome equal to the one that was saved. + + In the current implementation, a file is fetched with ``genome_kit.gk_data.data_manager``. + + Parameters + ---------- + variant_genome_id : :py:class:`str` + The id returned by :py:meth:`save`. + + Raises + ------ + :py:class:`~genome_kit.data_manager.GKDataFileNotFoundError` + No variant genome with this id was found. + :py:exc:`ValueError` + `variant_genome_id` is not an id produced by :py:meth:`save`, or the stored file does not match it. + """ + if not re.fullmatch("[0-9a-f]{64}", variant_genome_id): + raise ValueError("Invalid variant genome id: {!r}".format(variant_genome_id)) + path = gk_data.data_manager.get_file(cls._FILENAME_FORMAT.format(variant_genome_id)) + with gzip.open(path, "rt") as f: + genome = Genome(f.readline().rstrip("\n")) + variants = [] + for line in f: + chrom, pos, ref, alt = line.rstrip("\n").split(":") + variants.append(Variant(chrom, int(pos) - 1, ref, alt, genome)) + variant_genome = cls(genome, variants) + if variant_genome.id != variant_genome_id: + raise ValueError("Stored variant genome does not match id {}".format(variant_genome_id)) + return variant_genome + def __eq__(self, other): if not isinstance(other, VariantGenome): return False - if self.genome != other.genome: - return False - # sort and normalize both sides. - # `start` is ok to use here because VariantGenome doesn't support overlapping variants in any case. - vself = sorted((vv for n in (v._normalized_variant for v in self.variants) for vv in n), - key=attrgetter('start')) - vother = sorted((vv for n in (v._normalized_variant for v in other.variants) for vv in n), - key=attrgetter('start')) - if vself != vother: + if self.genome.config != other.genome.config: return False - return True + return self._canonical_variants() == other._canonical_variants() def __getattr__(self, name): # Intercept any attribute requests that weren't found on the VariantGenome object itself, diff --git a/src/gk_assert.h b/src/gk_assert.h index 8fbd7a22..11914aec 100644 --- a/src/gk_assert.h +++ b/src/gk_assert.h @@ -6,6 +6,7 @@ Copyright (C) 2016-2023 Deep Genomics Inc. All Rights Reserved. #define __GK_ASSERT_H__ #include "defines.h" +#include #include #include #include diff --git a/tests/test_variant_genome.py b/tests/test_variant_genome.py index c30457c9..b04230de 100644 --- a/tests/test_variant_genome.py +++ b/tests/test_variant_genome.py @@ -2,11 +2,34 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +import gzip +import os +import shutil +import tempfile import unittest +from unittest.mock import patch +from genome_kit import DataManager from genome_kit import VariantGenome from genome_kit import Interval from genome_kit import Variant -from . import MiniGenome +from genome_kit import gk_data +from genome_kit.data_manager import GKDataFileNotFoundError +from . import MiniGenome, get_test_file + + +class LocalDataManager(DataManager): + """Stores uploads in a local directory, serving other files from the test data.""" + + def get_file(self, filename): + path = os.path.join(self.data_dir, filename) + if os.path.exists(path): + return path + if filename.startswith(VariantGenome._FILENAME_FORMAT.split("{")[0]): + raise GKDataFileNotFoundError(filename, "not found") + return get_test_file(filename) + + def upload_file(self, filepath, filename, metadata=None): + shutil.copyfile(filepath, os.path.join(self.data_dir, filename)) class TestVariantGenome(unittest.TestCase): @@ -265,6 +288,67 @@ def test_string_variant(self): with self.assertRaises(TypeError): VariantGenome(self.genome, ["chr1:15:G:T"]) + def _variant(self, chrom, start, alt, ref_len=1, genome=None): + genome = genome or self.genome + ref = genome.dna(Interval(chrom, "+", start, start + ref_len, genome)) + return Variant(chrom, start, ref, alt, genome) + + def test_eq(self): + # Same start on different chromosomes, in different orders + a = self._variant("chr1", 10, "T") + b = self._variant("chr2", 10, "C") + self.assertEqual(VariantGenome(self.genome, [a, b]), VariantGenome(self.genome, [b, a])) + + # Un-normalized variant equals its normalized form + ref = self.genome.dna(Interval("chr1", "+", 10, 12, self.genome)) + # G-padded = "11:GT:GC", non-padded = "12:T:C" + padded = Variant("chr1", 10, ref, ref[0] + "C", self.genome) + non_padded = self._variant("chr1", 11, "C") + self.assertEqual(VariantGenome(self.genome, padded), VariantGenome(self.genome, non_padded)) + + # different variants + self.assertNotEqual(VariantGenome(self.genome, a), VariantGenome(self.genome, b)) + # different genome + self.assertNotEqual(VariantGenome(self.genome, a), VariantGenome(MiniGenome("hg19"), self._variant("chr1", 10, "T", genome=MiniGenome("hg19")))) + + def test_id(self): + a = self._variant("chr1", 10, "T") + b = self._variant("chr2", 10, "C") + vg = VariantGenome(self.genome, [a, b]) + self.assertRegex(vg.id, "^[0-9a-f]{64}$") + self.assertEqual(vg.id, VariantGenome(self.genome, [b, a]).id) + self.assertNotEqual(vg.id, VariantGenome(self.genome, a).id) + self.assertNotEqual(vg.id, VariantGenome(self.genome, []).id) + + def test_save_load(self): + variants = [self._variant("chr1", 10, "T"), self._variant("chr2", 10, "C"), self._variant("chr1", 20, "", ref_len=2)] + vg = VariantGenome(self.genome, variants) + with tempfile.TemporaryDirectory() as tmpdir, patch.object(gk_data, "data_manager", LocalDataManager(tmpdir)): + vg_id = vg.save() + self.assertEqual(vg_id, vg.id) + with open(os.path.join(tmpdir, VariantGenome._FILENAME_FORMAT.format(vg_id)), "rb") as f: + saved = f.read() + VariantGenome(self.genome, variants[::-1]).save() + with open(os.path.join(tmpdir, VariantGenome._FILENAME_FORMAT.format(vg_id)), "rb") as f: + self.assertEqual(f.read(), saved) + loaded = VariantGenome.load(vg_id) + self.assertEqual(loaded, vg) + self.assertEqual(loaded.id, vg_id) + self.assertEqual(loaded.dna(Interval("chr1", "+", 5, 25, loaded)), vg.dna(Interval("chr1", "+", 5, 25, vg))) + + with self.assertRaises(GKDataFileNotFoundError): + VariantGenome.load("0" * 64) + + for bad_id in ["", vg_id[:-1], vg_id.upper(), "../" + vg_id, vg_id + "\n"]: + with self.assertRaisesRegex(ValueError, "Invalid variant genome id"): + VariantGenome.load(bad_id) + + # Stored content that does not match its id is rejected + with gzip.open(os.path.join(tmpdir, VariantGenome._FILENAME_FORMAT.format(vg_id)), "wt") as f: + f.write(self.genome.config + "\n") + with self.assertRaisesRegex(ValueError, "does not match id"): + VariantGenome.load(vg_id) + def test_allow_outside_chromsome(self): variant = VariantGenome( self.genome, Variant.from_string("chr1:15:G:T", self.genome)