def setup_units(m, x): """ Sets up units for a ChaNGa simulation, defined by position and mass units. * time unit = :math:`\\sqrt{L_{unit}^3/G M_{unit}}` * velocity unit = :math:`L_{unit}/T_{unit}` * temperature units = Kelvin Parameters ---------- m, x : Unit, string, or SimArray Define the mass and position units Returns ------- units : dict A dictionary containing the units """ m_unit = get_units(m) pos_unit = get_units(x) # time units are sqrt(L^3/GM) t_unit = np.sqrt((pos_unit**3)*np.power((G*m_unit), -1)).units # velocity units are L/t v_unit = (pos_unit/t_unit).ratio('km s**-1') # Make it a unit v_unit = pynbody.units.Unit('{0} km s**-1'.format(v_unit)) # Temperature (just a hard-coded default) temp_unit = get_units('K') units = {'m': m_unit, 't': t_unit, 'v': v_unit, 'x': pos_unit, 'temp': temp_unit} return units
def setup_units(m, x): """ Sets up units for a ChaNGa simulation, defined by position and mass units. * time unit = :math:`\\sqrt{L_{unit}^3/G M_{unit}}` * velocity unit = :math:`L_{unit}/T_{unit}` * temperature units = Kelvin Parameters ---------- m, x : Unit, string, or SimArray Define the mass and position units Returns ------- units : dict A dictionary containing the units """ m_unit = get_units(m) pos_unit = get_units(x) # time units are sqrt(L^3/GM) t_unit = np.sqrt((pos_unit**3) * np.power((G * m_unit), -1)).units # velocity units are L/t v_unit = (pos_unit / t_unit).ratio('km s**-1') # Make it a unit v_unit = pynbody.units.Unit('{0} km s**-1'.format(v_unit)) # Temperature (just a hard-coded default) temp_unit = get_units('K') units = { 'm': m_unit, 't': t_unit, 'v': v_unit, 'x': pos_unit, 'temp': temp_unit } return units
def waterfall(power, t=None, m=None, log=True, normalize=True, cmap='cubehelix', vmin=None, vmax=None, colorbar=True): """ Generates a waterfall plot for power spectrum vs time for a simulation Parameters ---------- power : array, SimArray 2D shape (nt, nm) array of power spectrum t : array, SimArray (optional) time points. 1D array m : array, SimArray (optional) Fourier mode numbers. 1D array log: bool logscale the colors normalize : bool Normalize by the DC component at each time step cmap : str or colormap colormap to use vmin, vmax : float limits colorbar : bool Display colorbar Examples -------- >>> flist = diskpy.pychanga.get_fnames('snapshot') >>> m, power = diskpy.disk.powerspectrum_t(flist, spacing='log') >>> waterfall(power) """ # Initialize m and t nt, nm = power.shape if m is None: m = np.arange(nm) if t is None: t = np.arange(nt) mMesh, tMesh = np.meshgrid(m, t) # Normalize if normalize: # Power in DC component p0 = power[:, 0, None] # Make a 2D array (same shape as power) p0 = np.dot(p0, np.ones([1, nm])) # normalize power = power/p0 # Set up log-scale for plot if log: norm = LogNorm(vmin, vmax) else: norm = None plt.pcolormesh(tMesh[:, 1:], mMesh[:, 1:], power[:, 1:], norm=norm, \ cmap = cmap) plt.ylim(m[1], m[-1]) plt.ylabel('m') tUnits = get_units(t) if tUnits == 1: plt.xlabel('time step') else: plt.xlabel(' time $(' + tUnits.latex() + ')$') if colorbar: cbar = plt.colorbar() if normalize: text = r'$A_m/A_0$' else: units = get_units(power) text = r'$A_m (' + units.latex() + r') $' cbar.set_label(text)