rust-simulation-tools 0.2.2__tar.gz → 0.2.4__tar.gz

{rust_simulation_tools-0.2.2 → rust_simulation_tools-0.2.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: rust_simulation_tools
-Version: 0.2.2
+Version: 0.2.4
 Classifier: Programming Language :: Rust
 Classifier: Programming Language :: Python :: Implementation :: CPython
 Requires-Dist: numpy>=1.26

{rust_simulation_tools-0.2.2 → rust_simulation_tools-0.2.4}/crates/bindings/src/lib.rs

@@ -1539,7 +1539,7 @@ struct PyPbParams {
 #[pymethods]
 impl PyPbParams {
     #[new]
-    #[pyo3(signature = (grid_spacing=None, grid_buffer=None, solute_dielectric=None, solvent_dielectric=None, salt_concentration=None, temperature=None, probe_radius=None, ion_radius=None, tolerance=None, max_iterations=None))]
+    #[pyo3(signature = (grid_spacing=None, grid_buffer=None, solute_dielectric=None, solvent_dielectric=None, salt_concentration=None, temperature=None, probe_radius=None, ion_radius=None, tolerance=None, max_iterations=None, fillratio=None, fscale=None))]
     fn new(
         grid_spacing: Option<f64>,
         grid_buffer: Option<f64>,
@@ -1551,6 +1551,8 @@ impl PyPbParams {
         ion_radius: Option<f64>,
         tolerance: Option<f64>,
         max_iterations: Option<usize>,
+        fillratio: Option<f64>,
+        fscale: Option<usize>,
     ) -> Self {
         let mut params = PbParams::default();
         if let Some(v) = grid_spacing {
@@ -1583,6 +1585,12 @@ impl PyPbParams {
         if let Some(v) = max_iterations {
             params.max_iterations = v;
         }
+        if let Some(v) = fillratio {
+            params.fillratio = v;
+        }
+        if let Some(v) = fscale {
+            params.fscale = v;
+        }
         PyPbParams { inner: params }
     }
 

{rust_simulation_tools-0.2.2 → rust_simulation_tools-0.2.4}/crates/mmpbsa/src/gb_energy.rs

@@ -241,7 +241,7 @@ fn compute_born_radii(
             let psi2 = psi_rho * psi_rho;
             let psi3 = psi2 * psi_rho;
             let tanh_val = (alpha * psi_rho - beta * psi2 + gamma * psi3).tanh();
-            let inv_r = 1.0 / rho_i - tanh_val / rho_i;
+            let inv_r = 1.0 / rho_i - tanh_val / topology.radii[i];
             born_radii[i] = if inv_r > 0.0 {
                 1.0 / inv_r
             } else {
@@ -267,10 +267,14 @@ fn compute_born_radii(
 /// # Returns
 /// `GbEnergy` containing the total energy in kcal/mol and per-atom Born radii.
 ///
-/// The GB energy is computed as:
-///   E_GB = -½(1/ε_in - 1/ε_out) Σ_{i,j} q_i·q_j·exp(-κ·f_GB) / f_GB
+/// The GB energy is computed using the AMBER convention:
+///   E_GB = -½ Σ_{i,j} q_i·q_j · (1/ε_in - exp(-KSCALE·κ·f_GB)/ε_out) / f_GB
 ///
 /// where f_GB = sqrt(r²_ij + R_i·R_j·exp(-r²_ij/(4·R_i·R_j)))
+/// and KSCALE = 0.73 (Srinivasan et al.)
+///
+/// When salt_concentration = 0 (κ = 0), this reduces to:
+///   E_GB = -½(1/ε_in - 1/ε_out) Σ_{i,j} q_i·q_j / f_GB
 ///
 /// Charges are taken from `charges_amber` (AMBER internal units where
 /// q_amber = q_real × 18.2223, so q_i·q_j already yields kcal·Å/mol when
@@ -283,23 +287,34 @@ fn compute_gb_energy_from_radii(
     born_radii: &[f64],
 ) -> f64 {
     let n = topology.n_atoms;
-    let dielectric_factor =
-        -0.5 * (1.0 / params.solute_dielectric - 1.0 / params.solvent_dielectric);
     let kappa = compute_kappa(
         params.salt_concentration,
         params.solvent_dielectric,
         params.temperature,
     );
     let cutoff_sq = params.cutoff * params.cutoff;
+    let inv_ein = 1.0 / params.solute_dielectric;
+    let inv_eout = 1.0 / params.solvent_dielectric;
+
+    // AMBER Debye-Hückel scaling factor (Srinivasan et al.)
+    const KSCALE: f64 = 0.73;
+
+    // Dielectric factor depends on whether salt screening is active:
+    //   With salt:    E = -0.5 * q_i*q_j * (1/ε_in - exp(-KSCALE*κ*f)/ε_out) / f
+    //   Without salt: E = -0.5 * q_i*q_j * (1/ε_in - 1/ε_out) / f
+    let no_salt = kappa <= 0.0;
+    let dielectric_no_salt = -0.5 * (inv_ein - inv_eout);
 
     // Self-energy terms (O(N), kept serial)
     let mut energy = 0.0;
     for i in 0..n {
         let qi = topology.charges_amber[i];
-        let mut self_energy = dielectric_factor * qi * qi / born_radii[i];
-        if kappa > 0.0 {
-            self_energy *= (-kappa * born_radii[i]).exp();
-        }
+        let self_energy = if no_salt {
+            dielectric_no_salt * qi * qi / born_radii[i]
+        } else {
+            -0.5 * qi * qi * (inv_ein - (-KSCALE * kappa * born_radii[i]).exp() * inv_eout)
+                / born_radii[i]
+        };
         energy += self_energy;
     }
 
@@ -324,10 +339,11 @@ fn compute_gb_energy_from_radii(
                 let ri_rj = born_radii[i] * born_radii[j];
                 let f_gb = (r2 + ri_rj * (-r2 / (4.0 * ri_rj)).exp()).sqrt();
 
-                let mut pair_energy = dielectric_factor * qi * qj / f_gb;
-                if kappa > 0.0 {
-                    pair_energy *= (-kappa * f_gb).exp();
-                }
+                let pair_energy = if no_salt {
+                    dielectric_no_salt * qi * qj / f_gb
+                } else {
+                    -0.5 * qi * qj * (inv_ein - (-KSCALE * kappa * f_gb).exp() * inv_eout) / f_gb
+                };
 
                 sum += 2.0 * pair_energy;
             }

{rust_simulation_tools-0.2.2 → rust_simulation_tools-0.2.4}/crates/mmpbsa/src/mm_energy.rs

@@ -201,13 +201,14 @@ pub fn compute_mm_energy_with_nb(
         energy.angle += ka * (theta - theta_eq) * (theta - theta_eq);
     }
 
-    // Dihedral energy: E = (V_n/2) * [1 + cos(n*phi - gamma)]
+    // Dihedral energy: E = pk * [1 + cos(n*phi - gamma)]
+    // AMBER prmtop DIHEDRAL_FORCE_CONSTANT stores pk (already Vn/2, the half-barrier height)
     for &(i, j, k, l, type_idx, _ignore_14) in &topology.dihedrals {
-        let vn = topology.dihedral_force_constants[type_idx];
+        let pk = topology.dihedral_force_constants[type_idx];
         let n = topology.dihedral_periodicities[type_idx];
         let gamma = topology.dihedral_phases[type_idx];
         let phi = compute_dihedral(&coords[i], &coords[j], &coords[k], &coords[l]);
-        energy.dihedral += (vn / 2.0) * (1.0 + (n * phi - gamma).cos());
+        energy.dihedral += pk * (1.0 + (n * phi - gamma).cos());
     }
 
     // Non-bonded interactions (parallelized over outer atom index)
@@ -380,7 +381,8 @@ mod tests {
         top.lj_acoef = Arc::new(vec![0.0]);
         top.lj_bcoef = Arc::new(vec![0.0]);
 
-        // Trans dihedral (phi=PI): E = (2/2) * [1 + cos(2*PI - PI)] = 1 * [1 + cos(PI)] = 0
+        // Trans dihedral (phi=PI): E = 2 * [1 + cos(2*PI - PI)] = 2 * [1 + cos(PI)] = 0
+        // (pk=2.0 is already the half-barrier height in AMBER prmtop convention)
         let coords = [
             [1.0, 1.0, 0.0],
             [0.0, 0.0, 0.0],

{rust_simulation_tools-0.2.2 → rust_simulation_tools-0.2.4}/crates/mmpbsa/src/pb_energy.rs

@@ -7,8 +7,12 @@ use rayon;
 
 use crate::pb_grid::{
     assign_dielectrics, assign_kappa, auto_grid, map_charges, topology_charges, topology_radii,
+    PbGrid,
+};
+use crate::pb_solver::{
+    compute_elec_energy, interpolated_boundary, solve_lpbe_multigrid, BoundaryCondition,
+    PbSolveResult,
 };
-use crate::pb_solver::{compute_elec_energy, solve_lpbe_multigrid, BoundaryCondition};
 use crate::subsystem::{extract_coords, extract_subtopology};
 use rst_core::amber::prmtop::AmberTopology;
 
@@ -35,6 +39,11 @@ pub struct PbParams {
     pub tolerance: f64,
     /// Maximum SOR iterations. Default: 10000.
     pub max_iterations: usize,
+    /// Ratio of coarse grid extent to molecule extent for focusing. Default: 4.0.
+    /// Set to 0.0 to disable focusing (not recommended for large molecules).
+    pub fillratio: f64,
+    /// Ratio of coarse grid spacing to fine grid spacing. Default: 8.
+    pub fscale: usize,
 }
 
 impl Default for PbParams {
@@ -50,6 +59,8 @@ impl Default for PbParams {
             ion_radius: 2.0,
             tolerance: 1e-6,
             max_iterations: 10000,
+            fillratio: 4.0,
+            fscale: 8,
         }
     }
 }
@@ -74,13 +85,191 @@ pub(crate) fn compute_kappa(salt_conc: f64, solvent_dielectric: f64, temperature
     (factor * salt_conc / (solvent_dielectric * kb * temperature)).sqrt()
 }
 
+/// Build a coarse grid for focusing.
+///
+/// The coarse grid covers `fillratio` × the molecule extent, with spacing
+/// equal to `fscale` × the fine grid spacing.
+fn build_coarse_grid(
+    coords: &[[f64; 3]],
+    fine_spacing: f64,
+    fillratio: f64,
+    fscale: usize,
+) -> PbGrid {
+    let mut min = [f64::MAX; 3];
+    let mut max = [f64::MIN; 3];
+    for c in coords {
+        for d in 0..3 {
+            if c[d] < min[d] {
+                min[d] = c[d];
+            }
+            if c[d] > max[d] {
+                max[d] = c[d];
+            }
+        }
+    }
+
+    let coarse_spacing = fine_spacing * fscale as f64;
+    let mut dims = [0usize; 3];
+    let mut origin = [0.0f64; 3];
+    for d in 0..3 {
+        let mol_extent = max[d] - min[d];
+        let grid_extent = mol_extent * fillratio;
+        let n = (grid_extent / coarse_spacing).ceil() as usize + 1;
+        dims[d] = if n.is_multiple_of(2) { n + 1 } else { n };
+        let actual_extent = (dims[d] - 1) as f64 * coarse_spacing;
+        let center = 0.5 * (min[d] + max[d]);
+        origin[d] = center - 0.5 * actual_extent;
+    }
+
+    PbGrid::new(dims, [coarse_spacing; 3], origin)
+}
+
+/// Perform a single PB solve (solvated or reference) with optional focusing.
+///
+/// If `coarse_grid` is provided, first solves on the coarse grid, then
+/// interpolates boundary conditions for the fine grid.
+fn pb_solve_with_focusing(
+    fine_grid: &PbGrid,
+    coarse_grid: Option<&PbGrid>,
+    coords: &[[f64; 3]],
+    charges: &[f64],
+    radii: &[f64],
+    params: &PbParams,
+    eps_in: f64,
+    eps_out: f64,
+    kappa: f64,
+    use_ionic: bool,
+) -> (PbSolveResult, f64) {
+    // Determine DH BC parameters for the initial (or only) solve
+    let bc_kappa = if use_ionic { kappa } else { 0.0 };
+
+    match coarse_grid {
+        Some(cgrid) => {
+            // === Two-level focusing ===
+            // 1. Solve on coarse grid with DH/Coulomb BCs
+            let coarse_charges = map_charges(cgrid, coords, charges);
+            let coarse_diel =
+                assign_dielectrics(cgrid, coords, radii, params.probe_radius, eps_in, eps_out);
+            let coarse_kappa = if use_ionic {
+                assign_kappa(
+                    cgrid,
+                    coords,
+                    radii,
+                    params.ion_radius,
+                    kappa,
+                    params.solvent_dielectric,
+                )
+            } else {
+                vec![0.0; cgrid.len()]
+            };
+
+            let coarse_result = solve_lpbe_multigrid(
+                cgrid,
+                &coarse_charges,
+                &coarse_diel,
+                &coarse_kappa,
+                BoundaryCondition::DebyeHuckel,
+                coords,
+                charges,
+                bc_kappa,
+                eps_out,
+                params.tolerance * 10.0, // coarser tolerance for speed
+                params.max_iterations,
+            );
+
+            // 2. Interpolate coarse solution as BCs for fine grid
+            let fine_bc = interpolated_boundary(fine_grid, cgrid, &coarse_result.potential);
+
+            // 3. Solve on fine grid with interpolated BCs
+            let fine_charges = map_charges(fine_grid, coords, charges);
+            let fine_diel = assign_dielectrics(
+                fine_grid,
+                coords,
+                radii,
+                params.probe_radius,
+                eps_in,
+                eps_out,
+            );
+            let fine_kappa = if use_ionic {
+                assign_kappa(
+                    fine_grid,
+                    coords,
+                    radii,
+                    params.ion_radius,
+                    kappa,
+                    params.solvent_dielectric,
+                )
+            } else {
+                vec![0.0; fine_grid.len()]
+            };
+
+            let result = solve_lpbe_multigrid(
+                fine_grid,
+                &fine_charges,
+                &fine_diel,
+                &fine_kappa,
+                BoundaryCondition::Interpolated(fine_bc),
+                coords,
+                charges,
+                bc_kappa,
+                eps_out,
+                params.tolerance,
+                params.max_iterations,
+            );
+            let energy = compute_elec_energy(fine_grid, &result.potential, coords, charges);
+            (result, energy)
+        }
+        None => {
+            // === Single-level solve ===
+            let charge_map = map_charges(fine_grid, coords, charges);
+            let diel = assign_dielectrics(
+                fine_grid,
+                coords,
+                radii,
+                params.probe_radius,
+                eps_in,
+                eps_out,
+            );
+            let kappa_map = if use_ionic {
+                assign_kappa(
+                    fine_grid,
+                    coords,
+                    radii,
+                    params.ion_radius,
+                    kappa,
+                    params.solvent_dielectric,
+                )
+            } else {
+                vec![0.0; fine_grid.len()]
+            };
+
+            let result = solve_lpbe_multigrid(
+                fine_grid,
+                &charge_map,
+                &diel,
+                &kappa_map,
+                BoundaryCondition::DebyeHuckel,
+                coords,
+                charges,
+                bc_kappa,
+                eps_out,
+                params.tolerance,
+                params.max_iterations,
+            );
+            let energy = compute_elec_energy(fine_grid, &result.potential, coords, charges);
+            (result, energy)
+        }
+    }
+}
+
 /// Compute PB polar solvation energy for a molecular system.
 ///
-/// Performs two solves:
-/// 1. Solvated: ε_in inside molecule, ε_out outside, with ionic strength
-/// 2. Reference: ε_in everywhere (vacuum), no ionic strength
+/// Uses two-level focusing (when `fillratio > 0`):
+/// 1. Coarse solve on a large grid with DH/Coulomb boundary conditions
+/// 2. Fine solve with boundary conditions interpolated from the coarse solution
 ///
-/// Returns ΔG_PB = E_solvated - E_reference (reaction field energy).
+/// Performs paired solves (solvated + reference) and returns
+/// ΔG_PB = E_solvated - E_reference (reaction field energy).
 pub fn compute_pb_energy(
     topology: &AmberTopology,
     coords: &[[f64; 3]],
@@ -95,77 +284,60 @@ pub fn compute_pb_energy(
         params.temperature,
     );
 
-    // Build grid
-    let grid = auto_grid(coords, params.grid_spacing, params.grid_buffer);
-    log::info!(
-        "PB grid: {}x{}x{} ({} points) for {} atoms",
-        grid.dims[0],
-        grid.dims[1],
-        grid.dims[2],
-        grid.len(),
-        coords.len()
-    );
+    // Build fine grid (centered on molecule with buffer)
+    let fine_grid = auto_grid(coords, params.grid_spacing, params.grid_buffer);
 
-    // Map charges onto grid
-    let charge_map = map_charges(&grid, coords, &charges);
+    // Build coarse grid for focusing (if enabled)
+    let coarse_grid = if params.fillratio > 0.0 && params.fscale > 1 {
+        let cg = build_coarse_grid(coords, params.grid_spacing, params.fillratio, params.fscale);
+        log::info!(
+            "PB focusing: coarse {}x{}x{} (spacing {:.1}), fine {}x{}x{} (spacing {:.1}) for {} atoms",
+            cg.dims[0], cg.dims[1], cg.dims[2], cg.spacing[0],
+            fine_grid.dims[0], fine_grid.dims[1], fine_grid.dims[2], fine_grid.spacing[0],
+            coords.len()
+        );
+        Some(cg)
+    } else {
+        log::info!(
+            "PB grid: {}x{}x{} ({} points) for {} atoms",
+            fine_grid.dims[0],
+            fine_grid.dims[1],
+            fine_grid.dims[2],
+            fine_grid.len(),
+            coords.len()
+        );
+        None
+    };
 
     // Run solvated and reference solves in parallel
     let ((result_solv, e_solv), (result_ref, e_ref)) = rayon::join(
         || {
-            // === Solvated solve ===
-            let diel_solv = assign_dielectrics(
-                &grid,
+            pb_solve_with_focusing(
+                &fine_grid,
+                coarse_grid.as_ref(),
                 coords,
+                &charges,
                 &radii,
-                params.probe_radius,
+                params,
                 params.solute_dielectric,
                 params.solvent_dielectric,
-            );
-            let kappa_sq_solv = assign_kappa(&grid, coords, &radii, params.ion_radius, kappa);
-
-            let result = solve_lpbe_multigrid(
-                &grid,
-                &charge_map,
-                &diel_solv,
-                &kappa_sq_solv,
-                BoundaryCondition::DebyeHuckel,
-                coords,
-                &charges,
                 kappa,
-                params.solvent_dielectric,
-                params.tolerance,
-                params.max_iterations,
-            );
-            let energy = compute_elec_energy(&grid, &result.potential, coords, &charges);
-            (result, energy)
+                true, // use ionic strength
+            )
         },
         || {
-            // === Reference (vacuum) solve ===
-            let diel_ref = assign_dielectrics(
-                &grid,
-                coords,
-                &radii,
-                params.probe_radius,
-                params.solute_dielectric,
-                params.solute_dielectric,
-            );
-            let kappa_sq_ref = vec![0.0; grid.len()];
-
-            let result = solve_lpbe_multigrid(
-                &grid,
-                &charge_map,
-                &diel_ref,
-                &kappa_sq_ref,
-                BoundaryCondition::Zero,
+            pb_solve_with_focusing(
+                &fine_grid,
+                coarse_grid.as_ref(),
                 coords,
                 &charges,
-                0.0,
+                &radii,
+                params,
                 params.solute_dielectric,
-                params.tolerance,
-                params.max_iterations,
-            );
-            let energy = compute_elec_energy(&grid, &result.potential, coords, &charges);
-            (result, energy)
+                params.solute_dielectric, // uniform ε_in for reference
+                kappa,
+                false, // no ionic strength
+            )
         },
     );
 
@@ -186,11 +358,13 @@ pub fn compute_pb_energy(
     }
 
     log::debug!(
-        "PB energies: E_solv = {:.6} (iters={}), E_ref = {:.6} (iters={}), delta = {:.6}",
+        "PB energies: E_solv = {:.4} (iters={}, res={:.2e}), E_ref = {:.4} (iters={}, res={:.2e}), delta = {:.4}",
         e_solv,
         result_solv.iterations,
+        result_solv.final_residual,
         e_ref,
         result_ref.iterations,
+        result_ref.final_residual,
         e_solv - e_ref
     );
 

{rust_simulation_tools-0.2.2 → rust_simulation_tools-0.2.4}/crates/mmpbsa/src/pb_grid.rs

@@ -6,6 +6,7 @@
 
 use rayon::prelude::*;
 use rst_core::amber::prmtop::AmberTopology;
+use std::collections::VecDeque;
 
 /// A regular 3D grid for finite-difference PB calculations.
 #[derive(Debug, Clone)]
@@ -317,26 +318,48 @@ fn point_inside_molecule(
     false
 }
 
-/// Assign face-centered dielectric values based on molecular surface.
+/// Compute Solvent-Excluded Surface (SES) inside/outside classification.
 ///
-/// Uses a simple approach: for each grid edge, test both endpoints.
-/// If both are inside the molecule → eps_in, both outside → eps_out,
-/// mixed → harmonic mean.
-pub fn assign_dielectrics(
+/// The SES (molecular surface) is the boundary between regions accessible
+/// to a probe sphere of the given radius and regions that are not. This is
+/// tighter than the SAS (Solvent-Accessible Surface) because crevices and
+/// grooves between atoms are classified as interior even though they are
+/// outside the union of inflated spheres.
+///
+/// Algorithm:
+/// 1. Identify valid probe-center positions: grid points outside the SAS
+///    (where a probe sphere wouldn't overlap any atom) connected to the
+///    bulk (grid boundary) via a path of such positions.
+/// 2. A grid point is outside the SES if any valid probe center within
+///    `probe_radius` can reach it; otherwise it is inside.
+fn compute_inside_ses(
     grid: &PbGrid,
     coords: &[[f64; 3]],
     radii: &[f64],
     probe_radius: f64,
-    eps_in: f64,
-    eps_out: f64,
-) -> DielectricMaps {
+) -> Vec<bool> {
     let nx = grid.dims[0];
     let ny = grid.dims[1];
     let nz = grid.dims[2];
     let n = nx * ny * nz;
 
-    // First, compute per-node inside/outside flag
-    let inside: Vec<bool> = (0..n)
+    if probe_radius <= 0.0 {
+        // No probe: SES = VDW surface
+        return (0..n)
+            .into_par_iter()
+            .map(|idx| {
+                let iz = idx / (nx * ny);
+                let iy = (idx % (nx * ny)) / nx;
+                let ix = idx % nx;
+                let pt = grid.point(ix, iy, iz);
+                point_inside_molecule(&pt, coords, radii, 0.0)
+            })
+            .collect();
+    }
+
+    // Step 1: Compute inside_sas (SAS = union of spheres with r + probe).
+    // Outside SAS = valid probe-center positions (probe fits without overlap).
+    let inside_sas: Vec<bool> = (0..n)
         .into_par_iter()
         .map(|idx| {
             let iz = idx / (nx * ny);
@@ -347,14 +370,239 @@ pub fn assign_dielectrics(
         })
         .collect();
 
-    let eps_for = |a: bool, b: bool| -> f64 {
-        match (a, b) {
-            (true, true) => eps_in,
-            (false, false) => eps_out,
-            _ => 2.0 * eps_in * eps_out / (eps_in + eps_out), // harmonic mean
+    // Step 2: Flood-fill from boundary through outside_sas to find the
+    // connected exterior (probe centers reachable from the bulk).
+    let mut connected_exterior = vec![false; n];
+    let mut queue = VecDeque::new();
+
+    // Seed: boundary points that are outside SAS
+    for iz in 0..nz {
+        for iy in 0..ny {
+            for ix in 0..nx {
+                let is_boundary =
+                    ix == 0 || ix == nx - 1 || iy == 0 || iy == ny - 1 || iz == 0 || iz == nz - 1;
+                if is_boundary {
+                    let idx = grid.index(ix, iy, iz);
+                    if !inside_sas[idx] {
+                        connected_exterior[idx] = true;
+                        queue.push_back(idx);
+                    }
+                }
+            }
+        }
+    }
+
+    // BFS through outside_sas
+    while let Some(idx) = queue.pop_front() {
+        let iz = idx / (nx * ny);
+        let iy = (idx % (nx * ny)) / nx;
+        let ix = idx % nx;
+
+        let neighbors: [Option<usize>; 6] = [
+            if ix > 0 {
+                Some(grid.index(ix - 1, iy, iz))
+            } else {
+                None
+            },
+            if ix + 1 < nx {
+                Some(grid.index(ix + 1, iy, iz))
+            } else {
+                None
+            },
+            if iy > 0 {
+                Some(grid.index(ix, iy - 1, iz))
+            } else {
+                None
+            },
+            if iy + 1 < ny {
+                Some(grid.index(ix, iy + 1, iz))
+            } else {
+                None
+            },
+            if iz > 0 {
+                Some(grid.index(ix, iy, iz - 1))
+            } else {
+                None
+            },
+            if iz + 1 < nz {
+                Some(grid.index(ix, iy, iz + 1))
+            } else {
+                None
+            },
+        ];
+
+        for ni in neighbors.into_iter().flatten() {
+            if !connected_exterior[ni] && !inside_sas[ni] {
+                connected_exterior[ni] = true;
+                queue.push_back(ni);
+            }
         }
+    }
+
+    // Step 3: Dilate connected exterior by probe_radius.
+    // A point is outside SES if any connected_exterior point is within probe_radius.
+    let probe_sq = probe_radius * probe_radius;
+    let rx = (probe_radius / grid.spacing[0]).ceil() as isize;
+    let ry = (probe_radius / grid.spacing[1]).ceil() as isize;
+    let rz = (probe_radius / grid.spacing[2]).ceil() as isize;
+
+    (0..n)
+        .into_par_iter()
+        .map(|idx| {
+            let iz = (idx / (nx * ny)) as isize;
+            let iy = ((idx % (nx * ny)) / nx) as isize;
+            let ix = (idx % nx) as isize;
+            let pt = grid.point(ix as usize, iy as usize, iz as usize);
+
+            for dz in -rz..=rz {
+                let jz = iz + dz;
+                if jz < 0 || jz >= nz as isize {
+                    continue;
+                }
+                for dy in -ry..=ry {
+                    let jy = iy + dy;
+                    if jy < 0 || jy >= ny as isize {
+                        continue;
+                    }
+                    for dx in -rx..=rx {
+                        let jx = ix + dx;
+                        if jx < 0 || jx >= nx as isize {
+                            continue;
+                        }
+
+                        let jidx = grid.index(jx as usize, jy as usize, jz as usize);
+                        if connected_exterior[jidx] {
+                            let neighbor_pt = grid.point(jx as usize, jy as usize, jz as usize);
+                            let ddx = pt[0] - neighbor_pt[0];
+                            let ddy = pt[1] - neighbor_pt[1];
+                            let ddz = pt[2] - neighbor_pt[2];
+                            if ddx * ddx + ddy * ddy + ddz * ddz <= probe_sq {
+                                return false; // Probe can reach -> outside SES
+                            }
+                        }
+                    }
+                }
+            }
+            true // No probe can reach -> inside SES
+        })
+        .collect()
+}
+
+/// Compute the fraction of an axis-aligned edge segment that lies inside the
+/// molecular surface (union of atom spheres inflated by probe radius).
+///
+/// The edge runs from `lo` to `hi` along `axis` (0=x, 1=y, 2=z), at fixed
+/// perpendicular coordinates `perp1` and `perp2`.
+///
+/// Returns a value in [0, 1].
+fn fraction_inside_along_axis(
+    lo: f64,
+    hi: f64,
+    perp1: f64,
+    perp2: f64,
+    axis: usize,
+    coords: &[[f64; 3]],
+    radii: &[f64],
+    probe_radius: f64,
+) -> f64 {
+    let edge_len = hi - lo;
+    if edge_len <= 0.0 {
+        return 0.0;
+    }
+
+    // Perpendicular axis indices
+    let (a1, a2) = match axis {
+        0 => (1, 2),
+        1 => (0, 2),
+        _ => (0, 1),
     };
 
+    // Collect intervals along the edge that lie inside each atom's sphere
+    let mut intervals: Vec<(f64, f64)> = Vec::new();
+
+    for (c, &r) in coords.iter().zip(radii.iter()) {
+        let big_r = r + probe_radius;
+        let d1 = perp1 - c[a1];
+        let d2 = perp2 - c[a2];
+        let rho_sq = big_r * big_r - d1 * d1 - d2 * d2;
+        if rho_sq <= 0.0 {
+            continue;
+        }
+        let rho = rho_sq.sqrt();
+        let ilo = (c[axis] - rho).max(lo);
+        let ihi = (c[axis] + rho).min(hi);
+        if ilo < ihi {
+            intervals.push((ilo, ihi));
+        }
+    }
+
+    if intervals.is_empty() {
+        return 0.0;
+    }
+
+    // Sort by interval start
+    intervals.sort_unstable_by(|a, b| a.0.partial_cmp(&b.0).unwrap());
+
+    // Merge overlapping intervals and sum total inside length
+    let mut total = 0.0;
+    let mut cur_lo = intervals[0].0;
+    let mut cur_hi = intervals[0].1;
+    for &(ilo, ihi) in &intervals[1..] {
+        if ilo <= cur_hi {
+            cur_hi = cur_hi.max(ihi);
+        } else {
+            total += cur_hi - cur_lo;
+            cur_lo = ilo;
+            cur_hi = ihi;
+        }
+    }
+    total += cur_hi - cur_lo;
+
+    (total / edge_len).clamp(0.0, 1.0)
+}
+
+/// Weighted harmonic mean of dielectrics.
+///
+/// `1/ε = f/ε_in + (1-f)/ε_out` where `f` is the fraction inside.
+/// This is the standard approach used by AMBER and other PB solvers
+/// for face-centered dielectrics at boundary edges.
+#[inline]
+fn weighted_harmonic_eps(f: f64, eps_in: f64, eps_out: f64) -> f64 {
+    if f <= 0.0 {
+        return eps_out;
+    }
+    if f >= 1.0 {
+        return eps_in;
+    }
+    1.0 / (f / eps_in + (1.0 - f) / eps_out)
+}
+
+/// Assign face-centered dielectric values based on the Solvent-Excluded Surface.
+///
+/// Uses the SES (molecular surface) to classify grid points as inside or
+/// outside, then assigns edge dielectrics:
+/// - Both endpoints inside → ε_in
+/// - Both endpoints outside → ε_out
+/// - Mixed → weighted harmonic mean using the fraction of the edge inside
+///   the molecular surface (computed analytically from sphere geometry).
+pub fn assign_dielectrics(
+    grid: &PbGrid,
+    coords: &[[f64; 3]],
+    radii: &[f64],
+    probe_radius: f64,
+    eps_in: f64,
+    eps_out: f64,
+) -> DielectricMaps {
+    let nx = grid.dims[0];
+    let ny = grid.dims[1];
+    let nz = grid.dims[2];
+    let n = nx * ny * nz;
+
+    // Use the Solvent-Excluded Surface for inside/outside classification.
+    // This handles re-entrant regions correctly, unlike the simpler SAS
+    // (union of inflated spheres) which over-estimates the cavity size.
+    let inside = compute_inside_ses(grid, coords, radii, probe_radius);
+
     let mut eps_x = vec![eps_out; n];
     let mut eps_y = vec![eps_out; n];
     let mut eps_z = vec![eps_out; n];
@@ -364,14 +612,70 @@ pub fn assign_dielectrics(
             for ix in 0..nx {
                 let idx = grid.index(ix, iy, iz);
                 let a = inside[idx];
+
+                // X-edge: (ix,iy,iz) → (ix+1,iy,iz)
                 if ix + 1 < nx {
-                    eps_x[idx] = eps_for(a, inside[grid.index(ix + 1, iy, iz)]);
+                    let b = inside[grid.index(ix + 1, iy, iz)];
+                    if a == b {
+                        eps_x[idx] = if a { eps_in } else { eps_out };
+                    } else {
+                        // Use VDW sphere geometry (probe=0) for the fraction
+                        // since the SES contact surface is at the VDW boundary.
+                        let pt = grid.point(ix, iy, iz);
+                        let f = fraction_inside_along_axis(
+                            pt[0],
+                            pt[0] + grid.spacing[0],
+                            pt[1],
+                            pt[2],
+                            0,
+                            coords,
+                            radii,
+                            0.0, // VDW radii, no probe inflation
+                        );
+                        eps_x[idx] = weighted_harmonic_eps(f, eps_in, eps_out);
+                    }
                 }
+
+                // Y-edge: (ix,iy,iz) → (ix,iy+1,iz)
                 if iy + 1 < ny {
-                    eps_y[idx] = eps_for(a, inside[grid.index(ix, iy + 1, iz)]);
+                    let b = inside[grid.index(ix, iy + 1, iz)];
+                    if a == b {
+                        eps_y[idx] = if a { eps_in } else { eps_out };
+                    } else {
+                        let pt = grid.point(ix, iy, iz);
+                        let f = fraction_inside_along_axis(
+                            pt[1],
+                            pt[1] + grid.spacing[1],
+                            pt[0],
+                            pt[2],
+                            1,
+                            coords,
+                            radii,
+                            0.0,
+                        );
+                        eps_y[idx] = weighted_harmonic_eps(f, eps_in, eps_out);
+                    }
                 }
+
+                // Z-edge: (ix,iy,iz) → (ix,iy,iz+1)
                 if iz + 1 < nz {
-                    eps_z[idx] = eps_for(a, inside[grid.index(ix, iy, iz + 1)]);
+                    let b = inside[grid.index(ix, iy, iz + 1)];
+                    if a == b {
+                        eps_z[idx] = if a { eps_in } else { eps_out };
+                    } else {
+                        let pt = grid.point(ix, iy, iz);
+                        let f = fraction_inside_along_axis(
+                            pt[2],
+                            pt[2] + grid.spacing[2],
+                            pt[0],
+                            pt[1],
+                            2,
+                            coords,
+                            radii,
+                            0.0,
+                        );
+                        eps_z[idx] = weighted_harmonic_eps(f, eps_in, eps_out);
+                    }
                 }
             }
         }
@@ -385,7 +689,13 @@ pub fn assign_dielectrics(
     }
 }
 
-/// Assign κ² map: zero inside ion-exclusion surface, κ²_bulk outside.
+/// Assign κ̄² map: zero inside ion-exclusion surface, ε_s·κ² outside.
+///
+/// The linearized PBE is: ∇·[ε∇φ] - κ̄²φ = -4π·ec·ρ
+/// where κ̄² = ε_s · κ² (the Debye-Hückel κ² multiplied by the solvent
+/// dielectric). The factor of ε_s arises because the mobile ion charge
+/// density is -2c₀·φ/(kBT) (no ε_s), while κ² = 8π·ec·c₀·N_A/(ε_s·kBT)
+/// already has 1/ε_s from the Debye-Hückel derivation.
 ///
 /// The ion-exclusion surface is the molecular surface inflated by the ion radius.
 pub fn assign_kappa(
@@ -394,12 +704,14 @@ pub fn assign_kappa(
     radii: &[f64],
     ion_radius: f64,
     kappa_bulk: f64,
+    solvent_dielectric: f64,
 ) -> Vec<f64> {
     let nx = grid.dims[0];
     let ny = grid.dims[1];
     let nz = grid.dims[2];
     let n = nx * ny * nz;
-    let kappa_sq = kappa_bulk * kappa_bulk;
+    // κ̄² = ε_s · κ² is the coefficient in the linearized PBE
+    let kappa_bar_sq = solvent_dielectric * kappa_bulk * kappa_bulk;
 
     (0..n)
         .into_par_iter()
@@ -412,7 +724,7 @@ pub fn assign_kappa(
             if point_inside_molecule(&pt, coords, radii, ion_radius) {
                 0.0
             } else {
-                kappa_sq
+                kappa_bar_sq
             }
         })
         .collect()

{rust_simulation_tools-0.2.2 → rust_simulation_tools-0.2.4}/crates/mmpbsa/src/pb_solver.rs

@@ -10,12 +10,15 @@ use crate::pb_grid::{DielectricMaps, PbGrid};
 use rayon::prelude::*;
 
 /// Boundary condition for the PB solve.
-#[derive(Debug, Clone, Copy)]
+#[derive(Debug, Clone)]
 pub enum BoundaryCondition {
     /// phi = 0 on all boundaries.
     Zero,
     /// Debye-Huckel single-sphere approximation on boundaries.
     DebyeHuckel,
+    /// Interpolated from a coarse grid solution (for focusing).
+    /// Contains the precomputed boundary values for each boundary grid point.
+    Interpolated(Vec<f64>),
 }
 
 /// Set Debye-Huckel boundary conditions on the grid.
@@ -65,6 +68,59 @@ fn set_dh_boundary(
     }
 }
 
+/// Compute interpolated boundary conditions from a coarse grid solution.
+///
+/// For each boundary point of `fine_grid`, looks up the corresponding
+/// position in the coarse solution via trilinear interpolation.
+pub fn interpolated_boundary(
+    fine_grid: &PbGrid,
+    coarse_grid: &PbGrid,
+    coarse_potential: &[f64],
+) -> Vec<f64> {
+    let nx = fine_grid.dims[0];
+    let ny = fine_grid.dims[1];
+    let nz = fine_grid.dims[2];
+    let n = nx * ny * nz;
+    let mut boundary = vec![0.0f64; n];
+
+    for iz in 0..nz {
+        for iy in 0..ny {
+            for ix in 0..nx {
+                let is_boundary =
+                    ix == 0 || ix == nx - 1 || iy == 0 || iy == ny - 1 || iz == 0 || iz == nz - 1;
+                if !is_boundary {
+                    continue;
+                }
+                let pt = fine_grid.point(ix, iy, iz);
+                let idx = fine_grid.index(ix, iy, iz);
+                boundary[idx] = coarse_grid.interpolate_with_data(&pt, coarse_potential);
+            }
+        }
+    }
+    boundary
+}
+
+/// Set boundary conditions from precomputed interpolated values.
+fn set_interpolated_boundary(potential: &mut [f64], grid: &PbGrid, boundary: &[f64]) {
+    let nx = grid.dims[0];
+    let ny = grid.dims[1];
+    let nz = grid.dims[2];
+
+    for iz in 0..nz {
+        for iy in 0..ny {
+            for ix in 0..nx {
+                let is_boundary =
+                    ix == 0 || ix == nx - 1 || iy == 0 || iy == ny - 1 || iz == 0 || iz == nz - 1;
+                if !is_boundary {
+                    continue;
+                }
+                let idx = grid.index(ix, iy, iz);
+                potential[idx] = boundary[idx];
+            }
+        }
+    }
+}
+
 /// Estimate the optimal SOR relaxation parameter from grid dimensions.
 fn estimate_omega(dims: &[usize; 3]) -> f64 {
     // Spectral radius of Jacobi iteration for 3D Laplacian
@@ -264,11 +320,14 @@ pub fn solve_lpbe(
     let mut potential = vec![0.0f64; n];
 
     // Set boundary conditions
-    match boundary {
+    match &boundary {
         BoundaryCondition::Zero => {}
         BoundaryCondition::DebyeHuckel => {
             set_dh_boundary(&mut potential, grid, coords, charges, kappa_bulk, eps_out);
         }
+        BoundaryCondition::Interpolated(bvals) => {
+            set_interpolated_boundary(&mut potential, grid, bvals);
+        }
     }
 
     let omega = estimate_omega(&grid.dims);
@@ -703,11 +762,14 @@ pub fn solve_lpbe_multigrid(
     let mut potential = vec![0.0f64; n];
 
     // Set boundary conditions
-    match boundary {
+    match &boundary {
         BoundaryCondition::Zero => {}
         BoundaryCondition::DebyeHuckel => {
             set_dh_boundary(&mut potential, grid, coords, charges, kappa_bulk, eps_out);
         }
+        BoundaryCondition::Interpolated(bvals) => {
+            set_interpolated_boundary(&mut potential, grid, bvals);
+        }
     }
 
     let min_dim = grid.dims.iter().copied().min().unwrap_or(0);

{rust_simulation_tools-0.2.2 → rust_simulation_tools-0.2.4}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "maturin"
 
 [project]
 name = "rust_simulation_tools"
-version = "0.2.2"
+version = "0.2.4"
 description = "Fast algorithms for MD trajectories"
 readme = "README.md"
 authors = [