helper params for param opti with assert_eq!(opti_res, clean_cost)

phyne/data/tkf/params/tree.nwk

@@ -0,0 +1 @@
+((((((T19:0.18253564794288973,T38:0.11367942155605133)I0:0.16425146544541422,(((T9:0.22583938333180698,T64:0.22097885669157666)I1:0.054183821944129976,(T26:0.3340918865732189,T45:0.020802408576665286)I2:0.09063885549294241)I3:0.0047460824223079555,T53:0.32209213609679777)I4:0.1839292401884232)I5:0.07314241512109995,((T58:0.08394683607877096,(T27:0.0491489104689503,T20:0.1362122314698043)I6:0.01650547329500701)I7:0.48267922529669793,T28:0.3044049327795188)I8:0.008367001972542086)I9:0.010196170811873594,((((T55:0.12428046929819993,T60:0.050085343119415704)I10:0.11820270164771268,(T49:0.01848690698958188,T63:0.1774423979803579)I11:0.10838789088802715)I12:0.1376004245586173,(T11:0.1642276343063575,T16:0.23394546782410375)I13:0.24476506567139106)I14:0.004237620391101004,((T17:0.44997272536442356,T7:0.01393710522409032)I15:0.026726619505505467,(T43:0.010296842961598943,(T12:0.013285992238172823,(T61:0.04854232125532925,(T15:0.08008380273115935,(T2:0.1884842957547165,T62:0.008318333310841837)I16:0.04285393226529239)I17:0.059542193031121564)I18:0.22895126415363304)I19:0.06569601149001386)I20:0.056123869230948274)I21:0.21302543812529912)I22:0.04004580133067057)I23:0.012445655095494477,(T59:0.05057180362001219,((((((T34:0.009011379268137306,(T35:0.15092936735000784,T18:0.14421500117374628)I24:0.2195911543125767)I25:0.19482360545148897,(T3:0.13497597675010692,((T47:0.004173032352879809,T8:0.03698040578185024)I26:0.10354138339568661,(T33:0.07284176715819353,T4:0.006926480297948681)I27:0.003565525870755711)I28:0.07343030357240612)I29:0.12164092793060345)I30:0.015668637774335442,T6:0.023394260001709032)I31:0.04663931056287233,((T57:0.07536371686048421,T46:0.015844617929883884)I32:0.034610637639086665,T23:0.06056260240953739)I33:0.18284343535272585)I34:0.008756284636401615,T30:0.020400223534816533)I35:0.09026894762141248,((T24:0.01938171467253015,T41:0.0709037427939508)I36:0.10666773456595566,(T32:0.008738851914929117,(T52:0.014363125985653624,(T14:0.004537605881616293,(T10:0.1617130077530088,T42:0.027224014132103343)I37:0.11072710045883706)I38:0.02496444232084729)I39:0.2935687797149779)I40:0.14578540314359661)I41:0.02541892067309106)I42:0.16648929748105917)I43:0.07465518152493947)I44:0.2633651531166342,((((T5:0.0015642318939174669,T21:0.010929646468554873)I45:0.030608235578181143,T39:0.1676954040171613)I46:0.03168557455717944,((T48:0.060082658555633786,T36:0.04719974765845245)I47:0.008110770252420133,T54:0.08040962503719787)I48:0.23073639547292413)I49:0.029391830019327318,(((T13:0.12467911675223041,T31:0.07625879488099768)I50:0.04684824250198325,T37:0.07923244627200617)I51:0.016833079514477097,(((T22:0.10325235201624201,(T40:0.23917218416733682,T44:0.04287095335037999)I52:0.44144184404921705)I53:0.31356753624075046,(T25:0.16354471117510178,T51:0.03041603974875198)I54:0.07837158237710552)I55:0.03727556659982548,((T50:0.014078595495307702,T29:0.2852890552055393)I56:0.11119058202350303,(T56:0.035043706834188504,T1:0.1719311210997305)I57:0.021675960026272287)I58:0.00011928968311421468)I59:0.021631984247609603)I60:0.13365688129311915)I61:0.16167251826699555)I62:0.1924133290575341;

