Cabana_Grid_LocalMesh.hpp

Go to the documentation of this file.

/****************************************************************************
 * Copyright (c) 2018-2023 by the Cabana authors                            *
 * All rights reserved.                                                     *
 *                                                                          *
 * This file is part of the Cabana library. Cabana is distributed under a   *
 * BSD 3-clause license. For the licensing terms see the LICENSE file in    *
 * the top-level directory.                                                 *
 *                                                                          *
 * SPDX-License-Identifier: BSD-3-Clause                                    *
 ****************************************************************************/

#ifndef CABANA_GRID_LOCALMESH_HPP
#define CABANA_GRID_LOCALMESH_HPP
 
#include <Cabana_Grid_LocalGrid.hpp>
#include <Cabana_Grid_Types.hpp>
 
#include <Cabana_Utils.hpp> // FIXME: remove after next release.
 
#include <Kokkos_Core.hpp>
 
namespace Cabana
{
namespace Grid
{
//---------------------------------------------------------------------------//
// Forward declaration of local mesh.
template <class MemorySpace, class MeshType>
class LocalMesh;
 
//---------------------------------------------------------------------------//
template <class Scalar, class MemorySpace, std::size_t NumSpaceDim>
class LocalMesh<MemorySpace, UniformMesh<Scalar, NumSpaceDim>>
{
  public:
    using mesh_type = UniformMesh<Scalar, NumSpaceDim>;

    using scalar_type = Scalar;

    static constexpr std::size_t num_space_dim = NumSpaceDim;

    using memory_space = MemorySpace;
    static_assert( Kokkos::is_memory_space<MemorySpace>() );

    using execution_space = typename memory_space::execution_space;

    LocalMesh() = default;

    LocalMesh( const LocalGrid<UniformMesh<Scalar, num_space_dim>>& local_grid )
    {
        const auto& global_grid = local_grid.globalGrid();
        const auto& global_mesh = global_grid.globalMesh();
 
        // Get the cell size.
        for ( std::size_t d = 0; d < num_space_dim; ++d )
            _cell_size[d] = global_mesh.cellSize( d );
 
        // Compute face area.
        if ( 3 == num_space_dim )
        {
            _face_area[Dim::I] = _cell_size[Dim::J] * _cell_size[Dim::K];
            _face_area[Dim::J] = _cell_size[Dim::I] * _cell_size[Dim::K];
            _face_area[Dim::K] = _cell_size[Dim::I] * _cell_size[Dim::J];
        }
        else if ( 2 == num_space_dim )
        {
            _face_area[Dim::I] = _cell_size[Dim::J];
            _face_area[Dim::J] = _cell_size[Dim::I];
        }
 
        // Compute cell volume.
        _cell_volume = 1.0;
        for ( std::size_t d = 0; d < num_space_dim; ++d )
            _cell_volume *= _cell_size[d];
 
        // Compute the owned low corner
        for ( std::size_t d = 0; d < num_space_dim; ++d )
            _own_low_corner[d] = global_mesh.lowCorner( d ) +
                                 _cell_size[d] * global_grid.globalOffset( d );
 
        // Compute the owned high corner
        for ( std::size_t d = 0; d < num_space_dim; ++d )
            _own_high_corner[d] =
                global_mesh.lowCorner( d ) +
                _cell_size[d] * ( global_grid.globalOffset( d ) +
                                  global_grid.ownedNumCell( d ) );
 
        // Compute the ghosted low corner
        for ( std::size_t d = 0; d < num_space_dim; ++d )
            _ghost_low_corner[d] = lowCorner( Own(), d ) -
                                   local_grid.haloCellWidth() * _cell_size[d];
 
        // Compute the ghosted high corner
        for ( std::size_t d = 0; d < num_space_dim; ++d )
            _ghost_high_corner[d] = highCorner( Own(), d ) +
                                    local_grid.haloCellWidth() * _cell_size[d];
 
        // Periodicity
        for ( std::size_t d = 0; d < num_space_dim; ++d )
            _periodic[d] = global_grid.isPeriodic( d );
 
        // Determine if a block is on the low or high boundaries.
        for ( std::size_t d = 0; d < num_space_dim; ++d )
        {
            _boundary_lo[d] = global_grid.onLowBoundary( d );
            _boundary_hi[d] = global_grid.onHighBoundary( d );
        }
    }

