Skip to content

Commit

Permalink
Clean up Cabana usage and comments
Browse files Browse the repository at this point in the history
  • Loading branch information
@streeve
streeve committed Jun 1, 2023
1 parent ca0aee7 commit 8a168d6
Showing 1 changed file with 38 additions and 37 deletions.
75 changes: 38 additions & 37 deletions example/methods/finite_difference/finite_difference.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,34 +9,35 @@
* SPDX-License-Identifier: BSD-3-Clause *
****************************************************************************/

#include <Cajita.hpp>
#include <Kokkos_Core.hpp>
#include <array>
#include <math.h>

#include <mpi.h>

// we are currently assuming each boundary is a uniform Dirichlet condition
template <class ExecutionSpace, class grid_t, class array_t>
void updateBoundaries( ExecutionSpace exec_space, grid_t local_grid,
array_t& field )
#include <Kokkos_Core.hpp>

#include <Cajita.hpp>

template <class ExecutionSpace, class LocalGridType, class ArrayType>
void updateBoundaries( ExecutionSpace exec_space, LocalGridType local_grid,
ArrayType& field )
{
// Update the boundary on each face of the cube.
for ( int d = 0; d < 3; d++ )
{
for ( int dir = -1; dir < 2; dir += 2 )
{
std::array<int, 3> plane = { 0, 0, 0 };

plane[d] = dir;

// Get the boundary indices for this plane (each one is a separate,
// contiguous index space).
auto boundary_space = local_grid->boundaryIndexSpace(
Cajita::Own(), Cajita::Cell(), plane );

Cajita::grid_parallel_for(
"boundary_grid_for", exec_space, boundary_space,
"boundary_update", exec_space, boundary_space,
KOKKOS_LAMBDA( const int i, const int j, const int k ) {
// Dirichlet boundary condition example
// field( i, j, k, 0 ) = 1.0;

// Neumann boundary condition example
field( i, j, k, 0 ) =
field( i - plane[0], j - plane[1], k - plane[2], 0 );
Expand All @@ -45,14 +46,13 @@ void updateBoundaries( ExecutionSpace exec_space, grid_t local_grid,
}
}

void createGrid()
void finiteDifference()
{
int comm_rank;
MPI_Comm_rank( MPI_COMM_WORLD, &comm_rank );

using exec_space = Kokkos::DefaultExecutionSpace;
using device_type = exec_space::device_type;
using memory_space = device_type::memory_space;

// set up block decomposition
int comm_size;
Expand All @@ -66,15 +66,13 @@ void createGrid()
global_low_corner, global_high_corner, cell_size );

// Create the global grid
std::array<int, 3> ranks_per_dim = { comm_size, 1, 1 };
MPI_Dims_create( comm_size, 3, ranks_per_dim.data() );
Cajita::ManualBlockPartitioner<3> partitioner( ranks_per_dim );
Cajita::DimBlockPartitioner<3> partitioner;
std::array<bool, 3> periodic = { false, false, false };

auto global_grid =
createGlobalGrid( MPI_COMM_WORLD, global_mesh, periodic, partitioner );

// Create a local grid and local mesh with halo region
// Create a local grid and local mesh with halo region.
// Note this halo width needs to be large enough for the stencil used below.
unsigned halo_width = 1;
auto local_grid = Cajita::createLocalGrid( global_grid, halo_width );
auto local_mesh = Cajita::createLocalMesh<device_type>( *local_grid );
Expand All @@ -91,24 +89,23 @@ void createGrid()
Cajita::ArrayOp::assign( *T, 300.0, Cajita::Ghost() );
Cajita::ArrayOp::assign( *T, 300.0, Cajita::Own() );

auto T_view = T->view();

auto T_halo = createHalo( Cajita::NodeHaloPattern<3>(), halo_width, *T );

// update physical and processor boundaries
auto T_view = T->view();
updateBoundaries( exec_space(), local_grid, T_view );
T_halo->gather( exec_space(), *T );

// create and array to store previous temperature for explicit udpate
auto T0 = Cajita::createArray<double, device_type>( name, layout );
auto T0_view = T0->view();
auto T_prev = Cajita::createArray<double, device_type>( name, layout );
auto T_prev_view = T_prev->view();

// Solve heat conduction from point source
double dt = 1e-6;
double end_time = 1e-3;
int numSteps = static_cast<int>( end_time / dt );
int num_steps = static_cast<int>( end_time / dt );

// properties
// Material properties
double rho = 1000.0;
double specific_heat = 1000.0;
double kappa = 10.0;
Expand All @@ -126,14 +123,15 @@ void createGrid()

double I = 2.0 * eta * power / ( M_PI * sqrt( M_PI ) * r[0] * r[1] * r[2] );

for ( int step = 0; step < numSteps; ++step )
// Timestep loop.
for ( int step = 0; step < num_steps; ++step )
{
// store previous value for explicit update
Kokkos::deep_copy( T0_view, T_view );
// Store previous value for explicit update
Kokkos::deep_copy( T_prev_view, T_view );

// Solve finite difference
Cajita::grid_parallel_for(
"local_grid_for", exec_space(), owned_space,
"finite_difference", exec_space(), owned_space,
KOKKOS_LAMBDA( const int i, const int j, const int k ) {
double loc[3];
int idx[3] = { i, j, k };
Expand All @@ -145,23 +143,26 @@ void createGrid()

double Q = I * exp( -f ) * dt_rho_cp;

double laplacian =
( -6.0 * T0_view( i, j, k, 0 ) + T0_view( i - 1, j, k, 0 ) +
T0_view( i + 1, j, k, 0 ) + T0_view( i, j - 1, k, 0 ) +
T0_view( i, j + 1, k, 0 ) + T0_view( i, j, k - 1, 0 ) +
T0_view( i, j, k + 1, 0 ) ) *
alpha_dt_dx2;
double laplacian = ( -6.0 * T_prev_view( i, j, k, 0 ) +
T_prev_view( i - 1, j, k, 0 ) +
T_prev_view( i + 1, j, k, 0 ) +
T_prev_view( i, j - 1, k, 0 ) +
T_prev_view( i, j + 1, k, 0 ) +
T_prev_view( i, j, k - 1, 0 ) +
T_prev_view( i, j, k + 1, 0 ) ) *
alpha_dt_dx2;

T_view( i, j, k, 0 ) = T0_view( i, j, k, 0 ) + laplacian + Q;
T_view( i, j, k, 0 ) += laplacian + Q;
} );

// update the boundary conditions
// Update the boundary conditions
updateBoundaries( exec_space(), local_grid, T_view );

// Exchange halo values
T_halo->gather( exec_space(), *T );
}

// Write the final state.
Cajita::Experimental::BovWriter::writeTimeStep( 0, 0, *T );
}

Expand All @@ -171,7 +172,7 @@ int main( int argc, char* argv[] )
{
Kokkos::ScopeGuard scope_guard( argc, argv );

createGrid();
finiteDifference();
}
MPI_Finalize();

Expand Down

0 comments on commit 8a168d6

Please sign in to comment.