Add support for lifted edges

[constantinpape]

Supporting lifted edges would enable better modeling of long range information across time. Lifted edges only provide a contribution to the cost, but not to the connectivity.

Motivation

We sometimes have knowledge about long-range connectivity across time in tracking problems. For example:

• Transformer-based tracking models like the GTR predict affinities between detections across multiple timeframes.
• In heterogeneous cell populations (e.g. mixed cell lines or different cells with specific stains) only the same kinds should be connected via tracks. Since the attribution is probabilistic (e.g. cell type classification) and may not be available in each frame (e.g. cell line classification not possible during divisions) long-range edges area a good way to model this.

Details

The concept of edges was introduced for the multicut graph partitioning problem in Hornakova et al., and is explained in a less formal (and more approachable) way in Chapter 2.26 of my thesis.
The key idea behind lifted edges is that they only contribute to the cost of the solution, but not to the connectivity. This fits the problems described in the motivation well, because we want nodes connected by an active lifted edge to be also connected by "local" edges, i.e. edges that connect detections in adjacent frames.
I made this figure to illustrate the concept for tracking and the difference to normal edges:

In the optimal solution for a normal edge the nodes are not connected by a local path across time, because the edge induces conductivity. In the optimal solution for the lifted edge the nodes are connected, because the lifted edge only contributes to the cost of the solution and does not induce conductivity.

(Note: I assume that negative costs are attractive here, which I think matches the convention of motile.)

Implementation

I am not familiar with the implementation of motile / ilpy yet, but would be happy to help in the implementation of this. I already discussed this with @funkey and @ben, who both expressed interest in this feature and in helping to implement it.
It would probably best if you could give some high-level ideas on how to implement this @bentaculum or @funkey and then we can think about how to tackle this in more detail.

[bentaculum]

Hi @constantinpape, sounds great, and yes to all of the above! These tutorials on extending motile should be a good starting point for the technical implementation, using indicator variables.

As for concrete desired API, here are two suggestions

1. to pass in such costs as a dictionary of edges (node_from, node_to) and values.

import motile
from motile import data

graph = data.arlo_graph()

# create a motile solver
solver = motile.Solver(graph)

from motile.costs import NodeSelection, EdgeSelection

solver.add_costs(
    motile.costs.NodeSelection(
        weight=-1.0,
        attribute='score'))
solver.add_costs(
    motile.costs.EdgeSelection(
        weight=-1.0,
        attribute='score'))

lifted_costs = {
    (node_u, node_v): cost,
    ...
} 

solver.add_costs(
    LiftedCosts(lifted_costs, weight, constant)
)

solution = solver.solve()

This would require that these node ids are consistent with the ones in TrackGraph, a bit inconvenient.

2. to define edges as "lifted" when creating the graph object. Here's the graph creation from the quickstart example, with an extra lifted edge

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.3},
        {'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},
    {'source': 0, 'target': 4, 'score': 0.9, 'lifted': True},
]

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)

Then same as above plus

solver.add_costs(
    LiftedCosts(attribute='lifted', weight, constant)
)

We would then have to make sure in the initial `TrackGraph` setup that lifted edges are not included.

[constantinpape]

Thanks for the pointers @bentaculum. And I think I prefer the second suggestion for the API, which I find more explicit.
Will have a closer look at the possible implementation soon and let you know what I find!