    KOKKOS_INLINE_FUNCTION
    bool isPeriodic( const int dim ) const { return _periodic[dim]; }

    KOKKOS_INLINE_FUNCTION
    bool onLowBoundary( const int dim ) const { return _boundary_lo[dim]; }

    KOKKOS_INLINE_FUNCTION
    bool onHighBoundary( const int dim ) const { return _boundary_hi[dim]; }

    KOKKOS_INLINE_FUNCTION
    Scalar lowCorner( Own, const int dim ) const
    {
        return _own_low_corner[dim];
    }

    KOKKOS_INLINE_FUNCTION
    Scalar lowCorner( Ghost, const int dim ) const
    {
        return _ghost_low_corner[dim];
    }

    KOKKOS_INLINE_FUNCTION
    Scalar highCorner( Own, const int dim ) const
    {
        return _own_high_corner[dim];
    }

    KOKKOS_INLINE_FUNCTION
    Scalar highCorner( Ghost, const int dim ) const
    {
        return _ghost_high_corner[dim];
    }

    template <typename Decomposition>
    KOKKOS_FUNCTION Scalar extent( Decomposition d, const int dim ) const
    {
        return highCorner( d, dim ) - lowCorner( d, dim );
    }

    template <class Integer>
    KOKKOS_INLINE_FUNCTION void coordinates( Cell,
                                             const Integer index[num_space_dim],
                                             Scalar x[num_space_dim] ) const
    {
        for ( std::size_t d = 0; d < num_space_dim; ++d )
            x[d] = _ghost_low_corner[d] +
                   ( Scalar( index[d] ) + Scalar( 0.5 ) ) * _cell_size[d];
    }

    template <class Integer>
    KOKKOS_INLINE_FUNCTION void coordinates( Node,
                                             const Integer index[num_space_dim],
                                             Scalar x[num_space_dim] ) const
    {
        for ( std::size_t d = 0; d < num_space_dim; ++d )
            x[d] = _ghost_low_corner[d] + Scalar( index[d] ) * _cell_size[d];
    }

    template <int Dir, class Integer>
    KOKKOS_INLINE_FUNCTION void coordinates( Face<Dir>,
                                             const Integer index[num_space_dim],
                                             Scalar x[num_space_dim] ) const
    {
        static_assert( Dir < num_space_dim, "Face dimension out of bounds" );
        for ( std::size_t d = 0; d < num_space_dim; ++d )
            if ( Dir == d )
                x[d] =
                    _ghost_low_corner[d] + Scalar( index[d] ) * _cell_size[d];
            else
                x[d] = _ghost_low_corner[d] +
                       ( Scalar( index[d] ) + Scalar( 0.5 ) ) * _cell_size[d];
    }

    template <int Dir, class Integer, std::size_t NSD = num_space_dim>
    KOKKOS_INLINE_FUNCTION std::enable_if_t<3 == NSD, void>
    coordinates( Edge<Dir>, const Integer index[3], Scalar x[3] ) const
    {
        for ( std::size_t d = 0; d < 3; ++d )
            if ( Dir == d )
                x[d] = _ghost_low_corner[d] +
                       ( Scalar( index[d] ) + Scalar( 0.5 ) ) * _cell_size[d];
            else
                x[d] =
                    _ghost_low_corner[d] + Scalar( index[d] ) * _cell_size[d];
    }

    template <class Integer>
    KOKKOS_INLINE_FUNCTION Scalar measure( Node,
                                           const Integer[num_space_dim] ) const
    {
        return 0.0;
    }

    template <int Dir, class Integer, std::size_t NSD = num_space_dim>
    KOKKOS_INLINE_FUNCTION std::enable_if_t<3 == NSD, Scalar>
    measure( Edge<Dir>, const Integer[3] ) const
    {
        return _cell_size[Dir];
    }

