Skip to content

Commit

Permalink
tests: add and fix tests; mend blocks accordingly
Browse files Browse the repository at this point in the history
  • Loading branch information
@jeremy-baier
jeremy-baier committed Sep 9, 2024
1 parent 3d6250b commit 2e08082
Show file tree
Hide file tree
Showing 2 changed files with 38 additions and 27 deletions.
23 changes: 12 additions & 11 deletions enterprise_extensions/blocks.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -227,7 +227,8 @@ def red_noise_block(psd='powerlaw', prior='log-uniform', Tspan=None,

rn = gp_signals.FourierBasisGP(pl, components=components_low,
Tspan=Tspan, coefficients=coefficients,
combine=combine, selection=selection)
combine=combine, selection=selection,
name='red_noise')

rn_flat = gp_signals.FourierBasisGP(pl_flat,
modes=freqs[components_low:],
Expand All @@ -242,7 +243,8 @@ def red_noise_block(psd='powerlaw', prior='log-uniform', Tspan=None,
combine=combine,
coefficients=coefficients,
selection=selection,
modes=modes)
modes=modes,
name='red_noise')

if select == 'band+': # Add the common component as well
rn = rn + gp_signals.FourierBasisGP(pl, components=components,
Expand Down Expand Up @@ -419,13 +421,13 @@ def dm_noise_block(gp_kernel='diag', psd='powerlaw', nondiag_kernel='periodic',

if psd == 'spectrum':
if not vary:
log10_rho_dm = parameter.Constant()
log10_rho = parameter.Constant()
elif prior == 'uniform':
log10_rho_dm = parameter.LinearExp(-10, -4, size=components)
log10_rho = parameter.LinearExp(-10, -4, size=components)
elif prior == 'log-uniform':
log10_rho_dm = parameter.Uniform(-10, -4, size=components)
log10_rho = parameter.Uniform(-10, -4, size=components)

dm_prior = gpp.free_spectrum(log10_rho=log10_rho_dm)
dm_prior = gpp.free_spectrum(log10_rho=log10_rho)

dm_basis = utils.createfourierdesignmatrix_dm(nmodes=components,
Tspan=Tspan)
Expand Down Expand Up @@ -603,13 +605,12 @@ def chromatic_noise_block(gp_kernel='nondiag', psd='powerlaw',

if psd == 'spectrum':
if not vary:
log10_rho_chrom = parameter.Constant()
log10_rho = parameter.Constant()
elif prior == 'uniform':
log10_rho_chrom = parameter.LinearExp(-10, -4, size=components)
log10_rho = parameter.LinearExp(-10, -4, size=components)
elif prior == 'log-uniform':
log10_rho_chrom = parameter.Uniform(-10, -4, size=components)
print(log10_rho_chrom)
chm_prior = gpp.free_spectrum(log10_rho=log10_rho_chrom)
log10_rho = parameter.Uniform(-10, -4, size=components)
chm_prior = gpp.free_spectrum(log10_rho=log10_rho)

elif gp_kernel == 'nondiag':
if nondiag_kernel == 'periodic':
Expand Down
42 changes: 26 additions & 16 deletions tests/test_models.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
import logging
import os
import pickle
import numpy as np

import pytest
from enterprise import constants as const
Expand Down Expand Up @@ -194,48 +195,57 @@ def test_model_singlepsr_noise_chrom_nondiag(nodmx_psrs, caplog):
m.get_lnlikelihood(x0)


# FIXME: currently there are not any dm gp free spec params in the pulsars
def test_model_singlepsr_noise_dm_diag(nodmx_psrs, caplog):
# caplog.set_level(logging.CRITICAL)
m=models.model_singlepsr_noise(nodmx_psrs[0], vary_dm=True,
m=models.model_singlepsr_noise(nodmx_psrs[0], vary_dm=True, dm_var=True,
dmgp_kernel='diag')
assert hasattr(m, 'get_lnlikelihood')
x0 = {pname: p.sample() for pname, p in zip(m.param_names, m.params)}
x0 = np.hstack([p.sample() for p in m.params])
m.get_lnlikelihood(x0)
m=models.model_singlepsr_noise(nodmx_psrs[1], vary_dm=True,
m=models.model_singlepsr_noise(nodmx_psrs[1], vary_dm=True, dm_var=True,
dmgp_kernel='diag',
dm_psd='turnover')
assert hasattr(m, 'get_lnlikelihood')
x0 = {pname: p.sample() for pname, p in zip(m.param_names, m.params)}
x0 = np.hstack([p.sample() for p in m.params])
m.get_lnlikelihood(x0)
m=models.model_singlepsr_noise(nodmx_psrs[2], vary_dm=True,
m=models.model_singlepsr_noise(nodmx_psrs[2], vary_dm=True, dm_var=True,
dmgp_kernel='diag',
dm_psd='spectrum')
assert hasattr(m, 'get_lnlikelihood')
x0 = {pname: p.sample() for pname, p in zip(m.param_names, m.params)}
x0 = np.hstack([p.sample() for p in m.params])
print(x0)
m.get_lnlikelihood(x0)


def test_model_singlepsr_noise_chrom_diag(nodmx_psrs, caplog):
# caplog.set_level(logging.CRITICAL)
m=models.model_singlepsr_noise(nodmx_psrs[1], chrom_gp=True,
m=models.model_singlepsr_noise(nodmx_psrs[1], chrom_gp=True, vary_chrom=True,
chrom_gp_kernel='diag')
assert hasattr(m, 'get_lnlikelihood')
x0 = {pname: p.sample() for pname, p in zip(m.param_names, m.params)}
x0 = np.hstack([p.sample() for p in m.params])
m.get_lnlikelihood(x0)
m=models.model_singlepsr_noise(nodmx_psrs[1], chrom_gp=True,
m=models.model_singlepsr_noise(nodmx_psrs[1], chrom_gp=True, vary_chrom=True,
chrom_gp_kernel='diag',
chrom_psd='turnover')
assert hasattr(m, 'get_lnlikelihood')
x0 = {pname: p.sample() for pname, p in zip(m.param_names, m.params)}
x0 = np.hstack([p.sample() for p in m.params])
m.get_lnlikelihood(x0)
m=models.model_singlepsr_noise(nodmx_psrs[2], chrom_gp=True,
chrom_gp_kernel='diag',
chrom_psd='spectrum')
m=models.model_singlepsr_noise(nodmx_psrs[1],
chrom_gp=True, vary_chrom=True, chrom_gp_kernel='diag', chrom_psd='spectrum',
vary_dm=True, dm_var=True, dm_type='gp', dmgp_kernel='diag', dm_psd='spectrum',
red_var=True, psd='spectrum')
assert hasattr(m, 'get_lnlikelihood')
x0 = {pname: p.sample() for pname, p in zip(m.param_names, m.params)}
print(x0)
x0 = np.hstack([p.sample() for p in m.params])
m.get_lnlikelihood(x0)


def test_model_singlepsr_noise_chrom_vary_idx(nodmx_psrs, caplog):
# caplog.set_level(logging.CRITICAL)
m=models.model_singlepsr_noise(nodmx_psrs[1], chrom_gp=True, vary_chrom=True,
chrom_gp_kernel='diag', chrom_idx='vary')
assert hasattr(m, 'get_lnlikelihood')
assert "J1713+0747_chrom_gp_idx" in m.param_names
x0 = np.hstack([p.sample() for p in m.params])
m.get_lnlikelihood(x0)


Expand Down

0 comments on commit 2e08082

Please sign in to comment.