cosmolopy.utils

1 """Some utilities used by various CosmoloPy modules. 2 """ 3 import warnings 4 import math 5 import pickle 6 7 import numpy 8 import scipy 9 import scipy.integrate 10 import scipy.interpolate 11 12 import cosmolopy.distance as cd 13 import cosmolopy.constants as cc 14 from saveable import Saveable

15 16 -class AgeSpacedRedshift(Saveable):

17 """Set up uniform time array and corresponding redshift array. 18 """

19 - def __init__(self, z_min, z_max, dt_yr=2e6, **cosmo):

20 self.z_min = z_min 21 self.z_max = z_max 22 self.dt_yr = dt_yr 23 self.cosmo = cosmo 24 25 self.agefunc, self.redshiftfunc, e_f, e_t = \ 26 cd.quick_age_function(zmax=1.1 * z_max, 27 zmin=z_min-0.05, 28 zstep=0.01, 29 logspacing=True, 30 return_inverse=True, 31 **cosmo) 32 self.tmax = self.agefunc(z_min) 33 self.tmin = self.agefunc(z_max) 34 self.dt = self.dt_yr * cc.yr_s 35 self.t = numpy.arange(self.tmin, self.tmax + 1.01 * self.dt, self.dt) 36 self.t_yr = self.t / cc.yr_s 37 self.z = self.redshiftfunc(self.t) 38 print " Using %i points in t, dt = %.3g yr." % (len(self.t_yr), 39 self.dt_yr)

40

41 - def age_Gyr(self, z):

42 return self.agefunc(z)/cc.yr_s/1e9

43 - def __setstate__(self, dict):

44 """Unpickle.""" 45 self.__dict__.update(dict) 46 self.agefunc, self.redshiftfunc, e_f, e_t = \ 47 cd.quick_age_function(zmax=1.1 * self.z_max, 48 zmin=self.z_min-0.05, 49 zstep=0.01, 50 logspacing=True, 51 return_inverse=True, 52 **self.cosmo)

53

54 55 -class Extrapolate1d:

56 """Interpolate/Extrapolate 1d data. 57 """ 58

59 - def linear_coefficients(self, x=None, y=None, slope=None, 60 match_index=0):

61 if x is None: 62 x = self.x 63 if y is None: 64 y=self.y 65 if slope is None: 66 slope = (y[0]-y[-1]) / (x[0] - x[-1]) 67 intercept = y[match_index] - slope * x[match_index] 68 return slope, intercept

69

70 - def __init__(self, x, y, 71 bounds_behavior=['extrapolate', 'extrapolate'], 72 slopes=[None, None], 73 npoints = [2, 2], 74 **interpargs):

75 """ 76 77 Parameters 78 ---------- 79 80 x, y: 81 82 sequences of data. Will be sorted by x value before use. 83 84 bound_behavior: 85 86 length-2 sequence specifying behavior below the lower and 87 above the upper boungs of the data, respectively. Each 88 element can be 'extrapolate', 'constant', or a numerical 89 value. 90 91 npoints: 92 93 Linear extrapolation uses the slope between x[0] and 94 x[npoints-1] or x[-npoints] and x[-1]. Note: this is not a 95 linear fit over that range. It Ignores points within the 96 interval 97 98 interpargs: 99 100 Extra keywords passed to scipy.interpolate.interp1d. 101 102 """ 103 order = numpy.argsort(numpy.nan_to_num(x)) 104 self.x = x[order] 105 self.y = y[order] 106 107 self.bounds_behavior = bounds_behavior 108 109 self._interpfunc = scipy.interpolate.interp1d(self.x, self.y, 110 **interpargs) 111 112 if self.bounds_behavior[1] == 'constant': 113 self._exfuncHigh = lambda x1: self.y[-1] 114 elif self.bounds_behavior[1] == 'extrapolate': 115 n1 = npoints[1] 116 highSlope, highIntercept = self.linear_coefficients(self.x[-n1:], 117 self.y[-n1:], 118 slope=slopes[1], 119 match_index=-1) 120 self._exfuncHigh = lambda x1: x1 * highSlope + highIntercept 121 self.highSlope = highSlope 122 self.highIntercept = highIntercept 123 else: 124 self._exfuncHigh = lambda x1: self.bounds_behavior[1] 125 126 if self.bounds_behavior[0] == 'constant': 127 self._exfuncLow = lambda x1: self.y[0] 128 elif self.bounds_behavior[0] == 'extrapolate': 129 n0 = npoints[0] 130 lowSlope, lowIntercept = self.linear_coefficients(self.x[:n0], 131 self.y[:n0], 132 slope=slopes[0], 133 match_index=0) 134 self._exfuncLow = lambda x1: x1 * lowSlope + lowIntercept 135 self.lowSlope = lowSlope 136 self.lowIntercept = lowIntercept 137 else: 138 self._exfuncLow = lambda x1: self.bounds_behavior[0]

