diff --git a/grss/prop/__init__.py b/grss/prop/__init__.py index c7d0ace2..0d7db8ba 100644 --- a/grss/prop/__init__.py +++ b/grss/prop/__init__.py @@ -1,4 +1,5 @@ """GRSS orbit propagation subpackage""" +from .prop_parallel import * from .prop_simulation import * from .prop_unscented import * from .prop_utils import * diff --git a/grss/prop/prop_parallel.py b/grss/prop/prop_parallel.py new file mode 100644 index 00000000..ffbe77dc --- /dev/null +++ b/grss/prop/prop_parallel.py @@ -0,0 +1,55 @@ +"""Parallel computing utilities for the GRSS orbit propagation code""" +import pandas as pd + +__all__ = [ 'cluster_ca_or_impacts', +] + +def cluster_ca_or_impacts(full_list, max_duration=45, central_body=399): + """ + Cluster a list of close approaches by time and uniqueness. + + Parameters + ---------- + full_list : list of prop_simulation.CloseApproachParameters objects + List of close approaches to cluster. + max_duration : float + Maximum duration (in days) between close approaches in a cluster. + central_body : int + SPICE ID of the central body. + + Returns + ------- + all_clusters : tuple of list of prop_simulation.CloseApproachParameters objects + A tuple of close approach clusters (each cluster is a list of + close approaches). + """ + all_clusters = [] + full_list = [ca_or_impact for ca_or_impact in full_list + if ca_or_impact.centralBodySpiceId == central_body] + if not full_list: + return tuple(all_clusters) + times = [ca_or_impact.t for ca_or_impact in full_list] + bodies = [ca_or_impact.flybyBody for ca_or_impact in full_list] + idx_list = list(range(len(full_list))) + + df = pd.DataFrame({'time': times, 'body': bodies, 'idx': idx_list}) + df = df.sort_values(by=['time']) + + # create a new cluster if one of two conditions is met: + # 1. the time between the current close approach and the last close approach + # is greater than max_duration + # 2. the current close approach body is already in the current cluster + cluster = [full_list[df.iloc[0]['idx']]] + cluster_bodies = [df.iloc[0]['body']] + for i in range(1, len(df)): + time_condition = df.iloc[i]['time'] - df.iloc[i-1]['time'] > max_duration + body_condition = df.iloc[i]['body'] in cluster_bodies + if time_condition or body_condition: + all_clusters.append(cluster) + cluster = [full_list[df.iloc[i]['idx']]] + cluster_bodies = [df.iloc[i]['body']] + else: + cluster.append(full_list[df.iloc[i]['idx']]) + cluster_bodies.append(df.iloc[i]['body']) + all_clusters.append(cluster) + return tuple(all_clusters) diff --git a/grss/version.txt b/grss/version.txt index 58073ef8..fad066f8 100644 --- a/grss/version.txt +++ b/grss/version.txt @@ -1 +1 @@ -2.4.1 \ No newline at end of file +2.5.0 \ No newline at end of file