From 00f163a7046401b3c97f52088a1f50a8200c1815 Mon Sep 17 00:00:00 2001 From: Melody Tang Date: Tue, 15 Sep 2026 15:40:44 -0400 Subject: [PATCH] Report Tc and Tp in the observed BJD_TDB frame Adds tc_bjd/tp_bjd converting the reported conjunction/periastron times back to the observed frame (the Mahajan 2024 frame bug). Verified on GJ 1214: tc_bjd lands 0.6sigma from the paper. Model path unchanged. --- src/exozippy/components/orbit/defaults.yaml | 26 ++++ src/exozippy/components/orbit/orbit.py | 154 +++++++++++++++++++- src/exozippy/components/orbit/physics.py | 22 +++ tests/test_rm_ltt.py | 19 ++- tests/test_transit_ltt.py | 43 ++++-- 5 files changed, 246 insertions(+), 18 deletions(-) diff --git a/src/exozippy/components/orbit/defaults.yaml b/src/exozippy/components/orbit/defaults.yaml index ef6c88ff..3c0567f0 100644 --- a/src/exozippy/components/orbit/defaults.yaml +++ b/src/exozippy/components/orbit/defaults.yaml @@ -111,6 +111,32 @@ orbit: func_name: "calc_K" deps: ["m_companion", "m_total", "ecc", "a", "sini", "period"] + # Observed-frame (BJD_TDB) convert-back of tc/tp, when light-travel-time + # is active for this orbit (see Orbit._ltt_reporting_mask). tc/tp above + # are the TARGET-frame values the Kepler solve is evaluated against + # (components/ltt.py) -- correct for the model, but not the BJD_TDB + # timestamp an observer would see. `_ltt_delta` is a context node (the + # masked ltt.retarded_time delay, evaluated at tc), injected by + # Orbit.add_parameter -- never a manifest parameter. + tc_bjd: + unit: "d" + internal_unit: "d" + latex: "T_{C,\\rm obs}" + description: "Observed Conjunction Time" + expressions: + default: + func_name: "calc_bjd_shift" + deps: [ "tc", "_ltt_delta" ] + tp_bjd: + unit: "d" + internal_unit: "d" + latex: "T_{P,\\rm obs}" + description: "Observed Periastron Time" + expressions: + default: + func_name: "calc_bjd_shift" + deps: [ "tp", "_ltt_delta" ] + # --- Conditionally Sampled or Derived (Geometry Swap) --- # Sampled where the orbit uses the sqrt(e)cos/sin(omega) pair, and REPORTED # (derived from ecc/omega, manifest role 3) where it uses V_c/V_e -- so a diff --git a/src/exozippy/components/orbit/orbit.py b/src/exozippy/components/orbit/orbit.py index 61f551db..308123c6 100644 --- a/src/exozippy/components/orbit/orbit.py +++ b/src/exozippy/components/orbit/orbit.py @@ -15,6 +15,8 @@ from exozippy.outputs.prose import get_collector, join_names from exozippy.potentials import soft_lower_bound, soft_upper_bound +from .. import ltt + # this import is required even though it's not used explicitly # it registers all the mathematical relations from . import physics @@ -906,6 +908,27 @@ def register_parameters(self, system): } ) + # Observed-frame (BJD_TDB) convert-back of tc/tp -- see + # defaults.yaml's tc_bjd/tp_bjd and _ltt_delta_context. Needs + # a/m_primary/m_companion/m_total, so it is gated the same way + # they are; within that, `_ltt_reporting_mask` decides PER ORBIT + # whether any consumer actually retards this orbit's geometry + # (an orbit with light_travel_time off everywhere must report + # tc_bjd == tc exactly, not a shifted value -- see the mask's + # own docstring). Declared only where the mask is nonzero + # anywhere, mirroring RM's `if rm_enabled(system):` below: no + # manifest entry, no cost, for a system that never uses LTT. + self._ltt_report_mask = self._ltt_reporting_mask(system) + if self._ltt_report_mask.any(): + self.manifest["tc_bjd"] = { + "expr_key": "default", + "force_node": True, + } + self.manifest["tp_bjd"] = { + "expr_key": "default", + "force_node": True, + } + # Rossiter-McLaughlin: declare the spin-orbit params only when some # rvinstrument enables `rm:`. Samples the decorrelated # sqrt(vsini)cos/sin(lambda) pair and derives vsini/lam from them @@ -1418,6 +1441,67 @@ def _validate_bodies(self, system): ) return True + def _ltt_reporting_mask(self, system): + """Per-orbit 0.0/1.0 float array: does some consumer's model + actually retard THIS orbit's geometry, so `tc`/`tp` need the + +delay convert-back (`tc_bjd`/`tp_bjd`) to read as the observed + (BJD_TDB) frame rather than the target frame `ltt.py` evaluates + the Kepler solve in? + + Read from raw config, like `rm.rm_orbits_in_system` -- stage 3 + makes no promise that a sibling component has run its OWN stage 3 + yet, only that every component's `build_maps` (stage 2) has, which + is what `planet.orbit_map` needs. + + A transit file has no per-planet selection: build_likelihood + models every planet (hence every orbit with one) in every active + file's likelihood (see transit.py), so "this orbit is retarded by + transit" reduces to "this orbit has >=1 planet" AND "some transit + file has light_travel_time on" -- the per-file default there is + True, matching `Transit._light_travel_time_active`. RM is + per-orbit already, through its own `rm:` key. + + Mixing `light_travel_time` across files for the SAME orbit is a + pre-existing ambiguity in the model itself (that orbit's own tc/tp + posterior is then pulled toward the target frame by only some of + its data) -- not one this mask can resolve, so it warns once and + treats the orbit as active (closer to correct than leaving it at + the target-frame value outright). + """ + mask = np.zeros(self.n_elements, dtype=float) + cfg = getattr(system, "config", None) or {} + + transit_cfg = cfg.get("transit", []) or [] + transit_flags = [ + bool(c.get("light_travel_time", True)) for c in transit_cfg + ] + if any(transit_flags): + orbit_map = getattr( + getattr(system, "planet", None), "orbit_map", None + ) + if orbit_map is not None: + for o_idx in np.asarray(orbit_map, dtype=int): + if 0 <= o_idx < self.n_elements: + mask[o_idx] = 1.0 + if len(set(transit_flags)) > 1: + logger.warning( + "orbit: transit files disagree on light_travel_time; " + "tc_bjd/tp_bjd treat every orbit touched by an active " + "file as fully retarded, an approximation where they " + "mix on the same orbit's data." + ) + + rv_cfg = cfg.get("rvinstrument", []) or [] + name_to_idx = {n: i for i, n in enumerate(self.names)} + for entry in rv_cfg: + o_name = entry.get("rm") + if o_name and bool(entry.get("light_travel_time", True)): + o_idx = name_to_idx.get(o_name) + if o_idx is not None: + mask[o_idx] = 1.0 + + return mask + _GROUP_MASS_SIDE = {"m_primary": "primary", "m_companion": "companion"} # The chord expressions' deps that are NOT orbit parameters: the @@ -1426,8 +1510,11 @@ def _validate_bodies(self, system): # graph.py from looking for an `orbit.p` (the group masses avoid this by # naming `planet.mass`, a real parameter of a real component; there is no # such parameter for `chord_sign` at all, and `p`/`ar` are per PLANET, so - # the orbit could not consume them elementwise anyway). - context_dep_names = frozenset({"p", "ar", "chord_sign"}) + # the orbit could not consume them elementwise anyway). `_ltt_delta` is + # the same idiom for a different reason: it is the (masked) + # ltt.retarded_time delay `tc_bjd`/`tp_bjd` add back, not a manifest + # parameter of any component (see _ltt_delta_context). + context_dep_names = frozenset({"p", "ar", "chord_sign", "_ltt_delta"}) # ...and all three are built per ORBIT, so Component._element_expression # may slice them to a per-element mask. @@ -1549,8 +1636,71 @@ def add_parameter(self, model, param_name, system, context_nodes=None): context_nodes = dict(context_nodes or {}) for dep, node in self._chord_context(model, system).items(): context_nodes.setdefault(dep, node) + if param_name in self._LTT_REPORT_PARAMS and not context_nodes: + context_nodes = dict(context_nodes or {}) + context_nodes["_ltt_delta"] = self._ltt_delta_context( + model, system + ) return super().add_parameter(model, param_name, system, context_nodes) + # tc_bjd/tp_bjd (defaults.yaml) both consume the one `_ltt_delta` + # context node -- see _ltt_delta_context and calc_bjd_shift. + _LTT_REPORT_PARAMS = ("tc_bjd", "tp_bjd") + + def _ltt_delta_context(self, model, system): + """The `_ltt_delta` context node `tc_bjd`/`tp_bjd` consume: the + light-travel delay evaluated AT `tc`, reusing `ltt.retarded_time`'s + own delay output (never a hardcoded number), masked per orbit by + `_ltt_reporting_mask` so an orbit with light_travel_time off + everywhere gets exactly zero -- `tc_bjd == tc`, `tp_bjd == tp`. + + Factor is the OCCULTATION seam, `(m_primary - m_companion) / + m_total` -- matching transit.py's own `ltt_factor` (`tc`/`tp` are + defined by the transit/occultation geometry, not by either body's + own emission; see components/ltt.py's `factor` docs). Same lazy + same-component build as the group masses above, so this works + regardless of which of tc_bjd/tp_bjd the build order reaches first. + """ + for dep in ( + "tc", + "tp", + "n", + "ecc", + "sinw", + "cosw", + "inc", + "a", + "m_primary", + "m_companion", + "m_total", + ): + if not Component._parameter_is_current(self, dep, model): + self.add_parameter(model, dep, system) + + sin_i = pt.sin(self.inc.value) + factor = ( + self.m_primary.value - self.m_companion.value + ) / self.m_total.value + _, delay = ltt.retarded_time( + self.tc.value, + self.tp.value, + self.n.value, + self.ecc.value, + self.sinw.value, + self.cosw.value, + sin_i, + self.a.value, + factor=factor, + z0=0.0, + circular=self._all_circular(), + ) + mask = pt.as_tensor_variable( + getattr( + self, "_ltt_report_mask", np.zeros(self.n_elements) + ).astype("float64") + ) + return delay * mask + def build_likelihood(self, model, system): self._add_eccentricity_bound(system) self._add_vcve_terms(system) diff --git a/src/exozippy/components/orbit/physics.py b/src/exozippy/components/orbit/physics.py index 595f053c..fb3e928d 100644 --- a/src/exozippy/components/orbit/physics.py +++ b/src/exozippy/components/orbit/physics.py @@ -601,6 +601,28 @@ def calc_ts(ecc, omega, tc, period): return ts_from_ecc_omega(ecc, omega, tc, period, xp=pt) +@register_physics +def calc_bjd_shift(value, delta): + """`value + delta` -- the observed-frame (BJD_TDB) convert-back for a + target-frame timing quantity, shared by `tc_bjd` and `tp_bjd`. + + `delta` is the light-travel delay from `ltt.retarded_time`, evaluated at + `tc` and masked per orbit (see `Orbit._ltt_delta_context`) -- the SAME + delay for tc_bjd and tp_bjd, since `calc_tp`/`calc_tp_from_ecc` are + additive in `tc` (`tp = tc - M0/n`, with `M0` independent of `tc`), so + `tp + delta(tc) == (tc + delta(tc)) - M0/n` exactly: `tp_bjd` needs no + physics of its own beyond reusing tc_bjd's delay. `_ltt_delta_context` + builds that delay fresh per consumer (the same `_chord_context` idiom + `p`/`ar`/`chord_sign` already use for cosi/chord) rather than caching it + -- the two builds are structurally identical, so pytensor's merge + optimizer collapses them to one Kepler solve in the compiled graph + (confirmed: one `Kepler` Apply node for [tc_bjd, tp_bjd] together). All + the physics lives in the delay itself; this function is deliberately + trivial. + """ + return value + delta + + def mean_anomaly_at_true_anomaly(ecc, true_anomaly, xp=pt): """Mean anomaly at a true anomaly, in radians. diff --git a/tests/test_rm_ltt.py b/tests/test_rm_ltt.py index 79371103..3e831471 100644 --- a/tests/test_rm_ltt.py +++ b/tests/test_rm_ltt.py @@ -130,6 +130,15 @@ def test_wired_rm_ltt_delay_matches_a_over_c_through_real_accessors(tmp_path): m_total nodes the wiring itself used, read back independently) to 1e-6 relative, and is close to the known ~499 s/AU light time in absolute terms. + + `ltt.retarded_time` is patched on the shared `ltt` module, so + `orbit.tc_bjd`/`tp_bjd` (the observed/BJD_TDB-frame conjunction/ + periastron report this orbit's mass params make available) call it + too, at `tc`. Filtered out below by object identity (`t` is a + symbolic tensor at graph-build time, not a concrete array, so + comparing shapes is unreliable): orbit.py always passes + `system.orbit.tc.value` itself, so excluding calls whose `t` IS that + exact node leaves only the RM wiring's own call(s). """ rv_file = _write_two_row_rv(tmp_path / "rv.dat") real_retarded_time = rm.ltt.retarded_time @@ -137,7 +146,7 @@ def test_wired_rm_ltt_delay_matches_a_over_c_through_real_accessors(tmp_path): def _ltt_spy(*args, **kwargs): result = real_retarded_time(*args, **kwargs) - delay_calls.append(result[1]) + delay_calls.append((args[0], result[1])) # (t, delay) return result with mock.patch.object(rm.ltt, "retarded_time", side_effect=_ltt_spy): @@ -147,12 +156,16 @@ def _ltt_spy(*args, **kwargs): system.prepare() model = system.build_model() - assert len(delay_calls) >= 1 # build_likelihood (+ compile_plotters) + tc_node = ( + system.orbit.tc.value if "tc_bjd" in system.orbit.manifest else None + ) + rm_delay_calls = [delay for t, delay in delay_calls if t is not tc_node] + assert len(rm_delay_calls) >= 1 # build_likelihood (+ compile_plotters) with model: point = system.get_internal_point(model, system.get_raw_start(model)) - delay = _eval_at_point(delay_calls[0], model, point) + delay = _eval_at_point(rm_delay_calls[0], model, point) delay_primary = float(delay[0]) delay_secondary = float(delay[1]) diff --git a/tests/test_transit_ltt.py b/tests/test_transit_ltt.py index b8a1b7da..0d5a6449 100644 --- a/tests/test_transit_ltt.py +++ b/tests/test_transit_ltt.py @@ -131,10 +131,14 @@ def test_wired_ltt_delay_matches_a_over_c_through_real_accessors(tmp_path): the units/accessor contract end to end -- and separately is close to the known ~499 s/AU light time in absolute terms, guarding against a shared-formula bug that an internal-consistency check alone would miss. - Also confirms retarded_time is invoked exactly twice per build_model(): - once from build_likelihood's group loop (one oversample group here, - ninterp=1) and once from compile_plotters -- both paths are wired, live, - not a dead branch. + Also confirms retarded_time is invoked exactly twice per build_model() + FOR TRANSIT'S OWN WIRING: once from build_likelihood's group loop (one + oversample group here, ninterp=1) and once from compile_plotters -- + both paths are wired, live, not a dead branch. (`orbit.tc_bjd`/`tp_bjd` + -- the observed/BJD_TDB-frame conjunction/periastron report this + orbit's mass params make available -- call retarded_time too, at `tc` + (1-D, per-orbit, not per-observation); filtered out below since this + test is specifically about transit.py's own wiring.) """ lc = _write_two_row_lc(tmp_path / "lc.dat") real_retarded_time = ltt.retarded_time @@ -148,17 +152,18 @@ def _spy(*args, **kwargs): with mock.patch( "exozippy.components.transit.transit.ltt.retarded_time", side_effect=_spy, - ) as spy: + ): system = System( _ltt_wiring_config(lc), user_params=_ltt_wiring_params() ) system.prepare() model = system.build_model() - # One from build_likelihood's group loop (3-D t_grid, has the - # sub-exposure axis), one from compile_plotters (2-D, no - # sub-exposure axis -- smearing is applied outside by - # _smeared_full_lc instead). - assert spy.call_count == 2 + # transit.py's own wiring: one from build_likelihood's group loop + # (3-D t_grid, has the sub-exposure axis), one from compile_plotters + # (2-D, no sub-exposure axis). orbit.tc_bjd/tp_bjd's calls are 1-D + # (a time grid per orbit, not per observation) and excluded here. + transit_calls = [t_grid for t_grid, _ in calls if t_grid.ndim >= 2] + assert len(transit_calls) == 2 with model: point = system.get_internal_point(model, system.get_raw_start(model)) @@ -345,10 +350,18 @@ def test_mixed_group_ltt_gradient_is_finite(tmp_path): # circular orbit with no Kepler op to count. params = _ltt_wiring_params(eccentric=True) + real_retarded_time = ltt.retarded_time + calls = [] + + def _spy(*args, **kwargs): + result = real_retarded_time(*args, **kwargs) + calls.append(args[0]) # t_grid + return result + with mock.patch( "exozippy.components.transit.transit.ltt.retarded_time", - side_effect=ltt.retarded_time, - ) as spy: + side_effect=_spy, + ): system = System(config, user_params=params) system.prepare() model = system.build_model() @@ -357,7 +370,11 @@ def test_mixed_group_ltt_gradient_is_finite(tmp_path): # group actually took the "any() but not all()" branch and called # ltt.retarded_time at all (a bug that skipped the correction # entirely for a mixed group would also show 0 calls here). - assert spy.call_count == 2 + # orbit.tc_bjd/tp_bjd's calls are 1-D (per-orbit, not per- + # observation) and excluded, same reasoning as + # test_wired_ltt_delay_matches_a_over_c_through_real_accessors. + transit_calls = [t_grid for t_grid in calls if t_grid.ndim >= 2] + assert len(transit_calls) == 2 assert _count_kepler_solves(system.transit._model_flux_node) == 2