use std::{
fs,
fs::File,
io::{self, ErrorKind, Read, Seek, SeekFrom, Write},
path::Path,
process::Command,
};
use bio_apis::rcsb;
use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
use lin_alg::f64::{Mat3, Vec3};
const HEADER_SIZE: u64 = 1_024;
#[derive(Clone, Debug)]
pub struct DensityHeaderInner {
pub cell: UnitCell,
pub mapc: i32, pub mapr: i32,
pub maps: i32,
pub mx: i32,
pub my: i32,
pub mz: i32,
pub nxstart: i32,
pub nystart: i32,
pub nzstart: i32,
pub ispg: i32,
pub nsymbt: i32,
pub version: i32,
pub xorigin: Option<f32>,
pub yorigin: Option<f32>,
pub zorigin: Option<f32>, }
#[allow(unused)]
#[derive(Clone, Debug)]
pub struct MapHeader {
pub inner: DensityHeaderInner,
pub nx: i32,
pub ny: i32,
pub nz: i32,
pub mode: i32, pub dmin: f32,
pub dmax: f32,
pub dmean: f32,
}
fn read_map_header<R: Read + Seek>(mut r: R) -> io::Result<MapHeader> {
r.seek(SeekFrom::Start(0))?;
let nx = r.read_i32::<LittleEndian>()?;
let ny = r.read_i32::<LittleEndian>()?;
let nz = r.read_i32::<LittleEndian>()?;
let mode = r.read_i32::<LittleEndian>()?;
let nxstart = r.read_i32::<LittleEndian>()?;
let nystart = r.read_i32::<LittleEndian>()?;
let nzstart = r.read_i32::<LittleEndian>()?;
let mx = r.read_i32::<LittleEndian>()?;
let my = r.read_i32::<LittleEndian>()?;
let mz = r.read_i32::<LittleEndian>()?;
let mut cell = [0f32; 6];
for c in &mut cell {
*c = r.read_f32::<LittleEndian>()?;
}
let mapc = r.read_i32::<LittleEndian>()?;
let mapr = r.read_i32::<LittleEndian>()?;
let maps = r.read_i32::<LittleEndian>()?;
let dmin = r.read_f32::<LittleEndian>()?;
let dmax = r.read_f32::<LittleEndian>()?;
let dmean = r.read_f32::<LittleEndian>()?;
let ispg = r.read_i32::<LittleEndian>()?;
let nsymbt = r.read_i32::<LittleEndian>()?;
r.seek(SeekFrom::Start(27 * 4))?;
let version = r.read_i32::<LittleEndian>()?;
r.seek(SeekFrom::Start(49 * 4))?;
let xorigin_ = r.read_f32::<LittleEndian>()?;
let yorigin_ = r.read_f32::<LittleEndian>()?;
let zorigin_ = r.read_f32::<LittleEndian>()?;
let mut xorigin = None;
let mut yorigin = None;
let mut zorigin = None;
let mut tag = [0u8; 4];
r.seek(SeekFrom::Start(52 * 4))?;
r.read_exact(&mut tag)?;
if &tag != b"MAP " {
return Err(io::Error::new(
ErrorKind::InvalidData,
"Invalid MAP tag in header.",
));
}
const EPS: f32 = 0.0001;
if xorigin_.abs() > EPS {
xorigin = Some(xorigin_);
}
if yorigin_.abs() > EPS {
yorigin = Some(yorigin_);
}
if zorigin_.abs() > EPS {
zorigin = Some(zorigin_);
}
r.seek(SeekFrom::Start(HEADER_SIZE))?;
let cell = UnitCell::new(
cell[0] as f64,
cell[1] as f64,
cell[2] as f64,
cell[3] as f64,
cell[4] as f64,
cell[5] as f64,
);
let inner = DensityHeaderInner {
nxstart,
nystart,
nzstart,
mx,
my,
mz,
cell,
mapc,
mapr,
maps,
ispg,
nsymbt,
version,
xorigin,
yorigin,
zorigin,
};
Ok(MapHeader {
inner,
nx,
ny,
nz,
mode,
dmin,
dmax,
dmean,
})
}
#[derive(Clone, Debug)]
pub struct UnitCell {
pub a: f64,
pub b: f64,
pub c: f64,
pub alpha: f64,
pub beta: f64,
pub gamma: f64,
pub ortho: Mat3,
pub ortho_inv: Mat3,
}
impl UnitCell {
pub fn new(a: f64, b: f64, c: f64, alpha_deg: f64, beta_deg: f64, gamma_deg: f64) -> Self {
let (α, β, γ) = (
alpha_deg.to_radians(),
beta_deg.to_radians(),
gamma_deg.to_radians(),
);
let v_a = Vec3::new(a, 0.0, 0.0);
let v_b = Vec3::new(b * γ.cos(), b * γ.sin(), 0.0);
let cx = c * β.cos();
let cy = c * (α.cos() - β.cos() * γ.cos()) / γ.sin();
let cz = c * (1.0 - β.cos().powi(2) - cy.powi(2) / c.powi(2)).sqrt();
let v_c = Vec3::new(cx, cy, cz);
let ortho = Mat3::from_cols(v_a, v_b, v_c);
let ortho_inv = ortho.inverse().expect("unit-cell matrix is singular");
Self {
a,
b,
c,
alpha: α,
beta: β,
gamma: γ,
ortho,
ortho_inv,
}
}
pub fn fractional_to_cartesian(&self, f: Vec3) -> Vec3 {
self.ortho.clone() * f
}
pub fn cartesian_to_fractional(&self, c: Vec3) -> Vec3 {
self.ortho_inv.clone() * c
}
}
fn read_header_dens<R: Read + Seek>(data: &mut R) -> io::Result<(MapHeader, Vec<f32>)> {
let hdr = read_map_header(&mut *data)?;
if hdr.mode != 2 {
return Err(io::Error::new(
ErrorKind::InvalidData,
format!("Unsupported mode: {}", hdr.mode),
));
}
data.seek(SeekFrom::Start(HEADER_SIZE + hdr.inner.nsymbt as u64))?;
let n = (hdr.nx * hdr.ny * hdr.nz) as usize;
let mut dens = Vec::with_capacity(n);
for _ in 0..n {
dens.push(data.read_f32::<LittleEndian>()?);
}
Ok((hdr, dens))
}
pub(crate) fn get_origin_frac(hdr: &MapHeader, cell: &UnitCell) -> Vec3 {
if let (Some(ox), Some(oy), Some(oz)) =
(hdr.inner.xorigin, hdr.inner.yorigin, hdr.inner.zorigin)
{
cell.cartesian_to_fractional(Vec3::new(ox as f64, oy as f64, oz as f64))
} else {
Vec3::new(
hdr.inner.nxstart as f64 / hdr.inner.mx as f64,
hdr.inner.nystart as f64 / hdr.inner.my as f64,
hdr.inner.nzstart as f64 / hdr.inner.mz as f64,
)
}
}
#[derive(Clone, Debug)]
pub struct DensityMap {
pub hdr: MapHeader,
pub origin_frac: Vec3,
pub perm_f2c: [usize; 3],
pub perm_c2f: [usize; 3],
pub data: Vec<f32>,
pub(crate) mean: f32,
pub(crate) inv_sigma: f32,
}
impl DensityMap {
pub fn new(hdr: MapHeader, data: Vec<f32>) -> io::Result<Self> {
let perm_f2c = [
hdr.inner.mapc as usize - 1,
hdr.inner.mapr as usize - 1,
hdr.inner.maps as usize - 1,
];
let mut perm_c2f = [0; 3];
for (f, c) in perm_f2c.iter().enumerate() {
perm_c2f[*c] = f;
}
let origin_frac = get_origin_frac(&hdr, &hdr.inner.cell);
let n = data.len() as f32;
let mut mean = 0.;
for val in &data {
mean += val;
}
mean /= data.len() as f32;
println!("Map means. Hdr: {}, calculated: {}", hdr.dmean, mean);
let variance: f32 = data.iter().map(|v| (*v - mean).powi(2)).sum::<f32>() / n;
let sigma = variance.sqrt().max(1e-6); let inv_sigma = 1. / sigma;
Ok(Self {
hdr,
origin_frac,
perm_f2c,
perm_c2f,
data,
mean,
inv_sigma,
})
}
pub fn open<R: Read + Seek>(data: &mut R) -> io::Result<Self> {
let (hdr, data) = read_header_dens(data)?;
Self::new(hdr, data)
}
pub fn density_at_point(&self, cart: Vec3) -> f32 {
let mut frac = self.hdr.inner.cell.cartesian_to_fractional(cart);
frac.x -= frac.x.floor();
frac.y -= frac.y.floor();
frac.z -= frac.z.floor();
let ic = [
(frac.x * self.hdr.inner.mx as f64 - 0.5).round() as isize,
(frac.y * self.hdr.inner.my as f64 - 0.5).round() as isize,
(frac.z * self.hdr.inner.mz as f64 - 0.5).round() as isize,
];
let ifile = [
pmod(ic[self.perm_c2f[0]], self.hdr.nx as usize),
pmod(ic[self.perm_c2f[1]], self.hdr.ny as usize),
pmod(ic[self.perm_c2f[2]], self.hdr.nz as usize),
];
let offset = (ifile[2] * self.hdr.ny as usize + ifile[1]) * self.hdr.nx as usize + ifile[0];
self.data[offset]
}
pub fn density_at_point_trilinear(&self, cart: Vec3) -> f32 {
let mut frac = self.hdr.inner.cell.cartesian_to_fractional(cart);
frac.x -= frac.x.floor();
frac.y -= frac.y.floor();
frac.z -= frac.z.floor();
let gx = frac.x * self.hdr.inner.mx as f64 - 0.5;
let gy = frac.y * self.hdr.inner.my as f64 - 0.5;
let gz = frac.z * self.hdr.inner.mz as f64 - 0.5;
let ix0 = gx.floor() as isize;
let iy0 = gy.floor() as isize;
let iz0 = gz.floor() as isize;
let dx = (gx - ix0 as f64) as f32; let dy = (gy - iy0 as f64) as f32;
let dz = (gz - iz0 as f64) as f32;
let wx = [1.0 - dx, dx];
let wy = [1.0 - dy, dy];
let wz = [1.0 - dz, dz];
let mut rho = 0.;
for (cz, w_z) in [iz0, iz0 + 1].iter().zip(wz) {
let fz = pmod(*cz, self.hdr.nz as usize);
for (cy, w_y) in [iy0, iy0 + 1].iter().zip(wy) {
let fy = pmod(*cy, self.hdr.ny as usize);
for (cx, w_x) in [ix0, ix0 + 1].iter().zip(wx) {
let fx = pmod(*cx, self.hdr.nx as usize);
let ifile = [fx, fy, fz];
let voxel = [
ifile[self.perm_c2f[0]],
ifile[self.perm_c2f[1]],
ifile[self.perm_c2f[2]],
];
let offset = (voxel[2] * self.hdr.ny as usize + voxel[1])
* self.hdr.nx as usize
+ voxel[0];
rho += w_x * w_y * w_z * self.data[offset];
}
}
}
rho
}
pub fn density_to_sig(&self, val: f32) -> f32 {
(val - self.mean) * self.inv_sigma
}
pub fn load(path: &Path) -> io::Result<Self> {
let mut file = File::open(path)?;
Self::open(&mut file)
}
pub fn save(&self, path: &Path) -> io::Result<()> {
let len_expected = (self.hdr.nx as usize)
.saturating_mul(self.hdr.ny as usize)
.saturating_mul(self.hdr.nz as usize);
if self.data.len() != len_expected {
return Err(io::Error::new(
ErrorKind::InvalidData,
format!(
"data length ({}) does not match nx*ny*nz ({}).",
self.data.len(),
len_expected
),
));
}
let n = self.data.len() as f32;
let (mut dmin, mut dmax) = (f32::INFINITY, f32::NEG_INFINITY);
let mut sum = 0.0f32;
for &v in &self.data {
if v < dmin {
dmin = v;
}
if v > dmax {
dmax = v;
}
sum += v;
}
let dmean = if n > 0.0 { sum / n } else { 0.0 };
let mut var = 0.0f32;
for &v in &self.data {
let dv = v - dmean;
var += dv * dv;
}
let rms = if n > 0.0 { (var / n).sqrt() } else { 0.0 };
let mut hdr_buf = Vec::with_capacity(HEADER_SIZE as usize);
{
hdr_buf.write_i32::<LittleEndian>(self.hdr.nx)?; hdr_buf.write_i32::<LittleEndian>(self.hdr.ny)?; hdr_buf.write_i32::<LittleEndian>(self.hdr.nz)?; hdr_buf.write_i32::<LittleEndian>(2)?; hdr_buf.write_i32::<LittleEndian>(self.hdr.inner.nxstart)?; hdr_buf.write_i32::<LittleEndian>(self.hdr.inner.nystart)?; hdr_buf.write_i32::<LittleEndian>(self.hdr.inner.nzstart)?; hdr_buf.write_i32::<LittleEndian>(self.hdr.inner.mx)?; hdr_buf.write_i32::<LittleEndian>(self.hdr.inner.my)?; hdr_buf.write_i32::<LittleEndian>(self.hdr.inner.mz)?;
hdr_buf.write_f32::<LittleEndian>(self.hdr.inner.cell.a as f32)?;
hdr_buf.write_f32::<LittleEndian>(self.hdr.inner.cell.b as f32)?;
hdr_buf.write_f32::<LittleEndian>(self.hdr.inner.cell.c as f32)?;
hdr_buf.write_f32::<LittleEndian>(self.hdr.inner.cell.alpha as f32)?;
hdr_buf.write_f32::<LittleEndian>(self.hdr.inner.cell.beta as f32)?;
hdr_buf.write_f32::<LittleEndian>(self.hdr.inner.cell.gamma as f32)?;
hdr_buf.write_i32::<LittleEndian>(self.hdr.inner.mapc)?; hdr_buf.write_i32::<LittleEndian>(self.hdr.inner.mapr)?; hdr_buf.write_i32::<LittleEndian>(self.hdr.inner.maps)?; hdr_buf.write_f32::<LittleEndian>(dmin)?; hdr_buf.write_f32::<LittleEndian>(dmax)?; hdr_buf.write_f32::<LittleEndian>(dmean)?; hdr_buf.write_i32::<LittleEndian>(self.hdr.inner.ispg)?;
let nsymbt: i32 = 0;
hdr_buf.write_i32::<LittleEndian>(nsymbt)?;
for w in 25..=49 {
if w == 28 {
hdr_buf.write_i32::<LittleEndian>(if self.hdr.inner.version != 0 {
self.hdr.inner.version
} else {
20140
})?;
} else {
hdr_buf.write_i32::<LittleEndian>(0)?;
}
}
let xorig = self.hdr.inner.xorigin.unwrap_or(0.0);
let yorig = self.hdr.inner.yorigin.unwrap_or(0.0);
let zorig = self.hdr.inner.zorigin.unwrap_or(0.0);
hdr_buf.write_f32::<LittleEndian>(xorig)?; hdr_buf.write_f32::<LittleEndian>(yorig)?; hdr_buf.write_f32::<LittleEndian>(zorig)?;
hdr_buf.extend_from_slice(b"MAP ");
hdr_buf.extend_from_slice(&[0x44, 0x41, 0x00, 0x00]);
hdr_buf.write_f32::<LittleEndian>(rms)?;
let nlabel: i32 = 0;
hdr_buf.write_i32::<LittleEndian>(nlabel)?;
hdr_buf.resize(HEADER_SIZE as usize, 0);
}
let mut f = File::create(path)?;
f.write_all(&hdr_buf)?;
for &v in &self.data {
f.write_f32::<LittleEndian>(v)?;
}
Ok(())
}
pub fn load_sf_or_mtz(path: &Path, gemmi_path: Option<&Path>) -> io::Result<Self> {
let temp_file = "temp_map.map";
let program = match gemmi_path {
Some(p) => p.join("gemmi").to_str().unwrap().to_string(),
None => "gemmi".to_string(),
};
let out = Command::new(program)
.args(["sf2map", path.to_str().unwrap(), temp_file])
.output()?;
if !out.status.success() {
let stderr_str = String::from_utf8_lossy(&out.stderr);
return Err(io::Error::other(format!(
"Problem parsing SF or MTZ file: {}",
stderr_str
)));
}
let map = Self::load(Path::new(temp_file))?;
fs::remove_file(Path::new(temp_file))?;
Ok(map)
}
pub fn save_sf_or_mtz(&self, path: &Path, gemmi_path: Option<&Path>) -> io::Result<()> {
let ext = path
.extension()
.and_then(|s| s.to_str())
.map(|s| s.to_ascii_lowercase())
.ok_or_else(|| io::Error::other("Output path must have a file extension"))?;
if ext != "mtz" && ext != "cif" && ext != "mmcif" && ext != "sf" {
return Err(io::Error::other(format!(
"Unsupported output extension: .{ext} (expected .mtz, .cif, .mmcif, or .sf)"
)));
}
let temp_map = "temp_map.map";
self.save(Path::new(temp_map))?;
let program = match gemmi_path {
Some(p) => p.join("gemmi").to_str().unwrap().to_string(),
None => "gemmi".to_string(),
};
let out = Command::new(program)
.args(["map2sf", temp_map, path.to_str().unwrap(), "FWT", "PHWT"])
.output()?;
let _ = fs::remove_file(Path::new(temp_map));
if !out.status.success() {
let stderr_str = String::from_utf8_lossy(&out.stderr);
return Err(io::Error::other(format!(
"Problem writing SF or MTZ file via gemmi map2sf: {}",
stderr_str
)));
}
Ok(())
}
}
pub fn density_from_2fo_fc_rcsb_gemmi(
ident: &str,
gemmi_path: Option<&Path>,
) -> io::Result<DensityMap> {
println!("Downloading Map data for {ident}...");
let map_2fo_fc = rcsb::load_validation_2fo_fc_cif(ident)
.map_err(|_| io::Error::new(ErrorKind::InvalidData, "Problem loading 2fo-fc from RCSB"))?;
let path = Path::new("temp_map.cif");
fs::write(path, map_2fo_fc)?;
let result = DensityMap::load_sf_or_mtz(path, gemmi_path)?;
fs::remove_file(path)?;
Ok(result)
}
fn pmod(i: isize, n: usize) -> usize {
((i % n as isize) + n as isize) as usize % n
}