utils.py 9.93 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
8
from scipy.interpolate import interp1d
Yannick Roehlly's avatar
Yannick Roehlly committed
9
10
11
12
13


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

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

    Returns
    -------
21
    nu: float or array of floats
Yannick Roehlly's avatar
Yannick Roehlly committed
22
23
24
25
26
27
28
29
30
        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
31
    Parameters
Yannick Roehlly's avatar
Yannick Roehlly committed
32
    ----------
33
    frequency: float or numpy.array of floats
Yannick Roehlly's avatar
Yannick Roehlly committed
34
35
36
37
        The frequency(ies) in Hz.

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

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

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

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

    F = L / (4πDl2)

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

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

    """

65
    flux = luminosity / (4. * pi * dist * dist)
66

67
    return flux
Yannick Roehlly's avatar
Yannick Roehlly committed
68
69


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

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

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

    """
Yannick Roehlly's avatar
Yannick Roehlly committed
88
89
90
91
    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
92
    # Factor 1e-9 is to switch from nm to m (only one because the other nm
Yannick Roehlly's avatar
Yannick Roehlly committed
93
94
    # 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
95

Yannick Roehlly's avatar
Yannick Roehlly committed
96
    return fnu
Yannick Roehlly's avatar
Yannick Roehlly committed
97
98


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

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

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

    """
Yannick Roehlly's avatar
Yannick Roehlly committed
117
118
    wavelength = np.array(wavelength, dtype=float)
    fnu = np.array(fnu, dtype=float)
Yannick Roehlly's avatar
Yannick Roehlly committed
119

Yannick Roehlly's avatar
Yannick Roehlly committed
120
121
122
    # 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
123

Yannick Roehlly's avatar
Yannick Roehlly committed
124
    return flambda
Yannick Roehlly's avatar
Yannick Roehlly committed
125
126


Yannick Roehlly's avatar
Yannick Roehlly committed
127
128
def redshift_spectrum(wavelength, flux, redshift, is_fnu=False):
    """Redshit a spectrum
Yannick Roehlly's avatar
Yannick Roehlly committed
129

Yannick Roehlly's avatar
Yannick Roehlly committed
130
    Parameters
Yannick Roehlly's avatar
Yannick Roehlly committed
131
    ----------
132
    wavelength: array like of floats
Yannick Roehlly's avatar
Yannick Roehlly committed
133
        The wavelength in nm.
134
    flux: array like of floats
Yannick Roehlly's avatar
Yannick Roehlly committed
135
        The flux or luminosity density.
136
    redshift: float
Yannick Roehlly's avatar
Yannick Roehlly committed
137
        The redshift.
138
    is_fnu: boolean
Yannick Roehlly's avatar
Yannick Roehlly committed
139
140
141
        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
142
143
144

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

    """
Yannick Roehlly's avatar
Yannick Roehlly committed
150
151
152
    wavelength = np.array(wavelength, dtype=float)
    flux = np.array(flux, dtype=float)
    redshift = float(redshift)
Yannick Roehlly's avatar
Yannick Roehlly committed
153

Yannick Roehlly's avatar
Yannick Roehlly committed
154
155
156
157
    if redshift < 0:
        redshift_factor = 1. / (1. - redshift)
    else:
        redshift_factor = 1. + redshift
Yannick Roehlly's avatar
Yannick Roehlly committed
158

Yannick Roehlly's avatar
Yannick Roehlly committed
159
160
161
    if is_fnu:
        # Switch to Fλ
        flux = lambda_fnu_to_flambda(wavelength, flux)
Yannick Roehlly's avatar
Yannick Roehlly committed
162

Yannick Roehlly's avatar
Yannick Roehlly committed
163
164
    wavelength *= redshift_factor
    flux /= redshift_factor
Yannick Roehlly's avatar
Yannick Roehlly committed
165

Yannick Roehlly's avatar
Yannick Roehlly committed
166
167
168
    if is_fnu:
        # Switch back to Fλ
        flux = lambda_flambda_to_fnu(wavelength, flux)
Yannick Roehlly's avatar
Yannick Roehlly committed
169

Yannick Roehlly's avatar
Yannick Roehlly committed
170
    return wavelength, flux
171
172
173
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
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
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


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


@memoise_2var
def best_grid(wavelengths1, wavelengths2):
    """
    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'.

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

    """

    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:])
    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)