use std::fs;
use std::path::Path;
use std::time::Duration;
use serde::Deserialize;
const API_ROOT: &str = "https://api.esmatlas.com";
#[derive(Clone, Debug, Deserialize)]
pub struct AtlasConfidence {
pub pae: Vec<Vec<f32>>,
pub ptm: f32,
pub plddt: Vec<f32>,
}
impl AtlasConfidence {
pub fn size(&self) -> usize {
self.pae.len()
}
pub fn max_pae(&self) -> f32 {
self.pae.iter().flatten().copied().fold(0.0_f32, f32::max)
}
fn validate(self, id: &str) -> Result<Self, String> {
if self.pae.is_empty() || self.pae.iter().any(Vec::is_empty) {
return Err(format!(
"ESM Atlas confidence data for {id} has an empty PAE matrix"
));
}
let width = self.pae[0].len();
if self.pae.iter().any(|row| row.len() != width) {
return Err(format!(
"ESM Atlas confidence data for {id} has inconsistent PAE rows"
));
}
if !self.ptm.is_finite() || self.pae.iter().flatten().any(|value| !value.is_finite()) {
return Err(format!(
"ESM Atlas confidence data for {id} contains non-finite values"
));
}
Ok(self)
}
}
pub struct AtlasClient {
client: reqwest::blocking::Client,
}
impl AtlasClient {
pub fn new() -> Result<Self, String> {
let client = reqwest::blocking::Client::builder()
.user_agent(concat!("molx/", env!("CARGO_PKG_VERSION")))
.connect_timeout(Duration::from_secs(20))
.timeout(Duration::from_secs(300))
.build()
.map_err(|error| format!("cannot initialize HTTP client: {error}"))?;
Ok(Self { client })
}
pub fn fetch_structure(&self, id: &str) -> Result<String, String> {
let id = validate_atlas_id(id)?;
let url = format!("{API_ROOT}/fetchPredictedStructure/{id}");
let response = self
.client
.get(&url)
.send()
.map_err(|error| format!("ESM Atlas request failed: {error}"))?;
response_to_pdb(response, &format!("ESM Atlas ID {id}"))
}
pub fn fetch_confidence(&self, id: &str) -> Result<AtlasConfidence, String> {
let id = validate_atlas_id(id)?;
let url = format!("{API_ROOT}/fetchConfidencePrediction/{id}");
let response = self
.client
.get(&url)
.send()
.map_err(|error| format!("ESM Atlas confidence request failed: {error}"))?;
let status = response.status();
let text = response
.text()
.map_err(|error| format!("cannot read ESM Atlas confidence response: {error}"))?;
if !status.is_success() {
let detail = text.trim();
return Err(if detail.is_empty() {
format!("ESM Atlas confidence request failed with HTTP {status}")
} else {
format!("ESM Atlas confidence request failed with HTTP {status}: {detail}")
});
}
serde_json::from_str::<AtlasConfidence>(&text)
.map_err(|error| format!("invalid ESM Atlas confidence JSON for {id}: {error}"))?
.validate(&id)
}
pub fn fold_sequence(&self, sequence: &str) -> Result<String, String> {
let sequence = normalize_sequence(sequence)?;
let url = format!("{API_ROOT}/foldSequence/v1/pdb/");
let response = self
.client
.post(&url)
.header(reqwest::header::CONTENT_TYPE, "text/plain")
.body(sequence)
.send()
.map_err(|error| format!("ESMFold request failed: {error}"))?;
response_to_pdb(response, "ESMFold prediction")
}
}
fn response_to_pdb(
response: reqwest::blocking::Response,
operation: &str,
) -> Result<String, String> {
let status = response.status();
let text = response
.text()
.map_err(|error| format!("cannot read {operation} response: {error}"))?;
if !status.is_success() {
let detail = text.trim();
return Err(if detail.is_empty() {
format!("{operation} failed with HTTP {status}")
} else {
format!("{operation} failed with HTTP {status}: {detail}")
});
}
if !text.lines().any(|line| {
let record = line.get(..6).unwrap_or(line).trim();
record == "ATOM" || record == "HETATM"
}) {
return Err(format!(
"{operation} returned a successful response, but it was not PDB data"
));
}
Ok(text)
}
pub fn validate_atlas_id(input: &str) -> Result<String, String> {
let id = input.trim().to_ascii_uppercase();
let suffix = id.strip_prefix("MGYP").ok_or_else(|| {
format!("invalid ESM Atlas ID {input:?}: identifiers must start with MGYP")
})?;
if suffix.is_empty() || !suffix.chars().all(|character| character.is_ascii_digit()) {
return Err(format!(
"invalid ESM Atlas ID {input:?}: MGYP must be followed by digits"
));
}
Ok(id)
}
pub fn normalize_sequence(input: &str) -> Result<String, String> {
let mut sequence = String::new();
for line in input.lines() {
let trimmed = line.trim();
if trimmed.starts_with('>') {
continue;
}
for character in trimmed.chars() {
if character.is_ascii_whitespace() || character.is_ascii_digit() {
continue;
}
let amino_acid = character.to_ascii_uppercase();
if !matches!(
amino_acid,
'A' | 'R'
| 'N'
| 'D'
| 'C'
| 'E'
| 'Q'
| 'G'
| 'H'
| 'I'
| 'L'
| 'K'
| 'M'
| 'F'
| 'P'
| 'S'
| 'T'
| 'W'
| 'Y'
| 'V'
| 'B'
| 'J'
| 'O'
| 'U'
| 'X'
| 'Z'
) {
return Err(format!(
"invalid amino-acid character {character:?} in sequence"
));
}
sequence.push(amino_acid);
}
}
if sequence.len() < 4 {
return Err("protein sequence must contain at least 4 residues".to_owned());
}
if sequence.len() > 1_024 {
return Err(format!(
"sequence has {} residues; the public ESMFold endpoint is limited to small interactive jobs (maximum 1024 in MolX)",
sequence.len()
));
}
Ok(sequence)
}
pub fn read_sequence_file(path: impl AsRef<Path>) -> Result<String, String> {
let path = path.as_ref();
let text = fs::read_to_string(path)
.map_err(|error| format!("cannot read sequence file {}: {error}", path.display()))?;
normalize_sequence(&text)
}
pub fn write_pdb(path: impl AsRef<Path>, pdb: &str, force: bool) -> Result<(), String> {
let path = path.as_ref();
if path.exists() && !force {
return Err(format!(
"{} already exists; pass --force to replace it",
path.display()
));
}
fs::write(path, pdb).map_err(|error| format!("cannot write {}: {error}", path.display()))
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn validates_and_normalizes_atlas_id() {
assert_eq!(
validate_atlas_id("mgyp002537940442").unwrap(),
"MGYP002537940442"
);
assert!(validate_atlas_id("P12345").is_err());
assert!(validate_atlas_id("MGYP12X").is_err());
}
#[test]
fn cleans_fasta_sequence() {
assert_eq!(
normalize_sequence(">protein one\nMKT 12\nAYX\n").unwrap(),
"MKTAYX"
);
assert!(normalize_sequence("MKT*").is_err());
assert!(normalize_sequence("MK").is_err());
}
#[test]
fn parses_confidence_response() {
let confidence: AtlasConfidence =
serde_json::from_str(r#"{"pae":[[0,2],[3,0]],"ptm":0.75,"plddt":[0.9,0.8]}"#).unwrap();
let confidence = confidence.validate("MGYP1").unwrap();
assert_eq!(confidence.size(), 2);
assert_eq!(confidence.max_pae(), 3.0);
assert_eq!(confidence.plddt.len(), 2);
}
}