susan.ml.VolumePairs

class susan.ml.VolumePairs(tmp_base: str = 'vol_pair_tmp')

Bases: object

Buffer of independent half-map pairs for Noise2Noise training.

Warning

Experimental. This class is part of the machine-learning subpackage and may change or be removed in a future release.

Half-maps are produced by the SUSAN Averager with random half-set assignments, so each pair (half1, half2) shares the same underlying signal but has statistically independent noise realisations — the ideal N2N training condition.

The Averager already returns zero-mean, unit-std maps, so no further normalisation is applied here.

num_vol

Number of pairs currently stored in the buffer.

Type:

int

max_vol

Maximum number of pairs the buffer can hold.

Type:

int

box_size

Full edge length of the input volume in voxels.

Type:

int

buffer

Internal storage tensor of shape (max_vol, 2, Z, Y, X). None until set_size() or set_size_mask() is called.

Type:

torch.Tensor or None

tmp_base

Path prefix used for temporary files written by populate(). Default: 'vol_pair_tmp'.

Type:

str

mask

Cropped mask retained by set_size_mask(), floored at min_mask_value. None until set_size_mask() is called.

Type:

numpy.ndarray or None

apply_mask

When True, crop_vol() multiplies every cropped volume by mask before it is stored in the buffer. Default: False.

Type:

bool

min_mask_value

Lower bound applied to the mask when it is stored (mask = maximum(min_mask_value, input_mask)). A non-zero value attenuates — rather than zeroes — voxels outside the mask. Default: 0.0.

Type:

float

Setup

set_size(vol_size: int, num_vol: int, padding: int)

Initialise the buffer with a symmetric fixed-padding crop region.

Removes padding voxels from each face of the cubic volume. Use this when a simple isotropic margin is sufficient.

Parameters:

• vol_size (int) – Full edge length of the input volume in voxels.

• num_vol (int) – Number of half-map pairs to allocate space for.

• padding (int) – Number of voxels to crop from each face.

set_size_mask(vol_size: int, num_vol: int, mask: ndarray, extra_pad: int = 10, apply_mask: bool = False, min_mask_value: float = 0.0)

Initialise the buffer using the tight bounding box of a mask.

Computes the axis-aligned bounding box of the non-zero voxels in mask, expands it by extra_pad voxels on each side, and uses the result as the crop region. More memory-efficient than set_size() when the molecule occupies a small fraction of the box.

Parameters:

• vol_size (int) – Full edge length of the input volume in voxels.

• num_vol (int) – Number of half-map pairs to allocate space for.

• mask (numpy.ndarray) – 3-D binary (or non-negative) mask volume. Non-zero voxels define the region of interest.

• extra_pad (int, optional) – Extra voxels added around the bounding box on each side. Default: 10.

• apply_mask (bool, optional) – If True, crop_vol() multiplies every volume by the (cropped, floored) mask before it is stored in the buffer. Sets apply_mask. Default: False.

• min_mask_value (float, optional) – Lower bound for the retained mask: mask = maximum(min_mask_value, mask). With 0.0 voxels outside the mask are zeroed; a small positive value (e.g. 0.1) attenuates them instead. Sets min_mask_value. Default: 0.0.

Raises:

ValueError – If mask contains no non-zero voxels.

Population

init_with_halfmaps(vol1: ndarray, vol2: ndarray)

Store a single externally provided half-map pair.

Resets the buffer and stores the cropped pair at index 0. Use this when half-maps have been reconstructed outside SUSAN and only a single pair is needed (e.g. for inference or a quick training run).

Parameters:

• vol1 (numpy.ndarray) – First half-map (zero-mean, unit-std).

• vol2 (numpy.ndarray) – Second half-map (zero-mean, unit-std).

Raises:

RuntimeError – If the buffer has not been initialised; call set_size() or set_size_mask() first.

populate(avgr, ptcls, tomo_filename: str, num_entries: int = None, sigma_ang_3D=0.0, sigma_ang_2D=0.0, sigma_def=0.0)

Fill the buffer by re-randomising half-set assignments and running the Averager.

Each iteration re-randomises the half_id array in ptcls, saves a temporary particles file, runs the Averager to produce a pair of half-maps, and stores the cropped pair in the buffer. The loop repeats until the buffer is full or num_entries new pairs have been added.

When sigma_ang_3D or sigma_ang_2D is non-zero, the orientations of the half-1 particles are perturbed by an independent isotropic small rotation before reconstruction, so vol1 (buffer slot 0) becomes a deliberately misaligned reconstruction while vol2 (slot 1) stays cleanly aligned. Train with Noise2NoiseTrainer.train() directional=True on such pairs to teach the network to recover from small alignment errors (slot 0 = input, slot 1 = target).

