File size: 4,546 Bytes
3e62986 | 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 |
"""Functions relating to spatial locations of atoms within a structure."""
from collections.abc import Collection, Sequence
from flax_model.alphafold3 import structure
from flax_model.alphafold3.structure import mmcif
import numpy as np
import scipy
def _make_atom_has_clash_mask(
kd_query_result: np.ndarray,
struc: structure.Structure,
ignore_chains: Collection[str],
) -> np.ndarray:
"""Returns a boolean NumPy array representing whether each atom has a clash.
Args:
kd_query_result: NumPy array containing N-atoms arrays, each array
containing indices to atoms that clash with the N'th atom.
struc: Structure over which clashes were detected.
ignore_chains: Collection of chains that should not be considered clashing.
A boolean NumPy array of length N atoms.
"""
atom_is_clashing = np.zeros((struc.num_atoms,), dtype=bool)
for atom_index, clashes in enumerate(kd_query_result):
chain_i = struc.chain_id[atom_index]
if chain_i in ignore_chains:
continue
islig_i = struc.is_ligand_mask[atom_index]
for clashing_atom_index in clashes:
chain_c = struc.chain_id[clashing_atom_index]
if chain_c in ignore_chains:
continue
islig_c = struc.is_ligand_mask[clashing_atom_index]
if (
clashing_atom_index == atom_index
or chain_i == chain_c
or islig_i != islig_c
):
# Ignore clashes within chain or between ligand and polymer.
continue
atom_is_clashing[atom_index] = True
return atom_is_clashing
def find_clashing_chains(
struc: structure.Structure,
clash_thresh_angstrom: float = 1.7,
clash_thresh_fraction: float = 0.3,
) -> Sequence[str]:
"""Finds chains that clash with others.
Clashes are defined by polymer backbone atoms and all ligand atoms.
Ligand-polymer clashes are not dropped.
Will not find clashes if all coordinates are 0. Coordinates are all 0s if
the structure is generated from sequences only, as done for inference in
dendro for example.
Args:
struc: The structure defining the chains and atom positions.
clash_thresh_angstrom: Below this distance, atoms are considered clashing.
clash_thresh_fraction: Chains with more than this fraction of their atoms
considered clashing will be dropped. This value should be in the range (0,
1].
Returns:
A sequence of chain ids for chains that clash.
Raises:
ValueError: If `clash_thresh_fraction` is not in range (0,1].
"""
if not 0 < clash_thresh_fraction <= 1:
raise ValueError('clash_thresh_fraction must be in range (0,1]')
struc_backbone = struc.filter_polymers_to_single_atom_per_res()
if struc_backbone.num_chains == 0:
return []
# If the coordinates are all 0, do not search for clashes.
if not np.any(struc_backbone.coords):
return []
coord_kdtree = scipy.spatial.cKDTree(struc_backbone.coords)
# For each atom coordinate, find all atoms within the clash thresh radius.
clashing_per_atom = coord_kdtree.query_ball_point(
struc_backbone.coords, r=clash_thresh_angstrom
)
chain_ids = struc_backbone.chains
if struc_backbone.atom_occupancy is not None:
chain_occupancy = np.array([
np.mean(struc_backbone.atom_occupancy[start:end])
for start, end in struc_backbone.iter_chain_ranges()
])
else:
chain_occupancy = None
# Remove chains until no more significant clashing.
chains_to_remove = set()
for _ in range(len(chain_ids)):
# Calculate maximally clashing.
atom_has_clash = _make_atom_has_clash_mask(
clashing_per_atom, struc_backbone, chains_to_remove
)
clashes_per_chain = np.array([
atom_has_clash[start:end].mean()
for start, end in struc_backbone.iter_chain_ranges()
])
max_clash = np.max(clashes_per_chain)
if max_clash <= clash_thresh_fraction:
# None of the remaining chains exceed the clash fraction threshold, so
# we can exit.
break
# Greedily remove worst with the lowest occupancy.
most_clashes = np.nonzero(clashes_per_chain == max_clash)[0]
if chain_occupancy is not None:
occupancy_clashing = chain_occupancy[most_clashes]
last_lowest_occupancy = (
len(occupancy_clashing) - np.argmin(occupancy_clashing[::-1]) - 1
)
worst_and_last = most_clashes[last_lowest_occupancy]
else:
worst_and_last = most_clashes[-1]
chains_to_remove.add(chain_ids[worst_and_last])
return sorted(chains_to_remove, key=mmcif.str_id_to_int_id)
|