Skip to content

self energy and e3 model fix#52

Closed
Lonya0 wants to merge 3 commits into
DeePTB-Lab:mainfrom
Lonya0:main
Closed

self energy and e3 model fix#52
Lonya0 wants to merge 3 commits into
DeePTB-Lab:mainfrom
Lonya0:main

Conversation

@Lonya0

@Lonya0 Lonya0 commented May 13, 2026

Copy link
Copy Markdown
Contributor

Summary by CodeRabbit

  • New Features

    • Configurable parallel execution backend for self-energy computation and safer temporary-file lifecycle with atomic shard merging.
    • Optional override for overlap data in electronic-structure calculations.
  • Tests

    • Added an end-to-end NEGF test with reference structure and configuration and a transmission-value assertion.
    • New test input structure and JSON scenario files added.
  • Chores

    • Updated dependency constraint for dptb.

@AsymmetryChou AsymmetryChou self-assigned this May 13, 2026
@coderabbitai

coderabbitai Bot commented May 13, 2026

Copy link
Copy Markdown

Review Change Stack

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info
⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 693cab6e-1d1f-4baa-81d0-78062a2cc631

📥 Commits

Reviewing files that changed from the base of the PR and between b8bc3b1 and ff16e73.

📒 Files selected for processing (2)
  • dpnegf/tests/data/test_negf/test_negf_e3/negf_7_0.json
  • pyproject.toml
💤 Files with no reviewable changes (1)
  • dpnegf/tests/data/test_negf/test_negf_e3/negf_7_0.json
✅ Files skipped from review due to trivial changes (1)
  • pyproject.toml

📝 Walkthrough

Walkthrough

The PR adds configurable Joblib backends and temp-shard handling to parallel self-energy computation, propagates an override_overlap option through ElecStruCal and the NEGF runner, extends the argument schema with parallel_options, updates a dependency constraint, and adds end-to-end test fixtures and a validation test.

Changes

Parallel self-energy execution with override_overlap integration

Layer / File(s) Summary
Temp shard aggregation and parallel backend in lead_property.py
dpnegf/negf/lead_property.py
compute_all_self_energy now accepts a parallel_backend parameter (default "loky"), normalizes self_energy_save_path to an absolute path, creates a temp_dir for shard outputs, runs workers with Joblib using backend=parallel_backend, merges shard HDF5 files into final self_energy_leadL.h5/self_energy_leadR.h5 with error handling, removes temp_dir on success, and updates self_energy_worker to write shards into temp_dir. Imports for shutil and tempfile were added.
Configuration schema for parallel execution
dpnegf/utils/argcheck.py
Adds parallel_options() schema (fields: n_jobs, batch_size, backend) and extends negf() with override_overlap and parallel_options keys.
Override_overlap parameter and kwargs in ElecStruCal
dpnegf/utils/elec_struc_cal.py
ElecStruCal.__init__ accepts **kwargs and stores override_overlap. get_data and get_eigs accept override_overlap: Union[str,bool,None], normalize FalseNone, use provided string when truthy, and fallback to instance state when None.
NEGF runner integration of parallel and override options
dpnegf/runner/NEGF.py
NEGF.__init__ reads override_overlap and parallel_options from kwargs (defaults: n_jobs=-1, batch_size=200, backend="loky"). Fermi-level calculations construct ElecStruCal with override_overlap. prepare_self_energy converts save path to absolute and forwards n_jobs, batch_size, and parallel_backend to compute_all_self_energy for both SCF+Dirichlet and non-SCF code paths.
Test fixtures and end-to-end NEGF validation
dpnegf/tests/data/test_negf/test_negf_e3/7_0.vasp, dpnegf/tests/data/test_negf/test_negf_e3/negf_7_0.json, dpnegf/tests/test_negf_e3.py
Adds a 896-atom carbon VASP input, a NEGf scenario JSON with solver and parallel options, and a Pytest that runs the entrypoint and asserts on generated results and a transmission value.
Poetry dependency update
pyproject.toml
Bumps dptb dependency constraint from >=2.1.0 to >=2.2.0.

Sequence Diagram

