Add a ProjectionModel specification

README.md

@@ -8,6 +8,7 @@
 
 `cryotypes` defines a set of super-simple, extensible data structures for the fundamental types of cryoEM data and their relative metadata:
 - `PoseSet`: a set of particle poses, compatible with 2D and 3D data
+- `ProjectionModel`: a set of parameters for a projection model (tilt-series alignments)
 - `Tomogram`: a 3D image
 - `Micrograph`: a 2D image
 
@@ -49,6 +50,43 @@ Particle orientations are stored as
 [`scipy.spatial.transform.Rotation`](https://docs.scipy.org/doc/scipy/reference/generated/scipy.spatial.transform.Rotation.html) objects.
 These transformations should rotate the basis vectors of a reference such that they are correctly oriented in a tomogram.
 
+## `ProjectionModel`
+A `ProjectionModel` is a [pandas `DataFrame`](https://pandas.pydata.org/docs/) with specific column 
+headings for the parameters of a projection model. Together, this information constitues a 'tilt-series alignment'.
+
+| Heading         | Python name   | Semantics                                         |
+|:----------------|:--------------|:--------------------------------------------------|
+| `rotation_x`    | ROTATION_X    | specimen rotation around x-axis                   |
+| `rotation_y`    | ROTATION_Y    | specimen rotation around y-axis                   |
+| `rotation_z`    | ROTATION_Z    | specimen rotation around z-axis                   |
+| `dx`            | SHIFT_X       | specimen shift in x-dimension of the camera plane |
+| `dy`            | SHIFT_Y       | particle shift in y-dimension of the camera plane |
+| `experiment_id` | EXPERIMENT_ID | identifier for micrograph/tilt-series             |
+| `pixel_spacing` | PIXEL_SPACING | isotropic pixel/voxel spacing for shifts          |
+| `source`        | SOURCE        | reference to the file from which data came        |
+
+In the microsope reference frame, the z-axis is the beam direction.
+Extrinsic rotation of the tomogram around the x-axis, the y-axis, then the z-axis by 
+`rotation_x`, `rotation_y`, `rotation_z` followed by projection along the z-axis (beam direction)
+then shifting the 2D image in the camera plane by `dx` and `dy` produces the experimental projection
+image.
+
+A utility function is also provided for generating projection matrices from these data.
+These projection matrices can be used to calculate a 2D position in a tilt-image from a 3D position
+in the tomogram.
+
+```python
+from cryotypes.projectionmodel import projection_model_to_projection_matrices
+
+projection_matrices = projection_model_to_projection_matrices(
+    df=projection_model,  # ProjectionModel dataframe
+    tilt_image_center=(1919, 1355),  # tilt-image rotation center (xy)
+    tomogram_dimensions=(3838, 3710, 2000)  # dimensions of tomogram (xyz)
+)
+```
+
+**Note:** these projection matrices are only valid for positions in a tomogram of the dimensions
+provided in this function and must be recalculated for different tomogram dimensions.
 
 ## `Tomogram`
 A `Tomogram` is an object that follows a specific [python `Protocol`](https://docs.python.org/3/library/typing.html#typing.Protocol) for tomogram data. The protocol specifies the following attributes:

cryotypes/projectionmodel/__init__.py

@@ -0,0 +1,9 @@
+from ._data_labels import ProjectionModelDataLabels
+from ._typing import ProjectionModel
+from ._utils import projection_model_to_projection_matrices
+
+__all__ = [
+    "ProjectionModelDataLabels",
+    "ProjectionModel",
+    "projection_model_to_projection_matrices"
+]

cryotypes/projectionmodel/_data_labels.py

@@ -0,0 +1,15 @@
+class ProjectionModelDataLabels:
+    ROTATION_X = "rotation_x"
+    ROTATION_Y = "rotation_y"
+    ROTATION_Z = "rotation_z"
+    ROTATION = [ROTATION_X, ROTATION_Y, ROTATION_Z]
+
+    SHIFT_X = "dx"
+    SHIFT_Y = "dy"
+    SHIFT = [SHIFT_X, SHIFT_Y]
+
+    EXPERIMENT_ID = "experiment_id"
+
+    PIXEL_SPACING = "pixel_spacing"
+
+    SOURCE = "source"

cryotypes/projectionmodel/_transformations.py

@@ -0,0 +1,116 @@
+import numpy as np
+import einops
+
+
+def Rx(theta: np.ndarray) -> np.ndarray:
+    """Generate 4x4 matrices which rotate vectors around the X-axis by the angle `theta`.
+
+    The rotation angle `theta` is expected to be in degrees.
+    Application of these matrices is performed by left-multiplying xyzw column vectors.
+
+    Parameters
+    ----------
+    theta: np.ndarray
+        An `(n, )` array of rotation angles in degrees.
+
+    Returns
+    -------
+    matrices: np.ndarray
+        An `(n, 4, 4)` array of matrices which left-multiply xyzw column vectors.
+    """
+    theta = np.asarray(theta).reshape(-1)
+    c = np.cos(np.deg2rad(theta))
+    s = np.sin(np.deg2rad(theta))
+    matrices = einops.repeat(
+        np.eye(4), 'i j -> n i j', n=len(theta)
+    )
+    matrices[:, 1, 1] = c
+    matrices[:, 1, 2] = -s
+    matrices[:, 2, 1] = s
+    matrices[:, 2, 2] = c
+    return matrices
+
+
+def Ry(theta: np.ndarray) -> np.ndarray:
+    """Generate 4x4 matrices which rotate vectors around the Y-axis by the angle `theta`.
+
+    The rotation angle `theta` is expected to be in degrees.
+    Application of these matrices is performed by left-multiplying xyzw column vectors.
+
+    Parameters
+    ----------
+    theta: np.ndarray
+        An `(n, )` array of rotation angles in degrees.
+
+    Returns
+    -------
+    matrices: np.ndarray
+        An `(n, 4, 4)` array of matrices which left-multiply xyzw column vectors."""
+    theta = np.asarray(theta).reshape(-1)
+    c = np.cos(np.deg2rad(theta))
+    s = np.sin(np.deg2rad(theta))
+    matrices = einops.repeat(np.eye(4), 'i j -> n i j', n=len(theta))
+    matrices[:, 0, 0] = c
+    matrices[:, 0, 2] = s
+    matrices[:, 2, 0] = -s
+    matrices[:, 2, 2] = c
+    return matrices
+
+
+def Rz(theta: float) -> np.ndarray:
+    """Generate 4x4 matrices which rotate vectors around the Z-axis by the angle `theta`.
+
+    The rotation angle `theta` is expected to be in degrees.
+    Application of these matrices is performed by left-multiplying xyzw column vectors.
+
+    Parameters
+    ----------
+    theta: np.ndarray
+        An `(n, )` array of rotation angles in degrees.
+
+    Returns
+    -------
+    matrices: np.ndarray
+        An `(n, 4, 4)` array of matrices which left-multiply xyzw column vectors."""
+    theta = np.asarray(theta).reshape(-1)
+    c = np.cos(np.deg2rad(theta))
+    s = np.sin(np.deg2rad(theta))
+    matrices = einops.repeat(
+        np.eye(4), 'i j -> n i j', n=len(theta)
+    )
+    matrices[:, 0, 0] = c
+    matrices[:, 0, 1] = -s
+    matrices[:, 1, 0] = s
+    matrices[:, 1, 1] = c
+    return np.squeeze(matrices)
+
+
+def S(shifts: np.ndarray) -> np.ndarray:
+    """Generate 4x4 matrices for 2D (xy) or 3D (xyz) shifts.
+
+    Application of these matrices is performed by left-multiplying xyzw column vectors.
+
+    Parameters
+    ----------
+    shifts: np.ndarray
+        An `(n, 2)` or `(n, 3)` array of xy(z) shifts.
+
+    Returns
+    -------
+    matrices: np.ndarray
+        An `(n, 4, 4)` array of
+    """
+    shifts = np.asarray(shifts, dtype=float)
+    if shifts.shape[-1] == 2:
+        shifts = _promote_2d_to_3d(shifts)
+    shifts = np.array(shifts).reshape((-1, 3))
+    matrices = einops.repeat(np.eye(4), 'i j -> n i j', n=shifts.shape[0])
+    matrices[:, 0:3, 3] = shifts
+    return np.squeeze(matrices)
+
+
+def _promote_2d_to_3d(shifts: np.ndarray) -> np.ndarray:
+    """Promote 2D vectors to 3D with zeros in the last dimension."""
+    shifts = np.asarray(shifts).reshape(-1, 2)
+    shifts = np.c_[shifts, np.zeros(len(shifts))]
+    return np.squeeze(shifts)

cryotypes/projectionmodel/_typing.py

@@ -0,0 +1,4 @@
+import pandas as pd
+from typing_extensions import TypeAlias
+
+ProjectionModel: TypeAlias = pd.DataFrame

cryotypes/projectionmodel/_utils.py

@@ -0,0 +1,42 @@
+from typing import Tuple
+
+import numpy as np
+
+from ._typing import ProjectionModel
+from ._data_labels import ProjectionModelDataLabels as PMDL
+from ._transformations import Rx, Ry, Rz, S
+
+
+def projection_model_to_projection_matrices(
+        df: ProjectionModel,
+        tilt_image_center: Tuple[int, int],
+        tomogram_dimensions: Tuple[int, int, int]
+) -> np.ndarray:
+    """Calculate 4x4 projection matrices from a ProjectionModel dataframe.
+
+    The resulting 4x4 projection matrices left-multiply xyzw column vectors containing any position
+    in the tomogram to produce the projected position (xyzw) in the camera plane.
+
+
+    Parameters
+    ----------
+    df: ProjectionModel
+        A pandas DataFrame adhering to the `ProjectionModel` specification.
+    tilt_image_center: Tuple[int, int]
+        Rotation center in the 2D tilt-images, ordered xy.
+    tomogram_dimensions: Tuple[int, int, int]
+        dimensions of the tomogram, ordered xyz.
+
+    Returns
+    -------
+    projection_matrices: np.ndarray
+        An `(n, 4, 4)` array of projection matrices which left-multiply xyzw column vectors.
+    """
+    specimen_center = np.array(tomogram_dimensions) // 2
+    s0 = S(-specimen_center)
+    r0 = Rx(df[PMDL.ROTATION_X])
+    r1 = Ry(df[PMDL.ROTATION_Y])
+    r2 = Rz(df[PMDL.ROTATION_Z])
+    s1 = S(df[PMDL.SHIFT])
+    s2 = S(tilt_image_center)
+    return s2 @ s1 @ r2 @ r1 @ r0 @ s0

setup.cfg

@@ -27,6 +27,7 @@ install_requires =
     pandas
     scipy
     typing_extensions
+    einops
 python_requires = >=3.7
 setup_requires =
     setuptools_scm