load_asc Python fix (#1977)

arbor/morph/segment_tree.cpp

@@ -26,7 +26,7 @@ node_p yes = [](node_t) { return true; };
 std::map<msize_t, std::vector<msize_t>> tree_to_children(const segment_tree& tree) {
     const auto& parents = tree.parents();
     std::map<msize_t, std::vector<msize_t>> result;
-    for (auto ix = 0; ix < tree.size(); ++ix) result[parents[ix]].push_back(ix);
+    for (msize_t ix = 0; ix < tree.size(); ++ix) result[parents[ix]].push_back(ix);
     for (auto& [k, v]: result) std::sort(v.begin(), v.end());
     return result;
 }
@@ -105,7 +105,7 @@ equivalent(const segment_tree& a,
         auto bs = fetch_children(b_cursor, b.segments(), b_children_of);
         todo.pop_back();
         if (as.size() != bs.size()) return false;
-        for (auto ix = 0; ix < as.size(); ++ix) {
+        for (msize_t ix = 0; ix < as.size(); ++ix) {
             if ((as[ix].prox != bs[ix].prox) ||
                 (as[ix].dist != bs[ix].dist) ||
                 (as[ix].tag != bs[ix].tag)) return false;

arborio/include/arborio/neurolucida.hpp

@@ -45,8 +45,10 @@ struct asc_morphology {
 
 // Perform the parsing of the input as a string.
 ARB_ARBORIO_API asc_morphology parse_asc_string(const char* input);
+ARB_ARBORIO_API arb::segment_tree parse_asc_string_raw(const char* input);
 
 // Load asc morphology from file with name filename.
 ARB_ARBORIO_API asc_morphology load_asc(std::string filename);
+ARB_ARBORIO_API arb::segment_tree load_asc_raw(std::string filename);
 
 } // namespace arborio

arborio/neurolucida.cpp

@@ -610,7 +610,7 @@ parse_hopefully<sub_tree> parse_sub_tree(asc::lexer& L) {
 
 
 // Perform the parsing of the input as a string.
-asc_morphology parse_asc_string(const char* input) {
+ARB_ARBORIO_API arb::segment_tree parse_asc_string_raw(const char* input) {
     asc::lexer lexer(input);
 
     std::vector<sub_tree> sub_trees;
@@ -768,6 +768,14 @@ asc_morphology parse_asc_string(const char* input) {
         }
     }
 
+    return stree;
+}
+
+
+ARB_ARBORIO_API asc_morphology parse_asc_string(const char* input) {
+    // Parse segment tree
+    arb::segment_tree stree = parse_asc_string_raw(input);
+
     // Construct the morphology.
     arb::morphology morphology(stree);
 
@@ -781,7 +789,8 @@ asc_morphology parse_asc_string(const char* input) {
     return {stree, std::move(morphology), std::move(labels)};
 }
 
-ARB_ARBORIO_API asc_morphology load_asc(std::string filename) {
+
+inline std::string read_file(std::string filename) {
     std::ifstream fid(filename);
 
     if (!fid.good()) {
@@ -796,9 +805,21 @@ ARB_ARBORIO_API asc_morphology load_asc(std::string filename) {
 
     fstr.assign((std::istreambuf_iterator<char>(fid)),
                  std::istreambuf_iterator<char>());
+    return fstr;
+}
 
+
+ARB_ARBORIO_API asc_morphology load_asc(std::string filename) {
+    std::string fstr = read_file(filename);
     return parse_asc_string(fstr.c_str());
 }
 
+
+ARB_ARBORIO_API arb::segment_tree load_asc_raw(std::string filename) {
+    std::string fstr = read_file(filename);
+    return parse_asc_string_raw(fstr.c_str());
+}
+
+
 } // namespace arborio
 

python/morphology.cpp

@@ -356,18 +356,25 @@ void register_morphology(py::module& m) {
             [](const arborio::asc_morphology& m) {return label_dict_proxy(m.labels);},
             "The four canonical regions are labeled 'soma', 'axon', 'dend' and 'apic'.");
 
+    using asc_morph_or_tree = std::variant<arb::segment_tree, arborio::asc_morphology>;
+
     m.def("load_asc",
-        [](py::object fn) {
+        [](py::object fn, bool raw) -> asc_morph_or_tree {
             try {
                 auto contents = util::read_file_or_buffer(fn);
-                return arborio::load_asc(contents);
+                if (raw) {
+                    return arborio::parse_asc_string_raw(contents.c_str());
+                }
+                return arborio::parse_asc_string(contents.c_str());
             }
             catch (std::exception& e) {
                 // Try to produce helpful error messages for SWC parsing errors.
                 throw pyarb_error(util::pprintf("error loading neurolucida asc file: {}", e.what()));
             }
         },
-        "filename"_a, "Load a morphology and meta data from a Neurolucida ASCII .asc file.");
+        "filename_or_stream"_a,
+        pybind11::arg_v("raw", false, "Return a segment tree instead of a fully formed morphology"),
+        "Load a morphology or segment_tree and meta data from a Neurolucida ASCII .asc file.");
 
 
 #ifdef ARB_NEUROML_ENABLED

python/test/unit/test_io.py

@@ -5,6 +5,8 @@ import arbor as A
 from pathlib import Path
 from tempfile import TemporaryDirectory as TD
 from io import StringIO
+from functools import partial
+
 
 acc = """(arbor-component
   (meta-data
@@ -62,32 +64,141 @@ acc = """(arbor-component
 """
 
 
-class TestAccIo(unittest.TestCase):
-    def test_stringio(self):
-        sio = StringIO(acc)
-        A.load_component(sio)
+swc_arbor = """1 1 -5.0 0.0 0.0 5.0 -1
+2 1 0.0 0.0 0.0 5.0 1
+3 1 5.0 0.0 0.0 5.0 2
+"""
 
-    def test_fileio(self):
-        fn = "test.acc"
+
+swc_neuron = """1 1 0.1 0.2 0.3 0.4 -1
+"""
+
+
+asc = """((CellBody)\
+ (0 0 0 4)\
+)\
+((Dendrite)\
+ (0 2 0 2)\
+ (0 5 0 2)\
+ (\
+  (-5 5 0 2)\
+  |\
+  (6 5 0 2)\
+ )\
+)\
+((Axon)\
+ (0 -2 0 2)\
+ (0 -5 0 2)\
+ (\
+  (-5 -5 0 2)\
+  |\
+  (6 -5 0 2)\
+ )\
+)
+"""
+
+
+def load_string(loaders, morph_str):
+    for loader in loaders:
+        sio = StringIO(morph_str)
+        loader(sio)
+
+
+def load_file(loaders, morph_str, morph_fn):
+    for loader in loaders:
         with TD() as tmp:
             tmp = Path(tmp)
-            with open(tmp / fn, "w") as fd:
-                fd.write(acc)
-            with open(tmp / fn) as fd:
-                A.load_component(fd)
+            with open(tmp / morph_fn, "w") as fd:
+                fd.write(morph_str)
+            with open(tmp / morph_fn) as fd:
+                loader(fd)
 
-    def test_nameio(self):
-        fn = "test.acc"
+
+def load_name(loaders, morph_str, morph_fn):
+    for loader in loaders:
         with TD() as tmp:
             tmp = Path(tmp)
-            with open(tmp / fn, "w") as fd:
-                fd.write(acc)
-            A.load_component(str(tmp / fn))
+            with open(tmp / morph_fn, "w") as fd:
+                fd.write(morph_str)
+            loader(str(tmp / morph_fn))
 
-    def test_pathio(self):
-        fn = "test.acc"
+
+def load_pathio(loaders, morph_str, morph_fn):
+    for loader in loaders:
         with TD() as tmp:
             tmp = Path(tmp)
-            with open(tmp / fn, "w") as fd:
-                fd.write(acc)
-            A.load_component(tmp / fn)
+            with open(tmp / morph_fn, "w") as fd:
+                fd.write(morph_str)
+            loader(tmp / morph_fn)
+
+
+class TestAccIo(unittest.TestCase):
+    @staticmethod
+    def loaders():
+        return (A.load_component,)
+
+    def test_stringio(self):
+        load_string(self.loaders(), acc)
+
+    def test_fileio(self):
+        load_file(self.loaders(), acc, "test.acc")
+
+    def test_nameio(self):
+        load_name(self.loaders(), acc, "test.acc")
+
+    def test_pathio(self):
+        load_pathio(self.loaders(), acc, "test.acc")
+
+
+class TestSwcArborIo(unittest.TestCase):
+    @staticmethod
+    def loaders():
+        return (A.load_swc_arbor, partial(A.load_swc_arbor, raw=True))
+
+    def test_stringio(self):
+        load_string(self.loaders(), swc_arbor)
+
+    def test_fileio(self):
+        load_file(self.loaders(), swc_arbor, "test.swc")
+
+    def test_nameio(self):
+        load_name(self.loaders(), swc_arbor, "test.swc")
+
+    def test_pathio(self):
+        load_pathio(self.loaders(), swc_arbor, "test.swc")
+
+
+class TestSwcNeuronIo(unittest.TestCase):
+    @staticmethod
+    def loaders():
+        return (A.load_swc_neuron, partial(A.load_swc_neuron, raw=True))
+
+    def test_stringio(self):
+        load_string(self.loaders(), swc_neuron)
+
+    def test_fileio(self):
+        load_file(self.loaders(), swc_neuron, "test.swc")
+
+    def test_nameio(self):
+        load_name(self.loaders(), swc_neuron, "test.swc")
+
+    def test_pathio(self):
+        load_pathio(self.loaders(), swc_neuron, "test.swc")
+
+
+class TestAscIo(unittest.TestCase):
+    @staticmethod
+    def loaders():
+        return (A.load_asc, partial(A.load_asc, raw=True))
+
+    def test_stringio(self):
+        load_string(self.loaders(), asc)
+
+    def test_fileio(self):
+        load_file(self.loaders(), asc, "test.asc")
+
+    def test_nameio(self):
+        load_name(self.loaders(), asc, "test.asc")
+
+    def test_pathio(self):
+        load_pathio(self.loaders(), asc, "test.asc")