Reducing duplication in cellmesh integrators (#681)

* Polarization angle and disk blocking calculations separated to their own functions to decrease duplication.

* Unit tests for cellmesh tools added.

* Changelog and preliminary version update information added.

* The preliminary version syntax changed to make the installation to work.

* image_order_limit default fixed back to original

* Switch using the same new version number as already in the main branch.

* Debug print removed.

* Fixing a bug in the example without disk blocking.

* Likelihood check value updated to account for the fix in the Gaussian likelihood normalization sign done already in the past.

---------

Co-authored-by: Sebastien Guillot <sebguillot60@gmail.com>

Author: thjsal

changelog.d/681.attribution.rst

@@ -0,0 +1 @@
+Tuomo Salmi

changelog.d/681.changed.rst

@@ -0,0 +1 @@
+Reduced repetitive codes in the cellmesh integrators by calculating polarization angle and accretion disk blocking in their own functions. Unit tests were also added to these.

examples/examples_modeling_tutorial/TestRun_PolNum_split_1spot.py

@@ -77,7 +77,7 @@
 
 bounds = dict(spin_axis_position_angle = (None, None))
 photosphere = CustomPhotosphere_NumA5(hot = hot, elsewhere = elsewhere, stokes=True, bounds=bounds,
-                                values=dict(mode_frequency = spacetime['frequency']))
+                                values=dict(mode_frequency = spacetime['frequency']),disk_blocking=False)
 
 photosphere.hot_atmosphere = this_directory+'/model_data/Bobrikova_compton_slab_I.npz'
 photosphere.hot_atmosphere_Q = this_directory+'/model_data/Bobrikova_compton_slab_Q.npz'

examples/examples_modeling_tutorial/TestRun_PolNum_split_inference.py

@@ -714,7 +714,7 @@ def transform(self, p, **kwargs):
 likelihood.reinitialise()
 likelihood.clear_cache()
 
-true_logl = -1.5663833842e+06
+true_logl = -1.5663327873e+06
 
 if __name__ == '__main__': # sample from the posterior
     # inform source code that parameter objects updated when inverse sampling

examples/examples_modeling_tutorial/modules/CustomPhotosphere.py

@@ -194,7 +194,8 @@ def integrate(self, energies, threads):
                 except:
                     else_atm_ext = None
 
