18630926536

Committed 19 Oct 2025 01:12PM UTC coverage: 91.723%. First build

Build # 18630926536

Build Type

Pull #95

github

Committed by

web-flow

Commit Message

Merge 1ab293551 into c6da64f2f

Pull Request Pull Request #95: switch to uv

Run Details

156 of 173 new or added lines in 14 files covered. (90.17%)

1884 of 2054 relevant lines covered (91.72%)

3.67 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
from importlib.resources import files


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

jveitchmichaelis / rascal / 18630926536

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