PAVA for partially-ordered (lattice) x-axis

Question

Given an MxN matrix A, I want to find the nearest matrix B (least squares fit: ie minimize$\displaystyle\sum_{i,j} (B_{ij} - A_{ij})^2$) such that

$$ \forall i_1,j_1,i_2,j_2 : \\ i_1 \leq i_2 \land j_1 \leq j_2 \implies B_{i_1,j_1} \leq B_{i_2,j_2} $$

I'd rather not just throw a quadratic program solver at this as it shows up as a sub-problem in projecting arguments into the feasible set for another problem that I'm already using a nonlinear optimization solver for. Having to run one within another will likely be prohibitively slow.

This problem is essentially the 2D generalization of isotonic regression which admits the PAVA algorithm. I was hoping that might generalize to 2D, but I don't see how. Naively, just creating pools of equivalent values doesn't work, since eg. you could have a case where $A_{1,1}$ and $A_{1,2}$ violate the constraint, and separately $A_{1,1}$ and $A_{2,1}$ violate the constraint, but that doesn't mean $B_{1,2}$ should need to be equal to $B_{2,1}$

Answer 1

It looks like my worries were unfounded. Using an off-the-shelf quadratic-program solver, I get a solution ridiculously fast (for the sizes I care about, on the order of a tens of microseconds to about a millisecond or two and that's without even putting much effort into optimizing) so I can just do this as part of the projection function.

use std::borrow::Cow;
use ndarray::{Array1, Array2, ArrayView2};
use osqp::{CscMatrix, Problem};
pub fn nearest_2d_ascending(m: ArrayView2<f64>, settings: osqp::Settings) -> Array2<f64> {
    let (r, c) = m.dim();
    let nvars = r * c;
    // The quadratic part of the function being optimized is just the identity matrix
    let p = CscMatrix {
        nrows: nvars,
        ncols: nvars,
        indptr: Cow::Owned(Vec::from_iter(0..nvars + 1)),
        indices: Cow::Owned(Vec::from_iter(0..nvars)),
        data: Cow::Owned(vec![1.0; nvars]),
    };
    // The function being optimized is 0.5 * x^T * p * x + q^T * x. Since we care about the squared distance to our
    // input matrix, we can set q = -m (which will optimize for half the squared difference. We can ignore the
    // constant term).
    let neg_m = -m.as_standard_layout().into_owned();
    let q = neg_m.as_slice().unwrap();
    let mut sparse_constraints = vec![vec![]; nvars];
    let mut nconstraints: usize = 0;
    // Each of our constraints is of the form -x_i1_j1 + x_i2_j2 >= 0 so our within each rowcoefficients are going to
    // be in pairs of -1 and 1.
    for i in 0..r {
        for j in 0..c {
            if i < r - 1 {
                sparse_constraints[i * c + j].push((nconstraints, -1.0));
                sparse_constraints[(i + 1) * c + j].push((nconstraints, 1.0));
                nconstraints += 1;
            }
            if j < c - 1 {
                sparse_constraints[i * c + j].push((nconstraints, -1.0));
                sparse_constraints[i * c + (j + 1)].push((nconstraints, 1.0));
                nconstraints += 1;
            }
        }
    }
    // Unfortunately, the solver requires column-major order for the sparse matrix, so we need to transpose the
    // constraints
    let mut indptr = vec![];
    let mut indices = vec![];
    let mut data = vec![];
    let mut entry_ind = 0;
    for var in sparse_constraints {
        indptr.push(entry_ind);
        for &(index, value) in &var {
            indices.push(index);
            data.push(value);
        }
        entry_ind += var.len();
    }
    indptr.push(entry_ind);
    let a = CscMatrix {
        nrows: nconstraints,
        ncols: nvars,
        indptr: Cow::Owned(indptr),
        indices: Cow::Owned(indices),
        data: Cow::Owned(data),
    };
    let l = vec![0.0; nconstraints];
    let u = vec![f64::INFINITY; nconstraints];
let mut prob = Problem::new(p, q, a, &amp;l, &amp;u, &amp;settings).unwrap();
let result = prob.solve();
let Some(x) = result.x() else {
    panic!(&quot;No solution found: {:?}&quot;, result)
};

Array1::from_iter(x.iter().copied())
    .into_shape_with_order((r, c))
    .unwrap()

}