classy-szfast 0.0.3__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (27) hide show
  1. classy_szfast-0.0.3/.gitignore +7 -0
  2. classy_szfast-0.0.3/MANIFEST.in +1 -0
  3. classy_szfast-0.0.3/PKG-INFO +16 -0
  4. classy_szfast-0.0.3/README.md +1 -0
  5. classy_szfast-0.0.3/classy_szfast/__init__.py +10 -0
  6. classy_szfast-0.0.3/classy_szfast/classy_sz.py +585 -0
  7. classy_szfast-0.0.3/classy_szfast/classy_szfast.py +745 -0
  8. classy_szfast-0.0.3/classy_szfast/config.py +1 -0
  9. classy_szfast-0.0.3/classy_szfast/cosmopower.py +203 -0
  10. classy_szfast-0.0.3/classy_szfast/cosmosis_classy_szfast_interface.py +331 -0
  11. classy_szfast-0.0.3/classy_szfast/custom_bias/__init__.py +0 -0
  12. classy_szfast-0.0.3/classy_szfast/custom_bias/custom_bias.py +20 -0
  13. classy_szfast-0.0.3/classy_szfast/custom_profiles/__init__.py +0 -0
  14. classy_szfast-0.0.3/classy_szfast/custom_profiles/custom_profiles.py +45 -0
  15. classy_szfast-0.0.3/classy_szfast/pks_and_sigmas.py +150 -0
  16. classy_szfast-0.0.3/classy_szfast/suppress_warnings.py +10 -0
  17. classy_szfast-0.0.3/classy_szfast/utils.py +62 -0
  18. classy_szfast-0.0.3/classy_szfast.egg-info/PKG-INFO +16 -0
  19. classy_szfast-0.0.3/classy_szfast.egg-info/SOURCES.txt +26 -0
  20. classy_szfast-0.0.3/classy_szfast.egg-info/dependency_links.txt +1 -0
  21. classy_szfast-0.0.3/classy_szfast.egg-info/not-zip-safe +1 -0
  22. classy_szfast-0.0.3/classy_szfast.egg-info/requires.txt +4 -0
  23. classy_szfast-0.0.3/classy_szfast.egg-info/top_level.txt +1 -0
  24. classy_szfast-0.0.3/pypi_upload.sh +15 -0
  25. classy_szfast-0.0.3/pyproject.toml +21 -0
  26. classy_szfast-0.0.3/setup.cfg +7 -0
  27. classy_szfast-0.0.3/setup.py +17 -0
@@ -0,0 +1,7 @@
1
+
2
+ .DS_Store
3
+ *egg-info
4
+ test_classy_szfast.py
5
+ test_classy_szfast-mnu3state.py
6
+ __pycache__
7
+ requirements.txt
@@ -0,0 +1 @@
1
+
@@ -0,0 +1,16 @@
1
+ Metadata-Version: 2.1
2
+ Name: classy_szfast
3
+ Version: 0.0.3
4
+ Summary: The accelerator of the class_sz code from https://github.com/CLASS-SZ
5
+ Home-page: https://github.com/CLASS-SZ/classy_szfast
6
+ Download-URL: https://github.com/CLASS-SZ/classy_szfast
7
+ Author: Boris Bolliet, Ola Kusiak
8
+ Author-email: bb667@cam.ac.uk, akk2175@columbia.edu
9
+ Maintainer-email: Boris Bolliet <bb667@cam.ac.uk>
10
+ Project-URL: Homepage, https://github.com/CLASS-SZ
11
+ Project-URL: GitHub, https://github.com/CLASS-SZ
12
+ Description-Content-Type: text/markdown
13
+ Requires-Dist: tensorflow==2.13.0
14
+ Requires-Dist: tensorflow-probability==0.21.0
15
+ Requires-Dist: cosmopower
16
+ Requires-Dist: mcfit
@@ -0,0 +1 @@
1
+ Python module for the acceleration of class_sz and interface with likelihood codes.
@@ -0,0 +1,10 @@
1
+ from classy_szfast.classy_szfast import Class_szfast
2
+ from .config import *
3
+ from .utils import *
4
+ from .cosmopower import *
5
+ from .pks_and_sigmas import *
6
+
7
+ from pathlib import Path
8
+ import sys
9
+ from .custom_profiles import *
10
+ from .custom_bias import *
@@ -0,0 +1,585 @@
1
+ from cobaya.theories.classy import classy
2
+ from copy import deepcopy
3
+ from typing import NamedTuple, Sequence, Union, Optional
4
+ from cobaya.tools import load_module
5
+ import logging
6
+ import os
7
+ import numpy as np
8
+ import time
9
+
10
+ class classy_sz(classy):
11
+ use_class_sz_fast_mode = 0 # this is passed in the yaml file
12
+ use_class_sz_no_cosmo_mode = 0 # this is passed in the yaml file
13
+ lensing_lkl = "SOLikeT"
14
+ # skip_background_and_thermo = True
15
+ # ell_factor = False # True for pyactlite and bplike, False for clik
16
+
17
+ def initialize(self):
18
+ """Importing CLASS from the correct path, if given, and if not, globally."""
19
+ self.classy_module = self.is_installed()
20
+ if not self.classy_module:
21
+ raise NotInstalledError(
22
+ self.log, "Could not find CLASS_SZ. Check error message above.")
23
+ from classy_sz import Class, CosmoSevereError, CosmoComputationError
24
+ global CosmoComputationError, CosmoSevereError
25
+ self.classy = Class()
26
+ super(classy,self).initialize()
27
+ # Add general CLASS stuff
28
+ self.extra_args["output"] = self.extra_args.get("output", "")
29
+ if "sBBN file" in self.extra_args:
30
+ self.extra_args["sBBN file"] = (
31
+ self.extra_args["sBBN file"].format(classy=self.path))
32
+ # Derived parameters that may not have been requested, but will be necessary later
33
+ self.derived_extra = []
34
+ self.log.info("Initialized!")
35
+
36
+ if self.use_class_sz_no_cosmo_mode == 1:
37
+ self.log.info("Initializing cosmology part!")
38
+ initial_parameters = self.extra_args.copy()
39
+ # print("initial_parameters:",initial_parameters)
40
+
41
+ self.classy.set(initial_parameters)
42
+ self.classy.compute_class_szfast()
43
+ self.log.info("cosmology part initialized!")
44
+
45
+
46
+ # print(self.lensing_lkl)
47
+ # exit(0)
48
+
49
+ # # class_sz default params for lkl
50
+ # self.extra_args["output"] = 'tSZ_1h'
51
+ # self.extra_args["multipoles_sz"] = 'P15'
52
+ # self.extra_args['nlSZ'] = 18
53
+
54
+
55
+ # # here modify if you want to bypass stuff in the class computation
56
+ # def calculate(self, state, want_derived=True, **params_values_dict):
57
+ # print("Bypassing class_sz")
58
+
59
+
60
+
61
+
62
+ def must_provide(self, **requirements):
63
+
64
+ if "Cl_sz" in requirements:
65
+ # make sure cobaya still runs as it does for standard classy
66
+ requirements.pop("Cl_sz")
67
+ # specify the method to collect the new observable
68
+ self.collectors["Cl_sz"] = Collector(
69
+ method="cl_sz", # name of the method in classy.pyx
70
+ args_names=[],
71
+ args=[])
72
+
73
+ if "Cl_yxg" in requirements:
74
+ # make sure cobaya still runs as it does for standard classy
75
+ requirements.pop("Cl_yxg")
76
+ # specify the method to collect the new observable
77
+ self.collectors["Cl_yxg"] = Collector(
78
+ method="cl_yg", # name of the method in classy.pyx
79
+ args_names=[],
80
+ args=[])
81
+
82
+ if "Cl_gxg" in requirements:
83
+ # make sure cobaya still runs as it does for standard classy
84
+ requirements.pop("Cl_gxg")
85
+ # specify the method to collect the new observable
86
+ self.collectors["Cl_gxg"] = Collector(
87
+ method="cl_gg", # name of the method in classy.pyx
88
+ args_names=[],
89
+ args=[])
90
+
91
+
92
+ if "Cl_gxmu" in requirements:
93
+ # make sure cobaya still runs as it does for standard classy
94
+ requirements.pop("Cl_gxmu")
95
+ # specify the method to collect the new observable
96
+ self.collectors["Cl_gxmu"] = Collector(
97
+ method="cl_gm", # name of the method in classy.pyx
98
+ args_names=[],
99
+ args=[])
100
+
101
+
102
+ if "Cl_muxmu" in requirements:
103
+ # make sure cobaya still runs as it does for standard classy
104
+ requirements.pop("Cl_muxmu")
105
+ # specify the method to collect the new observable
106
+ self.collectors["Cl_muxmu"] = Collector(
107
+ method="cl_mm", # name of the method in classy.pyx
108
+ args_names=[],
109
+ args=[])
110
+
111
+
112
+ if "Cl_kxg" in requirements:
113
+ # make sure cobaya still runs as it does for standard classy
114
+ requirements.pop("Cl_kxg")
115
+ # specify the method to collect the new observable
116
+ self.collectors["Cl_kxg"] = Collector(
117
+ method="cl_kg", # name of the method in classy.pyx
118
+ args_names=[],
119
+ args=[])
120
+
121
+ if "Cl_kxmu" in requirements:
122
+ # make sure cobaya still runs as it does for standard classy
123
+ requirements.pop("Cl_kxmu")
124
+ # specify the method to collect the new observable
125
+ self.collectors["Cl_kxmu"] = Collector(
126
+ method="cl_km", # name of the method in classy.pyx
127
+ args_names=[],
128
+ args=[])
129
+
130
+
131
+ if "Cl_yxmu" in requirements:
132
+ # make sure cobaya still runs as it does for standard classy
133
+ requirements.pop("Cl_yxmu")
134
+ # specify the method to collect the new observable
135
+ self.collectors["Cl_yxmu"] = Collector(
136
+ method="cl_ym", # name of the method in classy.pyx
137
+ args_names=[],
138
+ args=[])
139
+ if "Cl_kgxg" in requirements:
140
+ # make sure cobaya still runs as it does for standard classy
141
+ requirements.pop("Cl_kgxg")
142
+ # specify the method to collect the new observable
143
+ self.collectors["Cl_kgxg"] = Collector(
144
+ method="cl_ggamma", # name of the method in classy.pyx
145
+ args_names=[],
146
+ args=[])
147
+ if "Cl_kgxmu" in requirements:
148
+ # make sure cobaya still runs as it does for standard classy
149
+ requirements.pop("Cl_kgxmu")
150
+ # specify the method to collect the new observable
151
+ self.collectors["Cl_kgxmu"] = Collector(
152
+ method="cl_kg_m", # name of the method in classy.pyx
153
+ args_names=[],
154
+ args=[])
155
+ if "Cl_IAxg" in requirements:
156
+ # make sure cobaya still runs as it does for standard classy
157
+ requirements.pop("Cl_IAxg")
158
+ # specify the method to collect the new observable
159
+ self.collectors["Cl_IAxg"] = Collector(
160
+ method="cl_IA_g", # name of the method in classy.pyx
161
+ args_names=[],
162
+ args=[])
163
+
164
+ if "sz_binned_cluster_counts" in requirements:
165
+ # make sure cobaya still runs as it does for standard classy
166
+ requirements.pop("sz_binned_cluster_counts")
167
+ # specify the method to collect the new observable
168
+ self.collectors["sz_binned_cluster_counts"] = Collector(
169
+ method="dndzdy_theoretical", # name of the method in classy.pyx
170
+ args_names=[],
171
+ args=[])
172
+
173
+ if "sz_unbinned_cluster_counts" in requirements:
174
+ # make sure cobaya still runs as it does for standard classy
175
+ requirements.pop("sz_unbinned_cluster_counts")
176
+ # specify the method to collect the new observable
177
+ self.collectors["sz_unbinned_cluster_counts"] = Collector(
178
+ method="szcounts_ntot_rates_loglike", # name of the method in classy.pyx
179
+ args_names=[],
180
+ args=[])
181
+
182
+ if "cl_cib_kappa" in requirements:
183
+ # make sure cobaya still runs as it does for standard classy
184
+ requirements.pop("cl_cib_kappa")
185
+ # specify the method to collect the new observable
186
+ self.collectors["cl_cib_kappa"] = Collector(
187
+ method="cl_lens_cib", # name of the method in classy.pyx
188
+ args_names=[],
189
+ args=[])
190
+ if "cl_galn_galn" in requirements:
191
+ # make sure cobaya still runs as it does for standard classy
192
+ requirements.pop("cl_galn_galn")
193
+ # specify the method to collect the new observable
194
+ self.collectors["cl_galn_galn"] = Collector(
195
+ method="cl_galn_galn", # name of the method in classy.pyx
196
+ args_names=[],
197
+ args=[])
198
+ if "cl_galn_lens" in requirements:
199
+ # make sure cobaya still runs as it does for standard classy
200
+ requirements.pop("cl_galn_lens")
201
+ # specify the method to collect the new observable
202
+ self.collectors["cl_galn_lens"] = Collector(
203
+ method="cl_galn_lens", # name of the method in classy.pyx
204
+ args_names=[],
205
+ args=[])
206
+
207
+ if "Cl_galnxtsz" in requirements:
208
+ # make sure cobaya still runs as it does for standard classy
209
+ requirements.pop("Cl_galnxtsz")
210
+ # specify the method to collect the new observable
211
+ self.collectors["Cl_galnxtsz"] = Collector(
212
+ method="cl_galn_tsz", # name of the method in classy.pyx
213
+ args_names=[],
214
+ args=[])
215
+ if "Cl_galnxgallens" in requirements:
216
+ # make sure cobaya still runs as it does for standard classy
217
+ requirements.pop("Cl_galnxgallens")
218
+ # specify the method to collect the new observable
219
+ self.collectors["Cl_galnxgallens"] = Collector(
220
+ method="cl_galn_gallens", # name of the method in classy.pyx
221
+ args_names=[],
222
+ args=[])
223
+ if "Cl_lensmagnxtsz" in requirements:
224
+ # make sure cobaya still runs as it does for standard classy
225
+ requirements.pop("Cl_lensmagnxtsz")
226
+ # specify the method to collect the new observable
227
+ self.collectors["Cl_lensmagnxtsz"] = Collector(
228
+ method="cl_lensmagn_tsz", # name of the method in classy.pyx
229
+ args_names=[],
230
+ args=[])
231
+ if "Cl_lensmagnxgallens" in requirements:
232
+ # make sure cobaya still runs as it does for standard classy
233
+ requirements.pop("Cl_lensmagnxgallens")
234
+ # specify the method to collect the new observable
235
+ self.collectors["Cl_lensmagnxgallens"] = Collector(
236
+ method="cl_lensmagn_gallens", # name of the method in classy.pyx
237
+ args_names=[],
238
+ args=[])
239
+ if "Cl_galnxIA" in requirements:
240
+ # make sure cobaya still runs as it does for standard classy
241
+ requirements.pop("Cl_galnxIA")
242
+ # specify the method to collect the new observable
243
+ self.collectors["Cl_galnxIA"] = Collector(
244
+ method="cl_galn_IA", # name of the method in classy.pyx
245
+ args_names=[],
246
+ args=[])
247
+ super().must_provide(**requirements)
248
+
249
+ # get the required new observable
250
+ def get_Cl(self, ell_factor=False, units="FIRASmuK2"):
251
+ if self.use_class_sz_fast_mode:
252
+ return self.get_Clfast(ell_factor=ell_factor)
253
+ else:
254
+ return self._get_Cl(ell_factor=ell_factor, units=units, lensed=True)
255
+
256
+ def get_Clfast(self,ell_factor = False):
257
+ # print('ell_factor:',self.ell_factor)
258
+ # exit(0)
259
+ cls = {}
260
+ cls = deepcopy(self._current_state["Cl"])
261
+ # ell_factor = self.ell_factor
262
+ # print(cls)
263
+ # exit(0)
264
+ # print('in get clfast:',ell_factor)
265
+ # print(cls)
266
+ lcp = np.asarray(cls['ell'])
267
+ # print(self.lensing_lkl)
268
+ if ell_factor==True:
269
+ cls['tt'] *= (2.7255e6)**2.*(lcp*(lcp+1.))/2./np.pi
270
+ cls['te'] *= (2.7255e6)**2.*(lcp*(lcp+1.))/2./np.pi
271
+ cls['ee'] *= (2.7255e6)**2.*(lcp*(lcp+1.))/2./np.pi
272
+
273
+ if self.lensing_lkl == "ACT" and ell_factor == False:
274
+ cls['tt'] *= (2.7255e6)**2.
275
+ cls['te'] *= (2.7255e6)**2.
276
+ cls['ee'] *= (2.7255e6)**2.
277
+ # print(cls['tt'][1230])
278
+ # print(cls['te'][1230])
279
+ # print(cls['ee'][1230])
280
+ # exit(0)
281
+ if self.lensing_lkl == "SOLikeT":
282
+ cls['pp'] *= (lcp*(lcp+1.))**2./4.
283
+ elif self.lensing_lkl == "ACT":
284
+ cls['pp'] *= 1.#(lcp*(lcp+1.))**2./4.
285
+ else: # here for the planck lensing lkl, using lfactor option gives:
286
+ cls['pp'] *= (lcp*(lcp+1.))**2.*1./2./np.pi
287
+ return cls
288
+ # get the required new observable
289
+ def get_Cl_sz(self):
290
+ cls = {}
291
+ cls = deepcopy(self._current_state["Cl_sz"])
292
+ return cls
293
+
294
+ def get_Cl_yxg(self):
295
+ cls = {}
296
+ cls = deepcopy(self._current_state["Cl_yxg"])
297
+ return cls
298
+ def get_Cl_kxg(self):
299
+ cls = {}
300
+ cls = deepcopy(self._current_state["Cl_kxg"])
301
+ return cls
302
+ def get_Cl_gxg(self):
303
+ cls = {}
304
+ cls = deepcopy(self._current_state["Cl_gxg"])
305
+ return cls
306
+ def get_Cl_muxmu(self):
307
+ cls = {}
308
+ cls = deepcopy(self._current_state["Cl_muxmu"])
309
+ return cls
310
+ def get_Cl_gxmu(self):
311
+ cls = {}
312
+ cls = deepcopy(self._current_state["Cl_gxmu"])
313
+ return cls
314
+ def get_Cl_kxmu(self):
315
+ cls = {}
316
+ cls = deepcopy(self._current_state["Cl_kxmu"])
317
+ return cls
318
+ def get_Cl_yxmu(self):
319
+ cls = {}
320
+ cls = deepcopy(self._current_state["Cl_yxmu"])
321
+ return cls
322
+ def get_Cl_kgxg(self):
323
+ cls = {}
324
+ cls = deepcopy(self._current_state["Cl_kgxg"])
325
+ return cls
326
+ def get_Cl_kgxmu(self):
327
+ cls = {}
328
+ cls = deepcopy(self._current_state["Cl_kgxmu"])
329
+ return cls
330
+ def get_Cl_IAxg(self):
331
+ cls = {}
332
+ cls = deepcopy(self._current_state["Cl_IAxg"])
333
+ return cls
334
+ def get_cl_galn_lens(self):
335
+ cls = {}
336
+ cls = deepcopy(self._current_state["cl_galn_lens"])
337
+ return cls
338
+ def get_cl_galn_galn(self):
339
+ cls = {}
340
+ cls = deepcopy(self._current_state["cl_galn_galn"])
341
+ return cls
342
+ def get_Cl_galnxtsz(self):
343
+ cls = {}
344
+ cls = deepcopy(self._current_state["Cl_galnxtsz"])
345
+ return cls
346
+ def get_Cl_galnxgallens(self):
347
+ cls = {}
348
+ cls = deepcopy(self._current_state["Cl_galnxgallens"])
349
+ return cls
350
+ def get_Cl_lensmagnxtsz(self):
351
+ cls = {}
352
+ cls = deepcopy(self._current_state["Cl_lensmagnxtsz"])
353
+ return cls
354
+ def get_Cl_lensmagnxgallens(self):
355
+ cls = {}
356
+ cls = deepcopy(self._current_state["Cl_lensmagnxgallens"])
357
+ return cls
358
+ def get_Cl_galnxIA(self):
359
+ cls = {}
360
+ cls = deepcopy(self._current_state["Cl_galnxIA"])
361
+ return cls
362
+
363
+ # get the required new observable
364
+ def get_sz_unbinned_cluster_counts(self):
365
+ cls = deepcopy(self._current_state["sz_unbinned_cluster_counts"])
366
+ # print(cls)
367
+ return cls['loglike'],cls['ntot'],cls['rates']
368
+
369
+
370
+ # get the required new observable
371
+ def get_sz_binned_cluster_counts(self):
372
+ cls = {}
373
+ cls = deepcopy(self._current_state["sz_binned_cluster_counts"])
374
+ return cls
375
+
376
+ def get_cl_cib_kappa(self):
377
+ cls = {}
378
+ cls = deepcopy(self._current_state["cl_cib_kappa"])
379
+ return cls
380
+
381
+ # IMPORTANT: this method is imported from cobaya and modified to accomodate the emulators
382
+ def calculate(self, state, want_derived=True, **params_values_dict):
383
+ # Set parameters
384
+ params_values = params_values_dict.copy()
385
+ # print('\n\n')
386
+ # print('>>> class_sz.py: class/class_sz using params:',params_values)
387
+ if 'N_ncdm' in self.extra_args.keys():
388
+ if self.extra_args['N_ncdm'] == 3:
389
+ str_mncdm = str(params_values['m_ncdm'])
390
+ params_values['m_ncdm'] = str_mncdm+','+str_mncdm+','+str_mncdm
391
+ # print('>>> class_sz.py: class/class_sz using params:',params_values)
392
+ # exit(0)
393
+ try:
394
+ params_values['ln10^{10}A_s'] = params_values.pop("logA")
395
+ self.set(params_values)
396
+ except KeyError:
397
+ self.set(params_values)
398
+ # Compute!
399
+ try:
400
+ if self.use_class_sz_fast_mode == 1:
401
+ # start = time.perf_counter()
402
+ if self.use_class_sz_no_cosmo_mode == 1:
403
+ # print(params_values)
404
+ self.classy.compute_class_sz(params_values)
405
+ else:
406
+ self.classy.compute_class_szfast()
407
+ # end = time.perf_counter()
408
+ # print('classy_szfast took:',end-start)
409
+ # self.classy.compute_class_szfast()
410
+ # elif self.use_class_sz_no_cosmo_mode == 1:
411
+ # self.classy.compute_class_sz(params_values)
412
+ else:
413
+ self.classy.compute()
414
+ # "Valid" failure of CLASS: parameters too extreme -> log and report
415
+ except self.classy_module.CosmoComputationError as e:
416
+ if self.stop_at_error:
417
+ self.log.error(
418
+ "Computation error (see traceback below)! "
419
+ "Parameters sent to CLASS: %r and %r.\n"
420
+ "To ignore this kind of error, make 'stop_at_error: False'.",
421
+ state["params"], dict(self.extra_args))
422
+ raise
423
+ else:
424
+ self.log.debug("Computation of cosmological products failed. "
425
+ "Assigning 0 likelihood and going on. "
426
+ "The output of the CLASS error was %s" % e)
427
+ return False
428
+ # CLASS not correctly initialized, or input parameters not correct
429
+ except self.classy_module.CosmoSevereError:
430
+ self.log.error("Serious error setting parameters or computing results. "
431
+ "The parameters passed were %r and %r. To see the original "
432
+ "CLASS' error traceback, make 'debug: True'.",
433
+ state["params"], self.extra_args)
434
+ raise # No LoggedError, so that CLASS traceback gets printed
435
+ # Gather products
436
+ for product, collector in self.collectors.items():
437
+ # print(product,collector)
438
+ # Special case: sigma8 needs H0, which cannot be known beforehand:
439
+ if "sigma8" in self.collectors:
440
+ self.collectors["sigma8"].args[0] = 8 / self.classy.h()
441
+ method = getattr(self.classy, collector.method)
442
+ arg_array = self.collectors[product].arg_array
443
+ if isinstance(arg_array, int):
444
+ arg_array = np.atleast_1d(arg_array)
445
+ if arg_array is None:
446
+ state[product] = method(
447
+ *self.collectors[product].args, **self.collectors[product].kwargs)
448
+ elif isinstance(arg_array, Sequence) or isinstance(arg_array, np.ndarray):
449
+ arg_array = np.array(arg_array)
450
+ if len(arg_array.shape) == 1:
451
+ # if more than one vectorised arg, assume all vectorised in parallel
452
+ n_values = len(self.collectors[product].args[arg_array[0]])
453
+ state[product] = np.zeros(n_values)
454
+ args = deepcopy(list(self.collectors[product].args))
455
+ for i in range(n_values):
456
+ for arg_arr_index in arg_array:
457
+ args[arg_arr_index] = \
458
+ self.collectors[product].args[arg_arr_index][i]
459
+ state[product][i] = method(
460
+ *args, **self.collectors[product].kwargs)
461
+ elif len(arg_array.shape) == 2:
462
+ if len(arg_array) > 2:
463
+ raise NotImplementedError("Only 2 array expanded vars so far.")
464
+ # Create outer combinations
465
+ x_and_y = np.array(np.meshgrid(
466
+ self.collectors[product].args[arg_array[0, 0]],
467
+ self.collectors[product].args[arg_array[1, 0]])).T
468
+ args = deepcopy(list(self.collectors[product].args))
469
+ result = np.empty(shape=x_and_y.shape[:2])
470
+ for i, row in enumerate(x_and_y):
471
+ for j, column_element in enumerate(x_and_y[i]):
472
+ args[arg_array[0, 0]] = column_element[0]
473
+ args[arg_array[1, 0]] = column_element[1]
474
+ result[i, j] = method(
475
+ *args, **self.collectors[product].kwargs)
476
+ state[product] = (
477
+ self.collectors[product].args[arg_array[0, 0]],
478
+ self.collectors[product].args[arg_array[1, 0]], result)
479
+ else:
480
+ raise ValueError("arg_array not correctly formatted.")
481
+ elif arg_array in self.collectors[product].kwargs:
482
+ value = np.atleast_1d(self.collectors[product].kwargs[arg_array])
483
+ state[product] = np.zeros(value.shape)
484
+ for i, v in enumerate(value):
485
+ kwargs = deepcopy(self.collectors[product].kwargs)
486
+ kwargs[arg_array] = v
487
+ state[product][i] = method(
488
+ *self.collectors[product].args, **kwargs)
489
+ else:
490
+ raise LoggedError(self.log, "Variable over which to do an array call "
491
+ f"not known: arg_array={arg_array}")
492
+ if collector.post:
493
+ state[product] = collector.post(*state[product])
494
+ # Prepare derived parameters
495
+ d, d_extra = self._get_derived_all(derived_requested=want_derived)
496
+ if want_derived:
497
+ state["derived"] = {p: d.get(p) for p in self.output_params}
498
+ # Prepare necessary extra derived parameters
499
+ state["derived_extra"] = deepcopy(d_extra)
500
+ # exit(0)
501
+
502
+
503
+ # # get the required new observable
504
+ # def get_Cl(self,ell_factor=True,units="FIRASmuK2"):
505
+ # ell_factor=self.ell_factor
506
+ # # if self.tsz.use_class_sz_fast_mode == 1:
507
+ # cls = {}
508
+ # cls['ell'] = np.arange(20000)
509
+ # # print(cls['ell'])
510
+ # cls['tt'] = np.zeros(20000)
511
+ # cls['te'] = np.zeros(20000)
512
+ # cls['ee'] = np.zeros(20000)
513
+ # cls['pp'] = np.zeros(20000)
514
+ # # if self.tt_spectra is not None:
515
+ # nl = len(self.classy.class_szfast.cp_predicted_tt_spectrum)
516
+ # # print('nl:',nl)
517
+ # cls['tt'][2:nl+2] = (2.7255e6)**2.*self.classy.class_szfast.cp_predicted_tt_spectrum
518
+ # if ell_factor==False:
519
+ # lcp = np.asarray(cls['ell'][2:nl+2])
520
+ # cls['tt'][2:nl+2] *= 1./(lcp*(lcp+1.)/2./np.pi)
521
+ #
522
+ # # if self.te_spectra is not None:
523
+ # cls['te'][2:nl+2] = (2.7255e6)**2.*self.classy.class_szfast.cp_predicted_te_spectrum
524
+ # if ell_factor==False:
525
+ # lcp = np.asarray(cls['ell'][2:nl+2])
526
+ # cls['te'][2:nl+2] *= 1./(lcp*(lcp+1.)/2./np.pi)
527
+ # # if self.ee_spectra is not None:
528
+ # cls['ee'][2:nl+2] = (2.7255e6)**2.*self.classy.class_szfast.cp_predicted_ee_spectrum
529
+ # if ell_factor==False:
530
+ # lcp = np.asarray(cls['ell'][2:nl+2])
531
+ # cls['ee'][2:nl+2] *= 1./(lcp*(lcp+1.)/2./np.pi)
532
+ # # if self.pp_spectra is not None:
533
+ # # nl = len(self.pp_spectra[0])
534
+ # if self.lensing_lkl == "SOLikeT":
535
+ # cls['pp'][2:nl+2] = self.classy.class_szfast.cp_predicted_pp_spectrum/4. ## this is clkk... works for so lensinglite lkl
536
+ # else:
537
+ # # here for the planck lensing lkl, using lfactor option gives:
538
+ # lcp = np.asarray(cls['ell'][2:nl+2])
539
+ # cls['pp'][2:nl+2] = self.classy.class_szfast.cp_predicted_pp_spectrum/(lcp*(lcp+1.))**2.
540
+ # cls['pp'][2:nl+2] *= (lcp*(lcp+1.))**2./2./np.pi
541
+ # return cls
542
+
543
+ # def check_ranges(self, z, k):
544
+ # return 1
545
+
546
+ # IMPORTANT: copied from cobaya and changed.
547
+ def get_param(self, p):
548
+ translated = self.translate_param(p)
549
+ for pool in ["params", "derived", "derived_extra"]:
550
+ value = (self.current_state[pool] or {}).get(translated, None)
551
+ if p == 'omegam':
552
+ print('getting omegam in get_param')
553
+ # print(translated)
554
+ # print(self.classy.Omega_m())
555
+ # exit(0)
556
+ return self.classy.Omega_m()
557
+ if value is not None:
558
+ return value
559
+
560
+ raise LoggedError(self.log, "Parameter not known: '%s'", p)
561
+ ### ORIGINAL FUNCTION:
562
+ # def get_param(self, p):
563
+ # translated = self.translate_param(p)
564
+ # for pool in ["params", "derived", "derived_extra"]:
565
+ # value = (self.current_state[pool] or {}).get(translated, None)
566
+ # if value is not None:
567
+ # return value
568
+ #
569
+ # raise LoggedError(self.log, "Parameter not known: '%s'", p)
570
+
571
+
572
+
573
+ @classmethod
574
+ def is_installed(cls, **kwargs):
575
+ return load_module('classy_sz')
576
+
577
+
578
+ # this just need to be there as it's used to fill-in self.collectors in must_provide:
579
+ class Collector(NamedTuple):
580
+ method: str
581
+ args: Sequence = []
582
+ args_names: Sequence = []
583
+ kwargs: dict = {}
584
+ arg_array: Union[int, Sequence] = None
585
+ post: Optional[callable] = None