Skip to content

Commit

Permalink
feat(subdomain): 2D -> 3D solver + tests
Browse files Browse the repository at this point in the history
  • Loading branch information
@rouson
rouson committed Dec 21, 2023
1 parent 54a8f65 commit d85487e
Show file tree
Hide file tree
Showing 3 changed files with 224 additions and 181 deletions.
11 changes: 9 additions & 2 deletions src/matcha/subdomain_m.f90
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ module subdomain_m

type subdomain_t
private
real, allocatable :: s_(:,:)
real, allocatable :: s_(:,:,:)
contains
procedure, pass(self) :: define
procedure, pass(self) :: step
Expand All @@ -20,6 +20,7 @@ module subdomain_m
generic :: assignment(=) => assign_and_sync
procedure dx
procedure dy
procedure dz
procedure values
end type

Expand All @@ -41,7 +42,7 @@ module subroutine step(alpha_dt, self)
pure module function values(self) result(my_values)
implicit none
class(subdomain_t), intent(in) :: self
real, allocatable :: my_values(:,:)
real, allocatable :: my_values(:,:,:)
end function

pure module function dx(self) result(my_dx)
Expand All @@ -56,6 +57,12 @@ pure module function dy(self) result(my_dy)
real my_dy
end function

pure module function dz(self) result(my_dz)
implicit none
class(subdomain_t), intent(in) :: self
real my_dz
end function

pure module function laplacian(rhs) result(laplacian_rhs)
implicit none
class(subdomain_t), intent(in) :: rhs
Expand Down
185 changes: 98 additions & 87 deletions src/matcha/subdomain_s.f90
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,13 +4,13 @@
use intrinsic_array_m, only : intrinsic_array_t
implicit none

real, allocatable :: halo_x(:,:)[:]
real, allocatable :: halo_x(:,:,:)[:]
integer, parameter :: west=1, east=2

type(data_partition_t) data_partition

real dx_, dy_
integer my_nx, nx, ny, me, num_subdomains, my_internal_left, my_internal_right
real dx_, dy_, dz_
integer my_nx, nx, ny, nz, me, num_subdomains, my_internal_west, my_internal_east

contains

Expand All @@ -20,8 +20,12 @@

nx = n
ny = nx
nz = nx

dx_ = side/(nx-1)
dy_ = dx_
dz_ = dx_

call assert(num_subdomains <= nx-nx_boundaries, &
"subdomain_t%define: num_subdomains <= nx-nx_boundaries", intrinsic_array_t([nx, num_subdomains]))
me = this_image()
Expand All @@ -31,32 +35,26 @@
my_nx = data_partition%last(me) - data_partition%first(me) + 1

if (allocated(self%s_)) deallocate(self%s_)
allocate(self%s_(my_nx, ny))

my_internal_left = merge(2, 1, me==1)
my_internal_right = merge(my_nx-1, my_nx, me==num_subdomains)

self%s_(my_internal_left:my_internal_right, 1) = boundary_val ! bottom subdomain boundary
self%s_(my_internal_left:my_internal_right, ny) = boundary_val ! top subdomain boundary
self%s_(my_internal_left:my_internal_right, 2:ny-1) = internal_val ! internal points

if (me == 1) then
self%s_(1, 1:ny) = boundary_val ! left domain boundary
else
self%s_(1, 2:ny-1) = internal_val ! left subdomain boundary
end if
if (me == num_subdomains) then
self%s_(my_nx, 1:ny) = boundary_val ! right domain boundary
else
self%s_(my_nx, 2:ny-1) = internal_val ! right subdomain boundary
end if
allocate(self%s_(my_nx, ny, nz))

my_internal_west = merge(2, 1, me==1)
my_internal_east = merge(my_nx-1, my_nx, me==num_subdomains)

self%s_(my_internal_west:my_internal_east, 2:ny-1, 2:nz-1) = internal_val ! internal points

self%s_(:, : , 1 ) = boundary_val ! minimum z boundary
self%s_(:, : , nz) = boundary_val ! maximum z boundary
self%s_(:, 1 , : ) = boundary_val ! minimum y boundary
self%s_(:, ny, : ) = boundary_val ! maximum y boundary

