WIP

Author: prisae

empymod/utils.py

@@ -258,13 +258,9 @@ 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')
-
-        # Check z
         inp_z = _check_var(inp[2], float, 1, name+'-z')
 
         # Expand x or y coordinate if necessary.
@@ -275,9 +271,11 @@ def chck_dipole(inp, name):
                 inp_y = np.repeat(inp_y, inp_x.size)
             else:
                 raise ValueError(
-                    f"Parameter {name}-x and {name}-y must have same shape or "
-                    f"(1,); provided: {inp_x.shape}; {inp_y.shape}."
+                    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)
@@ -310,21 +308,19 @@ def chck_dipole(inp, name):
         inp_a = _check_var(inp[3], float, 1, 'azimuth')
         inp_d = _check_var(inp[4], float, 1, 'dip')
 
-        # Expand x/y coordinate or azm if necessary.
-        if inp_a.size != 1 and out[0].shape != inp_a.shape:
-            if out[0].size == 1:
-                out[0] = np.repeat(out[0], inp_a.size)
-                out[1] = np.repeat(out[1], inp_a.size)
-            else:
-                raise ValueError("TODO")
-
-        # Expand x/y coordinate or azm if necessary.
-        if inp_d.size != 1 and out[0].shape != inp_d.shape:
-            if out[0].size == 1:
-                out[0] = np.repeat(out[0], inp_d.size)
-                out[1] = np.repeat(out[1], inp_d.size)
-            else:
-                raise ValueError("TODO")
+        # 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