A copy of my WFC implementation, presented as-is

wave_function_collapse.clj

;; you should be able to add svg-clj as a git dep: https://github.com/adam-james-v/svg-clj
;; there are some comment blocks and 'defs' littered in the file where I explored a few ideas
;; this is not a well-architected file but more of a 'scratchpad' approach to messing around with WFC,
;; so please just use and read it knowing it's not efficient or well designed, just scrappy :) 

(ns adam.tools.wave-function-collapse
  (:require
   [clojure.set :as set]
   [clojure.string :as str]
   [svg-clj.elements :as el]
   [svg-clj.layout :as lo]
   [svg-clj.path :as path]
   [svg-clj.parametric :as p]
   [svg-clj.tools :as tools]
   [svg-clj.transforms :as tf]
   [svg-clj.utils :as utils]))

(def tile-size 30)

(defn tile-style
  [el]
  (-> el
      (tf/style {:fill "none"
                 :stroke "slategray"
                 :stroke-width 1})))

(def base-tile
  (-> (el/rect tile-size tile-size)
      (tf/style {:fill "none"})))

(defn corner
  []
  (let [shape (rand-nth [(el/polyline [[(* tile-size 0.5) (* tile-size 1.0)]
                                       [(* tile-size 0.5) (* tile-size 0.5)]
                                       [(* tile-size 1.0) (* tile-size 0.5)]])
                         (el/polyline [[(* tile-size 0.5) (* tile-size 1.0)]
                                       [(* tile-size 1.0) (* tile-size 0.5)]])
                         (path/bezier [(* tile-size 0.5) (* tile-size 1.0)]
                                      [(* tile-size 0.5) (* tile-size 0.5)]
                                      [(* tile-size 1.0) (* tile-size 0.5)])])]
    (el/g
      base-tile
      (-> shape
          (tf/translate [(* tile-size -0.5) (* tile-size -0.5)])
          tile-style))))

(defn side
  []
  (let [shape (rand-nth [(el/polyline [[(* tile-size 0.0) (* tile-size 0.5)]
                                       [(* tile-size 1.0) (* tile-size 0.5)]])
                         (path/bezier [(* tile-size 0.0) (* tile-size 0.5)]
                                      [(* tile-size 0.5) (* tile-size 0.125)]
                                      [(* tile-size 1.0) (* tile-size 0.5)])])]
    (el/g
      base-tile
      (-> shape
          (tf/translate [(* tile-size -0.5) (* tile-size -0.5)])
          tile-style))))

(defn inner
  []
  (el/g
    base-tile
    (-> (el/circle (* tile-size 0.375 (rand)))
        tile-style
        (tf/style {:opacity (rand)}))
    (-> (el/circle (* tile-size 0.375 (rand)))
        tile-style
        (tf/style {:opacity (rand)}))))

