Fix for force_lever output issues

corelib/src/libs/SireCAS/lambdaschedule.cpp

@@ -238,10 +238,12 @@ QString LambdaSchedule::toString() const
         auto keys = this->stage_equations[i].keys();
         std::sort(keys.begin(), keys.end());
 
-        for (auto lever : keys)
+        for (const auto &lever : keys)
         {
+            auto output_name = lever;
+            output_name.replace("*::", "");
             lines.append(QString("    %1: %2")
-                             .arg(lever.replace("*::", ""))
+                             .arg(output_name)
                              .arg(this->stage_equations[i][lever].toOpenMMString()));
         }
     }
@@ -956,20 +958,35 @@ SireCAS::Expression LambdaSchedule::_getEquation(int stage,
                                      CODELOC);
 
     const auto default_lever = _get_lever_name("*", lever);
+    const auto default_force = _get_lever_name(force, "*");
+    const auto lever_name = _get_lever_name(force, lever);
 
-    if (force == "*")
+    const auto equations = this->stage_equations[stage];
+
+    // search from most specific to least specific
+    auto it = equations.find(lever_name);
+
+    if (it != equations.end())
     {
-        return this->stage_equations[stage].value(
-            default_lever, this->default_equations[stage]);
+        return it.value();
     }
-    else
+
+    it = equations.find(default_force);
+
+    if (it != equations.end())
+    {
+        return it.value();
+    }
+
+    it = equations.find(default_lever);
+
+    if (it != equations.end())
     {
-        return this->stage_equations[stage].value(
-            _get_lever_name(force, lever),
-            this->stage_equations[stage].value(
-                default_lever,
-                this->default_equations[stage]));
+        return it.value();
     }
+
+    // we don't have any match, so return the default equation for this stage
+    return this->default_equations[stage];
 }
 
 /** Return the equation used to control the specified 'lever'
@@ -1142,15 +1159,33 @@ QHash<QString, QVector<double>> LambdaSchedule::getLeverValues(
     QVector<double> values(lambda_values.count(), NAN);
 
     QHash<QString, QVector<double>> lever_values;
-    lever_values.reserve(this->lever_names.count() + 1);
+
+    // get all of the lever / force combinations in use
+    QSet<QString> all_levers;
+
+    for (const auto &equations : this->stage_equations)
+    {
+        for (const auto &lever : equations.keys())
+        {
+            all_levers.insert(lever);
+        }
+    }
+
+    QStringList levers = all_levers.values();
+    std::sort(levers.begin(), levers.end());
+
+    lever_values.reserve(levers.count() + 2);
 
     lever_values.insert("λ", values);
 
     lever_values.insert("default", values);
 
-    for (const auto &lever_name : this->lever_names)
+    for (const auto &lever : levers)
     {
-        lever_values.insert(lever_name, values);
+        if (lever.startsWith("*::"))
+            lever_values.insert(lever.mid(3), values);
+        else
+            lever_values.insert(lever, values);
     }
 
     if (this->nStages() == 0)
@@ -1174,12 +1209,16 @@ QHash<QString, QVector<double>> LambdaSchedule::getLeverValues(
         const auto equation = this->default_equations[stage];
         lever_values["default"][i] = equation(input_values);
 
-        for (const auto &lever_name : lever_names)
+        for (const auto &lever : levers)
         {
-            const auto equation = this->stage_equations[stage].value(
-                lever_name, this->default_equations[stage]);
+            auto parts = lever.split("::");
+
+            const auto equation = this->_getEquation(stage, parts[0], parts[1]);
 
-            lever_values[lever_name][i] = equation(input_values);
+            if (lever.startsWith("*::"))
+                lever_values[lever.mid(3)][i] = equation(input_values);
+            else
+                lever_values[lever][i] = equation(input_values);
         }
     }
 

tests/cas/test_lambdaschedule.py

@@ -10,6 +10,64 @@ def _assert_same_equation(x, eq1, eq2):
         assert eq1.evaluate(val) == pytest.approx(eq2.evaluate(val), 1e-5)
 
 
+def test_charge_scale():
+    l = sr.cas.LambdaSchedule.standard_morph()
+
+    morph_equation = l.get_equation(stage="morph")
+
+    l.add_charge_scale_stages()
+
+    _assert_same_equation(l.lam(), l.get_equation(stage="morph"), morph_equation)
+
+    assert l.get_stages() == ["decharge", "morph", "recharge"]
+
+    assert l.get_levers() == ["charge"]
+
+    assert l.get_constant("γ") == 0.2
+
+    gamma = l.get_constant_symbol("γ")
+
+    scaled_morph = gamma * morph_equation
+
+    _assert_same_equation(
+        l.lam(), l.get_equation(stage="morph", lever="charge"), scaled_morph
+    )
+
+    _assert_same_equation(
+        gamma, l.get_equation(stage="morph", lever="charge"), scaled_morph
+    )
+
+    l.set_equation(stage="recharge", force="ghost/ghost", lever="charge", equation=0.5)
+
+    l.set_equation(stage="*", force="ghost/ghost", lever="*", equation=1.5)
+
+    _assert_same_equation(
+        l.lam(),
+        l.get_equation(stage="recharge", force="ghost/ghost", lever="charge"),
+        sr.cas.Expression(0.5),
+    )
+
+    _assert_same_equation(
+        l.lam(),
+        l.get_equation(stage="recharge", lever="LJ"),
+        (1.0 - ((1.0 - gamma) * (1.0 - l.lam()))) * l.final(),
+    )
+
+    _assert_same_equation(l.lam(), l.get_equation(stage="recharge"), l.final())
+
+    _assert_same_equation(
+        l.lam(),
+        l.get_equation(stage="recharge", force="ghost/non-ghost", lever="LJ"),
+        l.final(),
+    )
+
+    _assert_same_equation(
+        l.lam(),
+        l.get_equation(stage="recharge", force="ghost/ghost", lever="LJ"),
+        sr.cas.Expression(1.5),
+    )
+
+
 def test_lambdaschedule():
     l = sr.cas.LambdaSchedule.standard_morph()