Restraints¶

Restraints are geometric regions (or half-spaces) that every atom of a target — or a chosen subset — must lie inside. molpack ships two families of built-in restraints: five geometric region restraints (below), and six collective distribution-matching restraints (further down). Both attach the same way, via .with_restraint().

Geometric built-ins¶

All arguments are standard Python floats / lists.

from molpack import (
    AbovePlaneRestraint,
    BelowPlaneRestraint,
    InsideBoxRestraint,
    InsideSphereRestraint,
    OutsideSphereRestraint,
)

box   = InsideBoxRestraint([0, 0, 0], [40, 40, 40])
ball  = InsideSphereRestraint([0, 0, 0], 20.0)
shell = OutsideSphereRestraint([0, 0, 0], 10.0)
above = AbovePlaneRestraint(normal=[0, 0, 1], distance=5.0)
below = BelowPlaneRestraint(normal=[0, 0, 1], distance=20.0)

Periodic boxes¶

InsideBoxRestraint doubles as the PBC declaration. Passing a periodic tuple turns any subset of axes periodic:

InsideBoxRestraint([0, 0, 0], [30, 30, 30], periodic=(True, True, True))

See Periodic boundaries for the full semantics and validation rules.

Collective (distribution-matching) restraints¶

Where a geometric restraint penalises each atom against a region, a collective restraint sees every copy of the target at once and drives the species' spatial distribution toward a target profile (via a squared 1-D Wasserstein penalty). Six built-ins cover Gaussian, exponential, and arbitrary-tabulated priors along either a plane or a radius — e.g. a Gaussian slab centred at z = 20:

from molpack import GaussianPlane

slab = GaussianPlane(normal=[0, 0, 1], offset=0.0, strength=1.0, mu=20.0, sigma=3.0)
target = Target(frame, count=200).with_restraint(slab)

The full list and constructor signatures are in the API reference.

Stacking multiple restraints¶

Apply several restraints to the same target by chaining .with_restraint():

target = (
    Target(frame, count=500)
    .with_name("water")
    .with_restraint(InsideBoxRestraint([0, 0, 0], [40, 40, 40]))
    .with_restraint(OutsideSphereRestraint([20, 20, 20], 5.0))
)

Each call attaches an independent restraint. All active restraints are evaluated at every optimizer step.

Scopes¶

A restraint can be applied at two scopes:

• Whole target — target.with_restraint(r) — penalises every atom.
• Atom subset — target.with_atom_restraint([0, 1, 2], r) — penalises only the listed atoms (0-based indices).

Example — a bilayer: pin heads above z=12, tails below z=2:

lipid = (
    Target(frame, count=20)
    .with_name("lipid")
    .with_restraint(InsideBoxRestraint([0, 0, 0], [40, 40, 14]))
    .with_atom_restraint([0, 1],   AbovePlaneRestraint([0, 0, 1], 12.0))
    .with_atom_restraint([30, 31], BelowPlaneRestraint([0, 0, 1], 2.0))
)

Global restraints¶

To apply one restraint to every target in a pack, attach it on the packer:

packer = (
    Molpack()
    .with_global_restraint(InsideBoxRestraint([0, 0, 0], [40, 40, 40]))
)

Semantically equivalent to calling .with_restraint(r) on every target, but avoids the duplication.

Custom restraints¶

Pass any object implementing f(x, scale, scale2) -> float and fg(x, scale, scale2) -> (float, (gx, gy, gz)). See the Restraint Protocol in molpack for the full contract.

class SphereRestraint:
    def __init__(self, center, radius):
        self.c = np.asarray(center)
        self.r = radius
    def f(self, x, scale, scale2):
        d = np.linalg.norm(np.asarray(x) - self.c) - self.r
        return scale2 * d * d if d > 0 else 0.0
    def fg(self, x, scale, scale2):
        rel = np.asarray(x) - self.c
        d = float(np.linalg.norm(rel))
        over = d - self.r
        if over <= 0:
            return 0.0, (0.0, 0.0, 0.0)
        factor = 2 * scale2 * over / d
        return scale2 * over * over, tuple(factor * rel)

target = Target(frame, count=10).with_restraint(SphereRestraint([0,0,0], 6.0))

Semantics¶

Every restraint contributes a continuously differentiable penalty \(f_{\text{rest}}(\mathbf{x})\) that is zero inside the allowed region and rises quadratically outside. The aggregate objective minimised by the packer is:

where \(f_{\text{dist}}\) is the pairwise distance-violation sum for the user-specified tolerance. Convergence is declared when both fall below precision.