merge conflict

.gitignore

@@ -1,3 +1,12 @@
+<<<<<<< HEAD
+=======
+scratch/
+# data/
+data/*.npy
+data/*.tsv
+data/*.pkl
+
+>>>>>>> 55804bf019051081df1543cee90ef5366e3049ae
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]

README.md

@@ -0,0 +1,16 @@
+# Protein conformal retrieval
+
+Code and notebooks from 2024 paper on conformal retrieval with proteins
+
+## Installation
+
+`pip install -e .`
+
+## Structure
+
+`./protein_conformal`: utility functions to creating confidence sets and assigning probabilities to any protein machine learning model for search
+`./scope`: experiments pertraining to SCOPe
+`./pfam`: notebooks demonstrating how to use our techniques to calibrate false discovery and false negative rates for different pfam classes
+`./ec`: experiments pertraining to EC number classification on uniprot
+`./data`: scripts and notebooks used to process data
+`./clean_selection`: scripts and notebooks used to process data

clean_selection/analyze_clean_hierarchical_loss_protein_vec.ipynb

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:17792a81e3b66606738ca77aa001077daa7f5c0063b3f956902996871561817c
-size 86942
+oid sha256:732b364272b85d567d6bb29dc0b1df8dcbd2afcc31a578a5e650e2a689ed9e55
+size 86932

clean_selection/clean_utils.py

@@ -39,8 +39,8 @@ def get_true_labels_test(file_name, test_idx: None):
     return true_label, all_label
 
 
-def infer_conformal(train_data, test_data, thresh, report_metrics = False, 
-                 pretrained=True, model_name=None, test_idx=None, name_id="1"):    
+def infer_conformal(train_data, test_data, thresh, report_metrics=False, 
+                 pretrained=True, model_name=None, test_idx=None, name_id="1"):
     use_cuda = torch.cuda.is_available()
     device = torch.device("cuda:0" if use_cuda else "cpu")
     dtype = torch.float32

clean_selection/process_clean_ec.ipynb

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:b6a7b8853fa3df8fe380cea2391bdace7b76605c3e80770b397c1ba5776aee6c
-size 5860
+oid sha256:744eae0442a80d617d24342f3e00cf49c3460edb260054e9f0dca173c87e5968
+size 5854

data/create_learn_then_test.ipynb

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:375209f78d75e7abe857572b15cccfc82fc6af4d99c93b501c39d88e9d66dcd8
+size 74759

data/create_new_learn_then_test.ipynb

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c321219b09969d691f09fad1c8cab75e0f4735d811c32871e69ee7e074ca8f10
+size 34530

data/create_smaller_ltt.ipynb

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:25f446c429ebf8cc67fa2bd395ac9fb3b363f24f01eef43e6213c38a42cc7e3f
+size 36171

ec/analyze_ec_hierarchical_loss_protein_vec.ipynb

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:136cfeed57224aa55620ad3ad776c89de7a7b24fd2012cc20fc2b1c915ebe6ce
-size 87826
+oid sha256:6af5a8524e4b92b94ceea5a925a4621760194f95639ddb3d964c39322ab14f90
+size 87812

ec/process_pfam_ec.ipynb

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:8762050ecbdc532dd9c962d8dc28ad54a6aa8fafb5058b99c1e2aab19ea594f8
-size 114148
+oid sha256:7a10ed21e5ed16e2de4871a50c53bf32cb0ea104c8f97b92a9b39970b7b2aece
+size 114134

pfam/analyze_protein_vec_results.ipynb

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:51a908707ef800db0c8be4aae8adc45f441123b511da5c9d272f88a48de2f607
+size 399639

pfam/multidomain_search.ipynb

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9fa68613561b4b7386628dd78f5f06b655cdc69bc493a517b79e92669d909a83
+size 2222

protein_conformal/util.py

@@ -0,0 +1,444 @@
+from sklearn.isotonic import IsotonicRegression
+import numpy as np
+from scipy.stats import binom, norm
+
+
+# get similarity scores and labels for a set of proteins
+def get_sims_labels(data, partial=False):
+    """
+    Get the similarities and labels from the given data.
+
+    Parameters:
+    - data: A list of query data.
+    - partial: A boolean indicating whether to use partial hits or exact hits.
+    exact: Pfam1001 == Pfam1001
+    partial: Pfam1001 in [Pfam1001,Pfam1002], where [Pfam1001,Pfam1002] is the set of Pfam domains in a protein.
+
+    Returns:
+    - sims: A list of similarity scores.
+    - labels: A list of labels.
+    """
+    sims = []
+    labels = []
+    for query in data:
+        similarity = query["S_i"]
+        sims += similarity.tolist()
+        if partial:
+            labels_to_append = np.logical_or.reduce(query["partial"], axis=1).tolist()
+        else:
+            labels_to_append = query["exact"]
+        labels += labels_to_append
+    return sims, labels
+
+def get_arbitrary_attribute(data, attribute: str):
+    # get an arbitrary attribute from the data
+    attributes = []
+    for query in data:
+        attribute = query[attribute]
+        attributes += attribute.tolist()
+    return attributes
+
+
+
+def get_thresh(data, alpha):
+    # conformal risk control
+    all_sim_exact = []
+    for query in data:
+        idx = query["exact"]
+        similarity = query["S_i"]
+        sims_to_append = similarity[idx]
+        all_sim_exact += list(sims_to_append)
+        n = len(all_sim_exact)
+        if n > 0:
+            lhat = np.quantile(
+                all_sim_exact,
+                np.maximum(alpha - (1 - alpha) / n, 0),
+                interpolation="lower",
+            )
+        else:
+            lhat = 0
+    return lhat
+
+
+# Bentkus p value
+def bentkus_p_value(r_hat, n, alpha):
+    return binom.cdf(np.ceil(n * r_hat), n, alpha / np.e)
+
+
+# def clt_p_value(r_hat,n,alpha):
+# TODO: we may have wanted to do a different implementation of this
+
+def clt_p_value(r_hat, std_hat, n, alpha):
+    z = (r_hat - alpha) / (std_hat / np.sqrt(n))
+    p_value = norm.cdf(z)
+    return p_value
+
+
+def percentage_of_discoveries(sims, labels, lam):
+    # FDR: Number of false matches / number of matches
+    total_discoveries = (sims >= lam).sum(axis=1)
+    return total_discoveries.mean() / len(labels)  # or sims.shape[1]
+
+
+def risk(sims, labels, lam):
+    # FDR: Number of false matches / number of matches
+    total_discoveries = (sims >= lam).sum(axis=1)
+    false_discoveries = ((1 - labels) * (sims >= lam)).sum(axis=1)
+    total_discoveries = np.maximum(total_discoveries, 1)
+    return (false_discoveries / total_discoveries).mean()
+
+
+def calculate_false_negatives(sims, labels, lam):
+    # FNR: Number of false non-matches / number of non-matches
+    total_non_matches = labels.sum(axis=1)
+    false_non_matches = (labels & (sims < lam)).sum(axis=1)
+    total_non_matches = np.maximum(total_non_matches, 1)
+    return (false_non_matches / total_non_matches).mean()
+
+
+def risk_no_empties(sims, labels, lam):
+    # FDR: Number of false matches / number of matches
+    total_discoveries = (sims >= lam).sum(axis=1)
+    false_discoveries = ((1 - labels) * (sims >= lam)).sum(axis=1)
+    idx = total_discoveries > 0
+    total_discoveries = total_discoveries[idx]
+    false_discoveries = false_discoveries[idx]
+    return (false_discoveries / total_discoveries).mean()
+
+
+def std_loss(sims, labels, lam):
+    # FDR: Number of false matches / number of matches
+    total_discoveries = (sims >= lam).sum(axis=1)
+    false_discoveries = ((1 - labels) * (sims >= lam)).sum(axis=1)
+    total_discoveries = np.maximum(total_discoveries, 1)
+    return (false_discoveries / total_discoveries).std()
+
+
+def get_thresh_FDR(labels, sims, alpha, delta=0.5, N=5000):
+    # FDR control with LTT
+    # labels = np.stack([query['exact'] for query in data], axis=0)
+    # sims = np.stack([query['S_i'] for query in data], axis=0)
+    print(f"sims.max: {sims.max()}")
+    n = len(labels)
+    lambdas = np.linspace(sims.min(), sims.max(), N)
+    risks = np.array([risk(sims, labels, lam) for lam in lambdas])
+    stds = np.array([std_loss(sims, labels, lam) for lam in lambdas])
+    # pvals = np.array( [bentkus_p_value(r,n,alpha) for r in risks] )
+    pvals = np.array([clt_p_value(r, s, n, alpha) for r, s in zip(risks, stds)])
+    below = pvals <= delta
+    # Pick the smallest lambda such that all lambda above it have p-value below delta
+    pvals_satisfy_condition = np.array([np.all(below[i:]) for i in range(N)])
+    lhat = lambdas[np.argmax(pvals_satisfy_condition)]
+    print(f"lhat: {lhat}")
+    print(f"risk: {risk(sims, labels, lhat)}")
+    return lhat
+
+
+# get the isotonic regression
+def get_isotone_regression(data, partial=False):
+    sims, labels = get_sims_labels(data, partial=partial)
+    ir = IsotonicRegression(out_of_bounds="clip")
+    ir.fit(sims, labels)
+    return ir
+
+
+# Simplified version of Venn Abers prediction
+def simplifed_venn_abers_prediction(calib_data, test_data_point):
+    """
+    Perform simplified Venn Abers prediction.
+
+    Args:
+        calib_data (list): List of calibration data points.
+        test_data_point: The test data point to be predicted.
+
+    Returns:
+        Tuple: A tuple containing the predicted probabilities for two isotonic regressions.
+    """
+
+    sims, labels = get_sims_labels(calib_data, partial=False)
+    print(sims)
+    print(labels)
+    print(len(sims))
+    print(len(labels))
+
+    sims.append(test_data_point)
+    labels.append(True)
+    print(len(sims))
+    print(len(labels))
+
+    ir_0 = IsotonicRegression(out_of_bounds="clip")
+    ir_1 = IsotonicRegression(out_of_bounds="clip")
+
+    ir_0.fit(sims, labels)
+
+    # run the second isotonic regression with the last point as a false label
+    labels[-1] = False
+    ir_1.fit(sims, labels)
+
+    p_0 = ir_0.predict([test_data_point])[0]
+    p_1 = ir_1.predict([test_data_point])[0]
+
+    return p_0, p_1
+
+
+def validate_lhat(data, lhat):
+    total_missed = 0
+    total_missed_partial = 0
+    total_exact = 0
+    total_inexact_identified = 0
+    total_identified = 0
+    total_partial = 0
+    total_partial_identified = 0
+    for query in data:
+        idx = query["exact"]
+        # if partial has multiple rows, we want to take the logical or of all of them. Otherwise just set it to the single row
+        # check if there is one or more rows
+        # query['partial'] = np.array(query['partial'])
+        if len(np.array(query["partial"]).shape) > 1:
+            idx_partial = np.logical_or.reduce(query["partial"], axis=1)
+        else:
+            idx_partial = query["partial"]
+
+        sims = query["S_i"]
+        sims_exact = sims[idx]
+        sims_partial = sims[idx_partial]
+        total_missed += (sims_exact < lhat).sum()
+
+        # TODO: are there any divisions by zero here?
+        total_missed_partial += (sims_partial < lhat).sum()
+        total_partial_identified += (sims_partial >= lhat).sum()
+        total_partial += len(sims_partial)
+
+        total_exact += len(sims_exact)
+        total_inexact_identified += (sims[~np.array(idx)] >= lhat).sum()
+        total_identified += (sims >= lhat).sum()
+    return (
+        total_missed / total_exact,
+        total_inexact_identified / total_identified,
+        total_missed_partial / total_partial,
+        total_partial_identified / total_identified,
+    )
+
+### FAISS related functions
+def load_database(lookup_database):
+    """
+    lookup_database: NxM matrix of embeddings
+    """
+    # On the fly FAISS import to prevent installation issues
+    import faiss
+
+    # Build an indexed database
+    d = lookup_database.shape[1]
+    index = faiss.IndexFlatIP(d)
+    faiss.normalize_L2(lookup_database)
+    index.add(lookup_database)
+
+    return index
+
+
+def query(index, queries, k=10):
+    # On the fly FAISS import to prevent installation issues
+    import faiss
+
+    # Search indexed database
+    faiss.normalize_L2(queries)
+    D, I = index.search(queries, k)
+
+    return (D, I)
+
+### functions for hierarchical conformal
+def build_scope_tree(list_sccs):
+    """
+    Build a scope tree from a list of SCCs
+    """
+    tree = {}
+    for sccs in list_sccs:
+        sccs = sccs.split(".")
+        node = tree
+        for s in sccs:
+            if s not in node:
+                node[s] = {}
+            node = node[s]
+    return tree
+
+def get_thresh_hierarchical(data, lambdas, alpha):
+    # get the worst case loss
+    wc_loss = max([np.sum(x["loss"]) for x in data])
+    wc_loss = wc_loss / 14777
+    loss_thresh = alpha - (wc_loss - alpha) / len(
+        data
+    )  # normalize by size of calib set
+    losses = []
+    best_lam = None
+    for lam in lambdas:
+        per_lam_loss = get_hierarchical_loss(data, lam)
+        if per_lam_loss > loss_thresh:
+            break
+        best_lam = lam
+        losses.append(per_lam_loss)
+    print("worst case loss: " + str(wc_loss))
+    print("Loss threshold: " + str(loss_thresh))
+    print("Best lambda: " + str(best_lam))
+    print("Loss of best lambda: " + str(losses[-1]))
+
+    return best_lam, losses
+
+
+def get_hierarchical_loss(data_, lambda_):
+    losses = []
+    for query in data_:
+        thresh_idx = query["Sum_Norm_S_i"] <= lambda_
+        if np.sum(thresh_idx) == 0:
+            loss = 0
+        else:
+            loss = np.sum(np.asarray(query["loss"])[thresh_idx]) / np.sum(
+                thresh_idx
+            )  # NOTE: have fixed denominator, but alpha has to change
+        losses.append(loss)  # average over all queries
+    return np.mean(losses)
+
+
+def get_thresh_max_hierarchical(data, lambdas, alpha, sim="cosine"):
+    # get the worst case loss
+    wc_loss = 4  # in the max case, the max_loss is simply retrieving a protein with different class
+    loss_thresh = alpha - (wc_loss - alpha) / len(
+        data
+    )  # normalize by size of calib set
+    losses = []
+    best_lam = None
+    if sim == "cosine":
+        ## reverse lambdas and return list
+        lambdas = list(reversed(lambdas))
+    for (
+        lam
+    ) in (
+        lambdas
+    ):  # start from the largest lambda since we are dealing with raw similarity scores
+        per_lam_loss = get_hierarchical_max_loss(data, lam, sim=sim)
+        if per_lam_loss > loss_thresh:
+            break
+        best_lam = lam
+        losses.append(per_lam_loss)
+    print("worst case loss: " + str(wc_loss))
+    print("Loss threshold: " + str(loss_thresh))
+    print("Best lambda: " + str(best_lam))
+    print("Loss of best lambda: " + str(losses[-1]))
+
+    return best_lam, losses
+
+
+def get_hierarchical_max_loss(data_, lambda_, sim="cosine"):
+    losses = []
+    for query in data_:
+        if sim == "cosine":
+            thresh_idx = query["S_i"] >= lambda_
+        else:
+            thresh_idx = query["S_i"] <= lambda_
+        if np.sum(thresh_idx) == 0:
+            loss = 0
+        else:
+            loss = np.max(np.asarray(query["loss"])[thresh_idx])  # monotonic loss
+            # loss = np.sum(np.asarray(query['loss'])[thresh_idx]) / np.sum(thresh_idx) # NOTE: have fixed denominator, but alpha has to change
+        losses.append(loss)  # average over all queries
+    return np.mean(losses)
+
+def get_scope_dict(true_test_idcs, test_df, lookup_idcs, lookup_df, D, I):
+    """
+    true_test_idcs: indices of the test set within the scope dataframe
+
+    test_df: dataframe containing the test set (indices are the same as true_test_idcs)
+
+    lookup_idcs: indices of the lookup set within the larger scope dataframe
+
+    lookup_df: dataframe containing the lookup set (indices are the same as lookup_idcs)
+
+    D: distances matrix (400 x 14777 or test x lookup by default)
+
+    I: indices matrix (400 x 14777 or test x lookup by default)
+
+    NOTE: Indices computed by FAISS are not the same as the indices of the dataframe, so
+    we use the lookup_idcs list to map FAISS indices in I to the indices of the dataframe
+    """
+
+    near_ids = []
+    min_sim = np.min(D)
+    max_sim = np.max(D)
+
+    for i in range(len(true_test_idcs)):
+        test_id = test_df.loc[true_test_idcs[i], "sid"]
+        test_sccs = test_df.loc[true_test_idcs[i], "sccs"]
+        query_ids = [lookup_df.loc[lookup_idcs[j], "sid"] for j in I[i]]
+        exact_loss = [
+            scope_hierarchical_loss(test_sccs, lookup_df.loc[lookup_idcs[j], "sccs"])
+            for j in I[i]
+        ]
+        # grab the 2nd element in the tuple belonging to each element of exact_loss as mask_exact
+        mask_exact = [x[1] for x in exact_loss]
+        loss = [x[0] for x in exact_loss]
+
+        # define mask_partial as 1 for any element of loss that is <=1 (tolerate retrieving homolog with diff family but same superfamily)
+        mask_partial = [l <= 1 for l in loss]
+
+        # create a row of size len(lookup_df) where each element is the sum of all entries in S_i until that index
+        sum = np.cumsum(D[i])
+        norm_sim = (D[i] - min_sim) / (
+            max_sim - min_sim
+        )  # convert similarities into a probability space (0, 1) based on (min_sim, max_sim)
+        # mask_exact = [test_sccs == lookup_df.loc[lookup_idcs[j], 'sccs'] for j in I[i]]
+
+        sum_norm_s_i = np.cumsum(norm_sim)
+        near_ids.append(
+            {
+                "test_id": test_id,
+                "query_ids": query_ids,
+                #'meta_query': meta_query,
+                "loss": loss,
+                "exact": mask_exact,
+                "partial": mask_partial,
+                "S_i": D[i],
+                "Sum_i": sum,
+                "Norm_S_i": norm_sim,
+                "Sum_Norm_S_i": sum_norm_s_i,
+                "I_i": I[i],
+            }
+        )
+    return near_ids
+
+"""
+def scope_hierarchical_loss(y_sccs, y_hat_sccs, slack = 0):
+
+    #Find the common ancestor of two sets of SCCs (0 if family, 1 if superfamily, 2 if fold, 3 if class)
+
+    # Find the common ancestor of the two sets of SCCs
+    y_sccs, y_hat_sccs = y_sccs.split('.'), y_hat_sccs.split('.')
+    assert len(y_sccs) == len(y_hat_sccs) == 4
+
+    loss, count = None, 0
+    while loss is None:
+        if y_sccs[-1] != y_hat_sccs[-1]:
+            count += 1
+            y_sccs.pop()
+            y_hat_sccs.pop()
+        else:
+            break
+    loss = count - slack
+    
+    return loss == 0
+"""
+
+def scope_hierarchical_loss(y_sccs, y_hat_sccs):
+    """
+    Find the common ancestor of two sets of SCCs (0 if family, 1 if superfamily, 2 if fold, 3 if class)
+    """
+    y_sccs, y_hat_sccs = y_sccs.split("."), y_hat_sccs.split(".")
+    first_non_matching_index = next(
+        (i for i, (x, y) in enumerate(zip(y_sccs, y_hat_sccs)) if x != y), len(y_sccs)
+    )
+
+    loss = (
+        len(y_sccs) - first_non_matching_index
+    )  # ex if the first mismatch is at idx 2 (0-indexed), that means that the second last label (superfamily) is wrong, which is a loss of two (wrong family & superfamily)
+    exact = True if len(y_sccs) == first_non_matching_index else False
+
+    return loss, exact
+

scope/analyze_scope_hierarchical_loss_protein_vec.ipynb

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:6ec68b7560cd591d661a8f77f9c2df9354aa99b2bf350ad7b96803a9f9eacece
-size 91708
+oid sha256:05e1080b7854b1b02b59c8913509a599d5d65f249f15712aad9d6dc58e877092
+size 91694

scope/analyze_scope_protein_vec.ipynb

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:5445b059ac49809729e487c50f60ed8ac2123f03126f56f77a7348313f26c11b
-size 449975
+oid sha256:15d00e9ddd6e3e23490a415f942065d9f485bac0d437f028eb400853aa75ffc2
+size 449919

scope/test_scope_conformal_retrieval.ipynb

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:c8f4006f8108efcd9cb99a08cabde82ccfda250dc4807f7ff5fbe4a9aaccbe90
-size 3232293
+oid sha256:34d3c6c5df4cef9235c33fd0c73e80507f8ba533d495d5c1f1df39323d52cb21
+size 3232279