phyne/src/likelihood/mod.rs

@@ -41,6 +41,32 @@ pub trait ModelSearchCost {
     fn set_param(&mut self, param: usize, value: f64);
     /// Returns the valid range for a model parameter [min, max], inclusive.
     fn param_range(&self, param: usize) -> ParamRange;
+
+    // These are not atomic parameters of the model but rather combinations of parameters like their
+    // product or ratio. They are useful during parameters optimisation for the case that the atomic
+    // parameters are not independent and the optimisation is more efficient in the space of the
+    // helper parameters.
+    fn helper_param_count(&self) -> usize {
+        0
+    }
+    fn helper_param(&self, _param: usize) -> f64 {
+        panic!(
+            "helper_param not implemented for this model. \
+             Make sure to implement this method if helper_param_count() > 0"
+        );
+    }
+    fn set_helper_param(&mut self, _param: usize, _value: f64) {
+        panic!(
+            "set_helper_param not implemented for this model. \
+             Make sure to implement this method if helper_param_count() > 0"
+        );
+    }
+    fn helper_param_range(&self, _param: usize) -> ParamRange {
+        panic!(
+            "helper_param_range not implemented for this model. \
+             Make sure to implement this method if helper_param_count() > 0"
+        );
+    }
     fn set_freqs(&mut self, freqs: FreqVector) -> Result<()>;
     fn empirical_freqs(&self) -> FreqVector;
     fn freqs(&self) -> &FreqVector;

phyne/src/optimisers/model_optimiser.rs

@@ -6,7 +6,7 @@ use argmin::solver::brent::BrentOpt;
 use log::{debug, info, warn};
 
 use crate::evolutionary_models::FrequencyOptimisation;
-use crate::likelihood::ModelSearchCost;
+use crate::likelihood::{ModelSearchCost, ParamRange};
 use crate::optimisers::{ModelOptimisationResult, SingleValOptResult, StopCondition};
 use crate::Result;
 
@@ -102,19 +102,64 @@ impl<C: ModelSearchCost + Display + Clone> ModelOptimiser<C> {
         Ok(cost)
     }
 
+    fn num_params(&self) -> usize {
+        self.c.param_count() + self.c.helper_param_count()
+    }
+
+    fn get_param_or_helper(&self, idx: usize) -> f64 {
+        if idx < self.c.param_count() {
+            self.c.param(idx)
+        } else {
+            self.c.helper_param(idx - self.c.param_count())
+        }
+    }
+
+    fn set_param_or_helper(&mut self, idx: usize, value: f64) {
+        if idx < self.c.param_count() {
+            self.c.set_param(idx, value);
+        } else {
+            self.c.set_helper_param(idx - self.c.param_count(), value);
+        }
+    }
+
+    fn param_range_or_helper(&self, idx: usize) -> ParamRange {
+        if idx < self.c.param_count() {
+            self.c.param_range(idx)
+        } else {
+            self.c.helper_param_range(idx - self.c.param_count())
+        }
+    }
+
     fn single_optimisation_iteration(&mut self) -> Result<f64> {
         let mut curr_cost = self.c.cost();
-        for param in 0..self.c.param_count() {
-            let start_value = self.c.param(param);
+        for param in 0..self.num_params() {
+            println!("Optimising parameter {param:?} with current cost {curr_cost}");
+            let start_value = self.get_param_or_helper(param);
             debug!("Optimising parameter {param:?} from value {start_value} with cost {curr_cost}");
             let param_opt = self.opt_parameter(param, start_value)?;
             if param_opt.final_cost < curr_cost {
                 // Parameter will have been reset by the optimiser, set it back to start value
-                self.c.set_param(param, start_value);
+                // self.c.set_param(param, start_value);
                 continue;
             }
-            self.c.set_param(param, param_opt.value);
+            self.set_param_or_helper(param, param_opt.value);
+
+            // The helper might set the lambda and mu correctly when setting a helper param,
+            // but not the before_helper_lambda and before_helper_mu.
+            // Alternatively we could have two trait methods one for setting helpers in the
+            // optimiser and one for setting helpers. Could be also called set_helper_repeatedly
+            // (for the optimiser that stays correct when calling the idx for the setter), and set
+            // helper that is not intended for repeated calls. Dont know what good names would be
+            for set_param_idx in 0..self.c.param_count() {
+                self.c.set_param(set_param_idx, self.c.param(set_param_idx));
+            }
             curr_cost = param_opt.final_cost;
+            debug_assert_eq!(
+                curr_cost,
+                self.c.cost(),
+                "Cost mismatch after optimisation of parameter {param:?}"
+            );
+
             debug!(
                 "Optimised parameter {param:?} to value {} with cost {curr_cost}",
                 param_opt.value
@@ -134,7 +179,8 @@ impl<C: ModelSearchCost + Display + Clone> ModelOptimiser<C> {
             cost: RefCell::new(self.c.clone()),
             param,
         };
-        let range = self.c.param_range(param);
+        //let range = self.c.param_range(param);
+        let range = self.param_range_or_helper(param);
         let min = range.0;
         let max = range.1.min(start_value * 100.0);
         debug_assert!(min <= start_value && start_value <= max);
@@ -155,6 +201,16 @@ pub(crate) struct ParamOptimiser<C: ModelSearchCost> {
     pub(crate) param: usize,
 }
 
+impl<C: ModelSearchCost> ParamOptimiser<C> {
+    fn set_param_or_helper(&self, cost: &mut C, param: usize, value: f64) {
+        if param < cost.param_count() {
+            cost.set_param(param, value);
+        } else {
+            cost.set_helper_param(param - cost.param_count(), value);
+        }
+    }
+}
+
 impl<C: ModelSearchCost> CostFunction for ParamOptimiser<C> {
     type Param = f64;
     type Output = f64;
@@ -165,7 +221,8 @@ impl<C: ModelSearchCost> CostFunction for ParamOptimiser<C> {
         } else {
             *value
         };
-        self.cost.borrow_mut().set_param(self.param, value);
+        // self.cost.borrow_mut().set_param(self.param, value);
+        self.set_param_or_helper(&mut self.cost.borrow_mut(), self.param, value);
         Ok(-self.cost.borrow().cost())
     }
 

phyne/src/optimisers/model_optimiser_tests.rs

@@ -602,7 +602,9 @@ fn tkf_model_opti_template<C: ModelSearchCost + Clone + Display>(c: C) {
         .unwrap();
     assert_eq!(intermediate_cost, o.initial_cost);
     assert_eq!(o.final_cost, o.cost.cost());
-    assert!(o.initial_cost < o.final_cost);
+    // The following assertion is not guaranteed to be true as the empirical optimisation may not
+    // improve the likelihood.
+    // assert!(o.initial_cost < o.final_cost);
     assert_eq!(o.cost.freqs(), &o.cost.empirical_freqs());
     for param in 0..c.param_count() {
         assert_ne!(c.param(param), o.cost.param(param));
@@ -640,3 +642,19 @@ fn tkf92_model_opti() {
         .unwrap();
     tkf_model_opti_template(tkf92);
 }
+
+#[cfg_attr(feature = "ci_coverage", ignore)]
+#[test]
+fn tkf92_model_opti_params() {
+    let fldr = Path::new("./data/tkf/params/");
+    let info = PIB::with_attrs(fldr.join("masa.fasta"), fldr.join("tree.nwk"))
+        .build_with_ancestors()
+        .unwrap();
+    // let subst_model = SubstModel::<HKY>::new(&[], &[2.0]);
+    let subst_model = SubstModel::<HKY>::default();
+    let tkf_model = TKF92IndelModel::new(&[0.01, 0.01025, 0.65]).unwrap();
+    let tkf92 = TKFCostBuilder::new(subst_model, tkf_model, info)
+        .build()
+        .unwrap();
+    tkf_model_opti_template(tkf92);
+}

phyne/src/tkf_model/mod.rs

@@ -90,6 +90,22 @@ impl<Q: QMatrix, T: TKFModel, AA: AncestralAlignment> ModelSearchCost for TKFCos
     fn freqs(&self) -> &FreqVector {
         self.subst_cost.freqs()
     }
+
+    fn helper_param_count(&self) -> usize {
+        self.indel_cost.helper_param_count()
+    }
+
+    fn helper_param(&self, param: usize) -> f64 {
+        self.indel_cost.helper_param(param)
+    }
+
+    fn set_helper_param(&mut self, param: usize, value: f64) {
+        self.indel_cost.set_helper_param(param, value);
+    }
+
+    fn helper_param_range(&self, param: usize) -> ParamRange {
+        self.indel_cost.helper_param_range(param)
+    }
 }
 
 impl<Q: QMatrix, T: TKFModel, AA: AncestralAlignment> TreeSearchCost for TKFCost<Q, T, AA> {

phyne/src/tkf_model/tkf92.rs

@@ -27,13 +27,28 @@ pub(crate) enum TKF92Parameters {
     Invalid(usize),
 }
 
+pub(crate) const TKF92_HELPER_PARAM_COUNT: usize = 2;
+
+#[derive(Debug, Eq, PartialEq, FromPrimitive, IntoPrimitive)]
+#[repr(usize)]
+pub(crate) enum TKF92HelperParameters {
+    Rate = 0,
+    Ratio = 1,
+    #[num_enum(catch_all)]
+    Invalid(usize),
+}
+
 #[derive(Clone, Debug, PartialEq)]
 pub struct TKF92IndelModel {
     params: Vec<f64>,
     /// precomputed r.ln()
     ln_r: f64,
     /// precomputed ((1 - r)/r).ln()
     ln_one_minus_r_over_r: f64,
+    /// original lambda before ratio or rate was set
+    before_helper_lambda: f64,
+    /// original mu before ratio or rate was set
+    before_helper_mu: f64,
 }
 
 impl TKF92IndelModel {
@@ -42,16 +57,17 @@ impl TKF92IndelModel {
     /// 0 < lambda < mu < infinity, 0 < r < 1.
     pub fn new(params: &[f64]) -> Result<Self> {
         valid_number_tkf_params("TKF92", "lambda, mu, r", params, TKF92_PARAM_COUNT)?;
-        valid_lambda_mu(
-            params[usize::from(TKF92Parameters::Lambda)],
-            params[usize::from(TKF92Parameters::Mu)],
-        )?;
-        valid_r(params[usize::from(TKF92Parameters::R)])?;
+        let lambda = params[usize::from(TKF92Parameters::Lambda)];
+        let mu = params[usize::from(TKF92Parameters::Mu)];
+        valid_lambda_mu(lambda, mu)?;
         let r = params[usize::from(TKF92Parameters::R)];
+        valid_r(r)?;
         Ok(Self {
             params: params.to_vec(),
             ln_r: r.ln(),
             ln_one_minus_r_over_r: ((1.0 - r) / r).ln(),
+            before_helper_lambda: lambda,
+            before_helper_mu: mu,
         })
     }
 
@@ -67,6 +83,8 @@ impl Default for TKF92IndelModel {
             params: vec![DEFAULT_LAMBDA, DEFAULT_MU, r],
             ln_r: r.ln(),
             ln_one_minus_r_over_r: (-r).ln_1p() - r.ln(),
+            before_helper_lambda: DEFAULT_LAMBDA,
+            before_helper_mu: DEFAULT_MU,
         }
     }
 }
@@ -92,10 +110,17 @@ impl TKFModel for TKF92IndelModel {
                 self.ln_r = value.ln();
                 self.ln_one_minus_r_over_r = (-value).ln_1p() - value.ln();
             }
-            _ => {
+            TKF92Parameters::Lambda => {
                 self.params[idx] = value;
+                self.before_helper_lambda = value;
             }
+            TKF92Parameters::Mu => {
+                self.params[idx] = value;
+                self.before_helper_mu = value;
+            }
+            _ => panic!("Invalid parameter index for TKF model: {param:?}"),
         };
+        println!("set param: lambda = {}, mu = {}", self.lambda(), self.mu(),);
     }
 
     fn param_range(&self, idx: usize) -> ParamRange {
@@ -134,6 +159,53 @@ impl TKFModel for TKF92IndelModel {
     fn get_blocks<AA: AncestralAlignment>(&self, msa: &AA) -> Vec<usize> {
         blocks_of_alignment(msa)
     }
+
+    fn helper_param_count(&self) -> usize {
+        // with that instead of zero the optimisation converges much faster and achieves better logl
+        TKF92_HELPER_PARAM_COUNT
+        // 0
+    }
+
+    fn helper_param(&self, idx: usize) -> f64 {
+        let helper_param = TKF92HelperParameters::from_primitive(idx);
+        match helper_param {
+            TKF92HelperParameters::Rate => (self.lambda() + self.mu()) / 2.0,
+            TKF92HelperParameters::Ratio => self.lambda() / self.mu(),
+            _ => panic!("Invalid helper parameter index for TKF model: {helper_param:?}"),
+        }
+    }
+
+    fn set_helper_param(&mut self, idx: usize, value: f64) {
+        let param = TKF92HelperParameters::from_primitive(idx);
+        match param {
+            TKF92HelperParameters::Rate => {
+                let lambda = self.before_helper_lambda;
+                let mu = self.before_helper_mu;
+                self.params[usize::from(TKF92Parameters::Lambda)] =
+                    2.0 * value / (mu / lambda + 1.0);
+                self.params[usize::from(TKF92Parameters::Mu)] = 2.0 * value / (lambda / mu + 1.0);
+            }
+            TKF92HelperParameters::Ratio => {
+                let lambda = self.before_helper_lambda;
+                let mu = self.before_helper_mu;
+                let x = (lambda + mu) / 2.0;
+                let y = (x * (1.0 - value)) / (1.0 + value);
+                self.params[usize::from(TKF92Parameters::Lambda)] = x - y;
+                self.params[usize::from(TKF92Parameters::Mu)] = x + y;
+            }
+            _ => panic!("Invalid helper parameter index for TKF model: {param:?}"),
+        }
+        println!("set helper: lambda = {}, mu = {}", self.lambda(), self.mu());
+    }
+
+    fn helper_param_range(&self, _idx: usize) -> ParamRange {
+        let param = TKF92HelperParameters::from_primitive(_idx);
+        match param {
+            TKF92HelperParameters::Rate => (f64::EPSILON, f64::MAX),
+            TKF92HelperParameters::Ratio => (f64::EPSILON, 1.0 - f64::EPSILON),
+            _ => panic!("Invalid helper parameter index for TKF model: {param:?}"),
+        }
+    }
 }
 
 impl Display for TKF92IndelModel {
@@ -297,6 +369,8 @@ mod private_tests {
             params: vec![0.5, 1.0, 0.3],
             ln_r: 0.0, // cache filled with dummy, since it is not needed here
             ln_one_minus_r_over_r: 0.0, // cache filled with dummy since it is not needed here
+            before_helper_lambda: 0.0, // cache filled with dummy since it is not needed here
+            before_helper_mu: 0.0, // cache filled with dummy since it is not
         };
         // Use an invalid index
         model.param_range(3);
@@ -308,6 +382,8 @@ mod private_tests {
             params: vec![1.1, 2.0, 0.3],
             ln_r: 0.0,                  // cache filled with dummy since it is not printed
             ln_one_minus_r_over_r: 0.0, // cache filled with dummy since it is not printed
+            before_helper_lambda: 0.0,  // cache filled with dummy since it is not needed here
+            before_helper_mu: 0.0,      // cache filled with dummy since it is not
         };
 
         let fmt = format!("{}", tkf_indel_model);
@@ -321,6 +397,8 @@ mod private_tests {
             params: vec![1.0, 2.0, 0.3],
             ln_r: 0.0,                  // dummy
             ln_one_minus_r_over_r: 0.0, // dummy
+            before_helper_lambda: 0.0,  // cache filled with dummy since it is not needed here
+            before_helper_mu: 0.0,      // cache filled with dummy since it is not
         };
         let new_lambda = 1.1;
         model.set_param(usize::from(TKF92Parameters::Lambda), new_lambda);