deepfold.utils package

Subpackages

Submodules

deepfold.utils.checkpoint_utils module

deepfold.utils.checkpoint_utils.map_init_state_dicts(alphafold_state_dict_keys: List[str], init_alphafold_state_dict: dict, init_optimizer_state_dict: dict) → Tuple[dict, dict]
deepfold.utils.checkpoint_utils.resume_checkpoint(alphafold: AlphaFold, optimizer: Optimizer | None, swa_alphafold: AlphaFoldSWA | None, checkpoint_dirpath: Path, device: device) → int
deepfold.utils.checkpoint_utils.resume_from_latest_checkpoint(alphafold: AlphaFold, optimizer: Optimizer | None, swa_alphafold: AlphaFoldSWA | None, checkpoints_dirpath: Path, device: device) → int
deepfold.utils.checkpoint_utils.save_checkpoint(alphafold: AlphaFold, optimizer: Optimizer, swa_alphafold: AlphaFoldSWA, iteration: int, checkpoint_dirpath: Path) → None
deepfold.utils.checkpoint_utils.save_last_checkpoint(alphafold: AlphaFold, optimizer: Optimizer, swa_alphafold: AlphaFoldSWA, iteration: int, checkpoints_dirpath: Path, keep_last_checkpoints: int) → None
deepfold.utils.checkpoint_utils.save_val_checkpoint(alphafold: AlphaFold | AlphaFoldSWA, iteration: int, checkpoints_dirpath: Path, keep_val_checkpoints: bool, val_avg_lddt_ca: float) → None

deepfold.utils.chunk_utils module

deepfold.utils.chunk_utils.chunk_layer(layer: Callable, inputs: Dict[str, Any], chunk_size: int, num_batch_dims: int, low_mem: bool = False, _out: Any | None = None, _add_into_out: bool = False) → Any

Implements the “chunking” procedure described in section 1.11.8.

Layer outputs and inputs are assumed to be simple “pytrees,” consisting only of (arbitrarily nested) lists, tuples, and dicts with torch.Tensor leaves.

Parameters:

  • layer – The layer to be applied chunk-wise

  • inputs – A (non-nested) dictionary of keyworded inputs. All leaves must be tensors and must share the same batch dimensions.

  • chunk_size – The number of sub-batches per chunk. If multiple batch dimensions are specified, a “sub-batch” is defined as a single indexing of all batch dimensions simultaneously (s.t. the number of sub-batches is the product of the batch dimensions).

  • no_batch_dims – How many of the initial dimensions of each input tensor can be considered batch dimensions.

  • low_mem – Avoids flattening potentially large input tensors. Unnecessary in most cases, and is ever so slightly slower than the default setting.

Returns:

The reassembled output of the layer on the inputs.

deepfold.utils.crop_utils module

deepfold.utils.crop_utils.unpad_to_schema_shape_(outputs: Dict[str, ndarray], output_schema_shapes: Dict[str, tuple], num_residues: int, num_clustered_msa_seq: int) → Dict[str, ndarray]

deepfold.utils.datetime_utils module

deepfold.utils.datetime_utils.datetime_from_string(datetime_string: str, datetime_format: str = '%Y-%m-%d %H:%M:%S') → datetime

Converts string to datetime object.

deepfold.utils.datetime_utils.datetime_to_string(datetime_object: datetime, string_format: str = '%Y-%m-%d %H:%M:%S') → str

Converts datetime object to string.

deepfold.utils.datetime_utils.get_timestamp_string() → str

Returns timestamp in YYYYmmdd_HHMMSS_ffffff format.

deepfold.utils.dist_utils module

deepfold.utils.dist_utils.get_pad_size(tensor: Tensor, dim: int, num_chunks: int) → int
deepfold.utils.dist_utils.pad_tensor(tensor: Tensor, dim: int, pad_size: int) → Tensor

deepfold.utils.feats_utils module

deepfold.utils.feats_utils.crop_a3m_string(a3m_string: str, start: int, end: int) → str
deepfold.utils.feats_utils.crop_features(feats: dict, start: int, end: int)
Parameters:

  • [start (start - 1, end)

  • range (end] ->)

deepfold.utils.file_utils module

