From 6fbf62536eed5b6cf1ab7aaa6af8d5f76def44de Mon Sep 17 00:00:00 2001
From: "Daniel R. Roe" <daniel.r.roe@gmail.com>
Date: Fri, 12 May 2023 12:54:14 -0400
Subject: [PATCH 01/15] Try adding avgucell to diffusion. Doesnt quite work
 yet, not sure why

---
 src/Action_Diffusion.cpp       | 46 +++++++++++++++++++++++++++++++---
 src/Action_Diffusion.h         |  2 ++
 src/Action_Unwrap.cpp          |  2 +-
 test/Test_Diffusion/RunTest.sh | 20 ++++++++++++++-
 4 files changed, 64 insertions(+), 6 deletions(-)

diff --git a/src/Action_Diffusion.cpp b/src/Action_Diffusion.cpp
index e21605b7bf..232be0c3f2 100644
--- a/src/Action_Diffusion.cpp
+++ b/src/Action_Diffusion.cpp
@@ -3,6 +3,7 @@
 #include "CpptrajStdio.h"
 #include "StringRoutines.h" // validDouble
 #include "DataSet_1D.h" // LinearRegression
+#include "DataSet_Mat3x3.h"
 #ifdef TIMER
 # include "Timer.h"
 #endif
@@ -24,12 +25,14 @@ Action_Diffusion::Action_Diffusion() :
   debug_(0),
   outputx_(0), outputy_(0), outputz_(0), outputr_(0), outputa_(0),
   diffout_(0),
-  masterDSL_(0)
+  masterDSL_(0),
+  avgucell_(0)
 {}
 
 static inline void ShortHelp() {
   mprintf("\t[{out <filename> | separateout <suffix>}] [time <time per frame>] [noimage]\n"
-          "\t[<mask>] [<set name>] [individual] [diffout <filename>] [nocalc]\n");
+          "\t[<mask>] [<set name>] [individual] [diffout <filename>] [nocalc]\n"
+          "\t[avgucell <avg ucell set>]\n");
 }
 
 void Action_Diffusion::Help() const {
@@ -110,6 +113,29 @@ Action::RetType Action_Diffusion::Init(ArgList& actionArgs, ActionInit& init, in
     }
   }
   if (diffout_ != 0) calcDiffConst_ = true;
+  // Get average box
+  std::string avgucellstr = actionArgs.GetStringKey("avgucell");
+  if (!avgucellstr.empty()) {
+    avgucell_ = init.DSL().GetDataSet( avgucellstr );
+    if (avgucell_ == 0) {
+      mprinterr("Error: No data set selected by '%s'\n", avgucellstr.c_str());
+      return Action::ERR;
+    }
+    if (avgucell_->Type() != DataSet::MAT3X3) {
+      mprinterr("Error: Average unit cell set '%s' is not a 3x3 matrix set.\n", avgucell_->legend());
+      return Action::ERR;
+    }
+    if (avgucell_->Size() < 1) {
+      mprinterr("Error: Average unit cell set '%s' is empty.\n", avgucell_->legend());
+      return Action::ERR;
+    }
+    DataSet_Mat3x3 const& matset = static_cast<DataSet_Mat3x3 const&>( *avgucell_ );
+    Matrix_3x3 const& ucell = matset[0];
+    if (avgbox_.SetupFromUcell( ucell )) {
+      mprinterr("Error: Could not set up box from unit cell parameters in '%s'\n", avgucell_->legend());
+      return Action::ERR;
+    }
+  }
   // Add DataSets
   dsname_ = actionArgs.GetStringNext();
   if (dsname_.empty())
@@ -205,6 +231,10 @@ Action::RetType Action_Diffusion::Init(ArgList& actionArgs, ActionInit& init, in
     mprintf("\tTo calculate diffusion constant from mean squared displacement plots,\n"
             "\t  calculate the slope of MSD vs time and multiply by 10.0/2*N (where N\n"
             "\t  is # of dimensions); this will give units of 1x10^-5 cm^2/s.\n");
+  if (avgucell_ != 0) {
+    mprintf("\tUsing average unit cell vectors from set '%s' to remove box fluctuations.\n", avgucell_->legend());
+    avgbox_.PrintInfo();
+  }
 # ifdef _OPENMP
 # pragma omp parallel
   {
@@ -245,8 +275,11 @@ Action::RetType Action_Diffusion::Setup(ActionSetup& setup) {
       return Action::ERR;
     }
 #   endif
-  } else
+  } else {
     mprintf("\tImaging disabled.\n");
+    if (avgucell_ != 0)
+      mprintf("Warning: 'avgucell' specified but trajectory has no unit cell.\n");
+  }
 
   // If imaging, reserve space for the previous fractional coordinates array
   if (imageOpt_.ImagingEnabled())
