from collections import Counter, defaultdict from ortools.linear_solver import pywraplp import random from omegaconf import DictConfig import pandas as pd from pathlib import Path import os import numpy as np from sklearn.model_selection import train_test_split from dpacman.data_tasks.fimo.post_fimo import get_reverse_complement import json import rootutils from dpacman.utils import pylogger root = rootutils.setup_root(__file__, indicator=".project-root", pythonpath=True) logger = pylogger.RankedLogger(__name__, rank_zero_only=True) def split_with_predefined_test( full_df = pd.DataFrame(), split_names=("train", "val", "test"), test_trs=None, test_dnas=None, ratios=(0.8, 0.1, 0.1), ): """ Method for splitting into train and val with a predefined test set. The proteins in the test set, and the DNA clusters of the DNAs they're associated with, must be excluded from train and val. The remaining rows for train and val are split to preserve 80/10/10 as best as possible. """ full_df[""] test = full_df.copy(deep=True) if test_trs is not None: test = test.loc[test["tr_seqid"].isin(test_trs)].reset_index(drop=True) if test_dnas is not None: test = test.loc[test["dna_seqid"].isin(test_dnas)].reset_index(drop=True) tr_clusters_to_exclude = test["tr_cluster_rep"].unique().tolist() dna_clusters_to_exclude = test["dna_cluster_rep"].unique().tolist() remaining = full_df.loc[ (~full_df["tr_cluster_rep"].isin(tr_clusters_to_exclude)) & (~full_df["dna_cluster_rep"].isin(dna_clusters_to_exclude)) ].reset_index(drop=True) test_ids = test["ID"].unique().tolist() remaining_ids = remaining["ID"].unique().tolist() remaining_clusters = remaining["dna_cluster_rep"].unique().tolis() lost_rows = full_df.loc[ (~full_df["ID"].isin(test_ids)) & (~full_df["ID"].isin(remaining_ids)) ] logger.info(f"Rows in test: {len(test)}") logger.info(f"Rows to be split between train and val: {len(remaining)}") total_rows = len(test) + len(remaining) logger.info(f"Total rows: {total_rows}. Test percentage: {100*len(test)/total_rows:.2f}%") logger.info(f"Lost rows: {len(lost_rows)}") train_ratio_from_remaining = round((0.8*total_rows)/len(remaining), 2) # use sklearn test_size_1 = 1 - train_ratio_from_remaining logger.info( f"\tPerforming first split: non-test clusters -> train clusters ({round(1-test_size_1,3)}) and val ({test_size_1})" ) X = remaining_clusters y = [0] * len(remaining_clusters) X_train, X_val, y_train, y_val = train_test_split( X, y, test_size=test_size_1, random_state=0 ) train = remaining.loc[remaining["dna_cluster_rep"].isin(X_train)] val = remaining.loc[remaining["dna_cluster_rep"].isin(X_val)] leaky_test = lost_rows splits = { "train": train, "val": val, "test": test, "leaky_test": leaky_test } return splits def split_bipartite_fast( dna_clusters, split_names=("train", "val", "test"), ratios=(0.8, 0.1, 0.1), ): # use sklearn test_size_1 = 0.2 test_size_2 = 0.5 logger.info( f"\tPerforming first split: all clusters -> train clusters ({round(1-test_size_1,3)}) and other ({test_size_1})" ) X = dna_clusters y = [0] * len(dna_clusters) X_train, X_test, y_train, y_test = train_test_split( X, y, test_size=test_size_1, random_state=0 ) logger.info( f"\tPerforming second split: other -> val clusters ({round(1-test_size_2,3)}) and test clusters ({test_size_2})" ) X_val, X_test, y_val, y_test = train_test_split( X_test, y_test, test_size=test_size_2, random_state=0 ) dna_assign = {} for x in X_train: dna_assign[x] = "train" for x in X_val: dna_assign[x] = "val" for x in X_test: dna_assign[x] = "test" kept_by_split = {"train": len(X_train), "val": len(X_val), "test": len(X_test)} return dna_assign, kept_by_split def convert_scores(scores): svec = [int(x) for x in scores.split(",")] max_score = max(svec) binary_svec = [0 if x weights). test_edges_must: None, list of pairs, or dict {(tf,dna): required_count}. - If a pair appears with required_count > 0, at least that many examples MUST be kept in TEST. - This implicitly pins both clusters of that pair to TEST (cluster exclusivity). Returns: tf_assign: {tf_cluster -> split} dna_assign: {dna_cluster -> split} kept_by_split: {split -> kept_count} (train/val/test only) total_kept: int split_to_indices: {split -> [input indices]} including 'leaky_test' split_to_edges: {split -> [(tf,dna), ...]} including 'leaky_test' """ # Aggregate counts per pair w = Counter(edges) tfs = {t for (t, _) in w} dnas = {d for (_, d) in w} S = list(split_names) rs = dict(zip(S, ratios)) N = sum(w.values()) if bigM is None: bigM = 1000 * max(1, N) # Index original edges so we can return a per-example split pair_to_indices = defaultdict(list) for idx, (c, d) in enumerate(edges): pair_to_indices[(c, d)].append(idx) if shuffle_within_pair: rng = random.Random(seed) for key in pair_to_indices: rng.shuffle(pair_to_indices[key]) # Parse required test edges req_test = Counter() if test_edges_must: if isinstance(test_edges_must, dict): for k, v in test_edges_must.items(): if not isinstance(k, tuple) or len(k) != 2: raise ValueError( "test_edges_must dict keys must be (tf_cluster, dna_cluster)" ) if v < 0: raise ValueError("required_count must be non-negative") if v: req_test[k] += int(v) else: # assume iterable of pairs req_test = Counter(test_edges_must) # Validate against available counts for pair, req in req_test.items(): if pair not in w: raise ValueError(f"Required test pair {pair} not present in edges.") if req > w[pair]: raise ValueError( f"Required count {req} for {pair} exceeds available {w[pair]}." ) # Build solver solver = pywraplp.Solver.CreateSolver("CBC") if solver is None: raise RuntimeError("Could not create CBC solver.") # Binary cluster assignments x = {(c, s): solver.BoolVar(f"x[{c},{s}]") for c in tfs for s in S} y = {(d, s): solver.BoolVar(f"y[{d},{s}]") for d in dnas for s in S} # Each cluster in exactly one split for c in tfs: solver.Add(sum(x[c, s] for s in S) == 1) for d in dnas: solver.Add(sum(y[d, s] for s in S) == 1) # Integer kept counts per pair and split (allow partial within-pair) k = { ((c, d), s): solver.IntVar(0, w[(c, d)], f"k[{c},{d},{s}]") for (c, d) in w for s in S } # Only keep in split s if both endpoint clusters are assigned to s for (c, d), wt in w.items(): for s in S: solver.Add(k[((c, d), s)] <= wt * x[c, s]) solver.Add(k[((c, d), s)] <= wt * y[d, s]) # Enforce minimum kept counts in TEST for required pairs for (c, d), req in req_test.items(): solver.Add(k[((c, d), "test")] >= req) # Optional: ensure each split has at least one cluster (feasibility depends on counts) if require_nonempty: for s in S: solver.Add(sum(x[c, s] for c in tfs) + sum(y[d, s] for d in dnas) >= 1) # Kept counts per split and total K = {s: solver.IntVar(0, N, f"K[{s}]") for s in S} for s in S: solver.Add(K[s] == sum(k[((c, d), s)] for (c, d) in w)) T = solver.IntVar(0, N, "T") solver.Add(T == sum(K[s] for s in S)) # Ratio deviation: K_s - r_s * T = d+ - d- dpos = {s: solver.NumVar(0, solver.infinity(), f"dpos[{s}]") for s in S} dneg = {s: solver.NumVar(0, solver.infinity(), f"dneg[{s}]") for s in S} for s in S: solver.Add(K[s] - rs[s] * T == dpos[s] - dneg[s]) # Optional hard band around target ratios if ratio_tolerance is not None: eps = float(ratio_tolerance) for s in S: solver.Add(K[s] >= (rs[s] - eps) * T) solver.Add(K[s] <= (rs[s] + eps) * T) # Objective: maximize T then minimize total deviation obj = solver.Objective() obj.SetMaximization() obj.SetCoefficient(T, float(bigM)) for s in S: obj.SetCoefficient(dpos[s], -1.0) obj.SetCoefficient(dneg[s], -1.0) status = solver.Solve() if status not in (pywraplp.Solver.OPTIMAL, pywraplp.Solver.FEASIBLE): raise RuntimeError( "No feasible solution (check ratio_tolerance vs. required test edges)." ) # Read cluster assignments tf_assign = {c: next(s for s in S if x[c, s].solution_value() > 0.5) for c in tfs} dna_assign = {d: next(s for s in S if y[d, s].solution_value() > 0.5) for d in dnas} # Kept counts per split kept_by_split = {s: int(round(K[s].solution_value())) for s in S} total_kept = int(round(T.solution_value())) # ---- Build per-example split assignment (including 'leaky_test') ---- split_to_indices = {s: [] for s in S} remaining_indices = {pair: list(pair_to_indices[pair]) for pair in pair_to_indices} # Allocate the kept examples per split (train/val/test) for (c, d), wt in w.items(): for s in S: cnt = int(round(k[((c, d), s)].solution_value())) if cnt > 0: take = remaining_indices[(c, d)][:cnt] split_to_indices[s].extend(take) remaining_indices[(c, d)] = remaining_indices[(c, d)][cnt:] # Everything left becomes leaky_test leaky_indices = [] for pair, idxs in remaining_indices.items(): if idxs: leaky_indices.extend(idxs) split_to_indices["leaky_test"] = leaky_indices split_to_edges = { s: [edges[i] for i in split_to_indices[s]] for s in split_to_indices } return ( tf_assign, dna_assign, kept_by_split, total_kept, split_to_indices, split_to_edges, ) class DSU: def __init__(self): self.p = {} def find(self, x): if x not in self.p: self.p[x] = x while self.p[x] != x: self.p[x] = self.p[self.p[x]] x = self.p[x] return x def union(self, a, b): ra, rb = self.find(a), self.find(b) if ra != rb: self.p[rb] = ra def split_bipartite_by_components( edges, split_names=("train", "val", "test"), ratios=(0.8, 0.1, 0.1), seed=0, require_nonempty=False, test_edges_must=None, # None, list[(tf,dna)], or dict{(tf,dna): count} ): """ Guarantees exclusivity: each TF cluster and DNA cluster appears in at most one split. Strategy: find connected components in the TF–DNA bipartite graph and assign components wholesale. """ rng = random.Random(seed) w = Counter(edges) # multiplicities per pair if not w: raise ValueError("No edges.") # 1) Build components with Union-Find (prefix to keep TF/DNA namespaces disjoint) dsu = DSU() for tf, dna in w: dsu.union(("T", tf), ("D", dna)) comp_pairs = defaultdict(list) comp_weight = defaultdict(int) for (tf, dna), cnt in w.items(): root = dsu.find(("T", tf)) # component id = root of TF endpoint comp_pairs[root].append((tf, dna)) comp_weight[root] += cnt comps = list(comp_pairs.keys()) C = len(comps) S = list(split_names) rs = dict(zip(S, ratios)) N = sum(comp_weight[c] for c in comps) target = {s: int(round(rs[s] * N)) for s in S} # 2) Pin components that contain required TEST pairs pinned = {} # comp_root -> pinned_split ("test") if test_edges_must: req = ( Counter(test_edges_must) if not isinstance(test_edges_must, dict) else Counter(test_edges_must) ) # Map each required pair to its component, ensure feasibility for (tf, dna), r in req.items(): if (tf, dna) not in w: raise ValueError(f"Required pair {(tf,dna)} not present.") if r > w[(tf, dna)]: raise ValueError( f"Required count {r} for {(tf,dna)} exceeds available {w[(tf,dna)]}." ) comp = dsu.find(("T", tf)) if comp in pinned and pinned[comp] != "test": raise ValueError( f"Component conflict: already pinned to {pinned[comp]}, but {(tf,dna)} demands test." ) pinned[comp] = "test" # NOTE: pinning a pair pins the WHOLE component to test (to keep exclusivity). # If you only want some edges kept in test and discard the rest, handle below when materializing. # 3) Assign components greedily by deficit kept_by_split = {s: 0 for s in S} comp_assign = {} # comp_root -> split # First assign pinned comps for comp, split in pinned.items(): comp_assign[comp] = split kept_by_split[split] += comp_weight[comp] # Sort remaining components by descending weight remaining = [c for c in comps if c not in comp_assign] remaining.sort(key=lambda c: comp_weight[c], reverse=True) # Ensure nonempty splits if requested (seed with largest remaining comps) if require_nonempty: seeds = remaining[: min(len(S), len(remaining))] for comp, s in zip(seeds, S): comp_assign[comp] = s kept_by_split[s] += comp_weight[comp] remaining = [c for c in remaining if c not in comp_assign] for comp in remaining: # choose split with largest deficit (target - current) deficits = {s: target[s] - kept_by_split[s] for s in S} best = max(deficits, key=lambda s: deficits[s]) comp_assign[comp] = best kept_by_split[best] += comp_weight[comp] total_kept = sum(kept_by_split.values()) # 4) Materialize per-example indices (and verify exclusivity) pair_to_indices = defaultdict(list) for idx, pair in enumerate(edges): pair_to_indices[pair].append(idx) split_to_indices = {s: [] for s in S} for comp, s in comp_assign.items(): for pair in comp_pairs[comp]: split_to_indices[s].extend(pair_to_indices[pair]) # Optional: if you pinned a comp due to a small 'must-test' count but # want to *discard* the rest instead of keeping them in test, uncomment: # for comp, s in comp_assign.items(): # if s == "test" and test_edges_must: # # Keep only the required counts; dump extras to 'leaky_test' # ... # (Left out for clarity; default is: keep the whole component in its split.) # 5) Build edge lists and simple cluster assignments split_to_edges = { s: [edges[i] for i in split_to_indices[s]] for s in split_to_indices } tf_assign, dna_assign = {}, {} for comp, s in comp_assign.items(): for tf, dna in comp_pairs[comp]: tf_assign[tf] = s dna_assign[dna] = s # 6) Safety check: no DNA/TF appears in multiple splits tf_in_split = defaultdict(set) dna_in_split = defaultdict(set) for s, elist in split_to_edges.items(): for tf, dna in elist: tf_in_split[tf].add(s) dna_in_split[dna].add(s) dup_tf = {tf: ss for tf, ss in tf_in_split.items() if len(ss) > 1} dup_dna = {dn: ss for dn, ss in dna_in_split.items() if len(ss) > 1} assert not dup_tf and not dup_dna, f"Exclusivity violated: {dup_tf} {dup_dna}" return ( tf_assign, dna_assign, kept_by_split, total_kept, split_to_indices, split_to_edges, ) def print_split_ratios(kept_by_split): total = sum(kept_by_split.values()) train_pcnt = 100 * kept_by_split["train"] / total val_pcnt = 100 * kept_by_split["val"] / total test_pcnt = 100 * kept_by_split["test"] / total logger.info( f"Cluster distribution - Train: {train_pcnt:.2f}%, Val: {val_pcnt:.2f}%, Test: {test_pcnt:.2f}%" ) def make_edges( processed_fimo_path: str, protein_cluster_path: str, dna_cluster_path: str ): """ Make edges for input to the splitting algorithm. Edges consist of: (tr_cluster_rep)_(dna_cluster_rep) where the cluster rep is the sequence ID """ # Read cluser data protein_clusters = pd.read_csv(protein_cluster_path, header=None, sep="\t") protein_clusters.columns = ["tr_cluster_rep", "tr_seqid"] dna_clusters = pd.read_csv(dna_cluster_path, header=None, sep="\t") dna_clusters.columns = ["dna_cluster_rep", "dna_seqid"] # Read datapoints edges = pd.read_parquet(processed_fimo_path) edges = pd.merge(edges, dna_clusters, on="dna_seqid", how="left") edges = pd.merge(edges, protein_clusters, on="tr_seqid", how="left") edges["edge"] = edges.apply( lambda row: (row["tr_cluster_rep"], row["dna_cluster_rep"]), axis=1 ) logger.info(f"Total unique edges: {len(edges['edge'].unique().tolist())}") dup_edges = edges.loc[edges.duplicated("edge")]["edge"].unique().tolist() logger.info(f"Total edges with >1 datapoint: {len(dup_edges)}") logger.info( f"Total datapoints belonging to a duplicate edge: {len(edges.loc[edges['edge'].isin(dup_edges)])}" ) return edges def check_validity(train, val, test, split_by="both"): """ Rigorous check for no overlap Columns = ["ID","dna_sequence","tr_sequence","tr_cluster_rep","dna_cluster_rep", "scores","split"] """ train_ids = set(train["ID"].unique().tolist()) val_ids = set(val["ID"].unique().tolist()) test_ids = set(test["ID"].unique().tolist()) assert len(train_ids.intersection(val_ids)) == 0 assert len(train_ids.intersection(test_ids)) == 0 assert len(val_ids.intersection(test_ids)) == 0 logger.info(f"Pass! No overlap in IDs") if split_by != "dna": train_tr_seqs = set(train["tr_sequence"].unique().tolist()) val_tr_seqs = set(val["tr_sequence"].unique().tolist()) test_tr_seqs = set(test["tr_sequence"].unique().tolist()) assert len(train_tr_seqs.intersection(val_tr_seqs)) == 0 assert len(train_tr_seqs.intersection(test_tr_seqs)) == 0 assert len(val_tr_seqs.intersection(test_tr_seqs)) == 0 logger.info(f"Pass! No overlap in TR sequences") train_tr_reps = set(train["tr_cluster_rep"].unique().tolist()) val_tr_reps = set(val["tr_cluster_rep"].unique().tolist()) test_tr_reps = set(test["tr_cluster_rep"].unique().tolist()) assert len(train_tr_reps.intersection(val_tr_reps)) == 0 assert len(train_tr_reps.intersection(test_tr_reps)) == 0 assert len(val_tr_reps.intersection(test_tr_reps)) == 0 logger.info(f"Pass! No overlap in TR cluster reps") if split_by != "protein": train_dna_seqs = set(train["dna_sequence"].unique().tolist()) val_dna_seqs = set(val["dna_sequence"].unique().tolist()) test_dna_seqs = set(test["dna_sequence"].unique().tolist()) assert len(train_dna_seqs.intersection(val_dna_seqs)) == 0 assert len(train_dna_seqs.intersection(test_dna_seqs)) == 0 assert len(val_dna_seqs.intersection(test_dna_seqs)) == 0 logger.info(f"Pass! No overlap in DNA sequences") train_dna_reps = set(train["dna_cluster_rep"].unique().tolist()) val_dna_reps = set(val["dna_cluster_rep"].unique().tolist()) test_dna_reps = set(test["dna_cluster_rep"].unique().tolist()) assert len(train_dna_reps.intersection(val_dna_reps)) == 0 assert len(train_dna_reps.intersection(test_dna_reps)) == 0 assert len(val_dna_reps.intersection(test_dna_reps)) == 0 logger.info(f"Pass! No overlap in DNA cluster reps") def augment_rc(df): """ Get the reverse complement and add it as a datapoint, effectively doubling the dataset. Also flip the orientation of the scores columns = ["ID","dna_sequence","tr_sequence","tr_cluster_rep","dna_cluster_rep", "scores","split"] """ df_rc = df.copy(deep=True) df_rc["dna_sequence"] = df_rc["dna_sequence"].apply( lambda x: get_reverse_complement(x) ) df_rc["ID"] = df_rc["ID"] + "_rc" df_rc["scores"] = df_rc["scores"].apply(lambda s: ",".join(s.split(",")[::-1])) final_df = pd.concat([df, df_rc]).reset_index(drop=True) return final_df def main(cfg: DictConfig): """ Take a set of DNA clusters + protein clusters, and create the best possible splits into train/val/test. """ # construct edges from training data edge_df = make_edges( processed_fimo_path=Path(root) / cfg.data_task.input_data_path, protein_cluster_path=Path(root) / cfg.data_task.cluster_output_paths.protein, dna_cluster_path=Path(root) / cfg.data_task.cluster_output_paths.dna, ) edges = edge_df["edge"].unique().tolist() # figure out if we actually even have a conflict total_proteins = len(edge_df["tr_seqid"].unique().tolist()) total_protein_clusters = len(edge_df["tr_cluster_rep"].unique().tolist()) no_protein_overlap = (total_proteins) == (total_protein_clusters) logger.info(f"All proteins are in their own clusters: {no_protein_overlap}") if cfg.data_task.split_by == "dna": if cfg.data_task.p_exclude: return else: logger.info(f"Easy split: all proteins are in their own clusters.") dna_clusters = edge_df["dna_cluster_rep"].unique().tolist() results = split_bipartite_fast( dna_clusters, split_names=("train", "val", "test"), ratios=( cfg.data_task.train_ratio, cfg.data_task.val_ratio, cfg.data_task.test_ratio, ), ) dna_assign, kept_by_split = results # assign datapoints to cluster by their DNA cluster rep edge_df["split"] = edge_df["dna_cluster_rep"].map(dna_assign) else: results = split_bipartite_by_components( edges, split_names=("train", "val", "test"), ratios=( cfg.data_task.train_ratio, cfg.data_task.val_ratio, cfg.data_task.test_ratio, ), require_nonempty=cfg.data_task.require_nonempty, seed=cfg.data_task.seed, test_edges_must=None, ) ( tf_assign, dna_assign, kept_by_split, total_kept, split_to_indices, split_to_edges, ) = results # Map each sample to its split print(tf_assign) print(dna_assign) edge_df["tr_split"] = edge_df["tr_cluster_rep"].map(tf_assign) edge_df["dna_split"] = edge_df["dna_cluster_rep"].map(dna_assign) edge_df["same_split"] = ( edge_df["tr_split"] == edge_df["dna_split"] ) # should always be true if easy cluster edge_df["split"] = edge_df["tr_split"] print(edge_df) edge_df["split"] = np.where( edge_df["same_split"], edge_df["split"], # keep existing split if same_split == True "leak", # otherwise leak ) print(edge_df) # Print ratios: hopefully close to desired (e.g. 80/10/10) print_split_ratios(kept_by_split) # Make train, val, test sets # make sure no ID is duplicate assert len(edge_df["ID"].unique()) == len(edge_df) split_cols = [ "ID", "dna_sequence", "tr_sequence", "tr_cluster_rep", "dna_cluster_rep", "scores", "split", ] train = edge_df.loc[edge_df["split"] == "train"].reset_index(drop=True)[split_cols] val = edge_df.loc[edge_df["split"] == "val"].reset_index(drop=True)[split_cols] test = edge_df.loc[edge_df["split"] == "test"].reset_index(drop=True)[split_cols] # ensure there is no overlap check_validity(train, val, test, split_by=cfg.data_task.split_by) total = sum([len(train), len(val), len(test)]) logger.info(f"Length of train dataset: {len(train)} ({100*len(train)/total:.2f}%)") logger.info(f"Length of val dataset: {len(val)} ({100*len(val)/total:.2f}%)") logger.info(f"Length of test dataset: {len(test)} ({100*len(test)/total:.2f}%)") logger.info(f"Total sequences = {total}. Same as edges size? {total==len(edge_df)}") og_unique_dna = pd.concat([train, val, test]) og_unique_dna = len(og_unique_dna["dna_sequence"].unique()) ## Now do RC data augmentation if asked if cfg.data_task.augment_rc: train = augment_rc(train) val = augment_rc(val) test = augment_rc(test) logger.info(f"Added reverse complement sequences to train, val, and test.") check_validity(train, val, test, split_by=cfg.data_task.split_by) total = sum([len(train), len(val), len(test)]) logger.info( f"Length of train dataset: {len(train)} ({100*len(train)/total:.2f}%)" ) logger.info(f"Length of val dataset: {len(val)} ({100*len(val)/total:.2f}%)") logger.info(f"Length of test dataset: {len(test)} ({100*len(test)/total:.2f}%)") logger.info( f"Total sequences = {total}. Same as edges size? {total==len(edge_df)}" ) # since we've added all these new DNA sequences, we do need a new apping of seq id to dna sequence all_data = pd.concat([train, val, test]) all_data["dna_seqid"] = all_data["ID"].str.split("_", n=1, expand=True)[1] dna_dict = dict(zip(all_data["dna_seqid"], all_data["dna_sequence"])) assert len(dna_dict) == len(all_data.drop_duplicates(["dna_sequence"])) new_map_path = str(Path(root) / cfg.data_task.dna_map_path).replace( ".json", "_with_rc.json" ) with open(new_map_path, "w") as f: json.dump(dna_dict, f, indent=2) logger.info( f"Saved DNA maps with reverse complements (len {len(dna_dict)}=2*original map of len {og_unique_dna}=={len(dna_dict)==2*og_unique_dna}) to {new_map_path}" ) # create the output dir split_out_dir = Path(root) / cfg.data_task.split_out_dir os.makedirs(split_out_dir, exist_ok=True) # add binary_scores to allow other training modes train["fimo_binary_sores"] = train["scores"].apply(lambda x: convert_scores(x)) val["fimo_binary_sores"] = val["scores"].apply(lambda x: convert_scores(x)) test["fimo_binary_sores"] = test["scores"].apply(lambda x: convert_scores(x)) # slect final cols and save split_final_cols = ["ID", "dna_sequence", "tr_sequence", "scores", "fimo_binary_sores", "split"] train[split_final_cols].to_csv(split_out_dir / "train.csv", index=False) val[split_final_cols].to_csv(split_out_dir / "val.csv", index=False) test[split_final_cols].to_csv(split_out_dir / "test.csv", index=False) logger.info(f"Saved all splits to {split_out_dir}")