-                R_in = self.disk['R_in'] * 1000
+                if self.disk is not None:
+                    R_in = self.disk['R_in'] * 1000
                 if self._stokes:
                     if self._disk_blocking:
                         self._signal, self._signalQ, self._signalU  = self._hot.integrate_stokes(self._spacetime,

pyproject.toml

@@ -12,7 +12,7 @@ build-backend = "setuptools.build_meta"
 # ---------------------------------------------------------
 [project]
 name = "xpsi"
-version = "3.2.1"
+version = "3.3.0"
 authors = [
     { name = "The X-PSI Core Team", email = "A.L.Watts@uva.nl" },
 ]

setup.py

@@ -235,6 +235,7 @@
     "xpsi.surface_radiation_field.elsewhere_user",
     "xpsi.surface_radiation_field.elsewhere_wrapper",
     "xpsi.cellmesh.integrator",
+    "xpsi.cellmesh.common_functions",
     "xpsi.cellmesh.integratorIQU",
     "xpsi.cellmesh.integrator_for_azimuthal_invariance",
     "xpsi.cellmesh.integrator_for_azimuthal_invariance_split",

xpsi/HotRegion.py

@@ -1078,6 +1078,9 @@ def integrate(self,
                   'smaller than 90 degrees. However, inclination > 90 degrees.'
                   ' Formula for r_psi_d in the integrator is incorrect in this'
                   ' case.')
+
+        if (not self.symmetry) and (R_in != 1e6):
+            print("WARNING: Inner disc radius was given, although the disk blocking model has not been implemented for the chosen integrator (symmetry=False).")
 
         super_pulse = self._integrator(threads,
                                        st.R,
@@ -1199,6 +1202,15 @@ def integrate_stokes(self,
                 raise AtmosError('The numerical atmosphere data were not preloaded, '
                                  'even though that is required by the current atmosphere extension.')
 
+        if st.i > _np.pi/2 and R_in != 1e6:
+            print('WARNING: disk blocking assumes that the inclination is '
+                  'smaller than 90 degrees. However, inclination > 90 degrees.'
+                  ' Formula for r_psi_d in the integrator is incorrect in this'
+                  ' case.')
+
+        if (not self.symmetry) and (R_in != 1e6):
+            print("WARNING: Inner disc radius was given, although the disk blocking model has not been implemented for the chosen integrator (symmetry=False).")
+
         all_pulses = self._integratorIQU(threads,
                                    st.R,
                                    st.Omega,

xpsi/cellmesh/common_functions.pxd

@@ -0,0 +1,29 @@
+cdef double compute_pol_ang(
+              double leaves_kdx,
+              double sin_psi,
+              double cos_psi,
+              double sin_alpha,
+              double cos_alpha,
+              double sin_theta_i,
+              double cos_theta_i,
+              double sin_i,
+              double cos_i,
+              double sin_gamma,
+              double cos_gamma,
+              double Grav_z,
+              double mu,
+              double eta,
+              double beta,
+              double Lorentz,
+              double cos_xi) noexcept nogil
+
+cdef int disk_block(
+              double R_in,
+              double cos_i,
+              double cos_psi,
+              double cos_theta_i,
+              double r_s_over_r,
+              double radius,
+              double sin_alpha,
+              double theta_i_over_pi
+              ) noexcept nogil

xpsi/cellmesh/common_functions.pyx

@@ -0,0 +1,165 @@
+from libc.math cimport M_PI, sqrt, sin, cos, acos, log10, pow, exp, fabs, ceil, log, atan2
+from ..tools.core cimport are_equal
+from libc.stdio cimport printf
+
+cdef double _2pi = 2.0 * M_PI
+
+cdef double compute_pol_ang(
+              double leaves_kdx,
+              double sin_psi,
+              double cos_psi,
+              double sin_alpha,
+              double cos_alpha,
+              double sin_theta_i,
+              double cos_theta_i,
+              double sin_i,
+              double cos_i,
+              double sin_gamma,
+              double cos_gamma,
+              double Grav_z,
+              double mu,
+              double eta,
+              double beta,
+              double Lorentz,
+              double cos_xi) noexcept nogil:
+    """
+    Calculation of polarization angle using the polarized Oblate Schwarcshild approximation as in Loktev et al. (2020).
+    """
+     
+    cdef:         
+        double sin_chi_0, cos_chi_0, chi_0, chi_1, chi_prime, chi
+        double sin_chi_1, cos_chi_1, sin_chi_prime, cos_chi_prime
+        double sin_lambda, cos_lambda, cos_eps
+        double sin_alpha_over_sin_psi        
+
+    sin_chi_0 = - sin_theta_i*sin(leaves_kdx) 
+    cos_chi_0 = sin_i*cos_theta_i - sin_theta_i*cos_i*cos(leaves_kdx)
+    chi_0 = atan2(sin_chi_0,cos_chi_0)
+
+    if not are_equal(sin_psi, 0.0):
+        sin_alpha_over_sin_psi = sin_alpha/sin_psi
+    else: #using small-angle limit of the Beloborodov (2002) approximation
+        sin_alpha_over_sin_psi = Grav_z
+
+    #Notes: mu = cos_sigma , Lorentz = 1/Gamma, mu0=eta*mu, cos_xi defined with no minus sign
+    sin_chi_1 = sin_gamma*sin_i*sin(leaves_kdx)*sin_alpha_over_sin_psi #times sin alpha sin sigma
+    cos_chi_1 = cos_gamma - cos_alpha*mu  #times sin alpha sin sigma 
+    chi_1 = atan2(sin_chi_1,cos_chi_1)
+
+    sin_lambda = sin_theta_i*cos_gamma - sin_gamma*cos_theta_i
+    cos_lambda = cos_theta_i*cos_gamma + sin_theta_i*sin_gamma
+    cos_eps = sin_alpha_over_sin_psi*(cos_i*sin_lambda - sin_i*cos_lambda*cos(leaves_kdx) + cos_psi*sin_gamma) - cos_alpha*sin_gamma
+
+    sin_chi_prime = cos_eps*eta*mu*beta/Lorentz
+    cos_chi_prime = (1. - mu**2 /(1. + beta*cos_xi))
+    chi_prime = atan2(sin_chi_prime,cos_chi_prime)
+
+    chi = chi_0+chi_1+chi_prime
+
+    #printf("leaves[_kdx] = %.6e ",leaves_kdx/_2pi)
+    #printf("chi_0 = %.6e ",chi_0)
+    #printf("chi_1 = %.6e ",chi_1)
+    #printf("chi_prime = %.6e ",chi_prime)
+    #printf("PA_tot = %.6e\n",chi)
+
+    return chi
+
+# Python wrapper for testing
+cpdef double compute_pol_ang_py(
+              double leaves_kdx,
+              double sin_psi,
+              double cos_psi,
+              double sin_alpha,
+              double cos_alpha,
+              double sin_theta_i,
+              double cos_theta_i,
+              double sin_i,
+              double cos_i,
+              double sin_gamma,
+              double cos_gamma,
+              double Grav_z,
+              double mu,
+              double eta,
+              double beta,
+              double Lorentz,
+              double cos_xi):
+    """
+    Python-callable wrapper for compute_pol_ang.
+    Useful for unit testing.
+    """
+    return compute_pol_ang(
+        leaves_kdx,
+        sin_psi,
+        cos_psi,
+        sin_alpha,
+        cos_alpha,
+        sin_theta_i,
+        cos_theta_i,
+        sin_i,
+        cos_i,
+        sin_gamma,
+        cos_gamma,
+        Grav_z,
+        mu,
+        eta,
+        beta,
+        Lorentz,
+        cos_xi
+    )
+
+cdef int disk_block(
+              double R_in,
+              double cos_i,
+              double cos_psi,
+              double cos_theta_i,
+              double r_s_over_r_i,
+              double radius,
+              double sin_alpha,
+              double theta_i_over_pi
+              ) noexcept nogil:
+    """
+    Checking whether an accretion disk blocks certain rays based on Ibragimov & Poutanen (2009). Returns 1 if the ray is not blocked.
+    """
+     
+    cdef:         
+        double cos_psi_d, sin_psi_d      # geometric quantities
+        double impact_b, r_psi_d         # impact parameter and radial coordinate
+        double r_s_i                     # schwarzschild radius at the cell location  
+
+    cos_psi_d = (cos_i * cos_psi - cos_theta_i) / sqrt(cos_i * cos_i + cos_theta_i * cos_theta_i - 2 * cos_i * cos_theta_i * cos_psi) #Ibragimov & Poutanen (2009), Equation (C2)
+    sin_psi_d = sqrt(1 - cos_psi_d * cos_psi_d)
+    r_s_i = r_s_over_r_i*radius
+    impact_b = radius * sin_alpha / sqrt(1 - r_s_over_r_i) # impact parameter
+    r_psi_d = sqrt((r_s_i * r_s_i * (1 - cos_psi_d) * (1 - cos_psi_d)) / (4 * (1 + cos_psi_d) * (1 + cos_psi_d)) +  ((impact_b * impact_b) / (sin_psi_d * sin_psi_d))) - (r_s_i * (1 - cos_psi_d)) / (2 * (1 + cos_psi_d)) ##Ibragimov & Poutanen (2009), Equation (B9)
+
+    if theta_i_over_pi < 0.5 or (theta_i_over_pi > 0.5 and r_psi_d < R_in):
+        return 1
+    else: #theta > pi/2 and (theta < pi/2 or r_psi_d > R_in), don't calculate.
+        return 0
+
+# Python wrapper for testing
+cpdef int disk_block_py(
+              double R_in,
+              double cos_i,
+              double cos_psi,
+              double cos_theta_i,
+              double r_s_over_r_i,
+              double radius,
+              double sin_alpha,
+              double theta_i_over_pi
+              ):
+    """
+    Python-callable wrapper for disk_ block.
+    Useful for unit testing.
+    """
+    return disk_block(
+            R_in,
+            cos_i,
+            cos_psi,
+            cos_theta_i,
+            r_s_over_r_i,
+            radius,
+            sin_alpha,
+            theta_i_over_pi
+    )
+