Skip to content

managers

Pydantic models for Leopard-EM program managers.

ConstrainedSearchManager

Bases: BaseModel2DTM

Model holding parameters necessary for running the constrained search program.

Attributes:

Name Type Description
template_volume_path str

Path to the template volume MRC file.

center_vector list[float]

The centre vector of the template volume.

particle_stack_reference ParticleStack

Particle stack object containing particle data reference particles.

particle_stack_constrained ParticleStack

Particle stack object containing particle data constrained particles.

defocus_refinement_config DefocusSearchConfig

Configuration for defocus refinement.

orientation_refinement_config RefineOrientationConfig

Configuration for orientation refinement.

preprocessing_filters PreprocessingFilters

Filters to apply to the particle images.

computational_config ComputationalConfig

What computational resources to allocate for the program.

template_volume ExcludedTensor

The template volume tensor (excluded from serialization).

false_positives float

The number of false positives to allow per particle.

Methods:

Name Description
TODO serialization/import methods
__init__

Initialize the constrained search manager.

make_backend_core_function_kwargs

Create the kwargs for the backend refine_template core function.

run_constrained_search

Run the constrained search program.

Source code in src/leopard_em/pydantic_models/managers/constrained_search_manager.py
 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
171
172
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
203
204
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
279
280
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
class ConstrainedSearchManager(BaseModel2DTM):
    """Model holding parameters necessary for running the constrained search program.

    Attributes
    ----------
    template_volume_path : str
        Path to the template volume MRC file.
    center_vector : list[float]
        The centre vector of the template volume.
    particle_stack_reference : ParticleStack
        Particle stack object containing particle data reference particles.
    particle_stack_constrained : ParticleStack
        Particle stack object containing particle data constrained particles.
    defocus_refinement_config : DefocusSearchConfig
        Configuration for defocus refinement.
    orientation_refinement_config : RefineOrientationConfig
        Configuration for orientation refinement.
    preprocessing_filters : PreprocessingFilters
        Filters to apply to the particle images.
    computational_config : ComputationalConfig
        What computational resources to allocate for the program.
    template_volume : ExcludedTensor
        The template volume tensor (excluded from serialization).
    false_positives : float
        The number of false positives to allow per particle.

    Methods
    -------
    TODO serialization/import methods
    __init__(self, skip_mrc_preloads: bool = False, **data: Any)
        Initialize the constrained search manager.
    make_backend_core_function_kwargs(self) -> dict[str, Any]
        Create the kwargs for the backend refine_template core function.
    run_constrained_search(self, orientation_batch_size: int = 64) -> None
        Run the constrained search program.
    """

    model_config: ClassVar = ConfigDict(arbitrary_types_allowed=True)

    template_volume_path: str  # In df per-particle, but ensure only one reference
    center_vector: list[float] = Field(default=[0.0, 0.0, 0.0])

    particle_stack_reference: ParticleStack
    particle_stack_constrained: ParticleStack
    defocus_refinement_config: DefocusSearchConfig
    orientation_refinement_config: ConstrainedOrientationConfig
    preprocessing_filters: PreprocessingFilters
    computational_config: ComputationalConfig

    # Excluded tensors
    template_volume: ExcludedTensor
    zdiffs: ExcludedTensor = torch.tensor([0.0])

    def __init__(self, skip_mrc_preloads: bool = False, **data: Any):
        super().__init__(**data)

        # Load the data from the MRC files
        if not skip_mrc_preloads:
            self.template_volume = load_mrc_volume(self.template_volume_path)

    # pylint: disable=too-many-locals
    def make_backend_core_function_kwargs(
        self, prefer_refined_angles: bool = True
    ) -> dict[str, Any]:
        """Create the kwargs for the backend constrained_template core function."""
        device_list = self.computational_config.gpu_devices

        template = load_template_tensor(
            template_volume=self.template_volume,
            template_volume_path=self.template_volume_path,
        )

        part_stk = self.particle_stack_reference

        euler_angles = part_stk.get_euler_angles(prefer_refined_angles)

        # The relative Euler angle offsets to search over
        euler_angle_offsets, _ = self.orientation_refinement_config.euler_angles_offsets

        # No pixel size refinement
        pixel_size_offsets = torch.tensor([0.0])

        # Extract and preprocess images and filters
        (
            particle_images_dft,
            template_dft,
            projective_filters,
        ) = setup_images_filters_particle_stack(
            part_stk, self.preprocessing_filters, template
        )

        # get z diff for each particle
        if not isinstance(self.center_vector, torch.Tensor):
            self.center_vector = torch.tensor(self.center_vector, dtype=torch.float32)
        rotation_matrices = roma.rotvec_to_rotmat(
            roma.euler_to_rotvec(convention="ZYZ", angles=euler_angles)
        ).to(torch.float32)
        rotated_vectors = rotation_matrices @ self.center_vector

        # Get z for each particle -> tensor shape [batch_size]
        new_z_diffs = rotated_vectors[:, 2]

        # The best defocus values for each particle (+ astigmatism)
        defocus_u, defocus_v = part_stk.get_absolute_defocus()
        defocus_u = defocus_u - new_z_diffs
        defocus_v = defocus_v - new_z_diffs
        # Store defocus values as instance attributes for later access
        self.zdiffs = new_z_diffs
        defocus_angle = torch.tensor(part_stk["astigmatism_angle"])

        # The relative defocus values to search over
        defocus_offsets = self.defocus_refinement_config.defocus_values

        ctf_kwargs = _setup_ctf_kwargs_from_particle_stack(
            part_stk, (template.shape[-2], template.shape[-1])
        )

        # Ger corr mean and variance
        # I want positions of reference but vals from constrained
        part_stk.set_column(
            "correlation_average_path",
            self.particle_stack_constrained["correlation_average_path"][0],
        )
        part_stk.set_column(
            "correlation_variance_path",
            self.particle_stack_constrained["correlation_variance_path"][0],
        )
        corr_mean_stack = part_stk.construct_cropped_statistic_stack(
            "correlation_average"
        )
        corr_std_stack = (
            part_stk.construct_cropped_statistic_stack(
                stat="correlation_variance",
                pos_reference="center",
                handle_bounds="pad",
                padding_mode="constant",
                padding_value=1e10,
            )
            ** 0.5
        )  # var to std

        return {
            "particle_stack_dft": particle_images_dft,
            "template_dft": template_dft,
            "euler_angles": euler_angles,
            "euler_angle_offsets": euler_angle_offsets,
            "defocus_u": defocus_u,
            "defocus_v": defocus_v,
            "defocus_angle": defocus_angle,
            "defocus_offsets": defocus_offsets,
            "pixel_size_offsets": pixel_size_offsets,
            "corr_mean": corr_mean_stack,
            "corr_std": corr_std_stack,
            "ctf_kwargs": ctf_kwargs,
            "projective_filters": projective_filters,
            "device": device_list,  # Pass all devices to core_refine_template
        }

    def run_constrained_search(
        self,
        output_dataframe_path: str,
        false_positives: float = 0.005,
        orientation_batch_size: int = 64,
    ) -> None:
        """Run the constrained search program and saves the resultant DataFrame to csv.

        Parameters
        ----------
        output_dataframe_path : str
            Path to save the constrained search results.
        false_positives : float
            The number of false positives to allow per particle.
        orientation_batch_size : int
            Number of orientations to process at once. Defaults to 64.
        """
        backend_kwargs = self.make_backend_core_function_kwargs()

        result = self.get_refine_result(backend_kwargs, orientation_batch_size)

        self.refine_result_to_dataframe(
            output_dataframe_path=output_dataframe_path,
            result=result,
            false_positives=false_positives,
        )

    def get_refine_result(
        self, backend_kwargs: dict, orientation_batch_size: int = 64
    ) -> dict[str, np.ndarray]:
        """Get refine template result.

        Parameters
        ----------
        backend_kwargs : dict
            Keyword arguments for the backend processing
        orientation_batch_size : int
            Number of orientations to process at once. Defaults to 64.

        Returns
        -------
        dict[str, np.ndarray]
            The result of the refine template program.
        """
        # Adjust batch size if orientation search is disabled
        if not self.orientation_refinement_config.enabled:
            orientation_batch_size = 1
        elif (
            self.orientation_refinement_config.euler_angles_offsets[0].shape[0]
            < orientation_batch_size
        ):
            orientation_batch_size = (
                self.orientation_refinement_config.euler_angles_offsets[0].shape[0]
            )

        result: dict[str, np.ndarray] = {}
        result = core_refine_template(
            batch_size=orientation_batch_size, **backend_kwargs
        )
        result = {k: v.cpu().numpy() for k, v in result.items()}
        return result

    # pylint: disable=too-many-locals
    def refine_result_to_dataframe(
        self,
        output_dataframe_path: str,
        result: dict[str, np.ndarray],
        false_positives: float = 0.005,
    ) -> None:
        """Convert refine template result to dataframe.

        Parameters
        ----------
        output_dataframe_path : str
            Path to save the refined particle data.
        result : dict[str, np.ndarray]
            The result of the refine template program.
        false_positives : float
            The number of false positives to allow per particle.
        """
        df_refined = self.particle_stack_reference.get_dataframe_copy()

        # x and y positions
        pos_offset_y = result["refined_pos_y"]
        pos_offset_x = result["refined_pos_x"]
        pos_offset_y_ang = pos_offset_y * df_refined["pixel_size"]
        pos_offset_x_ang = pos_offset_x * df_refined["pixel_size"]

        df_refined["refined_pos_y"] = pos_offset_y + df_refined["pos_y"]
        df_refined["refined_pos_x"] = pos_offset_x + df_refined["pos_x"]
        df_refined["refined_pos_y_img"] = pos_offset_y + df_refined["pos_y_img"]
        df_refined["refined_pos_x_img"] = pos_offset_x + df_refined["pos_x_img"]
        df_refined["refined_pos_y_img_angstrom"] = (
            pos_offset_y_ang + df_refined["pos_y_img_angstrom"]
        )
        df_refined["refined_pos_x_img_angstrom"] = (
            pos_offset_x_ang + df_refined["pos_x_img_angstrom"]
        )

        # Euler angles
        angle_idx = result["angle_idx"]
        df_refined["refined_psi"] = result["refined_euler_angles"][:, 2]
        df_refined["refined_theta"] = result["refined_euler_angles"][:, 1]
        df_refined["refined_phi"] = result["refined_euler_angles"][:, 0]

        _, euler_angle_offsets = self.orientation_refinement_config.euler_angles_offsets
        euler_angle_offsets_np = euler_angle_offsets.cpu().numpy()
        # Store the matched original offsets in the dataframe
        df_refined["original_offset_phi"] = euler_angle_offsets_np[angle_idx, 0]
        df_refined["original_offset_theta"] = euler_angle_offsets_np[angle_idx, 1]
        df_refined["original_offset_psi"] = euler_angle_offsets_np[angle_idx, 2]

        # Defocus
        df_refined["refined_relative_defocus"] = (
            result["refined_defocus_offset"]
            + df_refined["refined_relative_defocus"]
            - self.zdiffs.cpu().numpy()
        )

        # Pixel size
        df_refined["refined_pixel_size"] = (
            result["refined_pixel_size_offset"] + df_refined["pixel_size"]
        )

        # Cross-correlation statistics
        refined_mip = result["refined_cross_correlation"]
        refined_scaled_mip = result["refined_z_score"]
        df_refined["refined_mip"] = refined_mip
        df_refined["refined_scaled_mip"] = refined_scaled_mip

        # Reorder the columns
        df_refined = df_refined.reindex(columns=CONSTRAINED_DF_COLUMN_ORDER)

        # Save the refined DataFrame to disk
        df_refined.to_csv(output_dataframe_path)

        # Save a second dataframe
        # I also want the original user input offsets back somewhere
        # This one will have only those above threshold
        num_projections = (
            self.defocus_refinement_config.defocus_values.shape[0]
            * self.orientation_refinement_config.euler_angles_offsets[0].shape[0]
        )
        num_px = (
            self.particle_stack_reference.extracted_box_size[0]
            - self.particle_stack_reference.original_template_size[0]
            + 1
        ) * (
            self.particle_stack_reference.extracted_box_size[1]
            - self.particle_stack_reference.original_template_size[1]
            + 1
        )
        num_correlations = num_projections * num_px
        threshold = gaussian_noise_zscore_cutoff(
            num_correlations, float(false_positives)
        )

        # Save all parameters to CSV including false-positives
        params_df = pd.DataFrame(
            {
                "num_projections": [num_projections],
                "num_px": [num_px],
                "num_correlations": [num_correlations],
                "false_positives": [false_positives],
                "threshold": [threshold],
            }
        )
        params_df.to_csv(output_dataframe_path.replace(".csv", "_parameters.csv"))

        print(
            f"Threshold: {threshold} which gives {false_positives} "
            "false positives per particle"
        )
        df_refined_above_threshold = df_refined[
            df_refined["refined_scaled_mip"] > threshold
        ]
        # Also remove if refined_scaled_mip is inf or nan
        df_refined_above_threshold = df_refined_above_threshold[
            df_refined_above_threshold["refined_scaled_mip"] != np.inf
        ]
        df_refined_above_threshold = df_refined_above_threshold[
            df_refined_above_threshold["refined_scaled_mip"] != np.nan
        ]
        # Save the above threshold dataframe
        print(
            f"Saving above threshold dataframe to "
            f"{output_dataframe_path.replace('.csv', '_above_threshold.csv')}"
        )
        df_refined_above_threshold.to_csv(
            output_dataframe_path.replace(".csv", "_above_threshold.csv")
        )

