From 3bd53ecd6ee4766a6715fb221b450fd5c0b3041f Mon Sep 17 00:00:00 2001 From: Kiyam Lin Date: Wed, 23 Apr 2025 15:36:15 +0000 Subject: [PATCH 1/7] initial commit of levin non-limber implementation for glass, does the non-limber in comoving distance first and makes use of the geometric mean --- glass/spectra.py | 175 +++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 175 insertions(+) create mode 100644 glass/spectra.py diff --git a/glass/spectra.py b/glass/spectra.py new file mode 100644 index 000000000..56bee9040 --- /dev/null +++ b/glass/spectra.py @@ -0,0 +1,175 @@ +import numpy as np +import pylevin as levin +from tqdm import tqdm +from scipy.interpolate import RectBivariateSpline, UnivariateSpline + +def make_pk_interpolators(pk, k, z) -> RectBivariateSpline: + """ + Spline interpolator for the matter power spectrum. + + Parameters + ---------- + pk + 3D matter power spectrum. + k + Array of wavenumbers corresponding to 3D matter power spectrum values. + z + Array of redshifts corresponding to 3D matter power spectrum values. + + Returns + ------- + RectBivariateSpline object to interpolate the 3D matter power spectrum. + + """ + + pkz_interp = RectBivariateSpline(k, z, pk) + return pkz_interp + +def make_shell_interpolator(shell_z, + shell_weights, + cosmo + ) -> list[UnivariateSpline]: + """ + Spline interpolator for comoving shells. + + Parameters + ---------- + shell_z + The redshift boundaries for concentric matter shells. + shell_weights + The weights of the shells. + cosmo + Cosmology object to calculate the comoving distances. + + Returns + ------- + A list the length of the number of shells with each element being a UnivariateSpline for that shell. + + """ + + chi = cosmo.comoving_distance(shell_z) + dzdchi = UnivariateSpline(chi, shell_z, k=2, s=0).derivative()(chi) + shell_interp = [UnivariateSpline(chi, shell*dzdchi, k=2, s=0) for shell in shell_weights] + return shell_interp + +def make_glass_shell_interpolator(ws, cosmo) -> list[UnivariateSpline]: + """ + Spline interpolator for comoving shells. + + Parameters + ---------- + ws + Glass shells. + cosmo + Cosmology object to calculate the comoving distances. + + Returns + ------- + A list the length of the number of shells with each element being a UnivariateSpline for that shell. + + """ + return [UnivariateSpline(cosmo.comoving_distance(ws[i].za) * (1+ws[i].za), ws[i].wa, k=2, s=0, ext=1) for i in range(len(ws))] + +def get_cls(num_shells, + ell_min, # can be set to 1 + ell_max, + pk_interpolator, + shells_interpolate, + chi_lims, + N_thread = 4, #number of threads used for hyperthreading + n_sub = 8, #number of collocation points in each bisection + n_bisec_max = 64, #maximum number of bisections used + rel_acc = 1e-10, #relative accuracy target + boost_bessel = True, #should the bessel functions be calculated with boost instead of GSL, higher accuracy at high Bessel orders + verbose = False, #should the code talk to you? + lev_list = None + ): + """ + Calculates Cls using the geometric approximation and performing the integral in comoving distance first. + + Parameters + ---------- + num_shells + The number of shells. + ell_min + The minimum ell value. + ell_max + The maximum ell value. + pk_interpolator + The interpolator for the matter power spectrum. + shells_interpolate + The interpolator for the comoving shells. + chi_lims + The comoving distance limits for the shells. + N_thread + The number of threads used for hyperthreading. + n_sub + The number of collocation points in each bisection. + n_bisec_max + The maximum number of bisections used. + rel_acc + The relative accuracy target. + boost_bessel + Should the bessel functions be calculated with boost instead of GSL, higher accuracy at high Bessel orders. + verbose + Should the code talk to you? + lev_list + List of Levin objects to load. If None, new Levin objects are created. Useful for speeding up the calculation if we have at hand already calculated bisections. + + Returns + ------- + shell_cls + A 3D array of the shell cls. + shell_cls_dict + A dictionary with the shell cls labelled by the shell pair number. + lev_list + A list of Levin objects. If lev_list is None, this will be a list of new Levin objects created during the calculation. If lev_list is not None, this will be the same list passed in as lev_list. + """ + + # Some levin definitions + integral_type = 0 + logx = True # Tells the code to create a logarithmic spline in x for f(x) + logy = True # Tells the code to create a logarithmic spline in y for y = f(x) + + ell = np.unique((np.geomspace(ell_min, ell_max)).astype(int)) + + if lev_list is None: + lev_list = [] + load_levin = False + else: + lev_list + load_levin = True + + N_int = int(2e3) + k_int = np.geomspace(1e-4, 10, N_int) + inner_int = np.zeros((num_shells, len(ell), len(k_int))) + flat_idx = 0 + for idx_shell in (tqdm(np.arange(num_shells))): + chi_int = np.linspace(chi_lims[idx_shell], chi_lims[idx_shell + 1], 50) + pk_nl_new = np.zeros((len(chi_int), len(k_int))) + pk_nl_new[:,:] = np.sqrt(pk_interpolator(k_int, chi_int)).T*(shells_interpolate[idx_shell](chi_int))[:, None] + lower_limit = chi_int[0]*np.ones_like(k_int) + upper_limit = chi_int[-1]*np.ones_like(k_int) + + for i_ell, val_ell in enumerate(ell): + if load_levin is True: + lev = lev_list[flat_idx] + else: + lev = levin.pylevin(integral_type, chi_int, pk_nl_new, logx, logy, N_thread, True) + lev.set_levin(n_sub, n_bisec_max, rel_acc, boost_bessel, verbose) + + ell_values = (val_ell*np.ones_like(k_int)).astype(int) + lev.levin_integrate_bessel_single(lower_limit, upper_limit, k_int, ell_values, inner_int[idx_shell, i_ell,:]) + flat_idx += 1 + + if load_levin is False: + lev_list.append(lev) + + shell_cls = 2/np.pi*np.trapezoid(inner_int[:, None, :, :] * inner_int[None, : ,: ,:]*k_int**2,k_int, axis = -1) + + shell_cls_dict = {} + for i in range(len(chi_lims) - 1): + for j in range(len(chi_lims) - 1): + shell_cls_dict[f'W{i+1}xW{j+1}'] = shell_cls[i, j, :] + + return shell_cls, shell_cls_dict, lev_list \ No newline at end of file From 7056b8149d6174c5240b6f9f1b6fbd5faecf73f2 Mon Sep 17 00:00:00 2001 From: "Patrick J. Roddy" Date: Wed, 7 May 2025 11:55:55 +0100 Subject: [PATCH 2/7] Auto fix linting --- glass/spectra.py | 123 +++++++++++++++++++++++++++++------------------ 1 file changed, 76 insertions(+), 47 deletions(-) diff --git a/glass/spectra.py b/glass/spectra.py index 56bee9040..bfe3e0faa 100644 --- a/glass/spectra.py +++ b/glass/spectra.py @@ -1,7 +1,8 @@ import numpy as np import pylevin as levin -from tqdm import tqdm from scipy.interpolate import RectBivariateSpline, UnivariateSpline +from tqdm import tqdm + def make_pk_interpolators(pk, k, z) -> RectBivariateSpline: """ @@ -11,7 +12,7 @@ def make_pk_interpolators(pk, k, z) -> RectBivariateSpline: ---------- pk 3D matter power spectrum. - k + k Array of wavenumbers corresponding to 3D matter power spectrum values. z Array of redshifts corresponding to 3D matter power spectrum values. @@ -21,14 +22,11 @@ def make_pk_interpolators(pk, k, z) -> RectBivariateSpline: RectBivariateSpline object to interpolate the 3D matter power spectrum. """ - pkz_interp = RectBivariateSpline(k, z, pk) return pkz_interp - -def make_shell_interpolator(shell_z, - shell_weights, - cosmo - ) -> list[UnivariateSpline]: + + +def make_shell_interpolator(shell_z, shell_weights, cosmo) -> list[UnivariateSpline]: """ Spline interpolator for comoving shells. @@ -46,12 +44,14 @@ def make_shell_interpolator(shell_z, A list the length of the number of shells with each element being a UnivariateSpline for that shell. """ - chi = cosmo.comoving_distance(shell_z) dzdchi = UnivariateSpline(chi, shell_z, k=2, s=0).derivative()(chi) - shell_interp = [UnivariateSpline(chi, shell*dzdchi, k=2, s=0) for shell in shell_weights] + shell_interp = [ + UnivariateSpline(chi, shell * dzdchi, k=2, s=0) for shell in shell_weights + ] return shell_interp + def make_glass_shell_interpolator(ws, cosmo) -> list[UnivariateSpline]: """ Spline interpolator for comoving shells. @@ -68,22 +68,33 @@ def make_glass_shell_interpolator(ws, cosmo) -> list[UnivariateSpline]: A list the length of the number of shells with each element being a UnivariateSpline for that shell. """ - return [UnivariateSpline(cosmo.comoving_distance(ws[i].za) * (1+ws[i].za), ws[i].wa, k=2, s=0, ext=1) for i in range(len(ws))] - -def get_cls(num_shells, - ell_min, # can be set to 1 - ell_max, - pk_interpolator, - shells_interpolate, - chi_lims, - N_thread = 4, #number of threads used for hyperthreading - n_sub = 8, #number of collocation points in each bisection - n_bisec_max = 64, #maximum number of bisections used - rel_acc = 1e-10, #relative accuracy target - boost_bessel = True, #should the bessel functions be calculated with boost instead of GSL, higher accuracy at high Bessel orders - verbose = False, #should the code talk to you? - lev_list = None - ): + return [ + UnivariateSpline( + cosmo.comoving_distance(ws[i].za) * (1 + ws[i].za), + ws[i].wa, + k=2, + s=0, + ext=1, + ) + for i in range(len(ws)) + ] + + +def get_cls( + num_shells, + ell_min, # can be set to 1 + ell_max, + pk_interpolator, + shells_interpolate, + chi_lims, + N_thread=4, # number of threads used for hyperthreading + n_sub=8, # number of collocation points in each bisection + n_bisec_max=64, # maximum number of bisections used + rel_acc=1e-10, # relative accuracy target + boost_bessel=True, # should the bessel functions be calculated with boost instead of GSL, higher accuracy at high Bessel orders + verbose=False, # should the code talk to you? + lev_list=None, +): """ Calculates Cls using the geometric approximation and performing the integral in comoving distance first. @@ -125,51 +136,69 @@ def get_cls(num_shells, lev_list A list of Levin objects. If lev_list is None, this will be a list of new Levin objects created during the calculation. If lev_list is not None, this will be the same list passed in as lev_list. """ - # Some levin definitions integral_type = 0 - logx = True # Tells the code to create a logarithmic spline in x for f(x) - logy = True # Tells the code to create a logarithmic spline in y for y = f(x) - + logx = True # Tells the code to create a logarithmic spline in x for f(x) + logy = True # Tells the code to create a logarithmic spline in y for y = f(x) + ell = np.unique((np.geomspace(ell_min, ell_max)).astype(int)) - + if lev_list is None: lev_list = [] load_levin = False else: lev_list load_levin = True - + N_int = int(2e3) k_int = np.geomspace(1e-4, 10, N_int) inner_int = np.zeros((num_shells, len(ell), len(k_int))) flat_idx = 0 - for idx_shell in (tqdm(np.arange(num_shells))): + for idx_shell in tqdm(np.arange(num_shells)): chi_int = np.linspace(chi_lims[idx_shell], chi_lims[idx_shell + 1], 50) pk_nl_new = np.zeros((len(chi_int), len(k_int))) - pk_nl_new[:,:] = np.sqrt(pk_interpolator(k_int, chi_int)).T*(shells_interpolate[idx_shell](chi_int))[:, None] - lower_limit = chi_int[0]*np.ones_like(k_int) - upper_limit = chi_int[-1]*np.ones_like(k_int) - + pk_nl_new[:, :] = ( + np.sqrt(pk_interpolator(k_int, chi_int)).T + * (shells_interpolate[idx_shell](chi_int))[:, None] + ) + lower_limit = chi_int[0] * np.ones_like(k_int) + upper_limit = chi_int[-1] * np.ones_like(k_int) + for i_ell, val_ell in enumerate(ell): if load_levin is True: lev = lev_list[flat_idx] else: - lev = levin.pylevin(integral_type, chi_int, pk_nl_new, logx, logy, N_thread, True) + lev = levin.pylevin( + integral_type, chi_int, pk_nl_new, logx, logy, N_thread, True + ) lev.set_levin(n_sub, n_bisec_max, rel_acc, boost_bessel, verbose) - - ell_values = (val_ell*np.ones_like(k_int)).astype(int) - lev.levin_integrate_bessel_single(lower_limit, upper_limit, k_int, ell_values, inner_int[idx_shell, i_ell,:]) + + ell_values = (val_ell * np.ones_like(k_int)).astype(int) + lev.levin_integrate_bessel_single( + lower_limit, + upper_limit, + k_int, + ell_values, + inner_int[idx_shell, i_ell, :], + ) flat_idx += 1 - + if load_levin is False: lev_list.append(lev) - shell_cls = 2/np.pi*np.trapezoid(inner_int[:, None, :, :] * inner_int[None, : ,: ,:]*k_int**2,k_int, axis = -1) - + shell_cls = ( + 2 + / np.pi + * np.trapezoid( + inner_int[:, None, :, :] * inner_int[None, :, :, :] * k_int**2, + k_int, + axis=-1, + ) + ) + shell_cls_dict = {} for i in range(len(chi_lims) - 1): for j in range(len(chi_lims) - 1): - shell_cls_dict[f'W{i+1}xW{j+1}'] = shell_cls[i, j, :] - - return shell_cls, shell_cls_dict, lev_list \ No newline at end of file + shell_cls_dict[f"W{i + 1}xW{j + 1}"] = shell_cls[i, j, :] + + return shell_cls, shell_cls_dict, lev_list From 14d2eda7f7990fc123699c64795d0c8baaa42d6b Mon Sep 17 00:00:00 2001 From: "Patrick J. Roddy" Date: Wed, 7 May 2025 12:11:06 +0100 Subject: [PATCH 3/7] Fix line lengths --- glass/spectra.py | 39 +++++++++++++++++++++++++++------------ 1 file changed, 27 insertions(+), 12 deletions(-) diff --git a/glass/spectra.py b/glass/spectra.py index bfe3e0faa..b256c03d2 100644 --- a/glass/spectra.py +++ b/glass/spectra.py @@ -41,7 +41,8 @@ def make_shell_interpolator(shell_z, shell_weights, cosmo) -> list[UnivariateSpl Returns ------- - A list the length of the number of shells with each element being a UnivariateSpline for that shell. + A list the length of the number of shells with each element being a + UnivariateSpline for that shell. """ chi = cosmo.comoving_distance(shell_z) @@ -65,7 +66,8 @@ def make_glass_shell_interpolator(ws, cosmo) -> list[UnivariateSpline]: Returns ------- - A list the length of the number of shells with each element being a UnivariateSpline for that shell. + A list the length of the number of shells with each element being a + UnivariateSpline for that shell. """ return [ @@ -87,16 +89,24 @@ def get_cls( pk_interpolator, shells_interpolate, chi_lims, - N_thread=4, # number of threads used for hyperthreading - n_sub=8, # number of collocation points in each bisection - n_bisec_max=64, # maximum number of bisections used - rel_acc=1e-10, # relative accuracy target - boost_bessel=True, # should the bessel functions be calculated with boost instead of GSL, higher accuracy at high Bessel orders - verbose=False, # should the code talk to you? + # number of threads used for hyperthreading + N_thread: int = 4, + # number of collocation points in each bisection + n_sub: int = 8, + # maximum number of bisections used + n_bisec_max: int = 64, + # relative accuracy target + rel_acc: float = 1e-10, + # should the bessel functions be calculated with boost instead of GSL, + # higher accuracy at high Bessel orders + boost_bessel: bool = True, + # should the code talk to you? + verbose: bool = False, lev_list=None, ): """ - Calculates Cls using the geometric approximation and performing the integral in comoving distance first. + Calculates Cls using the geometric approximation and performing the integral + in comoving distance first. Parameters ---------- @@ -121,11 +131,14 @@ def get_cls( rel_acc The relative accuracy target. boost_bessel - Should the bessel functions be calculated with boost instead of GSL, higher accuracy at high Bessel orders. + Should the bessel functions be calculated with boost instead of GSL, + higher accuracy at high Bessel orders. verbose Should the code talk to you? lev_list - List of Levin objects to load. If None, new Levin objects are created. Useful for speeding up the calculation if we have at hand already calculated bisections. + List of Levin objects to load. If None, new Levin objects are created. + Useful for speeding up the calculation if we have at hand already + calculated bisections. Returns ------- @@ -134,7 +147,9 @@ def get_cls( shell_cls_dict A dictionary with the shell cls labelled by the shell pair number. lev_list - A list of Levin objects. If lev_list is None, this will be a list of new Levin objects created during the calculation. If lev_list is not None, this will be the same list passed in as lev_list. + A list of Levin objects. If lev_list is None, this will be a list of new + Levin objects created during the calculation. If lev_list is not None, + this will be the same list passed in as lev_list. """ # Some levin definitions integral_type = 0 From 4b892c0e89d27fa55be6ec1dcba83207ddf7b2f9 Mon Sep 17 00:00:00 2001 From: "Patrick J. Roddy" Date: Wed, 7 May 2025 12:13:27 +0100 Subject: [PATCH 4/7] Revert "Fix line lengths" This reverts commit 14d2eda7f7990fc123699c64795d0c8baaa42d6b. --- glass/spectra.py | 39 ++++++++++++--------------------------- 1 file changed, 12 insertions(+), 27 deletions(-) diff --git a/glass/spectra.py b/glass/spectra.py index b256c03d2..bfe3e0faa 100644 --- a/glass/spectra.py +++ b/glass/spectra.py @@ -41,8 +41,7 @@ def make_shell_interpolator(shell_z, shell_weights, cosmo) -> list[UnivariateSpl Returns ------- - A list the length of the number of shells with each element being a - UnivariateSpline for that shell. + A list the length of the number of shells with each element being a UnivariateSpline for that shell. """ chi = cosmo.comoving_distance(shell_z) @@ -66,8 +65,7 @@ def make_glass_shell_interpolator(ws, cosmo) -> list[UnivariateSpline]: Returns ------- - A list the length of the number of shells with each element being a - UnivariateSpline for that shell. + A list the length of the number of shells with each element being a UnivariateSpline for that shell. """ return [ @@ -89,24 +87,16 @@ def get_cls( pk_interpolator, shells_interpolate, chi_lims, - # number of threads used for hyperthreading - N_thread: int = 4, - # number of collocation points in each bisection - n_sub: int = 8, - # maximum number of bisections used - n_bisec_max: int = 64, - # relative accuracy target - rel_acc: float = 1e-10, - # should the bessel functions be calculated with boost instead of GSL, - # higher accuracy at high Bessel orders - boost_bessel: bool = True, - # should the code talk to you? - verbose: bool = False, + N_thread=4, # number of threads used for hyperthreading + n_sub=8, # number of collocation points in each bisection + n_bisec_max=64, # maximum number of bisections used + rel_acc=1e-10, # relative accuracy target + boost_bessel=True, # should the bessel functions be calculated with boost instead of GSL, higher accuracy at high Bessel orders + verbose=False, # should the code talk to you? lev_list=None, ): """ - Calculates Cls using the geometric approximation and performing the integral - in comoving distance first. + Calculates Cls using the geometric approximation and performing the integral in comoving distance first. Parameters ---------- @@ -131,14 +121,11 @@ def get_cls( rel_acc The relative accuracy target. boost_bessel - Should the bessel functions be calculated with boost instead of GSL, - higher accuracy at high Bessel orders. + Should the bessel functions be calculated with boost instead of GSL, higher accuracy at high Bessel orders. verbose Should the code talk to you? lev_list - List of Levin objects to load. If None, new Levin objects are created. - Useful for speeding up the calculation if we have at hand already - calculated bisections. + List of Levin objects to load. If None, new Levin objects are created. Useful for speeding up the calculation if we have at hand already calculated bisections. Returns ------- @@ -147,9 +134,7 @@ def get_cls( shell_cls_dict A dictionary with the shell cls labelled by the shell pair number. lev_list - A list of Levin objects. If lev_list is None, this will be a list of new - Levin objects created during the calculation. If lev_list is not None, - this will be the same list passed in as lev_list. + A list of Levin objects. If lev_list is None, this will be a list of new Levin objects created during the calculation. If lev_list is not None, this will be the same list passed in as lev_list. """ # Some levin definitions integral_type = 0 From 5e4d8e2bfaf9da0fe1d35594cd82e2f276502a56 Mon Sep 17 00:00:00 2001 From: "Patrick J. Roddy" Date: Wed, 7 May 2025 12:16:10 +0100 Subject: [PATCH 5/7] Fix line lengths --- glass/spectra.py | 39 +++++++++++++++++++++++++++------------ 1 file changed, 27 insertions(+), 12 deletions(-) diff --git a/glass/spectra.py b/glass/spectra.py index bfe3e0faa..dfe4cb879 100644 --- a/glass/spectra.py +++ b/glass/spectra.py @@ -41,7 +41,8 @@ def make_shell_interpolator(shell_z, shell_weights, cosmo) -> list[UnivariateSpl Returns ------- - A list the length of the number of shells with each element being a UnivariateSpline for that shell. + A list the length of the number of shells with each element being a + UnivariateSpline for that shell. """ chi = cosmo.comoving_distance(shell_z) @@ -65,7 +66,8 @@ def make_glass_shell_interpolator(ws, cosmo) -> list[UnivariateSpline]: Returns ------- - A list the length of the number of shells with each element being a UnivariateSpline for that shell. + A list the length of the number of shells with each element being a + UnivariateSpline for that shell. """ return [ @@ -87,16 +89,24 @@ def get_cls( pk_interpolator, shells_interpolate, chi_lims, - N_thread=4, # number of threads used for hyperthreading - n_sub=8, # number of collocation points in each bisection - n_bisec_max=64, # maximum number of bisections used - rel_acc=1e-10, # relative accuracy target - boost_bessel=True, # should the bessel functions be calculated with boost instead of GSL, higher accuracy at high Bessel orders - verbose=False, # should the code talk to you? + # number of threads used for hyperthreading + N_thread=4, + # number of collocation points in each bisection + n_sub=8, + # maximum number of bisections used + n_bisec_max=64, + # relative accuracy target + rel_acc=1e-10, + # should the bessel functions be calculated with boost instead of GSL, + # higher accuracy at high Bessel orders + boost_bessel=True, + # should the code talk to you? + verbose=False, lev_list=None, ): """ - Calculates Cls using the geometric approximation and performing the integral in comoving distance first. + Calculates Cls using the geometric approximation and performing the integral + in comoving distance first. Parameters ---------- @@ -121,11 +131,14 @@ def get_cls( rel_acc The relative accuracy target. boost_bessel - Should the bessel functions be calculated with boost instead of GSL, higher accuracy at high Bessel orders. + Should the bessel functions be calculated with boost instead of GSL, + higher accuracy at high Bessel orders. verbose Should the code talk to you? lev_list - List of Levin objects to load. If None, new Levin objects are created. Useful for speeding up the calculation if we have at hand already calculated bisections. + List of Levin objects to load. If None, new Levin objects are created. + Useful for speeding up the calculation if we have at hand already + calculated bisections. Returns ------- @@ -134,7 +147,9 @@ def get_cls( shell_cls_dict A dictionary with the shell cls labelled by the shell pair number. lev_list - A list of Levin objects. If lev_list is None, this will be a list of new Levin objects created during the calculation. If lev_list is not None, this will be the same list passed in as lev_list. + A list of Levin objects. If lev_list is None, this will be a list of new + Levin objects created during the calculation. If lev_list is not None, + this will be the same list passed in as lev_list. """ # Some levin definitions integral_type = 0 From 09e4faeddde9661838a758db153c0b85ef811642 Mon Sep 17 00:00:00 2001 From: "Patrick J. Roddy" Date: Wed, 7 May 2025 12:18:19 +0100 Subject: [PATCH 6/7] Use `itertools.product` --- glass/spectra.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/glass/spectra.py b/glass/spectra.py index dfe4cb879..081bd59d5 100644 --- a/glass/spectra.py +++ b/glass/spectra.py @@ -1,3 +1,5 @@ +import itertools + import numpy as np import pylevin as levin from scipy.interpolate import RectBivariateSpline, UnivariateSpline @@ -211,9 +213,10 @@ def get_cls( ) ) - shell_cls_dict = {} - for i in range(len(chi_lims) - 1): - for j in range(len(chi_lims) - 1): - shell_cls_dict[f"W{i + 1}xW{j + 1}"] = shell_cls[i, j, :] - + shell_cls_dict = { + f"W{i + 1}xW{j + 1}": shell_cls[i, j, :] + for i, j in itertools.product( + range(len(chi_lims) - 1), range(len(chi_lims) - 1) + ) + } return shell_cls, shell_cls_dict, lev_list From 946c527e45693061c26fc0fd229c316efd7417e3 Mon Sep 17 00:00:00 2001 From: "Patrick J. Roddy" Date: Wed, 7 May 2025 12:50:24 +0100 Subject: [PATCH 7/7] Add some typing --- glass/spectra.py | 36 +++++++++++++++++++++++------------- 1 file changed, 23 insertions(+), 13 deletions(-) diff --git a/glass/spectra.py b/glass/spectra.py index 081bd59d5..e3c45e01e 100644 --- a/glass/spectra.py +++ b/glass/spectra.py @@ -1,12 +1,20 @@ import itertools +from typing import TYPE_CHECKING import numpy as np import pylevin as levin from scipy.interpolate import RectBivariateSpline, UnivariateSpline from tqdm import tqdm +if TYPE_CHECKING: + from numpy.typing import NDArray -def make_pk_interpolators(pk, k, z) -> RectBivariateSpline: + from cosmology import Cosmology + + +def make_pk_interpolators( + pk, k: NDArray[float], z: NDArray[float] +) -> RectBivariateSpline: """ Spline interpolator for the matter power spectrum. @@ -28,7 +36,9 @@ def make_pk_interpolators(pk, k, z) -> RectBivariateSpline: return pkz_interp -def make_shell_interpolator(shell_z, shell_weights, cosmo) -> list[UnivariateSpline]: +def make_shell_interpolator( + shell_z, shell_weights, cosmo: Cosmology +) -> list[UnivariateSpline]: """ Spline interpolator for comoving shells. @@ -55,7 +65,7 @@ def make_shell_interpolator(shell_z, shell_weights, cosmo) -> list[UnivariateSpl return shell_interp -def make_glass_shell_interpolator(ws, cosmo) -> list[UnivariateSpline]: +def make_glass_shell_interpolator(ws, cosmo: Cosmology) -> list[UnivariateSpline]: """ Spline interpolator for comoving shells. @@ -84,26 +94,26 @@ def make_glass_shell_interpolator(ws, cosmo) -> list[UnivariateSpline]: ] -def get_cls( - num_shells, - ell_min, # can be set to 1 - ell_max, +def get_cls( # noqa: PLR0913 + num_shells: int, + ell_min: int, # can be set to 1 + ell_max: int, pk_interpolator, shells_interpolate, chi_lims, # number of threads used for hyperthreading - N_thread=4, + N_thread: int = 4, # number of collocation points in each bisection - n_sub=8, + n_sub: int = 8, # maximum number of bisections used - n_bisec_max=64, + n_bisec_max: int = 64, # relative accuracy target - rel_acc=1e-10, + rel_acc: float = 1e-10, # should the bessel functions be calculated with boost instead of GSL, # higher accuracy at high Bessel orders - boost_bessel=True, + boost_bessel: bool = True, # should the code talk to you? - verbose=False, + verbose: bool = False, lev_list=None, ): """