[WIP] Vectorized check_parallel_jacobian implementation that checks sparsity pattern

[Robbybp]

Inspired by your matmul-based implementation, I tried to re-create my original algorithm doing all computations in NumPy/SciPy. The main difference is that the original algorithm only compares vector-pairs that have the same (approximate) nonzero pattern. This branch contains calls to HierarchicalTimer to record timing information for profiling purposes.

Here is a script to profile on a Pyomo.DAE distillation column model:

from pyomo.environ import *
from pyomo.dae import *
from distill_DAE import model

instance = model.create_instance('distill.dat')

# Discretize using Finite Difference Approach
discretizer = TransformationFactory('dae.finite_difference')
discretizer.apply_to(instance, nfe=1000, scheme='BACKWARD')

from idaes.core.util.model_diagnostics import check_parallel_jacobian
from pyomo.common.timing import TicTocTimer, HierarchicalTimer
from pyomo.contrib.pynumero.interfaces.pyomo_nlp import PyomoNLP

instance.obj = Objective(expr=0)
nlp = PyomoNLP(instance)
jac = nlp.evaluate_jacobian()

print(f"Jacobian shape: {jac.shape}")

timer = TicTocTimer()
htimer = HierarchicalTimer()

timer.tic()
htimer.start("root")
htimer.start("vectorized")
check_parallel_jacobian(None, nlp=nlp, jac=jac, method="vectorized", timer=htimer)
htimer.stop("vectorized")
timer.toc("vectorized")
htimer.start("new")
check_parallel_jacobian(None, nlp=nlp, jac=jac, method="vectorized-with-filter", timer=htimer)
htimer.stop("new")
timer.toc("new")
htimer.stop("root")

print(htimer)

My results:

Jacobian shape: (100068, 101069)
[    0.00] Resetting the tic/toc delta timer
[+   0.10] vectorized
[+   0.73] new
Identifier                                  ncalls   cumtime   percall      %
-----------------------------------------------------------------------------
root                                             1     0.829     0.829  100.0
     ------------------------------------------------------------------------
     new                                         1     0.733     0.733   88.4
               --------------------------------------------------------------
               compare-nonzero-pattern           1     0.273     0.273   37.3
               compare-to-tolerance              1     0.000     0.000    0.0
               construct-initial-matrices        1     0.045     0.045    6.1
               differences                       1     0.000     0.000    0.0
               dot-products                      1     0.023     0.023    3.1
               extract-indices                   1     0.000     0.000    0.0
               filter-small-entries              1     0.364     0.364   49.6
               norms                             1     0.001     0.001    0.1
               upper-tri                         1     0.026     0.026    3.6
               other                           n/a     0.001       n/a    0.2
               ==============================================================
     vectorized                                  1     0.096     0.096   11.6
               --------------------------------------------------------------
               compare-to-tolerance              1     0.001     0.001    1.1
               dot-products                      1     0.013     0.013   13.2
               initial-matrices                  1     0.031     0.031   32.6
               inv-norms                         1     0.000     0.000    0.1
               norms                             1     0.001     0.001    1.1
               scale-by-norms                    1     0.018     0.018   18.4
               upper-tri                         1     0.031     0.031   31.9
               other                           n/a     0.001       n/a    1.5
               ==============================================================
     other                                     n/a     0.000       n/a    0.0
     ========================================================================
=============================================================================

And here is a script to profile on my singular-MBR example:

from idaes_examples.mod.diagnostics.gas_solid_contactors.example import (
    create_square_model_with_new_variable_and_constraint,
    create_corrected_square_model,
)
from idaes.core.util.model_diagnostics import (
    DiagnosticsToolbox,
    check_parallel_jacobian,
)

m = create_square_model_with_new_variable_and_constraint()
#m = create_corrected_square_model()
#dt = DiagnosticsToolbox(m)

#dt.report_structural_issues()
#dt.report_numerical_issues()
#dt.display_near_parallel_constraints()

from pyomo.contrib.pynumero.interfaces.pyomo_nlp import PyomoNLP
from pyomo.common.timing import TicTocTimer, HierarchicalTimer
nlp = PyomoNLP(m)
jac = nlp.evaluate_jacobian()

htimer = HierarchicalTimer()
timer = TicTocTimer()
timer.tic()

htimer.start("root")
htimer.start("vectorized")
parallel = check_parallel_jacobian(None, nlp=nlp, jac=jac, method="vectorized", timer=htimer)
print(f"N. parallel = {len(parallel)}")
htimer.stop("vectorized")
timer.toc("vectorized")
htimer.start("new")
parallel = check_parallel_jacobian(None, nlp=nlp, jac=jac, method="vectorized-with-filter", timer=htimer)
print(f"N. parallel = {len(parallel)}")
htimer.stop("new")
timer.toc("new")
htimer.stop("root")

print(htimer)

My results:

