diff --git a/.github/workflows/pip-build-lint-test-coverage.yml b/.github/workflows/pip-build-lint-test-coverage.yml index 7df15e1..352cda0 100644 --- a/.github/workflows/pip-build-lint-test-coverage.yml +++ b/.github/workflows/pip-build-lint-test-coverage.yml @@ -9,6 +9,7 @@ on: permissions: contents: read + pull-requests: read concurrency: group: ci-${{ github.ref }} @@ -94,6 +95,7 @@ jobs: needs: [ test ] permissions: contents: write + pull-requests: write defaults: run: shell: bash -l {0} diff --git a/pyproject.toml b/pyproject.toml index f956745..9070f38 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -28,6 +28,7 @@ dependencies = [ "ray", "spiceypy>=6.0.0", "rebound>=4.4.10", + "timezonefinder==8.0.0", ] [build-system] diff --git a/src/adam_assist/propagator.py b/src/adam_assist/propagator.py index 6a99982..6b7e9cf 100644 --- a/src/adam_assist/propagator.py +++ b/src/adam_assist/propagator.py @@ -155,6 +155,10 @@ def _propagate_orbits(self, orbits: OrbitType, times: TimestampType) -> OrbitTyp """ Propagate the orbits to the specified times. """ + # OPTIMIZATION: Fast path for single orbits + if len(orbits) == 1: + return self._propagate_single_orbit_optimized(orbits, times) + # The coordinate frame is the equatorial International Celestial Reference Frame (ICRF). # This is also the native coordinate system for the JPL binary files. # For units we use solar masses, astronomical units, and days. @@ -188,6 +192,144 @@ def _propagate_orbits(self, orbits: OrbitType, times: TimestampType) -> OrbitTyp return results + def _propagate_single_orbit_optimized( + self, orbit: OrbitType, times: TimestampType + ) -> OrbitType: + """ + Optimized propagation for a single orbit, bypassing grouping overhead. + """ + # Validate assumption + if len(orbit) != 1: + raise ValueError(f"Expected exactly 1 orbit, got {len(orbit)}") + + # Transform coordinates directly without grouping + transformed_coords = transform_coordinates( + orbit.coordinates, + origin_out=OriginCodes.SOLAR_SYSTEM_BARYCENTER, + frame_out="equatorial", + ) + transformed_input_orbit_times = transformed_coords.time.rescale("tdb") + transformed_coords = transformed_coords.set_column( + "time", transformed_input_orbit_times + ) + transformed_orbit = orbit.set_column("coordinates", transformed_coords) + + return self._propagate_single_orbit_inner_optimized(transformed_orbit, times) + + def _propagate_single_orbit_inner_optimized( + self, orbit: OrbitType, times: TimestampType + ) -> OrbitType: + """ + Inner propagation optimized for exactly one orbit. + """ + # Setup ephemeris and simulation + ephem = assist.Ephem(planets_path=de440, asteroids_path=de441_n16) + sim = rebound.Simulation() + + start_tdb_time = orbit.coordinates.time.jd().to_numpy()[0] + sim.t = start_tdb_time - ephem.jd_ref + + # Handle particle ID creation (optimized for single orbit) + is_variant = isinstance(orbit, VariantOrbits) + if is_variant: + orbit_id = str(orbit.orbit_id.to_numpy(zero_copy_only=False)[0]) + variant_id = str(orbit.variant_id.to_numpy(zero_copy_only=False)[0]) + particle_hash = uint32_hash(f"{orbit_id}\x1f{variant_id}") + else: + orbit_id = str(orbit.orbit_id.to_numpy(zero_copy_only=False)[0]) + particle_hash = uint32_hash(orbit_id) + + assist.Extras(sim, ephem) + + # Add single particle + coords = orbit.coordinates + position_arrays = coords.r + velocity_arrays = coords.v + + sim.add( + x=position_arrays[0, 0], + y=position_arrays[0, 1], + z=position_arrays[0, 2], + vx=velocity_arrays[0, 0], + vy=velocity_arrays[0, 1], + vz=velocity_arrays[0, 2], + hash=particle_hash, + ) + + # Set integrator parameters + sim.dt = self.initial_dt + sim.ri_ias15.min_dt = self.min_dt + sim.ri_ias15.adaptive_mode = self.adaptive_mode + sim.ri_ias15.epsilon = self.epsilon + + # Prepare integration times (numpy only) + integrator_times = times.rescale("tdb").jd().to_numpy() + integrator_times = integrator_times - ephem.jd_ref + + # Integration loop (preallocate state array) + N = len(integrator_times) + if N == 0: + return VariantOrbits.empty() if is_variant else Orbits.empty() + + xyzvxvyvz = np.zeros((N, 6), dtype="float64") + scratch = np.zeros((1, 6), dtype="float64") + + for i in range(N): + sim.integrate(integrator_times[i]) + scratch.fill(0.0) + sim.serialize_particle_data(xyzvxvyvz=scratch) + xyzvxvyvz[i, :] = scratch[0, :] + + # Build results + jd_times = integrator_times + ephem.jd_ref + times_out = Timestamp.from_jd(jd_times, scale="tdb") + origin_codes = Origin.from_kwargs( + code=pa.repeat("SOLAR_SYSTEM_BARYCENTER", xyzvxvyvz.shape[0]) + ) + + if is_variant: + orbit_ids_out = [orbit_id] * N + variant_ids_out = [variant_id] * N + object_id_out = np.tile(orbit.object_id.to_numpy(zero_copy_only=False), N) + + return VariantOrbits.from_kwargs( + orbit_id=orbit_ids_out, + variant_id=variant_ids_out, + object_id=object_id_out, + weights=orbit.weights, + weights_cov=orbit.weights_cov, + coordinates=CartesianCoordinates.from_kwargs( + x=xyzvxvyvz[:, 0], + y=xyzvxvyvz[:, 1], + z=xyzvxvyvz[:, 2], + vx=xyzvxvyvz[:, 3], + vy=xyzvxvyvz[:, 4], + vz=xyzvxvyvz[:, 5], + time=times_out, + origin=origin_codes, + frame="equatorial", + ), + ) + else: + orbit_ids_out = [orbit_id] * N + object_id_out = np.tile(orbit.object_id.to_numpy(zero_copy_only=False), N) + + return Orbits.from_kwargs( + coordinates=CartesianCoordinates.from_kwargs( + x=xyzvxvyvz[:, 0], + y=xyzvxvyvz[:, 1], + z=xyzvxvyvz[:, 2], + vx=xyzvxvyvz[:, 3], + vy=xyzvxvyvz[:, 4], + vz=xyzvxvyvz[:, 5], + time=times_out, + origin=origin_codes, + frame="equatorial", + ), + orbit_id=orbit_ids_out, + object_id=object_id_out, + ) + def _propagate_orbits_inner( self, orbits: OrbitType, times: TimestampType ) -> OrbitType: @@ -225,20 +367,24 @@ def _propagate_orbits_inner( particle_ids = np.array(particle_ids, dtype="object") orbit_id_mapping, uint_orbit_ids = hash_orbit_ids_to_uint32(particle_ids) + hash_to_index = {uint_orbit_ids[i].value: i for i in range(len(uint_orbit_ids))} # Add the orbits as particles to the simulation - coords_df = orbits.coordinates.to_dataframe() + # OPTIMIZED: Use direct array access instead of DataFrame conversion + coords = orbits.coordinates + position_arrays = coords.r # x, y, z columns + velocity_arrays = coords.v # vx, vy, vz columns assist.Extras(sim, ephem) - for i in range(len(coords_df)): + for i in range(len(position_arrays)): sim.add( - x=coords_df.x[i], - y=coords_df.y[i], - z=coords_df.z[i], - vx=coords_df.vx[i], - vy=coords_df.vy[i], - vz=coords_df.vz[i], + x=position_arrays[i, 0], + y=position_arrays[i, 1], + z=position_arrays[i, 2], + vx=velocity_arrays[i, 0], + vy=velocity_arrays[i, 1], + vz=velocity_arrays[i, 2], hash=uint_orbit_ids[i], ) @@ -271,18 +417,20 @@ def _propagate_orbits_inner( step_states.append(step_xyzvxvyvz) if is_variant: - particle_ids = [orbit_id_mapping[h] for h in orbit_id_hashes] - orbit_ids, variant_ids = zip( - *[pid.split(separator) for pid in particle_ids] + indices = np.fromiter( + (hash_to_index[h] for h in orbit_id_hashes), + dtype=np.int64, + count=sim.N, ) - step_orbit_ids.append(np.asarray(orbit_ids, dtype=object)) - step_variant_ids.append(np.asarray(variant_ids, dtype=object)) + step_orbit_ids.append(orbit_ids[indices]) + step_variant_ids.append(variant_ids[indices]) else: - step_orbit_ids.append( - np.asarray( - [orbit_id_mapping[h] for h in orbit_id_hashes], dtype=object - ) + indices = np.fromiter( + (hash_to_index[h] for h in orbit_id_hashes), + dtype=np.int64, + count=sim.N, ) + step_orbit_ids.append(particle_ids[indices]) # Build a single result table if len(step_states) == 0: @@ -409,20 +557,24 @@ def _detect_collisions( particle_ids = np.array(particle_ids, dtype="object") orbit_id_mapping, uint_orbit_ids = hash_orbit_ids_to_uint32(particle_ids) + {uint_orbit_ids[i].value: i for i in range(len(uint_orbit_ids))} # Add the orbits as particles to the simulation - coords_df = orbits.coordinates.to_dataframe() + # OPTIMIZED: Use direct array access instead of DataFrame conversion + coords = orbits.coordinates + position_arrays = coords.r # x, y, z columns + velocity_arrays = coords.v # vx, vy, vz columns assist.Extras(sim, ephem) - for i in range(len(coords_df)): + for i in range(len(position_arrays)): sim.add( - x=coords_df.x[i], - y=coords_df.y[i], - z=coords_df.z[i], - vx=coords_df.vx[i], - vy=coords_df.vy[i], - vz=coords_df.vz[i], + x=position_arrays[i, 0], + y=position_arrays[i, 1], + z=position_arrays[i, 2], + vx=velocity_arrays[i, 0], + vy=velocity_arrays[i, 1], + vz=velocity_arrays[i, 2], hash=uint_orbit_ids[i], )