deepfold.utils.file_utils.dump_pickle(obj: Any, path: PathLike, level: int = 5) → None
deepfold.utils.file_utils.find_paths(paths: Sequence[PathLike]) → List[Path]
deepfold.utils.file_utils.get_file_content_and_extension(file_path: PathLike) → Tuple[str, str]

Reads the content of a text file or compressed text file and returns its content along with the real extension (excluding compression extensions like .gz).

Parameters:

file_path (str) – Path to the text file or compressed text file.

Returns:

A tuple containing the file content (str) and the real extension (str).

Return type:

tuple

deepfold.utils.file_utils.load_pickle(path: PathLike) → Any
deepfold.utils.file_utils.read_text(path: PathLike) → str
deepfold.utils.file_utils.restore_wrapper(func)

deepfold.utils.hash_utils module

deepfold.utils.hash_utils.hash_string_into_number(s: str) → int

Hashes string into uint64-like integer number.

deepfold.utils.import_utils module

class deepfold.utils.import_utils.Param(param: Union[torch.Tensor, List[torch.Tensor]], param_type: deepfold.utils.import_utils.ParamType = <ParamType.Other: functools.partial(<function ParamType.<lambda> at 0x7f1642212560>)>, stacked: bool = False, swap: bool = False)

Bases: object

param: Tensor | List[Tensor]
param_type: ParamType = functools.partial(<function ParamType.<lambda>>)
stacked: bool = False
swap: bool = False
class deepfold.utils.import_utils.ParamType(value)

Bases: Enum

An enumeration.

LinearBiasMHA = functools.partial(<function ParamType.<lambda>>)
LinearMHAOutputWeight = functools.partial(<function ParamType.<lambda>>)
LinearWeight = functools.partial(<function ParamType.<lambda>>)
LinearWeightMHA = functools.partial(<function ParamType.<lambda>>)
LinearWeightOPM = functools.partial(<function ParamType.<lambda>>)
Other = functools.partial(<function ParamType.<lambda>>)
deepfold.utils.import_utils.assign(translation_dict, orig_weights)
deepfold.utils.import_utils.import_jax_weights_(model: AlphaFold, npz_path: str, is_multimer: bool = False, enable_ptm: bool = False, enable_templates: bool = False, fuse_projection_weights: bool = False) → None

Import AlphaFold JAX parameters.

Parameters:

  • model – AlphaFold model.

  • npz_path – Path to an NPZ file.

  • is_multimer – Whether multimer model or not.

  • enable_ptm – Enable predicted aligned error related modules.

  • enable_templates – Enable template related modules.

  • fuse_projection_weights – Whether triangular multiplicative layers are fused or not.

deepfold.utils.import_utils.reshape_weight(x: ndarray) → ndarray
deepfold.utils.import_utils.stacked(param_dict_list, out=None)
Parameters:

param_dict_list – A list of (nested) Param dicts to stack. The structure of each dict must be the identical (down to the ParamTypes of “parallel” Params). There must be at least one dict in the list.

deepfold.utils.iter_utils module

deepfold.utils.iter_utils.all_equal(values: list) → bool

Checks if all values in list are equal to each other.

deepfold.utils.iter_utils.flatten_dict(d: MutableMapping, parent_key: str = '', sep: str = '.')
deepfold.utils.iter_utils.list_zip(*arglists) → list

Transforms given columns into list of rows.

deepfold.utils.iter_utils.map_dict_values(fn: Callable, d: dict) → dict

Maps dictionary values using given function.

deepfold.utils.iter_utils.slice_generator(start: int, end: int, size: int) → Iterator[Tuple[int, int]]

Returns slice indices iterator from start to end.

deepfold.utils.log_utils module

class deepfold.utils.log_utils.TqdmHandler

Bases: StreamHandler

emit(record: LogRecord) → None

Emit a record.

If a formatter is specified, it is used to format the record. The record is then written to the stream with a trailing newline. If exception information is present, it is formatted using traceback.print_exception and appended to the stream. If the stream has an ‘encoding’ attribute, it is used to determine how to do the output to the stream.

deepfold.utils.log_utils.read_logs(filepath: Path, drop_overridden_iterations: bool = True) → DataFrame
deepfold.utils.log_utils.save_logs(logs: List[dict], outpath: Path, append: bool) → None
deepfold.utils.log_utils.setup_logging(filename: str | PathLike | Path, mode: str = 'w') → None

