Move norm argument onto wrapper function.

SWIFTSIM · Feb 12, 2025 · 06cc3ca · 06cc3ca
1 parent c4b5795
commit 06cc3ca
Show file tree

Hide file tree

Showing 5 changed files with 9 additions and 10 deletions.
diff --git a/swiftsimio/visualisation/_vistools.py b/swiftsimio/visualisation/_vistools.py
@@ -94,9 +94,8 @@ def _get_rotated_coordinates(data, rotation_matrix, rotation_center):
     return x, y, z
 
 
-def backend_restore_cosmo_and_units(backend_func):
+def backend_restore_cosmo_and_units(backend_func, norm=1.0):
     def wrapper(*args, **kwargs):
-        norm = kwargs.pop("norm")
         comoving = getattr(kwargs["m"], "comoving", None)
         if comoving is True:
             if kwargs["x"].comoving is False or kwargs["y"].comoving is False:

diff --git a/swiftsimio/visualisation/projection.py b/swiftsimio/visualisation/projection.py
@@ -120,10 +120,10 @@ def project_pixel_grid(
         res=resolution,
         box_x=region_info["periodic_box_x"],
         box_y=region_info["periodic_box_y"],
-        norm=(region_info["x_range"] * region_info["y_range"]),
     )
+    norm = region_info["x_range"] * region_info["y_range"]
     backend_func = (backends_parallel if parallel else backends)[backend]
-    image = backend_restore_cosmo_and_units(backend_func)(**kwargs)
+    image = backend_restore_cosmo_and_units(backend_func, norm=norm)(**kwargs)
 
     # determine the effective number of pixels for each dimension
     xres = int(

diff --git a/swiftsimio/visualisation/ray_trace.py b/swiftsimio/visualisation/ray_trace.py
@@ -37,7 +37,8 @@ def panel_pixel_grid(
 
     # There's a parallel version of core_panels but it seems
     # that it's never used anywhere.
-    return backend_restore_cosmo_and_units(backends["core_panels"])(
+    norm = region_info["x_range"] * region_info["y_range"]
+    return backend_restore_cosmo_and_units(backends["core_panels"], norm=norm)(
         x=x[mask] / region_info["max_range"],
         y=y[mask] / region_info["max_range"],
         z=z[mask],
@@ -47,7 +48,6 @@ def panel_pixel_grid(
         panels=panels,
         min_z=region_info["z_min"],
         max_z=region_info["z_max"],
-        norm=(region_info["x_range"] * region_info["y_range"]),
     )
 
 

diff --git a/swiftsimio/visualisation/slice.py b/swiftsimio/visualisation/slice.py
@@ -121,9 +121,9 @@ def slice_gas(
         box_x=region_info["periodic_box_x"],
         box_y=region_info["periodic_box_y"],
         box_z=region_info["periodic_box_z"],
-        norm=(region_info["x_range"] * region_info["y_range"] * region_info["z_range"]),
     )
+    norm = region_info["x_range"] * region_info["y_range"] * region_info["z_range"]
     backend_func = (backends_parallel if parallel else backends)[backend]
-    image = backend_restore_cosmo_and_units(backend_func)(**kwargs)
+    image = backend_restore_cosmo_and_units(backend_func, norm=norm)(**kwargs)
 
     return image
diff --git a/swiftsimio/visualisation/volume_render.py b/swiftsimio/visualisation/volume_render.py
@@ -104,10 +104,10 @@ def render_gas(
         box_x=region_info["periodic_box_x"],
         box_y=region_info["periodic_box_y"],
         box_z=region_info["periodic_box_z"],
-        norm=(region_info["x_range"] * region_info["y_range"] * region_info["z_range"]),
     )
+    norm = region_info["x_range"] * region_info["y_range"] * region_info["z_range"]
     backend_func = (backends_parallel if parallel else backends)["scatter"]
-    image = backend_restore_cosmo_and_units(backend_func)(**kwargs)
+    image = backend_restore_cosmo_and_units(backend_func, norm=norm)(**kwargs)
 
     return image