Perf: parallelize the construction of Jacobian and RHS-vector in Poisson solving problem#14
Conversation
|
|
||
| def _jflux_impl(nx, ny, nz, r, typ, sigma, eps, eps0, avgeps, with_index=False, | ||
| full_size=True) -> Tuple[np.ndarray, Optional[np.ndarray]]: | ||
| '''kernal implementation of the Jacobian calculation.''' |
There was a problem hiding this comment.
Variable name "kernal" should be "kernel" in the comment.
| '''kernal implementation of the Jacobian calculation.''' | |
| '''kernel implementation of the Jacobian calculation.''' |
|
Note Other AI code review bot(s) detectedCodeRabbit has detected other AI code review bot(s) in this pull request and will avoid duplicating their findings in the review comments. This may lead to a less comprehensive review. WalkthroughAdds a new Newton–Raphson Poisson solver module with flux and boundary helpers, exposes Changes
Sequence DiagramsequenceDiagram
participant SolverLoop as solve_poisson_NRcycle
participant NRConstruct as nr_construct
participant FluxImpl as Flux Helpers\n(_jflux_impl,_bflux_impl)
participant BoundImpl as Boundary Helpers\n(_impose_j_bound,_impose_b_bound)
participant LinearSolver as Linear Solver\n(scipy/pyamg)
SolverLoop->>NRConstruct: call(grid, phi, phi_, eps, charges, jinout?, binout?)
NRConstruct->>FluxImpl: compute jflux (interior stencil entries)
FluxImpl-->>NRConstruct: return per-neighbor flux arrays
NRConstruct->>BoundImpl: apply Dirichlet/Neumann masks to J and B
BoundImpl-->>NRConstruct: return masked/modified J and B
NRConstruct-->>SolverLoop: populate sparse Jacobian (J) and RHS (B)
SolverLoop->>LinearSolver: solve J · δφ = -B
LinearSolver-->>SolverLoop: return δφ
SolverLoop->>SolverLoop: compute norms, check convergence
alt RHS norm increased
SolverLoop->>SolverLoop: dampen δφ and rebuild via nr_construct
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
CC CLI report:
Verified Mathematical Equivalence
Performance Improvement The new vectorized implementation is significantly faster:
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
♻️ Duplicate comments (1)
dpnegf/negf/newton_raphson_speed_up.py (1)
63-63: Typo in docstring: "kernal" should be "kernel".This was already flagged by a previous reviewer.
🧹 Nitpick comments (7)
dpnegf/negf/poisson_init.py (4)
536-548: Consider using a function definition instead of a lambda assignment.The static analysis flags the lambda assignment (E731). Consider refactoring for clarity and consistency with Python style guidelines.
🔎 Proposed fix
- log.info(f'Solve Poisson equation by {method}') - solve = lambda jac, b: spsolve(jac, b) if method == 'scipy' \ - else self.solver_pyamg(jac, b, tolerance=1e-5) # hard-coded + log.info(f'Solve Poisson equation by {method}') + def solve(jac, b): + if method == 'scipy': + return spsolve(jac, b) + return self.solver_pyamg(jac, b, tolerance=1e-5) # hard-coded
538-541: Use exception chaining for better traceability.When re-raising an exception from within an
exceptblock, useraise ... from errto preserve the exception chain and improve debugging.🔎 Proposed fix
try: dtype = {'float64': np.float64, 'float32': np.float32}[dtype] - except KeyError: - raise ValueError(f"Unknown data type: {dtype}. Use 'float64' or 'float32'.") + except KeyError as err: + raise ValueError(f"Unknown data type: {dtype}. Use 'float64' or 'float32'.") from err
493-496: Unusedtoleranceparameter in method signature.The
toleranceparameter is declared but not used withinsolve_poisson_NRcycle. The hard-coded1e-5tolerance is used insolver_pyamgcall at line 548 instead.🔎 Proposed fix: use the parameter instead of hard-coding
- solve = lambda jac, b: spsolve(jac, b) if method == 'scipy' \ - else self.solver_pyamg(jac, b, tolerance=1e-5) # hard-coded + def solve(jac, b): + if method == 'scipy': + return spsolve(jac, b) + return self.solver_pyamg(jac, b, tolerance=tolerance)
691-768: Consider removing commented-out code before merging.The old implementation is preserved as comments. While this is useful during development and review, large blocks of commented-out code should typically be removed before merging, as version control preserves history.
dpnegf/tests/test_newton_raphson_speed_up.py (2)
62-68: Consider using logging instead of print statements in tests.Using
print()for timing output in tests can clutter CI logs. Consider usingloggingor making these timing statements conditional/optional.Also applies to: 131-137, 180-184, 230-236
170-170: Unused local variablenz.The variable
nzis unpacked but never used in these test methods. Consider using_as a placeholder if it's intentionally unused, or remove the unpacking.🔎 Proposed fix
- nx, ny, nz = self.nx, self.ny, self.nz + nx, ny, _ = self.nx, self.ny, self.nzAlso applies to: 214-214
dpnegf/negf/newton_raphson_speed_up.py (1)
180-180: Minor style issue: multiple statements on one line.Consider splitting for readability.
🔎 Proposed fix
- nx, ny, nz = grid_dim; nr = nx * ny * nz + nx, ny, nz = grid_dim + nr = nx * ny * nz
📜 Review details
Configuration used: defaults
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (4)
dpnegf/negf/newton_raphson_speed_up.pydpnegf/negf/poisson_init.pydpnegf/runner/NEGF.pydpnegf/tests/test_newton_raphson_speed_up.py
🧰 Additional context used
🧬 Code graph analysis (2)
dpnegf/runner/NEGF.py (1)
dpnegf/tests/test_poisson_init.py (1)
info(183-183)
dpnegf/negf/poisson_init.py (1)
dpnegf/negf/newton_raphson_speed_up.py (1)
nr_construct(109-251)
🪛 Ruff (0.14.8)
dpnegf/negf/newton_raphson_speed_up.py
9-9: Unused function argument: nz
(ARG001)
13-13: zip() without an explicit strict= parameter
Add explicit value for parameter strict=
(B905)
26-26: Unused function argument: nz
(ARG001)
30-30: zip() without an explicit strict= parameter
Add explicit value for parameter strict=
(B905)
47-47: Unused function argument: nz
(ARG001)
57-57: zip() without an explicit strict= parameter
Add explicit value for parameter strict=
(B905)
61-61: Unused function argument: nz
(ARG001)
91-91: zip() without an explicit strict= parameter
Add explicit value for parameter strict=
(B905)
177-177: Avoid specifying long messages outside the exception class
(TRY003)
179-179: Avoid specifying long messages outside the exception class
(TRY003)
180-180: Multiple statements on one line (semicolon)
(E702)
185-185: Avoid specifying long messages outside the exception class
(TRY003)
187-187: Avoid specifying long messages outside the exception class
(TRY003)
191-191: Avoid specifying long messages outside the exception class
(TRY003)
193-193: Avoid specifying long messages outside the exception class
(TRY003)
197-197: Avoid specifying long messages outside the exception class
(TRY003)
201-201: Avoid specifying long messages outside the exception class
(TRY003)
203-203: Avoid specifying long messages outside the exception class
(TRY003)
205-205: Avoid specifying long messages outside the exception class
(TRY003)
207-207: Avoid specifying long messages outside the exception class
(TRY003)
209-209: Avoid specifying long messages outside the exception class
(TRY003)
211-211: Avoid specifying long messages outside the exception class
(TRY003)
215-215: Avoid specifying long messages outside the exception class
(TRY003)
217-217: Avoid specifying long messages outside the exception class
(TRY003)
220-220: Avoid specifying long messages outside the exception class
(TRY003)
222-222: Avoid specifying long messages outside the exception class
(TRY003)
224-225: Avoid specifying long messages outside the exception class
(TRY003)
dpnegf/negf/poisson_init.py
495-495: Unused method argument: tolerance
(ARG002)
541-541: Within an except clause, raise exceptions with raise ... from err or raise ... from None to distinguish them from errors in exception handling
(B904)
541-541: Avoid specifying long messages outside the exception class
(TRY003)
545-545: Avoid specifying long messages outside the exception class
(TRY003)
547-548: Do not assign a lambda expression, use a def
Rewrite solve as a def
(E731)
dpnegf/tests/test_newton_raphson_speed_up.py
170-170: Local variable nz is assigned to but never used
Remove assignment to unused variable nz
(F841)
214-214: Local variable nz is assigned to but never used
Remove assignment to unused variable nz
(F841)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: run-tests
🔇 Additional comments (16)
dpnegf/runner/NEGF.py (3)
460-461: LGTM!Clean f-string formatting for the iteration log output. The explicit formatting with
:>12.8eprovides consistent alignment.
524-524: LGTM!Good use of f-string with list comprehension for consistent k-point tuple formatting.
578-578: LGTM!Consistent f-string formatting applied to energy and k-point logging throughout the file.
Also applies to: 702-704
dpnegf/negf/poisson_init.py (4)
1-17: LGTM!Appropriate imports added for timing instrumentation and the new
nr_constructfunction.
550-555: LGTM!Clear and informative iteration header for the Newton-Raphson solver output. Good use of structured logging with column alignment.
557-605: LGTM on the NR iteration loop with adaptive stability control.The control mechanism that monitors
norm_Band dampensdelta_phiwhen divergence is detected is a sensible addition for numerical stability. The timing instrumentation provides useful diagnostics.
667-689: LGTM!Clean delegation to
nr_constructwith proper parameter mapping. The permittivity averaging modes are correctly mapped from the object'saverage_modeattribute.dpnegf/tests/test_newton_raphson_speed_up.py (3)
1-18: LGTM!Good test setup with appropriate imports from the speed-up module and physical constants matching
poisson_init.py.
21-91: LGTM on boundary Jacobian imposition tests.Thorough testing of
_impose_j_boundagainst a manually constructed reference. The test validates diagonal and off-diagonal elements for all boundary types.
243-478: LGTM on comprehensive consistency test.Excellent test that reproduces the old implementation and compares it against
nr_construct. The test covers all three averaging modes (arithmetic, harmonic, geometric) and uses tight tolerances (rtol=1e-10,atol=1e-15).dpnegf/negf/newton_raphson_speed_up.py (6)
61-80: LGTM on vectorized flux implementation.Good use of NumPy vectorization for the flux calculations. The
full_sizeparameter provides flexibility for memory optimization.
82-94: LGTM on Jacobian assembly.Correct implementation of diagonal and off-diagonal element assignment with proper Boltzmann factor handling for free charge contribution.
96-107: LGTM on RHS vector assembly.Correct implementation with flux, free charge, and fixed charge contributions. The sign flip at the end is consistent with the Newton-Raphson formulation.
109-170: LGTM on comprehensive docstring.Excellent documentation with detailed parameter descriptions, types, and default values.
171-226: LGTM on input validation.Thorough validation of all input parameters with informative error messages. The handling of
gridpoint_typas either dict or array is a nice touch for flexibility.
227-251: LGTM on implementation section.Clean orchestration of the helper functions with appropriate assertions for shape validation.
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (1)
dpnegf/negf/newton_raphson_speed_up.py (1)
63-63: Fix typo: "kernal" should be "kernel".The docstring has a spelling error carried over from a previous review.
🔎 Proposed fix
- '''kernal implementation of the Jacobian calculation.''' + '''kernel implementation of the Jacobian calculation.'''
🧹 Nitpick comments (2)
dpnegf/negf/newton_raphson_speed_up.py (2)
14-14: Consider addingstrict=Truetozip()calls for safer iteration.The
zip()calls at lines 14, 31, 57, and 91 would benefit fromstrict=True(Python 3.10+) to catch length mismatches between paired sequences at runtime, preventing silent bugs if the displacement and type lists ever diverge.🔎 Proposed fix
- for t, d in zip(neumann_types_, displ_): + for t, d in zip(neumann_types_, displ_, strict=True):Apply similar changes to lines 31, 57, and 91.
Also applies to: 31-31, 57-57, 91-91
15-15: Minor efficiency: converttyponce before the loop.Line 15 recreates
np.array(typ)on every iteration. Sincetypis already validated as an ndarray innr_construct(line 230), this conversion is redundant. Consider converting once before the loop or passing a pre-converted array.🔎 Proposed fix
def _impose_j_bound(inout, nx, ny, nz, typ, val, mask): '''impose the special mask for boundary points''' + typ_arr = np.asarray(typ) # Neumann boundary types neumann_types_ = ['xmin', 'xmax', 'ymin', 'ymax', 'zmin', 'zmax'] displ_ = [+1, -1, +nx, -nx, +nx*ny, -nx*ny] for t, d in zip(neumann_types_, displ_): - i = np.where(np.array(typ) == t)[0] + i = np.where(typ_arr == t)[0] if len(i) == 0: log.warning(f'no grid point with type {t} found') # scipy sparse matrix also supports the parallelized assignment inout[i, i] = val inout[i, i + d] = mask # Dirichlet boundary - i = np.where(typ == 'Dirichlet')[0] + i = np.where(typ_arr == 'Dirichlet')[0] inout[i, i] = valSimilar pattern applies to
_impose_b_bound.
📜 Review details
Configuration used: defaults
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
dpnegf/negf/newton_raphson_speed_up.py
🧰 Additional context used
🪛 Ruff (0.14.8)
dpnegf/negf/newton_raphson_speed_up.py
9-9: Unused function argument: nz
(ARG001)
14-14: zip() without an explicit strict= parameter
Add explicit value for parameter strict=
(B905)
26-26: Unused function argument: nz
(ARG001)
31-31: zip() without an explicit strict= parameter
Add explicit value for parameter strict=
(B905)
47-47: Unused function argument: nz
(ARG001)
57-57: zip() without an explicit strict= parameter
Add explicit value for parameter strict=
(B905)
61-61: Unused function argument: nz
(ARG001)
91-91: zip() without an explicit strict= parameter
Add explicit value for parameter strict=
(B905)
177-177: Avoid specifying long messages outside the exception class
(TRY003)
179-179: Avoid specifying long messages outside the exception class
(TRY003)
180-180: Multiple statements on one line (semicolon)
(E702)
185-185: Avoid specifying long messages outside the exception class
(TRY003)
187-187: Avoid specifying long messages outside the exception class
(TRY003)
191-191: Avoid specifying long messages outside the exception class
(TRY003)
193-193: Avoid specifying long messages outside the exception class
(TRY003)
197-197: Avoid specifying long messages outside the exception class
(TRY003)
201-201: Avoid specifying long messages outside the exception class
(TRY003)
203-203: Avoid specifying long messages outside the exception class
(TRY003)
205-205: Avoid specifying long messages outside the exception class
(TRY003)
207-207: Avoid specifying long messages outside the exception class
(TRY003)
209-209: Avoid specifying long messages outside the exception class
(TRY003)
211-211: Avoid specifying long messages outside the exception class
(TRY003)
215-215: Avoid specifying long messages outside the exception class
(TRY003)
217-217: Avoid specifying long messages outside the exception class
(TRY003)
220-220: Avoid specifying long messages outside the exception class
(TRY003)
222-222: Avoid specifying long messages outside the exception class
(TRY003)
224-225: Avoid specifying long messages outside the exception class
(TRY003)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: run-tests
🔇 Additional comments (5)
dpnegf/negf/newton_raphson_speed_up.py (5)
41-45: LGTM! Spelling fix applied.The typo "coloumb" has been corrected to "coulomb". The implementation correctly converts charge density to Coulombs using the elementary charge constant.
47-59: LGTM!The flux calculation logic is correct for both full-size and indexed modes. The displacement signs correctly mirror the flux direction conventions used in
_jflux_impl.
82-94: LGTM!The Jacobian assembly correctly implements the finite-difference discretization with Boltzmann charge terms and boundary conditions. The sign conventions and displacement indexing are consistent with the flux calculations.
96-107: LGTM!The RHS vector assembly is mathematically consistent with the Jacobian. The sign flip on line 107 is correctly applied for the Newton-Raphson iteration (solving J·Δx = -B).
109-251: LGTM! Thorough validation and correct orchestration.The
nr_constructfunction provides a well-validated public API with:
- Comprehensive input type and shape checking
- Defensive handling of dict vs. array for
gridpoint_typ(lines 198-199)- Correct orchestration of flux calculation and matrix assembly
- Appropriate short-circuit for null inputs (lines 172-173)
The extensive validation is appropriate for a public API surface despite triggering style linters about long exception messages.
| raise ValueError("grid_dim must be a tuple or list of three integers") | ||
| if any(not isinstance(d, int) or d <= 0 for d in grid_dim): | ||
| raise ValueError("grid_dim must contain positive integers") | ||
| nx, ny, nz = grid_dim; nr = nx * ny * nz |
There was a problem hiding this comment.
Minor style: avoid semicolon for multiple statements.
Line 180 uses a semicolon to place multiple statements on one line, which violates PEP 8 style guidelines.
🔎 Proposed fix
- nx, ny, nz = grid_dim; nr = nx * ny * nz
+ nx, ny, nz = grid_dim
+ nr = nx * ny * nz📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| nx, ny, nz = grid_dim; nr = nx * ny * nz | |
| nx, ny, nz = grid_dim | |
| nr = nx * ny * nz |
🧰 Tools
🪛 Ruff (0.14.8)
180-180: Multiple statements on one line (semicolon)
(E702)
🤖 Prompt for AI Agents
In dpnegf/negf/newton_raphson_speed_up.py around line 180, the statement uses a
semicolon to combine two assignments; split it into two separate lines to comply
with PEP8: first unpack grid_dim into nx, ny, nz on one line, then compute nr =
nx * ny * nz on the following line (no semicolon).
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (2)
dpnegf/tests/test_poisson_init.py (2)
433-453: Misleading test name: no actual comparison of results.The test name claims to verify that different averaging modes produce different results, but:
- The final loop (lines 452-453) only checks shapes, not differences between results.
- The comment at lines 450-451 states that modes will give the same result for uniform
eps.This test effectively validates that all modes execute without error, not that they produce different outputs.
🔎 Suggested improvements
Option 1: Rename to reflect actual behavior
-def test_interface3d_different_average_modes_produce_different_results(): - """Test that different averaging modes produce different Jacobian matrices.""" +def test_interface3d_all_average_modes_work(): + """Test that all averaging modes execute without error."""Option 2: Add actual difference check with non-uniform epsilon
If you want to verify modes produce different results, use non-uniform dielectric regions and add assertions comparing the Jacobian values:
# After line 453, add: harmonic = results['harmonic'] arithmetic = results['arithmetic'] geometric = results['geometric'] # For non-uniform eps, modes should differ assert not np.allclose(harmonic, arithmetic), "Expected harmonic ≠ arithmetic" assert not np.allclose(arithmetic, geometric), "Expected arithmetic ≠ geometric"
505-516: Misleading test comment: both J and B are always computed.The docstring claims to test calling
nr_constructwith only Jacobian (binout=None), butto_scipy_Jac_B()always computes bothJandB. The test effectively validates that the scipy pathway produces non-empty outputs.🔎 Suggested fix
-def test_interface3d_NR_construct_only_jacobian(): - """Test that nr_construct can be called with only Jacobian (binout=None internally).""" +def test_interface3d_to_scipy_Jac_B_produces_nonempty_output(): + """Test that to_scipy_Jac_B produces valid non-empty Jacobian and RHS."""
📜 Review details
Configuration used: defaults
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
dpnegf/tests/test_poisson_init.py
🧰 Additional context used
🧬 Code graph analysis (1)
dpnegf/tests/test_poisson_init.py (1)
dpnegf/negf/poisson_init.py (8)
Grid(19-165)Dirichlet(192-216)Dielectric(218-242)Interface3D(244-768)get_potential_eps(394-436)NR_construct_Jac_B(643-768)solve_poisson_NRcycle(493-608)to_scipy_Jac_B(466-491)
🪛 Ruff (0.14.10)
dpnegf/tests/test_poisson_init.py
452-452: Loop control variable mode not used within loop body
(B007)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: run-tests
🔇 Additional comments (7)
dpnegf/tests/test_poisson_init.py (7)
360-369: LGTM: Clean helper for internal-point test cases.The 3×3×3 grid with corner-only atoms ensures at least one internal grid point exists, enabling tests that exercise flux calculations on non-boundary points.
372-378: LGTM: Default averaging mode verified.Test correctly validates that
eps_average_modedefaults to'harmonic'when not explicitly specified.
381-389: LGTM: Averaging mode options validated.Test correctly verifies all three averaging modes (
'harmonic','arithmetic','geometric') are accepted and stored.Note: The static analysis hint flagging
modeas unused at line 452 is a false positive—the variable is used in the loop body.
392-402: LGTM: Internal point detection verified.Test correctly validates that the 3×3×3 grid produces at least one internal point and that
internal_NPmatches the count of"in"entries inboudnary_points.
405-430: LGTM: Internal vs. boundary Jacobian diagonal handling verified.Test correctly validates that:
- Internal points have negative diagonals (flux sum + charge term).
- Boundary points have diagonal = 1.0 (Dirichlet enforcement).
This exercises the core vectorized NR construction logic.
456-469: LGTM: dtype handling validated.Test correctly verifies that both
float32andfloat64dtypes work without error and return numeric results.
472-502: LGTM: scipy solver and convergence validated.Both tests correctly exercise the scipy-based NR cycle:
test_interface3d_solve_poisson_NRcycle_scipy_methodvalidates basic operation.test_interface3d_solve_poisson_NRcycle_convergencetests with non-zero Dirichlet potential and checks finite convergence.
2 participants
It is expected to bring the speed-up for x20-50 on Jacobian and RHS-vector construction
Summary by CodeRabbit
New Features
Improvements
Tests
✏️ Tip: You can customize this high-level summary in your review settings.