utils.py 11.5 KB
Newer Older
Yannick Roehlly's avatar
Yannick Roehlly committed
1
# -*- coding: utf-8 -*-
2
3
# Copyright (C) 2012, 2013 Centre de données Astrophysiques de Marseille
# Licensed under the CeCILL-v2 licence - see Licence_CeCILL_V2-en.txt
Yannick Roehlly's avatar
Yannick Roehlly committed
4
# Authors: Yannick Roehlly, Médéric Boquien
Yannick Roehlly's avatar
Yannick Roehlly committed
5
6

import numpy as np
7
from scipy.constants import c, pi
Yannick Roehlly's avatar
Yannick Roehlly committed
8

9
10
# Cache of dx for integrate(y,dx) done by flux_trapz
dx_cache = {}
11
best_grid_cache = {}
12

Yannick Roehlly's avatar
Yannick Roehlly committed
13
14
15
16

def lambda_to_nu(wavelength):
    """Convert wavelength (nm) to frequency (Hz)

Yannick Roehlly's avatar
Yannick Roehlly committed
17
    Parameters
Yannick Roehlly's avatar
Yannick Roehlly committed
18
    ----------
19
    wavelength: float or array of floats
Yannick Roehlly's avatar
Yannick Roehlly committed
20
21
22
23
        The wavelength(s) in nm.

    Returns
    -------
24
    nu: float or array of floats
Yannick Roehlly's avatar
Yannick Roehlly committed
25
26
27
28
29
30
31
32
33
        The frequency(ies) in Hz.

    """
    return c / (wavelength * 1.e-9)


def nu_to_lambda(frequency):
    """Convert frequency (Hz) to wavelength (nm)

Yannick Roehlly's avatar
Yannick Roehlly committed
34
    Parameters
Yannick Roehlly's avatar
Yannick Roehlly committed
35
    ----------
36
    frequency: float or numpy.array of floats
Yannick Roehlly's avatar
Yannick Roehlly committed
37
38
39
40
        The frequency(ies) in Hz.

    Returns
    -------
41
    wavelength: float or numpy.array of floats
Yannick Roehlly's avatar
Yannick Roehlly committed
42
43
44
45
46
        The wavelength(s) in nm.

    """
    return 1.e-9 * c / frequency

Médéric Boquien's avatar
Médéric Boquien committed
47

48
def luminosity_to_flux(luminosity, dist):
Yannick Roehlly's avatar
Yannick Roehlly committed
49
50
51
52
53
    """
    Convert a luminosity (or luminosity density) to a flux (or flux density).

    F = L / (4πDl2)

Yannick Roehlly's avatar
Yannick Roehlly committed
54
    Parameters
Yannick Roehlly's avatar
Yannick Roehlly committed
55
    ----------
56
    luminosity: float or array of floats
Yannick Roehlly's avatar
Yannick Roehlly committed
57
        Luminosity (typically in W) or luminosity density (W/nm or W/Hz).
58
59
    dist: float
        Luminosity distance of the object in metres
Yannick Roehlly's avatar
Yannick Roehlly committed
60
61
62

    Returns
    -------
63
    flux: float or array of floats
Yannick Roehlly's avatar
Yannick Roehlly committed
64
65
66
67
        The flux (typically in W/m²) of flux density (W/m²/nm or W/m²/Hz).

    """

68
    flux = luminosity / (4. * pi * dist * dist)
69

70
    return flux
Yannick Roehlly's avatar
Yannick Roehlly committed
71
72


Yannick Roehlly's avatar
Yannick Roehlly committed
73
def lambda_flambda_to_fnu(wavelength, flambda):
Yannick Roehlly's avatar
Yannick Roehlly committed
74
75
76
    """
    Convert a Fλ vs λ spectrum to Fν vs λ

Yannick Roehlly's avatar
Yannick Roehlly committed
77
    Parameters
Yannick Roehlly's avatar
Yannick Roehlly committed
78
    ----------
79
    wavelength: list-like of floats
Yannick Roehlly's avatar
Yannick Roehlly committed
80
        The wavelengths in nm.
81
    flambda: list-like of floats
Yannick Roehlly's avatar
Yannick Roehlly committed
82
        Fλ flux density in W/m²/nm (or Lλ luminosity density in W/nm).
Yannick Roehlly's avatar
Yannick Roehlly committed
83
84
85

    Returns
    -------
86
    fnu: array of floats
Yannick Roehlly's avatar
Yannick Roehlly committed
87
88
        The Fν flux density in mJy (or the Lν luminosity density in
        1.e-29 W/Hz).
Yannick Roehlly's avatar
Yannick Roehlly committed
89
90

    """
