emmy.clj

;; # Exercise 1.44: The double pendulum

;; This namespace explores [Exercise
;; 1.44](https://tgvaughan.github.io/sicm/chapter001.html#Exe_1-44) from Sussman
;; and Wisdom's [Structure and Interpretation of Classical
;; Mechanics](https://tgvaughan.github.io/sicm/), using
;; the [Emmy](https://github.com/mentat-collective/emmy) Clojure library and
;; the Clerk rendering environment.

#_{:clj-kondo/ignore [:refer-all]}
(ns emmy
  (:refer-clojure
   :exclude [+ - * / partial ref zero? numerator denominator compare = run!
             abs infinite?])
  (:require [nextjournal.clerk :as clerk]
            [emmy.env :as e :refer :all]
            [emmy.expression.compile :as xc]
            [emmy.expression.render :as xr]))

;; ## Lagrangian
;;
;; Start with a coordinate transformation from `theta1`, `theta2` to rectangular
;; coordinates. We'll generate our Lagrangian by composing this with an rectangular
;; Lagrangian with the familiar form of `T - V`.

(defn angles->rect [l1 l2]
  (fn [[_ [theta1 theta2]]]
    (let [x1 (* l1 (sin theta1))
          y1 (- (* l1 (cos theta1)))
          x2 (+ x1 (* l2 (sin (+ theta1 theta2))))
          y2 (- y1 (* l2 (cos (+ theta1 theta2))))]
      (up x1 y1 x2 y2))))

;; `T` describes the sum of the kinetic energy of two particles in rectangular
;; coordinates.

(defn T [m1 m2]
  (fn [[_ _ [xdot1 ydot1 xdot2 ydot2]]]
    (+ (* 1/2 m1 (+ (square xdot1)
                    (square ydot1)))
       (* 1/2 m2 (+ (square xdot2)
                    (square ydot2))))))

;; `V` describes a uniform gravitational potential with coefficient `g`, acting
;; on two particles with masses of, respectively, `m1` and `m2`. Again, this is
;; written in rectangular coordinates.

(defn V [m1 m2 g]
  (fn [[_ [_ y1 _ y2]]]
    (+ (* m1 g y1)
       (* m2 g y2))))

;; Form the rectangular Lagrangian `L` by subtracting `(V m1 m2 g)` from `(T m1 m2)`:

(defn L-rect [m1 m2 g]
  (- (T m1 m2)
     (V m1 m2 g)))

;; Form the final Langrangian in generalized coordinates (the angles of each
;; segment) by composing `L-rect` with a properly transformed `angles->rect`
;; coordinate transform!

(defn L-double-pendulum [m1 m2 l1 l2 g]
  (compose (L-rect m1 m2 g)
           (F->C
            (angles->rect l1 l2))))

;; The Lagrangian is big and hairy:

(def symbolic-L
  ((L-double-pendulum 'm_1 'm_2 'l_1 'l_2 'g)
   (up 't
       (up 'theta_1 'theta_2)
       (up 'theta_1dot 'theta_2dot))))

;; Let's simplify that:

(simplify symbolic-L)

;; Better yet, let's render it as LaTeX, and create a helper function,
;; `render-eq` to make it easier to render simplified equations:

(def render-eq
  (comp clerk/tex ->TeX simplify))

(render-eq symbolic-L)

;; And here are the equations of motion for the system:

;; TODO this currently causes a notebook failure.
#_
(let [L (L-double-pendulum 'm_1 'm_2 'l_1 'l_2 'g)]
  (binding [xr/*TeX-vertical-down-tuples* true]
    (render-eq
     (((Lagrange-equations L)
       (up (literal-function 'theta_1)
           (literal-function 'theta_2)))
      't))))

;; What do these mean?
;;
;; - the system has two degrees of freedom: $\theta_1$ and $\theta_2$.
;; - at any point `t` in time, the two equations above, full of first and second
;; - order derivatives of the position functions, will stay true
;; - the system can use these equations to simulate the system, one tick at a time.

;; ## Simulation
;;
;; Next, let's run a simulation using those equations of motion and collect data
;; on each coordinate's evolution.
;;
;; Here are the constants specified in exercise 1.44:
;;
;; masses in kg:

(def m1 1.0)
(def m2 3.0)

;; lengths in meters:

(def l1 1.0)
(def l2 0.9)

;; `g` in units of m/s^2:

(def g 9.8)

;; And two sets of initial pairs of `theta1`, `theta2` angles corresponding to
;; chaotic and regular initial conditions:

(def chaotic-initial-q (up (/ Math/PI 2) Math/PI))
(def regular-initial-q (up (/ Math/PI 2) 0.0))

;; Composing `Lagrangian->state-derivative` with `L-double-pendulum` produces
;; a state derivative that we can use with our ODE solver:

(def state-derivative
  (compose
   Lagrangian->state-derivative
   L-double-pendulum))

;; Finally, two default parameters for our simulation. We'll record data in
;; steps of 0.01 seconds, and simulate to a horizon of 50 seconds.

(def step 0.01)
(def horizon 50)

;; `run!` will return a sequence of 5001 states, one for each measured point in
;; the simulation. The smaller-arity version simply passes in default masses and
;; lengths, but you can override those with the larger arity version if you like.

;; (The interface here could use some work: `integrate-state-derivative` tidies
;; this up a bit, but I want it out in the open for now.)

(defn run!
  ([step horizon initial-coords]
   (run! step horizon l1 l2 m1 m2 g initial-coords))
  ([step horizon l1 l2 m1 m2 g initial-coords]
   (let [collector     (atom (transient []))
         initial-state (up 0.0
                           initial-coords
                           (up 0.0 0.0))]
     ((evolve state-derivative m1 m2 l1 l2 g)
      initial-state
      step
      horizon
      {:compile? true
       :epsilon 1.0e-13
       :observe (fn [_ state]
                  (swap!
                   collector conj! state))})
     (persistent! @collector))))

;; Run the simulation for each set of initial conditions and show the final
;; state. Chaotic first:

(def raw-chaotic-data
  (time
   (run! step horizon chaotic-initial-q)))

;; Looks good:

(peek raw-chaotic-data)

;; Next, the regular initial condition:

(def raw-regular-data
  (time
   (run! step horizon regular-initial-q)))

;; Peek at the final state:

(peek raw-regular-data)


;; ## Measurements, Data Transformation

;; Next we'll chart the measurements trapped in those sequences of state tuples.

;; The exercise asks us to graph the energy of the system as a function of time.
;; Composing `Lagrangian->energy` with `L-double-pendulum` gives us a new
;; function (of a state tuple!) that will return the current energy in the
;; system.:

(def L-energy
  (compose
   Lagrangian->energy
   L-double-pendulum))

;; `energy-monitor` returns a function of `state` that stores an initial energy
;; value in its closure, and returns the delta of each new state's energy as
;; compared to the original.

(defn energy-monitor [energy-fn initial-state]
  (let [initial-energy (energy-fn initial-state)]
    (fn [state]
      (- (energy-fn state) initial-energy))))

;; Finally, the large `transform-data` function. The charting library we'll use
;; likes Clojure dictionaries; `transform-data` turns our raw data into a
;; sequence of dictionary with all values we might care to explore. This
;; includes:

;; - the generalized coordinate angles
;; - the generalized velocities of those angles
;; - the rectangular coordinates of each pendulum bob
;; - `:d-energy`, the error in energy at each point in the simulation

;; Here is `transform-data`:

#_{:clj-kondo/ignore [:unresolved-symbol]}
(defn transform-data [xs]
  (let [energy-fn (L-energy m1 m2 l1 l2 g)
        monitor   (xc/compile-state-fn
                   (energy-monitor energy-fn (first xs))
                   false
                   (first xs)
                   {:calling-convention :structure})
        xform     (xc/compile-state-fn
                   (angles->rect l1 l2)
                   false
                   (first xs)
                   {:calling-convention :structure})
        pv        (principal-value Math/PI)]
    (map (fn [[t [theta1 theta2] [thetadot1 thetadot2] :as state]]
           (let [[x1 y1 x2 y2] (xform state)]
             {:t t
              :theta1 (pv theta1)
              :x1 x1
              :y1 y1
              :theta2 (pv theta2)
              :x2 x2
              :y2 y2
              :thetadot1 thetadot1
              :thetadot2 thetadot2
              :d-energy (monitor state)}))
         xs)))

;; The following forms transform the raw data for each initial condition and
;; bind the results to `chaotic-data` and `regular-data` for exploration.

(defonce chaotic-data
  (doall
   (transform-data raw-chaotic-data)))

(defonce regular-data
  (doall
   (transform-data raw-regular-data)))

;; Here's the final, transformed chaotic state:

(last chaotic-data)

;; And the similar regular state:

(last regular-data)

;; ## Data Visualization Utilities

;; [Vega-Lite](https://vega.github.io/vega-lite/) allows us to
;; visualizing the system.

;; I am not a pro here, but this does the trick for now.

;; First, a function to transform the dictionaries we generated above into a
;; sequence of `x, y` coordinates tagged with distinct IDs for each pendulum
;; bob's points:

(defn points-data [data]
  (mapcat (fn [{:keys [t x1 y1 x2 y2]}]
            [{:t  t
              :x  x1
              :y  y1
              :id :p1}
             {:t  t
              :x  x2
              :y  y2
              :id :p2}])
          data))

;; `segments-data` generates endpoints of the pendulum segments, bob-to-bob:

(defn segments-data [data]
  (mapcat (fn [{:keys [t x1 y1 x2 y2]}]
            [{:t  t
              :x 0
              :y 0
              :x2 x1
              :y2 y1
              :id :p1}
             {:t  t
              :x  x1
              :y  y1
              :x2 x2
              :y2 y2
              :id :p2}])
          data))

;; ## Visualizations

;; a helper function that should be in clojure.core
(defn deep-merge [v & vs]
  (letfn [(rec-merge [v1 v2]
            (if (and (map? v1) (map? v2))
              (merge-with deep-merge v1 v2)
              v2))]
    (when (some identity vs)
      (reduce #(rec-merge %1 %2) v vs))))

;; With those tools in hand, let's make some charts. I'll call this first chart
;; the `system-inspector`; this function will return a chart that will let us
;; evolve the system with a time slider and monitor both angles, the energy
;; error, and the pendulum bob itself as it evolves through time.

(defn system-inspector [data]
  {:config {:bar {:binSpacing 1, :discreteBandSize 5, :continuousBandSize 5}},
   :datasets {:points data}
   :vconcat [{:hconcat [{:layer (mapv (partial deep-merge
                                               {:encoding {:y {:field "y", :type "quantitative"},
                                                           :color {:field :id, :type :nominal},
                                                           :opacity {:condition {:test "abs(selected_t - datum['t']) < 0.00001", :value 1},
                                                                     :value 0.3},
                                                           :x {:field "x", :type "quantitative"},
                                                           :y2 {:field "y2", :type "quantitative"},
                                                           :x2 {:field "x2", :type "quantitative"},
                                                           :tooltip [{:field "x", :type "quantitative"}
                                                                     {:field "y", :type "quantitative"}]},
                                                :height 300,
                                                :width 400})
                                      [{:data {:values (points-data data)}
                                        :encoding {:size {:condition
                                                          {:test "abs(selected_t - datum['t']) < 0.00001", :value 200},
                                                          :value 5}
                                                   :opacity {:value 0.3}}
                                        :selection {:selected {:fields [:t],
                                                               :type :single,
                                                               :bind {:t {:min 0.01, :max 49.99, :input :range, :step 0.01}}}}
                                        :mark {:type "circle"}}
                                       {:data {:values (segments-data data)}
                                        :encoding {:opacity {:value 0}}
                                        :mark "rule"}])}
                        {:encoding {:y {:field :d-energy, :type "quantitative"},
                                    :size {:condition
                                           {:test "abs(selected_t - datum['t']) < 0.00001", :value 200},
                                           :value 5},
                                    :x {:field :t, :type "quantitative"},
                                    :y2 {:field "y2", :type "quantitative"},
                                    :x2 {:field "x2", :type "quantitative"},
                                    :tooltip [{:field :t, :type "quantitative"}
                                              {:field :d-energy, :type "quantitative"}]},
                         :mark {:type "circle"},
                         :width 400,
                         :height 300,
                         :data {:name :points}}]}
             {:hconcat (mapv (partial deep-merge {:encoding {:y {:field :unassigned,
                                                                 :type "quantitative",
                                                                 :scale {:domain [(- Math/PI) Math/PI]}},
                                                             :size {:condition {:test "abs(selected_t - datum['t']) < 0.00001", :value 200},
                                                                    :value 5},
                                                             :x {:field :t, :type "quantitative"},
                                                             :y2 {:field "y2", :type "quantitative"},
                                                             :x2 {:field "x2", :type "quantitative"},
                                                             :tooltip [{:field :t, :type "quantitative"}
                                                                       {:field :unassigned, :type "quantitative"}]},
                                                  :mark {:type "circle"},
                                                  :width 400,
                                                  :height 300,
                                                  :data {:name :points}})
                             [{:encoding {:y {:field :theta1},
                                          :tooltip [{:field :theta1}]}}
                              {:encoding {:y {:field :theta2},
                                          :tooltip [{:field :theta2}]}}])}]})

;; Here's a system monitor for the chaotic initial condition:

(clerk/vl
 (system-inspector chaotic-data))

;; And again for the regular initial condition:

(clerk/vl
 (system-inspector regular-data))

;; ## Generalized coordinates, velocities

;; `angles-plot` generates a plot, with no animation, showing both theta angles
;; and their associated velocities:

(defn angles-plot [data]
  (let [quarter-spec {:encoding {:y {:field :unassigned
                                     :type "quantitative"}
                                 :size {:value 5}
                                 :x {:field :t :type "quantitative"}
                                 :y2 {:field "y2" :type "quantitative"}
                                 :x2 {:field "x2" :type "quantitative"}
                                 :tooltip [{:field :t :type "quantitative"}
                                           {:field :unassigned :type "quantitative"}]}
                      :mark {:type "circle"}
                      :width 400
                      :height 300
                      :data {:name :points}}]
    {:config {:bar {:binSpacing 1 :discreteBandSize 5 :continuousBandSize 5}}
     :datasets {:points data}
     :width 800
     :height 600
     :hconcat [{:vconcat (mapv (partial deep-merge quarter-spec)
                               [{:encoding {:y {:field :theta1
                                                :scale {:domain [(- Math/PI) Math/PI]}}
                                            :tooltip [{:field :theta1}]}}
                                {:encoding {:y {:field :theta2
                                                :scale {:domain [(- Math/PI) Math/PI]}}
                                            :tooltip [{:field :theta2}]}}])}
               {:vconcat (mapv (partial deep-merge quarter-spec)
                               [{:encoding {:y {:field :thetadot1}
                                            :tooltip [{:field :theta1}]}}
                                {:encoding {:y {:field :thetadot2}
                                            :tooltip [{:field :theta2}]}}])}]}))

;; Here are the angles for the chaotic initial condition:

(clerk/vl
 (angles-plot chaotic-data))

;; And the regular initial condition:

(clerk/vl
 (angles-plot regular-data))

;; ## Double Double Pendulum!

;; No visualizations here yet, but the code works well.

(defn L-double-double-pendulum [m1 m2 l1 l2 g]
  (fn [[t [thetas1 thetas2] [qdots1 qdots2]]]
    (let [s1 (up t thetas1 qdots1)
          s2 (up t thetas2 qdots2)]
      (+ ((L-double-pendulum m1 m2 l1 l2 g) s1)
         ((L-double-pendulum m1 m2 l1 l2 g) s2)))))

(def dd-state-derivative
  (compose
   Lagrangian->state-derivative
   L-double-double-pendulum))

(defn safe-log [x]
  (if (< x 1e-60)
    -138.0
    (Math/log x)))

(defn divergence-monitor []
  (let [pv (principal-value Math/PI)]
    (fn [[_ [thetas1 thetas2]]]
      (safe-log
       (Math/abs
        (pv
         (- (nth thetas1 1)
            (nth thetas2 1))))))))

(defn run-double-double!
  "Two different initializations, slightly kicked"
  [step horizon initial-q1]
  (let [initial-q2    (+ initial-q1 (up 0.0 1e-10))
        initial-qdot  (up 0.0 0.0)
        initial-state (up 0.0
                          (up initial-q1 initial-q2)
                          (up initial-qdot initial-qdot))
        collector     (atom (transient []))]
    ((evolve dd-state-derivative m1 m2 l1 l2 g)
     initial-state
     step
     horizon
     {:compile? true
      :epsilon 1.0e-13 ; = (max-norm 1.e-13)
      :observe (fn [_ state]
                 (swap! collector conj! state))})
    (persistent! @collector)))

(def raw-dd-chaotic-data
  (run-double-double! step horizon chaotic-initial-q))

;; Looks good:

(peek raw-dd-chaotic-data)

;; Next, the regular initial condition:

(def raw-dd-regular-data
  (run-double-double! step horizon regular-initial-q))

;; Peek at the final state:

(peek raw-dd-regular-data)