if (me == 1) self%s_(1 , :, :) = boundary_val ! minimum x boundary
if (me == num_subdomains) self%s_(my_nx, :, :) = boundary_val ! maximum x boundary

if (allocated(halo_x)) deallocate(halo_x)
allocate(halo_x(west:east, ny)[*])
if (me>1) halo_x(east,:)[me-1] = self%s_(1,:)
if (me<num_subdomains) halo_x(west,:)[me+1] = self%s_(my_nx,:)
allocate(halo_x(west:east, ny, nz)[*])
if (me>1) halo_x(east,:,:)[me-1] = self%s_(1,:,:)
if (me<num_subdomains) halo_x(west,:,:)[me+1] = self%s_(my_nx,:,:)
sync all

end procedure

module procedure dx
Expand All @@ -67,41 +65,51 @@
my_dy = dy_
end procedure

module procedure dz
my_dz = dz_
end procedure

module procedure laplacian

integer i, j
real, allocatable :: halo_left(:), halo_right(:)
integer i, j, k
real, allocatable :: halo_west(:,:), halo_east(:,:)

call assert(allocated(rhs%s_), "subdomain_t%laplacian: allocated(rhs%s_)")
call assert(allocated(halo_x), "subdomain_t%laplacian: allocated(halo_x)")

allocate(laplacian_rhs%s_(my_nx, ny))
allocate(laplacian_rhs%s_(my_nx, ny, nz))

halo_left = merge(halo_x(west,:), rhs%s_(1,:), me/=1)
i = my_internal_left
call assert(i+1<=my_nx,"laplacian: leftmost subdomain too small")
do concurrent(j=2:ny-1)
laplacian_rhs%s_(i,j) = (halo_left(j) - 2*rhs%s_(i,j) + rhs%s_(i+1,j))/dx_**2 + &
(rhs%s_(i,j-1) - 2*rhs%s_(i,j) + rhs%s_(i,j+1))/dy_**2
halo_west = merge(halo_x(west,:,:), rhs%s_(1,:,:), me/=1)
i = my_internal_west
call assert(i+1<=my_nx,"laplacian: westernmost subdomain too small")
do concurrent(j=2:ny-1, k=2:nz-1)
laplacian_rhs%s_(i,j,k) = ( halo_west(j,k ) - 2*rhs%s_(i,j,k) + rhs%s_(i+1,j ,k ))/dx_**2 + &
(rhs%s_(i,j-1,k ) - 2*rhs%s_(i,j,k) + rhs%s_(i ,j+1,k ))/dy_**2 + &
(rhs%s_(i,j ,k-1) - 2*rhs%s_(i,j,k) + rhs%s_(i ,j ,k+1))/dz_**2
end do

do concurrent(i=my_internal_left+1:my_internal_right-1, j=2:ny-1)
laplacian_rhs%s_(i,j) = (rhs%s_(i-1,j) - 2*rhs%s_(i,j) + rhs%s_(i+1,j))/dx_**2 + &
(rhs%s_(i,j-1) - 2*rhs%s_(i,j) + rhs%s_(i,j+1))/dy_**2
do concurrent(i=my_internal_west+1:my_internal_east-1, j=2:ny-1, k=2:nz-1)
laplacian_rhs%s_(i,j,k) = (rhs%s_(i-1,j ,k ) - 2*rhs%s_(i,j,k) + rhs%s_(i+1,j ,k ))/dx_**2 + &
(rhs%s_(i ,j-1,k ) - 2*rhs%s_(i,j,k) + rhs%s_(i ,j+1,k ))/dy_**2 + &
(rhs%s_(i ,j ,k-1) - 2*rhs%s_(i,j,k) + rhs%s_(i ,j ,k+1))/dz_**2
end do

