ortho_to_generax_workflow.sh

#!/usr/bin/env bash

##################################################################################
#
# Description: A pipeline for gene family tree inference using GeneRax, from
# multiple sequence alignments.
#
# Usage:
#   ./ortho_to_generax_workflow.sh --min_taxa <int> --ncores <int>
#
# Required arguments:
#   --min_taxa  Minimum number of unique species to retain a gene family
#   --ncores    Number of CPU cores for parallel steps
#
# Steps:
#   1. Filter MSAs with fewer than --min_taxa species
#   2. Trim alignments with trimAl
#   3. Infer gene trees with FastTreeMP
#   4. Build the GeneRax families file
#   5. Run GeneRax (UndatedDL model)
#
# Dependencies: parallel, trimal, FastTreeMP, generax, mpiexec
#
##################################################################################

## Argument parsing

set -euo pipefail

min_taxa=""
ncores=""

usage() {
    echo "Usage: $0 --min_taxa <int> --ncores <int>"
    echo "  --min_taxa   Minimum number of unique species to retain a gene family"
    echo "  --ncores     Number of CPU cores for parallel steps"
    exit 1
}

while [[ $# -gt 0 ]]; do
    case "$1" in
        --min_taxa) min_taxa="$2"; shift 2 ;;
        --ncores)   ncores="$2";   shift 2 ;;
        *) echo "Unknown argument: $1"; usage ;;
    esac
done

if [[ -z "$min_taxa" || -z "$ncores" ]]; then
    echo "Error: --min_taxa and --ncores are required."
    usage
fi

if ! [[ "$min_taxa" =~ ^[0-9]+$ ]] || ! [[ "$ncores" =~ ^[0-9]+$ ]]; then
    echo "Error: --min_taxa and --ncores must be positive integers."
    usage
fi

##################################################################################

## Configure paths

# Input MSAs produced by OrthoFinder
msa_path="./MultipleSequenceAlignments"

# Working directories
trim_path="./trimmedMSAs"
tree_path="./fastrees"

# GeneRax inputs / outputs
out_file="./families.txt"
species_tree="./species_tree.nw"
generax_out="./generax_OUT"

# Substitution model passed to GeneRax
model="JTT"

##################################################################################

## Filter MSAs: remove families with fewer than min_taxa species

echo "Filtering MSAs (min. ${min_taxa} species)"

for file in "$msa_path"/*.fa; do
    [ -e "$file" ] || continue
    n_taxa=$(grep ">" "$file" | awk -F'[>_]' '{print $2}' | sort -u | wc -l)
    if [ "$n_taxa" -lt "$min_taxa" ]; then
        echo "  Removing $file (${n_taxa} species)"
        rm "$file"
    fi
done

echo "Filtering complete."

##################################################################################

## Trim alignments with trimAl
# Replace stop codons (*) with gaps (-)
# Replace selenocysteine (U) with unknown amino acid (X) to avoid GeneRax errors
# Run trimAl with -automated1 

echo "Trimming alignments with trimAl"

mkdir -p "$trim_path"

export msa_path trim_path

ls "$msa_path" | parallel -j "$ncores" '
    echo "  Trimming: {}"
    sed -i "s/\*/-/g"          "${msa_path}/{}"
    sed -i "/^>/! s/U/X/g"     "${msa_path}/{}"
    trimal -in  "${msa_path}/{}" \
           -out "${trim_path}/{}" \
           -automated1
'

echo "Trimming complete"

##################################################################################

## Infer gene trees with FastTree

echo "Inferring gene trees with FastTreeMP"

mkdir -p "$tree_path"

fileN=$(ls "$trim_path" | grep -c ".fa")
i=0

for file in $(ls "$trim_path" | grep ".fa"); do
    ((i++))
    echo "  Processing file: $file  [${i}/${fileN}]"
    name=$(echo "$file" | cut -d"." -f1)
    FastTreeMP "${trim_path}/${file}" > "${tree_path}/${name}.tree"
done

echo "FastTreeMP complete"

##################################################################################

## Build the GeneRax families file

echo "Building GeneRax families file"

fileN=$(ls "$trim_path" | cut -f1 -d"." | sort -u | wc -l)
i=0

echo "[FAMILIES]" > "$out_file"

for name in $(ls "$trim_path" | cut -f1 -d"." | sort -u); do
    ((i++))
    echo "  Processing family: $name  [${i}/${fileN}]"

    echo "- ${name}"                        >> "$out_file"
    echo "starting_gene_tree = ${tree_path}/${name}.tree" >> "$out_file"
    echo "alignment = ${trim_path}/${name}.fa"            >> "$out_file"
    echo "subst_model = ${model}"           >> "$out_file"
done

echo "Families file written to: $out_file"

##################################################################################

## Run GeneRax (UndatedDL model)

echo "Running GeneRax"

mpiexec -oversubscribe -np "$ncores" generax \
    -f "$out_file"       \
    -s "$species_tree"   \
    -p "$generax_out"    \
    -r UndatedDL
 
echo "Pipeline complete. Results in: $generax_out"