parkererickson/LGGM-Text2Graph

1import torch2from src import utils3 4 5class ExtraMolecularFeatures:6    def __init__(self, dataset_infos):7        self.charge = ChargeFeature(remove_h=dataset_infos.remove_h, valencies=dataset_infos.valencies)8        self.valency = ValencyFeature()9        self.weight = WeightFeature(max_weight=dataset_infos.max_weight, atom_weights=dataset_infos.atom_weights)10 11    def __call__(self, noisy_data):12        charge = self.charge(noisy_data).unsqueeze(-1)      # (bs, n, 1)13        valency = self.valency(noisy_data).unsqueeze(-1)    # (bs, n, 1)14        weight = self.weight(noisy_data)                    # (bs, 1)15 16        extra_edge_attr = torch.zeros((*noisy_data['E_t'].shape[:-1], 0)).type_as(noisy_data['E_t'])17 18        return utils.PlaceHolder(X=torch.cat((charge, valency), dim=-1), E=extra_edge_attr, y=weight)19 20 21class ChargeFeature:22    def __init__(self, remove_h, valencies):23        self.remove_h = remove_h24        self.valencies = valencies25 26    def __call__(self, noisy_data):27        bond_orders = torch.tensor([0, 1, 2, 3, 1.5], device=noisy_data['E_t'].device).reshape(1, 1, 1, -1)28        weighted_E = noisy_data['E_t'] * bond_orders      # (bs, n, n, de)29        current_valencies = weighted_E.argmax(dim=-1).sum(dim=-1)   # (bs, n)30 31        valencies = torch.tensor(self.valencies, device=noisy_data['X_t'].device).reshape(1, 1, -1)32        X = noisy_data['X_t'] * valencies  # (bs, n, dx)33        normal_valencies = torch.argmax(X, dim=-1)               # (bs, n)34 35        return (normal_valencies - current_valencies).type_as(noisy_data['X_t'])36 37 38class ValencyFeature:39    def __init__(self):40        pass41 42    def __call__(self, noisy_data):43        orders = torch.tensor([0, 1, 2, 3, 1.5], device=noisy_data['E_t'].device).reshape(1, 1, 1, -1)44        E = noisy_data['E_t'] * orders      # (bs, n, n, de)45        valencies = E.argmax(dim=-1).sum(dim=-1)    # (bs, n)46        return valencies.type_as(noisy_data['X_t'])47 48 49class WeightFeature:50    def __init__(self, max_weight, atom_weights):51        self.max_weight = max_weight52        self.atom_weight_list = torch.tensor(list(atom_weights.values()))53 54    def __call__(self, noisy_data):55        X = torch.argmax(noisy_data['X_t'], dim=-1)     # (bs, n)56        X_weights = self.atom_weight_list.to(X.device)[X]           # (bs, n)57        return X_weights.sum(dim=-1).unsqueeze(-1).type_as(noisy_data['X_t']) / self.max_weight     # (bs, 1)58