ADIOS2Writers.h

// Copyright (C) 2021-2023 Jørgen S. Dokken and Garth N. Wells
//
// This file is part of DOLFINX (https://www.fenicsproject.org)
//
// SPDX-License-Identifier:    LGPL-3.0-or-later
 
#pragma once
 
#ifdef HAS_ADIOS2
 
#include "vtk_utils.h"
#include <adios2.h>
#include <algorithm>
#include <basix/mdspan.hpp>
#include <cassert>
#include <complex>
#include <concepts>
#include <dolfinx/common/IndexMap.h>
#include <dolfinx/fem/DofMap.h>
#include <dolfinx/fem/FiniteElement.h>
#include <dolfinx/fem/Function.h>
#include <dolfinx/mesh/Geometry.h>
#include <dolfinx/mesh/Mesh.h>
#include <filesystem>
#include <memory>
#include <mpi.h>
#include <string>
#include <type_traits>
#include <variant>
#include <vector>

 
namespace dolfinx::fem
{
template <dolfinx::scalar T, std::floating_point U>
class Function;
}
 
namespace dolfinx::io
{
namespace adios2_writer
{
template <std::floating_point T>
using U = std::vector<std::variant<
    std::shared_ptr<const fem::Function<float, T>>,
    std::shared_ptr<const fem::Function<double, T>>,
    std::shared_ptr<const fem::Function<std::complex<float>, T>>,
    std::shared_ptr<const fem::Function<std::complex<double>, T>>>>;
} // namespace adios2_writer

class ADIOS2Writer
{
protected:
  ADIOS2Writer(MPI_Comm comm, const std::filesystem::path& filename,
               std::string tag, std::string engine);

  ADIOS2Writer(ADIOS2Writer&& writer) = default;

  ADIOS2Writer(const ADIOS2Writer&) = delete;

  ~ADIOS2Writer();

  ADIOS2Writer& operator=(ADIOS2Writer&& writer) = default;
 
  // Copy assignment
  ADIOS2Writer& operator=(const ADIOS2Writer&) = delete;
 
public:
  void close();
 
protected:
  std::unique_ptr<adios2::ADIOS> _adios;
  std::unique_ptr<adios2::IO> _io;
  std::unique_ptr<adios2::Engine> _engine;
};

namespace impl_adios2
{
constexpr std::array field_ext = {"_real", "_imag"};

template <class T>
adios2::Attribute<T> define_attribute(adios2::IO& io, std::string name,
                                      const T& value, std::string var_name = "",
                                      std::string separator = "/")
{
  if (adios2::Attribute<T> attr = io.InquireAttribute<T>(name); attr)
    return attr;
  else
    return io.DefineAttribute<T>(name, value, var_name, separator);
}

template <class T>
adios2::Variable<T> define_variable(adios2::IO& io, std::string name,
                                    const adios2::Dims& shape = adios2::Dims(),
                                    const adios2::Dims& start = adios2::Dims(),
                                    const adios2::Dims& count = adios2::Dims())
{
  if (adios2::Variable v = io.InquireVariable<T>(name); v)
  {
    if (v.Count() != count and v.ShapeID() == adios2::ShapeID::LocalArray)
      v.SetSelection({start, count});
    return v;
  }
  else
    return io.DefineVariable<T>(name, shape, start, count);
}

template <std::floating_point T>
std::shared_ptr<const mesh::Mesh<T>>
extract_common_mesh(const typename adios2_writer::U<T>& u)
{
  // Extract mesh from first function
  assert(!u.empty());
  auto mesh = std::visit([](auto&& u) { return u->function_space()->mesh(); },
                         u.front());
  assert(mesh);
 
  // Check that all functions share the same mesh
  for (auto& v : u)
  {
    std::visit(
        [&mesh](auto&& u)
        {
          if (mesh != u->function_space()->mesh())
          {
            throw std::runtime_error(
                "ADIOS2Writer only supports functions sharing the same mesh");
          }
        },
        v);
  }
 
  return mesh;
}
 
} // namespace impl_adios2

namespace impl_vtx
{
std::stringstream create_vtk_schema(const std::vector<std::string>& point_data,
                                    const std::vector<std::string>& cell_data);

template <std::floating_point T>
std::tuple<std::vector<std::string>, std::vector<std::string>>
extract_function_names(const typename adios2_writer::U<T>& u)
{
  std::vector<std::string> names, dg0_names;
  for (auto& v : u)
  {
    std::visit(
        [&names, &dg0_names](auto&& u)
        {
          using U = std::decay_t<decltype(u)>;
          using X = typename U::element_type;
          std::vector<std::string>* fnames = &names;
          if (impl::is_cellwise(*(u->function_space()->element())))
          {
            fnames = &dg0_names;
          }
          if constexpr (std::is_floating_point_v<typename X::value_type>)
            fnames->push_back(u->name);
          else
          {
            fnames->push_back(u->name + impl_adios2::field_ext[0]);
            fnames->push_back(u->name + impl_adios2::field_ext[1]);
          }
        },
        v);
  }
 
  return {names, dg0_names};
}

template <typename T, std::floating_point X>
void vtx_write_data(adios2::IO& io, adios2::Engine& engine,
                    const fem::Function<T, X>& u)
{
  // Get function data array and information about layout
  assert(u.x());
  std::span<const T> u_vector = u.x()->array();
 
  // Pad to 3D if vector/tensor is product of dimensions is smaller than
  // 3**rank to ensure that we can visualize them correctly in Paraview
  std::span<const std::size_t> value_shape
      = u.function_space()->element()->value_shape();
  std::size_t rank = value_shape.size();
  std::size_t num_comp = std::reduce(value_shape.begin(), value_shape.end(), 1,
                                     std::multiplies{});
  if (num_comp < std::pow(3, rank))
    num_comp = std::pow(3, rank);
 
  std::shared_ptr<const fem::DofMap> dofmap = u.function_space()->dofmap();
  assert(dofmap);
  std::shared_ptr<const common::IndexMap> index_map = dofmap->index_map;
  assert(index_map);
  int index_map_bs = dofmap->index_map_bs();
  int dofmap_bs = dofmap->bs();
  std::uint32_t num_dofs = index_map_bs
                           * (index_map->size_local() + index_map->num_ghosts())
                           / dofmap_bs;
  if constexpr (std::is_scalar_v<T>)
  {
    // ---- Real
    std::vector<T> data(num_dofs * num_comp, 0);
    for (std::size_t i = 0; i < num_dofs; ++i)
      for (int j = 0; j < index_map_bs; ++j)
        data[i * num_comp + j] = u_vector[i * index_map_bs + j];
 
    adios2::Variable output = impl_adios2::define_variable<T>(
        io, u.name, {}, {}, {num_dofs, num_comp});
    engine.Put(output, data.data(), adios2::Mode::Sync);
  }
  else
  {
    // ---- Complex
    using U = typename T::value_type;
 
    std::vector<U> data(num_dofs * num_comp, 0);
    for (std::size_t i = 0; i < num_dofs; ++i)
      for (int j = 0; j < index_map_bs; ++j)
        data[i * num_comp + j] = std::real(u_vector[i * index_map_bs + j]);
 
    adios2::Variable output_real = impl_adios2::define_variable<U>(
        io, u.name + impl_adios2::field_ext[0], {}, {}, {num_dofs, num_comp});
    engine.Put(output_real, data.data(), adios2::Mode::Sync);
 
    std::ranges::fill(data, 0);
    for (std::size_t i = 0; i < num_dofs; ++i)
      for (int j = 0; j < index_map_bs; ++j)
        data[i * num_comp + j] = std::imag(u_vector[i * index_map_bs + j]);
    adios2::Variable output_imag = impl_adios2::define_variable<U>(
        io, u.name + impl_adios2::field_ext[1], {}, {}, {num_dofs, num_comp});
    engine.Put(output_imag, data.data(), adios2::Mode::Sync);
  }
}

template <std::floating_point T>
void vtx_write_mesh(adios2::IO& io, adios2::Engine& engine,
                    const mesh::Mesh<T>& mesh)
{
  const mesh::Geometry<T>& geometry = mesh.geometry();
  auto topology = mesh.topology();
  assert(topology);
 
  // "Put" geometry
  std::shared_ptr<const common::IndexMap> x_map = geometry.index_map();
  std::uint32_t num_vertices = x_map->size_local() + x_map->num_ghosts();
  adios2::Variable local_geometry = impl_adios2::define_variable<T>(
      io, "geometry", {}, {}, {num_vertices, 3});
  engine.Put(local_geometry, geometry.x().data());
 
  // Put number of nodes. The mesh data is written with local indices,
  // therefore we need the ghost vertices.
  adios2::Variable vertices = impl_adios2::define_variable<std::uint32_t>(
      io, "NumberOfNodes", {adios2::LocalValueDim});
  engine.Put<std::uint32_t>(vertices, num_vertices);
 
  auto [vtkcells, shape]
      = io::extract_vtk_connectivity(geometry.dofmap(), topology->cell_type());
 
  // Add cell metadata
  int tdim = topology->dim();
  adios2::Variable cell_var = impl_adios2::define_variable<std::uint32_t>(
      io, "NumberOfCells", {adios2::LocalValueDim});
  engine.Put<std::uint32_t>(cell_var, shape[0]);
  adios2::Variable celltype_var
      = impl_adios2::define_variable<std::uint32_t>(io, "types");
  engine.Put<std::uint32_t>(
      celltype_var, cells::get_vtk_cell_type(topology->cell_type(), tdim));
 
  // Pack mesh 'nodes'. Output is written as [N0, v0_0,...., v0_N0, N1,
  // v1_0,...., v1_N1,....], where N is the number of cell nodes and v0,
  // etc, is the node index
  std::vector<std::int64_t> cells(shape[0] * (shape[1] + 1), shape[1]);
  for (std::size_t c = 0; c < shape[0]; ++c)
  {
    std::span vtkcell(vtkcells.data() + c * shape[1], shape[1]);
    std::span cell(cells.data() + c * (shape[1] + 1), shape[1] + 1);
    std::ranges::copy(vtkcell, std::next(cell.begin()));
  }
 
  // Put topology (nodes)
  adios2::Variable local_topology = impl_adios2::define_variable<std::int64_t>(
      io, "connectivity", {}, {}, {shape[0], shape[1] + 1});
  engine.Put(local_topology, cells.data());
 
  // Vertex global ids and ghost markers
  adios2::Variable orig_id = impl_adios2::define_variable<std::int64_t>(
      io, "vtkOriginalPointIds", {}, {}, {num_vertices});
  engine.Put(orig_id, geometry.input_global_indices().data());
 
  std::vector<std::uint8_t> x_ghost(num_vertices, 0);
  std::fill(std::next(x_ghost.begin(), x_map->size_local()), x_ghost.end(), 1);
  adios2::Variable ghost = impl_adios2::define_variable<std::uint8_t>(
      io, "vtkGhostType", {}, {}, {x_ghost.size()});
  engine.Put(ghost, x_ghost.data());
  engine.PerformPuts();
}

template <std::floating_point T>
std::pair<std::vector<std::int64_t>, std::vector<std::uint8_t>>
vtx_write_mesh_from_space(adios2::IO& io, adios2::Engine& engine,
                          const fem::FunctionSpace<T>& V)
{
  auto mesh = V.mesh();
  assert(mesh);
  auto topology = mesh->topology();
  assert(topology);
  int tdim = topology->dim();
 
  // Get a VTK mesh with points at the 'nodes'
  auto [x, xshape, x_id, x_ghost, vtk, vtkshape] = io::vtk_mesh_from_space(V);
 
  std::uint32_t num_dofs = xshape[0];
 
  // -- Pack mesh 'nodes'. Output is written as [N0, v0_0,...., v0_N0, N1,
  // v1_0,...., v1_N1,....], where N is the number of cell nodes and v0,
  // etc, is the node index.
 
  // Create vector, setting all entries to nodes per cell (vtk.shape(1))
  std::vector<std::int64_t> cells(vtkshape[0] * (vtkshape[1] + 1), vtkshape[1]);
 
  // Set the [v0_0,...., v0_N0, v1_0,...., v1_N1,....] data
  for (std::size_t c = 0; c < vtkshape[0]; ++c)
  {
    std::span vtkcell(vtk.data() + c * vtkshape[1], vtkshape[1]);
    std::span cell(cells.data() + c * (vtkshape[1] + 1), vtkshape[1] + 1);
    std::ranges::copy(vtkcell, std::next(cell.begin()));
  }
 
  // Define ADIOS2 variables for geometry, topology, celltypes and
  // corresponding VTK data
  adios2::Variable local_geometry
      = impl_adios2::define_variable<T>(io, "geometry", {}, {}, {num_dofs, 3});
  adios2::Variable local_topology = impl_adios2::define_variable<std::int64_t>(
      io, "connectivity", {}, {}, {vtkshape[0], vtkshape[1] + 1});
  adios2::Variable cell_type
      = impl_adios2::define_variable<std::uint32_t>(io, "types");
  adios2::Variable vertices = impl_adios2::define_variable<std::uint32_t>(
      io, "NumberOfNodes", {adios2::LocalValueDim});
  adios2::Variable elements = impl_adios2::define_variable<std::uint32_t>(
      io, "NumberOfEntities", {adios2::LocalValueDim});
 
  // Write mesh information to file
  engine.Put<std::uint32_t>(vertices, num_dofs);
  engine.Put<std::uint32_t>(elements, vtkshape[0]);
  engine.Put<std::uint32_t>(
      cell_type, cells::get_vtk_cell_type(topology->cell_type(), tdim));
  engine.Put(local_geometry, x.data());
  engine.Put(local_topology, cells.data());
 
  // Node global ids
  adios2::Variable orig_id = impl_adios2::define_variable<std::int64_t>(
      io, "vtkOriginalPointIds", {}, {}, {x_id.size()});
  engine.Put(orig_id, x_id.data());
  adios2::Variable ghost = impl_adios2::define_variable<std::uint8_t>(
      io, "vtkGhostType", {}, {}, {x_ghost.size()});
  engine.Put(ghost, x_ghost.data());
 
  engine.PerformPuts();
  return {std::move(x_id), std::move(x_ghost)};
}
} // namespace impl_vtx

enum class VTXMeshPolicy : std::int8_t
{
  update, 
  reuse   
};

template <std::floating_point T>
class VTXWriter : public ADIOS2Writer
{
public:
  VTXWriter(MPI_Comm comm, const std::filesystem::path& filename,
            std::shared_ptr<const mesh::Mesh<T>> mesh,
            std::string engine = "BPFile")
      : ADIOS2Writer(comm, filename, "VTX mesh writer", engine), _mesh(mesh),
        _mesh_reuse_policy(VTXMeshPolicy::update),
        _has_piecewise_constant(false)
  {
    // Define VTK scheme attribute for mesh
    std::string vtk_scheme = impl_vtx::create_vtk_schema({}, {}).str();
    impl_adios2::define_attribute<std::string>(*_io, "vtk.xml", vtk_scheme);
  }

