def from_array(cls, x: Tensor, scale: float = 1.0) -> "Bounds":
"""
Instantiate a bounds compatible for bounding the given array. It also allows to set a \
margin for the high and low values.
The value of the high and low will be proportional to the maximum and minimum values of \
the array. Scale defines the proportion to make the bounds bigger and smaller. For \
example, if scale is 1.1 the higher bound will be 10% higher, and the lower bounds 10% \
smaller. If scale is 0.9 the higher bound will be 10% lower, and the lower bound 10% \
higher. If scale is one, `high` and `low` will be equal to the maximum and minimum values \
of the array.
Args:
x: Numpy array used to initialize the bounds.
scale: Value representing the tolerance in percentage from the current maximum and \
minimum values of the array.
Returns:
:class:`Bounds` instance.
Examples:
>>> import numpy
>>> x = numpy.ones((3, 3))
>>> x[1:-1, 1:-1] = -5
>>> bounds = Bounds.from_array(x, scale=1.5)
>>> print(bounds)
Bounds shape float64 dtype (3,) low [ 0.5 -7.5 0.5] high [1.5 1.5 1.5]
"""
xmin, xmax = API.min(x, axis=0), API.max(x, axis=0)
xmin_scaled, xmax_scaled = cls.get_scaled_intervals(xmin, xmax, scale)
return Bounds(low=xmin_scaled, high=xmax_scaled)
def test_split_states(self, states_class):
batch_size = 20
new_states = states_class(batch_size=batch_size, test="test")
for s in new_states.split_states(batch_size):
assert len(s) == 1
assert s.test == "test"
data = API.repeat(API.arange(5).reshape(1, -1), batch_size, 0)
new_states = states_class(batch_size=batch_size, test="test", data=data)
for s in new_states.split_states(batch_size):
assert len(s) == 1
assert s.test == "test"
assert bool((s.data == API.arange(5)).all()), s.data
chunk_len = 4
test_data = API.repeat(API.arange(5).reshape(1, -1), chunk_len, 0)
for s in new_states.split_states(5):
assert len(s) == chunk_len
assert s.test == "test"
assert (s.data == test_data).all(), (s.data.shape, test_data.shape)
batch_size = 21
data = API.repeat(API.arange(5).reshape(1, -1), batch_size, 0)
new_states = states_class(batch_size=batch_size, test="test", data=data)
chunk_len = 5
test_data = API.repeat(API.arange(5).reshape(1, -1), chunk_len, 0)
split_states = list(new_states.split_states(5))
for s in split_states[:-1]:
assert len(s) == chunk_len
assert s.test == "test"
assert (s.data == test_data).all(), (s.data.shape, test_data.shape)
assert len(split_states[-1]) == 1
assert split_states[-1].test == "test"
assert (split_states[-1].data == API.arange(5)).all(), (s.data.shape, test_data.shape)
def test_points_in_bounds(self, bounds_fixture):
zeros = API.zeros((3, 3))
assert all(bounds_fixture.points_in_bounds(zeros))
tens = API.ones((3, 3)) * 10.0
res = bounds_fixture.points_in_bounds(tens)
assert not res.any(), (res, tens)
tens = tensor([[-10, 0, 1], [0, 0, 0], [10, 10, 10]])
assert sum(bounds_fixture.points_in_bounds(tens)) == 1
def small_tree():
node_data = {"a": API.arange(10), "b": API.zeros(10)}
edge_data = {"c": API.ones(10)}
g = networkx.DiGraph()
for i in range(8):
g.add_node(to_node_id(i), **node_data)
pairs = [(0, 1), (1, 2), (2, 3), (2, 4), (2, 5), (3, 6), (3, 7)]
for a, b in pairs:
g.add_edge(to_node_id(a), to_node_id(b), **edge_data)
return g
def test_merge_states(self, states_class):
batch_size = 21
data = API.repeat(API.arange(5).reshape(1, -1), batch_size, 0)
new_states = states_class(batch_size=batch_size, test="test", data=data)
split_states = tuple(new_states.split_states(batch_size))
merged = new_states.merge_states(split_states)
assert len(merged) == batch_size
assert merged.test == "test"
assert (merged.data == data).all()
split_states = tuple(new_states.split_states(5))
merged = new_states.merge_states(split_states)
assert len(merged) == batch_size
assert merged.test == "test"
assert (merged.data == data).all()
def _ix(self, index: int):
# TODO(guillemdb): Allow slicing
data = {
k: API.unsqueeze(v[index], 0) if judo.is_tensor(v) else v
for k, v in self.items()
}
return self.__class__(batch_size=1, **data)
def params_to_arrays(param_dict: StateDict,
n_walkers: int) -> Dict[str, Tensor]:
"""
Create a dictionary containing the arrays specified by param_dict.
Args:
param_dict: Dictionary defining the attributes of the tensors.
n_walkers: Number items in the first dimension of the data tensors.
Returns:
Dictionary with the same keys as param_dict, containing arrays specified \
by `param_dict` values.
"""
tensor_dict = {}
for key, val in param_dict.items():
# Shape already includes the number of walkers. Remove walkers axis to create size.
shape = val.get("shape")
if shape is None:
val_size = val.get("size")
elif len(shape) > 1:
val_size = shape[1:]
else:
val_size = val.get("size")
# Create appropriate shapes with current state's number of walkers.
sizes = n_walkers if val_size is None else tuple([n_walkers
]) + val_size
if "size" in val:
del val["size"]
if "shape" in val:
del val["shape"]
tensor_dict[key] = API.zeros(sizes, **val)
return tensor_dict
def merge_one_name(states_list, name):
vals = []
for state in states_list:
data = state[name]
# Attributes that are not tensors are not stacked.
if not judo.is_tensor(data):
return data
state_len = len(state)
if len(data.shape) == 0 and state_len == 1:
# Name is scalar vector. Data is typing.Scalar value. Transform to array first
value = tensor([data]).flatten()
elif len(data.shape) == 1 and state_len == 1:
if data.shape[0] == 1:
# Name is typing.Scalar vector. Data already transformed to an array
value = data
else:
# Name is a matrix of vectors. Data needs an additional dimension
value = tensor([data])
elif len(data.shape) == 1 and state_len > 1:
# Name is a typing.Scalar vector. Data already has is a one dimensional array
value = data
elif (len(data.shape) > 1 and state_len > 1
or len(data.shape) > 1 and len(state) == 1):
# Name is a matrix of vectors. Data has the correct shape
value = data
else:
raise ValueError(
"Could not infer data concatenation for attribute %s with shape %s"
% (name, data.shape), )
vals.append(value)
return API.concatenate(vals)
def test_from_array_with_scale_positive(self):
array = tensor([[0, 0, 0], [10, 0, 0], [0, 10, 0], [10, 10, 10]],
dtype=dtype.float)
bounds = Bounds.from_array(array, scale=1.1)
assert (bounds.low == tensor([0, 0, 0], dtype=dtype.float)).all(), (
bounds.low,
array.min(axis=0),
)
assert (bounds.high == tensor([11, 11, 11],
dtype=dtype.float)).all(), (
bounds.high,
array.max(axis=0),
)
assert bounds.shape == (3, )
array = tensor(
[[-10, 0, 0], [-10, 0, 0], [0, -10, 0], [-10, -10, -10]],
dtype=dtype.float)
bounds = Bounds.from_array(array, scale=1.1)
assert (bounds.high == tensor([0, 0, 0], dtype=dtype.float)).all(), (
bounds.high,
array.max(axis=0),
)
assert (bounds.low == tensor([-11, -11, -11],
dtype=dtype.float)).all(), (
bounds.low,
array.min(axis=0),
)
assert bounds.shape == (3, )
array = tensor(
[[10, 10, 10], [100, 10, 10], [10, 100, 10], [100, 100, 100]],
dtype=dtype.float)
bounds = Bounds.from_array(array, scale=1.1)
assert API.allclose(bounds.low, tensor([9.0, 9.0, 9],
dtype=dtype.float)), (
bounds.low,
array.min(axis=0),
)
assert API.allclose(bounds.high,
tensor([110, 110, 110], dtype=dtype.float)), (
bounds.high,
array.max(axis=0),
)
assert bounds.shape == (3, )
def test_setitem(self, states_class):
name_1 = "miau"
val_1 = name_1
name_2 = "elephant"
val_2 = API.arange(10)
new_states = states_class(batch_size=2)
new_states[name_1] = val_1
new_states[name_2] = val_2
assert new_states[name_1] == val_1, type(new_states)
assert (new_states[name_2] == val_2).all(), type(new_states)
def test_clip(self):
tup = ((-1, 10), (-3, 4), (2, 5))
array = tensor([[-10, 0, 0], [11, 0, 0], [0, 11, 0], [11, 11, 11]],
dtype=dtype.float)
bounds = Bounds.from_tuples(tup)
clipped = bounds.clip(array)
target = tensor(
[[-1.0, 0.0, 2.0], [10.0, 0.0, 2.0], [0.0, 4.0, 2], [10, 4, 5]],
dtype=dtype.float)
assert API.allclose(clipped, target), (clipped.dtype, target.dtype)
def clip(self, x: Tensor) -> Tensor:
"""
Clip the values of the target array to fall inside the bounds (closed interval).
Args:
x: Numpy array to be clipped.
Returns:
Clipped numpy array with all its values inside the defined bounds.
"""
return API.clip(judo.astype(x, dtype.float), self.low, self.high)
def test_safe_margin(self, bounds_fixture: Bounds):
new_bounds = bounds_fixture.safe_margin()
assert API.allclose(new_bounds.low, bounds_fixture.low)
assert API.allclose(new_bounds.high, bounds_fixture.high)
low = API.full_like(bounds_fixture.low, -10)
new_bounds = bounds_fixture.safe_margin(low=low)
assert API.allclose(new_bounds.high, bounds_fixture.high)
assert API.allclose(new_bounds.low, low)
new_bounds = bounds_fixture.safe_margin(low=low, scale=2)
assert API.allclose(new_bounds.high, bounds_fixture.high * 2)
assert API.allclose(new_bounds.low, low * 2)
def test_append_leaf(self, tree):
node_data = {"node": API.arange(10)}
edge_data = {"edge": False}
leaf_id = to_node_id(-421)
epoch = 123
tree.append_leaf(
leaf_id=leaf_id,
parent_id=tree.root_id,
node_data=node_data,
edge_data=edge_data,
epoch=epoch,
)
assert (tree.data.nodes[leaf_id]["node"] == node_data["node"]).all()
assert tree.data.nodes[leaf_id]["epoch"] == epoch
assert tree.data.edges[(tree.root_id, leaf_id)] == edge_data
assert leaf_id in tree.leafs
assert tree.root_id not in tree.leafs
def split_states(self, n_chunks: int) -> Generator["States", None, None]:
"""
Return a generator for n_chunks different states, where each one \
contain only the data corresponding to one walker.
"""
def get_chunck_size(state, start, end):
for name in state._names:
attr = state[name]
if judo.is_tensor(attr):
return len(attr[start:end])
return int(numpy.ceil(self.n / n_chunks))
for start, end in API.similiar_chunks_indexes(self.n, n_chunks):
chunk_size = get_chunck_size(self, start, end)
data = {
k: val[start:end] if judo.is_tensor(val) else val
for k, val in self.items()
}
new_state = self.__class__(batch_size=chunk_size, **data)
yield new_state
def get_scaled_intervals(
low: Union[Tensor, float, int],
high: Union[Tensor, float, int],
scale: float,
) -> Tuple[Union[Tensor, float], Union[Tensor, float]]:
"""
Scale the high and low vectors by an scale factor.
The value of the high and low will be proportional to the maximum and minimum values of \
the array. Scale defines the proportion to make the bounds bigger and smaller. For \
example, if scale is 1.1 the higher bound will be 10% higher, and the lower bounds 10% \
smaller. If scale is 0.9 the higher bound will be 10% lower, and the lower bound 10% \
higher. If scale is one, `high` and `low` will be equal to the maximum and minimum values \
of the array.
Args:
high: Higher bound to be scaled.
low: Lower bound to be scaled.
scale: Value representing the tolerance in percentage from the current maximum and \
minimum values of the array.
Returns:
:class:`Bounds` instance.
"""
pct = tensor(scale - 1)
big_scale = 1 + API.abs(pct)
small_scale = 1 - API.abs(pct)
zero = judo.astype(tensor(0.0), low.dtype)
if pct > 0:
xmin_scaled = API.where(low < zero, low * big_scale,
low * small_scale)
xmax_scaled = API.where(high < zero, high * small_scale,
high * big_scale)
else:
xmin_scaled = API.where(low < zero, low * small_scale,
low * small_scale)
xmax_scaled = API.where(high < zero, high * big_scale,
high * small_scale)
return xmin_scaled, xmax_scaled
def __init__(
self,
high: Union[Tensor, Scalar] = numpy.inf,
low: Union[Tensor, Scalar] = numpy.NINF,
shape: Optional[tuple] = None,
dtype: Optional[type] = None,
):
"""
Initialize a :class:`Bounds`.
Args:
high: Higher value for the bound interval. If it is an typing_.Scalar \
it will be applied to all the coordinates of a target vector. \
If it is a vector, the bounds will be checked coordinate-wise. \
It defines and closed interval.
low: Lower value for the bound interval. If it is a typing_.Scalar it \
will be applied to all the coordinates of a target vector. \
If it is a vector, the bounds will be checked coordinate-wise. \
It defines and closed interval.
shape: Shape of the array that will be bounded. Only needed if `high` and `low` are \
vectors and it is used to define the dimensions that will be bounded.
dtype: Data type of the array that will be bounded. It can be inferred from `high` \
or `low` (the type of `high` takes priority).
Examples:
Initializing :class:`Bounds` using numpy arrays:
>>> import numpy
>>> high, low = numpy.ones(3, dtype=float), -1 * numpy.ones(3, dtype=int)
>>> bounds = Bounds(high=high, low=low)
>>> print(bounds)
Bounds shape float64 dtype (3,) low [-1 -1 -1] high [1. 1. 1.]
Initializing :class:`Bounds` using typing_.Scalars:
>>> import numpy
>>> high, low, shape = 4, 2.1, (5,)
>>> bounds = Bounds(high=high, low=low, shape=shape)
>>> print(bounds)
Bounds shape float64 dtype (5,) low [2.1 2.1 2.1 2.1 2.1] high [4. 4. 4. 4. 4.]
"""
# Infer shape if not specified
if shape is None and hasattr(high, "shape"):
shape = high.shape
elif shape is None and hasattr(low, "shape"):
shape = low.shape
elif shape is None:
raise TypeError(
"If shape is None high or low need to have .shape attribute.")
# High and low will be arrays of target shape
if not judo.is_tensor(high):
high = tensor(high) if isinstance(
high, _Iterable) else API.ones(shape) * high
if not judo.is_tensor(low):
low = tensor(low) if isinstance(
low, _Iterable) else API.ones(shape) * low
self.high = judo.astype(high, dtype)
self.low = judo.astype(low, dtype)
if dtype is not None:
self.dtype = dtype
elif hasattr(high, "dtype"):
self.dtype = high.dtype
elif hasattr(low, "dtype"):
self.dtype = low.dtype
else:
self.dtype = type(high) if high is not None else type(low)
def test_from_array_with_scale_negative(self):
# high +, low +, scale > 1
array = tensor(
[[-10, 0, 0], [-10, 0, 0], [0, -10, 0], [-10, -10, -10]],
dtype=dtype.float)
bounds = Bounds.from_array(array, scale=0.9)
assert (bounds.high == tensor([0, 0, 0], dtype=dtype.float)).all(), (
bounds.high,
array.max(axis=0),
)
assert (bounds.low == tensor([-9, -9, -9], dtype=dtype.float)).all(), (
bounds.low,
array.min(axis=0),
)
assert bounds.shape == (3, )
array = tensor([[0, 0, 0], [10, 0, 0], [0, 10, 0], [10, 10, 10]],
dtype=dtype.float)
bounds = Bounds.from_array(array, scale=0.9)
assert (bounds.low == tensor([0, 0, 0],
dtype=dtype.float)).all(), (bounds, array)
assert (bounds.high == tensor([9, 9, 9], dtype=dtype.float)).all()
assert bounds.shape == (3, )
# high +, low +, scale < 1
array = tensor(
[[10, 10, 10], [100, 10, 10], [10, 100, 10], [100, 100, 100]],
dtype=dtype.float)
bounds = Bounds.from_array(array, scale=0.9)
assert API.allclose(bounds.low,
tensor([9.0, 9.0, 9.0], dtype=dtype.float)), (
bounds.low,
array.min(axis=0),
)
assert API.allclose(bounds.high, tensor([90, 90, 90],
dtype=dtype.float)), (
bounds.high,
array.max(axis=0),
)
assert bounds.shape == (3, )
# high -, low -, scale > 1
array = tensor(
[[-100, -10, -10], [-100, -10, -10], [-10, -100, -10],
[-100, -100, -100]],
dtype=dtype.float,
)
bounds = Bounds.from_array(array, scale=1.1)
assert API.allclose(bounds.high, tensor([-9, -9, -9],
dtype=dtype.float)), (
bounds.high,
array.max(axis=0),
)
assert API.allclose(bounds.low,
tensor([-110, -110, -110], dtype=dtype.float)), (
bounds.low,
array.min(axis=0),
)
assert bounds.shape == (3, )
# high -, low -, scale < 1
array = tensor(
[[-100, -10, -10], [-100, -10, -10], [-10, -100, -10],
[-100, -100, -100]],
dtype=dtype.float,
)
bounds = Bounds.from_array(array, scale=0.9)
assert API.allclose(bounds.high,
tensor([-11, -11, -11], dtype=dtype.float)), (
bounds.high,
array.max(axis=0),
)
assert API.allclose(bounds.low,
tensor([-90, -90, -90], dtype=dtype.float)), (
bounds.low,
array.min(axis=0),
)
assert bounds.shape == (3, )