Search code examples
arraysjulia

Average neighbor cell values in 3D array


I have a 3D array that represents a lake basin, some of the 3D array is land (-999 values) and is not of interest. Within the lake, cells can contain either pelagic resource or littoral resources. These two resources grow at different rates and can be treated as independent, non-interacting parts of the lake.

Within the littoral and within the pelagic cells, I want to average cell values with neighboring cells (Chebyshev distance = 1), but only if the neighboring cells are of the same type, e.g. pelagic cells should only average with neighboring pelagic cells, not littoral or land cells and littoral cells should only average with neighboring littoral cells, not pelagic or land cells.

I have created a small hypothetical example.

using Distributions

# create hypothetical lake bathymetry
foo(N) = [min(i, j, N+1-i, N+1-j) for i in 1:N, j in 1:N]

# 10 x 10 surface area, with max depth of 5
bathy = foo(10)

# delineate land (-999) - area not interested in
bathy[bathy .< 2] .= -999

# delineate either littoral or pelagic part of lake
lake_type = ones(Int, size(bathy))
lake_type .= bathy

# if depth is between 0 and 4 cell is littoral (1)
lake_type[lake_type .> -1 .&& lake_type .< 4] .= 1

# if depth is > than 3 cell is pelagic (0)
lake_type[lake_type .> 3] .= 0

# create 3D lake resource array --------

# max depth 
mx_dpth = maximum(bathy)

# dims of basin
dims = (10,10, mx_dpth)

# 3d array of basin - to store resources
# [:, :, 1] = deepest layer of the lake
# [:, :, 5] = surface layer of the lake
basal_resource = zeros(dims)

# populate with resource - varies depending if littoral (0) or pelagic (1)
for i in 1:dims[1], j in 1:dims[2]
    if lake_type[i,j] == 0    # initial pelagic resource amount 
        basal_resource[i, j, (mx_dpth - bathy[i,j] + 1):mx_dpth] .= round.(rand(Uniform(100, 150), 1), sigdigits=3)
    end
    if lake_type[i,j] == 1 # initial littoral  resource amount
        basal_resource[i, j, (mx_dpth - bathy[i,j] + 1):mx_dpth] .= round.(rand(Uniform(1, 5), 1), sigdigits=3)
    end
end

# convert "land" cells to -999
basal_resource[basal_resource .== 0.0] .= -999

I'm now stuck trying to average only neighboring cells (in 3 dimensions) in basal_resource of the same type. Any advice would be much appreciated, hopefully my question makes sense.

EDIT

The 2D lake_type array is only relevant at the surface layer, as cells may be water at the surface, but then become land deeper down. Here is a 3D array indicating lake type through the whole water column.

# create 3d "map" of the lake indicating if cell is pelagic (0) or littoral (1)
lake_type_3d = copy(basal_resource)
replace!(x -> x>100.0 ? 0.0 : x, lake_type_3d)
replace!(x -> x>1.0 ? 1.0 : x, lake_type_3d)

Solution

  • Assuming the borders of the lake_type variable are always land (so the average doesn't really matter), you could do this with a loop like this:

    avg_same = zeros(size(basal_resource))
    
    for i in 2:size(basal_resource, 1)-1, j in 2:size(basal_resource, 2)-1
        neighbors2 = repeat(lake_type[i-1:i+1, j-1:j+1] .== lake_type[i, j], outer = [1, 1, 2])
        neighbors3 = repeat(lake_type[i-1:i+1, j-1:j+1] .== lake_type[i, j], outer = [1, 1, 3])
        avg_same[i, j, 1] = sum(basal_resource[i-1:i+1, j-1:j+1, 1:2] .* neighbors2) / sum(neighbors2)
        avg_same[i, j, end] = sum(basal_resource[i-1:i+1, j-1:j+1, end-1:end] .* neighbors2) / sum(neighbors2)
        for k in 2:size(basal_resource, 3)-1
            avg_same[i, j, k] = sum(basal_resource[i-1:i+1, j-1:j+1, k-1:k+1] .* neighbors3) / sum(neighbors3)
        end
    end
    

    This might not be the most performant way of accomplishing the task, but for decently sized arrays a nested loop shouldn't be too taxing.

    EDIT

    If you use a 3D version of lake_type, you can do the same kind of matrix multiplication logic:

    avg_same = zeros(size(basal_resource))
    
    for i in 2:size(basal_resource, 1)-1, j in 2:size(basal_resource, 2)-1, k in 1:size(basal_resource, 3)
        kmin = max(k-1, 1)
        kmax = min(k+1, size(basal_resource, 3))
        neighbors = lake_type_3d[i-1:i+1, j-1:j+1, kmin:kmax] .== lake_type_3d[i, j, k]
        avg_same[i, j, k] = sum(basal_resource[i-1:i+1, j-1:j+1, kmin:kmax] .* neighbors) / sum(neighbors)
    end