From aa3c9f5851a2189ac654e56cd7583730f8fc0847 Mon Sep 17 00:00:00 2001 From: Josh Borrow Date: Wed, 11 Sep 2024 17:35:57 -0400 Subject: [PATCH] Formatting --- swiftsimio/optional_packages.py | 2 - swiftsimio/visualisation/__init__.py | 1 - swiftsimio/visualisation/power_spectrum.py | 18 +++---- swiftsimio/visualisation/ray_trace.py | 50 +++++++++++++------ .../smoothing_length/__init__.py | 8 +-- swiftsimio/visualisation/volume_render.py | 18 ++++--- tests/test_visualisation.py | 33 +++++++----- 7 files changed, 76 insertions(+), 54 deletions(-) diff --git a/swiftsimio/optional_packages.py b/swiftsimio/optional_packages.py index f494da44..6f20082a 100644 --- a/swiftsimio/optional_packages.py +++ b/swiftsimio/optional_packages.py @@ -92,5 +92,3 @@ def x(func): # For additional CUDA API access cuda = None - - diff --git a/swiftsimio/visualisation/__init__.py b/swiftsimio/visualisation/__init__.py index b8bc3b34..05928f1b 100644 --- a/swiftsimio/visualisation/__init__.py +++ b/swiftsimio/visualisation/__init__.py @@ -6,4 +6,3 @@ from .slice import slice_scatter as slice from .slice import slice_gas, slice_gas_pixel_grid from .smoothing_length import generate_smoothing_lengths - diff --git a/swiftsimio/visualisation/power_spectrum.py b/swiftsimio/visualisation/power_spectrum.py index c9b0c54a..8f24e978 100644 --- a/swiftsimio/visualisation/power_spectrum.py +++ b/swiftsimio/visualisation/power_spectrum.py @@ -199,7 +199,7 @@ def render_to_deposit( """ # Get the positions and masses - folding = 2.0**folding + folding = 2.0 ** folding positions = data.coordinates quantity = getattr(data, project) @@ -244,10 +244,10 @@ def render_to_deposit( units = 1.0 / ( data.metadata.boxsize[0] * data.metadata.boxsize[1] * data.metadata.boxsize[2] ) - units.convert_to_units(1.0 / data.metadata.boxsize.units**3) + units.convert_to_units(1.0 / data.metadata.boxsize.units ** 3) units *= quantity.units - new_cosmo_factor = quantity.cosmo_factor / (coord_cosmo_factor**3) + new_cosmo_factor = quantity.cosmo_factor / (coord_cosmo_factor ** 3) return cosmo_array( deposition, comoving=comoving, cosmo_factor=new_cosmo_factor, units=units @@ -399,7 +399,7 @@ def folded_depositions_to_power_spectrum( if folding != final_folding: cutoff_wavenumber = ( - 2.0**folding * np.min(depositions[folding].shape) / np.min(box_size) + 2.0 ** folding * np.min(depositions[folding].shape) / np.min(box_size) ) if cutoff_above_wavenumber_fraction is not None: @@ -424,7 +424,7 @@ def folded_depositions_to_power_spectrum( corrected_wavenumber_centers[prefer_bins] = folded_wavenumber_centers[ prefer_bins ].to(corrected_wavenumber_centers.units) - folding_tracker[prefer_bins] = 2.0**folding + folding_tracker[prefer_bins] = 2.0 ** folding contributed_counts[prefer_bins] = folded_counts[prefer_bins] elif transition == "average": @@ -457,7 +457,7 @@ def folded_depositions_to_power_spectrum( # For debugging, we calculate an effective fold number. folding_tracker[use_bins] = ( - (folding_tracker * existing_weight + (2.0**folding) * new_weight) + (folding_tracker * existing_weight + (2.0 ** folding) * new_weight) / transition_norm )[use_bins] @@ -538,7 +538,7 @@ def deposition_to_power_spectrum( deposition.shape == cross_deposition.shape ), "Depositions must have the same shape" - folding = 2.0**folding + folding = 2.0 ** folding box_size_folded = box_size[0] / folding npix = deposition.shape[0] @@ -560,7 +560,7 @@ def deposition_to_power_spectrum( else: conj_fft = fft.conj() - fourier_amplitudes = (fft * conj_fft).real * box_size_folded**3 + fourier_amplitudes = (fft * conj_fft).real * box_size_folded ** 3 # Calculate k-value spacing (centered FFT) dk = 2 * np.pi / (box_size_folded) @@ -592,7 +592,7 @@ def deposition_to_power_spectrum( divisor[zero_mask] = 1 # Correct for folding - binned_amplitudes *= folding**3 + binned_amplitudes *= folding ** 3 # Correct units and names wavenumbers = unyt.unyt_array( diff --git a/swiftsimio/visualisation/ray_trace.py b/swiftsimio/visualisation/ray_trace.py index 3a263a1d..21c7229f 100644 --- a/swiftsimio/visualisation/ray_trace.py +++ b/swiftsimio/visualisation/ray_trace.py @@ -14,11 +14,15 @@ from swiftsimio.objects import cosmo_array from swiftsimio.reader import __SWIFTParticleDataset, SWIFTDataset -from swiftsimio.visualisation.projection_backends.kernels import kernel_gamma, kernel_double_precision as kernel +from swiftsimio.visualisation.projection_backends.kernels import ( + kernel_gamma, + kernel_double_precision as kernel, +) from swiftsimio.accelerated import jit, prange, NUM_THREADS import unyt + @jit(nopython=True, fastmath=True) def core_panels( x: np.float64, @@ -114,7 +118,9 @@ def core_panels( and particle_cell_y >= 0 and particle_cell_y <= maximal_array_index ): - output[particle_cell_x, particle_cell_y, panel] += (m[i] * inverse_cell_area) + output[particle_cell_x, particle_cell_y, panel] += ( + m[i] * inverse_cell_area + ) continue normalisation = 0.0 @@ -157,12 +163,11 @@ def core_panels( r = math.sqrt(distance_x_2 + distance_y_2) - output[cell_x, cell_y, panel] += ( - kernel(r, kernel_width) * normalisation - ) + output[cell_x, cell_y, panel] += kernel(r, kernel_width) * normalisation return output + @jit(nopython=True, fastmath=True) def core_panels_parallel( x: np.float64, @@ -250,9 +255,9 @@ def panel_pixel_grid( raise AttributeError( f'Comoving quantity "{project}" is not compatible with physical coordinates!' ) - m = m.value + m = m.value - # This provides a default 'slice it all' mask. + # This provides a default 'slice it all' mask. if mask is None: mask = np.s_[:] @@ -275,7 +280,6 @@ def panel_pixel_grid( z_min = unyt.unyt_quantity(0.0, units=box_z.units) z_max = box_z - x_range = x_max - x_min y_range = y_max - y_min @@ -283,7 +287,7 @@ def panel_pixel_grid( # we always use the maximum of x_range and y_range to normalise the coordinates # empty pixels in the resulting square image are trimmed afterwards max_range = max(x_range, y_range) - + try: hsml = data.smoothing_lengths except AttributeError: @@ -322,6 +326,7 @@ def panel_pixel_grid( max_z=z_max, ) + def panel_gas( data: SWIFTDataset, resolution: int, @@ -373,7 +378,6 @@ def panel_gas( ) - # --- Functions that actually perform the 'ray tracing'. @@ -381,10 +385,17 @@ def transfer_function(value, width, center): """ A simple gaussian transfer function centered around a specific value. """ - return 1 / (width * np.sqrt(2.0 * np.pi)) * np.exp(-0.5 * ((value - center) / width) ** 2) + return ( + 1 + / (width * np.sqrt(2.0 * np.pi)) + * np.exp(-0.5 * ((value - center) / width) ** 2) + ) + @jit(fastmath=True, nopython=True) -def integrate_ray_numba_specific(input: np.array, red: float, green: float, blue: float, center: float, width: float): +def integrate_ray_numba_specific( + input: np.array, red: float, green: float, blue: float, center: float, width: float +): """ Given a ray, integrate the transfer function along it """ @@ -393,10 +404,16 @@ def integrate_ray_numba_specific(input: np.array, red: float, green: float, blue color = np.array([red, green, blue], dtype=np.float32) for i in input: - value += color * 1 / (width * np.sqrt(2.0 * np.pi)) * np.exp(-0.5 * ((i - center) / width) ** 2) + value += ( + color + * 1 + / (width * np.sqrt(2.0 * np.pi)) + * np.exp(-0.5 * ((i - center) / width) ** 2) + ) return value / len(input) + @jit(fastmath=True, nopython=True) def integrate_ray_numba_nocolor(input: np.array, center: float, width: float): """ @@ -407,7 +424,11 @@ def integrate_ray_numba_nocolor(input: np.array, center: float, width: float): for i in input: value *= 0.99 - value += 1 / (width * np.sqrt(2.0 * np.pi)) * np.exp(-0.5 * ((i - center) / width) ** 2) + value += ( + 1 + / (width * np.sqrt(2.0 * np.pi)) + * np.exp(-0.5 * ((i - center) / width) ** 2) + ) return np.float32(value / len(input)) @@ -445,7 +466,6 @@ def integrate_ray_numba_nocolor(input: np.array, center: float, width: float): # return output - # # %% # import matplotlib.pyplot as plt # import swiftascmaps diff --git a/swiftsimio/visualisation/smoothing_length/__init__.py b/swiftsimio/visualisation/smoothing_length/__init__.py index 75ffd81f..7cccb0fe 100644 --- a/swiftsimio/visualisation/smoothing_length/__init__.py +++ b/swiftsimio/visualisation/smoothing_length/__init__.py @@ -1,10 +1,6 @@ from .sph import get_hsml as sph_get_hsml -from .nearest_neighbours import ( - get_hsml as nearest_neighbours_get_hsml, -) -from .generate import ( - generate_smoothing_lengths, -) +from .nearest_neighbours import get_hsml as nearest_neighbours_get_hsml +from .generate import generate_smoothing_lengths backends_get_hsml = { "sph": sph_get_hsml, diff --git a/swiftsimio/visualisation/volume_render.py b/swiftsimio/visualisation/volume_render.py index 55b99201..47fb9a1c 100644 --- a/swiftsimio/visualisation/volume_render.py +++ b/swiftsimio/visualisation/volume_render.py @@ -845,14 +845,14 @@ def render_gas( * data.metadata.boxsize[1] * data.metadata.boxsize[2] ) - units.convert_to_units(1.0 / data.metadata.boxsize.units**3) + units.convert_to_units(1.0 / data.metadata.boxsize.units ** 3) comoving = data.gas.coordinates.comoving coord_cosmo_factor = data.gas.coordinates.cosmo_factor if project is not None: units *= getattr(data.gas, project).units project_cosmo_factor = getattr(data.gas, project).cosmo_factor - new_cosmo_factor = project_cosmo_factor / coord_cosmo_factor**3 + new_cosmo_factor = project_cosmo_factor / coord_cosmo_factor ** 3 else: new_cosmo_factor = coord_cosmo_factor ** (-3) @@ -946,7 +946,8 @@ def visualise_render( ] images = [ - array([color[0] * x, color[1] * x, color[2] * x]).T for color, x in zip(colors, images) + array([color[0] * x, color[1] * x, color[2] * x]).T + for color, x in zip(colors, images) ] if return_type == "all": @@ -960,9 +961,7 @@ def visualise_render( def visualise_render_options( - centers: list[float], - widths: Union[list[float], float], - cmap: str = "viridis", + centers: list[float], widths: Union[list[float], float], cmap: str = "viridis" ) -> tuple["plt.Figure", "plt.Axes"]: """ Creates a figure of your rendering options. The y-axis is the output value @@ -998,10 +997,13 @@ def visualise_render_options( for center, width, color in zip(centers, widths, colors): xs = linspace(center - 5.0 * width, center + 5.0 * width, 100) - ys = [exp(-0.5 * ((center - x) / width)**2) / (width * sqrt(2.0 * pi)) for x in xs] + ys = [ + exp(-0.5 * ((center - x) / width) ** 2) / (width * sqrt(2.0 * pi)) + for x in xs + ] ax.axvline(center, color=color, linestyle="--") ax.plot(xs, ys, color=color) - return fig, ax \ No newline at end of file + return fig, ax diff --git a/tests/test_visualisation.py b/tests/test_visualisation.py index 14b32ff3..79b59f3a 100644 --- a/tests/test_visualisation.py +++ b/tests/test_visualisation.py @@ -416,25 +416,24 @@ def test_panel_rendering(filename): res = 1024 # Test the panel rendering - panel = panel_gas( - data, - resolution=res, - panels=N_depth, - project="masses", - ) + panel = panel_gas(data, resolution=res, panels=N_depth, project="masses") assert panel.shape[-1] == N_depth import matplotlib.pyplot as plt from matplotlib.colors import LogNorm - plt.imsave("panels_added.png", plt.get_cmap()(LogNorm(vmin=10**6, vmax=10**6.5)(np.sum(panel, axis=-1)))) - - projected = project_gas( - data, res, "masses", backend="renormalised" + plt.imsave( + "panels_added.png", + plt.get_cmap()(LogNorm(vmin=10 ** 6, vmax=10 ** 6.5)(np.sum(panel, axis=-1))), ) - plt.imsave("projected.png", plt.get_cmap()(LogNorm(vmin=10**6, vmax=10**6.5)(projected)),) + projected = project_gas(data, res, "masses", backend="renormalised") + + plt.imsave( + "projected.png", + plt.get_cmap()(LogNorm(vmin=10 ** 6, vmax=10 ** 6.5)(projected)), + ) fullstack = np.zeros((res, res)) @@ -443,11 +442,19 @@ def test_panel_rendering(filename): offset = 32 - plt.imsave("stacked.png", plt.get_cmap()(LogNorm()(fullstack[offset:-offset, offset:-offset]))) + plt.imsave( + "stacked.png", + plt.get_cmap()(LogNorm()(fullstack[offset:-offset, offset:-offset])), + ) - assert np.isclose(panel.sum(axis=-1)[offset:-offset, offset:-offset], projected[offset:-offset, offset:-offset], rtol=0.1).all() + assert np.isclose( + panel.sum(axis=-1)[offset:-offset, offset:-offset], + projected[offset:-offset, offset:-offset], + rtol=0.1, + ).all() return + def test_periodic_boundary_wrapping(): """ Test that periodic boundary wrapping works.