From d1e3cb9cca72edac588f4681cf1a1892f6d57f45 Mon Sep 17 00:00:00 2001
From: Ran Lu <ranl@princeton.edu>
Date: Tue, 10 Sep 2024 23:08:34 -0400
Subject: [PATCH] Support long range affinites

The long range affinities are collected based on a hard-coded grid for
now
---
 scripts/cut_chunk_agg.py    |  4 ++--
 src/seg/Utils.hpp           | 16 +++++++++++++---
 src/seg/atomic_chunk_ME.cpp |  5 +++--
 3 files changed, 18 insertions(+), 7 deletions(-)

diff --git a/scripts/cut_chunk_agg.py b/scripts/cut_chunk_agg.py
index 26bdca2..f1ba093 100644
--- a/scripts/cut_chunk_agg.py
+++ b/scripts/cut_chunk_agg.py
@@ -30,8 +30,8 @@ def write_metadata(fn, offset, size, ac_offset):
 ac_offset = param["ac_offset"]
 boundary_flags = param["boundary_flags"]
 
-extra_aff_backward = [1, 1, 1]
-extra_aff_forward = [0, 0, 0]
+extra_aff_backward = [5, 5, 2]
+extra_aff_forward = [5, 5, 2]
 
 extra_aff_before = [0, 0, 0]
 extra_aff_after = [x + y for x, y in zip(extra_aff_backward, extra_aff_forward)]
diff --git a/src/seg/Utils.hpp b/src/seg/Utils.hpp
index 5b9f3f7..a372d19 100644
--- a/src/seg/Utils.hpp
+++ b/src/seg/Utils.hpp
@@ -115,12 +115,22 @@ void traverseSegments(const Ts& seg, Ta& ... extractors)
     auto shape = seg.shape();
     auto c = Coord({0,0,0});
 
-    constexpr auto scan_grid = std::array<std::array<int, 3>, 3> {{
+    constexpr auto long_range = std::array<int, 3> {5, 5, 2};
+
+    constexpr auto scan_grid = std::array<std::array<int, 3>, 12> {{
                             {-1,0,0}, {0,-1,0}, {0,0,-1},
+                            {-long_range[0],0,0}, {0,-long_range[1],0},
+                            {-long_range[0],-long_range[1],0},
+                            {-long_range[0],long_range[1],0},
+                            {-long_range[0],0,-1},
+                            {0,-long_range[1],-1},
+                            {-long_range[0],0,1},
+                            {0,-long_range[1],1},
+                            {0,0,-long_range[2]},
     }};
 
-    constexpr std::array<int, 3> aff_backward = {1, 1, 1}
-    constexpr std::array<int, 3> aff_forward = {0, 0, 0};
+    constexpr std::array<int, 3> aff_backward = long_range;
+    constexpr std::array<int, 3> aff_forward = long_range;
 
     int z = skipType == 2 ? base[2]+1: base[2];
     for (; z != base[2]+shape[2]; z++) {
diff --git a/src/seg/atomic_chunk_ME.cpp b/src/seg/atomic_chunk_ME.cpp
index 5e2e2da..ff7ae22 100644
--- a/src/seg/atomic_chunk_ME.cpp
+++ b/src/seg/atomic_chunk_ME.cpp
@@ -54,10 +54,11 @@ int main(int argc, char * argv[])
     std::cout << "mmap seg data" << std::endl;
 
     bio::mapped_file_source aff_file;
-    size_t aff_bytes = sizeof(aff_t)*dim[0]*dim[1]*dim[2]*3;
+    size_t aff_dim = 12;
+    size_t aff_bytes = sizeof(aff_t)*dim[0]*dim[1]*dim[2]*aff_dim;
     aff_file.open("aff.raw", aff_bytes);
     assert(aff_file.is_open());
-    ConstChunkRef<aff_t, 4> aff_chunk (reinterpret_cast<const aff_t*>(aff_file.data()), boost::extents[Range(offset[0], offset[0]+dim[0])][Range(offset[1], offset[1]+dim[1])][Range(offset[2], offset[2]+dim[2])][3], boost::fortran_storage_order());
+    ConstChunkRef<aff_t, 4> aff_chunk (reinterpret_cast<const aff_t*>(aff_file.data()), boost::extents[Range(offset[0], offset[0]+dim[0])][Range(offset[1], offset[1]+dim[1])][Range(offset[2], offset[2]+dim[2])][aff_dim], boost::fortran_storage_order());
     std::cout << "mmap aff data" << std::endl;
 
     auto map = loadChunkMap<seg_t>("chunkmap.data");