Merge branch 'main' into linting

sarracen/readers/read_phantom.py

@@ -212,9 +212,18 @@ def _create_mass_column(df, header_vars):
     return df
 
 
-def read_phantom(filename: str,
-                 separate_types: str = 'sinks',
-                 ignore_inactive: bool = True):
+def _create_aprmass_column(df, header_vars):
+    """
+    Creates a mass column with the mass of each particle when there are
+    multiple refinement levels.
+    """
+    df['mass'] = header_vars['massoftype']
+    df['mass'] = df['mass']/(2**(df['apr_level'] - 1))
+
+    return df
+
+
+def read_phantom(filename: str, separate_types: str = 'sinks', ignore_inactive: bool = True):
     """
     Read data from a Phantom dump file.
 
@@ -280,6 +289,9 @@ def read_phantom(filename: str,
                 and 'itype' in df
                 and df['itype'].nunique() > 1):
             df = _create_mass_column(df, header_vars)
+        # create a column if APR is used and automatically scale masses
+        elif 'apr_level' in df:
+            df = _create_aprmass_column(df, header_vars)
         else:  # create global mass parameter
             header_vars['mass'] = header_vars['massoftype']
 

sarracen/render.py

@@ -16,7 +16,7 @@
 import seaborn as sns
 from matplotlib import pyplot as plt
 from matplotlib.axes import Axes
-from matplotlib.colors import Colormap, LogNorm
+from matplotlib.colors import Colormap, LogNorm, SymLogNorm
 
 from .interpolate import interpolate_2d_line, interpolate_2d, \
     interpolate_3d_proj, interpolate_3d_cross, interpolate_3d_vec, \
@@ -143,6 +143,7 @@ def render(data: 'SarracenDataFrame',  # noqa: F821
            rotation: Union[np.ndarray, list, Rotation] = None,
            rot_origin: Union[np.ndarray, list, str] = None,
            log_scale: bool = False,
+           symlog_scale: bool = False,
            dens_weight: bool = None,
            normalize: bool = True,
            hmin: bool = False,
@@ -205,17 +206,12 @@ def render(data: 'SarracenDataFrame',  # noqa: F821
         midpoint, that is, min + max / 2. Defaults to the midpoint.
     log_scale: bool
         Whether to use a logarithmic scale for color coding.
-    dens_weight: bool, optional
-        If True, will plot the target mutliplied by the density. Defaults to
-        True for column-integrated views, when the target is not density, and
-        False for everything else.
-    normalize: bool, optional
-        If True, will normalize the interpolation. Defaults to False (this may
-        change in future versions).
-    hmin: bool, optional
-        If True, a minimum smoothing length of 0.5 * pixel size will be
-        imposed. This ensures each particle contributes to at least one grid
-        cell / pixel. Defaults to False (this may change in a future verison).
+    symlog_scale: bool
+        Whether to use a symmetrical logarithmic scale for color coding (i.e., 
+        allows positive and negative values). Optionally add "linthresh" and 
+        "linscale" to kwargs to set the linear region and the scaling of linear 
+        values, respectively (defaults to 1e-9 and 1, respectevely). Only works 
+        if log_scale == True.
     cototation: list, optional
         Moves particles to the co-rotating frame of two location. corotation
         contains two lists which correspond to the two x, y, z coordinates
@@ -317,8 +313,16 @@ def render(data: 'SarracenDataFrame',  # noqa: F821
     kwargs.setdefault("origin", 'lower')
     kwargs.setdefault("extent", [xlim[0], xlim[1], ylim[0], ylim[1]])
     if log_scale:
-        kwargs.setdefault("norm", LogNorm(clip=True, vmin=kwargs.get('vmin'),
-                                          vmax=kwargs.get('vmax')))
+        if symlog_scale:
+            kwargs.setdefault("norm", 
+                              SymLogNorm(kwargs.pop("linthresh", 1e-9), 
+                                         linscale=kwargs.pop("linscale", 1.),
+                                         vmin=kwargs.get('vmin'), 
+                                         vmax=kwargs.get('vmax')))
+        else:
+            kwargs.setdefault("norm", LogNorm(clip=True, 
+                                              vmin=kwargs.get('vmin'), 
+                                              vmax=kwargs.get('vmax')))
         kwargs.pop("vmin", None)
         kwargs.pop("vmax", None)
 

sarracen/sarracen_dataframe.py

@@ -294,14 +294,17 @@ def render(self,
                rotation: Union[np.ndarray, list, Rotation] = None,
                rot_origin: Union[np.ndarray, list, str] = None,
                log_scale: bool = None,
+               symlog_scale: bool = None,
                dens_weight: bool = None,
                normalize: bool = False,
                hmin: bool = False,
                **kwargs) -> Axes:
         return render(self, target, x, y, z, xsec, kernel, x_pixels, y_pixels,
                       xlim, ylim, cmap, cbar, cbar_kws, cbar_ax, ax, exact,
                       backend, integral_samples, rotation, rot_origin,
-                      log_scale, dens_weight, normalize, hmin, **kwargs)
+                      log_scale, symlog_scale, dens_weight, normalize, hmin, 
+                      **kwargs)
+
 
     @_copy_doc(lineplot)
     def lineplot(self,

sarracen/tests/test_render.py

@@ -153,6 +153,9 @@ def test_kwargs(backend):
               'h': [1, 1], 'rho': [1, 1], 'm': [1, 1]}
     sdf_3 = SarracenDataFrame(data_3)
     sdf_3.backend = backend
+    df_4 = pd.DataFrame({'x': [-3, 6], 'y': [5, -1], 'z': [2, 1], 'P': [-1, 1], 'h': [1, 1], 'rho': [-1, -1], 'm': [1, 1]})
+    sdf_4 = SarracenDataFrame(df_4)
+    sdf_4.backend = backend
 
     for args in [{'data': sdf_2, 'xsec': None},
                  {'data': sdf_3, 'xsec': None},
@@ -161,6 +164,12 @@ def test_kwargs(backend):
         render(args['data'], 'P', xsec=args['xsec'], ax=ax, origin='upper')
         assert ax.images[0].origin == 'upper'
         plt.close(fig)
+
+    for arg in [True, False]:
+        fig, ax = plt.subplots()
+        render(sdf_4, 'P', ax=ax, log_scale=arg, symlog_scale=True, origin='upper', vmin=-1., vmax=1.)
+        assert ax.images[0].origin == 'upper'
+        plt.close(fig)
 
     fig, ax = plt.subplots()
     streamlines(sdf_2, ('Ax', 'Ay'), ax=ax, zorder=5)