Numerical Utilities#

The immlib library contains a number of utilities that are intended to make it easier to write scientific libraries that nicely handle numerical data, specifically NumPy arrays and PyTorch tensors. Although PyTorch and NumPy occasionally work together and though functions can often be written to handle either kind of input, a frequent scenario when writing scientific tools is the choice between writing (and supporting) a NumPy version of the tool, a PyTorch version of the tool, or both. The tools described here aim to make supporting a library that handles both much easier.

Note

PyTorch is not an explicit dependency of immlib because the tools in immlib are useful even if one does not use PyTorch. Most of the tools related to PyTorch described below will still function correctly if you do not have PyTorch installed (because, for example, they are written to fail when non-PyTorch inputs are provided). If you write an API with immlib’s tools, it should work for NumPy, and it should work for PyTorch when it is installed.

NumPy-specific Tools#

A few immlib tools are specific to NumPy arrays. The most critical of these are the is_array and to_array functions. These functions provide the simple utilities of testing whether something is a NumPy array or not and of converting something into a NumPy array, but with many additional options and features.

The is_array Function#

is_array(x) fundamentally tests whether the object x is either a NumPy array or a SciPy sprase array, but it allows pint quantities whose magnitudes are arrays to be considered arrays also (this behavior can be changed with the quant optional parameter). The optional parameters to is_array allow one to test for many additional constraints on the array object simultaneously. When the following descriptions refer to “the array” they are referring to x when x is an array and x’s magnitude when x is a pint quantity whose magnitude is an array.

  • Testing whether the object is a quantity.

    • is_array(x, quant=True) requires that x be a pint quantity with an array magnitude.

    • is_array(x, quant=False) requires that x be not be a pint quantity (instead it must be a plain NumPy array).

    • is_array(x, quant=None) (the default value) allows that x be either a NumPy array or a pint quantity with a NumPy array as its magnitude.

  • Testing the object’s physical units.

    • is_array(x, quant=True, unit='mm') requires that x be a pint quantity with a physical unit that is compatible with millimeters. Any length-based unit is compatible with millimeters, so a quantity in feet, for example, would be allowed.

    • is_array(x, unit=None) requires that x not have a unit, meaning that x must not be a pint quantity.

    • is_array(x, unit=...) (the default value) does not apply any requirements to the object’s units.

  • Testing the dtype.

    • is_array(x, dtype=dt) requires that the dtype of the array represented by x have a dtype that is a sub-dtype of dt. For example is_array(x, dtype=np.number) will require that x have a numerical dtype, but would allow integers, floating point, or complex dtypes.

    • is_array(x, dtype=(dt1, dt2)) then the dtype of the array must be equal to one of the listed dtypes.

    • is_array(x, dtype=None) (the default value) does not apply any restrictions to the array’s dtype.

  • Testing the shape.

    • is_array(x, ndim=d) requires that the number of dimensions of the array be equal to d.

    • is_array(x, ndim=(d1, d2)) requires that the number of dimensions of the array be equal to one of the listed dimensionalities.

    • is_array(x, shape=sh) requires that the shape of x match sh. When specifying the shape that the array must match, the object sh can take a few forms. First, it can be either a tuple of numbers of an integer. When an integer n is given, this is considered equivalent to (n,). Additionally, the shape pattern can contain entries equal to -1, which indicates that the relevant entry of the shape tuple can have any size. Finally, an Ellipsis object (...) can appear in at most one entry of the tuple. This entry matches any number of shape dimensions. For example, is_array(x, shape=(-1, 2, 3, ..., 4)) would return True for NumPy arrays with any of the following shapes: (100, 2, 3, 4), (10, 2, 3, 1, 4), (1, 2, 3, 1, 0, 4). It would not, however, match any of these shapes: (2, 3, 4), (10, 2, 3, 4, 5), (10, 2, 3).

    • is_array(x, numel=n) requires that the total number of elements in the array (after flattening) be equal to n.

    • Note that ndim, shape, and numel can be used simultaneously.

  • Testing immutability.

    • is_array(x, frozen=False) requires that x not be a read-only array.

    • is_array(x, frozen=True) requires that x be a read-only array.

    • is_array(x, frozen=None) does not apply any restrictions to the array’s immutability.

  • Testing sparsity.

    • is_array(x, sparse=True) requires that x be a SciPy sparse array (or a quantity whose magnitude is a sparse array).

    • is_array(x, sparse=False) requires that x not be a sparse array.

    • is_array(x, sparse=None) (the default value) does not require x to either be sparse or dense.

