File size: 3,945 Bytes
35cdf53 | 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 |
"""Post-processing utilities for AlphaFold inference results."""
import dataclasses
import datetime
import os
from flax_model.alphafold3 import version
from flax_model.alphafold3.model import confidence_types
from flax_model.alphafold3.model import mmcif_metadata
from flax_model.alphafold3.model import model
import numpy as np
@dataclasses.dataclass(frozen=True, slots=True, kw_only=True)
class ProcessedInferenceResult:
"""Stores attributes of a processed inference result.
Attributes:
cif: CIF file containing an inference result.
mean_confidence_1d: Mean 1D confidence calculated from confidence_1d.
ranking_score: Ranking score extracted from CIF metadata.
structure_confidence_summary_json: Content of JSON file with structure
confidences summary calculated from CIF file.
structure_full_data_json: Content of JSON file with structure full
confidences calculated from CIF file.
model_id: Identifier of the model that produced the inference result.
"""
cif: bytes
mean_confidence_1d: float
ranking_score: float
structure_confidence_summary_json: bytes
structure_full_data_json: bytes
model_id: bytes
def post_process_inference_result(
inference_result: model.InferenceResult,
) -> ProcessedInferenceResult:
"""Returns cif, confidence_1d_json, confidence_2d_json, mean_confidence_1d, and ranking confidence."""
# Add mmCIF metadata fields.
timestamp = datetime.datetime.now().isoformat(sep=' ', timespec='seconds')
cif_with_metadata = mmcif_metadata.add_metadata_to_mmcif(
old_cif=inference_result.predicted_structure.to_mmcif_dict(),
version=f'{version.__version__} @ {timestamp}',
model_id=inference_result.model_id,
)
cif = mmcif_metadata.add_legal_comment(cif_with_metadata.to_string())
cif = cif.encode('utf-8')
confidence_1d = confidence_types.AtomConfidence.from_inference_result(
inference_result
)
mean_confidence_1d = np.mean(confidence_1d.confidence)
structure_confidence_summary_json = (
confidence_types.StructureConfidenceSummary.from_inference_result(
inference_result
)
.to_json()
.encode('utf-8')
)
structure_full_data_json = (
confidence_types.StructureConfidenceFull.from_inference_result(
inference_result
)
.to_json()
.encode('utf-8')
)
return ProcessedInferenceResult(
cif=cif,
mean_confidence_1d=mean_confidence_1d,
ranking_score=float(inference_result.metadata['ranking_score']),
structure_confidence_summary_json=structure_confidence_summary_json,
structure_full_data_json=structure_full_data_json,
model_id=inference_result.model_id,
)
def write_output(
inference_result: model.InferenceResult,
output_dir: os.PathLike[str] | str,
terms_of_use: str | None = None,
name: str | None = None,
) -> None:
"""Writes processed inference result to a directory."""
processed_result = post_process_inference_result(inference_result)
prefix = f'{name}_' if name is not None else ''
with open(os.path.join(output_dir, f'{prefix}model.cif'), 'wb') as f:
f.write(processed_result.cif)
with open(
os.path.join(output_dir, f'{prefix}summary_confidences.json'), 'wb'
) as f:
f.write(processed_result.structure_confidence_summary_json)
with open(os.path.join(output_dir, f'{prefix}confidences.json'), 'wb') as f:
f.write(processed_result.structure_full_data_json)
if terms_of_use is not None:
with open(os.path.join(output_dir, 'TERMS_OF_USE.md'), 'wt') as f:
f.write(terms_of_use)
def write_embeddings(
embeddings: dict[str, np.ndarray],
output_dir: os.PathLike[str] | str,
name: str | None = None,
) -> None:
"""Writes embeddings to a directory."""
prefix = f'{name}_' if name is not None else ''
with open(os.path.join(output_dir, f'{prefix}embeddings.npz'), 'wb') as f:
np.savez_compressed(f, **embeddings)
|