Yannick Roehlly's avatar
Yannick Roehlly committed
91
92
93
94
    wavelength = np.array(wavelength, dtype=float)
    flambda = np.array(flambda, dtype=float)

    # Factor 1e+29 is to switch from W/m²/Hz to mJy
Yannick Roehlly's avatar
Yannick Roehlly committed
95
    # Factor 1e-9 is to switch from nm to m (only one because the other nm
Yannick Roehlly's avatar
Yannick Roehlly committed
96
97
    # wavelength goes with the Fλ in W/m²/nm).
    fnu = 1e+29 * 1e-9 * flambda * wavelength * wavelength / c
Yannick Roehlly's avatar
Yannick Roehlly committed
98

Yannick Roehlly's avatar
Yannick Roehlly committed
99
    return fnu
Yannick Roehlly's avatar
Yannick Roehlly committed
100
101


Yannick Roehlly's avatar
Yannick Roehlly committed
102
def lambda_fnu_to_flambda(wavelength, fnu):
Yannick Roehlly's avatar
Yannick Roehlly committed
103
104
105
    """
    Convert a Fν vs λ spectrum to Fλ vs λ

Yannick Roehlly's avatar
Yannick Roehlly committed
106
    Parameters
Yannick Roehlly's avatar
Yannick Roehlly committed
107
    ----------
108
    wavelength: list-like of floats
Yannick Roehlly's avatar
Yannick Roehlly committed
109
        The wavelengths in nm.
110
    fnu: list-like of floats
Yannick Roehlly's avatar
Yannick Roehlly committed
111
112
        The Fν flux density in mJy (of the  Lν luminosity density in
        1.e-29 W/Hz).
Yannick Roehlly's avatar
Yannick Roehlly committed
113
114
115

    Returns
    -------
116
    flambda: array of floats
Yannick Roehlly's avatar
Yannick Roehlly committed
117
        Fλ flux density in W/m²/nm (or Lλ luminosity density in W/nm).
Yannick Roehlly's avatar
Yannick Roehlly committed
118
119

    """
Yannick Roehlly's avatar
Yannick Roehlly committed
120
121
    wavelength = np.array(wavelength, dtype=float)
    fnu = np.array(fnu, dtype=float)
Yannick Roehlly's avatar
Yannick Roehlly committed
122

Yannick Roehlly's avatar
Yannick Roehlly committed
123
124
125
    # Factor 1e-29 is to switch from Jy to W/m²/Hz
    # Factor 1e+9 is to switch from m to nm
    flambda = 1e-29 * 1e+9 * fnu / (wavelength * wavelength) * c
Yannick Roehlly's avatar
Yannick Roehlly committed
126

Yannick Roehlly's avatar
Yannick Roehlly committed
127
    return flambda
Yannick Roehlly's avatar
Yannick Roehlly committed
128
129


Yannick Roehlly's avatar
Yannick Roehlly committed
130
131
def redshift_spectrum(wavelength, flux, redshift, is_fnu=False):
    """Redshit a spectrum
Yannick Roehlly's avatar
Yannick Roehlly committed
132

Yannick Roehlly's avatar
Yannick Roehlly committed
133
    Parameters
Yannick Roehlly's avatar
Yannick Roehlly committed
134
    ----------
135
    wavelength: array like of floats
Yannick Roehlly's avatar
Yannick Roehlly committed
136
        The wavelength in nm.
137
    flux: array like of floats
Yannick Roehlly's avatar
Yannick Roehlly committed
138
        The flux or luminosity density.
139
    redshift: float
Yannick Roehlly's avatar
Yannick Roehlly committed
140
        The redshift.
141
    is_fnu: boolean
Yannick Roehlly's avatar
Yannick Roehlly committed
142
143
144
        If false (default) the flux is a Fλ density in W/m²/nm (or a Lλ
        luminosity density in W/nm). If true, the flux is a Fν density in mJy
        (or a Lν luminosity density in 1.e-29 W/Hz).
Yannick Roehlly's avatar
Yannick Roehlly committed
145
146
147

    Results
    -------
148
    wavelength, flux: tuple of numpy arrays of floats
Yannick Roehlly's avatar
Yannick Roehlly committed
149
150
        The redshifted spectrum with the same kind of flux (or luminosity)
        density as the input.
Yannick Roehlly's avatar
Yannick Roehlly committed
151
152

    """
