From 99659a588752c501da68b4f5415f749f20718803 Mon Sep 17 00:00:00 2001
From: Scott MacLachlan <smaclachlan@mun.ca>
Date: Mon, 18 Nov 2024 12:37:02 -0330
Subject: [PATCH] Clean up constants for BC

---
 irksome/dirk_stepper.py     | 14 ++++++++------
 irksome/explicit_stepper.py |  5 ++++-
 2 files changed, 12 insertions(+), 7 deletions(-)

diff --git a/irksome/dirk_stepper.py b/irksome/dirk_stepper.py
index c32c419..be97597 100644
--- a/irksome/dirk_stepper.py
+++ b/irksome/dirk_stepper.py
@@ -71,7 +71,7 @@ def getFormDIRK(F, ks, butch, t, dt, u0, bcs=None):
         new_bc = bc.reconstruct(g=gdat)
         bcnew.append(new_bc)
 
-    return stage_F, (k, g, a, c, a_vals, d_val), bcnew
+    return stage_F, (k, g, a, c), bcnew, (a_vals, d_val)
 
 
 class DIRKTimeStepper:
@@ -105,7 +105,7 @@ def __init__(self, F, butcher_tableau, t, dt, u0, bcs=None,
         # that we update as we go.  We need to remember the
         # stage values we've computed earlier in the time step...
 
-        stage_F, (k, g, a, c, a_vals, d_val), bcnew = getFormDIRK(
+        stage_F, (k, g, a, c), bcnew, (a_vals, d_val) = getFormDIRK(
             F, self.ks, butcher_tableau, t, dt, u0, bcs=bcs)
 
         self.bcnew = bcnew
@@ -128,11 +128,13 @@ def __init__(self, F, butcher_tableau, t, dt, u0, bcs=None,
             self.problem, appctx=appctx, solver_parameters=solver_parameters,
             nullspace=nullspace)
 
-        self.kgac = k, g, a, c, a_vals, d_val
+        self.kgac = k, g, a, c
+        self.bc_constants = a_vals, d_val
 
-    def update_bc_constants(self, i, a_vals, d_val, c):
+    def update_bc_constants(self, i, c):
         AAb = self.AAb
         CCone = self.CCone
+        a_vals, d_val = self.bc_constants
         ns = AAb.shape[1]
         for j in range(i):
             a_vals[j].assign(AAb[i, j])
@@ -142,7 +144,7 @@ def update_bc_constants(self, i, a_vals, d_val, c):
         c.assign(CCone[i])
 
     def advance(self):
-        k, g, a, c, a_vals, d_val = self.kgac
+        k, g, a, c = self.kgac
         ks = self.ks
         u0 = self.u0
         dtc = float(self.dt)
@@ -160,7 +162,7 @@ def advance(self):
                     gbit += dtc * float(AA[i, j]) * kbit
 
             # update BC constants for the variational problem
-            self.update_bc_constants(i, a_vals, d_val, c)
+            self.update_bc_constants(i, c)
             a.assign(AA[i, i])
 
             # solve new variational problem, stash the computed
diff --git a/irksome/explicit_stepper.py b/irksome/explicit_stepper.py
index 67c55ea..734f559 100644
--- a/irksome/explicit_stepper.py
+++ b/irksome/explicit_stepper.py
@@ -25,7 +25,10 @@ def __init__(self, F, butcher_tableau, t, dt, u0, bcs=None,
     # we should impose the BCs so that they are satisfied for the next
     # stage for all but that last stage, and that they are satisfied
     # for the next timestep for the last stage.
-    def update_bc_constants(self, AAb, CCone, i, a_vals, d_val, c):
+    def update_bc_constants(self, i, c):
+        AAb = self.AAb
+        CCone = self.CCone
+        a_vals, d_val = self.bc_constants
         ns = AAb.shape[1]
         for j in range(i):
             a_vals[j].assign(AAb[i+1, j])