extract_distributions.awk

#!/usr/bin/env -S gawk -f
#
# SPDX-FileCopyrightText: 2024, 2025 Marcel Schilling <mschilling@ub.edu>
#
# SPDX-License-Identifier: AGPL-3.0-or-later

# RNA-seq data statistic extraction Awk script for `scr4eam` scRNA-seq
# read simulation pipeline.
#
# Copyright (C) 2024, 2025  Marcel Schilling
#
# This file is part of `scr4eam`.
#
# `scr4eam` is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as
# published by the Free Software Foundation, either version 3 of the
# License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.


#######################
# General information #
#######################

# File:     extract_distributions.awk
# Created:  2024-12-12
# Modified: 2025-07-11
# Author:   Marcel Schilling <@mschilling@ub.edu>
# License:  GNU Affero General Public License Version >= 3.0 (GNU AGPL v3+)
# Purpose:  Extract distributions required for read generation using `scr4eam`
#           from read and fragment mapping summary TSV generated by SCALPEL.


#####################################
# Changelog (reverse chronological) #
#####################################

# 2025-07-11: Email: IDIBELL -> UB.
# 2024-12-12: Initial working version.


#########
# Usage #
#########

# ```sh
# [gawk --file=]extract_distributions.awk \
#   [--assign=multi_isoform_genes_txt=<multi-isoform-genes-txt>] \
#   [--assign=read_dists_anchors_skipped_txt=<read-dists-anchors-skipped-txt>] \
#   [--assign=n_reads_per_fragment_txt=<n-reads-per-fragment-txt>] \
#   [--assign=read_distance_col=<read-distance-column>] \
#   [--assign=gene_col=<gene-column>] \
#   [--assign=isoform_col=<isoform-column>] \
#   [--assign=fragment_col=<fragment-column>] \
#   [--assign=read_col=<read-column>] \
#   [--assign=fragment_distance_col=<fragment-distance-column>] \
#   [< <fragments-and-reads-stats-tsv>] \
#   [> <unambiguous-fragment-dists-per-gene-tsv>]
# ```
#
# Command line arguments:
# * `<fragments-and-reads-stats-tsv>`: TSV file/stream to read read and
#                                      fragment mapping summary data (e.g. as
#                                      generated by [SCALPEL](scalpel)) from;
#                                      default: `STDIN`
# * `<unambiguous-fragment-dists-per-gene-tsv>`: TSV file/stream to write the
#                                                assigned gene and the
#                                                fragment's distance (0-based)
#                                                to the transcript's (3') end
#                                                for each unambiguously
#                                                assigned fragment to; default:
#                                                `STDOUT`
# * `multi_isoform_genes_txt`: Text file to write genes expressing more than
#                              just a single isoform to (one gene per line);
#                              default: `genes.multiple_expressed_isoforms.txt`
# * `read_dists_anchors_skipped_txt`: Text file to write distances (0-based) to
#                                     the fragment's 3' end for all non-anchor
#                                     reads (i.e. excluding the first
#                                     zero-distance read of each fragment) to
#                                     (one integer per line); default:
#                                     `read_dists.anchors_skipped.txt`
# * `n_reads_per_fragment_txt`: File name/path to write total read counts
#                               for each fragment to; default:
#                               `n_reads.per_fragment.txt`
# * `read_distance_col`: Column index (1-based) in the input TSV containing
#                        the read's distance (0-based) from the
#                        transcript's (3') end; default: `9`
# * `gene_col`: Column index (1-based) in the input TSV specifying the gene
#               the read was assigned to; default: `11`
# * `isoform_col`: Column index (1-based) in the input TSV specifying the
#                  isoform the read was assigned to; default: `12`
# * `fragment_col`: Column index (1-based) in the input TSV specifying the
#                   fragment the read was assigned to; default: `13`
# * `read_col`: Column index (1-based) in the input TSV identifying the
#               read the input line belongs to; default: `14`
# * `fragment_distance_col`: Column index (1-based) in the input TSV
#                            containing the fragment's distance (0-based)
#                            from the transcript's (3') end; default: `15`


