forked from enviPath/enviPy
1088 lines
38 KiB
Python
1088 lines
38 KiB
Python
import base64
|
|
import hashlib
|
|
import hmac
|
|
import json
|
|
import logging
|
|
import uuid
|
|
from collections import defaultdict
|
|
from datetime import datetime
|
|
from typing import Any, Dict, List, Optional, TYPE_CHECKING
|
|
|
|
from django.conf import settings as s
|
|
from django.db import transaction
|
|
from ninja import Schema
|
|
from pydantic import HttpUrl
|
|
|
|
from bayer.models import PESCompound
|
|
from epdb.models import (
|
|
AdditionalInformation,
|
|
Compound,
|
|
CompoundStructure,
|
|
Edge,
|
|
EnviFormer,
|
|
EPModel,
|
|
ExternalDatabase,
|
|
ExternalIdentifier,
|
|
License,
|
|
MLRelativeReasoning,
|
|
Node,
|
|
ParallelRule,
|
|
Pathway,
|
|
Reaction,
|
|
Rule,
|
|
RuleBasedRelativeReasoning,
|
|
Scenario,
|
|
Setting,
|
|
SimpleAmbitRule,
|
|
SimpleRDKitRule,
|
|
SimpleRule,
|
|
)
|
|
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 RefGroupExportSchema(RefExportSchema): ...
|
|
|
|
|
|
class RefCompoundExportSchema(RefExportSchema): ...
|
|
|
|
|
|
class RefCompoundStructureExportSchema(RefExportSchema): ...
|
|
|
|
|
|
class RefReactionExportSchema(RefExportSchema): ...
|
|
|
|
|
|
class RefRuleExportSchema(RefExportSchema): ...
|
|
|
|
|
|
class RefNodeExportSchema(RefExportSchema): ...
|
|
|
|
|
|
class RefEdgeExportSchema(RefExportSchema): ...
|
|
|
|
|
|
class RefPathwayExportSchema(RefExportSchema): ...
|
|
|
|
|
|
class RefScenarioExportSchema(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]
|
|
|
|
|
|
class PESCompoundExportSchema(CompoundExportSchema):
|
|
structures: List["PESCompoundStructureExportSchema"]
|
|
|
|
|
|
class PESCompoundStructureExportSchema(CompoundStructureExportSchema):
|
|
pes_link: HttpUrl
|
|
|
|
|
|
############
|
|
# 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 RuleExportSchema(RefRuleExportSchema):
|
|
name: str
|
|
description: str
|
|
aliases: List[str]
|
|
smirks: str
|
|
reactant_filter_smarts: Optional[str]
|
|
product_filter_smarts: Optional[str]
|
|
scenarios: List[RefScenarioExportSchema]
|
|
|
|
|
|
class ParallelRuleExportSchema(RefRuleExportSchema):
|
|
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
|
|
classification_level: str = "Internal"
|
|
data_pool: RefGroupExportSchema | None = None
|
|
compounds: List[CompoundExportSchema | PESCompoundExportSchema]
|
|
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_classification_level(obj):
|
|
if isinstance(obj, str):
|
|
return obj["classification_level"]
|
|
return obj.Classification(obj.classification_level).name
|
|
|
|
|
|
@staticmethod
|
|
def resolve_compounds(obj):
|
|
res = []
|
|
if isinstance(obj, dict):
|
|
for c in obj.get("compounds", []):
|
|
if "pes_link" in c:
|
|
res.append(PESCompoundExportSchema.model_validate(c))
|
|
else:
|
|
res.append(CompoundExportSchema.model_validate(c))
|
|
else:
|
|
for c in obj.compounds.all():
|
|
if isinstance(c, PESCompound):
|
|
res.append(PESCompoundExportSchema.from_orm(c))
|
|
else:
|
|
res.append(CompoundExportSchema.from_orm(c))
|
|
return res
|
|
|
|
@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 PackageExporter:
|
|
def __init__(self, package: Package):
|
|
self._raw_package = package
|
|
|
|
def do_export(self):
|
|
return PackageExporter._export_package_as_json(self._raw_package)
|
|
|
|
@staticmethod
|
|
def _export_package_as_json(package: Package) -> 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(package)
|
|
|
|
return data.model_dump(mode="json")
|
|
|
|
|
|
class PackageImporter:
|
|
"""
|
|
Imports package data from JSON export.
|
|
Handles object creation, relationship mapping, and dependency resolution.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
package: Dict[str, Any],
|
|
preserve_uuids: bool = False,
|
|
add_import_timestamp=True,
|
|
trust_reviewed=False,
|
|
):
|
|
"""
|
|
Initialize the importer.
|
|
|
|
Args:
|
|
preserve_uuids: If True, preserve original UUIDs. If False, generate new ones.
|
|
"""
|
|
self.preserve_uuids = preserve_uuids
|
|
self.add_import_timestamp = add_import_timestamp
|
|
self.trust_reviewed = trust_reviewed
|
|
self.uuid_mapping = {}
|
|
self.object_cache = {}
|
|
self._raw_package = package
|
|
|
|
def _get_or_generate_uuid(self, original_uuid: str) -> str:
|
|
"""Get mapped UUID or generate new one if not preserving UUIDs."""
|
|
if self.preserve_uuids:
|
|
return original_uuid
|
|
|
|
if original_uuid not in self.uuid_mapping:
|
|
self.uuid_mapping[original_uuid] = str(uuid.uuid4())
|
|
|
|
return self.uuid_mapping[original_uuid]
|
|
|
|
def _cache_object(self, model_name: str, uuid_str: str, obj):
|
|
"""Cache a created object for later reference."""
|
|
self.object_cache[(model_name, uuid_str)] = obj
|
|
|
|
def _get_cached_object(self, model_name: str, uuid_str: str):
|
|
"""Get a cached object by model name and UUID."""
|
|
return self.object_cache.get((model_name, uuid_str))
|
|
|
|
def do_import(self) -> Package:
|
|
return self._import_package_from_json(self._raw_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)
|
|
|
|
@transaction.atomic
|
|
def _import_package_from_json(self, package_data: Dict[str, Any]) -> Package:
|
|
"""
|
|
Import a complete package from JSON data.
|
|
|
|
Args:
|
|
package_data: Dictionary containing the package export data
|
|
|
|
Returns:
|
|
The created Package instance
|
|
"""
|
|
print(f"Starting import of package: {package_data['name']}")
|
|
|
|
# Create the main package
|
|
package = self._create_package(package_data)
|
|
|
|
# Import in dependency order
|
|
self._import_compounds(package, package_data.get("compounds", []))
|
|
self._import_structures(package, package_data.get("structures", []))
|
|
self._import_rules(package, package_data.get("rules", {}))
|
|
self._import_reactions(package, package_data.get("reactions", []))
|
|
self._import_pathways(package, package_data.get("pathways", []))
|
|
self._import_nodes(package, package_data.get("nodes", []))
|
|
self._import_edges(package, package_data.get("edges", []))
|
|
self._import_scenarios(package, package_data.get("scenarios", []))
|
|
|
|
if package_data.get("models"):
|
|
self._import_models(package, package_data["models"])
|
|
|
|
# Set default structures for compounds (after all structures are created)
|
|
self._set_default_structures(package_data.get("compounds", []))
|
|
|
|
print(f"Package import completed: {package.name}")
|
|
return package
|
|
|
|
def _create_package(self, package_data: Dict[str, Any]) -> Package:
|
|
"""Create the main package object."""
|
|
package_uuid = self._get_or_generate_uuid(package_data["uuid"])
|
|
|
|
# Handle license
|
|
license_obj = None
|
|
if package_data.get("license"):
|
|
license_data = package_data["license"]
|
|
license_obj, _ = License.objects.get_or_create(
|
|
name=license_data["name"],
|
|
defaults={
|
|
"cc_string": license_data.get("cc_string", ""),
|
|
"link": license_data.get("link", ""),
|
|
"image_link": license_data.get("image_link", ""),
|
|
},
|
|
)
|
|
|
|
new_name = package_data.get("name")
|
|
if self.add_import_timestamp:
|
|
new_name = f"{new_name} - Imported at {datetime.now()}"
|
|
|
|
new_reviewed = False
|
|
if self.trust_reviewed:
|
|
new_reviewed = package_data.get("reviewed", False)
|
|
|
|
package = Package.objects.create(
|
|
uuid=package_uuid,
|
|
name=new_name,
|
|
description=package_data["description"],
|
|
kv=package_data.get("kv", {}),
|
|
reviewed=new_reviewed,
|
|
license=license_obj,
|
|
)
|
|
|
|
self._cache_object("Package", package_data["uuid"], package)
|
|
print(f"Created package: {package.name}")
|
|
return package
|
|
|
|
def _import_compounds(self, package: Package, compounds_data: List[Dict[str, Any]]):
|
|
"""Import compounds."""
|
|
print(f"Importing {len(compounds_data)} compounds...")
|
|
|
|
for compound_data in compounds_data:
|
|
compound_uuid = self._get_or_generate_uuid(compound_data["uuid"])
|
|
|
|
compound = Compound.objects.create(
|
|
uuid=compound_uuid,
|
|
package=package,
|
|
name=compound_data["name"],
|
|
description=compound_data["description"],
|
|
kv=compound_data.get("kv", {}),
|
|
# default_structure will be set later
|
|
)
|
|
|
|
# Set aliases if present
|
|
if compound_data.get("aliases"):
|
|
compound.aliases = compound_data["aliases"]
|
|
compound.save()
|
|
|
|
self._cache_object("Compound", compound_data["uuid"], compound)
|
|
|
|
# Handle external identifiers
|
|
self._create_external_identifiers(
|
|
compound, compound_data.get("external_identifiers", [])
|
|
)
|
|
|
|
def _import_structures(self, package: Package, structures_data: List[Dict[str, Any]]):
|
|
"""Import compound structures."""
|
|
print(f"Importing {len(structures_data)} compound structures...")
|
|
|
|
for structure_data in structures_data:
|
|
structure_uuid = self._get_or_generate_uuid(structure_data["uuid"])
|
|
compound_uuid = structure_data["compound"]["uuid"]
|
|
compound = self._get_cached_object("Compound", compound_uuid)
|
|
|
|
if not compound:
|
|
print(f"Warning: Compound with UUID {compound_uuid} not found for structure")
|
|
continue
|
|
|
|
structure = CompoundStructure.objects.create(
|
|
uuid=structure_uuid,
|
|
compound=compound,
|
|
name=structure_data["name"],
|
|
description=structure_data["description"],
|
|
kv=structure_data.get("kv", {}),
|
|
smiles=structure_data["smiles"],
|
|
canonical_smiles=structure_data["canonical_smiles"],
|
|
inchikey=structure_data["inchikey"],
|
|
normalized_structure=structure_data.get("normalized_structure", False),
|
|
)
|
|
|
|
# Set aliases if present
|
|
if structure_data.get("aliases"):
|
|
structure.aliases = structure_data["aliases"]
|
|
structure.save()
|
|
|
|
self._cache_object("CompoundStructure", structure_data["uuid"], structure)
|
|
|
|
# Handle external identifiers
|
|
self._create_external_identifiers(
|
|
structure, structure_data.get("external_identifiers", [])
|
|
)
|
|
|
|
def _import_rules(self, package: Package, rules_data: Dict[str, Any]):
|
|
"""Import all types of rules."""
|
|
print("Importing rules...")
|
|
|
|
# Import simple rules first
|
|
simple_rules_data = rules_data.get("simple_rules", [])
|
|
print(f"Importing {len(simple_rules_data)} simple rules...")
|
|
|
|
for rule_data in simple_rules_data:
|
|
self._create_simple_rule(package, rule_data)
|
|
|
|
# Import parallel rules
|
|
parallel_rules_data = rules_data.get("parallel_rules", [])
|
|
print(f"Importing {len(parallel_rules_data)} parallel rules...")
|
|
|
|
for rule_data in parallel_rules_data:
|
|
self._create_parallel_rule(package, rule_data)
|
|
|
|
def _create_simple_rule(self, package: Package, rule_data: Dict[str, Any]):
|
|
"""Create a simple rule (SimpleAmbitRule or SimpleRDKitRule)."""
|
|
rule_uuid = self._get_or_generate_uuid(rule_data["uuid"])
|
|
rule_type = rule_data.get("rule_type", "SimpleRule")
|
|
|
|
common_fields = {
|
|
"uuid": rule_uuid,
|
|
"package": package,
|
|
"name": rule_data["name"],
|
|
"description": rule_data["description"],
|
|
"kv": rule_data.get("kv", {}),
|
|
}
|
|
|
|
if rule_type == "SimpleAmbitRule":
|
|
rule = SimpleAmbitRule.objects.create(
|
|
**common_fields,
|
|
smirks=rule_data.get("smirks", ""),
|
|
reactant_filter_smarts=rule_data.get("reactant_filter_smarts", ""),
|
|
product_filter_smarts=rule_data.get("product_filter_smarts", ""),
|
|
)
|
|
elif rule_type == "SimpleRDKitRule":
|
|
rule = SimpleRDKitRule.objects.create(
|
|
**common_fields, reaction_smarts=rule_data.get("reaction_smarts", "")
|
|
)
|
|
else:
|
|
rule = SimpleRule.objects.create(**common_fields)
|
|
|
|
# Set aliases if present
|
|
if rule_data.get("aliases"):
|
|
rule.aliases = rule_data["aliases"]
|
|
rule.save()
|
|
|
|
self._cache_object("SimpleRule", rule_data["uuid"], rule)
|
|
return rule
|
|
|
|
def _create_parallel_rule(self, package: Package, rule_data: Dict[str, Any]):
|
|
"""Create a parallel rule."""
|
|
rule_uuid = self._get_or_generate_uuid(rule_data["uuid"])
|
|
|
|
rule = ParallelRule.objects.create(
|
|
uuid=rule_uuid,
|
|
package=package,
|
|
name=rule_data["name"],
|
|
description=rule_data["description"],
|
|
kv=rule_data.get("kv", {}),
|
|
)
|
|
|
|
# Set aliases if present
|
|
if rule_data.get("aliases"):
|
|
rule.aliases = rule_data["aliases"]
|
|
rule.save()
|
|
|
|
# Add simple rules
|
|
for simple_rule_ref in rule_data.get("simple_rules", []):
|
|
simple_rule = self._get_cached_object("SimpleRule", simple_rule_ref["uuid"])
|
|
if simple_rule:
|
|
rule.simple_rules.add(simple_rule)
|
|
|
|
self._cache_object("ParallelRule", rule_data["uuid"], rule)
|
|
return rule
|
|
|
|
def _import_reactions(self, package: Package, reactions_data: List[Dict[str, Any]]):
|
|
"""Import reactions."""
|
|
print(f"Importing {len(reactions_data)} reactions...")
|
|
|
|
for reaction_data in reactions_data:
|
|
reaction_uuid = self._get_or_generate_uuid(reaction_data["uuid"])
|
|
|
|
reaction = Reaction.objects.create(
|
|
uuid=reaction_uuid,
|
|
package=package,
|
|
name=reaction_data["name"],
|
|
description=reaction_data["description"],
|
|
kv=reaction_data.get("kv", {}),
|
|
multi_step=reaction_data.get("multi_step", False),
|
|
medline_references=reaction_data.get("medline_references", []),
|
|
)
|
|
|
|
# Set aliases if present
|
|
if reaction_data.get("aliases"):
|
|
reaction.aliases = reaction_data["aliases"]
|
|
reaction.save()
|
|
|
|
# Add educts and products
|
|
for educt_ref in reaction_data.get("educts", []):
|
|
compound = self._get_cached_object("CompoundStructure", educt_ref["uuid"])
|
|
if compound:
|
|
reaction.educts.add(compound)
|
|
|
|
for product_ref in reaction_data.get("products", []):
|
|
compound = self._get_cached_object("CompoundStructure", product_ref["uuid"])
|
|
if compound:
|
|
reaction.products.add(compound)
|
|
|
|
# Add rules
|
|
for rule_ref in reaction_data.get("rules", []):
|
|
# Try to find rule in different caches
|
|
rule = self._get_cached_object(
|
|
"SimpleRule", rule_ref["uuid"]
|
|
) or self._get_cached_object("ParallelRule", rule_ref["uuid"])
|
|
if rule:
|
|
reaction.rules.add(rule)
|
|
|
|
self._cache_object("Reaction", reaction_data["uuid"], reaction)
|
|
|
|
# Handle external identifiers
|
|
self._create_external_identifiers(
|
|
reaction, reaction_data.get("external_identifiers", [])
|
|
)
|
|
|
|
def _import_pathways(self, package: Package, pathways_data: List[Dict[str, Any]]):
|
|
"""Import pathways."""
|
|
print(f"Importing {len(pathways_data)} pathways...")
|
|
|
|
for pathway_data in pathways_data:
|
|
pathway_uuid = self._get_or_generate_uuid(pathway_data["uuid"])
|
|
|
|
pathway = Pathway.objects.create(
|
|
uuid=pathway_uuid,
|
|
package=package,
|
|
name=pathway_data["name"],
|
|
description=pathway_data["description"],
|
|
kv=pathway_data.get("kv", {}),
|
|
# setting will be handled separately if needed
|
|
)
|
|
|
|
# Set aliases if present
|
|
if pathway_data.get("aliases"):
|
|
pathway.aliases = pathway_data["aliases"]
|
|
pathway.save()
|
|
|
|
self._cache_object("Pathway", pathway_data["uuid"], pathway)
|
|
|
|
def _import_nodes(self, package: Package, nodes_data: List[Dict[str, Any]]):
|
|
"""Import pathway nodes."""
|
|
print(f"Importing {len(nodes_data)} nodes...")
|
|
|
|
for node_data in nodes_data:
|
|
node_uuid = self._get_or_generate_uuid(node_data["uuid"])
|
|
pathway_uuid = node_data["pathway"]["uuid"]
|
|
pathway = self._get_cached_object("Pathway", pathway_uuid)
|
|
|
|
if not pathway:
|
|
print(f"Warning: Pathway with UUID {pathway_uuid} not found for node")
|
|
continue
|
|
|
|
# For now, we'll set default_node_label to None and handle it later
|
|
# as it requires compound structures to be fully imported
|
|
node = Node.objects.create(
|
|
uuid=node_uuid,
|
|
pathway=pathway,
|
|
name=node_data["name"],
|
|
description=node_data["description"],
|
|
kv=node_data.get("kv", {}),
|
|
depth=node_data.get("depth", 0),
|
|
default_node_label=self._get_cached_object(
|
|
"CompoundStructure", node_data["default_node_label"]["uuid"]
|
|
),
|
|
)
|
|
|
|
# Set aliases if present
|
|
if node_data.get("aliases"):
|
|
node.aliases = node_data["aliases"]
|
|
node.save()
|
|
|
|
self._cache_object("Node", node_data["uuid"], node)
|
|
# Store node_data for later processing of relationships
|
|
node._import_data = node_data
|
|
|
|
def _import_edges(self, package: Package, edges_data: List[Dict[str, Any]]):
|
|
"""Import pathway edges."""
|
|
print(f"Importing {len(edges_data)} edges...")
|
|
|
|
for edge_data in edges_data:
|
|
edge_uuid = self._get_or_generate_uuid(edge_data["uuid"])
|
|
pathway_uuid = edge_data["pathway"]["uuid"]
|
|
pathway = self._get_cached_object("Pathway", pathway_uuid)
|
|
|
|
if not pathway:
|
|
print(f"Warning: Pathway with UUID {pathway_uuid} not found for edge")
|
|
continue
|
|
|
|
# For now, we'll set edge_label to None and handle it later
|
|
edge = Edge.objects.create(
|
|
uuid=edge_uuid,
|
|
pathway=pathway,
|
|
name=edge_data["name"],
|
|
description=edge_data["description"],
|
|
kv=edge_data.get("kv", {}),
|
|
edge_label=self._get_cached_object("Reaction", edge_data["edge_label"]["uuid"]),
|
|
)
|
|
|
|
# Set aliases if present
|
|
if edge_data.get("aliases"):
|
|
edge.aliases = edge_data["aliases"]
|
|
edge.save()
|
|
|
|
# Add start and end nodes
|
|
for start_node_ref in edge_data.get("start_nodes", []):
|
|
node = self._get_cached_object("Node", start_node_ref["uuid"])
|
|
if node:
|
|
edge.start_nodes.add(node)
|
|
|
|
for end_node_ref in edge_data.get("end_nodes", []):
|
|
node = self._get_cached_object("Node", end_node_ref["uuid"])
|
|
if node:
|
|
edge.end_nodes.add(node)
|
|
|
|
self._cache_object("Edge", edge_data["uuid"], edge)
|
|
|
|
def _import_scenarios(self, package: Package, scenarios_data: List[Dict[str, Any]]):
|
|
"""Import scenarios."""
|
|
print(f"Importing {len(scenarios_data)} scenarios...")
|
|
|
|
# First pass: create scenarios without parent relationships
|
|
for scenario_data in scenarios_data:
|
|
scenario_uuid = self._get_or_generate_uuid(scenario_data["uuid"])
|
|
|
|
scenario_date = None
|
|
if scenario_data.get("scenario_date"):
|
|
scenario_date = scenario_data["scenario_date"]
|
|
|
|
scenario = Scenario.objects.create(
|
|
uuid=scenario_uuid,
|
|
package=package,
|
|
name=scenario_data["name"],
|
|
description=scenario_data["description"],
|
|
kv=scenario_data.get("kv", {}),
|
|
scenario_date=scenario_date,
|
|
scenario_type=scenario_data.get("scenario_type"),
|
|
additional_information=scenario_data.get("additional_information", {}),
|
|
)
|
|
|
|
self._cache_object("Scenario", scenario_data["uuid"], scenario)
|
|
# Store scenario_data for later processing of parent relationships
|
|
scenario._import_data = scenario_data
|
|
|
|
# Second pass: set parent relationships
|
|
for scenario_data in scenarios_data:
|
|
if scenario_data.get("parent"):
|
|
scenario = self._get_cached_object("Scenario", scenario_data["uuid"])
|
|
parent = self._get_cached_object("Scenario", scenario_data["parent"]["uuid"])
|
|
if scenario and parent:
|
|
scenario.parent = parent
|
|
scenario.save()
|
|
|
|
def _import_models(self, package: Package, models_data: List[Dict[str, Any]]):
|
|
"""Import EPModels."""
|
|
print(f"Importing {len(models_data)} models...")
|
|
|
|
for model_data in models_data:
|
|
model_uuid = self._get_or_generate_uuid(model_data["uuid"])
|
|
model_type = model_data.get("model_type", "EPModel")
|
|
|
|
common_fields = {
|
|
"uuid": model_uuid,
|
|
"package": package,
|
|
"name": model_data["name"],
|
|
"description": model_data["description"],
|
|
"kv": model_data.get("kv", {}),
|
|
}
|
|
|
|
# Add PackageBasedModel fields if present
|
|
if "threshold" in model_data:
|
|
common_fields.update(
|
|
{
|
|
"threshold": model_data.get("threshold"),
|
|
"eval_results": model_data.get("eval_results", {}),
|
|
"model_status": model_data.get("model_status", "INITIAL"),
|
|
}
|
|
)
|
|
|
|
# Create the appropriate model type
|
|
if model_type == "RuleBasedRelativeReasoning":
|
|
model = RuleBasedRelativeReasoning.objects.create(
|
|
**common_fields,
|
|
min_count=model_data.get("min_count", 1),
|
|
max_count=model_data.get("max_count", 10),
|
|
)
|
|
elif model_type == "MLRelativeReasoning":
|
|
model = MLRelativeReasoning.objects.create(**common_fields)
|
|
elif model_type == "EnviFormer":
|
|
model = EnviFormer.objects.create(**common_fields)
|
|
else:
|
|
model = EPModel.objects.create(**common_fields)
|
|
|
|
# Set aliases if present
|
|
if model_data.get("aliases"):
|
|
model.aliases = model_data["aliases"]
|
|
model.save()
|
|
|
|
# Add package relationships for PackageBasedModel
|
|
if hasattr(model, "rule_packages"):
|
|
for pkg_ref in model_data.get("rule_packages", []):
|
|
pkg = self._get_cached_object("Package", pkg_ref["uuid"])
|
|
if pkg:
|
|
model.rule_packages.add(pkg)
|
|
|
|
for pkg_ref in model_data.get("data_packages", []):
|
|
pkg = self._get_cached_object("Package", pkg_ref["uuid"])
|
|
if pkg:
|
|
model.data_packages.add(pkg)
|
|
|
|
for pkg_ref in model_data.get("eval_packages", []):
|
|
pkg = self._get_cached_object("Package", pkg_ref["uuid"])
|
|
if pkg:
|
|
model.eval_packages.add(pkg)
|
|
|
|
self._cache_object("EPModel", model_data["uuid"], model)
|
|
|
|
def _set_default_structures(self, compounds_data: List[Dict[str, Any]]):
|
|
"""Set default structures for compounds after all structures are created."""
|
|
print("Setting default structures for compounds...")
|
|
|
|
for compound_data in compounds_data:
|
|
if compound_data.get("default_structure"):
|
|
compound = self._get_cached_object("Compound", compound_data["uuid"])
|
|
structure = self._get_cached_object(
|
|
"CompoundStructure", compound_data["default_structure"]["uuid"]
|
|
)
|
|
if compound and structure:
|
|
compound.default_structure = structure
|
|
compound.save()
|
|
|
|
def _create_external_identifiers(self, obj, identifiers_data: List[Dict[str, Any]]):
|
|
"""Create external identifiers for an object."""
|
|
for identifier_data in identifiers_data:
|
|
# Get or create the external database
|
|
db_data = identifier_data["database"]
|
|
database, _ = ExternalDatabase.objects.get_or_create(
|
|
name=db_data["name"],
|
|
defaults={
|
|
"base_url": db_data.get("base_url", ""),
|
|
"full_name": db_data.get("name", ""),
|
|
"description": "",
|
|
"is_active": True,
|
|
},
|
|
)
|
|
|
|
# Create the external identifier
|
|
ExternalIdentifier.objects.create(
|
|
content_object=obj,
|
|
database=database,
|
|
identifier_value=identifier_data["identifier_value"],
|
|
url=identifier_data.get("url", ""),
|
|
is_primary=identifier_data.get("is_primary", False),
|
|
)
|
|
|
|
|
|
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
|