allow multiple susceptibility input formats

Author: OrtnerMichael

magpylib_material_response/demag.py

@@ -31,58 +31,50 @@
 logger.configure(**config)
 
 
-def get_susceptibilities(*sources, susceptibility):
+def get_susceptibilities(sources, susceptibility):
     """Return a list of length (len(sources)) with susceptibility values
     Priority is given at the source level, hovever if value is not found, it is searched
     up the parent tree, if available. Raises an error if no value is found when reached
     the top level of the tree."""
-
+    n = len(sources)
+    
     # susceptibilities from source attributes
     if susceptibility is None:
-        susceptibilities = []
+        susis = []
         for src in sources:
             susceptibility = getattr(src, "susceptibility", None)
             if susceptibility is None:
                 if src.parent is None:
                     raise ValueError("No susceptibility defined in any parent collection")
-                susceptibilities.extend(get_susceptibilities(src.parent))
+                susis.extend(get_susceptibilities(src.parent))
             elif not hasattr(susceptibility, "__len__"):
-                susceptibilities.append((susceptibility, susceptibility, susceptibility))
+                susis.append((susceptibility, susceptibility, susceptibility))
             elif len(susceptibility) == 3:
-                susceptibilities.append(susceptibility)
+                susis.append(susceptibility)
             else:
                 raise ValueError("susceptibility is not scalar or array fo length 3")
-        return susceptibilities
-
-    # susceptibilities as input to demag function
-    n = len(sources)
-    if np.isscalar(susceptibility):
-        susceptibility = np.ones((n,3))*susceptibility
-    elif len(susceptibility) == 3:
-        susceptibility = np.tile(susceptibility, (n,1))
-        if n==3:
-            raise ValueError(
-                "Apply_demag input susceptibility is ambiguous - either scalar list or vector single entry. "
-                "Please choose different means of input or change the number of cells in the Collection."
-            )
     else:
-        if len(susceptibility) != n:
-            raise ValueError(
-                "Apply_demag input susceptibility must be scalar, 3-vector, or same length as input Collection."
-            )
-        susceptibility = np.array(susceptibility)
-        if susceptibility.ndim == 1:
-            susceptibility = np.repeat(susceptibility,3).reshape(n,3)
-
-    susceptibility = np.reshape(susceptibility, 3 * n, order="F")
-
-
-
-
-
+        # susceptibilities as input to demag function
+        if np.isscalar(susceptibility):
+            susis = np.ones((n,3))*susceptibility
+        elif len(susceptibility) == 3:
+            susis = np.tile(susceptibility, (n,1))
+            if n==3:
+                raise ValueError(
+                    "Apply_demag input susceptibility is ambiguous - either scalar list or vector single entry. "
+                    "Please choose different means of input or change the number of cells in the Collection."
+                )
+        else:
+            if len(susceptibility) != n:
+                raise ValueError(
+                    "Apply_demag input susceptibility must be scalar, 3-vector, or same length as input Collection."
+                )
+            susis = np.array(susceptibility)
+            if susis.ndim == 1:
+                susis = np.repeat(susis,3).reshape(n,3)
 
-    
-    return np.array(susceptibilities)
+    susis = np.reshape(susis, 3 * n, order="F")
+    return np.array(susis)