Skip to content

refine_template_manager

Pydantic model for running the refine template program.

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
    )