model.py

class Model(hk.Module):
    def __init__(self, config: Config, name: str = 'diffuser')

class Config(base_config.BaseConfig):
    evoformer: evoformer_network.Evoformer.Config = base_config.autocreate()
    global_config: model_config.GlobalConfig = base_config.autocreate()
    heads: 'Model.HeadsConfig' = base_config.autocreate()
    num_recycles: int = 10
    return_embeddings: bool = False
    return_distogram: bool = False

class HeadsConfig(base_config.BaseConfig):
    diffusion: diffusion_head.DiffusionHead.Config = base_config.autocreate()
    confidence: confidence_head.ConfidenceHead.Config = base_config.autocreate()
    distogram: distogram_head.DistogramHead.Config = base_config.autocreate()

def __call__(
    self, 
    batch: features.BatchDict, 
    key: jax.Array | None = None
) -> ModelResult

ModelResult: TypeAlias = Mapping[str, Any]

@classmethod
def get_inference_result(
    cls,
    batch: features.BatchDict,
    result: ModelResult,
    target_name: str = '',
) -> Iterable[InferenceResult]

@dataclasses.dataclass(frozen=True, kw_only=True)
class InferenceResult:
    predicted_structure: structure.Structure
    numerical_data: Mapping[str, float | int | np.ndarray] = dataclasses.field(
        default_factory=dict
    )
    metadata: Mapping[str, float | int | np.ndarray] = dataclasses.field(
        default_factory=dict
    )
    debug_outputs: Mapping[str, Any] = dataclasses.field(default_factory=dict)
    model_id: bytes = b''

def get_predicted_structure(
    result: ModelResult, 
    batch: feat_batch.Batch
) -> structure.Structure

def create_target_feat_embedding(
    batch: feat_batch.Batch,
    config: evoformer_network.Evoformer.Config,
    global_config: model_config.GlobalConfig,
) -> jnp.ndarray

import jax
from alphafold3.model import model, features
from alphafold3.model import params
import haiku as hk

# Load model configuration
config = model.Model.Config()
config.num_recycles = 10
config.heads.diffusion.eval.num_samples = 5

# Create model
def forward_fn(batch):
    return model.Model(config)(batch)

model_fn = hk.transform(forward_fn)

# Load parameters
model_params = params.get_model_haiku_params(model_dir='/path/to/models')

# Run inference
rng_key = jax.random.PRNGKey(42)
result = model_fn.apply(model_params, rng_key, batch)

# Extract inference results
inference_results = list(
    model.Model.get_inference_result(
        batch=batch,
        result=result,
        target_name='my_protein'
    )
)

# Access first sample
first_result = inference_results[0]
print(f"Ranking score: {first_result.metadata['ranking_score']}")
print(f"pTM: {first_result.metadata['predicted_tm_score']}")
print(f"Has clash: {first_result.metadata['has_clash']}")

# Save structure
first_result.predicted_structure.save_to_mmcif('output.cif')

# Get confidence matrices
pae = first_result.numerical_data['full_pae']
pde = first_result.numerical_data['full_pde']
contacts = first_result.numerical_data['contact_probs']

# Chain-level metrics
chain_pair_iptm = first_result.metadata['chain_pair_iptm']
print(f"Chain pair interface scores:\n{chain_pair_iptm}")

# Overall quality
ranking_score = first_result.metadata['ranking_score']
fraction_disordered = first_result.metadata['fraction_disordered']
print(f"Quality: {ranking_score:.3f}, Disorder: {fraction_disordered:.3f}")

# Compare all samples by ranking score
samples = sorted(
    inference_results,
    key=lambda x: x.metadata['ranking_score'],
    reverse=True
)

best_sample = samples[0]
print(f"Best sample ranking score: {best_sample.metadata['ranking_score']}")

# Ensemble metrics
ptm_scores = [s.metadata['predicted_tm_score'] for s in samples]
print(f"pTM mean: {np.mean(ptm_scores):.3f}, std: {np.std(ptm_scores):.3f}")

# Configure model to return embeddings
config.return_embeddings = True
config.return_distogram = True

# Run inference
result = model_fn.apply(model_params, rng_key, batch)

# Access embeddings
single_emb = result['single_embeddings']  # [num_tokens, 384]
pair_emb = result['pair_embeddings']      # [num_tokens, num_tokens, 128]
distogram = result['distogram']['distogram']  # [num_tokens, num_tokens, 64]

print(f"Single embedding shape: {single_emb.shape}")
print(f"Pair embedding shape: {pair_emb.shape}")

# Simplified forward pass
def __call__(batch):
    # 1. Create target features
    target_feat = create_target_feat_embedding(batch, config, global_config)
    
    # 2. Run recycling through Evoformer trunk
    for i in range(num_recycles + 1):
        embeddings = evoformer(batch, prev_embeddings, target_feat)
    
    # 3. Sample diffusion (generates atom coordinates)
    samples = diffusion_head(batch, embeddings)
    
    # 4. Compute confidence metrics
    confidence_output = confidence_head(samples, embeddings)
    
    # 5. Compute distogram
    distogram = distogram_head(batch, embeddings)
    
    return {
        'diffusion_samples': samples,
        'distogram': distogram,
        **confidence_output
    }