Add finer controls for timestepping (#956)

src/amuse/community/mesa_r15140/examples/photo_load.py

@@ -0,0 +1,54 @@
+import matplotlib.pyplot as plt
+import numpy as np
+
+from amuse.units import units
+from amuse.community.mesa.interface import MESA
+from amuse import datamodel
+
+stellar_evolution = MESA(version='15140')
+
+masses=[2.0] | units.MSun
+stars = datamodel.Particles(len(masses), mass=masses)
+
+stars = stellar_evolution.native_stars.add_particles(stars)
+star = stars[0]
+
+star.evolve_one_step()
+
+print(star.time_step)
+
+star.time_step = 1. | units.yr
+
+star.evolve_one_step()
+
+star.save_photo('photo')
+star.save_model('model.mod')
+
+
+star.evolve_one_step()
+
+age = star.age
+
+model_number = star.get_history('model_number')
+
+
+masses=[1.0] | units.MSun
+stars = datamodel.Particles(len(masses), mass=masses,filename=['photo'])
+
+stars = stellar_evolution.photo_stars.add_particles(stars)
+star = stars[0]
+
+star.evolve_one_step() # Photos will give exactly the same answer compared to a run without a save/load
+
+print(star.age, star.get_history('model_number'), star.mass)
+
+
+masses=[1.0] | units.MSun
+stars = datamodel.Particles(len(masses), mass=masses,filename=['model.mod'])
+
+stars = stellar_evolution.pre_built_stars.add_particles(stars)
+star = stars[0]
+
+star.evolve_one_step() # Models do not give exactly the same answer as a photo
+
+print(star.age, star.get_history('model_number'), star.mass)

src/amuse/community/mesa_r15140/interface.f90

@@ -1306,6 +1306,78 @@ end function get_gyre
 
 
 
+!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
+! Procedures for manually controlling how a step gets taken
+!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
+
+
+   integer function solve_one_step(AMUSE_ID, first_try, result)
+      integer, intent(in) :: AMUSE_ID
+      logical, intent(in) :: first_try
+      integer, intent(out) :: result
+      integer :: ierr 
+
+      ierr = 0
+
+      call do_solve_one_step(AMUSE_ID, first_try, result, ierr)
+      solve_one_step = ierr
+
+   end function solve_one_step
+
+   integer function solve_one_step_pre(AMUSE_ID, result)
+      integer, intent(in) :: AMUSE_ID
+      integer, intent(out) :: result
+      integer :: ierr 
+
+      call do_solve_one_step_pre(AMUSE_ID, result, ierr)
+      solve_one_step_pre = ierr
+
+   end function solve_one_step_pre
+
+
+   integer function solve_one_step_post(AMUSE_ID, result)
+      integer, intent(in) :: AMUSE_ID
+      integer, intent(out) :: result
+      integer :: ierr 
+
+      call do_solve_one_step_post(AMUSE_ID, result, ierr)
+      solve_one_step_post = ierr
+
+   end function solve_one_step_post
+
+
+   integer function prepare_retry_step(AMUSE_ID, dt, result)
+   ! Prepare toretake a timestep with a different dt
+   ! Does not actualy retry the step
+      integer, intent(in) :: AMUSE_ID
+      double precision,intent(in) :: dt
+      integer, intent(out) :: result
+      integer :: ierr 
+
+
+      call set_current_dt(AMUSE_ID, dt, ierr) ! Sets dt not dt_next
+      if(ierr/=MESA_SUCESS) then
+         prepare_retry_step = ierr
+         return
+      end if
+
+      result = do_prepare_for_retry(AMUSE_id)
+      prepare_retry_step = 0
+
+   end function prepare_retry_step
+
+
+   integer function prepare_redo_step(AMUSE_ID, result)
+   ! Retake the same timestep with the same dt. Useful when mdot changes
+   ! Does not actualy redo the step 
+      integer, intent(in) :: AMUSE_ID
+      integer, intent(out) :: result
+      integer :: ierr 
+
+      result = do_prepare_for_redo(AMUSE_id)
+      prepare_redo_step = 0
+
+   end function prepare_redo_step
 
 
 end module AMUSE_mesa

src/amuse/community/mesa_r15140/interface.py

@@ -914,7 +914,7 @@ def get_gyre(self, index_of_the_star, mode_l=0,
         """
         Get gyre data. 
 
-        This returns a list of dicts where each element of the list coresponds to one mode
+        This returns a list of dicts where each element of the list corresponds to one mode
         
         Each dict contains the pg,p,g and complex frequency for the mode as well as
         arrays of r/R, xi_r, xi_h, and dwdx for the mode
@@ -963,7 +963,6 @@ def get_mixing_length_ratio(self,index_of_the_star):
         """
         Retrieve the current mixing_length_ratio of the star.
         """
-
         return self.get_control(index_of_the_star,'mixing_length_alpha')
 
     def set_mixing_length_ratio(self,index_of_the_star, mixing_length_ratio):
@@ -1126,6 +1125,81 @@ def set_accrete_composition_metals(self, index_of_the_star, **kwargs):
         return r
 
 
+    @legacy_function
+    def solve_one_step():
+        """
+        Takes one step forward, but does not handle retries or redo's. 
+        Must call solve_one_step_pre first and solve_one_step_post afterwards
+        """
+        function = LegacyFunctionSpecification()
+        function.addParameter('index_of_the_star', dtype='int32', direction=function.IN
+            , description="The index of the star to get the value of")
+        function.addParameter('first_try', dtype='bool', direction=function.IN
+            , description="If this is the first attempt at taking this timestep")
+        function.addParameter('result', dtype='int32', direction=function.OUT
+            , description="What the star should do next (keep going, redo, retry, terminate)")
+        function.result_type = 'int32'
+        return function
+
+
+    @legacy_function
+    def solve_one_step_pre():
+        """
+        Prepares to takes one step forward, but does not handle retries or redo's. 
+        Called before solve_one_step
+        """
+        function = LegacyFunctionSpecification()
+        function.addParameter('index_of_the_star', dtype='int32', direction=function.IN
+            , description="The index of the star to get the value of")
+        function.addParameter('result', dtype='int32', direction=function.OUT
+            , description="What the star should do next (keep going, redo, retry, terminate)")
+        function.result_type = 'int32'
+        return function
+
+    @legacy_function
+    def solve_one_step_post():
+        """
+        After taking one step forward cleanup the model, but does not handle retries or redo's. 
+        Called after solve_one_step
+        """
+        function = LegacyFunctionSpecification()
+        function.addParameter('index_of_the_star', dtype='int32', direction=function.IN
+            , description="The index of the star to get the value of")
+        function.addParameter('result', dtype='int32', direction=function.OUT
+            , description="What the star should do next (keep going, redo, retry, terminate)")
+        function.result_type = 'int32'
+        return function
+
+    @legacy_function
+    def prepare_retry_step():
+        """
+        Prepares to retry a step with a new dt.
+        Does not actually take the step
+        """
+        function = LegacyFunctionSpecification()
+        function.addParameter('index_of_the_star', dtype='int32', direction=function.IN
+            , description="The index of the star to get the value of")
+        function.addParameter('dt_next', dtype='float64', direction=function.IN
+            , description="New timestep to try")
+        function.addParameter('result', dtype='int32', direction=function.OUT
+            , description="What the star should do next (keep going, redo, retry, terminate)")
+        function.result_type = 'int32'
+        return function
+
+    @legacy_function
+    def prepare_redo_step():
+        """
+        Prepares to redo a step (a step with the same dt but where you have changed other options like the mdot)
+        Does not actually take the step
+        """
+        function = LegacyFunctionSpecification()
+        function.addParameter('index_of_the_star', dtype='int32', direction=function.IN
+            , description="The index of the star to get the value of")
+        function.addParameter('result', dtype='int32', direction=function.OUT
+            , description="What the star should do next (keep going, redo, retry, terminate)")
+        function.result_type = 'int32'
+        return function
+
 
 class MESA(StellarEvolution, InternalStellarStructure):
 
@@ -1333,6 +1407,13 @@ def define_particle_sets(self, handler):
             handler.add_method(particle_set_name, 'get_blocker_wind_efficiency')
             handler.add_method(particle_set_name, 'set_blocker_wind_efficiency')
 
+            handler.add_method(particle_set_name, 'solve_one_step')
+            handler.add_method(particle_set_name, 'solve_one_step_pre')
+            handler.add_method(particle_set_name, 'solve_one_step_post')
+            handler.add_method(particle_set_name, 'prepare_retry_step')
+            handler.add_method(particle_set_name, 'prepare_redo_step')
+
+
     def define_state(self, handler):
         StellarEvolution.define_state(self, handler)
         handler.add_method('EDIT', 'new_pre_ms_particle')
@@ -1776,6 +1857,36 @@ def define_methods(self, handler):
             (handler.ERROR_CODE)
         )
 
+        handler.add_method(
+            "solve_one_step",
+            (handler.INDEX, handler.NO_UNIT),
+            (handler.NO_UNIT,handler.ERROR_CODE)
+        )
+
+        handler.add_method(
+            "solve_one_step_pre",
+            (handler.INDEX),
+            (handler.NO_UNIT,handler.ERROR_CODE)
+        )
+
+        handler.add_method(
+            "solve_one_step_post",
+            (handler.INDEX),
+            (handler.NO_UNIT,handler.ERROR_CODE)
+        )
+
+        handler.add_method(
+            "prepare_retry_step",
+            (handler.INDEX, units.s),
+            (handler.NO_UNIT,handler.ERROR_CODE)
+        )
+
+        handler.add_method(
+            "prepare_redo_step",
+            (handler.INDEX),
+            (handler.NO_UNIT,handler.ERROR_CODE)
+        )
+
     def initialize_module_with_default_parameters(self):
         self.parameters.set_defaults()
         self.initialize_code()