-
Notifications
You must be signed in to change notification settings - Fork 9
/
correlation.py
134 lines (113 loc) · 4.42 KB
/
correlation.py
1
2
3
4
5
6
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
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
"""
Functions in this file are cpu/gpu agnostic.
"""
import numpy.typing as npt
import cupy.typing as cpt
from typing import Optional, Union
def mean_under_mask(
data: Union[npt.NDArray[float], cpt.NDArray[float]],
mask: Union[npt.NDArray[float], cpt.NDArray[float]],
mask_weight: Optional[float] = None,
) -> Union[float, cpt.NDArray[float]]:
"""Calculate mean of array in the mask region.
data and mask can be cupy or numpy arrays.
Parameters
----------
data: Union[npt.NDArray[float], cpt.NDArray[float]]
input array
mask: Union[npt.NDArray[float], cpt.NDArray[float]]
input mask, same dimensions as data
mask_weight: Optional[float], default None
optional weight of mask, if not provided mask.sum() is used to determine weight
Returns
-------
output: Union[float, cpt.NDArray[float]]
mean of data in the region of the mask
"""
output = (data * mask).sum() / (
mask_weight if mask_weight is not None else mask.sum()
)
return output
def std_under_mask(
data: Union[npt.NDArray[float], cpt.NDArray[float]],
mask: Union[npt.NDArray[float], cpt.NDArray[float]],
mean: float,
mask_weight: Optional[float] = None,
) -> Union[float, cpt.NDArray[float]]:
"""Calculate standard deviation of array in the mask region. Uses mean_under_mask()
to calculate the mean of data**2 within the mask.
data and mask can be cupy or numpy arrays.
Parameters
----------
data: Union[npt.NDArray[float], cpt.NDArray[float]]
input array
mask: Union[npt.NDArray[float], cpt.NDArray[float]]
input mask, same dimensions as data
mean: float
mean of array in masked region
mask_weight: Optional[float], default None
optional weight of mask, if not provided mask.sum() is used to determine weight
Returns
-------
output: Union[float, cpt.NDArray[float]]
standard deviation of data in the region of the mask
"""
output = (mean_under_mask(data**2, mask, mask_weight=mask_weight) - mean**2) ** 0.5
return output
def normalise(
data: Union[npt.NDArray[float], cpt.NDArray[float]],
mask: Optional[Union[npt.NDArray[float], cpt.NDArray[float]]] = None,
mask_weight: Optional[float] = None,
) -> Union[npt.NDArray[float], cpt.NDArray[float]]:
"""Normalise array by subtracting mean and dividing by standard deviation. If a mask
is provided the array is normalised with the mean and std calculated within the
mask.
data and mask can be cupy or numpy arrays.
Parameters
----------
data: Union[npt.NDArray[float], cpt.NDArray[float]]
input array to normalise
mask: Optional[Union[npt.NDArray[float], cpt.NDArray[float]]], default None
optional mask to normalise with mean and std in masked region
mask_weight: Optional[float], default None
optional float specifying mask weight, if not provided mask.sum() is used
Returns
-------
output: Union[npt.NDArray[float], cpt.NDArray[float]]
normalised array
"""
if mask is None:
mean, std = data.mean(), data.std()
else:
mean = mean_under_mask(data, mask, mask_weight=mask_weight)
std = std_under_mask(data, mask, mean, mask_weight=mask_weight)
output = (data - mean) / std
return output
def normalised_cross_correlation(
data1: Union[npt.NDArray[float], cpt.NDArray[float]],
data2: Union[npt.NDArray[float], cpt.NDArray[float]],
mask: Optional[Union[npt.NDArray[float], cpt.NDArray[float]]] = None,
) -> Union[float, cpt.NDArray[float]]:
"""Calculate normalised cross correlation between two arrays. Optionally only in a
masked region.
data1, data2, and mask can be cupy or numpy arrays.
Parameters
----------
data1: Union[npt.NDArray[float], cpt.NDArray[float]]
first array for correlation
data2: Union[npt.NDArray[float], cpt.NDArray[float]]
second array for correlation
mask: Optional[Union[npt.NDArray[float], cpt.NDArray[float]]], default None
optional mask to calculate the correlation under
Returns
-------
output: Union[float, cpt.NDArray[float]]
normalised cross correlation between the arrays
"""
if mask is None:
output = (normalise(data1) * normalise(data2)).sum() / data1.size
else:
output = (
normalise(data1, mask) * mask * normalise(data2, mask)
).sum() / mask.sum()
return output