hjxdl 0.0.1__py3-none-any.whl

hdl/metric_loss/loss.py

@@ -0,0 +1,79 @@
+import typing as t
+
+import torch
+from torch import nn
+
+from .multi_label import BPMLLLoss
+
+
+def get_lossfunc(
+    name: str,
+    *args,
+    **kwargs
+) -> t.Callable:
+    """Get loss function by name
+
+    Args:
+        name (str): The name of the loss function
+
+    Returns:
+        t.Callable: the loss function
+    """
+    name = name.lower()
+    if name == 'bce':
+        return nn.BCELoss(*args, **kwargs)
+    elif name == 'ce':
+        return nn.CrossEntropyLoss(*args, **kwargs)
+    elif name == 'mse':
+        return nn.MSELoss(*args, **kwargs)
+    elif name == 'bpmll':
+        return BPMLLLoss(*args, **kwargs)
+    elif name == 'nll':
+        return nn.GaussianNLLLoss(*args, **kwargs)
+
+
+def mtmc_loss(
+    y_preds: t.Iterable,
+    y_trues: t.Iterable,
+    loss_names: t.Iterable[str] = None,
+    individual: bool = False,
+    task_weights: t.List = None,
+    device=torch.device('cpu'),
+    **kwargs
+):
+    num_tasks = len(y_preds)
+    if loss_names is None: 
+        loss_func = nn.CrossEntropyLoss()
+        loss_funcs = [loss_func] * num_tasks
+    elif isinstance(loss_names, str):
+        loss_func = get_lossfunc(loss_names, **kwargs)
+        loss_funcs = [loss_func] * num_tasks
+    else:
+        loss_funcs = [
+            get_lossfunc(loss_str)
+            for loss_str in loss_names
+        ]
+
+    if task_weights is None:
+        task_weights = torch.ones(num_tasks).to(device)
+    else:
+        assert len(task_weights) == num_tasks
+        task_weights = torch.FloatTensor(task_weights).to(device)
+    
+    loss_values = [
+        loss_func(y_pred, y_true)
+        for y_pred, y_true, loss_func in zip(
+            y_preds, y_trues, loss_funcs
+        )
+    ]
+ 
+    loss_final = sum([
+        loss_value * task_weight
+        for loss_value, task_weight in zip(loss_values, task_weights)
+    ])
+    # loss_final = sum(loss_values) / num_tasks
+    if not individual:
+        return loss_final
+    else:
+        loss_list = [loss_value for loss_value in loss_values]
+        return (loss_final, loss_list)

hdl/metric_loss/metric.py