139

140 - def extrap_string(self):

141 extstr = "" 142 if hasattr(self, 'lowSlope'): 143 extstr += "y = %g x + %g for x <= %g" % (self.lowSlope, 144 self.lowIntercept, 145 self.x[0]) 146 if hasattr(self, 'highSlope'): 147 if hasattr(self, 'lowSlope'): 148 extstr += "\n" 149 extstr += "y = %g x + %g for x >= %g" % (self.highSlope, 150 self.highIntercept, 151 self.x[-1]) 152 return extstr

153

154 - def __call__(self, x1):

155 if numpy.isscalar(x1) or x1.shape==(): 156 if x1 <= self.x[0]: 157 return self._exfuncLow(x1) 158 elif x1 >= self.x[-1]: 159 return self._exfuncHigh(x1) 160 else: 161 return self._interpfunc(x1) 162 lowmask = x1 <= self.x[0] 163 highmask = x1 >= self.x[-1] 164 inmask = numpy.logical_not(numpy.logical_or(lowmask,highmask)) 165 if numpy.all(inmask): 166 return self._interpfunc(x1) 167 168 y1 = numpy.empty(x1.shape) 169 y1[inmask] = self._interpfunc(x1[inmask]) 170 y1[lowmask] = self._exfuncLow(x1[lowmask]) 171 y1[highmask] = self._exfuncHigh(x1[highmask]) 172 173 return y1

174

175 -class PiecewisePowerlaw(object):

176 """A piecewise powerlaw function. 177 178 You can specify the intervals and power indices, and this class 179 will figure out the coefficients needed to make the function 180 continuous and normalized to unit integral. 181 182 Notes 183 ----- 184 185 Intervals are defined by an array l 186 187 Powerlaw indicies by and array p 188 189 a_n are the coefficients. 190 191 f(x) = a_n x^{p_n} for l_{n-1} <= x < l_n 192 193 Recursion relation for continuity: 194 195 a_n = a_{n-1} l_n^{p_{n-1} - p_n} 196 197 Integral of a piece: 198 199 I_n = a_n p_n (l_{n+1}^{p_n - 1} - l_n^{p_n - 1}) 200 201 Total integral: 202 203 I_tot = Sum_0^N I_n 204 205 """ 206

207 - def __init__(self, limits, powers, 208 coefficients=None, 209 externalval=0.0, 210 norm=True):

