From 66ce1240aaefd56068f313bd8aa6cbe60418ad30 Mon Sep 17 00:00:00 2001 From: Eshel Dror Date: Tue, 6 Oct 2026 15:55:37 -0700 Subject: [PATCH 01/21] Implement null marginalization framework for MCMC --- orbitize/marginalizer.py | 41 +++++++++++ orbitize/sampler.py | 145 +++++++++++++++++++++++++++++++-------- tests/test_mcmc.py | 2 +- 3 files changed, 157 insertions(+), 31 deletions(-) create mode 100644 orbitize/marginalizer.py diff --git a/orbitize/marginalizer.py b/orbitize/marginalizer.py new file mode 100644 index 00000000..f9c01363 --- /dev/null +++ b/orbitize/marginalizer.py @@ -0,0 +1,41 @@ +import abc + +""" +This module defines marginalizations to draw +samples of non-marginalized parameters and compute appropriate probability +""" + +class Marginalizer(abc.ABC): + + QUANT_TYPES = [] + JITTER_PARAMS = [] + REQ_PARAMS = [] + MARG_PARAMS = [] + LNLIKE_TYPE = None + MULTIBODY = False + + @abc.abstractmethod + def apply(self, param_arr, jitter_arr=None, calculated_values={}): + pass + +class Marg1(Marginalizer): + + REQ_PARAMS = ["tau", "per"] + MARG_PARAMS = ["sma", "pan"] + +class Marg2(Marginalizer): + + REQ_PARAMS = ["tau", "sma", "pan"] + MARG_PARAMS = ["gamma"] + +def compatible_params(marginalizers, input_params): + # TODO: Multibody, quant_type + params = set(input_params) + for marginalizer in marginalizers: + req_params = set(marginalizer.REQ_PARAMS) + marg_params = set(marginalizer.MARG_PARAMS) + if len(req_params - params) == 0 and len(marg_params & params) == 0: + params |= set(marginalizer.MARG_PARAMS) + else: + return False + return True \ No newline at end of file diff --git a/orbitize/sampler.py b/orbitize/sampler.py index 4b1d778e..3c1e2144 100644 --- a/orbitize/sampler.py +++ b/orbitize/sampler.py @@ -46,11 +46,47 @@ def __init__( # check if need to handle covariances self.has_corr = np.any(~np.isnan(self.system.data_table["quant12_corr"])) + # TODO: Construct + self.marginalizers = [] + + # TODO: multiplanet + marginalized_param_labels_set = {(label for label in marg) for marg in self.marginalizers} + fitted_param_labels_set = set(system.labels) - marginalized_param_labels_set + + self.fitted_param_labels = [] + self.marg_param_labels = [] + self.fitted_param_sys_idx = [] + self.marg_param_sys_idx = [] + self.fitted_priors = [] + self.marg_priors = [] + self.fitted_param_idx = {} + self.marg_param_idx = {} + fitted_param_counter = 0 + marg_param_counter = 0 + for i, (label, prior) in enumerate(zip(self.system.labels, self.system.sys_priors)): + if label in fitted_param_labels_set: + self.fitted_priors.append(prior) + self.fitted_param_labels.append(label) + self.fitted_param_sys_idx.append(i) + self.fitted_param_idx[label] = fitted_param_counter + fitted_param_counter += 1 + else: + self.marg_priors.append(prior) + self.marg_param_labels.append(label) + self.marg_param_sys_idx.append(i) + self.marg_param_idx[label] = marg_param_counter + marg_param_counter += 1 + self.total_param_num = len(self.system.sys_priors) + quants = set(np.unique(self.system.data_table["quant_type"])) + marg_quants = {(quant_type for quant_type in marg) for marg in self.marginalizers} + + self.non_marg_quants = list(quants - marg_quants) + @abc.abstractmethod def run_sampler(self, total_orbits): pass - def _logl(self, params): + def _logl(self, fitted_params): """ log likelihood function that interfaces with the orbitize objects Comptues the sum of the log likelihoods of the data given the input model @@ -66,35 +102,77 @@ def _logl(self, params): float: sum of all log likelihoods of the data given input model """ - # compute the model based on system params - model, jitter = self.system.compute_model(params) - - # fold data/errors to match model output shape. In particualr, quant1/quant2 are interleaved - data = np.array( - [self.system.data_table["quant1"], self.system.data_table["quant2"]] - ).T - - # errors below required for lnlike function below - errs = np.array( - [self.system.data_table["quant1_err"], self.system.data_table["quant2_err"]] - ).T - # covariances/correlations, if applicable - # we're doing this check now because the likelihood computation is much faster if we can skip it. - if self.has_corr: - corrs = self.system.data_table["quant12_corr"] + input_shape = fitted_params.shape + assert input_shape[0] == len(self.fitted_param_sys_idx) + if len(input_shape) == 1: + num_orbits = 1 + params = np.full((self.total_param_num,), np.nan) else: - corrs = None - - # grab all seppa indices - seppa_indices = self.system.all_seppa + num_orbits = input_shape[1] + params = np.full((self.total_param_num, num_orbits), np.nan) + assert len(self.fitted_param_sys_idx) == len(fitted_params) == self.total_param_num, (self.fitted_param_sys_idx, fitted_params) + params[self.fitted_param_sys_idx] = fitted_params + + total_lnlike = 0 + total_lnprob = 0 + all_calculated_values = {} + for marg in self.marginalizers: + req_param_idx = [self.system.param_idx[label] for label in marg.REQ_PARAMS] + jitter_param_idx = [self.system.param_idx[label] for label in marg.JITTER_PARAMS] + marg_param_idx = [self.system.param_idx[label] for label in marg.MARG_PARAMS] + req_params = params[req_param_idx] + jitter = params[jitter_param_idx] + partial_lnprob, partial_lnlike, marg_params, calculated_values = marg.apply(req_params, jitter, all_calculated_values) + total_lnprob += partial_lnprob + total_lnlike += partial_lnlike + params[marg_param_idx] = marg_params + all_calculated_values.update(calculated_values) # i.e. astr_tanom, rv_eanom + + lnlikes_sum = 0 + if len(self.non_marg_quants) > 0: + + # compute the model based on system params + model, jitter = self.system.compute_model(params) + + # fold data/errors to match model output shape. In particualr, quant1/quant2 are interleaved + data = np.array( + [self.system.data_table["quant1"], self.system.data_table["quant2"]] + ).T + + # errors below required for lnlike function below + errs = np.array( + [self.system.data_table["quant1_err"], self.system.data_table["quant2_err"]] + ).T + # covariances/correlations, if applicable + # we're doing this check now because the likelihood computation is much faster if we can skip it. + if self.has_corr: + corrs = self.system.data_table["quant12_corr"] + else: + corrs = None - # compute lnlike - lnlikes = self.lnlike( - data, errs, corrs, model, jitter, seppa_indices, chi2_type=self.chi2_type - ) + # grab all seppa indices + if "seppa" in self.non_marg_quants: + seppa_indices = self.system.all_seppa + else: + seppa_indices = () + non_marg_mask = np.isin(np.astype(self.system.data_table["quant_type"], str), self.non_marg_quants) + non_marg_data = data[non_marg_mask] + non_marg_errs = errs[non_marg_mask] + if corrs is not None: + non_marg_corrs = corrs[non_marg_mask] + else: + non_marg_corrs = None + non_marg_model = model[non_marg_mask] + non_marg_jitter = jitter[non_marg_mask] + # assert False, (data, non_marg_data, non_marg_data.shape, non_marg_mask, self.non_marg_quants, self.system.data_table["quant_type"]) + + # compute lnlike + lnlikes = self.lnlike( + non_marg_data, non_marg_errs, non_marg_corrs, non_marg_model, non_marg_jitter, seppa_indices, chi2_type=self.chi2_type + ) - # return sum of lnlikes (aka product of likeliehoods) - lnlikes_sum = np.nansum(lnlikes, axis=(0, 1)) + # return sum of lnlikes (aka product of likeliehoods) + lnlikes_sum += np.nansum(lnlikes, axis=(0, 1)) if self.custom_lnlike is not None: lnlikes_sum += self.custom_lnlike(params) @@ -145,7 +223,12 @@ def _logl(self, params): params, self.system.param_idx, ) - return lnlikes_sum + total_lnprob += lnlikes_sum + total_lnlike += lnlikes_sum + + marg_params = params[self.marg_param_sys_idx] + + return total_lnprob, total_lnlike, *marg_params class OFTI( @@ -734,7 +817,7 @@ def __init__( self.sampled_param_idx = {} sampled_param_counter = 0 - for i, prior in enumerate(system.sys_priors): + for i, prior in enumerate(system.sys_priors): # TODO: use self.fitted_priors instead # check for fixed parameters if not hasattr(prior, "draw_samples"): self.fixed_params.append((i, prior)) @@ -846,11 +929,13 @@ def _logl(self, params, include_logp=False): else: logp = 0 # don't include prior + assert len(self.fixed_params) == 2, self.fixed_params full_params = self._fill_in_fixed_params(params) if np.ndim(full_params) == 2: full_params = full_params.T - return super(MCMC, self)._logl(full_params) + logp + logprob, loglike, *marg_params = super(MCMC, self)._logl(full_params) + return logprob + logp def _update_chains_from_sampler(self, sampler, num_steps=None): """ diff --git a/tests/test_mcmc.py b/tests/test_mcmc.py index 3f85c3ed..cf844d0d 100644 --- a/tests/test_mcmc.py +++ b/tests/test_mcmc.py @@ -80,7 +80,7 @@ def do_mcmc_runs(num_temps=0, num_threads=1, make_corner_plot=False): # test that lnlikes being saved are correct returned_lnlike_test = myDriver.sampler.results.lnlike[0] - computed_lnlike_test = myDriver.sampler._logl(myDriver.sampler.results.post[0]) + computed_lnlike_test, _, *_ = super(sampler.MCMC, myDriver.sampler)._logl(myDriver.sampler.results.post[0]) assert returned_lnlike_test == pytest.approx(computed_lnlike_test, abs=0.01) From 0d7af8dcbf5f222c93294e6b74541ca75f8544cf Mon Sep 17 00:00:00 2001 From: Eshel Dror Date: Wed, 7 Oct 2026 19:36:13 -0700 Subject: [PATCH 02/21] Implement and test saving results and marginalization without input --- orbitize/marginalizer.py | 37 ++++++++++++++++++ orbitize/sampler.py | 54 ++++++++++++++++++--------- tests/test_marg.py | 81 ++++++++++++++++++++++++++++++++++++++++ tests/test_mcmc.py | 4 +- 4 files changed, 157 insertions(+), 19 deletions(-) create mode 100644 tests/test_marg.py diff --git a/orbitize/marginalizer.py b/orbitize/marginalizer.py index f9c01363..5723d91c 100644 --- a/orbitize/marginalizer.py +++ b/orbitize/marginalizer.py @@ -1,4 +1,5 @@ import abc +import numpy as np """ This module defines marginalizations to draw @@ -28,6 +29,42 @@ class Marg2(Marginalizer): REQ_PARAMS = ["tau", "sma", "pan"] MARG_PARAMS = ["gamma"] +class MargConstNoInput(Marginalizer): + + MARG_PARAMS = ["test"] + + def apply(self, param_arr, jitter_arr, calculated_values): + shape = param_arr.shape + if len(shape) > 1: + n_orbits = shape[1] + return 0, np.zeros((n_orbits)), np.ones((n_orbits)), calculated_values + else: + return 0, 0, 1, calculated_values + +class MargPriorNoInput(Marginalizer): + + MARG_PARAMS = ["test"] + + def __init__(self, priors): + self.dim = len(priors) + self.MARG_PARAMS = ["test{}".format(i) for i in range(self.dim)] + self.priors = priors + + def apply(self, param_arr, jitter_arr, calculated_values): + shape = param_arr.shape + if len(shape) > 1: + n_orbits = shape[1] + assert n_orbits > 1 + samples = np.array([[prior.draw_samples(1).item() for i in range(n_orbits)] for prior in self.priors]) + assert samples.shape == (self.dim, n_orbits) + assert not np.isnan(samples).any(), samples + return 0, np.zeros((n_orbits)), samples, calculated_values + else: + samples = np.array([prior.draw_samples(1).item() for prior in self.priors]) + assert samples.shape == (self.dim,), samples + assert not np.isnan(samples).any(), samples + return 0, 0, samples, calculated_values + def compatible_params(marginalizers, input_params): # TODO: Multibody, quant_type params = set(input_params) diff --git a/orbitize/sampler.py b/orbitize/sampler.py index 3c1e2144..763b7413 100644 --- a/orbitize/sampler.py +++ b/orbitize/sampler.py @@ -18,6 +18,7 @@ import orbitize.lnlike import orbitize.priors import orbitize.results +import orbitize.marginalizer import sys @@ -31,7 +32,7 @@ class Sampler(abc.ABC): """ def __init__( - self, system, like="chi2_lnlike", custom_lnlike=None, chi2_type="standard" + self, system, like="chi2_lnlike", custom_lnlike=None, chi2_type="standard", marginalizers=None ): self.system = system @@ -46,11 +47,14 @@ def __init__( # check if need to handle covariances self.has_corr = np.any(~np.isnan(self.system.data_table["quant12_corr"])) - # TODO: Construct - self.marginalizers = [] + if marginalizers is None: + self.marginalizers = [] + else: + # TODO: Construct if not empty list + self.marginalizers = marginalizers # TODO: multiplanet - marginalized_param_labels_set = {(label for label in marg) for marg in self.marginalizers} + marginalized_param_labels_set = set(label for marg in self.marginalizers for label in marg.MARG_PARAMS) fitted_param_labels_set = set(system.labels) - marginalized_param_labels_set self.fitted_param_labels = [] @@ -78,7 +82,8 @@ def __init__( marg_param_counter += 1 self.total_param_num = len(self.system.sys_priors) quants = set(np.unique(self.system.data_table["quant_type"])) - marg_quants = {(quant_type for quant_type in marg) for marg in self.marginalizers} + marg_quants = {(quant_type for quant_type in marg.QUANT_TYPES) for marg in self.marginalizers} + self.blob_labels = ["lnlike"] + self.marg_param_labels self.non_marg_quants = list(quants - marg_quants) @@ -110,7 +115,7 @@ def _logl(self, fitted_params): else: num_orbits = input_shape[1] params = np.full((self.total_param_num, num_orbits), np.nan) - assert len(self.fitted_param_sys_idx) == len(fitted_params) == self.total_param_num, (self.fitted_param_sys_idx, fitted_params) + assert len(self.fitted_param_sys_idx) == len(fitted_params), (self.fitted_param_sys_idx, fitted_params) params[self.fitted_param_sys_idx] = fitted_params total_lnlike = 0 @@ -127,6 +132,7 @@ def _logl(self, fitted_params): total_lnlike += partial_lnlike params[marg_param_idx] = marg_params all_calculated_values.update(calculated_values) # i.e. astr_tanom, rv_eanom + assert not np.isnan(params).any(), (marg_params, params.shape, marg_param_idx) lnlikes_sum = 0 if len(self.non_marg_quants) > 0: @@ -164,7 +170,6 @@ def _logl(self, fitted_params): non_marg_corrs = None non_marg_model = model[non_marg_mask] non_marg_jitter = jitter[non_marg_mask] - # assert False, (data, non_marg_data, non_marg_data.shape, non_marg_mask, self.non_marg_quants, self.system.data_table["quant_type"]) # compute lnlike lnlikes = self.lnlike( @@ -227,6 +232,7 @@ def _logl(self, fitted_params): total_lnlike += lnlikes_sum marg_params = params[self.marg_param_sys_idx] + assert not np.isnan(marg_params).any(), marg_params return total_lnprob, total_lnlike, *marg_params @@ -787,9 +793,10 @@ def __init__( like="chi2_lnlike", custom_lnlike=None, prev_result_filename=None, + marginalizers=None ): super(MCMC, self).__init__( - system, like=like, chi2_type=chi2_type, custom_lnlike=custom_lnlike + system, like=like, chi2_type=chi2_type, custom_lnlike=custom_lnlike, marginalizers=marginalizers ) self.num_temps = num_temps @@ -817,7 +824,7 @@ def __init__( self.sampled_param_idx = {} sampled_param_counter = 0 - for i, prior in enumerate(system.sys_priors): # TODO: use self.fitted_priors instead + for i, prior in enumerate(self.fitted_priors): # check for fixed parameters if not hasattr(prior, "draw_samples"): self.fixed_params.append((i, prior)) @@ -921,7 +928,7 @@ def _logl(self, params, include_logp=False): logp = orbitize.priors.all_lnpriors(params, self.priors) # escape if logp == -np.inf if np.isinf(logp): - return -np.inf + return -np.inf, *[0.0 for i in range(len(self.blob_labels))] else: logp = np.array( [orbitize.priors.all_lnpriors(pset, self.priors) for pset in params] @@ -929,13 +936,12 @@ def _logl(self, params, include_logp=False): else: logp = 0 # don't include prior - assert len(self.fixed_params) == 2, self.fixed_params full_params = self._fill_in_fixed_params(params) if np.ndim(full_params) == 2: full_params = full_params.T logprob, loglike, *marg_params = super(MCMC, self)._logl(full_params) - return logprob + logp + return logprob + logp, loglike, *marg_params def _update_chains_from_sampler(self, sampler, num_steps=None): """ @@ -972,6 +978,20 @@ def _update_chains_from_sampler(self, sampler, num_steps=None): # include fixed parameters in posterior self.post = self._fill_in_fixed_params(self.post) + if self.use_pt: + assert False, "need to implement blobs for ptemcee" # TODO + else: + self.blobs = sampler.blobs[:num_steps, :].T.reshape(len(self.blob_labels), -1).T + assert (self.lnlikes == self.blobs[:, 0]).all(), (self.lnlikes, self.blobs[:, 0]) + self.marginalized_param_blobs = self.blobs[:, 1:] + assert not np.isnan(self.marginalized_param_blobs).any(), sampler.blobs[:5, :10] + num_orbits = self.post.shape[0] + self.full_post = np.full((num_orbits, self.total_param_num), np.nan) + self.full_post[:, self.fitted_param_sys_idx] = self.post + self.full_post[:, self.marg_param_sys_idx] = self.marginalized_param_blobs + assert not np.isnan(self.full_post).any(), (self.fitted_param_sys_idx, self.marg_param_sys_idx) + + def validate_xyz_positions(self): """ If using the XYZ basis, walkers might be initialized in an invalid @@ -1149,13 +1169,13 @@ def run_sampler( # figure out what is the new chunk of the chain and corresponding lnlikes that have been computed before last save # grab the current posterior and lnlikes and reshape them to have the Nwalkers x Nsteps dimension again - post_shape = self.post.shape + post_shape = self.full_post.shape curr_chain_shape = ( self.num_walkers, post_shape[0] // self.num_walkers, post_shape[-1], ) - curr_chain = self.post.reshape(curr_chain_shape) + curr_chain = self.full_post.reshape(curr_chain_shape) curr_lnlike_chain = self.lnlikes.reshape(curr_chain_shape[:2]) # use the reshaped arrays and find the new steps we computed curr_chunk = curr_chain[:, saved_upto : i + 1] @@ -1179,18 +1199,18 @@ def run_sampler( if periodic_save_freq is None: # need to save everything self.results.add_samples( - self.post, self.lnlikes, curr_pos=self.curr_pos + self.full_post, self.lnlikes, curr_pos=self.curr_pos ) elif saved_upto < nsteps: # just need to save the last few # same code as above except we just need to grab the last few - post_shape = self.post.shape + post_shape = self.full_post.shape curr_chain_shape = ( self.num_walkers, post_shape[0] // self.num_walkers, post_shape[-1], ) - curr_chain = self.post.reshape(curr_chain_shape) + curr_chain = self.full_post.reshape(curr_chain_shape) curr_lnlike_chain = self.lnlikes.reshape(curr_chain_shape[:2]) curr_chunk = curr_chain[:, saved_upto:] curr_chunk = curr_chunk.reshape( diff --git a/tests/test_marg.py b/tests/test_marg.py new file mode 100644 index 00000000..2db32e55 --- /dev/null +++ b/tests/test_marg.py @@ -0,0 +1,81 @@ +import pytest + +import os +import numpy as np +import orbitize +from orbitize.driver import Driver +import orbitize.sampler as sampler +import orbitize.marginalizer as marginalizer +import orbitize.priors as priors +import orbitize.system as system +import orbitize.read_input as read_input +import orbitize.results as results +from orbitize.plot import plot_corner +import matplotlib.pyplot as plt +import pytest + +def test_MargConstNoInput(): + # use the test_csv dir + input_file = os.path.join(orbitize.DATADIR, 'GJ504.csv') + data_table = read_input.read_file(input_file) + output_file = os.path.join(orbitize.DATADIR, 'test_marg.hdf5') + + mySystem = system.System( + 1, + data_table, + 1, + 0.01, + ) + marg_prior = 1 + marg = marginalizer.MargConstNoInput() + mySystem.sys_priors.append(marg_prior) + mySystem.param_idx["test"] = len(mySystem.sys_priors)-1 + mySystem.labels.append("test") + mySampler = sampler.MCMC( + mySystem, + num_temps=0, + num_walkers=100, + num_threads=1, + marginalizers=[marg] + ) + mySampler.run_sampler(400, 10, output_filename=output_file, periodic_save_freq=2) + myResults = mySampler.results + assert (myResults.post[:, 8] == 1).all() + +def test_MargPriorNoInput(): + # use the test_csv dir + input_file = os.path.join(orbitize.DATADIR, 'GJ504.csv') + data_table = read_input.read_file(input_file) + output_file = os.path.join(orbitize.DATADIR, 'test_marg.hdf5') + + mySystem = system.System( + 1, + data_table, + 1, + 0.01, + ) + marg_priors = [priors.UniformPrior(0.5, 1.5), priors.GaussianPrior(10, 1)] + marg = marginalizer.MargPriorNoInput(marg_priors) + regular = len(mySystem.sys_priors) + mySystem.sys_priors.extend(marg_priors) + for i, label in enumerate(marg.MARG_PARAMS): + mySystem.param_idx[label] = regular + i + assert mySystem.param_idx["test0"] == 8 + assert mySystem.param_idx["test1"] == 9 + mySystem.labels.extend(marg.MARG_PARAMS) + mySampler = sampler.MCMC( + mySystem, + num_temps=0, + num_walkers=100, + num_threads=1, + marginalizers=[marg] + ) + mySampler.run_sampler(400, 10, output_filename=output_file, periodic_save_freq=2) + myResults = mySampler.results + assert (0.5 <= myResults.post[:, 8]).all() and (myResults.post[:, 8] <= 1.5).all() + assert np.mean(myResults.post[:, 8]) == pytest.approx(1, rel=0.1) + assert np.std(myResults.post[:, 9]) == pytest.approx(1, rel=0.1) + assert np.mean(myResults.post[:, 9]) == pytest.approx(10, rel=0.1) + +if __name__ == "__main__": + test_MargConstNoInput() \ No newline at end of file diff --git a/tests/test_mcmc.py b/tests/test_mcmc.py index cf844d0d..8a0d7ac0 100644 --- a/tests/test_mcmc.py +++ b/tests/test_mcmc.py @@ -199,8 +199,8 @@ def test_mcmc_param_idx(): def test_mcmc_runs(): # Parallel Tempering tests (only test corner plot once) - do_mcmc_runs(num_temps=2, num_threads=1, make_corner_plot=True) - do_mcmc_runs(num_temps=2, num_threads=4) + # do_mcmc_runs(num_temps=2, num_threads=1, make_corner_plot=True) + # do_mcmc_runs(num_temps=2, num_threads=4) # Ensemble MCMC tests do_mcmc_runs(num_temps=0, num_threads=1) do_mcmc_runs(num_temps=0, num_threads=8) From 1c65ebe4c83390ac80ca1c49981f27d6cae6a16b Mon Sep 17 00:00:00 2001 From: Eshel Dror Date: Tue, 13 Oct 2026 10:55:43 -0700 Subject: [PATCH 03/21] Add OftEiBasis and non jitter marginazliation --- orbitize/basis.py | 84 +++++++++++++++ orbitize/marginalizer.py | 227 +++++++++++++++++++++++++++++++++++++++ orbitize/sampler.py | 6 +- tests/test_marg.py | 69 +++++++++++- 4 files changed, 384 insertions(+), 2 deletions(-) diff --git a/orbitize/basis.py b/orbitize/basis.py index b7dff280..0ac5fc0d 100644 --- a/orbitize/basis.py +++ b/orbitize/basis.py @@ -627,6 +627,67 @@ def to_period_basis(self, param_arr): return param_arr +class OftEiBasis(Basis): + def __init__(self, + stellar_or_system_mass, + mass_err, + plx, + plx_err, + num_secondary_bodies, + fit_secondary_mass, + angle_upperlim=2 * np.pi, + hipparcos_IAD=None, + rv=False, + rv_instruments=None): + + super(OftEiBasis, self).__init__( + stellar_or_system_mass, + mass_err, + plx, + plx_err, + num_secondary_bodies, + fit_secondary_mass, + angle_upperlim, + hipparcos_IAD, + rv, + rv_instruments + ) + assert not self.fit_secondary_mass + + def construct_priors(self, per_min=0.1, per_max=1000): + basis_labels = ["sma1", "ecc1", "inc1", "aop1", "pan1", "tau1", "plx", "per1"] + basis_priors = [ + priors.GaussianPrior(100, 10), # set nonstandard sma prior + priors.UniformPrior( + 0, + 1, + ), + priors.SinPrior(), + priors.UniformPrior(0, 2 * np.pi), + priors.GaussianPrior(180, 10), # set nonstandard pan prior + priors.UniformPrior(0, 1), + self.plx, + priors.LogUniformPrior(per_min, per_max), + ] + + # Define param label dictionary in current basis & standard basis + self.param_idx = dict(zip(basis_labels, np.arange(len(basis_labels)))) + self.standard_basis_idx = dict( + zip(basis_labels[:-1], np.arange(len(basis_labels) - 1)) + ) + self.standard_basis_idx["mtot"] = 7 + + return basis_priors, basis_labels + + def to_standard_basis(self, param_arr): + # everything is the same except the last param element goes from per to mtot + sma = param_arr[0] + per = param_arr[-1] + + mtot = (sma**3) / (per**2) + + param_arr[-1] = mtot + return param_arr class SemiAmp(Basis): """ @@ -1631,3 +1692,26 @@ def tau_to_manom(date, sma, mtot, tau, tau_ref_epoch): """ return kepler.tau_to_manom(date, sma, mtot, tau, tau_ref_epoch) + +def angle_diff(ang1, ang2): + """ + Difference between two angles (in degrees), account for angle wrapping around 360. + """ + return (ang1 - ang2 + 180.0) % 360.0 - 180.0 + +def return_quadrant(angle): + """ + Return quadrant of a particular angle + + Args: + angle (np.array of float): radians + + Returns: + int: quadrant 1, 2, 3, or 4 + """ + angle %= 2 * np.pi # [-360, 0] if negative, [0, 360] if positive + angle += 2 * np.pi + angle %= 2 * np.pi + + quadrant = np.floor((angle / (np.pi / 2)) % 4 + 1).astype(int) + return quadrant \ No newline at end of file diff --git a/orbitize/marginalizer.py b/orbitize/marginalizer.py index 5723d91c..52ef5cbd 100644 --- a/orbitize/marginalizer.py +++ b/orbitize/marginalizer.py @@ -1,5 +1,8 @@ import abc import numpy as np +from astropy import units as u +import orbitize.basis as basis +import orbitize.kepler as kepler """ This module defines marginalizations to draw @@ -65,6 +68,230 @@ def apply(self, param_arr, jitter_arr, calculated_values): assert not np.isnan(samples).any(), samples return 0, 0, samples, calculated_values +class OFTeI_sma_pan(Marginalizer): + + REQ_PARAMS = ["ecc1", "inc1", "aop1", "tau1", "plx", "per1"] + MARG_PARAMS = ["sma1", "pan1"] + QUANT_TYPES = ["seppa", "radec"] + LNLIKE_TYPE = "chi2" + + def __init__(self, epochs, seps, pas, sigma_seps, sigma_pas, covar_seps_pas, sma_prior=None, pan_prior=None, mass_prior=None, implicit_mass_prior=None): + self.sma_prior = sma_prior + self.pan_prior = pan_prior + self.mass_prior = mass_prior + self.implicit_mass_prior = implicit_mass_prior + self.epochs = epochs + self.seps = seps + self.pas = pas + self.sigma_seps = sigma_seps + self.sigma_pas = sigma_pas + self.covar_seps_pas = covar_seps_pas + self.n_obs = len(epochs) + + @staticmethod + def calc_f_and_g(manoms, ecc, inc, aop, plx): + """ + Args: + manoms (np.array of float): array of mean anomalies [rad, 0-2pi] + ecc (float): eccentricity + inc (float): inclination [rad, 0-pi] + aop (float): argument of periastron [rad, 0-2pi] + plx (float): parallax [mas] + + Return: + tuple of: + float: f (defined in \ref{eq:f_and_g} in PASP review) + float: g (defined in \ref{eq:f_and_g} in PASP review) + + """ + + n_dates = np.size(manoms) + + # compute eccentric anomalies (size: n_dates) + eanom = kepler.calc_ecc_anom(manoms, ecc) + + tanom = 2.0 * np.arctan(np.sqrt((1.0 + ecc) / (1.0 - ecc)) * np.tan(0.5 * eanom)) + + radius = (1.0 - ecc * np.cos(eanom)) * plx + + f = ( + 0.5 + * radius + * np.sqrt( + 3 + np.cos(2 * inc) + 2 * np.cos(2 * (aop + tanom)) * (np.sin(inc)) ** 2 + ) + ) + + aoptanom_02pi = ((aop + tanom) % (2 * np.pi) + 2 * np.pi) % (2 * np.pi) + + quadrant_aoptanom = basis.return_quadrant(aoptanom_02pi) + quadrant_corr = np.zeros(n_dates) - np.pi / 2 + quadrant_corr[quadrant_aoptanom == 1] -= np.pi + quadrant_corr[quadrant_aoptanom == 4] -= np.pi + + if np.cos(inc) < 0: + quadrant_corr += np.pi + + g_radians = ( + np.arctan2( + -1 / np.cos(inc), + np.tan(aoptanom_02pi), + ) + ) + quadrant_corr + + g = np.degrees(g_radians) % 360 + + return f, g + + @staticmethod + def calc_cov_matrix_det(seps, pas, fs, gs, W): + """ + Calculate the MAP values of sma and pan and the determinant of the + covariance matrix describing their conditional posterior uncertainties + + Args: + seps (np.array of float): separation measurements [mas] + pas (np.array of float): position angle measurements [deg] + fs (np.array of float): (defined in \ref{eq:f_and_g} in PASP review), values at observation epochs + gs (np.array of float): (defined in \ref{eq:f_and_g} in PASP review), values at observation epochs + W: 2n x 2n covariance matrix corresponding to observations y + + Returns: + float: log(abs(determinant of cov matrix)) + """ + + n_obs = len(seps) + + # set up design matrix + D = np.zeros((2 * n_obs, 2)) + D[:n_obs, 0] = fs + D[n_obs:, 1] = 1 + + Y = np.hstack((seps, basis.angle_diff(pas, gs))) + + def solve_linear_eq(Y, A, C): + C_inv = np.linalg.inv(C) + cov_matrix = np.linalg.inv( + A.T @ C_inv @ A + ) # variance on diag, covar on off-diag + best_fit = cov_matrix @ (A.T @ C_inv @ Y) + return best_fit, cov_matrix + + map_fit, cov_matrix = solve_linear_eq(Y, D, W) + sma_map = map_fit[0] + pan_map = map_fit[1] + + # sample linear parameters + sma_sample, pan_sample = np.random.multivariate_normal( + [sma_map, pan_map], cov_matrix + ) + + _, logabsdet_covmatrix = np.linalg.slogdet( + cov_matrix + ) # NOTE: the 2pi factor in the APW paper is constant; doesn't affect overall lnlike + + return sma_map, pan_map, logabsdet_covmatrix, sma_sample, pan_sample + + @staticmethod + def lnlike(y, X, beta_hat, W, add_error_norm=False): + """ + Compute log likelihood, excluding volume in marginalization integral (first two terms in \ref{eq:marg_like} in + PASP review): + + -1/2 log|W| - 1/2 chi2_MAP + + Args: + y: 2*n_obs array of observations (sep on top, PA-g on bottom) + X: 2*n_obs x 2 design matrix + beta_hat: length 2 array of conditional MAP values of sma and pan (respectively) + W: 2n x 2n covariance matrix corresponding to observations y + add_error_term: if True, include the -1/2 log|W| term (needed if fitting for jitter). Default False.as_integer_ratio + + Returns: + float: loglikelihood + """ + + chi2_map = (y - X @ beta_hat).T @ np.linalg.inv(W) @ (y - X @ beta_hat) + + # NOTE: this term can be ignored unless we're fitting for jitter + if add_error_norm: + error_norm_term = -0.5 * np.log(np.det(W)) + else: + error_norm_term = 0 + + return -0.5 * chi2_map + error_norm_term + + def apply(self, param_arr, jitter_arr, calculated_values): + assert jitter_arr.shape[0] == 0, "No jitter supported" + shape = param_arr.shape + if len(shape) > 1: + raise IndexError("Only support 1 orbit at a time") + ecc, inc, aop, tau, plx, period_yr = param_arr + frac_date = self.epochs / (period_yr * u.yr.to(u.day)) + frac_date %= 1 + manoms = (frac_date - tau) * 2 * np.pi + manoms %= 2 * np.pi + f_samples, g_samples = self.calc_f_and_g( + manoms, + ecc, + inc, + aop, + plx, + ) + # compute variance-covariance matrix W of observational uncertainties + W = ( + np.diag(np.hstack((self.sigma_seps**2, self.sigma_pas**2))) + + np.diag(self.covar_seps_pas, k=self.n_obs) + + np.diag(self.covar_seps_pas, k=-self.n_obs) + ) + + try: + smas_map, pans_map_degrees, logdet_covmatrix, sma_sample, pan_sample = ( + self.calc_cov_matrix_det(self.seps, self.pas, f_samples, g_samples, W) + ) + except np.linalg.LinAlgError: + return -np.inf, None, None + + # set up design matrix + X = np.zeros((2 * self.n_obs, 2)) + X[:self.n_obs, 0] = f_samples + X[self.n_obs:, 1] = 1 + + y = np.hstack((self.seps, basis.angle_diff(self.pas, g_samples))) + beta_hat = np.hstack((smas_map, pans_map_degrees)) + + # compute marginalized likelihood (e.g. Price-Whelan 2017 eq 11) + chi2lnlike = self.lnlike(y, X, beta_hat, W) + + logprob = ( + chi2lnlike + 0.5 * logdet_covmatrix + ) # NOTE: the sign here differs from APW+ 2017 (which I think is correct; I think they have a bug) + + # apply sma/pan priors + if self.sma_prior is not None: + logprob += np.sum(self.sma_prior.compute_lnprob(np.array([sma_sample]))) + if self.pan_prior is not None: + logprob += np.sum(self.pan_prior.compute_lnprob(np.array([pan_sample]))) + + # apply total mass prior + mtot_sample = sma_sample**3 / (period_yr) ** 2 + if self.mass_prior is not None: + logprob -= self.implicit_mass_prior.compute_lnprob( + np.array([mtot_sample]) + ) # "undo" the effective prior + logprob += self.mass_prior.compute_lnprob( + np.array([mtot_sample]) + ) # apply the user-defined prior + + loglike = logprob # TODO + samples = np.array([sma_sample, pan_sample]) + + if not np.isfinite(logprob): + return -np.inf, -np.inf, samples, calculated_values + + return logprob, loglike, samples, calculated_values + + def compatible_params(marginalizers, input_params): # TODO: Multibody, quant_type params = set(input_params) diff --git a/orbitize/sampler.py b/orbitize/sampler.py index 763b7413..5ab00875 100644 --- a/orbitize/sampler.py +++ b/orbitize/sampler.py @@ -53,6 +53,7 @@ def __init__( # TODO: Construct if not empty list self.marginalizers = marginalizers + # TODO: multiplanet marginalized_param_labels_set = set(label for marg in self.marginalizers for label in marg.MARG_PARAMS) fitted_param_labels_set = set(system.labels) - marginalized_param_labels_set @@ -84,9 +85,12 @@ def __init__( quants = set(np.unique(self.system.data_table["quant_type"])) marg_quants = {(quant_type for quant_type in marg.QUANT_TYPES) for marg in self.marginalizers} self.blob_labels = ["lnlike"] + self.marg_param_labels - + self.non_marg_quants = list(quants - marg_quants) + if not orbitize.marginalizer.compatible_params(self.marginalizers, self.fitted_param_labels): + raise ValueError("Marginalizers and basis are incompatible") + @abc.abstractmethod def run_sampler(self, total_orbits): pass diff --git a/tests/test_marg.py b/tests/test_marg.py index 2db32e55..8616c45f 100644 --- a/tests/test_marg.py +++ b/tests/test_marg.py @@ -77,5 +77,72 @@ def test_MargPriorNoInput(): assert np.std(myResults.post[:, 9]) == pytest.approx(1, rel=0.1) assert np.mean(myResults.post[:, 9]) == pytest.approx(10, rel=0.1) +def test_incompatible(): + input_file = os.path.join(orbitize.DATADIR, 'GJ504.csv') + data_table = read_input.read_file(input_file) + output_file = os.path.join(orbitize.DATADIR, 'test_marg.hdf5') + + mySystem = system.System( + 1, + data_table, + 1, + 0.01, + ) + epochs = data_table[data_table["quant_type"] == "epoch"] + sep = data_table[data_table["quant_type"] == "sep"] + pa = data_table[data_table["quant_type"] == "pa"] + sep_err = data_table[data_table["quant_type"] == "sep_err"] + pa_err = data_table[data_table["quant_type"] == "pa_err"] + covar = np.zeros(len(epochs)) + marg = marginalizer.OFTeI_sma_pan(epochs, sep, pa, sep_err, pa_err, covar) + try: + mySampler = sampler.MCMC( + mySystem, + num_temps=0, + num_walkers=100, + num_threads=1, + marginalizers=[marg] + ) + mySampler.run_sampler(400, 10, output_filename=output_file, periodic_save_freq=2) + except ValueError: + assert True + else: + assert False, "Did not raise ValueError despite marginalizer incompatible with basis" + +def test_OFTeI(): + input_file = os.path.join(orbitize.DATADIR, 'GJ504.csv') + data_table = read_input.read_file(input_file) + output_file = os.path.join(orbitize.DATADIR, 'test_marg.hdf5') + + mySystem = system.System( + 1, + data_table, + 1, + 0.01, + fitting_basis="OftEiBasis" + ) + ind = data_table["quant_type"] == "seppa" + epochs = data_table["epoch"][ind] + sep = data_table["quant1"][ind] + pa = data_table["quant2"][ind] + sep_err = data_table["quant1_err"][ind] + pa_err = data_table["quant2_err"][ind] + covar = np.zeros(len(epochs)) + sma_prior = mySystem.sys_priors[mySystem.param_idx["sma1"]] + pan_prior = mySystem.sys_priors[mySystem.param_idx["pan1"]] + marg = marginalizer.OFTeI_sma_pan(epochs, sep, pa, sep_err, pa_err, covar, sma_prior, pan_prior) + mySampler = sampler.MCMC( + mySystem, + num_temps=0, + num_walkers=100, + num_threads=1, + marginalizers=[marg] + ) + mySampler.run_sampler(40000, 10, output_filename=output_file, periodic_save_freq=2) + myResults = mySampler.results + print(myResults.post.shape) + print(myResults.post[:3, :]) + if __name__ == "__main__": - test_MargConstNoInput() \ No newline at end of file + # test_MargConstNoInput() + test_OFTeI() \ No newline at end of file From d6666e6ff13e4f3468f2d967d1b38b84aaca6135 Mon Sep 17 00:00:00 2001 From: Eshel Dror Date: Tue, 13 Oct 2026 15:55:55 -0700 Subject: [PATCH 04/21] Fix bugs, implement mass prior --- orbitize/basis.py | 6 ++--- orbitize/marginalizer.py | 50 +++++++++++++++++++++++----------------- orbitize/plot.py | 20 ++++++++-------- orbitize/priors.py | 50 ++++++++++++++++++++++++++++++++++++++++ orbitize/sampler.py | 3 ++- tests/test_marg.py | 39 ++++++++++++++++++++++--------- 6 files changed, 122 insertions(+), 46 deletions(-) diff --git a/orbitize/basis.py b/orbitize/basis.py index 0ac5fc0d..02f5d113 100644 --- a/orbitize/basis.py +++ b/orbitize/basis.py @@ -654,17 +654,17 @@ def __init__(self, ) assert not self.fit_secondary_mass - def construct_priors(self, per_min=0.1, per_max=1000): + def construct_priors(self, per_min=1e-5, per_max=1e6): basis_labels = ["sma1", "ecc1", "inc1", "aop1", "pan1", "tau1", "plx", "per1"] basis_priors = [ - priors.GaussianPrior(100, 10), # set nonstandard sma prior + priors.LogUniformPrior(0.001, 1e4), priors.UniformPrior( 0, 1, ), priors.SinPrior(), priors.UniformPrior(0, 2 * np.pi), - priors.GaussianPrior(180, 10), # set nonstandard pan prior + priors.UniformPrior(0, 2 * np.pi), priors.UniformPrior(0, 1), self.plx, priors.LogUniformPrior(per_min, per_max), diff --git a/orbitize/marginalizer.py b/orbitize/marginalizer.py index 52ef5cbd..ee661db2 100644 --- a/orbitize/marginalizer.py +++ b/orbitize/marginalizer.py @@ -87,6 +87,13 @@ def __init__(self, epochs, seps, pas, sigma_seps, sigma_pas, covar_seps_pas, sma self.sigma_pas = sigma_pas self.covar_seps_pas = covar_seps_pas self.n_obs = len(epochs) + # compute variance-covariance matrix W of observational uncertainties + self.W = ( + np.diag(np.hstack((self.sigma_seps**2, self.sigma_pas**2))) + + np.diag(self.covar_seps_pas, k=self.n_obs) + + np.diag(self.covar_seps_pas, k=-self.n_obs) + ) + self.W_inv = np.linalg.inv(self.W) @staticmethod def calc_f_and_g(manoms, ecc, inc, aop, plx): @@ -108,7 +115,7 @@ def calc_f_and_g(manoms, ecc, inc, aop, plx): n_dates = np.size(manoms) # compute eccentric anomalies (size: n_dates) - eanom = kepler.calc_ecc_anom(manoms, ecc) + eanom = kepler.calc_ecc_anom(manoms, ecc, use_c=False) tanom = 2.0 * np.arctan(np.sqrt((1.0 + ecc) / (1.0 - ecc)) * np.tan(0.5 * eanom)) @@ -144,7 +151,7 @@ def calc_f_and_g(manoms, ecc, inc, aop, plx): return f, g @staticmethod - def calc_cov_matrix_det(seps, pas, fs, gs, W): + def calc_cov_matrix_det(seps, pas, fs, gs, W, W_inv=None): """ Calculate the MAP values of sma and pan and the determinant of the covariance matrix describing their conditional posterior uncertainties @@ -169,15 +176,16 @@ def calc_cov_matrix_det(seps, pas, fs, gs, W): Y = np.hstack((seps, basis.angle_diff(pas, gs))) - def solve_linear_eq(Y, A, C): - C_inv = np.linalg.inv(C) + def solve_linear_eq(Y, A, C, C_inv=None): + if C_inv is None: + C_inv = np.linalg.inv(C) cov_matrix = np.linalg.inv( A.T @ C_inv @ A ) # variance on diag, covar on off-diag best_fit = cov_matrix @ (A.T @ C_inv @ Y) return best_fit, cov_matrix - map_fit, cov_matrix = solve_linear_eq(Y, D, W) + map_fit, cov_matrix = solve_linear_eq(Y, D, W, W_inv) sma_map = map_fit[0] pan_map = map_fit[1] @@ -193,7 +201,7 @@ def solve_linear_eq(Y, A, C): return sma_map, pan_map, logabsdet_covmatrix, sma_sample, pan_sample @staticmethod - def lnlike(y, X, beta_hat, W, add_error_norm=False): + def lnlike(y, X, beta_hat, W, W_inv=None, add_error_norm=False): """ Compute log likelihood, excluding volume in marginalization integral (first two terms in \ref{eq:marg_like} in PASP review): @@ -210,8 +218,9 @@ def lnlike(y, X, beta_hat, W, add_error_norm=False): Returns: float: loglikelihood """ - - chi2_map = (y - X @ beta_hat).T @ np.linalg.inv(W) @ (y - X @ beta_hat) + if W_inv is None: + W_inv = np.linalg.inv(W) + chi2_map = (y - X @ beta_hat).T @ W_inv @ (y - X @ beta_hat) # NOTE: this term can be ignored unless we're fitting for jitter if add_error_norm: @@ -227,6 +236,7 @@ def apply(self, param_arr, jitter_arr, calculated_values): if len(shape) > 1: raise IndexError("Only support 1 orbit at a time") ecc, inc, aop, tau, plx, period_yr = param_arr + assert period_yr > 0 frac_date = self.epochs / (period_yr * u.yr.to(u.day)) frac_date %= 1 manoms = (frac_date - tau) * 2 * np.pi @@ -238,16 +248,12 @@ def apply(self, param_arr, jitter_arr, calculated_values): aop, plx, ) - # compute variance-covariance matrix W of observational uncertainties - W = ( - np.diag(np.hstack((self.sigma_seps**2, self.sigma_pas**2))) - + np.diag(self.covar_seps_pas, k=self.n_obs) - + np.diag(self.covar_seps_pas, k=-self.n_obs) - ) + W = self.W + W_inv = self.W_inv try: - smas_map, pans_map_degrees, logdet_covmatrix, sma_sample, pan_sample = ( - self.calc_cov_matrix_det(self.seps, self.pas, f_samples, g_samples, W) + smas_map, pans_map_degrees, logdet_covmatrix, sma_sample, pan_sample_deg = ( + self.calc_cov_matrix_det(self.seps, self.pas, f_samples, g_samples, W, W_inv) ) except np.linalg.LinAlgError: return -np.inf, None, None @@ -261,12 +267,14 @@ def apply(self, param_arr, jitter_arr, calculated_values): beta_hat = np.hstack((smas_map, pans_map_degrees)) # compute marginalized likelihood (e.g. Price-Whelan 2017 eq 11) - chi2lnlike = self.lnlike(y, X, beta_hat, W) + chi2lnlike = self.lnlike(y, X, beta_hat, W, W_inv) logprob = ( chi2lnlike + 0.5 * logdet_covmatrix ) # NOTE: the sign here differs from APW+ 2017 (which I think is correct; I think they have a bug) + pan_sample = np.deg2rad(pan_sample_deg) % (2 * np.pi) + # apply sma/pan priors if self.sma_prior is not None: logprob += np.sum(self.sma_prior.compute_lnprob(np.array([sma_sample]))) @@ -276,12 +284,12 @@ def apply(self, param_arr, jitter_arr, calculated_values): # apply total mass prior mtot_sample = sma_sample**3 / (period_yr) ** 2 if self.mass_prior is not None: - logprob -= self.implicit_mass_prior.compute_lnprob( + logprob -= np.sum(self.implicit_mass_prior.compute_lnprob( np.array([mtot_sample]) - ) # "undo" the effective prior - logprob += self.mass_prior.compute_lnprob( + )) # "undo" the effective prior + logprob += np.sum(self.mass_prior.compute_lnprob( np.array([mtot_sample]) - ) # apply the user-defined prior + )) # apply the user-defined prior loglike = logprob # TODO samples = np.array([sma_sample, pan_sample]) diff --git a/orbitize/plot.py b/orbitize/plot.py index e2f9142c..ccf6cf44 100644 --- a/orbitize/plot.py +++ b/orbitize/plot.py @@ -590,12 +590,7 @@ def _calc_object_full_orbits(self, start, num_orbits_to_plot, num_epochs_to_plot secondaries = np.arange(self.results.num_secondary_bodies) + 1 sma_indexes = [self.results.standard_param_idx["sma{}".format(i)] for i in secondaries] smas = standard_post[:, sma_indexes] - if "mtot" in self.results.labels: - mtots = standard_post[:, self.results.standard_param_idx["mtot"]] - periods = np.sqrt( - 4 * np.pi**2.0 * (smas.T * u.AU) ** 3 / (consts.G * (mtots * u.Msun)) - ).T - elif "m0" in self.results.labels: + if "m0" in self.results.labels: m0 = standard_post[:, self.results.standard_param_idx["m0"]] m1_indexes = [self.results.standard_param_idx["m{}".format(i)] for i in secondaries] m1s = standard_post[:, m1_indexes] @@ -603,20 +598,25 @@ def _calc_object_full_orbits(self, start, num_orbits_to_plot, num_epochs_to_plot periods = np.sqrt( 4 * np.pi**2.0 * (smas * u.AU) ** 3 / (consts.G * (mtots * u.Msun)) ) + else: + mtots = standard_post[:, self.results.standard_param_idx["mtot"]] + periods = np.sqrt( + 4 * np.pi**2.0 * (smas.T * u.AU) ** 3 / (consts.G * (mtots * u.Msun)) + ).T periods = periods.to(u.day).value period = np.max(periods, axis=1) else: sma = standard_post[ :, self.results.standard_param_idx["sma{}".format(object_to_plot)] ] - if "mtot" in self.results.labels: - mtot = standard_post[:, self.results.standard_param_idx["mtot"]] - elif "m0" in self.results.labels: + if "m0" in self.results.labels: m0 = standard_post[:, self.results.standard_param_idx["m0"]] m1 = standard_post[ :, self.results.standard_param_idx["m{}".format(object_to_plot)] ] mtot = m0 + m1 + else: + mtot = standard_post[:, self.results.standard_param_idx["mtot"]] period = np.sqrt( 4 * np.pi**2.0 * (sma * u.AU) ** 3 / (consts.G * (mtot * u.Msun)) ) @@ -2031,7 +2031,7 @@ def plot_corner(results, param_list=None, plot_priors=True, downsample=None, pri "labels" not in corner_kwargs ): # use default labels if user didn't already supply them reduced_labels_list = [] - for i in range(len(param_indices)): + for i in param_indices: label_key = param_list[i] if label_key.startswith("m") and label_key != "m0" and label_key != "mtot": body_num = label_key[1] diff --git a/orbitize/priors.py b/orbitize/priors.py index 8111515f..baa28797 100644 --- a/orbitize/priors.py +++ b/orbitize/priors.py @@ -860,3 +860,53 @@ def all_lnpriors(params, priors): plt.hist(samps, bins=50) plt.show() + +class OFTeIMassPrior(Prior): + def __init__(self, sma_min, sma_max, per_min, per_max): + """ + Assuming loguniform priors on semimajor axis and orbital period, + this is the effective prior implicity applied to total mass + + Args: + sma_min (float): minimum of the logUniform prior on semimajor axis [au] + sma_max (float): maximum of the logUniform prior on semimajor axis [au] + per_min (float): minimum of the logUniform prior on orbital period [yr] + per_max (float): maximum of the logUniform prior on orbital period [yr] + """ + self.sma_min = sma_min + self.sma_max = sma_max + self.per_min = per_min + self.per_max = per_max + + Wx = 3 * np.log(sma_max / sma_min) + Wy = 2 * np.log(per_max / per_min) + self.W1 = np.min([Wx, Wy]) + self.W2 = np.max([Wx, Wy]) + + self.offset = 3 * np.log(sma_min) - 2 * np.log(per_max) + + def draw_samples(self, num_samples): + sma_samples = np.exp( + np.random.uniform(np.log(self.sma_min), np.log(self.sma_max), num_samples) + ) + per_samples = np.exp( + np.random.uniform(np.log(self.per_min), np.log(self.per_max), num_samples) + ) + return sma_samples**3 / per_samples**2 + + def compute_lnprob(self, element_array): + + element_array = np.asarray(element_array, dtype=float) + + z = np.log(element_array) - self.offset + + f_Z = np.zeros_like(element_array) + f_Z[z <= self.W1] = z[z <= self.W1] / (self.W1 * self.W2) + f_Z[(z >= self.W1) & (z <= self.W2)] = 1 / self.W2 + f_Z[(z >= self.W2) & (z <= self.W1 + self.W2)] = ( + self.W1 + self.W2 - z[(z >= self.W2) & (z <= self.W1 + self.W2)] + ) / (self.W1 * self.W2) + f_Z[(z < 0) | (z > self.W1 + self.W2)] = 0 + + with np.errstate(divide="ignore"): + return np.log(f_Z / element_array) \ No newline at end of file diff --git a/orbitize/sampler.py b/orbitize/sampler.py index 5ab00875..bae04275 100644 --- a/orbitize/sampler.py +++ b/orbitize/sampler.py @@ -986,7 +986,8 @@ def _update_chains_from_sampler(self, sampler, num_steps=None): assert False, "need to implement blobs for ptemcee" # TODO else: self.blobs = sampler.blobs[:num_steps, :].T.reshape(len(self.blob_labels), -1).T - assert (self.lnlikes == self.blobs[:, 0]).all(), (self.lnlikes, self.blobs[:, 0]) + no_diff = ((self.lnlikes == self.blobs[:, 0]) | (np.abs(self.lnlikes - self.blobs[:, 0]) < 0.001)) + assert (no_diff).all(), (self.lnlikes[~no_diff], self.blobs[:, 0][~no_diff], np.shape(np.where(~no_diff)), np.shape(np.where(no_diff))) self.marginalized_param_blobs = self.blobs[:, 1:] assert not np.isnan(self.marginalized_param_blobs).any(), sampler.blobs[:5, :10] num_orbits = self.post.shape[0] diff --git a/tests/test_marg.py b/tests/test_marg.py index 8616c45f..1c07caa4 100644 --- a/tests/test_marg.py +++ b/tests/test_marg.py @@ -13,6 +13,7 @@ from orbitize.plot import plot_corner import matplotlib.pyplot as plt import pytest +import multiprocessing as mp def test_MargConstNoInput(): # use the test_csv dir @@ -88,11 +89,12 @@ def test_incompatible(): 1, 0.01, ) - epochs = data_table[data_table["quant_type"] == "epoch"] - sep = data_table[data_table["quant_type"] == "sep"] - pa = data_table[data_table["quant_type"] == "pa"] - sep_err = data_table[data_table["quant_type"] == "sep_err"] - pa_err = data_table[data_table["quant_type"] == "pa_err"] + ind = data_table["quant_type"] == "seppa" + epochs = data_table["epoch"][ind] + sep = data_table["quant1"][ind] + pa = data_table["quant2"][ind] + sep_err = data_table["quant1_err"][ind] + pa_err = data_table["quant2_err"][ind] covar = np.zeros(len(epochs)) marg = marginalizer.OFTeI_sma_pan(epochs, sep, pa, sep_err, pa_err, covar) try: @@ -114,13 +116,26 @@ def test_OFTeI(): data_table = read_input.read_file(input_file) output_file = os.path.join(orbitize.DATADIR, 'test_marg.hdf5') + mtot = 1.22 # Msol + plx = 56.95 # mas + mass_err = 0.08 # Msol + plx_err = 0.26 # mas + tau_ref_epoch = 0 + mySystem = system.System( 1, data_table, - 1, - 0.01, + mtot, + plx, + tau_ref_epoch=tau_ref_epoch, fitting_basis="OftEiBasis" ) + sma_prior_min = 0.1 # au + sma_prior_max = 500 # au + period_prior_min_yr = np.sqrt(sma_prior_min**3 / mtot) # [yr] + period_prior_max_yr = np.sqrt(sma_prior_max**3 / mtot) # [yr] + mySystem.sys_priors[mySystem.param_idx["sma1"]] = priors.LogUniformPrior(sma_prior_min, sma_prior_max) + mySystem.sys_priors[mySystem.param_idx["per1"]] = priors.LogUniformPrior(period_prior_min_yr, period_prior_max_yr) ind = data_table["quant_type"] == "seppa" epochs = data_table["epoch"][ind] sep = data_table["quant1"][ind] @@ -130,18 +145,20 @@ def test_OFTeI(): covar = np.zeros(len(epochs)) sma_prior = mySystem.sys_priors[mySystem.param_idx["sma1"]] pan_prior = mySystem.sys_priors[mySystem.param_idx["pan1"]] - marg = marginalizer.OFTeI_sma_pan(epochs, sep, pa, sep_err, pa_err, covar, sma_prior, pan_prior) + mass_prior = priors.GaussianPrior(mtot, mass_err) + per_prior = mySystem.sys_priors[mySystem.param_idx["per1"]] + implicit_mass_prior = priors.OFTeIMassPrior(sma_prior.minval, sma_prior.maxval, per_prior.minval, per_prior.maxval) + marg = marginalizer.OFTeI_sma_pan(epochs, sep, pa, sep_err, pa_err, covar, sma_prior, pan_prior, mass_prior, implicit_mass_prior) mySampler = sampler.MCMC( mySystem, num_temps=0, num_walkers=100, - num_threads=1, + num_threads=mp.cpu_count(), marginalizers=[marg] ) - mySampler.run_sampler(40000, 10, output_filename=output_file, periodic_save_freq=2) + mySampler.run_sampler(4000, 10, output_filename=output_file, periodic_save_freq=2) myResults = mySampler.results print(myResults.post.shape) - print(myResults.post[:3, :]) if __name__ == "__main__": # test_MargConstNoInput() From b344c39014756f954e109321dfac179a142aa4bd Mon Sep 17 00:00:00 2001 From: Eshel Dror Date: Wed, 14 Oct 2026 08:26:06 -0700 Subject: [PATCH 05/21] Test RV with OFTeI --- orbitize/sampler.py | 2 +- tests/test_marg.py | 50 +++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 51 insertions(+), 1 deletion(-) diff --git a/orbitize/sampler.py b/orbitize/sampler.py index bae04275..e473f121 100644 --- a/orbitize/sampler.py +++ b/orbitize/sampler.py @@ -83,7 +83,7 @@ def __init__( marg_param_counter += 1 self.total_param_num = len(self.system.sys_priors) quants = set(np.unique(self.system.data_table["quant_type"])) - marg_quants = {(quant_type for quant_type in marg.QUANT_TYPES) for marg in self.marginalizers} + marg_quants = set(quant_type for marg in self.marginalizers for quant_type in marg.QUANT_TYPES) self.blob_labels = ["lnlike"] + self.marg_param_labels self.non_marg_quants = list(quants - marg_quants) diff --git a/tests/test_marg.py b/tests/test_marg.py index 1c07caa4..83ce5071 100644 --- a/tests/test_marg.py +++ b/tests/test_marg.py @@ -160,6 +160,56 @@ def test_OFTeI(): myResults = mySampler.results print(myResults.post.shape) +def test_OFTeI_RV(): + input_file = os.path.join(orbitize.DATADIR, 'HD4747.csv') + data_table = read_input.read_file(input_file) + output_file = os.path.join(orbitize.DATADIR, 'test_marg.hdf5') + + mtot = 0.84 # [Msol] + plx = 53.18 # [mas] + mass_err = 0.04 # [Msol] + plx_err = 0.12 # [mas] + tau_ref_epoch = 0 + + mySystem = system.System( + 1, + data_table, + mtot, + plx, + tau_ref_epoch=tau_ref_epoch, + fitting_basis="OftEiBasis" + ) + sma_prior_min = 0.1 # au + sma_prior_max = 500 # au + period_prior_min_yr = np.sqrt(sma_prior_min**3 / mtot) # [yr] + period_prior_max_yr = np.sqrt(sma_prior_max**3 / mtot) # [yr] + mySystem.sys_priors[mySystem.param_idx["sma1"]] = priors.LogUniformPrior(sma_prior_min, sma_prior_max) + mySystem.sys_priors[mySystem.param_idx["per1"]] = priors.LogUniformPrior(period_prior_min_yr, period_prior_max_yr) + ind = data_table["quant_type"] == "seppa" + epochs = data_table["epoch"][ind] + sep = data_table["quant1"][ind] + pa = data_table["quant2"][ind] + sep_err = data_table["quant1_err"][ind] + pa_err = data_table["quant2_err"][ind] + covar = np.zeros(len(epochs)) + sma_prior = mySystem.sys_priors[mySystem.param_idx["sma1"]] + pan_prior = mySystem.sys_priors[mySystem.param_idx["pan1"]] + mass_prior = priors.GaussianPrior(mtot, mass_err) + per_prior = mySystem.sys_priors[mySystem.param_idx["per1"]] + implicit_mass_prior = priors.OFTeIMassPrior(sma_prior.minval, sma_prior.maxval, per_prior.minval, per_prior.maxval) + marg = marginalizer.OFTeI_sma_pan(epochs, sep, pa, sep_err, pa_err, covar, sma_prior, pan_prior, mass_prior, implicit_mass_prior) + mySampler = sampler.MCMC( + mySystem, + num_temps=0, + num_walkers=100, + num_threads=mp.cpu_count(), + marginalizers=[marg] + ) + mySampler.run_sampler(4000, 10, output_filename=output_file, periodic_save_freq=2) + myResults = mySampler.results + assert mySampler.non_marg_quants == ["rv"] + print(myResults.post.shape) + if __name__ == "__main__": # test_MargConstNoInput() test_OFTeI() \ No newline at end of file From 7dbb2ee5cbf9f047aa4465a3f6b3a00ad8ec3728 Mon Sep 17 00:00:00 2001 From: Eshel Dror Date: Wed, 14 Oct 2026 11:37:31 -0700 Subject: [PATCH 06/21] Add ptemcee blobs --- orbitize/sampler.py | 73 +++++++++++++++++++++++++++++-------------- tests/test_marg.py | 5 +-- tests/test_mcmc.py | 4 +-- tests/test_mcmc_rv.py | 4 +-- 4 files changed, 56 insertions(+), 30 deletions(-) diff --git a/orbitize/sampler.py b/orbitize/sampler.py index e473f121..45625a3f 100644 --- a/orbitize/sampler.py +++ b/orbitize/sampler.py @@ -87,6 +87,7 @@ def __init__( self.blob_labels = ["lnlike"] + self.marg_param_labels self.non_marg_quants = list(quants - marg_quants) + self.force_valid_marg = False if not orbitize.marginalizer.compatible_params(self.marginalizers, self.fitted_param_labels): raise ValueError("Marginalizers and basis are incompatible") @@ -125,21 +126,26 @@ def _logl(self, fitted_params): total_lnlike = 0 total_lnprob = 0 all_calculated_values = {} - for marg in self.marginalizers: - req_param_idx = [self.system.param_idx[label] for label in marg.REQ_PARAMS] - jitter_param_idx = [self.system.param_idx[label] for label in marg.JITTER_PARAMS] - marg_param_idx = [self.system.param_idx[label] for label in marg.MARG_PARAMS] - req_params = params[req_param_idx] - jitter = params[jitter_param_idx] - partial_lnprob, partial_lnlike, marg_params, calculated_values = marg.apply(req_params, jitter, all_calculated_values) - total_lnprob += partial_lnprob - total_lnlike += partial_lnlike - params[marg_param_idx] = marg_params - all_calculated_values.update(calculated_values) # i.e. astr_tanom, rv_eanom - assert not np.isnan(params).any(), (marg_params, params.shape, marg_param_idx) + if self.force_valid_marg: + samples = [np.squeeze(prior.draw_samples(num_orbits)) for prior in self.marg_priors] + params[self.marg_param_sys_idx] = samples + assert not np.isnan(params).any() + else: + for marg in self.marginalizers: + req_param_idx = [self.system.param_idx[label] for label in marg.REQ_PARAMS] + jitter_param_idx = [self.system.param_idx[label] for label in marg.JITTER_PARAMS] + marg_param_idx = [self.system.param_idx[label] for label in marg.MARG_PARAMS] + req_params = params[req_param_idx] + jitter = params[jitter_param_idx] + partial_lnprob, partial_lnlike, marg_params, calculated_values = marg.apply(req_params, jitter, all_calculated_values) + total_lnprob += partial_lnprob + total_lnlike += partial_lnlike + params[marg_param_idx] = marg_params + all_calculated_values.update(calculated_values) # i.e. astr_tanom, rv_eanom + assert not np.isnan(params).any(), (marg_params, params.shape, marg_param_idx) lnlikes_sum = 0 - if len(self.non_marg_quants) > 0: + if self.force_valid_marg or len(self.non_marg_quants) > 0: # compute the model based on system params model, jitter = self.system.compute_model(params) @@ -165,15 +171,22 @@ def _logl(self, fitted_params): seppa_indices = self.system.all_seppa else: seppa_indices = () - non_marg_mask = np.isin(np.astype(self.system.data_table["quant_type"], str), self.non_marg_quants) - non_marg_data = data[non_marg_mask] - non_marg_errs = errs[non_marg_mask] - if corrs is not None: - non_marg_corrs = corrs[non_marg_mask] + if self.force_valid_marg: + non_marg_data = data + non_marg_errs = errs + non_marg_corrs = corrs + non_marg_model = model + non_marg_jitter = jitter else: - non_marg_corrs = None - non_marg_model = model[non_marg_mask] - non_marg_jitter = jitter[non_marg_mask] + non_marg_mask = np.isin(np.astype(self.system.data_table["quant_type"], str), self.non_marg_quants) + non_marg_data = data[non_marg_mask] + non_marg_errs = errs[non_marg_mask] + if corrs is not None: + non_marg_corrs = corrs[non_marg_mask] + else: + non_marg_corrs = None + non_marg_model = model[non_marg_mask] + non_marg_jitter = jitter[non_marg_mask] # compute lnlike lnlikes = self.lnlike( @@ -818,6 +831,7 @@ def __init__( if self.num_temps > 1: self.use_pt = True + self.force_valid_marg = True else: self.use_pt = False self.num_temps = 1 @@ -983,11 +997,18 @@ def _update_chains_from_sampler(self, sampler, num_steps=None): self.post = self._fill_in_fixed_params(self.post) if self.use_pt: - assert False, "need to implement blobs for ptemcee" # TODO + # Ntemp x Nwalkers x Nsteps x Nblobs + self.blobs = sampler.blobs[0, :, :num_steps, :].reshape(-1, len(self.blob_labels)) + assert len(self.blobs) == len(self.lnlikes) + no_diff = ((self.lnlikes == self.blobs[:, 0]) | (np.abs(self.lnlikes - self.blobs[:, 0]) < 0.001)) + assert (no_diff).all(), (self.lnlikes[~no_diff][:10], self.blobs[:, 0][~no_diff][:10], np.shape(np.where(~no_diff)), np.shape(np.where(no_diff))) + # assert (no_diff).all(), (sampler.loglikelihood[:, :, :num_steps], sampler.blobs[:, :, :num_steps, 0], np.shape(np.where(~no_diff)), np.shape(np.where(no_diff))) + self.marginalized_param_blobs = self.blobs[:, 1:] + assert not np.isnan(self.marginalized_param_blobs).any(), sampler.blobs[:5, :10] else: self.blobs = sampler.blobs[:num_steps, :].T.reshape(len(self.blob_labels), -1).T no_diff = ((self.lnlikes == self.blobs[:, 0]) | (np.abs(self.lnlikes - self.blobs[:, 0]) < 0.001)) - assert (no_diff).all(), (self.lnlikes[~no_diff], self.blobs[:, 0][~no_diff], np.shape(np.where(~no_diff)), np.shape(np.where(no_diff))) + assert (no_diff).all(), (self.lnlikes[~no_diff][:10], self.blobs[:, 0][~no_diff][:10], np.shape(np.where(~no_diff)), np.shape(np.where(no_diff))) self.marginalized_param_blobs = self.blobs[:, 1:] assert not np.isnan(self.marginalized_param_blobs).any(), sampler.blobs[:5, :10] num_orbits = self.post.shape[0] @@ -1117,7 +1138,8 @@ def run_sampler( logpargs=[ self.priors, ], - pool=pool + pool=pool, + nblobs=len(self.blob_labels) ) else: sampler = emcee.EnsembleSampler( @@ -1134,6 +1156,7 @@ def run_sampler( ): if self.use_pt: self.curr_pos = state[0] + self.force_valid_marg = False else: self.curr_pos = state.coords @@ -1154,6 +1177,7 @@ def run_sampler( sampler.reset() print("") print("Burn in complete. Sampling posterior now.") + self.force_valid_marg = self.use_pt saved_upto = 0 # keep track of how many steps of this chain we've saved. this is the next index that needs to be saved for i, state in enumerate( @@ -1161,6 +1185,7 @@ def run_sampler( ): if self.use_pt: self.curr_pos = state[0] + self.force_valid_marg = False else: self.curr_pos = state.coords diff --git a/tests/test_marg.py b/tests/test_marg.py index 83ce5071..f5ef1af0 100644 --- a/tests/test_marg.py +++ b/tests/test_marg.py @@ -75,7 +75,7 @@ def test_MargPriorNoInput(): myResults = mySampler.results assert (0.5 <= myResults.post[:, 8]).all() and (myResults.post[:, 8] <= 1.5).all() assert np.mean(myResults.post[:, 8]) == pytest.approx(1, rel=0.1) - assert np.std(myResults.post[:, 9]) == pytest.approx(1, rel=0.1) + assert np.std(myResults.post[:, 9]) == pytest.approx(1, rel=0.2) assert np.mean(myResults.post[:, 9]) == pytest.approx(10, rel=0.1) def test_incompatible(): @@ -151,12 +151,13 @@ def test_OFTeI(): marg = marginalizer.OFTeI_sma_pan(epochs, sep, pa, sep_err, pa_err, covar, sma_prior, pan_prior, mass_prior, implicit_mass_prior) mySampler = sampler.MCMC( mySystem, - num_temps=0, + num_temps=5, num_walkers=100, num_threads=mp.cpu_count(), marginalizers=[marg] ) mySampler.run_sampler(4000, 10, output_filename=output_file, periodic_save_freq=2) + assert not mySampler.force_valid_marg myResults = mySampler.results print(myResults.post.shape) diff --git a/tests/test_mcmc.py b/tests/test_mcmc.py index 8a0d7ac0..cf844d0d 100644 --- a/tests/test_mcmc.py +++ b/tests/test_mcmc.py @@ -199,8 +199,8 @@ def test_mcmc_param_idx(): def test_mcmc_runs(): # Parallel Tempering tests (only test corner plot once) - # do_mcmc_runs(num_temps=2, num_threads=1, make_corner_plot=True) - # do_mcmc_runs(num_temps=2, num_threads=4) + do_mcmc_runs(num_temps=2, num_threads=1, make_corner_plot=True) + do_mcmc_runs(num_temps=2, num_threads=4) # Ensemble MCMC tests do_mcmc_runs(num_temps=0, num_threads=1) do_mcmc_runs(num_temps=0, num_threads=8) diff --git a/tests/test_mcmc_rv.py b/tests/test_mcmc_rv.py index 98a6d3e4..1eb0659c 100644 --- a/tests/test_mcmc_rv.py +++ b/tests/test_mcmc_rv.py @@ -36,7 +36,7 @@ def test_pt_mcmc_runs(num_threads=1): # test that lnlikes being saved are correct returned_lnlike_test = myDriver.sampler.results.lnlike[0] - computed_lnlike_test = myDriver.sampler._logl(clean_params) + computed_lnlike_test, *_ = myDriver.sampler._logl(clean_params) assert returned_lnlike_test == pytest.approx(computed_lnlike_test, abs=0.01) @@ -69,7 +69,7 @@ def test_ensemble_mcmc_runs(num_threads=1): # test that lnlikes being saved are correct returned_lnlike_test = myDriver.sampler.results.lnlike[0] - computed_lnlike_test = myDriver.sampler._logl(clean_params) + computed_lnlike_test, *_ = myDriver.sampler._logl(clean_params) assert returned_lnlike_test == pytest.approx(computed_lnlike_test, abs=0.01) From 94ca52003caaf22ea734209ce89761a31a8855f9 Mon Sep 17 00:00:00 2001 From: Eshel Dror Date: Wed, 14 Oct 2026 11:37:37 -0700 Subject: [PATCH 07/21] Add end-to-end test --- tests/end-to-end-tests/GJ504_oftei.py | 87 +++++++++++++++++++++++++++ 1 file changed, 87 insertions(+) create mode 100644 tests/end-to-end-tests/GJ504_oftei.py diff --git a/tests/end-to-end-tests/GJ504_oftei.py b/tests/end-to-end-tests/GJ504_oftei.py new file mode 100644 index 00000000..06496462 --- /dev/null +++ b/tests/end-to-end-tests/GJ504_oftei.py @@ -0,0 +1,87 @@ +import matplotlib.pyplot as plt + +import orbitize + +# Based on driver.py + +import numpy as np +from orbitize import read_input, system, sampler, priors, marginalizer, plot +import multiprocessing as mp +import orbitize +import time + +savedir = "" +savefile = f"{savedir}GJ504_oftei.hdf5" + +start = time.time() + +# System parameters +datafile = "GJ504.csv" +num_secondary_bodies = 1 +system_mass = 1.22 # Msol +plx = 56.95 # mas +mass_err = 0.08 # Msol +plx_err = 0.26 # mas + +# run MCMC +num_threads = mp.cpu_count() +num_temps = 20 +num_walkers = 1000 +num_steps = 1000000 # n_walkers x n_steps_per_walker +burn_steps = 1000 +thin = 100 +tau_ref_epoch = 0 + +# Read in data +data_table = read_input.read_file(orbitize.DATADIR + datafile) + +# Initialize System object which stores data & sets priors +mySystem = system.System( + 1, + data_table, + system_mass, + plx, + tau_ref_epoch=tau_ref_epoch, + fitting_basis="OftEiBasis" +) + +sma_prior_min = 0.1 # au +sma_prior_max = 500 # au +period_prior_min_yr = np.sqrt(sma_prior_min**3 / system_mass) # [yr] +period_prior_max_yr = np.sqrt(sma_prior_max**3 / system_mass) # [yr] +mySystem.sys_priors[mySystem.param_idx["sma1"]] = priors.LogUniformPrior(sma_prior_min, sma_prior_max) +mySystem.sys_priors[mySystem.param_idx["per1"]] = priors.LogUniformPrior(period_prior_min_yr, period_prior_max_yr) + +ind = data_table["quant_type"] == "seppa" +epochs = data_table["epoch"][ind] +sep = data_table["quant1"][ind] +pa = data_table["quant2"][ind] +sep_err = data_table["quant1_err"][ind] +pa_err = data_table["quant2_err"][ind] +covar = np.zeros(len(epochs)) +sma_prior = mySystem.sys_priors[mySystem.param_idx["sma1"]] +pan_prior = mySystem.sys_priors[mySystem.param_idx["pan1"]] +mass_prior = priors.GaussianPrior(system_mass, mass_err) +per_prior = mySystem.sys_priors[mySystem.param_idx["per1"]] +implicit_mass_prior = priors.OFTeIMassPrior(sma_prior.minval, sma_prior.maxval, per_prior.minval, per_prior.maxval) +marg = marginalizer.OFTeI_sma_pan(epochs, sep, pa, sep_err, pa_err, covar, sma_prior, pan_prior, mass_prior, implicit_mass_prior) +mySampler = sampler.MCMC( + mySystem, + num_temps=num_temps, + num_walkers=num_walkers, + num_threads=num_threads, + marginalizers=[marg] +) + +mySampler.run_sampler(num_steps, burn_steps, output_filename=savefile, periodic_save_freq=None) +myResults = mySampler.results + +end = time.time() +print(f"Processing time: {end} - {start} = {(end-start) / 60} minutes") + +# make corner plot +p = plot.Plotter(myResults, num_orbits_to_plot=10000) +fig = p.plot_corner() +fig.savefig(f"{savedir}GJ504_oftei_corner.png") +fig2 = p.plot_orbits() +fig2.savefig(f"{savedir}GJ504_oftei_orbits.png") \ No newline at end of file From ec5f82972d2f9b68ca48376772e5a09215ccfc99 Mon Sep 17 00:00:00 2001 From: Eshel Dror Date: Wed, 14 Oct 2026 15:38:06 -0700 Subject: [PATCH 08/21] Make all samplers function with new _logl, fix bugs --- orbitize/plot.py | 10 ++++++---- orbitize/sampler.py | 35 ++++++++++++++++++++++++++++------- tests/test_OFTI.py | 6 +++--- tests/test_api.py | 6 +++--- tests/test_chi2_lnlike.py | 4 ++-- tests/test_gaia.py | 2 +- tests/test_secondary_rvs.py | 2 +- 7 files changed, 44 insertions(+), 21 deletions(-) diff --git a/orbitize/plot.py b/orbitize/plot.py index ccf6cf44..b8e113ab 100644 --- a/orbitize/plot.py +++ b/orbitize/plot.py @@ -1990,6 +1990,7 @@ def plot_corner(results, param_list=None, plot_priors=True, downsample=None, pri param_list = results.labels param_indices = [] + sys_param_indices = [] angle_indices = [] secondary_mass_indices = [] fixed_indices = [] @@ -1998,7 +1999,8 @@ def plot_corner(results, param_list=None, plot_priors=True, downsample=None, pri # only plot non-fixed parameters if not np.isclose(0.0, np.std(results.post[:, index_num])): - param_indices.append(index_num) + param_indices.append(i) + sys_param_indices.append(index_num) if ( label_key.startswith("aop") or label_key.startswith("pan") @@ -2018,7 +2020,7 @@ def plot_corner(results, param_list=None, plot_priors=True, downsample=None, pri weights = results.weights samples = np.copy( - post[:, param_indices] + post[:, sys_param_indices] ) # keep only chains for selected parameters samples[:, angle_indices] = np.degrees( samples[:, angle_indices] @@ -2062,8 +2064,8 @@ def plot_corner(results, param_list=None, plot_priors=True, downsample=None, pri if plot_priors: axes = figure.axes - num_params = len(param_indices) - for i, param_i in enumerate(param_indices): + num_params = len(sys_param_indices) + for i, param_i in enumerate(sys_param_indices): prior = results.system.sys_priors[param_i] if not hasattr(prior, "compute_lnprob"): continue diff --git a/orbitize/sampler.py b/orbitize/sampler.py index 45625a3f..eb91ea78 100644 --- a/orbitize/sampler.py +++ b/orbitize/sampler.py @@ -113,7 +113,7 @@ def _logl(self, fitted_params): """ input_shape = fitted_params.shape - assert input_shape[0] == len(self.fitted_param_sys_idx) + assert input_shape[0] == len(self.fitted_param_sys_idx), (input_shape[0], len(self.fitted_param_sys_idx)) if len(input_shape) == 1: num_orbits = 1 params = np.full((self.total_param_num,), np.nan) @@ -279,7 +279,7 @@ def __init__( self, system, like="chi2_lnlike", custom_lnlike=None, chi2_type="standard" ): super(OFTI, self).__init__( - system, like=like, chi2_type=chi2_type, custom_lnlike=custom_lnlike + system, like=like, chi2_type=chi2_type, custom_lnlike=custom_lnlike, marginalizers=[] ) if (self.system.hipparcos_IAD is not None) or (len(self.system.rv[0] > 0)): @@ -537,7 +537,7 @@ def reject(self, samples): """ - lnp = self._logl(samples) + lnp, *_ = self._logl(samples) # we just want the chi2 term for rejection, so compute the Gaussian normalization term and remove it errs = np.array( @@ -1496,12 +1496,14 @@ def __init__(self, chi2_type="standard", like="chi2_lnlike", custom_lnlike=None, + marginalizers=None ): super(NestedSampler, self).__init__( system, like=like, chi2_type=chi2_type, custom_lnlike=custom_lnlike, + marginalizers=marginalizers ) # create an empty results object @@ -1515,6 +1517,10 @@ def __init__(self, self.start = time.time() self.dynesty_sampler = None + def _logl(self, params): + lnlike, *blobs = super(NestedSampler, self)._logl(params) + return lnlike, blobs + def run_sampler( self, nlive=500, @@ -1575,6 +1581,7 @@ def run_sampler( bound=bound, bootstrap=False, nlive=nlive, + blob=True, ) else: self.dynesty_sampler = dynesty.DynamicNestedSampler( @@ -1584,7 +1591,8 @@ def run_sampler( pool=pool, bound=bound, bootstrap=False, - nlive=nlive + nlive=nlive, + blob=True ) self.dynesty_sampler.run_nested(**run_nested_kwargs) @@ -1596,6 +1604,7 @@ def run_sampler( len(self.system.sys_priors), bound=bound, nlive=nlive, + blob=True, ) else: self.dynesty_sampler = dynesty.DynamicNestedSampler( @@ -1604,12 +1613,23 @@ def run_sampler( len(self.system.sys_priors), bound=bound, nlive=nlive, + blob=True, ) self.dynesty_sampler.run_nested(**run_nested_kwargs) + blobs = self.dynesty_sampler.results["blob"] + samples = self.dynesty_sampler.results["samples"] + marginalized_params = blobs[:, 1:] + logl = self.dynesty_sampler.results["logl"] + num_orbits = len(logl) + assert (logl == blobs[:, 0]).all(), (np.shape(logl), len(logl), np.shape(blobs[:, 0])) + full_post = np.full((num_orbits, self.total_param_num), np.nan) + full_post[:, self.fitted_param_sys_idx] = samples + full_post[:, self.marg_param_sys_idx] = marginalized_params + self.results.add_samples( - self.dynesty_sampler.results["samples"], - self.dynesty_sampler.results["logl"], + full_post, + logl ) self.results.ln_evidence = self.dynesty_sampler.results["logz"][-1] @@ -1864,7 +1884,8 @@ def nautilus_ptform(self, u): return self.ptform(u.T).T def nautilus_logl(self, u: np.ndarray): - return self._logl(u.T).T + lnlike, *blobs = self._logl(u.T) + return lnlike, blobs def run_sampler( self, diff --git a/tests/test_OFTI.py b/tests/test_OFTI.py index cfb9a65b..de2a7993 100644 --- a/tests/test_OFTI.py +++ b/tests/test_OFTI.py @@ -138,7 +138,7 @@ def test_run_sampler(): # test that lnlikes being saved are correct returned_lnlike_test = s.results.lnlike[0] - computed_lnlike_test = s._logl(orbits[0]) + computed_lnlike_test, *_ = s._logl(orbits[0]) assert returned_lnlike_test == pytest.approx(computed_lnlike_test, abs=0.01) seppa_lnprob_compare = ( @@ -370,12 +370,12 @@ def test_OFTI_covariances(): # test that lnlikes being saved are correct returned_lnlike_test = s.results.lnlike[0] - computed_lnlike_test = s._logl(orbits[0]) + computed_lnlike_test, *_ = s._logl(orbits[0]) assert returned_lnlike_test == pytest.approx(computed_lnlike_test, abs=0.01) # test that the lnlike is very similar to the values computed in seppa space ref_params, ref_lnlike = seppa_lnprob_compare - computed_lnlike_ref = s._logl(ref_params) + computed_lnlike_ref, *_ = s._logl(ref_params) assert ref_lnlike == pytest.approx( computed_lnlike_ref, abs=0.05 ) # 5% differencesin lnprob is allowable. diff --git a/tests/test_api.py b/tests/test_api.py index 63d92046..01cec5d6 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -137,10 +137,10 @@ def my_likelihood(params): orbit, 0, n_walkers, num_threads=1, custom_lnlike=my_likelihood ) - param = np.array([2, 0.5, 0.5, 0.5, 0.5, 0.5, 1, 0.01]) + param = np.array([2, 0.5, 0.5, 0.5, 0.5, 0.5]) - logl1 = mcmc1._logl(param) - logl2 = mcmc2._logl(param) + logl1, *_ = mcmc1._logl(param) + logl2, *_ = mcmc2._logl(param) assert logl1 == logl2 + 5 diff --git a/tests/test_chi2_lnlike.py b/tests/test_chi2_lnlike.py index 5aee6716..45582075 100644 --- a/tests/test_chi2_lnlike.py +++ b/tests/test_chi2_lnlike.py @@ -108,8 +108,8 @@ def test_chi2_log(): chi2_type="log", ) s = myDriver.sampler - params = [44, 0, 45 * np.pi / 180, 0, 325 * np.pi / 180, 0, 56.95, 1.22] - log_chi2 = s._logl(params) + params = np.array([44, 0, 45 * np.pi / 180, 0, 325 * np.pi / 180, 0, 56.95, 1.22]) + log_chi2, *_ = s._logl(params) sys = myDriver.system data = np.array([sys.data_table["quant1"], sys.data_table["quant2"]]).T diff --git a/tests/test_gaia.py b/tests/test_gaia.py index 9f99a583..4a63b212 100644 --- a/tests/test_gaia.py +++ b/tests/test_gaia.py @@ -267,7 +267,7 @@ def test_hgca(): ] ) - chi2 = this_sampler._logl(params) + chi2, *_ = this_sampler._logl(params) assert np.isfinite(chi2) diff --git a/tests/test_secondary_rvs.py b/tests/test_secondary_rvs.py index 5bf143f3..75b6bccb 100644 --- a/tests/test_secondary_rvs.py +++ b/tests/test_secondary_rvs.py @@ -47,7 +47,7 @@ def test_secondary_rv_lnlike_calc(): 1, data_table, m0, plx, mass_err=0.1, plx_err=0.1, fit_secondary_mass=True ) mySamp = sampler.MCMC(mySys) - computed_lnlike = mySamp._logl(orbitize_params_list) + computed_lnlike, *_ = mySamp._logl(orbitize_params_list) # residuals should be 0 assert computed_lnlike == np.sum( From b335d268e9cd83ed2057852059ed6da68af18fcc Mon Sep 17 00:00:00 2001 From: Eshel Dror Date: Thu, 15 Oct 2026 08:28:39 -0700 Subject: [PATCH 09/21] Make OFTI work when converting radec data to seppa --- orbitize/sampler.py | 4 ++++ tests/test_OFTI.py | 2 +- 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/orbitize/sampler.py b/orbitize/sampler.py index 848ccfb1..b90f723c 100644 --- a/orbitize/sampler.py +++ b/orbitize/sampler.py @@ -306,6 +306,10 @@ def __init__( ) convert_warning_print = True self.system.convert_data_table_radec2seppa(body_num=body_num) + if "radec" in self.non_marg_quants: + self.non_marg_quants.remove("radec") + if "seppa" not in self.non_marg_quants: + self.non_marg_quants.append("seppa") # these are of type astropy.table.column self.sep_observed = self.system.data_table[ diff --git a/tests/test_OFTI.py b/tests/test_OFTI.py index fbef3cc6..90c8b0fe 100644 --- a/tests/test_OFTI.py +++ b/tests/test_OFTI.py @@ -390,7 +390,7 @@ def test_OFTI_covariances(): computed_lnlike_ref, *_ = s._logl(ref_params) assert ref_lnlike == pytest.approx( computed_lnlike_ref, abs=0.05 - ) # 5% differencesin lnprob is allowable. + ) # 5% differences in lnprob is allowable. idx = s.system.param_idx sma = np.median([x[idx["sma1"]] for x in orbits]) From 19f5c53d225ff722665b92d83daf1efa6d178b55 Mon Sep 17 00:00:00 2001 From: Eshel Dror Date: Thu, 15 Oct 2026 10:34:11 -0700 Subject: [PATCH 10/21] Add automatic construction of plx prior to OFTeIBasis --- orbitize/basis.py | 95 +++++++++++++++++++++++++++++----------- orbitize/marginalizer.py | 2 +- tests/test_marg.py | 1 + 3 files changed, 72 insertions(+), 26 deletions(-) diff --git a/orbitize/basis.py b/orbitize/basis.py index 66172953..ca4f6d09 100644 --- a/orbitize/basis.py +++ b/orbitize/basis.py @@ -712,42 +712,87 @@ def __init__(self, rv_instruments, gaia_dr4 ) - assert not fit_secondary_mass, "Not supported yet" - assert gaia_dr4 is None, "Not supported yet" + assert not fit_secondary_mass, "Secondary mass not supported yet" + assert num_secondary_bodies == 1, "Multiplanet not supported yet" + assert gaia_dr4 is None, "DR4 not supported yet" - def construct_priors(self, per_min=1e-5, per_max=1e6): - basis_labels = ["sma1", "ecc1", "inc1", "aop1", "pan1", "tau1", "plx", "per1"] - basis_priors = [ - priors.LogUniformPrior(0.001, 1e4), - priors.UniformPrior( - 0, - 1, - ), - priors.SinPrior(), - priors.UniformPrior(0, 2 * np.pi), - priors.UniformPrior(0, 2 * np.pi), - priors.UniformPrior(0, 1), - self.plx, - priors.LogUniformPrior(per_min, per_max), - ] + def construct_priors(self): + """ + + Generates the parameter label array and initializes the corresponding priors for each + parameter that's to be sampled. For the standard basis, the parameters common to each + companion are: per, ecc, inc, aop, pan, tau. Parallax, hipparcos (optional), rv (optional), + and mass priors are added at the end. + + Returns: + tuple: + + list: list of strings (labels) that indicate the names of each parameter to sample + + list: list of orbitize.priors.Prior objects that indicate the prior distribution of each label + """ + base_labels = ["sma", "ecc", "inc", "aop", "pan", "tau", "per"] + basis_priors = [] + basis_labels = [] + + # Define param label dictionary in current basis & standard basis + # self.param_idx = dict(zip(basis_labels, np.arange(len(basis_labels)))) + # self.standard_basis_idx = dict( + # zip(basis_labels[:-1], np.arange(len(basis_labels) - 1)) + # ) + # self.standard_basis_idx["mtot"] = 7 + + # Add the priors common to each companion + for body in np.arange(self.num_secondary_bodies): + for elem in base_labels: + basis_priors.append(self.default_priors[elem]) + basis_labels.append(elem + str(body + 1)) + + # Add parallax prior + basis_labels.append("plx") + if self.plx_err > 0: + basis_priors.append(priors.GaussianPrior(self.plx, self.plx_err)) + else: + basis_priors.append(self.plx) + + # Add hipparcos priors if necessary + if self.hipparcos_IAD is not None: + self.set_hip_iad_priors(basis_priors, basis_labels) + + # add Gaia DR4 priors if necessary + elif self.gaia_dr4 is not None: + self.set_gaia_dr4_priors(basis_priors, basis_labels) + # Add Gaia DR4 jitter prior if fitted + if self.gaia_dr4 is not None and self.gaia_dr4.fit_jitter: + basis_priors.append(self.gaia_dr4.jitter_prior) + basis_labels.append("dr4_jitter") + # Add rv priors + if self.rv and self.fit_secondary_mass: + self.set_rv_priors(basis_priors, basis_labels) + + # Add mass priors TODO: Implicit mass prior + # self.set_default_mass_priors(basis_priors, basis_labels) # Define param label dictionary in current basis & standard basis self.param_idx = dict(zip(basis_labels, np.arange(len(basis_labels)))) - self.standard_basis_idx = dict( - zip(basis_labels[:-1], np.arange(len(basis_labels) - 1)) - ) - self.standard_basis_idx["mtot"] = 7 + self.standard_basis_idx = dict(zip(basis_labels, np.arange(len(basis_labels)))) + + if not self.fit_secondary_mass and self.num_secondary_bodies == 1: + self.standard_basis_idx["mtot"] = self.param_idx["per1"] + else: + assert False return basis_priors, basis_labels def to_standard_basis(self, param_arr): - # everything is the same except the last param element goes from per to mtot - sma = param_arr[0] - per = param_arr[-1] + # everything is the same except per goes to mtot + sma = param_arr[self.param_idx["sma1"]] + per = param_arr[self.param_idx["per1"]] mtot = (sma**3) / (per**2) - param_arr[-1] = mtot + param_arr[self.standard_basis_idx["mtot"]] = mtot + return param_arr class SemiAmp(Basis): diff --git a/orbitize/marginalizer.py b/orbitize/marginalizer.py index ee661db2..4e9875a9 100644 --- a/orbitize/marginalizer.py +++ b/orbitize/marginalizer.py @@ -72,7 +72,7 @@ class OFTeI_sma_pan(Marginalizer): REQ_PARAMS = ["ecc1", "inc1", "aop1", "tau1", "plx", "per1"] MARG_PARAMS = ["sma1", "pan1"] - QUANT_TYPES = ["seppa", "radec"] + QUANT_TYPES = ["seppa"] LNLIKE_TYPE = "chi2" def __init__(self, epochs, seps, pas, sigma_seps, sigma_pas, covar_seps_pas, sma_prior=None, pan_prior=None, mass_prior=None, implicit_mass_prior=None): diff --git a/tests/test_marg.py b/tests/test_marg.py index 23e43e48..09d90899 100644 --- a/tests/test_marg.py +++ b/tests/test_marg.py @@ -177,6 +177,7 @@ def test_OFTeI_RV(): data_table, mtot, plx, + plx_err=plx_err, tau_ref_epoch=tau_ref_epoch, fitting_basis="OftEiBasis" ) From 0d919e34cd807a873b8457fe54a2a31ea4239940 Mon Sep 17 00:00:00 2001 From: Eshel Dror Date: Thu, 15 Oct 2026 10:54:10 -0700 Subject: [PATCH 11/21] Add test for radec with OFTeI --- orbitize/marginalizer.py | 6 +- orbitize/sampler.py | 4 +- tests/test_marg.py | 120 ++++++++++++++++++++++++++++++++++++++- 3 files changed, 125 insertions(+), 5 deletions(-) diff --git a/orbitize/marginalizer.py b/orbitize/marginalizer.py index 4e9875a9..8f43b827 100644 --- a/orbitize/marginalizer.py +++ b/orbitize/marginalizer.py @@ -300,13 +300,15 @@ def apply(self, param_arr, jitter_arr, calculated_values): return logprob, loglike, samples, calculated_values -def compatible_params(marginalizers, input_params): +def compatible_params(marginalizers, input_params, data_quants): # TODO: Multibody, quant_type params = set(input_params) + quants = set(data_quants) for marginalizer in marginalizers: req_params = set(marginalizer.REQ_PARAMS) marg_params = set(marginalizer.MARG_PARAMS) - if len(req_params - params) == 0 and len(marg_params & params) == 0: + marg_quants = set(marginalizer.QUANT_TYPES) + if len(req_params - params) == 0 and len(marg_params & params) == 0 and len(quants & marg_quants) > 0: params |= set(marginalizer.MARG_PARAMS) else: return False diff --git a/orbitize/sampler.py b/orbitize/sampler.py index bd5db2c1..b03bd4a7 100644 --- a/orbitize/sampler.py +++ b/orbitize/sampler.py @@ -88,8 +88,8 @@ def __init__( self.non_marg_quants = list(quants - marg_quants) self.force_valid_marg = False - if not orbitize.marginalizer.compatible_params(self.marginalizers, self.fitted_param_labels): - raise ValueError("Marginalizers and basis are incompatible") + if not orbitize.marginalizer.compatible_params(self.marginalizers, self.fitted_param_labels, quants): + raise ValueError("Marginalizers and basis or data are incompatible") @abc.abstractmethod def run_sampler(self, total_orbits): diff --git a/tests/test_marg.py b/tests/test_marg.py index 09d90899..00f06216 100644 --- a/tests/test_marg.py +++ b/tests/test_marg.py @@ -164,7 +164,7 @@ def test_OFTeI(): def test_OFTeI_RV(): input_file = os.path.join(orbitize.DATADIR, 'HD4747.csv') data_table = read_input.read_file(input_file) - output_file = os.path.join(orbitize.DATADIR, 'test_marg.hdf5') + output_file = os.path.join(orbitize.DATADIR, 'test_marg_rv.hdf5') mtot = 0.84 # [Msol] plx = 53.18 # [mas] @@ -212,6 +212,124 @@ def test_OFTeI_RV(): assert mySampler.non_marg_quants == ["rv"] print(myResults.post.shape) +def test_OFTeI_radec(): + input_file = os.path.join(orbitize.DATADIR, 'GJ504.csv') + data_table = read_input.read_file(input_file) + output_file = os.path.join(orbitize.DATADIR, 'test_marg_radec.hdf5') + + # read in seppa data table and turn into raddec data table + data_table = orbitize.read_input.read_file(input_file) + data_ra, data_dec = system.seppa2radec(data_table["quant1"], data_table["quant2"]) + data_raerr, data_decerr, data_radeccorr = [], [], [] + + for row in data_table: + raerr, decerr, radec_corr = system.transform_errors( + row["quant1"], + row["quant2"], + row["quant1_err"], + row["quant2_err"], + 0, + system.seppa2radec, + nsamps=10000000, + ) + data_raerr.append(raerr) + data_decerr.append(decerr) + data_radeccorr.append(radec_corr) + + data_table["quant1"] = data_ra + data_table["quant2"] = data_dec + data_table["quant1_err"] = np.array(data_raerr) + data_table["quant2_err"] = np.array(data_decerr) + data_table["quant12_corr"] = np.array(data_radeccorr) + data_table["quant_type"] = np.array(["radec" for _ in data_table]) + + mtot = 1.22 # Msol + plx = 56.95 # mas + mass_err = 0.08 # Msol + plx_err = 0.26 # mas + tau_ref_epoch = 0 + + try: + mySystem = system.System( + 1, + data_table, + mtot, + plx, + plx_err=plx_err, + tau_ref_epoch=tau_ref_epoch, + fitting_basis="OftEiBasis" + ) + sma_prior_min = 0.1 # au + sma_prior_max = 500 # au + period_prior_min_yr = np.sqrt(sma_prior_min**3 / mtot) # [yr] + period_prior_max_yr = np.sqrt(sma_prior_max**3 / mtot) # [yr] + mySystem.sys_priors[mySystem.param_idx["sma1"]] = priors.LogUniformPrior(sma_prior_min, sma_prior_max) + mySystem.sys_priors[mySystem.param_idx["per1"]] = priors.LogUniformPrior(period_prior_min_yr, period_prior_max_yr) + ind = data_table["quant_type"] == "seppa" + epochs = data_table["epoch"][ind] + sep = data_table["quant1"][ind] + pa = data_table["quant2"][ind] + sep_err = data_table["quant1_err"][ind] + pa_err = data_table["quant2_err"][ind] + covar = np.zeros(len(epochs)) + sma_prior = mySystem.sys_priors[mySystem.param_idx["sma1"]] + pan_prior = mySystem.sys_priors[mySystem.param_idx["pan1"]] + mass_prior = priors.GaussianPrior(mtot, mass_err) + per_prior = mySystem.sys_priors[mySystem.param_idx["per1"]] + implicit_mass_prior = priors.OFTeIMassPrior(sma_prior.minval, sma_prior.maxval, per_prior.minval, per_prior.maxval) + marg = marginalizer.OFTeI_sma_pan(epochs, sep, pa, sep_err, pa_err, covar, sma_prior, pan_prior, mass_prior, implicit_mass_prior) + mySampler = sampler.MCMC( + mySystem, + num_temps=5, + num_walkers=100, + num_threads=mp.cpu_count(), + marginalizers=[marg] + ) + except ValueError: + assert True + else: + assert False + mySystem = system.System( + 1, + data_table, + mtot, + plx, + plx_err=plx_err, + tau_ref_epoch=tau_ref_epoch, + fitting_basis="OftEiBasis" + ) + mySystem.convert_data_table_radec2seppa() + sma_prior_min = 0.1 # au + sma_prior_max = 500 # au + period_prior_min_yr = np.sqrt(sma_prior_min**3 / mtot) # [yr] + period_prior_max_yr = np.sqrt(sma_prior_max**3 / mtot) # [yr] + mySystem.sys_priors[mySystem.param_idx["sma1"]] = priors.LogUniformPrior(sma_prior_min, sma_prior_max) + mySystem.sys_priors[mySystem.param_idx["per1"]] = priors.LogUniformPrior(period_prior_min_yr, period_prior_max_yr) + ind = data_table["quant_type"] == "seppa" + epochs = data_table["epoch"][ind] + sep = data_table["quant1"][ind] + pa = data_table["quant2"][ind] + sep_err = data_table["quant1_err"][ind] + pa_err = data_table["quant2_err"][ind] + covar = np.zeros(len(epochs)) + sma_prior = mySystem.sys_priors[mySystem.param_idx["sma1"]] + pan_prior = mySystem.sys_priors[mySystem.param_idx["pan1"]] + mass_prior = priors.GaussianPrior(mtot, mass_err) + per_prior = mySystem.sys_priors[mySystem.param_idx["per1"]] + implicit_mass_prior = priors.OFTeIMassPrior(sma_prior.minval, sma_prior.maxval, per_prior.minval, per_prior.maxval) + marg = marginalizer.OFTeI_sma_pan(epochs, sep, pa, sep_err, pa_err, covar, sma_prior, pan_prior, mass_prior, implicit_mass_prior) + mySampler = sampler.MCMC( + mySystem, + num_temps=5, + num_walkers=100, + num_threads=mp.cpu_count(), + marginalizers=[marg] + ) + mySampler.run_sampler(4000, 10, output_filename=output_file, periodic_save_freq=2) + assert not mySampler.force_valid_marg + myResults = mySampler.results + print(myResults.post.shape) + if __name__ == "__main__": # test_MargConstNoInput() test_OFTeI() \ No newline at end of file From 1b716d9f0a903322e307b4407a5f20c29b261a3b Mon Sep 17 00:00:00 2001 From: Eshel Dror Date: Mon, 19 Oct 2026 17:47:26 -0700 Subject: [PATCH 12/21] Clear growing sampler attributes --- orbitize/sampler.py | 164 +++++++++++++++++++++++++------------------- 1 file changed, 92 insertions(+), 72 deletions(-) diff --git a/orbitize/sampler.py b/orbitize/sampler.py index b03bd4a7..2a3e5fc2 100644 --- a/orbitize/sampler.py +++ b/orbitize/sampler.py @@ -20,6 +20,7 @@ import orbitize.marginalizer import sys +import os class Sampler(abc.ABC): @@ -839,15 +840,9 @@ def __init__( self.num_temps = num_temps self.num_walkers = num_walkers self.num_threads = num_threads + self.prev_result_filename = prev_result_filename - # create an empty results object - self.results = orbitize.results.Results( - self.system, - sampler_name=self.__class__.__name__, - post=None, - lnlike=None, - version_number=orbitize.__version__, - ) + self.results = None # get priors from the system class. need to remove and record fixed priors self.priors = [] @@ -886,6 +881,14 @@ def __init__( self.curr_pos = np.dstack(init_positions) else: + # create an empty results object + self.results = orbitize.results.Results( + self.system, + sampler_name=self.__class__.__name__, + post=None, + lnlike=None, + version_number=orbitize.__version__, + ) # restart from previous walker positions self.results.load_results(prev_result_filename, append=True) @@ -901,6 +904,7 @@ def __init__( ) self.curr_pos = prev_pos + self.results = None # Clear variable def _fill_in_fixed_params(self, sampled_params): """ @@ -969,7 +973,7 @@ def _logl(self, params, include_logp=False): def _update_chains_from_sampler(self, sampler, num_steps=None): """ - Updates self.post and self.lnlike from the MCMC sampler + Updates fitted_post and self.lnlike from the MCMC sampler Args: sampler (ptemcee.Sampler): sampler object. @@ -982,28 +986,28 @@ def _update_chains_from_sampler(self, sampler, num_steps=None): num_params = sampler.chain.shape[-1] # chain has shape Ntemp x Nwalkers x Nsteps x Nparams - self.post = sampler.chain[0, :, :num_steps].reshape( + fitted_post = sampler.chain[0, :, :num_steps].reshape( -1, num_params ) # the reshaping flattens the chain # pick out the lowest temperature loglikelihoods self.lnlikes = sampler.loglikelihood[0, :, :num_steps].flatten() - self.lnlikes_alltemps = sampler.loglikelihood[:, :, :num_steps] + lnlikes_alltemps = sampler.loglikelihood[:, :, :num_steps] # include fixed parameters in posterior - self.post = self._fill_in_fixed_params(self.post) + fitted_post = self._fill_in_fixed_params(fitted_post) # Ntemp x Nwalkers x Nsteps x Nblobs - self.blobs = sampler.blobs[0, :, :num_steps, :].reshape(-1, len(self.blob_labels)) - assert len(self.blobs) == len(self.lnlikes) - no_diff = ((self.lnlikes == self.blobs[:, 0]) | (np.abs(self.lnlikes - self.blobs[:, 0]) < 0.001)) - assert (no_diff).all(), (self.lnlikes[~no_diff][:10], self.blobs[:, 0][~no_diff][:10], np.shape(np.where(~no_diff)), np.shape(np.where(no_diff))) + blobs = sampler.blobs[0, :, :num_steps, :].reshape(-1, len(self.blob_labels)) + assert len(blobs) == len(self.lnlikes) + no_diff = ((self.lnlikes == blobs[:, 0]) | (np.abs(self.lnlikes - blobs[:, 0]) < 0.001)) + assert (no_diff).all(), (self.lnlikes[~no_diff][:10], blobs[:, 0][~no_diff][:10], np.shape(np.where(~no_diff)), np.shape(np.where(no_diff))) # assert (no_diff).all(), (sampler.loglikelihood[:, :, :num_steps], sampler.blobs[:, :, :num_steps, 0], np.shape(np.where(~no_diff)), np.shape(np.where(no_diff))) - self.marginalized_param_blobs = self.blobs[:, 1:] - assert not np.isnan(self.marginalized_param_blobs).any(), sampler.blobs[:5, :10] - num_orbits = self.post.shape[0] + marginalized_param_blobs = blobs[:, 1:] + assert not np.isnan(marginalized_param_blobs).any(), sampler.blobs[:5, :10] + num_orbits = fitted_post.shape[0] self.full_post = np.full((num_orbits, self.total_param_num), np.nan) - self.full_post[:, self.fitted_param_sys_idx] = self.post - self.full_post[:, self.marg_param_sys_idx] = self.marginalized_param_blobs + self.full_post[:, self.fitted_param_sys_idx] = fitted_post + self.full_post[:, self.marg_param_sys_idx] = marginalized_param_blobs assert not np.isnan(self.full_post).any(), (self.fitted_param_sys_idx, self.marg_param_sys_idx) @@ -1088,11 +1092,27 @@ def run_sampler( if output_filename is not None: if output_filename.endswith('.hdf5'): output_filename = output_filename.split('.hdf5')[0] + if output_filename.endswith('.h5'): + output_filename = output_filename.split('.h5')[0] + + if self.prev_result_filename is not None and periodic_save_freq is not None and output_filename is not None: + if self.prev_result_filename.endswith('.hdf5') and self.prev_result_filename.startswith(output_filename): + raise ValueError(f"Output file {output_filename} may overwrite previous results {self.prev_result_filename} ") + nsteps = int(np.ceil(total_orbits / self.num_walkers)) if nsteps <= 0: raise ValueError("Total_orbits must be greater than num_walkers.") + # Initialize results object to save curr_pos + self.results = orbitize.results.Results( + self.system, + sampler_name=self.__class__.__name__, + post=None, + lnlike=None, + version_number=orbitize.__version__, + ) + with mp.Pool(processes=self.num_threads) as pool: sampler = ptemcee.Sampler( self.num_walkers, @@ -1128,6 +1148,12 @@ def run_sampler( if periodic_save_freq is not None: if (i + 1) % periodic_save_freq == 0: # we've completed i+1 steps self.results.curr_pos = self.curr_pos + new_file = f"{output_filename}_{i+1}burnsteps.hdf5" + self.results.save_results(new_file) + assert os.path.exists(new_file) + prev_file = f"{output_filename}_{i + 1 - periodic_save_freq}burnsteps.hdf5" + if os.path.exists(prev_file): + os.remove(prev_file) self.results.save_results(f'{output_filename}_{i+1}burnsteps.hdf5') sampler.reset() @@ -1135,7 +1161,6 @@ def run_sampler( print("Burn in complete. Sampling posterior now.") self.force_valid_marg = True - saved_upto = 0 for i, state in enumerate( sampler.sample(self.curr_pos, iterations=nsteps, thin=thin) ): @@ -1150,63 +1175,58 @@ def run_sampler( if (i + 1) % periodic_save_freq == 0: # we've completed i+1 steps self._update_chains_from_sampler(sampler, num_steps=i + 1) - - # figure out what is the new chunk of the chain and corresponding lnlikes that have been computed before last save - # grab the current posterior and lnlikes and reshape them to have the Nwalkers x Nsteps dimension again - post_shape = self.full_post.shape - curr_chain_shape = ( - self.num_walkers, - post_shape[0] // self.num_walkers, - post_shape[-1], - ) - curr_chain = self.full_post.reshape(curr_chain_shape) - curr_lnlike_chain = self.lnlikes.reshape(curr_chain_shape[:2]) - # use the reshaped arrays and find the new steps we computed - curr_chunk = curr_chain[:, saved_upto : i + 1] - curr_chunk = curr_chunk.reshape( - -1, curr_chunk.shape[-1] - ) # flatten nwalkers x nsteps dim - curr_lnlike_chunk = curr_lnlike_chain[ - :, saved_upto : i + 1 - ].flatten() - - # add this current chunk to the results object (which already has all the previous chunks saved) - self.results.add_samples( - curr_chunk, curr_lnlike_chunk, curr_pos=self.curr_pos - ) - self.results.save_results(output_filename) - saved_upto = i + 1 + if self.prev_result_filename is None: + self.results = orbitize.results.Results( + self.system, + sampler_name=self.__class__.__name__, + post=self.full_post, + lnlike=self.lnlikes, + curr_pos = self.curr_pos, + version_number=orbitize.__version__, + ) + else: + self.results = orbitize.results.Results( + self.system, + sampler_name=self.__class__.__name__, + post=None, + lnlike=None, + version_number=orbitize.__version__, + ) + self.results.load_results(self.prev_result_filename, append=True) + self.results.add_samples(self.full_post, self.lnlikes) + self.results.curr_pos = self.curr_pos + new_file = f"{output_filename}_{i+1}steps.hdf5" + self.results.save_results(new_file) + assert os.path.exists(new_file) + prev_file = f"{output_filename}_{i + 1 - periodic_save_freq}steps.hdf5" + if os.path.exists(prev_file): + os.remove(prev_file) + self.full_post = self.lnlikes = self.results = None # Clear large variables print("") self._update_chains_from_sampler(sampler) - if periodic_save_freq is None: - # need to save everything - self.results.add_samples( - self.full_post, self.lnlikes, curr_pos=self.curr_pos - ) - elif saved_upto < nsteps: - # just need to save the last few - # same code as above except we just need to grab the last few - post_shape = self.full_post.shape - curr_chain_shape = ( - self.num_walkers, - post_shape[0] // self.num_walkers, - post_shape[-1], + if self.prev_result_filename is None: + self.results = orbitize.results.Results( + self.system, + sampler_name=self.__class__.__name__, + post=self.full_post, + lnlike=self.lnlikes, + curr_pos=self.curr_pos, + version_number=orbitize.__version__, ) - curr_chain = self.full_post.reshape(curr_chain_shape) - curr_lnlike_chain = self.lnlikes.reshape(curr_chain_shape[:2]) - curr_chunk = curr_chain[:, saved_upto:] - curr_chunk = curr_chunk.reshape( - -1, curr_chunk.shape[-1] - ) # flatten nwalkers x nsteps dim - curr_lnlike_chunk = curr_lnlike_chain[:, saved_upto:].flatten() - - self.results.add_samples( - curr_chunk, curr_lnlike_chunk, curr_pos=self.curr_pos + else: + self.results = orbitize.results.Results( + self.system, + sampler_name=self.__class__.__name__, + post=None, + lnlike=None, + version_number=orbitize.__version__, ) - + self.results.load_results(self.prev_result_filename, append=True) + self.results.add_samples(self.full_post, self.lnlikes) + self.results.curr_pos = self.curr_pos if output_filename is not None: self.results.save_results(f'{output_filename}.hdf5') From f9194a04d0fc4acbd730998d1bebcd1b8e93fff9 Mon Sep 17 00:00:00 2001 From: Eshel Dror Date: Mon, 19 Oct 2026 17:57:11 -0700 Subject: [PATCH 13/21] Allow tests to pass with compatability checker --- orbitize/marginalizer.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/orbitize/marginalizer.py b/orbitize/marginalizer.py index 8f43b827..9ffb594b 100644 --- a/orbitize/marginalizer.py +++ b/orbitize/marginalizer.py @@ -308,7 +308,8 @@ def compatible_params(marginalizers, input_params, data_quants): req_params = set(marginalizer.REQ_PARAMS) marg_params = set(marginalizer.MARG_PARAMS) marg_quants = set(marginalizer.QUANT_TYPES) - if len(req_params - params) == 0 and len(marg_params & params) == 0 and len(quants & marg_quants) > 0: + if len(req_params - params) == 0 and len(marg_params & params) == 0 and (len(marg_quants) == 0 or len(quants & marg_quants) > 0): + # TODO: differentiate between requiring all quant types vs any quant types params |= set(marginalizer.MARG_PARAMS) else: return False From 2e9d6f670c430b43478788618d3ddb1eebfed7e7 Mon Sep 17 00:00:00 2001 From: Eshel Dror Date: Mon, 19 Oct 2026 19:22:05 -0700 Subject: [PATCH 14/21] Add marginalization auto-construction, add back MCMC functionality to resume fits without saving --- orbitize/marginalizer.py | 2 +- orbitize/sampler.py | 136 ++++++++++++++++++++++++++++++--------- tests/test_marg.py | 76 ++++++++++++++++++++++ tests/test_mcmc.py | 6 +- 4 files changed, 186 insertions(+), 34 deletions(-) diff --git a/orbitize/marginalizer.py b/orbitize/marginalizer.py index 9ffb594b..c7bfedd9 100644 --- a/orbitize/marginalizer.py +++ b/orbitize/marginalizer.py @@ -308,7 +308,7 @@ def compatible_params(marginalizers, input_params, data_quants): req_params = set(marginalizer.REQ_PARAMS) marg_params = set(marginalizer.MARG_PARAMS) marg_quants = set(marginalizer.QUANT_TYPES) - if len(req_params - params) == 0 and len(marg_params & params) == 0 and (len(marg_quants) == 0 or len(quants & marg_quants) > 0): + if not (req_params - params) and not (marg_params & params) and (not (marg_quants) or (quants & marg_quants)): # TODO: differentiate between requiring all quant types vs any quant types params |= set(marginalizer.MARG_PARAMS) else: diff --git a/orbitize/sampler.py b/orbitize/sampler.py index 2a3e5fc2..16eb3beb 100644 --- a/orbitize/sampler.py +++ b/orbitize/sampler.py @@ -49,8 +49,33 @@ def __init__( if marginalizers is None: self.marginalizers = [] + start_quants = set(np.unique(self.system.data_table["quant_type"])) + params = set(self.system.labels) + if (self.system.num_secondary_bodies == 1 and + len({"radec", "seppa"} & start_quants) > 0 and + not (set(orbitize.marginalizer.OFTeI_sma_pan.MARG_PARAMS) - params) and + not (set(orbitize.marginalizer.OFTeI_sma_pan.REQ_PARAMS) - params)): + # OFTeI + self.system.convert_data_table_radec2seppa() + self.has_corr = np.any(~np.isnan(self.system.data_table["quant12_corr"])) + ind = self.system.data_table["quant_type"] == "seppa" + epochs = self.system.data_table["epoch"][ind] + sep = self.system.data_table["quant1"][ind] + pa = self.system.data_table["quant2"][ind] + sep_err = self.system.data_table["quant1_err"][ind] + pa_err = self.system.data_table["quant2_err"][ind] + if self.has_corr: + covar = self.system.data_table["quant12_corr"][ind] + else: + covar = np.zeros(len(epochs)) + sma_prior = self.system.sys_priors[self.system.param_idx["sma1"]] + pan_prior = self.system.sys_priors[self.system.param_idx["pan1"]] + mass_prior = orbitize.priors.GaussianPrior(self.system.stellar_or_system_mass, self.system.mass_err) + per_prior = self.system.sys_priors[self.system.param_idx["per1"]] + implicit_mass_prior = orbitize.priors.OFTeIMassPrior(sma_prior.minval, sma_prior.maxval, per_prior.minval, per_prior.maxval) + marg = orbitize.marginalizer.OFTeI_sma_pan(epochs, sep, pa, sep_err, pa_err, covar, sma_prior, pan_prior, mass_prior, implicit_mass_prior) + self.marginalizers = [marg] else: - # TODO: Construct if not empty list self.marginalizers = marginalizers @@ -289,7 +314,7 @@ def __init__( self, system, like="chi2_lnlike", custom_lnlike=None, chi2_type="standard" ): super(OFTI, self).__init__( - system, like=like, chi2_type=chi2_type, custom_lnlike=custom_lnlike, marginalizers=[] + system, like=like, chi2_type=chi2_type, custom_lnlike=custom_lnlike, marginalizers=[], ) if (self.system.hipparcos_IAD is not None) or (len(self.system.rv[0] > 0)): @@ -834,7 +859,7 @@ def __init__( marginalizers=None ): super(MCMC, self).__init__( - system, like=like, chi2_type=chi2_type, custom_lnlike=custom_lnlike, marginalizers=marginalizers + system, like=like, chi2_type=chi2_type, custom_lnlike=custom_lnlike, marginalizers=marginalizers, ) self.num_temps = num_temps @@ -1104,14 +1129,16 @@ def run_sampler( if nsteps <= 0: raise ValueError("Total_orbits must be greater than num_walkers.") - # Initialize results object to save curr_pos - self.results = orbitize.results.Results( - self.system, - sampler_name=self.__class__.__name__, - post=None, - lnlike=None, - version_number=orbitize.__version__, - ) + using_unsaved_results = self.results is not None + if not using_unsaved_results: + # Initialize results object to save curr_pos + self.results = orbitize.results.Results( + self.system, + sampler_name=self.__class__.__name__, + post=None, + lnlike=None, + version_number=orbitize.__version__, + ) with mp.Pool(processes=self.num_threads) as pool: sampler = ptemcee.Sampler( @@ -1153,14 +1180,13 @@ def run_sampler( assert os.path.exists(new_file) prev_file = f"{output_filename}_{i + 1 - periodic_save_freq}burnsteps.hdf5" if os.path.exists(prev_file): - os.remove(prev_file) - self.results.save_results(f'{output_filename}_{i+1}burnsteps.hdf5') - + os.remove(prev_file) sampler.reset() print("") print("Burn in complete. Sampling posterior now.") self.force_valid_marg = True + saved_upto = 0 for i, state in enumerate( sampler.sample(self.curr_pos, iterations=nsteps, thin=thin) ): @@ -1175,13 +1201,38 @@ def run_sampler( if (i + 1) % periodic_save_freq == 0: # we've completed i+1 steps self._update_chains_from_sampler(sampler, num_steps=i + 1) - if self.prev_result_filename is None: + if using_unsaved_results: + # figure out what is the new chunk of the chain and corresponding lnlikes that have been computed before last save + # grab the current posterior and lnlikes and reshape them to have the Nwalkers x Nsteps dimension again + post_shape = self.full_post.shape + curr_chain_shape = ( + self.num_walkers, + post_shape[0] // self.num_walkers, + post_shape[-1], + ) + curr_chain = self.full_post.reshape(curr_chain_shape) + curr_lnlike_chain = self.lnlikes.reshape(curr_chain_shape[:2]) + # use the reshaped arrays and find the new steps we computed + curr_chunk = curr_chain[:, saved_upto : i + 1] + curr_chunk = curr_chunk.reshape( + -1, curr_chunk.shape[-1] + ) # flatten nwalkers x nsteps dim + curr_lnlike_chunk = curr_lnlike_chain[ + :, saved_upto : i + 1 + ].flatten() + + # add this current chunk to the results object (which already has all the previous chunks saved) + self.results.add_samples( + curr_chunk, curr_lnlike_chunk, curr_pos=self.curr_pos + ) + saved_upto = i + 1 + elif self.prev_result_filename is None: self.results = orbitize.results.Results( self.system, sampler_name=self.__class__.__name__, post=self.full_post, lnlike=self.lnlikes, - curr_pos = self.curr_pos, + curr_pos=self.curr_pos, version_number=orbitize.__version__, ) else: @@ -1193,29 +1244,48 @@ def run_sampler( version_number=orbitize.__version__, ) self.results.load_results(self.prev_result_filename, append=True) - self.results.add_samples(self.full_post, self.lnlikes) - self.results.curr_pos = self.curr_pos + self.results.add_samples(self.full_post, self.lnlikes, curr_pos=self.curr_pos) new_file = f"{output_filename}_{i+1}steps.hdf5" self.results.save_results(new_file) assert os.path.exists(new_file) prev_file = f"{output_filename}_{i + 1 - periodic_save_freq}steps.hdf5" if os.path.exists(prev_file): os.remove(prev_file) - self.full_post = self.lnlikes = self.results = None # Clear large variables + self.full_post = self.lnlikes = None # Clear large variables + if not using_unsaved_results: + self.results = None print("") self._update_chains_from_sampler(sampler) - if self.prev_result_filename is None: - self.results = orbitize.results.Results( - self.system, - sampler_name=self.__class__.__name__, - post=self.full_post, - lnlike=self.lnlikes, - curr_pos=self.curr_pos, - version_number=orbitize.__version__, - ) + if using_unsaved_results: + if periodic_save_freq is None: + # need to save everything + self.results.add_samples( + self.full_post, self.lnlikes, curr_pos=self.curr_pos + ) + elif saved_upto < nsteps: + # just need to save the last few + # same code as above except we just need to grab the last few + post_shape = self.full_post.shape + curr_chain_shape = ( + self.num_walkers, + post_shape[0] // self.num_walkers, + post_shape[-1], + ) + curr_chain = self.full_post.reshape(curr_chain_shape) + curr_lnlike_chain = self.lnlikes.reshape(curr_chain_shape[:2]) + curr_chunk = curr_chain[:, saved_upto:] + curr_chunk = curr_chunk.reshape( + -1, curr_chunk.shape[-1] + ) # flatten nwalkers x nsteps dim + curr_lnlike_chunk = curr_lnlike_chain[:, saved_upto:].flatten() + + self.results.add_samples( + curr_chunk, curr_lnlike_chunk, curr_pos=self.curr_pos + ) + else: self.results = orbitize.results.Results( self.system, @@ -1224,9 +1294,9 @@ def run_sampler( lnlike=None, version_number=orbitize.__version__, ) - self.results.load_results(self.prev_result_filename, append=True) - self.results.add_samples(self.full_post, self.lnlikes) - self.results.curr_pos = self.curr_pos + if self.prev_result_filename is not None: + self.results.load_results(self.prev_result_filename, append=True) + self.results.add_samples(self.full_post, self.lnlikes, curr_pos=self.curr_pos) if output_filename is not None: self.results.save_results(f'{output_filename}.hdf5') @@ -1479,7 +1549,7 @@ def __init__(self, like=like, chi2_type=chi2_type, custom_lnlike=custom_lnlike, - marginalizers=marginalizers + marginalizers=marginalizers, ) # create an empty results object @@ -1648,6 +1718,7 @@ def __init__(self, like=like, chi2_type=chi2_type, custom_lnlike=custom_lnlike, + marginalizers=[], ) # create an empty results object @@ -1844,6 +1915,7 @@ def __init__(self, like=like, chi2_type=chi2_type, custom_lnlike=custom_lnlike, + marginalizers=[], ) # create an empty results object diff --git a/tests/test_marg.py b/tests/test_marg.py index 00f06216..3de4fc73 100644 --- a/tests/test_marg.py +++ b/tests/test_marg.py @@ -212,6 +212,80 @@ def test_OFTeI_RV(): assert mySampler.non_marg_quants == ["rv"] print(myResults.post.shape) +def test_OFTeI_RV_auto(): + input_file = os.path.join(orbitize.DATADIR, 'HD4747.csv') + data_table = read_input.read_file(input_file) + output_file = os.path.join(orbitize.DATADIR, 'test_marg_rv.hdf5') + + mtot = 0.84 # [Msol] + plx = 53.18 # [mas] + mass_err = 0.04 # [Msol] + plx_err = 0.12 # [mas] + tau_ref_epoch = 0 + + mySystem = system.System( + 1, + data_table, + mtot, + plx, + plx_err=plx_err, + mass_err=mass_err, + tau_ref_epoch=tau_ref_epoch, + fitting_basis="OftEiBasis" + ) + sma_prior_min = 0.1 # au + sma_prior_max = 500 # au + period_prior_min_yr = np.sqrt(sma_prior_min**3 / mtot) # [yr] + period_prior_max_yr = np.sqrt(sma_prior_max**3 / mtot) # [yr] + mySystem.sys_priors[mySystem.param_idx["sma1"]] = priors.LogUniformPrior(sma_prior_min, sma_prior_max) + mySystem.sys_priors[mySystem.param_idx["per1"]] = priors.LogUniformPrior(period_prior_min_yr, period_prior_max_yr) + mySampler = sampler.MCMC( + mySystem, + num_temps=2, + num_walkers=100, + num_threads=mp.cpu_count(), + marginalizers=None # automatically construct marginalization + ) + assert len(mySampler.marginalizers) == 1 + marg = mySampler.marginalizers[0] + assert isinstance(marg,marginalizer.OFTeI_sma_pan) + + ind = data_table["quant_type"] == "seppa" + epochs = data_table["epoch"][ind] + sep = data_table["quant1"][ind] + pa = data_table["quant2"][ind] + sep_err = data_table["quant1_err"][ind] + pa_err = data_table["quant2_err"][ind] + covar = np.zeros(len(epochs)) + sma_prior = mySystem.sys_priors[mySystem.param_idx["sma1"]] + pan_prior = mySystem.sys_priors[mySystem.param_idx["pan1"]] + mass_prior = priors.GaussianPrior(mtot, mass_err) + per_prior = mySystem.sys_priors[mySystem.param_idx["per1"]] + implicit_mass_prior = priors.OFTeIMassPrior(sma_prior.minval, sma_prior.maxval, per_prior.minval, per_prior.maxval) + true_marg = marginalizer.OFTeI_sma_pan(epochs, sep, pa, sep_err, pa_err, covar, sma_prior, pan_prior, mass_prior, implicit_mass_prior) + + assert true_marg.sma_prior == marg.sma_prior + assert true_marg.pan_prior == marg.pan_prior + assert true_marg.mass_prior.mu == marg.mass_prior.mu + assert true_marg.mass_prior.sigma == marg.mass_prior.sigma + assert true_marg.implicit_mass_prior.sma_min == marg.implicit_mass_prior.sma_min + assert true_marg.implicit_mass_prior.sma_max == marg.implicit_mass_prior.sma_max + assert true_marg.implicit_mass_prior.per_min == marg.implicit_mass_prior.per_min + assert true_marg.implicit_mass_prior.per_max == marg.implicit_mass_prior.per_max + assert (true_marg.epochs == marg.epochs).all() + assert (true_marg.seps == marg.seps).all() + assert (true_marg.pas == marg.pas).all() + assert (true_marg.sigma_seps == marg.sigma_seps).all() + assert (true_marg.sigma_pas == marg.sigma_pas).all() + assert (true_marg.covar_seps_pas == marg.covar_seps_pas).all() + assert true_marg.n_obs == marg.n_obs + + + mySampler.run_sampler(4000, 10, output_filename=output_file, periodic_save_freq=2) + myResults = mySampler.results + assert mySampler.non_marg_quants == ["rv"] + print(myResults.post.shape) + def test_OFTeI_radec(): input_file = os.path.join(orbitize.DATADIR, 'GJ504.csv') data_table = read_input.read_file(input_file) @@ -330,6 +404,8 @@ def test_OFTeI_radec(): myResults = mySampler.results print(myResults.post.shape) +# TODO: test_OFTeI_radec_auto + if __name__ == "__main__": # test_MargConstNoInput() test_OFTeI() \ No newline at end of file diff --git a/tests/test_mcmc.py b/tests/test_mcmc.py index 867a333d..ffea8a4f 100644 --- a/tests/test_mcmc.py +++ b/tests/test_mcmc.py @@ -68,6 +68,7 @@ def do_mcmc_runs(num_temps=1, num_threads=1, make_corner_plot=False): assert np.all(saved_results.curr_pos == myDriver.sampler.curr_pos) # also check it is consistent with the internal results object in myDriver assert myDriver.sampler.results.post.shape[0] == 1500 + end_pos = myDriver.sampler.results.curr_pos # run it a little more testing that everything gets saved even if prediodic_save_freq is not a multiple of the number of steps output_filename_2 = os.path.join(orbitize.DATADIR, 'test_mcmc_v1.hdf5') @@ -86,8 +87,11 @@ def do_mcmc_runs(num_temps=1, num_threads=1, make_corner_plot=False): # test resuming and restarting from a prevous save new_sampler = sampler.MCMC(myDriver.system, num_temps=num_temps, num_walkers=n_walkers, num_threads=num_threads, prev_result_filename=output_filename) - assert new_sampler.results.post.shape[0] == 1500 + # Check that the curr_pos was extracted and the results are not saved in memory + assert (new_sampler.curr_pos == end_pos).all() + assert new_sampler.results is None new_sampler.run_sampler(500, burn_steps=1) + # Results should be added on to previous ones assert new_sampler.results.post.shape[0] == 2000 assert new_sampler.results.post[0,0] == myDriver.sampler.results.post[0,0] From 9e133d3f10fe41d08383418524a5cac3f6a89095 Mon Sep 17 00:00:00 2001 From: Eshel Dror Date: Tue, 20 Oct 2026 14:57:49 -0700 Subject: [PATCH 15/21] Add warning if using unsaved results --- orbitize/sampler.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/orbitize/sampler.py b/orbitize/sampler.py index 16eb3beb..e2e0a315 100644 --- a/orbitize/sampler.py +++ b/orbitize/sampler.py @@ -1130,7 +1130,10 @@ def run_sampler( raise ValueError("Total_orbits must be greater than num_walkers.") using_unsaved_results = self.results is not None - if not using_unsaved_results: + if using_unsaved_results: + warnings.warn("Calling `run_sampler` multiple times with the same MCMC sampler forces orbitize! to keep a large Results object in memory." \ + " The fit will run much faster if you reinitialize the sampler and use the `prev_result_filename` kwarg to load the previous results.", stacklevel=2) + else: # Initialize results object to save curr_pos self.results = orbitize.results.Results( self.system, From f662895f7fef14bedd87cbf9443d1fe4cac1af92 Mon Sep 17 00:00:00 2001 From: Eshel Dror Date: Tue, 20 Oct 2026 15:56:48 -0700 Subject: [PATCH 16/21] Update ptemcee requirement --- pyproject.toml | 2 +- requirements.txt | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index bee81966..32315a4b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -15,7 +15,7 @@ dependencies = [ "astropy>=4", "scipy", "emcee>=3", - "ptemcee_for_orbitize", + "ptemcee_for_orbitize>=1.1.0", "matplotlib", "corner", "h5py", diff --git a/requirements.txt b/requirements.txt index 877ee7b3..7ee764d9 100644 --- a/requirements.txt +++ b/requirements.txt @@ -2,7 +2,7 @@ astropy>=4 numpy scipy emcee>=3 -ptemcee_for_orbitize +ptemcee_for_orbitize>=1.1.0 matplotlib corner h5py From ccd1823a84965e39f2c8c8baab35cf51042e30b3 Mon Sep 17 00:00:00 2001 From: Eshel Dror Date: Wed, 21 Oct 2026 12:21:20 -0700 Subject: [PATCH 17/21] Fix string data type for quant types --- orbitize/sampler.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/orbitize/sampler.py b/orbitize/sampler.py index e2e0a315..3a6a1e87 100644 --- a/orbitize/sampler.py +++ b/orbitize/sampler.py @@ -49,7 +49,7 @@ def __init__( if marginalizers is None: self.marginalizers = [] - start_quants = set(np.unique(self.system.data_table["quant_type"])) + start_quants = set(np.array(np.unique(self.system.data_table["quant_type"]), dtype=np.dtypes.StringDType)) params = set(self.system.labels) if (self.system.num_secondary_bodies == 1 and len({"radec", "seppa"} & start_quants) > 0 and @@ -107,7 +107,7 @@ def __init__( self.marg_param_idx[label] = marg_param_counter marg_param_counter += 1 self.total_param_num = len(self.system.sys_priors) - quants = set(np.unique(self.system.data_table["quant_type"])) + quants = set(np.array(np.unique(self.system.data_table["quant_type"]), dtype=np.dtypes.StringDType)) marg_quants = set(quant_type for marg in self.marginalizers for quant_type in marg.QUANT_TYPES) self.blob_labels = ["lnlike"] + self.marg_param_labels From 636ec68bb2c4fbbbbc28024a84c8d5a9c6c7b99a Mon Sep 17 00:00:00 2001 From: Eshel Dror Date: Wed, 21 Oct 2026 19:44:47 -0700 Subject: [PATCH 18/21] Calculate correct lnlike except angle wrapping --- orbitize/marginalizer.py | 17 +++++++++++------ orbitize/sampler.py | 4 ++-- tests/test_marg.py | 14 ++++++++++++++ 3 files changed, 27 insertions(+), 8 deletions(-) diff --git a/orbitize/marginalizer.py b/orbitize/marginalizer.py index c7bfedd9..adb38cc3 100644 --- a/orbitize/marginalizer.py +++ b/orbitize/marginalizer.py @@ -3,6 +3,7 @@ from astropy import units as u import orbitize.basis as basis import orbitize.kepler as kepler +import scipy.stats """ This module defines marginalizations to draw @@ -198,7 +199,10 @@ def solve_linear_eq(Y, A, C, C_inv=None): cov_matrix ) # NOTE: the 2pi factor in the APW paper is constant; doesn't affect overall lnlike - return sma_map, pan_map, logabsdet_covmatrix, sma_sample, pan_sample + dist = scipy.stats.multivariate_normal([sma_map, pan_map], cov_matrix) + chi2_sample_norm = dist.logpdf([sma_sample, pan_sample]) + + return sma_map, pan_map, logabsdet_covmatrix, sma_sample, pan_sample, chi2_sample_norm @staticmethod def lnlike(y, X, beta_hat, W, W_inv=None, add_error_norm=False): @@ -224,7 +228,7 @@ def lnlike(y, X, beta_hat, W, W_inv=None, add_error_norm=False): # NOTE: this term can be ignored unless we're fitting for jitter if add_error_norm: - error_norm_term = -0.5 * np.log(np.det(W)) + error_norm_term = -0.5 * np.linalg.det(W) else: error_norm_term = 0 @@ -252,11 +256,11 @@ def apply(self, param_arr, jitter_arr, calculated_values): W_inv = self.W_inv try: - smas_map, pans_map_degrees, logdet_covmatrix, sma_sample, pan_sample_deg = ( + smas_map, pans_map_degrees, logdet_covmatrix, sma_sample, pan_sample_deg, chi2samples = ( self.calc_cov_matrix_det(self.seps, self.pas, f_samples, g_samples, W, W_inv) ) except np.linalg.LinAlgError: - return -np.inf, None, None + return -np.inf, -np.inf, None, None # set up design matrix X = np.zeros((2 * self.n_obs, 2)) @@ -267,7 +271,7 @@ def apply(self, param_arr, jitter_arr, calculated_values): beta_hat = np.hstack((smas_map, pans_map_degrees)) # compute marginalized likelihood (e.g. Price-Whelan 2017 eq 11) - chi2lnlike = self.lnlike(y, X, beta_hat, W, W_inv) + chi2lnlike = self.lnlike(y, X, beta_hat, W, W_inv, add_error_norm=True) logprob = ( chi2lnlike + 0.5 * logdet_covmatrix @@ -275,6 +279,8 @@ def apply(self, param_arr, jitter_arr, calculated_values): pan_sample = np.deg2rad(pan_sample_deg) % (2 * np.pi) + loglike = logprob + chi2samples + # apply sma/pan priors if self.sma_prior is not None: logprob += np.sum(self.sma_prior.compute_lnprob(np.array([sma_sample]))) @@ -291,7 +297,6 @@ def apply(self, param_arr, jitter_arr, calculated_values): np.array([mtot_sample]) )) # apply the user-defined prior - loglike = logprob # TODO samples = np.array([sma_sample, pan_sample]) if not np.isfinite(logprob): diff --git a/orbitize/sampler.py b/orbitize/sampler.py index 3a6a1e87..e47de08f 100644 --- a/orbitize/sampler.py +++ b/orbitize/sampler.py @@ -1024,8 +1024,8 @@ def _update_chains_from_sampler(self, sampler, num_steps=None): # Ntemp x Nwalkers x Nsteps x Nblobs blobs = sampler.blobs[0, :, :num_steps, :].reshape(-1, len(self.blob_labels)) assert len(blobs) == len(self.lnlikes) - no_diff = ((self.lnlikes == blobs[:, 0]) | (np.abs(self.lnlikes - blobs[:, 0]) < 0.001)) - assert (no_diff).all(), (self.lnlikes[~no_diff][:10], blobs[:, 0][~no_diff][:10], np.shape(np.where(~no_diff)), np.shape(np.where(no_diff))) + # no_diff = ((self.lnlikes == blobs[:, 0]) | (np.abs(self.lnlikes - blobs[:, 0]) < 0.001)) + # assert (no_diff).all(), (self.lnlikes[~no_diff][:10], blobs[:, 0][~no_diff][:10], np.shape(np.where(~no_diff)), np.shape(np.where(no_diff))) # assert (no_diff).all(), (sampler.loglikelihood[:, :, :num_steps], sampler.blobs[:, :, :num_steps, 0], np.shape(np.where(~no_diff)), np.shape(np.where(no_diff))) marginalized_param_blobs = blobs[:, 1:] assert not np.isnan(marginalized_param_blobs).any(), sampler.blobs[:5, :10] diff --git a/tests/test_marg.py b/tests/test_marg.py index 3de4fc73..a774a658 100644 --- a/tests/test_marg.py +++ b/tests/test_marg.py @@ -159,6 +159,20 @@ def test_OFTeI(): mySampler.run_sampler(4000, 10, output_filename=output_file, periodic_save_freq=2) assert not mySampler.force_valid_marg myResults = mySampler.results + + myNoMargSampler = sampler.MCMC( + mySystem, + marginalizers=[] + ) + + lnlike_true = [] + for params, lnlike in zip(mySampler.results.post, mySampler.results.lnlike): + lnlike_true_i = super(sampler.MCMC, myNoMargSampler)._logl(params)[0] + # assert lnlike_true_i == pytest.approx(lnlike, abs=.01), (lnlike_true_i, lnlike) + lnlike_true.append(lnlike_true_i) + lnlike_true = np.array(lnlike_true) + assert (np.abs(lnlike_true - mySampler.results.lnlike) < .01).all(), np.where(np.abs(lnlike_true - mySampler.results.lnlike) < .01) + print(myResults.post.shape) def test_OFTeI_RV(): From 4f41a9e7f45f0c9b24df07307f7c2e5f7b4916de Mon Sep 17 00:00:00 2001 From: Eshel Dror Date: Wed, 21 Oct 2026 19:51:58 -0700 Subject: [PATCH 19/21] Wrap pan residual for sample --- orbitize/marginalizer.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/orbitize/marginalizer.py b/orbitize/marginalizer.py index adb38cc3..fc82bef5 100644 --- a/orbitize/marginalizer.py +++ b/orbitize/marginalizer.py @@ -195,6 +195,9 @@ def solve_linear_eq(Y, A, C, C_inv=None): [sma_map, pan_map], cov_matrix ) + pan_res = (pan_sample - pan_map + 180.0) % 360.0 - 180.0 + pan_sample = pan_map + pan_res + _, logabsdet_covmatrix = np.linalg.slogdet( cov_matrix ) # NOTE: the 2pi factor in the APW paper is constant; doesn't affect overall lnlike From d096baff4524aae36ab9b14bb41a6b31d41cfa42 Mon Sep 17 00:00:00 2001 From: Eshel Dror Date: Tue, 27 Oct 2026 17:59:13 -0700 Subject: [PATCH 20/21] Revert "Wrap pan residual for sample" This reverts commit 4f41a9e7f45f0c9b24df07307f7c2e5f7b4916de. --- orbitize/marginalizer.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/orbitize/marginalizer.py b/orbitize/marginalizer.py index fc82bef5..adb38cc3 100644 --- a/orbitize/marginalizer.py +++ b/orbitize/marginalizer.py @@ -195,9 +195,6 @@ def solve_linear_eq(Y, A, C, C_inv=None): [sma_map, pan_map], cov_matrix ) - pan_res = (pan_sample - pan_map + 180.0) % 360.0 - 180.0 - pan_sample = pan_map + pan_res - _, logabsdet_covmatrix = np.linalg.slogdet( cov_matrix ) # NOTE: the 2pi factor in the APW paper is constant; doesn't affect overall lnlike From f30bfbc81914209709db5b0cb6b885863d5c59da Mon Sep 17 00:00:00 2001 From: Eshel Dror Date: Tue, 27 Oct 2026 17:59:24 -0700 Subject: [PATCH 21/21] Revert "Calculate correct lnlike except angle wrapping" This reverts commit 636ec68bb2c4fbbbbc28024a84c8d5a9c6c7b99a. --- orbitize/marginalizer.py | 17 ++++++----------- orbitize/sampler.py | 4 ++-- tests/test_marg.py | 14 -------------- 3 files changed, 8 insertions(+), 27 deletions(-) diff --git a/orbitize/marginalizer.py b/orbitize/marginalizer.py index adb38cc3..c7bfedd9 100644 --- a/orbitize/marginalizer.py +++ b/orbitize/marginalizer.py @@ -3,7 +3,6 @@ from astropy import units as u import orbitize.basis as basis import orbitize.kepler as kepler -import scipy.stats """ This module defines marginalizations to draw @@ -199,10 +198,7 @@ def solve_linear_eq(Y, A, C, C_inv=None): cov_matrix ) # NOTE: the 2pi factor in the APW paper is constant; doesn't affect overall lnlike - dist = scipy.stats.multivariate_normal([sma_map, pan_map], cov_matrix) - chi2_sample_norm = dist.logpdf([sma_sample, pan_sample]) - - return sma_map, pan_map, logabsdet_covmatrix, sma_sample, pan_sample, chi2_sample_norm + return sma_map, pan_map, logabsdet_covmatrix, sma_sample, pan_sample @staticmethod def lnlike(y, X, beta_hat, W, W_inv=None, add_error_norm=False): @@ -228,7 +224,7 @@ def lnlike(y, X, beta_hat, W, W_inv=None, add_error_norm=False): # NOTE: this term can be ignored unless we're fitting for jitter if add_error_norm: - error_norm_term = -0.5 * np.linalg.det(W) + error_norm_term = -0.5 * np.log(np.det(W)) else: error_norm_term = 0 @@ -256,11 +252,11 @@ def apply(self, param_arr, jitter_arr, calculated_values): W_inv = self.W_inv try: - smas_map, pans_map_degrees, logdet_covmatrix, sma_sample, pan_sample_deg, chi2samples = ( + smas_map, pans_map_degrees, logdet_covmatrix, sma_sample, pan_sample_deg = ( self.calc_cov_matrix_det(self.seps, self.pas, f_samples, g_samples, W, W_inv) ) except np.linalg.LinAlgError: - return -np.inf, -np.inf, None, None + return -np.inf, None, None # set up design matrix X = np.zeros((2 * self.n_obs, 2)) @@ -271,7 +267,7 @@ def apply(self, param_arr, jitter_arr, calculated_values): beta_hat = np.hstack((smas_map, pans_map_degrees)) # compute marginalized likelihood (e.g. Price-Whelan 2017 eq 11) - chi2lnlike = self.lnlike(y, X, beta_hat, W, W_inv, add_error_norm=True) + chi2lnlike = self.lnlike(y, X, beta_hat, W, W_inv) logprob = ( chi2lnlike + 0.5 * logdet_covmatrix @@ -279,8 +275,6 @@ def apply(self, param_arr, jitter_arr, calculated_values): pan_sample = np.deg2rad(pan_sample_deg) % (2 * np.pi) - loglike = logprob + chi2samples - # apply sma/pan priors if self.sma_prior is not None: logprob += np.sum(self.sma_prior.compute_lnprob(np.array([sma_sample]))) @@ -297,6 +291,7 @@ def apply(self, param_arr, jitter_arr, calculated_values): np.array([mtot_sample]) )) # apply the user-defined prior + loglike = logprob # TODO samples = np.array([sma_sample, pan_sample]) if not np.isfinite(logprob): diff --git a/orbitize/sampler.py b/orbitize/sampler.py index e47de08f..3a6a1e87 100644 --- a/orbitize/sampler.py +++ b/orbitize/sampler.py @@ -1024,8 +1024,8 @@ def _update_chains_from_sampler(self, sampler, num_steps=None): # Ntemp x Nwalkers x Nsteps x Nblobs blobs = sampler.blobs[0, :, :num_steps, :].reshape(-1, len(self.blob_labels)) assert len(blobs) == len(self.lnlikes) - # no_diff = ((self.lnlikes == blobs[:, 0]) | (np.abs(self.lnlikes - blobs[:, 0]) < 0.001)) - # assert (no_diff).all(), (self.lnlikes[~no_diff][:10], blobs[:, 0][~no_diff][:10], np.shape(np.where(~no_diff)), np.shape(np.where(no_diff))) + no_diff = ((self.lnlikes == blobs[:, 0]) | (np.abs(self.lnlikes - blobs[:, 0]) < 0.001)) + assert (no_diff).all(), (self.lnlikes[~no_diff][:10], blobs[:, 0][~no_diff][:10], np.shape(np.where(~no_diff)), np.shape(np.where(no_diff))) # assert (no_diff).all(), (sampler.loglikelihood[:, :, :num_steps], sampler.blobs[:, :, :num_steps, 0], np.shape(np.where(~no_diff)), np.shape(np.where(no_diff))) marginalized_param_blobs = blobs[:, 1:] assert not np.isnan(marginalized_param_blobs).any(), sampler.blobs[:5, :10] diff --git a/tests/test_marg.py b/tests/test_marg.py index a774a658..3de4fc73 100644 --- a/tests/test_marg.py +++ b/tests/test_marg.py @@ -159,20 +159,6 @@ def test_OFTeI(): mySampler.run_sampler(4000, 10, output_filename=output_file, periodic_save_freq=2) assert not mySampler.force_valid_marg myResults = mySampler.results - - myNoMargSampler = sampler.MCMC( - mySystem, - marginalizers=[] - ) - - lnlike_true = [] - for params, lnlike in zip(mySampler.results.post, mySampler.results.lnlike): - lnlike_true_i = super(sampler.MCMC, myNoMargSampler)._logl(params)[0] - # assert lnlike_true_i == pytest.approx(lnlike, abs=.01), (lnlike_true_i, lnlike) - lnlike_true.append(lnlike_true_i) - lnlike_true = np.array(lnlike_true) - assert (np.abs(lnlike_true - mySampler.results.lnlike) < .01).all(), np.where(np.abs(lnlike_true - mySampler.results.lnlike) < .01) - print(myResults.post.shape) def test_OFTeI_RV():