From 16bc157785ed6e118fe60ec3f8269ac6d2021bb1 Mon Sep 17 00:00:00 2001
From: DirkEilander <15379728+DirkEilander@users.noreply.github.com>
Date: Thu, 14 Mar 2024 13:45:34 +0100
Subject: [PATCH] simplify ind method (#189)

* simplify ind method

* add unittest for index

* bump version to 1.0.4.dev
---
 envs/hydromt-sfincs-dev.yml   |  2 +-
 hydromt_sfincs/__init__.py    |  2 +-
 hydromt_sfincs/regulargrid.py |  4 +---
 tests/conftest.py             | 24 ++++++++++++++++++++++++
 tests/test_regulargrid.py     |  8 ++++++++
 5 files changed, 35 insertions(+), 5 deletions(-)
 create mode 100644 tests/test_regulargrid.py

diff --git a/envs/hydromt-sfincs-dev.yml b/envs/hydromt-sfincs-dev.yml
index fc949877..4458ad0c 100644
--- a/envs/hydromt-sfincs-dev.yml
+++ b/envs/hydromt-sfincs-dev.yml
@@ -9,7 +9,7 @@ dependencies:
 - descartes
 - ffmpeg
 - geopandas>=0.8
-- hydromt>=0.9.1
+- hydromt>=0.9.1, <0.10
 - jupyterlab
 - matplotlib
 - nbsphinx
diff --git a/hydromt_sfincs/__init__.py b/hydromt_sfincs/__init__.py
index e4d47b46..9c051e8d 100644
--- a/hydromt_sfincs/__init__.py
+++ b/hydromt_sfincs/__init__.py
@@ -3,7 +3,7 @@
 from os.path import dirname, join, abspath
 
 
-__version__ = "1.0.3"
+__version__ = "1.0.4.dev"
 
 DATADIR = join(dirname(abspath(__file__)), "data")
 
diff --git a/hydromt_sfincs/regulargrid.py b/hydromt_sfincs/regulargrid.py
index b8db21df..e4fe71e0 100644
--- a/hydromt_sfincs/regulargrid.py
+++ b/hydromt_sfincs/regulargrid.py
@@ -103,9 +103,7 @@ def empty_mask(self) -> xr.DataArray:
     def ind(self, mask: np.ndarray) -> np.ndarray:
         """Return indices of active cells in mask."""
         assert mask.shape == (self.nmax, self.mmax)
-        iok = np.where(np.transpose(mask) > 0)
-        iok = (iok[1], iok[0])
-        ind = np.ravel_multi_index(iok, (self.nmax, self.mmax), order="F")
+        ind = np.where(mask.ravel(order="F"))[0]
         return ind
 
     def write_ind(
diff --git a/tests/conftest.py b/tests/conftest.py
index 247d693b..9f95fe29 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -3,13 +3,37 @@
 from os.path import abspath, dirname, join
 
 import pytest
+import numpy as np
 
 from hydromt_sfincs.sfincs import SfincsModel
+from hydromt_sfincs.regulargrid import RegularGrid
 
 TESTDATADIR = join(dirname(abspath(__file__)), "data")
 TESTMODELDIR = join(TESTDATADIR, "sfincs_test")
 
 
+@pytest.fixture
+def reggrid():
+    # create a simple regular grid
+    grid = RegularGrid(
+        x0=318650,
+        y0=5040000,
+        dx=150,
+        dy=150,
+        nmax=84,  # height
+        mmax=36,  # width
+    )
+    return grid
+
+
+@pytest.fixture
+def mask(reggrid):
+    # create a simple mask
+    mask = np.zeros((reggrid.nmax, reggrid.mmax), dtype="u1")
+    mask[2:, 3:-1] = 1
+    return mask
+
+
 @pytest.fixture
 def weirs():
     feats = [
diff --git a/tests/test_regulargrid.py b/tests/test_regulargrid.py
new file mode 100644
index 00000000..c64f88b2
--- /dev/null
+++ b/tests/test_regulargrid.py
@@ -0,0 +1,8 @@
+import numpy as np
+
+
+def test_ind(reggrid, mask):
+    ind = reggrid.ind(mask)
+    assert ind[0] == 254
+    assert ind[-1] == 2939
+    assert ind.size == np.sum(mask > 0)