Proximal Optimal Transport Modeling of Population Dynamics

Consider a population of particles evolving with time, monitored through snapshots, using particles sampled within the population at successive timestamps. Given only access to these snapshots, can we reconstruct individual trajectories for these particles? This question arises in many crucial scientific challenges of our time, notably single-cell genomics. In this paper, we propose to model population dynamics as realizations of a causal Jordan-Kinderlehrer-Otto (JKO) flow of measures: The JKO scheme posits that the new configuration taken by a population at time t+1 is one that trades off a better configuration for the population, in the sense that it decreases an energy, while remaining close (in Wasserstein distance) to the previous configuration observed at t. Our goal in this work is to learn such an energy given data. To that end, we propose JKOnet, a neural architecture that computes (in end-to-end differentiable fashion) the JKO flow given a parametric energy and initial configuration of points. We demonstrate the good performance and robustness of the JKOnet fitting procedure, compared to a more direct forward method.