Fix extremely slow sparse_keys() call for unsorted sparse params

[idelder]

Fixes # .

Fix a very slow call from sparse_keys() to keys() for unsorted params.

Summary/Motivation:

The sparse_keys() function was updated to return ordered/sorted sparse indices when sorting was specified, going via the keys() function. The change, however, also routed unsorted parameters through the keys() function, collapsing to the final block where filtering was performed over the entire index cross product, keeping only sparse keys, but in index-insertion order. For very large but sparsely indexed params, this dramatically increased (by many orders of magnitude) the time to return the sparse keys set. For unsorted parameters, it makes sense to default back to old behaviour of just returning the unordered sparse keys immediately.

Changes proposed in this PR:

• When the sort argument is not set (defaults to UNSORTED) or explicitly set to UNSORTED, just return the unordered sparse keys set as in previous pyomo versions.
• When set to SORTED_INDICES, this will go via keys() to an explicit sorting of the sparse keys as was established in 6.10
• When set to ORDERED_INDICES, collapses to the final block in keys() where the sparse keys set is slowly returned in index insertion order.

AI-Use Disclosure

• AI tools were NOT used during the preparation of this PR

or

• AI tools contributed to the development of this PR

  • AI tools generated documentation (including the PR description/comments, code comments, and/or Sphinx documentation)
  • AI tools generated tests (baselines, examples, and/or code)
  • AI tools generated code (apart from tests)

  Review process (select ONE):

  • Rewritten: All AI-generated content was rewritten by me before being committed.
  • Reviewed/verified: I retained AI-generated content and verified it before committing. Verification included (as applicable):
    • Ran the code and fixed issues
    • Added and ran tests
    • Checked correctness/logic of code and tests
    • Checked for alignment with the contribution guide
    • Considered security implications
  • As-is: AI-generated content was commited directly to the repository

Notes for reviewers (optional):
Does this correctly cover the different cases of sorting?
(This is my first PR to the pyomo codebase, forgive me if I am missing something)

Legal Acknowledgement

By contributing to this software project, I have read the contribution guide and agree to the following terms and conditions for my contribution:

1. I agree my contributions are submitted under the BSD license.
2. I represent I am authorized to make the contributions and grant the license. If my employer has rights to intellectual property that includes these contributions, I represent that I have received permission to make contributions and grant the required license on behalf of that employer.

[jsiirola]

Initial thoughts on this PR:

• This potentially changes things so that keys from sparse_keys and keys come in a different order (for the same arguments). The current implementation can be slow because we are making sure that the order that you see keys is consistent between the two methods. This is a potentially "breaking change" that will need to be discussed (likely at a developer call)
• At a bare minimum, the ordering from sparse_keys(), sparse_values(), and sparse_items() must all match (for equivalent parameters) [this is for consistency with dict.{keys, values, items}]

[blnicho]

@idelder this is waiting on the second point from @jsiirola's comment to be addressed.

the ordering from sparse_keys(), sparse_values(), and sparse_items() must all match (for equivalent parameters)

So your change to sparse_keys needs to be propagated to those other "sparse" methods for consistency. I'm going to mark this PR as a draft for now. Please mark it as "ready for review" after you've made that change.

[idelder]

Okay done. The sparse_values() and sparse_items() calls get values and items via a keys() call anyway, so by moving the fix into the keys() function they automatically follow the same order.

Some testing of calls to sparse_keys() below, for an nxn sparse param with 100 randomly ordered diagonal entries:

Before this fix ~ O(N²) (gave up after n = 100,000)

n	sparse_keys() (ms)
100	0.214
1000	22.052
10000	2033.067
100000	199673.438

After this fix

n	sparse_keys() (ms)
100	0.016
1000	0.022
10000	0.07
100000	1.144
1000000	14.446
10000000	145.957
100000000	1907.136

Checking ordering

Checking 10_000 x 10_000 param for matching order...