    template <int Dir, class Integer>
    KOKKOS_INLINE_FUNCTION Scalar measure( Face<Dir>,
                                           const Integer[num_space_dim] ) const
    {
        static_assert( Dir < num_space_dim, "Face dimension out of bounds" );
        return _face_area[Dir];
    }

    template <class Integer>
    KOKKOS_INLINE_FUNCTION Scalar measure( Cell,
                                           const Integer[num_space_dim] ) const
    {
        return _cell_volume;
    }
 
  private:
    Kokkos::Array<Scalar, num_space_dim> _cell_size;
    Kokkos::Array<Scalar, num_space_dim> _face_area;
    Scalar _cell_volume;
    Kokkos::Array<Scalar, num_space_dim> _own_low_corner;
    Kokkos::Array<Scalar, num_space_dim> _own_high_corner;
    Kokkos::Array<Scalar, num_space_dim> _ghost_low_corner;
    Kokkos::Array<Scalar, num_space_dim> _ghost_high_corner;
    Kokkos::Array<bool, num_space_dim> _periodic;
    Kokkos::Array<bool, num_space_dim> _boundary_lo;
    Kokkos::Array<bool, num_space_dim> _boundary_hi;
};
 
//---------------------------------------------------------------------------//
template <class Scalar, class MemorySpace, std::size_t NumSpaceDim>
class LocalMesh<MemorySpace, NonUniformMesh<Scalar, NumSpaceDim>>
{
  public:
    using mesh_type = NonUniformMesh<Scalar, NumSpaceDim>;

    using scalar_type = Scalar;

    static constexpr std::size_t num_space_dim = NumSpaceDim;

    using memory_space = MemorySpace;
    static_assert( Kokkos::is_memory_space<MemorySpace>() );

    using execution_space = typename memory_space::execution_space;

