hg.stevelosh.com > dotfiles

(eval-when (:compile-toplevel :load-toplevel :execute)
  (ql:quickload '(:adopt :str :alexandria :losh :iterate :boots) :silent t))

(defpackage :sav
  (:use :cl :iterate :losh)
  (:export :toplevel :*ui*))

(in-package :sav)


;;;; Configuration ------------------------------------------------------------

(defparameter *score-match* 1)
(defparameter *score-mismatch* -1.5)
(defparameter *score-gap* -2)


;;;; Utilities ----------------------------------------------------------------
(defclass* seq ()
  (label bases (annotations :initform nil)))


(defun make-buffer (&key initial-contents
                         (element-type t)
                         (initial-capacity (max 64 (length initial-contents))))
  (let ((buffer (make-array initial-capacity
                  :element-type element-type
                  :adjustable t
                  :fill-pointer (length initial-contents))))
    (when initial-contents
      (replace buffer initial-contents))
    buffer))

(defun buffer-push (buffer element)
  (vector-push-extend element buffer)
  element)

(defun buffer-append (buffer sequence)
  (let* ((l1 (length buffer))
         (l2 (length sequence))
         (needed (+ l1 l2)))
    (when (< (array-total-size buffer) needed)
      (adjust-array buffer (max needed (* l1 2))))
    (setf (fill-pointer buffer) needed)
    (replace buffer sequence :start1 l1)
    sequence))

(defmacro-clause (BUFFERING expr &optional
                  APPEND (append nil)
                  INTO (var iterate::*result-var*)
                  INITIAL-CONTENTS (initial-contents '())
                  ELEMENT-TYPE (element-type t))
  `(progn
     (with ,var = (make-buffer :initial-contents ,initial-contents
                               :element-type ,element-type))
     (,(if append 'buffer-append 'buffer-push) ,var ,expr)))


(defun read-fasta (stream)
  (iterate
    (with label = nil)
    (case (peek-char nil stream nil :eof)
      (:eof (finish))
      (#\Newline (read-char stream))
      ((#\> #\@) (if label
             (finish)
             (setf label (subseq (read-line stream) 1))))
      (t (buffering (read-line stream) :into data
                    :append t
                    :element-type 'character)))
    (finally (return (make-instance 'seq :label label :bases data)))))

(defun read-fastq (stream)
  (let ((header-line (read-line stream))
        (sequence-line (read-line stream))
        (sep-line (read-line stream))
        (quality-line (read-line stream)))
    (declare (ignore sep-line quality-line))
    (make-instance 'seq
      :label (subseq header-line 1)
      :bases sequence-line)))

(defun read-bare (stream)
  (make-instance 'seq :label "Query" :bases (read-line stream)))

(defun read-fast* (stream)
  (ecase (peek-char nil stream nil :eof)
    (:eof nil)
    ((#\A #\C #\T #\G #\n #\a #\c #\t #\g #\n) (read-bare stream))
    (#\> (read-fasta stream))
    (#\@ (read-fastq stream))))

(defun parse-file (path)
  (with-open-file (s path)
    (iterate
      (for seq = (read-fast* s))
      (while seq)
      (collect seq))))

(defun read-data (path &key default-name)
  ;; todo handle -
  (declare (ignore default-name))
  (parse-file path))

(defun dna-reverse (seq)
  (reverse seq))

(defun dna-complement (seq)
  (map (type-of seq) (lambda (ch)
                       (ecase ch
                         (#\A #\T)
                         (#\T #\A)
                         (#\C #\G)
                         (#\G #\C)
                         (#\N #\N)))
       seq))

(defun dna-reverse-complement (seq)
  (dna-reverse (dna-complement seq)))


;;;; Alignment ----------------------------------------------------------------
(defun make-smith-waterman-grid (rows cols)
  (make-array (list rows cols) :initial-element 0))

(defun extract-smith-waterman-path (grid prev best-row best-col)
  (iterate (for coord
                :first (cons best-row best-col)
                :then (aref prev (car coord) (cdr coord)))
           (until (< (aref grid (car coord) (cdr coord)) 0.1f0))
           (collect (cons (1- (car coord)) (1- (cdr coord))) :at start)))

(defun print-grid (seq1 seq2 grid &key (pad 4))
  (flet ((print-padded (seq)
           (map nil (lambda (el) (format t "~V@A" pad el)) seq))
         (row-of (array row) ; ugly hack, but just for debugging…
           (iterate
             (for col :below (array-dimension array 1))
             (collect (aref array row col)))))
    (print-padded "  ")
    (print-padded seq1)
    (terpri)
    (iterate
      (for row :below (array-dimension grid 0))
      (print-padded (list (if (plusp row)
                            (aref seq2 (1- row))
                            #\Space)))
      (print-padded (row-of grid row))
      (terpri))))

(defun render-annotations (seq pad gaps anns)
  (if (null anns)
    (list (list) "")
    (let ((chars "123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz")
          (n 0)
          (legend (list)))
      (flet ((annotate (text)
               (let ((ch (aref chars n)))
                 (incf n)
                 (push (format nil "~A - ~A" ch text) legend)
                 ch)))
        (iterate
          (for i :from 0 :below (length seq))
          (collect (if-found (text (gethash i anns))
                     (annotate text)
                     #\space)
                   :into line :result-type 'string)
          (loop :while (eql i (car gaps))
                :do (progn (collect #\space :into line :result-type 'string)
                           (pop gaps)))
          (returning (list (nreverse legend)
                           (str:concat (str:repeat pad " ") line))))))))

(defclass* alignment ()
  (lab1 seq1 pad1 leg1 ann1
   m padm
   lab2 seq2 pad2 leg2 ann2))

(defun render-alignment (rseq qseq path)
  (iterate
    (with seq1 = (bases rseq))
    (with seq2 = (bases qseq))
    (with (lo2 . lo1) = (first path))
    (with (hi2 . hi1) = (car (last path)))
    (with gaps1 = (list))
    (with gaps2 = (list))
    (with pad1 = (if (>= lo1 lo2) 0 (- lo2 lo1)))
    (with pad2 = (if (>= lo2 lo1) 0 (- lo1 lo2)))
    (with padM = (max lo1 lo2))
    (with p1 = -1)
    (with p2 = -1)
    (for (i2 . i1) :in path)
    (for ch1 = (aref seq1 i1))
    (for ch2 = (aref seq2 i2))
    (collect (cond ((or (= i1 p1) (= i2 p2)) #\g)
                   ((char= ch1 ch2) #\|)
                   (t #\x))
             :into m :result-type 'string)
    (collect (if (> i1 p1)
               (progn (setf p1 i1)
                      ch1)
               (progn (push i1 gaps1)
                      #\-))
             :into l1 :result-type 'string)
    (collect (if (> i2 p2)
               (progn (setf p2 i2)
                      ch2)
               (progn (push i2 gaps2)
                      #\-))
             :into l2 :result-type 'string)
    (returning
      (destructuring-bind (leg1 ann1) (render-annotations seq1 pad1 gaps1 (annotations rseq))
        (destructuring-bind (leg2 ann2) (render-annotations seq2 pad2 gaps2 (annotations qseq))
          (make-instance 'alignment
            :lab1 (label rseq)
            :seq1 (str:concat (str:repeat pad1 " ")
                              (subseq seq1 0 lo1)
                              l1
                              (subseq seq1 (1+ hi1) nil))
            :pad1 pad1 :leg1 leg1 :ann1 ann1
            :m (str:concat (str:repeat padM " ")
                           m)
            :padm padm
            :lab2 (label qseq)
            :seq2 (str:concat (str:repeat pad2 " ")
                              (subseq seq2 0 lo2)
                              l2
                              (subseq seq2 (1+ hi2) nil))
            :pad2 pad2 :leg2 leg2 :ann2 ann2))))))


;; (defun print-alignment (seq1 seq2 path &optional ann1 ann2)
;;   (let ((a (render-alignment seq1 seq2 path ann1 ann2)))
;;     (map nil #'write-line (leg1 a))
;;     (when (leg1 a) (write-line ""))
;;     (write-line (ann1 a))
;;     (write-line (seq1 a))
;;     (write-line (m a))
;;     (write-line (seq2 a))
;;     (write-line (ann2 a))
;;     (when (leg2 a) (write-line ""))
;;     (map nil #'write-line (leg2 a))))

(defun smith-waterman (rseq qseq)
  (iterate
    (with seq1 = (coerce (bases rseq) 'simple-base-string))
    (with seq2 = (coerce (bases qseq) 'simple-base-string))
    (with best-row = 0)
    (with best-col = 0)
    (with best-score = 0)
    (with rows = (1+ (length seq2)))
    (with cols = (1+ (length seq1)))
    (with grid = (make-smith-waterman-grid rows cols))
    (with prev = (make-array (list rows cols) :initial-element nil))
    (for row :from 1 :below rows)
    (for ch2 :in-string seq2)
    (iterate
      (for col :from 1 :below cols)
      (for ch1 :in-string seq1)
      (for row-1 = (1- row))
      (for col-1 = (1- col))
      (for match = (+ (aref grid row-1 col-1)
                      (if (char= ch1 ch2) *score-match* *score-mismatch*)))
      (for gap-1 = (+ (aref grid row col-1) *score-gap*))
      (for gap-2 = (+ (aref grid row-1 col) *score-gap*))
      (for best = (max match gap-1 gap-2 0))
      (setf (aref grid row col) best
            (aref prev row col) (cond ((= best 0) nil)
                                      ((= best match) (cons row-1 col-1))
                                      ((= best gap-1) (cons row col-1))
                                      ((= best gap-2) (cons row-1 col))))
      (if (> best best-score)
        (setf best-score best
              best-row row
              best-col col)))
    (returning
      (let ((path (extract-smith-waterman-path grid prev best-row best-col)))
        (render-alignment rseq qseq path)))))


;;;; UI -----------------------------------------------------------------------
(defparameter *selected-index* 0)
(defparameter *r-seqs* nil)
(defparameter *q-seq* nil)
(defparameter *x-offset* 5)

(defparameter *alignments* nil)

(defconstant +attr/match+    (boots:attr :bold t :fg (boots:rgb 50 255 50)))
(defconstant +attr/mismatch+ (boots:attr :bold t :fg (boots:rgb 255 50 50)))
(defconstant +attr/not-selected+ (boots:attr :bold t))
(defconstant +attr/selected+ (boots:attr :bold t :fg (boots:rgb 0 0 0) :bg (boots:rgb 255 255 255)))

(defun draw/status (pad)
  (boots:draw pad 0 0 "SAV v1"))

(defun draw/alignment (index pad &aux (alignment (nth index *alignments*)))
  (boots:draw pad 0 0 (format nil "~A vs ~A" (lab1 alignment) (lab2 alignment))
              (if (= index *selected-index*) +attr/selected+ +attr/not-selected+))
  (boots:draw pad *x-offset* 2 (seq1 alignment))
  (boots:draw pad *x-offset* 4 (seq2 alignment))
  (loop :for i :from 0
        :for ch :across (m alignment)
        :do (multiple-value-bind (c a) (case ch
                                         (#\| (values #\| +attr/match+))
                                         (#\g (values #\space +attr/mismatch+))
                                         (#\x (values #\x +attr/mismatch+))
                                         (#\space (values nil nil)))
              (when c
                (boots:draw pad (+ *x-offset* i) 2 (aref (seq1 alignment) i) a)
                (boots:draw pad (+ *x-offset* i) 3 c a)
                (boots:draw pad (+ *x-offset* i) 4 (aref (seq2 alignment) i) a)))))

(defun press-any-key ()
  (boots:read-event)
  nil)

(defun recompute-alignments ()
  (setf *alignments* (loop :for r-seq :in *r-seqs*
                           :collect (smith-waterman r-seq *q-seq*))))

(defun incf-selected (delta)
  (zapf *selected-index* (mod (+ % delta) (length *alignments*))))

(defun swap-alignments (a b)
  (let ((ai (mod a (length *alignments*)))
        (bi (mod b (length *alignments*))))
    (rotatef (nth ai *alignments*)
             (nth bi *alignments*))
    (setf *selected-index* bi)))

(defun ui/main ()
  (setf (boots:root boots:*screen*)
        (boots:make-stack
          :children (list*
                      (boots:make-canvas :height 1 :draw #'draw/status)
                      (loop :for a :in *alignments*
                            :for i :from 0
                            :collect (boots:make-canvas
                                       :border-top t :fill-char #\space
                                       :draw (curry #'draw/alignment i))))))
  (loop
    (boots:redraw)
    (boots:event-case (boots:read-event)
      ((#\k :up)   (incf-selected -1))
      ((#\j :down) (incf-selected  1))
      (#\K (swap-alignments *selected-index* (1- *selected-index*)))
      (#\J (swap-alignments *selected-index* (1+ *selected-index*)))
      ((#\h :left)  (incf *x-offset* 10))
      ((#\l :right) (decf *x-offset* 10))
      ((#\escape #\q) (return)))))


;;;; Entry Point --------------------------------------------------------------
(defun run% (r-seqs q-seq)
  (boots/terminals/ansi:with-ansi-terminal (terminal :truecolor t)
    (boots:with-screen (boots:*screen* terminal)
      (boots:with-light-borders
        (let* ((*r-seqs* r-seqs)
               (*q-seq* q-seq))
          (recompute-alignments)
          (ui/main))))))

(defun run (reference-path query-path)
  (let ((r-seqs (read-data reference-path :default-name "Reference"))
        (q-seqs (read-data query-path :default-name "Query")))
    (run% r-seqs (first q-seqs))))


(defun r ()
  (run "/home/sjl/scratch/sav-foo.fasta" "/home/sjl/scratch/sav-bar.fasta"))


;;;; User Interface -----------------------------------------------------------
(defparameter *option/help*
  (adopt:make-option 'help
    :help "Display help and exit."
    :long "help"
    :short #\h
    :reduce (constantly t)))

(defparameter *option/version*
  (adopt:make-option 'version
    :help "Display version information and exit."
    :long "version"
    :reduce (constantly t)))

(adopt:define-string *help-text*
  "sav is a Sequence Alignment Visualizer")

(defparameter *examples*
  '(("Align the clipboard to the L1HS consensus sequence:"
     . "sav path/to/l1.fa <(pbpaste)")))

(defparameter *ui*
  (adopt:make-interface
    :name "sav"
    :usage "[OPTIONS] REF QUERY"
    :summary "sequence alignment visualizer"
    :help *help-text*
    :examples *examples*
    :contents (list
                *option/help*
                *option/version*)))


(defun toplevel ()
  (adopt::quit-on-ctrl-c ()
    (multiple-value-bind (arguments options) (adopt:parse-options-or-exit *ui*)
      (when (gethash 'help options)
        (adopt:print-help-and-exit *ui*))
      (let () ; todo options
        (handler-case
            (progn
              (assert (= (length arguments) 2) () "REF and QUERY arguments are both required.")
              (run (first arguments) (second arguments)))
          (error (e) (adopt:print-error-and-exit e)))))))


#; Scratch --------------------------------------------------------------------

;; (run "~/scratch/sav-foo.fasta" "~/scratch/sav-bar.fasta"
;;      (alexandria:alist-hash-table '())
;;      (alexandria:alist-hash-table '()))
author	Steve Losh <steve@stevelosh.com>
date	Tue, 13 Oct 2026 10:31:54 -0400
parents	(none)
children	(none)