Allow limiting scale of ic

Ceyron · Mar 18, 2024 · baf3c7b · baf3c7b
1 parent 847933a
commit baf3c7b
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 6 deletions.
diff --git a/exponax/ic/_diffused_noise.py b/exponax/ic/_diffused_noise.py
@@ -12,14 +12,16 @@ class DiffusedNoise(BaseRandomICGenerator):
     domain_extent: float
     intensity: float
     zero_mean: bool
+    max_one: bool
 
     def __init__(
         self,
         num_spatial_dims: int,
         *,
         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
@@ -15,23 +15,26 @@ class GaussianRandomField(BaseRandomICGenerator):
     num_spatial_dims: int
     domain_extent: float
     powerlaw_exponent: float
-    normalize: bool
+    zero_mean: bool
+    max_one: bool
 
     def __init__(
         self,
         num_spatial_dims: int,
         *,
         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.
         """
         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
@@ -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,13 +28,15 @@ 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
 
         self.cutoff = cutoff
         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