Coverage for python/lsst/meas/photoz/base/estimate_photoz_task.py: 80%

1# This file is part of meas_photoz_base. 

2# 

3# Developed for the LSST Data Management System. 

4# This product includes software developed by the LSST Project 

5# (https://www.lsst.org). 

6# See the COPYRIGHT file at the top-level directory of this distribution 

7# for details of code ownership. 

8# 

9# This program is free software: you can redistribute it and/or modify 

10# it under the terms of the GNU General Public License as published by 

11# the Free Software Foundation, either version 3 of the License, or 

12# (at your option) any later version. 

13# 

14# This program is distributed in the hope that it will be useful, 

15# but WITHOUT ANY WARRANTY; without even the implied warranty of 

16# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

17# GNU General Public License for more details. 

18# 

19# You should have received a copy of the GNU General Public License 

20# along with this program. If not, see <https://www.gnu.org/licenses/>. 

21 

22from __future__ import annotations 

23 

24__all__ = [ 

25 "EstimatePhotozAlgoConfigBase", 

26 "EstimatePhotozAlgoTask", 

27 "EstimatePhotozTask", 

28 "EstimatePhotozTaskConfig", 

29 "photozAlgoRegistry", 

30] 

31 

32from abc import ABC, abstractmethod 

33from typing import Any 

34 

35import numpy as np 

36import qp 

37from astropy.table import Table 

38from ceci.config import StageConfig as CeciStageConfig 

39from ceci.config import StageParameter as CeciParam 

40from rail.core.model import Model 

41from rail.estimation.estimator import CatEstimator 

42from rail.interfaces import PZFactory 

43 

44import lsst.pex.config as pexConfig 

45import lsst.pipe.base.connectionTypes as cT 

46from lsst.pipe.base import ( 

47 PipelineTask, 

48 PipelineTaskConfig, 

49 PipelineTaskConnections, 

50 Struct, 

51 Task, 

52) 

53 

54from .utils import chunk_iterator 

55 

56 

57class EstimatePhotozConnections( 

58 PipelineTaskConnections, 

59 dimensions=("skymap", "tract"), 

60 defaultTemplates={"algo": "trainz"}, 

61): 

62 """Connections for tasks that make p(z) estimates. 

63 

64 These will take pickled model file as a "calibration-like" input, 

65 an object table as input, and create a p(z) file in 'qp' format. 

66 """ 

67 

68 photoz_model = cT.PrerequisiteInput( 

69 doc="Model for PZ Estimation", 

70 name="photoz_model_{algo}", 

71 storageClass="PhotozModel", 

72 dimensions=["instrument"], 

73 isCalibration=True, 

74 ) 

75 objects = cT.Input( 

76 doc="Object table", 

77 name="object", 

78 storageClass="ArrowAstropy", 

79 dimensions=("skymap", "tract"), 

80 deferLoad=True, 

81 ) 

82 

83 photoz_ensemble = cT.Output( 

84 doc="Per-object p(z) estimates", 

85 name="photoz_ensemble_{algo}", 

86 storageClass="QPEnsemble", 

87 dimensions=("skymap", "tract"), 

88 ) 

89 

90 

91class EstimatePhotozAlgoConfigBase( 

92 pexConfig.Config, 

93): 

94 """Base class for configurations of algorithm-specific p(z) 

95 estimation tasks. 

96 

97 This class mostly just translates the RAIL configuration 

98 parameters to pex.config parameters. 

99 

100 Subclasses will just have to set `estimator_class` and `stage_name` 

101 and invoke `_make_fields` once in the module. 

102 """ 

103 

104 @classmethod 

105 @abstractmethod 

106 def estimator_class(cls) -> type[CatEstimator]: 

107 """Return the type of the estimator's RAIL class.""" 

108 raise NotImplementedError("Subclasses must specify an estimator class") 

109 

110 # This should be a property but py3.13+ don't allow it 

111 @classmethod 

112 @abstractmethod 

113 def stage_name(cls) -> str: 

114 """Return the RAIL stage name for the estimator.""" 

115 raise NotImplementedError("Subclasses must define a RAIL stage name") 

116 

117 # Extinction coefficients; see https://ui.adsabs.harvard.edu/abs/1989ApJ...345..245C/abstract 

118 # Also in rail.utils.catalog_utils.RubinCatalogConfig.a_env 

119 default_a_env_values = dict( 

120 u=4.81, 

121 g=3.64, 

122 r=2.70, 

123 i=2.06, 

124 z=1.58, 

125 y=1.31, 

126 ) 

127 

128 # These appear in many DESC repos, for example: 

