Exemplo n.º 1
0
def run_secularmultiple(particle_set, end_time, start_time=(0 |units.Myr), \
                        N_output=100, debug_mode=False, genT4System=False, \
                        exportData=True, useAMD=True, GCode = None, \
                        KeySystemID=None, SEVCode=None):
    '''Does what it says on the tin.'''
    try:
        hierarchical_test = [x for x in particle_set if x.is_binary == True]
        print("The supplied set has", len(hierarchical_test),
              "node particles and is a tree.")
        py_particles = particle_set
    except:
        print("The supplied set is NOT a tree set! Building tree ...")
        py_particles = get_full_hierarchical_structure(particle_set,
                                                       KeySystemID=KeySystemID,
                                                       SEVCode=SEVCode)
        hierarchical_test = [x for x in py_particles if x.is_binary == True]
        print("Tree has been built with", len(hierarchical_test),
              "node particles.")
    nodes = py_particles.select(lambda x: x == True, ["is_binary"])
    Num_nodes = len(nodes)
    stellarCollisionOccured = False

    if GCode == None:
        code = SecularMultiple()
    else:
        code = GCode

    if exportData:
        plot_a_AU = defaultdict(list)
        plot_e = defaultdict(list)
        plot_peri_AU = defaultdict(list)
        plot_stellar_inc_deg = defaultdict(list)
        if useAMD:
            plot_AMDBeta = defaultdict(list)
        plot_times_Myr = []

    if genT4System:
        print(nodes.inclination)
        nodes.inclination = [-18.137, 0.0, 23.570] | units.deg
        print(nodes.inclination)

    if debug_mode:
        print('=' * 50)
        print('t/kyr', 0.00)
        print('a/AU', nodes.semimajor_axis)
        print('p/day', nodes.period)
        print('e', nodes.eccentricity)
        print('i/deg', nodes.inclination)
        print('AP/deg', \
            nodes.argument_of_pericenter)
        print('LAN/deg', \
            nodes.longitude_of_ascending_node)
    #print(py_particles)
    code.particles.add_particles(py_particles)
    #code.commit_particles()
    #print(code.particles.semimajor_axis)

    code.model_time = start_time
    #print(py_particles.id, py_particles.semimajor_axis)
    channel_from_particles_to_code = py_particles.new_channel_to(
        code.particles)
    channel_from_code_to_particles = code.particles.new_channel_to(
        py_particles)
    #print(py_particles.id, py_particles.semimajor_axis)

    channel_from_particles_to_code.copy(
    )  #copy_attributes(['semimajor_axis', 'eccentricity', \
    #'longitude_of_ascending_node', 'argument_of_pericenter', 'inclination'])
    #print('This is After the First Channel Copy:', code.particles.semimajor_axis)
    time = start_time
    if useAMD:
        #print(py_particles.id, py_particles.mass)
        jovianParent = get_jovian_parent(py_particles)
        output_time_step = 1000 * jovianParent.period.value_in(
            units.Myr) | units.Myr
        PS = initialize_PlanetarySystem_from_HierarchicalSet(py_particles)
        PS.get_SystemBetaValues()
        if exportData:
            for planet in PS.planets:
                plot_AMDBeta[planet.id].append(planet.AMDBeta)
    else:
        output_time_step = end_time / float(N_output)

    if exportData:
        plot_times_Myr.append(time.value_in(units.Myr))
        for i, node in enumerate(nodes):
            plot_a_AU[node.child2.id].append(
                node.semimajor_axis.value_in(units.AU))
            plot_e[node.child2.id].append(node.eccentricity)
            plot_peri_AU[node.child2.id].append(
                node.semimajor_axis.value_in(units.AU) *
                (1.0 - node.eccentricity))
            plot_stellar_inc_deg[node.child2.id].append(
                node.inclination.value_in(units.deg))
    counter = 0
    while time <= end_time:
        #print('Start of Time Loop')
        #print(output_time_step)
        time += output_time_step
        counter += 1
        #print(time)
        #print(code.model_time)
        #print(code.particles.semimajor_axis)
        code.evolve_model(time)
        #print('Evolved model to:', time.value_in(units.Myr), "Myr")
        #print(code.particles.semimajor_axis)
        channel_from_code_to_particles.copy()
        #channel_from_code_to_particles.copy_attributes(['semimajor_axis', 'eccentricity', \
        #'longitude_of_ascending_node', 'argument_of_pericenter', 'inclination'])
        #print('Hello')
        py_particles.time = time
        if exportData:
            plot_times_Myr.append(time.value_in(units.Myr))
            nodes = py_particles.select(lambda x: x == True, ["is_binary"])
            for i, node in enumerate(nodes):
                plot_a_AU[node.child2.id].append(
                    node.semimajor_axis.value_in(units.AU))
                plot_e[node.child2.id].append(node.eccentricity)
                plot_peri_AU[node.child2.id].append(
                    node.semimajor_axis.value_in(units.AU) *
                    (1.0 - node.eccentricity))
                plot_stellar_inc_deg[node.child2.id].append(
                    node.inclination.value_in(units.deg))

        if time == end_time + output_time_step:
            map_node_oe_to_lilsis(py_particles)

        if debug_mode:
            if time == end_time or time == output_time_step:
                print('=' * 50)
                print('t/kyr', time.value_in(units.kyr))
                print('a/AU', nodes.semimajor_axis)
                print('p/day', nodes.period)
                print('e', nodes.eccentricity)
                print('i/deg', nodes.inclination)
                print('AP/deg', \
                    nodes.argument_of_pericenter)
                print('LAN/deg', \
                    nodes.longitude_of_ascending_node)
        # Check for Planet Destruction from Star
        #print(py_particles.id, py_particles.semimajor_axis)
        temp = check_for_stellar_collision(py_particles)
        #print(temp)
        # Returns 'None' if No Destruction!
        if temp != None:
            code.stop()
            code = SecularMultiple()
            code.model_time = time
            py_particles = Particles()
            py_particles.add_particles(temp)
            code.particles.add_particles(py_particles)
            py_particles.time = time
            #code.commit_particles()
            channel_from_particles_to_code = py_particles.new_channel_to(
                code.particles)
            channel_from_code_to_particles = code.particles.new_channel_to(
                py_particles)
            channel_from_particles_to_code.copy()
            nodes = py_particles.select(lambda x: x == True, ["is_binary"])
            stellarCollisionOccured = True
            if useAMD:
                PS = initialize_PlanetarySystem_from_HierarchicalSet(
                    py_particles)
            #channel_from_code_to_particles.copy_attributes(['semimajor_axis', 'eccentricity', \
            #'longitude_of_ascending_node', 'argument_of_pericenter', 'inclination'])
        #print(code.particles.semimajor_axis)
        # AMD Checking
        if useAMD:
            PS = update_oe_for_PlanetarySystem(PS, py_particles)
            PS.get_SystemBetaValues()
            if exportData:
                for planet in PS.planets:
                    plot_AMDBeta[planet.id].append(planet.AMDBeta)
            if counter % 100 == 0 and len(
                    PS.planets.select(lambda x: x < 1.0, ["AMDBeta"])) > 1:
                break

    if GCode == None:
        code.stop()
    else:
        code = reset_secularmultiples(code)
    if exportData:
        if useAMD:
            data = plot_times_Myr, plot_a_AU, plot_e, plot_peri_AU, plot_stellar_inc_deg, plot_AMDBeta
        else:
            data = plot_times_Myr, plot_a_AU, plot_e, plot_peri_AU, plot_stellar_inc_deg
    else:
        data = None
    # Set the Output Code to be the New Code if a Stellar Collision Occured.
    if stellarCollisionOccured:
        newcode = code
    else:
        newcode = None
    return py_particles, data, newcode
