diff --git a/.github/workflows/matlab-ci.yml b/.github/workflows/matlab-ci.yml
index ae026ed..3a3cd5d 100644
--- a/.github/workflows/matlab-ci.yml
+++ b/.github/workflows/matlab-ci.yml
@@ -1,6 +1,5 @@
 name: MATLAB CI checks
 on:
-  push:
   pull_request:
     branches: [ "main", "dev" ]
   workflow_dispatch:
diff --git a/.github/workflows/release-please.yml b/.github/workflows/release-please.yml
index 9194151..326f10c 100644
--- a/.github/workflows/release-please.yml
+++ b/.github/workflows/release-please.yml
@@ -1,7 +1,7 @@
 on:
   push:
     branches:
-      - dev
+      - main
   workflow_dispatch:
 
 permissions:
@@ -20,4 +20,4 @@ jobs:
           token: ${{ secrets.GITHUB_TOKEN }}
 
           release-type: simple
-          target-branch: dev
+          target-branch: main
diff --git a/buildfile.m b/buildfile.m
index 0f49fa0..979bdb8 100644
--- a/buildfile.m
+++ b/buildfile.m
@@ -1,9 +1,12 @@
 function plan = buildfile()
 import matlab.buildtool.tasks.CodeIssuesTask
+import matlab.buildtool.tasks.TestTask
 
 plan = buildplan(localfunctions);
 
 plan("check") = CodeIssuesTask;
 
-plan.DefaultTasks = "check";
+plan("test") = TestTask;
+
+plan.DefaultTasks = ["check"];
 end
diff --git a/demo.m b/demo.m
index b373e2e..272df45 100644
--- a/demo.m
+++ b/demo.m
@@ -23,12 +23,10 @@
 
 mask = rand(grid.cells.num,1) < 0.001; % process only a fraction of cells
 
-enable_waitbar = true;
-num_par_workers = Inf; % use all parallel workers from the pool
+sub_rock = downscale_all(grid,rock,mask,downscale_params,num_par_workers=Inf);
 
-sub_rock = downscale_all(grid,rock,mask,downscale_params,num_par_workers);
-
-strata_trapped = strata_trapper(grid, sub_rock, mask, params, options, enable_waitbar, num_par_workers);
+strata_trapped = strata_trapper(grid, sub_rock, mask, params, options, ...
+enable_waitbar=true, num_par_workers=Inf);
 
 %% Visualize saturation functions
 
diff --git a/downscale_all.m b/downscale_all.m
index 76ad154..8b59054 100644
--- a/downscale_all.m
+++ b/downscale_all.m
@@ -1,4 +1,11 @@
-function [sub_rock] = downscale_all(grid,rock,mask,params,num_par_workers)
+function [sub_rock] = downscale_all(grid,rock,mask,params,args)
+    arguments
+        grid
+        rock
+        mask
+        params
+        args.num_par_workers (1,1) uint32  = Inf;
+    end
 
 coarse_idx = 1:length(mask);
 sub_rock(coarse_idx) = struct('poro',[],'perm',[]);
@@ -7,7 +14,7 @@
 perm = rock.perm;
 poro = rock.poro;
 
-parfor (cell_index = coarse_idx, num_par_workers)
+parfor (cell_index = coarse_idx, args.num_par_workers)
     if ~mask(cell_index)
         continue;
     end
@@ -16,4 +23,3 @@
     poro(cell_index), perm(cell_index,:), DR(cell_index,:), params);
 end
 end
-
diff --git a/src/calc_percolation.m b/src/calc_percolation.m
index f2e9742..91432e6 100644
--- a/src/calc_percolation.m
+++ b/src/calc_percolation.m
@@ -6,8 +6,10 @@
 Nz = size(p_entry,3);
 h(1,1,1:Nz) = linspace(Lz/Nz/2,Lz - Lz/Nz/2,Nz);
 
-hydrostatic_correction = include_gravity * isfinite(rho_gas) * std_gravity() * ...
-    (rho_water - rho_gas) * (h - h_ref);
+hydrostatic_correction = 0;
+if ~ismissing(rho_gas)
+    hydrostatic_correction = include_gravity * std_gravity() * (rho_water - rho_gas) * (h - h_ref);
+end
 
 invasion = false(size(p_entry));
 invadable = (p_entry + hydrostatic_correction) < p_boundary;
diff --git a/strata_trapper.m b/strata_trapper.m
index bf7c148..79de546 100644
--- a/strata_trapper.m
+++ b/strata_trapper.m
@@ -1,12 +1,12 @@
-function strata_trapped = strata_trapper(grid, sub_rock, mask, params, options, enable_waitbar, num_par_workers)
+function strata_trapped = strata_trapper(grid, sub_rock, mask, params, options, args)
 arguments
     grid            (1,1) struct
     sub_rock        (1,:) struct
     mask            (:,1) logical
     params          (1,1) Params
     options         (1,1) Options = Options();
-    enable_waitbar  (1,1) logical = false;
-    num_par_workers (1,1) uint32  = Inf;
+    args.enable_waitbar  (1,1) logical = false;
+    args.num_par_workers (1,1) uint32  = Inf;
 end
 
 perm_upscaled = zeros(grid.cells.num, 3);
@@ -21,14 +21,14 @@
 mask = mask(1:cells_num);
 
 wb_queue = parallel.pool.DataQueue;
-if enable_waitbar
+if args.enable_waitbar
     parforWaitbar(0,sum(mask));
     afterEach(wb_queue,@parforWaitbar);
 end
 
 DR = [grid.DX,grid.DY,grid.DZ];
 
-parfor (cell_index = 1:cells_num, num_par_workers)
+parfor (cell_index = 1:cells_num, args.num_par_workers)
     if ~mask(cell_index)
         continue;
     end
@@ -45,7 +45,7 @@
     krw(cell_index,:,:) = krw_cell;
     krg(cell_index,:,:) = krg_cell;
 
-    if enable_waitbar
+    if args.enable_waitbar
         send(wb_queue,cell_index);
     end
 end
@@ -62,7 +62,7 @@
     'rel_perm_gas', krg ...
     );
 
-if enable_waitbar
+if args.enable_waitbar
     parforWaitbar(0,0,'ready');
 end
 
diff --git a/test.m b/test.m
new file mode 100644
index 0000000..7cbe0b5
--- /dev/null
+++ b/test.m
@@ -0,0 +1,2 @@
+%%
+demo;