nikhil-sarin/redback | Build 13137940639 | test/model_test.py | Coveralls

1

import unittest

1✔

2

from os import listdir

1✔

3

from os.path import dirname

1✔

4

from unittest import mock

1✔

6

import bilby

1✔

7

import numpy as np

1✔

9

import redback.model_library

1✔

11

_dirname = dirname(__file__)

1✔

14

class TestModels(unittest.TestCase):

1✔

16

    def setUp(self) -> None:

1✔

17

        self.path_to_files = f"{_dirname}/../redback/priors/"

1✔

18

        self.prior_files = listdir(self.path_to_files)

1✔

20

    def tearDown(self) -> None:

1✔

21

        pass

1✔

23

    def get_prior(self, file):

1✔

24

        prior_dict = bilby.prior.PriorDict()

1✔

25

        prior_dict.from_file(f"{self.path_to_files}{file}")

1✔

26

        return prior_dict

1✔

28

    def test_models(self):

1✔

29

        kwargs = dict(frequency=2e14)

1✔

30

        times = np.array([1, 2, 3])

1✔

31

        for f in self.prior_files:

1✔

32

            print(f)

1✔

33

            model_name = f.replace(".prior", "")

1✔

34

            if model_name == 'trapped_magnetar':

1✔

35

                kwargs['output_format'] = 'luminosity'

1✔

36

            elif model_name in ['fitted', 'jetsimpy_tophat', 'jetsimpy_gaussian', 'jetsimpy_powerlaw']:

1✔

39

                kwargs['output_format'] = 'flux_density'

1✔

40

            prior = self.get_prior(file=f)

1✔

41

            sample = prior.sample()

1✔

42

            if model_name == 'polytrope_eos_two_component_bns':

1✔

43

                sample['gamma_1'] = 4.04

1✔

44

                sample['gamma_2'] = 2.159

1✔

45

                sample['gamma_3'] = 3.688

1✔

46

                sample['log_p'] = 33.72

1✔

47

            function = redback.model_library.all_models_dict[model_name]

1✔

48

            ys = function(times, **sample, **kwargs)

1✔

49

            self.assertEqual(len(times), len(ys))

1✔

51

class TestPhaseModels(unittest.TestCase):

1✔

52

    def setUp(self) -> None:

1✔

53

        self.path_to_files = f"{_dirname}/../redback/priors/"

1✔

54

        self.prior_files = listdir(self.path_to_files)

1✔

56

    def tearDown(self) -> None:

1✔

57

        pass

1✔

59

    def get_prior(self, file):

1✔

60

        prior_dict = bilby.prior.PriorDict()

1✔

61

        prior_dict.from_file(f"{self.path_to_files}{file}")

1✔

62

        return prior_dict

1✔

64

    def test_models(self):

1✔

65

        kwargs = dict(frequency=2e14)

1✔

66

        times = np.array([1, 2, 3]) + 55855

1✔

67

        for f in self.prior_files:

1✔

68

            print(f)

1✔

69

            model_name = f.replace(".prior", "")

1✔

70

            if model_name == 'trapped_magnetar':

1✔

71

                kwargs['output_format'] = 'luminosity'

1✔

72

            elif model_name in ['fitted', 'jetsimpy_tophat', 'jetsimpy_gaussian', 'jetsimpy_powerlaw']:

1✔

75

                kwargs['output_format'] = 'flux_density'

1✔

76

            prior = self.get_prior(file=f)

1✔

77

            function = redback.model_library.all_models_dict['t0_base_model']

1✔

78

            kwargs['base_model'] = model_name

1✔

79

            kwargs['t0'] = 55855

1✔

80

            sample = prior.sample()

1✔

81

            if model_name == 'polytrope_eos_two_component_bns':

1✔

82

                sample['gamma_1'] = 4.04

1✔

83

                sample['gamma_2'] = 2.159

1✔

84

                sample['gamma_3'] = 3.688

1✔

85

                sample['log_p'] = 33.72

1✔

86

            ys = function(times, **sample, **kwargs)

1✔

87

            self.assertEqual(len(times), len(ys))

1✔

89

class TestMagnitudeOutput(unittest.TestCase):

1✔

90

    def setUp(self) -> None:

1✔

91

        self.path_to_files = f"{_dirname}/../redback/priors/"

1✔

92

        self.prior_files = listdir(self.path_to_files)

1✔

94

    def tearDown(self) -> None:

1✔

95

        pass

1✔

97

    def get_prior(self, file):

1✔

98

        prior_dict = bilby.prior.PriorDict()

1✔

99

        prior_dict.from_file(f"{self.path_to_files}{file}")

1✔

100

        return prior_dict

1✔

102

    def test_models(self):

1✔

103

        kwargs = dict(frequency=2e14, bands='ztfg')

1✔

104

        times = np.array([1, 2, 3])

1✔

105

        for f in self.prior_files:

1✔

106

            print(f)

1✔

107

            model_name = f.replace(".prior", "")

1✔

108

            if model_name == 'trapped_magnetar':

1✔

109

                kwargs['output_format'] = 'luminosity'

1✔

