First implementation

Author: prisae

CHANGELOG.rst

@@ -13,9 +13,9 @@ 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).
+- Modelling routines: The definition of the coordinates (and, where it applies,
+  azimuths and dips/elevations) is now more flexible (see corresponding
+  function docstrings).
 
 - Bug fixes, small improvements and maintenance
 

empymod/model.py

@@ -96,15 +96,8 @@ def bipole(src, rec, depth, res, freqtime, signal=None, aniso=None,
         - [x0, x1, y0, y1, z0, z1] (bipole of finite length)
         - [x, y, z, azimuth, dip]  (dipole, infinitesimal small)
 
-        Dimensions:
-
-        - The coordinates x, y, and z (dipole) or x0, x1, y0, y1, z0, and z1
-          (bipole) can be single values or arrays.
-        - The variables x and y (dipole) or x0, x1, y0, and y1 (bipole) must
-          have the same dimensions.
-        - The variables z, azimuth, and dip (dipole) or z0 and z1 (bipole) must
-          either be single values or having the same dimension as the other
-          coordinates.
+        For `N` sources or receivers, all variables must be of size `N` or 1
+        (in the latter case it will be expanded to `N`).
 
         Angles (coordinate system is either left-handed with positive z down or
         right-handed with positive z up; East-North-Depth):
@@ -835,15 +828,8 @@ def loop(src, rec, depth, res, freqtime, signal=None, aniso=None, epermH=None,
         - [x0, x1, y0, y1, z0, z1] (bipole of finite length)
         - [x, y, z, azimuth, dip]  (dipole, infinitesimal small)
 
-        Dimensions:
-
-        - The coordinates x, y, and z (dipole) or x0, x1, y0, y1, z0, and z1
-          (bipole) can be single values or arrays.
-        - The variables x and y (dipole) or x0, x1, y0, and y1 (bipole) must
-          have the same dimensions.
-        - The variables z, azimuth, and dip (dipole) or z0 and z1 (bipole) must
-          either be single values or having the same dimension as the other
-          coordinates.
+        For `N` sources or receivers, all variables must be of size `N` or 1
+        (in the latter case it will be expanded to `N`).
 
         Angles (coordinate system is either left-handed with positive z down or
         right-handed with positive z up; East-North-Depth):

empymod/utils.py

@@ -258,14 +258,32 @@ def check_bipole(inp, name):
 
     def chck_dipole(inp, name):
         r"""Check inp for shape and type."""
-        # Check x
+        # Check x, y, and z
         inp_x = _check_var(inp[0], float, 1, name+'-x')
-
-        # Check y and ensure it has same dimension as x
-        inp_y = _check_var(inp[1], float, 1, name+'-y', inp_x.shape)
-
-        # Check z
-        inp_z = _check_var(inp[2], float, 1, name+'-z', (1,), inp_x.shape)
+        inp_y = _check_var(inp[1], float, 1, name+'-y')
+        inp_z = _check_var(inp[2], float, 1, name+'-z')
+
+        # 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 be a single value; provided: "
+                    f"{inp_x.shape}; {inp_y.shape}."
+                )
+
+        # If x/y=1, but z not => expand
+        if inp_z.size > 1 and inp_x.size == 1:
+            inp_x = np.repeat(inp_x, inp_z.size)
+            inp_y = np.repeat(inp_y, inp_z.size)
+
+        # If x/y!=1 and z!=1, check!
+        else:
+            _check_shape(inp_z, name+'-z', (1,), inp_x.shape)
 
         # Check if all depths are the same, if so replace by one value
         if np.all(np.isclose(inp_z-inp_z[0], 0)):
@@ -287,18 +305,32 @@ def chck_dipole(inp, name):
         out = chck_dipole(inp, name)
 
         # Check azimuth and dip
-        inp_a = _check_var(inp[3], float, 1, 'azimuth', (1,), out[0].shape)
-        inp_d = _check_var(inp[4], float, 1, 'dip', (1,), out[0].shape)
+        inp_a = _check_var(inp[3], float, 1, 'azimuth')
+        inp_d = _check_var(inp[4], float, 1, 'dip')
+
+        # Expand x/y coordinate or azm/dip if necessary.
+        # THINK ABOUT IT!
+        for inp_ad, ad_name in ([inp_a, 'azimuth'], [inp_d, 'dip']):
+            if inp_ad.size != 1 and out[0].shape != inp_ad.shape:
+                if out[0].size == 1:
+                    out[0] = np.repeat(out[0], inp_ad.size)
+                    out[1] = np.repeat(out[1], inp_ad.size)
+                else:
+                    raise ValueError(
+                        f"Parameter {name}-x and {name}-y must be of same size "
+                        "or be a single value; provided: "
+                        f"{inp_x.shape}; {inp_y.shape}."
+                    )
 
         # How many different depths
         nz = out[2].size
 
         # Expand azimuth and dip to match number of depths
         if nz > 1:
             if inp_a.size == 1:
-                inp_a = np.ones(nz)*inp_a
+                inp_a = np.repeat(inp_a, nz)
             if inp_d.size == 1:
-                inp_d = np.ones(nz)*inp_d
+                inp_d = np.repeat(inp_d, nz)
 
         out = [*out, inp_a, inp_d]
 
@@ -310,11 +342,19 @@ def chck_dipole(inp, name):
         # If one pole has a single depth, but the other has various
         # depths, we have to repeat the single depth, as we will have
         # to loop over them.
-        if out0[2].size != out1[2].size:
-            if out0[2].size == 1:
-                out0[2] = np.repeat(out0[2], out1[2].size)
-            else:
-                out1[2] = np.repeat(out1[2], out0[2].size)
+        for i in range(3):
+            if out0[i].size != out1[i].size:
+                if out0[i].size == 1:
+                    out0[i] = np.repeat(out0[i], out1[i].size)
+                elif out1[i].size == 1:
+                    out1[i] = np.repeat(out1[i], out0[i].size)
+                else:
+                    raise ValueError("TODO")
+
+        # Think about it. Same as above, but for the coordinates.
+        # 1. Test that x0, x1 have same dimension
+        # 2. If x0.shape != 1 and z0.shape != 1: check
+        # 3. if z.shape > x0.shape -> expand
 
         # Check if inp is a dipole instead of a bipole
         # (This is a problem, as we would could not define the angles then.)