    LocalMesh(
        const LocalGrid<NonUniformMesh<Scalar, num_space_dim>>& local_grid )
    {
        const auto& global_grid = local_grid.globalGrid();
        const auto& global_mesh = global_grid.globalMesh();
 
        // Compute the owned low corner.
        for ( std::size_t d = 0; d < num_space_dim; ++d )
            _own_low_corner[d] =
                global_mesh.nonUniformEdge( d )[global_grid.globalOffset( d )];
 
        // Compute the owned high corner
        for ( std::size_t d = 0; d < num_space_dim; ++d )
            _own_high_corner[d] =
                global_mesh.nonUniformEdge( d )[global_grid.globalOffset( d ) +
                                                global_grid.ownedNumCell( d )];
 
        // Compute the ghosted low corner
        for ( std::size_t d = 0; d < num_space_dim; ++d )
        {
            // Interior domain case.
            if ( !global_grid.onLowBoundary( d ) )
            {
                _ghost_low_corner[d] = global_mesh.nonUniformEdge(
                    d )[global_grid.globalOffset( d ) -
                        local_grid.haloCellWidth()];
            }
 
            // Periodic boundary. Use cells on other side of boundary to
            // generate geometry.
            else if ( global_grid.isPeriodic( d ) )
            {
                int nedge = global_mesh.nonUniformEdge( d ).size();
                _ghost_low_corner[d] =
                    global_mesh.nonUniformEdge( d ).front() -
                    ( global_mesh.nonUniformEdge( d ).back() -
                      global_mesh.nonUniformEdge(
                          d )[nedge - local_grid.haloCellWidth() - 1] );
            }
 
            // In the non-periodic boundary case we extrapolate halo cells to
            // have the same width as the boundary cell.
            else
            {
                Scalar dx = global_mesh.nonUniformEdge( d )[1] -
                            global_mesh.nonUniformEdge( d )[0];
                _ghost_low_corner[d] = global_mesh.nonUniformEdge( d ).front() -
                                       dx * local_grid.haloCellWidth();
            }
        }
 
        // Compute the ghosted high corner
        for ( std::size_t d = 0; d < num_space_dim; ++d )
        {
            // Interior domain case.
            if ( !global_grid.onHighBoundary( d ) )
            {
                _ghost_high_corner[d] = global_mesh.nonUniformEdge(
                    d )[global_grid.globalOffset( d ) +
                        global_grid.ownedNumCell( d ) +
                        local_grid.haloCellWidth()];
            }
 
            // Periodic boundary. Use cells on other side of boundary to
            // generate geometry.
            else if ( global_grid.isPeriodic( d ) )
            {
                _ghost_high_corner[d] =
                    global_mesh.nonUniformEdge( d ).back() +
                    ( global_mesh.nonUniformEdge(
                          d )[local_grid.haloCellWidth()] -
                      global_mesh.nonUniformEdge( d ).front() );
            }
 
            // In the non-periodic boundary case we extrapolate halo cells to
            // have the same width as the boundary cell.
            else
            {
                int nedge = global_mesh.nonUniformEdge( d ).size();
                Scalar dx = global_mesh.nonUniformEdge( d )[nedge - 1] -
                            global_mesh.nonUniformEdge( d )[nedge - 2];
                _ghost_high_corner[d] = global_mesh.nonUniformEdge( d ).back() +
                                        dx * local_grid.haloCellWidth();
            }
        }
 
        // Get the node index spaces.
        auto owned_nodes_local =
            local_grid.indexSpace( Own(), Node(), Local() );
        auto ghosted_nodes_local =
            local_grid.indexSpace( Ghost(), Node(), Local() );
        auto owned_nodes_global =
            local_grid.indexSpace( Own(), Node(), Global() );
        for ( std::size_t d = 0; d < num_space_dim; ++d )
        {
            // Allocate edges on the device for this dimension.
            const auto& global_edge = global_mesh.nonUniformEdge( d );
            int nedge = ghosted_nodes_local.extent( d );
            int nedge_global = global_edge.size();
            _local_edges[d] = Kokkos::View<Scalar*, MemorySpace>(
                Kokkos::ViewAllocateWithoutInitializing( "local_edges" ),
                nedge );
 
            // Compute edges on the host.
            auto edge_mirror = Kokkos::create_mirror_view( Kokkos::HostSpace(),
                                                           _local_edges[d] );
 
            // Compute the owned edges.
            for ( int n = owned_nodes_local.min( d );
                  n < owned_nodes_local.max( d ); ++n )
            {
                edge_mirror( n ) = global_edge[owned_nodes_global.min( d ) + n -
                                               owned_nodes_local.min( d )];
            }
 
            // Compute the lower boundary edges.
            if ( !global_grid.onLowBoundary( d ) )
            {
                // Interior block gets edges from neighbors.
                for ( int n = 0; n < owned_nodes_local.min( d ); ++n )
                {
                    edge_mirror( n ) =
                        global_edge[owned_nodes_global.min( d ) + n -
                                    owned_nodes_local.min( d )];
                }
            }
            else if ( global_grid.isPeriodic( d ) )
            {
                // Periodic boundary block gets edges from neighbor on
                // opposite side of boundary.
                for ( int n = 0; n < owned_nodes_local.min( d ); ++n )
                {
                    edge_mirror( n ) =
                        global_edge.front() - global_edge.back() +
                        global_edge[global_edge.size() - 1 -
                                    local_grid.haloCellWidth() + n];
                }
            }
            else
            {
                // Non-periodic boundary block extrapolates edges using
                // boundary cell width.
                for ( int n = 0; n < owned_nodes_local.min( d ); ++n )
                {
                    Scalar dx = global_edge[1] - global_edge[0];
                    edge_mirror( n ) = global_edge.front() -
                                       ( owned_nodes_local.min( d ) - n ) * dx;
                }
            }
 
            // Compute the upper boundary edges.
            if ( !global_grid.onHighBoundary( d ) )
            {
                // Interior block gets edges from neighbors.
                for ( int n = owned_nodes_local.max( d );
                      n < ghosted_nodes_local.max( d ); ++n )
                {
                    edge_mirror( n ) =
                        global_edge[owned_nodes_global.min( d ) + n -
                                    owned_nodes_local.min( d )];
                }
            }
            else if ( global_grid.isPeriodic( d ) )
            {
                // Periodic boundary block gets edges from neighbor on
                // opposite side of boundary.
                for ( int n = 0; n < ghosted_nodes_local.max( d ) -
                                         owned_nodes_local.max( d );
                      ++n )
                {
                    edge_mirror( owned_nodes_local.max( d ) + n ) =
                        global_edge.back() + global_edge[n] -
                        global_edge.front();
                }
            }
            else
            {
                // Non-periodic boundary block extrapolates edges using
                // boundary cell width.
                for ( int n = owned_nodes_local.max( d );
                      n < ghosted_nodes_local.max( d ); ++n )
                {
                    Scalar dx = global_edge[nedge_global - 1] -
                                global_edge[nedge_global - 2];
                    edge_mirror( n ) =
                        global_edge.back() +
                        ( n - owned_nodes_local.max( d ) + 1 ) * dx;
                }
            }
 
            // Copy edges to the device.
            Kokkos::deep_copy( _local_edges[d], edge_mirror );
        }
 
        // Periodicity
        for ( std::size_t d = 0; d < num_space_dim; ++d )
            _periodic[d] = global_grid.isPeriodic( d );
 
        // Determine if a block is on the low or high boundaries.
        for ( std::size_t d = 0; d < num_space_dim; ++d )
        {
            _boundary_lo[d] = global_grid.onLowBoundary( d );
            _boundary_hi[d] = global_grid.onHighBoundary( d );
        }
    }

