upload first version

app.py

@@ -0,0 +1,128 @@
+import streamlit as st
+import jax
+import jax.numpy as jnp
+import matplotlib.pyplot as plt
+
+
+def parabola_fn(x):
+    return x**0.5
+
+
+def circle_fn(x):
+    return (1 - x**2) ** 0.5
+
+
+d_parabola_fn = jax.grad(parabola_fn)
+d_circle_fn = jax.grad(circle_fn)
+
+
+def loss_fn(params):
+    x1 = params["x1"]
+    x2 = params["x2"]
+    # parpendicular line to the tangent of the parabola: y = m1 * x + c1
+    m1 = -1 / d_parabola_fn(x1)
+    c1 = parabola_fn(x1) - m1 * x1
+
+    def perpendicular_parabola_fn(x):
+        return m1 * x + c1
+
+    # parpendicular line to the tangent of the circle: y = m2 * x + c2
+    m2 = -1 / d_circle_fn(x2)
+    c2 = circle_fn(x2) - m2 * x2
+
+    def perpendicular_circle_fn(x):
+        return m2 * x + c2
+
+    # x_star and y_star are the intersection of the two lines
+    x_star = (c2 - c1) / (m1 - m2)
+    y_star = m1 * x_star + c1
+
+    # three quantities should be equal to each other
+    # 1. distance between intersection and parabola
+    # 2. distance between intersection and circle
+    # 3. distance between intersection and x=0 line
+
+    d1 = (x_star - x1) ** 2 + (y_star - parabola_fn(x1)) ** 2
+    d2 = (x_star - x2) ** 2 + (y_star - circle_fn(x2)) ** 2
+    d3 = x_star**2
+
+    aux = {
+        "x_star": x_star,
+        "y_star": y_star,
+        "perpendicular_parabola_fn": perpendicular_parabola_fn,
+        "perpendicular_circle_fn": perpendicular_circle_fn,
+        "r": d1**0.5,
+    }
+    # final loss
+    loss = (d1 - d2) ** 2 + (d1 - d3) ** 2 + (d2 - d3) ** 2
+
+    return loss, aux
+
+
+x = jnp.linspace(0, 1, 100)
+
+st.title("Radius of the Circle: Optimization Playground")
+
+col1, col2 = st.columns(2)
+
+x1 = col1.slider("initial x1 (x intersection with parabola)", 0.0, 1.0, 0.5)
+x2 = col1.slider("initial x2 (x intersection with the circle)", 0.0, 1.0, 0.5)
+n_epochs = col2.slider("n_epochs", 0, 1000, 50)
+lr = col2.slider("lr", 0.0, 1.0, value=0.1, step=0.01)
+
+# submit button
+submit = st.button("submit")
+
+# when submit button is clicked run the following code
+params = {"x1": x1, "x2": x2}
+losses = []
+value_and_grad_fn = jax.value_and_grad(loss_fn, has_aux=True)
+
+# initialize plot
+fig, axes = plt.subplots(1, 2, figsize=(15, 5))
+axes[0].set_xlim(0, 1)
+axes[0].set_ylim(0, 1)
+
+value, aux = loss_fn(params)
+(pbola_plot,) = axes[0].plot(x, parabola_fn(x), color="red")
+(pbola_perpendicular_plot,) = axes[0].plot(x, aux["perpendicular_parabola_fn"](x), color="red", linestyle="--")
+(cicle_plot,) = axes[0].plot(x, circle_fn(x), color="blue")
+(circle_perpendicular_plot,) = axes[0].plot(x, aux["perpendicular_circle_fn"](x), color="blue", linestyle="--")
+x_star, y_star = aux["x_star"], aux["y_star"]
+radius = aux["r"]
+axes[0].add_patch(plt.Circle((x_star, y_star), radius, fill=False))
+
+axes[1].set_xlim(0, n_epochs)
+axes[1].set_ylim(0, value)
+(loss_plot,) = axes[1].plot(losses, color="black")
+
+pbar = st.progress(0)
+
+with st.empty():
+    st.pyplot(fig)
+    if submit:
+        for i in range(n_epochs):
+            (value, _), grad = value_and_grad_fn(params)
+
+            params["x1"] -= lr * grad["x1"]
+            params["x2"] -= lr * grad["x2"]
+            losses.append(value)
+
+            _, aux = loss_fn(params)
+            print(params, grad, lr)
+            pbola_plot.set_data(x, parabola_fn(x))
+            pbola_perpendicular_plot.set_data(x, aux["perpendicular_parabola_fn"](x))
+            cicle_plot.set_data(x, circle_fn(x))
+            circle_perpendicular_plot.set_data(x, aux["perpendicular_circle_fn"](x))
+            x_star, y_star = aux["x_star"], aux["y_star"]
+            radius = aux["r"]
+            axes[0].add_patch(plt.Circle((x_star, y_star), radius, fill=False))
+
+            loss_plot.set_data(range(len(losses)), losses)
+
+            pbar.progress(i / n_epochs)
+
+            axes[0].set_title(f"x1: {params['x1']:.3f}, x2: {params['x2']:.3f} \n r: {radius:.4f}")
+            axes[1].set_title(f"epoch: {i}, loss: {value:.5f}")
+
+            st.pyplot(fig)

requirements.txt

@@ -0,0 +1,5 @@
+streamlit
+matplotlib
+numpy
+jaxlib
+jax