Yannick Roehlly's avatar
Yannick Roehlly committed
153
154
155
    wavelength = np.array(wavelength, dtype=float)
    flux = np.array(flux, dtype=float)
    redshift = float(redshift)
Yannick Roehlly's avatar
Yannick Roehlly committed
156

Yannick Roehlly's avatar
Yannick Roehlly committed
157
158
159
160
    if redshift < 0:
        redshift_factor = 1. / (1. - redshift)
    else:
        redshift_factor = 1. + redshift
Yannick Roehlly's avatar
Yannick Roehlly committed
161

Yannick Roehlly's avatar
Yannick Roehlly committed
162
163
164
    if is_fnu:
        # Switch to Fλ
        flux = lambda_fnu_to_flambda(wavelength, flux)
Yannick Roehlly's avatar
Yannick Roehlly committed
165

Yannick Roehlly's avatar
Yannick Roehlly committed
166
167
    wavelength *= redshift_factor
    flux /= redshift_factor
Yannick Roehlly's avatar
Yannick Roehlly committed
168

Yannick Roehlly's avatar
Yannick Roehlly committed
169
170
171
    if is_fnu:
        # Switch back to Fλ
        flux = lambda_flambda_to_fnu(wavelength, flux)
Yannick Roehlly's avatar
Yannick Roehlly committed
172

Yannick Roehlly's avatar
Yannick Roehlly committed
173
    return wavelength, flux
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236


def memoise_1var(f):
    """
    Memoisation helper for a function taking 1 numpy array as input. As it is
    not hashable, we cannot use the standard memoisation function. Here we
    we use the array size, minimum, and maximum values as hashes.

    Parameters
    ----------
    f: function
        Function to memoise.

    Returns
    -------
    memoise_helper: function
        Meomoised best_grid function

    """
    memo = {}

    def memoise_helper(x):
        sx = x.size
        minx = x[0]
        maxx = x[-1]
        if (sx, minx, maxx) not in memo:
            memo[(sx, minx, maxx)] = f(x)
        return memo[(sx, minx, maxx)]
    return memoise_helper


def memoise_2var(f):
    """
    Memoisation helper for a function taking 2 numpy arrays as input. As they
    are not hashable, we cannot use the standard memoisation function. Here
    we use the array sizes, minimum, and maximum values as hashes.

    Parameters
    ----------
    f: function
        Function to memoise.

    Returns
    -------
    memoise_helper: function
        Meomoised best_grid function

    """
    memo = {}

    def memoise_helper(x, y):
        sx = x.size
        minx = x[0]
        maxx = x[-1]
        sy = y.size
        miny = y[0]
        maxy = y[-1]
        if (sx, minx, maxx, sy, miny, maxy) not in memo:
            memo[(sx, minx, maxx, sy, miny, maxy)] = f(x, y)
        return memo[(sx, minx, maxx, sy, miny, maxy)]
    return memoise_helper


237
def best_grid(wavelengths1, wavelengths2, key):
238
239
240
241
242
243
244
245
246
247
248
249
    """
    Return the best wavelength grid to regrid to arrays

    Considering the two wavelength grids passed in parameters, this function
    compute the best new grid that will be used to regrid the two spectra
    before combining them. We do not use np.unique as it is much slowe than
    finding the unique elements by hand.

    Parameters
    ----------
    wavelengths1, wavelengths2: array of floats
        The wavelength grids to be 'regridded'.
250
251
    key: tuple
        Key to key the results in cache.
252
253
254
255
256
257
258
259

    Returns
    -------
    new_grid: array of floats
        Array containing all the wavelengths found in the input arrays.

    """

260
261
    if key in best_grid_cache:
        return best_grid_cache[key]
262
263
264
265
    wl = np.concatenate((wavelengths1, wavelengths2))
    wl.sort(kind='mergesort')
    flag = np.ones(len(wl), dtype=bool)
    np.not_equal(wl[1:], wl[:-1], out=flag[1:])
266
    best_grid_cache[key] = wl[flag]
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
    return wl[flag]


