Correct format

Change-Id: I7b29fc4e4104c5afb567948232ce0c5be293cc8b

tests/test_extract_input.py

@@ -0,0 +1,45 @@
+from nose import tools as nt
+from nose.tools import assert_dict_equal
+import neurom
+import numpy as np
+from numpy.testing import assert_array_almost_equal
+from tns import extract_input
+import os
+
+_path = os.path.dirname(os.path.abspath(__file__))
+POP_PATH = os.path.join(_path, '../test_data/bio/')
+
+POPUL = neurom.load_neurons(POP_PATH)
+
+def test_num_trees():
+    target_numBAS = {'num_trees': {'data': {'bins': [4, 5, 6, 7, 8, 9],
+                                      'weights': [1, 0, 0, 0, 0, 1]}}}
+    target_numAX =  {'num_trees': {'data': {'bins': [1], 'weights': [2]}}}
+
+    numBAS = extract_input.from_neurom.number_neurites(POPUL)
+    numAX = extract_input.from_neurom.number_neurites(POPUL, neurite_type=neurom.AXON)
+    assert_dict_equal(numBAS, target_numBAS)
+    assert_dict_equal(numAX, target_numAX)
+
+def test_trunk_distr():
+    target_trunkBAS = {'trunk':
+                      {'azimuth': {'uniform': {'max': 0.0, 'min': np.pi}},
+                                   'orientation_deviation': {'data': {'bins': [0.19391773616376634,
+                                                                               0.4880704446023673,
+                                                                               0.7822231530409682,
+                                                                               1.076375861479569,
+                                                                               1.3705285699181702,
+                                                                               1.6646812783567713,
+                                                                               1.9588339867953721,
+                                                                               2.2529866952339734,
+                                                                               2.547139403672574,
+                                                                               2.841292112111175],
+                                                                       'weights': [4, 3, 1, 2, 0, 1, 0, 0, 0, 2]}}}}
+    target_trunkAPIC = {'trunk': {'azimuth': {'uniform': {'max': 0.0, 'min': np.pi}},
+                                  'orientation_deviation': {'data': {'bins': [0.0], 'weights': [2]}}}}
+
+    trunkAP = extract_input.from_neurom.trunk_neurite(POPUL, neurite_type=neurom.APICAL_DENDRITE, bins=1)
+    trunkBAS = extract_input.from_neurom.trunk_neurite(POPUL, bins=10)
+    assert_dict_equal(trunkBAS, target_trunkBAS)
+    assert_dict_equal(trunkAP, target_trunkAPIC)
+

tns/extract_input/from_neurom.py

@@ -29,18 +29,19 @@ def soma_data(pop):
     return {"size": transform_distr(ss)}
 
 
-def trunk_neurite(pop, neurite_type=nm.BASAL_DENDRITE):
+def trunk_neurite(pop, neurite_type=nm.BASAL_DENDRITE, bins=30):
     '''Extracts the trunk data for a specific tree type'''
 
     angles = [nm.get('trunk_angles', neuron, neurite_type=neurite_type) for neuron in pop]
     angles = [i for a in angles for i in a]
-    heights, bins = np.histogram(angles, bins=30)
+    heights, bins = np.histogram(angles, bins=bins)
 
     # Extract trunk relative orientations to reshample
+    actual_bins = (bins[1:] + bins[:-1]) / 2.
 
     return {"trunk": {"orientation_deviation": {"data":
-                                                {"bins": (bins[1:] + bins[:-1]) / 2.,
-                                                 "weights": heights}},
+                                                {"bins": actual_bins.tolist(),
+                                                 "weights": heights.tolist()}},
                       "azimuth": {"uniform": {"min": np.pi, "max": 0.0}}}}
 
 
@@ -53,5 +54,5 @@ def number_neurites(pop, neurite_type=nm.BASAL_DENDRITE):
     heights, bins = np.histogram(nneurites, bins=np.arange(np.min(nneurites),
                                                            np.max(nneurites) + 2))
 
-    return {"num_trees": {"data": {"bins": bins[:-1],
-                                   "weights": heights}}}
+    return {"num_trees": {"data": {"bins": bins[:-1].tolist(),
+                                   "weights": heights.tolist()}}}