diff --git a/examples/Baseline.jl b/examples/Baseline.jl
new file mode 100644
index 0000000..5ad7dca
--- /dev/null
+++ b/examples/Baseline.jl
@@ -0,0 +1,211 @@
+using Revise
+using EcoSISTEM
+using EcoSISTEM.ClimatePref
+using EcoSISTEM.Units
+using Peatland
+using Unitful, Unitful.DefaultSymbols
+using Distributions
+using Plots
+using JLD2
+using DataFrames
+using Missings
+using CSV
+using Distances
+using Shapefile
+using Diversity
+
+### HISTORIC RAINFALL DATA 2010 - 2020 (REPEATED) ###
+function runDrying(; save = false, save_path = pwd())
+    JLD2.@load("data/Peat_30_spp.jld2", peat_spp)
+    JLD2.@load("data/Peat_30_moss.jld2", moss_spp)
+
+    file = "data/CF_outline.tif"
+    cf = readfile(file, 261000.0m, 266000.0m, 289000.0m, 293000.0m)
+    active = Array(.!isnan.(Array(cf)))
+    #heatmap(active)
+
+    grid = size(cf); individuals=1_000_000; area = step(cf.axes[1]) * step(cf.axes[2]) * prod(grid);
+    numMoss = nrow(moss_spp); numShrub = nrow(peat_spp); numSpecies = numMoss + numShrub
+
+    mosses = 1:numMoss
+    shrubs = findall((peat_spp[!, :Type] .== "Shrub") .| (peat_spp[!, :Type] .== "Grass") .| (peat_spp[!, :Type] .== "Herb")) .+ numMoss
+    trees = findall((peat_spp[!, :Type] .== "Tree")) .+ numMoss
+
+    height = peat_spp[!, :Plant_height] .* m
+    moss_height = moss_spp[!, :Len] .* mm
+
+    #Set up how much energy each species consumes
+    req1 = moss_height .* rand(Normal(1.0, 0.1), numMoss) .* mm ./m
+    req2 = height .* rand(Normal(10.0, 0.1), numShrub) .* mm ./m
+    energy_vec = WaterRequirement([req1; req2])
+
+    # Set rates for birth and death
+    birth = 0.1/year
+    death = 0.1/year
+    longevity = 0.1
+    survival = 0.1
+    boost = 1.0
+    # Collect model parameters together
+    param = EqualPop(birth, death, longevity, survival, boost)
+
+    # Create kernel for movement
+    kernel = fill(GaussianKernel(20.0m, 1e-5), numSpecies)
+    movement = BirthOnlyMovement(kernel, NoBoundary())
+
+    # Create species list, including their temperature preferences, seed abundance and native status
+    pref1 = rand(Normal(90.0, 13.0), numMoss) .* m^3 
+    pref2 = rand(Normal(70.0, 18.0), numShrub) .* m^3 
+    opts = [pref1; pref2]
+    vars = [fill(10.0m^3, numMoss); fill(10.0m^3, numShrub)]
+    water_traits = GaussTrait(opts, vars)
+    ele_traits = GaussTrait(fill(1.0, numSpecies), fill(20.0, numSpecies))
+    soilDict = Dict("hygrophilous" => [8, 11], "terrestrial" => [1, 4, 5], "terrestrial/hygrophilous" => [1, 4, 5, 8, 11])
+    soil_pref1 = fill([8, 11], numMoss); soil_pref2 = [soilDict[x] for x in peat_spp[!, :Habitat]]
+    soil_pref = soiltrait([soil_pref1; soil_pref2])
+    traits = TraitCollection3(water_traits, ele_traits, soil_pref)
+    native = fill(true, numSpecies)
+    p = 1 ./ energy_vec.energy 
+    abun = rand(Multinomial(individuals, p ./ sum(p)))
+    # abun = [fill(div(individuals, numMoss), numMoss); fill(0, numShrub)]
+    sppl = SpeciesList(numSpecies, traits, abun, energy_vec,
+        movement, param, native)
+
+    # Create abiotic environment - even grid of one temperature
+    file = "data/CF_elevation_smooth.tif"
+    ele = readfile(file, 261000.0m, 266000.0m, 289000.0m, 293000.0m)
+    ele.data[ele.data .< 0] .= 0
+
+    file = "data/CF_TPI_smooth.tif"
+    tpi = readfile(file, 261000.0m, 266000.0m, 289000.0m, 293000.0m)
+    tpi.data[tpi.data .< 0] .= 0
+
+    @load "data/RainfallBudget_burnin.jld2"
+    file = "data/LCM.tif"
+    soil = readfile(file, 261000.0m, 266000.0m, 289000.0m, 293000.0m)
+    soil = Int.(soil)
+    abenv = peatlandAE(ele, soil, 10.0m^3, bud, active) 
+    abenv.habitat.h1.matrix .*= tpi.data
+    abenv.habitat.h1.matrix[abenv.habitat.h1.matrix .< 10.0m^3] .= 10.0m^3
+
+    # Set relationship between species and environment (gaussian)
+    rel = multiplicativeTR3(Gauss{typeof(1.0m^3)}(), NoRelContinuous{Float64}(), soilmatch{Int64}())
+
+    # Create new transition list
+    transitions = TransitionList(true)
+    addtransition!(transitions, UpdateEnergy(update_energy_usage!))
+    addtransition!(transitions, UpdateEnvironment(update_peat_environment!))
+    active_squares = findall(active[1:end])
+    peat_squares = findall(soil .== 11)
+    # Add in species based transitions (only for active squares)
+    for spp in eachindex(sppl.species.names) 
+        for loc in active_squares
+            addtransition!(transitions, GenerateSeed(spp, loc, sppl.params.birth[spp]))
+            addtransition!(transitions, DeathProcess(spp, loc, sppl.params.death[spp]))
+            addtransition!(transitions, SeedDisperse(spp, loc))
+            addtransition!(transitions, WaterUse(spp, loc, 0.01))
+            if spp > numMoss
+                addtransition!(transitions, Invasive(spp, loc, 10.0/28days))
+            end
+        end
+    end
+    # Add in location based transitions and ditches
+    for loc in eachindex(abenv.habitat.h1.matrix)
+        # if loc ∈ ditch_locations
+        #     drainage = 1.0/month
+        # else
+            drainage = 0.001/month
+        # end
+        κ = 0.01/month
+        λ = 0.01/month
+
+        addtransition!(transitions, LateralFlow(loc, κ, λ))
+        addtransition!(transitions, WaterFlux(loc, drainage, 150.0m^3))
+    end
+
+    # Create ecosystem
+    eco = PeatSystem(sppl, abenv, rel, transitions = transitions)
+
+    # Run simulation
+    # Simulation Parameters
+    burnin = 30year; times1 = 10year; times2 = 10year; timestep = 1month;
+    record_interval = 3months
+    lensim = length(0years:record_interval:burnin)
+    lensim1 = length(0years:record_interval:times1)
+    lensim2 = length(0years:record_interval:times2)
+    abuns1 = generate_storage(eco, lensim1, 1)[:, :, :, 1]
+    abuns2 = generate_storage(eco, lensim2, 1)[:, :, :, 1]
+    # Burnin
+    abuns = generate_storage(eco, lensim, 1)[:, :, :, 1]
+    @time simulate_record!(abuns, eco, burnin, record_interval, timestep, save = save, save_path = save_path, specialise = true)
+    println(mean(eco.abenv.habitat.h1.matrix[active]))
+
+    for loc in active_squares
+        addtransition!(transitions, Peatland.Dry(loc, 0.05, 1year))
+    end
+    @time simulate_record!(abuns1, eco, times1, record_interval, timestep, save = save, save_path = save_path, specialise = true);
+    println(mean(eco.abenv.habitat.h1.matrix[active]))
+
+    for loc in active_squares
+        addtransition!(transitions, Peatland.Rewet(loc, 0.05, 1year, 150.0m^3))
+    end
+    @time simulate_record!(abuns2, eco, times2, record_interval, timestep, save = save, save_path = save_path, specialise = true);
+    println(mean(eco.abenv.habitat.h1.matrix[active]))
+
+    abuns = cat(abuns, abuns1, abuns2,  dims = 3)
+    abuns = reshape(abuns, (numSpecies, grid[1], grid[2], lensim + lensim1 + lensim2))
+    
+    return abuns
+end
+
+abuns = runDrying();
+@save "/home/claireh/sdc/Peatland/Peatland_baseline.jld2" abuns=abuns[:, :, :, [12,end]]
+
+
+plot(grid = false, label = "", layout = (3, 1), left_margin = 1.0*Plots.inch, size = (1000, 1200))
+for i in mosses
+    plot!(0:3/12:30.75,sum(abuns[i, :, :, 80:end], dims = (1,2))[1, 1, :,], grid = false, label = "", subplot = 1,
+    title = "Mosses", colour = :grey, alpha = 0.1)
+end
+annotate!([-10], [7e6], ["A"])
+plot!(0:3/12:30.75, mean(sum(abuns[mosses, :, :, 80:end], dims = (2, 3))[:, 1, 1, :], dims = 1)[1, :], grid = false, label = "", subplot = 1,
+    colour = :black)
+vline!([10, 20], subplot = 1, color = :black, lty = :dot, label = "")
+for i in shrubs
+    plot!(0:3/12:30.75,sum(abuns[i, :, :, 80:end], dims = (1, 2))[1, 1, :,], grid = false, label = "", subplot = 2,
+    title = "Shrubs, Herbs and Grasses", colour = :grey, alpha = 0.1)
+end
+plot!(0:3/12:30.75, mean(sum(abuns[shrubs, :, :, 80:end], dims = (2, 3))[:, 1, 1, :], dims = 1)[1, :], grid = false, label = "", subplot = 2,
+     colour = :black)
+vline!([10, 20], subplot = 2, color = :black, lty = :dot, label = "")
+for i in trees
+    plot!(0:3/12:30.75,sum(abuns[i, :, :, 80:end], dims = (1, 2))[1, 1, :,], grid = false, label = "", subplot = 3, 
+    title = "Trees", colour = :grey, alpha = 0.1)
+end
+plot!(0:3/12:30.75, mean(sum(abuns[trees, :, :, 80:end], dims = (2, 3))[:, 1, 1, :], dims = 1)[1, :], grid = false, label = "", subplot = 3,
+   colour = :black)
+vline!([10, 20], subplot = 3, color = :black, lty = :dot, label = "")
+Plots.pdf("Abuns_baseline.pdf")
+
+
+@load "data/Peatland_baseline.jld2"
+JLD2.@load("data/Peat_30_spp.jld2", peat_spp)
+JLD2.@load("data/Peat_30_moss.jld2", moss_spp)
+numMoss = nrow(moss_spp); numShrub = nrow(peat_spp); numSpecies = numMoss + numShrub
+mosses = 1:numMoss
+shrubs = findall((peat_spp[!, :Type] .== "Shrub") .| (peat_spp[!, :Type] .== "Grass") .| (peat_spp[!, :Type] .== "Herb")) .+ numMoss
+trees = findall((peat_spp[!, :Type] .== "Tree")) .+ numMoss
+file = "data/CF_outline.tif"
+cf = readfile(file, 261000.0m, 266000.0m, 289000.0m, 293000.0m)
+active = Array(.!isnan.(Array(cf)))
+
+heatmap(layout = (1,2), size = (1200, 1000), grid = false, aspect_ratio = 1, clim = (0, 10), titlelocation = :left,
+left_margin = 10*Plots.mm, guidefont = "Helvetica Bold", guidefontsize = 16, titlefontsize = 18)
+sumabuns = Float64.(reshape(sum(abuns[shrubs, :, :, 1], dims = 1), size(active)))
+sumabuns[.!active] .= NaN
+heatmap!(261000.0:10:266000.0, 289000.0:10:293000.0, sumabuns', subplot = 1, c = :viridis,
+title = "A")
+sumabuns = Float64.(reshape(sum(abuns[shrubs, :, :, 2], dims = 1), size(active)))
+sumabuns[.!active] .= NaN
+heatmap!(261000.0:10:266000.0, 289000.0:10:293000.0, sumabuns', subplot = 2, c = :viridis,
+title = "B")
+Plots.pdf("plots/Shrubs.pdf")
\ No newline at end of file
diff --git a/examples/Baseline_ditches.jl b/examples/Baseline_ditches.jl
new file mode 100644
index 0000000..13baed5
--- /dev/null
+++ b/examples/Baseline_ditches.jl
@@ -0,0 +1,283 @@
+using Revise
+using EcoSISTEM
+using EcoSISTEM.ClimatePref
+using EcoSISTEM.Units
+using Peatland
+using Unitful, Unitful.DefaultSymbols
+using Distributions
+using Plots
+using JLD2
+using DataFrames
+using Missings
+using CSV
+using Distances
+using Shapefile
+using Diversity
+
+### HISTORIC RAINFALL DATA 2010 - 2020 (REPEATED) ###
+function runDrying(burnin::Unitful.Time, times1::Unitful.Time, times2::Unitful.Time, timestep::Unitful.Time, record_interval::Unitful.Time; save = false, save_path = pwd())
+    JLD2.@load("data/Peat_30_spp.jld2", peat_spp)
+    JLD2.@load("data/Peat_30_moss.jld2", moss_spp)
+
+    file = "data/CF_outline.tif"
+    cf = readfile(file, 261000.0m, 266000.0m, 289000.0m, 293000.0m)
+    active = Array(.!isnan.(Array(cf)))
+    #heatmap(active)
+
+    grid = size(cf); individuals=1_000_000; area = step(cf.axes[1]) * step(cf.axes[2]) * prod(grid);
+    numMoss = nrow(moss_spp); numShrub = nrow(peat_spp); numSpecies = numMoss + numShrub
+
+    mosses = 1:numMoss
+    shrubs = findall((peat_spp[!, :Type] .== "Shrub") .| (peat_spp[!, :Type] .== "Grass") .| (peat_spp[!, :Type] .== "Herb")) .+ numMoss
+    trees = findall((peat_spp[!, :Type] .== "Tree")) .+ numMoss
+
+    height = peat_spp[!, :Plant_height] .* m
+    moss_height = moss_spp[!, :Len] .* mm
+
+    #Set up how much energy each species consumes
+    req1 = moss_height .* rand(Normal(1.0, 0.1), numMoss) .* mm ./m
+    req2 = height .* rand(Normal(10.0, 0.1), numShrub) .* mm ./m
+    energy_vec = WaterRequirement([req1; req2])
+
+    # Set rates for birth and death
+    birth = 0.1/year
+    death = 0.1/year
+    longevity = 0.1
+    survival = 0.1
+    boost = 1.0
+    # Collect model parameters together
+    param = EqualPop(birth, death, longevity, survival, boost)
+
+    # Create kernel for movement
+    kernel = fill(GaussianKernel(20.0m, 1e-5), numSpecies)
+    movement = BirthOnlyMovement(kernel, NoBoundary())
+
+    # Create species list, including their temperature preferences, seed abundance and native status
+    pref1 = rand(Normal(90.0, 13.0), numMoss) .* m^3 
+    pref2 = rand(Normal(70.0, 18.0), numShrub) .* m^3 
+    opts = [pref1; pref2]
+    vars = [fill(10.0m^3, numMoss); fill(10.0m^3, numShrub)]
+    water_traits = GaussTrait(opts, vars)
+    ele_traits = GaussTrait(fill(1.0, numSpecies), fill(20.0, numSpecies))
+    soilDict = Dict("hygrophilous" => [8, 11], "terrestrial" => [1, 4, 5], "terrestrial/hygrophilous" => [1, 4, 5, 8, 11])
+    soil_pref1 = fill([8, 11], numMoss); soil_pref2 = [soilDict[x] for x in peat_spp[!, :Habitat]]
+    soil_pref = soiltrait([soil_pref1; soil_pref2])
+    traits = TraitCollection3(water_traits, ele_traits, soil_pref)
+    native = fill(true, numSpecies)
+    p = 1 ./ energy_vec.energy 
+    abun = rand(Multinomial(individuals, p ./ sum(p)))
+    # abun = [fill(div(individuals, numMoss), numMoss); fill(0, numShrub)]
+    sppl = SpeciesList(numSpecies, traits, abun, energy_vec,
+        movement, param, native)
+
+    # Create abiotic environment - even grid of one temperature
+    file = "data/CF_elevation_smooth.tif"
+    ele = readfile(file, 261000.0m, 266000.0m, 289000.0m, 293000.0m)
+    ele.data[ele.data .< 0] .= 0
+
+    file = "data/CF_TPI_smooth.tif"
+    tpi = readfile(file, 261000.0m, 266000.0m, 289000.0m, 293000.0m)
+    tpi.data[tpi.data .< 0] .= 0
+
+    @load "data/RainfallBudget_burnin.jld2"
+    file = "data/LCM.tif"
+    soil = readfile(file, 261000.0m, 266000.0m, 289000.0m, 293000.0m)
+    soil = Int.(soil)
+    abenv = peatlandAE(ele, soil, 10.0m^3, bud, active) 
+    abenv.habitat.h1.matrix .*= tpi.data
+    abenv.habitat.h1.matrix[abenv.habitat.h1.matrix .< 10.0m^3] .= 10.0m^3
+
+    # Set relationship between species and environment (gaussian)
+    rel = multiplicativeTR3(Gauss{typeof(1.0m^3)}(), NoRelContinuous{Float64}(), soilmatch{Int64}())
+
+    # Create new transition list
+    transitions = TransitionList(true)
+    addtransition!(transitions, UpdateEnergy(update_energy_usage!))
+    addtransition!(transitions, UpdateEnvironment(update_peat_environment!))
+    active_squares = findall(active[1:end])
+    peat_squares = findall(soil .== 11)
+    # Add in species based transitions (only for active squares)
+    for spp in eachindex(sppl.species.names) 
+        for loc in active_squares
+            addtransition!(transitions, GenerateSeed(spp, loc, sppl.params.birth[spp]))
+            addtransition!(transitions, DeathProcess(spp, loc, sppl.params.death[spp]))
+            addtransition!(transitions, SeedDisperse(spp, loc))
+            addtransition!(transitions, WaterUse(spp, loc, 0.01))
+            if spp > numMoss
+                addtransition!(transitions, Invasive(spp, loc, 10.0/28days))
+            end
+        end
+    end
+    # Add in location based transitions and ditches
+    for loc in eachindex(abenv.habitat.h1.matrix)
+        # if loc ∈ ditch_locations
+        #     drainage = 1.0/month
+        # else
+            drainage = 0.001/month
+        # end
+        κ = 0.01/month
+        λ = 0.01/month
+
+        addtransition!(transitions, LateralFlow(loc, κ, λ))
+        addtransition!(transitions, WaterFlux(loc, drainage, 150.0m^3))
+    end
+
+    # Create ecosystem
+    eco = PeatSystem(sppl, abenv, rel, transitions = transitions)
+
+    # Run simulation
+    # Simulation Parameters
+    lensim = length(0years:record_interval:burnin)
+    lensim1 = length(0years:record_interval:times1)
+    lensim2 = length(0years:record_interval:times2)
+    abuns1 = generate_storage(eco, lensim1, 1)[:, :, :, 1]
+    abuns2 = generate_storage(eco, lensim2, 1)[:, :, :, 1]
+    # Burnin
+    abuns = generate_storage(eco, lensim, 1)[:, :, :, 1]
+    @time simulate_record!(abuns, eco, burnin, record_interval, timestep, save = save, save_path = save_path, specialise = true)
+    println(mean(eco.abenv.habitat.h1.matrix[active]))
+
+    file = "data/Ditches_full.tif"
+    ditches = readfile(file, 289000.0m, 293000.0m, 261000.0m, 266000.0m)
+    ditch_locations = findall(.!isnan.(ditches))
+    locs = [convert_coords(d[1], d[2], size(cf, 1)) for d in ditch_locations]
+    abenv.habitat.h1.matrix[locs] .= 0.0m^3
+    abenv.habitat.h2.matrix[locs] .= 0
+    abenv.active[ditch_locations] .= false
+
+    @time simulate_record!(abuns1, eco, times1, record_interval, timestep, save = save, save_path = save_path, specialise = true);
+    println(mean(eco.abenv.habitat.h1.matrix[active]))
+
+    abenv.habitat.h2.matrix[locs] .= ele[locs]
+    abenv.active[ditch_locations] .= false
+
+    @time simulate_record!(abuns2, eco, times2, record_interval, timestep, save = save, save_path = save_path, specialise = true);
+    println(mean(eco.abenv.habitat.h1.matrix[active]))
+
+    abuns = cat(abuns, abuns1, abuns2,  dims = 3)
+    abuns = reshape(abuns, (numSpecies, grid[1], grid[2], lensim + lensim1 + lensim2))
+    
+    return abuns
+end
+
+burnin = 70year; times1 = 30year; times2 = 30year; timestep = 1month; record_interval = 3months
+abuns = runDrying(burnin, times1, times2, timestep, record_interval);
+@save "/home/claireh/sdc/Peatland/Peatland_baseline_ditches.jld2" abuns=abuns[:, :, :, [50*4, 70*4, 90*4, end]]
+
+
+times = ustrip.(uconvert.(year, 20years:record_interval:((burnin+times1+times2)+record_interval*3)))
+plot(grid = false, label = "", layout = (3, 1), left_margin = 1.0*Plots.inch, size = (1000, 1200))
+for i in mosses
+    plot!(times,sum(abuns[i, :, :, 80:end], dims = (1,2))[1, 1, :,], grid = false, label = "", subplot = 1,
+    title = "Mosses", colour = :grey, alpha = 0.1)
+end
+plot!(times, mean(sum(abuns[mosses, :, :, 80:end], dims = (2, 3))[:, 1, 1, :], dims = 1)[1, :], grid = false, label = "", subplot = 1,
+    colour = :black)
+vline!([burnin/year, (burnin+times1)/year], subplot = 1, color = :black, lty = :dot, label = "")
+for i in shrubs
+    plot!(times,sum(abuns[i, :, :, 80:end], dims = (1, 2))[1, 1, :,], grid = false, label = "", subplot = 2,
+    title = "Shrubs, Herbs and Grasses", colour = :grey, alpha = 0.1)
+end
+plot!(times, mean(sum(abuns[shrubs, :, :, 80:end], dims = (2, 3))[:, 1, 1, :], dims = 1)[1, :], grid = false, label = "", subplot = 2,
+     colour = :black)
+vline!([burnin/year, (burnin+times1)/year], subplot = 2, color = :black, lty = :dot, label = "")
+for i in trees
+    plot!(times,sum(abuns[i, :, :, 80:end], dims = (1, 2))[1, 1, :,], grid = false, label = "", subplot = 3, 
+    title = "Trees", colour = :grey, alpha = 0.1)
+end
+plot!(times, mean(sum(abuns[trees, :, :, 80:end], dims = (2, 3))[:, 1, 1, :], dims = 1)[1, :], grid = false, label = "", subplot = 3,
+   colour = :black)
+vline!([burnin/year, (burnin+times1)/year], subplot = 3, color = :black, lty = :dot, label = "")
+Plots.pdf("Abuns_baseline_ditches.pdf")
+
+
+@load "data/Peatland_baseline_ditches.jld2"
+JLD2.@load("data/Peat_30_spp.jld2", peat_spp)
+JLD2.@load("data/Peat_30_moss.jld2", moss_spp)
+numMoss = nrow(moss_spp); numShrub = nrow(peat_spp); numSpecies = numMoss + numShrub
+mosses = 1:numMoss
+shrubs = findall((peat_spp[!, :Type] .== "Shrub") .| (peat_spp[!, :Type] .== "Grass") .| (peat_spp[!, :Type] .== "Herb")) .+ numMoss
+trees = findall((peat_spp[!, :Type] .== "Tree")) .+ numMoss
+file = "data/CF_outline.tif"
+cf = readfile(file, 261000.0m, 266000.0m, 289000.0m, 293000.0m)
+active = Array(.!isnan.(Array(cf)))
+
+heatmap(layout = (2,2), size = (1200, 1000), grid = false, aspect_ratio = 1, clim = (0, 600), titlelocation = :left,
+left_margin = 10*Plots.mm, guidefont = "Helvetica Bold", guidefontsize = 16, titlefontsize = 18)
+sumabuns = Float64.(reshape(sum(abuns[mosses, :, :, 1], dims = 1), size(active)))
+sumabuns[.!active] .= NaN
+heatmap!(261000.0:10:266000.0, 289000.0:10:293000.0, sumabuns', subplot = 1, c = :viridis,
+title = "A")
+sumabuns = Float64.(reshape(sum(abuns[1:numMoss, :, :, 2], dims = 1), size(active)))
+sumabuns[.!active] .= NaN
+heatmap!(261000.0:10:266000.0, 289000.0:10:293000.0, sumabuns', subplot = 2, c = :viridis,
+title = "B")
+sumabuns = Float64.(reshape(sum(abuns[1:numMoss, :, :, 3], dims = 1), size(active)))
+sumabuns[.!active] .= NaN
+heatmap!(261000.0:10:266000.0, 289000.0:10:293000.0, sumabuns', subplot = 3, c = :viridis,
+title = "C")
+sumabuns = Float64.(reshape(sum(abuns[1:numMoss, :, :, end], dims = 1), size(active)))
+sumabuns[.!active] .= NaN
+heatmap!(261000.0:10:266000.0, 289000.0:10:293000.0, sumabuns', subplot = 4, c = :viridis,
+title = "D")
+Plots.pdf("plots/Mosses_baseline.pdf")
+
+heatmap(layout = (2,2), size = (1200, 1000), grid = false, aspect_ratio = 1, clim = (0, 10), titlelocation = :left,
+left_margin = 10*Plots.mm, guidefont = "Helvetica Bold", guidefontsize = 16, titlefontsize = 18)
+sumabuns = Float64.(reshape(sum(abuns[shrubs, :, :, 1], dims = 1), size(active)))
+sumabuns[.!active] .= NaN
+heatmap!(261000.0:10:266000.0, 289000.0:10:293000.0, sumabuns', subplot = 1, c = :viridis,
+title = "A")
+sumabuns = Float64.(reshape(sum(abuns[shrubs, :, :, 2], dims = 1), size(active)))
+sumabuns[.!active] .= NaN
+heatmap!(261000.0:10:266000.0, 289000.0:10:293000.0, sumabuns', subplot = 2, c = :viridis,
+title = "B")
+sumabuns = Float64.(reshape(sum(abuns[shrubs, :, :, 3], dims = 1), size(active)))
+sumabuns[.!active] .= NaN
+heatmap!(261000.0:10:266000.0, 289000.0:10:293000.0, sumabuns', subplot = 3, c = :viridis,
+title = "C")
+sumabuns = Float64.(reshape(sum(abuns[shrubs, :, :, end], dims = 1), size(active)))
+sumabuns[.!active] .= NaN
+heatmap!(261000.0:10:266000.0, 289000.0:10:293000.0, sumabuns', subplot = 4, c = :viridis,
+title = "D")
+Plots.pdf("plots/Shrubs_baseline.pdf")
+
+heatmap(layout = (2,2), size = (1200, 1000), grid = false, aspect_ratio = 1, clim = (0, 20), titlelocation = :left,
+left_margin = 10*Plots.mm, guidefont = "Helvetica Bold", guidefontsize = 16, titlefontsize = 18)
+sumabuns = Float64.(reshape(sum(abuns[shrubs, :, :, 1], dims = 1), size(active)))
+sumabuns[.!active] .= NaN
+heatmap!(261000.0:10:266000.0, 289000.0:10:293000.0, sumabuns', subplot = 1, c = :viridis,
+title = "A")
+sumabuns = Float64.(reshape(sum(abuns[shrubs, :, :, 2], dims = 1), size(active)))
+sumabuns[.!active] .= NaN
+heatmap!(261000.0:10:266000.0, 289000.0:10:293000.0, sumabuns', subplot = 2, c = :viridis,
+title = "B")
+sumabuns = Float64.(reshape(sum(abuns[shrubs, :, :, 3], dims = 1), size(active)))
+sumabuns[.!active] .= NaN
+heatmap!(261000.0:10:266000.0, 289000.0:10:293000.0, sumabuns', subplot = 3, c = :viridis,
+title = "C")
+sumabuns = Float64.(reshape(sum(abuns[shrubs, :, :, end], dims = 1), size(active)))
+sumabuns[.!active] .= NaN
+heatmap!(261000.0:10:266000.0, 289000.0:10:293000.0, sumabuns', subplot = 4, c = :viridis,
+title = "D")
+Plots.pdf("plots/Shrubs_baseline.pdf")
+
+heatmap(layout = (2,2), size = (1200, 1000), grid = false, aspect_ratio = 1, clim = (0, 2), titlelocation = :left,
+left_margin = 10*Plots.mm, guidefont = "Helvetica Bold", guidefontsize = 16, titlefontsize = 18)
+sumabuns = Float64.(reshape(sum(abuns[trees, :, :, 1], dims = 1), size(active)))
+sumabuns[.!active] .= NaN
+heatmap!(261000.0:10:266000.0, 289000.0:10:293000.0, sumabuns', subplot = 1, c = :viridis,
+title = "A")
+sumabuns = Float64.(reshape(sum(abuns[trees, :, :, 2], dims = 1), size(active)))
+sumabuns[.!active] .= NaN
+heatmap!(261000.0:10:266000.0, 289000.0:10:293000.0, sumabuns', subplot = 2, c = :viridis,
+title = "B")
+sumabuns = Float64.(reshape(sum(abuns[trees, :, :, 3], dims = 1), size(active)))
+sumabuns[.!active] .= NaN
+heatmap!(261000.0:10:266000.0, 289000.0:10:293000.0, sumabuns', subplot = 3, c = :viridis,
+title = "C")
+sumabuns = Float64.(reshape(sum(abuns[trees, :, :, end], dims = 1), size(active)))
+sumabuns[.!active] .= NaN
+heatmap!(261000.0:10:266000.0, 289000.0:10:293000.0, sumabuns', subplot = 4, c = :viridis,
+title = "D")
+Plots.pdf("plots/Trees_baseline.pdf")
\ No newline at end of file
diff --git a/examples/Hydrology.jl b/examples/Hydrology.jl
new file mode 100644
index 0000000..3e57d92
--- /dev/null
+++ b/examples/Hydrology.jl
@@ -0,0 +1,171 @@
+using Revise
+using EcoSISTEM
+using EcoSISTEM.ClimatePref
+using EcoSISTEM.Units
+using Peatland
+using Unitful, Unitful.DefaultSymbols
+using Distributions
+using Plots
+using JLD2
+using DataFrames
+using Missings
+using CSV
+using Distances
+using Shapefile
+using Diversity
+
+function buildEco(timestep::Unitful.Time; ditch = true)
+    JLD2.@load("data/Peat_30_spp.jld2", peat_spp)
+    JLD2.@load("data/Peat_30_moss.jld2", moss_spp)
+
+    moss_spp = filter(f -> (f.Len == findmax(moss_spp.Len)[1]) || (f.Len == findmin(moss_spp.Len)[1]), moss_spp)
+    peat_spp = filter(f -> (f.Plant_height == findmax(peat_spp.Plant_height)[1]) || (f.Plant_height == findmin(peat_spp.Plant_height)[1]), peat_spp)
+
+    file = "data/CF_outline.tif"
+    cf = readfile(file, 261000.0m, 266000.0m, 289000.0m, 293000.0m)
+    active = Array(.!isnan.(Array(cf)))
+    #heatmap(active)
+
+    grid = size(cf); individuals=1_000_000; area = step(cf.axes[1]) * step(cf.axes[2]) * prod(grid);
+    numMoss = nrow(moss_spp); numShrub = nrow(peat_spp); numSpecies = numMoss + numShrub
+
+    mosses = 1:numMoss
+    shrubs = findall((peat_spp[!, :Type] .== "Shrub") .| (peat_spp[!, :Type] .== "Grass") .| (peat_spp[!, :Type] .== "Herb")) .+ numMoss
+    trees = findall((peat_spp[!, :Type] .== "Tree")) .+ numMoss
+
+    height = peat_spp[!, :Plant_height] .* m
+    moss_height = moss_spp[!, :Len] .* mm
+
+    #Set up how much energy each species consumes
+    req1 = moss_height .* rand(Normal(1, 0.1), numMoss) .* mm ./m
+    req2 = height .* rand(Normal(10.0, 0.1), numShrub) .* mm ./m
+    energy_vec = WaterRequirement([req1; req2])
+
+    # Set rates for birth and death
+    birth = 0.1/year
+    death = 0.1/year
+    longevity = 0.1
+    survival = 0.01
+    boost = 1.0
+    # Collect model parameters together
+    param = EqualPop(birth, death, longevity, survival, boost)
+
+    # Create kernel for movement
+    kernel = fill(GaussianKernel(20.0m, 1e-5), numSpecies)
+    movement = BirthOnlyMovement(kernel, NoBoundary())
+
+    # Create species list, including their temperature preferences, seed abundance and native status
+    pref1 = rand(Normal(9.0, 1.30), numMoss) .* m^3 
+    pref2 = rand(Normal(7.0, 1.80), numShrub) .* m^3 
+    opts = [pref1; pref2]
+    vars = [fill(1.0m^3, numMoss); fill(1.0m^3, numShrub)]
+    water_traits = GaussTrait(opts, vars)
+    ele_traits = GaussTrait(fill(1.0, numSpecies), fill(20.0, numSpecies))
+    soilDict = Dict("hygrophilous" => [8, 11], "terrestrial" => [1, 4, 5], "terrestrial/hygrophilous" => [1, 4, 5, 8, 11])
+    soil_pref1 = fill([8, 11], numMoss); soil_pref2 = [soilDict[x] for x in peat_spp[!, :Habitat]]
+    soil_pref = soiltrait([soil_pref1; soil_pref2])
+    traits = TraitCollection3(water_traits, ele_traits, soil_pref)
+    native = fill(true, numSpecies)
+    p = 1 ./ energy_vec.energy 
+    abun = rand(Multinomial(individuals, p ./ sum(p)))
+    # abun = [fill(div(individuals, numMoss), numMoss); fill(0, numShrub)]
+    sppl = SpeciesList(numSpecies, traits, abun, energy_vec,
+        movement, param, native)
+
+    # Create abiotic environment - even grid of one temperature
+    file = "data/CF_elevation_smooth.tif"
+    ele = readfile(file, 261000.0m, 266000.0m, 289000.0m, 293000.0m)
+    ele.data[ele.data .< 0] .= 0
+
+    file = "data/CF_TPI_smooth.tif"
+    tpi = readfile(file, 261000.0m, 266000.0m, 289000.0m, 293000.0m)
+    tpi.data[tpi.data .< 0] .= 0
+
+    @load "data/RainfallBudget_burnin.jld2"
+    bud.matrix .*= (timestep / month)
+    file = "data/LCM.tif"
+    soil = readfile(file, 261000.0m, 266000.0m, 289000.0m, 293000.0m)
+    soil = Int.(soil)
+    abenv = peatlandAE(ele, soil, 1.0m^3, bud, active) 
+    abenv.habitat.h1.matrix .*= tpi.data
+    abenv.habitat.h1.matrix[abenv.habitat.h1.matrix .< 1.0m^3] .= 1.0m^3
+
+    # If there are ditches, make sure they are lower than everything else around and filled with water
+    if ditch
+        file = "data/Ditches_full.tif"
+        ditches = readfile(file, 289000.0m, 293000.0m, 261000.0m, 266000.0m)
+        ditch_locations = findall(.!isnan.(ditches))
+        ditch_locs = [convert_coords(d[1], d[2], size(cf, 1)) for d in ditch_locations]
+        file = "data/MainRivers.tif"
+        rivers = readfile(file, 289000.0m, 293000.0m, 261000.0m, 266000.0m)
+        river_locations = findall(.!isnan.(rivers))
+        river_locs = [convert_coords(r[1], r[2], size(cf, 1)) for r in river_locations]
+        locs = [ditch_locs; river_locs ...]
+        # locs = ditch_locs
+        abenv.habitat.h1.matrix[locs] .= 0.0m^3
+        abenv.habitat.h2.matrix[locs] .= 0
+        abenv.active[ditch_locations] .= false
+        abenv.active[river_locations] .= false
+    else
+        ditch_locs = []
+        river_locs = []
+    end
+    # heatmap(ustrip.(abenv.habitat.h1.matrix)')
+
+    # Set relationship between species and environment (gaussian)
+    rel = multiplicativeTR3(Gauss{typeof(1.0m^3)}(), NoRelContinuous{Float64}(), soilmatch{Int64}())
+
+    # Create new transition list
+    transitions = TransitionList(true)
+    addtransition!(transitions, UpdateEnergy(update_energy_usage!))
+    addtransition!(transitions, UpdateEnvironment(update_peat_environment!))
+    active_squares = findall(active[1:end])
+    peat_squares = findall(soil .== 11)
+    # Add in species based transitions (only for active squares)
+    for spp in eachindex(sppl.species.names) 
+        for loc in active_squares
+            addtransition!(transitions, GenerateSeed(spp, loc, sppl.params.birth[spp]))
+            addtransition!(transitions, DeathProcess(spp, loc, sppl.params.death[spp]))
+            addtransition!(transitions, SeedDisperse(spp, loc))
+            addtransition!(transitions, WaterUse(spp, loc, 0.01))
+            if spp > numMoss
+                addtransition!(transitions, Invasive(spp, loc, 10.0/28days))
+            end
+        end
+    end
+    # Add in location based transitions and ditches
+    for loc in eachindex(abenv.habitat.h1.matrix)
+        if loc ∈ ditch_locs
+            drainage = 1.0/month
+        elseif loc ∈ river_locs
+            drainage = 1.0/month
+        else
+            drainage = 0.0001/month
+        end
+        infiltration = 0.2/month
+        evaporation = 0.7/month
+        κ = 0.001m^2/month
+        ν = 0.001m^2/month
+
+        addtransition!(transitions, LateralFlow(loc, κ, ν, 100.0m^3))
+        addtransition!(transitions, WaterFlux(loc, drainage, infiltration, evaporation, 100.0m^3))
+    end
+
+    transitions = specialise_transition_list(transitions)
+    # Create ecosystem
+    eco = PeatSystem(sppl, abenv, rel, transitions = transitions)
+    return eco
+end
+timestep = 1month
+eco = buildEco(timestep)
+envs = zeros(10)
+for i in 1:10
+    EcoSISTEM.update!(eco, timestep, eco.transitions)
+    envs[i] = ustrip.(mean(eco.abenv.habitat.h1.matrix[eco.abenv.active]))
+end
+heatmap(ustrip.(eco.abenv.habitat.h1.matrix)')#, size = (1000, 600), layout = 2)
+# heatmap!(ustrip.(eco.cache.surfacewater)', clim = (0, 10), subplot = 2)
+plot(envs)
+
+eco.abenv.habitat.h1.matrix[.!eco.abenv.active] .= NaN*m^3
+heatmap(ustrip.(eco.abenv.habitat.h1.matrix)', clim = (0, 20))
\ No newline at end of file
diff --git a/examples/Landcover_plot.jl b/examples/Landcover_plot.jl
new file mode 100644
index 0000000..e8fb07b
--- /dev/null
+++ b/examples/Landcover_plot.jl
@@ -0,0 +1,38 @@
+using EcoSISTEM
+using EcoSISTEM.ClimatePref
+using EcoSISTEM.Units
+using Peatland
+using Unitful, Unitful.DefaultSymbols
+using Distributions
+using Plots
+using JLD2
+using DataFrames
+using Missings
+using CSV
+using Distances
+using Shapefile
+using Diversity
+
+gr()
+file = "data/LCM.tif"
+colours = ["red", "darkgreen", "brown", "green", "lightgreen", "green3", "yellow4", "yellow",
+            "purple", "pink", "turquoise3", "lavender", "navy", "blue", "gold", "gold", "lightyellow","lightyellow", "lightblue", "black", "grey"]
+labels = ["Broadleaved woodland", "Coniferous woodland", "Arable",
+"Improved grassland", "Neutral grassland", "Calcareous grassland", "Acid grassland",
+"Fen, Marsh, Swamp", "Heather",
+"Heather grassland", "Bog",
+"Inland Rock", "Saltwater",
+"Freshwater", "Supra-littoral rock",
+"Supra-littoral sediment", "Littoral rock",
+"Littoral sediment", "Saltmarsh", "Urban",
+"Suburban"]
+soil = readfile(file, 261000.0m, 266000.0m, 289000.0m, 293000.0m)
+heatmap(soil', c = colours)
+heatmap!(ditches', c = :white)
+
+file = "data/CF_elevation.tif"
+ele = readfile(file, 261000.0m, 266000.0m, 289000.0m, 293000.0m)
+heatmap(log.(1 .+ ele)',  cbar = true)
+file = "data/Ditches2.tif"
+ditches = readfile(file, 261000.0m, 266000.0m, 289000.0m, 293000.0m)
+heatmap!(ditches', c = :white)
diff --git a/examples/Sensitivity.jl b/examples/Sensitivity.jl
new file mode 100644
index 0000000..1b71b90
--- /dev/null
+++ b/examples/Sensitivity.jl
@@ -0,0 +1,160 @@
+using Revise
+using EcoSISTEM
+using EcoSISTEM.ClimatePref
+using EcoSISTEM.Units
+using Peatland
+using Unitful, Unitful.DefaultSymbols
+using Distributions
+using Plots
+using JLD2
+using DataFrames
+using Missings
+using CSV
+using Distances
+using Shapefile
+using Diversity
+
+### HISTORIC RAINFALL DATA 2010 - 2020 (REPEATED) ###
+function runPast(paramDict::Dict, ditch = false; reps = 10, save = false, save_path = pwd())
+    JLD2.@load("data/Peat_30_spp.jld2", peat_spp)
+    JLD2.@load("data/Peat_30_moss.jld2", moss_spp)
+
+    file = "data/CF_outline.tif"
+    cf = readfile(file, 261000.0m, 266000.0m, 289000.0m, 293000.0m)
+    active = Array(.!isnan.(Array(cf)))
+    #heatmap(active)
+
+    grid = size(cf); individuals=1_000_000; area = step(cf.axes[1]) * step(cf.axes[2]) * prod(grid);
+    numMoss = nrow(moss_spp); numShrub = nrow(peat_spp); numSpecies = numMoss + numShrub
+
+    mosses = 1:numMoss
+    shrubs = findall((peat_spp[!, :Type] .== "Shrub") .| (peat_spp[!, :Type] .== "Grass") .| (peat_spp[!, :Type] .== "Herb")) .+ numMoss
+    trees = findall((peat_spp[!, :Type] .== "Tree")) .+ numMoss
+
+    height = peat_spp[!, :Plant_height] .* m
+    moss_height = moss_spp[!, :Len] .* mm
+
+    #Set up how much energy each species consumes
+    req1 = moss_height .* rand(Normal(1.0, 0.1), numMoss) .* mm ./m
+    req2 = height .* rand(Normal(10.0, 0.1), numShrub) .* mm ./m
+    energy_vec = WaterRequirement([req1; req2])
+
+
+    # Set rates for birth and death
+    birth = paramDict.birth
+    death = paramDict.death
+    longevity = paramDict.longevity
+    survival = paramDict.survival
+    boost = 1.0
+    # Collect model parameters together
+    param = EqualPop(birth, death, longevity, survival, boost)
+
+    # Create kernel for movement
+    kernel = fill(GaussianKernel(paramDict.move, 1e-5), numSpecies)
+    movement = BirthOnlyMovement(kernel, NoBoundary())
+
+    # Create species list, including their temperature preferences, seed abundance and native status
+    pref1 = rand(Normal(90.0, 13.0), numMoss) .* m^3 
+    pref2 = rand(Normal(70.0, 18.0), numShrub) .* m^3 
+    opts = [pref1; pref2]
+    vars = [fill(10.0m^3, numMoss); fill(10.0m^3, numShrub)]
+    water_traits = GaussTrait(opts, vars)
+    ele_traits = GaussTrait(fill(1.0, numSpecies), fill(20.0, numSpecies))
+    soilDict = Dict("hygrophilous" => [8, 11], "terrestrial" => [1, 4, 5], "terrestrial/hygrophilous" => [1, 4, 5, 8, 11])
+    soil_pref1 = fill([8, 11], numMoss); soil_pref2 = [soilDict[x] for x in peat_spp[!, :Habitat]]
+    soil_pref = soiltrait([soil_pref1; soil_pref2])
+    traits = TraitCollection3(water_traits, ele_traits, soil_pref)
+    native = fill(true, numSpecies)
+    p = 1 ./ energy_vec.energy 
+    abun = rand(Multinomial(individuals, p ./ sum(p)))
+    # abun = [fill(div(individuals, numMoss), numMoss); fill(0, numShrub)]
+    sppl = SpeciesList(numSpecies, traits, abun, energy_vec,
+        movement, param, native)
+    for i in reps
+        # Create abiotic environment - even grid of one temperature
+        file = "data/CF_elevation_smooth.tif"
+        ele = readfile(file, 261000.0m, 266000.0m, 289000.0m, 293000.0m)
+        ele.data[ele.data .< 0] .= 0
+
+        file = "data/CF_TPI_smooth.tif"
+        tpi = readfile(file, 261000.0m, 266000.0m, 289000.0m, 293000.0m)
+        tpi.data[tpi.data .< 0] .= 0
+
+        @load "data/RainfallBudget_burnin.jld2"
+        file = "data/LCM.tif"
+        soil = readfile(file, 261000.0m, 266000.0m, 289000.0m, 293000.0m)
+        soil = Int.(soil)
+        abenv = peatlandAE(ele, soil, 10.0m^3, bud, active) 
+        abenv.habitat.h1.matrix .*= tpi.data
+        abenv.habitat.h1.matrix[abenv.habitat.h1.matrix .< 10.0m^3] .= 10.0m^3
+
+        # If there are ditches, make sure they are lower than everything else around and filled with water
+        if ditch
+            file = "data/Ditches_full.tif"
+            ditches = readfile(file, 289000.0m, 293000.0m, 261000.0m, 266000.0m)
+            ditch_locations = findall(.!isnan.(ditches))
+            locs = [convert_coords(d[1], d[2], size(cf, 1)) for d in ditch_locations]
+            abenv.habitat.h1.matrix[locs] .= 0.0m^3
+            abenv.habitat.h2.matrix[locs] .= 0
+            abenv.active[ditch_locations] .= false
+        else
+            ditch_locations = []
+        end
+        # heatmap(ustrip.(abenv.habitat.h1.matrix)')
+
+        # Set relationship between species and environment (gaussian)
+        rel = multiplicativeTR3(Gauss{typeof(1.0m^3)}(), NoRelContinuous{Float64}(), soilmatch{Int64}())
+
+        # Create new transition list
+        transitions = TransitionList(true)
+        addtransition!(transitions, UpdateEnergy(update_energy_usage!))
+        addtransition!(transitions, UpdateEnvironment(update_peat_environment!))
+        active_squares = findall(active[1:end])
+        peat_squares = findall(soil .== 11)
+        # Add in species based transitions (only for active squares)
+        for spp in eachindex(sppl.species.names) 
+            for loc in active_squares
+                addtransition!(transitions, GenerateSeed(spp, loc, sppl.params.birth[spp]))
+                addtransition!(transitions, DeathProcess(spp, loc, sppl.params.death[spp]))
+                addtransition!(transitions, SeedDisperse(spp, loc))
+                addtransition!(transitions, WaterUse(spp, loc, 0.01))
+                if spp > numMoss
+                    addtransition!(transitions, Invasive(spp, loc, 10.0/28days))
+                end
+            end
+        end
+        # Add in location based transitions and ditches
+        for loc in eachindex(abenv.habitat.h1.matrix)
+            # if loc ∈ ditch_locations
+            #     drainage = 1.0/month
+            # else
+                drainage = 0.029/month
+            # end
+            κ = 0.01m^2/month
+            λ = 0.01m^2/month
+
+            addtransition!(transitions, LateralFlow(loc, κ, λ))
+            addtransition!(transitions, WaterFlux(loc, drainage, 150.0m^3))
+        end
+
+        transitions = specialise_transition_list(transitions)
+        # Create ecosystem
+        eco = PeatSystem(sppl, abenv, rel, transitions = transitions)
+
+        # Run simulation
+        # Simulation Parameters
+        burnin = 30year; times = 70year; timestep = 1year;
+        record_interval = 1year
+        lensim = length(0years:record_interval:times)
+        # Burnin
+        abuns = generate_storage(eco, lensim, 10)[:, :, :, 1]
+        @time simulate!(eco, burnin, timestep, specialise = true);
+        @time simulate_record!(abuns[:, :, i], eco, times, record_interval, timestep, save = save, save_path = save_path, specialise = true);
+    end
+    return abuns
+end
+
+## WITHOUT DITCHES ##
+for i in 1:10
+    paramDict = Dict()
+    abuns = runPast();
\ No newline at end of file
diff --git a/examples/Shocks.jl b/examples/Shocks.jl
new file mode 100644
index 0000000..ac3e3ee
--- /dev/null
+++ b/examples/Shocks.jl
@@ -0,0 +1,370 @@
+using Revise
+using EcoSISTEM
+using EcoSISTEM.ClimatePref
+using EcoSISTEM.Units
+using Peatland
+using Unitful, Unitful.DefaultSymbols
+using Distributions
+using Plots
+using JLD2
+using DataFrames
+using Missings
+using CSV
+using Distances
+using Shapefile
+using Diversity
+### FUTURE RAINFALL SCENARIO 2010 - 2080 plus summer drought in year 2060 ###
+
+function runFuture(ditch::Bool = false; save = false, save_path = pwd())
+    JLD2.@load("data/Peat_30_spp.jld2", peat_spp)
+    JLD2.@load("data/Peat_30_moss.jld2", moss_spp)
+
+    file = "data/CF_outline.tif"
+    cf = readfile(file, 261000.0m, 266000.0m, 289000.0m, 293000.0m)
+    active = Array(.!isnan.(Array(cf)))
+    #heatmap(active)
+
+    grid = size(cf); individuals=1_000_000; area = step(cf.axes[1]) * step(cf.axes[2]) * prod(grid);
+    numMoss = nrow(moss_spp); numShrub = nrow(peat_spp); numSpecies = numMoss + numShrub
+
+    mosses = 1:numMoss
+    shrubs = findall((peat_spp[!, :Type] .== "Shrub") .| (peat_spp[!, :Type] .== "Grass") .| (peat_spp[!, :Type] .== "Herb")) .+ numMoss
+    trees = findall((peat_spp[!, :Type] .== "Tree")) .+ numMoss
+
+    height = peat_spp[!, :Plant_height] .* m
+    moss_height = moss_spp[!, :Len] .* mm
+
+    #Set up how much energy each species consumes
+    req1 = moss_height .* rand(Normal(1.0, 0.1), numMoss) .* mm ./m
+    req2 = height .* rand(Normal(10.0, 0.1), numShrub) .* mm ./m
+    energy_vec = WaterRequirement([req1; req2])
+
+    # Set rates for birth and death
+    birth = 0.1/year
+    death = 0.1/year
+    longevity = 0.1
+    survival = 0.01
+    boost = 1.0
+    # Collect model parameters together
+    param = EqualPop(birth, death, longevity, survival, boost)
+
+    # Create kernel for movement
+    kernel = fill(GaussianKernel(20.0m, 1e-5), numSpecies)
+    movement = BirthOnlyMovement(kernel, NoBoundary())
+
+    # Create species list, including their temperature preferences, seed abundance and native status
+    pref1 = rand(Normal(90.0, 13.0), numMoss) .* m^3 
+    pref2 = rand(Normal(70.0, 18.0), numShrub) .* m^3 
+    opts = [pref1; pref2]
+    vars = [fill(10.0m^3, numMoss); fill(10.0m^3, numShrub)]
+    water_traits = GaussTrait(opts, vars)
+    ele_traits = GaussTrait(fill(1.0, numSpecies), fill(20.0, numSpecies))
+    soilDict = Dict("hygrophilous" => [8, 11], "terrestrial" => [1, 4, 5], "terrestrial/hygrophilous" => [1, 4, 5, 8, 11])
+    soil_pref1 = fill([8, 11], numMoss); soil_pref2 = [soilDict[x] for x in peat_spp[!, :Habitat]]
+    soil_pref = soiltrait([soil_pref1; soil_pref2])
+    traits = TraitCollection3(water_traits, ele_traits, soil_pref)
+    native = fill(true, numSpecies)
+    p = 1 ./ energy_vec.energy 
+    abun = rand(Multinomial(individuals, p ./ sum(p)))
+    # abun = [fill(div(individuals, numMoss), numMoss); fill(0, numShrub)]
+    sppl = SpeciesList(numSpecies, traits, abun, energy_vec,
+        movement, param, native)
+
+    # Create abiotic environment - even grid of one temperature
+    file = "data/CF_elevation_smooth.tif"
+    ele = readfile(file, 261000.0m, 266000.0m, 289000.0m, 293000.0m)
+    ele.data[ele.data .< 0] .= 0
+
+    file = "data/CF_TPI_smooth.tif"
+    tpi = readfile(file, 261000.0m, 266000.0m, 289000.0m, 293000.0m)
+    tpi.data[tpi.data .< 0] .= 0
+
+    @load "data/RainfallBudget_burnin.jld2"
+    file = "data/LCM.tif"
+    soil = readfile(file, 261000.0m, 266000.0m, 289000.0m, 293000.0m)
+    soil = Int.(soil)
+    abenv = peatlandAE(ele, soil, 10.0m^3, bud, active) 
+    abenv.habitat.h1.matrix .*= tpi.data
+    abenv.habitat.h1.matrix[abenv.habitat.h1.matrix .< 10.0m^3] .= 10.0m^3
+
+    # If there are ditches, make sure they are lower than everything else around and filled with water
+    if ditch
+        file = "data/Ditches_full.tif"
+        ditches = readfile(file, 289000.0m, 293000.0m, 261000.0m, 266000.0m)
+        ditch_locations = findall(.!isnan.(ditches))
+        locs = [convert_coords(d[1], d[2], size(cf, 1)) for d in ditch_locations]
+        abenv.habitat.h1.matrix[locs] .= 0.0m^3
+        abenv.habitat.h2.matrix[locs] .= 0
+        abenv.active[ditch_locations] .= false
+    else
+        ditch_locations = []
+    end
+    # heatmap(ustrip.(abenv.habitat.h1.matrix)')
+
+    # Set relationship between species and environment (gaussian)
+    rel = multiplicativeTR3(Gauss{typeof(1.0m^3)}(), NoRelContinuous{Float64}(), soilmatch{Int64}())
+
+    # Create new transition list
+    transitions = TransitionList(true)
+    addtransition!(transitions, UpdateEnergy(update_energy_usage!))
+    addtransition!(transitions, UpdateEnvironment(update_peat_environment!))
+    active_squares = findall(active[1:end])
+    peat_squares = findall(soil .== 11)
+    # Add in species based transitions
+    for spp in eachindex(sppl.species.names) 
+        for loc in active_squares
+            addtransition!(transitions, GenerateSeed(spp, loc, sppl.params.birth[spp]))
+            addtransition!(transitions, DeathProcess(spp, loc, sppl.params.death[spp]))
+            addtransition!(transitions, SeedDisperse(spp, loc))
+            addtransition!(transitions, WaterUse(spp, loc, 0.01))
+            if spp > numMoss
+                addtransition!(transitions, Invasive(spp, loc, 10.0/28days))
+            end
+        end
+    end
+    # Add in location based transitions and ditches
+    for loc in eachindex(abenv.habitat.h1.matrix)
+        # if loc ∈ ditch_locations
+        #     drainage = 1.0/month
+        # else
+            drainage = 0.001/month
+        # end
+        κ = 0.01/month
+        λ = 0.01/month
+
+        addtransition!(transitions, LateralFlow(loc, κ, λ))
+        addtransition!(transitions, WaterFlux(loc, drainage, 150.0m^3))
+    end
+
+    transitions = specialise_transition_list(transitions)
+    # Create ecosystem
+    eco = PeatSystem(sppl, abenv, rel, transitions = transitions)
+
+    # Run simulation
+    # Simulation Parameters
+    burnin = 30year; times = 70year; timestep = 1month;
+    record_interval = 3month
+    lensim = length(0years:record_interval:times)
+    # Burnin
+    abuns = generate_storage(eco, lensim, 1)[:, :, :, 1]
+    abuns1 = generate_storage(eco, lensim, 1)[:, :, :, 1]
+    abuns2 = generate_storage(eco, lensim, 1)[:, :, :, 1]
+    @time simulate!(eco, burnin, timestep, specialise = true)
+
+    @load "data/RainfallBudget_future.jld2"
+    eco.abenv.budget = bud
+    bud.matrix[:, :, 721:732] .*= 0.1
+
+    @time simulate_record!(abuns, eco, times, record_interval, timestep, save = save, save_path = save_path, specialise = true);
+
+    return abuns
+end
+
+## WITHOUT DITCHES ##
+abuns = runFuture();
+@save "/home/claireh/sdc/Peatland/Peatland_future_shock.jld2" abuns=abuns[:, :, [12,end]]
+
+plot(grid = false, label = "", layout = (3, 1), left_margin = 1.0*Plots.inch, size = (1000, 1200))
+for i in mosses
+    plot!(0:3/12:70,sum(abuns[i, :, :], dims = 1)[1, :,], grid = false, label = "", subplot = 1,
+    title = "Mosses", colour = :grey, alpha = 0.1)
+end
+annotate!([-10], [7e6], ["A"])
+vline!([60], subplot = 1, color = :black, lty = :dot, label = "")
+plot!(0:3/12:70, mean(sum(abuns[mosses, :, :], dims = 2)[:, 1, :], dims = 1)[1, :], grid = false, label = "", subplot = 1,
+    colour = :black)
+for i in shrubs
+    plot!(0:3/12:70,sum(abuns[i, :, :], dims = 1)[1, :,], grid = false, label = "", subplot = 2,
+    title = "Shrubs, Herbs and Grasses", colour = :grey, alpha = 0.1)
+end
+plot!(0:3/12:70, mean(sum(abuns[shrubs, :, :], dims = 2)[:, 1, :], dims = 1)[1, :], grid = false, label = "", subplot = 2,
+     colour = :black)
+vline!([60], subplot = 2, color = :black, lty = :dot, label = "")
+for i in trees
+    plot!(0:3/12:70,sum(abuns[i, :, :], dims = 1)[1, :,], grid = false, label = "", subplot = 3, 
+    title = "Trees", colour = :grey, alpha = 0.1)
+end
+plot!(0:3/12:70, mean(sum(abuns[trees, :, :], dims = 2)[:, 1, :], dims = 1)[1, :], grid = false, label = "", subplot = 3,
+   colour = :black)
+vline!([60], subplot = 3, color = :black, lty = :dot, label = "")
+Plots.pdf("Abuns_future_shock.pdf")
+
+
+### Past RAINFALL SCENARIO 2010 - 2080 plus summer drought in year 2060 ###
+
+function runPast(ditch::Bool = false; save = false, save_path = pwd())
+    JLD2.@load("data/Peat_30_spp.jld2", peat_spp)
+    JLD2.@load("data/Peat_30_moss.jld2", moss_spp)
+
+    file = "data/CF_outline.tif"
+    cf = readfile(file, 261000.0m, 266000.0m, 289000.0m, 293000.0m)
+    active = Array(.!isnan.(Array(cf)))
+    #heatmap(active)
+
+    grid = size(cf); individuals=1_000_000; area = step(cf.axes[1]) * step(cf.axes[2]) * prod(grid);
+    numMoss = nrow(moss_spp); numShrub = nrow(peat_spp); numSpecies = numMoss + numShrub
+
+    mosses = 1:numMoss
+    shrubs = findall((peat_spp[!, :Type] .== "Shrub") .| (peat_spp[!, :Type] .== "Grass") .| (peat_spp[!, :Type] .== "Herb")) .+ numMoss
+    trees = findall((peat_spp[!, :Type] .== "Tree")) .+ numMoss
+
+    height = peat_spp[!, :Plant_height] .* m
+    moss_height = moss_spp[!, :Len] .* mm
+
+    #Set up how much energy each species consumes
+    req1 = moss_height .* rand(Normal(1.0, 0.1), numMoss) .* mm ./m
+    req2 = height .* rand(Normal(10.0, 0.1), numShrub) .* mm ./m
+    energy_vec = WaterRequirement([req1; req2])
+
+    # Set rates for birth and death
+    birth = 0.1/year
+    death = 0.1/year
+    longevity = 0.1
+    survival = 0.01
+    boost = 1.0
+    # Collect model parameters together
+    param = EqualPop(birth, death, longevity, survival, boost)
+
+    # Create kernel for movement
+    kernel = fill(GaussianKernel(20.0m, 1e-5), numSpecies)
+    movement = BirthOnlyMovement(kernel, NoBoundary())
+
+    # Create species list, including their temperature preferences, seed abundance and native status
+    pref1 = rand(Normal(90.0, 13.0), numMoss) .* m^3 
+    pref2 = rand(Normal(70.0, 18.0), numShrub) .* m^3 
+    opts = [pref1; pref2]
+    vars = [fill(10.0m^3, numMoss); fill(10.0m^3, numShrub)]
+    water_traits = GaussTrait(opts, vars)
+    ele_traits = GaussTrait(fill(1.0, numSpecies), fill(20.0, numSpecies))
+    soilDict = Dict("hygrophilous" => [8, 11], "terrestrial" => [1, 4, 5], "terrestrial/hygrophilous" => [1, 4, 5, 8, 11])
+    soil_pref1 = fill([8, 11], numMoss); soil_pref2 = [soilDict[x] for x in peat_spp[!, :Habitat]]
+    soil_pref = soiltrait([soil_pref1; soil_pref2])
+    traits = TraitCollection3(water_traits, ele_traits, soil_pref)
+    native = fill(true, numSpecies)
+    p = 1 ./ energy_vec.energy 
+    abun = rand(Multinomial(individuals, p ./ sum(p)))
+    # abun = [fill(div(individuals, numMoss), numMoss); fill(0, numShrub)]
+    sppl = SpeciesList(numSpecies, traits, abun, energy_vec,
+        movement, param, native)
+
+    # Create abiotic environment - even grid of one temperature
+    file = "data/CF_elevation_smooth.tif"
+    ele = readfile(file, 261000.0m, 266000.0m, 289000.0m, 293000.0m)
+    ele.data[ele.data .< 0] .= 0
+
+    file = "data/CF_TPI_smooth.tif"
+    tpi = readfile(file, 261000.0m, 266000.0m, 289000.0m, 293000.0m)
+    tpi.data[tpi.data .< 0] .= 0
+
+    @load "data/RainfallBudget_burnin.jld2"
+    file = "data/LCM.tif"
+    soil = readfile(file, 261000.0m, 266000.0m, 289000.0m, 293000.0m)
+    soil = Int.(soil)
+    abenv = peatlandAE(ele, soil, 10.0m^3, bud, active) 
+    abenv.habitat.h1.matrix .*= tpi.data
+    abenv.habitat.h1.matrix[abenv.habitat.h1.matrix .< 10.0m^3] .= 10.0m^3
+
+    # If there are ditches, make sure they are lower than everything else around and filled with water
+    if ditch
+        file = "data/Ditches_full.tif"
+        ditches = readfile(file, 289000.0m, 293000.0m, 261000.0m, 266000.0m)
+        ditch_locations = findall(.!isnan.(ditches))
+        locs = [convert_coords(d[1], d[2], size(cf, 1)) for d in ditch_locations]
+        abenv.habitat.h1.matrix[locs] .= 0.0m^3
+        abenv.habitat.h2.matrix[locs] .= 0
+        abenv.active[ditch_locations] .= false
+    else
+        ditch_locations = []
+    end
+    # heatmap(ustrip.(abenv.habitat.h1.matrix)')
+
+    # Set relationship between species and environment (gaussian)
+    rel = multiplicativeTR3(Gauss{typeof(1.0m^3)}(), NoRelContinuous{Float64}(), soilmatch{Int64}())
+
+    # Create new transition list
+    transitions = TransitionList(true)
+    addtransition!(transitions, UpdateEnergy(update_energy_usage!))
+    addtransition!(transitions, UpdateEnvironment(update_peat_environment!))
+    active_squares = findall(active[1:end])
+    peat_squares = findall(soil .== 11)
+    # Add in species based transitions
+    for spp in eachindex(sppl.species.names) 
+        for loc in active_squares
+            addtransition!(transitions, GenerateSeed(spp, loc, sppl.params.birth[spp]))
+            addtransition!(transitions, DeathProcess(spp, loc, sppl.params.death[spp]))
+            addtransition!(transitions, SeedDisperse(spp, loc))
+            addtransition!(transitions, WaterUse(spp, loc, 0.01))
+            if spp > numMoss
+                addtransition!(transitions, Invasive(spp, loc, 10.0/28days))
+            end
+        end
+    end
+    # Add in location based transitions and ditches
+    for loc in eachindex(abenv.habitat.h1.matrix)
+        # if loc ∈ ditch_locations
+        #     drainage = 1.0/month
+        # else
+            drainage = 0.001/month
+        # end
+        κ = 0.01/month
+        λ = 0.01/month
+
+        addtransition!(transitions, LateralFlow(loc, κ, λ))
+        addtransition!(transitions, WaterFlux(loc, drainage, 150.0m^3))
+    end
+
+    transitions = specialise_transition_list(transitions)
+    # Create ecosystem
+    eco = PeatSystem(sppl, abenv, rel, transitions = transitions)
+
+    # Run simulation
+    # Simulation Parameters
+    burnin = 30year; times = 70year; timestep = 1month;
+    record_interval = 3month
+    lensim = length(0years:record_interval:times)
+    # Burnin
+    abuns = generate_storage(eco, lensim, 1)[:, :, :, 1]
+    abuns1 = generate_storage(eco, lensim, 1)[:, :, :, 1]
+    abuns2 = generate_storage(eco, lensim, 1)[:, :, :, 1]
+    @time simulate!(eco, burnin, timestep, specialise = true)
+
+    fullbud = Array{typeof(bud.matrix[1]), 3}(undef, size(bud.matrix, 1), size(bud.matrix, 2), size(bud.matrix, 3) * 7)
+    for i in 1:7
+        fullbud[:, :, (i-1)*120+1:(i*120)] .= bud.matrix
+    end
+    bud.matrix = fullbud
+    bud.matrix[:, :, 725:728] .*= 0.1
+
+    @time simulate_record!(abuns, eco, times, record_interval, timestep, save = save, save_path = save_path, specialise = true);
+
+    return abuns
+end
+
+abuns = runPast();
+@save "/home/claireh/sdc/Peatland/Peatland_past_shock.jld2" abuns=abuns[:, :, [12,end]]
+
+plot(grid = false, label = "", layout = (3, 1), left_margin = 1.0*Plots.inch, size = (1000, 1200))
+for i in mosses
+    plot!(0:3/12:70,sum(abuns[i, :, :], dims = 1)[1, :,], grid = false, label = "", subplot = 1,
+    title = "Mosses", colour = :grey, alpha = 0.1)
+end
+annotate!([-10], [7e6], ["A"])
+vline!([60], subplot = 1, color = :black, lty = :dot, label = "")
+plot!(0:3/12:70, mean(sum(abuns[mosses, :, :], dims = 2)[:, 1, :], dims = 1)[1, :], grid = false, label = "", subplot = 1,
+    colour = :black)
+for i in shrubs
+    plot!(0:3/12:70,sum(abuns[i, :, :], dims = 1)[1, :,], grid = false, label = "", subplot = 2,
+    title = "Shrubs, Herbs and Grasses", colour = :grey, alpha = 0.1)
+end
+plot!(0:3/12:70, mean(sum(abuns[shrubs, :, :], dims = 2)[:, 1, :], dims = 1)[1, :], grid = false, label = "", subplot = 2,
+     colour = :black)
+vline!([60], subplot = 2, color = :black, lty = :dot, label = "")
+for i in trees
+    plot!(0:3/12:70,sum(abuns[i, :, :], dims = 1)[1, :,], grid = false, label = "", subplot = 3, 
+    title = "Trees", colour = :grey, alpha = 0.1)
+end
+plot!(0:3/12:70, mean(sum(abuns[trees, :, :], dims = 2)[:, 1, :], dims = 1)[1, :], grid = false, label = "", subplot = 3,
+   colour = :black)
+vline!([60], subplot = 3, color = :black, lty = :dot, label = "")
+Plots.pdf("Abuns_past_shock.pdf")
+
diff --git a/examples/SoilMoisture.jl b/examples/SoilMoisture.jl
new file mode 100644
index 0000000..038feea
--- /dev/null
+++ b/examples/SoilMoisture.jl
@@ -0,0 +1,19 @@
+using NetCDF
+using EcoSISTEM.ClimatePref
+using EcoSISTEM.Units
+using Unitful
+using Unitful.DefaultSymbols
+using Peatland
+using Statistics
+
+times = collect(2015years:1month:2016years)
+eastings = unique(reverse(ncread("data/CHESSLandMonHydEn2015.nc", "eastings"))) .* m
+northings = unique(ncread("data/CHESSLandMonHydEn2015.nc", "northings")) .* m
+xs = collect(261000.0m:1000m:266000.0m); ys = collect(289000.0m:1000m:293000.0m)
+soilm = readSoilCHESS("data/CHESSLandMonHydEn2015.nc", "smcl", times, xs, ys)
+mean(soilm) .* 10m^2 / 100kg
+
+## Approx. 3.77 m^3
+
+gmax = 0.05mm/(g*day*m^2)
+gmax * 10m^2 * 1000g/m
\ No newline at end of file
diff --git a/examples/Test.jl b/examples/Test.jl
new file mode 100644
index 0000000..1f1de86
--- /dev/null
+++ b/examples/Test.jl
@@ -0,0 +1,78 @@
+using EcoSISTEM
+using EcoSISTEM.Units
+using Peatland
+using Unitful, Unitful.DefaultSymbols
+using Distributions
+using Diversity
+using DataPipeline
+using Plots
+
+handle = initialise()
+# Set up initial parameters for ecosystem
+numSpecies = 10; grid = (10, 10); req= 10.0kJ; individuals=100_000; area = 100.0*km^2; totalK = 1_000.0kJ/km^2
+
+# Set up how much energy each species consumes
+# energy_vec = SolarRequirement(collect(1.0:10) .* kJ)
+energy_vec = SolarRequirement([collect(0.1:0.1:0.5); collect(1:5)] .* kJ)
+
+# Set rates for birth and death
+birth = fill(0.1/year, numSpecies)
+death = fill(0.1/year, numSpecies)
+longevity = 0.1
+survival = 0.01
+boost = 1.0
+# Collect model parameters together
+param = PopGrowth{typeof(unit(birth[1]))}(birth, death, longevity, survival, boost)
+
+# Create kernel for movement
+kernel = fill(GaussianKernel(5.0km, 1e-9), numSpecies)
+movement = BirthOnlyMovement(kernel, NoBoundary())
+
+# Create species list, including their temperature preferences, seed abundance and native status
+#opts = fill(274.0K, numSpecies)
+opts = collect(270.0:279) .* K
+vars = fill(0.5K, numSpecies)
+traits = GaussTrait(opts, vars)
+native = fill(true, numSpecies)
+p = 1 ./ energy_vec.energy 
+abun = rand(Multinomial(individuals, p ./ sum(p)))
+sppl = SpeciesList(numSpecies, traits, abun, energy_vec,
+    movement, param, native)
+
+# Create abiotic environment - even grid of one temperature
+abenv = tempgradAE(270.0K, 279.0K, grid, totalK, area, 0.0K/month)
+
+# Set relationship between species and environment (gaussian)
+rel = Gauss{typeof(1.0K)}()
+
+# Create transition list
+transitions = TransitionList(true)
+addtransition!(transitions, UpdateEnergy(EcoSISTEM.update_energy_usage!))
+addtransition!(transitions, UpdateEnvironment(update_environment!))
+for spp in eachindex(sppl.species.names)
+    for loc in eachindex(abenv.habitat.matrix)
+        addtransition!(transitions, GenerateSeed(spp, loc, sppl.params.birth[spp]))
+        addtransition!(transitions, DeathProcess(spp, loc, sppl.params.death[spp]))
+        addtransition!(transitions, SeedDisperse(spp, loc))
+    end
+end
+transitions = specialise_transition_list(transitions)
+
+#Create ecosystem
+eco = Ecosystem(sppl, abenv, rel, transitions = transitions)
+
+# Simulation Parameters
+burnin = 100years; times = 100years; timestep = 1month; record_interval = 3months; repeats = 1
+lensim = length(0years:record_interval:times)
+abuns = zeros(Int64, numSpecies, prod(grid), lensim)
+# Burnin
+@time simulate!(eco, burnin, timestep, specialise = true)
+@time simulate_record!(abuns, eco, times, record_interval, timestep, specialise = true);
+
+p = plot(sum(abuns[1, :, :], dims = 1)[1, :], grid = false, label = "", left_margin = 1.0*Plots.inch)
+for i in 2:10
+    p = plot!(sum(abuns[i, :, :], dims = 1)[1, :], grid = false, label = "", left_margin = 1.0*Plots.inch)
+end
+path = link_write!(handle, "PeatlandTest")
+Plots.pdf(path)
+finalise(handle)
\ No newline at end of file