def align_v_positive_lon(
data: coord.BaseCoordinateFrame,
fit_values: T.Dict[str, T.Any],
subsel: T.Union[type(Ellipsis), T.Sequence, slice] = Ellipsis,
):
"""Align the velocity along the positive Longitudinal direction.
Parameters
----------
data : Coordinate
Must have differentials
fit_values : dict
The rotation and origin. Output of `~minimize`
subsel : slice
sub-select a portion of the `pm_lon_coslat` for determining
the average velocity.
Returns
-------
values : dict
`fit_values` with "rotation" adjusted.
"""
values = copy.deepcopy(fit_values) # copy for safety
rotation = values["rotation"]
rot_data = data.transform_to(_make_frame(**values))
# rot_datarot_data.represent_as(coord.SphericalRepresentation)
# # all this to get the rotated velocity
# # TODO faster!
# rot_matrix = reference_to_skyoffset_matrix(
# lon=origin.lon, lat=origin.lat, rotation=rotation
# )
# rot_data = data.transform(rot_matrix).represent_as(
# coord.SphericalRepresentation,
# differential_class=coord.SphericalCosLatDifferential,
# )
# rot_vel = rot_data.differentials["s"]
# get average velocity to determine whether need to rotate.
# TODO determine whether
avg = np.median(rot_data.pm_lon_coslat[subsel])
if avg < 0: # need to flip
rotation = rotation + 180 * u.deg
return values
def test_non_sunpy_frame_to_wcs():
# For a non-SunPy frame, our mapping should return None
frame = BaseCoordinateFrame()
assert solar_frame_to_wcs_mapping(frame) is None
def fit_frame(
data: coord.BaseCoordinateFrame,
origin: coord.BaseCoordinateFrame,
rot0: u.Quantity = 0 * u.deg,
bounds: T.Sequence = (-np.inf, np.inf),
*,
fix_origin: bool = _minimize_defaults["fix_origin"],
use_lmfit: T.Optional[bool] = _minimize_defaults["use_lmfit"],
leastsquares: bool = _minimize_defaults["leastsquares"],
align_v: bool = _minimize_defaults["align_v"],
**fit_kwargs,
):
"""Find Best-Fit Rotated Frame.
Parameters
----------
data : :class:`~astropy.coordinates.CartesianRepresentation`
If `align_v`, must have attached differentials
rot0 : |Quantity|
Initial guess for rotation
origin : :class:`~astropy.coordinates.BaseCoordinateFrame`
location of point on sky about which to rotate
bounds : array-like, optional
Parameter bounds.
See :func:`~trackstream.preprocess.fit_rotated_frame.make_bounds`
::
[[rot_low, rot_up],
[lon_low, lon_up],
[lat_low, lat_up]]
Returns
-------
res : Any
The result of the minimization. Depends on arguments.
Dict[str, Any]
Has fields "rotation" and "origin".
Other Parameters
----------------
use_lmfit : bool, optional, kwarg only
Whether to use ``lmfit`` package
leastsquares : bool, optional, kwarg only
If `use_lmfit` is False, whether to to use
:func:`~scipy.optimize.least_square` or
:func:`~scipy.optimize.minimize` (default)
align_v : bool, optional, kwarg only
Whether to align velocity to be in positive direction
fit_kwargs:
Into whatever minimization package / function is used.
Raises
------
ValueError
If `use_lmfit` and lmfit is not installed.
"""
# ------------------------
# Inputs
# Data
# need to make sure Cartesian representation
# data = data.represent_as(
# coord.CartesianRepresentation,
# differential_class=coord.CartesianDifferential,
# )
# Origin
# We work with a SphericalRepresentation, but
# if isinstance(origin, coord.SkyCoord):
# raise TypeError
origin_frame = origin.__class__
origin = origin.represent_as(coord.SphericalRepresentation)
if use_lmfit is None:
use_lmfit = conf.use_lmfit
x0 = u.Quantity([rot0, origin.lon, origin.lat]).to_value(u.deg)
subsel = fit_kwargs.pop("subsel", Ellipsis)
# ------------------------
# Fitting
if use_lmfit: # lmfit
if not HAS_LMFIT:
raise ValueError("`lmfit` package not available.")
res, values = _fit_representation_lmfit(
data.cartesian,
x0=x0,
bounds=bounds,
fix_origin=fix_origin,
**fit_kwargs,
)
else: # scipy
res, values = _fit_representation_scipy(
data.cartesian,
x0=x0,
bounds=bounds,
fix_origin=fix_origin,
use_leastsquares=leastsquares,
**fit_kwargs,
)
# /def
# ------------------------
# Return
best_rot = values[0]
best_origin = coord.UnitSphericalRepresentation(
lon=values[1],
lat=values[2], # TODO re-add distance
)
best_origin = origin_frame(best_origin)
values = dict(rotation=best_rot, origin=best_origin)
if align_v:
values = align_v_positive_lon(data, values, subsel=subsel)
return res, values