    KOKKOS_INLINE_FUNCTION
    bool isPeriodic( const int dim ) const { return _periodic[dim]; }

    KOKKOS_INLINE_FUNCTION
    bool onLowBoundary( const int dim ) const { return _boundary_lo[dim]; }

    KOKKOS_INLINE_FUNCTION
    bool onHighBoundary( const int dim ) const { return _boundary_hi[dim]; }

    KOKKOS_INLINE_FUNCTION
    Scalar lowCorner( Own, const int dim ) const
    {
        return _own_low_corner[dim];
    }

    KOKKOS_INLINE_FUNCTION
    Scalar lowCorner( Ghost, const int dim ) const
    {
        return _ghost_low_corner[dim];
    }

    KOKKOS_INLINE_FUNCTION
    Scalar highCorner( Own, const int dim ) const
    {
        return _own_high_corner[dim];
    }

    KOKKOS_INLINE_FUNCTION
    Scalar highCorner( Ghost, const int dim ) const
    {
        return _ghost_high_corner[dim];
    }

    template <typename Decomposition>
    KOKKOS_FUNCTION Scalar extent( Decomposition d, const int dim ) const
    {
        return highCorner( d, dim ) - lowCorner( d, dim );
    }

    template <class Integer>
    KOKKOS_INLINE_FUNCTION void coordinates( Cell,
                                             const Integer index[num_space_dim],
                                             Scalar x[num_space_dim] ) const
    {
        for ( std::size_t d = 0; d < num_space_dim; ++d )
            x[d] = ( _local_edges[d]( index[d] + 1 ) +
                     _local_edges[d]( index[d] ) ) /
                   Scalar( 2.0 );
    }

    template <class Integer>
    KOKKOS_INLINE_FUNCTION void coordinates( Node,
                                             const Integer index[num_space_dim],
                                             Scalar x[num_space_dim] ) const
    {
        for ( std::size_t d = 0; d < num_space_dim; ++d )
            x[d] = _local_edges[d]( index[d] );
    }

    template <int Dir, class Integer>
    KOKKOS_INLINE_FUNCTION void coordinates( Face<Dir>,
                                             const Integer index[num_space_dim],
                                             Scalar x[num_space_dim] ) const
    {
        static_assert( Dir < num_space_dim, "Face dimension out of bounds" );
 
        for ( std::size_t d = 0; d < num_space_dim; ++d )
            if ( Dir == d )
                x[d] = _local_edges[d]( index[d] );
            else
                x[d] = ( _local_edges[d]( index[d] + 1 ) +
                         _local_edges[d]( index[d] ) ) /
                       Scalar( 2.0 );
    }

