Skip to content

orientation_search

Serialization and validation of orientation search parameters for 2DTM.

ConstrainedOrientationConfig

Bases: BaseModel2DTM

Serialization and validation of constrained orientation parameters.

Attributes:

Name Type Description
enabled bool

Whether to enable constrained orientation search.
phi_step float

Angular step size for phi in degrees. Must be greater than or equal to 0.
theta_step float

Angular step size for theta in degrees. Must be greater than or equal to 0.
psi_step float

Angular step size for psi in degrees. Must be greater than or equal to 0.
rotation_axis_euler_angles list[float]

List of Euler angles (phi, theta, psi) for the rotation axis.
phi_min float

Minimum value for the phi angle in degrees.
phi_max float

Maximum value for the phi angle in degrees.
theta_min float

Minimum value for the theta angle in degrees.
theta_max float

Maximum value for the theta angle in degrees.
psi_min float

Minimum value for the psi angle in degrees.
psi_max float

Maximum value for the psi angle in degrees.
Source code in src/leopard_em/pydantic_models/config/orientation_search.py 
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
class ConstrainedOrientationConfig(BaseModel2DTM):
    """Serialization and validation of constrained orientation parameters.

    Attributes
    ----------
    enabled: bool
        Whether to enable constrained orientation search.
    phi_step: float
        Angular step size for phi in degrees.
        Must be greater than or equal to 0.
    theta_step: float
        Angular step size for theta in degrees.
        Must be greater than or equal to 0.
    psi_step: float
        Angular step size for psi in degrees.
        Must be greater than or equal to 0.
    rotation_axis_euler_angles: list[float]
        List of Euler angles (phi, theta, psi) for the rotation axis.
    phi_min: float
        Minimum value for the phi angle in degrees.
    phi_max: float
        Maximum value for the phi angle in degrees.
    theta_min: float
        Minimum value for the theta angle in degrees.
    theta_max: float
        Maximum value for the theta angle in degrees.
    psi_min: float
        Minimum value for the psi angle in degrees.
    psi_max: float
        Maximum value for the psi angle in degrees.
    """

    enabled: bool = True
    phi_step: Optional[float] = None
    theta_step: float = 2.5
    psi_step: float = 1.5
    rotation_axis_euler_angles: tuple[float, float, float] = Field(
        default=[0.0, 0.0, 0.0]
    )
    phi_min: float = 0.0
    phi_max: float = 360.0
    theta_min: float = 0.0
    theta_max: float = 180.0
    psi_min: float = 0.0
    psi_max: float = 360.0
    base_grid_method: Literal["uniform", "healpix", "basic", "roll"] = "uniform"

    search_roll_axis: bool = True
    roll_axis: Optional[tuple[float, float]] = Field(default=[0, 1])
    roll_step: float = 2.0

    @property
    def euler_angles_offsets(self) -> tuple[torch.Tensor, torch.Tensor]:
        """Return the Euler angle offsets to search over.

        Note that this method uses a uniform grid search which approximates SO(3) space
        well when the angular ranges are small.

        Returns
        -------
        tuple[torch.Tensor, torch.Tensor]
            A tuple of two tensors of shape (N, 3) where N is the number of
            orientations to search over. The first tensor represents the Euler
            angles of the rotated template, and the second tensor represents
            the Euler angles of the rotation axis. The columns represent the
            phi, theta, and psi angles, respectively, in the 'ZYZ' convention.
        """
        if not self.enabled:
            return torch.zeros((1, 3)), torch.zeros((1, 3))

        if self.search_roll_axis:
            self.roll_axis = None
        roll_axis = None
        if self.roll_axis is not None:
            roll_axis = torch.tensor(self.roll_axis)

        if self.base_grid_method == "roll":
            euler_angles_offsets = get_roll_angles(
                psi_step=self.psi_step,
                psi_min=self.psi_min,
                psi_max=self.psi_max,
                theta_step=self.theta_step,
                theta_min=self.theta_min,
                theta_max=self.theta_max,
                roll_axis=roll_axis,
                roll_axis_step=self.roll_step,
            )
        else:
            euler_angles_offsets = get_uniform_euler_angles(
                phi_step=self.phi_step,
                theta_step=self.theta_step,
                psi_step=self.psi_step,
                phi_min=self.phi_min,
                phi_max=self.phi_max,
                theta_min=self.theta_min,
                theta_max=self.theta_max,
                psi_min=self.psi_min,
                psi_max=self.psi_max,
                base_grid_method=self.base_grid_method,
            )
        # Convert to rotation matrix
        rot_z_matrix = roma.euler_to_rotmat(
            "ZYZ",
            euler_angles_offsets,
            degrees=True,
            device=euler_angles_offsets.device,
        ).to(torch.float32)
        # Apply rotation to the template
        rot_axis_matrix = roma.euler_to_rotmat(
            "ZYZ",
            torch.tensor(self.rotation_axis_euler_angles),
            degrees=True,
            device=euler_angles_offsets.device,
        ).to(torch.float32)

        rot_matrix_batch = roma.rotmat_composition(
            sequence=(rot_axis_matrix, rot_z_matrix, rot_axis_matrix.transpose(-1, -2))
        )

        # Convert back to Euler angles
        euler_angles_offsets_rotated = roma.rotmat_to_euler(
            "ZYZ", rot_matrix_batch, degrees=True
        ).to(torch.float32)
        return euler_angles_offsets_rotated, euler_angles_offsets