@@ -308,8 +341,13 @@ Action::RetType Action_Diffusion::Setup(ActionSetup& setup) {
 
 // Action_Diffusion::DoAction()
 Action::RetType Action_Diffusion::DoAction(int frameNum, ActionFrame& frm) {
+  Matrix_3x3 const* ucell = 0;
   if (imageOpt_.ImagingEnabled()) {
     imageOpt_.SetImageType( frm.Frm().BoxCrd().Is_X_Aligned_Ortho() );
+    if (avgucell_ == 0)
+      ucell = &(frm.Frm().BoxCrd().UnitCell());
+    else
+      ucell = &(avgbox_.UnitCell());
   }
 
   // ----- Load initial frame if necessary -------
@@ -363,7 +401,7 @@ Action::RetType Action_Diffusion::DoAction(int frameNum, ActionFrame& frm) {
       ixyz[2] = ixyz[2] - round(ixyz[2]);
       xyz_frac = previousFrac_[imask] + ixyz;
       // Back to Cartesian
-      Vec3 xyz_cart1 = frm.Frm().BoxCrd().UnitCell().TransposeMult( xyz_frac );
+      Vec3 xyz_cart1 = ucell->TransposeMult( xyz_frac );
       // Update reference frac coords
       previousFrac_[imask] = xyz_frac;
       // Calculate the distance between the fixed coordinates
diff --git a/src/Action_Diffusion.h b/src/Action_Diffusion.h
index e0f50136c6..09f602cef3 100644
--- a/src/Action_Diffusion.h
+++ b/src/Action_Diffusion.h
@@ -52,6 +52,8 @@ class Action_Diffusion : public Action {
     DataSetList* masterDSL_;
     std::string dsname_;
     Dimension Xdim_;
+    DataSet* avgucell_;   ///< Hold average unit cell parameters for removing box fluctuations
+    Box avgbox_;          ///< Hold average box for removing box fluctuations
 #   ifdef MPI
     Parallel::Comm trajComm_;
 #   endif
diff --git a/src/Action_Unwrap.cpp b/src/Action_Unwrap.cpp
index 638af63939..fe92f9360a 100644
--- a/src/Action_Unwrap.cpp
+++ b/src/Action_Unwrap.cpp
@@ -27,7 +27,7 @@ void Action_Unwrap::Help() const {
           "\t[ %s ] [<mask>]\n", DataSetList::RefArgs);
   mprintf("  Reverse of 'image'; unwrap coordinates in <mask> according\n"
           "  to the first frame, or optionally a reference structure. Can\n"
-          " unwrap by atom (default), by residue, or by molecule.\n");
+          "  unwrap by atom (default), by residue, or by molecule.\n");
 }
 
 // Action_Unwrap::Init()
diff --git a/test/Test_Diffusion/RunTest.sh b/test/Test_Diffusion/RunTest.sh
index 461c245507..6d79a0ee74 100755
--- a/test/Test_Diffusion/RunTest.sh
+++ b/test/Test_Diffusion/RunTest.sh
@@ -5,7 +5,8 @@
 INPUT="diffusion.in"
 
 CleanFiles $INPUT diff_?.xmgr diff.dat diff.1.dat diff.2.dat diff.3.dat nw.dat \
-           WAT_O.agr DC.dat Nonortho.agr Nonortho.dat noimage.agr noimage.dat
+           WAT_O.agr DC.dat Nonortho.agr Nonortho.dat noimage.agr noimage.dat \
+           tz2.ortho.wato.dat
 TESTNAME='Diffusion tests'
 Requires netcdf maxthreads 10
 
@@ -47,6 +48,20 @@ EOF
   DoTest Nonortho.agr.save Nonortho.agr
 }
 
+Test_diffusion_avgucell() {
+  TOP=../tz2.ortho.parm7
+  cat > $INPUT <<EOF
+trajin ../tz2.ortho.nc
+avgbox MyBox
+run
+
+diffusion Water :WAT@O out tz2.ortho.wato.dat avgucell MyBox[avg]
+run
+EOF
+  RunCpptraj "Diffusion with average box correction test"
+  DoTest ../Test_Unwrap/tz2.ortho.wato.dat.save tz2.ortho.wato.dat
+}
+
 Test_diffusion_noImage() {
   TOP=../tz2.truncoct.parm7
   cat > $INPUT <<EOF
@@ -100,6 +115,7 @@ if [ $? -eq 0 ] ; then
   Test_diffusion_oldSyntax
   Test_diffusion_newSyntax
   Test_diffusion_nonOrtho
+  Test_diffusion_avgucell
 fi
 UNITNAME='STFC diffusion tests'
 CheckFor notparallel
@@ -107,6 +123,8 @@ if [ $? -eq 0 ] ; then
   Test_stfc_diffusion
 fi
 
+
+
 EndTest
 
 exit 0

From 5a3b8efe1e98a8d17b49249a86983dcd6ce0453e Mon Sep 17 00:00:00 2001
From: "Daniel R. Roe" <daniel.r.roe@gmail.com>
Date: Mon, 15 May 2023 10:11:28 -0400
Subject: [PATCH 02/15] Do not use average unit cell on initial frame

---
 src/Action_Diffusion.cpp | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/src/Action_Diffusion.cpp b/src/Action_Diffusion.cpp
index 232be0c3f2..f8cfef4a71 100644
--- a/src/Action_Diffusion.cpp
+++ b/src/Action_Diffusion.cpp
@@ -368,6 +368,12 @@ Action::RetType Action_Diffusion::DoAction(int frameNum, ActionFrame& frm) {
       {
         previousFrac_.push_back( frm.Frm().BoxCrd().FracCell() * Vec3( initial_.XYZ(*atom) ) );
       }
+      if (avgucell_ != 0) {
+        // If using the average unit cell, we do not want it on the 
+        // initial frame since that can shift coordinates and the
+        // first frame should have a distance of 0.
+        ucell = &(frm.Frm().BoxCrd().UnitCell());
+      }
     }
   } // END initial frame load
   // ----- Diffusion calculation -----------------

From 0ffce0f9ef8dbf785f686505571d00a5cfc66f3b Mon Sep 17 00:00:00 2001
From: "Daniel R. Roe" <daniel.r.roe@gmail.com>
Date: Mon, 15 May 2023 10:13:42 -0400
Subject: [PATCH 03/15] Change name to avoid confusion

---
 src/Action_Diffusion.cpp | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/src/Action_Diffusion.cpp b/src/Action_Diffusion.cpp
index f8cfef4a71..7e0cc60694 100644
--- a/src/Action_Diffusion.cpp
+++ b/src/Action_Diffusion.cpp
@@ -394,7 +394,7 @@ Action::RetType Action_Diffusion::DoAction(int frameNum, ActionFrame& frm) {
     // Get current coords for this atom.
     const double* XYZ = frm.Frm().XYZ(at);
     // Get initial coords for this atom
-    const double* iXYZ = initial_.XYZ(at);
+    const double* initialXYZ = initial_.XYZ(at);
     // Calculate distance from initial position. 
     double delx, dely, delz;
     if ( imageOpt_.ImagingEnabled() ) {
@@ -412,14 +412,14 @@ Action::RetType Action_Diffusion::DoAction(int frameNum, ActionFrame& frm) {
       previousFrac_[imask] = xyz_frac;
       // Calculate the distance between the fixed coordinates
       // and reference (initial) frame coordinates.
-      delx = xyz_cart1[0] - iXYZ[0];
-      dely = xyz_cart1[1] - iXYZ[1];
-      delz = xyz_cart1[2] - iXYZ[2];
+      delx = xyz_cart1[0] - initialXYZ[0];
+      dely = xyz_cart1[1] - initialXYZ[1];
+      delz = xyz_cart1[2] - initialXYZ[2];
     } else {
       // No imaging. Calculate distance from current position to initial position.
-      delx = XYZ[0] - iXYZ[0];
-      dely = XYZ[1] - iXYZ[1];
-      delz = XYZ[2] - iXYZ[2];
+      delx = XYZ[0] - initialXYZ[0];
+      dely = XYZ[1] - initialXYZ[1];
+      delz = XYZ[2] - initialXYZ[2];
     }
     // Calc distances for this atom
     double distx = delx * delx;

From 14a5e74b201fdf570dc4ce318e3ffeee40608823 Mon Sep 17 00:00:00 2001
From: "Daniel R. Roe" <daniel.r.roe@gmail.com>
Date: Mon, 15 May 2023 10:14:14 -0400
Subject: [PATCH 04/15] Add 'noimage' keyword

---
 test/Test_Unwrap/RunTest.sh | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/test/Test_Unwrap/RunTest.sh b/test/Test_Unwrap/RunTest.sh
index cb105eb8ba..234480dbc5 100755
--- a/test/Test_Unwrap/RunTest.sh
+++ b/test/Test_Unwrap/RunTest.sh
@@ -31,7 +31,7 @@ avgbox MyBox
 run
 
 unwrap bymol avgucell MyBox[avg]
-diffusion Water :WAT@O out tz2.ortho.wato.dat
+diffusion Water :WAT@O out tz2.ortho.wato.dat noimage
 run
 EOF
 RunCpptraj "Unwrap with average box correction test"

From 87d9d16a9f538923f7cb8cd0813fe6f0ebebc162 Mon Sep 17 00:00:00 2001
From: "Daniel R. Roe" <daniel.r.roe@gmail.com>
Date: Mon, 15 May 2023 10:15:01 -0400
Subject: [PATCH 05/15] Update dependencies

---
 src/cpptrajdepend | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/cpptrajdepend b/src/cpptrajdepend
index 1b0f1e8201..db9400b0ff 100644
--- a/src/cpptrajdepend
+++ b/src/cpptrajdepend
@@ -24,7 +24,7 @@ Action_CreateReservoir.o : Action_CreateReservoir.cpp Action.h ActionState.h Act
 Action_DNAionTracker.o : Action_DNAionTracker.cpp Action.h ActionState.h Action_DNAionTracker.h ArgList.h AssociatedData.h Atom.h AtomMask.h AtomType.h BaseIOtype.h Box.h Constants.h CoordinateInfo.h CpptrajFile.h CpptrajStdio.h DataFile.h DataFileList.h DataSet.h DataSetList.h DataSet_Coords.h DataSet_Coords_REF.h Dimension.h DispatchObject.h DistRoutines.h FileIO.h FileName.h FileTypes.h Frame.h ImageOption.h MaskToken.h Matrix_3x3.h MetaData.h Molecule.h NameType.h Parallel.h ParameterHolders.h ParameterSet.h ParameterTypes.h Range.h ReferenceFrame.h ReplicaDimArray.h Residue.h Segment.h SymbolExporting.h TextFormat.h Timer.h Topology.h TypeNameHolder.h Unit.h Vec3.h
 Action_DSSP.o : Action_DSSP.cpp Action.h ActionState.h Action_DSSP.h ArgList.h AssociatedData.h Atom.h AtomMask.h AtomType.h BaseIOtype.h Box.h CharMask.h Constants.h CoordinateInfo.h CpptrajFile.h CpptrajStdio.h DataFile.h DataFileList.h DataSet.h DataSetList.h DataSet_Coords.h DataSet_Coords_REF.h Dimension.h DispatchObject.h DistRoutines.h FileIO.h FileName.h FileTypes.h Frame.h ImageOption.h MaskToken.h Matrix_3x3.h MetaData.h Molecule.h NameType.h Parallel.h ParameterHolders.h ParameterSet.h ParameterTypes.h Range.h ReferenceFrame.h ReplicaDimArray.h Residue.h Segment.h SymbolExporting.h TextFormat.h Timer.h Topology.h TypeNameHolder.h Unit.h Vec3.h
 Action_Density.o : Action_Density.cpp Action.h ActionState.h Action_Density.h ArgList.h AssociatedData.h Atom.h AtomMask.h AtomType.h BaseIOtype.h Box.h Constants.h CoordinateInfo.h CpptrajFile.h CpptrajStdio.h DataFile.h DataFileList.h DataSet.h DataSetList.h DataSet_1D.h DataSet_Coords.h DataSet_Coords_REF.h Dimension.h DispatchObject.h FileIO.h FileName.h FileTypes.h Frame.h MaskToken.h Matrix_3x3.h MetaData.h Molecule.h NameType.h OnlineVarT.h Parallel.h ParameterHolders.h ParameterSet.h ParameterTypes.h Range.h ReferenceFrame.h ReplicaDimArray.h Residue.h Segment.h StringRoutines.h SymbolExporting.h TextFormat.h Timer.h Topology.h TypeNameHolder.h Unit.h Vec3.h
-Action_Diffusion.o : Action_Diffusion.cpp Action.h ActionState.h Action_Diffusion.h ArgList.h AssociatedData.h Atom.h AtomMask.h AtomType.h BaseIOtype.h Box.h Constants.h CoordinateInfo.h CpptrajFile.h CpptrajStdio.h DataFile.h DataFileList.h DataSet.h DataSetList.h DataSet_1D.h DataSet_Coords.h DataSet_Coords_REF.h Dimension.h DispatchObject.h FileIO.h FileName.h FileTypes.h Frame.h ImageOption.h MaskToken.h Matrix_3x3.h MetaData.h Molecule.h NameType.h Parallel.h ParameterHolders.h ParameterSet.h ParameterTypes.h Range.h ReferenceFrame.h ReplicaDimArray.h Residue.h Segment.h StringRoutines.h SymbolExporting.h TextFormat.h Timer.h Topology.h TypeNameHolder.h Unit.h Vec3.h
+Action_Diffusion.o : Action_Diffusion.cpp Action.h ActionState.h Action_Diffusion.h ArgList.h AssociatedData.h Atom.h AtomMask.h AtomType.h BaseIOtype.h Box.h Constants.h CoordinateInfo.h CpptrajFile.h CpptrajStdio.h DataFile.h DataFileList.h DataSet.h DataSetList.h DataSet_1D.h DataSet_Coords.h DataSet_Coords_REF.h DataSet_Mat3x3.h Dimension.h DispatchObject.h FileIO.h FileName.h FileTypes.h Frame.h ImageOption.h MaskToken.h Matrix_3x3.h MetaData.h Molecule.h NameType.h Parallel.h ParameterHolders.h ParameterSet.h ParameterTypes.h Range.h ReferenceFrame.h ReplicaDimArray.h Residue.h Segment.h StringRoutines.h SymbolExporting.h TextFormat.h Timer.h Topology.h TypeNameHolder.h Unit.h Vec3.h
 Action_Dihedral.o : Action_Dihedral.cpp Action.h ActionState.h Action_Dihedral.h ArgList.h AssociatedData.h Atom.h AtomMask.h AtomType.h BaseIOtype.h Box.h Constants.h CoordinateInfo.h CpptrajFile.h CpptrajStdio.h DataFile.h DataFileList.h DataSet.h DataSetList.h DataSet_1D.h DataSet_Coords.h DataSet_Coords_REF.h DataSet_double.h Dimension.h DispatchObject.h FileIO.h FileName.h FileTypes.h Frame.h MaskToken.h Matrix_3x3.h MetaData.h Molecule.h NameType.h Parallel.h ParameterHolders.h ParameterSet.h ParameterTypes.h Range.h ReferenceFrame.h ReplicaDimArray.h Residue.h Segment.h SymbolExporting.h TextFormat.h Timer.h Topology.h TorsionRoutines.h TypeNameHolder.h Unit.h Vec3.h
 Action_DihedralRMS.o : Action_DihedralRMS.cpp Action.h ActionState.h Action_DihedralRMS.h ArgList.h AssociatedData.h Atom.h AtomMask.h AtomType.h BaseIOtype.h Box.h Constants.h CoordinateInfo.h CpptrajFile.h CpptrajStdio.h DataFile.h DataFileList.h DataSet.h DataSetList.h DataSet_Coords.h DataSet_Coords_REF.h DihedralSearch.h Dimension.h DispatchObject.h FileIO.h FileName.h FileTypes.h Frame.h MaskToken.h Matrix_3x3.h MetaData.h Molecule.h NameType.h Parallel.h ParameterHolders.h ParameterSet.h ParameterTypes.h Range.h ReferenceAction.h ReferenceFrame.h ReplicaDimArray.h Residue.h Segment.h SymbolExporting.h TextFormat.h Timer.h Topology.h TorsionRoutines.h TypeNameHolder.h Unit.h Vec3.h
 Action_Dipole.o : Action_Dipole.cpp Action.h ActionState.h Action_Dipole.h ArgList.h ArrayIterator.h AssociatedData.h Atom.h AtomMask.h AtomType.h BaseIOtype.h Box.h CharMask.h Constants.h CoordinateInfo.h CpptrajFile.h CpptrajStdio.h DataFile.h DataFileList.h DataSet.h DataSetList.h DataSet_3D.h DataSet_Coords.h DataSet_Coords_REF.h DataSet_GridFlt.h Dimension.h DispatchObject.h FileIO.h FileName.h FileTypes.h Frame.h Grid.h GridAction.h GridBin.h MaskToken.h Matrix_3x3.h MetaData.h Molecule.h NameType.h Parallel.h ParameterHolders.h ParameterSet.h ParameterTypes.h Range.h ReferenceFrame.h ReplicaDimArray.h Residue.h Segment.h SymbolExporting.h TextFormat.h Timer.h Topology.h TypeNameHolder.h Unit.h Vec3.h

From e158bc2d20c383e2a9236d4786871fa182a46860 Mon Sep 17 00:00:00 2001
From: "Daniel R. Roe" <daniel.r.roe@gmail.com>
Date: Tue, 16 May 2023 08:47:05 -0400
Subject: [PATCH 06/15] Ensure all processes get a copy of the average box.

---
 src/Action_AvgBox.cpp | 18 +++++++++++++++++-
 1 file changed, 17 insertions(+), 1 deletion(-)

diff --git a/src/Action_AvgBox.cpp b/src/Action_AvgBox.cpp
index b475c1c73c..8b04a2ebd5 100644
--- a/src/Action_AvgBox.cpp
+++ b/src/Action_AvgBox.cpp
@@ -60,6 +60,7 @@ Action::RetType Action_AvgBox::DoAction(int frameNum, ActionFrame& frm)
 
 # ifdef MPI
 int Action_AvgBox::SyncAction() {
+  Matrix_3x3 ucell;
   // Transfer all processes avg to master and combine.
   if (trajComm_.Master()) {
     // Master
@@ -70,6 +71,9 @@ int Action_AvgBox::SyncAction() {
       for (int i = 0; i < 9; i++, j+= 3)
         avgbox_[i].Combine( Stats<double>(recvbuf[j], recvbuf[j+1], recvbuf[j+2]) );
     }
+    for (int i = 0; i < 9; i++)
+      ucell[i] = avgbox_[i].mean();
+    ucell.Print("Average Unit Cell Vectors:");
   } else {
     // Children
     double sendbuf[27];
@@ -81,19 +85,31 @@ int Action_AvgBox::SyncAction() {
     }
     trajComm_.Send( sendbuf, 27, MPI_DOUBLE, 0, 2100 );
   }
-  // TODO broadcast?
+  // Broadcast the unit cell
+  ucell.BroadcastMatrix( trajComm_ );
+  // Add to DataSet on all processes 
+  ((DataSet_Mat3x3*)boxMatrix_)->AddMat3x3( ucell );
   return 0;
 }
 #endif
 
+/** Do averaging in serial. In parallel just print box since averaging
+  * is already done in SyncAction().
+  */
 void Action_AvgBox::Print() {
   mprintf("    AVGBOX:\n");
   if (avgbox_[0].nData() > 0) {
+#   ifdef MPI
+    // Unit cell data was already set up in SyncAction
+    DataSet_Mat3x3 const& dset = static_cast<DataSet_Mat3x3 const&>( *boxMatrix_ );
+    Matrix_3x3 const& ucell = dset[0];
+#   else
     Matrix_3x3 ucell;
     for (int i = 0; i < 9; i++)
       ucell[i] = avgbox_[i].mean();
     ucell.Print("Average Unit Cell Vectors:");
     ((DataSet_Mat3x3*)boxMatrix_)->AddMat3x3( ucell );
+#   endif
     Box thisBox;
     if (thisBox.SetupFromUcell( ucell )) {
       mprintf("Warning: Box appears to be invalid.\n");

From d185af514262d49003ff7041642202c1d1b0e917 Mon Sep 17 00:00:00 2001
From: "Daniel R. Roe" <daniel.r.roe@gmail.com>
Date: Tue, 16 May 2023 16:06:28 -0400
Subject: [PATCH 07/15] Allow calculation of diffusion in parallel with imaging
 using multiple time origins.

---
 src/Action_Diffusion.cpp | 127 +++++++++++++++++++++++++++++++++++++--
 src/Action_Diffusion.h   |   6 ++
 src/Parallel.cpp         |   2 +-
 src/Parallel.h           |   2 +-
 src/cpptrajdepend        |   2 +-
 5 files changed, 132 insertions(+), 7 deletions(-)

diff --git a/src/Action_Diffusion.cpp b/src/Action_Diffusion.cpp
index 7e0cc60694..b862088804 100644
--- a/src/Action_Diffusion.cpp
+++ b/src/Action_Diffusion.cpp
@@ -10,6 +10,9 @@
 #ifdef _OPENMP
 # include <omp.h>
 #endif
+#ifdef MPI
+#include "DataSet_double.h"
+#endif
 
 // CONSTRUCTOR
 Action_Diffusion::Action_Diffusion() :
@@ -27,6 +30,9 @@ Action_Diffusion::Action_Diffusion() :
   diffout_(0),
   masterDSL_(0),
   avgucell_(0)
+# ifdef MPI
+  ,multipleTimeOrigins_(false)
+# endif
 {}
 
 static inline void ShortHelp() {
@@ -56,6 +62,7 @@ Action::RetType Action_Diffusion::Init(ArgList& actionArgs, ActionInit& init, in
 {
 # ifdef MPI
   trajComm_ = init.TrajComm();
+  multipleTimeOrigins_ = false;
 # endif
   debug_ = debugIn;
   imageOpt_.InitImaging( !(actionArgs.hasKey("noimage")) );
@@ -130,7 +137,13 @@ Action::RetType Action_Diffusion::Init(ArgList& actionArgs, ActionInit& init, in
       return Action::ERR;
     }
     DataSet_Mat3x3 const& matset = static_cast<DataSet_Mat3x3 const&>( *avgucell_ );
+#   ifdef MPI
+    Matrix_3x3 ucell = matset[0];
+    // Broadcast the master unit cell
+    ucell.BroadcastMatrix( trajComm_ );
+#   else
     Matrix_3x3 const& ucell = matset[0];
+#   endif
     if (avgbox_.SetupFromUcell( ucell )) {
       mprinterr("Error: Could not set up box from unit cell parameters in '%s'\n", avgucell_->legend());
       return Action::ERR;
@@ -145,8 +158,10 @@ Action::RetType Action_Diffusion::Init(ArgList& actionArgs, ActionInit& init, in
   avg_z_ = init.DSL().AddSet(DataSet::DOUBLE, MetaData(dsname_, "Z"));
   avg_r_ = init.DSL().AddSet(DataSet::DOUBLE, MetaData(dsname_, "R"));
   avg_a_ = init.DSL().AddSet(DataSet::DOUBLE, MetaData(dsname_, "A"));
-  if (avg_x_ == 0 || avg_y_ == 0 || avg_z_ == 0 || avg_r_ == 0 || avg_a_ == 0)
+  if (avg_x_ == 0 || avg_y_ == 0 || avg_z_ == 0 || avg_r_ == 0 || avg_a_ == 0) {
+    mprinterr("Error: Could not allocate one or more average sets.\n");
     return Action::ERR;
+  }
   if (outputr_ != 0) outputr_->AddDataSet( avg_r_ );
   if (outputx_ != 0) outputx_->AddDataSet( avg_x_ );
   if (outputy_ != 0) outputy_->AddDataSet( avg_y_ );
@@ -169,7 +184,10 @@ Action::RetType Action_Diffusion::Init(ArgList& actionArgs, ActionInit& init, in
     diffCorrl_ = init.DSL().AddSet(DataSet::DOUBLE, MetaData(dsname_, "Corr", ts));
     if (diffConst_ == 0 || diffLabel_ == 0 || diffSlope_ == 0 || diffInter_ == 0 ||
         diffCorrl_ == 0)
+    {
+      mprinterr("Error: Could not allocate 1 or more diffusion constant sets.\n");
       return Action::ERR;
+    }
 #   ifdef MPI
     // No sync needed since these are not time series
     diffConst_->SetNeedsSync( false );
@@ -265,7 +283,7 @@ Action::RetType Action_Diffusion::Setup(ActionSetup& setup) {
   imageOpt_.SetupImaging( setup.CoordInfo().TrajBox().HasBox() );
   if (imageOpt_.ImagingEnabled()) {
     mprintf("\tImaging enabled.\n");
-#   ifdef MPI
+/*#   ifdef MPI
     if (trajComm_.Size() > 1) {
       mprinterr("Error: Imaging for 'diffusion' is not supported in parallel as there is\n"
                 "Error:   no way to correct for imaging that has taken place on preceding\n"
@@ -274,7 +292,7 @@ Action::RetType Action_Diffusion::Setup(ActionSetup& setup) {
                 "Error:   specify the 'noimage' keyword.\n");
       return Action::ERR;
     }
-#   endif
+#   endif*/
   } else {
     mprintf("\tImaging disabled.\n");
     if (avgucell_ != 0)
@@ -285,6 +303,14 @@ Action::RetType Action_Diffusion::Setup(ActionSetup& setup) {
   if (imageOpt_.ImagingEnabled())
     previousFrac_.reserve( mask_.Nselected() );
 
+# ifdef MPI
+  if (initial_.empty() && imageOpt_.ImagingEnabled() && trajComm_.Size() > 1) {
+    mprintf("Warning: Calculating diffusion in parallel with imaging turned on.\n"
+            "Warning:   Mean-squared distance calculation will be averaged starting\n"
+            "Warning:   from %i independent time origins.\n", trajComm_.Size());
+  }
+# endif
+
   // If initial frame already set and current # atoms > # atoms in initial
   // frame this will probably cause an error.
   if (!initial_.empty() && setup.Top().Natom() > initial_.Natom()) {
@@ -295,6 +321,12 @@ Action::RetType Action_Diffusion::Setup(ActionSetup& setup) {
 
   // Set up sets for individual atoms if necessary
   if (printIndividual_) {
+#   ifdef MPI
+    if (imageOpt_.ImagingEnabled() && trajComm_.Size() > 1) {
+      mprinterr("Error: Cannot track individual atom diffusion in parallel with imaging.\n");
+      return Action::ERR;
+    }
+#   endif
     // Create as many spots for sets as needed. All do not have to be used.
     if (mask_.back() >= (int)atom_x_.size()) {
       int newSize = mask_.back() + 1;
@@ -341,6 +373,7 @@ Action::RetType Action_Diffusion::Setup(ActionSetup& setup) {
 
 // Action_Diffusion::DoAction()
 Action::RetType Action_Diffusion::DoAction(int frameNum, ActionFrame& frm) {
+  //rprintf("DEBUG: DIFFUSION FRAME %i\n", frameNum);
   Matrix_3x3 const* ucell = 0;
   if (imageOpt_.ImagingEnabled()) {
     imageOpt_.SetImageType( frm.Frm().BoxCrd().Is_X_Aligned_Ortho() );
@@ -354,7 +387,10 @@ Action::RetType Action_Diffusion::DoAction(int frameNum, ActionFrame& frm) {
   if (initial_.empty()) {
     initial_ = frm.Frm();
 #   ifdef MPI
-    if (trajComm_.Size() > 1) {
+    if (imageOpt_.ImagingEnabled()) {
+      // Indicate data sets will have multiple time origins
+      multipleTimeOrigins_ = true;
+    } else if (trajComm_.Size() > 1) {
       if (trajComm_.Master())
         for (int rank = 1; rank < trajComm_.Size(); ++rank)
           initial_.SendFrame( rank, trajComm_ );
@@ -465,6 +501,89 @@ Action::RetType Action_Diffusion::DoAction(int frameNum, ActionFrame& frm) {
   return Action::OK;
 }
 
+#ifdef MPI
+void Action_Diffusion::average_multiple_time_origins(DataSet* set,
+                                                     std::vector<int> const& nvals_on_rank,
+                                                     int max)
+const
+{
+  // Sanity check - underlying set should be double
+  if (set->Type() != DataSet::DOUBLE) {
+    mprinterr("Internal Error: Action_Diffusion::average_multiple_time_origins(): Set %s is not double.\n", set->legend());
+    return;
+  }
+  DataSet_double const& dset = static_cast<DataSet_double const&>( *set );
+  // Average over all the values
+  if (trajComm_.Master()) {
+    std::vector<double> Sum_all( max, 0.0 );
+    std::vector<double> Ntotal( max, 0.0 );
+    std::vector<double> fromChild( max, 0.0 );
+    // Sum master rank
+    for (unsigned int idx = 0; idx < set->Size(); idx++) {
+      Sum_all[idx] = dset[idx];
+      Ntotal[idx] = 1.0;
+    }
+    for (int rank = 1; rank < trajComm_.Size(); rank++) {
+      // Receive from child
+      trajComm_.Recv( &fromChild[0], nvals_on_rank[rank], MPI_DOUBLE, rank, 10101 ); // FIXME tag
+      // Sum child rank
+      for (int idx = 0; idx < nvals_on_rank[rank]; idx++) {
+        Sum_all[idx] += fromChild[idx];
+        Ntotal[idx] += 1.0;
+      }
+    }
+    //for (int idx = 0; idx < max; idx++) 
+    //  mprintf("DEBUG sum=%g total=%g avg=%g\n", Sum_all[idx], Ntotal[idx], Sum_all[idx] / Ntotal[idx]);
+    // Ensure master set is big enough to hold result
+    DataSet_double& updateSet = static_cast<DataSet_double&>( *set );
+    if (dset.Size() < (unsigned int)max)
+      updateSet.Resize( max );
+    // Update master set
+    for (int idx = 0; idx < max; idx++)
+      updateSet[idx] = Sum_all[idx] / Ntotal[idx];
+  } else {
+    // Send the values to master
+    trajComm_.Send(dset.DvalPtr(), dset.Size(), MPI_DOUBLE, 0, 10101); // FIXME tag
+  }
+  // Set no longer needs to be synced
+  set->SetNeedsSync( false );
+}
+
+/** See if we need to average over multiple time origins. */
+int Action_Diffusion::SyncAction() {
+  if (multipleTimeOrigins_) {
+    mprintf("DEBUG: Multiple time origins.\n");
+    int max = 0;
+    std::vector<int> nvals_on_rank;
+    if (trajComm_.Master()) {
+      // Get how many values each rank has
+      nvals_on_rank.assign( trajComm_.Size(), 0 );
+      int nvals = (int)avg_a_->Size();
+      trajComm_.GatherMaster( &nvals, 1, MPI_INT, &nvals_on_rank[0] );
+      mprintf("DEBUG: %s array sizes:", avg_a_->legend());
+      for (std::vector<int>::const_iterator it = nvals_on_rank.begin(); it != nvals_on_rank.end(); ++it)
+        mprintf(" %i", *it);
+      mprintf("\n");
+      max = nvals;
+      for (unsigned int idx = 1; idx < nvals_on_rank.size(); idx++)
+        if (nvals_on_rank[idx] > max) max = nvals_on_rank[idx];
+      mprintf("DEBUG: Max is %i\n", max);
+
+    } else {
+      // Tell master how many values we have
+      int nvals = (int)avg_a_->Size();
+      trajComm_.GatherMaster( &nvals, 1, MPI_INT, 0);
+    }
+    average_multiple_time_origins( avg_x_, nvals_on_rank, max );
+    average_multiple_time_origins( avg_y_, nvals_on_rank, max );
+    average_multiple_time_origins( avg_z_, nvals_on_rank, max );
+    average_multiple_time_origins( avg_r_, nvals_on_rank, max );
+    average_multiple_time_origins( avg_a_, nvals_on_rank, max );
+  }
+  return 0;
+}
+#endif /* MPI */
+
 // Action_Diffusion::Print()
 void Action_Diffusion::Print() {
   if (!calcDiffConst_) return;
diff --git a/src/Action_Diffusion.h b/src/Action_Diffusion.h
index 09f602cef3..8b62666c5c 100644
--- a/src/Action_Diffusion.h
+++ b/src/Action_Diffusion.h
@@ -12,6 +12,11 @@ class Action_Diffusion : public Action {
     Action::RetType Setup(ActionSetup&);
     Action::RetType DoAction(int, ActionFrame&);
     void Print();
+#   ifdef MPI
+    int SyncAction();
+
+    void average_multiple_time_origins(DataSet*, std::vector<int> const&, int) const;
+#   endif
 
     typedef DataSetList::DataListType Dlist;
     typedef std::vector<Vec3> Varray;
@@ -56,6 +61,7 @@ class Action_Diffusion : public Action {
     Box avgbox_;          ///< Hold average box for removing box fluctuations
 #   ifdef MPI
     Parallel::Comm trajComm_;
+    bool multipleTimeOrigins_; ///< If true, parallel diffusion calc with imaging, multiple time originsi
 #   endif
 };
 #endif
diff --git a/src/Parallel.cpp b/src/Parallel.cpp
index 52f3bf63df..c2c7c692c6 100644
--- a/src/Parallel.cpp
+++ b/src/Parallel.cpp
@@ -321,7 +321,7 @@ int Parallel::Comm::DivideAmongProcesses(int& my_start, int& my_stop, int maxElt
 /** Use MPI_REDUCE to OP the values in sendbuffer and place them in
   * recvbuffer on master.
   */
-int Parallel::Comm::ReduceMaster(void* recvBuffer, void* sendBuffer, int N,
+int Parallel::Comm::ReduceMaster(void* recvBuffer, const void* sendBuffer, int N,
                            MPI_Datatype datatype, MPI_Op op) const
 {
   int err = MPI_Reduce(sendBuffer, recvBuffer, N, datatype, op, 0, comm_);
diff --git a/src/Parallel.h b/src/Parallel.h
index 29498c6d77..d1d017d6f3 100644
--- a/src/Parallel.h
+++ b/src/Parallel.h
@@ -131,7 +131,7 @@ class Parallel::Comm {
     /// my_start, my_stop, maxElts
     int DivideAmongProcesses(int&, int&, int) const;
     /// RecvBuffer, SendBuffer, Count, DataType, Op
-    int ReduceMaster(void*, void*, int, MPI_Datatype, MPI_Op) const;
+    int ReduceMaster(void*, const void*, int, MPI_Datatype, MPI_Op) const;
     /// Rank, RecvBuffer, SendBuffer, Count, DataType, Op
     int Reduce(int, void*, void*, int, MPI_Datatype, MPI_Op) const;
     /// Buffer, Count, Rank, DataType 
diff --git a/src/cpptrajdepend b/src/cpptrajdepend
index db9400b0ff..f09eb41671 100644
--- a/src/cpptrajdepend
+++ b/src/cpptrajdepend
@@ -24,7 +24,7 @@ Action_CreateReservoir.o : Action_CreateReservoir.cpp Action.h ActionState.h Act
 Action_DNAionTracker.o : Action_DNAionTracker.cpp Action.h ActionState.h Action_DNAionTracker.h ArgList.h AssociatedData.h Atom.h AtomMask.h AtomType.h BaseIOtype.h Box.h Constants.h CoordinateInfo.h CpptrajFile.h CpptrajStdio.h DataFile.h DataFileList.h DataSet.h DataSetList.h DataSet_Coords.h DataSet_Coords_REF.h Dimension.h DispatchObject.h DistRoutines.h FileIO.h FileName.h FileTypes.h Frame.h ImageOption.h MaskToken.h Matrix_3x3.h MetaData.h Molecule.h NameType.h Parallel.h ParameterHolders.h ParameterSet.h ParameterTypes.h Range.h ReferenceFrame.h ReplicaDimArray.h Residue.h Segment.h SymbolExporting.h TextFormat.h Timer.h Topology.h TypeNameHolder.h Unit.h Vec3.h
 Action_DSSP.o : Action_DSSP.cpp Action.h ActionState.h Action_DSSP.h ArgList.h AssociatedData.h Atom.h AtomMask.h AtomType.h BaseIOtype.h Box.h CharMask.h Constants.h CoordinateInfo.h CpptrajFile.h CpptrajStdio.h DataFile.h DataFileList.h DataSet.h DataSetList.h DataSet_Coords.h DataSet_Coords_REF.h Dimension.h DispatchObject.h DistRoutines.h FileIO.h FileName.h FileTypes.h Frame.h ImageOption.h MaskToken.h Matrix_3x3.h MetaData.h Molecule.h NameType.h Parallel.h ParameterHolders.h ParameterSet.h ParameterTypes.h Range.h ReferenceFrame.h ReplicaDimArray.h Residue.h Segment.h SymbolExporting.h TextFormat.h Timer.h Topology.h TypeNameHolder.h Unit.h Vec3.h
 Action_Density.o : Action_Density.cpp Action.h ActionState.h Action_Density.h ArgList.h AssociatedData.h Atom.h AtomMask.h AtomType.h BaseIOtype.h Box.h Constants.h CoordinateInfo.h CpptrajFile.h CpptrajStdio.h DataFile.h DataFileList.h DataSet.h DataSetList.h DataSet_1D.h DataSet_Coords.h DataSet_Coords_REF.h Dimension.h DispatchObject.h FileIO.h FileName.h FileTypes.h Frame.h MaskToken.h Matrix_3x3.h MetaData.h Molecule.h NameType.h OnlineVarT.h Parallel.h ParameterHolders.h ParameterSet.h ParameterTypes.h Range.h ReferenceFrame.h ReplicaDimArray.h Residue.h Segment.h StringRoutines.h SymbolExporting.h TextFormat.h Timer.h Topology.h TypeNameHolder.h Unit.h Vec3.h
-Action_Diffusion.o : Action_Diffusion.cpp Action.h ActionState.h Action_Diffusion.h ArgList.h AssociatedData.h Atom.h AtomMask.h AtomType.h BaseIOtype.h Box.h Constants.h CoordinateInfo.h CpptrajFile.h CpptrajStdio.h DataFile.h DataFileList.h DataSet.h DataSetList.h DataSet_1D.h DataSet_Coords.h DataSet_Coords_REF.h DataSet_Mat3x3.h Dimension.h DispatchObject.h FileIO.h FileName.h FileTypes.h Frame.h ImageOption.h MaskToken.h Matrix_3x3.h MetaData.h Molecule.h NameType.h Parallel.h ParameterHolders.h ParameterSet.h ParameterTypes.h Range.h ReferenceFrame.h ReplicaDimArray.h Residue.h Segment.h StringRoutines.h SymbolExporting.h TextFormat.h Timer.h Topology.h TypeNameHolder.h Unit.h Vec3.h
+Action_Diffusion.o : Action_Diffusion.cpp Action.h ActionState.h Action_Diffusion.h ArgList.h AssociatedData.h Atom.h AtomMask.h AtomType.h BaseIOtype.h Box.h Constants.h CoordinateInfo.h CpptrajFile.h CpptrajStdio.h DataFile.h DataFileList.h DataSet.h DataSetList.h DataSet_1D.h DataSet_Coords.h DataSet_Coords_REF.h DataSet_Mat3x3.h DataSet_double.h Dimension.h DispatchObject.h FileIO.h FileName.h FileTypes.h Frame.h ImageOption.h MaskToken.h Matrix_3x3.h MetaData.h Molecule.h NameType.h Parallel.h ParameterHolders.h ParameterSet.h ParameterTypes.h Range.h ReferenceFrame.h ReplicaDimArray.h Residue.h Segment.h StringRoutines.h SymbolExporting.h TextFormat.h Timer.h Topology.h TypeNameHolder.h Unit.h Vec3.h
 Action_Dihedral.o : Action_Dihedral.cpp Action.h ActionState.h Action_Dihedral.h ArgList.h AssociatedData.h Atom.h AtomMask.h AtomType.h BaseIOtype.h Box.h Constants.h CoordinateInfo.h CpptrajFile.h CpptrajStdio.h DataFile.h DataFileList.h DataSet.h DataSetList.h DataSet_1D.h DataSet_Coords.h DataSet_Coords_REF.h DataSet_double.h Dimension.h DispatchObject.h FileIO.h FileName.h FileTypes.h Frame.h MaskToken.h Matrix_3x3.h MetaData.h Molecule.h NameType.h Parallel.h ParameterHolders.h ParameterSet.h ParameterTypes.h Range.h ReferenceFrame.h ReplicaDimArray.h Residue.h Segment.h SymbolExporting.h TextFormat.h Timer.h Topology.h TorsionRoutines.h TypeNameHolder.h Unit.h Vec3.h
 Action_DihedralRMS.o : Action_DihedralRMS.cpp Action.h ActionState.h Action_DihedralRMS.h ArgList.h AssociatedData.h Atom.h AtomMask.h AtomType.h BaseIOtype.h Box.h Constants.h CoordinateInfo.h CpptrajFile.h CpptrajStdio.h DataFile.h DataFileList.h DataSet.h DataSetList.h DataSet_Coords.h DataSet_Coords_REF.h DihedralSearch.h Dimension.h DispatchObject.h FileIO.h FileName.h FileTypes.h Frame.h MaskToken.h Matrix_3x3.h MetaData.h Molecule.h NameType.h Parallel.h ParameterHolders.h ParameterSet.h ParameterTypes.h Range.h ReferenceAction.h ReferenceFrame.h ReplicaDimArray.h Residue.h Segment.h SymbolExporting.h TextFormat.h Timer.h Topology.h TorsionRoutines.h TypeNameHolder.h Unit.h Vec3.h
 Action_Dipole.o : Action_Dipole.cpp Action.h ActionState.h Action_Dipole.h ArgList.h ArrayIterator.h AssociatedData.h Atom.h AtomMask.h AtomType.h BaseIOtype.h Box.h CharMask.h Constants.h CoordinateInfo.h CpptrajFile.h CpptrajStdio.h DataFile.h DataFileList.h DataSet.h DataSetList.h DataSet_3D.h DataSet_Coords.h DataSet_Coords_REF.h DataSet_GridFlt.h Dimension.h DispatchObject.h FileIO.h FileName.h FileTypes.h Frame.h Grid.h GridAction.h GridBin.h MaskToken.h Matrix_3x3.h MetaData.h Molecule.h NameType.h Parallel.h ParameterHolders.h ParameterSet.h ParameterTypes.h Range.h ReferenceFrame.h ReplicaDimArray.h Residue.h Segment.h SymbolExporting.h TextFormat.h Timer.h Topology.h TypeNameHolder.h Unit.h Vec3.h

From 1a6fad09f55b6607838f3600c9613f18d3477b49 Mon Sep 17 00:00:00 2001
From: "Daniel R. Roe" <daniel.roe@nih.gov>
Date: Wed, 17 May 2023 13:08:48 -0400
Subject: [PATCH 08/15] Add 3x3 matrix array data netcdf write

---
 src/DataIO_NetCDF.cpp | 42 +++++++++++++++++++++++++++++++++++++++++-
 src/cpptrajdepend     |  2 +-
 2 files changed, 42 insertions(+), 2 deletions(-)

diff --git a/src/DataIO_NetCDF.cpp b/src/DataIO_NetCDF.cpp
index fac7827d6b..df0e05aa5a 100644
--- a/src/DataIO_NetCDF.cpp
+++ b/src/DataIO_NetCDF.cpp
@@ -18,6 +18,7 @@
 # include "DataSet_Vector_Scalar.h"
 # include "DataSet_PairwiseCache_MEM.h"
 # include "DataSet_unsignedInt.h"
+# include "DataSet_Mat3x3.h"
 #endif /* BINTRAJ */
 
 /// CONSTRUCTOR
@@ -1421,6 +1422,7 @@ int DataIO_NetCDF::writeData_1D(DataSet const* ds, Dimension const& dim, SetArra
       case DataSet::VECTOR_SCALAR :
       case DataSet::VECTOR  :
       case DataSet::DOUBLE  :
+      case DataSet::MAT3X3  :
       case DataSet::XYMESH  :          dtype = NC_DOUBLE ; break;
       case DataSet::PH      :
       case DataSet::UNSIGNED_INTEGER : // NOTE; NC_UINT is invalid for NetCDF classic (v3)
@@ -1474,6 +1476,26 @@ int DataIO_NetCDF::writeData_1D(DataSet const* ds, Dimension const& dim, SetArra
         NcVar originVar = defineVar( DimIds, dtype, vecSet.Meta().PrintName(), "origins", vecVar.VID());
         set_vars.push_back( originVar );
       }
+    } else if (it->DS()->Type() == DataSet::MAT3X3) {
+      // ----- Array of 3x3 Matrices ---
+      DataSet_Mat3x3 const& matSet = static_cast<DataSet_Mat3x3 const&>( *(it->DS()) );
+      // Define vecxyz dimension (unit cell vectors)
+      dimIdx = defineDim("vecxyz", 3, matSet.Meta().Legend() + " vecxyz" );
+      if (dimIdx < 0) return 1;
+      NcDim vecxyzDim = Dimensions_[dimIdx];
+      // Define xyz dimension (for each unit cell vector)
+      dimIdx = defineDim("xyz", 3, matSet.Meta().Legend() + " xyz" );
+      if (dimIdx < 0) return 1;
+      NcDim xyzDim = Dimensions_[dimIdx];
+      // Define mat3x3 variable
+      std::vector<int> DimIds(3);
+      DimIds[0] = lengthDim.DID();
+      DimIds[1] = vecxyzDim.DID();
+      DimIds[2] = xyzDim.DID();
+      NcVar mat3x3Var = defineVar( DimIds, dtype, matSet.Meta().PrintName(), "mat3x3", -1);
+      if (mat3x3Var.Empty()) return 1;
+      if (AddDataSetInfo( it->DS(), ncid_, mat3x3Var.VID() )) return 1;
+      set_vars.push_back( mat3x3Var );
     } else if (it->DS()->Type() == DataSet::XYMESH) {
       // ----- XY mesh set -------------
       // Define the Y variable
@@ -1558,6 +1580,23 @@ int DataIO_NetCDF::writeData_1D(DataSet const* ds, Dimension const& dim, SetArra
           }
         }
       }
+    } else if (dset->DS()->Type() == DataSet::MAT3X3) {
+      // ----- Array of 3x3 matrices ---
+      VarArray const& set_vars = *it;
+      DataSet_Mat3x3 const& matSet = static_cast<DataSet_Mat3x3 const&>( *(dset->DS()) );
+      size_t start[3], count[3];
+      start[1] = 0;
+      start[2] = 0;
+      count[0] = 1;
+      count[1] = 3;
+      count[2] = 3;
+      for (unsigned int idx = 0; idx < matSet.Size(); idx++) {
+        start[0] = idx;
+        if (NC::CheckErr(nc_put_vara(ncid_, set_vars[0].VID(), start, count, matSet[idx].Dptr()))) {
+          mprinterr("Error: Coult not write 3x3 matrix %u from '%s'\n", idx, matSet.legend());
+          return 1;
+        }
+      }
     } else if (dset->DS()->Type() == DataSet::XYMESH) {
       // ----- XY mesh set -------------
       VarArray const& set_vars = *it;
@@ -1912,7 +1951,8 @@ int DataIO_NetCDF::writeData_modes(DataSet const* ds) {
 static inline bool isValid1dSetToWrite(DataSet const* ds) {
   return (ds->Group() == DataSet::SCALAR_1D ||
           ds->Type() == DataSet::STRING ||
-          ds->Type() == DataSet::VECTOR
+          ds->Type() == DataSet::VECTOR ||
+          ds->Type() == DataSet::MAT3X3
          );
 }
 
diff --git a/src/cpptrajdepend b/src/cpptrajdepend
index f09eb41671..295eed367d 100644
--- a/src/cpptrajdepend
+++ b/src/cpptrajdepend
@@ -213,7 +213,7 @@ DataIO_Evecs.o : DataIO_Evecs.cpp ArgList.h ArrayIterator.h AssociatedData.h Ato
 DataIO_Gnuplot.o : DataIO_Gnuplot.cpp ArgList.h Array1D.h AssociatedData.h Atom.h AtomMask.h AtomType.h BaseIOtype.h Box.h BufferedLine.h Constants.h CoordinateInfo.h CpptrajFile.h CpptrajStdio.h DataIO.h DataIO_Gnuplot.h DataSet.h DataSetList.h DataSet_1D.h DataSet_2D.h DataSet_Coords.h DataSet_Coords_REF.h Dimension.h FileIO.h FileName.h Frame.h MaskToken.h Matrix_3x3.h MetaData.h Molecule.h NameType.h Parallel.h ParameterHolders.h ParameterSet.h ParameterTypes.h Range.h ReferenceFrame.h ReplicaDimArray.h Residue.h Segment.h SymbolExporting.h TextFormat.h Timer.h Topology.h TypeNameHolder.h Unit.h Vec3.h
 DataIO_Grace.o : DataIO_Grace.cpp ArgList.h AssociatedData.h Atom.h AtomMask.h AtomType.h BaseIOtype.h Box.h BufferedLine.h Constants.h CoordinateInfo.h CpptrajFile.h CpptrajStdio.h DataIO.h DataIO_Grace.h DataSet.h DataSetList.h DataSet_1D.h DataSet_Coords.h DataSet_Coords_REF.h DataSet_double.h DataSet_string.h Dimension.h FileIO.h FileName.h Frame.h MaskToken.h Matrix_3x3.h MetaData.h Molecule.h NameType.h Parallel.h ParameterHolders.h ParameterSet.h ParameterTypes.h Range.h ReferenceFrame.h ReplicaDimArray.h Residue.h Segment.h SymbolExporting.h TextFormat.h Timer.h Topology.h TypeNameHolder.h Unit.h Vec3.h
 DataIO_Mdout.o : DataIO_Mdout.cpp ArgList.h AssociatedData.h Atom.h AtomMask.h AtomType.h BaseIOtype.h Box.h BufferedLine.h Constants.h CoordinateInfo.h CpptrajFile.h CpptrajStdio.h DataIO.h DataIO_Mdout.h DataSet.h DataSetList.h DataSet_1D.h DataSet_Coords.h DataSet_Coords_REF.h DataSet_double.h Dimension.h FileIO.h FileName.h Frame.h MaskToken.h Matrix_3x3.h MetaData.h Molecule.h NameType.h Parallel.h ParameterHolders.h ParameterSet.h ParameterTypes.h Range.h ReferenceFrame.h ReplicaDimArray.h Residue.h Segment.h StringRoutines.h SymbolExporting.h TextFormat.h Timer.h Topology.h TypeNameHolder.h Unit.h Vec3.h
-DataIO_NetCDF.o : DataIO_NetCDF.cpp ArgList.h ArrayIterator.h AssociatedData.h Atom.h AtomMask.h AtomType.h BaseIOtype.h Box.h Cluster/Cframes.h ComplexArray.h Constants.h CoordinateInfo.h CpptrajFile.h CpptrajStdio.h DataIO.h DataIO_NetCDF.h DataSet.h DataSetList.h DataSet_1D.h DataSet_2D.h DataSet_3D.h DataSet_Coords.h DataSet_Coords_REF.h DataSet_MatrixDbl.h DataSet_Mesh.h DataSet_Modes.h DataSet_PairwiseCache.h DataSet_PairwiseCache_MEM.h DataSet_Vector.h DataSet_Vector_Scalar.h DataSet_string.h DataSet_unsignedInt.h Dimension.h FileIO.h FileName.h Frame.h GridBin.h MaskToken.h Matrix.h Matrix_3x3.h MetaData.h Molecule.h NC_Routines.h NameType.h Parallel.h ParameterHolders.h ParameterSet.h ParameterTypes.h Range.h ReferenceFrame.h ReplicaDimArray.h Residue.h Segment.h Spline.h StringRoutines.h SymbolExporting.h TextFormat.h Timer.h Topology.h TypeNameHolder.h Unit.h Vec3.h Version.h
+DataIO_NetCDF.o : DataIO_NetCDF.cpp ArgList.h ArrayIterator.h AssociatedData.h Atom.h AtomMask.h AtomType.h BaseIOtype.h Box.h Cluster/Cframes.h ComplexArray.h Constants.h CoordinateInfo.h CpptrajFile.h CpptrajStdio.h DataIO.h DataIO_NetCDF.h DataSet.h DataSetList.h DataSet_1D.h DataSet_2D.h DataSet_3D.h DataSet_Coords.h DataSet_Coords_REF.h DataSet_Mat3x3.h DataSet_MatrixDbl.h DataSet_Mesh.h DataSet_Modes.h DataSet_PairwiseCache.h DataSet_PairwiseCache_MEM.h DataSet_Vector.h DataSet_Vector_Scalar.h DataSet_string.h DataSet_unsignedInt.h Dimension.h FileIO.h FileName.h Frame.h GridBin.h MaskToken.h Matrix.h Matrix_3x3.h MetaData.h Molecule.h NC_Routines.h NameType.h Parallel.h ParameterHolders.h ParameterSet.h ParameterTypes.h Range.h ReferenceFrame.h ReplicaDimArray.h Residue.h Segment.h Spline.h StringRoutines.h SymbolExporting.h TextFormat.h Timer.h Topology.h TypeNameHolder.h Unit.h Vec3.h Version.h
 DataIO_OpenDx.o : DataIO_OpenDx.cpp ArgList.h ArrayIterator.h AssociatedData.h Atom.h AtomMask.h AtomType.h BaseIOtype.h Box.h BufferedLine.h Constants.h CoordinateInfo.h CpptrajFile.h CpptrajStdio.h DataIO.h DataIO_OpenDx.h DataSet.h DataSetList.h DataSet_3D.h DataSet_Coords.h DataSet_Coords_REF.h DataSet_GridDbl.h DataSet_GridFlt.h Dimension.h FileIO.h FileName.h Frame.h Grid.h GridBin.h MaskToken.h Matrix_3x3.h MetaData.h Molecule.h NameType.h Parallel.h ParameterHolders.h ParameterSet.h ParameterTypes.h ProgressBar.h Range.h ReferenceFrame.h ReplicaDimArray.h Residue.h Segment.h SymbolExporting.h TextFormat.h Timer.h Topology.h TypeNameHolder.h Unit.h Vec3.h
 DataIO_Peaks.o : DataIO_Peaks.cpp ArgList.h AssociatedData.h Atom.h AtomMask.h AtomType.h BaseIOtype.h Box.h Constants.h CoordinateInfo.h CpptrajFile.h CpptrajStdio.h DataIO.h DataIO_Peaks.h DataSet.h DataSetList.h DataSet_Coords.h DataSet_Coords_REF.h DataSet_Vector_Scalar.h Dimension.h FileIO.h FileName.h Frame.h MaskToken.h Matrix_3x3.h MetaData.h Molecule.h NameType.h Parallel.h ParameterHolders.h ParameterSet.h ParameterTypes.h Range.h ReferenceFrame.h ReplicaDimArray.h Residue.h Segment.h StringRoutines.h SymbolExporting.h TextFormat.h Timer.h Topology.h TypeNameHolder.h Unit.h Vec3.h
 DataIO_RemLog.o : DataIO_RemLog.cpp ArgList.h AssociatedData.h Atom.h AtomMask.h AtomType.h BaseIOtype.h Box.h BufferedLine.h Constants.h CoordinateInfo.h CpptrajFile.h CpptrajStdio.h DataIO.h DataIO_RemLog.h DataSet.h DataSetList.h DataSet_Coords.h DataSet_Coords_REF.h DataSet_RemLog.h Dimension.h FileIO.h FileName.h Frame.h MaskToken.h Matrix_3x3.h MetaData.h Molecule.h NameType.h Parallel.h ParameterHolders.h ParameterSet.h ParameterTypes.h ProgressBar.h Range.h ReferenceFrame.h ReplicaDimArray.h Residue.h Segment.h StringRoutines.h SymbolExporting.h TextFormat.h Timer.h Topology.h TypeNameHolder.h Unit.h Vec3.h

From 8f8ccb0829dbc4ab3d49581deb070e39b2c20be1 Mon Sep 17 00:00:00 2001
From: "Daniel R. Roe" <daniel.roe@nih.gov>
Date: Wed, 17 May 2023 13:20:41 -0400
Subject: [PATCH 09/15] Add mat3x3 netcdf read

---
 src/DataIO_NetCDF.cpp | 31 +++++++++++++++++++++++++++++++
 src/DataIO_NetCDF.h   |  2 ++
 src/DataSet_Mat3x3.h  |  2 ++
 3 files changed, 35 insertions(+)

diff --git a/src/DataIO_NetCDF.cpp b/src/DataIO_NetCDF.cpp
index df0e05aa5a..d20e982812 100644
--- a/src/DataIO_NetCDF.cpp
+++ b/src/DataIO_NetCDF.cpp
@@ -29,6 +29,7 @@ DataIO_NetCDF::DataIO_NetCDF() :
 {
   SetValid( DataSet::MODES );
   SetValid( DataSet::VECTOR_SCALAR );
+  SetValid( DataSet::MAT3X3 );
   SetValid( DataSet::PMATRIX_MEM );
 }
 #else
@@ -425,6 +426,33 @@ int DataIO_NetCDF::readData_1D_unsignedInt(DataSet* ds, NcVar const& yVar, VarAr
   return 0;
 }
 
+/** Read 3x3 matrix array set with CPPTRAJ conventions. */
+int DataIO_NetCDF::readData_Mat3x3(DataSet* ds, NcVar const& yVar, VarArray& Vars) const {
+  size_t start[3];
+  start[1] = 0;
+  start[2] = 0;
+  size_t count[3];
+  count[0] = 1;
+  count[1] = 3;
+  count[2] = 3;
+
+  unsigned int nmats = dimLen(yVar.DimId(0));
+
+  DataSet_Mat3x3& matSet = static_cast<DataSet_Mat3x3&>( *ds );
+  // Get 3x3 matrices
+  matSet.Resize( nmats );
+  for (unsigned int idx = 0; idx < nmats; idx++) {
+    start[0] = idx;
+    if (NC::CheckErr(nc_get_vara(ncid_, yVar.VID(), start, count, matSet[idx].Dptr()))) {
+      mprinterr("Error: Could not get 3x3 matrix values for %u '%s'\n", idx, matSet.legend());
+      return 1;
+    }
+  }
+  Vars[yVar.VID()].MarkRead();
+
+  return 0;
+}
+
 /** Read vector set with CPPTRAJ conventions. */
 int DataIO_NetCDF::readData_1D_vector(DataSet* ds, NcVar const& yVar, VarArray& Vars) const {
   size_t start[2];
@@ -957,6 +985,9 @@ const
     } else if (dtype == DataSet::VECTOR_SCALAR) {
       if (readData_1D_vector_scalar(ds, *var, Vars))
         return 1;
+    } else if (dtype == DataSet::MAT3X3) {
+      if (readData_Mat3x3(ds, *var, Vars))
+        return 1;
     } else if (dtype == DataSet::STRING) {
       if (readData_1D_string(ds, *var, Vars))
         return 1;
diff --git a/src/DataIO_NetCDF.h b/src/DataIO_NetCDF.h
index 59dc3001a1..dc1c6356c4 100644
--- a/src/DataIO_NetCDF.h
+++ b/src/DataIO_NetCDF.h
@@ -41,6 +41,8 @@ class DataIO_NetCDF : public DataIO {
 
     int readData_1D_unsignedInt(DataSet*, NcVar const&, VarArray&) const;
 
+    int readData_Mat3x3(DataSet*, NcVar const&, VarArray&) const;
+
     int readData_1D_vector(DataSet*, NcVar const&, VarArray&) const;
 
     int readData_1D_vector_scalar(DataSet*, NcVar const&, VarArray&) const;
diff --git a/src/DataSet_Mat3x3.h b/src/DataSet_Mat3x3.h
index 9f810f3388..063c7b9195 100644
--- a/src/DataSet_Mat3x3.h
+++ b/src/DataSet_Mat3x3.h
@@ -20,6 +20,7 @@ class DataSet_Mat3x3 : public DataSet {
     int Append(DataSet*);
     size_t MemUsageInBytes() const { return (9 * data_.size() * sizeof(double)); }
     // -------------------------------------------
+    void Resize(size_t s)               { data_.resize( s );    }
     void AddMat3x3( Matrix_3x3 const& m) { data_.push_back( m ); }
     typedef Marray::const_iterator const_iterator;
     const_iterator begin() const { return data_.begin(); }
@@ -28,6 +29,7 @@ class DataSet_Mat3x3 : public DataSet {
     iterator begin() { return data_.begin(); }
     iterator end()   { return data_.end();   }
     Matrix_3x3 const& operator[](unsigned int i) const { return data_[i]; }
+    Matrix_3x3& operator[](unsigned int i) { return data_[i]; }
   private:
     Marray data_;
 };

From de4233677ca4e6c4338facdd83fe08fe1f3570b2 Mon Sep 17 00:00:00 2001
From: "Daniel R. Roe" <daniel.roe@nih.gov>
Date: Thu, 18 May 2023 11:58:57 -0400
Subject: [PATCH 10/15] Clean up output

---
 src/Action_Diffusion.cpp | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/src/Action_Diffusion.cpp b/src/Action_Diffusion.cpp
index b862088804..e3db7c9cb3 100644
--- a/src/Action_Diffusion.cpp
+++ b/src/Action_Diffusion.cpp
@@ -552,7 +552,7 @@ const
 /** See if we need to average over multiple time origins. */
 int Action_Diffusion::SyncAction() {
   if (multipleTimeOrigins_) {
-    mprintf("DEBUG: Multiple time origins.\n");
+    mprintf("    DIFFUSION: Calculating diffusion by averaging over %i time origins.\n", trajComm_.Size());
     int max = 0;
     std::vector<int> nvals_on_rank;
     if (trajComm_.Master()) {
@@ -560,14 +560,14 @@ int Action_Diffusion::SyncAction() {
       nvals_on_rank.assign( trajComm_.Size(), 0 );
       int nvals = (int)avg_a_->Size();
       trajComm_.GatherMaster( &nvals, 1, MPI_INT, &nvals_on_rank[0] );
-      mprintf("DEBUG: %s array sizes:", avg_a_->legend());
-      for (std::vector<int>::const_iterator it = nvals_on_rank.begin(); it != nvals_on_rank.end(); ++it)
-        mprintf(" %i", *it);
-      mprintf("\n");
+      //mprintf("DEBUG: %s array sizes:", avg_a_->legend());
+      //for (std::vector<int>::const_iterator it = nvals_on_rank.begin(); it != nvals_on_rank.end(); ++it)
+      //  mprintf(" %i", *it);
+      //mprintf("\n");
       max = nvals;
       for (unsigned int idx = 1; idx < nvals_on_rank.size(); idx++)
         if (nvals_on_rank[idx] > max) max = nvals_on_rank[idx];
-      mprintf("DEBUG: Max is %i\n", max);
+      mprintf("\tMax time is %g ps.\n", (double)(max-1) * time_);
 
     } else {
       // Tell master how many values we have

From b68c9530d323062a3b3f02282930c1a0408c2bd3 Mon Sep 17 00:00:00 2001
From: "Daniel R. Roe" <daniel.roe@nih.gov>
Date: Thu, 18 May 2023 12:07:08 -0400
Subject: [PATCH 11/15] Add MPI imaged diffusion test

---
 .../2procs.tz2.ortho.wato.dat.save            |  6 +++++
 test/Test_Diffusion/RunTest.sh                | 25 +++++++++++++++----
 2 files changed, 26 insertions(+), 5 deletions(-)
 create mode 100644 test/Test_Diffusion/2procs.tz2.ortho.wato.dat.save

diff --git a/test/Test_Diffusion/2procs.tz2.ortho.wato.dat.save b/test/Test_Diffusion/2procs.tz2.ortho.wato.dat.save
new file mode 100644
index 0000000000..8fc514c07a
--- /dev/null
+++ b/test/Test_Diffusion/2procs.tz2.ortho.wato.dat.save
@@ -0,0 +1,6 @@
+#Time        Water[R]     Water[X]     Water[Y]     Water[Z]     Water[A]
+       0       0.0000       0.0000       0.0000       0.0000       0.0000
+       1       2.2328       0.7769       0.7325       0.7233       1.4942
+       2       4.2339       1.4700       1.3595       1.4044       2.0576
+       3       6.0225       2.1225       1.8734       2.0265       2.4540
+       4       7.8191       2.7637       2.5268       2.5286       2.7962
diff --git a/test/Test_Diffusion/RunTest.sh b/test/Test_Diffusion/RunTest.sh
index 6d79a0ee74..da584e2b7c 100755
--- a/test/Test_Diffusion/RunTest.sh
+++ b/test/Test_Diffusion/RunTest.sh
@@ -49,8 +49,11 @@ EOF
 }
 
 Test_diffusion_avgucell() {
-  TOP=../tz2.ortho.parm7
-  cat > $INPUT <<EOF
+  UNITNAME='Diffusion with average box correction test'
+  CheckFor maxthreads 2
+  if [ $? -eq 0 ] ; then
+    TOP=../tz2.ortho.parm7
+    cat > $INPUT <<EOF
 trajin ../tz2.ortho.nc
 avgbox MyBox
 run
@@ -58,8 +61,20 @@ run
 diffusion Water :WAT@O out tz2.ortho.wato.dat avgucell MyBox[avg]
 run
 EOF
-  RunCpptraj "Diffusion with average box correction test"
-  DoTest ../Test_Unwrap/tz2.ortho.wato.dat.save tz2.ortho.wato.dat
+    RunCpptraj "$UNITNAME"
+    if [ -z "$N_THREADS" ] ; then
+      testsave=../Test_Unwrap/tz2.ortho.wato.dat.save
+    elif [ $N_THREADS -eq 1 ] ; then
+      testsave=../Test_Unwrap/tz2.ortho.wato.dat.save
+    elif [ $N_THREADS -eq 2 ] ; then
+      testsave=2procs.tz2.ortho.wato.dat.save
+    else
+      testsave=''
+    fi
+    if [ ! -z "$testsave" ] ; then
+      DoTest $testsave tz2.ortho.wato.dat
+    fi
+  fi
 }
 
 Test_diffusion_noImage() {
@@ -109,13 +124,13 @@ EOF
 }
 
 Test_diffusion_noImage
+Test_diffusion_avgucell
 UNITNAME='Imaged diffusion tests'
 CheckFor maxthreads 1
 if [ $? -eq 0 ] ; then
   Test_diffusion_oldSyntax
   Test_diffusion_newSyntax
   Test_diffusion_nonOrtho
-  Test_diffusion_avgucell
 fi
 UNITNAME='STFC diffusion tests'
 CheckFor notparallel

From 60ef10c7e8fb8a63a542ab455810098bf00663ed Mon Sep 17 00:00:00 2001
From: "Daniel R. Roe" <daniel.roe@nih.gov>
Date: Thu, 18 May 2023 12:25:49 -0400
Subject: [PATCH 12/15] Add allowmultipleorigins keyword to allow calculating
 diffusion by averaging over multiple time origins in parallel. Only want this
 done if user specifies since it is not as accurate as using the full
 trajectory.

---
 src/Action_Diffusion.cpp       | 46 +++++++++++++++++++++++++---------
 src/Action_Diffusion.h         |  1 +
 test/Test_Diffusion/RunTest.sh |  2 +-
 3 files changed, 36 insertions(+), 13 deletions(-)

diff --git a/src/Action_Diffusion.cpp b/src/Action_Diffusion.cpp
index e3db7c9cb3..6e4216eef0 100644
--- a/src/Action_Diffusion.cpp
+++ b/src/Action_Diffusion.cpp
@@ -29,7 +29,8 @@ Action_Diffusion::Action_Diffusion() :
   outputx_(0), outputy_(0), outputz_(0), outputr_(0), outputa_(0),
   diffout_(0),
   masterDSL_(0),
-  avgucell_(0)
+  avgucell_(0),
+  allowMultipleTimeOrigins_(false)
 # ifdef MPI
   ,multipleTimeOrigins_(false)
 # endif
@@ -39,6 +40,9 @@ static inline void ShortHelp() {
   mprintf("\t[{out <filename> | separateout <suffix>}] [time <time per frame>] [noimage]\n"
           "\t[<mask>] [<set name>] [individual] [diffout <filename>] [nocalc]\n"
           "\t[avgucell <avg ucell set>]\n");
+#ifdef MPI
+  mprintf("\t[allowmultipleorigins]\n");
+#endif
 }
 
 void Action_Diffusion::Help() const {
@@ -120,6 +124,7 @@ Action::RetType Action_Diffusion::Init(ArgList& actionArgs, ActionInit& init, in
     }
   }
   if (diffout_ != 0) calcDiffConst_ = true;
+  allowMultipleTimeOrigins_ = actionArgs.hasKey("allowmultipleorigins");
   // Get average box
   std::string avgucellstr = actionArgs.GetStringKey("avgucell");
   if (!avgucellstr.empty()) {
@@ -253,6 +258,17 @@ Action::RetType Action_Diffusion::Init(ArgList& actionArgs, ActionInit& init, in
     mprintf("\tUsing average unit cell vectors from set '%s' to remove box fluctuations.\n", avgucell_->legend());
     avgbox_.PrintInfo();
   }
+  if (allowMultipleTimeOrigins_) {
+# ifdef MPI
+    if (imageOpt_.UseImage()) {
+      if (trajComm_.Size() > 1)
+        mprintf("\tTrajectories that require unit cell imaging will be averaged over %i time origins.\n", trajComm_.Size());
+    } else
+      mprintf("\tImaging is disabled, ignoring 'allowmultipleorigins'\n");
+# else
+    mprintf("\tThe 'allowmultipleorigins' keyword is only relevant in parallel, ignoring.\n");
+# endif /* MPI */
+  }
 # ifdef _OPENMP
 # pragma omp parallel
   {
@@ -283,16 +299,23 @@ Action::RetType Action_Diffusion::Setup(ActionSetup& setup) {
   imageOpt_.SetupImaging( setup.CoordInfo().TrajBox().HasBox() );
   if (imageOpt_.ImagingEnabled()) {
     mprintf("\tImaging enabled.\n");
-/*#   ifdef MPI
+#   ifdef MPI
     if (trajComm_.Size() > 1) {
-      mprinterr("Error: Imaging for 'diffusion' is not supported in parallel as there is\n"
-                "Error:   no way to correct for imaging that has taken place on preceding\n"
-                "Error:   MPI ranks. To use 'diffusion' in parallel, the trajectory should\n"
-                "Error:   be unwrapped. If this trajectory has already been unwrapped please\n"
-                "Error:   specify the 'noimage' keyword.\n");
-      return Action::ERR;
+      if (!allowMultipleTimeOrigins_) {
+        mprinterr("Error: Imaging for 'diffusion' is not supported in parallel as there is\n"
+                  "Error:   no way to correct for imaging that has taken place on preceding\n"
+                  "Error:   MPI ranks. To use 'diffusion' in parallel, the trajectory should\n"
+                  "Error:   be unwrapped. If this trajectory has already been unwrapped please\n"
+                  "Error:   specify the 'noimage' keyword.\n");
+        mprinterr("Error: To calculate diffusion in parallel by dividing the trajectory among\n"
+                  "Error:   multiple time origins, specify the 'allowmultipleorigins' keyword.\n");
+        return Action::ERR;
+      } else {
+        // Indicate data sets will have multiple time origins
+        multipleTimeOrigins_ = true;
+      }
     }
-#   endif*/
+#   endif
   } else {
     mprintf("\tImaging disabled.\n");
     if (avgucell_ != 0)
@@ -304,7 +327,7 @@ Action::RetType Action_Diffusion::Setup(ActionSetup& setup) {
     previousFrac_.reserve( mask_.Nselected() );
 
 # ifdef MPI
-  if (initial_.empty() && imageOpt_.ImagingEnabled() && trajComm_.Size() > 1) {
+  if (multipleTimeOrigins_ && initial_.empty()) {
     mprintf("Warning: Calculating diffusion in parallel with imaging turned on.\n"
             "Warning:   Mean-squared distance calculation will be averaged starting\n"
             "Warning:   from %i independent time origins.\n", trajComm_.Size());
@@ -388,8 +411,7 @@ Action::RetType Action_Diffusion::DoAction(int frameNum, ActionFrame& frm) {
     initial_ = frm.Frm();
 #   ifdef MPI
     if (imageOpt_.ImagingEnabled()) {
-      // Indicate data sets will have multiple time origins
-      multipleTimeOrigins_ = true;
+
     } else if (trajComm_.Size() > 1) {
       if (trajComm_.Master())
         for (int rank = 1; rank < trajComm_.Size(); ++rank)
diff --git a/src/Action_Diffusion.h b/src/Action_Diffusion.h
index 8b62666c5c..be83e3dab2 100644
--- a/src/Action_Diffusion.h
+++ b/src/Action_Diffusion.h
@@ -59,6 +59,7 @@ class Action_Diffusion : public Action {
     Dimension Xdim_;
     DataSet* avgucell_;   ///< Hold average unit cell parameters for removing box fluctuations
     Box avgbox_;          ///< Hold average box for removing box fluctuations
+    bool allowMultipleTimeOrigins_; ///< If true, allow multiple time origins with imaging
 #   ifdef MPI
     Parallel::Comm trajComm_;
     bool multipleTimeOrigins_; ///< If true, parallel diffusion calc with imaging, multiple time originsi
diff --git a/test/Test_Diffusion/RunTest.sh b/test/Test_Diffusion/RunTest.sh
index da584e2b7c..a5303cbd95 100755
--- a/test/Test_Diffusion/RunTest.sh
+++ b/test/Test_Diffusion/RunTest.sh
@@ -58,7 +58,7 @@ trajin ../tz2.ortho.nc
 avgbox MyBox
 run
 
-diffusion Water :WAT@O out tz2.ortho.wato.dat avgucell MyBox[avg]
+diffusion Water :WAT@O out tz2.ortho.wato.dat avgucell MyBox[avg] allowmultipleorigins
 run
 EOF
     RunCpptraj "$UNITNAME"

From 0be41df4d0a2ae56fda39ec7dad09d3c5912ddce Mon Sep 17 00:00:00 2001
From: "Daniel R. Roe" <daniel.roe@nih.gov>
Date: Thu, 18 May 2023 13:40:05 -0400
Subject: [PATCH 13/15] Improve the check for parallel 'individual'

---
 src/Action_Diffusion.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/Action_Diffusion.cpp b/src/Action_Diffusion.cpp
index 6e4216eef0..baf6d394c2 100644
--- a/src/Action_Diffusion.cpp
+++ b/src/Action_Diffusion.cpp
@@ -345,7 +345,7 @@ Action::RetType Action_Diffusion::Setup(ActionSetup& setup) {
   // Set up sets for individual atoms if necessary
   if (printIndividual_) {
 #   ifdef MPI
-    if (imageOpt_.ImagingEnabled() && trajComm_.Size() > 1) {
+    if (multipleTimeOrigins_ && trajComm_.Size() > 1) {
       mprinterr("Error: Cannot track individual atom diffusion in parallel with imaging.\n");
       return Action::ERR;
     }

From ac852f366fe0f778e4caf47187339e1ac663492b Mon Sep 17 00:00:00 2001
From: "Daniel R. Roe" <daniel.roe@nih.gov>
Date: Thu, 18 May 2023 13:46:04 -0400
Subject: [PATCH 14/15] Update diffusion manual entry. Add avgucell and
 allowmultipleorigins keyword descriptions. Add aspects created for
 'individual' keyword. Add units.

---
 doc/cpptraj.lyx | 139 +++++++++++++++++++++++++++++++++++++++++++++---
 1 file changed, 133 insertions(+), 6 deletions(-)

diff --git a/doc/cpptraj.lyx b/doc/cpptraj.lyx
index c30e9ea981..b05321de32 100644
--- a/doc/cpptraj.lyx
+++ b/doc/cpptraj.lyx
@@ -24005,6 +24005,14 @@ diffusion [{out <filename> | separateout <suffix>}] [time <time per frame>]
 
 \end_layout
 
+\begin_layout LyX-Code
+          [avgucell <avg ucell set>]
+\end_layout
+
+\begin_layout LyX-Code
+          [allowmultipleorigins]
+\end_layout
+
 \begin_deeper
 \begin_layout Description
 [out
@@ -24068,6 +24076,18 @@ name>] MSD data set name.
 [nocalc] Do not calculate diffusion constants.
 \end_layout
 
+\begin_layout Description
+[avgucell] Remove periodic box fluctuations from imaged NPT trajectories
+ using average unit cell vectors.
+\end_layout
+
+\begin_layout Description
+[allowmultipleorigins] (MPI only).
+ For imaged trajectories in parallel, calculate diffusion by averaging over
+ multiple time origins.
+ Should be used with caution.
+\end_layout
+
 \begin_layout Standard
 DataSet Aspects:
 \end_layout
@@ -24112,10 +24132,50 @@ DataSet Aspects:
 [Corr] Linear regression correlation coefficients.
 \end_layout
 
+\begin_layout Description
+[aX]:<atomN> (
+\series bold
+individual
+\series default
+ only) Atom <N> MSD in the X direction.
+\end_layout
+
+\begin_layout Description
+[aY]:<atomN> (
+\series bold
+individual
+\series default
+ only) Atom <N> MSD in the Y direction.
+\end_layout
+
+\begin_layout Description
+[aZ]:<atomN> (
+\series bold
+individual
+\series default
+ only) Atom <N> MSD in the Z direction.
+\end_layout
+
+\begin_layout Description
+[aR]:<atomN> (
+\series bold
+individual
+\series default
+ only) Atom <N> overall MSD.
+\end_layout
+
+\begin_layout Description
+[aA]:<atomN> (
+\series bold
+individual
+\series default
+ only) Atom <N> overall displacement.
+\end_layout
+
 \end_deeper
 \begin_layout Standard
-Compute mean square displacement (MSD) plots (using distance traveled from
- initial position) for the atoms in 
+Compute mean-squared displacement (MSD, in Angstroms squared) plots (using
+ distance traveled from initial position) for the atoms in 
 \series bold
 <mask>
 \series default
@@ -24141,14 +24201,43 @@ In order to correctly calculate diffusion molecules should take continuous
 noimage
 \series default
  keyword can be used.
- In order to remove the noise caused by box fluctuations in NPT trajectories,
- the trajectory should be unwrapped with the average box (see the 
+ To remove the 
+\begin_inset Quotes eld
+\end_inset
+
+noise
+\begin_inset Quotes erd
+\end_inset
+
+ caused by box fluctuations in NPT trajectories, the average unit cell vectors
+ describing the average box can be provided with the 
+\series bold
+avgucell
+\series default
+ keyword; see the 
+\series bold
+\emph on
+avgbox
+\series default
+\emph default
+ command (
+\begin_inset CommandInset ref
+LatexCommand vref
+reference "subsec:cpptraj_avgbox"
+plural "false"
+caps "false"
+noprefix "false"
+
+\end_inset
+
+).
+ Alternatively, the trajectory can be unwrapped prior using the 
 \series bold
 \emph on
 unwrap
 \series default
 \emph default
- command for an example, 
+ command, 
 \begin_inset CommandInset ref
 LatexCommand vref
 reference "subsec:cpptraj_unwrap"
@@ -24158,7 +24247,45 @@ noprefix "false"
 
 \end_inset
 
-).
+.
+ If the trajectory is unwrapped, the 
+\series bold
+noimage
+\series default
+ keyword should be specified.
+\end_layout
+
+\begin_layout Standard
+Note that in parallel, imaging becomes difficult because there is no way
+ to correct for any wrapping that has been done on preceding MPI ranks.
+ Therefore this command will not work on imaged trajectories in parallel
+ by default.
+ There are two workarounds.
+ 1) Unwrap the trajectory prior to 
+\series bold
+\emph on
+diffusion
+\series default
+\emph default
+ with the 
+\series bold
+\emph on
+unwrap
+\series default
+\emph default
+ command, then run the diffusion calculation with the 
+\series bold
+noimage
+\series default
+ keyword (this is the recommended way), or 2) specify the 
+\series bold
+allowmultipleorigins
+\series default
+ keyword to calculate MSD separately on each MPI rank, then averaging over
+ all MSD plots.
+ This means the maximum length of any given MSD plot will be <# frames>
+ / <# MPI ranks>, and the calculated diffusion constants will not be as
+ accurate.
 \end_layout
 
 \begin_layout Standard

From f73432ca52692a1455e24dc139de448d32d5dbbb Mon Sep 17 00:00:00 2001
From: "Daniel R. Roe" <daniel.roe@nih.gov>
Date: Thu, 18 May 2023 13:48:00 -0400
Subject: [PATCH 15/15] 6.19.2. Revision bump for diffusion avgucell and
 allowmultipleorigins keywords, fixed avgbox command in parallel, and addition
 of MAT3X3 DataSet read/write to NetCDF data file.

---
 src/Version.h | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/Version.h b/src/Version.h
index e724fb8ab8..bea2dba8cb 100644
--- a/src/Version.h
+++ b/src/Version.h
@@ -12,7 +12,7 @@
  * Whenever a number that precedes <revision> is incremented, all subsequent
  * numbers should be reset to 0.
  */
-#define CPPTRAJ_INTERNAL_VERSION "V6.19.1"
+#define CPPTRAJ_INTERNAL_VERSION "V6.19.2"
 /// PYTRAJ relies on this
 #define CPPTRAJ_VERSION_STRING CPPTRAJ_INTERNAL_VERSION
 #endif