[    0.00] Resetting the tic/toc delta timer
N. parallel = 1012
[+   0.04] vectorized
N. parallel = 11
[+   0.12] new
Identifier                                  ncalls   cumtime   percall      %
-----------------------------------------------------------------------------
root                                             1     0.159     0.159  100.0
     ------------------------------------------------------------------------
     new                                         1     0.123     0.123   77.6
               --------------------------------------------------------------
               compare-nonzero-pattern           1     0.050     0.050   40.1
               compare-to-tolerance              1     0.000     0.000    0.0
               construct-initial-matrices        1     0.005     0.005    3.7
               differences                       1     0.000     0.000    0.0
               dot-products                      1     0.018     0.018   14.2
               extract-indices                   1     0.000     0.000    0.0
               filter-small-entries              1     0.038     0.038   30.5
               norms                             1     0.000     0.000    0.2
               upper-tri                         1     0.014     0.014   11.0
               other                           n/a     0.000       n/a    0.3
               ==============================================================
     vectorized                                  1     0.036     0.036   22.4
               --------------------------------------------------------------
               compare-to-tolerance              1     0.000     0.000    1.4
               dot-products                      1     0.004     0.004   11.6
               initial-matrices                  1     0.003     0.003    9.2
               inv-norms                         1     0.000     0.000    0.1
               norms                             1     0.000     0.000    0.9
               scale-by-norms                    1     0.009     0.009   24.3
               upper-tri                         1     0.018     0.018   49.3
               other                           n/a     0.001       n/a    3.2
               ==============================================================
     other                                     n/a     0.000       n/a    0.0
     ========================================================================
=============================================================================

This implementation is still a factor of 3-10 slower than your vectorized implementation, but it does significantly reduce the bottleneck due to sorting and comparing nonzero structures compared to the original, pure-Python implementation. I have not spent much time optimizing compare-nonzero-pattern and filter-small-entries, so further improvements may be possible.

What do you think about this algorithm, if it is possible to reduce the run time a bit more?

[dallan-keylogic]

The algorithm adds considerable complexity and index management. I don't fully understand what's going on it, especially the compare-nonzero-pattern section. Is the intention to reduce the runtime further, is it to reduce the number of parallel entries that are presented to the user to the most important ones, or is it to apply some different criteria for parallelness than the one that my method has implemented?

I had assumed that the following model might get past it as not having parallel rows because one row has a column entry where the other does not:

        m = ConcreteModel()

        m.v1 = Var(initialize=1e-8)
        m.v2 = Var()

        m.c1 = Constraint(expr=m.v1 == 0)
        m.c2 = Constraint(expr=m.v1 + 1e-8 * m.v2 == 0)

However, the algorithm correctly detects that the two rows are almost parallel:

(Pdb) check_parallel_jacobian(m, direction="row") == [(m.c1, m.c2)]
True

So if there's an example of two rows/columns for which the two implementations give different answers, I'd love to see it.

[Robbybp]

These give different results on the singular MBR example above. The purpose of only comparing same-nonzero-pattern vectors was initially performance, but it actually appears to be irrelevant for performance, in practice. However, based on the MBR example, only comparing vectors with the same nonzero pattern appears to give fewer false positives. This is why I'd like to find a solution that considers sparsity patterns when deciding which pairs of vectors can be parallel. It also might be useful to have an option that allows users to decide whether they want to restrict pairs to same-sparsity-pattern or not.

[dallan-keylogic]

I don't see the additional constraint pairs returned by the unfiltered method as "false positives". Through the geometric interpretation, those constraints are "nearly parallel".

See this example:

import numpy as np
from numpy.linalg import norm
indices = nlp.get_constraint_indices([model2.fs.MB.solid_super_vel[0.0], model2.fs.MB.density_flowrate_constraint[0.0,0.2]])
print(jac.getrow(indices[0]))
print("========================")
print(jac.getrow(indices[1]))
u = jac.getrow(indices[0]).todense()
v = jac.getrow(indices[1]).todense()
print("========================")
np.arccos(u.dot(v.T)/(norm(u)*norm(v)))

which returns

  (0, 1508)	0.1326566278709479
  (0, 1497)	148139.4724110699
  (0, 1000)	17.847046656540535
========================
  (0, 3510)	-1.0
  (0, 1633)	0.1326566278709479
  (0, 1497)	146006.08981376316
  (0, 1000)	17.590028198659624
========================
matrix([[6.96679884e-06]])

So MB.solid_super_vel[0.0] has a tiny contribution from 1508, density_flowrate_constraint[0.0,0.2] has a tiny contribution from variables 1633 and 3510, but their contributions from variables 1000 and 1497 are nearly identical. Accordingly, the angle between them is 7e-6 radians.

I still don't understand what's going on in the sparsity matching algorithm and why it says these two constraints are "not parallel" but that the two constraints in my toy model are parallel.

[dallan-keylogic]

There's good news, though. After using the autoscalers provided in Andrew's PR (AutoScaler.variables_by_magnitude and constraints_by_jacobian_norm), the unfiltered method only returns the 11 constraints you want.

Funnily enough, if you autoscale before the second round of optimization, IPOPT now converges where it diverged before. That goes to show the strength of IPOPT's regularization in the face of degeneracy.

[Robbybp]

There's good news, though. After using the autoscalers provided in Andrew's PR (AutoScaler.variables_by_magnitude and constraints_by_jacobian_norm), the unfiltered method only returns the 11 constraints you want.

Interesting. I wouldn't be opposed to scaling the model in the example, and using the sparsity-pattern-agnostic implementation. What parallel component tolerance did you use there?

Funnily enough, if you autoscale before the second round of optimization, IPOPT now converges where it diverged before. That goes to show the strength of IPOPT's regularization in the face of degeneracy.

Thanks for the info, I will look into this. Maybe I have to make the optimization problem harder.

[dallan-keylogic]

I also have an SVD-based method to diagnose extra degrees of freedom that I tried to apply to your model. Unfortunately, all that's done is reveal the inadequacies of the existing autoscaling. Any variable that can pass through zero doesn't get scaled well at present (so basically all derivative variables and accumulation terms), and enthalpy is a perennial bugbear. For example, enthalpy flow terms have magnitudes of 533902 J/s and 533921 J/s at 0.9 and 1 in the length domain, but when you take the difference to calculate heat transfer, you end up with a 20 J/s difference.