Merge pull request #28 from FormingWorlds/ls/fixed-rosenbrock

Ls/fixed rosenbrock

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "fwl-mors"
-version = "24.10.27"
+version = "24.11.18"
 description = "Stellar rotation and activity evolution model"
 readme = "README.md"
 authors = [
@@ -70,7 +70,7 @@ testpaths = ["tests"]
 
 [tool.bumpversion]
 # https://callowayproject.github.io/bump-my-version/howtos/calver/
-current_version = "24.10.27"
+current_version = "24.11.18"
 parse = """(?x)                     # Verbose mode
     (?P<release>                    # The release part
         (?:[1-9][0-9])\\.           # YY.

src/mors/parameters.py

@@ -18,15 +18,15 @@ def SetDefaultParameters(paramsDefault):
     """Sets up the default parameters."""
 
     # Time integration solver parameters
-    paramsDefault['TimeIntegrationMethod'] = 'ForwardEuler'   # options are 'ForwardEuler', 'RungeKutta4', 'RungeKuttaFehlberg', 'Rosenbrock'
+    paramsDefault['TimeIntegrationMethod'] = 'RosenbrockFixed'# options are 'ForwardEuler', 'RungeKutta4', 'RungeKuttaFehlberg', 'Rosenbrock' and 'RosenbrockFixed'
     paramsDefault['nStepMax'] = 10**6                         # maximum number of timesteps to allow before error
     paramsDefault['DeltaDesired'] = 1.0e-5                    # desired Delta value to use in numerical solvers such as Runge-Kutta-Fehlberg and Rosenbrock
     paramsDefault['AgeMinDefault'] = 1.0                      # Myr - default age to start evolutionary tracks
     paramsDefault['AgeMaxDefault'] = 5000.0                   # Myr - default age to end evolutionary tracks
     paramsDefault['deltaJac'] = 1.0e-5                        # pertubation to use when calculating Jacobian for Rosenbrock solver
     paramsDefault['CoefficientsRB'] = _CoefficientsRB()       # all coefficents needed for the Rosenbrock solver
     paramsDefault['dAgeMin'] = 1.0e-5                         # Myr - minimum timestep to allow
-    paramsDefault['dAgeMax'] = 5000.0                         # Myr - maximum timestep to allow
+    paramsDefault['dAgeMax'] = 50.0                           # Myr - maximum timestep to allow
 
     # Physical processes to include
     paramsDefault['CoreEnvelopeDecoupling'] = True            # should core envelope decoupling be included

src/mors/rotevo.py

@@ -420,6 +420,12 @@ def EvolveRotationStep(Mstar=None,Age=None,OmegaEnv=None,OmegaCore=None,dAgeMax=
                                                                         OmegaCore=OmegaCore , dAge=dAge , dAgeMax=dAgeMax ,
                                                                         params=params , StarEvo=StarEvo )
 
+    elif ( params['TimeIntegrationMethod'] == 'RosenbrockFixed' ):
+
+        dAge , dAgeNew , OmegaEnv , OmegaCore =  _EvolveRotationStepRBFixed( Mstar=Mstar , Age=Age , OmegaEnv=OmegaEnv ,
+                                                                             OmegaCore=OmegaCore , dAge=dAge , dAgeMax=dAgeMax ,
+                                                                             params=params , StarEvo=StarEvo )
+
     else:
         raise Exception("invalid value of TimeIntegrationMethod in parameters")
 
@@ -528,7 +534,7 @@ def _EvolveRotationStepRKF(Mstar=None,Age=None,OmegaEnv=None,OmegaCore=None,dAge
     return dAge , dAgeNew , OmegaEnv , OmegaCore
 
 def _EvolveRotationStepRB(Mstar=None,Age=None,OmegaEnv=None,OmegaCore=None,dAgeMax=None,dAge=None,params=params.paramsDefault,StarEvo=None):
-    """Takes basic stellar parameters, evolves by timestep using Runge-Kutta-Fehlberg method."""
+    """Takes basic stellar parameters, evolves by timestep using Rosenbrock method."""
 
     # Get coefficients for solver
     CoefficientsRB = params['CoefficientsRB']
@@ -587,6 +593,39 @@ def _EvolveRotationStepRB(Mstar=None,Age=None,OmegaEnv=None,OmegaCore=None,dAgeM
 
     return dAge , dAgeNew , OmegaEnv , OmegaCore
 
+def _EvolveRotationStepRBFixed(Mstar=None,Age=None,OmegaEnv=None,OmegaCore=None,dAgeMax=None,dAge=None,params=params.paramsDefault,StarEvo=None):
+    """Takes basic stellar parameters, evolves by timestep using Rosenbrock method with fixed time grid."""
+
+    # Get coefficients for solver
+    CoefficientsRB = params['CoefficientsRB']
+
+    # Get time step
+    dAge = 0.1 * np.power(Age, 0.75)
+    dAge = min( dAge , dAgeMax )
+    dAge = min( dAge , params['dAgeMax'] )
+    dAgeNew = dAge
+
+    # Setup array to hold integration values
+    X = np.array([OmegaEnv,OmegaCore])
+    nVar = len(X)
+
+    # Get Jacobian and assume it is a constant over timestep
+    Jac = _JacobianRB( Mstar , Age+dAge , X , nVar , params , StarEvo )
+
+    # Calculate the k coefficients
+    kCoeff = _kCoeffRB( Mstar , Age , dAge , X , Jac , nVar , CoefficientsRB , params , StarEvo )
+
+    # Get updated values
+    Xnew = copy.deepcopy(X)
+    for i in range(0,CoefficientsRB['s']):
+        Xnew += CoefficientsRB['b'][i] * kCoeff[:,i]
+
+    # Save new estimate
+    OmegaEnv = Xnew[0]
+    OmegaCore = Xnew[1]
+
+    return dAge , dAgeNew , OmegaEnv , OmegaCore
+
 def _JacobianRB(Mstar,Age,X,nVar,params,StarEvo):
     """Calculates Jacobian d(dOmega/dt)/dOmega terms for Rosenbrock solver."""
 

tests/test_spada.py → tests/test_spada_FE.py

@@ -10,14 +10,16 @@
 )
 
 @pytest.mark.parametrize("inp,expected", TEST_DATA)
-def test_spada(inp, expected):
+def test_spada_FE(inp, expected):
 
     mors.DownloadEvolutionTracks('Spada')
-    star = mors.Star(Mstar=inp[0], Omega=inp[1])
+    params = mors.NewParams()
+    params['TimeIntegrationMethod'] = 'ForwardEuler'
+    star = mors.Star(Mstar=inp[0], Omega=inp[1], params = params)
     ret = (
          star.Value(inp[2], 'Rstar'),
          star.Value(inp[2], 'Lbol'),
          star.Value(inp[2], 'Leuv'),
          )
 
-    assert_allclose(ret, expected, rtol=1e-5, atol=0)
+    assert_allclose(ret, expected, rtol=1e-6, atol=0)

tests/test_spada_RB.py

@@ -0,0 +1,23 @@
+import mors
+import pytest
+from numpy.testing import assert_allclose
+
+TEST_DATA = (
+    ((0.128, 45.2, 8.5e1),(0.21264087, 1.68629162e+31, 1.74475018e+28)),
+    ((1.113, 17.7, 3.2e3),(1.16581971, 6.81767904e+33, 7.94621466e+28)),
+    ((0.995, 1.005, 1.0e4),(1.48910765e+00, 7.80659378e+33, 3.16581356e+28)),
+    ((1.000, 1.00, 1.0e4),(1.52588718e+00, 8.13197065e+33, 3.29155751e+28)),
+)
+
+@pytest.mark.parametrize("inp,expected", TEST_DATA)
+def test_spada_RB(inp, expected):
+
+    mors.DownloadEvolutionTracks('Spada')
+    star = mors.Star(Mstar=inp[0], Omega=inp[1])
+    ret = (
+         star.Value(inp[2], 'Rstar'),
+         star.Value(inp[2], 'Lbol'),
+         star.Value(inp[2], 'Leuv'),
+         )
+
+    assert_allclose(ret, expected, rtol=1e-6, atol=0)