euler_angles_offsets property

Return the Euler angle offsets to search over.

Note that this method uses a uniform grid search which approximates SO(3) space well when the angular ranges are small.

Returns:

Type Description
tuple[Tensor, Tensor]

A tuple of two tensors of shape (N, 3) where N is the number of orientations to search over. The first tensor represents the Euler angles of the rotated template, and the second tensor represents the Euler angles of the rotation axis. The columns represent the phi, theta, and psi angles, respectively, in the 'ZYZ' convention.

MultipleOrientationConfig

Bases: BaseModel2DTM

Configuration for multiple orientation search ranges.

This class allows specifying multiple complete orientation search configurations and concatenates their Euler angles.

Attributes:

Name Type Description
orientation_configs list[OrientationSearchConfig]

List of orientation search configurations to combine.
Source code in src/leopard_em/pydantic_models/config/orientation_search.py 
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
class MultipleOrientationConfig(BaseModel2DTM):
    """Configuration for multiple orientation search ranges.

    This class allows specifying multiple complete orientation search
    configurations and concatenates their Euler angles.

    Attributes
    ----------
    orientation_configs : list[OrientationSearchConfig]
        List of orientation search configurations to combine.
    """

    orientation_configs: list[OrientationSearchConfig]

    @property
    def euler_angles(self) -> torch.Tensor:
        """Returns the concatenated Euler angles from all orientation configs.

        Returns
        -------
        torch.Tensor
            A tensor of shape (N, 3) where N is the total number of orientations
            from all configurations. The columns represent the psi, theta, and phi
            angles respectively.
        """
        all_euler_angles = []
        for config in self.orientation_configs:
            all_euler_angles.append(config.euler_angles)

        return torch.cat(all_euler_angles, dim=0)

euler_angles property

Returns the concatenated Euler angles from all orientation configs.

Returns:

Type Description
Tensor

A tensor of shape (N, 3) where N is the total number of orientations from all configurations. The columns represent the psi, theta, and phi angles respectively.

OrientationSearchConfig

Bases: BaseModel2DTM

Serialization and validation of orientation search parameters for 2DTM.

The angles -- phi, theta, and psi -- represent Euler angles in the 'ZYZ' convention in units of degrees between 0 and 360 (for phi and psi) or between 0 and 180 (for theta).

This model effectively acts as a connector into the torch_so3.uniform_so3_sampling.get_uniform_euler_angles function from the torch-so3 package.

TODO: Implement indexing to get the i-th or range of orientations in the search space (need to be ordered).

Attributes:

Name Type Description
psi_step float

Angular step size for psi in degrees. Must be greater than 0.
theta_step float

Angular step size for theta in degrees. Must be greater than 0.
phi_min float

Minimum value for the phi angle in degrees.
phi_max float

Maximum value for the phi angle in degrees.
theta_min float

Minimum value for the theta angle in degrees.
theta_max float

Maximum value for the theta angle in degrees.
psi_min float

Minimum value for the psi angle in degrees.
psi_max float

Maximum value for the psi angle in degrees.
base_grid_method str

Method for sampling orientations. Default is 'uniform'. Currently only 'uniform' is supported.
symmetry str

