From bb36d88c19cff5004ca605aa58f5965fb5418030 Mon Sep 17 00:00:00 2001
From: Aaron LI <aly@aaronly.me>
Date: Fri, 6 Oct 2017 13:25:00 +0800
Subject: clusters/solver.py: Adjust parameter tc,uc orders

---
 fg21sim/extragalactic/clusters/solver.py | 17 +++++++----------
 1 file changed, 7 insertions(+), 10 deletions(-)

(limited to 'fg21sim/extragalactic/clusters/solver.py')

diff --git a/fg21sim/extragalactic/clusters/solver.py b/fg21sim/extragalactic/clusters/solver.py
index 0fb220d..9117bfb 100644
--- a/fg21sim/extragalactic/clusters/solver.py
+++ b/fg21sim/extragalactic/clusters/solver.py
@@ -271,7 +271,7 @@ class FokkerPlanckSolver:
         Wminus = W * np.exp(-ww/2)
         return Wminus
 
-    def tridiagonal_coefs(self, tc, uc, tstep=None):
+    def tridiagonal_coefs(self, uc, tc, tstep):
         """
         Calculate the coefficients for the tridiagonal system of linear
         equations corresponding to the original Fokker-Planck equation.
@@ -287,10 +287,7 @@ class FokkerPlanckSolver:
 
         References: Ref.[park1996],Eqs.(16,18,34)
         """
-        if tstep is None:
-            dt = self.tstep
-        else:
-            dt = tstep
+        dt = tstep
         x = self.x
         dx = self.dx
         dx_phalf = self.dx_phalf
@@ -367,13 +364,13 @@ class FokkerPlanckSolver:
         """
         pass
 
-    def solve_step(self, tc, uc, tstep=None):
+    def solve_step(self, uc, tc, tstep=None):
         """
         Solve the Fokker-Planck equation by a single step.
         """
         if tstep is None:
             tstep = self.tstep
-        a, b, c, r = self.tridiagonal_coefs(tc=tc, uc=uc, tstep=tstep)
+        a, b, c, r = self.tridiagonal_coefs(uc=uc, tc=tc, tstep=tstep)
         TDM_a = -a[1:]  # Also drop the first element
         TDM_b = b
         TDM_c = -c[:-1]  # Also drop the last element
@@ -381,7 +378,7 @@ class FokkerPlanckSolver:
         t2 = tc + tstep
         u2 = TDMAsolver(TDM_a, TDM_b, TDM_c, TDM_rhs)
         u2 = self.fix_boundary(u2)
-        return (t2, u2)
+        return (u2, t2)
 
     def solve(self, u0, tstart, tstop):
         """
@@ -399,9 +396,9 @@ class FokkerPlanckSolver:
         while tc+tstep < tstop:
             i += 1
             logger.debug("[%d/%d] t=%.3f ..." % (i, nstep, tc))
-            tc, uc = self.solve_step(tc, uc, tstep=tstep)
+            uc, tc = self.solve_step(uc=uc, tc=tc, tstep=tstep)
         # Last step
         tstep = tstop - tc
         logger.debug("Last step: t=%.3f (tstep=%.3f) ..." % (tc, tstep))
-        tc, uc = self.solve_step(tc, uc, tstep=tstep)
+        uc, __ = self.solve_step(uc=uc, tc=tc, tstep=tstep)
         return uc
-- 
cgit v1.2.2