Extending motile

written by Jan Funke

Download this page as a Jupyter notebook

motile ships with some basic tracking variables, constraints, and costs (see the API Reference for a complete list). In some situations, you might want to extend motile to fit your specific application. This tutorial will show you how to implement new constraints, variables, and costs.

Consider the following example track graph, where each node is a potential object and each edge a potential link of objects between frames:

(click here to see how to create the example track graph)

import motile
import networkx as nx

cells = [
        {"id": 0, "t": 0, "x": 1, "score": 0.8},
        {"id": 1, "t": 0, "x": 25, "score": 0.1},
        {"id": 2, "t": 1, "x": 0, "score": 0.3},
        {"id": 3, "t": 1, "x": 26, "score": 0.4},
        {"id": 4, "t": 2, "x": 2, "score": 0.6},
        {"id": 5, "t": 2, "x": 24, "score": 0.6},
        {"id": 6, "t": 2, "x": 35, "score": 0.7}
]

edges = [
    {"source": 0, "target": 2, "score": 0.9},
    {"source": 1, "target": 3, "score": 0.9},
    {"source": 0, "target": 3, "score": 0.5},
    {"source": 1, "target": 2, "score": 0.5},
    {"source": 2, "target": 4, "score": 0.7},
    {"source": 3, "target": 5, "score": 0.7},
    {"source": 2, "target": 5, "score": 0.3},
    {"source": 3, "target": 4, "score": 0.3},
    {"source": 3, "target": 6, "score": 0.8}
]

graph = nx.DiGraph()
graph.add_nodes_from([
    (cell["id"], cell)
    for cell in cells
])
graph.add_edges_from([
    (edge["source"], edge["target"], edge)
    for edge in edges
])

graph = motile.TrackGraph(graph)

from motile_toolbox.visualization import draw_track_graph

draw_track_graph(graph, alpha_attribute="score", label_attribute="score")

We will start with a simple model: the costs for selecting nodes and edges should depend on their score attribute, and the start of a track should have a positive cost as well. Furthermore, we are only interested in single-object tracks, i.e., tracks don’t merge or split. The solver for this model is:

from motile.constraints import MaxParents, MaxChildren
from motile.costs import NodeSelection, EdgeSelection, Appear

solver = motile.Solver(graph)

solver.add_constraint(MaxParents(1))
solver.add_constraint(MaxChildren(1))

solver.add_cost(NodeSelection(weight=-1, attribute="score"))
solver.add_cost(EdgeSelection(weight=-1, attribute="score"))
solver.add_cost(Appear(constant=1))

solver.solve()

…and our initial solution looks like this:

from motile_toolbox.visualization import draw_solution

draw_solution(graph, solver, label_attribute="score")

Adding Constraints

New constraints are introduced by subclassing Constraint and implementing the instantiate method. This method should return a list of ilpy.Constraint.

Imagine we know precisely that we want to track at most \(k\) objects, but we don’t know beforehand which of the many objects in the track graph those are. We are thus interested in finding at most \(k\) tracks with minimal costs.

This can be done with a constraint as follows:

import ilpy
from motile.variables import NodeAppear


class LimitNumTracks(motile.constraints.Constraint):

  def __init__(self, num_tracks):
      self.num_tracks = num_tracks

  def instantiate(self, solver):

      appear_indicators = solver.get_variables(NodeAppear)

      constraint = ilpy.Constraint()
      for appear_indicator, index in appear_indicators.items():
        constraint.set_coefficient(index, 1.0)
      constraint.set_relation(ilpy.Relation.LessEqual)
      constraint.set_value(self.num_tracks)

      return [constraint]

The instantiate method gets access to the solver the constraint is added to. Through the solver, we can then access variables to formulate constraints on them. Here, we make use of the NodeAppear variables: Simply counting how many nodes are marked as “appear” gives us the number of selected tracks. We then require that this number be less than or equal to num_tracks.