Symmetry group of the template. Default is 'C1'. Note that if symmetry is provided, then the angle min/max values must be all set to None (validation will set these automatically based on the symmetry group).
Source code in src/leopard_em/pydantic_models/config/orientation_search.py 
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
class OrientationSearchConfig(BaseModel2DTM):
    """Serialization and validation of orientation search parameters for 2DTM.

    The angles -- phi, theta, and psi -- represent Euler angles in the 'ZYZ'
    convention in units of degrees between 0 and 360 (for phi and psi) or
    between 0 and 180 (for theta).

    This model effectively acts as a connector into the
    `torch_so3.uniform_so3_sampling.get_uniform_euler_angles` function from the
    [torch-so3](https://github.com/teamtomo/torch-so3) package.

    TODO: Implement indexing to get the i-th or range of orientations in the
    search space (need to be ordered).

    Attributes
    ----------
    psi_step : float
        Angular step size for psi in degrees. Must be greater
        than 0.
    theta_step : float
        Angular step size for theta in degrees. Must be
        greater than 0.
    phi_min : float
        Minimum value for the phi angle in degrees.
    phi_max : float
        Maximum value for the phi angle in degrees.
    theta_min : float
        Minimum value for the theta angle in degrees.
    theta_max : float
        Maximum value for the theta angle in degrees.
    psi_min : float
        Minimum value for the psi angle in degrees.
    psi_max : float
        Maximum value for the psi angle in degrees.
    base_grid_method : str
        Method for sampling orientations. Default is 'uniform'.
        Currently only 'uniform' is supported.
    symmetry : str
        Symmetry group of the template. Default is 'C1'. Note that if symmetry is
        provided, then the angle min/max values must be all set to None (validation
        will set these automatically based on the symmetry group).
    """

    psi_step: Annotated[float, Field(ge=0.0)] = 1.5
    theta_step: Annotated[float, Field(ge=0.0)] = 2.5
    phi_min: Optional[float] = None
    phi_max: Optional[float] = None
    theta_min: Optional[float] = None
    theta_max: Optional[float] = None
    psi_min: Optional[float] = None
    psi_max: Optional[float] = None

    base_grid_method: Literal["uniform", "healpix", "cartesian"] = "uniform"
    symmetry: Optional[str] = "C1"

    @model_validator(mode="after")  # type: ignore
    def validate_angle_ranges_and_symmetry(self) -> Self:
        """Validate that angle ranges are consistent with symmetry.

        There should be only two valid cases for combinations of manually defined
        angle ranges and the symmetry group:
        1. Symmetry argument is *not* None, and all angle min/max values
           are set to None. In this case, the angle ranges will be set based on
           the symmetry group.
        2. Symmetry argument is None, and all angle min/max values are not None.

        If any other combination is provided, a ValueError will be raised.
        """
        _all_none = all(
            x is None
            for x in [
                self.phi_min,
                self.phi_max,
                self.theta_min,
                self.theta_max,
                self.psi_min,
                self.psi_max,
            ]
        )

        # Check that both symmetry and angle ranges are not None
        if not self.symmetry and _all_none:
            raise ValueError(
                "Either a symmetry group must be provided, or all angle ranges must "
                "not be None. Both symmetry and angle ranges were set to None."
            )

        # Case where both symmetry and angle ranges are provided
        if self.symmetry and not _all_none:
            raise ValueError(
                "Symmetry group is provided, but angle ranges are also set. "
                "Please set all angle ranges to None when using symmetry."
            )

        # Case where symmetry group is provided, validate the symmetry
        if self.symmetry:
            match = re.match(r"([A-Za-z]+)(\d*)$", self.symmetry)
            if not match:
                raise ValueError(f"Invalid symmetry format: {self.symmetry}")

        # If we reach here, it means that either symmetry is set or angle ranges are set
        # but not both, so we can proceed.
        return self

    @property
    def euler_angles(self) -> torch.Tensor:
        """Returns the Euler angles ('ZYZ' convention) to search over.

        Returns
        -------
        torch.Tensor
            A tensor of shape (N, 3) where N is the number of orientations to
            search over. The columns represent the psi, theta, and phi angles
            respectively.
        """
        # If the symmetry used for the angular ranges, calculate the angular ranges
        # based on the symmetry group.
        if self.symmetry is not None:
            match = re.match(r"([A-Za-z]+)(\d*)", self.symmetry)
            if match is None:
                raise ValueError(f"Invalid symmetry format: {self.symmetry}")

            sym_group = match.group(1)
            sym_order = int(match.group(2)) if match.group(2) else 1
            (phi_min, phi_max, theta_min, theta_max, psi_min, psi_max) = (
                get_symmetry_ranges(sym_group, sym_order)
            )
        # Otherwise, use the provided angular ranges replacing with default values if
        # any are set to None.
        else:
            phi_min = self.phi_min if self.phi_min is not None else 0.0
            phi_max = self.phi_max if self.phi_max is not None else 360.0
            theta_min = self.theta_min if self.theta_min is not None else 0.0
            theta_max = self.theta_max if self.theta_max is not None else 180.0
            psi_min = self.psi_min if self.psi_min is not None else 0.0
            psi_max = self.psi_max if self.psi_max is not None else 360.0

        # Generate angles
        return get_uniform_euler_angles(
            psi_step=self.psi_step,
            theta_step=self.theta_step,
            phi_min=phi_min,
            phi_max=phi_max,
            theta_min=theta_min,
            theta_max=theta_max,
            psi_min=psi_min,
            psi_max=psi_max,
            base_grid_method=self.base_grid_method,
        )

euler_angles property

Returns the Euler angles ('ZYZ' convention) to search over.

Returns:

Type Description
Tensor

A tensor of shape (N, 3) where N is the number of orientations to search over. The columns represent the psi, theta, and phi angles respectively.

validate_angle_ranges_and_symmetry()

Validate that angle ranges are consistent with symmetry.

There should be only two valid cases for combinations of manually defined angle ranges and the symmetry group: 1. Symmetry argument is not None, and all angle min/max values are set to None. In this case, the angle ranges will be set based on the symmetry group. 2. Symmetry argument is None, and all angle min/max values are not None.

If any other combination is provided, a ValueError will be raised.

Source code in src/leopard_em/pydantic_models/config/orientation_search.py 
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
@model_validator(mode="after")  # type: ignore
def validate_angle_ranges_and_symmetry(self) -> Self:
    """Validate that angle ranges are consistent with symmetry.

    There should be only two valid cases for combinations of manually defined
    angle ranges and the symmetry group:
    1. Symmetry argument is *not* None, and all angle min/max values
       are set to None. In this case, the angle ranges will be set based on
       the symmetry group.
    2. Symmetry argument is None, and all angle min/max values are not None.

    If any other combination is provided, a ValueError will be raised.
    """
    _all_none = all(
        x is None
        for x in [
            self.phi_min,
            self.phi_max,
            self.theta_min,
            self.theta_max,
            self.psi_min,
            self.psi_max,
        ]
    )

    # Check that both symmetry and angle ranges are not None
    if not self.symmetry and _all_none:
        raise ValueError(
            "Either a symmetry group must be provided, or all angle ranges must "
            "not be None. Both symmetry and angle ranges were set to None."
        )

    # Case where both symmetry and angle ranges are provided
    if self.symmetry and not _all_none:
        raise ValueError(
            "Symmetry group is provided, but angle ranges are also set. "
            "Please set all angle ranges to None when using symmetry."
        )

    # Case where symmetry group is provided, validate the symmetry
    if self.symmetry:
        match = re.match(r"([A-Za-z]+)(\d*)$", self.symmetry)
        if not match:
            raise ValueError(f"Invalid symmetry format: {self.symmetry}")

    # If we reach here, it means that either symmetry is set or angle ranges are set
    # but not both, so we can proceed.
    return self

RefineOrientationConfig

Bases: BaseModel2DTM

Serialization and validation of orientation refinement parameters.

Angles will be sampled from [-coarse_step, coarse_step] in increments of 'fine_step' for the orientation refinement search.

Attributes:

Name Type Description
orientation_sampling_method str

Method for sampling orientations. Default is 'Hopf Fibration'. Currently only 'Hopf Fibration' is supported.
template_symmetry str

Symmetry group of the template. Default is 'C1'. Currently only 'C1' is supported.
phi_step_coarse float

Angular step size for phi in degrees for previous, coarse search. This corresponds to the 'OrientationSearchConfig.phi_step' value for the match template program. Must be greater than or equal to 0.
phi_step_fine float

