original_docstring: Compute the pattern speed of a snapshot. Uses continuity equation in cylindrical coordinates. Integrating over
      all z and taking fourier transform in angular direction gives i m R F(Sigma) = d/dR{F(Sigma vR)} + i m F(sigma vT)
      where F(y) is the angular fourier transform of y. This provides an equation at each radius and m for Sigma. By
      default combine m=2,4,6,8,10 over the radial range 1->4 into one measurement.
docstring: Compute the pattern speed of a snapshot. Uses continuity equation in cylindrical coordinates. Integrating over
  all z and taking fourier transform in angular direction gives i m R F(Sigma) = d/dR{F(Sigma vR)} + i m F(sigma vT)
  where F(y) is the angular fourier transform of y. This provides an equation at each radius and m for Sigma. By
  default combine m=2,4,6,8,10 over the radial range 1->4 into one measurement.
docstring_tokens: Compute the pattern speed of a snapshot . Uses continuity equation in cylindrical coordinates . Integrating over all z and taking fourier transform in angular direction gives i m R F ( Sigma ) = d / dR { F ( Sigma vR ) } + i m F ( sigma vT ) where F ( y ) is the angular fourier transform of y . This provides an equation at each radius and m for Sigma . By default combine m = 2 4 6 8 10 over the radial range 1 - > 4 into one measurement .
code: def patternspeed(snap, rrange=(1, 4), n=range(2, 12, 2), combine=True,
                   plot=None):
      """Compute the pattern speed of a snapshot. Uses continuity equation in cylindrical coordinates. Integrating over
      all z and taking fourier transform in angular direction gives i m R F(Sigma) = d/dR{F(Sigma vR)} + i m F(sigma vT)
      where F(y) is the angular fourier transform of y. This provides an equation at each radius and m for Sigma. By
      default combine m=2,4,6,8,10 over the radial range 1->4 into one measurement."""
      rcyl = torch.norm(snap.positions[:, 0:2], dim=-1)
      vr = torch.einsum('...i,...i->...', snap.positions[:, 0:2],
                        snap.velocities[:, 0:2]) / rcyl
      vt = (snap.positions[:, 1] * snap.velocities[:, 0] -
            snap.positions[:, 0] * snap.velocities[:, 1]) / rcyl
  
      rbins = np.linspace(0., 10, 50)
      rmid, (surfdensft, surfdensrvrft, surfdensvtft) = bar_cyl_fft(snap,
                                                                    rbins=rbins,
                                                                    weights=(None,
                                                                             vr * rcyl,
                                                                             vt))
      dr = rbins[1] - rbins[0]
      dsurfdensrvr_drft = np.zeros_like(surfdensrvrft)
      dsurfdensrvr_drft[0, :] = (surfdensrvrft[1, :] - surfdensrvrft[0, :]) / dr
      dsurfdensrvr_drft[-1, :] = (surfdensrvrft[-1, :] - surfdensrvrft[-2, :]) / dr
      dsurfdensrvr_drft[1:-1, :] = (surfdensrvrft[2:, :] - surfdensrvrft[:-2]) / (
                  2 * dr)
  
      omegas = []
      omegaerrs = []
      idx_good_r = (rmid > rrange[0]) & (rmid < rrange[1])
      for i in n:
          omega = surfdensvtft[:, i] / surfdensft[:, i] / rmid + \
                  1j * dsurfdensrvr_drft[:, i] / surfdensft[:, i] / i / rmid
          omegas.append(np.abs(np.mean(omega[idx_good_r])))
          omegaerrs.append(
              np.std(np.abs(omega[idx_good_r])) / np.sqrt(np.sum(idx_good_r)))
      omegas, omegaerrs = np.array(omegas), np.array(omegaerrs)
  
      # combine by weighted mean
      omega = np.sum(omegas / omegaerrs ** 2) / np.sum(1 / omegaerrs ** 2)
      omegaerr = 1 / np.sqrt(np.sum(1 / omegaerrs ** 2))
  
      if plot is not None:
          omega, omegaerr = np.sum(omegas / omegaerrs ** 2) / np.sum(
              1 / omegaerrs ** 2), 1 / np.sqrt(
              np.sum(1 / omegaerrs ** 2))
          plot.errorbar(np.arange(2, 12, 2), y=omegas, yerr=omegaerrs, fmt='ko',
                        markersize=3)
          plot.axhline(y=omega, color='r')
          plot.axhspan(ymin=omega - omegaerr, ymax=omega + omegaerr, color='r',
                       alpha=0.2)
          plot.set_ylim([omega - 10 * omegaerr, omega + 10 * omegaerr])
          plot.set_ylabel(r'$\Omega$')
          plot.set_xlabel('Mode $m$')
      if not combine:
          return omegas, omegaerrs
      else:
          return omega, omegaerr