@@ -0,0 +1,178 @@
+import math
+from typing import Callable, List, Union
+from functools import partial
+
+import numpy as np
+from sklearn.metrics import (
+    auc, mean_absolute_error, mean_squared_error,
+    precision_recall_curve, r2_score,
+    roc_auc_score, accuracy_score, log_loss, matthews_corrcoef,
+    # top_k_accuracy_score
+) 
+import torch
+import torch.nn as nn
+import scipy
+
+
+def prc_auc(targets: List[int], preds: List[float]) -> float:
+    """
+    Computes the area under the precision-recall curve.
+
+    :param targets: A list of binary targets.
+    :param preds: A list of prediction probabilities.
+    :return: The computed prc-auc.
+    """
+    precision, recall, _ = precision_recall_curve(targets, preds)
+    return auc(recall, precision)
+
+
+def bce(targets: List[int], preds: List[float]) -> float:
+    """
+    Computes the binary cross entropy loss.
+
+    :param targets: A list of binary targets.
+    :param preds: A list of prediction probabilities.
+    :return: The computed binary cross entropy.
+    """
+    # Don't use logits because the sigmoid is added in all places except training itself
+    bce_func = nn.BCELoss(reduction='mean')
+    loss = bce_func(target=torch.Tensor(targets), input=torch.Tensor(preds)).item()
+
+    return loss
+
+
+def rmse(targets: List[float], preds: List[float]) -> float:
+    """
+    Computes the root mean squared error.
+
+    :param targets: A list of targets.
+    :param preds: A list of predictions.
+    :return: The computed rmse.
+    """
+    return math.sqrt(mean_squared_error(targets, preds))
+
+
+def mse(targets: List[float], preds: List[float]) -> float:
+    """
+    Computes the mean squared error.
+
+    :param targets: A list of targets.
+    :param preds: A list of predictions.
+    :return: The computed mse.
+    """
+    return mean_squared_error(targets, preds)
+
+
+def accuracy(targets: List[int], preds: Union[List[float], List[List[float]]], threshold: float = 0.5) -> float:
+    """
+    Computes the accuracy of a binary prediction task using a given threshold for generating hard predictions.
+
+    Alternatively, computes accuracy for a multiclass prediction task by picking the largest probability.
+
+    :param targets: A list of binary targets.
+    :param preds: A list of prediction probabilities.
+    :param threshold: The threshold above which a prediction is a 1 and below which (inclusive) a prediction is a 0.
+    :return: The computed accuracy.
+    """
+    if type(preds[0]) == list:  # multiclass
+        hard_preds = [p.index(max(p)) for p in preds]
+    else:
+        hard_preds = [1 if p > threshold else 0 for p in preds]  # binary prediction
+
+    return accuracy_score(targets, hard_preds)
+
+
+def rsquared(x, y):
+    """ Return R^2 where x and y are array-like."""
+
+    _, _, r_value, _, _ = scipy.stats.linregress(x, y)
+    return r_value ** 2
+
+
+def mcc(y_true, y_pred):
+    y_true = np.array(y_true).astype(int)
+    # y_true = np.where(y_true == 1, 1, -1).astype(int)
+    y_pred = np.array(y_pred)
+    y_pred = (y_pred >= 0.5).astype(int)
+
+    return matthews_corrcoef(y_true, y_pred)
+
+
+def topk(y_true, y_pred, k=1):
+
+    y_true = np.array(y_true).astype(int)
+
+    y_pred = np.array(y_pred)
+
+    sorted_pred = np.argsort(y_pred, axis=1, kind='mergesort')[:, ::-1]
+    hits = (y_true == sorted_pred[:, :k].T).any(axis=0)
+    num_hits = np.sum(hits)
+
+    return num_hits / len(y_true) 
+
+
+def get_metric(metric: str) -> Callable[[Union[List[int], List[float]], List[float]], float]:
+    r"""
+    Gets the metric function corresponding to a given metric name.
+
+    Supports:
+
+    * :code:`auc`: Area under the receiver operating characteristic curve
+    * :code:`prc-auc`: Area under the precision recall curve
+    * :code:`rmse`: Root mean squared error
+    * :code:`mse`: Mean squared error
+    * :code:`mae`: Mean absolute error
+    * :code:`r2`: Coefficient of determination R\ :superscript:`2`
+    * :code:`accuracy`: Accuracy (using a threshold to binarize predictions)
+    * :code:`cross_entropy`: Cross entropy
+    * :code:`binary_cross_entropy`: Binary cross entropy
+
+    :param metric: Metric name.
+    :return: A metric function which takes as arguments a list of targets and a list of predictions and returns.
+    """
+    if metric == 'mcc':
+        return mcc
+
+    if metric == 'rsquared':
+        return rsquared
+
+    if metric == 'auc':
+        return roc_auc_score
+
+    if metric == 'prc-auc':
+        return prc_auc
+
+    if metric == 'rmse':
+        return rmse
+
+    if metric == 'mse':
+        return mse
+
+    if metric == 'mae':
+        return mean_absolute_error
+
+    if metric == 'r2':
+        return r2_score
+
+    if metric == 'acc':
+        return accuracy
+
+    if metric == 'ce':
+        return log_loss
+
+    if metric == 'bce':
+        return bce
+    
+    if metric == 'topk':
+        return topk 
+    
+    if metric == 'top3':
+        return partial(topk, k=3)
+    
+    if metric == 'top5':
+        return partial(topk, k=5)
+    
+    if metric == 'top10':
+        return partial(topk, k=10)
+
+    raise ValueError(f'Metric "{metric}" not supported.')

