def setup_class(cls): """Setup fixtures for testing.""" num = 40 cls.arclength = np.linspace(0, 10, num=num) * u.deg lon = np.linspace(0, 25, num=num) * u.deg lat = np.linspace(-10, 10, num=num) * u.deg distance = np.linspace(8, 15, num=num) * u.kpc cls.data = coord.ICRS( coord.SphericalRepresentation(lon=lon, lat=lat, distance=distance), ) cls.interps = dict( lon=IUSU(cls.arclength, lon), lat=IUSU(cls.arclength, lat), distance=IUSU(cls.arclength, distance), ) # origin i = num // 2 cls.origin = coord.ICRS(ra=lon[i], dec=lat[i]) # track cls.track = StreamTrack( cls.interps, stream_data=cls.data, origin=cls.origin, )
def _mcmc_sample_to_coord(self, p): _, orbit_pars = self._unpack_orbit(0, p) rep = coord.SphericalRepresentation(lon=[0.] * u.radian, lat=[orbit_pars['phi2']] * u.radian, distance=[orbit_pars['d']] * u.kpc) return coord.Galactic(orbitfit.rotate_sph_coordinate(rep, self.R.T))
def test_run(self): """Test method ``run``.""" # ------------------- # defaults resid = self.inst.run(fit_potential=self.original_potential, original_potential=None, observable=None, representation_type=None, batch=True, **self.kwargs) assert resid == coord.CartesianRepresentation((0, 0, 0)) # ------------------- # passing in values resid = self.inst.run( fit_potential=self.original_potential, original_potential=self.original_potential, observable=self.observable, representation_type="spherical", batch=True, ) assert resid == coord.SphericalRepresentation(0 * u.rad, 0 * u.rad, 0)
def test_xpercenterror_factory(): """Test :fun:`~discO.core.measurement.xpercenterror_factory`.""" rep = coord.SphericalRepresentation(1 * u.rad, 1 * u.rad, 1 * u.kpc) crd = coord.ICRS(rep.reshape(-1, 1)) # -------------------------- # percent input func = measurement.xpercenterror_factory(10 * u.percent) res = func(crd) assert callable(func) assert np.allclose(res, [0.1, 0.1, 0.1]) # -------------------------- # fractional error input func2 = measurement.xpercenterror_factory(0.2) res2 = func2(crd) assert callable(func2) assert np.allclose(res2, [0.2, 0.2, 0.2]) # -------------------------- # caching assert measurement.xpercenterror_factory(10 * u.percent) is func assert measurement.xpercenterror_factory(0.2) is func2 # -------------------------- # docstring editing assert "10.0%" in func.__doc__ assert "20.0%" in func2.__doc__
def plot_data(data, err, R, wrap_angle=180 * u.degree, fig=None, gal=True, **style_kw): if fig is None: fig, axes = pl.subplots(2, 3, figsize=(12, 8), sharex=True) else: axes = np.array(fig.axes).reshape(2, 3) if gal: rep = coord.SphericalRepresentation(lon=data['phi1'], lat=data['phi2'], distance=data['distance']) g = coord.Galactic(orbitfit.rotate_sph_coordinate(rep, R.T)) lon = g.l lat = g.b axes[1, 1].set_xlabel('$l$ [deg]') else: lon = coord.Angle(data['phi1']) lat = coord.Angle(data['phi2']) axes[1, 1].set_xlabel('$\phi_1$ [deg]') style = dict(linestyle='none', marker='.', ecolor='#aaaaaa') for k, v in style_kw.items(): style[k] = v lon = lon.wrap_at(wrap_angle) axes[0, 0].errorbar(lon.degree, lat.degree, 1E-8 * lat.degree, **style) axes[1, 0].errorbar(lon.degree, data['distance'].value, err['distance'].value, **style) axes[0, 1].errorbar(lon.degree, galactic.decompose(data['mul']).value, galactic.decompose(err['mul']).value, **style) axes[1, 1].errorbar(lon.degree, galactic.decompose(data['mub']).value, galactic.decompose(err['mub']).value, **style) axes[0, 2].errorbar(lon.degree, galactic.decompose(data['vr']).value, galactic.decompose(err['vr']).value, **style) try: fig.tight_layout() except AttributeError: # already called pass axes[1, 2].set_visible(False) return fig, axes
def setup_class(cls): """Setup fixtures for testing.""" # test class class TestClass(metaclass=cls.obj): _frame = coord.Galactocentric() cls.subclass = TestClass cls.frame = coord.ICRS cls.rep = coord.SphericalRepresentation( lon=[0, 1, 2] * u.deg, lat=[3, 4, 5] * u.deg, distance=[6, 7, 8] * u.kpc, ) cls.points = coord.SkyCoord(cls.frame(cls.rep), copy=False) # the potential cls.potential = object()
def test_representation_is_BaseCoordinateFrame(): """Test when representation is a BaseCoordinateFrame.""" # basic usage assert (resolve_representationlike( representation=coord.CartesianRepresentation, ) == coord.CartesianRepresentation) assert (resolve_representationlike( representation=coord.CartesianRepresentation( x=(1, 2, 3)), ) == coord.CartesianRepresentation) # replicates without data c = coord.SphericalRepresentation( lon=1 * u.deg, lat=2 * u.deg, distance=3 * u.kpc, ) assert (resolve_representationlike( representation=c) == coord.SphericalRepresentation)
def eq2cart(r, lat, lon): """ Equatorial coordinates to Cartesian coordinates. Parameters ---------- r : float or :py:class:`~numpy.ndarray` Radius. lat : :py:class:`~numpy.ndarray` Elevation angle [rad]. lon : :py:class:`~numpy.ndarray` Longitude angle [rad]. Returns ------- XYZ : :py:class:`~numpy.ndarray` (3, ...) Cartesian XYZ coordinates. Examples -------- .. testsetup:: import numpy as np from imot_tools.math.sphere.transform import eq2cart .. doctest:: >>> xyz = eq2cart(1, 0, 0) >>> np.around(xyz, 2) array([[1.], [0.], [0.]]) """ r = np.array([r]) if chk.is_scalar(r) else np.array(r, copy=False) if np.any(r < 0): raise ValueError("Parameter[r] must be non-negative.") XYZ = (coord.SphericalRepresentation(lon * u.rad, lat * u.rad, r).to_cartesian().xyz.to_value( u.dimensionless_unscaled)) return XYZ
def setup_class(cls): """Setup fixtures for testing.""" # subclasses can define an attribute "potential" that is used by the # wrapper. Else, let's just make an object if not hasattr(cls, "potential"): cls.potential = object() cls.inst = cls.obj(cls.potential, frame="galactocentric") cls.frame = coord.ICRS cls.rep = coord.SphericalRepresentation( lon=[0, 1, 2] * u.deg, lat=[3, 4, 5] * u.deg, distance=[6, 7, 8] * u.kpc, ) cls.points = coord.SkyCoord(cls.frame(cls.rep), copy=False) class SubClass(cls.obj, key=cls.obj.__name__): pass cls.subclass = SubClass
def test_representation_representation(): """ Test that Representations are represented correctly. """ # With no unit we get "None" in the unit row c = coordinates.CartesianRepresentation([0], [1], [0], unit=u.one) t = Table([c]) assert t.pformat() == [' col0 ', '------------', '(0., 1., 0.)'] c = coordinates.CartesianRepresentation([0], [1], [0], unit='m') t = Table([c]) assert t.pformat() == [' col0 ', ' m ', '------------', '(0., 1., 0.)'] c = coordinates.SphericalRepresentation([10]*u.deg, [20]*u.deg, [1]*u.pc) t = Table([c]) assert t.pformat() == [' col0 ', ' deg, deg, pc ', '--------------', '(10., 20., 1.)'] c = coordinates.UnitSphericalRepresentation([10]*u.deg, [20]*u.deg) t = Table([c]) assert t.pformat() == [' col0 ', ' deg ', '----------', '(10., 20.)'] c = coordinates.SphericalCosLatDifferential( [10]*u.mas/u.yr, [2]*u.mas/u.yr, [10]*u.km/u.s) t = Table([c]) assert t.pformat() == [' col0 ', 'mas / yr, mas / yr, km / s', '--------------------------', ' (10., 2., 10.)']
def setup_class(cls): """Setup fixtures for testing.""" # Make Data in rotated frame, to ensure linearity num = 40 lon = np.linspace(0, 25, num=num) * u.deg lat = np.linspace(0, 0, num=num) * u.deg distance = np.linspace(8, 15, num=num) * u.kpc i = num // 2 # index of origin cls.origin = coord.ICRS(ra=lon[i], dec=lat[i], distance=distance[i]) cls.rotframe = coord.SkyOffsetFrame( origin=cls.origin, rotation=-37 * u.deg, ) rotdata = cls.rotframe.realize_frame( coord.SphericalRepresentation(lon=lon, lat=lat, distance=distance), ) # given the construction, the lon variable is a good arclength, # just need to adjust for the origin cls.arclength = lon - lon[i] # make data in ICRS from rotated data cls.data = rotdata.transform_to(coord.ICRS) cls.lon = cls.data.ra cls.lat = cls.data.dec cls.distance = cls.data.distance cls.data_err = QTable( cls.data.cartesian.xyz.T / 10, names=["x_err", "y_err", "z_err"], ) cls.tracker = cls.klass( data=cls.data, origin=cls.origin, data_err=cls.data_err, frame=cls.rotframe, )
def __call__(self, n, *, random=None, **kwargs): representation_type = self.representation_type # can be None if random is None: random = np.random elif isinstance(random, int): random = np.random.RandomState(random) # return rep = coord.SphericalRepresentation( lon=random.uniform(size=n) * u.deg, lat=2 * random.uniform(size=n) * u.deg, distance=10 * u.kpc, ) sample = coord.SkyCoord( self.frame.realize_frame( rep, representation_type=representation_type or rep.__class__, ), copy=False, ) sample.cache["mass"] = np.ones(n) sample.cache["potential"] = object() return sample
def setup_class(cls): """Setup fixtures for testing.""" cls.representation_type = coord.CartesianRepresentation cls.observable = "acceleration" cls.kwargs = {} _r = np.linspace(0.1, 10, num=50) _lon = np.linspace(0, 360, num=10) _lat = np.linspace(-90, 90, num=10) r, lon, lat = np.meshgrid(_r, _lon, _lat) cls.points = coord.SphericalRepresentation( lon=lon * u.deg, lat=lat * u.deg, distance=r * u.kpc, ) if cls.obj is residual.ResidualMethod: cls.original_potential = object() class SubClass(cls.obj): def evaluate_potential(self, *args, **kwargs): return coord.CartesianRepresentation(x=1, y=2, z=3) cls.klass = SubClass # /class # have to go the long way around cls.inst = cls.klass(original_potential=cls.original_potential, observable=cls.observable, representation_type=cls.representation_type, **cls.kwargs) # need to assign points for run cls.inst.points = cls.points else: pass # need to do in subclass
# Stuff for testing mixin columns MIXIN_COLS = { 'quantity': [0, 1, 2, 3] * u.m, 'longitude': coordinates.Longitude([0., 1., 5., 6.] * u.deg, wrap_angle=180. * u.deg), 'latitude': coordinates.Latitude([5., 6., 10., 11.] * u.deg), 'time': time.Time([2000, 2001, 2002, 2003], format='jyear'), 'skycoord': coordinates.SkyCoord(ra=[0, 1, 2, 3] * u.deg, dec=[0, 1, 2, 3] * u.deg), 'sphericalrep': coordinates.SphericalRepresentation([0, 1, 2, 3] * u.deg, [0, 1, 2, 3] * u.deg, 1 * u.kpc), 'cartesianrep': coordinates.CartesianRepresentation([0, 1, 2, 3] * u.pc, [4, 5, 6, 7] * u.pc, [9, 8, 8, 6] * u.pc), 'sphericaldiff': coordinates.SphericalCosLatDifferential([0, 1, 2, 3] * u.mas / u.yr, [0, 1, 2, 3] * u.mas / u.yr, 10 * u.km / u.s), 'arraywrap': table_helpers.ArrayWrapper([0, 1, 2, 3]), 'ndarray': np.array([(7, 'a'), (8, 'b'), (9, 'c'), (9, 'c')], dtype='<i4,|S1').view(table.NdarrayMixin), } MIXIN_COLS['earthlocation'] = coordinates.EarthLocation(
# Licensed under a 3-clause BSD style license - see LICENSE.rst """ Mixin columns for use in ascii/tests/test_ecsv.py, fits/tests/test_connect.py, and misc/tests/test_hdf5.py """ from astropy import coordinates, table, time, units as u el = coordinates.EarthLocation(x=[1, 2] * u.km, y=[3, 4] * u.km, z=[5, 6] * u.km) sr = coordinates.SphericalRepresentation([0, 1] * u.deg, [2, 3] * u.deg, 1 * u.kpc) cr = coordinates.CartesianRepresentation([0, 1] * u.pc, [4, 5] * u.pc, [8, 6] * u.pc) sd = coordinates.SphericalCosLatDifferential([0, 1] * u.mas / u.yr, [0, 1] * u.mas / u.yr, 10 * u.km / u.s) srd = coordinates.SphericalRepresentation(sr, differentials=sd) sc = coordinates.SkyCoord([1, 2], [3, 4], unit='deg,deg', frame='fk4', obstime='J1990.5') scd = coordinates.SkyCoord([1, 2], [3, 4], [5, 6], unit='deg,deg,m', frame='fk4', obstime=['J1990.5'] * 2) scdc = scd.copy() scdc.representation_type = 'cartesian' scpm = coordinates.SkyCoord([1, 2], [3, 4], [5, 6], unit='deg,deg,pc',
def test_highlevel_api(): J2001 = time.Time('J2001') # <--------------------------"High-level" class--------------------------------> # The "high-level" class is intended to wrap the lower-level classes in such a # way that they can be round-tripped, as well as providing a variety of # convenience functionality. This document is not intended to show *all* of the # possible high-level functionality, rather how the high-level classes are # initialized and interact with the low-level classes # this creates an object that contains an `ICRS` low-level class, initialized # identically to the first ICRS example further up. sc = coords.SkyCoord(coords.SphericalRepresentation(lon=8 * u.hour, lat=5 * u.deg, distance=1 * u.kpc), frame='icrs') # Other representations and `system` keywords delegate to the appropriate # low-level class. The already-existing registry for user-defined coordinates # will be used by `SkyCoordinate` to figure out what various the `system` # keyword actually means. sc = coords.SkyCoord(ra=8 * u.hour, dec=5 * u.deg, frame='icrs') sc = coords.SkyCoord(l=120 * u.deg, b=5 * u.deg, frame='galactic') # High-level classes can also be initialized directly from low-level objects sc = coords.SkyCoord(coords.ICRS(ra=8 * u.hour, dec=5 * u.deg)) # The next example raises an error because the high-level class must always # have position data. with pytest.raises(ValueError): sc = coords.SkyCoord(coords.FK5(equinox=J2001)) # raises ValueError # similarly, the low-level object can always be accessed # this is how it's supposed to look, but sometimes the numbers get rounded in # funny ways # assert repr(sc.frame) == '<ICRS Coordinate: ra=120.0 deg, dec=5.0 deg>' rscf = repr(sc.frame) assert rscf.startswith('<ICRS Coordinate: (ra, dec) in deg') # and the string representation will be inherited from the low-level class. # same deal, should loook like this, but different archituectures/ python # versions may round the numbers differently # assert repr(sc) == '<SkyCoord (ICRS): ra=120.0 deg, dec=5.0 deg>' rsc = repr(sc) assert rsc.startswith('<SkyCoord (ICRS): (ra, dec) in deg') # Supports a variety of possible complex string formats sc = coords.SkyCoord('8h00m00s +5d00m00.0s', frame='icrs') # In the next example, the unit is only needed b/c units are ambiguous. In # general, we *never* accept ambiguity sc = coords.SkyCoord('8:00:00 +5:00:00.0', unit=(u.hour, u.deg), frame='icrs') # The next one would yield length-2 array coordinates, because of the comma sc = coords.SkyCoord(['8h 5d', '2°2′3″ 0.3rad'], frame='icrs') # It should also interpret common designation styles as a coordinate # NOT YET # sc = coords.SkyCoord('SDSS J123456.89-012345.6', frame='icrs') # but it should also be possible to provide formats for outputting to strings, # similar to `Time`. This can be added right away or at a later date. # transformation is done the same as for low-level classes, which it delegates to sc_fk5_j2001 = sc.transform_to(coords.FK5(equinox=J2001)) assert sc_fk5_j2001.equinox == J2001 # The key difference is that the high-level class remembers frame information # necessary for round-tripping, unlike the low-level classes: sc1 = coords.SkyCoord(ra=8 * u.hour, dec=5 * u.deg, equinox=J2001, frame='fk5') sc2 = sc1.transform_to('icrs') # The next assertion succeeds, but it doesn't mean anything for ICRS, as ICRS # isn't defined in terms of an equinox assert sc2.equinox == J2001 # But it *is* necessary once we transform to FK5 sc3 = sc2.transform_to('fk5') assert sc3.equinox == J2001 assert_allclose(sc1.ra, sc3.ra) # `SkyCoord` will also include the attribute-style access that is in the # v0.2/0.3 coordinate objects. This will *not* be in the low-level classes sc = coords.SkyCoord(ra=8 * u.hour, dec=5 * u.deg, frame='icrs') scgal = sc.galactic assert str(scgal).startswith('<SkyCoord (Galactic): (l, b)') # the existing `from_name` and `match_to_catalog_*` methods will be moved to the # high-level class as convenience functionality. # in remote-data test below! # m31icrs = coords.SkyCoord.from_name('M31', frame='icrs') # assert str(m31icrs) == '<SkyCoord (ICRS) RA=10.68471 deg, Dec=41.26875 deg>' if HAS_SCIPY: cat1 = coords.SkyCoord(ra=[1, 2] * u.hr, dec=[3, 4.01] * u.deg, distance=[5, 6] * u.kpc, frame='icrs') cat2 = coords.SkyCoord(ra=[1, 2, 2.01] * u.hr, dec=[3, 4, 5] * u.deg, distance=[5, 200, 6] * u.kpc, frame='icrs') idx1, sep2d1, dist3d1 = cat1.match_to_catalog_sky(cat2) idx2, sep2d2, dist3d2 = cat1.match_to_catalog_3d(cat2) assert np.any(idx1 != idx2)
# -*- coding: utf-8 -*- from astropy import coordinates as coords from astropy import units as u #<-----------------Classes for representation of coordinate data---------------> # These classes inherit from a common base class and internally contain Quantity # objects, which are arrays (although they may act as scalars, like numpy's # length-0 "arrays") # They can be initialized with a variety of ways that make intuitive sense. # Distance is optional. coords.SphericalRepresentation(lon=8*u.hour, lat=5*u.deg) coords.SphericalRepresentation(lon=8*u.hourangle, lat=5*u.deg) coords.SphericalRepresentation(lon=8*u.hourangle, lat=5*u.deg, distance=10*u.kpc) # In the initial implementation, the lat/lon/distance arguments to the # initializer must be in order. A *possible* future change will be to allow # smarter guessing of the order. E.g. `Latitude` and `Longitude` objects can be # given in any order. coords.SphericalRepresentation(coords.Longitude(8, u.hour), coords.Latitude(5, u.deg)) coords.SphericalRepresentation(coords.Longitude(8, u.hour), coords.Latitude(5, u.deg), coords.Distance(10, u.kpc)) # Arrays of any of the inputs are fine coords.SphericalRepresentation(lon=[8, 9]*u.hourangle, lat=[5, 6]*u.deg) # Default is to copy arrays, but optionally, it can be a reference coords.SphericalRepresentation(lon=[8, 9]*u.hourangle, lat=[5, 6]*u.deg, copy=False) # strings are parsed by `Latitude` and `Longitude` constructors, so no need to # implement parsing in the Representation classes
def main(true_potential_name, fit_potential_name, index, pool, frac_distance_err=1, n_stars=32, n_walkers=None, n_burn=0, n_iterations=1024, overwrite=False, dont_optimize=False, name=None): true_potential = potentials[true_potential_name] _path, _ = os.path.split(os.path.abspath(__file__)) top_path = os.path.abspath(os.path.join(_path, "..")) simulation_path = os.path.join(top_path, "output", "simulations", true_potential_name) output_path = os.path.join(top_path, "output", "orbitfit", true_potential_name, fit_potential_name, "d_{:.1f}percent".format(frac_distance_err)) plot_path = os.path.join(output_path, "plots") sampler_file = os.path.join(output_path, "{}-emcee-{}.h5".format(name, index)) model_file = os.path.join(output_path, "{}-model-{}.pickle".format(name, index)) if os.path.exists(sampler_file) and not overwrite: logger.info("Orbit index {} already complete.".format(index)) return if not os.path.exists(output_path): os.makedirs(output_path) if not os.path.exists(plot_path): os.mkdir(plot_path) logger.info("Potential: {}".format(fit_potential_name)) if fit_potential_name == 'plummer': Model = orbitfit.PlummerOrbitfitModel kw = dict() freeze = dict() freeze['potential_b'] = 10. potential_guess = [6E11] mcmc_potential_std = [1E8] potential_truth = [true_potential.parameters['m'].value] elif fit_potential_name == 'scf': Model = orbitfit.SCFOrbitfitModel kw = dict(nmax=8) freeze = dict() freeze['potential_r_s'] = 10. potential_guess = [6E11] + [1.3, 0, 0, 0, 0, 0, 0, 0, 0 ] # HACK: try this first mcmc_potential_std = [1E8] + [1E-3] * 9 potential_truth = [ true_potential.parameters['m'].value ] + true_potential.parameters['Snlm'].ravel().tolist() elif fit_potential_name == 'triaxialnfw': Model = orbitfit.TriaxialNFWOrbitfitModel kw = dict() freeze = dict() freeze['potential_r_s'] = 20. freeze['potential_a'] = 1. potential_guess = [(200 * u.km / u.s).decompose(galactic).value, 0.8, 0.6] mcmc_potential_std = [1E-5, 1E-3, 1E-3] potential_truth = [ true_potential.parameters['v_c'].value, true_potential.parameters['b'].value, true_potential.parameters['c'].value ] else: raise ValueError( "Invalid potential name '{}'".format(fit_potential_name)) with h5py.File(os.path.join(simulation_path, "mock_stream_data.h5"), "r") as f: g = f[str(index)] pos = g['pos'][:] * u.Unit(f[str(index)]['pos'].attrs['unit']) vel = g['vel'][:] * u.Unit(f[str(index)]['vel'].attrs['unit']) R = g['R'][:] dt = g.attrs['dt'] n_steps = g.attrs['n_steps'] stream = gd.CartesianPhaseSpacePosition(pos=pos, vel=vel) idx = np.concatenate( ([0], np.random.permutation(pos.shape[1])[:n_stars - 1])) true_stream_c, true_stream_v = stream[idx].to_frame( coord.Galactic, **FRAME) true_stream_rot = rotate_sph_coordinate(true_stream_c, R) # intrinsic widths - all are smaller than errors except sky pos, distance rtide = 0.5 * u.kpc phi2_sigma = (rtide / true_stream_rot.distance.mean()).decompose().value d_sigma = rtide.to(u.kpc).value stream_rot = coord.SphericalRepresentation( lon=true_stream_rot.lon, lat=np.random.normal(true_stream_rot.lat.radian, phi2_sigma) * u.radian, distance=np.random.normal(true_stream_rot.distance.value, d_sigma) * u.kpc) # set all proper motions to zero because they shouldn't matter stream_v = [true_stream_v[0] * 0., true_stream_v[1] * 0., true_stream_v[2]] data, err = observe_data(stream_rot, stream_v, frac_distance_err=frac_distance_err, vr_err=10 * u.km / u.s) # freeze all intrinsic widths freeze['phi2_sigma'] = phi2_sigma freeze['d_sigma'] = d_sigma freeze['vr_sigma'] = 5E-4 freeze['mu_sigma'] = 1000. model = Model(data=data, err=err, R=R, dt=dt, n_steps=int(1.5 * n_steps), freeze=freeze, **kw) # save the truth model.true_p = ([stream_rot[0].lat.radian, stream_rot[0].distance.value] + [v[0].decompose(galactic).value for v in stream_v] + potential_truth) # pickle the model with open(model_file, 'wb') as f: pickle.dump(model, f) # starting position for optimization p_guess = ( [true_stream_rot[0].lat.radian, true_stream_rot[0].distance.value] + [v[0].decompose(galactic).value for v in true_stream_v] + potential_guess) logger.debug("ln_posterior at initialization: {}".format(model(p_guess))) if n_walkers is None: n_walkers = 8 * len(p_guess) if not dont_optimize: logger.debug("optimizing ln_posterior first...") res = so.minimize(lambda p: -model(p), x0=p_guess, method='powell') p_best = res.x logger.debug("...done. optimization returned: {}".format(p_best)) if not res.success: pool.close() raise ValueError("Failed to optimize!") logger.debug("ln_posterior at optimized p: {}".format(model(p_best))) # plot the orbit of the optimized parameters orbit = true_potential.integrate_orbit(model._mcmc_sample_to_w0(res.x), dt=dt, nsteps=n_steps) fig, _ = plot_data(data, err, R, gal=False) fig, _ = plot_orbit(orbit, fig=fig, R=R, gal=False) fig.savefig( os.path.join(plot_path, "{}-optimized-{}.png".format(name, index))) mcmc_p0 = emcee.utils.sample_ball(res.x, 1E-3 * np.array(p_best), size=n_walkers) else: # mcmc_std = ([freeze['phi2_sigma'], freeze['d_sigma'], freeze['mu_sigma']] + # [freeze['mu_sigma'], freeze['vr_sigma']] + mcmc_potential_std) # HACK: mcmc_std = ([freeze['phi2_sigma'], freeze['d_sigma'], 1E-4] + [1E-4, freeze['vr_sigma']] + mcmc_potential_std) mcmc_p0 = emcee.utils.sample_ball(p_guess, mcmc_std, size=n_walkers) # now, create initial conditions for MCMC walkers in a small ball around the # optimized parameter vector sampler = emcee.EnsembleSampler(nwalkers=n_walkers, dim=len(p_guess), lnpostfn=model, pool=pool) if n_burn > 0: logger.info("burning in sampler for {} steps".format(n_burn)) pos, _, _ = sampler.run_mcmc(mcmc_p0, N=n_burn) logger.debug("finished burn-in") sampler.reset() else: pos = mcmc_p0 logger.info("running mcmc sampler with {} walkers for {} steps".format( n_walkers, n_iterations)) # restart_p = np.median(sampler.chain[:,-1], axis=0) # mcmc_p0 = emcee.utils.sample_ball(restart_p, 1E-3*restart_p, size=n_walkers) # sampler.reset() _ = sampler.run_mcmc(pos, N=n_iterations) logger.info("finished sampling") pool.close() logger.debug("saving sampler data") with h5py.File(sampler_file, 'w') as g: g['chain'] = sampler.chain g['acceptance_fraction'] = sampler.acceptance_fraction g['lnprobability'] = sampler.lnprobability g.attrs['n_stars'] = n_stars g.attrs['frac_distance_err'] = frac_distance_err if n_iterations > 256: logger.debug("plotting...") flatchain = np.vstack(sampler.chain[:, -256::4]) fig, _ = plot_data(data, err, R, gal=False) for i, link in enumerate(flatchain): orbit = true_potential.integrate_orbit( model._mcmc_sample_to_w0(link), dt=dt, nsteps=n_steps) fig, _ = plot_orbit(orbit, fig=fig, R=R, gal=False, alpha=0.25) if i == 32: break fig.savefig( os.path.join(plot_path, "{}-mcmc-{}.png".format(name, index))) sys.exit(0)
def test__parse_c_err(self): """Test method ``_parse_c_err```.""" expected_dpos = np.array([[0.1, 0.2, 1.0], [0.2, 0.3, 1.0]]) # -------------------------- # with c_err = None # c_err -> <SkyCoord (ICRS): (ra, dec) in deg # [(0.1, 0.2), (0.2, 0.3)]> d_pos = self.inst._parse_c_err(None, self.c) assert np.allclose(d_pos, expected_dpos) # -------------------------- # BaseCoordinateFrame, SkyCoord r_err = coord.SphericalRepresentation( (0.1, 0.2) * u.deg, (0.2, 0.3) * u.deg, 1, ) c_err = coord.ICRS(r_err) d_pos = self.inst._parse_c_err(c_err, self.c) assert np.allclose(d_pos, expected_dpos) # Now with the wrong representation type with pytest.raises(TypeError, match="matching `representation_type`"): self.inst._parse_c_err( coord.ICRS( r_err.to_cartesian(), representation_type=coord.CartesianRepresentation, ), self.c, ) # -------------------------- # BaseRepresentation d_pos = self.inst._parse_c_err(r_err, self.c) assert np.allclose(d_pos, expected_dpos) # Now with the wrong representation type with pytest.raises(TypeError, match="`c_err` must be the same"): self.inst._parse_c_err(r_err.to_cartesian(), self.c) # -------------------------- # Mapping with pytest.raises(NotImplementedError): self.inst._parse_c_err({}, self.c) # -------------------------- # percent error d_pos = self.inst._parse_c_err(10 * u.percent, self.c) assert np.allclose(d_pos, expected_dpos[:, :-1]) # -------------------------- # number d_pos = self.inst._parse_c_err(0.1, self.c) assert d_pos == 0.1 # -------------------------- # callable d_pos = self.inst._parse_c_err(lambda c: 0.1, self.c) assert d_pos == 0.1 # -------------------------- # unrecognized with pytest.raises(NotImplementedError, match="is not yet supported."): self.inst._parse_c_err(NotImplementedError(), self.c)
def test__fix_branch_cuts(self): """Test method ``_fix_branch_cuts``. .. todo:: graphical proof via mpl_test that the point hasn't moved. """ # ------------------------------- # no angular units rep = coord.CartesianRepresentation( x=[1, 2] * u.kpc, y=[3, 4] * u.kpc, z=[5, 6] * u.kpc, ) array = rep._values.view(dtype=np.float64).reshape(rep.shape[0], -1).T got = self.inst._fix_branch_cuts(array, rep.__class__, rep._units) assert got is array # ------------------------------- # UnitSphericalRepresentation # 1) all good rep = coord.UnitSphericalRepresentation( lon=[1, 2] * u.deg, lat=[3, 4] * u.deg, ) array = rep._values.view(dtype=np.float64).reshape(rep.shape[0], -1).T got = self.inst._fix_branch_cuts( array.copy(), rep.__class__, rep._units, ) assert np.allclose(got, array) # 2) needs correction array = np.array([[-360, 0, 360], [-91, 0, 91]]) got = self.inst._fix_branch_cuts( array.copy(), rep.__class__, rep._units, ) assert np.allclose(got, np.array([[-180, 0, 540], [-89, 0, 89]])) # ------------------------------- # SphericalRepresentation # 1) all good rep = coord.SphericalRepresentation( lon=[1, 2] * u.deg, lat=[3, 4] * u.deg, distance=[5, 6] * u.kpc, ) array = rep._values.view(dtype=np.float64).reshape(rep.shape[0], -1).T got = self.inst._fix_branch_cuts( array.copy(), rep.__class__, rep._units, ) assert np.allclose(got, array) # 2) needs correction array = np.array([[-360, 0, 360], [-91, 0, 91], [5, 6, 7]]) got = self.inst._fix_branch_cuts( array.copy(), rep.__class__, rep._units, ) assert np.allclose( got, np.array([[-180, 0, 540], [-89, 0, 89], [5, 6, 7]]), ) # 3) needs correction array = np.array([[-360, 0, 360], [-91, 0, 91], [-5, 6, -7]]) got = self.inst._fix_branch_cuts( array.copy(), rep.__class__, rep._units, ) assert np.allclose( got, np.array([[0, 0, 720], [89, 0, -89], [5, 6, 7]]), ) # ------------------------------- # CylindricalRepresentation # 1) all good rep = coord.CylindricalRepresentation( rho=[5, 6] * u.kpc, phi=[1, 2] * u.deg, z=[3, 4] * u.parsec, ) array = rep._values.view(dtype=np.float64).reshape(rep.shape[0], -1).T got = self.inst._fix_branch_cuts( array.copy(), rep.__class__, rep._units, ) assert np.allclose(got, array) # 2) needs correction array = np.array([[-5, 6, -7], [-180, 0, 180], [-4, 4, 4]]) got = self.inst._fix_branch_cuts( array.copy(), rep.__class__, rep._units, ) assert np.allclose(got, np.array([[5, 6, 7], [0, 0, 360], [-4, 4, 4]])) # ------------------------------- # NotImplementedError with pytest.raises(NotImplementedError): rep = coord.PhysicsSphericalRepresentation( phi=[1, 2] * u.deg, theta=[3, 4] * u.deg, r=[5, 6] * u.kpc, ) array = (rep._values.view(dtype=np.float64).reshape( rep.shape[0], -1).T) self.inst._fix_branch_cuts( array.copy(), coord.PhysicsSphericalRepresentation, rep._units, )