From 1378ec0d6ee35c5ae5101e6e8cd45ab07219b1b9 Mon Sep 17 00:00:00 2001 From: Rob Farmer Date: Fri, 12 May 2023 13:03:43 +0200 Subject: [PATCH] Add finer controls for timestepping --- .../mesa_r15140/examples/photo_load.py | 54 ++++++++ src/amuse/community/mesa_r15140/interface.f90 | 72 +++++++++++ src/amuse/community/mesa_r15140/interface.py | 115 +++++++++++++++++- .../community/mesa_r15140/mesa_interface.f90 | 103 ++++++++++++++++ 4 files changed, 342 insertions(+), 2 deletions(-) create mode 100644 src/amuse/community/mesa_r15140/examples/photo_load.py diff --git a/src/amuse/community/mesa_r15140/examples/photo_load.py b/src/amuse/community/mesa_r15140/examples/photo_load.py new file mode 100644 index 0000000000..9ae637f77b --- /dev/null +++ b/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) \ No newline at end of file diff --git a/src/amuse/community/mesa_r15140/interface.f90 b/src/amuse/community/mesa_r15140/interface.f90 index a642684326..85e5b54bf4 100644 --- a/src/amuse/community/mesa_r15140/interface.f90 +++ b/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 \ No newline at end of file diff --git a/src/amuse/community/mesa_r15140/interface.py b/src/amuse/community/mesa_r15140/interface.py index 38a5f81dcf..980c1b3ef6 100644 --- a/src/amuse/community/mesa_r15140/interface.py +++ b/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() diff --git a/src/amuse/community/mesa_r15140/mesa_interface.f90 b/src/amuse/community/mesa_r15140/mesa_interface.f90 index 5c4c26a482..a3c4d804c7 100644 --- a/src/amuse/community/mesa_r15140/mesa_interface.f90 +++ b/src/amuse/community/mesa_r15140/mesa_interface.f90 @@ -569,6 +569,95 @@ subroutine do_evolve_one_step(id, result, ierr) end subroutine do_evolve_one_step + integer function do_prepare_for_redo(id) + integer, intent(in) :: id + do_prepare_for_redo = star_prepare_to_redo(id) + end function do_prepare_for_redo + + integer function do_prepare_for_retry(id) + integer, intent(in) :: id + do_prepare_for_retry = star_prepare_to_retry(id) + end function do_prepare_for_retry + + + subroutine do_solve_one_step_pre(id, result, ierr) + integer, intent(in) :: id + integer, intent(out) :: result, ierr + type (star_info), pointer :: s + + integer :: model_number + logical :: continue_evolve_loop, first_try + logical,parameter :: dbg=.false. + + ierr = MESA_SUCESS + + call star_ptr(id, s, ierr) + if (failed('star_ptr',ierr)) return + + call before_step_loop(s% id, ierr) + if (failed('before_step_loop',ierr)) return + + result = s% extras_start_step(id) + if (result /= keep_going) return + + + end subroutine do_solve_one_step_pre + + + subroutine do_solve_one_step(id, first_try, result, ierr) + integer, intent(in) :: id + logical, intent(in) :: first_try + integer, intent(out) :: result, ierr + type (star_info), pointer :: s + + integer :: model_number + logical :: continue_evolve_loop + logical,parameter :: dbg=.false. + + ierr = MESA_SUCESS + + call star_ptr(id, s, ierr) + if (failed('star_ptr',ierr)) return + + + result = star_evolve_step(id, first_try) + if (result == keep_going) result = star_check_model(id) + if (result == keep_going) result = s% extras_check_model(id) + if (result == keep_going) result = star_pick_next_timestep(id) + if (result == keep_going) return + + end subroutine do_solve_one_step + + + subroutine do_solve_one_step_post(id, result, ierr) + integer, intent(in) :: id + integer, intent(out) :: result, ierr + type (star_info), pointer :: s + + integer :: model_number + logical :: continue_evolve_loop, first_try + logical,parameter :: dbg=.false. + + ierr = MESA_SUCESS + + call star_ptr(id, s, ierr) + if (failed('star_ptr',ierr)) return + + + s% doing_first_model_of_run = .false. + + call after_step_loop(s% id, s% inlist_fname, & + dbg, result, ierr) + if (failed('after_step_loop',ierr)) return + + call do_saves(id, ierr) + if (failed('do_saves',ierr)) return + + call flush() + + end subroutine do_solve_one_step_post + + subroutine evolve_until(id, delta_t, ierr) integer, intent(in) :: id real(dp), intent(in) :: delta_t @@ -639,6 +728,7 @@ subroutine set_model_number(id, new_number, ierr) s% model_number = new_number end subroutine set_model_number + subroutine set_timestep(id, dt_next, ierr) integer, intent(in) :: id real(dp), intent(in) :: dt_next @@ -649,6 +739,19 @@ subroutine set_timestep(id, dt_next, ierr) end subroutine set_timestep + subroutine set_current_dt(id, dt, ierr) + integer, intent(in) :: id + real(dp), intent(in) :: dt + integer, intent(out) :: ierr + type (star_info), pointer :: s + + call star_ptr(id, s, ierr) + if (failed('set_dt',ierr)) return + + s% dt = dt + + end subroutine set_current_dt + subroutine set_new_mass(id, new_mass, ierr) integer, intent(in) :: id real(dp), intent(in) :: new_mass