Skip to content

Commit

Permalink
Add meshtags write
Browse files Browse the repository at this point in the history
  • Loading branch information
ampdes committed Jul 26, 2024
1 parent d7e8e3a commit ec05f4c
Show file tree
Hide file tree
Showing 3 changed files with 157 additions and 11 deletions.
37 changes: 36 additions & 1 deletion cpp/demo/checkpointing/main.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -21,16 +21,51 @@ int main(int argc, char* argv[])
dolfinx::init_logging(argc, argv);
MPI_Init(&argc, &argv);

// Create mesh and function space
// Create mesh
auto part = mesh::create_cell_partitioner(mesh::GhostMode::shared_facet);
auto mesh = std::make_shared<mesh::Mesh<float>>(mesh::create_rectangle<float>(
MPI_COMM_WORLD, {{{0.0, 0.0}, {1.0, 1.0}}}, {4, 4},
mesh::CellType::quadrilateral, part));

// Create cell meshtags
int dim = geometry.dim();
topology->create_entities(dim);
const std::shared_ptr<const dolfinx::common::IndexMap> topo_imap
= topology->index_map(dim);

std::uint32_t num_entities = topo_imap->size_local();

auto cmap = mesh->geometry().cmap();
auto geom_layout = cmap.create_dof_layout();
std::uint32_t num_dofs_per_entity = geom_layout.num_entity_closure_dofs(dim);

std::vector<int32_t> entities_array(num_entities * num_dofs_per_entity);
std::vector<int32_t> entities_offsets(num_entities + 1);
std::uint64_t offset = topo_imap->local_range()[0];
std::vector<double> values(num_entities);

for (int i = 0; i < values.size(); ++i)
{
values[i] = (double)(i + offset);
}

for (int i = 0; i < values.size() + 1; ++i)
entities_offsets[i] = entities->offsets()[i];

for (int i = 0; i < (int)(num_entities * num_dofs_per_entity); ++i)
entities_array[i] = entities->array()[i];

graph::AdjacencyList<std::int32_t> entities_local(entities_array,
entities_offsets);

auto meshtags = std::make_shared<mesh::MeshTags<U>>(
mesh::create_meshtags<U>(topology, dim, entities_local, values));

auto writer = ADIOS2Engine(mesh->comm(), "mesh.bp", "mesh-write", "BP5",
adios2::Mode::Write);

io::checkpointing::write(writer, mesh);
io::checkpointing::write(writer, meshtags);

MPI_Finalize();
return 0;
Expand Down
109 changes: 103 additions & 6 deletions cpp/dolfinx/io/checkpointing.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -27,8 +27,8 @@ std::map<basix::element::lagrange_variant, std::string> lagrange_variants{
};

template <std::floating_point T>
void _write(ADIOS2Engine& adios2engine,
std::shared_ptr<dolfinx::mesh::Mesh<T>> mesh)
void _write_mesh(ADIOS2Engine& adios2engine,
std::shared_ptr<dolfinx::mesh::Mesh<T>> mesh)
{

auto io = adios2engine.io();
Expand Down Expand Up @@ -126,26 +126,123 @@ void _write(ADIOS2Engine& adios2engine,
writer->EndStep();
}

template <std::floating_point T>
void _write_meshtags(ADIOS2Engine& adios2engine,
std::shared_ptr<dolfinx::mesh::MeshTags<T>> meshtags)
{
// meshtagsdata
auto tag_entities = meshtags->indices();
auto dim = meshtags->dim();
std::uint32_t num_tag_entities_local
= meshtags->topology()->index_map(dim)->size_local();

int num_tag_entities = tag_entities.size();

std::vector<std::int32_t> local_tag_entities;
local_tag_entities.reserve(num_tag_entities);

std::uint64_t num_saved_tag_entities = 0;
for (int i = 0; i < num_tag_entities; ++i)
{
if (tag_entities[i] < (int)num_tag_entities_local)
{
num_saved_tag_entities += 1;
local_tag_entities.push_back(tag_entities[i]);
}
}
local_tag_entities.resize(num_saved_tag_entities);

// Compute the global offset for owned (local) vertex indices
std::uint64_t local_start = 0;
{
MPI_Exscan(&num_saved_tag_entities, &local_start, 1, MPI_UINT64_T, MPI_SUM,
mesh->comm());
}

std::uint64_t num_tag_entities_global = 0;
MPI_Allreduce(&num_saved_tag_entities, &num_tag_entities_global, 1,
MPI_UINT64_T, MPI_SUM, mesh->comm());

auto values = meshtags->values();
const std::span<const double> local_values(values.begin(),
num_saved_tag_entities);

std::vector<std::int32_t> entities_to_geometry
= mesh::entities_to_geometry(*mesh, dim, tag_entities, false);

auto imap = mesh->geometry().index_map();
std::vector<std::int64_t> topology_array(entities_to_geometry.size());

std::iota(topology_array.begin(), topology_array.end(), 0);

imap->local_to_global(entities_to_geometry, topology_array);

std::string name = "meshtags_" + meshtags->name;

io->DefineAttribute<std::string>("meshtags_name", meshtags->name);
io->DefineAttribute<std::string>("meshtags_dim", dim);

adios2::Variable<std::uint64_t> var_num_tag_entities_global
= io->DefineVariable<std::uint64_t>("num_tag_entities_global");
adios2::Variable<std::uint32_t> var_num_dofs_per_entity
= io->DefineVariable<std::uint32_t>("num_dofs_per_entity");

adios2::Variable<std::int64_t> var_topology
= io->DefineVariable<std::int64_t>(
name + "_topology", {num_tag_entities_global, num_dofs_per_entity},
{local_start, 0}, {num_tag_entities_local, num_dofs_per_entity},
adios2::ConstantDims);

adios2::Variable<T> var_values = io->DefineVariable<T>(
name + "_values", {num_tag_entities_global}, {local_start},
{num_saved_tag_entities}, adios2::ConstantDims);

writer->BeginStep();
writer->Put(var_num_dofs_per_entity, num_dofs_per_entity);
writer->Put(var_num_tag_entities_global, num_tag_entities_global);
writer->Put(var_topology, topology_array.data());
writer->Put(var_values, local_values.data());
writer->EndStep();
}

} // namespace

