import base64 import hashlib import hmac import json import logging import uuid from abc import ABC, abstractmethod from collections import defaultdict from datetime import datetime from typing import Any, Dict, List, Optional, TYPE_CHECKING, Type from django.conf import settings as s from envipy_additional_information import EnviPyModel, UIConfig from ninja import Schema from pydantic import HttpUrl, ValidationError from epdb.exceptions import PackageImportException from epdb.models import ( AdditionalInformation, Compound, CompoundStructure, Edge, License, Node, ParallelRule, Pathway, Reaction, Rule, Scenario, Setting, SimpleAmbitRule, ) from utilities.chem import FormatConverter logger = logging.getLogger(__name__) Package = s.GET_PACKAGE_MODEL() if TYPE_CHECKING: from epdb.logic import SPathway class LicenseExportSchema(Schema): cc_string: str link: HttpUrl image_link: HttpUrl ############## # RefSchemas # ############## class RefExportSchema(Schema): uuid: str url: str @staticmethod def resolve_uuid(obj): value = obj.get("uuid") if isinstance(obj, dict) else obj.uuid return str(value) class RefCompoundExportSchema(RefExportSchema): ... class RefCompoundStructureExportSchema(RefExportSchema): ... class RefReactionExportSchema(RefExportSchema): ... class RefRuleExportSchema(RefExportSchema): ... class RefNodeExportSchema(RefExportSchema): ... class RefEdgeExportSchema(RefExportSchema): ... class RefPathwayExportSchema(RefExportSchema): ... class RefScenarioExportSchema(RefExportSchema): ... class RefEnzymeExportSchema(RefExportSchema): ... ############ # Compound # ############ class CompoundExportSchema(RefCompoundExportSchema): name: str description: str aliases: List[str] default_structure: RefCompoundStructureExportSchema structures: List["CompoundStructureExportSchema"] scenarios: List[RefScenarioExportSchema] class CompoundStructureExportSchema(RefCompoundStructureExportSchema): name: str description: str aliases: List[str] smiles: str molfile: Optional[str] normalized_structure: bool scenarios: List[RefScenarioExportSchema] ############ # Reaction # ############ class ReactionExportSchema(RefReactionExportSchema): name: str description: str aliases: List[str] educts: List[RefCompoundStructureExportSchema] products: List[RefCompoundStructureExportSchema] rules: List[RefRuleExportSchema] multi_step: bool medline_references: List[str] | None scenarios: List[RefScenarioExportSchema] ######### # Rules # ######### class EnzymeExportSchema(RefEnzymeExportSchema): ec_number: str classification_level: int linking_method: str reaction_evidence: List[RefReactionExportSchema] edge_evidence: List[RefEdgeExportSchema] class EnzymeRuleExportSchema(RefRuleExportSchema): enzymes: List[EnzymeExportSchema] = [] @staticmethod def resolve_enzymes(obj): if isinstance(obj, dict): res = [] for e in obj.get("enzymes", []): res.append(EnzymeExportSchema.model_validate(e)) return res return obj.enzymelink_set.all() class RuleExportSchema(EnzymeRuleExportSchema): name: str description: str aliases: List[str] smirks: str reactant_filter_smarts: Optional[str] product_filter_smarts: Optional[str] scenarios: List[RefScenarioExportSchema] class ParallelRuleExportSchema(EnzymeRuleExportSchema): name: str description: str aliases: List[str] simple_rules: List[RefRuleExportSchema] scenarios: List[RefScenarioExportSchema] ########################### # Pathway / Nodes / Edges # ########################### class NodeExportSchema(RefNodeExportSchema): name: str description: str aliases: List[str] default_node_label: RefCompoundStructureExportSchema node_labels: List[RefCompoundStructureExportSchema] depth: int stereo_removed: bool scenarios: List[RefScenarioExportSchema] class EdgeExportSchema(RefEdgeExportSchema): name: str description: str aliases: List[str] edge_label: RefReactionExportSchema start_nodes: List[RefNodeExportSchema] end_nodes: List[RefNodeExportSchema] scenarios: List[RefScenarioExportSchema] class PathwayExportSchema(RefPathwayExportSchema): name: str description: str aliases: List[str] predicted: bool nodes: List[NodeExportSchema] edges: List[EdgeExportSchema] scenarios: List[RefScenarioExportSchema] class ScenarioExportSchema(RefScenarioExportSchema): name: str description: str scenario_date: str scenario_type: str class AdditionalInformationExportSchema(RefExportSchema): type: str data: dict scenario: RefScenarioExportSchema | None = None attach_object: RefExportSchema | None = None @staticmethod def resolve_attach_object(obj): if isinstance(obj, dict): if obj.get("attach_object") is None: return None return RefExportSchema.model_validate(obj["attach_object"]) return obj.content_object ########### # Package # ########### class PackageExportSchema(Schema): name: str description: str uuid: str reviewed: bool license: LicenseExportSchema | None = None compounds: List[CompoundExportSchema] reactions: List[ReactionExportSchema] simple_rules: List[RuleExportSchema] composite_rules: List[ParallelRuleExportSchema] pathways: List[PathwayExportSchema] scenarios: List[ScenarioExportSchema] additional_information: List[AdditionalInformationExportSchema] @staticmethod def resolve_uuid(obj): value = obj.get("uuid") if isinstance(obj, dict) else obj.uuid return str(value) @staticmethod def resolve_compounds(obj): res = [] if isinstance(obj, dict): for c in obj.get("compounds", []): res.append(CompoundExportSchema.model_validate(c)) return res return obj.compounds.all() @staticmethod def resolve_simple_rules(obj): if isinstance(obj, dict): result = [] for r in obj.get("simple_rules", []): result.append(RuleExportSchema.model_validate(r)) return result return SimpleAmbitRule.objects.filter(package=obj) @staticmethod def resolve_composite_rules(obj): if isinstance(obj, dict): result = [] for r in obj.get("composite_rules", []): result.append(ParallelRuleExportSchema.model_validate(r)) return result return ParallelRule.objects.filter(package=obj) @staticmethod def resolve_additional_information(obj): if isinstance(obj, dict): result = [] for ai in obj.get("additional_information", []): result.append(AdditionalInformationExportSchema.model_validate(ai)) return result return AdditionalInformation.objects.filter(package=obj) class Exporter(ABC): def __init__(self, package: Package): self._raw_package = package def do_export(self): return self._export() @abstractmethod def _export(self): pass class PackageExporter(Exporter): def _export(self) -> Dict[str, Any]: """ Dumps a Package and all its related objects as JSON. Args: package: The Package instance to dump Returns: Dict containing the complete package data as JSON-serializable structure """ data = PackageExportSchema.from_orm(self._raw_package) return data.model_dump(mode="json") class PathwayExporter(Exporter): def __init__(self, package: Package, add_infs_to_export: List[str] = []): super().__init__(package) self._add_infs_to_export = add_infs_to_export def _flatten_additional_information(self, ai: AdditionalInformation) -> dict[str, Any]: model_cls: Type[EnviPyModel] = type(ai.get()) def _flatten(d: dict, parent_key: str = "") -> dict[str, Any]: items: dict[str, Any] = {} for key, value in d.items(): new_key = f"{parent_key}__{key.lower()}" if parent_key else key.lower() if isinstance(value, dict): items.update(_flatten(value, new_key)) else: items[new_key] = value return items flat = _flatten(ai.data, ai.type) ui_class = getattr(model_cls, "UI", None) if ui_class is None: return flat for f in model_cls.model_fields: ui_info = getattr(ui_class, f, None) if not isinstance(ui_info, UIConfig) or ui_info.unit is None: continue flat[f"{model_cls.__name__}__{f}__unit"] = ui_info.unit return flat def _export(self): from io import StringIO from csv import DictWriter rows = [] for pw in self._raw_package.pathways.all(): for n in pw.nodes: for scen in pw.scenarios.all(): row = { "pathway_name": pw.name, "pathway_id": str(pw.url), "node_depth": n.depth, "compound_id": str(n.default_node_label.compound.url), "pubchem_ID": n.default_node_label.pubchem_compound_id, "compound_name": n.default_node_label.compound.name, "compound_smiles": n.default_node_label.smiles, "scenario_id": str(scen.url), "scenario_name": scen.name, "scenario_type": scen.scenario_type, "scenario_description": scen.description, } if self._add_infs_to_export: ai_qs = AdditionalInformation.objects.filter( scenario=scen, type__in=self._add_infs_to_export ) else: ai_qs = AdditionalInformation.objects.filter(scenario=scen) for ai in ai_qs: if ai.type == "ProposedIntermediate" and ai.content_object == n: row.update({"proposed_intermediate": True}) elif ai.type == "SpikeCompound": spike = {"SpikeCompound__url": ai.get().url} try: struc = CompoundStructure.objects.get( compound__package=self._raw_package, url=ai.get().url ) spike["SpikeCompound__smiles"] = struc.smiles except Exception: spike["SpikeCompound__smiles"] = None row.update(**spike) else: row.update(self._flatten_additional_information(ai)) rows.append(row) # Get all header fields all_header_fields = set() for row in rows: all_header_fields.update(row.keys()) # Per request the CSV should start with these fields header = [ "pathway_name", "pathway_id", "node_depth", "compound_id", "pubchem_ID", "compound_name", "compound_smiles", "scenario_id", "scenario_name", "scenario_type", "scenario_description", ] # User remaining fields and place them after the predefined values in a sorted manner remainder = sorted(list(all_header_fields.difference(set(header)))) header.extend(remainder) buffer = StringIO() writer = DictWriter(buffer, fieldnames=header, delimiter="\t") writer.writeheader() writer.writerows(rows) buffer.seek(0) return buffer.getvalue() class PackageImporter: def __init__(self, package: Dict[str, Any], preserve_uuids: bool = False): self.preserve_uuids = preserve_uuids self._raw_package = package self._cache = {} def do_import(self) -> Package: return self._import_package_from_json() def _import_compounds( self, package: Package, compounds: List[CompoundExportSchema], ): for c in compounds: new_c = Compound() new_c.uuid = str(uuid.uuid4()) if not self.preserve_uuids else c.uuid new_c.package = package new_c.name = c.name new_c.description = c.description new_c.aliases = c.aliases new_c.save() self._cache[c.uuid] = new_c for cs in c.structures: new_cs = CompoundStructure() new_cs.uuid = str(uuid.uuid4()) if not self.preserve_uuids else cs.uuid new_cs.compound = new_c new_cs.name = cs.name new_cs.description = cs.description new_cs.aliases = cs.aliases new_cs.smiles = cs.smiles new_cs.molfile = cs.molfile new_cs.normalized_structure = cs.normalized_structure new_cs.save() self._cache[cs.uuid] = new_cs # Now we can assigne the default_structure to the compound new_c.default_structure = self._cache[c.default_structure.uuid] new_c.save() def _import_simple_rules(self, package: Package, rules: List[RuleExportSchema]): for r in rules: new_r = SimpleAmbitRule() new_r.uuid = str(uuid.uuid4()) if not self.preserve_uuids else r.uuid new_r.package = package new_r.name = r.name new_r.description = r.description new_r.aliases = r.aliases new_r.smirks = r.smirks new_r.reactant_filter_smarts = r.reactant_filter_smarts new_r.product_filter_smarts = r.product_filter_smarts new_r.save() self._cache[r.uuid] = new_r def _import_composite_rules(self, package: Package, rules: List[ParallelRuleExportSchema]): for r in rules: new_r = ParallelRule() new_r.uuid = str(uuid.uuid4()) if not self.preserve_uuids else r.uuid new_r.package = package new_r.name = r.name new_r.description = r.description new_r.aliases = r.aliases new_r.save() for sr in r.simple_rules: new_r.simple_rules.add(self._cache[sr.uuid]) self._cache[r.uuid] = new_r def _import_reactions( self, package: Package, reactions: List[ReactionExportSchema], ): for r in reactions: new_r = Reaction() new_r.uuid = str(uuid.uuid4()) if not self.preserve_uuids else r.uuid new_r.package = package new_r.name = r.name new_r.description = r.description new_r.aliases = r.aliases new_r.multi_step = r.multi_step new_r.medline_references = r.medline_references new_r.save() for educt in r.educts: new_r.educts.add(self._cache[educt.uuid]) for product in r.products: new_r.products.add(self._cache[product.uuid]) for rule in r.rules: new_r.rules.add(self._cache[rule.uuid]) self._cache[r.uuid] = new_r def _import_pathways(self, package: Package, pathways: List[PathwayExportSchema]): for pw in pathways: new_pw = Pathway() new_pw.uuid = str(uuid.uuid4()) if not self.preserve_uuids else pw.uuid new_pw.package = package new_pw.name = pw.name new_pw.description = pw.description new_pw.aliases = pw.aliases new_pw.predicted = pw.predicted new_pw.save() self._cache[pw.uuid] = new_pw for n in pw.nodes: new_n = Node() new_n.uuid = str(uuid.uuid4()) if not self.preserve_uuids else n.uuid new_n.pathway = new_pw new_n.name = n.name new_n.description = n.description new_n.aliases = n.aliases new_n.default_node_label = self._cache[n.default_node_label.uuid] new_n.depth = n.depth new_n.stereo_removed = n.stereo_removed new_n.save() for nl in n.node_labels: new_n.node_labels.add(self._cache[nl.uuid]) self._cache[n.uuid] = new_n for e in pw.edges: new_e = Edge() new_e.uuid = str(uuid.uuid4()) if not self.preserve_uuids else e.uuid new_e.pathway = new_pw new_e.name = e.name new_e.description = e.description new_e.aliases = e.aliases new_e.edge_label = self._cache[e.edge_label.uuid] new_e.save() for sn in e.start_nodes: new_e.start_nodes.add(self._cache[sn.uuid]) for en in e.end_nodes: new_e.end_nodes.add(self._cache[en.uuid]) self._cache[e.uuid] = new_e def _import_scenarios(self, package: Package, scenarios: List[ScenarioExportSchema]): for scen in scenarios: new_s = Scenario() new_s.uuid = str(uuid.uuid4()) if not self.preserve_uuids else scen.uuid new_s.package = package new_s.name = scen.name new_s.description = scen.description new_s.scenario_date = scen.scenario_date new_s.scenario_type = scen.scenario_type new_s.save() self._cache[scen.uuid] = new_s def _import_additional_information( self, package: Package, additional_information: List[AdditionalInformationExportSchema] ): for ai in additional_information: new_ai = AdditionalInformation() new_ai.uuid = str(uuid.uuid4()) if not self.preserve_uuids else ai.uuid new_ai.package = package new_ai.type = ai.type new_ai.data = ai.data if ai.scenario: new_ai.scenario = self._cache[ai.scenario.uuid] if ai.attach_object: new_ai.content_object = self._cache[ai.attach_object.uuid] new_ai.save() self._cache[ai.uuid] = new_ai def _link_scenarios_after_import(self, data: PackageExportSchema): collections = [ data.compounds, data.simple_rules, data.composite_rules, data.reactions, data.pathways, [n for pw in data.pathways for n in pw.nodes], [e for pw in data.pathways for e in pw.edges], ] for coll in collections: for elem in coll: elem_obj = self._cache[elem.uuid] for scen in elem.scenarios: elem_obj.scenarios.add(self._cache[scen.uuid]) def _import_package_from_json( self, ) -> Package: try: parsed = PackageExportSchema.model_validate(self._raw_package) except ValidationError as e: logger.error(f"Error validating package data: {e}") raise PackageImportException("Deserialization failed") package_uuid = str(uuid.uuid4()) if self.preserve_uuids: package_uuid = parsed.uuid package_license = License.objects.get(link=parsed.license.link) if parsed.license else None package = Package() package.uuid = package_uuid package.name = f"{parsed.name} - Imported at {datetime.now()}" package.description = parsed.description package.reviewed = False package.license = package_license package.save() # Import Elements self._import_compounds(package, parsed.compounds) self._import_simple_rules(package, parsed.simple_rules) self._import_composite_rules(package, parsed.composite_rules) self._import_reactions(package, parsed.reactions) self._import_pathways(package, parsed.pathways) self._import_scenarios(package, parsed.scenarios) self._import_additional_information(package, parsed.additional_information) self._link_scenarios_after_import(parsed) return package @staticmethod def sign(data: Dict[str, Any], key: str) -> Dict[str, Any]: json_str = json.dumps(data, sort_keys=True, separators=(",", ":")) signature = hmac.new(key.encode(), json_str.encode(), hashlib.sha256).digest() data["_signature"] = base64.b64encode(signature).decode() return data @staticmethod def verify(data: Dict[str, Any], key: str) -> bool: copied_data = data.copy() sig = copied_data.pop("_signature") signature = base64.b64decode(sig, validate=True) json_str = json.dumps(copied_data, sort_keys=True, separators=(",", ":")) expected = hmac.new(key.encode(), json_str.encode(), hashlib.sha256).digest() return hmac.compare_digest(signature, expected) class PathwayUtils: def __init__(self, pathway: "Pathway"): self.pathway = pathway @staticmethod def _get_products(smiles: str, rules: List["Rule"]): educt_rule_products: Dict[str, Dict[str, List[str]]] = defaultdict( lambda: defaultdict(list) ) for r in rules: product_sets = r.apply(smiles) for product_set in product_sets: for product in product_set: educt_rule_products[smiles][r.url].append(product) return educt_rule_products def find_missing_rules(self, rules: List["Rule"]): print(f"Processing {self.pathway.name}") # compute products for each node / rule combination in the pathway educt_rule_products = defaultdict(lambda: defaultdict(list)) for node in self.pathway.nodes: educt_rule_products.update(**self._get_products(node.default_node_label.smiles, rules)) # loop through edges and determine reactions that can't be constructed by # any of the rules or a combination of two rules in a chained fashion res: Dict[str, List["Rule"]] = dict() for edge in self.pathway.edges: found = False reaction = edge.edge_label educts = [cs for cs in reaction.educts.all()] products = [cs.smiles for cs in reaction.products.all()] rule_chain = [] for educt in educts: educt = educt.smiles triggered_rules = list(educt_rule_products.get(educt, {}).keys()) for triggered_rule in triggered_rules: if rule_products := educt_rule_products[educt][triggered_rule]: # check if this rule covers the reaction if FormatConverter.smiles_covered_by( products, rule_products, standardize=True, canonicalize_tautomers=True ): found = True else: # Check if another prediction step would cover the reaction for product in rule_products: prod_rule_products = self._get_products(product, rules) prod_triggered_rules = list( prod_rule_products.get(product, {}).keys() ) for prod_triggered_rule in prod_triggered_rules: if second_step_products := prod_rule_products[product][ prod_triggered_rule ]: if FormatConverter.smiles_covered_by( products, second_step_products, standardize=True, canonicalize_tautomers=True, ): rule_chain.append( ( triggered_rule, Rule.objects.get(url=triggered_rule).name, ) ) rule_chain.append( ( prod_triggered_rule, Rule.objects.get(url=prod_triggered_rule).name, ) ) res[edge.url] = rule_chain if not found: res[edge.url] = rule_chain return res def engineer(self, setting: "Setting"): from epdb.logic import SPathway from utilities.chem import FormatConverter from utilities.ml import graph_from_pathway, get_shortest_path # get a fresh copy pw = Pathway.objects.get(id=self.pathway.pk) root_nodes = [n.default_node_label.smiles for n in pw.root_nodes] if len(root_nodes) != 1: logger.warning(f"Pathway {pw.name} has {len(root_nodes)} root nodes") # spw, mapping, intermediates return None, {}, [] # Predict the Pathway in memory spw = SPathway(root_nodes[0], None, setting) level = 0 while not spw.done: spw.predict_step(from_depth=level) level += 1 # Generate SNode -> Node mapping node_mapping = {} for node in pw.nodes: for snode in spw.smiles_to_node.values(): data_smiles = node.default_node_label.smiles pred_smiles = snode.smiles # "~" denotes any bond remove and use implicit single bond for comparison data_key = FormatConverter.InChIKey(data_smiles.replace("~", "")) pred_key = FormatConverter.InChIKey(pred_smiles.replace("~", "")) if data_key == pred_key: node_mapping[snode] = node reverse_mapping = {v: k for k, v in node_mapping.items()} graph = graph_from_pathway(spw) intermediate_mapping = [] # loop through each edge and each reactant <-> product pair # and compute the shortest path on the predicted pathway for e in pw.edges: for start in e.start_nodes.all(): if start not in reverse_mapping: continue start_snode = reverse_mapping[start] for end in e.end_nodes.all(): if end not in reverse_mapping: continue end_snode = reverse_mapping[end] # If res is non-empty, we've found intermediates intermediate_smiles = get_shortest_path( graph, FormatConverter.standardize(start_snode.smiles, remove_stereo=True), FormatConverter.standardize(end_snode.smiles, remove_stereo=True), ) if intermediate_smiles: intermediates = [] prev = start_snode.smiles for smi in intermediate_smiles + [end_snode.smiles]: for e in spw.get_edge_for_educt_smiles(prev): if smi in e.product_smiles(): intermediates.append(e) prev = smi intermediate_mapping.append( (start, end, start_snode, end_snode, intermediates) ) return spw, reverse_mapping, intermediate_mapping @staticmethod def spathway_to_pathway( package: "Package", spw: "SPathway", name: str = None, description: str = None ): snode_to_node_mapping = dict() root_nodes = spw.root_nodes pw = Pathway.create( package=package, smiles=root_nodes[0].smiles, name=name, description=description, predicted=True, ) pw.setting = spw.prediction_setting pw.save() snode_to_node_mapping[root_nodes[0]] = pw.root_nodes[0] if len(root_nodes) > 1: for rn in root_nodes[1:]: n = Node.create(pw, rn.smiles, depth=0) snode_to_node_mapping[rn] = n for snode, node in snode_to_node_mapping.items(): spw.snode_persist_lookup[snode] = node spw.persist = pw spw._sync_to_pathway() return pw