(Check all 100 but print first three entries only)
Sparse keys: [(8101, 8101), (4421, 4421), (222, 222)]
Sparse values: [8101, 4421, 222]
Sparse items: [((8101, 8101), 8101), ((4421, 4421), 4421), ((222, 222), 222)]

PASS: all three iterators return elements in the same order.

Testing script

"""
Script to test:
1. Time to call param.sparse_keys() with a fixed number of entries
2. That sparse_keys(), sparse_values(), and sparse_items() return elements in the same order
"""

import time
import pyomo.environ as pyo
import random

SIZES = [100, 1_000, 10_000, 100_000, 1_000_000, 10_000_000, 100_000_000]
N_ENTRIES = 100  # number of diagonals that are populated

print("=" * 60)
print(f"Timing sparse_keys() with {N_ENTRIES} entries")
print(f"{'n':>10}  {'entries':>10}  {'sparse_keys() ms':>18}")
print("-" * 60)

def construct_param(n):
    rand_elements = random.sample(range(0, n), N_ENTRIES)
    data = {(i, i): i for i in rand_elements} # put some values down the diagonal

    m = pyo.ConcreteModel()
    m.I = pyo.Set(initialize=range(n))
    m.J = pyo.Set(initialize=range(n))
    m.P = pyo.Param(m.I, m.J, initialize=data, default=0.0)
    m.P.construct()

    return m, data

for n in SIZES:
    m, _ = construct_param(n)

    t0 = time.perf_counter()
    sparse_keys = list(m.P.sparse_keys())
    elapsed_ms = (time.perf_counter() - t0) * 1000

    print(f"{n:>10}  {N_ENTRIES:>10}  {elapsed_ms:>18.3f}")

m, data = construct_param(10_000)

print("\nChecking 10_000 x 10_000 param for matching order...\n")
match = all(
    [
        list(m.P.sparse_keys()) == list(data.keys()),
        list(m.P.sparse_values()) == list(data.values()),
        list(m.P.sparse_items()) == list(data.items())
    ]
)
print("Sparse keys: ", list(m.P.sparse_keys())[0:3])
print("Sparse values: ", list(m.P.sparse_values())[0:3])
print("Sparse items: ", list(m.P.sparse_items())[0:3])

if match:
    print("\nPASS: all three iterators return elements in the same order.")
else:
    print("\nFAIL: order mismatch detected.")

[jsiirola]

I have been debating this, and I think I agree this is a reasonable initial approach. You do need to update the condition and resolve the failing test (I believe that the test should be updated to iterate over a list (i.e., list(m.b)).

We should consider adding additional special-case handling for ORDERED_INDICES, too (but that can be a separate PR).

[jsiirola]

This needs to be either

        elif sort == SortComponents.UNSORTED:

or

        elif not sort:

(SortComponents.UNSORTED is 0, so SortComponents.UNSORTED in sort is always True)

[idelder]

I see okay that makes sense. I went with elif not sort

I'd also noticed the test_flatten bug but wasn't sure if I should mix that into this PR. Did that too.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 90.14%. Comparing base (e099f64) to head (5734120).

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #3971   +/-   ##
=======================================
  Coverage   90.14%   90.14%           
=======================================
  Files         917      917           
  Lines      109261   109263    +2     
=======================================
+ Hits        98494    98496    +2     
  Misses      10767    10767           

Flag	Coverage Δ
builders	29.13% <100.00%> (+0.01%)	⬆️
default	86.14% <100.00%> (?)
expensive	35.10% <100.00%> (?)
linux	87.67% <100.00%> (-1.98%)	⬇️
linux_other	87.67% <100.00%> (-0.01%)	⬇️
oldsolvers	28.05% <100.00%> (+<0.01%)	⬆️
osx	83.12% <100.00%> (+<0.01%)	⬆️
win	85.44% <100.00%> (+<0.01%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.