4515862454

Build Type

push

github

Committed by cylammarco

Commit Message

added remarklint to pre-commit. linted everything.

Run Details

1884 of 2056 relevant lines covered (91.63%)

3.65 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

90.12

/src/rascal/util.py

import os

import numpy as np
from scipy.optimize import curve_fit
import pkg_resources


def get_vapour_pressure(temperature):
    """
    Appendix A.I of https://emtoolbox.nist.gov/Wavelength/Documentation.asp
    """
    K1 = 1.16705214528e03
    K2 = -7.24213167032e05
    K3 = -1.70738469401e01
    K4 = 1.20208247025e04
    K5 = -3.23255503223e06
    K6 = 1.49151086135e01
    K7 = -4.82326573616e03
    K8 = 4.05113405421e05
    K9 = -2.38555575678e-01
    K10 = 6.50175348448e02
    omega = temperature + K9 / (temperature - K10)
    A = omega**2.0 + K1 * omega + K2
    B = K3 * omega**2.0 + K4 * omega + K5
    C = K6 * omega**2.0 + K7 * omega + K8
    X = -B + np.sqrt(B**2.0 - 4 * A * C)
    vapour_pressure = 10**6.0 * (2.0 * C / X) ** 4.0
    return vapour_pressure


def get_vapour_partial_pressure(relative_humidity, vapour_pressure):
    """
    Appendix A.II of https://emtoolbox.nist.gov/Wavelength/Documentation.asp
    """
    partial_pressure = relative_humidity / 100.0 * vapour_pressure
    return partial_pressure


def edlen_refraction(
    wavelengths, temperature, pressure, vapour_partial_pressure
):
    """
    Appendix A.IV of https://emtoolbox.nist.gov/Wavelength/Documentation.asp

    Parameters
    ----------
    wavelengths: float
        In unit of Angstrom
    temperature: float
        In unit of Celcius
    pressure: float
        In unit of Pascal

    """

    # Convert to micron for computing variable S
    w = np.array(wavelengths) / 1e4

    t = temperature
    T = temperature + 273.15

    A = 8342.54
    B = 2406147.0
    C = 15998.0
    D = 96095.43
    E = 0.601
    F = 0.00972
    G = 0.003661
    S = np.array(w) ** -2.0
    n_s = 1.0 + 1e-8 * (A + B / (130 - S) + C / (38.9 - S))
    X = (1.0 + 1e-8 * (E - F * t) * pressure) / (1.0 + G * t)
    n_tp = 1.0 + pressure * (n_s - 1.0) * X / D
    n = (
        n_tp
        - 1e-10
        * (292.75 / T)
        * (3.7345 - 0.0401 * S)
        * vapour_partial_pressure
    )
    return n


def vacuum_to_air_wavelength(
    wavelengths, temperature=273.15, pressure=101325, relative_humidity=0
):
    """

    The conversion follows the Modified Edlén Equations

    https://emtoolbox.nist.gov/Wavelength/Documentation.asp

    pressure drops by ~10% per 1000m above sea level
    temperature depends heavily on the location
    relative humidity is between 0-100, depends heavily on the location


    Parameters
    ----------
    wavelengths: float or numpy.array
        Wavelengths in vacuum
    temperature: float
        In unit of Kelvin
    pressure: float
        In unit of Pa
    relative_humidity: float
        Unitless in percentage (i.e. 0 - 100)

    Returns
    -------
    air wavelengths: float or numpy.array
        The wavelengths in air given the condition
    """

    # Convert to celcius
    t = temperature - 273.15

    vapour_pressure = get_vapour_pressure(t)
    vapour_partial_pressure = get_vapour_partial_pressure(
        relative_humidity, vapour_pressure
    )
    return np.array(wavelengths) / edlen_refraction(
        wavelengths, t, pressure, vapour_partial_pressure
    )


