plantcad/evolutionary-constraint

import pandas as pd

CENTER_POSITION = 255
NUCLEOTIDES = list("ACGT")

df = pd.read_csv("hf://datasets/kuleshov-group/cross-species-single-nucleotide-annotation/Evolutionary_constraint/valid.tsv", sep="\t")
df = df.rename(columns={"sequences": "seq"})
df["chrom_pos"] = df.pos
df["pos"] = CENTER_POSITION
df["ref"] = df.seq.str[CENTER_POSITION]

def subsample(df, n, seed):
    return (
        df.groupby('label').apply(lambda x: x.sample(n=n // 2, random_state=seed))
        .reset_index(drop=True)
        .sort_values(["chrom", "chrom_pos"])
    )

df.drop(columns=["ref"]).to_parquet("valid_full.parquet", index=False)
subsample(df[df.ref.isin(NUCLEOTIDES)].drop(columns=["ref"]), 10_000, 42).to_parquet("valid_10k.parquet", index=False)