Replace temperature script with version that uses swiftsimio

examples/HydroTests/IdealizedAGNJet/analysis_scripts/temperature_slices.py

 import pylab
+import unyt
 from pylab import *
 from scipy import stats
 import h5py as h5
-from sphviewer.tools import QuickView
 import matplotlib
 from matplotlib.colors import LogNorm
 import swiftsimio as sw
 import glob
 import re
 import os
+from swiftsimio.visualisation.projection import scatter
+from swiftsimio import mask, load
 
 matplotlib.use("Agg")
 
@@ -64,6 +66,67 @@ def c1(theta, gama, beta):
 def D_jet(theta, gama, beta, rho_0, P_j, t):
     return c1(theta, gama, beta) * (P_j / rho_0 * t ** 3) ** (1 / 5)
 
+def wrap(dx,box):
+    dx[dx>0.5*box]-=box
+    dx[dx<-0.5*box]+=box
+    return dx
+
+def temperature_scatter(center, snapFile, slice_factor):
+    
+    #set center of box
+    xChoice=center[0]
+    yChoice=center[1]
+    zChoice=center[2]
+    mapRes = 1024
+    
+    xCen = unyt.unyt_quantity(xChoice,'kpc')
+    yCen = unyt.unyt_quantity(yChoice,'kpc')
+    zCen = unyt.unyt_quantity(zChoice,'kpc')
+
+    #Max. region in both z and y-> needs to be a box
+    maxRegion = unyt.unyt_quantity(250,'kpc')
+    
+    #depth of slice
+    depth =  unyt.unyt_quantity(slice_factor*maxRegion,'kpc')
+    maskRegion = mask(snapFile)
+
+    #spatially mask the snapshot data around the center of jet
+    region=[[xCen-maxRegion,xCen+maxRegion],
+            [yCen-depth,yCen+depth],
+            [zCen-maxRegion,zCen+maxRegion]]
+
+    maskRegion.constrain_spatial(region)
+
+    #load the data for only the masked region                                                                                    
+    data = load(snapFile,mask=maskRegion)
+    dx=data.gas.coordinates.value[:,0]-xChoice
+    dy=data.gas.coordinates.value[:,1]-yChoice
+    dz=data.gas.coordinates.value[:,2]-zChoice
+
+    h=data.gas.smoothing_lengths.value
+    m=data.gas.masses.value
+    t=data.gas.temperatures.value
+    d=data.gas.densities.value
+
+    #mask the data
+    ind_mask=np.where((dx>-maxRegion.value)&(dx<maxRegion.value)& 
+                      (dy>-depth)&(dy<depth)&
+                      (dz>-maxRegion.value)&(dz<maxRegion.value))
+    
+    dx=(dx[ind_mask]+maxRegion.value)/(2.*maxRegion.value)
+    dy=(dy[ind_mask]+maxRegion.value)/(2.*maxRegion.value)
+    dz=(dz[ind_mask]+maxRegion.value)/(2.*maxRegion.value)
+    h=h[ind_mask]/(2.*maxRegion.value)
+    m=m[ind_mask]
+    t=t[ind_mask]
+    d=d[ind_mask]
+
+    #scatter the particles, scatter function inverts x and y axes 
+    mapUp=scatter(x=dz,y=dx,h=h,m=m*t,res=mapRes)
+    mapDo=scatter(x=dz,y=dx,h=h,m=m,res=mapRes)
+    map_final=mapUp/mapDo
+
+    return np.log10(map_final)
 
 # Make plot
 snapshots = np.array([5, 10, 15, 20])
@@ -83,110 +146,26 @@ fig = plt.figure(figsize=(23, 10.5))
 gs = gridspec.GridSpec(1, 4, hspace=0, wspace=0)
 
 for i in range(4):
+    
     # Read in the data