Exemplo n.º 2
0
class CloseEncounters():
    def __init__(self, Star_EncounterHistory, KeplerWorkerList = None, \
                 NBodyWorkerList = None, SecularWorker = None, SEVWorker = None):
        '''EncounterHistory should be a List of the Format {RotationKey: [Encounter0_FilePath, ...]}'''
        # Find the Main System's Host Star's ID and Assign it to 'KeySystemID'
        self.doEncounterPatching = True
        self.KeySystemID = int(Star_EncounterHistory[list(
            Star_EncounterHistory)[0]][0].split("/")[-2])
        self.ICs = defaultdict(list)
        self.StartTimes = defaultdict(list)
        self.desired_endtime = 1.0 | units.Gyr
        self.max_end_time = 0.1 | units.Myr
        self.kep = KeplerWorkerList
        self.NBodyCodes = NBodyWorkerList
        self.SecularCode = SecularWorker
        self.SEVCode = SEVWorker
        self.getOEData = True
        self.OEData = defaultdict(list)
        # Create a List of StartingTimes and Encounter Initial Conditions (ICs) for all Orientations
        for RotationKey in Star_EncounterHistory.keys():
            for i, Encounter in enumerate(Star_EncounterHistory[RotationKey]):
                self.ICs[RotationKey].append(
                    read_set_from_file(Encounter,
                                       format="hdf5",
                                       version='2.0',
                                       close_file=True))
                self.StartTimes[RotationKey].append(
                    np.max(np.unique(self.ICs[RotationKey][i].time)))
        #print(self.KeySystemID)
        self.FinalStates = defaultdict(list)

    def SimAllEncounters(self):
        ''' Call this function to run all Encounters for a System.'''
        # Start up Kepler Functions if Needed
        if self.kep == None:
            self.kep = []
            bodies = self.ICs[next(iter(self.ICs))][0]
            converter = nbody_system.nbody_to_si(
                bodies.mass.sum(),
                2 * np.max(bodies.radius.number) | bodies.radius.unit)
            self.kep.append(
                Kepler(unit_converter=converter, redirection='none'))
            self.kep.append(
                Kepler(unit_converter=converter, redirection='none'))
            self.kep[0].initialize_code()
            self.kep[1].initialize_code()
        # Start up NBodyCodes if Needed
        if self.NBodyCodes == None:
            self.NBodyCodes = [
                initialize_GravCode(ph4),
                initialize_isOverCode()
            ]
        # Start up SecularCode if Needed
        if self.SecularCode == None:
            self.SecularCode = SecularMultiple()
        if self.SEVCode == None:
            self.SEVCode = SSE()

        # Begin Looping over Rotation Keys ...
        for RotationKey in self.ICs.keys():
            for i in range(len(self.ICs[RotationKey])):
                try:
                    print(util.timestamp(), "!!! UPDATE: Starting the following encounter", \
                          "Star", self.KeySystemID, RotationKey, "-", i)

                    # Identify the Current Encounter in the List for This Rotation
                    CurrentEncounter = self.ICs[RotationKey][i]
                    #print(CurrentEncounter[0].position)

                    # Create the Encounter Instance with the Current Encounter
                    Encounter_Inst = self.SingleEncounter(CurrentEncounter)

                    # Simulate the Encounter till the Encounter is Over via N-Body Integrator
                    # -OR- the time to the Next Encounter is Reached
                    if len(self.StartTimes[RotationKey]) == 1 or i + 1 == len(
                            self.StartTimes[RotationKey]):
                        current_max_endtime = self.max_end_time
                    else:
                        current_max_endtime = self.StartTimes[RotationKey][i +
                                                                           1]
                    EndingState = Encounter_Inst.SimSingleEncounter(current_max_endtime, \
                                                                    start_time = self.StartTimes[RotationKey][i], \
                                                                    GCodes = self.NBodyCodes)
                    EndingStateTime = np.max(np.unique(EndingState.time))

                    #print(Encounter_Inst.particles[0].position)
                    print("The Encounter was over after:",
                          (EndingStateTime -
                           self.StartTimes[RotationKey][i]).value_in(
                               units.Myr))

                    #print(EndingState.id, EndingState.x)
                    print('----------')
                    #print(Encounter_Inst.particles.id, Encounter_Inst.particles.x)

                    # Strip off Anything Not Associated with the Key System
                    systems_in_current_encounter = stellar_systems.get_heirarchical_systems_from_set(
                        EndingState, kepler_workers=self.kep)

                    # Reassign the EndingState to include the Primary System ONLY
                    EndingState = systems_in_current_encounter[
                        self.KeySystemID]
                    print("Before Secular:", EndingState.id, EndingState.x)
                    #print(EndingState[0].position)

                    #print(len(self.ICs[RotationKey])-1)
                    # If Encounter Patching is Desired -AND- it isn't the last Encounter
                    if i + 1 < len(self.ICs[RotationKey]
                                   ) and self.doEncounterPatching:

                        # Identify the Next Encounter in the List
                        NextEncounter = self.ICs[RotationKey][i + 1]

                        # Simulate System till the Next Encounter's Start Time
                        Encounter_Inst = self.SingleEncounter(EndingState)
                        FinalState, data, newcode = Encounter_Inst.SimSecularSystem(self.StartTimes[RotationKey][i+1], \
                                                                     start_time = EndingStateTime, \
                                                                     GCode = self.SecularCode, getOEData=self.getOEData, \
                                                                     KeySystemID = self.KeySystemID, SCode=self.SEVCode)
                        if newcode != None:
                            self.SecularCode = newcode
                        print("After Secular:", FinalState.id)
                        # Begin Patching of the End State to the Next Encounter
                        self.ICs[RotationKey][i + 1] = self.PatchedEncounter(
                            FinalState, NextEncounter)
                    else:
                        # Simulate System till Desired Global Endtime
                        #print(CurrentEncounter[0].time.value_in(units.Myr))
                        #print(EndingState[0].time.value_in(units.Myr))
                        Encounter_Inst = self.SingleEncounter(EndingState)
                        FinalState, data, newcode = Encounter_Inst.SimSecularSystem(self.desired_endtime, \
                                                                     start_time = EndingStateTime, \
                                                                     GCode = self.SecularCode, getOEData=self.getOEData, \
                                                                     KeySystemID = self.KeySystemID, SCode=self.SEVCode)
                        if newcode != None:
                            self.SecularCode = newcode
                        print("After Secular:", FinalState.id)

                    # Append the FinalState of Each Encounter to its Dictionary
                    self.FinalStates[RotationKey].append(FinalState)
                    if self.getOEData and data != None:
                        self.OEData[RotationKey].append(data)
                except:
                    print("!!!! Alert: Skipping", RotationKey, "-", i,
                          "for Star", self.KeySystemID,
                          "due to unforseen issues!")
                    print("!!!!        The Particle Set's IDs are as follows:",
                          self.ICs[RotationKey][i].id)

        # Stop the NBody Codes if not Provided
        if self.kep == None:
            self.kep[0].stop()
            self.kep[1].stop()
        # Start up NBodyCodes if Needed
        if self.NBodyCodes == None:
            self.NBodyCodes[0].stop()
            self.NBodyCodes[1].stop()
        # Start up SecularCode if Needed
        if self.SecularCode == None:
            self.SecularCode.stop()
        return None

    def PatchedEncounter(self, EndingState, NextEncounter):
        ''' Call this function to Patch Encounter Endstates to the Next Encounter'''
        # Determine Time to Next Encounter
        current_time = max(EndingState.time)
        final_time = max(NextEncounter.time)

        # Map the Orbital Elements to the Child2 Particles (LilSis) [MUST BE A TREE SET]
        enc_patching.map_node_oe_to_lilsis(EndingState)

        # Seperate Next Encounter Systems to Locate the Primary System
        systems_at_next_encounter = stellar_systems.get_heirarchical_systems_from_set(
            NextEncounter)
        sys_1 = systems_at_next_encounter[self.KeySystemID]
        # Note: This was changed to handle encounters of which result in one
        #       bound object of multiple subsystems. ~ Joe G. | 8/24/20
        BoundObjOnly = False
        if len(systems_at_next_encounter.keys()) == 1:
            BoundObjOnly = True
        else:
            secondary_sysID = [
                key for key in list(systems_at_next_encounter.keys())
                if key != int(self.KeySystemID)
            ][0]
            sys_2 = systems_at_next_encounter[secondary_sysID]

        # Get Planet and Star Subsets for the Current and Next Encounter
        children_at_EndingState = EndingState.select(lambda x: x == False,
                                                     ["is_binary"])
        planets_at_current_encounter = util.get_planets(
            children_at_EndingState)
        hoststar_at_current_encounter = util.get_stars(
            children_at_EndingState).select(lambda x: x == self.KeySystemID,
                                            ["id"])[0]
        planets_at_next_encounter = util.get_planets(sys_1)
        print("Planets at Next Encount:", planets_at_next_encounter.id)
        hoststar_at_next_encounter = util.get_stars(sys_1).select(
            lambda x: x == self.KeySystemID, ["id"])[0]
        #print(hoststar_at_next_encounter)

        # Update Current Positions & Velocitys to Relative Coordinates from Orbital Parameters!!
        # TO-DO: Does not handle Binary Star Systems
        for planet in planets_at_current_encounter:
            #print(planet.id, planet.position)
            nbody_PlanetStarPair = \
            new_binary_from_orbital_elements(hoststar_at_current_encounter.mass, planet.mass, planet.semimajor_axis, \
                                             eccentricity = planet.eccentricity, inclination=planet.inclination, \
                                             longitude_of_the_ascending_node=planet.longitude_of_ascending_node, \
                                             argument_of_periapsis=planet.argument_of_pericenter, G=units.constants.G, \
                                             true_anomaly = 360*rp.uniform(0.0,1.0) | units.deg) # random point in the orbit
            planet.position = nbody_PlanetStarPair[1].position
            planet.velocity = nbody_PlanetStarPair[1].velocity
            #print(planet.id, planet.position)

        for planet in planets_at_current_encounter:
            print(planet.id, planet.position)

        # Release a Warning when Odd Planet Number Combinations Occur (Very Unlikely, More of a Safe Guard)
        if len(planets_at_current_encounter) != len(planets_at_next_encounter):
            print("!!!! Expected",
                  len(planets_at_next_encounter), "planets but recieved only",
                  len(planets_at_current_encounter))

        # Move Planets to Host Star in the Next Encounter
        for next_planet in planets_at_next_encounter:
            for current_planet in planets_at_current_encounter:
                if next_planet.id == current_planet.id:
                    next_planet.position = current_planet.position + hoststar_at_next_encounter.position
                    next_planet.velocity = current_planet.velocity + hoststar_at_next_encounter.velocity
                    break

        #for planet in planets_at_next_encounter:
        #    print(planet.id, planet.position)
        #for particle in sys_1:
        #    print(particle.id, particle.position)

        # Recombine Seperated Systems to Feed into SimSingleEncounter
        UpdatedNextEncounter = Particles()
        print("IDs in System 1", sys_1.id)
        UpdatedNextEncounter.add_particles(sys_1)
        if not BoundObjOnly:
            UpdatedNextEncounter.add_particles(sys_2)
            print("IDs in System 2", sys_2.id)

        # Return the Updated and Patched Encounter as a Partcile Set for the N-Body Simulation
        return UpdatedNextEncounter

    class SingleEncounter():
        def __init__(self, EncounterBodies):
            self.particles = EncounterBodies

        def SimSecularSystem(self, desired_end_time, **kwargs):
            start_time = kwargs.get("start_time", 0 | units.Myr)
            getOEData = kwargs.get("getOEData", False)
            KeySystemID = kwargs.get("KeySystemID", None)
            GCode = kwargs.get("GCode", None)
            SEVCode = kwargs.get("SCode", None)

            self.particles, data, newcode = enc_patching.run_secularmultiple(self.particles, desired_end_time, \
                                                                  start_time = start_time, N_output=1, \
                                                                  GCode=GCode, exportData=getOEData, \
                                                                  KeySystemID=KeySystemID, SEVCode=SEVCode)
            return self.particles, data, newcode

        def SimSingleEncounter(self, max_end_time, **kwargs):
            delta_time = kwargs.get("delta_time", 100 | units.yr)
            converter = kwargs.get("converter", None)
            start_time = kwargs.get("start_time", 0 | units.yr)
            doStateSaves = kwargs.get("doStateSaves", False)
            doVerbosSaves = kwargs.get("doEncPatching", False)
            GCodes = kwargs.get("GCodes", None)

            GravitatingBodies = self.particles

            # Set Up the Integrators
            if GCodes == None:
                if converter == None:
                    converter = nbody_system.nbody_to_si(GravitatingBodies.mass.sum(), \
                                                         2 * np.max(GravitatingBodies.radius.number) | GravitatingBodies.radius.unit)
                gravity = initialize_GravCode(ph4, converter=converter)
                over_grav = initialize_isOverCode(converter=converter)
            else:
                gravity = GCodes[0]
                over_grav = GCodes[1]

            # Set up SaveState if Requested
            if doStateSaves:
                pass  # Currently Massive Memory Leak Present

            # Store Initial Center of Mass Information
            rCM_i = GravitatingBodies.center_of_mass()
            vCM_i = GravitatingBodies.center_of_mass_velocity()

            # Remove Attributes that Cause Issues with SmallN
            if 'child1' in GravitatingBodies.get_attribute_names_defined_in_store(
            ):
                del GravitatingBodies.child1, GravitatingBodies.child2

            # Moving the Encounter's Center of Mass to the Origin and Setting it at Rest
            GravitatingBodies.position -= rCM_i
            GravitatingBodies.velocity -= vCM_i

            # Add and Commit the Scattering Particles
            gravity.particles.add_particles(
                GravitatingBodies)  # adds bodies to gravity calculations
            gravity.commit_particles()
            #gravity.begin_time = start_time

            # Create the Channel to Python Set & Copy it Over
            channel_from_grav_to_python = gravity.particles.new_channel_to(
                GravitatingBodies)
            channel_from_grav_to_python.copy()

            # Get Free-Fall Time for the Collision
            s = util.get_stars(GravitatingBodies)
            t_freefall = s.dynamical_timescale()
            #print(gravity.begin_time)
            #print(t_freefall)

            # Setting Coarse Timesteps
            list_of_times = np.arange(0.0, start_time.value_in(units.yr)+max_end_time.value_in(units.yr), \
                                      delta_time.value_in(units.yr)) | units.yr
            stepNumber = 0

            # Loop through the List of Coarse Timesteps
            for current_time in list_of_times:
                #print(current_time)
                #print(GravitatingBodies.time)
                #print(gravity.sync_time)
                # Evolve the Model to the Desired Current Time
                gravity.evolve_model(current_time)

                # Update Python Set in In-Code Set
                channel_from_grav_to_python.copy()  # original
                channel_from_grav_to_python.copy_attribute(
                    "index_in_code", "id")

                # Check to See if the Encounter is Over After the Freefall Time and Every 25 Steps After That
                if current_time > 1.25 * t_freefall and stepNumber % 25 == 0:
                    over = util.check_isOver(gravity.particles, over_grav)
                    print("Is it Over?", over)
                    if over:
                        #print(gravity.particles[0].position)
                        #print(GravitatingBodies[0].position)
                        current_time += 100 | units.yr
                        # Get to a Final State After Several Planet Orbits
                        gravity.evolve_model(current_time)
                        gravity.update_particle_set()
                        gravity.particles.synchronize_to(GravitatingBodies)
                        channel_from_grav_to_python.copy()
                        GravitatingBodies.time = start_time + current_time
                        # Create a Save State
                        break
                if current_time == list_of_times[-1]:
                    # Create a Save State
                    pass
                stepNumber += 1
            if GCodes == None:
                # Stop the Gravity Code Once the Encounter Finishes
                gravity.stop()
                over_grav.stop()
            else:
                # Reset the Gravity Codes Once Encounter Finishes
                gravity.reset()
                over_grav.reset()

            # Return the GravitatingBodies as they are at the End of the Simulation
            return GravitatingBodies