(defn render-tiles
  [tiles]
  (let [n-tiles (count tiles)
        ny (int (Math/sqrt n-tiles))
        nx (Math/ceil (/ n-tiles ny))
        grid (p/rect-grid nx ny tile-size tile-size)]
    (lo/distribute-on-pts (map #(%) tiles) grid)))

#_(defn rotate-module
  "Rotates a tile `deg` degrees counter-clockwise"
  [{:keys [sockets tile]} deg]
  (let [n (/ deg 90)]
    {:sockets (if (= n 0)
                sockets
                (vec (take 4 (drop n (cycle sockets)))))
     :tile (fn [] (mdl/rotate (tile) [0 0 (* -1 deg)]))}))

(defn rotate-module
  "Rotates a tile 90 degrees counter clockwise"
  ([tile-map] (rotate-module tile-map 90))
  ([{:keys [sockets tile]} deg]
   (let [n (/ deg 90)]
     {:sockets (if (= n 0)
                 sockets
                 (vec (take 4 (drop n (cycle sockets)))))
      :tile    (fn [] (tf/rotate (tile) (- deg)))})))

(defn module-rotations
  [module n]
  (when (#{2 4} n)
    (let [[b c d a] (take 4 (iterate rotate-module module))]
      (case n
        2 [a c]
        4 [a b c d]
        nil))))

;; sockets follow [:north :east :south :west]
(def basic-module-set
  (concat
    (module-rotations {:sockets [:a :b :b :a] :tile corner} 4)
    (module-rotations {:sockets [:a :b :a :b] :tile side} 4)
    [{:sockets [:a :a :a :a] :tile inner}]))

(def empty-module-set
  (concat
    (module-rotations {:sockets [0 1 1 0] :tile []} 4)
    (module-rotations {:sockets [0 1 0 1] :tile []} 4)
    [{:sockets [0 0 0 0] :tile []}]))

(defn neighbour-keys
  [[col row]]
  [[col (dec row)]
   [(inc col) row]
   [col (inc row)]
   [(dec col) row]])

(defn grid-rect
  ([w h] (grid-rect 0 0 w h))
  ([x y w h]
   (for [b (range y (+ h y))
         a (range x (+ w x))]
     [a b])))

(defn initial-grid
  [w h module-set]
  (zipmap (grid-rect w h) (repeat (vec module-set))))

(defn valid-sockets-for-module
  [{:keys [sockets]} module-set]
  (let [[n e s w] sockets
        all-sockets (set (map :sockets module-set))]
    [(set (filter #(= n (nth % 2)) all-sockets))
     (set (filter #(= e (nth % 3)) all-sockets))
     (set (filter #(= s (nth % 0)) all-sockets))
     (set (filter #(= w (nth % 1)) all-sockets))]))

(defn merge-socket-sets
  [socket-sets]
  (let [[ns es ss ws] (map (fn [n] (map #(get % n) socket-sets)) [0 1 2 3])]
    (mapv #(apply set/union %) [ns es ss ws])))

(defn valid-sockets-for-neighbours
  [pos gridmap module-set]
  (merge-socket-sets (map #(valid-sockets-for-module % module-set) (get gridmap pos))))

(defn lowest-entropy-pos
  [gridmap]
  (let [entropies (dissoc (group-by second (update-vals gridmap count)) 0 1)
        [le he]   (when (seq entropies)
                  (apply (juxt min max) (keys entropies)))
        choices   (cond
                  ;; all cells have the same entropy
                    (= 1 (count entropies))
                    (first (vals entropies))

                    ;; only one cell left un-collapsed
                    (and (= 2 (count entropies)) (= 1 (count (get entropies he))))
                    (get entropies he)

                    ;; cells either have full entropy or minimum possible entropy
                    (and (= 2 (count entropies)) (= 1 (count (get entropies 1))))
                    (get entropies he)

                    ;; some cells have greater than minimum possible entropy, pick those
                    :else (get entropies (apply min (keys (dissoc entropies 1 0)))))
        [k _] (rand-nth choices)]
    (when le
      (if (< le 1)
        (println (str "Entropy Too low -> Propagation Error perhaps? Entropy: " le))
        k))))

(defn collapse-one-at
  [pos gridmap]
  (update gridmap pos #(vector (rand-nth %))))

(defn update-neighbours
  [pos gridmap module-set]
  (let [neighbours     (neighbour-keys pos)
        socket-sets    (zipmap neighbours (valid-sockets-for-neighbours pos gridmap module-set))
        new-neighbours (for [k neighbours]
                         (when-let [modules (get gridmap k)]
                           (let [valid-modules (filter #((get socket-sets k) (:sockets %)) modules)]
                             [k (if (seq valid-modules)
                                  valid-modules
                                  (take 1 modules))])))]
    (into {} new-neighbours)))

(defn propagate
  [gridmap module-set]
  (let [seed (lowest-entropy-pos gridmap)]
    (loop [gm (collapse-one-at seed gridmap)
           stack [seed]]
      (if (< (count stack) 1)
        gm
        (let [[pos stack] ((juxt peek pop) stack)
              nks (neighbour-keys pos)
              old-neighbours (into {} (mapv #(vector % (get gm %)) nks))
              new-neighbours (update-neighbours pos gm module-set)
              new-stack (into stack (comp
                                     (filter #(and
                                               (pos? (count (get new-neighbours %)))
                                               (not=
                                                (count (get new-neighbours %))
                                                (count (get old-neighbours %)))))
                                     (map conj)) nks)]
          (recur (merge gm new-neighbours) new-stack))))))

(defn collapsed?
  [gridmap]
  (let [counts (set (map count (vals gridmap)))]
    (or (= #{1} counts)
        (= #{0 1} counts))))

(defn collapse
  [gridmap module-set]
  (if (collapsed? gridmap)
    gridmap
    #_(let [gridmap (-> gridmap
                      (propagate module-set))])
    (recur (propagate (update-vals gridmap shuffle) module-set) module-set)))

(defn generate
  [ncols nrows module-set]
  (collapse (initial-grid ncols nrows module-set) module-set))

(defn render-sockets
  [[n e s w]]
  (let [cols {0 "red"
              1 "green"
              2 "blue"
              :c "cyan"
              :o "red"}]
    (-> (el/g
         (when (= n :c)
           (-> (el/circle 2)
               (tf/translate [(* tile-size 0.45) (* tile-size 0.05)])
               (tf/style {:fill (get cols n)
                          :stroke "none"})))
         (when (= e :c)
           (-> (el/circle 2)
               (tf/translate [(* tile-size 0.95) (* tile-size 0.45)])
               (tf/style {:fill (get cols e)
                          :stroke "none"})))
         (when (= s :c)
           (-> (el/circle 2)
               (tf/translate [(* tile-size 0.45) (* tile-size 0.95)])
               (tf/style {:fill (get cols s)
                          :stroke "none"})))
         (when (= w :c)
           (-> (el/circle 2)
               (tf/translate [(* tile-size 0.05) (* tile-size 0.45)])
               (tf/style {:fill (get cols w)
                          :stroke "none"}))))
        (tf/translate [(* tile-size -0.5) (* tile-size -0.5)]))))

(defn render-gridmap
  [gridmap]
  (let [rfn (fn [[pos modules]]
              (let [{:keys [tile sockets]} (first modules)]
                (-> (el/g
                     (tile)
                     #_(-> (el/rect tile-size tile-size)
                         (tf/style {:fill "none"
                                    :stroke "gray"
                                    :stroke-width 0.375})))
                    #_(el/g (render-sockets sockets))
                    (tf/translate (utils/v* [tile-size tile-size] pos)))))]
    (el/g (pmap rfn gridmap))))




(comment
  (def basic-module-set
    (concat
     (map #(rotate-module {:sockets [0 1 1 0] :tile corner} %) [0 90 180 270])
     (map #(rotate-module {:sockets [0 1 0 1] :tile side} %) [#_0 90])
     (repeat 1 {:sockets [0 0 0 0] :tile inner})))


  (-> (generate 30 30 basic-module-set)
      render-gridmap
      tools/cider-show)

  ;; set some tiles yourself
  (let [modules basic-module-set
        gr (initial-grid 35 35 modules)
        gr2 (zipmap (grid-rect 5 15 25 15) (repeat [{:sockets [0 0 0 0] :tile (fn [] )}]))]
    (-> (merge gr gr2)
        (collapse modules)
        render-gridmap
        tools/cider-show))




  )



(defn side2
  []
  (let [pts   (mapv
               (fn [pt]
                 (mapv * pt (repeat tile-size)))
               [[0 0.5] [1 0.5]
                [1 1] [0 1]])]
    (-> (el/g
         (-> (el/polygon pts)
             (tf/style {:fill "slategray" :stroke "none"}))
         (-> (el/polyline (take 2 pts))
             (tf/style {:fill "none" :stroke "white"})))
        (tf/translate [(* tile-size -0.5) (* tile-size -0.5)]))))

(defn ground
  []
  (let [pts   (mapv
               (fn [pt]
                 (mapv * pt (repeat tile-size)))
               [[0 0.5] [1 0.5]
                [1 1] [0 1]])]
    (-> (el/g
         (-> (el/polygon pts)
             (tf/style {:fill "green" :stroke "none"}))
         (-> (el/polyline (take 2 pts))
             (tf/style {:fill "none" :stroke "white"})))
        (tf/translate [(* tile-size -0.5) (* tile-size -0.5)]))))

(defn underground
  []
  (let [pts   (mapv
               (fn [pt]
                 (mapv * pt (repeat tile-size)))
               [[0 0] [1 0]
                [1 1] [0 1]])]
    (-> (el/g
         (-> (el/polygon pts)
             (tf/style {:fill "brown" :stroke "none"})))
        (tf/translate [(* tile-size -0.5) (* tile-size -0.5)]))))

(defn inside
  []
  (let [pts   (mapv
               (fn [pt]
                 (mapv * pt (repeat tile-size)))
               [[0 0] [1 0]
                [1 1] [0 1]])]
    (-> (el/g
         (-> (el/polygon pts)
             (tf/style {:fill "slategray" :stroke "none"})))
        (tf/translate [(* tile-size -0.5) (* tile-size -0.5)]))))

(defn outside
  []
  #_(let [pts   (mapv
               (fn [pt]
                 (mapv * pt (repeat tile-size)))
               [[0 0] [1 0]
                [0 1] [1 1]])]
    (-> (el/g
         (-> (el/polygon pts)
             (tf/style {:fill "slategray" :stroke "none"}))
         (-> (el/polyline (take 2 pts))
             (tf/style {:fill "none" :stroke "white"})))
        (tf/translate [(* tile-size -0.5) (* tile-size -0.5)]))))

(defn corner2
  []
  (let [pts   (mapv
               (fn [pt]
                 (mapv * pt (repeat tile-size)))
               [[0.5 1] [0.5 0.5]
                [1 0.5] [1 1]])]
    (-> (el/g
         (-> (el/polygon pts)
             (tf/style {:fill "slategray" :stroke "none"}))
         (-> (el/polyline (take 3 pts))
             (tf/style {:fill "none" :stroke "white"})))
        (tf/translate [(* tile-size -0.5) (* tile-size -0.5)]))))

(defn corner3
  []
  (let [pts   (mapv
               (fn [pt]
                 (mapv * pt (repeat tile-size)))
               [[0 0]
                [0.5 0]
                [0.5 0.5]
                [1 0.5]
                [1 1]
                [0 1]])]
    (-> (el/g
         (-> (el/polygon pts)
             (tf/style {:fill "slategray" :stroke "none"}))
         (-> (el/polyline (take 3 (rest pts)))
             (tf/style {:fill "none" :stroke "white"})))
        (tf/translate [(* tile-size -0.5) (* tile-size -0.5)]))))

(defn side-ground1
  []
  (el/g
   (-> (side2) (tf/rotate -90))
   (ground)))

(defn side-ground2
  []
  (el/g
   (-> (side2) (tf/rotate 90))
   (ground)))

(defn inside-ground
  []
  (el/g
   (inside)
   (ground)))

(def building-module-set
  (concat
   (repeat 1 {:sockets [:o :og :g :og] :tile ground})
   (repeat 1 {:sockets [:g :g :g :g] :tile underground})
   (repeat 1 {:sockets [:o :o :o :o] :tile outside})
   (repeat 1 {:sockets [:oi :ig :g :og] :tile side-ground1})
   (repeat 1 {:sockets [:oi :og :g :ig] :tile side-ground2})
   (repeat 1 {:sockets [:i :ig :g :ig] :tile inside-ground})
   (repeat 1 {:sockets [:i :i :i :i] :tile inside})
   (apply concat (repeat 1 (map #(rotate-module {:sockets [:o :oi :oi :o] :tile corner2} %) [0 270])))
   (apply concat (repeat 1 (map #(rotate-module {:sockets [:oi :oi :i :i] :tile corner3} %) [0 90])))
   (apply concat (repeat 1 (map #(rotate-module {:sockets [:o :oi :i :oi] :tile side2} %) [0 90 270])))))

(comment
  (let [modules building-module-set
        [w h] [40 20]
        gr (initial-grid w h modules)]
    (-> (merge
         gr
         {[0 (- h 2)] (repeat 2 {:sockets [:i :ig :g :ig] :tile inside-ground})})
        (collapse modules)
        render-gridmap
        tools/cider-show))



  )