deepfold.utils.precision module

deepfold.utils.precision.is_fp16_enabled() → bool

deepfold.utils.random module

deepfold.utils.random.get_seed_from_string(s: str) → int

Hashes input string and returns uint64-like integer seed value.

deepfold.utils.random.get_seed_randomly() → int

Returns truly pseduorandom uint64-like integer seed value.

deepfold.utils.random.numpy_seed(seed: int | None = None, *additional_seeds, key: str = '')

Context manager which seeds the Numpy PRNG with the seed and restores the state.

deepfold.utils.random.str_hash(string: str)

deepfold.utils.rigid_utils module

class deepfold.utils.rigid_utils.Rigid(rots: Rotation | None, trans: Tensor | None)

Bases: object

A class representing a rigid transformation. Little more than a wrapper around two objects: a Rotation object and a [*, 3] translation Designed to behave approximately like a single torch tensor with the shape of the shared batch dimensions of its component parts.

apply(pts: Tensor) → Tensor

Applies the transformation to a coordinate tensor.

Parameters:

pts – A [*, 3] coordinate tensor.

Returns:

The transformed points.

apply_rot_fn(fn: Callable[[Rotation], Rotation]) → Rigid

Applies a Rotation -> Rotation function to the stored rotation object.

Parameters:

fn – A function of type Rotation -> Rotation

Returns:

A transformation object with a transformed rotation.

apply_trans_fn(fn: Callable[[Tensor], Tensor]) → Rigid

Applies a Tensor -> Tensor function to the stored translation.

Parameters:

fn – A function of type Tensor -> Tensor to be applied to the translation

Returns:

A transformation object with a transformed translation.

static cat(ts: Sequence[Rigid], dim: int) → Rigid

Concatenates transformations along a new dimension.

Parameters:

  • ts – A list of T objects

  • dim – The dimension along which the transformations should be concatenated

Returns:

A concatenated transformation object

compose(r: Rigid) → Rigid

Composes the current rigid object with another.

Parameters:

r – Another Rigid object

Returns:

The composition of the two transformations

compose_q_update_vec(q_update_vec: Tensor) → Rigid

Composes the transformation with a quaternion update vector of shape [*, 6], where the final 6 columns represent the x, y, and z values of a quaternion of form (1, x, y, z) followed by a 3D translation.

Parameters:

q_vec – The quaternion update vector.

Returns:

The composed transformation.

cuda() → Rigid

Moves the transformation object to GPU memory

Returns:

A version of the transformation on GPU

property device: device

Returns the device on which the Rigid’s tensors are located.

Returns:

The device on which the Rigid’s tensors are located

property dtype: dtype

Returns the dtype of the Rigid tensors.

Returns:

The dtype of the Rigid tensors

static from_3_points(p_neg_x_axis: Tensor, origin: Tensor, p_xy_plane: Tensor, eps: float = 1e-08) → Rigid

Implements algorithm 21. Constructs transformations from sets of 3 points using the Gram-Schmidt algorithm.

Parameters:

  • p_neg_x_axis – [*, 3] coordinates

  • origin – [*, 3] coordinates used as frame origins

  • p_xy_plane – [*, 3] coordinates

  • eps – Small epsilon value

Returns:

A transformation object of shape [*]

static from_tensor_4x4(t: Tensor) → Rigid

Constructs a transformation from a homogenous transformation tensor.

Parameters:

t – [*, 4, 4] homogenous transformation tensor

Returns:

T object with shape [*]

static from_tensor_7(t: Tensor, normalize_quats: bool = False) → Rigid
get_rots() → Rotation

Getter for the rotation.

Returns:

The rotation object

get_trans() → Tensor

Getter for the translation.

Returns:

The stored translation

static identity(shape: Tuple[int], dtype: dtype | None = None, device: device | None = None, requires_grad: bool = True, fmt: str = 'quat') → Rigid

Constructs an identity transformation.

Parameters:

  • shape – The desired shape

  • dtype – The dtype of both internal tensors

  • device – The device of both internal tensors

  • requires_grad – Whether grad should be enabled for the internal tensors

Returns:

The identity transformation

invert() → Rigid

