Import and expose emodel traces

alembic/versions/20250602_145126_8f0d631d2bd4_default_migration_message.py

@@ -0,0 +1,45 @@
+"""Default migration message
+
+Revision ID: 8f0d631d2bd4
+Revises: 1589bff44728
+Create Date: 2025-06-02 14:51:26.859717
+
+"""
+
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+
+
+from sqlalchemy import Text
+import app.db.types
+
+# revision identifiers, used by Alembic.
+revision: str = "8f0d631d2bd4"
+down_revision: Union[str, None] = "1589bff44728"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table(
+        "derivation",
+        sa.Column("used_id", sa.Uuid(), nullable=False),
+        sa.Column("generated_id", sa.Uuid(), nullable=False),
+        sa.ForeignKeyConstraint(
+            ["generated_id"], ["entity.id"], name=op.f("fk_derivation_generated_id_entity")
+        ),
+        sa.ForeignKeyConstraint(
+            ["used_id"], ["entity.id"], name=op.f("fk_derivation_used_id_entity")
+        ),
+        sa.PrimaryKeyConstraint("used_id", "generated_id", name=op.f("pk_derivation")),
+    )
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_table("derivation")
+    # ### end Alembic commands ###

app/cli/import_data.py

@@ -7,7 +7,6 @@
 from abc import ABC, abstractmethod
 from collections import Counter, defaultdict
 from contextlib import closing
-from operator import attrgetter
 from pathlib import Path
 from typing import Any
 
@@ -18,7 +17,7 @@
 
 from app.cli import curate, utils
 from app.cli.brain_region_data import BRAIN_ATLAS_REGION_VOLUMES
