From baf3c7b279a958885eea8357c3b81d886e5df1f2 Mon Sep 17 00:00:00 2001
From: Felix Koehler <f.koehler@tum.de>
Date: Mon, 18 Mar 2024 10:03:33 +0100
Subject: [PATCH] Allow limiting scale of ic

---
 exponax/ic/_diffused_noise.py           |  8 +++++++-
 exponax/ic/_gaussian_random_field.py    | 15 ++++++++++-----
 exponax/ic/_truncated_fourier_series.py |  6 ++++++
 3 files changed, 23 insertions(+), 6 deletions(-)

diff --git a/exponax/ic/_diffused_noise.py b/exponax/ic/_diffused_noise.py
index 031e89e..615ce8f 100644
--- a/exponax/ic/_diffused_noise.py
+++ b/exponax/ic/_diffused_noise.py
@@ -12,6 +12,7 @@ class DiffusedNoise(BaseRandomICGenerator):
     domain_extent: float
     intensity: float
     zero_mean: bool
+    max_one: bool
 
     def __init__(
         self,
@@ -19,7 +20,8 @@ def __init__(
         *,
         domain_extent: float = 1.0,
         intensity=0.001,
-        zero_mean: bool = False,
+        zero_mean: bool = True,
+        max_one: bool = False,
     ):
         """
         Randomly generated initial condition consisting of a diffused noise field.
@@ -37,6 +39,7 @@ def __init__(
         self.domain_extent = domain_extent
         self.intensity = intensity
         self.zero_mean = zero_mean
+        self.max_one = max_one
 
     def __call__(
         self, num_points: int, *, key: PRNGKeyArray
@@ -56,4 +59,7 @@ def __call__(
         if self.zero_mean:
             ic = ic - jnp.mean(ic)
 
+        if self.max_one:
+            ic = ic / jnp.max(ic)
+
         return ic
diff --git a/exponax/ic/_gaussian_random_field.py b/exponax/ic/_gaussian_random_field.py
index 6714952..9b3d65a 100644
--- a/exponax/ic/_gaussian_random_field.py
+++ b/exponax/ic/_gaussian_random_field.py
@@ -15,7 +15,8 @@ class GaussianRandomField(BaseRandomICGenerator):
     num_spatial_dims: int
     domain_extent: float
     powerlaw_exponent: float
-    normalize: bool
+    zero_mean: bool
+    max_one: bool
 
     def __init__(
         self,
@@ -23,7 +24,8 @@ def __init__(
         *,
         domain_extent: float = 1.0,
         powerlaw_exponent: float = 3.0,
-        normalize: bool = True,
+        zero_mean: bool = True,
+        max_one: bool = False,
     ):
         """
         Randomly generated initial condition consisting of a Gaussian random field.
@@ -31,7 +33,8 @@ def __init__(
         self.num_spatial_dims = num_spatial_dims
         self.domain_extent = domain_extent
         self.powerlaw_exponent = powerlaw_exponent
-        self.normalize = normalize
+        self.zero_mean = zero_mean
+        self.max_one = max_one
 
     def __call__(
         self, num_points: int, *, key: PRNGKeyArray
@@ -62,8 +65,10 @@ def __call__(
             axes=space_indices(self.num_spatial_dims),
         )
 
-        if self.normalize:
+        if self.zero_mean:
             ic = ic - jnp.mean(ic)
-            ic = ic / jnp.std(ic)
+
+        if self.max_one:
+            ic = ic / jnp.max(jnp.abs(ic))
 
         return ic
diff --git a/exponax/ic/_truncated_fourier_series.py b/exponax/ic/_truncated_fourier_series.py
index 2d26446..feec51e 100644
--- a/exponax/ic/_truncated_fourier_series.py
+++ b/exponax/ic/_truncated_fourier_series.py
@@ -18,6 +18,7 @@ class RandomTruncatedFourierSeries(BaseRandomICGenerator):
     amplitude_range: tuple[int, int]
     angle_range: tuple[int, int]
     offset_range: tuple[int, int]
+    max_one: bool
 
     def __init__(
         self,
@@ -27,6 +28,7 @@ def __init__(
         amplitude_range: tuple[int, int] = (-1.0, 1.0),
         angle_range: tuple[int, int] = (0.0, 2.0 * jnp.pi),
         offset_range: tuple[int, int] = (0.0, 0.0),  # no offset by default
+        max_one: bool = False,
     ):
         self.num_spatial_dims = num_spatial_dims
 
@@ -34,6 +36,7 @@ def __init__(
         self.amplitude_range = amplitude_range
         self.angle_range = angle_range
         self.offset_range = offset_range
+        self.max_one = max_one
 
     def __call__(
         self, num_points: int, *, key: PRNGKeyArray
@@ -82,4 +85,7 @@ def __call__(
             axes=space_indices(self.num_spatial_dims),
         )
 
+        if self.max_one:
+            u /= jnp.max(jnp.abs(u))
+
         return u