feat: type hints for 3d eulerian numeric kernels

sopht/numeric/eulerian_grid_ops/stencil_ops_3d/advection_flux_3d.py

@@ -1,22 +1,24 @@
 """Kernels for computing advection flux in 3D."""
+import numpy as np
 import pystencils as ps
-
 import sympy as sp
-
-from sopht.utils.pyst_kernel_config import get_pyst_dtype, get_pyst_kernel_config
+import sopht.utils as spu
+from typing import Callable
 
 
 def gen_advection_flux_conservative_eno3_pyst_kernel_3d(
-    real_t, num_threads=False, fixed_grid_size=False
-):
+    real_t: type,
+    num_threads: bool | int = False,
+    fixed_grid_size: tuple[int, int, int] | bool = False,
+) -> Callable:
     # TODO expand docs
     """3D conservative ENO3 advection flux kernel generator."""
-    pyst_dtype = get_pyst_dtype(real_t)
-    kernel_config = get_pyst_kernel_config(real_t, num_threads)
+    pyst_dtype = spu.get_pyst_dtype(real_t)
+    kernel_config = spu.get_pyst_kernel_config(real_t, num_threads)
     # we can add dtype checks later
     grid_info = (
         f"{fixed_grid_size[0]}, {fixed_grid_size[1]}, {fixed_grid_size[2]}"
-        if fixed_grid_size
+        if type(fixed_grid_size) is tuple[int, ...]
         else "3D"
     )
 
@@ -171,10 +173,16 @@ def _advection_flux_z_back_conservative_eno3_stencil_3d():
     _advection_flux_z_back_conservative_eno3_kernel_3d = ps.create_kernel(
         _advection_flux_z_back_conservative_eno3_stencil_3d, config=kernel_config
     ).compile()
