Add simple streamlit notebook to understand difficulties

examples/understanding_normalized_and_difficulty.py

@@ -0,0 +1,154 @@
+"""
+This is a streamlit app.
+"""
+import jax
+import matplotlib.pyplot as plt
+import streamlit as st
+
+import exponax as ex
+
+jax.config.update("jax_platform_name", "cpu")
+
+with st.sidebar:
+    num_points = st.slider("Number of points", 16, 256, 48)
+    num_steps = st.slider("Number of steps", 1, 300, 50)
+    num_modes_init = st.slider("Number of modes in the initial condition", 1, 40, 5)
+    num_substeps = st.slider("Number of substeps", 1, 100, 1)
+
+    use_difficulty = st.toggle("Use difficulty", value=True)
+
+    overall_scale = st.slider("Overall scale", 0.1, 50.0, 1.0)
+
+    a_0_cols = st.columns(3)
+    with a_0_cols[0]:
+        a_0_mantissa = st.slider("a_0 mantissa", 0.0, 10.0, 0.0)
+    with a_0_cols[1]:
+        a_0_exponent = st.slider("a_0 exponent", -5, 5, 0)
+    with a_0_cols[2]:
+        a_0_sign = st.select_slider("a_0 sign", options=["-", "+"])
+    a_0 = float(f"{a_0_sign}{a_0_mantissa}e{a_0_exponent}")
+
+    a_1_cols = st.columns(3)
+    with a_1_cols[0]:
+        a_1_mantissa = st.slider("a_1 mantissa", 0.0, 10.0, 0.1)
+    with a_1_cols[1]:
+        a_1_exponent = st.slider("a_1 exponent", -5, 5, 0)
+    with a_1_cols[2]:
+        a_1_sign = st.select_slider("a_1 sign", options=["-", "+"])
+    a_1 = float(f"{a_1_sign}{a_1_mantissa}e{a_1_exponent}")
+
+    a_2_cols = st.columns(3)
+    with a_2_cols[0]:
+        a_2_mantissa = st.slider("a_2 mantissa", 0.0, 10.0, 0.0)
+    with a_2_cols[1]:
+        a_2_exponent = st.slider("a_2 exponent", -5, 5, 0)
+    with a_2_cols[2]:
+        a_2_sign = st.select_slider("a_2 sign", options=["-", "+"])
+    a_2 = float(f"{a_2_sign}{a_2_mantissa}e{a_2_exponent}")
+
+    a_3_cols = st.columns(3)
+    with a_3_cols[0]:
+        a_3_mantissa = st.slider("a_3 mantissa", 0.0, 10.0, 0.0)
+    with a_3_cols[1]:
+        a_3_exponent = st.slider("a_3 exponent", -5, 5, 0)
+    with a_3_cols[2]:
+        a_3_sign = st.select_slider("a_3 sign", options=["-", "+"])
+    a_3 = float(f"{a_3_sign}{a_3_mantissa}e{a_3_exponent}")
+
+    a_4_cols = st.columns(3)
+    with a_4_cols[0]:
+        a_4_mantissa = st.slider("a_4 mantissa", 0.0, 10.0, 0.0)
+    with a_4_cols[1]:
+        a_4_exponent = st.slider("a_4 exponent", -5, 5, 0)
+    with a_4_cols[2]:
+        a_4_sign = st.select_slider("a_4 sign", options=["-", "+"])
+    a_4 = float(f"{a_4_sign}{a_4_mantissa}e{a_4_exponent}")
+
+    b_0_cols = st.columns(3)
+    with b_0_cols[0]:
+        b_0_mantissa = st.slider("b_0 mantissa", 0.0, 10.0, 0.0)
+    with b_0_cols[1]:
+        b_0_exponent = st.slider("b_0 exponent", -5, 5, 0)
+    with b_0_cols[2]:
+        b_0_sign = st.select_slider("b_0 sign", options=["-", "+"])
+    b_0 = float(f"{b_0_sign}{b_0_mantissa}e{b_0_exponent}")
+
+    b_1_cols = st.columns(3)
+    with b_1_cols[0]:
+        b_1_mantissa = st.slider("b_1 mantissa", 0.0, 10.0, 0.0)
+    with b_1_cols[1]:
+        b_1_exponent = st.slider("b_1 exponent", -5, 5, 0)
+    with b_1_cols[2]:
+        b_1_sign = st.select_slider("b_1 sign", options=["-", "+"])
+    b_1 = float(f"{b_1_sign}{b_1_mantissa}e{b_1_exponent}")
+
+    b_2_cols = st.columns(3)
+    with b_2_cols[0]:
+        b_2_mantissa = st.slider("b_2 mantissa", 0.0, 10.0, 0.0)
+    with b_2_cols[1]:
+        b_2_exponent = st.slider("b_2 exponent", -5, 5, 0)
+    with b_2_cols[2]:
+        b_2_sign = st.select_slider("b_2 sign", options=["-", "+"])
+    b_2 = float(f"{b_2_sign}{b_2_mantissa}e{b_2_exponent}")
+
+    # a_0 = st.slider("a_0", -10.0, 10.0, 0.0)
+    # a_1 = st.slider("a_1", -10.0, 10.0, 0.1)
+    # a_2 = st.slider("a_2", -10.0, 10.0, 0.0)
+    # a_3 = st.slider("a_3", -10.0, 10.0, 0.0)
+    # a_4 = st.slider("a_4", -10.0, 10.0, 0.0)
+    # b_0 = st.slider("b_0", -10.0, 10.0, 0.0)
+    # b_1 = st.slider("b_1", -10.0, 10.0, 0.0)
+    # b_2 = st.slider("b_2", -10.0, 10.0, 0.0)
+
+linear_tuple = (a_0, a_1, a_2, a_3, a_4)
+nonlinear_tuple = (b_0, b_1, b_2)
+
+linear_tuple = tuple([overall_scale * x for x in linear_tuple])
+nonlinear_tuple = tuple([overall_scale * x for x in nonlinear_tuple])
+
+if use_difficulty:
+    stepper = ex.RepeatedStepper(
+        ex.normalized.DifficultyGeneralNonlinearStepper(
+            1,
+            num_points,
+            linear_difficulties=tuple(x / num_substeps for x in linear_tuple),
+            nonlinear_difficulties=tuple(x / num_substeps for x in nonlinear_tuple),
+        ),
+        num_substeps,
+    )
+else:
+    stepper = ex.RepeatedStepper(
+        ex.normalized.NormlizedGeneralNonlinearStepper(
+            1,
+            num_points,
+            normalized_coefficients_linear=tuple(
+                x / num_substeps for x in linear_tuple
+            ),
+            normalized_coefficients_nonlinear=tuple(
+                x / num_substeps for x in nonlinear_tuple
+            ),
+        ),
+        num_substeps,
+    )
+
+ic_gen = ex.ic.RandomSineWaves1d(1, cutoff=num_modes_init, max_one=True)
+u_0 = ic_gen(num_points, key=jax.random.PRNGKey(0))
+
+trj = ex.rollout(stepper, num_steps, include_init=True)(u_0)
+
+v_range = st.slider("Colorbar range", 0.1, 10.0, 1.0)
+
+fig, ax = plt.subplots()
+ax.imshow(
+    trj[:, 0, :].T,
+    aspect="auto",
+    vmin=-v_range,
+    vmax=v_range,
+    cmap="RdBu_r",
+    origin="lower",
+)
+
+st.write(f"Linear: {linear_tuple}")
+st.write(f"Nonlinear: {nonlinear_tuple}")
+
+st.pyplot(fig)