skbio.tree.BPTree.to_device

BPTree.to_device(xp, device=None)

Return this tree with its data in an array backend and on a device.

This is how a tree is moved to a GPU, e.g. after it was read from a file, so that its batch operations and cophenet can run there (see Notes of BPTree). It is also how a tree is moved back to NumPy, or between devices.

Parameters:

xpmodule

The array backend, e.g. cupy, torch, jax.numpy or numpy, or its array API namespace.

devicedevice, optional

The device, in the terms of xp: e.g. "cuda" or "cuda:1" for PyTorch, a device number or a cupy.cuda.Device for CuPy, or the device of an existing array of xp. Defaults to the default device of xp: the current GPU for CuPy, and the CPU for PyTorch and NumPy.

Returns:

BPTree

A new tree whose data is a copy of the topology in xp on device. It shares this tree’s navigation index and node attributes (lengths, names, edges), which stay on the host, so nothing is rebuilt. Setting an attribute of either tree (e.g. with set_names) does not change the other.

Raises:

TypeError

If xp is not a supported array backend.

See also

data

Notes

The topology is copied from the host, one byte per parenthesis. On a GPU, the navigation index follows at the first operation that needs it there, and is then kept on the device with the tree.

Examples

>>> import numpy as np
>>> from skbio.tree import BPTree
>>> tree = BPTree.read(["((a:1,b:2)c:3,(d:4,e:5)f:6)root;"])

On a GPU with CuPy, for example, tree.to_device(cupy) returns the tree with its data there, and tree.to_device(torch, "cuda") does the same with PyTorch. Moved to NumPy, the tree is on the host:

>>> host = tree.to_device(np)
>>> type(host.data).__name__
'ndarray'
>>> host.lca_batch([2], [4])  # the lowest common ancestor of a and b: c
array([1])