@memoise_2var
def diff_wl(wavelengths1, wavelengths2):
    """
    Memoised version of numpy.setdiff1d. It is used to find the values that
    are in wavelengths1 but not in wavelengths2.

    Parameters
    ----------
    wavelengths1, wavelengths2: array of floats
        The wavelength grids to be compared

    Returns
    -------
    setdiff1d: array
        Array containing the wavelengths in wavelengths1 but not in
        wavelengths2.

    """

    return np.setdiff1d(wavelengths1, wavelengths2, assume_unique=True)


@memoise_1var
def argsort_wl(wavelengths):
    """
    Memoised version of numpy.argsort. It is used to sort the wavelengths.

    Parameters
    ----------
    wavelengths: array
        Array containing the wavelengths to be sorted.

    Returns
    -------
    argsort: array
        Array containing the sorting indices.

    """

    return np.argsort(wavelengths,
                      kind='mergesort')


def interpolate_lumin(wl, lumin, wl_new, lumin_new):
    """
    Procedure to interpolate the luminosity components given the wavelengths
    grid and the luminosity of a new component.

    Parameters
    ----------
    wl: 1D array
        Wavelengths grid of the luminosity components.
    lumin: 2D array
        Luminosity components. The first index gives a given component and the
        second gives the wavelength.
    wl_new: 1D array
        Wavelengths grid of he new component.
    lumin_new: 1D array
        Luminosity of the new component.

    Returns
    -------
    wl_best: 1D array
        The best merged wavelengths grid
    lumin_out: 2D array
        New luminosities array resampled if necessary to the best merge
        wavelengths grid, and including the new component.

    """
    # First we remove from the wavelengths grid of the new component the
    # wavelengths that are already present for the other components. We do
    # that to avoid having the interpolate on a wavelength for which we
    # already know the answer.
    wl_unique = diff_wl(wl_new, wl)

    # Output 2D array. We increase the number of components by one to include
    # the new directly rather than vstack-ing later. We also increase the size
    # to include possible new wavelengths
    lumin_out = np.zeros((lumin.shape[0]+1, lumin.shape[1]+wl_unique.size))

    # If the new component has non already existing wavelengths then we
    # interpolate on these new wavelengths and only those ones. The first
    # lumin.shape[0] elements are the same as the input luminosity components.
    # We only carry out the interpolation of the new components. Finally, we
    # reorder the output array with increasing wavelength.
    if wl_unique.size > 0:
        lumin_out[:-1, :lumin.shape[1]] = lumin

        # We interpolate only on the wavelengths where the components are
        # already defined.
        w = np.where((wl_unique > wl[0]) & (wl_unique < wl[-1]))
        for i in range(lumin.shape[0]):
            lumin_out[i, lumin.shape[1]+w[0]] = np.interp(wl_unique[w], wl,
                                                          lumin[i, :])
        wl_best = np.concatenate((wl, wl_unique))
        s = argsort_wl(wl_best)
        wl_best = wl_best[s]
        lumin_out = lumin_out[:, s]
    else:
        wl_best = wl
        lumin_out[:-1, :] = lumin

    # We interpolate the new component on the new merged wavelength grid.
    lumin_out[-1, :] = np.interp(wl_best, wl_new, lumin_new, left=0., right=0.)

    return (wl_best, lumin_out)
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412


def flux_trapz(y, x, key):
    """
    Light weight trapezoidal rule used to compute the flux in filters. It
    assumes that all the x and y input arrays are already numpy arrays to save
    time. Also the width between x elements are saved in cache using the "key"
    parameter as for a given x sampling it will always be the same. We do not
    compute they key ourselves because we do not have a proper way to hash it.
    Using x[0], x[-1], and x.size is not sufficient as if the redshift is small
    it is very likely that not of these elements will change. However the
    calling function has this knowledge and actually uses them to get the
    resampled filters from the cache. That way the two cache are sure to remain
    consistent.

    Parameters
    ----------
    y: 1D array
        The values to be integrated, typically fluxes
    x: 1D array
        Sampling of y, typically wavelengths
    key: tuple
        Contains the size of the x array, the name of the filter, and the
        redshift. Those elements point to a unique x grid. This is used to
        cache some computations are the x sampling will not change.

    Returns
    -------
    flux_trapz: float
        integral(y, dx)
    """
    if key in dx_cache:
        dx = dx_cache[key]
    else:
        dx = np.diff(x)
        dx_cache[key] = dx
    return (dx*(y[1:]+y[:-1])).sum()/2.