<jupyter_start><jupyter_text>Data Duplication<jupyter_code>len(sets['train']['sequence'].unique()) != len(sets['train'])<jupyter_output><empty_output><jupyter_text>Alignments<jupyter_code>from Bio import pairwise2

alignment = pairwise2.align.globalxx(sets['train']['sequence'][0], sets['train']['sequence'][1])
print(pairwise2.format_alignment(*alignment[0]))<jupyter_output>GSSQETIEVE-DEEEARRVAKEL--R-KKGYEVK----D------E--R---RGNKWHVHRT------
     |     |  || |   ||  | ||  |      |      |  |   |   | | |       
-----T----LD--EA-R---ELVERAKK--E--GTGMDVNGQRFEDWREAER---W-V-R-EQEKNK
  Score=18<jupyter_text>Score per metrics* More unique values => more stability?* Does sequence length matter?<jupyter_code>def append_to_key(d, k, v):
    if k in d:
        values = d[k]
        values.append(v)
        d[k] = values
        
    else:
        d[k] = [v]
        
def compute_mean_scores(d):
    new_d = dict()
    for k in d:
        mean_score = np.mean(d[k])
        new_d[k] = (mean_score, len(d[k]))
        
    return  OrderedDict(sorted(new_d.items()))
score_per_length = dict()
score_per_uniqueness = dict()

for idx, (seq, score) in sets['train'].iterrows():
    # score_per_length
    size = len(seq)
    append_to_key(score_per_length, size, score)
    
    
    # score_per_uniqueness => score per number of unique aminoacids in each seq
    uniques = set(seq)
    unique_size = len(uniques)
    
    append_to_key(score_per_uniqueness, unique_size, score)
    
    # count unique aminoacids
    for amino in uniques:
        uniques.add(amino)

compute_mean_scores(score_per_length)
compute_mean_scores(score_per_uniqueness)<jupyter_output><empty_output><jupyter_text>By Score<jupyter_code>sorted_train = sets['train'].sort_values(by='consensus_stability_score').reset_index()
alignment = pairwise2.align.globalxx(sorted_train['sequence'][0], sorted_train['sequence'][1])
print(pairwise2.format_alignment(*alignment[0]))

sorted_train.head(5)
alignment = pairwise2.align.globalxx(
    sorted_train['sequence'][len(sorted_train)-1], 
    sorted_train['sequence'][len(sorted_train)-2]
)

print(pairwise2.format_alignment(*alignment[0]))

sorted_train.tail(5)<jupyter_output>GSSET-RYRFTDEEE-ARRAAKEWARRGYQVHVTQNGTYWEVEVR
||| | |||||| || |||||||||||||||||||||||||||||
GSS-TTRYRFTD-EEKARRAAKEWARRGYQVHVTQNGTYWEVEVR
  Score=41<jupyter_text>Subsampling techniques * Diversity Maximizing: This is a greedy strategy which starts from the reference and adds the sequence with highest average hamming distance to current set of sequences.* Diversity Minimizing: This strategy is equivalent to the Diversity Maximizing strategy, but adds the se- quence with lowest average hamming distance. It is used to explore the effects of diversity on model per- formance.* HHFilter: This strategy applies hhfilter (Steinegger etal.,2019)withthe-diff Mparameter,whichre- turns M or more sequences that maximize diversity (the result is usually close to M). If more than M sequences are returned we apply the Diversity Maxi- mizing strategy on top of the output.<jupyter_code>import jax.numpy as jnp
import numpy as np

def onehot_pair(u, v):
    vocab = set(u).union(set(v))
    mapping = {k: v for v, k in enumerate(vocab)}
    
    for idx, val in enumerate(u):
        u[idx] = mapping[val]
    
    for idx, val in enumerate(v):
        v[idx] = mapping[val]
        
    return u, v

def onehot_arrays(arrays):
    vocab = set()
    for array in arrays[1:]:
        vocab = vocab.union(set([char for char in array]))
    
    mapping = {k: v for v, k in enumerate(vocab)}
    
    array_final = np.empty_like(arrays)
    for idx, string in enumerate(arrays):
        array = np.array([char for char in string])
        for arr_idx, val in enumerate(array):
            array[arr_idx] = mapping[val]
        array_final[idx] = array
        
    return array_final

