Rotation conventions agree at runtime; no ambiguity in rotation convention

Author: michael-0brien

cryojax/jax_util/_errors.py

@@ -2,12 +2,12 @@
 from typing import TypeVar
 
 import equinox as eqx
-from jaxtyping import ArrayLike, Bool, PyTree
+from jaxtyping import ArrayLike, Bool
 
 from .._config import CRYOJAX_ENABLE_CHECKS
 
 
-T = TypeVar("T", bound="PyTree")
+T = TypeVar("T")
 
 
 def maybe_error_if(x: T, pred_fn: Callable[[T], Bool[ArrayLike, "..."]], msg: str) -> T:

cryojax/simulator/__init__.py

@@ -73,6 +73,7 @@
     IndependentAtomVolume as IndependentAtomVolume,
     LobatoScatteringFactor as LobatoScatteringFactor,
     PengScatteringFactor as PengScatteringFactor,
+    RealVoxelCloudVolume as RealVoxelCloudVolume,
     RealVoxelGridVolume as RealVoxelGridVolume,
     RealVoxelProjection as RealVoxelProjection,
 )

cryojax/simulator/_api_utils.py

@@ -2,7 +2,6 @@
 from collections.abc import Callable
 from typing import Any, Literal, overload
 
-import equinox as eqx
 import equinox.internal as eqxi
 import jax.numpy as jnp
 import mmdf
@@ -42,28 +41,14 @@
     FourierVoxelSplineVolume,
     GaussianMixtureVolume,
     IndependentAtomVolume,
+    RealVoxelCloudVolume,
     RealVoxelGridVolume,
 )
 
 
 identity_fn = eqxi.doc_repr(lambda x, _: x, "identity_fn")
 
 
