File size: 6,685 Bytes
75b53c6 73d658a 75b53c6 73d658a 75b53c6 73d658a 75b53c6 | 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 | # Copyright 2026 The HuggingFace Team. All rights reserved.
# RFD4-derived preprocessing: Copyright (c) 2026, Institute for Protein Design,
# University of Washington (BSD-3-Clause). See THIRD_PARTY_LICENSES.txt.
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0
"""AtomWorks adapter: annotated CIF -> prepared tensors for `RFD4ProteinaPipeline`, and generated tensors -> CIF."""
import os
import random
import tempfile
from pathlib import Path
import numpy as np
import torch
def prepare_cif(path: str | Path, *, seed: int, device: torch.device, path_length_cutoff: int = 64):
"""Return a single-example feature batch and its atom mapping, without changing the caller's RNG state.
`device` is where the returned tensors live; pass the pipeline's device so no copy is needed at call time.
"""
try:
from atomworks.io.parser import parse
except ImportError as error:
raise ImportError(
"Annotated CIF inputs require the optional AtomWorks package. Install atomworks[ml]; RFD4 is not required."
) from error
# Import before parsing to register the model's condition and provenance annotations.
from ._cif.datapipes_pipeline import build_design_validation_pipeline
from ._cif.datapipes_schema_proteina import SPECS
from ._cif.utils_constants import LOCAL_LATENT_RESIDUES
from ._cif.utils_parse_presets import ANNOTATED_CIF
path = Path(path)
if not path.is_file():
raise FileNotFoundError(path)
numpy_state, python_state = np.random.get_state(), random.getstate()
try:
with torch.random.fork_rng(devices=[]):
np.random.seed(seed % 2**32)
random.seed(seed)
torch.random.default_generator.manual_seed(seed)
parsed = parse(path, config=ANNOTATED_CIF)
data = {"atom_array": parsed["assemblies"]["1"][0], "example_id": path.stem}
# The RFD4 datapipe itself, in its inference mode: there is one pipeline here, not a
# second one maintained alongside it. `build_design_validation_pipeline` is
# `build_design_train_pipeline(is_inference=True)`.
#
# `na_ss_conditioning` turns on nucleic-acid secondary-structure featurization. RFD4
# gates that on the metrics it is about to compute, and its own inference default is
# no tags at all -- but an annotated CIF can carry NA secondary-structure conditioning
# in the file, and without this it is parsed and then silently discarded, leaving the
# model's `na_ss_state` and `na_ss_pair` inputs at their defaults. This is RFD4's own
# switch, set the way RFD4's `ss_conditioned` config sets it, not a flag invented here.
pipeline = build_design_validation_pipeline(
metrics_tags=["na_ss_conditioning"],
residue_names_supporting_local_latents=list(LOCAL_LATENT_RESIDUES),
path_length_search_cutoff=path_length_cutoff,
expandable_segment_rng=np.random.default_rng(seed),
)
data = pipeline(data)
finally:
np.random.set_state(numpy_state)
random.setstate(python_state)
# Match the released schema, including padding of auxiliary fields and absent conditions.
tensors = {
key: torch.as_tensor(value, dtype=SPECS[key][0] if key in SPECS else None)
for key, value in data["feats"].items()
if isinstance(value, (torch.Tensor, np.ndarray))
}
length = max(
tensor.shape[axis]
for key, tensor in tensors.items()
for axis, dimension in enumerate(SPECS[key][2].split() if key in SPECS else ["I"])
if dimension == "I" and axis < tensor.ndim
)
features = {}
for key, tensor in tensors.items():
dtype, default, dimensions = SPECS.get(key, (tensor.dtype, 0, "I"))
axes = dimensions.split()
padding = []
for axis in reversed(range(tensor.ndim)):
padding.extend((0, length - tensor.shape[axis] if axis < len(axes) and axes[axis] == "I" else 0))
if any(padding):
tensor = torch.nn.functional.pad(tensor, padding, value=default)
features[key] = tensor.unsqueeze(0).to(device)
for key, (dtype, default, dimensions) in SPECS.items():
if key not in features:
shape = (1, *(length if dimension == "I" else int(dimension) for dimension in dimensions.split()))
features[key] = torch.full(shape, default, dtype=dtype, device=device)
features["padding_mask"] = torch.ones((1, length), dtype=torch.bool, device=device)
features["coords_nm"] = features["coords"] / 10.0
features["coord_mask"] = features["atom37_mask"]
return features, data["atom_array"]
def write_cif(output, features, atom_array, path: str | Path, row: int = 0) -> str:
"""Write the generated atoms of an `RFD4ProteinaPipelineOutput` to a new annotated CIF.
Conditioning annotations and provenance are preserved and predicted unindexed-motif insertion mappings are added.
`features` and `atom_array` are the values returned by `prepare_cif`. Coordinates are converted from nanometers to
angstroms. For a batched output (the features repeated for several designs), `row` selects the design to write.
An existing `path` is never overwritten.
"""
from ._cif.transforms_unindexing import annotate_unindexed_insertion
from ._cif.utils_structure_export import process_single_sample_to_atom_array, write_generated_structure_cif
path = Path(path)
sample = process_single_sample_to_atom_array(
coords=output.coords_nm[row].detach().cpu() * 10,
residue_type=output.residue_type[row].detach().cpu(),
padding_mask=features["padding_mask"][row].detach().cpu(),
atom_mask=output.atom_mask[row].detach().cpu(),
input_atom_array=atom_array,
insertion_logits=output.insertion_logits[row].detach().cpu(),
)
array = sample["full_array_with_unindexed"]
if array is None:
array = sample["generated_atom_array"]
array = annotate_unindexed_insertion(array, sample["predicted_insertion_mapping"])
# Publish a complete file without overwriting any existing output, including a racing writer.
with tempfile.TemporaryDirectory(dir=path.parent, prefix=".rfd4-cif-") as temporary:
staged = Path(temporary) / path.name
write_generated_structure_cif(array, staged)
os.link(staged, path)
return str(path)
|