Merge pull request #173 from reflectometry/contract_profile_fix

use separate rho_M_para and rho_M_perp profiles during contract_mag

refl1d/lib_numba/contract_profile.py

@@ -120,31 +120,37 @@ def align_magnetic(d, sigma, rho, irho, dM, sigmaM, rhoM, thetaM, output_flat):
 @numba.njit(cache=True)
 def contract_mag(d, sigma, rho, irho, rhoM, thetaM, dA):
     n = len(d)
+    m = n - 1 # /* last middle layer */
     i = newi = 1  # /* Skip the substrate */
-    while (i < n):
-
+    while (i < m):
         # /* Get ready for the next layer */
         # /* Accumulation of the first row happens in the inner loop */
-        dz = weighted_dz = 0
-        rhoarea = irhoarea = rhoMarea = thetaMarea = 0.0
+        dz = 0
+        rhoarea = irhoarea = rhoMpara_area = rhoMperp_area = 0.0
         rholo = rhohi = rho[i]
         irholo = irhohi = irho[i]
-        maglo = maghi = rhoM[i] * math.cos(thetaM[i] * math.pi / 180.0)
+        # /* Pre-calculate projections */
+        thetaM_radians = thetaM[i] * math.pi / 180.0
+        rhoMpara = rhoM[i] * math.cos(thetaM_radians)
+        rhoMperp = rhoM[i] * math.sin(thetaM_radians)
+        # /*
+        #  * Note that mpara indicates M when theta_M = 0,
+        #  * and mperp indicates M when theta_M = 90,
+        #  * but in most cases M is parallel to H when theta_M = 270
+        #  * and Aguide = 270
+        #  */
+        mparalo = mparahi = rhoMpara
+        mperplo = mperphi = rhoMperp
 
         # /* Accumulate slices into layer */
-        while True:
+        while (i < m):
             # /* Accumulate next slice */
             dz += d[i]
             rhoarea += d[i] * rho[i]
             irhoarea += d[i] * irho[i]
-
-            # /* Weight the magnetic signal by the in -plane contribution
-            # * when accumulating rhoM and thetaM. */
-            weight = math.cos(thetaM[i]*math.pi/180.)
-            mag = rhoM[i]*weight
-            rhoMarea += d[i]*rhoM[i]*weight
-            thetaMarea += d[i]*thetaM[i]*weight
-            weighted_dz += d[i]*weight
+            # /* Use pre-calculated next values */
+            rhoMpara_area += rhoMpara * d[i]
+            rhoMperp_area += rhoMperp * d[i]
 
             # /* If no more slices or sigma != 0, break immediately */
             i += 1
@@ -166,33 +172,54 @@ def contract_mag(d, sigma, rho, irho, rhoM, thetaM, dA):
             if ((irhohi-irholo)*(dz+d[i]) > dA):
                 break
 
-            if (mag < maglo):
-                maglo = mag
-            if (mag > maghi):
-                maghi = mag
-            if ((maghi-maglo)*(dz+d[i]) > dA):
+            # /* Pre-calculate projections of next layer */
+            thetaM_radians = thetaM[i] * math.pi / 180.0
+            rhoMpara = rhoM[i] * math.cos(thetaM_radians)
+            rhoMperp = rhoM[i] * math.sin(thetaM_radians)
+
+            if (rhoMpara < mparalo):
+                mparalo = rhoMpara
+            if (rhoMpara > mparahi):
+                mparahi = rhoMpara
+            if ((mparahi-mparalo)*(dz+d[i]) > dA):
                 break
 
-        # /* Save the layer */
-        assert(newi < n)
+            if (rhoMperp < mperplo):
+                mperplo = rhoMperp
+            if (rhoMperp > mperphi):
+                mperphi = rhoMperp
+            if ((mperphi-mperplo)*(dz+d[i]) > dA):
+                break
+
+        assert(newi < m)
         d[newi] = dz
-        if (i == n):
-            # /* Last layer uses surface values */
-            rho[newi] = rho[n-1]
-            irho[newi] = irho[n-1]
-            rhoM[newi] = rhoM[n-1]
-            thetaM[newi] = thetaM[n-1]
-            # /* No interface for final layer */
+        if (dz == 0):
+            rho[newi] = rho[i-1]
+            irho[newi] = irho[i-1]
+            rhoM[newi] = rhoM[i-1]
+            thetaM[newi] = thetaM[i-1]
         else:
-            # /* Middle layers uses average values */
             rho[newi] = rhoarea / dz
             irho[newi] = irhoarea / dz
-            rhoM[newi] = rhoMarea / weighted_dz
-            thetaM[newi] = thetaMarea / weighted_dz
-            sigma[newi] = sigma[i-1]
-            # /* First layer uses substrate values */
+            mean_rhoMpara = rhoMpara_area / dz
+            mean_rhoMperp = rhoMperp_area / dz
+            mean_rhoM = math.sqrt(mean_rhoMpara**2 + mean_rhoMperp**2)
+            thetaM_from_mean = math.atan2(mean_rhoMperp, mean_rhoMpara) * 180.0 / math.pi
+            rhoM[newi] = mean_rhoM
+            thetaM[newi] = thetaM_from_mean
+        sigma[newi] = sigma[i-1]
+
         newi += 1
 
+    # /* Save the last layer */
+    # /* Last layer uses surface values */
+    rho[newi] = rho[n-1]
+    irho[newi] = irho[n-1]
+    rhoM[newi] = rhoM[n-1]
+    thetaM[newi] = thetaM[n-1]
+    # /* No interface for final layer */
+    newi += 1
+
     return newi