From 7806aec79e967785316eacd00940f9a3d0189a29 Mon Sep 17 00:00:00 2001 From: Jo Bovy Date: Tue, 8 Oct 2024 13:58:19 -0400 Subject: [PATCH] Allow plotting of potentials on physical axes --- HISTORY.txt | 3 +- galpy/potential/Potential.py | 127 ++++++++++++++--------------------- 2 files changed, 52 insertions(+), 78 deletions(-) diff --git a/HISTORY.txt b/HISTORY.txt index 124e278a6..e9b57185c 100644 --- a/HISTORY.txt +++ b/HISTORY.txt @@ -1,7 +1,8 @@ v1.10.1 (Expected around 2024-11-01) ==================== -- Propagate general plotting keywords in Potential.plot/plotPotentials. +- Propagate general plotting keywords in Potential.plot/plotPotentials. Also allow + plotting potentials on physical axes. - Added the generalized particle-spray model as galpy.df.basestreamspraydf with two subclasses: chen24spraydf and fadal15spraydf. Enabled integrating orbits of stream diff --git a/galpy/potential/Potential.py b/galpy/potential/Potential.py index 929e599ad..faf93da53 100644 --- a/galpy/potential/Potential.py +++ b/galpy/potential/Potential.py @@ -974,79 +974,26 @@ def plot( - 2024-07-09 - Propagate other plotting keywords - Bovy (UofT) """ - if effective and xy: - raise RuntimeError("xy and effective cannot be True at the same time") - rmin = conversion.parse_length(rmin, ro=self._ro) - rmax = conversion.parse_length(rmax, ro=self._ro) - zmin = conversion.parse_length(zmin, ro=self._ro) - zmax = conversion.parse_length(zmax, ro=self._ro) - Lz = conversion.parse_angmom(Lz, ro=self._ro, vo=self._vo) - if xrange is None: - xrange = [rmin, rmax] - if yrange is None: - yrange = [zmin, zmax] - if savefilename is not None and os.path.exists(savefilename): - print("Restoring savefile " + savefilename + " ...") - savefile = open(savefilename, "rb") - potRz = pickle.load(savefile) - Rs = pickle.load(savefile) - zs = pickle.load(savefile) - savefile.close() - else: - if effective and Lz is None: - raise RuntimeError("When effective=True, you need to specify Lz=") - Rs = numpy.linspace(xrange[0], xrange[1], nrs) - zs = numpy.linspace(yrange[0], yrange[1], nzs) - potRz = numpy.zeros((nrs, nzs)) - for ii in range(nrs): - for jj in range(nzs): - if xy: - R, phi, z = coords.rect_to_cyl(Rs[ii], zs[jj], 0.0) - else: - R, z = Rs[ii], zs[jj] - potRz[ii, jj] = evaluatePotentials( - self, R, z, t=t, phi=phi, use_physical=False - ) - if effective: - potRz[ii, :] += 0.5 * Lz**2 / Rs[ii] ** 2.0 - # Don't plot outside of the desired range - potRz[Rs < rmin, :] = numpy.nan - potRz[Rs > rmax, :] = numpy.nan - potRz[:, zs < zmin] = numpy.nan - potRz[:, zs > zmax] = numpy.nan - # Infinity is bad for plotting - potRz[~numpy.isfinite(potRz)] = numpy.nan - if savefilename is not None: - print("Writing savefile " + savefilename + " ...") - savefile = open(savefilename, "wb") - pickle.dump(potRz, savefile) - pickle.dump(Rs, savefile) - pickle.dump(zs, savefile) - savefile.close() - if xy: - xlabel = r"$x/R_0$" - ylabel = r"$y/R_0$" - else: - xlabel = r"$R/R_0$" - ylabel = r"$z/R_0$" - if levels is None: - levels = numpy.linspace(numpy.nanmin(potRz), numpy.nanmax(potRz), ncontours) - if cntrcolors is None: - cntrcolors = "k" - return plot.dens2d( - potRz.T, - origin="lower", - cmap="gist_gray", - contours=True, - xlabel=xlabel, - ylabel=ylabel, + return plotPotentials( + self, + t=t, + rmin=rmin, + rmax=rmax, + nrs=nrs, + zmin=zmin, + zmax=zmax, + nzs=nzs, + phi=phi, + xy=xy, + effective=effective, + Lz=Lz, xrange=xrange, yrange=yrange, - aspect=0.75 * (rmax - rmin) / (zmax - zmin), - cntrls="-", justcontours=justcontours, levels=levels, cntrcolors=cntrcolors, + ncontours=ncontours, + savefilename=savefilename, **kwargs, ) @@ -2829,6 +2776,7 @@ def plotPotentials( ----- - 2010-07-09 - Written by Bovy (NYU). - 2024-07-09 - Propagate keyword arguments to `galpy.util.plot.dens2d` - Bovy (UofT) + - 2024-10-08 - Allow plotting in physical coordinates - Bovy (UofT) See Also -------- @@ -2838,15 +2786,32 @@ def plotPotentials( if effective and xy: raise RuntimeError("xy and effective cannot be True at the same time") Pot = flatten(Pot) - rmin = conversion.parse_length(rmin, **get_physical(Pot)) - rmax = conversion.parse_length(rmax, **get_physical(Pot)) - zmin = conversion.parse_length(zmin, **get_physical(Pot)) - zmax = conversion.parse_length(zmax, **get_physical(Pot)) - Lz = conversion.parse_angmom(Lz, **get_physical(Pot)) + physical_kwargs = conversion.extract_physical_kwargs(kwargs) + use_physical, ro, vo = conversion.physical_output(Pot, physical_kwargs, "density") + if ro is None: + ro = get_physical(Pot)["ro"] + if vo is None: + vo = get_physical(Pot)["vo"] + physical_kwargs["quantity"] = False # make sure to not use quantity output + rmin = conversion.parse_length(rmin, ro=ro) + rmax = conversion.parse_length(rmax, ro=ro) + zmin = conversion.parse_length(zmin, ro=ro) + zmax = conversion.parse_length(zmax, ro=ro) + Lz = conversion.parse_angmom(Lz, ro=ro, vo=vo) if xrange is None: xrange = [rmin, rmax] + else: + xrange = [ + conversion.parse_length(xrange[0], ro=ro), + conversion.parse_length(xrange[1], ro=ro), + ] if yrange is None: yrange = [zmin, zmax] + else: + yrange = [ + conversion.parse_length(yrange[0], ro=ro), + conversion.parse_length(yrange[1], ro=ro), + ] if savefilename is not None and os.path.exists(savefilename): print("Restoring savefile " + savefilename + " ...") savefile = open(savefilename, "rb") @@ -2888,11 +2853,19 @@ def plotPotentials( if aspect is None: aspect = 0.75 * (rmax - rmin) / (zmax - zmin) if xy: - xlabel = r"$x/R_0$" - ylabel = r"$y/R_0$" + xlabel = r"$x" + ylabel = r"$y" else: - xlabel = r"$R/R_0$" - ylabel = r"$z/R_0$" + xlabel = r"$R" + ylabel = r"$z" + if use_physical: + xlabel += r"\,(\mathrm{kpc})$" + ylabel += r"\,(\mathrm{kpc})$" + xrange = [xrange[0] * ro, xrange[1] * ro] + yrange = [yrange[0] * ro, yrange[1] * ro] + else: + xlabel += "/R_0$" + ylabel += "/R_0$" if levels is None: levels = numpy.linspace(numpy.nanmin(potRz), numpy.nanmax(potRz), ncontours) if cntrcolors is None: