Add eigenvalue solver option to structure calculations#46
Conversation
WalkthroughThis pull request introduces configurable eigenvalue solver selection and external overlap matrix override capabilities throughout the NEGF calculation pipeline. The Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes
Poem
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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (5)
dpnegf/utils/argcheck.py (1)
1079-1090: Consider adding validation for allowedeig_solvervalues.The argument accepts any string, but only "torch" and "numpy" are valid values. Consider using an explicit list of allowed values to fail-fast on invalid input.
doc_eig_solver="The eigenvalue solver to use." + doc_eig_solver="The eigenvalue solver to use. Options: 'torch' or 'numpy'." doc_nel_atom = "The number of electrons in each element." return [ Argument("device", dict, optional=False, sub_fields=device(), doc=doc_device), Argument("lead_L", dict, optional=False, sub_fields=lead(), doc=doc_lead_L), Argument("lead_R", dict, optional=False, sub_fields=lead(), doc=doc_lead_R), Argument("kmesh", list, optional=True, default=[1,1,1], doc=doc_kmesh), Argument("pbc", list, optional=True, default=[False, False, False], doc=doc_pbc), Argument("gamma_center", list, optional=True, default=True, doc=doc_gamma_center), Argument("time_reversal_symmetry", list, optional=True, default=True, doc=doc_time_reversal_symmetry), Argument("e_fermi_smearing", str, optional=True, default="FD", doc=doc_e_fermi_smearing), - Argument("eig_solver", str, optional=True, default="torch", doc=doc_eig_solver), + Argument("eig_solver", str, optional=True, default="torch", alias=["torch", "numpy"], doc=doc_eig_solver), Argument("nel_atom", [dict,None], optional=True, default=None, doc=doc_nel_atom) ]Note: Verify if
dargs.Argumentsupports analiasor similar parameter for value validation. If not, validation could be added in the consuming code.dpnegf/tests/test_get_fermi.py (1)
11-37: Tests provide good coverage for both solvers.The test properly exercises both
eig_solveroptions ("torch" and "numpy") with both smearing methods. All four test cases correctly expect the same Fermi level value, validating solver equivalence.Consider using
pytest.mark.parametrizeto reduce code duplication and make test failures more specific.import pytest @pytest.mark.parametrize("smearing_method,eig_solver", [ ("FD", "torch"), ("Gaussian", "torch"), ("FD", "numpy"), ("Gaussian", "numpy"), ]) def test_get_fermi(smearing_method, eig_solver): ckpt = f"{rootdir}/test_get_fermi/nnsk.best.pth" stru_data = f"{rootdir}/test_get_fermi/PRIMCELL.vasp" model = build_model(checkpoint=ckpt) nel_atom = {"Au":11} elec_cal = ElecStruCal(model=model, device='cpu') _, efermi = elec_cal.get_fermi_level( data=stru_data, nel_atom=nel_atom, smearing_method=smearing_method, meshgrid=[30,30,30], eig_solver=eig_solver ) assert abs(efermi + 3.2262574434280395) < 1e-3dpnegf/utils/elec_struc_cal.py (3)
71-76: Update type hints to use explicitOptionalsyntax.Several parameters use implicit
Optional(e.g.,pbc: Union[bool, list] = None). Per PEP 484, use explicitOptional[T]orT | Nonefor nullable parameters.def get_data(self, data: Union[AtomicData, ase.Atoms, str], - pbc:Union[bool,list]=None, - device: Union[str, torch.device]=None, - AtomicData_options:dict=None, + pbc: Optional[Union[bool, list]] = None, + device: Optional[Union[str, torch.device]] = None, + AtomicData_options: Optional[dict] = None, override_overlap:Optional[str]=None):
204-214: Backup/restore pattern for overlap data is fragile.The code saves
EDGE_OVERLAP_KEYandNODE_OVERLAP_KEYbeforeself.model(data)and restores them after, suggesting the model call overwrites these fields. This workaround is necessary but could break if model behavior changes.Consider adding a comment explaining why this is needed.
if isinstance(override_overlap, str): + # Store override overlap data before model call, as model may overwrite these fields override_overlap_edge = data[AtomicDataDict.EDGE_OVERLAP_KEY] override_overlap_node = data[AtomicDataDict.NODE_OVERLAP_KEY] # get the eigenvalues data = self.model(data) if isinstance(override_overlap, str): + # Restore override overlap data after model call data[AtomicDataDict.EDGE_OVERLAP_KEY] = override_overlap_edge data[AtomicDataDict.NODE_OVERLAP_KEY] = override_overlap_node
218-221: Consider adding validation foreig_solverparameter.The
eig_solverparameter is passed directly toself.eigv()without validation. If an invalid value is passed, the error will occur deep in the eigenvalue calculation rather than at the interface.def get_fermi_level(self, data: Union[AtomicData, ase.Atoms, str], nel_atom: dict, \ - meshgrid: list = None, klist: np.ndarray=None, pbc:Union[bool,list]=None, - AtomicData_options:dict=None, q_tol:float=1e-10, smearing_method:str='FD', - temp:float=300,eig_solver:Optional[str]='torch'): + meshgrid: Optional[list] = None, klist: Optional[np.ndarray] = None, + pbc: Optional[Union[bool, list]] = None, + AtomicData_options: Optional[dict] = None, q_tol: float = 1e-10, + smearing_method: str = 'FD', temp: float = 300, + eig_solver: str = 'torch'): '''...''' + if eig_solver not in ('torch', 'numpy'): + raise ValueError(f"eig_solver must be 'torch' or 'numpy', got '{eig_solver}'")
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (5)
dpnegf/runner/NEGF.py(1 hunks)dpnegf/tests/test_get_fermi.py(1 hunks)dpnegf/utils/argcheck.py(2 hunks)dpnegf/utils/elec_struc_cal.py(7 hunks)examples/atomic_chain_api/input_files/negf_chain_new.json(1 hunks)
🧰 Additional context used
🧬 Code graph analysis (1)
dpnegf/utils/elec_struc_cal.py (1)
dpnegf/data/interfaces/ham_to_feature.py (1)
block_to_feature(11-173)
🪛 Ruff (0.14.6)
dpnegf/utils/elec_struc_cal.py
6-6: Redefinition of unused np from line 3
Remove definition: np
(F811)
73-73: PEP 484 prohibits implicit Optional
Convert to T | None
(RUF013)
75-75: PEP 484 prohibits implicit Optional
Convert to T | None
(RUF013)
176-176: PEP 484 prohibits implicit Optional
Convert to T | None
(RUF013)
177-177: PEP 484 prohibits implicit Optional
Convert to T | None
(RUF013)
219-219: PEP 484 prohibits implicit Optional
Convert to T | None
(RUF013)
219-219: PEP 484 prohibits implicit Optional
Convert to T | None
(RUF013)
220-220: PEP 484 prohibits implicit Optional
Convert to T | None
(RUF013)
⏰ 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 (3)
examples/atomic_chain_api/input_files/negf_chain_new.json (1)
22-22: LGTM!The new
eig_solverconfiguration option is correctly added to the example file, demonstrating how users can select the numpy solver.dpnegf/runner/NEGF.py (1)
180-186: LGTM!The
eig_solverparameter is correctly retrieved fromstru_optionswith a sensible default of"torch"for backward compatibility, and properly passed toget_fermi_level.dpnegf/tests/test_get_fermi.py (1)
1-1: I'm unable to access the repository (DeePTB-Lab/dpnegf) to verify the import path change. Both direct cloning and GitHub CLI methods failed, likely due to repository access restrictions or availability issues.Since I cannot verify the concern about the import path change from
dpnegf.utils.elec_struc_caltodptb.postprocess.elec_struc_calthrough automated means, here is the rewritten review comment:
Verify import path change is intentional and correct.
The import changed from
dpnegf.utils.elec_struc_caltodptb.postprocess.elec_struc_cal. Please confirm this reflects an intended refactoring (e.g., moving the module between packages) and that theElecStruCalclass is now located in thedptb.postprocessmodule rather thandpnegf.utils.
| import numpy as np | ||
| from ase.io import read | ||
| import ase | ||
| import numpy as np |
There was a problem hiding this comment.
🛠️ Refactor suggestion | 🟠 Major
Remove duplicate numpy import.
numpy is imported twice: on line 3 and line 6. The second import is redundant.
import numpy as np
from ase.io import read
import ase
-import numpy as np
from typing import Union📝 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.
| import numpy as np | |
| from ase.io import read | |
| import ase | |
| import numpy as np | |
| import numpy as np | |
| from ase.io import read | |
| import ase |
🧰 Tools
🪛 Ruff (0.14.6)
6-6: Redefinition of unused np from line 3
Remove definition: np
(F811)
🤖 Prompt for AI Agents
In dpnegf/utils/elec_struc_cal.py around lines 3 to 6, numpy is imported twice
(lines 3 and 6); remove the redundant second "import numpy as np" so numpy is
only imported once, keeping the existing first import and leaving the other
imports unchanged.
| if isinstance(override_overlap, str): | ||
| assert os.path.exists(override_overlap), "Overlap file not found." | ||
| 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() |
There was a problem hiding this comment.
Potential state mutation issue with self.eigv modification.
When override_overlap is provided and self.overlap is False, a new Eigenvalues instance is assigned to self.eigv. This modifies the instance state, which could cause unexpected behavior in subsequent calls if the caller expects the original non-overlap eigv.
Consider using a local variable instead of modifying instance state, or document this behavior clearly.
+ eigv_to_use = self.eigv # default to instance eigv
if isinstance(override_overlap, str):
assert os.path.exists(override_overlap), "Overlap file not found."
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(
+ eigv_to_use = 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()Note: This would require returning the eigv_to_use from get_data or restructuring the code flow.
Committable suggestion skipped: line range outside the PR's diff.
🤖 Prompt for AI Agents
In dpnegf/utils/elec_struc_cal.py around lines 142 to 163, the code assigns a
new Eigenvalues instance to self.eigv when override_overlap is provided and
self.overlap is False, mutating object state; change this to create a local
variable (e.g., eigv_to_use) instead of reassigning self.eigv and wire that
local through the remainder of the function (or return it from get_data) so the
instance state is not mutated unexpectedly; ensure any downstream calls use the
local eigv_to_use and remove/avoid modifying self.eigv unless intentional and
documented.
1 participant
Introduce an eigenvalue solver option in the structure calculation module, allowing users to choose between 'torch' and 'numpy' solvers. Update example JSON and related tests to reflect this new functionality. Note that this function should be compatible with the corresponding version of DeePTB with numpy solver.
Summary by CodeRabbit
New Features
Tests
✏️ Tip: You can customize this high-level summary in your review settings.