def air_to_vacuum_wavelength(
    wavelengths, temperature=273.15, pressure=101325, relative_humidity=0
):
    """
    The conversion follows the Modified Edlén Equations
    https://emtoolbox.nist.gov/Wavelength/Documentation.asp
    https://iopscience.iop.org/article/10.1088/0026-1394/35/2/8
    pressure drops by ~10% per 1000m above sea level
    temperature depends heavily on the location
    relative humidity is between 0-100, depends heavily on the location
    Parameters
    ----------
    wavelengths: float or numpy.array
        Wavelengths in vacuum in unit of Angstrom
    temperature: float
        In unit of Kelvin
    pressure: float
        In unit of Pa
    relative_humidity: float
        Unitless in percentage (i.e. 0 - 100)
    Returns
    -------
    air wavelengths: float or numpy.array
        The wavelengths in air given the condition
    """

    # Convert to celcius
    t = temperature - 273.15

    # get_vapour_pressure takes temperature in Celcius
    vapour_pressure = get_vapour_pressure(t)
    vapour_partial_pressure = get_vapour_partial_pressure(
        relative_humidity, vapour_pressure
    )
    return np.array(wavelengths) * edlen_refraction(
        wavelengths, t, pressure, vapour_partial_pressure
    )


def filter_wavelengths(lines, min_atlas_wavelength, max_atlas_wavelength):
    """
    Filters a wavelength list to a minimum and maximum range.

    Parameters
    ----------

    lines: list
        List of input wavelengths
    min_atlas_wavelength: int
        Min wavelength, Ansgtrom
    max_atlas_wavelength: int
        Max wavelength, Angstrom

    Returns
    -------

    lines: list
        Filtered wavelengths within specified range limit

    """

    wavelengths = lines[:, 1].astype(np.float64)

    _, index, _ = np.unique(wavelengths, return_counts=True, return_index=True)

    wavelength_mask = (wavelengths[index] >= min_atlas_wavelength) & (
        wavelengths[index] <= max_atlas_wavelength
    )

    return lines[index][wavelength_mask]


def filter_separation(wavelengths, min_separation=0):
    """
    Filters a wavelength list by a separation threshold.

    Parameters
    ----------

    wavelengths: list
        List of input wavelengths
    min_separation: int
        Separation threshold, Ansgtrom

    Returns
    -------

    distance_mask: list
        Mask of values which satisfy the separation criteria

    """

    left_dists = np.zeros_like(wavelengths)
    left_dists[1:] = wavelengths[1:] - wavelengths[:-1]

    right_dists = np.zeros_like(wavelengths)
    right_dists[:-1] = wavelengths[1:] - wavelengths[:-1]
    distances = np.minimum(right_dists, left_dists)
    distances[0] = right_dists[0]
    distances[-1] = left_dists[-1]

    distance_mask = np.abs(distances) >= min_separation

    return distance_mask


def filter_intensity(elements, lines, min_intensity=None):
    """
    Filters a line list by an intensity threshold
    Parameters
    ----------
    lines: list[tuple (str, float, float)]
        A list of input lines where 1st parameter is the name of the element
        the 2nd parameter is the wavelength, the 3rd is the intensities
    min_intensity: float
        Intensity threshold
    Returns
    -------
    lines: list
        Filtered line list
    """

    if min_intensity is None:
        min_intensity_dict = {}

        for i, e in enumerate(elements):
            min_intensity_dict[e] = 0.0

    elif isinstance(min_intensity, (int, float)):
        min_intensity_dict = {}

        for i, e in enumerate(elements):
            min_intensity_dict[e] = float(min_intensity)

    elif isinstance(min_intensity, (list, np.ndarray)):
        assert len(min_intensity) == len(elements), (
            "min_intensity has to be in, float of list/array "
            "the same size as the elements. min_intensity is {}"
            "and elements is {}.".format(min_intensity, elements)
        )

        min_intensity_dict = {}

        for i, e in enumerate(elements):
            min_intensity_dict[e] = min_intensity[i]

    else:
        raise ValueError(
            "min_intensity has to be in, float of list/array "
            "the same size as the elements. min_intensity is {}"
            "and elements is {}.".format(min_intensity, elements)
        )

    out = []

    for line in lines:
        element = line[0]
        intensity = float(line[2])

        if intensity >= min_intensity_dict[element]:
            out.append(True)
        else:
            out.append(False)

    return np.array(out).astype(bool)


