From 43954c712dc504125f41e5f02fdddd6444e1aa9a Mon Sep 17 00:00:00 2001
From: kirk0830 <yike.huang@aliyun.com>
Date: Wed, 9 Jul 2025 18:17:45 +0800
Subject: [PATCH 01/14] the program can run now

---
 dpnegf/negf/newton_raphson_speed_up.py | 448 +++++++++++++++++++++++++
 dpnegf/negf/poisson_init.py            |  94 +-----
 2 files changed, 465 insertions(+), 77 deletions(-)
 create mode 100644 dpnegf/negf/newton_raphson_speed_up.py

diff --git a/dpnegf/negf/newton_raphson_speed_up.py b/dpnegf/negf/newton_raphson_speed_up.py
new file mode 100644
index 0000000..078635f
--- /dev/null
+++ b/dpnegf/negf/newton_raphson_speed_up.py
@@ -0,0 +1,448 @@
+import time
+import unittest
+from typing import Optional, Tuple, List, Callable
+
+import numpy as np
+from scipy.sparse import lil_matrix
+
+def _impose_j_bound(inout, nx, ny, nz, typ, val, mask):
+    '''impose the special mask for boundary points'''
+    my_boundary_types_ = ['xmin', 'xmax', 'ymin', 'ymax', 'zmin', 'zmax', 'Dirichlet']
+    my_displ_ = [+1, -1, +nx, -nx, +nx*ny, -nx*ny]
+    for t, d in zip(my_boundary_types_, my_displ_):
+        # perf bottle-neck: search
+        i = np.where(np.array(typ) == t)[0] # all indices of the boundary type
+        assert len(i) > 0, f'no grid point with type {t} found'
+        # scipy sparse matrix also supports the parallized assignment
+        inout[i, i] = val
+        inout[i, i + d] = mask
+
+    # there is also a special type "Dirichlet" for the Dirichlet boundary condition
+    i = np.where(typ == 'Dirichlet')[0]
+    inout[i, i] = val
+
+class TestImposeJacobianBoundaryPerf(unittest.TestCase):
+    def setUp(self):
+        # assuming there are 1e6 grid points, each direction has 1e3
+        self.nx, self.ny, self.nz = 100, 100, 100
+        self.nr = self.nx * self.ny * self.nz
+        # say we have 6 boundary types, each has 300 points
+        typ = ['xmin'] * 30 + ['xmax'] * 30 + ['ymin'] * 30 \
+            + ['ymax'] * 30 + ['zmin'] * 30 + ['zmax'] * 30
+        # and all others are of type "in"
+        typ += ['in'] * (self.nr - 180)
+        self.typ = np.array(typ, dtype=str)
+        
+    def test_impose_bound(self):
+        # create a dummy inout matrix
+        ref = lil_matrix((self.nr, self.nr), dtype=float)
+        nx, ny, nz = self.nx, self.ny, self.nz
+        t = time.time()
+        
+        for i in range(self.nr):
+            if self.typ[i] == "xmin":
+                ref[i, i] = -1.0
+                ref[i, i + 1] = 1.0
+            elif self.typ[i] == "xmax":
+                ref[i, i] = -1.0
+                ref[i, i - 1] = 1.0
+            elif self.typ[i] == "ymin":
+                ref[i, i] = -1.0
+                ref[i, i + nx] = 1.0
+            elif self.typ[i] == "ymax":
+                ref[i, i] = -1.0
+                ref[i, i - nx] = 1.0
+            elif self.typ[i] == "zmin":
+                ref[i, i] = -1.0
+                ref[i, i + nx*ny] = 1.0
+            elif self.typ[i] == "zmax":
+                ref[i, i] = -1.0
+                ref[i, i - nx*ny] = 1.0
+        print(f'old boundary impose method took {time.time() - t:.4f} seconds')
+        
+        inout = lil_matrix((self.nr, self.nr), dtype=float)
+        t = time.time()
+        # call the function
+        _impose_j_bound(inout, nx, ny, nz, self.typ, -1.0, 1.0)
+        print(f'_impose_j_bound took {time.time() - t:.4f} seconds')
+    
+        # check if the inout matrix is equal to the reference matrix
+        self.assertTrue(all(inout[i, i] == ref[i, i] for i in range(self.nr)))
+        self.assertTrue(all(inout[i, i + 1] == ref[i, i + 1] for i in range(30)))
+        self.assertTrue(all(inout[i, i - 1] == ref[i, i - 1] for i in range(30, 60)))
+        self.assertTrue(all(inout[i, i + nx] == ref[i, i + nx] for i in range(60, 90)))
+        self.assertTrue(all(inout[i, i - nx] == ref[i, i - nx] for i in range(90, 120)))
+        self.assertTrue(all(inout[i, i + nx*ny] == ref[i, i + nx*ny] for i in range(120, 150)))
+        self.assertTrue(all(inout[i, i - nx*ny] == ref[i, i - nx*ny] for i in range(150, 180)))
+
+def _impose_b_bound(inout, nx, ny, nz, typ, phi, dirichlet_pot):
+    '''impose the special mask for boundary points in b vector'''
+    my_boundary_types_ = ['xmin', 'xmax', 'ymin', 'ymax', 'zmin', 'zmax', 'Dirichlet']
+    my_displ_ = [+1, -1, +nx, -nx, +nx*ny, -nx*ny]
+    for t, d in zip(my_boundary_types_, my_displ_):
+        # perf bottle-neck: search
+        i = np.where(typ == t)[0]
+        assert len(i) > 0, f'no grid point with type {t} found'
+        inout[i] = phi[i] - phi[i + d]
+        
+    # there is also a special type "Dirichlet" for the Dirichlet boundary condition
+    i = np.where(typ == 'Dirichlet')[0]
+    inout[i] = phi[i] - dirichlet_pot[i]
+
+class TestImposeRHSVecBoundaryPerf(unittest.TestCase):
+    def setUp(self):
+        # assuming there are 1e6 grid points, each direction has 1e3
+        self.nx, self.ny, self.nz = 100, 100, 100
+        self.nr = self.nx * self.ny * self.nz
+        # say we have 6 boundary types, each has 300 points
+        typ = ['xmin'] * 30 + ['xmax'] * 30 + ['ymin'] * 30 \
+            + ['ymax'] * 30 + ['zmin'] * 30 + ['zmax'] * 30 + ['Dirichlet'] * 30
+        # and all others are of type "in"
+        typ += ['in'] * (self.nr - 210)
+        self.typ = np.array(typ, dtype=str)
+        self.phi = np.random.rand(self.nr).astype(float)
+        self.fixed_pot = np.random.rand(self.nr).astype(float)
+        
+    def test_impose_bound(self):
+        # create a dummy inout vector
+        ref = np.zeros(self.nr, dtype=float)
+        nx, ny, nz = self.nx, self.ny, self.nz
+        t = time.time()
+        
+        for i in range(self.nr):
+            if self.typ[i] == "xmin":
+                ref[i] = self.phi[i] - self.phi[i + 1]
+            elif self.typ[i] == "xmax":
+                ref[i] = self.phi[i] - self.phi[i - 1]
+            elif self.typ[i] == "ymin":
+                ref[i] = self.phi[i] - self.phi[i + nx]
+            elif self.typ[i] == "ymax":
+                ref[i] = self.phi[i] - self.phi[i - nx]
+            elif self.typ[i] == "zmin":
+                ref[i] = self.phi[i] - self.phi[i + nx*ny]
+            elif self.typ[i] == "zmax":
+                ref[i] = self.phi[i] - self.phi[i - nx*ny]
+            elif self.typ[i] == "Dirichlet":
+                ref[i] = self.phi[i] - self.fixed_pot[i]
+        print(f'old boundary impose method took {time.time() - t:.4f} seconds')
+        
+        inout = np.zeros(self.nr, dtype=float)
+        t = time.time()
+        # call the function
+        _impose_b_bound(inout, nx, ny, nz, self.typ, self.phi, self.fixed_pot)
+        print(f'_impose_b_bound took {time.time() - t:.4f} seconds')
+    
+        # check if the inout vector is equal to the reference vector
+        self.assertTrue(np.all(inout == ref))
+
+def coloumb(chr: float|np.ndarray) -> float|np.ndarray:
+    '''convert the charge to unit of Coulomb'''
+    assert isinstance(chr, (float, np.ndarray)), "chr must be a float or numpy array"
+    from dpnegf.utils.constants import elementary_charge as e
+    return chr * e
+
+def _bflux_impl(jflux, i, nx, ny, nz, phi, full_size=True) -> np.ndarray:
+    '''calculate the b vector from the jflux and phi'''
+    mydisp_ = [-1, +1, -nx, +nx, -nx*ny, +nx*ny]
+    if full_size:
+        assert jflux.shape == (6, len(phi))
+        bflux = np.zeros((6, len(phi)), dtype=float)
+        for j, d in enumerate(mydisp_):
+            bflux[j, i] = jflux[j, i] * (phi[i + d] - phi[i])
+    else:
+        assert jflux.shape == (6, len(i))
+        bflux = np.array([jf * (phi[i + d] - phi[i]) for jf, d in zip(jflux, mydisp_)])
+        assert bflux.shape == (6, len(i))
+    return bflux
+
+class TestBFluxImplPerf(unittest.TestCase):
+    def setUp(self):
+        # assuming there are 1e6 grid points, each direction has 1e3
+        self.nx, self.ny, self.nz = 100, 100, 100
+        self.nr = self.nx * self.ny * self.nz
+        # set the boundary types
+        typ = np.array(['in'] * self.nr, dtype=str).reshape((self.nx, self.ny, self.nz))
+        typ[ 0,  :,  :] = 'xmin'
+        typ[-1,  :,  :] = 'xmax'
+        typ[ :,  0,  :] = 'ymin'
+        typ[ :, -1,  :] = 'ymax'
+        typ[ :,  :,  0] = 'zmin'
+        typ[ :,  :, -1] = 'zmax'
+        typ = typ.flatten().tolist()
+        self.typ = np.array(typ, dtype=str)
+        self.r = np.random.rand(self.nr, 3).astype(float)
+        self.sigma = np.random.rand(self.nr, 3).astype(float)
+        self.eps = np.random.rand(self.nr).astype(float)
+        self.eps0 = 8.854187817e-12
+        self.avgeps = lambda eps1, eps2: (eps1 + eps2) / 2.0  # arithmetic mean
+    def test_bflux_impl(self):
+        ind = np.where(self.typ == "in")[0]
+        jflux = np.random.rand(6, self.nr).astype(float)
+        phi = np.random.rand(self.nr).astype(float)
+        # calculate the reference result
+        ref = np.zeros((6, self.nr), dtype=float)
+        nx, ny, nz = self.nx, self.ny, self.nz
+        t = time.time()
+        for i in range(self.nr):
+            if self.typ[i] == 'in':
+                ref[0, i] = jflux[0, i] * (phi[i     - 1] - phi[i])
+                ref[1, i] = jflux[1, i] * (phi[i     + 1] - phi[i])
+                ref[2, i] = jflux[2, i] * (phi[i    - nx] - phi[i])
+                ref[3, i] = jflux[3, i] * (phi[i    + nx] - phi[i])
+                ref[4, i] = jflux[4, i] * (phi[i - nx*ny] - phi[i])
+                ref[5, i] = jflux[5, i] * (phi[i + nx*ny] - phi[i])
+        print(f'old bflux_impl method took {time.time() - t:.4f} seconds')
+        # call the function
+        t = time.time()
+        bflux = _bflux_impl(jflux, ind, self.nx, self.ny, self.nz, phi, full_size=True)
+        print(f'_bflux_impl took {time.time() - t:.4f} seconds')
+        # check if the result is equal to the reference result
+        self.assertTrue(bflux.shape == (6, self.nr))
+        self.assertTrue(np.allclose(bflux, ref, rtol=1e-5, atol=1e-8))
+
+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.'''
+
+    # get all indices of the grid points of type "in"
+    i = np.where(typ == "in")[0]
+    # j stands for flux, <listcomp> faster than for loop
+    mydisp_ = [-1, +1, -nx, +nx, -nx*ny, +nx*ny]
+    if not full_size:
+        jflux = np.array([avgeps(eps[i + d], eps[i]) * eps0 * sigma[i, j // 2] \
+                          / np.abs(r[i + d, j // 2] - r[i, j // 2])
+                          for j, d in enumerate(mydisp_)]).reshape((6, -1))
+        assert jflux.shape == (6, len(i))
+    else:
+        jflux = np.zeros((6, len(typ)), dtype=float)
+        for j, d in enumerate(mydisp_):
+            jflux[j, i] =   avgeps(eps[i + d], eps[i]) * eps0 * sigma[i, j // 2] \
+                          / np.abs(r[i + d, j // 2] - r[i, j // 2])
+
+    return (jflux, i if with_index else None)
+
+class TestJFluxImplPerf(unittest.TestCase):
+    def setUp(self):
+        # assuming there are 1e6 grid points, each direction has 1e3
+        self.nx, self.ny, self.nz = 100, 100, 100
+        self.nr = self.nx * self.ny * self.nz
+        # set the boundary types
+        typ = np.array(['in'] * self.nr, dtype=str).reshape((self.nx, self.ny, self.nz))
+        typ[ 0,  :,  :] = 'xmin'
+        typ[-1,  :,  :] = 'xmax'
+        typ[ :,  0,  :] = 'ymin'
+        typ[ :, -1,  :] = 'ymax'
+        typ[ :,  :,  0] = 'zmin'
+        typ[ :,  :, -1] = 'zmax'
+        typ = typ.flatten().tolist()
+        self.typ = np.array(typ, dtype=str)
+        self.r = np.random.rand(self.nr, 3).astype(float)
+        self.sigma = np.random.rand(self.nr, 3).astype(float)
+        self.eps = np.random.rand(self.nr).astype(float)
+        self.eps0 = 8.854187817e-12  # vacuum permittivity
+        self.avgeps = lambda eps1, eps2: (eps1 + eps2) / 2.0  # arithmetic mean
+        
+    def test_jflux_impl(self):
+        # calculate the reference result
+        ref = np.zeros((6, self.nr), dtype=float)
+        nx, ny, nz = self.nx, self.ny, self.nz
+        t = time.time()
+        for i in range(self.nr):
+            if self.typ[i] == "in":
+                ref[0, i] = self.sigma[i, 0] * self.eps0 * self.avgeps(self.eps[i - 1], self.eps[i]) \
+                            / np.abs(self.r[i - 1, 0] - self.r[i, 0])
+                ref[1, i] = self.sigma[i, 0] * self.eps0 * self.avgeps(self.eps[i + 1], self.eps[i]) \
+                            / np.abs(self.r[i + 1, 0] - self.r[i, 0])
+                ref[2, i] = self.sigma[i, 1] * self.eps0 * self.avgeps(self.eps[i - nx], self.eps[i]) \
+                            / np.abs(self.r[i - nx, 1] - self.r[i, 1])
+                ref[3, i] = self.sigma[i, 1] * self.eps0 * self.avgeps(self.eps[i + nx], self.eps[i]) \
+                            / np.abs(self.r[i + nx, 1] - self.r[i, 1])
+                ref[4, i] = self.sigma[i, 2] * self.eps0 * self.avgeps(self.eps[i - nx*ny], self.eps[i]) \
+                            / np.abs(self.r[i - nx*ny, 2] - self.r[i, 2])
+                ref[5, i] = self.sigma[i, 2] * self.eps0 * self.avgeps(self.eps[i + nx*ny], self.eps[i]) \
+                            / np.abs(self.r[i + nx*ny, 2] - self.r[i, 2])
+        print(f'old jflux_impl method took {time.time() - t:.4f} seconds')
+        # call the function
+        t = time.time()
+        jflux, i = _jflux_impl(self.nx, self.ny, self.nz, self.r, self.typ, 
+                               self.sigma, self.eps, self.eps0, self.avgeps, 
+                               with_index=True, full_size=True)
+        print(f'_jflux_impl took {time.time() - t:.4f} seconds')
+        # check if the result is equal to the reference result
+        self.assertTrue(jflux.shape == (6, self.nr))
+        self.assertTrue(i.shape == (len(self.typ[self.typ == "in"]),))
+        self.assertTrue(np.allclose(jflux, ref, rtol=1e-5, atol=1e-8))
+
+def _jacobian(inout, flux, i, nx, ny, nz, typ, zfree, phi, phi_, beta):
+    '''calculate the Jacobian matrix for the Poisson equation'''
+    mydisp_ = [-1, +1, -nx, +nx, -nx*ny, +nx*ny]
+    # diagonal elements
+    assert flux.shape == (6, len(i))
+    inout[i, i]  = -np.sum(flux, axis=0)
+    inout[i, i] -= beta * \
+        np.abs(coloumb(zfree[i])) * np.exp(-beta * np.sign(zfree[i]) * (phi[i] - phi_[i]))
+    # non-diagonal elements
+    for jf, d in zip(flux, mydisp_):
+        inout[i, i + d] = jf
+    # impose the boundary conditions
+    _impose_j_bound(inout, nx, ny, nz, typ, val=1.0, mask=-1.0)
+
+def _rhsvec(inout, flux, i, nx, ny, nz, typ, zfree, zfixed, phi, phi_, dirichlet_pot, beta):
+    '''calculate the right-hand side vector for the Poisson equation'''
+    # calculate the b vector
+    assert flux.shape == (6, len(i))
+    inout[i]  = np.sum(flux, axis=0)
+    inout[i] += coloumb(zfree[i]) * \
+        np.exp(-beta * np.sign(zfree[i]) * (phi[i] - phi_[i]))
+    inout[i] += coloumb(zfixed[i])
+    # impose the boundary conditions
+    _impose_b_bound(inout, nx, ny, nz, typ, phi, dirichlet_pot)
+    
+    inout *= -1 # for convenience change the sign of B in later NR iteration
+
+def calculate(jinout: Optional[lil_matrix], 
+              binout: Optional[np.ndarray],
+              grid_dim: Tuple[int, int, int]|List[int],
+              gridpoint_coords: np.ndarray, 
+              gridpoint_typ: List[str]|np.ndarray,
+              gridpoint_surfarea: np.ndarray, 
+              eps: np.ndarray, 
+              phi: np.ndarray, 
+              phi_: np.ndarray, 
+              free_chr: np.ndarray,
+              fixed_chr: Optional[np.ndarray] = None,
+              dirichlet_pot: Optional[np.ndarray] = None,
+              eps0: float = 8.854187817e-12, 
+              beta: float = 1/(1.380649e-23 * 300), 
+              feps: Callable[[float, float], float] = lambda eps1, eps2: (eps1 + eps2) / 2.0) -> None:
+    '''
+    refactored version of the implementation to calculate the Jacobian for the Poisson equation
+    using the Newton-Raphson method (finit difference method).
+    
+    NOTE: the input would be modified in-place
+    
+    Parameters
+    ----------
+    jinout : scipy.sparse.lil_matrix
+        The Jacobian matrix to be calculated. It is a sparse matrix in LIL format. If is None,
+        then the Jacobian will not be calculated.
+    binout : np.ndarray
+        The b vector to be calculated, which is the right-hand side of the linear system.
+        It is a dense numpy array. If is None, then the b vector will not be calculated.
+    grid_dim : tuple or list of int
+        The dimensions of the grid as (nx, ny, nz), where nx, ny, and nz are the number of grid 
+        points in the x, y, and z directions, respectively.
+    gridpoint_coords : np.ndarray
+        The coordinates of the grid points, shape (nr, 3), where nr is the total number of grid
+        points.
+    gridpoint_typ : list of str
+        The type of grid point, with length nr
+    gridpoint_surfarea : np.ndarray
+        The area of Voronoi surface in three dimension, shape (nr, 3). For example the gridpoint_
+        area[ir, 0] indicates the area of surface whose norm vector is along x-axis
+    eps : np.ndarray
+        The relative permittivity of each grid point
+    phi : np.ndarray
+        The electrostatic potential mapped on real-space grid point
+    phi\_ : np.ndarray
+        The "old" electrostatic potential mapped on real-space grid point
+    free_chr : np.ndarray
+        The free charge density mapped on real-space grid point
+    fixed_chr : np.ndarray, optional
+        The fixed charge density mapped on real-space grid point, shape (nr,). 
+    dirichlet_pot : np.ndarray, optional
+        The Dirichlet potential mapped on real-space grid point, shape (nr,). 
+    eps0 : float, optional
+        The vacuum permittivity, a constant value. Default is 8.854187817e-12 F/m.
+    beta : float, optional
+        The inverse temperature (1/kBT), used in the Boltzmann factor for charge density.
+        Default is 1/(1.380649e-23 * 300), which corresponds to room temperature (300 K). 
+    feps : callable, optional
+        the functor to calculate the average permittivity on a pair of grid points.
+        It should take two arguments, eps1 and eps2, and return the average permittivity.
+        If not provided, it defaults to the arithmetic mean.
+    '''
+    # shortcut: nothing to do
+    if jinout is None and binout is None:
+        return
+    
+    # check dimensions
+    if not isinstance(grid_dim, (tuple, list)) or len(grid_dim) != 3:
+        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
+    
+    # simple sanity check
+    if jinout is not None:
+        if not isinstance(jinout, lil_matrix):
+            raise TypeError("jinout (jacobian) must be a scipy.sparse.lil_matrix")
+        if jinout.shape != (nr, nr):
+            raise ValueError(f"jinout (jacobian) must have shape ({nr}, {nr}), but got {jinout.shape}")
+
+    if binout is not None:
+        if not isinstance(binout, np.ndarray):
+            raise TypeError("binout must be a numpy array")
+        if binout.shape != (nr,):
+            raise ValueError(f"binout must have shape ({nr},), but got {binout.shape}")
+    
+    # check the common input parameters
+    if not isinstance(gridpoint_coords, np.ndarray) or gridpoint_coords.shape != (nr, 3):
+        raise ValueError(f"gridpoint_coords must be a numpy array of shape ({nr}, 3)")
+    gridpoint_typ = np.array(list(gridpoint_typ.values()), dtype=str) \
+        if isinstance(gridpoint_typ, dict) else np.array(gridpoint_typ, dtype=str)
+    if not isinstance(gridpoint_typ, np.ndarray) or len(gridpoint_typ) != nr:
+        raise ValueError(f"gridpoint_typ must be a list of length {nr}")
+    if not isinstance(gridpoint_surfarea, np.ndarray) or gridpoint_surfarea.shape != (nr, 3):
+        raise ValueError(f"gridpoint_surfarea must be a numpy array of shape ({nr}, 3)")
+    if not isinstance(eps, np.ndarray) or eps.shape != (nr,):
+        raise ValueError(f"eps must be a numpy array of shape ({nr},)")
+    if not isinstance(phi, np.ndarray) or phi.shape != (nr,):
+        raise ValueError(f"phi must be a numpy array of shape ({nr},)")
+    if not isinstance(phi_, np.ndarray) or phi_.shape != (nr,):
+        raise ValueError(f"phi_ must be a numpy array of shape ({nr},)")
+    if not isinstance(free_chr, np.ndarray) or free_chr.shape != (nr,):
+        raise ValueError(f"free_chr must be a numpy array of shape ({nr},)")
+    
+    if binout is not None:
+        if not isinstance(fixed_chr, np.ndarray) or fixed_chr.shape != (nr,):
+            raise ValueError(f"fixed_chr must be a numpy array of shape ({nr},)")
+        if not isinstance(dirichlet_pot, np.ndarray) or dirichlet_pot.shape != (nr,):
+            raise ValueError(f"dirichlet_pot must be a numpy array of shape ({nr},)")
+
+    if not isinstance(eps0, (float, int)):
+        raise TypeError("eps0 must be a float or int")
+    if not isinstance(beta, (float, int)):
+        raise TypeError("beta must be a float or int")
+    if not callable(feps):
+        raise TypeError("feps must be a callable function with strictly two arguments: "
+                        "eps1 and eps2 to calculate the average permittivity")
+    
+    # ================================ implementation =======================================
+    # making alias for more succint code
+    r      = gridpoint_coords
+    typ    = np.array(gridpoint_typ, dtype=str)
+    sigma  = gridpoint_surfarea
+    zfree  = free_chr
+    zfixed = fixed_chr
+    
+    # calculate the flux of jacobian for all grid points of type "in"
+    jflux, ind = _jflux_impl(nx, ny, nz, r, typ, sigma, eps, eps0, feps, 
+                             with_index=True, full_size=False)
+    # the "ind" is the indices of the grid points of type "in"
+    assert jflux.shape == (6, len(ind)) # ensure the full_size is disabled
+    
+    if jinout is not None: # only calculate jacobian when jinout is provided (allocated)
+        _jacobian(jinout, jflux, ind, nx, ny, nz, typ, zfree, phi, phi_, beta)
+        assert jinout.shape == (nr, nr)
+
+    if binout is not None: # only calculate b vector when binout is provided (allocated)
+        # calculate bflux with jflux
+        bflux = _bflux_impl(jflux, ind, nx, ny, nz, phi, full_size=False)
+        assert bflux.shape == (6, len(ind))
+        _rhsvec(binout, bflux, ind, nx, ny, nz, typ, zfree, zfixed, phi, phi_, 
+                dirichlet_pot, beta)
+        assert binout.shape == (nr,)
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/dpnegf/negf/poisson_init.py b/dpnegf/negf/poisson_init.py
index a34417c..ceade33 100644
--- a/dpnegf/negf/poisson_init.py
+++ b/dpnegf/negf/poisson_init.py
@@ -638,80 +638,20 @@ def NR_construct_Jac_B(self,J,B):
         - Uses constants such as eps0 and elementary_charge, which must be defined in the scope.
         - The method modifies J and B in place.
         """
-        # construct the Jacobian and B for the Poisson equation
-        def average_eps(eps1, eps2, mode:str='harmonic'):
-            if mode == 'arithmetic':
-                return 0.5 * (eps1 + eps2)
-            elif mode == 'harmonic':
-                return 2.0 * eps1 * eps2 / (eps1 + eps2)
-        average_mode = self.average_mode
-        Nx = self.grid.shape[0];Ny = self.grid.shape[1];Nz = self.grid.shape[2]
-        for gp_index in range(self.grid.Np):
-            if self.boudnary_points[gp_index] == "in":
-                flux_xm_J = self.grid.surface_grid[gp_index,0]*eps0*average_eps(self.eps[gp_index-1],self.eps[gp_index],mode = average_mode)\
-                /abs(self.grid.grid_coord[gp_index,0]-self.grid.grid_coord[gp_index-1,0])
-                flux_xm_B = flux_xm_J*(self.phi[gp_index-1]-self.phi[gp_index])
-
-                flux_xp_J = self.grid.surface_grid[gp_index,0]*eps0*average_eps(self.eps[gp_index+1],self.eps[gp_index],mode = average_mode)\
-                /abs(self.grid.grid_coord[gp_index+1,0]-self.grid.grid_coord[gp_index,0])
-                flux_xp_B = flux_xp_J*(self.phi[gp_index+1]-self.phi[gp_index])
-                
-                flux_ym_J = self.grid.surface_grid[gp_index,1]*eps0*average_eps(self.eps[gp_index-Nx],self.eps[gp_index],mode = average_mode)\
-                /abs(self.grid.grid_coord[gp_index-Nx,1]-self.grid.grid_coord[gp_index,1])
-                flux_ym_B = flux_ym_J*(self.phi[gp_index-Nx]-self.phi[gp_index])
-
-                flux_yp_J = self.grid.surface_grid[gp_index,1]*eps0*average_eps(self.eps[gp_index+Nx],self.eps[gp_index],mode = average_mode)\
-                /abs(self.grid.grid_coord[gp_index+Nx,1]-self.grid.grid_coord[gp_index,1])
-                flux_yp_B = flux_yp_J*(self.phi[gp_index+Nx]-self.phi[gp_index])
-
-                flux_zm_J = self.grid.surface_grid[gp_index,2]*eps0*average_eps(self.eps[gp_index-Nx*Ny],self.eps[gp_index],mode = average_mode)\
-                /abs(self.grid.grid_coord[gp_index-Nx*Ny,2]-self.grid.grid_coord[gp_index,2])
-                flux_zm_B = flux_zm_J*(self.phi[gp_index-Nx*Ny]-self.phi[gp_index])
-
-                flux_zp_J = self.grid.surface_grid[gp_index,2]*eps0*average_eps(self.eps[gp_index+Nx*Ny],self.eps[gp_index],mode = average_mode)\
-                /abs(self.grid.grid_coord[gp_index+Nx*Ny,2]-self.grid.grid_coord[gp_index,2])
-                flux_zp_B = flux_zp_J*(self.phi[gp_index+Nx*Ny]-self.phi[gp_index])
-
-                # add flux term to matrix J
-                J[gp_index,gp_index] = -(flux_xm_J+flux_xp_J+flux_ym_J+flux_yp_J+flux_zm_J+flux_zp_J)\
-                    +elementary_charge*self.free_charge[gp_index]*(-np.sign(self.free_charge[gp_index]))/self.kBT*\
-                    np.exp(-np.sign(self.free_charge[gp_index])*(self.phi[gp_index]-self.phi_old[gp_index])/self.kBT)
-                J[gp_index,gp_index-1] = flux_xm_J
-                J[gp_index,gp_index+1] = flux_xp_J
-                J[gp_index,gp_index-Nx] = flux_ym_J
-                J[gp_index,gp_index+Nx] = flux_yp_J
-                J[gp_index,gp_index-Nx*Ny] = flux_zm_J
-                J[gp_index,gp_index+Nx*Ny] = flux_zp_J
-
-
-                # add flux term to matrix B
-                B[gp_index] = (flux_xm_B+flux_xp_B+flux_ym_B+flux_yp_B+flux_zm_B+flux_zp_B)
-                B[gp_index] += elementary_charge*self.free_charge[gp_index]*np.exp(-np.sign(self.free_charge[gp_index])\
-                    *(self.phi[gp_index]-self.phi_old[gp_index])/self.kBT)+elementary_charge*self.fixed_charge[gp_index]
-
-            else:# boundary points
-                J[gp_index,gp_index] = 1.0 # correct for both Dirichlet and Neumann boundary condition
-                
-                if self.boudnary_points[gp_index] == "xmin":   
-                    J[gp_index,gp_index+1] = -1.0
-                    B[gp_index] = (self.phi[gp_index]-self.phi[gp_index+1])
-                elif self.boudnary_points[gp_index] == "xmax":
-                    J[gp_index,gp_index-1] = -1.0
-                    B[gp_index] = (self.phi[gp_index]-self.phi[gp_index-1])
-                elif self.boudnary_points[gp_index] == "ymin":
-                    J[gp_index,gp_index+Nx] = -1.0
-                    B[gp_index] = (self.phi[gp_index]-self.phi[gp_index+Nx])
-                elif self.boudnary_points[gp_index] == "ymax":
-                    J[gp_index,gp_index-Nx] = -1.0
-                    B[gp_index] = (self.phi[gp_index]-self.phi[gp_index-Nx])
-                elif self.boudnary_points[gp_index] == "zmin":
-                    J[gp_index,gp_index+Nx*Ny] = -1.0
-                    B[gp_index] = (self.phi[gp_index]-self.phi[gp_index+Nx*Ny])
-                elif self.boudnary_points[gp_index] == "zmax":
-                    J[gp_index,gp_index-Nx*Ny] = -1.0
-                    B[gp_index] = (self.phi[gp_index]-self.phi[gp_index-Nx*Ny])
-                elif self.boudnary_points[gp_index] == "Dirichlet":
-                    B[gp_index] = (self.phi[gp_index]-self.lead_gate_potential[gp_index])
-
-            if B[gp_index]!=0: # for convenience change the sign of B in later NR iteration
-                B[gp_index] = -B[gp_index]           
\ No newline at end of file
+        from dpnegf.negf.newton_raphson_speed_up import calculate as my_speed_up_kernel
+        nx, ny, nz = self.grid.shape[:3]
+        my_speed_up_kernel(jinout=J, 
+                           binout=B, 
+                           grid_dim=(nx, ny, nz),
+                           gridpoint_coords=self.grid.grid_coord,
+                           gridpoint_typ=self.boudnary_points,
+                           gridpoint_surfarea=self.grid.surface_grid,
+                           eps=self.eps,
+                           phi=self.phi,
+                           phi_=self.phi_old,
+                           free_chr=self.free_charge,
+                           fixed_chr=self.fixed_charge,
+                           dirichlet_pot=self.lead_gate_potential,
+                           eps0=eps0,
+                           beta=1.0/self.kBT)
+

From 6215bce4d491884c5fb84cb8e2f5a46fa2771f61 Mon Sep 17 00:00:00 2001
From: kirk0830 <yike.huang@aliyun.com>
Date: Wed, 9 Jul 2025 18:21:59 +0800
Subject: [PATCH 02/14] recover the functionality of epsilon averaging

---
 dpnegf/negf/poisson_init.py | 37 ++++++++++++++++++++++---------------
 1 file changed, 22 insertions(+), 15 deletions(-)

diff --git a/dpnegf/negf/poisson_init.py b/dpnegf/negf/poisson_init.py
index ceade33..9dc9e59 100644
--- a/dpnegf/negf/poisson_init.py
+++ b/dpnegf/negf/poisson_init.py
@@ -638,20 +638,27 @@ def NR_construct_Jac_B(self,J,B):
         - Uses constants such as eps0 and elementary_charge, which must be defined in the scope.
         - The method modifies J and B in place.
         """
-        from dpnegf.negf.newton_raphson_speed_up import calculate as my_speed_up_kernel
+        from dpnegf.negf.newton_raphson_speed_up import calculate as construct
         nx, ny, nz = self.grid.shape[:3]
-        my_speed_up_kernel(jinout=J, 
-                           binout=B, 
-                           grid_dim=(nx, ny, nz),
-                           gridpoint_coords=self.grid.grid_coord,
-                           gridpoint_typ=self.boudnary_points,
-                           gridpoint_surfarea=self.grid.surface_grid,
-                           eps=self.eps,
-                           phi=self.phi,
-                           phi_=self.phi_old,
-                           free_chr=self.free_charge,
-                           fixed_chr=self.fixed_charge,
-                           dirichlet_pot=self.lead_gate_potential,
-                           eps0=eps0,
-                           beta=1.0/self.kBT)
+        feps = lambda eps1, eps2: (eps1 + eps2) / 2.0
+        if self.average_mode == 'harmonic':
+            feps = lambda eps1, eps2: 2.0 * eps1 * eps2 / (eps1 + eps2)
+        else:
+            assert self.average_mode == 'arithmetic'
+        
+        construct(jinout=J, 
+                  binout=B, 
+                  grid_dim=(nx, ny, nz),
+                  gridpoint_coords=self.grid.grid_coord,
+                  gridpoint_typ=self.boudnary_points,
+                  gridpoint_surfarea=self.grid.surface_grid,
+                  eps=self.eps,
+                  phi=self.phi,
+                  phi_=self.phi_old,
+                  free_chr=self.free_charge,
+                  fixed_chr=self.fixed_charge,
+                  dirichlet_pot=self.lead_gate_potential,
+                  eps0=eps0,
+                  beta=1.0/self.kBT,
+                  feps=feps)
 

From 7d4bd2e6811bffe27a7b0db75c7ad9295e4262f3 Mon Sep 17 00:00:00 2001
From: kirk0830 <yike.huang@aliyun.com>
Date: Wed, 9 Jul 2025 18:38:34 +0800
Subject: [PATCH 03/14] make unittest to be more practical

---
 dpnegf/negf/newton_raphson_speed_up.py | 90 ++++++++++++++++----------
 1 file changed, 55 insertions(+), 35 deletions(-)

diff --git a/dpnegf/negf/newton_raphson_speed_up.py b/dpnegf/negf/newton_raphson_speed_up.py
index 078635f..98238b3 100644
--- a/dpnegf/negf/newton_raphson_speed_up.py
+++ b/dpnegf/negf/newton_raphson_speed_up.py
@@ -27,11 +27,15 @@ def setUp(self):
         self.nx, self.ny, self.nz = 100, 100, 100
         self.nr = self.nx * self.ny * self.nz
         # say we have 6 boundary types, each has 300 points
-        typ = ['xmin'] * 30 + ['xmax'] * 30 + ['ymin'] * 30 \
-            + ['ymax'] * 30 + ['zmin'] * 30 + ['zmax'] * 30
-        # and all others are of type "in"
-        typ += ['in'] * (self.nr - 180)
-        self.typ = np.array(typ, dtype=str)
+        typ = np.array(['****'] * self.nr, dtype=str).reshape((self.nx, self.ny, self.nz))
+        typ[typ == '****'] = 'in'
+        typ[   0,    :,  :] = 'xmin'
+        typ[  -1,    :,  :] = 'xmax'
+        typ[1:-1,    0,  :] = 'ymin'
+        typ[1:-1,   -1,  :] = 'ymax'
+        typ[1:-1, 1:-1,  0] = 'zmin'
+        typ[1:-1, 1:-1, -1] = 'zmax'
+        self.typ = typ.flatten()
         
     def test_impose_bound(self):
         # create a dummy inout matrix
@@ -67,13 +71,26 @@ def test_impose_bound(self):
         print(f'_impose_j_bound took {time.time() - t:.4f} seconds')
     
         # check if the inout matrix is equal to the reference matrix
+        # check all diagonal elements
         self.assertTrue(all(inout[i, i] == ref[i, i] for i in range(self.nr)))
-        self.assertTrue(all(inout[i, i + 1] == ref[i, i + 1] for i in range(30)))
-        self.assertTrue(all(inout[i, i - 1] == ref[i, i - 1] for i in range(30, 60)))
-        self.assertTrue(all(inout[i, i + nx] == ref[i, i + nx] for i in range(60, 90)))
-        self.assertTrue(all(inout[i, i - nx] == ref[i, i - nx] for i in range(90, 120)))
-        self.assertTrue(all(inout[i, i + nx*ny] == ref[i, i + nx*ny] for i in range(120, 150)))
-        self.assertTrue(all(inout[i, i - nx*ny] == ref[i, i - nx*ny] for i in range(150, 180)))
+        # check xmin points
+        ind_xmin = np.where(self.typ == 'xmin')[0]
+        self.assertTrue(all(inout[i, i + 1] == ref[i, i + 1] for i in ind_xmin))
+        # check xmax points
+        ind_xmax = np.where(self.typ == 'xmax')[0]
+        self.assertTrue(all(inout[i, i - 1] == ref[i, i - 1] for i in ind_xmax))
+        # check ymin points
+        ind_ymin = np.where(self.typ == 'ymin')[0]
+        self.assertTrue(all(inout[i, i + nx] == ref[i, i + nx] for i in ind_ymin))
+        # check ymax points
+        ind_ymax = np.where(self.typ == 'ymax')[0]
+        self.assertTrue(all(inout[i, i - nx] == ref[i, i - nx] for i in ind_ymax))
+        # check zmin points
+        ind_zmin = np.where(self.typ == 'zmin')[0]
+        self.assertTrue(all(inout[i, i + nx*ny] == ref[i, i + nx*ny] for i in ind_zmin))
+        # check zmax points
+        ind_zmax = np.where(self.typ == 'zmax')[0]
+        self.assertTrue(all(inout[i, i - nx*ny] == ref[i, i - nx*ny] for i in ind_zmax))
 
 def _impose_b_bound(inout, nx, ny, nz, typ, phi, dirichlet_pot):
     '''impose the special mask for boundary points in b vector'''
@@ -94,12 +111,15 @@ def setUp(self):
         # assuming there are 1e6 grid points, each direction has 1e3
         self.nx, self.ny, self.nz = 100, 100, 100
         self.nr = self.nx * self.ny * self.nz
-        # say we have 6 boundary types, each has 300 points
-        typ = ['xmin'] * 30 + ['xmax'] * 30 + ['ymin'] * 30 \
-            + ['ymax'] * 30 + ['zmin'] * 30 + ['zmax'] * 30 + ['Dirichlet'] * 30
-        # and all others are of type "in"
-        typ += ['in'] * (self.nr - 210)
-        self.typ = np.array(typ, dtype=str)
+        typ = np.array(['****'] * self.nr, dtype=str).reshape((self.nx, self.ny, self.nz))
+        typ[typ == '****'] = 'in'
+        typ[   0,    :,  :] = 'xmin'
+        typ[  -1,    :,  :] = 'xmax'
+        typ[1:-1,    0,  :] = 'ymin'
+        typ[1:-1,   -1,  :] = 'ymax'
+        typ[1:-1, 1:-1,  0] = 'zmin'
+        typ[1:-1, 1:-1, -1] = 'zmax'
+        self.typ = typ.flatten()
         self.phi = np.random.rand(self.nr).astype(float)
         self.fixed_pot = np.random.rand(self.nr).astype(float)
         
@@ -161,15 +181,15 @@ def setUp(self):
         self.nx, self.ny, self.nz = 100, 100, 100
         self.nr = self.nx * self.ny * self.nz
         # set the boundary types
-        typ = np.array(['in'] * self.nr, dtype=str).reshape((self.nx, self.ny, self.nz))
-        typ[ 0,  :,  :] = 'xmin'
-        typ[-1,  :,  :] = 'xmax'
-        typ[ :,  0,  :] = 'ymin'
-        typ[ :, -1,  :] = 'ymax'
-        typ[ :,  :,  0] = 'zmin'
-        typ[ :,  :, -1] = 'zmax'
-        typ = typ.flatten().tolist()
-        self.typ = np.array(typ, dtype=str)
+        typ = np.array(['****'] * self.nr, dtype=str).reshape((self.nx, self.ny, self.nz))
+        typ[typ == '****'] = 'in'
+        typ[   0,    :,  :] = 'xmin'
+        typ[  -1,    :,  :] = 'xmax'
+        typ[1:-1,    0,  :] = 'ymin'
+        typ[1:-1,   -1,  :] = 'ymax'
+        typ[1:-1, 1:-1,  0] = 'zmin'
+        typ[1:-1, 1:-1, -1] = 'zmax'
+        self.typ = typ.flatten()
         self.r = np.random.rand(self.nr, 3).astype(float)
         self.sigma = np.random.rand(self.nr, 3).astype(float)
         self.eps = np.random.rand(self.nr).astype(float)
@@ -227,15 +247,15 @@ def setUp(self):
         self.nx, self.ny, self.nz = 100, 100, 100
         self.nr = self.nx * self.ny * self.nz
         # set the boundary types
-        typ = np.array(['in'] * self.nr, dtype=str).reshape((self.nx, self.ny, self.nz))
-        typ[ 0,  :,  :] = 'xmin'
-        typ[-1,  :,  :] = 'xmax'
-        typ[ :,  0,  :] = 'ymin'
-        typ[ :, -1,  :] = 'ymax'
-        typ[ :,  :,  0] = 'zmin'
-        typ[ :,  :, -1] = 'zmax'
-        typ = typ.flatten().tolist()
-        self.typ = np.array(typ, dtype=str)
+        typ = np.array(['****'] * self.nr, dtype=str).reshape((self.nx, self.ny, self.nz))
+        typ[typ == '****'] = 'in'
+        typ[   0,    :,  :] = 'xmin'
+        typ[  -1,    :,  :] = 'xmax'
+        typ[1:-1,    0,  :] = 'ymin'
+        typ[1:-1,   -1,  :] = 'ymax'
+        typ[1:-1, 1:-1,  0] = 'zmin'
+        typ[1:-1, 1:-1, -1] = 'zmax'
+        self.typ = typ.flatten()
         self.r = np.random.rand(self.nr, 3).astype(float)
         self.sigma = np.random.rand(self.nr, 3).astype(float)
         self.eps = np.random.rand(self.nr).astype(float)

From b0e149f35fce83f5f23b38ad7dc1221f9ac0aff8 Mon Sep 17 00:00:00 2001
From: kirk0830 <yike.huang@aliyun.com>
Date: Wed, 9 Jul 2025 18:43:44 +0800
Subject: [PATCH 04/14] update the annotation of function NR_construct_Jac_B

---
 dpnegf/negf/poisson_init.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/dpnegf/negf/poisson_init.py b/dpnegf/negf/poisson_init.py
index 9dc9e59..423c1ed 100644
--- a/dpnegf/negf/poisson_init.py
+++ b/dpnegf/negf/poisson_init.py
@@ -625,12 +625,14 @@ def NR_construct_Jac_B(self,J,B):
         For boundary points, the method applies appropriate boundary conditions (Dirichlet or Neumann) by 
         modifying J and B accordingly, based on the type of boundary (xmin, xmax, ymin, ymax, zmin, zmax, or Dirichlet).
         After assembling the contributions, the sign of B is flipped for nonzero entries for the Newton-Raphson iteration.
+        
         Parameters
         ----------
         J : numpy.ndarray
             The Jacobian matrix to be constructed/updated (shape: [Np, Np], where Np is the number of grid points).
         B : numpy.ndarray
             The right-hand side vector to be constructed/updated (shape: [Np]).
+            
         Notes
         -----
         - Assumes that self.grid, self.eps, self.phi, self.phi_old, self.free_charge, self.fixed_charge, 

From fa97971f967e8f8165c3abe22ba925f8834868ad Mon Sep 17 00:00:00 2001
From: kirk0830 <yike.huang@aliyun.com>
Date: Wed, 9 Jul 2025 18:45:11 +0800
Subject: [PATCH 05/14] add one empty line

---
 dpnegf/negf/newton_raphson_speed_up.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/dpnegf/negf/newton_raphson_speed_up.py b/dpnegf/negf/newton_raphson_speed_up.py
index 98238b3..888fc19 100644
--- a/dpnegf/negf/newton_raphson_speed_up.py
+++ b/dpnegf/negf/newton_raphson_speed_up.py
@@ -195,6 +195,7 @@ def setUp(self):
         self.eps = np.random.rand(self.nr).astype(float)
         self.eps0 = 8.854187817e-12
         self.avgeps = lambda eps1, eps2: (eps1 + eps2) / 2.0  # arithmetic mean
+
     def test_bflux_impl(self):
         ind = np.where(self.typ == "in")[0]
         jflux = np.random.rand(6, self.nr).astype(float)

From fdcdcc51319da196cc1fb6b26bb26411fd62c4eb Mon Sep 17 00:00:00 2001
From: kirk0830 <yike.huang@aliyun.com>
Date: Thu, 10 Jul 2025 14:47:05 +0800
Subject: [PATCH 06/14] format some iteration output

---
 dpnegf/negf/poisson_init.py | 21 ++++++++++++++-------
 dpnegf/runner/NEGF.py       | 11 ++++++-----
 2 files changed, 20 insertions(+), 12 deletions(-)

diff --git a/dpnegf/negf/poisson_init.py b/dpnegf/negf/poisson_init.py
index 423c1ed..46efde0 100644
--- a/dpnegf/negf/poisson_init.py
+++ b/dpnegf/negf/poisson_init.py
@@ -1,3 +1,6 @@
+import time
+import logging
+
 import numpy as np 
 # import pyamg #TODO: later add it to optional dependencies,like sisl
 # from pyamg.gallery import poisson
@@ -6,11 +9,10 @@
 from scipy.constants import epsilon_0 as eps0  #TODO:later add to untils.constants.py
 from scipy.sparse import csr_matrix
 from scipy.sparse.linalg import spsolve
-import logging
+
 #eps0 = 8.854187817e-12 # in the unit of F/m
 # As length in deeptb is in the unit of Angstrom, the unit of eps0 is F/Angstrom
 eps0 = eps0*1e-10 # in the unit of F/Angstrom
-
 log = logging.getLogger(__name__)
 
 class Grid(object):
@@ -543,7 +545,12 @@ def solve_poisson_NRcycle(self,method='pyamg',tolerance=1e-7,dtype:str='float64'
         else:
             raise ValueError('Unknown data type: ',dtype)
 
+        log.info("Start solving Poisson equation (Newton-Raphson)")
+        log.info("-"*(8+1+15+1+15+1+6))
+        log.info(f"{'ITERSTEP':>8s} {'|| DPHI ||':>15s} {'MAX(DPHI)':>15s} {'TIME/s':>6s}")
+        log.info("-"*(8+1+15+1+15+1+6))
         while norm_delta_phi > 1e-3 and NR_cycle_step < 100:
+            t = time.time()
             # obtain the Jacobian and B for the Poisson equation
             Jacobian,B = self.to_scipy_Jac_B(dtype=dtype)
             norm_B = np.linalg.norm(B)
@@ -551,7 +558,7 @@ def solve_poisson_NRcycle(self,method='pyamg',tolerance=1e-7,dtype:str='float64'
             if method == 'scipy':   #TODO: rename to 'Direct
                 if NR_cycle_step == 0:
                     log.info(msg="Solve Poisson equation by scipy")
-                delta_phi = spsolve(Jacobian,B)
+                delta_phi = spsolve(Jacobian, B)
             elif method == 'pyamg': #TODO: rename to 'AMG'
                 if NR_cycle_step == 0:
                     log.info(msg="Solve Poisson equation by pyamg")
@@ -564,7 +571,7 @@ def solve_poisson_NRcycle(self,method='pyamg',tolerance=1e-7,dtype:str='float64'
             self.phi += delta_phi
 
             if norm_delta_phi > 1e-3:
-                _,B = self.to_scipy_Jac_B()
+                _, B = self.to_scipy_Jac_B()
                 norm_B_new = np.linalg.norm(B)
                 control_count = 1
                 # control the norm of B to avoid larger norm_B after one NR cycle
@@ -572,13 +579,13 @@ def solve_poisson_NRcycle(self,method='pyamg',tolerance=1e-7,dtype:str='float64'
                     if control_count==1: 
                         log.warning(msg="norm_B increase after this  NR cycle, contorler starts!")
                     self.phi -= delta_phi/np.power(2,control_count)
-                    _,B = self.to_scipy_Jac_B()
+                    _, B = self.to_scipy_Jac_B()
                     norm_B_new = np.linalg.norm(B)
                     control_count += 1
                     log.info(msg="    control_count: {:.1f}   norm_B_new: {:.5f}".format(float(control_count),norm_B_new))    
-                               
+
             NR_cycle_step += 1
-            log.info(msg="  NR cycle step: {:d}   norm_delta_phi: {:.8f}   max_delta_phi: {:.8f}".format(int(NR_cycle_step),norm_delta_phi,max_delta_phi))
+            log.info(f"{NR_cycle_step:>8} {norm_delta_phi:>15.8e} {max_delta_phi:>15.8e} {time.time()-t:>6.2f}")
         
         max_diff = np.max(abs(self.phi-self.phi_old))
         return max_diff
diff --git a/dpnegf/runner/NEGF.py b/dpnegf/runner/NEGF.py
index ab2d4d4..b576096 100644
--- a/dpnegf/runner/NEGF.py
+++ b/dpnegf/runner/NEGF.py
@@ -457,7 +457,8 @@ def poisson_negf_scf(self,interface_poisson,atom_gridpoint_index,err=1e-6,max_it
                 interface_poisson.phi = mixer.update(interface_poisson.phi.copy(), interface_poisson.phi_old.copy())
 
             iter_count += 1 # Gummel type iteration
-            log.info(msg="Poisson-NEGF iteration: {}    Potential Diff Maximum: {}\n".format(iter_count,max_diff_phi))
+            log.info(f"Poisson-NEGF iteration: {iter_count}")
+            log.info(f"Potential Diff Maximum: {max_diff_phi:>12.8e}\n")
             max_diff_list.append(max_diff_phi)
 
             if max_diff_phi <= err:
@@ -520,7 +521,7 @@ def negf_compute(self,scf_require=False,Vbias=None):
             self.out['k'].append(k)
             self.out['wk'].append(self.wk[ik])
             self.free_charge.update({str(k):torch.zeros_like(torch.tensor(self.device_atom_norbs),dtype=torch.complex128)})
-            log.info(msg="Properties computation at k = [{:.4f},{:.4f},{:.4f}]".format(float(k[0]),float(k[1]),float(k[2])))
+            log.info(f"Properties computation at k = ({', '.join([f'{kk:>6.4f}' for kk in k])})")
 
             if scf_require:
                 if self.density_options["method"] == "Fiori":                    
@@ -574,7 +575,7 @@ def negf_compute(self,scf_require=False,Vbias=None):
                     if output_freq == 0: output_freq = 1
                     for ie, e in enumerate(self.uni_grid):
                         if ie % output_freq == 0:
-                            log.info(msg="computing green's function at e = {:.3f}".format(float(e)))
+                            log.info(f" computing green's function at e = {float(e):>6.3f}")
                         if self.scf:
                             if not self.poisson_options["with_Dirichlet_leads"]:
                                 for ll in self.stru_options.keys():
@@ -698,9 +699,9 @@ def negf_compute(self,scf_require=False,Vbias=None):
                     # TODO: check the following code for multiple k points calculation
                     if self.out_lcurrent:
                         lcurrent = 0
-                        log.info(msg="computing local current at k = [{:.4f},{:.4f},{:.4f}]".format(float(k[0]),float(k[1]),float(k[2])))
+                        log.info(f"Properties computation at k = ({', '.join([f'{kk:>6.4f}' for kk in k])})")
                         for i, e in enumerate(self.int_grid):
-                            log.info(msg=" computing green's function at e = {:.3f}".format(float(e)))
+                            log.info(f" computing green's function at e = {float(e):>6.3f}")
                             for ll in self.stru_options.keys():
                                 if ll.startswith("lead"):
                                     getattr(self.deviceprop, ll).self_energy(

From eb9bd64ccd0c71dd47412c88cf170d3b961e1ce8 Mon Sep 17 00:00:00 2001
From: kirk0830 <yike.huang@aliyun.com>
Date: Thu, 10 Jul 2025 23:24:11 +0800
Subject: [PATCH 07/14] format code

---
 dpnegf/negf/poisson_init.py | 192 +++++++++++++++++++++++++++---------
 1 file changed, 143 insertions(+), 49 deletions(-)

diff --git a/dpnegf/negf/poisson_init.py b/dpnegf/negf/poisson_init.py
index 46efde0..c1e5df9 100644
--- a/dpnegf/negf/poisson_init.py
+++ b/dpnegf/negf/poisson_init.py
@@ -86,7 +86,8 @@ def __init__(self,xg,yg,zg,xa,ya,za):
         assert self.Np == len(xmesh)
         assert self.grid_coord.shape[0] == self.Np
         
-        log.info(msg="Number of grid points: {:.1f}   Number of atoms: {:.1f}".format(float(self.Np),self.Na))
+        log.info(f"Number of grid points: {self.Np}")
+        log.info(f"Number of atoms: {self.Na}")
         # print('Number of grid points: ',self.Np,' grid shape: ',self.grid_coord.shape,' Number of atoms: ',self.Na)
 
         # find the index of the atoms in the grid
@@ -162,7 +163,6 @@ def cal_vorlen(self,x):
             xd[i] = (abs(x[i]-x[i-1])+abs(x[i]-x[i+1]))/2
         return xd
 
-
 class region(object):
     """
     A class representing a 3D rectangular region defined by ranges along the x, y, and z axes.
@@ -214,7 +214,6 @@ def __init__(self,x_range,y_range,z_range):
         # Fermi_level of gate (in unit eV)
         self.Ef = 0.0        
 
-
 class Dielectric(region):
     """
     Represents a dielectric region with a specified permittivity.
@@ -241,9 +240,6 @@ def __init__(self,x_range,y_range,z_range):
         # dielectric permittivity
         self.eps = 1.0
 
-
-
-
 class Interface3D(object):
     """
     Interface3D(grid, Dirichlet_group, dielectric_group)
@@ -436,7 +432,7 @@ def get_potential_eps(self, region_list):
             else:
                 raise ValueError('Unknown region type: ',region_list[i].__class__.__name__)
         
-        log.info(msg="Number of Dirichlet points: {:.1f}".format(float(Dirichlet_point)))
+        log.info(f"Number of Dirichlet points: {Dirichlet_point}")
         
         
     def to_pyamg_Jac_B(self,dtype=np.float64):
@@ -461,9 +457,9 @@ def to_pyamg_Jac_B(self,dtype=np.float64):
         B = np.zeros(Jacobian.shape[0],dtype=Jacobian.dtype)
 
         Jacobian_lil = Jacobian.tolil()
-        self.NR_construct_Jac_B(Jacobian_lil,B)
+        self.NR_construct_Jac_B(Jacobian_lil, B)
         Jacobian = Jacobian_lil.tocsr() 
-        return Jacobian,B
+        return Jacobian, B
     
     
     def to_scipy_Jac_B(self,dtype=np.float64):
@@ -488,14 +484,15 @@ def to_scipy_Jac_B(self,dtype=np.float64):
         B = np.zeros(Jacobian.shape[0],dtype=Jacobian.dtype)
 
         Jacobian_lil = Jacobian.tolil()
-        self.NR_construct_Jac_B(Jacobian_lil,B)
+        self.NR_construct_Jac_B(Jacobian_lil, B)
         Jacobian = Jacobian_lil.tocsr() 
         # self.construct_poisson(A,b)
-        return Jacobian,B
-    
+        return Jacobian, B
 
-
-    def solve_poisson_NRcycle(self,method='pyamg',tolerance=1e-7,dtype:str='float64'):
+    def solve_poisson_NRcycle(self,
+                              method='pyamg',
+                              tolerance=1e-7,
+                              dtype:str='float64'):
         """
         Solve the Poisson equation using the Newton-Raphson (NR) iterative method.
         This NR method is inspired by NanoTCAD ViDES (http://vides.nanotcad.com/vides/),iteratively solving 
@@ -533,40 +530,45 @@ def solve_poisson_NRcycle(self,method='pyamg',tolerance=1e-7,dtype:str='float64'
         # solve the Poisson equation with Newton-Raphson method
         # delta_phi: the correction on the potential
         # It has been tested that dtype='float64' is a more stable SCF choice.
-      
-        
         norm_delta_phi = 1.0 #  Euclidean norm of delta_phi in each step
         NR_cycle_step = 0
+        
+        # initialize the precision
+        try:
+            dtype = {'float64': np.float64, 'float32': np.float32}[dtype]
+        except KeyError:
+            raise ValueError(f"Unknown data type: {dtype}. Use 'float64' or 'float32'.")
 
-        if dtype == 'float64':
-            dtype = np.float64
-        elif dtype == 'float32':
-            dtype = np.float32
-        else:
-            raise ValueError('Unknown data type: ',dtype)
+        # initialize the solver
+        assert method in ['scipy', 'pyamg']
+        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
 
+        # print iteration header
         log.info("Start solving Poisson equation (Newton-Raphson)")
         log.info("-"*(8+1+15+1+15+1+6))
         log.info(f"{'ITERSTEP':>8s} {'|| DPHI ||':>15s} {'MAX(DPHI)':>15s} {'TIME/s':>6s}")
         log.info("-"*(8+1+15+1+15+1+6))
+
+        # start iteration
         while norm_delta_phi > 1e-3 and NR_cycle_step < 100:
+            # start / refresh the timer
             t = time.time()
+            t0 = t # for timing the whole step
+
             # obtain the Jacobian and B for the Poisson equation
-            Jacobian,B = self.to_scipy_Jac_B(dtype=dtype)
+            Jacobian, B = self.to_scipy_Jac_B(dtype=dtype)
+            log.info(f'  matrix construction time: {time.time()-t:.2f} s')
+            t = time.time()
             norm_B = np.linalg.norm(B)
-           
-            if method == 'scipy':   #TODO: rename to 'Direct
-                if NR_cycle_step == 0:
-                    log.info(msg="Solve Poisson equation by scipy")
-                delta_phi = spsolve(Jacobian, B)
-            elif method == 'pyamg': #TODO: rename to 'AMG'
-                if NR_cycle_step == 0:
-                    log.info(msg="Solve Poisson equation by pyamg")
-                delta_phi = self.solver_pyamg(Jacobian,B,tolerance=1e-5)
-            else:
-                raise ValueError('Unknown Poisson solver: ',method)
-                        
-            max_delta_phi = np.max(abs(delta_phi))
+            
+            # solve
+            delta_phi = solve(Jacobian, B)
+            log.info(f'  sparse matrix solve time: {time.time()-t:.2f} s')
+            t = time.time()
+            
+            max_delta_phi = np.max(np.abs(delta_phi))
             norm_delta_phi = np.linalg.norm(delta_phi)
             self.phi += delta_phi
 
@@ -576,18 +578,31 @@ def solve_poisson_NRcycle(self,method='pyamg',tolerance=1e-7,dtype:str='float64'
                 control_count = 1
                 # control the norm of B to avoid larger norm_B after one NR cycle
                 while norm_B_new > norm_B and control_count < 2:
-                    if control_count==1: 
-                        log.warning(msg="norm_B increase after this  NR cycle, contorler starts!")
-                    self.phi -= delta_phi/np.power(2,control_count)
+                    if control_count == 1: 
+                        log.warning("norm_B increase after this cycle, controller starts!")
+                    # reduce the correction on the potential
+                    self.phi -= delta_phi/np.power(2, control_count)
+                    
+                    # re-construct the Jacobian and B
                     _, B = self.to_scipy_Jac_B()
+                    
+                    # re-calculate the norm of B
                     norm_B_new = np.linalg.norm(B)
                     control_count += 1
-                    log.info(msg="    control_count: {:.1f}   norm_B_new: {:.5f}".format(float(control_count),norm_B_new))    
-
+                    
+                    log.info(f"    ic = {control_count}"            # control_count
+                             f"    || B_new || = {norm_B_new:.5e}") # norm_B_new
+                
+                # print
+                log.info(f'  RHS-vector regulation time: {time.time()-t:.2f} s')
+                t = time.time()
+                
             NR_cycle_step += 1
-            log.info(f"{NR_cycle_step:>8} {norm_delta_phi:>15.8e} {max_delta_phi:>15.8e} {time.time()-t:>6.2f}")
+            # print
+            log.info(f"{NR_cycle_step:>8} {norm_delta_phi:>15.8e} "
+                     f"{max_delta_phi:>15.8e} {time.time()-t0:>6.2f}")
         
-        max_diff = np.max(abs(self.phi-self.phi_old))
+        max_diff = np.max(np.abs(self.phi - self.phi_old))
         return max_diff
 
     def solver_pyamg(self,A,b,tolerance=1e-7,accel=None):
@@ -649,12 +664,12 @@ def NR_construct_Jac_B(self,J,B):
         """
         from dpnegf.negf.newton_raphson_speed_up import calculate as construct
         nx, ny, nz = self.grid.shape[:3]
-        feps = lambda eps1, eps2: (eps1 + eps2) / 2.0
-        if self.average_mode == 'harmonic':
-            feps = lambda eps1, eps2: 2.0 * eps1 * eps2 / (eps1 + eps2)
-        else:
-            assert self.average_mode == 'arithmetic'
-        
+        # https://manual.gromacs.org/documentation/current/reference-manual/topologies/parameter-files.html#nbpar
+        feps = {'harmonic'  : lambda eps1, eps2: 2.0 * eps1 * eps2 / (eps1 + eps2),
+                'arithmetic': lambda eps1, eps2: 0.5 * (eps1 + eps2),
+                'geometric' : lambda eps1, eps2: np.sqrt(eps1 * eps2), # gromacs comb-rule 2
+                }[self.average_mode]
+
         construct(jinout=J, 
                   binout=B, 
                   grid_dim=(nx, ny, nz),
@@ -671,3 +686,82 @@ def NR_construct_Jac_B(self,J,B):
                   beta=1.0/self.kBT,
                   feps=feps)
 
+        # ===================================== OLD IMPLEMENTATION =========================================
+        # construct the Jacobian and B for the Poisson equation
+        # def average_eps(eps1, eps2, mode:str='harmonic'):
+        #     if mode == 'arithmetic':
+        #         return 0.5 * (eps1 + eps2)
+        #     elif mode == 'harmonic':
+        #         return 2.0 * eps1 * eps2 / (eps1 + eps2)
+        # average_mode = self.average_mode
+        # Nx = self.grid.shape[0];Ny = self.grid.shape[1];Nz = self.grid.shape[2]
+        # for gp_index in range(self.grid.Np):
+        #     if self.boudnary_points[gp_index] == "in":
+        #         flux_xm_J = self.grid.surface_grid[gp_index,0]*eps0*average_eps(self.eps[gp_index-1],self.eps[gp_index],mode = average_mode)\
+        #         /abs(self.grid.grid_coord[gp_index,0]-self.grid.grid_coord[gp_index-1,0])
+        #         flux_xm_B = flux_xm_J*(self.phi[gp_index-1]-self.phi[gp_index])
+
+        #         flux_xp_J = self.grid.surface_grid[gp_index,0]*eps0*average_eps(self.eps[gp_index+1],self.eps[gp_index],mode = average_mode)\
+        #         /abs(self.grid.grid_coord[gp_index+1,0]-self.grid.grid_coord[gp_index,0])
+        #         flux_xp_B = flux_xp_J*(self.phi[gp_index+1]-self.phi[gp_index])
+                
+        #         flux_ym_J = self.grid.surface_grid[gp_index,1]*eps0*average_eps(self.eps[gp_index-Nx],self.eps[gp_index],mode = average_mode)\
+        #         /abs(self.grid.grid_coord[gp_index-Nx,1]-self.grid.grid_coord[gp_index,1])
+        #         flux_ym_B = flux_ym_J*(self.phi[gp_index-Nx]-self.phi[gp_index])
+
+        #         flux_yp_J = self.grid.surface_grid[gp_index,1]*eps0*average_eps(self.eps[gp_index+Nx],self.eps[gp_index],mode = average_mode)\
+        #         /abs(self.grid.grid_coord[gp_index+Nx,1]-self.grid.grid_coord[gp_index,1])
+        #         flux_yp_B = flux_yp_J*(self.phi[gp_index+Nx]-self.phi[gp_index])
+
+        #         flux_zm_J = self.grid.surface_grid[gp_index,2]*eps0*average_eps(self.eps[gp_index-Nx*Ny],self.eps[gp_index],mode = average_mode)\
+        #         /abs(self.grid.grid_coord[gp_index-Nx*Ny,2]-self.grid.grid_coord[gp_index,2])
+        #         flux_zm_B = flux_zm_J*(self.phi[gp_index-Nx*Ny]-self.phi[gp_index])
+
+        #         flux_zp_J = self.grid.surface_grid[gp_index,2]*eps0*average_eps(self.eps[gp_index+Nx*Ny],self.eps[gp_index],mode = average_mode)\
+        #         /abs(self.grid.grid_coord[gp_index+Nx*Ny,2]-self.grid.grid_coord[gp_index,2])
+        #         flux_zp_B = flux_zp_J*(self.phi[gp_index+Nx*Ny]-self.phi[gp_index])
+
+        #         # add flux term to matrix J
+        #         J[gp_index,gp_index] = -(flux_xm_J+flux_xp_J+flux_ym_J+flux_yp_J+flux_zm_J+flux_zp_J)\
+        #             +elementary_charge*self.free_charge[gp_index]*(-np.sign(self.free_charge[gp_index]))/self.kBT*\
+        #             np.exp(-np.sign(self.free_charge[gp_index])*(self.phi[gp_index]-self.phi_old[gp_index])/self.kBT)
+        #         J[gp_index,gp_index-1] = flux_xm_J
+        #         J[gp_index,gp_index+1] = flux_xp_J
+        #         J[gp_index,gp_index-Nx] = flux_ym_J
+        #         J[gp_index,gp_index+Nx] = flux_yp_J
+        #         J[gp_index,gp_index-Nx*Ny] = flux_zm_J
+        #         J[gp_index,gp_index+Nx*Ny] = flux_zp_J
+
+
+        #         # add flux term to matrix B
+        #         B[gp_index] = (flux_xm_B+flux_xp_B+flux_ym_B+flux_yp_B+flux_zm_B+flux_zp_B)
+        #         B[gp_index] += elementary_charge*self.free_charge[gp_index]*np.exp(-np.sign(self.free_charge[gp_index])\
+        #             *(self.phi[gp_index]-self.phi_old[gp_index])/self.kBT)+elementary_charge*self.fixed_charge[gp_index]
+
+        #     else:# boundary points
+        #         J[gp_index,gp_index] = 1.0 # correct for both Dirichlet and Neumann boundary condition
+                
+        #         if self.boudnary_points[gp_index] == "xmin":   
+        #             J[gp_index,gp_index+1] = -1.0
+        #             B[gp_index] = (self.phi[gp_index]-self.phi[gp_index+1])
+        #         elif self.boudnary_points[gp_index] == "xmax":
+        #             J[gp_index,gp_index-1] = -1.0
+        #             B[gp_index] = (self.phi[gp_index]-self.phi[gp_index-1])
+        #         elif self.boudnary_points[gp_index] == "ymin":
+        #             J[gp_index,gp_index+Nx] = -1.0
+        #             B[gp_index] = (self.phi[gp_index]-self.phi[gp_index+Nx])
+        #         elif self.boudnary_points[gp_index] == "ymax":
+        #             J[gp_index,gp_index-Nx] = -1.0
+        #             B[gp_index] = (self.phi[gp_index]-self.phi[gp_index-Nx])
+        #         elif self.boudnary_points[gp_index] == "zmin":
+        #             J[gp_index,gp_index+Nx*Ny] = -1.0
+        #             B[gp_index] = (self.phi[gp_index]-self.phi[gp_index+Nx*Ny])
+        #         elif self.boudnary_points[gp_index] == "zmax":
+        #             J[gp_index,gp_index-Nx*Ny] = -1.0
+        #             B[gp_index] = (self.phi[gp_index]-self.phi[gp_index-Nx*Ny])
+        #         elif self.boudnary_points[gp_index] == "Dirichlet":
+        #             B[gp_index] = (self.phi[gp_index]-self.lead_gate_potential[gp_index])
+
+        #     if B[gp_index]!=0: # for convenience change the sign of B in later NR iteration
+        #         B[gp_index] = -B[gp_index]
+    
\ No newline at end of file

From 53c5f6e5a657e71c97fc6d85b4214235e8a8baf8 Mon Sep 17 00:00:00 2001
From: AsymmetryChou <181240085@smail.nju.edu.cn>
Date: Mon, 22 Dec 2025 22:05:50 +0800
Subject: [PATCH 08/14] Fix: remove Assertion bacause not every BC would appear
 in each case

---
 dpnegf/negf/newton_raphson_speed_up.py | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/dpnegf/negf/newton_raphson_speed_up.py b/dpnegf/negf/newton_raphson_speed_up.py
index 888fc19..712eb78 100644
--- a/dpnegf/negf/newton_raphson_speed_up.py
+++ b/dpnegf/negf/newton_raphson_speed_up.py
@@ -1,10 +1,13 @@
 import time
+import logging
 import unittest
 from typing import Optional, Tuple, List, Callable
 
 import numpy as np
 from scipy.sparse import lil_matrix
 
+log = logging.getLogger(__name__)
+
 def _impose_j_bound(inout, nx, ny, nz, typ, val, mask):
     '''impose the special mask for boundary points'''
     my_boundary_types_ = ['xmin', 'xmax', 'ymin', 'ymax', 'zmin', 'zmax', 'Dirichlet']
@@ -12,7 +15,8 @@ def _impose_j_bound(inout, nx, ny, nz, typ, val, mask):
     for t, d in zip(my_boundary_types_, my_displ_):
         # perf bottle-neck: search
         i = np.where(np.array(typ) == t)[0] # all indices of the boundary type
-        assert len(i) > 0, f'no grid point with type {t} found'
+        if len(i) == 0:
+            log.warning(f'no grid point with type {t} found')
         # scipy sparse matrix also supports the parallized assignment
         inout[i, i] = val
         inout[i, i + d] = mask
@@ -99,7 +103,8 @@ def _impose_b_bound(inout, nx, ny, nz, typ, phi, dirichlet_pot):
     for t, d in zip(my_boundary_types_, my_displ_):
         # perf bottle-neck: search
         i = np.where(typ == t)[0]
-        assert len(i) > 0, f'no grid point with type {t} found'
+        if len(i) == 0:
+            log.warning(f'no grid point with type {t} found')
         inout[i] = phi[i] - phi[i + d]
         
     # there is also a special type "Dirichlet" for the Dirichlet boundary condition

From 722540af19f02fd68eef15b40babf066b19e7837 Mon Sep 17 00:00:00 2001
From: AsymmetryChou <181240085@smail.nju.edu.cn>
Date: Mon, 22 Dec 2025 22:48:22 +0800
Subject: [PATCH 09/14] Feat: add unit tests for newton_raphson_speed_up.py

---
 dpnegf/tests/test_newton_raphson_speed_up.py | 492 +++++++++++++++++++
 1 file changed, 492 insertions(+)
 create mode 100644 dpnegf/tests/test_newton_raphson_speed_up.py

diff --git a/dpnegf/tests/test_newton_raphson_speed_up.py b/dpnegf/tests/test_newton_raphson_speed_up.py
new file mode 100644
index 0000000..d8702d6
--- /dev/null
+++ b/dpnegf/tests/test_newton_raphson_speed_up.py
@@ -0,0 +1,492 @@
+import time
+import unittest
+
+import numpy as np
+from scipy.sparse import lil_matrix
+
+from scipy.constants import epsilon_0, elementary_charge
+
+from dpnegf.negf.newton_raphson_speed_up import (
+    _impose_j_bound,
+    _impose_b_bound,
+    _bflux_impl,
+    _jflux_impl,
+    nr_construct,
+)
+
+# eps0 in units of F/Angstrom (same as poisson_init.py)
+eps0 = epsilon_0 * 1e-10
+
+
+class TestImposeJacobianBoundaryPerf(unittest.TestCase):
+    def setUp(self):
+        # assuming there are 1e6 grid points, each direction has 1e2
+        self.nx, self.ny, self.nz = 100, 100, 100
+        self.nr = self.nx * self.ny * self.nz
+        # say we have 6 boundary types, each has 300 points
+        typ = np.array(['****'] * self.nr, dtype=str).reshape((self.nx, self.ny, self.nz))
+        typ[typ == '****'] = 'in'
+        typ[   0,    :,  :] = 'xmin'
+        typ[  -1,    :,  :] = 'xmax'
+        typ[1:-1,    0,  :] = 'ymin'
+        typ[1:-1,   -1,  :] = 'ymax'
+        typ[1:-1, 1:-1,  0] = 'zmin'
+        typ[1:-1, 1:-1, -1] = 'zmax'
+        self.typ = typ.flatten()
+
+    def test_impose_bound(self):
+        # create a dummy inout matrix
+        ref = lil_matrix((self.nr, self.nr), dtype=float)
+        nx, ny, nz = self.nx, self.ny, self.nz
+        t = time.time()
+
+        for i in range(self.nr):
+            if self.typ[i] == "xmin":
+                ref[i, i] = -1.0
+                ref[i, i + 1] = 1.0
+            elif self.typ[i] == "xmax":
+                ref[i, i] = -1.0
+                ref[i, i - 1] = 1.0
+            elif self.typ[i] == "ymin":
+                ref[i, i] = -1.0
+                ref[i, i + nx] = 1.0
+            elif self.typ[i] == "ymax":
+                ref[i, i] = -1.0
+                ref[i, i - nx] = 1.0
+            elif self.typ[i] == "zmin":
+                ref[i, i] = -1.0
+                ref[i, i + nx*ny] = 1.0
+            elif self.typ[i] == "zmax":
+                ref[i, i] = -1.0
+                ref[i, i - nx*ny] = 1.0
+        print(f'old boundary impose method took {time.time() - t:.4f} seconds')
+
+        inout = lil_matrix((self.nr, self.nr), dtype=float)
+        t = time.time()
+        # call the function
+        _impose_j_bound(inout, nx, ny, nz, self.typ, -1.0, 1.0)
+        print(f'_impose_j_bound took {time.time() - t:.4f} seconds')
+
+        # check if the inout matrix is equal to the reference matrix
+        # check all diagonal elements
+        self.assertTrue(all(inout[i, i] == ref[i, i] for i in range(self.nr)))
+        # check xmin points
+        ind_xmin = np.where(self.typ == 'xmin')[0]
+        self.assertTrue(all(inout[i, i + 1] == ref[i, i + 1] for i in ind_xmin))
+        # check xmax points
+        ind_xmax = np.where(self.typ == 'xmax')[0]
+        self.assertTrue(all(inout[i, i - 1] == ref[i, i - 1] for i in ind_xmax))
+        # check ymin points
+        ind_ymin = np.where(self.typ == 'ymin')[0]
+        self.assertTrue(all(inout[i, i + nx] == ref[i, i + nx] for i in ind_ymin))
+        # check ymax points
+        ind_ymax = np.where(self.typ == 'ymax')[0]
+        self.assertTrue(all(inout[i, i - nx] == ref[i, i - nx] for i in ind_ymax))
+        # check zmin points
+        ind_zmin = np.where(self.typ == 'zmin')[0]
+        self.assertTrue(all(inout[i, i + nx*ny] == ref[i, i + nx*ny] for i in ind_zmin))
+        # check zmax points
+        ind_zmax = np.where(self.typ == 'zmax')[0]
+        self.assertTrue(all(inout[i, i - nx*ny] == ref[i, i - nx*ny] for i in ind_zmax))
+
+
+class TestImposeRHSVecBoundaryPerf(unittest.TestCase):
+    def setUp(self):
+        # assuming there are 1e6 grid points, each direction has 1e2
+        self.nx, self.ny, self.nz = 100, 100, 100
+        self.nr = self.nx * self.ny * self.nz
+        typ = np.array(['****'] * self.nr, dtype=str).reshape((self.nx, self.ny, self.nz))
+        typ[typ == '****'] = 'in'
+        typ[   0,    :,  :] = 'xmin'
+        typ[  -1,    :,  :] = 'xmax'
+        typ[1:-1,    0,  :] = 'ymin'
+        typ[1:-1,   -1,  :] = 'ymax'
+        typ[1:-1, 1:-1,  0] = 'zmin'
+        typ[1:-1, 1:-1, -1] = 'zmax'
+        self.typ = typ.flatten()
+        self.phi = np.random.rand(self.nr).astype(float)
+        self.fixed_pot = np.random.rand(self.nr).astype(float)
+
+    def test_impose_bound(self):
+        # create a dummy inout vector
+        ref = np.zeros(self.nr, dtype=float)
+        nx, ny, nz = self.nx, self.ny, self.nz
+        t = time.time()
+
+        for i in range(self.nr):
+            if self.typ[i] == "xmin":
+                ref[i] = self.phi[i] - self.phi[i + 1]
+            elif self.typ[i] == "xmax":
+                ref[i] = self.phi[i] - self.phi[i - 1]
+            elif self.typ[i] == "ymin":
+                ref[i] = self.phi[i] - self.phi[i + nx]
+            elif self.typ[i] == "ymax":
+                ref[i] = self.phi[i] - self.phi[i - nx]
+            elif self.typ[i] == "zmin":
+                ref[i] = self.phi[i] - self.phi[i + nx*ny]
+            elif self.typ[i] == "zmax":
+                ref[i] = self.phi[i] - self.phi[i - nx*ny]
+            elif self.typ[i] == "Dirichlet":
+                ref[i] = self.phi[i] - self.fixed_pot[i]
+        print(f'old boundary impose method took {time.time() - t:.4f} seconds')
+
+        inout = np.zeros(self.nr, dtype=float)
+        t = time.time()
+        # call the function
+        _impose_b_bound(inout, nx, ny, nz, self.typ, self.phi, self.fixed_pot)
+        print(f'_impose_b_bound took {time.time() - t:.4f} seconds')
+
+        # check if the inout vector is equal to the reference vector
+        self.assertTrue(np.all(inout == ref))
+
+
+class TestBFluxImplPerf(unittest.TestCase):
+    def setUp(self):
+        # assuming there are 1e6 grid points, each direction has 1e2
+        self.nx, self.ny, self.nz = 100, 100, 100
+        self.nr = self.nx * self.ny * self.nz
+        # set the boundary types
+        typ = np.array(['****'] * self.nr, dtype=str).reshape((self.nx, self.ny, self.nz))
+        typ[typ == '****'] = 'in'
+        typ[   0,    :,  :] = 'xmin'
+        typ[  -1,    :,  :] = 'xmax'
+        typ[1:-1,    0,  :] = 'ymin'
+        typ[1:-1,   -1,  :] = 'ymax'
+        typ[1:-1, 1:-1,  0] = 'zmin'
+        typ[1:-1, 1:-1, -1] = 'zmax'
+        self.typ = typ.flatten()
+        self.r = np.random.rand(self.nr, 3).astype(float)
+        self.sigma = np.random.rand(self.nr, 3).astype(float)
+        self.eps = np.random.rand(self.nr).astype(float)
+        self.eps0 = 8.854187817e-12
+        self.avgeps = lambda eps1, eps2: (eps1 + eps2) / 2.0  # arithmetic mean
+
+    def test_bflux_impl(self):
+        ind = np.where(self.typ == "in")[0]
+        jflux = np.random.rand(6, self.nr).astype(float)
+        phi = np.random.rand(self.nr).astype(float)
+        # calculate the reference result
+        ref = np.zeros((6, self.nr), dtype=float)
+        nx, ny, nz = self.nx, self.ny, self.nz
+        t = time.time()
+        for i in range(self.nr):
+            if self.typ[i] == 'in':
+                ref[0, i] = jflux[0, i] * (phi[i     - 1] - phi[i])
+                ref[1, i] = jflux[1, i] * (phi[i     + 1] - phi[i])
+                ref[2, i] = jflux[2, i] * (phi[i    - nx] - phi[i])
+                ref[3, i] = jflux[3, i] * (phi[i    + nx] - phi[i])
+                ref[4, i] = jflux[4, i] * (phi[i - nx*ny] - phi[i])
+                ref[5, i] = jflux[5, i] * (phi[i + nx*ny] - phi[i])
+        print(f'old bflux_impl method took {time.time() - t:.4f} seconds')
+        # call the function
+        t = time.time()
+        bflux = _bflux_impl(jflux, ind, self.nx, self.ny, self.nz, phi, full_size=True)
+        print(f'_bflux_impl took {time.time() - t:.4f} seconds')
+        # check if the result is equal to the reference result
+        self.assertTrue(bflux.shape == (6, self.nr))
+        self.assertTrue(np.allclose(bflux, ref, rtol=1e-5, atol=1e-8))
+
+
+class TestJFluxImplPerf(unittest.TestCase):
+    def setUp(self):
+        # assuming there are 1e6 grid points, each direction has 1e2
+        self.nx, self.ny, self.nz = 100, 100, 100
+        self.nr = self.nx * self.ny * self.nz
+        # set the boundary types
+        typ = np.array(['****'] * self.nr, dtype=str).reshape((self.nx, self.ny, self.nz))
+        typ[typ == '****'] = 'in'
+        typ[   0,    :,  :] = 'xmin'
+        typ[  -1,    :,  :] = 'xmax'
+        typ[1:-1,    0,  :] = 'ymin'
+        typ[1:-1,   -1,  :] = 'ymax'
+        typ[1:-1, 1:-1,  0] = 'zmin'
+        typ[1:-1, 1:-1, -1] = 'zmax'
+        self.typ = typ.flatten()
+        self.r = np.random.rand(self.nr, 3).astype(float)
+        self.sigma = np.random.rand(self.nr, 3).astype(float)
+        self.eps = np.random.rand(self.nr).astype(float)
+        self.eps0 = 8.854187817e-12  # vacuum permittivity
+        self.avgeps = lambda eps1, eps2: (eps1 + eps2) / 2.0  # arithmetic mean
+
+    def test_jflux_impl(self):
+        # calculate the reference result
+        ref = np.zeros((6, self.nr), dtype=float)
+        nx, ny, nz = self.nx, self.ny, self.nz
+        t = time.time()
+        for i in range(self.nr):
+            if self.typ[i] == "in":
+                ref[0, i] = self.sigma[i, 0] * self.eps0 * self.avgeps(self.eps[i - 1], self.eps[i]) \
+                            / np.abs(self.r[i - 1, 0] - self.r[i, 0])
+                ref[1, i] = self.sigma[i, 0] * self.eps0 * self.avgeps(self.eps[i + 1], self.eps[i]) \
+                            / np.abs(self.r[i + 1, 0] - self.r[i, 0])
+                ref[2, i] = self.sigma[i, 1] * self.eps0 * self.avgeps(self.eps[i - nx], self.eps[i]) \
+                            / np.abs(self.r[i - nx, 1] - self.r[i, 1])
+                ref[3, i] = self.sigma[i, 1] * self.eps0 * self.avgeps(self.eps[i + nx], self.eps[i]) \
+                            / np.abs(self.r[i + nx, 1] - self.r[i, 1])
+                ref[4, i] = self.sigma[i, 2] * self.eps0 * self.avgeps(self.eps[i - nx*ny], self.eps[i]) \
+                            / np.abs(self.r[i - nx*ny, 2] - self.r[i, 2])
+                ref[5, i] = self.sigma[i, 2] * self.eps0 * self.avgeps(self.eps[i + nx*ny], self.eps[i]) \
+                            / np.abs(self.r[i + nx*ny, 2] - self.r[i, 2])
+        print(f'old jflux_impl method took {time.time() - t:.4f} seconds')
+        # call the function
+        t = time.time()
+        jflux, i = _jflux_impl(self.nx, self.ny, self.nz, self.r, self.typ,
+                               self.sigma, self.eps, self.eps0, self.avgeps,
+                               with_index=True, full_size=True)
+        print(f'_jflux_impl took {time.time() - t:.4f} seconds')
+        # check if the result is equal to the reference result
+        self.assertTrue(jflux.shape == (6, self.nr))
+        self.assertTrue(i.shape == (len(self.typ[self.typ == "in"]),))
+        self.assertTrue(np.allclose(jflux, ref, rtol=1e-5, atol=1e-8))
+
+
+class TestNRConstructConsistency(unittest.TestCase):
+    """
+    Test consistency between new nr_construct implementation and the old
+    implementation from poisson_init.py (NR_construct_Jac_B).
+
+    The old implementation is reproduced here based on the commented-out code
+    in poisson_init.py for direct comparison.
+    """
+
+    def setUp(self):
+        # Use a smaller grid for faster testing
+        self.nx, self.ny, self.nz = 10, 10, 10
+        self.nr = self.nx * self.ny * self.nz
+
+        # Set up boundary types
+        # Use Fortran order (x varies fastest) to match grid coordinate indexing
+        typ = np.array(['****'] * self.nr, dtype=str).reshape((self.nx, self.ny, self.nz))
+        typ[typ == '****'] = 'in'
+        typ[   0,    :,  :] = 'xmin'
+        typ[  -1,    :,  :] = 'xmax'
+        typ[1:-1,    0,  :] = 'ymin'
+        typ[1:-1,   -1,  :] = 'ymax'
+        typ[1:-1, 1:-1,  0] = 'zmin'
+        typ[1:-1, 1:-1, -1] = 'zmax'
+        self.typ = typ.flatten(order='F')
+
+        # Set up grid coordinates (ensuring no zero distances)
+        # Grid indexing: i+1 is x neighbor, i+nx is y neighbor, i+nx*ny is z neighbor
+        # So we need x to vary fastest when flattening (Fortran order)
+        x = np.linspace(0, 1, self.nx)
+        y = np.linspace(0, 1, self.ny)
+        z = np.linspace(0, 1, self.nz)
+        xx, yy, zz = np.meshgrid(x, y, z, indexing='ij')
+        # Use Fortran order to make x vary fastest, then y, then z
+        self.grid_coord = np.stack([
+            xx.flatten(order='F'),
+            yy.flatten(order='F'),
+            zz.flatten(order='F')
+        ], axis=1)
+
+        # Set up surface areas (random positive values)
+        self.surface_grid = np.abs(np.random.rand(self.nr, 3)) + 0.1
+
+        # Set up permittivity (random positive values)
+        self.eps = np.abs(np.random.rand(self.nr)) + 1.0
+
+        # Set up potentials
+        self.phi = np.random.rand(self.nr)
+        self.phi_old = np.random.rand(self.nr)
+
+        # Set up charges (can be positive or negative)
+        self.free_charge = (np.random.rand(self.nr) - 0.5) * 1e-10
+        self.fixed_charge = (np.random.rand(self.nr) - 0.5) * 1e-10
+
+        # Set up Dirichlet potential for boundary
+        self.lead_gate_potential = np.random.rand(self.nr)
+
+        # Temperature
+        self.kBT = 0.0259  # ~300K in eV
+
+        # Average mode
+        self.average_mode = 'arithmetic'
+
+    def _old_implementation(self, J, B):
+        """
+        Reproduce the old implementation from poisson_init.py (commented out code).
+        This is a direct port of the loop-based implementation.
+        """
+        def average_eps(eps1, eps2, mode='arithmetic'):
+            if mode == 'arithmetic':
+                return 0.5 * (eps1 + eps2)
+            elif mode == 'harmonic':
+                return 2.0 * eps1 * eps2 / (eps1 + eps2)
+            elif mode == 'geometric':
+                return np.sqrt(eps1 * eps2)
+
+        average_mode = self.average_mode
+        Nx = self.nx
+        Ny = self.ny
+
+        for gp_index in range(self.nr):
+            if self.typ[gp_index] == "in":
+                # x-direction fluxes
+                flux_xm_J = self.surface_grid[gp_index, 0] * eps0 * \
+                    average_eps(self.eps[gp_index - 1], self.eps[gp_index], mode=average_mode) / \
+                    abs(self.grid_coord[gp_index, 0] - self.grid_coord[gp_index - 1, 0])
+                flux_xm_B = flux_xm_J * (self.phi[gp_index - 1] - self.phi[gp_index])
+
+                flux_xp_J = self.surface_grid[gp_index, 0] * eps0 * \
+                    average_eps(self.eps[gp_index + 1], self.eps[gp_index], mode=average_mode) / \
+                    abs(self.grid_coord[gp_index + 1, 0] - self.grid_coord[gp_index, 0])
+                flux_xp_B = flux_xp_J * (self.phi[gp_index + 1] - self.phi[gp_index])
+
+                # y-direction fluxes
+                flux_ym_J = self.surface_grid[gp_index, 1] * eps0 * \
+                    average_eps(self.eps[gp_index - Nx], self.eps[gp_index], mode=average_mode) / \
+                    abs(self.grid_coord[gp_index - Nx, 1] - self.grid_coord[gp_index, 1])
+                flux_ym_B = flux_ym_J * (self.phi[gp_index - Nx] - self.phi[gp_index])
+
+                flux_yp_J = self.surface_grid[gp_index, 1] * eps0 * \
+                    average_eps(self.eps[gp_index + Nx], self.eps[gp_index], mode=average_mode) / \
+                    abs(self.grid_coord[gp_index + Nx, 1] - self.grid_coord[gp_index, 1])
+                flux_yp_B = flux_yp_J * (self.phi[gp_index + Nx] - self.phi[gp_index])
+
+                # z-direction fluxes
+                flux_zm_J = self.surface_grid[gp_index, 2] * eps0 * \
+                    average_eps(self.eps[gp_index - Nx * Ny], self.eps[gp_index], mode=average_mode) / \
+                    abs(self.grid_coord[gp_index - Nx * Ny, 2] - self.grid_coord[gp_index, 2])
+                flux_zm_B = flux_zm_J * (self.phi[gp_index - Nx * Ny] - self.phi[gp_index])
+
+                flux_zp_J = self.surface_grid[gp_index, 2] * eps0 * \
+                    average_eps(self.eps[gp_index + Nx * Ny], self.eps[gp_index], mode=average_mode) / \
+                    abs(self.grid_coord[gp_index + Nx * Ny, 2] - self.grid_coord[gp_index, 2])
+                flux_zp_B = flux_zp_J * (self.phi[gp_index + Nx * Ny] - self.phi[gp_index])
+
+                # Jacobian diagonal element
+                J[gp_index, gp_index] = -(flux_xm_J + flux_xp_J + flux_ym_J + flux_yp_J + flux_zm_J + flux_zp_J) \
+                    + elementary_charge * self.free_charge[gp_index] * (-np.sign(self.free_charge[gp_index])) / self.kBT * \
+                    np.exp(-np.sign(self.free_charge[gp_index]) * (self.phi[gp_index] - self.phi_old[gp_index]) / self.kBT)
+
+                # Jacobian off-diagonal elements
+                J[gp_index, gp_index - 1] = flux_xm_J
+                J[gp_index, gp_index + 1] = flux_xp_J
+                J[gp_index, gp_index - Nx] = flux_ym_J
+                J[gp_index, gp_index + Nx] = flux_yp_J
+                J[gp_index, gp_index - Nx * Ny] = flux_zm_J
+                J[gp_index, gp_index + Nx * Ny] = flux_zp_J
+
+                # B vector
+                B[gp_index] = (flux_xm_B + flux_xp_B + flux_ym_B + flux_yp_B + flux_zm_B + flux_zp_B)
+                B[gp_index] += elementary_charge * self.free_charge[gp_index] * \
+                    np.exp(-np.sign(self.free_charge[gp_index]) * (self.phi[gp_index] - self.phi_old[gp_index]) / self.kBT) \
+                    + elementary_charge * self.fixed_charge[gp_index]
+
+            else:  # boundary points
+                J[gp_index, gp_index] = 1.0
+
+                if self.typ[gp_index] == "xmin":
+                    J[gp_index, gp_index + 1] = -1.0
+                    B[gp_index] = (self.phi[gp_index] - self.phi[gp_index + 1])
+                elif self.typ[gp_index] == "xmax":
+                    J[gp_index, gp_index - 1] = -1.0
+                    B[gp_index] = (self.phi[gp_index] - self.phi[gp_index - 1])
+                elif self.typ[gp_index] == "ymin":
+                    J[gp_index, gp_index + Nx] = -1.0
+                    B[gp_index] = (self.phi[gp_index] - self.phi[gp_index + Nx])
+                elif self.typ[gp_index] == "ymax":
+                    J[gp_index, gp_index - Nx] = -1.0
+                    B[gp_index] = (self.phi[gp_index] - self.phi[gp_index - Nx])
+                elif self.typ[gp_index] == "zmin":
+                    J[gp_index, gp_index + Nx * Ny] = -1.0
+                    B[gp_index] = (self.phi[gp_index] - self.phi[gp_index + Nx * Ny])
+                elif self.typ[gp_index] == "zmax":
+                    J[gp_index, gp_index - Nx * Ny] = -1.0
+                    B[gp_index] = (self.phi[gp_index] - self.phi[gp_index - Nx * Ny])
+                elif self.typ[gp_index] == "Dirichlet":
+                    B[gp_index] = (self.phi[gp_index] - self.lead_gate_potential[gp_index])
+
+            # Sign flip for NR iteration
+            if B[gp_index] != 0:
+                B[gp_index] = -B[gp_index]
+
+    def test_nr_construct_vs_old_implementation(self):
+        """Test that nr_construct produces the same results as the old implementation."""
+        # Create matrices for old implementation
+        J_old = lil_matrix((self.nr, self.nr), dtype=float)
+        B_old = np.zeros(self.nr, dtype=float)
+
+        # Run old implementation
+        t = time.time()
+        self._old_implementation(J_old, B_old)
+        print(f'Old implementation took {time.time() - t:.4f} seconds')
+
+        # Create matrices for new implementation
+        J_new = lil_matrix((self.nr, self.nr), dtype=float)
+        B_new = np.zeros(self.nr, dtype=float)
+
+        # Select average function based on mode
+        feps = {'harmonic': lambda eps1, eps2: 2.0 * eps1 * eps2 / (eps1 + eps2),
+                'arithmetic': lambda eps1, eps2: 0.5 * (eps1 + eps2),
+                'geometric': lambda eps1, eps2: np.sqrt(eps1 * eps2)}[self.average_mode]
+
+        # Run new implementation
+        t = time.time()
+        nr_construct(
+            jinout=J_new,
+            binout=B_new,
+            grid_dim=(self.nx, self.ny, self.nz),
+            gridpoint_coords=self.grid_coord,
+            gridpoint_typ=self.typ,
+            gridpoint_surfarea=self.surface_grid,
+            eps=self.eps,
+            phi=self.phi,
+            phi_=self.phi_old,
+            free_chr=self.free_charge,
+            fixed_chr=self.fixed_charge,
+            dirichlet_pot=self.lead_gate_potential,
+            eps0=eps0,
+            beta=1.0 / self.kBT,
+            feps=feps
+        )
+        print(f'New implementation took {time.time() - t:.4f} seconds')
+
+        # Compare Jacobian matrices
+        J_old_dense = J_old.toarray()
+        J_new_dense = J_new.toarray()
+
+        # Check diagonal elements
+        diag_diff = np.abs(np.diag(J_old_dense) - np.diag(J_new_dense))
+        max_diag_diff = np.max(diag_diff)
+        print(f'Max diagonal difference: {max_diag_diff}')
+
+        # Check off-diagonal elements
+        off_diag_diff = np.abs(J_old_dense - J_new_dense)
+        np.fill_diagonal(off_diag_diff, 0)
+        max_off_diag_diff = np.max(off_diag_diff)
+        print(f'Max off-diagonal difference: {max_off_diag_diff}')
+
+        # Check B vectors
+        B_diff = np.abs(B_old - B_new)
+        max_B_diff = np.max(B_diff)
+        print(f'Max B vector difference: {max_B_diff}')
+
+        # Assert with reasonable tolerance
+        rtol = 1e-10
+        atol = 1e-15
+
+        self.assertTrue(
+            np.allclose(J_old_dense, J_new_dense, rtol=rtol, atol=atol),
+            f'Jacobian matrices differ: max diagonal diff={max_diag_diff}, '
+            f'max off-diagonal diff={max_off_diag_diff}'
+        )
+        self.assertTrue(
+            np.allclose(B_old, B_new, rtol=rtol, atol=atol),
+            f'B vectors differ: max diff={max_B_diff}'
+        )
+
+    def test_nr_construct_harmonic_average(self):
+        """Test with harmonic average mode."""
+        self.average_mode = 'harmonic'
+        self.test_nr_construct_vs_old_implementation()
+
+    def test_nr_construct_geometric_average(self):
+        """Test with geometric average mode."""
+        self.average_mode = 'geometric'
+        self.test_nr_construct_vs_old_implementation()
+
+
+if __name__ == '__main__':
+    unittest.main()

From c33fd56a015eb2e1a37eacf4d8e95b7db76ec186 Mon Sep 17 00:00:00 2001
From: AsymmetryChou <181240085@smail.nju.edu.cn>
Date: Mon, 22 Dec 2025 22:48:49 +0800
Subject: [PATCH 10/14] Feat: replace construct with nr_construct for Poisson
 solver

---
 dpnegf/negf/poisson_init.py | 36 +++++++++++++++++++-----------------
 1 file changed, 19 insertions(+), 17 deletions(-)

diff --git a/dpnegf/negf/poisson_init.py b/dpnegf/negf/poisson_init.py
index c1e5df9..73ed829 100644
--- a/dpnegf/negf/poisson_init.py
+++ b/dpnegf/negf/poisson_init.py
@@ -6,6 +6,7 @@
 # from pyamg.gallery import poisson
 from dpnegf.utils.constants import elementary_charge
 from dpnegf.utils.constants import Boltzmann, eV2J
+from dpnegf.negf.newton_raphson_speed_up import nr_construct
 from scipy.constants import epsilon_0 as eps0  #TODO:later add to untils.constants.py
 from scipy.sparse import csr_matrix
 from scipy.sparse.linalg import spsolve
@@ -540,7 +541,8 @@ def solve_poisson_NRcycle(self,
             raise ValueError(f"Unknown data type: {dtype}. Use 'float64' or 'float32'.")
 
         # initialize the solver
-        assert method in ['scipy', 'pyamg']
+        if method not in ['scipy', 'pyamg']:
+            raise ValueError(f"Unknown Poisson solver method: {method}. Use 'scipy' or 'pyamg'.")
         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
@@ -662,7 +664,7 @@ def NR_construct_Jac_B(self,J,B):
         - Uses constants such as eps0 and elementary_charge, which must be defined in the scope.
         - The method modifies J and B in place.
         """
-        from dpnegf.negf.newton_raphson_speed_up import calculate as construct
+        
         nx, ny, nz = self.grid.shape[:3]
         # https://manual.gromacs.org/documentation/current/reference-manual/topologies/parameter-files.html#nbpar
         feps = {'harmonic'  : lambda eps1, eps2: 2.0 * eps1 * eps2 / (eps1 + eps2),
@@ -670,21 +672,21 @@ def NR_construct_Jac_B(self,J,B):
                 'geometric' : lambda eps1, eps2: np.sqrt(eps1 * eps2), # gromacs comb-rule 2
                 }[self.average_mode]
 
-        construct(jinout=J, 
-                  binout=B, 
-                  grid_dim=(nx, ny, nz),
-                  gridpoint_coords=self.grid.grid_coord,
-                  gridpoint_typ=self.boudnary_points,
-                  gridpoint_surfarea=self.grid.surface_grid,
-                  eps=self.eps,
-                  phi=self.phi,
-                  phi_=self.phi_old,
-                  free_chr=self.free_charge,
-                  fixed_chr=self.fixed_charge,
-                  dirichlet_pot=self.lead_gate_potential,
-                  eps0=eps0,
-                  beta=1.0/self.kBT,
-                  feps=feps)
+        nr_construct(jinout=J, 
+                    binout=B, 
+                    grid_dim=(nx, ny, nz),
+                    gridpoint_coords=self.grid.grid_coord,
+                    gridpoint_typ=self.boudnary_points,
+                    gridpoint_surfarea=self.grid.surface_grid,
+                    eps=self.eps,
+                    phi=self.phi,
+                    phi_=self.phi_old,
+                    free_chr=self.free_charge,
+                    fixed_chr=self.fixed_charge,
+                    dirichlet_pot=self.lead_gate_potential,
+                    eps0=eps0,
+                    beta=1.0/self.kBT,
+                    feps=feps)
 
         # ===================================== OLD IMPLEMENTATION =========================================
         # construct the Jacobian and B for the Poisson equation

From e1b01954e128dad40164ca7e02a1f7f28856da2b Mon Sep 17 00:00:00 2001
From: AsymmetryChou <181240085@smail.nju.edu.cn>
Date: Mon, 22 Dec 2025 22:51:11 +0800
Subject: [PATCH 11/14] Refactor: move unit tests to isolated file and
 formating variables

---
 dpnegf/negf/newton_raphson_speed_up.py | 253 ++-----------------------
 1 file changed, 15 insertions(+), 238 deletions(-)

diff --git a/dpnegf/negf/newton_raphson_speed_up.py b/dpnegf/negf/newton_raphson_speed_up.py
index 712eb78..3e96b91 100644
--- a/dpnegf/negf/newton_raphson_speed_up.py
+++ b/dpnegf/negf/newton_raphson_speed_up.py
@@ -1,6 +1,4 @@
-import time
 import logging
-import unittest
 from typing import Optional, Tuple, List, Callable
 
 import numpy as np
@@ -10,9 +8,9 @@
 
 def _impose_j_bound(inout, nx, ny, nz, typ, val, mask):
     '''impose the special mask for boundary points'''
-    my_boundary_types_ = ['xmin', 'xmax', 'ymin', 'ymax', 'zmin', 'zmax', 'Dirichlet']
-    my_displ_ = [+1, -1, +nx, -nx, +nx*ny, -nx*ny]
-    for t, d in zip(my_boundary_types_, my_displ_):
+    boundary_types_ = ['xmin', 'xmax', 'ymin', 'ymax', 'zmin', 'zmax', 'Dirichlet']
+    displ_ = [+1, -1, +nx, -nx, +nx*ny, -nx*ny]
+    for t, d in zip(boundary_types_, displ_):
         # perf bottle-neck: search
         i = np.where(np.array(typ) == t)[0] # all indices of the boundary type
         if len(i) == 0:
@@ -25,82 +23,11 @@ def _impose_j_bound(inout, nx, ny, nz, typ, val, mask):
     i = np.where(typ == 'Dirichlet')[0]
     inout[i, i] = val
 
-class TestImposeJacobianBoundaryPerf(unittest.TestCase):
-    def setUp(self):
-        # assuming there are 1e6 grid points, each direction has 1e3
-        self.nx, self.ny, self.nz = 100, 100, 100
-        self.nr = self.nx * self.ny * self.nz
-        # say we have 6 boundary types, each has 300 points
-        typ = np.array(['****'] * self.nr, dtype=str).reshape((self.nx, self.ny, self.nz))
-        typ[typ == '****'] = 'in'
-        typ[   0,    :,  :] = 'xmin'
-        typ[  -1,    :,  :] = 'xmax'
-        typ[1:-1,    0,  :] = 'ymin'
-        typ[1:-1,   -1,  :] = 'ymax'
-        typ[1:-1, 1:-1,  0] = 'zmin'
-        typ[1:-1, 1:-1, -1] = 'zmax'
-        self.typ = typ.flatten()
-        
-    def test_impose_bound(self):
-        # create a dummy inout matrix
-        ref = lil_matrix((self.nr, self.nr), dtype=float)
-        nx, ny, nz = self.nx, self.ny, self.nz
-        t = time.time()
-        
-        for i in range(self.nr):
-            if self.typ[i] == "xmin":
-                ref[i, i] = -1.0
-                ref[i, i + 1] = 1.0
-            elif self.typ[i] == "xmax":
-                ref[i, i] = -1.0
-                ref[i, i - 1] = 1.0
-            elif self.typ[i] == "ymin":
-                ref[i, i] = -1.0
-                ref[i, i + nx] = 1.0
-            elif self.typ[i] == "ymax":
-                ref[i, i] = -1.0
-                ref[i, i - nx] = 1.0
-            elif self.typ[i] == "zmin":
-                ref[i, i] = -1.0
-                ref[i, i + nx*ny] = 1.0
-            elif self.typ[i] == "zmax":
-                ref[i, i] = -1.0
-                ref[i, i - nx*ny] = 1.0
-        print(f'old boundary impose method took {time.time() - t:.4f} seconds')
-        
-        inout = lil_matrix((self.nr, self.nr), dtype=float)
-        t = time.time()
-        # call the function
-        _impose_j_bound(inout, nx, ny, nz, self.typ, -1.0, 1.0)
-        print(f'_impose_j_bound took {time.time() - t:.4f} seconds')
-    
-        # check if the inout matrix is equal to the reference matrix
-        # check all diagonal elements
-        self.assertTrue(all(inout[i, i] == ref[i, i] for i in range(self.nr)))
-        # check xmin points
-        ind_xmin = np.where(self.typ == 'xmin')[0]
-        self.assertTrue(all(inout[i, i + 1] == ref[i, i + 1] for i in ind_xmin))
-        # check xmax points
-        ind_xmax = np.where(self.typ == 'xmax')[0]
-        self.assertTrue(all(inout[i, i - 1] == ref[i, i - 1] for i in ind_xmax))
-        # check ymin points
-        ind_ymin = np.where(self.typ == 'ymin')[0]
-        self.assertTrue(all(inout[i, i + nx] == ref[i, i + nx] for i in ind_ymin))
-        # check ymax points
-        ind_ymax = np.where(self.typ == 'ymax')[0]
-        self.assertTrue(all(inout[i, i - nx] == ref[i, i - nx] for i in ind_ymax))
-        # check zmin points
-        ind_zmin = np.where(self.typ == 'zmin')[0]
-        self.assertTrue(all(inout[i, i + nx*ny] == ref[i, i + nx*ny] for i in ind_zmin))
-        # check zmax points
-        ind_zmax = np.where(self.typ == 'zmax')[0]
-        self.assertTrue(all(inout[i, i - nx*ny] == ref[i, i - nx*ny] for i in ind_zmax))
-
 def _impose_b_bound(inout, nx, ny, nz, typ, phi, dirichlet_pot):
     '''impose the special mask for boundary points in b vector'''
-    my_boundary_types_ = ['xmin', 'xmax', 'ymin', 'ymax', 'zmin', 'zmax', 'Dirichlet']
-    my_displ_ = [+1, -1, +nx, -nx, +nx*ny, -nx*ny]
-    for t, d in zip(my_boundary_types_, my_displ_):
+    boundary_types_ = ['xmin', 'xmax', 'ymin', 'ymax', 'zmin', 'zmax', 'Dirichlet']
+    displ_ = [+1, -1, +nx, -nx, +nx*ny, -nx*ny]
+    for t, d in zip(boundary_types_, displ_):
         # perf bottle-neck: search
         i = np.where(typ == t)[0]
         if len(i) == 0:
@@ -111,55 +38,6 @@ def _impose_b_bound(inout, nx, ny, nz, typ, phi, dirichlet_pot):
     i = np.where(typ == 'Dirichlet')[0]
     inout[i] = phi[i] - dirichlet_pot[i]
 
-class TestImposeRHSVecBoundaryPerf(unittest.TestCase):
-    def setUp(self):
-        # assuming there are 1e6 grid points, each direction has 1e3
-        self.nx, self.ny, self.nz = 100, 100, 100
-        self.nr = self.nx * self.ny * self.nz
-        typ = np.array(['****'] * self.nr, dtype=str).reshape((self.nx, self.ny, self.nz))
-        typ[typ == '****'] = 'in'
-        typ[   0,    :,  :] = 'xmin'
-        typ[  -1,    :,  :] = 'xmax'
-        typ[1:-1,    0,  :] = 'ymin'
-        typ[1:-1,   -1,  :] = 'ymax'
-        typ[1:-1, 1:-1,  0] = 'zmin'
-        typ[1:-1, 1:-1, -1] = 'zmax'
-        self.typ = typ.flatten()
-        self.phi = np.random.rand(self.nr).astype(float)
-        self.fixed_pot = np.random.rand(self.nr).astype(float)
-        
-    def test_impose_bound(self):
-        # create a dummy inout vector
-        ref = np.zeros(self.nr, dtype=float)
-        nx, ny, nz = self.nx, self.ny, self.nz
-        t = time.time()
-        
-        for i in range(self.nr):
-            if self.typ[i] == "xmin":
-                ref[i] = self.phi[i] - self.phi[i + 1]
-            elif self.typ[i] == "xmax":
-                ref[i] = self.phi[i] - self.phi[i - 1]
-            elif self.typ[i] == "ymin":
-                ref[i] = self.phi[i] - self.phi[i + nx]
-            elif self.typ[i] == "ymax":
-                ref[i] = self.phi[i] - self.phi[i - nx]
-            elif self.typ[i] == "zmin":
-                ref[i] = self.phi[i] - self.phi[i + nx*ny]
-            elif self.typ[i] == "zmax":
-                ref[i] = self.phi[i] - self.phi[i - nx*ny]
-            elif self.typ[i] == "Dirichlet":
-                ref[i] = self.phi[i] - self.fixed_pot[i]
-        print(f'old boundary impose method took {time.time() - t:.4f} seconds')
-        
-        inout = np.zeros(self.nr, dtype=float)
-        t = time.time()
-        # call the function
-        _impose_b_bound(inout, nx, ny, nz, self.typ, self.phi, self.fixed_pot)
-        print(f'_impose_b_bound took {time.time() - t:.4f} seconds')
-    
-        # check if the inout vector is equal to the reference vector
-        self.assertTrue(np.all(inout == ref))
-
 def coloumb(chr: float|np.ndarray) -> float|np.ndarray:
     '''convert the charge to unit of Coulomb'''
     assert isinstance(chr, (float, np.ndarray)), "chr must be a float or numpy array"
@@ -168,64 +46,18 @@ def coloumb(chr: float|np.ndarray) -> float|np.ndarray:
 
 def _bflux_impl(jflux, i, nx, ny, nz, phi, full_size=True) -> np.ndarray:
     '''calculate the b vector from the jflux and phi'''
-    mydisp_ = [-1, +1, -nx, +nx, -nx*ny, +nx*ny]
+    disp_ = [-1, +1, -nx, +nx, -nx*ny, +nx*ny]
     if full_size:
         assert jflux.shape == (6, len(phi))
         bflux = np.zeros((6, len(phi)), dtype=float)
-        for j, d in enumerate(mydisp_):
+        for j, d in enumerate(disp_):
             bflux[j, i] = jflux[j, i] * (phi[i + d] - phi[i])
     else:
         assert jflux.shape == (6, len(i))
-        bflux = np.array([jf * (phi[i + d] - phi[i]) for jf, d in zip(jflux, mydisp_)])
+        bflux = np.array([jf * (phi[i + d] - phi[i]) for jf, d in zip(jflux, disp_)])
         assert bflux.shape == (6, len(i))
     return bflux
 
-class TestBFluxImplPerf(unittest.TestCase):
-    def setUp(self):
-        # assuming there are 1e6 grid points, each direction has 1e3
-        self.nx, self.ny, self.nz = 100, 100, 100
-        self.nr = self.nx * self.ny * self.nz
-        # set the boundary types
-        typ = np.array(['****'] * self.nr, dtype=str).reshape((self.nx, self.ny, self.nz))
-        typ[typ == '****'] = 'in'
-        typ[   0,    :,  :] = 'xmin'
-        typ[  -1,    :,  :] = 'xmax'
-        typ[1:-1,    0,  :] = 'ymin'
-        typ[1:-1,   -1,  :] = 'ymax'
-        typ[1:-1, 1:-1,  0] = 'zmin'
-        typ[1:-1, 1:-1, -1] = 'zmax'
-        self.typ = typ.flatten()
-        self.r = np.random.rand(self.nr, 3).astype(float)
-        self.sigma = np.random.rand(self.nr, 3).astype(float)
-        self.eps = np.random.rand(self.nr).astype(float)
-        self.eps0 = 8.854187817e-12
-        self.avgeps = lambda eps1, eps2: (eps1 + eps2) / 2.0  # arithmetic mean
-
-    def test_bflux_impl(self):
-        ind = np.where(self.typ == "in")[0]
-        jflux = np.random.rand(6, self.nr).astype(float)
-        phi = np.random.rand(self.nr).astype(float)
-        # calculate the reference result
-        ref = np.zeros((6, self.nr), dtype=float)
-        nx, ny, nz = self.nx, self.ny, self.nz
-        t = time.time()
-        for i in range(self.nr):
-            if self.typ[i] == 'in':
-                ref[0, i] = jflux[0, i] * (phi[i     - 1] - phi[i])
-                ref[1, i] = jflux[1, i] * (phi[i     + 1] - phi[i])
-                ref[2, i] = jflux[2, i] * (phi[i    - nx] - phi[i])
-                ref[3, i] = jflux[3, i] * (phi[i    + nx] - phi[i])
-                ref[4, i] = jflux[4, i] * (phi[i - nx*ny] - phi[i])
-                ref[5, i] = jflux[5, i] * (phi[i + nx*ny] - phi[i])
-        print(f'old bflux_impl method took {time.time() - t:.4f} seconds')
-        # call the function
-        t = time.time()
-        bflux = _bflux_impl(jflux, ind, self.nx, self.ny, self.nz, phi, full_size=True)
-        print(f'_bflux_impl took {time.time() - t:.4f} seconds')
-        # check if the result is equal to the reference result
-        self.assertTrue(bflux.shape == (6, self.nr))
-        self.assertTrue(np.allclose(bflux, ref, rtol=1e-5, atol=1e-8))
-
 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.'''