hdl/metric_loss/multi_label.py

@@ -0,0 +1,42 @@
+import torch
+from torch import Tensor
+
+
+class BPMLLLoss(torch.nn.Module):
+    def __init__(self, reduction: str = 'mean', bias=(1, 1)):
+        super(BPMLLLoss, self).__init__()
+        self.bias = bias
+        self.reduction = reduction
+        assert len(self.bias) == 2 \
+            and all(map(lambda x: isinstance(x, int) and x > 0, bias)), \
+            "bias must be positive integers"
+
+    def forward(self, c: Tensor, y: Tensor) -> Tensor:
+        r"""
+        compute the loss, which has the form:
+        L = \sum_{i=1}^{m} \frac{1}{|Y_i| \cdot |\bar{Y}_i|} \sum_{(k, l) \in Y_i \times \bar{Y}_i} \exp{-c^i_k+c^i_l}
+        :param c: prediction tensor, size: batch_size * n_labels
+        :param y: target tensor, size: batch_size * n_labels
+        :return: size: scalar tensor
+        """
+        y = y.float()
+        y_bar = -y + 1
+        y_norm = torch.pow(y.sum(dim=(1,)), self.bias[0])
+        y_bar_norm = torch.pow(y_bar.sum(dim=(1,)), self.bias[1])
+        assert torch.all(y_norm != 0) or torch.all(y_bar_norm != 0), \
+            "an instance cannot have none or all the labels"
+        loss = 1 / torch.mul(y_norm, y_bar_norm) \
+            * self.pairwise_sub_exp(y, y_bar, c)
+        
+        if self.reduction == 'mean':
+            return torch.mean(loss)
+        elif self.reduction == 'none':
+            return loss  
+
+    def pairwise_sub_exp(self, y: Tensor, y_bar: Tensor, c: Tensor) -> Tensor:
+        r"""
+        compute \sum_{(k, l) \in Y_i \times \bar{Y}_i} \exp{-c^i_k+c^i_l}
+        """
+        truth_matrix = y.unsqueeze(2).float() @ y_bar.unsqueeze(1).float()
+        exp_matrix = torch.exp(c.unsqueeze(1) - c.unsqueeze(2))
+        return (torch.mul(truth_matrix, exp_matrix)).sum(dim=(1, 2))

hdl/metric_loss/nt_xent.py

