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
[docs]
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ʹ)
[docs]
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ʹ)}"
[docs]
class Rho(BaseRho):
"""
A container for the density tensor Rho
:param terms: list of all terms.
"""
def __init__(self, terms: List[Term]):
self.data = dict()
self.terms = {term.term_id: term for term in terms}
[docs]
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]
[docs]
class RhoMatrixBuilder:
"""
This class helps to build the matrix for rhos.
All possible rhos are defined by terms.
:param terms: list of all terms.
"""
def __init__(self, terms: List[Term]):
# 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
[docs]
@log_method
def reset_matrix(self):
"""
Reset the matrix to fill it from scratch later.
"""
self.rho_matrix = self.rho_matrix * 0
[docs]
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
[docs]
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
[docs]
@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