Skip to content

process_results

Functions related to result processing after backend functions.

aggregate_distributed_results(results)

Combine the 2DTM results from multiple devices.

NOTE: This assumes that all tensors have been passed back to the CPU and are in the form of numpy arrays.

Parameters:

Name Type Description Default
results list[dict[str, ndarray]]

List of dictionaries containing the results from each device. Each dictionary contains the following keys: - "mip": Maximum intensity projection of the cross-correlation values. - "best_phi": Best phi angle for each pixel. - "best_theta": Best theta angle for each pixel. - "best_psi": Best psi angle for each pixel. - "best_defocus": Best defocus value for each pixel. - "best_pixel_size": Best pixel size value for each pixel. - "correlation_sum": Sum of cross-correlation values for each pixel. - "correlation_squared_sum": Sum of squared cross-correlation values for each pixel. - "total_projections": Total number of projections calculated.

 required
Source code in src/leopard_em/backend/process_results.py 
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
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
def aggregate_distributed_results(
    results: list[dict[str, torch.Tensor | np.ndarray]],
) -> dict[str, torch.Tensor]:
    """Combine the 2DTM results from multiple devices.

    NOTE: This assumes that all tensors have been passed back to the CPU and are in
    the form of numpy arrays.

    Parameters
    ----------
    results : list[dict[str, np.ndarray]]
        List of dictionaries containing the results from each device. Each dictionary
        contains the following keys:
            - "mip": Maximum intensity projection of the cross-correlation values.
            - "best_phi": Best phi angle for each pixel.
            - "best_theta": Best theta angle for each pixel.
            - "best_psi": Best psi angle for each pixel.
            - "best_defocus": Best defocus value for each pixel.
            - "best_pixel_size": Best pixel size value for each pixel.
            - "correlation_sum": Sum of cross-correlation values for each pixel.
            - "correlation_squared_sum": Sum of squared cross-correlation values for
              each pixel.
            - "total_projections": Total number of projections calculated.
    """
    # Ensure all the tensors are passed back to CPU as numpy arrays
    # Not sure why cannot sync across devices, but this is a workaround
    results = [
        {
            key: value.cpu().numpy() if isinstance(value, torch.Tensor) else value
            for key, value in result.items()
        }
        for result in results
    ]

    # Find which device had the highest MIP for each pixel and index stats accordingly.
    # Results after 'take_along_axis' have extra dimension at idx 0.
    mips = np.stack([result["mip"] for result in results], axis=0)
    best_phi = np.stack([result["best_phi"] for result in results], axis=0)
    best_theta = np.stack([result["best_theta"] for result in results], axis=0)
    best_psi = np.stack([result["best_psi"] for result in results], axis=0)
    best_defocus = np.stack([result["best_defocus"] for result in results], axis=0)
    best_pixel_size = np.stack(
        [result["best_pixel_size"] for result in results], axis=0
    )
    mip_max = mips.max(axis=0)
    mip_argmax = mips.argmax(axis=0)

    best_phi = np.take_along_axis(best_phi, mip_argmax[None, ...], axis=0)
    best_theta = np.take_along_axis(best_theta, mip_argmax[None, ...], axis=0)
    best_psi = np.take_along_axis(best_psi, mip_argmax[None, ...], axis=0)
    best_defocus = np.take_along_axis(best_defocus, mip_argmax[None, ...], axis=0)
    best_pixel_size = np.take_along_axis(best_pixel_size, mip_argmax[None, ...], axis=0)
    best_phi = best_phi[0]
    best_theta = best_theta[0]
    best_psi = best_psi[0]
    best_defocus = best_defocus[0]
    best_pixel_size = best_pixel_size[0]
    # Sum the sums and squared sums of the cross-correlation values
    correlation_sum = np.stack(
        [result["correlation_sum"] for result in results], axis=0
    ).sum(axis=0)
    correlation_squared_sum = np.stack(
        [result["correlation_squared_sum"] for result in results], axis=0
    ).sum(axis=0)

    # NOTE: Currently only tracking total number of projections for statistics,
    # but could be future case where number of projections calculated on each
    # device is necessary for some statistical computation.
    total_projections = sum(result["total_projections"] for result in results)

    # Cast back to torch tensors on the CPU
    mip_max = torch.from_numpy(mip_max)
    best_phi = torch.from_numpy(best_phi)
    best_theta = torch.from_numpy(best_theta)
    best_psi = torch.from_numpy(best_psi)
    best_defocus = torch.from_numpy(best_defocus)
    best_pixel_size = torch.from_numpy(best_pixel_size)
    correlation_sum = torch.from_numpy(correlation_sum)
    correlation_squared_sum = torch.from_numpy(correlation_squared_sum)

    return {
        "mip": mip_max,
        "best_phi": best_phi,
        "best_theta": best_theta,
        "best_psi": best_psi,
        "best_defocus": best_defocus,
        "correlation_sum": correlation_sum,
        "correlation_squared_sum": correlation_squared_sum,
        "total_projections": total_projections,
    }

correlation_sum_and_squared_sum_to_mean_and_variance(correlation_sum, correlation_squared_sum, total_correlation_positions)