halo_right = merge(halo_x(east,:), rhs%s_(my_nx,:), me/=num_subdomains)
i = my_internal_right
call assert(i-1>0,"laplacian: rightmost subdomain too small")
do concurrent(j=2:ny-1)
laplacian_rhs%s_(i,j) = (rhs%s_(i-1,j) - 2*rhs%s_(i,j) + halo_right(j))/dx_**2 + &
(rhs%s_(i,j-1) - 2*rhs%s_(i,j) + rhs%s_(i,j+1))/dy_**2
halo_east = merge(halo_x(east,:,:), rhs%s_(my_nx,:,:), me/=num_subdomains)
i = my_internal_east
call assert(i-1>0,"laplacian: easternmost subdomain too small")
do concurrent(j=2:ny-1, k=2:nz-1)
laplacian_rhs%s_(i,j,k) = (rhs%s_(i-1,j ,k ) - 2*rhs%s_(i,j,k) + halo_east(j ,k ))/dx_**2 + &
(rhs%s_(i ,j-1,k ) - 2*rhs%s_(i,j,k) + rhs%s_(i ,j+1,k ))/dy_**2 + &
(rhs%s_(i ,j ,k-1) - 2*rhs%s_(i,j,k) + rhs%s_(i ,j ,k+1))/dz_**2
end do

laplacian_rhs%s_(:, 1) = 0.
laplacian_rhs%s_(:,ny) = 0.
if (me==1) laplacian_rhs%s_(1,:) = 0.
if (me==num_subdomains) laplacian_rhs%s_(my_nx,:) = 0.
laplacian_rhs%s_(:, 1,:) = 0.
laplacian_rhs%s_(:,ny,:) = 0.
laplacian_rhs%s_(:,:, 1) = 0.
laplacian_rhs%s_(:,:,nz) = 0.
if (me==1) laplacian_rhs%s_(1,:,:) = 0.
if (me==num_subdomains) laplacian_rhs%s_(my_nx,:,:) = 0.

end procedure

module procedure multiply
Expand All @@ -118,8 +126,8 @@
call assert(allocated(rhs%s_), "subdomain_t%assign_and_sync: allocated(rhs%s_)")
sync all
lhs%s_ = rhs%s_
if (me>1) halo_x(east,:)[me-1] = rhs%s_(1,:)
if (me<num_subdomains) halo_x(west,:)[me+1] = rhs%s_(my_nx,:)
if (me>1) halo_x(east,:,:)[me-1] = rhs%s_(1,:,:)
if (me<num_subdomains) halo_x(west,:,:)[me+1] = rhs%s_(my_nx,:,:)
sync all
end procedure

Expand All @@ -130,14 +138,14 @@

module procedure step

real, allocatable :: increment(:,:)
real, allocatable :: increment(:,:,:)

call assert(allocated(self%s_), "subdomain_t%laplacian: allocated(rhs%s_)")
call assert(allocated(halo_x), "subdomain_t%laplacian: allocated(halo_x)")
call assert(my_internal_left+1<=my_nx,"laplacian: leftmost subdomain too small")
call assert(my_internal_right-1>0,"laplacian: rightmost subdomain too small")
call assert(my_internal_west+1<=my_nx,"laplacian: westernmost subdomain too small")
call assert(my_internal_east-1>0,"laplacian: easternmost subdomain too small")

allocate(increment(my_nx,ny))
allocate(increment(my_nx,ny,nz))

call internal_points(increment)
call edge_points(increment)
Expand All @@ -151,58 +159,61 @@
contains

subroutine internal_points(ds)
real, intent(inout) :: ds(:,:)
integer i, j