In other words, we formulated a constraint by enforcing a certain relation to hold between a subset of variables. Note that we didn’t have to do anything special to create the node appear variables: we simply ask the solver for them, and if they haven’t been added to the solver already, they will be added automatically.

Once implemented, the constraint can simply be added to the solver by instantiating the new class:

solver.add_constraint(LimitNumTracks(1))
solver.solve()

draw_solution(graph, solver, label_attribute="score")

Indeed, we have limited the solution to contain at most one track.

Adding Costs

We might want to add custom costs to motile that are not already covered by the existing ones. For the sake of this tutorial, let’s say we want to make the selection of nodes cheaper the higher up the node is in space. For obvious reasons, let’s call this new cost SillyCost.

Costs in motile are added by subclassing Cost and implementing the apply method:

from motile.variables import NodeSelected


class SillyCost(motile.costs.Cost):

    def __init__(self, position_attribute, weight=1.0):
        self.position_attribute = position_attribute
        self.weight = motile.costs.Weight(weight)

    def apply(self, solver):
        node_indicators = solver.get_variables(NodeSelected)

        for node_indicator, index in node_indicators.items():

            # x position of the node
            x = solver.graph.nodes[node_indicator][self.position_attribute]

            # costs can be negative (i.e., a reward)
            costs = -x

            solver.add_variable_cost(index, costs, weight=self.weight)

Similar to adding constraints, we simply ask the solver for the variables we want to add costs for and add them to the inference problem via Solver.add_variable_cost.

We can now add those costs in the same way others are added, i.e.:

print("Before adding silly cost:")
print(solver.get_variables(NodeSelected))

solver.add_cost(SillyCost('x', weight=0.02))

print("After adding silly cost:")
print(solver.get_variables(NodeSelected))

Before adding silly cost:
NodeSelected(0): cost=-0.800000011920929 value=1.0
NodeSelected(1): cost=-0.10000000149011612 value=0.0
NodeSelected(2): cost=-0.30000001192092896 value=1.0
NodeSelected(3): cost=-0.4000000059604645 value=0.0
NodeSelected(4): cost=-0.6000000238418579 value=1.0
NodeSelected(5): cost=-0.6000000238418579 value=0.0
NodeSelected(6): cost=-0.699999988079071 value=0.0
After adding silly cost:
NodeSelected(0): cost=-0.8199999928474426 value=1.0
NodeSelected(1): cost=-0.6000000238418579 value=0.0
NodeSelected(2): cost=-0.30000001192092896 value=1.0
NodeSelected(3): cost=-0.9199999570846558 value=0.0
NodeSelected(4): cost=-0.64000004529953 value=1.0
NodeSelected(5): cost=-1.0800000429153442 value=0.0
NodeSelected(6): cost=-1.399999976158142 value=0.0

As we can see, our new cost have been applied to each NodeSelected variable and nodes with a larger x value are now cheaper than others.

Solving again will now select the upper one of the possible tracks in the track graph:

solver.solve()

draw_solution(graph, solver, label_attribute="score")

Adding Variables

Variables in motile are added by subclassing the Variable class. Subclasses need to implement at least a static method instantiate. This method should return keys for what kind of things we would like to create a variable for. This should be a list of anything that is hashable (the keys will be used in a dictionary).

If, for example, we would like to create a new variable for each edge in the track graph, we would simply return a list of tuples, each tuple representing an edge (graph.edges would fit that, in this case).

Variables can also be introduced for other things that are not nodes or edges. We could create variables only for a subset of nodes, for pairs of edges, or the number of division events.

Let’s say we want to add new variables for pairs of edges that leave and enter the same node. Having such a variable might be useful to, for example, measure the curvature of a track and put a cost on that.

To create our new variables, we simply return a list of all pairs of edges we wish to have a variable for in instantiate. However, declaring those variables alone is not sufficient. To give them semantic meaning, we also have to make sure that our edge-pair variables are set to 1 if the two edges they represent have been selected, and 0 otherwise. To that end, we also add constraints that are specific to our variables by overriding the instantiate_constraints method, such that our variables are linked to the already existing EdgeSelected variables.