129 # https://github.com/LSSTDESC/TXPipe/blob/00ebe7476fd5d9529f5bbc4d73fcef0629d134c7/examples/dp0.2/config.yml#L47 

130 # They seem to be 10y WFD limits. Origin unclear. 

131 default_mag_limit_10y_values = dict( 

132 u=27.79, 

133 g=29.04, 

134 r=29.06, 

135 i=28.62, 

136 z=27.98, 

137 y=27.05, 

138 ) 

139 

140 # These appear to be from Roman-Rubin simulations: 

141 # https://github.com/LSSTDESC/rail_base/blob/v1.2.1/src/rail/utils/catalog_utils.py#L207 

142 # Presumably max 5y depth, and more useful for now 

143 default_mag_limit_values = dict( 

144 u=24.0, 

145 g=27.66, 

146 r=27.25, 

147 i=26.6, 

148 z=26.24, 

149 y=25.35, 

150 ) 

151 

152 def get_band_a_env_dict(self): 

153 """Return the set of a_envs to use.""" 

154 return {band_: self.default_a_env_values[band_] for band_ in self.bands_to_convert} 

155 

156 def get_mag_lim_dict(self): 

157 """Return the set of maglims to use.""" 

158 return { 

159 self.mag_template.format(band=band): self.default_mag_limit_values[band] 

160 for band in self.bands_to_convert 

161 } 

162 

163 def get_flux_names(self) -> dict[str, str]: 

164 """Return a dict mapping band to flux column name.""" 

165 return {band: self.flux_column_template.format(band=band) for band in self.bands_to_convert} 

166 

167 def get_flux_err_names(self) -> dict[str, str]: 

168 """Return a dict mapping band to flux error column name.""" 

169 return {band: self.flux_err_column_template.format(band=band) for band in self.bands_to_convert} 

170 

171 def get_mag_names(self) -> dict[str, str]: 

172 """Return a dict mapping band to mag column name.""" 

173 return {band: self.mag_template.format(band=band) for band in self.bands_to_convert} 

174 

175 def get_mag_err_names(self) -> dict[str, str]: 

176 """Return a dict mapping band to mag error column name.""" 

177 return {band: self.mag_err_template.format(band=band) for band in self.bands_to_convert} 

178 

179 mag_offset = pexConfig.Field(doc="Magnitude offset", dtype=float, default=31.4) 

180 deredden = pexConfig.Field[bool]( 

181 doc="Apply dereddening", 

182 default=True, 

183 ) 

184 band_ref = pexConfig.Field[str]( 

185 doc="Name of the most reliable reference band, if needed", 

186 default="i", 

187 ) 

188 bands_to_convert = pexConfig.ListField[str]( 

189 doc="Names of bands to convert fluxs to mags for RAIL", 

190 default=["u", "g", "r", "i", "z", "y"], 

191 ) 

192 flux_column_template = pexConfig.Field[str]( 

193 doc="Template for flux column names", 

194 default="{band}_gaap1p0Flux", 

195 # default="{band}_cModelFlux", 

196 ) 

197 flux_err_column_template = pexConfig.Field[str]( 

198 doc="Template for flux error column names", 

199 default="{band}_gaap1p0FluxErr", 

200 # default="{band}_cModelFluxErr", 

201 ) 

202 mag_template = pexConfig.Field[str]( 

203 doc="Template for magnitude names", 

204 default="{band}_gaap1p0Mag", 

205 # default="{band}_cModelMag", 

206 ) 

207 mag_err_template = pexConfig.Field[str]( 

208 doc="Template for magntitude error names", 

209 default="{band}_gaap1p0MagErr", 

210 # default="{band}_cModelMagErr", 

211 ) 

212 nondetect_val = pexConfig.Field[float]( 

213 doc="Magnitude to set for non-detections", 

214 default=np.nan, 

215 ) 

216 band_a_env = pexConfig.DictField[str, float]( 

217 doc="Reddening parameters", 

218 default=default_a_env_values, 

219 ) 

220 

221 def freeze(self): 

222 if not self._frozen: 222 ↛ 224line 222 didn't jump to line 224 because the condition on line 222 was always true

223 self._finalize() 

224 super().freeze() 

225 

226 def _finalize(self): 

227 # These calls will fail if it's already frozen. 

228 if hasattr(self, "ref_band"): 

229 self.ref_band = self.mag_template.format(band=self.band_ref) 

230 if hasattr(self, "bands"): 

231 # This is a list of mag columns in RAIL, not bands 

232 self.bands = list(self.get_mag_names().values()) 

233 if hasattr(self, "err_bands"): 

234 self.err_bands = list(self.get_mag_err_names().values()) 

