Coverage for src/galsbi/galsbi.py: 98%
133 statements
« prev ^ index » next coverage.py v7.6.9, created at 2024-12-13 03:24 +0000
1# Copyright (C) 2024 ETH Zurich
2# Institute for Particle Physics and Astrophysics
3# Author: Silvan Fischbacher
4# created: Thu Aug 08 2024
6import contextlib
7import importlib
9from astropy.io import fits
10from astropy.table import Table
11from cosmic_toolbox import arraytools as at
12from cosmic_toolbox import logger
13from ufig import run_util
15from . import citations, load
17LOGGER = logger.get_logger(__name__)
20class GalSBI:
21 """
22 This class is the main interface to the model. It provides methods to generate
23 mock galaxy catalogs and to cite the model.
24 """
26 def __init__(self, name):
27 """
28 :param name: name of the model to use
29 """
30 self.name = name
31 self.mode = None
32 self.filters = None
34 def generate_catalog(
35 self,
36 mode="intrinsic",
37 config_file=None,
38 model_index=0,
39 file_name="GalSBI_sim",
40 **kwargs,
41 ):
42 """
43 Generates a mock galaxy catalog using the model and configuration specified. The
44 parameter model_index is used to select a specific set of model parameters from
45 the ABC posterior. If a list of model parameters is provided, catalogs are
46 generated for each set of parameters. The saved catalogs and images are named
47 according to the file_name and model_index.
49 Names of the files
50 ------------------
51 Intrinsic ucat galaxy catalog: f"{file_name}_{index}_{band}_ucat.gal.cat"
52 Intrinsic ucat star catalog: f"{file_name}_{index}_{band}_ucat.star.cat"
53 Output catalog: f"{file_name}_{index}_{band}_se.cat"
54 Output image: f"{file_name}_{index}_{band}_image.fits"
55 Segmentation map: f"{file_name}_{index}_{band}_se_seg.h5"
56 Background map: f"{file_name}_{index}_{band}_se_bkg.h5"
58 :param mode: mode to use for generating the catalog, either "intrinsic", "emu",
59 "image", "image+SE", "config_file"
60 :param config_file: dictionary or path to a configuration file to use for
61 generating the catalog (only used if mode="config_file")
62 :param model_index: index of the model parameters to use for generating the
63 catalog
64 :param catalog_name: filename of the catalog and images to generate
65 :param kwargs: additional keyword arguments to pass to the workflow (overwrites
66 the values from the model parameters and config file)
67 """
68 self.mode = mode
69 model_parameters = load.load_abc_posterior(self.name)
70 config = load.load_config(self.name, mode, config_file)
72 if isinstance(model_index, int):
73 model_index = [model_index]
74 for index in model_index:
75 LOGGER.info(
76 "Generating catalog for model"
77 f" {self.name} and mode {mode} with index {index}"
78 )
79 kwargs["galaxy_catalog_name_format"] = (
80 f"{file_name}_{index}_{ } { } _ucat.gal.cat"
81 )
82 kwargs["star_catalog_name_format"] = (
83 f"{file_name}_{index}_{ } { } _ucat.star.cat"
84 )
85 kwargs["sextractor_forced_photo_catalog_name_format"] = (
86 f"{file_name}_{index}_{ } { } _se.cat"
87 )
88 kwargs["image_name_format"] = f"{file_name}_{index}_{ } { } _image.fits"
89 kwargs["tile_name"] = ""
90 self.catalog_name = file_name
91 self._run(config, model_parameters[index], **kwargs)
93 __call__ = generate_catalog
95 def _run(self, config, model_parameters, **kwargs):
96 """
97 Runs the workflow with the given configuration and model parameters
99 :param config: configuration to use for generating the catalog
100 :param model_parameters: model parameters to use for generating the catalog
101 :param kwargs: additional keyword arguments to pass to the workflow (overwrites
102 the values from the model parameters and config file)
103 """
104 kargs = {}
105 for col in model_parameters.dtype.names:
106 kargs[col] = model_parameters[col]
107 if ("moffat_beta1" in model_parameters.dtype.names) and (
108 "moffat_beta2" in model_parameters.dtype.names
109 ):
110 kargs["psf_beta"] = [
111 model_parameters["moffat_beta1"][0],
112 model_parameters["moffat_beta2"][0],
113 ]
114 kargs.update(kwargs)
115 if "filters" in kargs:
116 self.filters = kargs["filters"]
117 else:
118 config_module = importlib.import_module(config)
119 self.filters = config_module.filters
121 run_util.run_ufig_from_config(config, **kargs)
123 def cite(self):
124 """
125 Prints all the papers that should be cited when using the configuration
126 specified
127 """
128 print("\033[1mPlease cite the following papers\033[0m")
129 print("=================================")
130 print("\033[1mFor using the GalSBI model:\033[0m")
131 citations.cite_galsbi_release()
132 print("\033[1mFor using the galsbi python package:\033[0m")
133 citations.cite_code_release(self.mode)
134 print("")
136 print(
137 "\033[1mFor the galaxy population model and redshift distribution:\033[0m"
138 )
139 citations.cite_abc_posterior(self.name)
140 print("")
141 print("Example:")
142 print("--------")
143 print(
144 "We use the GalSBI framework (PAPERS GalSBI release) to generate mock"
145 " galaxy catalogs. The galaxy population model corresponds to the"
146 " posterior from (PAPER model). (...) "
147 "Acknowledgements: We acknowledge the use of the following software:"
148 "(numpy), (scipy), (PAPERS code release), (...)"
149 )
151 def load_catalogs(self, output_format="rec", model_index=0, combine=False):
152 """
153 Loads the catalogs generated by the model.
155 :param output_format: format of the output, either "rec", "df" or "fits"
156 :param model_index: index of the model parameters to use for loading the
157 catalogs
158 :param combine: if True, combines the catalogs from all bands into a single
159 catalog
160 :return: catalogs in the specified format
161 """
162 if self.filters is None:
163 raise RuntimeError("please generate catalogs first")
165 if output_format == "rec":
166 convert = lambda x: x # noqa: E731
167 elif output_format == "df":
168 convert = at.rec2pd
169 elif output_format == "fits":
170 convert = Table
171 else:
172 raise ValueError(f"Unknown output format {output_format}")
174 output = {}
175 for band in self.filters:
176 catalog_name = f"{self.catalog_name}_{model_index}_{band}_ucat.gal.cat"
177 with contextlib.suppress(FileNotFoundError):
178 output[f"ucat galaxies {band}"] = at.load_hdf(catalog_name)
179 catalog_name = f"{self.catalog_name}_{model_index}_{band}_ucat.star.cat"
180 with contextlib.suppress(FileNotFoundError):
181 output[f"ucat stars {band}"] = at.load_hdf(catalog_name)
182 catalog_name = f"{self.catalog_name}_{model_index}_{band}_se.cat"
183 with contextlib.suppress(FileNotFoundError):
184 output[f"sextractor {band}"] = at.load_hdf(catalog_name)
186 if not combine:
187 catalogs = {key: convert(value) for key, value in output.items()}
188 return catalogs
190 combined_catalogs = self._build_combined_catalogs(output)
191 return {key: convert(value) for key, value in combined_catalogs.items()}
193 def load_images(self, model_index=0):
194 """
195 Loads the images generated by the model. This include the actual image,
196 the segmentation map and the background map.
198 param model_index: index of the model parameters to use for loading the images
199 return: images as numpy arrays
200 """
201 output = {}
202 for band in self.filters:
203 image_name = f"{self.catalog_name}_{model_index}_{band}_image.fits"
204 try:
205 hdul = fits.open(image_name)
206 image = hdul[0].data
207 hdul.close()
208 output[f"image {band}"] = image
209 except FileNotFoundError:
210 pass
211 segmap_name = f"{self.catalog_name}_{model_index}_{band}_se_seg.h5"
212 with contextlib.suppress(FileNotFoundError):
213 output[f"segmentation {band}"] = at.load_hdf_cols(segmap_name)[
214 "SEGMENTATION"
215 ]
216 bkgmap_name = f"{self.catalog_name}_{model_index}_{band}_se_bkg.h5"
217 with contextlib.suppress(FileNotFoundError):
218 output[f"background {band}"] = at.load_hdf_cols(bkgmap_name)[
219 "BACKGROUND"
220 ]
221 return output
223 def _build_combined_catalogs(self, catalogs):
224 band_dep_params = ["int_mag", "mag", "abs_mag", "bkg_noise_amp"]
225 combined_catalogs = {}
226 filter = self.filters[0]
227 if f"ucat galaxies {filter}" in catalogs:
228 new_cat = {}
229 for f in self.filters:
230 cat = catalogs[f"ucat galaxies {f}"]
231 for par in cat.dtype.names:
232 if par not in band_dep_params:
233 new_cat[par] = cat[par]
234 else:
235 new_cat[f"{par} {f}"] = cat[par]
236 combined_catalogs["ucat galaxies"] = at.dict2rec(new_cat)
237 if f"ucat stars {filter}" in catalogs:
238 new_cat = {}
239 for f in self.filters:
240 cat = catalogs[f"ucat stars {f}"]
241 for par in cat.dtype.names:
242 if par not in band_dep_params:
243 new_cat[par] = cat[par]
244 else:
245 new_cat[f"{par} {f}"] = cat[par]
246 combined_catalogs["ucat stars"] = at.dict2rec(new_cat)
247 if f"sextractor {filter}" in catalogs:
248 band_ind_params = [
249 "dec",
250 "ra",
251 "z",
252 "e1",
253 "e2",
254 "r50",
255 "sersic_n",
256 "galaxy_type",
257 "id",
258 "x",
259 "y",
260 ]
261 new_cat = {}
262 for f in self.filters:
263 cat = catalogs[f"sextractor {f}"]
264 for par in cat.dtype.names:
265 if par in band_ind_params:
266 new_cat[par] = cat[par]
267 else:
268 new_cat[f"{par} {f}"] = cat[par]
269 combined_catalogs["sextractor"] = at.dict2rec(new_cat)
270 return combined_catalogs