def broadcast_to(self, shape, **kw):
"""
Return a view of the array broadcasted to another shape. See
numpy.broadcast_to.
"""
numpy.broadcast_shapes(self.shape,
shape) # raises if not broadcastable
d = {
name: broadcast_to(x, shape + self.dtype[name].shape, **kw)
for name, x in self._dict.items()
}
return self._array(shape, self.dtype, d)
def _set_values(cls, lower, upper, size, shape, initval):
if size is None:
size = shape
lower = np.asarray(lower)
lower = floatX(np.where(lower == None, -np.inf, lower))
upper = np.asarray(upper)
upper = floatX(np.where(upper == None, np.inf, upper))
if initval is None:
_size = np.broadcast_shapes(to_tuple(size), np.shape(lower),
np.shape(upper))
_lower = np.broadcast_to(lower, _size)
_upper = np.broadcast_to(upper, _size)
initval = np.where(
(_lower == -np.inf) & (_upper == np.inf),
0,
np.where(
_lower == -np.inf,
_upper - 1,
np.where(_upper == np.inf, _lower + 1,
(_lower + _upper) / 2),
),
)
lower = as_tensor_variable(floatX(lower))
upper = as_tensor_variable(floatX(upper))
return lower, upper, initval
def __mul__(self, other):
if isinstance(other, gpuarray.GPUArray):
result = type(self)(np.broadcast_shapes(self.shape, other.shape),
gpuarray._get_common_dtype(self, other))
return self._elwise_multiply(other, result)
else:
return super().__mul__(other)
def copy_from_usm_ndarray_to_usm_ndarray(dst, src):
if type(dst) is not dpt.usm_ndarray or type(src) is not dpt.usm_ndarray:
raise TypeError
if dst.ndim == src.ndim and dst.shape == src.shape:
copy_same_shape(dst, src)
try:
common_shape = np.broadcast_shapes(dst.shape, src.shape)
except ValueError:
raise ValueError
if dst.size < src.size:
raise ValueError
if len(common_shape) > dst.ndim:
ones_count = len(common_shape) - dst.ndim
for k in range(ones_count):
if common_shape[k] != 1:
raise ValueError
common_shape = common_shape[ones_count:]
if src.ndim < len(common_shape):
new_src_strides = (0, ) * (len(common_shape) - src.ndim) + src.strides
src_same_shape = dpt.usm_ndarray(common_shape,
dtype=src.dtype,
buffer=src,
strides=new_src_strides)
else:
src_same_shape = src
src_same_shape.shape = common_shape
copy_same_shape(dst, src_same_shape)
def __init__(self, ir: array, ishape: Shape, dim: int, name: str = "Conv"):
"""ND convolution on last `N` axis.
Parameters
----------
ir : array
The impulse responses. Must be at least of `ndim==dim`.
ishape : tuple of int
The shape of the input images. Images are on the last two axis.
dim : int
The last `dim` axis where convolution apply.
"""
super().__init__(
ishape=ishape,
oshape=np.broadcast_shapes(ishape, ir.shape[:-dim] + ishape[-dim:]),
name=name,
)
self.dim = dim
self.imp_resp = ir
self.freq_resp = udft.ir2fr(ir, self.ishape[-dim:])
self.margins = ir.shape[-dim:]
self._slices = [slice(None) for _ in range(len(ishape) - dim)]
for idx in reversed(range(dim)):
self._slices.append(
slice(
ir.shape[idx] // 2,
ishape[idx] - ir.shape[idx] // 2 + ir.shape[idx] % 2,
)
)
def _make_casadi_types_broadcastable(x1, x2):
def shape_2D(object: Union[float, int, Iterable, _onp.ndarray]) -> Tuple:
shape = _onp.shape(object)
if len(shape) == 0:
return (1, 1)
elif len(shape) == 1:
return (1, shape[0])
elif len(shape) == 2:
return shape
else:
raise ValueError(
"CasADi can't handle arrays with >2 dimensions, unfortunately."
)
x1_shape = shape_2D(x1)
x2_shape = shape_2D(x2)
shape = _onp.broadcast_shapes(x1_shape, x2_shape)
x1_tiled = _cas.repmat(
x1,
shape[0] // x1_shape[0],
shape[1] // x1_shape[1],
)
x2_tiled = _cas.repmat(
x2,
shape[0] // x2_shape[0],
shape[1] // x2_shape[1],
)
return x1_tiled, x2_tiled
def broadcasted_axes(*dfs):
"""
Helper function which, from a collection of arrays, series, frames and other
values, retrieves the axes of series and frames which result from
broadcasting operations. It checks whether index and columns of given
series and frames, repespectively, are aligned. Using this function allows
to subsequently use pure numpy operations and keep the axes in the
background.
"""
axes = []
shape = (1, )
if set(map(type, dfs)) == {tuple}:
dfs = sum(dfs, ())
for df in dfs:
shape = np.broadcast_shapes(shape, np.asarray(df).shape)
if isinstance(df, (pd.Series, pd.DataFrame)):
if len(axes):
assert (axes[-1] == df.axes[-1]).all(), (
'Series or DataFrames '
'are not aligned. Please make sure that all indexes and '
'columns of Series and DataFrames going into the linear '
'expression are equally sorted.')
axes = df.axes if len(df.axes) > len(axes) else axes
return axes, shape
def broadcast_arrays(*arrays, **kw):
"""
Version of numpy.broadcast_arrays that works with StructuredArray and JAX
arrays.
"""
shapes = [a.shape for a in arrays]
shape = numpy.broadcast_shapes(*shapes)
return [broadcast_to(a, shape, **kw) for a in arrays]
def test_string_comparisons_empty(op, ufunc, sym, dtypes):
arr = np.empty((1, 0, 1, 5), dtype=dtypes[0])
arr2 = np.empty((100, 1, 0, 1), dtype=dtypes[1])
expected = np.empty(np.broadcast_shapes(arr.shape, arr2.shape), dtype=bool)
assert_array_equal(op(arr, arr2), expected)
assert_array_equal(ufunc(arr, arr2), expected)
assert_array_equal(np.compare_chararrays(arr, arr2, sym, False), expected)
def _check_shapes(
self,
x0_shape: ShapeType,
x1_shape: Optional[ShapeType] = None,
) -> ShapeType:
"""Checks input argument shapes and computes the broadcast batch shape of both
inputs.
This function checks the shapes of the inputs to the :meth:`__call__` method and
it computes the `bcast_batch_shape` mentioned in the docstring.
Parameters
----------
x0_shape :
Shape of the first input to the covariance function.
x1_shape :
Shape of the (optional) second input to the covariance function.
Returns
-------
broadcast_batch_shape :
The `batch_shape` after broadcasting the inputs to a common shape.
Raises
-------
ValueError
If one of the input shapes is not of the form ``batch_shape_{0,1} +``
:attr:`input_shape`.
ValueError
If the inputs can not be broadcast to a common shape.
"""
err_msg = (
"The shape of the input array `{argname}` must match the `input_shape` "
f"`{self.input_shape}` of the kernel along its last dimension, but an "
"array with shape `{shape}` was given.")
if x0_shape[len(x0_shape) - self._input_ndim:] != self.input_shape:
raise ValueError(err_msg.format(argname="x0", shape=x0_shape))
broadcast_batch_shape = x0_shape[:len(x0_shape) - self._input_ndim]
if x1_shape is not None:
if x1_shape[len(x1_shape) - self._input_ndim:] != self.input_shape:
raise ValueError(err_msg.format(argname="x1", shape=x1_shape))
try:
broadcast_batch_shape = np.broadcast_shapes(
broadcast_batch_shape,
x1_shape[:len(x1_shape) - self._input_ndim],
)
except ValueError as ve:
err_msg = (
f"The input arrays `x0` and `x1` with shapes {x0_shape} and "
f"{x1_shape} can not be broadcast to a common shape.")
raise ValueError(err_msg) from ve
return broadcast_batch_shape
def branch_lengths(self) -> np.ndarray:
batch_shape = np.broadcast_shapes(
self.heights.shape[:-1], self.topology.parent_indices.shape[:-1]
)
heights_b = np.broadcast_to(self.heights, batch_shape + self.heights.shape[-1:])
parent_indices_b = np.broadcast_to(
self.topology.parent_indices,
batch_shape + self.topology.parent_indices.shape[-1:],
)
return (
np.take_along_axis(heights_b, parent_indices_b, axis=-1)
- heights_b[..., :-1]
)
def __pow__(
self, other
) -> "Tensor": # With power, there are some limitations. The exponent must not be differentiable. Future versions of my library will correct for this, however.
if isinstance(other, Tensor):
common_shape = np.broadcast_shapes(self.shape, other.shape)
x, y = broadcast(self, common_shape), broadcast(other, common_shape)
output = Tensor(
x.data ** y.data, x.requires_grad or y.requires_grad, Pow(x, y),
)
return output
return self ** Tensor(other)
def __mul__(self, other) -> "Tensor":
if isinstance(other, Tensor):
common_shape = np.broadcast_shapes(self.shape, other.shape)
x, y = broadcast(self, common_shape), broadcast(other, common_shape)
output = Tensor(
x.data * y.data,
x.requires_grad or y.requires_grad,
Mul(x, y), # Pass the multiplication operation
)
return output
return self * Tensor(other)
def __add__(self, other, sub=False, rsub=False):
"""Add an array with an array or an array with a scalar."""
if isinstance(other, gpuarray.GPUArray):
# add another vector
result = type(self)(np.broadcast_shapes(self.shape, other.shape),
gpuarray._get_common_dtype(self, other))
return self._axpbyz(-1 if rsub else 1, other, -1 if sub else 1,
result)
else:
if sub:
return super().__sub__(other)
elif rsub:
return super().__rsub__(other)
else:
return super().__add__(other)
def __sub__(self, other) -> "Tensor":
if isinstance(other, Tensor):
common_shape = np.broadcast_shapes(self.shape, other.shape)
x, y = (
broadcast(self, common_shape),
broadcast(other, common_shape),
) # Same pattern here, like with addition
output = Tensor(
x.data - y.data,
x.requires_grad or y.requires_grad,
Sub(x, y), # Pass the subtraction operation
)
return output
return self - Tensor(other)
def sample(
self,
l: ArrayLike,
m: ArrayLike,
frequency: u.Quantity, # noqa: E741
frame: Union[AltAzFrame, RADecFrame],
output_type: OutputType,
*,
out: Optional[np.ndarray] = None) -> np.ndarray:
_check_out(out, output_type)
l_ = _asarray(l)
m_ = _asarray(m)
l_, m_ = np.broadcast_arrays(l_, m_)
# numba seems to trigger a FutureWarning when it checks the writeable
# flag on these broadcast arrays. Suppress it by making them explicitly
# readonly. Explicitly take views so that we don't modify the input
# arrays (which could otherwise happen if the transformations above are
# no-ops).
l_ = l_.view()
m_ = m_.view()
l_.flags.writeable = False
m_.flags.writeable = False
if isinstance(frame, RADecFrame):
transform = frame.lm_to_hv()
# Broadcast transform with l, m
shape = np.broadcast_shapes(transform.shape[:-2], l_.shape)
l_, m_ = self._transform_lm(
np.broadcast_to(transform, shape + (2, 2)),
np.broadcast_to(l_, shape), np.broadcast_to(m_, shape))
elif not isinstance(frame, AltAzFrame):
raise TypeError(
f'frame must be RADecFrame or AltAzFrame, not {type(frame)}')
l_ = _asarray(l_, np.float32)
m_ = _asarray(m_, np.float32)
if output_type in {OutputType.JONES_XY, OutputType.JONES_HV}:
return self._sample_altaz_jones(l_,
m_,
frequency,
frame,
output_type,
out=out)
else:
jones = self._sample_altaz_jones(l_, m_, frequency, frame,
output_type)
return self._finalize(jones, output_type, out=out)
def align_shape(*polys: PolyLike) -> Tuple[ndpoly, ...]:
"""
Align polynomial by shape.
Args:
polys:
Polynomial to make adjustment to.
Returns:
Same as ``polys``, but internal adjustments made to make them
compatible for further operations.
Examples:
>>> q0, q1 = numpoly.variable(2)
>>> poly1 = 4*q0
>>> poly2 = numpoly.polynomial([[2*q0+1, 3*q0-q1]])
>>> poly1.shape
()
>>> poly2.shape
(1, 2)
>>> poly1, poly2 = numpoly.align_shape(poly1, poly2)
>>> poly1
polynomial([[4*q0, 4*q0]])
>>> poly2
polynomial([[2*q0+1, -q1+3*q0]])
>>> poly1.shape
(1, 2)
>>> poly2.shape
(1, 2)
"""
# return tuple(numpoly.broadcast_arrays(*polys))
polys_ = [numpoly.aspolynomial(poly) for poly in polys]
common = numpy.ones(
numpy.broadcast_shapes(*[poly.shape for poly in polys_]), dtype=int)
for idx, poly in enumerate(polys_):
if poly.shape != common.shape:
polys_[idx] = poly.from_attributes(
exponents=poly.exponents,
coefficients=tuple(coeff * common
for coeff in poly.coefficients),
names=poly.indeterminants,
)
return tuple(polys_)
def __add__(
self, other
) -> "Tensor": # Add tensors together, in these element-wise functions, we'll have a broadcast precursor
if isinstance(other, Tensor):
common_shape = np.broadcast_shapes(self.shape, other.shape)
x, y = (
broadcast(self, common_shape),
broadcast(other, common_shape),
) # This broadcasting is essential to make sure gradients are accumulated and scatted properly
output = Tensor( # If either operand required gradients, the output will too.
x.data + y.data,
x.requires_grad or y.requires_grad,
Add(x, y), # Create the output tensor and pass the addition operation.
)
return output
return self + Tensor(other) # If it's not a tensor, make it one.
def sample(self,
l: ArrayLike,
m: ArrayLike,
frequency: u.Quantity,
frame: Union[AltAzFrame, RADecFrame],
output_type: OutputType,
*,
out: Optional[np.ndarray] = None) -> np.ndarray:
l_ = np.asarray(l)
m_ = np.asarray(m)
if output_type != OutputType.UNPOLARIZED_POWER:
raise NotImplementedError(
'Only UNPOLARIZED_POWER is currently implemented')
in_shape = np.broadcast_shapes(l_.shape, m_.shape, frame.shape)
out_shape = frequency.shape + in_shape
if out is not None:
if out.shape != out_shape:
raise ValueError(
f'out must have shape {out_shape}, not {out.shape}')
if out.dtype != np.float32:
raise TypeError(
f'out must have dtype float32, not {out.dtype}')
if not out.flags.c_contiguous:
raise ValueError('out must be C contiguous')
else:
out = np.empty(out_shape, np.float32)
frequency_Hz = frequency.to_value(u.Hz).astype(np.float32,
copy=False,
casting='same_kind')
samples = self._interp_samples(frequency_Hz)
# Create view with l/m axis flattened to 1D for benefit of numba
l_view = np.broadcast_to(l_, in_shape).ravel()
m_view = np.broadcast_to(m_, in_shape).ravel()
out_view = out.view()
out_view.shape = frequency.shape + l_view.shape
_sample_impl(l_view, m_view, samples, self._step, out_view)
return out
def multiply(
x1: PolyLike,
x2: PolyLike,
out: Optional[ndpoly] = None,
where: numpy.typing.ArrayLike = True,
**kwargs: Any,
) -> ndpoly:
"""
Multiply arguments element-wise.
Args:
x1, x2:
Input arrays to be multiplied. If ``x1.shape != x2.shape``, they
must be broadcastable to a common shape (which becomes the shape of
the output).
out:
A location into which the result is stored. If provided, it must
have a shape that the inputs broadcast to. If not provided or
`None`, a freshly-allocated array is returned. A tuple (possible
only as a keyword argument) must have length equal to the number of
outputs.
where:
This condition is broadcast over the input. At locations where the
condition is True, the `out` array will be set to the ufunc result.
Elsewhere, the `out` array will retain its original value. Note
that if an uninitialized `out` array is created via the default
``out=None``, locations within it where the condition is False will
remain uninitialized.
kwargs:
Keyword args passed to numpy.ufunc.
Returns:
The product of `x1` and `x2`, element-wise. This is a scalar if
both `x1` and `x2` are scalars.
Examples:
>>> poly = numpy.arange(9.0).reshape((3, 3))
>>> q0q1q2 = numpoly.variable(3)
>>> numpoly.multiply(poly, q0q1q2)
polynomial([[0.0, q1, 2.0*q2],
[3.0*q0, 4.0*q1, 5.0*q2],
[6.0*q0, 7.0*q1, 8.0*q2]])
"""
x1, x2 = numpoly.align_indeterminants(x1, x2)
dtype = numpy.find_common_type([x1.dtype, x2.dtype], [])
shape = numpy.broadcast_shapes(x1.shape, x2.shape)
where = numpy.asarray(where)
exponents = numpy.unique(numpy.tile(x1.exponents, (len(x2.exponents), 1)) +
numpy.repeat(x2.exponents, len(x1.exponents), 0),
axis=0)
out_ = numpoly.ndpoly(
exponents=exponents,
shape=shape,
names=x1.indeterminants,
dtype=dtype,
) if out is None else out
seen = set()
for expon1, coeff1 in zip(x1.exponents, x1.coefficients):
for expon2, coeff2 in zip(x2.exponents, x2.coefficients):
key = (expon1 + expon2 + x1.KEY_OFFSET).ravel()
key = key.view(f"U{len(expon1)}").item()
if key in seen:
out_.values[key] += numpy.multiply(coeff1,
coeff2,
where=where,
**kwargs)
else:
numpy.multiply(coeff1,
coeff2,
out=out_.values[key],
where=where,
**kwargs)
seen.add(key)
if out is None:
out_ = numpoly.clean_attributes(out_)
return out_
def numpy_broadcast_shapes(*args):
return np.broadcast_shapes(*args)
from typing import List, Dict, Any
import numpy as np
import numpy.typing as npt
AR_f8: npt.NDArray[np.float64]
AR_LIKE_f: List[float]
interface_dict: Dict[str, Any]
reveal_type(np.lib.stride_tricks.DummyArray(interface_dict)) # E: numpy.lib.stride_tricks.DummyArray
reveal_type(np.lib.stride_tricks.as_strided(AR_f8)) # E: numpy.ndarray[Any, numpy.dtype[{float64}]]
reveal_type(np.lib.stride_tricks.as_strided(AR_LIKE_f)) # E: numpy.ndarray[Any, numpy.dtype[Any]]
reveal_type(np.lib.stride_tricks.as_strided(AR_f8, strides=(1, 5))) # E: numpy.ndarray[Any, numpy.dtype[{float64}]]
reveal_type(np.lib.stride_tricks.as_strided(AR_f8, shape=[9, 20])) # E: numpy.ndarray[Any, numpy.dtype[{float64}]]
reveal_type(np.lib.stride_tricks.sliding_window_view(AR_f8, 5)) # E: numpy.ndarray[Any, numpy.dtype[{float64}]]
reveal_type(np.lib.stride_tricks.sliding_window_view(AR_LIKE_f, (1, 5))) # E: numpy.ndarray[Any, numpy.dtype[Any]]
reveal_type(np.lib.stride_tricks.sliding_window_view(AR_f8, [9], axis=1)) # E: numpy.ndarray[Any, numpy.dtype[{float64}]]
reveal_type(np.broadcast_to(AR_f8, 5)) # E: numpy.ndarray[Any, numpy.dtype[{float64}]]
reveal_type(np.broadcast_to(AR_LIKE_f, (1, 5))) # E: numpy.ndarray[Any, numpy.dtype[Any]]
reveal_type(np.broadcast_to(AR_f8, [4, 6], subok=True)) # E: numpy.ndarray[Any, numpy.dtype[{float64}]]
reveal_type(np.broadcast_shapes((1, 2), [3, 1], (3, 2))) # E: tuple[builtins.int]
reveal_type(np.broadcast_shapes((6, 7), (5, 6, 1), 7, (5, 1, 7))) # E: tuple[builtins.int]
reveal_type(np.broadcast_arrays(AR_f8, AR_f8)) # E: list[numpy.ndarray[Any, numpy.dtype[Any]]]
reveal_type(np.broadcast_arrays(AR_f8, AR_LIKE_f)) # E: list[numpy.ndarray[Any, numpy.dtype[Any]]]
def rng_fn(
cls,
rng: np.random.RandomState,
mu: Union[np.ndarray, float],
sigma: Union[np.ndarray, float],
init: float,
steps: int,
size: Tuple[int],
) -> np.ndarray:
"""Gaussian Random Walk generator.
The init value is drawn from the Normal distribution with the same sigma as the
innovations.
Notes
-----
Currently does not support custom init distribution
Parameters
----------
rng: np.random.RandomState
Numpy random number generator
mu: array_like
Random walk mean
sigma: np.ndarray
Standard deviation of innovation (sigma > 0)
init: float
Initialization value for GaussianRandomWalk
steps: int
Length of random walk, must be greater than 1. Returned array will be of size+1 to
account as first value is initial value
size: int
The number of Random Walk time series generated
Returns
-------
ndarray
"""
if steps < 1:
raise ValueError("Steps must be greater than 0")
# If size is None then the returned series should be (*implied_dims, 1+steps)
if size is None:
# broadcast parameters with each other to find implied dims
bcast_shape = np.broadcast_shapes(
np.asarray(mu).shape,
np.asarray(sigma).shape,
np.asarray(init).shape,
)
dist_shape = (*bcast_shape, int(steps))
# If size is None then the returned series should be (*size, 1+steps)
else:
dist_shape = (*size, int(steps))
# Add one dimension to the right, so that mu and sigma broadcast safely along
# the steps dimension
innovations = rng.normal(loc=mu[..., None], scale=sigma[..., None], size=dist_shape)
grw = np.concatenate([init[..., None], innovations], axis=-1)
return np.cumsum(grw, axis=-1)
def chol_solve_numpy(t, b, diageps=None):
"""
t (..., n)
b (..., n, m) or (n,)
t[0] += diageps
m = toeplitz(t)
l = chol(m)
return solve(l, b)
pure numpy, object arrays supported
"""
t = numpy.copy(t, subok=True)
n = t.shape[-1]
b = numpy.asanyarray(b)
vec = b.ndim < 2
if vec:
b = b[:, None]
assert b.shape[-2] == n
if n == 0:
shape = numpy.broadcast_shapes(t.shape[:-1], b.shape[:-2])
shape += (n, ) if vec else b.shape[-2:]
dtype = numpy.result_type(t.dtype, b.dtype)
return numpy.empty(shape, dtype)
if diageps is not None:
t[..., 0] += diageps
if numpy.any(t[..., 0] <= 0):
msg = '1-th leading minor is not positive definite'
raise numpy.linalg.LinAlgError(msg)
norm = numpy.copy(t[..., 0, None], subok=True)
t /= norm
invLb = numpy.copy(numpy.broadcast_arrays(b, t[..., None])[0], subok=True)
prevLi = t
g = numpy.stack([numpy.roll(t, 1, -1), t], -2)
for i in range(1, n):
assert numpy.all(g[..., 0, i] > 0)
rho = -g[..., 1, i, None, None] / g[..., 0, i, None, None]
if numpy.any(numpy.abs(rho) >= 1):
msg = '{}-th leading minor is not positive definite'.format(i + 1)
raise numpy.linalg.LinAlgError(msg)
gamma = numpy.sqrt((1 - rho) * (1 + rho))
g[..., :, i:] += g[..., ::-1, i:] * rho
g[..., :, i:] /= gamma
Li = g[..., 0, i:] # i-th column of L from row i
invLb[..., i:, :] -= invLb[..., i - 1, None, :] * prevLi[..., i:, None]
invLb[..., i, :] /= Li[..., 0, None]
prevLi[..., i:] = Li
g[..., 0, i:] = numpy.roll(g[..., 0, i:], 1, -1)
invLb /= numpy.sqrt(norm[..., None])
if vec:
invLb = numpy.squeeze(invLb, -1)
return invLb
def __init__(self, *arrays):
self.shape = numpy.broadcast_shapes(*(a.shape for a in arrays))
def share_figure(
x,
probability,
choices=None,
weights=1,
xlabel=None,
bins=None,
pct_bins=20,
figsize=(12, 4),
style='stacked',
discrete=None,
xlim=None,
xscale=None,
xmajorticks=None,
xminorticks=None,
include_nests=False,
exclude_alts=None,
format='figure',
**kwargs,
):
"""
Generate a figure of variables over a range of variable values.
Parameters
----------
x : array-like, 1-d
An array giving values for some variable.
probability : array-like, 2-d
The pre-calculated probability array for all cases in this analysis.
First dimension must be the same shape as `x`. The second dimension
represents the alternatives (or similar).
choices : array-like, optional
The observed choices array for all cases in this analysis. If provided,
the first dimension must be the same shape as `x`. The second dimension
represents the alternatives (or similar).
weights : array-like, 1-d, optional
The case weights for all cases in this analysis. If provided,
the shape must be the same shape as `x`.
xlabel : str, optional
A label to use for the x-axis of the resulting figure. If not given,
the value of `x.name` is used if it exists. Set to `False` to omit the
x-axis label.
bins : int, optional
The number of equal-sized bins to use.
pct_bins : int or array-like, default 20
The number of equal-mass bins to use.
style : {'stacked', 'dataframe', 'many'}
The type of output to generate.
discrete : bool, default False
Whether to treat the data values explicitly as discrete (vs continuous)
data. This will change the styling and automatic bin generation. If
there are very few unique values, the data will be assumed to be
discrete anyhow.
xlim : 2-tuple, optional
Explicitly set the range of values shown on the x axis of generated
figures. This can truncate long tails. The actual histogram bins
are not changed.
include_nests : bool, default False
Whether to include nests in the figure.
exclude_alts : Collection, optional
Alternatives to exclude from the figure.
filter : str, optional
A filter that will be used to select only a subset of cases.
format : {'figure','svg'}, default 'figure'
How to return the result if it is a figure. The default is to return
the raw matplotlib Figure instance, ot set to `svg` to get a SVG
rendering as an xmle.Elem.
Returns
-------
Figure, DataFrame, or Elem
"""
if style not in {'stacked', 'dataframe', 'many'}:
raise ValueError("style must be in {'stacked', 'dataframe', 'many'}")
if include_nests and style == 'stacked' and exclude_alts is None:
import warnings
warnings.warn(
"including nests in a stacked figure is likely to give "
"misleading results unless constituent alternatives are omitted")
if exclude_alts is None:
exclude_alts = set()
if xlabel is None:
try:
xlabel = x.name
except AttributeError:
pass
filter_ = slice(None)
h_pr = {}
h_ch = {}
discrete_values = None
if discrete:
discrete_values = numpy.unique(x)
elif discrete is None:
from .histograms import seems_like_discrete_data
discrete, discrete_values = seems_like_discrete_data(
x, return_uniques=True)
pr = numpy.asarray(probability)
if choices is not None:
ch = numpy.asarray(choices)
else:
ch = None
wt = numpy.asarray(weights)
x_discrete_labels = None if discrete_values is None else [
str(i) for i in discrete_values
]
if bins is None:
if isinstance(x.dtype, pandas.CategoricalDtype):
discrete_values = numpy.arange(len(x_discrete_labels))
bins = numpy.arange(len(x_discrete_labels) + 1)
x = x.cat.codes
elif isinstance(pct_bins, int):
bins = numpy.percentile(x, numpy.linspace(0, 100, pct_bins + 1))
else:
bins = numpy.percentile(x, pct_bins)
try:
columns = probability.columns
except AttributeError:
columns = None
else:
columns = dict(enumerate(columns))
# check for correct array shapes, raise helpful message if not compatible
pr_w_shape = numpy.broadcast_shapes(pr[:, 0].shape, wt.shape)
if x.shape != pr_w_shape:
raise ValueError(f"incompatible shapes, "
f"x.shape={x.shape}, "
f"pr.shape={pr.shape}, "
f"wt.shape={wt.shape}, "
f"(pr[:,i]*wt).shape={pr_w_shape}")
if ch is not None:
ch_w_shape = numpy.broadcast_shapes(ch[:, 0].shape, wt.shape)
if x.shape != ch_w_shape:
raise ValueError(f"incompatible shapes, "
f"x.shape={x.shape}, "
f"ch.shape={ch.shape}, "
f"wt.shape={wt.shape}, "
f"(ch[:,i]*wt).shape={ch_w_shape}")
for i in range(pr.shape[1]):
h_pr[i], _ = numpy.histogram(
x,
weights=pr[:, i] * wt,
bins=bins,
)
if ch is not None:
h_ch[i], _ = numpy.histogram(
x,
weights=ch[:, i] * wt,
bins=bins,
)
h_pr = pandas.DataFrame(h_pr)
h_pr.index = pandas.IntervalIndex.from_breaks(bins) # bins[:-1]
h_pr.rename(columns=columns, inplace=True)
_denominator, _ = numpy.histogram(
x,
weights=pr.sum(1) * wt,
bins=bins,
)
h_pr_share = (h_pr / _denominator.reshape(-1, 1))
if ch is not None:
_denominator_ch, _ = numpy.histogram(
x,
weights=ch.sum(1) * wt,
bins=bins,
)
h_ch = pandas.DataFrame(h_ch)
h_ch.index = h_pr.index
h_ch.rename(columns=columns, inplace=True)
h_ch_share = (h_ch / _denominator_ch.reshape(-1, 1))
else:
h_ch_share = None
if discrete:
x_placement = numpy.arange(len(bins) - 1)
x_alignment = 'center'
bin_widths = 0.8
else:
x_placement = bins[:-1]
x_alignment = 'edge'
bin_widths = bins[1:] - bins[:-1]
if xlabel is False:
xlabel = None
if xlim is None:
xlim = (bins[0], bins[-1])
if style == 'dataframe':
if ch is not None:
result = pandas.concat(
{
'Modeled Shares': h_pr_share,
'Observed Shares': h_ch_share,
},
axis=1,
sort=False)
else:
result = pandas.concat({
'Modeled Shares': h_pr_share,
},
axis=1,
sort=False)
result['Count', '*'] = h_pr.sum(1)
if xlabel:
result.index.name = xlabel
elif style == 'stacked':
fig, (ax0, ax1) = plt.subplots(1, 2, figsize=figsize)
bottom0 = 0
bottom1 = 0
for i in h_pr_share.columns:
ax0.bar(
x_placement,
height=h_pr_share[i],
bottom=bottom0,
width=bin_widths,
align=x_alignment,
label=i,
)
bottom0 = h_pr_share[i].fillna(0).values + bottom0
ax1.bar(
x_placement,
height=h_ch_share[i],
bottom=bottom1,
width=bin_widths,
align=x_alignment,
)
bottom1 = h_ch_share[i].fillna(0).values + bottom1
ax0.set_ylim(0, 1)
if not discrete:
ax0.set_xlim(*xlim)
if xscale:
if isinstance(xscale, str):
ax0.set_xscale(xscale)
elif isinstance(xscale, dict):
ax0.set_xscale(**xscale)
else:
raise ValueError(
f"xscale must be str or dict, not {type(xscale)}")
if xmajorticks is not None:
ax0.set_xticks(xmajorticks)
ax0.set_xticklabels(xmajorticks)
if xminorticks is not None:
ax0.set_xticks(xminorticks, minor=True)
if x_discrete_labels is not None:
ax0.set_xticks(numpy.arange(len(x_discrete_labels)))
ax0.set_xticklabels(x_discrete_labels)
ax0.set_title('Modeled Shares')
ax1.set_ylim(0, 1)
if not discrete:
ax1.set_xlim(*xlim)
if xscale:
if isinstance(xscale, str):
ax1.set_xscale(xscale)
elif isinstance(xscale, dict):
ax1.set_xscale(**xscale)
else:
raise ValueError(
f"xscale must be str or dict, not {type(xscale)}")
if xmajorticks is not None:
ax1.set_xticks(xmajorticks)
ax1.set_xticklabels(xmajorticks)
if xminorticks is not None:
ax1.set_xticks(xminorticks, minor=True)
if x_discrete_labels is not None:
ax1.set_xticks(numpy.arange(len(x_discrete_labels)))
ax1.set_xticklabels(x_discrete_labels)
ax1.set_title('Observed Shares')
if xlabel:
ax0.set_xlabel(xlabel)
ax1.set_xlabel(xlabel)
fig.legend(loc='center right', )
# fig.tight_layout(pad=0.5)
if format == 'svg':
result = plot_as_svg_xhtml(fig, **kwargs)
fig.clf()
plt.close(fig)
elif format == 'png':
from .png import make_png
result = make_png(fig, **kwargs)
fig.clf()
plt.close(fig)
else:
result = fig
else:
fig, axes = plt.subplots(len(h_pr_share.columns), 1, figsize=figsize)
shift = 0.4 if discrete else 0
for n, i in enumerate(h_pr_share.columns):
x_, y_ = pseudo_bar_data(bins - shift,
h_pr_share[i],
gap=0.2 if discrete else 0)
axes[n].plot(x_, y_, label='Modeled' if n == 0 else None, lw=1.5)
x_ch_, y_ch_ = pseudo_bar_data(bins - shift,
h_ch_share[i],
gap=0.2 if discrete else 0)
axes[n].fill_between(
x_ch_,
y_ch_,
label='Observed' if n == 0 else None,
step=None,
facecolor='#ffbe4d',
edgecolor='#ffa200',
lw=1.5,
)
if not discrete:
axes[n].set_xlim(*xlim)
if xscale:
if isinstance(xscale, str):
axes[n].set_xscale(xscale)
elif isinstance(xscale, dict):
axes[n].set_xscale(**xscale)
else:
raise ValueError(
f"xscale must be str or dict, not {type(xscale)}")
if xmajorticks is not None:
axes[n].set_xticks(xmajorticks)
axes[n].set_xticklabels(xmajorticks)
if xminorticks is not None:
axes[n].set_xticks(xminorticks, minor=True)
if x_discrete_labels is not None:
axes[n].set_xticks(numpy.arange(len(x_discrete_labels)))
axes[n].set_xticklabels(x_discrete_labels)
axes[n].set_ylabel(i)
# axes[n].legend(
# # loc='center right',
# )
legnd = axes[0].legend(loc='lower center',
ncol=2,
borderaxespad=0,
bbox_to_anchor=(0.5, 1.08))
if xlabel:
axes[-1].set_xlabel(xlabel)
#fig.tight_layout(pad=0.5)
if format == 'svg':
result = plot_as_svg_xhtml(fig,
bbox_extra_artists=[legnd],
**kwargs)
fig.clf()
plt.close(fig)
elif format == 'png':
from .png import make_png
result = make_png(fig, **kwargs)
fig.clf()
plt.close(fig)
else:
result = fig
return result