235 if hasattr(self, "mag_limits"): 

236 self.mag_limits = self.get_mag_lim_dict() 

237 if hasattr(self, "band_a_env"): 237 ↛ exitline 237 didn't return from function '_finalize' because the condition on line 237 was always true

238 self.band_a_env = self.get_band_a_env_dict() 

239 

240 @classmethod 

241 def _make_fields(cls) -> None: 

242 """Import the RAIL estimation stage. 

243 

244 This method loops through the stage config parameters and converts 

245 RAIL/Ceci parameters to corresponding pex.config parameters. 

246 

247 It should be called exactly once, immediately after the definition 

248 of every subclass of this base class. 

249 """ 

250 if hasattr(cls, "__fields_made__"): 250 ↛ 251line 250 didn't jump to line 251 because the condition on line 250 was never true

251 if cls.__fields_made__ is not True: 

252 raise RuntimeError(f"{cls.__fields_made__=} exists but is not True") 

253 raise RuntimeError(f"{cls=} called _make_fields twice") 

254 stage_class = cls.estimator_class() 

255 for key, val in stage_class.config_options.items(): 

256 if isinstance(val, CeciStageConfig): 

257 val = val.get(key) 

258 if isinstance(val, CeciParam): 258 ↛ 255line 258 didn't jump to line 255 because the condition on line 258 was always true

259 if val.dtype in [bool, int, float, str]: 

260 if (attr := getattr(cls, key, None)) is not None: 

261 if not isinstance(attr, pexConfig.Field): 261 ↛ 262line 261 didn't jump to line 262 because the condition on line 261 was never true

262 raise RuntimeError(f"{cls=} {key=} exists but is of {type(key)=}, not Field") 

263 elif attr.dtype != val.dtype: 263 ↛ 264line 263 didn't jump to line 264 because the condition on line 263 was never true

264 raise RuntimeError(f"{cls=} {key=} exists but {attr.dtype=} != {val.dtype=}") 

265 attr.default = val.default 

266 attr.doc = f"{val.msg} (overriding base doc='{attr.doc}')" 

267 else: 

268 setattr( 

269 cls, 

270 key, 

271 pexConfig.Field(doc=val.msg, dtype=val.dtype, default=val.default), 

272 ) 

273 elif val.dtype in [list]: 

274 # this is a hack, but it works. 

275 if val.default: 

276 item_type = type(val.default[0]) 

277 else: 

278 item_type = str 

279 setattr( 

280 cls, 

281 key, 

282 pexConfig.ListField(doc=val.msg, dtype=item_type, default=val.default), 

283 ) 

284 elif val.dtype in [dict]: 284 ↛ 255line 284 didn't jump to line 255 because the condition on line 284 was always true

285 setattr( 

286 cls, 

287 key, 

288 pexConfig.DictField(doc=val.msg, keytype=str, default=val.default), 

289 ) 

290 cls.__fields_made__ = True 

291 

292 

293photozAlgoRegistry = pexConfig.makeRegistry( 

294 doc="A registry of photometric redshift estimation algorithm subtasks", 

295) 

296 

297 

298class EstimatePhotozAlgoTask(Task, ABC): 

299 """Task for algorithm-specific p(z) estimation. 

300 

301 This provides almost all of the functionality 

302 needed to run RAIL p(z) algorithms. 

303 

304 Parameters 

305 ---------- 

306 **kwargs 

307 Additional keyword arguments to pass to super().__init__. 

308 """ 

309 

310 ConfigClass = EstimatePhotozAlgoConfigBase 

311 

312 mag_conv = np.log(10) * 0.4 

313 

314 def __init__(self, **kwargs: Any): 

315 super().__init__(**kwargs) 

316 

317 @staticmethod 

318 def _flux_to_mag( 

319 flux_vals: np.ndarray, 

320 mag_offset: float, 

321 nondetect_val: float, 

322 ) -> np.ndarray: 

323 """Convert flux to magnitude. 

324 

325 Parameters 

326 ---------- 

327 flux_vals : np.array 

328 Input flux values (units?) 

329 

330 mag_offset : float 

331 Magnitude offset (corresponding to a flux of 1.) 

332 

333 nondetect_val : float 

334 Value to set for non-detections 

335 

336 Returns 

337 ------- 

338 mags : np.array 

339 Magnitude values 

340 """ 

341 vals = np.empty_like(flux_vals) 

342 positive = flux_vals >= 0 

343 vals[positive] = -2.5 * np.log10(flux_vals[positive]) + mag_offset 

344 vals[~positive] = nondetect_val 

345 return vals 

346 

