diff --git a/examples/populate_matrix.jl b/examples/populate_matrix.jl
new file mode 100644
index 0000000..d8a8c87
--- /dev/null
+++ b/examples/populate_matrix.jl
@@ -0,0 +1,50 @@
+# In this example, we populate a simple matrix based on a geometry definition.
+# Specifically, we loop over a series of grid points and decide what kind of
+# material should occupy that cell by searching through the geometry
+# specification. We show how one can forward the machinery underneath
+# GeometryPrimitives.jl so that we can define our own "material" types.
+
+using GeometryPrimitives
+using Lazy
+
+# create new material type
+struct Material <: Shape{3,9}
+    shape::Shape   # material shape
+    ε::Float64     # material permittivity
+end
+
+# forward relevant methods from our new type
+@forward Material.shape Base.in, GeometryPrimitives.bounds
+
+function populate_matrix()
+    # set up geometry
+    geometry = [
+        Material(Cylinder([0.0, 0.0, 0.0], 2.5, 100.0), 3.4)
+    ]
+
+    # set up domain
+    sx, sy = 10.0, 10.0
+    resolution = 20
+
+    # build matrix
+    Nx = Int(sx * resolution)
+    Ny = Int(sy * resolution)
+    M = zeros(Nx,Ny)
+
+    # loop over grid points
+    x = range(-sx/2,stop=sx/2,length=Nx)
+    y = range(-sy/2,stop=sy/2,length=Ny)
+    for ix = 1:Nx, iy = 1:Ny
+        p = [x[ix], y[iy], 0.0] # current point on grid
+        idx = findfirst(p,geometry) # search through geometry tree
+        a = 1.0 # default material
+        if !isnothing(idx)
+            a = geometry[idx].ε
+        end
+        M[ix,iy] = a # fill in matrix
+    end
+
+    return M
+end
+
+M = populate_matrix()
\ No newline at end of file
diff --git a/src/util/kdtree.jl b/src/util/kdtree.jl
index 8c803e3..819f789 100644
--- a/src/util/kdtree.jl
+++ b/src/util/kdtree.jl
@@ -16,14 +16,16 @@ not shapes of nonzero size.)
 """
 mutable struct KDTree{K,S<:Shape{K}}
     s::Vector{S}
+    s_index::Vector{Int}
     ix::Int
     x::Float64
     left::KDTree{K,S}  # shapes ≤ x in coordinate ix
     right::KDTree{K,S} # shapes > x in coordinate ix
-    KDTree{K,S}(s::AbstractVector{S}) where {K,S<:Shape{K}} = new(s, 0)
+    KDTree{K,S}(s::AbstractVector{S}) where {K,S<:Shape{K}} = new(s, collect(eachindex(s)), 0)
+    KDTree{K,S}(s::AbstractVector{S},s_index::Vector{<:Int}) where {K,S<:Shape{K}} = new(s, s_index, 0)
     function KDTree{K,S}(ix::Integer, x::Real, left::KDTree{K,S}, right::KDTree{K,S}) where {K,S<:Shape{K}}
         1 ≤ ix ≤ K || throw(BoundsError())
-        new(S[], ix, x, left, right)
+        new(S[], Int[], ix, x, left, right)
     end
 end
 
@@ -38,8 +40,15 @@ Construct a K-D tree (`KDTree`) representation of a list of
 When searching the tree, shapes that appear earlier in `s`
 take precedence over shapes that appear later.
 """
+
 function KDTree(s::AbstractVector{S}) where {K,S<:Shape{K}}
-    (length(s) ≤ 4 || K == 0) && return KDTree{K,S}(s)
+    # If no list of indicies is provided, simply enumerate by the number of
+    # shapes in `s`.
+    return KDTree(s,collect(eachindex(s)))
+end
+
+function KDTree(s::AbstractVector{S}, s_index::AbstractVector{<:Integer}) where {K,S<:Shape{K}}
+    (length(s) ≤ 4 || K == 0) && return KDTree{K,S}(s, s_index)
 
     # figure out the best dimension ix to divide over,
     # the dividing plane x, and the number (nl,nr) of
@@ -61,22 +70,26 @@ function KDTree(s::AbstractVector{S}) where {K,S<:Shape{K}}
     end
 
     # don't bother subdividing if it doesn't reduce the # of shapes much
-    4*max(nl,nr) > 3*length(s) && return KDTree{K,S}(s)
+    4*max(nl,nr) > 3*length(s) && return KDTree{K,S}(s,s_index)
 
     # create the arrays of shapes in each subtree
     sl = Vector{S}(undef, nl)
+    sl_idx = Vector{Int}(undef, nl)
     sr = Vector{S}(undef, nr)
+    sr_idx = Vector{Int}(undef, nr)
     il = ir = 0
     for k in eachindex(s)
         if b[k][1][ix] ≤ x
             sl[il += 1] = s[k]
+            sl_idx[il] = s_index[k]
         end
         if b[k][2][ix] > x
             sr[ir += 1] = s[k]
+            sr_idx[ir] = s_index[k]
         end
     end
 
-    return KDTree{K,S}(ix, x, KDTree(sl), KDTree(sr))
+    return KDTree{K,S}(ix, x, KDTree(sl,sl_idx), KDTree(sr,sr_idx))
 end
 
 depth(kd::KDTree) = kd.ix == 0 ? 0 : max(depth(kd.left), depth(kd.right)) + 1
@@ -104,7 +117,7 @@ function Base.findfirst(p::SVector{N}, s::Vector{S}) where {N,S<:Shape{N}}
     for i in eachindex(s)
         b = bounds(s[i])
         if all(b[1] .< p .< b[2]) && p ∈ s[i]  # check if p is within bounding box is faster
-            return s[i]
+            return i
         end
     end
     return nothing
@@ -118,7 +131,12 @@ function Base.findfirst(p::SVector{N}, kd::KDTree{N}) where {N}
             return findfirst(p, kd.right)
         end
     else
-        return findfirst(p, kd.s)
+        idx = findfirst(p, kd.s)
+        if isnothing(idx)
+            return idx
+        else
+            return kd.s_index[idx]
+        end
     end
 end
 
diff --git a/test/kdtree.jl b/test/kdtree.jl
index 5ad5497..390caa9 100644
--- a/test/kdtree.jl
+++ b/test/kdtree.jl
@@ -7,7 +7,7 @@
     s = Shape2[Ball([i,0], 1) for i in 0:20]
     kd = KDTree(s)
     @test GeometryPrimitives.depth(kd) == 3
-    @test findfirst([10.1,0], kd).c[1] == 10
+    @test s[findfirst([10.1,0], kd)].c[1] == 10
     @test findfirst([10.1,1], kd) == nothing
     @test checktree(kd, s)