fix referenced morphs/biophysics: load from included files, add tests, add some error checking

.pre-commit-config.yaml

@@ -5,7 +5,9 @@ repos:
     rev: v4.4.0
     hooks:
     -   id: trailing-whitespace
+        exclude: NeuroML.*xsd
     -   id: end-of-file-fixer
+        exclude: NeuroML.*xsd
     -   id: check-added-large-files
 - repo: https://github.com/astral-sh/ruff-pre-commit
   rev: v0.4.1

neuroml/test/test_utils.py

@@ -7,14 +7,17 @@
 Copyright 2023 NeuroML contributors
 """
 
+import os
 import unittest
 
 import neuroml
 from neuroml.utils import (
     component_factory,
+    fix_external_morphs_biophys_in_cell,
     get_relative_component_path,
     print_hierarchy,
 )
+from neuroml.writers import NeuroMLWriter
 
 
 class UtilsTestCase(unittest.TestCase):
@@ -55,3 +58,129 @@ def test_networkx_hier_graph(self):
         path, graph = get_relative_component_path("Input", "Instance")
         self.assertIsNotNone(graph)
         self.assertEqual(path, "../Population/Instance")
+
+    def test_fix_external_morphs_biophys_in_cell(self):
+        """Test fix_external_morphs_biophys_in_cell function"""
+        # document that includes cell and morphology with cell referring to
+        # morphology
+        nml_doc = component_factory("NeuroMLDocument", id="testdoc")
+        nml_doc.add("Morphology", id="test_morph_1", validate=False)
+        test_cell_1 = nml_doc.add(
+            "Cell", id="test_cell_1", morphology_attr="test_morph_1"
+        )
+        test_cell_1.morphology = None
+
+        fix_external_morphs_biophys_in_cell(nml_doc)
+        self.assertIsNotNone(nml_doc.cells[0].morphology)
+        self.assertIsNone(nml_doc.cells[0].morphology_attr)
+        self.assertEqual(nml_doc.cells[0].morphology.id, "test_morph_1")
+        print(nml_doc)
+
+        # check that a key error is raised if the referenced morph cannot be
+        # found
+        nml_doc_2 = component_factory("NeuroMLDocument", id="testdoc")
+        nml_doc_2.add("Morphology", id="test_morph_5", validate=False)
+        test_cell_2 = nml_doc_2.add(
+            "Cell", id="test_cell_1", morphology_attr="test_morph_1"
+        )
+        test_cell_2.morphology = None
+
+        with self.assertRaises(KeyError):
+            fix_external_morphs_biophys_in_cell(nml_doc_2)
+            print(nml_doc_2)
+
+    def test_fix_external_morphs_biophys_in_cell_2(self):
+        """Test fix_external_morphs_biophys_in_cell function"""
+        # document that includes cell and biophysical properties with cell
+        # referring to biophysical properties
+        nml_doc = component_factory("NeuroMLDocument", id="testdoc")
+        nml_doc.add("BiophysicalProperties", id="test_biophys_1", validate=False)
+        test_cell_1 = nml_doc.add(
+            "Cell", id="test_cell_1", biophysical_properties_attr="test_biophys_1"
+        )
+        test_cell_1.biophysical_properties = None
+
+        fix_external_morphs_biophys_in_cell(nml_doc)
+        self.assertIsNotNone(nml_doc.cells[0].biophysical_properties)
+        self.assertIsNone(nml_doc.cells[0].biophysical_properties_attr)
+        self.assertEqual(nml_doc.cells[0].biophysical_properties.id, "test_biophys_1")
+        print(nml_doc)
+
+        # check that a key error is raised if the referenced biophysical
+        # property cannot be found
+        nml_doc_2 = component_factory("NeuroMLDocument", id="testdoc")
+        nml_doc_2.add("BiophysicalProperties", id="test_biophys_5", validate=False)
+        test_cell_2 = nml_doc_2.add(
+            "Cell", id="test_cell_1", biophysical_properties_attr="test_biophys_1"
+        )
+        test_cell_2.biophysical_properties = None
+        print(nml_doc_2)
+
+        with self.assertRaises(KeyError):
+            fix_external_morphs_biophys_in_cell(nml_doc_2)
+
+    def test_fix_external_morphs_biophys_in_cell_3(self):
+        """Test fix_external_morphs_biophys_in_cell function"""
+        filename = "nml_morph_doc.nml"
+        nml_doc = component_factory("NeuroMLDocument", id="testdoc")
+        nml_doc.add("Morphology", id="test_morph_1", validate=False)
+        NeuroMLWriter.write(nml_doc, file=filename)
+
+        # doc that includes
+        nml_doc_2 = component_factory("NeuroMLDocument", id="testdoc")
+        nml_doc_2.add("IncludeType", href=filename)
+        test_cell_1 = nml_doc_2.add(
+            "Cell", id="test_cell_1", morphology_attr="test_morph_1"
+        )
+        test_cell_1.morphology = None
+
+        fix_external_morphs_biophys_in_cell(nml_doc_2)
+        self.assertIsNotNone(nml_doc_2.cells[0].morphology)
+        self.assertIsNone(nml_doc_2.cells[0].morphology_attr)
+        self.assertEqual(nml_doc_2.cells[0].morphology.id, "test_morph_1")
+        print(nml_doc_2)
+        os.unlink(filename)
+
+        # doc that does not include, and should fail
+        nml_doc_3 = component_factory("NeuroMLDocument", id="testdoc")
+        test_cell_2 = nml_doc_3.add(
+            "Cell", id="test_cell_1", morphology_attr="test_morph_1"
+        )
+        test_cell_2.morphology = None
+        print(nml_doc_3)
+        with self.assertRaises(KeyError):
+            fix_external_morphs_biophys_in_cell(nml_doc_3)
+
+    def test_fix_external_morphs_biophys_in_cell_4(self):
+        """Test fix_external_morphs_biophys_in_cell function"""
+        # document that includes cell and biophysical properties with cell
+        # referring to biophysical properties
+        filename = "nml_biophys_doc.nml"
+        nml_doc = component_factory("NeuroMLDocument", id="testdoc")
+        nml_doc.add("BiophysicalProperties", id="test_biophys_1", validate=False)
+        NeuroMLWriter.write(nml_doc, file=filename)
+
+        # doc that includes
+        nml_doc_2 = component_factory("NeuroMLDocument", id="testdoc")
+        nml_doc_2.add("IncludeType", href=filename)
+        test_cell_1 = nml_doc_2.add(
+            "Cell", id="test_cell_1", biophysical_properties_attr="test_biophys_1"
+        )
+        test_cell_1.biophysical_properties = None
+
+        fix_external_morphs_biophys_in_cell(nml_doc_2)
+        self.assertIsNotNone(nml_doc_2.cells[0].biophysical_properties)
+        self.assertIsNone(nml_doc_2.cells[0].biophysical_properties_attr)
+        self.assertEqual(nml_doc_2.cells[0].biophysical_properties.id, "test_biophys_1")
+        print(nml_doc_2)
+        os.unlink(filename)
+
+        # doc that does not include, and should fail
+        nml_doc_3 = component_factory("NeuroMLDocument", id="testdoc")
+        test_cell_2 = nml_doc_3.add(
+            "Cell", id="test_cell_1", biophysical_properties_attr="test_biophys_1"
+        )
+        test_cell_2.biophysical_properties = None
+        print(nml_doc_3)
+        with self.assertRaises(KeyError):
+            fix_external_morphs_biophys_in_cell(nml_doc_3)

neuroml/utils.py

@@ -4,11 +4,13 @@
 
 """
 