@@ -0,0 +1,65 @@
+import torch
+import numpy as np
+
+
+class NTXentLoss(torch.nn.Module):
+
+    def __init__(self, device, batch_size, temperature, use_cosine_similarity):
+        super(NTXentLoss, self).__init__()
+        self.batch_size = batch_size
+        self.temperature = temperature
+        self.device = device
+        self.softmax = torch.nn.Softmax(dim=-1)
+        self.mask_samples_from_same_repr = self._get_correlated_mask().type(torch.bool)
+        self.similarity_function = self._get_similarity_function(use_cosine_similarity)
+        self.criterion = torch.nn.CrossEntropyLoss(reduction="sum")
+
+    def _get_similarity_function(self, use_cosine_similarity):
+        if use_cosine_similarity:
+            self._cosine_similarity = torch.nn.CosineSimilarity(dim=-1)
+            return self._cosine_simililarity
+        else:
+            return self._dot_simililarity
+
+    def _get_correlated_mask(self):
+        diag = np.eye(2 * self.batch_size)
+        l1 = np.eye((2 * self.batch_size), 2 * self.batch_size, k=-self.batch_size)
+        l2 = np.eye((2 * self.batch_size), 2 * self.batch_size, k=self.batch_size)
+        mask = torch.from_numpy((diag + l1 + l2))
+        mask = (1 - mask).type(torch.bool)
+        return mask.to(self.device)
+
+    @staticmethod
+    def _dot_simililarity(x, y):
+        v = torch.tensordot(x.unsqueeze(1), y.T.unsqueeze(0), dims=2)
+        # x shape: (N, 1, C)
+        # y shape: (1, C, 2N)
+        # v shape: (N, 2N)
+        return v
+
+    def _cosine_simililarity(self, x, y):
+        # x shape: (N, 1, C)
+        # y shape: (1, 2N, C)
+        # v shape: (N, 2N)
+        v = self._cosine_similarity(x.unsqueeze(1), y.unsqueeze(0))
+        return v
+
+    def forward(self, zis, zjs):
+        representations = torch.cat([zjs, zis], dim=0)
+
+        similarity_matrix = self.similarity_function(representations, representations)
+
+        # filter out the scores from the positive samples
+        l_pos = torch.diag(similarity_matrix, self.batch_size)
+        r_pos = torch.diag(similarity_matrix, -self.batch_size)
+        positives = torch.cat([l_pos, r_pos]).view(2 * self.batch_size, 1)
+
+        negatives = similarity_matrix[self.mask_samples_from_same_repr].view(2 * self.batch_size, -1)
+
+        logits = torch.cat((positives, negatives), dim=1)
+        logits /= self.temperature
+
+        labels = torch.zeros(2 * self.batch_size).to(self.device).long()
+        loss = self.criterion(logits, labels)
+
+        return loss / (2 * self.batch_size)

hdl/models/chiral_gnn.py

