refactoring + sva

.gitignore

@@ -1,3 +1,9 @@
+scratch/
+# data/
+data/*.npy
+data/*.tsv
+data/*.pkl
+
 # 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:fbd2453a7cca1d1420476aa527873c065cff16e25af7296bfb18e1ba510ed7b8
+size 5853

create_new_learn_then_test.ipynb → data/create_new_learn_then_test.ipynb

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

Interactive-1.ipynb → data/create_smaller_ltt.ipynb

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:ad0aca44189d80ee7752b2fa0509d027d12e7d18d47bfe9737a50131054c1d96
-size 402714
+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

analyze_protein_vec_results.ipynb → pfam/analyze_protein_vec_results.ipynb

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:10e65ba7a87a4a3bfebbd30007f0b8ed4173ac59f6d47e14cf8c850f1910a021
-size 13863852
+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/{scope_utils.py → util.py}

@@ -1,23 +1,32 @@
-import numpy as np
-import matplotlib.pyplot as plt
-import pdb
-import ipdb
 from sklearn.isotonic import IsotonicRegression
-import seaborn as sns
+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']
+        similarity = query["S_i"]
         sims += similarity.tolist()
         if partial:
-            #labels_to_append = np.logical_or.reduce(query['partial'], axis=1).tolist()
-            # NOTE: no need to do the above for scope - i already handle the pre-processing
-            labels_to_append = query['partial']
+            labels_to_append = np.logical_or.reduce(query["partial"], axis=1).tolist()
         else:
-            labels_to_append = query['exact']
+            labels_to_append = query["exact"]
         labels += labels_to_append
     return sims, labels
 
@@ -29,67 +38,81 @@ def get_arbitrary_attribute(data, attribute: str):
         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']
+        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')
+            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 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]
+    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)
+    false_discoveries = ((1 - labels) * (sims >= lam)).sum(axis=1)
     total_discoveries = np.maximum(total_discoveries, 1)
-    return (false_discoveries/total_discoveries).mean()
+    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()
+    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)
+    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()
+    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)
+    false_discoveries = ((1 - labels) * (sims >= lam)).sum(axis=1)
     total_discoveries = np.maximum(total_discoveries, 1)
-    return (false_discoveries/total_discoveries).std()
+    return (false_discoveries / total_discoveries).std()
+
 
 def get_thresh_FDR(labels, sims, alpha, delta=0.5, N=5000):
     # FDR control with LTT
@@ -97,114 +120,66 @@ def get_thresh_FDR(labels, sims, alpha, delta=0.5, N=5000):
     # 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)] )
+    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) ])
+    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
 
-def get_isotone_regression(data):
-    sims, labels = get_sims_labels(data, partial=True)
-    ir = IsotonicRegression(out_of_bounds='clip')
+
+# 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
 
-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
-
-"""
-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 get_scope_dict(true_test_idcs, test_df, lookup_idcs, lookup_df, D, I):
+# Simplified version of Venn Abers prediction
+def simplifed_venn_abers_prediction(calib_data, test_data_point):
     """
-    true_test_idcs: indices of the test set within the scope dataframe
+    Perform simplified Venn Abers prediction.
 
-    test_df: dataframe containing the test set (indices are the same as true_test_idcs)
+    Args:
+        calib_data (list): List of calibration data points.
+        test_data_point: The test data point to be predicted.
 
-    lookup_idcs: indices of the lookup set within the larger scope dataframe
+    Returns:
+        Tuple: A tuple containing the predicted probabilities for two isotonic regressions.
+    """
 
-    lookup_df: dataframe containing the lookup set (indices are the same as lookup_idcs)
+    sims, labels = get_sims_labels(calib_data, partial=False)
+    print(sims)
+    print(labels)
+    print(len(sims))
+    print(len(labels))
 
-    D: distances matrix (400 x 14777 or test x lookup by default)
+    sims.append(test_data_point)
+    labels.append(True)
+    print(len(sims))
+    print(len(labels))
 
-    I: indices matrix (400 x 14777 or test x lookup by default)
+    ir_0 = IsotonicRegression(out_of_bounds="clip")
+    ir_1 = IsotonicRegression(out_of_bounds="clip")
 
-    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
-    """
+    ir_0.fit(sims, labels)
 
-    near_ids = []
-    min_sim = np.min(D)
-    max_sim = np.max(D)
+    # run the second isotonic regression with the last point as a false label
+    labels[-1] = False
+    ir_1.fit(sims, labels)
 
-    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]
+    p_0 = ir_0.predict([test_data_point])[0]
+    p_1 = ir_1.predict([test_data_point])[0]
 
-        # 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]]
+    return p_0, p_1
 
-        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 validate_lhat(data, lhat):
     total_missed = 0
@@ -215,16 +190,16 @@ def validate_lhat(data, lhat):
     total_partial = 0
     total_partial_identified = 0
     for query in data:
-        idx = query['exact']
+        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)
+        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']
+            idx_partial = query["partial"]
+
+        sims = query["S_i"]
         sims_exact = sims[idx]
         sims_partial = sims[idx_partial]
         total_missed += (sims_exact < lhat).sum()
@@ -237,9 +212,14 @@ def validate_lhat(data, lhat):
         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
-
-
+    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
@@ -255,11 +235,39 @@ def load_database(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 = 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
+    loss_thresh = alpha - (wc_loss - alpha) / len(
+        data
+    )  # normalize by size of calib set
     losses = []
     best_lam = None
     for lam in lambdas:
@@ -275,28 +283,37 @@ def get_thresh_hierarchical(data, lambdas, alpha):
 
     return best_lam, losses
 
+
 def get_hierarchical_loss(data_, lambda_):
     losses = []
     for query in data_:
-        thresh_idx = query['Sum_Norm_S_i'] <= lambda_
+        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
+            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
+    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
+    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
@@ -309,42 +326,119 @@ def get_thresh_max_hierarchical(data, lambdas, alpha, sim="cosine"):
 
     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_
+            thresh_idx = query["S_i"] >= lambda_
         else:
-            thresh_idx = query['S_i'] <= lambda_
+            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
+            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
 
-def query(index, queries, k=10):
-    # On the fly FAISS import to prevent installation issues
-    import faiss
+    test_df: dataframe containing the test set (indices are the same as true_test_idcs)
 
-    # Search indexed database
-    faiss.normalize_L2(queries)
-    D, I = index.search(queries, k)
+    lookup_idcs: indices of the lookup set within the larger scope dataframe
 
-    return (D, I)
+    lookup_df: dataframe containing the lookup set (indices are the same as lookup_idcs)
 
-def build_scope_tree(list_sccs):
+    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
     """
-    Build a scope tree from a list of SCCs
+
+    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):
     """
-    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
+    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

util.py

@@ -1,183 +0,0 @@
-from sklearn.isotonic import IsotonicRegression
-import numpy as np
-from scipy.stats import binom, norm
-
-
-def get_sims_labels(data, partial=False):
-    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_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):
-
-
-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
-
-
-def get_isotone_regression(data):
-    sims, labels = get_sims_labels(data, partial=True)
-    ir = IsotonicRegression(out_of_bounds="clip")
-    ir.fit(sims, labels)
-    return ir
-
-
-def simplifed_venn_abers_prediction(calib_data, test_data_point):
-    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)
-
-    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,
-    )