-from app.cli.curation import electrical_cell_recording
+from app.cli.curation import cell_composition, electrical_cell_recording
 from app.cli.types import ContentType
 from app.cli.utils import (
     AUTHORIZED_PUBLIC,
@@ -37,6 +36,7 @@
     BrainRegionHierarchy,
     CellComposition,
     DataMaturityAnnotationBody,
+    Derivation,
     ElectricalCellRecording,
     EModel,
     Entity,
@@ -71,10 +71,6 @@
 from app.logger import L
 from app.schemas.base import ProjectContext
 
-from app.cli.curation import electrical_cell_recording, cell_composition
-from app.cli.types import ContentType
-
-
 BRAIN_ATLAS_NAME = "BlueBrain Atlas"
 
 REQUIRED_PATH = click.Path(exists=True, readable=True, dir_okay=False, resolve_path=True)
@@ -561,6 +557,75 @@ def ingest(
         db.commit()
 
 
+class ImportEModelDerivations(Import):
+    name = "EModelWorkflow"
+
+    @staticmethod
+    def is_correct_type(data):
+        types = ensurelist(data["@type"])
+        return "EModelWorkflow" in types
+
+    @staticmethod
+    def ingest(
+        db,
+        project_context,
+        data_list: list[dict],
+        all_data_by_id: dict[str, dict],
+        hierarchy_name: str,
+    ):
+        """Import emodel derivations from EModelWorkflow."""
+        legacy_emodel_ids = set()
+        derivations = {}
+        for data in tqdm(data_list, desc="EModelWorkflow"):
+            legacy_emodel_id = utils.find_id_in_entity(data, "EModel", "generates")
+            legacy_etc_id = utils.find_id_in_entity(
+                data, "ExtractionTargetsConfiguration", "hasPart"
+            )
+            if not legacy_emodel_id:
+                L.warning("Not found EModel id in EModelWorkflow: {}", data["@id"])
+                continue
+            if not legacy_etc_id:
+                L.warning(
+                    "Not found ExtractionTargetsConfiguration id in EModelWorkflow: {}", data["@id"]
+                )
+                continue
+            if not (etc := all_data_by_id.get(legacy_etc_id)):
+                L.warning("Not found ExtractionTargetsConfiguration with id {}", legacy_etc_id)
+                continue
+            if not (legacy_trace_ids := list(utils.find_ids_in_entity(etc, "Trace", "uses"))):
+                L.warning(
+                    "Not found traces in ExtractionTargetsConfiguration with id {}", legacy_etc_id
+                )
+                continue
+            if legacy_emodel_id in legacy_emodel_ids:
+                L.warning("Duplicated and ignored traces for EModel id {}", legacy_emodel_id)
+                continue
+            legacy_emodel_ids.add(legacy_emodel_id)
+            if not (emodel := utils._find_by_legacy_id(legacy_emodel_id, EModel, db)):
+                L.warning("Not found EModel with legacy id {}", legacy_emodel_id)
+                continue
+            if emodel.id in derivations:
+                L.warning("Duplicated and ignored traces for EModel uuid {}", emodel.id)
+            derivations[emodel.id] = [
+                utils._find_by_legacy_id(legacy_trace_id, ElectricalCellRecording, db).id
+                for legacy_trace_id in legacy_trace_ids
+            ]
+
+        rows = [
+            Derivation(used_id=trace_id, generated_id=emodel_id)
+            for emodel_id, trace_ids in derivations.items()
+            for trace_id in trace_ids
+        ]
+        L.info(
+            "Imported derivations for {} EModels from {} records", len(derivations), len(data_list)
+        )
+        # delete everything from derivation table before adding the records
+        query = sa.delete(Derivation)
+        db.execute(query)
+        db.add_all(rows)
+        db.commit()
+
+
 class ImportBrainAtlas(Import):
     name = "BrainAtlas"
 
@@ -1370,6 +1435,7 @@ def _do_import(db, input_dir, project_context, hierarchy_name):
         ImportBrainAtlas,
         ImportDistribution,
         ImportNeuronMorphologyFeatureAnnotation,
+        ImportEModelDerivations,
     ]
 
     for importer in importers:

app/cli/utils.py

@@ -361,13 +361,14 @@ def get_or_create_distribution(
 def find_id_in_entity(entity: dict | None, type_: str, entity_list_key: str):
     if not entity:
         return None
-    return next(
-        (
-            part.get("@id")
-            for part in ensurelist(entity.get(entity_list_key, []))
-            if is_type(part, type_)
-        ),
-        None,
+    return next(find_ids_in_entity(entity, type_, entity_list_key), None)
+
+
+def find_ids_in_entity(entity: dict, type_: str, entity_list_key: str):
+    return (
+        part.get("@id")
+        for part in ensurelist(entity.get(entity_list_key, []))
+        if is_type(part, type_)
     )
 
 

app/db/model.py

@@ -919,3 +919,11 @@ class CellComposition(NameDescriptionVectorMixin, LocationMixin, SpeciesMixin, E
     __tablename__ = EntityType.cell_composition
     id: Mapped[uuid.UUID] = mapped_column(ForeignKey("entity.id"), primary_key=True)
     __mapper_args__ = {"polymorphic_identity": __tablename__}  # noqa: RUF012
+
+
+class Derivation(Base):
+    __tablename__ = "derivation"
+    used_id: Mapped[uuid.UUID] = mapped_column(ForeignKey("entity.id"), primary_key=True)
+    generated_id: Mapped[uuid.UUID] = mapped_column(ForeignKey("entity.id"), primary_key=True)
+    used: Mapped["Entity"] = relationship(foreign_keys=[used_id])
+    generated: Mapped["Entity"] = relationship(foreign_keys=[generated_id])

app/filters/entity.py

@@ -0,0 +1,20 @@
+from typing import Annotated
+
+from fastapi_filter import FilterDepends
+
+from app.db.model import Entity
+from app.db.types import EntityType
+from app.filters.base import CustomFilter
+
+
+class BasicEntityFilter(CustomFilter):
+    type: EntityType | None = None
+
+    order_by: list[str] = ["-creation_date"]  # noqa: RUF012
+
+    class Constants(CustomFilter.Constants):
+        model = Entity
+        ordering_model_fields = ["creation_date", "update_date", "name"]  # noqa: RUF012
+
+
+BasicEntityFilterDep = Annotated[BasicEntityFilter, FilterDepends(BasicEntityFilter)]

app/routers/__init__.py

@@ -10,6 +10,7 @@
     brain_region_hierarchy,
     cell_composition,
     contribution,
+    derivation,
     electrical_cell_recording,
     emodel,
     etype,
@@ -45,6 +46,7 @@
     brain_region_hierarchy.router,
     cell_composition.router,
     contribution.router,
+    derivation.router,
     electrical_cell_recording.router,
     emodel.router,
     etype.router,

app/routers/asset.py

@@ -1,16 +1,15 @@
 """Generic asset routes."""
 
 import uuid
-from enum import StrEnum
 from http import HTTPStatus
 from pathlib import Path
-from typing import TYPE_CHECKING, Annotated
+from typing import Annotated
 
 from fastapi import APIRouter, Form, HTTPException, UploadFile, status
 from starlette.responses import RedirectResponse
 
 from app.config import settings
-from app.db.types import AssetLabel, EntityType
+from app.db.types import AssetLabel
 from app.dependencies.auth import UserContextDep, UserContextWithProjectIdDep
 from app.dependencies.db import RepoGroupDep
 from app.dependencies.s3 import S3ClientDep
@@ -19,6 +18,7 @@
 from app.schemas.types import ListResponse, PaginationResponse
 from app.service import asset as asset_service
 from app.utils.files import calculate_sha256_digest, get_content_type
+from app.utils.routers import EntityRoute, entity_route_to_type
 from app.utils.s3 import (
     delete_from_s3,
     generate_presigned_url,
@@ -32,18 +32,6 @@
     tags=["assets"],
 )
 
-if not TYPE_CHECKING:
-    # EntityRoute (hyphen-separated) <-> EntityType (underscore_separated)
-    EntityRoute = StrEnum(
-        "EntityRoute", {item.name: item.name.replace("_", "-") for item in EntityType}
-    )
-else:
-    EntityRoute = StrEnum
-
-
-def _entity_route_to_type(entity_route: EntityRoute) -> EntityType:
-    return EntityType[entity_route.name]
-
 
 @router.get("/{entity_route}/{entity_id}/assets")
 def get_entity_assets(
@@ -56,7 +44,7 @@ def get_entity_assets(
     assets = asset_service.get_entity_assets(
         repos,
         user_context=user_context,
-        entity_type=_entity_route_to_type(entity_route),
+        entity_type=entity_route_to_type(entity_route),
         entity_id=entity_id,
     )
     # TODO: proper pagination
@@ -76,7 +64,7 @@ def get_entity_asset(
     return asset_service.get_entity_asset(
         repos,
         user_context=user_context,
-        entity_type=_entity_route_to_type(entity_route),
+        entity_type=entity_route_to_type(entity_route),
         entity_id=entity_id,
         asset_id=asset_id,
     )
@@ -117,7 +105,7 @@ def upload_entity_asset(
     asset_read = asset_service.create_entity_asset(
         repos=repos,
         user_context=user_context,
-        entity_type=_entity_route_to_type(entity_route),
+        entity_type=entity_route_to_type(entity_route),
         entity_id=entity_id,
         filename=file.filename,
         content_type=content_type,
@@ -149,7 +137,7 @@ def download_entity_asset(
     asset = asset_service.get_entity_asset(
         repos,
         user_context=user_context,
-        entity_type=_entity_route_to_type(entity_route),
+        entity_type=entity_route_to_type(entity_route),
         entity_id=entity_id,
         asset_id=asset_id,
     )
@@ -205,7 +193,7 @@ def delete_entity_asset(
     asset = asset_service.delete_entity_asset(
         repos,
         user_context=user_context,
-        entity_type=_entity_route_to_type(entity_route),
+        entity_type=entity_route_to_type(entity_route),
         entity_id=entity_id,
         asset_id=asset_id,
     )

app/routers/derivation.py

@@ -0,0 +1,12 @@
+"""Generic derivation routes."""
+
+from fastapi import APIRouter
+
+import app.service.derivation
+
+router = APIRouter(
+    prefix="",
+    tags=["derivation"],
+)
+
+router.get("/{entity_route}/{entity_id}/derived-from")(app.service.derivation.read_many)

app/routers/emodel.py

@@ -1,5 +1,6 @@
 from fastapi import APIRouter
 
+import app.service.electrical_cell_recording
 import app.service.emodel
 
 router = APIRouter(

app/schemas/base.py

@@ -109,3 +109,7 @@ class LicensedCreateMixin(BaseModel):
 class LicensedReadMixin(BaseModel):
     model_config = ConfigDict(from_attributes=True)
     license: LicenseRead | None
+
+
+class BasicEntityRead(IdentifiableMixin, EntityTypeMixin):
+    pass