RefinementRestraints: AtomsNear - the atoms near a point, symmetry mates and lattice translates included

A point query on a populated gemmi NeighborSearch that returns each neighbour's copy nearest the
point with its distance and the Cartesian image transform, nearest first. The nearest-image shift and
the Cartesian form of a symmetry image, which Contacts() computed inline, are now two small helpers
that both use, so the restraint path and the new query share one image computation. The solvent audit
needs it to measure the surroundings of a water or an ion outside the restraint machinery.

Contacts() is unchanged in what it computes: a refinement of 1HR2 with --mode refine gives the
same refined model file, byte for byte.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01HVeDQqN9qRD57vnShzTC37

image_analysis/structure_refinement/RefinementRestraints.cpp

@@ -7,6 +7,7 @@
 #include <cmath>
 #include <functional>
 #include <set>
+#include <tuple>
 #include <unordered_map>
 
 #include "gemmi/neighbor.hpp"
@@ -576,6 +577,21 @@ RefinementTopology BuildRefinementTopology(gemmi::Structure &st, bool hydrogens)
 
 namespace {
 
+// The whole cells that bring the image of the atom at `pos` under `ft` nearest to the point at fractional `fi`.
+gemmi::Vec3 NearestShift(const gemmi::UnitCell &cell, const gemmi::FTransform &ft, const gemmi::Fractional &fi,
+                         const gemmi::Position &pos) {
+    const gemmi::Fractional fj = ft.apply(cell.fractionalize(pos));
+    return {std::round(fi.x - fj.x), std::round(fi.y - fj.y), std::round(fi.z - fj.z)};
+}
+
+// The image `ft` moved by `shift` whole cells, as a Cartesian x -> R x + t.
+gemmi::Transform CartesianImage(const gemmi::UnitCell &cell, const gemmi::FTransform &ft, const gemmi::Vec3 &shift) {
+    gemmi::Transform t;
+    t.mat = cell.orth.mat.multiply(ft.mat).multiply(cell.frac.mat);
+    t.vec = cell.orth.mat.multiply(ft.vec + shift);
+    return t;
+}
+
 // The pairs of atoms of `model`, symmetry mates and lattice translates included, closer than `reach` and not excluded, where at least one of
 // the two is refinable, neither is a hydrogen and the two are not alternative conformations of each other: `make`
 // turns each into a record (the image of j as x -> rot x + tran, Cartesian, and the distance), or says to leave it out.
@@ -617,8 +633,7 @@ RestraintSet Contacts(const RefinementTopology &topo, gemmi::Model &model, const
                     return;
                 // The image of j nearest to i, as a Cartesian x -> R x + t.
                 const gemmi::FTransform ft = ns.get_image_transformation(m.image_idx);
-                const gemmi::Fractional fj = ft.apply(cell.fractionalize(pos[j]));
-                const gemmi::Vec3 shift(std::round(fi.x - fj.x), std::round(fi.y - fj.y), std::round(fi.z - fj.z));
+                const gemmi::Vec3 shift = NearestShift(cell, ft, fi, pos[j]);
                 const bool same_copy = m.image_idx == 0 && shift.x == 0 && shift.y == 0 && shift.z == 0;
                 // Within the copy itself: each pair once, the excluded ones left out. With any other copy
                 // (a symmetry mate or a lattice translate): the pair is met again from j, so it is kept from
@@ -641,16 +656,15 @@ RestraintSet Contacts(const RefinementTopology &topo, gemmi::Model &model, const
                             break;
                     }
                 }
-                const gemmi::Mat33 rot = cell.orth.mat.multiply(ft.mat).multiply(cell.frac.mat);
-                const gemmi::Vec3 tran = cell.orth.mat.multiply(ft.vec + shift);
-                const gemmi::Position image(rot.multiply(pos[j]) + tran);
+                const gemmi::Transform t = CartesianImage(cell, ft, shift);
+                const gemmi::Position image(t.apply(pos[j]));
                 RestraintRecord r;
                 r.atom[0] = static_cast<int>(i);
                 r.atom[1] = j;
                 for (int a = 0; a < 3; a++) {
                     for (int b = 0; b < 3; b++)
-                        r.rot[3 * a + b] = static_cast<float>(rot[a][b]);
-                    r.tran[a] = static_cast<float>(tran.at(a));
+                        r.rot[3 * a + b] = static_cast<float>(t.mat[a][b]);
+                    r.tran[a] = static_cast<float>(t.vec.at(a));
                 }
                 if (make(static_cast<int>(i), j, pos[i].dist(image), r))
                     found[block].push_back(r);
@@ -667,6 +681,30 @@ RestraintSet Contacts(const RefinementTopology &topo, gemmi::Model &model, const
 
 }  // namespace
 
+std::vector<ModelNeighbour> AtomsNear(gemmi::NeighborSearch &ns, const gemmi::Model &model, const gemmi::Position &point,
+                                      double radius) {
+    const gemmi::UnitCell &cell = ns.grid.unit_cell;
+    const gemmi::Fractional fp = cell.fractionalize(point);
+    std::vector<ModelNeighbour> out;
+    ns.for_each(point, '\0', radius, [&](gemmi::NeighborSearch::Mark &m, double) {
+        const gemmi::Atom &a = model.chains[m.chain_idx].residues[m.residue_idx].atoms[m.atom_idx];
+        const gemmi::FTransform ft = ns.get_image_transformation(m.image_idx);
+        ModelNeighbour n;
+        n.chain = m.chain_idx;
+        n.residue = m.residue_idx;
+        n.atom = m.atom_idx;
+        n.image = CartesianImage(cell, ft, NearestShift(cell, ft, fp, a.pos));
+        n.pos = gemmi::Position(n.image.apply(a.pos));
+        n.distance = point.dist(n.pos);
+        out.push_back(n);
+    });
+    // nearest first; equal distances in model order, so the list does not depend on the search's cells
+    std::sort(out.begin(), out.end(), [](const ModelNeighbour &a, const ModelNeighbour &b) {
+        return std::tie(a.distance, a.chain, a.residue, a.atom) < std::tie(b.distance, b.chain, b.residue, b.atom);
+    });
+    return out;
+}
+
 RestraintSet BuildRepulsions(const RefinementTopology &topo, gemmi::Model &model, const gemmi::UnitCell &cell,
                              const gemmi::SpaceGroup &sg, double margin, size_t nthreads) {
     float max_radius = 0;

image_analysis/structure_refinement/RefinementRestraints.h

@@ -12,6 +12,7 @@
 
 #include "gemmi/model.hpp"
 #include "gemmi/monlib.hpp"
+#include "gemmi/neighbor.hpp"
 #include "gemmi/topo.hpp"
 #include "gemmi/unitcell.hpp"
 
@@ -97,6 +98,22 @@ void PlaceRidingHydrogens(RefinementTopology &topo);
 RestraintSet BuildRepulsions(const RefinementTopology &topo, gemmi::Model &model, const gemmi::UnitCell &cell,
                              const gemmi::SpaceGroup &sg, double margin, size_t nthreads = 1);
 
+// One atom of a model near a point (AtomsNear): which atom, and its copy - symmetry mate or lattice translate -
+// nearest the point.
+struct ModelNeighbour {
+    int chain = 0, residue = 0, atom = 0;   // indices into the model
+    gemmi::Position pos;                    // the copy's position
+    double distance = 0;                    // A, from the point
+    gemmi::Transform image;                 // x -> R x + t (Cartesian) taking the atom to its copy
+};
+
+// The atoms of `model` within `radius` of `point`, symmetry mates and lattice translates included, nearest first;
+// the point's own atom, where it is one, among them at distance 0. `ns` is a search populated over `model` in a
+// cell with its symmetry images set (UnitCell::set_cell_images_from_spacegroup), made with a radius of at least
+// `radius`; gemmi's NeighborSearch takes it non-const.
+std::vector<ModelNeighbour> AtomsNear(gemmi::NeighborSearch &ns, const gemmi::Model &model, const gemmi::Position &point,
+                                      double radius);
+
 // The B similarity of the atoms in contact (closer than 4 A, symmetry mates included) that no bond or angle relates,
 // waters among them, under the same conditions as BuildRepulsions().
 RestraintSet BuildBContacts(const RefinementTopology &topo, gemmi::Model &model, const gemmi::UnitCell &cell,