211 """Defined a piecewise powerlaw. 212 213 If coefficients is None then the coefficients are determined 214 by requiring the function to be continuous and normalized to 215 an integral of one. 216 217 The function is composed of N powerlaws, where N = len(powers). 218 219 len(limits) must be one greated than len(powers) 220 221 Parameters 222 ---------- 223 224 limits: array (length n+1) 225 boundaries of the specified powerlaws. Must be one greater in 226 length than coefficents and powers. Specify -numpy.infty for 227 the first limit or numpy.infty for the last limit for 228 unbounded powerlaws. 229 230 coefficients: optional array (length n) 231 values of the coefficient a_i 232 233 powers: array (length n) 234 values of the powerlaw indices p_i 235 236 externalval: scalar 237 Value to return outside the defined domain. None 238 correspons to 'NaN'. 239 240 norm: boolean 241 Whether to normalize the integral of the function over the 242 defined domain to unity. 243 244 The resulting function takes a single, one-dimensional array of 245 values on which to operate. 246 247 """ 248 249 limits = numpy.atleast_1d(limits) 250 powers = numpy.atleast_1d(powers) 251 252 if not len(limits) == len(powers)+1: 253 raise ValueError("limits must be one longer than powers.") 254 255 if coefficients is None: 256 coefficients = numpy.ones(len(powers)) 257 258 # Leaving a_0 = 1, apply the recurence relation. 259 for n in range(1,len(powers)): 260 coefficients[n] = (coefficients[n-1] * 261 limits[n]**(powers[n-1] - powers[n])) 262 else: 263 coefficients = numpy.atleast_1d(coefficients) 264 if not len(coefficients) == len(powers): 265 raise ValueError("coefficients and powers must be"+ 266 " the same length.") 267 268 # Find the integral of each piece. 269 integrals = ((coefficients / (powers + 1.)) * 270 (limits[1:]**(powers + 1.) - 271 limits[:-1]**(powers + 1.))) 272 if norm: 273 # The total integral over the function. 274 integralTot = numpy.sum(integrals) 275 276 coefficients = coefficients / integralTot 277 integrals = integrals / integralTot 278 279 for array in [limits, coefficients, powers]: 280 if array.ndim > 1: 281 raise ValueError("arguments must be a 1D arrays or scalars.") 282 self._integrals = integrals 283 self._limits = limits.reshape((-1,1)) 284 self._coefficients = coefficients.reshape((-1,1)) 285 self._powers = powers.reshape((-1,1)) 286 self._externalval = externalval

287

288 - def __call__(self, x):

289 """Evaluate the powerlaw at values x. 290 """ 291 x = numpy.atleast_1d(x) 292 if x.ndim > 1: 293 raise ValueError("argument must be a 1D array or scalar.") 294 y = numpy.sum((self._coefficients * x**self._powers) * 295 (x >= self._limits[0:-1]) * (x < self._limits[1:]), 296 axis=0) 297 y[x < self._limits[0]] = self._externalval 298 y[x >= self._limits[-1]] = self._externalval 299 return y

300

301 - def integrate(self, low, high, weight_power=None):