347 @staticmethod 

348 def _flux_err_to_mag_err( 

349 flux_vals: np.ndarray, 

350 flux_err_vals: np.ndarray, 

351 mag_conv: float, 

352 nondetect_val: float = np.nan, 

353 ) -> np.ndarray: 

354 """Config flux error to magnitude error. 

355 

356 Parameters 

357 ---------- 

358 flux_vals : np.array 

359 Input flux values (units?) 

360 

361 flux_err_vals : np.array 

362 Input flux errors (units?) 

363 

364 mag_conv : float 

365 Magnitude to flux conversion (typically np.log(10)*0.4) 

366 

367 nondetect_val : float 

368 Value to set for non-detections 

369 

370 Returns 

371 ------- 

372 mags_errs : np.array 

373 Magnitude errors 

374 """ 

375 vals = np.empty_like(flux_vals) 

376 positive = flux_vals >= 0 

377 vals[positive] = flux_err_vals[positive] / (flux_vals[positive] * mag_conv) 

378 vals[~positive] = nondetect_val 

379 return vals 

380 

381 @staticmethod 

382 def _deredden_mags( 

383 data: dict[str, np.ndarray], 

384 a_env_dict: dict[str, float], 

385 mag_names: dict[str, str], 

386 nondetect_val: float, 

387 ) -> dict[str, np.ndarray]: 

388 """Deredden the magnitdues 

389 

390 Parameters 

391 ---------- 

392 data: dict[str, np.array] 

393 Input data 

394 

395 a_env_dict: dict[str, float], 

396 Redenning parameters for bands 

397 

398 mag_names: dict[str, str] 

399 Mapping from bands to magnitudes 

400 

401 nondetect_val : float 

402 Value to set for non-detections 

403 

404 Returns 

405 ------- 

406 mags: dict[str, np.array] 

407 Udpated dict with dereddened mags 

408 """ 

409 ebv = data["ebv"] 

410 for band_, a_env_ in a_env_dict.items(): 

411 mag_name = mag_names[band_] 

412 raw_mag = data[mag_name] 

413 dered_mag = np.where( 

414 np.isfinite(raw_mag), 

415 raw_mag - ebv * a_env_, 

416 nondetect_val, 

417 ) 

418 data[mag_name] = dered_mag 

419 return data 

420 

421 def _get_mags_and_errs( 

422 self, 

423 fluxes: Table, 

424 mag_offset: float, 

425 ) -> dict[str, np.ndarray]: 

426 """Fill and return a numpy dict with mags and mag errors. 

427 

428 Parameters 

429 ---------- 

430 fluxes : Table 

431 Input fluxes and flux errors 

432 

433 mag_offset : float 

434 Magnitude offset (corresponding to a flux of 1.) 

435 

436 Returns 

437 ------- 

438 mags: dict[str, np.array] 

439 Numpy dict with mags and mag errors 

440 """ 

441 # get all the column names we will use 

442 flux_names = self.config.get_flux_names() 

443 mag_names = self.config.get_mag_names() 

444 flux_err_names = self.config.get_flux_err_names() 

445 mag_err_names = self.config.get_mag_err_names() 

446 nondetect_val = self.config.nondetect_val 

447 # output dict 

448 mag_dict = {} 

449 if object_id_name := self.config.id_col: 

450 mag_dict[object_id_name] = fluxes[object_id_name] 

451 

452 # loop over bands, make mags and mag errors and fill dict 

453 for band in flux_names.keys(): 

454 fluxVals = np.asarray(fluxes[flux_names[band]]) 

455 fluxErrVals = np.asarray(fluxes[flux_err_names[band]]) 

456 mag_dict[mag_names[band]] = self._flux_to_mag( 

457 fluxVals, 

458 mag_offset, 

459 nondetect_val, 

460 ) 

461 if flux_err_names: 

462 mag_dict[mag_err_names[band]] = self._flux_err_to_mag_err( 

463 fluxVals, 

464 fluxErrVals, 

465 self.mag_conv, 

466 nondetect_val, 

467 ) 

468 

469 # return the dict with the mags 

470 return mag_dict 

471 

472 def init( 

473 self, 

474 photoz_model: Model, 

475 ) -> None: 

476 """Set up the RAIL stage to compute photo-zs. 

477 

478 Parameters 

479 ---------- 

480 photoz_model : Model 

481 Model used by the p(z) estimation algorithm. 

482 """ 

483 # pop the pipeline task config options 

484 # so that we can pass the rest to RAIL 

485 rail_kwargs = self.config.toDict().copy() 

486 for key in ["saveLogOutput", "stage_name", "mag_offset", "connections"]: 

