Synthyra/ESMFold2
0505
1from __future__ import annotations
2
3import os
4from collections import defaultdict
5from contextlib import nullcontext
6from dataclasses import is_dataclass
7from io import BytesIO
8from typing import (
9 Any,
10 ContextManager,
11 Generator,
12 Iterable,
13 Protocol,
14 Sequence,
15 TypeVar,
16 runtime_checkable,
17)
18from warnings import warn
19
20import huggingface_hub
21import numpy as np
22import torch
23import zstd
24
25from .esmfold2_constants_esm3 import CHAIN_BREAK_STR
26from .esmfold2_utils_types import FunctionAnnotation
27
28MAX_SUPPORTED_DISTANCE = 1e6
29
30
31TSequence = TypeVar("TSequence", bound=Sequence)
32
33
34@runtime_checkable
35class Concatable(Protocol):
36 @classmethod
37 def concat(cls, objs: list[Concatable]) -> Concatable: ...
38
39
40def slice_python_object_as_numpy(
41 obj: TSequence, idx: int | list[int] | slice | np.ndarray
42) -> TSequence:
43 """
44 Slice a python object (like a list, string, or tuple) as if it was a numpy object.
45
46 Example:
47 >>> obj = "ABCDE"
48 >>> slice_python_object_as_numpy(obj, [1, 3, 4])
49 "BDE"
50
51 >>> obj = [1, 2, 3, 4, 5]
52 >>> slice_python_object_as_numpy(obj, np.arange(5) < 3)
53 [1, 2, 3]
54 """
55 if np.isscalar(idx):
56 idx = [int(idx)] # type: ignore
57
58 if isinstance(idx, np.ndarray) and idx.dtype == bool:
59 sliced_obj = [obj[i] for i in np.where(idx)[0]]
60 elif isinstance(idx, slice):
61 sliced_obj = obj[idx]
62 else:
63 sliced_obj = [obj[i] for i in idx] # type: ignore
64
65 match obj, sliced_obj:
66 case str(), list():
67 sliced_obj = "".join(sliced_obj)
68 case _:
69 sliced_obj = obj.__class__(sliced_obj) # type: ignore
70
71 return sliced_obj # type: ignore
72
73
74def slice_any_object(
75 obj: TSequence, idx: int | list[int] | slice | np.ndarray
76) -> TSequence:
77 """
78 Slice a arbitrary object (like a list, string, or tuple) as if it was a numpy object. Similar to `slice_python_object_as_numpy`, but detects if it's a numpy array or Tensor and uses the existing slice method if so.
79
80 If the object is a dataclass, it will simply apply the index to the object, under the assumption that the object has correcty implemented numpy indexing.
81
82 Example:
83 >>> obj = "ABCDE"
84 >>> slice_any_object(obj, [1, 3, 4])
85 "BDE"
86
87 >>> obj = np.array([1, 2, 3, 4, 5])
88 >>> slice_any_object(obj, np.arange(5) < 3)
89 np.array([1, 2, 3])
90
91 >>> obj = ProteinChain.from_rcsb("1a3a", "A")
92 >>> slice_any_object(obj, np.arange(len(obj)) < 10)
93 # ProteinChain w/ length 10
94
95 """
96 if isinstance(obj, (np.ndarray, torch.Tensor)):
97 return obj[idx] # type: ignore
98 elif is_dataclass(obj):
99 # if passing a dataclass, assume it implements a custom slice
100 return obj[idx] # type: ignore
101 else:
102 return slice_python_object_as_numpy(obj, idx)
103
104
105def rbf(values, v_min, v_max, n_bins=16):
106 """
107 Returns RBF encodings in a new dimension at the end.
108 """
109 rbf_centers = torch.linspace(
110 v_min, v_max, n_bins, device=values.device, dtype=values.dtype
111 )
112 rbf_centers = rbf_centers.view([1] * len(values.shape) + [-1])
113 rbf_std = (v_max - v_min) / n_bins
114 z = (values.unsqueeze(-1) - rbf_centers) / rbf_std
115 return torch.exp(-(z**2))
116
117
118def batched_gather(data, inds, dim=0, no_batch_dims=0):
119 ranges = []
120 for i, s in enumerate(data.shape[:no_batch_dims]):
121 r = torch.arange(s)
122 r = r.view(*(*((1,) * i), -1, *((1,) * (len(inds.shape) - i - 1))))
123 ranges.append(r)
124
125 remaining_dims = [slice(None) for _ in range(len(data.shape) - no_batch_dims)]
126 remaining_dims[dim - no_batch_dims if dim >= 0 else dim] = inds
127 ranges.extend(remaining_dims)
128 return data[ranges]
129
130
131def node_gather(s: torch.Tensor, edges: torch.Tensor) -> torch.Tensor:
132 return batched_gather(s.unsqueeze(-3), edges, -2, no_batch_dims=len(s.shape) - 1)
133
134
135def knn_graph(
136 coords: torch.Tensor,
137 coord_mask: torch.Tensor,
138 padding_mask: torch.Tensor,
139 sequence_id: torch.Tensor,
140 *,
141 no_knn: int,
142):
143 L = coords.shape[-2]
144 num_by_dist = min(no_knn, L)
145 device = coords.device
146
147 coords = coords.nan_to_num()
148 coord_mask = ~(coord_mask[..., None, :] & coord_mask[..., :, None])
149 padding_pairwise_mask = padding_mask[..., None, :] | padding_mask[..., :, None]
150 if sequence_id is not None:
151 padding_pairwise_mask |= torch.unsqueeze(sequence_id, 1) != torch.unsqueeze(
152 sequence_id, 2
153 )
154 dists = (coords.unsqueeze(-2) - coords.unsqueeze(-3)).norm(dim=-1)
155 arange = torch.arange(L, device=device)
156 seq_dists = (arange.unsqueeze(-1) - arange.unsqueeze(-2)).abs()
157 # We only support up to a certain distance, above that, we use sequence distance
158 # instead. This is so that when a large portion of the structure is masked out,
159 # the edges are built according to sequence distance.
160 max_dist = MAX_SUPPORTED_DISTANCE
161 if not (dists[~coord_mask] < max_dist).all():
162 raise ValueError(
163 f"Coordinate pairwise distances exceed max supported distance ({max_dist}). "
164 )
165 struct_then_seq_dist = (
166 seq_dists.to(dists.dtype)
167 .mul(1e2)
168 .add(max_dist)
169 .where(coord_mask, dists)
170 .masked_fill(padding_pairwise_mask, torch.inf)
171 )
172 dists, edges = struct_then_seq_dist.sort(dim=-1, descending=False)
173 # This is a L x L tensor, where we index by rows first,
174 # and columns are the edges we should pick.
175 chosen_edges = edges[..., :num_by_dist]
176 chosen_mask = dists[..., :num_by_dist].isfinite()
177 return chosen_edges, chosen_mask
178
179
180def stack_variable_length_tensors(
181 sequences: Sequence[torch.Tensor],
182 constant_value: int | float = 0,
183 dtype: torch.dtype | None = None,
184) -> torch.Tensor:
185 """Automatically stack tensors together, padding variable lengths with the
186 value in constant_value. Handles an arbitrary number of dimensions.
187
188 Examples:
189 >>> tensor1, tensor2 = torch.ones([2]), torch.ones([5])
190 >>> stack_variable_length_tensors(tensor1, tensor2)
191 tensor of shape [2, 5]. First row is [1, 1, 0, 0, 0]. Second row is all ones.
192
193 >>> tensor1, tensor2 = torch.ones([2, 4]), torch.ones([5, 3])
194 >>> stack_variable_length_tensors(tensor1, tensor2)
195 tensor of shape [2, 5, 4]
196 """
197 batch_size = len(sequences)
198 shape = [batch_size] + np.max([seq.shape for seq in sequences], 0).tolist()
199
200 if dtype is None:
201 dtype = sequences[0].dtype
202 device = sequences[0].device
203
204 array = torch.full(shape, constant_value, dtype=dtype, device=device)
205 for arr, seq in zip(array, sequences):
206 arrslice = tuple(slice(dim) for dim in seq.shape)
207 arr[arrslice] = seq
208
209 return array
210
211
212def binpack(
213 tensor: torch.Tensor, sequence_id: torch.Tensor | None, pad_value: int | float
214):
215 """
216 Args:
217 tensor (Tensor): [B, L, ...]
218
219 Returns:
220 Tensor: [B_binpacked, L_binpacked, ...]
221 """
222 if sequence_id is None:
223 return tensor
224
225 num_sequences = sequence_id.max(dim=-1).values + 1
226
227 dims = sequence_id.shape + tensor.shape[2:]
228 output_tensor = torch.full(
229 dims, fill_value=pad_value, dtype=tensor.dtype, device=tensor.device
230 )
231
232 idx = 0
233 for batch_idx, (batch_seqid, batch_num_sequences) in enumerate(
234 zip(sequence_id, num_sequences)
235 ):
236 for seqid in range(batch_num_sequences):
237 mask = batch_seqid == seqid
238 output_tensor[batch_idx, mask] = tensor[idx, : mask.sum()]
239 idx += 1
240 return output_tensor
241
242
243def unbinpack(
244 tensor: torch.Tensor, sequence_id: torch.Tensor | None, pad_value: int | float
245):
246 """
247 Args:
248 tensor (Tensor): [B, L, ...]
249
250 Returns:
251 Tensor: [B_unbinpacked, L_unbinpack, ...]
252 """
253 if sequence_id is None:
254 return tensor
255
256 unpacked_tensors = []
257 num_sequences = sequence_id.max(dim=-1).values + 1
258 for batch_idx, (batch_seqid, batch_num_sequences) in enumerate(
259 zip(sequence_id, num_sequences)
260 ):
261 for seqid in range(batch_num_sequences):
262 mask = batch_seqid == seqid
263 unpacked = tensor[batch_idx, mask]
264 unpacked_tensors.append(unpacked)
265 return stack_variable_length_tensors(unpacked_tensors, pad_value)
266
267
268def fp32_autocast_context(device_type: str) -> ContextManager[Any]: # type: ignore
269 """
270 Returns an autocast context manager that disables downcasting by AMP.
271
272 Args:
273 device_type: The device type ('cpu' or 'cuda')
274
275 Returns:
276 An autocast context manager with the specified behavior.
277 """
278 if device_type == "cpu":
279 return torch.amp.autocast(device_type, enabled=False) # type: ignore
280 elif device_type == "mps":
281 # For MPS, just return a no-op context manager (nullcontext) since MPS does not support autocast.
282 return nullcontext()
283 elif device_type == "cuda":
284 return torch.amp.autocast(device_type, dtype=torch.float32) # type: ignore
285 else:
286 raise ValueError(f"Unsupported device type: {device_type}")
287
288
289def merge_ranges(ranges: list[range], merge_gap_max: int | None = None) -> list[range]:
290 """Merge overlapping ranges into sorted, non-overlapping segments.
291
292 Args:
293 ranges: collection of ranges to merge.
294 merge_gap_max: optionally merge neighboring ranges that are separated by a gap
295 no larger than this size.
296 Returns:
297 non-overlapping ranges merged from the inputs, sorted by position.
298 """
299 ranges = sorted(ranges, key=lambda r: r.start)
300 merge_gap_max = merge_gap_max if merge_gap_max is not None else 0
301 assert merge_gap_max >= 0, f"Invalid merge_gap_max: {merge_gap_max}"
302
303 merged = []
304 for r in ranges:
305 if not merged:
306 merged.append(r)
307 else:
308 last = merged[-1]
309 if last.stop + merge_gap_max >= r.start:
310 merged[-1] = range(last.start, max(last.stop, r.stop))
311 else:
312 merged.append(r)
313 return merged
314
315
316def merge_annotations(
317 annotations: list[FunctionAnnotation], merge_gap_max: int | None = None
318) -> list[FunctionAnnotation]:
319 """Merges annotations into non-overlapping segments.
320
321 Args:
322 annotations: annotations to merge.
323 merge_gap_max: optionally merge neighboring ranges that are separated by a gap
324 no larger than this size.
325 Returns:
326 non-overlapping annotations with gaps merged.
327 """
328 grouped: dict[str, list[range]] = defaultdict(list)
329 for a in annotations:
330 # +1 since FunctionAnnotation.end is inlcusive.
331 grouped[a.label].append(range(a.start, a.end + 1))
332
333 merged = []
334 for label, ranges in grouped.items():
335 merged_ranges = merge_ranges(ranges, merge_gap_max=merge_gap_max)
336 for range_ in merged_ranges:
337 annotation = FunctionAnnotation(
338 label=label,
339 start=range_.start,
340 end=range_.stop - 1, # convert range.stop exclusive -> inclusive.
341 )
342 merged.append(annotation)
343 return merged
344
345
346def replace_inf(data):
347 if data is None:
348 return None
349 array = np.asarray(data, dtype=np.float32)
350 array = np.where(np.isinf(array), 1000, array)
351 return array.tolist()
352
353
354def maybe_tensor(x, convert_none_to_nan: bool = False) -> torch.Tensor | None:
355 if x is None:
356 return None
357 if isinstance(x, torch.Tensor):
358 return x
359 if isinstance(x, list) and all(isinstance(t, torch.Tensor) for t in x):
360 return torch.stack(x)
361 if convert_none_to_nan:
362 x = np.asarray(x, dtype=np.float32)
363 x = np.where(x is None, np.nan, x)
364 return torch.tensor(x)
365
366
367def maybe_list(x, convert_nan_to_none: bool = False) -> list | None:
368 if x is None:
369 return None
370 if not convert_nan_to_none:
371 return x.tolist()
372
373 # Handle both torch.tensor and np.ndarray input.
374 if isinstance(x, torch.Tensor):
375 nan_mask = torch.isnan(x).cpu().numpy()
376 np_arr = x.cpu().numpy().astype(object)
377 elif isinstance(x, np.ndarray):
378 nan_mask = np.isnan(x)
379 np_arr = x.astype(object)
380 else:
381 raise TypeError("maybe_list can only work with torch.tensor or np.ndarray.")
382
383 np_arr[nan_mask] = None
384 return np_arr.tolist()
385
386
387def huggingfacehub_login():
388 """Authenticates with the Hugging Face Hub using the HF_TOKEN environment
389 variable, else by prompting the user"""
390 token = os.environ.get("HF_TOKEN")
391 huggingface_hub.login(token=token)
392
393
394def get_chainbreak_boundaries_from_sequence(sequence: Sequence[str]) -> np.ndarray:
395 chain_boundaries = [0]
396 for i, aa in enumerate(sequence):
397 if aa == CHAIN_BREAK_STR:
398 if i == (len(sequence) - 1):
399 raise ValueError(
400 "Encountered chain break token at end of sequence, this is unexpected."
401 )
402 if i == (len(sequence) - 2):
403 warn(
404 "Encountered chain break token at penultimate position, this is unexpected."
405 )
406 chain_boundaries.append(i)
407 chain_boundaries.append(i + 1)
408 chain_boundaries.append(len(sequence))
409 assert len(chain_boundaries) % 2 == 0
410 chain_boundaries = np.array(chain_boundaries).reshape(-1, 2)
411 return chain_boundaries
412
413
414def deserialize_tensors(b: bytes) -> Any:
415 buf = BytesIO(zstd.ZSTD_uncompress(b))
416 d = torch.load(buf, map_location="cpu", weights_only=False)
417 return d
418
419
420def join_lists(
421 lists: Sequence[Sequence[Any]], separator: Sequence[Any] | None = None
422) -> list[Any]:
423 """Joins multiple lists with separator element. Like str.join but for lists.
424
425 Example: [[1, 2], [3], [4]], separator=[0] -> [1, 2, 0, 3, 0, 4]
426
427 Args:
428 lists: Lists of elements to chain
429 separator: separators to intsert between chained output.
430 Returns:
431 Joined lists.
432 """
433 if not lists:
434 return []
435 joined = []
436 joined.extend(lists[0])
437 for l in lists[1:]:
438 if separator:
439 joined.extend(separator)
440 joined.extend(l)
441 return joined
442
443
444def iterate_with_intermediate(
445 lists: Iterable, intermediate
446) -> Generator[Any, None, None]:
447 """
448 Iterate over the iterable, yielding the intermediate value between
449 every element of the intermediate. Useful for joining objects with
450 separator tokens.
451 """
452 it = iter(lists)
453 yield next(it)
454 for l in it:
455 yield intermediate
456 yield l
457
458
459def concat_objects(objs: Sequence[Any], separator: Any | None = None):
460 """
461 Concat objects with each other using a separator token.
462
463 Supports:
464 - Concatable (objects that implement `concat` classmethod)
465 - strings
466 - lists
467 - numpy arrays
468 - torch Tensors
469
470 Example:
471 >>> foo = "abc"
472 >>> bar = "def"
473 >>> concat_objects([foo, bar], "|")
474 "abc|def"
475 """
476 match objs[0]:
477 case Concatable():
478 return objs[0].__class__.concat(objs) # type: ignore
479 case str():
480 assert isinstance(
481 separator, str
482 ), "Trying to join strings but separator is not a string"
483 return separator.join(objs)
484 case list():
485 if separator is not None:
486 return join_lists(objs, [separator])
487 else:
488 return join_lists(objs)
489 case np.ndarray():
490 if separator is not None:
491 return np.concatenate(
492 list(iterate_with_intermediate(objs, np.array([separator])))
493 )
494 else:
495 return np.concatenate(objs)
496 case torch.Tensor():
497 if separator is not None:
498 return torch.cat(
499 list(iterate_with_intermediate(objs, torch.tensor([separator])))
500 )
501 else:
502 return torch.cat(objs) # type: ignore
503 case _:
504 raise TypeError(type(objs[0]))
505