Computes a soft distance function to the atom centers of a protein. Implements Eq. (1) of the paper in a fast and numerically stable way. Args: x (Tensor): (N,3) atom centers. y (Tensor): (M,3) sampling locations. batch_x (integer Tensor): (N,) batch vector for x, a
(x, y, batch_x, batch_y, smoothness=0.01, atomtypes=None)
| 154 | |
| 155 | |
| 156 | def soft_distances(x, y, batch_x, batch_y, smoothness=0.01, atomtypes=None): |
| 157 | """Computes a soft distance function to the atom centers of a protein. |
| 158 | |
| 159 | Implements Eq. (1) of the paper in a fast and numerically stable way. |
| 160 | |
| 161 | Args: |
| 162 | x (Tensor): (N,3) atom centers. |
| 163 | y (Tensor): (M,3) sampling locations. |
| 164 | batch_x (integer Tensor): (N,) batch vector for x, as in PyTorch_geometric. |
| 165 | batch_y (integer Tensor): (M,) batch vector for y, as in PyTorch_geometric. |
| 166 | smoothness (float, optional): atom radii if atom types are not provided. Defaults to .01. |
| 167 | atomtypes (integer Tensor, optional): (N,6) one-hot encoding of the atom chemical types. Defaults to None. |
| 168 | |
| 169 | Returns: |
| 170 | Tensor: (M,) values of the soft distance function on the points `y`. |
| 171 | """ |
| 172 | # Build the (N, M, 1) symbolic matrix of squared distances: |
| 173 | x_i = LazyTensor(x[:, None, :]) # (N, 1, 3) atoms |
| 174 | y_j = LazyTensor(y[None, :, :]) # (1, M, 3) sampling points |
| 175 | D_ij = ((x_i - y_j) ** 2).sum(-1) # (N, M, 1) squared distances |
| 176 | |
| 177 | # Use a block-diagonal sparsity mask to support heterogeneous batch processing: |
| 178 | D_ij.ranges = diagonal_ranges(batch_x, batch_y) |
| 179 | |
| 180 | if atomtypes is not None: |
| 181 | # Turn the one-hot encoding "atomtypes" into a vector of diameters "smoothness_i": |
| 182 | # (N, 6) -> (N, 1, 1) (There are 6 atom types) |
| 183 | atomic_radii = torch.tensor( |
| 184 | [170, 110, 152, 155, 180, 190], dtype=torch.float, device=x.device |
| 185 | ) |
| 186 | |
| 187 | atomic_radii = atomic_radii / atomic_radii.min() |
| 188 | atomtype_radii = atomtypes * atomic_radii[None, :] # n_atoms, n_atomtypes |
| 189 | # smoothness = atomtypes @ atomic_radii # (N, 6) @ (6,) = (N,) |
| 190 | smoothness = torch.sum( |
| 191 | smoothness * atomtype_radii, dim=1, keepdim=False |
| 192 | ) # n_atoms, 1 |
| 193 | smoothness_i = LazyTensor(smoothness[:, None, None]) |
| 194 | |
| 195 | # Compute an estimation of the mean smoothness in a neighborhood |
| 196 | # of each sampling point: |
| 197 | # density = (-D_ij.sqrt()).exp().sum(0).view(-1) # (M,) local density of atoms |
| 198 | # smooth = (smoothness_i * (-D_ij.sqrt()).exp()).sum(0).view(-1) # (M,) |
| 199 | # mean_smoothness = smooth / density # (M,) |
| 200 | |
| 201 | # soft_dists = -mean_smoothness * ( |
| 202 | # (-D_ij.sqrt() / smoothness_i).logsumexp(dim=0) |
| 203 | # ).view(-1) |
| 204 | mean_smoothness = (-D_ij.sqrt()).exp() |
| 205 | mean_smoothness = mean_smoothness.sum(0) |
| 206 | mean_smoothness_j = LazyTensor(mean_smoothness[None, :, :]) |
| 207 | mean_smoothness = ( |
| 208 | smoothness_i * (-D_ij.sqrt()).exp() / mean_smoothness_j |
| 209 | ) # n_atoms, n_points, 1 |
| 210 | mean_smoothness = mean_smoothness.sum(0).view(-1) |
| 211 | soft_dists = -mean_smoothness * ( |
| 212 | (-D_ij.sqrt() / smoothness_i).logsumexp(dim=0) |
| 213 | ).view(-1) |
no test coverage detected