23865392830

Committed 01 Apr 2026 06:53PM UTC coverage: 97.562% (-0.3%) from 97.852%

Build # 23865392830

Build Type

Pull #8

github

Committed by

web-flow

Commit Message

Merge 09564b5e0 into 70a6e3a94

Pull Request Pull Request #8: Refactoring: add thread-safe logging; add typevars.

Coverage Stats

210 of 214 branches covered (98.13%)

Branch coverage included in aggregate %.

2391 of 2452 relevant lines covered (97.51%)

0.98 hits per line

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

98.53

solrat/atom_model/multi_term_atom_model/object/rho_matrix_builder.py

from typing import List

import numpy as np
import pandas as pd
from numpy import sqrt

from solrat.atom_model.base_atom_model.object.rho import BaseRho
from solrat.atom_model.multi_term_atom_model.object.level_registry import Term
from solrat.engine.functions.decorators import log_method
from solrat.engine.functions.general import half_int_to_str
from solrat.engine.functions.looping import PROJECTION, TRIANGULAR
from solrat.engine.generators.nested_loops import nested_loops


def construct_coherence_id(term: Term, K: float, Q: float, J: float, Jʹ: float):
    """
    Construct a unique ID for a coherence
    """
    return construct_coherence_id_from_term_id(term_id=term.term_id, K=K, Q=Q, J=J, Jʹ=Jʹ)


def construct_coherence_id_from_term_id(term_id: str, K: float, Q: float, J: float, Jʹ: float):
    """
    Construct a unique ID for a coherence
    """
    return f"{term_id}_K={half_int_to_str(K)}_Q={half_int_to_str(Q)}_J={half_int_to_str(J)}_Jʹ={half_int_to_str(Jʹ)}"


class Rho(BaseRho):
    def __init__(self, terms: List[Term]):
        """
        A container for the density tensor Rho
        :param terms: list of all terms.
        """
        self.data = dict()
        self.terms = {term.term_id: term for term in terms}

    def set_from_term_id(self, term_id: str, K: float, Q: float, J: float, Jʹ: float, value: complex):
        """
        Set the value
        """
        self.data[construct_coherence_id_from_term_id(term_id=term_id, K=K, Q=Q, J=J, Jʹ=Jʹ)] = value

    def __call__(self, K: float, Q: float, J: float, Jʹ: float, term_id: str) -> np.complex128:
        """
        Get the value
        """
        coherence_id = construct_coherence_id_from_term_id(term_id=term_id, K=K, Q=Q, J=J, Jʹ=Jʹ)
        return self.data[coherence_id]


