Merge pull request NeuroML#404 from AdityaBITMESRA/feat/ExportSWC

Feat/export swc-Adds function to export SWC graph to a new SWC file

pyneuroml/swc/LoadSWC.py

@@ -276,6 +276,24 @@ def get_branch_points(
                 branch_points[type_id] = self.get_nodes_with_multiple_children(type_id)
             return branch_points
 
+    def export_to_swc_file(self, filename: str) -> None:
+        """
+        Export the SWCGraph to a new SWC file.
+
+        :param filename: The path to the output SWC file
+        :type filename: str
+        """
+        with open(filename, "w") as file:
+            # Write metadata
+            for key, value in self.metadata.items():
+                file.write(f"# {key} {value}\n")
+
+            # Write node data
+            for node in sorted(self.nodes, key=lambda n: n.id):
+                file.write(
+                    f"{node.id} {node.type} {node.x:.4f} {node.y:.4f} {node.z:.4f} {node.radius:.4f} {node.parent_id}\n"
+                )
+
 
 def parse_header(line: str) -> typing.Optional[typing.Tuple[str, str]]:
     """

tests/swc/Case1_new.swc

@@ -0,0 +1,13 @@
+# Case 1: Single contour soma case
+# Fixed: 3 point soma centred at origin, length 20um, radius 10um
+
+1   1   0   0   0   10  -1
+2   1   0   -10 0   10  1
+3   1   0   10  0   10  1
+
+# Branching dendrite starting at "edge" of soma
+
+4   3   10  0   0   2   1
+5   3   30  0   0   2   4
+6   3   40  10   0   2   5
+7   3   40  -10   0   2   5

tests/swc/Case2_new.swc

@@ -0,0 +1,14 @@
+# Case 2: No soma case
+# NO CHANGE FROM ORIGINAL FORMAT
+
+# Soma not reconstructed...
+# 3 axons emanate from the same point
+
+1   2   0  0   0   2   -1
+2   2   20  0   0   2   1
+
+3   2   0  20   0   2   1
+4   2   0  30   0   2   3
+
+5   2   0  -20  0   2   1
+6   2   0  -30  0   2   5

tests/swc/Case3_new.swc

@@ -0,0 +1,12 @@
+# Case 3: Multiple contours soma case
+# Fixed: 3 point soma centred at origin, length 20um, radius 10um
+
+1   1   0   0   0   10  -1
+2   1   0   -10 0   10  1
+3   1   0   10  0   10  1
+
+# Single dendrite starting at "edge" of soma
+
+4  3   10  0   0   2   1
+5   3   20  0   0   2   4
+6   3   30  0   0   2   5

tests/swc/Case4_new.swc

@@ -0,0 +1,20 @@
+# Case 4: Multiple cylinder soma case
+# NO CHANGE FROM ORIGINAL FORMAT
+
+# 4 real soma points, i.e. 3 real segments starting at origin
+
+1   1   0   0   0   5  -1
+2   1   0   5   0   10  1
+3   1   0   10  0   10  2
+4   1   0   15  0   5  3
+
+# One dendrite starting at top of soma, one starting at the bottom, one from side
+
+5   3   0   20  0   5   4
+6   3   0   30  0   5   5
+
+7   3   0   -5  0   5   1
+8   3   0   -15  0   2.5   7
+
+9   3   10   10  0   5   2
+10  3   20   10  0   5   9

tests/swc/Case5_new.swc

@@ -0,0 +1,17 @@
+# Case 5: Spherical soma case
+# Fixed: 3 point soma centred at origin, length 20um, radius 10um
+
+1   1   0   0   0   10  -1
+2   1   0   -10 0   10  1
+3   1   0   10  0   10  1
+
+# 3 dendrites emanate from the soma, starting at a distance of 10um from the origin
+
+4   3   10  0   0   2   1
+5   3   30  0   0   2   4
+
+6   3   0  10   0   2   1
+7   3   0  30   0   2   6
+
+8   3   0  -10  0   2   1
+9   3   0  -30  0   2   8

tests/swc/test_LoadSWC.py

@@ -1,10 +1,14 @@
+import os
 import unittest
 
-from pyneuroml.swc.LoadSWC import SWCGraph, SWCNode
+from pyneuroml.swc.LoadSWC import SWCGraph, SWCNode, load_swc
 
 
 class TestSWCNode(unittest.TestCase):
+    """Test cases for the SWCNode class."""
+
     def test_init(self):
+        """Test the initialization of an SWCNode object."""
         node = SWCNode(1, 1, 0.0, 0.0, 0.0, 1.0, -1)
         self.assertEqual(node.id, 1)
         self.assertEqual(node.type, 1)
@@ -15,12 +19,16 @@ def test_init(self):
         self.assertEqual(node.parent_id, -1)
 
     def test_invalid_init(self):
+        """Test that initializing an SWCNode with invalid data raises a ValueError."""
         with self.assertRaises(ValueError):
             SWCNode("a", 1, 0.0, 0.0, 0.0, 1.0, -1)
 
 
 class TestSWCGraph(unittest.TestCase):
+    """Test cases for the SWCGraph class."""
+
     def setUp(self):
+        """Set up a sample SWCGraph for testing."""
         self.tree = SWCGraph()
         self.node1 = SWCNode(1, 1, 0.0, 0.0, 0.0, 1.0, -1)
         self.node2 = SWCNode(2, 3, 1.0, 0.0, 0.0, 0.5, 1)
@@ -30,39 +38,95 @@ def setUp(self):
         self.tree.add_node(self.node3)
 
     def test_duplicate_node(self):
+        """Test that adding a duplicate node raises a ValueError."""
         with self.assertRaises(ValueError):
             self.tree.add_node(SWCNode(1, 1, 0.0, 0.0, 0.0, 1.0, -1))
 
     def test_add_metadata(self):
+        """Test adding valid metadata to the SWCGraph."""
         self.tree.add_metadata("ORIGINAL_SOURCE", "file.swc")
         self.assertEqual(self.tree.metadata["ORIGINAL_SOURCE"], "file.swc")
 
     def test_invalid_metadata(self):
+        """Test that adding invalid metadata does not modify the metadata dictionary."""
         self.tree.add_metadata("INVALID_FIELD", "value")
         self.assertEqual(self.tree.metadata, {})
 
     def test_get_parent(self):
+        """Test getting the parent node of a given node."""
         self.assertIsNone(self.tree.get_parent(self.node1.id))
         self.assertEqual(self.tree.get_parent(self.node2.id), self.node1)
         self.assertEqual(self.tree.get_parent(self.node3.id), self.node2)
         with self.assertRaises(ValueError):
             self.tree.get_parent(4)
 
     def test_get_children(self):
+        """Test getting the children of a given node."""
         self.assertEqual(self.tree.get_children(self.node1.id), [self.node2])
         self.assertEqual(self.tree.get_children(self.node2.id), [self.node3])
         with self.assertRaises(ValueError):
             self.tree.get_parent(4)
 
     def test_get_nodes_with_multiple_children(self):
+        """Test getting nodes with multiple children."""
         node4 = SWCNode(4, 3, 3.0, 0.0, 0.0, 0.5, 2)
         self.tree.add_node(node4)
         self.assertEqual(self.tree.get_nodes_with_multiple_children(), [self.node2])
 
     def test_get_nodes_by_type(self):
+        """Test getting nodes by their type."""
         self.assertEqual(self.tree.get_nodes_by_type(1), [self.node1])
         self.assertEqual(self.tree.get_nodes_by_type(3), [self.node2, self.node3])
 
 
+class TestSWCExport(unittest.TestCase):
+    """Test cases for exporting SWC files."""
+
+    def setUp(self):
+        """Set up file paths for testing."""
+        current_dir = os.path.dirname(os.path.abspath(__file__))
+        self.input_file = os.path.join(current_dir, "Case1_new.swc")
+        self.output_file = os.path.join(current_dir, "Case1_exported.swc")
+
+    def test_load_export_compare(self):
+        """Test loading an SWC file, exporting it, and comparing the results."""
+        # Load the original file
+        original_tree = load_swc(self.input_file)
+
+        # Export the loaded tree
+        original_tree.export_to_swc_file(self.output_file)
+
+        # Check if the exported file was created
+        self.assertTrue(os.path.exists(self.output_file))
+
+        # Load the exported file
+        exported_tree = load_swc(self.output_file)
+
+        # Compare the original and exported trees
+        self.assertEqual(len(original_tree.nodes), len(exported_tree.nodes))
+
+        # Compare a few key properties of the first and last nodes
+        self.compare_nodes(original_tree.nodes[0], exported_tree.nodes[0])
+        self.compare_nodes(original_tree.nodes[-1], exported_tree.nodes[-1])
+
+        # Compare metadata
+        self.assertEqual(original_tree.metadata, exported_tree.metadata)
+
+    def compare_nodes(self, node1, node2):
+        """Compare two SWCNode objects for equality."""
+        self.assertEqual(node1.id, node2.id)
+        self.assertEqual(node1.type, node2.type)
+        self.assertEqual(node1.parent_id, node2.parent_id)
+        self.assertAlmostEqual(node1.x, node2.x, places=4)
+        self.assertAlmostEqual(node1.y, node2.y, places=4)
+        self.assertAlmostEqual(node1.z, node2.z, places=4)
+        self.assertAlmostEqual(node1.radius, node2.radius, places=4)
+
+    def tearDown(self):
+        """Clean up by removing the exported file if it exists."""
+        if os.path.exists(self.output_file):
+            os.remove(self.output_file)
+
+
 if __name__ == "__main__":
     unittest.main()