    template <int Dir, class Integer, std::size_t NSD = num_space_dim>
    KOKKOS_INLINE_FUNCTION std::enable_if_t<3 == NSD, void>
    coordinates( Edge<Dir>, const Integer index[3], Scalar x[3] ) const
    {
        for ( std::size_t d = 0; d < 3; ++d )
            if ( Dir == d )
                x[d] = ( _local_edges[d]( index[d] + 1 ) +
                         _local_edges[d]( index[d] ) ) /
                       Scalar( 2.0 );
            else
                x[d] = _local_edges[d]( index[d] );
    }

    template <class Integer>
    KOKKOS_INLINE_FUNCTION Scalar measure( Node,
                                           const Integer[num_space_dim] ) const
    {
        return 0.0;
    }

    template <int Dir, class Integer, std::size_t NSD = num_space_dim>
    KOKKOS_INLINE_FUNCTION std::enable_if_t<3 == NSD, Scalar>
    measure( Edge<Dir>, const Integer index[3] ) const
    {
        return _local_edges[Dir][index[Dir] + 1] -
               _local_edges[Dir][index[Dir]];
    }

    template <class Integer, std::size_t NSD = num_space_dim>
    KOKKOS_INLINE_FUNCTION std::enable_if_t<3 == NSD, Scalar>
    measure( Face<Dim::I>, const Integer index[3] ) const
    {
        return measure( Edge<Dim::J>(), index ) *
               measure( Edge<Dim::K>(), index );
    }

    template <class Integer, std::size_t NSD = num_space_dim>
    KOKKOS_INLINE_FUNCTION std::enable_if_t<3 == NSD, Scalar>
    measure( Face<Dim::J>, const Integer index[3] ) const
    {
        return measure( Edge<Dim::I>(), index ) *
               measure( Edge<Dim::K>(), index );
    }

    template <class Integer, std::size_t NSD = num_space_dim>
    KOKKOS_INLINE_FUNCTION std::enable_if_t<3 == NSD, Scalar>
    measure( Face<Dim::K>, const Integer index[3] ) const
    {
        return measure( Edge<Dim::I>(), index ) *
               measure( Edge<Dim::J>(), index );
    }

    template <class Integer, std::size_t NSD = num_space_dim>
    KOKKOS_INLINE_FUNCTION std::enable_if_t<2 == NSD, Scalar>
    measure( Face<Dim::I>, const Integer index[2] ) const
    {
        return _local_edges[Dim::J][index[Dim::J] + 1] -
               _local_edges[Dim::J][index[Dim::J]];
    }

    template <class Integer, std::size_t NSD = num_space_dim>
    KOKKOS_INLINE_FUNCTION std::enable_if_t<2 == NSD, Scalar>
    measure( Face<Dim::J>, const Integer index[2] ) const
    {
        return _local_edges[Dim::I][index[Dim::I] + 1] -
               _local_edges[Dim::I][index[Dim::I]];
    }

    template <class Integer>
    KOKKOS_INLINE_FUNCTION Scalar
    measure( Cell, const Integer index[num_space_dim] ) const
    {
        Scalar m = 1.0;
        for ( std::size_t d = 0; d < num_space_dim; ++d )
            m *= _local_edges[d][index[d] + 1] - _local_edges[d][index[d]];
        return m;
    }
 
  private:
    Kokkos::Array<Scalar, num_space_dim> _own_low_corner;
    Kokkos::Array<Scalar, num_space_dim> _own_high_corner;
    Kokkos::Array<Scalar, num_space_dim> _ghost_low_corner;
    Kokkos::Array<Scalar, num_space_dim> _ghost_high_corner;
    Kokkos::Array<Kokkos::View<Scalar*, MemorySpace>, num_space_dim>
        _local_edges;
    Kokkos::Array<bool, num_space_dim> _periodic;
    Kokkos::Array<bool, num_space_dim> _boundary_lo;
    Kokkos::Array<bool, num_space_dim> _boundary_hi;
};
 
//---------------------------------------------------------------------------//
template <class MemorySpace, class MeshType>
LocalMesh<MemorySpace, MeshType>
createLocalMesh( const LocalGrid<MeshType>& local_grid )
{
    return LocalMesh<MemorySpace, MeshType>( local_grid );
}
 
//---------------------------------------------------------------------------//
 
} // namespace Grid
} // namespace Cabana
 
#endif // end CABANA_GRID_LOCALMESH_HPP