diff --git a/src/discovery/matrix.py b/src/discovery/matrix.py index 13b29b6..04b7620 100644 --- a/src/discovery/matrix.py +++ b/src/discovery/matrix.py @@ -11,7 +11,7 @@ def config(**kwargs): global jnp, jsp, jnparray, jnpzeros, intarray, jnpkey, jnpsplit, jnpnormal - global matrix_factor, matrix_solve, matrix_norm + global matrix_factor, matrix_solve, matrix_norm, partial np.logdet = lambda a: np.sum(np.log(np.abs(a))) jax.numpy.logdet = lambda a: jax.numpy.sum(jax.numpy.log(jax.numpy.abs(a))) diff --git a/src/discovery/models.py b/src/discovery/models.py new file mode 100644 index 0000000..e279663 --- /dev/null +++ b/src/discovery/models.py @@ -0,0 +1,109 @@ +from functools import partial +import discovery as ds +from typing import Optional, Callable + + +def make_array_likelihood( + psrs: list[ds.Pulsar], + correlation_orf: Optional[Callable]=None, + spectrum: Callable=ds.powerlaw, + intrinsic_red_components: int=30, + common_components: int=14, + gamma_common: Optional[float]=None + ) -> ds.ArrayLikelihood: + """Construct discovery array likelihood object from a list of Pulsar objects which may contain a noise dictionary. + + Parameters: + psrs (list): List of pulsar objects + correlation_orf (function, optional) [None]: Correlation overlap reduction function (ORF). Default gives intrinsic red noise only (IRN) model. + Options include `ds.uncorrelated_orf`, `ds.hd_orf`, `ds.dipole_orf`, `ds.monopole_orf`, or a custom ORF function + spectrum (function): Spectrum function for all red noise parameters + Options include `ds.powerlaw`, `ds.freespectrum`, or a custom spectrum function + intrinsic_red_components (int) [30]: Number of individual pulsar red noise components + common_components (int) [14]: Number of common red noise components + gamma_common (float) [None]: Common red noise spectral index + + Returns: + array_likelihood (object): Discovery ArrayLikelihood object + """ + + tspan = ds.getspan(psrs) + + if gamma_common is not None and (correlation_orf is None or spectrum.__name__ != 'powerlaw'): + raise ValueError('gamma_common is set but correlation_orf is None or spectrum is not a powerlaw.') + + # generate pulsar likelihood objects which are common to all ArrayLikelihood objects + pulsar_likelihood_generator = (ds.PulsarLikelihood((psr.residuals, + ds.makenoise_measurement(psr, psr.noisedict), + ds.makegp_ecorr(psr, psr.noisedict), + ds.makegp_timing(psr, svd=True), + )) for psr in psrs) + + if correlation_orf is None: # intrinsic red noise only + + if spectrum.__name__ == 'powerlaw': + components = intrinsic_red_components + elif spectrum.__name__ == 'freespectrum': + components = {'log10_rho': intrinsic_red_components} + else: + raise ValueError('Power spectral density function not recognized.') + + print('Creating intrinsic red noise only model (IRN).') + al = ds.ArrayLikelihood(pulsar_likelihood_generator, + ds.makecommongp_fourier(psrs, spectrum, components=components, T=tspan, name='red_noise')) + + elif correlation_orf.__name__ == 'uncorrelated_orf': # CURN model + + if spectrum.__name__ == 'powerlaw': + if gamma_common is not None: + common_spectrum = ds.makepowerlaw_crn(common_components, crn_gamma=gamma_common) + gamma_common_name = [] + else: + common_spectrum = ds.makepowerlaw_crn(common_components) + gamma_common_name = ['crn_gamma'] + + common_names = ['crn_log10_A'] + gamma_common_name + + elif spectrum.__name__ == 'freespectrum': + common_spectrum = ds.makefreespectrum_crn(common_components) + components = {'log10_rho': intrinsic_red_components, 'crn_log10_rho': common_components} + + else: + raise ValueError('Power spectral density function not recognized.') + + print('Creating uncorrelated red noise model (CURN).') + al = ds.ArrayLikelihood(pulsar_likelihood_generator, + ds.makecommongp_fourier(psrs, common_spectrum, intrinsic_red_components, T=tspan, common=common_names, name='red_noise')) + + else: + + if spectrum.__name__ == 'powerlaw': + if gamma_common is not None: + common_spectrum = partial(spectrum, gamma=gamma_common) + gamma_common_name = [] + else: + common_spectrum = spectrum + gamma_common_name = ['gw_gamma'] + + common_names = ['gw_log10_A'] + gamma_common_name + + print(f'Creating {correlation_orf.__name__} correlated red noise model.') + al = ds.ArrayLikelihood(pulsar_likelihood_generator, + ds.makecommongp_fourier(psrs, spectrum, intrinsic_red_components, T=tspan, name='red_noise'), + ds.makegp_fourier_global(psrs, spectrum, correlation_orf, common_components, T=tspan, name='gw')) + + # TODO: Implement free spectrum model for correlated models in ArrayLikelihood + # elif spectrum.__name__ == 'freespectrum': + # common_spectrum = spectrum + # intrinsic_components = {'log10_rho': intrinsic_red_components} + # global_components = {'log10_rho': intrinsic_red_components, 'gw_log10_rho': common_components} + + # print(f'Creating {correlation_orf.__name__} correlated red noise model.') + # al = ds.ArrayLikelihood(pulsar_likelihood_generator, + # ds.makecommongp_fourier(psrs, spectrum, intrinsic_components, T=tspan, name='red_noise'), + # ds.makegp_fourier_global(psrs, spectrum, correlation_orf, global_components, T=tspan, name='gw')) + + else: + raise ValueError('Power spectral density function not recognized.') + + return al