diff --git a/psrqpy/search.py b/psrqpy/search.py index 2f78447b..d853cdd8 100644 --- a/psrqpy/search.py +++ b/psrqpy/search.py @@ -10,6 +10,7 @@ from collections import OrderedDict import re import six +import sys from six.moves import cPickle as pickle from six import string_types @@ -2643,11 +2644,232 @@ def __repr__(self): return repr(self.table) + def plot(self, param1, param2, logx=False, logy=False, excludeAssoc=[], showtypes=[], + showGCs=False, showSNRs=False, markertypes={}, usecondition=True, usepsrs=True, + rcparams={}, usealtair=False, xlabel=None, ylabel=None, intrinsicpdot=False, + xlims=None, ylims=None, **kwargs): + """ + Produce a scatter plot of pairs of parameters using :module:`matplotlib.pyplot`. + + Args: + param1 (str): a pulsar parameter to be plotted along the x-axis. The parameter values + must be floating point numbers. + param2 (str): a pulsar parameter to be plotted along the y-axis. The parameter values + must be floating point numbers. + logx (bool): set to True to use a logarithmic x-axis. Defaults to False. + logy (bool): set to True to use a logarithmic y-axis. Defaults to False. + showtypes (list, str): a list of pulsar types to highlight with markers in the plot. + These can contain any of the following: ``BINARY``, ``HE``, ``NRAD``, ``RRAT``, + ``XINS``, ``AXP`` or ``SGR``, or ``ALL`` to show all types. Default to showing no + types. + showGCs (bool): show markers to denote the pulsars in globular clusters. Defaults to + False. + showSNRs (bool): show markers to denote the pulsars with supernova remnants associated + with them. Defaults to False. + markertypes (dict): a dictionary of marker styles and colors keyed to the pulsar types + above. + usecondition (bool): if True create the diagram only with + pulsars that conform to the original query condition values. + Defaults to True. + usepsrs (bool): if True create the P-Pdot diagram only with pulsars + specified in the original query. Defaults to True. + excludeAssoc (list): a list of pulsar association type (only 'GC' or 'SNR' at the + moment) to exclude from the plot. + usealtair (bool): create an altair interactive figure rather than a + :class:`matplotlib.figure.Figure`. + xlabel (str): a label to be used for the x-axis if not wanting to use the standard + parameter name. + ylabel (str): a label to be used for the y-axis if not wanting to use the standard + parameter name. + xlims (list): the [min, max] x-limits to plot with + ylims (array_like): the [min, max] y-limits to plot with + intrinsicpdot (bool): use the intrinsic period derivative corrected + for the `Shklovskii effect `_ + rather than the observed value. Defaults to False. + """ + + try: + import matplotlib as mpl + from matplotlib import pyplot as pl + except ImportError: + raise Exception('Cannot produce P-Pdot plot as Matplotlib is not available') + + if usealtair: + # import altair + try: + import altair as alt + except ImportError: + raise ImportError('Cannot create plot using altair as it is not available') + + # check that both input parameters have been queried + orig_query = deepcopy(self.query_params) + if self.query_params is None: + self.query_params = [param1, param2] + else: + if param1 not in self.query_params: + self.query_params = self.query_params + [param1] + if param2 not in self.query_params: + self.query_params = self.query_params + [param2] + + for param in ["JNAME", "ASSOC", "BINARY", "TYPE", "P1_I"]: + if param not in self.query_params: + if param == "P1_I" and not intrinsicpdot: + continue + self.query_params = self.query_params + [param] + + # get table containing all required parameters + table = self.query_table(usecondition=usecondition, + usepsrs=usepsrs) + + # reset query parameters + self.query_params = orig_query + + if not self.num_pulsars: + print("No pulsars found, so no plot has been produced") + return None + + if isinstance(showtypes, string_types): + nshowtypes = [showtypes] + else: + nshowtypes = showtypes + + for stype in list(nshowtypes): + if 'ALL' == stype.upper(): + nshowtypes = list(PSR_TYPE) + # remove radio as none are returned as this + del nshowtypes[nshowtypes.index('RADIO')] + break + elif stype.upper() not in list(PSR_TYPE): + warnings.warn('"TYPE" {} is not recognised, so will not be ' + 'included'.format(stype)) + del nshowtypes[nshowtypes.index(stype)] + if 'SGR' == stype.upper(): # synonym for AXP + nshowtypes[nshowtypes.index(stype)] = 'AXP' + + # extract periods and period derivatives + p1values = table[param1] + p2values = table[param2] + + if "P1" in [param1, param2] and intrinsicpdot: + ipdotidx = np.isfinite(table['P1_I']) + if param1 == "P1": + p1values[ipdotidx] = table['P1_I'][ipdotidx] + else: + p2values[ipdotidx] = table['P1_I'][ipdotidx] + + # get only finite values + pidx = (np.isfinite(p1values)) & (np.isfinite(p2values)) + p1values = p1values[pidx] + p2values = p2values[pidx] + + if 'ASSOC' in table.columns: + assocs = table['ASSOC'][pidx] # associations + if 'TYPE' in table.columns: + types = table['TYPE'][pidx] # pulsar types + if 'BINARY' in nshowtypes: + binaries = table['BINARY'][pidx] # binary pulsars + + # check whether to exclude GC or SNR pulsars + for exc in excludeAssoc: + nongcidxs = np.flatnonzero( + np.char.find(np.array(assocs.tolist(), + dtype=np.str), '{}:'.format(exc)) == -1) + p1values = p1values[nongcidxs] + p2values = p2values[nongcidxs] + if 'ASSOC' in table.columns: + assocs = assocs[nongcidxs] + if 'TYPE' in table.columns: + types = types[nongcidxs] + if 'BINARY' in nshowtypes: + binaries = binaries[nongcidxs] + + # set axes scales + scalex = 'linear' if not logx else 'log' + scaley = 'linear' if not logy else 'log' + scales = [scalex, scaley] + + # if scale is log, but data contains negative values, switch to a + # symlog scale + for scale, param in zip(scales, [p1values, p2values]): + if scale == "log": + if np.any(param <= 0.0): + scale = "symlog" + + if usealtair: + # set tool tip values + tooltip = [] + if 'JNAME' in table.columns: + tooltip.append('JNAME') + if 'F0' in table.columns and 'F0' not in [param1, param2]: + tooltip.append('F0') + elif 'P0' in table.columns and 'P0' not in [param1, param2]: + tooltip.append('P0') + tooltip.append(param1) + tooltip.append(param2) + + psrtable = {} + if "BINARY" in nshowtypes: + psrtable["BINARY"] = binaries.to_pandas() + + text = alt.TextConfig(font='Helvetica') + config = alt.Config(text=text) + + chart = alt.Chart(table.to_pandas(), height=600, width=800, + config=config).mark_circle().encode( + alt.X(param1+':Q', scale=alt.Scale(type=scales[0]), + axis=alt.Axis(format='~g'), + title=(param1 if xlabel is None else xlabel)), + alt.Y(param2+':Q', scale=alt.Scale(type=scales[1]), + axis=alt.Axis(format='~g'), + title=(param2 if ylabel is None else ylabel)), + tooltip=tooltip + ).configure_axis( + titleFont="Helvetica", + titleFontSize=12, + labelFont="Helvetica" + ) + + return chart.interactive() + else: + rcparams['figure.figsize'] = rcparams['figure.figsize'] if 'figure.figsize' in rcparams else (9, 9.5) + rcparams['figure.dpi'] = rcparams['figure.dpi'] if 'figure.dpi' in rcparams else 250 + rcparams['text.usetex'] = rcparams['text.usetex'] if 'text.usetex' in rcparams else True + rcparams['axes.linewidth'] = rcparams['axes.linewidth'] if 'axes.linewidth' in rcparams else 0.5 + rcparams['axes.grid'] = rcparams['axes.grid'] if 'axes.grid' in rcparams else False + rcparams['font.family'] = rcparams['font.family'] if 'font.family' in rcparams else 'sans-serif' + rcparams['font.sans-serif'] = rcparams['font.sans-serif'] if 'font.sans-serif' in rcparams else 'Avant Garde, Helvetica, Computer Modern Sans serif' + rcparams['font.size'] = rcparams['font.size'] if 'font.size' in rcparams else 20 + rcparams['legend.fontsize'] = rcparams['legend.fontsize'] if 'legend.fontsize' in rcparams else 16 + rcparams['legend.frameon'] = rcparams['legend.frameon'] if 'legend.frameon' in rcparams else False + + mpl.rcParams.update(rcparams) + + fig, ax = pl.subplots() + + ax.plot(p1values, p2values, **kwargs) + ax.set_xscale(scales[0]) + ax.set_yscale(scales[1]) + ax.set_xlabel(param1 if xlabel is None else xlabel) + ax.set_ylabel(param2 if ylabel is None else ylabel) + + # get limits + if xlims is None and scales[0] == "log": + xlims = [10**np.floor(np.min(np.log10(p1values))), + 10.*int(np.ceil(np.max(p1values)/10.))] + if ylims is None and scales[1] == "log": + ylims = [10**np.floor(np.min(np.log10(p2values))), + 10**np.ceil(np.max(np.log10(p2values)))] + ax.set_xlim(xlims) + ax.set_ylim(ylims) + + return fig, ax + def ppdot(self, intrinsicpdot=False, excludeGCs=False, showtypes=[], showGCs=False, showSNRs=False, markertypes={}, deathline=True, deathmodel='Ip', filldeath=True, filldeathtype={}, showtau=True, brakingidx=3, tau=None, showB=True, Bfield=None, pdotlims=None, - periodlims=None, usecondition=True, usepsrs=True, rcparams={}): + periodlims=None, usecondition=True, usepsrs=True, rcparams={}, + usealtair=False): """ Draw a lovely period vs period derivative diagram. @@ -2697,6 +2919,8 @@ def ppdot(self, intrinsicpdot=False, excludeGCs=False, showtypes=[], specified in the original query. Defaults to True. rcparams (dict): a dictionary of :py:obj:`matplotlib.rcParams` setup parameters for the plot. + usealtair (bool): use `altair `_ to + produce the plot. Returns: :class:`matplotlib.figure.Figure`: the figure object @@ -2711,268 +2935,208 @@ def ppdot(self, intrinsicpdot=False, excludeGCs=False, showtypes=[], from .utils import death_line, label_line - # get table containing all required parameters - table = self.query_table(usecondition=usecondition, - usepsrs=usepsrs, - query_params=['P0', 'P1', 'P1_I', 'ASSOC', - 'BINARY', 'TYPE']) - - if len(table) == 0: - print("No pulsars found, so no P-Pdot plot has been produced") - return None - - if isinstance(showtypes, string_types): - nshowtypes = [showtypes] + orig_condition = self.condition + if usecondition: + if "P1 > 0" not in self.condition: + # only use positive period derivaties + if self.condition is None: + self.condition = "P1 > 0" + else: + self.condition += " && P1 > 0" else: - nshowtypes = showtypes - - for stype in list(nshowtypes): - if 'ALL' == stype.upper(): - nshowtypes = list(PSR_TYPE) - # remove radio as none are returned as this - del nshowtypes[nshowtypes.index('RADIO')] - break - elif stype.upper() not in list(PSR_TYPE): - warnings.warn('"TYPE" {} is not recognised, so will not be ' - 'included'.format(stype)) - del nshowtypes[nshowtypes.index(stype)] - if 'SGR' == stype.upper(): # synonym for AXP - nshowtypes[nshowtypes.index(stype)] = 'AXP' - - # set plot parameters - rcparams['figure.figsize'] = rcparams['figure.figsize'] if \ - 'figure.figsize' in rcparams else (9, 9.5) - rcparams['figure.dpi'] = rcparams['figure.dpi'] if \ - 'figure.dpi' in rcparams else 250 - rcparams['text.usetex'] = rcparams['text.usetex'] if \ - 'text.usetex' in rcparams else True - rcparams['axes.linewidth'] = rcparams['axes.linewidth'] if \ - 'axes.linewidth' in rcparams else 0.5 - rcparams['axes.grid'] = rcparams['axes.grid'] if \ - 'axes.grid' in rcparams else False - rcparams['font.family'] = rcparams['font.family'] if \ - 'font.family' in rcparams else 'sans-serif' - rcparams['font.sans-serif'] = rcparams['font.sans-serif'] if \ - 'font.sans-serif' in rcparams else \ - 'Avant Garde, Helvetica, Computer Modern Sans serif' - rcparams['font.size'] = rcparams['font.size'] if \ - 'font.size' in rcparams else 20 - rcparams['legend.fontsize'] = rcparams['legend.fontsize'] if \ - 'legend.fontsize' in rcparams else 16 - rcparams['legend.frameon'] = rcparams['legend.frameon'] if \ - 'legend.frameon' in rcparams else False - - mpl.rcParams.update(rcparams) - - fig, ax = pl.subplots() - - # extract periods and period derivatives - periods = table['P0'] - pdots = table['P1'] - if intrinsicpdot: # use instrinsic period derivatives if requested - ipdotidx = np.isfinite(table['P1_I']) - pdots[ipdotidx] = table['P1_I'][ipdotidx] - - # get only finite values - pidx = (np.isfinite(periods)) & (np.isfinite(pdots)) - periods = periods[pidx] - pdots = pdots[pidx] + self.condition = "P1 > 0" + usecondition = True # must use this condition - if 'ASSOC' in table.columns: - assocs = table['ASSOC'][pidx] # associations - if 'TYPE' in table.columns: - types = table['TYPE'][pidx] # pulsar types - if 'BINARY' in nshowtypes: - binaries = table['BINARY'][pidx] # binary pulsars + if usealtair: + # import altair + try: + import altair + except ImportError: + raise ImportError('Cannot create plot using altair as it is ' + 'not available') + fig = self.plot( + "P0", "P1", + usealtair=True, + usepsrs=usepsrs, + usecondition=usecondition, + xlabel="Period (s)", + ylabel="Period Derivative", + logx=True, + logy=True, + intrinsicpdot=intrinsicpdot + ) + else: + figax = self.plot( + "P0", "P1", + usepsrs=usepsrs, + usecondition=usecondition, + xlabel="Period (s)", + ylabel="Period Derivative", + logx=True, + logy=True, + intrinsicpdot=intrinsicpdot, + xlims=periodlims, + ylims=pdotlims + ) + + if figax is None: + return None + else: + fig, ax = figax + + if periodlims is None: + periodlims = ax.get_xlims() + + if pdotlims is None: + pdotlims = ax.get_ylims() + + if deathline: + deathpdots = 10**death_line(np.log10(periodlims), + linemodel=deathmodel) + ax.loglog(periodlims, deathpdots, 'k--', linewidth=0.5) + + if filldeath: + if not filldeathtype: + filldeathtype = {} + + filldeathtype['linestyle'] = filldeathtype['linestyle'] if \ + 'linestyle' in filldeathtype else '-' + filldeathtype['alpha'] = filldeathtype['alpha'] if \ + 'alpha' in filldeathtype else 0.15 + filldeathtype['facecolor'] = filldeathtype['facecolor'] if \ + 'facecolor' in filldeathtype else 'darkorange' + filldeathtype['hatch'] = filldeathtype['hatch'] if \ + 'hatch' in filldeathtype else '' + ax.fill_between(periodlims, deathpdots, pdotlims[0], + **filldeathtype) + + # add markers for each pulsar type + if not markertypes: + markertypes = {} + + # check if markers have been defined by the user or not + markertypes['AXP'] = {'marker': 's', 'markeredgecolor': 'red'} if \ + 'AXP' not in markertypes else markertypes['AXP'] + markertypes['BINARY'] = {'marker': 'o', 'markeredgecolor': 'grey'} if \ + 'BINARY' not in markertypes else markertypes['BINARY'] + markertypes['HE'] = {'marker': 'D', 'markeredgecolor': 'orange'} if \ + 'HE' not in markertypes else markertypes['HE'] + markertypes['RRAT'] = {'marker': 'h', 'markeredgecolor': 'green'} if \ + 'RRAT' not in markertypes else markertypes['RRAT'] + markertypes['NRAD'] = {'marker': 'v', 'markeredgecolor': 'blue'} if \ + 'NRAD' not in markertypes else markertypes['NRAD'] + markertypes['XINS'] = {'marker': '^', 'markeredgecolor': 'magenta'} if \ + 'XINS' not in markertypes else markertypes['XINS'] + markertypes['GC'] = {'marker': '8', 'markeredgecolor': 'cyan'} if \ + 'GC' not in markertypes else markertypes['GC'] + markertypes['SNR'] = {'marker': '*', 'markeredgecolor': 'darkorchid'} if \ + 'SNR' not in markertypes else markertypes['SNR'] + + # legend strings for different types + typelegstring = {} + typelegstring['AXP'] = r'SGR/AXP' + typelegstring['NRAD'] = r'"Radio-Quiet"' + typelegstring['XINS'] = r'Pulsed Thermal X-ray' + typelegstring['BINARY'] = r'Binary' + typelegstring['HE'] = r'Radio-IR Emission' + typelegstring['GC'] = r'Globular Cluster' + typelegstring['SNR'] = r'SNR' + typelegstring['RRAT'] = r'RRAT' + + # show globular cluster pulsars + if showGCs and not excludeGCs: + nshowtypes.append('GC') + + # show pulsars with associated supernova remnants + if showSNRs: + nshowtypes.append('SNR') + + handles = OrderedDict() + + for stype in nshowtypes: + if stype.upper() in PSR_TYPE + ['GC', 'SNR']: + thistype = stype.upper() + if thistype == 'BINARY': + # for binaries used the 'BINARY' column in the table + typeidx = np.flatnonzero(~binaries.mask) + elif thistype in ['GC', 'SNR']: + typeidx = np.flatnonzero( + np.char.find(np.array(assocs.tolist(), + dtype=np.str), thistype) != -1) + else: + typeidx = np.flatnonzero( + np.char.find(np.array(types.tolist(), + dtype=np.str), thistype) != -1) + + if len(typeidx) == 0: + continue + + # default to empty markers with no lines between them + if 'markerfacecolor' not in markertypes[thistype]: + markertypes[thistype]['markerfacecolor'] = 'none' + if 'linestyle' not in markertypes[thistype]: + markertypes[thistype]['linestyle'] = 'none' + typehandle, = ax.loglog(periods[typeidx], pdots[typeidx], + label=typelegstring[thistype], + **markertypes[thistype]) + if thistype in typelegstring: + handles[typelegstring[thistype]] = typehandle + else: + handles[thistype] = typehandle - # now get only positive pdot values - pidx = pdots > 0. - periods = periods[pidx] - pdots = pdots[pidx] - if 'ASSOC' in table.columns: - assocs = assocs[pidx] # associations - if 'TYPE' in table.columns: - types = types[pidx] # pulsar types - if 'BINARY' in nshowtypes: - binaries = binaries[pidx] # binary pulsars + ax.legend(handles.values(), handles.keys(), loc='upper left', + numpoints=1) - # check whether to exclude globular cluster pulsars that could have - # contaminated spin-down value - if excludeGCs: - # use '!=' to find GC indexes - nongcidxs = np.flatnonzero( - np.char.find(np.array(assocs.tolist(), - dtype=np.str), 'GC:') == -1) - periods = periods[nongcidxs] - pdots = pdots[nongcidxs] - if 'ASSOC' in table.columns: - assocs = assocs[nongcidxs] - if 'TYPE' in table.columns: - types = types[nongcidxs] - if 'BINARY' in nshowtypes: - binaries = binaries[nongcidxs] - - # plot pulsars - ax.loglog(periods, pdots, marker='.', color='dimgrey', - linestyle='none') - ax.set_xlabel(r'Period (s)') - ax.set_ylabel(r'Period Derivative') - - # get limits - if periodlims is None: - periodlims = [10**np.floor(np.min(np.log10(periods))), - 10.*int(np.ceil(np.max(pdots)/10.))] - if pdotlims is None: - pdotlims = [10**np.floor(np.min(np.log10(pdots))), - 10**np.ceil(np.max(np.log10(pdots)))] - ax.set_xlim(periodlims) - ax.set_ylim(pdotlims) - - if deathline: - deathpdots = 10**death_line(np.log10(periodlims), - linemodel=deathmodel) - ax.loglog(periodlims, deathpdots, 'k--', linewidth=0.5) - - if filldeath: - if not filldeathtype: - filldeathtype = {} - - filldeathtype['linestyle'] = filldeathtype['linestyle'] if \ - 'linestyle' in filldeathtype else '-' - filldeathtype['alpha'] = filldeathtype['alpha'] if \ - 'alpha' in filldeathtype else 0.15 - filldeathtype['facecolor'] = filldeathtype['facecolor'] if \ - 'facecolor' in filldeathtype else 'darkorange' - filldeathtype['hatch'] = filldeathtype['hatch'] if \ - 'hatch' in filldeathtype else '' - ax.fill_between(periodlims, deathpdots, pdotlims[0], - **filldeathtype) - - # add markers for each pulsar type - if not markertypes: - markertypes = {} - - # check if markers have been defined by the user or not - markertypes['AXP'] = {'marker': 's', 'markeredgecolor': 'red'} if \ - 'AXP' not in markertypes else markertypes['AXP'] - markertypes['BINARY'] = {'marker': 'o', 'markeredgecolor': 'grey'} if \ - 'BINARY' not in markertypes else markertypes['BINARY'] - markertypes['HE'] = {'marker': 'D', 'markeredgecolor': 'orange'} if \ - 'HE' not in markertypes else markertypes['HE'] - markertypes['RRAT'] = {'marker': 'h', 'markeredgecolor': 'green'} if \ - 'RRAT' not in markertypes else markertypes['RRAT'] - markertypes['NRAD'] = {'marker': 'v', 'markeredgecolor': 'blue'} if \ - 'NRAD' not in markertypes else markertypes['NRAD'] - markertypes['XINS'] = {'marker': '^', 'markeredgecolor': 'magenta'} if \ - 'XINS' not in markertypes else markertypes['XINS'] - markertypes['GC'] = {'marker': '8', 'markeredgecolor': 'cyan'} if \ - 'GC' not in markertypes else markertypes['GC'] - markertypes['SNR'] = {'marker': '*', 'markeredgecolor': 'darkorchid'} if \ - 'SNR' not in markertypes else markertypes['SNR'] - - # legend strings for different types - typelegstring = {} - typelegstring['AXP'] = r'SGR/AXP' - typelegstring['NRAD'] = r'"Radio-Quiet"' - typelegstring['XINS'] = r'Pulsed Thermal X-ray' - typelegstring['BINARY'] = r'Binary' - typelegstring['HE'] = r'Radio-IR Emission' - typelegstring['GC'] = r'Globular Cluster' - typelegstring['SNR'] = r'SNR' - typelegstring['RRAT'] = r'RRAT' - - # show globular cluster pulsars - if showGCs and not excludeGCs: - nshowtypes.append('GC') - - # show pulsars with associated supernova remnants - if showSNRs: - nshowtypes.append('SNR') - - handles = OrderedDict() - - for stype in nshowtypes: - if stype.upper() in PSR_TYPE + ['GC', 'SNR']: - thistype = stype.upper() - if thistype == 'BINARY': - # for binaries used the 'BINARY' column in the table - typeidx = np.flatnonzero(~binaries.mask) - elif thistype in ['GC', 'SNR']: - typeidx = np.flatnonzero( - np.char.find(np.array(assocs.tolist(), - dtype=np.str), thistype) != -1) + # add characteristic age lines + tlines = OrderedDict() + if showtau: + if tau is None: + taus = [1e5, 1e6, 1e7, 1e8, 1e9] # default characteristic ages else: - typeidx = np.flatnonzero( - np.char.find(np.array(types.tolist(), - dtype=np.str), thistype) != -1) - - if len(typeidx) == 0: - continue - - # default to empty markers with no lines between them - if 'markerfacecolor' not in markertypes[thistype]: - markertypes[thistype]['markerfacecolor'] = 'none' - if 'linestyle' not in markertypes[thistype]: - markertypes[thistype]['linestyle'] = 'none' - typehandle, = ax.loglog(periods[typeidx], pdots[typeidx], - label=typelegstring[thistype], - **markertypes[thistype]) - if thistype in typelegstring: - handles[typelegstring[thistype]] = typehandle + taus = tau + + nbrake = brakingidx + for tauv in taus: + pdots_tc = age_pdot(periodlims, tau=tauv, braking_idx=nbrake) + tline, = ax.loglog(periodlims, pdots_tc, 'k-.', linewidth=0.5) + # check if taus are powers of 10 + taupow = np.floor(np.log10(tauv)) + numv = tauv/10**taupow + if numv == 1.: + tlines[r'$10^{{{0:d}}}\,{{\rm yr}}$'.format(int(taupow))] = tline + else: + tlines[r'${{{0:.1f}}}!\times\!10^{{{1:d}}}\,{{\rm yr}}$' + .format(numv, taupow)] = tline + + # add magnetic field lines + Blines = OrderedDict() + if showB: + if Bfield is None: + Bs = [1e10, 1e11, 1e12, 1e13, 1e14] else: - handles[thistype] = typehandle - - ax.legend(handles.values(), handles.keys(), loc='upper left', - numpoints=1) + Bs = Bfield + for B in Bs: + pdots_B = B_field_pdot(periodlims, Bfield=B) + bline, = ax.loglog(periodlims, pdots_B, 'k:', linewidth=0.5) + # check if Bs are powers of 10 + Bpow = np.floor(np.log10(B)) + numv = B/10**Bpow + if numv == 1.: + Blines[r'$10^{{{0:d}}}\,{{\rm G}}$'.format(int(Bpow))] = bline + else: + Blines[r'${{{0:.1f}}}!\times\!10^{{{1:d}}}\,{{\rm G}}$' + .format(numv, Bpow)] = bline - # add characteristic age lines - tlines = OrderedDict() - if showtau: - if tau is None: - taus = [1e5, 1e6, 1e7, 1e8, 1e9] # default characteristic ages - else: - taus = tau - - nbrake = brakingidx - for tauv in taus: - pdots_tc = age_pdot(periodlims, tau=tauv, braking_idx=nbrake) - tline, = ax.loglog(periodlims, pdots_tc, 'k-.', linewidth=0.5) - # check if taus are powers of 10 - taupow = np.floor(np.log10(tauv)) - numv = tauv/10**taupow - if numv == 1.: - tlines[r'$10^{{{0:d}}}\,{{\rm yr}}$'.format(int(taupow))] = tline - else: - tlines[r'${{{0:.1f}}}!\times\!10^{{{1:d}}}\,{{\rm yr}}$' - .format(numv, taupow)] = tline - - # add magnetic field lines - Blines = OrderedDict() - if showB: - if Bfield is None: - Bs = [1e10, 1e11, 1e12, 1e13, 1e14] - else: - Bs = Bfield - for B in Bs: - pdots_B = B_field_pdot(periodlims, Bfield=B) - bline, = ax.loglog(periodlims, pdots_B, 'k:', linewidth=0.5) - # check if Bs are powers of 10 - Bpow = np.floor(np.log10(B)) - numv = B/10**Bpow - if numv == 1.: - Blines[r'$10^{{{0:d}}}\,{{\rm G}}$'.format(int(Bpow))] = bline - else: - Blines[r'${{{0:.1f}}}!\times\!10^{{{1:d}}}\,{{\rm G}}$' - .format(numv, Bpow)] = bline + fig.tight_layout() - fig.tight_layout() + # add text for characteristic age lines and magnetic field strength lines + for l in tlines: + _ = label_line(ax, tlines[l], l, color='k', fs=18, frachoffset=0.05) - # add text for characteristic age lines and magnetic field strength lines - for l in tlines: - _ = label_line(ax, tlines[l], l, color='k', fs=18, frachoffset=0.05) + for l in Blines: + _ = label_line(ax, Blines[l], l, color='k', fs=18, frachoffset=0.90) - for l in Blines: - _ = label_line(ax, Blines[l], l, color='k', fs=18, frachoffset=0.90) + # reset original conditions + self.condition = orig_condition - # return the figure return fig