Repository: emsig/emg3d

Issue:
ENH: Add option to remove empty Tx-Rx-f-pairs
`Survey.select` can be very useful to reduce a dataset. However, currently it will, e.g., keep all receivers if we select a particular source, even if this receiver has no data from that source.

- Should be a keyword to switch off/on (by default on?)
- This only applies if the survey has observed data.

To find empty slices is straight forward:
```
# empty sources
np.all(np.isnan(survey.data.observed.data), axis=(1, 2))

# empty receivers
np.all(np.isnan(survey.data.observed.data), axis=(0, 2))

# empty frequencies
np.all(np.isnan(survey.data.observed.data), axis=(0, 1))
```

We can get then a list of it and use recursion to get the survey.

Discussion:
coveralls: ## Pull Request Test Coverage Report for [Build 1456708057](https://coveralls.io/builds/44242768)

* **16** of **16**   **(100.0%)**  changed or added relevant lines in **1** file are covered.
* No unchanged relevant lines lost coverage.
* Overall coverage increased (+**0.0005%**) to **99.856%**

---



|  Totals | [![Coverage Status](https://coveralls.io/builds/44242768/badge)](https://coveralls.io/builds/44242768) |
| :-- | --: |
| Change from base [Build 1389981510](https://coveralls.io/builds/43808971): |  0.0005% |
| Covered Lines: | 4165 |
| Relevant Lines: | 4171 |

---
##### 💛  - [Coveralls](https://coveralls.io)


Patch:
diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index dcaddd5..47d1eb7 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 f90c48e..b34a1fe 100644
--- a/emg3d/surveys.py
+++ b/emg3d/surveys.py
@@ -309,7 +309,8 @@ class Survey:
         """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 @@ class Survey:
             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 @@ class Survey:
         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):


Test patch:
diff --git a/tests/test_cli.py b/tests/test_cli.py
index c768766..93b4e7d 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -588,5 +588,5 @@ class TestRun:
 
         # 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 cd67f95..0e5d99e 100644
--- a/tests/test_surveys.py
+++ b/tests/test_surveys.py
@@ -185,6 +185,42 @@ class TestSurvey():
         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=[