def load_calibration_lines(
    elements=[],
    min_atlas_wavelength=3000.0,
    max_atlas_wavelength=15000.0,
    min_intensity=10,
    min_distance=10,
    vacuum=False,
    pressure=101325.0,
    temperature=273.15,
    relative_humidity=0.0,
    linelist="nist",
    brightest_n_lines=None,
):
    """
    Load calibration lines from the standard NIST atlas.
    Rascal provides a cleaned set of NIST lines that can be
    used for general purpose calibration. It is recommended
    however that for repeated and robust calibration, the
    user should specify an instrument-specific atlas.

    Provide a wavelength range suitable to your calibration
    source. You can also specify a minimum intensity that
    corresponds to the values listed in the NIST tables.

    If you want air wavelengths (default), you can provide
    atmospheric conditions for your system. In most cases
    the default values of standard temperature and pressure
    should be sufficient.


    Parameters
    ----------
    elements: list
        List of short element names, e.g. He as per NIST
    linelist: str
        Either 'nist' to use the default lines or path to a linelist file.
    min_atlas_wavelength: int
        Minimum wavelength to search, Angstrom
    max_atlas_wavelength: int
        Maximum wavelength to search, Angstrom
    min_intensity: int
        Minimum intensity to search, per NIST
    max_intensity: int
        Maximum intensity to search, per NIST
    vacuum: bool
        Return vacuum wavelengths
    pressure: float
        Atmospheric pressure, Pascal
    temperature: float
        Temperature in Kelvin, default room temp
    relative_humidity: float
        Relative humidity, percent

    Returns
    -------
    out: list
        Emission lines corresponding to the parameters specified
    """

    if isinstance(elements, str):
        elements = [elements]

    # Element, wavelength, intensity
    if isinstance(linelist, str):
        if linelist.lower() == "nist":
            file_path = pkg_resources.resource_filename(
                "rascal", "arc_lines/nist_clean.csv"
            )
            lines = np.loadtxt(file_path, delimiter=",", dtype=">U12")
        elif os.path.exists(linelist):
            lines = np.loadtxt(linelist, delimiter=",", dtype=">U12")
        else:
            raise ValueError(
                f"Unknown string is provided as linelist: {linelist}."
            )
    else:
        raise ValueError("Please provide a valid format of line list.")

    # Mask elements
    mask = [(li[0] in elements) for li in lines]
    lines = lines[mask]

    # update the wavelength limit
    if not vacuum:
        min_atlas_wavelength, max_atlas_wavelength = air_to_vacuum_wavelength(
            (min_atlas_wavelength, max_atlas_wavelength),
            temperature,
            pressure,
            relative_humidity,
        )

    # Filter wavelengths
    lines = filter_wavelengths(
        lines, min_atlas_wavelength, max_atlas_wavelength
    )

    # Filter intensities
    if isinstance(min_intensity, (float, int, list, np.ndarray)):
        intensity_mask = filter_intensity(elements, lines, min_intensity)
    else:
        intensity_mask = np.ones_like(lines[:, 0]).astype(bool)

    element_list = lines[:, 0][intensity_mask]
    wavelength_list = lines[:, 1][intensity_mask].astype("float64")
    intensity_list = lines[:, 2][intensity_mask].astype("float64")

    if brightest_n_lines is not None:
        to_keep = np.argsort(np.array(intensity_list))[::-1][
            :brightest_n_lines
        ]
        element_list = element_list[to_keep]
        intensity_list = intensity_list[to_keep]
        wavelength_list = wavelength_list[to_keep]

    # Calculate peak separation
    if min_distance > 0:
        distance_mask = filter_separation(wavelength_list, min_distance)
    else:
        distance_mask = np.ones_like(wavelength_list.astype(bool))

    element_list = element_list[distance_mask]
    wavelength_list = wavelength_list[distance_mask]
    intensity_list = intensity_list[distance_mask]

    # Vacuum to air conversion
    if not vacuum:
        wavelength_list = vacuum_to_air_wavelength(
            wavelength_list, temperature, pressure, relative_humidity
        )

    return element_list, wavelength_list, intensity_list