Inverts the transformation.

Returns:

The inverse transformation.

invert_apply(pts: Tensor) → Tensor

Applies the inverse of the transformation to a coordinate tensor.

Parameters:

pts – A [*, 3] coordinate tensor

Returns:

The transformed points.

static make_transform_from_reference(n_xyz, ca_xyz, c_xyz, eps=1e-20)

Returns a transformation object from reference coordinates.

Note that this method does not take care of symmetries. If you provide the atom positions in the non-standard way, the N atom will end up not at [-0.527250, 1.359329, 0.0] but instead at [-0.527250, -1.359329, 0.0]. You need to take care of such cases in your code.

Parameters:

  • n_xyz – A [*, 3] tensor of nitrogen xyz coordinates.

  • ca_xyz – A [*, 3] tensor of carbon alpha xyz coordinates.

  • c_xyz – A [*, 3] tensor of carbon xyz coordinates.

Returns:

A transformation object. After applying the translation and rotation to the reference backbone, the coordinates will approximately equal to the input coordinates.

map_tensor_fn(fn: Callable[[Tensor], Tensor]) → Rigid

Apply a Tensor -> Tensor function to underlying translation and rotation tensors, mapping over the translation/rotation dimensions respectively.

Parameters:

fn – A Tensor -> Tensor function to be mapped over the Rigid

Returns:

The transformed Rigid object

scale_translation(trans_scale_factor: float) → Rigid

Scales the translation by a constant factor.

Parameters:

trans_scale_factor – The constant factor

Returns:

A transformation object with a scaled translation.

property shape: Size

Returns the shape of the shared dimensions of the rotation and the translation.

Returns:

The shape of the transformation

stop_rot_gradient() → Rigid

Detaches the underlying rotation object

Returns:

A transformation object with detached rotations

to_tensor_4x4() → Tensor

Converts a transformation to a homogenous transformation tensor.

Returns:

A [*, 4, 4] homogenous transformation tensor

to_tensor_7() → Tensor

Converts a transformation to a tensor with 7 final columns, four for the quaternion followed by three for the translation.

Returns:

A [*, 7] tensor representation of the transformation

unsqueeze(dim: int) → Rigid

Analogous to torch.unsqueeze. The dimension is relative to the shared dimensions of the rotation/translation.

Parameters:

dim – A positive or negative dimension index.

Returns:

The unsqueezed transformation.

class deepfold.utils.rigid_utils.Rotation(rot_mats: Tensor | None = None, quats: Tensor | None = None, normalize_quats: bool = True)

Bases: object

A 3D rotation. Depending on how the object is initialized, the rotation is represented by either a rotation matrix or a quaternion, though both formats are made available by helper functions. To simplify gradient computation, the underlying format of the rotation cannot be changed in-place. Like Rigid, the class is designed to mimic the behavior of a torch Tensor, almost as if each Rotation object were a tensor of rotations, in one format or another.

apply(pts: Tensor) → Tensor

Apply the current Rotation as a rotation matrix to a set of 3D coordinates.

Parameters:

pts – A [*, 3] set of points

Returns:

[*, 3] rotated points

static cat(rs: Sequence[Rotation], dim: int) → Rotation

Concatenates rotations along one of the batch dimensions. Analogous to torch.cat().

Note that the output of this operation is always a rotation matrix, regardless of the format of input rotations.

Parameters:

  • rs – A list of rotation objects

  • dim – The dimension along which the rotations should be concatenated

Returns:

A concatenated Rotation object in rotation matrix format

compose_q(r: Rotation, normalize_quats: bool = True) → Rotation

Compose the quaternions of the current Rotation object with those of another.

Depending on whether either Rotation was initialized with quaternions, this function may call torch.linalg.eigh.

Parameters:

r – An update rotation object

Returns:

An updated rotation object

compose_q_update_vec(q_update_vec: Tensor, normalize_quats: bool = True) → Rotation

Returns a new quaternion Rotation after updating the current object’s underlying rotation with a quaternion update, formatted as a [*, 3] tensor whose final three columns represent x, y, z such that (1, x, y, z) is the desired (not necessarily unit) quaternion update.

Parameters:

  • q_update_vec – A [*, 3] quaternion update tensor

  • normalize_quats – Whether to normalize the output quaternion

Returns:

An updated Rotation

compose_r(r: Rotation) → Rotation

Compose the rotation matrices of the current Rotation object with those of another.

Parameters:

r – An update rotation object

Returns:

An updated rotation object

cuda() → Rotation

Analogous to the cuda() method of torch Tensors

Returns:

A copy of the Rotation in CUDA memory

detach() → Rotation

Returns a copy of the Rotation whose underlying Tensor has been detached from its torch graph.

Returns:

A copy of the Rotation whose underlying Tensor has been detached from its torch graph

property device: device

The device of the underlying rotation

Returns:

The device of the underlying rotation

property dtype: dtype

Returns the dtype of the underlying rotation.

Returns:

The dtype of the underlying rotation

get_cur_rot() → Tensor

Return the underlying rotation in its current form

Returns:

The stored rotation

get_quats() → Tensor

Returns the underlying rotation as a quaternion tensor.

Depending on whether the Rotation was initialized with a quaternion, this function may call torch.linalg.eigh.

Returns:

The rotation as a quaternion tensor.

get_rot_mats() → Tensor

Returns the underlying rotation as a rotation matrix tensor.

Returns:

The rotation as a rotation matrix tensor

static identity(shape, dtype: dtype | None = None, device: device | None = None, requires_grad: bool = True, fmt: str = 'quat') → Rotation

Returns an identity Rotation.

Parameters:

  • shape – The “shape” of the resulting Rotation object. See documentation for the shape property

  • dtype – The torch dtype for the rotation

  • device – The torch device for the new rotation

  • requires_grad – Whether the underlying tensors in the new rotation object should require gradient computation

  • fmt – One of “quat” or “rot_mat”. Determines the underlying format of the new object’s rotation

Returns:

A new identity rotation

invert() → Rotation

Returns the inverse of the current Rotation.

Returns:

The inverse of the current Rotation

invert_apply(pts: Tensor) → Tensor

The inverse of the apply() method.

Parameters:

pts – A [*, 3] set of points

Returns:

[*, 3] inverse-rotated points

map_tensor_fn(fn: Callable[[Tensor], Tensor]) → Rotation

Apply a Tensor -> Tensor function to underlying rotation tensors, mapping over the rotation dimension(s). Can be used e.g. to sum out a one-hot batch dimension.

Parameters:

fn – A Tensor -> Tensor function to be mapped over the Rotation

Returns:

The transformed Rotation object

property requires_grad: bool

Returns the requires_grad property of the underlying rotation

Returns:

The requires_grad property of the underlying tensor

property shape: Size

Returns the virtual shape of the rotation object. This shape is defined as the batch dimensions of the underlying rotation matrix or quaternion. If the Rotation was initialized with a [10, 3, 3] rotation matrix tensor, for example, the resulting shape would be [10].

Returns:

The virtual shape of the rotation object

to(device: device | None, dtype: dtype | None) → Rotation

Analogous to the to() method of torch Tensors

Parameters:

  • device – A torch device

  • dtype – A torch dtype

Returns:

A copy of the Rotation using the new device and dtype

unsqueeze(dim: int) → Rotation

Analogous to torch.unsqueeze. The dimension is relative to the shape of the Rotation object.

Parameters:

dim – A positive or negative dimension index.

Returns:

The unsqueezed Rotation.

deepfold.utils.rigid_utils.identity_quats(batch_dims: Tuple[int], dtype: dtype | None = None, device: device | None = None, requires_grad: bool = True) → Tensor
deepfold.utils.rigid_utils.identity_rot_mats(batch_dims: Tuple[int], dtype: dtype | None = None, device: device | None = None, requires_grad: bool = True) → Tensor
deepfold.utils.rigid_utils.identity_trans(batch_dims: Tuple[int], dtype: dtype | None = None, device: device | None = None, requires_grad: bool = True) → Tensor
deepfold.utils.rigid_utils.invert_quat(quat: Tensor)
deepfold.utils.rigid_utils.invert_rot_mat(rot_mat: Tensor)
deepfold.utils.rigid_utils.quat_multiply(quat1, quat2)

Multiply a quaternion by another quaternion.

deepfold.utils.rigid_utils.quat_multiply_by_vec(quat, vec)

