Improvements to JuMP code per @IainNZ.