gillespie.clj

(ns mca.gillespie
  (:require [clojure.contrib.generic.math-functions :as math]))

(defn choice [m]
  (let [r (rand (apply + (vals m)))]
    (loop [[[k w] & kws] (seq m)
           sum 0]
      (if (or (nil? kws) (< r (+ sum w))) k
          (recur kws (+ sum w))))))

(defn rxn-commands [rxn]
  (:cmds rxn))

(defn execute-command [mixture [cmd n specie]]
  (update-in mixture [specie] (partial cmd n)))

(defn apply-rxn [rxn mixture]
  (reduce execute-command mixture (rxn-commands rxn)))

(defn get-activities [state]
  (let [{:keys [mixture rxns]} state
        as (for [{:keys [lhs rate]} rxns]
             (apply * rate (for [[mol freq] lhs]
                             (* (mixture mol) freq))))]
    (zipmap rxns as)))

(defn step [state]
  (let [activities (get-activities state)
        rxn (choice activities)
        total-activity (apply + (vals activities))
        dt (/ (math/log (/ 1 (rand))) total-activity)]
    (-> state
        (update-in [:time] (partial + dt))
        (update-in [:num-steps] inc)
        (update-in [:mixture] (partial apply-rxn rxn))
        (update-in [:executed-rxn] rxn))))

(defn compile-rxn [[name rxn rate]]
  (let [[lhs arrow rhs] (partition-by #{"<->" "->"} rxn)
        [lhs-counts rhs-counts] (map frequencies [lhs rhs])
        mols (distinct (concat lhs rhs))
        cmds (for [mol mols
                   :let [diff (- (or (rhs-counts mol) 0)
                                 (or (lhs-counts mol) 0))
                         op (cond
                              (> diff 0) +
                              (< diff 0) -)]
                   :when op]
               [op (math/abs diff) mol])]
    {:name name :lhs lhs-counts :cmds cmds :rate (val rxn)}))

(defn simulate* [rxns init & {:keys [num-steps time until perturbations]}]
  (let [state {:time 0 :num-steps 0 :mixture init :rxns (mapcat compile-rxn rxns)}
        sim (iterate (fn [{rxns :rxns :as state}]
                       (reduce #(%2 %1) (step state) perturbations)) state)]
    (cond
      num-steps (take num-steps sim)
      time (take-while #(= (:time %) time) sim)
      until (take-while until sim)
      :else sim)))

(defn replace-keys [f m]
  (into {} (for [[k v] m] [(f k) v])))

(defn transform-rxn [rxn]
  (:rxn (reduce (fn [state a]
                  (if (number? a)
                    (assoc state :n a)
                    {:n 1 :rxn (concat (:rxn state) (repeat (:n state) (str a)))}))
                {:n 1 :rxn []} rxn)))

(defn or-comb [& preds]
  (fn [x] (some #(% x) preds)))

;; TODO make regular expressions for sequences... it'll make recognizing
;;      patterns in macros much more easy
(defn encode-rxn-set [rxn-set-spec]
  (let [splitted (take-nth 2 (partition-by #{'at} rxn-set-spec))
        rates (map (fn [[rxn1 rxn2]]
                     (cond
                       (some #{ '->} rxn1) (take 1 rxn2)
                       (some #{'<->} rxn1) (take 2 rxn2)))
                   (partition 2 1 splitted))
        rxns (cons (first splitted)
                   (drop-last (map #(drop (count %1) %2) rates (rest splitted))))
        names (map (fn [[name]] (if (string? name) name "")) rxns)
        rxns-seq (map (comp transform-rxn (partial remove (or-comb string? #{'* '+}))) rxns)]
    (map vector names rxns-seq rates)))

(defn make-perturbation [pred mod]
  (fn [state]
    (if (pred state) (mod state) state)))

(defn make-modification [[op n-or-rxn specie-or-n]]
  (case op
    'add #(update-in % [:mixture (str specie-or-n)] (partial + n-or-rxn))
    'del #(update-in % [:mixture (str specie-or-n)] (partial - n-or-rxn))
    'change-rate #(assoc-in % [:rxns n-or-rxn :rate] specie-or-n)
    'inc-rate #(update-in % [:rxns n-or-rxn :rate] (partial + specie-or-n))
    'dec-rate #(update-in % [:rxns n-or-rxn :rate] (partial - specie-or-n))))

(defn get-qty-if-molecule [x state]
  (let [{:keys [mixture]} state]
    (if (symbol? x) (mixture (str x)) x)))

(defn make-pred [op pred-part]
  (case op
    'at (let [[wait-amount wait-type] pred-part]
          #(= wait-amount ((case wait-type
                             'time-units :time
                             'events :num-steps) %)))
    'every (let [[wait-amount wait-type] pred-part]
             #(zero? (mod ((case wait-type
                             'time-units :time
                             'events :num-steps) %) wait-amount)))
    'when (cond
            (= (count pred-part) 3) (let [[a op b] pred-part]
                                      #(op (get-qty-if-molecule a %)
                                           (get-qty-if-molecule b %)))
            (= (last pred-part) 'executes) #(= (-> % :executed-rxn :name)
                                               (second pred-part)))))

(defn parse-perturbations [perturbations]
  (for [p perturbations
        :let [mod-part (take-last 3 p)
              pred-part (drop-last 3 (rest p))]]
    (make-perturbation (make-pred (first p) pred-part)
                       (make-modification mod-part))))

(defn parse-opts [opts]
  (let [groups (map (partial apply concat) (partition 2 (partition-by keyword? opts)))
        [[_ & perturbations]] (filter (comp #{:perturbations} first) groups)]
    (into {} (cons [:perturbations (parse-perturbations perturbations)]
                   (map vec (remove (comp #{:perturbations} first) groups))))))

(defmacro simulate [rxn-set-spec init & opts]
  `(simulate* ~(encode-rxn-set rxn-set-spec) ~(replace-keys str init) ~(parse-opts opts)))


(comment
(simulate ["r1" X0 + E1 <-> X0-E1 at 1, 1
           "r2" X0-E1 <-> X1-E1 at 1, 1
           "r3" X1-E1 <-> X1 + E1 at 1, 1]
          {X0 1000
           E1 20}
          :num-steps 1000 ;; or :time t, or :until pred (over the state)
          :perturbations
          (at 10 time-units add 1 X0)
          (every 20 events del 1 X1)
          (when X0 = 990 change-rate "r1" 10)
          (when "r2" executes inc-rate "r3" 1)))