+import copy
 import inspect
+import logging
 import os
 import sys
 import warnings
-from typing import Any, Dict, Optional, Type, Union
+from typing import Any, Dict, List, Optional, Set, Type, Union
 
 import networkx
 
@@ -17,6 +19,9 @@
 
 from . import loaders
 
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)
+
 
 def validate_neuroml2(file_name: str) -> None:
     """Validate a NeuroML document against the NeuroML schema specification.
@@ -310,18 +315,108 @@ def get_relative_component_path(
     return (path, graph)
 
 
-def fix_external_morphs_biophys_in_cell(nml2_doc: NeuroMLDocument) -> None:
-    """
-    Only used in the case where a cell element has a morphology (or biophysicalProperties) attribute, as opposed to a
-    subelement morphology/biophysicalProperties. This will substitute the external element into the cell element for ease of access
+def fix_external_morphs_biophys_in_cell(
+    nml2_doc: NeuroMLDocument, overwrite: bool = True
+) -> NeuroMLDocument:
+    """Handle externally referenced morphologies and biophysics in cells.
+
+    This is only used in the case where a cell element has a morphology (or
+    biophysicalProperties) attribute, as opposed to a subelement
+    morphology/biophysicalProperties. This will substitute the external element
+    into the cell element for ease of access
+
+    The referenced morphologies can be included in the same document directly,
+    or in other documents included using the "IncludeType". This function will
+    load the included documents and attempt to read referenced bits from them.
+
+    Note that if a cell already includes Morphology and BiophysicalProperties,
+    we just use those. Any references to other Morphology/BiophysicalProperties
+    elements will be ignored.
+
+    :param nml2_doc: NeuroML document
+    :type nml2_doc: neuroml.NeuroMLDocument
+    :param overwrite: toggle whether the document is overwritten or a deep copy
+        created
+    :type overwrite: bool
+    :returns: neuroml document
+    :raises KeyError: if referenced morphologies/biophysics cannot be found
     """
