Prepare batch for forward pass. :param batch: A batch dictionary.
(batch: MODEL_BATCH)
| 1544 | |
| 1545 | |
| 1546 | def prepare_batch(batch: MODEL_BATCH): |
| 1547 | """Prepare batch for forward pass. |
| 1548 | |
| 1549 | :param batch: A batch dictionary. |
| 1550 | """ |
| 1551 | if "outputs" not in batch: |
| 1552 | batch["outputs"] = {} |
| 1553 | if "indexer" in batch and "gather_idx_a_cotherid" not in batch["indexer"]: |
| 1554 | cother_mask = batch["features"]["res_atom_mask"].bool().clone() |
| 1555 | cother_mask[:, 1] = False |
| 1556 | atom37_mask = torch.zeros_like(cother_mask, dtype=torch.long) |
| 1557 | atom37_mask += torch.arange(0, atom37_mask.size(0), device=atom37_mask.device).unsqueeze( |
| 1558 | -1 |
| 1559 | ) |
| 1560 | batch["indexer"]["gather_idx_a_cotherid"] = atom37_mask[cother_mask] |
| 1561 | if "features" in batch and "apo_res_alignment_mask" not in batch["features"]: |
| 1562 | batch["features"]["apo_res_alignment_mask"] = torch.ones_like( |
| 1563 | batch["features"]["res_atom_mask"][:, 1], dtype=torch.bool |
| 1564 | ) |
| 1565 | if "num_molid" in batch["metadata"].keys() and batch["metadata"]["num_molid"] > 0: |
| 1566 | batch["misc"]["protein_only"] = False |
| 1567 | else: |
| 1568 | batch["misc"]["protein_only"] = True |
| 1569 | |
| 1570 | |
| 1571 | def centralize_complex_graph(complex_graph: Dict[str, Any]) -> Dict[str, Any]: |