def test_TransitionQuery():
# direct initialization
test0 = {"ch1": {"P": [-1, -1]}, "ch2": {"P": [-1], "D": [0]}}
obj1 = TransitionQuery(**test0)
assert obj1.name is None
assert obj1.description is None
assert obj1.label is None
assert obj1.ch1 == SymmetryQuery(
P=[-1, -1]), "TransitionQuery ch1 not equal."
assert obj1.ch2 == SymmetryQuery(P=[-1],
D=[0]), "TransitionQuery ch2 not equal."
assert obj1.ch3 is None, "TransitionQuery ch3 not equal."
test1 = {"ch2": {"P": [-2], "D": [2]}}
obj2 = TransitionQuery(**test1)
assert obj2.name is None
assert obj2.description is None
assert obj2.label is None
assert obj2.ch1 == SymmetryQuery(P=[-1]), "TransitionQuery ch1 not equal."
assert obj2.ch2 == SymmetryQuery(P=[-2],
D=[2]), "TransitionQuery ch2 not equal."
assert obj2.ch3 is None, "TransitionQuery ch3 not equal."
test2 = {"ch1": None, "ch2": {"P": [-1, 1], "D": [2]}}
obj3 = TransitionQuery(**test2)
assert obj3.name is None
assert obj3.description is None
assert obj3.label is None
assert obj3.ch1 is None, "TransitionQuery ch1 not equal."
assert obj3.ch2 == SymmetryQuery(P=[-1, 1],
D=[2]), "TransitionQuery ch2 not equal."
assert obj3.ch3 is None, "TransitionQuery ch3 not equal."
def test_5Q_VAS_general():
"""5Q-VAS method test"""
mth = FiveQ_VAS(
channels=["17O"], spectral_dimensions=[SpectralDimension(), SpectralDimension()]
)
assert mth.name == "FiveQ_VAS"
assert mth.description == "Simulate a 5Q variable-angle spinning spectrum."
assert mth.spectral_dimensions[0].events[0].transition_queries == [
TransitionQuery(ch1={"P": [-5], "D": [0]})
]
assert mth.spectral_dimensions[1].events[0].transition_queries == [
TransitionQuery(ch1={"P": [-1], "D": [0]})
]
assert FiveQ_VAS.parse_dict_with_units(mth.json()) == mth
assert np.allclose(
mth.affine_matrix, [0.3243243243243243, 0.6756756756756757, 0.0, 1.0]
)
serialize = mth.json()
_ = serialize.pop("affine_matrix")
assert serialize == {
"channels": ["17O"],
"description": "Simulate a 5Q variable-angle spinning spectrum.",
"name": "FiveQ_VAS",
**sample_test_output(-5),
}
def test_7Q_VAS_general():
"""7Q-VAS method test"""
mth = SevenQ_VAS(channels=["51V"], spectral_dimensions=[{}, {}])
assert mth.name == "SevenQ_VAS"
assert mth.description == "Simulate a 7Q variable-angle spinning spectrum."
assert mth.spectral_dimensions[0].events[0].transition_queries == [
TransitionQuery(ch1={"P": [-7], "D": [0]})
]
assert mth.spectral_dimensions[1].events[0].transition_queries == [
TransitionQuery(ch1={"P": [-1], "D": [0]})
]
assert SevenQ_VAS.parse_dict_with_units(mth.json()) == mth
assert np.allclose(mth.affine_matrix, [0.2184466, 0.7815534, 0.0, 1.0])
serialize = mth.json()
_ = serialize.pop("affine_matrix")
assert serialize == {
"channels": ["51V"],
"description": "Simulate a 7Q variable-angle spinning spectrum.",
"name": "SevenQ_VAS",
**sample_test_output(-7),
}
def test_3Q_VAS_general():
"""3Q-VAS method test"""
mth = ThreeQ_VAS(channels=["87Rb"], spectral_dimensions=[{}, {}])
assert mth.name == "ThreeQ_VAS"
assert mth.description == "Simulate a 3Q variable-angle spinning spectrum."
assert mth.spectral_dimensions[0].events[0].transition_query == [
TransitionQuery(ch1={
"P": [-3],
"D": [0]
})
]
assert mth.spectral_dimensions[1].events[0].transition_query == [
TransitionQuery(ch1={
"P": [-1],
"D": [0]
})
]
assert ThreeQ_VAS.parse_dict_with_units(mth.json()) == mth
assert np.allclose(mth.affine_matrix, [0.5625, 0.4375, 0.0, 1.0])
serialize = mth.json()
_ = serialize.pop("affine_matrix")
assert serialize == {
"channels": ["87Rb"],
"description": "Simulate a 3Q variable-angle spinning spectrum.",
"name": "ThreeQ_VAS",
**sample_test_output(-3),
}
def test_ST1_VAS_general():
"""Inner satellite-transition variable-angle spinning method"""
mth = ST1_VAS(
channels=["87Rb"],
magnetic_flux_density=11.7, # in T
spectral_dimensions=[
{"count": 1024, "spectral_width": 3e4},
{"count": 1024, "spectral_width": 2e4},
],
)
assert mth.name == "ST1_VAS"
des = (
"Simulate a 1.5 -> 0.5 and -0.5 -> -1.5 satellite-transition variable-angle "
"spinning spectrum."
)
assert mth.description == des
assert mth.spectral_dimensions[0].events[0].transition_queries == [
TransitionQuery(ch1={"P": [-1], "D": [2]}),
TransitionQuery(ch1={"P": [-1], "D": [-2]}),
]
assert mth.spectral_dimensions[1].events[0].transition_queries == [
TransitionQuery(ch1={"P": [-1], "D": [0]})
]
assert ST1_VAS.parse_dict_with_units(mth.json()) == mth
assert np.allclose(mth.affine_matrix, [0.52941176, 0.47058824, 0.0, 1.0])
serialize = mth.json()
_ = serialize.pop("affine_matrix")
assert serialize == {
"channels": ["87Rb"],
"description": des,
"magnetic_flux_density": "11.7 T",
"name": "ST1_VAS",
"rotor_angle": "0.9553166181245 rad",
"rotor_frequency": "1000000000000.0 Hz",
"spectral_dimensions": [
{
"count": 1024,
"spectral_width": "30000.0 Hz",
"events": [
{
"transition_queries": [
{"ch1": {"P": [-1], "D": [i]}} for i in [2, -2]
]
}
],
},
{
"count": 1024,
"spectral_width": "20000.0 Hz",
"events": [{"transition_queries": [{"ch1": {"P": [-1], "D": [0]}}]}],
},
],
}
def test_ST2_VAS_general():
"""Second to inner satellite-transition variable-angle spinning method"""
mth = ST2_VAS(
channels=["17O"],
magnetic_flux_density=9.4, # in T
spectral_dimensions=[
{
"count": 1024,
"spectral_width": 5e4
},
{
"count": 1024,
"spectral_width": 5e4
},
],
)
assert mth.name == "ST2_VAS"
des = (
"Simulate a 2.5 -> 1.5 and -1.5 -> -2.5 satellite-transition variable-angle "
"spinning spectrum.")
assert mth.description == des
assert mth.spectral_dimensions[0].events[0].transition_query == [
TransitionQuery(ch1={
"P": [-1],
"D": [4]
}),
TransitionQuery(ch1={
"P": [-1],
"D": [-4]
}),
]
assert mth.spectral_dimensions[1].events[0].transition_query == [
TransitionQuery(ch1={
"P": [-1],
"D": [0]
})
]
assert ST2_VAS.parse_dict_with_units(mth.json()) == mth
assert np.allclose(mth.affine_matrix, [0.35294118, 0.64705882, 0.0, 1.0])
serialize = mth.json()
_ = serialize.pop("affine_matrix")
assert serialize == {
"channels": ["17O"],
"description": des,
"name": "ST2_VAS",
**sample_test_output(4),
}
def test_SSB_general():
"""Inner satellite-transition variable-angle spinning method"""
mth = SSB2D(
channels=["87Rb"],
magnetic_flux_density=9.4, # in T
rotor_frequency=1200,
spectral_dimensions=[
{"count": 1024, "spectral_width": 5e4},
{"count": 1024, "spectral_width": 5e4},
],
)
assert mth.name == "SSB2D"
assert mth.description == "Simulate a 2D sideband separation method."
# test transition query
tq = TransitionQuery(ch1={"P": [-1], "D": [0]})
assert mth.spectral_dimensions[0].events[0].transition_query[0] == tq
assert mth.spectral_dimensions[1].events[0].transition_query[0] == tq
# test rotor_frequency
assert mth.spectral_dimensions[0].events[0].rotor_frequency == 1200
assert mth.spectral_dimensions[1].events[0].rotor_frequency == 1e12
# check serialization
assert SSB2D.parse_dict_with_units(mth.json()) == mth
assert np.allclose(mth.affine_matrix, [1, -1, 0.0, 1.0])
serialize = mth.json()
_ = serialize.pop("affine_matrix")
assert serialize == {
"channels": ["87Rb"],
"description": mth.description,
"magnetic_flux_density": "9.4 T",
"name": "SSB2D",
"rotor_angle": "0.9553166181245 rad",
"rotor_frequency": "1200.0 Hz",
"spectral_dimensions": [
{
"count": 1024,
"spectral_width": "50000.0 Hz",
"events": [{"transition_query": [{"ch1": {"P": [-1], "D": [0]}}]}],
},
{
"count": 1024,
"events": [
{
"rotor_frequency": "1000000000000.0 Hz",
"transition_query": [{"ch1": {"P": [-1], "D": [0]}}],
}
],
"spectral_width": "50000.0 Hz",
},
],
}
def check_equal(query, isotopes, channels, desired):
test = TransitionQuery(**query).combination(isotopes, channels)
for element in ["P", "D"]:
assert len(test[element]) == len(desired[element])
for test_item in test[element]:
assert np.any([
np.allclose(test_item, desired_item, equal_nan=True)
for desired_item in desired[element]
]), f"Error in element {element}"