get_refine_result(backend_kwargs, orientation_batch_size=64)

Get refine template result.

Parameters:

Name Type Description Default
backend_kwargs dict

Keyword arguments for the backend processing

required
orientation_batch_size int

Number of orientations to process at once. Defaults to 64.

64

Returns:

Type Description
dict[str, ndarray]

The result of the refine template program.

Source code in src/leopard_em/pydantic_models/managers/constrained_search_manager.py
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
def get_refine_result(
    self, backend_kwargs: dict, orientation_batch_size: int = 64
) -> dict[str, np.ndarray]:
    """Get refine template result.

    Parameters
    ----------
    backend_kwargs : dict
        Keyword arguments for the backend processing
    orientation_batch_size : int
        Number of orientations to process at once. Defaults to 64.

    Returns
    -------
    dict[str, np.ndarray]
        The result of the refine template program.
    """
    # Adjust batch size if orientation search is disabled
    if not self.orientation_refinement_config.enabled:
        orientation_batch_size = 1
    elif (
        self.orientation_refinement_config.euler_angles_offsets[0].shape[0]
        < orientation_batch_size
    ):
        orientation_batch_size = (
            self.orientation_refinement_config.euler_angles_offsets[0].shape[0]
        )

    result: dict[str, np.ndarray] = {}
    result = core_refine_template(
        batch_size=orientation_batch_size, **backend_kwargs
    )
    result = {k: v.cpu().numpy() for k, v in result.items()}
    return result

make_backend_core_function_kwargs(prefer_refined_angles=True)

Create the kwargs for the backend constrained_template core function.

Source code in src/leopard_em/pydantic_models/managers/constrained_search_manager.py
 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
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
def make_backend_core_function_kwargs(
    self, prefer_refined_angles: bool = True
) -> dict[str, Any]:
    """Create the kwargs for the backend constrained_template core function."""
    device_list = self.computational_config.gpu_devices

    template = load_template_tensor(
        template_volume=self.template_volume,
        template_volume_path=self.template_volume_path,
    )

    part_stk = self.particle_stack_reference

    euler_angles = part_stk.get_euler_angles(prefer_refined_angles)

    # The relative Euler angle offsets to search over
    euler_angle_offsets, _ = self.orientation_refinement_config.euler_angles_offsets

    # No pixel size refinement
    pixel_size_offsets = torch.tensor([0.0])

    # Extract and preprocess images and filters
    (
        particle_images_dft,
        template_dft,
        projective_filters,
    ) = setup_images_filters_particle_stack(
        part_stk, self.preprocessing_filters, template
    )

    # get z diff for each particle
    if not isinstance(self.center_vector, torch.Tensor):
        self.center_vector = torch.tensor(self.center_vector, dtype=torch.float32)
    rotation_matrices = roma.rotvec_to_rotmat(
        roma.euler_to_rotvec(convention="ZYZ", angles=euler_angles)
    ).to(torch.float32)
    rotated_vectors = rotation_matrices @ self.center_vector

    # Get z for each particle -> tensor shape [batch_size]
    new_z_diffs = rotated_vectors[:, 2]

    # The best defocus values for each particle (+ astigmatism)
    defocus_u, defocus_v = part_stk.get_absolute_defocus()
    defocus_u = defocus_u - new_z_diffs
    defocus_v = defocus_v - new_z_diffs
    # Store defocus values as instance attributes for later access
    self.zdiffs = new_z_diffs
    defocus_angle = torch.tensor(part_stk["astigmatism_angle"])

    # The relative defocus values to search over
    defocus_offsets = self.defocus_refinement_config.defocus_values

    ctf_kwargs = _setup_ctf_kwargs_from_particle_stack(
        part_stk, (template.shape[-2], template.shape[-1])
    )

    # Ger corr mean and variance
    # I want positions of reference but vals from constrained
    part_stk.set_column(
        "correlation_average_path",
        self.particle_stack_constrained["correlation_average_path"][0],
    )
    part_stk.set_column(
        "correlation_variance_path",
        self.particle_stack_constrained["correlation_variance_path"][0],
    )
    corr_mean_stack = part_stk.construct_cropped_statistic_stack(
        "correlation_average"
    )
    corr_std_stack = (
        part_stk.construct_cropped_statistic_stack(
            stat="correlation_variance",
            pos_reference="center",
            handle_bounds="pad",
            padding_mode="constant",
            padding_value=1e10,
        )
        ** 0.5
    )  # var to std

    return {
        "particle_stack_dft": particle_images_dft,
        "template_dft": template_dft,
        "euler_angles": euler_angles,
        "euler_angle_offsets": euler_angle_offsets,
        "defocus_u": defocus_u,
        "defocus_v": defocus_v,
        "defocus_angle": defocus_angle,
        "defocus_offsets": defocus_offsets,
        "pixel_size_offsets": pixel_size_offsets,
        "corr_mean": corr_mean_stack,
        "corr_std": corr_std_stack,
        "ctf_kwargs": ctf_kwargs,
        "projective_filters": projective_filters,
        "device": device_list,  # Pass all devices to core_refine_template
    }

refine_result_to_dataframe(output_dataframe_path, result, false_positives=0.005)

Convert refine template result to dataframe.

Parameters:

Name Type Description Default
output_dataframe_path str

Path to save the refined particle data.

required
result dict[str, ndarray]

The result of the refine template program.

required
false_positives float

The number of false positives to allow per particle.

0.005
Source code in src/leopard_em/pydantic_models/managers/constrained_search_manager.py
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
279
280
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
def refine_result_to_dataframe(
    self,
    output_dataframe_path: str,
    result: dict[str, np.ndarray],
    false_positives: float = 0.005,
) -> None:
    """Convert refine template result to dataframe.

    Parameters
    ----------
    output_dataframe_path : str
        Path to save the refined particle data.
    result : dict[str, np.ndarray]
        The result of the refine template program.
    false_positives : float
        The number of false positives to allow per particle.
    """
    df_refined = self.particle_stack_reference.get_dataframe_copy()

    # x and y positions
    pos_offset_y = result["refined_pos_y"]
    pos_offset_x = result["refined_pos_x"]
    pos_offset_y_ang = pos_offset_y * df_refined["pixel_size"]
    pos_offset_x_ang = pos_offset_x * df_refined["pixel_size"]

    df_refined["refined_pos_y"] = pos_offset_y + df_refined["pos_y"]
    df_refined["refined_pos_x"] = pos_offset_x + df_refined["pos_x"]
    df_refined["refined_pos_y_img"] = pos_offset_y + df_refined["pos_y_img"]
    df_refined["refined_pos_x_img"] = pos_offset_x + df_refined["pos_x_img"]
    df_refined["refined_pos_y_img_angstrom"] = (
        pos_offset_y_ang + df_refined["pos_y_img_angstrom"]
    )
    df_refined["refined_pos_x_img_angstrom"] = (
        pos_offset_x_ang + df_refined["pos_x_img_angstrom"]
    )

    # Euler angles
    angle_idx = result["angle_idx"]
    df_refined["refined_psi"] = result["refined_euler_angles"][:, 2]
    df_refined["refined_theta"] = result["refined_euler_angles"][:, 1]
    df_refined["refined_phi"] = result["refined_euler_angles"][:, 0]

    _, euler_angle_offsets = self.orientation_refinement_config.euler_angles_offsets
    euler_angle_offsets_np = euler_angle_offsets.cpu().numpy()
    # Store the matched original offsets in the dataframe
    df_refined["original_offset_phi"] = euler_angle_offsets_np[angle_idx, 0]
    df_refined["original_offset_theta"] = euler_angle_offsets_np[angle_idx, 1]
    df_refined["original_offset_psi"] = euler_angle_offsets_np[angle_idx, 2]

    # Defocus
    df_refined["refined_relative_defocus"] = (
        result["refined_defocus_offset"]
        + df_refined["refined_relative_defocus"]
        - self.zdiffs.cpu().numpy()
    )

    # Pixel size
    df_refined["refined_pixel_size"] = (
        result["refined_pixel_size_offset"] + df_refined["pixel_size"]
    )

    # Cross-correlation statistics
    refined_mip = result["refined_cross_correlation"]
    refined_scaled_mip = result["refined_z_score"]
    df_refined["refined_mip"] = refined_mip
    df_refined["refined_scaled_mip"] = refined_scaled_mip

    # Reorder the columns
    df_refined = df_refined.reindex(columns=CONSTRAINED_DF_COLUMN_ORDER)

    # Save the refined DataFrame to disk
    df_refined.to_csv(output_dataframe_path)

    # Save a second dataframe
    # I also want the original user input offsets back somewhere
    # This one will have only those above threshold
    num_projections = (
        self.defocus_refinement_config.defocus_values.shape[0]
        * self.orientation_refinement_config.euler_angles_offsets[0].shape[0]
    )
    num_px = (
        self.particle_stack_reference.extracted_box_size[0]
        - self.particle_stack_reference.original_template_size[0]
        + 1
    ) * (
        self.particle_stack_reference.extracted_box_size[1]
        - self.particle_stack_reference.original_template_size[1]
        + 1
    )
    num_correlations = num_projections * num_px
    threshold = gaussian_noise_zscore_cutoff(
        num_correlations, float(false_positives)
    )

    # Save all parameters to CSV including false-positives
    params_df = pd.DataFrame(
        {
            "num_projections": [num_projections],
            "num_px": [num_px],
            "num_correlations": [num_correlations],
            "false_positives": [false_positives],
            "threshold": [threshold],
        }
    )
    params_df.to_csv(output_dataframe_path.replace(".csv", "_parameters.csv"))

    print(
        f"Threshold: {threshold} which gives {false_positives} "
        "false positives per particle"
    )
    df_refined_above_threshold = df_refined[
        df_refined["refined_scaled_mip"] > threshold
    ]
    # Also remove if refined_scaled_mip is inf or nan
    df_refined_above_threshold = df_refined_above_threshold[
        df_refined_above_threshold["refined_scaled_mip"] != np.inf
    ]
    df_refined_above_threshold = df_refined_above_threshold[
        df_refined_above_threshold["refined_scaled_mip"] != np.nan
    ]
    # Save the above threshold dataframe
    print(
        f"Saving above threshold dataframe to "
        f"{output_dataframe_path.replace('.csv', '_above_threshold.csv')}"
    )
    df_refined_above_threshold.to_csv(
        output_dataframe_path.replace(".csv", "_above_threshold.csv")
    )