Convert the sum and squared sum of the correlation values to mean and variance.

Parameters:

Name Type Description Default
correlation_sum Tensor

Sum of the correlation values.

 required
correlation_squared_sum Tensor

Sum of the squared correlation values.

 required
total_correlation_positions int

Total number cross-correlograms calculated.

 required

Returns:

Type Description
tuple[Tensor, Tensor]

Tuple containing the mean and variance of the correlation values.
Source code in src/leopard_em/backend/process_results.py 
 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
def correlation_sum_and_squared_sum_to_mean_and_variance(
    correlation_sum: torch.Tensor,
    correlation_squared_sum: torch.Tensor,
    total_correlation_positions: int,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Convert the sum and squared sum of the correlation values to mean and variance.

    Parameters
    ----------
    correlation_sum : torch.Tensor
        Sum of the correlation values.
    correlation_squared_sum : torch.Tensor
        Sum of the squared correlation values.
    total_correlation_positions : int
        Total number cross-correlograms calculated.

    Returns
    -------
    tuple[torch.Tensor, torch.Tensor]
        Tuple containing the mean and variance of the correlation values.
    """
    correlation_mean = correlation_sum / total_correlation_positions
    correlation_variance = correlation_squared_sum / total_correlation_positions
    correlation_variance -= correlation_mean**2
    correlation_variance = torch.sqrt(torch.clamp(correlation_variance, min=0))
    return correlation_mean, correlation_variance

scale_mip(mip, mip_scaled, correlation_sum, correlation_squared_sum, total_correlation_positions)

Scale the MIP to Z-score map by the mean and variance of the correlation values.

Z-score is accounting for the variation in image intensity and spurious correlations by subtracting the mean and dividing by the standard deviation pixel-wise. Since cross-correlation values are roughly normally distributed for pure noise, Z-score effectively becomes a measure of how unexpected (highly correlated to the reference template) a region is in the image. Note that we are looking at maxima of millions of Gaussian distributions, so Z-score has to be compared with a generalized extreme value distribution (GEV) to determine significance (done elsewhere).

NOTE: This method also updates the correlation_sum and correlation_squared_sum tensors in-place into the mean and variance, respectively. Likely should reflect conversions in variable names...

Parameters:

Name Type Description Default
mip Tensor

MIP of the correlation values.

 required
mip_scaled Tensor

Scaled MIP of the correlation values.

 required
correlation_sum Tensor

Sum of the correlation values. Updated to mean of the correlation values.

 required
correlation_squared_sum Tensor

Sum of the squared correlation values. Updated to variance of the correlation.

 required
total_correlation_positions int

Total number cross-correlograms calculated.

 required

Returns:

Type Description
tuple[Tensor, Tensor]

Tuple containing, in order, the MIP, scaled MIP, correlation mean, and correlation variance.
Source code in src/leopard_em/backend/process_results.py 
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
def scale_mip(
    mip: torch.Tensor,
    mip_scaled: torch.Tensor,
    correlation_sum: torch.Tensor,
    correlation_squared_sum: torch.Tensor,
    total_correlation_positions: int,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
    """Scale the MIP to Z-score map by the mean and variance of the correlation values.

    Z-score is accounting for the variation in image intensity and spurious correlations
    by subtracting the mean and dividing by the standard deviation pixel-wise. Since
    cross-correlation values are roughly normally distributed for pure noise, Z-score
    effectively becomes a measure of how unexpected (highly correlated to the reference
    template) a region is in the image. Note that we are looking at maxima of millions
    of Gaussian distributions, so Z-score has to be compared with a generalized extreme
    value distribution (GEV) to determine significance (done elsewhere).

    NOTE: This method also updates the correlation_sum and correlation_squared_sum
    tensors in-place into the mean and variance, respectively. Likely should reflect
    conversions in variable names...

    Parameters
    ----------
    mip : torch.Tensor
        MIP of the correlation values.
    mip_scaled : torch.Tensor
        Scaled MIP of the correlation values.
    correlation_sum : torch.Tensor
        Sum of the correlation values. Updated to mean of the correlation values.
    correlation_squared_sum : torch.Tensor
        Sum of the squared correlation values. Updated to variance of the correlation.
    total_correlation_positions : int
        Total number cross-correlograms calculated.

    Returns
    -------
    tuple[torch.Tensor, torch.Tensor]
        Tuple containing, in order, the MIP, scaled MIP, correlation mean, and
        correlation variance.
    """
    corr_mean, corr_variance = correlation_sum_and_squared_sum_to_mean_and_variance(
        correlation_sum, correlation_squared_sum, total_correlation_positions
    )

    # Calculate normalized MIP
    mip_scaled = mip - corr_mean
    torch.where(
        corr_variance != 0,  # preventing zero division error, albeit unlikely
        mip_scaled / corr_variance,
        torch.zeros_like(mip_scaled),
        out=mip_scaled,
    )

    # # Update correlation_sum and correlation_squared_sum to mean and variance
    # correlation_sum.copy_(corr_mean)
    # correlation_squared_sum.copy_(corr_variance)

    return mip, mip_scaled, corr_mean, corr_variance