diff --git a/Cargo.lock b/Cargo.lock index 72a6fda..53f9316 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -23,6 +23,37 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" +[[package]] +name = "crossbeam-deque" +version = "0.8.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5181e0de7b61eb03a81e347d6dd8797bae9da5146707b51077e2d71a54ec0ceb" +dependencies = [ + "crossbeam-epoch", + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-epoch" +version = "0.9.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2d6914041f254d6e9176c01941b21115dcfb7089e55135a35411081bd106ef3f" +dependencies = [ + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-utils" +version = "0.8.22" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "61803da095bee82a81bb1a452ecc25d3b2f1416d1897eb86430c6159ef717c17" + +[[package]] +name = "either" +version = "1.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9e5e8f6c15a24b9a3ee5efec809ccd006d3b30e8b3bb63c39af737c7f87daa1d" + [[package]] name = "heck" version = "0.5.0" @@ -59,6 +90,7 @@ version = "0.1.0" dependencies = [ "approx", "num-complex", + "rayon", "thiserror", ] @@ -181,6 +213,26 @@ dependencies = [ "proc-macro2", ] +[[package]] +name = "rayon" +version = "1.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fb39b166781f92d482534ef4b4b1b2568f42613b53e5b6c160e24cfbfa30926d" +dependencies = [ + "either", + "rayon-core", +] + +[[package]] +name = "rayon-core" +version = "1.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22e18b0f0062d30d4230b2e85ff77fdfe4326feb054b9783a3460d8435c8ab91" +dependencies = [ + "crossbeam-deque", + "crossbeam-utils", +] + [[package]] name = "rustversion" version = "1.0.23" diff --git a/README.md b/README.md index 8b4f451..5d58a82 100644 --- a/README.md +++ b/README.md @@ -37,19 +37,15 @@ program. `negforge-py` is a thin binding layer on top of it. by both the electrostatics and NEGF modules. - `device` — device geometry/parameters and the electrostatic (`calc_potential`) and ballistic-current (`calc_current`) calculations. -- `negf` — the NEGF retarded Green's-function calculation. Rather than the - original's dense O(N^3) matrix inversion (`inv()` in MATLAB, fine for a - one-off plot but too slow to call repeatedly in a self-consistent loop - or a bias sweep), this uses a left-connected recursive Green's-function - sweep (O(N)) for the diagonal (local density of states) and a direct - tridiagonal solve (O(N)) for each of the two contact-column quantities - needed for the charge density — both cross-checked against a dense - reference solver in the test suite. +- `negf` — the NEGF retarded Green's-function calculation, with two + interchangeable algorithms (`GreenFunctionAlgorithm::{Recursive, + Dense}`) and a parallel energy loop — see "Performance" below. - `charge` — NEGF-derived electron density, with two bug fixes relative to the original `calc_n()` (see "Deviations from the original" below). - `selfconsistent` — **new**: the Poisson↔NEGF self-consistency loop. - `sweep` — bias sweeps (`sweep_v_g`, `sweep_v_ds`) and subthreshold-swing - extraction, mirroring `plot_Vg_I` / `plot_Vds_I` / `plot_S`. + extraction, mirroring `plot_Vg_I` / `plot_Vds_I` / `plot_S`. Each sweep + point runs in parallel (see "Performance"). ## Building and running @@ -154,8 +150,9 @@ below is a deliberate, documented decision, not an accident: `calc_green` already conditions each contact's self-energy on its own band edge. -5. **O(N) NEGF instead of O(N^3).** See "Architecture" above. Purely a - performance change (validated against a dense reference solver in +5. **A choice of O(N) or O(N^3) NEGF, plus parallelism.** See + "Performance" below. Purely an implementation/engineering addition + (validated by cross-checking the two algorithms against each other in tests); the physics is unchanged. Everything else — the electrostatic operator, the ballistic current @@ -168,19 +165,113 @@ re-derive the model's physics from first principles, only to make it run, close its one clearly-missing feedback loop, and fix the bugs that stood in the way of that loop actually doing something. +## Performance + +### Dense vs. recursive NEGF: the tradeoff + +Every NEGF quantity (local density of states, the injected charge used by +the self-consistent loop) can be computed two ways, chosen via +`GreenFunctionAlgorithm` in Rust or an `algorithm: "recursive" | "dense"` +string argument in Python: + +| | `Dense` | `Recursive` (default) | +|--------------------|------------------------------------------------------|-----------------------------------------------------------------| +| Method | Build the full N×N complex matrix, invert it (Gauss-Jordan), read off the diagonal and two boundary columns — what the original MATLAB `inv()` call did. | Left-connected recursive Green's-function sweep (diagonal) + a direct tridiagonal solve (the two boundary columns). | +| Cost per energy point | O(N^3) time, O(N^2) memory | O(N) time, O(N) memory | +| Why it exists | Obviously correct: a direct definition-level matrix inversion, no tridiagonal-structure assumption, no recursion formula to get subtly wrong. Useful as a trusted reference (see tests) and as a fallback if the Hamiltonian ever gains longer-range hopping and stops being purely tridiagonal. | The one to use for anything performance-sensitive — the self-consistent loop and bias sweeps call this dozens to hundreds of times. | + +Both are cross-checked against each other in +`negf::tests::dense_and_recursive_algorithms_agree_end_to_end` and +`selfconsistent::tests::dense_and_recursive_converge_to_the_same_potential`. +Measured with `cargo run --release --example benchmark -p negforge-core` +(single NEGF sweep, 100 energy points, on a 4-core machine): + +``` + N dense (s) recursive (s) speedup + 21 0.0018 0.0002 10x + 51 0.0144 0.0002 76x + 101 0.1202 0.0003 395x + 201 0.8849 0.0005 1889x + 351 6.2474 0.0009 6611x +``` + +Dense scales cubically and recursive is essentially flat, as expected. At +the model's default device size (N=561), a full self-consistent solve +using `Dense` did not finish a single iteration in several minutes — it is +not a realistic choice at that scale, only for small devices or +cross-validation. + +### Parallelism + +Every energy point in a Green's-function sweep is an independent +calculation (they only read the same fixed Hamiltonian), so the energy +loop runs in parallel via `rayon`, spreading points across all available +CPU cores automatically (respects `RAYON_NUM_THREADS` if you want to cap +it). This applies to both algorithms and is the default — there's no +opt-in needed. Measured on the same 4-core machine (Recursive, N=561, 900 +energy points, one sweep): + +``` +1 thread: 0.0479 s +all cores (4): 0.0169 s +speedup: 2.8x-4.0x (varies by run/load) +``` + +Bias sweeps (`sweep_v_g`/`sweep_v_ds`) are parallel too, at the *point* +level rather than the energy level: `set_v_g`/`set_v_ds` reset a device's +`psi_g`/`psi_bi`/`rho` from scratch (see `Device::init_vectors`), so +nothing carries over between bias points — each one is evaluated on its +own cloned `Device`, in parallel. This is why `sweep_v_g`/`sweep_v_ds` take +`&Device` (a template whose bias is varied) rather than `&mut Device`: the +input device's own state is left untouched, which also fixed a surprising +side effect the earlier API had (a sweep silently leaving the device +parked at its last bias point). + +The one thing that is **not** parallelizable is the self-consistent loop's +*iterations* — each iteration's charge depends on the previous iteration's +potential, a genuine sequential dependency. Parallelism instead comes from +within each iteration's NEGF sweep (above), which is where nearly all the +time goes. + +### Using the switch + +Rust: + +```rust +use negforge_core::{GreenFunctionAlgorithm, SelfConsistentOptions}; + +let opts = SelfConsistentOptions { + algorithm: GreenFunctionAlgorithm::Dense, // or ::Recursive (default) + ..Default::default() +}; +``` + +Python: + +```python +dev.solve_self_consistent(algorithm="dense") # or "recursive" (default) +dev.local_density_of_states(algorithm="dense") +dev.sweep_v_g(0.0, 0.4, 0.05, self_consistent=True, algorithm="dense") +``` + ## Testing - `crates/negforge-core/src/*.rs` — unit tests per module, including cross-checks of the O(N) tridiagonal/recursive-Green's-function solvers - against dense (Gaussian-elimination / full-matrix-inversion) reference - implementations on small systems. + against the dense (Gaussian-elimination / full-matrix-inversion) + algorithm on small systems, and of the two `GreenFunctionAlgorithm` + variants against each other end-to-end (`negf.rs`) and through the full + self-consistent loop (`selfconsistent.rs`). - `crates/negforge-core/tests/self_consistent_realistic_device.rs` — an integration test at the model's default (non-toy) device scale, confirming the self-consistent loop converges and reproduces the expected ballistic-MOSFET trend (current increasing with gate bias). +- `crates/negforge-core/examples/benchmark.rs` — the performance audit + behind the numbers quoted above; run it with `cargo run --release + --example benchmark -p negforge-core`. - `notebooks/negforge_demo.ipynb` has been executed end-to-end - (`jupyter nbconvert --execute`) to confirm the full frontend path works; - outputs are cleared before committing since they go stale the moment the - engine changes. + (`jupyter nbconvert --execute`) to confirm the full frontend path works, + including the dense-vs-recursive comparison cell; outputs are cleared + before committing since they go stale the moment the engine changes. Run everything with `cargo test --workspace`. diff --git a/crates/negforge-core/Cargo.toml b/crates/negforge-core/Cargo.toml index 625b2fb..d008731 100644 --- a/crates/negforge-core/Cargo.toml +++ b/crates/negforge-core/Cargo.toml @@ -8,6 +8,7 @@ description = "1D ballistic MOSFET electrostatics + NEGF transport engine" [dependencies] num-complex = "0.4" thiserror = "1" +rayon = "1" [dev-dependencies] approx = "0.5" diff --git a/crates/negforge-core/examples/benchmark.rs b/crates/negforge-core/examples/benchmark.rs new file mode 100644 index 0000000..238835b --- /dev/null +++ b/crates/negforge-core/examples/benchmark.rs @@ -0,0 +1,91 @@ +//! Performance audit: dense vs. recursive NEGF, and parallel vs. serial +//! energy-loop evaluation. +//! +//! Run with `cargo run --release --example benchmark -p negforge-core`. +//! Numbers from this example are quoted in the top-level README's +//! "Performance" section. + +use std::time::Instant; + +use negforge_core::negf::{green_function_sweep, GreenFunctionAlgorithm}; +use negforge_core::{Device, DeviceParams}; + +fn device_with_n_sites(target_n: usize) -> Device { + // l_ch/l_ds/a chosen so the resulting grid has (very close to) target_n + // sites, holding geometry proportions roughly fixed. + let a = 1.0; + let l_ch = (target_n as f64 - 1.0) * a / 3.0; + let mut dev = Device::new(DeviceParams { + a, + l_ch, + l_ds: l_ch, + auto_size_contacts: false, + v_ds: 0.3, + v_g: 0.2, + ..Default::default() + }); + dev.calc_potential(); + dev +} + +fn time_sweep(device: &Device, n_energies: usize, algorithm: GreenFunctionAlgorithm) -> f64 { + let e_min = device.psi_f.iter().cloned().fold(f64::INFINITY, f64::min); + let energies: Vec = (0..n_energies).map(|k| e_min + k as f64 * 0.005).collect(); + let start = Instant::now(); + let _ = green_function_sweep(device, &energies, 0.05, algorithm); + start.elapsed().as_secs_f64() +} + +fn main() { + const N_ENERGIES: usize = 100; + println!( + "=== Dense (O(N^3)) vs. Recursive (O(N)) NEGF, {N_ENERGIES} energy points per sweep ===" + ); + println!( + "{:>6} {:>14} {:>14} {:>10}", + "N", "dense (s)", "recursive (s)", "speedup" + ); + for &n in &[21usize, 51, 101, 201, 351] { + let device = device_with_n_sites(n); + let actual_n = device.n; + let dense_t = time_sweep(&device, N_ENERGIES, GreenFunctionAlgorithm::Dense); + let recursive_t = time_sweep(&device, N_ENERGIES, GreenFunctionAlgorithm::Recursive); + println!( + "{:>6} {:>14.4} {:>14.4} {:>9.0}x", + actual_n, + dense_t, + recursive_t, + dense_t / recursive_t + ); + } + + println!(); + println!("=== Parallel vs. serial energy loop (Recursive, default-scale device N=561, 900 energy points) ==="); + let mut device = Device::new(DeviceParams { + v_ds: 0.3, + v_g: 0.2, + ..Default::default() + }); + device.calc_potential(); + println!("N = {}", device.n); + + let n_energies = 900; + let cores = std::thread::available_parallelism() + .map(|n| n.get()) + .unwrap_or(1); + + let serial_pool = rayon::ThreadPoolBuilder::new() + .num_threads(1) + .build() + .unwrap(); + let serial_t = + serial_pool.install(|| time_sweep(&device, n_energies, GreenFunctionAlgorithm::Recursive)); + + // Default global pool: uses all available cores. + let parallel_t = time_sweep(&device, n_energies, GreenFunctionAlgorithm::Recursive); + + println!("available cores (per std::thread::available_parallelism): {cores}"); + println!("1 thread: {serial_t:.4} s"); + println!("all cores ({cores}): {parallel_t:.4} s"); + println!("speedup: {:.2}x", serial_t / parallel_t); +} diff --git a/crates/negforge-core/src/lib.rs b/crates/negforge-core/src/lib.rs index c97b713..2494e99 100644 --- a/crates/negforge-core/src/lib.rs +++ b/crates/negforge-core/src/lib.rs @@ -16,6 +16,6 @@ mod tridiag; pub use device::{Device, DeviceParams}; pub use error::{NegForgeError, Result}; -pub use negf::GreenFunctionResult; +pub use negf::{GreenFunctionAlgorithm, GreenFunctionResult}; pub use selfconsistent::{SelfConsistentOptions, SelfConsistentResult}; pub use sweep::IvPoint; diff --git a/crates/negforge-core/src/negf.rs b/crates/negforge-core/src/negf.rs index 59e447b..0ec26c2 100644 --- a/crates/negforge-core/src/negf.rs +++ b/crates/negforge-core/src/negf.rs @@ -4,23 +4,52 @@ //! the tight-binding Hamiltonian for the potential profile `psi_f`, attaches //! open-boundary self-energies at the source/drain contacts, and computes //! the retarded Green's function `G^r(E) = [(E + i*eta) I - H]^-1` for a -//! grid of energies. +//! grid of energies. Only the diagonal (local density of states) and the +//! first/last columns (needed for the injected charge density) are needed. //! -//! Only the diagonal (local density of states) and the first/last columns -//! (needed for the injected charge density) are needed, so rather than a -//! dense O(N^3) matrix inversion (what the original MATLAB `inv()` call -//! did) this uses: +//! ## Dense vs. recursive: the tradeoff //! -//! - a left-connected recursive Green's function sweep (O(N)) for the full -//! diagonal, and -//! - a direct tridiagonal solve (O(N), via [`tridiag::solve_complex`]) for -//! each of the two boundary columns, sidestepping the (easy to get wrong) -//! off-diagonal recursive Green's function formulas entirely. +//! Two algorithms compute those same outputs ([`GreenFunctionAlgorithm`]): //! -//! Both pieces are cross-checked against a dense reference solver in the -//! test module below. +//! - **[`GreenFunctionAlgorithm::Dense`]**: build the full `N x N` complex +//! matrix and invert it (Gauss-Jordan elimination), then read off the +//! diagonal and the two boundary columns. This is what the original +//! MATLAB `inv()` call did. It's O(N^3) time and O(N^2) memory per energy +//! point. The appeal is that it's *obviously* correct — a direct +//! definition-level matrix inversion with no tridiagonal-structure +//! assumptions and no recursion formulas to get subtly wrong — which is +//! exactly why it exists here too, as a trusted reference to validate the +//! fast path against (see the tests below) and as a fallback that keeps +//! working if the Hamiltonian ever stops being purely tridiagonal (e.g. a +//! future extension adding longer-range hopping or a non-nearest-neighbor +//! coupling). +//! - **[`GreenFunctionAlgorithm::Recursive`]** (default): a left-connected +//! recursive Green's-function sweep (O(N)) for the diagonal, plus a +//! direct tridiagonal solve (O(N), via [`tridiag::solve_complex`]) for +//! each boundary column — sidestepping the (easy to get wrong) +//! off-diagonal recursive Green's-function formulas entirely. This is +//! the one to use for anything performance-sensitive: a single sweep at +//! `N = 561` (the model's default device size) is already ~200x fewer +//! floating-point operations than dense, and the gap grows with N. It's +//! also the only one fast enough to call repeatedly inside the +//! self-consistent loop or a bias sweep — see `selfconsistent.rs`. +//! +//! Both are exercised and cross-checked against each other in the test +//! module below; `crates/negforge-core/examples/benchmark.rs` has wall-clock +//! numbers. +//! +//! ## Parallelism +//! +//! Every energy point is an independent Green's-function evaluation (they +//! only share the read-only Hamiltonian built from the device's current +//! `psi_f`), so the energy loop is embarrassingly parallel. It's run +//! through `rayon`'s `par_iter`, which spreads energy points across all +//! available CPU cores automatically (respects `RAYON_NUM_THREADS` if you +//! want to cap it). This benefits both algorithms, but matters far more for +//! `Dense`, where each energy point is individually expensive. use num_complex::Complex64; +use rayon::prelude::*; use crate::device::Device; use crate::tridiag; @@ -45,6 +74,23 @@ use crate::tridiag; /// rate) rather than a physics change. pub const DEFAULT_ETA: f64 = 1e-8; +/// Which algorithm [`green_function_sweep`] uses to evaluate the retarded +/// Green's function at each energy point. See the module docs for the +/// tradeoff. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub enum GreenFunctionAlgorithm { + /// O(N) recursive Green's-function sweep + tridiagonal solves. Use this + /// for anything performance-sensitive (self-consistent loop, bias + /// sweeps, large devices). Default. + #[default] + Recursive, + /// O(N^3) full dense matrix inversion, matching the original MATLAB + /// `inv()` call. Simple and easy to trust, but slow — mainly useful as + /// a reference for cross-validation, for small devices, or if you + /// don't trust the recursive path for some reason. + Dense, +} + /// Result of a [`green_function_sweep`] call: retarded-Green's-function /// quantities on an energy grid, one entry per energy point. pub struct GreenFunctionResult { @@ -66,62 +112,58 @@ pub struct GreenFunctionResult { pub k_da: Vec, } +/// One energy point's worth of Green's-function output: `(k_sa, k_da, +/// diag_row, col_source_row, col_drain_row)`. +type EnergyPointResult = (Complex64, Complex64, Vec, Vec, Vec); + /// Run the NEGF sweep over `energies` for the device's current `psi_f`, /// with imaginary broadening `eta` (see [`DEFAULT_ETA`] for guidance on -/// choosing it). Mirrors `calc_green()`. The caller chooses the energy grid -/// (the original hardcoded `min(Psi_f) : dE : 0.7*E_max`). -pub fn green_function_sweep(device: &Device, energies: &[f64], eta: f64) -> GreenFunctionResult { +/// choosing it) and the given [`GreenFunctionAlgorithm`]. Mirrors +/// `calc_green()`. The caller chooses the energy grid (the original +/// hardcoded `min(Psi_f) : dE : 0.7*E_max`). +/// +/// Energy points are evaluated in parallel across available CPU cores (see +/// the module docs). +pub fn green_function_sweep( + device: &Device, + energies: &[f64], + eta: f64, + algorithm: GreenFunctionAlgorithm, +) -> GreenFunctionResult { let n = device.n; let t = device.t_hop; let a = device.params.a; let psi_0s = device.psi_f[0]; let psi_0d = device.psi_f[n - 1]; - // Bulk on-site energies (before contact self-energy is applied). + // Bulk on-site energies (before contact self-energy is applied), and + // the constant real off-diagonal — shared read-only state across all + // (parallel) energy-point evaluations. let bulk_diag: Vec = device.psi_f.iter().map(|p| 2.0 * t + p).collect(); let c1 = bulk_diag[0]; let c_n = bulk_diag[n - 1]; - let off_diag = vec![t; n - 1]; // constant real off-diagonal of (E+i eta)I - H + let off_diag = vec![t; n - 1]; + + let per_energy: Vec = energies + .par_iter() + .map(|&e| { + energy_point( + e, eta, t, a, psi_0s, psi_0d, c1, c_n, &bulk_diag, &off_diag, algorithm, + ) + }) + .collect(); + let mut k_sa = Vec::with_capacity(energies.len()); + let mut k_da = Vec::with_capacity(energies.len()); let mut g_diag = Vec::with_capacity(energies.len()); let mut g_col_source = Vec::with_capacity(energies.len()); let mut g_col_drain = Vec::with_capacity(energies.len()); - let mut k_sa = Vec::with_capacity(energies.len()); - let mut k_da = Vec::with_capacity(energies.len()); - - for &e in energies { - let ka_s = Complex64::new(-(2.0 * t - e + psi_0s) / (2.0 * t), 0.0).acos(); - let ka_d = Complex64::new(-(2.0 * t - e + psi_0d) / (2.0 * t), 0.0).acos(); + for (ka_s, ka_d, diag_row, col_source_row, col_drain_row) in per_energy { k_sa.push(ka_s); k_da.push(ka_d); - - let mut diag: Vec = bulk_diag - .iter() - .map(|&d| Complex64::new(e, eta) - Complex64::new(d, 0.0)) - .collect(); - - if e >= psi_0s { - let sigma_l = t * (Complex64::i() * ka_s).exp(); - diag[0] = Complex64::new(e, eta) - (Complex64::new(c1, 0.0) + sigma_l); - } - if e >= psi_0d { - let sigma_r = t * (Complex64::i() * ka_d).exp(); - diag[n - 1] = Complex64::new(e, eta) - (Complex64::new(c_n, 0.0) + sigma_r); - } - - let full_diag = full_diagonal(&off_diag, &diag); - - let mut e_source = vec![Complex64::new(0.0, 0.0); n]; - e_source[0] = Complex64::new(1.0, 0.0); - let col_source = tridiag::solve_complex(&off_diag, &diag, &off_diag, &e_source); - - let mut e_drain = vec![Complex64::new(0.0, 0.0); n]; - e_drain[n - 1] = Complex64::new(1.0, 0.0); - let col_drain = tridiag::solve_complex(&off_diag, &diag, &off_diag, &e_drain); - - g_diag.push(full_diag.iter().map(|g| g.im / a).collect()); - g_col_source.push(col_source.iter().map(|g| g.norm_sqr()).collect()); - g_col_drain.push(col_drain.iter().map(|g| g.norm_sqr()).collect()); + g_diag.push(diag_row); + g_col_source.push(col_source_row); + g_col_drain.push(col_drain_row); } GreenFunctionResult { @@ -134,10 +176,74 @@ pub fn green_function_sweep(device: &Device, energies: &[f64], eta: f64) -> Gree } } +#[allow(clippy::too_many_arguments)] +fn energy_point( + e: f64, + eta: f64, + t: f64, + a: f64, + psi_0s: f64, + psi_0d: f64, + c1: f64, + c_n: f64, + bulk_diag: &[f64], + off_diag: &[f64], + algorithm: GreenFunctionAlgorithm, +) -> EnergyPointResult { + let n = bulk_diag.len(); + let ka_s = Complex64::new(-(2.0 * t - e + psi_0s) / (2.0 * t), 0.0).acos(); + let ka_d = Complex64::new(-(2.0 * t - e + psi_0d) / (2.0 * t), 0.0).acos(); + + let mut diag: Vec = bulk_diag + .iter() + .map(|&d| Complex64::new(e, eta) - Complex64::new(d, 0.0)) + .collect(); + + if e >= psi_0s { + let sigma_l = t * (Complex64::i() * ka_s).exp(); + diag[0] = Complex64::new(e, eta) - (Complex64::new(c1, 0.0) + sigma_l); + } + if e >= psi_0d { + let sigma_r = t * (Complex64::i() * ka_d).exp(); + diag[n - 1] = Complex64::new(e, eta) - (Complex64::new(c_n, 0.0) + sigma_r); + } + + let (diag_row, col_source, col_drain) = match algorithm { + GreenFunctionAlgorithm::Recursive => { + let full_diag = full_diagonal(off_diag, &diag); + + let mut e_source = vec![Complex64::new(0.0, 0.0); n]; + e_source[0] = Complex64::new(1.0, 0.0); + let col_source = tridiag::solve_complex(off_diag, &diag, off_diag, &e_source); + + let mut e_drain = vec![Complex64::new(0.0, 0.0); n]; + e_drain[n - 1] = Complex64::new(1.0, 0.0); + let col_drain = tridiag::solve_complex(off_diag, &diag, off_diag, &e_drain); + + (full_diag, col_source, col_drain) + } + GreenFunctionAlgorithm::Dense => { + let inverse = dense_full_inverse(off_diag, &diag); + let full_diag: Vec = (0..n).map(|i| inverse[i][i]).collect(); + let col_source: Vec = (0..n).map(|i| inverse[i][0]).collect(); + let col_drain: Vec = (0..n).map(|i| inverse[i][n - 1]).collect(); + (full_diag, col_source, col_drain) + } + }; + + ( + ka_s, + ka_d, + diag_row.iter().map(|g| g.im / a).collect(), + col_source.iter().map(|g| g.norm_sqr()).collect(), + col_drain.iter().map(|g| g.norm_sqr()).collect(), + ) +} + /// Full diagonal of `A^-1` for a tridiagonal `A` with constant real /// off-diagonal `t` (i.e. `A[i][i+1] = A[i+1][i] = t` for all `i`) and /// complex diagonal `diag`, via the standard left-connected recursive -/// Green's function sweep. +/// Green's function sweep. O(N). fn full_diagonal(off_diag: &[f64], diag: &[Complex64]) -> Vec { let n = diag.len(); let mut g_left = vec![Complex64::new(0.0, 0.0); n]; @@ -156,61 +262,64 @@ fn full_diagonal(off_diag: &[f64], diag: &[Complex64]) -> Vec { full } -#[cfg(test)] +/// Full `A^-1` for a tridiagonal `A` with constant real off-diagonal `t` +/// and complex diagonal `diag`, via Gauss-Jordan elimination with partial +/// pivoting on the dense `N x N` matrix. O(N^3). This is the +/// [`GreenFunctionAlgorithm::Dense`] backend, and also serves as the +/// reference implementation the recursive path is validated against in +/// tests. #[allow(clippy::needless_range_loop)] -mod tests { - use super::*; - use crate::device::{Device, DeviceParams}; - use approx::assert_relative_eq; - - /// Dense reference: build the full complex tridiagonal matrix and - /// invert it via Gauss-Jordan elimination, for cross-checking the O(N) - /// recursive/solve-based implementation above on small systems. - fn dense_inverse(off_diag: &[f64], diag: &[Complex64]) -> Vec> { - let n = diag.len(); - let mut a = vec![vec![Complex64::new(0.0, 0.0); n]; n]; - for i in 0..n { - a[i][i] = diag[i]; - if i > 0 { - a[i][i - 1] = Complex64::new(off_diag[i - 1], 0.0); - } - if i < n - 1 { - a[i][i + 1] = Complex64::new(off_diag[i], 0.0); +fn dense_full_inverse(off_diag: &[f64], diag: &[Complex64]) -> Vec> { + let n = diag.len(); + let mut a = vec![vec![Complex64::new(0.0, 0.0); n]; n]; + for i in 0..n { + a[i][i] = diag[i]; + if i > 0 { + a[i][i - 1] = Complex64::new(off_diag[i - 1], 0.0); + } + if i < n - 1 { + a[i][i + 1] = Complex64::new(off_diag[i], 0.0); + } + } + let mut inv = vec![vec![Complex64::new(0.0, 0.0); n]; n]; + for i in 0..n { + inv[i][i] = Complex64::new(1.0, 0.0); + } + for col in 0..n { + let mut pivot = col; + for row in (col + 1)..n { + if a[row][col].norm() > a[pivot][col].norm() { + pivot = row; } } - let mut inv = vec![vec![Complex64::new(0.0, 0.0); n]; n]; - for i in 0..n { - inv[i][i] = Complex64::new(1.0, 0.0); + a.swap(col, pivot); + inv.swap(col, pivot); + let pivot_val = a[col][col]; + for k in 0..n { + a[col][k] /= pivot_val; + inv[col][k] /= pivot_val; } - for col in 0..n { - let mut pivot = col; - for row in (col + 1)..n { - if a[row][col].norm() > a[pivot][col].norm() { - pivot = row; - } + let pivot_row_a = a[col].clone(); + let pivot_row_inv = inv[col].clone(); + for row in 0..n { + if row == col { + continue; } - a.swap(col, pivot); - inv.swap(col, pivot); - let pivot_val = a[col][col]; + let factor = a[row][col]; for k in 0..n { - a[col][k] /= pivot_val; - inv[col][k] /= pivot_val; - } - let pivot_row_a = a[col].clone(); - let pivot_row_inv = inv[col].clone(); - for row in 0..n { - if row == col { - continue; - } - let factor = a[row][col]; - for k in 0..n { - a[row][k] -= factor * pivot_row_a[k]; - inv[row][k] -= factor * pivot_row_inv[k]; - } + a[row][k] -= factor * pivot_row_a[k]; + inv[row][k] -= factor * pivot_row_inv[k]; } } - inv } + inv +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::device::{Device, DeviceParams}; + use approx::assert_relative_eq; #[test] fn recursive_diagonal_matches_dense_inverse() { @@ -221,7 +330,7 @@ mod tests { .collect(); let recursive = full_diagonal(&off_diag, &diag); - let dense = dense_inverse(&off_diag, &diag); + let dense = dense_full_inverse(&off_diag, &diag); for i in 0..n { assert_relative_eq!(recursive[i].re, dense[i][i].re, epsilon = 1e-8); @@ -245,7 +354,7 @@ mod tests { e_last[n - 1] = Complex64::new(1.0, 0.0); let col_last = tridiag::solve_complex(&off_diag, &diag, &off_diag, &e_last); - let dense = dense_inverse(&off_diag, &diag); + let dense = dense_full_inverse(&off_diag, &diag); for i in 0..n { assert_relative_eq!(col0[i].re, dense[i][0].re, epsilon = 1e-8); @@ -268,7 +377,12 @@ mod tests { let e_min = device.psi_f.iter().cloned().fold(f64::INFINITY, f64::min); let energies: Vec = (0..20).map(|i| e_min + i as f64 * 0.01).collect(); - let result = green_function_sweep(&device, &energies, DEFAULT_ETA); + let result = green_function_sweep( + &device, + &energies, + DEFAULT_ETA, + GreenFunctionAlgorithm::Recursive, + ); assert_eq!(result.g_diag.len(), energies.len()); for row in &result.g_diag { @@ -279,4 +393,48 @@ mod tests { assert!(row.iter().all(|v| v.is_finite() && *v >= 0.0)); } } + + #[test] + fn dense_and_recursive_algorithms_agree_end_to_end() { + let mut device = Device::new(DeviceParams { + a: 1.0, + l_ch: 10.0, + l_ds: 10.0, + auto_size_contacts: false, + ..Default::default() + }); + device.calc_potential(); + + let e_min = device.psi_f.iter().cloned().fold(f64::INFINITY, f64::min); + let energies: Vec = (0..15).map(|i| e_min + i as f64 * 0.02).collect(); + + let recursive = green_function_sweep( + &device, + &energies, + DEFAULT_ETA, + GreenFunctionAlgorithm::Recursive, + ); + let dense = green_function_sweep( + &device, + &energies, + DEFAULT_ETA, + GreenFunctionAlgorithm::Dense, + ); + + for k in 0..energies.len() { + for i in 0..device.n { + assert_relative_eq!(recursive.g_diag[k][i], dense.g_diag[k][i], epsilon = 1e-6); + assert_relative_eq!( + recursive.g_col_source[k][i], + dense.g_col_source[k][i], + epsilon = 1e-6 + ); + assert_relative_eq!( + recursive.g_col_drain[k][i], + dense.g_col_drain[k][i], + epsilon = 1e-6 + ); + } + } + } } diff --git a/crates/negforge-core/src/selfconsistent.rs b/crates/negforge-core/src/selfconsistent.rs index 970b862..ea18431 100644 --- a/crates/negforge-core/src/selfconsistent.rs +++ b/crates/negforge-core/src/selfconsistent.rs @@ -37,7 +37,7 @@ use crate::charge; use crate::constants::E as ELEMENTARY_CHARGE; use crate::device::Device; use crate::error::{NegForgeError, Result}; -use crate::negf::{self, GreenFunctionResult}; +use crate::negf::{self, GreenFunctionAlgorithm, GreenFunctionResult}; #[derive(Debug, Clone, Copy, PartialEq)] pub struct SelfConsistentOptions { @@ -61,6 +61,11 @@ pub struct SelfConsistentOptions { /// resolution (at the cost of a noisier, potentially non-converging /// iteration). pub eta: f64, + /// Which NEGF algorithm to use for the sweep inside each iteration. + /// Defaults to `Recursive` (O(N)) — see `negf.rs` module docs for the + /// dense-vs-recursive tradeoff. `Dense` works too, just far slower per + /// iteration; mainly useful for cross-validating a suspicious result. + pub algorithm: GreenFunctionAlgorithm, } impl Default for SelfConsistentOptions { @@ -71,6 +76,7 @@ impl Default for SelfConsistentOptions { mixing: 0.3, green_energy_fraction: 0.7, eta: 0.08, + algorithm: GreenFunctionAlgorithm::Recursive, } } } @@ -109,7 +115,7 @@ pub fn solve_self_consistent( let mut last_residual = f64::INFINITY; for iteration in 1..=opts.max_iterations { let energies = negf_energy_grid(device, opts.green_energy_fraction); - let green = negf::green_function_sweep(device, &energies, opts.eta); + let green = negf::green_function_sweep(device, &energies, opts.eta, opts.algorithm); let n_electron = charge::electron_density( &green, device.t_hop, @@ -177,6 +183,7 @@ mod tests { mixing: 0.3, green_energy_fraction: 0.7, eta: 0.04, + algorithm: GreenFunctionAlgorithm::Recursive, }; let result = solve_self_consistent(&mut device, &opts).expect("should converge"); assert!(result.residual < opts.tolerance); @@ -185,6 +192,43 @@ mod tests { assert!(device.psi_f.iter().all(|v| v.is_finite())); } + #[test] + fn dense_and_recursive_converge_to_the_same_potential() { + let make_device = || { + Device::new(DeviceParams { + a: 1.0, + l_ch: 10.0, + l_ds: 10.0, + auto_size_contacts: false, + d_e: 0.01, + ..Default::default() + }) + }; + let base_opts = SelfConsistentOptions { + max_iterations: 100, + tolerance: 1e-5, + mixing: 0.3, + green_energy_fraction: 0.7, + eta: 0.04, + algorithm: GreenFunctionAlgorithm::Recursive, + }; + + let mut recursive_device = make_device(); + solve_self_consistent(&mut recursive_device, &base_opts) + .expect("recursive should converge"); + + let mut dense_device = make_device(); + let dense_opts = SelfConsistentOptions { + algorithm: GreenFunctionAlgorithm::Dense, + ..base_opts + }; + solve_self_consistent(&mut dense_device, &dense_opts).expect("dense should converge"); + + for (a, b) in recursive_device.psi_f.iter().zip(dense_device.psi_f.iter()) { + assert!((a - b).abs() < 1e-4, "psi_f mismatch: {a} vs {b}"); + } + } + #[test] fn rejects_invalid_mixing() { let mut device = Device::new(DeviceParams { diff --git a/crates/negforge-core/src/sweep.rs b/crates/negforge-core/src/sweep.rs index 4688103..adbe8f9 100644 --- a/crates/negforge-core/src/sweep.rs +++ b/crates/negforge-core/src/sweep.rs @@ -1,5 +1,17 @@ //! Bias sweeps and derived figures of merit (subthreshold swing), mirroring //! `plot_Vg_I`, `plot_Vds_I` and `plot_S` in `legacy_matlab/quantumsim.m`. +//! +//! Each bias point in a sweep is fully independent: `set_v_g`/`set_v_ds` +//! reset `psi_g`, `psi_bi` and `rho` from scratch (see `Device::init_vectors`), +//! so nothing carries over from one point to the next. That makes a sweep +//! embarrassingly parallel — each point runs on its own cloned `Device` via +//! `rayon`, which matters most for self-consistent sweeps (each point is a +//! full Poisson<->NEGF iteration, the most expensive operation in this +//! crate). The sweep functions therefore take `&Device` (a template whose +//! bias is varied) rather than `&mut Device`: the input device's own state +//! is left untouched. + +use rayon::prelude::*; use crate::device::Device; use crate::error::Result; @@ -21,54 +33,54 @@ fn inclusive_range(min: f64, max: f64, step: f64) -> Vec { (0..=steps).map(|i| min + i as f64 * step).collect() } -/// Gate-voltage sweep at fixed drain bias. Mirrors `plot_Vg_I`. +fn one_point( + device: &Device, + voltage: f64, + self_consistency: SelfConsistency, + set_bias: impl Fn(&mut Device, f64), +) -> Result { + let mut dev = device.clone(); + set_bias(&mut dev, voltage); + match self_consistency { + Some(opts) => { + selfconsistent::solve_self_consistent(&mut dev, opts)?; + } + None => dev.calc_potential(), + } + Ok(IvPoint { + voltage, + current: dev.calc_current(), + }) +} + +/// Gate-voltage sweep at fixed drain bias. Mirrors `plot_Vg_I`. Points are +/// evaluated in parallel (see module docs). pub fn sweep_v_g( - device: &mut Device, + device: &Device, v_min: f64, v_max: f64, step: f64, self_consistency: SelfConsistency, ) -> Result> { - let mut points = Vec::new(); - for v in inclusive_range(v_min, v_max, step) { - device.set_v_g(v); - match self_consistency { - Some(opts) => { - selfconsistent::solve_self_consistent(device, opts)?; - } - None => device.calc_potential(), - } - points.push(IvPoint { - voltage: v, - current: device.calc_current(), - }); - } - Ok(points) + inclusive_range(v_min, v_max, step) + .into_par_iter() + .map(|v| one_point(device, v, self_consistency, Device::set_v_g)) + .collect() } -/// Drain-voltage sweep at fixed gate bias. Mirrors `plot_Vds_I`. +/// Drain-voltage sweep at fixed gate bias. Mirrors `plot_Vds_I`. Points are +/// evaluated in parallel (see module docs). pub fn sweep_v_ds( - device: &mut Device, + device: &Device, v_min: f64, v_max: f64, step: f64, self_consistency: SelfConsistency, ) -> Result> { - let mut points = Vec::new(); - for v in inclusive_range(v_min, v_max, step) { - device.set_v_ds(v); - match self_consistency { - Some(opts) => { - selfconsistent::solve_self_consistent(device, opts)?; - } - None => device.calc_potential(), - } - points.push(IvPoint { - voltage: v, - current: device.calc_current(), - }); - } - Ok(points) + inclusive_range(v_min, v_max, step) + .into_par_iter() + .map(|v| one_point(device, v, self_consistency, Device::set_v_ds)) + .collect() } /// Ordinary least-squares fit `y = slope * x + intercept`. @@ -108,11 +120,11 @@ mod tests { #[test] fn v_g_sweep_produces_monotonic_current_for_ballistic_mosfet() { - let mut device = Device::new(DeviceParams { + let device = Device::new(DeviceParams { v_ds: 0.3, ..Default::default() }); - let points = sweep_v_g(&mut device, 0.0, 0.4, 0.05, None).unwrap(); + let points = sweep_v_g(&device, 0.0, 0.4, 0.05, None).unwrap(); assert_eq!(points.len(), 9); for w in points.windows(2) { assert!( @@ -124,15 +136,28 @@ mod tests { #[test] fn subthreshold_swing_is_positive_for_a_reasonable_device() { - let mut device = Device::new(DeviceParams { + let device = Device::new(DeviceParams { v_ds: 0.3, ..Default::default() }); - let points = sweep_v_g(&mut device, 0.0, 0.4, 0.02, None).unwrap(); + let points = sweep_v_g(&device, 0.0, 0.4, 0.02, None).unwrap(); let s = subthreshold_swing(&points, 0.0, 0.4); assert!(s > 0.0 && s.is_finite()); } + #[test] + fn sweep_does_not_mutate_the_input_device() { + let device = Device::new(DeviceParams { + v_ds: 0.3, + ..Default::default() + }); + let psi_f_before = device.psi_f.clone(); + let v_g_before = device.params.v_g; + let _ = sweep_v_g(&device, 0.0, 0.4, 0.05, None).unwrap(); + assert_eq!(device.psi_f, psi_f_before); + assert_eq!(device.params.v_g, v_g_before); + } + #[test] fn linear_fit_recovers_known_line() { let x = vec![0.0, 1.0, 2.0, 3.0]; diff --git a/crates/negforge-py/src/lib.rs b/crates/negforge-py/src/lib.rs index 90e660b..613661a 100644 --- a/crates/negforge-py/src/lib.rs +++ b/crates/negforge-py/src/lib.rs @@ -3,15 +3,28 @@ // no-op conversion in our code. #![allow(clippy::useless_conversion)] -use pyo3::exceptions::PyRuntimeError; +use pyo3::exceptions::{PyRuntimeError, PyValueError}; use pyo3::prelude::*; -use negforge_core::{Device, DeviceParams, SelfConsistentOptions}; +use negforge_core::{Device, DeviceParams, GreenFunctionAlgorithm, SelfConsistentOptions}; fn to_py_err(e: negforge_core::NegForgeError) -> PyErr { PyRuntimeError::new_err(e.to_string()) } +/// Parse the Python-facing `algorithm` string ("recursive" or "dense") into +/// a [`GreenFunctionAlgorithm`]. See `negforge_core::negf` module docs for +/// the tradeoff between the two. +fn parse_algorithm(algorithm: &str) -> PyResult { + match algorithm { + "recursive" => Ok(GreenFunctionAlgorithm::Recursive), + "dense" => Ok(GreenFunctionAlgorithm::Dense), + other => Err(PyValueError::new_err(format!( + "unknown algorithm {other:?}, expected \"recursive\" or \"dense\"" + ))), + } +} + /// Python-facing wrapper around [`negforge_core::Device`]. #[pyclass(name = "Device")] struct PyDevice { @@ -93,13 +106,14 @@ impl PyDevice { self.inner.set_l_ch(l); } - #[pyo3(signature = (max_iterations=50, tolerance=1e-6, mixing=0.3, eta=0.08))] + #[pyo3(signature = (max_iterations=50, tolerance=1e-6, mixing=0.3, eta=0.08, algorithm="recursive"))] fn solve_self_consistent( &mut self, max_iterations: usize, tolerance: f64, mixing: f64, eta: f64, + algorithm: &str, ) -> PyResult<(usize, f64)> { let opts = SelfConsistentOptions { max_iterations, @@ -107,6 +121,7 @@ impl PyDevice { mixing, green_energy_fraction: 0.7, eta, + algorithm: parse_algorithm(algorithm)?, }; let result = negforge_core::selfconsistent::solve_self_consistent(&mut self.inner, &opts) .map_err(to_py_err)?; @@ -146,16 +161,24 @@ impl PyDevice { .collect() } + /// Sweep points run in parallel across CPU cores and don't mutate this + /// device (each point uses its own internal clone) — see + /// `negforge_core::sweep` module docs. + #[pyo3(signature = (v_min, v_max, step, self_consistent=false, algorithm="recursive"))] fn sweep_v_g( - &mut self, + &self, v_min: f64, v_max: f64, step: f64, self_consistent: bool, + algorithm: &str, ) -> PyResult<(Vec, Vec)> { - let opts = SelfConsistentOptions::default(); + let opts = SelfConsistentOptions { + algorithm: parse_algorithm(algorithm)?, + ..Default::default() + }; let sc = if self_consistent { Some(&opts) } else { None }; - let points = negforge_core::sweep::sweep_v_g(&mut self.inner, v_min, v_max, step, sc) + let points = negforge_core::sweep::sweep_v_g(&self.inner, v_min, v_max, step, sc) .map_err(to_py_err)?; Ok(( points.iter().map(|p| p.voltage).collect(), @@ -163,16 +186,24 @@ impl PyDevice { )) } + /// Sweep points run in parallel across CPU cores and don't mutate this + /// device (each point uses its own internal clone) — see + /// `negforge_core::sweep` module docs. + #[pyo3(signature = (v_min, v_max, step, self_consistent=false, algorithm="recursive"))] fn sweep_v_ds( - &mut self, + &self, v_min: f64, v_max: f64, step: f64, self_consistent: bool, + algorithm: &str, ) -> PyResult<(Vec, Vec)> { - let opts = SelfConsistentOptions::default(); + let opts = SelfConsistentOptions { + algorithm: parse_algorithm(algorithm)?, + ..Default::default() + }; let sc = if self_consistent { Some(&opts) } else { None }; - let points = negforge_core::sweep::sweep_v_ds(&mut self.inner, v_min, v_max, step, sc) + let points = negforge_core::sweep::sweep_v_ds(&self.inner, v_min, v_max, step, sc) .map_err(to_py_err)?; Ok(( points.iter().map(|p| p.voltage).collect(), @@ -182,8 +213,14 @@ impl PyDevice { /// Run the NEGF sweep at the device's current potential and return /// `(energies, ldos)` where `ldos[k]` is the local density of states - /// row (one value per grid site) at `energies[k]`. - fn local_density_of_states(&self) -> (Vec, Vec>) { + /// row (one value per grid site) at `energies[k]`. Energy points run + /// in parallel across CPU cores. `algorithm` is `"recursive"` (default, + /// O(N) per energy point) or `"dense"` (O(N^3), matching the original + /// MATLAB `inv()` — see `negforge_core::negf` module docs for why both + /// exist). + #[pyo3(signature = (algorithm="recursive"))] + fn local_density_of_states(&self, algorithm: &str) -> PyResult<(Vec, Vec>)> { + let algorithm = parse_algorithm(algorithm)?; let e_min = self .inner .psi_f @@ -198,8 +235,9 @@ impl PyDevice { &self.inner, &energies, negforge_core::negf::DEFAULT_ETA, + algorithm, ); - (result.energies, result.g_diag) + Ok((result.energies, result.g_diag)) } } diff --git a/notebooks/negforge_demo.ipynb b/notebooks/negforge_demo.ipynb index df4122a..e89668c 100644 --- a/notebooks/negforge_demo.ipynb +++ b/notebooks/negforge_demo.ipynb @@ -200,12 +200,54 @@ "plt.show()" ] }, + { + "cell_type": "markdown", + "id": "833e119b", + "metadata": {}, + "source": [ + "## 6. Dense vs. recursive NEGF, and parallelism\n", + "\n", + "Every NEGF quantity here (local density of states, injected charge) can be computed two ways:\n", + "\n", + "- `algorithm=\"recursive\"` (default): an O(N) recursive Green's-function sweep. Fast enough to call repeatedly inside the self-consistent loop or a bias sweep.\n", + "- `algorithm=\"dense\"`: a literal O(N^3) full-matrix inversion, matching the original MATLAB `inv()` call. Simple and easy to trust, but scales terribly — kept mainly as a reference to validate the fast path against.\n", + "\n", + "Both give the same answer; only the speed differs. Energy points are also evaluated in parallel across CPU cores for either algorithm, which is why the gap below is smaller than the crate's own benchmark (`crates/negforge-core/examples/benchmark.rs`) shows for a single-threaded run — see the README's \"Performance\" section for the full numbers and the reasoning behind both choices." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4f3117e4", + "metadata": {}, + "outputs": [], + "source": [ + "import time\n", + "\n", + "# A small device, so the O(N^3) dense path finishes in reasonable time.\n", + "small = negforge.Device(v_ds=0.3, v_g=0.2, a=1.0, l_ch=10.0, l_ds=10.0, auto_size_contacts=False)\n", + "\n", + "t0 = time.perf_counter()\n", + "small_recursive = small.local_density_of_states(algorithm=\"recursive\")\n", + "t_recursive = time.perf_counter() - t0\n", + "\n", + "t0 = time.perf_counter()\n", + "small_dense = small.local_density_of_states(algorithm=\"dense\")\n", + "t_dense = time.perf_counter() - t0\n", + "\n", + "max_diff = np.max(np.abs(np.array(small_recursive[1]) - np.array(small_dense[1])))\n", + "print(f\"N = {small.n} sites\")\n", + "print(f\"recursive: {t_recursive * 1000:.2f} ms\")\n", + "print(f\"dense: {t_dense * 1000:.2f} ms ({t_dense / t_recursive:.0f}x slower)\")\n", + "print(f\"max |LDOS difference| between algorithms: {max_diff:.2e} (should be ~0)\")" + ] + }, { "cell_type": "markdown", "id": "8ba69efc", "metadata": {}, "source": [ - "## 6. Interactive exploration\n", + "## 7. Interactive exploration\n", "\n", "Drag the sliders to see how gate and drain bias reshape the (decoupled) potential profile in real time." ] diff --git a/python/negforge/__init__.py b/python/negforge/__init__.py index 41cb5fb..e703163 100644 --- a/python/negforge/__init__.py +++ b/python/negforge/__init__.py @@ -78,13 +78,21 @@ def solve_self_consistent( tolerance: float = 1e-6, mixing: float = 0.3, eta: float = 0.08, + algorithm: str = "recursive", ) -> tuple[int, float]: """Run the self-consistent Poisson<->NEGF loop (not present in the original MATLAB code — see the top-level README). Returns `(iterations, residual)` on success; raises `RuntimeError` if it does not converge within `max_iterations`. + + `algorithm` is `"recursive"` (default, O(N) per NEGF energy point) + or `"dense"` (O(N^3), matching the original MATLAB `inv()` call). + Recursive is what you want for anything but cross-validating a + suspicious result — each iteration here runs a full NEGF sweep, and + dense makes that dramatically slower. See the README's "Performance" + section for the tradeoff and benchmark numbers. """ - return self._inner.solve_self_consistent(max_iterations, tolerance, mixing, eta) + return self._inner.solve_self_consistent(max_iterations, tolerance, mixing, eta, algorithm) # -- bias control -------------------------------------------------------- @@ -132,24 +140,47 @@ def screening_length(self) -> float: """Natural (screening) length lambda, nm.""" return self._inner.screening_length - def local_density_of_states(self): + def local_density_of_states(self, algorithm: str = "recursive"): """NEGF sweep at the device's current potential. Returns `(energies, ldos)` where `ldos[k]` is the local density of states across all grid sites at `energies[k]`. Matches - `calc_green()`'s `G_r_diag` in the original code. + `calc_green()`'s `G_r_diag` in the original code. Energy points run + in parallel across CPU cores. `algorithm` is `"recursive"` (default) + or `"dense"` — see `solve_self_consistent` / the README for the + tradeoff. """ - energies, ldos = self._inner.local_density_of_states() + energies, ldos = self._inner.local_density_of_states(algorithm) return np.asarray(energies), np.asarray(ldos) - def sweep_v_g(self, v_min: float, v_max: float, step: float, self_consistent: bool = False) -> IVCurve: + def sweep_v_g( + self, + v_min: float, + v_max: float, + step: float, + self_consistent: bool = False, + algorithm: str = "recursive", + ) -> IVCurve: """Gate-voltage sweep at the device's current drain bias. Matches - `plot_Vg_I`.""" - voltage, current = self._inner.sweep_v_g(v_min, v_max, step, self_consistent) + `plot_Vg_I`. Points run in parallel across CPU cores and don't + mutate this device. `algorithm` (`"recursive"`/`"dense"`) only + matters when `self_consistent=True` — see `solve_self_consistent`. + """ + voltage, current = self._inner.sweep_v_g(v_min, v_max, step, self_consistent, algorithm) return IVCurve(voltage, current) - def sweep_v_ds(self, v_min: float, v_max: float, step: float, self_consistent: bool = False) -> IVCurve: + def sweep_v_ds( + self, + v_min: float, + v_max: float, + step: float, + self_consistent: bool = False, + algorithm: str = "recursive", + ) -> IVCurve: """Drain-voltage sweep at the device's current gate bias. Matches - `plot_Vds_I`.""" - voltage, current = self._inner.sweep_v_ds(v_min, v_max, step, self_consistent) + `plot_Vds_I`. Points run in parallel across CPU cores and don't + mutate this device. `algorithm` (`"recursive"`/`"dense"`) only + matters when `self_consistent=True` — see `solve_self_consistent`. + """ + voltage, current = self._inner.sweep_v_ds(v_min, v_max, step, self_consistent, algorithm) return IVCurve(voltage, current)