From ec8605894a4ccb0e5c5cd2bcc0735423d312177a Mon Sep 17 00:00:00 2001 From: yeager7 Date: Thu, 18 Jul 2024 15:58:44 -0700 Subject: [PATCH 01/55] new file to handle coordinate conversions --- ssapy/coordinates.py | 653 +++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 653 insertions(+) create mode 100644 ssapy/coordinates.py diff --git a/ssapy/coordinates.py b/ssapy/coordinates.py new file mode 100644 index 0000000..3198912 --- /dev/null +++ b/ssapy/coordinates.py @@ -0,0 +1,653 @@ +# flake8: noqa: E501 + +from .constants import EARTH_RADIUS, WGS84_EARTH_OMEGA +from .accel import AccelKepler +from .body import get_body, MoonPosition +from .compute import groundTrack, rv +from .constants import RGEO +from .orbit import Orbit +from .propagator import RK78Propagator +from .utils import hms_to_dd, dd_to_hms, dd_to_dms + +import numpy as np +from astropy.time import Time + +# VECTOR FUNCTIONS FOR COORDINATE MATH +def unit_vector(vector): + """ Returns the unit vector of the vector.""" + return vector / np.linalg.norm(vector) + + +def getAngle(a, b, c): # a,b,c where b is the vertex + a = np.atleast_2d(a) + b = np.atleast_2d(b) + c = np.atleast_2d(c) + ba = np.subtract(a, b) + bc = np.subtract(c, b) + cosine_angle = np.sum(ba * bc, axis=-1) / (np.linalg.norm(ba, axis=-1) * np.linalg.norm(bc, axis=-1)) + return np.arccos(cosine_angle) + + +def angle_between_vectors(vector1, vector2): + return np.arccos(np.clip(np.dot(vector1, vector2) / (np.linalg.norm(vector1) * np.linalg.norm(vector2)), -1.0, 1.0)) + + +def rotation_matrix_from_vectors(vec1, vec2): + """ Find the rotation matrix that aligns vec1 to vec2 + :param vec1: A 3d "source" vector + :param vec2: A 3d "destination" vector + :return mat: A transform matrix (3x3) which when applied to vec1, aligns it with vec2. + """ + a, b = (vec1 / np.linalg.norm(vec1)).reshape(3), (vec2 / np.linalg.norm(vec2)).reshape(3) + v = np.cross(a, b) + c = np.dot(a, b) + s = np.linalg.norm(v) + kmat = np.array([[0, -v[2], v[1]], [v[2], 0, -v[0]], [-v[1], v[0], 0]]) + rotation_matrix = np.eye(3) + kmat + kmat.dot(kmat) * ((1 - c) / (s**2)) + return rotation_matrix + + +def normed(arr): + return arr / np.sqrt(np.einsum("...i,...i", arr, arr))[..., None] + + +def einsum_norm(a, indices='ij,ji->i'): + return np.sqrt(np.einsum(indices, a, a)) + + +def normSq(arr): + return np.einsum("...i,...i", arr, arr) + + +def norm(arr): + return np.sqrt(np.einsum("...i,...i", arr, arr)) + + +def rotate_vector(v_unit, theta, phi): + v_unit = v_unit / np.linalg.norm(v_unit, axis=-1) + if np.all(np.abs(v_unit) != np.max(np.abs(v_unit))): + perp_vector = np.cross(v_unit, np.array([1, 0, 0])) + else: + perp_vector = np.cross(v_unit, np.array([0, 1, 0])) + perp_vector /= np.linalg.norm(perp_vector) + + theta = np.radians(theta) + phi = np.radians(phi) + cos_theta = np.cos(theta) + sin_theta = np.sin(theta) + cos_phi = np.cos(phi) + sin_phi = np.sin(phi) + + R1 = np.array([ + [cos_theta + (1 - cos_theta) * perp_vector[0]**2, + (1 - cos_theta) * perp_vector[0] * perp_vector[1] - sin_theta * perp_vector[2], + (1 - cos_theta) * perp_vector[0] * perp_vector[2] + sin_theta * perp_vector[1]], + [(1 - cos_theta) * perp_vector[1] * perp_vector[0] + sin_theta * perp_vector[2], + cos_theta + (1 - cos_theta) * perp_vector[1]**2, + (1 - cos_theta) * perp_vector[1] * perp_vector[2] - sin_theta * perp_vector[0]], + [(1 - cos_theta) * perp_vector[2] * perp_vector[0] - sin_theta * perp_vector[1], + (1 - cos_theta) * perp_vector[2] * perp_vector[1] + sin_theta * perp_vector[0], + cos_theta + (1 - cos_theta) * perp_vector[2]**2] + ]) + + # Apply the rotation matrix to v_unit to get the rotated unit vector + v1 = np.dot(R1, v_unit) + + # Rotation matrix for rotation about v_unit + R2 = np.array([[cos_phi + (1 - cos_phi) * v_unit[0]**2, + (1 - cos_phi) * v_unit[0] * v_unit[1] - sin_phi * v_unit[2], + (1 - cos_phi) * v_unit[0] * v_unit[2] + sin_phi * v_unit[1]], + [(1 - cos_phi) * v_unit[1] * v_unit[0] + sin_phi * v_unit[2], + cos_phi + (1 - cos_phi) * v_unit[1]**2, + (1 - cos_phi) * v_unit[1] * v_unit[2] - sin_phi * v_unit[0]], + [(1 - cos_phi) * v_unit[2] * v_unit[0] - sin_phi * v_unit[1], + (1 - cos_phi) * v_unit[2] * v_unit[1] + sin_phi * v_unit[0], + cos_phi + (1 - cos_phi) * v_unit[2]**2]]) + + v2 = np.dot(R2, v1) + return v2 / np.linalg.norm(v2, axis=-1) + + +def rotate_points_3d(points, axis=np.array([0, 0, 1]), theta=-np.pi / 2): + """ + Rotate a set of 3D points about a 3D axis by an angle theta in radians. + + Args: + points (np.ndarray): The set of 3D points to rotate, as an Nx3 array. + axis (np.ndarray): The 3D axis to rotate about, as a length-3 array. Default is the z-axis. + theta (float): The angle to rotate by, in radians. Default is pi/2. + + Returns: + np.ndarray: The rotated set of 3D points, as an Nx3 array. + """ + # Normalize the axis to be a unit vector + axis = axis / np.linalg.norm(axis) + + # Compute the quaternion representing the rotation + qw = np.cos(theta / 2) + qx, qy, qz = axis * np.sin(theta / 2) + + # Construct the rotation matrix from the quaternion + R = np.array([ + [1 - 2 * qy**2 - 2 * qz**2, 2 * qx * qy - 2 * qz * qw, 2 * qx * qz + 2 * qy * qw], + [2 * qx * qy + 2 * qz * qw, 1 - 2 * qx**2 - 2 * qz**2, 2 * qy * qz - 2 * qx * qw], + [2 * qx * qz - 2 * qy * qw, 2 * qy * qz + 2 * qx * qw, 1 - 2 * qx**2 - 2 * qy**2] + ]) + + # Apply the rotation matrix to the set of points + rotated_points = np.dot(R, points.T).T + + return rotated_points + + +def perpendicular_vectors(v): + """Returns two vectors that are perpendicular to v and each other.""" + # Check if v is the zero vector + if np.allclose(v, np.zeros_like(v)): + raise ValueError("Input vector cannot be the zero vector.") + + # Choose an arbitrary non-zero vector w that is not parallel to v + w = np.array([1., 0., 0.]) + if np.allclose(v, w) or np.allclose(v, -w): + w = np.array([0., 1., 0.]) + u = np.cross(v, w) + if np.allclose(u, np.zeros_like(u)): + w = np.array([0., 0., 1.]) + u = np.cross(v, w) + w = np.cross(v, u) + + return u, w + + +def points_on_circle(r, v, rad, num_points=4): + # Convert inputs to numpy arrays + r = np.array(r) + v = np.array(v) + + # Find the perpendicular vectors to the given vector v + if np.all(v[:2] == 0): + if np.all(v[2] == 0): + raise ValueError("The given vector v must not be the zero vector.") + else: + u = np.array([1, 0, 0]) + else: + u = np.array([-v[1], v[0], 0]) + u = u / np.linalg.norm(u) + w = np.cross(u, v) + w_norm = np.linalg.norm(w) + if w_norm < 1e-15: + # v is parallel to z-axis + w = np.array([0, 1, 0]) + else: + w = w / w_norm + # Generate a sequence of angles for equally spaced points + angles = np.linspace(0, 2 * np.pi, num_points, endpoint=False) + + # Compute the x, y, z coordinates of each point on the circle + x = rad * np.cos(angles) * u[0] + rad * np.sin(angles) * w[0] + y = rad * np.cos(angles) * u[1] + rad * np.sin(angles) * w[1] + z = rad * np.cos(angles) * u[2] + rad * np.sin(angles) * w[2] + + # Apply rotation about z-axis by 90 degrees + rot_matrix = np.array([[0, 1, 0], [-1, 0, 0], [0, 0, 1]]) + rotated_points = np.dot(rot_matrix, np.column_stack((x, y, z)).T).T + + # Translate the rotated points to the center point r + points_rotated = rotated_points + r.reshape(1, 3) + + return points_rotated + + +def dms_to_rad(coords): + from astropy.coordinates import Angle + if isinstance(coords, (list, tuple)): + return [Angle(coord).radian for coord in coords] + else: + return Angle(coords).radian + return + + +def dms_to_deg(coords): + from astropy.coordinates import Angle + if isinstance(coords, (list, tuple)): + return [Angle(coord).deg for coord in coords] + else: + return Angle(coords).deg + return + + +def rad0to2pi(angles): + return (2 * np.pi + angles) * (angles < 0) + angles * (angles > 0) + + +def deg0to360(array_): + try: + return [i % 360 for i in array_] + except TypeError: + return array_ % 360 + + +def deg0to360array(array_): + return [i % 360 for i in array_] + + +def deg90to90(val_in): + if hasattr(val_in, "__len__"): + val_out = [] + for i, v in enumerate(val_in): + while v < -90: + v += 90 + while v > 90: + v -= 90 + val_out.append(v) + else: + while val_in < -90: + val_in += 90 + while val_in > 90: + val_in -= 90 + val_out = val_in + return val_out + + +def deg90to90array(array_): + return [i % 90 for i in array_] + + +def cart2sph_deg(x, y, z): + hxy = np.hypot(x, y) + r = np.hypot(hxy, z) + el = np.arctan2(z, hxy) * (180 / np.pi) + az = (np.arctan2(y, x)) * (180 / np.pi) + return az, el, r + + +def cart_to_cyl(x, y, z): + r = np.linalg.norm([x, y]) + theta = np.arctan2(y, x) + return r, theta, z + + +def inert2rot(x, y, xe, ye, xs=0, ys=0): # Places Earth at (-1,0) + earth_theta = np.arctan2(ye - ys, xe - xs) + theta = np.arctan2(y - ys, x - xs) + distance = np.sqrt(np.power((x - xs), 2) + np.power((y - ys), 2)) + xrot = distance * np.cos(np.pi + (theta - earth_theta)) + yrot = distance * np.sin(np.pi + (theta - earth_theta)) + return xrot, yrot + + +def sim_lonlatrad(x, y, z, xe, ye, ze, xs, ys, zs): + # convert all to geo coordinates + x = x - xe + y = y - ye + z = z - ze + xs = xs - xe + ys = ys - ye + zs = zs - ze + # convert x y z to lon lat radius + longitude, latitude, radius = cart2sph_deg(x, y, z) + slongitude, slatitude, sradius = cart2sph_deg(xs, ys, zs) + # correct so that Sun is at (0,0) + longitude = deg0to360(slongitude - longitude) + latitude = latitude - slatitude + return longitude, latitude, radius + + +def sun_ra_dec(time_): + out = get_body(Time(time_, format='mjd')) + return out.ra.to('rad').value, out.dec.to('rad').value + + +def ra_dec(r=None, v=None, x=None, y=None, z=None, vx=None, vy=None, vz=None, r_earth=np.array([0, 0, 0]), v_earth=np.array([0, 0, 0]), input_unit='si'): + if r is None or v is None: + if x is not None and y is not None and z is not None and vx is not None and vy is not None and vz is not None: + r = np.array([x, y, z]) + v = np.array([vx, vy, vz]) + else: + raise ValueError("Either provide r and v arrays or individual coordinates (x, y, z) and velocities (vx, vy, vz)") + + # Subtract Earth's position and velocity from the input arrays + r = r - r_earth + v = v - v_earth + + d_earth_mag = einsum_norm(r, 'ij,ij->i') + ra = rad0to2pi(np.arctan2(r[:, 1], r[:, 0])) # in radians + dec = np.arcsin(r[:, 2] / d_earth_mag) + return ra, dec + + +def lonlat_distance(lat1, lat2, lon1, lon2): + # Haversine formula + dlon = lon2 - lon1 + dlat = lat2 - lat1 + a = np.sin(dlat / 2)**2 + np.cos(lat1) * np.cos(lat2) * np.sin(dlon / 2)**2 + c = 2 * np.arcsin(np.sqrt(a)) + # Radius of earth in kilometers. Use 3956 for miles + # calculate the result + return (c * EARTH_RADIUS) + + +def altitude2zenithangle(altitude, deg=True): + if deg: + out = 90 - altitude + else: + out = np.pi / 2 - altitude + return out + + +def zenithangle2altitude(zenith_angle, deg=True): + if deg: + out = 90 - zenith_angle + else: + out = np.pi / 2 - zenith_angle + return out + + +def rightasension2hourangle(right_ascension, local_time): + if type(right_ascension) is not str: + right_ascension = dd_to_hms(right_ascension) + if type(local_time) is not str: + local_time = dd_to_dms(local_time) + _ra = float(right_ascension.split(':')[0]) + _lt = float(local_time.split(':')[0]) + if _ra > _lt: + __ltm, __lts = local_time.split(':')[1:] + local_time = f'{24 + _lt}:{__ltm}:{__lts}' + + return dd_to_dms(hms_to_dd(local_time) - hms_to_dd(right_ascension)) + + +def equatorial_to_horizontal(observer_latitude, declination, right_ascension=None, hour_angle=None, local_time=None, hms=False): + if right_ascension is not None: + hour_angle = rightasension2hourangle(right_ascension, local_time) + if hms: + hour_angle = hms_to_dd(hour_angle) + elif hour_angle is not None: + if hms: + hour_angle = hms_to_dd(hour_angle) + elif right_ascension is not None and hour_angle is not None: + print('Both right_ascension and hour_angle parameters are provided.\nUsing hour_angle for calculations.') + if hms: + hour_angle = hms_to_dd(hour_angle) + else: + print('Either right_ascension or hour_angle must be provided.') + + observer_latitude, hour_angle, declination = np.radians([observer_latitude, hour_angle, declination]) + + zenith_angle = np.arccos(np.sin(observer_latitude) * np.sin(declination) + np.cos(observer_latitude) * np.cos(declination) * np.cos(hour_angle)) + + altitude = zenithangle2altitude(zenith_angle, deg=False) + + _num = np.sin(declination) - np.sin(observer_latitude) * np.cos(zenith_angle) + _den = np.cos(observer_latitude) * np.sin(zenith_angle) + azimuth = np.arccos(_num / _den) + + if observer_latitude < 0: + azimuth = np.pi - azimuth + altitude, azimuth = np.degrees([altitude, azimuth]) + + return azimuth, altitude + + +def horizontal_to_equatorial(observer_latitude, azimuth, altitude): + altitude, azimuth, latitude = np.radians([altitude, azimuth, observer_latitude]) + zenith_angle = zenithangle2altitude(altitude) + + zenith_angle = [-zenith_angle if latitude < 0 else zenith_angle][0] + + declination = np.sin(latitude) * np.cos(zenith_angle) + declination = declination + (np.cos(latitude) * np.sin(zenith_angle) * np.cos(azimuth)) + declination = np.arcsin(declination) + + _num = np.cos(zenith_angle) - np.sin(latitude) * np.sin(declination) + _den = np.cos(latitude) * np.cos(declination) + hour_angle = np.arccos(_num / _den) + + if (latitude > 0 > declination) or (latitude < 0 < declination): + hour_angle = 2 * np.pi - hour_angle + + declination, hour_angle = np.degrees([declination, hour_angle]) + + return hour_angle, declination + + +_ecliptic = 0.409092601 # np.radians(23.43927944) +cos_ec = 0.9174821430960974 +sin_ec = 0.3977769690414367 + + +def equatorial_xyz_to_ecliptic_xyz(xq, yq, zq): + xc = xq + yc = cos_ec * yq + sin_ec * zq + zc = -sin_ec * yq + cos_ec * zq + return xc, yc, zc + + +def ecliptic_xyz_to_equatorial_xyz(xc, yc, zc): + xq = xc + yq = cos_ec * yc - sin_ec * zc + zq = sin_ec * yc + cos_ec * zc + return xq, yq, zq + + +def xyz_to_ecliptic(xc, yc, zc, xe=0, ye=0, ze=0, degrees=False): + x_ast_to_earth = xc - xe + y_ast_to_earth = yc - ye + z_ast_to_earth = zc - ze + d_earth_mag = np.sqrt(np.power(x_ast_to_earth, 2) + np.power(y_ast_to_earth, 2) + np.power(z_ast_to_earth, 2)) + ec_longitude = rad0to2pi(np.arctan2(y_ast_to_earth, x_ast_to_earth)) # in radians + ec_latitude = np.arcsin(z_ast_to_earth / d_earth_mag) + if degrees: + return np.degrees(ec_longitude), np.degrees(ec_latitude) + else: + return ec_longitude, ec_latitude + + +def xyz_to_equatorial(xq, yq, zq, xe=0, ye=0, ze=0, degrees=False): + # RA / DEC calculation - assumes XY plane to be celestial equator, and -x axis to be vernal equinox + x_ast_to_earth = xq - xe + y_ast_to_earth = yq - ye + z_ast_to_earth = zq - ze + d_earth_mag = np.sqrt(np.power(x_ast_to_earth, 2) + np.power(y_ast_to_earth, 2) + np.power(z_ast_to_earth, 2)) + ra = rad0to2pi(np.arctan2(y_ast_to_earth, x_ast_to_earth)) # in radians + dec = np.arcsin(z_ast_to_earth / d_earth_mag) + if degrees: + return np.degrees(ra), np.degrees(dec) + else: + return ra, dec + + +def ecliptic_xyz_to_equatorial(xc, yc, zc, xe=0, ye=0, ze=0, degrees=False): + # Convert ecliptic cartesian into equitorial cartesian + x_ast_to_earth, y_ast_to_earth, z_ast_to_earth = ecliptic_xyz_to_equatorial_xyz(xc - xe, yc - ye, zc - ze) + d_earth_mag = np.sqrt(np.power(x_ast_to_earth, 2) + np.power(y_ast_to_earth, 2) + np.power(z_ast_to_earth, 2)) + ra = rad0to2pi(np.arctan2(y_ast_to_earth, x_ast_to_earth)) # in radians + dec = np.arcsin(z_ast_to_earth / d_earth_mag) + if degrees: + return np.degrees(ra), np.degrees(dec) + else: + return ra, dec + + +def equatorial_to_ecliptic(right_ascension, declination, degrees=False): + ra, dec = np.radians(right_ascension), np.radians(declination) + ec_latitude = np.arcsin(cos_ec * np.sin(dec) - sin_ec * np.cos(dec) * np.sin(ra)) + ec_longitude = np.arctan((cos_ec * np.cos(dec) * np.sin(ra) + sin_ec * np.sin(dec)) / (np.cos(dec) * np.cos(ra))) + if degrees: + return deg0to360(np.degrees(ec_longitude)), np.degrees(ec_latitude) + else: + return rad0to2pi(ec_longitude), ec_latitude + + +def ecliptic_to_equatorial(lon, lat, degrees=False): + lon, lat = np.radians(lon), np.radians(lat) + ra = np.arctan((cos_ec * np.cos(lat) * np.sin(lon) - sin_ec * np.sin(lat)) / (np.cos(lat) * np.cos(lon))) + dec = np.arcsin(cos_ec * np.sin(lat) + sin_ec * np.cos(lat) * np.sin(lon)) + if degrees: + return np.degrees(ra), np.degrees(dec) + else: + return ra, dec + + +def proper_motion_ra_dec(r=None, v=None, x=None, y=None, z=None, vx=None, vy=None, vz=None, r_earth=np.array([0, 0, 0]), v_earth=np.array([0, 0, 0]), input_unit='si'): + if r is None or v is None: + if x is not None and y is not None and z is not None and vx is not None and vy is not None and vz is not None: + r = np.array([x, y, z]) + v = np.array([vx, vy, vz]) + else: + raise ValueError("Either provide r and v arrays or individual coordinates (x, y, z) and velocities (vx, vy, vz)") + + # Subtract Earth's position and velocity from the input arrays + r = r - r_earth + v = v - v_earth + + # Distances to Earth and Sun + d_earth_mag = einsum_norm(r, 'ij,ij->i') + + # RA / DEC calculation + ra = rad0to2pi(np.arctan2(r[:, 1], r[:, 0])) # in radians + dec = np.arcsin(r[:, 2] / d_earth_mag) + ra_unit_vector = np.array([-np.sin(ra), np.cos(ra), np.zeros(np.shape(ra))]).T + dec_unit_vector = -np.array([np.cos(np.pi / 2 - dec) * np.cos(ra), np.cos(np.pi / 2 - dec) * np.sin(ra), -np.sin(np.pi / 2 - dec)]).T + pmra = (np.einsum('ij,ij->i', v, ra_unit_vector)) / d_earth_mag * 206265 # arcseconds / second + pmdec = (np.einsum('ij,ij->i', v, dec_unit_vector)) / d_earth_mag * 206265 # arcseconds / second + + if input_unit == 'si': + return pmra, pmdec + elif input_unit == 'rebound': + pmra = pmra / (31557600 * 2 * np.pi) + pmdec = pmdec / (31557600 * 2 * np.pi) # arcseconds * (au/sim_time)/au, convert to arcseconds / second + return pmra, pmdec + else: + print('Error - units provided not available, provide either SI or rebound units.') + return + + + +def gcrf_to_lunar(r, t, v=None): + class MoonRotator: + def __init__(self): + self.mpm = MoonPosition() + + def __call__(self, r, t): + rmoon = self.mpm(t) + vmoon = (self.mpm(t + 5.0) - self.mpm(t - 5.0)) / 10. + xhat = normed(rmoon.T).T + vpar = np.einsum("ab,ab->b", xhat, vmoon) * xhat + vperp = vmoon - vpar + yhat = normed(vperp.T).T + zhat = np.cross(xhat, yhat, axisa=0, axisb=0).T + R = np.empty((3, 3, len(t))) + R[0] = xhat + R[1] = yhat + R[2] = zhat + return np.einsum("abc,cb->ca", R, r) + rotator = MoonRotator() + if v is None: + return rotator(r, t) + else: + r_lunar = rotator(r, t) + v_lunar = v_from_r(r_lunar, t) + return r_lunar, v_lunar + + +def gcrf_to_lunar_fixed(r, t, v=None): + r_lunar = gcrf_to_lunar(r, t) - gcrf_to_lunar(get_body('moon').position(t).T, t) + if v is None: + return r_lunar + else: + v = v_from_r(r_lunar, t) + return r_lunar, v + + +def gcrf_to_radec(gcrf_coords): + x, y, z = gcrf_coords + # Calculate right ascension in radians + ra = np.arctan2(y, x) + # Convert right ascension to degrees + ra_deg = np.degrees(ra) + # Normalize right ascension to the range [0, 360) + ra_deg = ra_deg % 360 + # Calculate declination in radians + dec_rad = np.arctan2(z, np.sqrt(x**2 + y**2)) + # Convert declination to degrees + dec_deg = np.degrees(dec_rad) + return (ra_deg, dec_deg) + + +def gcrf_to_ecef_bad(r_gcrf, t): + if isinstance(t, Time): + t = t.gps + r_gcrf = np.atleast_2d(r_gcrf) + rotation_angles = WGS84_EARTH_OMEGA * (t - Time("1980-3-20T11:06:00", format='isot').gps) + cos_thetas = np.cos(rotation_angles) + sin_thetas = np.sin(rotation_angles) + + # Create an array of 3x3 rotation matrices + Rz = np.array([[cos_thetas, -sin_thetas, np.zeros_like(cos_thetas)], + [sin_thetas, cos_thetas, np.zeros_like(cos_thetas)], + [np.zeros_like(cos_thetas), np.zeros_like(cos_thetas), np.ones_like(cos_thetas)]]).T + + # Apply the rotation matrices to all rows of r_gcrf simultaneously + r_ecef = np.einsum('ijk,ik->ij', Rz, r_gcrf) + return r_ecef + + +def gcrf_to_lat_lon(r, t): + lon, lat, height = groundTrack(r, t) + return lon, lat, height + + +def gcrf_to_itrf(r_gcrf, t, v=None): + x, y, z = groundTrack(r_gcrf, t, format='cartesian') + _ = np.array([x, y, z]).T + if v is None: + return _ + else: + return _, v_from_r(_, t) + + +def gcrf_to_sim_geo(r_gcrf, t, h=10): + if np.min(np.diff(t.gps)) < h: + h = np.min(np.diff(t.gps)) + r_gcrf = np.atleast_2d(r_gcrf) + r_geo, v_geo = rv(Orbit.fromKeplerianElements(*[RGEO, 0, 0, 0, 0, 0], t=t[0]), t, propagator=RK78Propagator(AccelKepler(), h=h)) + angle_geo_to_x = np.arctan2(r_geo[:, 1], r_geo[:, 0]) + c = np.cos(angle_geo_to_x) + s = np.sin(angle_geo_to_x) + rotation = np.array([[c, -s, np.zeros_like(c)], [s, c, np.zeros_like(c)], [np.zeros_like(c), np.zeros_like(c), np.ones_like(c)]]).T + return np.einsum('ijk,ik->ij', rotation, r_gcrf) + + +# Function still in development, not 100% accurate. +def gcrf_to_itrf_astropy(state_vectors, t): + import astropy.units as u + from astropy.coordinates import GCRS, ITRS, SkyCoord, get_body_barycentric, solar_system_ephemeris, ICRS + + sc = SkyCoord(x=state_vectors[:, 0] * u.m, y=state_vectors[:, 1] * u.m, z=state_vectors[:, 2] * u.m, representation_type='cartesian', frame=GCRS(obstime=t)) + sc_itrs = sc.transform_to(ITRS(obstime=t)) + with solar_system_ephemeris.set('de430'): # other options: builtin, de432s + earth = get_body_barycentric('earth', t) + earth_center_itrs = SkyCoord(earth.x, earth.y, earth.z, representation_type='cartesian', frame=ICRS()).transform_to(ITRS(obstime=t)) + itrs_coords = SkyCoord( + sc_itrs.x.value - earth_center_itrs.x.to_value(u.m), + sc_itrs.y.value - earth_center_itrs.y.to_value(u.m), + sc_itrs.z.value - earth_center_itrs.z.to_value(u.m), + representation_type='cartesian', + frame=ITRS(obstime=t) + ) + # Extract Cartesian coordinates and convert to meters + itrs_coords_meters = np.array([itrs_coords.x, + itrs_coords.y, + itrs_coords.z]).T + return itrs_coords_meters + + +def v_from_r(r, t): + if isinstance(t[0], Time): + t = t.gps + delta_r = np.diff(r, axis=0) + delta_t = np.diff(t) + v = delta_r / delta_t[:, np.newaxis] + v = np.vstack((v, v[-1])) + return v From 22cb94901fae45718142c5e064480cabea155b49 Mon Sep 17 00:00:00 2001 From: yeager7 Date: Thu, 18 Jul 2024 16:10:53 -0700 Subject: [PATCH 02/55] moved most of the coordinate functions to .coordinates, some remain as they will require much more refactoring --- ssapy/utils.py | 88 ++++++++++++++++++++++---------------------------- 1 file changed, 39 insertions(+), 49 deletions(-) diff --git a/ssapy/utils.py b/ssapy/utils.py index 620bac4..17b16aa 100644 --- a/ssapy/utils.py +++ b/ssapy/utils.py @@ -5,7 +5,12 @@ import astropy.units as u from . import datadir -from .constants import MOON_RADIUS, WGS84_EARTH_OMEGA +from .accel import AccelKepler +from .body import get_body, MoonPosition +from .constants import RGEO, MOON_RADIUS, WGS84_EARTH_OMEGA +from .compute import groundTrack, rv +from .orbit import Orbit +from .propagator import RK78Propagator try: import erfa @@ -1157,54 +1162,6 @@ def teme_to_gcrf(t): return erfa.tr(gcrf_to_teme(t)) -def gcrf_to_lunar(r, times): - from .body import MoonPosition - - class MoonRotator: - def __init__(self): - self.mpm = MoonPosition() - - def __call__(self, r, t): - rmoon = self.mpm(t) - vmoon = (self.mpm(t + 5.0) - self.mpm(t - 5.0)) / 10. - xhat = normed(rmoon.T).T - vpar = np.einsum("ab,ab->b", xhat, vmoon) * xhat - vperp = vmoon - vpar - yhat = normed(vperp.T).T - zhat = np.cross(xhat, yhat, axisa=0, axisb=0).T - R = np.empty((3, 3, len(t))) - R[0] = xhat - R[1] = yhat - R[2] = zhat - return np.einsum("abc,cb->ca", R, r) - rotator = MoonRotator() - if isinstance(times, Time): - times = times.gps - return rotator(r, times) - - -def gcrf_to_stationary_lunar(r, times): - from .body import get_body - return gcrf_to_lunar(r, times) - gcrf_to_lunar(get_body('moon').position(times).T, times) - - -def gcrf_to_ecef(r_gcrf, t): - if isinstance(t, Time): - t = t.gps - rotation_angles = WGS84_EARTH_OMEGA * (t - t[0]) - cos_thetas = np.cos(rotation_angles) - sin_thetas = np.sin(rotation_angles) - - # Create an array of 3x3 rotation matrices - Rz = np.array([[cos_thetas, -sin_thetas, np.zeros_like(cos_thetas)], - [sin_thetas, cos_thetas, np.zeros_like(cos_thetas)], - [np.zeros_like(cos_thetas), np.zeros_like(cos_thetas), np.ones_like(cos_thetas)]]).T - - # Apply the rotation matrices to all rows of r_gcrf simultaneously - r_ecef = np.einsum('ijk,ik->ij', Rz, r_gcrf) - return r_ecef - - # Stolen from https://github.com/lsst/utils/blob/main/python/lsst/utils/wrappers.py INTRINSIC_SPECIAL_ATTRIBUTES = frozenset( ( @@ -1485,3 +1442,36 @@ def integrate_orbit_best_model(r=None, v=None, t=None, koe=None, duration=None, except (RuntimeError, ValueError) as err: print(err) return np.nan, np.nan + + +def find_smallest_bounding_cube(r): + """ + Find the smallest bounding cube for a set of 3D coordinates. + + Parameters: + r (np.ndarray): An array of shape (n, 3) containing the 3D coordinates. + + Returns: + tuple: A tuple containing the lower and upper bounds of the bounding cube. + """ + # Find the minimum and maximum coordinates for each axis + min_coords = np.min(r, axis=0) + max_coords = np.max(r, axis=0) + + # Determine the range for each axis + ranges = max_coords - min_coords + + # Calculate the maximum range to ensure the bounding cube + max_range = np.max(ranges) + + # Calculate the center of the bounding cube + center = (max_coords + min_coords) / 2 + + # Calculate the half side length of the cube + half_side_length = max_range / 2 + + # Find the limits of the cube + lower_bound = center - half_side_length + upper_bound = center + half_side_length + + return lower_bound, upper_bound From 932fb95a873c282ffced70aca9bac4c1aeef7c0a Mon Sep 17 00:00:00 2001 From: yeager7 Date: Thu, 18 Jul 2024 16:11:31 -0700 Subject: [PATCH 03/55] added several functions for use in plotUtils --- ssapy/compute.py | 137 +++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 137 insertions(+) diff --git a/ssapy/compute.py b/ssapy/compute.py index d6d6a84..eb56825 100644 --- a/ssapy/compute.py +++ b/ssapy/compute.py @@ -3,6 +3,7 @@ import astropy.units as u from .constants import EARTH_RADIUS, EARTH_MU, MOON_RADIUS +from .coordinates import rotation_matrix_from_vectors from .propagator import KeplerianPropagator from .utils import ( norm, normed, unitAngle3, LRU_Cache, lb_to_unit, sunPos, _gpsToTT, @@ -1336,3 +1337,139 @@ def M_v(r_sat, r_earth, r_sun, r_moon=False, radius=0.4, albedo=0.20, sun_Mag=4. return Mag_v, merged_dict else: return Mag_v + + +def calc_gamma(r, t): + """ + Calculate the gamma angle between position and velocity vectors in the ITRF frame. + + Parameters: + r (numpy.ndarray): The position vectors in the GCRF frame, shaped (n, 3), where n is the number of vectors. + t (numpy.ndarray or astropy.time.Time): The times corresponding to the position vectors. Can be an array of GPS seconds or an Astropy Time object. + + Returns: + numpy.ndarray: An array of gamma angles in degrees between the position and velocity vectors for each time point. + + Notes: + - This function first converts the given position vectors from the GCRF frame to the ITRF frame. + - It then calculates the angle between the position and velocity vectors at each time point in the ITRF frame. + - If the input time array is an Astropy Time object, it converts it to GPS time before processing. + """ + r_itrf, v_itrf = gcrf_to_itrf(r, t, v=True) + if isinstance(t[0], Time): + t = t.gps + gamma = np.degrees(np.apply_along_axis(lambda x: angle_between_vectors(x[:3], x[3:]), 1, np.concatenate((r_itrf, v_itrf), axis=1))) - 90 + return gamma + + +def moon_normal_vector(t): + r = get_body("moon").position(t).T + r_random = get_body("moon").position(t.gps + 604800).T + return np.cross(r, r_random) / np.linalg.norm(r, axis=-1) + + +def lunar_lagrange_points(t): + r = get_body("moon").position(t).T + d = np.linalg.norm(r) # Distance between Earth and Moon + unit_vector_moon = r / np.linalg.norm(r, axis=-1) + # plane_vector = np.cross(r, r_random) + lunar_period_seconds = 2.3605915968e6 + + # Coefficients of the quadratic equation + a = EARTH_MU - MOON_MU + b = 2 * MOON_MU * d + c = -MOON_MU * d**2 + + # Solve the quadratic equation + discriminant = b**2 - 4*a*c + + if discriminant >= 0: + L1_from_moon = (-b - np.sqrt(discriminant)) / (2*a) * unit_vector_moon + L2_from_moon = (-b + np.sqrt(discriminant)) / (2*a) * unit_vector_moon + else: + print("Discriminate is less than 0! THAT'S WEIRD FIX IT.") + L1_from_moon = None + L2_from_moon = None + + return { + "L1": L1_from_moon + r, + "L2": L2_from_moon + r, + "L3": -r, + "L4": get_body("moon").position(t.gps + lunar_period_seconds / 6).T, + "L5": get_body("moon").position(t.gps - lunar_period_seconds / 6).T + } + + +def lunar_lagrange_points_circular(t): + r = get_body("moon").position(t).T + d = np.linalg.norm(r) # Distance between Earth and Moon + unit_vector_moon = r / np.linalg.norm(r, axis=-1) + + # Coefficients of the quadratic equation + a = EARTH_MU - MOON_MU + b = 2 * MOON_MU * d + c = -MOON_MU * d**2 + + # Solve the quadratic equation + discriminant = b**2 - 4*a*c + + if discriminant >= 0: + L1_from_moon = (-b - np.sqrt(discriminant)) / (2*a) * unit_vector_moon + L2_from_moon = (-b + np.sqrt(discriminant)) / (2*a) * unit_vector_moon + else: + print("Discriminate is less than 0! THAT'S WEIRD FIX IT.") + L1_from_moon = None + L2_from_moon = None + + # L45 + # Create the rotation matrix to align z-axis with the normal vector + normal_vector = moon_normal_vector(t) + rotation_matrix = rotation_matrix_from_vectors(np.array([0, 0, 1]), normal_vector) + theta = np.radians(60) + np.arctan2(unit_vector_moon[1], unit_vector_moon[0]) + L4 = np.vstack((d * np.cos(theta), d * np.sin(theta), np.zeros_like(theta))).T + L4 = np.squeeze(L4 @ rotation_matrix.T) + theta = -np.radians(60) + np.arctan2(unit_vector_moon[1], unit_vector_moon[0]) + L5 = np.vstack((d * np.cos(theta), d * np.sin(theta), np.zeros_like(theta))).T + L5 = np.squeeze(L5 @ rotation_matrix.T) + + return { + "L1": L1_from_moon + r, + "L2": L2_from_moon + r, + "L3": -r, + "L4": L4, + "L5": L5 + } + + +def lagrange_points_lunar_frame(): + r = np.array([LD / RGEO, 0, 0]) + d = np.linalg.norm(r) # Distance between Earth and Moon + unit_vector_moon = r / np.linalg.norm(r, axis=-1) + + # Coefficients of the quadratic equation + a = EARTH_MU - MOON_MU + b = 2 * MOON_MU * d + c = -MOON_MU * d**2 + + # Solve the quadratic equation + discriminant = b**2 - 4*a*c + + if discriminant >= 0: + L1_from_moon = (-b - np.sqrt(discriminant)) / (2*a) * unit_vector_moon + L2_from_moon = (-b + np.sqrt(discriminant)) / (2*a) * unit_vector_moon + else: + print("Discriminate is less than 0! THAT'S WEIRD FIX IT.") + L1_from_moon = None + L2_from_moon = None + + # L45 + theta = np.radians(60) + L4 = np.squeeze(np.vstack((d * np.cos(theta), d * np.sin(theta), np.zeros_like(theta))).T) + L5 = np.squeeze(np.vstack((d * np.cos(-theta), d * np.sin(-theta), np.zeros_like(-theta))).T) + return { + "L1": L1_from_moon + r, + "L2": L2_from_moon + r, + "L3": -r, + "L4": L4, + "L5": L5 + } From ee1c03e6f29ebd75a35b67794f645caa33ab2020 Mon Sep 17 00:00:00 2001 From: yeager7 Date: Thu, 18 Jul 2024 16:15:57 -0700 Subject: [PATCH 04/55] improved the brightness functions and added a gamma angle calculation --- ssapy/compute.py | 75 ++++++++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 72 insertions(+), 3 deletions(-) diff --git a/ssapy/compute.py b/ssapy/compute.py index eb56825..6078456 100644 --- a/ssapy/compute.py +++ b/ssapy/compute.py @@ -2,8 +2,9 @@ from astropy.time import Time import astropy.units as u -from .constants import EARTH_RADIUS, EARTH_MU, MOON_RADIUS -from .coordinates import rotation_matrix_from_vectors +from .body import get_body +from .constants import RGEO, LD, EARTH_RADIUS, EARTH_MU, MOON_RADIUS, MOON_MU +from .coordinates import rotation_matrix_from_vectors, angle_between_vectors, gcrf_to_itrf from .propagator import KeplerianPropagator from .utils import ( norm, normed, unitAngle3, LRU_Cache, lb_to_unit, sunPos, _gpsToTT, @@ -1321,7 +1322,7 @@ def sun_shine(r_sat, r_earth, r_sun, radius, albedo, albedo_front, area_panels): return {'sun_bus': flux_bus, 'sun_panels': flux_front} -def M_v(r_sat, r_earth, r_sun, r_moon=False, radius=0.4, albedo=0.20, sun_Mag=4.80, albedo_earth=0.30, albedo_moon=0.12, albedo_back=0.50, albedo_front=0.05, area_panels=100, return_components=False): +def calc_M_v(r_sat, r_earth, r_sun, r_moon=False, radius=0.4, albedo=0.20, sun_Mag=4.80, albedo_earth=0.30, albedo_moon=0.12, albedo_back=0.50, albedo_front=0.05, area_panels=100, return_components=False): r_sun_sat = np.linalg.norm(r_sat - r_sun, axis=-1) frac_flux_sun = {'sun_bus': 0, 'sun_panels': 0} frac_flux_earth = {'earth_bus': 0, 'earth_panels': 0} @@ -1339,6 +1340,73 @@ def M_v(r_sat, r_earth, r_sun, r_moon=False, radius=0.4, albedo=0.20, sun_Mag=4. return Mag_v +def M_v_lambertian(r_sat, times, radius=1.0, albedo=0.20, sun_Mag=4.80, albedo_earth=0.30, albedo_moon=0.12, plot=False): + pc_to_m = 3.085677581491367e+16 + r_sun = get_body('Sun').position(times).T + r_moon = get_body('Moon').position(times).T + r_earth = np.zeros_like(r_sun) + + r_sun_sat = np.linalg.norm(r_sat - r_sun, axis=-1) + r_earth_sat = np.linalg.norm(r_sat, axis=-1) + r_moon_sat = np.linalg.norm(r_sat - r_moon, axis=-1) + + sun_angle = getAngle(r_sun, r_sat, r_earth) + earth_angle = np.pi - sun_angle + moon_phase_angle = getAngle(r_sun, r_moon, r_sat) # Phase of the moon as viewed from the sat. + moon_to_earth_angle = getAngle(r_moon, r_sat, r_earth) + + flux_moon_to_sat = 2 / 3 * albedo_moon * MOON_RADIUS**2 / (np.pi * (r_moon_sat)**2) * (np.sin(moon_phase_angle) + (np.pi - moon_phase_angle) * np.cos(moon_phase_angle)) # Fraction of sunlight reflected from the Moon to satellite + flux_earth_to_sat = 2 / 3 * albedo_earth * EARTH_RADIUS**2 / (np.pi * (r_earth_sat)**2) * (np.sin(earth_angle) + (np.pi - earth_angle) * np.cos(earth_angle)) # Fraction of sunlight reflected from the Earth to satellite + + frac_flux_sun = 2 / 3 * albedo * radius**2 / (np.pi * (r_earth_sat)**2) * (np.sin(sun_angle) + (np.pi - sun_angle) * np.cos(sun_angle)) # Fraction of light reflected off satellite from Sun + frac_flux_earth = 2 / 3 * albedo * radius**2 / (np.pi * r_earth_sat**2) * flux_earth_to_sat + frac_flux_moon = 2 / 3 * albedo * radius**2 / (np.pi * r_earth_sat**2) * flux_moon_to_sat + Mag_v = (2.5 * np.log10((r_sun_sat / (10 * pc_to_m))**2) + sun_Mag) - 2.5 * np.log10(frac_flux_sun + frac_flux_earth + frac_flux_moon) + if plot: + import matplotlib.pyplot as plt + sun_scale = 149597870700.0 * (RGEO / np.max(r_earth_sat) ) * 0.75 + color_map ='inferno_r' + fig = plt.figure(figsize=(18, 4)) + ax = fig.add_subplot(1, 4, 1) + ax.scatter(r_earth[:, 0], r_earth[:, 1], c='Blue', s=10) + scatter = ax.scatter(r_sat[:, 0] / RGEO, r_sat[:, 1] / RGEO, c=sun_angle, cmap=color_map) + colorbar = plt.colorbar(scatter) + ax.scatter(r_sun[:, 0] / sun_scale, r_sun[:, 1] / sun_scale, c=plt.cm.Oranges(np.linspace(0.25, 0.75, len(r_sat[:, 0]))), s=10) + ax.set_title('Solar Phase') + ax.set_xlabel('X [GEO]') + ax.set_ylabel('Y [GEO]') + ax.axis('equal') + + ax = fig.add_subplot(1, 4, 2) + ax.scatter(r_earth[0], r_earth[1], c='Blue', s=10) + scatter = ax.scatter(r_sat[:, 0] / RGEO, r_sat[:, 1] / RGEO, c=(2.5 * np.log10((r_sun_sat / (10 * pc_to_m))**2) + sun_Mag) - 2.5 * np.log10(frac_flux_sun), cmap=color_map) + colorbar = plt.colorbar(scatter) + ax.scatter(r_sun[:, 0] / sun_scale, r_sun[:, 1] / sun_scale, c=plt.cm.Oranges(np.linspace(0.25, 0.75, len(r_sat[:, 0]))), s=10) + ax.set_title('Solar M_v') + ax.axis('equal') + + ax = fig.add_subplot(1, 4, 3) + ax.scatter(r_earth[:, 0], r_earth[:, 1], c='Blue', s=10) + scatter = ax.scatter(r_sat[:, 0] / RGEO, r_sat[:, 1] / RGEO, c=(2.5 * np.log10((r_sun_sat / (10 * pc_to_m))**2) + sun_Mag) - 2.5 * np.log10(frac_flux_earth), cmap=color_map) + colorbar = plt.colorbar(scatter) + ax.scatter(r_sun[:, 0] / sun_scale, r_sun[:, 1] / sun_scale, c=plt.cm.Oranges(np.linspace(0.25, 0.75, len(r_sat[:, 0]))), s=10) + + ax.set_title('Earth M_v') + ax.axis('equal') + + ax = fig.add_subplot(1, 4, 4) + ax.scatter(r_earth[:, 0], r_earth[:, 1], c='Blue', s=10) + scatter = ax.scatter(r_sat[:, 0] / RGEO, r_sat[:, 1] / RGEO, c=(2.5 * np.log10((r_sun_sat / (10 * pc_to_m))**2) + sun_Mag) - 2.5 * np.log10(frac_flux_moon), cmap=color_map) + ax.scatter(r_moon[:, 0] / RGEO, r_moon[:, 1] / RGEO, c=plt.cm.Greys(np.linspace(0.5, 1, len(r_sat[:, 0]))), s=5) + + colorbar = plt.colorbar(scatter) + ax.set_title('Lunar M_v') + ax.axis('equal') + plt.show() + + return Mag_v + + def calc_gamma(r, t): """ Calculate the gamma angle between position and velocity vectors in the ITRF frame. @@ -1362,6 +1430,7 @@ def calc_gamma(r, t): return gamma + def moon_normal_vector(t): r = get_body("moon").position(t).T r_random = get_body("moon").position(t.gps + 604800).T From c1f59d3134008327dbae059800a60a0b978f62b6 Mon Sep 17 00:00:00 2001 From: yeager7 Date: Thu, 18 Jul 2024 16:21:46 -0700 Subject: [PATCH 05/55] updated the calculate_orbital_elements function --- ssapy/compute.py | 23 +- ssapy/plotUtils.py | 814 ++++++++++++++++++++------------------------- 2 files changed, 371 insertions(+), 466 deletions(-) diff --git a/ssapy/compute.py b/ssapy/compute.py index 6078456..33d8990 100644 --- a/ssapy/compute.py +++ b/ssapy/compute.py @@ -1201,7 +1201,7 @@ def bracket(f, x, dx, xmax): } -def _nby3shape(arr_): +def nby3shape(arr_): if arr_.ndim == 1: return np.reshape(arr_, (1, 3)) if arr_.ndim == 2: @@ -1211,12 +1211,12 @@ def _nby3shape(arr_): return arr_.T -def keplerian_orbital_elements(r_, v_, mu_barycenter=EARTH_MU): +def calculate_orbital_elements(r_, v_, mu_barycenter=EARTH_MU): # mu_barycenter - all bodies interior to Earth # 1.0013415732186798 #All bodies of solar system mu_ = mu_barycenter - rarr = _nby3shape(r_) - varr = _nby3shape(v_) + rarr = nby3shape(r_) + varr = nby3shape(v_) aarr = [] earr = [] incarr = [] @@ -1224,6 +1224,7 @@ def keplerian_orbital_elements(r_, v_, mu_barycenter=EARTH_MU): argument_of_periapsisarr = [] longitude_of_ascending_nodearr = [] true_anomalyarr = [] + hmagarr = [] for r, v in zip(rarr, varr): r = np.array(r) # print(f'r: {r}') v = np.array(v) # print(f'v: {v}') @@ -1235,7 +1236,7 @@ def keplerian_orbital_elements(r_, v_, mu_barycenter=EARTH_MU): hmag = np.sqrt(h.dot(h)) n = np.cross(np.array([0, 0, 1]), h) - a = 1 / ((2 / rmag) - (vmag**2) / mu_) + a = 1 / ((2 / rmag) - (vmag ** 2) / mu_) evector = np.cross(v, h) / (mu_) - r / rmag e = np.sqrt(evector.dot(evector)) @@ -1263,7 +1264,17 @@ def keplerian_orbital_elements(r_, v_, mu_barycenter=EARTH_MU): argument_of_periapsisarr.append(argument_of_periapsis) longitude_of_ascending_nodearr.append(longitude_of_ascending_node) true_anomalyarr.append(true_anomaly) - return {'a': aarr, 'e': earr, 'i': incarr, 'lv': true_longitudearr, 'pa': argument_of_periapsisarr, 'raan': longitude_of_ascending_nodearr, 'ta': true_anomalyarr} + hmagarr.append(hmag) + return { + 'a': aarr, + 'e': earr, + 'i': incarr, + 'tl': true_longitudearr, + 'ap': argument_of_periapsisarr, + 'raan': longitude_of_ascending_nodearr, + 'ta': true_anomalyarr, + 'L': hmagarr + } ###################################################################################### diff --git a/ssapy/plotUtils.py b/ssapy/plotUtils.py index 2030440..615823b 100644 --- a/ssapy/plotUtils.py +++ b/ssapy/plotUtils.py @@ -1,9 +1,13 @@ from .body import get_body -from .compute import groundTrack -from .constants import RGEO, EARTH_RADIUS -from .utils import find_file, Time, norm, gcrf_to_ecef, gcrf_to_lunar, gcrf_to_stationary_lunar +from .compute import groundTrack, lagrange_points_lunar_frame, lagrange_points_lunar_frame_ +from .constants import RGEO, LD, EARTH_RADIUS, MOON_RADIUS +from .coordinates import gcrf_to_itrf, gcrf_to_lunar_fixed, gcrf_to_lunar +from .utils import find_file, Time, find_smallest_bounding_cube import numpy as np +import os +import re +import io import matplotlib.pyplot as plt import matplotlib.cm as cm @@ -77,18 +81,69 @@ def load_moon_file(): return moon -def groundTrackPlot(r, time, ground_stations=None): +def drawMoon(time, ngrid=100, R=MOON_RADIUS, rfactor=1): + """ + Parameters + ---------- + time : array_like or astropy.time.Time (n,) + If float (array), then should correspond to GPS seconds; + i.e., seconds since 1980-01-06 00:00:00 UTC + ngrid: int + Number of grid points in Earth model. + R: float + Earth radius in meters. Default is WGS84 value. + rfactor: float + Factor by which to enlarge Earth (for visualization purposes) + + """ + moon = load_moon_file() + + from numbers import Real + from erfa import gst94 + lat = np.linspace(-np.pi / 2, np.pi / 2, ngrid) + lon = np.linspace(-np.pi, np.pi, ngrid) + lat, lon = np.meshgrid(lat, lon) + x = np.cos(lat) * np.cos(lon) + y = np.cos(lat) * np.sin(lon) + z = np.sin(lat) + u = np.linspace(0, 1, ngrid) + v, u = np.meshgrid(u, u) + + # Need earth rotation angle for t + # Just use erfa.gst94. + # This ignores precession/nutation, ut1-tt and polar motion, but should + # be good enough for visualization. + if isinstance(time, Time): + time = time.gps + if isinstance(time, Real): + time = np.array([time]) + + mjd_tt = 44244.0 + (time + 51.184) / 86400 + gst = gst94(2400000.5, mjd_tt) + + u = u - (gst / (2 * np.pi))[:, None, None] + v = np.broadcast_to(v, u.shape) + + return ipv.plot_mesh( + x * R * rfactor, y * R * rfactor, z * R * rfactor, + u=u, v=v, + wireframe=False, + texture=moon + ) + + +def groundTrackPlot(r, time, ground_stations=None, save_path=False): """ Parameters ---------- r : (3,) array_like - Orbit positions in meters. - times: (n,) array_like - array of Astropy Time objects or time in gps seconds. + t: (n,) array_like - array of Astropy Time objects or time in gps seconds. optional - ground_stations: (n,2) array of of ground station (lat,lon) in degrees """ lon, lat, height = groundTrack(r, time) - plt.figure(figsize=(15, 12)) + fig = plt.figure(figsize=(15, 12)) plt.imshow(load_earth_file(), extent=[-180, 180, -90, 90]) plt.plot(np.rad2deg(lon), np.rad2deg(lat)) if ground_stations is not None: @@ -97,6 +152,8 @@ def groundTrackPlot(r, time, ground_stations=None): plt.ylim(-90, 90) plt.xlim(-180, 180) plt.show() + if save_path: + save_plot(fig, save_path) def groundTrackVideo(r, time): @@ -151,474 +208,311 @@ def check_numpy_array(variable): return "not numpy" -def _make_scatter(x, y, z, xm=False, ym=False, zm=False, limits=False, title='', figsize=(7, 7), orbit_index='', save_path=False, plot_type=False): - # DETERMINE WHAT TYPE OF SCATTER THIS IS - if limits is False: - limits = np.nanmax(np.abs([x, y, z])) * 1.2 - - dotcolors = cm.rainbow(np.linspace(0, 1, len(x))) - if np.size(xm) > 1: - gradient_colors = cm.Greys(np.linspace(0.5, 1, len(xm))) - else: - gradient_colors = "grey" - - # Central dot color, central dot size, dashed line radius - plot_settings = { - "gcrf": ("blue", 50, 1), - "ecef": ("blue", 50, 1), - "lunar": ("blue", 50, 1), - "stationary_lunar": ("grey", 25, 1.3) - } - - # Check if the plot_type is in the dictionary, and set central_dot accordingly - if plot_type in plot_settings: - stn = plot_settings[plot_type] - else: - raise ValueError("Unknown plot type provided. Accepted: gcrf, ecef, lunar, stationary_lunar") - - # Creating plot - plt.rcParams.update({'font.size': 9, 'figure.facecolor': 'w'}) - fig = plt.figure(dpi=100, figsize=figsize) - - ax1 = fig.add_subplot(2, 2, 1, projection='3d') - ax1.scatter3D(x, y, z, color=dotcolors, s=1) - ax1.scatter3D(0, 0, 0, color=stn[0], s=stn[1]) - if xm is not False: - ax1.scatter3D(xm, ym, zm, color=gradient_colors, s=5) - ax1.view_init(elev=30, azim=60) - ax1.set_xlim([-limits, limits]) - ax1.set_ylim([-limits, limits]) - ax1.set_zlim([-limits, limits]) - ax1.set_aspect('equal') # aspect ratio is 1:1:1 in data space - ax1.set_xlabel('x [GEO]') - ax1.set_ylabel('y [GEO]') - ax1.set_zlabel('z [GEO]') - - ax2 = fig.add_subplot(2, 2, 2) - ax2.add_patch(plt.Circle((0, 0), stn[2], color='black', linestyle='dashed', fill=False)) - ax2.scatter(x, y, color=dotcolors, s=1) - ax2.scatter(0, 0, color=stn[0], s=stn[1]) - if xm is not False: - ax2.scatter(xm, ym, color=gradient_colors, s=5) - ax2.set_aspect('equal') - ax2.set_xlabel('x [GEO]') - ax2.set_ylabel('y [GEO]') - ax2.set_xlim((-limits, limits)) - ax2.set_ylim((-limits, limits)) - ax2.text(x[0], y[0], f'← start {orbit_index}') - ax2.text(x[-1], y[-1], f'← end {orbit_index}') - ax2.set_title(f'{title}') - - ax3 = fig.add_subplot(2, 2, 3) - ax3.add_patch(plt.Circle((0, 0), stn[2], color='black', linestyle='dashed', fill=False)) - ax3.scatter(x, z, color=dotcolors, s=1) - ax3.scatter(0, 0, color=stn[0], s=stn[1]) - if xm is not False: - ax3.scatter(xm, zm, color=gradient_colors, s=5) - ax3.set_aspect('equal') - ax3.set_xlabel('x [GEO]') - ax3.set_ylabel('z [GEO]') - ax3.set_xlim((-limits, limits)) - ax3.set_ylim((-limits, limits)) - ax3.text(x[0], z[0], f'← start {orbit_index}') - ax3.text(x[-1], z[-1], f'← end {orbit_index}') - - ax4 = fig.add_subplot(2, 2, 4) - ax4.add_patch(plt.Circle((0, 0), stn[2], color='black', linestyle='dashed', fill=False)) - ax4.scatter(y, z, color=dotcolors, s=1) - ax4.scatter(0, 0, color=stn[0], s=stn[1]) - if xm is not False: - ax4.scatter(ym, zm, color=gradient_colors, s=5) - ax4.set_aspect('equal') - ax4.set_xlabel('y [GEO]') - ax4.set_ylabel('z [GEO]') - ax4.set_xlim((-limits, limits)) - ax4.set_ylim((-limits, limits)) - ax4.text(y[0], z[0], f'← start {orbit_index}') - ax4.text(y[-1], z[-1], f'← end {orbit_index}') - - plt.tight_layout(pad=2.0, h_pad=2.0) - if save_path: - if save_path.lower().endswith('.png'): - save_plot_to_png(fig, save_path) - else: - save_plot_to_pdf(fig, save_path) - return fig - - -def gcrf_plot(r, times=[], limits=False, title='', save_path=False, figsize=(7, 7)): - """ - Parameters - ---------- - r : (n,3) or array of [(n,3), ..., (n,3)] array_like - Position of orbiting object(s) in meters. - times: optional - times when r was calculated. - limits: optional - x and y limits of the plot - title: optional - title of the plot - """ - - if np.size(times) < 1: - r_moon = np.atleast_2d(get_body("moon").position(Time("2000-1-1"))) - else: - r_moon = get_body("moon").position(times).T - xm = r_moon[:, 0] / RGEO - ym = r_moon[:, 1] / RGEO - zm = r_moon[:, 2] / RGEO - input_type = check_numpy_array(r) - if input_type == "numpy array": - x = r[:, 0] / RGEO - y = r[:, 1] / RGEO - z = r[:, 2] / RGEO - fig = _make_scatter(x=x, y=y, z=z, xm=xm, ym=ym, zm=zm, limits=limits, title=title, figsize=figsize, save_path=save_path, plot_type="gcrf") - if input_type == "list of numpy array": - limits_plot = 0 - for i, row in enumerate(r): - if limits is False and limits_plot < np.nanmax(norm(row) / RGEO) * 1.2: - limits_plot = np.nanmax(norm(row) / RGEO) * 1.2 - else: - limits_plot = limits - x = row[:, 0] / RGEO - y = row[:, 1] / RGEO - z = row[:, 2] / RGEO - fig = _make_scatter(x=x, y=y, z=z, xm=xm, ym=ym, zm=zm, limits=limits_plot, title=title, orbit_index=i, figsize=figsize, save_path=save_path, plot_type="gcrf") - return fig - - -def ecef_plot(r, times, limits=False, title='', save_path=False, figsize=(7, 7)): +def orbit_plot(r, t=[], limits=False, title='', figsize=(7, 7), save_path=False, frame="gcrf", show=False): """ Parameters ---------- r : (n,3) or array of [(n,3), ..., (n,3)] array_like Position of orbiting object(s) in meters. - times: times when r was calculated. + t: optional - t when r was calculated. limits: optional - x and y limits of the plot title: optional - title of the plot """ - input_type = check_numpy_array(r) - r = gcrf_to_ecef(r, times) - if input_type == "numpy array": - x = r[:, 0] / RGEO - y = r[:, 1] / RGEO - z = r[:, 2] / RGEO - fig = _make_scatter(x=x, y=y, z=z, limits=limits, title=title, save_path=save_path, plot_type="ecef") - if input_type == "list of numpy array": - limits_plot = 0 - for i, row in enumerate(r): - if limits is False and limits_plot < np.nanmax(norm(row) / RGEO) * 1.2: - limits_plot = np.nanmax(norm(row) / RGEO) * 1.2 - else: - limits_plot = limits - x = row[:, 0] / RGEO - y = row[:, 1] / RGEO - z = row[:, 2] / RGEO - fig = _make_scatter(x=x, y=y, z=z, limits_plot=limits_plot, title=title, orbit_index=i, save_path=save_path, plot_type="ecef") - return fig - - -def lunar_plot(r, times, limits=False, title='', save_path=False, figsize=(7, 7)): - """ - Parameters - ---------- - r : (n,3) or array of [(n,3), ..., (n,3)] array_like - Position of orbiting object(s) in meters. - times: array of Astropy time objects - limits: optional - x and y limits of the plot - title: optional - title of the plot - """ - - input_type = check_numpy_array(r) - r = gcrf_to_lunar(r, times) - r_moon = gcrf_to_lunar(get_body("moon").position(times).T, times) - xm = r_moon[:, 0] / RGEO - ym = r_moon[:, 1] / RGEO - zm = r_moon[:, 2] / RGEO - if input_type == "numpy array": - x = r[:, 0] / RGEO - y = r[:, 1] / RGEO - z = r[:, 2] / RGEO - fig = _make_scatter(x, y, z, xm, ym, zm, limits=limits, title=title, save_path=save_path, plot_type="lunar") - if input_type == "list of numpy array": - limits_plot = 0 - for i, row in enumerate(r): - if limits is False and limits_plot < np.nanmax(norm(row) / RGEO) * 1.2: - limits_plot = np.nanmax(norm(row) / RGEO) * 1.2 - else: - limits_plot = limits - x = row[:, 0] / RGEO - y = row[:, 1] / RGEO - z = row[:, 2] / RGEO - fig = _make_scatter(x, y, z, xm, ym, zm, limits=limits_plot, title=title, orbit_index=i, save_path=save_path, plot_type="lunar") - return fig - - -def lunar_stationary_plot(r, times, limits=False, title='', save_path=False, figsize=(7, 7)): - """ - Parameters - ---------- - r : (n,3) or array of [(n,3), ..., (n,3)] array_like - Position of orbiting object(s) in meters. - times: array of Astropy time objects - limits: optional - x and y limits of the plot - title: optional - title of the plot - """ - input_type = check_numpy_array(r) - r = gcrf_to_stationary_lunar(r, times) - if input_type == "numpy array": - x = r[:, 0] / RGEO - y = r[:, 1] / RGEO - z = r[:, 2] / RGEO - fig = _make_scatter(x, y, z, limits=limits, title=title, save_path=save_path, plot_type="stationary_lunar") - if input_type == "list of numpy array": - limits_plot = 0 - for i, row in enumerate(r): - if limits is False and limits_plot < np.nanmax(norm(row) / RGEO) * 1.2: - limits_plot = np.nanmax(norm(row) / RGEO) * 1.2 - else: - limits_plot = limits - x = row[:, 0] / RGEO - y = row[:, 1] / RGEO - z = row[:, 2] / RGEO - fig = _make_scatter(x, y, z, limts=limits_plot, title=title, orbit_index=i, save_path=save_path, plot_type="stationary_lunar") - return fig - - -def tracking_plot(r, times, ground_stations=None, limits=False, title='', figsize=(7, 8), save_path=False, elev=30, azim=90, scale=5): - """ - Create a 3D tracking plot of satellite positions over time on Earth's surface. - - Parameters - ---------- - r : numpy.ndarray or list of numpy.ndarray - Satellite positions in GCRF coordinates. If a single numpy array, it represents the satellite's position vector over time. If a list of numpy arrays, it represents multiple satellite position vectors. - - times : numpy.ndarray - Timestamps corresponding to the satellite positions. - - ground_stations : list of tuples, optional - List of ground stations represented as (latitude, longitude) pairs. Default is None. - - limits : float or bool, optional - The plot limits for x, y, and z axes. If a float, it sets the limits for all axes. If False, the limits are automatically determined based on the data. Default is False. - - title : str, optional - Title for the plot. Default is an empty string. - - figsize : tuple, optional - Figure size in inches (width, height). Default is (7, 8). - - save_path : str or bool, optional - Path to save the plot as an image or PDF. If False, the plot is not saved. Default is False. - - elev : float, optional - Elevation angle for the 3D plot view. Default is 30 degrees. - - azim : float, optional - Azimuthal angle for the 3D plot view. Default is 90 degrees. - - scale : int, optional - Scaling factor for the Earth's image. Default is 5. - - frame : str, optional - Coordinate frame for the satellite positions, "gcrf" or "ecef". Default is "gcrf". - - Returns - ------- - matplotlib.figure.Figure - The created tracking plot figure. - - Notes - ----- - - The function supports plotting the positions of one or multiple satellites over time. - - Ground station locations can be optionally displayed on the plot. - - The limits parameter can be set to specify the plot's axis limits or automatically determined if set to False. - - The frame parameter determines the coordinate frame for the satellite positions, "gcrf" (default) or "ecef". - - Example Usage - ------------- - - Single satellite tracking plot: - tracking_plot(r_satellite, times, ground_stations=[(40, -75)], title="Satellite Tracking") - - - Multiple satellite tracking plot: - tracking_plot([r_satellite_1, r_satellite_2], times, title="Multiple Satellite Tracking") - - - Save the plot as a PNG image: - tracking_plot(r_satellite, times, save_path="satellite_tracking.png") - - - Customize the plot view: - tracking_plot(r_satellite, times, elev=45, azim=120) + def _make_scatter(fig, ax1, ax2, ax3, ax4, r, t, limits, title='', orbit_index='', num_orbits=1, frame=False): + if np.size(t) < 1: + if frame in ["itrf", "lunar", "lunar_fixed"]: + raise("Need to provide t for itrf, lunar or lunar fixed frames") + r_moon = np.atleast_2d(get_body("moon").position(Time("2000-1-1"))) + else: + r_moon = get_body("moon").position(t).T + + # Dictionary of frame transformations and titles + def get_main_category(frame): + variant_mapping = { + "gcrf": "gcrf", + "gcrs": "gcrf", + "itrf": "itrf", + "itrs": "itrf", + "lunar": "lunar", + "lunar_fixed": "lunar", + "lunar fixed": "lunar", + "lunar_centered": "lunar", + "lunar centered": "lunar", + "lunarearthfixed": "lunar axis", + "lunarearth": "lunar axis", + "lunar axis": "lunar axis", + "lunar_axis": "lunar axis", + "lunaraxis": "lunar axis", + } + return variant_mapping.get(frame.lower()) + + frame_transformations = { + "gcrf": ("GCRF", None), + "itrf": ("ITRF", gcrf_to_itrf), + "lunar": ("Lunar Frame", gcrf_to_lunar_fixed), + "lunar axis": ("Moon on x-axis Frame", gcrf_to_lunar), + } + + # Check if the frame is in the dictionary, and set central_dot accordingly + frame = get_main_category(frame) + if frame in frame_transformations: + title2, transform_func = frame_transformations[frame] + if transform_func: + r = transform_func(r, t) + r_moon = transform_func(r_moon, t) + else: + raise ValueError("Unknown plot type provided. Accepted: gcrf, itrf, lunar, lunar fixed") - - Set custom axis limits: - tracking_plot(r_satellite, times, limits=500) - """ - def _make_plot(r, times, ground_stations, limits, title, figsize, save_path, elev, azim, scale, orbit_index=''): - lon, lat, height = groundTrack(r, times) - r = gcrf_to_ecef(r, times) x = r[:, 0] / RGEO y = r[:, 1] / RGEO z = r[:, 2] / RGEO + xm = r_moon[:, 0] / RGEO + ym = r_moon[:, 1] / RGEO + zm = r_moon[:, 2] / RGEO + + if np.size(xm) > 1: + gradient_colors = cm.Greys(np.linspace(0, .8, len(xm)))[::-1] + blues = cm.Blues(np.linspace(.4, .9, len(xm)))[::-1] + else: + gradient_colors = "grey" + blues = 'Blue' + plot_settings = { + "gcrf": ("blue", 50, 1, xm, ym, zm, gradient_colors), + "itrf": ("blue", 50, 1, xm, ym, zm, gradient_colors), + "lunar": ("grey", 25, 1.3, xm, ym, zm, blues), + "lunar axis": ("blue", 50, 1, -xm, -ym, -zm, gradient_colors) + } + try: + stn = plot_settings[frame] + except KeyError: + raise ValueError("Unknown plot type provided. Accepted: 'gcrf', 'itrf', 'lunar', 'lunar fixed'") if limits is False: - limits = np.nanmax(np.abs([x, y, z])) * 1.2 - - dotcolors = cm.rainbow(np.linspace(0, 1, len(x))) - - # Creating plot - plt.rcParams.update({'font.size': 9, 'figure.facecolor': 'w'}) - fig = plt.figure(dpi=100, figsize=figsize) - earth_png = PILImage.open(find_file("earth", ext=".png")) - earth_png = earth_png.resize((5400 // scale, 2700 // scale)) - ax_gt = fig.add_subplot(2, 2, (3, 4)) - ax_gt.imshow(earth_png, extent=[-180, 180, -90, 90]) - ax_gt.plot(np.rad2deg(lon), np.rad2deg(lat)) - if ground_stations is not None: - for ground_station in ground_stations: - ax_gt.scatter(ground_station[1], ground_station[0], s=50, color='Red') - ax_gt.set_ylim(-90, 90) - ax_gt.set_xlim(-180, 180) - ax_gt.set_xlabel('longitude [degrees]') - ax_gt.set_ylabel('latitude [degrees]') - - bm = np.array(earth_png.resize([int(d) for d in earth_png.size])) / 256. - lons = np.linspace(-180, 180, bm.shape[1]) * np.pi / 180 - lats = np.linspace(-90, 90, bm.shape[0])[::-1] * np.pi / 180 - mesh_x = np.outer(np.cos(lons), np.cos(lats)).T * 0.15126911409197252 - mesh_y = np.outer(np.sin(lons), np.cos(lats)).T * 0.15126911409197252 - mesh_z = np.outer(np.ones(np.size(lons)), np.sin(lats)).T * 0.15126911409197252 - - ax_3d = fig.add_subplot(2, 2, 1, projection='3d') - ax_3d.scatter3D(x, y, z, color=dotcolors, s=1) - ax_3d.plot_surface(mesh_x, mesh_y, mesh_z, rstride=4, cstride=4, facecolors=bm, shade=False) - ax_3d.view_init(elev=elev, azim=azim) - ax_3d.set_xlim([-limits, limits]) - ax_3d.set_ylim([-limits, limits]) - ax_3d.set_zlim([-limits, limits]) - ax_3d.set_aspect('equal') # aspect ratio is 1:1:1 in data space - ax_3d.set_xlabel('x [GEO]') - ax_3d.set_ylabel('y [GEO]') - ax_3d.set_zlabel('z [GEO]') - - ax_3d_r = fig.add_subplot(2, 2, 2, projection='3d') - ax_3d_r.scatter3D(x, y, z, color=dotcolors, s=1) - ax_3d_r.plot_surface(mesh_x, mesh_y, mesh_z, rstride=4, cstride=4, facecolors=bm, shade=False) - ax_3d_r.view_init(elev=elev, azim=180 + azim) - ax_3d_r.set_xlim([-limits, limits]) - ax_3d_r.set_ylim([-limits, limits]) - ax_3d_r.set_zlim([-limits, limits]) - ax_3d_r.set_aspect('equal') # aspect ratio is 1:1:1 in data space - ax_3d_r.set_xlabel('x [GEO]') - ax_3d_r.set_ylabel('y [GEO]') - ax_3d_r.set_zlabel('z [GEO]') - - plt.tight_layout() - if save_path: - if save_path.lower().endswith('.png'): - save_plot_to_png(fig, save_path) - else: - save_plot_to_pdf(fig, save_path) - return fig + lower_bound, upper_bound = find_smallest_bounding_cube(r / RGEO) + lower_bound = lower_bound * 1.2 + upper_bound = upper_bound * 1.2 + if orbit_index == '': + angle = 0 + dotcolors = cm.rainbow(np.linspace(0, 1, len(x))) + else: + angle = orbit_index * 10 + dotcolors = cm.rainbow(np.linspace(0, 1, num_orbits))[orbit_index] + ax1.add_patch(plt.Circle((0, 0), stn[2], color='white', linestyle='dashed', fill=False)) + ax1.scatter(x, y, color=dotcolors, s=1) + ax1.scatter(0, 0, color=stn[0], s=stn[1]) + if xm is not False: + ax1.scatter(stn[3], stn[4], color=stn[6], s=5) + ax1.set_aspect('equal') + ax1.set_xlabel('x [GEO]') + ax1.set_ylabel('y [GEO]') + ax1.set_xlim((lower_bound[0], upper_bound[0])) + ax1.set_ylim((lower_bound[1], upper_bound[1])) + ax1.set_title(f'Frame: {title2}', color='white') + if 'lunar' in frame: + colors = ['red', 'green', 'purple', 'orange', 'cyan'] + for (point, pos), color in zip(lagrange_points_lunar_frame().items(), colors): + if 'axis' in frame: + pass + else: + pos[0] = pos[0] - LD / RGEO + ax1.scatter(pos[0], pos[1], color=color, label=point) + ax1.text(pos[0], pos[1], point, color=color) + + ax2.add_patch(plt.Circle((0, 0), stn[2], color='white', linestyle='dashed', fill=False)) + ax2.scatter(x, z, color=dotcolors, s=1) + ax2.scatter(0, 0, color=stn[0], s=stn[1]) + if xm is not False: + ax2.scatter(stn[3], stn[5], color=stn[6], s=5) + ax2.set_aspect('equal') + ax2.set_xlabel('x [GEO]') + ax2.set_ylabel('z [GEO]') + ax2.set_xlim((lower_bound[0], upper_bound[0])) + ax2.set_ylim((lower_bound[2], upper_bound[2])) + ax2.set_title(f'{title}', color='white') + if 'lunar' in frame: + colors = ['red', 'green', 'purple', 'orange', 'cyan'] + for (point, pos), color in zip(lagrange_points_lunar_frame().items(), colors): + if 'axis' in frame: + pass + else: + pos[0] = pos[0] - LD / RGEO + ax2.scatter(pos[0], pos[2], color=color, label=point) + ax2.text(pos[0], pos[2], point, color=color) + + ax3.add_patch(plt.Circle((0, 0), stn[2], color='white', linestyle='dashed', fill=False)) + ax3.scatter(y, z, color=dotcolors, s=1) + ax3.scatter(0, 0, color=stn[0], s=stn[1]) + if xm is not False: + ax3.scatter(stn[4], stn[5], color=stn[6], s=5) + ax3.set_aspect('equal') + ax3.set_xlabel('y [GEO]') + ax3.set_ylabel('z [GEO]') + ax3.set_xlim((lower_bound[1], upper_bound[1])) + ax3.set_ylim((lower_bound[2], upper_bound[2])) + if 'lunar' in frame: + colors = ['red', 'green', 'purple', 'orange', 'cyan'] + for (point, pos), color in zip(lagrange_points_lunar_frame().items(), colors): + if 'axis' in frame: + pass + else: + pos[0] = pos[0] - LD / RGEO + print(pos) + ax3.scatter(pos[1], pos[2], color=color, label=point) + ax3.text(pos[1], pos[2], point, color=color) + + ax4.scatter3D(x, y, z, color=dotcolors, s=1) + ax4.scatter3D(0, 0, 0, color=stn[0], s=stn[1]) + if xm is not False: + ax4.scatter3D(stn[3], stn[4], stn[5], color=stn[6], s=5) + ax4.set_xlim([lower_bound[0], upper_bound[0]]) + ax4.set_ylim([lower_bound[1], upper_bound[1]]) + ax4.set_zlim([lower_bound[2], upper_bound[2]]) + ax4.set_aspect('equal') # aspect ratio is 1:1:1 in data space + ax4.set_xlabel('x [GEO]') + ax4.set_ylabel('y [GEO]') + ax4.set_zlabel('z [GEO]') + if 'lunar' in frame: + colors = ['red', 'green', 'purple', 'orange', 'cyan'] + for (point, pos), color in zip(lagrange_points_lunar_frame().items(), colors): + if 'axis' in frame: + pass + else: + pos[0] = pos[0] - LD / RGEO + ax4.scatter(pos[0], pos[1], pos[2], color=color, label=point) + ax4.text(pos[0], pos[1], pos[2], point, color=color) + + return fig, ax1, ax2, ax3, ax4 input_type = check_numpy_array(r) + + fig = plt.figure(dpi=100, figsize=figsize, facecolor='black') + ax1 = fig.add_subplot(2, 2, 1) + ax2 = fig.add_subplot(2, 2, 2) + ax3 = fig.add_subplot(2, 2, 3) + ax4 = fig.add_subplot(2, 2, 4, projection='3d') + if input_type == "numpy array": - fig = _make_plot( - r, times, ground_stations=ground_stations, - limits=limits, title=title, figsize=figsize, - save_path=save_path, elev=elev, azim=azim, - scale=scale) + fig, ax1, ax2, ax3, ax4 = _make_scatter(fig, ax1, ax2, ax3, ax4, r=r, t=t, limits=limits, title=title, frame=frame) if input_type == "list of numpy array": - limits_plot = 0 + num_orbits = np.shape(r)[0] for i, row in enumerate(r): - if limits is False and limits_plot < np.nanmax(norm(row) / RGEO) * 1.2: - limits_plot = np.nanmax(norm(row) / RGEO) * 1.2 - else: - limits_plot = limits - fig = _make_plot( - row, times, ground_stations=ground_stations, - limits=limits_plot, title=title, figsize=figsize, - elev=elev, azim=azim, save_path=save_path, - scale=scale, orbit_index=i - ) - return fig - - -def plotRVTrace(chain, lnprob, lnprior, start=0, fig=None): - """Plot rv MCMC traces. - - Parameters - ---------- - chain : array_like (nStep, nChain, 6) - MCMC sample chain. - lnprob : array_like (nStep, nChain) - Log posterior probability of sample chain. - lnprior : array_like (nStep, nChain) - Log posterior probability of sample chain. - start : int, optional - Optional starting index for plot - fig : matplotlib Figure, optional - Where to make plots. Will be created if not explicitly supplied. - """ - import matplotlib.pyplot as plt - if fig is None: - fig, axes = plt.subplots(nrows=4, ncols=2, figsize=(8.5, 11)) - axes = axes.ravel() - else: - try: - axes = np.array(fig.axes).reshape((4, 2)).ravel() - except Exception: - raise ValueError("Provided figure has {} axes, but plot requires " - "dimensions K={}".format(np.array(fig.axes).shape, (4, 2))) - - labels = ["x (km)", "y (km)", "z (km)", - "vx (km/s)", "vy (km/s)", "vz (km/s)"] - xs = np.arange(start, chain.shape[0]) - for i, (ax, label) in enumerate(zip(axes, labels)): - ax.plot(xs, chain[start:, :, i] / 1000, alpha=0.1, color='k') - ax.set_ylabel(label) - axes[-2].plot(xs, lnprob[start:], alpha=0.1, color='k') - axes[-2].set_ylabel("lnprob") - axes[-1].plot(xs, lnprior[start:], alpha=0.1, color='k') - axes[-1].set_ylabel("lnprior") - return fig - - -def plotRVCorner(chain, lnprob, lnprior, start=0, **kwargs): - """Make rv samples corner plot. - - Parameters - ---------- - chain : array_like (nStep, nChain, 6) - MCMC sample chain. - lnprob : array_like (nStep, nChain) - Log posterior probability of sample chain. - lnprior : array_like (nStep, nChain) - Log prior probability of sample chain. - start : int, optional - Optional starting index for selecting samples - kwargs : optional kwargs to pass on to corner.corner - - Returns - ------- - fig : matplotlib.Figure - Figure instance from corner.corner. - """ - import corner - - chain = chain[start:] / 1000 - lnprob = np.atleast_3d(lnprob[start:]) - lnprior = np.atleast_3d(lnprior[start:]) - data = np.concatenate([chain, lnprob, lnprior], axis=2) - data = data.reshape((-1, 8)) - labels = ["x (km)", "y (km)", "z (km)", - "vx (km/s)", "vy (km/s)", "vz (km/s)", - "lnprob", "lnprior"] - - return corner.corner(data, labels=labels, **kwargs) + fig, ax1, ax2, ax3, ax4 = _make_scatter(fig, ax1, ax2, ax3, ax4, r=row, t=t, limits=limits, title=title, orbit_index=i, num_orbits=num_orbits, frame=frame) + + # Set axis color to white + for i, ax in enumerate([ax1, ax2, ax3, ax4]): + ax.set_facecolor('black') + ax.spines['bottom'].set_color('white') + ax.spines['top'].set_color('white') + ax.spines['left'].set_color('white') + ax.spines['right'].set_color('white') + ax.xaxis.label.set_color('white') + ax.yaxis.label.set_color('white') + ax.tick_params(axis='x', colors='white') + ax.tick_params(axis='y', colors='white') + if i == 3: + ax.tick_params(axis='z', colors='white') + + # Set text color to white + for ax in [ax1, ax2, ax3, ax4]: + for text in ax.get_xticklabels() + ax.get_yticklabels() + [ax.xaxis.label, ax.yaxis.label]: + text.set_color('white') + + #Save the plot + fig.patch.set_facecolor('black') + if show: + plt.show() + if save_path: + save_plot(fig, save_path) + return fig, [ax1, ax2, ax3, ax4] -###################################################################### +def globe_plot(r, t, limits=False, title='', figsize=(7, 8), save_path=False, el=30, az=0, scale=1): + x = r[:, 0] / RGEO + y = r[:, 1] / RGEO + z = r[:, 2] / RGEO + if limits is False: + limits = np.nanmax(np.abs([x, y, z])) * 1.2 + + earth_png = PILImage.open(find_file("earth", ext=".png")) + earth_png = earth_png.resize((5400 // scale, 2700 // scale)) + bm = np.array(earth_png.resize([int(d) for d in earth_png.size])) / 256. + lons = np.linspace(-180, 180, bm.shape[1]) * np.pi / 180 + lats = np.linspace(-90, 90, bm.shape[0])[::-1] * np.pi / 180 + mesh_x = np.outer(np.cos(lons), np.cos(lats)).T * EARTH_RADIUS / RGEO + mesh_y = np.outer(np.sin(lons), np.cos(lats)).T * EARTH_RADIUS / RGEO + mesh_z = np.outer(np.ones(np.size(lons)), np.sin(lats)).T * EARTH_RADIUS / RGEO + + dotcolors = plt.cm.rainbow(np.linspace(0, 1, len(x))) + fig = plt.figure(dpi=100, figsize=figsize) + ax = fig.add_subplot(111, projection='3d') + fig.patch.set_facecolor('black') + ax.tick_params(axis='both', colors='white') + ax.grid(True, color='grey', linestyle='--', linewidth=0.5) + ax.set_facecolor('black') # Set plot background color to black + ax.scatter(x, y, z, color=dotcolors, s=1) + ax.plot_surface(mesh_x, mesh_y, mesh_z, rstride=4, cstride=4, facecolors=bm, shade=False) + ax.view_init(elev=el, azim=az) + ax.set_xlim([-limits, limits]) + ax.set_ylim([-limits, limits]) + ax.set_zlim([-limits, limits]) + ax.set_xlabel('x [GEO]', color='white') # Set x-axis label color to white + ax.set_ylabel('y [GEO]', color='white') # Set y-axis label color to white + ax.set_zlabel('z [GEO]', color='white') # Set z-axis label color to white + ax.tick_params(axis='x', colors='white') # Set x-axis tick color to white + ax.tick_params(axis='y', colors='white') # Set y-axis tick color to white + ax.tick_params(axis='z', colors='white') # Set z-axis tick color to white + ax.set_aspect('equal') + fig, ax = make_black(fig, ax) + if save_path: + save_plot(fig, save_path) + return fig, ax + + +def make_white(fig, *axes): + fig.patch.set_facecolor('white') + + for ax in axes: + ax.set_facecolor('white') + ax_items = [ax.title, ax.xaxis.label, ax.yaxis.label] + if hasattr(ax, 'zaxis'): + ax_items.append(ax.zaxis.label) + ax_items += ax.get_xticklabels() + ax.get_yticklabels() + if hasattr(ax, 'get_zticklabels'): + ax_items += ax.get_zticklabels() + ax_items += ax.get_xticklines() + ax.get_yticklines() + if hasattr(ax, 'get_zticklines'): + ax_items += ax.get_zticklines() + for item in ax_items: + item.set_color('black') + + return fig, axes + + +def make_black(fig, *axes): + fig.patch.set_facecolor('black') + + for ax in axes: + ax.set_facecolor('black') + ax_items = [ax.title, ax.xaxis.label, ax.yaxis.label] + if hasattr(ax, 'zaxis'): + ax_items.append(ax.zaxis.label) + ax_items += ax.get_xticklabels() + ax.get_yticklabels() + if hasattr(ax, 'get_zticklabels'): + ax_items += ax.get_zticklabels() + ax_items += ax.get_xticklines() + ax.get_yticklines() + if hasattr(ax, 'get_zticklines'): + ax_items += ax.get_zticklines() + for item in ax_items: + item.set_color('white') + + return fig, axes + + +# ##################################################################### # Formatting x axis -###################################################################### -def format_xaxis_decimal_year(time_array, ax): - n = 5 # Number of nearly evenly spaced points to select +# ##################################################################### +def date_format(time_array, ax): + n = 6 # Number of nearly evenly spaced points to select time_span_in_months = (time_array[-1].datetime - time_array[0].datetime).days / 30 if time_span_in_months < 1: # Get the time span in hours @@ -642,7 +536,7 @@ def format_xaxis_decimal_year(time_array, ax): selected_month_year_strings = [t.strftime('%d-%b-%Y') for t in selected_times] else: # Get the first of n nearly evenly spaced months in the time - step = int(len(time_array) / (n - 1)) + step = int(len(time_array) / (n - 1)) - 1 selected_times = time_array[::step] selected_month_year_strings = [t.strftime('%b-%Y') for t in selected_times] selected_decimal_years = [t.decimalyear for t in selected_times] @@ -658,12 +552,6 @@ def format_xaxis_decimal_year(time_array, ax): def save_plot_to_pdf(figure, pdf_path): - ###################################################################### - # Save figures appended to a pdf. - ###################################################################### - import io - import os - import re global save_plot_to_pdf_call_count save_plot_to_pdf_call_count += 1 if '~' == pdf_path[0]: @@ -702,7 +590,7 @@ def save_plot_to_pdf(figure, pdf_path): return -def save_plot_to_png(figure, save_path, dpi=200): +def save_plot(figure, save_path, dpi=200): """ Save a Python figure as a PNG image. @@ -713,7 +601,13 @@ def save_plot_to_png(figure, save_path, dpi=200): Returns: None """ + if save_path.lower().endswith('.pdf'): + save_plot_to_pdf(figure, save_path) + return try: + base_name, extension = os.path.splitext(save_path) + if extension.lower() != '.png': + save_path = base_name + '.png' # Save the figure as a PNG image figure.savefig(save_path, dpi=dpi, bbox_inches='tight') plt.close(figure) # Close the figure to release resources From 64f78ac9f2dc4473a2f9f8ffe7e5a5417c33507d Mon Sep 17 00:00:00 2001 From: yeager7 Date: Thu, 18 Jul 2024 16:22:04 -0700 Subject: [PATCH 06/55] added keplerian orbital element plots --- ssapy/plotUtils.py | 239 ++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 237 insertions(+), 2 deletions(-) diff --git a/ssapy/plotUtils.py b/ssapy/plotUtils.py index 615823b..97c6f68 100644 --- a/ssapy/plotUtils.py +++ b/ssapy/plotUtils.py @@ -1,6 +1,6 @@ from .body import get_body -from .compute import groundTrack, lagrange_points_lunar_frame, lagrange_points_lunar_frame_ -from .constants import RGEO, LD, EARTH_RADIUS, MOON_RADIUS +from .compute import groundTrack, lagrange_points_lunar_frame, calculate_orbital_elements +from .constants import RGEO, LD, EARTH_RADIUS, MOON_RADIUS, EARTH_MU, MOON_MU from .coordinates import gcrf_to_itrf, gcrf_to_lunar_fixed, gcrf_to_lunar from .utils import find_file, Time, find_smallest_bounding_cube @@ -11,6 +11,7 @@ import matplotlib.pyplot as plt import matplotlib.cm as cm +from matplotlib import mplcolors from PyPDF2 import PdfMerger from matplotlib.backends.backend_pdf import PdfPages @@ -468,6 +469,240 @@ def globe_plot(r, t, limits=False, title='', figsize=(7, 8), save_path=False, el return fig, ax +def koe_plot(r, v, t=Time("2025-01-01", scale='utc') + np.linspace(0, int(1 * 365.25), int(365.25 * 24)), elements=['a', 'e', 'i'], save_path=False, body='Earth'): + if 'earth' in body.lower(): + orbital_elements = calculate_orbital_elements(r, v, mu_barycenter=EARTH_MU) + else: + orbital_elements = calculate_orbital_elements(r, v, mu_barycenter=MOON_MU) + fig, ax1 = plt.subplots(dpi=100) + ax1.plot([], [], label='semi-major axis [GEO]', c='C0', linestyle='-') + ax2 = ax1.twinx() + make_white(fig, *[ax1, ax2]) + + ax1.plot(Time(t).decimalyear, [x for x in orbital_elements['e']], label='eccentricity', c='C1') + ax1.plot(Time(t).decimalyear, [x for x in orbital_elements['i']], label='inclination [rad]', c='C2') + ax1.set_xlabel('Year') + ax1.set_ylim((0, np.pi / 2)) + ylabel = ax1.set_ylabel('', color='black') + x = ylabel.get_position()[0] + 0.05 + y = ylabel.get_position()[1] + fig.text(x - 0.001, y - 0.225, 'Eccentricity', color='C1', rotation=90) + fig.text(x, y - 0.05, '/', color='k', rotation=90) + fig.text(x, y - 0.025, 'Inclination [Radians]', color='C2', rotation=90) + + ax1.legend(loc='upper left') + a = [x / RGEO for x in orbital_elements['a']] + ax2.plot(Time(t).decimalyear, a, label='semi-major axis [GEO]', c='C0', linestyle='-') + ax2.set_ylabel('semi-major axis [GEO]', color='C0') + ax2.yaxis.label.set_color('C0') + ax2.tick_params(axis='y', colors='C0') + ax2.spines['right'].set_color('C0') + + if np.abs(np.max(a) - np.min(a)) < 2: + ax2.set_ylim((np.min(a) - 0.5, np.max(a) + 0.5)) + date_format(t, ax1) + + plt.show(block=False) + if save_path: + save_plot(fig, save_path) + return fig, ax1 + + +def koe_2dhist(stable_data, title="Initial orbital elements of\n1 year stable cislunar orbits", limits=[1, 50], bins=200, logscale=False, cmap='coolwarm', save_path=False): + if logscale or logscale == 'log': + norm = mplcolors.LogNorm(limits[0], limits[1]) + else: + norm = mplcolors.Normalize(limits[0], limits[1]) + fig, axes = plt.subplots(dpi=100, figsize=(10, 8), nrows=3, ncols=3) + st = fig.suptitle(title, fontsize=12) + st.set_x(0.46) + st.set_y(0.9) + ax = axes.flat[0] + ax.hist2d([x / RGEO for x in stable_data.a], [x for x in stable_data.e], bins=bins, norm=norm, cmap=cmap) + ax.set_xlabel("") + ax.set_ylabel("eccentricity") + ax.set_xticks(np.arange(1, 20, 2)) + ax.set_yticks(np.arange(0, 1, 0.2)) + ax.set_xlim((1, 18)) + axes.flat[1].set_axis_off() + axes.flat[2].set_axis_off() + + ax = axes.flat[3] + ax.hist2d([x / RGEO for x in stable_data.a], [np.degrees(x) for x in stable_data.i], bins=bins, norm=norm, cmap=cmap) + ax.set_xlabel("") + ax.set_ylabel("inclination [deg]") + ax.set_xticks(np.arange(1, 20, 2)) + ax.set_yticks(np.arange(0, 91, 15)) + ax.set_xlim((1, 18)) + ax = axes.flat[4] + ax.hist2d([x for x in stable_data.e], [np.degrees(x) for x in stable_data.i], bins=bins, norm=norm, cmap=cmap) + ax.set_xlabel("") + ax.set_ylabel("") + ax.set_xticks(np.arange(0, 1, 0.2)) + ax.set_yticks(np.arange(0, 91, 15)) + axes.flat[5].set_axis_off() + + ax = axes.flat[6] + ax.hist2d([x / RGEO for x in stable_data.a], [np.degrees(x) for x in stable_data.ta], bins=bins, norm=norm, cmap=cmap) + ax.set_xlabel("semi-major axis [GEO]") + ax.set_ylabel("True Anomaly [deg]") + ax.set_xticks(np.arange(1, 20, 2)) + ax.set_yticks(np.arange(0, 361, 60)) + ax.set_xlim((1, 18)) + ax = axes.flat[7] + ax.hist2d([x for x in stable_data.e], [np.degrees(x) for x in stable_data.ta], bins=bins, norm=norm, cmap=cmap) + ax.set_xlabel("eccentricity") + ax.set_ylabel("") + ax.set_xticks(np.arange(0, 1, 0.2)) + ax.set_yticks(np.arange(0, 361, 60)) + ax = axes.flat[8] + ax.hist2d([np.degrees(x) for x in stable_data.i], [np.degrees(x) for x in stable_data.ta], bins=bins, norm=norm, cmap=cmap) + ax.set_xlabel("inclination [deg]") + ax.set_ylabel("") + ax.set_xticks(np.arange(0, 91, 15)) + ax.set_yticks(np.arange(0, 361, 60)) + + im = fig.subplots_adjust(right=0.8) + cbar_ax = fig.add_axes([0.82, 0.15, 0.01, 0.7]) + fig.colorbar(im, cax=cbar_ax, norm=norm, cmap=cmap) + fig, ax = make_white(fig, ax) + if save_path: + save_plot(fig, save_path) + return fig + + + +def scatter2d(x, y, cs, xlabel='x', ylabel='y', title='', cbar_label='', dotsize=1, colorsMap='jet', colorscale='linear', colormin=False, colormax=False, save_path=False): + fig = plt.figure() + ax = fig.add_subplot(111) + if colormax is False: + colormax = np.max(cs) + if colormin is False: + colormin = np.min(cs) + cm = plt.get_cmap(colorsMap) + if colorscale == 'linear': + cNorm = mplcolors.Normalize(vmin=colormin, vmax=colormax) + elif colorscale == 'log': + cNorm = mplcolors.LogNorm(vmin=colormin, vmax=colormax) + scalarMap = cm.ScalarMappable(norm=cNorm, cmap=cm) + ax.scatter(x, y, c=scalarMap.to_rgba(cs), s=dotsize) + ax.set_xlabel(xlabel) + ax.set_ylabel(ylabel) + ax.set_title(title) + scalarMap.set_array(cs) + fig.colorbar(scalarMap, shrink=.5, label=f'{cbar_label}', pad=0.04) + plt.tight_layout() + fig, ax = make_black(fig, ax) + plt.show(block=False) + if save_path: + save_plot(fig, save_path) + return + + +def scatter3d(x, y=None, z=None, cs=None, xlabel='x', ylabel='y', zlabel='z', cbar_label='', dotsize=1, colorsMap='jet', title='', save_path=False): + fig = plt.figure() + ax = fig.add_subplot(111, projection='3d') + if x.ndim > 1: + r = x + x = r[:, 0] + y = r[:, 1] + z = r[:, 2] + if cs is None: + ax.scatter(x, y, z, s=dotsize) + else: + cm = plt.get_cmap(colorsMap) + cNorm = mplcolors.Normalize(vmin=min(cs), vmax=max(cs)) + scalarMap = cm.ScalarMappable(norm=cNorm, cmap=cm) + ax.scatter(x, y, z, c=scalarMap.to_rgba(cs), s=dotsize) + scalarMap.set_array(cs) + fig.colorbar(scalarMap, shrink=.5, label=f'{cbar_label}', pad=0.075) + ax.set_xlabel(xlabel) + ax.set_ylabel(ylabel) + ax.set_zlabel(zlabel) + plt.title(title) + plt.tight_layout() + fig, ax = make_black(fig, ax) + plt.show(block=False) + if save_path: + save_plot(fig, save_path) + return fig, ax + + +def dotcolors_scaled(num_colors): + return cm.rainbow(np.linspace(0, 1, num_colors)) + + +# Make a plot of multiple cislunar orbit in GCRF frame. +def orbit_divergence_plot(rs, r_moon=[], t=False, limits=False, title='', save_path=False): + if limits is False: + limits = np.nanmax(np.linalg.norm(rs, axis=1) / RGEO) * 1.2 + print(f'limits: {limits}') + if np.size(r_moon) < 1: + moon = get_body("moon") + r_moon = moon.position(t) + else: + # print('Lunar position(s) provided.') + if r_moon.ndim != 2: + raise IndexError(f"input moon data shape: {np.shape(r_moon)}, input should be 2 dimensions.") + return None + if np.shape(r_moon)[1] == 3: + r_moon = r_moon.T + # print(f"Tranposed input to {np.shape(r_moon)}") + fig = plt.figure(dpi=100, figsize=(15, 4)) + for i in range(rs.shape[-1]): + r = rs[:, :, i] + x = r[:, 0] / RGEO + y = r[:, 1] / RGEO + z = r[:, 2] / RGEO + xm = r_moon[0] / RGEO + ym = r_moon[1] / RGEO + zm = r_moon[2] / RGEO + dotcolors = cm.rainbow(np.linspace(0, 1, len(x))) + + # Creating plot + plt.subplot(1, 3, 1) + plt.scatter(x, y, color=dotcolors, s=1) + plt.scatter(0, 0, color="blue", s=50) + plt.scatter(xm, ym, color="grey", s=5) + plt.axis('scaled') + plt.xlabel('x [GEO]') + plt.ylabel('y [GEO]') + plt.xlim((-limits, limits)) + plt.ylim((-limits, limits)) + plt.text(x[0], y[0], '$\leftarrow$ start') + plt.text(x[-1], y[-1], '$\leftarrow$ end') + + plt.subplot(1, 3, 2) + plt.scatter(x, z, color=dotcolors, s=1) + plt.scatter(0, 0, color="blue", s=50) + plt.scatter(xm, zm, color="grey", s=5) + plt.axis('scaled') + plt.xlabel('x [GEO]') + plt.ylabel('z [GEO]') + plt.xlim((-limits, limits)) + plt.ylim((-limits, limits)) + plt.text(x[0], z[0], '$\leftarrow$ start') + plt.text(x[-1], z[-1], '$\leftarrow$ end') + plt.title(f'{title}') + + plt.subplot(1, 3, 3) + plt.scatter(y, z, color=dotcolors, s=1) + plt.scatter(0, 0, color="blue", s=50) + plt.scatter(ym, zm, color="grey", s=5) + plt.axis('scaled') + plt.xlabel('y [GEO]') + plt.ylabel('z [GEO]') + plt.xlim((-limits, limits)) + plt.ylim((-limits, limits)) + plt.text(y[0], z[0], '$\leftarrow$ start') + plt.text(y[-1], z[-1], '$\leftarrow$ end') + plt.tight_layout() + plt.show(block=False) + if save_path: + save_plot(fig, save_path) + return + + def make_white(fig, *axes): fig.patch.set_facecolor('white') From 0d9c443a94290f241715dddd5a4c278a7f14fc60 Mon Sep 17 00:00:00 2001 From: yeager7 Date: Thu, 18 Jul 2024 16:29:30 -0700 Subject: [PATCH 07/55] simply.py is a useful file to add functions that combine many other functionalities within ssapy. --- ssapy/simple.py | 137 ++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 137 insertions(+) create mode 100644 ssapy/simple.py diff --git a/ssapy/simple.py b/ssapy/simple.py new file mode 100644 index 0000000..5c3c2b1 --- /dev/null +++ b/ssapy/simple.py @@ -0,0 +1,137 @@ +# flake8: noqa: E501 + +""" +Functions to simplify certain aspects of SSAPy. +For when you want a quick answer and do not need want to use high fidelity defaults. +""" + +from .utils import get_times, points_on_circle +from . import Orbit, rv +from .accel import AccelKepler, AccelSolRad, AccelEarthRad, AccelDrag +from .body import get_body +from .gravity import AccelHarmonic, AccelThirdBody +from .propagator import RK78Propagator + +from astropy.time import Time +import numpy as np + + +def keplerian_prop(integration_timestep=10): + return RK78Propagator(AccelKepler(), h=integration_timestep) + + +accel_3_cache = None +def threebody_prop(integration_timestep=10): + global accel_3_cache + if accel_3_cache is None: + accel_3_cache = AccelKepler() + AccelThirdBody(get_body("moon")) + return RK78Propagator(accel_3_cache, h=integration_timestep) + + +accel_4_cache = None +def fourbody_prop(integration_timestep=10): + global accel_4_cache + if accel_4_cache is None: + accel_4_cache = AccelKepler() + AccelThirdBody(get_body("moon")) + AccelThirdBody(get_body("Sun")) + return RK78Propagator(accel_4_cache, h=integration_timestep) + + +accel_best_cache = None +def best_prop(integration_timestep=10, kwargs=dict(mass=250, area=.022, CD=2.3, CR=1.3)): + global accel_best_cache + if accel_best_cache is None: + aEarth = AccelKepler() + AccelHarmonic(get_body("Earth", model="EGM2008"), 140, 140) + aMoon = AccelThirdBody(get_body("moon")) + AccelHarmonic(get_body("moon"), 20, 20) + aSun = AccelThirdBody(get_body("Sun")) + aMercury = AccelThirdBody(get_body("Mercury")) + aVenus = AccelThirdBody(get_body("Venus")) + aMars = AccelThirdBody(get_body("Mars")) + aJupiter = AccelThirdBody(get_body("Jupiter")) + aSaturn = AccelThirdBody(get_body("Saturn")) + aUranus = AccelThirdBody(get_body("Uranus")) + aNeptune = AccelThirdBody(get_body("Neptune")) + nonConservative = AccelSolRad(**kwargs) + AccelEarthRad(**kwargs) + AccelDrag(**kwargs) + planets = aMercury + aVenus + aMars + aJupiter + aSaturn + aUranus + aNeptune + accel_best_cache = aEarth + aMoon + aSun + planets + nonConservative + return RK78Propagator(accel_best_cache, h=integration_timestep) + + +def ssapy_kwargs(mass=250, area=0.022, CD=2.3, CR=1.3): + # Asteroid parameters + kwargs = dict( + mass=mass, # [kg] + area=area, # [m^2] + CD=CD, # Drag coefficient + CR=CR, # Radiation pressure coefficient + ) + return kwargs + + +def ssapy_prop(integration_timestep=60, propkw=ssapy_kwargs()): + # Accelerations - pass a body object or string of body name. + moon = get_body("moon") + sun = get_body("Sun") + Mercury = get_body("Mercury") + Venus = get_body("Venus") + Earth = get_body("Earth", model="EGM2008") + Mars = get_body("Mars") + Jupiter = get_body("Jupiter") + Saturn = get_body("Saturn") + Uranus = get_body("Uranus") + Neptune = get_body("Neptune") + aEarth = AccelKepler() + AccelHarmonic(Earth, 140, 140) + aSun = AccelThirdBody(sun) + aMoon = AccelThirdBody(moon) + AccelHarmonic(moon, 20, 20) + aSolRad = AccelSolRad(**propkw) + aEarthRad = AccelEarthRad(**propkw) + accel = aEarth + aMoon + aSun + aSolRad + aEarthRad + # Build propagator + prop = RK78Propagator(accel, h=integration_timestep) + return prop + + +# Uses the current best propagator and acceleration models in ssapy +def ssapy_orbit(orbit=None, a=None, e=0, i=0, pa=0, raan=0, ta=0, r=None, v=None, duration=(30, 'day'), freq=(1, 'hr'), start_date="2025-01-01", t=None, integration_timestep=10, mass=250, area=0.022, CD=2.3, CR=1.3, prop=ssapy_prop()): + # Everything is in SI units, except time. + # density #kg/m^3 --> density + t0 = Time(start_date, scale='utc') + if t is None: + time_is_None = True + t = get_times(duration=duration, freq=freq, t=t) + else: + time_is_None = False + + if orbit is not None: + pass + elif a is not None: + kElements = [a, e, i, pa, raan, ta] + orbit = Orbit.fromKeplerianElements(*kElements, t0) + elif r is not None and v is not None: + orbit = Orbit(r, v, t0) + else: + raise ValueError("Either Keplerian elements (a, e, i, pa, raan, ta) or position and velocity vectors (r, v) must be provided.") + + try: + r, v = rv(orbit, t, prop) + if time_is_None: + return r, v, t + else: + return r, v + except (RuntimeError, ValueError) as err: + print(err) + return np.nan, np.nan, np.nan + + +# Generate orbits near stable orbit. +def get_similar_orbits(r0, v0, rad=1e5, num_orbits=4, duration=(90, 'days'), freq=(1, 'hour'), start_date="2025-1-1", mass=250): + r0 = np.reshape(r0, (1, 3)) + v0 = np.reshape(v0, (1, 3)) + print(r0, v0) + for idx, point in enumerate(points_on_circle(r0, v0, rad=rad, num_points=num_orbits)): + # Calculate entire satellite trajectory + r, v = ssapy_orbit(r=point, v=v0, duration=duration, freq=freq, start_date=start_date, integration_timestep=10, mass=mass, area=mass / 19000 + 0.01, CD=2.3, CR=1.3) + if idx == 0: + trajectories = np.concatenate((r0, v0), axis=1)[:len(r)] + rv = np.concatenate((r, v), axis=1) + trajectories = np.dstack((trajectories, rv)) + return trajectories From 1e95465796b4ac407320a83e34812293aecc8726 Mon Sep 17 00:00:00 2001 From: yeager7 Date: Thu, 18 Jul 2024 16:30:16 -0700 Subject: [PATCH 08/55] added points_on_circle for use in simply.py --- ssapy/utils.py | 154 +++++++++++++------------------------------------ 1 file changed, 39 insertions(+), 115 deletions(-) diff --git a/ssapy/utils.py b/ssapy/utils.py index 17b16aa..14801e3 100644 --- a/ssapy/utils.py +++ b/ssapy/utils.py @@ -1329,121 +1329,6 @@ def find_nearest_indices(A, B): return nearest_indices -def integrate_orbit_best_model(r=None, v=None, t=None, koe=None, duration=None, freq=None, start_date=None, mass=250, area=0.22): - """ - Integrate satellite orbit trajectory using either (r, v) or Keplerian orbital elements. - - Parameters: - ----------- - r : array_like, optional - Initial position vector [x, y, z] in meters. - v : array_like, optional - Initial velocity vector [vx, vy, vz] in meters per second. - t : array_like, optional - Array of time values for integration in seconds since the start date. - koe : dict, optional - Dictionary containing Keplerian orbital elements: - - 'a': Semi-major axis in meters - - 'e': Eccentricity - - 'i': Inclination in radians - - 'trueAnomaly': True anomaly in radians - - 'pa': Argument of perigee in radians - - 'raan': Right ascension of ascending node in radians - duration : tuple, optional - Duration of integration (value, unit), e.g., (30, 'day'). - freq : tuple, optional - Frequency of integration steps (value, unit), e.g., (1, 'hr'). - start_date : str, optional - Start date in the format "YYYY-MM-DD" for time calculations. - mass: float, optional - Mass of the orbiting object in kg. - area: float, optional - Cross sectional area of the orbiting object in m^2. - Returns: - -------- - r : ndarray - Array of position vectors [x, y, z] over the integrated trajectory. - v : ndarray - Array of velocity vectors [vx, vy, vz] corresponding to the position vectors. - - Notes: - ------ - The function integrates the satellite orbit trajectory using either initial position and - velocity vectors or Keplerian orbital elements. - - Either (r, v) or 'koe' must be provided, but not both. - If 'koe' are provided, Keplerian-to-Cartesian conversion must be implemented. - - If 't' is not provided, 'duration', 'freq', and 'start_date' must be provided for time calculation. - """ - - from .accel import AccelKepler, AccelEarthRad, AccelSolRad, AccelDrag - from .body import get_body - from .compute import rv - from .gravity import AccelThirdBody, AccelHarmonic - from .orbit import Orbit - from .propagator import RK78Propagator - - # Time span of integration #Only takes integer number of days. Unless providing your own time object. - if t is None and start_date is None: - raise ValueError("Either an array of times 't' or a 'start_date', 'duration' and 'freq' must be provided.") - if t is None: - t = get_times(duration=duration, freq=freq, t=start_date) - - if (r is None or v is None) and koe is None: - raise ValueError("Either (r, v) or 'koe' must be provided.") - - if (r is not None and v is not None) and koe is not None: - raise ValueError("Both (r, v) and 'koe' cannot be provided simultaneously.") - - # Properties of sat - kwargs = dict( - mass=mass, # [kg] --> was 1e4 - area=area, # [m^2] - CD=2.3, # Drag coefficient - CR=1.3, # Radiation pressure coefficient - ) - - if koe is not None: - # Extract Keplerian orbital elements from the dictionary - a = koe.get('a', None) - e = koe.get('e', None) - i = koe.get('i', None) - trueAnomaly = koe.get('trueAnomaly', None) - pa = koe.get('pa', None) - raan = koe.get('raan', None) - if None in (a, e, i, trueAnomaly, pa, raan): - raise ValueError("Incomplete koe dictionary.") - kElements = [a, e, i, pa, raan, trueAnomaly] - orbit = Orbit.fromKeplerianElements(*kElements, t[0], propkw=kwargs) - elif (r is not None or v is not None): - orbit = Orbit(r, v, t[0], propkw=kwargs) - - # Most accurate trajectory - earth = get_body("earth") - moon = get_body("moon") - sun = get_body("sun") - - # Accelerations - pass a body object or string of bo - # dy name. - aEarth = AccelKepler() + AccelHarmonic(earth) - aMoon = AccelThirdBody(moon) + AccelHarmonic(moon) - aSun = AccelThirdBody(sun) - aSolRad = AccelSolRad() - aEarthRad = AccelEarthRad() - aDrag = AccelDrag() - accel = aEarth + aMoon + aSun + aSolRad + aEarthRad + aDrag - prop = RK78Propagator(accel, h=10.0) - - # Calculate entire satellite trajectory - try: - r, v = rv(orbit, t, prop) - return r, v - except (RuntimeError, ValueError) as err: - print(err) - return np.nan, np.nan - - def find_smallest_bounding_cube(r): """ Find the smallest bounding cube for a set of 3D coordinates. @@ -1475,3 +1360,42 @@ def find_smallest_bounding_cube(r): upper_bound = center + half_side_length return lower_bound, upper_bound + + +def points_on_circle(r, v, rad, num_points=4): + # Convert inputs to numpy arrays + r = np.array(r) + v = np.array(v) + + # Find the perpendicular vectors to the given vector v + if np.all(v[:2] == 0): + if np.all(v[2] == 0): + raise ValueError("The given vector v must not be the zero vector.") + else: + u = np.array([1, 0, 0]) + else: + u = np.array([-v[1], v[0], 0]) + u = u / np.linalg.norm(u) + w = np.cross(u, v) + w_norm = np.linalg.norm(w) + if w_norm < 1e-15: + # v is parallel to z-axis + w = np.array([0, 1, 0]) + else: + w = w / w_norm + # Generate a sequence of angles for equally spaced points + angles = np.linspace(0, 2 * np.pi, num_points, endpoint=False) + + # Compute the x, y, z coordinates of each point on the circle + x = rad * np.cos(angles) * u[0] + rad * np.sin(angles) * w[0] + y = rad * np.cos(angles) * u[1] + rad * np.sin(angles) * w[1] + z = rad * np.cos(angles) * u[2] + rad * np.sin(angles) * w[2] + + # Apply rotation about z-axis by 90 degrees + rot_matrix = np.array([[0, 1, 0], [-1, 0, 0], [0, 0, 1]]) + rotated_points = np.dot(rot_matrix, np.column_stack((x, y, z)).T).T + + # Translate the rotated points to the center point r + points_rotated = rotated_points + r.reshape(1, 3) + + return points_rotated \ No newline at end of file From 4a08a1877c7f94dec8669906e6826f1539189d9e Mon Sep 17 00:00:00 2001 From: SuperdoerTrav Date: Fri, 19 Jul 2024 10:44:54 -0700 Subject: [PATCH 09/55] added ipython module for interactive plots --- setup.py | 1 + 1 file changed, 1 insertion(+) diff --git a/setup.py b/setup.py index 2410a39..29d3fee 100644 --- a/setup.py +++ b/setup.py @@ -73,6 +73,7 @@ def build_extension(self, ext): 'pandas', 'h5py', 'pypdf2', + 'ipython', 'ipyvolume', 'ipython_genutils', 'jplephem', From ad95095ff477799d5c42b22e76888701937ac7eb Mon Sep 17 00:00:00 2001 From: SuperdoerTrav Date: Fri, 19 Jul 2024 10:59:08 -0700 Subject: [PATCH 10/55] added imageio for generation of gifs. --- setup.py | 1 + 1 file changed, 1 insertion(+) diff --git a/setup.py b/setup.py index 29d3fee..a0f3a0f 100644 --- a/setup.py +++ b/setup.py @@ -73,6 +73,7 @@ def build_extension(self, ext): 'pandas', 'h5py', 'pypdf2', + 'imageio', 'ipython', 'ipyvolume', 'ipython_genutils', From ce561a8827e5476f0ba93217908d938338df85ab Mon Sep 17 00:00:00 2001 From: SuperdoerTrav Date: Fri, 19 Jul 2024 12:52:22 -0700 Subject: [PATCH 11/55] added tqdm for progress bars --- setup.py | 1 + 1 file changed, 1 insertion(+) diff --git a/setup.py b/setup.py index a0f3a0f..74505fd 100644 --- a/setup.py +++ b/setup.py @@ -85,6 +85,7 @@ def build_extension(self, ext): 'sphinx-autobuild', 'sphinx-tabs', 'sphinx-automodapi', + 'tqdm', 'myst-parser', 'graphviz', ], From a1412f2a56a7ce884c8fabcf85ea01e9cccb2379 Mon Sep 17 00:00:00 2001 From: SuperdoerTrav Date: Fri, 19 Jul 2024 12:55:42 -0700 Subject: [PATCH 12/55] added faster io functions and updated old io functions added by me --- ssapy/io.py | 398 ++++++++++++++++++++++++++++++++-------------------- 1 file changed, 242 insertions(+), 156 deletions(-) diff --git a/ssapy/io.py b/ssapy/io.py index 0927a60..6dce75b 100644 --- a/ssapy/io.py +++ b/ssapy/io.py @@ -1,5 +1,8 @@ import datetime import numpy as np +import os +import csv +import pandas as pd from astropy.time import Time import astropy.units as u from .constants import EARTH_MU @@ -586,87 +589,13 @@ def load_b3obs_file(file_name): } return catalog -###################################################################### -# MPI -###################################################################### - - -def mpi_scatter(scatter_array): - from mpi4py import MPI - - comm = MPI.COMM_WORLD # Defines the default communicator - num_procs = comm.Get_size() # Stores the number of processes in size. - rank = comm.Get_rank() # Stores the rank (pid) of the current process - # stat = MPI.Status() - print(f'Number of procs: {num_procs}, rank: {rank}') - remainder = np.size(scatter_array) % num_procs - base_load = np.size(scatter_array) // num_procs - if rank == 0: - print('All processors will process at least {0} simulations.'.format( - base_load)) - print('{0} processors will process an additional simulations'.format( - remainder)) - load_list = np.concatenate((np.ones(remainder) * (base_load + 1), - np.ones(num_procs - remainder) * base_load)) - if rank == 0: - print('load_list={0}'.format(load_list)) - if rank < remainder: - scatter_array_local = np.zeros(base_load + 1, dtype=np.int64) - else: - scatter_array_local = np.zeros(base_load, dtype=np.int64) - disp = np.zeros(num_procs) - for i in range(np.size(load_list)): - if i == 0: - disp[i] = 0 - else: - disp[i] = disp[i - 1] + load_list[i - 1] - comm.Scatterv([scatter_array, load_list, disp, MPI.DOUBLE], scatter_array_local) - print(f"Process {rank} received the scattered arrays: {scatter_array_local}") - return scatter_array_local, rank - - -def mpi_scatter_exclude_rank_0(scatter_array): - # Function is for rank 0 to be used as a saving processor - all other processors will complete tasks. - from mpi4py import MPI - - comm = MPI.COMM_WORLD - num_procs = comm.Get_size() - rank = comm.Get_rank() - print(f'Number of procs: {num_procs}, rank: {rank}') - - num_workers = num_procs - 1 - remainder = np.size(scatter_array) % num_workers - base_load = np.size(scatter_array) // num_workers - - if rank == 0: - print(f'All processors will process at least {base_load} simulations.') - print(f'{remainder} processors will process an additional simulation.') - - load_list = np.concatenate((np.zeros(1), np.ones(remainder) * (base_load + 1), - np.ones(num_workers - remainder) * base_load)) - - if rank == 0: - print(f'load_list={load_list}') - - scatter_array_local = np.zeros(int(load_list[rank]), dtype=np.int64) - - disp = np.zeros(num_procs) - for i in range(1, num_procs): - disp[i] = disp[i - 1] + load_list[i - 1] - - if rank == 0: - dummy_recvbuf = np.zeros(1, dtype=np.int64) - comm.Scatterv([scatter_array, load_list, disp, MPI.INT64_T], dummy_recvbuf) - else: - comm.Scatterv([scatter_array, load_list, disp, MPI.INT64_T], scatter_array_local) - print(f"Process {rank} received the {len(scatter_array_local)} element scattered array: {scatter_array_local}") - - return scatter_array_local, rank - # ============================================================================= # File Handling Functions # ============================================================================= -import os +def file_exists(filename): + from glob import glob + name, _ = os.path.splitext(filename) + return bool(glob(f"{name}.*")) def exists(pathname): @@ -703,8 +632,26 @@ def rmfile(pathname): return -def listdir(dir_path='*', files_only=False, exclude=None): +def sortbynum(files): + import re + if len(files[0].split('/')) > 1: + files_shortened = [] + file_prefix = '/'.join(files[0].split('/')[:-1]) + for file in files: + files_shortened.append(file.split('/')[-1]) + files_sorted = sorted(files_shortened, key=lambda x: float(re.findall("(\d+)", x)[0])) + sorted_files = [] + for file in files_sorted: + sorted_files.append(f'{file_prefix}/{file}') + else: + sorted_files = sorted(files, key=lambda x: float(re.findall("(\d+)", x)[0])) + return sorted_files + + +def listdir(dir_path='*', files_only=False, exclude=None, sorted=False): from glob import glob + if '*' not in dir_path: + dir_path = os.path.join(dir_path, '*') expanded_paths = glob(dir_path) if files_only: @@ -717,7 +664,10 @@ def listdir(dir_path='*', files_only=False, exclude=None): if exclude: new_files = [file for file in files if exclude not in os.path.basename(file)] files = new_files - return sorted(files) + if sorted: + return sortbynum(files) + else: + return files def get_memory_usage(): @@ -762,6 +712,8 @@ def merge_dicts(file_names, save_path): print('Beginning final save.') psave(save_path, master_dict) return + + ###################################################################### # Sliceable Numpys save and load ###################################################################### @@ -794,7 +746,7 @@ def npload(filename_): ###################################################################### -def h5append(filename, pathname, append_data): +def append_h5(filename, pathname, append_data): """ Append data to key in HDF5 file. @@ -807,20 +759,21 @@ def h5append(filename, pathname, append_data): None """ try: - with h5py.File(filename, "r+") as f: + with h5py.File(filename, "a") as f: if pathname in f: path_data_old = np.array(f.get(pathname)) - new_data = np.append(path_data_old, np.array(append_data)) - f[pathname] = new_data - else: - f.create_dataset(pathname, data=np.array(append_data), maxshape=None) + append_data = np.append(path_data_old, np.array(append_data)) + del f[pathname] + f.create_dataset(pathname, data=np.array(append_data), maxshape=None) except FileNotFoundError: - print(f"File not found: {filename}") + print(f"File not found: {filename}\nCreating new dataset: {filename}") + save_h5(filename, pathname, append_data) + except (ValueError, KeyError) as err: print(f"Error: {err}") -def h5overwrite(filename, pathname, new_data): +def overwrite_h5(filename, pathname, new_data): """ Overwrite key in HDF5 file. @@ -832,93 +785,157 @@ def h5overwrite(filename, pathname, new_data): Returns: None """ + try: - with h5py.File(filename, "a") as f: - f.create_dataset(pathname, data=new_data, maxshape=None) - f.close() - except (FileNotFoundError, ValueError, KeyError): - try: - with h5py.File(filename, 'r+') as f: - del f[pathname] - f.close() - except (FileNotFoundError, ValueError, KeyError) as err: - print(f'Error: {err}') try: with h5py.File(filename, "a") as f: f.create_dataset(pathname, data=new_data, maxshape=None) f.close() - except (FileNotFoundError, ValueError, KeyError) as err: - print(f'File: {filename}{pathname}, Error: {err}') - - -def h5save(filename, pathname, data): + except (FileNotFoundError, ValueError, KeyError): + try: + with h5py.File(filename, 'r+') as f: + del f[pathname] + f.close() + except (FileNotFoundError, ValueError, KeyError) as err: + print(f'Error: {err}') + try: + with h5py.File(filename, "a") as f: + f.create_dataset(pathname, data=new_data, maxshape=None) + f.close() + except (FileNotFoundError, ValueError, KeyError) as err: + print(f'File: {filename}{pathname}, Error: {err}') + except (BlockingIOError, OSError) as err: + print(f"\n{err}\nPath: {pathname}\nFile: {filename}\n") + return None + + +def save_h5(filename, pathname, data): """ - Save data to HDF5 file. + Save data to HDF5 file with recursive attempt in case of write errors. Args: filename (str): The filename of the HDF5 file. pathname (str): The path to the data in the HDF5 file. data (any): The data to be saved. + max_retries (int): Maximum number of recursive retries in case of write errors. + retry_delay (tuple): A tuple representing the range of delay (in seconds) between retries. Returns: None """ - with h5py.File(filename, "a") as f: + try: try: - f.create_dataset(pathname, data=data, maxshape=None) + with h5py.File(filename, "a") as f: + f.create_dataset(pathname, data=data, maxshape=None) + f.flush() + return except ValueError as err: print(f"Did not save, key: {pathname} exists in file: {filename}. {err}") + return # If the key already exists, no need to retry + except (BlockingIOError, OSError) as err: + print(f"\n{err}\nPath: {pathname}\nFile: {filename}\n") + return None -def h5load(filename, pathname): +def read_h5(filename, pathname): """ Load data from HDF5 file. Args: - filename_ (str): The filename of the HDF5 file. - pathname_ (str): The path to the data in the HDF5 file. + filename (str): The filename of the HDF5 file. + pathname (str): The path to the data in the HDF5 file. Returns: The data loaded from the HDF5 file. """ - with h5py.File(filename, 'r') as f: - data = np.array(f.get(pathname)) - f.close() - return data + try: + with h5py.File(filename, 'r') as f: + data = f.get(pathname) + if data is None: + return None + else: + return np.array(data) + except (ValueError, KeyError, TypeError): + return None + except FileNotFoundError: + print(f'File not found. {filename}') + raise + except (BlockingIOError, OSError) as err: + print(f"\n{err}\nPath: {pathname}\nFile: {filename}\n") + raise -def h5loadall(filename_): - """ - Load all data from HDF5 file. +def read_h5_all(file_path): + data_dict = {} - Args: - filename_ (str): The filename of the HDF5 file. + with h5py.File(file_path, 'r') as file: + # Recursive function to traverse the HDF5 file and populate the dictionary + def traverse(group, path=''): + for key, item in group.items(): + new_path = f"{path}/{key}" if path else key - Returns: - A dictionary of data loaded from the HDF5 file. - """ - with h5py.File(filename_, "r") as f: - # List all groups - keys = list(f.keys()) - return_data = {key: np.array(f.get(key)) for key in keys} + if isinstance(item, h5py.Group): + traverse(item, path=new_path) + else: + data_dict[new_path] = item[()] + + traverse(file) + return data_dict - return return_data, keys +def combine_h5(filename, files, verbose=False, overwrite=False): + if verbose: + from tqdm import tqdm + iterable = enumerate(tqdm(files)) + else: + iterable = enumerate(files) + if overwrite: + rmfile(filename) + for idx, file in iterable: + if verbose: + iterable2 = tqdm(h5_keys(file)) + else: + iterable2 = files + for key in iterable2: + try: + if h5_key_exists(filename, key): + continue + save_h5(filename, key, read_h5(file, key)) + except TypeError as err: + print(read_h5(file, key)) + print(f'{err}, key: {key}, file: {file}') + print('Completed HDF5 merge.') -def h5keys(filename): + +def h5_keys(file_path): """ List all groups in HDF5 file. Args: - filename_ (str): The filename of the HDF5 file. + file_path (str): The file_path of the HDF5 file. Returns: A list of group keys in the HDF5 file. """ - with h5py.File(filename, "r") as f: - # List all groups - group_keys = list(f.keys()) - return group_keys + keys_list = [] + with h5py.File(file_path, 'r') as file: + # Recursive function to traverse the HDF5 file and collect keys + def traverse(group, path=''): + for key, item in group.items(): + new_path = f"{path}/{key}" if path else key + if isinstance(item, h5py.Group): + traverse(item, path=new_path) + else: + keys_list.append(new_path) + traverse(file) + return keys_list + + +def h5_root_keys(file_path): + with h5py.File(file_path, 'r') as file: + keys_in_root = list(file.keys()) + # print("Keys in the root group:", keys_in_root) + return keys_in_root def h5_key_exists(filename, key): @@ -940,7 +957,6 @@ def h5_key_exists(filename, key): return False -import pandas as pd ###################################################################### # CSV ###################################################################### @@ -953,7 +969,7 @@ def makedf(df): return df -def header_csv(file_name, sep=None): +def read_csv_header(file_name, sep=None): """ Get the header of a CSV file. @@ -964,7 +980,6 @@ def header_csv(file_name, sep=None): Returns: A list of the header fields. """ - import csv if sep is None: sep = guess_csv_delimiter(file_name) # Guess the delimiter with open(file_name, 'r') as infile: @@ -1024,8 +1039,32 @@ def read_csv(file_name, sep=None, dtypes=None, col=False, to_np=False, drop_nan= return df +def append_dict_to_csv(file_name, data_dict, delimiter='\t'): + # Extract keys and values from the dictionary + keys = list(data_dict.keys()) + values = list(data_dict.values()) + + # Determine the length of the arrays + array_length = len(values[0]) + + # Determine if file exists + file_exists = os.path.exists(file_name) + + # Open the CSV file in append mode + with open(file_name, 'a', newline='') as csvfile: + writer = csv.writer(csvfile, delimiter=delimiter) + + # Write header if file doesn't exist + if not file_exists: + writer.writerow(keys) + + # Write each element from arrays as a new row + for i in range(array_length): + row = [values[j][i] for j in range(len(keys))] + writer.writerow(row) + + def guess_csv_delimiter(file_name): - import csv """ Guess the delimiter used in a CSV file. @@ -1065,7 +1104,7 @@ def save_csv(file_name, df, sep='\t', dtypes=None): return -def append_csv(file_names, save_path='combined_data.csv', sep=None, dtypes=None, progress=None): +def append_csv(file_names, save_path='combined_data.csv', sep=None, dtypes=False, progress=None): """ Appends multiple CSV files into a single CSV file. @@ -1090,11 +1129,11 @@ def append_csv(file_names, save_path='combined_data.csv', sep=None, dtypes=None, if progress is not None: get_memory_usage() print(f"Appended {i+1} of {len(file_names)}.") - except FileNotFoundError: + except (FileNotFoundError, pd.errors.EmptyDataError, pd.errors.ParserError) as e: error_files.append(file) + print(f"Error processing file {file}: {e}") combined_df = pd.concat(dataframes, ignore_index=True) - if dtypes: combined_df = combined_df.astype(dtypes) @@ -1105,11 +1144,11 @@ def append_csv(file_names, save_path='combined_data.csv', sep=None, dtypes=None, print(f'The final dataframe has {combined_df.shape[0]} rows and {combined_df.shape[1]} columns.') if error_files: - print(f'The following files could not be found: {error_files}') + print(f'The following files ERRORED and were not included: {error_files}') + return def append_csv_on_disk(csv_files, output_file): - import csv # Assumes each file has the same delimiters delimiter = guess_csv_delimiter(csv_files[0]) # Open the output file for writing @@ -1137,6 +1176,42 @@ def append_csv_on_disk(csv_files, output_file): print(f'Completed appending of: {output_file}.') +def save_csv_header(filename, header, delimiter='\t'): + """ + Saves a header row to a CSV file with a specified delimiter. + + Parameters: + filename (str): The name of the file where the header will be saved. + header (list): A list of strings representing the column names. + delimiter (str, optional): The delimiter to use between columns in the CSV file. + Default is tab ('\t'). + + Example: + save_csv_header('output.csv', ['Name', 'Age', 'City'], delimiter=',') + """ + with open(filename, 'w', newline='') as csvfile: + writer = csv.writer(csvfile, delimiter=delimiter) + writer.writerow(header) + + +def save_csv_array_to_line(filename, array, delimiter='\t'): + """ + Appends a single row of data to a CSV file with a specified delimiter. + + Parameters: + filename (str): The name of the file to which the row will be appended. + array (list): A list of values representing a single row of data to be appended to the CSV file. + delimiter (str, optional): The delimiter to use between columns in the CSV file. + Default is tab ('\t'). + + Example: + save_csv_array_to_line('output.csv', ['Alice', 30, 'New York'], delimiter=',') + """ + with open(filename, 'a', newline='') as csvfile: + writer = csv.writer(csvfile, delimiter=delimiter) + writer.writerow(array) + + def save_csv_line(file_name, df, sep='\t', dtypes=None): """ Save a Pandas DataFrame to a CSV file, appending the DataFrame to the file if it exists. @@ -1159,23 +1234,26 @@ def save_csv_line(file_name, df, sep='\t', dtypes=None): return -# Create a lock to synchronize access to the file -import threading -file_lock = threading.Lock() +_column_data = None +def exists_in_csv(csv_file, column, number, sep='\t'): + try: + global _column_data + if _column_data is None: + _column_data = read_csv(csv_file, sep=sep, col=column, to_np=True) + return np.isin(number, _column_data) + except IOError: + return False -def exists_in_csv(csv_file, column_name, number, sep='\t'): - import csv - with file_lock: - try: - with open(csv_file, 'r') as f: - reader = csv.DictReader(f, delimiter=sep) - for row in reader: - if row[column_name] == str(number): - return True - except IOError: - return False - return False +def exists_in_csv_old(csv_file, column, number, sep='\t'): + try: + with open(csv_file, 'r') as f: + reader = csv.DictReader(f, delimiter=sep) + for row in reader: + if row[column] == str(number): + return True + except IOError: + return False def pd_flatten(data, factor=1): @@ -1198,3 +1276,11 @@ def str_to_array(s): def pdstr_to_arrays(df): return df.apply(str_to_array).to_numpy() + + +def allfiles(dirName=os.getcwd()): + # Get the list of all files in directory tree at given path + listOfFiles = list() + for (dirpath, dirnames, filenames) in os.walk(dirName): + listOfFiles += [os.path.join(dirpath, file) for file in filenames] + return listOfFiles From 87c37bacc10386cb460580c1aef018bb1f54f458 Mon Sep 17 00:00:00 2001 From: SuperdoerTrav Date: Fri, 19 Jul 2024 13:00:18 -0700 Subject: [PATCH 13/55] added functions to calculate the position of lagrange points --- ssapy/compute.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/ssapy/compute.py b/ssapy/compute.py index 33d8990..118c4ee 100644 --- a/ssapy/compute.py +++ b/ssapy/compute.py @@ -4,11 +4,10 @@ from .body import get_body from .constants import RGEO, LD, EARTH_RADIUS, EARTH_MU, MOON_RADIUS, MOON_MU -from .coordinates import rotation_matrix_from_vectors, angle_between_vectors, gcrf_to_itrf from .propagator import KeplerianPropagator from .utils import ( norm, normed, unitAngle3, LRU_Cache, lb_to_unit, sunPos, _gpsToTT, - iers_interp + iers_interp, rotation_matrix_from_vectors, angle_between_vectors, gcrf_to_itrf ) from .orbit import Orbit from .ellipsoid import Ellipsoid From 34da31f8f9fb8454d2cc34b10cf501ee5db90dbf Mon Sep 17 00:00:00 2001 From: SuperdoerTrav Date: Fri, 19 Jul 2024 13:01:01 -0700 Subject: [PATCH 14/55] removed due to circular import conflicts, coordinates are all together in utils.py where they initially were --- ssapy/coordinates.py | 653 ------------------------------------------- 1 file changed, 653 deletions(-) delete mode 100644 ssapy/coordinates.py diff --git a/ssapy/coordinates.py b/ssapy/coordinates.py deleted file mode 100644 index 3198912..0000000 --- a/ssapy/coordinates.py +++ /dev/null @@ -1,653 +0,0 @@ -# flake8: noqa: E501 - -from .constants import EARTH_RADIUS, WGS84_EARTH_OMEGA -from .accel import AccelKepler -from .body import get_body, MoonPosition -from .compute import groundTrack, rv -from .constants import RGEO -from .orbit import Orbit -from .propagator import RK78Propagator -from .utils import hms_to_dd, dd_to_hms, dd_to_dms - -import numpy as np -from astropy.time import Time - -# VECTOR FUNCTIONS FOR COORDINATE MATH -def unit_vector(vector): - """ Returns the unit vector of the vector.""" - return vector / np.linalg.norm(vector) - - -def getAngle(a, b, c): # a,b,c where b is the vertex - a = np.atleast_2d(a) - b = np.atleast_2d(b) - c = np.atleast_2d(c) - ba = np.subtract(a, b) - bc = np.subtract(c, b) - cosine_angle = np.sum(ba * bc, axis=-1) / (np.linalg.norm(ba, axis=-1) * np.linalg.norm(bc, axis=-1)) - return np.arccos(cosine_angle) - - -def angle_between_vectors(vector1, vector2): - return np.arccos(np.clip(np.dot(vector1, vector2) / (np.linalg.norm(vector1) * np.linalg.norm(vector2)), -1.0, 1.0)) - - -def rotation_matrix_from_vectors(vec1, vec2): - """ Find the rotation matrix that aligns vec1 to vec2 - :param vec1: A 3d "source" vector - :param vec2: A 3d "destination" vector - :return mat: A transform matrix (3x3) which when applied to vec1, aligns it with vec2. - """ - a, b = (vec1 / np.linalg.norm(vec1)).reshape(3), (vec2 / np.linalg.norm(vec2)).reshape(3) - v = np.cross(a, b) - c = np.dot(a, b) - s = np.linalg.norm(v) - kmat = np.array([[0, -v[2], v[1]], [v[2], 0, -v[0]], [-v[1], v[0], 0]]) - rotation_matrix = np.eye(3) + kmat + kmat.dot(kmat) * ((1 - c) / (s**2)) - return rotation_matrix - - -def normed(arr): - return arr / np.sqrt(np.einsum("...i,...i", arr, arr))[..., None] - - -def einsum_norm(a, indices='ij,ji->i'): - return np.sqrt(np.einsum(indices, a, a)) - - -def normSq(arr): - return np.einsum("...i,...i", arr, arr) - - -def norm(arr): - return np.sqrt(np.einsum("...i,...i", arr, arr)) - - -def rotate_vector(v_unit, theta, phi): - v_unit = v_unit / np.linalg.norm(v_unit, axis=-1) - if np.all(np.abs(v_unit) != np.max(np.abs(v_unit))): - perp_vector = np.cross(v_unit, np.array([1, 0, 0])) - else: - perp_vector = np.cross(v_unit, np.array([0, 1, 0])) - perp_vector /= np.linalg.norm(perp_vector) - - theta = np.radians(theta) - phi = np.radians(phi) - cos_theta = np.cos(theta) - sin_theta = np.sin(theta) - cos_phi = np.cos(phi) - sin_phi = np.sin(phi) - - R1 = np.array([ - [cos_theta + (1 - cos_theta) * perp_vector[0]**2, - (1 - cos_theta) * perp_vector[0] * perp_vector[1] - sin_theta * perp_vector[2], - (1 - cos_theta) * perp_vector[0] * perp_vector[2] + sin_theta * perp_vector[1]], - [(1 - cos_theta) * perp_vector[1] * perp_vector[0] + sin_theta * perp_vector[2], - cos_theta + (1 - cos_theta) * perp_vector[1]**2, - (1 - cos_theta) * perp_vector[1] * perp_vector[2] - sin_theta * perp_vector[0]], - [(1 - cos_theta) * perp_vector[2] * perp_vector[0] - sin_theta * perp_vector[1], - (1 - cos_theta) * perp_vector[2] * perp_vector[1] + sin_theta * perp_vector[0], - cos_theta + (1 - cos_theta) * perp_vector[2]**2] - ]) - - # Apply the rotation matrix to v_unit to get the rotated unit vector - v1 = np.dot(R1, v_unit) - - # Rotation matrix for rotation about v_unit - R2 = np.array([[cos_phi + (1 - cos_phi) * v_unit[0]**2, - (1 - cos_phi) * v_unit[0] * v_unit[1] - sin_phi * v_unit[2], - (1 - cos_phi) * v_unit[0] * v_unit[2] + sin_phi * v_unit[1]], - [(1 - cos_phi) * v_unit[1] * v_unit[0] + sin_phi * v_unit[2], - cos_phi + (1 - cos_phi) * v_unit[1]**2, - (1 - cos_phi) * v_unit[1] * v_unit[2] - sin_phi * v_unit[0]], - [(1 - cos_phi) * v_unit[2] * v_unit[0] - sin_phi * v_unit[1], - (1 - cos_phi) * v_unit[2] * v_unit[1] + sin_phi * v_unit[0], - cos_phi + (1 - cos_phi) * v_unit[2]**2]]) - - v2 = np.dot(R2, v1) - return v2 / np.linalg.norm(v2, axis=-1) - - -def rotate_points_3d(points, axis=np.array([0, 0, 1]), theta=-np.pi / 2): - """ - Rotate a set of 3D points about a 3D axis by an angle theta in radians. - - Args: - points (np.ndarray): The set of 3D points to rotate, as an Nx3 array. - axis (np.ndarray): The 3D axis to rotate about, as a length-3 array. Default is the z-axis. - theta (float): The angle to rotate by, in radians. Default is pi/2. - - Returns: - np.ndarray: The rotated set of 3D points, as an Nx3 array. - """ - # Normalize the axis to be a unit vector - axis = axis / np.linalg.norm(axis) - - # Compute the quaternion representing the rotation - qw = np.cos(theta / 2) - qx, qy, qz = axis * np.sin(theta / 2) - - # Construct the rotation matrix from the quaternion - R = np.array([ - [1 - 2 * qy**2 - 2 * qz**2, 2 * qx * qy - 2 * qz * qw, 2 * qx * qz + 2 * qy * qw], - [2 * qx * qy + 2 * qz * qw, 1 - 2 * qx**2 - 2 * qz**2, 2 * qy * qz - 2 * qx * qw], - [2 * qx * qz - 2 * qy * qw, 2 * qy * qz + 2 * qx * qw, 1 - 2 * qx**2 - 2 * qy**2] - ]) - - # Apply the rotation matrix to the set of points - rotated_points = np.dot(R, points.T).T - - return rotated_points - - -def perpendicular_vectors(v): - """Returns two vectors that are perpendicular to v and each other.""" - # Check if v is the zero vector - if np.allclose(v, np.zeros_like(v)): - raise ValueError("Input vector cannot be the zero vector.") - - # Choose an arbitrary non-zero vector w that is not parallel to v - w = np.array([1., 0., 0.]) - if np.allclose(v, w) or np.allclose(v, -w): - w = np.array([0., 1., 0.]) - u = np.cross(v, w) - if np.allclose(u, np.zeros_like(u)): - w = np.array([0., 0., 1.]) - u = np.cross(v, w) - w = np.cross(v, u) - - return u, w - - -def points_on_circle(r, v, rad, num_points=4): - # Convert inputs to numpy arrays - r = np.array(r) - v = np.array(v) - - # Find the perpendicular vectors to the given vector v - if np.all(v[:2] == 0): - if np.all(v[2] == 0): - raise ValueError("The given vector v must not be the zero vector.") - else: - u = np.array([1, 0, 0]) - else: - u = np.array([-v[1], v[0], 0]) - u = u / np.linalg.norm(u) - w = np.cross(u, v) - w_norm = np.linalg.norm(w) - if w_norm < 1e-15: - # v is parallel to z-axis - w = np.array([0, 1, 0]) - else: - w = w / w_norm - # Generate a sequence of angles for equally spaced points - angles = np.linspace(0, 2 * np.pi, num_points, endpoint=False) - - # Compute the x, y, z coordinates of each point on the circle - x = rad * np.cos(angles) * u[0] + rad * np.sin(angles) * w[0] - y = rad * np.cos(angles) * u[1] + rad * np.sin(angles) * w[1] - z = rad * np.cos(angles) * u[2] + rad * np.sin(angles) * w[2] - - # Apply rotation about z-axis by 90 degrees - rot_matrix = np.array([[0, 1, 0], [-1, 0, 0], [0, 0, 1]]) - rotated_points = np.dot(rot_matrix, np.column_stack((x, y, z)).T).T - - # Translate the rotated points to the center point r - points_rotated = rotated_points + r.reshape(1, 3) - - return points_rotated - - -def dms_to_rad(coords): - from astropy.coordinates import Angle - if isinstance(coords, (list, tuple)): - return [Angle(coord).radian for coord in coords] - else: - return Angle(coords).radian - return - - -def dms_to_deg(coords): - from astropy.coordinates import Angle - if isinstance(coords, (list, tuple)): - return [Angle(coord).deg for coord in coords] - else: - return Angle(coords).deg - return - - -def rad0to2pi(angles): - return (2 * np.pi + angles) * (angles < 0) + angles * (angles > 0) - - -def deg0to360(array_): - try: - return [i % 360 for i in array_] - except TypeError: - return array_ % 360 - - -def deg0to360array(array_): - return [i % 360 for i in array_] - - -def deg90to90(val_in): - if hasattr(val_in, "__len__"): - val_out = [] - for i, v in enumerate(val_in): - while v < -90: - v += 90 - while v > 90: - v -= 90 - val_out.append(v) - else: - while val_in < -90: - val_in += 90 - while val_in > 90: - val_in -= 90 - val_out = val_in - return val_out - - -def deg90to90array(array_): - return [i % 90 for i in array_] - - -def cart2sph_deg(x, y, z): - hxy = np.hypot(x, y) - r = np.hypot(hxy, z) - el = np.arctan2(z, hxy) * (180 / np.pi) - az = (np.arctan2(y, x)) * (180 / np.pi) - return az, el, r - - -def cart_to_cyl(x, y, z): - r = np.linalg.norm([x, y]) - theta = np.arctan2(y, x) - return r, theta, z - - -def inert2rot(x, y, xe, ye, xs=0, ys=0): # Places Earth at (-1,0) - earth_theta = np.arctan2(ye - ys, xe - xs) - theta = np.arctan2(y - ys, x - xs) - distance = np.sqrt(np.power((x - xs), 2) + np.power((y - ys), 2)) - xrot = distance * np.cos(np.pi + (theta - earth_theta)) - yrot = distance * np.sin(np.pi + (theta - earth_theta)) - return xrot, yrot - - -def sim_lonlatrad(x, y, z, xe, ye, ze, xs, ys, zs): - # convert all to geo coordinates - x = x - xe - y = y - ye - z = z - ze - xs = xs - xe - ys = ys - ye - zs = zs - ze - # convert x y z to lon lat radius - longitude, latitude, radius = cart2sph_deg(x, y, z) - slongitude, slatitude, sradius = cart2sph_deg(xs, ys, zs) - # correct so that Sun is at (0,0) - longitude = deg0to360(slongitude - longitude) - latitude = latitude - slatitude - return longitude, latitude, radius - - -def sun_ra_dec(time_): - out = get_body(Time(time_, format='mjd')) - return out.ra.to('rad').value, out.dec.to('rad').value - - -def ra_dec(r=None, v=None, x=None, y=None, z=None, vx=None, vy=None, vz=None, r_earth=np.array([0, 0, 0]), v_earth=np.array([0, 0, 0]), input_unit='si'): - if r is None or v is None: - if x is not None and y is not None and z is not None and vx is not None and vy is not None and vz is not None: - r = np.array([x, y, z]) - v = np.array([vx, vy, vz]) - else: - raise ValueError("Either provide r and v arrays or individual coordinates (x, y, z) and velocities (vx, vy, vz)") - - # Subtract Earth's position and velocity from the input arrays - r = r - r_earth - v = v - v_earth - - d_earth_mag = einsum_norm(r, 'ij,ij->i') - ra = rad0to2pi(np.arctan2(r[:, 1], r[:, 0])) # in radians - dec = np.arcsin(r[:, 2] / d_earth_mag) - return ra, dec - - -def lonlat_distance(lat1, lat2, lon1, lon2): - # Haversine formula - dlon = lon2 - lon1 - dlat = lat2 - lat1 - a = np.sin(dlat / 2)**2 + np.cos(lat1) * np.cos(lat2) * np.sin(dlon / 2)**2 - c = 2 * np.arcsin(np.sqrt(a)) - # Radius of earth in kilometers. Use 3956 for miles - # calculate the result - return (c * EARTH_RADIUS) - - -def altitude2zenithangle(altitude, deg=True): - if deg: - out = 90 - altitude - else: - out = np.pi / 2 - altitude - return out - - -def zenithangle2altitude(zenith_angle, deg=True): - if deg: - out = 90 - zenith_angle - else: - out = np.pi / 2 - zenith_angle - return out - - -def rightasension2hourangle(right_ascension, local_time): - if type(right_ascension) is not str: - right_ascension = dd_to_hms(right_ascension) - if type(local_time) is not str: - local_time = dd_to_dms(local_time) - _ra = float(right_ascension.split(':')[0]) - _lt = float(local_time.split(':')[0]) - if _ra > _lt: - __ltm, __lts = local_time.split(':')[1:] - local_time = f'{24 + _lt}:{__ltm}:{__lts}' - - return dd_to_dms(hms_to_dd(local_time) - hms_to_dd(right_ascension)) - - -def equatorial_to_horizontal(observer_latitude, declination, right_ascension=None, hour_angle=None, local_time=None, hms=False): - if right_ascension is not None: - hour_angle = rightasension2hourangle(right_ascension, local_time) - if hms: - hour_angle = hms_to_dd(hour_angle) - elif hour_angle is not None: - if hms: - hour_angle = hms_to_dd(hour_angle) - elif right_ascension is not None and hour_angle is not None: - print('Both right_ascension and hour_angle parameters are provided.\nUsing hour_angle for calculations.') - if hms: - hour_angle = hms_to_dd(hour_angle) - else: - print('Either right_ascension or hour_angle must be provided.') - - observer_latitude, hour_angle, declination = np.radians([observer_latitude, hour_angle, declination]) - - zenith_angle = np.arccos(np.sin(observer_latitude) * np.sin(declination) + np.cos(observer_latitude) * np.cos(declination) * np.cos(hour_angle)) - - altitude = zenithangle2altitude(zenith_angle, deg=False) - - _num = np.sin(declination) - np.sin(observer_latitude) * np.cos(zenith_angle) - _den = np.cos(observer_latitude) * np.sin(zenith_angle) - azimuth = np.arccos(_num / _den) - - if observer_latitude < 0: - azimuth = np.pi - azimuth - altitude, azimuth = np.degrees([altitude, azimuth]) - - return azimuth, altitude - - -def horizontal_to_equatorial(observer_latitude, azimuth, altitude): - altitude, azimuth, latitude = np.radians([altitude, azimuth, observer_latitude]) - zenith_angle = zenithangle2altitude(altitude) - - zenith_angle = [-zenith_angle if latitude < 0 else zenith_angle][0] - - declination = np.sin(latitude) * np.cos(zenith_angle) - declination = declination + (np.cos(latitude) * np.sin(zenith_angle) * np.cos(azimuth)) - declination = np.arcsin(declination) - - _num = np.cos(zenith_angle) - np.sin(latitude) * np.sin(declination) - _den = np.cos(latitude) * np.cos(declination) - hour_angle = np.arccos(_num / _den) - - if (latitude > 0 > declination) or (latitude < 0 < declination): - hour_angle = 2 * np.pi - hour_angle - - declination, hour_angle = np.degrees([declination, hour_angle]) - - return hour_angle, declination - - -_ecliptic = 0.409092601 # np.radians(23.43927944) -cos_ec = 0.9174821430960974 -sin_ec = 0.3977769690414367 - - -def equatorial_xyz_to_ecliptic_xyz(xq, yq, zq): - xc = xq - yc = cos_ec * yq + sin_ec * zq - zc = -sin_ec * yq + cos_ec * zq - return xc, yc, zc - - -def ecliptic_xyz_to_equatorial_xyz(xc, yc, zc): - xq = xc - yq = cos_ec * yc - sin_ec * zc - zq = sin_ec * yc + cos_ec * zc - return xq, yq, zq - - -def xyz_to_ecliptic(xc, yc, zc, xe=0, ye=0, ze=0, degrees=False): - x_ast_to_earth = xc - xe - y_ast_to_earth = yc - ye - z_ast_to_earth = zc - ze - d_earth_mag = np.sqrt(np.power(x_ast_to_earth, 2) + np.power(y_ast_to_earth, 2) + np.power(z_ast_to_earth, 2)) - ec_longitude = rad0to2pi(np.arctan2(y_ast_to_earth, x_ast_to_earth)) # in radians - ec_latitude = np.arcsin(z_ast_to_earth / d_earth_mag) - if degrees: - return np.degrees(ec_longitude), np.degrees(ec_latitude) - else: - return ec_longitude, ec_latitude - - -def xyz_to_equatorial(xq, yq, zq, xe=0, ye=0, ze=0, degrees=False): - # RA / DEC calculation - assumes XY plane to be celestial equator, and -x axis to be vernal equinox - x_ast_to_earth = xq - xe - y_ast_to_earth = yq - ye - z_ast_to_earth = zq - ze - d_earth_mag = np.sqrt(np.power(x_ast_to_earth, 2) + np.power(y_ast_to_earth, 2) + np.power(z_ast_to_earth, 2)) - ra = rad0to2pi(np.arctan2(y_ast_to_earth, x_ast_to_earth)) # in radians - dec = np.arcsin(z_ast_to_earth / d_earth_mag) - if degrees: - return np.degrees(ra), np.degrees(dec) - else: - return ra, dec - - -def ecliptic_xyz_to_equatorial(xc, yc, zc, xe=0, ye=0, ze=0, degrees=False): - # Convert ecliptic cartesian into equitorial cartesian - x_ast_to_earth, y_ast_to_earth, z_ast_to_earth = ecliptic_xyz_to_equatorial_xyz(xc - xe, yc - ye, zc - ze) - d_earth_mag = np.sqrt(np.power(x_ast_to_earth, 2) + np.power(y_ast_to_earth, 2) + np.power(z_ast_to_earth, 2)) - ra = rad0to2pi(np.arctan2(y_ast_to_earth, x_ast_to_earth)) # in radians - dec = np.arcsin(z_ast_to_earth / d_earth_mag) - if degrees: - return np.degrees(ra), np.degrees(dec) - else: - return ra, dec - - -def equatorial_to_ecliptic(right_ascension, declination, degrees=False): - ra, dec = np.radians(right_ascension), np.radians(declination) - ec_latitude = np.arcsin(cos_ec * np.sin(dec) - sin_ec * np.cos(dec) * np.sin(ra)) - ec_longitude = np.arctan((cos_ec * np.cos(dec) * np.sin(ra) + sin_ec * np.sin(dec)) / (np.cos(dec) * np.cos(ra))) - if degrees: - return deg0to360(np.degrees(ec_longitude)), np.degrees(ec_latitude) - else: - return rad0to2pi(ec_longitude), ec_latitude - - -def ecliptic_to_equatorial(lon, lat, degrees=False): - lon, lat = np.radians(lon), np.radians(lat) - ra = np.arctan((cos_ec * np.cos(lat) * np.sin(lon) - sin_ec * np.sin(lat)) / (np.cos(lat) * np.cos(lon))) - dec = np.arcsin(cos_ec * np.sin(lat) + sin_ec * np.cos(lat) * np.sin(lon)) - if degrees: - return np.degrees(ra), np.degrees(dec) - else: - return ra, dec - - -def proper_motion_ra_dec(r=None, v=None, x=None, y=None, z=None, vx=None, vy=None, vz=None, r_earth=np.array([0, 0, 0]), v_earth=np.array([0, 0, 0]), input_unit='si'): - if r is None or v is None: - if x is not None and y is not None and z is not None and vx is not None and vy is not None and vz is not None: - r = np.array([x, y, z]) - v = np.array([vx, vy, vz]) - else: - raise ValueError("Either provide r and v arrays or individual coordinates (x, y, z) and velocities (vx, vy, vz)") - - # Subtract Earth's position and velocity from the input arrays - r = r - r_earth - v = v - v_earth - - # Distances to Earth and Sun - d_earth_mag = einsum_norm(r, 'ij,ij->i') - - # RA / DEC calculation - ra = rad0to2pi(np.arctan2(r[:, 1], r[:, 0])) # in radians - dec = np.arcsin(r[:, 2] / d_earth_mag) - ra_unit_vector = np.array([-np.sin(ra), np.cos(ra), np.zeros(np.shape(ra))]).T - dec_unit_vector = -np.array([np.cos(np.pi / 2 - dec) * np.cos(ra), np.cos(np.pi / 2 - dec) * np.sin(ra), -np.sin(np.pi / 2 - dec)]).T - pmra = (np.einsum('ij,ij->i', v, ra_unit_vector)) / d_earth_mag * 206265 # arcseconds / second - pmdec = (np.einsum('ij,ij->i', v, dec_unit_vector)) / d_earth_mag * 206265 # arcseconds / second - - if input_unit == 'si': - return pmra, pmdec - elif input_unit == 'rebound': - pmra = pmra / (31557600 * 2 * np.pi) - pmdec = pmdec / (31557600 * 2 * np.pi) # arcseconds * (au/sim_time)/au, convert to arcseconds / second - return pmra, pmdec - else: - print('Error - units provided not available, provide either SI or rebound units.') - return - - - -def gcrf_to_lunar(r, t, v=None): - class MoonRotator: - def __init__(self): - self.mpm = MoonPosition() - - def __call__(self, r, t): - rmoon = self.mpm(t) - vmoon = (self.mpm(t + 5.0) - self.mpm(t - 5.0)) / 10. - xhat = normed(rmoon.T).T - vpar = np.einsum("ab,ab->b", xhat, vmoon) * xhat - vperp = vmoon - vpar - yhat = normed(vperp.T).T - zhat = np.cross(xhat, yhat, axisa=0, axisb=0).T - R = np.empty((3, 3, len(t))) - R[0] = xhat - R[1] = yhat - R[2] = zhat - return np.einsum("abc,cb->ca", R, r) - rotator = MoonRotator() - if v is None: - return rotator(r, t) - else: - r_lunar = rotator(r, t) - v_lunar = v_from_r(r_lunar, t) - return r_lunar, v_lunar - - -def gcrf_to_lunar_fixed(r, t, v=None): - r_lunar = gcrf_to_lunar(r, t) - gcrf_to_lunar(get_body('moon').position(t).T, t) - if v is None: - return r_lunar - else: - v = v_from_r(r_lunar, t) - return r_lunar, v - - -def gcrf_to_radec(gcrf_coords): - x, y, z = gcrf_coords - # Calculate right ascension in radians - ra = np.arctan2(y, x) - # Convert right ascension to degrees - ra_deg = np.degrees(ra) - # Normalize right ascension to the range [0, 360) - ra_deg = ra_deg % 360 - # Calculate declination in radians - dec_rad = np.arctan2(z, np.sqrt(x**2 + y**2)) - # Convert declination to degrees - dec_deg = np.degrees(dec_rad) - return (ra_deg, dec_deg) - - -def gcrf_to_ecef_bad(r_gcrf, t): - if isinstance(t, Time): - t = t.gps - r_gcrf = np.atleast_2d(r_gcrf) - rotation_angles = WGS84_EARTH_OMEGA * (t - Time("1980-3-20T11:06:00", format='isot').gps) - cos_thetas = np.cos(rotation_angles) - sin_thetas = np.sin(rotation_angles) - - # Create an array of 3x3 rotation matrices - Rz = np.array([[cos_thetas, -sin_thetas, np.zeros_like(cos_thetas)], - [sin_thetas, cos_thetas, np.zeros_like(cos_thetas)], - [np.zeros_like(cos_thetas), np.zeros_like(cos_thetas), np.ones_like(cos_thetas)]]).T - - # Apply the rotation matrices to all rows of r_gcrf simultaneously - r_ecef = np.einsum('ijk,ik->ij', Rz, r_gcrf) - return r_ecef - - -def gcrf_to_lat_lon(r, t): - lon, lat, height = groundTrack(r, t) - return lon, lat, height - - -def gcrf_to_itrf(r_gcrf, t, v=None): - x, y, z = groundTrack(r_gcrf, t, format='cartesian') - _ = np.array([x, y, z]).T - if v is None: - return _ - else: - return _, v_from_r(_, t) - - -def gcrf_to_sim_geo(r_gcrf, t, h=10): - if np.min(np.diff(t.gps)) < h: - h = np.min(np.diff(t.gps)) - r_gcrf = np.atleast_2d(r_gcrf) - r_geo, v_geo = rv(Orbit.fromKeplerianElements(*[RGEO, 0, 0, 0, 0, 0], t=t[0]), t, propagator=RK78Propagator(AccelKepler(), h=h)) - angle_geo_to_x = np.arctan2(r_geo[:, 1], r_geo[:, 0]) - c = np.cos(angle_geo_to_x) - s = np.sin(angle_geo_to_x) - rotation = np.array([[c, -s, np.zeros_like(c)], [s, c, np.zeros_like(c)], [np.zeros_like(c), np.zeros_like(c), np.ones_like(c)]]).T - return np.einsum('ijk,ik->ij', rotation, r_gcrf) - - -# Function still in development, not 100% accurate. -def gcrf_to_itrf_astropy(state_vectors, t): - import astropy.units as u - from astropy.coordinates import GCRS, ITRS, SkyCoord, get_body_barycentric, solar_system_ephemeris, ICRS - - sc = SkyCoord(x=state_vectors[:, 0] * u.m, y=state_vectors[:, 1] * u.m, z=state_vectors[:, 2] * u.m, representation_type='cartesian', frame=GCRS(obstime=t)) - sc_itrs = sc.transform_to(ITRS(obstime=t)) - with solar_system_ephemeris.set('de430'): # other options: builtin, de432s - earth = get_body_barycentric('earth', t) - earth_center_itrs = SkyCoord(earth.x, earth.y, earth.z, representation_type='cartesian', frame=ICRS()).transform_to(ITRS(obstime=t)) - itrs_coords = SkyCoord( - sc_itrs.x.value - earth_center_itrs.x.to_value(u.m), - sc_itrs.y.value - earth_center_itrs.y.to_value(u.m), - sc_itrs.z.value - earth_center_itrs.z.to_value(u.m), - representation_type='cartesian', - frame=ITRS(obstime=t) - ) - # Extract Cartesian coordinates and convert to meters - itrs_coords_meters = np.array([itrs_coords.x, - itrs_coords.y, - itrs_coords.z]).T - return itrs_coords_meters - - -def v_from_r(r, t): - if isinstance(t[0], Time): - t = t.gps - delta_r = np.diff(r, axis=0) - delta_t = np.diff(t) - v = delta_r / delta_t[:, np.newaxis] - v = np.vstack((v, v[-1])) - return v From 2705c11e3ae8a848d07745f9eb9655e0c59b8a71 Mon Sep 17 00:00:00 2001 From: SuperdoerTrav Date: Fri, 19 Jul 2024 13:01:27 -0700 Subject: [PATCH 15/55] added many more coordinate conversions useful for plotting --- ssapy/utils.py | 637 ++++++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 631 insertions(+), 6 deletions(-) diff --git a/ssapy/utils.py b/ssapy/utils.py index 14801e3..3cd5692 100644 --- a/ssapy/utils.py +++ b/ssapy/utils.py @@ -5,12 +5,7 @@ import astropy.units as u from . import datadir -from .accel import AccelKepler -from .body import get_body, MoonPosition -from .constants import RGEO, MOON_RADIUS, WGS84_EARTH_OMEGA -from .compute import groundTrack, rv -from .orbit import Orbit -from .propagator import RK78Propagator +from .constants import MOON_RADIUS try: import erfa @@ -1162,6 +1157,636 @@ def teme_to_gcrf(t): return erfa.tr(gcrf_to_teme(t)) + +# VECTOR FUNCTIONS FOR COORDINATE MATH +def unit_vector(vector): + """ Returns the unit vector of the vector.""" + return vector / np.linalg.norm(vector) + + +def getAngle(a, b, c): # a,b,c where b is the vertex + a = np.atleast_2d(a) + b = np.atleast_2d(b) + c = np.atleast_2d(c) + ba = np.subtract(a, b) + bc = np.subtract(c, b) + cosine_angle = np.sum(ba * bc, axis=-1) / (np.linalg.norm(ba, axis=-1) * np.linalg.norm(bc, axis=-1)) + return np.arccos(cosine_angle) + + +def angle_between_vectors(vector1, vector2): + return np.arccos(np.clip(np.dot(vector1, vector2) / (np.linalg.norm(vector1) * np.linalg.norm(vector2)), -1.0, 1.0)) + + +def rotation_matrix_from_vectors(vec1, vec2): + """ Find the rotation matrix that aligns vec1 to vec2 + :param vec1: A 3d "source" vector + :param vec2: A 3d "destination" vector + :return mat: A transform matrix (3x3) which when applied to vec1, aligns it with vec2. + """ + a, b = (vec1 / np.linalg.norm(vec1)).reshape(3), (vec2 / np.linalg.norm(vec2)).reshape(3) + v = np.cross(a, b) + c = np.dot(a, b) + s = np.linalg.norm(v) + kmat = np.array([[0, -v[2], v[1]], [v[2], 0, -v[0]], [-v[1], v[0], 0]]) + rotation_matrix = np.eye(3) + kmat + kmat.dot(kmat) * ((1 - c) / (s**2)) + return rotation_matrix + + +def rotate_vector(v_unit, theta, phi): + v_unit = v_unit / np.linalg.norm(v_unit, axis=-1) + if np.all(np.abs(v_unit) != np.max(np.abs(v_unit))): + perp_vector = np.cross(v_unit, np.array([1, 0, 0])) + else: + perp_vector = np.cross(v_unit, np.array([0, 1, 0])) + perp_vector /= np.linalg.norm(perp_vector) + + theta = np.radians(theta) + phi = np.radians(phi) + cos_theta = np.cos(theta) + sin_theta = np.sin(theta) + cos_phi = np.cos(phi) + sin_phi = np.sin(phi) + + R1 = np.array([ + [cos_theta + (1 - cos_theta) * perp_vector[0]**2, + (1 - cos_theta) * perp_vector[0] * perp_vector[1] - sin_theta * perp_vector[2], + (1 - cos_theta) * perp_vector[0] * perp_vector[2] + sin_theta * perp_vector[1]], + [(1 - cos_theta) * perp_vector[1] * perp_vector[0] + sin_theta * perp_vector[2], + cos_theta + (1 - cos_theta) * perp_vector[1]**2, + (1 - cos_theta) * perp_vector[1] * perp_vector[2] - sin_theta * perp_vector[0]], + [(1 - cos_theta) * perp_vector[2] * perp_vector[0] - sin_theta * perp_vector[1], + (1 - cos_theta) * perp_vector[2] * perp_vector[1] + sin_theta * perp_vector[0], + cos_theta + (1 - cos_theta) * perp_vector[2]**2] + ]) + + # Apply the rotation matrix to v_unit to get the rotated unit vector + v1 = np.dot(R1, v_unit) + + # Rotation matrix for rotation about v_unit + R2 = np.array([[cos_phi + (1 - cos_phi) * v_unit[0]**2, + (1 - cos_phi) * v_unit[0] * v_unit[1] - sin_phi * v_unit[2], + (1 - cos_phi) * v_unit[0] * v_unit[2] + sin_phi * v_unit[1]], + [(1 - cos_phi) * v_unit[1] * v_unit[0] + sin_phi * v_unit[2], + cos_phi + (1 - cos_phi) * v_unit[1]**2, + (1 - cos_phi) * v_unit[1] * v_unit[2] - sin_phi * v_unit[0]], + [(1 - cos_phi) * v_unit[2] * v_unit[0] - sin_phi * v_unit[1], + (1 - cos_phi) * v_unit[2] * v_unit[1] + sin_phi * v_unit[0], + cos_phi + (1 - cos_phi) * v_unit[2]**2]]) + + v2 = np.dot(R2, v1) + return v2 / np.linalg.norm(v2, axis=-1) + + +def rotate_points_3d(points, axis=np.array([0, 0, 1]), theta=-np.pi / 2): + """ + Rotate a set of 3D points about a 3D axis by an angle theta in radians. + + Args: + points (np.ndarray): The set of 3D points to rotate, as an Nx3 array. + axis (np.ndarray): The 3D axis to rotate about, as a length-3 array. Default is the z-axis. + theta (float): The angle to rotate by, in radians. Default is pi/2. + + Returns: + np.ndarray: The rotated set of 3D points, as an Nx3 array. + """ + # Normalize the axis to be a unit vector + axis = axis / np.linalg.norm(axis) + + # Compute the quaternion representing the rotation + qw = np.cos(theta / 2) + qx, qy, qz = axis * np.sin(theta / 2) + + # Construct the rotation matrix from the quaternion + R = np.array([ + [1 - 2 * qy**2 - 2 * qz**2, 2 * qx * qy - 2 * qz * qw, 2 * qx * qz + 2 * qy * qw], + [2 * qx * qy + 2 * qz * qw, 1 - 2 * qx**2 - 2 * qz**2, 2 * qy * qz - 2 * qx * qw], + [2 * qx * qz - 2 * qy * qw, 2 * qy * qz + 2 * qx * qw, 1 - 2 * qx**2 - 2 * qy**2] + ]) + + # Apply the rotation matrix to the set of points + rotated_points = np.dot(R, points.T).T + + return rotated_points + + +def perpendicular_vectors(v): + """Returns two vectors that are perpendicular to v and each other.""" + # Check if v is the zero vector + if np.allclose(v, np.zeros_like(v)): + raise ValueError("Input vector cannot be the zero vector.") + + # Choose an arbitrary non-zero vector w that is not parallel to v + w = np.array([1., 0., 0.]) + if np.allclose(v, w) or np.allclose(v, -w): + w = np.array([0., 1., 0.]) + u = np.cross(v, w) + if np.allclose(u, np.zeros_like(u)): + w = np.array([0., 0., 1.]) + u = np.cross(v, w) + w = np.cross(v, u) + + return u, w + + +def points_on_circle(r, v, rad, num_points=4): + # Convert inputs to numpy arrays + r = np.array(r) + v = np.array(v) + + # Find the perpendicular vectors to the given vector v + if np.all(v[:2] == 0): + if np.all(v[2] == 0): + raise ValueError("The given vector v must not be the zero vector.") + else: + u = np.array([1, 0, 0]) + else: + u = np.array([-v[1], v[0], 0]) + u = u / np.linalg.norm(u) + w = np.cross(u, v) + w_norm = np.linalg.norm(w) + if w_norm < 1e-15: + # v is parallel to z-axis + w = np.array([0, 1, 0]) + else: + w = w / w_norm + # Generate a sequence of angles for equally spaced points + angles = np.linspace(0, 2 * np.pi, num_points, endpoint=False) + + # Compute the x, y, z coordinates of each point on the circle + x = rad * np.cos(angles) * u[0] + rad * np.sin(angles) * w[0] + y = rad * np.cos(angles) * u[1] + rad * np.sin(angles) * w[1] + z = rad * np.cos(angles) * u[2] + rad * np.sin(angles) * w[2] + + # Apply rotation about z-axis by 90 degrees + rot_matrix = np.array([[0, 1, 0], [-1, 0, 0], [0, 0, 1]]) + rotated_points = np.dot(rot_matrix, np.column_stack((x, y, z)).T).T + + # Translate the rotated points to the center point r + points_rotated = rotated_points + r.reshape(1, 3) + + return points_rotated + + +def dms_to_rad(coords): + from astropy.coordinates import Angle + if isinstance(coords, (list, tuple)): + return [Angle(coord).radian for coord in coords] + else: + return Angle(coords).radian + return + + +def dms_to_deg(coords): + from astropy.coordinates import Angle + if isinstance(coords, (list, tuple)): + return [Angle(coord).deg for coord in coords] + else: + return Angle(coords).deg + return + + +def rad0to2pi(angles): + return (2 * np.pi + angles) * (angles < 0) + angles * (angles > 0) + + +def deg0to360(array_): + try: + return [i % 360 for i in array_] + except TypeError: + return array_ % 360 + + +def deg0to360array(array_): + return [i % 360 for i in array_] + + +def deg90to90(val_in): + if hasattr(val_in, "__len__"): + val_out = [] + for i, v in enumerate(val_in): + while v < -90: + v += 90 + while v > 90: + v -= 90 + val_out.append(v) + else: + while val_in < -90: + val_in += 90 + while val_in > 90: + val_in -= 90 + val_out = val_in + return val_out + + +def deg90to90array(array_): + return [i % 90 for i in array_] + + +def cart2sph_deg(x, y, z): + hxy = np.hypot(x, y) + r = np.hypot(hxy, z) + el = np.arctan2(z, hxy) * (180 / np.pi) + az = (np.arctan2(y, x)) * (180 / np.pi) + return az, el, r + + +def cart_to_cyl(x, y, z): + r = np.linalg.norm([x, y]) + theta = np.arctan2(y, x) + return r, theta, z + + +def inert2rot(x, y, xe, ye, xs=0, ys=0): # Places Earth at (-1,0) + earth_theta = np.arctan2(ye - ys, xe - xs) + theta = np.arctan2(y - ys, x - xs) + distance = np.sqrt(np.power((x - xs), 2) + np.power((y - ys), 2)) + xrot = distance * np.cos(np.pi + (theta - earth_theta)) + yrot = distance * np.sin(np.pi + (theta - earth_theta)) + return xrot, yrot + + +def sim_lonlatrad(x, y, z, xe, ye, ze, xs, ys, zs): + # convert all to geo coordinates + x = x - xe + y = y - ye + z = z - ze + xs = xs - xe + ys = ys - ye + zs = zs - ze + # convert x y z to lon lat radius + longitude, latitude, radius = cart2sph_deg(x, y, z) + slongitude, slatitude, sradius = cart2sph_deg(xs, ys, zs) + # correct so that Sun is at (0,0) + longitude = deg0to360(slongitude - longitude) + latitude = latitude - slatitude + return longitude, latitude, radius + + +def sun_ra_dec(time_): + out = get_body(Time(time_, format='mjd')) + return out.ra.to('rad').value, out.dec.to('rad').value + + +def ra_dec(r=None, v=None, x=None, y=None, z=None, vx=None, vy=None, vz=None, r_earth=np.array([0, 0, 0]), v_earth=np.array([0, 0, 0]), input_unit='si'): + if r is None or v is None: + if x is not None and y is not None and z is not None and vx is not None and vy is not None and vz is not None: + r = np.array([x, y, z]) + v = np.array([vx, vy, vz]) + else: + raise ValueError("Either provide r and v arrays or individual coordinates (x, y, z) and velocities (vx, vy, vz)") + + # Subtract Earth's position and velocity from the input arrays + r = r - r_earth + v = v - v_earth + + d_earth_mag = einsum_norm(r, 'ij,ij->i') + ra = rad0to2pi(np.arctan2(r[:, 1], r[:, 0])) # in radians + dec = np.arcsin(r[:, 2] / d_earth_mag) + return ra, dec + + +def lonlat_distance(lat1, lat2, lon1, lon2): + # Haversine formula + dlon = lon2 - lon1 + dlat = lat2 - lat1 + a = np.sin(dlat / 2)**2 + np.cos(lat1) * np.cos(lat2) * np.sin(dlon / 2)**2 + c = 2 * np.arcsin(np.sqrt(a)) + # Radius of earth in kilometers. Use 3956 for miles + # calculate the result + return (c * EARTH_RADIUS) + + +def altitude2zenithangle(altitude, deg=True): + if deg: + out = 90 - altitude + else: + out = np.pi / 2 - altitude + return out + + +def zenithangle2altitude(zenith_angle, deg=True): + if deg: + out = 90 - zenith_angle + else: + out = np.pi / 2 - zenith_angle + return out + + +def rightasension2hourangle(right_ascension, local_time): + if type(right_ascension) is not str: + right_ascension = dd_to_hms(right_ascension) + if type(local_time) is not str: + local_time = dd_to_dms(local_time) + _ra = float(right_ascension.split(':')[0]) + _lt = float(local_time.split(':')[0]) + if _ra > _lt: + __ltm, __lts = local_time.split(':')[1:] + local_time = f'{24 + _lt}:{__ltm}:{__lts}' + + return dd_to_dms(hms_to_dd(local_time) - hms_to_dd(right_ascension)) + + +def equatorial_to_horizontal(observer_latitude, declination, right_ascension=None, hour_angle=None, local_time=None, hms=False): + if right_ascension is not None: + hour_angle = rightasension2hourangle(right_ascension, local_time) + if hms: + hour_angle = hms_to_dd(hour_angle) + elif hour_angle is not None: + if hms: + hour_angle = hms_to_dd(hour_angle) + elif right_ascension is not None and hour_angle is not None: + print('Both right_ascension and hour_angle parameters are provided.\nUsing hour_angle for calculations.') + if hms: + hour_angle = hms_to_dd(hour_angle) + else: + print('Either right_ascension or hour_angle must be provided.') + + observer_latitude, hour_angle, declination = np.radians([observer_latitude, hour_angle, declination]) + + zenith_angle = np.arccos(np.sin(observer_latitude) * np.sin(declination) + np.cos(observer_latitude) * np.cos(declination) * np.cos(hour_angle)) + + altitude = zenithangle2altitude(zenith_angle, deg=False) + + _num = np.sin(declination) - np.sin(observer_latitude) * np.cos(zenith_angle) + _den = np.cos(observer_latitude) * np.sin(zenith_angle) + azimuth = np.arccos(_num / _den) + + if observer_latitude < 0: + azimuth = np.pi - azimuth + altitude, azimuth = np.degrees([altitude, azimuth]) + + return azimuth, altitude + + +def horizontal_to_equatorial(observer_latitude, azimuth, altitude): + altitude, azimuth, latitude = np.radians([altitude, azimuth, observer_latitude]) + zenith_angle = zenithangle2altitude(altitude) + + zenith_angle = [-zenith_angle if latitude < 0 else zenith_angle][0] + + declination = np.sin(latitude) * np.cos(zenith_angle) + declination = declination + (np.cos(latitude) * np.sin(zenith_angle) * np.cos(azimuth)) + declination = np.arcsin(declination) + + _num = np.cos(zenith_angle) - np.sin(latitude) * np.sin(declination) + _den = np.cos(latitude) * np.cos(declination) + hour_angle = np.arccos(_num / _den) + + if (latitude > 0 > declination) or (latitude < 0 < declination): + hour_angle = 2 * np.pi - hour_angle + + declination, hour_angle = np.degrees([declination, hour_angle]) + + return hour_angle, declination + + +_ecliptic = 0.409092601 # np.radians(23.43927944) +cos_ec = 0.9174821430960974 +sin_ec = 0.3977769690414367 + + +def equatorial_xyz_to_ecliptic_xyz(xq, yq, zq): + xc = xq + yc = cos_ec * yq + sin_ec * zq + zc = -sin_ec * yq + cos_ec * zq + return xc, yc, zc + + +def ecliptic_xyz_to_equatorial_xyz(xc, yc, zc): + xq = xc + yq = cos_ec * yc - sin_ec * zc + zq = sin_ec * yc + cos_ec * zc + return xq, yq, zq + + +def xyz_to_ecliptic(xc, yc, zc, xe=0, ye=0, ze=0, degrees=False): + x_ast_to_earth = xc - xe + y_ast_to_earth = yc - ye + z_ast_to_earth = zc - ze + d_earth_mag = np.sqrt(np.power(x_ast_to_earth, 2) + np.power(y_ast_to_earth, 2) + np.power(z_ast_to_earth, 2)) + ec_longitude = rad0to2pi(np.arctan2(y_ast_to_earth, x_ast_to_earth)) # in radians + ec_latitude = np.arcsin(z_ast_to_earth / d_earth_mag) + if degrees: + return np.degrees(ec_longitude), np.degrees(ec_latitude) + else: + return ec_longitude, ec_latitude + + +def xyz_to_equatorial(xq, yq, zq, xe=0, ye=0, ze=0, degrees=False): + # RA / DEC calculation - assumes XY plane to be celestial equator, and -x axis to be vernal equinox + x_ast_to_earth = xq - xe + y_ast_to_earth = yq - ye + z_ast_to_earth = zq - ze + d_earth_mag = np.sqrt(np.power(x_ast_to_earth, 2) + np.power(y_ast_to_earth, 2) + np.power(z_ast_to_earth, 2)) + ra = rad0to2pi(np.arctan2(y_ast_to_earth, x_ast_to_earth)) # in radians + dec = np.arcsin(z_ast_to_earth / d_earth_mag) + if degrees: + return np.degrees(ra), np.degrees(dec) + else: + return ra, dec + + +def ecliptic_xyz_to_equatorial(xc, yc, zc, xe=0, ye=0, ze=0, degrees=False): + # Convert ecliptic cartesian into equitorial cartesian + x_ast_to_earth, y_ast_to_earth, z_ast_to_earth = ecliptic_xyz_to_equatorial_xyz(xc - xe, yc - ye, zc - ze) + d_earth_mag = np.sqrt(np.power(x_ast_to_earth, 2) + np.power(y_ast_to_earth, 2) + np.power(z_ast_to_earth, 2)) + ra = rad0to2pi(np.arctan2(y_ast_to_earth, x_ast_to_earth)) # in radians + dec = np.arcsin(z_ast_to_earth / d_earth_mag) + if degrees: + return np.degrees(ra), np.degrees(dec) + else: + return ra, dec + + +def equatorial_to_ecliptic(right_ascension, declination, degrees=False): + ra, dec = np.radians(right_ascension), np.radians(declination) + ec_latitude = np.arcsin(cos_ec * np.sin(dec) - sin_ec * np.cos(dec) * np.sin(ra)) + ec_longitude = np.arctan((cos_ec * np.cos(dec) * np.sin(ra) + sin_ec * np.sin(dec)) / (np.cos(dec) * np.cos(ra))) + if degrees: + return deg0to360(np.degrees(ec_longitude)), np.degrees(ec_latitude) + else: + return rad0to2pi(ec_longitude), ec_latitude + + +def ecliptic_to_equatorial(lon, lat, degrees=False): + lon, lat = np.radians(lon), np.radians(lat) + ra = np.arctan((cos_ec * np.cos(lat) * np.sin(lon) - sin_ec * np.sin(lat)) / (np.cos(lat) * np.cos(lon))) + dec = np.arcsin(cos_ec * np.sin(lat) + sin_ec * np.cos(lat) * np.sin(lon)) + if degrees: + return np.degrees(ra), np.degrees(dec) + else: + return ra, dec + + +def proper_motion_ra_dec(r=None, v=None, x=None, y=None, z=None, vx=None, vy=None, vz=None, r_earth=np.array([0, 0, 0]), v_earth=np.array([0, 0, 0]), input_unit='si'): + if r is None or v is None: + if x is not None and y is not None and z is not None and vx is not None and vy is not None and vz is not None: + r = np.array([x, y, z]) + v = np.array([vx, vy, vz]) + else: + raise ValueError("Either provide r and v arrays or individual coordinates (x, y, z) and velocities (vx, vy, vz)") + + # Subtract Earth's position and velocity from the input arrays + r = r - r_earth + v = v - v_earth + + # Distances to Earth and Sun + d_earth_mag = einsum_norm(r, 'ij,ij->i') + + # RA / DEC calculation + ra = rad0to2pi(np.arctan2(r[:, 1], r[:, 0])) # in radians + dec = np.arcsin(r[:, 2] / d_earth_mag) + ra_unit_vector = np.array([-np.sin(ra), np.cos(ra), np.zeros(np.shape(ra))]).T + dec_unit_vector = -np.array([np.cos(np.pi / 2 - dec) * np.cos(ra), np.cos(np.pi / 2 - dec) * np.sin(ra), -np.sin(np.pi / 2 - dec)]).T + pmra = (np.einsum('ij,ij->i', v, ra_unit_vector)) / d_earth_mag * 206265 # arcseconds / second + pmdec = (np.einsum('ij,ij->i', v, dec_unit_vector)) / d_earth_mag * 206265 # arcseconds / second + + if input_unit == 'si': + return pmra, pmdec + elif input_unit == 'rebound': + pmra = pmra / (31557600 * 2 * np.pi) + pmdec = pmdec / (31557600 * 2 * np.pi) # arcseconds * (au/sim_time)/au, convert to arcseconds / second + return pmra, pmdec + else: + print('Error - units provided not available, provide either SI or rebound units.') + return + + + +def gcrf_to_lunar(r, t, v=None): + class MoonRotator: + def __init__(self): + self.mpm = MoonPosition() + + def __call__(self, r, t): + rmoon = self.mpm(t) + vmoon = (self.mpm(t + 5.0) - self.mpm(t - 5.0)) / 10. + xhat = normed(rmoon.T).T + vpar = np.einsum("ab,ab->b", xhat, vmoon) * xhat + vperp = vmoon - vpar + yhat = normed(vperp.T).T + zhat = np.cross(xhat, yhat, axisa=0, axisb=0).T + R = np.empty((3, 3, len(t))) + R[0] = xhat + R[1] = yhat + R[2] = zhat + return np.einsum("abc,cb->ca", R, r) + rotator = MoonRotator() + if v is None: + return rotator(r, t) + else: + r_lunar = rotator(r, t) + v_lunar = v_from_r(r_lunar, t) + return r_lunar, v_lunar + + +def gcrf_to_lunar_fixed(r, t, v=None): + r_lunar = gcrf_to_lunar(r, t) - gcrf_to_lunar(get_body('moon').position(t).T, t) + if v is None: + return r_lunar + else: + v = v_from_r(r_lunar, t) + return r_lunar, v + + +def gcrf_to_radec(gcrf_coords): + x, y, z = gcrf_coords + # Calculate right ascension in radians + ra = np.arctan2(y, x) + # Convert right ascension to degrees + ra_deg = np.degrees(ra) + # Normalize right ascension to the range [0, 360) + ra_deg = ra_deg % 360 + # Calculate declination in radians + dec_rad = np.arctan2(z, np.sqrt(x**2 + y**2)) + # Convert declination to degrees + dec_deg = np.degrees(dec_rad) + return (ra_deg, dec_deg) + + +def gcrf_to_ecef_bad(r_gcrf, t): + if isinstance(t, Time): + t = t.gps + r_gcrf = np.atleast_2d(r_gcrf) + rotation_angles = WGS84_EARTH_OMEGA * (t - Time("1980-3-20T11:06:00", format='isot').gps) + cos_thetas = np.cos(rotation_angles) + sin_thetas = np.sin(rotation_angles) + + # Create an array of 3x3 rotation matrices + Rz = np.array([[cos_thetas, -sin_thetas, np.zeros_like(cos_thetas)], + [sin_thetas, cos_thetas, np.zeros_like(cos_thetas)], + [np.zeros_like(cos_thetas), np.zeros_like(cos_thetas), np.ones_like(cos_thetas)]]).T + + # Apply the rotation matrices to all rows of r_gcrf simultaneously + r_ecef = np.einsum('ijk,ik->ij', Rz, r_gcrf) + return r_ecef + + +def gcrf_to_lat_lon(r, t): + from .compute import groundTrack + lon, lat, height = groundTrack(r, t) + return lon, lat, height + + +def gcrf_to_itrf(r_gcrf, t, v=None): + from .compute import groundTrack + x, y, z = groundTrack(r_gcrf, t, format='cartesian') + _ = np.array([x, y, z]).T + if v is None: + return _ + else: + return _, v_from_r(_, t) + + +def gcrf_to_sim_geo(r_gcrf, t, h=10): + from .accel import AccelKepler + from .compute import rv + if np.min(np.diff(t.gps)) < h: + h = np.min(np.diff(t.gps)) + r_gcrf = np.atleast_2d(r_gcrf) + r_geo, v_geo = rv(Orbit.fromKeplerianElements(*[RGEO, 0, 0, 0, 0, 0], t=t[0]), t, propagator=RK78Propagator(AccelKepler(), h=h)) + angle_geo_to_x = np.arctan2(r_geo[:, 1], r_geo[:, 0]) + c = np.cos(angle_geo_to_x) + s = np.sin(angle_geo_to_x) + rotation = np.array([[c, -s, np.zeros_like(c)], [s, c, np.zeros_like(c)], [np.zeros_like(c), np.zeros_like(c), np.ones_like(c)]]).T + return np.einsum('ijk,ik->ij', rotation, r_gcrf) + + +# Function still in development, not 100% accurate. +def gcrf_to_itrf_astropy(state_vectors, t): + import astropy.units as u + from astropy.coordinates import GCRS, ITRS, SkyCoord, get_body_barycentric, solar_system_ephemeris, ICRS + + sc = SkyCoord(x=state_vectors[:, 0] * u.m, y=state_vectors[:, 1] * u.m, z=state_vectors[:, 2] * u.m, representation_type='cartesian', frame=GCRS(obstime=t)) + sc_itrs = sc.transform_to(ITRS(obstime=t)) + with solar_system_ephemeris.set('de430'): # other options: builtin, de432s + earth = get_body_barycentric('earth', t) + earth_center_itrs = SkyCoord(earth.x, earth.y, earth.z, representation_type='cartesian', frame=ICRS()).transform_to(ITRS(obstime=t)) + itrs_coords = SkyCoord( + sc_itrs.x.value - earth_center_itrs.x.to_value(u.m), + sc_itrs.y.value - earth_center_itrs.y.to_value(u.m), + sc_itrs.z.value - earth_center_itrs.z.to_value(u.m), + representation_type='cartesian', + frame=ITRS(obstime=t) + ) + # Extract Cartesian coordinates and convert to meters + itrs_coords_meters = np.array([itrs_coords.x, + itrs_coords.y, + itrs_coords.z]).T + return itrs_coords_meters + + +def v_from_r(r, t): + if isinstance(t[0], Time): + t = t.gps + delta_r = np.diff(r, axis=0) + delta_t = np.diff(t) + v = delta_r / delta_t[:, np.newaxis] + v = np.vstack((v, v[-1])) + return v + + # Stolen from https://github.com/lsst/utils/blob/main/python/lsst/utils/wrappers.py INTRINSIC_SPECIAL_ATTRIBUTES = frozenset( ( From 2bf8b8cf4e06a9139b67ebd736bd78b74999a0b8 Mon Sep 17 00:00:00 2001 From: SuperdoerTrav Date: Fri, 19 Jul 2024 13:01:59 -0700 Subject: [PATCH 16/55] added a gif writing function --- ssapy/plotUtils.py | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/ssapy/plotUtils.py b/ssapy/plotUtils.py index 97c6f68..4e88790 100644 --- a/ssapy/plotUtils.py +++ b/ssapy/plotUtils.py @@ -1,8 +1,7 @@ from .body import get_body from .compute import groundTrack, lagrange_points_lunar_frame, calculate_orbital_elements from .constants import RGEO, LD, EARTH_RADIUS, MOON_RADIUS, EARTH_MU, MOON_MU -from .coordinates import gcrf_to_itrf, gcrf_to_lunar_fixed, gcrf_to_lunar -from .utils import find_file, Time, find_smallest_bounding_cube +from .utils import find_file, Time, find_smallest_bounding_cube, gcrf_to_itrf, gcrf_to_lunar_fixed, gcrf_to_lunar import numpy as np import os @@ -11,7 +10,7 @@ import matplotlib.pyplot as plt import matplotlib.cm as cm -from matplotlib import mplcolors +from matplotlib import colors as mplcolors from PyPDF2 import PdfMerger from matplotlib.backends.backend_pdf import PdfPages @@ -849,3 +848,14 @@ def save_plot(figure, save_path, dpi=200): print(f"Figure saved at: {save_path}") except Exception as e: print(f"Error occurred while saving the figure: {e}") + + +def write_gif(gif_name, frames, fps=30): + import imageio + print(f'Writing gif: {gif_name}') + with imageio.get_writer(gif_name, mode='I', duration=1 / fps) as writer: + for i, filename in enumerate(frames): + image = imageio.imread(filename) + writer.append_data(image) + print(f'Wrote {gif_name}') + return From 0d3751b7d78ee1fdc523cfa618d4ccdfbcd4e005 Mon Sep 17 00:00:00 2001 From: SuperdoerTrav Date: Fri, 19 Jul 2024 13:25:55 -0700 Subject: [PATCH 17/55] added built in plot to rotate_vector - used for validation --- ssapy/utils.py | 27 ++++++++++++++++++++++++++- 1 file changed, 26 insertions(+), 1 deletion(-) diff --git a/ssapy/utils.py b/ssapy/utils.py index 3cd5692..92b6538 100644 --- a/ssapy/utils.py +++ b/ssapy/utils.py @@ -1193,7 +1193,7 @@ def rotation_matrix_from_vectors(vec1, vec2): return rotation_matrix -def rotate_vector(v_unit, theta, phi): +def rotate_vector(v_unit, theta, phi, plot=False, save_idx=False): v_unit = v_unit / np.linalg.norm(v_unit, axis=-1) if np.all(np.abs(v_unit) != np.max(np.abs(v_unit))): perp_vector = np.cross(v_unit, np.array([1, 0, 0])) @@ -1235,6 +1235,31 @@ def rotate_vector(v_unit, theta, phi): cos_phi + (1 - cos_phi) * v_unit[2]**2]]) v2 = np.dot(R2, v1) + + if plot: + import matplotlib.pyplot as plt + plt.rcParams.update({'font.size': 9, 'figure.facecolor': 'black'}) + fig = plt.figure() + ax = fig.add_subplot(111, projection='3d') + ax.quiver(0, 0, 0, v_unit[0], v_unit[1], v_unit[2], color='b') + ax.quiver(0, 0, 0, v1[0], v1[1], v1[2], color='g') + ax.quiver(0, 0, 0, v2[0], v2[1], v2[2], color='r') + ax.set_xlabel('X', color='white') + ax.set_ylabel('Y', color='white') + ax.set_zlabel('Z', color='white') + ax.set_facecolor('black') # Set plot background color to black + ax.tick_params(axis='x', colors='white') # Set x-axis tick color to white + ax.tick_params(axis='y', colors='white') # Set y-axis tick color to white + ax.tick_params(axis='z', colors='white') # Set z-axis tick color to white + ax.set_title('Vector Plot', color='white') + ax.set_xlim(-1, 1) + ax.set_ylim(-1, 1) + ax.set_zlim(-1, 1) + plt.grid(True) + if save_idx is not False: + from .plotUtils import save_plot + ax.set_title(f'Vector Plot\ntheta: {np.degrees(theta):.0f}, phi: {np.degrees(phi):.0f}', color='white') + save_plot(fig, f'/p/lustre1/yeager7/plots_gif/rotate_vector_frames/{save_idx}.png') return v2 / np.linalg.norm(v2, axis=-1) From a61ceb893e0a1bb6e2ee8aa435002d6f3e2f2075 Mon Sep 17 00:00:00 2001 From: SuperdoerTrav Date: Fri, 19 Jul 2024 13:51:16 -0700 Subject: [PATCH 18/55] fixed remaining circular imports --- ssapy/utils.py | 72 +++++++++++++++++++++++++++++++++++++++++++++++--- 1 file changed, 69 insertions(+), 3 deletions(-) diff --git a/ssapy/utils.py b/ssapy/utils.py index 92b6538..f02fa69 100644 --- a/ssapy/utils.py +++ b/ssapy/utils.py @@ -5,7 +5,7 @@ import astropy.units as u from . import datadir -from .constants import MOON_RADIUS +from .constants import RGEO, EARTH_RADIUS, MOON_RADIUS, WGS84_EARTH_OMEGA try: import erfa @@ -184,6 +184,10 @@ def normed(arr): return arr / np.sqrt(np.einsum("...i,...i", arr, arr))[..., None] +def einsum_norm(a, indices='ij,ji->i'): + return np.sqrt(np.einsum(indices, a, a)) + + def unitAngle3(r1, r2): """Robustly compute angle between unit vectors r1 and r2. Vectorized for multiple triplets. @@ -1259,7 +1263,7 @@ def rotate_vector(v_unit, theta, phi, plot=False, save_idx=False): if save_idx is not False: from .plotUtils import save_plot ax.set_title(f'Vector Plot\ntheta: {np.degrees(theta):.0f}, phi: {np.degrees(phi):.0f}', color='white') - save_plot(fig, f'/p/lustre1/yeager7/plots_gif/rotate_vector_frames/{save_idx}.png') + save_plot(fig, f"{os.path.expanduser('~/ssapy_test/rotate_vector_frames/')}{save_idx}.png") return v2 / np.linalg.norm(v2, axis=-1) @@ -1449,6 +1453,7 @@ def sim_lonlatrad(x, y, z, xe, ye, ze, xs, ys, zs): def sun_ra_dec(time_): + from .body import get_body out = get_body(Time(time_, format='mjd')) return out.ra.to('rad').value, out.dec.to('rad').value @@ -1680,6 +1685,7 @@ def proper_motion_ra_dec(r=None, v=None, x=None, y=None, z=None, vx=None, vy=Non def gcrf_to_lunar(r, t, v=None): + from .body import MoonPosition class MoonRotator: def __init__(self): self.mpm = MoonPosition() @@ -1707,6 +1713,7 @@ def __call__(self, r, t): def gcrf_to_lunar_fixed(r, t, v=None): + from .body import get_body r_lunar = gcrf_to_lunar(r, t) - gcrf_to_lunar(get_body('moon').position(t).T, t) if v is None: return r_lunar @@ -1767,6 +1774,8 @@ def gcrf_to_itrf(r_gcrf, t, v=None): def gcrf_to_sim_geo(r_gcrf, t, h=10): from .accel import AccelKepler from .compute import rv + from .orbit import Orbit + from .propagator import RK78Propagator if np.min(np.diff(t.gps)) < h: h = np.min(np.diff(t.gps)) r_gcrf = np.atleast_2d(r_gcrf) @@ -1878,6 +1887,64 @@ def run(self): return orig +def dms_to_dd(dms): # Degree minute second to Degree decimal + dms, out = [[dms] if type(dms) is str else dms][0], [] + for i in dms: + deg, minute, sec = [float(j) for j in i.split(':')] + if deg < 0: + minute, sec = float(f'-{minute}'), float(f'-{sec}') + out.append(deg + minute / 60 + sec / 3600) + return [out[0] if type(dms) is str or len(dms) == 1 else out][0] + + +def dd_to_dms(degree_decimal): + _d, __d = np.trunc(degree_decimal), degree_decimal - np.trunc(degree_decimal) + __d = [-__d if degree_decimal < 0 else __d][0] + _m, __m = np.trunc(__d * 60), __d * 60 - np.trunc(__d * 60) + _s = round(__m * 60, 4) + _s = [int(_s) if int(_s) == _s else _s][0] + if _s == 60: + _m, _s = _m + 1, '00' + elif _s > 60: + _m, _s = _m + 1, _s - 60 + + return f'{int(_d)}:{int(_m)}:{_s}' + + +def hms_to_dd(hms): + _type = type(hms) + hms, out = [[hms] if _type == str else hms][0], [] + for i in hms: + if i[0] != '-': + hour, minute, sec = i.split(':') + hour, minute, sec = float(hour), float(minute), float(sec) + out.append(hour * 15 + (minute / 4) + (sec / 240)) + else: + print('hms cannot be negative.') + + return [out[0] if _type == str or len(hms) == 1 else out][0] + + +def dd_to_hms(degree_decimal): + if type(degree_decimal) is str: + degree_decimal = dms_to_dd(degree_decimal) + if degree_decimal < 0: + print('dd for HMS conversion cannot be negative, assuming positive.') + _dd = -degree_decimal / 15 + else: + _dd = degree_decimal / 15 + _h, __h = np.trunc(_dd), _dd - np.trunc(_dd) + _m, __m = np.trunc(__h * 60), __h * 60 - np.trunc(__h * 60) + _s = round(__m * 60, 4) + _s = [int(_s) if int(_s) == _s else _s][0] + if _s == 60: + _m, _s = _m + 1, '00' + elif _s > 60: + _m, _s = _m + 1, _s - 60 + + return f'{int(_h)}:{int(_m)}:{_s}' + + def get_times(duration=(30, 'day'), freq=(1, 'hr'), t=Time("2025-01-01", scale='utc')): """ Calculate a list of times spaced equally apart over a specified duration. @@ -1940,7 +2007,6 @@ def interpolate_points_between(r, m): def check_lunar_collision(r, times, m=1000): - from .body import get_body """ Checks if the trajectory of a particle intersects with the Moon. From 66ee3b03d5e5cf154c99798911e5764a415630ab Mon Sep 17 00:00:00 2001 From: SuperdoerTrav Date: Fri, 19 Jul 2024 14:49:58 -0700 Subject: [PATCH 19/55] cleaned up the get_times function, added weeks as a unit --- ssapy/utils.py | 7 ++++--- tests/test_plots.py | 30 ++++++++++++++++++++++++++++++ 2 files changed, 34 insertions(+), 3 deletions(-) create mode 100644 tests/test_plots.py diff --git a/ssapy/utils.py b/ssapy/utils.py index f02fa69..a687fb6 100644 --- a/ssapy/utils.py +++ b/ssapy/utils.py @@ -1263,7 +1263,7 @@ def rotate_vector(v_unit, theta, phi, plot=False, save_idx=False): if save_idx is not False: from .plotUtils import save_plot ax.set_title(f'Vector Plot\ntheta: {np.degrees(theta):.0f}, phi: {np.degrees(phi):.0f}', color='white') - save_plot(fig, f"{os.path.expanduser('~/ssapy_test/rotate_vector_frames/')}{save_idx}.png") + save_plot(fig, f"{os.path.expanduser('~/ssapy_test_plots/rotate_vector_frames/')}{save_idx}.png") return v2 / np.linalg.norm(v2, axis=-1) @@ -1945,7 +1945,7 @@ def dd_to_hms(degree_decimal): return f'{int(_h)}:{int(_m)}:{_s}' -def get_times(duration=(30, 'day'), freq=(1, 'hr'), t=Time("2025-01-01", scale='utc')): +def get_times(duration, freq, t): """ Calculate a list of times spaced equally apart over a specified duration. @@ -1957,7 +1957,8 @@ def get_times(duration=(30, 'day'), freq=(1, 'hr'), t=Time("2025-01-01", scale=' frequency of time outputs in units provided t: ssapy.utils.Time, optional The starting time. Default is "2025-01-01". - + example input: + duration=(30, 'day'), freq=(1, 'hr'), t=Time("2025-01-01", scale='utc') Returns ------- times: array-like diff --git a/tests/test_plots.py b/tests/test_plots.py new file mode 100644 index 0000000..f81c607 --- /dev/null +++ b/tests/test_plots.py @@ -0,0 +1,30 @@ +from ssapy.plotUtils import * +from ssapy.simple import ssapy_orbit +from ssapy.io import listdir, sortbynum +from ssapy.utils import rotate_vector + +import os +import shutil +import numpy as np +from IPython.display import clear_output + +# Example usage: +v_unit = np.array([1, 0, 0]) # Replace this with your actual unit vector + +figs = [] + +save_directory = os.path.expanduser('~/ssapy_test_plots/rotate_vector_frames/') +os.makedirs(save_directory, exist_ok=True) + +i = 0 +for theta in range(0, 181, 10): + for phi in range(0, 361, 10): + clear_output(wait=True) + new_unit_vector = rotate_vector(v_unit, theta, phi, plot=True, save_idx=i) + i += 1 + +gif_path = f"{os.path.expanduser('~/ssapy_test_plots/')}rotate_vectors_{v_unit[0]:.0f}_{v_unit[1]:.0f}_{v_unit[2]:.0f}.gif" +write_gif(gif_name=gif_path, frames=sortbynum(listdir(f'{save_directory}*')), fps=20) +shutil.rmtree(save_directory) + +print(f"Rotate vector plot successfully created.") \ No newline at end of file From 7700a2b73b3a59f2fbe406076ecbfbf7b0890709 Mon Sep 17 00:00:00 2001 From: SuperdoerTrav Date: Fri, 19 Jul 2024 14:51:13 -0700 Subject: [PATCH 20/55] added doc string to ssapy_orbit() --- ssapy/simple.py | 35 +++++++++++++++++++++++++++++++++-- 1 file changed, 33 insertions(+), 2 deletions(-) diff --git a/ssapy/simple.py b/ssapy/simple.py index 5c3c2b1..c9a7421 100644 --- a/ssapy/simple.py +++ b/ssapy/simple.py @@ -92,8 +92,39 @@ def ssapy_prop(integration_timestep=60, propkw=ssapy_kwargs()): # Uses the current best propagator and acceleration models in ssapy def ssapy_orbit(orbit=None, a=None, e=0, i=0, pa=0, raan=0, ta=0, r=None, v=None, duration=(30, 'day'), freq=(1, 'hr'), start_date="2025-01-01", t=None, integration_timestep=10, mass=250, area=0.022, CD=2.3, CR=1.3, prop=ssapy_prop()): - # Everything is in SI units, except time. - # density #kg/m^3 --> density + """ + Compute the orbit of a spacecraft given either Keplerian elements or position and velocity vectors. + + Parameters: + - orbit (Orbit, optional): An Orbit object if you already have an orbit defined. + - a (float, optional): Semi-major axis of the orbit in meters. + - e (float, optional): Eccentricity of the orbit (default is 0, i.e., circular orbit). + - i (float, optional): Inclination of the orbit in degrees. + - pa (float, optional): Argument of perigee in degrees. + - raan (float, optional): Right ascension of the ascending node in degrees. + - ta (float, optional): True anomaly in degrees. + - r (array-like, optional): Position vector in meters. + - v (array-like, optional): Velocity vector in meters per second. + - duration (tuple, optional): Duration of the simulation as a tuple (value, unit), where unit is 'day', 'hour', etc. Default is 30 days. + - freq (tuple, optional): Frequency of the output as a tuple (value, unit), where unit is 'day', 'hour', etc. Default is 1 hour. + - start_date (str, optional): Start date of the simulation in 'YYYY-MM-DD' format. Default is "2025-01-01". + - t (array-like, optional): Specific times at which to compute the orbit. If None, times will be generated based on duration and frequency. + - integration_timestep (float, optional): Time step for numerical integration in seconds. Default is 10 seconds. + - mass (float, optional): Mass of the spacecraft in kilograms. Default is 250 kg. + - area (float, optional): Cross-sectional area of the spacecraft in square meters. Default is 0.022 m². + - CD (float, optional): Drag coefficient. Default is 2.3. + - CR (float, optional): Reflectivity coefficient. Default is 1.3. + - prop (function, optional): A function to compute the perturbation effects. Default is `ssapy_prop()`. + + Returns: + - r (array-like): Position vectors of the spacecraft at the specified times. + - v (array-like): Velocity vectors of the spacecraft at the specified times. + - t (array-like): Times at which the orbit was computed. Returned only if `t` was None. + + Raises: + - ValueError: If neither Keplerian elements nor position and velocity vectors are provided. + - RuntimeError or ValueError: If an error occurs during computation. + """ t0 = Time(start_date, scale='utc') if t is None: time_is_None = True From e7b913f1ebcdfea79432af658abebcd7283355eb Mon Sep 17 00:00:00 2001 From: SuperdoerTrav Date: Fri, 19 Jul 2024 15:09:26 -0700 Subject: [PATCH 21/55] added plot_path as optional input to rotate_vector() --- ssapy/utils.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/ssapy/utils.py b/ssapy/utils.py index a687fb6..81a9745 100644 --- a/ssapy/utils.py +++ b/ssapy/utils.py @@ -1197,7 +1197,7 @@ def rotation_matrix_from_vectors(vec1, vec2): return rotation_matrix -def rotate_vector(v_unit, theta, phi, plot=False, save_idx=False): +def rotate_vector(v_unit, theta, phi, plot_path=False, save_idx=False): v_unit = v_unit / np.linalg.norm(v_unit, axis=-1) if np.all(np.abs(v_unit) != np.max(np.abs(v_unit))): perp_vector = np.cross(v_unit, np.array([1, 0, 0])) @@ -1240,7 +1240,7 @@ def rotate_vector(v_unit, theta, phi, plot=False, save_idx=False): v2 = np.dot(R2, v1) - if plot: + if plot_path is False: import matplotlib.pyplot as plt plt.rcParams.update({'font.size': 9, 'figure.facecolor': 'black'}) fig = plt.figure() @@ -1263,7 +1263,7 @@ def rotate_vector(v_unit, theta, phi, plot=False, save_idx=False): if save_idx is not False: from .plotUtils import save_plot ax.set_title(f'Vector Plot\ntheta: {np.degrees(theta):.0f}, phi: {np.degrees(phi):.0f}', color='white') - save_plot(fig, f"{os.path.expanduser('~/ssapy_test_plots/rotate_vector_frames/')}{save_idx}.png") + save_plot(fig, f"{plot_path}{save_idx}.png") return v2 / np.linalg.norm(v2, axis=-1) From 27b39e5cf959d1b8d41ef6e9a8192e385a1bc383 Mon Sep 17 00:00:00 2001 From: SuperdoerTrav Date: Fri, 19 Jul 2024 16:14:09 -0700 Subject: [PATCH 22/55] added simply.py which will contain simplified functions for using ssapy --- ssapy/__init__.py | 1 + 1 file changed, 1 insertion(+) diff --git a/ssapy/__init__.py b/ssapy/__init__.py index 504f1ba..b505ddf 100644 --- a/ssapy/__init__.py +++ b/ssapy/__init__.py @@ -39,6 +39,7 @@ from . import plotUtils from . import io from . import utils +from . import simple from astropy.time import Time, TimeDelta import astropy.units as u from datetime import timedelta From dd9423e9b86634a44a69c6b8c0297c962f63e950 Mon Sep 17 00:00:00 2001 From: SuperdoerTrav Date: Fri, 19 Jul 2024 16:51:45 -0700 Subject: [PATCH 23/55] adding a testing plot testing function --- ssapy/plotUtils.py | 36 +++++++++- ssapy/simple.py | 12 +--- ssapy/utils.py | 3 +- tests/test_plots.py | 165 +++++++++++++++++++++++++++++++++++++++----- 4 files changed, 188 insertions(+), 28 deletions(-) diff --git a/ssapy/plotUtils.py b/ssapy/plotUtils.py index 4e88790..a6b52e1 100644 --- a/ssapy/plotUtils.py +++ b/ssapy/plotUtils.py @@ -474,6 +474,7 @@ def koe_plot(r, v, t=Time("2025-01-01", scale='utc') + np.linspace(0, int(1 * 36 else: orbital_elements = calculate_orbital_elements(r, v, mu_barycenter=MOON_MU) fig, ax1 = plt.subplots(dpi=100) + fig.patch.set_facecolor('white') ax1.plot([], [], label='semi-major axis [GEO]', c='C0', linestyle='-') ax2 = ax1.twinx() make_white(fig, *[ax1, ax2]) @@ -513,6 +514,7 @@ def koe_2dhist(stable_data, title="Initial orbital elements of\n1 year stable ci else: norm = mplcolors.Normalize(limits[0], limits[1]) fig, axes = plt.subplots(dpi=100, figsize=(10, 8), nrows=3, ncols=3) + fig.patch.set_facecolor('white') st = fig.suptitle(title, fontsize=12) st.set_x(0.46) st.set_y(0.9) @@ -742,6 +744,38 @@ def make_black(fig, *axes): return fig, axes +def draw_dashed_circle(ax, normal_vector, radius, dashes, dash_length=0.1, label='Dashed Circle'): + from .utils import rotation_matrix_from_vectors + # Define the circle in the xy-plane + theta = np.linspace(0, 2 * np.pi, 1000) + x_circle = radius * np.cos(theta) + y_circle = radius * np.sin(theta) + z_circle = np.zeros_like(theta) + + # Stack the coordinates into a matrix + circle_points = np.vstack((x_circle, y_circle, z_circle)).T + + # Create the rotation matrix to align z-axis with the normal vector + normal_vector = normal_vector / np.linalg.norm(normal_vector) + rotation_matrix = rotation_matrix_from_vectors(np.array([0, 0, 1]), normal_vector) + + # Rotate the circle points + rotated_points = circle_points @ rotation_matrix.T + + # Create dashed effect + dash_points = [] + dash_gap = int(len(theta) / dashes) + for i in range(dashes): + start_idx = i * dash_gap + end_idx = start_idx + int(dash_length * len(theta)) + dash_points.append(rotated_points[start_idx:end_idx]) + + # Plot the dashed circle in 3D + for points in dash_points: + ax.plot(points[:, 0], points[:, 1], points[:, 2], 'k--', label=label) + label = None # Only one label + + # ##################################################################### # Formatting x axis # ##################################################################### @@ -845,7 +879,7 @@ def save_plot(figure, save_path, dpi=200): # Save the figure as a PNG image figure.savefig(save_path, dpi=dpi, bbox_inches='tight') plt.close(figure) # Close the figure to release resources - print(f"Figure saved at: {save_path}") + # print(f"Figure saved at: {save_path}") except Exception as e: print(f"Error occurred while saving the figure: {e}") diff --git a/ssapy/simple.py b/ssapy/simple.py index c9a7421..1e86a88 100644 --- a/ssapy/simple.py +++ b/ssapy/simple.py @@ -91,7 +91,7 @@ def ssapy_prop(integration_timestep=60, propkw=ssapy_kwargs()): # Uses the current best propagator and acceleration models in ssapy -def ssapy_orbit(orbit=None, a=None, e=0, i=0, pa=0, raan=0, ta=0, r=None, v=None, duration=(30, 'day'), freq=(1, 'hr'), start_date="2025-01-01", t=None, integration_timestep=10, mass=250, area=0.022, CD=2.3, CR=1.3, prop=ssapy_prop()): +def ssapy_orbit(orbit=None, a=None, e=0, i=0, pa=0, raan=0, ta=0, r=None, v=None, duration=(30, 'day'), freq=(1, 'hr'), t0=Time("2025-01-01", scale='utc'), t=None, prop=ssapy_prop()): """ Compute the orbit of a spacecraft given either Keplerian elements or position and velocity vectors. @@ -107,13 +107,8 @@ def ssapy_orbit(orbit=None, a=None, e=0, i=0, pa=0, raan=0, ta=0, r=None, v=None - v (array-like, optional): Velocity vector in meters per second. - duration (tuple, optional): Duration of the simulation as a tuple (value, unit), where unit is 'day', 'hour', etc. Default is 30 days. - freq (tuple, optional): Frequency of the output as a tuple (value, unit), where unit is 'day', 'hour', etc. Default is 1 hour. - - start_date (str, optional): Start date of the simulation in 'YYYY-MM-DD' format. Default is "2025-01-01". + - t0 (str, optional): Start date of the simulation in 'YYYY-MM-DD' format. Default is "2025-01-01". - t (array-like, optional): Specific times at which to compute the orbit. If None, times will be generated based on duration and frequency. - - integration_timestep (float, optional): Time step for numerical integration in seconds. Default is 10 seconds. - - mass (float, optional): Mass of the spacecraft in kilograms. Default is 250 kg. - - area (float, optional): Cross-sectional area of the spacecraft in square meters. Default is 0.022 m². - - CD (float, optional): Drag coefficient. Default is 2.3. - - CR (float, optional): Reflectivity coefficient. Default is 1.3. - prop (function, optional): A function to compute the perturbation effects. Default is `ssapy_prop()`. Returns: @@ -125,10 +120,9 @@ def ssapy_orbit(orbit=None, a=None, e=0, i=0, pa=0, raan=0, ta=0, r=None, v=None - ValueError: If neither Keplerian elements nor position and velocity vectors are provided. - RuntimeError or ValueError: If an error occurs during computation. """ - t0 = Time(start_date, scale='utc') if t is None: time_is_None = True - t = get_times(duration=duration, freq=freq, t=t) + t = get_times(duration=duration, freq=freq, t=t0) else: time_is_None = False diff --git a/ssapy/utils.py b/ssapy/utils.py index 81a9745..84bb304 100644 --- a/ssapy/utils.py +++ b/ssapy/utils.py @@ -1240,7 +1240,7 @@ def rotate_vector(v_unit, theta, phi, plot_path=False, save_idx=False): v2 = np.dot(R2, v1) - if plot_path is False: + if plot_path is not False: import matplotlib.pyplot as plt plt.rcParams.update({'font.size': 9, 'figure.facecolor': 'black'}) fig = plt.figure() @@ -2008,6 +2008,7 @@ def interpolate_points_between(r, m): def check_lunar_collision(r, times, m=1000): + from .body import get_body """ Checks if the trajectory of a particle intersects with the Moon. diff --git a/tests/test_plots.py b/tests/test_plots.py index f81c607..f14c372 100644 --- a/tests/test_plots.py +++ b/tests/test_plots.py @@ -1,30 +1,161 @@ -from ssapy.plotUtils import * -from ssapy.simple import ssapy_orbit -from ssapy.io import listdir, sortbynum -from ssapy.utils import rotate_vector - +import ssapy import os import shutil import numpy as np -from IPython.display import clear_output +import matplotlib.pyplot as plt + +save_folder = os.path.expanduser('~/ssapy_test_plots/') +print(f"Putting test_plot.py output in: {save_folder}") -# Example usage: +# Testing rotate_vector() in utils. v_unit = np.array([1, 0, 0]) # Replace this with your actual unit vector figs = [] -save_directory = os.path.expanduser('~/ssapy_test_plots/rotate_vector_frames/') -os.makedirs(save_directory, exist_ok=True) +temp_directory = f'{save_folder}rotate_vector_frames/' +os.makedirs(temp_directory, exist_ok=True) i = 0 -for theta in range(0, 181, 10): - for phi in range(0, 361, 10): - clear_output(wait=True) - new_unit_vector = rotate_vector(v_unit, theta, phi, plot=True, save_idx=i) +for theta in range(0, 181, 20): + for phi in range(0, 361, 20): + new_unit_vector = ssapy.utils.rotate_vector(v_unit, theta, phi, plot_path=temp_directory, save_idx=i) i += 1 -gif_path = f"{os.path.expanduser('~/ssapy_test_plots/')}rotate_vectors_{v_unit[0]:.0f}_{v_unit[1]:.0f}_{v_unit[2]:.0f}.gif" -write_gif(gif_name=gif_path, frames=sortbynum(listdir(f'{save_directory}*')), fps=20) -shutil.rmtree(save_directory) +gif_path = f"{save_folder}rotate_vectors_{v_unit[0]:.0f}_{v_unit[1]:.0f}_{v_unit[2]:.0f}.gif" +ssapy.plotUtils.write_gif(gif_name=gif_path, frames=ssapy.io.sortbynum(ssapy.io.listdir(f'{temp_directory}*')), fps=20) +shutil.rmtree(temp_directory) + + +# Creating orbit plots +times = ssapy.utils.get_times(duration=(1, 'year'), freq=(1, 'hour'), t='2025-3-1') +moon = ssapy.get_body("moon").position(times).T + + +def DRO(t, delta_r=7.52064e7, delta_v=344): + moon = ssapy.get_body("moon") + + unit_vector_moon = moon.position(t) / np.linalg.norm(moon.position(t)) + moon_v = (moon.position(t.gps) - moon.position(t.gps - 1)) / 1 + unit_vector_moon_velocity = moon_v / np.linalg.norm(moon_v) + ssapy.compute.lunar_lagrange_points(t=times[0]) + + r = (np.linalg.norm(moon.position(t)) - delta_r) * unit_vector_moon + v = (np.linalg.norm(moon_v) + delta_v) * unit_vector_moon_velocity + + orbit = ssapy.Orbit(r=r, v=v, t=t) + return orbit + + +def Lunar_L4(t, delta_r=7.52064e7, delta_v=344): + moon = ssapy.get_body("moon") + + unit_vector_moon = moon.position(t) / np.linalg.norm(moon.position(t)) + moon_v = (moon.position(t.gps) - moon.position(t.gps - 1)) / 1 + unit_vector_moon_velocity = moon_v / np.linalg.norm(moon_v) + + r = (np.linalg.norm(moon.position(t)) - delta_r) * unit_vector_moon + v = (np.linalg.norm(moon_v) + delta_v) * unit_vector_moon_velocity + + orbit = ssapy.Orbit(r=r, v=v, t=t) + return orbit + + +# DRO +dro_orbit = DRO(t=times[0]) +r, v = ssapy.simple.ssapy_orbit(orbit=dro_orbit, t=times) +ssapy.plotUtils.orbit_plot(r=r, t=times, save_path=f"{save_folder}DRO_orbit", frame='Lunar', show=False) +r_lunar, v_lunar = ssapy.utils.gcrf_to_lunar_fixed(r, t=times, v=True) +ssapy.plotUtils.koe_plot(r, v, t=times, body='Earth', save_path=f"{save_folder}Keplerian_orbital_elements.png") + +L4_orbit = Lunar_L4(t=times[0], delta_r=7.52064e7, delta_v=344) +r, v = ssapy.simple.ssapy_orbit(orbit=L4_orbit, t=times) +ssapy.plotUtils.orbit_plot(r=r, t=times, save_path=f"{save_folder}gcrf_plot.png", frame='gcrf', show=True) +ssapy.plotUtils.orbit_plot(r=r, t=times, save_path=f"{save_folder}itrf_plot", frame='itrf', show=True) +ssapy.plotUtils.orbit_plot(r=r, t=times, save_path=f"{save_folder}lunar_plot", frame='lunar', show=True) +ssapy.plotUtils.orbit_plot(r=r, t=times, save_path=f"{save_folder}lunar_axis_lot", frame='lunar axis', show=True) + + +# Example usage +earth_pos = np.array([0, 0, 0]) # Earth at the origin +moon_pos = ssapy.get_body("moon").position(times[0]).T + +# Plotting +fig = plt.figure(figsize=(8, 8)) +fig.patch.set_facecolor('white') +ax = fig.add_subplot(111, projection='3d') + +# Plot Earth +ax.scatter(earth_pos[0], earth_pos[1], earth_pos[2], color='blue', label='Earth') +ax.text(earth_pos[0], earth_pos[1], earth_pos[2], 'Earth', color='blue') + +# Plot Moon +ax.scatter(moon_pos[0], moon_pos[1], moon_pos[2], color='grey', label='Moon') +ax.text(moon_pos[0], moon_pos[1], moon_pos[2], 'Moon', color='grey') + +# Plot Lagrange points +colors = ['red', 'green', 'purple', 'orange', 'cyan'] +for (point, pos), color in zip(ssapy.compute.lunar_lagrange_points(t=times[0]).items(), colors): + ax.scatter(pos[0], pos[1], pos[2], color=color, label=point) + ax.text(pos[0], pos[1], pos[2], point, color=color) + +# Add a dashed black circle at distance LD +current_LD = np.linalg.norm(moon_pos, axis=-1) +normal_vector = ssapy.compute.moon_normal_vector(t=times[0]) +ssapy.plotUtils.draw_dashed_circle(ax, normal_vector, current_LD, 12) +ax.quiver(0, 0, 0, normal_vector[0], normal_vector[1], normal_vector[2], color='r', length=1) + +# Labels and legend +ax.set_xlabel('X (m)') +ax.set_ylabel('Y (m)') +ax.set_zlabel('Z (m)') +ax.set_title("Lunar Lagrange Points using Moon's true position") +ax.axis('equal') +ax.legend() +plt.show() +ssapy.plotUtils.save_plot(fig, save_path=f"{save_folder}lagrange_points") + +# Plotting +fig = plt.figure(figsize=(8, 8)) +fig.patch.set_facecolor('white') +ax = fig.add_subplot(111, projection='3d') + +# Plot Earth +ax.scatter(earth_pos[0], earth_pos[1], earth_pos[2], color='blue', label='Earth') +ax.text(earth_pos[0], earth_pos[1], earth_pos[2], 'Earth', color='blue') + +# Plot Moon +ax.scatter(moon_pos[0], moon_pos[1], moon_pos[2], color='grey', label='Moon') +ax.text(moon_pos[0], moon_pos[1], moon_pos[2], 'Moon', color='grey') + +# Plot Lagrange points +colors = ['red', 'green', 'purple', 'orange', 'cyan'] +for (point, pos), color in zip(ssapy.compute.lunar_lagrange_points_circular(t=times[0]).items(), colors): + ax.scatter(pos[0], pos[1], pos[2], color=color, label=point) + ax.text(pos[0], pos[1], pos[2], point, color=color) + +# Add a dashed black circle at distance LD +current_LD = np.linalg.norm(moon_pos, axis=-1) +normal_vector = ssapy.compute.moon_normal_vector(t=times[0]) +ssapy.plotUtils.draw_dashed_circle(ax, normal_vector, current_LD, 12) +ax.quiver(0, 0, 0, normal_vector[0], normal_vector[1], normal_vector[2], color='r', length=1) + +# Labels and legend +ax.set_xlabel('X (m)') +ax.set_ylabel('Y (m)') +ax.set_zlabel('Z (m)') +ax.set_title('Lunar Lagrange Points assuming circular orbit.') +ax.axis('equal') +ax.legend() +plt.show() +ssapy.plotUtils.save_plot(fig, save_path=f"{save_folder}lagrange_points") + + -print(f"Rotate vector plot successfully created.") \ No newline at end of file +print(f"Created a GCRF orbit plot.") +print(f"Created a ITRF orbit plot.") +print(f"Created a Lunar orbit plot.") +print(f"Created a Lunar axis orbit plot.") +print(f"Lagrange points were calculated correctly.") +print(f"Rotate vector plot successfully created.") +print(f"save_plot() executed succesfully.") +print(f"write_gif() executed succesfully.") \ No newline at end of file From ff406d5723e5ecf23c5b7e9f1284b720b94d5421 Mon Sep 17 00:00:00 2001 From: SuperdoerTrav Date: Mon, 22 Jul 2024 10:35:58 -0700 Subject: [PATCH 24/55] changed the output of ssapy_orbit() to always give back the time array used. --- ssapy/simple.py | 8 +------- 1 file changed, 1 insertion(+), 7 deletions(-) diff --git a/ssapy/simple.py b/ssapy/simple.py index 1e86a88..967d0cc 100644 --- a/ssapy/simple.py +++ b/ssapy/simple.py @@ -121,10 +121,7 @@ def ssapy_orbit(orbit=None, a=None, e=0, i=0, pa=0, raan=0, ta=0, r=None, v=None - RuntimeError or ValueError: If an error occurs during computation. """ if t is None: - time_is_None = True t = get_times(duration=duration, freq=freq, t=t0) - else: - time_is_None = False if orbit is not None: pass @@ -138,10 +135,7 @@ def ssapy_orbit(orbit=None, a=None, e=0, i=0, pa=0, raan=0, ta=0, r=None, v=None try: r, v = rv(orbit, t, prop) - if time_is_None: - return r, v, t - else: - return r, v + return r, v, t except (RuntimeError, ValueError) as err: print(err) return np.nan, np.nan, np.nan From 265d383bc427ae68b9c180685cd44e19859a48ae Mon Sep 17 00:00:00 2001 From: SuperdoerTrav Date: Mon, 22 Jul 2024 15:19:07 -0700 Subject: [PATCH 25/55] changed save directory of test plots --- ssapy/plotUtils.py | 4 ++-- tests/test_plots.py | 42 +++++++++++++++++------------------------- 2 files changed, 19 insertions(+), 27 deletions(-) diff --git a/ssapy/plotUtils.py b/ssapy/plotUtils.py index a6b52e1..8ed91a3 100644 --- a/ssapy/plotUtils.py +++ b/ssapy/plotUtils.py @@ -132,7 +132,7 @@ def drawMoon(time, ngrid=100, R=MOON_RADIUS, rfactor=1): ) -def groundTrackPlot(r, time, ground_stations=None, save_path=False): +def groundTrackPlot(r, t, ground_stations=None, save_path=False): """ Parameters ---------- @@ -141,7 +141,7 @@ def groundTrackPlot(r, time, ground_stations=None, save_path=False): optional - ground_stations: (n,2) array of of ground station (lat,lon) in degrees """ - lon, lat, height = groundTrack(r, time) + lon, lat, height = groundTrack(r, t) fig = plt.figure(figsize=(15, 12)) plt.imshow(load_earth_file(), extent=[-180, 180, -90, 90]) diff --git a/tests/test_plots.py b/tests/test_plots.py index f14c372..0fb4ffe 100644 --- a/tests/test_plots.py +++ b/tests/test_plots.py @@ -4,7 +4,7 @@ import numpy as np import matplotlib.pyplot as plt -save_folder = os.path.expanduser('~/ssapy_test_plots/') +save_folder = os.path.expanduser('./ssapy_test_plots/') print(f"Putting test_plot.py output in: {save_folder}") # Testing rotate_vector() in utils. @@ -46,34 +46,30 @@ def DRO(t, delta_r=7.52064e7, delta_v=344): return orbit -def Lunar_L4(t, delta_r=7.52064e7, delta_v=344): - moon = ssapy.get_body("moon") - - unit_vector_moon = moon.position(t) / np.linalg.norm(moon.position(t)) - moon_v = (moon.position(t.gps) - moon.position(t.gps - 1)) / 1 - unit_vector_moon_velocity = moon_v / np.linalg.norm(moon_v) - - r = (np.linalg.norm(moon.position(t)) - delta_r) * unit_vector_moon - v = (np.linalg.norm(moon_v) + delta_v) * unit_vector_moon_velocity - - orbit = ssapy.Orbit(r=r, v=v, t=t) - return orbit - - # DRO dro_orbit = DRO(t=times[0]) -r, v = ssapy.simple.ssapy_orbit(orbit=dro_orbit, t=times) +r, v, t = ssapy.simple.ssapy_orbit(orbit=dro_orbit, t=times) ssapy.plotUtils.orbit_plot(r=r, t=times, save_path=f"{save_folder}DRO_orbit", frame='Lunar', show=False) r_lunar, v_lunar = ssapy.utils.gcrf_to_lunar_fixed(r, t=times, v=True) +print("Succesfully converted gcrf to lunar frame.") ssapy.plotUtils.koe_plot(r, v, t=times, body='Earth', save_path=f"{save_folder}Keplerian_orbital_elements.png") -L4_orbit = Lunar_L4(t=times[0], delta_r=7.52064e7, delta_v=344) -r, v = ssapy.simple.ssapy_orbit(orbit=L4_orbit, t=times) ssapy.plotUtils.orbit_plot(r=r, t=times, save_path=f"{save_folder}gcrf_plot.png", frame='gcrf', show=True) ssapy.plotUtils.orbit_plot(r=r, t=times, save_path=f"{save_folder}itrf_plot", frame='itrf', show=True) ssapy.plotUtils.orbit_plot(r=r, t=times, save_path=f"{save_folder}lunar_plot", frame='lunar', show=True) ssapy.plotUtils.orbit_plot(r=r, t=times, save_path=f"{save_folder}lunar_axis_lot", frame='lunar axis', show=True) +print(f"Created a GCRF orbit plot.") +print(f"Created a ITRF orbit plot.") +print(f"Created a Lunar orbit plot.") +print(f"Created a Lunar axis orbit plot.") + +# Globe plot of GTO Orbit +r_geo, _, t_geo = ssapy.simple.ssapy_orbit(a=ssapy.constants.RGEO, e=0.3, t=times) +ssapy.plotUtils.globe_plot(r=r_geo, t=t_geo, save_path=f"{save_folder}globe_plot", scale=5) +print('Created a globe plot.') +ssapy.plotUtils.groundTrackPlot(r=r_geo, t=t_geo, ground_stations=None, save_path=f"{save_folder}ground_track_plot") +print('Created a ground track plot.') # Example usage earth_pos = np.array([0, 0, 0]) # Earth at the origin @@ -149,13 +145,9 @@ def Lunar_L4(t, delta_r=7.52064e7, delta_v=344): plt.show() ssapy.plotUtils.save_plot(fig, save_path=f"{save_folder}lagrange_points") - - -print(f"Created a GCRF orbit plot.") -print(f"Created a ITRF orbit plot.") -print(f"Created a Lunar orbit plot.") -print(f"Created a Lunar axis orbit plot.") print(f"Lagrange points were calculated correctly.") print(f"Rotate vector plot successfully created.") print(f"save_plot() executed succesfully.") -print(f"write_gif() executed succesfully.") \ No newline at end of file +print(f"write_gif() executed succesfully.") + +print(f"\nFinished plot testing!\n") \ No newline at end of file From 8a185240efcea97ebf030ebaacbf4e70f78c43cc Mon Sep 17 00:00:00 2001 From: SuperdoerTrav Date: Mon, 22 Jul 2024 15:20:58 -0700 Subject: [PATCH 26/55] added an ignore for the ssapy_test_plots/ folder which will be generated if user runs test_plots.py --- .gitignore | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.gitignore b/.gitignore index 195140c..64c631b 100644 --- a/.gitignore +++ b/.gitignore @@ -46,6 +46,8 @@ coverage.xml pytest.xml pytest_session.txt cover/ +# Ignore ssapy_test_plots folder in tests subfolder +tests/ssapy_test_plots/ # Sphinx documentation docs/_build/ From 8f36d9e874256ac3391c976632a00f02234d2394 Mon Sep 17 00:00:00 2001 From: SuperdoerTrav Date: Mon, 22 Jul 2024 16:14:02 -0700 Subject: [PATCH 27/55] added docstrings to all plot functions --- ssapy/plotUtils.py | 430 ++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 423 insertions(+), 7 deletions(-) diff --git a/ssapy/plotUtils.py b/ssapy/plotUtils.py index 8ed91a3..83a2345 100644 --- a/ssapy/plotUtils.py +++ b/ssapy/plotUtils.py @@ -210,15 +210,52 @@ def check_numpy_array(variable): def orbit_plot(r, t=[], limits=False, title='', figsize=(7, 7), save_path=False, frame="gcrf", show=False): """ - Parameters + Plot 2D and 3D projections of the orbit of one or more objects. + + Parameters: ---------- - r : (n,3) or array of [(n,3), ..., (n,3)] array_like - Position of orbiting object(s) in meters. - t: optional - t when r was calculated. - limits: optional - x and y limits of the plot - title: optional - title of the plot - """ + r : (n,3) or list of [(n,3), ..., (n,3)] array_like + Position of orbiting object(s) in meters. If a list is provided, each element represents the orbit of a different object. + t : array_like, optional + Time array corresponding to the position vectors. Required for frames "itrf", "lunar", or "lunar fixed". + limits : float or tuple, optional + Limits for the x and y axes of the 2D plots. If not provided, it is calculated based on the data. + title : str, optional + Title of the plot. Default is an empty string. + figsize : tuple of int, optional + Size of the figure (width, height) in inches. Default is (7, 7). + save_path : str, optional + Path to save the generated plot. If not provided, the plot will not be saved. Default is False. + frame : str, optional + The coordinate frame for the plot. Accepted values are "gcrf", "itrf", "lunar", "lunar fixed". + show : bool, optional + Whether to display the plot. Default is False. + + Returns: + ------- + fig : matplotlib.figure.Figure + The figure object containing the plot. + ax : list of matplotlib.axes.Axes + The list of axes objects used in the plot. Includes 2D axes (ax1, ax2, ax3) and 3D axis (ax4). + + Notes: + ------ + - The function creates a set of subplots showing the projection of the orbit data in 2D (x-y, x-z, y-z) and 3D (x-y-z). + - For the "lunar" and "lunar fixed" frames, the plot also includes the positions of Lagrange points. + - If a list of orbits is provided, each orbit is plotted with a different color. + - The plot background is set to black, and axis labels and ticks are set to white. + + Example usage: + ---------- + import numpy as np + from your_module import orbit_plot + + # Example data + r = np.array([[1e7, 2e7, 3e7], [4e7, 5e7, 6e7]]) # Replace with actual data + t = np.linspace(0, 10, len(r)) # Replace with actual time data + fig, axes = orbit_plot(r, t, limits=1e8, title='Orbit Plot', frame='gcrf', show=True) + """ def _make_scatter(fig, ax1, ax2, ax3, ax4, r, t, limits, title='', orbit_index='', num_orbits=1, frame=False): if np.size(t) < 1: if frame in ["itrf", "lunar", "lunar_fixed"]: @@ -427,6 +464,38 @@ def get_main_category(frame): def globe_plot(r, t, limits=False, title='', figsize=(7, 8), save_path=False, el=30, az=0, scale=1): + """ + Plot a 3D scatter plot of position vectors on a globe representation. + + Parameters: + - r (array-like): Position vectors with shape (n, 3), where n is the number of points. + - t (array-like): Time array corresponding to the position vectors. This parameter is not used in the current function implementation but is included for consistency. + - limits (float, optional): The limit for the plot axes. If not provided, it is calculated based on the data. Default is False. + - title (str, optional): Title of the plot. Default is an empty string. + - figsize (tuple of int, optional): Figure size (width, height) in inches. Default is (7, 8). + - save_path (str, optional): Path to save the generated plot. If not provided, the plot will not be saved. Default is False. + - el (int, optional): Elevation angle (in degrees) for the view of the plot. Default is 30. + - az (int, optional): Azimuth angle (in degrees) for the view of the plot. Default is 0. + - scale (int, optional): Scale factor for resizing the Earth image. Default is 1. + + Returns: + - fig (matplotlib.figure.Figure): The figure object containing the plot. + - ax (matplotlib.axes._subplots.Axes3DSubplot): The 3D axis object used in the plot. + + The function creates a 3D scatter plot of the position vectors on a globe. The globe is represented using a textured Earth image, and the scatter points are colored using a rainbow colormap. The plot's background is set to black, and the plot is displayed with customizable elevation and azimuth angles. + + Example usage: + ``` + import numpy as np + from your_module import globe_plot + + # Example data + r = np.array([[1, 2, 3], [4, 5, 6]]) # Replace with actual data + t = np.arange(len(r)) # Replace with actual time data + + globe_plot(r, t, save_path='globe_plot.png') + ``` + """ x = r[:, 0] / RGEO y = r[:, 1] / RGEO z = r[:, 2] / RGEO @@ -469,6 +538,37 @@ def globe_plot(r, t, limits=False, title='', figsize=(7, 8), save_path=False, el def koe_plot(r, v, t=Time("2025-01-01", scale='utc') + np.linspace(0, int(1 * 365.25), int(365.25 * 24)), elements=['a', 'e', 'i'], save_path=False, body='Earth'): + """ + Plot orbital elements over time for a given trajectory. + + Parameters: + - r (array-like): Position vectors for the orbit. + - v (array-like): Velocity vectors for the orbit. + - t (array-like, optional): Time array for the plot, given as a sequence of `astropy.time.Time` objects or a `Time` object with `np.linspace`. Default is one year of hourly intervals starting from "2025-01-01". + - elements (list of str, optional): List of orbital elements to plot. Options include 'a' (semi-major axis), 'e' (eccentricity), and 'i' (inclination). Default is ['a', 'e', 'i']. + - save_path (str, optional): Path to save the generated plot. If not provided, the plot will not be saved. Default is False. + - body (str, optional): The celestial body for which to calculate the orbital elements. Options are 'Earth' or 'Moon'. Default is 'Earth'. + + Returns: + - fig (matplotlib.figure.Figure): The figure object containing the plot. + - ax1 (matplotlib.axes.Axes): The primary axis object used in the plot. + + The function calculates orbital elements for the given position and velocity vectors, and plots these elements over time. It creates a plot with two y-axes: one for the eccentricity and inclination, and the other for the semi-major axis. The x-axis represents time in decimal years. + + Example usage: + ``` + import numpy as np + from astropy.time import Time + from your_module import koe_plot + + # Example data + r = np.array([[[1, 0, 0], [0, 1, 0]]]) # Replace with actual data + v = np.array([[[0, 1, 0], [-1, 0, 0]]]) # Replace with actual data + t = Time("2025-01-01", scale='utc') + np.linspace(0, int(1 * 365.25), int(365.25 * 24)) + + koe_plot(r, v, t, save_path='orbital_elements_plot.png') + ``` + """ if 'earth' in body.lower(): orbital_elements = calculate_orbital_elements(r, v, mu_barycenter=EARTH_MU) else: @@ -509,6 +609,40 @@ def koe_plot(r, v, t=Time("2025-01-01", scale='utc') + np.linspace(0, int(1 * 36 def koe_2dhist(stable_data, title="Initial orbital elements of\n1 year stable cislunar orbits", limits=[1, 50], bins=200, logscale=False, cmap='coolwarm', save_path=False): + """ + Create a 2D histogram plot for various orbital elements of stable cislunar orbits. + + Parameters: + - stable_data (object): An object with attributes `a`, `e`, `i`, and `ta`, which are arrays of semi-major axis, eccentricity, inclination, and true anomaly, respectively. + - title (str, optional): Title of the figure. Default is "Initial orbital elements of\n1 year stable cislunar orbits". + - limits (list, optional): Color scale limits for the histogram. Default is [1, 50]. + - bins (int, optional): Number of bins for the 2D histograms. Default is 200. + - logscale (bool or str, optional): Whether to use logarithmic scaling for the color bar. Default is False. Can also be 'log' to apply logarithmic scaling. + - cmap (str, optional): Colormap to use for the histograms. Default is 'coolwarm'. + - save_path (str, optional): Path to save the generated plot. If not provided, the plot will not be saved. Default is False. + + Returns: + - fig (matplotlib.figure.Figure): The figure object containing the 2D histograms. + + This function creates a 3x3 grid of 2D histograms showing the relationships between various orbital elements, including semi-major axis, eccentricity, inclination, and true anomaly. The color scale of the histograms can be adjusted with a logarithmic or linear normalization. The plot is customized with labels and a color bar. + + Example usage: + ``` + import numpy as np + from your_module import koe_2dhist + + # Example data + class StableData: + def __init__(self): + self.a = np.random.uniform(1, 20, 1000) + self.e = np.random.uniform(0, 1, 1000) + self.i = np.radians(np.random.uniform(0, 90, 1000)) + self.ta = np.radians(np.random.uniform(0, 360, 1000)) + + stable_data = StableData() + koe_2dhist(stable_data, save_path='orbit_histograms.pdf') + ``` + """ if logscale or logscale == 'log': norm = mplcolors.LogNorm(limits[0], limits[1]) else: @@ -574,6 +708,45 @@ def koe_2dhist(stable_data, title="Initial orbital elements of\n1 year stable ci def scatter2d(x, y, cs, xlabel='x', ylabel='y', title='', cbar_label='', dotsize=1, colorsMap='jet', colorscale='linear', colormin=False, colormax=False, save_path=False): + """ + Create a 2D scatter plot with optional color mapping. + + Parameters: + - x (numpy.ndarray): Array of x-coordinates. + - y (numpy.ndarray): Array of y-coordinates. + - cs (numpy.ndarray): Array of values for color mapping. + - xlabel (str, optional): Label for the x-axis. Default is 'x'. + - ylabel (str, optional): Label for the y-axis. Default is 'y'. + - title (str, optional): Title of the plot. Default is an empty string. + - cbar_label (str, optional): Label for the color bar. Default is an empty string. + - dotsize (int, optional): Size of the dots in the scatter plot. Default is 1. + - colorsMap (str, optional): Colormap to use for the color mapping. Default is 'jet'. + - colorscale (str, optional): Scale for the color mapping, either 'linear' or 'log'. Default is 'linear'. + - colormin (float, optional): Minimum value for color scaling. If False, it is set to the minimum value of `cs`. Default is False. + - colormax (float, optional): Maximum value for color scaling. If False, it is set to the maximum value of `cs`. Default is False. + - save_path (str, optional): File path to save the plot. If not provided, the plot is not saved. Default is False. + + Returns: + - fig (matplotlib.figure.Figure): The figure object. + - ax (matplotlib.axes._subplots.AxesSubplot): The 2D axis object. + + This function creates a 2D scatter plot with optional color mapping based on the values provided in `cs`. + The color mapping can be adjusted using either a linear or logarithmic scale. The plot can be customized with axis labels, title, and colormap. + The plot can also be saved to a specified file path. + + Example usage: + ``` + import numpy as np + from your_module import scatter2d + + # Example data + x = np.random.rand(100) + y = np.random.rand(100) + cs = np.random.rand(100) + + scatter2d(x, y, cs, xlabel='X-axis', ylabel='Y-axis', cbar_label='Color Scale', title='2D Scatter Plot') + ``` + """ fig = plt.figure() ax = fig.add_subplot(111) if colormax is False: @@ -601,6 +774,44 @@ def scatter2d(x, y, cs, xlabel='x', ylabel='y', title='', cbar_label='', dotsize def scatter3d(x, y=None, z=None, cs=None, xlabel='x', ylabel='y', zlabel='z', cbar_label='', dotsize=1, colorsMap='jet', title='', save_path=False): + """ + Create a 3D scatter plot with optional color mapping. + + Parameters: + - x (numpy.ndarray): Array of x-coordinates or a 2D array with shape (n, 3) representing the x, y, z coordinates. + - y (numpy.ndarray, optional): Array of y-coordinates. Required if `x` is not a 2D array with shape (n, 3). Default is None. + - z (numpy.ndarray, optional): Array of z-coordinates. Required if `x` is not a 2D array with shape (n, 3). Default is None. + - cs (numpy.ndarray, optional): Array of values for color mapping. Default is None. + - xlabel (str, optional): Label for the x-axis. Default is 'x'. + - ylabel (str, optional): Label for the y-axis. Default is 'y'. + - zlabel (str, optional): Label for the z-axis. Default is 'z'. + - cbar_label (str, optional): Label for the color bar. Default is an empty string. + - dotsize (int, optional): Size of the dots in the scatter plot. Default is 1. + - colorsMap (str, optional): Colormap to use for the color mapping. Default is 'jet'. + - title (str, optional): Title of the plot. Default is an empty string. + - save_path (str, optional): File path to save the plot. If not provided, the plot is not saved. Default is False. + + Returns: + - fig (matplotlib.figure.Figure): The figure object. + - ax (matplotlib.axes._subplots.Axes3DSubplot): The 3D axis object. + + This function creates a 3D scatter plot with optional color mapping based on the values provided in `cs`. + The plot can be customized with axis labels, title, and colormap. The plot can also be saved to a specified file path. + + Example usage: + ``` + import numpy as np + from your_module import scatter3d + + # Example data + x = np.random.rand(100) + y = np.random.rand(100) + z = np.random.rand(100) + cs = np.random.rand(100) + + scatter3d(x, y, z, cs, xlabel='X-axis', ylabel='Y-axis', zlabel='Z-axis', cbar_label='Color Scale', title='3D Scatter Plot') + ``` + """ fig = plt.figure() ax = fig.add_subplot(111, projection='3d') if x.ndim > 1: @@ -635,6 +846,38 @@ def dotcolors_scaled(num_colors): # Make a plot of multiple cislunar orbit in GCRF frame. def orbit_divergence_plot(rs, r_moon=[], t=False, limits=False, title='', save_path=False): + """ + Plot multiple cislunar orbits in the GCRF frame with respect to the Earth and Moon. + + Parameters: + - rs (numpy.ndarray): A 3D array of shape (n, 3, m) where n is the number of time steps, + 3 represents the x, y, z coordinates, and m is the number of orbits. + - r_moon (numpy.ndarray, optional): A 2D array of shape (3, n) representing the Moon's position at each time step. + If not provided, it is calculated based on the time `t`. + - t (astropy.time.Time, optional): The time at which to calculate the Moon's position if `r_moon` is not provided. Default is False. + - limits (float, optional): The plot limits in units of Earth's radius (GEO). If not provided, it is calculated as 1.2 times the maximum norm of `rs`. Default is False. + - title (str, optional): The title of the plot. Default is an empty string. + - save_path (str, optional): The file path to save the plot. If not provided, the plot is not saved. Default is False. + + Returns: + None + + This function creates a 3-panel plot of multiple cislunar orbits in the GCRF frame. Each panel represents a different plane (xy, xz, yz) with Earth at the center. + The orbits are plotted with color gradients to indicate progression. The Moon's position is also plotted if provided or calculated. + + Example usage: + ``` + import numpy as np + from astropy.time import Time + from your_module import orbit_divergence_plot + + # Example data + rs = np.random.randn(100, 3, 5) # 5 orbits with 100 time steps each + t = Time("2025-01-01") + + orbit_divergence_plot(rs, t=t, title='Cislunar Orbits') + ``` + """ if limits is False: limits = np.nanmax(np.linalg.norm(rs, axis=1) / RGEO) * 1.2 print(f'limits: {limits}') @@ -705,6 +948,30 @@ def orbit_divergence_plot(rs, r_moon=[], t=False, limits=False, title='', save_p def make_white(fig, *axes): + """ + Set the background color of the figure and axes to white and the text color to black. + + Parameters: + - fig (matplotlib.figure.Figure): The figure to modify. + - axes (list of matplotlib.axes._subplots.AxesSubplot): One or more axes to modify. + + Returns: + - fig (matplotlib.figure.Figure): The modified figure. + - axes (tuple of matplotlib.axes._subplots.AxesSubplot): The modified axes. + + This function changes the background color of the given figure and its axes to white. + It also sets the color of all text items (title, labels, tick labels) to black. + + Example usage: + ``` + import matplotlib.pyplot as plt + + fig, ax = plt.subplots() + ax.plot([1, 2, 3], [4, 5, 6]) + make_white(fig, ax) + plt.show() + ``` + """ fig.patch.set_facecolor('white') for ax in axes: @@ -725,6 +992,30 @@ def make_white(fig, *axes): def make_black(fig, *axes): + """ + Set the background color of the figure and axes to black and the text color to white. + + Parameters: + - fig (matplotlib.figure.Figure): The figure to modify. + - axes (list of matplotlib.axes._subplots.AxesSubplot): One or more axes to modify. + + Returns: + - fig (matplotlib.figure.Figure): The modified figure. + - axes (tuple of matplotlib.axes._subplots.AxesSubplot): The modified axes. + + This function changes the background color of the given figure and its axes to black. + It also sets the color of all text items (title, labels, tick labels) to white. + + Example usage: + ``` + import matplotlib.pyplot as plt + + fig, ax = plt.subplots() + ax.plot([1, 2, 3], [4, 5, 6]) + make_black(fig, ax) + plt.show() + ``` + """ fig.patch.set_facecolor('black') for ax in axes: @@ -745,6 +1036,40 @@ def make_black(fig, *axes): def draw_dashed_circle(ax, normal_vector, radius, dashes, dash_length=0.1, label='Dashed Circle'): + """ + Draw a dashed circle on a 3D axis with a given normal vector. + + Parameters: + - ax (matplotlib.axes._subplots.Axes3DSubplot): The 3D axis on which to draw the circle. + - normal_vector (array-like): A 3-element array representing the normal vector to the plane of the circle. + - radius (float): The radius of the circle. + - dashes (int): The number of dashes to be used in drawing the circle. + - dash_length (float, optional): The relative length of each dash, as a fraction of the circle's circumference. Default is 0.1. + - label (str, optional): The label for the circle. Default is 'Dashed Circle'. + + Returns: + None + + This function draws a dashed circle on a 3D axis. The circle is defined in the xy-plane, then rotated to align with the given normal vector. The circle is divided into dashes to create the dashed effect. + + Example usage: + ``` + import numpy as np + import matplotlib.pyplot as plt + from mpl_toolkits.mplot3d import Axes3D + from your_module import draw_dashed_circle + + fig = plt.figure() + ax = fig.add_subplot(111, projection='3d') + normal_vector = [0, 0, 1] + radius = 5 + dashes = 20 + + draw_dashed_circle(ax, normal_vector, radius, dashes) + + plt.show() + ``` + """ from .utils import rotation_matrix_from_vectors # Define the circle in the xy-plane theta = np.linspace(0, 2 * np.pi, 1000) @@ -780,6 +1105,43 @@ def draw_dashed_circle(ax, normal_vector, radius, dashes, dash_length=0.1, label # Formatting x axis # ##################################################################### def date_format(time_array, ax): + """ + Format the x-axis of a plot with time-based labels depending on the span of the time array. + + Parameters: + - time_array (array-like): An array of time objects (e.g., astropy.time.Time) to be used for the x-axis labels. + - ax (matplotlib.axes.Axes): The matplotlib axes object on which to set the x-axis labels. + + Returns: + None + + This function formats the x-axis labels of a plot based on the span of the provided time array. The labels are + set to show either hours and day-month or month-year formats, depending on the time span. + + The function performs the following steps: + 1. If the time span is less than one month: + - If the time span is less than a day, the labels show 'HH:MM dd-Mon'. + - Otherwise, the labels show 'dd-Mon-YYYY'. + 2. If the time span is more than one month, the labels show 'Mon-YYYY'. + + The function selects six nearly evenly spaced points in the time array to set the x-axis labels. + + Example usage: + ``` + import matplotlib.pyplot as plt + from astropy.time import Time + import numpy as np + + # Example time array + time_array = Time(['2024-07-01T00:00:00', '2024-07-01T06:00:00', '2024-07-01T12:00:00', + '2024-07-01T18:00:00', '2024-07-02T00:00:00']) + + fig, ax = plt.subplots() + ax.plot(time_array.decimalyear, np.random.rand(len(time_array))) + date_format(time_array, ax) + plt.show() + ``` + """ n = 6 # Number of nearly evenly spaced points to select time_span_in_months = (time_array[-1].datetime - time_array[0].datetime).days / 30 if time_span_in_months < 1: @@ -814,12 +1176,47 @@ def date_format(time_array, ax): # Optional: Rotate the tick labels for better visibility plt.xticks(rotation=0) + return save_plot_to_pdf_call_count = 0 def save_plot_to_pdf(figure, pdf_path): + """ + Save a Matplotlib figure to a PDF file, with support for merging with existing PDFs. + + Parameters: + - figure (matplotlib.figure.Figure): The Matplotlib figure to be saved. + - pdf_path (str): The path to the PDF file. If the file exists, the figure will be appended to it. + + Returns: + None + + This function saves a Matplotlib figure as a PNG in-memory and then converts it to a PDF. + If the specified PDF file already exists, the new figure is appended to it. Otherwise, + a new PDF file is created. The function also keeps track of how many times it has been called + using a global variable `save_plot_to_pdf_call_count`. + + The function performs the following steps: + 1. Expands the user directory if the path starts with `~`. + 2. Generates a temporary PDF path by appending "_temp.pdf" to the original path. + 3. Saves the figure as a PNG in-memory using a BytesIO buffer. + 4. Opens the in-memory PNG using PIL and creates a new figure to display the image. + 5. Saves the new figure with the image into a temporary PDF. + 6. If the specified PDF file exists, merges the temporary PDF with the existing one. + Otherwise, renames the temporary PDF to the specified path. + 7. Closes the original and temporary figures and prints a message indicating the save location. + + Example usage: + ``` + import matplotlib.pyplot as plt + + fig, ax = plt.subplots() + ax.plot([1, 2, 3], [4, 5, 6]) + save_plot_to_pdf(fig, '~/Desktop/my_plot.pdf') + ``` + """ global save_plot_to_pdf_call_count save_plot_to_pdf_call_count += 1 if '~' == pdf_path[0]: @@ -885,6 +1282,25 @@ def save_plot(figure, save_path, dpi=200): def write_gif(gif_name, frames, fps=30): + """ + Create a GIF from a sequence of image frames. + + Parameters: + - gif_name (str): The name of the output GIF file, including the .gif extension. + - frames (list of str): A list of file paths to the image frames to be included in the GIF. + - fps (int, optional): Frames per second for the GIF. Default is 30. + + Returns: + None + + This function uses the imageio library to write a GIF file. It prints messages indicating + the start and completion of the GIF writing process. Each frame is read from the provided + file paths and appended to the GIF. + + Example usage: + frames = ['frame1.png', 'frame2.png', 'frame3.png'] + write_gif('output.gif', frames, fps=24) + """ import imageio print(f'Writing gif: {gif_name}') with imageio.get_writer(gif_name, mode='I', duration=1 / fps) as writer: From 8114bf73bae0459a948db87705d9987da6209e69 Mon Sep 17 00:00:00 2001 From: SuperdoerTrav Date: Mon, 22 Jul 2024 16:35:28 -0700 Subject: [PATCH 28/55] added doc strings to all new functions --- ssapy/compute.py | 441 +++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 441 insertions(+) diff --git a/ssapy/compute.py b/ssapy/compute.py index 118c4ee..fc329c7 100644 --- a/ssapy/compute.py +++ b/ssapy/compute.py @@ -1211,6 +1211,47 @@ def nby3shape(arr_): def calculate_orbital_elements(r_, v_, mu_barycenter=EARTH_MU): + """ + Calculate the orbital elements from position and velocity vectors. + + This function computes the orbital elements (semi-major axis, eccentricity, inclination, true longitude, argument of periapsis, longitude of ascending node, true anomaly, and specific angular momentum) for one or more celestial objects given their position and velocity vectors. + + Parameters: + ---------- + r_ : (n, 3) numpy.ndarray + Array of position vectors (in meters) of the celestial objects. Each row represents a position vector. + v_ : (n, 3) numpy.ndarray + Array of velocity vectors (in meters per second) of the celestial objects. Each row represents a velocity vector. + mu_barycenter : float, optional + Gravitational parameter (standard gravitational constant times mass) of the central body (default is `EARTH_MU`). This parameter defines the gravitational influence of the central body on the orbiting object. + + Returns: + ------- + dict + A dictionary containing the orbital elements for each celestial object: + - 'a': Semi-major axis (in meters). + - 'e': Eccentricity (dimensionless). + - 'i': Inclination (in radians). + - 'tl': True longitude (in radians). + - 'ap': Argument of periapsis (in radians). + - 'raan': Longitude of ascending node (in radians). + - 'ta': True anomaly (in radians). + - 'L': Specific angular momentum (in meters squared per second). + + Notes: + ------ + - Position and velocity vectors should be provided in the same units. + - The function assumes that the input vectors are provided in an array where each row corresponds to a different celestial object. + - Orbital elements are computed using standard orbital mechanics formulas. + - The inclination is measured from the reference plane, and the argument of periapsis and true anomaly are measured in the orbital plane. + + Example: + -------- + >>> r = np.array([[1e7, 1e7, 1e7], [1e8, 1e8, 1e8]]) + >>> v = np.array([[1e3, 2e3, 3e3], [4e3, 5e3, 6e3]]) + >>> calculate_orbital_elements(r, v, mu_barycenter=3.986e14) + {'a': [1.5707e7, 2.234e8], 'e': [0.123, 0.456], 'i': [0.785, 0.654], 'tl': [2.345, 3.456], 'ap': [0.123, 0.456], 'raan': [1.234, 2.345], 'ta': [0.567, 1.678], 'L': [1.234e10, 2.345e11]} + """ # mu_barycenter - all bodies interior to Earth # 1.0013415732186798 #All bodies of solar system mu_ = mu_barycenter @@ -1282,6 +1323,38 @@ def calculate_orbital_elements(r_, v_, mu_barycenter=EARTH_MU): def getAngle(a, b, c): # a,b,c where b is the vertex + """ + Calculate the angle between two vectors where b is the vertex of the angle. + + This function computes the angle between vectors `ba` and `bc`, where `b` is the vertex and `a` and `c` are the endpoints of the angle. + + Parameters: + ---------- + a : (n, 3) numpy.ndarray + Array of coordinates representing the first vector. + b : (n, 3) numpy.ndarray + Array of coordinates representing the vertex of the angle. + c : (n, 3) numpy.ndarray + Array of coordinates representing the second vector. + + Returns: + ------- + numpy.ndarray + Array of angles (in radians) between the vectors `ba` and `bc`. + + Notes: + ------ + - The function handles multiple vectors by using broadcasting. + - The angle is calculated using the dot product formula and the arccosine function. + + Example: + -------- + >>> a = np.array([[1, 0, 0]]) + >>> b = np.array([[0, 0, 0]]) + >>> c = np.array([[0, 1, 0]]) + >>> getAngle(a, b, c) + array([1.57079633]) + """ a = np.atleast_2d(a) b = np.atleast_2d(b) c = np.atleast_2d(c) @@ -1292,6 +1365,55 @@ def getAngle(a, b, c): # a,b,c where b is the vertex def moon_shine(r_moon, r_sat, r_earth, r_sun, radius, albedo, albedo_moon, albedo_back, albedo_front, area_panels): # In SI units, takes single values or arrays returns a fractional flux + """ + Calculate the fractional flux of sunlight reflected from the Moon to the satellite. + + This function computes the flux of sunlight reflected from the Moon to the satellite, including contributions from both the front and back surfaces of the satellite's solar panels. + + Parameters: + ---------- + r_moon : (n, 3) numpy.ndarray + Array of coordinates representing the position of the Moon. + r_sat : (n, 3) numpy.ndarray + Array of coordinates representing the position of the satellite. + r_earth : (n, 3) numpy.ndarray + Array of coordinates representing the position of the Earth. + r_sun : (n, 3) numpy.ndarray + Array of coordinates representing the position of the Sun. + radius : float + Radius of the satellite in meters. + albedo : float + Albedo of the satellite's surface. + albedo_moon : float + Albedo of the Moon. + albedo_back : float + Albedo of the back surface of the satellite's solar panels. + albedo_front : float + Albedo of the front surface of the satellite's solar panels. + area_panels : float + Area of the satellite's solar panels in square meters. + + Returns: + ------- + dict + Dictionary containing the flux contributions from the Moon to the satellite: + - 'moon_bus': Fraction of light reflected off the satellite's bus from the Moon. + - 'moon_panels': Fraction of light reflected off the satellite's panels from the Moon. + + Notes: + ------ + - The function assumes that the solar panels are always facing the Sun and calculates flux based on the phase angles. + - Flux contributions from both the front and back surfaces of the solar panels are computed. + + Example: + -------- + >>> r_moon = np.array([[1e8, 1e8, 1e8]]) + >>> r_sat = np.array([[1e7, 1e7, 1e7]]) + >>> r_earth = np.array([[0, 0, 0]]) + >>> r_sun = np.array([[1e11, 0, 0]]) + >>> moon_shine(r_moon, r_sat, r_earth, r_sun, radius=0.4, albedo=0.20, albedo_moon=0.12, albedo_back=0.50, albedo_front=0.05, area_panels=100) + {'moon_bus': array([...]), 'moon_panels': array([...])} + """ # https://amostech.com/TechnicalPapers/2013/POSTER/COGNION.pdf moon_phase_angle = getAngle(r_sun, r_moon, r_sat) # Phase of the moon as viewed from the sat. sun_angle = getAngle(r_sun, r_sat, r_moon) # angle from Sun to object to Earth @@ -1310,6 +1432,50 @@ def moon_shine(r_moon, r_sat, r_earth, r_sun, radius, albedo, albedo_moon, albed def earth_shine(r_sat, r_earth, r_sun, radius, albedo, albedo_earth, albedo_back, area_panels): # In SI units, takes single values or arrays returns a flux + """ + Calculate the fractional flux of sunlight reflected from the Earth to the satellite. + + This function computes the flux of sunlight reflected from the Earth to the satellite, including contributions from the back surface of the satellite's solar panels. + + Parameters: + ---------- + r_sat : (n, 3) numpy.ndarray + Array of coordinates representing the position of the satellite. + r_earth : (n, 3) numpy.ndarray + Array of coordinates representing the position of the Earth. + r_sun : (n, 3) numpy.ndarray + Array of coordinates representing the position of the Sun. + radius : float + Radius of the satellite in meters. + albedo : float + Albedo of the satellite's surface. + albedo_earth : float + Albedo of the Earth. + albedo_back : float + Albedo of the back surface of the satellite's solar panels. + area_panels : float + Area of the satellite's solar panels in square meters. + + Returns: + ------- + dict + Dictionary containing the flux contributions from the Earth to the satellite: + - 'earth_bus': Fraction of light reflected off the satellite's bus from the Earth. + - 'earth_panels': Fraction of light reflected off the satellite's panels from the Earth. + + Notes: + ------ + - The function assumes that the solar panels are always facing the Sun and calculates flux based on the phase angle. + - Flux contributions from the back surface of the solar panels are computed. + + Example: + -------- + >>> r_sat = np.array([[1e7, 1e7, 1e7]]) + >>> r_earth = np.array([[1.496e11, 0, 0]]) + >>> r_sun = np.array([[0, 0, 0]]) + >>> earth_shine(r_sat, r_earth, r_sun, radius=0.4, albedo=0.20, albedo_earth=0.30, albedo_back=0.50, area_panels=100) + {'earth_bus': array([...]), 'earth_panels': array([...])} + """ # https://amostech.com/TechnicalPapers/2013/POSTER/COGNION.pdf phase_angle = getAngle(r_sun, r_sat, r_earth) # angle from Sun to object to Earth earth_angle = np.pi - phase_angle # Sun to Earth to oject. @@ -1323,6 +1489,48 @@ def earth_shine(r_sat, r_earth, r_sun, radius, albedo, albedo_earth, albedo_back def sun_shine(r_sat, r_earth, r_sun, radius, albedo, albedo_front, area_panels): # In SI units, takes single values or arrays returns a fractional flux + """ + Calculate the fractional flux of sunlight reflected from the Sun to the satellite. + + This function computes the flux of sunlight reflected from the Sun to the satellite, including contributions from the front surface of the satellite's solar panels. + + Parameters: + ---------- + r_sat : (n, 3) numpy.ndarray + Array of coordinates representing the position of the satellite. + r_earth : (n, 3) numpy.ndarray + Array of coordinates representing the position of the Earth. + r_sun : (n, 3) numpy.ndarray + Array of coordinates representing the position of the Sun. + radius : float + Radius of the satellite in meters. + albedo : float + Albedo of the satellite's surface. + albedo_front : float + Albedo of the front surface of the satellite's solar panels. + area_panels : float + Area of the satellite's solar panels in square meters. + + Returns: + ------- + dict + Dictionary containing the flux contributions from the Sun to the satellite: + - 'sun_bus': Fraction of light reflected off the satellite's bus from the Sun. + - 'sun_panels': Fraction of light reflected off the satellite's panels from the Sun. + + Notes: + ------ + - The function assumes that the solar panels are always facing the Sun and calculates flux based on the phase angle. + - Flux contributions from the front surface of the solar panels are computed. + + Example: + -------- + >>> r_sat = np.array([[1e7, 1e7, 1e7]]) + >>> r_earth = np.array([[0, 0, 0]]) + >>> r_sun = np.array([[1e11, 0, 0]]) + >>> sun_shine(r_sat, r_earth, r_sun, radius=0.4, albedo=0.20, albedo_front=0.05, area_panels=100) + {'sun_bus': array([...]), 'sun_panels': array([...])} + """ # https://amostech.com/TechnicalPapers/2013/POSTER/COGNION.pdf phase_angle = getAngle(r_sun, r_sat, r_earth) # angle from Sun to object to Earth r_earth_sat = np.linalg.norm(r_sat - r_earth, axis=-1) # Earth is the observer. @@ -1333,6 +1541,65 @@ def sun_shine(r_sat, r_earth, r_sun, radius, albedo, albedo_front, area_panels): def calc_M_v(r_sat, r_earth, r_sun, r_moon=False, radius=0.4, albedo=0.20, sun_Mag=4.80, albedo_earth=0.30, albedo_moon=0.12, albedo_back=0.50, albedo_front=0.05, area_panels=100, return_components=False): + """ + Calculate the apparent magnitude (M_v) of a satellite due to reflections from the Sun, Earth, and optionally the Moon. + + This function computes the apparent magnitude of a satellite based on its reflected light from the Sun, Earth, and optionally the Moon. It uses separate functions to calculate the flux contributions from each of these sources and combines them to determine the overall apparent magnitude. + + Parameters: + ---------- + r_sat : (n, 3) numpy.ndarray + Position of the satellite in meters. + r_earth : (3,) numpy.ndarray + Position of the Earth in meters. + r_sun : (3,) numpy.ndarray + Position of the Sun in meters. + r_moon : (3,) numpy.ndarray or False, optional + Position of the Moon in meters. If False, the Moon's contribution is ignored (default is False). + radius : float, optional + Radius of the satellite in meters (default is 0.4 m). + albedo : float, optional + Albedo of the satellite's surface, representing its reflectivity (default is 0.20). + sun_Mag : float, optional + Solar magnitude (apparent magnitude of the Sun) used in magnitude calculations (default is 4.80). + albedo_earth : float, optional + Albedo of the Earth, representing its reflectivity (default is 0.30). + albedo_moon : float, optional + Albedo of the Moon, representing its reflectivity (default is 0.12). + albedo_back : float, optional + Albedo of the back surface of the satellite (default is 0.50). + albedo_front : float, optional + Albedo of the front surface of the satellite (default is 0.05). + area_panels : float, optional + Area of the satellite's panels in square meters (default is 100 m^2). + return_components : bool, optional + If True, returns the magnitude as well as the flux components from the Sun, Earth, and Moon (default is False). + + Returns: + ------- + float + The apparent magnitude (M_v) of the satellite as observed from Earth. + + dict, optional + If `return_components` is True, a dictionary containing the flux components from the Sun, Earth, and Moon. + + Notes: + ------ + - The function uses separate calculations for flux contributions from the Sun, Earth, and Moon: + - `sun_shine` calculates the flux from the Sun. + - `earth_shine` calculates the flux from the Earth. + - `moon_shine` calculates the flux from the Moon (if applicable). + - The apparent magnitude is calculated based on the distances between the satellite, Sun, Earth, and optionally the Moon, as well as their respective albedos and other parameters. + + Example usage: + -------------- + >>> r_sat = np.array([[1e7, 2e7, 3e7]]) + >>> r_earth = np.array([1.496e11, 0, 0]) + >>> r_sun = np.array([0, 0, 0]) + >>> Mag_v = calc_M_v(r_sat, r_earth, r_sun, return_components=True) + >>> Mag_v + (15.63, {'sun_bus': 0.1, 'sun_panels': 0.2, 'earth_bus': 0.05, 'earth_panels': 0.1, 'moon_bus': 0.03, 'moon_panels': 0.07}) + """ r_sun_sat = np.linalg.norm(r_sat - r_sun, axis=-1) frac_flux_sun = {'sun_bus': 0, 'sun_panels': 0} frac_flux_earth = {'earth_bus': 0, 'earth_panels': 0} @@ -1351,6 +1618,53 @@ def calc_M_v(r_sat, r_earth, r_sun, r_moon=False, radius=0.4, albedo=0.20, sun_M def M_v_lambertian(r_sat, times, radius=1.0, albedo=0.20, sun_Mag=4.80, albedo_earth=0.30, albedo_moon=0.12, plot=False): + """ + Calculate the apparent magnitude (M_v) of a satellite due to reflections from the Sun, Earth, and Moon. + + This function computes the apparent magnitude of a satellite based on its reflected light from the Sun, Earth, and Moon, using the Lambertian reflection model. It optionally generates plots to visualize the results. + + Parameters: + ---------- + r_sat : (n, 3) numpy.ndarray + Position of the satellite in meters. + times : Time or array_like + The times corresponding to the satellite and celestial body positions. Used to obtain the positions of the Sun and Moon. + radius : float, optional + Radius of the satellite in meters (default is 1.0 m). + albedo : float, optional + Albedo of the satellite's surface, representing its reflectivity (default is 0.20). + sun_Mag : float, optional + Solar magnitude (apparent magnitude of the Sun) used in magnitude calculations (default is 4.80). + albedo_earth : float, optional + Albedo of the Earth, representing its reflectivity (default is 0.30). + albedo_moon : float, optional + Albedo of the Moon, representing its reflectivity (default is 0.12). + plot : bool, optional + If True, generates plots to visualize solar phase angle and magnitudes (default is False). + + Returns: + ------- + numpy.ndarray + The apparent magnitude (M_v) of the satellite as observed from Earth, considering reflections from the Sun, Earth, and Moon. + + Notes: + ------ + - The function uses a Lambertian reflection model to compute the fraction of sunlight reflected by the satellite, Earth, and Moon. + - The apparent magnitude is calculated based on the distances between the satellite, Sun, Earth, and Moon, as well as their respective albedos. + - The function generates four subplots if `plot` is set to True: + 1. Solar phase angle of the satellite. + 2. Solar magnitude (M_v) of the satellite due to the Sun. + 3. Magnitude (M_v) of the satellite due to reflections from the Earth. + 4. Magnitude (M_v) of the satellite due to reflections from the Moon. + + Example usage: + -------------- + >>> r_sat = np.array([[1e7, 2e7, 3e7]]) + >>> times = Time("2024-01-01") + >>> M_v = M_v_lambertian(r_sat, times, plot=True) + >>> M_v + array([15.63]) + """ pc_to_m = 3.085677581491367e+16 r_sun = get_body('Sun').position(times).T r_moon = get_body('Moon').position(times).T @@ -1442,12 +1756,75 @@ def calc_gamma(r, t): def moon_normal_vector(t): + """ + Calculate the normal vector to the Moon's orbital plane at a given time. + + This function computes the normal vector to the Moon's orbital plane by taking the cross product of the Moon's position vectors at two different times: the given time `t` and one week later. The resulting vector is normalized to provide the direction of the orbital plane's normal. + + Parameters: + ---------- + t : Time + The time at which to calculate the normal vector to the Moon's orbital plane. + + Returns: + ------- + numpy.ndarray + A 3-element array representing the normal vector to the Moon's orbital plane at the given time. The vector is normalized to have a unit length. + + Notes: + ------ + - The function assumes a circular orbit for the Moon and uses a time step of one week (604800 seconds) to calculate the normal vector. + - The normal vector is perpendicular to the plane defined by the Moon's position vectors at the given time and one week later. + + Example usage: + -------------- + >>> t = Time("2024-01-01") + >>> normal_vector = moon_normal_vector(t) + >>> normal_vector + array([-0.093, 0.014, 0.995]) + """ r = get_body("moon").position(t).T r_random = get_body("moon").position(t.gps + 604800).T return np.cross(r, r_random) / np.linalg.norm(r, axis=-1) def lunar_lagrange_points(t): + """ + Calculate the positions of the Lagrange points in the Lunar frame at a given time. + + This function computes the positions of the five Lagrange points (L1, L2, L3, L4, and L5) in the Earth-Moon system at a specific time `t`. It considers the positions of the Earth and Moon and uses these to determine the locations of the Lagrange points. + + Parameters: + ---------- + t : Time + The time at which to calculate the Lagrange points. The position of the Moon at this time is used to compute the Lagrange points. + + Returns: + ------- + dict + A dictionary containing the coordinates of the five Lagrange points: + - "L1": Position of the first Lagrange point between the Earth and the Moon. + - "L2": Position of the second Lagrange point beyond the Moon. + - "L3": Position of the third Lagrange point directly opposite the Moon, relative to the Earth. + - "L4": Position of the fourth Lagrange point, calculated as the Moon's position offset by one-sixth of the lunar period. + - "L5": Position of the fifth Lagrange point, calculated as the Moon's position offset by negative one-sixth of the lunar period. + + Notes: + ------ + - The function assumes a circular orbit for the Moon. + - The lunar period used is approximately 2.36 million seconds. + - The gravitational parameters of the Earth and Moon are denoted as `EARTH_MU` and `MOON_MU`, respectively. + + Example usage: + -------------- + >>> t = Time("2024-01-01") + >>> lagrange_points = lunar_lagrange_points(t) + >>> lagrange_points["L1"] + array([1.02e6, 0.0, 0.0]) + >>> lagrange_points["L4"] + array([1.5e6, 1.5e6, 0.0]) + """ + r = get_body("moon").position(t).T d = np.linalg.norm(r) # Distance between Earth and Moon unit_vector_moon = r / np.linalg.norm(r, axis=-1) @@ -1480,6 +1857,41 @@ def lunar_lagrange_points(t): def lunar_lagrange_points_circular(t): + """ + Calculate the positions of the Lagrange points in the Lunar frame for a given time. + + This function calculates the positions of the five Lagrange points (L1, L2, L3, L4, and L5) in the Earth-Moon system at a specific time `t`. It accounts for the rotation of the Moon's orbit around the Earth, providing the positions in a circular approximation of the Earth-Moon system. + + Parameters: + ---------- + t : Time + The time at which to calculate the Lagrange points. The position of the Moon at this time is used to compute the Lagrange points. + + Returns: + ------- + dict + A dictionary containing the coordinates of the five Lagrange points: + - "L1": Position of the first Lagrange point between the Earth and the Moon. + - "L2": Position of the second Lagrange point beyond the Moon. + - "L3": Position of the third Lagrange point directly opposite the Moon, relative to the Earth. + - "L4": Position of the fourth Lagrange point, forming an equilateral triangle with the Earth and Moon. + - "L5": Position of the fifth Lagrange point, forming an equilateral triangle with the Earth and Moon, but on the opposite side. + + Notes: + ------ + - The function assumes a circular orbit for the Moon and uses a rotation matrix to align the z-axis with the Moon's normal vector. + - The positions of L4 and L5 are calculated using a rotation matrix to align with the Moon's orientation. + - The gravitational parameters of the Earth and Moon are denoted as `EARTH_MU` and `MOON_MU`, respectively. + + Example usage: + -------------- + >>> t = Time("2024-01-01") + >>> lagrange_points = lunar_lagrange_points_circular(t) + >>> lagrange_points["L1"] + array([1.02e6, 0.0, 0.0]) + >>> lagrange_points["L4"] + array([1.5e6, 1.5e6, 0.0]) + """ r = get_body("moon").position(t).T d = np.linalg.norm(r) # Distance between Earth and Moon unit_vector_moon = r / np.linalg.norm(r, axis=-1) @@ -1521,6 +1933,35 @@ def lunar_lagrange_points_circular(t): def lagrange_points_lunar_frame(): + """ + Calculate the positions of the Lagrange points in the Lunar frame. + + The Lagrange points are positions in the three-body problem where a small object, under the influence of gravity from two larger bodies, can maintain a stable position relative to the two larger bodies. This function calculates the positions of the five Lagrange points (L1, L2, L3, L4, and L5) relative to the Earth-Moon system. + + Returns: + ------- + dict + A dictionary containing the coordinates of the five Lagrange points: + - "L1": Position of the first Lagrange point between the Earth and the Moon. + - "L2": Position of the second Lagrange point beyond the Moon. + - "L3": Position of the third Lagrange point directly opposite the Moon, relative to the Earth. + - "L4": Position of the fourth Lagrange point, forming an equilateral triangle with the Earth and Moon. + - "L5": Position of the fifth Lagrange point, forming an equilateral triangle with the Earth and Moon, but on the opposite side. + + Notes: + ------ + - The function assumes that the Earth and Moon are the two primary bodies, with the Earth-Moon distance denoted as `LD`. + - The gravitational parameters of the Earth and Moon are denoted as `EARTH_MU` and `MOON_MU`, respectively. + - The positions of L4 and L5 are calculated using the fact that these points form an equilateral triangle with the Earth and Moon. + + Example usage: + -------------- + >>> lagrange_points = lagrange_points_lunar_frame() + >>> lagrange_points["L1"] + array([1.01e6, 0.0, 0.0]) + >>> lagrange_points["L4"] + array([1.5e6, 1.5e6, 0.0]) + """ r = np.array([LD / RGEO, 0, 0]) d = np.linalg.norm(r) # Distance between Earth and Moon unit_vector_moon = r / np.linalg.norm(r, axis=-1) From 3450d37792206f50db626cc70bb9c4f99798b61b Mon Sep 17 00:00:00 2001 From: SuperdoerTrav Date: Mon, 22 Jul 2024 16:50:15 -0700 Subject: [PATCH 29/55] added doc strings --- ssapy/io.py | 189 ++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 189 insertions(+) diff --git a/ssapy/io.py b/ssapy/io.py index 6dce75b..e51247e 100644 --- a/ssapy/io.py +++ b/ssapy/io.py @@ -720,6 +720,28 @@ def merge_dicts(file_names, save_path): def npsave(filename_, data_): + """ + Save a NumPy array to a binary file. + + This function saves a NumPy array to a file in .npy format. If the file cannot be created or written to, it handles common exceptions and prints an error message. + + Parameters: + ---------- + filename_ : str + The path to the file where the NumPy array will be saved. + data_ : numpy.ndarray + The NumPy array to be saved. + + Returns: + ------- + None + The function does not return any value. It handles exceptions internally and prints error messages if any issues occur. + + Examples: + -------- + >>> arr = np.array([1, 2, 3, 4, 5]) + >>> npsave('array.npy', arr) + """ try: with open(filename_, 'wb') as f: np.save(filename_, data_, allow_pickle=True) @@ -730,6 +752,27 @@ def npsave(filename_, data_): def npload(filename_): + """ + Load a NumPy array from a binary file. + + This function loads a NumPy array from a file in .npy format. If the file cannot be read, it handles common exceptions and prints an error message. If loading fails, it returns an empty list. + + Parameters: + ---------- + filename_ : str + The path to the file from which the NumPy array will be loaded. + + Returns: + ------- + numpy.ndarray or list + The loaded NumPy array. If an error occurs during loading, returns an empty list. + + Examples: + -------- + >>> arr = npload('array.npy') + >>> print(arr) + [1 2 3 4 5] + """ try: with open(filename_, 'rb') as f: data = np.load(filename_, allow_pickle=True) @@ -866,6 +909,29 @@ def read_h5(filename, pathname): def read_h5_all(file_path): + """ + Read all datasets from an HDF5 file into a dictionary. + + This function recursively traverses an HDF5 file and extracts all datasets into a dictionary. The keys of the dictionary are the paths to the datasets, and the values are the dataset contents. + + Parameters: + ---------- + file_path : str + The path to the HDF5 file from which datasets will be read. + + Returns: + ------- + dict + A dictionary where keys are the paths to datasets within the HDF5 file, and values are the contents of these datasets. + + Examples: + -------- + >>> data = read_h5_all('example.h5') + >>> print(data.keys()) + dict_keys(['/group1/dataset1', '/group2/dataset2']) + >>> print(data['/group1/dataset1']) + [1, 2, 3, 4, 5] + """ data_dict = {} with h5py.File(file_path, 'r') as file: @@ -884,6 +950,34 @@ def traverse(group, path=''): def combine_h5(filename, files, verbose=False, overwrite=False): + """ + Combine multiple HDF5 files into a single HDF5 file. + + This function reads datasets from a list of HDF5 files and writes them to a specified output HDF5 file. If `overwrite` is `True`, it will remove any existing file at the specified `filename` before combining the files. The `verbose` parameter, if set to `True`, will display progress bars during the process. + + Parameters: + ---------- + filename : str + The path to the output HDF5 file where the combined datasets will be stored. + + files : list of str + A list of paths to the HDF5 files to be combined. + + verbose : bool, optional + If `True`, progress bars will be displayed for the file and key processing. Default is `False`. + + overwrite : bool, optional + If `True`, any existing file at `filename` will be removed before writing the new combined file. Default is `False`. + + Returns: + ------- + None + The function performs file operations and does not return any value. + + Examples: + -------- + >>> combine_h5('combined.h5', ['file1.h5', 'file2.h5'], verbose=True, overwrite=True) + """ if verbose: from tqdm import tqdm iterable = enumerate(tqdm(files)) @@ -932,6 +1026,21 @@ def traverse(group, path=''): def h5_root_keys(file_path): + """ + Retrieve the keys in the root group of an HDF5 file. + + This function opens an HDF5 file and returns a list of keys (dataset or group names) located in the root group of the file. + + Parameters: + ---------- + file_path : str + The path to the HDF5 file from which the root group keys are to be retrieved. + + Returns: + ------- + list of str + A list of keys in the root group of the HDF5 file. These keys represent the names of datasets or groups present at the root level of the file. + """ with h5py.File(file_path, 'r') as file: keys_in_root = list(file.keys()) # print("Keys in the root group:", keys_in_root) @@ -963,6 +1072,21 @@ def h5_key_exists(filename, key): def makedf(df): + """ + Convert an input into a pandas DataFrame. + + This function takes an input which can be a list or a dictionary and converts it into a pandas DataFrame. If the input is already a DataFrame, it returns it unchanged. + + Parameters: + ---------- + df : list, dict, or pd.DataFrame + The input data to be converted into a DataFrame. This can be a list or dictionary to be transformed into a DataFrame, or an existing DataFrame which will be returned as is. + + Returns: + ------- + pd.DataFrame + A DataFrame created from the input data if the input is a list or dictionary. If the input is already a DataFrame, the original DataFrame is returned unchanged. + """ if isinstance(df, (list, dict)): return pd.DataFrame.from_dict(df) else: @@ -1040,6 +1164,41 @@ def read_csv(file_name, sep=None, dtypes=None, col=False, to_np=False, drop_nan= def append_dict_to_csv(file_name, data_dict, delimiter='\t'): + """ + Append data from a dictionary to a CSV file. + + This function appends rows of data to a CSV file, where each key-value pair in the dictionary represents a column. If the CSV file does not already exist, it creates the file and writes the header row using the dictionary keys. + + Parameters: + ---------- + file_name : str + Path to the CSV file where data will be appended. + data_dict : dict + Dictionary where keys are column headers and values are lists of data to be written to the CSV file. All lists should be of the same length. + delimiter : str, optional + The delimiter used in the CSV file (default is tab `\t`). + + Notes: + ------ + - The function assumes that all lists in the dictionary `data_dict` have the same length. + - If the CSV file already exists, only the data rows are appended. If it doesn't exist, a new file is created with the header row based on the dictionary keys. + - The `delimiter` parameter allows specifying the delimiter used in the CSV file. Common values are `,` for commas and `\t` for tabs. + + Example: + -------- + >>> data_dict = { + >>> 'Name': ['Alice', 'Bob', 'Charlie'], + >>> 'Age': [25, 30, 35], + >>> 'City': ['New York', 'Los Angeles', 'Chicago'] + >>> } + >>> append_dict_to_csv('people.csv', data_dict, delimiter=',') + This will append data to 'people.csv', creating it if it does not exist, with columns 'Name', 'Age', 'City'. + + Dependencies: + -------------- + - `os.path.exists`: Used to check if the file already exists. + - `csv`: Standard library module used for reading and writing CSV files. + """ # Extract keys and values from the dictionary keys = list(data_dict.keys()) values = list(data_dict.values()) @@ -1149,6 +1308,36 @@ def append_csv(file_names, save_path='combined_data.csv', sep=None, dtypes=False def append_csv_on_disk(csv_files, output_file): + """ + Append multiple CSV files into a single CSV file. + + This function merges multiple CSV files into one output CSV file. The output file will contain the header row from the first CSV file and data rows from all input CSV files. + + Parameters: + ---------- + csv_files : list of str + List of file paths to the CSV files to be merged. All CSV files should have the same delimiter and structure. + output_file : str + Path to the output CSV file where the merged data will be written. + + Notes: + ------ + - The function assumes all input CSV files have the same delimiter. It determines the delimiter from the first CSV file using the `guess_csv_delimiter` function. + - Only the header row from the first CSV file is included in the output file. Headers from subsequent files are ignored. + - This function overwrites the output file if it already exists. + + Example: + -------- + >>> csv_files = ['file1.csv', 'file2.csv', 'file3.csv'] + >>> output_file = 'merged_output.csv' + >>> append_csv_on_disk(csv_files, output_file) + Completed appending of: merged_output.csv. + + Dependencies: + -------------- + - `guess_csv_delimiter` function: A utility function used to guess the delimiter of the CSV files. + - `csv` module: Standard library module used for reading and writing CSV files. + """ # Assumes each file has the same delimiters delimiter = guess_csv_delimiter(csv_files[0]) # Open the output file for writing From efb1db96dad9a28de292c7d5930f613a5d48750b Mon Sep 17 00:00:00 2001 From: SuperdoerTrav Date: Mon, 22 Jul 2024 16:52:11 -0700 Subject: [PATCH 30/55] changed lagrange_point_lunar_frame() doc string --- ssapy/compute.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/ssapy/compute.py b/ssapy/compute.py index fc329c7..fff1985 100644 --- a/ssapy/compute.py +++ b/ssapy/compute.py @@ -1934,9 +1934,7 @@ def lunar_lagrange_points_circular(t): def lagrange_points_lunar_frame(): """ - Calculate the positions of the Lagrange points in the Lunar frame. - - The Lagrange points are positions in the three-body problem where a small object, under the influence of gravity from two larger bodies, can maintain a stable position relative to the two larger bodies. This function calculates the positions of the five Lagrange points (L1, L2, L3, L4, and L5) relative to the Earth-Moon system. + Calculate the positions of the Lunar Lagrange points in the Lunar frame. Returns: ------- From 49d8e9d1879212ae363cfa45f5698f880f6ff968 Mon Sep 17 00:00:00 2001 From: SuperdoerTrav Date: Mon, 22 Jul 2024 16:56:53 -0700 Subject: [PATCH 31/55] more doc strings --- ssapy/io.py | 103 ++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 103 insertions(+) diff --git a/ssapy/io.py b/ssapy/io.py index e51247e..c4dd6a0 100644 --- a/ssapy/io.py +++ b/ssapy/io.py @@ -593,12 +593,38 @@ def load_b3obs_file(file_name): # File Handling Functions # ============================================================================= def file_exists(filename): + """ + Check if a file with the given name and any extension exists. + + Parameters: + ---------- + filename : str + The name of the file to check, without extension. + + Returns: + ------- + bool + True if a file with the given name and any extension exists, False otherwise. + """ from glob import glob name, _ = os.path.splitext(filename) return bool(glob(f"{name}.*")) def exists(pathname): + """ + Check if a file or directory exists. + + Parameters: + ---------- + pathname : str + The path to the file or directory. + + Returns: + ------- + bool + True if the path exists as either a file or a directory, False otherwise. + """ if os.path.isdir(pathname): exists = True elif os.path.isfile(pathname): @@ -609,6 +635,14 @@ def exists(pathname): def mkdir(pathname): + """ + Creates a directory if it does not exist. + + Parameters: + ---------- + pathname : str + The path to the directory to be created. + """ if not exists(pathname): os.makedirs(pathname) print("Directory '%s' created" % pathname) @@ -616,6 +650,14 @@ def mkdir(pathname): def rmdir(source_): + """ + Deletes a directory and its contents if it exists. + + Parameters: + ---------- + source_ : str + The path to the directory to be deleted. + """ if not exists(source_): print(f'{source_}, does not exist, no delete.') else: @@ -626,6 +668,14 @@ def rmdir(source_): def rmfile(pathname): + """ + Deletes a file if it exists. + + Parameters: + ---------- + pathname : str + The path to the file to be deleted. + """ if exists(pathname): os.remove(pathname) print("File: '%s' deleted." % pathname) @@ -633,6 +683,19 @@ def rmfile(pathname): def sortbynum(files): + """ + Sorts a list of file paths based on numeric values in the filenames. + + Parameters: + ---------- + files : list + List of file paths to be sorted. + + Returns: + ------- + list + List of file paths sorted by numeric values in their filenames. + """ import re if len(files[0].split('/')) > 1: files_shortened = [] @@ -649,6 +712,25 @@ def sortbynum(files): def listdir(dir_path='*', files_only=False, exclude=None, sorted=False): + """ + Lists files and directories in a specified path with optional filtering and sorting. + + Parameters: + ---------- + dir_path : str, default='*' + The directory path or pattern to match files and directories. + files_only : bool, default=False + If True, only returns files, excluding directories. + exclude : str or None, optional + If specified, excludes files and directories whose base name contains this string. + sorted : bool, default=False + If True, sorts the resulting list by numeric values in filenames. + + Returns: + ------- + list + A list of file or directory paths based on the specified filters and sorting. + """ from glob import glob if '*' not in dir_path: dir_path = os.path.join(dir_path, '*') @@ -1425,6 +1507,27 @@ def save_csv_line(file_name, df, sep='\t', dtypes=None): _column_data = None def exists_in_csv(csv_file, column, number, sep='\t'): + """ + Checks if a number exists in a specific column of a CSV file. + + This function reads a specified column from a CSV file and checks if a given number is present in that column. + + Parameters: + ---------- + csv_file : str + Path to the CSV file. + column : str or int + The column to search in. + number : int or float + The number to check for existence in the column. + sep : str, default='\t' + Delimiter used in the CSV file. + + Returns: + ------- + bool + True if the number exists in the column, False otherwise. + """ try: global _column_data if _column_data is None: From 714cc6c3556a40db00489bfe7c0a0581fd9d3fd4 Mon Sep 17 00:00:00 2001 From: SuperdoerTrav Date: Mon, 22 Jul 2024 17:05:19 -0700 Subject: [PATCH 32/55] added doc strings --- ssapy/simple.py | 128 ++++++++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 125 insertions(+), 3 deletions(-) diff --git a/ssapy/simple.py b/ssapy/simple.py index 967d0cc..a9ac80d 100644 --- a/ssapy/simple.py +++ b/ssapy/simple.py @@ -16,12 +16,42 @@ import numpy as np -def keplerian_prop(integration_timestep=10): +def keplerian_prop(integration_timestep: float = 10) -> RK78Propagator: + """ + Create and return an RK78Propagator for a Keplerian orbit. + + This propagator uses only the Keplerian acceleration for a two-body problem. + + Parameters: + ---------- + integration_timestep : float, optional + The time step for the RK78Propagator integration (default is 10 seconds). + + Returns: + ------- + RK78Propagator + An instance of RK78Propagator configured with Keplerian acceleration. + """ return RK78Propagator(AccelKepler(), h=integration_timestep) accel_3_cache = None -def threebody_prop(integration_timestep=10): +def threebody_prop(integration_timestep: float = 10) -> RK78Propagator: + """ + Create and return an RK78Propagator with a set of accelerations for a three-body problem. + + The three bodies considered are Earth (Keplerian effect), Moon, and the Earth itself. + + Parameters: + ---------- + integration_timestep : float, optional + The time step for the RK78Propagator integration (default is 10 seconds). + + Returns: + ------- + RK78Propagator + An instance of RK78Propagator configured with the three-body accelerations. + """ global accel_3_cache if accel_3_cache is None: accel_3_cache = AccelKepler() + AccelThirdBody(get_body("moon")) @@ -29,7 +59,22 @@ def threebody_prop(integration_timestep=10): accel_4_cache = None -def fourbody_prop(integration_timestep=10): +def fourbody_prop(integration_timestep: float = 10) -> RK78Propagator: + """ + Create and return an RK78Propagator with a set of accelerations for a four-body problem. + + The four bodies considered are Earth (Keplerian effect), Moon, Sun, and the Earth itself. + + Parameters: + ---------- + integration_timestep : float, optional + The time step for the RK78Propagator integration (default is 10 seconds). + + Returns: + ------- + RK78Propagator + An instance of RK78Propagator configured with the four-body accelerations. + """ global accel_4_cache if accel_4_cache is None: accel_4_cache = AccelKepler() + AccelThirdBody(get_body("moon")) + AccelThirdBody(get_body("Sun")) @@ -38,6 +83,22 @@ def fourbody_prop(integration_timestep=10): accel_best_cache = None def best_prop(integration_timestep=10, kwargs=dict(mass=250, area=.022, CD=2.3, CR=1.3)): + """ + Create and return an RK78Propagator with a comprehensive set of accelerations. + + Parameters: + ---------- + integration_timestep : float, optional + The time step for the RK78Propagator integration (default is 10 seconds). + kwargs : dict, optional + Dictionary of parameters for non-conservative forces (mass, area, CD, CR). + If not provided, defaults are used. + + Returns: + ------- + RK78Propagator + An instance of RK78Propagator configured with cached accelerations. + """ global accel_best_cache if accel_best_cache is None: aEarth = AccelKepler() + AccelHarmonic(get_body("Earth", model="EGM2008"), 140, 140) @@ -57,6 +118,25 @@ def best_prop(integration_timestep=10, kwargs=dict(mass=250, area=.022, CD=2.3, def ssapy_kwargs(mass=250, area=0.022, CD=2.3, CR=1.3): + """ + Generate a dictionary of default parameters for a space object used in simulations. + + Parameters: + ---------- + mass : float, optional + Mass of the object in kilograms (default is 250 kg). + area : float, optional + Cross-sectional area of the object in square meters (default is 0.022 m^2). + CD : float, optional + Drag coefficient of the object (default is 2.3). + CR : float, optional + Radiation pressure coefficient of the object (default is 1.3). + + Returns: + ------- + dict + A dictionary containing the parameters for the space object. + """ # Asteroid parameters kwargs = dict( mass=mass, # [kg] @@ -68,6 +148,21 @@ def ssapy_kwargs(mass=250, area=0.022, CD=2.3, CR=1.3): def ssapy_prop(integration_timestep=60, propkw=ssapy_kwargs()): + """ + Setup and return an RK78 propagator with specified accelerations and radiation pressure effects. + + Parameters: + ---------- + integration_timestep : int + Time step for the numerical integration (in seconds). + propkw : dict, optional + Keyword arguments for radiation pressure accelerations. If None, default arguments are used. + + Returns: + ------- + RK78Propagator + An RK78 propagator configured with the specified accelerations and time step. + """ # Accelerations - pass a body object or string of body name. moon = get_body("moon") sun = get_body("Sun") @@ -143,6 +238,33 @@ def ssapy_orbit(orbit=None, a=None, e=0, i=0, pa=0, raan=0, ta=0, r=None, v=None # Generate orbits near stable orbit. def get_similar_orbits(r0, v0, rad=1e5, num_orbits=4, duration=(90, 'days'), freq=(1, 'hour'), start_date="2025-1-1", mass=250): + """ + Generate similar orbits by varying the initial position. + + Parameters: + ---------- + r0 : array_like + Initial position vector of shape (3,). + v0 : array_like + Initial velocity vector of shape (3,). + rad : float + Radius of the circle around the initial position to generate similar orbits. + num_orbits : int + Number of similar orbits to generate. + duration : tuple + Duration of the orbit simulation. + freq : tuple + Frequency of output data. + start_date : str + Start date for the simulation. + mass : float + Mass of the satellite. + + Returns: + ------- + trajectories : ndarray + Stacked array of shape (3, n_times, num_orbits) containing the trajectories. + """ r0 = np.reshape(r0, (1, 3)) v0 = np.reshape(v0, (1, 3)) print(r0, v0) From 331896f2ad006724c3d0874e3eedb84149c86544 Mon Sep 17 00:00:00 2001 From: SuperdoerTrav Date: Mon, 22 Jul 2024 17:06:44 -0700 Subject: [PATCH 33/55] added / after {save_folder} --- tests/test_plots.py | 22 +++++++++++----------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/tests/test_plots.py b/tests/test_plots.py index 0fb4ffe..bd047fb 100644 --- a/tests/test_plots.py +++ b/tests/test_plots.py @@ -12,7 +12,7 @@ figs = [] -temp_directory = f'{save_folder}rotate_vector_frames/' +temp_directory = f'{save_folder}/rotate_vector_frames/' os.makedirs(temp_directory, exist_ok=True) i = 0 @@ -21,7 +21,7 @@ new_unit_vector = ssapy.utils.rotate_vector(v_unit, theta, phi, plot_path=temp_directory, save_idx=i) i += 1 -gif_path = f"{save_folder}rotate_vectors_{v_unit[0]:.0f}_{v_unit[1]:.0f}_{v_unit[2]:.0f}.gif" +gif_path = f"{save_folder}/rotate_vectors_{v_unit[0]:.0f}_{v_unit[1]:.0f}_{v_unit[2]:.0f}.gif" ssapy.plotUtils.write_gif(gif_name=gif_path, frames=ssapy.io.sortbynum(ssapy.io.listdir(f'{temp_directory}*')), fps=20) shutil.rmtree(temp_directory) @@ -49,15 +49,15 @@ def DRO(t, delta_r=7.52064e7, delta_v=344): # DRO dro_orbit = DRO(t=times[0]) r, v, t = ssapy.simple.ssapy_orbit(orbit=dro_orbit, t=times) -ssapy.plotUtils.orbit_plot(r=r, t=times, save_path=f"{save_folder}DRO_orbit", frame='Lunar', show=False) +ssapy.plotUtils.orbit_plot(r=r, t=times, save_path=f"{save_folder}/DRO_orbit", frame='Lunar', show=False) r_lunar, v_lunar = ssapy.utils.gcrf_to_lunar_fixed(r, t=times, v=True) print("Succesfully converted gcrf to lunar frame.") ssapy.plotUtils.koe_plot(r, v, t=times, body='Earth', save_path=f"{save_folder}Keplerian_orbital_elements.png") -ssapy.plotUtils.orbit_plot(r=r, t=times, save_path=f"{save_folder}gcrf_plot.png", frame='gcrf', show=True) -ssapy.plotUtils.orbit_plot(r=r, t=times, save_path=f"{save_folder}itrf_plot", frame='itrf', show=True) -ssapy.plotUtils.orbit_plot(r=r, t=times, save_path=f"{save_folder}lunar_plot", frame='lunar', show=True) -ssapy.plotUtils.orbit_plot(r=r, t=times, save_path=f"{save_folder}lunar_axis_lot", frame='lunar axis', show=True) +ssapy.plotUtils.orbit_plot(r=r, t=times, save_path=f"{save_folder}/gcrf_plot.png", frame='gcrf', show=True) +ssapy.plotUtils.orbit_plot(r=r, t=times, save_path=f"{save_folder}/itrf_plot", frame='itrf', show=True) +ssapy.plotUtils.orbit_plot(r=r, t=times, save_path=f"{save_folder}/lunar_plot", frame='lunar', show=True) +ssapy.plotUtils.orbit_plot(r=r, t=times, save_path=f"{save_folder}/lunar_axis_lot", frame='lunar axis', show=True) print(f"Created a GCRF orbit plot.") print(f"Created a ITRF orbit plot.") print(f"Created a Lunar orbit plot.") @@ -65,10 +65,10 @@ def DRO(t, delta_r=7.52064e7, delta_v=344): # Globe plot of GTO Orbit r_geo, _, t_geo = ssapy.simple.ssapy_orbit(a=ssapy.constants.RGEO, e=0.3, t=times) -ssapy.plotUtils.globe_plot(r=r_geo, t=t_geo, save_path=f"{save_folder}globe_plot", scale=5) +ssapy.plotUtils.globe_plot(r=r_geo, t=t_geo, save_path=f"{save_folder}/globe_plot", scale=5) print('Created a globe plot.') -ssapy.plotUtils.groundTrackPlot(r=r_geo, t=t_geo, ground_stations=None, save_path=f"{save_folder}ground_track_plot") +ssapy.plotUtils.groundTrackPlot(r=r_geo, t=t_geo, ground_stations=None, save_path=f"{save_folder}/ground_track_plot") print('Created a ground track plot.') # Example usage @@ -108,7 +108,7 @@ def DRO(t, delta_r=7.52064e7, delta_v=344): ax.axis('equal') ax.legend() plt.show() -ssapy.plotUtils.save_plot(fig, save_path=f"{save_folder}lagrange_points") +ssapy.plotUtils.save_plot(fig, save_path=f"{save_folder}/lagrange_points") # Plotting fig = plt.figure(figsize=(8, 8)) @@ -143,7 +143,7 @@ def DRO(t, delta_r=7.52064e7, delta_v=344): ax.axis('equal') ax.legend() plt.show() -ssapy.plotUtils.save_plot(fig, save_path=f"{save_folder}lagrange_points") +ssapy.plotUtils.save_plot(fig, save_path=f"{save_folder}/lagrange_points") print(f"Lagrange points were calculated correctly.") print(f"Rotate vector plot successfully created.") From 5fe524d609674d7f435d068db442ce7f2d1c8a40 Mon Sep 17 00:00:00 2001 From: SuperdoerTrav Date: Mon, 22 Jul 2024 17:07:49 -0700 Subject: [PATCH 34/55] empty line at end --- tests/test_plots.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_plots.py b/tests/test_plots.py index bd047fb..fe075d0 100644 --- a/tests/test_plots.py +++ b/tests/test_plots.py @@ -150,4 +150,4 @@ def DRO(t, delta_r=7.52064e7, delta_v=344): print(f"save_plot() executed succesfully.") print(f"write_gif() executed succesfully.") -print(f"\nFinished plot testing!\n") \ No newline at end of file +print(f"\nFinished plot testing!\n") From 98c28e2dfbf9515a1a2ecc00b52e2866ce11e223 Mon Sep 17 00:00:00 2001 From: SuperdoerTrav Date: Mon, 22 Jul 2024 17:08:34 -0700 Subject: [PATCH 35/55] doc string to DRO() --- tests/test_plots.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/tests/test_plots.py b/tests/test_plots.py index fe075d0..3aa7885 100644 --- a/tests/test_plots.py +++ b/tests/test_plots.py @@ -32,6 +32,23 @@ def DRO(t, delta_r=7.52064e7, delta_v=344): + """ + Calculate an orbit with adjustments based on the Moon's position and velocity. + + Parameters: + ---------- + t : Time + The time at which to calculate the orbit. + delta_r : float, optional + The adjustment to the Moon's position (default is 7.52064e7 meters). + delta_v : float, optional + The adjustment to the Moon's velocity (default is 344 meters/second). + + Returns: + ------- + Orbit + SSAPy Orbit object. + """ moon = ssapy.get_body("moon") unit_vector_moon = moon.position(t) / np.linalg.norm(moon.position(t)) From c8480443a410d3f35bc6f9bffd15dc97ea3db02b Mon Sep 17 00:00:00 2001 From: SuperdoerTrav Date: Thu, 25 Jul 2024 13:38:37 -0700 Subject: [PATCH 36/55] new line added --- .gitignore | 1 + 1 file changed, 1 insertion(+) diff --git a/.gitignore b/.gitignore index 64c631b..6d806e6 100644 --- a/.gitignore +++ b/.gitignore @@ -46,6 +46,7 @@ coverage.xml pytest.xml pytest_session.txt cover/ + # Ignore ssapy_test_plots folder in tests subfolder tests/ssapy_test_plots/ From 10ae164f048749ef4f36a1cb156f9cf5dca49fac Mon Sep 17 00:00:00 2001 From: SuperdoerTrav Date: Thu, 25 Jul 2024 13:40:00 -0700 Subject: [PATCH 37/55] mentioned Kelerian in the doc string --- ssapy/compute.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/ssapy/compute.py b/ssapy/compute.py index fff1985..8f6913b 100644 --- a/ssapy/compute.py +++ b/ssapy/compute.py @@ -1212,9 +1212,9 @@ def nby3shape(arr_): def calculate_orbital_elements(r_, v_, mu_barycenter=EARTH_MU): """ - Calculate the orbital elements from position and velocity vectors. + Calculate the Keplerian orbital elements from position and velocity vectors. - This function computes the orbital elements (semi-major axis, eccentricity, inclination, true longitude, argument of periapsis, longitude of ascending node, true anomaly, and specific angular momentum) for one or more celestial objects given their position and velocity vectors. + This function computes the Keplerian orbital elements (semi-major axis, eccentricity, inclination, true longitude, argument of periapsis, longitude of ascending node, true anomaly, and specific angular momentum) for one or more celestial objects given their position and velocity vectors. Parameters: ---------- From 4f224ac8071bda5b7491ba69920979c2b9381000 Mon Sep 17 00:00:00 2001 From: SuperdoerTrav Date: Thu, 25 Jul 2024 13:41:06 -0700 Subject: [PATCH 38/55] file_exists renamed --- ssapy/io.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ssapy/io.py b/ssapy/io.py index c4dd6a0..f6d2a08 100644 --- a/ssapy/io.py +++ b/ssapy/io.py @@ -592,7 +592,7 @@ def load_b3obs_file(file_name): # ============================================================================= # File Handling Functions # ============================================================================= -def file_exists(filename): +def file_exists_extension_agnostic(filename): """ Check if a file with the given name and any extension exists. From 0ff5aa759741c0b92a914674828c102b71594f7a Mon Sep 17 00:00:00 2001 From: SuperdoerTrav Date: Thu, 25 Jul 2024 13:44:43 -0700 Subject: [PATCH 39/55] renamed getAngle and moved get_angle to utils.py where a copy was already. --- ssapy/compute.py | 22 +++++++++++----------- ssapy/utils.py | 34 +++++++++++++++++++++++++++++++++- 2 files changed, 44 insertions(+), 12 deletions(-) diff --git a/ssapy/compute.py b/ssapy/compute.py index 8f6913b..2a53fcd 100644 --- a/ssapy/compute.py +++ b/ssapy/compute.py @@ -6,7 +6,7 @@ from .constants import RGEO, LD, EARTH_RADIUS, EARTH_MU, MOON_RADIUS, MOON_MU from .propagator import KeplerianPropagator from .utils import ( - norm, normed, unitAngle3, LRU_Cache, lb_to_unit, sunPos, _gpsToTT, + norm, normed, unitAngle3, get_angle, LRU_Cache, lb_to_unit, sunPos, _gpsToTT, iers_interp, rotation_matrix_from_vectors, angle_between_vectors, gcrf_to_itrf ) from .orbit import Orbit @@ -1322,7 +1322,7 @@ def calculate_orbital_elements(r_, v_, mu_barycenter=EARTH_MU): ###################################################################################### -def getAngle(a, b, c): # a,b,c where b is the vertex +def get_angle(a, b, c): # a,b,c where b is the vertex """ Calculate the angle between two vectors where b is the vertex of the angle. @@ -1352,7 +1352,7 @@ def getAngle(a, b, c): # a,b,c where b is the vertex >>> a = np.array([[1, 0, 0]]) >>> b = np.array([[0, 0, 0]]) >>> c = np.array([[0, 1, 0]]) - >>> getAngle(a, b, c) + >>> get_angle(a, b, c) array([1.57079633]) """ a = np.atleast_2d(a) @@ -1415,9 +1415,9 @@ def moon_shine(r_moon, r_sat, r_earth, r_sun, radius, albedo, albedo_moon, albed {'moon_bus': array([...]), 'moon_panels': array([...])} """ # https://amostech.com/TechnicalPapers/2013/POSTER/COGNION.pdf - moon_phase_angle = getAngle(r_sun, r_moon, r_sat) # Phase of the moon as viewed from the sat. - sun_angle = getAngle(r_sun, r_sat, r_moon) # angle from Sun to object to Earth - moon_to_earth_angle = getAngle(r_moon, r_sat, r_earth) + moon_phase_angle = get_angle(r_sun, r_moon, r_sat) # Phase of the moon as viewed from the sat. + sun_angle = get_angle(r_sun, r_sat, r_moon) # angle from Sun to object to Earth + moon_to_earth_angle = get_angle(r_moon, r_sat, r_earth) r_moon_sat = np.linalg.norm(r_sat - r_moon, axis=-1) r_earth_sat = np.linalg.norm(r_sat - r_earth, axis=-1) # Earth is the observer. flux_moon_to_sat = 2 / 3 * albedo_moon * MOON_RADIUS**2 / (np.pi * (r_moon_sat)**2) * (np.sin(moon_phase_angle) + (np.pi - moon_phase_angle) * np.cos(moon_phase_angle)) # Fraction of sunlight reflected from the Moon to satellite @@ -1477,7 +1477,7 @@ def earth_shine(r_sat, r_earth, r_sun, radius, albedo, albedo_earth, albedo_back {'earth_bus': array([...]), 'earth_panels': array([...])} """ # https://amostech.com/TechnicalPapers/2013/POSTER/COGNION.pdf - phase_angle = getAngle(r_sun, r_sat, r_earth) # angle from Sun to object to Earth + phase_angle = get_angle(r_sun, r_sat, r_earth) # angle from Sun to object to Earth earth_angle = np.pi - phase_angle # Sun to Earth to oject. r_earth_sat = np.linalg.norm(r_sat - r_earth, axis=-1) # Earth is the observer. flux_earth_to_sat = 2 / 3 * albedo_earth * EARTH_RADIUS**2 / (np.pi * (r_earth_sat)**2) * (np.sin(earth_angle) + (np.pi - earth_angle) * np.cos(earth_angle)) # Fraction of sunlight reflected from the Earth to satellite @@ -1532,7 +1532,7 @@ def sun_shine(r_sat, r_earth, r_sun, radius, albedo, albedo_front, area_panels): {'sun_bus': array([...]), 'sun_panels': array([...])} """ # https://amostech.com/TechnicalPapers/2013/POSTER/COGNION.pdf - phase_angle = getAngle(r_sun, r_sat, r_earth) # angle from Sun to object to Earth + phase_angle = get_angle(r_sun, r_sat, r_earth) # angle from Sun to object to Earth r_earth_sat = np.linalg.norm(r_sat - r_earth, axis=-1) # Earth is the observer. flux_front = np.zeros_like(phase_angle) flux_front[phase_angle < np.pi / 2] = albedo_front * area_panels / (np.pi * r_earth_sat[phase_angle < np.pi / 2]**2) * np.cos(phase_angle[phase_angle < np.pi / 2]) # Fraction of Sun light scattered off front of the solar panels - which are assumed to be always facing the Sun. Angle: Sun - Sat - Observer @@ -1674,10 +1674,10 @@ def M_v_lambertian(r_sat, times, radius=1.0, albedo=0.20, sun_Mag=4.80, albedo_e r_earth_sat = np.linalg.norm(r_sat, axis=-1) r_moon_sat = np.linalg.norm(r_sat - r_moon, axis=-1) - sun_angle = getAngle(r_sun, r_sat, r_earth) + sun_angle = get_angle(r_sun, r_sat, r_earth) earth_angle = np.pi - sun_angle - moon_phase_angle = getAngle(r_sun, r_moon, r_sat) # Phase of the moon as viewed from the sat. - moon_to_earth_angle = getAngle(r_moon, r_sat, r_earth) + moon_phase_angle = get_angle(r_sun, r_moon, r_sat) # Phase of the moon as viewed from the sat. + moon_to_earth_angle = get_angle(r_moon, r_sat, r_earth) flux_moon_to_sat = 2 / 3 * albedo_moon * MOON_RADIUS**2 / (np.pi * (r_moon_sat)**2) * (np.sin(moon_phase_angle) + (np.pi - moon_phase_angle) * np.cos(moon_phase_angle)) # Fraction of sunlight reflected from the Moon to satellite flux_earth_to_sat = 2 / 3 * albedo_earth * EARTH_RADIUS**2 / (np.pi * (r_earth_sat)**2) * (np.sin(earth_angle) + (np.pi - earth_angle) * np.cos(earth_angle)) # Fraction of sunlight reflected from the Earth to satellite diff --git a/ssapy/utils.py b/ssapy/utils.py index 84bb304..cfb6405 100644 --- a/ssapy/utils.py +++ b/ssapy/utils.py @@ -1168,7 +1168,39 @@ def unit_vector(vector): return vector / np.linalg.norm(vector) -def getAngle(a, b, c): # a,b,c where b is the vertex +def get_angle(a, b, c): # a,b,c where b is the vertex + """ + Calculate the angle between two vectors where b is the vertex of the angle. + + This function computes the angle between vectors `ba` and `bc`, where `b` is the vertex and `a` and `c` are the endpoints of the angle. + + Parameters: + ---------- + a : (n, 3) numpy.ndarray + Array of coordinates representing the first vector. + b : (n, 3) numpy.ndarray + Array of coordinates representing the vertex of the angle. + c : (n, 3) numpy.ndarray + Array of coordinates representing the second vector. + + Returns: + ------- + numpy.ndarray + Array of angles (in radians) between the vectors `ba` and `bc`. + + Notes: + ------ + - The function handles multiple vectors by using broadcasting. + - The angle is calculated using the dot product formula and the arccosine function. + + Example: + -------- + >>> a = np.array([[1, 0, 0]]) + >>> b = np.array([[0, 0, 0]]) + >>> c = np.array([[0, 1, 0]]) + >>> get_angle(a, b, c) + array([1.57079633]) + """ a = np.atleast_2d(a) b = np.atleast_2d(b) c = np.atleast_2d(c) From 978380ec5724eb3520a4e5d88865df7bb8eb3807 Mon Sep 17 00:00:00 2001 From: SuperdoerTrav Date: Thu, 25 Jul 2024 13:50:28 -0700 Subject: [PATCH 40/55] renamed file_exists --- ssapy/plotUtils.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/ssapy/plotUtils.py b/ssapy/plotUtils.py index 83a2345..8c1dfea 100644 --- a/ssapy/plotUtils.py +++ b/ssapy/plotUtils.py @@ -600,7 +600,7 @@ def koe_plot(r, v, t=Time("2025-01-01", scale='utc') + np.linspace(0, int(1 * 36 if np.abs(np.max(a) - np.min(a)) < 2: ax2.set_ylim((np.min(a) - 0.5, np.max(a) + 0.5)) - date_format(t, ax1) + format_date_axis(t, ax1) plt.show(block=False) if save_path: @@ -1104,7 +1104,7 @@ def draw_dashed_circle(ax, normal_vector, radius, dashes, dash_length=0.1, label # ##################################################################### # Formatting x axis # ##################################################################### -def date_format(time_array, ax): +def format_date_axis(time_array, ax): """ Format the x-axis of a plot with time-based labels depending on the span of the time array. @@ -1138,7 +1138,7 @@ def date_format(time_array, ax): fig, ax = plt.subplots() ax.plot(time_array.decimalyear, np.random.rand(len(time_array))) - date_format(time_array, ax) + format_date_axis(time_array, ax) plt.show() ``` """ From 1db35bcd16b0e47bf5725b14bb70839409f4b504 Mon Sep 17 00:00:00 2001 From: SuperdoerTrav Date: Thu, 25 Jul 2024 13:53:37 -0700 Subject: [PATCH 41/55] modified save_plot to accept any extension --- ssapy/plotUtils.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/ssapy/plotUtils.py b/ssapy/plotUtils.py index 8c1dfea..eb10615 100644 --- a/ssapy/plotUtils.py +++ b/ssapy/plotUtils.py @@ -1257,11 +1257,11 @@ def save_plot_to_pdf(figure, pdf_path): def save_plot(figure, save_path, dpi=200): """ - Save a Python figure as a PNG image. + Save a Python figure as a PNG/JPG/PDF/ect. image. If no extension is given in the save_path, a .png is defaulted. Parameters: figure (matplotlib.figure.Figure): The figure object to be saved. - save_path (str): The file path where the PNG image will be saved. + save_path (str): The file path where the image will be saved. Returns: None @@ -1271,7 +1271,7 @@ def save_plot(figure, save_path, dpi=200): return try: base_name, extension = os.path.splitext(save_path) - if extension.lower() != '.png': + if extension.lower() == '': save_path = base_name + '.png' # Save the figure as a PNG image figure.savefig(save_path, dpi=dpi, bbox_inches='tight') From 64843a23b0787dc3b62743825d6bdb5cec69bd53 Mon Sep 17 00:00:00 2001 From: SuperdoerTrav Date: Thu, 25 Jul 2024 13:55:44 -0700 Subject: [PATCH 42/55] renamed write gif to save_animated_gif --- ssapy/plotUtils.py | 2 +- tests/test_plots.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/ssapy/plotUtils.py b/ssapy/plotUtils.py index eb10615..2b78f93 100644 --- a/ssapy/plotUtils.py +++ b/ssapy/plotUtils.py @@ -1281,7 +1281,7 @@ def save_plot(figure, save_path, dpi=200): print(f"Error occurred while saving the figure: {e}") -def write_gif(gif_name, frames, fps=30): +def save_animated_gif(gif_name, frames, fps=30): """ Create a GIF from a sequence of image frames. diff --git a/tests/test_plots.py b/tests/test_plots.py index 3aa7885..eb569cc 100644 --- a/tests/test_plots.py +++ b/tests/test_plots.py @@ -22,7 +22,7 @@ i += 1 gif_path = f"{save_folder}/rotate_vectors_{v_unit[0]:.0f}_{v_unit[1]:.0f}_{v_unit[2]:.0f}.gif" -ssapy.plotUtils.write_gif(gif_name=gif_path, frames=ssapy.io.sortbynum(ssapy.io.listdir(f'{temp_directory}*')), fps=20) +ssapy.plotUtils.save_animated_gif(gif_name=gif_path, frames=ssapy.io.sortbynum(ssapy.io.listdir(f'{temp_directory}*')), fps=20) shutil.rmtree(temp_directory) @@ -165,6 +165,6 @@ def DRO(t, delta_r=7.52064e7, delta_v=344): print(f"Lagrange points were calculated correctly.") print(f"Rotate vector plot successfully created.") print(f"save_plot() executed succesfully.") -print(f"write_gif() executed succesfully.") +print(f"save_animated_gif() executed succesfully.") print(f"\nFinished plot testing!\n") From 1945c4bf1bca67df0ce164bd635f2989c773f23f Mon Sep 17 00:00:00 2001 From: SuperdoerTrav Date: Thu, 25 Jul 2024 13:56:35 -0700 Subject: [PATCH 43/55] removed a comment --- ssapy/simple.py | 1 - 1 file changed, 1 deletion(-) diff --git a/ssapy/simple.py b/ssapy/simple.py index a9ac80d..f6a1e40 100644 --- a/ssapy/simple.py +++ b/ssapy/simple.py @@ -236,7 +236,6 @@ def ssapy_orbit(orbit=None, a=None, e=0, i=0, pa=0, raan=0, ta=0, r=None, v=None return np.nan, np.nan, np.nan -# Generate orbits near stable orbit. def get_similar_orbits(r0, v0, rad=1e5, num_orbits=4, duration=(90, 'days'), freq=(1, 'hour'), start_date="2025-1-1", mass=250): """ Generate similar orbits by varying the initial position. From 12a37b7554bac60d03fd9c342e9a6332c7ce14c1 Mon Sep 17 00:00:00 2001 From: SuperdoerTrav Date: Thu, 25 Jul 2024 13:57:56 -0700 Subject: [PATCH 44/55] removed get_angle() --- ssapy/compute.py | 42 ------------------------------------------ 1 file changed, 42 deletions(-) diff --git a/ssapy/compute.py b/ssapy/compute.py index 2a53fcd..f92cf38 100644 --- a/ssapy/compute.py +++ b/ssapy/compute.py @@ -1322,48 +1322,6 @@ def calculate_orbital_elements(r_, v_, mu_barycenter=EARTH_MU): ###################################################################################### -def get_angle(a, b, c): # a,b,c where b is the vertex - """ - Calculate the angle between two vectors where b is the vertex of the angle. - - This function computes the angle between vectors `ba` and `bc`, where `b` is the vertex and `a` and `c` are the endpoints of the angle. - - Parameters: - ---------- - a : (n, 3) numpy.ndarray - Array of coordinates representing the first vector. - b : (n, 3) numpy.ndarray - Array of coordinates representing the vertex of the angle. - c : (n, 3) numpy.ndarray - Array of coordinates representing the second vector. - - Returns: - ------- - numpy.ndarray - Array of angles (in radians) between the vectors `ba` and `bc`. - - Notes: - ------ - - The function handles multiple vectors by using broadcasting. - - The angle is calculated using the dot product formula and the arccosine function. - - Example: - -------- - >>> a = np.array([[1, 0, 0]]) - >>> b = np.array([[0, 0, 0]]) - >>> c = np.array([[0, 1, 0]]) - >>> get_angle(a, b, c) - array([1.57079633]) - """ - a = np.atleast_2d(a) - b = np.atleast_2d(b) - c = np.atleast_2d(c) - ba = np.subtract(a, b) - bc = np.subtract(c, b) - cosine_angle = np.sum(ba * bc, axis=-1) / (np.linalg.norm(ba, axis=-1) * np.linalg.norm(bc, axis=-1)) - return np.arccos(cosine_angle) - - def moon_shine(r_moon, r_sat, r_earth, r_sun, radius, albedo, albedo_moon, albedo_back, albedo_front, area_panels): # In SI units, takes single values or arrays returns a fractional flux """ Calculate the fractional flux of sunlight reflected from the Moon to the satellite. From c20ae295e742187e61a930673e3eba91a2216dfb Mon Sep 17 00:00:00 2001 From: SuperdoerTrav Date: Thu, 25 Jul 2024 14:00:59 -0700 Subject: [PATCH 45/55] changed doc strings to lagrange point functions --- ssapy/compute.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/ssapy/compute.py b/ssapy/compute.py index f92cf38..5f21283 100644 --- a/ssapy/compute.py +++ b/ssapy/compute.py @@ -1748,9 +1748,9 @@ def moon_normal_vector(t): def lunar_lagrange_points(t): """ - Calculate the positions of the Lagrange points in the Lunar frame at a given time. + Calculate the positions of the Lagrange points in the GCRF frame at a given time. - This function computes the positions of the five Lagrange points (L1, L2, L3, L4, and L5) in the Earth-Moon system at a specific time `t`. It considers the positions of the Earth and Moon and uses these to determine the locations of the Lagrange points. + This function computes the positions of the five Lagrange points (L1, L2, L3, L4, and L5) in the Earth-Moon system at a specific time `t`. It considers the positions of the Earth and Moon and uses the orbital period of the Moon to find the Lagrange points. Parameters: ---------- @@ -1816,7 +1816,7 @@ def lunar_lagrange_points(t): def lunar_lagrange_points_circular(t): """ - Calculate the positions of the Lagrange points in the Lunar frame for a given time. + Calculate the positions of the Lagrange points in the GCRF frame for a given time. This function calculates the positions of the five Lagrange points (L1, L2, L3, L4, and L5) in the Earth-Moon system at a specific time `t`. It accounts for the rotation of the Moon's orbit around the Earth, providing the positions in a circular approximation of the Earth-Moon system. @@ -1892,7 +1892,7 @@ def lunar_lagrange_points_circular(t): def lagrange_points_lunar_frame(): """ - Calculate the positions of the Lunar Lagrange points in the Lunar frame. + Calculate the positions of the Lunar Lagrange points in the Lunar frame, This frame is defined by the coordinate transformation in utils.py gcrf_to_lunar(). Returns: ------- From 45535f5ba06ce26ee897eb8858e9c80e6762fa2b Mon Sep 17 00:00:00 2001 From: SuperdoerTrav Date: Thu, 25 Jul 2024 14:29:31 -0700 Subject: [PATCH 46/55] changed some function names that were changed in .utils and .compute --- tests/test_plots.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_plots.py b/tests/test_plots.py index eb569cc..37e8eac 100644 --- a/tests/test_plots.py +++ b/tests/test_plots.py @@ -31,7 +31,7 @@ moon = ssapy.get_body("moon").position(times).T -def DRO(t, delta_r=7.52064e7, delta_v=344): +def initialize_DRO(t, delta_r=7.52064e7, delta_v=344): """ Calculate an orbit with adjustments based on the Moon's position and velocity. @@ -64,7 +64,7 @@ def DRO(t, delta_r=7.52064e7, delta_v=344): # DRO -dro_orbit = DRO(t=times[0]) +dro_orbit = initialize_DRO(t=times[0]) r, v, t = ssapy.simple.ssapy_orbit(orbit=dro_orbit, t=times) ssapy.plotUtils.orbit_plot(r=r, t=times, save_path=f"{save_folder}/DRO_orbit", frame='Lunar', show=False) r_lunar, v_lunar = ssapy.utils.gcrf_to_lunar_fixed(r, t=times, v=True) From 92e773c991ae67d5858c9b8bcb73c746bb3e1958 Mon Sep 17 00:00:00 2001 From: SuperdoerTrav Date: Thu, 25 Jul 2024 14:38:00 -0700 Subject: [PATCH 47/55] merged make_black and make_white to set_color_theme --- ssapy/plotUtils.py | 80 +++++++++++++--------------------------------- 1 file changed, 22 insertions(+), 58 deletions(-) diff --git a/ssapy/plotUtils.py b/ssapy/plotUtils.py index 2b78f93..c56b4ca 100644 --- a/ssapy/plotUtils.py +++ b/ssapy/plotUtils.py @@ -531,7 +531,7 @@ def globe_plot(r, t, limits=False, title='', figsize=(7, 8), save_path=False, el ax.tick_params(axis='y', colors='white') # Set y-axis tick color to white ax.tick_params(axis='z', colors='white') # Set z-axis tick color to white ax.set_aspect('equal') - fig, ax = make_black(fig, ax) + fig, ax = set_color_theme(fig, ax, theme='black') if save_path: save_plot(fig, save_path) return fig, ax @@ -577,7 +577,7 @@ def koe_plot(r, v, t=Time("2025-01-01", scale='utc') + np.linspace(0, int(1 * 36 fig.patch.set_facecolor('white') ax1.plot([], [], label='semi-major axis [GEO]', c='C0', linestyle='-') ax2 = ax1.twinx() - make_white(fig, *[ax1, ax2]) + set_color_theme(fig, *[ax1, ax2], theme='white') ax1.plot(Time(t).decimalyear, [x for x in orbital_elements['e']], label='eccentricity', c='C1') ax1.plot(Time(t).decimalyear, [x for x in orbital_elements['i']], label='inclination [rad]', c='C2') @@ -700,7 +700,7 @@ def __init__(self): im = fig.subplots_adjust(right=0.8) cbar_ax = fig.add_axes([0.82, 0.15, 0.01, 0.7]) fig.colorbar(im, cax=cbar_ax, norm=norm, cmap=cmap) - fig, ax = make_white(fig, ax) + fig, ax = set_color_theme(fig, ax, theme='white') if save_path: save_plot(fig, save_path) return fig @@ -766,7 +766,7 @@ def scatter2d(x, y, cs, xlabel='x', ylabel='y', title='', cbar_label='', dotsize scalarMap.set_array(cs) fig.colorbar(scalarMap, shrink=.5, label=f'{cbar_label}', pad=0.04) plt.tight_layout() - fig, ax = make_black(fig, ax) + fig, ax = set_color_theme(fig, ax, theme='black') plt.show(block=False) if save_path: save_plot(fig, save_path) @@ -833,7 +833,7 @@ def scatter3d(x, y=None, z=None, cs=None, xlabel='x', ylabel='y', zlabel='z', cb ax.set_zlabel(zlabel) plt.title(title) plt.tight_layout() - fig, ax = make_black(fig, ax) + fig, ax = set_color_theme(fig, ax, theme='black') plt.show(block=False) if save_path: save_plot(fig, save_path) @@ -947,64 +947,21 @@ def orbit_divergence_plot(rs, r_moon=[], t=False, limits=False, title='', save_p return -def make_white(fig, *axes): +def set_color_theme(fig, *axes, theme): """ - Set the background color of the figure and axes to white and the text color to black. + Set the color theme of the figure and axes to white or black and the text color to white or black. Parameters: - fig (matplotlib.figure.Figure): The figure to modify. - axes (list of matplotlib.axes._subplots.AxesSubplot): One or more axes to modify. - - Returns: - - fig (matplotlib.figure.Figure): The modified figure. - - axes (tuple of matplotlib.axes._subplots.AxesSubplot): The modified axes. - - This function changes the background color of the given figure and its axes to white. - It also sets the color of all text items (title, labels, tick labels) to black. - - Example usage: - ``` - import matplotlib.pyplot as plt - - fig, ax = plt.subplots() - ax.plot([1, 2, 3], [4, 5, 6]) - make_white(fig, ax) - plt.show() - ``` - """ - fig.patch.set_facecolor('white') - - for ax in axes: - ax.set_facecolor('white') - ax_items = [ax.title, ax.xaxis.label, ax.yaxis.label] - if hasattr(ax, 'zaxis'): - ax_items.append(ax.zaxis.label) - ax_items += ax.get_xticklabels() + ax.get_yticklabels() - if hasattr(ax, 'get_zticklabels'): - ax_items += ax.get_zticklabels() - ax_items += ax.get_xticklines() + ax.get_yticklines() - if hasattr(ax, 'get_zticklines'): - ax_items += ax.get_zticklines() - for item in ax_items: - item.set_color('black') - - return fig, axes - - -def make_black(fig, *axes): - """ - Set the background color of the figure and axes to black and the text color to white. - - Parameters: - - fig (matplotlib.figure.Figure): The figure to modify. - - axes (list of matplotlib.axes._subplots.AxesSubplot): One or more axes to modify. - + - theme (str) either black/dark or white. + Returns: - fig (matplotlib.figure.Figure): The modified figure. - axes (tuple of matplotlib.axes._subplots.AxesSubplot): The modified axes. - This function changes the background color of the given figure and its axes to black. - It also sets the color of all text items (title, labels, tick labels) to white. + This function changes the background color of the given figure and its axes to black or white. + It also sets the color of all text items (title, labels, tick labels) to white or black. Example usage: ``` @@ -1012,14 +969,21 @@ def make_black(fig, *axes): fig, ax = plt.subplots() ax.plot([1, 2, 3], [4, 5, 6]) - make_black(fig, ax) + set_color_theme(fig, ax, theme='black') plt.show() ``` """ - fig.patch.set_facecolor('black') + if theme == 'black' or theme == 'dark': + background = 'black' + secondary = 'white' + else: + background = 'white' + secondary = 'black' + + fig.patch.set_facecolor(background) for ax in axes: - ax.set_facecolor('black') + ax.set_facecolor(background) ax_items = [ax.title, ax.xaxis.label, ax.yaxis.label] if hasattr(ax, 'zaxis'): ax_items.append(ax.zaxis.label) @@ -1030,7 +994,7 @@ def make_black(fig, *axes): if hasattr(ax, 'get_zticklines'): ax_items += ax.get_zticklines() for item in ax_items: - item.set_color('white') + item.set_color(secondary) return fig, axes From 2307b32b39223c60c5ea0a4681b7413aa1a09c0d Mon Sep 17 00:00:00 2001 From: SuperdoerTrav Date: Thu, 25 Jul 2024 14:43:15 -0700 Subject: [PATCH 48/55] renamed groundTrackPlot and added Keplerian to several dot strings --- ssapy/plotUtils.py | 51 ++++++++++++++++++++++----------------------- tests/test_plots.py | 2 +- 2 files changed, 26 insertions(+), 27 deletions(-) diff --git a/ssapy/plotUtils.py b/ssapy/plotUtils.py index c56b4ca..8d7a7f1 100644 --- a/ssapy/plotUtils.py +++ b/ssapy/plotUtils.py @@ -132,7 +132,7 @@ def drawMoon(time, ngrid=100, R=MOON_RADIUS, rfactor=1): ) -def groundTrackPlot(r, t, ground_stations=None, save_path=False): +def ground_track_plot(r, t, ground_stations=None, save_path=False): """ Parameters ---------- @@ -331,12 +331,12 @@ def get_main_category(frame): if orbit_index == '': angle = 0 - dotcolors = cm.rainbow(np.linspace(0, 1, len(x))) + scatter_dot_colors = cm.rainbow(np.linspace(0, 1, len(x))) else: angle = orbit_index * 10 - dotcolors = cm.rainbow(np.linspace(0, 1, num_orbits))[orbit_index] + scatter_dot_colors = cm.rainbow(np.linspace(0, 1, num_orbits))[orbit_index] ax1.add_patch(plt.Circle((0, 0), stn[2], color='white', linestyle='dashed', fill=False)) - ax1.scatter(x, y, color=dotcolors, s=1) + ax1.scatter(x, y, color=scatter_dot_colors, s=1) ax1.scatter(0, 0, color=stn[0], s=stn[1]) if xm is not False: ax1.scatter(stn[3], stn[4], color=stn[6], s=5) @@ -357,7 +357,7 @@ def get_main_category(frame): ax1.text(pos[0], pos[1], point, color=color) ax2.add_patch(plt.Circle((0, 0), stn[2], color='white', linestyle='dashed', fill=False)) - ax2.scatter(x, z, color=dotcolors, s=1) + ax2.scatter(x, z, color=scatter_dot_colors, s=1) ax2.scatter(0, 0, color=stn[0], s=stn[1]) if xm is not False: ax2.scatter(stn[3], stn[5], color=stn[6], s=5) @@ -378,7 +378,7 @@ def get_main_category(frame): ax2.text(pos[0], pos[2], point, color=color) ax3.add_patch(plt.Circle((0, 0), stn[2], color='white', linestyle='dashed', fill=False)) - ax3.scatter(y, z, color=dotcolors, s=1) + ax3.scatter(y, z, color=scatter_dot_colors, s=1) ax3.scatter(0, 0, color=stn[0], s=stn[1]) if xm is not False: ax3.scatter(stn[4], stn[5], color=stn[6], s=5) @@ -398,7 +398,7 @@ def get_main_category(frame): ax3.scatter(pos[1], pos[2], color=color, label=point) ax3.text(pos[1], pos[2], point, color=color) - ax4.scatter3D(x, y, z, color=dotcolors, s=1) + ax4.scatter3D(x, y, z, color=scatter_dot_colors, s=1) ax4.scatter3D(0, 0, 0, color=stn[0], s=stn[1]) if xm is not False: ax4.scatter3D(stn[3], stn[4], stn[5], color=stn[6], s=5) @@ -511,14 +511,14 @@ def globe_plot(r, t, limits=False, title='', figsize=(7, 8), save_path=False, el mesh_y = np.outer(np.sin(lons), np.cos(lats)).T * EARTH_RADIUS / RGEO mesh_z = np.outer(np.ones(np.size(lons)), np.sin(lats)).T * EARTH_RADIUS / RGEO - dotcolors = plt.cm.rainbow(np.linspace(0, 1, len(x))) + scatter_dot_colors = plt.cm.rainbow(np.linspace(0, 1, len(x))) fig = plt.figure(dpi=100, figsize=figsize) ax = fig.add_subplot(111, projection='3d') fig.patch.set_facecolor('black') ax.tick_params(axis='both', colors='white') ax.grid(True, color='grey', linestyle='--', linewidth=0.5) ax.set_facecolor('black') # Set plot background color to black - ax.scatter(x, y, z, color=dotcolors, s=1) + ax.scatter(x, y, z, color=scatter_dot_colors, s=1) ax.plot_surface(mesh_x, mesh_y, mesh_z, rstride=4, cstride=4, facecolors=bm, shade=False) ax.view_init(elev=el, azim=az) ax.set_xlim([-limits, limits]) @@ -539,7 +539,7 @@ def globe_plot(r, t, limits=False, title='', figsize=(7, 8), save_path=False, el def koe_plot(r, v, t=Time("2025-01-01", scale='utc') + np.linspace(0, int(1 * 365.25), int(365.25 * 24)), elements=['a', 'e', 'i'], save_path=False, body='Earth'): """ - Plot orbital elements over time for a given trajectory. + Plot Keplerian orbital elements over time for a given trajectory. Parameters: - r (array-like): Position vectors for the orbit. @@ -608,9 +608,9 @@ def koe_plot(r, v, t=Time("2025-01-01", scale='utc') + np.linspace(0, int(1 * 36 return fig, ax1 -def koe_2dhist(stable_data, title="Initial orbital elements of\n1 year stable cislunar orbits", limits=[1, 50], bins=200, logscale=False, cmap='coolwarm', save_path=False): +def koe_hist_2d(stable_data, title="Initial orbital elements of\n1 year stable cislunar orbits", limits=[1, 50], bins=200, logscale=False, cmap='coolwarm', save_path=False): """ - Create a 2D histogram plot for various orbital elements of stable cislunar orbits. + Create a 2D histogram plot for various Keplerian orbital elements of stable cislunar orbits. Parameters: - stable_data (object): An object with attributes `a`, `e`, `i`, and `ta`, which are arrays of semi-major axis, eccentricity, inclination, and true anomaly, respectively. @@ -629,7 +629,7 @@ def koe_2dhist(stable_data, title="Initial orbital elements of\n1 year stable ci Example usage: ``` import numpy as np - from your_module import koe_2dhist + from your_module import koe_hist_2d # Example data class StableData: @@ -640,7 +640,7 @@ def __init__(self): self.ta = np.radians(np.random.uniform(0, 360, 1000)) stable_data = StableData() - koe_2dhist(stable_data, save_path='orbit_histograms.pdf') + koe_hist_2d(stable_data, save_path='orbit_histograms.pdf') ``` """ if logscale or logscale == 'log': @@ -707,7 +707,7 @@ def __init__(self): -def scatter2d(x, y, cs, xlabel='x', ylabel='y', title='', cbar_label='', dotsize=1, colorsMap='jet', colorscale='linear', colormin=False, colormax=False, save_path=False): +def scatter_2d(x, y, cs, xlabel='x', ylabel='y', title='', cbar_label='', dotsize=1, colorsMap='jet', colorscale='linear', colormin=False, colormax=False, save_path=False): """ Create a 2D scatter plot with optional color mapping. @@ -737,14 +737,14 @@ def scatter2d(x, y, cs, xlabel='x', ylabel='y', title='', cbar_label='', dotsize Example usage: ``` import numpy as np - from your_module import scatter2d + from your_module import scatter_2d # Example data x = np.random.rand(100) y = np.random.rand(100) cs = np.random.rand(100) - scatter2d(x, y, cs, xlabel='X-axis', ylabel='Y-axis', cbar_label='Color Scale', title='2D Scatter Plot') + scatter_2d(x, y, cs, xlabel='X-axis', ylabel='Y-axis', cbar_label='Color Scale', title='2D Scatter Plot') ``` """ fig = plt.figure() @@ -773,7 +773,7 @@ def scatter2d(x, y, cs, xlabel='x', ylabel='y', title='', cbar_label='', dotsize return -def scatter3d(x, y=None, z=None, cs=None, xlabel='x', ylabel='y', zlabel='z', cbar_label='', dotsize=1, colorsMap='jet', title='', save_path=False): +def scatter_3d(x, y=None, z=None, cs=None, xlabel='x', ylabel='y', zlabel='z', cbar_label='', dotsize=1, colorsMap='jet', title='', save_path=False): """ Create a 3D scatter plot with optional color mapping. @@ -801,7 +801,7 @@ def scatter3d(x, y=None, z=None, cs=None, xlabel='x', ylabel='y', zlabel='z', cb Example usage: ``` import numpy as np - from your_module import scatter3d + from your_module import scatter_3d # Example data x = np.random.rand(100) @@ -809,7 +809,7 @@ def scatter3d(x, y=None, z=None, cs=None, xlabel='x', ylabel='y', zlabel='z', cb z = np.random.rand(100) cs = np.random.rand(100) - scatter3d(x, y, z, cs, xlabel='X-axis', ylabel='Y-axis', zlabel='Z-axis', cbar_label='Color Scale', title='3D Scatter Plot') + scatter_3d(x, y, z, cs, xlabel='X-axis', ylabel='Y-axis', zlabel='Z-axis', cbar_label='Color Scale', title='3D Scatter Plot') ``` """ fig = plt.figure() @@ -840,11 +840,10 @@ def scatter3d(x, y=None, z=None, cs=None, xlabel='x', ylabel='y', zlabel='z', cb return fig, ax -def dotcolors_scaled(num_colors): +def scatter_dot_colors_scaled(num_colors): return cm.rainbow(np.linspace(0, 1, num_colors)) -# Make a plot of multiple cislunar orbit in GCRF frame. def orbit_divergence_plot(rs, r_moon=[], t=False, limits=False, title='', save_path=False): """ Plot multiple cislunar orbits in the GCRF frame with respect to the Earth and Moon. @@ -901,11 +900,11 @@ def orbit_divergence_plot(rs, r_moon=[], t=False, limits=False, title='', save_p xm = r_moon[0] / RGEO ym = r_moon[1] / RGEO zm = r_moon[2] / RGEO - dotcolors = cm.rainbow(np.linspace(0, 1, len(x))) + scatter_dot_colors = cm.rainbow(np.linspace(0, 1, len(x))) # Creating plot plt.subplot(1, 3, 1) - plt.scatter(x, y, color=dotcolors, s=1) + plt.scatter(x, y, color=scatter_dot_colors, s=1) plt.scatter(0, 0, color="blue", s=50) plt.scatter(xm, ym, color="grey", s=5) plt.axis('scaled') @@ -917,7 +916,7 @@ def orbit_divergence_plot(rs, r_moon=[], t=False, limits=False, title='', save_p plt.text(x[-1], y[-1], '$\leftarrow$ end') plt.subplot(1, 3, 2) - plt.scatter(x, z, color=dotcolors, s=1) + plt.scatter(x, z, color=scatter_dot_colors, s=1) plt.scatter(0, 0, color="blue", s=50) plt.scatter(xm, zm, color="grey", s=5) plt.axis('scaled') @@ -930,7 +929,7 @@ def orbit_divergence_plot(rs, r_moon=[], t=False, limits=False, title='', save_p plt.title(f'{title}') plt.subplot(1, 3, 3) - plt.scatter(y, z, color=dotcolors, s=1) + plt.scatter(y, z, color=scatter_dot_colors, s=1) plt.scatter(0, 0, color="blue", s=50) plt.scatter(ym, zm, color="grey", s=5) plt.axis('scaled') diff --git a/tests/test_plots.py b/tests/test_plots.py index 37e8eac..0cc0733 100644 --- a/tests/test_plots.py +++ b/tests/test_plots.py @@ -85,7 +85,7 @@ def initialize_DRO(t, delta_r=7.52064e7, delta_v=344): ssapy.plotUtils.globe_plot(r=r_geo, t=t_geo, save_path=f"{save_folder}/globe_plot", scale=5) print('Created a globe plot.') -ssapy.plotUtils.groundTrackPlot(r=r_geo, t=t_geo, ground_stations=None, save_path=f"{save_folder}/ground_track_plot") +ssapy.plotUtils.ground_track_plot(r=r_geo, t=t_geo, ground_stations=None, save_path=f"{save_folder}/ground_track_plot") print('Created a ground track plot.') # Example usage From f3ab81c7ab6747382aee169d6aaedceca8e0cc9a Mon Sep 17 00:00:00 2001 From: SuperdoerTrav Date: Thu, 25 Jul 2024 14:45:01 -0700 Subject: [PATCH 49/55] drawmoon and drawearth renamed --- ssapy/plotUtils.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/ssapy/plotUtils.py b/ssapy/plotUtils.py index 8d7a7f1..3267764 100644 --- a/ssapy/plotUtils.py +++ b/ssapy/plotUtils.py @@ -24,7 +24,7 @@ def load_earth_file(): return earth -def drawEarth(time, ngrid=100, R=EARTH_RADIUS, rfactor=1): +def draw_earth(time, ngrid=100, R=EARTH_RADIUS, rfactor=1): """ Parameters ---------- @@ -81,7 +81,7 @@ def load_moon_file(): return moon -def drawMoon(time, ngrid=100, R=MOON_RADIUS, rfactor=1): +def draw_moon(time, ngrid=100, R=MOON_RADIUS, rfactor=1): """ Parameters ---------- @@ -171,7 +171,7 @@ def groundTrackVideo(r, time): ipv.style.box_off() ipv.style.axes_off() widgets = [] - widgets.append(drawEarth(time)) + widgets.append(draw_earth(time)) widgets.append( ipv.scatter( r[:, 0, None], From 7207768d7e027c1fb99efbf1b8a9a7ed3f009fec Mon Sep 17 00:00:00 2001 From: SuperdoerTrav Date: Thu, 25 Jul 2024 14:48:00 -0700 Subject: [PATCH 50/55] standardized npload and npsave and pload psave --- ssapy/io.py | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/ssapy/io.py b/ssapy/io.py index f6d2a08..71d806a 100644 --- a/ssapy/io.py +++ b/ssapy/io.py @@ -765,7 +765,7 @@ def get_memory_usage(): ###################################################################### -def psave(filename_, data_): +def save_pickle(filename_, data_): from six.moves import cPickle as pickle # for performance with open(filename_, 'wb') as f: pickle.dump(data_, f) @@ -773,7 +773,7 @@ def psave(filename_, data_): return -def pload(filename_): +def load_pickle(filename_): from six.moves import cPickle as pickle # for performance try: # print('Openning: ' + current_filename) @@ -790,9 +790,9 @@ def merge_dicts(file_names, save_path): number_of_files = len(file_names); master_dict = {} for count, file in enumerate(file_names): print(f'Merging dict: {count+1} of {number_of_files}, name: {file}, num of master keys: {len(master_dict.keys())}, num of new keys: {len(master_dict.keys())}') - master_dict.update(pload(file)) + master_dict.update(load_pickle(file)) print('Beginning final save.') - psave(save_path, master_dict) + save_pickle(save_path, master_dict) return @@ -801,7 +801,7 @@ def merge_dicts(file_names, save_path): ###################################################################### -def npsave(filename_, data_): +def save_np(filename_, data_): """ Save a NumPy array to a binary file. @@ -822,7 +822,7 @@ def npsave(filename_, data_): Examples: -------- >>> arr = np.array([1, 2, 3, 4, 5]) - >>> npsave('array.npy', arr) + >>> save_np('array.npy', arr) """ try: with open(filename_, 'wb') as f: @@ -833,7 +833,7 @@ def npsave(filename_, data_): return -def npload(filename_): +def load_np(filename_): """ Load a NumPy array from a binary file. @@ -851,7 +851,7 @@ def npload(filename_): Examples: -------- - >>> arr = npload('array.npy') + >>> arr = load_np('array.npy') >>> print(arr) [1 2 3 4 5] """ @@ -1570,9 +1570,9 @@ def pdstr_to_arrays(df): return df.apply(str_to_array).to_numpy() -def allfiles(dirName=os.getcwd()): +def get_all_files_recursive(path_name=os.getcwd()): # Get the list of all files in directory tree at given path listOfFiles = list() - for (dirpath, dirnames, filenames) in os.walk(dirName): + for (dirpath, dirnames, filenames) in os.walk(path_name): listOfFiles += [os.path.join(dirpath, file) for file in filenames] return listOfFiles From 0abf65588d01bf9a63a9287df2db862683d65e71 Mon Sep 17 00:00:00 2001 From: SuperdoerTrav Date: Thu, 25 Jul 2024 14:51:54 -0700 Subject: [PATCH 51/55] more renamed and changes --- ssapy/io.py | 29 +++++++++++++++++++++++++---- 1 file changed, 25 insertions(+), 4 deletions(-) diff --git a/ssapy/io.py b/ssapy/io.py index 71d806a..9add072 100644 --- a/ssapy/io.py +++ b/ssapy/io.py @@ -625,9 +625,7 @@ def exists(pathname): bool True if the path exists as either a file or a directory, False otherwise. """ - if os.path.isdir(pathname): - exists = True - elif os.path.isfile(pathname): + if os.path.isdir(pathname) or os.path.isfile(pathname): exists = True else: exists = False @@ -686,15 +684,38 @@ def sortbynum(files): """ Sorts a list of file paths based on numeric values in the filenames. + This function assumes that each filename contains at least one numeric value + and sorts the files based on the first numeric value found in the filename. + Parameters: ---------- files : list - List of file paths to be sorted. + List of file paths to be sorted. Each file path can be a full path or just a filename. Returns: ------- list List of file paths sorted by numeric values in their filenames. + + Notes: + ----- + - This function extracts the first numeric value it encounters in each filename. + - If no numeric value is found in a filename, the function may raise an error. + - The numeric value can appear anywhere in the filename. + - The function does not handle cases where filenames have different directory prefixes. + + Raises: + ------ + ValueError: + If a filename does not contain any numeric value. + + Examples: + -------- + >>> sortbynum(['file2.txt', 'file10.txt', 'file1.txt']) + ['file1.txt', 'file2.txt', 'file10.txt'] + + >>> sortbynum(['/path/to/file2.txt', '/path/to/file10.txt', '/path/to/file1.txt']) + ['/path/to/file1.txt', '/path/to/file2.txt', '/path/to/file10.txt'] """ import re if len(files[0].split('/')) > 1: From 8e131eac0bb45eacb05497ad482ceab1aaf22a4f Mon Sep 17 00:00:00 2001 From: SuperdoerTrav Date: Thu, 25 Jul 2024 16:17:24 -0700 Subject: [PATCH 52/55] lunar decapitalized --- ssapy/compute.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/ssapy/compute.py b/ssapy/compute.py index 5f21283..f1f05cc 100644 --- a/ssapy/compute.py +++ b/ssapy/compute.py @@ -1748,7 +1748,7 @@ def moon_normal_vector(t): def lunar_lagrange_points(t): """ - Calculate the positions of the Lagrange points in the GCRF frame at a given time. + Calculate the positions of the lunar Lagrange points in the GCRF frame at a given time. This function computes the positions of the five Lagrange points (L1, L2, L3, L4, and L5) in the Earth-Moon system at a specific time `t`. It considers the positions of the Earth and Moon and uses the orbital period of the Moon to find the Lagrange points. @@ -1816,7 +1816,7 @@ def lunar_lagrange_points(t): def lunar_lagrange_points_circular(t): """ - Calculate the positions of the Lagrange points in the GCRF frame for a given time. + Calculate the positions of the lunar Lagrange points in the GCRF frame for a given time. This function calculates the positions of the five Lagrange points (L1, L2, L3, L4, and L5) in the Earth-Moon system at a specific time `t`. It accounts for the rotation of the Moon's orbit around the Earth, providing the positions in a circular approximation of the Earth-Moon system. @@ -1892,7 +1892,7 @@ def lunar_lagrange_points_circular(t): def lagrange_points_lunar_frame(): """ - Calculate the positions of the Lunar Lagrange points in the Lunar frame, This frame is defined by the coordinate transformation in utils.py gcrf_to_lunar(). + Calculate the positions of the lunar Lagrange points in the lunar frame, This frame is defined by the coordinate transformation in utils.py gcrf_to_lunar(). Returns: ------- From 9c9a3d988c95e3f2c1b539518b9332cf7ac18039 Mon Sep 17 00:00:00 2001 From: SuperdoerTrav Date: Thu, 25 Jul 2024 16:19:53 -0700 Subject: [PATCH 53/55] bool instantly returned in exists --- ssapy/io.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/ssapy/io.py b/ssapy/io.py index 9add072..a54c50c 100644 --- a/ssapy/io.py +++ b/ssapy/io.py @@ -626,10 +626,9 @@ def exists(pathname): True if the path exists as either a file or a directory, False otherwise. """ if os.path.isdir(pathname) or os.path.isfile(pathname): - exists = True + return True else: - exists = False - return exists + return False def mkdir(pathname): From 018100c7c67b3fa54c4a9b05d81f700e0cffc906 Mon Sep 17 00:00:00 2001 From: SuperdoerTrav Date: Thu, 25 Jul 2024 16:21:35 -0700 Subject: [PATCH 54/55] index added to sortbynum to allow user to chose which number to sort by --- ssapy/io.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/ssapy/io.py b/ssapy/io.py index a54c50c..08fcc36 100644 --- a/ssapy/io.py +++ b/ssapy/io.py @@ -679,7 +679,7 @@ def rmfile(pathname): return -def sortbynum(files): +def sortbynum(files, index=0): """ Sorts a list of file paths based on numeric values in the filenames. @@ -690,6 +690,8 @@ def sortbynum(files): ---------- files : list List of file paths to be sorted. Each file path can be a full path or just a filename. + index: int + Index of the number in the string do you want to sort on. Returns: ------- @@ -722,12 +724,12 @@ def sortbynum(files): file_prefix = '/'.join(files[0].split('/')[:-1]) for file in files: files_shortened.append(file.split('/')[-1]) - files_sorted = sorted(files_shortened, key=lambda x: float(re.findall("(\d+)", x)[0])) + files_sorted = sorted(files_shortened, key=lambda x: float(re.findall("(\d+)", x)[index])) sorted_files = [] for file in files_sorted: sorted_files.append(f'{file_prefix}/{file}') else: - sorted_files = sorted(files, key=lambda x: float(re.findall("(\d+)", x)[0])) + sorted_files = sorted(files, key=lambda x: float(re.findall("(\d+)", x)[index])) return sorted_files From 107dc5c101c36cc9ec73637d111ab991f94f8d49 Mon Sep 17 00:00:00 2001 From: SuperdoerTrav Date: Thu, 25 Jul 2024 16:33:07 -0700 Subject: [PATCH 55/55] sortbynum renamed to _sortbynum --- ssapy/io.py | 12 +++++++----- tests/test_plots.py | 2 +- 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/ssapy/io.py b/ssapy/io.py index 08fcc36..68a3953 100644 --- a/ssapy/io.py +++ b/ssapy/io.py @@ -679,7 +679,7 @@ def rmfile(pathname): return -def sortbynum(files, index=0): +def _sortbynum(files, index=0): """ Sorts a list of file paths based on numeric values in the filenames. @@ -712,10 +712,10 @@ def sortbynum(files, index=0): Examples: -------- - >>> sortbynum(['file2.txt', 'file10.txt', 'file1.txt']) + >>> _sortbynum(['file2.txt', 'file10.txt', 'file1.txt']) ['file1.txt', 'file2.txt', 'file10.txt'] - >>> sortbynum(['/path/to/file2.txt', '/path/to/file10.txt', '/path/to/file1.txt']) + >>> _sortbynum(['/path/to/file2.txt', '/path/to/file10.txt', '/path/to/file1.txt']) ['/path/to/file1.txt', '/path/to/file2.txt', '/path/to/file10.txt'] """ import re @@ -733,7 +733,7 @@ def sortbynum(files, index=0): return sorted_files -def listdir(dir_path='*', files_only=False, exclude=None, sorted=False): +def listdir(dir_path='*', files_only=False, exclude=None, sorted=False, index=0): """ Lists files and directories in a specified path with optional filtering and sorting. @@ -747,6 +747,8 @@ def listdir(dir_path='*', files_only=False, exclude=None, sorted=False): If specified, excludes files and directories whose base name contains this string. sorted : bool, default=False If True, sorts the resulting list by numeric values in filenames. + index : int, default=0 + sorted required to be true. Index of the digit used for sorting. Returns: ------- @@ -769,7 +771,7 @@ def listdir(dir_path='*', files_only=False, exclude=None, sorted=False): new_files = [file for file in files if exclude not in os.path.basename(file)] files = new_files if sorted: - return sortbynum(files) + return _sortbynum(files, index=index) else: return files diff --git a/tests/test_plots.py b/tests/test_plots.py index 0cc0733..f5716b7 100644 --- a/tests/test_plots.py +++ b/tests/test_plots.py @@ -22,7 +22,7 @@ i += 1 gif_path = f"{save_folder}/rotate_vectors_{v_unit[0]:.0f}_{v_unit[1]:.0f}_{v_unit[2]:.0f}.gif" -ssapy.plotUtils.save_animated_gif(gif_name=gif_path, frames=ssapy.io.sortbynum(ssapy.io.listdir(f'{temp_directory}*')), fps=20) +ssapy.plotUtils.save_animated_gif(gif_name=gif_path, frames=ssapy.io.listdir(f'{temp_directory}*', sorted=True), fps=20) shutil.rmtree(temp_directory)