  VTXWriter(MPI_Comm comm, const std::filesystem::path& filename,
            const typename adios2_writer::U<T>& u, std::string engine,
            VTXMeshPolicy mesh_policy = VTXMeshPolicy::update)
      : ADIOS2Writer(comm, filename, "VTX function writer", engine),
        _mesh(impl_adios2::extract_common_mesh<T>(u)), _u(u),
        _mesh_reuse_policy(mesh_policy), _has_piecewise_constant(false)
  {
    if (u.empty())
      throw std::runtime_error("VTXWriter fem::Function list is empty.");
 
    // Extract space from first function
    auto V0 = std::visit([](auto& u) { return u->function_space().get(); },
                         u.front());
    assert(V0);
 
    // Replace first function space if it is a space for piecewise constants
    bool has_V0_changed = false;
    for (auto& v : u)
    {
      std::visit(
          [&V0, &has_V0_changed](auto& u)
          {
            auto V = u->function_space().get();
            assert(V);
            if (!impl::is_cellwise(*V->element()))
            {
              V0 = V;
              has_V0_changed = true;
            }
          },
          v);
      if (has_V0_changed)
        break;
    }
    auto element0 = V0->element().get();
    assert(element0);
 
    // Check if function is mixed
    if (element0->is_mixed())
    {
      throw std::runtime_error(
          "Mixed functions are not supported by VTXWriter.");
    }
 
    // FIXME: is the below check adequate for detecting a Lagrange
    // element?
    // Check that element is Lagrange
    if (!element0->interpolation_ident())
    {
      throw std::runtime_error(
          "Only (discontinuous) Lagrange functions are "
          "supported. Interpolate Functions before output.");
    }
 
    // Check that all functions come from same element type
    for (auto& v : _u)
    {
      std::visit(
          [V0, this](auto& u)
          {
            auto element = u->function_space()->element();
            assert(element);
            bool is_piecewise_constant = impl::is_cellwise(*element);
            _has_piecewise_constant
                = _has_piecewise_constant || is_piecewise_constant;
            if (*element != *V0->element().get() and !is_piecewise_constant)
            {
              throw std::runtime_error("All functions in VTXWriter must have "
                                       "the same element type.");
            }
#ifndef NDEBUG
            auto dmap0 = V0->dofmap()->map();
            auto dmap = u->function_space()->dofmap()->map();
            if ((dmap0.size() != dmap.size()
                 or !std::equal(dmap0.data_handle(),
                                dmap0.data_handle() + dmap0.size(),
                                dmap.data_handle()))
                and !is_piecewise_constant)
            {
              throw std::runtime_error(
                  "All functions must have the same dofmap for VTXWriter.");
            }
#endif
          },
          v);
    }
 
    // Define VTK scheme attribute for set of functions
    auto [names, dg0_names] = impl_vtx::extract_function_names<T>(u);
    std::string vtk_scheme;
    vtk_scheme = impl_vtx::create_vtk_schema(names, dg0_names).str();
 
    impl_adios2::define_attribute<std::string>(*_io, "vtk.xml", vtk_scheme);
  }

