v1.4.5

build/lib/mwahpy/orbit_fitter.py

@@ -99,15 +99,15 @@ def calc_lam_bet(self, normal, point):
 '''
 
 #do this in one place so it's uniform across the file and easily changeable
-def make_orbit(params, int_ts=None):
+def make_orbit(params, int_ts=None, pot=None):
 
     #negate the x velocity because galpy is in left-handed frame and we aren't barbarians
     o = Orbit(vxvv=[params[0]*u.deg, params[1]*u.deg, params[2]*u.kpc, -1*params[3]*u.km/u.s, params[4]*u.km/u.s, params[5]*u.km/u.s], uvw=True, lb=True, ro=8., vo=220., zo=0., solarmotion=[0, -220, 0])
 
     if int_ts: #use the global ts unless explicitly provided with a different ts
-        o.integrate(int_ts, mwahpy_default_pot)
+        o.integrate(int_ts, pot)
     else:
-        o.integrate(ts, mwahpy_default_pot)
+        o.integrate(ts, pot)
 
     return o
 
@@ -165,14 +165,14 @@ def get_closest_index(val, Lam):
 
 #get_model_from_orbit: data, orbit, vector, vector -> list(int) x3
 #take in data, orbit, and plane info: Output model data corresponding to each data point
-def get_model_from_orbit(data, params):
+def get_model_from_orbit(data, params, pot=None):
     #data: the data that the orbit is being fit to
     #o: the test orbit that we are calculating the goodness-of-fit of
     #normal: the normal vector to the plane of the Great Circle we are estimating for the orbit
     #point: parameter for the axis generation of the Great Circle coordinates
 
     #initialize the orbit we are fitting
-    o = make_orbit(params)
+    o = make_orbit(params, pot=pot)
 
     #we flip it around so that we are fitting both the forwards and the backwards orbit
     o_rev = o.flip()
@@ -221,13 +221,13 @@ def get_model_from_orbit(data, params):
 
 #chi_squared: data, galpy.Orbit() --> float
 #takes in the observed data and a test orbit and calculates the goodness-of-fit using a chi-squared method
-def chi_squared(params, data=[]):
+def chi_squared(params, data=[], pot=None):
     #data: the data that the orbit is being fit to
     #o: the test orbit that we are calculating the goodness-of-fit of
     #normal: the normal vector to the plane of the Great Circle we are estimating for the orbit
     #point: parameter for the axis generation of the Great Circle coordinates
 
-    B_model, d_model, vx_model, vy_model, vz_model, vgsr_model, costs = get_model_from_orbit(data, params) #get model data from orbit
+    B_model, d_model, vx_model, vy_model, vz_model, vgsr_model, costs = get_model_from_orbit(data, params, pot=pot) #get model data from orbit
 
     #B_model sometimes has different length than data.b, no idea why
     #I think it might be a race condition
@@ -282,7 +282,7 @@ def chi_squared(params, data=[]):
 
 #optimize: data -> [float, float, float, float, float], (float, float, float), (float, float, float)
 #takes in data, then fits a Great Circle to that data and minimizes the chi_squared to fit an orbit to the data
-def optimize(data_opt, max_it, bounds, guess, mode, **kwargs):
+def optimize(data_opt, max_it, bounds, guess, mode, method='Nelder-Mead', pot=None, **kwargs):
 
     '''
     ============================================================================
@@ -303,7 +303,7 @@ def optimize(data_opt, max_it, bounds, guess, mode, **kwargs):
             print(RuntimeWarning('Keyword `bounds` was not provided in `fit_orbit()` (required for differential evolution): using default bounds of [(0, 360), (-90, 90), (0, 100), (-500, 500), (-500, 500), (-500, 500)].'))
             bounds = [(0, 360), (-90, 90), (0, 100), (-500, 500), (-500, 500), (-500, 500)]
 
-        params = scopt.differential_evolution(chi_squared, bounds, args=(data_opt,), maxiter=max_it, popsize=pop_size, mutation=diff_scaling_factor, recombination=crossover_rate, workers=-1, disp=not(verbose), **kwargs).x
+        params = scopt.differential_evolution(chi_squared, bounds, args=(data_opt, pot,), maxiter=max_it, popsize=pop_size, mutation=diff_scaling_factor, recombination=crossover_rate, workers=-1, disp=not(verbose), **kwargs).x
 
     elif mode == 'gd': #do gradient descent
 
@@ -313,14 +313,14 @@ def optimize(data_opt, max_it, bounds, guess, mode, **kwargs):
 
         #have to do it this way because you can't pass 'bounds' as a kwarg
         if bounds is not None:
-            params = scopt.minimize(chi_squared, guess, args=(data_opt,), bounds=bounds, options={'maxiter':max_it, 'disp':verbose}, **kwargs).x
+            params = scopt.minimize(chi_squared, guess, args=(data_opt, pot,), bounds=bounds, options={'maxiter':max_it, 'disp':verbose}, method=method, **kwargs).x
         else:
-            params = scopt.minimize(chi_squared, guess, args=(data_opt,), options={'maxiter':max_it, 'disp':verbose}, **kwargs).x
+            params = scopt.minimize(chi_squared, guess, args=(data_opt, pot,), options={'maxiter':max_it, 'disp':verbose}, method=method, **kwargs).x
 
     else:
         raise BaseException('`mode` for `fit_orbit()` must be either `de` for differential evolution, or `gd` for gradient descent.')
 
-    x2 = chi_squared(params, data_opt)
+    x2 = chi_squared(params, data=data_opt, pot=pot)
 
     return params, x2
 
@@ -332,7 +332,11 @@ def optimize(data_opt, max_it, bounds, guess, mode, **kwargs):
 
 def fit_orbit(l, b, b_err, d, d_err, vx=None, vy=None, vz=None, vgsr=None, \
                  vx_err=None, vy_err=None, vz_err=None, vgsr_err=None, max_it=100, \
-                 bounds=None, guess=None, t_len=None, mode='de', **kwargs):
+                 bounds=None, guess=None, t_len=None, mode='de', pot=None, **kwargs):
+
+    if pot is None:
+        print(RuntimeWarning("Didn't provide potential for constructing orbits! Using mwahpy_default_pot for potential model."))
+        pot = mwahpy_default_pot
 
     #construct data
     #set proper flags based on input data
@@ -367,7 +371,7 @@ def fit_orbit(l, b, b_err, d, d_err, vx=None, vy=None, vz=None, vgsr=None, \
         print('===================================')
 
     #optimization
-    params, x2 = optimize(data_opt, max_it, bounds, guess, mode, **kwargs)
+    params, x2 = optimize(data_opt, max_it, bounds, guess, mode, pot=pot, **kwargs)
 
     print('===================================')
     print('Params: l, b, d, vx, vy, vz')
@@ -385,9 +389,9 @@ def fit_orbit(l, b, b_err, d, d_err, vx=None, vy=None, vz=None, vgsr=None, \
 ================================================================================
 '''
 
-def plot_orbit_gal(l, b, d, params):
+def plot_orbit_gal(l, b, d, params, pot=None):
 
-    o = make_orbit(params)
+    o = make_orbit(params, pot=pot)
 
     o_rev = o.flip()
     o_rev.integrate(ts, mwahpy_default_pot)
@@ -417,14 +421,14 @@ def plot_orbit_gal(l, b, d, params):
 
     plt.show()
 
-def plot_orbit_icrs(l, b, d, params):
+def plot_orbit_icrs(l, b, d, params, pot=None):
 
     s = SkyCoord(l, b, frame='galactic', unit=(u.deg, u.deg))
     s = s.transform_to('icrs')
     ra = s.ra
     dec = s.dec
 
-    o = make_orbit(params)
+    o = make_orbit(params, pot=pot)
 
     o_rev = o.flip()
     o_rev.integrate(ts, mwahpy_default_pot)
@@ -473,7 +477,7 @@ def test():
     ts = np.linspace(0, 0.25, 1000)*u.Gyr
     sample = np.array([100, 250, 400, 500, 600, 750, 850])
 
-    o = make_orbit([test_o_l, test_o_b, test_o_d, test_o_vx, test_o_vy, test_o_vz])
+    o = make_orbit([test_o_l, test_o_b, test_o_d, test_o_vx, test_o_vy, test_o_vz], pot=mwahpy_default_pot)
 
     l = np.take(o.ll(ts), sample)
     b = np.take(o.bb(ts), sample)
@@ -493,11 +497,11 @@ def test():
     test_orbit_data = OrbitData(l, b, d, None, None, None, None, b_err, d_err, None, None, None, None)
 
     print('Goodness-of-Fit of actual values:')
-    print(chi_squared([test_o_l, test_o_b, test_o_d, test_o_vx, test_o_vy, test_o_vz], data=test_orbit_data))
+    print(chi_squared([test_o_l, test_o_b, test_o_d, test_o_vx, test_o_vy, test_o_vz], data=test_orbit_data, pot=mwahpy_default_pot))
 
-    params, x2 = fit_orbit(l, b, b_err, d, d_err)
-    plot_orbit_gal(l, b, d, params)
-    plot_orbit_icrs(l, b, d, params)
+    params, x2 = fit_orbit(l, b, b_err, d, d_err, pot=mwahpy_default_pot)
+    plot_orbit_gal(l, b, d, params, pot=mwahpy_default_pot)
+    plot_orbit_icrs(l, b, d, params, pot=mwahpy_default_pot)
 
 '''
 ================================================================================

build/lib/mwahpy/timestep.py

@@ -716,7 +716,7 @@ def get_self_energies(t):
 
     #compute the kinetic energy of each particle (within reference frame of
     #the COM's of the particles)
-    kin_energies = t.vx**2 + t.vy**2 + t.vz**2
+    kin_energies = 0.5*(t.vx**2 + t.vy**2 + t.vz**2)
 
     energies = pot_energies + kin_energies
 

changelog.txt

@@ -1,3 +1,15 @@
+v1.4.5: (MINOR PATCH)
+ - fixed bug where orbit fitter plots would use the mwahpy default potential even
+   if another was provided
+ - output_handler.read_output() now works for files with the <bodies> </bodies> tags format
+ - changed name of sky_to_pole() to pole_rotation(), because it was confusing. 
+   sky_to_pole() is now an alias for pole_rotation() and causes a DeprecationWarning
+ - added pole_rotation_inv() function
+ - timestep.scatter() now has a 'cbar' keyword to turn on a colorbar & shade particles appropriately
+ - improved how .coords was checking types
+ - added eq_to_OC() transformation from ra/dec to Orphan-Chenab lambda/beta, 
+   and OC_to_eq() inverse transformation
+
 v1.4.4: (MINOR PATCH)
  - orbit_fitter gradient descent now uses Nelder-Mead optimizer by default, BFGS
    was terminating the optimization early

mwahpy.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: mwahpy
-Version: 1.4.3
+Version: 1.4.4
 Summary: A python package for easily parsing and processing data from MilkyWay@home
 Home-page: https://github.com/thomasdonlon/mwahpy
 Author: Tom Donlon

mwahpy/coords.py

@@ -15,29 +15,29 @@
 #These functions aren't meant to be accessed by outside files or by end users,
 #and as such they are not well documented, named, or tested
 
-def wrap_long(long, rad=False):
+def wrap_long(lon, rad=False):
 
     if rad:
-        long = long * 180/np.pi
+        lon = lon * 180/np.pi
 
     use_array = True
-    if type(long) != type(np.array([])):
+    if not(isinstance(lon, np.ndarray)):
         use_array = False
-        long = np.array([long])
+        lon = np.array([lon])
 
-    for i in range(len(long)):
-        if long[i] < 0:
-            long[i] += 360
-        elif long[i] > 360:
-            long[i] -= 360
+    for i in range(len(lon)):
+        if lon[i] < 0:
+            lon[i] += 360
+        elif lon[i] > 360:
+            lon[i] -= 360
 
     if rad:
-        long = long * np.pi/180
+        lon = lon * np.pi/180
 
     if not use_array:
-        long = long[0]
+        lon = lon[0]
 
-    return long
+    return lon
 
 #rotate the given data around the provided axis
 #TODO: Allow array-like input
@@ -72,7 +72,7 @@ def rot_around_arb_axis(x, y, z, ux, uy, uz, theta):
 def long_lat_to_unit_vec(l, b, left_handed=False, rad=False):
 
     use_array = True
-    if type(l) != type(np.array([])):
+    if not(isinstance(l, np.ndarray)):
         use_array = False
         l = np.array([l])
         b = np.array([b])
@@ -126,7 +126,7 @@ def comp_wise_dot(M, v, normalize=False):
 def cart_to_gal(x, y, z, left_handed=False):
     #get l, b, r (helio) from galactocentric X, Y, Z coords
     use_array = True
-    if type(x) != type(np.array([])):
+    if not(isinstance(x, np.ndarray)):
         use_array = False
         x = np.array([x])
         y = np.array([y])
@@ -152,7 +152,7 @@ def cart_to_gal(x, y, z, left_handed=False):
 def gal_to_cart(l, b, r, left_handed=False, rad=False):
 
     use_array = True
-    if type(l) != type(np.array([])):
+    if not(isinstance(l, np.ndarray)):
         use_array = False
         l = np.array([l])
         b = np.array([b])
@@ -187,7 +187,7 @@ def gal_to_cart(l, b, r, left_handed=False, rad=False):
 def cart_to_cyl(x, y, z):
 
     use_array = True
-    if type(x) != type(np.array([])):
+    if not(isinstance(x, np.ndarray)):
         use_array = False
         x = np.array([x])
         y = np.array([y])
@@ -206,7 +206,7 @@ def cart_to_cyl(x, y, z):
 def cyl_to_cart(R, z, phi):
 
     use_array = True
-    if type(R) != type(np.array([])):
+    if not(isinstance(R, np.ndarray)):
         use_array = False
         R = np.array([R])
         phi = np.array([phi])
@@ -227,7 +227,7 @@ def cyl_to_cart(R, z, phi):
 def cart_to_sph(x, y, z):
 
     use_array = True
-    if type(x) != type(np.array([])):
+    if not(isinstance(x, np.ndarray)):
         use_array = False
         x = np.array([x])
         y = np.array([y])
@@ -250,7 +250,7 @@ def cart_to_sph(x, y, z):
 def sph_to_cart(phi, theta, r):
 
     use_array = True
-    if type(phi) != type(np.array([])):
+    if not(isinstance(phi, np.ndarray)):
         use_array = False
         phi = np.array([phi])
         theta = np.array([theta])
@@ -343,7 +343,7 @@ def get_vxvyvz(ra, dec, dist, rv, pmra, pmde):
 def get_rvpm(ra, dec, dist, U, V, W):
 
     use_array = True
-    if type(ra) != type(np.array([])):
+    if not(isinstance(ra, np.ndarray)):
         use_array = False
         ra = np.array([ra])
         dec = np.array([dec])
@@ -438,7 +438,7 @@ def vlos_to_vgsr(l, b, vlos):
     #TODO: Allow ra, dec as inputs
 
     use_array = True
-    if type(l) != type(np.array([])):
+    if not(isinstance(l, np.ndarray)):
         use_array = False
         l = np.array([l])
         b = np.array([b])
@@ -459,7 +459,7 @@ def vgsr_to_vlos(l, b, vgsr):
     #TODO: Allow ra, dec as inputs
 
     use_array = True
-    if type(l) != type(np.array([])):
+    if not(isinstance(l, np.ndarray)):
         use_array = False
         l = np.array([l])
         b = np.array([b])
@@ -724,7 +724,7 @@ def sgr_to_gal(Lam, Bet):
 
 #-------------------------------------------------------------------------------
 
-#sky_to_pole: array(float), array(float), tuple(float, float), tuple(float, float) -> array(float), array(float)
+#pole_rotation: array(float), array(float), tuple(float, float), tuple(float, float) -> array(float), array(float)
 #rotate the positions on the sky (sky1, sky2) into a new frame determined by
 #the pole of the new frame (pole1, pole2) and the origin of the new frame (origin1, origin2)
 #The output is the longitude and latitude in the new coordinate frame
@@ -734,7 +734,9 @@ def sgr_to_gal(Lam, Bet):
 #e.g. if sky1 is a list of RAs and sky2 is a list of Decs, then
 #the pole and origin arguments must also be specified in RA Dec.
 
-def sky_to_pole(sky1, sky2, pole, origin, wrap=False, rad=False):
+#sky_to_pole() is deprecated but it is an alias of this function
+
+def pole_rotation(sky1, sky2, pole, origin, wrap=False, rad=False):
     #sky1, sky2: positions of the data on the sky (e.g. sky1 = array(RA), sky2 = array(Dec), etc.)
     #pole: position of the pole of the new coordinate system, tuple
     #origin: position of the origin of the new coordinate system, tuple
@@ -801,6 +803,36 @@ def sky_to_pole(sky1, sky2, pole, origin, wrap=False, rad=False):
 
     return Lam, Bet
 
+#given the origin and pole for a pole_rotation (in Eq, galactic, etc.), and the new Lambda/Beta,
+#returns the sky coordinates in whatever coordinates pole and origin were specified in
+def pole_rotation_inv(Lam, Bet, pole, origin, **kwargs):
+
+    new_lon, new_lat = pole_rotation(np.array([0, 0]), np.array([90, 0]), pole, origin, **kwargs)
+    new_pole = (new_lon[0], new_lat[0])
+    new_origin = (new_lon[1], new_lat[1])
+    sky1, sky2 = pole_rotation(Lam, Bet, new_pole, new_origin, **kwargs)
+
+    return sky1, sky2
+
+def sky_to_pole(sky1, sky2, pole, origin, wrap=False, rad=False):
+
+    print(DeprecationWarning('sky_to_pole() is deprecated as of mwahpy v1.4.5. Please use pole_rotation() instead (sky_to_pole is now an alias for this). In future updates, this function may be removed.'))
+
+    return pole_rotation(sky1, sky2, pole, origin, wrap=False, rad=False)
+
+#-------------------------------------------------------------------------------
+#hard-coded pole_rotation transformations
+
+oc_pole = (72, -14) #koposov et al. (2019) orphan-chenab stream transformation (ra, dec)
+oc_origin = (191.10487, 62.86084)
+
+def eq_to_OC(ra, dec, **kwargs):
+    return pole_rotation(ra, dec, oc_pole, oc_origin, **kwargs)
+
+def OC_to_eq(Lam, Bet, **kwargs):
+    return pole_rotation_inv(Lam, Bet, oc_pole, oc_origin, **kwargs)
+
+
 #===============================================================================
 # RUNTIME
 #===============================================================================