Run the constrained search program and saves the resultant DataFrame to csv.

Parameters:

Name Type Description Default
output_dataframe_path str

Path to save the constrained search results.

required
false_positives float

The number of false positives to allow per particle.

0.005
orientation_batch_size int

Number of orientations to process at once. Defaults to 64.

64
Source code in src/leopard_em/pydantic_models/managers/constrained_search_manager.py
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
def run_constrained_search(
    self,
    output_dataframe_path: str,
    false_positives: float = 0.005,
    orientation_batch_size: int = 64,
) -> None:
    """Run the constrained search program and saves the resultant DataFrame to csv.

    Parameters
    ----------
    output_dataframe_path : str
        Path to save the constrained search results.
    false_positives : float
        The number of false positives to allow per particle.
    orientation_batch_size : int
        Number of orientations to process at once. Defaults to 64.
    """
    backend_kwargs = self.make_backend_core_function_kwargs()

    result = self.get_refine_result(backend_kwargs, orientation_batch_size)

    self.refine_result_to_dataframe(
        output_dataframe_path=output_dataframe_path,
        result=result,
        false_positives=false_positives,
    )

MatchTemplateManager

Bases: BaseModel2DTM

Model holding parameters necessary for running full orientation 2DTM.

Attributes:

Name Type Description
micrograph_path str

Path to the micrograph .mrc file.

template_volume_path str

Path to the template volume .mrc file.

micrograph ExcludedTensor

Image to run template matching on. Not serialized.

template_volume ExcludedTensor

Template volume to match against. Not serialized.

optics_group OpticsGroup

Optics group parameters for the imaging system on the microscope.

defocus_search_config DefocusSearchConfig

Parameters for searching over defocus values.

orientation_search_config OrientationSearchConfig

Parameters for searching over orientation angles.

preprocessing_filters PreprocessingFilters

Configurations for the preprocessing filters to apply during correlation.

match_template_result MatchTemplateResult

Result of the match template program stored as an instance of the MatchTemplateResult class.

computational_config ComputationalConfig

Parameters for controlling computational resources.

Methods:

Name Description
validate_micrograph_path

Ensure the micrograph file exists.

validate_template_volume_path

Ensure the template volume file exists.

__init__

Constructor which also loads the micrograph and template volume from disk. The 'preload_mrc_files' parameter controls whether to read the MRC files immediately upon initialization.

make_backend_core_function_kwargs

Generates the keyword arguments for backend 'core_match_template' call from held parameters. Does the necessary pre-processing steps to filter the image and template.

run_match_template

Runs the base match template program in PyTorch.

results_to_dataframe

half_template_width_pos_shift: bool = True, exclude_columns: Optional[list] = None, locate_peaks_kwargs: Optional[dict] = None,

) -> pd.DataFrame

Converts the basic extracted peak info DataFrame (from the result object) to a DataFrame with additional information about reference files, microscope parameters, etc.

save_config

Save this Pydantic model config to disk.

Source code in src/leopard_em/pydantic_models/managers/match_template_manager.py
 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
