From 436a20b9183830b8e39c046160531b5dfb5970bc Mon Sep 17 00:00:00 2001 From: Phil Smith Date: Thu, 25 Sep 2025 17:17:37 +0000 Subject: [PATCH 1/9] Allow toggling of additional simulator behaviour Allow selecting linopt function, disabling chromaticity aand disabling radiation --- src/atip/load_sim.py | 60 ++++++++++++++++++++---- src/atip/simulator.py | 104 +++++++++++++++++++++++++++++++++++------- src/atip/utils.py | 19 +++++++- 3 files changed, 154 insertions(+), 29 deletions(-) diff --git a/src/atip/load_sim.py b/src/atip/load_sim.py index 698f66e..2c00ab2 100644 --- a/src/atip/load_sim.py +++ b/src/atip/load_sim.py @@ -13,7 +13,13 @@ def load_from_filepath( - pytac_lattice, at_lattice_filepath, callback=None, disable_emittance=False + pytac_lattice, + at_lattice_filepath, + linopt_function="linopt6", + disable_emittance=False, + disable_chromaticity=False, + disable_radiation=False, + callback=None, ): """Load simulator data sources onto the lattice and its elements. @@ -21,9 +27,15 @@ def load_from_filepath( pytac_lattice (pytac.lattice.Lattice): An instance of a Pytac lattice. at_lattice_filepath (str): The path to a .mat file from which the Accelerator Toolbox lattice can be loaded. + linopt_function (str): Which pyAT linear optics function to use: linopt2, + linopt4, linopt6. + disable_emittance (bool): Whether the emittance calculations should be + disabled. + disable_chromaticity (bool): Whether the chromaticity calculations should be + disabled. + disable_radiation (bool): Whether radiation calculations should be disabled. callback (typing.Callable): To be called after completion of each round of - physics calculations. - disable_emittance (bool): Whether the emittance should be calculated. + physics calculations. Returns: pytac.lattice.Lattice: The same Pytac lattice object, but now with a @@ -34,19 +46,40 @@ def load_from_filepath( name=pytac_lattice.name, energy=pytac_lattice.get_value("energy", units=pytac.PHYS), ) - return load(pytac_lattice, at_lattice, callback, disable_emittance) + return load( + pytac_lattice, + at_lattice, + linopt_function, + disable_emittance, + disable_chromaticity, + disable_radiation, + callback, + ) -def load(pytac_lattice, at_lattice, callback=None, disable_emittance=False): +def load( + pytac_lattice, + at_lattice, + linopt_function="linopt6", + disable_emittance=False, + disable_chromaticity=False, + disable_radiation=False, + callback=None, +): """Load simulator data sources onto the lattice and its elements. Args: pytac_lattice (pytac.lattice.Lattice): An instance of a Pytac lattice. - at_lattice (at.lattice_object.Lattice): An instance of an Accelerator - Toolbox lattice object. + at_lattice (at.lattice_object.Lattice): An instance of an AT lattice object. + linopt_function (str): Which pyAT linear optics function to use: linopt2, + linopt4, linopt6. + disable_emittance (bool): Whether the emittance calculations should be + disabled. + disable_chromaticity (bool): Whether the chromaticity calculations should be + disabled. + disable_radiation (bool): Whether radiation calculations should be disabled. callback (typing.Callable): To be called after completion of each round of - physics calculations. - disable_emittance (bool): Whether the emittance should be calculated. + physics calculations. Returns: pytac.lattice.Lattice: The same Pytac lattice object, but now with a @@ -58,7 +91,14 @@ def load(pytac_lattice, at_lattice, callback=None, disable_emittance=False): f"(AT:{len(at_lattice)} Pytac:{len(pytac_lattice)})." ) # Initialise an instance of the ATSimulator Object. - atsim = ATSimulator(at_lattice, callback, disable_emittance) + atsim = ATSimulator( + at_lattice, + linopt_function, + disable_emittance, + disable_chromaticity, + disable_radiation, + callback, + ) # Set the simulator data source on the Pytac lattice. pytac_lattice.set_data_source(ATLatticeDataSource(atsim), pytac.SIM) # Load the sim onto each element. diff --git a/src/atip/simulator.py b/src/atip/simulator.py index da82665..1b292b3 100644 --- a/src/atip/simulator.py +++ b/src/atip/simulator.py @@ -24,7 +24,10 @@ class LatticeData: def calculate_optics( at_lattice: at.lattice_object.Lattice, refpts: ArrayLike, + linopt_function: str = "linopt6", disable_emittance: bool = False, + disable_chromaticity: bool = False, + disable_radiation: bool = False, ) -> LatticeData: """Perform the physics calculations on the lattice. @@ -43,13 +46,47 @@ def calculate_optics( LatticeData: The calculated lattice data. """ logging.debug("Starting physics calculations.") - - orbit0, _ = at_lattice.find_orbit6() - logging.debug("Completed orbit calculation.") - - _, beamdata, twiss = at_lattice.linopt6( - refpts=refpts, get_chrom=True, orbit=orbit0, keep_lattice=True + logging.debug( + f"Using simulation params: {linopt_function}, disable_emittance=" + f"{disable_emittance}, disable_chromaticity={disable_chromaticity}, " + f"disable_radiation={disable_radiation}" ) + if linopt_function == "linopt6": + orbit0, _ = at_lattice.find_orbit6() + logging.debug("Completed orbit calculation.") + + _, beamdata, twiss = at_lattice.linopt6( + refpts=refpts, + get_chrom=not disable_chromaticity, + orbit=orbit0, + keep_lattice=True, + ) + elif linopt_function == "linopt4": + orbit0, _ = at_lattice.find_orbit4() + logging.debug("Completed orbit calculation.") + + _, beamdata, twiss = at_lattice.linopt6( + refpts=refpts, + get_chrom=not disable_chromaticity, + orbit=orbit0, + keep_lattice=True, + ) + elif linopt_function == "linopt2": + orbit0, _ = at_lattice.find_orbit() + logging.debug("Completed orbit calculation.") + + _, beamdata, twiss = at_lattice.linopt2( + refpts=refpts, + get_chrom=not disable_chromaticity, + orbit=orbit0, + keep_lattice=True, + ) + else: + raise ValueError( + f"Error. Invalid linopt function selected: {linopt_function}. Simulation " + "data not calculated." + ) + logging.debug("Completed linear optics calculation.") if not disable_emittance: @@ -57,7 +94,13 @@ def calculate_optics( logging.debug("Completed emittance calculation") else: emitdata = () - radint = at_lattice.get_radiation_integrals(twiss=twiss) + + if not disable_radiation: + radint = at_lattice.get_radiation_integrals(twiss=twiss) + logging.debug("Completed radiation calculation") + else: + radint = () + logging.debug("All calculation complete.") return LatticeData(twiss, beamdata.tune, beamdata.chromaticity, emitdata, radint) @@ -98,7 +141,15 @@ class ATSimulator: physics data upon a change. """ - def __init__(self, at_lattice, callback=None, disable_emittance=False): + def __init__( + self, + at_lattice, + linopt_function="linopt6", + disable_emittance=False, + disable_chromaticity=False, + disable_radiation=False, + callback=None, + ): """ .. Note:: To avoid errors, the physics data must be initially calculated here, during creation, otherwise it could be accidentally @@ -107,12 +158,16 @@ def __init__(self, at_lattice, callback=None, disable_emittance=False): the thread. Args: - at_lattice (at.lattice_object.Lattice): An instance of an AT - lattice object. - callback (typing.Callable): Optional, if passed it is called on completion - of each round of physics calculations. - disable_emittance (bool): Whether or not to perform the beam - envelope based emittance calculations. + at_lattice (at.lattice_object.Lattice): An instance of an AT lattice object. + linopt_function (str): Which pyAT linear optics function to use: linopt2, + linopt4, linopt6. + disable_emittance (bool): Whether the emittance calculations should be + disabled. + disable_chromaticity (bool): Whether the chromaticity calculations should be + disabled. + disable_radiation (bool): Whether radiation calculations should be disabled. + callback (typing.Callable): To be called after completion of each round of + physics calculations. **Methods:** """ @@ -122,12 +177,22 @@ def __init__(self, at_lattice, callback=None, disable_emittance=False): ) self._at_lat = at_lattice self._rp = numpy.ones(len(at_lattice) + 1, dtype=bool) + self._linopt_function = linopt_function self._disable_emittance = disable_emittance - self._at_lat.radiation_on() + self._disable_chromaticity = disable_chromaticity + self._disable_radiation = disable_radiation + + if not self._disable_radiation: + self._at_lat.radiation_on() # Initial phys data calculation. self._lattice_data = calculate_optics( - self._at_lat, self._rp, self._disable_emittance + self._at_lat, + self._rp, + self._linopt_function, + self._disable_emittance, + self._disable_chromaticity, + self._disable_radiation, ) # Threading stuff initialisation. @@ -196,7 +261,12 @@ def _recalculate_phys_data(self, callback): if bool(self._paused) is False: try: self._lattice_data = calculate_optics( - self._at_lat, self._rp, self._disable_emittance + self._at_lat, + self._rp, + self._linopt_function, + self._disable_emittance, + self._disable_chromaticity, + self._disable_radiation, ) except Exception as e: warn(at.AtWarning(e), stacklevel=1) diff --git a/src/atip/utils.py b/src/atip/utils.py index 5712480..564d696 100644 --- a/src/atip/utils.py +++ b/src/atip/utils.py @@ -29,7 +29,14 @@ def load_at_lattice(mode="I04", **kwargs): return at_lattice -def loader(mode="I04", callback=None, disable_emittance=False): +def loader( + mode="I04", + linopt_function="linopt6", + disable_emittance=False, + disable_chromaticity=False, + disable_radiation=False, + callback=None, +): """Load a unified lattice of the specifed mode. .. Note:: A unified lattice is a Pytac lattice where the corresponding AT @@ -52,7 +59,15 @@ def loader(mode="I04", callback=None, disable_emittance=False): periodicity=1, energy=pytac_lattice.get_value("energy", units=pytac.PHYS), ) - lattice = atip.load_sim.load(pytac_lattice, at_lattice, callback, disable_emittance) + lattice = atip.load_sim.load( + pytac_lattice, + at_lattice, + linopt_function, + disable_emittance, + disable_chromaticity, + disable_radiation, + callback, + ) return lattice From 0a2748340a8eca08f1905afa8358e18c87e24060 Mon Sep 17 00:00:00 2001 From: Phil Smith Date: Mon, 13 Oct 2025 10:43:30 +0000 Subject: [PATCH 2/9] Add tests for new sim params Also fix some tests which were broken by reordering of function arguments --- tests/conftest.py | 2 +- tests/test_at_simulator_object.py | 4 +- tests/test_load.py | 64 ++++++++++++++++++++++++++++++- 3 files changed, 65 insertions(+), 5 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index 4add592..a5f2ffd 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -80,7 +80,7 @@ def at_lattice(request): return atip.utils.load_at_lattice(request.param) -@pytest.fixture(scope="function", params=["DIAD"]) +@pytest.fixture(scope="function", params=["I04"]) def lattice_filepath(request): here = os.path.dirname(__file__) filepath = os.path.realpath( diff --git a/tests/test_at_simulator_object.py b/tests/test_at_simulator_object.py index 6f8dbe6..afe3a8f 100644 --- a/tests/test_at_simulator_object.py +++ b/tests/test_at_simulator_object.py @@ -214,9 +214,9 @@ def test_recalculate_phys_data_callback(at_lattice): atip.simulator.ATSimulator(at_lattice) # Check non-callable callback argument raises TypeError. with pytest.raises(TypeError): - atip.simulator.ATSimulator(at_lattice, "") + atip.simulator.ATSimulator(at_lattice, callback="") callback_func = mock.Mock() - atsim = atip.simulator.ATSimulator(at_lattice, callback_func) + atsim = atip.simulator.ATSimulator(at_lattice, callback=callback_func) atsim.queue_set(mock.Mock(), "f", 0) atsim.wait_for_calculations() callback_func.assert_called_once_with() diff --git a/tests/test_load.py b/tests/test_load.py index 3929bc1..6a11118 100644 --- a/tests/test_load.py +++ b/tests/test_load.py @@ -42,14 +42,16 @@ def test_load_atip_and_pytac_lattices(pytac_lattice, lattice_filepath): @pytest.mark.parametrize("at_and_pytac_lattices", RINGMODES_TO_TEST, indirect=True) def test_load_with_non_callable_callback_raises_TypeError(at_and_pytac_lattices): with pytest.raises(TypeError): - atip.load_sim.load(at_and_pytac_lattices[0], at_and_pytac_lattices[1], "") + atip.load_sim.load( + at_and_pytac_lattices[0], at_and_pytac_lattices[1], callback="" + ) @pytest.mark.parametrize("at_and_pytac_lattices", RINGMODES_TO_TEST, indirect=True) def test_load_with_callback(at_and_pytac_lattices): callback_func = mock.Mock() lat = atip.load_sim.load( - at_and_pytac_lattices[0], at_and_pytac_lattices[1], callback_func + at_and_pytac_lattices[0], at_and_pytac_lattices[1], callback=callback_func ) atsim = lat._data_source_manager._data_sources[pytac.SIM]._atsim atip.utils.trigger_calc(at_and_pytac_lattices[0]) @@ -60,3 +62,61 @@ def test_load_with_callback(at_and_pytac_lattices): def test_load_raises_ValueError_if_incompatible_lattices(): with pytest.raises(ValueError): atip.load_sim.load([1], [1, 2]) # length mismatch + + +@mock.patch("atip.simulator.calculate_optics") +def test_load_with_default_sim_params( + mocked_calc_optics, + pytac_lattice, + lattice_filepath, +): + pytac_lattice = atip.load_sim.load_from_filepath(pytac_lattice, lattice_filepath) + + mocked_calc_optics.assert_called_with( + mock.ANY, + mock.ANY, + "linopt6", + False, + False, + False, + ) + + +@pytest.mark.parametrize( + "linopt, disable_emittance, disable_chromaticity, disable_radiation", + [ + ("linopt6", False, True, False), + ("linopt6", True, True, False), + ("linopt4", True, False, True), + ("linopt4", True, True, True), + ("linopt2", True, False, True), + ("linopt2", True, True, True), + ], +) +@mock.patch("atip.simulator.calculate_optics") +def test_load_with_non_default_sim_params( + mocked_calc_optics, + pytac_lattice, + lattice_filepath, + linopt, + disable_emittance, + disable_chromaticity, + disable_radiation, +): + pytac_lattice = atip.load_sim.load_from_filepath( + pytac_lattice, + lattice_filepath, + linopt, + disable_emittance, + disable_chromaticity, + disable_radiation, + ) + + mocked_calc_optics.assert_called_with( + mock.ANY, + mock.ANY, + linopt, + disable_emittance, + disable_chromaticity, + disable_radiation, + ) From 64954ac1eafe6abe9682623b563746425ad7c352 Mon Sep 17 00:00:00 2001 From: Phil Smith Date: Mon, 13 Oct 2025 11:19:38 +0000 Subject: [PATCH 3/9] Remove support for python 3.12 and 3.13 due to cothread issues The CI is failing after python 3.11 due to a cothread seg fault. These version of cothread are known to be unreliable so, remove support for python 3.12 and 3.13 until after we have switched from cothread to asyncio --- .github/workflows/ci.yml | 10 ++++------ pyproject.toml | 2 -- 2 files changed, 4 insertions(+), 8 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 0195c6c..5cca8ae 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -5,7 +5,6 @@ on: pull_request: jobs: - lint: uses: ./.github/workflows/_tox.yml with: @@ -15,11 +14,11 @@ jobs: strategy: matrix: runs-on: ["ubuntu-latest"] # can add windows-latest, macos-latest - python-version: ["3.11", "3.12", "3.13"] + python-version: ["3.11"] include: # Include one that runs in the dev environment - runs-on: "ubuntu-latest" - python-version: "dev" + python-version: "3.11" fail-fast: false uses: ./.github/workflows/_test.yml with: @@ -41,17 +40,16 @@ jobs: docs: uses: ./.github/workflows/_docs.yml - dist: uses: ./.github/workflows/_dist.yml - + pypi: needs: [dist, test] if: github.ref_type == 'tag' uses: ./.github/workflows/_pypi.yml permissions: id-token: write - + release: needs: [dist, test, docs] if: github.ref_type == 'tag' diff --git a/pyproject.toml b/pyproject.toml index 8ddfdba..92c06fd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -8,8 +8,6 @@ classifiers = [ "Development Status :: 3 - Alpha", "License :: OSI Approved :: Apache Software License", "Programming Language :: Python :: 3.11", - "Programming Language :: Python :: 3.12", - "Programming Language :: Python :: 3.13", ] description = "Accelerator Toolbox Interface for Pytac" dependencies = [ From 078d3548f9719f59a0eb164178c5a0e18716f03a Mon Sep 17 00:00:00 2001 From: Phil Smith Date: Tue, 14 Oct 2025 11:50:15 +0000 Subject: [PATCH 4/9] Refactor sim params into a dataclass There will be a corresponding update to Virtac --- src/atip/load_sim.py | 38 ++----- src/atip/simulator.py | 160 +++++++++++++++--------------- src/atip/utils.py | 13 +-- tests/test_at_simulator_object.py | 14 +-- tests/test_load.py | 50 ++++------ 5 files changed, 118 insertions(+), 157 deletions(-) diff --git a/src/atip/load_sim.py b/src/atip/load_sim.py index 2c00ab2..4759e7e 100644 --- a/src/atip/load_sim.py +++ b/src/atip/load_sim.py @@ -15,10 +15,7 @@ def load_from_filepath( pytac_lattice, at_lattice_filepath, - linopt_function="linopt6", - disable_emittance=False, - disable_chromaticity=False, - disable_radiation=False, + sim_params=None, callback=None, ): """Load simulator data sources onto the lattice and its elements. @@ -27,13 +24,8 @@ def load_from_filepath( pytac_lattice (pytac.lattice.Lattice): An instance of a Pytac lattice. at_lattice_filepath (str): The path to a .mat file from which the Accelerator Toolbox lattice can be loaded. - linopt_function (str): Which pyAT linear optics function to use: linopt2, - linopt4, linopt6. - disable_emittance (bool): Whether the emittance calculations should be - disabled. - disable_chromaticity (bool): Whether the chromaticity calculations should be - disabled. - disable_radiation (bool): Whether radiation calculations should be disabled. + sim_params (SimParams | None): An optional dataclass containing the pyAT + simulation parameters to use. callback (typing.Callable): To be called after completion of each round of physics calculations. @@ -49,10 +41,7 @@ def load_from_filepath( return load( pytac_lattice, at_lattice, - linopt_function, - disable_emittance, - disable_chromaticity, - disable_radiation, + sim_params, callback, ) @@ -60,10 +49,7 @@ def load_from_filepath( def load( pytac_lattice, at_lattice, - linopt_function="linopt6", - disable_emittance=False, - disable_chromaticity=False, - disable_radiation=False, + sim_params=None, callback=None, ): """Load simulator data sources onto the lattice and its elements. @@ -71,13 +57,8 @@ def load( Args: pytac_lattice (pytac.lattice.Lattice): An instance of a Pytac lattice. at_lattice (at.lattice_object.Lattice): An instance of an AT lattice object. - linopt_function (str): Which pyAT linear optics function to use: linopt2, - linopt4, linopt6. - disable_emittance (bool): Whether the emittance calculations should be - disabled. - disable_chromaticity (bool): Whether the chromaticity calculations should be - disabled. - disable_radiation (bool): Whether radiation calculations should be disabled. + sim_params (SimParams | None): An optional dataclass containing the pyAT + simulation parameters to use. callback (typing.Callable): To be called after completion of each round of physics calculations. @@ -93,10 +74,7 @@ def load( # Initialise an instance of the ATSimulator Object. atsim = ATSimulator( at_lattice, - linopt_function, - disable_emittance, - disable_chromaticity, - disable_radiation, + sim_params, callback, ) # Set the simulator data source on the Pytac lattice. diff --git a/src/atip/simulator.py b/src/atip/simulator.py index 1b292b3..4e17fab 100644 --- a/src/atip/simulator.py +++ b/src/atip/simulator.py @@ -2,6 +2,7 @@ import logging from dataclasses import dataclass +from enum import StrEnum, auto from warnings import warn import at @@ -12,6 +13,38 @@ from scipy.constants import speed_of_light +class LinoptType(StrEnum): + LINOPT2 = auto() + LINOPT4 = auto() + LINOPT6 = auto() + + +@dataclass +class SimParams: + linopt: LinoptType = LinoptType.LINOPT6 + emittance: bool = True + chromaticity: bool = True + radiation: bool = True + + def __post_init__(self): + """Check that we have a valid combination of simulation parameters.""" + if self.radiation: + if self.linopt == LinoptType.LINOPT2 or self.linopt == LinoptType.LINOPT4: + raise ValueError( + f"You must disable radiation to use linopt function: {self.linopt}", + ) + else: + if self.linopt == LinoptType.LINOPT6: + raise ValueError( + f"You cannot use linopt function: {self.linopt} with radiation " + f"disabled.", + ) + elif self.emittance: + raise ValueError( + "You cannot calculate emittance with radiation disabled", + ) + + @dataclass class LatticeData: twiss: ArrayLike @@ -22,12 +55,7 @@ class LatticeData: def calculate_optics( - at_lattice: at.lattice_object.Lattice, - refpts: ArrayLike, - linopt_function: str = "linopt6", - disable_emittance: bool = False, - disable_chromaticity: bool = False, - disable_radiation: bool = False, + at_lattice: at.lattice_object.Lattice, refpts: ArrayLike, sp: SimParams ) -> LatticeData: """Perform the physics calculations on the lattice. @@ -39,63 +67,55 @@ def calculate_optics( Args: at_lattice (at.lattice_object.Lattice): AT lattice definition. refpts (numpy.typing.NDArray): A boolean array specifying the points at which - to calculate physics data. - disable_emittance (bool): whether to calculate emittance. + to calculate physics data. + sp (SimParams): An optional dataclass containing the pyAT simulation + parameters to use. Returns: LatticeData: The calculated lattice data. """ logging.debug("Starting physics calculations.") logging.debug( - f"Using simulation params: {linopt_function}, disable_emittance=" - f"{disable_emittance}, disable_chromaticity={disable_chromaticity}, " - f"disable_radiation={disable_radiation}" + f"Using simulation params: {sp.linopt}, emittance={sp.emittance}, chromaticity=" + f"{sp.chromaticity}, radiation={sp.radiation}" ) - if linopt_function == "linopt6": - orbit0, _ = at_lattice.find_orbit6() - logging.debug("Completed orbit calculation.") - - _, beamdata, twiss = at_lattice.linopt6( - refpts=refpts, - get_chrom=not disable_chromaticity, - orbit=orbit0, - keep_lattice=True, - ) - elif linopt_function == "linopt4": - orbit0, _ = at_lattice.find_orbit4() - logging.debug("Completed orbit calculation.") - - _, beamdata, twiss = at_lattice.linopt6( - refpts=refpts, - get_chrom=not disable_chromaticity, - orbit=orbit0, - keep_lattice=True, - ) - elif linopt_function == "linopt2": - orbit0, _ = at_lattice.find_orbit() - logging.debug("Completed orbit calculation.") - - _, beamdata, twiss = at_lattice.linopt2( - refpts=refpts, - get_chrom=not disable_chromaticity, - orbit=orbit0, - keep_lattice=True, - ) - else: - raise ValueError( - f"Error. Invalid linopt function selected: {linopt_function}. Simulation " - "data not calculated." - ) + match sp.linopt: + case LinoptType.LINOPT2: + orbit_func = at_lattice.find_orbit + linopt_func = at_lattice.linopt2 + case LinoptType.LINOPT4: + orbit_func = at_lattice.find_orbit4 + linopt_func = at_lattice.linopt4 + case LinoptType.LINOPT6: + orbit_func = at_lattice.find_orbit6 + linopt_func = at_lattice.linopt6 + case _: + raise ValueError( + f"Error. Invalid linopt function selected: {sp.linopt}. Simulation " + "data not calculated." + ) + + # Perform pyAT orbit calculation + orbit0, _ = orbit_func() + logging.debug("Completed orbit calculation.") + + # Perform pyAT linear optics calculation + _, beamdata, twiss = linopt_func( + refpts=refpts, + get_chrom=sp.chromaticity, + orbit=orbit0, + keep_lattice=True, + ) logging.debug("Completed linear optics calculation.") - if not disable_emittance: + if sp.emittance: emitdata = at_lattice.ohmi_envelope(orbit=orbit0, keep_lattice=True) logging.debug("Completed emittance calculation") else: emitdata = () - if not disable_radiation: + if sp.radiation: radint = at_lattice.get_radiation_integrals(twiss=twiss) logging.debug("Completed radiation calculation") else: @@ -126,8 +146,6 @@ class ATSimulator: physics data is calculated. _rp (numpy.typing.NDArray): A boolean array to be used as refpts for the physics calculations. - _disable_emittance (bool): Whether or not to perform the beam - envelope based emittance calculations. _lattice_data (LatticeData): calculated physics data function linopt (see at.lattice.linear.py). _queue (cothread.EventQueue): A queue of changes to be applied to @@ -144,10 +162,7 @@ class ATSimulator: def __init__( self, at_lattice, - linopt_function="linopt6", - disable_emittance=False, - disable_chromaticity=False, - disable_radiation=False, + sim_params=None, callback=None, ): """ @@ -159,13 +174,8 @@ def __init__( Args: at_lattice (at.lattice_object.Lattice): An instance of an AT lattice object. - linopt_function (str): Which pyAT linear optics function to use: linopt2, - linopt4, linopt6. - disable_emittance (bool): Whether the emittance calculations should be - disabled. - disable_chromaticity (bool): Whether the chromaticity calculations should be - disabled. - disable_radiation (bool): Whether radiation calculations should be disabled. + sim_params (SimParams | None): An optional dataclass containing the pyAT + simulation parameters to use. callback (typing.Callable): To be called after completion of each round of physics calculations. @@ -177,23 +187,16 @@ def __init__( ) self._at_lat = at_lattice self._rp = numpy.ones(len(at_lattice) + 1, dtype=bool) - self._linopt_function = linopt_function - self._disable_emittance = disable_emittance - self._disable_chromaticity = disable_chromaticity - self._disable_radiation = disable_radiation - if not self._disable_radiation: + if sim_params is None: + sim_params = SimParams() + self._sim_params = sim_params + + if self._sim_params.radiation: self._at_lat.radiation_on() # Initial phys data calculation. - self._lattice_data = calculate_optics( - self._at_lat, - self._rp, - self._linopt_function, - self._disable_emittance, - self._disable_chromaticity, - self._disable_radiation, - ) + self._lattice_data = calculate_optics(self._at_lat, self._rp, self._sim_params) # Threading stuff initialisation. self._queue = cothread.EventQueue() @@ -261,12 +264,7 @@ def _recalculate_phys_data(self, callback): if bool(self._paused) is False: try: self._lattice_data = calculate_optics( - self._at_lat, - self._rp, - self._linopt_function, - self._disable_emittance, - self._disable_chromaticity, - self._disable_radiation, + self._at_lat, self._rp, self._sim_params ) except Exception as e: warn(at.AtWarning(e), stacklevel=1) @@ -563,7 +561,7 @@ def get_emittance(self, field=None): Raises: pytac.FieldException: if the specified field is not valid for emittance. """ - if not self._disable_emittance: + if self._sim_params.emittance: if field is None: return self._lattice_data.emittance[0]["emitXY"] elif field == "x": diff --git a/src/atip/utils.py b/src/atip/utils.py index 564d696..c28771a 100644 --- a/src/atip/utils.py +++ b/src/atip/utils.py @@ -31,10 +31,7 @@ def load_at_lattice(mode="I04", **kwargs): def loader( mode="I04", - linopt_function="linopt6", - disable_emittance=False, - disable_chromaticity=False, - disable_radiation=False, + sim_params=None, callback=None, ): """Load a unified lattice of the specifed mode. @@ -45,9 +42,10 @@ def loader( Args: mode (str): The lattice operation mode. + sim_params (SimParams | None): An optional dataclass containing the pyAT + simulation parameters to use. callback (typing.Callable): Callable to be called after completion of each round of physics calculations in ATSimulator. - disable_emittance (bool): Whether the emittance should be calculated. Returns: pytac.lattice.Lattice: A Pytac lattice object with the simulator data @@ -62,10 +60,7 @@ def loader( lattice = atip.load_sim.load( pytac_lattice, at_lattice, - linopt_function, - disable_emittance, - disable_chromaticity, - disable_radiation, + sim_params, callback, ) return lattice diff --git a/tests/test_at_simulator_object.py b/tests/test_at_simulator_object.py index afe3a8f..3310b0b 100644 --- a/tests/test_at_simulator_object.py +++ b/tests/test_at_simulator_object.py @@ -169,18 +169,18 @@ def test_recalculate_phys_data(atsim, initial_phys_data): numpy.testing.assert_almost_equal(emit, [1.34308653e-10, 3.74339964e-13], decimal=3) -def test_disable_emittance_flag(atsim, initial_phys_data): +def test_emittance_flag(atsim, initial_phys_data): # Check emittance data is intially there - assert not atsim._disable_emittance + assert atsim._sim_params.emittance assert len(atsim._lattice_data.emittance) == 3 - # Check that ohmi_envelope is called when disable_emittance is False + # Check that ohmi_envelope is called when emittance is True atsim._at_lat.ohmi_envelope = mock.Mock() atsim.trigger_calculation() cothread.Sleep(0.1) atsim._at_lat.ohmi_envelope.assert_called_once() - # Check that ohmi_envelope isn't called when disable_emittance is True and that + # Check that ohmi_envelope isn't called when emittance is False and that # there isn't any emittance data - atsim._disable_emittance = True + atsim._sim_params.emittance = False atsim._at_lat.ohmi_envelope.reset_mock() atsim.trigger_calculation() cothread.Sleep(0.1) @@ -317,13 +317,13 @@ def test_get_m66(mocked_atsim, at_lattice): def test_get_emittance(mocked_atsim): - assert not mocked_atsim._disable_emittance + assert mocked_atsim._sim_params.emittance numpy.testing.assert_equal(mocked_atsim.get_emittance(), [1.4, 0.45]) assert mocked_atsim.get_emittance("x") == 1.4 assert mocked_atsim.get_emittance("y") == 0.45 with pytest.raises(FieldException): mocked_atsim.get_emittance("not_a_field") - mocked_atsim._disable_emittance = True + mocked_atsim._sim_params.emittance = False with pytest.raises(DataSourceException): mocked_atsim.get_emittance() diff --git a/tests/test_load.py b/tests/test_load.py index 6a11118..4cad808 100644 --- a/tests/test_load.py +++ b/tests/test_load.py @@ -65,7 +65,7 @@ def test_load_raises_ValueError_if_incompatible_lattices(): @mock.patch("atip.simulator.calculate_optics") -def test_load_with_default_sim_params( +def test_load_from_filepath_with_default_sim_params( mocked_calc_optics, pytac_lattice, lattice_filepath, @@ -73,24 +73,19 @@ def test_load_with_default_sim_params( pytac_lattice = atip.load_sim.load_from_filepath(pytac_lattice, lattice_filepath) mocked_calc_optics.assert_called_with( - mock.ANY, - mock.ANY, - "linopt6", - False, - False, - False, + mock.ANY, mock.ANY, atip.simulator.SimParams() ) @pytest.mark.parametrize( - "linopt, disable_emittance, disable_chromaticity, disable_radiation", + "linopt, emittance, chromaticity, radiation", [ - ("linopt6", False, True, False), - ("linopt6", True, True, False), - ("linopt4", True, False, True), - ("linopt4", True, True, True), - ("linopt2", True, False, True), - ("linopt2", True, True, True), + ("linopt6", True, False, True), + ("linopt6", False, False, True), + ("linopt4", False, True, False), + ("linopt4", False, False, False), + ("linopt2", False, True, False), + ("linopt2", False, False, False), ], ) @mock.patch("atip.simulator.calculate_optics") @@ -99,24 +94,19 @@ def test_load_with_non_default_sim_params( pytac_lattice, lattice_filepath, linopt, - disable_emittance, - disable_chromaticity, - disable_radiation, + emittance, + chromaticity, + radiation, ): - pytac_lattice = atip.load_sim.load_from_filepath( - pytac_lattice, - lattice_filepath, + sim_params = atip.simulator.SimParams( linopt, - disable_emittance, - disable_chromaticity, - disable_radiation, + emittance, + chromaticity, + radiation, ) - mocked_calc_optics.assert_called_with( - mock.ANY, - mock.ANY, - linopt, - disable_emittance, - disable_chromaticity, - disable_radiation, + pytac_lattice = atip.load_sim.load_from_filepath( + pytac_lattice, lattice_filepath, sim_params ) + + mocked_calc_optics.assert_called_with(mock.ANY, mock.ANY, sim_params) From 1f4e53c69f97b056c97bd62fdffa34ed5d08309f Mon Sep 17 00:00:00 2001 From: Phil Smith Date: Tue, 14 Oct 2025 11:52:39 +0000 Subject: [PATCH 5/9] Dont error on sphinx nitpicky warning --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 92c06fd..95f71db 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -101,7 +101,7 @@ commands = pre-commit: pre-commit run --all-files --show-diff-on-failure {posargs} type-checking: mypy src tests {posargs} tests: pytest --cov=atip --cov-report term --cov-report xml:cov.xml {posargs} - docs: sphinx-{posargs:build -EW --keep-going} -T docs build/html + docs: sphinx-{posargs:build -E --keep-going} -T docs build/html """ # Add -W flag to sphinx-build if you want to fail on warnings From f5c21d820975c1a9c4efcfb01d0a0eb3f321d961 Mon Sep 17 00:00:00 2001 From: Phil Smith Date: Tue, 14 Oct 2025 14:15:13 +0000 Subject: [PATCH 6/9] Adjust pyAT sim param setting We now use enable_6d instead of the misleading radiation_on, enable_6d is only called for linopt6. Also modify the validation exceptions to be more accurate --- src/atip/simulator.py | 20 ++++++++------------ 1 file changed, 8 insertions(+), 12 deletions(-) diff --git a/src/atip/simulator.py b/src/atip/simulator.py index 4e17fab..a374157 100644 --- a/src/atip/simulator.py +++ b/src/atip/simulator.py @@ -28,18 +28,14 @@ class SimParams: def __post_init__(self): """Check that we have a valid combination of simulation parameters.""" - if self.radiation: - if self.linopt == LinoptType.LINOPT2 or self.linopt == LinoptType.LINOPT4: + if self.linopt == LinoptType.LINOPT2 or self.linopt == LinoptType.LINOPT4: + if self.emittance or self.radiation: raise ValueError( - f"You must disable radiation to use linopt function: {self.linopt}", + "Emittance and radiation calculations must be disabled when using " + f"{self.linopt}", ) - else: - if self.linopt == LinoptType.LINOPT6: - raise ValueError( - f"You cannot use linopt function: {self.linopt} with radiation " - f"disabled.", - ) - elif self.emittance: + if self.linopt == LinoptType.LINOPT6: + if not self.radiation and self.emittance: raise ValueError( "You cannot calculate emittance with radiation disabled", ) @@ -192,8 +188,8 @@ def __init__( sim_params = SimParams() self._sim_params = sim_params - if self._sim_params.radiation: - self._at_lat.radiation_on() + if self._sim_params.linopt == LinoptType.LINOPT6: + self._at_lat.enable_6d() # Initial phys data calculation. self._lattice_data = calculate_optics(self._at_lat, self._rp, self._sim_params) From 63855d59790aebe6baa9aa877d205e396a67bf75 Mon Sep 17 00:00:00 2001 From: Phil Smith Date: Tue, 14 Oct 2025 15:16:22 +0000 Subject: [PATCH 7/9] Make dataclass frozen This will mean that post_init is called if anyone later modifies the attributes of the dataclass. This ensure validation always --- src/atip/simulator.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/atip/simulator.py b/src/atip/simulator.py index a374157..f6d810b 100644 --- a/src/atip/simulator.py +++ b/src/atip/simulator.py @@ -19,7 +19,7 @@ class LinoptType(StrEnum): LINOPT6 = auto() -@dataclass +@dataclass(frozen=True) class SimParams: linopt: LinoptType = LinoptType.LINOPT6 emittance: bool = True @@ -34,7 +34,7 @@ def __post_init__(self): "Emittance and radiation calculations must be disabled when using " f"{self.linopt}", ) - if self.linopt == LinoptType.LINOPT6: + elif self.linopt == LinoptType.LINOPT6: if not self.radiation and self.emittance: raise ValueError( "You cannot calculate emittance with radiation disabled", From 48f0e7df2e62665251ff62a801e35af37c8c2a13 Mon Sep 17 00:00:00 2001 From: Phil Smith Date: Tue, 14 Oct 2025 15:26:29 +0000 Subject: [PATCH 8/9] Move linopt validation check to dataclass --- src/atip/simulator.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/src/atip/simulator.py b/src/atip/simulator.py index f6d810b..ca8b61a 100644 --- a/src/atip/simulator.py +++ b/src/atip/simulator.py @@ -28,6 +28,14 @@ class SimParams: def __post_init__(self): """Check that we have a valid combination of simulation parameters.""" + try: + LinoptType(self.linopt) + except ValueError as e: + raise ValueError( + f"{self.linopt} is not a valid linopt function. Choose from: " + f"{[sp.value for sp in LinoptType]}" + ) from e + if self.linopt == LinoptType.LINOPT2 or self.linopt == LinoptType.LINOPT4: if self.emittance or self.radiation: raise ValueError( @@ -86,11 +94,6 @@ def calculate_optics( case LinoptType.LINOPT6: orbit_func = at_lattice.find_orbit6 linopt_func = at_lattice.linopt6 - case _: - raise ValueError( - f"Error. Invalid linopt function selected: {sp.linopt}. Simulation " - "data not calculated." - ) # Perform pyAT orbit calculation orbit0, _ = orbit_func() From eec4a53a4d5a5a16a87856d86160349ffab1bea1 Mon Sep 17 00:00:00 2001 From: Phil Smith Date: Tue, 14 Oct 2025 16:14:55 +0000 Subject: [PATCH 9/9] Split up disable_emittance flag tests Now that disabling emittance on an existing ATSimulator isnt supported, it makes more sense to have two seperate tests one which tests a simulator with and one without emittance disabled. Rather than testing both at once --- tests/test_at_simulator_object.py | 23 ++++++++++++++++------- 1 file changed, 16 insertions(+), 7 deletions(-) diff --git a/tests/test_at_simulator_object.py b/tests/test_at_simulator_object.py index 3310b0b..cf540a2 100644 --- a/tests/test_at_simulator_object.py +++ b/tests/test_at_simulator_object.py @@ -169,21 +169,25 @@ def test_recalculate_phys_data(atsim, initial_phys_data): numpy.testing.assert_almost_equal(emit, [1.34308653e-10, 3.74339964e-13], decimal=3) -def test_emittance_flag(atsim, initial_phys_data): +def test_ohmi_envelope_with_emittance_enabled(atsim, initial_phys_data): # Check emittance data is intially there - assert atsim._sim_params.emittance assert len(atsim._lattice_data.emittance) == 3 + assert atsim._sim_params.emittance # Check that ohmi_envelope is called when emittance is True atsim._at_lat.ohmi_envelope = mock.Mock() atsim.trigger_calculation() cothread.Sleep(0.1) atsim._at_lat.ohmi_envelope.assert_called_once() - # Check that ohmi_envelope isn't called when emittance is False and that - # there isn't any emittance data - atsim._sim_params.emittance = False - atsim._at_lat.ohmi_envelope.reset_mock() + + +def test_ohmi_envelope_with_emittance_disabled(atsim, initial_phys_data): + atsim._sim_params = atip.simulator.SimParams(emittance=False) + assert not atsim._sim_params.emittance + atsim._at_lat.ohmi_envelope = mock.Mock() atsim.trigger_calculation() cothread.Sleep(0.1) + # Check that ohmi_envelope isn't called when emittance is False and that + # there isn't any emittance data atsim._at_lat.ohmi_envelope.assert_not_called() assert len(atsim._lattice_data.emittance) == 0 @@ -323,7 +327,12 @@ def test_get_emittance(mocked_atsim): assert mocked_atsim.get_emittance("y") == 0.45 with pytest.raises(FieldException): mocked_atsim.get_emittance("not_a_field") - mocked_atsim._sim_params.emittance = False + + +def test_get_emittance_with_emittance_disabled(mocked_atsim): + assert mocked_atsim._sim_params.emittance + mocked_atsim._sim_params = atip.simulator.SimParams(emittance=False) + assert not mocked_atsim._sim_params.emittance with pytest.raises(DataSourceException): mocked_atsim.get_emittance()