25191188779

Committed 30 Apr 2026 09:53PM UTC coverage: 95.704% (-1.1%) from 96.759%

Build # 25191188779

Build Type

Pull #16

github

Committed by

web-flow

Commit Message

Merge 7e267c2c9 into 90cd517a0

Pull Request Pull Request #16: Use Frame engine in SEE

Coverage Stats

205 of 221 branches covered (92.76%)

Branch coverage included in aggregate %.

2580 of 2689 relevant lines covered (95.95%)

0.96 hits per line

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

58.82

solrat/atom_model/multi_term_atom_model/object/precomputed_data.py

import os
from dataclasses import dataclass, field
from typing import Union

import pandas as pd

from solrat.engine.generators.merge_frame import Frame, FrameFactor


def _frame_to_csv(frame: Frame, path: str) -> None:
    """Save a Frame to CSV by evaluating all unmerged factors into a single 'coefficient' column."""
    f = frame.copy()
    # Merge any unmerged factors
    for name in list(f.factors.keys()):
        if not f.factors[name].merged:
            f.merge_factor(name)
    # Combine multiple merged factors into one column
    if len(f.factors) > 1:
        combined_name = f.combine_all_merged_factors()
    else:
        combined_name = next(iter(f.factors)) if f.factors else None
    # Rename to 'coefficient' if needed
    if combined_name is not None and combined_name != "coefficient":
        f.frame = f.frame.rename(columns={combined_name: "coefficient"})
    f.frame.to_csv(path, index=False)


def _frame_from_csv(path: str) -> Frame:
    """Load a Frame from CSV. The 'coefficient' column is registered as a merged factor."""
    df = pd.read_csv(path)
    frame = Frame(base_frame=df)
    loop_cols = [c for c in df.columns if c != "coefficient"]
    frame.factors["coefficient"] = FrameFactor(
        name="coefficient",
        dependencies=loop_cols,
        merged=True,
    )
    return frame


_FILE_MAP = {
    "coherence_decay_frame": "coherence_decay_n0.csv",
    "coherence_decay_frame_n_1": "coherence_decay_n1.csv",
    "absorption_frame": "absorption.csv",
    "emission_e_frame": "emission_e.csv",
    "emission_s_frame": "emission_s.csv",
    "relaxation_e_frame": "relaxation_e.csv",
    "relaxation_a_frame": "relaxation_a.csv",
    "relaxation_s_frame": "relaxation_s.csv",
}


@dataclass
class PrecomputedData:
    """
    Container for precomputed atom-specific SEE frames.

    Frames that do not depend on the radiation tensor or atmosphere parameters can be
    precomputed once, saved to CSV via :meth:`save_to_directory`, and reloaded via
    :meth:`load_from_directory`.  Frames that still require the radiation tensor (absorption,
    emission_s, relaxation_a, relaxation_s) are stored with the atom-specific factor already
    evaluated; the radiation-tensor factor is applied per :meth:`fill_all_equations` call.
    """

    coherence_decay_frame: Frame
    absorption_frame: Frame
    emission_e_frame: Frame
    emission_s_frame: Frame
    relaxation_e_frame: Frame
    relaxation_a_frame: Frame
    relaxation_s_frame: Frame
    coherence_decay_frame_n_1: Union[Frame, None] = field(default=None)

    def save_to_directory(self, directory: str) -> None:
        """Save all frames to CSV files in *directory*."""
        os.makedirs(directory, exist_ok=True)
        for attr, filename in _FILE_MAP.items():
            frame = getattr(self, attr, None)
            if frame is not None:
                _frame_to_csv(frame, os.path.join(directory, filename))

    @classmethod
    def load_from_directory(cls, directory: str) -> "PrecomputedData":
        """Load all frames from CSV files in *directory*."""

        def _load(filename: str) -> Frame:
            return _frame_from_csv(os.path.join(directory, filename))

        n1_path = os.path.join(directory, _FILE_MAP["coherence_decay_frame_n_1"])
        return cls(
            coherence_decay_frame=_load(_FILE_MAP["coherence_decay_frame"]),
            absorption_frame=_load(_FILE_MAP["absorption_frame"]),
            emission_e_frame=_load(_FILE_MAP["emission_e_frame"]),
            emission_s_frame=_load(_FILE_MAP["emission_s_frame"]),
            relaxation_e_frame=_load(_FILE_MAP["relaxation_e_frame"]),
            relaxation_a_frame=_load(_FILE_MAP["relaxation_a_frame"]),
            relaxation_s_frame=_load(_FILE_MAP["relaxation_s_frame"]),
            coherence_decay_frame_n_1=(
                _load(_FILE_MAP["coherence_decay_frame_n_1"]) if os.path.exists(n1_path) else None
            ),
        )

1	import os	1✔
2	from dataclasses import dataclass, field	1✔
3	from typing import Union	1✔
4
5	import pandas as pd	1✔
6
7	from solrat.engine.generators.merge_frame import Frame, FrameFactor	1✔
8
9
10	def _frame_to_csv(frame: Frame, path: str) -> None:	1✔
11	"""Save a Frame to CSV by evaluating all unmerged factors into a single 'coefficient' column."""
12	f = frame.copy()	×
13	# Merge any unmerged factors
14	for name in list(f.factors.keys()):	×
15	if not f.factors[name].merged:	×
16	f.merge_factor(name)	×
17	# Combine multiple merged factors into one column
18	if len(f.factors) > 1:	×
19	combined_name = f.combine_all_merged_factors()	×
20	else:
21	combined_name = next(iter(f.factors)) if f.factors else None	×
22	# Rename to 'coefficient' if needed
23	if combined_name is not None and combined_name != "coefficient":	×
24	f.frame = f.frame.rename(columns={combined_name: "coefficient"})	×
25	f.frame.to_csv(path, index=False)	×
26
27
28	def _frame_from_csv(path: str) -> Frame:	1✔
29	"""Load a Frame from CSV. The 'coefficient' column is registered as a merged factor."""
30	df = pd.read_csv(path)	1✔
31	frame = Frame(base_frame=df)	1✔
32	loop_cols = [c for c in df.columns if c != "coefficient"]	1✔
33	frame.factors["coefficient"] = FrameFactor(	1✔
34	name="coefficient",
35	dependencies=loop_cols,
36	merged=True,
37	)
38	return frame	1✔
39
40
41	_FILE_MAP = {	1✔
42	"coherence_decay_frame": "coherence_decay_n0.csv",
43	"coherence_decay_frame_n_1": "coherence_decay_n1.csv",
44	"absorption_frame": "absorption.csv",
45	"emission_e_frame": "emission_e.csv",
46	"emission_s_frame": "emission_s.csv",
47	"relaxation_e_frame": "relaxation_e.csv",
48	"relaxation_a_frame": "relaxation_a.csv",
49	"relaxation_s_frame": "relaxation_s.csv",
50	}
51
52
53	@dataclass	1✔
54	class PrecomputedData:	1✔
55	"""
56	Container for precomputed atom-specific SEE frames.
57
58	Frames that do not depend on the radiation tensor or atmosphere parameters can be
59	precomputed once, saved to CSV via :meth:`save_to_directory`, and reloaded via
60	:meth:`load_from_directory`. Frames that still require the radiation tensor (absorption,
61	emission_s, relaxation_a, relaxation_s) are stored with the atom-specific factor already
62	evaluated; the radiation-tensor factor is applied per :meth:`fill_all_equations` call.
63	"""
64
65	coherence_decay_frame: Frame	1✔
66	absorption_frame: Frame	1✔
67	emission_e_frame: Frame	1✔
68	emission_s_frame: Frame	1✔
69	relaxation_e_frame: Frame	1✔
70	relaxation_a_frame: Frame	1✔
71	relaxation_s_frame: Frame	1✔
72	coherence_decay_frame_n_1: Union[Frame, None] = field(default=None)	1✔
73
74	def save_to_directory(self, directory: str) -> None:	1✔
75	"""Save all frames to CSV files in directory."""
76	os.makedirs(directory, exist_ok=True)	×
77	for attr, filename in _FILE_MAP.items():	×
78	frame = getattr(self, attr, None)	×
79	if frame is not None:	×
80	_frame_to_csv(frame, os.path.join(directory, filename))	×
81
82	@classmethod	1✔
83	def load_from_directory(cls, directory: str) -> "PrecomputedData":	1✔
84	"""Load all frames from CSV files in directory."""
85
86	def _load(filename: str) -> Frame:	1✔
87	return _frame_from_csv(os.path.join(directory, filename))	1✔
88
89	n1_path = os.path.join(directory, _FILE_MAP["coherence_decay_frame_n_1"])	1✔
90	return cls(	1✔
91	coherence_decay_frame=_load(_FILE_MAP["coherence_decay_frame"]),
92	absorption_frame=_load(_FILE_MAP["absorption_frame"]),
93	emission_e_frame=_load(_FILE_MAP["emission_e_frame"]),
94	emission_s_frame=_load(_FILE_MAP["emission_s_frame"]),
95	relaxation_e_frame=_load(_FILE_MAP["relaxation_e_frame"]),
96	relaxation_a_frame=_load(_FILE_MAP["relaxation_a_frame"]),
97	relaxation_s_frame=_load(_FILE_MAP["relaxation_s_frame"]),
98	coherence_decay_frame_n_1=(
99	_load(_FILE_MAP["coherence_decay_frame_n_1"]) if os.path.exists(n1_path) else None
100	),
101	)

yakovkinii / SolRaT / 25191188779

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