import numpy as np def div0(a, b, fill=0): """a / b, divide by 0 -> `fill` div0( [-1, 0, 1], 0, fill=np.nan) -> [nan nan nan] div0( 1, 0, fill=np.inf ) -> inf """ with np.errstate(divide="ignore", invalid="ignore"): c = np.true_divide(a, b) if np.isscalar(c): return c if np.isfinite(c) else fill else: c[~np.isfinite(c)] = fill return c