do concurrent(i=my_internal_left+1:my_internal_right-1, j=2:ny-1)
ds(i,j) = alpha_dt*( &
(self%s_(i-1,j) - 2*self%s_(i,j) + self%s_(i+1,j))/dx_**2 + &
(self%s_(i,j-1) - 2*self%s_(i,j) + self%s_(i,j+1))/dy_**2 &
real, intent(inout) :: ds(:,:,:)
integer i, j, k

do concurrent(i=my_internal_west+1:my_internal_east-1, j=2:ny-1, k=2:nz-1)
ds(i,j,k) = alpha_dt*( &
(self%s_(i-1,j ,k ) - 2*self%s_(i,j,k) + self%s_(i+1,j,k ))/dx_**2 + &
(self%s_(i ,j-1,k ) - 2*self%s_(i,j,k) + self%s_(i,j+1,k ))/dy_**2 + &
(self%s_(i ,j ,k-1) - 2*self%s_(i,j,k) + self%s_(i,j ,k+1))/dz_**2 &
)
end do
end subroutine

subroutine edge_points(ds)
real, intent(inout) :: ds(:,:)
real, allocatable :: halo_left(:), halo_right(:)
integer i, j

halo_left = merge(halo_x(west,:), self%s_(1,:), me/=1)
halo_right = merge(halo_x(east,:), self%s_(my_nx,:), me/=num_subdomains)

i = my_internal_left
do concurrent(j=2:ny-1)
ds(i,j) = alpha_dt*( &
(halo_left(j) - 2*self%s_(i,j) + self%s_(i+1,j))/dx_**2 + &
(self%s_(i,j-1) - 2*self%s_(i,j) + self%s_(i,j+1))/dy_**2 &
real, intent(inout) :: ds(:,:,:)
real, allocatable :: halo_west(:,:), halo_east(:,:)
integer i, j, k

halo_west = merge(halo_x(west,:,:), self%s_(1, :,:), me/=1)
halo_east = merge(halo_x(east,:,:), self%s_(my_nx,:,:), me/=num_subdomains)

i = my_internal_west
do concurrent(j=2:ny-1,k=2:nz-1)
ds(i,j,k) = alpha_dt*( &
(halo_west(j ,k ) - 2*self%s_(i,j,k) + self%s_(i+1,j ,k ))/dx_**2 + &
(self%s_(i,j-1,k ) - 2*self%s_(i,j,k) + self%s_(i ,j+1,k ))/dy_**2 + &
(self%s_(i,j ,k-1) - 2*self%s_(i,j,k) + self%s_(i ,j ,k+1))/dz_**2 &
)
end do

i = my_internal_right
do concurrent(j=2:ny-1)
ds(i,j) = alpha_dt*( &
(self%s_(i-1,j) - 2*self%s_(i,j) + halo_right(j))/dx_**2 + &
(self%s_(i,j-1) - 2*self%s_(i,j) + self%s_(i,j+1))/dy_**2 &
i = my_internal_east
do concurrent(j=2:ny-1, k=2:nz-1)
ds(i,j,k) = alpha_dt*( &
(self%s_(i-1,j ,k ) - 2*self%s_(i,j,k) + halo_east(j ,k ))/dx_**2 + &
(self%s_(i ,j-1,k ) - 2*self%s_(i,j,k) + self%s_(i,j+1,k ))/dy_**2 + &
(self%s_(i ,j ,k-1) - 2*self%s_(i,j,k) + self%s_(i,j ,k+1))/dz_**2 &
)
end do
end subroutine

subroutine apply_boundary_condition(ds)
real, intent(inout) :: ds(:,:)
real, intent(inout) :: ds(:,:,:)
integer i, j

ds(:, 1) = 0.
ds(:,ny) = 0.
if (me==1) ds(1,:) = 0.
if (me==num_subdomains) ds(my_nx,:) = 0.
ds(:,1:ny:ny-1,: ) = 0.
ds(:, : ,1:nz:nz-1) = 0.
if (me==1) ds(1,:,:) = 0.
if (me==num_subdomains) ds(my_nx,:,:) = 0.
end subroutine

subroutine exchange_halo(s)
real, intent(in) :: s(:,:)
if (me>1) halo_x(east,:)[me-1] = s(1,:)
if (me<num_subdomains) halo_x(west,:)[me+1] = s(my_nx,:)
real, intent(in) :: s(:,:,:)
if (me>1) halo_x(east,:,:)[me-1] = s(1,:,:)
if (me<num_subdomains) halo_x(west,:,:)[me+1] = s(my_nx,:,:)
end subroutine

end procedure

end submodule subdomain_s
end submodule subdomain_s
Loading

0 comments on commit d85487e

Please sign in to comment.