From 4e918b34a0f083ef4262d73b0b5900948dee30c8 Mon Sep 17 00:00:00 2001 From: Lucas Sawade Date: Fri, 17 Jul 2026 16:26:51 +0200 Subject: [PATCH] StiffnessAssembler --- core/specfem/linear_system/dof_map.hpp | 165 ++++++++++ .../linear_system/tpetra_assembler.cpp | 231 +++++++++++++ .../linear_system/tpetra_assembler.hpp | 121 +++++++ docs/Doxyfile.in | 3 +- .../api/specfem/linear_system/index.rst | 22 +- .../stiffness_assembler_tests.cpp | 306 ++++++++++++++++++ tests/unit-tests/serial.cmake | 24 ++ 7 files changed, 866 insertions(+), 6 deletions(-) create mode 100644 core/specfem/linear_system/dof_map.hpp create mode 100644 core/specfem/linear_system/tpetra_assembler.cpp create mode 100644 core/specfem/linear_system/tpetra_assembler.hpp create mode 100644 tests/unit-tests/linear_system/stiffness_assembler_tests.cpp diff --git a/core/specfem/linear_system/dof_map.hpp b/core/specfem/linear_system/dof_map.hpp new file mode 100644 index 000000000..6393275db --- /dev/null +++ b/core/specfem/linear_system/dof_map.hpp @@ -0,0 +1,165 @@ +#pragma once + +#ifdef SPECFEM_ENABLE_TRILINOS + +#include "specfem/assembly/assembly.hpp" +#include "specfem/element.hpp" +#include "specfem/enums.hpp" +#include "specfem/setup.hpp" +#include +#include +#include +#include +#include +#include + +namespace specfem { +namespace linear_system { + +/** + * @brief Scalar used for assembled matrices and vectors. + * + * Aliased to `type_real` so the linear system follows the build's precision: + * `float` by default (matching the float-only Tpetra installs on the + * clusters), `double` with `SPECFEM_ENABLE_DOUBLE_PRECISION`. Single switch + * point if the two ever need to diverge. + */ +using scalar_type = type_real; + +/// Tpetra map with the default local/global ordinals and node type +using map_type = Tpetra::Map<>; + +/// Global ordinal used for degree-of-freedom ids +using global_ordinal_type = map_type::global_ordinal_type; + +/** + * @brief Maps SPECFEM++ (iglob, icomp) degrees of freedom of one medium to + * Tpetra global ids and owns the row (owned) and column (overlap) maps. + * + * The GID layout is component-blocked: `gid = icomp * nglob + iglob`, + * deliberately matching SPECFEM++ field storage + * (`Kokkos::View` of shape + * `(nglob, ncomp)`), so a Tpetra vector maps 1:1 onto field memory with no + * permutation at solve time. It also matches the element-local ordering + * `specfem::linear_system::local_dof_index`. The layout lives ONLY in + * @ref gid -- change it there to swap the ordering. + * + * One DofMap instance describes one medium: `nglob` and `ncomp` are + * per-medium quantities (a future multi-medium system holds one DofMap and + * one matrix block per medium). + * + * Serial-only in this milestone: SPECFEM++ numbers global points per rank + * with no globally consistent numbering across ranks, so the constructor + * throws for communicators with more than one rank. The owned and overlap + * maps are kept separate in the API so distributed assembly (owned iglob -> + * owned GIDs, shared-interface points -> overlap map, Export(ADD) into the + * owned matrix) can be added without changing callers; at one rank they are + * the same map. + */ +class DofMap { +public: + /** + * @brief Construct the map for `nglob` mesh points with `ncomp` components + * per point. + * + * @param nglob Number of unique global mesh points of the medium + * @param ncomp Number of field components per point (3 for 3D elastic) + * @param comm Teuchos communicator; must have exactly one rank + */ + DofMap(const int nglob, const int ncomp, + const Teuchos::RCP> &comm) + : nglob_(nglob), ncomp_(ncomp), comm_(comm) { + if (comm_->getSize() > 1) { + throw std::runtime_error( + "specfem::linear_system::DofMap: distributed assembly on " + + std::to_string(comm_->getSize()) + + " ranks is not implemented yet. SPECFEM++ numbers global points " + "per rank, and the cross-rank GID negotiation is a follow-up of " + "issue #1982. Run on a single rank."); + } + const Tpetra::global_size_t num_entries = + static_cast(num_global_dofs()); + const global_ordinal_type index_base = 0; + owned_map_ = Teuchos::rcp(new map_type(num_entries, index_base, comm_)); + // At one rank every dof is owned; a distributed build replaces this with + // a map that additionally holds the shared-interface dofs of neighbor + // ranks. + overlap_map_ = owned_map_; + } + + /** + * @brief Build the DofMap of one medium from an assembled 3D simulation. + * + * Reads `nglob` from the forward simulation field and the component count + * from the element attributes, using `Tpetra::getDefaultComm()`. + * + * @tparam MediumTag Medium whose degrees of freedom the map describes + * @param assembly Assembly with constructed fields + * @return DofMap for the medium + */ + template + static DofMap from_assembly( + const specfem::assembly::assembly + &assembly) { + const auto &field = + assembly.fields + .get_simulation_field(); + const int nglob = field.get_nglob(); + constexpr int ncomp = + specfem::element::attributes::components; + return DofMap(nglob, ncomp, Tpetra::getDefaultComm()); + } + + /** + * @brief Global id of component `icomp` at mesh point `iglob`. + * + * Component-blocked layout `gid = icomp * nglob + iglob` -- the single + * source of truth for the global dof ordering (see the class docs). + * + * @param iglob Per-medium global point index in `[0, nglob())` + * @param icomp Field component in `[0, ncomp())` + * @return Tpetra global dof id in `[0, num_global_dofs())` + */ + inline global_ordinal_type gid(const int iglob, const int icomp) const { + return static_cast(icomp) * nglob_ + iglob; + } + + /// Number of unique global mesh points of the medium + inline int nglob() const { return nglob_; } + + /// Number of field components per mesh point + inline int ncomp() const { return ncomp_; } + + /// Total number of degrees of freedom: `ncomp * nglob` + inline global_ordinal_type num_global_dofs() const { + return static_cast(ncomp_) * nglob_; + } + + /// Uniquely-owned row map: contiguous `[0, num_global_dofs())` + inline Teuchos::RCP owned_map() const { return owned_map_; } + + /** + * @brief Overlap (column) map: owned dofs plus, in a future distributed + * build, the shared-interface dofs of neighbor ranks. Equal to + * @ref owned_map at one rank. + */ + inline Teuchos::RCP overlap_map() const { + return overlap_map_; + } + + /// Communicator the maps are defined on + inline Teuchos::RCP> comm() const { return comm_; } + +private: + int nglob_ = 0; ///< Points of the medium + int ncomp_ = 0; ///< Components per point + Teuchos::RCP> comm_; ///< Communicator + Teuchos::RCP owned_map_; ///< Uniquely-owned rows + Teuchos::RCP overlap_map_; ///< Owned + shared dofs +}; + +} // namespace linear_system +} // namespace specfem + +#endif // SPECFEM_ENABLE_TRILINOS diff --git a/core/specfem/linear_system/tpetra_assembler.cpp b/core/specfem/linear_system/tpetra_assembler.cpp new file mode 100644 index 000000000..d4f084e27 --- /dev/null +++ b/core/specfem/linear_system/tpetra_assembler.cpp @@ -0,0 +1,231 @@ +#include "specfem/linear_system/tpetra_assembler.hpp" + +#ifdef SPECFEM_ENABLE_TRILINOS + +#include "specfem/assembly/assembly.hpp" +#include "specfem/datatype/element_index_range.hpp" +#include "specfem/element/to_string.hpp" +#include "specfem/linear_system/element_stiffness.hpp" +#include "specfem/tags.hpp" +#include +#include +#include +#include +#include +#include +#include + +template +specfem::linear_system::StiffnessAssembler::StiffnessAssembler( + const AssemblyType &assembly, const int batch_size) + : assembly_(assembly), + dof_map_(DofMap::from_assembly(assembly)), + batch_size_(batch_size) { + + if (batch_size_ < 1) { + throw std::runtime_error( + "specfem::linear_system::StiffnessAssembler: batch_size must be at " + "least 1."); + } + + specfem::linear_system::validate_stiffness_scope(assembly_); + + // Single-medium milestone: matrix blocks coupling different media + // (fluid-solid) are deferred, so reject mixed meshes outright rather than + // silently assembling an operator that ignores the coupling. + const int nspec = assembly_.element_types.nspec; + for (int ispec = 0; ispec < nspec; ++ispec) { + const auto element_medium = assembly_.element_types.get_medium_tag(ispec); + if (element_medium != medium_tag) { + std::ostringstream message; + message << "specfem::linear_system::StiffnessAssembler: element " << ispec + << " has medium '" << specfem::element::to_string(element_medium) + << "'; only single-medium '" + << specfem::element::to_string(medium_tag) + << "' meshes are supported. Fluid-solid coupling blocks are " + "deferred (issue #1982)."; + throw std::runtime_error(message.str()); + } + } +} + +template +std::vector +specfem::linear_system::StiffnessAssembler::element_column_gids( + const int ispec) const { + const auto &field = assembly_.fields.template get_simulation_field< + specfem::simulation::field_type::forward>(); + const int ngllz = field.ngllz; + const int nglly = field.nglly; + const int ngllx = field.ngllx; + const int npoints = ngllz * nglly * ngllx; + + std::vector cols(static_cast(ncomp) * + npoints); + for (int iz = 0; iz < ngllz; ++iz) { + for (int iy = 0; iy < nglly; ++iy) { + for (int ix = 0; ix < ngllx; ++ix) { + const int iglob = + field.template get_iglob(ispec, iz, iy, ix); + // Same ordering as local_dof_index on the cubic GLL grid. + const int point = (iz * nglly + iy) * ngllx + ix; + for (int icomp = 0; icomp < ncomp; ++icomp) { + cols[static_cast(icomp) * npoints + point] = + dof_map_.gid(iglob, icomp); + } + } + } + } + return cols; +} + +template +Teuchos::RCP +specfem::linear_system::StiffnessAssembler::build_graph() const { + const auto &field = assembly_.fields.template get_simulation_field< + specfem::simulation::field_type::forward>(); + const int ngllz = field.ngllz; + const int nglly = field.nglly; + const int ngllx = field.ngllx; + const int npoints = ngllz * nglly * ngllx; + const int ndof_e = ncomp * npoints; + + const auto elements = + assembly_.element_types.get_elements_on_host(medium_tag); + + // Pass 1: per-row allocation upper bound. A row interacts with every dof + // of every element sharing its mesh point; per-element duplicates are + // inserted raw in pass 2 (fillComplete merges them), so the bound must -- + // and does exactly -- cover the raw insert count. + std::vector adjacent(dof_map_.nglob(), 0); + for (int i = 0; i < elements.size(); ++i) { + const int ispec = elements(i); + for (int iz = 0; iz < ngllz; ++iz) { + for (int iy = 0; iy < nglly; ++iy) { + for (int ix = 0; ix < ngllx; ++ix) { + ++adjacent[field.template get_iglob(ispec, iz, iy, + ix)]; + } + } + } + } + + std::vector entries_per_row( + static_cast(dof_map_.num_global_dofs()), 0); + for (int iglob = 0; iglob < dof_map_.nglob(); ++iglob) { + for (int icomp = 0; icomp < ncomp; ++icomp) { + entries_per_row[static_cast(dof_map_.gid(iglob, icomp))] = + adjacent[iglob] * static_cast(ndof_e); + } + } + + auto graph = Teuchos::rcp( + new crs_graph_type(dof_map_.overlap_map(), + Teuchos::ArrayView( + entries_per_row.data(), entries_per_row.size()))); + + // Pass 2: every dof of an element couples to every other dof of that + // element, so each element contributes its full column list to each of its + // rows. + for (int i = 0; i < elements.size(); ++i) { + const auto cols = element_column_gids(elements(i)); + for (int r = 0; r < ndof_e; ++r) { + graph->insertGlobalIndices(cols[r], ndof_e, cols.data()); + } + } + + graph->fillComplete(dof_map_.owned_map(), dof_map_.owned_map()); + return graph; +} + +template +void specfem::linear_system::StiffnessAssembler::fill_matrix( + crs_matrix_type &matrix) const { + const auto &field = assembly_.fields.template get_simulation_field< + specfem::simulation::field_type::forward>(); + const int npoints = field.ngllz * field.nglly * field.ngllx; + const int ndof_e = ncomp * npoints; + + const auto elements = + assembly_.element_types.get_elements_on_host(medium_tag); + const int nelements = elements.size(); + + // One block buffer reused across batches; only batch_size_ dense element + // blocks ever exist at a time -- no global dense matrix. + Kokkos::View + k_e("specfem::linear_system::element_stiffness_blocks", + std::min(batch_size_, std::max(nelements, 1)), ndof_e, ndof_e); + auto h_k_e = Kokkos::create_mirror_view(k_e); + + for (int offset = 0; offset < nelements; offset += batch_size_) { + const int batch_count = std::min(batch_size_, nelements - offset); + const auto batch = specfem::datatype::subview( + elements, Kokkos::pair(offset, offset + batch_count)); + + specfem::linear_system::compute_element_stiffness(assembly_, batch, + k_e); + Kokkos::deep_copy(h_k_e, k_e); + + for (int e = 0; e < batch_count; ++e) { + const int ispec = batch(e); + const auto cols = element_column_gids(ispec); + for (int r = 0; r < ndof_e; ++r) { + const int updated = matrix.sumIntoGlobalValues( + cols[r], ndof_e, &h_k_e(e, r, 0), cols.data()); + if (updated != ndof_e) { + std::ostringstream message; + message << "specfem::linear_system::StiffnessAssembler: row " + "scatter of element " + << ispec << " updated " << updated << " of " << ndof_e + << " entries; the graph does not match the element " + "connectivity."; + throw std::runtime_error(message.str()); + } + } + } + } +} + +template +Teuchos::RCP +specfem::linear_system::StiffnessAssembler::assemble() const { + const auto graph = build_graph(); + + // Matrix on the fill-complete static graph: values start at zero and only + // sumIntoGlobalValues is allowed, which is exactly what the scatter uses. + auto matrix = Teuchos::rcp(new crs_matrix_type(graph)); + + fill_matrix(*matrix); + + matrix->fillComplete(dof_map_.owned_map(), dof_map_.owned_map()); + + // MPI-ready seam: a distributed build fills the matrix on the overlap map + // and Export(ADD)s into the owned map here, reproducing the matrix-free + // assembly sum across ranks. Unreachable at one rank (DofMap rejects + // larger communicators), where overlap == owned and the matrix is final. + if (!dof_map_.overlap_map()->isSameAs(*dof_map_.owned_map())) { + throw std::runtime_error( + "specfem::linear_system::StiffnessAssembler: overlap map differs " + "from owned map, but the distributed Export(ADD) step is not " + "implemented yet (follow-up of issue #1982)."); + } + + return matrix; +} + +namespace specfem::linear_system_impl { +/// Tag bundle for the only combination explicitly instantiated for the +/// linear system (issue #1982). +using elastic_isotropic_tags = + specfem::tags::Tags; +} // namespace specfem::linear_system_impl + +// Explicit instantiation: 3D elastic isotropic +template class specfem::linear_system::StiffnessAssembler< + specfem::linear_system_impl::elastic_isotropic_tags>; + +#endif // SPECFEM_ENABLE_TRILINOS diff --git a/core/specfem/linear_system/tpetra_assembler.hpp b/core/specfem/linear_system/tpetra_assembler.hpp new file mode 100644 index 000000000..ebaac5d6e --- /dev/null +++ b/core/specfem/linear_system/tpetra_assembler.hpp @@ -0,0 +1,121 @@ +#pragma once + +#ifdef SPECFEM_ENABLE_TRILINOS + +#include "specfem/enums.hpp" +#include "specfem/linear_system/dof_map.hpp" +#include +#include +#include +#include + +namespace specfem { +namespace linear_system { + +/// Tpetra sparsity graph with the default ordinals and node type +using crs_graph_type = Tpetra::CrsGraph<>; + +/// Assembled sparse matrix type (see @ref scalar_type for the precision) +using crs_matrix_type = Tpetra::CrsMatrix; + +/** + * @brief Assembles the global stiffness matrix \f$ K \f$ of one medium as a + * `Tpetra::CrsMatrix` from dense element blocks. + * + * The assembled operator satisfies \f$ K u = \f$ internal force + * \f$ = -\mathrm{accel} \f$ of the matrix-free + * `compute_stiffness_interaction` kernel (before mass division) -- see + * @ref compute_element_stiffness for the sign convention. Assembly never + * materializes a global dense matrix: element blocks are computed in batches + * on the (Kokkos) device by the stiffness probe kernel, mirrored to the + * host, and scattered into the sparse matrix with batched row updates. + * + * One assembler produces the matrix block of one medium + * (`Tags::medium_tag`); a future multi-medium system holds one assembler and + * one matrix per medium. This milestone additionally requires the mesh to be + * single-medium (fluid-solid coupling blocks are deferred) and serial (see + * @ref DofMap). + * + * @tparam Tags Compile-time tags (dimension, medium, property, attenuation); + * only `dim3, elastic, isotropic, none` is instantiated + */ +template class StiffnessAssembler { +public: + constexpr static auto dimension_tag = Tags::dimension_tag; + constexpr static auto medium_tag = Tags::medium_tag; + + static_assert(dimension_tag == specfem::element::dimension_tag::dim3, + "StiffnessAssembler takes a dim3 assembly; Tags must be a " + "dim3 bundle."); + + /// Field components per mesh point of the medium + constexpr static int ncomp = + specfem::element::attributes::components; + + /** + * @brief Elements whose stiffness blocks are formed per probe-kernel + * launch. Bounds the block buffer: `batch * ndof_e^2` scalars (~36 MB for + * 64 elastic NGLL = 5 elements in single precision). + */ + constexpr static int default_batch_size = 64; + + using AssemblyType = specfem::assembly::assembly; + + /** + * @brief Validate scope and set up the dof map. + * + * Throws `std::runtime_error` if any element is outside the supported + * scope (see @ref validate_stiffness_scope), if the mesh contains elements + * of a medium other than `Tags::medium_tag` (single-medium milestone; + * coupling blocks are deferred), or if the communicator has more than one + * rank (see @ref DofMap). + * + * @param assembly Assembled mesh, jacobian matrix, material properties, + * and fields; must outlive the assembler + * @param batch_size Elements per probe-kernel launch (>= 1) + */ + StiffnessAssembler(const AssemblyType &assembly, + const int batch_size = default_batch_size); + + /** + * @brief Assemble the stiffness matrix. + * + * Pipeline: build the `Tpetra::CrsGraph` from element connectivity (two + * host passes over the index mapping), construct the matrix on the static + * graph, fill it batch-by-batch with `sumIntoGlobalValues` row updates + * from the probed element blocks, and `fillComplete()`. Row/column ids + * follow @ref DofMap::gid. + * + * @return Fill-complete stiffness matrix on the owned map + */ + Teuchos::RCP assemble() const; + + /// Dof map shared by the matrix and any right-hand-side/solution vectors + const DofMap &dof_map() const { return dof_map_; } + +private: + /// Build and fill-complete the sparsity graph from element connectivity + Teuchos::RCP build_graph() const; + + /// Probe element blocks in batches and scatter them into the matrix + void fill_matrix(crs_matrix_type &matrix) const; + + /** + * @brief Global column ids of one element in element-local dof order. + * + * Single source of truth for the ldof <-> gid correspondence used by both + * the graph build and the block scatter; entry `ldof` (see + * @ref local_dof_index) holds `dof_map_.gid(iglob(ispec, iz, iy, ix), + * icomp)`. + */ + std::vector element_column_gids(const int ispec) const; + + const AssemblyType &assembly_; ///< Borrowed assembly (not owned) + DofMap dof_map_; ///< Per-medium dof numbering and maps + int batch_size_; ///< Elements per probe-kernel launch +}; + +} // namespace linear_system +} // namespace specfem + +#endif // SPECFEM_ENABLE_TRILINOS diff --git a/docs/Doxyfile.in b/docs/Doxyfile.in index 0ba2edc9b..eae18c4df 100644 --- a/docs/Doxyfile.in +++ b/docs/Doxyfile.in @@ -2292,7 +2292,8 @@ INCLUDE_FILE_PATTERNS = PREDEFINED = KOKKOS_INLINE_FUNCTION= \ KOKKOS_FORCEINLINE_FUNCTION= \ - KOKKOS_FUNCTION= + KOKKOS_FUNCTION= \ + SPECFEM_ENABLE_TRILINOS diff --git a/docs/sections/api/specfem/linear_system/index.rst b/docs/sections/api/specfem/linear_system/index.rst index 81c420780..9622e3139 100644 --- a/docs/sections/api/specfem/linear_system/index.rst +++ b/docs/sections/api/specfem/linear_system/index.rst @@ -4,11 +4,23 @@ ========================== Utilities for assembling the spectral-element operator into an explicit -linear system (issue #1982). Currently provides dense element stiffness -extraction -- probing the matrix-free element operator with local unit -vectors -- for the 3D elastic isotropic medium, plus a scope validator that -rejects meshes outside the supported tag combination (attenuation, Stacey -boundaries). +linear system (issue #1982). Provides dense element stiffness extraction -- +probing the matrix-free element operator with local unit vectors -- for the +3D elastic isotropic medium, plus a scope validator that rejects meshes +outside the supported tag combination (attenuation, Stacey boundaries). + +When SPECFEM++ is built with Trilinos (``SPECFEM_ENABLE_TRILINOS=ON``), the +module additionally provides ``DofMap`` -- the per-medium mapping from +SPECFEM++ ``(iglob, icomp)`` degrees of freedom to Tpetra global ids, with +component-blocked layout ``gid = icomp * nglob + iglob`` matching field +storage -- and ``StiffnessAssembler``, which assembles the global stiffness +matrix :math:`K` of one medium as a ``Tpetra::CrsMatrix`` by scattering +batched dense element blocks into a ``Tpetra::CrsGraph`` built from element +connectivity. The assembled operator satisfies :math:`K u =` internal force +:math:`= -\mathrm{accel}` of the matrix-free +``compute_stiffness_interaction`` kernel (before mass division). Assembly is +serial-only in this milestone; the owned/overlap map split in ``DofMap`` +keeps the API ready for distributed Export(ADD) assembly. .. doxygennamespace:: specfem::linear_system :members: diff --git a/tests/unit-tests/linear_system/stiffness_assembler_tests.cpp b/tests/unit-tests/linear_system/stiffness_assembler_tests.cpp new file mode 100644 index 000000000..018304ce6 --- /dev/null +++ b/tests/unit-tests/linear_system/stiffness_assembler_tests.cpp @@ -0,0 +1,306 @@ +#include "../SPECFEM_Environment.hpp" +#include + +#ifdef SPECFEM_ENABLE_TRILINOS + +#include "specfem/assembly/assembly.hpp" +#include "specfem/compute/impl/compute_stiffness_interaction.hpp" +#include "specfem/io.hpp" +#include "specfem/linear_system/tpetra_assembler.hpp" +#include "specfem/mesh.hpp" +#include "specfem/quadrature.hpp" +#include "specfem/runtime_configuration.hpp" +#include "specfem/tags.hpp" +#include +#include +#include +#include +#include +#include +#include + +namespace stiffness_assembler_test { + +constexpr auto dim3_tag = specfem::element::dimension_tag::dim3; +constexpr auto elastic_tag = specfem::element::medium_tag::elastic; +constexpr auto forward_tag = specfem::simulation::field_type::forward; +constexpr int NGLL = 5; +constexpr int ncomp = 3; + +constexpr bool single_precision = sizeof(type_real) == sizeof(float); + +using AssemblyType = specfem::assembly::assembly; +using StiffnessTags = + specfem::tags::Tags; +using AssemblerType = specfem::linear_system::StiffnessAssembler; +using VectorType = + Tpetra::Vector; + +// Build a full assembly from a Newmark displacement-test dataset. Paths are +// relative to TEST_OUTPUT_DIR, where the displacement_tests data tree is +// linked (see SERIAL_LINK_DIRS in serial.cmake). +std::unique_ptr build_assembly_3d(const std::string &test_name) { + const std::string test_path = + "displacement_tests/Newmark/serial/dim3/" + test_name; + + specfem::runtime_configuration::setup setup(test_path + + "/specfem_config.yaml"); + + const auto database_filename = setup.get_databases(); + const auto &source_entries = setup.get_source_entries(); + const auto stations_node = setup.get_stations(); + const auto quadratures = setup.instantiate_quadrature(); + + auto mesh = specfem::io::read_3d_mesh(database_filename, + setup.get_attenuation_setup()); + + const type_real dt = setup.get_dt(); + const int nsteps = setup.get_nsteps(); + + auto [sources, t0, starttime] = specfem::io::read_sources( + source_entries, nsteps, setup.get_t0(), dt, setup.get_simulation_type()); + (void)starttime; + setup.update_t0(t0); + + auto receivers = specfem::io::read_3d_receivers(stations_node); + + return std::make_unique( + mesh, quadratures, sources, receivers, setup.get_seismogram_types(), + setup.get_t0(), dt, nsteps, setup.get_max_seismogram_step(), + setup.get_nstep_between_samples(), setup.get_simulation_type(), + setup.allocate_boundary_values(), setup.instantiate_property_reader()); +} + +// Shared fixture: assemble the matrix of the small homogeneous elastic mesh +// once for the whole suite; assembly happens on the first access so that a +// throwing constructor fails inside a test body, not in SetUpTestSuite. +class StiffnessAssembler3D : public ::testing::Test { +protected: + static void TearDownTestSuite() { + matrix_ = Teuchos::null; + assembler_.reset(); + delete assembly_; + assembly_ = nullptr; + } + + static AssemblyType &assembly() { + if (assembly_ == nullptr) { + assembly_ = build_assembly_3d("HomogeneousHalfspaceSmallNoABCForceSource") + .release(); + } + return *assembly_; + } + + static AssemblerType &assembler() { + if (!assembler_) { + assembler_ = std::make_unique(assembly()); + } + return *assembler_; + } + + static Teuchos::RCP matrix() { + if (matrix_.is_null()) { + matrix_ = assembler().assemble(); + } + return matrix_; + } + + static AssemblyType *assembly_; + static std::unique_ptr assembler_; + static Teuchos::RCP matrix_; +}; + +AssemblyType *StiffnessAssembler3D::assembly_ = nullptr; +std::unique_ptr StiffnessAssembler3D::assembler_; +Teuchos::RCP + StiffnessAssembler3D::matrix_; + +// Largest absolute matrix entry -- the natural scale for tolerances. +type_real max_abs_entry( + const Teuchos::RCP &matrix) { + const auto values = matrix->getLocalMatrixHost().values; + type_real scale = 0; + for (std::size_t k = 0; k < values.extent(0); ++k) { + scale = std::max(scale, std::abs(values(k))); + } + return scale; +} + +TEST_F(StiffnessAssembler3D, GlobalDimensionsMatchDofMap) { + const auto matrix = this->matrix(); + const auto &dof_map = assembler().dof_map(); + + const auto num_dofs = + static_cast(dof_map.num_global_dofs()); + EXPECT_EQ(num_dofs, + static_cast(ncomp) * dof_map.nglob()); + EXPECT_EQ(matrix->getGlobalNumRows(), num_dofs); + EXPECT_EQ(matrix->getGlobalNumCols(), num_dofs); + EXPECT_TRUE(matrix->isFillComplete()); + EXPECT_GT(matrix->getGlobalNumEntries(), 0u); + EXPECT_GT(max_abs_entry(matrix), static_cast(0)); +} + +TEST_F(StiffnessAssembler3D, SymmetricBilinearForm) { + const auto matrix = this->matrix(); + const auto map = assembler().dof_map().owned_map(); + + // K is symmetric, so x' K z == z' K x for any x, z. The bilinear form + // probes symmetry of the whole assembled matrix at matrix-vector cost. + std::mt19937 generator(2026); + std::uniform_real_distribution distribution(-1, 1); + + VectorType x(map), z(map), k_x(map), k_z(map); + { + auto x_view = x.getLocalViewHost(Tpetra::Access::OverwriteAll); + auto z_view = z.getLocalViewHost(Tpetra::Access::OverwriteAll); + for (std::size_t i = 0; i < x_view.extent(0); ++i) { + x_view(i, 0) = distribution(generator); + z_view(i, 0) = distribution(generator); + } + } + + matrix->apply(x, k_x); + matrix->apply(z, k_z); + + const type_real x_k_z = static_cast(x.dot(k_z)); + const type_real z_k_x = static_cast(z.dot(k_x)); + + const type_real scale = std::max(std::max(std::abs(x_k_z), std::abs(z_k_x)), + max_abs_entry(matrix)); + const type_real tol = (single_precision ? static_cast(1e-3) + : static_cast(1e-10)) * + scale; + EXPECT_NEAR(x_k_z, z_k_x, tol); +} + +TEST_F(StiffnessAssembler3D, RigidBodyTranslationNullspace) { + const auto matrix = this->matrix(); + const auto &dof_map = assembler().dof_map(); + const auto map = dof_map.owned_map(); + + const type_real scale = max_abs_entry(matrix); + ASSERT_GT(scale, static_cast(0)); + + // With natural boundary conditions only, a rigid translation of one + // component produces zero strain, so K annihilates it. Rows sum over up to + // ~8 * 375 float entries, hence the loose relative tolerance. + const type_real tol = (single_precision ? static_cast(5e-3) + : static_cast(1e-10)) * + scale; + + for (int icomp = 0; icomp < ncomp; ++icomp) { + VectorType translation(map), response(map); + { + auto view = translation.getLocalViewHost(Tpetra::Access::OverwriteAll); + for (int iglob = 0; iglob < dof_map.nglob(); ++iglob) { + view(static_cast(dof_map.gid(iglob, icomp)), 0) = 1; + } + } + + matrix->apply(translation, response); + EXPECT_LE(response.normInf(), tol) + << "translation nullspace violated for component " << icomp; + } +} + +TEST_F(StiffnessAssembler3D, MatchesMatrixFreeOperatorGlobally) { + const auto matrix = this->matrix(); + const auto &dof_map = assembler().dof_map(); + + auto &field = assembly().fields.template get_simulation_field(); + const auto &field_impl = field.template get_field(); + const auto h_u = field_impl.get_host_field(); + const auto h_v = field_impl.get_host_field_dot(); + const auto h_a = field_impl.get_host_field_dot_dot(); + const int nglob = field_impl.nglob; + ASSERT_EQ(nglob, dof_map.nglob()); + + // Random displacement over the whole mesh; velocity and acceleration zero. + std::mt19937 generator(54321); + std::uniform_real_distribution distribution(-1, 1); + for (int iglob = 0; iglob < nglob; ++iglob) { + for (int icomp = 0; icomp < ncomp; ++icomp) { + h_u(iglob, icomp) = distribution(generator); + h_v(iglob, icomp) = 0; + h_a(iglob, icomp) = 0; + } + } + + // Production matrix-free operator (no mass division): accel = -K u. + assembly().fields.copy_to_device(); + using BaseTags = specfem::tags::Tags; + specfem::compute::impl::compute_stiffness_interaction< + NGLL, specfem::tags::expand>( + assembly(), 0); + specfem::compute::impl::compute_stiffness_interaction< + NGLL, specfem::tags::expand>( + assembly(), 0); + assembly().fields.copy_to_host(); + + // Assembled operator applied to the same displacement. + VectorType u(dof_map.owned_map()), k_u(dof_map.owned_map()); + { + auto view = u.getLocalViewHost(Tpetra::Access::OverwriteAll); + for (int iglob = 0; iglob < nglob; ++iglob) { + for (int icomp = 0; icomp < ncomp; ++icomp) { + view(static_cast(dof_map.gid(iglob, icomp)), 0) = + h_u(iglob, icomp); + } + } + } + matrix->apply(u, k_u); + + type_real scale = 0; + type_real max_diff = 0; + { + auto view = k_u.getLocalViewHost(Tpetra::Access::ReadOnly); + for (int iglob = 0; iglob < nglob; ++iglob) { + for (int icomp = 0; icomp < ncomp; ++icomp) { + const type_real expected = -h_a(iglob, icomp); + const type_real actual = + view(static_cast(dof_map.gid(iglob, icomp)), 0); + scale = std::max(scale, std::abs(expected)); + max_diff = std::max(max_diff, std::abs(expected - actual)); + } + } + } + ASSERT_GT(scale, static_cast(0)); + + const type_real rel_tol = single_precision ? static_cast(2e-3) + : static_cast(1e-10); + EXPECT_LE(max_diff, rel_tol * scale) + << "assembled K u disagrees with the matrix-free operator"; +} + +TEST(StiffnessAssemblerScope3D, RejectsMultiMediumMeshes) { + const auto mixed_assembly = build_assembly_3d("AcousticElasticForce"); + EXPECT_THROW(AssemblerType assembler(*mixed_assembly), std::runtime_error); +} + +TEST(StiffnessAssemblerScope3D, RejectsStaceyBoundaries) { + const auto stacey_assembly = build_assembly_3d("HomogeneousHalfSpaceStacey"); + EXPECT_THROW(AssemblerType assembler(*stacey_assembly), std::runtime_error); +} + +} // namespace stiffness_assembler_test + +#else // !SPECFEM_ENABLE_TRILINOS + +TEST(StiffnessAssembler3D, SkippedWithoutTrilinos) { + GTEST_SKIP() << "SPECFEM++ was built without Trilinos " + "(SPECFEM_ENABLE_TRILINOS=OFF); the stiffness assembler is " + "unavailable."; +} + +#endif // SPECFEM_ENABLE_TRILINOS + +int main(int argc, char *argv[]) { + ::testing::InitGoogleTest(&argc, argv); + ::testing::AddGlobalTestEnvironment(new SPECFEMEnvironment); + return RUN_ALL_TESTS(); +} diff --git a/tests/unit-tests/serial.cmake b/tests/unit-tests/serial.cmake index 010292934..9a4b1103c 100644 --- a/tests/unit-tests/serial.cmake +++ b/tests/unit-tests/serial.cmake @@ -979,6 +979,29 @@ target_link_libraries( -lpthread -lm ) +add_executable( + stiffness_assembler_tests + linear_system/stiffness_assembler_tests.cpp +) + +target_link_libraries( + stiffness_assembler_tests + specfem::linear_system + specfem::quadrature + specfem::mesh + yaml-cpp + specfem_environment + specfem::assembly + specfem::runtime_configuration + timescheme + point + specfem::algorithms + specfem::solver + specfem::periodic_tasks + ${BOOST_LIBS} + -lpthread -lm +) + # Register serial tests for discovery set(SERIAL_TEST_TARGETS serial_mpi_tests @@ -1033,6 +1056,7 @@ set(SERIAL_TEST_TARGETS surface_elevation_tests trilinos_smoke_tests element_stiffness_tests + stiffness_assembler_tests ) if (NOT SPECFEM_ENABLE_MPI)