diff --git a/hippo/designdb/components/compound.py b/hippo/designdb/components/compound.py index b2a0be0..d4b0893 100644 --- a/hippo/designdb/components/compound.py +++ b/hippo/designdb/components/compound.py @@ -18,6 +18,7 @@ ReactantModel, ReactionModel, ScaffoldModel, + TargetModel, ) from django.db.models import Exists, OuterRef from IPython.display import display @@ -376,17 +377,24 @@ def add_stock( return quote return quote.pk - def get_tags(self) -> list[str]: - """Get the tags assigned to this compound""" - return list(self._instance.tags.values_list('compound_tag_name', flat=True)) + def get_tags(self, target: TargetModel | None = None) -> list[str]: + """Get the tags assigned to this compound - def add_tag(self, tag: str) -> None: + :param target: Optionally restrict to tags assigned under this + :class:`.TargetModel`, defaults to ``None`` (tags across all targets) + """ + junctions = CompoundTagJunctionModel.objects.filter(compound=self._instance) + if target is not None: + junctions = junctions.filter(target=target) + return list(junctions.values_list('compound_tag__compound_tag_name', flat=True)) + + def add_tag(self, tag: str, *, target: TargetModel) -> None: """Add a tag to this compound""" assert isinstance(tag, str) tag_obj, _ = CompoundTagModel.objects.get_or_create(compound_tag_name=tag) CompoundTagJunctionModel.objects.get_or_create( - compound=self._instance, compound_tag=tag_obj + compound=self._instance, compound_tag=tag_obj, target=target ) self._tags = None # invalidate cache diff --git a/hippo/designdb/models.py b/hippo/designdb/models.py index 8d94fdc..281d439 100644 --- a/hippo/designdb/models.py +++ b/hippo/designdb/models.py @@ -641,12 +641,17 @@ class Meta(BaseModel.Meta): class CompoundTagJunctionModel(BaseModel): - pk = models.CompositePrimaryKey('compound_id', 'compound_tag_id') + pk = models.CompositePrimaryKey('compound_id', 'compound_tag_id', 'target_id') compound = models.ForeignKey( CompoundModel, on_delete=models.CASCADE, db_column='compound_id', ) + target = models.ForeignKey( + TargetModel, + on_delete=models.CASCADE, + db_column='target_id', + ) compound_tag = models.ForeignKey( CompoundTagModel, on_delete=models.CASCADE, @@ -691,12 +696,17 @@ class Meta(BaseModel.Meta): class CompoundEnumerationMethodJunctionModel(BaseModel): - pk = models.CompositePrimaryKey('compound_id', 'enumeration_method_id') + pk = models.CompositePrimaryKey('compound_id', 'enumeration_method_id', 'target_id') compound = models.ForeignKey( CompoundModel, on_delete=models.CASCADE, db_column='compound_id', ) + target = models.ForeignKey( + TargetModel, + on_delete=models.CASCADE, + db_column='target_id', + ) enumeration_method = models.ForeignKey( EnumerationMethodModel, on_delete=models.CASCADE, diff --git a/hippo/designdb/services/ingestion.py b/hippo/designdb/services/ingestion.py index c6263f8..09eb260 100644 --- a/hippo/designdb/services/ingestion.py +++ b/hippo/designdb/services/ingestion.py @@ -10,7 +10,9 @@ import pandas as pd from designdb.components.compound import Ingredient from designdb.models import ( + CompoundEnumerationMethodJunctionModel, CompoundModel, + CompoundTagJunctionModel, EnumerationMethodModel, PoseMethodModel, PoseModel, @@ -371,7 +373,12 @@ def ingest_filesystem( smiles=smiles, # inchikey=sane_inchikey, ) - compound.tags.add(*compound_tags) + for compound_tag in compound_tags: + CompoundTagJunctionModel.objects.get_or_create( + compound=compound, + compound_tag=compound_tag, + target=target, + ) if compound_created: result.compounds_created += 1 @@ -514,9 +521,18 @@ def ingest_sdf( # smiles=sane_smiles, # inchikey=sane_inchikey, ) - compound.tags.add(*compound_tags) + for compound_tag in compound_tags: + CompoundTagJunctionModel.objects.get_or_create( + compound=compound, + compound_tag=compound_tag, + target=target, + ) if enumeration_method_obj is not None: - compound.enumeration_methods.add(enumeration_method_obj) + CompoundEnumerationMethodJunctionModel.objects.get_or_create( + compound=compound, + enumeration_method=enumeration_method_obj, + target=target, + ) if compound_created: result.compounds_created += 1 @@ -1080,7 +1096,12 @@ def ingest_syndirella_elabs( products = CompoundModel.objects.filter(pk__in=product_ids) product_tags = CompoundTagService.tags_from_list(product_tag_list) for compound in products: - compound.tags.add(*product_tags) + for compound_tag in product_tags: + CompoundTagJunctionModel.objects.get_or_create( + compound=compound, + compound_tag=compound_tag, + target=target, + ) # bulk register scaffold relationships diff --git a/hippo/designdb/sets/compound.py b/hippo/designdb/sets/compound.py index eee7553..ce4b1d2 100644 --- a/hippo/designdb/sets/compound.py +++ b/hippo/designdb/sets/compound.py @@ -19,6 +19,7 @@ ReactionModel, RouteModel, ScaffoldModel, + TargetModel, ) from django.db.models import Count, Exists, OuterRef, Q, QuerySet from django.db.models.query import ModelIterable @@ -1287,6 +1288,8 @@ def write_CAR_csv( def add_tag( self, tag: str, + *, + target: TargetModel, ) -> None: """Add this tag to every member of the set""" @@ -1296,7 +1299,9 @@ def add_tag( CompoundTagJunctionModel.objects.bulk_create( [ - CompoundTagJunctionModel(compound=compound, compound_tag=compound_tag) + CompoundTagJunctionModel( + compound=compound, compound_tag=compound_tag, target=target + ) for compound in self._queryset ], ignore_conflicts=True, @@ -1503,14 +1508,25 @@ def inchikeys(self) -> list[str]: """Returns the inchikeys of compounds in this set""" return list(self._queryset.values_list('compound_inchikey', flat=True)) + def get_tags(self, target: TargetModel | None = None) -> set[str]: + """Returns the set of unique tags present in this compound set + + :param target: Optionally restrict to tags assigned under this + :class:`.TargetModel`, defaults to ``None`` (tags across all targets) + """ + junctions = CompoundTagJunctionModel.objects.filter(compound_id__in=self.ids) + if target is not None: + junctions = junctions.filter(target=target) + return set( + junctions.values_list( + 'compound_tag__compound_tag_name', flat=True + ).distinct() + ) + @property def tags(self) -> set[str]: """Returns the set of unique tags present in this compound set""" - return set( - CompoundTagJunctionModel.objects.filter(compound_id__in=self.ids) - .values_list('compound_tag__compound_tag_name', flat=True) - .distinct() - ) + return self.get_tags() @property def num_poses(self) -> int: diff --git a/images/xchem-designdb/init-db/01_schema.sql b/images/xchem-designdb/init-db/01_schema.sql index d32048b..964a23a 100644 --- a/images/xchem-designdb/init-db/01_schema.sql +++ b/images/xchem-designdb/init-db/01_schema.sql @@ -206,10 +206,11 @@ CREATE TABLE IF NOT EXISTS designdb.compound_tags ( -- New table CREATE TABLE IF NOT EXISTS designdb.has_compound_tags ( compound_id BIGINT NOT NULL REFERENCES designdb.compounds (id) ON DELETE CASCADE, + target_id BIGINT NOT NULL REFERENCES designdb.targets (id) ON DELETE CASCADE, compound_tag_id BIGINT NOT NULL REFERENCES designdb.compound_tags (id) ON DELETE CASCADE, created_on TIMESTAMPTZ DEFAULT now(), updated_on TIMESTAMPTZ DEFAULT now(), - PRIMARY KEY (compound_id, compound_tag_id) + PRIMARY KEY (compound_id, compound_tag_id, target_id) ); -- New table @@ -229,10 +230,11 @@ CREATE TABLE IF NOT EXISTS designdb.enumeration_methods ( -- New table CREATE TABLE IF NOT EXISTS designdb.has_enumeration_methods ( compound_id BIGINT NOT NULL REFERENCES designdb.compounds (id) ON DELETE CASCADE, + target_id BIGINT NOT NULL REFERENCES designdb.targets (id) ON DELETE CASCADE, enumeration_method_id BIGINT NOT NULL REFERENCES designdb.enumeration_methods (id) ON DELETE CASCADE, created_on TIMESTAMPTZ DEFAULT now(), updated_on TIMESTAMPTZ DEFAULT now(), - PRIMARY KEY (compound_id, enumeration_method_id) + PRIMARY KEY (compound_id, enumeration_method_id, target_id) ); -- New table diff --git a/uv.lock b/uv.lock index 2c48a96..6759362 100644 --- a/uv.lock +++ b/uv.lock @@ -4279,7 +4279,7 @@ wheels = [ [[package]] name = "xchem-hippo" version = "1.0.0" -source = { virtual = "." } +source = { editable = "." } dependencies = [ { name = "apsw" }, { name = "chardet" },