171
172
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
203
204
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
279
280
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
class MatchTemplateManager(BaseModel2DTM):
    """Model holding parameters necessary for running full orientation 2DTM.

    Attributes
    ----------
    micrograph_path : str
        Path to the micrograph .mrc file.
    template_volume_path : str
        Path to the template volume .mrc file.
    micrograph : ExcludedTensor
        Image to run template matching on. Not serialized.
    template_volume : ExcludedTensor
        Template volume to match against. Not serialized.
    optics_group : OpticsGroup
        Optics group parameters for the imaging system on the microscope.
    defocus_search_config : DefocusSearchConfig
        Parameters for searching over defocus values.
    orientation_search_config : OrientationSearchConfig
        Parameters for searching over orientation angles.
    preprocessing_filters : PreprocessingFilters
        Configurations for the preprocessing filters to apply during
        correlation.
    match_template_result : MatchTemplateResult
        Result of the match template program stored as an instance of the
        `MatchTemplateResult` class.
    computational_config : ComputationalConfig
        Parameters for controlling computational resources.

    Methods
    -------
    validate_micrograph_path(v: str) -> str
        Ensure the micrograph file exists.
    validate_template_volume_path(v: str) -> str
        Ensure the template volume file exists.
    __init__(preload_mrc_files: bool = False , **data: Any)
        Constructor which also loads the micrograph and template volume from disk.
        The 'preload_mrc_files' parameter controls whether to read the MRC files
        immediately upon initialization.
    make_backend_core_function_kwargs() -> dict[str, Any]
        Generates the keyword arguments for backend 'core_match_template' call from
        held parameters. Does the necessary pre-processing steps to filter the image
        and template.
    run_match_template(orientation_batch_size: int = 1, do_result_export: bool = True)
        Runs the base match template program in PyTorch.
    results_to_dataframe(
        half_template_width_pos_shift: bool = True,
        exclude_columns: Optional[list] = None,
        locate_peaks_kwargs: Optional[dict] = None,
    ) -> pd.DataFrame
        Converts the basic extracted peak info DataFrame (from the result object) to a
        DataFrame with additional information about reference files, microscope
        parameters, etc.
    save_config(path: str, mode: Literal["yaml", "json"] = "yaml") -> None
        Save this Pydantic model config to disk.
    """

    model_config: ClassVar = ConfigDict(arbitrary_types_allowed=True)

    # Serialized attributes
    micrograph_path: str
    template_volume_path: str
    optics_group: OpticsGroup
    defocus_search_config: DefocusSearchConfig
    orientation_search_config: OrientationSearchConfig
    preprocessing_filters: PreprocessingFilters
    match_template_result: MatchTemplateResult
    computational_config: ComputationalConfig

    # Non-serialized large array-like attributes
    micrograph: ExcludedTensor
    template_volume: ExcludedTensor

    ###########################
    ### Pydantic Validators ###
    ###########################

    @field_validator("micrograph_path")  # type: ignore
    def validate_micrograph_path(cls, v) -> str:
        """Ensure the micrograph file exists."""
        if not os.path.exists(v):
            raise ValueError(f"File '{v}' for micrograph does not exist.")

        return str(v)

    @field_validator("template_volume_path")  # type: ignore
    def validate_template_volume_path(cls, v) -> str:
        """Ensure the template volume file exists."""
        if not os.path.exists(v):
            raise ValueError(f"File '{v}' for template volume does not exist.")

        return str(v)

    def __init__(self, preload_mrc_files: bool = False, **data: Any):
        super().__init__(**data)

        if preload_mrc_files:
            # Load the data from the MRC files
            self.micrograph = load_mrc_image(self.micrograph_path)
            self.template_volume = load_mrc_volume(self.template_volume_path)

    ############################################
    ### Functional (data processing) methods ###
    ############################################

    def make_backend_core_function_kwargs(self) -> dict[str, Any]:
        """Generates the keyword arguments for backend call from held parameters."""
        # Ensure the micrograph and template are loaded and in the correct format
        if self.micrograph is None:
            self.micrograph = load_mrc_image(self.micrograph_path)
        if self.template_volume is None:
            self.template_volume = load_mrc_volume(self.template_volume_path)

        # Ensure the micrograph and template are both Tensors before proceeding
        if not isinstance(self.micrograph, torch.Tensor):
            image = torch.from_numpy(self.micrograph)
        else:
            image = self.micrograph

        if not isinstance(self.template_volume, torch.Tensor):
            template = torch.from_numpy(self.template_volume)
        else:
            template = self.template_volume

        # Fourier transform the image (RFFT, unshifted)
        image_dft = torch.fft.rfftn(image)  # pylint: disable=E1102
        image_dft[0, 0] = 0 + 0j  # zero out the constant term

        # Get the bandpass filter individually
        bp_config = self.preprocessing_filters.bandpass_filter
        bandpass_filter = bp_config.calculate_bandpass_filter(image_dft.shape)

        # Calculate the cumulative filters for both the image and the template.
        cumulative_filter_image = self.preprocessing_filters.get_combined_filter(
            ref_img_rfft=image_dft,
            output_shape=image_dft.shape,
        )
        # NOTE: Here, manually accounting for the RFFT in output shape since we have not
        # RFFT'd the template volume yet. Also, this is 2-dimensional, not 3-dimensional
        cumulative_filter_template = self.preprocessing_filters.get_combined_filter(
            ref_img_rfft=image_dft,
            output_shape=(template.shape[-2], template.shape[-1] // 2 + 1),
        )

        # Apply the pre-processing and normalization
        image_preprocessed_dft = preprocess_image(
            image_rfft=image_dft,
            cumulative_fourier_filters=cumulative_filter_image,
            bandpass_filter=bandpass_filter,
        )

        # Calculate the CTF filters at each defocus value
        defocus_values = self.defocus_search_config.defocus_values

        # set pixel search to 0.0 for match template
        pixel_size_offsets = torch.tensor([0.0], dtype=torch.float32)

        ctf_filters = calculate_ctf_filter_stack(
            template_shape=(template.shape[0], template.shape[0]),
            optics_group=self.optics_group,
            defocus_offsets=defocus_values,
            pixel_size_offsets=pixel_size_offsets,
        )

        # Grab the Euler angles from the orientation search configuration
        # (phi, theta, psi) for ZYZ convention
        euler_angles = self.orientation_search_config.euler_angles
        euler_angles = euler_angles.to(torch.float32)

        template_dft = volume_to_rfft_fourier_slice(template)

        return {
            "image_dft": image_preprocessed_dft,
            "template_dft": template_dft,
            "ctf_filters": ctf_filters,
            "whitening_filter_template": cumulative_filter_template,
            "euler_angles": euler_angles,
            "defocus_values": defocus_values,
            "pixel_values": pixel_size_offsets,
            "device": self.computational_config.gpu_devices,
        }

    def run_match_template(
        self,
        orientation_batch_size: int = 16,
        do_result_export: bool = True,
        do_valid_cropping: bool = True,
    ) -> None:
        """Runs the base match template in pytorch.

        Parameters
        ----------
        orientation_batch_size : int
            The number of projections to process in a single batch. Default is 1.
        do_result_export : bool
            If True, call the `MatchTemplateResult.export_results` method to save the
            results to disk directly after running the match template. Default is True.
        do_valid_cropping : bool
            If True, apply the valid cropping mode to the results. Default is True.

        Returns
        -------
        None
        """
        core_kwargs = self.make_backend_core_function_kwargs()
        results = core_match_template(
            **core_kwargs, orientation_batch_size=orientation_batch_size
        )

        # Place results into the `MatchTemplateResult` object and save it.
        self.match_template_result.mip = results["mip"]
        self.match_template_result.scaled_mip = results["scaled_mip"]

        self.match_template_result.correlation_average = results["correlation_mean"]
        self.match_template_result.correlation_variance = results[
            "correlation_variance"
        ]
        self.match_template_result.orientation_psi = results["best_psi"]
        self.match_template_result.orientation_theta = results["best_theta"]
        self.match_template_result.orientation_phi = results["best_phi"]
        self.match_template_result.relative_defocus = results["best_defocus"]

        self.match_template_result.total_projections = results["total_projections"]
        self.match_template_result.total_orientations = results["total_orientations"]
        self.match_template_result.total_defocus = results["total_defocus"]

        # Apply the valid cropping mode to the results
        if do_valid_cropping:
            nx = self.template_volume.shape[-1]
            self.match_template_result.apply_valid_cropping((nx, nx))

        if do_result_export:
            self.match_template_result.export_results()

    def results_to_dataframe(
        self,
        half_template_width_pos_shift: bool = True,
        exclude_columns: Optional[list] = None,
        locate_peaks_kwargs: Optional[dict] = None,
    ) -> pd.DataFrame:
        """Converts the match template results to a DataFrame with additional info.

        Data included in this dataframe should be sufficient to do cross-correlation on
        the extracted peaks, that is, all the microscope parameters, defocus parameters,
        etc. are included in the dataframe. Run-specific filter information is *not*
        included in this dataframe; use the YAML configuration file to replicate a
        match_template run.

        Parameters
        ----------
        half_template_width_pos_shift : bool, optional
            If True, columns for the image peak position are shifted by half a template
            width to correspond to the center of the particle. This should be done when
            the position of a peak corresponds to the top-left corner of the template
            rather than the center. Default is True. This should generally be left as
            True unless you know what you are doing.
        exclude_columns : list, optional
            List of columns to exclude from the DataFrame. Default is None and no
            columns are excluded.
        locate_peaks_kwargs : dict, optional
            Keyword arguments to pass to the 'MatchTemplateResult.locate_peaks' method.
            Default is None and no additional keyword arguments are passed.

        Returns
        -------
        pd.DataFrame
            DataFrame containing the match template results.
        """
        # Short circuit if no kwargs and peaks have already been located
        if locate_peaks_kwargs is None:
            if self.match_template_result.match_template_peaks is None:
                self.match_template_result.locate_peaks()
        else:
            self.match_template_result.locate_peaks(**locate_peaks_kwargs)

        # DataFrame comes with the following columns :
        # ['mip', 'scaled_mip', 'correlation_mean', 'correlation_variance',
        # 'total_correlations'. 'pos_y', 'pos_x', 'psi', 'theta', 'phi',
        # 'relative_defocus', ]
        df = self.match_template_result.peaks_to_dataframe()

        # DataFrame currently contains pixel coordinates for results. Coordinates in
        # image correspond with upper left corner of the template. Need to translate
        # coordinates by half template width to get to particle center in image.
        # NOTE: We are assuming the template is cubic
        nx = mrcfile.open(self.template_volume_path).header.nx
        if half_template_width_pos_shift:
            df["pos_y_img"] = df["pos_y"] + nx // 2
            df["pos_x_img"] = df["pos_x"] + nx // 2
        else:
            df["pos_y_img"] = df["pos_y"]
            df["pos_x_img"] = df["pos_x"]

        # Also, the positions are in terms of pixels. Also add columns for particle
        # positions in terms of Angstroms.
        pixel_size = self.optics_group.pixel_size
        df["pos_y_img_angstrom"] = df["pos_y_img"] * pixel_size
        df["pos_x_img_angstrom"] = df["pos_x_img"] * pixel_size

        # Add microscope (CTF) parameters
        df["defocus_u"] = self.optics_group.defocus_u
        df["defocus_v"] = self.optics_group.defocus_v
        df["astigmatism_angle"] = self.optics_group.astigmatism_angle
        df["pixel_size"] = pixel_size
        df["refined_pixel_size"] = pixel_size
        df["voltage"] = self.optics_group.voltage
        df["spherical_aberration"] = self.optics_group.spherical_aberration
        df["amplitude_contrast_ratio"] = self.optics_group.amplitude_contrast_ratio
        df["phase_shift"] = self.optics_group.phase_shift
        df["ctf_B_factor"] = self.optics_group.ctf_B_factor

        # Add paths to the micrograph and reference template
        df["micrograph_path"] = self.micrograph_path
        df["template_path"] = self.template_volume_path

        # Add paths to the output statistic files
        df["mip_path"] = self.match_template_result.mip_path
        df["scaled_mip_path"] = self.match_template_result.scaled_mip_path
        df["psi_path"] = self.match_template_result.orientation_psi_path
        df["theta_path"] = self.match_template_result.orientation_theta_path
        df["phi_path"] = self.match_template_result.orientation_phi_path
        df["defocus_path"] = self.match_template_result.relative_defocus_path
        df["correlation_average_path"] = (
            self.match_template_result.correlation_average_path
        )
        df["correlation_variance_path"] = (
            self.match_template_result.correlation_variance_path
        )

        # Add particle index
        df["particle_index"] = df.index

        # Reorder columns
        df = df.reindex(columns=MATCH_TEMPLATE_DF_COLUMN_ORDER)

        # Drop columns if requested
        if exclude_columns is not None:
            df = df.drop(columns=exclude_columns)

        return df

    def save_config(self, path: str, mode: Literal["yaml", "json"] = "yaml") -> None:
        """Save this Pydandic model to disk. Wrapper around the serialization methods.

        Parameters
        ----------
        path : str
            Path to save the configuration file.
        mode : Literal["yaml", "json"], optional
            Serialization format to use. Default is 'yaml'.

        Returns
        -------
        None

        Raises
        ------
        ValueError
            If an invalid serialization mode is provided.
        """
        if mode == "yaml":
            self.to_yaml(path)
        elif mode == "json":
            self.to_json(path)
        else:
            raise ValueError(f"Invalid serialization mode '{mode}'.")

make_backend_core_function_kwargs()

Generates the keyword arguments for backend call from held parameters.

Source code in src/leopard_em/pydantic_models/managers/match_template_manager.py
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
171
172
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
203
204
205
206
207
208
209
210
def make_backend_core_function_kwargs(self) -> dict[str, Any]:
    """Generates the keyword arguments for backend call from held parameters."""
    # Ensure the micrograph and template are loaded and in the correct format
    if self.micrograph is None:
        self.micrograph = load_mrc_image(self.micrograph_path)
    if self.template_volume is None:
        self.template_volume = load_mrc_volume(self.template_volume_path)

    # Ensure the micrograph and template are both Tensors before proceeding
    if not isinstance(self.micrograph, torch.Tensor):
        image = torch.from_numpy(self.micrograph)
    else:
        image = self.micrograph

    if not isinstance(self.template_volume, torch.Tensor):
        template = torch.from_numpy(self.template_volume)
    else:
        template = self.template_volume

    # Fourier transform the image (RFFT, unshifted)
    image_dft = torch.fft.rfftn(image)  # pylint: disable=E1102
    image_dft[0, 0] = 0 + 0j  # zero out the constant term

    # Get the bandpass filter individually
    bp_config = self.preprocessing_filters.bandpass_filter
    bandpass_filter = bp_config.calculate_bandpass_filter(image_dft.shape)

    # Calculate the cumulative filters for both the image and the template.
    cumulative_filter_image = self.preprocessing_filters.get_combined_filter(
        ref_img_rfft=image_dft,
        output_shape=image_dft.shape,
    )
    # NOTE: Here, manually accounting for the RFFT in output shape since we have not
    # RFFT'd the template volume yet. Also, this is 2-dimensional, not 3-dimensional
    cumulative_filter_template = self.preprocessing_filters.get_combined_filter(
        ref_img_rfft=image_dft,
        output_shape=(template.shape[-2], template.shape[-1] // 2 + 1),
    )

    # Apply the pre-processing and normalization
    image_preprocessed_dft = preprocess_image(
        image_rfft=image_dft,
        cumulative_fourier_filters=cumulative_filter_image,
        bandpass_filter=bandpass_filter,
    )

    # Calculate the CTF filters at each defocus value
    defocus_values = self.defocus_search_config.defocus_values

    # set pixel search to 0.0 for match template
    pixel_size_offsets = torch.tensor([0.0], dtype=torch.float32)

    ctf_filters = calculate_ctf_filter_stack(
        template_shape=(template.shape[0], template.shape[0]),
        optics_group=self.optics_group,
        defocus_offsets=defocus_values,
        pixel_size_offsets=pixel_size_offsets,
    )

    # Grab the Euler angles from the orientation search configuration
    # (phi, theta, psi) for ZYZ convention
    euler_angles = self.orientation_search_config.euler_angles
    euler_angles = euler_angles.to(torch.float32)

    template_dft = volume_to_rfft_fourier_slice(template)

    return {
        "image_dft": image_preprocessed_dft,
        "template_dft": template_dft,
        "ctf_filters": ctf_filters,
        "whitening_filter_template": cumulative_filter_template,
        "euler_angles": euler_angles,
        "defocus_values": defocus_values,
        "pixel_values": pixel_size_offsets,
        "device": self.computational_config.gpu_devices,
    }

results_to_dataframe(half_template_width_pos_shift=True, exclude_columns=None, locate_peaks_kwargs=None)

Converts the match template results to a DataFrame with additional info.

Data included in this dataframe should be sufficient to do cross-correlation on the extracted peaks, that is, all the microscope parameters, defocus parameters, etc. are included in the dataframe. Run-specific filter information is not included in this dataframe; use the YAML configuration file to replicate a match_template run.

Parameters:

Name Type Description Default
half_template_width_pos_shift bool

If True, columns for the image peak position are shifted by half a template width to correspond to the center of the particle. This should be done when the position of a peak corresponds to the top-left corner of the template rather than the center. Default is True. This should generally be left as True unless you know what you are doing.

True
exclude_columns list

List of columns to exclude from the DataFrame. Default is None and no columns are excluded.

None
locate_peaks_kwargs dict

Keyword arguments to pass to the 'MatchTemplateResult.locate_peaks' method. Default is None and no additional keyword arguments are passed.

None

Returns:

Type Description
DataFrame

DataFrame containing the match template results.

Source code in src/leopard_em/pydantic_models/managers/match_template_manager.py
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
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
def results_to_dataframe(
    self,
    half_template_width_pos_shift: bool = True,
    exclude_columns: Optional[list] = None,
    locate_peaks_kwargs: Optional[dict] = None,
) -> pd.DataFrame:
    """Converts the match template results to a DataFrame with additional info.

    Data included in this dataframe should be sufficient to do cross-correlation on
    the extracted peaks, that is, all the microscope parameters, defocus parameters,
    etc. are included in the dataframe. Run-specific filter information is *not*
    included in this dataframe; use the YAML configuration file to replicate a
    match_template run.

    Parameters
    ----------
    half_template_width_pos_shift : bool, optional
        If True, columns for the image peak position are shifted by half a template
        width to correspond to the center of the particle. This should be done when
        the position of a peak corresponds to the top-left corner of the template
        rather than the center. Default is True. This should generally be left as
        True unless you know what you are doing.
    exclude_columns : list, optional
        List of columns to exclude from the DataFrame. Default is None and no
        columns are excluded.
    locate_peaks_kwargs : dict, optional
        Keyword arguments to pass to the 'MatchTemplateResult.locate_peaks' method.
        Default is None and no additional keyword arguments are passed.

    Returns
    -------
    pd.DataFrame
        DataFrame containing the match template results.
    """
    # Short circuit if no kwargs and peaks have already been located
    if locate_peaks_kwargs is None:
        if self.match_template_result.match_template_peaks is None:
            self.match_template_result.locate_peaks()
    else:
        self.match_template_result.locate_peaks(**locate_peaks_kwargs)

    # DataFrame comes with the following columns :
    # ['mip', 'scaled_mip', 'correlation_mean', 'correlation_variance',
    # 'total_correlations'. 'pos_y', 'pos_x', 'psi', 'theta', 'phi',
    # 'relative_defocus', ]
    df = self.match_template_result.peaks_to_dataframe()

    # DataFrame currently contains pixel coordinates for results. Coordinates in
    # image correspond with upper left corner of the template. Need to translate
    # coordinates by half template width to get to particle center in image.
    # NOTE: We are assuming the template is cubic
    nx = mrcfile.open(self.template_volume_path).header.nx
    if half_template_width_pos_shift:
        df["pos_y_img"] = df["pos_y"] + nx // 2
        df["pos_x_img"] = df["pos_x"] + nx // 2
    else:
        df["pos_y_img"] = df["pos_y"]
        df["pos_x_img"] = df["pos_x"]

    # Also, the positions are in terms of pixels. Also add columns for particle
    # positions in terms of Angstroms.
    pixel_size = self.optics_group.pixel_size
    df["pos_y_img_angstrom"] = df["pos_y_img"] * pixel_size
    df["pos_x_img_angstrom"] = df["pos_x_img"] * pixel_size

    # Add microscope (CTF) parameters
    df["defocus_u"] = self.optics_group.defocus_u
    df["defocus_v"] = self.optics_group.defocus_v
    df["astigmatism_angle"] = self.optics_group.astigmatism_angle
    df["pixel_size"] = pixel_size
    df["refined_pixel_size"] = pixel_size
    df["voltage"] = self.optics_group.voltage
    df["spherical_aberration"] = self.optics_group.spherical_aberration
    df["amplitude_contrast_ratio"] = self.optics_group.amplitude_contrast_ratio
    df["phase_shift"] = self.optics_group.phase_shift
    df["ctf_B_factor"] = self.optics_group.ctf_B_factor

    # Add paths to the micrograph and reference template
    df["micrograph_path"] = self.micrograph_path
    df["template_path"] = self.template_volume_path

    # Add paths to the output statistic files
    df["mip_path"] = self.match_template_result.mip_path
    df["scaled_mip_path"] = self.match_template_result.scaled_mip_path
    df["psi_path"] = self.match_template_result.orientation_psi_path
    df["theta_path"] = self.match_template_result.orientation_theta_path
    df["phi_path"] = self.match_template_result.orientation_phi_path
    df["defocus_path"] = self.match_template_result.relative_defocus_path
    df["correlation_average_path"] = (
        self.match_template_result.correlation_average_path
    )
    df["correlation_variance_path"] = (
        self.match_template_result.correlation_variance_path
    )

    # Add particle index
    df["particle_index"] = df.index

    # Reorder columns
    df = df.reindex(columns=MATCH_TEMPLATE_DF_COLUMN_ORDER)

    # Drop columns if requested
    if exclude_columns is not None:
        df = df.drop(columns=exclude_columns)

    return df

run_match_template(orientation_batch_size=16, do_result_export=True, do_valid_cropping=True)

Runs the base match template in pytorch.

Parameters:

Name Type Description Default
orientation_batch_size int

The number of projections to process in a single batch. Default is 1.

16
do_result_export bool

If True, call the MatchTemplateResult.export_results method to save the results to disk directly after running the match template. Default is True.

True
do_valid_cropping bool

If True, apply the valid cropping mode to the results. Default is True.

True

Returns:

Type Description
None
Source code in src/leopard_em/pydantic_models/managers/match_template_manager.py
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
def run_match_template(
    self,
    orientation_batch_size: int = 16,
    do_result_export: bool = True,
    do_valid_cropping: bool = True,
) -> None:
    """Runs the base match template in pytorch.

    Parameters
    ----------
    orientation_batch_size : int
        The number of projections to process in a single batch. Default is 1.
    do_result_export : bool
        If True, call the `MatchTemplateResult.export_results` method to save the
        results to disk directly after running the match template. Default is True.
    do_valid_cropping : bool
        If True, apply the valid cropping mode to the results. Default is True.

    Returns
    -------
    None
    """
    core_kwargs = self.make_backend_core_function_kwargs()
    results = core_match_template(
        **core_kwargs, orientation_batch_size=orientation_batch_size
    )

    # Place results into the `MatchTemplateResult` object and save it.
    self.match_template_result.mip = results["mip"]
    self.match_template_result.scaled_mip = results["scaled_mip"]

    self.match_template_result.correlation_average = results["correlation_mean"]
    self.match_template_result.correlation_variance = results[
        "correlation_variance"
    ]
    self.match_template_result.orientation_psi = results["best_psi"]
    self.match_template_result.orientation_theta = results["best_theta"]
    self.match_template_result.orientation_phi = results["best_phi"]
    self.match_template_result.relative_defocus = results["best_defocus"]

    self.match_template_result.total_projections = results["total_projections"]
    self.match_template_result.total_orientations = results["total_orientations"]
    self.match_template_result.total_defocus = results["total_defocus"]

    # Apply the valid cropping mode to the results
    if do_valid_cropping:
        nx = self.template_volume.shape[-1]
        self.match_template_result.apply_valid_cropping((nx, nx))

    if do_result_export:
        self.match_template_result.export_results()

save_config(path, mode='yaml')

Save this Pydandic model to disk. Wrapper around the serialization methods.

Parameters:

Name Type Description Default
path str

Path to save the configuration file.

required
mode Literal['yaml', 'json']

Serialization format to use. Default is 'yaml'.

'yaml'

Returns:

Type Description
None

Raises:

Type Description
ValueError

If an invalid serialization mode is provided.

Source code in src/leopard_em/pydantic_models/managers/match_template_manager.py
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
def save_config(self, path: str, mode: Literal["yaml", "json"] = "yaml") -> None:
    """Save this Pydandic model to disk. Wrapper around the serialization methods.

    Parameters
    ----------
    path : str
        Path to save the configuration file.
    mode : Literal["yaml", "json"], optional
        Serialization format to use. Default is 'yaml'.

    Returns
    -------
    None

    Raises
    ------
    ValueError
        If an invalid serialization mode is provided.
    """
    if mode == "yaml":
        self.to_yaml(path)
    elif mode == "json":
        self.to_json(path)
    else:
        raise ValueError(f"Invalid serialization mode '{mode}'.")

validate_micrograph_path(v)

Ensure the micrograph file exists.

Source code in src/leopard_em/pydantic_models/managers/match_template_manager.py
107
108
109
110
111
112
113
@field_validator("micrograph_path")  # type: ignore
def validate_micrograph_path(cls, v) -> str:
    """Ensure the micrograph file exists."""
    if not os.path.exists(v):
        raise ValueError(f"File '{v}' for micrograph does not exist.")

    return str(v)

validate_template_volume_path(v)

Ensure the template volume file exists.

Source code in src/leopard_em/pydantic_models/managers/match_template_manager.py
115
116
117
118
119
120
121
@field_validator("template_volume_path")  # type: ignore
def validate_template_volume_path(cls, v) -> str:
    """Ensure the template volume file exists."""
    if not os.path.exists(v):
        raise ValueError(f"File '{v}' for template volume does not exist.")

    return str(v)

OptimizeTemplateManager

Bases: BaseModel2DTM

Model holding parameters necessary for running the optimize template program.

Attributes:

Name Type Description
particle_stack ParticleStack

Particle stack object containing particle data.

pixel_size_coarse_search PixelSizeSearchConfig

Configuration for pixel size coarse search.

pixel_size_fine_search PixelSizeSearchConfig

Configuration for pixel size fine search.

preprocessing_filters PreprocessingFilters

Filters to apply to the particle images.

computational_config ComputationalConfig

What computational resources to allocate for the program.

simulator Simulator

The simulator object.

Methods:

Name Description
TODO serialization/import methods
__init__

Initialize the optimize template manager.

make_backend_core_function_kwargs

Create the kwargs for the backend optimize_template core function.

run_optimize_template

Run the optimize template program.

Source code in src/leopard_em/pydantic_models/managers/optimize_template_manager.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
171
172
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
203
204
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
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
class OptimizeTemplateManager(BaseModel2DTM):
    """Model holding parameters necessary for running the optimize template program.

    Attributes
    ----------
    particle_stack : ParticleStack
        Particle stack object containing particle data.
    pixel_size_coarse_search : PixelSizeSearchConfig
        Configuration for pixel size coarse search.
    pixel_size_fine_search : PixelSizeSearchConfig
        Configuration for pixel size fine search.
    preprocessing_filters : PreprocessingFilters
        Filters to apply to the particle images.
    computational_config : ComputationalConfig
        What computational resources to allocate for the program.
    simulator : Simulator
        The simulator object.

    Methods
    -------
    TODO serialization/import methods
    __init__(self, skip_mrc_preloads: bool = False, **data: Any)
        Initialize the optimize template manager.
    make_backend_core_function_kwargs(self) -> dict[str, Any]
        Create the kwargs for the backend optimize_template core function.
    run_optimize_template(self, output_text_path: str) -> None
        Run the optimize template program.
    """

    model_config: ClassVar = ConfigDict(arbitrary_types_allowed=True)

    particle_stack: ParticleStack
    pixel_size_coarse_search: PixelSizeSearchConfig
    pixel_size_fine_search: PixelSizeSearchConfig
    preprocessing_filters: PreprocessingFilters
    computational_config: ComputationalConfig
    simulator: Simulator

    # Excluded tensors
    template_volume: ExcludedTensor

    def make_backend_core_function_kwargs(
        self, prefer_refined_angles: bool = True
    ) -> dict[str, Any]:
        """Create the kwargs for the backend refine_template core function.

        Parameters
        ----------
        prefer_refined_angles : bool
            Whether to use refined angles or not. Defaults to True.
        """
        # simulate template volume
        template = self.simulator.run(gpu_ids=self.computational_config.gpu_ids)

        # The set of "best" euler angles from match template search
        # Check if refined angles exist, otherwise use the original angles
        euler_angles = self.particle_stack.get_euler_angles(prefer_refined_angles)

        # The relative Euler angle offsets to search over (none for optimization)
        euler_angle_offsets = torch.zeros((1, 3))

        # The relative defocus values to search over (none for optimization)
        defocus_offsets = torch.tensor([0.0])

        # The relative pixel size values to search over (none for optimization)
        pixel_size_offsets = torch.tensor([0.0])

        # Use the common utility function to set up the backend kwargs
        # pylint: disable=duplicate-code
        return setup_particle_backend_kwargs(
            particle_stack=self.particle_stack,
            template=template,
            preprocessing_filters=self.preprocessing_filters,
            euler_angles=euler_angles,
            euler_angle_offsets=euler_angle_offsets,
            defocus_offsets=defocus_offsets,
            pixel_size_offsets=pixel_size_offsets,
            device_list=self.computational_config.gpu_devices,
        )

    def run_optimize_template(self, output_text_path: str) -> None:
        """Run the refine template program and saves the resultant DataFrame to csv.

        Parameters
        ----------
        output_text_path : str
            Path to save the optimized template pixel size.
        """
        if self.pixel_size_coarse_search.enabled:
            # Create a file for logging all iterations
            all_results_path = self._get_all_results_path(output_text_path)
            # Create the file and write header
            with open(all_results_path, "w", encoding="utf-8") as f:
                f.write("Pixel Size (Ã…),SNR\n")

            optimal_template_px = self.optimize_pixel_size(all_results_path)
            print(f"Optimal template px: {optimal_template_px:.3f} Ã…")
            # print this to the text file
            with open(output_text_path, "w", encoding="utf-8") as f:
                f.write(f"Optimal template px: {optimal_template_px:.3f} Ã…")

    def optimize_pixel_size(self, all_results_path: str) -> float:
        """Optimize the pixel size of the template volume.

        Parameters
        ----------
        all_results_path : str
            Path to the file for logging all iterations

        Returns
        -------
        float
            The optimal pixel size.
        """
        initial_template_px = self.simulator.pixel_spacing
        print(f"Initial template px: {initial_template_px:.3f} Ã…")

        best_snr = float("-inf")
        best_px = float(initial_template_px)

        print("Starting coarse search...")

        pixel_size_offsets_coarse = self.pixel_size_coarse_search.pixel_size_values
        coarse_px_values = pixel_size_offsets_coarse + initial_template_px

        consecutive_decreases = 0
        consecutive_threshold = 2
        previous_snr = float("-inf")
        for px in coarse_px_values:
            snr = self.evaluate_template_px(px=px.item())
            print(f"Pixel size: {px:.3f}, SNR: {snr:.3f}")

            # Log to file
            with open(all_results_path, "a", encoding="utf-8") as f:
                f.write(f"{px:.3f},{snr:.3f}\n")

            if snr > best_snr:
                best_snr = snr
                best_px = px.item()
            if snr > previous_snr:
                consecutive_decreases = 0
            else:
                consecutive_decreases += 1
                if consecutive_decreases >= consecutive_threshold:
                    print(
                        f"SNR decreased for {consecutive_threshold} iterations. "
                        f"Stopping coarse px search."
                    )
                    break
            previous_snr = snr

        if self.pixel_size_fine_search.enabled:
            pixel_size_offsets_fine = self.pixel_size_fine_search.pixel_size_values
            fine_px_values = pixel_size_offsets_fine + best_px

            consecutive_decreases = 0
            previous_snr = float("-inf")
            for px in fine_px_values:
                snr = self.evaluate_template_px(px=px.item())
                print(f"Pixel size: {px:.3f}, SNR: {snr:.3f}")

                # Log to file
                with open(all_results_path, "a", encoding="utf-8") as f:
                    f.write(f"{px:.3f},{snr:.3f}\n")

                if snr > best_snr:
                    best_snr = snr
                    best_px = px.item()
                if snr > previous_snr:
                    consecutive_decreases = 0
                else:
                    consecutive_decreases += 1
                    if consecutive_decreases >= consecutive_threshold:
                        print(
                            f"SNR decreased for {consecutive_threshold} iterations. "
                            "Stopping fine px search."
                        )
                        break
                previous_snr = snr

        return best_px

    def evaluate_template_px(self, px: float) -> float:
        """Evaluate the template pixel size.

        Parameters
        ----------
        px : float
            The pixel size to evaluate.

        Returns
        -------
        float
            The mean SNR of the template.
        """
        self.simulator.pixel_spacing = px
        backend_kwargs = self.make_backend_core_function_kwargs()
        result = self.get_correlation_result(backend_kwargs, 1)
        mean_snr = self.results_to_snr(result)
        return mean_snr

    def get_correlation_result(
        self, backend_kwargs: dict, orientation_batch_size: int = 64
    ) -> dict[str, np.ndarray]:
        """Get correlation result.

        Parameters
        ----------
        backend_kwargs : dict
            Keyword arguments for the backend processing
        orientation_batch_size : int
            Number of orientations to process at once. Defaults to 64.

        Returns
        -------
        dict[str, np.ndarray]
            The result of the refine template program.
        """
        # pylint: disable=duplicate-code
        result: dict[str, np.ndarray] = {}
        result = core_refine_template(
            batch_size=orientation_batch_size, **backend_kwargs
        )
        result = {k: v.cpu().numpy() for k, v in result.items()}

        return result

    def results_to_snr(self, result: dict[str, np.ndarray]) -> float:
        """Convert optimize template result to mean SNR.

        Parameters
        ----------
        result : dict[str, np.ndarray]
            The result of the optimize template program.

        Returns
        -------
        float
            The mean SNR of the template.
        """
        # Filter out any infinite or NaN values
        # NOTE: There should not be NaNs or infs, will follow up later
        refined_scaled_mip = result["refined_z_score"]
        refined_scaled_mip = refined_scaled_mip[np.isfinite(refined_scaled_mip)]

        # If more than n values, keep only the top n highest SNRs
        best_n = 6
        if len(refined_scaled_mip) > best_n:
            refined_scaled_mip = np.sort(refined_scaled_mip)[-best_n:]

        # Printing out the results to console
        print(
            f"max snr: {refined_scaled_mip.max()}, min snr: {refined_scaled_mip.min()}"
        )

        mean_snr = float(refined_scaled_mip.mean())

        return mean_snr

    def _get_all_results_path(self, output_text_path: str) -> str:
        """Generate the results file path from the output text path.

        Parameters
        ----------
        output_text_path : str
            Path to the output text file

        Returns
        -------
        str
            Path to the file with _all.txt extension
        """
        base, _ = os.path.splitext(output_text_path)
        return f"{base}_all.csv"

evaluate_template_px(px)

Evaluate the template pixel size.

Parameters:

Name Type Description Default
px float

The pixel size to evaluate.

required

Returns:

Type Description
float

The mean SNR of the template.

Source code in src/leopard_em/pydantic_models/managers/optimize_template_manager.py
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
def evaluate_template_px(self, px: float) -> float:
    """Evaluate the template pixel size.

    Parameters
    ----------
    px : float
        The pixel size to evaluate.

    Returns
    -------
    float
        The mean SNR of the template.
    """
    self.simulator.pixel_spacing = px
    backend_kwargs = self.make_backend_core_function_kwargs()
    result = self.get_correlation_result(backend_kwargs, 1)
    mean_snr = self.results_to_snr(result)
    return mean_snr

get_correlation_result(backend_kwargs, orientation_batch_size=64)

Get correlation result.

Parameters:

Name Type Description Default
backend_kwargs dict

Keyword arguments for the backend processing

required
orientation_batch_size int

Number of orientations to process at once. Defaults to 64.

64

Returns:

Type Description
dict[str, ndarray]

The result of the refine template program.

Source code in src/leopard_em/pydantic_models/managers/optimize_template_manager.py
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
def get_correlation_result(
    self, backend_kwargs: dict, orientation_batch_size: int = 64
) -> dict[str, np.ndarray]:
    """Get correlation result.

    Parameters
    ----------
    backend_kwargs : dict
        Keyword arguments for the backend processing
    orientation_batch_size : int
        Number of orientations to process at once. Defaults to 64.

    Returns
    -------
    dict[str, np.ndarray]
        The result of the refine template program.
    """
    # pylint: disable=duplicate-code
    result: dict[str, np.ndarray] = {}
    result = core_refine_template(
        batch_size=orientation_batch_size, **backend_kwargs
    )
    result = {k: v.cpu().numpy() for k, v in result.items()}

    return result

make_backend_core_function_kwargs(prefer_refined_angles=True)

Create the kwargs for the backend refine_template core function.

Parameters:

Name Type Description Default
prefer_refined_angles bool

Whether to use refined angles or not. Defaults to True.

True
Source code in src/leopard_em/pydantic_models/managers/optimize_template_manager.py
 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
def make_backend_core_function_kwargs(
    self, prefer_refined_angles: bool = True
) -> dict[str, Any]:
    """Create the kwargs for the backend refine_template core function.

    Parameters
    ----------
    prefer_refined_angles : bool
        Whether to use refined angles or not. Defaults to True.
    """
    # simulate template volume
    template = self.simulator.run(gpu_ids=self.computational_config.gpu_ids)

    # The set of "best" euler angles from match template search
    # Check if refined angles exist, otherwise use the original angles
    euler_angles = self.particle_stack.get_euler_angles(prefer_refined_angles)

    # The relative Euler angle offsets to search over (none for optimization)
    euler_angle_offsets = torch.zeros((1, 3))

    # The relative defocus values to search over (none for optimization)
    defocus_offsets = torch.tensor([0.0])

    # The relative pixel size values to search over (none for optimization)
    pixel_size_offsets = torch.tensor([0.0])

    # Use the common utility function to set up the backend kwargs
    # pylint: disable=duplicate-code
    return setup_particle_backend_kwargs(
        particle_stack=self.particle_stack,
        template=template,
        preprocessing_filters=self.preprocessing_filters,
        euler_angles=euler_angles,
        euler_angle_offsets=euler_angle_offsets,
        defocus_offsets=defocus_offsets,
        pixel_size_offsets=pixel_size_offsets,
        device_list=self.computational_config.gpu_devices,
    )

optimize_pixel_size(all_results_path)

Optimize the pixel size of the template volume.

Parameters:

Name Type Description Default
all_results_path str

Path to the file for logging all iterations

required

Returns:

Type Description
float

The optimal pixel size.

Source code in src/leopard_em/pydantic_models/managers/optimize_template_manager.py
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
171
172
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
def optimize_pixel_size(self, all_results_path: str) -> float:
    """Optimize the pixel size of the template volume.

    Parameters
    ----------
    all_results_path : str
        Path to the file for logging all iterations

    Returns
    -------
    float
        The optimal pixel size.
    """
    initial_template_px = self.simulator.pixel_spacing
    print(f"Initial template px: {initial_template_px:.3f} Ã…")

    best_snr = float("-inf")
    best_px = float(initial_template_px)

    print("Starting coarse search...")

    pixel_size_offsets_coarse = self.pixel_size_coarse_search.pixel_size_values
    coarse_px_values = pixel_size_offsets_coarse + initial_template_px

    consecutive_decreases = 0
    consecutive_threshold = 2
    previous_snr = float("-inf")
    for px in coarse_px_values:
        snr = self.evaluate_template_px(px=px.item())
        print(f"Pixel size: {px:.3f}, SNR: {snr:.3f}")

        # Log to file
        with open(all_results_path, "a", encoding="utf-8") as f:
            f.write(f"{px:.3f},{snr:.3f}\n")

        if snr > best_snr:
            best_snr = snr
            best_px = px.item()
        if snr > previous_snr:
            consecutive_decreases = 0
        else:
            consecutive_decreases += 1
            if consecutive_decreases >= consecutive_threshold:
                print(
                    f"SNR decreased for {consecutive_threshold} iterations. "
                    f"Stopping coarse px search."
                )
                break
        previous_snr = snr

    if self.pixel_size_fine_search.enabled:
        pixel_size_offsets_fine = self.pixel_size_fine_search.pixel_size_values
        fine_px_values = pixel_size_offsets_fine + best_px

        consecutive_decreases = 0
        previous_snr = float("-inf")
        for px in fine_px_values:
            snr = self.evaluate_template_px(px=px.item())
            print(f"Pixel size: {px:.3f}, SNR: {snr:.3f}")

            # Log to file
            with open(all_results_path, "a", encoding="utf-8") as f:
                f.write(f"{px:.3f},{snr:.3f}\n")

            if snr > best_snr:
                best_snr = snr
                best_px = px.item()
            if snr > previous_snr:
                consecutive_decreases = 0
            else:
                consecutive_decreases += 1
                if consecutive_decreases >= consecutive_threshold:
                    print(
                        f"SNR decreased for {consecutive_threshold} iterations. "
                        "Stopping fine px search."
                    )
                    break
            previous_snr = snr

    return best_px

results_to_snr(result)

Convert optimize template result to mean SNR.

Parameters:

Name Type Description Default
result dict[str, ndarray]

The result of the optimize template program.

required

Returns:

Type Description
float

The mean SNR of the template.

Source code in src/leopard_em/pydantic_models/managers/optimize_template_manager.py
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
279
def results_to_snr(self, result: dict[str, np.ndarray]) -> float:
    """Convert optimize template result to mean SNR.

    Parameters
    ----------
    result : dict[str, np.ndarray]
        The result of the optimize template program.

    Returns
    -------
    float
        The mean SNR of the template.
    """
    # Filter out any infinite or NaN values
    # NOTE: There should not be NaNs or infs, will follow up later
    refined_scaled_mip = result["refined_z_score"]
    refined_scaled_mip = refined_scaled_mip[np.isfinite(refined_scaled_mip)]

    # If more than n values, keep only the top n highest SNRs
    best_n = 6
    if len(refined_scaled_mip) > best_n:
        refined_scaled_mip = np.sort(refined_scaled_mip)[-best_n:]

    # Printing out the results to console
    print(
        f"max snr: {refined_scaled_mip.max()}, min snr: {refined_scaled_mip.min()}"
    )

    mean_snr = float(refined_scaled_mip.mean())

    return mean_snr

run_optimize_template(output_text_path)

Run the refine template program and saves the resultant DataFrame to csv.

Parameters:

Name Type Description Default
output_text_path str

Path to save the optimized template pixel size.

required
Source code in src/leopard_em/pydantic_models/managers/optimize_template_manager.py
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
def run_optimize_template(self, output_text_path: str) -> None:
    """Run the refine template program and saves the resultant DataFrame to csv.

    Parameters
    ----------
    output_text_path : str
        Path to save the optimized template pixel size.
    """
    if self.pixel_size_coarse_search.enabled:
        # Create a file for logging all iterations
        all_results_path = self._get_all_results_path(output_text_path)
        # Create the file and write header
        with open(all_results_path, "w", encoding="utf-8") as f:
            f.write("Pixel Size (Ã…),SNR\n")

        optimal_template_px = self.optimize_pixel_size(all_results_path)
        print(f"Optimal template px: {optimal_template_px:.3f} Ã…")
        # print this to the text file
        with open(output_text_path, "w", encoding="utf-8") as f:
            f.write(f"Optimal template px: {optimal_template_px:.3f} Ã…")

RefineTemplateManager

Bases: BaseModel2DTM

Model holding parameters necessary for running the refine template program.

Attributes:

Name Type Description
template_volume_path str

Path to the template volume MRC file.

particle_stack ParticleStack

Particle stack object containing particle data.

defocus_refinement_config DefocusSearchConfig

Configuration for defocus refinement.

pixel_size_refinement_config PixelSizeSearchConfig

Configuration for pixel size refinement.

orientation_refinement_config RefineOrientationConfig

Configuration for orientation refinement.

preprocessing_filters PreprocessingFilters

Filters to apply to the particle images.

computational_config ComputationalConfig

What computational resources to allocate for the program.

template_volume ExcludedTensor

The template volume tensor (excluded from serialization).

Methods:

Name Description
TODO serialization/import methods
__init__

Initialize the refine template manager.

make_backend_core_function_kwargs

Create the kwargs for the backend refine_template core function.

run_refine_template

Run the refine template program.

Source code in src/leopard_em/pydantic_models/managers/refine_template_manager.py
 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
171
172
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
203
204
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
class RefineTemplateManager(BaseModel2DTM):
    """Model holding parameters necessary for running the refine template program.

    Attributes
    ----------
    template_volume_path : str
        Path to the template volume MRC file.
    particle_stack : ParticleStack
        Particle stack object containing particle data.
    defocus_refinement_config : DefocusSearchConfig
        Configuration for defocus refinement.
    pixel_size_refinement_config : PixelSizeSearchConfig
        Configuration for pixel size refinement.
    orientation_refinement_config : RefineOrientationConfig
        Configuration for orientation refinement.
    preprocessing_filters : PreprocessingFilters
        Filters to apply to the particle images.
    computational_config : ComputationalConfig
        What computational resources to allocate for the program.
    template_volume : ExcludedTensor
        The template volume tensor (excluded from serialization).

    Methods
    -------
    TODO serialization/import methods
    __init__(self, skip_mrc_preloads: bool = False, **data: Any)
        Initialize the refine template manager.
    make_backend_core_function_kwargs(self) -> dict[str, Any]
        Create the kwargs for the backend refine_template core function.
    run_refine_template(self, orientation_batch_size: int = 64) -> None
        Run the refine template program.
    """

    model_config: ClassVar = ConfigDict(arbitrary_types_allowed=True)

    template_volume_path: str  # In df per-particle, but ensure only one reference
    particle_stack: ParticleStack
    defocus_refinement_config: DefocusSearchConfig
    pixel_size_refinement_config: PixelSizeSearchConfig
    orientation_refinement_config: RefineOrientationConfig
    preprocessing_filters: PreprocessingFilters
    computational_config: ComputationalConfig

    # Excluded tensors
    template_volume: ExcludedTensor

    def __init__(self, skip_mrc_preloads: bool = False, **data: Any):
        super().__init__(**data)

        # Load the data from the MRC files
        if not skip_mrc_preloads:
            self.template_volume = load_mrc_volume(self.template_volume_path)

    def make_backend_core_function_kwargs(
        self, prefer_refined_angles: bool = True
    ) -> dict[str, Any]:
        """Create the kwargs for the backend refine_template core function.

        Parameters
        ----------
        prefer_refined_angles : bool
            Whether to use the refined angles from the particle stack. Defaults to
            False.
        """
        # Ensure the template is loaded in as a Tensor object
        template = load_template_tensor(
            template_volume=self.template_volume,
            template_volume_path=self.template_volume_path,
        )

        # The set of "best" euler angles from match template search
        # Check if refined angles exist, otherwise use the original angles
        euler_angles = self.particle_stack.get_euler_angles(prefer_refined_angles)

        # The relative Euler angle offsets to search over
        euler_angle_offsets = self.orientation_refinement_config.euler_angles_offsets

        # The relative defocus values to search over
        defocus_offsets = self.defocus_refinement_config.defocus_values

        # The relative pixel size values to search over
        pixel_size_offsets = self.pixel_size_refinement_config.pixel_size_values

        # Use the common utility function to set up the backend kwargs
        # pylint: disable=duplicate-code
        return setup_particle_backend_kwargs(
            particle_stack=self.particle_stack,
            template=template,
            preprocessing_filters=self.preprocessing_filters,
            euler_angles=euler_angles,
            euler_angle_offsets=euler_angle_offsets,
            defocus_offsets=defocus_offsets,
            pixel_size_offsets=pixel_size_offsets,
            device_list=self.computational_config.gpu_devices,
        )

    def run_refine_template(
        self, output_dataframe_path: str, orientation_batch_size: int = 64
    ) -> None:
        """Run the refine template program and saves the resultant DataFrame to csv.

        Parameters
        ----------
        output_dataframe_path : str
            Path to save the refined particle data.
        orientation_batch_size : int
            Number of orientations to process at once. Defaults to 64.
        """
        backend_kwargs = self.make_backend_core_function_kwargs()

        result = self.get_refine_result(backend_kwargs, orientation_batch_size)

        self.refine_result_to_dataframe(
            output_dataframe_path=output_dataframe_path, result=result
        )

    def get_refine_result(
        self, backend_kwargs: dict, orientation_batch_size: int = 64
    ) -> dict[str, np.ndarray]:
        """Get refine template result.

        Parameters
        ----------
        backend_kwargs : dict
            Keyword arguments for the backend processing
        orientation_batch_size : int
            Number of orientations to process at once. Defaults to 64.

        Returns
        -------
        dict[str, np.ndarray]
            The result of the refine template program.
        """
        # Adjust batch size if orientation search is disabled
        if not self.orientation_refinement_config.enabled:
            orientation_batch_size = 1
        elif (
            self.orientation_refinement_config.euler_angles_offsets.shape[0]
            < orientation_batch_size
        ):
            orientation_batch_size = (
                self.orientation_refinement_config.euler_angles_offsets.shape[0]
            )

        # pylint: disable=duplicate-code
        result: dict[str, np.ndarray] = {}
        result = core_refine_template(
            batch_size=orientation_batch_size, **backend_kwargs
        )
        result = {k: v.cpu().numpy() for k, v in result.items()}

        return result

    def refine_result_to_dataframe(
        self, output_dataframe_path: str, result: dict[str, np.ndarray]
    ) -> None:
        """Convert refine template result to dataframe.

        Parameters
        ----------
        output_dataframe_path : str
            Path to save the refined particle data.
        result : dict[str, np.ndarray]
            The result of the refine template program.
        """
        # pylint: disable=duplicate-code
        df_refined = self.particle_stack._df.copy()  # pylint: disable=protected-access

        # x and y positions
        pos_offset_y = result["refined_pos_y"]
        pos_offset_x = result["refined_pos_x"]
        pos_offset_y_ang = pos_offset_y * df_refined["pixel_size"]
        pos_offset_x_ang = pos_offset_x * df_refined["pixel_size"]

        df_refined["refined_pos_y"] = pos_offset_y + df_refined["pos_y"]
        df_refined["refined_pos_x"] = pos_offset_x + df_refined["pos_x"]
        df_refined["refined_pos_y_img"] = pos_offset_y + df_refined["pos_y_img"]
        df_refined["refined_pos_x_img"] = pos_offset_x + df_refined["pos_x_img"]
        df_refined["refined_pos_y_img_angstrom"] = (
            pos_offset_y_ang + df_refined["pos_y_img_angstrom"]
        )
        df_refined["refined_pos_x_img_angstrom"] = (
            pos_offset_x_ang + df_refined["pos_x_img_angstrom"]
        )

        # Euler angles
        df_refined["refined_psi"] = result["refined_euler_angles"][:, 2]
        df_refined["refined_theta"] = result["refined_euler_angles"][:, 1]
        df_refined["refined_phi"] = result["refined_euler_angles"][:, 0]

        # Defocus
        # Check if refined_relative_defocus already exists in the dataframe
        if "refined_relative_defocus" in df_refined.columns:
            df_refined["refined_relative_defocus"] = (
                result["refined_defocus_offset"]
                + df_refined["refined_relative_defocus"]
            )
        else:
            # If not, create it from relative_defocus
            df_refined["refined_relative_defocus"] = (
                result["refined_defocus_offset"] + df_refined["relative_defocus"]
            )

        # Pixel size
        df_refined["refined_pixel_size"] = (
            result["refined_pixel_size_offset"] + df_refined["pixel_size"]
        )

        # Cross-correlation statistics
        # Check if correlation statistic files exist and use them if available
        # This allows for shifts during refinement

        # if (
        #    "correlation_average_path" in df_refined.columns
        #    and "correlation_variance_path" in df_refined.columns
        # ):
        # Check if files exist for at least the first entry
        #    if (
        #        df_refined["correlation_average_path"].iloc[0]
        #        and df_refined["correlation_variance_path"].iloc[0]
        #    ):
        # Load the correlation statistics from the files
        #        correlation_average = read_mrc_to_numpy(
        #            df_refined["correlation_average_path"].iloc[0]
        #        )
        #        correlation_variance = read_mrc_to_numpy(
        #            df_refined["correlation_variance_path"].iloc[0]
        #        )
        #        df_refined["correlation_mean"] = correlation_average[
        #            df_refined["refined_pos_y"], df_refined["refined_pos_x"]
        #           ]
        #        df_refined["correlation_variance"] = correlation_variance[
        #            df_refined["refined_pos_y"], df_refined["refined_pos_x"]
        #        ]
        refined_mip = result["refined_cross_correlation"]
        refined_scaled_mip = result["refined_z_score"]
        df_refined["refined_mip"] = refined_mip
        df_refined["refined_scaled_mip"] = refined_scaled_mip

        # Reorder the columns
        df_refined = df_refined.reindex(columns=REFINED_DF_COLUMN_ORDER)

        # Save the refined DataFrame to disk
        df_refined.to_csv(output_dataframe_path)

get_refine_result(backend_kwargs, orientation_batch_size=64)

Get refine template result.

Parameters:

Name Type Description Default
backend_kwargs dict

Keyword arguments for the backend processing

required
orientation_batch_size int

Number of orientations to process at once. Defaults to 64.

64

Returns:

Type Description
dict[str, ndarray]

The result of the refine template program.

Source code in src/leopard_em/pydantic_models/managers/refine_template_manager.py
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
171
172
173
174
def get_refine_result(
    self, backend_kwargs: dict, orientation_batch_size: int = 64
) -> dict[str, np.ndarray]:
    """Get refine template result.

    Parameters
    ----------
    backend_kwargs : dict
        Keyword arguments for the backend processing
    orientation_batch_size : int
        Number of orientations to process at once. Defaults to 64.

    Returns
    -------
    dict[str, np.ndarray]
        The result of the refine template program.
    """
    # Adjust batch size if orientation search is disabled
    if not self.orientation_refinement_config.enabled:
        orientation_batch_size = 1
    elif (
        self.orientation_refinement_config.euler_angles_offsets.shape[0]
        < orientation_batch_size
    ):
        orientation_batch_size = (
            self.orientation_refinement_config.euler_angles_offsets.shape[0]
        )

    # pylint: disable=duplicate-code
    result: dict[str, np.ndarray] = {}
    result = core_refine_template(
        batch_size=orientation_batch_size, **backend_kwargs
    )
    result = {k: v.cpu().numpy() for k, v in result.items()}

    return result

make_backend_core_function_kwargs(prefer_refined_angles=True)

Create the kwargs for the backend refine_template core function.

Parameters:

Name Type Description Default
prefer_refined_angles bool

Whether to use the refined angles from the particle stack. Defaults to False.

True
Source code in src/leopard_em/pydantic_models/managers/refine_template_manager.py
 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
def make_backend_core_function_kwargs(
    self, prefer_refined_angles: bool = True
) -> dict[str, Any]:
    """Create the kwargs for the backend refine_template core function.

    Parameters
    ----------
    prefer_refined_angles : bool
        Whether to use the refined angles from the particle stack. Defaults to
        False.
    """
    # Ensure the template is loaded in as a Tensor object
    template = load_template_tensor(
        template_volume=self.template_volume,
        template_volume_path=self.template_volume_path,
    )

    # The set of "best" euler angles from match template search
    # Check if refined angles exist, otherwise use the original angles
    euler_angles = self.particle_stack.get_euler_angles(prefer_refined_angles)

    # The relative Euler angle offsets to search over
    euler_angle_offsets = self.orientation_refinement_config.euler_angles_offsets

    # The relative defocus values to search over
    defocus_offsets = self.defocus_refinement_config.defocus_values

    # The relative pixel size values to search over
    pixel_size_offsets = self.pixel_size_refinement_config.pixel_size_values

    # Use the common utility function to set up the backend kwargs
    # pylint: disable=duplicate-code
    return setup_particle_backend_kwargs(
        particle_stack=self.particle_stack,
        template=template,
        preprocessing_filters=self.preprocessing_filters,
        euler_angles=euler_angles,
        euler_angle_offsets=euler_angle_offsets,
        defocus_offsets=defocus_offsets,
        pixel_size_offsets=pixel_size_offsets,
        device_list=self.computational_config.gpu_devices,
    )

refine_result_to_dataframe(output_dataframe_path, result)

Convert refine template result to dataframe.

Parameters:

Name Type Description Default
output_dataframe_path str

Path to save the refined particle data.

required
result dict[str, ndarray]

The result of the refine template program.

required
Source code in src/leopard_em/pydantic_models/managers/refine_template_manager.py
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
203
204
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
def refine_result_to_dataframe(
    self, output_dataframe_path: str, result: dict[str, np.ndarray]
) -> None:
    """Convert refine template result to dataframe.

    Parameters
    ----------
    output_dataframe_path : str
        Path to save the refined particle data.
    result : dict[str, np.ndarray]
        The result of the refine template program.
    """
    # pylint: disable=duplicate-code
    df_refined = self.particle_stack._df.copy()  # pylint: disable=protected-access

    # x and y positions
    pos_offset_y = result["refined_pos_y"]
    pos_offset_x = result["refined_pos_x"]
    pos_offset_y_ang = pos_offset_y * df_refined["pixel_size"]
    pos_offset_x_ang = pos_offset_x * df_refined["pixel_size"]

    df_refined["refined_pos_y"] = pos_offset_y + df_refined["pos_y"]
    df_refined["refined_pos_x"] = pos_offset_x + df_refined["pos_x"]
    df_refined["refined_pos_y_img"] = pos_offset_y + df_refined["pos_y_img"]
    df_refined["refined_pos_x_img"] = pos_offset_x + df_refined["pos_x_img"]
    df_refined["refined_pos_y_img_angstrom"] = (
        pos_offset_y_ang + df_refined["pos_y_img_angstrom"]
    )
    df_refined["refined_pos_x_img_angstrom"] = (
        pos_offset_x_ang + df_refined["pos_x_img_angstrom"]
    )

    # Euler angles
    df_refined["refined_psi"] = result["refined_euler_angles"][:, 2]
    df_refined["refined_theta"] = result["refined_euler_angles"][:, 1]
    df_refined["refined_phi"] = result["refined_euler_angles"][:, 0]

    # Defocus
    # Check if refined_relative_defocus already exists in the dataframe
    if "refined_relative_defocus" in df_refined.columns:
        df_refined["refined_relative_defocus"] = (
            result["refined_defocus_offset"]
            + df_refined["refined_relative_defocus"]
        )
    else:
        # If not, create it from relative_defocus
        df_refined["refined_relative_defocus"] = (
            result["refined_defocus_offset"] + df_refined["relative_defocus"]
        )

    # Pixel size
    df_refined["refined_pixel_size"] = (
        result["refined_pixel_size_offset"] + df_refined["pixel_size"]
    )

    # Cross-correlation statistics
    # Check if correlation statistic files exist and use them if available
    # This allows for shifts during refinement

    # if (
    #    "correlation_average_path" in df_refined.columns
    #    and "correlation_variance_path" in df_refined.columns
    # ):
    # Check if files exist for at least the first entry
    #    if (
    #        df_refined["correlation_average_path"].iloc[0]
    #        and df_refined["correlation_variance_path"].iloc[0]
    #    ):
    # Load the correlation statistics from the files
    #        correlation_average = read_mrc_to_numpy(
    #            df_refined["correlation_average_path"].iloc[0]
    #        )
    #        correlation_variance = read_mrc_to_numpy(
    #            df_refined["correlation_variance_path"].iloc[0]
    #        )
    #        df_refined["correlation_mean"] = correlation_average[
    #            df_refined["refined_pos_y"], df_refined["refined_pos_x"]
    #           ]
    #        df_refined["correlation_variance"] = correlation_variance[
    #            df_refined["refined_pos_y"], df_refined["refined_pos_x"]
    #        ]
    refined_mip = result["refined_cross_correlation"]
    refined_scaled_mip = result["refined_z_score"]
    df_refined["refined_mip"] = refined_mip
    df_refined["refined_scaled_mip"] = refined_scaled_mip

    # Reorder the columns
    df_refined = df_refined.reindex(columns=REFINED_DF_COLUMN_ORDER)

    # Save the refined DataFrame to disk
    df_refined.to_csv(output_dataframe_path)

run_refine_template(output_dataframe_path, orientation_batch_size=64)

Run the refine template program and saves the resultant DataFrame to csv.

Parameters:

Name Type Description Default
output_dataframe_path str

Path to save the refined particle data.

required
orientation_batch_size int

Number of orientations to process at once. Defaults to 64.

64
Source code in src/leopard_em/pydantic_models/managers/refine_template_manager.py
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
def run_refine_template(
    self, output_dataframe_path: str, orientation_batch_size: int = 64
) -> None:
    """Run the refine template program and saves the resultant DataFrame to csv.

    Parameters
    ----------
    output_dataframe_path : str
        Path to save the refined particle data.
    orientation_batch_size : int
        Number of orientations to process at once. Defaults to 64.
    """
    backend_kwargs = self.make_backend_core_function_kwargs()

    result = self.get_refine_result(backend_kwargs, orientation_batch_size)

    self.refine_result_to_dataframe(
        output_dataframe_path=output_dataframe_path, result=result
    )