def gauss(x, a, x0, sigma):
    """
    1D Gaussian

    Parameters
    ----------
    x:
        value or values to evaluate the Gaussian at
    a: float
        Magnitude
    x0: float
        Gaussian centre
    sigma: float
        Standard deviation (spread)

    Returns
    -------
    out: list
        The Gaussian function evaluated at provided x
    """

    return a * np.exp(-((x - x0) ** 2) / (2 * sigma**2 + 1e-9))


def refine_peaks(spectrum, peaks, window_width=10, distance=None):
    """
    Refine peak locations in a spectrum from a set of initial estimates.

    This function attempts to fit a Gaussian to each peak in the provided
    list. It returns a list of sub-pixel refined peaks. If two peaks are
    very close, they can be refined to the same location. In this case
    only one of the peaks will be returned - i.e. this function will return
    a unique set of peak locations.

    Parameters
    ----------
    spectrum: list
        Input spectrum (list of intensities)
    peaks: list
        A list of peak locations in pixels
    window_width: int
        Size of window to consider in fit either side of
        initial peak location

    Returns
    -------
    refined_peaks: list
        A list of refined peak locations

    """

    refined_peaks = []

    spectrum = np.array(spectrum)
    length = len(spectrum)

    for peak in peaks:
        y = spectrum[
            max(0, int(peak) - window_width) : min(
                int(peak) + window_width, length
            )
        ]
        y /= np.nanmax(y)
        x = np.arange(len(y))

        mask = np.isfinite(y) & ~np.isnan(y)
        n = np.sum(mask)

        if n == 0:
            continue

        mean = np.sum(x * y) / n
        sigma = np.sum(y * (x - mean) ** 2) / n

        try:
            popt, _ = curve_fit(gauss, x[mask], y[mask], p0=[1, mean, sigma])
            height, centre, _ = popt

            if height < 0:
                continue

            if centre > len(spectrum) or centre < 0:
                continue

            refined_peaks.append(peak - window_width + centre)

        except RuntimeError:
            continue

    refined_peaks = np.array(refined_peaks)
    mask = (refined_peaks > 0) & (refined_peaks < len(spectrum))

    # Remove peaks that are within rounding errors from each other from the
    # curve_fit
    distance_mask = np.isclose(refined_peaks[:-1], refined_peaks[1:])
    distance_mask = np.insert(distance_mask, 0, False)

    return refined_peaks[mask & ~distance_mask]


def _derivative(p):
    """
    Compute the derivative of a polynomial function.

    Parameters
    ----------
    p: list
        Polynomial coefficients, in increasing order
        (e.g. 0th coefficient first)

    Returns
    -------
    derv: list
        Derivative coefficients, i * p[i]

    """
    derv = []
    for i in range(1, len(p)):
        derv.append(i * p[i])
    return derv

