Skip to content

Implement batched RGF refactorization and related unit tests#27

Merged
AsymmetryChou merged 8 commits into
deepmodeling:mainfrom
AsymmetryChou:rgf_acc
Jun 12, 2026
Merged

Implement batched RGF refactorization and related unit tests#27
AsymmetryChou merged 8 commits into
deepmodeling:mainfrom
AsymmetryChou:rgf_acc

Conversation

@AsymmetryChou

Copy link
Copy Markdown
Collaborator

This pull request introduces robust support for batched (vectorized) energy calculations throughout the NEGF device and recursive Green's function code. The changes update both the core algorithms and their documentation to consistently handle tensors with a leading batch dimension, enabling efficient multi-energy evaluations. Several tensor operations have been generalized to work in both batched and scalar modes, and some mathematical expressions have been modernized for clarity and performance.

Batched energy and tensor support:

  • Added batched energy support to recursive_gf_cal and related device property routines, ensuring all tensors (Hamiltonians, self-energies, Green's functions) can carry a leading batch dimension and updating all relevant tensor operations accordingly. [1] [2] [3] [4]
  • Introduced the _build_s_in_batched helper in device_property.py to efficiently construct per-block self-energy tensors for batched calculations.

Tensor operation improvements and consistency:

  • Replaced .conj().T with .mH (Hermitian transpose) for clarity and consistency in all Green's function and self-energy calculations. [1] [2]
  • Updated all diagonal extraction and concatenation operations (e.g., .diag() to .diagonal(dim1=-2, dim2=-1), torch.cat(..., dim=-1)) to support batched tensors and maintain correct output shapes. [1] [2]

Algorithmic and codebase enhancements:

  • Generalized the calculation of transmission (_cal_tc_), density of states (_cal_dos_), and local density of states (_cal_ldos_) to work seamlessly in both batched and scalar modes, improving performance and code maintainability. [1] [2] [3]
  • Improved documentation in recursive_gf_cal to clearly describe batched tensor shapes and requirements for all arguments and outputs. [1] [2]

Minor and maintenance changes:

  • Added comments and TODOs to highlight areas that still expect scalar inputs and may require further vectorization.

These changes collectively modernize the NEGF codebase, making it more efficient and ready for high-throughput device simulations across multiple energies.

@AsymmetryChou AsymmetryChou merged commit 74a8649 into deepmodeling:main Jun 12, 2026
1 check passed
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.

1 participant