allow bias(k,z) for source window functions

camb/__init__.py

@@ -7,7 +7,7 @@
 __author__ = "Antony Lewis"
 __contact__ = "antony at cosmologist dot info"
 __url__ = "https://camb.readthedocs.io"
-__version__ = "1.3.7"
+__version__ = "1.3.8"
 
 from . import baseconfig
 

camb/baseconfig.py

@@ -504,7 +504,7 @@ class F2003:
                             "sized fields only allowed for ctypes Arrays")
                     if dic["size"] not in [x[0] for x in _fields]:
                         raise CAMBFortranError(
-                            "size must be name of field in same structure (%s for %s)" % (
+                            "size must be the name of a field in same structure (%s for %s)" % (
                                 dic["size"], field_name))
                     new_field = SizedArrayField(field_name, dic["size"])
                     ctypes_fields.append(("_" + field_name, field_type))

camb/sources.py

@@ -1,4 +1,4 @@
-from .baseconfig import F2003Class, fortran_class, numpy_1d, np, numpy_1d_or_null
+from .baseconfig import F2003Class, fortran_class, numpy_1d, numpy_2d, np, numpy_1d_or_null
 from ctypes import POINTER, c_int, c_double, byref
 
 
@@ -35,30 +35,45 @@ class SplinedSourceWindow(SourceWindow):
     """
     _fortran_class_name_ = "TSplinedSourceWindow"
 
-    _methods_ = [("SetTable", [POINTER(c_int), numpy_1d, numpy_1d, numpy_1d_or_null])]
+    _methods_ = [("SetTable", [POINTER(c_int), numpy_1d, numpy_1d, numpy_1d_or_null]),
+                 ("SetTable2DBias", [POINTER(c_int), POINTER(c_int), numpy_1d, numpy_1d, numpy_1d, numpy_2d])
+                 ]
 
     def __init__(self, **kwargs):
         z = kwargs.pop('z', None)
         if z is not None:
-            self.set_table(z, kwargs.pop('W'), kwargs.pop('bias_z', None))
+            self.set_table(z, kwargs.pop('W'), kwargs.pop('bias_z', None),
+                           kwargs.pop('k_bias', None), kwargs.pop('bias_kz', None))
         super().__init__(**kwargs)
 
-    def set_table(self, z, W, bias_z=None):
+    def set_table(self, z, W, bias_z=None, k_bias=None, bias_kz=None):
         """
         Set arrays of z and W(z) for cublic spline interpolation. Note that W(z) is the total count distribution
         observed, not a fractional selection function on an underlying distribution.
 
         :param z: array of redshift values (monotonically increasing)
         :param W: array of window function values. It must be well enough sampled to smoothly cubic-spline interpolate
-        :param bias_z: optional array of bias values at each z
+        :param bias_z: optional array of bias values at each z for scale-independent bias
+        :param k_bias: optional array of k values for bias
+        :param bias_kz: optional 2D contiguous array for space-dependent bias(k, z).
+                        Must ensure range of k is large enough to cover required vaules.
+
         """
         if len(W) != len(z) or z[-1] < z[1] or len(z) < 5:
             raise ValueError(
                 "Redshifts must be well sampled and in ascending order, with window function the same length as z")
+        z = np.ascontiguousarray(z, dtype=np.float64)
+        W = np.ascontiguousarray(W, dtype=np.float64)
         if bias_z is not None:
+            if bias_kz is not None:
+                raise ValueError("set bias_k or bias_zk")
             bias_z = np.ascontiguousarray(bias_z, dtype=np.float64)
             if len(bias_z) != len(z):
                 raise ValueError("bias array must be same size as the redshift array")
-
-        self.f_SetTable(byref(c_int(len(z))), np.ascontiguousarray(z, dtype=np.float64),
-                        np.ascontiguousarray(W, dtype=np.float64), bias_z)
+        if bias_kz is not None:
+            k = np.ascontiguousarray(k_bias, dtype=np.float64)
+            if bias_kz.shape[0] != len(k) or bias_kz.shape[1] != len(z):
+                raise ValueError('Bias array does not match shape of k,z arrays')
+            self.f_SetTable2DBias(byref(c_int(len(z))), byref(c_int(len(k))), z, k, W, bias_kz)
+        else:
+            self.f_SetTable(byref(c_int(len(z))), z, W, bias_z)

camb/tests/camb_test.py

@@ -625,10 +625,18 @@ def testSources(self):
         cls = results.get_source_cls_dict()
         zs = np.arange(0, 0.5, 0.02)
         W = np.exp(-(zs - 0.17) ** 2 / 2 / 0.04 ** 2) / np.sqrt(2 * np.pi) / 0.04
-        pars.SourceWindows[0] = SplinedSourceWindow(bias=1.2, dlog10Ndm=-0.2, z=zs, W=W)
-        results = camb.get_results(pars)
-        cls2 = results.get_source_cls_dict()
-        self.assertTrue(np.allclose(cls2["W1xW1"][2:1200], cls["W1xW1"][2:1200], rtol=1e-3))
+
+        ks = np.logspace(-4, 3, 50)
+        bias_kz = 1.2 * np.ones((len(ks), len(zs)))
+        test_windows = [SplinedSourceWindow(bias=1.2, dlog10Ndm=-0.2, z=zs, W=W),
+                        SplinedSourceWindow(bias_z=1.2 * np.ones_like(zs), dlog10Ndm=-0.2, z=zs, W=W),
+                        SplinedSourceWindow(k_bias=ks, bias_kz=bias_kz, dlog10Ndm=-0.2, z=zs, W=W)]
+        for window in test_windows:
+            pars.SourceWindows[0] = window
+            results = camb.get_results(pars)
+            cls2 = results.get_source_cls_dict()
+            self.assertTrue(np.allclose(cls2["W1xW1"][2:1200], cls["W1xW1"][2:1200], rtol=1e-3))
+
         pars.SourceWindows = [GaussianSourceWindow(redshift=1089, source_type='lensing', sigma=30)]
         results = camb.get_results(pars)
         cls = results.get_source_cls_dict()

fortran/SourceWindows.f90

@@ -2,6 +2,7 @@ module SourceWindows
     use precision
     use Classes
     use MpiUtils
+    use Interpolation, only : TCubicSpline, TInterpGrid2D
     implicit none
 
     integer, parameter :: window_21cm = 1, window_counts = 2, window_lensing = 3
@@ -31,12 +32,14 @@ module SourceWindows
     Type, extends(TSourceWindow) :: TSplinedSourceWindow
         Type(TCubicSpline), allocatable :: Window
         Type(TCubicSpline), allocatable :: Bias_z
+        Type(TInterpGrid2D), allocatable :: Bias_zk
         real(dl) :: maxwin
     contains
     procedure, nopass :: SelfPointer => TSplinedSourceWindow_SelfPointer
     procedure :: count_obs_window_z => TSplinedSourceWindow_count_obs_window_z
     procedure :: GetScales => TSplinedSourceWindow_GetScales
     procedure :: SetTable => TSplinedSourceWindow_SetTable
+    procedure :: SetTable2DBias => TSplinedSourceWindow_SetTable2DBias
     procedure :: GetBias => TSplinedSourceWindow_GetBias
     end Type TSplinedSourceWindow
 
@@ -215,7 +218,18 @@ real(dl) function TSplinedSourceWindow_GetBias(this,k,a)
     class(TSplinedSourceWindow) :: this
     real(dl), intent(in) :: k,a
     real(dl) z
-    if (allocated(this%bias_z))  then
+    integer error
+
+    if (allocated(this%bias_zk)) then
+        z = 1/a-1
+        if (z > this%Window%X(this%Window%n) .or. z < this%Window%X(1)) then
+            TSplinedSourceWindow_GetBias = 0
+        else
+            error = 0
+            TSplinedSourceWindow_GetBias = this%bias_zk%value(z, k, error)
+            if (error /= 0) TSplinedSourceWindow_GetBias = 0
+        end if
+    elseif (allocated(this%bias_z))  then
         z = 1/a-1
         if (z > this%Window%X(this%Window%n) .or. z < this%Window%X(1)) then
             TSplinedSourceWindow_GetBias = 0
@@ -242,21 +256,42 @@ subroutine  TSplinedSourceWindow_SetTable(this, n, z, W, bias_z)
     end if
     if (present(bias_z)) then
         if (allocated(this%Bias_z)) deallocate(this%Bias_z)
+        if (allocated(this%Bias_zk)) deallocate(this%Bias_zk)
         if (n>0) then
             allocate(this%Bias_z)
             call this%Bias_z%Init(z,bias_z)
         end if
     end if
-
     end subroutine TSplinedSourceWindow_SetTable
 
+    subroutine  TSplinedSourceWindow_SetTable2DBias(this, n, nk, z, k, W, bias_zk)
+    class(TSplinedSourceWindow) :: this
+    integer, intent(in) :: n, nk
+    real(dl), intent(in) :: z(n), W(n), k(nk)
+    real(dl), intent(in) :: bias_zk(n,nk)
+
+    if (allocated(this%Window)) deallocate(this%Window)
+    if (n>0) then
+        allocate(this%Window)
+        call this%Window%Init(z,W)
+        this%maxwin = maxval(this%Window%F)
+    end if
+    if (allocated(this%Bias_zk)) deallocate(this%Bias_zk)
+    if (n>0 .and. nk>0) then
+        allocate(this%Bias_zk)
+        call this%Bias_zk%Init(z,k,bias_zk)
+    end if
+
+    end subroutine TSplinedSourceWindow_SetTable2DBias
+
+
     function TSplinedSourceWindow_count_obs_window_z(this, z, winamp)
     !distribution function W(z) for the observed sources, used for lensing and number count spectrum
     !Winamp is amplitude normalized to 1 so the code can tell when the window is very small
     !note this is the total count distribution observed, not a fractional selection function on an underlying distribution
     class(TSplinedSourceWindow) :: this
     real(dl), intent(in) :: z
-    real(dl) TSplinedSourceWindow_count_obs_window_z, dz,winamp
+    real(dl) TSplinedSourceWindow_count_obs_window_z, winamp
 
     if (z > this%Window%X(this%Window%n) .or. z < this%Window%X(1)) then
         TSplinedSourceWindow_count_obs_window_z =0
@@ -270,7 +305,7 @@ end function TSplinedSourceWindow_count_obs_window_z
     subroutine TSplinedSourceWindow_GetScales(this, zpeak, sigma_z, zpeakstart, zpeakend)
     class(TSplinedSourceWindow) :: this
     real(dl), intent(out) :: zpeak, sigma_z, zpeakstart, zpeakend
-    integer ix, i, j
+    integer i, j
     real(dl) z1, zstart, zend, targ
 
     associate(z => this%Window%X, W => this%Window%F, n=>this%Window%n)

fortran/config.f90

@@ -3,7 +3,7 @@ module config
     use constants, only: const_twopi
     implicit none
 
-    character(LEN=*), parameter :: version = '1.3.7'
+    character(LEN=*), parameter :: version = '1.3.8'
 
     integer :: FeedbackLevel = 0 !if >0 print out useful information about the model
 

setup.py

@@ -305,8 +305,9 @@ def run(self):
               'Programming Language :: Python :: 3.6',
               'Programming Language :: Python :: 3.7',
               'Programming Language :: Python :: 3.8',
-              'Programming Language :: Python :: 3.9'
-              'Programming Language :: Python :: 3.10'
+              'Programming Language :: Python :: 3.9',
+              'Programming Language :: Python :: 3.10',
+              'Programming Language :: Python :: 3.11'
           ],
           keywords=['cosmology', 'CAMB', 'CMB'],
           install_requires=['scipy>=1.0', 'sympy>=1.0'],