110

            elif model_name in ['tophat_and_twocomponent', 'tophat_and_twolayerstratified',

1✔

113

                pass

1✔

115

                kwargs['output_format'] = 'magnitude'

1✔

116

            prior = self.get_prior(file=f)

1✔

117

            sample = prior.sample()

1✔

118

            if model_name == 'polytrope_eos_two_component_bns':

1✔

119

                sample['gamma_1'] = 4.04

1✔

120

                sample['gamma_2'] = 2.159

1✔

121

                sample['gamma_3'] = 3.688

1✔

122

                sample['log_p'] = 33.72

1✔

123

            function = redback.model_library.all_models_dict[model_name]

1✔

124

            ys = function(times, **sample, **kwargs)

1✔

125

            self.assertEqual(len(times), len(ys))

1✔

127

class TestFluxOutput(unittest.TestCase):

1✔

128

    def setUp(self) -> None:

1✔

129

        self.path_to_files = f"{_dirname}/../redback/priors/"

1✔

130

        self.prior_files = listdir(self.path_to_files)

1✔

132

    def tearDown(self) -> None:

1✔

133

        pass

1✔

135

    def get_prior(self, file):

1✔

136

        prior_dict = bilby.prior.PriorDict()

1✔

137

        prior_dict.from_file(f"{self.path_to_files}{file}")

1✔

138

        return prior_dict

1✔

140

    def test_models(self):

1✔

141

        kwargs = dict(frequency=2e14, bands='ztfg')

1✔

142

        times = np.array([1, 2, 3])

1✔

143

        valid_models = ['arnett', 'one_component_kilonova_model', 'slsn',

1✔

145

        for f in self.prior_files:

1✔

146

            model_name = f.replace(".prior", "")

1✔

147

            if model_name in valid_models:

1✔

148

                print(f)

1✔

149

                kwargs['output_format'] = 'flux'

1✔

150

                prior = self.get_prior(file=f)

1✔

151

                sample = prior.sample()

1✔

152

                if model_name == 'polytrope_eos_two_component_bns':

1✔

157

                function = redback.model_library.all_models_dict[model_name]

1✔

158

                ys = function(times, **sample, **kwargs)

1✔

160

                ys = np.ones(len(times))

1✔

161

            self.assertEqual(len(times), len(ys))

1✔

164

class TestIntegratedFluxModelFlux(unittest.TestCase):

1✔

165

    def setUp(self) -> None:

1✔

166

        self.path_to_files = f"{_dirname}/../redback/priors/"

1✔

167

        self.prior_files = listdir(self.path_to_files)

1✔

169

    def tearDown(self) -> None:

1✔

170

        pass

1✔

172

    def get_prior(self, file):

1✔

173

        prior_dict = bilby.prior.PriorDict()

1✔

174

        prior_dict.from_file(f"{self.path_to_files}{file}")

1✔

175

        return prior_dict

1✔

177

    def test_models(self):

1✔

178

        kwargs = dict(frequency=np.array([2e14,6e14]), bands='ztfg')

1✔

179

        times = np.array([1, 2, 3])

1✔

180

        valid_models = ['gaussiancore']

1✔

181

        for f in self.prior_files:

1✔

182

            model_name = f.replace(".prior", "")

1✔

183

            if model_name in valid_models:

1✔

184

                print(f)

1✔

185

                prior = self.get_prior(file=f)

1✔

186

                kwargs['base_model'] = model_name

1✔

187

                sample = prior.sample()

1✔

188

                function = redback.model_library.all_models_dict['integrated_flux_afterglowpy_base_model']

1✔

189

                ys = function(times, **sample, **kwargs)

1✔

191

                ys = np.ones(len(times))

1✔

192

            self.assertEqual(len(times), len(ys))

1✔

194

class TestIntegratedFluxModelRate(unittest.TestCase):

1✔

195

    def setUp(self) -> None:

1✔

196

        self.path_to_files = f"{_dirname}/../redback/priors/"

1✔

197

        self.prior_files = listdir(self.path_to_files)

1✔

199

    def tearDown(self) -> None:

1✔

200

        pass

1✔

202

    def get_prior(self, file):

1✔

203

        prior_dict = bilby.prior.PriorDict()

1✔

204

        prior_dict.from_file(f"{self.path_to_files}{file}")

1✔

205

        return prior_dict

1✔

207

    def test_models(self):

1✔

208

        kwargs = dict(frequency=np.array([2e14,6e14]), bands='ztfg')

1✔

209

        times = np.array([1, 2, 3])

1✔

210

        valid_models = ['smoothpowerlaw']

1✔

211

        for f in self.prior_files:

1✔

212

            model_name = f.replace(".prior", "")

1✔

213

            if model_name in valid_models:

1✔

214

                print(f)

1✔

215

                prior = self.get_prior(file=f)

1✔

216

                kwargs['base_model'] = model_name

1✔

217

                function = redback.model_library.all_models_dict['integrated_flux_rate_model']

1✔

218

                ys = function(times, **prior.sample(), **kwargs)

1✔

220

                ys = np.ones(len(times))

1✔

221

            self.assertEqual(len(times), len(ys))

1✔

223

class TestExtinctionModelsFluxDensity(unittest.TestCase):

1✔

224

    def setUp(self) -> None:

1✔

225

        self.path_to_files = f"{_dirname}/../redback/priors/"

1✔

226

        self.prior_files = listdir(self.path_to_files)

1✔

228

    def tearDown(self) -> None:

1✔

229

        pass

1✔

231

    def get_prior(self, file):

1✔

232

        prior_dict = bilby.prior.PriorDict()

1✔

233

        prior_dict.from_file(f"{self.path_to_files}{file}")

1✔

234

        return prior_dict

1✔

236

    def test_models(self):

1✔

237

        kwargs = dict(frequency=2e14, bands='ztfg')

1✔

238

        times = np.array([1, 2, 3]) + 58555

1✔

239

        valid_models = ['arnett', 'slsn','one_component_kilonova_model'

1✔

242

        supernova_models = ['arnett', 'slsn']

1✔

243

        kilonova_models = ['one_component_kilonova_model']

1✔

244

        tde_models = ['tde_analytical']

1✔

245

        magnetardriven_models = ['basic_mergernova']

1✔

246

        for f in self.prior_files:

1✔

247

            model_name = f.replace(".prior", "")

1✔

248

            if model_name in valid_models:

1✔

249

                print(f)

1✔

250

                kwargs['output_format'] = 'flux_density'

1✔

251

                kwargs['base_model'] = model_name

1✔

252

                prior = self.get_prior(file=f)

1✔

253

                prior['av'] = 0.5

1✔

254

                prior['t0'] = 58555

1✔

255

                if model_name in supernova_models:

1✔

256

                    function = 't0_supernova_extinction'

1✔

257

                elif model_name in kilonova_models:

1✔

259

                elif model_name in tde_models:

1✔

261

                elif model_name in magnetardriven_models:

1✔

262

                    function = 't0_magnetar_driven_extinction'

1✔

263

                if model_name == 'smoothpowerlaw':

1✔

264

                    function = 't0_afterglow_extinction'

1✔

265

                if model_name == 'gaussiancore':

1✔

266

                    function = 't0_afterglow_extinction_model_d2g'

1✔

267

                    kwargs['lognh'] = 22

1✔

268

                    kwargs['factor'] = 5

1✔

269

                    prior.pop('av')

1✔

270

                if model_name == 'shock_cooling':

1✔

271

                    function = 't0_shock_powered_extinction'

1✔

272

                function = redback.model_library.all_models_dict[function]

1✔

273

                ys = function(times, **prior.sample(), **kwargs)

1✔

275

                ys = np.ones(len(times))

1✔

276

            self.assertEqual(len(times), len(ys))

1✔

278

class TestExtinctionModelsMagnitude(unittest.TestCase):

1✔

279

    def setUp(self) -> None:

1✔

280

        self.path_to_files = f"{_dirname}/../redback/priors/"

1✔

281

        self.prior_files = listdir(self.path_to_files)

1✔

283

    def tearDown(self) -> None:

1✔

284

        pass

1✔

286

    def get_prior(self, file):

1✔

287

        prior_dict = bilby.prior.PriorDict()

1✔

288

        prior_dict.from_file(f"{self.path_to_files}{file}")

1✔

289

        return prior_dict

1✔

291

    def test_models(self):

1✔

292

        kwargs = dict(frequency=2e14, bands='ztfg')

1✔

293

        times = np.array([1, 2, 3]) + 58555

1✔

294

        valid_models = ['arnett', 'slsn','one_component_kilonova_model'

1✔

296

        supernova_models = ['arnett', 'slsn']

1✔

297

        kilonova_models = ['one_component_kilonova_model']

1✔

298

        tde_models = ['tde_analytical']

1✔

299

        magnetardriven_models = ['basic_mergernova']

1✔

300

        for f in self.prior_files:

1✔

301

            model_name = f.replace(".prior", "")

1✔

302

            if model_name in valid_models:

1✔

303

                print(f)

1✔

304

                kwargs['output_format'] = 'magnitude'

1✔

305

                kwargs['base_model'] = model_name

1✔

306

                prior = self.get_prior(file=f)

1✔

307

                prior['av'] = 0.5

1✔

308

                prior['t0'] = 58555

1✔

309

                if model_name in supernova_models:

1✔

310

                    function = 't0_supernova_extinction'

1✔

311

                elif model_name in kilonova_models:

1✔

313

                elif model_name in tde_models:

1✔

315

                elif model_name in magnetardriven_models:

1✔

316

                    function = 't0_magnetar_driven_extinction'

1✔

317

                if model_name == 'shock_cooling':

1✔

318

                    function = 't0_shock_powered_extinction'

1✔

319

                function = redback.model_library.all_models_dict[function]

1✔

320

                ys = function(times, **prior.sample(), **kwargs)

1✔

322

                ys = np.ones(len(times))

1✔

323

            self.assertEqual(len(times), len(ys))

1✔

325

class TestHomologousExpansion(unittest.TestCase):

1✔

326

    def setUp(self):

1✔

327

        self.time = np.array([1, 2, 3])

1✔

328

        self.base_model = 'arnett_bolometric'

1✔

330

    def tearDown(self) -> None:

1✔

331

        pass

1✔

333

    def get_prior(self):

1✔

334

        prior_dict = redback.priors.get_priors(self.base_model)

1✔

335

        return prior_dict

1✔

337

    def test_base_model(self):

1✔

338

        with mock.patch('redback.transient_models.supernova_models.homologous_expansion_supernova') as m:

1✔

339

            kwargs = dict(frequency=2e14, bands='ztfg')

1✔

340

            prior = self.get_prior()

1✔

341

            kwargs['output_format'] = 'flux_density'

1✔

342

            prior.pop('vej')

1✔

343

            prior['ek'] =  bilby.core.prior.LogUniform(1e50, 1e51, 'ek')

1✔

344

            prior['redshift'] = 0.075

1✔

345

            prior['temperature_floor'] = bilby.core.prior.LogUniform(1e3,1e5,name = 'temperature_floor')

1✔

346

            function = redback.transient_models.supernova_models.homologous_expansion_supernova

1✔

347

            m.base_model = self.base_model

1✔

348

            sample = prior.sample()

1✔

349

            actual = function(self.time, base_model=self.base_model, **sample, **kwargs)

1✔

350

            m.assert_called_with(self.time, base_model=self.base_model, **sample, **kwargs)

1✔

352

    def test_fluxdensity_output(self):

1✔

353

        kwargs = dict(frequency=2e14, bands='ztfg')

1✔

354

        prior = self.get_prior()

1✔

355

        kwargs['base_model'] = self.base_model

1✔

356

        kwargs['output_format'] = 'flux_density'

1✔

357

        function = redback.model_library.all_models_dict['homologous_expansion_supernova']

1✔

358

        prior.pop('vej')

1✔

359

        prior['ek'] =  bilby.core.prior.LogUniform(1e50, 1e51, 'ek')

1✔

360

        prior['redshift'] = 0.075

1✔

361

        prior['temperature_floor'] = bilby.core.prior.LogUniform(1e3,1e5,name = 'temperature_floor')

1✔

362

        ys = function(self.time, **prior.sample(), **kwargs)

1✔

363

        self.assertEqual(len(self.time), len(ys))

1✔

365

    def test_magnitude_output(self):

1✔

366

        kwargs = dict(frequency=2e14, bands='ztfg')

1✔

367

        prior = self.get_prior()

1✔

368

        kwargs['base_model'] = self.base_model

1✔

369

        kwargs['output_format'] = 'magnitude'

1✔

370

        function = redback.model_library.all_models_dict['homologous_expansion_supernova']

1✔

371

        prior.pop('vej')

1✔

372

        prior['ek'] =  bilby.core.prior.LogUniform(1e50, 1e51, 'ek')

1✔

373

        prior['redshift'] = 0.075

1✔

374

        prior['temperature_floor'] = bilby.core.prior.LogUniform(1e3,1e5,name = 'temperature_floor')

1✔

375

        ys = function(self.time, **prior.sample(), **kwargs)

1✔

376

        self.assertEqual(len(self.time), len(ys))

1✔

378

class TestThinShellExpansion(unittest.TestCase):

1✔

379

    def setUp(self):

1✔

380

        self.time = np.array([1, 2, 3])

1✔

381

        self.base_model = 'arnett_bolometric'

1✔

383

    def tearDown(self) -> None:

1✔

384

        pass

1✔

386

    def get_prior(self):

1✔

387

        prior_dict = redback.priors.get_priors(self.base_model)

1✔

388

        return prior_dict

1✔

390

    def test_base_model(self):

1✔

391

        with mock.patch('redback.transient_models.supernova_models.thin_shell_supernova') as m:

1✔

392

            kwargs = dict(frequency=2e14, bands='ztfg')

1✔

393

            prior = self.get_prior()

1✔

394

            kwargs['output_format'] = 'flux_density'

1✔

395

            prior.pop('vej')

1✔

396

            prior['ek'] =  bilby.core.prior.LogUniform(1e50, 1e51, 'ek')

1✔

397

            prior['redshift'] = 0.075

1✔

398

            prior['temperature_floor'] = bilby.core.prior.LogUniform(1e3,1e5,name = 'temperature_floor')

1✔

399

            function = redback.transient_models.supernova_models.thin_shell_supernova

1✔

400

            m.base_model = self.base_model

1✔

401

            sample = prior.sample()

1✔

402

            actual = function(self.time, base_model=self.base_model, **sample, **kwargs)

1✔

403

            m.assert_called_with(self.time, base_model=self.base_model, **sample, **kwargs)

1✔

405

    def test_fluxdensity_output(self):

1✔

406

        kwargs = dict(frequency=2e14, bands='ztfg')

1✔

407

        prior = self.get_prior()

1✔

408

        kwargs['base_model'] = self.base_model

1✔

409

        kwargs['output_format'] = 'flux_density'

1✔

410

        function = redback.model_library.all_models_dict['thin_shell_supernova']

1✔

411

        prior.pop('vej')

1✔

412

        prior['ek'] =  bilby.core.prior.LogUniform(1e50, 1e51, 'ek')

1✔

413

        prior['redshift'] = 0.075

1✔

414

        prior['temperature_floor'] = bilby.core.prior.LogUniform(1e3,1e5,name = 'temperature_floor')

1✔

415

        ys = function(self.time, **prior.sample(), **kwargs)

1✔

416

        self.assertEqual(len(self.time), len(ys))

1✔

418

    def test_magnitude_output(self):

1✔

419

        kwargs = dict(frequency=2e14, bands='ztfg')

1✔

420

        prior = self.get_prior()

1✔

421

        kwargs['base_model'] = self.base_model

1✔

422

        kwargs['output_format'] = 'magnitude'

1✔

423

        function = redback.model_library.all_models_dict['thin_shell_supernova']

1✔

424

        prior.pop('vej')

1✔

425

        prior['ek'] =  bilby.core.prior.LogUniform(1e50, 1e51, 'ek')

1✔

426

        prior['redshift'] = 0.075

1✔

427

        prior['temperature_floor'] = bilby.core.prior.LogUniform(1e3,1e5,name = 'temperature_floor')

1✔

428

        ys = function(self.time, **prior.sample(), **kwargs)

1✔

429

        self.assertEqual(len(self.time), len(ys))

1✔

nikhil-sarin / redback / 13137940639

Source File
Press 'n' to go to next uncovered line, 'b' for previous

1	import unittest	1✔
2	from os import listdir	1✔
3	from os.path import dirname	1✔
4	from unittest import mock	1✔
5
6	import bilby	1✔
7	import numpy as np	1✔
8
9	import redback.model_library	1✔
10
11	_dirname = dirname(__file__)	1✔
12
13
14	class TestModels(unittest.TestCase):	1✔
15
16	def setUp(self) -> None:	1✔
17	self.path_to_files = f"{_dirname}/../redback/priors/"	1✔
18	self.prior_files = listdir(self.path_to_files)	1✔
19
20	def tearDown(self) -> None:	1✔
21	pass	1✔
22
23	def get_prior(self, file):	1✔
24	prior_dict = bilby.prior.PriorDict()	1✔
25	prior_dict.from_file(f"{self.path_to_files}{file}")	1✔
26	return prior_dict	1✔
27
28	def test_models(self):	1✔
29	kwargs = dict(frequency=2e14)	1✔
30	times = np.array([1, 2, 3])	1✔
31	for f in self.prior_files:	1✔
32	print(f)	1✔
33	model_name = f.replace(".prior", "")	1✔
34	if model_name == 'trapped_magnetar':	1✔
35	kwargs['output_format'] = 'luminosity'	1✔
36	elif model_name in ['fitted', 'jetsimpy_tophat', 'jetsimpy_gaussian', 'jetsimpy_powerlaw']:	1✔
NEW 37	pass	×
38	else:
39	kwargs['output_format'] = 'flux_density'	1✔
40	prior = self.get_prior(file=f)	1✔
41	sample = prior.sample()	1✔
42	if model_name == 'polytrope_eos_two_component_bns':	1✔
43	sample['gamma_1'] = 4.04	1✔
44	sample['gamma_2'] = 2.159	1✔
45	sample['gamma_3'] = 3.688	1✔
46	sample['log_p'] = 33.72	1✔
47	function = redback.model_library.all_models_dict[model_name]	1✔
48	ys = function(times, sample, kwargs)	1✔
49	self.assertEqual(len(times), len(ys))	1✔
50
51	class TestPhaseModels(unittest.TestCase):	1✔
52	def setUp(self) -> None:	1✔
53	self.path_to_files = f"{_dirname}/../redback/priors/"	1✔
54	self.prior_files = listdir(self.path_to_files)	1✔
55
56	def tearDown(self) -> None:	1✔
57	pass	1✔
58
59	def get_prior(self, file):	1✔
60	prior_dict = bilby.prior.PriorDict()	1✔
61	prior_dict.from_file(f"{self.path_to_files}{file}")	1✔
62	return prior_dict	1✔
63
64	def test_models(self):	1✔
65	kwargs = dict(frequency=2e14)	1✔
66	times = np.array([1, 2, 3]) + 55855	1✔
67	for f in self.prior_files:	1✔
68	print(f)	1✔
69	model_name = f.replace(".prior", "")	1✔
70	if model_name == 'trapped_magnetar':	1✔
71	kwargs['output_format'] = 'luminosity'	1✔
72	elif model_name in ['fitted', 'jetsimpy_tophat', 'jetsimpy_gaussian', 'jetsimpy_powerlaw']:	1✔
NEW 73	pass	×
74	else:
75	kwargs['output_format'] = 'flux_density'	1✔
76	prior = self.get_prior(file=f)	1✔
77	function = redback.model_library.all_models_dict['t0_base_model']	1✔
78	kwargs['base_model'] = model_name	1✔
79	kwargs['t0'] = 55855	1✔
80	sample = prior.sample()	1✔
81	if model_name == 'polytrope_eos_two_component_bns':	1✔
82	sample['gamma_1'] = 4.04	1✔
83	sample['gamma_2'] = 2.159	1✔
84	sample['gamma_3'] = 3.688	1✔
85	sample['log_p'] = 33.72	1✔
86	ys = function(times, sample, kwargs)	1✔
87	self.assertEqual(len(times), len(ys))	1✔
88
89	class TestMagnitudeOutput(unittest.TestCase):	1✔
90	def setUp(self) -> None:	1✔
91	self.path_to_files = f"{_dirname}/../redback/priors/"	1✔
92	self.prior_files = listdir(self.path_to_files)	1✔
93
94	def tearDown(self) -> None:	1✔
95	pass	1✔
96
97	def get_prior(self, file):	1✔
98	prior_dict = bilby.prior.PriorDict()	1✔
99	prior_dict.from_file(f"{self.path_to_files}{file}")	1✔
100	return prior_dict	1✔
101
102	def test_models(self):	1✔
103	kwargs = dict(frequency=2e14, bands='ztfg')	1✔
104	times = np.array([1, 2, 3])	1✔
105	for f in self.prior_files:	1✔
106	print(f)	1✔
107	model_name = f.replace(".prior", "")	1✔
108	if model_name == 'trapped_magnetar':	1✔
109	kwargs['output_format'] = 'luminosity'	1✔
110	elif model_name in ['tophat_and_twocomponent', 'tophat_and_twolayerstratified',	1✔
111	'tophat_and_arnett', 'tophatredback_and_twolayerstratified',
112	'fitted', 'jetsimpy_tophat', 'jetsimpy_gaussian', 'jetsimpy_powerlaw']:
113	pass	1✔
114	else:
115	kwargs['output_format'] = 'magnitude'	1✔
116	prior = self.get_prior(file=f)	1✔
117	sample = prior.sample()	1✔
118	if model_name == 'polytrope_eos_two_component_bns':	1✔
119	sample['gamma_1'] = 4.04	1✔
120	sample['gamma_2'] = 2.159	1✔
121	sample['gamma_3'] = 3.688	1✔
122	sample['log_p'] = 33.72	1✔
123	function = redback.model_library.all_models_dict[model_name]	1✔
124	ys = function(times, sample, kwargs)	1✔
125	self.assertEqual(len(times), len(ys))	1✔
126
127	class TestFluxOutput(unittest.TestCase):	1✔
128	def setUp(self) -> None:	1✔
129	self.path_to_files = f"{_dirname}/../redback/priors/"	1✔
130	self.prior_files = listdir(self.path_to_files)	1✔
131
132	def tearDown(self) -> None:	1✔
133	pass	1✔
134
135	def get_prior(self, file):	1✔
136	prior_dict = bilby.prior.PriorDict()	1✔
137	prior_dict.from_file(f"{self.path_to_files}{file}")	1✔
138	return prior_dict	1✔
139
140	def test_models(self):	1✔
141	kwargs = dict(frequency=2e14, bands='ztfg')	1✔
142	times = np.array([1, 2, 3])	1✔
143	valid_models = ['arnett', 'one_component_kilonova_model', 'slsn',	1✔
144	'tde_analytical', 'basic_mergernova', 'csm_nickel']
145	for f in self.prior_files:	1✔
146	model_name = f.replace(".prior", "")	1✔
147	if model_name in valid_models:	1✔
148	print(f)	1✔
149	kwargs['output_format'] = 'flux'	1✔
150	prior = self.get_prior(file=f)	1✔
151	sample = prior.sample()	1✔
152	if model_name == 'polytrope_eos_two_component_bns':	1✔
153	sample['gamma_1'] = 4.04	×
154	sample['gamma_2'] = 2.159	×
155	sample['gamma_3'] = 3.688	×
156	sample['log_p'] = 33.72	×
157	function = redback.model_library.all_models_dict[model_name]	1✔
158	ys = function(times, sample, kwargs)	1✔
159	else:
160	ys = np.ones(len(times))	1✔
161	self.assertEqual(len(times), len(ys))	1✔
162
163
164	class TestIntegratedFluxModelFlux(unittest.TestCase):	1✔
165	def setUp(self) -> None:	1✔
166	self.path_to_files = f"{_dirname}/../redback/priors/"	1✔
167	self.prior_files = listdir(self.path_to_files)	1✔
168
169	def tearDown(self) -> None:	1✔
170	pass	1✔
171
172	def get_prior(self, file):	1✔
173	prior_dict = bilby.prior.PriorDict()	1✔
174	prior_dict.from_file(f"{self.path_to_files}{file}")	1✔
175	return prior_dict	1✔
176
177	def test_models(self):	1✔
178	kwargs = dict(frequency=np.array([2e14,6e14]), bands='ztfg')	1✔
179	times = np.array([1, 2, 3])	1✔
180	valid_models = ['gaussiancore']	1✔
181	for f in self.prior_files:	1✔
182	model_name = f.replace(".prior", "")	1✔
183	if model_name in valid_models:	1✔
184	print(f)	1✔
185	prior = self.get_prior(file=f)	1✔
186	kwargs['base_model'] = model_name	1✔
187	sample = prior.sample()	1✔
188	function = redback.model_library.all_models_dict['integrated_flux_afterglowpy_base_model']	1✔
189	ys = function(times, sample, kwargs)	1✔
190	else:
191	ys = np.ones(len(times))	1✔
192	self.assertEqual(len(times), len(ys))	1✔
193
194	class TestIntegratedFluxModelRate(unittest.TestCase):	1✔
195	def setUp(self) -> None:	1✔
196	self.path_to_files = f"{_dirname}/../redback/priors/"	1✔
197	self.prior_files = listdir(self.path_to_files)	1✔
198
199	def tearDown(self) -> None:	1✔
200	pass	1✔
201
202	def get_prior(self, file):	1✔
203	prior_dict = bilby.prior.PriorDict()	1✔
204	prior_dict.from_file(f"{self.path_to_files}{file}")	1✔
205	return prior_dict	1✔
206
207	def test_models(self):	1✔
208	kwargs = dict(frequency=np.array([2e14,6e14]), bands='ztfg')	1✔
209	times = np.array([1, 2, 3])	1✔
210	valid_models = ['smoothpowerlaw']	1✔
211	for f in self.prior_files:	1✔
212	model_name = f.replace(".prior", "")	1✔
213	if model_name in valid_models:	1✔
214	print(f)	1✔
215	prior = self.get_prior(file=f)	1✔
216	kwargs['base_model'] = model_name	1✔
217	function = redback.model_library.all_models_dict['integrated_flux_rate_model']	1✔
218	ys = function(times, prior.sample(), kwargs)	1✔
219	else:
220	ys = np.ones(len(times))	1✔
221	self.assertEqual(len(times), len(ys))	1✔
222
223	class TestExtinctionModelsFluxDensity(unittest.TestCase):	1✔
224	def setUp(self) -> None:	1✔
225	self.path_to_files = f"{_dirname}/../redback/priors/"	1✔
226	self.prior_files = listdir(self.path_to_files)	1✔
227
228	def tearDown(self) -> None:	1✔
229	pass	1✔
230
231	def get_prior(self, file):	1✔
232	prior_dict = bilby.prior.PriorDict()	1✔
233	prior_dict.from_file(f"{self.path_to_files}{file}")	1✔
234	return prior_dict	1✔
235
236	def test_models(self):	1✔
237	kwargs = dict(frequency=2e14, bands='ztfg')	1✔
238	times = np.array([1, 2, 3]) + 58555	1✔
239	valid_models = ['arnett', 'slsn','one_component_kilonova_model'	1✔
240	'tde_analytical', 'basic_mergernova', 'gaussiancore',
241	'smoothpowerlaw', 'shock_cooling']
242	supernova_models = ['arnett', 'slsn']	1✔
243	kilonova_models = ['one_component_kilonova_model']	1✔
244	tde_models = ['tde_analytical']	1✔
245	magnetardriven_models = ['basic_mergernova']	1✔
246	for f in self.prior_files:	1✔
247	model_name = f.replace(".prior", "")	1✔
248	if model_name in valid_models:	1✔
249	print(f)	1✔
250	kwargs['output_format'] = 'flux_density'	1✔
251	kwargs['base_model'] = model_name	1✔
252	prior = self.get_prior(file=f)	1✔
253	prior['av'] = 0.5	1✔
254	prior['t0'] = 58555	1✔
255	if model_name in supernova_models:	1✔
256	function = 't0_supernova_extinction'	1✔
257	elif model_name in kilonova_models:	1✔
258	function = 't0_kilonova_extinction'	×
259	elif model_name in tde_models:	1✔
260	function = 't0_tde_extinction'	×
261	elif model_name in magnetardriven_models:	1✔
262	function = 't0_magnetar_driven_extinction'	1✔
263	if model_name == 'smoothpowerlaw':	1✔
264	function = 't0_afterglow_extinction'	1✔
265	if model_name == 'gaussiancore':	1✔
266	function = 't0_afterglow_extinction_model_d2g'	1✔
267	kwargs['lognh'] = 22	1✔
268	kwargs['factor'] = 5	1✔
269	prior.pop('av')	1✔
270	if model_name == 'shock_cooling':	1✔
271	function = 't0_shock_powered_extinction'	1✔
272	function = redback.model_library.all_models_dict[function]	1✔
273	ys = function(times, prior.sample(), kwargs)	1✔
274	else:
275	ys = np.ones(len(times))	1✔
276	self.assertEqual(len(times), len(ys))	1✔
277
278	class TestExtinctionModelsMagnitude(unittest.TestCase):	1✔
279	def setUp(self) -> None:	1✔
280	self.path_to_files = f"{_dirname}/../redback/priors/"	1✔
281	self.prior_files = listdir(self.path_to_files)	1✔
282
283	def tearDown(self) -> None:	1✔
284	pass	1✔
285
286	def get_prior(self, file):	1✔
287	prior_dict = bilby.prior.PriorDict()	1✔
288	prior_dict.from_file(f"{self.path_to_files}{file}")	1✔
289	return prior_dict	1✔
290
291	def test_models(self):	1✔
292	kwargs = dict(frequency=2e14, bands='ztfg')	1✔
293	times = np.array([1, 2, 3]) + 58555	1✔
294	valid_models = ['arnett', 'slsn','one_component_kilonova_model'	1✔
295	'tde_analytical', 'basic_mergernova', 'shock_cooling']
296	supernova_models = ['arnett', 'slsn']	1✔
297	kilonova_models = ['one_component_kilonova_model']	1✔
298	tde_models = ['tde_analytical']	1✔
299	magnetardriven_models = ['basic_mergernova']	1✔
300	for f in self.prior_files:	1✔
301	model_name = f.replace(".prior", "")	1✔
302	if model_name in valid_models:	1✔
303	print(f)	1✔
304	kwargs['output_format'] = 'magnitude'	1✔
305	kwargs['base_model'] = model_name	1✔
306	prior = self.get_prior(file=f)	1✔
307	prior['av'] = 0.5	1✔
308	prior['t0'] = 58555	1✔
309	if model_name in supernova_models:	1✔
310	function = 't0_supernova_extinction'	1✔
311	elif model_name in kilonova_models:	1✔
312	function = 't0_kilonova_extinction'	×
313	elif model_name in tde_models:	1✔
314	function = 't0_tde_extinction'	×
315	elif model_name in magnetardriven_models:	1✔
316	function = 't0_magnetar_driven_extinction'	1✔
317	if model_name == 'shock_cooling':	1✔
318	function = 't0_shock_powered_extinction'	1✔
319	function = redback.model_library.all_models_dict[function]	1✔
320	ys = function(times, prior.sample(), kwargs)	1✔
321	else:
322	ys = np.ones(len(times))	1✔
323	self.assertEqual(len(times), len(ys))	1✔
324
325	class TestHomologousExpansion(unittest.TestCase):	1✔
326	def setUp(self):	1✔
327	self.time = np.array([1, 2, 3])	1✔
328	self.base_model = 'arnett_bolometric'	1✔
329
330	def tearDown(self) -> None:	1✔
331	pass	1✔
332
333	def get_prior(self):	1✔
334	prior_dict = redback.priors.get_priors(self.base_model)	1✔
335	return prior_dict	1✔
336
337	def test_base_model(self):	1✔
338	with mock.patch('redback.transient_models.supernova_models.homologous_expansion_supernova') as m:	1✔
339	kwargs = dict(frequency=2e14, bands='ztfg')	1✔
340	prior = self.get_prior()	1✔
341	kwargs['output_format'] = 'flux_density'	1✔
342	prior.pop('vej')	1✔
343	prior['ek'] = bilby.core.prior.LogUniform(1e50, 1e51, 'ek')	1✔
344	prior['redshift'] = 0.075	1✔
345	prior['temperature_floor'] = bilby.core.prior.LogUniform(1e3,1e5,name = 'temperature_floor')	1✔
346	function = redback.transient_models.supernova_models.homologous_expansion_supernova	1✔
347	m.base_model = self.base_model	1✔
348	sample = prior.sample()	1✔
349	actual = function(self.time, base_model=self.base_model, sample, kwargs)	1✔
350	m.assert_called_with(self.time, base_model=self.base_model, sample, kwargs)	1✔
351
352	def test_fluxdensity_output(self):	1✔
353	kwargs = dict(frequency=2e14, bands='ztfg')	1✔
354	prior = self.get_prior()	1✔
355	kwargs['base_model'] = self.base_model	1✔
356	kwargs['output_format'] = 'flux_density'	1✔
357	function = redback.model_library.all_models_dict['homologous_expansion_supernova']	1✔
358	prior.pop('vej')	1✔
359	prior['ek'] = bilby.core.prior.LogUniform(1e50, 1e51, 'ek')	1✔
360	prior['redshift'] = 0.075	1✔
361	prior['temperature_floor'] = bilby.core.prior.LogUniform(1e3,1e5,name = 'temperature_floor')	1✔
362	ys = function(self.time, prior.sample(), kwargs)	1✔
363	self.assertEqual(len(self.time), len(ys))	1✔
364
365	def test_magnitude_output(self):	1✔
366	kwargs = dict(frequency=2e14, bands='ztfg')	1✔
367	prior = self.get_prior()	1✔
368	kwargs['base_model'] = self.base_model	1✔
369	kwargs['output_format'] = 'magnitude'	1✔
370	function = redback.model_library.all_models_dict['homologous_expansion_supernova']	1✔
371	prior.pop('vej')	1✔
372	prior['ek'] = bilby.core.prior.LogUniform(1e50, 1e51, 'ek')	1✔
373	prior['redshift'] = 0.075	1✔
374	prior['temperature_floor'] = bilby.core.prior.LogUniform(1e3,1e5,name = 'temperature_floor')	1✔
375	ys = function(self.time, prior.sample(), kwargs)	1✔
376	self.assertEqual(len(self.time), len(ys))	1✔
377
378	class TestThinShellExpansion(unittest.TestCase):	1✔
379	def setUp(self):	1✔
380	self.time = np.array([1, 2, 3])	1✔
381	self.base_model = 'arnett_bolometric'	1✔
382
383	def tearDown(self) -> None:	1✔
384	pass	1✔
385
386	def get_prior(self):	1✔
387	prior_dict = redback.priors.get_priors(self.base_model)	1✔
388	return prior_dict	1✔
389
390	def test_base_model(self):	1✔
391	with mock.patch('redback.transient_models.supernova_models.thin_shell_supernova') as m:	1✔
392	kwargs = dict(frequency=2e14, bands='ztfg')	1✔
393	prior = self.get_prior()	1✔
394	kwargs['output_format'] = 'flux_density'	1✔
395	prior.pop('vej')	1✔
396	prior['ek'] = bilby.core.prior.LogUniform(1e50, 1e51, 'ek')	1✔
397	prior['redshift'] = 0.075	1✔
398	prior['temperature_floor'] = bilby.core.prior.LogUniform(1e3,1e5,name = 'temperature_floor')	1✔
399	function = redback.transient_models.supernova_models.thin_shell_supernova	1✔
400	m.base_model = self.base_model	1✔
401	sample = prior.sample()	1✔
402	actual = function(self.time, base_model=self.base_model, sample, kwargs)	1✔
403	m.assert_called_with(self.time, base_model=self.base_model, sample, kwargs)	1✔
404
405	def test_fluxdensity_output(self):	1✔
406	kwargs = dict(frequency=2e14, bands='ztfg')	1✔
407	prior = self.get_prior()	1✔
408	kwargs['base_model'] = self.base_model	1✔
409	kwargs['output_format'] = 'flux_density'	1✔
410	function = redback.model_library.all_models_dict['thin_shell_supernova']	1✔
411	prior.pop('vej')	1✔
412	prior['ek'] = bilby.core.prior.LogUniform(1e50, 1e51, 'ek')	1✔
413	prior['redshift'] = 0.075	1✔
414	prior['temperature_floor'] = bilby.core.prior.LogUniform(1e3,1e5,name = 'temperature_floor')	1✔
415	ys = function(self.time, prior.sample(), kwargs)	1✔
416	self.assertEqual(len(self.time), len(ys))	1✔
417
418	def test_magnitude_output(self):	1✔
419	kwargs = dict(frequency=2e14, bands='ztfg')	1✔
420	prior = self.get_prior()	1✔
421	kwargs['base_model'] = self.base_model	1✔
422	kwargs['output_format'] = 'magnitude'	1✔
423	function = redback.model_library.all_models_dict['thin_shell_supernova']	1✔
424	prior.pop('vej')	1✔
425	prior['ek'] = bilby.core.prior.LogUniform(1e50, 1e51, 'ek')	1✔
426	prior['redshift'] = 0.075	1✔
427	prior['temperature_floor'] = bilby.core.prior.LogUniform(1e3,1e5,name = 'temperature_floor')	1✔
428	ys = function(self.time, prior.sample(), kwargs)	1✔
429	self.assertEqual(len(self.time), len(ys))	1✔
430

nikhil-sarin / redback / 13137940639

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous