24916259816

Committed 24 Apr 2026 11:20PM UTC coverage: 96.759% (-0.8%) from 97.562%

Build # 24916259816

Build Type

push

github

Committed by

web-flow

Commit Message

Clean up RTE interface. (#10)

Clean up RTE interface:

compute all Stokes parameters in one Frame
access Einstein coefficients inside TransitionRegistry
add caching for Paschen-Back eigenvalue problem solving and custom PaschenBack container class
add caching/precomputing of eta_a and eta_s frames
remove Einstein coefficients from loopers, merge them together with other factors instead
remove Paschen-Back explicit pre-computation

General clean-up:

automatically adjust for air index when plotting Stokes profiles
unify and refactor plotters
add clean normalization to plotters
improve argument names
unify hz vs sm1
use VERBOSE logging level for debugging
make loopers thread-safe
use custom logging init instead of yatools

New tests/demos:

add Hanle effect demo
add demo for LTE magnetic sensitivity of Mn I 5432
add demo for quick-start of SolRaT
add test to check that precomputing of SEE works as intended
add test to check the disable_r_s option of current vs legacy SEE implementation
add test to compare current vs legacy implementation of SEE for non-zero magnetic field
add test for sanity checks in SEE of some real atoms

Coverage Stats

214 of 220 branches covered (97.27%)

Branch coverage included in aggregate %.

2503 of 2588 relevant lines covered (96.72%)

0.97 hits per line

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

0.0

solrat/atom_model/shared/utility/plot_stokes_profiles.py

from enum import Enum, auto

import numpy as np
from matplotlib import pyplot as plt

from solrat.atom_model.shared.object.stokes import Stokes
from solrat.atom_model.shared.utility.functions import (
    frequency_sm1_to_lambda_A,
    lambda_air_to_vacuum,
    lambda_vacuum_to_air,
)
from solrat.engine.functions.decorators import log_method


class StokesNorm(Enum):
    """Normalization mode for Stokes profile plots."""

    NONE = auto()  # raw values, no normalization
    MAX_I = auto()  # divide all components by max(I)
    BY_REFERENCE = auto()  # divide all components by reference Stokes I (continuum)
    MAX_IpV_ImV = auto()  # normalize I+V and I-V each by their own max (IpmV plotter only)


def _compute_wavelength_axis(nu, reference_lambda_A_air, use_air_wavelengths):
    lambda_A_vac = frequency_sm1_to_lambda_A(nu)
    if use_air_wavelengths:
        wavelength = lambda_vacuum_to_air(lambda_A_vac)
        ref = reference_lambda_A_air
    else:
        wavelength = lambda_A_vac
        ref = lambda_air_to_vacuum(reference_lambda_A_air) if reference_lambda_A_air is not None else None
    if ref is not None:
        label = (
            r"$\Delta\lambda_\mathrm{air}$ ($\AA$)" if use_air_wavelengths else r"$\Delta\lambda_\mathrm{vac}$ ($\AA$)"
        )
        return wavelength - ref, label
    label = r"$\lambda_\mathrm{air}$ ($\AA$)" if use_air_wavelengths else r"$\lambda_\mathrm{vac}$ ($\AA$)"
    return wavelength, label


class PlotterBase:  # pragma: no cover
    Norm = StokesNorm

    def __init__(
        self, title, use_air_wavelengths, reference_lambda_A_air, n_axes, y_labels, figsize=(8, 8), x_label=None
    ):
        self.colors = plt.rcParams["axes.prop_cycle"].by_key()["color"]
        self.next_color_index = 0
        self.vacuum_to_air = use_air_wavelengths
        self.reference_lambda_A_air = reference_lambda_A_air
        self.fig, self.axs = plt.subplots(n_axes, 1, sharex=True, constrained_layout=True, figsize=figsize, num=title)
        self.fig.suptitle(title)
        for ax, label in zip(self.axs, y_labels):
            ax.set_ylabel(label)
        self._wavelength_label = x_label

    def _next_color(self, color):
        if color == "auto":
            color = self.colors[self.next_color_index % len(self.colors)]
            self.next_color_index += 1
        return color

    def _wavelength_axis(self, nu):
        wavelength, label = _compute_wavelength_axis(nu, self.reference_lambda_A_air, self.vacuum_to_air)
        if self._wavelength_label is None:
            self._wavelength_label = label
        return wavelength

    @log_method
    def show(self):
        for ax in self.axs:
            ax.grid(True)
            ax.legend(loc="upper left", bbox_to_anchor=(1.02, 1), fontsize="x-small")
        self.axs[-1].set_xlabel(self._wavelength_label or r"$\lambda_\mathrm{vac}$ ($\AA$)")
        plt.show()


class StokesPlotter_IV(PlotterBase):  # pragma: no cover
    """
    Stokes plotter class for Stokes I and V profiles.
    """

    def __init__(self, title="", use_air_wavelengths=True, reference_lambda_A_air=None):
        super().__init__(
            title=title,
            use_air_wavelengths=use_air_wavelengths,
            reference_lambda_A_air=reference_lambda_A_air,
            n_axes=2,
            y_labels=[r"Stokes $I/I_{max}$", r"Stokes $V/I_{max}$"],
        )

    @log_method
    def add(self, nu, stokes_I, stokes_V, color=None, label="", linewidth=1.5):
        color = self._next_color(color)
        wavelength = self._wavelength_axis(nu)
        if stokes_I is not None:
            self.axs[0].plot(wavelength, stokes_I / np.max(stokes_I), label=label, color=color, linewidth=linewidth)
        if stokes_V is not None:
            self.axs[1].plot(wavelength, stokes_V / np.max(stokes_I), label=label, color=color, linewidth=linewidth)

    @log_method
    def add_stokes(
        self,
        nu,
        stokes: Stokes,
        stokes_reference: Stokes = None,
        norm: StokesNorm = StokesNorm.NONE,
        color=None,
        label="",
        linewidth=1.5,
    ):
        if norm == StokesNorm.MAX_I:
            scale = np.max(stokes.I)
        elif norm == StokesNorm.BY_REFERENCE:
            scale = stokes_reference.I
        elif norm == StokesNorm.NONE:
            scale = 1
        else:
            raise ValueError(f"Did not recognize normalization option {norm}")

        self.add(
            nu=nu,
            stokes_I=stokes.I / scale,
            stokes_V=stokes.V / scale,
            color=color,
            label=label,
            linewidth=linewidth,
        )


class StokesPlotter_IV_IpmV(PlotterBase):  # pragma: no cover
    r"""
    Stokes plotter class for Stokes :math:`I, V, I\pm V` profiles.
    """

    def __init__(self, title="", use_air_wavelengths=False, reference_lambda_A_air=None):
        super().__init__(
            title=title,
            use_air_wavelengths=use_air_wavelengths,
            reference_lambda_A_air=reference_lambda_A_air,
            n_axes=3,
            y_labels=[r"Stokes $I$", r"Stokes $V$", r"Stokes $(I\pm V)$"],
        )

    @log_method
    def add(self, nu, stokes_I, stokes_V, color=None, label="", linewidth=1.5):
        color = self._next_color(color)
        wavelength = self._wavelength_axis(nu)
        self.axs[0].plot(wavelength, stokes_I, label=label, color=color, linewidth=linewidth)
        self.axs[1].plot(wavelength, stokes_V, label=label, color=color, linewidth=linewidth)
        self.axs[2].plot(wavelength, stokes_I + stokes_V, "-", label=label + " $I+V$", color=color, linewidth=linewidth)
        self.axs[2].plot(
            wavelength, stokes_I - stokes_V, "--", label=label + " $I-V$", color=color, linewidth=linewidth
        )

    @log_method
    def add_stokes(
        self,
        nu,
        stokes: Stokes,
        stokes_reference: Stokes = None,
        norm: StokesNorm = StokesNorm.NONE,
        color=None,
        label="",
        linewidth=1.5,
    ):
        color = self._next_color(color)
        wavelength = self._wavelength_axis(nu)

        if norm == StokesNorm.NONE:
            I, V = stokes.I, stokes.V
        elif norm == StokesNorm.MAX_I:
            scale = np.max(stokes.I)
            I, V = stokes.I / scale, stokes.V / scale
        elif norm == StokesNorm.BY_REFERENCE:
            scale = stokes_reference.I
            I, V = stokes.I / scale, stokes.V / scale
        elif norm == StokesNorm.MAX_IpV_ImV:
            I, V = stokes.I, stokes.V
        else:
            raise ValueError(f"Did not recognize normalization option {norm}")

        self.axs[0].plot(wavelength, I, label=label, color=color, linewidth=linewidth)
        self.axs[1].plot(wavelength, V, label=label, color=color, linewidth=linewidth)

        if norm == StokesNorm.MAX_IpV_ImV:
            IpV = stokes.I + stokes.V
            ImV = stokes.I - stokes.V
            IpV = IpV / np.max(np.abs(IpV))
            ImV = ImV / np.max(np.abs(ImV))
        else:
            IpV = I + V
            ImV = I - V

        self.axs[2].plot(wavelength, IpV, "-", label=label + " $I+V$", color=color, linewidth=linewidth)
        self.axs[2].plot(wavelength, ImV, "--", label=label + " $I-V$", color=color, linewidth=linewidth)


class StokesPlotter(PlotterBase):  # pragma: no cover
    r"""
    Stokes plotter class for Stokes :math:`I, Q, U, V` profiles.
    """

    def __init__(
        self,
        title="",
        use_air_wavelengths=False,
        reference_lambda_A_air=None,
        x_label=None,
        y_label_I=r"Stokes $I$",
        y_label_Q=r"Stokes $Q$",
        y_label_U=r"Stokes $U$",
        y_label_V=r"Stokes $V$",
    ):
        super().__init__(
            title=title,
            use_air_wavelengths=use_air_wavelengths,
            reference_lambda_A_air=reference_lambda_A_air,
            n_axes=4,
            y_labels=[y_label_I, y_label_Q, y_label_U, y_label_V],
            x_label=x_label,
        )

    @log_method
    def add(
        self,
        nu,
        stokes_I,
        stokes_Q,
        stokes_U,
        stokes_V,
        color=None,
        label="",
        style="-",
        linewidth=1.5,
    ):
        color = self._next_color(color)
        wavelength = self._wavelength_axis(nu)

        if stokes_I is not None:
            self.axs[0].plot(wavelength, stokes_I, style, label=label, color=color, linewidth=linewidth)
        if stokes_Q is not None:
            self.axs[1].plot(wavelength, stokes_Q, style, label=label, color=color, linewidth=linewidth)
        if stokes_U is not None:
            self.axs[2].plot(wavelength, stokes_U, style, label=label, color=color, linewidth=linewidth)
        if stokes_V is not None:
            self.axs[3].plot(wavelength, stokes_V, style, label=label, color=color, linewidth=linewidth)

    @log_method
    def add_stokes(
        self,
        nu,
        stokes: Stokes,
        stokes_reference: Stokes = None,
        norm: StokesNorm = StokesNorm.NONE,
        color=None,
        label="",
        style="-",
        linewidth=1.5,
    ):
        if norm == StokesNorm.MAX_I:
            scale = np.max(stokes.I)
        elif norm == StokesNorm.BY_REFERENCE:
            scale = stokes_reference.I
        elif norm == StokesNorm.NONE:
            scale = 1
        else:
            raise ValueError(f"Did not recognize normalization option {norm}")
        self.add(
            nu=nu,
            stokes_I=stokes.I / scale,
            stokes_Q=stokes.Q / scale,
            stokes_U=stokes.U / scale,
            stokes_V=stokes.V / scale,
            color=color,
            label=label,
            style=style,
            linewidth=linewidth,
        )

1	from enum import Enum, auto	×
2
3	import numpy as np	×
4	from matplotlib import pyplot as plt	×
5
6	from solrat.atom_model.shared.object.stokes import Stokes	×
7	from solrat.atom_model.shared.utility.functions import (	×
8	frequency_sm1_to_lambda_A,
9	lambda_air_to_vacuum,
10	lambda_vacuum_to_air,
11	)
12	from solrat.engine.functions.decorators import log_method	×
13
14
15	class StokesNorm(Enum):	×
16	"""Normalization mode for Stokes profile plots."""
17
18	NONE = auto() # raw values, no normalization	×
19	MAX_I = auto() # divide all components by max(I)	×
20	BY_REFERENCE = auto() # divide all components by reference Stokes I (continuum)	×
21	MAX_IpV_ImV = auto() # normalize I+V and I-V each by their own max (IpmV plotter only)	×
22
23
24	def _compute_wavelength_axis(nu, reference_lambda_A_air, use_air_wavelengths):	×
25	lambda_A_vac = frequency_sm1_to_lambda_A(nu)	×
26	if use_air_wavelengths:	×
27	wavelength = lambda_vacuum_to_air(lambda_A_vac)	×
28	ref = reference_lambda_A_air	×
29	else:
30	wavelength = lambda_A_vac	×
31	ref = lambda_air_to_vacuum(reference_lambda_A_air) if reference_lambda_A_air is not None else None	×
32	if ref is not None:	×
33	label = (	×
34	r"$\Delta\lambda_\mathrm{air}$ ($\AA$)" if use_air_wavelengths else r"$\Delta\lambda_\mathrm{vac}$ ($\AA$)"
35	)
36	return wavelength - ref, label	×
37	label = r"$\lambda_\mathrm{air}$ ($\AA$)" if use_air_wavelengths else r"$\lambda_\mathrm{vac}$ ($\AA$)"	×
38	return wavelength, label	×
39
40
41	class PlotterBase: # pragma: no cover
42	Norm = StokesNorm
43
44	def __init__(
45	self, title, use_air_wavelengths, reference_lambda_A_air, n_axes, y_labels, figsize=(8, 8), x_label=None
46	):
47	self.colors = plt.rcParams["axes.prop_cycle"].by_key()["color"]
48	self.next_color_index = 0
49	self.vacuum_to_air = use_air_wavelengths
50	self.reference_lambda_A_air = reference_lambda_A_air
51	self.fig, self.axs = plt.subplots(n_axes, 1, sharex=True, constrained_layout=True, figsize=figsize, num=title)
52	self.fig.suptitle(title)
53	for ax, label in zip(self.axs, y_labels):
54	ax.set_ylabel(label)
55	self._wavelength_label = x_label
56
57	def _next_color(self, color):
58	if color == "auto":
59	color = self.colors[self.next_color_index % len(self.colors)]
60	self.next_color_index += 1
61	return color
62
63	def _wavelength_axis(self, nu):
64	wavelength, label = _compute_wavelength_axis(nu, self.reference_lambda_A_air, self.vacuum_to_air)
65	if self._wavelength_label is None:
66	self._wavelength_label = label
67	return wavelength
68
69	@log_method
70	def show(self):
71	for ax in self.axs:
72	ax.grid(True)
73	ax.legend(loc="upper left", bbox_to_anchor=(1.02, 1), fontsize="x-small")
74	self.axs[-1].set_xlabel(self._wavelength_label or r"$\lambda_\mathrm{vac}$ ($\AA$)")
75	plt.show()
76
77
78	class StokesPlotter_IV(PlotterBase): # pragma: no cover
79	"""
80	Stokes plotter class for Stokes I and V profiles.
81	"""
82
83	def __init__(self, title="", use_air_wavelengths=True, reference_lambda_A_air=None):
84	super().__init__(
85	title=title,
86	use_air_wavelengths=use_air_wavelengths,
87	reference_lambda_A_air=reference_lambda_A_air,
88	n_axes=2,
89	y_labels=[r"Stokes $I/I_{max}$", r"Stokes $V/I_{max}$"],
90	)
91
92	@log_method
93	def add(self, nu, stokes_I, stokes_V, color=None, label="", linewidth=1.5):
94	color = self._next_color(color)
95	wavelength = self._wavelength_axis(nu)
96	if stokes_I is not None:
97	self.axs[0].plot(wavelength, stokes_I / np.max(stokes_I), label=label, color=color, linewidth=linewidth)
98	if stokes_V is not None:
99	self.axs[1].plot(wavelength, stokes_V / np.max(stokes_I), label=label, color=color, linewidth=linewidth)
100
101	@log_method
102	def add_stokes(
103	self,
104	nu,
105	stokes: Stokes,
106	stokes_reference: Stokes = None,
107	norm: StokesNorm = StokesNorm.NONE,
108	color=None,
109	label="",
110	linewidth=1.5,
111	):
112	if norm == StokesNorm.MAX_I:
113	scale = np.max(stokes.I)
114	elif norm == StokesNorm.BY_REFERENCE:
115	scale = stokes_reference.I
116	elif norm == StokesNorm.NONE:
117	scale = 1
118	else:
119	raise ValueError(f"Did not recognize normalization option {norm}")
120
121	self.add(
122	nu=nu,
123	stokes_I=stokes.I / scale,
124	stokes_V=stokes.V / scale,
125	color=color,
126	label=label,
127	linewidth=linewidth,
128	)
129
130
131	class StokesPlotter_IV_IpmV(PlotterBase): # pragma: no cover
132	r"""
133	Stokes plotter class for Stokes :math:`I, V, I\pm V` profiles.
134	"""
135
136	def __init__(self, title="", use_air_wavelengths=False, reference_lambda_A_air=None):
137	super().__init__(
138	title=title,
139	use_air_wavelengths=use_air_wavelengths,
140	reference_lambda_A_air=reference_lambda_A_air,
141	n_axes=3,
142	y_labels=[r"Stokes $I$", r"Stokes $V$", r"Stokes $(I\pm V)$"],
143	)
144
145	@log_method
146	def add(self, nu, stokes_I, stokes_V, color=None, label="", linewidth=1.5):
147	color = self._next_color(color)
148	wavelength = self._wavelength_axis(nu)
149	self.axs[0].plot(wavelength, stokes_I, label=label, color=color, linewidth=linewidth)
150	self.axs[1].plot(wavelength, stokes_V, label=label, color=color, linewidth=linewidth)
151	self.axs[2].plot(wavelength, stokes_I + stokes_V, "-", label=label + " $I+V$", color=color, linewidth=linewidth)
152	self.axs[2].plot(
153	wavelength, stokes_I - stokes_V, "--", label=label + " $I-V$", color=color, linewidth=linewidth
154	)
155
156	@log_method
157	def add_stokes(
158	self,
159	nu,
160	stokes: Stokes,
161	stokes_reference: Stokes = None,
162	norm: StokesNorm = StokesNorm.NONE,
163	color=None,
164	label="",
165	linewidth=1.5,
166	):
167	color = self._next_color(color)
168	wavelength = self._wavelength_axis(nu)
169
170	if norm == StokesNorm.NONE:
171	I, V = stokes.I, stokes.V
172	elif norm == StokesNorm.MAX_I:
173	scale = np.max(stokes.I)
174	I, V = stokes.I / scale, stokes.V / scale
175	elif norm == StokesNorm.BY_REFERENCE:
176	scale = stokes_reference.I
177	I, V = stokes.I / scale, stokes.V / scale
178	elif norm == StokesNorm.MAX_IpV_ImV:
179	I, V = stokes.I, stokes.V
180	else:
181	raise ValueError(f"Did not recognize normalization option {norm}")
182
183	self.axs[0].plot(wavelength, I, label=label, color=color, linewidth=linewidth)
184	self.axs[1].plot(wavelength, V, label=label, color=color, linewidth=linewidth)
185
186	if norm == StokesNorm.MAX_IpV_ImV:
187	IpV = stokes.I + stokes.V
188	ImV = stokes.I - stokes.V
189	IpV = IpV / np.max(np.abs(IpV))
190	ImV = ImV / np.max(np.abs(ImV))
191	else:
192	IpV = I + V
193	ImV = I - V
194
195	self.axs[2].plot(wavelength, IpV, "-", label=label + " $I+V$", color=color, linewidth=linewidth)
196	self.axs[2].plot(wavelength, ImV, "--", label=label + " $I-V$", color=color, linewidth=linewidth)
197
198
199	class StokesPlotter(PlotterBase): # pragma: no cover
200	r"""
201	Stokes plotter class for Stokes :math:`I, Q, U, V` profiles.
202	"""
203
204	def __init__(
205	self,
206	title="",
207	use_air_wavelengths=False,
208	reference_lambda_A_air=None,
209	x_label=None,
210	y_label_I=r"Stokes $I$",
211	y_label_Q=r"Stokes $Q$",
212	y_label_U=r"Stokes $U$",
213	y_label_V=r"Stokes $V$",
214	):
215	super().__init__(
216	title=title,
217	use_air_wavelengths=use_air_wavelengths,
218	reference_lambda_A_air=reference_lambda_A_air,
219	n_axes=4,
220	y_labels=[y_label_I, y_label_Q, y_label_U, y_label_V],
221	x_label=x_label,
222	)
223
224	@log_method
225	def add(
226	self,
227	nu,
228	stokes_I,
229	stokes_Q,
230	stokes_U,
231	stokes_V,
232	color=None,
233	label="",
234	style="-",
235	linewidth=1.5,
236	):
237	color = self._next_color(color)
238	wavelength = self._wavelength_axis(nu)
239
240	if stokes_I is not None:
241	self.axs[0].plot(wavelength, stokes_I, style, label=label, color=color, linewidth=linewidth)
242	if stokes_Q is not None:
243	self.axs[1].plot(wavelength, stokes_Q, style, label=label, color=color, linewidth=linewidth)
244	if stokes_U is not None:
245	self.axs[2].plot(wavelength, stokes_U, style, label=label, color=color, linewidth=linewidth)
246	if stokes_V is not None:
247	self.axs[3].plot(wavelength, stokes_V, style, label=label, color=color, linewidth=linewidth)
248
249	@log_method
250	def add_stokes(
251	self,
252	nu,
253	stokes: Stokes,
254	stokes_reference: Stokes = None,
255	norm: StokesNorm = StokesNorm.NONE,
256	color=None,
257	label="",
258	style="-",
259	linewidth=1.5,
260	):
261	if norm == StokesNorm.MAX_I:
262	scale = np.max(stokes.I)
263	elif norm == StokesNorm.BY_REFERENCE:
264	scale = stokes_reference.I
265	elif norm == StokesNorm.NONE:
266	scale = 1
267	else:
268	raise ValueError(f"Did not recognize normalization option {norm}")
269	self.add(
270	nu=nu,
271	stokes_I=stokes.I / scale,
272	stokes_Q=stokes.Q / scale,
273	stokes_U=stokes.U / scale,
274	stokes_V=stokes.V / scale,
275	color=color,
276	label=label,
277	style=style,
278	linewidth=linewidth,
279	)

yakovkinii / SolRaT / 24916259816

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