302 """Integrate the function from low to high. 303 304 Optionally weight the integral by x^weight_power. 305 306 """ 307 limits = self._limits.flatten() 308 coefficients = self._coefficients.flatten() 309 powers = self._powers.flatten() 310 if weight_power is not None: 311 powers += weight_power 312 # Integral of each piece over its domain. 313 integrals = ((coefficients / (powers + 1.)) * 314 (limits[1:]**(powers + 1.) - 315 limits[:-1]**(powers + 1.))) 316 else: 317 integrals = self._integrals 318 319 pairs = numpy.broadcast(low, high) 320 integral = numpy.empty(pairs.shape) 321 for (i, (x0, x1)) in enumerate(pairs): 322 # Sort the integral limits. 323 x0, x1 = list(numpy.sort([x0,x1])) 324 325 # Select the pieces contained entirely in the interval. 326 mask = numpy.logical_and(x0 < limits[:-1], 327 x1 >= limits[1:]).flatten() 328 indices = numpy.where(mask) 329 if not numpy.any(mask): 330 integral.flat[i] = 0 331 332 # If the interval is outside the domain. 333 if x0 > limits[-1] or x1 < limits[0]: 334 integral.flat[i] = 0 335 continue 336 337 # Find out if any piece contains the entire interval: 338 containedmask = numpy.logical_and(x0 >= limits[:-1], 339 x1 < limits[1:]) 340 # Three possibilites: 341 if numpy.any(containedmask): 342 # The interval is contained in a single segment. 343 index = numpy.where(containedmask)[0][0] 344 integral.flat[i] = ((coefficients[index] / 345 (powers[index] + 1.)) * 346 (x1**(powers[index] + 1.) - 347 x0**(powers[index] + 1.))) 348 continue 349 elif x1 >= limits[0] and x1 < limits[1]: 350 # x1 is in the first segment. 351 highi = 0 352 lowi = -1 353 elif x0 < limits[-1] and x0 >= limits[-2]: 354 # x0 is in the last segment: 355 lowi = len(limits) - 2 356 highi = len(limits) 357 else: 358 # We must be spanning the division between a pair of pieces. 359 lowi = numpy.max(numpy.where(x0 >= limits[:-1])) 360 highi = numpy.min(numpy.where(x1 < limits[1:])) 361 insideintegral = 0 362 else: 363 # Add up the integrals of the pieces totally inside the interval. 364 insideintegral = numpy.sum(integrals[indices]) 365 366 lowi = numpy.min(indices) - 1 367 highi = numpy.max(indices) + 1 368 369 # Check that the integral limits are inside our domain. 370 if x0 < limits[0] or lowi < 0: 371 lowintegral = 0. 372 else: 373 lowintegral = ((coefficients[lowi] / (powers[lowi] + 1.)) * 374 (limits[lowi + 1]**(powers[lowi] + 1.) - 375 x0**(powers[lowi] + 1.))) 376 if x1 > limits[-1] or highi > len(coefficients) - 1: 377 highintegral = 0. 378 else: 379 highintegral = ((coefficients[highi] / (powers[highi] + 1.)) * 380 (x1**(powers[highi] + 1.) - 381 limits[highi]**(powers[highi] + 1.))) 382 integral.flat[i] = highintegral + insideintegral + lowintegral 383 return integral

384

385 -def ccumulate(function, x, max=None, **kwargs):

386 """Integrate a function from x to max, where x can be an array. 387 388 Parameters 389 ---------- 390 391 function: callable 392 393 x: array-like 394 395 max: float 396 defaults to max(x) 397 398 Notes 399 ----- 400 401 This can be used to find the complementary cumulative distribution 402 function (CCDF) given the probability distribution function (PDF). 403 404 Unlike integrate_piecewise, the x values don't have to be in 405 order, though a warning will be issued if any are greater than 406 max, if max is specified. 407 """ 408 409 x = numpy.atleast_1d(x) 410 # Sort the x values and get a list of non-NaN values in order. 411 order = numpy.argsort(numpy.nan_to_num(x)) 412 bad_mask = numpy.isnan(x[order]) 413 bad_order = order[bad_mask] 414 good_order = order[numpy.logical_not(bad_mask)] 415 x0 = x[good_order] 416 417 if len(x0) == 1: 418 integral = [0] 419 else: 420 integral = integrate_piecewise(function, x0, **kwargs) 421 422 # If we have a max value, we need to add the integral from max(x0) to max. 423 postintegral = 0. 424 if max is not None: 425 sign = 0. 426 if max >= x0[-1]: 427 sign = 1. 428 x1 = [x0[-1], max] 429 elif(max < x0[-1]): 430 warnings.warn("max %s is less than maximum x %s" % (max, x0[-1])) 431 sign = -1. 432 x1 = [max, x0[-1]] 433 postintegral += sign * integrate_piecewise(function, x1, **kwargs)[-1] 434 435 # Reverse the direction of the integration. 436 cintegral = postintegral + integral[-1] - integral 437 438 # Put the integral values back in the right order. 439 ordintegral = numpy.empty(len(x)) 440 ordintegral[good_order] = cintegral 441 ordintegral[bad_order] = numpy.nan 442 return ordintegral

443

444 -def integrate_piecewise(function, x, method='romberg', return_pieces=False, 445 **kwargs):

