High Performance GeoSpatial Data Interpolation With H3

An example of doing high perf geo-spatial computations using Clojure and Uber's H3 Java bindings.

Recently did a project where I needed to do a poor-man's version of Inverse Distance Weighting for data points on a map powered by Uber's H3 Hexagonal Hierarchical Spatial Index.

At the beginning I attempted to do this all in the browser with the H3-JS bindings but computation was taking a VERY long time and when I tried to distribute the compute to WebWorkers, I learned the h3-js library isn't Worker compatible yet (uses document and in the code quite a bit). Luckily Uber has released quite a couple bindings for the native C compiled binary including ones for Java! Using these I was able to pretty easily write some Clojure around them and get a much needed performance boost.

Here be dragons: Do not take this snippet as an example of how to do good IDW, this snippet was a proof of concept on the viability of doing high-perf geo-spatial compute with Clojure using H3.

(ns org.providence.blend
  (:require [clojure.spec.alpha :as s]
            [clojure.spec.test.alpha :as stest]
            [clojure.test.check]
            [clojure.pprint]
            [clojure.edn :as edn]
            [clojure.java.io]
            [cheshire.core :as json])
  (:import [com.uber.h3core H3Core]))

#_(set! *warn-on-reflection* true)

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; Specs
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;


(s/def ::h3-id (s/and string? #(re-matches #"[a-z\d]{15}" %)))
(s/def ::value (s/and number?))
(s/def ::h3-value-map (s/map-of ::h3-id ::value))
(s/def ::pos-num (s/and number? pos?))

(s/fdef exp
  :args (s/cat :x ::pos-num
               :n ::pos-num)
  :ret ::pos-num)

(s/fdef get-distance
  :args (s/cat :origin ::h3-id
               :destination ::h3-id)
  :ret ::pos-num)

(s/fdef get-neighbours
  :args (s/cat :h3-id ::h3-id)
  :ret (s/coll-of ::h3-id))

(s/fdef get-value
  :args (s/cat :h3-id ::h3-id
               :root-id ::h3-id
               :root-value ::pos-num)
  :ret number?)

(s/fdef compact
  :args (s/cat :h3-value-map ::h3-value-map)
  :ret (s/coll-of map? :kind (s/and :value :hex-ids)))

(s/fdef krig
  :args (s/cat :root-id ::h3-id
               :root-value ::pos-num)
  :ret ::h3-value-map)


;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; Main
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;


(def ^com.uber.h3core.H3Core h3 (H3Core/newInstance))

(defn ^:private exp
  "Calculate `x` to the power of `n`
  @see https://stackoverflow.com/a/5058544"
  [x n]
  (loop [acc 1.0
         n   n]
    (if (zero? n)
      acc
      (recur (* x acc) (dec n)))))

(defn ^:private get-distance
  "Calculate the distance between the `origin` h3-index and the `destination` h3 index."
  [^String origin ^String destination]
  (let [origin-res      (.h3GetResolution h3 origin)
        destination-res (.h3GetResolution h3 origin)]
    (when (not= origin-res destination-res)
      (throw (ex-info "Resolution of origin and destination must be the same to calculate distance"
                      {:origin                 origin
                       :destination            destination
                       :origin-resolution      origin-res
                       :destination-resolution destination-res})))
    (.h3Distance h3 origin destination)))

(defn ^:private get-neighbours
  [^String h3-id]
  (->> (.kRing h3 h3-id 1)
       (remove (fn [self] (= h3-id self)))))

(defn ^:private get-value
  [h3-id root-id root-value]
  (if (= h3-id root-id)
    root-value
    (let [distance         (get-distance root-id h3-id)
          decay-constant   0.9
          calculated-value (* root-value (Math/pow decay-constant distance))]
      calculated-value)))

(defn krig
  [root-id root-value]
  (loop [h3-value-map (transient {})
         queue        #{root-id}]
    (if (zero? (count queue))
      ;; base case: return the h3 map
      (persistent! h3-value-map)
      ;; recurssive case: add the first in queue to the map, add neighbors to queue
      (let [next-id         (first queue)
            next-value      (get-value next-id root-id root-value)
            updated-map     (conj! h3-value-map {next-id next-value})
            krig-floor      0.001
            next-neighbours (if (< next-value krig-floor)
                              #{}
                              (->> (get-neighbours next-id)
                                   (remove (fn [neighbour] (updated-map neighbour)))))
            updated-queue   (into (disj queue next-id) next-neighbours)]
        (recur updated-map updated-queue)))))

(defn compact
  "Groups a flat map of {[h3-id: string]: number} by value."
  [h3-value-map]
  (loop [value-map h3-value-map
         compacted (transient {})]
    (if (zero? (count value-map))
      ;; base case - convert value-to-id-list map to list of value-to-id-list
      (let [compacted (persistent! compacted)]
        (loop [head (first compacted)
               tail (rest compacted)
               acc  (transient [])]
          (if (nil? head)
            ;; basecase - return the acc
            (persistent! acc)
            ;; recursive - conj new value/id-list entry into acc
            (recur (first tail) (rest tail) (conj! acc {:value   (first head)
                                                        :hex-ids (second head)})))))

      ;; recusive - take the first of the value-map and add it to the key list
      (let [[key val]               (first value-map)
            key-list-for-val        (or (compacted val) [])
            updated-h3-list-for-val (conj key-list-for-val key)]
        (recur (dissoc value-map key)
               (assoc! compacted val updated-h3-list-for-val))))))

(defn ^:private main
  [& args]
  {:pre [(s/valid? (s/coll-of string?) args)]}
  (let [[lat lng val]   (map edn/read-string args)
        are-all-numbers (not (some false? (map number? [lat lng val])))]
    (when (not are-all-numbers)
      (throw (ex-info "Provided values could not be parsed as numbers"
                      {:args args})))

    (let [hex-id              (.geoToH3Address h3 lat lng 7)
          krigged             (krig hex-id val)
          compacted           (compact krigged)
          tmp-dir             (System/getProperty "java.io.tmpdir")
          krigged-out         (clojure.java.io/file tmp-dir "providence" "krigged.json")
          krigged-grouped-out (clojure.java.io/file tmp-dir "providence" "krigged-grouped.json")
          krigged-grouped-edn (clojure.java.io/file tmp-dir "providence" "krigged-grouped.json.edn")]
      (clojure.java.io/make-parents krigged-out)
      (println "Writing to:" (-> krigged-out .getParentFile .getCanonicalPath))
      (spit krigged-out (json/generate-string krigged {:pretty true}))
      (println "Krigged:" (.getPath krigged-out))
      (spit krigged-grouped-out (json/generate-string compacted {:pretty true}))
      (println "Grouped:" (.getPath krigged-grouped-out))
      (spit krigged-grouped-edn (pr-str compacted) #_(with-out-str (clojure.pprint/pprint compacted)))
      (println "Grouped:" (.getPath krigged-grouped-edn))
      (println "Final map size:" (count krigged)))))

(defn -main
  "Sample usage: `clj -m org.providence.core 37.775938728915946 -122.41795063018799 0.8`"
  [& args]
  #_(s/check-asserts true)
  #_(stest/instrument)
  (dotimes [_ 20]
    (time (apply main args))))

#_(dotimes [_ 20] (time (main "37.775938728915946" "-122.41795063018799" "0.8")))