def collate_lep(batch):
"""
Collates LEP datapoints into the batch format for Cormorant.
:param batch: The data to be collated.
:type batch: list of datapoints
:param batch: The collated data.
:type batch: dict of Pytorch tensors
"""
batch = {prop: batch_stack([mol[prop] for mol in batch]) for prop in batch[0].keys()}
# Define which fields to keep
to_keep1 = (batch['charges_active'].sum(0) > 0)
to_keep2 = (batch['charges_inactive'].sum(0) > 0)
# Start building the new batch
new_batch = {}
# Copy label data.
new_batch['label'] = batch['label']
# Split structural data and drop zeros
for key in ['charges','positions','one_hot']:
new_batch[key+'1'] = drop_zeros( batch[key+'_active'], key+'_active', to_keep1 )
new_batch[key+'2'] = drop_zeros( batch[key+'_inactive'], key+'_inactive', to_keep2 )
# Define the atom masks
atom_mask1 = new_batch['charges1'] > 0
atom_mask2 = new_batch['charges2'] > 0
new_batch['atom_mask1'] = atom_mask1
new_batch['atom_mask2'] = atom_mask2
# Define the edge masks
edge_mask1 = atom_mask1.unsqueeze(1) * atom_mask1.unsqueeze(2)
edge_mask2 = atom_mask2.unsqueeze(1) * atom_mask2.unsqueeze(2)
new_batch['edge_mask1'] = edge_mask1
new_batch['edge_mask2'] = edge_mask2
return new_batch
def collate_lba(batch):
"""
Collates LBA datapoints into the batch format for Cormorant.
:param batch: The data to be collated.
:type batch: list of datapoints
:param batch: The collated data.
:type batch: dict of Pytorch tensors
"""
batch = {prop: batch_stack([mol[prop] for mol in batch]) for prop in batch[0].keys()}
# Define which fields to keep
to_keep = (batch['charges'].sum(0) > 0)
# Start building the new batch
new_batch = {}
# Copy label data.
new_batch['neglog_aff'] = batch['neglog_aff']
# Split structural data and drop zeros
for key in ['charges','positions','one_hot']:
new_batch[key] = drop_zeros( batch[key], key, to_keep )
# Define the atom mask
atom_mask = new_batch['charges'] > 0
new_batch['atom_mask'] = atom_mask
# Define the edge mask
edge_mask = atom_mask.unsqueeze(1) * atom_mask.unsqueeze(2)
new_batch['edge_mask'] = edge_mask
return new_batch