Alignment Model

An alignment model defines the protocols for pre-transformation and alignment scoring.

Alignment Workflows 

There are three types of alignment workflows.

Single template

If the alignment model is created with a single template image, the workflow is very simple. It masks and pre-transforms both the sub-volume and the template, and then aligns them.

        graph LR
    vol(sub-volume\nat <i>i</i>-th molecule)
    tmp[[template image]]
    vol_t(transformed\nsub volume)
    tmp_t[[transformed\ntemplate]]
    aln{alignment}
    result[alignment results]

    vol--masking &<br>pre-transformation-->vol_t-->aln
    tmp--masking &<br>pre-transformation-->tmp_t-->aln
    aln-->result

Multiple templates

If the alignment model is created with multiple template images, masking, pre-transformation and alignment are performed for each template separately. The alignment result with the best score will be considered as the optimal result.

        graph LR

    subgraph Subvolumes
        vol(sub-volume\nat <i>i</i>-th molecule)
        vol_t(transformed\nsub volume)
        vol--masking &<br>pre-transformation-->vol_t
    end

    subgraph Templates
        tmp0[[template image A]]
        tmp1[[template image B]]
        tmp2[[template image C]]
        tmp0_t[[transformed\ntemplate A]]
        tmp1_t[[transformed\ntemplate B]]
        tmp2_t[[transformed\ntemplate C]]
        tmp0--masking &<br>pre-transformation-->tmp0_t
        tmp1--masking &<br>pre-transformation-->tmp1_t
        tmp2--masking &<br>pre-transformation-->tmp2_t
    end

    subgraph Alignment
        aln0{alignment}
        aln1{alignment}
        aln2{alignment}
    end
    result[best alignment results]

    vol_t-->aln0
    vol_t-->aln1
    vol_t-->aln2

    tmp0_t-->aln0
    tmp1_t-->aln1
    tmp2_t-->aln2

    aln0-->result
    aln1-->result
    aln2-->result

Single template with rotation

Many alignment methods do not search for the optimal rotation of the template image. In this case, rotated templates will be generated and used for alignment. Essentially, it is the same as the multiple-template workflow.

        graph LR

    subgraph Subvolumes
        vol(sub-volume\nat <i>i</i>-th molecule)
        vol_t(transformed\nsub volume)
        vol--masking &<br>pre-transformation-->vol_t
    end

    subgraph Templates
        tmp[[template image]]
        tmp_t[[transformed\ntemplate]]
        tmp--masking &<br>pre-transformation-->tmp_t

        tmp0[[template image A]]
        tmp1[[template image B]]
        tmp2[[template image C]]
        rot{image rotation}
        tmp_t-->rot
        rot-->tmp0
        rot-->tmp1
        rot-->tmp2
    end

    subgraph Alignment
        aln0{alignment}
        aln1{alignment}
        aln2{alignment}
    end
    result[best alignment results]

    vol_t-->aln0
    vol_t-->aln1
    vol_t-->aln2

    tmp0-->aln0
    tmp1-->aln1
    tmp2-->aln2

    aln0-->result
    aln1-->result
    aln2-->result

Ready-to-use Models 

There are now two alignment models that can be used directly.

ZNCCAlignment
Model that align subvolumes using ZNCC (Zero-mean Normalized Cross Correlation) score.
PCCAlignment
Model that align subvolumes using PCC (Phase Cross Correlation) score.

Both models are implemented with low-pass filtering, template rotation and missing wedge masking.

Model construction

from acryo.alignment import ZNCCAlignment

model = ZNCCAlignment(
    template,  # template image
    mask,      # mask image
    rotations=[(10, 5), (4, 2), (8, 4)],
    cutoff=0.5,
    tilt=(-60, 60),
)

Shape of mask must be the same as template. template * mask and subvolume * mask will be used for alignment.
rotations can be three tuples or a scipy.spatial.transform.Rotation object.
- If three tuples are given, each tuple defines the maximum rotation angle and the increment around z, y or x (external) axis. The unit is degree. For example, the first (10, 5) means that the rotation angles -10, -5, 0, 5, 10 will be used for the rotation around z axis.
- If a scipy.spatial.transform.Rotation object is given, all the rotations in the object will be used for alignment. Make sure that the identity rotation is included.
cutoff is the relative cutoff frequency for low-pass filtering. The Nyquist frequency is \(0.5 \times \sqrt{3} = 0.866\) for 3D images.
tilt is the tilt series model.
- If a (float, float) tuple is given, it will be interpreted as the minimum and maximum tilt angle in degree. The rotation axis is assumed to be the y axis.
- You can use classes provided in acryo.tilt to define more complex tilt series models. For instance, single_axis() for single-axis tilt series and dual_axis() for dual-axis tilt series.

Align images