sequenceDiagram
  participant NEGF
  participant prepare_self_energy
  participant compute_all_self_energy
  participant Joblib as Joblib (parallel_backend)
  participant self_energy_worker
  participant merge as HDF5 merge
  NEGF->>prepare_self_energy: Invoke with parallel_options
  prepare_self_energy->>compute_all_self_energy: Pass n_jobs, batch_size, parallel_backend
  compute_all_self_energy->>compute_all_self_energy: Create temp_dir
  compute_all_self_energy->>Joblib: Execute workers with backend
  Joblib->>self_energy_worker: Run per (k, e) point
  self_energy_worker->>self_energy_worker: Write shard HDF5 to temp_dir
  compute_all_self_energy->>merge: Aggregate shards
  merge->>merge: Finalize leadL.h5 / leadR.h5
  compute_all_self_energy->>compute_all_self_energy: Remove temp_dir
Loading

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

Possibly related PRs

Suggested reviewers

  • AsymmetryChou

Poem

A rabbit hops through temp_dir's lair,
Backends chatter in parallel air,
Shards gather, then merge with a cheer,
Override flags whisper: "I am here."
🐇✨

🚥 Pre-merge checks | ✅ 3 | ❌ 2

❌ Failed checks (1 warning, 1 inconclusive)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 58.33% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
Title check ❓ Inconclusive The title 'self energy and e3 model fix' is vague and generic, failing to clearly specify what was actually fixed or changed in the self-energy computation or e3 model. Replace with a more specific title such as 'Add configurable Joblib backend and temp directory management for self-energy computation' or 'Implement parallel backend configuration and safer temp file handling in NEGF self-energy computation'.
✅ Passed checks (3 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.

✏️ Tip: You can configure your own custom pre-merge checks in the settings.

✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests

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.

❤️ Share

Comment @coderabbitai help to get the list of available commands and usage tips.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 4

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)
dpnegf/utils/elec_struc_cal.py (1)

153-172: ⚠️ Potential issue | 🟡 Minor | ⚡ Quick win

Use a context manager for h5py.File to avoid leaked handles on exceptions.

Line 153 opens the file and Line 172 closes it manually; any exception in between can skip close.

Suggested fix
-            overlap_blocks = h5py.File(override_overlap, "r")
-            if len(overlap_blocks) != 1:
-                log.info('Overlap file contains more than one overlap matrix, only first will be used.')
-            if self.overlap:
-                log.warning('override_overlap is enabled while model contains overlap, override_overlap will be used.')
-            if "0" in overlap_blocks:
-                overlaps = overlap_blocks["0"]
-            else:
-                overlaps = overlap_blocks["1"]
-            block_to_feature(data, self.model.idp, blocks=False, overlap_blocks=overlaps)
-            if not self.overlap:
-                self.eigv = Eigenvalues(
-                    idp=self.model.idp,
-                    device=self.device,
-                    s_edge_field=AtomicDataDict.EDGE_OVERLAP_KEY,
-                    s_node_field=AtomicDataDict.NODE_OVERLAP_KEY,
-                    s_out_field=AtomicDataDict.OVERLAP_KEY,
-                    dtype=self.model.dtype,
-                )
-            overlap_blocks.close()
+            with h5py.File(override_overlap, "r") as overlap_blocks:
+                if len(overlap_blocks) != 1:
+                    log.info('Overlap file contains more than one overlap matrix, only first will be used.')
+                if self.overlap:
+                    log.warning('override_overlap is enabled while model contains overlap, override_overlap will be used.')
+                overlaps = overlap_blocks["0"] if "0" in overlap_blocks else overlap_blocks["1"]
+                block_to_feature(data, self.model.idp, blocks=False, overlap_blocks=overlaps)
+                if not self.overlap:
+                    self.eigv = Eigenvalues(
+                        idp=self.model.idp,
+                        device=self.device,
+                        s_edge_field=AtomicDataDict.EDGE_OVERLAP_KEY,
+                        s_node_field=AtomicDataDict.NODE_OVERLAP_KEY,
+                        s_out_field=AtomicDataDict.OVERLAP_KEY,
+                        dtype=self.model.dtype,
+                    )
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@dpnegf/utils/elec_struc_cal.py` around lines 153 - 172, The code opens an
h5py.File into overlap_blocks and closes it manually, risking leaked file
handles on exceptions; wrap the open in a context manager (with
h5py.File(override_overlap, "r") as overlap_blocks:) and move the logic that
reads overlap_blocks (the log/info checks, the "0"/"1" selection, call to
block_to_feature(... overlap_blocks=overlaps) and the conditional creation of
Eigenvalues) inside that with-block, then remove the explicit
overlap_blocks.close() call; keep existing references to override_overlap,
overlaps, block_to_feature, self.overlap, and the Eigenvalues constructor
unchanged.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@dpnegf/runner/NEGF.py`:
- Around line 125-129: The default parallel_options dict uses the wrong key
"n_job" which causes KeyError when code later accesses
self.parallel_options["n_jobs"]; update the default to use the correct key name
"n_jobs" (i.e., set "n_jobs": -1) or alternatively change the later accesses to
read "n_job" consistently, but prefer changing the default so
self.parallel_options and all reads use "n_jobs"; ensure any other keys (e.g.,
"batch_size", "backend") remain unchanged.
- Around line 113-115: The validation currently asserts override_overlap is a
str which rejects the valid False boolean; update the check in the __init__ (or
where override_overlap is read) so it accepts either False or a str (e.g. use
isinstance(kwargs['override_overlap'], (str, bool)) or explicitly allow
kwargs['override_overlap'] is False), then assign self.override_overlap =
kwargs['override_overlap'] unchanged; reference the override_overlap handling
around the current assert and self.override_overlap assignment.

