Dataset model

clickhouse_search/management/commands/set_saved_variant_key.py

@@ -6,7 +6,7 @@
 
 from clickhouse_search.search import get_clickhouse_key_lookup
 from reference_data.models import GENOME_VERSION_GRCh38, GENOME_VERSION_GRCh37
-from seqr.models import SavedVariant, Sample
+from seqr.models import SavedVariant, Dataset
 from seqr.utils.xpos_utils import parse_variant_id
 
 logger = logging.getLogger(__name__)
@@ -22,7 +22,7 @@ def handle(self, *args, **options):
             saved_variant_json__populations__isnull=False, # Omit manual variants
         ).values_list('variant_id', flat=True).distinct()
         ids_by_dataset_type = {
-            Sample.DATASET_TYPE_VARIANT_CALLS: [], Sample.DATASET_TYPE_MITO_CALLS: [], Sample.DATASET_TYPE_SV_CALLS: [],
+            Dataset.DATASET_TYPE_VARIANT_CALLS: [], Dataset.DATASET_TYPE_MITO_CALLS: [], Dataset.DATASET_TYPE_SV_CALLS: [],
         }
         for variant_id in variant_ids:
             parsed_id = parse_variant_id(variant_id)
@@ -31,31 +31,31 @@ def handle(self, *args, **options):
                 parsed_id = parse_variant_id(variant_id[:-1])
             if parsed_id:
                 is_mito = parsed_id[0].replace('chr', '').startswith('M')
-                dataset_type = Sample.DATASET_TYPE_MITO_CALLS if is_mito else Sample.DATASET_TYPE_VARIANT_CALLS
+                dataset_type = Dataset.DATASET_TYPE_MITO_CALLS if is_mito else Dataset.DATASET_TYPE_VARIANT_CALLS
             else:
-                dataset_type = Sample.DATASET_TYPE_SV_CALLS
+                dataset_type = Dataset.DATASET_TYPE_SV_CALLS
             ids_by_dataset_type[dataset_type].append(variant_id)
 
-        no_key_mito = self._set_variant_keys(ids_by_dataset_type[Sample.DATASET_TYPE_MITO_CALLS], Sample.DATASET_TYPE_MITO_CALLS)
+        no_key_mito = self._set_variant_keys(ids_by_dataset_type[Dataset.DATASET_TYPE_MITO_CALLS], Dataset.DATASET_TYPE_MITO_CALLS)
 
         no_key_snv_indel = self._set_variant_keys(
-            ids_by_dataset_type[Sample.DATASET_TYPE_VARIANT_CALLS] + list(no_key_mito), Sample.DATASET_TYPE_VARIANT_CALLS,
+            ids_by_dataset_type[Dataset.DATASET_TYPE_VARIANT_CALLS] + list(no_key_mito), Dataset.DATASET_TYPE_VARIANT_CALLS,
         )
         if no_key_snv_indel:
             self._resolve_missing_variants(no_key_snv_indel, GENOME_VERSION_GRCh38)
 
         no_keys_svs = self._set_variant_keys(
-            ids_by_dataset_type[Sample.DATASET_TYPE_SV_CALLS], f'{Sample.DATASET_TYPE_SV_CALLS}_{Sample.SAMPLE_TYPE_WGS}',
+            ids_by_dataset_type[Dataset.DATASET_TYPE_SV_CALLS], f'{Dataset.DATASET_TYPE_SV_CALLS}_{Dataset.SAMPLE_TYPE_WGS}',
         )
-        no_keys_svs = self._set_variant_keys(list(no_keys_svs), f'{Sample.DATASET_TYPE_SV_CALLS}_{Sample.SAMPLE_TYPE_WES}')
+        no_keys_svs = self._set_variant_keys(list(no_keys_svs), f'{Dataset.DATASET_TYPE_SV_CALLS}_{Dataset.SAMPLE_TYPE_WES}')
         if no_keys_svs:
             self._resolve_reloaded_svs(no_keys_svs)
 
         variant_ids_37 = SavedVariant.objects.filter(
             key__isnull=True, family__project__genome_version=GENOME_VERSION_GRCh37,
         ).values_list('variant_id', flat=True).distinct()
 
-        no_keys_37 = self._set_variant_keys(variant_ids_37, Sample.DATASET_TYPE_VARIANT_CALLS, genome_version=GENOME_VERSION_GRCh37)
+        no_keys_37 = self._set_variant_keys(variant_ids_37, Dataset.DATASET_TYPE_VARIANT_CALLS, genome_version=GENOME_VERSION_GRCh37)
         if no_keys_37:
             self._resolve_missing_variants(no_keys_37, GENOME_VERSION_GRCh37)
 
