From 8a6d2634923f2e14f7d310f5417671d509e15579 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Dieter=20Werthm=C3=BCller?=
 <prisae@users.noreply.github.com>
Date: Sat, 13 Nov 2021 15:40:49 +0100
Subject: [PATCH] Remove empty src-rec-freq pairs (#259)

---
 CHANGELOG.rst         |  9 +++++++++
 emg3d/surveys.py      | 40 +++++++++++++++++++++++++++++++++++++---
 tests/test_cli.py     |  2 +-
 tests/test_surveys.py | 36 ++++++++++++++++++++++++++++++++++++
 4 files changed, 83 insertions(+), 4 deletions(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index dcaddd5b..47d1eb75 100644
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -6,6 +6,15 @@ Changelog
 """"""""""
 
 
+latest
+------
+
+- ``survey``:
+
+  - ``select`` removes now empty source-receiver-frequency pairs. If you want
+    the old behaviour set ``remove_empty=False``.
+
+
 v1.3.0: File-based computations
 -------------------------------
 
diff --git a/emg3d/surveys.py b/emg3d/surveys.py
index f90c48ed..b34a1fe2 100644
--- a/emg3d/surveys.py
+++ b/emg3d/surveys.py
@@ -309,7 +309,8 @@ def data(self):
         """Data, a :class:`xarray.Dataset` instance."""
         return self._data
 
-    def select(self, sources=None, receivers=None, frequencies=None):
+    def select(self, sources=None, receivers=None, frequencies=None,
+               remove_empty=True):
         """Return a Survey with selected sources, receivers, and frequencies.
 
 
@@ -319,6 +320,11 @@ def select(self, sources=None, receivers=None, frequencies=None):
             Lists containing the wanted sources, receivers, and frequencies.
             If None, all are selected.
 
+        remove_empty : bool, default: True
+            If True, and self.data.observed has finite entries, it removes
+            empty source-receiver-frequency entries and according sources,
+            receivers, and frequencies.
+
 
         Returns
         -------
@@ -358,8 +364,36 @@ def select(self, sources=None, receivers=None, frequencies=None):
         for key in survey['data'].keys():
             survey['data'][key] = self.data[key].sel(**selection)
 
-        # Return new, reduced survey.
-        return Survey.from_dict(survey)
+            # Check if there are any finite observed data.
+            if remove_empty and key == 'observed':
+                data = survey['data'][key].data
+                remove_empty = np.isfinite(data).any()
+
+        # Create new, reduced survey.
+        red_survey = Survey.from_dict(survey)
+
+        # Remove empty source-receiver-frequency pairs.
+        if remove_empty:
+
+            def get_names(name, i0, i1, i2):
+                """Return non-NaN names."""
+                ibool = np.isnan(data).all(axis=(i1, i2))
+                ind = np.arange(data.shape[i0])[~ibool]
+                keys = survey[name].keys()
+                return [n for i, n in enumerate(keys) if i in ind]
+
+            # Get names.
+            srcnames = get_names('sources', 0, 1, 2)
+            recnames = get_names('receivers', 1, 0, 2)
+            freqnames = get_names('frequencies', 2, 0, 1)
+
+            # Use recursion to remove empty pairs.
+            red_survey = red_survey.select(
+                    sources=srcnames, receivers=recnames,
+                    frequencies=freqnames, remove_empty=False,
+            )
+
+        return red_survey
 
     @property
     def shape(self):
diff --git a/tests/test_cli.py b/tests/test_cli.py
index c768766f..93b4e7d3 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -588,5 +588,5 @@ def test_data(self, tmpdir, capsys):
 
         # Ensure dry_run returns same shaped data as the real thing.
         res = emg3d.load(os.path.join(tmpdir, 'output.npz'))
-        assert_allclose(res['data'].shape, (1, 5, 1))
+        assert_allclose(res['data'].shape, (1, 4, 1))
         assert res['n_observations'] == 4
diff --git a/tests/test_surveys.py b/tests/test_surveys.py
index cd67f955..0e5d99e8 100644
--- a/tests/test_surveys.py
+++ b/tests/test_surveys.py
@@ -185,6 +185,42 @@ def test_select(self):
         t2 = survey.select(frequencies=[], receivers='RxEP-1')
         assert t2.shape == (3, 1, 0)
 
+    def test_select_remove_empty(self):
+        survey = surveys.Survey(
+            sources=surveys.txrx_coordinates_to_dict(
+                emg3d.TxElectricDipole,
+                ([0, 1, 2], 0, 0, 0, 0)
+            ),
+            receivers=surveys.txrx_coordinates_to_dict(
+                emg3d.RxElectricPoint,
+                ([100, 101, 102], 0, 0, 0, 0)
+            ),
+            frequencies=np.array([1, 2, 3]),
+        )
+
+        selection = {'sources': ['TxED-2', 'TxED-3'],
+                     'receivers': ['RxEP-2', 'RxEP-3'],
+                     'frequencies': ['f-2', 'f-3']}
+
+        # Without 'observed', it should have no effect.
+        new = survey.select(**selection)
+        assert new.shape == (2, 2, 2)
+        new = survey.select(**selection, remove_empty=False)
+        assert new.shape == (2, 2, 2)
+
+        # Create and add 'observed' data.
+        data = np.arange(27.).reshape(3, 3, 3)
+        data = data + 1j*data
+        data[1:, 1:, 1:] = np.nan
+        data[2, 2, 2] = 26.0+26.0j
+        survey.data.observed[...] = data
+
+        # With observed, it should remove if True.
+        new = survey.select(**selection)
+        assert new.shape == (1, 1, 1)
+        new = survey.select(**selection, remove_empty=False)
+        assert new.shape == (2, 2, 2)
+
     def test_src_rec_coordinates(self):
         survey = surveys.Survey(
             sources=[