From f71be270aa114fe9c32d645bc458ecb38b93c279 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=B8rgen=20Schartum=20Dokken?= Date: Mon, 15 Jan 2024 20:37:15 +0100 Subject: [PATCH] Add `IndexMap::global_to_local` to the Python interface (#2972) * Add `global_to_local` to the Python interface. Make it required to send in output array, as it makes sense from a performance perspective to let the user send in this array, as it might be used in loops. * Add global indices * Make indexmap local_to_global and global_to_local pure functions in python interface * Simplify --- python/dolfinx/wrappers/common.cpp | 14 ++++++++++++-- python/test/unit/common/test_index_map.py | 14 +++++++++++++- 2 files changed, 25 insertions(+), 3 deletions(-) diff --git a/python/dolfinx/wrappers/common.cpp b/python/dolfinx/wrappers/common.cpp index 58e9ad99b0d..0d4aaf2b5dc 100644 --- a/python/dolfinx/wrappers/common.cpp +++ b/python/dolfinx/wrappers/common.cpp @@ -131,8 +131,18 @@ void common(nb::module_& m) self.local_to_global(std::span(local.data(), local.size()), global); return dolfinx_wrappers::as_nbarray(std::move(global)); }, - nb::arg("local")); - + nb::arg("local")) + .def( + "global_to_local", + [](const dolfinx::common::IndexMap& self, + nb::ndarray, nb::c_contig> global) + { + std::vector local(global.size()); + self.global_to_local(std::span(global.data(), global.size()), + local); + return dolfinx_wrappers::as_nbarray(std::move(local)); + }, + nb::arg("global")); // dolfinx::common::Timer nb::class_(m, "Timer", "Timer class") .def(nb::init<>()) diff --git a/python/test/unit/common/test_index_map.py b/python/test/unit/common/test_index_map.py index 2f9a0256ea5..120708f21e5 100644 --- a/python/test/unit/common/test_index_map.py +++ b/python/test/unit/common/test_index_map.py @@ -91,6 +91,19 @@ def test_index_map_ghost_lifetime(): map = dolfinx.common.IndexMap(comm, local_size, [dest, src], map_ghosts, src) assert map.size_global == local_size * comm.size + # Test global to local map + global_indices = np.arange(map.size_global, dtype=np.int64) + local_indices = map.global_to_local(global_indices) + for i, ghost in enumerate(map_ghosts): + assert local_indices[ghost] == local_size + i + + # Create marker for all indices that are on process (local or ghost) + on_process = np.zeros(map.size_global, dtype=bool) + on_process[map.local_range[0]:map.local_range[1]] = True + on_process[map.ghosts] = True + assert (local_indices[on_process] >= 0).all() + assert np.allclose(local_indices[np.invert(on_process)], -1) + # Test lifetime management ghosts = map.ghosts assert np.array_equal(ghosts, map_ghosts) @@ -167,6 +180,5 @@ def test_create_submap_owner_change(): assert np.array_equal(sub_imap.ghosts, [2 * (comm.rank + 1)]) assert np.array_equal(sub_imap.owners, [comm.rank + 1]) assert np.array_equal(sub_imap_to_imap, [0, 2, 3]) - global_indices = sub_imap.local_to_global(np.arange(sub_imap.size_local + sub_imap.num_ghosts, dtype=np.int32)) assert np.array_equal(global_indices, np.arange(comm.rank * 2, comm.rank * 2 + 3))