From 63537702ac9e3d6a22476278c21563fe28837af4 Mon Sep 17 00:00:00 2001 From: austin-hoover Date: Thu, 22 Oct 2026 14:33:42 -0400 Subject: [PATCH] Test matrix lattice tunes match fractional tunes from transfer matrix --- tests/py/orbit/test_matrix_lattice.py | 53 +++++++++++++++----- tests/py/orbit/test_teapot_matrix_lattice.py | 31 ------------ 2 files changed, 40 insertions(+), 44 deletions(-) delete mode 100644 tests/py/orbit/test_teapot_matrix_lattice.py diff --git a/tests/py/orbit/test_matrix_lattice.py b/tests/py/orbit/test_matrix_lattice.py index 01104b69..9e6e9cb0 100644 --- a/tests/py/orbit/test_matrix_lattice.py +++ b/tests/py/orbit/test_matrix_lattice.py @@ -1,12 +1,20 @@ from pathlib import Path -from pprint import pprint - -import pytest from orbit.core.bunch import Bunch from orbit.parsers import MADX_Parser from orbit.teapot import TEAPOT_Lattice from orbit.teapot import TEAPOT_MATRIX_Lattice +from orbit.utils.consts import mass_proton + + +def matrix_is_uncoupled(matrix_lattice: TEAPOT_MATRIX_Lattice, atol: float = 1e-12): + for i in range(4): + for j in range(4): + matrix_element = matrix_lattice.oneTurnMatrix.get(i, j) + if (i >= 2 and j < 2) or (j >= 2 and i < 2): + if abs(matrix_element) > atol: + return False + return True def test_sns_ring_eigtunes_coupled(): @@ -25,12 +33,13 @@ def test_sns_ring_eigtunes_coupled(): bunch.getSyncParticle().kinEnergy(1.3) matrix_lattice = TEAPOT_MATRIX_Lattice(lattice, bunch) - ring_params = matrix_lattice.getRingParametersDict() + assert not matrix_is_uncoupled(matrix_lattice) - nux = ring_params["fractional tune x"] - nuy = ring_params["fractional tune y"] - nu1 = ring_params["fractional tune 1"] - nu2 = ring_params["fractional tune 2"] + matrix_params = matrix_lattice.getRingParametersDict() + nux = matrix_params["fractional tune x"] + nuy = matrix_params["fractional tune y"] + nu1 = matrix_params["fractional tune 1"] + nu2 = matrix_params["fractional tune 2"] assert abs(nux - nu1) > 0.01 assert abs(nuy - nu2) > 0.01 @@ -47,13 +56,31 @@ def test_sns_ring_eigtunes_uncoupled(): bunch.getSyncParticle().kinEnergy(1.3) matrix_lattice = TEAPOT_MATRIX_Lattice(lattice, bunch) - ring_params = matrix_lattice.getRingParametersDict() + assert matrix_is_uncoupled(matrix_lattice) - nux = ring_params["fractional tune x"] - nuy = ring_params["fractional tune y"] - nu1 = ring_params["fractional tune 1"] - nu2 = ring_params["fractional tune 2"] + matrix_params = matrix_lattice.getRingParametersDict() + nux = matrix_params["fractional tune x"] + nuy = matrix_params["fractional tune y"] + nu1 = matrix_params["fractional tune 1"] + nu2 = matrix_params["fractional tune 2"] assert abs(nux - nu1) < 1e-8 assert abs(nuy - nu2) < 1e-8 +def test_frac_tunes_match_matrix_params(): + madx_file = Path(__file__).parent / "inputs" / "sns_ring.lat" + + lattice = TEAPOT_Lattice() + lattice.readMADX(str(madx_file), "rnginjsol") + lattice.initialize() + + bunch = Bunch() + bunch.mass(mass_proton) + bunch.getSyncParticle().kinEnergy(1.3) + + matrix_lattice = TEAPOT_MATRIX_Lattice(lattice, bunch) + matrix_params = matrix_lattice.getRingParametersDict() + tune_x, tune_y = matrix_lattice.getTunes() + assert matrix_is_uncoupled(matrix_lattice) + assert abs(tune_x % 1.0 - matrix_params["fractional tune x"]) < 1e-8 + assert abs(tune_y % 1.0 - matrix_params["fractional tune y"]) < 1e-8 diff --git a/tests/py/orbit/test_teapot_matrix_lattice.py b/tests/py/orbit/test_teapot_matrix_lattice.py deleted file mode 100644 index 6d13a8a7..00000000 --- a/tests/py/orbit/test_teapot_matrix_lattice.py +++ /dev/null @@ -1,31 +0,0 @@ -from pathlib import Path -from pprint import pprint - -from orbit.core.bunch import Bunch -from orbit.parsers import MADX_Parser -from orbit.teapot import TEAPOT_Lattice -from orbit.teapot import TEAPOT_MATRIX_Lattice -from orbit.utils.consts import mass_proton - -madx_file = Path(__file__).parent / "inputs" / "sns_ring.lat" - -lattice = TEAPOT_Lattice() -lattice.readMADX(str(madx_file), "rnginjsol") -lattice.initialize() - -bunch = Bunch() -bunch.mass(mass_proton) -bunch.getSyncParticle().kinEnergy(1.3) - -matrix_lattice = TEAPOT_MATRIX_Lattice(lattice, bunch) -ring_params = matrix_lattice.getRingParametersDict() - -pprint(ring_params) - -tune_x, tune_y = matrix_lattice.getTunes() - -print("tune_x (from integration):", tune_x) -print("tune_y (from integration):", tune_y) -print("frac_tune_x (from one-turn matrix):", ring_params["fractional tune x"]) -print("frac_tune_y (from one-turn matrix):", ring_params["fractional tune y"]) -print()