Viterbi Algorithm for Connecting Fragments in Long-Range Paths

The Viterbi Algorithm is a dynamic programming approach to finding an optimal path between a starting state and a goal state. For our neuron reconstruction problem, the states are defined as fragments, while the traversals are the connections made between endpoints using evidence observed from image intensity data. This notebook illustrates a simple example of the Viterbi Algorithm connecting synthetic fragments in a 10x10 grid example.

[1]:
import numpy as np
from brainlit.algorithms.connect_fragments.tests.grid_generator import grid_gen
from brainlit.algorithms.connect_fragments.dynamic_programming_viterbi import viterbi_algorithm
import matplotlib.pyplot as plt
/opt/buildhome/python3.7/lib/python3.7/site-packages/tqdm/auto.py:22: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html
  from .autonotebook import tqdm as notebook_tqdm
---------------------------------------------------------------------------
ModuleNotFoundError                       Traceback (most recent call last)
<ipython-input-1-aeab7efaa678> in <module>
      1 import numpy as np
----> 2 from brainlit.algorithms.connect_fragments.tests.grid_generator import grid_gen
      3 from brainlit.algorithms.connect_fragments.dynamic_programming_viterbi import viterbi_algorithm
      4 import matplotlib.pyplot as plt

ModuleNotFoundError: No module named 'brainlit.algorithms.connect_fragments.tests.grid_generator'

Running a 10x10 Grid Example

[2]:
# Load the 10x10 example from grid_generator
img, lbls, _ , somas = grid_gen(10)
plt.figure()
plt.title("Image data")
plt.imshow(img[:,:,0])
plt.figure()
plt.title("Pre-generated Labels")
plt.imshow(lbls[:,:,0])

# Initiate the algorithm class
alg = viterbi_algorithm(img, lbls, somas, [1,1,1])
print("Soma label: ", somas)
---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
<ipython-input-2-7a15ed9b4288> in <module>
      1 # Load the 10x10 example from grid_generator
----> 2 img, lbls, _ , somas = grid_gen(10)
      3 plt.figure()
      4 plt.title("Image data")
      5 plt.imshow(img[:,:,0])

NameError: name 'grid_gen' is not defined
[3]:
# Manually identify the endpoints. Note that Labels 2 and 5 are "blobs" and do not have endpoints
endpoints = {}
endpoints[1] = ((0,0,0),(0,2,0))
endpoints[3] = ((3,3,0),(6,5,0))
endpoints[4] = ((7,7,0),(7,8,0))
# Assign the endpoints
alg.end_points = endpoints
---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
<ipython-input-3-d10bf5a39ca2> in <module>
      5 endpoints[4] = ((7,7,0),(7,8,0))
      6 # Assign the endpoints
----> 7 alg.end_points = endpoints

NameError: name 'alg' is not defined

Running Viterbi Algorithm

Within the Viterbi class object, we first need to compute the distance matrices. After that, we can run the Viterbi algorithm to find the best path.

[4]:
def print_path(alg, path):
    c = alg.connection_mat
    path_lbls = path[1]
    for i in range(len(path_lbls)-1):
        from_lbl = path_lbls[i]
        to_lbl = path_lbls[i+1]

        print(f"From {from_lbl} to {to_lbl}: {c[0][from_lbl][to_lbl]}, {c[1][from_lbl][to_lbl]}")
[5]:
alg.compute_all_dists()
top_path, sorted_paths = alg.viterbi_frag(1, K=4, somas=alg.somas)
print(top_path)
print_path(alg, top_path)
---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
<ipython-input-5-eb9e1927d59a> in <module>
----> 1 alg.compute_all_dists()
      2 top_path, sorted_paths = alg.viterbi_frag(1, K=4, somas=alg.somas)
      3 print(top_path)
      4 print_path(alg, top_path)

NameError: name 'alg' is not defined

From the path found, we can observe that the algorithm traversed from the closest endpoints on each fragment to the goal state soma with label 5, starting from fragment label 1. We can also observe that the algorithm chose to ignore the blob fragment label 2, as it is not a soma state nor is it biologically sound to have a blob as an intermediate state.