##############
# Parameters #
##############

BEGIN {
  # Default output file names:
  if(!multi_isoform_genes_txt) {
    multi_isoform_genes_txt = "genes.multiple_expressed_isoforms.txt"
  }
  if(!read_dists_anchors_skipped_txt) {
    read_dists_anchors_skipped_txt = "read_dists.anchors_skipped.txt"
  }
  if(!n_reads_per_fragment_txt) {
    n_reads_per_fragment_txt = "n_reads.per_fragment.txt"
  }

  # Default input column indices:
  if(!read_distance_col) { read_distance_col = 9 }
  if(!gene_col) { gene_col = 11 }
  if(!isoform_col) { isoform_col = 12 }
  if(!fragment_col) { fragment_col = 13 }
  if(!read_col) { read_col = 14 }
  if(!fragment_distance_col) { fragment_distance_col = 15 }


#####################
# In-/Output format #
#####################

  # Read & write TSV.
  OFS = FS = "\t"


#############
# Constants #
#############

  # Boolean helpers for readabilty:
  false = 0
  true = !false
}


##########################
# Input parsing (i): IDs #
##########################

# Parse isoform & fragment/read IDs.
{
  next_isoform = $isoform_col
  next_fragment = $fragment_col
  next_read = $read_col
}


##################
# Initialization #
##################

# Initialise first read/fragment.
NR == 1 {
  fragment = next_fragment
  fragment_isoform = isoform = next_isoform
  read = next_read
  read_distance = $read_distance_col
  fragment_distance = $fragment_distance_col
}


############################
# Read/fragment processing #
############################

# Process reads after gathering all related data.
# Note: The same read can follow itself due to splicing or because it is the
# only read assigned to a multi-isoform gene. Only in the latter case the read
# should be considered processed as we are moving fragment by fragment, isoform
# by isoform, and read by read.
isoform != next_isoform || read != next_read { process_read() }

# Process fragments after gathering all related data.
fragment != next_fragment { process_fragment() }


############################
# Input parsing (ii): Data #
############################

# Parse new read/fragment data.
{
  read_distance = $read_distance_col
  gene = $gene_col
  fragment_distance = $fragment_distance_col
}


##################################
# Multi-isoform genes accounting #
##################################

# Keep track of genes expressing more than one isoform to exclude from fragment
# distance sampling to avoid bias towards shorter fragments.
!(gene in reported_multi_iso_genes) {
  if (gene in gene2isoform) {
    if (isoform != gene2isoform[gene]) {
      print gene > multi_isoform_genes_txt
      reported_multi_iso_genes[gene]
    }
  } else {
    gene2isoform[gene] = isoform
  }
}


########################
# Finalization/Cleanup #
########################

# Don't forget to process the final lines as well.
END {
  process_read()
  process_fragment()
}


#############
# Functions #
#############

# Define what to do when a read's data has been read.
function process_read() {

  # Only consider read for the first isoform (all read-level data should be
  # identical) to avoid double-counting.
  if (isoform == fragment_isoform) {

    # Keep track of read counts per fragment.
    n_reads++

    # Keep track of the read distances within the fragment, but skip one read
    # anchoring the fragment itself as its position won't require sampling.
    # Note: If several reads share the same position, zero-distance reads could
    # still be recorded and sampled, as intended.
    if (anchor_read_skipped || read_distance != fragment_distance) {
      print read_distance - fragment_distance > read_dists_anchors_skipped_txt
    } else { # !anchors_skipped && read_distance == fragment_distance
      anchor_read_skipped = true
    }
  }
  read = next_read
  if(fragment == next_fragment) { isoform = next_isoform }
}

# Define what to do when a fragment's data has been read.
function process_fragment() {

  # Keep track of distances of unambiguously assigned fragments.
  if (isoform == fragment_isoform) {
    print gene, fragment_distance
  }

  # Keep track of read counts per fragment.
  print n_reads > n_reads_per_fragment_txt

  # Initialise next fragment.
  fragment = next_fragment
  fragment_isoform = isoform = next_isoform
  n_reads = anchor_read_skipped = false
}