In `@dpnegf/tests/test_negf_e3.py`:
- Around line 36-37: Replace the shell-based cleanup call that uses
os.system("rm -r "+output+"/results") with a platform-safe Python removal:
import shutil if not already imported, and call
shutil.rmtree(os.path.join(output, "results")) (optionally wrapped in a
try/except or check with os.path.exists) instead of os.system; update the block
around the existing os.path.exists(output+"/results") check and remove the
os.system usage to ensure safe, cross-platform test cleanup.

In `@dpnegf/utils/argcheck.py`:
- Line 1070: The schema for Argument("override_overlap", str, optional=True,
doc=doc_override_overlap) currently only accepts str and thus blocks the
explicit disable path; update the Argument type to accept both strings and
boolean False (e.g., allow (str, bool) or an equivalent union that permits
False) so callers can pass False to disable overlap override, and ensure any
downstream validation logic that checks override_overlap handles the False
sentinel correctly while keeping doc_override_overlap unchanged.

---

Outside diff comments:
In `@dpnegf/utils/elec_struc_cal.py`:
- Around line 153-172: The code opens an h5py.File into overlap_blocks and
closes it manually, risking leaked file handles on exceptions; wrap the open in
a context manager (with h5py.File(override_overlap, "r") as overlap_blocks:) and
move the logic that reads overlap_blocks (the log/info checks, the "0"/"1"
selection, call to block_to_feature(... overlap_blocks=overlaps) and the
conditional creation of Eigenvalues) inside that with-block, then remove the
explicit overlap_blocks.close() call; keep existing references to
override_overlap, overlaps, block_to_feature, self.overlap, and the Eigenvalues
constructor unchanged.
🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

  • Push a commit to this branch (recommended)
  • Create a new PR with the fixes

ℹ️ Review info
⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 9634520d-38e1-4c60-abe6-9674b91abc5f

📥 Commits

Reviewing files that changed from the base of the PR and between 25b9b5f and b8bc3b1.

📒 Files selected for processing (7)
  • dpnegf/negf/lead_property.py
  • dpnegf/runner/NEGF.py
  • dpnegf/tests/data/test_negf/test_negf_e3/7_0.vasp
  • dpnegf/tests/data/test_negf/test_negf_e3/negf_7_0.json
  • dpnegf/tests/test_negf_e3.py
  • dpnegf/utils/argcheck.py
  • dpnegf/utils/elec_struc_cal.py

Comment thread dpnegf/runner/NEGF.py
Comment on lines +113 to +115
if 'override_overlap' in kwargs:
assert isinstance(kwargs['override_overlap'], str)
self.override_overlap = kwargs['override_overlap']

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major | ⚡ Quick win

override_overlap=False cannot pass through runner validation.

Line 114 asserts str, which rejects the explicit False disable behavior supported downstream.

Suggested fix
-        if 'override_overlap' in kwargs:
-            assert isinstance(kwargs['override_overlap'], str)
+        if 'override_overlap' in kwargs:
+            assert isinstance(kwargs['override_overlap'], (str, bool, type(None)))
             self.override_overlap = kwargs['override_overlap']
📝 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.

