First version of Zstd dump atom

cmake/Modules/FindZstd.cmake

@@ -0,0 +1,28 @@
+# Find Zstd library
+#
+#  Zstd_INCLUDE_DIRS - where to find zstd.h, etc.
+#  Zstd_LIBRARIES    - List of libraries when using libzstd
+#  Zstd_FOUND        - True if libzstd is found.
+
+find_path(Zstd_INCLUDE_DIR NAMES zstd.h)
+find_library(Zstd_LIBRARY NAMES zstd)
+
+# handle the QUIET and REQUIRED arguments and
+# set Zstd_FOUND to TRUE if all variables are non-zero
+include(FindPackageHandleStandardArgs)
+find_package_handle_standard_args(Zstd DEFAULT_MSG Zstd_LIBRARY Zstd_INCLUDE_DIR)
+
+# Copy the results to the output variables and target.
+if(Zstd_FOUND)
+  set(Zstd_LIBRARIES ${Zstd_LIBRARY})
+  set(Zstd_INCLUDE_DIRS ${Zstd_INCLUDE_DIR})
+
+  if(NOT TARGET Zstd::Zstd)
+    add_library(Zstd::Zstd UNKNOWN IMPORTED)
+    set_target_properties(Zstd::Zstd PROPERTIES
+      IMPORTED_LOCATION "${Zstd_LIBRARY}"
+      INTERFACE_INCLUDE_DIRECTORIES "${Zstd_INCLUDE_DIR}")
+  endif()
+endif()
+
+mark_as_advanced(Zstd_INCLUDE_DIR Zstd_LIBRARY)

cmake/Modules/Packages/COMPRESS.cmake

@@ -1,2 +1,5 @@
 find_package(ZLIB REQUIRED)
 target_link_libraries(lammps PRIVATE ZLIB::ZLIB)
+
+find_package(Zstd REQUIRED)
+target_link_libraries(lammps PRIVATE Zstd::Zstd)

src/COMPRESS/dump_atom_zstd.cpp

@@ -0,0 +1,221 @@
+/* ----------------------------------------------------------------------
+   LAMMPS - Large-scale Atomic/Molecular Massively Parallel Simulator
+   http://lammps.sandia.gov, Sandia National Laboratories
+   Steve Plimpton, sjplimp@sandia.gov
+
+   Copyright (2003) Sandia Corporation.  Under the terms of Contract
+   DE-AC04-94AL85000 with Sandia Corporation, the U.S. Government retains
+   certain rights in this software.  This software is distributed under
+   the GNU General Public License.
+
+   See the README file in the top-level LAMMPS directory.
+------------------------------------------------------------------------- */
+
+#include "dump_atom_zstd.h"
+#include "domain.h"
+#include "error.h"
+#include "update.h"
+
+#include <fmt/format.h>
+#include <cstring>
+
+using namespace LAMMPS_NS;
+
+DumpAtomZstd::DumpAtomZstd(LAMMPS *lmp, int narg, char **arg) :
+  DumpAtom(lmp, narg, arg)
+{
+  cctx = nullptr;
+  zstdFp = nullptr;
+  fp = nullptr;
+  out_buffer_size = ZSTD_CStreamOutSize();
+  out_buffer = new char[out_buffer_size];
+
+  if (!compressed)
+    error->all(FLERR,"Dump atom/zstd only writes compressed files");
+}
+
+/* ---------------------------------------------------------------------- */
+
+DumpAtomZstd::~DumpAtomZstd()
+{
+  if(cctx && zstdFp) zstd_close();
+
+  delete [] out_buffer;
+  out_buffer = nullptr;
+  out_buffer_size = 0;
+}
+
+/* ----------------------------------------------------------------------
+   generic opening of a dump file
+   ASCII or binary or zstdipped
+   some derived classes override this function
+------------------------------------------------------------------------- */
+
+void DumpAtomZstd::openfile()
+{
+  // single file, already opened, so just return
+
+  if (singlefile_opened) return;
+  if (multifile == 0) singlefile_opened = 1;
+
+  // if one file per timestep, replace '*' with current timestep
+
+  char *filecurrent = filename;
+  if (multiproc) filecurrent = multiname;
+
+  if (multifile) {
+    char *filestar = filecurrent;
+    filecurrent = new char[strlen(filestar) + 16];
+    char *ptr = strchr(filestar,'*');
+    *ptr = '\0';
+    if (padflag == 0)
+      sprintf(filecurrent,"%s" BIGINT_FORMAT "%s",
+              filestar,update->ntimestep,ptr+1);
+    else {
+      char bif[8],pad[16];
+      strcpy(bif,BIGINT_FORMAT);
+      sprintf(pad,"%%s%%0%d%s%%s",padflag,&bif[1]);
+      sprintf(filecurrent,pad,filestar,update->ntimestep,ptr+1);
+    }
+    *ptr = '*';
+    if (maxfiles > 0) {
+      if (numfiles < maxfiles) {
+        nameslist[numfiles] = new char[strlen(filecurrent)+1];
+        strcpy(nameslist[numfiles],filecurrent);
+        ++numfiles;
+      } else {
+        remove(nameslist[fileidx]);
+        delete[] nameslist[fileidx];
+        nameslist[fileidx] = new char[strlen(filecurrent)+1];
+        strcpy(nameslist[fileidx],filecurrent);
+        fileidx = (fileidx + 1) % maxfiles;
+      }
+    }
+  }
+
+  // each proc with filewriter = 1 opens a file
+
+  if (filewriter) {
+    if (append_flag) {
+      zstdFp = fopen(filecurrent,"ab9");
+    } else {
+      zstdFp = fopen(filecurrent,"wb9");
+    }
+
+    if (zstdFp == nullptr) error->one(FLERR,"Cannot open dump file");
+
+    cctx = ZSTD_createCCtx();
+
+    if (cctx == nullptr) error->one(FLERR,"Cannot create Zstd context");
+  } else zstdFp = nullptr;
+
+  // delete string with timestep replaced
+
+  if (multifile) delete [] filecurrent;
+}
+
+/* ---------------------------------------------------------------------- */
+
+void DumpAtomZstd::write_header(bigint ndump)
+{
+  std::string header;
+
+  if ((multiproc) || (!multiproc && me == 0)) {
+    if (unit_flag && !unit_count) {
+      ++unit_count;
+      header = fmt::format("ITEM: UNITS\n%s\n",update->unit_style);
+    }
+
+    if (time_flag) {
+      header += fmt::format("ITEM: TIME\n%.16g\n", compute_time());
+    }
+
+    header += fmt::format("ITEM: TIMESTEP\n{}\n", update->ntimestep);
+    header += fmt::format("ITEM: NUMBER OF ATOMS\n{}\n", ndump);
+    if (domain->triclinic == 0) {
+      header += fmt::format("ITEM: BOX BOUNDS {}\n", boundstr);
+      header += fmt::format("{0:-1.16e} {1:-1.16e}\n", boxxlo, boxxhi);
+      header += fmt::format("{0:-1.16e} {1:-1.16e}\n", boxylo, boxyhi);
+      header += fmt::format("{0:-1.16e} {1:-1.16e}\n", boxzlo, boxzhi);
+    } else {
+      header += fmt::format("ITEM: BOX BOUNDS xy xz yz {}\n", boundstr);
+      header += fmt::format("{0:-1.16e} {1:-1.16e} {2:-1.16e}\n", boxxlo, boxxhi, boxxy);
+      header += fmt::format("{0:-1.16e} {1:-1.16e} {2:-1.16e}\n", boxylo, boxyhi, boxxz);
+      header += fmt::format("{0:-1.16e} {1:-1.16e} {2:-1.16e}\n", boxzlo, boxzhi, boxyz);
+    }
+    header += fmt::format("ITEM: ATOMS {}\n", columns);
+
+    zstd_write(header.c_str(), header.length());
+  }
+}
+
+/* ---------------------------------------------------------------------- */
+
+void DumpAtomZstd::write_data(int n, double *mybuf)
+{
+  ZSTD_inBuffer input = { mybuf, (size_t)n, 0 };
+  ZSTD_EndDirective mode = ZSTD_e_continue;
+
+  do {
+      ZSTD_outBuffer output = { out_buffer, out_buffer_size, 0 };
+      size_t const remaining = ZSTD_compressStream2(cctx, &output, &input, mode);
+      fwrite(out_buffer, sizeof(char), output.pos, zstdFp);
+  } while(input.pos < input.size);
+}
+
+/* ---------------------------------------------------------------------- */
+
+void DumpAtomZstd::write()
+{
+  DumpAtom::write();
+  if (filewriter) {
+    if (multifile) {
+      zstd_close();
+    } else {
+      if (flush_flag && zstdFp) {
+        zstd_flush();
+        fflush(zstdFp);
+      }
+    }
+  }
+}
+
+void DumpAtomZstd::zstd_write(const void * buffer, size_t length) {
+  ZSTD_inBuffer input = { buffer, length, 0 };
+  ZSTD_EndDirective mode = ZSTD_e_continue;
+  
+  do {
+    ZSTD_outBuffer output = { out_buffer, out_buffer_size, 0 };
+    size_t const remaining = ZSTD_compressStream2(cctx, &output, &input, mode);
+    fwrite(out_buffer, sizeof(char), output.pos, zstdFp);
+  } while(input.pos < input.size);
+}
+
+void DumpAtomZstd::zstd_flush() {
+  size_t remaining;
+  ZSTD_inBuffer input = { nullptr, 0, 0 };
+  ZSTD_EndDirective mode = ZSTD_e_flush;
+
+  do {
+    ZSTD_outBuffer output = { out_buffer, out_buffer_size, 0 };
+    remaining = ZSTD_compressStream2(cctx, &output, &input, mode);
+    fwrite(out_buffer, sizeof(char), output.pos, zstdFp);
+  } while(remaining);
+}
+
+void DumpAtomZstd::zstd_close() {
+  size_t remaining;
+  ZSTD_inBuffer input = { nullptr, 0, 0 };
+  ZSTD_EndDirective mode = ZSTD_e_end;
+
+  do {
+    ZSTD_outBuffer output = { out_buffer, out_buffer_size, 0 };
+    remaining = ZSTD_compressStream2(cctx, &output, &input, mode);
+    fwrite(out_buffer, sizeof(char), output.pos, zstdFp);
+  } while(remaining);
+
+  ZSTD_freeCCtx(cctx);
+  cctx = nullptr;
+  if (zstdFp) fclose(zstdFp);
+  zstdFp = nullptr;
+}

src/COMPRESS/dump_atom_zstd.h

@@ -0,0 +1,65 @@
+/* -*- c++ -*- ----------------------------------------------------------
+   LAMMPS - Large-scale Atomic/Molecular Massively Parallel Simulator
+   http://lammps.sandia.gov, Sandia National Laboratories
+   Steve Plimpton, sjplimp@sandia.gov
+
+   Copyright (2003) Sandia Corporation.  Under the terms of Contract
+   DE-AC04-94AL85000 with Sandia Corporation, the U.S. Government retains
+   certain rights in this software.  This software is distributed under
+   the GNU General Public License.
+
+   See the README file in the top-level LAMMPS directory.
+------------------------------------------------------------------------- */
+
+#ifdef DUMP_CLASS
+
+DumpStyle(atom/zstd,DumpAtomZstd)
+
+#else
+
+#ifndef LMP_DUMP_ATOM_ZSTD_H
+#define LMP_DUMP_ATOM_ZSTD_H
+
+#include "dump_atom.h"
+#include <zstd.h>
+#include <stdio.h>
+
+namespace LAMMPS_NS {
+
+class DumpAtomZstd : public DumpAtom {
+ public:
+  DumpAtomZstd(class LAMMPS *, int, char **);
+  virtual ~DumpAtomZstd();
+
+ protected:
+  ZSTD_CCtx * cctx;
+  FILE * zstdFp;
+  char * out_buffer;
+  size_t out_buffer_size;
+
+  virtual void openfile();
+  virtual void write_header(bigint);
+  virtual void write_data(int, double *);
+  virtual void write();
+
+  void zstd_write(const void * buffer, size_t length);
+  void zstd_flush();
+  void zstd_close();
+};
+
+}
+
+#endif
+#endif
+
+/* ERROR/WARNING messages:
+
+E: Dump atom/zstd only writes compressed files
+
+The dump atom/zstd output file name must have a .zst suffix.
+
+E: Cannot open dump file
+
+Self-explanatory.
+
+*/

src/dump.cpp

@@ -116,6 +116,7 @@ Dump::Dump(LAMMPS *lmp, int /*narg*/, char **arg) : Pointers(lmp)
   // check file suffixes
   //   if ends in .bin = binary file
   //   else if ends in .gz = gzipped text file
+  //   else if ends in .zst = Zstd compressed text file
   //   else ASCII text file
 
   fp = NULL;
@@ -153,6 +154,8 @@ Dump::Dump(LAMMPS *lmp, int /*narg*/, char **arg) : Pointers(lmp)
   if (suffix > filename && strcmp(suffix,".bin") == 0) binary = 1;
   suffix = filename + strlen(filename) - strlen(".gz");
   if (suffix > filename && strcmp(suffix,".gz") == 0) compressed = 1;
+  suffix = filename + strlen(filename) - strlen(".zst");
+  if (suffix > filename && strcmp(suffix,".zst") == 0) compressed = 1;
 }
 
 /* ---------------------------------------------------------------------- */