diff --git a/tardis/montecarlo/spectrum.py b/tardis/montecarlo/spectrum.py index 551490c8091..a97fc77f1dd 100644 --- a/tardis/montecarlo/spectrum.py +++ b/tardis/montecarlo/spectrum.py @@ -37,16 +37,19 @@ def __init__(self, _frequency, luminosity): self._frequency = _frequency.to("Hz", u.spectral()) self.luminosity = luminosity.to("erg / s") + l_nu_unit = u.def_unit('erg\ s^{-1}\ Hz^{-1}', u.Unit('erg/(s Hz)')) + l_lambda_unit = u.def_unit('erg\ s^{-1}\ \\AA^{-1}', u.Unit('erg/(s AA)')) + self.frequency = self._frequency[:-1] self.delta_frequency = self._frequency[1] - self._frequency[0] self.wavelength = self.frequency.to("angstrom", u.spectral()) self.luminosity_density_nu = ( self.luminosity / self.delta_frequency - ).to("erg / (s Hz)") + ).to(l_nu_unit) self.luminosity_density_lambda = self.f_nu_to_f_lambda( self.luminosity_density_nu, - ) + ).to(l_lambda_unit) @property def flux_nu(self): @@ -100,19 +103,34 @@ def f_nu_to_f_lambda(self, f_nu): return f_nu * self.frequency / self.wavelength def plot(self, ax=None, mode="wavelength", **kwargs): + """plots the TARDISSpectrum + + Parameters + ---------- + ax : matplotlib.axes.Axes, optional + axes to plot on, by default None + mode : str, optional + plot against wavelength or frequency, by default "wavelength" + """ if ax is None: from matplotlib.pyplot import gca ax = gca() if mode == "wavelength": - ax.plot( - self.wavelength.value, - self.luminosity_density_lambda.value, - **kwargs, + ax.plot(self.wavelength.value, self.luminosity_density_lambda.value, **kwargs) + ax.set_xlabel( + f"Wavelength [{self.wavelength.unit.to_string('latex_inline')}]" + ) + ax.set_ylabel( + f"$L_\\lambda$ [{self.luminosity_density_lambda.unit.to_string('latex_inline')}]" + ) + elif mode == "frequency": + ax.plot(self.frequency.value, self.luminosity_density_nu.value, **kwargs) + ax.set_xlabel( + f"Frequency [{self.frequency.unit.to_string('latex_inline')}]" ) - ax.set_xlabel(f"Wavelength [{self.wavelength.unit._repr_latex_()}]") ax.set_ylabel( - f"Flux [{self.luminosity_density_lambda.unit._repr_latex_():s}]" + f"$L_\\nu$ [{self.luminosity_density_nu.unit.to_string('latex_inline')}]" ) else: warnings.warn(f"Did not find plotting mode {mode}, doing nothing.")