  VTXWriter(MPI_Comm comm, const std::filesystem::path& filename,
            const typename adios2_writer::U<T>& u,
            VTXMeshPolicy mesh_policy = VTXMeshPolicy::update)
      : VTXWriter(comm, filename, u, "BPFile", mesh_policy)
  {
  }
 
  // Copy constructor
  VTXWriter(const VTXWriter&) = delete;

  VTXWriter(VTXWriter&& file) = default;

  ~VTXWriter() = default;

  VTXWriter& operator=(VTXWriter&&) = default;
 
  // Copy assignment
  VTXWriter& operator=(const VTXWriter&) = delete;

  void write(double t)
  {
    assert(_io);
    adios2::Variable var_step
        = impl_adios2::define_variable<double>(*_io, "step");
 
    assert(_engine);
    _engine->BeginStep();
    _engine->template Put<double>(var_step, t);
 
    // If we have no non-constant functions write the mesh to file
    auto [names, dg0_names] = impl_vtx::extract_function_names<T>(_u);
    if ((names.size() == 0) or _u.empty())
    {
      impl_vtx::vtx_write_mesh(*_io, *_engine, *_mesh);
    }
    else
    {
      if (_mesh_reuse_policy == VTXMeshPolicy::update
          or !(_io->template InquireVariable<std::int64_t>("connectivity")))
      {
        // Write a single mesh for functions as they share finite
        // element
        std::tie(_x_id, _x_ghost) = std::visit(
            [&](auto& u)
            {
              return impl_vtx::vtx_write_mesh_from_space(*_io, *_engine,
                                                         *u->function_space());
            },
            _u[0]);
      }
      else
      {
        // Node global ids
        adios2::Variable orig_id = impl_adios2::define_variable<std::int64_t>(
            *_io, "vtkOriginalPointIds", {}, {}, {_x_id.size()});
        _engine->Put(orig_id, _x_id.data());
        adios2::Variable ghost = impl_adios2::define_variable<std::uint8_t>(
            *_io, "vtkGhostType", {}, {}, {_x_ghost.size()});
        _engine->Put(ghost, _x_ghost.data());
        _engine->PerformPuts();
      }
    }
    // Write function data for each function to file
    for (auto& v : _u)
    {
      std::visit([&](auto& u) { impl_vtx::vtx_write_data(*_io, *_engine, *u); },
                 v);
    }
 
    _engine->EndStep();
  }
 
private:
  std::shared_ptr<const mesh::Mesh<T>> _mesh;
  adios2_writer::U<T> _u;
 
  // Control whether the mesh is written to file once or at every time
  // step
  VTXMeshPolicy _mesh_reuse_policy;
  std::vector<std::int64_t> _x_id;
  std::vector<std::uint8_t> _x_ghost;
 
  // Special handling of piecewise constant functions
  bool _has_piecewise_constant;
};

template <typename U, typename T>
VTXWriter(MPI_Comm comm, U filename, T mesh)
    -> VTXWriter<typename std::remove_cvref<
        typename T::element_type>::type::geometry_type::value_type>;
 
} // namespace dolfinx::io
 
#endif