def _rotate_polygon(lon, lat, lon0, lat0): """ Given a polygon with vertices defined by (lon, lat), rotate the polygon such that the North pole of the spherical coordinates is now at (lon0, lat0). Therefore, to end up with a polygon centered on (lon0, lat0), the polygon should initially be drawn around the North pole. """ # Create a representation object polygon = UnitSphericalRepresentation(lon=lon, lat=lat) # Determine rotation matrix to make it so that the circle is centered # on the correct longitude/latitude. m1 = rotation_matrix(-(0.5 * np.pi * u.radian - lat0), axis='y') m2 = rotation_matrix(-lon0, axis='z') transform_matrix = m2 * m1 # Apply 3D rotation polygon = polygon.to_cartesian() try: polygon = polygon.transform(transform_matrix) except: # TODO: remove once Astropy 1.1 is no longer supported polygon = _transform_cartesian(polygon, transform_matrix) polygon = UnitSphericalRepresentation.from_cartesian(polygon) return polygon.lon, polygon.lat
def align_particles(ws, ix=-1): # assumes ws[:,0] is the "progenitor" orbit new_cen_x = ws[:,0,:3].copy() new_cen_v = ws[:,0,3:].copy() end_w = ws[ix] # put endpoint on x axis in x-z plane # first about y theta = np.arctan2(new_cen_x[ix,2],new_cen_x[ix,0]) * u.radian R1 = rotation_matrix(-theta, 'y') new_cen_x = R1.dot(new_cen_x.T).T new_cen_v = R1.dot(new_cen_v.T).T # then about z theta = np.arctan2(new_cen_x[ix,1],new_cen_x[ix,0]) * u.radian R2 = rotation_matrix(theta, 'z') new_cen_x = R2.dot(new_cen_x.T).T new_cen_v = R2.dot(new_cen_v.T).T # now align L L = np.cross(new_cen_x[ix], new_cen_v[ix])[0] theta = np.arccos(L[2] / np.sqrt(np.sum(L**2))) * u.radian R3 = rotation_matrix(-theta, 'x') new_cen_x = np.array(R3.dot(new_cen_x.T).T) R = R3*R2*R1 new_end_ptcl_x = np.array(R.dot(end_w[:,:3].T).T) return new_cen_x, new_end_ptcl_x
def rotate_w(w, phi, theta): w = np.atleast_1d(w) if w.ndim > 1: raise ValueError("w must be 1D") R_alt = rotation_matrix(theta-90.*u.deg, "y") R_azi = rotation_matrix(phi, "z") r = w[:w.size//2] v = w[w.size//2:] r_rot = np.array(r.dot(R_alt).dot(R_azi)) v_rot = np.array(v.dot(R_alt.T).dot(R_azi.T)) return np.hstack((r_rot, v_rot)).squeeze()
def icrs_to_astrometric(icrs_coord, astrometric_frame): """Convert an ICRS coordinate to an Astrometric frame.""" # Define rotation matricies along the position angle vector, and # relative to the origin. mat2 = rotation_matrix(-astrometric_frame.origin.dec, 'y') mat3 = rotation_matrix(astrometric_frame.origin.ra, 'z') R = mat2 * mat3 xyz = icrs_coord.cartesian.xyz orig_shape = xyz.shape xyz = R.dot(xyz.reshape(xyz.shape[0], np.prod(xyz.shape[1:]))).reshape(orig_shape) representation = CartesianRepresentation(xyz) return astrometric_frame.realize_frame(representation)
def compute_stream_rotation_matrix(coordinate, wxyz0=None, align_lon='mean'): """ Compute the rotation matrix to go from the frame of the input coordinate to closely align the equator with the stream. Parameters ---------- coordinate : :class:`astropy.coordinate.SkyCoord`, :class:`astropy.coordinate.BaseCoordinateFrame` The coordinates of the stream stars. wxyz0 : array_like (optional) Initial guess for the quaternion vector that represents the rotation. align_lon : str, int (optional) Can specify either 'min', 'max', or an integer index. This picks the 'pirvot' star, whose longitude is set to 0. Returns ------- R : :class:`~numpy.ndarray` A 3 by 3 rotation matrix (has shape ``(3,3)``) to convert heliocentric, Cartesian coordinates in the input coordinate frame to stream coordinates. """ if wxyz0 is None: wxyz0 = Quaternion.random().wxyz res = minimize(_rotation_opt_func, x0=wxyz0, args=(coordinate.cartesian.xyz,)) R = Quaternion(res.x).rotation_matrix new_xyz = R.dot(coordinate.cartesian.xyz.value) lon = np.arctan2(new_xyz[1], new_xyz[0]) if align_lon == 'mean': _lon = np.mean(lon) R3 = 1. elif align_lon == 'min': ix = lon.argmin() _lon = lon[ix] R3 = 1. elif align_lon == 'max': ix = lon.argmax() _lon = lon[ix] R3 = rotation_matrix(np.pi*u.radian, 'x') else: ix = int(align_lon) _lon = lon[ix] R3 = 1. R2 = rotation_matrix(_lon*u.radian, 'z') R = R3*R2*R return R
def astrometric_to_icrs(astrometric_coord, icrs_frame): """Convert an OSIRIS frame coordinaate to an ICRS""" # Define rotation matricies along the position angle vector, and # relative to the origin. mat2 = rotation_matrix(-astrometric_coord.origin.dec, 'y') mat3 = rotation_matrix(astrometric_coord.origin.ra, 'z') R = mat2 * mat3 Rinv = np.linalg.inv(R) xyz = astrometric_coord.cartesian.xyz orig_shape = xyz.shape xyz = Rinv.dot(xyz.reshape(xyz.shape[0], np.prod(xyz.shape[1:]))).reshape(orig_shape) representation = CartesianRepresentation(xyz) return icrs_frame.realize_frame(representation)
def main(plot=False): all_data = OphiuchusData() # stars from both 2015a and 2015b data = OphiuchusData(expr="source == 'Sesar2015a'") # just use the original BHB stars # compute the rotation matrix to go from Galacic to Stream coordinates, first without # an extra rotation to put the center at lon=0 np.random.seed(42) R1 = orbitfit.compute_stream_rotation_matrix(data.coord) if plot: # rotate all data to plot rot_rep = orbitfit.rotate_sph_coordinate(all_data.coord, R1) pl.figure() pl.suptitle("Before rotating in longitude", fontsize=18) pl.plot(rot_rep.lon.wrap_at(180*u.deg).degree, rot_rep.lat.degree, ls='none', marker='o', ms=5.) rot_rep = orbitfit.rotate_sph_coordinate(data.coord, R1) # now rotate by the median longitude to put the 'center' of the stream at lon=0 extra_rot = np.median(rot_rep.lon.wrap_at(180*u.deg)) R2 = rotation_matrix(extra_rot, 'z') R = np.asarray(R2*R1) print("Rotation matrix:") for i in range(3): print("{:>20.15f} {:>20.15f} {:>20.15f}".format(*R[i])) all_rot_rep = orbitfit.rotate_sph_coordinate(all_data.coord, R) if plot: pl.figure() pl.suptitle("After rotating in longitude", fontsize=18) pl.plot(all_rot_rep.lon.wrap_at(180*u.deg).degree, all_rot_rep.lat.degree, ls='none', marker='o', ms=5.) pl.show()
def compute_align_matrix(w): """ Given a single phase-space position, compute the rotation matrix that orients the angular momentum with the z axis and places the point along the x axis. Parameters ---------- w : array_like The point to transform. Returns ------- R : :class:`numpy.ndarray` A 2D numpy array (rotation matrix). """ w = _validate_nd_array(w, expected_ndim=1) x = w[:3].copy() v = w[3:].copy() # first rotate about z to put on x-z plane theta = np.arctan2(x[1], x[0]) * u.radian R1 = rotation_matrix(theta, 'z') x = np.asarray(R1.dot(x))[0] v = np.asarray(R1.dot(v))[0] # now rotate about y to put on x axis theta = np.arctan2(x[2], x[0]) * u.radian R2 = rotation_matrix(-theta, 'y') x = np.asarray(R2.dot(x))[0] v = np.asarray(R2.dot(v))[0] # now align L with z axis # theta = np.arccos(L[2] / np.sqrt(np.sum(L**2))) * u.radian L = np.cross(x, v) theta = np.arctan2(L[2], L[1]) * u.radian R3 = rotation_matrix(theta - 90*u.deg, 'x') x = np.asarray(R3.dot(x))[0] v = np.asarray(R3.dot(v))[0] return R3*R2*R1
def _icrs_gctc_velocity_matrix(galactocentric_frame): """ Construct a transformation matrix to go from heliocentric ICRS to a galactocentric frame. This is just a rotation and tilt which makes it approximately the same as transforming to Galactic coordinates. This only works for velocity because there is no shift due to the position of the Sun. """ # define rotation matrix to align x(ICRS) with the vector to the Galactic center M1 = rotation_matrix(-galactocentric_frame.galcen_dec, 'y') M2 = rotation_matrix(galactocentric_frame.galcen_ra, 'z') # extra roll away from the Galactic x-z plane M3 = rotation_matrix(ROLL0 - galactocentric_frame.roll, 'x') # rotate about y' to account for tilt due to Sun's height above the plane z_d = (galactocentric_frame.z_sun / galactocentric_frame.galcen_distance).decompose() M4 = rotation_matrix(-np.arcsin(z_d), 'y') return M4*M3*M1*M2 # this is right: 4,3,1,2
def _icrs_gctc_velocity_matrix(galactocentric_frame): """ Construct a transformation matrix to go from heliocentric ICRS to a galactocentric frame. This is just a rotation and tilt which makes it approximately the same as transforming to Galactic coordinates. This only works for velocity because there is no shift due to the position of the Sun. """ # define rotation matrix to align x(ICRS) with the vector to the Galactic center M1 = rotation_matrix(-galactocentric_frame.galcen_dec, 'y') M2 = rotation_matrix(galactocentric_frame.galcen_ra, 'z') # extra roll away from the Galactic x-z plane M3 = rotation_matrix(ROLL0 - galactocentric_frame.roll, 'x') # rotate about y' to account for tilt due to Sun's height above the plane z_d = (galactocentric_frame.z_sun / galactocentric_frame.galcen_distance).decompose() M4 = rotation_matrix(-np.arcsin(z_d), 'y') return M4 * M3 * M1 * M2 # this is right: 4,3,1,2
def rotate(vector, angle, axis='z', unit=None): """Rotates a vector around axis a right-handed positive angle. This is just a convenience function around `astropy.coordinates.angles.rotation_matrix`. Parameters ---------- vector : array_like Dimension 3 vector. angle : convertible to Angle Angle of rotation. axis : str or 3-sequence Either 'x','y', 'z', or a (x,y,z) specifying an axis to rotate about. If 'x','y', or 'z', the rotation sense is counterclockwise looking down the + axis (e.g. positive rotations obey left-hand-rule). unit : UnitBase, optional If `angle` does not have associated units, they are in this unit. If neither are provided, it is assumed to be degrees. Notes ----- This is just a convenience function around `astropy.coordinates.angles.rotation_matrix`. This performs a so-called *active* or *alibi* transformation: rotates the vector while the coordinate system remains unchanged. To do the opposite operation (*passive* or *alias* transformation) call the function as ``rotate(vec, ax, -angle, unit)`` or use the convenience function :py:func:`transform`, see [1]. References ---------- .. [1] http://en.wikipedia.org/wiki/Rotation_matrix#Ambiguities """ rot = np.asarray(angles.rotation_matrix(-angle, axis, unit)) return np.dot(rot, vector)
@property def lonangle(self): return self.Lambda @property def latangle(self): return self.Beta # Define the Euler angles (from Law & Majewski 2010) phi = radians(180+3.75) theta = radians(90-13.46) psi = radians(180+14.111534) # Generate the rotation matrix using the x-convention (see Goldstein) D = rotation_matrix(phi, "z", degrees=False) C = rotation_matrix(theta, "x", degrees=False) B = rotation_matrix(psi, "z", degrees=False) sgr_matrix = np.array(B.dot(C).dot(D)) # Galactic to Sgr coordinates @transformations.transform_function(coord.GalacticCoordinates, SgrCoordinates) def galactic_to_sgr(galactic_coord): """ Compute the transformation from Galactic spherical to Sgr coordinates. """ l = galactic_coord.l.radians b = galactic_coord.b.radians X = cos(b)*cos(l) Y = cos(b)*sin(l)
# In this case we override the names in the spherical representations but don't # do anything with other representations like cartesian or cylindrical. # # Next we have to define the transformation from this coordinate system to some # other built-in coordinate system; we will use Galactic coordinates. We can do # this by defining functions that return transformation matrices, or by simply # defining a function that accepts a coordinate and returns a new coordinate in # the new system. We'll start by constructing the rotation matrix, using the # ``rotation_matrix`` helper function: SGR_PHI = np.radians(180 + 3.75) # Euler angles (from Law & Majewski 2010) SGR_THETA = np.radians(90 - 13.46) SGR_PSI = np.radians(180 + 14.111534) # Generate the rotation matrix using the x-convention (see Goldstein) D = rotation_matrix(SGR_PHI, "z", unit=u.radian) C = rotation_matrix(SGR_THETA, "x", unit=u.radian) B = rotation_matrix(SGR_PSI, "z", unit=u.radian) SGR_MATRIX = np.array(B.dot(C).dot(D)) ############################################################################## # This is done at the module level, since it will be used by both the # transformation from Sgr to Galactic as well as the inverse from Galactic to # Sgr. Now we can define our first transformation function: @frame_transform_graph.transform(coord.FunctionTransform, coord.Galactic, Sagittarius) def galactic_to_sgr(gal_coord, sgr_frame): """ Compute the transformation from Galactic spherical to heliocentric Sgr coordinates.
def _rotation_opt_func(th, xyz): if th < -1. or th > 1.: return np.inf R = rotation_matrix(np.arccos(th)*u.radian, 'x') xyz2 = coord.CartesianRepresentation(R.dot(xyz)*u.one).xyz return np.sum(xyz2[2]**2).value
def compute_stream_rotation_matrix(coords, zero_pt='mean'): """ Compute the rotation matrix to go from the frame of the input coordinate to closely align the equator with the stream. .. note:: if using zero_pt='mean' or 'median, the coordinates in the new system might not hit (0,0) because the mean / median is taken across each dimension separately. Parameters ---------- coords : :class:`astropy.coordinate.SkyCoord`, :class:`astropy.coordinate.BaseCoordinateFrame` The coordinates of the stream stars. zero_pt : str,:class:`astropy.coordinate.SkyCoord`, :class:`astropy.coordinate.BaseCoordinateFrame` The origin of the new coordinates in the old coordinates. Returns ------- R : :class:`~numpy.ndarray` A 3 by 3 rotation matrix (has shape ``(3,3)``) to convert heliocentric, Cartesian coordinates in the input coordinate frame to stream coordinates. """ if hasattr(coords, 'spherical'): sph = coords.spherical else: sph = coords if zero_pt == 'mean' or zero_pt == 'median': lon = sph.lon.wrap_at(360*u.degree) lat = sph.lat.wrap_at(90*u.degree) if np.any(lon < 10*u.degree) and np.any(lon > 350*u.degree): # it's wrapping lon = lon.wrap_at(180*u.degree) if np.any(lat < -80*u.degree) and np.any(lat > 80*u.degree): # it's wrapping lat = lat.wrap_at(180*u.degree) zero_pt = coord.UnitSphericalRepresentation(lon=getattr(np, zero_pt)(lon), lat=getattr(np, zero_pt)(lat)) elif hasattr(zero_pt, 'spherical'): zero_pt = zero_pt.spherical # first determine rotation matrix to put zero_pt at (0,0) R1 = rotation_matrix(zero_pt.lon, 'z') R2 = rotation_matrix(-zero_pt.lat, 'y') xyz2 = (R2*R1).dot(sph.represent_as(coord.CartesianRepresentation).xyz) # determine initial guess for angle with some math trickery _r = np.sqrt(xyz2[1]**2 + xyz2[2]**2) ix = _r.argmin() if ix == 0: ix += 1 elif ix == (xyz2[1].size-1): ix -= 1 guess = 180*u.degree - np.arctan2(xyz2[2][ix+1]-xyz2[2][ix], xyz2[1][ix+1]-xyz2[1][ix]).to(u.degree) res = minimize(_rotation_opt_func, x0=np.cos(guess), args=(xyz2,), method="powell") if not res.success: raise ValueError("Failed to compute final alignment angle.") if np.allclose(np.abs(res.x), 1., atol=1E-5): guess = 180*u.degree - guess res = minimize(_rotation_opt_func, x0=np.cos(guess), args=(xyz2,), method="powell") R3 = rotation_matrix(np.arccos(res.x)*u.radian, 'x') R = R3*R2*R1 return R
def compute_stream_rotation_matrix(coords, zero_pt='mean'): """ Compute the rotation matrix to go from the frame of the input coordinate to closely align the equator with the stream. .. note:: if using zero_pt='mean' or 'median, the coordinates in the new system might not hit (0,0) because the mean / median is taken across each dimension separately. Parameters ---------- coords : :class:`astropy.coordinate.SkyCoord`, :class:`astropy.coordinate.BaseCoordinateFrame` The coordinates of the stream stars. zero_pt : str,:class:`astropy.coordinate.SkyCoord`, :class:`astropy.coordinate.BaseCoordinateFrame` The origin of the new coordinates in the old coordinates. Returns ------- R : :class:`~numpy.ndarray` A 3 by 3 rotation matrix (has shape ``(3,3)``) to convert heliocentric, Cartesian coordinates in the input coordinate frame to stream coordinates. """ if hasattr(coords, 'spherical'): sph = coords.spherical else: sph = coords if zero_pt == 'mean' or zero_pt == 'median': lon = sph.lon.wrap_at(360 * u.degree) lat = sph.lat.wrap_at(90 * u.degree) if np.any(lon < 10 * u.degree) and np.any( lon > 350 * u.degree): # it's wrapping lon = lon.wrap_at(180 * u.degree) if np.any(lat < -80 * u.degree) and np.any( lat > 80 * u.degree): # it's wrapping lat = lat.wrap_at(180 * u.degree) zero_pt = coord.UnitSphericalRepresentation(lon=getattr(np, zero_pt)(lon), lat=getattr(np, zero_pt)(lat)) elif hasattr(zero_pt, 'spherical'): zero_pt = zero_pt.spherical # first determine rotation matrix to put zero_pt at (0,0) R1 = rotation_matrix(zero_pt.lon, 'z') R2 = rotation_matrix(-zero_pt.lat, 'y') xyz2 = (R2 * R1).dot(sph.represent_as(coord.CartesianRepresentation).xyz) # determine initial guess for angle with some math trickery _r = np.sqrt(xyz2[1]**2 + xyz2[2]**2) ix = _r.argmin() if ix == 0: ix += 1 elif ix == (xyz2[1].size - 1): ix -= 1 guess = 180 * u.degree - np.arctan2(xyz2[2][ix + 1] - xyz2[2][ix], xyz2[1][ix + 1] - xyz2[1][ix]).to( u.degree) res = minimize(_rotation_opt_func, x0=np.cos(guess), args=(xyz2, ), method="powell") if not res.success: raise ValueError("Failed to compute final alignment angle.") if np.allclose(np.abs(res.x), 1., atol=1E-5): guess = 180 * u.degree - guess res = minimize(_rotation_opt_func, x0=np.cos(guess), args=(xyz2, ), method="powell") R3 = rotation_matrix(np.arccos(res.x) * u.radian, 'x') R = R3 * R2 * R1 return R
@property def lonangle(self): return self.Lambda @property def latangle(self): return self.Beta # Define the Euler angles (from Law & Majewski 2010) phi = radians(180 + 3.75) theta = radians(90 - 13.46) psi = radians(180 + 14.111534) # Generate the rotation matrix using the x-convention (see Goldstein) D = rotation_matrix(phi, "z", degrees=False) C = rotation_matrix(theta, "x", degrees=False) B = rotation_matrix(psi, "z", degrees=False) sgr_matrix = np.array(B.dot(C).dot(D)) # Galactic to Sgr coordinates @transformations.transform_function(coord.GalacticCoordinates, SgrCoordinates) def galactic_to_sgr(galactic_coord): """ Compute the transformation from Galactic spherical to Sgr coordinates. """ l = galactic_coord.l.radians b = galactic_coord.b.radians X = cos(b) * cos(l)
def _rotation_opt_func(th, xyz): if th < -1. or th > 1.: return np.inf R = rotation_matrix(np.arccos(th) * u.radian, 'x') xyz2 = coord.CartesianRepresentation(R.dot(xyz) * u.one).xyz return np.sum(xyz2[2]**2).value
def tomin(ang, inrep=rep[hostidx], axis=rep[targetidx].xyz, target=galactic_center.icrs): newr = rotate_repr(inrep, rotation_matrix(ang[0]*u.deg, axis)) return ICRS(newr).separation(target).radian
def _ecliptic_rotation_matrix_pulsar(obl): """Here we only do the obliquity angle rotation. Astropy will add the precession-nutation correction. """ return rotation_matrix(obl, 'x')
# In this case we override the names in the spherical representations but don't # do anything with other representations like cartesian or cylindrical. # # Next we have to define the transformation from this coordinate system to some # other built-in coordinate system; we will use Galactic coordinates. We can do # this by defining functions that return transformation matrices, or by simply # defining a function that accepts a coordinate and returns a new coordinate in # the new system. We'll start by constructing the rotation matrix, using the # ``rotation_matrix`` helper function: SGR_PHI = np.radians(180+3.75) # Euler angles (from Law & Majewski 2010) SGR_THETA = np.radians(90-13.46) SGR_PSI = np.radians(180+14.111534) # Generate the rotation matrix using the x-convention (see Goldstein) D = rotation_matrix(SGR_PHI, "z", unit=u.radian) C = rotation_matrix(SGR_THETA, "x", unit=u.radian) B = rotation_matrix(SGR_PSI, "z", unit=u.radian) SGR_MATRIX = np.array(B.dot(C).dot(D)) ############################################################################## # This is done at the module level, since it will be used by both the # transformation from Sgr to Galactic as well as the inverse from Galactic to # Sgr. Now we can define our first transformation function: @frame_transform_graph.transform(coord.FunctionTransform, coord.Galactic, Sagittarius) def galactic_to_sgr(gal_coord, sgr_frame): """ Compute the transformation from Galactic spherical to heliocentric Sgr coordinates. """
""" default_representation = coord.SphericalRepresentation _frame_specific_representation_info = { 'spherical': {'names': ('Lambda', 'Beta', 'distance'), 'units': (u.degree, u.degree, None)}, 'unitspherical': {'names': ('Lambda', 'Beta'), 'units': (u.degree, u.degree)}} # Define the Euler angles (from Law & Majewski 2010) phi = np.radians(180+3.75) theta = np.radians(90-13.46) psi = np.radians(180+14.111534) # Generate the rotation matrix using the x-convention (see Goldstein) D = rotation_matrix(phi, "z", unit=u.radian) C = rotation_matrix(theta, "x", unit=u.radian) B = rotation_matrix(psi, "z", unit=u.radian) sgr_matrix = np.array(B.dot(C).dot(D)) # Galactic to Sgr coordinates @frame_transform_graph.transform(coord.FunctionTransform, coord.Galactic, Sagittarius) def galactic_to_sgr(gal_coord, sgr_frame): """ Compute the transformation from Galactic spherical to heliocentric Sgr coordinates. """ l = np.atleast_1d(gal_coord.l.radian) b = np.atleast_1d(gal_coord.b.radian)
def tomin(ang, inrep=rep[hostidx], axis=rep[targetidx].xyz, target=galactic_center.icrs): newr = rotate_repr(inrep, rotation_matrix(ang[0] * u.deg, axis)) return ICRS(newr).separation(target).radian
def main(progenitor_mass, n_stars, seed=42): np.random.seed(seed) _path, _ = os.path.split(os.path.abspath(__file__)) top_path = os.path.abspath(os.path.join(_path, "..")) output_path = os.path.join(top_path, "output", "simulations") if not os.path.exists(output_path): os.makedirs(output_path) for potential_name, potential in potentials.items(): logger.info("Potential: {}".format(potential_name)) this_output_path = os.path.join(output_path, potential_name) this_plot_path = os.path.join(this_output_path, 'plots') if not os.path.exists(this_output_path): os.mkdir(this_output_path) if not os.path.exists(this_plot_path): os.mkdir(this_plot_path) with h5py.File(os.path.join(this_output_path, "mock_stream_data.h5"), "w") as f: for i, (per, apo) in enumerate(per_apo): g = f.create_group(str(i)) g.attrs['apocenter'] = apo g.attrs['pericenter'] = per # get random initial conditions for given pericenter, apocenter w0, T = peri_apo_to_random_w0(per, apo, potential, frac_r_start=0.1) # integration time t1 = T / 2 * 0.9 n_steps = 10000 dt = t1 / n_steps g.attrs['n_steps'] = n_steps g.attrs['dt'] = dt logger.debug("dt: {:.2f}, N steps: {}".format(dt, n_steps)) # integrate orbit prog_orbit = potential.integrate_orbit(w0, dt=dt, nsteps=n_steps) sph, _ = prog_orbit.represent_as(coord.SphericalRepresentation) logger.debug("Data distance min,max = {}, {}".format( sph.distance.min(), sph.distance.max())) m = progenitor_mass * u.Msun rtide = (m / potential.mass_enclosed(w0.pos))**( 1 / 3.) * np.sqrt(np.sum(w0.pos**2)) vdisp = np.sqrt(G * m / (2 * rtide)).to(u.km / u.s) logger.debug("rtide, vdisp: {}, {}".format(rtide, vdisp)) # Option 1: generate mock stream # stream = mockstream.fardal_stream(potential, prog_orbit=prog_orbit, # prog_mass=m, release_every=1, # Integrator=gi.DOPRI853Integrator) # Option 2: integrate a ball of test particle orbits # std = gd.CartesianPhaseSpacePosition(pos=[rtide.value]*3*u.kpc, # vel=[vdisp.value]*3*u.km/u.s) # ball_w = w0.w(galactic)[:,0] # ball_std = std.w(galactic)[:,0] # ball_w0 = np.random.normal(ball_w, ball_std, size=(n_stars,6)) # ball_w0 = gd.CartesianPhaseSpacePosition.from_w(ball_w0.T, units=galactic) # stream_orbits = potential.integrate_orbit(ball_w0, dt=1., nsteps=n_steps) # stream = stream_orbits[-1] # Option 3: just take single orbit, convolve with uncertainties prog_orbit = prog_orbit[-n_steps::2] stream = gd.CartesianPhaseSpacePosition(pos=prog_orbit.pos, vel=prog_orbit.vel) # save simulated stream data g.attrs['mass'] = progenitor_mass g.create_dataset('pos', shape=stream.pos.shape, dtype=np.float64, data=stream.pos.decompose(galactic).value) g['pos'].attrs['unit'] = 'kpc' g.create_dataset('vel', shape=stream.vel.shape, dtype=np.float64, data=stream.vel.decompose(galactic).value) g['vel'].attrs['unit'] = 'kpc/Myr' # plot the orbit in cartesian coords fig = prog_orbit.plot(color='lightblue', alpha=0.5) fig = stream.plot(axes=fig.axes, marker='.', alpha=0.5) for ax in fig.axes: ax.set_xlim(-apo - 10, apo + 10) ax.set_ylim(-apo - 10, apo + 10) fig.savefig( os.path.join(this_plot_path, "orbit-{}.png".format(i))) # convert to sky coordinates and compute the stream rotation matrix stream_c, stream_v = stream.to_frame(coord.Galactic, **FRAME) R = compute_stream_rotation_matrix(stream_c, zero_pt=stream_c[0]) stream_rot = rotate_sph_coordinate(stream_c, R) if stream_rot.lon.wrap_at(180 * u.degree).degree[-1] < 0: logger.debug("flipping stream...") flip = rotation_matrix(180 * u.degree, 'x') stream_rot = rotate_sph_coordinate(stream_rot, flip) R = flip * R g['R'] = R # plot the orbit on the sky in galactic and in stream coordinates fig_gal, _ = plot_stream_obs(stream_c, stream_v) fig_rot, _ = plot_stream_obs(stream_rot, stream_v) fig_gal.savefig( os.path.join(this_plot_path, "stream-{}-gal.png".format(i))) fig_rot.savefig( os.path.join(this_plot_path, "stream-{}-rot.png".format(i))) pl.close('all')
def add_oriented_radecs(elvis_tab, hostidx=0, targetidx=1, target_coord=SkyCoord(0 * u.deg, 0 * u.deg), earth_distance=8.5 * u.kpc, earth_vrot=220 * u.km / u.s, roll_angle=0 * u.deg): """ Computes a spherical coordinate system centered on the `hostidx` halo, re-oriented so that `targetidx` is at the `target_coord` coordinate location. Note that this adds columns 'host<n>_*' to `elvis_tab`, and will *overwrite* them if they already exist. """ if hasattr(target_coord, 'spherical'): target_lat = target_coord.spherical.lat target_lon = target_coord.spherical.lon else: target_lat = target_coord.lat target_lon = target_coord.lon def offset_repr(rep, vector, newrep=None): if newrep is None: newrep = rep.__class__ newxyz = rep.to_cartesian().xyz + vector.reshape(3, 1) return CartesianRepresentation(newxyz).represent_as(newrep) def rotate_repr(rep, matrix, newrep=None): if newrep is None: newrep = rep.__class__ newxyz = np.dot(matrix.view(np.ndarray), rep.to_cartesian().xyz) return CartesianRepresentation(newxyz).represent_as(newrep) rep = CartesianRepresentation(elvis_tab['X'], elvis_tab['Y'], elvis_tab['Z']) # first we offset the catalog to have its origin at host rep = offset_repr(rep, -rep.xyz[:, hostidx]) # now rotate so that host1 is along the z-axis, and apply the arbitrary roll angle usph = rep.represent_as(UnitSphericalRepresentation) M1 = rotation_matrix(usph.lon[targetidx], 'z') M2 = rotation_matrix(90 * u.deg - usph.lat[targetidx], 'y') M3 = rotation_matrix(roll_angle, 'z') Mfirst = M3 * M2 * M1 rep = rotate_repr(rep, Mfirst) # now determine the location of the earth in this system # need diagram to explain this, but it uses SSA formula theta = target_coord.separation(galactic_center) # target to GC angle D = rep.z[targetidx] # distance to the target host R = earth_distance # srho = (R/D) * np.sin(theta) # sdelta_p = (srho * np.cos(theta) + (1 - srho**2)**0.5) # sdelta_m = (srho * np.cos(theta) - (1 - srho**2)**0.5) d1, d2 = R * np.cos(theta), (D**2 - (R * np.sin(theta))**2)**0.5 dp, dm = d1 + d2, d1 - d2 sdelta = (dp / D) * np.sin(theta) x = R * sdelta z = R * (1 - sdelta**2)**0.5 earth_location = u.Quantity([x, 0 * u.kpc, z]) # now offset to put earth at the origin rep = offset_repr(rep, -earth_location) sph = rep.represent_as(SphericalRepresentation) # rotate to put the target at its correct spot # first sent the target host to 0,0 M1 = rotation_matrix(sph[targetidx].lon, 'z') M2 = rotation_matrix(-sph[targetidx].lat, 'y') # now rotate from origin to target lat,lon M3 = rotation_matrix(target_lat, 'y') M4 = rotation_matrix(-target_lon, 'z') Mmiddle = M4 * M3 * M2 * M1 rep = rotate_repr(rep, Mmiddle) # now one more rotation about the target to stick the GC in the right place def tomin(ang, inrep=rep[hostidx], axis=rep[targetidx].xyz, target=galactic_center.icrs): newr = rotate_repr(inrep, rotation_matrix(ang[0] * u.deg, axis)) return ICRS(newr).separation(target).radian rot_angle = optimize.minimize(tomin, np.array(0).ravel(), method='Nelder-Mead')['x'][0] Mlast = rotation_matrix(rot_angle * u.deg, rep[targetidx].xyz) rep = rotate_repr(rep, Mlast) sph = rep.represent_as(SphericalRepresentation) elvis_tab['host{}_lat'.format(hostidx)] = sph.lat.to(u.deg) elvis_tab['host{}_lon'.format(hostidx)] = sph.lon.to(u.deg) elvis_tab['host{}_dist'.format(hostidx)] = sph.distance # now compute velocities # host galactocentric dvxg = u.Quantity((elvis_tab['Vx']) - elvis_tab['Vx'][hostidx]) dvyg = u.Quantity((elvis_tab['Vy']) - elvis_tab['Vy'][hostidx]) dvzg = u.Quantity((elvis_tab['Vz']) - elvis_tab['Vz'][hostidx]) earth_location_in_xyz = np.dot(Mfirst.T, earth_location) dxg = elvis_tab['X'] - elvis_tab['X'][0] - earth_location_in_xyz[0] dyg = elvis_tab['Y'] - elvis_tab['Y'][0] - earth_location_in_xyz[0] dzg = elvis_tab['Z'] - elvis_tab['Z'][0] - earth_location_in_xyz[0] vrg = (dvxg * dxg + dvyg * dyg + dvzg * dzg) * (dxg**2 + dyg**2 + dzg**2)**-0.5 elvis_tab['host{}_galvr'.format(hostidx)] = vrg.to(u.km / u.s) # "vLSR-like" # first figure out the rotation direction earth_rotdir = SkyCoord(90 * u.deg, 0 * u.deg, frame='galactic').icrs #now apply the component from that everywhere offset_angle = earth_rotdir.separation(ICRS(sph)) vrlsr = vrg - earth_vrot * np.cos(offset_angle) elvis_tab['host{}_vrlsr'.format(hostidx)] = vrlsr.to(u.km / u.s)
@property def lonangle(self): return self.Lambda @property def latangle(self): return self.Beta # Define the Euler angles (from Law & Majewski 2010) phi = radians(180 + 3.75) theta = radians(90 - 13.46) psi = radians(180 + 14.111534) # Generate the rotation matrix using the x-convention (see Goldstein) D = rotation_matrix(phi, "z", unit=u.radian) C = rotation_matrix(theta, "x", unit=u.radian) B = rotation_matrix(psi, "z", unit=u.radian) sgr_matrix = np.array(B.dot(C).dot(D)) # Galactic to Sgr coordinates @transformations.transform_function(coord.GalacticCoordinates, SgrCoordinates) def galactic_to_sgr(galactic_coord): """ Compute the transformation from Galactic spherical to Sgr coordinates. """ l = galactic_coord.l.radian b = galactic_coord.b.radian X = cos(b) * cos(l)
def main(progenitor_mass, n_stars, seed=42): np.random.seed(seed) _path,_ = os.path.split(os.path.abspath(__file__)) top_path = os.path.abspath(os.path.join(_path, "..")) output_path = os.path.join(top_path, "output", "simulations") if not os.path.exists(output_path): os.makedirs(output_path) for potential_name,potential in potentials.items(): logger.info("Potential: {}".format(potential_name)) this_output_path = os.path.join(output_path, potential_name) this_plot_path = os.path.join(this_output_path, 'plots') if not os.path.exists(this_output_path): os.mkdir(this_output_path) if not os.path.exists(this_plot_path): os.mkdir(this_plot_path) with h5py.File(os.path.join(this_output_path, "mock_stream_data.h5"), "w") as f: for i,(per,apo) in enumerate(per_apo): g = f.create_group(str(i)) g.attrs['apocenter'] = apo g.attrs['pericenter'] = per # get random initial conditions for given pericenter, apocenter w0,T = peri_apo_to_random_w0(per, apo, potential, frac_r_start=0.1) # integration time t1 = T/2*0.9 n_steps = 10000 dt = t1/n_steps g.attrs['n_steps'] = n_steps g.attrs['dt'] = dt logger.debug("dt: {:.2f}, N steps: {}".format(dt, n_steps)) # integrate orbit prog_orbit = potential.integrate_orbit(w0, dt=dt, nsteps=n_steps) sph,_ = prog_orbit.represent_as(coord.SphericalRepresentation) logger.debug("Data distance min,max = {}, {}".format(sph.distance.min(), sph.distance.max())) m = progenitor_mass*u.Msun rtide = (m/potential.mass_enclosed(w0.pos))**(1/3.) * np.sqrt(np.sum(w0.pos**2)) vdisp = np.sqrt(G*m/(2*rtide)).to(u.km/u.s) logger.debug("rtide, vdisp: {}, {}".format(rtide, vdisp)) # Option 1: generate mock stream # stream = mockstream.fardal_stream(potential, prog_orbit=prog_orbit, # prog_mass=m, release_every=1, # Integrator=gi.DOPRI853Integrator) # Option 2: integrate a ball of test particle orbits # std = gd.CartesianPhaseSpacePosition(pos=[rtide.value]*3*u.kpc, # vel=[vdisp.value]*3*u.km/u.s) # ball_w = w0.w(galactic)[:,0] # ball_std = std.w(galactic)[:,0] # ball_w0 = np.random.normal(ball_w, ball_std, size=(n_stars,6)) # ball_w0 = gd.CartesianPhaseSpacePosition.from_w(ball_w0.T, units=galactic) # stream_orbits = potential.integrate_orbit(ball_w0, dt=1., nsteps=n_steps) # stream = stream_orbits[-1] # Option 3: just take single orbit, convolve with uncertainties prog_orbit = prog_orbit[-n_steps::2] stream = gd.CartesianPhaseSpacePosition(pos=prog_orbit.pos, vel=prog_orbit.vel) # save simulated stream data g.attrs['mass'] = progenitor_mass g.create_dataset('pos', shape=stream.pos.shape, dtype=np.float64, data=stream.pos.decompose(galactic).value) g['pos'].attrs['unit'] = 'kpc' g.create_dataset('vel', shape=stream.vel.shape, dtype=np.float64, data=stream.vel.decompose(galactic).value) g['vel'].attrs['unit'] = 'kpc/Myr' # plot the orbit in cartesian coords fig = prog_orbit.plot(color='lightblue', alpha=0.5) fig = stream.plot(axes=fig.axes, marker='.', alpha=0.5) for ax in fig.axes: ax.set_xlim(-apo-10, apo+10) ax.set_ylim(-apo-10, apo+10) fig.savefig(os.path.join(this_plot_path, "orbit-{}.png".format(i))) # convert to sky coordinates and compute the stream rotation matrix stream_c,stream_v = stream.to_frame(coord.Galactic, **FRAME) R = compute_stream_rotation_matrix(stream_c, zero_pt=stream_c[0]) stream_rot = rotate_sph_coordinate(stream_c, R) if stream_rot.lon.wrap_at(180*u.degree).degree[-1] < 0: logger.debug("flipping stream...") flip = rotation_matrix(180*u.degree, 'x') stream_rot = rotate_sph_coordinate(stream_rot, flip) R = flip*R g['R'] = R # plot the orbit on the sky in galactic and in stream coordinates fig_gal,_ = plot_stream_obs(stream_c, stream_v) fig_rot,_ = plot_stream_obs(stream_rot, stream_v) fig_gal.savefig(os.path.join(this_plot_path, "stream-{}-gal.png".format(i))) fig_rot.savefig(os.path.join(this_plot_path, "stream-{}-rot.png".format(i))) pl.close('all')
def add_oriented_radecs(elvis_tab, hostidx=0, targetidx=1, target_coord=SkyCoord(0*u.deg, 0*u.deg), earth_distance=8.5*u.kpc, earth_vrot=220*u.km/u.s, roll_angle=0*u.deg): """ Computes a spherical coordinate system centered on the `hostidx` halo, re-oriented so that `targetidx` is at the `target_coord` coordinate location. Note that this adds columns 'host<n>_*' to `elvis_tab`, and will *overwrite* them if they already exist. """ if hasattr(target_coord, 'spherical'): target_lat = target_coord.spherical.lat target_lon = target_coord.spherical.lon else: target_lat = target_coord.lat target_lon = target_coord.lon def offset_repr(rep, vector, newrep=None): if newrep is None: newrep = rep.__class__ newxyz = rep.to_cartesian().xyz + vector.reshape(3, 1) return CartesianRepresentation(newxyz).represent_as(newrep) def rotate_repr(rep, matrix, newrep=None): if newrep is None: newrep = rep.__class__ newxyz = np.dot(matrix.view(np.ndarray), rep.to_cartesian().xyz) return CartesianRepresentation(newxyz).represent_as(newrep) rep = CartesianRepresentation(elvis_tab['X'], elvis_tab['Y'], elvis_tab['Z']) # first we offset the catalog to have its origin at host rep = offset_repr(rep, -rep.xyz[:, hostidx]) # now rotate so that host1 is along the z-axis, and apply the arbitrary roll angle usph = rep.represent_as(UnitSphericalRepresentation) M1 = rotation_matrix(usph.lon[targetidx], 'z') M2 = rotation_matrix(90*u.deg-usph.lat[targetidx], 'y') M3 = rotation_matrix(roll_angle, 'z') Mfirst = M3*M2*M1 rep = rotate_repr(rep, Mfirst) # now determine the location of the earth in this system # need diagram to explain this, but it uses SSA formula theta = target_coord.separation(galactic_center) # target to GC angle D = rep.z[targetidx] # distance to the target host R = earth_distance # srho = (R/D) * np.sin(theta) # sdelta_p = (srho * np.cos(theta) + (1 - srho**2)**0.5) # sdelta_m = (srho * np.cos(theta) - (1 - srho**2)**0.5) d1, d2 = R * np.cos(theta), (D**2 - (R * np.sin(theta))**2)**0.5 dp, dm = d1 + d2, d1 - d2 sdelta = (dp/D) * np.sin(theta) x = R * sdelta z = R * (1-sdelta**2)**0.5 earth_location = u.Quantity([x, 0*u.kpc, z]) # now offset to put earth at the origin rep = offset_repr(rep, -earth_location) sph = rep.represent_as(SphericalRepresentation) # rotate to put the target at its correct spot # first sent the target host to 0,0 M1 = rotation_matrix(sph[targetidx].lon, 'z') M2 = rotation_matrix(-sph[targetidx].lat, 'y') # now rotate from origin to target lat,lon M3 = rotation_matrix(target_lat, 'y') M4 = rotation_matrix(-target_lon, 'z') Mmiddle = M4*M3*M2*M1 rep = rotate_repr(rep, Mmiddle) # now one more rotation about the target to stick the GC in the right place def tomin(ang, inrep=rep[hostidx], axis=rep[targetidx].xyz, target=galactic_center.icrs): newr = rotate_repr(inrep, rotation_matrix(ang[0]*u.deg, axis)) return ICRS(newr).separation(target).radian rot_angle = optimize.minimize(tomin, np.array(0).ravel(), method='Nelder-Mead')['x'][0] Mlast = rotation_matrix(rot_angle*u.deg, rep[targetidx].xyz) rep = rotate_repr(rep, Mlast) sph = rep.represent_as(SphericalRepresentation) elvis_tab['host{}_lat'.format(hostidx)] = sph.lat.to(u.deg) elvis_tab['host{}_lon'.format(hostidx)] = sph.lon.to(u.deg) elvis_tab['host{}_dist'.format(hostidx)] = sph.distance # now compute velocities # host galactocentric dvxg = u.Quantity((elvis_tab['Vx'])-elvis_tab['Vx'][hostidx]) dvyg = u.Quantity((elvis_tab['Vy'])-elvis_tab['Vy'][hostidx]) dvzg = u.Quantity((elvis_tab['Vz'])-elvis_tab['Vz'][hostidx]) earth_location_in_xyz = np.dot(Mfirst.T, earth_location) dxg = elvis_tab['X'] - elvis_tab['X'][0] - earth_location_in_xyz[0] dyg = elvis_tab['Y'] - elvis_tab['Y'][0] - earth_location_in_xyz[0] dzg = elvis_tab['Z'] - elvis_tab['Z'][0] - earth_location_in_xyz[0] vrg = (dvxg*dxg + dvyg*dyg + dvzg*dzg) * (dxg**2+dyg**2+dzg**2)**-0.5 elvis_tab['host{}_galvr'.format(hostidx)] = vrg.to(u.km/u.s) # "vLSR-like" # first figure out the rotation direction earth_rotdir = SkyCoord(90*u.deg, 0*u.deg, frame='galactic').icrs #now apply the component from that everywhere offset_angle = earth_rotdir.separation(ICRS(sph)) vrlsr = vrg - earth_vrot*np.cos(offset_angle) elvis_tab['host{}_vrlsr'.format(hostidx)] = vrlsr.to(u.km/u.s)