From 7996cd6fdbc98794bed166a674136039be02b4b2 Mon Sep 17 00:00:00 2001
From: Karamaz0V1 <florent.guiotte@gmail.com>
Date: Sun, 30 Aug 2020 15:33:04 +0200
Subject: [PATCH] Add DTM filter and 2D grids

---
 idefix/helpers.py                 | 48 +++++++++++++++++++++++++++++++
 idefix/vxl.py                     |  9 ++++--
 test/test_helpers.py              |  5 ++++
 test/test_vxl.py                  |  1 +
 test/test_vxl/pc0_grid_s1-1-n.txt |  3 ++
 5 files changed, 63 insertions(+), 3 deletions(-)
 create mode 100644 test/test_vxl/pc0_grid_s1-1-n.txt

diff --git a/idefix/helpers.py b/idefix/helpers.py
index ab22eaf..d2cfd35 100644
--- a/idefix/helpers.py
+++ b/idefix/helpers.py
@@ -15,6 +15,7 @@ import numpy as np
 from scipy.interpolate import griddata
 from rasterio import fill
 import sap
+import higra as hg
 
 from .vxl import get_grid, bin, squash
 
@@ -87,3 +88,50 @@ def dsm(pcloud, cell_size=1., last=False):
     rstr = interpolate(rstr, 'idw')
 
     return rstr
+
+def dtm_dh_filter(dsm, sigma=.5, epsilon=20000, alpha=2):
+    """Compute a digital terrain model (DTM) from a DSM.
+
+    Work best with DSM of last echo.
+
+    Parameters
+    ----------
+    dsm : ndarray
+        The DSM.
+    sigma : scalar
+        The height theshold to trigger object detection. Default is 
+        0.5 m.
+    epsilon : scalar
+        The area theshold for ground objects. All objects with surface
+        greater than epsilon are forcedto be ground. Default is 20 km².
+    alpha : scalar
+        The area threshold for horizontal noise filter. Area variations
+        smaller than alpha are removed for the detection of height
+        threshold sigma. Default is 2 m².
+
+    Returns
+    -------
+    dtm : ndarray
+        The DTM computed from the DSM.
+
+    """
+    mt = sap.MaxTree(dsm)
+    area = mt.get_attribute('area')
+    
+    area_child = hg.accumulate_parallel(mt._tree, area, hg.Accumulators.max)
+    pruned = (area - area_child) <= alpha
+
+    pruned_tree, pruned_map = hg.simplify_tree(mt._tree, pruned)
+    
+    dh = mt._alt[pruned_map] - mt._alt[pruned_map][pruned_tree.parents()]
+    remove = dh > sigma
+
+    original_map = np.zeros(mt.num_nodes(), dtype=np.int)
+    original_map[pruned_map] = np.arange(pruned_map.size)
+    original_map = hg.accumulate_and_max_sequential(mt._tree, original_map, np.arange(mt._tree.num_leaves()), hg.Accumulators.max)
+    original_remove = remove[original_map] & (area < epsilon)
+    
+    dtm = mt.reconstruct(original_remove, filtering='min')
+    
+    return dtm
+
diff --git a/idefix/vxl.py b/idefix/vxl.py
index fcee31a..ba6342b 100644
--- a/idefix/vxl.py
+++ b/idefix/vxl.py
@@ -29,7 +29,7 @@ def _ui_step(step, spatial):
         out_step = [step] * spatial.shape[-1]
 
     for s in out_step:
-        if s <= 0:
+        if s and s <= 0:
             msg = 'Step should be greater than 0, steps = \'{}\'.'.format(step)
             log.error(msg)
             raise ValueError(msg)
@@ -62,8 +62,11 @@ def get_grid(spatial, step):
     grid = []
     for a_min, a_max, a_s in zip(bb[0], bb[1], step):
         # Beware of float underflow
-        bins = np.trunc((a_max - a_min) / a_s).astype(int) + 1
-        grid += [np.linspace(a_min, a_min + bins * a_s, bins + 1)]
+        if a_s:
+            bins = np.trunc((a_max - a_min) / a_s).astype(int) + 1
+            grid += [np.linspace(a_min, a_min + bins * a_s, bins + 1)]
+        else:
+            grid += [np.array((a_min, a_max + 1))]
 
     return grid
 
diff --git a/test/test_helpers.py b/test/test_helpers.py
index 378f7bf..03e53ca 100644
--- a/test/test_helpers.py
+++ b/test/test_helpers.py
@@ -39,3 +39,8 @@ def test_dsm(datadir, params):
 
     assert dsm is not None, 'Did not return anything...'
     assert not np.isnan(dsm).any(), 'Some missing values in DSM'
+
+def test_dtm(ma_raster):
+    dtm = helpers.dtm_dh_filter(ma_raster)
+
+    assert dtm is not None, 'Did not return anything...'
diff --git a/test/test_vxl.py b/test/test_vxl.py
index 914ec5a..99c01cf 100644
--- a/test/test_vxl.py
+++ b/test/test_vxl.py
@@ -75,6 +75,7 @@ def data_grid(datadir, set_id, step_id):
     ('0', .7, '0_7'),
     ('0', .15, '0_15'),
     ('0', [1.,1.,2.] , '1-1-2'),
+    ('0', [1.,1.,None] , '1-1-n'),
 ])
 def test_get_grid(datadir, set_id, step, grid_id):
     spatial = data_pc(datadir, set_id).spatial
diff --git a/test/test_vxl/pc0_grid_s1-1-n.txt b/test/test_vxl/pc0_grid_s1-1-n.txt
new file mode 100644
index 0000000..1382f54
--- /dev/null
+++ b/test/test_vxl/pc0_grid_s1-1-n.txt
@@ -0,0 +1,3 @@
+1.0 2.0 3.0 4.0 5.0 6.0 7.0 8.0 9.0 10.0 11.0
+1.0 2.0 3.0 4.0 5.0 6.0 7.0 8.0 9.0 10.0 11.0
+1.0 11.0