Skip to content

Extending the Crate

Tutorials for writing your own AtomRestraint / Region / Handler / Relaxer types. Every extension trait in this crate follows the same shape (direction-3 rule — see concepts):

pub trait X + N concrete pub struct types implementing it. User types impl X identically. No built-in/plugin type-level distinction.

Custom AtomRestraint

Goal: pull atoms toward a target plane with a quadratic attractive well.

Step 1 — define the struct

use molrs::types::F;
# use molpack::AtomRestraint;

#[derive(Debug, Clone, Copy)]
pub struct PlaneTether {
    pub normal: [F; 3],
    pub offset: F,
    pub k: F,
}
# impl AtomRestraint for PlaneTether {
#     fn f(&self, _x: &[F; 3], _s: F, _s2: F) -> F { 0.0 }
#     fn fg(&self, _x: &[F; 3], _s: F, _s2: F, _g: &mut [F; 3]) -> F { 0.0 }
# }
  • pub fields — users construct with PlaneTether { normal, offset, k }, no builder.
  • Debug required because AtomRestraint has a Debug supertrait bound (so Target's derived Debug keeps working).

Step 2 — implement AtomRestraint

# use molrs::types::F;
# use molpack::AtomRestraint;
# #[derive(Debug)]
# pub struct PlaneTether { pub normal: [F; 3], pub offset: F, pub k: F }
impl AtomRestraint for PlaneTether {
    fn f(&self, pos: &[F; 3], _scale: F, _scale2: F) -> F {
        let d = self.normal[0] * pos[0]
              + self.normal[1] * pos[1]
              + self.normal[2] * pos[2]
              - self.offset;
        0.5 * self.k * d * d
    }
    fn fg(&self, pos: &[F; 3], scale: F, scale2: F, g: &mut [F; 3]) -> F {
        let d = self.normal[0] * pos[0]
              + self.normal[1] * pos[1]
              + self.normal[2] * pos[2]
              - self.offset;
        g[0] += self.k * d * self.normal[0];
        g[1] += self.k * d * self.normal[1];
        g[2] += self.k * d * self.normal[2];
        self.f(pos, scale, scale2)
    }
}

Three contracts that all restraints must obey:

  1. Gradient accumulates with +=. Multiple restraints may touch the same atom.
  2. fg returns the value. The hot path uses the returned value for the fdist/frest accumulation — don't return 0.0 just because the caller might discard it.
  3. Scale/scale2 usage is your choice. Linear-penalty restraints typically use scale; quadratic-penalty ones use scale2. Your tether uses its own k — ignore both knobs if you prefer.

Step 3 — write a gradient test

# use molrs::types::F;
# use molpack::AtomRestraint;
# #[derive(Debug)] pub struct PlaneTether { pub normal: [F; 3], pub offset: F, pub k: F }
# impl AtomRestraint for PlaneTether {
#     fn f(&self, _x: &[F; 3], _s: F, _s2: F) -> F { 0.0 }
#     fn fg(&self, _x: &[F; 3], _s: F, _s2: F, _g: &mut [F; 3]) -> F { 0.0 }
# }
#[test]
fn plane_tether_gradient_matches_fd() {
    let r = PlaneTether { normal: [0.0, 0.0, 1.0], offset: 5.0, k: 2.0 };
    let x = [1.0, 2.0, 7.0];
    let mut g = [0.0; 3];
    let _ = r.fg(&x, 1.0, 1.0, &mut g);
    let h: F = 1e-5;
    for k in 0..3 {
        let mut xp = x; xp[k] += h;
        let mut xm = x; xm[k] -= h;
        let fd = (r.f(&xp, 1.0, 1.0) - r.f(&xm, 1.0, 1.0)) / (2.0 * h);
        assert!(
            (g[k] - fd).abs() < 1e-4,
            "axis {k}: analytic={}, fd={}", g[k], fd,
        );
    }
}

Convention: ε = 1e-5, tolerance 1e-3 (looser if the restraint has kinks).

Step 4 — use it

# use molrs::types::F;
# use molpack::AtomRestraint;
# #[derive(Debug, Clone, Copy)] pub struct PlaneTether { pub normal: [F; 3], pub offset: F, pub k: F }
# impl AtomRestraint for PlaneTether {
#     fn f(&self, _x: &[F; 3], _s: F, _s2: F) -> F { 0.0 }
#     fn fg(&self, _x: &[F; 3], _s: F, _s2: F, _g: &mut [F; 3]) -> F { 0.0 }
# }
use molpack::{InsideBoxRestraint, Target};
# let (pos, rad) = (&[[0.0; 3]][..], &[1.0][..]);

let target = Target::from_coords(pos, rad, 100)
    .with_restraint(InsideBoxRestraint::new([0.0; 3], [40.0; 3], [false; 3]))
    .with_restraint(PlaneTether { normal: [0.0, 0.0, 1.0], offset: 20.0, k: 1.0 });

Built-in InsideBoxRestraint and user PlaneTether take the same code path — direction-3 in action.

In Python — the same restraint, duck-typed

The interfaces are also exposed to Python through a duck-typed protocol: a restraint is any object exposing f and fg, and the packer consumes it on the same code path as a built-in. That makes Python the natural place to prototype a restraint the input grammar can't express — a few lines, no recompile — and then, once it earns its place, swap in the native equivalent with the driver script unchanged.

A duck-typed restraint is any object exposing f and fg. The Python protocol differs from the Rust trait in exactly one way: fg returns (energy, (gx, gy, gz)) instead of accumulating into a &mut [F; 3] — the binding does the += for you. The contracts are otherwise identical: fg returns the energy, the gradient is the chain rule (dU/dξ)·∇ξ, a linear-energy penalty rides scale (a quadratic one scale2). A missing f or fg is rejected at attach time with a TypeError, and an exception raised inside fg propagates back out of pack. The PlaneTether above, duck-typed in Python:

import numpy as np

class PlaneTether:
    """Pull a site toward the plane z = offset (a per-site geometric bias)."""
    def __init__(self, normal, offset, k):
        n = np.asarray(normal, float)
        self.n, self.offset, self.k = n / np.linalg.norm(n), offset, k

    def f(self, x, scale, scale2):
        d = float(self.n @ np.asarray(x, float)) - self.offset
        return scale * 0.5 * self.k * d * d

    def fg(self, x, scale, scale2):
        d = float(self.n @ np.asarray(x, float)) - self.offset
        g = scale * self.k * d * self.n
        return scale * 0.5 * self.k * d * d, (float(g[0]), float(g[1]), float(g[2]))

Attach it with with_atom_restraint (a site subset) or with_restraint (every copy), and pack.

A per-site field cannot reproduce a distribution. It is tempting to build a per-site penalty from a target density ρ* (e.g. Boltzmann inversion U = −kT·ln ρ*) to make sites follow ρ*. This fails under packing's energy minimisation: ∑ᵢ U(xᵢ) is minimised by driving every site to the single minimum of U — the mode of ρ* — so the sites collapse onto the peak instead of spreading over the distribution. To drive a whole species onto a target profile, use a collective restraint (Target.with_collective_restraint, e.g. the built-in ProfileMatch), whose penalty is a function of the entire group and whose gradient couples the copies, so the fixed point is empirical distribution = target.

Custom Region

Goal: a conical region with apex at origin, axis along +z, half-angle 30°.

use molrs::types::F;
# use molpack::Region;

#[derive(Debug, Clone, Copy)]
pub struct ConeRegion {
    pub apex: [F; 3],
    pub axis: [F; 3],
    pub half_angle_cos: F,
}

impl Region for ConeRegion {
    fn contains(&self, x: &[F; 3]) -> bool {
        self.signed_distance(x) <= 0.0
    }
    fn signed_distance(&self, x: &[F; 3]) -> F {
        let dx = x[0] - self.apex[0];
        let dy = x[1] - self.apex[1];
        let dz = x[2] - self.apex[2];
        let r = (dx * dx + dy * dy + dz * dz).sqrt();
        if r < 1e-12 { return 0.0; }
        let axis_dot =
            (dx * self.axis[0] + dy * self.axis[1] + dz * self.axis[2]) / r;
        self.half_angle_cos - axis_dot
    }
    // Default FD gradient is OK for prototypes. Override analytically
    // for hot-path use — see below.
}

Compose with built-ins

# use molrs::types::F;
# use molpack::Region;
# #[derive(Debug, Clone, Copy)]
# pub struct ConeRegion { pub apex: [F; 3], pub axis: [F; 3], pub half_angle_cos: F }
# impl Region for ConeRegion {
#     fn contains(&self, _x: &[F; 3]) -> bool { true }
#     fn signed_distance(&self, _x: &[F; 3]) -> F { 0.0 }
# }
use molpack::{InsideSphereRegion, RegionExt, RegionRestraint, Target};
# let (pos, rad) = (&[[0.0; 3]][..], &[1.0][..]);

let cone = ConeRegion {
    apex: [0.0; 3],
    axis: [0.0, 0.0, 1.0],
    half_angle_cos: (std::f64::consts::PI / 6.0).cos(),
};
let sphere = InsideSphereRegion::new([0.0; 3], 10.0);
let region = cone.and(sphere);

let target = Target::from_coords(pos, rad, 100)
    .with_restraint(RegionRestraint(region));

RegionExt::and / or / not come from a blanket impl on every Region. The resulting type And<ConeRegion, InsideSphereRegion> is static-dispatch — no heap.

Analytic gradient override

For hot-path use, override signed_distance_grad analytically. The cone above:

# use molrs::types::F;
# use molpack::Region;
# #[derive(Debug, Clone, Copy)]
# pub struct ConeRegion { pub apex: [F; 3], pub axis: [F; 3], pub half_angle_cos: F }
# impl Region for ConeRegion {
#     fn contains(&self, _x: &[F; 3]) -> bool { true }
#     fn signed_distance(&self, _x: &[F; 3]) -> F { 0.0 }
fn signed_distance_grad(&self, x: &[F; 3]) -> [F; 3] {
    let dx = x[0] - self.apex[0];
    let dy = x[1] - self.apex[1];
    let dz = x[2] - self.apex[2];
    let r2 = dx * dx + dy * dy + dz * dz;
    let r = r2.sqrt();
    if r < 1e-12 { return [0.0; 3]; }
    let axis_dot = (dx * self.axis[0] + dy * self.axis[1] + dz * self.axis[2]) / r;
    let inv_r = 1.0 / r;
    // signed_distance = cos(α) - axis_dot ⇒ grad = -∂axis_dot/∂x
    [
        -(self.axis[0] * inv_r - axis_dot * dx * inv_r * inv_r),
        -(self.axis[1] * inv_r - axis_dot * dy * inv_r * inv_r),
        -(self.axis[2] * inv_r - axis_dot * dz * inv_r * inv_r),
    ]
}
# }

Then finite-difference check it — same pattern as the AtomRestraint test.

Custom Handler

Goal: a handler that writes a CSV row per step so you can plot the objective evolution.

use std::fs::File;
use std::io::{BufWriter, Write};
use molpack::{F, Handler, PackContext, StepInfo};

pub struct CsvHandler { writer: BufWriter<File> }

impl CsvHandler {
    pub fn new(path: &str) -> std::io::Result<Self> {
        let mut w = BufWriter::new(File::create(path)?);
        writeln!(w, "phase,loop_idx,fdist,frest,improvement_pct")?;
        Ok(Self { writer: w })
    }
}

impl Handler for CsvHandler {
    fn on_step(&mut self, info: &StepInfo, _sys: &PackContext) {
        let _ = writeln!(
            self.writer,
            "{},{},{},{},{}",
            info.phase.phase,
            info.loop_idx,
            info.fdist,
            info.frest,
            info.improvement_pct,
        );
    }
}

Handler notes:

  • on_step is the only required method. Everything else has a default no-op.
  • sys is &PackContext, never &mut. Handlers cannot mutate packer state — use a Relaxer if you need to.
  • Multiple handlers run in registration order. Register your CSV handler before ProgressHandler to get a row on every step, vice-versa otherwise.
  • should_stop is polled every iteration. Return true to break the outer loop early. Useful for time budgets or custom convergence criteria.

Custom Relaxer

Goal: a relaxer that tries random rigid-body translations and accepts if the objective decreases.

use molrs::Frame;
use molrs::types::F;
use molpack::{Relaxer, RelaxerRunner};
use rand::{Rng, RngCore};

#[derive(Debug, Clone)]
pub struct JiggleRelaxer {
    pub steps: usize,
    pub max_delta: F,
}

impl Relaxer for JiggleRelaxer {
    fn spawn(&self, _frame: Option<&Frame>, _ref_coords: &[[F; 3]]) -> Box<dyn RelaxerRunner> {
        Box::new(JiggleRunner {
            steps: self.steps,
            max_delta: self.max_delta,
            accepted: 0,
            total: 0,
        })
    }
}

pub struct JiggleRunner {
    steps: usize,
    max_delta: F,
    accepted: usize,
    total: usize,
}

impl RelaxerRunner for JiggleRunner {
    fn on_iter(
        &mut self,
        coords: &[[F; 3]],
        f_current: F,
        evaluate: &mut dyn FnMut(&[[F; 3]]) -> F,
        rng: &mut dyn RngCore,
    ) -> Option<Vec<[F; 3]>> {
        let mut best = coords.to_vec();
        let mut best_f = f_current;
        let mut accepted_any = false;
        for _ in 0..self.steps {
            let dx = (rng.next_u32() as f64 / u32::MAX as f64 * 2.0 - 1.0) * self.max_delta;
            let dy = (rng.next_u32() as f64 / u32::MAX as f64 * 2.0 - 1.0) * self.max_delta;
            let dz = (rng.next_u32() as f64 / u32::MAX as f64 * 2.0 - 1.0) * self.max_delta;
            let trial: Vec<[F; 3]> = best.iter()
                .map(|p| [p[0] + dx, p[1] + dy, p[2] + dz])
                .collect();
            let f_trial = evaluate(&trial);
            self.total += 1;
            if f_trial < best_f {
                self.accepted += 1;
                best_f = f_trial;
                best = trial;
                accepted_any = true;
            }
        }
        if accepted_any { Some(best) } else { None }
    }
    fn acceptance_rate(&self) -> F {
        if self.total == 0 { 0.0 } else { self.accepted as F / self.total as F }
    }
}

Relaxer notes:

  • Two-part design. Relaxer is the immutable builder; RelaxerRunner holds per-pack state. spawn() is called once per target type at pack() entry.
  • evaluate closure tests trial coords against the full objective without mutating the reference — use it as often as you like.
  • Return Some(new_coords) only if something changed. The packer skips unnecessary cache invalidation when you return None.
  • count == 1 required. Multi-copy targets share reference coords; a relaxer that mutates them would silently change all copies.

Testing discipline

Kind Location Convention
Unit test #[cfg(test)] mod tests in the same file One #[test] fn per behavior
Integration test tests/<name>.rs use molpack::{…}; only public API
Gradient finite-difference alongside unit test ε=1e-5, tol=1e-3
Regression vs Packmol tests/examples_batch.rs (#[ignore]) Run with --ignored --release

Run all:

cargo test --all-features
cargo test --release --test examples_batch -- --ignored

Rules:

  • Every new AtomRestraint gets an FD gradient test.
  • Every new Region gets a boolean-algebra + signed-distance sign test and (for hot-path use) an analytic-gradient FD test.

Common pitfalls

  • Gradient sign. Every AtomRestraint accumulates ∂penalty/∂x. Optimizer negates for descent. If your molecules fly out of the region, the gradient has the wrong sign — penalty should point toward the violation boundary.
  • Rotation convention. Single-atom tests pass with both LEFT and RIGHT Euler multiplication; multi-atom tests don't. Always test Euler changes with ≥ 2 atoms.
  • Cell<f64> is not Sync. Use AtomicU64 + f64::to_bits / from_bits for interior mutability in Send + Sync contexts.
  • 0-based atom indexing. Target::with_atom_restraint uses Rust-native 0-based indices. &[0, 1] selects the first two atoms. Packmol .inp files use 1-based — subtract 1 at the parse boundary.
  • count == 1 required for relaxers. Multi-copy targets share reference coords.
  • radscale is phase-dependent. Don't hard-code atomic radii — always go through sys.radius[i]. evaluate_unscaled temporarily swaps radius with radius_ini for user-facing numbers.
  • PBC boxes must be valid. Zero-length axis returns PackError::InvalidPBCBox.

Contributing flow

  1. Write a failing test.
  2. Implement until it passes.
  3. Run the full gate: 
    cargo test --all-features
cargo clippy -- -D warnings
cargo fmt --all --check