All of the above tests can be combined simultaneously in a single function call. Here are a few examples:

import immlib as il, numpy as np

# Make a NumPy array we can query.
arr = np.random.randn(100, 5, 2, 1, 8)

# Is this an array?
il.is_array(arr)
True

Is arr an array with exactly 8000 elements?

il.is_array(arr, numel=8000)
True

Is arr an array that has a quantity and exactly 8000 elements?

il.is_array(arr, numel=8000, quant=True)
False

Is arr an array that has a quantity, exactly 8000 elements, and an shape that ends with 5, 2, _, 8?

il.is_array(arr, numel=8000, shape=(..., 5, 2, -1, 8))
True

The to_array Function#

The to_array function is at its code similar to NumPy’s asarray function in that it attempts to convert the argument into a NumPy array, if possible without copying it. Like is_array, the to_array function considers SciPy sprase arrays and pint quantities whose magnitudes are arrays to be arrays (though this behavior can be modified with options). Overall, the to_array has many of the same options as the is_array function, though the options are applied to the returned array rather than used as a query.

  • Creating a quantity.

    • to_array(x, quant=True) will raise an error if x is not already a quantity (because it does not know what unit to give x).

    • is_array(x, quant=False) will return an array only, stripping the unit off of x if it is a quantity object.

    • to_array(x, quant=None) (the default value) does not attemt to change whether x is a quantity or not and will return a quantity if x is already a quantity (but it will always ensure that the quantity’s magnitude is an array).

  • Changing the array’s physical units.

    • to_array(x, quant=True, unit='mm') returns a quantity with the unit of millimeters. If x is a quantity whose unit is not compatible with millimeters, then an error is raised. Otherwise, the unit and magnitude are converted.

    • to_array(x, unit=None) and to_array(x, quant=False) will both return an array that is not a quantity.

    • to

    • to_array(x, unit=...) and to_array(x, quant=None) (the default values) do not attept to change whether the object is a quantity or its unit.

  • Setting the dtype.

    • to_array(x, dtype=dt) returns an array with the given dtype dt.

  • Immutability.

    • to_array(x, frozen=False) returns an array that is not read-only.

    • to_array(x, frozen=True) returns a copy of x that is read-only unless x is already a read-only array (in which case it is returned).

    • to_array(x, frozen=None) does not apply any restrictions to the array’s immutability, leaving it unchanged.

  • Creating sparse or dense arrays.

    • to_array(x, sparse=True) converts x into a SciPy sparse array and returs it.

    • to_array(x, sparse=False) converts x into a dense array and returns it.

    • to_array(x, sparse=None) (the default value) does not change the sparsity of x.

  • Additional options.

    • copy=True can be passed to to_array() to force the data to be copied into the returned array. The default, copy=False only copies the data when needed.

    • The order parameter is passed along to the numpy.array or numpy.asarray function.

    • detach (default: True) specifies that when a PyTorch tensor that requires gradient tracking is passed to to_array, immlib should detach the tensor and return its NumPy representation. If this is set to False, then an error is raised whenever such a tensor is encountered.

PyTorch Tensor Functions#

In addition to the is_array() and to_array() functions described above, there are two very similar functions for PyTorch tensors: is_tensor() and to_tensor(). These two functions are largely the same as their NumPy equivalents and thus are not documented again. They do, however have a few differences. Primarily, both is_tensor and to_tensor accept options for requires_grad and device that allow querying or manipulating the PyTorch device and gradient tracking requirement. Additionally, because PyTorch does not support read-only tensors, the frozen options are not available.

Generic NumPy/PyTorch Functions#