class RhoMatrixBuilder:
    def __init__(self, terms: List[Term]):
        """
        This class helps to build the matrix for rhos.
        All possible rhos are defined by terms.

        :param terms: list of all terms.
        """

        # Create mapping [term_id, K, Q, J, Jʹ] <-> matrix index
        self.index_to_parameters = dict()
        self.coherence_id_to_index = dict()
        self.trace_indexes = []
        self.trace_weights = []
        index = 0
        for term in terms:
            for J, Jʹ, K, Q in nested_loops(
                J=TRIANGULAR(term.L, term.S),
                Jʹ=TRIANGULAR(term.L, term.S),
                K=TRIANGULAR("J", "Jʹ"),
                Q=PROJECTION("K"),
            ):
                coherence_id = construct_coherence_id(term=term, K=K, Q=Q, J=J, Jʹ=Jʹ)
                self.coherence_id_to_index[coherence_id] = index
                self.index_to_parameters[index] = (term.term_id, K, Q, J, Jʹ)
                if K == 0 and Q == 0 and J == Jʹ:
                    self.trace_indexes.append(index)
                    self.trace_weights.append(sqrt(2 * J + 1))
                index += 1

        # create the matrix
        matrix_size = index
        self.rho_matrix = np.zeros((matrix_size, matrix_size), dtype=np.complex128)
        self.selected_coherence = None

    @log_method
    def reset_matrix(self):
        """
        Reset the matrix to fill it from scratch later.
        """
        self.rho_matrix = self.rho_matrix * 0

    def select_equation(self, term: Term, K: int, Q: int, J: float, Jʹ: float):
        r"""
        Selects the equation to add coefficients to.

        The equation is of a form :math:`M_{ij} \rho_j = 0`. Here we select i.
        """
        coherence_id = construct_coherence_id(term=term, K=K, Q=Q, J=J, Jʹ=Jʹ)
        self.selected_coherence = coherence_id

    def add_coefficient(self, term: Term, K: int, Q: int, J: float, Jʹ: float, coefficient: complex):
        r"""
        Adds a coefficient to the selected equation.

        The equation is of a form :math:`M_{ij} \rho_j = 0`. We have already selected i.
        Now we do :math:`M_{ij} += coefficient`, where j is defined by {term, :math:`K, Q, J, Jʹ`}.
        """
        coefficient = complex(coefficient)
        assert isinstance(coefficient, complex), "Coefficient must be complex"
        coherence_id = construct_coherence_id(term=term, K=K, Q=Q, J=J, Jʹ=Jʹ)

        if coefficient == 0:
            return

        index0 = self.coherence_id_to_index[self.selected_coherence]
        if coherence_id not in self.coherence_id_to_index.keys():
            raise ValueError(f"Trying to add coefficient to non-existing coherence {coherence_id}")
        index1 = self.coherence_id_to_index[coherence_id]
        self.rho_matrix[index0, index1] += coefficient

    @log_method
    def add_coefficient_from_df(self, df: pd.DataFrame):
        r"""
        Adds a coefficient to the selected equation from the dataframe of the form "index0", "index1", "coefficient".

        The equation is of a :math:`M_{ij} \rho_j = 0`.
        For each df row, we do :math:`M_{ij} += coefficient`, where i=index0 and j=index1.
        """
        df = df[["index0", "index1", "coefficient"]].groupby(["index0", "index1"]).sum().reset_index()
        self.rho_matrix[df.index0, df.index1] += df.coefficient

1	from typing import List	1✔
2
3	import numpy as np	1✔
4	import pandas as pd	1✔
5	from numpy import sqrt	1✔
6
7	from solrat.atom_model.base_atom_model.object.rho import BaseRho	1✔
8	from solrat.atom_model.multi_term_atom_model.object.level_registry import Term	1✔
9	from solrat.engine.functions.decorators import log_method	1✔
10	from solrat.engine.functions.general import half_int_to_str	1✔
11	from solrat.engine.functions.looping import PROJECTION, TRIANGULAR	1✔
12	from solrat.engine.generators.nested_loops import nested_loops	1✔
13
14
15	def construct_coherence_id(term: Term, K: float, Q: float, J: float, Jʹ: float):	1✔
16	"""
17	Construct a unique ID for a coherence
18	"""
19	return construct_coherence_id_from_term_id(term_id=term.term_id, K=K, Q=Q, J=J, Jʹ=Jʹ)	1✔
20
21
22	def construct_coherence_id_from_term_id(term_id: str, K: float, Q: float, J: float, Jʹ: float):	1✔
23	"""
24	Construct a unique ID for a coherence
25	"""
26	return f"{term_id}_K={half_int_to_str(K)}_Q={half_int_to_str(Q)}_J={half_int_to_str(J)}_Jʹ={half_int_to_str(Jʹ)}"	1✔
27
28
29	class Rho(BaseRho):	1✔
30	def __init__(self, terms: List[Term]):	1✔
31	"""
32	A container for the density tensor Rho
33	:param terms: list of all terms.
34	"""
35	self.data = dict()	1✔
36	self.terms = {term.term_id: term for term in terms}	1✔
37
38	def set_from_term_id(self, term_id: str, K: float, Q: float, J: float, Jʹ: float, value: complex):	1✔
39	"""
40	Set the value
41	"""
42	self.data[construct_coherence_id_from_term_id(term_id=term_id, K=K, Q=Q, J=J, Jʹ=Jʹ)] = value	1✔
43
44	def __call__(self, K: float, Q: float, J: float, Jʹ: float, term_id: str) -> np.complex128:	1✔
45	"""
46	Get the value
47	"""
48	coherence_id = construct_coherence_id_from_term_id(term_id=term_id, K=K, Q=Q, J=J, Jʹ=Jʹ)	1✔
49	return self.data[coherence_id]	1✔
50
51
52	class RhoMatrixBuilder:	1✔
53	def __init__(self, terms: List[Term]):	1✔
54	"""
55	This class helps to build the matrix for rhos.
56	All possible rhos are defined by terms.
57
58	:param terms: list of all terms.
59	"""
60
61	# Create mapping [term_id, K, Q, J, Jʹ] <-> matrix index
62	self.index_to_parameters = dict()	1✔
63	self.coherence_id_to_index = dict()	1✔
64	self.trace_indexes = []	1✔
65	self.trace_weights = []	1✔
66	index = 0	1✔
67	for term in terms:	1✔
68	for J, Jʹ, K, Q in nested_loops(	1✔
69	J=TRIANGULAR(term.L, term.S),
70	Jʹ=TRIANGULAR(term.L, term.S),
71	K=TRIANGULAR("J", "Jʹ"),
72	Q=PROJECTION("K"),
73	):
74	coherence_id = construct_coherence_id(term=term, K=K, Q=Q, J=J, Jʹ=Jʹ)	1✔
75	self.coherence_id_to_index[coherence_id] = index	1✔
76	self.index_to_parameters[index] = (term.term_id, K, Q, J, Jʹ)	1✔
77	if K == 0 and Q == 0 and J == Jʹ:	1✔
78	self.trace_indexes.append(index)	1✔
79	self.trace_weights.append(sqrt(2 * J + 1))	1✔
80	index += 1	1✔
81
82	# create the matrix
83	matrix_size = index	1✔
84	self.rho_matrix = np.zeros((matrix_size, matrix_size), dtype=np.complex128)	1✔
85	self.selected_coherence = None	1✔
86
87	@log_method	1✔
88	def reset_matrix(self):	1✔
89	"""
90	Reset the matrix to fill it from scratch later.
91	"""
92	self.rho_matrix = self.rho_matrix * 0	1✔
93
94	def select_equation(self, term: Term, K: int, Q: int, J: float, Jʹ: float):	1✔
95	r"""
96	Selects the equation to add coefficients to.
97
98	The equation is of a form :math:`M_{ij} \rho_j = 0`. Here we select i.
99	"""
100	coherence_id = construct_coherence_id(term=term, K=K, Q=Q, J=J, Jʹ=Jʹ)	1✔
101	self.selected_coherence = coherence_id	1✔
102
103	def add_coefficient(self, term: Term, K: int, Q: int, J: float, Jʹ: float, coefficient: complex):	1✔
104	r"""
105	Adds a coefficient to the selected equation.
106
107	The equation is of a form :math:`M_{ij} \rho_j = 0`. We have already selected i.
108	Now we do :math:`M_{ij} += coefficient`, where j is defined by {term, :math:`K, Q, J, Jʹ`}.
109	"""
110	coefficient = complex(coefficient)	1✔
111	assert isinstance(coefficient, complex), "Coefficient must be complex"	1✔
112	coherence_id = construct_coherence_id(term=term, K=K, Q=Q, J=J, Jʹ=Jʹ)	1✔
113
114	if coefficient == 0:	1✔
115	return	1✔
116
117	index0 = self.coherence_id_to_index[self.selected_coherence]	1✔
118	if coherence_id not in self.coherence_id_to_index.keys():	1✔
119	raise ValueError(f"Trying to add coefficient to non-existing coherence {coherence_id}")	×
120	index1 = self.coherence_id_to_index[coherence_id]	1✔
121	self.rho_matrix[index0, index1] += coefficient	1✔
122
123	@log_method	1✔
124	def add_coefficient_from_df(self, df: pd.DataFrame):	1✔
125	r"""
126	Adds a coefficient to the selected equation from the dataframe of the form "index0", "index1", "coefficient".
127
128	The equation is of a :math:`M_{ij} \rho_j = 0`.
129	For each df row, we do :math:`M_{ij} += coefficient`, where i=index0 and j=index1.
130	"""
131	df = df[["index0", "index1", "coefficient"]].groupby(["index0", "index1"]).sum().reset_index()	1✔
132	self.rho_matrix[df.index0, df.index1] += df.coefficient	1✔

yakovkinii / SolRaT / 23865392830

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