From 5d2967ae62403f24fd5f1fd2f3225672f06a5045 Mon Sep 17 00:00:00 2001 From: stevertaylor Date: Fri, 23 Aug 2024 20:27:48 -0500 Subject: [PATCH 1/7] adding new models module --- src/discovery/models.py | 51 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 51 insertions(+) create mode 100644 src/discovery/models.py diff --git a/src/discovery/models.py b/src/discovery/models.py new file mode 100644 index 0000000..daef476 --- /dev/null +++ b/src/discovery/models.py @@ -0,0 +1,51 @@ +from jax.tree_util import Partial +import discovery as ds + +def lhood_maker(psrs, noisedict=None, gamma_common=None, + red_components=30, common_components=14, + common_type='curn', array_like=False): + """ + Construct discovery likelihood object from list of pulsar objects. + Parameters: + - psrs (list): List of pulsar objects. + - noisedict (dict, optional): Dictionary containing noise properties for each pulsar. Defaults to None. + - gamma_common (float, optional): Common red noise spectral index. Defaults to None. + - red_components (int, optional): Number of red noise components. Defaults to 30. + - common_components (int, optional): Number of common noise components. Defaults to 14. + - common_type (str, optional): Type of common noise model. Defaults to 'curn'. + - array_like (bool, optional): Whether to implement `batched` GPs (experimental). Defaults to False. [Not implemented yet] + + Returns: + - gl (object): Discovery likelihood object. + """ + + tspan = ds.getspan(psrs) + + if gamma_common is not None: + common_powerlaw = Partial(ds.powerlaw, gamma=gamma_common) + gamma_common_name = [] + else: + common_powerlaw = ds.powerlaw + gamma_common_name = ['gw_gamma'] + + if common_type == 'curn': + gl = ds.GlobalLikelihood((ds.PulsarLikelihood([psr.residuals, + ds.makenoise_measurement(psr, noisedict), + ds.makegp_ecorr(psr, noisedict), + ds.makegp_timing(psr, svd=True), + ds.makegp_fourier(psr, ds.powerlaw, red_components, T=tspan, name='red_noise'), + ds.makegp_fourier(psr, common_powerlaw, common_components, T=tspan, + common=['gw_log10_A']+gamma_common_name, name='gw') + ]) for psr in psrs)) + elif common_type == 'hd': + gl = ds.GlobalLikelihood((ds.PulsarLikelihood([psr.residuals, + ds.makenoise_measurement(psr, noisedict), + ds.makegp_ecorr(psr, noisedict), + ds.makegp_timing(psr, svd=True), + ds.makegp_fourier(psr, ds.powerlaw, red_components, T=tspan, name="red_noise") + ]) for psr in psrs), + ds.makegp_fourier_global(psrs, common_powerlaw, + ds.hd_orf, common_components, + T=tspan, name="gw")) + + return gl \ No newline at end of file From ed1633031f5e4f68c09d4384c7ff9271399a24f7 Mon Sep 17 00:00:00 2001 From: stevertaylor Date: Sat, 24 Aug 2024 11:22:40 -0500 Subject: [PATCH 2/7] feat: adding a samplers module with a basic wrapper script for ptmcmc --- src/discovery/samplers/ptmcmc_wrapper.py | 278 +++++++++++++++++++++++ 1 file changed, 278 insertions(+) create mode 100644 src/discovery/samplers/ptmcmc_wrapper.py diff --git a/src/discovery/samplers/ptmcmc_wrapper.py b/src/discovery/samplers/ptmcmc_wrapper.py new file mode 100644 index 0000000..b18fcf8 --- /dev/null +++ b/src/discovery/samplers/ptmcmc_wrapper.py @@ -0,0 +1,278 @@ +import os, jax +import numpy as np + +import discovery as ds +from PTMCMCSampler import PTMCMCSampler as ptmcmc + + +class JumpProposal(object): + + def __init__(self, param_names, empirical_distr=None, + save_ext_dists=False, outdir='./chains'): + """Set up some custom jump proposals""" + + self.params = param_names + + # parameter map + self.pmap = {} + ct = 0 + for p in self.params: + self.pmap[str(p)] = slice(ct, ct+1) + ct += 1 + + def draw_from_prior(self, x, iter, beta): + """Draw a sample from the prior distribution. + The function signature is specific to PTMCMCSampler. + + Parameters: + - x: The current state of the chain. + - iter: The current iteration number. + - beta: The current inverse temperature. + + Returns: + - q: The new state drawn from the prior distribution. + - lqxy: The log probability ratio of the forward-backward jump. + + """ + + q = x.copy() + lqxy = 0 + + # randomly choose parameter + param = np.random.choice(self.params) + + q[self.pmap[str(param)]] = list(ds.prior.sample_uniform([param]).values()) + + # forward-backward jump probability + lqxy += 0 + + return q, float(lqxy) + + def draw_from_red_prior(self, x, iter, beta): + """ + Draw a sample from the red noise prior distribution. + + Parameters: + - x: numpy.ndarray + The current state of the parameters. + - iter: int + The current iteration number. + - beta: float + The inverse temperature parameter. + + Returns: + - q: numpy.ndarray + The new state of the parameters after drawing from the red noise prior. + - lqxy: float + The log of the forward-backward jump probability ratio. + + """ + + q = x.copy() + lqxy = 0 + + signal_name = 'red_noise' + red_pars = [p for p in self.params if signal_name in p] + + # draw parameter from signal model + param = np.random.choice(red_pars) + q[self.pmap[str(param)]] = list(ds.prior.sample_uniform([param]).values()) + + # forward-backward jump probability + lqxy += 0 + + return q, float(lqxy) + + def draw_from_gwb_log_uniform_distribution(self, x, iter, beta): + """ + Draws a sample from the log-uniform distribution for the GWB. + + Parameters: + - x: The current state of the parameters. + - iter: The current iteration number. + - beta: The inverse temperature parameter. + + Returns: + - q: The new state of the parameters after the draw. + - lqxy: The log of the forward-backward jump probability ratio. + """ + + q = x.copy() + lqxy = 0 + + # draw parameter from signal model + param = [p for p in self.params + if ('gw' in p and 'log10_A' in p)][0] + + q[self.pmap[str(param)]] = list(ds.prior.sample_uniform([param]).values()) + + # forward-backward jump probability + lqxy += 0 + + return q, float(lqxy) + + +class InferenceModel(object): + """ + A class representing an inference model. + Parameters: + - model: The model object. + - pnames: A list of parameter names. + + Methods: + - __init__(self, model, pnames=None): Initializes the InferenceModel object. + - x2p(self, x): Converts a list of values to a dictionary of parameter-value pairs. + - p2x(self, p): Converts a dictionary of parameter-value pairs to a list of values. + - get_parameter_groups(self): Utility function to get parameter groupings for sampling. + - setup_sampler(self, outdir='chains', resume=False, empirical_distr=None, groups=None, loglkwargs={}, logpkwargs={}): Sets up the sampler for MCMC sampling. + """ + + def __init__(self, model, pnames=None): + """ + Initializes an instance of the `disco_ptmcmc` class. + Parameters: + - model: The model object. + - pnames: A list of pulsar names. + Returns: + None + """ + + self.pnames = pnames + self.param_names = model.logL.params + + loglike = model.logL + logprior = ds.prior.makelogprior_uniform(self.param_names, + ds.prior.priordict_standard) + + jlogl = jax.jit(loglike) + jlogp = jax.jit(logprior) + + self.get_lnlikelihood = lambda x: float(jlogl(self.x2p(x))) + self.get_lnprior = lambda x: float(jlogp(self.x2p(x))) + + def x2p(self, x): + """ + Converts a list of parameter values `x` to a dictionary representation. + + Args: + x (list): A list of parameter values. + + Returns: + dict: A dictionary representation of the parameter values, where the keys are the parameter names and the values are the corresponding values from `x`. + """ + # does not handle vector parameters + return {par: val for par, val in zip(self.param_names, x)} + + def p2x(self, p): + """ + Convert a dictionary of values to a NumPy array. + + Parameters: + p (dict): A dictionary containing values. + + Returns: + numpy.ndarray: A NumPy array containing the values from the dictionary. + """ + return np.array(list(p.values()), 'd') + + def get_parameter_groups(self): + """Utility function to get parameter groupings for sampling.""" + params = self.param_names + ndim = len(params) + groups = [list(np.arange(0, ndim))] + + # get global and individual parameters + gpars = [p for p in params if 'gw' in p or 'curn' in p] + ipars = [p for p in params if p not in gpars] + if gpars: + # add a group of all global parameters + groups.append([params.index(gp) for gp in gpars]) + if ipars: + #groups.append([params.index(ip) for ip in ipars]) + groups += [[params.index(ip) for ip in ipars if p in ip] + for p in self.pnames] + + return groups + + def setup_sampler(self, outdir='chains', resume=False, + empirical_distr=None, groups=None, + loglkwargs={}, logpkwargs={}): + """ + Set up the sampler for performing MCMC (Markov Chain Monte Carlo) sampling. + + Parameters: + - outdir (str): The output directory for saving the chains. Default is 'chains'. + - resume (bool): Whether to resume from a previous run. Default is False. + - empirical_distr (None or str): The path to the empirical distribution file. Default is None. + - groups (None or list): The parameter groups for the sampler. Default is None. + - loglkwargs (dict): Additional keyword arguments for the log-likelihood function. Default is an empty dictionary. + - logpkwargs (dict): Additional keyword arguments for the log-prior function. Default is an empty dictionary. + + Returns: + - sampler: The initialized PTMCMCSampler object. + """ + + # dimension of parameter space + ndim = len(self.param_names) + + # initial jump covariance matrix + if os.path.exists(outdir+'/cov.npy') and resume: + cov = np.load(outdir+'/cov.npy') + + # check that the one we load is the same shape as our data + cov_new = np.diag(np.ones(ndim) * 1.0**2) + if cov.shape != cov_new.shape: + msg = 'The covariance matrix (cov.npy) in the output folder is ' + msg += 'the wrong shape for the parameters given. ' + msg += 'Start with a different output directory or ' + msg += 'change resume to False to overwrite the run that exists.' + + raise ValueError(msg) + else: + cov = np.diag(np.ones(ndim) * 1.0**2) # used to be 0.1 + + # parameter groupings + if groups is None: + groups = self.get_parameter_groups() + + sampler = ptmcmc.PTSampler(ndim, self.get_lnlikelihood, self.get_lnprior, cov, + groups=groups, outDir=outdir, resume=resume, + loglkwargs=loglkwargs, logpkwargs=logpkwargs) + + # additional jump proposals + jp = JumpProposal(param_names=self.param_names, empirical_distr=None, + save_ext_dists=False, outdir=outdir) + sampler.jp = jp + + # always add draw from prior + sampler.addProposalToCycle(jp.draw_from_prior, 5) + + # try adding empirical proposals + #if empirical_distr is not None: + # print('Adding empirical proposals...\n') + # sampler.addProposalToCycle(jp.draw_from_empirical_distr, 25) + + # Red noise prior draw + if any('red_noise' in s for s in self.param_names): + print('Adding red noise prior draws...\n') + sampler.addProposalToCycle(jp.draw_from_red_prior, 10) + + # GWB uniform distribution draw + if np.any([('gw' in par and 'log10_A' in par) for par in self.param_names]): + print('Adding GWB uniform distribution draws...\n') + sampler.addProposalToCycle(jp.draw_from_gwb_log_uniform_distribution, 10) + + # free spectrum prior draw + #if np.any(['log10_rho' in par for par in self.param_names]): + # print('Adding free spectrum prior draws...\n') + # sampler.addProposalToCycle(jp.draw_from_gw_rho_prior, 25) + + # Prior distribution draw for parameters named GW + #if any([str(p).split(':')[0] for p in list(self.params) if 'gw' in str(p)]): + # print('Adding gw param prior draws...\n') + # sampler.addProposalToCycle(jp.draw_from_par_prior( + # par_names=[str(p).split(':')[0] for + # p in list(self.params) + # if 'gw' in str(p)]), 10) + + return sampler \ No newline at end of file From 7aed3c07d5984cd5158483976986c24d82c2b9f3 Mon Sep 17 00:00:00 2001 From: stevertaylor Date: Sat, 24 Aug 2024 11:32:24 -0500 Subject: [PATCH 3/7] Revert "feat: adding a samplers module with a basic wrapper script for ptmcmc" This reverts commit ed1633031f5e4f68c09d4384c7ff9271399a24f7. --- src/discovery/samplers/ptmcmc_wrapper.py | 278 ----------------------- 1 file changed, 278 deletions(-) delete mode 100644 src/discovery/samplers/ptmcmc_wrapper.py diff --git a/src/discovery/samplers/ptmcmc_wrapper.py b/src/discovery/samplers/ptmcmc_wrapper.py deleted file mode 100644 index b18fcf8..0000000 --- a/src/discovery/samplers/ptmcmc_wrapper.py +++ /dev/null @@ -1,278 +0,0 @@ -import os, jax -import numpy as np - -import discovery as ds -from PTMCMCSampler import PTMCMCSampler as ptmcmc - - -class JumpProposal(object): - - def __init__(self, param_names, empirical_distr=None, - save_ext_dists=False, outdir='./chains'): - """Set up some custom jump proposals""" - - self.params = param_names - - # parameter map - self.pmap = {} - ct = 0 - for p in self.params: - self.pmap[str(p)] = slice(ct, ct+1) - ct += 1 - - def draw_from_prior(self, x, iter, beta): - """Draw a sample from the prior distribution. - The function signature is specific to PTMCMCSampler. - - Parameters: - - x: The current state of the chain. - - iter: The current iteration number. - - beta: The current inverse temperature. - - Returns: - - q: The new state drawn from the prior distribution. - - lqxy: The log probability ratio of the forward-backward jump. - - """ - - q = x.copy() - lqxy = 0 - - # randomly choose parameter - param = np.random.choice(self.params) - - q[self.pmap[str(param)]] = list(ds.prior.sample_uniform([param]).values()) - - # forward-backward jump probability - lqxy += 0 - - return q, float(lqxy) - - def draw_from_red_prior(self, x, iter, beta): - """ - Draw a sample from the red noise prior distribution. - - Parameters: - - x: numpy.ndarray - The current state of the parameters. - - iter: int - The current iteration number. - - beta: float - The inverse temperature parameter. - - Returns: - - q: numpy.ndarray - The new state of the parameters after drawing from the red noise prior. - - lqxy: float - The log of the forward-backward jump probability ratio. - - """ - - q = x.copy() - lqxy = 0 - - signal_name = 'red_noise' - red_pars = [p for p in self.params if signal_name in p] - - # draw parameter from signal model - param = np.random.choice(red_pars) - q[self.pmap[str(param)]] = list(ds.prior.sample_uniform([param]).values()) - - # forward-backward jump probability - lqxy += 0 - - return q, float(lqxy) - - def draw_from_gwb_log_uniform_distribution(self, x, iter, beta): - """ - Draws a sample from the log-uniform distribution for the GWB. - - Parameters: - - x: The current state of the parameters. - - iter: The current iteration number. - - beta: The inverse temperature parameter. - - Returns: - - q: The new state of the parameters after the draw. - - lqxy: The log of the forward-backward jump probability ratio. - """ - - q = x.copy() - lqxy = 0 - - # draw parameter from signal model - param = [p for p in self.params - if ('gw' in p and 'log10_A' in p)][0] - - q[self.pmap[str(param)]] = list(ds.prior.sample_uniform([param]).values()) - - # forward-backward jump probability - lqxy += 0 - - return q, float(lqxy) - - -class InferenceModel(object): - """ - A class representing an inference model. - Parameters: - - model: The model object. - - pnames: A list of parameter names. - - Methods: - - __init__(self, model, pnames=None): Initializes the InferenceModel object. - - x2p(self, x): Converts a list of values to a dictionary of parameter-value pairs. - - p2x(self, p): Converts a dictionary of parameter-value pairs to a list of values. - - get_parameter_groups(self): Utility function to get parameter groupings for sampling. - - setup_sampler(self, outdir='chains', resume=False, empirical_distr=None, groups=None, loglkwargs={}, logpkwargs={}): Sets up the sampler for MCMC sampling. - """ - - def __init__(self, model, pnames=None): - """ - Initializes an instance of the `disco_ptmcmc` class. - Parameters: - - model: The model object. - - pnames: A list of pulsar names. - Returns: - None - """ - - self.pnames = pnames - self.param_names = model.logL.params - - loglike = model.logL - logprior = ds.prior.makelogprior_uniform(self.param_names, - ds.prior.priordict_standard) - - jlogl = jax.jit(loglike) - jlogp = jax.jit(logprior) - - self.get_lnlikelihood = lambda x: float(jlogl(self.x2p(x))) - self.get_lnprior = lambda x: float(jlogp(self.x2p(x))) - - def x2p(self, x): - """ - Converts a list of parameter values `x` to a dictionary representation. - - Args: - x (list): A list of parameter values. - - Returns: - dict: A dictionary representation of the parameter values, where the keys are the parameter names and the values are the corresponding values from `x`. - """ - # does not handle vector parameters - return {par: val for par, val in zip(self.param_names, x)} - - def p2x(self, p): - """ - Convert a dictionary of values to a NumPy array. - - Parameters: - p (dict): A dictionary containing values. - - Returns: - numpy.ndarray: A NumPy array containing the values from the dictionary. - """ - return np.array(list(p.values()), 'd') - - def get_parameter_groups(self): - """Utility function to get parameter groupings for sampling.""" - params = self.param_names - ndim = len(params) - groups = [list(np.arange(0, ndim))] - - # get global and individual parameters - gpars = [p for p in params if 'gw' in p or 'curn' in p] - ipars = [p for p in params if p not in gpars] - if gpars: - # add a group of all global parameters - groups.append([params.index(gp) for gp in gpars]) - if ipars: - #groups.append([params.index(ip) for ip in ipars]) - groups += [[params.index(ip) for ip in ipars if p in ip] - for p in self.pnames] - - return groups - - def setup_sampler(self, outdir='chains', resume=False, - empirical_distr=None, groups=None, - loglkwargs={}, logpkwargs={}): - """ - Set up the sampler for performing MCMC (Markov Chain Monte Carlo) sampling. - - Parameters: - - outdir (str): The output directory for saving the chains. Default is 'chains'. - - resume (bool): Whether to resume from a previous run. Default is False. - - empirical_distr (None or str): The path to the empirical distribution file. Default is None. - - groups (None or list): The parameter groups for the sampler. Default is None. - - loglkwargs (dict): Additional keyword arguments for the log-likelihood function. Default is an empty dictionary. - - logpkwargs (dict): Additional keyword arguments for the log-prior function. Default is an empty dictionary. - - Returns: - - sampler: The initialized PTMCMCSampler object. - """ - - # dimension of parameter space - ndim = len(self.param_names) - - # initial jump covariance matrix - if os.path.exists(outdir+'/cov.npy') and resume: - cov = np.load(outdir+'/cov.npy') - - # check that the one we load is the same shape as our data - cov_new = np.diag(np.ones(ndim) * 1.0**2) - if cov.shape != cov_new.shape: - msg = 'The covariance matrix (cov.npy) in the output folder is ' - msg += 'the wrong shape for the parameters given. ' - msg += 'Start with a different output directory or ' - msg += 'change resume to False to overwrite the run that exists.' - - raise ValueError(msg) - else: - cov = np.diag(np.ones(ndim) * 1.0**2) # used to be 0.1 - - # parameter groupings - if groups is None: - groups = self.get_parameter_groups() - - sampler = ptmcmc.PTSampler(ndim, self.get_lnlikelihood, self.get_lnprior, cov, - groups=groups, outDir=outdir, resume=resume, - loglkwargs=loglkwargs, logpkwargs=logpkwargs) - - # additional jump proposals - jp = JumpProposal(param_names=self.param_names, empirical_distr=None, - save_ext_dists=False, outdir=outdir) - sampler.jp = jp - - # always add draw from prior - sampler.addProposalToCycle(jp.draw_from_prior, 5) - - # try adding empirical proposals - #if empirical_distr is not None: - # print('Adding empirical proposals...\n') - # sampler.addProposalToCycle(jp.draw_from_empirical_distr, 25) - - # Red noise prior draw - if any('red_noise' in s for s in self.param_names): - print('Adding red noise prior draws...\n') - sampler.addProposalToCycle(jp.draw_from_red_prior, 10) - - # GWB uniform distribution draw - if np.any([('gw' in par and 'log10_A' in par) for par in self.param_names]): - print('Adding GWB uniform distribution draws...\n') - sampler.addProposalToCycle(jp.draw_from_gwb_log_uniform_distribution, 10) - - # free spectrum prior draw - #if np.any(['log10_rho' in par for par in self.param_names]): - # print('Adding free spectrum prior draws...\n') - # sampler.addProposalToCycle(jp.draw_from_gw_rho_prior, 25) - - # Prior distribution draw for parameters named GW - #if any([str(p).split(':')[0] for p in list(self.params) if 'gw' in str(p)]): - # print('Adding gw param prior draws...\n') - # sampler.addProposalToCycle(jp.draw_from_par_prior( - # par_names=[str(p).split(':')[0] for - # p in list(self.params) - # if 'gw' in str(p)]), 10) - - return sampler \ No newline at end of file From fb3b5ed9ec0784c0c5f59ed6a034a556a1e49b4f Mon Sep 17 00:00:00 2001 From: stevertaylor Date: Mon, 26 Aug 2024 14:45:18 -0500 Subject: [PATCH 4/7] updating function name to make_likelihood --- src/discovery/models.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/src/discovery/models.py b/src/discovery/models.py index daef476..60afa35 100644 --- a/src/discovery/models.py +++ b/src/discovery/models.py @@ -1,7 +1,7 @@ from jax.tree_util import Partial import discovery as ds -def lhood_maker(psrs, noisedict=None, gamma_common=None, +def make_likelihood(psrs, noisedict=None, gamma_common=None, red_components=30, common_components=14, common_type='curn', array_like=False): """ @@ -18,7 +18,7 @@ def lhood_maker(psrs, noisedict=None, gamma_common=None, Returns: - gl (object): Discovery likelihood object. """ - + tspan = ds.getspan(psrs) if gamma_common is not None: @@ -34,7 +34,7 @@ def lhood_maker(psrs, noisedict=None, gamma_common=None, ds.makegp_ecorr(psr, noisedict), ds.makegp_timing(psr, svd=True), ds.makegp_fourier(psr, ds.powerlaw, red_components, T=tspan, name='red_noise'), - ds.makegp_fourier(psr, common_powerlaw, common_components, T=tspan, + ds.makegp_fourier(psr, common_powerlaw, common_components, T=tspan, common=['gw_log10_A']+gamma_common_name, name='gw') ]) for psr in psrs)) elif common_type == 'hd': @@ -44,8 +44,8 @@ def lhood_maker(psrs, noisedict=None, gamma_common=None, ds.makegp_timing(psr, svd=True), ds.makegp_fourier(psr, ds.powerlaw, red_components, T=tspan, name="red_noise") ]) for psr in psrs), - ds.makegp_fourier_global(psrs, common_powerlaw, - ds.hd_orf, common_components, + ds.makegp_fourier_global(psrs, common_powerlaw, + ds.hd_orf, common_components, T=tspan, name="gw")) return gl \ No newline at end of file From 644d08bb72e84387c948cf9c48c131b99deb2c42 Mon Sep 17 00:00:00 2001 From: Aaron Date: Fri, 3 Jan 2025 12:38:58 -0800 Subject: [PATCH 5/7] Move to ArrayLikelihoods --- src/discovery/models.py | 120 ++++++++++++++++++++++++++-------------- 1 file changed, 79 insertions(+), 41 deletions(-) diff --git a/src/discovery/models.py b/src/discovery/models.py index 60afa35..27e2827 100644 --- a/src/discovery/models.py +++ b/src/discovery/models.py @@ -1,51 +1,89 @@ -from jax.tree_util import Partial +from functools import partial import discovery as ds +from typing import Optional, Callable + + +def make_array_likelihood( + psrs: list[ds.Pulsar], + gamma_common: Optional[float]=None, + correlation_orf: Optional[Callable]=None, + intrinsic_red_components: int=30, + common_components: int=14, + spectrum: Callable=ds.powerlaw, + ) -> ds.ArrayLikelihood: + """Construct discovery array likelihood object from a list of Pulsar objects which may contain a noise dictionary. -def make_likelihood(psrs, noisedict=None, gamma_common=None, - red_components=30, common_components=14, - common_type='curn', array_like=False): - """ - Construct discovery likelihood object from list of pulsar objects. Parameters: - - psrs (list): List of pulsar objects. - - noisedict (dict, optional): Dictionary containing noise properties for each pulsar. Defaults to None. - - gamma_common (float, optional): Common red noise spectral index. Defaults to None. - - red_components (int, optional): Number of red noise components. Defaults to 30. - - common_components (int, optional): Number of common noise components. Defaults to 14. - - common_type (str, optional): Type of common noise model. Defaults to 'curn'. - - array_like (bool, optional): Whether to implement `batched` GPs (experimental). Defaults to False. [Not implemented yet] + psrs (list): List of pulsar objects + gamma_common (float) [None]: Common red noise spectral index + 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 + intrinsic_red_components (int) [30]: Number of individual pulsar red noise components + common_components (int) [14]: Number of common red noise components + spectrum (function): Spectrum function for all red noise parameters + Options include `ds.powerlaw`, `ds.freespectrum`, or a custom spectrum function + tnequad (bool): Use temponest measurement noise definition Returns: - - gl (object): Discovery likelihood object. + array_likelihood (object): Discovery ArrayLikelihood object """ tspan = ds.getspan(psrs) - if gamma_common is not None: - common_powerlaw = Partial(ds.powerlaw, gamma=gamma_common) - gamma_common_name = [] + if spectrum.__name__ == 'powerlaw': # powerlaw pieces + if gamma_common is not None and correlation_orf is None: + print('Warning: gamma_common is set but correlation_orf is None. Setting gamma_common to None.') + gamma_common_name = [] + common_spectrum = None # just use spectrum in this case + elif gamma_common is not None: + print(f'Creating powerlaw spectrum red noise model with gamma={gamma_common}.') + common_spectrum = ds.makepowerlaw_crn(common_components, crn_gamma=gamma_common) + gamma_common_name = [] + else: + print('Creating powerlaw spectrum red noise model with varied gamma.') + common_spectrum = ds.makepowerlaw_crn(common_components) + gamma_common_name = ['crn_gamma'] + + common_names = ['crn_log10_A'] + gamma_common_name + components = intrinsic_red_components + + elif spectrum.__name__ == 'freespectrum': # free spectrum pieces + if correlation_orf is None: + common_spectrum = ds.makefreespectrum_crn(common_components) + components = {'log10_rho': intrinsic_red_components} + else: + common_spectrum = ds.makefreespectrum_crn(common_components) + components = {'log10_rho': intrinsic_red_components, 'crn_log10_rho': common_components} + + common_names = ['crn_log10_rho'] + else: - common_powerlaw = ds.powerlaw - gamma_common_name = ['gw_gamma'] - - if common_type == 'curn': - gl = ds.GlobalLikelihood((ds.PulsarLikelihood([psr.residuals, - ds.makenoise_measurement(psr, noisedict), - ds.makegp_ecorr(psr, noisedict), - ds.makegp_timing(psr, svd=True), - ds.makegp_fourier(psr, ds.powerlaw, red_components, T=tspan, name='red_noise'), - ds.makegp_fourier(psr, common_powerlaw, common_components, T=tspan, - common=['gw_log10_A']+gamma_common_name, name='gw') - ]) for psr in psrs)) - elif common_type == 'hd': - gl = ds.GlobalLikelihood((ds.PulsarLikelihood([psr.residuals, - ds.makenoise_measurement(psr, noisedict), - ds.makegp_ecorr(psr, noisedict), - ds.makegp_timing(psr, svd=True), - ds.makegp_fourier(psr, ds.powerlaw, red_components, T=tspan, name="red_noise") - ]) for psr in psrs), - ds.makegp_fourier_global(psrs, common_powerlaw, - ds.hd_orf, common_components, - T=tspan, name="gw")) - - return gl \ No newline at end of file + raise ValueError('Spectrum function not recognized.') + + # correlation block + if correlation_orf is None: # intrinsic red noise only + print('Creating intrinsic red noise only model (IRN).') + al = ds.ArrayLikelihood((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), + ds.makecommongp_fourier(psrs, spectrum, components=components, T=tspan, name='red_noise')) + elif correlation_orf.__name__ == 'uncorrelated_orf': + print('Creating uncorrelated red noise model (CURN).') + al = ds.ArrayLikelihood((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), + ds.makecommongp_fourier(psrs, common_spectrum, intrinsic_red_components, T=tspan, common=common_names, name='red_noise')) + else: + print(f'Creating {correlation_orf.__name__} correlated red noise model.') + al = ds.ArrayLikelihood((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), + ds.makecommongp_fourier(psrs, spectrum, intrinsic_red_components, T=tspan, name='red_noise'), + ds.makegp_fourier_global(psrs, spectrum, correlation_orf, components, T=tspan, name='gw')) + + return al From 60e53dbc694222dd17ad573dfd861d6a6947d669 Mon Sep 17 00:00:00 2001 From: Aaron Date: Fri, 3 Jan 2025 16:17:41 -0800 Subject: [PATCH 6/7] add partial to global --- src/discovery/matrix.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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))) From ecdf16c7e7832623d89316fc76107174a8f9d6de Mon Sep 17 00:00:00 2001 From: Aaron Date: Sat, 4 Jan 2025 01:29:03 -0800 Subject: [PATCH 7/7] IRN, CURN, correlated ArrayLikelihood models added --- src/discovery/models.py | 122 +++++++++++++++++++++++----------------- 1 file changed, 71 insertions(+), 51 deletions(-) diff --git a/src/discovery/models.py b/src/discovery/models.py index 27e2827..e279663 100644 --- a/src/discovery/models.py +++ b/src/discovery/models.py @@ -5,24 +5,23 @@ def make_array_likelihood( psrs: list[ds.Pulsar], - gamma_common: Optional[float]=None, correlation_orf: Optional[Callable]=None, + spectrum: Callable=ds.powerlaw, intrinsic_red_components: int=30, common_components: int=14, - spectrum: Callable=ds.powerlaw, + 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 - gamma_common (float) [None]: Common red noise spectral index 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 - intrinsic_red_components (int) [30]: Number of individual pulsar red noise components - common_components (int) [14]: Number of common red noise components spectrum (function): Spectrum function for all red noise parameters Options include `ds.powerlaw`, `ds.freespectrum`, or a custom spectrum function - tnequad (bool): Use temponest measurement noise definition + 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 @@ -30,60 +29,81 @@ def make_array_likelihood( tspan = ds.getspan(psrs) - if spectrum.__name__ == 'powerlaw': # powerlaw pieces - if gamma_common is not None and correlation_orf is None: - print('Warning: gamma_common is set but correlation_orf is None. Setting gamma_common to None.') - gamma_common_name = [] - common_spectrum = None # just use spectrum in this case - elif gamma_common is not None: - print(f'Creating powerlaw spectrum red noise model with gamma={gamma_common}.') - common_spectrum = ds.makepowerlaw_crn(common_components, crn_gamma=gamma_common) - gamma_common_name = [] - else: - print('Creating powerlaw spectrum red noise model with varied gamma.') - common_spectrum = ds.makepowerlaw_crn(common_components) - gamma_common_name = ['crn_gamma'] + 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.') - common_names = ['crn_log10_A'] + gamma_common_name - components = intrinsic_red_components + # 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) - elif spectrum.__name__ == 'freespectrum': # free spectrum pieces - if correlation_orf is None: - common_spectrum = ds.makefreespectrum_crn(common_components) + 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} - common_names = ['crn_log10_rho'] - - else: - raise ValueError('Spectrum function not recognized.') + else: + raise ValueError('Power spectral density function not recognized.') - # correlation block - if correlation_orf is None: # intrinsic red noise only - print('Creating intrinsic red noise only model (IRN).') - al = ds.ArrayLikelihood((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), - ds.makecommongp_fourier(psrs, spectrum, components=components, T=tspan, name='red_noise')) - elif correlation_orf.__name__ == 'uncorrelated_orf': print('Creating uncorrelated red noise model (CURN).') - al = ds.ArrayLikelihood((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), - ds.makecommongp_fourier(psrs, common_spectrum, intrinsic_red_components, T=tspan, common=common_names, name='red_noise')) + al = ds.ArrayLikelihood(pulsar_likelihood_generator, + ds.makecommongp_fourier(psrs, common_spectrum, intrinsic_red_components, T=tspan, common=common_names, name='red_noise')) + else: - print(f'Creating {correlation_orf.__name__} correlated red noise model.') - al = ds.ArrayLikelihood((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), - ds.makecommongp_fourier(psrs, spectrum, intrinsic_red_components, T=tspan, name='red_noise'), - ds.makegp_fourier_global(psrs, spectrum, correlation_orf, components, T=tspan, name='gw')) + + 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