molgfx-geometry 0.3.5

// The index-to-float casts below build fixture coordinates from a loop
// counter whose range is a fixed literal in this file; nothing truncates.

use super::*;
use molgfx_core::{
    AtomProperty, AtomPropertyMeaning, ColorScheme, PropertyAppearance, RepresentationKind,
    ScalarFieldSemantics, Scene, SurfaceKind,
};
use molgfx_math::Rgba8;
use std::sync::Arc;

fn scene_table() -> (Scene, molgfx_core::RepresentationHandle) {
    let structure = fixture_structure();
    let mut scene = match Scene::from_structure(&structure) {
        Ok(scene) => scene,
        Err(e) => panic!("fixture scene builds: {e}"),
    };
    let sel = scene.add_selection(AtomSelection::All);
    let rep = match scene.represent(sel, RepresentationKind::Spacefill) {
        Ok(rep) => rep,
        Err(e) => panic!("spacefill applies: {e}"),
    };
    (scene, rep)
}

pub(crate) fn fixture_structure() -> molframe::Structure {
    let cif = "\
data_test
loop_
_atom_site.group_PDB
_atom_site.id
_atom_site.type_symbol
_atom_site.label_atom_id
_atom_site.label_alt_id
_atom_site.label_comp_id
_atom_site.label_asym_id
_atom_site.label_entity_id
_atom_site.label_seq_id
_atom_site.Cartn_x
_atom_site.Cartn_y
_atom_site.Cartn_z
_atom_site.occupancy
_atom_site.B_iso_or_equiv
_atom_site.auth_seq_id
_atom_site.auth_asym_id
_atom_site.pdbx_PDB_model_num
ATOM 1 N N  . GLY A 1 1 0.0 0.0 0.0 1.00 10.0 1 A 1
ATOM 2 C CA . GLY A 1 1 1.5 0.0 0.0 1.00 10.0 1 A 1
ATOM 3 O O  . GLY A 1 1 3.0 1.0 0.0 1.00 10.0 1 A 1
loop_
_struct_conn.id
_struct_conn.conn_type_id
_struct_conn.ptnr1_label_asym_id
_struct_conn.ptnr1_label_seq_id
_struct_conn.ptnr1_label_comp_id
_struct_conn.ptnr1_label_atom_id
_struct_conn.ptnr2_label_asym_id
_struct_conn.ptnr2_label_seq_id
_struct_conn.ptnr2_label_comp_id
_struct_conn.ptnr2_label_atom_id
_struct_conn.pdbx_value_order
1 covale A 1 GLY N A 1 GLY CA SING
2 covale A 1 GLY CA A 1 GLY O AROM
";
    let options = molframe::ReadOptions::new();
    match molframe::read_bytes(cif.as_bytes().to_vec(), Some("t.cif"), &options) {
        Ok((structure, _)) => structure,
        Err(diagnostics) => panic!("fixture parses: {diagnostics:?}"),
    }
}

#[test]
fn packing_all_atoms_produces_one_record_per_atom_in_order() {
    let (scene, rep_handle) = scene_table();
    let Some(table) = scene.first_atoms() else {
        panic!("scene has atoms")
    };
    let Some(rep) = scene.representation(rep_handle) else {
        panic!("representation resolves")
    };
    let mut out = Vec::new();
    pack_atoms(table, rep, &AtomSelection::All, &mut out).unwrap_or_else(|error| panic!("{error}"));
    assert_eq!(out.len(), 3);
    // Records preserve atom order and name their source row; the shader gathers
    // the position from the coordinate column through that identity.
    assert_eq!(out[1].entity_id.unpack(), Some((EntityKind::Atom, 1)));
    let expected = [1.5f32, 0.0, 0.0];
    let position = table.coords().slice().get(1).copied().unwrap_or_default();
    for (got, want) in position.iter().zip(expected) {
        assert!((got - want).abs() < f32::EPSILON);
    }
    // Spacefill draws full van der Waals radii.
    assert!(
        (out[0].radius - 1.55).abs() < 1e-6,
        "nitrogen radius unscaled"
    );
}

#[test]
fn ball_and_stick_scales_radii_down() {
    let (mut scene, _) = scene_table();
    let sel = scene.add_selection(AtomSelection::All);
    let Ok(bs) = scene.represent(sel, RepresentationKind::BallAndStick) else {
        panic!("ball-and-stick applies")
    };
    let Some(table) = scene.first_atoms() else {
        panic!("scene has atoms")
    };
    let Some(rep) = scene.representation(bs) else {
        panic!("representation resolves")
    };
    let mut out = Vec::new();
    pack_atoms(table, rep, &AtomSelection::All, &mut out).unwrap_or_else(|error| panic!("{error}"));
    assert!(out[0].radius < 1.0, "ball radii are a fraction of vdW");
}

#[test]
fn negative_radius_scale_cannot_create_inverted_gpu_spheres() {
    let (mut scene, rep_handle) = scene_table();
    let Some(representation) = scene.representation_mut(rep_handle) else {
        panic!("representation resolves")
    };
    representation.params.radius_scale = -4.0;
    let Some(table) = scene.first_atoms() else {
        panic!("scene has atoms")
    };
    let Some(representation) = scene.representation(rep_handle) else {
        panic!("representation resolves")
    };
    let mut atoms = Vec::new();
    pack_atoms(table, representation, &AtomSelection::All, &mut atoms)
        .unwrap_or_else(|error| panic!("{error}"));
    assert!(atoms.iter().all(|atom| atom.radius == 0.0));
}

#[test]
fn exact_sas_impostors_pack_the_probe_once() {
    let (mut scene, rep_handle) = scene_table();
    let Some(representation) = scene.representation_mut(rep_handle) else {
        panic!("representation resolves")
    };
    representation.kind = RepresentationKind::Surface;
    representation.params.surface_kind = SurfaceKind::SolventAccessible;
    let probe = representation.params.probe_radius;
    let mut atoms = Vec::new();
    let Some(table) = scene.first_atoms() else {
        panic!("scene has atoms")
    };
    let Some(representation) = scene.representation(rep_handle) else {
        panic!("representation resolves")
    };
    let expected = table.radius().values()[0] * representation.params.radius_scale + probe;
    pack_atoms(table, representation, &AtomSelection::All, &mut atoms)
        .unwrap_or_else(|error| panic!("{error}"));
    assert_eq!(atoms[0].radius.to_bits(), expected.to_bits());
}

#[test]
fn representation_opacity_stays_out_of_shared_atom_records() {
    let (mut scene, rep_handle) = scene_table();
    let Some(representation) = scene.representation_mut(rep_handle) else {
        panic!("representation resolves")
    };
    representation.material.opacity = 0.4;
    let Some(table) = scene.first_atoms() else {
        panic!("scene has atoms")
    };
    let Some(representation) = scene.representation(rep_handle) else {
        panic!("representation resolves")
    };
    let mut atoms = Vec::new();
    pack_atoms(table, representation, &AtomSelection::All, &mut atoms)
        .unwrap_or_else(|error| panic!("{error}"));
    assert!(atoms.iter().all(|atom| atom.color.a == u8::MAX));
}

#[test]
fn no_colour_scheme_reaches_the_shared_atom_records() {
    // The record carries the element colour so that changing a scheme is a
    // uniform write rather than a repack. A uniform scheme must therefore leave
    // the records byte-identical to the element-coloured ones.
    let (mut scene, rep_handle) = scene_table();
    let mut element_colored = Vec::new();
    {
        let Some(table) = scene.first_atoms() else {
            panic!("scene has atoms")
        };
        let Some(representation) = scene.representation(rep_handle) else {
            panic!("representation resolves")
        };
        pack_atoms(
            table,
            representation,
            &AtomSelection::All,
            &mut element_colored,
        )
        .unwrap_or_else(|error| panic!("{error}"));
    }
    let Some(representation) = scene.representation_mut(rep_handle) else {
        panic!("representation resolves")
    };
    representation.color = ColorScheme::Uniform(Rgba8::opaque(12, 34, 56));
    representation.material.opacity = 0.5;
    let Some(table) = scene.first_atoms() else {
        panic!("scene has atoms")
    };
    let Some(representation) = scene.representation(rep_handle) else {
        panic!("representation resolves")
    };
    let mut atoms = Vec::new();
    pack_atoms(table, representation, &AtomSelection::All, &mut atoms)
        .unwrap_or_else(|error| panic!("{error}"));
    assert_eq!(atoms, element_colored, "a scheme change must not repack");
    assert!(atoms.iter().all(|atom| atom.color.a == u8::MAX));
}

#[test]
fn a_property_colour_scheme_leaves_records_element_coloured() {
    // The ramp is applied in the shader against the property arena, so packing
    // under a property scheme must not bake ramp colours into the records.
    let (mut scene, rep_handle) = scene_table();
    let Some((owner, _)) = scene.structures().next() else {
        panic!("structure exists")
    };
    let property = AtomProperty::new(
        owner,
        "confidence",
        Arc::from([0.0, 0.5, f32::NAN]),
        AtomPropertyMeaning::Confidence,
        ScalarFieldSemantics::UncalibratedRank,
    )
    .unwrap_or_else(|error| panic!("property validates: {error}"));
    let property_handle = scene
        .add_atom_property(property)
        .unwrap_or_else(|error| panic!("property attaches: {error}"));
    let Some(property) = scene.atom_property(property_handle) else {
        panic!("property resolves")
    };
    let scheme = ColorScheme::property(property_handle, property);
    let Some(representation) = scene.representation_mut(rep_handle) else {
        panic!("representation resolves")
    };
    representation.color = scheme;
    let Some((_, placed)) = scene.structures().next() else {
        panic!("structure exists")
    };
    let Some(representation) = scene.representation(rep_handle) else {
        panic!("representation resolves")
    };
    let mut atoms = Vec::new();
    pack_atoms(
        &placed.atoms,
        representation,
        &AtomSelection::All,
        &mut atoms,
    )
    .unwrap_or_else(|error| panic!("{error}"));
    let elements = placed.atoms.color().values();
    for (atom, element) in atoms.iter().zip(elements) {
        assert_eq!(atom.color, *element, "records stay element-coloured");
    }
}

#[test]
fn a_property_appearance_leaves_records_element_coloured() {
    // Opacity and edge softness are presentation uniforms; only the element
    // colour belongs in the shared record.
    let (mut scene, rep_handle) = scene_table();
    let Some((owner, _)) = scene.structures().next() else {
        panic!("structure exists")
    };
    let property = AtomProperty::new(
        owner,
        "confidence",
        Arc::from([0.0, 1.0, f32::NAN]),
        AtomPropertyMeaning::Confidence,
        ScalarFieldSemantics::UncalibratedRank,
    )
    .unwrap_or_else(|error| panic!("property validates: {error}"));
    let property_handle = scene
        .add_atom_property(property)
        .unwrap_or_else(|error| panic!("property attaches: {error}"));
    let appearance = PropertyAppearance::confidence(property_handle, [0.0, 1.0])
        .unwrap_or_else(|error| panic!("appearance validates: {error}"));
    let Some(representation) = scene.representation_mut(rep_handle) else {
        panic!("representation resolves")
    };
    representation.appearance = Some(appearance);
    let Some((_, placed)) = scene.structures().next() else {
        panic!("structure exists")
    };
    let Some(representation) = scene.representation(rep_handle) else {
        panic!("representation resolves")
    };
    let mut atoms = Vec::new();
    pack_atoms(
        &placed.atoms,
        representation,
        &AtomSelection::All,
        &mut atoms,
    )
    .unwrap_or_else(|error| panic!("{error}"));
    let elements = placed.atoms.color().values();
    for (atom, element) in atoms.iter().zip(elements) {
        assert_eq!(atom.color, *element);
        assert_eq!(atom.color.a, u8::MAX, "opacity is not baked in");
        assert_eq!(atom.semantic >> 24, 0, "softness is not baked in");
    }
}

#[test]
fn a_partial_selection_packs_only_its_rows() {
    let (scene, rep_handle) = scene_table();
    let Some(table) = scene.first_atoms() else {
        panic!("scene has atoms")
    };
    let Some(rep) = scene.representation(rep_handle) else {
        panic!("representation resolves")
    };
    let mut out = Vec::new();
    pack_atoms(table, rep, &AtomSelection::Sparse(vec![0, 2]), &mut out)
        .unwrap_or_else(|error| panic!("{error}"));
    assert_eq!(out.len(), 2);
    let (Some((_, i0)), Some((_, i2))) = (out[0].entity_id.unpack(), out[1].entity_id.unpack())
    else {
        panic!("entity ids unpack")
    };
    assert_eq!((i0, i2), (0, 2), "entity ids keep original row indices");
}

