1# Licensed under a 3-clause BSD style license - see LICENSE.rst 2""" 3Convenience functions for `astropy.cosmology`. 4""" 5 6import warnings 7 8import numpy as np 9 10from astropy.units import Quantity 11from astropy.utils.exceptions import AstropyUserWarning 12 13from . import units as cu 14from .core import CosmologyError 15 16__all__ = ['z_at_value'] 17 18__doctest_requires__ = {'*': ['scipy']} 19 20 21def _z_at_scalar_value(func, fval, zmin=1e-8, zmax=1000, ztol=1e-8, maxfun=500, 22 method='Brent', bracket=None, verbose=False): 23 """ 24 Find the redshift ``z`` at which ``func(z) = fval``. 25 See :func:`astropy.cosmology.funcs.z_at_value`. 26 """ 27 from scipy.optimize import minimize_scalar 28 29 opt = {'maxiter': maxfun} 30 # Assume custom methods support the same options as default; otherwise user 31 # will see warnings. 32 if str(method).lower() == 'bounded': 33 opt['xatol'] = ztol 34 if bracket is not None: 35 warnings.warn(f"Option 'bracket' is ignored by method {method}.") 36 bracket = None 37 else: 38 opt['xtol'] = ztol 39 40 # fval falling inside the interval of bracketing function values does not 41 # guarantee it has a unique solution, but for Standard Cosmological 42 # quantities normally should (being monotonic or having a single extremum). 43 # In these cases keep solver from returning solutions outside of bracket. 44 fval_zmin, fval_zmax = func(zmin), func(zmax) 45 nobracket = False 46 if np.sign(fval - fval_zmin) != np.sign(fval_zmax - fval): 47 if bracket is None: 48 nobracket = True 49 else: 50 fval_brac = func(np.asanyarray(bracket)) 51 if np.sign(fval - fval_brac[0]) != np.sign(fval_brac[-1] - fval): 52 nobracket = True 53 else: 54 zmin, zmax = bracket[0], bracket[-1] 55 fval_zmin, fval_zmax = fval_brac[[0, -1]] 56 if nobracket: 57 warnings.warn(f"fval is not bracketed by func(zmin)={fval_zmin} and " 58 f"func(zmax)={fval_zmax}. This means either there is no " 59 "solution, or that there is more than one solution " 60 "between zmin and zmax satisfying fval = func(z).", 61 AstropyUserWarning) 62 63 if isinstance(fval_zmin, Quantity): 64 val = fval.to_value(fval_zmin.unit) 65 else: 66 val = fval 67 68 # 'Brent' and 'Golden' ignore `bounds`, force solution inside zlim 69 def f(z): 70 if z > zmax: 71 return 1.e300 * (1.0 + z - zmax) 72 elif z < zmin: 73 return 1.e300 * (1.0 + zmin - z) 74 elif isinstance(fval_zmin, Quantity): 75 return abs(func(z).value - val) 76 else: 77 return abs(func(z) - val) 78 79 res = minimize_scalar(f, method=method, bounds=(zmin, zmax), 80 bracket=bracket, options=opt) 81 82 # Scipy docs state that `OptimizeResult` always has 'status' and 'message' 83 # attributes, but only `_minimize_scalar_bounded()` seems to have really 84 # implemented them. 85 if not res.success: 86 warnings.warn(f"Solver returned {res.get('status')}: {res.get('message', 'Unsuccessful')}\n" 87 f"Precision {res.fun} reached after {res.nfev} function calls.", 88 AstropyUserWarning) 89 90 if verbose: 91 print(res) 92 93 if np.allclose(res.x, zmax): 94 raise CosmologyError( 95 f"Best guess z={res.x} is very close to the upper z limit {zmax}." 96 "\nTry re-running with a different zmax.") 97 elif np.allclose(res.x, zmin): 98 raise CosmologyError( 99 f"Best guess z={res.x} is very close to the lower z limit {zmin}." 100 "\nTry re-running with a different zmin.") 101 return res.x 102 103 104def z_at_value(func, fval, zmin=1e-8, zmax=1000, ztol=1e-8, maxfun=500, 105 method='Brent', bracket=None, verbose=False): 106 """Find the redshift ``z`` at which ``func(z) = fval``. 107 108 This finds the redshift at which one of the cosmology functions or 109 methods (for example Planck13.distmod) is equal to a known value. 110 111 .. warning:: 112 Make sure you understand the behavior of the function that you are 113 trying to invert! Depending on the cosmology, there may not be a 114 unique solution. For example, in the standard Lambda CDM cosmology, 115 there are two redshifts which give an angular diameter distance of 116 1500 Mpc, z ~ 0.7 and z ~ 3.8. To force ``z_at_value`` to find the 117 solution you are interested in, use the ``zmin`` and ``zmax`` keywords 118 to limit the search range (see the example below). 119 120 Parameters 121 ---------- 122 func : function or method 123 A function that takes a redshift as input. 124 125 fval : `~astropy.units.Quantity` 126 The (scalar or array) value of ``func(z)`` to recover. 127 128 zmin : float or array-like['dimensionless'] or quantity-like, optional 129 The lower search limit for ``z``. Beware of divergences 130 in some cosmological functions, such as distance moduli, 131 at z=0 (default 1e-8). 132 133 zmax : float or array-like['dimensionless'] or quantity-like, optional 134 The upper search limit for ``z`` (default 1000). 135 136 ztol : float or array-like['dimensionless'], optional 137 The relative error in ``z`` acceptable for convergence. 138 139 maxfun : int or array-like, optional 140 The maximum number of function evaluations allowed in the 141 optimization routine (default 500). 142 143 method : str or callable, optional 144 Type of solver to pass to the minimizer. The built-in options provided 145 by :func:`~scipy.optimize.minimize_scalar` are 'Brent' (default), 146 'Golden' and 'Bounded' with names case insensitive - see documentation 147 there for details. It also accepts a custom solver by passing any 148 user-provided callable object that meets the requirements listed 149 therein under the Notes on "Custom minimizers" - or in more detail in 150 :doc:`scipy:tutorial/optimize` - although their use is currently 151 untested. 152 153 .. versionadded:: 4.3 154 155 bracket : sequence or object array[sequence], optional 156 For methods 'Brent' and 'Golden', ``bracket`` defines the bracketing 157 interval and can either have three items (z1, z2, z3) so that 158 z1 < z2 < z3 and ``func(z2) < func (z1), func(z3)`` or two items z1 159 and z3 which are assumed to be a starting interval for a downhill 160 bracket search. For non-monotonic functions such as angular diameter 161 distance this may be used to start the search on the desired side of 162 the maximum, but see Examples below for usage notes. 163 164 .. versionadded:: 4.3 165 166 verbose : bool, optional 167 Print diagnostic output from solver (default `False`). 168 169 .. versionadded:: 4.3 170 171 Returns 172 ------- 173 z : `~astropy.units.Quantity` ['redshift'] 174 The redshift ``z`` satisfying ``zmin < z < zmax`` and ``func(z) = 175 fval`` within ``ztol``. Has units of cosmological redshift. 176 177 Warns 178 ----- 179 :class:`~astropy.utils.exceptions.AstropyUserWarning` 180 If ``fval`` is not bracketed by ``func(zmin)=fval(zmin)`` and 181 ``func(zmax)=fval(zmax)``. 182 183 If the solver was not successful. 184 185 Raises 186 ------ 187 :class:`astropy.cosmology.CosmologyError` 188 If the result is very close to either ``zmin`` or ``zmax``. 189 ValueError 190 If ``bracket`` is not an array nor a 2 (or 3) element sequence. 191 TypeError 192 If ``bracket`` is not an object array. 2 (or 3) element sequences will 193 be turned into object arrays, so this error should only occur if a 194 non-object array is used for ``bracket``. 195 196 Notes 197 ----- 198 This works for any arbitrary input cosmology, but is inefficient if you 199 want to invert a large number of values for the same cosmology. In this 200 case, it is faster to instead generate an array of values at many 201 closely-spaced redshifts that cover the relevant redshift range, and then 202 use interpolation to find the redshift at each value you are interested 203 in. For example, to efficiently find the redshifts corresponding to 10^6 204 values of the distance modulus in a Planck13 cosmology, you could do the 205 following: 206 207 >>> import astropy.units as u 208 >>> from astropy.cosmology import Planck13, z_at_value 209 210 Generate 10^6 distance moduli between 24 and 44 for which we 211 want to find the corresponding redshifts: 212 213 >>> Dvals = (24 + np.random.rand(1000000) * 20) * u.mag 214 215 Make a grid of distance moduli covering the redshift range we 216 need using 50 equally log-spaced values between zmin and 217 zmax. We use log spacing to adequately sample the steep part of 218 the curve at low distance moduli: 219 220 >>> zmin = z_at_value(Planck13.distmod, Dvals.min()) 221 >>> zmax = z_at_value(Planck13.distmod, Dvals.max()) 222 >>> zgrid = np.geomspace(zmin, zmax, 50) 223 >>> Dgrid = Planck13.distmod(zgrid) 224 225 Finally interpolate to find the redshift at each distance modulus: 226 227 >>> zvals = np.interp(Dvals.value, Dgrid.value, zgrid) 228 229 Examples 230 -------- 231 >>> import astropy.units as u 232 >>> from astropy.cosmology import Planck13, Planck18, z_at_value 233 234 The age and lookback time are monotonic with redshift, and so a 235 unique solution can be found: 236 237 >>> z_at_value(Planck13.age, 2 * u.Gyr) # doctest: +FLOAT_CMP 238 <Quantity 3.19812268 redshift> 239 240 The angular diameter is not monotonic however, and there are two 241 redshifts that give a value of 1500 Mpc. You can use the zmin and 242 zmax keywords to find the one you are interested in: 243 244 >>> z_at_value(Planck18.angular_diameter_distance, 245 ... 1500 * u.Mpc, zmax=1.5) # doctest: +FLOAT_CMP 246 <Quantity 0.68044452 redshift> 247 >>> z_at_value(Planck18.angular_diameter_distance, 248 ... 1500 * u.Mpc, zmin=2.5) # doctest: +FLOAT_CMP 249 <Quantity 3.7823268 redshift> 250 251 Alternatively the ``bracket`` option may be used to initialize the 252 function solver on a desired region, but one should be aware that this 253 does not guarantee it will remain close to this starting bracket. 254 For the example of angular diameter distance, which has a maximum near 255 a redshift of 1.6 in this cosmology, defining a bracket on either side 256 of this maximum will often return a solution on the same side: 257 258 >>> z_at_value(Planck18.angular_diameter_distance, 259 ... 1500 * u.Mpc, bracket=(1.0, 1.2)) # doctest: +FLOAT_CMP +IGNORE_WARNINGS 260 <Quantity 0.68044452 redshift> 261 262 But this is not ascertained especially if the bracket is chosen too wide 263 and/or too close to the turning point: 264 265 >>> z_at_value(Planck18.angular_diameter_distance, 266 ... 1500 * u.Mpc, bracket=(0.1, 1.5)) # doctest: +SKIP 267 <Quantity 3.7823268 redshift> # doctest: +SKIP 268 269 Likewise, even for the same minimizer and same starting conditions different 270 results can be found depending on architecture or library versions: 271 272 >>> z_at_value(Planck18.angular_diameter_distance, 273 ... 1500 * u.Mpc, bracket=(2.0, 2.5)) # doctest: +SKIP 274 <Quantity 3.7823268 redshift> # doctest: +SKIP 275 276 >>> z_at_value(Planck18.angular_diameter_distance, 277 ... 1500 * u.Mpc, bracket=(2.0, 2.5)) # doctest: +SKIP 278 <Quantity 0.68044452 redshift> # doctest: +SKIP 279 280 It is therefore generally safer to use the 3-parameter variant to ensure 281 the solution stays within the bracketing limits: 282 283 >>> z_at_value(Planck18.angular_diameter_distance, 1500 * u.Mpc, 284 ... bracket=(0.1, 1.0, 1.5)) # doctest: +FLOAT_CMP 285 <Quantity 0.68044452 redshift> 286 287 Also note that the luminosity distance and distance modulus (two 288 other commonly inverted quantities) are monotonic in flat and open 289 universes, but not in closed universes. 290 291 All the arguments except ``func``, ``method`` and ``verbose`` accept array 292 inputs. This does NOT use interpolation tables or any method to speed up 293 evaluations, rather providing a convenient means to broadcast arguments 294 over an element-wise scalar evaluation. 295 296 The most common use case for non-scalar input is to evaluate 'func' for an 297 array of ``fval``: 298 299 >>> z_at_value(Planck13.age, [2, 7] * u.Gyr) # doctest: +FLOAT_CMP 300 <Quantity [3.19812061, 0.75620443] redshift> 301 302 ``fval`` can be any shape: 303 304 >>> z_at_value(Planck13.age, [[2, 7], [1, 3]]*u.Gyr) # doctest: +FLOAT_CMP 305 <Quantity [[3.19812061, 0.75620443], 306 [5.67661227, 2.19131955]] redshift> 307 308 Other arguments can be arrays. For non-monotic functions -- for example, 309 the angular diameter distance -- this can be useful to find all solutions. 310 311 >>> z_at_value(Planck13.angular_diameter_distance, 1500 * u.Mpc, 312 ... zmin=[0, 2.5], zmax=[2, 4]) # doctest: +FLOAT_CMP 313 <Quantity [0.68127747, 3.79149062] redshift> 314 315 The ``bracket`` argument can likewise be be an array. However, since 316 bracket must already be a sequence (or None), it MUST be given as an 317 object `numpy.ndarray`. Importantly, the depth of the array must be such 318 that each bracket subsequence is an object. Errors or unexpected results 319 will happen otherwise. A convenient means to ensure the right depth is by 320 including a length-0 tuple as a bracket and then truncating the object 321 array to remove the placeholder. This can be seen in the following 322 example: 323 324 >>> bracket=np.array([(1.0, 1.2),(2.0, 2.5), ()], dtype=object)[:-1] 325 >>> z_at_value(Planck18.angular_diameter_distance, 1500 * u.Mpc, 326 ... bracket=bracket) # doctest: +SKIP 327 <Quantity [0.68044452, 3.7823268] redshift> 328 """ 329 # `fval` can be a Quantity, which isn't (yet) compatible w/ `numpy.nditer` 330 # so we strip it of units for broadcasting and restore the units when 331 # passing the elements to `_z_at_scalar_value`. 332 fval = np.asanyarray(fval) 333 unit = getattr(fval, 'unit', 1) # can be unitless 334 zmin = Quantity(zmin, cu.redshift).value # must be unitless 335 zmax = Quantity(zmax, cu.redshift).value 336 337 # bracket must be an object array (assumed to be correct) or a 'scalar' 338 # bracket: 2 or 3 elt sequence 339 if not isinstance(bracket, np.ndarray): # 'scalar' bracket 340 if bracket is not None and len(bracket) not in (2, 3): 341 raise ValueError("`bracket` is not an array " 342 "nor a 2 (or 3) element sequence.") 343 else: # munge bracket into a 1-elt object array 344 bracket = np.array([bracket, ()], dtype=object)[:1].squeeze() 345 if bracket.dtype != np.object_: 346 raise TypeError(f"`bracket` has dtype {bracket.dtype}, not 'O'") 347 348 # make multi-dimensional iterator for all but `method`, `verbose` 349 with np.nditer( 350 [fval, zmin, zmax, ztol, maxfun, bracket, None], 351 flags = ['refs_ok'], 352 op_flags = [*[['readonly']] * 6, # ← inputs output ↓ 353 ['writeonly', 'allocate', 'no_subtype']], 354 op_dtypes = (*(None,)*6, fval.dtype), 355 casting="no", 356 ) as it: 357 for fv, zmn, zmx, zt, mfe, bkt, zs in it: # ← eltwise unpack & eval ↓ 358 zs[...] = _z_at_scalar_value(func, fv * unit, zmin=zmn, zmax=zmx, 359 ztol=zt, maxfun=mfe, bracket=bkt.item(), 360 # not broadcasted 361 method=method, verbose=verbose) 362 # since bracket is an object array, the output will be too, so it is 363 # cast to the same type as the function value. 364 result = it.operands[-1] # zs 365 366 return result << cu.redshift 367