@@ -103,7 +103,7 @@ def _query_missing_variants(cls, variant_ids, variant_fields, genome_version=GEN
             variant_id__in=variant_ids, family__project__genome_version=genome_version,
         )
         num_missing = missing_variants.count()
-        missing_with_data_qs = missing_variants.filter(family__individual__sample__is_active=True).distinct()
+        missing_with_data_qs = missing_variants.filter(family__individual__active_datasets__isnull=False).distinct()
         missing_with_search_data = missing_with_data_qs.values(
             'variant_id', *variant_fields,
         ).annotate(family_ids=ArrayAgg('family__family_id', distinct=True)).order_by('variant_id')
@@ -128,7 +128,7 @@ def _resolve_missing_variants(cls, variant_ids, genome_version):
     @classmethod
     def _resolve_reloaded_svs(cls, variant_ids):
         missing_with_search_data, num_missing = cls._query_missing_variants(
-            list(variant_ids), ['family__individual__sample__sample_type'],
+            list(variant_ids), ['family__individual__active_datasets__sample_type'],
         )
         logger.info(
             f'{num_missing} SV variants have no key, {num_missing - len(missing_with_search_data)} of which have no search data'
@@ -139,7 +139,7 @@ def _resolve_reloaded_svs(cls, variant_ids):
         missing_by_sample_type = defaultdict(list)
         for variant in missing_with_search_data:
             variant_id = variant['variant_id']
-            sample_type = variant['family__individual__sample__sample_type']
+            sample_type = variant['family__individual__active_datasets__sample_type']
             missing_by_sample_type[sample_type].append(f"{variant_id} - {'; '.join(variant['family_ids'])}" )
 
         for sample_type, variants in missing_by_sample_type.items():

clickhouse_search/management/tests/set_saved_variant_key_tests.py

@@ -1,7 +1,7 @@
 from django.core.management import call_command
 import mock
 
-from seqr.models import Project, Sample, SavedVariant
+from seqr.models import Project, Dataset, SavedVariant
 from seqr.views.utils.test_utils import AnvilAuthenticationTestCase
 
 MOCK_GCNV_DATA = [
@@ -18,7 +18,10 @@ class SetSavedVariantKeyTest(AnvilAuthenticationTestCase):
     @classmethod
     def setUpTestData(cls):
         Project.objects.filter(id=3).update(genome_version='38')
-        Sample.objects.filter(guid='S000154_na20889').update(dataset_type='SV', is_active=True)
+        dataset = Dataset.objects.get(guid='S000154_na20889')
+        dataset.dataset_type = 'SV'
+        dataset.save()
+        dataset.active_individuals.set({17})
         for sv in SavedVariant.objects.filter(key__isnull=False):
             sv.saved_variant_json = {
                 'genotypes': sv.genotypes, 'populations': {'gnomad': {'af': 0.01}},

clickhouse_search/models/gt_stats_models.py

@@ -5,7 +5,7 @@
 from clickhouse_search.backend.table_models import RefreshableMaterializedView, RefreshableMaterializedViewMeta, \
     IncrementalMaterializedView, Dictionary
 from reference_data.models import GENOME_VERSION_GRCh38, GENOME_VERSION_GRCh37
-from seqr.models import Sample
+from seqr.models import Dataset
 
 
 class BaseProjectGtStats(models.ClickhouseModel):
@@ -308,10 +308,10 @@ class Meta(GtStatsDictMeta):
         layout = 'FLAT(MAX_ARRAY_SIZE 5000000)'
 
 PROJECT_GT_STATS_VIEW_CLASS_MAP = {
-    GENOME_VERSION_GRCh37: {Sample.DATASET_TYPE_VARIANT_CALLS: ProjectsToGtStatsGRCh37SnvIndel},
+    GENOME_VERSION_GRCh37: {Dataset.DATASET_TYPE_VARIANT_CALLS: ProjectsToGtStatsGRCh37SnvIndel},
     GENOME_VERSION_GRCh38: {
-        Sample.DATASET_TYPE_VARIANT_CALLS: ProjectsToGtStatsSnvIndel,
-        Sample.DATASET_TYPE_MITO_CALLS: ProjectsToGtStatsMito,
-        Sample.DATASET_TYPE_SV_CALLS: ProjectsToGtStatsSv,
+        Dataset.DATASET_TYPE_VARIANT_CALLS: ProjectsToGtStatsSnvIndel,
+        Dataset.DATASET_TYPE_MITO_CALLS: ProjectsToGtStatsMito,
+        Dataset.DATASET_TYPE_SV_CALLS: ProjectsToGtStatsSv,
     },
 }

clickhouse_search/models/search_models.py

@@ -10,7 +10,7 @@
 from clickhouse_search.models.reference_data_models import GnomadNonCodingConstraintDict, BaseSpliceAi, \
     ScreenDict
 from reference_data.models import GENOME_VERSION_GRCh38, GENOME_VERSION_GRCh37
-from seqr.models import Sample
+from seqr.models import Dataset
 from seqr.utils.xpos_utils import CHROMOSOME_CHOICES
 from settings import CLICKHOUSE_IN_MEMORY_DIR, CLICKHOUSE_DATA_DIR
 
@@ -414,7 +414,7 @@ class Meta(BaseEntries.Meta):
         )
 
 class EntriesSv(BaseEntries):
-    SAMPLE_TYPE = Sample.SAMPLE_TYPE_WGS
+    SAMPLE_TYPE = Dataset.SAMPLE_TYPE_WGS
     CALL_FIELDS = [
         ('sampleId', models.StringField()),
         ('gt', models.Enum8Field(null=True, blank=True, choices=[(0, 'REF'), (1, 'HET'), (2, 'HOM')])),
@@ -436,7 +436,7 @@ class Meta(BaseEntries.Meta):
         db_table = 'GRCh38/SV/entries'
 
 class EntriesGcnv(BaseEntries):
-    SAMPLE_TYPE = Sample.SAMPLE_TYPE_WES
+    SAMPLE_TYPE = Dataset.SAMPLE_TYPE_WES
     CALL_FIELDS = [
         ('sampleId', models.StringField()),
         ('gt', models.Enum8Field(null=True, blank=True, choices=[(0, 'REF'), (1, 'HET'), (2, 'HOM')])),
@@ -621,21 +621,21 @@ class Meta:
 
 
 ENTRY_CLASS_MAP = {
-    GENOME_VERSION_GRCh37: {Sample.DATASET_TYPE_VARIANT_CALLS: EntriesGRCh37SnvIndel},
+    GENOME_VERSION_GRCh37: {Dataset.DATASET_TYPE_VARIANT_CALLS: EntriesGRCh37SnvIndel},
     GENOME_VERSION_GRCh38: {
-        Sample.DATASET_TYPE_VARIANT_CALLS: EntriesSnvIndel,
-        Sample.DATASET_TYPE_MITO_CALLS: EntriesMito,
-        f'{Sample.DATASET_TYPE_SV_CALLS}_{Sample.SAMPLE_TYPE_WGS}': EntriesSv,
-        f'{Sample.DATASET_TYPE_SV_CALLS}_{Sample.SAMPLE_TYPE_WES}': EntriesGcnv,
+        Dataset.DATASET_TYPE_VARIANT_CALLS: EntriesSnvIndel,
+        Dataset.DATASET_TYPE_MITO_CALLS: EntriesMito,
+        f'{Dataset.DATASET_TYPE_SV_CALLS}_{Dataset.SAMPLE_TYPE_WGS}': EntriesSv,
+        f'{Dataset.DATASET_TYPE_SV_CALLS}_{Dataset.SAMPLE_TYPE_WES}': EntriesGcnv,
     },
 }
 VARIANTS_CLASS_MAP = {
-    GENOME_VERSION_GRCh37: {Sample.DATASET_TYPE_VARIANT_CALLS: VariantsGRCh37SnvIndel},
+    GENOME_VERSION_GRCh37: {Dataset.DATASET_TYPE_VARIANT_CALLS: VariantsGRCh37SnvIndel},
     GENOME_VERSION_GRCh38: {
-        Sample.DATASET_TYPE_VARIANT_CALLS: VariantsSnvIndel,
-        Sample.DATASET_TYPE_MITO_CALLS: VariantsMito,
-        f'{Sample.DATASET_TYPE_SV_CALLS}_{Sample.SAMPLE_TYPE_WGS}': VariantsSv,
-        f'{Sample.DATASET_TYPE_SV_CALLS}_{Sample.SAMPLE_TYPE_WES}': VariantsGcnv,
+        Dataset.DATASET_TYPE_VARIANT_CALLS: VariantsSnvIndel,
+        Dataset.DATASET_TYPE_MITO_CALLS: VariantsMito,
+        f'{Dataset.DATASET_TYPE_SV_CALLS}_{Dataset.SAMPLE_TYPE_WGS}': VariantsSv,
+        f'{Dataset.DATASET_TYPE_SV_CALLS}_{Dataset.SAMPLE_TYPE_WES}': VariantsGcnv,
     },
 }
 VARIANT_DETAILS_CLASS_MAP = {