def _validate_vector(u, dtype=None):
    u = jnp.asarray(u, dtype=dtype)
    if u.ndim == 1:
        return u

    # Ensure values such as u=1 and u=[1] still return 1-D arrays.
    u = jnp.atleast_1d(u.squeeze())
    if u.ndim > 1:
        raise ValueError("Input vector should be 1-D.")
    warnings.warn(
        "scipy.spatial.distance metrics ignoring length-1 dimensions is "
        "deprecated in SciPy 1.7 and will raise an error in SciPy 1.9.",
        DeprecationWarning)
    return u

def hamming(u, v):
    if np.array(u).dtype == "<U1" and np.array(v).dtype == "<U1":
        u, v = onehot_pair(u, v)
    u = _validate_vector(u)
    v = _validate_vector(v)
    
    u_ne_v = u != v
    return np.average(u_ne_v)
import random

from numba import jit

from tqdm import tqdm
import jax

def compute_diversity(seq1, seq2, method="hamming"):
    if method == "alignment":
        method = pairwise2.align.globalxx
        diversity = method(seq1, seq2)
        diversity = diversity[0].score / max(len(seq1), len(seq2))

    elif method == "hamming":
        if len(seq1) != len(seq2):
            return np.nan

        method = hamming
        diversity = method(list(seq1), list(seq2))

    return diversity

def dataset_diversity(sequences, method="hamming", reduce="mean", verbose=True):
    if reduce != "mean":
        raise NotImplementedError
    else:
        # nan results are due to different lengths
        reducer = np.nanmean

    reduced_divs = []
    if verbose:
        pbar = tqdm(total=len(sequences), miniters=1, smoothing=1)

    all_other_sequences = [
        np.concatenate((sequences[:idx], sequences[idx + 1 :]))
        for idx in range(len(sequences))
    ]
    
    for seq_idx, seq in enumerate(sequences):
        other_sequences = onehot_arrays(all_other_sequences[seq_idx])
        print(other_sequences)
        if seq in other_sequences:
            raise IndexError(
                f"{seq_idx} is in other_sequences, e.g. it is not correctly indexed"
            )

        v_diversity = jax.vmap(lambda x: compute_diversity(seq, x))
        if len(other_sequences) > 1:
            
            div_vs_all = v_diversity()
        else:
            print(f"Skipping sequence {seq_idx}")
            continue

        reduced_div_vs_all = reducer(div_vs_all) if len(div_vs_all) >= 1 else np.nan
        reduced_divs.append(reduced_div_vs_all)

        if verbose:
            pbar.update(1)
            pbar.refresh()

    if verbose:
        pbar.close()

    return reduced_divs
compute_diversity(
    sets["train"]["sequence"].values[0],
    sets["train"]["sequence"].values[11],
    "alignment",
)
compute_diversity(sets["train"]["sequence"].values[0], sets["train"]["sequence"].values[11])
diversities_train = dataset_diversity(sets["train"]["sequence"].values)
diversities_train[:10]
pd.read_csv("data/stability_diversity_train.csv", index_col=0)[:10]
# ds_div_train = dataset_diversity(sets["train"]["sequence"])
# sets["train"]["diversity"] = np.nanmean(np.array(ds_div_train), 1)
# np.nanmean(np.array(ds_div_train))
# 0.8xxx
# ds_div_test = dataset_diversity(sets["test"]["sequence"])
# sets["test"]["diversity"] = np.nanmean(np.array(ds_div_test), 1)
# np.nanmean(np.array(ds_div_test))
# sets["train"].to_csv(f"{data_path}/stability_diversity_train.csv")
# sets["test"].to_csv(f"{data_path}/stability_diversity_test.csv")
sets = {
    'train' : pd.read_csv(f"{data_path}/stability_diversity_train.csv"),
    'test'  : pd.read_csv(f"{data_path}/stability_diversity_test.csv") 
}

sorted_div = {
    "train" : sets["train"].sort_values(by="diversity", axis=0, ascending=False),
    "test"  : sets["test"].sort_values(by="diversity", axis=0, ascending=False),
}
import torch
from torch.utils.data import Dataset