446 """Integrate function and return the integral at a sequence of points. 447 448 Useful when you want to efficiently calculate a cumulative integral. 449 450 Also useful for piecewise-defined functions where discontinuities 451 or critical points cause quadrature routines to complain or become 452 inaccurate. 453 454 Integration methods available are: quad, romberg. 455 456 Parameters 457 ---------- 458 function : callable 459 User defined function. Should take a single vector argument 460 and return q vector of the same shape. 461 462 x : array_like 463 Array of points at which to evaluate the integral. 464 465 method : str, optional 466 Name of the method to use to integrate each segment. 'quad' or 467 'romberg'. 468 469 return_pieces : bool, optional 470 Return the individual segments rather than the sum of all 471 preceding segments. 472 473 Returns 474 ------- 475 integral : ndarray 476 The value of the integral at each x. The first value is always zero. 477 """ 478 479 x = numpy.asarray(x) 480 if numpy.any(x[1:] - x[:-1] < 0): 481 raise ValueError, "Array x must increase monotonically." 482 if numpy.any(numpy.isnan(x)): 483 raise ValueError, "Array x must not include NaN values." 484 integral_list = [0.0] 485 if method is None: 486 method = 'quad' 487 if method=='quad': 488 args = {'limit':200} 489 args.update(kwargs) 490 for i in xrange(1, len(x)): 491 a, b = x[i-1], x[i] 492 integral, error = scipy.integrate.quad(function, a, b, 493 **args) 494 integral_list.append(integral) 495 elif method=='romberg': 496 args = {'divmax':100, 'vec_func':True} 497 args.update(kwargs) 498 for i in xrange(1, len(x)): 499 a, b = x[i-1], x[i] 500 integral = scipy.integrate.romberg(function, a, b, 501 **args) 502 integral_list.append(integral) 503 else: 504 raise ValueError, "Method '%s' unknown." % method 505 506 integrals = numpy.asarray(integral_list) 507 if return_pieces: 508 return integrals 509 510 integral = numpy.cumsum(numpy.nan_to_num(integrals)) 511 return integral

512

513 @numpy.vectorize 514 -def _vecquad(function, low, high, **kwargs):

515 integral, error = scipy.integrate.quad(function, 516 low, 517 high, 518 **kwargs) 519 return integral, error

520

521 -def vecquad(function, low, high, **kwargs):

522 """Integrate a function from low to high (vectorized). 523 524 Vectorized convenience function. 525 526 """ 527 return _vecquad(function,low,high, **kwargs)

528

529 @numpy.vectorize 530 -def _logquad(function, low, high, **kwargs):

531 # Transform the function to log space. 532 def func_dlnx (lnx): 533 x = numpy.exp(lnx) 534 return function(x) * x

535 integral, error = scipy.integrate.quad(func_dlnx, 536 math.log(low), 537 math.log(high), 538 **kwargs) 539 return integral, error 540

541 -def logquad(function, low, high, **kwargs):

542 """Integrate a function from low to high using a log transform (vectorized). 543 544 The log transform is applied to the variable over which the 545 integration is being performed. 546 547 """ 548 return _logquad(function, low, high, **kwargs)

549

550 -class Normalize:

551 """A decorator that normalizes a function. 552 553 Only works for functions of a single variable. 554 555 The new function is normalized over the interval from min to max, 556 i.e. the integral of the new function from low to high is one. 557 """ 558

559 - def __init__(self, min, max, quiet=False, **kwargs):

560 self.min = min 561 self.max = max 562 self.quiet = quiet 563 self.kwargs = kwargs

564

565 - def __call__(self, function):

566 integral = logquad(function, self.min, self.max, **self.kwargs)[0] 567 newfunction = lambda x: function(x)/integral 568 569 if not self.quiet: 570 print "Normalization factor for %s is %.3g" % (function.__name__, 571 1./integral) 572 # inspired by 573 # http://wiki.python.org/moin/PythonDecoratorLibrary#DifferentDecoratorForms 574 575 newfunction.__name__ = function.__name__ 576 newfunction.__dict__.update(function.__dict__) 577 newfunction.__doc__ = function.__doc__ 578 newfunction.min = self.min 579 newfunction.max = self.max 580 return newfunction

581

Source Code for Module cosmolopy.utils