Finally, the functions is_numeric() and to_numeric() are generic versions of is_array() and to_array() that work for either PyTorch tensors or NumPy arrays. The idea behind these functions is that sometimes all that matters is that an input variable be a numeric collection with a certain shape and type, and that can be easily detected with a call to is_numeric(x, dtype=int, shape=(2,-1)) (potentially with different options). The to_numeric() function always returns either a tensor or an array (or a quantity whose magnitude is a tensor or array, keeping in mind that SciPy sprase arrays are considered arrays).

Additionally, immlib defines a few utility functions related to the above that should behave similarly for array and tensor inputs. All of these functions accept a subset of the options that the is_numeric and to_numeric functions accept.

  • is_sparse(x) detects whether a tensor or array is sparse.

  • to_sparse(x) converts x into a sparse array or tensor.

  • is_dense(x) detects whether a tensor or array is dense.

  • to_dense(x) converts x into a dense array or tensor.

The @numapi Decorator#

A common issue when designing scientific libraries is the need to choose between NumPy and PyTorch. Because the two libraries use slightly different APIs and have slightly different ways of expressing the same formulae, the two can’t always be programmed using the same code. To simplify the job of handling both cases, a function can be decorated using the @numapi decorator, then two different versions of the function can be declared: one for NumPy arrays and one for PyTorch tensors.

import immlib as il

# Declare a function that uses different code for numpy arrays and pytorch
# tensors:
@il.numapi
def betafn(alpha, beta):
   """Calculates and returns the Beta function B(a, b).
   
   If any of the parameters are PyTorch tensors, then a PyTorch tensor is
   returned. Otherwise, a NumPy array is returned.
   """
   pass

# The NumPy version of this function:
@betafn.array
def _(a, b):
    from scipy.special import beta
    return beta(a, b)

# The PyTorch version of this function:
@betafn.tensor
def _(a, b):
    import torch
    la = torch.lgamma(a)
    lb = torch.lgamma(b)
    lab = torch.lgamma(a + b)
    return torch.exp(la + lb - lab)

# When we give this function a PyTorch tensor, it returns the same type:
betafn(1.0, il.to_tensor([2.5, 3.2]))
tensor([0.4000, 0.3125])
# But if we give it any other data, it returns NumPy arrays:
betafn(1.0, [2.5, 3.2])
array([0.4   , 0.3125])

The @array_args, @tensor_args, and @numeric_args Decorators#

Finally, the @array_args, @tensor_args, and @numeric_args decorators ensure that any arguments to the decorated function are automatically promoted to the associated type. Optionally, a list of the names of the arguments that should be promoted can be given, causing the decorator to skip any argument that isn’t listed. The numeric_args decorator promotes all arguments (or all listed arguments) to arrays or all arguments to tensors but it will never promote some to arrays and others to tensors; tensors are used if any of the listed arguments are tensors.

The @tensor_args deocrator has an extra option, keep_arrays, whose default value is False. If this value is set to True, then all returned tensors are converted back into NumPy arrays if none of the arguments were tensors. This effectively allows one to write a PyTorch implementation of a numerical function that can then operate over PyTorch tensors or NumPy arrays.

The following example demonstrates this method of writing a PyTorch implementation only:

import immlib as il

@il.tensor_args(keep_arrays=True)
def betafn(a, b):
    import torch
    print('Type of a:', type(a))
    print('Type of b:', type(a))
    la = torch.lgamma(a)
    lb = torch.lgamma(b)
    lab = torch.lgamma(a + b)
    return torch.exp(la + lb - lab)
 
# Call the function using PyTorch tensors:
betafn(1.0, il.to_tensor([2.5, 3.2]))
Type of a: <class 'torch.Tensor'>
Type of b: <class 'torch.Tensor'>
tensor([0.4000, 0.3125])
# Call the function without PyTorch tensors:
betafn(1.0, [2.5, 3.2])
Type of a: <class 'torch.Tensor'>
Type of b: <class 'torch.Tensor'>
array([0.39999998, 0.3125001 ], dtype=float32)