1"""Assorted helper methods""" 

2from collections.abc import Iterable 

3import numpy as np 

4 

5 

6def veclinkedfn(linkedfn, i): 

7 "Generate an indexed linking function." 

8 def newlinkedfn(c): 

9 "Linked function that pulls out a particular index" 

10 return np.array(linkedfn(c))[i] 

11 return newlinkedfn 

12 

13 

14def initsolwarning(result, category="uncategorized"): 

15 "Creates a results dictionary for a particular category of warning." 

16 if "warnings" not in result: 

17 result["warnings"] = {} 

18 if category not in result["warnings"]: 

19 result["warnings"][category] = [] 

20 

21 

22def appendsolwarning(msg, data, result, category="uncategorized"): 

23 "Append a particular category of warnings to a solution." 

24 result["warnings"][category].append((msg, data)) 

25 

26 

27@np.vectorize 

28def isnan(element): 

29 "Determine if something of arbitrary type is a numpy nan." 

30 try: 

31 return np.isnan(element) 

32 except TypeError: 

33 return False 

34 

35 

36def maybe_flatten(value): 

37 "Extract values from 0-d numpy arrays, if necessary" 

38 if hasattr(value, "shape") and not value.shape: 

39 return value.flatten()[0] # 0-d numpy arrays 

40 return value 

41 

42 

43def try_str_without(item, excluded, *, latex=False): 

44 "Try to call item.str_without(excluded); fall back to str(item)" 

45 if latex and hasattr(item, "latex"): 

46 return item.latex(excluded) 

47 if hasattr(item, "str_without"): 

48 return item.str_without(excluded) 

49 return str(item) 

50 

51 

52def mag(c): 

53 "Return magnitude of a Number or Quantity" 

54 return getattr(c, "magnitude", c) 

55 

56 

57def is_sweepvar(sub): 

58 "Determines if a given substitution indicates a sweep." 

59 return splitsweep(sub)[0] 

60 

61 

62def splitsweep(sub): 

63 "Splits a substitution into (is_sweepvar, sweepval)" 

64 try: 

65 sweep, value = sub 

66 if sweep is "sweep" and (isinstance(value, Iterable) or # pylint: disable=literal-comparison 

67 hasattr(value, "__call__")): 

68 return True, value 

69 except (TypeError, ValueError): 

70 pass 

71 return False, None