+    x_axis_idx = spu.VectorField.x_axis_idx()
+    y_axis_idx = spu.VectorField.y_axis_idx()
+    z_axis_idx = spu.VectorField.z_axis_idx()
 
     def advection_flux_conservative_eno3_pyst_kernel_3d(
-        advection_flux, field, velocity, inv_dx
-    ):
+        advection_flux: np.ndarray,
+        field: np.ndarray,
+        velocity: np.ndarray,
+        inv_dx: float,
+    ) -> None:
         # TODO expand docs
         """3D conservative ENO3 advection flux kernel.
 
@@ -185,37 +193,37 @@ def advection_flux_conservative_eno3_pyst_kernel_3d(
         _advection_flux_x_front_conservative_eno3_kernel_3d(
             advection_flux=advection_flux,
             field=field,
-            velocity_x=velocity[0],
+            velocity_x=velocity[x_axis_idx],
             inv_dx=inv_dx,
         )
         _advection_flux_x_back_conservative_eno3_kernel_3d(
             advection_flux=advection_flux,
             field=field,
-            velocity_x=velocity[0],
+            velocity_x=velocity[x_axis_idx],
             inv_dx=inv_dx,
         )
         _advection_flux_y_front_conservative_eno3_kernel_3d(
             advection_flux=advection_flux,
             field=field,
-            velocity_y=velocity[1],
+            velocity_y=velocity[y_axis_idx],
             inv_dx=inv_dx,
         )
         _advection_flux_y_back_conservative_eno3_kernel_3d(
             advection_flux=advection_flux,
             field=field,
-            velocity_y=velocity[1],
+            velocity_y=velocity[y_axis_idx],
             inv_dx=inv_dx,
         )
         _advection_flux_z_front_conservative_eno3_kernel_3d(
             advection_flux=advection_flux,
             field=field,
-            velocity_z=velocity[2],
+            velocity_z=velocity[z_axis_idx],
             inv_dx=inv_dx,
         )
         _advection_flux_z_back_conservative_eno3_kernel_3d(
             advection_flux=advection_flux,
             field=field,
-            velocity_z=velocity[2],
+            velocity_z=velocity[z_axis_idx],
             inv_dx=inv_dx,
         )
 

sopht/numeric/eulerian_grid_ops/stencil_ops_3d/advection_timestep_3d.py

@@ -1,43 +1,42 @@
 """Kernels for performing advection timestep in 3D."""
-from sopht.numeric.eulerian_grid_ops.stencil_ops_3d.advection_flux_3d import (
-    gen_advection_flux_conservative_eno3_pyst_kernel_3d,
-)
-from sopht.numeric.eulerian_grid_ops.stencil_ops_3d.elementwise_ops_3d import (
-    gen_elementwise_sum_pyst_kernel_3d,
-    gen_set_fixed_val_pyst_kernel_3d,
-)
+import numpy as np
+import sopht.numeric.eulerian_grid_ops as spne
+import sopht.utils as spu
+from typing import Callable, Literal
 
 
 def gen_advection_timestep_euler_forward_conservative_eno3_pyst_kernel_3d(
-    real_t,
-    num_threads=False,
-    fixed_grid_size=False,
-    field_type="scalar",
-):
+    real_t: type,
+    num_threads: bool | int = False,
+    fixed_grid_size: tuple[int, int, int] | bool = False,
+    field_type: Literal["scalar", "vector"] = "scalar",
+) -> Callable:
     # TODO expand docs
     """3D Advection (ENO3 stencil) Euler forward timestep generator."""
-    assert field_type == "scalar" or field_type == "vector", "Invalid field type"
-    elementwise_sum_pyst_kernel_3d = gen_elementwise_sum_pyst_kernel_3d(
+    elementwise_sum_pyst_kernel_3d = spne.gen_elementwise_sum_pyst_kernel_3d(
         real_t=real_t,
         fixed_grid_size=fixed_grid_size,
         num_threads=num_threads,
     )
-    set_fixed_val_pyst_kernel_3d = gen_set_fixed_val_pyst_kernel_3d(
+    set_fixed_val_pyst_kernel_3d = spne.gen_set_fixed_val_pyst_kernel_3d(
         real_t=real_t,
         fixed_grid_size=fixed_grid_size,
         num_threads=num_threads,
     )
     advection_flux_conservative_eno3_pyst_kernel_3d = (
-        gen_advection_flux_conservative_eno3_pyst_kernel_3d(
+        spne.gen_advection_flux_conservative_eno3_pyst_kernel_3d(
             real_t=real_t,
             fixed_grid_size=fixed_grid_size,
             num_threads=num_threads,
         )
     )
 
     def advection_timestep_euler_forward_conservative_eno3_pyst_kernel_3d(
-        field, advection_flux, velocity, dt_by_dx
-    ):
+        field: np.ndarray,
+        advection_flux: np.ndarray,
+        velocity: np.ndarray,
+        dt_by_dx: float,
+    ) -> None:
         """3D Advection (ENO3 stencil) Euler forward timestep (scalar field).
 
         Performs an inplace advection timestep via ENO3 in 3D using Euler forward,
@@ -54,35 +53,44 @@ def advection_timestep_euler_forward_conservative_eno3_pyst_kernel_3d(
             sum_field=field, field_1=field, field_2=advection_flux
         )
 
-    if field_type == "scalar":
-        return advection_timestep_euler_forward_conservative_eno3_pyst_kernel_3d
-    elif field_type == "vector":
+    match field_type:
+        case "scalar":
+            return advection_timestep_euler_forward_conservative_eno3_pyst_kernel_3d
+        case "vector":
+            x_axis_idx = spu.VectorField.x_axis_idx()
+            y_axis_idx = spu.VectorField.y_axis_idx()
+            z_axis_idx = spu.VectorField.z_axis_idx()
 
-        def vector_field_advection_timestep_euler_forward_conservative_eno3_pyst_kernel_3d(
-            vector_field, advection_flux, velocity, dt_by_dx
-        ):
-            """3D Advection (ENO3 stencil) Euler forward timestep (vector field).
+            def vector_field_advection_timestep_euler_forward_conservative_eno3_pyst_kernel_3d(
+                vector_field: np.ndarray,
+                advection_flux: np.ndarray,
+                velocity: np.ndarray,
+                dt_by_dx: float,
+            ) -> None:
+                """3D Advection (ENO3 stencil) Euler forward timestep (vector field).
 
-            Performs an inplace advection timestep via ENO3 in 3D using Euler forward,
-            for a 3D vector field (3, n, n, n).
-            """
-            advection_timestep_euler_forward_conservative_eno3_pyst_kernel_3d(
-                field=vector_field[0],
-                advection_flux=advection_flux,
-                velocity=velocity,
-                dt_by_dx=dt_by_dx,
-            )
-            advection_timestep_euler_forward_conservative_eno3_pyst_kernel_3d(
-                field=vector_field[1],
-                advection_flux=advection_flux,
-                velocity=velocity,
-                dt_by_dx=dt_by_dx,
-            )
-            advection_timestep_euler_forward_conservative_eno3_pyst_kernel_3d(
-                field=vector_field[2],
-                advection_flux=advection_flux,
-                velocity=velocity,
-                dt_by_dx=dt_by_dx,
-            )
+                Performs an inplace advection timestep via ENO3 in 3D using Euler forward,
+                for a 3D vector field (3, n, n, n).
+                """
+                advection_timestep_euler_forward_conservative_eno3_pyst_kernel_3d(
+                    field=vector_field[x_axis_idx],
+                    advection_flux=advection_flux,
+                    velocity=velocity,
+                    dt_by_dx=dt_by_dx,
+                )
+                advection_timestep_euler_forward_conservative_eno3_pyst_kernel_3d(
+                    field=vector_field[y_axis_idx],
+                    advection_flux=advection_flux,
+                    velocity=velocity,
+                    dt_by_dx=dt_by_dx,
+                )
+                advection_timestep_euler_forward_conservative_eno3_pyst_kernel_3d(
+                    field=vector_field[z_axis_idx],
+                    advection_flux=advection_flux,
+                    velocity=velocity,
+                    dt_by_dx=dt_by_dx,
+                )
 
-        return vector_field_advection_timestep_euler_forward_conservative_eno3_pyst_kernel_3d
+            return vector_field_advection_timestep_euler_forward_conservative_eno3_pyst_kernel_3d
+        case _:
+            raise ValueError("Invalid field type")

sopht/numeric/eulerian_grid_ops/stencil_ops_3d/brinkmann_penalise_3d.py

@@ -1,25 +1,24 @@
 """Kernels for Brinkmann penalisation in 3D."""
+import numpy as np
 import pystencils as ps
-
 import sympy as sp
-
-from sopht.utils.pyst_kernel_config import get_pyst_dtype, get_pyst_kernel_config
+import sopht.utils as spu
+from typing import Callable, Literal
 
 
 def gen_brinkmann_penalise_pyst_kernel_3d(
-    real_t,
-    num_threads=False,
-    fixed_grid_size=False,
-    field_type="scalar",
-):
+    real_t: type,
+    num_threads: bool | int = False,
+    fixed_grid_size: tuple[int, int, int] | bool = False,
+    field_type: Literal["scalar", "vector"] = "scalar",
+) -> Callable:
     """Brinkmann penalisation 3D kernel generator."""
-    assert field_type == "scalar" or field_type == "vector", "Invalid field type"
-    pyst_dtype = get_pyst_dtype(real_t)
-    kernel_config = get_pyst_kernel_config(real_t, num_threads)
+    pyst_dtype = spu.get_pyst_dtype(real_t)
+    kernel_config = spu.get_pyst_kernel_config(real_t, num_threads)
     # we can add dtype checks later
     grid_info = (
         f"{fixed_grid_size[0]}, {fixed_grid_size[1]}, {fixed_grid_size[2]}"
-        if fixed_grid_size
+        if type(fixed_grid_size) is tuple[int, ...]
         else "3D"
     )
 
@@ -37,38 +36,44 @@ def _brinkmann_penalise_stencil_3d():
     brinkmann_penalise_pyst_kernel_3d = ps.create_kernel(
         _brinkmann_penalise_stencil_3d, config=kernel_config
     ).compile()
-    if field_type == "scalar":
-        return brinkmann_penalise_pyst_kernel_3d
-    elif field_type == "vector":
+    match field_type:
+        case "scalar":
+            return brinkmann_penalise_pyst_kernel_3d
+        case "vector":
+            x_axis_idx = spu.VectorField.x_axis_idx()
+            y_axis_idx = spu.VectorField.y_axis_idx()
+            z_axis_idx = spu.VectorField.z_axis_idx()
 
-        def brinkmann_penalise_vector_field_pyst_kernel_3d(
-            penalised_vector_field,
-            penalty_factor,
-            char_field,
-            penalty_vector_field,
-            vector_field,
-        ):
-            """Brinkmann penalises a vector field in 3D."""
-            brinkmann_penalise_pyst_kernel_3d(
-                penalised_field=penalised_vector_field[0],
-                penalty_factor=penalty_factor,
-                char_field=char_field,
-                penalty_field=penalty_vector_field[0],
-                field=vector_field[0],
-            )
-            brinkmann_penalise_pyst_kernel_3d(
-                penalised_field=penalised_vector_field[1],
-                penalty_factor=penalty_factor,
-                char_field=char_field,
-                penalty_field=penalty_vector_field[1],
-                field=vector_field[1],
-            )
-            brinkmann_penalise_pyst_kernel_3d(
-                penalised_field=penalised_vector_field[2],
-                penalty_factor=penalty_factor,
-                char_field=char_field,
-                penalty_field=penalty_vector_field[2],
-                field=vector_field[2],
-            )
+            def brinkmann_penalise_vector_field_pyst_kernel_3d(
+                penalised_vector_field: np.ndarray,
+                penalty_factor: float,
+                char_field: np.ndarray,
+                penalty_vector_field: np.ndarray,
+                vector_field: np.ndarray,
+            ) -> None:
+                """Brinkmann penalises a vector field in 3D."""
+                brinkmann_penalise_pyst_kernel_3d(
+                    penalised_field=penalised_vector_field[x_axis_idx],
+                    penalty_factor=penalty_factor,
+                    char_field=char_field,
+                    penalty_field=penalty_vector_field[x_axis_idx],
+                    field=vector_field[x_axis_idx],
+                )
+                brinkmann_penalise_pyst_kernel_3d(
+                    penalised_field=penalised_vector_field[y_axis_idx],
+                    penalty_factor=penalty_factor,
+                    char_field=char_field,
+                    penalty_field=penalty_vector_field[y_axis_idx],
+                    field=vector_field[y_axis_idx],
+                )
+                brinkmann_penalise_pyst_kernel_3d(
+                    penalised_field=penalised_vector_field[z_axis_idx],
+                    penalty_factor=penalty_factor,
+                    char_field=char_field,
+                    penalty_field=penalty_vector_field[z_axis_idx],
+                    field=vector_field[z_axis_idx],
+                )
 
-        return brinkmann_penalise_vector_field_pyst_kernel_3d
+            return brinkmann_penalise_vector_field_pyst_kernel_3d
+        case _:
+            raise ValueError("Invalid field type")

sopht/numeric/eulerian_grid_ops/stencil_ops_3d/char_func_from_level_set_3d.py

@@ -1,31 +1,29 @@
 """Kernels for computing characteristic function from level set field in 3D."""
 import numpy as np
-
 import pystencils as ps
-
 import sympy as sp
-
-from sopht.utils.pyst_kernel_config import get_pyst_dtype, get_pyst_kernel_config
+import sopht.utils as spu
+from typing import Callable
 
 
 def gen_char_func_from_level_set_via_sine_heaviside_pyst_kernel_3d(
-    blend_width,
-    real_t,
-    num_threads=False,
-    fixed_grid_size=False,
-):
+    blend_width: float,
+    real_t: type,
+    num_threads: bool | int = False,
+    fixed_grid_size: tuple[int, int, int] | bool = False,
+) -> Callable:
     """Level set --> Characteristic function 3D kernel generator.
 
     Generate function that computes characteristic function field
     from the level set field, via a smooth sine Heaviside function.
     """
     grid_info = (
         f"{fixed_grid_size[0]}, {fixed_grid_size[1]}, {fixed_grid_size[2]}"
-        if fixed_grid_size
+        if type(fixed_grid_size) is tuple[int, ...]
         else "3D"
     )
-    pyst_dtype = get_pyst_dtype(real_t)
-    kernel_config = get_pyst_kernel_config(real_t, num_threads)
+    pyst_dtype = spu.get_pyst_dtype(real_t)
+    kernel_config = spu.get_pyst_kernel_config(real_t, num_threads)
     sine_prefactor = np.pi / blend_width
 
     @ps.kernel