class ProteinStabilityDataset(Dataset):
    """Protein1D Stability dataset."""

    def __init__(self, proteins_path, ret_dict = True):
        """
        Args:
            proteins_path (string): Path to the H5Py file that contains sequences and embeddings.
        """
        self.stability_path = proteins_path
        self.ret_dict = ret_dict

    def __len__(self):
        return len(h5py.File(self.stability_path, "r")['sequences'])
    
    def __getitem__(self, idx):
        if torch.is_tensor(idx):
            idx = idx.tolist()
        
        with h5py.File(self.stability_path, "r") as dset:
            sequences = dset['sequences'][idx]
            embeddings = dset['embeddings'][idx]
            labels = dset['labels'][idx]
            
            sample = {
                'sequences': sequences, 
                'embeddings': torch.from_numpy(embeddings), 
                'labels': torch.Tensor([labels])
            }

        return sample if self.ret_dict else (sample['embeddings'], sample['labels'])
from torch.utils.data import Sampler, DataLoader, SubsetRandomSampler
from typing import List, Sized, Iterator

import numpy as np

class SubsetDiversitySampler(Sampler):
    """Samples elements given their diversity w.r.t. the rest of the dataset"""
    
    def __init__(self, dataset: Sized, diversity_path : str, diversity_cutoff : float, strategy : str = "maximize", seed : int = 123) -> None:
        """
        Args:
          diversity_path (string): Path to the csv with sequences and diversity.
        """
        self.diversity_path = diversity_path
        self.diversity_data = pd.read_csv(self.diversity_path, index_col=0)
        
        if strategy == "maximize":
            sorting_order = False
        elif strategy == "minimize":
            sorting_order = True
        else: 
            raise ValueError(f"Strategy {strategy} is not supported")
        
        self.diversity_data = self.diversity_data.sort_values(
            by='diversity', 
            ascending=sorting_order
        )
        
        self.cutoff = diversity_cutoff
        self.indices = []
        self.strategy = strategy
        self.seed = seed
        
        if strategy == 'maximize':
            self.cutoff_lambda = lambda x, cutoff : x < cutoff
        elif strategy == 'minimize':
            self.cutoff_lambda = lambda x, cutoff : x > cutoff
            
        self.subset_by_cutoff(self.cutoff, self.cutoff_lambda)
            
    def subset_by_cutoff(self, cutoff, cutoff_func) -> None:
        for row in self.diversity_data.itertuples():
            is_cutoff = cutoff_func(row.diversity, cutoff)
            if is_cutoff:
                break
            else:
                self.indices.append(row.Index)
        
    def __iter__(self) -> Iterator[int]:
        rng = np.random.default_rng(self.seed)
        return iter(rng.permutation(self.indices))

    def __len__(self) -> int:
        return len(self.indices)
dataset = ProteinStabilityDataset("data/stability.h5", ret_dict=True)
sampler = SubsetDiversitySampler(dataset, "data/stability_diversity_full.csv", 0.85)
dl = DataLoader(dataset, batch_size=128, sampler=sampler)

# check diversity sampler
df = pd.read_csv("data/stability_diversity_full.csv", index_col=0)
for batch in dl:
    for seq in batch["sequences"]:
        row = df[df.sequence == str(seq, 'utf-8')]
        if float(row.diversity.values) < 0.85:
            raise AssertionError
# for s in [1000, 1500, 2000, 3000, 4000, 5000, 6000]:
#     for kind in ["train"]:
#         chosen_by_diversity = select_by_diversity(
#             sequences=sorted_div[kind]["sequence"].values, 
#             diversities=sorted_div[kind]["diversity"].values,
#             min_size=s, 
#             min_score=0.78
#         )
    
#     subset = sets["train"][sets["train"]["sequence"].isin(chosen_by_diversity)]
#     subset.to_csv(f"{data_path}/stability_diversity_train_{s}.csv", index=False)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.125784541312011, 'fraction_numerical': 0.010919095520591523, 'mean_word_length': 3.019350380096752, 'pattern_counts': {'":': 6, '<': 23, '<?xml version=': 0, '>': 29, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '5', 'dir': 'jupyter-structured-clean-dedup', 'id': '66574', 'license': 'MIT', 'n_tokens_mistral': 3974, 'n_tokens_neox': 3810, 'n_words': 813, 'path': 'project/parallel_synthesis/experiments/exploration.ipynb', 'repo_name': 'stepp1/protera-stability', 'size': '11631'}