diff --git a/reaction_diffusion_dynamics.py b/reaction_diffusion_dynamics.py new file mode 100644 index 0000000..f8fd35f --- /dev/null +++ b/reaction_diffusion_dynamics.py @@ -0,0 +1,335 @@ +""" +This is a streamlit app. +""" +import base64 +import dataclasses +import io +import json +import random +from dataclasses import dataclass +from typing import Optional + +import jax +import matplotlib.pyplot as plt +import numpy as np +import streamlit as st +import streamlit.components.v1 as components +from IPython.display import DisplayObject +from matplotlib.colors import Colormap, LinearSegmentedColormap, ListedColormap + +import exponax as ex + +st.set_page_config(layout="wide") +jax.config.update("jax_platform_name", "cpu") + +with st.sidebar: + st.title("Exponax Dynamics Brewer") + dimension_type = st.select_slider( + "Number of Spatial Dimensions (ST=Spatio-Temporal plot)", + options=["1d ST", "1d", "2d", "2d ST", "3d"], + ) + 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 (not GS)", 1, 40, 5 + ) + num_substeps = st.slider("Number of substeps", 1, 100, 1) + + v_range = st.slider("Value range", 0.1, 10.0, 1.0) + + st.divider() + + rd_type = st.select_slider( + "Type of Equation", ["FisherKPP", "GrayScott", "SwiftHohenberg"] + ) + + domain_extent_cols = st.columns(3) + with domain_extent_cols[0]: + domain_extent_mantissa = st.slider("domain_extent mantissa", 0.0, 1.0, 0.1) + with domain_extent_cols[1]: + domain_extent_exponent = st.slider("domain_extent exponent", -5, 5, 1) + domain_extent_sign = "+" + domain_extent = float( + f"{domain_extent_sign}{domain_extent_mantissa}e{domain_extent_exponent}" + ) + + dt_cols = st.columns(3) + with dt_cols[0]: + dt_mantissa = st.slider("dt mantissa", 0.0, 1.0, 0.5) + with dt_cols[1]: + dt_exponent = st.slider("dt exponent", -5, 5, 0) + dt_sign = "+" + dt = float(f"{dt_sign}{dt_mantissa}e{dt_exponent}") + + if rd_type == "FisherKPP": + diffusivity_cols = st.columns(3) + with diffusivity_cols[0]: + diffusivity_mantissa = st.slider("diffusivity mantissa", 0.0, 1.0, 0.1) + with diffusivity_cols[1]: + diffusivity_exponent = st.slider("diffusivity exponent", -5, 5, -2) + diffusivity_sign = "+" + diffusivity = float( + f"{diffusivity_sign}{diffusivity_mantissa}e{diffusivity_exponent}" + ) + + reactivity_cols = st.columns(3) + with reactivity_cols[0]: + reactivity_mantissa = st.slider("reactivity mantissa", 0.0, 1.0, 0.1) + with reactivity_cols[1]: + reactivity_exponent = st.slider("reactivity exponent", -5, 5, 0) + reactivity_sign = "+" + reactivity = float( + f"{reactivity_sign}{reactivity_mantissa}e{reactivity_exponent}" + ) + elif rd_type == "GrayScott": + feed_rate = st.slider("feed_rate", 0.00, 0.08, 0.008, 0.004, format="%.3f") + kill_rate = st.slider("kill_rate", 0.03, 0.07, 0.046, 0.002, format="%.3f") + + st.write("Only displays the second channel") + elif rd_type == "SwiftHohenberg": + reactivity = st.slider("reactivity", 0.0, 1.0, 0.7) + critical_number = st.slider("critical_number", 0.0, 1.0, 1.0) + + # st.write(f"dt: {dt}") + # st.write(f"diffusivity: {diffusivity}") + # st.write(f"domain_extent: {domain_extent}") + + # st.write(f"Linear: {linear_tuple}") + # st.write(f"Nonlinear: {nonlinear_tuple}") + +if dimension_type in ["1d ST", "1d"]: + num_spatial_dims = 1 +elif dimension_type in ["2d ST", "2d"]: + num_spatial_dims = 2 +elif dimension_type == "3d": + num_spatial_dims = 3 + +if rd_type == "FisherKPP": + stepper = ex.RepeatedStepper( + ex.reaction.FisherKPP( + num_spatial_dims, + domain_extent, + num_points, + dt / num_substeps, + diffusivity=diffusivity, + reactivity=reactivity, + ), + num_substeps, + ) +elif rd_type == "GrayScott": + stepper = ex.RepeatedStepper( + ex.reaction.GrayScott( + num_spatial_dims, + domain_extent, + num_points, + dt / num_substeps, + feed_rate=feed_rate, + kill_rate=kill_rate, + ), + num_substeps, + ) +elif rd_type == "SwiftHohenberg": + stepper = ex.RepeatedStepper( + ex.reaction.SwiftHohenberg( + num_spatial_dims, + domain_extent, + num_points, + dt / num_substeps, + reactivity=reactivity, + critical_number=critical_number, + ), + num_substeps, + ) + +if rd_type == "GrayScott": + ic_gen = ex.ic.RandomMultiChannelICGenerator( + [ + ex.ic.RandomGaussianBlobs(num_spatial_dims, one_complement=True), + ex.ic.RandomGaussianBlobs(num_spatial_dims), + ] + ) +else: + if num_spatial_dims == 1: + ic_gen = ex.ic.RandomSineWaves1d( + num_spatial_dims, cutoff=num_modes_init, max_one=True + ) + else: + ic_gen = ex.ic.RandomTruncatedFourierSeries( + num_spatial_dims, cutoff=num_modes_init, max_one=True + ) + + if rd_type == "FisherKPP": + ic_gen = ex.ic.ClampingICGenerator(ic_gen) + +u_0 = ic_gen(num_points, key=jax.random.PRNGKey(0)) + +trj = ex.rollout(stepper, num_steps, include_init=True)(u_0) + +if rd_type == "GrayScott": + trj = trj[:, 1:] + + +TEMPLATE_IFRAME = """ +