Feature/stageprocessing

doc/arkode/guide/source/Constants.rst

@@ -470,9 +470,9 @@ contains the ARKODE output constants.
    Commented-out table rows:
 
       +-------------------------------------+------+------------------------------------------------------------+
-      | :index:`ARK_POSTPROCESS_STEP_FAIL`  | -37  | An error occurred when calling the user-provided           |
+      | :index:`ARK_POSTPROCESS_STEP_FAIL`  | -37  | An error occurred when calling a user-provided.            |
       |                                     |      | step-based :c:func:`ARKPostProcessFn` routine.             |
       +-------------------------------------+------+------------------------------------------------------------+
-      | :index:`ARK_POSTPROCESS_STAGE_FAIL` | -38  | An error occurred when calling the user-provided           |
+      | :index:`ARK_POSTPROCESS_STAGE_FAIL` | -38  | An error occurred when calling a user-provided             |
       |                                     |      | stage-based :c:func:`ARKPostProcessFn` routine.            |
       +-------------------------------------+------+------------------------------------------------------------+

include/arkode/arkode.h

@@ -124,7 +124,8 @@ extern "C" {
 #define ARK_INNERTOOUTER_FAIL    -36
 
 /* ARK_POSTPROCESS_FAIL equals ARK_POSTPROCESS_STEP_FAIL
-   for backwards compatibility */
+   for backwards compatibility. Note that we use these
+   same constants for step and stage preprocessing errors */
 #define ARK_POSTPROCESS_FAIL       -37
 #define ARK_POSTPROCESS_STEP_FAIL  -37
 #define ARK_POSTPROCESS_STAGE_FAIL -38
@@ -276,8 +277,14 @@ SUNDIALS_EXPORT int ARKodeClearStopTime(void* arkode_mem);
 SUNDIALS_EXPORT int ARKodeSetFixedStep(void* arkode_mem, sunrealtype hfixed);
 SUNDIALS_EXPORT int ARKodeSetStepDirection(void* arkode_mem, sunrealtype stepdir);
 SUNDIALS_EXPORT int ARKodeSetUserData(void* arkode_mem, void* user_data);
+SUNDIALS_EXPORT int ARKodeSetPreprocessStepFn(void* arkode_mem,
+                                              ARKPostProcessFn ProcessStep);
 SUNDIALS_EXPORT int ARKodeSetPostprocessStepFn(void* arkode_mem,
                                                ARKPostProcessFn ProcessStep);
+SUNDIALS_EXPORT int ARKodeSetPostprocessStepFailFn(void* arkode_mem,
+                                                   ARKPostProcessFn ProcessStep);
+SUNDIALS_EXPORT int ARKodeSetPreprocessRHSFn(void* arkode_mem,
+                                             ARKPostProcessFn ProcessRHS);
 SUNDIALS_EXPORT int ARKodeSetPostprocessStageFn(void* arkode_mem,
                                                 ARKPostProcessFn ProcessStage);
 

src/arkode/arkode.c

@@ -904,6 +904,18 @@ int ARKodeEvolve(void* arkode_mem, sunrealtype tout, N_Vector yout,
                  "step = %li, tn = " SUN_FORMAT_G ", h = " SUN_FORMAT_G,
                  ark_mem->nst + 1, ark_mem->tn, ark_mem->h);
 
+      /* fill (tcur,ycur) with the last accepted step solution */
+      ark_mem->tcur = ark_mem->tn;
+      N_VScale(ONE, ark_mem->yn, ark_mem->ycur);
+
+      /* call the user-supplied step preprocessing function (if it exists) */
+      if (ark_mem->PreProcessStep != NULL)
+      {
+        retval = ark_mem->PreProcessStep(ark_mem->tcur, ark_mem->ycur,
+                                         ark_mem->ps_data);
+        if (retval != 0) { return (ARK_POSTPROCESS_STEP_FAIL); }
+      }
+
       /* Call time stepper module to attempt a step:
             0 => step completed successfully
            >0 => step encountered recoverable failure; reduce step if possible
@@ -1001,6 +1013,14 @@ int ARKodeEvolve(void* arkode_mem, sunrealtype tout, N_Vector yout,
       ark_mem->h *= ark_mem->eta;
       ark_mem->next_h = ark_mem->hprime = ark_mem->h;
 
+      /* since the previous step attempt failed, call the user-supplied step failure postprocessing function (if it exists) */
+      if (ark_mem->PostProcessStepFail != NULL)
+      {
+        retval = ark_mem->PostProcessStepFail(ark_mem->tcur, ark_mem->ycur,
+                                              ark_mem->ps_data);
+        if (retval != 0) { return (ARK_POSTPROCESS_STEP_FAIL); }
+      }
+
     } /* end looping for step attempts */
 
     /* If step attempt loop succeeded, complete step (update current time, solution,
@@ -1587,12 +1607,15 @@ ARKodeMem arkCreate(SUNContext sunctx)
   ark_mem->VRabstolMallocDone = SUNFALSE;
   ark_mem->MallocDone         = SUNFALSE;
 
-  /* No user-supplied step postprocessing function yet */
-  ark_mem->ProcessStep = NULL;
-  ark_mem->ps_data     = NULL;
+  /* No user-supplied step pre- or post-processing functions yet */
+  ark_mem->PreProcessStep      = NULL;
+  ark_mem->PostProcessStep     = NULL;
+  ark_mem->PostProcessStepFail = NULL;
+  ark_mem->ps_data             = NULL;
 