Angular step size for phi in degrees for current, fine search. Must be greater than or equal to 0.
theta_step_coarse float

Angular step size for theta in degrees for previous, coarse search. This corresponds to the 'OrientationSearchConfig.theta_step' value for the match template program. Must be greater than or equal to 0.
theta_step_fine float

Angular step size for theta in degrees for current, fine search. Must be greater than or equal to 0.
psi_step_coarse float

Angular step size for psi in degrees for previous, coarse search. This corresponds to the 'OrientationSearchConfig.psi_step' value for the match template program. Must be greater than or equal to 0.
psi_step_fine float

Angular step size for psi in degrees for current, fine search. Must be greater than or equal to 0.
Source code in src/leopard_em/pydantic_models/config/orientation_search.py 
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
class RefineOrientationConfig(BaseModel2DTM):
    """Serialization and validation of orientation refinement parameters.

    Angles will be sampled from [-coarse_step, coarse_step] in increments of
    'fine_step' for the orientation refinement search.

    Attributes
    ----------
    orientation_sampling_method : str
        Method for sampling orientations. Default is 'Hopf Fibration'.
        Currently only 'Hopf Fibration' is supported.
    template_symmetry : str
        Symmetry group of the template. Default is 'C1'.
        Currently only 'C1' is supported.
    phi_step_coarse : float
        Angular step size for phi in degrees for previous, coarse search.
        This corresponds to the 'OrientationSearchConfig.phi_step' value
        for the match template program. Must be greater than or equal to 0.
    phi_step_fine : float
        Angular step size for phi in degrees for current, fine search.
        Must be greater than or equal to 0.
    theta_step_coarse : float
        Angular step size for theta in degrees for previous, coarse
        search. This corresponds to the
        'OrientationSearchConfig.theta_step' value for the match template
        program. Must be greater than or equal to 0.
    theta_step_fine : float
        Angular step size for theta in degrees for current, fine search.
        Must be greater than or equal to 0.
    psi_step_coarse : float
        Angular step size for psi in degrees for previous, coarse search.
        This corresponds to the 'OrientationSearchConfig.psi_step' value
        for the match template program. Must be greater than or equal to 0.
    psi_step_fine : float
        Angular step size for psi in degrees for current, fine search.
        Must be greater than or equal to 0.

    """

    enabled: bool = True
    phi_step_coarse: Annotated[float, Field(..., ge=0.0)] = 2.5
    phi_step_fine: Annotated[float, Field(..., ge=0.0)] = 0.25
    theta_step_coarse: Annotated[float, Field(..., ge=0.0)] = 2.5
    theta_step_fine: Annotated[float, Field(..., ge=0.0)] = 0.25
    psi_step_coarse: Annotated[float, Field(..., ge=0.0)] = 1.5
    psi_step_fine: Annotated[float, Field(..., ge=0.0)] = 0.1
    base_grid_method: Literal["uniform", "healpix", "basic"] = "uniform"

    @property
    def euler_angles_offsets(self) -> torch.Tensor:
        """Return the Euler angle offsets to search over.

        Note that this method uses a uniform grid search which approximates SO(3) space
        well when the angular ranges are small (e.g. ±2.5 degrees).

        Returns
        -------
        torch.Tensor
            A tensor of shape (N, 3) where N is the number of orientations to
            search over. The columns represent the phi, theta, and psi angles,
            respectively, in the 'ZYZ' convention.
        """
        if not self.enabled:
            return torch.zeros((1, 3))

        return get_local_high_resolution_angles(
            coarse_phi_step=self.phi_step_coarse,
            coarse_theta_step=self.theta_step_coarse,
            coarse_psi_step=self.psi_step_coarse,
            fine_phi_step=self.phi_step_fine,
            fine_theta_step=self.theta_step_fine,
            fine_psi_step=self.psi_step_fine,
            base_grid_method=self.base_grid_method,
        )

euler_angles_offsets property

Return the Euler angle offsets to search over.

Note that this method uses a uniform grid search which approximates SO(3) space well when the angular ranges are small (e.g. ±2.5 degrees).

Returns:

Type Description
Tensor

A tensor of shape (N, 3) where N is the number of orientations to search over. The columns represent the phi, theta, and psi angles, respectively, in the 'ZYZ' convention.