From ed5fcc70754c08fc00877326c060ac5dac6df87a Mon Sep 17 00:00:00 2001 From: Felix Koehler Date: Mon, 6 May 2024 08:28:37 +0200 Subject: [PATCH 1/3] Add polynomial difficulty stepper --- exponax/normalized/__init__.py | 3 +- exponax/normalized/_polynomial.py | 51 +++++++++++++++++++++++++++++++ 2 files changed, 53 insertions(+), 1 deletion(-) diff --git a/exponax/normalized/__init__.py b/exponax/normalized/__init__.py index 109a1dc..d44ffa6 100644 --- a/exponax/normalized/__init__.py +++ b/exponax/normalized/__init__.py @@ -15,7 +15,7 @@ DiffultyLinearStepperSimple, NormalizedLinearStepper, ) -from ._polynomial import NormalizedPolynomialStepper +from ._polynomial import DifficultyPolynomialStepper, NormalizedPolynomialStepper from ._utils import ( denormalize_coefficients, denormalize_convection_scale, @@ -39,6 +39,7 @@ "DiffultyLinearStepperSimple", "DifficultyConvectionStepper", "DifficultyGradientNormStepper", + "DifficultyPolynomialStepper", "NormalizedConvectionStepper", "NormlizedGeneralNonlinearStepper", "NormalizedGradientNormStepper", diff --git a/exponax/normalized/_polynomial.py b/exponax/normalized/_polynomial.py index 3692f7e..0ff1402 100644 --- a/exponax/normalized/_polynomial.py +++ b/exponax/normalized/_polynomial.py @@ -3,6 +3,7 @@ from .._base_stepper import BaseStepper from ..nonlin_fun import PolynomialNonlinearFun +from ._utils import extract_normalized_coefficients_from_difficulty class NormalizedPolynomialStepper(BaseStepper): @@ -72,3 +73,53 @@ def _build_nonlinear_fun( coefficients=self.normalized_polynomial_scales, dealiasing_fraction=self.dealiasing_fraction, ) + + +class DifficultyPolynomialStepper(NormalizedPolynomialStepper): + linear_difficulties: tuple[float, ...] + polynomial_difficulties: tuple[float, ...] + + def __init__( + self, + num_spatial_dims: int = 1, + num_points: int = 48, + *, + linear_difficulties: tuple[float, ...] = ( + 10.0 * 0.001 / (10.0**0) * 48**0, + 0.0, + 1.0 * 0.001 / (10.0**2) * 48**2 * 2**1, + ), + polynomial_difficulties: tuple[float, ...] = ( + 0.0, + 0.0, + -10.0 * 0.001, + ), + order: int = 2, + dealiasing_fraction: float = 2 / 3, + num_circle_points: int = 16, + circle_radius: float = 1.0, + ): + """ + By default: Fisher-KPP + """ + self.linear_difficulties = linear_difficulties + self.polynomial_difficulties = polynomial_difficulties + + normalized_coefficients = extract_normalized_coefficients_from_difficulty( + linear_difficulties, + num_spatial_dims=num_spatial_dims, + num_points=num_points, + ) + # For polynomial nonlinearities, we have difficulties == normalized scales + normalized_polynomial_scales = polynomial_difficulties + + super().__init__( + num_spatial_dims=num_spatial_dims, + num_points=num_points, + normalized_coefficients=normalized_coefficients, + normalized_polynomial_scales=normalized_polynomial_scales, + order=order, + dealiasing_fraction=dealiasing_fraction, + num_circle_points=num_circle_points, + circle_radius=circle_radius, + ) From b7a0e90d497e3315034eb9e4b12711a89b017501 Mon Sep 17 00:00:00 2001 From: Felix Koehler Date: Mon, 6 May 2024 08:38:34 +0200 Subject: [PATCH 2/3] Implement difficulty-based general nonlinear stepper --- exponax/normalized/__init__.py | 6 ++- exponax/normalized/_general_nonlinear.py | 65 +++++++++++++++++++++++- exponax/normalized/_utils.py | 50 ++++++++++++++++++ 3 files changed, 119 insertions(+), 2 deletions(-) diff --git a/exponax/normalized/__init__.py b/exponax/normalized/__init__.py index d44ffa6..691df66 100644 --- a/exponax/normalized/__init__.py +++ b/exponax/normalized/__init__.py @@ -8,7 +8,10 @@ slightly differently. """ from ._convection import DifficultyConvectionStepper, NormalizedConvectionStepper -from ._general_nonlinear import NormlizedGeneralNonlinearStepper +from ._general_nonlinear import ( + DifficultyGeneralNonlinearStepper, + NormlizedGeneralNonlinearStepper, +) from ._gradient_norm import DifficultyGradientNormStepper, NormalizedGradientNormStepper from ._linear import ( DifficultyLinearStepper, @@ -40,6 +43,7 @@ "DifficultyConvectionStepper", "DifficultyGradientNormStepper", "DifficultyPolynomialStepper", + "DifficultyGeneralNonlinearStepper", "NormalizedConvectionStepper", "NormlizedGeneralNonlinearStepper", "NormalizedGradientNormStepper", diff --git a/exponax/normalized/_general_nonlinear.py b/exponax/normalized/_general_nonlinear.py index 16bbd59..84f162d 100644 --- a/exponax/normalized/_general_nonlinear.py +++ b/exponax/normalized/_general_nonlinear.py @@ -3,6 +3,10 @@ from .._base_stepper import BaseStepper from ..nonlin_fun import GeneralNonlinearFun +from ._utils import ( + extract_normalized_coefficients_from_difficulty, + extract_normalized_nonlinear_scales_from_difficulty, +) class NormlizedGeneralNonlinearStepper(BaseStepper): @@ -15,7 +19,7 @@ def __init__( num_spatial_dims: int, num_points: int, *, - normalized_coefficients_linear: tuple[float, ...] = (0.0, 0.0, 0.01 * 0.1), + normalized_coefficients_linear: tuple[float, ...] = (0.0, 0.0, 0.1 * 0.1), normalized_coefficients_nonlinear: tuple[float, ...] = (0.0, -1.0 * 0.1, 0.0), order=2, dealiasing_fraction: float = 2 / 3, @@ -68,3 +72,62 @@ def _build_nonlinear_fun( scale_list=self.normalized_coefficients_nonlinear, zero_mode_fix=True, # ToDo: check this ) + + +class DifficultyGeneralNonlinearStepper(NormlizedGeneralNonlinearStepper): + linear_difficulties: tuple[float, ...] + nonlinear_difficulties: tuple[float, ...] + + def __init__( + self, + num_spatial_dims: int = 1, + num_points: int = 48, + *, + linear_difficulties: tuple[float, ...] = ( + 0.0, + 0.0, + 0.1 * 0.1 / 1.0 * 48**2 * 2, + ), + nonlinear_difficulties: tuple[float, ...] = ( + 0.0, + -1.0 * 0.1 / 1.0 * 48, + 0.0, + ), + maximum_absolute: float = 1.0, + order: int = 2, + dealiasing_fraction: float = 2 / 3, + num_circle_points: int = 16, + circle_radius: float = 1.0, + ): + """ + By default Burgers. + """ + self.linear_difficulties = linear_difficulties + self.nonlinear_difficulties = nonlinear_difficulties + + normalized_coefficients_linear = ( + extract_normalized_coefficients_from_difficulty( + linear_difficulties, + num_spatial_dims=num_spatial_dims, + num_points=num_points, + ) + ) + normalized_coefficients_nonlinear = ( + extract_normalized_nonlinear_scales_from_difficulty( + nonlinear_difficulties, + num_spatial_dims=num_spatial_dims, + num_points=num_points, + maximum_absolute=maximum_absolute, + ) + ) + + super().__init__( + num_spatial_dims=num_spatial_dims, + num_points=num_points, + normalized_coefficients_linear=normalized_coefficients_linear, + normalized_coefficients_nonlinear=normalized_coefficients_nonlinear, + order=order, + dealiasing_fraction=dealiasing_fraction, + num_circle_points=num_circle_points, + circle_radius=circle_radius, + ) diff --git a/exponax/normalized/_utils.py b/exponax/normalized/_utils.py index c48eabc..6049e34 100644 --- a/exponax/normalized/_utils.py +++ b/exponax/normalized/_utils.py @@ -218,3 +218,53 @@ def extract_normalized_gradient_norm_scale_from_difficulty( maximum_absolute * jnp.square(num_points) * num_spatial_dims ) return normalized_gradient_norm_scale + + +def reduce_normalized_nonlinear_scales_to_difficulty( + normalized_nonlinear_scales: tuple[float], + *, + num_spatial_dims: int, + num_points: int, + maximum_absolute: float, +): + nonlinear_difficulties = ( + normalized_nonlinear_scales[0], # Polynomial: normalized == difficulty + reduce_normalized_convection_scale_to_difficulty( + normalized_nonlinear_scales[1], + num_spatial_dims=num_spatial_dims, + num_points=num_points, + maximum_absolute=maximum_absolute, + ), + reduce_normalized_gradient_norm_scale_to_difficulty( + normalized_nonlinear_scales[2], + num_spatial_dims=num_spatial_dims, + num_points=num_points, + maximum_absolute=maximum_absolute, + ), + ) + return nonlinear_difficulties + + +def extract_normalized_nonlinear_scales_from_difficulty( + nonlinear_difficulties: tuple[float], + *, + num_spatial_dims: int, + num_points: int, + maximum_absolute: float, +): + normalized_nonlinear_scales = ( + nonlinear_difficulties[0], # Polynomial: normalized == difficulty + extract_normalized_convection_scale_from_difficulty( + nonlinear_difficulties[1], + num_spatial_dims=num_spatial_dims, + num_points=num_points, + maximum_absolute=maximum_absolute, + ), + extract_normalized_gradient_norm_scale_from_difficulty( + nonlinear_difficulties[2], + num_spatial_dims=num_spatial_dims, + num_points=num_points, + maximum_absolute=maximum_absolute, + ), + ) + return normalized_nonlinear_scales From 8f3d5e1c93129db90799d18d75d4179972030527 Mon Sep 17 00:00:00 2001 From: Felix Koehler Date: Mon, 6 May 2024 09:16:03 +0200 Subject: [PATCH 3/3] Add simple streamlit notebook to understand difficulties --- ...understanding_normalized_and_difficulty.py | 154 ++++++++++++++++++ 1 file changed, 154 insertions(+) create mode 100644 examples/understanding_normalized_and_difficulty.py diff --git a/examples/understanding_normalized_and_difficulty.py b/examples/understanding_normalized_and_difficulty.py new file mode 100644 index 0000000..9a07415 --- /dev/null +++ b/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)