487 rail_kwargs.pop(key, None) 

488 rail_kwargs["output_mode"] = "return" 

489 

490 # Build the RAIL stage 

491 self._stage = PZFactory.build_stage_instance( 

492 self.config.stage_name(), 

493 self.config.estimator_class(), 

494 model_path=photoz_model.data, 

495 input_path="dummy.in", 

496 **rail_kwargs, 

497 ) 

498 self._stage._initialize_run() 

499 

500 def col_names( 

501 self, 

502 ) -> list[str]: 

503 """Get the list of column names to read from the input data.""" 

504 columns = list(self.config.get_flux_names().values()) + list( 

505 self.config.get_flux_err_names().values() 

506 ) 

507 if self.config.deredden: 

508 columns += ["ebv"] 

509 

510 return columns 

511 

512 def run( 

513 self, 

514 fluxes: Table, 

515 ) -> Struct: 

516 """Run a p(z) estimation algorithm. 

517 

518 Parameters 

519 ---------- 

520 fluxes : Table 

521 Fluxes used to compute the redshifts. 

522 

523 Returns 

524 ------- 

525 photoz_pdfs : qp.Ensemble 

526 Object with the p(z) PDFs. 

527 """ 

528 n_obj = len(fluxes) 

529 # Convert fluxes to mags 

530 mags = self._get_mags_and_errs(fluxes, self.config.mag_offset) 

531 nondetect_val = self.config.nondetect_val if hasattr(self.config, "nondetect_val") else np.nan 

532 

533 # De-redden 

534 if self.config.deredden: 534 ↛ 544line 534 didn't jump to line 544 because the condition on line 534 was always true

535 # asarray will convert an astropy column to an array w/o units 

536 mags["ebv"] = np.asarray(fluxes["ebv"]) 

537 mags = self._deredden_mags( 

538 mags, 

539 self.config.band_a_env, 

540 self.config.get_mag_names(), 

541 nondetect_val, 

542 ) 

543 

544 pz_ensembles: list[qp.Ensemble] = [] 

545 

546 # split processing into chunks to avoid memory issues in rail 

547 iterator = chunk_iterator(self.config.chunk_size, mags, n_obj) 

548 for a_chunk, chunk_size in iterator: 

549 # Pass the mags to RAIL and get back the p(z) pdfs 

550 # as a qp.Ensemble object 

551 pz_ensembles.append(PZFactory.estimate_single_pz(self._stage, a_chunk, chunk_size)) 

552 

553 photoz_pdfs = qp.concatenate(pz_ensembles) 

554 return Struct(photoz_ensemble=photoz_pdfs) 

555 

556 

557class EstimatePhotozTaskConfig(PipelineTaskConfig, pipelineConnections=EstimatePhotozConnections): 

558 """Configuration for EstimatePhotozTask PipelineTask.""" 

559 

560 photoz_algo = photozAlgoRegistry.makeField( 

561 doc="Algorithm specific configuration p(z) estimation task", 

562 ) 

563 

564 

565class EstimatePhotozTask(PipelineTask): 

566 """PipelineTask for p(z) estimation. 

567 

568 Parameters 

569 ---------- 

570 initInputs 

571 Initialization inputs to pass to super().__init__. 

572 **kwargs 

573 Additional keyword arguments to pass to super().__init__. 

574 """ 

575 

576 ConfigClass = EstimatePhotozTaskConfig 

577 _DefaultName = "estimatePhotoz" 

578 

579 def __init__(self, initInputs: dict, **kwargs): 

580 super().__init__(initInputs=initInputs, **kwargs) 

581 self._initialized = False 

582 self.makeSubtask("photoz_algo") 

583 

584 def runQuantum(self, butlerQC, inputRefs, outputRefs): 

585 inputs = butlerQC.get(inputRefs) 

586 inputs["fluxes"] = inputs.pop("objects").get( 

587 parameters=dict(columns=self.photoz_algo.col_names()), 

588 ) 

589 outputs = self.run(**inputs, skip_init=self._initialized) 

590 butlerQC.put(outputs, outputRefs) 

591 

592 def run( 

593 self, 

594 *, 

595 photoz_model: Model, 

596 fluxes: Table, 

597 skip_init: bool = False, 

598 ) -> Struct: 

599 if not skip_init: 599 ↛ 603line 599 didn't jump to line 603 because the condition on line 599 was always true

600 self._initialized = True 

601 self.photoz_algo.init(photoz_model) 

602 

603 ret_struct = self.photoz_algo.run(fluxes) 

604 return Struct(photoz_ensemble=ret_struct.photoz_ensemble)