2626try :
2727 from mpi4py import MPI
2828except :
29- print (
30- "This example requires mpi4py. Please install with conda install mpi4py -c conda-forge"
31- )
29+ print ("This example requires mpi4py. Please install with conda install mpi4py -c conda-forge" )
3230
3331import pykep as pk
3432import paseos
4442# Now we will initialize MPI, for more details please refer to the mpi4py docs.
4543# In MPI "rank" indicates the index of the compute node (so 0-3 in our example).
4644comm = MPI .COMM_WORLD
47- assert (
48- comm .Get_size () == 4
49- ), "Please run the example with mpiexec -n 4 python mpi_example.py"
45+ assert comm .Get_size () == 4 , "Please run the example with mpiexec -n 4 python mpi_example.py"
5046rank = comm .Get_rank ()
5147other_ranks = [x for x in range (4 ) if x != rank ]
5248print (f"Started rank { rank } , other ranks are { other_ranks } " )
6561planet_list , sats_pos_and_v , _ = get_constellation (
6662 altitude , inclination , nSats , nPlanes , t0 , verbose = False
6763)
68- print (
69- f"Rank { rank } set up its orbit with altitude={ altitude } m and inclination={ inclination } deg"
70- )
64+ print (f"Rank { rank } set up its orbit with altitude={ altitude } m and inclination={ inclination } deg" )
7165
7266############ PASEOS INIT #############
7367# We will now initialize the PASEOS instance on each rank
7872local_actor = ActorBuilder .get_actor_scaffold (
7973 name = "Sat_" + str (rank ), actor_type = SpacecraftActor , epoch = t0
8074)
81- ActorBuilder .set_orbit (
82- actor = local_actor , position = pos , velocity = v , epoch = t0 , central_body = earth
83- )
75+ ActorBuilder .set_orbit (actor = local_actor , position = pos , velocity = v , epoch = t0 , central_body = earth )
8476
8577paseos_instance = paseos .init_sim (local_actor = local_actor )
8678print (f"Rank { rank } set up its PASEOS instance for its local actor { local_actor } " )
9789# Let's define the variable to track the actors we see
9890total_seen_actors = 0
9991
92+
10093# We will (ab)use PASEOS constraint function to track all the actors
10194# we see in an evaluation window (see timestep below).
10295# Turn on SHOW_ALL_WINDOWS if you want to see each window
@@ -135,7 +128,6 @@ def constraint_func(verbose=SHOW_ALL_WINDOWS):
135128
136129# Run until end of simulation
137130while t <= simulation_time :
138-
139131 # Advance the simulation state of this rank
140132 # Note how we pass the "constraint_func" to tell paseos
141133 # to track windows
@@ -147,9 +139,7 @@ def constraint_func(verbose=SHOW_ALL_WINDOWS):
147139 sys .stdout .flush () # update prints to better see parallelism
148140
149141 # Exchange actors between all ranks
150- exchange_actors (
151- comm , paseos_instance , local_actor , other_ranks , rank , verbose = SHOW_ALL_COMMS
152- )
142+ exchange_actors (comm , paseos_instance , local_actor , other_ranks , rank , verbose = SHOW_ALL_COMMS )
153143
154144# Wait until all ranks finished
155145print (f"Rank { rank } finished the simulation. Waiting for all to finish." )
0 commit comments