Loading...
Searching...
No Matches
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});
spdlog::debug("Output data size={}", data.size());
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});
spdlog::debug("Put local_geometry: {}x3", num_vertices);
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]);
spdlog::debug("Put local_cells: {}", 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});
spdlog::debug("Put local_topology: {}x{}", 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);
spdlog::debug("x={}, xshape={}x{}, x_id={}, x_ghost={}", x.size(), xshape[0],
xshape[1], x_id.size(), x_ghost.size());
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.
spdlog::debug("Create cells: [{}x{}]", vtkshape[0], vtkshape[1]);
// 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
spdlog::debug("vertices={}, elements={}, local_geom={}, local_cells={}",
num_dofs, vtkshape[0], x.size(), cells.size());
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(), 1});
engine.Put(orig_id, x_id.data());
adios2::Variable ghost = impl_adios2::define_variable<std::uint8_t>(
io, "vtkGhostType", {}, {}, {x_ghost.size(), 1});
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");
spdlog::debug("ADIOS2: 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())
{
spdlog::debug("ADIOS2: write_mesh");
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)
{
spdlog::debug("ADIOS2: write_mesh_from_space");
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();
}
}
spdlog::debug("Write function data");
// 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
DofMap.h
Degree-of-freedom map representations and tools.
dolfinx::fem::Function
Definition Function.h:47
dolfinx::MPI::rank
int rank(MPI_Comm comm)
Return process rank for the communicator.
Definition MPI.cpp:64
dolfinx::fem::sparsitybuild::cells
void cells(la::SparsityPattern &pattern, std::array< std::span< const std::int32_t >, 2 > cells, std::array< std::reference_wrapper< const DofMap >, 2 > dofmaps)
Iterate over cells and insert entries into sparsity pattern.
Definition sparsitybuild.cpp:16
dolfinx::fem
Finite element method functionality.
Definition assemble_expression_impl.h:23
dolfinx::io
Support for file IO.
Definition cells.h:119