fluid.clj

(ns com.jalat.fluid
  (:import (java.awt Color Graphics Dimension)
	   (java.awt.image BufferedImage)
	   (javax.swing JPanel JFrame)))



(def grid-size (int 100))

(defn make-grid []
  (make-array Float/TYPE (* (+ grid-size 2)
			    (+ grid-size 2))))
(def density-grid 
     (make-grid))
(def tmp-density-grid
     (make-grid))
(def density-source-grid 
     (make-grid))

(def x-vel-grid 
     (make-grid))
(def tmp-x-vel-grid
     (make-grid))
(def x-vel-source-grid 
     (make-grid))
(def y-vel-grid
     (make-grid))
(def tmp-y-vel-grid
     (make-grid))
(def y-vel-source-grid
     (make-grid))



(defn cell [#^Integer x #^Integer y]
  (+ x
     (* y
	(+ (int grid-size) (int 2)))))


(defn add-source [#^floats grid #^floats source-grid #^Float dt]
  (dotimes [n (alength grid)]
    (aset-float grid n (+ (aget grid n)
			  (* (aget source-grid n)
			     dt)))))

(defn diffuse [type #^floats grid #^floats diffused-grid diff-ratio dt]
  (let [a (float (* dt diff-ratio grid-size grid-size))
	a4-1 (float (+ 1 (* 4 a)))]
    (dotimes [n 20]
      (dotimes [x grid-size]
       (dotimes [y  grid-size]
	 (let [c (int (cell (inc x) (inc y)))]
	   (aset diffused-grid c
		 (float (/ (+ (float (aget grid c))
			      (float (*  a
					 (+ (+ (aget diffused-grid (int (dec c)))
					       (aget diffused-grid (int (inc c))))
					    (+ (aget diffused-grid (int (- c (+ grid-size (int 2)))))
					       (aget diffused-grid (int (+ c (+ grid-size (int 2))))))))))
			   a4-1)))))))
    (set-bnd type diffused-grid)
    (dotimes [n (alength grid)]
      (aset grid n (aget diffused-grid n)))))

(defn advect [type #^floats grid #^floats tmp-grid #^floats x-grid #^floats y-grid dt]
  (let [a (float (* grid-size dt))]
    (dotimes [x grid-size]
      (dotimes [y grid-size]
	(let [c (int (cell x y))
	      i (max (float 0.5)
		     (min (float (+ (float grid-size) (float 0.5)))
			  (- x (* a (aget x-grid c)))))
	      j (max (float 0.5)
		     (min (float (+ (float grid-size) (float 0.5)))
			  (- y (* a (aget y-grid c)))))
	      i0 (int i)
	      i1 (inc i0)
	      j0 (int j)
	      j1 (inc j0)
	      s1 (- i i0)
	      s0 (- 1 s1)
	      t1 (- j j0)
	      t0 (- 1 t1)]
	  (aset tmp-grid c
		(+ (* s0 
		      (+ (* t0 (aget grid (int (+ i0 
						  (* j0
						     (int 102))))))
			 (* t1 (aget grid (int (+ i0
						  (* j1 
						     (int 102))))))))
		   (* s1 
		      (+ (* t0 (aget grid (int (+ i1
						  (* j0
						     (int 102))))))
			 (* t1 (aget grid (int (+ i1 
						  (* j1
						     (int 102)))))))))))))
    (set-bnd type tmp-grid)
    (dotimes [n (alength grid)]
      (aset grid n (aget tmp-grid n)))))

(defn project [#^floats x-vel #^floats y-vel]
  (let [h (float (/ 1 grid-size))
	div #^floats (make-grid)
	p #^floats (make-grid)]
    (dotimes [x grid-size]
      (dotimes [y (int grid-size)]
	(let [c (int (cell (inc x)
			   (inc y)))]
	  (aset div c
		(float (* (float -0.5) 
			  h
			  (+ (- (aget x-vel (inc c))
				(aget x-vel (dec c)))
			     (- (aget y-vel (+ c
					       (+ (int 2)
						  (int grid-size))))
				(aget y-vel (- c 
					       (+ (int 2)
						  (int grid-size)))))))))
	  
	  (aset p c (float 0)))))
    (set-bnd 0 div)
    (set-bnd 0 p)
    (dotimes [n 20]
      (dotimes [x grid-size]
	(dotimes [y (int grid-size)]
	  (let [c (int (cell (inc x)
			  (inc y)))]
	    (aset p c
		  (/ (float (+ (+ (+ (aget div c)
				     (aget p (dec c)))
				  (+ (aget p (inc c)) 
				     (aget p (+ (int c)
						(+ (int 2)
						   (int grid-size))))))
			       (aget p (- (int c) 
					  (+ (int 2)
					     (int grid-size))))))
		     (float 4.0))))))
      (set-bnd 0 p))
    (dotimes [x grid-size]
      (dotimes [y (int grid-size)]
	(let [c (int (cell (inc x)
			   (inc y)))]
	  (aset x-vel c
		(- (aget x-vel c)
		   (float (/ (* (- (aget p (inc c))
				   (aget p (dec c)))
				(float 0.5))
			     h))))
	  (aset y-vel c
		(- (aget y-vel c)
		   (float (/ (* (- (aget p (+ c (int 102)))
				   (aget p (- c (int 102))))
				(float 0.5))
			     h)))))))
    (set-bnd 1 x-vel)
    (set-bnd 2 y-vel)))
  
(defn set-bnd [type grid]
  (dotimes [n grid-size]
    (aset grid (int (cell 0 n))
	  (if (= 1 type)
	    (- (aget grid (int (cell 1 n))))
	    (aget grid (int (cell 1 n)))))
    (aset grid (int (cell (+ grid-size 1) n))
	  (if (= 1 type)
	    (- (aget grid (int (cell (int grid-size) n))))
	    (aget grid (int (cell (int grid-size) n)))))    
    (aset grid (int (cell n 0))
	  (if (= 2 type)
	    (- (aget grid (int (cell n 1))))
	    (aget grid (int (cell n 1)))))
    (aset grid (int (cell n (+ (int grid-size) 1)))
	  (if (= 2 type)
	    (- (aget grid (int (cell n (int grid-size)))))
	    (aget grid (int (cell n (int grid-size))))))))

(defn print-grid [grid dim]
  (dotimes [x dim]
    (dotimes [y dim]
      (print (aget grid (cell x y)) " "))
    (println)))

(def scale 5)
(defn fill-cell [#^Graphics g x y c]
  (doto g
    (.setColor c)
    (.fillRect (* x scale) (* y scale) scale scale)))

(defn render-grid [#^Graphics g grid]
  (let [scale (int scale)
	img (new BufferedImage (* scale grid-size) (* scale grid-size)
		 (. BufferedImage TYPE_INT_ARGB))
	bg (. img (getGraphics))]
    (doto bg
      (.setColor (. Color black))
      (.fillRect 0 0 (. img (getWidth))
		 (. img getHeight)))
    (dotimes [x grid-size]
      (dotimes [y grid-size]
	(doto bg
	  (.setColor (Color. (int (min 255 
				       (max 0 (* 10000 (aget grid (+ x (* y (+ (int  2) (int grid-size)))))))))
			     (int 0) (int 0)))
	  (.fillRect (* x scale)
		     (* y scale)
		     scale scale))))
    (. g (drawImage img 0 0 nil))
    (. bg (dispose))))


(def panel (doto (proxy [JPanel] []
		   (paint [g] 
			  (render-grid g density-grid)))
	     (.setPreferredSize (new Dimension 
				     (* scale grid-size)
				     (* scale grid-size)))))

(def frame (doto (new JFrame)
	     (.add panel)
	     (.pack)
	     (.show)))

(def boring-grid
     (let [x (make-grid)]
       (dotimes [n (alength x)]
	 (aset x n (float 1)))
       x))

(def animator (agent nil))
(def running true)
(def animation-sleep-ms 100)
(defn animation [x]
  (when running
    (send-off *agent* #'animation))
					;density solver
  (add-source density-grid density-source-grid (float 0.001))
  (diffuse 0 density-grid tmp-density-grid (float 0.0025) (float 0.001))
  (advect 0 density-grid tmp-density-grid x-vel-grid y-vel-grid (float 0.001))
					;velocity evolver
  (add-source x-vel-grid x-vel-source-grid (float 0.001))
  (add-source y-vel-grid y-vel-source-grid (float 0.001))
  (diffuse 1 x-vel-grid tmp-x-vel-grid (float 0.00025) (float 0.001))
  (diffuse 2 y-vel-grid tmp-y-vel-grid (float 0.00025) (float 0.001))  
  (project x-vel-grid y-vel-grid )
  (advect 1 x-vel-grid tmp-x-vel-grid x-vel-grid y-vel-grid (float 0.001))
  (advect 2 y-vel-grid tmp-y-vel-grid x-vel-grid y-vel-grid (float 0.001))
  (project x-vel-grid y-vel-grid )

  (. panel (repaint))
  (. Thread (sleep animation-sleep-ms))
  nil)

(defn restart []
  (def running nil)
  (. Thread (sleep 5000))
  (dotimes [n (alength density-grid)]
    (aset density-grid n (float 0))
    (aset tmp-density-grid n (float 0))
    (aset density-source-grid n (float 0))
    (aset x-vel-grid n (float 0))
    (aset tmp-x-vel-grid n (float 0))
    (aset x-vel-source-grid n (float 0))
    (aset y-vel-grid n (float 0))
    (aset tmp-y-vel-grid n (float 0))
    (aset y-vel-source-grid n (float 0)))
  (def running true)

  (dotimes [n (alength x-vel-grid)]
    (aset y-vel-source-grid n (float -10))) ; general updraft

  (dotimes [n (alength x-vel-grid)]     ; little bit of sidewind to make it a little more interesting.
    (aset x-vel-source-grid n (float (* (float -0.5) (. Math sin (* 0.001 n))))))

  (dotimes [y 20] (dotimes [n 10]    ;stronger updraft around the "fire"
		    (aset y-vel-source-grid (cell (+ 45 n) (+ 75 y)) (float -100))))

  (aset density-source-grid (cell 50 90) (float 80)) ; smoke..

  (send-off animator animation))

(restart)