@@ -233,82 +65,30 @@ def _jflux_impl(nx, ny, nz, r, typ, sigma, eps, eps0, avgeps, with_index=False,
     # get all indices of the grid points of type "in"
     i = np.where(typ == "in")[0]
     # j stands for flux, <listcomp> faster than for loop
-    mydisp_ = [-1, +1, -nx, +nx, -nx*ny, +nx*ny]
+    disp_ = [-1, +1, -nx, +nx, -nx*ny, +nx*ny]
     if not full_size:
         jflux = np.array([avgeps(eps[i + d], eps[i]) * eps0 * sigma[i, j // 2] \
                           / np.abs(r[i + d, j // 2] - r[i, j // 2])
-                          for j, d in enumerate(mydisp_)]).reshape((6, -1))
+                          for j, d in enumerate(disp_)]).reshape((6, -1))
         assert jflux.shape == (6, len(i))
     else:
         jflux = np.zeros((6, len(typ)), dtype=float)
-        for j, d in enumerate(mydisp_):
+        for j, d in enumerate(disp_):
             jflux[j, i] =   avgeps(eps[i + d], eps[i]) * eps0 * sigma[i, j // 2] \
                           / np.abs(r[i + d, j // 2] - r[i, j // 2])
 
     return (jflux, i if with_index else None)
 
-class TestJFluxImplPerf(unittest.TestCase):
-    def setUp(self):
-        # assuming there are 1e6 grid points, each direction has 1e3
-        self.nx, self.ny, self.nz = 100, 100, 100
-        self.nr = self.nx * self.ny * self.nz
-        # set the boundary types
-        typ = np.array(['****'] * self.nr, dtype=str).reshape((self.nx, self.ny, self.nz))
-        typ[typ == '****'] = 'in'
-        typ[   0,    :,  :] = 'xmin'
-        typ[  -1,    :,  :] = 'xmax'
-        typ[1:-1,    0,  :] = 'ymin'
-        typ[1:-1,   -1,  :] = 'ymax'
-        typ[1:-1, 1:-1,  0] = 'zmin'
-        typ[1:-1, 1:-1, -1] = 'zmax'
-        self.typ = typ.flatten()
-        self.r = np.random.rand(self.nr, 3).astype(float)
-        self.sigma = np.random.rand(self.nr, 3).astype(float)
-        self.eps = np.random.rand(self.nr).astype(float)
-        self.eps0 = 8.854187817e-12  # vacuum permittivity
-        self.avgeps = lambda eps1, eps2: (eps1 + eps2) / 2.0  # arithmetic mean
-        
-    def test_jflux_impl(self):
-        # calculate the reference result
-        ref = np.zeros((6, self.nr), dtype=float)
-        nx, ny, nz = self.nx, self.ny, self.nz
-        t = time.time()
-        for i in range(self.nr):
-            if self.typ[i] == "in":
-                ref[0, i] = self.sigma[i, 0] * self.eps0 * self.avgeps(self.eps[i - 1], self.eps[i]) \
-                            / np.abs(self.r[i - 1, 0] - self.r[i, 0])
-                ref[1, i] = self.sigma[i, 0] * self.eps0 * self.avgeps(self.eps[i + 1], self.eps[i]) \
-                            / np.abs(self.r[i + 1, 0] - self.r[i, 0])
-                ref[2, i] = self.sigma[i, 1] * self.eps0 * self.avgeps(self.eps[i - nx], self.eps[i]) \
-                            / np.abs(self.r[i - nx, 1] - self.r[i, 1])
-                ref[3, i] = self.sigma[i, 1] * self.eps0 * self.avgeps(self.eps[i + nx], self.eps[i]) \
-                            / np.abs(self.r[i + nx, 1] - self.r[i, 1])
-                ref[4, i] = self.sigma[i, 2] * self.eps0 * self.avgeps(self.eps[i - nx*ny], self.eps[i]) \
-                            / np.abs(self.r[i - nx*ny, 2] - self.r[i, 2])
-                ref[5, i] = self.sigma[i, 2] * self.eps0 * self.avgeps(self.eps[i + nx*ny], self.eps[i]) \
-                            / np.abs(self.r[i + nx*ny, 2] - self.r[i, 2])
-        print(f'old jflux_impl method took {time.time() - t:.4f} seconds')
-        # call the function
-        t = time.time()
-        jflux, i = _jflux_impl(self.nx, self.ny, self.nz, self.r, self.typ, 
-                               self.sigma, self.eps, self.eps0, self.avgeps, 
-                               with_index=True, full_size=True)
-        print(f'_jflux_impl took {time.time() - t:.4f} seconds')
-        # check if the result is equal to the reference result
-        self.assertTrue(jflux.shape == (6, self.nr))
-        self.assertTrue(i.shape == (len(self.typ[self.typ == "in"]),))
-        self.assertTrue(np.allclose(jflux, ref, rtol=1e-5, atol=1e-8))
-
 def _jacobian(inout, flux, i, nx, ny, nz, typ, zfree, phi, phi_, beta):
     '''calculate the Jacobian matrix for the Poisson equation'''
-    mydisp_ = [-1, +1, -nx, +nx, -nx*ny, +nx*ny]
+    disp_ = [-1, +1, -nx, +nx, -nx*ny, +nx*ny]
     # diagonal elements
     assert flux.shape == (6, len(i))
     inout[i, i]  = -np.sum(flux, axis=0)
     inout[i, i] -= beta * \
         np.abs(coloumb(zfree[i])) * np.exp(-beta * np.sign(zfree[i]) * (phi[i] - phi_[i]))
     # non-diagonal elements
-    for jf, d in zip(flux, mydisp_):
+    for jf, d in zip(flux, disp_):
         inout[i, i + d] = jf
     # impose the boundary conditions
     _impose_j_bound(inout, nx, ny, nz, typ, val=1.0, mask=-1.0)
@@ -326,7 +106,7 @@ def _rhsvec(inout, flux, i, nx, ny, nz, typ, zfree, zfixed, phi, phi_, dirichlet
     
     inout *= -1 # for convenience change the sign of B in later NR iteration
 
-def calculate(jinout: Optional[lil_matrix], 
+def nr_construct(jinout: Optional[lil_matrix], 
               binout: Optional[np.ndarray],
               grid_dim: Tuple[int, int, int]|List[int],
               gridpoint_coords: np.ndarray, 
@@ -469,6 +249,3 @@ def calculate(jinout: Optional[lil_matrix],
         _rhsvec(binout, bflux, ind, nx, ny, nz, typ, zfree, zfixed, phi, phi_, 
                 dirichlet_pot, beta)
         assert binout.shape == (nr,)
-
-if __name__ == '__main__':
-    unittest.main()

From f3ca2ae1144df75fce5d9f11ad09849e392df8ee Mon Sep 17 00:00:00 2001
From: AsymmetryChou <181240085@smail.nju.edu.cn>
Date: Mon, 22 Dec 2025 23:36:39 +0800
Subject: [PATCH 12/14] Refactor: update boundary type handling in impose
 functions

---
 dpnegf/negf/newton_raphson_speed_up.py | 22 +++++++++++-----------
 1 file changed, 11 insertions(+), 11 deletions(-)

diff --git a/dpnegf/negf/newton_raphson_speed_up.py b/dpnegf/negf/newton_raphson_speed_up.py
index 3e96b91..c60fcd4 100644
--- a/dpnegf/negf/newton_raphson_speed_up.py
+++ b/dpnegf/negf/newton_raphson_speed_up.py
@@ -8,33 +8,33 @@
 
 def _impose_j_bound(inout, nx, ny, nz, typ, val, mask):
     '''impose the special mask for boundary points'''
-    boundary_types_ = ['xmin', 'xmax', 'ymin', 'ymax', 'zmin', 'zmax', 'Dirichlet']
+    # 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(boundary_types_, displ_):
-        # perf bottle-neck: search
-        i = np.where(np.array(typ) == t)[0] # all indices of the boundary type
+    for t, d in zip(neumann_types_, displ_):
+        i = np.where(np.array(typ) == t)[0]
         if len(i) == 0:
             log.warning(f'no grid point with type {t} found')
-        # scipy sparse matrix also supports the parallized assignment
+        # scipy sparse matrix also supports the parallelized assignment
         inout[i, i] = val
         inout[i, i + d] = mask
 
-    # there is also a special type "Dirichlet" for the Dirichlet boundary condition
+    # Dirichlet boundary
     i = np.where(typ == 'Dirichlet')[0]
     inout[i, i] = val
 
 def _impose_b_bound(inout, nx, ny, nz, typ, phi, dirichlet_pot):
     '''impose the special mask for boundary points in b vector'''
-    boundary_types_ = ['xmin', 'xmax', 'ymin', 'ymax', 'zmin', 'zmax', 'Dirichlet']
+    # 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(boundary_types_, displ_):
-        # perf bottle-neck: search
+    for t, d in zip(neumann_types_, displ_):
         i = np.where(typ == t)[0]
         if len(i) == 0:
             log.warning(f'no grid point with type {t} found')
         inout[i] = phi[i] - phi[i + d]
-        
-    # there is also a special type "Dirichlet" for the Dirichlet boundary condition
+
+    # Dirichlet boundary
     i = np.where(typ == 'Dirichlet')[0]
     inout[i] = phi[i] - dirichlet_pot[i]
 

From b5593085d7fcff0544188d15c80b2ce8ed748b9c Mon Sep 17 00:00:00 2001
From: AsymmetryChou <181240085@smail.nju.edu.cn>
Date: Mon, 22 Dec 2025 23:37:45 +0800
Subject: [PATCH 13/14] Fix: correct spelling of 'coulomb' function in
 newton_raphson_speed_up.py

---
 dpnegf/negf/newton_raphson_speed_up.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/dpnegf/negf/newton_raphson_speed_up.py b/dpnegf/negf/newton_raphson_speed_up.py
index c60fcd4..ae2bdcb 100644
--- a/dpnegf/negf/newton_raphson_speed_up.py
+++ b/dpnegf/negf/newton_raphson_speed_up.py
@@ -38,7 +38,7 @@ def _impose_b_bound(inout, nx, ny, nz, typ, phi, dirichlet_pot):
     i = np.where(typ == 'Dirichlet')[0]
     inout[i] = phi[i] - dirichlet_pot[i]
 
-def coloumb(chr: float|np.ndarray) -> float|np.ndarray:
+def coulomb(chr: float|np.ndarray) -> float|np.ndarray:
     '''convert the charge to unit of Coulomb'''
     assert isinstance(chr, (float, np.ndarray)), "chr must be a float or numpy array"
     from dpnegf.utils.constants import elementary_charge as e
@@ -86,7 +86,7 @@ def _jacobian(inout, flux, i, nx, ny, nz, typ, zfree, phi, phi_, beta):
     assert flux.shape == (6, len(i))
     inout[i, i]  = -np.sum(flux, axis=0)
     inout[i, i] -= beta * \
-        np.abs(coloumb(zfree[i])) * np.exp(-beta * np.sign(zfree[i]) * (phi[i] - phi_[i]))
+        np.abs(coulomb(zfree[i])) * np.exp(-beta * np.sign(zfree[i]) * (phi[i] - phi_[i]))
     # non-diagonal elements
     for jf, d in zip(flux, disp_):
         inout[i, i + d] = jf
@@ -98,9 +98,9 @@ def _rhsvec(inout, flux, i, nx, ny, nz, typ, zfree, zfixed, phi, phi_, dirichlet
     # calculate the b vector
     assert flux.shape == (6, len(i))
     inout[i]  = np.sum(flux, axis=0)
-    inout[i] += coloumb(zfree[i]) * \
+    inout[i] += coulomb(zfree[i]) * \
         np.exp(-beta * np.sign(zfree[i]) * (phi[i] - phi_[i]))
-    inout[i] += coloumb(zfixed[i])
+    inout[i] += coulomb(zfixed[i])
     # impose the boundary conditions
     _impose_b_bound(inout, nx, ny, nz, typ, phi, dirichlet_pot)
     

From eba14c2ed8ea4ffb0f6e138c9d61fa4ddd637988 Mon Sep 17 00:00:00 2001
From: AsymmetryChou <181240085@smail.nju.edu.cn>
Date: Tue, 23 Dec 2025 10:15:03 +0800
Subject: [PATCH 14/14] =?UTF-8?q?Feat:=20=E6=B7=BB=E5=8A=A0=E5=A4=9A?=
 =?UTF-8?q?=E4=B8=AA=E9=92=88=E5=AF=B9=20Interface3D=20=E7=9A=84=E6=96=B0?=
 =?UTF-8?q?=E6=B5=8B=E8=AF=95=E7=94=A8=E4=BE=8B=EF=BC=8C=E6=B6=B5=E7=9B=96?=
 =?UTF-8?q?=E5=86=85=E9=83=A8=E7=82=B9=E5=92=8C=E4=B8=8D=E5=90=8C=E5=B9=B3?=
 =?UTF-8?q?=E5=9D=87=E6=A8=A1=E5=BC=8F?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 dpnegf/tests/test_poisson_init.py | 161 ++++++++++++++++++++++++++++++
 1 file changed, 161 insertions(+)

diff --git a/dpnegf/tests/test_poisson_init.py b/dpnegf/tests/test_poisson_init.py
index 650488e..6b4c373 100644
--- a/dpnegf/tests/test_poisson_init.py
+++ b/dpnegf/tests/test_poisson_init.py
@@ -355,6 +355,167 @@ def test_interface3d_NR_construct_Jac_B_boundary_and_internal():
     assert np.allclose(B, B_std)
 
 
+# ========================= New tests after refactoring =========================
+
+def make_grid_with_internal_points():
+    """Create a 3x3x3 grid that has internal points (not just boundary)."""
+    xg = np.array([0.0, 1.0, 2.0])
+    yg = np.array([0.0, 1.0, 2.0])
+    zg = np.array([0.0, 1.0, 2.0])
+    # Place atoms at corners only
+    xa = np.array([0.0, 2.0])
+    ya = np.array([0.0, 2.0])
+    za = np.array([0.0, 2.0])
+    return Grid(xg, yg, zg, xa, ya, za)
+
+
+def test_interface3d_eps_average_mode_default():
+    """Test that default eps_average_mode is 'harmonic'."""
+    grid = make_simple_grid()
+    d = Dirichlet((0, 0), (0, 1), (0, 1))
+    diel = Dielectric((0, 1), (0, 1), (0, 1))
+    iface = Interface3D(grid, [d], [diel])
+    assert iface.average_mode == 'harmonic'
+
+
+def test_interface3d_eps_average_mode_options():
+    """Test that all eps_average_mode options are accepted."""
+    grid = make_simple_grid()
+    d = Dirichlet((0, 0), (0, 1), (0, 1))
+    diel = Dielectric((0, 1), (0, 1), (0, 1))
+
+    for mode in ['harmonic', 'arithmetic', 'geometric']:
+        iface = Interface3D(grid, [d], [diel], eps_average_mode=mode)
+        assert iface.average_mode == mode
+
+
+def test_interface3d_with_internal_points():
+    """Test Interface3D with a grid that has internal points."""
+    grid = make_grid_with_internal_points()
+    d = Dirichlet((0, 0), (0, 2), (0, 2))
+    diel = Dielectric((0, 2), (0, 2), (0, 2))
+    iface = Interface3D(grid, [d], [diel])
+
+    # Should have internal points (center point of 3x3x3 grid)
+    internal_count = sum(1 for v in iface.boudnary_points.values() if v == "in")
+    assert internal_count > 0, "Expected at least one internal point in 3x3x3 grid"
+    assert iface.internal_NP == internal_count
+
+
+def test_interface3d_NR_construct_with_internal_points():
+    """Test NR_construct_Jac_B with internal points exercises flux calculations."""
+    grid = make_grid_with_internal_points()
+    d = Dirichlet((0, 0), (0, 2), (0, 2))
+    diel = Dielectric((0, 2), (0, 2), (0, 2))
+    diel.eps = 3.9  # Set non-trivial permittivity
+    iface = Interface3D(grid, [d], [diel])
+    iface.get_potential_eps([d, diel])
+
+    J = lil_matrix((grid.Np, grid.Np))
+    B = np.zeros(grid.Np)
+    iface.NR_construct_Jac_B(J, B)
+
+    # Check matrix properties
+    assert J.shape == (grid.Np, grid.Np)
+    assert B.shape == (grid.Np,)
+
+    # For internal points, diagonal should be negative (sum of fluxes)
+    # For boundary points, diagonal should be 1.0
+    for i in range(grid.Np):
+        if iface.boudnary_points[i] == "in":
+            # Internal point: diagonal is negative (flux sum + charge term)
+            assert J[i, i] < 0, f"Internal point {i} should have negative diagonal"
+        else:
+            # Boundary point: diagonal is 1.0
+            assert J[i, i] == 1.0, f"Boundary point {i} should have diagonal = 1.0"
+
+
+def test_interface3d_different_average_modes_produce_different_results():
+    """Test that different averaging modes produce different Jacobian matrices."""
+    grid = make_grid_with_internal_points()
+    d = Dirichlet((0, 0), (0, 2), (0, 2))
+    diel = Dielectric((0, 2), (0, 2), (0, 2))
+    diel.eps = 3.9
+
+    results = {}
+    for mode in ['harmonic', 'arithmetic', 'geometric']:
+        iface = Interface3D(grid, [d], [diel], eps_average_mode=mode)
+        iface.get_potential_eps([d, diel])
+        J = lil_matrix((grid.Np, grid.Np))
+        B = np.zeros(grid.Np)
+        iface.NR_construct_Jac_B(J, B)
+        results[mode] = J.toarray()
+
+    # Different modes should produce different matrices (at least for internal points)
+    # Since all have same eps, the averaging modes will give same result for uniform eps
+    # This test verifies the mode is being applied (no runtime error)
+    for mode, J_arr in results.items():
+        assert J_arr.shape == (grid.Np, grid.Np)
+
+
+def test_interface3d_solve_poisson_NRcycle_valid_dtypes():
+    """Test that valid dtypes (float32, float64) work without error."""
+    # Use grid with internal points to avoid singular matrix
+    grid = make_grid_with_internal_points()
+    d = Dirichlet((0, 0), (0, 2), (0, 2))
+    diel = Dielectric((0, 2), (0, 2), (0, 2))
+    diel.eps = 3.9
+
+    for dtype in ['float32', 'float64']:
+        iface = Interface3D(grid, [d], [diel])
+        iface.get_potential_eps([d, diel])
+        # Should not raise - just verify it runs
+        result = iface.solve_poisson_NRcycle(method='scipy', dtype=dtype)
+        assert isinstance(result, (float, np.floating))
+
+
+def test_interface3d_solve_poisson_NRcycle_scipy_method():
+    """Test that scipy method works correctly."""
+    # Use grid with internal points to avoid singular matrix
+    grid = make_grid_with_internal_points()
+    d = Dirichlet((0, 0), (0, 2), (0, 2))
+    diel = Dielectric((0, 2), (0, 2), (0, 2))
+    diel.eps = 3.9
+    iface = Interface3D(grid, [d], [diel])
+    iface.get_potential_eps([d, diel])
+
+    # Should complete without error
+    max_diff = iface.solve_poisson_NRcycle(method='scipy')
+    assert isinstance(max_diff, (float, np.floating))
+    assert max_diff >= 0  # max_diff should be non-negative
+
+
+def test_interface3d_solve_poisson_NRcycle_convergence():
+    """Test that NR cycle converges for a simple case."""
+    grid = make_grid_with_internal_points()
+    d = Dirichlet((0, 0), (0, 2), (0, 2))
+    d.Ef = 1.0  # Set non-zero potential
+    diel = Dielectric((0, 2), (0, 2), (0, 2))
+    diel.eps = 3.9
+    iface = Interface3D(grid, [d], [diel])
+    iface.get_potential_eps([d, diel])
+
+    # Run NR cycle
+    max_diff = iface.solve_poisson_NRcycle(method='scipy', dtype='float64')
+
+    # Should converge (max_diff is finite)
+    assert np.isfinite(max_diff)
+
+
+def test_interface3d_NR_construct_only_jacobian():
+    """Test that nr_construct can be called with only Jacobian (binout=None internally)."""
+    grid = make_simple_grid()
+    d = Dirichlet((0, 0), (0, 1), (0, 1))
+    diel = Dielectric((0, 1), (0, 1), (0, 1))
+    iface = Interface3D(grid, [d], [diel])
+
+    # to_scipy_Jac_B always computes both, but we verify the output is valid
+    J, B = iface.to_scipy_Jac_B()
+    assert J is not None
+    assert B is not None
+    assert J.nnz > 0  # Jacobian should have non-zero entries
+
+