Suggested change
if 'override_overlap' in kwargs:
assert isinstance(kwargs['override_overlap'], str)
self.override_overlap = kwargs['override_overlap']
if 'override_overlap' in kwargs:
assert isinstance(kwargs['override_overlap'], (str, bool, type(None)))
self.override_overlap = kwargs['override_overlap']
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@dpnegf/runner/NEGF.py` around lines 113 - 115, The validation currently
asserts override_overlap is a str which rejects the valid False boolean; update
the check in the __init__ (or where override_overlap is read) so it accepts
either False or a str (e.g. use isinstance(kwargs['override_overlap'], (str,
bool)) or explicitly allow kwargs['override_overlap'] is False), then assign
self.override_overlap = kwargs['override_overlap'] unchanged; reference the
override_overlap handling around the current assert and self.override_overlap
assignment.

Comment thread dpnegf/runner/NEGF.py
Comment on lines +125 to +129
self.parallel_options = {
"n_job": -1,
"batch_size": 200,
"backend": "loky"
}

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🔴 Critical | ⚡ Quick win

Default parallel options will crash with KeyError.

Line 126 defines n_job, but Lines 564 and 571 read n_jobs. If parallel_options is not provided, this path fails at runtime.

Suggested fix
         else:
             self.parallel_options = {
-                "n_job": -1,
+                "n_jobs": -1,
                 "batch_size": 200,
                 "backend": "loky"
             }

Also applies to: 564-565, 571-572

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@dpnegf/runner/NEGF.py` around lines 125 - 129, The default parallel_options
dict uses the wrong key "n_job" which causes KeyError when code later accesses
self.parallel_options["n_jobs"]; update the default to use the correct key name
"n_jobs" (i.e., set "n_jobs": -1) or alternatively change the later accesses to
read "n_job" consistently, but prefer changing the default so
self.parallel_options and all reads use "n_jobs"; ensure any other keys (e.g.,
"batch_size", "backend") remain unchanged.

Comment on lines +36 to +37
if os.path.exists(output+"/results"):
os.system("rm -r "+output+"/results")

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor | ⚡ Quick win

Avoid shell-based cleanup in tests.

Line 37 uses os.system("rm -r ..."); prefer shutil.rmtree for safer, platform-independent cleanup.

Suggested fix
+import shutil
...
-    if os.path.exists(output+"/results"):
-        os.system("rm -r "+output+"/results")
+    if os.path.exists(output + "/results"):
+        shutil.rmtree(output + "/results")
📝 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.

Suggested change
if os.path.exists(output+"/results"):
os.system("rm -r "+output+"/results")
if os.path.exists(output + "/results"):
shutil.rmtree(output + "/results")
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@dpnegf/tests/test_negf_e3.py` around lines 36 - 37, Replace the shell-based
cleanup call that uses os.system("rm -r "+output+"/results") with a
platform-safe Python removal: import shutil if not already imported, and call
shutil.rmtree(os.path.join(output, "results")) (optionally wrapped in a
try/except or check with os.path.exists) instead of os.system; update the block
around the existing os.path.exists(output+"/results") check and remove the
os.system usage to ensure safe, cross-platform test cleanup.

Comment thread dpnegf/utils/argcheck.py
Argument("out_ldos", bool, optional=True, default=False, doc=doc_out_ldos),
Argument("out_lcurrent", bool, optional=True, default=False, doc=doc_out_lcurrent)
Argument("out_lcurrent", bool, optional=True, default=False, doc=doc_out_lcurrent),
Argument("override_overlap", str, optional=True, doc=doc_override_overlap),

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major | ⚡ Quick win

override_overlap schema blocks the explicit disable path.

Line 1070 only accepts str, but the downstream API supports False to explicitly disable overlap override. This prevents that behavior from being expressed in config.

Suggested fix
-        Argument("override_overlap", str, optional=True, doc=doc_override_overlap),
+        Argument("override_overlap", [str, bool, None], optional=True, default=None, doc=doc_override_overlap),
📝 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.

Suggested change
Argument("override_overlap", str, optional=True, doc=doc_override_overlap),
Argument("override_overlap", [str, bool, None], optional=True, default=None, doc=doc_override_overlap),
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@dpnegf/utils/argcheck.py` at line 1070, The schema for
Argument("override_overlap", str, optional=True, doc=doc_override_overlap)
currently only accepts str and thus blocks the explicit disable path; update the
Argument type to accept both strings and boolean False (e.g., allow (str, bool)
or an equivalent union that permits False) so callers can pass False to disable
overlap override, and ensure any downstream validation logic that checks
override_overlap handles the False sentinel correctly while keeping
doc_override_overlap unchanged.

@Lonya0 Lonya0 closed this Jun 11, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants