⚡ Optimize correlator loading in operator_loader.py#81
Conversation
This commit optimizes the `correlators()` method in `OperatorModule` to eliminate N+1 query-like inefficiencies in nested loops. Key improvements: - Restructured loops to iterate over correlator types first, minimizing factory function lookups. - Implemented a mapping dictionary for O(1) factory function retrieval, replacing long `if/elif` chains. - Pre-calculates index pair metadata (`sites` and `key`) once per pair instead of for every correlator type. - Eliminated redundant `kwargs.copy()` calls by passing parameters directly to factory functions. - Optimized both `SPIN_1_2` and `SPIN_1` local space types. These changes result in a measurable reduction in Python-level overhead, especially for large systems with many index pairs. Benchmarks show ~1.4x speedup in operator generation logic. Co-authored-by: makskliczkowski <48489493+makskliczkowski@users.noreply.github.com>
|
👋 Jules, reporting for duty! I'm here to lend a hand with this pull request. When you start a review, I'll add a 👀 emoji to each comment to let you know I've read it. I'll focus on feedback directed at me and will do my best to stay out of conversations between you and other bots or reviewers to keep the noise down. I'll push a commit with your requested changes shortly after. Please note there might be a delay between these steps, but rest assured I'm on the job! For more direct control, you can switch me to Reactive Mode. When this mode is on, I will only act on comments where you specifically mention me with New to Jules? Learn more at jules.google/docs. For security, I will only act on instructions from the user who triggered this task. |
There was a problem hiding this comment.
Pull request overview
This PR optimizes OperatorLoader.correlators() by reducing per-operator overhead during correlator operator instantiation, which is used across the algebra/operator layer (and by NQS code paths that request correlator kernels).
Changes:
- Replaces inner-loop
if/elifcorrelator dispatch with a pre-built correlator→factory mapping. - Precomputes
(key, sites)information for each(i, j)pair once and reuses it across correlator types. - Avoids per-operator
kwargs.copy()and builds correlator dictionaries via comprehensions.
Comments suppressed due to low confidence (1)
Python/QES/Algebra/Operator/operator_loader.py:498
- Same issue as in the spin-1/2 branch: calling
factory(sites=sites, **kwargs)will error ifkwargscontainssites. Previously any caller-providedsiteswas overwritten by the localkwargs_op["sites"] = sites, so this should be made robust by popping/overridingsitesbefore invoking the operator factory.
ops = {}
for corr in correlators:
factory = mapping.get(corr)
if factory is None:
raise ValueError(
f"Unsupported spin-1 correlator type '{corr}'. Supported: 'xx', 'yy', 'zz'."
)
ops[corr] = {key: factory(sites=sites, **kwargs) for key, sites in pairs_info}
return ops
| ops = {} | ||
| for corr in correlators: | ||
| for i, j in indices_pairs: | ||
| sites = [i, j] if i is not None and j is not None else None | ||
| key = (i, j) if sites is not None else "i,j" | ||
| kwargs_op = kwargs.copy() | ||
| kwargs_op["sites"] = sites | ||
|
|
||
| if "xx" == corr: | ||
| op = ops_module.sig_x(**kwargs_op) | ||
| elif "xy" == corr: | ||
| op = ops_module.sig_xy(**kwargs_op) | ||
| elif "xz" == corr: | ||
| op = ops_module.sig_xz(**kwargs_op) | ||
| elif "yx" == corr: | ||
| op = ops_module.sig_yx(**kwargs_op) | ||
| elif "yy" == corr: | ||
| op = ops_module.sig_y(**kwargs_op) | ||
| elif "yz" == corr: | ||
| op = ops_module.sig_yz(**kwargs_op) | ||
| elif "zx" == corr: | ||
| op = ops_module.sig_zx(**kwargs_op) | ||
| elif "zy" == corr: | ||
| op = ops_module.sig_zy(**kwargs_op) | ||
| elif "zz" == corr: | ||
| op = ops_module.sig_z(**kwargs_op) | ||
| else: | ||
| raise ValueError(f"Unknown correlator type: {corr}") | ||
| ops[corr][key] = op | ||
| factory = mapping.get(corr) | ||
| if factory is None: | ||
| raise ValueError(f"Unknown correlator type: {corr}") | ||
|
|
||
| ops[corr] = {key: factory(sites=sites, **kwargs) for key, sites in pairs_info} | ||
| return ops |
There was a problem hiding this comment.
Fixed in 7121f8e. I now remove any caller-provided sites from kwargs before correlator construction, so the computed pair-specific sites always applies and factory(sites=..., **kwargs) no longer raises duplicate-keyword errors.
There was a problem hiding this comment.
I have further refactored the correlators method to unify the optimization logic across both SPIN_1_2 and SPIN_1 types. The loops and pre-calculations are now shared, with only the factory mapping being type-specific. This significantly reduces redundancy and improves maintainability while retaining the 1.4x performance boost.
Agent-Logs-Url: https://github.com/makskliczkowski/pyqusolver/sessions/89875efe-47a5-4f6c-a46c-2f533e0e7947 Co-authored-by: makskliczkowski <48489493+makskliczkowski@users.noreply.github.com>
Refactored the `correlators()` method to share the optimized loop structure and pre-calculations between different local space types. Key changes: - Unified logic for `SPIN_1_2` and `SPIN_1` to reduce code duplication. - Extracted factory mappings to be type-specific while keeping core generation logic shared. - Maintained the 1.4x performance improvement by preserving loop restructuring and batch instantiation. - Cleaned up auxiliary benchmarking scripts. Co-authored-by: makskliczkowski <48489493+makskliczkowski@users.noreply.github.com>
The
correlators()method inoperator_loader.pycontained nested loops that performed redundant operations like string comparisons, dictionary lookups, and dictionary copying for every single operator instantiation. By restructuring the loops and pre-calculating loop-invariant data, the overhead is significantly reduced.Specifically:
siteslist and dictionarykeyare pre-calculated for each index pair once.kwargs.copy()is avoided in favor of direct keyword argument passing.Verification was performed using standalone benchmark and unit test scripts with mocks to account for the restricted environment. Benchmark results showed a 1.4x speedup for 100 sites (10,000 pairs).
PR created automatically by Jules for task 16919638002894485980 started by @makskliczkowski