From c26edc6edefabf37bb169b75ce70fdd9f5f29813 Mon Sep 17 00:00:00 2001
From: Alexandre Magueresse <alexandre.magueresse@monash.edu>
Date: Wed, 20 Dec 2023 16:12:14 +1100
Subject: [PATCH] Fixed constructors of TransientFEOperator and only assemble
 jacobian matrices if they have contributions

---
 src/ODEs/ODEOpsFromTFEOps.jl     | 36 +++++++++++++++++++-------------
 src/ODEs/TransientFEOperators.jl | 14 ++++++-------
 2 files changed, 28 insertions(+), 22 deletions(-)

diff --git a/src/ODEs/ODEOpsFromTFEOps.jl b/src/ODEs/ODEOpsFromTFEOps.jl
index f9d34eccb..5f1182b21 100644
--- a/src/ODEs/ODEOpsFromTFEOps.jl
+++ b/src/ODEs/ODEOpsFromTFEOps.jl
@@ -107,7 +107,8 @@ function allocate_odeopcache(
       const_form = nothing
       if is_form_constant(odeop, k)
         jac = jacs[k+1]
-        matdata = collect_cell_matrix(Ut, V, jac(t, uh, du, v))
+        dc = jac(t, uh, du, v)
+        matdata = collect_cell_matrix(Ut, V, dc)
         const_form = assemble_matrix(assembler, matdata)
       end
       const_forms = (const_forms..., const_form)
@@ -117,7 +118,8 @@ function allocate_odeopcache(
     k = order
     if is_form_constant(odeop, k)
       jac = jacs[k+1]
-      matdata = collect_cell_matrix(Ut, V, jac(t, uh, du, v))
+      dc = jac(t, uh, du, v)
+      matdata = collect_cell_matrix(Ut, V, dc)
       const_form = assemble_matrix(assembler, matdata)
     end
     const_forms = (const_forms..., const_form)
@@ -167,7 +169,8 @@ function Algebra.residual!(
   !add && fill!(r, zero(eltype(r)))
 
   res = get_res(odeop.tfeop)
-  vecdata = collect_cell_vector(V, res(t, uh, v))
+  dc = res(t, uh, v)
+  vecdata = collect_cell_vector(V, dc)
   assemble_vector_add!(r, assembler, vecdata)
 
   r
@@ -184,11 +187,10 @@ function Algebra.residual!(
   assembler = get_assembler(odeop.tfeop)
 
   !add && fill!(r, zero(eltype(r)))
-  dc = DomainContribution()
 
   # Residual
   res = get_res(odeop.tfeop)
-  dc = dc + res(t, uh, v)
+  dc = res(t, uh, v)
 
   # Mass
   order = get_order(odeop)
@@ -213,11 +215,10 @@ function Algebra.residual!(
   assembler = get_assembler(odeop.tfeop)
 
   !add && fill!(r, zero(eltype(r)))
-  dc = DomainContribution()
 
   # Residual
   res = get_res(odeop.tfeop)
-  dc = dc + res(t, uh, v)
+  dc = res(t, uh, v)
 
   # Mass
   order = get_order(odeop)
@@ -242,11 +243,10 @@ function Algebra.residual!(
   assembler = get_assembler(odeop.tfeop)
 
   !add && fill!(r, zero(eltype(r)))
-  dc = DomainContribution()
 
   # Residual
   res = get_res(odeop.tfeop)
-  dc = dc + res(t, uh, v)
+  dc = res(t, uh, v)
 
   # Forms
   order = get_order(odeop)
@@ -309,8 +309,10 @@ function jacobian_add!(
     dc = dc + w * jac(t, uh, du, v)
   end
 
-  matdata = collect_cell_matrix(Ut, V, dc)
-  assemble_matrix_add!(J, assembler, matdata)
+  if num_domains(dc) > 0
+    matdata = collect_cell_matrix(Ut, V, dc)
+    assemble_matrix_add!(J, assembler, matdata)
+  end
 
   J
 end
@@ -349,8 +351,10 @@ function jacobian_add!(
     end
   end
 
-  matdata = collect_cell_matrix(Ut, V, dc)
-  assemble_matrix_add!(J, assembler, matdata)
+  if num_domains(dc) > 0
+    matdata = collect_cell_matrix(Ut, V, dc)
+    assemble_matrix_add!(J, assembler, matdata)
+  end
 
   J
 end
@@ -380,8 +384,10 @@ function jacobian_add!(
     end
   end
 
-  matdata = collect_cell_matrix(Ut, V, dc)
-  assemble_matrix_add!(J, assembler, matdata)
+  if num_domains(dc) > 0
+    matdata = collect_cell_matrix(Ut, V, dc)
+    assemble_matrix_add!(J, assembler, matdata)
+  end
 
   J
 end
diff --git a/src/ODEs/TransientFEOperators.jl b/src/ODEs/TransientFEOperators.jl
index 43371c80c..3244a42d1 100644
--- a/src/ODEs/TransientFEOperators.jl
+++ b/src/ODEs/TransientFEOperators.jl
@@ -236,7 +236,7 @@ function TransientFEOperator(
   jac::Function,
   trial, test
 )
-  TransientFEOpFromWeakForm(
+  TransientFEOperator(
     res, (jac,),
     trial, test
   )
@@ -247,7 +247,7 @@ function TransientFEOperator(
   jac::Function, jac_t::Function,
   trial, test
 )
-  TransientFEOpFromWeakForm(
+  TransientFEOperator(
     res, (jac, jac_t),
     trial, test
   )
@@ -258,7 +258,7 @@ function TransientFEOperator(
   jac::Function, jac_t::Function, jac_tt::Function,
   trial, test
 )
-  TransientFEOpFromWeakForm(
+  TransientFEOperator(
     res, (jac, jac_t, jac_tt),
     trial, test
   )
@@ -368,7 +368,7 @@ function TransientQuasilinearFEOperator(
   jac::Function, jac_t::Function,
   trial, test
 )
-  TransientQuasilinearFEOpFromWeakForm(
+  TransientQuasilinearFEOperator(
     mass, res, (jac, jac_t),
     trial, test
   )
@@ -379,7 +379,7 @@ function TransientQuasilinearFEOperator(
   jac::Function, jac_t::Function, jac_tt::Function,
   trial, test
 )
-  TransientQuasilinearFEOpFromWeakForm(
+  TransientQuasilinearFEOperator(
     mass, res, (jac, jac_t, jac_tt),
     trial, test
   )
@@ -514,7 +514,7 @@ function TransientSemilinearFEOperator(
   trial, test;
   constant_mass::Bool=false,
 )
-  TransientSemilinearFEOpFromWeakForm(
+  TransientSemilinearFEOperator(
     mass, res, (jac, jac_t),
     trial, test;
     constant_mass
@@ -527,7 +527,7 @@ function TransientSemilinearFEOperator(
   trial, test;
   constant_mass::Bool=false,
 )
-  TransientSemilinearFEOpFromWeakForm(
+  TransientSemilinearFEOperator(
     mass, res, (jac, jac_t, jac_tt),
     trial, test;
     constant_mass