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 thatxbe apintquantity with an array magnitude.is_array(x, quant=False)requires thatxbe not be apintquantity (instead it must be a plain NumPy array).is_array(x, quant=None)(the default value) allows thatxbe either a NumPy array or apintquantity with a NumPy array as its magnitude.
Testing the object’s physical units.
is_array(x, quant=True, unit='mm')requires thatxbe apintquantity 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 thatxnot have a unit, meaning thatxmust not be apintquantity.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 byxhave a dtype that is a sub-dtype ofdt. For exampleis_array(x, dtype=np.number)will require thatxhave 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 tod.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 ofxmatchsh. When specifying the shape that the array must match, the objectshcan take a few forms. First, it can be either a tuple of numbers of an integer. When an integernis 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, anEllipsisobject (...) 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 returnTruefor 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 ton.Note that
ndim,shape, andnumelcan be used simultaneously.
Testing immutability.
is_array(x, frozen=False)requires thatxnot be a read-only array.is_array(x, frozen=True)requires thatxbe 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 thatxbe a SciPy sparse array (or a quantity whose magnitude is a sparse array).is_array(x, sparse=False)requires thatxnot be a sparse array.is_array(x, sparse=None)(the default value) does not requirexto 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 ifxis not already a quantity (because it does not know what unit to givex).is_array(x, quant=False)will return an array only, stripping the unit off ofxif it is a quantity object.to_array(x, quant=None)(the default value) does not attemt to change whetherxis a quantity or not and will return a quantity ifxis 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. Ifxis 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)andto_array(x, quant=False)will both return an array that is not a quantity.to
to_array(x, unit=...)andto_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 dtypedt.
Immutability.
to_array(x, frozen=False)returns an array that is not read-only.to_array(x, frozen=True)returns a copy ofxthat is read-only unlessxis 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)convertsxinto a SciPy sparse array and returs it.to_array(x, sparse=False)convertsxinto a dense array and returns it.to_array(x, sparse=None)(the default value) does not change the sparsity ofx.
Additional options.
copy=Truecan be passed toto_array()to force the data to be copied into the returned array. The default,copy=Falseonly copies the data when needed.The
orderparameter is passed along to thenumpy.arrayornumpy.asarrayfunction.detach(default:True) specifies that when a PyTorch tensor that requires gradient tracking is passed toto_array,immlibshould detach the tensor and return its NumPy representation. If this is set toFalse, 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)convertsxinto a sparse array or tensor.is_dense(x)detects whether a tensor or array is dense.to_dense(x)convertsxinto 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)