boltzfaa_chunker.sh

#!/usr/bin/bash

# WARN: multi-fasta file and a3m dir name have to share the same name
# take input fasta
input_file="$1"
# set path to MSAs generated by colabfold_search
PATHTOMSAS='./msas/nife_msas'

# Check if the input file is provided
if [ -z "$1" ]; then
    echo "Error: must provide one argument"
    echo "Usage: $(basename $0) [INPUT.fasta]"
    exit 1
fi

# create output_dir
output_dir=${input_file%.*}
file_name=${output_dir##*/}
mkdir -p "${output_dir}"

# awk funciton to split fastas
awk -v outdir="${output_dir}" '
  BEGIN {
    prev_seqname = ""
  }
  /^>/{
    split($1, arr, ">")
    seqname = arr[2]
    gsub(/[\.\|]/, "_", seqname)
    gsub(/[#%^\*\\+!={}?]/, "", seqname)
    if (seqname != prev_seqname && outfile) {
      close(outfile)
    }
    outfile = outdir "/" seqname ".fasta"
    prev_seqname = seqname
  }
  {
    print $0 >> outfile
  }
' "${input_file}"

# rename the fasta headers for boltz input
for file in "${output_dir}"/*.fasta; do
    name=${file##*/}
    name=${name%%.*}

    sed -i "/>/ s#^>.*#>A|protein|${PATHTOMSAS}/${name}.a3m#" "${file}"

done

# split the fasta files into chunks of 5000 between new dirs
chunksize=2500
file_count=0
for file in "${output_dir}"/*.fasta; do
    chunk_dir=CHUNK$(printf %03d $((file_count / chunksize)))
    dir="${output_dir}"/"${chunk_dir}"
    mkdir -p "$dir"
    mv "$file" "$dir"
    ((file_count++))
done

# mv *.faa ${output_dir}
# mv ${output_dir}/${input_file} .