#[test]
fn licorice_uses_the_bond_radius_for_atom_junctions() {
    let structure = fixture_structure();
    let mut scene = match Scene::from_structure(&structure) {
        Ok(scene) => scene,
        Err(error) => panic!("fixture scene builds: {error}"),
    };
    let selection = scene.add_selection(AtomSelection::All);
    let representation = match scene.represent(selection, RepresentationKind::Licorice) {
        Ok(handle) => handle,
        Err(error) => panic!("licorice applies: {error}"),
    };
    let Some(table) = scene.first_atoms() else {
        panic!("scene has atoms")
    };
    let Some(representation) = scene.representation(representation) else {
        panic!("representation resolves")
    };
    let mut atoms = Vec::new();
    pack_atoms(table, representation, &AtomSelection::All, &mut atoms)
        .unwrap_or_else(|error| panic!("{error}"));
    assert!(
        atoms
            .iter()
            .all(|atom| (atom.radius - representation.params.bond_radius).abs() < f32::EPSILON)
    );
}

#[test]
fn repacking_reuses_the_scratch_allocation() {
    let (scene, rep_handle) = scene_table();
    let Some(table) = scene.first_atoms() else {
        panic!("scene has atoms")
    };
    let Some(rep) = scene.representation(rep_handle) else {
        panic!("representation resolves")
    };
    let mut out = Vec::new();
    pack_atoms(table, rep, &AtomSelection::All, &mut out).unwrap_or_else(|error| panic!("{error}"));
    let capacity = out.capacity();
    let pointer = out.as_ptr();
    pack_atoms(table, rep, &AtomSelection::All, &mut out).unwrap_or_else(|error| panic!("{error}"));
    assert_eq!(out.capacity(), capacity);
    assert_eq!(
        out.as_ptr(),
        pointer,
        "steady-state repack allocates nothing"
    );
}

#[test]
fn a_bead_encloses_its_residue_and_keeps_one_sphere_per_residue() {
    use molgfx_core::{AtomSelection, RepresentationTarget, Scene};

    let structure = fixture_structure();
    let mut scene = Scene::new();
    let Ok(handle) = scene.add_structure(&structure) else {
        panic!("fixture places")
    };
    let selection = AtomSelection::All;
    let selection_handle = scene.add_selection(AtomSelection::All);
    let mut representation = Representation::new(
        RepresentationTarget::Selection(selection_handle),
        RepresentationKind::Beads,
    );
    representation.params.radius_scale = 1.0;
    let Some(placed) = scene.structure(handle) else {
        panic!("placed structure resolves")
    };

    let mut beads = Vec::new();
    pack_residue_beads(
        &placed.atoms,
        &placed.hierarchy,
        ColorContext::new(&scene, handle),
        &representation,
        &selection,
        &mut beads,
    )
    .unwrap_or_else(|error| panic!("{error}"));

    let residues = placed.atoms.residue().values();
    let distinct = {
        let mut seen: Vec<u32> = Vec::new();
        for residue in residues {
            if !seen.contains(residue) {
                seen.push(*residue);
            }
        }
        seen.len()
    };
    assert_eq!(beads.len(), distinct, "one bead per residue");

    // Each bead answers for its own residue: it names the residue's first atom,
    // and its radius must enclose every one of that residue's atoms, measured
    // from the first selected atom whose coordinate the shader gathers.
    let coords = placed.atoms.coords().slice();
    let radii = placed.atoms.radius().values();
    for bead in &beads {
        let Some((kind, first)) = bead.entity_id.unpack() else {
            panic!("bead names a source row");
        };
        assert_eq!(kind, EntityKind::Atom);
        let Some(residue) = residues.get(first as usize) else {
            panic!("bead source row is a real atom");
        };
        let members: Vec<usize> = residues
            .iter()
            .enumerate()
            .filter(|(_, candidate)| candidate == &residue)
            .map(|(index, _)| index)
            .collect();
        let centre = coords
            .get(first as usize)
            .map_or(molgfx_math::Vec3::ZERO, |value| {
                molgfx_math::Vec3::from_array(*value)
            });
        for member in members {
            let Some(point) = coords.get(member).copied() else {
                continue;
            };
            let point = molgfx_math::Vec3::from_array(point);
            let extent = radii.get(member).copied().unwrap_or_default();
            assert!(
                centre.distance(point) + extent <= bead.radius + 1.0e-3,
                "residue {residue} atom {member} lies inside its bead"
            );
        }
    }
}