use crate::constants::codata2018::EV_ANGSTROM_FACTOR;
use crate::parameters::ParameterModel;
use crate::types::{AlignedMatrix, BasisType, MolecularBatch};
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct ExternalCharge {
pub x: f64,
pub y: f64,
pub z: f64,
pub charge: f64,
}
impl ExternalCharge {
pub fn new(x: f64, y: f64, z: f64, charge: f64) -> Self {
Self { x, y, z, charge }
}
}
#[inline(always)]
pub fn klopman_radius(gss: f64) -> f64 {
if gss > 1.0e-6 {
EV_ANGSTROM_FACTOR / (2.0 * gss)
} else {
0.50
}
}
pub fn apply_external_charges_to_hcore(
batch: &MolecularBatch,
model: &dyn ParameterModel,
external_charges: &[ExternalCharge],
h_core: &mut AlignedMatrix<f64>,
) {
if external_charges.is_empty() {
return;
}
let natoms = batch.natoms;
for i in 0..natoms {
let za = batch.atomic_numbers[i];
let p_a = match model.get_element(za) {
Some(p) => p,
None => continue,
};
let rho = klopman_radius(p_a.gss);
let rho_sq = rho * rho;
let xi = batch.x[i];
let yi = batch.y[i];
let zi = batch.z[i];
let mut v_ext = 0.0;
for ext in external_charges {
let dx = xi - ext.x;
let dy = yi - ext.y;
let dz = zi - ext.z;
let r2 = dx * dx + dy * dy + dz * dz;
let denom = (r2 + rho_sq).sqrt();
v_ext -= ext.charge * EV_ANGSTROM_FACTOR / denom;
}
let orb_start = batch.orbital_offsets[i];
let norbs_atom = match batch.basis_types[i] {
BasisType::S => 1,
BasisType::SP => 4,
BasisType::SPD => 9,
};
for mu in 0..norbs_atom {
let idx = orb_start + mu;
let old_val = h_core.get(idx, idx);
h_core.set(idx, idx, old_val + v_ext);
}
}
}
pub fn compute_external_charges_core_energy(
batch: &MolecularBatch,
model: &dyn ParameterModel,
external_charges: &[ExternalCharge],
) -> f64 {
if external_charges.is_empty() {
return 0.0;
}
let natoms = batch.natoms;
let mut e_core_ext = 0.0;
for i in 0..natoms {
let za = batch.atomic_numbers[i];
let p_a = match model.get_element(za) {
Some(p) => p,
None => continue,
};
let core_charge = p_a.core_charge;
let xi = batch.x[i];
let yi = batch.y[i];
let zi = batch.z[i];
for ext in external_charges {
let dx = xi - ext.x;
let dy = yi - ext.y;
let dz = zi - ext.z;
let r = (dx * dx + dy * dy + dz * dz).sqrt();
if r > 1.0e-8 {
e_core_ext += core_charge * ext.charge * EV_ANGSTROM_FACTOR / r;
}
}
}
e_core_ext
}
#[allow(clippy::needless_range_loop)]
pub fn compute_external_charges_gradients(
batch: &MolecularBatch,
density: &AlignedMatrix<f64>,
model: &dyn ParameterModel,
external_charges: &[ExternalCharge],
gradients: &mut [[f64; 3]],
) {
if external_charges.is_empty() {
return;
}
let natoms = batch.natoms;
for i in 0..natoms {
let za = batch.atomic_numbers[i];
let p_a = match model.get_element(za) {
Some(p) => p,
None => continue,
};
let core_charge = p_a.core_charge;
let rho = klopman_radius(p_a.gss);
let rho_sq = rho * rho;
let orb_start = batch.orbital_offsets[i];
let norbs_atom = match batch.basis_types[i] {
BasisType::S => 1,
BasisType::SP => 4,
BasisType::SPD => 9,
};
let mut pop_a = 0.0;
for mu in 0..norbs_atom {
let idx = orb_start + mu;
pop_a += density.get(idx, idx);
}
let xi = batch.x[i];
let yi = batch.y[i];
let zi = batch.z[i];
for ext in external_charges {
let dx = xi - ext.x;
let dy = yi - ext.y;
let dz = zi - ext.z;
let r2 = dx * dx + dy * dy + dz * dz;
let r = r2.sqrt();
if r < 1.0e-8 {
continue;
}
let denom_core = r2 * r; let denom_elec = (r2 + rho_sq).powf(1.5);
let factor =
ext.charge * EV_ANGSTROM_FACTOR * (-core_charge / denom_core + pop_a / denom_elec);
gradients[i][0] += factor * dx;
gradients[i][1] += factor * dy;
gradients[i][2] += factor * dz;
}
}
}