Multiply a quaternion by a pure-vector quaternion.

deepfold.utils.rigid_utils.quat_to_rot(quat: Tensor) → Tensor

Converts a quaternion to a rotation matrix.

Parameters:

quat – [*, 4] quaternions

Returns:

[*, 3, 3] rotation matrices

deepfold.utils.rigid_utils.rot_matmul(a: Tensor, b: Tensor) → Tensor

Performs matrix multiplication of two rotation matrix tensors. Written out by hand to avoid AMP downcasting.

Parameters:

  • a – [*, 3, 3] left multiplicand

  • b – [*, 3, 3] right multiplicand

Returns:

The product ab

deepfold.utils.rigid_utils.rot_to_quat(rot: Tensor) → Tensor
deepfold.utils.rigid_utils.rot_vec_mul(r: Tensor, t: Tensor) → Tensor

Applies a rotation to a vector. Written out by hand to avoid transfer to avoid AMP downcasting.

Parameters:

  • r – [*, 3, 3] rotation matrices

  • t – [*, 3] coordinate tensors

Returns:

[*, 3] rotated coordinates

deepfold.utils.script_utils module

deepfold.utils.script_utils.get_basename(path: PathLike)
deepfold.utils.script_utils.make_output_directory(output_dir: PathLike, model_name: str)
deepfold.utils.script_utils.parse_fasta(fasta_str: str)

deepfold.utils.swa_utils module

class deepfold.utils.swa_utils.AlphaFoldSWA(alphafold: AlphaFold, enabled: bool, decay_rate: float)

Bases: PatchedModule

Weight averaging wrapper.

forward(batch)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

update(model: AlphaFold) → None
class deepfold.utils.swa_utils.swa_avg_fn(decay_rate: float)

Bases: object

Averaging function for EMA with configurable decay rate.

Suppl. ‘1.11.7 Evaluator setup’.

deepfold.utils.tensor_utils module

Utilities related to tensor operations.

deepfold.utils.tensor_utils.add(m1: Tensor, m2: Tensor, inplace: bool) → Tensor
deepfold.utils.tensor_utils.array_tree_map(fn, tree, *, leaf_type=<class 'numpy.ndarray'>)
deepfold.utils.tensor_utils.batched_gather(data: Tensor, inds: Tensor, dim: int = 0, num_batch_dims: int = 0) → Tensor
deepfold.utils.tensor_utils.collate(samples: List[dict]) → dict

Converts list of samples into a batch dict.

deepfold.utils.tensor_utils.dict_map(fn: Callable[[T], Any], dic: Dict[Any, dict | list | tuple | T], leaf_type: Type[T]) → Dict[Any, dict | list | tuple | Any]
deepfold.utils.tensor_utils.dict_multimap(fn: Callable, dicts: List[Dict[Any, Any]]) → Dict[Any, Any]
deepfold.utils.tensor_utils.flatten_final_dims(t: Tensor, num_dims: int) → Tensor
deepfold.utils.tensor_utils.masked_mean(mask: Tensor, value: Tensor, dim: int | Tuple[int, ...], eps: float = 0.0001, keepdim: bool = False) → Tensor
deepfold.utils.tensor_utils.one_hot(x: Tensor, v_bins: Tensor) → Tensor
deepfold.utils.tensor_utils.permute_final_dims(tensor: Tensor, inds: List[int]) → Tensor
deepfold.utils.tensor_utils.pts_to_distogram(pts: Tensor, min_bin: float = 2.2325, max_bin: float = 21.6875, num_bins: int = 64) → Tensor
deepfold.utils.tensor_utils.tensor_tree_map(fn, tree, *, leaf_type=<class 'torch.Tensor'>)
deepfold.utils.tensor_utils.tree_map(fn: Callable[[T], Any], tree: T, leaf_type: Type[T]) → Any
deepfold.utils.tensor_utils.tree_map(fn: Callable[[T], Any], tree: dict, leaf_type: Type[T]) → dict
deepfold.utils.tensor_utils.tree_map(fn: Callable[[T], Any], tree: list, leaf_type: Type[T]) → list
deepfold.utils.tensor_utils.tree_map(fn: Callable[[T], Any], tree: tuple, leaf_type: Type[T]) → tuple

Module contents