code_tokens: def patternspeed ( snap , rrange = ( 1 , 4 ) , n = range ( 2 , 12 , 2 ) , combine = True , plot = None ) : rcyl = torch . norm ( snap . positions [ : , 0 : 2 ] , dim = - 1 ) vr = torch . einsum ( '...i,...i->...' , snap . positions [ : , 0 : 2 ] , snap . velocities [ : , 0 : 2 ] ) / rcyl vt = ( snap . positions [ : , 1 ] * snap . velocities [ : , 0 ] - snap . positions [ : , 0 ] * snap . velocities [ : , 1 ] ) / rcyl rbins = np . linspace ( 0. , 10 , 50 ) rmid , ( surfdensft , surfdensrvrft , surfdensvtft ) = bar_cyl_fft ( snap , rbins = rbins , weights = ( None , vr * rcyl , vt ) ) dr = rbins [ 1 ] - rbins [ 0 ] dsurfdensrvr_drft = np . zeros_like ( surfdensrvrft ) dsurfdensrvr_drft [ 0 , : ] = ( surfdensrvrft [ 1 , : ] - surfdensrvrft [ 0 , : ] ) / dr dsurfdensrvr_drft [ - 1 , : ] = ( surfdensrvrft [ - 1 , : ] - surfdensrvrft [ - 2 , : ] ) / dr dsurfdensrvr_drft [ 1 : - 1 , : ] = ( surfdensrvrft [ 2 : , : ] - surfdensrvrft [ : - 2 ] ) / ( 2 * dr ) omegas = [ ] omegaerrs = [ ] idx_good_r = ( rmid > rrange [ 0 ] ) & ( rmid < rrange [ 1 ] ) for i in n : omega = surfdensvtft [ : , i ] / surfdensft [ : , i ] / rmid + 1j * dsurfdensrvr_drft [ : , i ] / surfdensft [ : , i ] / i / rmid omegas . append ( np . abs ( np . mean ( omega [ idx_good_r ] ) ) ) omegaerrs . append ( np . std ( np . abs ( omega [ idx_good_r ] ) ) / np . sqrt ( np . sum ( idx_good_r ) ) ) omegas , omegaerrs = np . array ( omegas ) , np . array ( omegaerrs ) omega = np . sum ( omegas / omegaerrs ** 2 ) / np . sum ( 1 / omegaerrs ** 2 ) omegaerr = 1 / np . sqrt ( np . sum ( 1 / omegaerrs ** 2 ) ) if plot is not None : omega , omegaerr = np . sum ( omegas / omegaerrs ** 2 ) / np . sum ( 1 / omegaerrs ** 2 ) , 1 / np . sqrt ( np . sum ( 1 / omegaerrs ** 2 ) ) plot . errorbar ( np . arange ( 2 , 12 , 2 ) , y = omegas , yerr = omegaerrs , fmt = 'ko' , markersize = 3 ) plot . axhline ( y = omega , color = 'r' ) plot . axhspan ( ymin = omega - omegaerr , ymax = omega + omegaerr , color = 'r' , alpha = 0.2 ) plot . set_ylim ( [ omega - 10 * omegaerr , omega + 10 * omegaerr ] ) plot . set_ylabel ( r'$\Omega$' ) plot . set_xlabel ( 'Mode $m$' ) if not combine : return omegas , omegaerrs else : return omega , omegaerr
short_docstring: Compute the pattern speed of a snapshot.
short_docstring_tokens: Compute the pattern speed of a snapshot .
comment: """Compute the pattern speed of a snapshot. Uses continuity equation in cylindrical coordinates. Integrating over
      all z and taking fourier transform in angular direction gives i m R F(Sigma) = d/dR{F(Sigma vR)} + i m F(sigma vT)
      where F(y) is the angular fourier transform of y. This provides an equation at each radius and m for Sigma. By
      default combine m=2,4,6,8,10 over the radial range 1->4 into one measurement."""$SEP$# combine by weighted mean