-  /* No user-supplied stage postprocessing function yet */
-  ark_mem->ProcessStage = NULL;
+  /* No user-supplied stage pre- or post-processing functions yet */
+  ark_mem->PreProcessRHS    = NULL;
+  ark_mem->PostProcessStage = NULL;
 
   /* No user_data pointer yet */
   ark_mem->user_data = NULL;
@@ -2727,9 +2750,10 @@ int arkCompleteStep(ARKodeMem ark_mem, sunrealtype dsm)
   }
 
   /* apply user-supplied step postprocessing function (if supplied) */
-  if (ark_mem->ProcessStep != NULL)
+  if (ark_mem->PostProcessStep != NULL)
   {
-    retval = ark_mem->ProcessStep(ark_mem->tcur, ark_mem->ycur, ark_mem->ps_data);
+    retval = ark_mem->PostProcessStep(ark_mem->tcur, ark_mem->ycur,
+                                      ark_mem->ps_data);
     if (retval != 0) { return (ARK_POSTPROCESS_STEP_FAIL); }
   }
 

src/arkode/arkode_arkstep.c

@@ -1330,22 +1330,29 @@ int arkStep_FullRHS(ARKodeMem ark_mem, sunrealtype t, N_Vector y, N_Vector f,
     /* compute the full RHS */
     if (!(ark_mem->fn_is_current))
     {
-      /* compute the explicit component */
-      if (step_mem->explicit)
+      /* apply user-supplied stage preprocessing function (if supplied) */
+      if (ark_mem->PreProcessRHS != NULL)
       {
-        retval = step_mem->fe(t, y, step_mem->Fe[0], ark_mem->user_data);
-        step_mem->nfe++;
+        retval = ark_mem->PreProcessRHS(t, y, ark_mem->user_data);
+        if (retval != 0) { return (ARK_POSTPROCESS_STAGE_FAIL); }
+      }
+
+      /* compute the implicit component */
+      if (step_mem->implicit)
+      {
+        retval = step_mem->fi(t, y, step_mem->Fi[0], ark_mem->user_data);
+        step_mem->nfi++;
         if (retval != 0)
         {
           arkProcessError(ark_mem, ARK_RHSFUNC_FAIL, __LINE__, __func__,
                           __FILE__, MSG_ARK_RHSFUNC_FAILED, t);
           return (ARK_RHSFUNC_FAIL);
         }
 
-        /* compute and store M(t)^{-1} fe */
+        /* compute and store M(t)^{-1} fi */
         if (step_mem->mass_type == MASS_TIMEDEP)
         {
-          retval = step_mem->msolve((void*)ark_mem, step_mem->Fe[0],
+          retval = step_mem->msolve((void*)ark_mem, step_mem->Fi[0],
                                     step_mem->nlscoef / ark_mem->h);
           if (retval)
           {
@@ -1356,22 +1363,22 @@ int arkStep_FullRHS(ARKodeMem ark_mem, sunrealtype t, N_Vector y, N_Vector f,
         }
       }
 
-      /* compute the implicit component */
-      if (step_mem->implicit)
+      /* compute the explicit component */
+      if (step_mem->explicit)
       {
-        retval = step_mem->fi(t, y, step_mem->Fi[0], ark_mem->user_data);
-        step_mem->nfi++;
+        retval = step_mem->fe(t, y, step_mem->Fe[0], ark_mem->user_data);
+        step_mem->nfe++;
         if (retval != 0)
         {
           arkProcessError(ark_mem, ARK_RHSFUNC_FAIL, __LINE__, __func__,
                           __FILE__, MSG_ARK_RHSFUNC_FAILED, t);
           return (ARK_RHSFUNC_FAIL);
         }
 
-        /* compute and store M(t)^{-1} fi */
+        /* compute and store M(t)^{-1} fe */
         if (step_mem->mass_type == MASS_TIMEDEP)
         {
-          retval = step_mem->msolve((void*)ark_mem, step_mem->Fi[0],
+          retval = step_mem->msolve((void*)ark_mem, step_mem->Fe[0],
                                     step_mem->nlscoef / ark_mem->h);
           if (retval)
           {
@@ -1455,11 +1462,18 @@ int arkStep_FullRHS(ARKodeMem ark_mem, sunrealtype t, N_Vector y, N_Vector f,
       /* recompute RHS functions */
       if (recomputeRHS)
       {
-        /* compute the explicit component */
-        if (step_mem->explicit)
+        /* apply user-supplied stage preprocessing function (if supplied) */
+        if (ark_mem->PreProcessRHS != NULL)
         {
-          retval = step_mem->fe(t, y, step_mem->Fe[0], ark_mem->user_data);
-          step_mem->nfe++;
+          retval = ark_mem->PreProcessRHS(t, y, ark_mem->user_data);
+          if (retval != 0) { return (ARK_POSTPROCESS_STAGE_FAIL); }
+        }
+
+        /* compute the implicit component */
+        if (step_mem->implicit)
+        {
+          retval = step_mem->fi(t, y, step_mem->Fi[0], ark_mem->user_data);
+          step_mem->nfi++;
           if (retval != 0)
           {
             arkProcessError(ark_mem, ARK_RHSFUNC_FAIL, __LINE__, __func__,
@@ -1470,7 +1484,7 @@ int arkStep_FullRHS(ARKodeMem ark_mem, sunrealtype t, N_Vector y, N_Vector f,
           /* compute and store M(t)^{-1} fi */
           if (step_mem->mass_type == MASS_TIMEDEP)
           {
-            retval = step_mem->msolve((void*)ark_mem, step_mem->Fe[0],
+            retval = step_mem->msolve((void*)ark_mem, step_mem->Fi[0],
                                       step_mem->nlscoef / ark_mem->h);
             if (retval)
             {
@@ -1481,11 +1495,11 @@ int arkStep_FullRHS(ARKodeMem ark_mem, sunrealtype t, N_Vector y, N_Vector f,
           }
         }
 
-        /* compute the implicit component */
-        if (step_mem->implicit)
+        /* compute the explicit component */
+        if (step_mem->explicit)
         {
-          retval = step_mem->fi(t, y, step_mem->Fi[0], ark_mem->user_data);
-          step_mem->nfi++;
+          retval = step_mem->fe(t, y, step_mem->Fe[0], ark_mem->user_data);
+          step_mem->nfe++;
           if (retval != 0)
           {
             arkProcessError(ark_mem, ARK_RHSFUNC_FAIL, __LINE__, __func__,
@@ -1496,7 +1510,7 @@ int arkStep_FullRHS(ARKodeMem ark_mem, sunrealtype t, N_Vector y, N_Vector f,
           /* compute and store M(t)^{-1} fi */
           if (step_mem->mass_type == MASS_TIMEDEP)
           {
-            retval = step_mem->msolve((void*)ark_mem, step_mem->Fi[0],
+            retval = step_mem->msolve((void*)ark_mem, step_mem->Fe[0],
                                       step_mem->nlscoef / ark_mem->h);
             if (retval)
             {
@@ -1564,11 +1578,18 @@ int arkStep_FullRHS(ARKodeMem ark_mem, sunrealtype t, N_Vector y, N_Vector f,
 
   case ARK_FULLRHS_OTHER:
 
-    /* compute the explicit component and store in ark_tempv2 */
-    if (step_mem->explicit)
+    /* apply user-supplied stage preprocessing function (if supplied) */
+    if (ark_mem->PreProcessRHS != NULL)
     {
-      retval = step_mem->fe(t, y, ark_mem->tempv2, ark_mem->user_data);
-      step_mem->nfe++;
+      retval = ark_mem->PreProcessRHS(t, y, ark_mem->user_data);
+      if (retval != 0) { return (ARK_POSTPROCESS_STAGE_FAIL); }
+    }
+
+    /* compute the implicit component and store in sdata */
+    if (step_mem->implicit)
+    {
+      retval = step_mem->fi(t, y, step_mem->sdata, ark_mem->user_data);
+      step_mem->nfi++;
       if (retval != 0)
       {
         arkProcessError(ark_mem, ARK_RHSFUNC_FAIL, __LINE__, __func__, __FILE__,
@@ -1577,11 +1598,11 @@ int arkStep_FullRHS(ARKodeMem ark_mem, sunrealtype t, N_Vector y, N_Vector f,
       }
     }
 
-    /* compute the implicit component and store in sdata */
-    if (step_mem->implicit)
+    /* compute the explicit component and store in ark_tempv2 */
+    if (step_mem->explicit)
     {
-      retval = step_mem->fi(t, y, step_mem->sdata, ark_mem->user_data);
-      step_mem->nfi++;
+      retval = step_mem->fe(t, y, ark_mem->tempv2, ark_mem->user_data);
+      step_mem->nfe++;
       if (retval != 0)
       {
         arkProcessError(ark_mem, ARK_RHSFUNC_FAIL, __LINE__, __func__, __FILE__,
@@ -1845,6 +1866,13 @@ int arkStep_TakeStep_Z(ARKodeMem ark_mem, sunrealtype* dsmPtr, int* nflagPtr)
       }
       else
       {
+        /* apply user-supplied stage preprocessing function (if supplied) */
+        if (ark_mem->PreProcessRHS != NULL)
+        {
+          retval = ark_mem->PreProcessRHS(ark_mem->tn, ark_mem->yn,
+                                          ark_mem->user_data);
+          if (retval != 0) { return (ARK_POSTPROCESS_STAGE_FAIL); }
+        }
         retval = step_mem->fi(ark_mem->tn, ark_mem->yn, step_mem->Fi[0],
                               ark_mem->user_data);
         step_mem->nfi++;
@@ -2015,10 +2043,10 @@ int arkStep_TakeStep_Z(ARKodeMem ark_mem, sunrealtype* dsmPtr, int* nflagPtr)
     /* apply user-supplied stage postprocessing function (if supplied) */
     /* NOTE: with internally inconsistent IMEX methods (c_i^E != c_i^I) the value
        of tcur corresponds to the stage time from the implicit table (c_i^I). */
-    if (ark_mem->ProcessStage != NULL)
+    if (ark_mem->PostProcessStage != NULL)
     {
-      retval = ark_mem->ProcessStage(ark_mem->tcur, ark_mem->ycur,
-                                     ark_mem->user_data);
+      retval = ark_mem->PostProcessStage(ark_mem->tcur, ark_mem->ycur,
+                                         ark_mem->user_data);
       if (retval != 0)
       {
         SUNLogInfo(ARK_LOGGER, "end-stages-list",
@@ -2068,6 +2096,24 @@ int arkStep_TakeStep_Z(ARKodeMem ark_mem, sunrealtype* dsmPtr, int* nflagPtr)
       }
     }
 
+    /* apply user-supplied stage preprocessing function (if supplied) */
+    /* NOTE: with internally inconsistent IMEX methods (c_i^E != c_i^I) the value
+       of tcur corresponds to the stage time from the implicit table (c_i^I). */
+    if (ark_mem->PreProcessRHS != NULL)
+    {
+      if ((step_mem->implicit && !deduce_stage) || (step_mem->explicit))
+      {
+        retval = ark_mem->PreProcessRHS(ark_mem->tcur, ark_mem->ycur,
+                                        ark_mem->user_data);
+        if (retval != 0)
+        {
+          SUNLogInfo(ARK_LOGGER, "end-stages-list",
+                     "status = failed preprocess stage, retval = %i", retval);
+          return (ARK_POSTPROCESS_STAGE_FAIL);
+        }
+      }
+    }
+
     /*    store implicit RHS (value in Fi[is] is from preceding nonlinear iteration) */
     if (step_mem->implicit)
     {