From a0f82a4b18f86a690ed2428631d287631a8c838a Mon Sep 17 00:00:00 2001
From: Richard Berger <richard.berger@temple.edu>
Date: Mon, 31 Aug 2020 15:00:57 -0400
Subject: [PATCH] Add compression_level parameter to dump xyz/gz and local/gz

---
 src/COMPRESS/dump_local_gz.cpp | 27 +++++++++++++++++++++++++--
 src/COMPRESS/dump_local_gz.h   |  3 +++
 src/COMPRESS/dump_xyz_gz.cpp   | 28 ++++++++++++++++++++++++++--
 src/COMPRESS/dump_xyz_gz.h     |  3 +++
 4 files changed, 57 insertions(+), 4 deletions(-)

diff --git a/src/COMPRESS/dump_local_gz.cpp b/src/COMPRESS/dump_local_gz.cpp
index 0642e52b2b..280b66bb2a 100644
--- a/src/COMPRESS/dump_local_gz.cpp
+++ b/src/COMPRESS/dump_local_gz.cpp
@@ -17,6 +17,7 @@
 #include "update.h"
 
 #include <cstring>
+#include <fmt/format.h>
 
 using namespace LAMMPS_NS;
 
@@ -25,6 +26,8 @@ DumpLocalGZ::DumpLocalGZ(LAMMPS *lmp, int narg, char **arg) :
 {
   gzFp = NULL;
 
+  compression_level = Z_BEST_COMPRESSION;
+
   if (!compressed)
     error->all(FLERR,"Dump local/gz only writes compressed files");
 }
@@ -91,12 +94,15 @@ void DumpLocalGZ::openfile()
   // each proc with filewriter = 1 opens a file
 
   if (filewriter) {
+    std::string mode;
     if (append_flag) {
-      gzFp = gzopen(filecurrent,"ab9");
+      mode = fmt::format("ab{}", compression_level);
     } else {
-      gzFp = gzopen(filecurrent,"wb9");
+      mode = fmt::format("wb{}", compression_level);
     }
 
+    gzFp = gzopen(filecurrent, mode.c_str());
+
     if (gzFp == NULL) error->one(FLERR,"Cannot open dump file");
   } else gzFp = NULL;
 
@@ -171,3 +177,20 @@ void DumpLocalGZ::write()
   }
 }
 
+/* ---------------------------------------------------------------------- */
+
+int DumpLocalGZ::modify_param(int narg, char **arg)
+{
+  int consumed = DumpLocal::modify_param(narg, arg);
+  if(consumed == 0) {
+    if (strcmp(arg[0],"compression_level") == 0) {
+      if (narg < 2) error->all(FLERR,"Illegal dump_modify command");
+      int min_level = Z_DEFAULT_COMPRESSION;
+      int max_level = Z_BEST_COMPRESSION;
+      if (compression_level < 0 || compression_level > max_level)
+        error->all(FLERR, fmt::format("Illegal dump_modify command: compression level must in the range of [{}, {}]", min_level, max_level));
+      return 2;
+    }
+  }
+  return consumed;
+}
diff --git a/src/COMPRESS/dump_local_gz.h b/src/COMPRESS/dump_local_gz.h
index cc788863de..b3f7c7dcf8 100644
--- a/src/COMPRESS/dump_local_gz.h
+++ b/src/COMPRESS/dump_local_gz.h
@@ -31,12 +31,15 @@ class DumpLocalGZ : public DumpLocal {
   virtual ~DumpLocalGZ();
 
  protected:
+  int compression_level;
   gzFile gzFp;  // file pointer for the compressed output stream
 
   virtual void openfile();
   virtual void write_header(bigint);
   virtual void write_data(int, double *);
   virtual void write();
+
+  virtual int modify_param(int, char **);
 };
 
 }
diff --git a/src/COMPRESS/dump_xyz_gz.cpp b/src/COMPRESS/dump_xyz_gz.cpp
index c76b7afa17..bd6000591b 100644
--- a/src/COMPRESS/dump_xyz_gz.cpp
+++ b/src/COMPRESS/dump_xyz_gz.cpp
@@ -16,6 +16,7 @@
 #include "update.h"
 
 #include <cstring>
+#include <fmt/format.h>
 
 using namespace LAMMPS_NS;
 
@@ -24,6 +25,8 @@ DumpXYZGZ::DumpXYZGZ(LAMMPS *lmp, int narg, char **arg) :
 {
   gzFp = NULL;
 
+  compression_level = Z_BEST_COMPRESSION;
+
   if (!compressed)
     error->all(FLERR,"Dump xyz/gz only writes compressed files");
 }
@@ -90,12 +93,15 @@ void DumpXYZGZ::openfile()
   // each proc with filewriter = 1 opens a file
 
   if (filewriter) {
+    std::string mode;
     if (append_flag) {
-      gzFp = gzopen(filecurrent,"ab9");
+      mode = fmt::format("ab{}", compression_level);
     } else {
-      gzFp = gzopen(filecurrent,"wb9");
+      mode = fmt::format("wb{}", compression_level);
     }
 
+    gzFp = gzopen(filecurrent, mode.c_str());
+
     if (gzFp == NULL) error->one(FLERR,"Cannot open dump file");
   } else gzFp = NULL;
 
@@ -134,3 +140,21 @@ void DumpXYZGZ::write()
     }
   }
 }
+
+/* ---------------------------------------------------------------------- */
+
+int DumpXYZGZ::modify_param(int narg, char **arg)
+{
+  int consumed = DumpXYZ::modify_param(narg, arg);
+  if(consumed == 0) {
+    if (strcmp(arg[0],"compression_level") == 0) {
+      if (narg < 2) error->all(FLERR,"Illegal dump_modify command");
+      int min_level = Z_DEFAULT_COMPRESSION;
+      int max_level = Z_BEST_COMPRESSION;
+      if (compression_level < 0 || compression_level > max_level)
+        error->all(FLERR, fmt::format("Illegal dump_modify command: compression level must in the range of [{}, {}]", min_level, max_level));
+      return 2;
+    }
+  }
+  return consumed;
+}
diff --git a/src/COMPRESS/dump_xyz_gz.h b/src/COMPRESS/dump_xyz_gz.h
index 45d5456ddb..834db488a5 100644
--- a/src/COMPRESS/dump_xyz_gz.h
+++ b/src/COMPRESS/dump_xyz_gz.h
@@ -31,12 +31,15 @@ class DumpXYZGZ : public DumpXYZ {
   virtual ~DumpXYZGZ();
 
  protected:
+  int compression_level;
   gzFile gzFp;  // file pointer for the compressed output stream
 
   virtual void openfile();
   virtual void write_header(bigint);
   virtual void write_data(int, double *);
   virtual void write();
+
+  virtual int modify_param(int, char **);
 };
 
 }