-def _maybe_invert_rotation(
-    pose: AbstractPose,
-    volume: AbstractVolumeParametrization,
-    rotation_convention: Literal["object", "frame"],
-) -> AbstractPose:
-    jaxpr_fn = eqx.filter_make_jaxpr(lambda vol: vol.to_representation())
-    _, out_dynamic, out_static = jaxpr_fn(volume)
-    out_struct = eqx.combine(out_dynamic, out_static)
-    conventions = [rotation_convention, out_struct.rotation_convention]
-    if conventions[0] != conventions[1]:
-        pose = pose.to_inverse_rotation()
-
-    return pose
-
-
 @overload
 def make_image_model(
     volume: AbstractVolumeParametrization,
@@ -79,7 +64,6 @@ def make_image_model(
     signal_centering: Literal["bg", "mean"] = "mean",
     translate_mode: Literal["fft", "atom", "none"] = "fft",
     quantity_mode: Literal["none"] = "none",
-    rotation_convention: Literal["object", "frame"] = "object",
 ) -> ProjectionImageModel: ...
 
 
@@ -98,7 +82,6 @@ def make_image_model(  # pyright: ignore[reportOverlappingOverload]
     signal_centering: Literal["bg", "mean"] = "mean",
     translate_mode: Literal["fft", "atom", "none"] = "fft",
     quantity_mode: Literal["none"] = "none",
-    rotation_convention: Literal["object", "frame"] = "object",
 ) -> LinearImageModel: ...
 
 
@@ -117,7 +100,6 @@ def make_image_model(
     signal_centering: Literal["bg", "mean"] = "mean",
     translate_mode: Literal["fft", "atom", "none"] = "fft",
     quantity_mode: Literal["contrast"] = "contrast",
-    rotation_convention: Literal["object", "frame"] = "object",
 ) -> ContrastImageModel: ...
 
 
@@ -136,7 +118,6 @@ def make_image_model(
     signal_centering: Literal["bg", "mean"] = "mean",
     translate_mode: Literal["fft", "atom", "none"] = "fft",
     quantity_mode: Literal["intensity"] = "intensity",
-    rotation_convention: Literal["object", "frame"] = "object",
 ) -> IntensityImageModel: ...
 
 
@@ -155,7 +136,6 @@ def make_image_model(
     signal_centering: Literal["bg", "mean"] = "mean",
     translate_mode: Literal["fft", "atom", "none"] = "fft",
     quantity_mode: Literal["counts"] = "counts",
-    rotation_convention: Literal["object", "frame"] = "object",
 ) -> ElectronCountsImageModel: ...
 
 
@@ -173,7 +153,6 @@ def make_image_model(
     signal_centering: Literal["bg", "mean"] = "mean",
     translate_mode: Literal["fft", "atom", "none"] = "fft",
     quantity_mode: Literal["contrast", "intensity", "counts", "none"] = "none",
-    rotation_convention: Literal["object", "frame"] = "object",
 ) -> AbstractImageModel:
     """Construct an [`cryojax.simulator.AbstractImageModel`][] for
     most common use-cases.
@@ -281,20 +260,6 @@ def make_image_model(
             Uses the [`cryojax.simulator.ElectronCountsImageModel`][]
             to simulate electron counts.
             If this is passed, a `detector` must also be passed.
-    - `rotation_convention`:
-        If `'object'`, the rotation given by `pose` is of the object.
-        If `'frame'`, the rotation given by `pose` is of the frame. These
-        are related by transpose.
-
-    !!! info
-        The `make_image_model` function enforces agreement between
-        rotation conventions of different volumes via the
-        `rotation_convention` argument. Lower level `cryojax` APIs
-        will not enforce this agreement, such as if the user instantiates
-        an [`cryojax.simulator.AbstractImageModel`][] directly.
-
-        In these cases, agreement can be acheived with a manual transpose
-        via `pose.to_inverse_rotation()`.
 
     **Returns:**
 
@@ -309,13 +274,6 @@ def make_image_model(
     [`cryojax.simulator.ElectronCountsImageModel`][] depending on
     the value of `quantity_mode`.
     """
-    # Invert pose if
-    if rotation_convention not in ["object", "frame"]:
-        raise ValueError(
-            f"Found `rotation_convention = {rotation_convention}`, but valid "
-            "values are 'object' and 'frame'."
-        )
-    pose = _maybe_invert_rotation(pose, volume, rotation_convention)
     options = dict(
         normalizes_signal=normalizes_signal,
         signal_centering=signal_centering,
@@ -581,14 +539,31 @@ def render_voxel_volume(
 ) -> RealVoxelGridVolume: ...
 
 
+@overload
+def render_voxel_volume(
+    atom_volume: AbstractAtomVolume,
+    render_fn: AbstractVolumeRenderFn,
+    *,
+    output_type: type[RealVoxelCloudVolume] = RealVoxelCloudVolume,
+) -> RealVoxelCloudVolume: ...
+
+
 def render_voxel_volume(
     atom_volume: AbstractAtomVolume,
     render_fn: AbstractVolumeRenderFn,
     *,
     output_type: type[
-        FourierVoxelGridVolume | FourierVoxelSplineVolume | RealVoxelGridVolume
+        FourierVoxelGridVolume
+        | FourierVoxelSplineVolume
+        | RealVoxelGridVolume
+        | RealVoxelCloudVolume
     ] = FourierVoxelGridVolume,
-) -> FourierVoxelGridVolume | FourierVoxelSplineVolume | RealVoxelGridVolume:
+) -> (
+    FourierVoxelGridVolume
+    | FourierVoxelSplineVolume
+    | RealVoxelGridVolume
+    | RealVoxelCloudVolume
+):
     """Render a voxel volume representation from an atomistic one.
 
     !!! example "Simulate an image with Fourier slice extraction"
@@ -624,7 +599,8 @@ def render_voxel_volume(
         Either [`cryojax.simulator.FourierVoxelGridVolume`][] /
         [`cryojax.simulator.FourierVoxelSplineVolume`][] for
         fourier-space representations, or
-        [`cryojax.simulator.RealVoxelGridVolume`][] for real-space.
+        [`cryojax.simulator.RealVoxelGridVolume`][] /
+        [`cryojax.simulator.RealVoxelCloudVolume`][] for real-space.
 
 
     **Returns:**
@@ -650,14 +626,20 @@ def render_voxel_volume(
         else:
             spline_coefficients = compute_spline_coefficients(fourier_voxel_grid)
             return FourierVoxelSplineVolume(spline_coefficients, frequency_slice)
-    elif output_type == RealVoxelGridVolume:
+    elif output_type == RealVoxelGridVolume or output_type == RealVoxelCloudVolume:
         coordinate_grid = make_coordinate_grid(render_fn.shape)
         real_voxel_grid = render_fn(atom_volume, outputs_real_space=True)
-        return RealVoxelGridVolume(real_voxel_grid, coordinate_grid)
+        if output_type == RealVoxelGridVolume:
+            return RealVoxelGridVolume(real_voxel_grid, coordinate_grid)
+        else:
+            return RealVoxelCloudVolume.from_real_voxel_grid(
+                real_voxel_grid, coordinate_grid_in_pixels=coordinate_grid
+            )
     else:
         raise ValueError(
-            "Only `output_type` equal to `FourierVoxelGridVolume`, "
-            "`FourierVoxelSplineVolume`, or `RealVoxelGridVolume` "
-            "are supported."
-            f"Got `output_type = {output_type}`."
+            f"Got `output_type = {output_type}`, but this is "
+            "not supported by `render_voxel_volume(..., output_type=...)`."
+            "Valid values for `output_type` are `FourierVoxelGridVolume`, "
+            "`FourierVoxelSplineVolume`, `RealVoxelGridVolume`, or "
+            "`RealVoxelCloudVolume`."
         )

cryojax/simulator/_volume/__init__.py

@@ -29,6 +29,7 @@
     PengScatteringFactor as PengScatteringFactor,
 )
 from .real_voxels import (
+    RealVoxelCloudVolume as RealVoxelCloudVolume,
     RealVoxelGridVolume as RealVoxelGridVolume,
     RealVoxelProjection as RealVoxelProjection,
 )

cryojax/simulator/_volume/auto_select.py

@@ -28,7 +28,7 @@
     FFTAtomRenderFn,
     IndependentAtomVolume,
 )
-from .real_voxels import RealVoxelGridVolume, RealVoxelProjection
+from .real_voxels import RealVoxelCloudVolume, RealVoxelProjection
 
 
 class AutoVolumeProjection(
@@ -47,10 +47,13 @@ class AutoVolumeProjection(
         | [`cryojax.simulator.GaussianMixtureVolume`][] | [`cryojax.simulator.GaussianMixtureProjection`][] | atom |
         | [`cryojax.simulator.IndependentAtomVolume`][] | [`cryojax.simulator.FFTAtomProjection`][] | atom |
         | [`cryojax.simulator.FourierVoxelGridVolume`][] or [`cryojax.simulator.FourierVoxelSplineVolume`][] | [`cryojax.simulator.FourierSliceExtraction`][] | voxel |
-        | [`cryojax.simulator.RealVoxelGridVolume`][] | [`cryojax.simulator.RealVoxelProjection`][] | voxel |
+        | [`cryojax.simulator.RealVoxelCloudVolume`][] | [`cryojax.simulator.RealVoxelProjection`][] | voxel |
+
+        Note that [`cryojax.simulator.RealVoxelGridVolume`][] does not have an associated projection method. To
+        compute projections from real-space voxels, use [`cryojax.simulator.RealVoxelCloudVolume`][].
 
         To use advanced options for a given projection method,
-        see each respective class.
+        instantiate each respective class directly.
 
     !!! warning
         If using [`cryojax.simulator.FFTAtomRenderFn`][] or [`cryojax.simulator.RealVoxelProjection`][], [`jax-finufft`](https://github.com/flatironinstitute/jax-finufft)
@@ -67,14 +70,15 @@ def _select_projection_method(
             integrator = FourierSliceExtraction()
         elif isinstance(volume, GaussianMixtureVolume):
             integrator = GaussianMixtureProjection()
-        elif isinstance(volume, RealVoxelGridVolume):
+        elif isinstance(volume, RealVoxelCloudVolume):
             integrator = RealVoxelProjection()
         elif isinstance(volume, IndependentAtomVolume):
             integrator = FFTAtomProjection()
         else:
             raise ValueError(
                 "Could not use `AutoVolumeProjection` for volume of "
-                f"type {type(volume).__name__}. If using a custom volume, "
+                f"type {type(volume).__name__}. See the documentation for "
+                "supported types. If using a custom volume, "
                 "please directly pass an integrator."
             )
         return integrator

cryojax/simulator/_volume/base_volume.py

@@ -1,5 +1,5 @@
 import abc
-from typing import Generic, Literal, TypeVar
+from typing import Generic, TypeVar
 from typing_extensions import Self, override
 
 import equinox as eqx
@@ -118,10 +118,10 @@ class AbstractVolumeRepresentation(AbstractVolumeParametrization, strict=True):
     classes for imaging.
     """
 
-    rotation_convention: eqx.AbstractClassVar[Literal["object", "frame"]]
+    is_frame_rotation: eqx.AbstractClassVar[bool]
 
     @abc.abstractmethod
-    def rotate_to_pose(self, pose: AbstractPose, inverse: bool = False) -> Self:
+    def rotate_to_pose(self, pose: AbstractPose) -> Self:
         """Rotate the coordinate system of the volume."""
         raise NotImplementedError
 
@@ -174,7 +174,7 @@ class AbstractVoxelVolume(AbstractVolumeRepresentation, strict=True):
     @property
     @abc.abstractmethod
     def shape(self) -> tuple[int, ...]:
-        """The shape of the voxel array."""
+        """The shape of the voxel grid."""
         raise NotImplementedError
 
     @classmethod

cryojax/simulator/_volume/fourier_voxels.py

@@ -2,7 +2,7 @@
 Fourier voxel-based representations of a volume.
 """
 
-from typing import ClassVar, Literal, cast
+from typing import ClassVar, cast
 from typing_extensions import Self, override
 
 import equinox as eqx
@@ -41,12 +41,14 @@ class AbstractFourierVoxelVolume(AbstractVoxelVolume, strict=True):
     frequency_slice_in_pixels: eqx.AbstractVar[Float[Array, "1 dim dim 3"]]
 
     @override
-    def rotate_to_pose(self, pose: AbstractPose, inverse: bool = False) -> Self:
+    def rotate_to_pose(self, pose: AbstractPose) -> Self:
         """Return a new volume with a rotated `frequency_slice_in_pixels`."""
         return eqx.tree_at(
             lambda d: d.frequency_slice_in_pixels,
             self,
-            pose.rotate_coordinates(self.frequency_slice_in_pixels, inverse=inverse),
+            pose.rotate_coordinates(
+                self.frequency_slice_in_pixels, inverse=self.is_frame_rotation
+            ),
         )
 
 
@@ -56,7 +58,7 @@ class FourierVoxelGridVolume(AbstractFourierVoxelVolume, strict=True):
     fourier_voxel_grid: Complex[Array, "dim dim dim"]
     frequency_slice_in_pixels: Float[Array, "1 dim dim 3"]
 
-    rotation_convention: ClassVar[Literal["frame"]] = "frame"
+    is_frame_rotation: ClassVar[bool] = True
 
     def __init__(
         self,
@@ -146,7 +148,7 @@ class FourierVoxelSplineVolume(AbstractFourierVoxelVolume, strict=True):
     spline_coefficients: Complex[Array, "coeff_dim coeff_dim coeff_dim"]
     frequency_slice_in_pixels: Float[Array, "1 dim dim 3"]
 
-    rotation_convention: ClassVar[Literal["frame"]] = "frame"
+    is_frame_rotation: ClassVar[bool] = True
 
     def __init__(
         self,
@@ -317,8 +319,10 @@ def integrate(
             )
         else:
             raise ValueError(
-                "Supported types for `volume_representation` are "
-                "`FourierVoxelGridVolume` and FourierVoxelSplineVolume`."
+                "Got unsupported type for `volume_representation` in "
+                "`FourierSliceExtraction.integrate`. Expected `FourierVoxelGridVolume` "
+                "or `FourierVoxelSplineVolume`, "
+                f"but got `{volume_representation.__class__.__name__}`."
             )
 
         # Resize the image to match the AbstractImageConfig.padded_shape
@@ -435,8 +439,10 @@ def integrate(
             )
         else:
             raise ValueError(
-                "Supported types for `volume_representation` are "
-                "`FourierVoxelGridVolume` and `FourierVoxelSplineVolume`."
+                "Got unsupported type for `volume_representation` in "
+                "`EwaldSphereExtraction.integrate`. Expected `FourierVoxelGridVolume` "
+                "or `FourierVoxelSplineVolume`, "
+                f"but got `{volume_representation.__class__.__name__}`."
             )
 
         # Resize the image to match the AbstractImageConfig.padded_shape

cryojax/simulator/_volume/gaussian_volume.py

@@ -65,7 +65,7 @@ class GaussianMixtureVolume(AbstractAtomVolume, strict=True):
     amplitudes: Float[Array, "n_positions n_gaussians"]
     variances: Float[Array, " n_positions n_gaussians"]
 
-    rotation_convention: ClassVar[Literal["object"]] = "object"
+    is_frame_rotation: ClassVar[bool] = False
 
     def __init__(
         self,
@@ -197,12 +197,12 @@ def from_tabulated_parameters(
         return cls(atom_positions, amplitudes, b_factor_to_variance(b_factors))
 
     @override
-    def rotate_to_pose(self, pose: AbstractPose, inverse: bool = False) -> Self:
+    def rotate_to_pose(self, pose: AbstractPose) -> Self:
         """Return a new potential with rotated `positions`."""
         return eqx.tree_at(
             lambda d: d.positions,
             self,
-            pose.rotate_coordinates(self.positions, inverse=inverse),
+            pose.rotate_coordinates(self.positions, inverse=self.is_frame_rotation),
         )
 
     @override