File size: 1,590 Bytes
6aab6b3 | 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 | import numpy as np
import torch
from torch.nn.functional import one_hot
from boltzgen.data import const
def load_dummy_templates(tdim: int, num_tokens: int) -> list[dict]:
"""Load dummy templates"""
# Allocate features
res_type = np.zeros((tdim, num_tokens), dtype=np.int64)
frame_rot = np.zeros((tdim, num_tokens, 3, 3), dtype=np.float32)
frame_t = np.zeros((tdim, num_tokens, 3), dtype=np.float32)
cb_coords = np.zeros((tdim, num_tokens, 3), dtype=np.float32)
ca_coords = np.zeros((tdim, num_tokens, 3), dtype=np.float32)
frame_mask = np.zeros((tdim, num_tokens), dtype=np.float32)
cb_mask = np.zeros((tdim, num_tokens), dtype=np.float32)
template_mask = np.zeros((tdim, num_tokens), dtype=np.float32)
query_to_template = np.zeros((tdim, num_tokens), dtype=np.int64)
visibility_ids = np.zeros((tdim, num_tokens), dtype=np.float32)
# Convert to one-hot
res_type = torch.from_numpy(res_type)
res_type = one_hot(res_type, num_classes=const.num_tokens)
return {
"template_restype": res_type,
"template_frame_rot": torch.from_numpy(frame_rot),
"template_frame_t": torch.from_numpy(frame_t),
"template_cb": torch.from_numpy(cb_coords),
"template_ca": torch.from_numpy(ca_coords),
"template_mask_cb": torch.from_numpy(cb_mask),
"template_mask_frame": torch.from_numpy(frame_mask),
"template_mask": torch.from_numpy(template_mask),
"query_to_template": torch.from_numpy(query_to_template),
"visibility_ids": torch.from_numpy(visibility_ids),
} |