Coverage for tests/test_estimate_photoz_task_trainz.py: 100%

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# (http://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 <http://www.gnu.org/licenses/>. 

21 

22import os 

23 

24import astropy.units as u 

25import numpy as np 

26from astropy.table import Table 

27from rail.core.model import Model as PhotoZModel 

28 

29from lsst.meas.photoz.base.estimate_photoz_task_trainz import ( 

30 EstimatePhotozTrainZConfig, 

31 EstimatePhotozTrainZTask, 

32) 

33 

34TESTDIR = os.path.abspath(os.path.dirname(__file__)) 

35 

36 

37def test_algo_config(): 

38 """Test default initialization of TrainZ config class.""" 

39 config = EstimatePhotozTrainZConfig() 

40 config.validate() 

41 

42 

43def test_run_trainz() -> None: 

44 """Test that trainz runs on plausible data.""" 

45 path = os.path.join(TESTDIR, "data", "model_inform_train_z_wrap.pickle") 

46 model = PhotoZModel.read(path) 

47 config = EstimatePhotozTrainZConfig() 

48 task = EstimatePhotozTrainZTask(config=config, initInputs={}) 

49 

50 algo_config = task.photoz_algo.config 

51 n_data = 173 

52 algo_config.chunk_size = int(round(n_data / 3.1)) 

53 config.freeze() 

54 mags = np.linspace(19, 26, n_data) 

55 mag_names = algo_config.get_mag_names() 

56 data = {column: mags + 0.5 * idx for idx, column in enumerate(mag_names.values())} 

57 for band, column in algo_config.get_flux_names().items(): 

58 data[column] = (data[mag_names[band]] * u.ABmag).to(u.nJy).value 

59 

60 data["ebv"] = 0.05 * (1 + np.sin(np.arange(n_data))) 

61 data = Table(data) 

62 

63 result = task.run(photoz_model=model, fluxes=data) 

64 yvals = result.photoz_ensemble.objdata["yvals"] 

65 assert yvals.shape[0] == n_data 

66 np.testing.assert_array_compare(np.greater, np.sum(yvals, axis=1), 0)