The Averager’s verbosity and rec_halfsets attributes, and the original half_id, ali_eu and prj_eu arrays, are fully restored on exit even if an exception is raised.

Parameters:

• avgr (Averager) – Configured Averager instance. rec_halfsets will be temporarily forced to True; verbosity will be silenced.

• ptcls (Particles) – Particle stack. The half_id array (and, when perturbation is enabled, ali_eu / prj_eu) is temporarily modified and restored on exit.

• tomo_filename (str) – Path to the .tomostxt tomograms file.

• num_entries (int, optional) – Maximum number of new pairs to add in this call. None fills the buffer to capacity. Default: None.

• sigma_ang_3D (float or (float, float), optional) – Std. dev. (in degrees) of the isotropic rotation applied to the half-1 particles’ 3-D orientations (ali_eu). Models 3-D alignment error. A scalar uses a fixed sigma for every entry; a (lo, hi) pair sweeps the range across the entries generated in this call (evenly spread and shuffled), which trains the network to be robust over a range of blur levels rather than a single one. Include lo == 0 to also see cleanly-aligned pairs and avoid over-sharpening at inference. 0.0 disables it. Default: 0.0.

• sigma_ang_2D (float or (float, float), optional) – Same as sigma_ang_3D, but applied independently to each half-1 projection orientation (prj_eu). Models per-tilt / tilt-series alignment error. 0.0 disables it. Default: 0.0.

• sigma_def (float or (float, float), optional) – Std. dev. (in Angstroms) of an additive Gaussian shift applied independently to each half-1 projection’s defocus. The same shift is added to def_U and def_V so the mean defocus moves while the astigmatism is preserved. Models defocus-estimation error, which degrades CTF correction. Scalar or (lo, hi) range, same semantics as sigma_ang_3D. 0.0 disables it. Default: 0.0.

Raises:

RuntimeError – If the buffer has not been initialised; call set_size() or set_size_mask() first.

Examples

Plain Noise2Noise denoising pairs (both halves cleanly aligned):

>>> vp = VolumePairs()
>>> vp.set_size_mask(box_size, num_vol=50, mask=mask)
>>> vp.populate(avgr, ptcls, 'tomos.tomostxt')
>>> trainer.train(vp, n_epochs=100)          # symmetric, denoising only

Misalignment-recovery pairs: perturb the half-1 orientations so slot 0 is a misaligned reconstruction and slot 1 is cleanly aligned, then train directionally (slot 0 -> slot 1) so the network learns to undo the alignment blur while still denoising:

>>> vp.populate(avgr, ptcls, 'tomos.tomostxt', sigma_ang_3D=3.0)
>>> trainer.train(vp, n_epochs=100, directional=True)

For robustness across blur levels (recommended), pass a (lo, hi) range including 0 instead of a single sigma, so the buffer spans cleanly-aligned to strongly-misaligned pairs:

>>> vp.populate(avgr, ptcls, 'tomos.tomostxt', sigma_ang_3D=(0.0, 4.0))
>>> trainer.train(vp, n_epochs=100, directional=True)

sigma_ang_3D is the std. dev. of the applied rotation magnitude in degrees, so the typical geodesic error is roughly 0.8 * sigma (about 2.3 deg for sigma_ang_3D=3.0). Add sigma_ang_2D to also perturb each projection orientation independently (per-tilt alignment error), or sigma_def (Angstroms) to perturb the per-projection defocus and model CTF-correction error. Enable and validate these one at a time rather than all at once.

reset()

Clear the buffer without reallocating memory.

Sets num_vol to 0, resets the write head to slot 0, and zeroes the internal tensor. Subsequent calls to push() or populate() will refill from the beginning.

Utilities

crop_vol(vol: ndarray) → ndarray

Crop a volume to the active region set by set_size() or set_size_mask().

When apply_mask is True and a mask has been stored, the cropped volume is additionally multiplied by mask (floored at min_mask_value) before being returned — so volumes are masked as push() writes them into the buffer.

Parameters:

vol (numpy.ndarray) – Input volume of shape (box_size, box_size, box_size).

Returns:

Cropped sub-volume, mask-multiplied when apply_mask is set.

Return type:

numpy.ndarray

pad_vol(vol: ndarray) → ndarray

Zero-pad a cropped volume back to the full box size.

Inverse of crop_vol().

Parameters:

vol (numpy.ndarray) – Cropped volume as returned by crop_vol().

Returns:

Full-size volume with the cropped region placed at the correct position and zeros elsewhere.

Return type:

numpy.ndarray