diff --git a/riskmapjnr/benchmark/vulnerability_map.py b/riskmapjnr/benchmark/vulnerability_map.py
index ff7871b..51e0feb 100644
--- a/riskmapjnr/benchmark/vulnerability_map.py
+++ b/riskmapjnr/benchmark/vulnerability_map.py
@@ -7,6 +7,7 @@
 
 
 def vulnerability_map(
+        forest_file,
         dist_file,
         dist_bins,
         subj_file,
@@ -20,6 +21,10 @@ def vulnerability_map(
     values indicate higher vulnerability. Raster type is UInt16 ([0,
     65535]). NoData value is set to 0.
 
+    :param forest_file: Input forest cover file. Necessary to have
+        forest extent within country borders (info not provided by
+        dist_file).
+
     :param dist_file: Input file of distance to forest edge.
 
     :param dist_bins: The distance bins used to convert distance to
@@ -44,19 +49,23 @@ def vulnerability_map(
     # Input rasters
     # ==============================================================
 
+    # Forest cover
+    forest_ds = gdal.Open(forest_file)
+    forest_band = forest_ds.GetRasterBand(1)
+    # Raster size
+    xsize = forest_band.XSize
+    ysize = forest_band.YSize
+
     # Distance to forest edge raster file
     dist_ds = gdal.Open(dist_file)
     dist_band = dist_ds.GetRasterBand(1)
-    # Raster size
-    xsize = dist_band.XSize
-    ysize = dist_band.YSize
 
     # Subjurisdictions raster file
     subj_ds = gdal.Open(subj_file)
     subj_band = subj_ds.GetRasterBand(1)
 
     # Make blocks
-    blockinfo = makeblock(dist_file, blk_rows=blk_rows)
+    blockinfo = makeblock(forest_file, blk_rows=blk_rows)
     nblock = blockinfo[0]
     nblock_x = blockinfo[1]
     x = blockinfo[3]
@@ -96,6 +105,7 @@ def vulnerability_map(
         px = b % nblock_x
         py = b // nblock_x
         # Data
+        forest_data = forest_band.ReadAsArray(x[px], y[py], nx[px], ny[py])
         dist_data = dist_band.ReadAsArray(x[px], y[py], nx[px], ny[py])
         subj_data = subj_band.ReadAsArray(x[px], y[py], nx[px], ny[py])
         # Categorize
@@ -104,9 +114,9 @@ def vulnerability_map(
                           right=True)
         cat_data = cat_data.reshape(dist_data.shape)
         cat_data = n_classes + 1 - cat_data
-        # Set classes 0 and 1
-        cat_data[dist_data == 0] = 0
+        # Set classes 1 beyond distance threshold
         cat_data[dist_data > dist_thresh] = 1
+        cat_data[forest_data == 0] = 0  # Mask with forest
         # Adding subjurisdiction info
         cat_data = cat_data * 1000 + subj_data
         cat_data[cat_data <= 1000] = 0