Line data    Source code

       1             : /* +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
       2             :    Copyright (c) 2026 Pengchao Zhang
       3             : 
       4             :    See http://www.plumed.org for more information.
       5             : 
       6             :    This file is part of plumed, version 2.
       7             : 
       8             :    plumed is free software: you can redistribute it and/or modify
       9             :    it under the terms of the GNU Lesser General Public License as published by
      10             :    the Free Software Foundation, either version 3 of the License, or
      11             :    (at your option) any later version.
      12             : 
      13             :    plumed is distributed in the hope that it will be useful,
      14             :    but WITHOUT ANY WARRANTY; without even the implied warranty of
      15             :    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
      16             :    GNU Lesser General Public License for more details.
      17             : 
      18             :    You should have received a copy of the GNU Lesser General Public License
      19             :    along with plumed.  If not, see <http://www.gnu.org/licenses/>.
      20             : +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ */
      21             : #include "function/Function.h"
      22             : #include "core/ActionRegister.h"
      23             : #include "tools/Communicator.h"
      24             : #include <algorithm>
      25             : #include <cmath>
      26             : #include <vector>
      27             : 
      28             : namespace PLMD {
      29             : namespace pathbias {
      30             : 
      31             : //+PLUMEDOC FUNCTION PATH_LOGMEANEXP
      32             : /*
      33             : Compute a stable log arithmetic mean across synchronized replicas.
      34             : 
      35             : ARG is one finite, nonperiodic dimensionless scalar ell_b on each replica.
      36             : The result is log(sum_b exp(ell_b)/P). Its local derivative is
      37             : exp(ell_b)/sum_c exp(ell_c), with no additional factor of 1/P. Only replica
      38             : roots communicate across replicas; spatial MPI ranks are not extra beads.
      39             : EXPECTED_REPLICAS is required to catch a missing or incorrect replica group.
      40             : All replicas must execute the same graph at the same step and time.
      41             : Numerical derivatives at the action level are unsupported: finite-difference
      42             : tests must perturb one replica's input at a time.
      43             : 
      44             : For a common frozen energy field v(s), use ell=-v/kBT and apply -kBT times
      45             : this result as a complete-path bias. This averages the effective target to
      46             : unbiased probability ratio, not the density estimate or its logarithm.
      47             : 
      48             : ```plumed
      49             : #SETTINGS NREPLICAS=2
      50             : d: DISTANCE ATOMS=1,2 NOPBC
      51             : v: CUSTOM ARG=d FUNC=x*x PERIODIC=NO
      52             : ell: CUSTOM ARG=v FUNC=-x/2.5 PERIODIC=NO
      53             : lm: PATH_LOGMEANEXP ARG=ell EXPECTED_REPLICAS=2
      54             : energy: CUSTOM ARG=lm FUNC=-2.5*x PERIODIC=NO
      55             : bias: BIASVALUE ARG=energy
      56             : ```
      57             : 
      58             : Verify identical field parameters/state hashes before a run or restart.
      59             : This action has no bias history and does not make adaptive OPES learning
      60             : compatible with this Hamiltonian. If v is an active frozen bias action,
      61             : apply the correction V_path-v instead of adding V_path on top of it.
      62             : The molecular dynamics adapter owns dynamical force scaling and single-path
      63             : energy accounting. Uniform bead observables still use one path weight
      64             : exp(V_path/kBT); the local derivative weights are not observable weights.
      65             : */
      66             : //+ENDPLUMEDOC
      67             : 
      68             : class PathLogMeanExp : public function::Function {
      69             :   int replicas=1;
      70             :   int replica=0;
      71             :   void collectiveCheck(int invalid);
      72             : public:
      73             :   explicit PathLogMeanExp(const ActionOptions&);
      74             :   static void registerKeywords(Keywords&);
      75             :   void calculate() override;
      76             : };
      77             : 
      78             : PLUMED_REGISTER_ACTION(PathLogMeanExp,"PATH_LOGMEANEXP")
      79             : 
      80         152 : void PathLogMeanExp::registerKeywords(Keywords& keys) {
      81         152 :   Function::registerKeywords(keys);
      82         152 :   keys.add("compulsory","EXPECTED_REPLICAS","number of synchronized replicas in this path");
      83         304 :   keys.setValueDescription("scalar","the dimensionless log arithmetic mean of replica ratios");
      84         152 : }
      85             : 
      86        1172 : void PathLogMeanExp::collectiveCheck(int invalid) {
      87        1172 :   comm.Sum(invalid);
      88        1172 :   if(comm.Get_rank()==0) {
      89         841 :     multi_sim_comm.Sum(invalid);
      90             :   }
      91        1172 :   comm.Bcast(invalid,0);
      92        1172 :   if(invalid) {
      93          24 :     error("invalid scalar, replica count, numerical derivative request, or nonfinite input on at least one replica");
      94             :   }
      95        1148 : }
      96             : 
      97         148 : PathLogMeanExp::PathLogMeanExp(const ActionOptions& ao):Action(ao),Function(ao) {
      98         148 :   if(comm.Get_rank()==0) {
      99          89 :     replicas=multi_sim_comm.Get_size();
     100          89 :     replica=multi_sim_comm.Get_rank();
     101             :   }
     102         148 :   comm.Bcast(replicas,0);
     103         148 :   comm.Bcast(replica,0);
     104         148 :   int expected=0;
     105         148 :   parse("EXPECTED_REPLICAS",expected);
     106         148 :   int invalid=(expected<1 || expected!=replicas || getNumberOfArguments()!=1);
     107         148 :   if(getNumberOfArguments()==1) {
     108         148 :     invalid=invalid || getPntrToArgument(0)->getRank()!=0 || getPntrToArgument(0)->isPeriodic();
     109             :   }
     110           0 :   invalid=invalid || checkNumericalDerivatives();
     111         148 :   collectiveCheck(invalid);
     112         136 :   checkRead();
     113         136 :   addValueWithDerivatives();
     114         136 :   setNotPeriodic();
     115         136 :   log.printf("  stable log mean over %d synchronized replicas\n",replicas);
     116         160 : }
     117             : 
     118        1024 : void PathLogMeanExp::calculate() {
     119        1024 :   const double value=getArgument(0);
     120        1024 :   collectiveCheck(!std::isfinite(value));
     121        1012 :   std::vector<double> values(3*replicas,0.0);
     122        1012 :   if(comm.Get_rank()==0) {
     123         745 :     values[replica]=value;
     124         745 :     values[replicas+replica]=static_cast<double>(getStep());
     125         745 :     values[2*replicas+replica]=getTime();
     126         745 :     multi_sim_comm.Sum(values);
     127             :   }
     128        1012 :   comm.Bcast(values,0);
     129        3976 :   for(int b=0; b<replicas; ++b) {
     130        2972 :     if(values[replicas+b]!=values[replicas] || values[2*replicas+b]!=values[2*replicas]) {
     131           8 :       error("replicas must evaluate at identical steps and times");
     132             :     }
     133             :   }
     134        1004 :   const double maximum=*std::max_element(values.begin(),values.begin()+replicas);
     135             :   double sum=0;
     136        3960 :   for(int b=0; b<replicas; ++b) {
     137        2956 :     sum+=std::exp(values[b]-maximum);
     138             :   }
     139        1012 :   setValue(maximum+std::log(sum/replicas));
     140        1004 :   setDerivative(0,std::exp(value-maximum)/sum);
     141        1004 : }
     142             : 
     143             : }
     144             : }