From 00c96900414b82806c7be75fdc7759f55dc8aaab Mon Sep 17 00:00:00 2001
From: ThomasBaycroft <thombaycroft@gmail.com>
Date: Wed, 31 Jul 2024 13:39:14 +0100
Subject: [PATCH] ETVmodel ref_time as a parameter

---
 src/kima/ETVmodel.cpp | 33 ++++++++++++++++++++++++++-------
 src/kima/ETVmodel.h   |  3 +++
 src/kima/kepler.cpp   |  2 +-
 3 files changed, 30 insertions(+), 8 deletions(-)

diff --git a/src/kima/ETVmodel.cpp b/src/kima/ETVmodel.cpp
index e318ab6..6348634 100644
--- a/src/kima/ETVmodel.cpp
+++ b/src/kima/ETVmodel.cpp
@@ -36,6 +36,12 @@ void ETVmodel::setPriors()  // BUG: should be done by only one thread!
         ephem2_prior = make_prior<Gaussian>(0.0,pow(10,-10.));
     if (ephemeris >= 3 && !ephem3_prior)
         ephem3_prior = make_prior<Gaussian>(0.0,pow(10,-12.));
+
+    if (!ref_time_prior)
+    {
+        ref_time_prior = make_prior<Gaussian>(data.M0_epoch,0.1);
+        printf("# No prior on reference time specified, taken as Gaussian with range of 0.1 days\n");
+    }
         
     if (known_object) { // KO mode!
         for (size_t i = 0; i < n_known_object; i++)
@@ -64,6 +70,7 @@ void ETVmodel::from_prior(RNG& rng)
     planets.consolidate_diff();
     
     jitter = Jprior->generate(rng);
+    M0_epoch = ref_time_prior->generate(rng);
     
     
     if (ephemeris >= 1) ephem1 = ephem1_prior->generate(rng);
@@ -116,7 +123,7 @@ void ETVmodel::calculate_mu()
     // Zero the signal
     if(!update) // not updating, means recalculate everything
     {
-        mu.assign(mu.size(), data.M0_epoch);
+        mu.assign(mu.size(), M0_epoch);
         
         staleness = 0;
         
@@ -234,7 +241,7 @@ double ETVmodel::perturb(RNG& rng)
         //subtract ephemeris
         for(size_t i=0; i<mu.size(); i++)
         {
-            mu[i] += -data.M0_epoch - ephem1*data.epochs[i]- 0.5*ephem2*ephem1*pow(data.epochs[i],2.0) - ephem3*pow(ephem1,2.0)*pow(data.epochs[i],3.0)/6.0;
+            mu[i] += -M0_epoch - ephem1*data.epochs[i]- 0.5*ephem2*ephem1*pow(data.epochs[i],2.0) - ephem3*pow(ephem1,2.0)*pow(data.epochs[i],3.0)/6.0;
         }
         
         // propose new ephemeris
@@ -242,10 +249,12 @@ double ETVmodel::perturb(RNG& rng)
         if (ephemeris >= 2) ephem2_prior->perturb(ephem2, rng);
         if (ephemeris == 3) ephem3_prior->perturb(ephem3, rng);
 
+        ref_time_prior->perturb(M0_epoch, rng);
+
         //add ephemeris back in
         for(size_t i=0; i<mu.size(); i++)
         {
-            mu[i] += data.M0_epoch + ephem1*data.epochs[i]+ 0.5*ephem2*ephem1*pow(data.epochs[i],2.0) + ephem3*pow(ephem1,2.0)*pow(data.epochs[i],3.0)/6.0;
+            mu[i] += M0_epoch + ephem1*data.epochs[i]+ 0.5*ephem2*ephem1*pow(data.epochs[i],2.0) + ephem3*pow(ephem1,2.0)*pow(data.epochs[i],3.0)/6.0;
         }
     }
 
@@ -330,7 +339,11 @@ void ETVmodel::print(std::ostream& out) const
     out<<jitter<<'\t';
 
 
+
     out.precision(24);
+
+    out << M0_epoch << '\t';
+
     if (ephemeris >= 1) out << ephem1 << '\t';
     if (ephemeris >= 2) out << ephem2 << '\t';
     if (ephemeris == 3) out << ephem3 << '\t';
@@ -361,7 +374,9 @@ string ETVmodel::description() const
     string desc;
     string sep = "   ";
 
-    desc += "jitter   ";
+    desc += "jitter" + sep;
+
+    desc += "ref_time" + sep;
 
     if (ephemeris >= 1) desc += "ephem1" + sep;
     if (ephemeris >= 2) desc += "ephem2" + sep;
@@ -432,14 +447,13 @@ void ETVmodel::save_setup() {
     fout << "skip: " << data._skip << endl;
 
 
-    fout.precision(15);
-    fout << "M0_epoch: " << data.M0_epoch << endl;
-    fout.precision(6);
+    fout.precision(12);
 
     fout << endl;
 
     fout << "[priors.general]" << endl;
     fout << "Jprior: " << *Jprior << endl;
+    fout << "ref_time_prior: " << *ref_time_prior << endl;
 
     if (ephemeris >= 1) fout << "ephem1_prior: " << *ephem1_prior << endl;
     if (ephemeris >= 2) fout << "ephem2_prior: " << *ephem2_prior << endl;
@@ -545,6 +559,11 @@ NB_MODULE(ETVmodel, m) {
             [](ETVmodel &m, distribution &d) { m.Jprior = d; },
             "Prior for the extra white noise (jitter)")
 
+        .def_prop_rw("ref_time_prior",
+            [](ETVmodel &m) { return m.ref_time_prior; },
+            [](ETVmodel &m, distribution &d) { m.ref_time_prior = d; },
+            "Prior for the reference time (where epoch = 0)")
+
         .def_prop_rw("ephem1_prior",
             [](ETVmodel &m) { return m.ephem1_prior; },
             [](ETVmodel &m, distribution &d) { m.ephem1_prior = d; },
diff --git a/src/kima/ETVmodel.h b/src/kima/ETVmodel.h
index 9d4226b..d38efc0 100644
--- a/src/kima/ETVmodel.h
+++ b/src/kima/ETVmodel.h
@@ -45,6 +45,7 @@ class KIMA_API ETVmodel
             DNest4::RJObject<ETVConditionalPrior>(5, npmax, fix, ETVConditionalPrior());
 
         double ephem1, ephem2=0.0, ephem3=0.0;
+        double M0_epoch;
         double nu;
         double jitter;
 
@@ -81,6 +82,8 @@ class KIMA_API ETVmodel
         distribution Jprior;
         /// Prior for the stellar jitter (common to all instruments)
         distribution stellar_jitter_prior;
+        ///Prior for the reference time
+        distribution ref_time_prior;
         /// Prior for the slope
         distribution ephem1_prior;
         /// Prior for the quadratic coefficient of the trend
diff --git a/src/kima/kepler.cpp b/src/kima/kepler.cpp
index 53fa205..743bbe4 100644
--- a/src/kima/kepler.cpp
+++ b/src/kima/kepler.cpp
@@ -1182,7 +1182,7 @@ NB_MODULE(kepler, m) {
     m.def("contour_solver", [](double M, double ecc) { return contour::solver(M, ecc); }, "M"_a, "ecc"_a);
 
 
-    m.def("keplerian2", &brandt::keplerian,
+    m.def("keplerian", &brandt::keplerian,
           "t"_a, "P"_a, "K"_a, "ecc"_a, "w"_a, "M0"_a, "M0_epoch"_a,
           KEPLERIAN_DOC);
     m.def("keplerian2", &brandt::keplerian2,