Move functional logic out of stepping function.

RadostW · Dec 4, 2024 · eb97474 · eb97474
1 parent bea391e
commit eb97474
Showing 1 changed file with 13 additions and 13 deletions.
diff --git a/pychastic/sde_solver.py b/pychastic/sde_solver.py
@@ -170,6 +170,19 @@ def L(f, idx):
                     raise ValueError
             return f
 
+        id_ = lambda x: x
+
+        f_w = L(id_, "w")
+
+        f_t = L(id_, "t")
+        f_ww = L(id_, "ww")
+
+        f_tw = L(id_, "tw")
+        f_wt = L(id_, "wt")
+        f_www = L(id_, "www")
+
+        f_tt = L(id_, "tt")
+
         def step(
             x,
             d_t,
@@ -181,19 +194,6 @@ def step(
             scheme="euler",
         ):
 
-            id_ = lambda x: x
-
-            f_w = L(id_, "w")
-
-            f_t = L(id_, "t")
-            f_ww = L(id_, "ww")
-
-            f_tw = L(id_, "tw")
-            f_wt = L(id_, "wt")
-            f_www = L(id_, "www")
-
-            f_tt = L(id_, "tt")
-
             new_x = x
             new_x += (f_t(x)*d_t).squeeze() + contract_all(f_w(x), d_w)