-    for cell in nml2_doc.cells:
-        if cell.morphology_attr != None:
-            ext_morph = nml2_doc.get_by_id(cell.morphology_attr)
-            cell.morphology = ext_morph
-        if cell.biophysical_properties_attr != None:
-            ext_bp = nml2_doc.get_by_id(cell.biophysical_properties_attr)
-            cell.biophysical_properties = ext_bp
+    if overwrite is False:
+        newdoc = copy.deepcopy(nml2_doc)
+    else:
+        newdoc = nml2_doc
+
+    # get a list of morph/biophys ids being referred to by cells
+    referenced_ids = []
+    for cell in newdoc.cells:
+        if cell.morphology_attr is not None:
+            if cell.morphology is None:
+                referenced_ids.append(cell.morphology_attr)
+            else:
+                logger.warning(
+                    f"Cell ({cell}) already contains a Morphology, ignoring reference."
+                )
+                logger.warning("Please check/correct your cell description")
+        if cell.biophysical_properties_attr is not None:
+            if cell.biophysical_properties is None:
+                referenced_ids.append(cell.biophysical_properties_attr)
+            else:
+                logger.warning(
+                    f"Cell ({cell}) already contains a BiophysicalProperties element, ignoring reference."
+                )
+                logger.warning("Please check/correct your cell description")
+
+    # load referenced ids from included files and store them in dicts
+    ext_morphs = {}
+    ext_biophys = {}
+    for inc in newdoc.includes:
+        incdoc = loaders.read_neuroml2_file(inc.href, verbose=False, optimized=True)
+        for morph in incdoc.morphology:
+            if morph.id in referenced_ids:
+                ext_morphs[morph.id] = morph
+        for biophys in incdoc.biophysical_properties:
+            if biophys.id in referenced_ids:
+                ext_biophys[biophys.id] = biophys
+
+    # also include morphs/biophys that are in the same document
+    for morph in newdoc.morphology:
+        if morph.id in referenced_ids:
+            ext_morphs[morph.id] = morph
+    for biophys in newdoc.biophysical_properties:
+        if biophys.id in referenced_ids:
+            ext_biophys[biophys.id] = biophys
+
+    # update cells by placing the morphology/biophys in them:
+    # if referenced ids are not found, throw errors
+    for cell in newdoc.cells:
+        if cell.morphology_attr is not None and cell.morphology is None:
+            try:
+                # TODO: do we need a deepcopy here?
+                cell.morphology = copy.deepcopy(ext_morphs[cell.morphology_attr])
+                cell.morphology_attr = None
+            except KeyError as e:
+                logger.error(
+                    f"Morphology with id {cell.morphology_attr} was not found in included/external morphologies."
+                )
+                raise e
+
+        if (
+            cell.biophysical_properties_attr is not None
+            and cell.biophysical_properties is None
+        ):
+            try:
+                # TODO: do we need a deepcopy here?
+                cell.biophysical_properties = copy.deepcopy(
+                    ext_biophys[cell.biophysical_properties_attr]
+                )
+                cell.biophysical_properties_attr = None
+            except KeyError as e:
+                logger.error(
+                    f"Biophysics with id {cell.biophysical_properties_attr} was not found in included/external biophysics."
+                )
+                raise e
+
+    return newdoc
 
 
 def main():