From 3199ac7871c082d2ba70b5f2b08b5314f0ceae97 Mon Sep 17 00:00:00 2001
From: Hendrik Ranocha <mail@ranocha.de>
Date: Wed, 9 Nov 2022 11:44:20 +0100
Subject: [PATCH] fix parabolic terms

---
 src/Trixi.jl                                  |  2 +-
 src/semidiscretization/semidiscretization.jl  | 26 +++++++++++++------
 ...semidiscretization_hyperbolic_parabolic.jl |  2 +-
 3 files changed, 20 insertions(+), 10 deletions(-)

diff --git a/src/Trixi.jl b/src/Trixi.jl
index 2ec92d67532..e09d1f25eb8 100644
--- a/src/Trixi.jl
+++ b/src/Trixi.jl
@@ -28,7 +28,7 @@ using Reexport: @reexport
 using SciMLBase: SciMLBase, unwrapped_f,
                  CallbackSet, DiscreteCallback,
                  ODEProblem, ODESolution, ODEFunction,
-                 SplitODEProblem
+                 SplitODEProblem, SplitFunction
 import SciMLBase: get_du, get_tmp_cache, u_modified!,
                   AbstractODEIntegrator, init, step!, check_error,
                   get_proposed_dt, set_proposed_dt!,
diff --git a/src/semidiscretization/semidiscretization.jl b/src/semidiscretization/semidiscretization.jl
index f38c929306c..4bcc5878ef8 100644
--- a/src/semidiscretization/semidiscretization.jl
+++ b/src/semidiscretization/semidiscretization.jl
@@ -94,12 +94,14 @@ end
 
 # get the semidiscretization from an `ODEIntegrator`
 function extract_semidiscretization(integrator)
-  f = unwrapped_f(integrator.f.f)
-  if f isa RHSWrapper
-    return f.semi
-  else
-    return integrator.p
-  end
+  # TODO: check alternatives such as
+  # f = unwrapped_f(integrator.f.f)
+  # if f isa RHSWrapper
+  #   return f.semi
+  # else
+  #   return integrator.p
+  # end
+  extract_semidiscretization(integrator.sol)
 end
 
 function extract_semidiscretization(sol::ODESolution)
@@ -107,14 +109,22 @@ function extract_semidiscretization(sol::ODESolution)
 end
 
 function extract_semidiscretization(ode::ODEProblem)
-  f = unwrapped_f(ode.f.f)
+  _extract_semidiscretization(ode.f, ode.p)
+end
+
+function _extract_semidiscretization(f, p)
+  f = unwrapped_f(f.f)
   if f isa RHSWrapper
     return f.semi
   else
-    return ode.p
+    return p
   end
 end
 
+function _extract_semidiscretization(f::SplitFunction, p)
+  _extract_semidiscretization(f.f1, p)
+end
+
 
 """
     semidiscretize(semi::AbstractSemidiscretization, tspan, restart_file::AbstractString)
diff --git a/src/semidiscretization/semidiscretization_hyperbolic_parabolic.jl b/src/semidiscretization/semidiscretization_hyperbolic_parabolic.jl
index 29fbad30592..07a3ed7a33c 100644
--- a/src/semidiscretization/semidiscretization_hyperbolic_parabolic.jl
+++ b/src/semidiscretization/semidiscretization_hyperbolic_parabolic.jl
@@ -225,7 +225,7 @@ function semidiscretize(semi::SemidiscretizationHyperbolicParabolic, tspan)
   #       See https://github.com/trixi-framework/Trixi.jl/issues/328
   iip = true # is-inplace, i.e., we modify a vector when calling rhs!, rhs_parabolic!
   wrapper_hyperbolic = RHSWrapper(semi, rhs!) # wrap `rhs!`` and `semi` to reduce latency, see `RHSWrapper`
-  wrapper_parabolic = RHSWrapper(rhs_parabolic!, rhs!)
+  wrapper_parabolic = RHSWrapper(semi, rhs_parabolic!)
   return SplitODEProblem{iip}(wrapper_hyperbolic, wrapper_parabolic, u0_ode, tspan, semi)
 end