Neaten stokes code

ratt-ru · Nov 23, 2023 · 220b578 · 220b578
1 parent f1442ef
commit 220b578
Showing 1 changed file with 28 additions and 104 deletions.
diff --git a/pfb/utils/stokes.py b/pfb/utils/stokes.py
@@ -351,67 +351,29 @@ def stokes_funcs(data, jones, product, pol):
     if jones.ndim == 6:  # Full mode
         if product == literal('I'):
             i = 0
-            Wsymb = lambdify((gp00, gp01, gp10, gp11,
-                              gq00, gq01, gq10, gq11,
-                              w0, w1, w2, w3),
-                              sm.simplify(W[i,i]))
-            Wjfn = njit(nogil=True, fastmath=True, inline='always')(Wsymb)
-
-
-            Dsymb = lambdify((gp00, gp01, gp10, gp11,
-                              gq00, gq01, gq10, gq11,
-                              w0, w1, w2, w3,
-                              v00, v01, v10, v11),
-                              sm.simplify(C[i]))
-            Djfn = njit(nogil=True, fastmath=True, inline='always')(Dsymb)
-
         elif product == literal('Q'):
             i = 1
-            Wsymb = lambdify((gp00, gp01, gp10, gp11,
-                              gq00, gq01, gq10, gq11,
-                              w0, w1, w2, w3),
-                              sm.simplify(W[i,i]))
-            Wjfn = njit(nogil=True, fastmath=True, inline='always')(Wsymb)
-
-            Dsymb = lambdify((gp00, gp01, gp10, gp11,
-                              gq00, gq01, gq10, gq11,
-                              w0, w1, w2, w3,
-                              v00, v01, v10, v11),
-                              sm.simplify(C[i]))
-            Djfn = njit(nogil=True, fastmath=True, inline='always')(Dsymb)
-
         elif product == literal('U'):
             i = 2
-            Wsymb = lambdify((gp00, gp01, gp10, gp11,
-                              gq00, gq01, gq10, gq11,
-                              w0, w1, w2, w3),
-                              sm.simplify(W[i,i]))
-            Wjfn = njit(nogil=True, fastmath=True, inline='always')(Wsymb)
-            Dsymb = lambdify((gp00, gp01, gp10, gp11,
-                              gq00, gq01, gq10, gq11,
-                              w0, w1, w2, w3,
-                              v00, v01, v10, v11),
-                              sm.simplify(C[i]))
-            Djfn = njit(nogil=True, fastmath=True, inline='always')(Dsymb)
-
         elif product == literal('V'):
             i = 3
-            Wsymb = lambdify((gp00, gp01, gp10, gp11,
-                              gq00, gq01, gq10, gq11,
-                              w0, w1, w2, w3),
-                              sm.simplify(W[i,i]))
-            Wjfn = njit(nogil=True, fastmath=True, inline='always')(Wsymb)
-
-            Dsymb = lambdify((gp00, gp01, gp10, gp11,
-                              gq00, gq01, gq10, gq11,
-                              w0, w1, w2, w3,
-                              v00, v01, v10, v11),
-                              sm.simplify(D[i]))
-            Djfn = njit(nogil=True, fastmath=True, inline='always')(Dsymb)
-
         else:
             raise ValueError(f"Unknown polarisation product {product}")
 
+        Wsymb = lambdify((gp00, gp01, gp10, gp11,
+                            gq00, gq01, gq10, gq11,
+                            w0, w1, w2, w3),
+                            sm.simplify(sm.expand(W[i,i])))
+        Wjfn = njit(nogil=True, fastmath=True, inline='always')(Wsymb)
+
+
+        Dsymb = lambdify((gp00, gp01, gp10, gp11,
+                            gq00, gq01, gq10, gq11,
+                            w0, w1, w2, w3,
+                            v00, v01, v10, v11),
+                            sm.simplify(smexpand(C[i])))
+        Djfn = njit(nogil=True, fastmath=True, inline='always')(Dsymb)
+
         @njit(nogil=True, fastmath=True, inline='always')
         def wfunc(gp, gq, W):
             gp00 = gp[0,0]
@@ -465,67 +427,29 @@ def vfunc(gp, gq, W, V):
 
         if product == literal('I'):
             i = 0
-            Wsymb = lambdify((gp00, gp11,
-                              gq00, gq11,
-                              w0, w1, w2, w3),
-                              sm.simplify(W[i,i]))
-            Wjfn = njit(nogil=True, fastmath=True, inline='always')(Wsymb)
-
-
-            Dsymb = lambdify((gp00, gp11,
-                              gq00, gq11,
-                              w0, w1, w2, w3,
-                              v00, v01, v10, v11),
-                              sm.simplify(C[i]))
-            Djfn = njit(nogil=True, fastmath=True, inline='always')(Dsymb)
-
         elif product == literal('Q'):
             i = 1
-            Wsymb = lambdify((gp00, gp11,
-                              gq00, gq11,
-                              w0, w1, w2, w3),
-                              sm.simplify(W[i,i]))
-            Wjfn = njit(nogil=True, fastmath=True, inline='always')(Wsymb)
-
-            Dsymb = lambdify((gp00, gp11,
-                              gq00, gq11,
-                              w0, w1, w2, w3,
-                              v00, v01, v10, v11),
-                              sm.simplify(C[i]))
-            Djfn = njit(nogil=True, fastmath=True, inline='always')(Dsymb)
-
         elif product == literal('U'):
             i = 2
-            Wsymb = lambdify((gp00, gp11,
-                              gq00, gq11,
-                              w0, w1, w2, w3),
-                             sm.simplify(W[i,i]))
-            Wjfn = njit(nogil=True, fastmath=True, inline='always')(Wsymb)
-
-            Dsymb = lambdify((gp00, gp11,
-                              gq00, gq11,
-                              w0, w1, w2, w3,
-                              v00, v01, v10, v11),
-                              sm.simplify(C[i]))
-            Djfn = njit(nogil=True, fastmath=True, inline='always')(Dsymb)
-
         elif product == literal('V'):
             i = 3
-            Wsymb = lambdify((gp00, gp11,
-                              gq00, gq11,
-                              w0, w1, w2, w3),
-                              sm.simplify(W[i,i]))
-            Wjfn = njit(nogil=True, fastmath=True, inline='always')(Wsymb)
-
-            Dsymb = lambdify((gp00, gp11,
-                              gq00, gq11,
-                              w0, w1, w2, w3,
-                              v00, v01, v10, v11),
-                              sm.simplify(C[i]))
-            Djfn = njit(nogil=True, fastmath=True, inline='always')(Dsymb)
         else:
             raise ValueError(f"Unknown polarisation product {product}")
 
+        Wsymb = lambdify((gp00, gp11,
+                            gq00, gq11,
+                            w0, w1, w2, w3),
+                            sm.simplify(sm.expand(W[i,i])))
+        Wjfn = njit(nogil=True, fastmath=True, inline='always')(Wsymb)
+
+
+        Dsymb = lambdify((gp00, gp11,
+                            gq00, gq11,
+                            w0, w1, w2, w3,
+                            v00, v01, v10, v11),
+                            sm.simplify(sm.expand(C[i])))
+        Djfn = njit(nogil=True, fastmath=True, inline='always')(Dsymb)
+
         @njit(nogil=True, fastmath=True, inline='always')
         def wfunc(gp, gq, W):
             gp00 = gp[0]