| 1967 | |
| 1968 | @derived_from(np) |
| 1969 | def compress(condition, a, axis=None): |
| 1970 | if not is_arraylike(condition): |
| 1971 | # Allow `condition` to be anything array-like, otherwise ensure `condition` |
| 1972 | # is a numpy array. |
| 1973 | condition = np.asarray(condition) |
| 1974 | condition = condition.astype(bool) |
| 1975 | a = asarray(a) |
| 1976 | |
| 1977 | if condition.ndim != 1: |
| 1978 | raise ValueError("Condition must be one dimensional") |
| 1979 | |
| 1980 | if axis is None: |
| 1981 | a = a.ravel() |
| 1982 | axis = 0 |
| 1983 | axis = validate_axis(axis, a.ndim) |
| 1984 | |
| 1985 | # Treat `condition` as filled with `False` (if it is too short) |
| 1986 | a = a[ |
| 1987 | tuple( |
| 1988 | slice(None, len(condition)) if i == axis else slice(None) |
| 1989 | for i in range(a.ndim) |
| 1990 | ) |
| 1991 | ] |
| 1992 | |
| 1993 | # Use `condition` to select along 1 dimension |
| 1994 | a = a[tuple(condition if i == axis else slice(None) for i in range(a.ndim))] |
| 1995 | |
| 1996 | return a |
| 1997 | |
| 1998 | |
| 1999 | @derived_from(np) |