Flexible coordinates for dipoles

CHANGELOG.rst

@@ -13,6 +13,10 @@ v2.2.x
 latest
 ------
 
+- Modelling routines: The definition of the coordinates for ``dipole``,
+  ``dipole_k``, and ``analytical`` is now more flexible (``x`` and ``y``
+  coordinate can now have different dimension, as long as one is a scalar).
+
 - Bug fixes, small improvements and maintenance
 
   - Cleaned-up the namespace by setting ``dir()`` explicitly.

empymod/model.py

@@ -593,8 +593,9 @@ def dipole(src, rec, depth, res, freqtime, signal=None, ab=11, aniso=None,
     src, rec : list of floats or arrays
         Source and receiver coordinates [x, y, z] (m):
 
-        - The x- and y-coordinates can be arrays, z is a single value.
-        - The x- and y-coordinates must have the same dimension.
+        For `N` sources or receivers, the x- and y-coordinates must be of size
+        `N` or 1 (in the latter case it will be expanded to `N`); z is always a
+        single value.
 
         Sources or receivers placed on a layer interface are considered in the
         upper layer.
@@ -1235,8 +1236,9 @@ def analytical(src, rec, res, freqtime, solution='fs', signal=None, ab=11,
     src, rec : list of floats or arrays
         Source and receiver coordinates [x, y, z] (m):
 
-        - The x- and y-coordinates can be arrays, z is a single value.
-        - The x- and y-coordinates must have the same dimension.
+        For `N` sources or receivers, the x- and y-coordinates must be of size
+        `N` or 1 (in the latter case it will be expanded to `N`); z is always a
+        single value.
 
     res : float
         Horizontal resistivity rho_h (Ohm.m).
@@ -1562,8 +1564,9 @@ def dipole_k(src, rec, depth, res, freq, wavenumber, ab=11, aniso=None,
     src, rec : list of floats or arrays
         Source and receiver coordinates [x, y, z] (m):
 
-        - The x- and y-coordinates can be arrays, z is a single value.
-        - The x- and y-coordinates must have the same dimension.
+        - For `N` sources or receivers, the x- and y-coordinates must be of
+          size `N` or 1 (in the latter case it will be expanded to `N`); z is
+          always a single value.
         - The x- and y-coordinates only matter for the angle-dependent factor.
 
         Sources or receivers placed on a layer interface are considered in the

empymod/utils.py

@@ -361,10 +361,23 @@ def check_dipole(inp, name, verb):
 
     """
 
-    # Check inp for x, y, and z; x & y must have same length, z is a float
+    # Check inp for x, y, and z; x & y must be of size N or 1, z is a float
     _check_shape(np.squeeze(np.asarray(inp, dtype=object)), name, (3,))
     inp_x = _check_var(inp[0], float, 1, name+'-x')
-    inp_y = _check_var(inp[1], float, 1, name+'-y', inp_x.shape)
+    inp_y = _check_var(inp[1], float, 1, name+'-y')
+
+    # Expand x or y coordinate if necessary.
+    if inp_x.size != inp_y.size:
+        if inp_x.size == 1:
+            inp_x = np.repeat(inp_x, inp_y.size)
+        elif inp_y.size == 1:
+            inp_y = np.repeat(inp_y, inp_x.size)
+        else:
+            raise ValueError(
+                f"Parameters {name}-x and {name}-y must be of same size or "
+                f"be a single value; provided: {inp_x.shape}; {inp_y.shape}."
+            )
+
     inp_z = _check_var(inp[2], float, 1, name+'-z', (1,))
     out = [inp_x, inp_y, inp_z]
 

tests/test_fdesign.py

@@ -593,3 +593,7 @@ def test_print_count(capsys):
                                 'time': default_timer(), 'warn-r': 0})
     out, _ = capsys.readouterr()
     assert out == ""
+
+
+def test_all_dir():
+    assert set(fdesign.__all__) == set(dir(fdesign))

tests/test_filters.py

@@ -77,3 +77,7 @@ def test_fourier_dlf():                                    # 3. Fourier filters
         assert len(dlf.cos) == nr
         assert len(dlf.sin) == nr
         assert_allclose(dlf.factor, fact)
+
+
+def test_all_dir():
+    assert set(filters.__all__) == set(dir(filters))

tests/test_io.py

@@ -125,3 +125,7 @@ def test_ComplexNumPyEncoder():
     # Error
     with pytest.raises(TypeError, match="Object of type module"):
         test.default(io)
+
+
+def test_all_dir():
+    assert set(io.__all__) == set(dir(io))

tests/test_kernel.py

@@ -170,3 +170,7 @@ def test_halfspace():                                            # 7. halfspace
     hsbp['11']['aniso'] = hsbp['11']['aniso'][1]
     hs_res = bipole(**hsbp['11'])
     assert_allclose(direct, hs_res, atol=1e-2)
+
+
+def test_all_dir():
+    assert set(kernel.__all__) == set(dir(kernel))

tests/test_model.py

@@ -6,6 +6,7 @@
 
 # Import main modelling routines from empymod directly to ensure they are in
 # the __init__.py-file.
+from empymod import model
 from empymod import bipole, dipole, analytical, loop
 # Import rest from model
 from empymod.model import gpr, dipole_k, fem, tem
@@ -1140,3 +1141,7 @@ def test_analytical(self):
 
         assert a.shape == (2, )
         assert b.shape == (1, 2, 1)
+
+
+def test_all_dir():
+    assert set(model.__all__) == set(dir(model))

tests/test_tmtemod.py

@@ -110,3 +110,7 @@ def test_fields():
             # Check
             assert_allclose(out[0], TMTE[0] + TMTE[1])
             assert_allclose(out[1], TMTE[2] + TMTE[3])
+
+
+def test_all_dir():
+    assert set(tmtemod.__all__) == set(dir(tmtemod))

tests/test_transform.py

@@ -489,3 +489,7 @@ def test_dlf():                                                       # 10. dlf
 
         # Compare
         assert_allclose(np.squeeze(fEM3), np.squeeze(freq2), rtol=1e-3)
+
+
+def test_all_dir():
+    assert set(transform.__all__) == set(dir(transform))

tests/test_utils.py

@@ -200,12 +200,21 @@ def test_check_dipole(capsys):
     outstr += "     > z       [m] :  0\n"
     assert out == outstr
 
+    # Check expansion
+    rec, nrec = utils.check_dipole([[0, 1], 0, 0], 'rec', 0)
+    assert_allclose(rec[1], [0, 0])
+    rec, nrec = utils.check_dipole([3, [0, 1, 2], 0], 'rec', 0)
+    assert_allclose(rec[0], [3, 3, 3])
+
     # Check Errors: more than one z
     with pytest.raises(ValueError, match='Parameter src has wrong shape'):
         utils.check_dipole([[0, 0], [0, 0], [0, 0]], 'src', 3)
     # Check Errors: wrong number of elements
     with pytest.raises(ValueError, match='Parameter rec has wrong shape'):
         utils.check_dipole([0, 0, 0, 0], 'rec', 3)
+    # Check Errors: x- and y- of different sizes, != 1
+    with pytest.raises(ValueError, match='Parameters rec-x and rec-y must be'):
+        utils.check_dipole([[0, 1], [0, 1, 2], 0], 'rec', 3)
 
 
 def test_check_frequency(capsys):
@@ -1213,3 +1222,7 @@ def test_import_time():
     # Currently we check t < 1.2s.
     # => That should come down to t < 0.5s in the future!
     assert float(out.stderr.decode("utf-8")[:-1]) < 1.2
+
+
+def test_all_dir():
+    assert set(utils.__all__) == set(dir(utils))