using namespace dolfinx::io::checkpointing;

//-----------------------------------------------------------------------------
void dolfinx::io::checkpointing::write(
void dolfinx::io::checkpointing::write_mesh(
ADIOS2Engine& adios2engine,
std::shared_ptr<dolfinx::mesh::Mesh<float>> mesh)
{

_write(adios2engine, mesh);
_write_mesh(adios2engine, mesh);
}

//-----------------------------------------------------------------------------
void dolfinx::io::checkpointing::write(
void dolfinx::io::checkpointing::write_mesh(
ADIOS2Engine& adios2engine,
std::shared_ptr<dolfinx::mesh::Mesh<double>> mesh)
{

_write(adios2engine, mesh);
_write_mesh(adios2engine, mesh);
}

//-----------------------------------------------------------------------------
void dolfinx::io::checkpointing::write_meshtags(
ADIOS2Engine& adios2engine,
std::shared_ptr<dolfinx::mesh::MeshTags<float>> meshtags)
{

_write_meshtags(adios2engine, meshtags);
}

//-----------------------------------------------------------------------------
void dolfinx::io::checkpointing::write_meshtags(
ADIOS2Engine& adios2engine,
std::shared_ptr<dolfinx::mesh::MeshTags<double>> meshtags)
{

_write_meshtags(adios2engine, meshtags);
}

#endif
22 changes: 18 additions & 4 deletions cpp/dolfinx/io/checkpointing.h
Original file line number Diff line number Diff line change
Expand Up @@ -24,15 +24,29 @@ namespace dolfinx::io::checkpointing
///
/// @param[in] adios2engine ADIOS2Engine
/// @param[in] mesh Mesh of type float to write to the file
void write(ADIOS2Engine& adios2engine,
std::shared_ptr<dolfinx::mesh::Mesh<float>> mesh);
void write_mesh(ADIOS2Engine& adios2engine,
std::shared_ptr<dolfinx::mesh::Mesh<float>> mesh);

/// @brief Write mesh to a file.
///
/// @param[in] adios2engine ADIOS2Engine
/// @param[in] mesh Mesh of type double to write to the file
void write(ADIOS2Engine& adios2engine,
std::shared_ptr<dolfinx::mesh::Mesh<double>> mesh);
void write_mesh(ADIOS2Engine& adios2engine,
std::shared_ptr<dolfinx::mesh::Mesh<double>> mesh);

/// @brief Write meshtags to a file.
///
/// @param[in] adios2engine ADIOS2Engine
/// @param[in] meshtags MeshTags of type float to write to the file
void write_meshtags(ADIOS2Engine& adios2engine,
std::shared_ptr<dolfinx::mesh::MeshTags<float>> meshtags);

/// @brief Write meshtags to a file.
///
/// @param[in] adios2engine ADIOS2Engine
/// @param[in] meshtags MeshTags of type double to write to the file
void write_meshtags(ADIOS2Engine& adios2engine,
std::shared_ptr<dolfinx::mesh::MeshTags<double>> meshtags);

} // namespace dolfinx::io::checkpointing

Expand Down

0 comments on commit ec05f4c

Please sign in to comment.