@@ -0,0 +1,176 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from torch_geometric.nn import global_add_pool, global_mean_pool, global_max_pool, GlobalAttention, Set2Set
+from hdl.layers.graph.chiral_graph import (
+    GCNConv,
+    GINEConv,
+    DMPNNConv,
+    # get_tetra_update,
+)
+
+from hdl.layers.graph.tetra import (
+    # get_tetra_update,
+    TETRA_UPDATE_DICT
+)
+
+
+class GNN(nn.Module):
+    def __init__(
+        self,
+        # args,
+        num_node_features: int = 48,
+        num_edge_features: int = 7,
+        depth: int = 15,
+        hidden_size: int = 128,
+        dropout: float = 0.1,
+        gnn_type: str = 'dmpnn',
+        graph_pool: str = 'mean',
+        tetra: bool = True,
+        task: str = 'classification', 
+        output_num: int = None,
+        message: str = 'tetra_permute_concat',
+        include_vars: bool = False,
+    ):
+        super(GNN, self).__init__()
+
+        self.init_args = {
+            "num_node_features": num_node_features,
+            "num_edge_features": num_edge_features,
+            "depth": depth,
+            "hidden_size": hidden_size,
+            "dropout": dropout,
+            "gnn_type": gnn_type,
+            "graph_pool": graph_pool,
+            "tetra": tetra,
+            "task": task,
+            "message": message,
+            "include_vars": include_vars
+        }
+
+        self.depth = depth
+        self.hidden_size = hidden_size
+        self.dropout = dropout
+        self.gnn_type = gnn_type
+        self.graph_pool = graph_pool
+        self.tetra = tetra
+        self.task = task
+        self.out_dim = output_num
+        self.include_vars = include_vars
+
+        if self.gnn_type == 'dmpnn':
+            self.edge_init = nn.Linear(61, self.hidden_size)
+            self.edge_to_node = DMPNNConv(
+                hidden_size=hidden_size,
+                tetra=tetra,
+                message=message
+            )
+        else:
+            self.node_init = nn.Linear(num_node_features, self.hidden_size)
+            self.edge_init = nn.Linear(13, self.hidden_size)
+
+        # layers
+        self.convs = torch.nn.ModuleList()
+
+        for _ in range(self.depth):
+            if self.gnn_type == 'gin':
+                self.convs.append(GINEConv(
+                    hidden_size=hidden_size,
+                    tetra=tetra,
+                    message=message
+                ))
+            elif self.gnn_type == 'gcn':
+                self.convs.append(GCNConv(
+                    hidden_size=hidden_size,
+                    tetra=tetra,
+                    message=message
+                ))
+            elif self.gnn_type == 'dmpnn':
+                self.convs.append(DMPNNConv(
+                    hidden_size=hidden_size,
+                    tetra=tetra,
+                    message=message
+                ))
+            else:
+                ValueError('Undefined GNN type called {}'.format(self.gnn_type))
+
+        # graph pooling
+        if self.tetra:
+            self.tetra_update = TETRA_UPDATE_DICT[message](hidden_size)
+            # self.tetra_update = get_tetra_update(args)
+
+        if self.graph_pool == "sum":
+            self.pool = global_add_pool
+        elif self.graph_pool == "mean":
+            self.pool = global_mean_pool
+        elif self.graph_pool == "max":
+            self.pool = global_max_pool
+        elif self.graph_pool == "attn":
+            self.pool = GlobalAttention(
+                gate_nn=torch.nn.Sequential(torch.nn.Linear(self.hidden_size, 2 * self.hidden_size),
+                                            torch.nn.BatchNorm1d(2 * self.hidden_size),
+                                            torch.nn.ReLU(),
+                                            torch.nn.Linear(2 * self.hidden_size, 1)))
+        elif self.graph_pool == "set2set":
+            self.pool = Set2Set(self.hidden_size, processing_steps=2)
+        else:
+            raise ValueError("Invalid graph pooling type.")
+
+        # ffn
+        self.mult = 2 if self.graph_pool == "set2set" else 1
+        if self.include_vars:
+            out_dim = 2
+        elif self.out_dim:
+            out_dim = self.out_dim
+        else:
+            out_dim = 1
+        self.ffn = nn.Linear(self.mult * self.hidden_size, out_dim)
+
+    def forward(self, data):
+        x, edge_index, edge_attr, batch, parity_atoms, parity_bond_index = data.x, data.edge_index, data.edge_attr, data.batch, data.parity_atoms, data.parity_bond_index
+
+        if self.gnn_type == 'dmpnn':
+            row, col = edge_index
+            edge_attr = torch.cat([x[row], edge_attr], dim=1)
+            edge_attr = F.relu(self.edge_init(edge_attr))
+        else:
+            x = F.relu(self.node_init(x))
+            edge_attr = F.relu(self.edge_init(edge_attr))
+
+        x_list = [x]
+        edge_attr_list = [edge_attr]
+
+        # convolutions
+        for layer_idx in range(self.depth):
+
+            x_h, edge_attr_h = self.convs[layer_idx](x_list[-1], edge_index, edge_attr_list[-1], parity_atoms, parity_bond_index)
+            h = edge_attr_h if self.gnn_type == 'dmpnn' else x_h
+
+            if layer_idx == self.depth - 1:
+                h = F.dropout(h, self.dropout, training=self.training)
+            else:
+                h = F.dropout(F.relu(h), self.dropout, training=self.training)
+
+            if self.gnn_type == 'dmpnn':
+                h += edge_attr_h
+                edge_attr_list.append(h)
+            else:
+                h += x_h
+                x_list.append(h)
+
+        # dmpnn edge -> node aggregation
+        if self.gnn_type == 'dmpnn':
+            h, _ = self.edge_to_node(x_list[-1], edge_index, h, parity_atoms, parity_bond_index)
+
+        if self.task == 'regression':
+            output = torch.sigmoid(self.ffn(self.pool(h, batch)))
+        elif self.task == 'classification':
+
+            output = torch.sigmoid(self.ffn(self.pool(h, batch)))
+        # mean = output[:, 0]       
+        if not self.include_vars:
+            return output
+        else:
+            mean, var = F.softplus(output[:, 1])
+            return mean, var