File size: 3,928 Bytes
6cc35b0 | 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 | """Dataclass support for aligned residue-level fields."""
from __future__ import annotations
from abc import ABC, abstractmethod
from collections.abc import Iterable
from dataclasses import Field, dataclass, fields, replace
from typing import Any, Self
import numpy as np
from .esmfold2_misc import concat_objects, slice_any_object
Index = int | list[int] | slice | np.ndarray
def _is_sequential(field: Field[Any]) -> bool:
return bool(field.metadata.get("sequence", False))
def _sequence_axis(field: Field[Any]) -> int:
axis = int(field.metadata.get("sequence_dim", 0))
if axis not in (0, 1):
raise NotImplementedError("SequentialDataclass supports sequence_dim values zero and one.")
return axis
def _slice_value(value: Any, index: Index, axis: int) -> Any:
if axis == 0:
return slice_any_object(value, index)
sliced = [slice_any_object(track, index) for track in value]
return value.__class__(sliced)
def _iter_sequence_lengths(value: Any, axis: int) -> Iterable[int]:
if axis == 0:
yield len(value)
else:
yield from (len(track) for track in value)
@dataclass(frozen=True)
class SequentialDataclass(ABC):
"""Keep dataclass fields aligned along a shared residue dimension.
A subclass marks aligned fields with ``metadata={"sequence": True}``.
``sequence_dim`` may be zero for a direct sequence or one for a collection
of aligned tracks. ``join_token`` is passed to the package concatenation
helper when instances are joined.
"""
def __post_init__(self) -> None:
expected = len(self)
for field in fields(self):
if not _is_sequential(field) or field.name == "complex":
continue
value = getattr(self, field.name)
if value is None:
continue
for actual in _iter_sequence_lengths(value, _sequence_axis(field)):
if actual != expected:
raise ValueError(
f"Mismatch in sequence length for field: {field.name}. "
f"Expected {expected}, received {actual}"
)
@abstractmethod
def __len__(self) -> int:
"""Return the shared sequence length."""
raise NotImplementedError
def __getitem__(self, index: Index) -> Self:
"""Apply one sequence index to every aligned field."""
normalized_index: Index = [index] if isinstance(index, int) else index
updates: dict[str, Any] = {}
for field in fields(self):
if not _is_sequential(field):
continue
value = getattr(self, field.name)
if value is not None:
updates[field.name] = _slice_value(value, normalized_index, _sequence_axis(field))
return replace(self, **updates)
@classmethod
def concat(cls, items: list[Self], **overrides: Any) -> Self:
"""Join aligned fields and retain non-sequential values from the first item."""
if not items:
raise ValueError("SequentialDataclass.concat requires at least one item.")
updates: dict[str, Any] = {}
for field in fields(cls):
if not _is_sequential(field):
continue
first_value = getattr(items[0], field.name)
if first_value is None:
continue
values = [getattr(item, field.name) for item in items]
join_token = field.metadata.get("join_token")
if _sequence_axis(field) == 0:
updates[field.name] = concat_objects(values, join_token)
else:
tracks = [concat_objects(track, join_token) for track in zip(*values, strict=True)]
updates[field.name] = first_value.__class__(tracks)
updates.update(overrides)
return replace(items[0], **updates)
__all__ = ["SequentialDataclass"]
|