-    if snapshots[i] < 10:
-        data = sw.load(
-            path_to_file + snapshot_base + "000" + str(snapshots[i]) + ".hdf5"
-        )
-    if snapshots[i] >= 10 and snapshots[i] < 100:
-        data = sw.load(
-            path_to_file + snapshot_base + "00" + str(snapshots[i]) + ".hdf5"
-        )
-    if snapshots[i] >= 100 and snapshots[i] < 1000:
-        data = sw.load(path_to_file + snapshot_base + "0" + str(snapshots[i]) + ".hdf5")
-    if snapshots[i] >= 1000:
-        data = sw.load(path_to_file + snapshot_base + "" + str(snapshots[i]) + ".hdf5")
-
-    # Box size
-    boxsize = data.metadata.boxsize.value
-
-    # Properties we need
-    gas_positions = data.gas.coordinates.value
-    gas_temperatures = data.gas.temperatures.value
-    gas_masses = (
-        data.gas.masses.value * 1e10
-    )  # quickview complains later if this value is too small
-    gas_smoothing_lengths = data.gas.smoothing_lengths.value
-    gas_ids = data.gas.particle_ids.value
-
-    # Recenter the positions
-    gas_positions[:, 0] -= boxsize[0] / 2.0
-    gas_positions[:, 1] -= boxsize[1] / 2.0
-    gas_positions[:, 2] -= boxsize[2] / 2.0
-
-    # Apply selection to exclude particles in the reservoir
-    selection = gas_ids > 1e7
-
-    # Also exclude particles outside the slice
-    selection = (
-        selection
-        & (gas_positions[:, 1] < slice_factor * r_size)
-        & (gas_positions[:, 1] > -slice_factor * r_size)
-    )
-
-    # Also exclude particles outside visualisation box
-    selection = (
-        selection
-        & (gas_positions[:, 0] < aspect_ratio * r_size)
-        & (gas_positions[:, 0] > -aspect_ratio * r_size)
-    )
-    selection = (
-        selection & (gas_positions[:, 2] < r_size) & (gas_positions[:, 2] > -r_size)
-    )
+    snapFile = path_to_file + snapshot_base + f"{snapshots[i]:04d}.hdf5"
+    print(f'snap {i}')
+    data = sw.load(snapFile)
 
-    gas_masses = gas_masses[selection]
-    gas_temperatures = gas_temperatures[selection]
-    gas_smoothing_lengths = gas_smoothing_lengths[selection]
-    gas_positions = gas_positions[selection]
-
-    # Surface density
-    qv_gas = QuickView(
-        gas_positions,
-        hsml=gas_smoothing_lengths,
-        mass=gas_masses,
-        plot=False,
-        logscale=True,
-        r="infinity",
-        p=0,
-        t=90,
-        extent=[-r_size, r_size, -r_size, r_size],
-        x=0,
-        y=0,
-        z=0,
-        xsize=1500,
-        ysize=1500,
-    )
-    img_gas = qv_gas.get_image()
-    ext_gas = qv_gas.get_extent()
-
-    # Surface density times temperature
-    qv_gas2 = QuickView(
-        gas_positions,
-        hsml=gas_smoothing_lengths,
-        mass=gas_masses * gas_temperatures,
-        plot=False,
-        logscale=True,
-        r="infinity",
-        p=0,
-        t=90,
-        extent=[-r_size, r_size, -r_size, r_size],
-        x=0,
-        y=0,
-        z=0,
-        xsize=1500,
-        ysize=1500,
-    )
-    img_gas2 = qv_gas2.get_image()
+    boxsize = data.metadata.boxsize.value
 
-    # Projection-averaged temperature (logged)
-    img_gas_temp = img_gas2 - img_gas
+    #setup limits for imshow
+    ext_gas=[-125, 125, -250, 250]
+    
+    # Projection-averaged temperature (logged)    
+    img_gas_temp = temperature_scatter([boxsize[0] / 2.0, boxsize[1] / 2.0, boxsize[2] / 2.0], snapFile, slice_factor)
 
     plt.subplot(gs[i])
     # plot the temperature
     im = plt.imshow(
         img_gas_temp, vmin=mintemp, vmax=maxtemp, cmap="inferno", extent=ext_gas
     )
-
+    
     # Some jet parameters for an analytical estimate
     opening_angle_in_degrees = 10
     theta = opening_angle_in_degrees / 180 * np.pi  # in radians
@@ -249,4 +228,4 @@ cb.ax.get_yaxis().labelpad = 30
 cb.ax.tick_params(labelsize=26)
 
 # Save figure
-plt.savefig("temperature_slices.png", bbox_inches="tight", pad_inches=0.1)
+plt.savefig("temperature_slices.png", bbox_inches="tight", pad_inches=0.1)
\ No newline at end of file