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,
)