The align() method is used to align a sub-volume to the template image of the model. Note that this method does not actually transform the sub-volume to the template. It only calculate the optimal shift/rotation parameters. To transform the sub-volume, use fit().

result = model.align(
    img,
    max_shifts,
    quaternion,
    pos,
    backend,
)

img is the sub-volume to be aligned. It must be a 3D array with the same shape as the template.
max_shifts is a tuple of maximum shifts in z, y and x direction. The unit is pixel but it can be a float number.
quaternion is the rotation of the sub-volume in the original tomogram. It must be a (4,) numpy.ndarray object of quaternion. If you are using acryo.Molecules, its quaternions can directly be used here. This is basically used to mask the missing wedge.
pos is the position of the sub-volume in the original tomogram. It must be a (3,) numpy.ndarray object. Default alignment models does not use this parameter.
backend is the array API backend. It can be "numpy" or "cupy".

The return value result is a named-tuple AlignmentResult object. It contains the following fields.

class AlignmentResult(NamedTuple):
    label: int
    shift: NDArray[np.float32]
    quat: NDArray[np.float32]
    score: float

label is the integer label of the best alignment if multiple templates are used.
shift is the optimal shift in z, y and x direction.
quat is the optimal rotation in quaternion.
score is the alignment score of the best alignment.

Fit images

The fit() method is used to transform the sub-volume to fit the template image of the model. It is essentially the same as calling align() for every rotation and then Affine transform the sub-volume to the best alignment result, but fit() is faster because it parallelizes the rotation and alignment processes.

result = model.fit(
    img,
    max_shifts,
    cval=0.0,
    backend=None,
)

img and max_shifts is the same as align().
cval is the constant value used for Affine transformations. 1% percentile will be used by default.

Correlation landscape

The word “correlation landscape” came from “energy landscape” in the context of protein folding. It is a 3D array of the correlation scores between the sub-volume and the template image.

arr = model.landscape(
    img,
    max_shifts: tuple[float, float, float],
    quaternion: NDArray[np.float32] | None = None,
    pos: NDArray[np.float32] | None = None,
    upsample: int = 1,
    backend: Backend | None = None,
)

img is the sub-volume to be aligned. It must be a 3D array with the same shape as the template.
max_shifts is a tuple of maximum shifts in z, y and x direction. The unit is pixel but it can be a float number.
quaternion is the rotation of the sub-volume in the original tomogram. It must be a (4,) numpy.ndarray object of quaternion. If you are using acryo.Molecules, its quaternions can directly be used here. This is basically used to mask the missing wedge.
pos is the position of the sub-volume in the original tomogram. It must be a (3,) numpy.ndarray object. Default alignment models does not use this parameter.

Define Custom Alignment Model 

In acryo.alignment, there are several abstract base classes that can be used to efficiently define custom alignment models.

BaseAlignmentModel … The most basic one that provides the minimum interface. Need to override _optimize() and pre_transform().
RotationImplemented … Rotated templates will be generated even if the optimization algorithm does not optimize the rotation. Need to override _optimize() and pre_transform().
TomographyInput … Rotation, low-pass filtering and missing wedge masking is already implemented. Only need to override _optimize().

When you override methods, the following should be noted.

pre_transform()
This method must have the following signature.
def pre_transform( self, image: NDArray[np.float32], backend=None, ) -> NDArray[np.complex64]: ...
The input image could be either the sub-volume or the template image. It is masked by the input mask image but is not masked by the missing wedge mask in TomographyInput. The output image will be directly passed to the _optimize() method, so the data type depends on the implementation. backend is the backend array API object. You don’t have to use it unless you want to implement same method for GPU.
_optimize()
This method must have the following signature.
def _optimize( self, subvolume: NDArray[T], template: NDArray[T], max_shifts: tuple[float, float, float], quaternion: NDArray[np.float32], pos: NDArray[np.float32], backend=None, ) -> tuple[NDArray[np.float32], NDArray[np.float32], float]: ...
This method is called for every set of sub-volume and template images.
- subvolume and template is the images after pre-transformation. Thus, they could be Fourier transformed.
- max_shift is directly passed from align() or fit() method.
- quaternion is the rotation of the sub-volume. This parameter can be used to mask the missing wedge.
- pos is the position of the sub-volume in the original tomogram. Its unit is pixel. This parameter can be used for CTF correction of defocusing.
- backend is the backend array API object. You don’t have to use it unless you want to implement same method for GPU.
- The return value must be a tuple of (shift, rotation, score).
  - shift is the optimal shift in z, y and x direction. More precisely, ndi.shift(img, -shift) will properly align the image to the template.
  - rotation is the optimal rotation in quaternion. If the alignment model does not optimize the rotation, this value should be array([0, 0, 0, 1]).
  - score is the score of the alignment. Larger score means better alignment.