The complete variable declaration looks like this:

import ilpy
from motile.variables import EdgeSelected


class EdgePairs(motile.variables.Variable):

  @staticmethod
  def instantiate(solver):

      edge_pairs = [
          (in_edge, out_edge)
          # for each node
          for node in solver.graph.nodes
          # for each pair of incoming and outgoing edge
          for in_edge in solver.graph.prev_edges[node]
          for out_edge in solver.graph.next_edges[node]
      ]

      return edge_pairs

  @staticmethod
  def instantiate_constraints(solver):

      edge_indicators = solver.get_variables(EdgeSelected)
      edge_pair_indicators = solver.get_variables(EdgePairs)

      constraints = []
      for (in_edge, out_edge), pair_index in edge_pair_indicators.items():

          in_edge_index = edge_indicators[in_edge]
          out_edge_index = edge_indicators[out_edge]

          # edge pair indicator = 1 <=> in edge = 1 and out edge = 1
          constraint = ilpy.Constraint()
          constraint.set_coefficient(pair_index, 2)
          constraint.set_coefficient(in_edge_index, -1)
          constraint.set_coefficient(out_edge_index, -1)
          constraint.set_relation(ilpy.Relation.LessEqual)
          constraint.set_value(0)
          constraints.append(constraint)

          constraint = ilpy.Constraint()
          constraint.set_coefficient(pair_index, -1)
          constraint.set_coefficient(in_edge_index, 1)
          constraint.set_coefficient(out_edge_index, 1)
          constraint.set_relation(ilpy.Relation.LessEqual)
          constraint.set_value(1)
          constraints.append(constraint)

      return constraints

Variables on their own, however, don’t do anything yet. They only start to affect the solution if they are involved in constraints or have a cost.

The following defines a cost on our new variables, which loosely approximate the local curvature of the track:

class CurvatureCost(motile.costs.Cost):

    def __init__(self, position_attribute, weight=1.0):
        self.position_attribute = position_attribute
        self.weight = motile.costs.Weight(weight)

    def apply(self, solver):

        # get edge pair variables
        edge_pair_indicators = solver.get_variables(EdgePairs)

        for (in_edge, out_edge), index in edge_pair_indicators.items():

            in_offset = self.get_edge_offset(solver.graph, in_edge)
            out_offset = self.get_edge_offset(solver.graph, out_edge)

            curvature_cost = abs(out_offset - in_offset)

            solver.add_variable_cost(index, curvature_cost, self.weight)

    def get_edge_offset(self, graph, edge):

        pos_v = graph.nodes[edge[1]][self.position_attribute]
        pos_u = graph.nodes[edge[0]][self.position_attribute]

        return pos_v - pos_u

As before, we don’t have to do anything special to add the variables: we simply ask the solver for them, and if they don’t exist yet, they will be instantiated.

We can now add those costs to our solvers, which in turn makes use of our newly added variables:

solver.add_cost(CurvatureCost('x', weight=0.1))
solver.solve()

Let’s inspect the solution!

print(solver.get_variables(EdgePairs))
draw_solution(graph, solver, label_attribute="score")

EdgePairs(((0, 2), (2, 4))): cost=0.30000001192092896 value=0.0
EdgePairs(((0, 2), (2, 5))): cost=2.5 value=0.0
EdgePairs(((1, 2), (2, 4))): cost=2.700000047683716 value=0.0
EdgePairs(((1, 2), (2, 5))): cost=4.900000095367432 value=0.0
EdgePairs(((0, 3), (3, 5))): cost=2.700000047683716 value=0.0
EdgePairs(((0, 3), (3, 4))): cost=4.900000095367432 value=0.0
EdgePairs(((0, 3), (3, 6))): cost=1.600000023841858 value=0.0
EdgePairs(((1, 3), (3, 5))): cost=0.30000001192092896 value=1.0
EdgePairs(((1, 3), (3, 4))): cost=2.5 value=0.0
EdgePairs(((1, 3), (3, 6))): cost=0.800000011920929 value=0.0