From e356b020f5d0a36786a562aac6098adc62087829 Mon Sep 17 00:00:00 2001 From: Joachim Moeyens Date: Mon, 28 Aug 2023 13:46:59 -0700 Subject: [PATCH 1/3] Refactor TestOrbit class --- thor/orbit.py | 438 ++++++++++++-------------------------------------- 1 file changed, 104 insertions(+), 334 deletions(-) diff --git a/thor/orbit.py b/thor/orbit.py index e5121a68..b868dec2 100644 --- a/thor/orbit.py +++ b/thor/orbit.py @@ -1,361 +1,131 @@ -import logging - -import numpy as np -from numba import jit - -from .coordinates import _convertSphericalToCartesian, transformCoordinates - -# from .projections import cartesianToGnomonic - -X_AXIS = np.array([1.0, 0.0, 0.0]) -Y_AXIS = np.array([0.0, 1.0, 0.0]) -Z_AXIS = np.array([0.0, 0.0, 1.0]) - -logger = logging.getLogger(__name__) - -__all__ = [ - "calcNae", - "calcDelta", - "calcXae", - "calcXa", - "calcNhat", - "calcR1", - "calcR2", - "TestOrbit", -] +import uuid +from typing import Optional, TypeVar, Union + +import quivr as qv +from adam_core.coordinates import ( + CartesianCoordinates, + CometaryCoordinates, + KeplerianCoordinates, + SphericalCoordinates, +) +from adam_core.observers import Observers +from adam_core.orbits import Ephemeris, Orbits +from adam_core.propagator import PYOORB + +CoordinateType = TypeVar( + "CoordinateType", + bound=Union[ + CartesianCoordinates, + SphericalCoordinates, + KeplerianCoordinates, + CometaryCoordinates, + ], +) class TestOrbit: - """ - TestOrbit: Class that calculates and stores the rotation matrices - for a guess of heliocentric distance and velocity. To be used in - tandem with the Cell class. - - Parameters - ---------- - cartesian : `~numpy.ndarray` (1, 6) - Cartesian ecliptic orbital elements with postions in units of au - and velocities in units of au per day. - t0 : `~astropy.time.core.Time` (1) - Epoch at which orbital elements are defined. - """ - - def __init__(self, cartesian, epoch): - self.cartesian = cartesian - self.epoch = epoch - - def prepare(self): + def __init__( + self, + coordinates: CoordinateType, + orbit_id: Optional[str] = None, + object_id: Optional[str] = None, + ): """ - Calculate rotation matrices. - - Populates the following class properties: - n_hat : vector normal to the plane of orbit - R1 : rotation matrix to rotate towards x-y plane - R2 : rotation matrix to rotate towards x-axis - M : final rotation matrix + Create a test orbit from a set of orbital elements. - Returns - ------- - None + Parameters + ---------- + coordinates : `~adam_core.coordinates.cartesian.CartesianCoordinates`, + `~adam_core.coordinates.spherical.SphericalCoordinates`, + `~adam_core.coordinates.keplerian.KeplerianCoordinates`, + `~adam_core.coordinates.cometary.CometaryCoordinates` + The orbital elements that define this test orbit. Can be any representation but will + be stored internally as Cartesian elements. + orbit_id : str, optional + Orbit ID. If not provided, a random UUID will be generated. + object_id : str, optional + Object ID, if it exists. """ - logger.debug("Calculating vector normal to plane of orbit...") - self.n_hat = calcNhat(self.cartesian.reshape(1, -1))[0] - - logger.debug("Calculating R1 rotation matrix...") - self.R1 = calcR1(self.n_hat) - self.x_a_xy = np.array(self.R1 @ self.cartesian[:3]) + # Test orbits should be singletons + assert len(coordinates) == 1 + + # Test orbit selection will likely occur in a non-Cartesian coordinate system + # so we should accept any coordinate system and convert to Cartesian as the + # most stable representation + if not isinstance(coordinates, CartesianCoordinates): + cartesian_coordinates = coordinates.to_cartesian() + else: + cartesian_coordinates = coordinates + + if orbit_id is not None: + self.orbit_id = orbit_id + else: + self.orbit_id = uuid.uuid4().hex + + if object_id is not None: + self.object_id = object_id + + self._orbit = Orbits.from_kwargs( + orbit_id=[self.orbit_id], + object_id=[self.object_id], + coordinates=cartesian_coordinates, + ) - logger.debug("Calculating R2 rotation matrix...") - self.R2 = calcR2(self.x_a_xy) + @classmethod + def from_orbits(cls, orbits): + assert len(orbits) == 1 + return cls( + orbits.coordinates, orbits.orbit_id[0].as_py(), orbits.object_id[0].as_py() + ) - logger.debug("Calculating final rotation matrix...") - self.M = self.R2 @ self.R1 - return + @property + def orbit(self): + return self._orbit - def applyToObservations(self, observations): + def propagate(self, times, propagator=PYOORB(), max_processes=1) -> Orbits: """ - Apply the prepared rotations to the given observations. Adds the gnomonic - plane coordinates to observations (columns: theta_x_deg, theta_y_deg) + Propagate this test orbit to the given times. Parameters ---------- - observations : `~pandas.DataFrame` - DataFrame of observations defined at the same epoch as this test orbit, - to project into the test orbit's frame. + times : `~astropy.time.core.Time` + Times to which to propagate the orbit. + propagator : `~adam_core.propagator.propagator.Propagator`, optional + Propagator to use to propagate the orbit. Defaults to PYOORB. + num_processes : int, optional + Number of processes to use to propagate the orbit. Defaults to 1. Returns ------- - None + propagated_orbit : `~adam_core.orbits.orbits.Orbits` + The test orbit propagated to the given times. """ - logger.debug("Applying rotation matrices to observations...") - logger.debug("Converting to ecliptic coordinates...") - coords_eq = observations[["RA_deg", "Dec_deg"]].values - coords_eq = np.hstack([np.ones((len(coords_eq), 1)), coords_eq]) - coords_ec = transformCoordinates( - coords_eq, - "equatorial", - "ecliptic", - representation_in="spherical", - representation_out="spherical", - ) - - logger.debug("Calculating object to observer unit vector...") - n_ae = calcNae(coords_ec[:, 1:3]) - x_e = observations[["obs_x", "obs_y", "obs_z"]].values - - r = np.linalg.norm(self.cartesian[:3]) - logger.debug(f"Calculating object to observer distance assuming r = {r} AU...") - delta = calcDelta(r, x_e, n_ae) - - logger.debug("Calculating object to observer position vector...") - x_ae = calcXae(delta, n_ae) - - logger.debug("Calculating heliocentric object position vector...") - x_a = calcXa(x_ae, x_e) - - logger.debug( - "Applying rotation matrix M to heliocentric object position vector..." + return propagator.propagate_orbits( + self.orbit, times, max_processes=max_processes, chunk_size=1 ) - coords_cart_rotated = np.array(self.M @ x_a.T).T - - logger.debug("Performing gnomonic projection...") - gnomonic_coords = cartesianToGnomonic(coords_cart_rotated) - - observations["obj_x'"] = x_a[:, 0] - observations["obj_y'"] = x_a[:, 1] - observations["obj_z'"] = x_a[:, 2] - observations["obj_x''"] = coords_cart_rotated[:, 0] - observations["obj_y''"] = coords_cart_rotated[:, 1] - observations["obj_z''"] = coords_cart_rotated[:, 2] - observations["theta_x_deg"] = np.degrees(gnomonic_coords[:, 0]) - observations["theta_y_deg"] = np.degrees(gnomonic_coords[:, 1]) - observations["test_obj_x"] = self.cartesian[0] - observations["test_obj_y"] = self.cartesian[1] - observations["test_obj_z"] = self.cartesian[2] - observations["test_obj_vx"] = self.cartesian[3] - observations["test_obj_vy"] = self.cartesian[4] - observations["test_obj_vz"] = self.cartesian[5] - - coords_rot_test = self.M @ self.cartesian[:3].T - coords_rot_testv = self.M @ self.cartesian[3:].T - observations["test_obj_x''"] = coords_rot_test[0] - observations["test_obj_y''"] = coords_rot_test[1] - observations["test_obj_z''"] = coords_rot_test[2] - observations["test_obj_vx''"] = coords_rot_testv[0] - observations["test_obj_vy''"] = coords_rot_testv[1] - observations["test_obj_vz''"] = coords_rot_testv[2] - return - def applyToEphemeris(self, ephemeris): + def generate_ephemeris( + self, observers, propagator=PYOORB(), max_processes=1 + ) -> qv.MultiKeyLinkage[Ephemeris, Observers]: """ - Apply the prepared rotations to the given ephemerides. Adds the gnomonic - plane coordinates to observations (columns: theta_x_deg, theta_y_deg, vtheta_x, and vtheta_y) + Generate ephemeris for this test orbit at the given observers. Parameters ---------- - ephemeris : `~pandas.DataFrame` - DataFrame of ephemeris generated by a THOR backend defined at the same epoch as this test orbit, - to project into the test orbit's frame. + observers : `~adam_core.observers.Observers` + Observers from which to generate ephemeris. + propagator : `~adam_core.propagator.propagator.Propagator`, optional + Propagator to use to propagate the orbit. Defaults to PYOORB. + num_processes : int, optional + Number of processes to use to propagate the orbit. Defaults to 1. Returns ------- - None + ephemeris : qv.MultiKeyLinkage[ + `~adam_core.orbits.ephemeris.Ephemeris`, + `~adam_core.observers.observers.Observers`] + The ephemeris of the test orbit at the given observers. """ - raise NotImplementedError - - -@jit("f8[:,:](f8[:,:])", nopython=True, cache=True) -def calcNae(coords_ec_ang): - """ - Convert angular ecliptic coordinates to - to a cartesian unit vector. - - Input coordinate array should have shape (N, 2): - np.array([[lon1, lat1], - [lon2, lat2], - [lon3, lat3], - ..... - [lonN, latN]]) - - Parameters - ---------- - coords_ec_ang : `~numpy.ndarray` (N, 2) - Ecliptic longitude and latitude in degrees. - - Returns - ------- - N_ae : `~numpy.ndarray` (N, 3) - Cartesian unit vector in direction of provided - angular coordinates. - """ - N = len(coords_ec_ang) - rho = np.ones(N) - lon = np.radians(coords_ec_ang[:, 0]) - lat = np.radians(coords_ec_ang[:, 1]) - velocities = np.zeros(len(rho)) - x, y, z, vx, vy, vz = _convertSphericalToCartesian( - rho, lon, lat, velocities, velocities, velocities - ) - - n_ae = np.zeros((N, 3)) - n_ae[:, 0] = x - n_ae[:, 1] = y - n_ae[:, 2] = z - return n_ae - - -@jit("f8[:](f8, f8[:,:], f8[:,:])", nopython=True, cache=True) -def calcDelta(r, x_e, n_ae): - """ - Calculate topocentric distance to the asteroid. - - Parameters - ---------- - r : float (1) - Heliocentric/barycentric distance in arbitrary units. - x_e : `~numpy.ndarray` (N, 3) - Topocentric position vector in same units as r. - n_ae : `~numpy.ndarray` (N, 3) - Unit vector in direction of asteroid from the topocentric position - in same units as r. - - Returns - ------- - delta : `~numpy.ndarray` (N) - Distance from topocenter to asteroid in units of r. - """ - N = len(x_e) - delta = np.zeros(N) - rsq = r**2 - for i in range(N): - n_ae_i = np.ascontiguousarray(n_ae[i]) - x_e_i = np.ascontiguousarray(x_e[i]) - ndotxe = np.dot(n_ae_i, x_e_i) - delta[i] = -ndotxe + np.sqrt(ndotxe**2 + rsq - np.linalg.norm(x_e_i) ** 2) - return delta - - -@jit("f8[:,:](f8[:], f8[:,:])", nopython=True, cache=True) -def calcXae(delta, n_ae): - """ - Calculate the topocenter to asteroid position vector. - - Parameters - ---------- - delta : float - Distance from the topocenter to asteroid in arbitrary units. - n_ae : `~numpy.ndarray` (3) - Unit vector in direction of asteroid from the topocentric position - in same units as delta. - - Returns - ------- - x_ae : `~numpy.ndarray` (N, 3) - Topocenter to asteroid position vector in units of delta. - """ - x_ae = np.zeros_like(n_ae) - for i, (delta_i, n_ae_i) in enumerate(zip(delta, n_ae)): - x_ae[i] = delta_i * n_ae_i - return x_ae - - -@jit("f8[:,:](f8[:,:],f8[:,:])", nopython=True, cache=True) -def calcXa(x_ae, x_e): - """ - Calculate the asteroid position vector. - - Parameters - ---------- - x_ae : `~numpy.ndarray` (3) - Topocenter to asteroid position vector in arbitrary units. - x_e : `~numpy.ndarray` (3) - Topocentric position vector in same units as x_ae. - - Returns - ------- - x_a : `~numpy.ndarray` (3) - Asteroid position vector in units of x_ae. - """ - return x_ae + x_e - - -@jit("f8[:,:](f8[:,:])", nopython=True, cache=True) -def calcNhat(state_vector): - """ - Calculate the unit vector normal to the plane of the orbit. - - Parameters - ---------- - state_vector : `~numpy.ndarray` (N, 6) - Cartesian state vectors in units of au and au p day. - - Returns - ------- - n_hat : `~numpy.ndarray` (N, 3) - Unit vector normal to plane of orbit. - """ - # Cross product of r and v is equal to the angular momentum - # vector of the orbit (which is always perpendicular to the plane of the - # orbit at the epoch) - rv = np.cross(state_vector[:, 0:3], state_vector[:, 3:6]) - # Ideally, using rv = np.linalg.norm(rv, axis=1, keepdims=True) would be cleaner. - # However, both the axis and keepdims kwargs are not supported by numba - # so reshape and normalize manually. - rv_norm = np.zeros((len(rv), 1)) - for i in range(len(rv)): - rv_norm[i] = np.linalg.norm(rv[i]) - n_hat = rv / rv_norm - return n_hat - - -@jit("f8[:,:](f8[:])", nopython=True, cache=True) -def calcR1(n_hat): - """ - Calculate the rotation matrix that would rotate the - position vector x_ae to the x-y plane. - - Parameters - ---------- - n_hat : `~numpy.ndarray` (3) - Unit vector normal to plane of orbit. - - Returns - ------- - R1 : `~numpy.matrix` (3, 3) - Rotation matrix. - """ - n_hat_ = np.ascontiguousarray(n_hat) - # Find the rotation axis v - v = np.cross(n_hat_, Z_AXIS) - # Calculate the cosine of the rotation angle, equivalent to the cosine of the - # inclination - c = np.dot(n_hat_, Z_AXIS) - # Compute the skew-symmetric cross-product of the rotation axis vector v - vp = np.array([[0, -v[2], v[1]], [v[2], 0, -v[0]], [-v[1], v[0], 0]]) - # Calculate R1 - R1 = np.identity(3) + vp + np.linalg.matrix_power(vp, 2) * (1 / (1 + c)) - return R1 - - -@jit("f8[:,:](f8[:])", nopython=True, cache=True) -def calcR2(x_a_xy): - """ - Calculate the rotation matrix that would rotate a vector in - the x-y plane to the x-axis. - - Parameters - ---------- - x_a_xy : `~numpy.ndarray` (3) - Barycentric asteroid position vector rotated to the x-y plane. - - Returns - ------- - R2 : `~numpy.ndarray` (3, 3) - Rotation ndarray. - """ - x_a_xy = x_a_xy / np.linalg.norm(x_a_xy) - # Assuming the vector x_a_xy has been normalized, and is in the xy plane. - ca = x_a_xy[0] - sa = x_a_xy[1] - R2 = np.array([[ca, sa, 0.0], [-sa, ca, 0.0], [0.0, 0.0, 1.0]]) - return R2 + return propagator.generate_ephemeris( + self.orbit, observers, max_processes=max_processes, chunk_size=1 + ) From 302d627978bb11d39cbe329547ccc247db1f855f Mon Sep 17 00:00:00 2001 From: Joachim Moeyens Date: Mon, 28 Aug 2023 15:47:58 -0700 Subject: [PATCH 2/3] Add assume_heliocentric_distance --- thor/orbit.py | 57 +++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 57 insertions(+) diff --git a/thor/orbit.py b/thor/orbit.py index b868dec2..3b9d06de 100644 --- a/thor/orbit.py +++ b/thor/orbit.py @@ -1,12 +1,15 @@ import uuid from typing import Optional, TypeVar, Union +import numpy as np import quivr as qv from adam_core.coordinates import ( CartesianCoordinates, CometaryCoordinates, KeplerianCoordinates, + OriginCodes, SphericalCoordinates, + transform_coordinates, ) from adam_core.observers import Observers from adam_core.orbits import Ephemeris, Orbits @@ -129,3 +132,57 @@ def generate_ephemeris( return propagator.generate_ephemeris( self.orbit, observers, max_processes=max_processes, chunk_size=1 ) + + +def assume_heliocentric_distance( + r_mag: float, coords: SphericalCoordinates, origin_coords: CartesianCoordinates +) -> SphericalCoordinates: + """ + Given a heliocentric distance, for all coordinates that do not have a topocentric distance defined (rho), calculate + the topocentric distance assuming the coordinates are located at the given heliocentric distance. + + Parameters + ---------- + r_mag : float + Heliocentric distance to assume for the coordinates with missing topocentric distance. This is + typically the same distance as the heliocentric distance of test orbit at the time + of the coordinates. + coords : `~adam_core.coordinates.spherical.SphericalCoordinates` + Coordinates to assume the heliocentric distance for. + origin_coords : `~adam_core.coordinates.cartesian.CartesianCoordinates` + Heliocentric coordinates of the origin of the topocentric coordinates. + + Returns + ------- + coords : `~adam_core.coordinates.spherical.SphericalCoordinates` + Coordinates with the missing topocentric distance replaced with the calculated topocentric distance. + """ + assert len(origin_coords) == 1 + assert np.all(origin_coords.origin == OriginCodes.SUN) + + # Extract the topocentric distance and topocentric radial velocity from the coordinates + rho = coords.rho.to_numpy(zero_copy_only=False) + vrho = coords.vrho.to_numpy(zero_copy_only=False) + + # Transform the coordinates to the ecliptic frame by assuming they lie on a unit sphere + # (this assumption will only be applied to coordinates with missing rho values) + coords_eq_unit = coords.to_unit_sphere(only_missing=True) + coords_ec = transform_coordinates( + coords_eq_unit, SphericalCoordinates, frame_out="ecliptic" + ) + + # Transform the coordinates to cartesian and calculate the unit vectors pointing + # from the origin to the coordinates + coords_ec_xyz = coords_ec.to_cartesian() + unit_vectors = coords_ec_xyz.r_hat + + # Calculate the topocentric distance such that the heliocentric distance to the coordinate + # is r_mag + dotprod = np.sum(unit_vectors * origin_coords.r, axis=1) + delta = -dotprod + np.sqrt(dotprod**2 + r_mag**2 - origin_coords.r_mag**2) + + # Where rho was not defined, replace it with the calculated topocentric distance + coords_ec = coords_ec.set_column("rho", np.where(np.isnan(rho), delta, rho)) + coords_ec = coords_ec.set_column("vrho", vrho) + + return coords_ec From 6f7b1a9317457226f54cd12cacbae3c30c3ee572 Mon Sep 17 00:00:00 2001 From: Joachim Moeyens Date: Mon, 28 Aug 2023 15:58:37 -0700 Subject: [PATCH 3/3] Add TestOrbit.range_observations --- thor/orbit.py | 117 ++++++++++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 114 insertions(+), 3 deletions(-) diff --git a/thor/orbit.py b/thor/orbit.py index 3b9d06de..ad103735 100644 --- a/thor/orbit.py +++ b/thor/orbit.py @@ -2,18 +2,24 @@ from typing import Optional, TypeVar, Union import numpy as np +import pyarrow as pa import quivr as qv from adam_core.coordinates import ( CartesianCoordinates, CometaryCoordinates, + CoordinateCovariances, KeplerianCoordinates, + Origin, OriginCodes, SphericalCoordinates, transform_coordinates, ) +from adam_core.observations.detections import PointSourceDetections +from adam_core.observations.exposures import Exposures from adam_core.observers import Observers from adam_core.orbits import Ephemeris, Orbits -from adam_core.propagator import PYOORB +from adam_core.propagator import PYOORB, Propagator +from astropy.time import Time CoordinateType = TypeVar( "CoordinateType", @@ -26,6 +32,13 @@ ) +class RangedPointSourceDetections(qv.Table): + + id = qv.StringColumn() + exposure_id = qv.StringColumn() + coordinates = SphericalCoordinates.as_column() + + class TestOrbit: def __init__( self, @@ -85,7 +98,12 @@ def from_orbits(cls, orbits): def orbit(self): return self._orbit - def propagate(self, times, propagator=PYOORB(), max_processes=1) -> Orbits: + def propagate( + self, + times: Time, + propagator: Propagator = PYOORB(), + max_processes: Optional[int] = 1, + ) -> Orbits: """ Propagate this test orbit to the given times. @@ -108,7 +126,10 @@ def propagate(self, times, propagator=PYOORB(), max_processes=1) -> Orbits: ) def generate_ephemeris( - self, observers, propagator=PYOORB(), max_processes=1 + self, + observers: Observers, + propagator: Propagator = PYOORB(), + max_processes: Optional[int] = 1, ) -> qv.MultiKeyLinkage[Ephemeris, Observers]: """ Generate ephemeris for this test orbit at the given observers. @@ -133,6 +154,96 @@ def generate_ephemeris( self.orbit, observers, max_processes=max_processes, chunk_size=1 ) + def range_observations( + self, + observations: qv.Linkage[PointSourceDetections, Exposures], + max_processes: Optional[int] = 1, + ) -> RangedPointSourceDetections: + """ + Given a set of observations, propagate this test orbit to the times of the observations and calculate the + topocentric distance (range) assuming they lie at the same heliocentric distance as the test orbit. + + Parameters + ---------- + observations : qv.Linkage[PointSourceDetections, Exposures] + Observations to range. + max_processes : int, optional + Number of processes to use to propagate the orbit. Defaults to 1. + + Returns + ------- + ranged_point_source_detections : `~thor.orbit.RangedPointSourceDetections` + The ranged detections. + """ + exposures = observations.right_table + ephemeris = self.generate_ephemeris( + exposures.observers(), max_processes=max_processes + ) + # Get the light-time corrected state vector: the state vector + # at the time where the light reflected/emitted from the object + # would have reached the observer. + # We will use this state vector to get the heliocentric distance of + # the object at the time of the exposure. + # TODO: We could use other adam_core functionality to calculate this if need + # be and we may need to if we plan on mapping covariances. + propagated_orbit = ephemeris.left_table.aberrated_coordinates + + # TODO: We could investigate using concurrent futures here to parallelize + # this loop + rpsds = [] + for propagated_orbit_i, exposure_i in zip(propagated_orbit, exposures): + + # Select the detections that belong to this exposure + detections_i = observations.select_left(exposure_i.id[0]) + + # Get the heliocentric distance of the object at the time of the exposure + r_mag = propagated_orbit_i.r_mag[0] + + # Get the observer's heliocentric coordinates + observer_i = exposure_i.observers() + + # Create an array of observatory codes for the detections + num_detections = len(detections_i) + observatory_codes = np.repeat( + exposure_i.observatory_code[0].as_py(), num_detections + ) + + # The following can be replaced with: + # coords = detections_i.to_spherical(observatory_codes) + # Start replacement: + sigma_data = np.vstack( + [ + pa.nulls(num_detections, pa.float64()), + detections_i.ra_sigma.to_numpy(zero_copy_only=False), + detections_i.dec_sigma.to_numpy(zero_copy_only=False), + pa.nulls(num_detections, pa.float64()), + pa.nulls(num_detections, pa.float64()), + pa.nulls(num_detections, pa.float64()), + ] + ).T + coords = SphericalCoordinates.from_kwargs( + lon=detections_i.ra, + lat=detections_i.dec, + time=detections_i.time, + covariance=CoordinateCovariances.from_sigmas(sigma_data), + origin=Origin.from_kwargs(code=observatory_codes), + frame="equatorial", + ) + # End replacement (only once + # https://github.com/B612-Asteroid-Institute/adam_core/pull/45 is merged) + + rpsds.append( + RangedPointSourceDetections.from_kwargs( + id=detections_i.id, + exposure_id=detections_i.exposure_id, + coordinates=assume_heliocentric_distance( + r_mag, coords, observer_i.coordinates + ), + ) + ) + + return qv.concatenate(rpsds) + def assume_heliocentric_distance( r_mag: float, coords: SphericalCoordinates, origin_coords: CartesianCoordinates