diff --git a/sarracen/readers/read_phantom.py b/sarracen/readers/read_phantom.py index 69126a3..482c915 100644 --- a/sarracen/readers/read_phantom.py +++ b/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'] diff --git a/sarracen/render.py b/sarracen/render.py index 84a1032..ccf48bf 100644 --- a/sarracen/render.py +++ b/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) diff --git a/sarracen/sarracen_dataframe.py b/sarracen/sarracen_dataframe.py index 3721045..71ab7db 100644 --- a/sarracen/sarracen_dataframe.py +++ b/sarracen/sarracen_dataframe.py @@ -294,6 +294,7 @@ 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, @@ -301,7 +302,9 @@ def render(self, 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, diff --git a/sarracen/tests/test_render.py b/sarracen/tests/test_render.py index 54e4a19..c935361 100644 --- a/sarracen/tests/test_render.py +++ b/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)