1	import os	4✔
2
3	import numpy as np	4✔
4	from scipy.optimize import curve_fit	4✔
5	import pkg_resources	4✔
6
7
8	def get_vapour_pressure(temperature):	4✔
9	"""
10	Appendix A.I of https://emtoolbox.nist.gov/Wavelength/Documentation.asp
11	"""
12	K1 = 1.16705214528e03	4✔
13	K2 = -7.24213167032e05	4✔
14	K3 = -1.70738469401e01	4✔
15	K4 = 1.20208247025e04	4✔
16	K5 = -3.23255503223e06	4✔
17	K6 = 1.49151086135e01	4✔
18	K7 = -4.82326573616e03	4✔
19	K8 = 4.05113405421e05	4✔
20	K9 = -2.38555575678e-01	4✔
21	K10 = 6.50175348448e02	4✔
22	omega = temperature + K9 / (temperature - K10)	4✔
23	A = omega*2.0 + K1 omega + K2	4✔
24	B = K3 * omega*2.0 + K4 omega + K5	4✔
25	C = K6 * omega*2.0 + K7 omega + K8	4✔
26	X = -B + np.sqrt(B*2.0 - 4 A * C)	4✔
27	vapour_pressure = 10*6.0 (2.0 * C / X) ** 4.0	4✔
28	return vapour_pressure	4✔
29
30
31	def get_vapour_partial_pressure(relative_humidity, vapour_pressure):	4✔
32	"""
33	Appendix A.II of https://emtoolbox.nist.gov/Wavelength/Documentation.asp
34	"""
35	partial_pressure = relative_humidity / 100.0 * vapour_pressure	4✔
36	return partial_pressure	4✔
37
38
39	def edlen_refraction(	4✔
40	wavelengths, temperature, pressure, vapour_partial_pressure
41	):
42	"""
43	Appendix A.IV of https://emtoolbox.nist.gov/Wavelength/Documentation.asp
44
45	Parameters
46	----------
47	wavelengths: float
48	In unit of Angstrom
49	temperature: float
50	In unit of Celcius
51	pressure: float
52	In unit of Pascal
53
54	"""
55
56	# Convert to micron for computing variable S
57	w = np.array(wavelengths) / 1e4	4✔
58
59	t = temperature	4✔
60	T = temperature + 273.15	4✔
61
62	A = 8342.54	4✔
63	B = 2406147.0	4✔
64	C = 15998.0	4✔
65	D = 96095.43	4✔
66	E = 0.601	4✔
67	F = 0.00972	4✔
68	G = 0.003661	4✔
69	S = np.array(w) ** -2.0	4✔
70	n_s = 1.0 + 1e-8 * (A + B / (130 - S) + C / (38.9 - S))	4✔
71	X = (1.0 + 1e-8 * (E - F * t) * pressure) / (1.0 + G * t)	4✔
72	n_tp = 1.0 + pressure * (n_s - 1.0) * X / D	4✔
73	n = (	4✔
74	n_tp
75	- 1e-10
76	* (292.75 / T)
77	* (3.7345 - 0.0401 * S)
78	* vapour_partial_pressure
79	)
80	return n	4✔
81
82
83	def vacuum_to_air_wavelength(	4✔
84	wavelengths, temperature=273.15, pressure=101325, relative_humidity=0
85	):
86	"""
87
88	The conversion follows the Modified Edlén Equations
89
90	https://emtoolbox.nist.gov/Wavelength/Documentation.asp
91
92	pressure drops by ~10% per 1000m above sea level
93	temperature depends heavily on the location
94	relative humidity is between 0-100, depends heavily on the location
95
96
97	Parameters
98	----------
99	wavelengths: float or numpy.array
100	Wavelengths in vacuum
101	temperature: float
102	In unit of Kelvin
103	pressure: float
104	In unit of Pa
105	relative_humidity: float
106	Unitless in percentage (i.e. 0 - 100)
107
108	Returns
109	-------
110	air wavelengths: float or numpy.array
111	The wavelengths in air given the condition
112	"""
113
114	# Convert to celcius
115	t = temperature - 273.15	4✔
116
117	vapour_pressure = get_vapour_pressure(t)	4✔
118	vapour_partial_pressure = get_vapour_partial_pressure(	4✔
119	relative_humidity, vapour_pressure
120	)
121	return np.array(wavelengths) / edlen_refraction(	4✔
122	wavelengths, t, pressure, vapour_partial_pressure
123	)
124
125
126	def air_to_vacuum_wavelength(	4✔
127	wavelengths, temperature=273.15, pressure=101325, relative_humidity=0
128	):
129	"""
130	The conversion follows the Modified Edlén Equations
131	https://emtoolbox.nist.gov/Wavelength/Documentation.asp
132	https://iopscience.iop.org/article/10.1088/0026-1394/35/2/8
133	pressure drops by ~10% per 1000m above sea level
134	temperature depends heavily on the location
135	relative humidity is between 0-100, depends heavily on the location
136	Parameters
137	----------
138	wavelengths: float or numpy.array
139	Wavelengths in vacuum in unit of Angstrom
140	temperature: float
141	In unit of Kelvin
142	pressure: float
143	In unit of Pa
144	relative_humidity: float
145	Unitless in percentage (i.e. 0 - 100)
146	Returns
147	-------
148	air wavelengths: float or numpy.array
149	The wavelengths in air given the condition
150	"""
151
152	# Convert to celcius
153	t = temperature - 273.15	4✔
154
155	# get_vapour_pressure takes temperature in Celcius
156	vapour_pressure = get_vapour_pressure(t)	4✔
157	vapour_partial_pressure = get_vapour_partial_pressure(	4✔
158	relative_humidity, vapour_pressure
159	)
160	return np.array(wavelengths) * edlen_refraction(	4✔
161	wavelengths, t, pressure, vapour_partial_pressure
162	)
163
164
165	def filter_wavelengths(lines, min_atlas_wavelength, max_atlas_wavelength):	4✔
166	"""
167	Filters a wavelength list to a minimum and maximum range.
168
169	Parameters
170	----------
171
172	lines: list
173	List of input wavelengths
174	min_atlas_wavelength: int
175	Min wavelength, Ansgtrom
176	max_atlas_wavelength: int
177	Max wavelength, Angstrom
178
179	Returns
180	-------
181
182	lines: list
183	Filtered wavelengths within specified range limit
184
185	"""
186
187	wavelengths = lines[:, 1].astype(np.float64)	4✔
188
189	_, index, _ = np.unique(wavelengths, return_counts=True, return_index=True)	4✔
190
191	wavelength_mask = (wavelengths[index] >= min_atlas_wavelength) & (	4✔
192	wavelengths[index] <= max_atlas_wavelength
193	)
194
195	return lines[index][wavelength_mask]	4✔
196
197
198	def filter_separation(wavelengths, min_separation=0):	4✔
199	"""
200	Filters a wavelength list by a separation threshold.
201
202	Parameters
203	----------
204
205	wavelengths: list
206	List of input wavelengths
207	min_separation: int
208	Separation threshold, Ansgtrom
209
210	Returns
211	-------
212
213	distance_mask: list
214	Mask of values which satisfy the separation criteria
215
216	"""
217
218	left_dists = np.zeros_like(wavelengths)	4✔
219	left_dists[1:] = wavelengths[1:] - wavelengths[:-1]	4✔
220
221	right_dists = np.zeros_like(wavelengths)	4✔
222	right_dists[:-1] = wavelengths[1:] - wavelengths[:-1]	4✔
223	distances = np.minimum(right_dists, left_dists)	4✔
224	distances[0] = right_dists[0]	4✔
225	distances[-1] = left_dists[-1]	4✔
226
227	distance_mask = np.abs(distances) >= min_separation	4✔
228
229	return distance_mask	4✔
230
231
232	def filter_intensity(elements, lines, min_intensity=None):	4✔
233	"""
234	Filters a line list by an intensity threshold
235	Parameters
236	----------
237	lines: list[tuple (str, float, float)]
238	A list of input lines where 1st parameter is the name of the element
239	the 2nd parameter is the wavelength, the 3rd is the intensities
240	min_intensity: float
241	Intensity threshold
242	Returns
243	-------
244	lines: list
245	Filtered line list
246	"""
247
248	if min_intensity is None:	4✔
249	min_intensity_dict = {}	×
250
251	for i, e in enumerate(elements):	×
252	min_intensity_dict[e] = 0.0	×
253
254	elif isinstance(min_intensity, (int, float)):	4✔
255	min_intensity_dict = {}	4✔
256
257	for i, e in enumerate(elements):	4✔
258	min_intensity_dict[e] = float(min_intensity)	4✔
259
260	elif isinstance(min_intensity, (list, np.ndarray)):	×
261	assert len(min_intensity) == len(elements), (	×
262	"min_intensity has to be in, float of list/array "
263	"the same size as the elements. min_intensity is {}"
264	"and elements is {}.".format(min_intensity, elements)
265	)
266
267	min_intensity_dict = {}	×
268
269	for i, e in enumerate(elements):	×
270	min_intensity_dict[e] = min_intensity[i]	×
271
272	else:
273	raise ValueError(	×
274	"min_intensity has to be in, float of list/array "
275	"the same size as the elements. min_intensity is {}"
276	"and elements is {}.".format(min_intensity, elements)
277	)
278
279	out = []	4✔
280
281	for line in lines:	4✔
282	element = line[0]	4✔
283	intensity = float(line[2])	4✔
284
285	if intensity >= min_intensity_dict[element]:	4✔
286	out.append(True)	4✔
287	else:
288	out.append(False)	4✔
289
290	return np.array(out).astype(bool)	4✔
291
292
293	def load_calibration_lines(	4✔
294	elements=[],
295	min_atlas_wavelength=3000.0,
296	max_atlas_wavelength=15000.0,
297	min_intensity=10,
298	min_distance=10,
299	vacuum=False,
300	pressure=101325.0,
301	temperature=273.15,
302	relative_humidity=0.0,
303	linelist="nist",
304	brightest_n_lines=None,
305	):
306	"""
307	Load calibration lines from the standard NIST atlas.
308	Rascal provides a cleaned set of NIST lines that can be
309	used for general purpose calibration. It is recommended
310	however that for repeated and robust calibration, the
311	user should specify an instrument-specific atlas.
312
313	Provide a wavelength range suitable to your calibration
314	source. You can also specify a minimum intensity that
315	corresponds to the values listed in the NIST tables.
316
317	If you want air wavelengths (default), you can provide
318	atmospheric conditions for your system. In most cases
319	the default values of standard temperature and pressure
320	should be sufficient.
321
322
323	Parameters
324	----------
325	elements: list
326	List of short element names, e.g. He as per NIST
327	linelist: str
328	Either 'nist' to use the default lines or path to a linelist file.
329	min_atlas_wavelength: int
330	Minimum wavelength to search, Angstrom
331	max_atlas_wavelength: int
332	Maximum wavelength to search, Angstrom
333	min_intensity: int
334	Minimum intensity to search, per NIST
335	max_intensity: int
336	Maximum intensity to search, per NIST
337	vacuum: bool
338	Return vacuum wavelengths
339	pressure: float
340	Atmospheric pressure, Pascal
341	temperature: float
342	Temperature in Kelvin, default room temp
343	relative_humidity: float
344	Relative humidity, percent
345
346	Returns
347	-------
348	out: list
349	Emission lines corresponding to the parameters specified
350	"""
351
352	if isinstance(elements, str):	4✔
353	elements = [elements]	4✔
354
355	# Element, wavelength, intensity
356	if isinstance(linelist, str):	4✔
357	if linelist.lower() == "nist":	4✔
358	file_path = pkg_resources.resource_filename(	4✔
359	"rascal", "arc_lines/nist_clean.csv"
360	)
361	lines = np.loadtxt(file_path, delimiter=",", dtype=">U12")	4✔
362	elif os.path.exists(linelist):	×
363	lines = np.loadtxt(linelist, delimiter=",", dtype=">U12")	×
364	else:
365	raise ValueError(	×
366	f"Unknown string is provided as linelist: {linelist}."
367	)
368	else:
369	raise ValueError("Please provide a valid format of line list.")	×
370
371	# Mask elements
372	mask = [(li[0] in elements) for li in lines]	4✔
373	lines = lines[mask]	4✔
374
375	# update the wavelength limit
376	if not vacuum:	4✔
377	min_atlas_wavelength, max_atlas_wavelength = air_to_vacuum_wavelength(	4✔
378	(min_atlas_wavelength, max_atlas_wavelength),
379	temperature,
380	pressure,
381	relative_humidity,
382	)
383
384	# Filter wavelengths
385	lines = filter_wavelengths(	4✔
386	lines, min_atlas_wavelength, max_atlas_wavelength
387	)
388
389	# Filter intensities
390	if isinstance(min_intensity, (float, int, list, np.ndarray)):	4✔
391	intensity_mask = filter_intensity(elements, lines, min_intensity)	4✔
392	else:
393	intensity_mask = np.ones_like(lines[:, 0]).astype(bool)	×
394
395	element_list = lines[:, 0][intensity_mask]	4✔
396	wavelength_list = lines[:, 1][intensity_mask].astype("float64")	4✔
397	intensity_list = lines[:, 2][intensity_mask].astype("float64")	4✔
398
399	if brightest_n_lines is not None:	4✔
400	to_keep = np.argsort(np.array(intensity_list))[::-1][	4✔
401	:brightest_n_lines
402	]
403	element_list = element_list[to_keep]	4✔
404	intensity_list = intensity_list[to_keep]	4✔
405	wavelength_list = wavelength_list[to_keep]	4✔
406
407	# Calculate peak separation
408	if min_distance > 0:	4✔
409	distance_mask = filter_separation(wavelength_list, min_distance)	4✔
410	else:
411	distance_mask = np.ones_like(wavelength_list.astype(bool))	4✔
412
413	element_list = element_list[distance_mask]	4✔
414	wavelength_list = wavelength_list[distance_mask]	4✔
415	intensity_list = intensity_list[distance_mask]	4✔
416
417	# Vacuum to air conversion
418	if not vacuum:	4✔
419	wavelength_list = vacuum_to_air_wavelength(	4✔
420	wavelength_list, temperature, pressure, relative_humidity
421	)
422
423	return element_list, wavelength_list, intensity_list	4✔
424
425
426	def gauss(x, a, x0, sigma):	4✔
427	"""
428	1D Gaussian
429
430	Parameters
431	----------
432	x:
433	value or values to evaluate the Gaussian at
434	a: float
435	Magnitude
436	x0: float
437	Gaussian centre
438	sigma: float
439	Standard deviation (spread)
440
441	Returns
442	-------
443	out: list
444	The Gaussian function evaluated at provided x
445	"""
446
447	return a * np.exp(-((x - x0) ** 2) / (2 * sigma**2 + 1e-9))	4✔
448
449
450	def refine_peaks(spectrum, peaks, window_width=10, distance=None):	4✔
451	"""
452	Refine peak locations in a spectrum from a set of initial estimates.
453
454	This function attempts to fit a Gaussian to each peak in the provided
455	list. It returns a list of sub-pixel refined peaks. If two peaks are
456	very close, they can be refined to the same location. In this case
457	only one of the peaks will be returned - i.e. this function will return
458	a unique set of peak locations.
459
460	Parameters
461	----------
462	spectrum: list
463	Input spectrum (list of intensities)
464	peaks: list
465	A list of peak locations in pixels
466	window_width: int
467	Size of window to consider in fit either side of
468	initial peak location
469
470	Returns
471	-------
472	refined_peaks: list
473	A list of refined peak locations
474
475	"""
476
477	refined_peaks = []	4✔
478
479	spectrum = np.array(spectrum)	4✔
480	length = len(spectrum)	4✔
481
482	for peak in peaks:	4✔
483	y = spectrum[	4✔
484	max(0, int(peak) - window_width) : min(
485	int(peak) + window_width, length
486	)
487	]
488	y /= np.nanmax(y)	4✔
489	x = np.arange(len(y))	4✔
490
491	mask = np.isfinite(y) & ~np.isnan(y)	4✔
492	n = np.sum(mask)	4✔
493
494	if n == 0:	4✔
495	continue	×
496
497	mean = np.sum(x * y) / n	4✔
498	sigma = np.sum(y * (x - mean) ** 2) / n	4✔
499
500	try:	4✔
501	popt, _ = curve_fit(gauss, x[mask], y[mask], p0=[1, mean, sigma])	4✔
502	height, centre, _ = popt	4✔
503
504	if height < 0:	4✔
505	continue	×
506
507	if centre > len(spectrum) or centre < 0:	4✔
508	continue	4✔
509
510	refined_peaks.append(peak - window_width + centre)	4✔
511
512	except RuntimeError:	4✔
513	continue	4✔
514
515	refined_peaks = np.array(refined_peaks)	4✔
516	mask = (refined_peaks > 0) & (refined_peaks < len(spectrum))	4✔
517
518	# Remove peaks that are within rounding errors from each other from the
519	# curve_fit
520	distance_mask = np.isclose(refined_peaks[:-1], refined_peaks[1:])	4✔
521	distance_mask = np.insert(distance_mask, 0, False)	4✔
522
523	return refined_peaks[mask & ~distance_mask]	4✔
524
525
526	def _derivative(p):	4✔
527	"""
528	Compute the derivative of a polynomial function.
529
530	Parameters
531	----------
532	p: list
533	Polynomial coefficients, in increasing order
534	(e.g. 0th coefficient first)
535
536	Returns
537	-------
538	derv: list
539	Derivative coefficients, i * p[i]
540
541	"""
542	derv = []	4✔
543	for i in range(1, len(p)):	4✔
544	derv.append(i * p[i])	4✔
545	return derv	4✔

jveitchmichaelis / rascal / 4515862454

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous