Module rfm 

Rectified Flow Matching (RFM) helpers.

In the FM literature, “rectification” is largely about the coupling (\pi(x_0,x_1)): choosing (or improving) which base samples are paired with which target samples so that trajectories become straighter / less “curly”.

This module provides small, testable coupling utilities built on top of wass (OT).

Functions

apply_pairing

Apply the selected pairing strategy to a minibatch of source and target points.

greedy_bipartite_match_from_weights

Build a greedy one-to-one matching (a permutation) from a nonnegative n×n weight matrix.

minibatch_exp_greedy_pairing

Fast minibatch pairing: convert costs to weights via exp(-cost / temp) then greedy match.

minibatch_ot_greedy_pairing

Compute a minibatch OT coupling between two equal-sized point clouds x and y, then return a greedy one-to-one matching derived from the transport plan.

minibatch_ot_greedy_pairing_normalized

Like minibatch_ot_greedy_pairing but normalizes the cost matrix before Sinkhorn.

minibatch_ot_selective_pairing

Partial Sinkhorn pairing:

minibatch_partial_rowwise_pairing

Partial pairing heuristic:

minibatch_rowwise_nearest_pairing

Fast minibatch pairing: greedy row-wise nearest neighbor on the Euclidean cost matrix.