Merge pull request #315 from sciris/rc1.4.0-psl-280-numdigits

Rc1.4.0 psl 280 numdigits

sciris/sc_math.py

@@ -17,8 +17,8 @@
 #%% Find and approximation functions
 ##############################################################################
 
-__all__ = ['approx', 'safedivide', 'findinds', 'findfirst', 'findlast', 'findnearest',
-           'count', 'dataindex', 'getvalidinds', 'sanitize', 'getvaliddata', 'isprime']
+__all__ = ['approx', 'safedivide', 'findinds', 'findfirst', 'findlast', 'findnearest', 'count',
+           'dataindex', 'getvalidinds', 'sanitize', 'getvaliddata', 'isprime', 'numdigits']
 
 
 def approx(val1=None, val2=None, eps=None, **kwargs):
@@ -365,6 +365,65 @@ def isprime(n, verbose=False):
     return True
 
 
+def numdigits(n, *args, count_minus=False, count_decimal=False):
+    """
+    Count the number of digits in a number (or list of numbers).
+
+    Useful for e.g. knowing how long a string needs to be to fit a given number.
+
+    If a number is less than 1, return the number of digits until the decimal
+    place.
+
+    Reference: https://stackoverflow.com/questions/22656345/how-to-count-the-number-of-digits-in-python
+
+    Args:
+        n (int/float/list/array): number or list of numbers
+        args (list): additional numbers
+        count_minus (bool): whether to count the minus sign as a digit
+        count_decimal (bool): whether to count the decimal point as a digit
+
+    **Examples**::
+
+        sc.numdigits(12345) # Returns 5
+        sc.numdigits(12345.5) # Returns 5
+        sc.numdigits(0) # Returns 1
+        sc.numdigits(-12345) # Returns 5
+        sc.numdigits(-12345, count_minus=True) # Returns 6
+        sc.numdigits(12, 123, 12345) # Returns [2, 3, 5]
+        sc.numdigits(0.01) # Returns -2
+        sc.numdigits(0.01, count_decimal=True) # Returns -4
+
+    New in version 2.0.0.
+    """
+    is_scalar = True if scu.isnumber(n) and len(args) == 0 else False
+
+    vals = cat(n, *args)
+
+    output = []
+    for n in vals:
+        abs_n = abs(n)
+        is_decimal = 0 < abs_n < 1
+        n_digits = 1
+        if n < 0 and count_minus:
+            n_digits += 1
+        if is_decimal:
+            if count_decimal:
+                n_digits += 1
+            else:
+                n_digits -= 1
+
+        if abs_n > 0:
+            if is_decimal:
+                n_digits = -n_digits
+            n_digits += int(np.floor(np.log10(abs_n)))
+        output.append(n_digits)
+    output = np.array(output)
+    if is_scalar:
+        output = output[0]
+
+    return output
+
+
 
 ##############################################################################
 #%% Other functions
@@ -1049,4 +1108,4 @@ def calc(xi, yi):
     # Convert back to original size and dtype
     zi = np.array(zi.reshape(orig_shape), dtype=orig_dtype)
 
-    return zi
+    return zi

sciris/sc_plotting.py

@@ -193,7 +193,7 @@ def surf3d(data, x=None, y=None, fig=None, returnfig=False, colorbar=True, figkw
 
     # Create figure
     fig,ax = ax3d(returnfig=True, fig=fig, figkwargs=figkwargs, **axkwargs)
-    ny,nx = pl.array(data).shape
+    ny,nx = np.array(data).shape
 
     if x is None:
         x = np.arange(nx)

tests/test_math.py

@@ -67,6 +67,10 @@ def test_find():
     found.vals = sc.findinds([2,3,6,3], 6)
     assert found.vals[0] == 2
 
+    print('Testing sc.count()')
+    found.count = sc.count([1,2,2,3], 2.0)
+    assert found.count == 2
+
     print('Testing sc.findfirst(), sc.findlast()')
     found.first = sc.findfirst(pl.rand(10))
     found.last = sc.findlast(pl.rand(10))
@@ -91,6 +95,14 @@ def test_find():
     sanitized      = sc.sanitize(np.array([3, 4, np.nan, 8, 2, np.nan, np.nan, np.nan, 8]), replacenans=0)
     found.sanitized = sanitized
 
+    print('Testing sc.numdigits()')
+    found.numdigits = sc.numdigits(1234)
+    found.numdigits_max = max(sc.numdigits([10, 200, 30000]))
+    found.numdigits_dec = sc.numdigits(0.01)
+    assert found.numdigits == 4
+    assert found.numdigits_max == 5
+    assert found.numdigits_dec == -2
+
     return found
 
 
@@ -153,18 +165,18 @@ def test_gauss2d(doplot=doplot):
     sc.heading('Testing Gaussian 2D smoothing')
 
     # Parameters
-    x = pl.rand(40)
-    y = pl.rand(40)
+    x = np.random.rand(40)
+    y = np.random.rand(40)
     z = 1-(x-0.5)**2 + (y-0.5)**2
 
     # Method 1 -- form grid
-    xi = pl.linspace(0,1,20)
-    yi = pl.linspace(0,1,20)
-    zi = sc.gauss2d(x, y, z, xi, yi, scale=0.1)
+    xi = np.linspace(0,1,20)
+    yi = np.linspace(0,1,20)
+    zi = sc.gauss2d(x, y, z, xi, yi, scale=0.1, grid=True)
 
     # Method 2 -- use points directly
-    xi2 = pl.rand(400)
-    yi2 = pl.rand(400)
+    xi2 = np.random.rand(400)
+    yi2 = np.random.rand(400)
     zi2 = sc.gauss2d(x, y, z, xi2, yi2, scale=0.1, grid=False)
 
     if doplot: