From b20c1676ff504dc1f126604d5b24dc005032d364 Mon Sep 17 00:00:00 2001
From: Maksim Elizarev <m.elizarev@imperial.ac.uk>
Date: Fri, 11 Oct 2024 14:45:37 +0100
Subject: [PATCH] WIP optimized sparse input/output

---
 strata_trapper.m | 31 ++++++++++++++++---------------
 1 file changed, 16 insertions(+), 15 deletions(-)

diff --git a/strata_trapper.m b/strata_trapper.m
index bf7c148..551c316 100644
--- a/strata_trapper.m
+++ b/strata_trapper.m
@@ -9,16 +9,18 @@
     num_par_workers (1,1) uint32  = Inf;
 end
 
-perm_upscaled = zeros(grid.cells.num, 3);
+cells_num = min(length(mask),grid.cells.num);
+cell_idxs = 1:cells_num;
+mask = mask(cell_idxs);
+subset_len = sum(mask);
+
+perm_upscaled = zeros(subset_len, 3);
 
 saturations = linspace(params.sw_resid,1,options.sat_num_points);
 
-cap_pres_upscaled = nan(grid.cells.num,length(saturations));
-krw = nan(grid.cells.num,3,length(saturations));
-krg = nan(grid.cells.num,3,length(saturations));
-
-cells_num = min(length(mask),grid.cells.num);
-mask = mask(1:cells_num);
+cap_pres_upscaled = nan(subset_len,length(saturations));
+krw = nan(subset_len,3,length(saturations));
+krg = nan(subset_len,3,length(saturations));
 
 wb_queue = parallel.pool.DataQueue;
 if enable_waitbar
@@ -26,13 +28,11 @@
     afterEach(wb_queue,@parforWaitbar);
 end
 
-DR = [grid.DX,grid.DY,grid.DZ];
+DR = [grid.DX(mask),grid.DY(mask),grid.DZ(mask)];
 
-parfor (cell_index = 1:cells_num, num_par_workers)
-    if ~mask(cell_index)
-        continue;
-    end
+sub_rock = sub_rock(mask);
 
+parfor (cell_index = 1:subset_len, num_par_workers)
     sub_porosity = sub_rock(cell_index).poro;
     sub_permeability = sub_rock(cell_index).perm;
 
@@ -50,16 +50,17 @@
     end
 end
 
-krw(mask,:,saturations<=params.sw_resid) = 0;
+krw(:,:,saturations<=params.sw_resid) = 0;
 krg(krg<0) = 0;
-krg(mask,:,saturations>=1)=0;
+krg(:,:,saturations>=1)=0;
 
 strata_trapped = struct(...
     'permeability', perm_upscaled, ...
     'saturation', saturations,...
     'capillary_pressure', cap_pres_upscaled, ...
     'rel_perm_wat', krw, ...
-    'rel_perm_gas', krg ...
+    'rel_perm_gas', krg, ...
+    'idx', cell_idxs(mask) ...
     );
 
 if enable_waitbar