diff --git a/.github/workflows/ruff_formatting.yml b/.github/workflows/ruff_formatting.yml new file mode 100644 index 00000000..bbce2ccb --- /dev/null +++ b/.github/workflows/ruff_formatting.yml @@ -0,0 +1,24 @@ +name: Ruff Format Check + +on: + push: + branches: [ main ] + pull_request: + branches: [ main ] + +jobs: + ruff-check: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v4 + with: + python-version: "3.11" # or your preferred version + + - name: Install Ruff + run: pip install ruff + + - name: Run Ruff Check + run: ruff check --diff . diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 00000000..f3a3141a --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,21 @@ +repos: + - repo: https://github.com/pre-commit/pre-commit-hooks + rev: v6.0.0 + hooks: + - id: trailing-whitespace + - id: end-of-file-fixer + + - repo: https://github.com/astral-sh/ruff-pre-commit + rev: v0.14.5 + hooks: + - id: ruff + args: [--fix] + - id: ruff-format + + # - repo: https://github.com/pre-commit/mirrors-mypy + # rev: v1.18.2 + # hooks: + # - id: mypy + # additional_dependencies: [ + # "pydantic>=2.0", + # ] diff --git a/docs/conf.py b/docs/conf.py index 94e77854..5272bf42 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -1,6 +1,6 @@ # -*- coding: utf-8 -*- # -# pyAML documentation build configuration file, +# pyAML documentation build configuration file, # # This file is execfile()d with the current directory set to its # containing dir. @@ -17,6 +17,7 @@ # import pathlib import sys + # ignore numpy warnings, see: # https://stackoverflow.com/questions/40845304/runtimewarning-numpy-dtype-size-changed-may-indicate-binary-incompatibility import warnings @@ -44,18 +45,19 @@ # Add any Sphinx extension module names here, as strings. They can be # extensions coming with Sphinx (named 'sphinx.ext.*') or your custom # ones. -extensions = ['sphinx.ext.autodoc', - 'sphinx.ext.doctest', - 'sphinx.ext.todo', - 'sphinx.ext.coverage', - 'sphinx.ext.mathjax', - 'sphinx.ext.viewcode', - 'sphinx.ext.githubpages', - 'sphinx.ext.napoleon', - 'sphinx.ext.autosectionlabel', - 'sphinx.ext.autosummary', - 'myst_nb' - ] +extensions = [ + "sphinx.ext.autodoc", + "sphinx.ext.doctest", + "sphinx.ext.todo", + "sphinx.ext.coverage", + "sphinx.ext.mathjax", + "sphinx.ext.viewcode", + "sphinx.ext.githubpages", + "sphinx.ext.napoleon", + "sphinx.ext.autosectionlabel", + "sphinx.ext.autosummary", + "myst_nb", +] autosectionlabel_prefix_document = True autosectionlabel_maxdepth = 2 @@ -66,18 +68,18 @@ # You can specify multiple suffix as a list of string: # # source_suffix = ['.rst', '.md'] -source_suffix = '.rst' +source_suffix = ".rst" # The master toctree document. -master_doc = 'index' +master_doc = "index" # General information about the project. project = ABOUT_PYAML["__title__"] -copyright_ = '2024, pyAML collaboration' +copyright_ = "2024, pyAML collaboration" author = ABOUT_PYAML["__author__"] rst_prolog = f""" -:github_url: {ABOUT_PYAML['__url__']} +:github_url: {ABOUT_PYAML["__url__"]} """ # The version info for the project you're documenting, acts as replacement for @@ -94,7 +96,7 @@ # # This is also used if you do content translation via gettext catalogs. # Usually you set "language" from the command line for these cases. -language = 'en' +language = "en" # List of patterns, relative to source directory, that match files and # directories to ignore when looking for source files. @@ -102,7 +104,7 @@ exclude_patterns = [] # The name of the Pygments (syntax highlighting) style to use. -pygments_style = 'sphinx' +pygments_style = "sphinx" # If true, `todo` and `todoList` produce output, else they produce nothing. todo_include_todos = True @@ -120,12 +122,12 @@ # further. For a list of options available for each theme, see the # documentation. html_theme_options = { - 'collapse_navigation': False, - 'display_version': True, - 'logo_only': True, - 'navigation_depth': 6, + "collapse_navigation": False, + "display_version": True, + "logo_only": True, + "navigation_depth": 6, "rightsidebar": True, - "relbarbgcolor": "black" + "relbarbgcolor": "black", } @@ -133,31 +135,31 @@ # that is the logo of the docs, or URL that points an image file for the logo. # It is placed at the top of the sidebar; # its width should therefore not exceed 200 pixels. -html_logo = '_static/img/logo.png' +html_logo = "_static/img/logo.png" html_copy_source = False html_theme_options = { "github_url": "https://github.com/atcollab/at", "logo": { - "image_light": '_static/img/logo.png', - "image_dark": '_static/img/dark.png', - } + "image_light": "_static/img/logo.png", + "image_dark": "_static/img/dark.png", + }, } # Add any paths that contain custom static files (such as style sheets) here, # relative to this directory. They are copied after the builtin static files, # so a file named "default.css" will overwrite the builtin "default.css". # -html_static_path = ['_static'] +html_static_path = ["_static"] # A dictionary of values to pass into the template engine’s context for all # pages. Single values can also be put in this dictionary using the # -A command-line option of sphinx-build. html_context = { - 'display_github': True, + "display_github": True, # the following are only needed if :github_url: is not set - 'github_user': author, - 'github_repo': project, - 'github_version': 'main/docs/', + "github_user": author, + "github_repo": project, + "github_version": "main/docs/", } # A list of CSS files. The entry must be a filename string or a tuple @@ -169,7 +171,9 @@ # html_css_files = ["css/custom.css"] -smartquotes_action = "qe" # renders only quotes and ellipses (...) but not dashes (option: D) +smartquotes_action = ( + "qe" # renders only quotes and ellipses (...) but not dashes (option: D) +) # Custom sidebar templates, must be a dictionary that maps document names # to template names. @@ -177,9 +181,9 @@ # This is required for the alabaster theme # refs: http://alabaster.readthedocs.io/en/latest/installation.html#sidebars html_sidebars = { - '**': [ - 'relations.html', # needs 'show_related': True theme option to display - 'searchbox.html', + "**": [ + "relations.html", # needs 'show_related': True theme option to display + "searchbox.html", ], "index": [], "common/about": [], @@ -189,7 +193,7 @@ # -- Options for HTMLHelp output ------------------------------------------ # Output file base name for HTML help builder. -htmlhelp_basename = 'pyamldoc' +htmlhelp_basename = "pyamldoc" # -- Options for LaTeX output --------------------------------------------- @@ -197,15 +201,12 @@ # The paper size ('letterpaper' or 'a4paper'). # # 'papersize': 'letterpaper', - # The font size ('10pt', '11pt' or '12pt'). # # 'pointsize': '10pt', - # Additional stuff for the LaTeX preamble. # # 'preamble': '', - # Latex figure (float) alignment # # 'figure_align': 'htbp', @@ -215,18 +216,14 @@ # (source start file, target name, title, # author, documentclass [howto, manual, or own class]). latex_documents = [ - (master_doc, 'pyaml.tex', u'pyAML Documentation', - u'pyAML collaboration', 'manual'), + (master_doc, "pyaml.tex", "pyAML Documentation", "pyAML collaboration", "manual"), ] # -- Options for manual page output --------------------------------------- # One entry per manual page. List of tuples # (source start file, name, description, authors, manual section). -man_pages = [ - (master_doc, 'pyaml', u'pyAML Documentation', - [author], 1) -] +man_pages = [(master_doc, "pyaml", "pyAML Documentation", [author], 1)] # -- Options for Texinfo output ------------------------------------------- @@ -234,18 +231,28 @@ # (source start file, target name, title, author, # dir menu entry, description, category) texinfo_documents = [ - (master_doc, 'pyaml', u'pyAML Documentation', - author, 'pyaml', 'One line description of project.', - 'Miscellaneous'), + ( + master_doc, + "pyaml", + "pyAML Documentation", + author, + "pyaml", + "One line description of project.", + "Miscellaneous", + ), ] # -- Autodoc Configuration --------------------------------------------------- # Add here all modules to be mocked up. When the dependencies are not met # at building time. Here used to have PyQT mocked. -autodoc_mock_imports = ['PyQt5', 'PyQt5.QtGui', 'PyQt5.QtCore', 'PyQt5.QtWidgets', - "matplotlib.backends.backend_qt5agg", - ] +autodoc_mock_imports = [ + "PyQt5", + "PyQt5.QtGui", + "PyQt5.QtCore", + "PyQt5.QtWidgets", + "matplotlib.backends.backend_qt5agg", +] # -- Options for the myst markdown parser ------------------------------------ @@ -257,5 +264,5 @@ "deflist", ] myst_heading_anchors = 3 -nb_execution_mode = "off" #"auto" +nb_execution_mode = "off" # "auto" nb_execution_allow_errors = True diff --git a/docs/notebooks/control_system.ipynb b/docs/notebooks/control_system.ipynb index a0bcc52e..dd9f04e2 100644 --- a/docs/notebooks/control_system.ipynb +++ b/docs/notebooks/control_system.ipynb @@ -15,7 +15,7 @@ "metadata": {}, "outputs": [], "source": [ - "from pyaml.pyaml import pyaml,PyAML" + "from pyaml.pyaml import PyAML, pyaml" ] } ], diff --git a/docs/notebooks/live_design.ipynb b/docs/notebooks/live_design.ipynb index 4185bc23..f25c3f93 100644 --- a/docs/notebooks/live_design.ipynb +++ b/docs/notebooks/live_design.ipynb @@ -33,10 +33,11 @@ "metadata": {}, "outputs": [], "source": [ - "from pyaml.pyaml import pyaml,PyAML\n", + "import numpy as np\n", "from pyaml.instrument import Instrument\n", - "from pyaml.configuration.factory import Factory\n", - "import numpy as np" + "from pyaml.pyaml import PyAML, pyaml\n", + "\n", + "from pyaml.configuration.factory import Factory" ] }, { @@ -54,8 +55,8 @@ "metadata": {}, "outputs": [], "source": [ - "ml:PyAML = pyaml('../../tests/config/EBSTune.yaml')\n", - "SR:Instrument = ml.get('sr')" + "ml: PyAML = pyaml(\"../../tests/config/EBSTune.yaml\")\n", + "SR: Instrument = ml.get(\"sr\")" ] }, { @@ -103,7 +104,7 @@ } ], "source": [ - "sr.design.get_lattice() #Direct access to underlying pyAT lattice" + "sr.design.get_lattice() # Direct access to underlying pyAT lattice" ] }, { @@ -115,7 +116,7 @@ "source": [ "quadForTune = sr.get_magnets(\"QForTune\")\n", "\n", - "tune_device = sr.get_betatron_tune_monitor('BETATRON_TUNE')" + "tune_device = sr.get_betatron_tune_monitor(\"BETATRON_TUNE\")" ] }, { @@ -140,14 +141,14 @@ "print(f\"Tune directly from the lattice {tune_design}\")\n", "initial_tune = tune_device.tune.get()\n", "print(f\"Tune via pyAML interface for design mode {initial_tune}\")\n", - "tunemat = np.zeros((len(quadForTune),2))\n", + "tunemat = np.zeros((len(quadForTune), 2))\n", "\n", - "for idx,m in enumerate(quadForTune):\n", - " str = m.strength.get()\n", - " m.strength.set(str+1e-4)\n", - " dq = tune_device.tune.get() - initial_tune\n", - " tunemat[idx] = dq*1e4\n", - " m.strength.set(str)" + "for idx, m in enumerate(quadForTune):\n", + " str = m.strength.get()\n", + " m.strength.set(str + 1e-4)\n", + " dq = tune_device.tune.get() - initial_tune\n", + " tunemat[idx] = dq * 1e4\n", + " m.strength.set(str)" ] }, { @@ -165,10 +166,12 @@ "metadata": {}, "outputs": [], "source": [ - "print(tune_device.tune.get()) # what ever the globally set mode whas\n", + "print(tune_device.tune.get()) # what ever the globally set mode has\n", "\n", - "SR.design.get_betatron_tune_monitor('BETATRON_TUNE').tune.get() # tune from design simulations \n", - "# SR.live.get_betatron_tune_monitor('BETATRON_TUNE').tune.get() # tune from real accelerator" + "# tune from design simulations\n", + "SR.design.get_betatron_tune_monitor(\"BETATRON_TUNE\").tune.get()\n", + "# tune from real accelerator\n", + "# SR.live.get_betatron_tune_monitor('BETATRON_TUNE').tune.get()" ] }, { diff --git a/docs/notebooks/load_configuration.ipynb b/docs/notebooks/load_configuration.ipynb index 0832a439..4dba3bd3 100644 --- a/docs/notebooks/load_configuration.ipynb +++ b/docs/notebooks/load_configuration.ipynb @@ -15,9 +15,9 @@ "metadata": {}, "outputs": [], "source": [ - "from pyaml.pyaml import pyaml,PyAML\n", + "from pyaml.pyaml import PyAML, pyaml\n", "\n", - "ml:PyAML = pyaml('../../tests/config/EBSTune.yaml')" + "ml: PyAML = pyaml(\"../../tests/config/EBSTune.yaml\")" ] }, { diff --git a/examples/ESRF_tune_example/esrf_tune_example.py b/examples/ESRF_tune_example/esrf_tune_example.py index 9d1e9c15..77400362 100644 --- a/examples/ESRF_tune_example/esrf_tune_example.py +++ b/examples/ESRF_tune_example/esrf_tune_example.py @@ -1,18 +1,19 @@ -from pyaml.accelerator import Accelerator -import numpy as np import os +import numpy as np + +from pyaml.accelerator import Accelerator # Get the directory of the current script script_dir = os.path.dirname(__file__) # Go up one level and then into 'data' -relative_path = os.path.join(script_dir, '..', '..', 'tests', 'config','EBSTune.yaml') +relative_path = os.path.join(script_dir, "..", "..", "tests", "config", "EBSTune.yaml") # Normalize the path (resolves '..') absolute_path = os.path.abspath(relative_path) -sr:Accelerator = Accelerator.load(absolute_path) +sr: Accelerator = Accelerator.load(absolute_path) sr.design.get_lattice().disable_6d() quadForTuneDesign = sr.design.get_magnets("QForTune") @@ -21,31 +22,31 @@ # Build tune response matrix tune = sr.design.get_lattice().get_tune() print(tune) -tunemat = np.zeros((len(quadForTuneDesign),2)) +tunemat = np.zeros((len(quadForTuneDesign), 2)) -for idx,m in enumerate(quadForTuneDesign): - str = m.strength.get() - m.strength.set(str+1e-4) - dq = sr.design.get_lattice().get_tune() - tune - tunemat[idx] = dq*1e4 - m.strength.set(str) +for idx, m in enumerate(quadForTuneDesign): + str = m.strength.get() + m.strength.set(str + 1e-4) + dq = sr.design.get_lattice().get_tune() - tune + tunemat[idx] = dq * 1e4 + m.strength.set(str) # Compute correction matrix correctionmat = np.linalg.pinv(tunemat.T) # Correct tune strs = quadForTuneDesign.strengths.get() -strs += np.matmul(correctionmat,[0.1,0.05]) # Ask for correction [dqx,dqy] +strs += np.matmul(correctionmat, [0.1, 0.05]) # Ask for correction [dqx,dqy] quadForTuneDesign.strengths.set(strs) newTune = sr.design.get_lattice().get_tune() -diffTune = newTune-tune +diffTune = newTune - tune print(diffTune) -assert( np.abs(diffTune[0]-0.1) < 1e-3 ) -assert( np.abs(diffTune[1]-0.05) < 1e-3 ) +assert np.abs(diffTune[0] - 0.1) < 1e-3 +assert np.abs(diffTune[1] - 0.05) < 1e-3 if False: # Correct the tune on live (need a Virutal Accelerator) quadForTuneLive = sr.live.get_magnets("QForTune") strs = quadForTuneLive.strengths.get() - strs += np.matmul(correctionmat,[0.1,0.05]) # Ask for correction [dqx,dqy] + strs += np.matmul(correctionmat, [0.1, 0.05]) # Ask for correction [dqx,dqy] quadForTuneLive.strengths.set(strs) diff --git a/examples/ESRF_tune_example/esrf_tune_example_no_yaml.py b/examples/ESRF_tune_example/esrf_tune_example_no_yaml.py index abbca888..f80aeb66 100644 --- a/examples/ESRF_tune_example/esrf_tune_example_no_yaml.py +++ b/examples/ESRF_tune_example/esrf_tune_example_no_yaml.py @@ -1,24 +1,45 @@ -from pyaml.instrument import Instrument,ConfigModel as InstrumentConfigModel -from pyaml.magnet.quadrupole import Quadrupole,ConfigModel as QuadrupoleConfig -from pyaml.configuration.csvcurve import CSVCurve,ConfigModel as CSVCureveConfig -from pyaml.magnet.linear_model import LinearMagnetModel,ConfigModel as LinearMagnetModelConfig -from pyaml.lattice.simulator import Simulator,ConfigModel as SimulatorConfigModel -from pyaml.arrays.magnet import Magnet,ConfigModel as MagnetArrayConfigModel -from tango.pyaml.controlsystem import TangoControlSystem,ConfigModel as ControlSystemConfig -from tango.pyaml.attribute import Attribute,ConfigModel as AttributeConfig -from tango.pyaml.attribute_read_only import AttributeReadOnly,ConfigModel as AttributeReadOnlyConfig -from pyaml.configuration import set_root_folder import os - +import time import numpy as np -import time +from pyaml.instrument import ConfigModel as InstrumentConfigModel +from pyaml.instrument import Instrument +from tango.pyaml.attribute import Attribute +from tango.pyaml.attribute import ConfigModel as AttributeConfig +from tango.pyaml.attribute_read_only import ( + AttributeReadOnly, +) +from tango.pyaml.attribute_read_only import ( + ConfigModel as AttributeReadOnlyConfig, +) +from tango.pyaml.controlsystem import ( + ConfigModel as ControlSystemConfig, +) +from tango.pyaml.controlsystem import ( + TangoControlSystem, +) + +from pyaml.arrays.magnet import ConfigModel as MagnetArrayConfigModel +from pyaml.arrays.magnet import Magnet +from pyaml.configuration import set_root_folder +from pyaml.configuration.csvcurve import ConfigModel as CSVCureveConfig +from pyaml.configuration.csvcurve import CSVCurve +from pyaml.lattice.simulator import ConfigModel as SimulatorConfigModel +from pyaml.lattice.simulator import Simulator +from pyaml.magnet.linear_model import ( + ConfigModel as LinearMagnetModelConfig, +) +from pyaml.magnet.linear_model import ( + LinearMagnetModel, +) +from pyaml.magnet.quadrupole import ConfigModel as QuadrupoleConfig +from pyaml.magnet.quadrupole import Quadrupole # Get the directory of the current script script_dir = os.path.dirname(__file__) # Go up one level and then into 'data' -relative_path = os.path.join(script_dir, '..', '..', 'tests') +relative_path = os.path.join(script_dir, "..", "..", "tests") # Normalize the path (resolves '..') absolute_path = os.path.abspath(relative_path) @@ -27,33 +48,73 @@ # Configuration -tangocs = ControlSystemConfig(name="live",tango_host="ebs-simu-3:10000") +tangocs = ControlSystemConfig(name="live", tango_host="ebs-simu-3:10000") control = TangoControlSystem(tangocs) control.init_cs() qfCurve = CSVCurve(CSVCureveConfig(file="config/sr/magnet_models/QF1_strength.csv")) qdCurve = CSVCurve(CSVCureveConfig(file="config/sr/magnet_models/QD2_strength.csv")) -elemConfig = [ {"name":"QF1A-C01", "attname":"srmag/vps-qf1/c01-a/current", "calibration_factor":1.00504, "curve":qfCurve}, - {"name":"QF1E-C01", "attname":"srmag/vps-qf1/c01-e/current", "calibration_factor":0.998212, "curve":qfCurve}, - {"name":"QD2A-C01", "attname":"srmag/vps-qd2/c01-a/current", "calibration_factor":1.00504, "curve":qdCurve}, - {"name":"QD2E-C01", "attname":"srmag/vps-qd2/c01-e/current", "calibration_factor":1.003485189, "curve":qdCurve} ] - -devices=[] -names=[] +elemConfig = [ + { + "name": "QF1A-C01", + "attname": "srmag/vps-qf1/c01-a/current", + "calibration_factor": 1.00504, + "curve": qfCurve, + }, + { + "name": "QF1E-C01", + "attname": "srmag/vps-qf1/c01-e/current", + "calibration_factor": 0.998212, + "curve": qfCurve, + }, + { + "name": "QD2A-C01", + "attname": "srmag/vps-qd2/c01-a/current", + "calibration_factor": 1.00504, + "curve": qdCurve, + }, + { + "name": "QD2E-C01", + "attname": "srmag/vps-qd2/c01-e/current", + "calibration_factor": 1.003485189, + "curve": qdCurve, + }, +] + +devices = [] +names = [] for cfg in elemConfig: - qAtt = Attribute(AttributeConfig(attribute=cfg["attname"],unit="A")) - qModel = LinearMagnetModel(LinearMagnetModelConfig(curve=cfg["curve"],calibration_factor=cfg["calibration_factor"],powerconverter=qAtt,unit="1/m")) - devices.append( Quadrupole(QuadrupoleConfig(name=cfg["name"],model=qModel)) ) + qAtt = Attribute(AttributeConfig(attribute=cfg["attname"], unit="A")) + qModel = LinearMagnetModel( + LinearMagnetModelConfig( + curve=cfg["curve"], + calibration_factor=cfg["calibration_factor"], + powerconverter=qAtt, + unit="1/m", + ) + ) + devices.append(Quadrupole(QuadrupoleConfig(name=cfg["name"], model=qModel))) names.append(cfg["name"]) -simulator = Simulator(SimulatorConfigModel(name="design",lattice="config/sr/lattices/ebs.mat")) - -quads = Magnet(MagnetArrayConfigModel(name="quadsForTune",elements=names)) +simulator = Simulator( + SimulatorConfigModel(name="design", lattice="config/sr/lattices/ebs.mat") +) -sr = Instrument(InstrumentConfigModel(name="sr",energy=6e9,controls=[control],simulators=[simulator],devices=devices,arrays=[quads],data_folder="/tmp")) +quads = Magnet(MagnetArrayConfigModel(name="quadsForTune", elements=names)) +sr = Instrument( + InstrumentConfigModel( + name="sr", + energy=6e9, + controls=[control], + simulators=[simulator], + devices=devices, + arrays=[quads], + data_folder="/tmp", + ) +) # Usage exmaple @@ -64,27 +125,29 @@ # Compute tune response matrix for the 4 quads from simulator sr.design.get_lattice().disable_6d() tune = sr.design.get_lattice().get_tune() -tunemat = np.zeros((len(quadForTuneDesign),2)) -for idx,m in enumerate(quadForTuneDesign): +tunemat = np.zeros((len(quadForTuneDesign), 2)) +for idx, m in enumerate(quadForTuneDesign): str = m.strength.get() - m.strength.set(str+1e-4) + m.strength.set(str + 1e-4) dq = sr.design.get_lattice().get_tune() - tune - tunemat[idx] = dq*1e4 + tunemat[idx] = dq * 1e4 m.strength.set(str) # Compute correction matrix correctionmat = np.linalg.pinv(tunemat.T) # Correct tune on live -qxAtt = AttributeReadOnly(AttributeReadOnlyConfig(attribute="sys/ringsimulator/ebs/Tune_h",unit="")) -qyAtt = AttributeReadOnly(AttributeReadOnlyConfig(attribute="sys/ringsimulator/ebs/Tune_v",unit="")) +qxAtt = AttributeReadOnly( + AttributeReadOnlyConfig(attribute="sys/ringsimulator/ebs/Tune_h", unit="") +) +qyAtt = AttributeReadOnly( + AttributeReadOnlyConfig(attribute="sys/ringsimulator/ebs/Tune_v", unit="") +) print(f"Tune={qxAtt.readback()}, {qyAtt.readback()}") strs = quadForTuneLive.strengths.get() -strs += np.matmul(correctionmat,[0.1,0.05]) # Ask for correction [dqx,dqy] +strs += np.matmul(correctionmat, [0.1, 0.05]) # Ask for correction [dqx,dqy] quadForTuneLive.strengths.set(strs) time.sleep(3) print(f"Tune={qxAtt.readback()}, {qyAtt.readback()}") - - diff --git a/pyaml/__init__.py b/pyaml/__init__.py index d9538786..bf958687 100644 --- a/pyaml/__init__.py +++ b/pyaml/__init__.py @@ -14,8 +14,8 @@ import logging.config import os -from pyaml.common.exception import PyAMLException -from pyaml.common.exception import PyAMLConfigException + +from pyaml.common.exception import PyAMLConfigException, PyAMLException __all__ = [__version__, PyAMLException, PyAMLConfigException] @@ -27,5 +27,5 @@ logger = logging.getLogger("pyaml") level = os.getenv("PYAML_LOG_LEVEL", "").upper() -if len(level)>0: +if len(level) > 0: logger.setLevel(getattr(logging, level, logging.WARNING)) diff --git a/pyaml/accelerator.py b/pyaml/accelerator.py index fd625f4a..63a70b5a 100644 --- a/pyaml/accelerator.py +++ b/pyaml/accelerator.py @@ -2,26 +2,29 @@ Instrument class """ -from .control.controlsystem import ControlSystem -from .common.element import Element -from .lattice.simulator import Simulator +from pydantic import BaseModel, ConfigDict + from .arrays.array import ArrayConfig -from pydantic import BaseModel,ConfigDict +from .common.element import Element from .configuration import load_accelerator +from .control.controlsystem import ControlSystem +from .lattice.simulator import Simulator # Define the main class name for this module PYAMLCLASS = "Accelerator" -class ConfigModel(BaseModel): - model_config = ConfigDict(arbitrary_types_allowed=True,extra="forbid") +class ConfigModel(BaseModel): + model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") name: str """Instrument name""" energy: float - """Instrument nominal energy, for ramped machine, this value can be dynamically set""" + """Instrument nominal energy, for ramped machine, + this value can be dynamically set""" controls: list[ControlSystem] = None - """List of control system used, an instrument can access several control systems""" + """List of control system used, an instrument + can access several control systems""" simulators: list[Simulator] = None """Simulator list""" data_folder: str @@ -34,7 +37,7 @@ class ConfigModel(BaseModel): class Accelerator(object): """PyAML top level class""" - + def __init__(self, cfg: ConfigModel): self._cfg = cfg __design = None @@ -43,20 +46,20 @@ def __init__(self, cfg: ConfigModel): if cfg.controls is not None: for c in cfg.controls: if c.name() == "live": - self.__live = c + self.__live = c else: - # Add as dynacmic attribute - setattr(self,c.name(),c) + # Add as dynacmic attribute + setattr(self, c.name(), c) c.init_cs() c.fill_device(cfg.devices) if cfg.simulators is not None: for s in cfg.simulators: if s.name() == "design": - self.__design = s + self.__design = s else: - # Add as dynacmic attribute - setattr(self,s.name(),s) + # Add as dynacmic attribute + setattr(self, s.name(), s) s.fill_device(cfg.devices) if cfg.arrays is not None: @@ -71,7 +74,7 @@ def __init__(self, cfg: ConfigModel): if cfg.energy is not None: self.set_energy(cfg.energy) - def set_energy(self,E:float): + def set_energy(self, E: float): if self._cfg.simulators is not None: for s in self._cfg.simulators: s.set_energy(E) @@ -88,7 +91,7 @@ def design(self) -> Simulator: return self.__design @staticmethod - def load(filename:str, use_fast_loader:bool = False) -> "Accelerator": + def load(filename: str, use_fast_loader: bool = False) -> "Accelerator": """ Load an accelerator from a config file. @@ -97,7 +100,9 @@ def load(filename:str, use_fast_loader:bool = False) -> "Accelerator": filename : str Configuration file name, yaml or json. use_fast_loader : bool - Use fast yaml loader. When specified, no line number are reported in case of error, - only the element name that triggered the error will be reported in the exception) + Use fast yaml loader. When specified, + no line number are reported in case of error, + only the element name that triggered the error + will be reported in the exception) """ - return load_accelerator(filename,use_fast_loader) + return load_accelerator(filename, use_fast_loader) diff --git a/pyaml/arrays/array.py b/pyaml/arrays/array.py index 711ca970..f5a4db3b 100644 --- a/pyaml/arrays/array.py +++ b/pyaml/arrays/array.py @@ -1,25 +1,30 @@ """ Magnet array configuration """ + +from pydantic import BaseModel, ConfigDict + +from pyaml.common.exception import PyAMLException + from ..common.element_holder import ElementHolder -from pydantic import BaseModel,ConfigDict class ArrayConfigModel(BaseModel): - - model_config = ConfigDict(arbitrary_types_allowed=True,extra="forbid") + model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") name: str """Family name""" elements: list[str] """List of pyaml element names""" + class ArrayConfig(object): """ Class that implements configuration for access to arrays (families) """ + def __init__(self, cfg: ArrayConfigModel): self._cfg = cfg - def fill_array(self,holder:ElementHolder): - raise "Array.fill_array() is not subclassed" + def fill_array(self, holder: ElementHolder): + raise PyAMLException("Array.fill_array() is not subclassed") diff --git a/pyaml/arrays/bpm.py b/pyaml/arrays/bpm.py index 2355cd7f..f7fd7e16 100644 --- a/pyaml/arrays/bpm.py +++ b/pyaml/arrays/bpm.py @@ -1,15 +1,16 @@ -from .array import ArrayConfigModel,ArrayConfig from ..common.element_holder import ElementHolder +from .array import ArrayConfig, ArrayConfigModel # Define the main class name for this module PYAMLCLASS = "BPM" -class ConfigModel(ArrayConfigModel):... -class BPM(ArrayConfig): +class ConfigModel(ArrayConfigModel): ... + - def __init__(self, cfg: ArrayConfigModel): +class BPM(ArrayConfig): + def __init__(self, cfg: ArrayConfigModel): super().__init__(cfg) - def fill_array(self,holder:ElementHolder): - holder.fill_bpm_array(self._cfg.name,self._cfg.elements) + def fill_array(self, holder: ElementHolder): + holder.fill_bpm_array(self._cfg.name, self._cfg.elements) diff --git a/pyaml/arrays/bpm_array.py b/pyaml/arrays/bpm_array.py index 70123f22..a268b2db 100644 --- a/pyaml/arrays/bpm_array.py +++ b/pyaml/arrays/bpm_array.py @@ -1,40 +1,39 @@ -from ..common.abstract import ReadFloatArray +import numpy as np + from ..bpm.bpm import BPM +from ..common.abstract import ReadFloatArray from ..control.deviceaccesslist import DeviceAccessList from .element_array import ElementArray -import numpy as np class RWBPMPosition(ReadFloatArray): - - def __init__(self, name:str, bpms:list[BPM]): + def __init__(self, name: str, bpms: list[BPM]): self.__bpms = bpms self.__name = name - self.__aggregator:DeviceAccessList = None + self.__aggregator: DeviceAccessList = None # Gets the values def get(self) -> np.array: if not self.__aggregator: return np.array([b.positions.get() for b in self.__bpms]) - else: - return self.__aggregator.get().reshape(len(self.__bpms),2) + else: + return self.__aggregator.get().reshape(len(self.__bpms), 2) # Gets the unit of the values def unit(self) -> list[str]: return [b.positions.unit() for b in self.__bpms] # Set the aggregator (Control system only) - def set_aggregator(self,agg:DeviceAccessList): + def set_aggregator(self, agg: DeviceAccessList): self.__aggregator = agg class RWBPMSinglePosition(ReadFloatArray): - - def __init__(self, name:str, bpms:list[BPM],idx: int): + def __init__(self, name: str, bpms: list[BPM], idx: int): self.__bpms = bpms self.__name = name self.__idx = idx - self.__aggregator:DeviceAccessList = None + self.__aggregator: DeviceAccessList = None # Gets the values def get(self) -> np.array: @@ -48,17 +47,16 @@ def unit(self) -> list[str]: return [b.positions.unit() for b in self.__bpms] # Set the aggregator (Control system only) - def set_aggregator(self,agg:DeviceAccessList): + def set_aggregator(self, agg: DeviceAccessList): self.__aggregator = agg - class BPMArray(ElementArray): """ Class that implements access to a BPM array """ - def __init__(self,arrayName:str,bpms:list[BPM],use_aggregator = True): + def __init__(self, arrayName: str, bpms: list[BPM], use_aggregator=True): """ Construct a BPM array @@ -67,18 +65,19 @@ def __init__(self,arrayName:str,bpms:list[BPM],use_aggregator = True): arrayName : str Array name bpms: list[BPM] - BPM list, all elements must be attached to the same instance of + BPM list, all elements must be attached to the same instance of either a Simulator or a ControlSystem. use_aggregator : bool - Use aggregator to increase performance by using paralell access to underlying devices. + Use aggregator to increase performance by using paralell + access to underlying devices. """ - super().__init__(arrayName,bpms,use_aggregator) + super().__init__(arrayName, bpms, use_aggregator) - self.__hvpos = RWBPMPosition(arrayName,bpms) - self.__hpos = RWBPMSinglePosition(arrayName,bpms,0) - self.__vpos = RWBPMSinglePosition(arrayName,bpms,1) + self.__hvpos = RWBPMPosition(arrayName, bpms) + self.__hpos = RWBPMSinglePosition(arrayName, bpms, 0) + self.__vpos = RWBPMSinglePosition(arrayName, bpms, 1) - if use_aggregator: + if use_aggregator: aggs = self.get_peer().create_bpm_aggregators(bpms) self.__hvpos.set_aggregator(aggs[0]) self.__hpos.set_aggregator(aggs[1]) @@ -104,8 +103,3 @@ def v(self) -> RWBPMSinglePosition: Give access to bpm V posttions of each bpm of this array """ return self.__vpos - - - - - \ No newline at end of file diff --git a/pyaml/arrays/cfm_magnet.py b/pyaml/arrays/cfm_magnet.py index e0aabab7..4ac75ef8 100644 --- a/pyaml/arrays/cfm_magnet.py +++ b/pyaml/arrays/cfm_magnet.py @@ -1,15 +1,16 @@ -from .array import ArrayConfigModel,ArrayConfig from ..common.element_holder import ElementHolder +from .array import ArrayConfig, ArrayConfigModel # Define the main class name for this module PYAMLCLASS = "CombinedFunctionMagnet" -class ConfigModel(ArrayConfigModel):... -class CombinedFunctionMagnet(ArrayConfig): +class ConfigModel(ArrayConfigModel): ... + - def __init__(self, cfg: ArrayConfigModel): +class CombinedFunctionMagnet(ArrayConfig): + def __init__(self, cfg: ArrayConfigModel): super().__init__(cfg) - def fill_array(self,holder:ElementHolder): - holder.fill_cfm_magnet_array(self._cfg.name,self._cfg.elements) + def fill_array(self, holder: ElementHolder): + holder.fill_cfm_magnet_array(self._cfg.name, self._cfg.elements) diff --git a/pyaml/arrays/cfm_magnet_array.py b/pyaml/arrays/cfm_magnet_array.py index 30b1cd93..8e362588 100644 --- a/pyaml/arrays/cfm_magnet_array.py +++ b/pyaml/arrays/cfm_magnet_array.py @@ -1,14 +1,15 @@ +import numpy as np + from ..common.abstract import ReadWriteFloatArray +from ..common.exception import PyAMLException from ..magnet.cfm_magnet import CombinedFunctionMagnet from .element_array import ElementArray -from ..common.exception import PyAMLException -import numpy as np -#TODO handle aggregator for CFM +# TODO handle aggregator for CFM -class RWMagnetStrengths(ReadWriteFloatArray): - def __init__(self, name:str, magnets:list[CombinedFunctionMagnet]): +class RWMagnetStrengths(ReadWriteFloatArray): + def __init__(self, name: str, magnets: list[CombinedFunctionMagnet]): self.__name = name self.__magnets = magnets self.__nb = sum(m.nb_multipole() for m in magnets) @@ -18,20 +19,20 @@ def get(self) -> np.array: r = np.zeros(self.__nb) idx = 0 for m in self.__magnets: - r[idx:idx+m.nb_multipole()] = m.strengths.get() - idx+=m.nb_multipole() + r[idx : idx + m.nb_multipole()] = m.strengths.get() + idx += m.nb_multipole() return r # Sets the values - def set(self, value:np.array): - nvalue = np.ones(self.__nb) * value if isinstance(value,float) else value + def set(self, value: np.array): + nvalue = np.ones(self.__nb) * value if isinstance(value, float) else value idx = 0 for m in self.__magnets: - m.strengths.set(nvalue[idx:idx+m.nb_multipole()]) - idx+=m.nb_multipole() + m.strengths.set(nvalue[idx : idx + m.nb_multipole()]) + idx += m.nb_multipole() # Sets the values and waits that the read values reach their setpoint - def set_and_wait(self, value:np.array): + def set_and_wait(self, value: np.array): raise NotImplementedError("Not implemented yet.") # Gets the unit of the values @@ -41,9 +42,9 @@ def unit(self) -> list[str]: r.extend(m.strengths.unit()) return r -class RWMagnetHardwares(ReadWriteFloatArray): - def __init__(self, name:str, magnets:list[CombinedFunctionMagnet]): +class RWMagnetHardwares(ReadWriteFloatArray): + def __init__(self, name: str, magnets: list[CombinedFunctionMagnet]): self.__name = name self.__magnets = magnets self.__nb = sum(m.nb_multipole() for m in magnets) @@ -53,20 +54,20 @@ def get(self) -> np.array: r = np.zeros(self.__nb) idx = 0 for m in self.__magnets: - r[idx:idx+m.nb_multipole()] = m.hardwares.get() - idx+=m.nb_multipole() + r[idx : idx + m.nb_multipole()] = m.hardwares.get() + idx += m.nb_multipole() return r # Sets the values - def set(self, value:np.array): - nvalue = np.ones(self.__nb) * value if isinstance(value,float) else value + def set(self, value: np.array): + nvalue = np.ones(self.__nb) * value if isinstance(value, float) else value idx = 0 for m in self.__magnets: - m.hardwares.set(nvalue[idx:idx+m.nb_multipole()]) - idx+=m.nb_multipole() - + m.hardwares.set(nvalue[idx : idx + m.nb_multipole()]) + idx += m.nb_multipole() + # Sets the values and waits that the read values reach their setpoint - def set_and_wait(self, value:np.array): + def set_and_wait(self, value: np.array): raise NotImplementedError("Not implemented yet.") # Gets the unit of the values @@ -82,7 +83,12 @@ class CombinedFunctionMagnetArray(ElementArray): Class that implements access to a magnet array """ - def __init__(self,arrayName:str,magnets:list[CombinedFunctionMagnet],use_aggregator = False): + def __init__( + self, + arrayName: str, + magnets: list[CombinedFunctionMagnet], + use_aggregator=False, + ): """ Construct a magnet array @@ -91,20 +97,25 @@ def __init__(self,arrayName:str,magnets:list[CombinedFunctionMagnet],use_aggrega arrayName : str Array name magnets: list[Magnet] - Magnet list, all elements must be attached to the same instance of + Magnet list, all elements must be attached to the same instance of either a Simulator or a ControlSystem. use_aggregator : bool - Use aggregator to increase performance by using paralell access to underlying devices. + Use aggregator to increase performance by using paralell + access to underlying devices. """ - super().__init__(arrayName,magnets,use_aggregator) - - self.__rwstrengths = RWMagnetStrengths(arrayName,magnets) - self.__rwhardwares = RWMagnetHardwares(arrayName,magnets) + super().__init__(arrayName, magnets, use_aggregator) + + self.__rwstrengths = RWMagnetStrengths(arrayName, magnets) + self.__rwhardwares = RWMagnetHardwares(arrayName, magnets) if use_aggregator: - raise(PyAMLException("Aggregator not implemented for CombinedFunctionMagnetArray")) + raise ( + PyAMLException( + "Aggregator not implemented for CombinedFunctionMagnetArray" + ) + ) - @property + @property def strengths(self) -> RWMagnetStrengths: """ Give access to strength of each magnet of this array @@ -117,9 +128,3 @@ def hardwares(self) -> RWMagnetHardwares: Give access to hardware value of each magnet of this array """ return self.__rwhardwares - - - - - - \ No newline at end of file diff --git a/pyaml/arrays/element.py b/pyaml/arrays/element.py index 0d4ced0f..efb919c3 100644 --- a/pyaml/arrays/element.py +++ b/pyaml/arrays/element.py @@ -1,15 +1,16 @@ -from .array import ArrayConfigModel,ArrayConfig from ..common.element_holder import ElementHolder +from .array import ArrayConfig, ArrayConfigModel # Define the main class name for this module PYAMLCLASS = "Element" -class ConfigModel(ArrayConfigModel):... -class Element(ArrayConfig): +class ConfigModel(ArrayConfigModel): ... + - def __init__(self, cfg: ArrayConfigModel): +class Element(ArrayConfig): + def __init__(self, cfg: ArrayConfigModel): super().__init__(cfg) - def fill_array(self,holder:ElementHolder): - holder.fill_element_array(self._cfg.name,self._cfg.elements) + def fill_array(self, holder: ElementHolder): + holder.fill_element_array(self._cfg.name, self._cfg.elements) diff --git a/pyaml/arrays/element_array.py b/pyaml/arrays/element_array.py index d05cf5be..2736843d 100644 --- a/pyaml/arrays/element_array.py +++ b/pyaml/arrays/element_array.py @@ -1,18 +1,19 @@ -from ..common.element import Element -from ..magnet.magnet import Magnet +import fnmatch +import importlib + from ..bpm.bpm import BPM -from ..magnet.cfm_magnet import CombinedFunctionMagnet +from ..common.element import Element from ..common.exception import PyAMLException +from ..magnet.cfm_magnet import CombinedFunctionMagnet +from ..magnet.magnet import Magnet -import importlib -import fnmatch class ElementArray(list[Element]): """ Class that implements access to a element array """ - def __init__(self,arrayName:str,elements:list[Element],use_aggregator = True): + def __init__(self, arrayName: str, elements: list[Element], use_aggregator=True): """ Construct an element array @@ -21,17 +22,22 @@ def __init__(self,arrayName:str,elements:list[Element],use_aggregator = True): arrayName : str Array name elements: list[Element] - Element list, all elements must be attached to the same instance of + Element list, all elements must be attached to the same instance of either a Simulator or a ControlSystem. use_aggregator : bool - Use aggregator to increase performance by using paralell access to underlying devices. + Use aggregator to increase performance by using paralell + access to underlying devices. """ super().__init__(i for i in elements) self.__name = arrayName - self.__peer = self[0]._peer if len(self)>0 else None + self.__peer = self[0]._peer if len(self) > 0 else None self.__use_aggretator = use_aggregator - if self.__peer is None or any([m._peer!=self.__peer for m in self]): - raise PyAMLException(f"{self.__class__.__name__} {self.get_name()}: All elements must be attached to the same instance of either a Simulator or a ControlSystem") + if self.__peer is None or any([m._peer != self.__peer for m in self]): + raise PyAMLException( + f"{self.__class__.__name__} {self.get_name()}: " + "All elements must be attached to the same instance " + "of either a Simulator or a ControlSystem" + ) def get_peer(self): """ @@ -44,54 +50,50 @@ def get_name(self) -> str: def names(self) -> list[str]: return [e.get_name() for e in self] - - def __create_array(self,arrName:str,eltType:type,elements:list): - if len(elements)==0: + def __create_array(self, arrName: str, eltType: type, elements: list): + if len(elements) == 0: return [] - if issubclass(eltType,Magnet): + if issubclass(eltType, Magnet): m = importlib.import_module("pyaml.arrays.magnet_array") - arrayClass = getattr(m, "MagnetArray", None) - return arrayClass("",elements,self.__use_aggretator) - elif issubclass(eltType,BPM): + arrayClass = getattr(m, "MagnetArray", None) + return arrayClass("", elements, self.__use_aggretator) + elif issubclass(eltType, BPM): m = importlib.import_module("pyaml.arrays.bpm_array") - arrayClass = getattr(m, "BPMArray", None) - return arrayClass("",elements,self.__use_aggretator) - elif issubclass(eltType,CombinedFunctionMagnet): + arrayClass = getattr(m, "BPMArray", None) + return arrayClass("", elements, self.__use_aggretator) + elif issubclass(eltType, CombinedFunctionMagnet): m = importlib.import_module("pyaml.arrays.cfm_magnet_array") - arrayClass = getattr(m, "CombinedFunctionMagnetArray", None) - return arrayClass("",elements,self.__use_aggretator) - elif issubclass(eltType,Element): - return ElementArray("",elements,self.__use_aggretator) + arrayClass = getattr(m, "CombinedFunctionMagnetArray", None) + return arrayClass("", elements, self.__use_aggretator) + elif issubclass(eltType, Element): + return ElementArray("", elements, self.__use_aggretator) else: raise PyAMLException(f"Unsupported sliced array for type {str(eltType)}") - - def __eval_field(self,attName:str,e:Element) -> str: + + def __eval_field(self, attName: str, e: Element) -> str: funcName = "get_" + attName - func = getattr(e,funcName, None) + func = getattr(e, funcName, None) return func() if func is not None else "" - def __getitem__(self,key): - - if isinstance(key,slice): - + def __getitem__(self, key): + if isinstance(key, slice): # Slicing eltType = None r = [] for i in range(*key.indices(len(self))): if eltType is None: eltType = type(self[i]) - elif not isinstance(self[i],eltType): - eltType = Element # Fall back to element + elif not isinstance(self[i], eltType): + eltType = Element # Fall back to element r.append(self[i]) - return self.__create_array("",eltType,r) - - elif isinstance(key,str): + return self.__create_array("", eltType, r) - fields = key.split(':') + elif isinstance(key, str): + fields = key.split(":") - if len(fields)<=1: + if len(fields) <= 1: # Selection by name eltType = None r = [] @@ -99,23 +101,23 @@ def __getitem__(self,key): if fnmatch.fnmatch(e.get_name(), key): if eltType is None: eltType = type(e) - elif not isinstance(e,eltType): - eltType = Element # Fall back to element + elif not isinstance(e, eltType): + eltType = Element # Fall back to element r.append(e) else: # Selection by fields eltType = None r = [] for e in self: - txt = self.__eval_field(fields[0],e) - if fnmatch.fnmatch(txt , fields[1]): + txt = self.__eval_field(fields[0], e) + if fnmatch.fnmatch(txt, fields[1]): if eltType is None: eltType = type(e) - elif not isinstance(e,eltType): - eltType = Element # Fall back to element + elif not isinstance(e, eltType): + eltType = Element # Fall back to element r.append(e) - return self.__create_array("",eltType,r) + return self.__create_array("", eltType, r) else: # Default to super selection diff --git a/pyaml/arrays/magnet.py b/pyaml/arrays/magnet.py index 0e7c26ec..067e2ba2 100644 --- a/pyaml/arrays/magnet.py +++ b/pyaml/arrays/magnet.py @@ -1,15 +1,16 @@ -from .array import ArrayConfigModel,ArrayConfig from ..common.element_holder import ElementHolder +from .array import ArrayConfig, ArrayConfigModel # Define the main class name for this module PYAMLCLASS = "Magnet" -class ConfigModel(ArrayConfigModel):... -class Magnet(ArrayConfig): +class ConfigModel(ArrayConfigModel): ... + - def __init__(self, cfg: ArrayConfigModel): +class Magnet(ArrayConfig): + def __init__(self, cfg: ArrayConfigModel): super().__init__(cfg) - def fill_array(self,holder:ElementHolder): - holder.fill_magnet_array(self._cfg.name,self._cfg.elements) + def fill_array(self, holder: ElementHolder): + holder.fill_magnet_array(self._cfg.name, self._cfg.elements) diff --git a/pyaml/arrays/magnet_array.py b/pyaml/arrays/magnet_array.py index a67519e5..b5c992f0 100644 --- a/pyaml/arrays/magnet_array.py +++ b/pyaml/arrays/magnet_array.py @@ -1,16 +1,17 @@ +import numpy as np + from ..common.abstract import ReadWriteFloatArray -from ..magnet.magnet import Magnet from ..common.abstract_aggregator import ScalarAggregator +from ..magnet.magnet import Magnet from .element_array import ElementArray -import numpy as np -class RWMagnetStrength(ReadWriteFloatArray): - def __init__(self, name:str, magnets:list[Magnet]): +class RWMagnetStrength(ReadWriteFloatArray): + def __init__(self, name: str, magnets: list[Magnet]): self.__name = name self.__magnets = magnets self.__nb = len(self.__magnets) - self.__aggregator:ScalarAggregator = None + self.__aggregator: ScalarAggregator = None # Gets the values def get(self) -> np.array: @@ -20,16 +21,16 @@ def get(self) -> np.array: return self.__aggregator.get() # Sets the values - def set(self, value:np.array): - nvalue = np.ones(self.__nb) * value if isinstance(value,float) else value + def set(self, value: np.array): + nvalue = np.ones(self.__nb) * value if isinstance(value, float) else value if not self.__aggregator: - for idx,m in enumerate(self.__magnets): + for idx, m in enumerate(self.__magnets): m.strength.set(nvalue[idx]) else: self.__aggregator.set(nvalue) - + # Sets the values and waits that the read values reach their setpoint - def set_and_wait(self, value:np.array): + def set_and_wait(self, value: np.array): raise NotImplementedError("Not implemented yet.") # Gets the unit of the values @@ -37,16 +38,16 @@ def unit(self) -> list[str]: return [m.strength.unit() for m in self.__magnets] # Set the aggregator (Control system only) - def set_aggregator(self,agg:ScalarAggregator): + def set_aggregator(self, agg: ScalarAggregator): self.__aggregator = agg -class RWMagnetHardware(ReadWriteFloatArray): - def __init__(self, name:str, magnets:list[Magnet]): +class RWMagnetHardware(ReadWriteFloatArray): + def __init__(self, name: str, magnets: list[Magnet]): self.__name = name self.__magnets = magnets self.__nb = len(self.__magnets) - self.__aggregator:ScalarAggregator = None + self.__aggregator: ScalarAggregator = None # Gets the values def get(self) -> np.array: @@ -56,16 +57,16 @@ def get(self) -> np.array: return self.__aggregator.get() # Sets the values - def set(self, value:np.array): - nvalue = np.ones(self.__nb) * value if isinstance(value,float) else value + def set(self, value: np.array): + nvalue = np.ones(self.__nb) * value if isinstance(value, float) else value if not self.__aggregator: - for idx,m in enumerate(self.__magnets): + for idx, m in enumerate(self.__magnets): m.hardware.set(value[idx]) else: self.__aggregator.set(value) - + # Sets the values and waits that the read values reach their setpoint - def set_and_wait(self, value:np.array): + def set_and_wait(self, value: np.array): raise NotImplementedError("Not implemented yet.") # Gets the unit of the values @@ -73,15 +74,16 @@ def unit(self) -> list[str]: return [m.hardware.unit() for m in self.__magnets] # Set the aggregator - def set_aggregator(self,agg:ScalarAggregator): + def set_aggregator(self, agg: ScalarAggregator): self.__aggregator = agg + class MagnetArray(ElementArray): """ Class that implements access to a magnet array """ - def __init__(self,arrayName:str,magnets:list[Magnet],use_aggregator = True): + def __init__(self, arrayName: str, magnets: list[Magnet], use_aggregator=True): """ Construct a magnet array @@ -90,15 +92,16 @@ def __init__(self,arrayName:str,magnets:list[Magnet],use_aggregator = True): arrayName : str Array name magnets: list[Magnet] - Magnet list, all elements must be attached to the same instance of + Magnet list, all elements must be attached to the same instance of either a Simulator or a ControlSystem. use_aggregator : bool - Use aggregator to increase performance by using paralell access to underlying devices. + Use aggregator to increase performance by using + paralell access to underlying devices. """ - super().__init__(arrayName,magnets,use_aggregator) - - self.__rwstrengths = RWMagnetStrength(arrayName,magnets) - self.__rwhardwares = RWMagnetHardware(arrayName,magnets) + super().__init__(arrayName, magnets, use_aggregator) + + self.__rwstrengths = RWMagnetStrength(arrayName, magnets) + self.__rwhardwares = RWMagnetHardware(arrayName, magnets) if use_aggregator: aggs = self.get_peer().create_magnet_strength_aggregator(magnets) @@ -106,7 +109,7 @@ def __init__(self,arrayName:str,magnets:list[Magnet],use_aggregator = True): self.__rwstrengths.set_aggregator(aggs) self.__rwhardwares.set_aggregator(aggh) - @property + @property def strengths(self) -> RWMagnetStrength: """ Give access to strength of each magnet of this array @@ -119,9 +122,3 @@ def hardwares(self) -> RWMagnetHardware: Give access to hardware value of each magnet of this array """ return self.__rwhardwares - - - - - - \ No newline at end of file diff --git a/pyaml/bpm/bpm.py b/pyaml/bpm/bpm.py index f6a8e3f3..1c68560d 100644 --- a/pyaml/bpm/bpm.py +++ b/pyaml/bpm/bpm.py @@ -1,7 +1,8 @@ -from ..common.element import Element, ElementConfigModel -from ..lattice.abstract_impl import RBpmArray, RWBpmOffsetArray, RWBpmTiltScalar from ..bpm.bpm_model import BPMModel +from ..common.element import Element, ElementConfigModel from ..common.exception import PyAMLException +from ..lattice.abstract_impl import RBpmArray, RWBpmOffsetArray, RWBpmTiltScalar + try: from typing import Self # Python 3.11+ except ImportError: @@ -9,11 +10,10 @@ PYAMLCLASS = "BPM" -class ConfigModel(ElementConfigModel): - - model: BPMModel | None = None - """Object in charge of BPM modeling""" +class ConfigModel(ElementConfigModel): + model: BPMModel | None = None + """Object in charge of BPM modeling""" class BPM(Element): @@ -34,7 +34,7 @@ def __init__(self, cfg: ConfigModel): model : BPMModel BPM model in charge of computing beam position """ - + super().__init__(cfg.name) self.__model = cfg.model if hasattr(cfg, "model") else None @@ -45,12 +45,12 @@ def __init__(self, cfg: ConfigModel): @property def model(self) -> BPMModel: - return self.__model + return self.__model @property def positions(self) -> RBpmArray: if self.__positions is None: - raise PyAMLException(f"{str(self)} has no attached positions") + raise PyAMLException(f"{str(self)} has no attached positions") return self.__positions @property @@ -65,8 +65,13 @@ def tilt(self) -> RWBpmTiltScalar: raise PyAMLException(f"{str(self)} has no attached tilt") return self.__tilt - def attach(self, peer, positions: RBpmArray , offset: RWBpmOffsetArray, - tilt: RWBpmTiltScalar) -> Self: + def attach( + self, + peer, + positions: RBpmArray, + offset: RWBpmOffsetArray, + tilt: RWBpmTiltScalar, + ) -> Self: # Attach positions, offset and tilt attributes and returns a new # reference obj = self.__class__(self._cfg) @@ -76,4 +81,3 @@ def attach(self, peer, positions: RBpmArray , offset: RWBpmOffsetArray, obj.__tilt = tilt obj._peer = peer return obj - \ No newline at end of file diff --git a/pyaml/bpm/bpm_model.py b/pyaml/bpm/bpm_model.py index a19f5205..6d4dfd76 100644 --- a/pyaml/bpm/bpm_model.py +++ b/pyaml/bpm/bpm_model.py @@ -1,13 +1,17 @@ from abc import ABCMeta, abstractmethod + import numpy as np from numpy.typing import NDArray + from ..control.deviceaccess import DeviceAccess + class BPMModel(metaclass=ABCMeta): """ Abstract class providing interface to accessing BPM positions, offsets, tilts. """ + @abstractmethod def read_position(self) -> NDArray[np.float64]: """ @@ -19,7 +23,7 @@ def read_position(self) -> NDArray[np.float64]: positions """ pass - + @abstractmethod def read_tilt(self) -> float: """ @@ -73,7 +77,7 @@ def set_offset(self, offset: NDArray[np.float64]): def get_pos_devices(self) -> list[DeviceAccess]: """ Get device handles used for position reading - + Returns ------- list[DeviceAccess] @@ -85,7 +89,7 @@ def get_pos_devices(self) -> list[DeviceAccess]: def get_tilt_device(self) -> DeviceAccess: """ Get device handle used for tilt access - + Returns ------- list[DeviceAccess] @@ -97,7 +101,7 @@ def get_tilt_device(self) -> DeviceAccess: def get_offset_devices(self) -> list[DeviceAccess]: """ Get device handles used for offset access - + Returns ------- list[DeviceAccess] diff --git a/pyaml/bpm/bpm_simple_model.py b/pyaml/bpm/bpm_simple_model.py index 5365526e..4f55076e 100644 --- a/pyaml/bpm/bpm_simple_model.py +++ b/pyaml/bpm/bpm_simple_model.py @@ -1,30 +1,34 @@ -from pyaml.bpm.bpm_model import BPMModel -from pydantic import BaseModel,ConfigDict import numpy as np -from ..control.deviceaccess import DeviceAccess +from numpy.typing import NDArray +from pydantic import BaseModel, ConfigDict + +from pyaml.bpm.bpm_model import BPMModel + from ..common.element import __pyaml_repr__ +from ..control.deviceaccess import DeviceAccess -from numpy.typing import NDArray # Define the main class name for this module PYAMLCLASS = "BPMSimpleModel" -class ConfigModel(BaseModel): - model_config = ConfigDict(arbitrary_types_allowed=True,extra="forbid") +class ConfigModel(BaseModel): + model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") x_pos: DeviceAccess """Horizontal position""" y_pos: DeviceAccess """Vertical position""" + class BPMSimpleModel(BPMModel): """ Concrete implementation of BPMModel that simulates a BPM with tilt and offset values. """ + def __init__(self, cfg: ConfigModel): - self._cfg = cfg - + self._cfg = cfg + self.__x_pos = cfg.x_pos self.__y_pos = cfg.y_pos @@ -38,7 +42,7 @@ def read_position(self) -> NDArray: positions """ return np.array([self.__x_pos.get(), self.__y_pos.get()]) - + def read_tilt(self) -> float: """ Simulate reading the tilt value from a BPM. @@ -48,7 +52,7 @@ def read_tilt(self) -> float: The tilt value of the BPM """ raise NotImplementedError("Tilt reading not implemented in this model.") - + def read_offset(self) -> NDArray: """ Simulate reading the offset values from a BPM. @@ -59,6 +63,7 @@ def read_offset(self) -> NDArray: offsets """ raise NotImplementedError("Offset reading not implemented in this model.") + def set_tilt(self, tilt: float): """ Simulate setting the tilt value of a BPM. @@ -71,7 +76,7 @@ def set_tilt(self, tilt: float): None """ raise NotImplementedError("Tilt setting not implemented in this model.") - + def set_offset(self, offset_values: np.ndarray): """ Simulate setting the offset values of a BPM @@ -82,22 +87,22 @@ def set_offset(self, offset_values: np.ndarray): offsets to set for the BPM """ raise NotImplementedError("Offset setting not implemented in this model.") - + def get_pos_devices(self) -> list[DeviceAccess]: """ Get device handles used for position reading - + Returns ------- list[DeviceAccess] Array of DeviceAcess """ - return [self.__x_pos,self.__y_pos] + return [self.__x_pos, self.__y_pos] def get_tilt_device(self) -> DeviceAccess: """ Get device handle used for tilt access - + Returns ------- list[DeviceAccess] @@ -108,7 +113,7 @@ def get_tilt_device(self) -> DeviceAccess: def get_offset_devices(self) -> list[DeviceAccess]: """ Get device handles used for offset access - + Returns ------- list[DeviceAccess] @@ -118,4 +123,3 @@ def get_offset_devices(self) -> list[DeviceAccess]: def __repr__(self): return __pyaml_repr__(self) - diff --git a/pyaml/bpm/bpm_tiltoffset_model.py b/pyaml/bpm/bpm_tiltoffset_model.py index 92bb2b0f..ab940832 100644 --- a/pyaml/bpm/bpm_tiltoffset_model.py +++ b/pyaml/bpm/bpm_tiltoffset_model.py @@ -1,17 +1,19 @@ +import numpy as np +from numpy.typing import NDArray +from pydantic import BaseModel, ConfigDict + from pyaml.bpm.bpm_model import BPMModel from pyaml.bpm.bpm_simple_model import BPMSimpleModel -from pydantic import BaseModel,ConfigDict -import numpy as np -from ..control.deviceaccess import DeviceAccess + from ..common.element import __pyaml_repr__ +from ..control.deviceaccess import DeviceAccess -from numpy.typing import NDArray # Define the main class name for this module PYAMLCLASS = "BPMTiltOffsetModel" -class ConfigModel(BaseModel): - model_config = ConfigDict(arbitrary_types_allowed=True,extra="forbid") +class ConfigModel(BaseModel): + model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") x_pos: DeviceAccess """Horizontal position""" @@ -24,13 +26,15 @@ class ConfigModel(BaseModel): tilt: DeviceAccess """BPM tilt""" + class BPMTiltOffsetModel(BPMSimpleModel): """ Concrete implementation of BPMModel that simulates a BPM with tilt and offset values. """ + def __init__(self, cfg: ConfigModel): - super().__init__(cfg) + super().__init__(cfg) self.__x_pos = cfg.x_pos self.__y_pos = cfg.y_pos self.__x_offset = cfg.x_offset @@ -46,7 +50,7 @@ def read_tilt(self) -> float: The tilt value of the BPM """ return self.__tilt.get() - + def read_offset(self) -> NDArray: """ Simulate reading the offset values from a BPM. @@ -57,7 +61,7 @@ def read_offset(self) -> NDArray: offsets """ return np.array([self.__x_offset.get(), self.__y_offset.get()]) - + def set_tilt(self, tilt: float): """ Simulate setting the tilt value of a BPM. @@ -70,7 +74,7 @@ def set_tilt(self, tilt: float): None """ self.__tilt.set(tilt) - + def set_offset(self, offset_values: np.ndarray): """ Simulate setting the offset values of a BPM @@ -82,22 +86,22 @@ def set_offset(self, offset_values: np.ndarray): """ self.__x_offset.set(offset_values[0]) self.__y_offset.set(offset_values[1]) - + def get_pos_devices(self) -> list[DeviceAccess]: """ Get device handles used for position reading - + Returns ------- list[DeviceAccess] Array of DeviceAcess """ - return [self.__x_pos,self.__y_pos] + return [self.__x_pos, self.__y_pos] def get_tilt_device(self) -> DeviceAccess: """ Get device handle used for tilt access - + Returns ------- DeviceAccess @@ -108,13 +112,13 @@ def get_tilt_device(self) -> DeviceAccess: def get_offset_devices(self) -> list[DeviceAccess]: """ Get device handles used for offset access - + Returns ------- list[DeviceAccess] Array of DeviceAcess """ - return [self.__x_offset,self.__y_offset] + return [self.__x_offset, self.__y_offset] def __repr__(self): return __pyaml_repr__(self) diff --git a/pyaml/common/abstract.py b/pyaml/common/abstract.py index dfe75db8..487ba6f9 100644 --- a/pyaml/common/abstract.py +++ b/pyaml/common/abstract.py @@ -1,9 +1,10 @@ -from numpy import double from abc import ABCMeta, abstractmethod -from numpy import array + +from numpy import array, double # Float ---------------------------------------------------------------- + class ReadFloatScalar(metaclass=ABCMeta): """ Abstract class providing read access to a scalar double @@ -14,22 +15,25 @@ def get(self) -> double: """Get the value""" pass + @abstractmethod def unit(self) -> str: """Get the unit of the value""" pass - + + class ReadWriteFloatScalar(ReadFloatScalar): """ Abstract class providing read write access to a scalar double - """ + """ + @abstractmethod - def set(self, value:double): + def set(self, value: double): """Set the value""" pass - + # Sets the value and wait that the read value reach the setpoint @abstractmethod - def set_and_wait(self, value:double): + def set_and_wait(self, value: double): """Set the value and wait that setpoint is reached""" pass @@ -44,6 +48,7 @@ def get(self) -> array: """Get the value""" pass + @abstractmethod def unit(self) -> list[str]: """Get the unit of the values""" pass @@ -53,22 +58,25 @@ class ReadWriteFloatArray(ReadFloatScalar): """ Abstract class providing read write access to a vector of double """ + @abstractmethod - def set(self, value:array): + def set(self, value: array): """Set the values""" pass - + # Sets the value and waits that the read value reach the setpoint @abstractmethod - def set_and_wait(self, value:array): + def set_and_wait(self, value: array): """Set the values and wait that setpoints are reached""" pass + class RWMapper(ReadWriteFloatScalar): """ Class mapping a scalar to an element of an array """ - def __init__(self, bind, idx:int): + + def __init__(self, bind, idx: int): self.bind = bind self.idx = idx @@ -77,20 +85,19 @@ def get(self) -> float: return self.bind.get()[self.idx] # Sets the value - def set(self, value:float): + def set(self, value: float): arr = self.bind.get() arr[self.idx] = value self.bind.set(arr) # Sets the value and wait that the read value reach the setpoint - def set_and_wait(self, value:float): + def set_and_wait(self, value: float): raise NotImplementedError("Not implemented yet.") - + # Return the unit def unit(self) -> str: return self.bind.unit()[self.idx] - + # Return the mapped index def index(self) -> int: return self.idx - diff --git a/pyaml/common/abstract_aggregator.py b/pyaml/common/abstract_aggregator.py index 18b62a66..65ddb7cb 100644 --- a/pyaml/common/abstract_aggregator.py +++ b/pyaml/common/abstract_aggregator.py @@ -1,7 +1,8 @@ from abc import ABCMeta, abstractmethod -import numpy.typing as npt import numpy as np +import numpy.typing as npt + class ScalarAggregator(metaclass=ABCMeta): """ diff --git a/pyaml/common/element.py b/pyaml/common/element.py index c5b94d4d..ec810950 100644 --- a/pyaml/common/element.py +++ b/pyaml/common/element.py @@ -1,29 +1,33 @@ -from .exception import PyAMLException +from pydantic import BaseModel, ConfigDict -from pydantic import BaseModel,ConfigDict +from .exception import PyAMLException def __pyaml_repr__(obj): """ Returns a string representation of a pyaml object """ - if hasattr(obj,"_cfg"): - if isinstance(obj,Element): - return repr(obj._cfg).replace("ConfigModel(",obj.__class__.__name__ + "(peer='" + obj.get_peer() + "', ") + if hasattr(obj, "_cfg"): + if isinstance(obj, Element): + return repr(obj._cfg).replace( + "ConfigModel(", + obj.__class__.__name__ + "(peer='" + obj.get_peer() + "', ", + ) else: # no peer - return repr(obj._cfg).replace("ConfigModel",obj.__class__.__name__ ) + return repr(obj._cfg).replace("ConfigModel", obj.__class__.__name__) else: # Default to repr return repr(obj) -class ElementConfigModel(BaseModel): - model_config = ConfigDict(arbitrary_types_allowed=True,extra="forbid") +class ElementConfigModel(BaseModel): + model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") - name : str + name: str """Element name""" + class Element(object): """ Class providing access to one element of a physical or simulated lattice @@ -32,9 +36,10 @@ class Element(object): name: str The unique name identifying the element in the configuration file """ - def __init__(self,name:str): + + def __init__(self, name: str): self.__name: str = name - self._peer = None # Peer: ControlSystem, Simulator + self._peer = None # Peer: ControlSystem, Simulator def get_name(self): """ @@ -42,7 +47,7 @@ def get_name(self): """ return self.__name - def set_energy(self,E:float): + def set_energy(self, E: float): """ Set the instrument energy on this element """ @@ -50,20 +55,23 @@ def set_energy(self,E:float): def check_peer(self): """ - Throws an exception if the element is not attacched to a simulator or to a control system + Throws an exception if the element is not attacched + to a simulator or to a control system """ if self._peer is None: - raise PyAMLException(f"{str(self)} is not attached to a control system or the a simulator") - + raise PyAMLException( + f"{str(self)} is not attachedto a control system or the a simulator" + ) + def get_peer(self) -> str: """ Returns a string representation of peer simulator or control system """ - return "None" if self._peer is None else f"{self._peer.__class__.__name__}:{self._peer.name()}" - + return ( + "None" + if self._peer is None + else f"{self._peer.__class__.__name__}:{self._peer.name()}" + ) + def __repr__(self): return __pyaml_repr__(self) - - - - \ No newline at end of file diff --git a/pyaml/common/element_holder.py b/pyaml/common/element_holder.py index 0287251e..7f54eb24 100644 --- a/pyaml/common/element_holder.py +++ b/pyaml/common/element_holder.py @@ -1,22 +1,25 @@ """ Module handling element references for simulators and control system """ -from .element import Element -from ..magnet.magnet import Magnet -from ..magnet.cfm_magnet import CombinedFunctionMagnet -from ..bpm.bpm import BPM -from ..rf.rf_plant import RFPlant -from ..rf.rf_transmitter import RFTransmitter -from ..arrays.magnet_array import MagnetArray -from ..arrays.cfm_magnet_array import CombinedFunctionMagnetArray + from ..arrays.bpm_array import BPMArray +from ..arrays.cfm_magnet_array import CombinedFunctionMagnetArray from ..arrays.element_array import ElementArray +from ..arrays.magnet_array import MagnetArray +from ..bpm.bpm import BPM from ..common.exception import PyAMLException from ..diagnostics.tune_monitor import BetatronTuneMonitor +from ..magnet.cfm_magnet import CombinedFunctionMagnet +from ..magnet.magnet import Magnet +from ..rf.rf_plant import RFPlant +from ..rf.rf_transmitter import RFTransmitter +from .element import Element + class ElementHolder(object): """ - Class that store references of objects used from both simulators and control system + Class that store references of objects used from both + simulators and control system """ def __init__(self): @@ -35,116 +38,142 @@ def __init__(self): self.__BPM_ARRAYS: dict = {} self.__ELEMENT_ARRAYS: dict = {} - def fill_device(self,elements:list[Element]): - raise "ElementHolder.fill_device() is not subclassed" - - def fill_array(self,arrayName:str,elementNames:list[str],get_func,constructor,ARR:dict): - a = [] - for name in elementNames: - try: - m = get_func(name) - except Exception as err: - raise PyAMLException(f"{constructor.__name__} {arrayName} : {err} @index {len(a)}") from None - if m in a: - raise PyAMLException(f"{constructor.__name__} {arrayName} : duplicate name {name} @index {len(a)}") from None - a.append(m) - ARR[arrayName] = constructor(arrayName,a) - - - def __add(self,array,element:Element): - if element.get_name() in self.__ALL: # Ensure name unicity - raise PyAMLException(f"Duplicate element {element.__class__.__name__} name {element.get_name()}") from None - array[element.get_name()] = element - self.__ALL[element.get_name()] = element - - def __get(self,what,name,array) -> Element: - if name not in array: - raise PyAMLException(f"{what} {name} not defined") - return array[name] - + def fill_device(self, elements: list[Element]): + raise PyAMLException("ElementHolder.fill_device() is not subclassed") + + def fill_array( + self, arrayName: str, elementNames: list[str], get_func, constructor, ARR: dict + ): + a = [] + for name in elementNames: + try: + m = get_func(name) + except Exception as err: + raise PyAMLException( + f"{constructor.__name__} {arrayName} : {err} @index {len(a)}" + ) from None + if m in a: + raise PyAMLException( + f"{constructor.__name__} {arrayName} : " + f"duplicate name {name} @index {len(a)}" + ) from None + a.append(m) + ARR[arrayName] = constructor(arrayName, a) + + def __add(self, array, element: Element): + if element.get_name() in self.__ALL: # Ensure name unicity + raise PyAMLException( + f"Duplicate element {element.__class__.__name__} " + "name {element.get_name()}" + ) from None + array[element.get_name()] = element + self.__ALL[element.get_name()] = element + + def __get(self, what, name, array) -> Element: + if name not in array: + raise PyAMLException(f"{what} {name} not defined") + return array[name] + # Generic elements - def fill_element_array(self,arrayName:str,elementNames:list[str]): - self.fill_array(arrayName,elementNames,self.get_element,ElementArray,self.__ELEMENT_ARRAYS) + def fill_element_array(self, arrayName: str, elementNames: list[str]): + self.fill_array( + arrayName, + elementNames, + self.get_element, + ElementArray, + self.__ELEMENT_ARRAYS, + ) + + def get_element(self, name: str) -> Element: + return self.__get("Element", name, self.__ALL) - def get_element(self,name:str) -> Element: - return self.__get("Element",name,self.__ALL) - - def get_elemens(self,name:str) -> ElementArray: - return self.__get("Element array",name,self.__ELEMENT_ARRAYS) + def get_elemens(self, name: str) -> ElementArray: + return self.__get("Element array", name, self.__ELEMENT_ARRAYS) def get_all_elements(self) -> list[Element]: - return [value for key, value in self.__ALL.items()] - + return [value for key, value in self.__ALL.items()] + # Magnets - - def fill_magnet_array(self,arrayName:str,elementNames:list[str]): - self.fill_array(arrayName,elementNames,self.get_magnet,MagnetArray,self.__MAGNET_ARRAYS) - def get_magnet(self,name:str) -> Magnet: - return self.__get("Magnet",name,self.__MAGNETS) - - def add_magnet(self,m:Magnet): - self.__add(self.__MAGNETS,m) + def fill_magnet_array(self, arrayName: str, elementNames: list[str]): + self.fill_array( + arrayName, elementNames, self.get_magnet, MagnetArray, self.__MAGNET_ARRAYS + ) + + def get_magnet(self, name: str) -> Magnet: + return self.__get("Magnet", name, self.__MAGNETS) + + def add_magnet(self, m: Magnet): + self.__add(self.__MAGNETS, m) - def get_magnets(self,name:str) -> MagnetArray: - return self.__get("Magnet array",name,self.__MAGNET_ARRAYS) + def get_magnets(self, name: str) -> MagnetArray: + return self.__get("Magnet array", name, self.__MAGNET_ARRAYS) def get_all_magnets(self) -> list[Magnet]: - return [value for key, value in self.__MAGNETS.items()] + return [value for key, value in self.__MAGNETS.items()] # Combined Function Magnets - - def fill_cfm_magnet_array(self,arrayName:str,elementNames:list[str]): - self.fill_array(arrayName,elementNames,self.get_cfm_magnet,CombinedFunctionMagnetArray,self.__CFM_MAGNET_ARRAYS) - def get_cfm_magnet(self,name:str) -> Magnet: - return self.__get("CombinedFunctionMagnet",name,self.__CFM_MAGNETS) - - def add_cfm_magnet(self,m:Magnet): - self.__add(self.__CFM_MAGNETS,m) + def fill_cfm_magnet_array(self, arrayName: str, elementNames: list[str]): + self.fill_array( + arrayName, + elementNames, + self.get_cfm_magnet, + CombinedFunctionMagnetArray, + self.__CFM_MAGNET_ARRAYS, + ) - def get_cfm_magnets(self,name:str) -> CombinedFunctionMagnetArray: - return self.__get("CombinedFunctionMagnet array",name,self.__CFM_MAGNET_ARRAYS) + def get_cfm_magnet(self, name: str) -> Magnet: + return self.__get("CombinedFunctionMagnet", name, self.__CFM_MAGNETS) + + def add_cfm_magnet(self, m: Magnet): + self.__add(self.__CFM_MAGNETS, m) + + def get_cfm_magnets(self, name: str) -> CombinedFunctionMagnetArray: + return self.__get( + "CombinedFunctionMagnet array", name, self.__CFM_MAGNET_ARRAYS + ) def get_all_cfm_magnets(self) -> list[CombinedFunctionMagnet]: - return [value for key, value in self.__CFM_MAGNETS.items()] - + return [value for key, value in self.__CFM_MAGNETS.items()] + # BPMs - def fill_bpm_array(self,arrayName:str,elementNames:list[str]): - self.fill_array(arrayName,elementNames,self.get_bpm,BPMArray,self.__BPM_ARRAYS) + def fill_bpm_array(self, arrayName: str, elementNames: list[str]): + self.fill_array( + arrayName, elementNames, self.get_bpm, BPMArray, self.__BPM_ARRAYS + ) - def get_bpm(self,name:str) -> Element: - return self.__get("BPM",name,self.__BPMS) + def get_bpm(self, name: str) -> Element: + return self.__get("BPM", name, self.__BPMS) - def add_bpm(self,bpm:BPM): - self.__add(self.__BPMS,bpm) + def add_bpm(self, bpm: BPM): + self.__add(self.__BPMS, bpm) - def get_bpms(self,name:str) -> BPMArray: - return self.__get("BPM array",name,self.__BPM_ARRAYS) + def get_bpms(self, name: str) -> BPMArray: + return self.__get("BPM array", name, self.__BPM_ARRAYS) def get_all_bpms(self) -> list[BPM]: - return [value for key, value in self.__BPMS.items()] + return [value for key, value in self.__BPMS.items()] # RF - def get_rf_plant(self,name:str) -> RFPlant: - return self.__get("RFPlant",name,self.__RFPLANT) + def get_rf_plant(self, name: str) -> RFPlant: + return self.__get("RFPlant", name, self.__RFPLANT) - def add_rf_plant(self,rf:RFPlant): - self.__add(self.__RFPLANT,rf) + def add_rf_plant(self, rf: RFPlant): + self.__add(self.__RFPLANT, rf) - def add_rf_transnmitter(self,rf:RFTransmitter): - self.__add(self.__RFTRANSMITTER,rf) + def add_rf_transnmitter(self, rf: RFTransmitter): + self.__add(self.__RFTRANSMITTER, rf) - def get_rf_trasnmitter(self,name:str) -> RFTransmitter: - return self.__get("RFTransmitter",name,self.__RFTRANSMITTER) + def get_rf_trasnmitter(self, name: str) -> RFTransmitter: + return self.__get("RFTransmitter", name, self.__RFTRANSMITTER) + # Tune monitor - # Tune monitor - - def get_betatron_tune_monitor(self, name:str) -> BetatronTuneMonitor: - return self.__get("Diagnostic",name,self.__DIAG) + def get_betatron_tune_monitor(self, name: str) -> BetatronTuneMonitor: + return self.__get("Diagnostic", name, self.__DIAG) - def add_betatron_tune_monitor(self, tune_monitor:Element): - self.__add(self.__DIAG,tune_monitor) + def add_betatron_tune_monitor(self, tune_monitor: Element): + self.__add(self.__DIAG, tune_monitor) diff --git a/pyaml/common/exception.py b/pyaml/common/exception.py index ce512c04..3bf9cb78 100644 --- a/pyaml/common/exception.py +++ b/pyaml/common/exception.py @@ -4,6 +4,7 @@ class PyAMLException(Exception): Attributes: message -- explanation of the error """ + def __init__(self, message): super().__init__(message) self.message = message @@ -15,6 +16,7 @@ class PyAMLConfigException(Exception): Attributes: message -- explanation of the error """ + def __init__(self, message): super().__init__(message) self.message = message diff --git a/pyaml/configuration/__init__.py b/pyaml/configuration/__init__.py index 83bbc7ac..2bdf00aa 100644 --- a/pyaml/configuration/__init__.py +++ b/pyaml/configuration/__init__.py @@ -2,5 +2,5 @@ PyAML configuration module """ -from .fileloader import load_accelerator, set_root_folder, get_root_folder from .factory import Factory +from .fileloader import get_root_folder, load_accelerator, set_root_folder diff --git a/pyaml/configuration/csvcurve.py b/pyaml/configuration/csvcurve.py index 004227ab..b802eef0 100644 --- a/pyaml/configuration/csvcurve.py +++ b/pyaml/configuration/csvcurve.py @@ -1,22 +1,23 @@ -from ..configuration import get_root_folder -from ..common.exception import PyAMLException -from .curve import Curve - from pathlib import Path -from pydantic import BaseModel,ConfigDict + import numpy as np +from pydantic import BaseModel, ConfigDict +from ..common.exception import PyAMLException +from ..configuration import get_root_folder +from .curve import Curve # Define the main class name for this module PYAMLCLASS = "CSVCurve" -class ConfigModel(BaseModel): - model_config = ConfigDict(arbitrary_types_allowed=True,extra="forbid") +class ConfigModel(BaseModel): + model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") file: str """CSV file that contains the curve (n rows,2 columns)""" + class CSVCurve(Curve): """ Class for load CSV (x,y) curve @@ -26,18 +27,23 @@ def __init__(self, cfg: ConfigModel): self._cfg = cfg # Load CSV curve - path:Path = get_root_folder() / cfg.file + path: Path = get_root_folder() / cfg.file try: self._curve = np.genfromtxt(path, delimiter=",", dtype=float, loose=False) except ValueError as e: - raise PyAMLException(f"CSVCurve(file='{cfg.file}',dtype=float): {str(e)}") from None + raise PyAMLException( + f"CSVCurve(file='{cfg.file}',dtype=float): {str(e)}" + ) from None _s = np.shape(self._curve) if len(_s) != 2 or _s[1] != 2: - raise PyAMLException(f"CSVCurve(file='{cfg.file}',dtype=float): wrong shape (2,2) expected but got {str(_s)}") + raise PyAMLException( + f"CSVCurve(file='{cfg.file}',dtype=float):" + f"wrong shape (2,2) expected but got {str(_s)}" + ) def get_curve(self) -> np.array: return self._curve def __repr__(self): - return repr(self._cfg).replace("ConfigModel",self.__class__.__name__) + return repr(self._cfg).replace("ConfigModel", self.__class__.__name__) diff --git a/pyaml/configuration/csvmatrix.py b/pyaml/configuration/csvmatrix.py index a70e162e..f8b283e3 100644 --- a/pyaml/configuration/csvmatrix.py +++ b/pyaml/configuration/csvmatrix.py @@ -1,22 +1,23 @@ -from ..configuration import get_root_folder -from .matrix import Matrix -from ..common.exception import PyAMLException - -from pydantic import BaseModel,ConfigDict from pathlib import Path + import numpy as np +from pydantic import BaseModel, ConfigDict +from ..common.exception import PyAMLException +from ..configuration import get_root_folder +from .matrix import Matrix # Define the main class name for this module PYAMLCLASS = "CSVMatrix" -class ConfigModel(BaseModel): - model_config = ConfigDict(arbitrary_types_allowed=True,extra="forbid") +class ConfigModel(BaseModel): + model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") file: str """CSV file that contains the matrix""" + class CSVMatrix(Matrix): """ Class for loading CSV matrix @@ -25,14 +26,16 @@ class CSVMatrix(Matrix): def __init__(self, cfg: ConfigModel): self._cfg = cfg # Load CSV matrix - path:Path = get_root_folder() / cfg.file + path: Path = get_root_folder() / cfg.file try: self._mat = np.genfromtxt(path, delimiter=",", dtype=float, loose=False) except ValueError as e: - raise PyAMLException(f"CSVMatrix(file='{cfg.file}',dtype=float): {str(e)}") from None + raise PyAMLException( + f"CSVMatrix(file='{cfg.file}',dtype=float): {str(e)}" + ) from None def get_matrix(self) -> np.array: return self._mat def __repr__(self): - return repr(self._cfg).replace("ConfigModel",self.__class__.__name__) + return repr(self._cfg).replace("ConfigModel", self.__class__.__name__) diff --git a/pyaml/configuration/curve.py b/pyaml/configuration/curve.py index 6f0ba62a..df73ec01 100644 --- a/pyaml/configuration/curve.py +++ b/pyaml/configuration/curve.py @@ -1,6 +1,8 @@ from abc import ABCMeta, abstractmethod + import numpy as np + class Curve(metaclass=ABCMeta): """ Abstract class providing access to a curve @@ -10,12 +12,12 @@ class Curve(metaclass=ABCMeta): def get_curve(self) -> np.array: """ Returns the curve (n rows,2 columns). - Curve is expected to be monotonic (non-decreasing or non-increasing). + Curve is expected to be monotonic (non-decreasing or non-increasing). """ pass @classmethod - def inverse(cls, curve:np.array) -> np.array: + def inverse(cls, curve: np.array) -> np.array: """ Returns the inverse curve. Basically swap x and y and sort y in ascending order. @@ -26,5 +28,5 @@ def inverse(cls, curve:np.array) -> np.array: Curve to be inverted """ __curve = curve - __sortedCurve = __curve[__curve[:,1].argsort()] - return __sortedCurve[:,[1,0]] + __sortedCurve = __curve[__curve[:, 1].argsort()] + return __sortedCurve[:, [1, 0]] diff --git a/pyaml/configuration/factory.py b/pyaml/configuration/factory.py index 5c9ae23d..757fb33b 100644 --- a/pyaml/configuration/factory.py +++ b/pyaml/configuration/factory.py @@ -1,12 +1,14 @@ # PyAML factory (construct AML objects from config files) +import fnmatch import importlib from threading import Lock -import fnmatch -from ..common.exception import PyAMLConfigException -from ..common.element import Element from pydantic import ValidationError +from ..common.element import Element +from ..common.exception import PyAMLConfigException + + class BuildStrategy: def can_handle(self, module: object, config_dict: dict) -> bool: """Return True if this strategy can handle the module/config.""" @@ -16,6 +18,7 @@ def build(self, module: object, config_dict: dict): """Build the object according to custom logic.""" raise NotImplementedError + class PyAMLFactory: """Singleton factory to build PyAML elements with future compatibility logic.""" @@ -24,7 +27,8 @@ class PyAMLFactory: def __new__(cls): """ - No matter how many times you call PyAMLFactory(), it will be created only once. + No matter how many times you call PyAMLFactory(), + it will be created only once. """ with cls._lock: if cls._instance is None: @@ -41,47 +45,55 @@ def remove_strategy(self, strategy: BuildStrategy): """Register a plugin-based strategy for object creation.""" self._strategies.remove(strategy) - def handle_validation_error(self, e, type_str:str, location_str:str, field_locations:dict): - # Handle pydantic errors - globalMessage = "" - for err in e.errors(): - msg = err['msg'] - field = "" - if len(err['loc'])==2: - field, fieldIdx = err['loc'] - message = f"'{field}.{fieldIdx}': {msg}" - else: - field = err['loc'][0] - message = f"'{field}': {msg}" - if field_locations and field in field_locations: - file, line, col = field_locations[field] - loc = f"{file} at line {line}, colum {col}" - message += f" {loc}" - globalMessage += message - globalMessage += ", " - # Discard pydantic stack trace - raise PyAMLConfigException(f"{globalMessage} for object: '{type_str}' {location_str}") from None - - def build_object(self, d:dict): + def handle_validation_error( + self, e, type_str: str, location_str: str, field_locations: dict + ): + # Handle pydantic errors + globalMessage = "" + for err in e.errors(): + msg = err["msg"] + field = "" + if len(err["loc"]) == 2: + field, fieldIdx = err["loc"] + message = f"'{field}.{fieldIdx}': {msg}" + else: + field = err["loc"][0] + message = f"'{field}': {msg}" + if field_locations and field in field_locations: + file, line, col = field_locations[field] + loc = f"{file} at line {line}, colum {col}" + message += f" {loc}" + globalMessage += message + globalMessage += ", " + # Discard pydantic stack trace + raise PyAMLConfigException( + f"{globalMessage} for object: '{type_str}' {location_str}" + ) from None + + def build_object(self, d: dict): """Build an object from the dict""" - location = d.pop('__location__', None) - field_locations = d.pop('__fieldlocations__', None) + location = d.pop("__location__", None) + field_locations = d.pop("__fieldlocations__", None) location_str = "" if location: file, line, col = location location_str = f"{file} at line {line}, column {col}." - if not isinstance(d,dict): + if not isinstance(d, dict): raise PyAMLConfigException(f"Unexpected object {str(d)} {location_str}") - if not "type" in d: - raise PyAMLConfigException(f"No type specified for {str(type(d))}:{str(d)} {location_str}") + if "type" not in d: + raise PyAMLConfigException( + f"No type specified for {str(type(d))}:{str(d)} {location_str}" + ) type_str = d.pop("type") try: module = importlib.import_module(type_str) except ModuleNotFoundError as ex: # Discard module not found stack trace - raise PyAMLConfigException(f"Module referenced in type cannot be found: '{type_str}' {location_str}") from None + raise PyAMLConfigException( + f"Module referenced in type cannot be found:'{type_str}' {location_str}" + ) from None # Try plugin strategies first for strategy in self._strategies: @@ -91,24 +103,30 @@ def build_object(self, d:dict): self.register_element(obj) return obj except Exception as e: - raise PyAMLConfigException(f"Custom strategy failed {location_str}") from e + raise PyAMLConfigException( + f"Custom strategy failed {location_str}" + ) from e # Default loading strategy # Get the config object config_cls = getattr(module, "ConfigModel", None) if config_cls is None: - raise PyAMLConfigException(f"ConfigModel class '{type_str}.ConfigModel' not found {location_str}") + raise PyAMLConfigException( + f"ConfigModel class '{type_str}.ConfigModel' not found {location_str}" + ) # Get the class name cls_name = getattr(module, "PYAMLCLASS", None) if cls_name is None: - raise PyAMLConfigException(f"PYAMLCLASS definition not found in '{type_str}' {location_str}") + raise PyAMLConfigException( + f"PYAMLCLASS definition not found in '{type_str}' {location_str}" + ) try: # Validate the model cfg = config_cls.model_validate(d) except ValidationError as e: - self.handle_validation_error(e,type_str,location_str,field_locations) + self.handle_validation_error(e, type_str, location_str, field_locations) # Construct and return the object elem_cls = getattr(module, cls_name, None) @@ -119,58 +137,62 @@ def build_object(self, d:dict): obj = elem_cls(cfg) self.register_element(obj) except Exception as e: - raise PyAMLConfigException(f"{str(e)} when creating '{type_str}.{cls_name}' {location_str}") + raise PyAMLConfigException( + f"{str(e)} when creating '{type_str}.{cls_name}' {location_str}" + ) from e return obj - def depth_first_build(self, d): - """Main factory function (Depth-first factory)""" - - if isinstance(d,list): - # list can be a list of objects or a list of native types - l = [] - for index, e in enumerate(d): - if isinstance(e,dict) or isinstance(e,list): - obj = self.depth_first_build(e) - l.append(obj) - else: - l.append(e) - return l - - elif isinstance(d,dict): - for key, value in d.items(): - if not key == "__fieldlocations__": - if isinstance(value,dict) or isinstance(value,list): - obj = self.depth_first_build(value) - # Replace the inner dict by the object itself - d[key]=obj - - # We are now on leaf (no nested object), we can construct - return self.build_object(d) - - raise PyAMLConfigException(f"Unexpected element found. 'dict' or 'list' expected but got '{d.__class__.__name__}'") + """Main factory function (Depth-first factory)""" + + if isinstance(d, list): + # list can be a list of objects or a list of native types + l = [] + for _index, e in enumerate(d): + if isinstance(e, dict) or isinstance(e, list): + obj = self.depth_first_build(e) + l.append(obj) + else: + l.append(e) + return l + + elif isinstance(d, dict): + for key, value in d.items(): + if not key == "__fieldlocations__": + if isinstance(value, dict) or isinstance(value, list): + obj = self.depth_first_build(value) + # Replace the inner dict by the object itself + d[key] = obj + + # We are now on leaf (no nested object), we can construct + return self.build_object(d) + + raise PyAMLConfigException( + "Unexpected element found. 'dict' or 'list' expected " + "but got '{d.__class__.__name__}'" + ) def register_element(self, elt): - if isinstance(elt,Element): + if isinstance(elt, Element): name = elt.get_name() if name in self._elements: raise PyAMLConfigException(f"element {name} already defined") self._elements[name] = elt - - def get_element(self, name:str): + def get_element(self, name: str): if name not in self._elements: raise PyAMLConfigException(f"element {name} not defined") return self._elements[name] - - def get_elements_by_name(self,wildcard:str) -> list[Element]: - return [e for k,e in self._elements.items() if fnmatch.fnmatch(k, wildcard)] - def get_elements_by_type(self,type) -> list[Element]: - return [e for k,e in self._elements.items() if isinstance(e,type)] + def get_elements_by_name(self, wildcard: str) -> list[Element]: + return [e for k, e in self._elements.items() if fnmatch.fnmatch(k, wildcard)] + + def get_elements_by_type(self, type) -> list[Element]: + return [e for k, e in self._elements.items() if isinstance(e, type)] def clear(self): self._elements.clear() + Factory = PyAMLFactory() diff --git a/pyaml/configuration/fileloader.py b/pyaml/configuration/fileloader.py index 5b0d92b7..8f5686d4 100644 --- a/pyaml/configuration/fileloader.py +++ b/pyaml/configuration/fileloader.py @@ -1,20 +1,21 @@ # PyAML config file loader -import logging -import json -from typing import Union, TYPE_CHECKING -from pathlib import Path +import collections.abc import io +import json +import logging import os +from pathlib import Path +from typing import TYPE_CHECKING, Union import yaml -from yaml.loader import SafeLoader from yaml import CLoader from yaml.constructor import ConstructorError -import collections.abc +from yaml.loader import SafeLoader -from .. import PyAMLException from pyaml.configuration.factory import Factory +from .. import PyAMLException + if TYPE_CHECKING: from pyaml.accelerator import Accelerator @@ -40,58 +41,72 @@ def get_root_folder() -> Path: """ return ROOT["path"] + class PyAMLConfigCyclingException(PyAMLException): - - def __init__(self, error_filename:str, path_stack:list[Path]): + def __init__(self, error_filename: str, path_stack: list[Path]): self.error_filename = error_filename parent_file_stack = [parent_path.name for parent_path in path_stack] - super().__init__(f"Circular file inclusion of {error_filename}. File list before reaching it: {parent_file_stack}") + super().__init__( + f"Circular file inclusion of {error_filename}. " + f"File list before reaching it: {parent_file_stack}" + ) + pass -def load_accelerator(filename:str, use_fast_loader:bool = False) -> "Accelerator": - """ Load an accelerator from file.""" - # Asume that all files are referenced from folder where main AML file is stored +def load_accelerator(filename: str, use_fast_loader: bool = False) -> "Accelerator": + """Load an accelerator from file.""" + + # Asume that all files are referenced from + # folder where main AML file is stored if not os.path.exists(filename): - raise PyAMLException(f"{filename} file not found") + raise PyAMLException(f"{filename} file not found") rootfolder = os.path.abspath(os.path.dirname(filename)) set_root_folder(rootfolder) - config_dict = load(os.path.basename(filename),None,use_fast_loader) + config_dict = load(os.path.basename(filename), None, use_fast_loader) aml = Factory.depth_first_build(config_dict) Factory.clear() return aml -def load(filename:str, paths_stack:list=None, use_fast_loader:bool = False) -> Union[dict,list]: + +def load( + filename: str, paths_stack: list = None, use_fast_loader: bool = False +) -> Union[dict, list]: """Load recursively a configuration setup""" if filename.endswith(".yaml") or filename.endswith(".yml"): l = YAMLLoader(filename, paths_stack, use_fast_loader) elif filename.endswith(".json"): l = JSONLoader(filename, paths_stack, use_fast_loader) else: - raise PyAMLException(f"{filename} File format not supported (only .yaml .yml or .json)") + raise PyAMLException( + f"{filename} File format not supported (only .yaml .yml or .json)" + ) return l.load() + # Expand condition def hasToExpand(value): - return isinstance(value,str) and any(value.endswith(suffix) for suffix in accepted_suffixes) + return isinstance(value, str) and any( + value.endswith(suffix) for suffix in accepted_suffixes + ) # Loader base class (nested files expansion) class Loader: - - def __init__(self, filename:str, parent_path_stack:list[Path]): - self.path:Path = get_root_folder() / filename - self.files_stack:list[Path] = [] + def __init__(self, filename: str, parent_path_stack: list[Path]): + self.path: Path = get_root_folder() / filename + self.files_stack: list[Path] = [] if parent_path_stack: - if any(self.path.samefile(parent_path) for parent_path in parent_path_stack): + if any( + self.path.samefile(parent_path) for parent_path in parent_path_stack + ): raise PyAMLConfigCyclingException(filename, parent_path_stack) self.files_stack.extend(parent_path_stack) self.files_stack.append(self.path) - # Recursively expand a dict - def expand_dict(self,d:dict): + def expand_dict(self, d: dict): for key, value in d.items(): try: if hasToExpand(value): @@ -99,8 +114,8 @@ def expand_dict(self,d:dict): else: self.expand(value) except PyAMLConfigCyclingException as pyaml_ex: - location = d.pop('__location__', None) - field_locations = d.pop('__fieldlocations__', None) + location = d.pop("__location__", None) + field_locations = d.pop("__fieldlocations__", None) location_str = "" if location: file, line, col = location @@ -108,33 +123,36 @@ def expand_dict(self,d:dict): location = field_locations[key] file, line, col = location location_str = f" in {file} at line {line}, column {col}" - raise PyAMLException(f"Circular file inclusion of {pyaml_ex.error_filename}{location_str}") from pyaml_ex + raise PyAMLException( + "Circular file inclusion " + f"of {pyaml_ex.error_filename}{location_str}" + ) from pyaml_ex # Recursively expand a list - def expand_list(self,l:list): - for idx,value in enumerate(l): + def expand_list(self, l: list): + for idx, value in enumerate(l): if hasToExpand(value): l[idx] = load(value, self.files_stack) else: self.expand(value) # Recursively expand an object - def expand(self,obj: Union[dict,list]): - if isinstance(obj,dict): + def expand(self, obj: Union[dict, list]): + if isinstance(obj, dict): self.expand_dict(obj) - elif isinstance(obj,list): + elif isinstance(obj, list): self.expand_list(obj) - return obj + return obj # Load a file - def load(self) -> Union[dict,list]: + def load(self) -> Union[dict, list]: raise PyAMLException(str(self.path) + ": load() method not implemented") -class SafeLineLoader(SafeLoader): +class SafeLineLoader(SafeLoader): def __init__(self, stream): super().__init__(stream) - self.filename = stream.name if isinstance(stream,io.TextIOWrapper) else "" + self.filename = stream.name if isinstance(stream, io.TextIOWrapper) else "" def construct_mapping(self, node, deep=False): mapping = {} @@ -143,39 +161,53 @@ def construct_mapping(self, node, deep=False): for key_node, value_node in node.value: key = self.construct_object(key_node, deep=deep) if not isinstance(key, collections.abc.Hashable): - raise ConstructorError("while constructing a mapping", node.start_mark, - "found unhashable key", key_node.start_mark) + raise ConstructorError( + "while constructing a mapping", + node.start_mark, + "found unhashable key", + key_node.start_mark, + ) value = self.construct_object(value_node, deep=deep) mapping[key] = value - field_mapping[key] = (self.filename, value_node.start_mark.line + 1 , value_node.start_mark.column + 1) + field_mapping[key] = ( + self.filename, + value_node.start_mark.line + 1, + value_node.start_mark.column + 1, + ) # Add location information inside the dict - mapping['__location__'] = (self.filename, node.start_mark.line + 1 , node.start_mark.column + 1) - mapping['__fieldlocations__'] = field_mapping + mapping["__location__"] = ( + self.filename, + node.start_mark.line + 1, + node.start_mark.column + 1, + ) + mapping["__fieldlocations__"] = field_mapping return mapping + # YAML loader class YAMLLoader(Loader): - def __init__(self, filename: str, parent_paths_stack:list,use_fast_loader:bool): + def __init__(self, filename: str, parent_paths_stack: list, use_fast_loader: bool): super().__init__(filename, parent_paths_stack) self._loader = SafeLineLoader if not use_fast_loader else CLoader self.use_fast_loader = use_fast_loader - def load(self) -> Union[dict,list]: + def load(self) -> Union[dict, list]: logger.log(logging.DEBUG, f"Loading YAML file '{self.path}'") with open(self.path) as file: try: - return self.expand(yaml.load(file,Loader=self._loader)) + return self.expand(yaml.load(file, Loader=self._loader)) except yaml.YAMLError as e: raise PyAMLException(str(self.path) + ": " + str(e)) from e + # JSON loader class JSONLoader(Loader): - def __init__(self, filename: str, parent_paths_stack:list,use_fast_loader:bool): + def __init__(self, filename: str, parent_paths_stack: list, use_fast_loader: bool): super().__init__(filename, parent_paths_stack) self.use_fast_loader = False - def load(self) -> Union[dict,list]: + def load(self) -> Union[dict, list]: logger.log(logging.DEBUG, f"Loading JSON file '{self.path}'") with open(self.path) as file: try: diff --git a/pyaml/configuration/matrix.py b/pyaml/configuration/matrix.py index bdff486f..dc5b9ac7 100644 --- a/pyaml/configuration/matrix.py +++ b/pyaml/configuration/matrix.py @@ -1,10 +1,13 @@ from abc import ABCMeta, abstractmethod + from numpy import array + class Matrix(metaclass=ABCMeta): """ Abstract class providing access to a matrix """ + @abstractmethod def get_matrix(self) -> array: """Returns the matrix""" diff --git a/pyaml/configuration/models.py b/pyaml/configuration/models.py index 48e3b5dd..8a45de5c 100644 --- a/pyaml/configuration/models.py +++ b/pyaml/configuration/models.py @@ -1,8 +1,8 @@ -from pathlib import Path -from typing import Dict, List, Tuple -import json import importlib +import json import types +from pathlib import Path +from typing import Dict, List, Tuple import yaml from pydantic import BaseModel, Field, model_validator @@ -14,15 +14,14 @@ class PyAmlBaseModel(BaseModel): module_path: str = Field(validation_alias="type", default="") model_config = { - "populate_by_name": True, # allows setting "module_path" instead of `type` directly + "populate_by_name": True, + # allows setting "module_path" instead of `type` directly } class ConfigBase(PyAmlBaseModel): - @staticmethod def load_config_from_file(path: str | Path): - path = Path(path).resolve() match path.suffix: @@ -33,7 +32,8 @@ def load_config_from_file(path: str | Path): data = json.loads(path.read_text()) case _: raise ValueError( - f"Unsupported file type: {path.suffix}. Expected .yaml, .yml, or .json." + f"Unsupported file type: {path.suffix}." + "Expected .yaml, .yml, or .json." ) if isinstance(data, list): @@ -63,17 +63,18 @@ def validate_sub_config( cfg = cls.load_config_from_file(full_path) if not isinstance(cfg, expected_class): raise ValueError( - f"Invalid type for '{field_name}': {type(cfg)}. Expected {expected_class}." + f"Invalid type for '{field_name}': {type(cfg)}." + "Expected {expected_class}." ) return cfg else: raise ValueError( - f"Invalid type for '{field_name}': {type(v)}. Expected str or a type inherited from {expected_class}." + f"Invalid type for '{field_name}': {type(v)}." + "Expected str or a type inherited from {expected_class}." ) def _parse_yaml_file(path: str | Path) -> Tuple[Path, Dict | List]: - path = get_config_file_path(path) with path.open() as f: data = yaml.safe_load(f) @@ -82,7 +83,6 @@ def _parse_yaml_file(path: str | Path) -> Tuple[Path, Dict | List]: def _parse_json_file(path: str | Path) -> Tuple[Path, Dict | List]: - path = get_config_file_path(path) data = json.loads(Path(path).read_text()) @@ -92,7 +92,6 @@ def _parse_json_file(path: str | Path) -> Tuple[Path, Dict | List]: def construct_config_from_dict( path: Path, data: Dict ) -> Tuple[str, types.ModuleType, ConfigBase]: - module_path = data.get("type") if not module_path: raise ValueError(f"No 'type' field found in {path}") @@ -110,7 +109,6 @@ def construct_config_from_dict( def construct_element( path: Path, module_path: str, module: types.ModuleType, cfg: ConfigBase ): - elem_class_name = module_path.split(".")[-1] elem_cls = getattr(module, elem_class_name, None) if elem_cls is None: @@ -129,7 +127,6 @@ def recursively_construct_element_from_cfg(cfg: ConfigBase): def construct_element_from_dict(path: Path, data: Dict) -> BaseModel: - module_path, module, cfg = construct_config_from_dict(path, data) elem = construct_element(path, module_path, module, cfg) diff --git a/pyaml/control/abstract_impl.py b/pyaml/control/abstract_impl.py index 275e42fa..2ebfaf3f 100644 --- a/pyaml/control/abstract_impl.py +++ b/pyaml/control/abstract_impl.py @@ -1,27 +1,29 @@ +import numpy as np +from numpy import double +from numpy.typing import NDArray + +from ..bpm.bpm_model import BPMModel from ..common import abstract -from ..control.deviceaccesslist import DeviceAccessList +from ..common.abstract_aggregator import ScalarAggregator from ..control.deviceaccess import DeviceAccess -from ..magnet.model import MagnetModel +from ..control.deviceaccesslist import DeviceAccessList from ..magnet.magnet import Magnet -from ..bpm.bpm_model import BPMModel +from ..magnet.model import MagnetModel from ..rf.rf_plant import RFPlant from ..rf.rf_transmitter import RFTransmitter -from ..common.abstract_aggregator import ScalarAggregator -from numpy import double -import numpy as np -from numpy.typing import NDArray -#------------------------------------------------------------------------------ +# ------------------------------------------------------------------------------ + class CSScalarAggregator(ScalarAggregator): """ Basic control system aggregator for a list of scalar values """ - def __init__(self, devs:DeviceAccessList): + def __init__(self, devs: DeviceAccessList): self._devs = devs - def add_devices(self, devices:DeviceAccess | list[DeviceAccess] ): + def add_devices(self, devices: DeviceAccess | list[DeviceAccess]): self._devs.add_devices(devices) def set(self, value: NDArray[np.float64]): @@ -38,51 +40,63 @@ def readback(self) -> np.array: def unit(self) -> str: return self._devs.unit() - + def nb_device(self) -> int: return self._devs.__len__() -#------------------------------------------------------------------------------ + +# ------------------------------------------------------------------------------ + class CSStrengthScalarAggregator(CSScalarAggregator): """ Control system aggregator for a list of magnet strengths. - This aggregator is in charge of computing hardware setpoints and applying them without overlap. - When virtual magnets exported from combined function mangets are present (RWMapper), + This aggregator is in charge of computing hardware setpoints + and applying them without overlap. + When virtual magnets exported from combined function mangets are present (RWMapper), the aggregator prevents to apply several times the same power supply setpoint. """ - def __init__(self, peer:CSScalarAggregator): - CSScalarAggregator.__init__(self,peer._devs) - self.__models: list[MagnetModel] = [] # List of magnet model - self.__modelToMagnet: list[list[tuple[int,int]]] = [] # strengths indexing - self.__nbMagnet = 0 # Number of magnet strengths + def __init__(self, peer: CSScalarAggregator): + CSScalarAggregator.__init__(self, peer._devs) + self.__models: list[MagnetModel] = [] # List of magnet model + self.__modelToMagnet: list[list[tuple[int, int]]] = [] # strengths indexing + self.__nbMagnet = 0 # Number of magnet strengths - def add_magnet(self, magnet:Magnet): - # Incoming magnet can be a magnet exported from a CombinedFunctionMagnet or simple magnet. + def add_magnet(self, magnet: Magnet): + # Incoming magnet can be a magnet exported from + # a CombinedFunctionMagnet or simple magnet. # All magnets exported from a same CombinedFunctionMagnet share the same model # TODO: check that strength is supported (m.strength may be None) - strengthIndex = magnet.strength.index() if isinstance(magnet.strength,abstract.RWMapper) else 0 + strengthIndex = ( + magnet.strength.index() + if isinstance(magnet.strength, abstract.RWMapper) + else 0 + ) if magnet.model not in self.__models: index = len(self.__models) self.__models.append(magnet.model) - self.__modelToMagnet.append([(self.__nbMagnet,strengthIndex)]) + self.__modelToMagnet.append([(self.__nbMagnet, strengthIndex)]) self._devs.add_devices(magnet.model.get_devices()) else: index = self.__models.index(magnet.model) - self.__modelToMagnet[index].append((self.__nbMagnet,strengthIndex)) + self.__modelToMagnet[index].append((self.__nbMagnet, strengthIndex)) self.__nbMagnet += 1 def set(self, value: NDArray[np.float64]): - allHardwareValues = self._devs.get() # Read all hardware setpoints + allHardwareValues = self._devs.get() # Read all hardware setpoints newHardwareValues = np.zeros(self.nb_device()) hardwareIndex = 0 - for modelIndex,model in enumerate(self.__models): + for modelIndex, model in enumerate(self.__models): nbDev = len(model.get_devices()) - mStrengths = model.compute_strengths( allHardwareValues[hardwareIndex:hardwareIndex+nbDev] ) - for (valueIdx,strengthIdx) in self.__modelToMagnet[modelIndex]: + mStrengths = model.compute_strengths( + allHardwareValues[hardwareIndex : hardwareIndex + nbDev] + ) + for valueIdx, strengthIdx in self.__modelToMagnet[modelIndex]: mStrengths[strengthIdx] = value[valueIdx] - newHardwareValues[hardwareIndex:hardwareIndex+nbDev] = model.compute_hardware_values(mStrengths) + newHardwareValues[hardwareIndex : hardwareIndex + nbDev] = ( + model.compute_hardware_values(mStrengths) + ) hardwareIndex += nbDev self._devs.set(newHardwareValues) @@ -90,25 +104,29 @@ def set_and_wait(self, value: NDArray[np.float64]): raise NotImplementedError("Not implemented yet.") def get(self) -> NDArray[np.float64]: - allHardwareValues = self._devs.get() # Read all hardware setpoints + allHardwareValues = self._devs.get() # Read all hardware setpoints allStrength = np.zeros(self.__nbMagnet) hardwareIndex = 0 - for modelIndex,model in enumerate(self.__models): + for modelIndex, model in enumerate(self.__models): nbDev = len(model.get_devices()) - mStrengths = model.compute_strengths( allHardwareValues[hardwareIndex:hardwareIndex+nbDev] ) - for (valueIdx,strengthIdx) in self.__modelToMagnet[modelIndex]: + mStrengths = model.compute_strengths( + allHardwareValues[hardwareIndex : hardwareIndex + nbDev] + ) + for valueIdx, strengthIdx in self.__modelToMagnet[modelIndex]: allStrength[valueIdx] = mStrengths[strengthIdx] hardwareIndex += nbDev return allStrength def readback(self) -> np.array: - allHardwareValues = self._devs.readback() # Read all hardware readback + allHardwareValues = self._devs.readback() # Read all hardware readback allStrength = np.zeros(self.__nbMagnet) hardwareIndex = 0 - for modelIndex,model in enumerate(self.__models): + for modelIndex, model in enumerate(self.__models): nbDev = len(model.get_devices()) - mStrengths = model.compute_strengths( allHardwareValues[hardwareIndex:hardwareIndex+nbDev] ) - for (valueIdx,strengthIdx) in self.__modelToMagnet[modelIndex]: + mStrengths = model.compute_strengths( + allHardwareValues[hardwareIndex : hardwareIndex + nbDev] + ) + for valueIdx, strengthIdx in self.__modelToMagnet[modelIndex]: allStrength[valueIdx] = mStrengths[strengthIdx] hardwareIndex += nbDev return allStrength @@ -117,20 +135,22 @@ def unit(self) -> str: return self._devs.unit() -#------------------------------------------------------------------------------ +# ------------------------------------------------------------------------------ + class RWHardwareScalar(abstract.ReadWriteFloatScalar): """ - Class providing read write access to a magnet of a control system (in hardware units) + Class providing read write access to a magnet + of a control system (in hardware units) """ - def __init__(self, model:MagnetModel): + def __init__(self, model: MagnetModel): self.__model = model def get(self) -> float: return self.__model.read_hardware_values()[0] - - def set(self, value:float): + + def set(self, value: float): self.__model.send_hardware_values(np.array([value])) def set_and_wait(self, value: double): @@ -139,14 +159,16 @@ def set_and_wait(self, value: double): def unit(self) -> str: return self.__model.get_hardware_units()[0] -#------------------------------------------------------------------------------ + +# ------------------------------------------------------------------------------ + class RWStrengthScalar(abstract.ReadWriteFloatScalar): """ Class providing read write access to a strength of a control system """ - def __init__(self, model:MagnetModel): + def __init__(self, model: MagnetModel): self.__model = model # Gets the value @@ -155,25 +177,29 @@ def get(self) -> float: return self.__model.compute_strengths(currents)[0] # Sets the value - def set(self, value:float): + def set(self, value: float): current = self.__model.compute_hardware_values([value]) self.__model.send_hardware_values(current) # Sets the value and wait that the read value reach the setpoint - def set_and_wait(self, value:float): + def set_and_wait(self, value: float): raise NotImplementedError("Not implemented yet.") # Gets the unit of the value def unit(self) -> str: return self.__model.get_strength_units()[0] -#------------------------------------------------------------------------------ + +# ------------------------------------------------------------------------------ + class RWHardwareArray(abstract.ReadWriteFloatArray): """ - Class providing read write access to a magnet array of a control system (in hardware units) + Class providing read write access to a magnet array + of a control system (in hardware units) """ - def __init__(self, model:MagnetModel): + + def __init__(self, model: MagnetModel): self.__model = model # Gets the value @@ -181,25 +207,27 @@ def get(self) -> np.array: return self.__model.read_hardware_values() # Sets the value - def set(self, value:np.array): + def set(self, value: np.array): self.__model.send_hardware_values(value) - + # Sets the value and waits that the read value reach the setpoint - def set_and_wait(self, value:np.array): + def set_and_wait(self, value: np.array): raise NotImplementedError("Not implemented yet.") - # Gets the unit of the value def unit(self) -> list[str]: return self.__model.get_hardware_units() -#------------------------------------------------------------------------------ + +# ------------------------------------------------------------------------------ + class RWStrengthArray(abstract.ReadWriteFloatArray): """ Class providing read write access to magnet strengths of a control system """ - def __init__(self, model:MagnetModel): + + def __init__(self, model: MagnetModel): self.__model = model # Gets the value @@ -209,25 +237,28 @@ def get(self) -> np.array: return str # Sets the value - def set(self, value:np.array): + def set(self, value: np.array): cur = self.__model.compute_hardware_values(value) self.__model.send_hardware_values(cur) # Sets the value and waits that the read value reach the setpoint - def set_and_wait(self, value:np.array): + def set_and_wait(self, value: np.array): raise NotImplementedError("Not implemented yet.") # Gets the unit of the value def unit(self) -> list[str]: return self.__model.get_strength_units() -#------------------------------------------------------------------------------ + +# ------------------------------------------------------------------------------ + class RBpmArray(abstract.ReadFloatArray): """ Class providing read access to a BPM array of a control system """ - def __init__(self, model:BPMModel): + + def __init__(self, model: BPMModel): self.__model = model # Gets the value @@ -238,34 +269,41 @@ def get(self) -> np.array: def unit(self) -> str: return self.__model.get_pos_devices()[0].unit() -#------------------------------------------------------------------------------ + +# ------------------------------------------------------------------------------ + class RWBpmTiltScalar(abstract.ReadFloatScalar): """ Class providing read access to a BPM tilt of a control system """ - def __init__(self, model:BPMModel): + + def __init__(self, model: BPMModel): self.__model = model # Gets the value def get(self) -> float: return self.__model.read_tilt() - def set(self, value:float): + def set(self, value: float): self.__model.set_tilt(value) def set_and_wait(self, value: NDArray[np.float64]): raise NotImplementedError("Not implemented yet.") + # Gets the unit of the value def unit(self) -> str: return self.__model.get_tilt_device().unit() -#------------------------------------------------------------------------------ + +# ------------------------------------------------------------------------------ + class RWBpmOffsetArray(abstract.ReadWriteFloatArray): """ Class providing read write access to a BPM offset of a control system """ + def __init__(self, model: BPMModel): self.__model = model @@ -275,81 +313,93 @@ def get(self) -> NDArray[np.float64]: # Sets the value def set(self, value: NDArray[np.float64]): - self.__model.set_offset(value) + self.__model.set_offset(value) + def set_and_wait(self, value: NDArray[np.float64]): raise NotImplementedError("Not implemented yet.") + # Gets the unit of the value def unit(self) -> str: return self.__model.get_offset_devices()[0].unit() -#------------------------------------------------------------------------------ + +# ------------------------------------------------------------------------------ + class RWRFVoltageScalar(abstract.ReadWriteFloatScalar): """ - Class providing read write access to cavity voltage for a transmitter of a control system. + Class providing read write access to cavity voltage + for a transmitter of a control system. """ - def __init__(self, transmitter:RFTransmitter): + def __init__(self, transmitter: RFTransmitter): self.__transmitter = transmitter - def get(self) -> float: + def get(self) -> float: return self.__transmitter._cfg.voltage.get() - - def set(self,value:float): + + def set(self, value: float): return self.__transmitter._cfg.voltage.set(value) - def set_and_wait(self, value:float): + def set_and_wait(self, value: float): raise NotImplementedError("Not implemented yet.") - + def unit(self) -> str: return self.__transmitter._cfg.voltage.unit() -#------------------------------------------------------------------------------ + +# ------------------------------------------------------------------------------ + class RWRFPhaseScalar(abstract.ReadWriteFloatScalar): """ - Class providing read write access to cavity phase for a transmitter of a control system. + Class providing read write access to cavity phase + for a transmitter of a control system. """ - def __init__(self, transmitter:RFTransmitter): + def __init__(self, transmitter: RFTransmitter): self.__transmitter = transmitter - def get(self) -> float: + def get(self) -> float: return self.__transmitter._cfg.phase.get() - - def set(self,value:float): + + def set(self, value: float): return self.__transmitter._cfg.phase.set(value) - def set_and_wait(self, value:float): + def set_and_wait(self, value: float): raise NotImplementedError("Not implemented yet.") - + def unit(self) -> str: return self.__transmitter._cfg.phase.unit() - -#------------------------------------------------------------------------------ + + +# ------------------------------------------------------------------------------ + class RWRFFrequencyScalar(abstract.ReadWriteFloatScalar): """ Class providing read write access to RF frequency of a control system. """ - def __init__(self, rf:RFPlant ): + def __init__(self, rf: RFPlant): self.__rf = rf def get(self) -> float: # Serialized cavity has the same frequency return self.__rf._cfg.masterclock.get() - def set(self,value:float): + def set(self, value: float): return self.__rf._cfg.masterclock.set(value) - def set_and_wait(self, value:float): + def set_and_wait(self, value: float): raise NotImplementedError("Not implemented yet.") def unit(self) -> str: return self.__rf._cfg.masterclock.unit() -#------------------------------------------------------------------------------ + +# ------------------------------------------------------------------------------ + class RBetatronTuneArray(abstract.ReadFloatArray): """ @@ -361,9 +411,12 @@ def __init__(self, tune_monitor): def get(self) -> NDArray: # Return horizontal and vertical betatron tunes as a NumPy array - return np.array([self.__tune_monitor._cfg.tune_h.get(), - self.__tune_monitor._cfg.tune_v.get()]) + return np.array( + [ + self.__tune_monitor._cfg.tune_h.get(), + self.__tune_monitor._cfg.tune_v.get(), + ] + ) def unit(self) -> str: return self.__tune_monitor._cfg.tune_v.unit() - diff --git a/pyaml/control/controlsystem.py b/pyaml/control/controlsystem.py index d2ab5c4b..816d163f 100644 --- a/pyaml/control/controlsystem.py +++ b/pyaml/control/controlsystem.py @@ -1,22 +1,34 @@ from abc import ABCMeta, abstractmethod -from ..common.element_holder import ElementHolder + +from ..bpm.bpm import BPM from ..common.abstract import RWMapper +from ..common.abstract_aggregator import ScalarAggregator from ..common.element import Element -from ..control.abstract_impl import RWHardwareScalar,RWHardwareArray,RWStrengthScalar,RWStrengthArray -from ..bpm.bpm import BPM +from ..common.element_holder import ElementHolder +from ..configuration.factory import Factory +from ..control.abstract_impl import ( + CSScalarAggregator, + CSStrengthScalarAggregator, + RBetatronTuneArray, + RBpmArray, + RWBpmOffsetArray, + RWBpmTiltScalar, + RWHardwareArray, + RWHardwareScalar, + RWRFFrequencyScalar, + RWRFPhaseScalar, + RWRFVoltageScalar, + RWStrengthArray, + RWStrengthScalar, +) from ..diagnostics.tune_monitor import BetatronTuneMonitor -from ..control.abstract_impl import RWBpmTiltScalar,RWBpmOffsetArray, RBpmArray -from ..control.abstract_impl import RWRFFrequencyScalar,RWRFVoltageScalar,RWRFPhaseScalar -from ..control.abstract_impl import CSScalarAggregator,CSStrengthScalarAggregator -from ..common.abstract_aggregator import ScalarAggregator -from ..control.abstract_impl import RBetatronTuneArray -from ..magnet.magnet import Magnet from ..magnet.cfm_magnet import CombinedFunctionMagnet -from ..rf.rf_plant import RFPlant,RWTotalVoltage +from ..magnet.magnet import Magnet +from ..rf.rf_plant import RFPlant, RWTotalVoltage from ..rf.rf_transmitter import RFTransmitter -from ..configuration.factory import Factory -class ControlSystem(ElementHolder,metaclass=ABCMeta): + +class ControlSystem(ElementHolder, metaclass=ABCMeta): """ Abstract class providing access to a control system float variable """ @@ -33,7 +45,7 @@ def init_cs(self): def name(self) -> str: """Return control system name (i.e. live)""" pass - + @abstractmethod def scalar_aggregator(self) -> str | None: """Returns the module name used for handling aggregator of DeviceAccess""" @@ -46,26 +58,32 @@ def vector_aggregator(self) -> str | None: def create_scalar_aggregator(self) -> ScalarAggregator: mod = self.scalar_aggregator() - agg = Factory.build_object({"type":mod}) if mod is not None else None + agg = Factory.build_object({"type": mod}) if mod is not None else None return CSScalarAggregator(agg) - - def create_magnet_strength_aggregator(self,magnets:list[Magnet]) -> ScalarAggregator: + + def create_magnet_strength_aggregator( + self, magnets: list[Magnet] + ) -> ScalarAggregator: agg = CSStrengthScalarAggregator(self.create_scalar_aggregator()) for m in magnets: agg.add_magnet(m) return agg - def create_magnet_harddware_aggregator(self,magnets:list[Magnet]) -> ScalarAggregator: - # When working in hardware space, 1 single power supply device per multipolar strength is required + def create_magnet_harddware_aggregator( + self, magnets: list[Magnet] + ) -> ScalarAggregator: + """When working in hardware space, 1 single power + supply device per multipolar strength is required + """ agg = self.create_scalar_aggregator() for m in magnets: if not m.model.has_hardware(): - return None - psIndex = m.hardware.index() if isinstance(m.hardware,RWMapper) else 0 + return None + psIndex = m.hardware.index() if isinstance(m.hardware, RWMapper) else 0 agg.add_devices(m.model.get_devices()[psIndex]) return agg - - def create_bpm_aggregators(self,bpms:list[BPM]) -> list[ScalarAggregator]: + + def create_bpm_aggregators(self, bpms: list[BPM]) -> list[ScalarAggregator]: agg = self.create_scalar_aggregator() aggh = self.create_scalar_aggregator() aggv = self.create_scalar_aggregator() @@ -74,13 +92,12 @@ def create_bpm_aggregators(self,bpms:list[BPM]) -> list[ScalarAggregator]: agg.add_devices(devs) aggh.add_devices(devs[0]) aggv.add_devices(devs[1]) - return [agg,aggh,aggv] - + return [agg, aggh, aggv] - def set_energy(self,E:float): + def set_energy(self, E: float): """ Sets the energy on magnets belonging to this control system - + Parameters ---------- E : float @@ -89,57 +106,59 @@ def set_energy(self,E:float): # Needed by energy dependant element (i.e. magnet coil current calculation) for m in self.get_all_elements(): m.set_energy(E) - - def fill_device(self,elements:list[Element]): + + def fill_device(self, elements: list[Element]): """ - Fill device of this control system with Element coming from the configuration file - + Fill device of this control system with Element + coming from the configuration file + Parameters ---------- elements : list[Element] - List of elements coming from the configuration file to attach to this control system - """ + List of elements coming from the configuration + file to attach to this control system + """ for e in elements: - if isinstance(e,Magnet): - current = RWHardwareScalar(e.model) if e.model.has_hardware() else None - strength = RWStrengthScalar(e.model) if e.model.has_physics() else None - # Create a unique ref for this control system - m = e.attach(self,strength, current) - self.add_magnet(m) - - elif isinstance(e,CombinedFunctionMagnet): - currents = RWHardwareArray(e.model) if e.model.has_hardware() else None - strengths = RWStrengthArray(e.model) if e.model.has_physics() else None - # Create unique refs the cfm and each of its function for this control system - ms = e.attach(self,strengths,currents) - self.add_cfm_magnet(ms[0]) - for m in ms[1:]: - self.add_magnet(m) - - elif isinstance(e,BPM): - tilt = RWBpmTiltScalar(e.model) - offsets = RWBpmOffsetArray(e.model) - positions = RBpmArray(e.model) - e = e.attach(self,positions, offsets, tilt) - self.add_bpm(e) - - - elif isinstance(e,RFPlant): - attachedTrans: list[RFTransmitter] = [] - if e._cfg.transmitters: - for t in e._cfg.transmitters: - voltage = RWRFVoltageScalar(t) - phase = RWRFPhaseScalar(t) - nt = t.attach(self,voltage,phase) - self.add_rf_transnmitter(nt) - attachedTrans.append(nt) - - frequency = RWRFFrequencyScalar(e) - voltage = RWTotalVoltage(attachedTrans) if e._cfg.transmitters else None - ne = e.attach(self,frequency,voltage) - self.add_rf_plant(ne) - - elif isinstance(e,BetatronTuneMonitor): - betatron_tune = RBetatronTuneArray(e) - e = e.attach(self,betatron_tune) - self.add_betatron_tune_monitor(e) + if isinstance(e, Magnet): + current = RWHardwareScalar(e.model) if e.model.has_hardware() else None + strength = RWStrengthScalar(e.model) if e.model.has_physics() else None + # Create a unique ref for this control system + m = e.attach(self, strength, current) + self.add_magnet(m) + + elif isinstance(e, CombinedFunctionMagnet): + currents = RWHardwareArray(e.model) if e.model.has_hardware() else None + strengths = RWStrengthArray(e.model) if e.model.has_physics() else None + # Create unique refs the cfm and + # each of its function for this control system + ms = e.attach(self, strengths, currents) + self.add_cfm_magnet(ms[0]) + for m in ms[1:]: + self.add_magnet(m) + + elif isinstance(e, BPM): + tilt = RWBpmTiltScalar(e.model) + offsets = RWBpmOffsetArray(e.model) + positions = RBpmArray(e.model) + e = e.attach(self, positions, offsets, tilt) + self.add_bpm(e) + + elif isinstance(e, RFPlant): + attachedTrans: list[RFTransmitter] = [] + if e._cfg.transmitters: + for t in e._cfg.transmitters: + voltage = RWRFVoltageScalar(t) + phase = RWRFPhaseScalar(t) + nt = t.attach(self, voltage, phase) + self.add_rf_transnmitter(nt) + attachedTrans.append(nt) + + frequency = RWRFFrequencyScalar(e) + voltage = RWTotalVoltage(attachedTrans) if e._cfg.transmitters else None + ne = e.attach(self, frequency, voltage) + self.add_rf_plant(ne) + + elif isinstance(e, BetatronTuneMonitor): + betatron_tune = RBetatronTuneArray(e) + e = e.attach(self, betatron_tune) + self.add_betatron_tune_monitor(e) diff --git a/pyaml/control/device.py b/pyaml/control/device.py index 3e5c5872..5df9858f 100644 --- a/pyaml/control/device.py +++ b/pyaml/control/device.py @@ -1,5 +1,5 @@ import numpy as np -from pydantic import BaseModel,Field +from pydantic import BaseModel, Field from .deviceaccess import DeviceAccess from .readback_value import Value @@ -7,8 +7,8 @@ # Define the main class name for this module PYAMLCLASS = "Device" -class ConfigModel(BaseModel): +class ConfigModel(BaseModel): setpoint: str """Name of control system device value (i.e. a power supply current)""" readback: str @@ -16,9 +16,10 @@ class ConfigModel(BaseModel): unit: str """Value unit""" + class Device(DeviceAccess): """ - Class that implements a default device class that just prints out + Class that implements a default device class that just prints out values (Debugging purpose) """ diff --git a/pyaml/control/deviceaccess.py b/pyaml/control/deviceaccess.py index 5d506827..41218c68 100644 --- a/pyaml/control/deviceaccess.py +++ b/pyaml/control/deviceaccess.py @@ -2,8 +2,10 @@ from typing import Union import numpy.typing as npt + from .readback_value import Value + class DeviceAccess(metaclass=ABCMeta): """ Abstract class providing access to a control system float variable diff --git a/pyaml/control/deviceaccesslist.py b/pyaml/control/deviceaccesslist.py index f3363201..de61b539 100644 --- a/pyaml/control/deviceaccesslist.py +++ b/pyaml/control/deviceaccesslist.py @@ -1,28 +1,31 @@ from abc import ABCMeta, abstractmethod -import numpy.typing as npt import numpy as np -from .readback_value import Value +import numpy.typing as npt + from .deviceaccess import DeviceAccess +from .readback_value import Value + -class DeviceAccessList(list[DeviceAccess],metaclass=ABCMeta): +class DeviceAccessList(list[DeviceAccess], metaclass=ABCMeta): """ Abstract class providing access to a list of control system float variable """ @abstractmethod - def add_devices(self, devices:DeviceAccess | list[DeviceAccess]): - """Add a DeviceAccess to this list""" + def add_devices(self, devices: DeviceAccess | list[DeviceAccess]): + """Add a DeviceAccess to this list""" pass @abstractmethod def get_devices(self) -> DeviceAccess | list[DeviceAccess]: - """Get the DeviceAccess list""" + """Get the DeviceAccess list""" pass @abstractmethod def set(self, value: npt.NDArray[np.float64]): - """Write a list of control system device variable (i.e. a power supply currents)""" + """Write a list of control system device variable + (i.e. a power supply currents)""" pass @abstractmethod diff --git a/pyaml/control/readback_value.py b/pyaml/control/readback_value.py index 3a32a0e8..c0dc08aa 100644 --- a/pyaml/control/readback_value.py +++ b/pyaml/control/readback_value.py @@ -1,20 +1,21 @@ - from datetime import datetime -from typing import Union - from enum import Enum, auto +from typing import Union import numpy as np + class Quality(Enum): VALID = auto() INVALID = auto() ALARM = auto() CHANGING = auto() WARNING = auto() + def __str__(self): return self.name + class Value: """ Represents a numerical value with associated quality and timestamp. @@ -32,7 +33,12 @@ class Value: Timestamp associated with the value. Defaults to current time. """ - def __init__(self, value: Union[float, int, np.ndarray], quality: Quality = Quality.VALID, timestamp: datetime = None): + def __init__( + self, + value: Union[float, int, np.ndarray], + quality: Quality = Quality.VALID, + timestamp: datetime = None, + ): self.value = value self.quality = quality self.timestamp = timestamp or datetime.now() @@ -46,7 +52,10 @@ def __repr__(self): str Human-readable representation. """ - return f"Value({self.value}, quality='{self.quality}', timestamp='{self.timestamp}')" + return ( + f"Value({self.value}, quality='{self.quality}'," + + f"timestamp='{self.timestamp}')" + ) def __float__(self): """ diff --git a/pyaml/diagnostics/tune_monitor.py b/pyaml/diagnostics/tune_monitor.py index 12c66a5a..6a851c32 100644 --- a/pyaml/diagnostics/tune_monitor.py +++ b/pyaml/diagnostics/tune_monitor.py @@ -1,6 +1,7 @@ -from ..common.element import Element, ElementConfigModel from ..common.abstract import ReadFloatArray +from ..common.element import Element, ElementConfigModel from ..control.deviceaccess import DeviceAccess + try: from typing import Self # Python 3.11+ except ImportError: @@ -9,18 +10,20 @@ PYAMLCLASS = "BetatronTuneMonitor" -class ConfigModel(ElementConfigModel): - model_config = ConfigDict(arbitrary_types_allowed=True,extra="forbid") +class ConfigModel(ElementConfigModel): + model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") tune_h: DeviceAccess """Horizontal betatron tune""" tune_v: DeviceAccess """Vertical betatron tune""" + class BetatronTuneMonitor(Element): """ - Class providing access to a betatron tune monitor of a physical or simulated lattice. + Class providing access to a betatron tune monitor + of a physical or simulated lattice. The monitor provides horizontal and vertical betatron tune measurements. """ @@ -30,10 +33,10 @@ def __init__(self, cfg: ConfigModel): Parameters ---------- cfg : ConfigModel - Configuration for the BetatronTuneMonitor, including + Configuration for the BetatronTuneMonitor, including device access for horizontal and vertical tunes. """ - + super().__init__(cfg.name) self._cfg = cfg self.__tune = None @@ -48,4 +51,3 @@ def attach(self, peer, betatron_tune: ReadFloatArray) -> Self: obj.__tune = betatron_tune obj._peer = peer return obj - diff --git a/pyaml/lattice/abstract_impl.py b/pyaml/lattice/abstract_impl.py index f7bc2519..b8591ef2 100644 --- a/pyaml/lattice/abstract_impl.py +++ b/pyaml/lattice/abstract_impl.py @@ -1,19 +1,20 @@ +import at +import numpy as np +from numpy.typing import NDArray +from scipy.constants import speed_of_light + from ..common import abstract +from ..common.abstract_aggregator import ScalarAggregator from ..common.exception import PyAMLException from ..magnet.model import MagnetModel -from .polynom_info import PolynomInfo from ..rf.rf_plant import RFPlant from ..rf.rf_transmitter import RFTransmitter -from ..common.abstract_aggregator import ScalarAggregator - -import numpy as np -import at -from scipy.constants import speed_of_light -from numpy.typing import NDArray +from .polynom_info import PolynomInfo # TODO handle serialized magnets -#------------------------------------------------------------------------------ +# ------------------------------------------------------------------------------ + class RWHardwareScalar(abstract.ReadWriteFloatScalar): """ @@ -21,7 +22,9 @@ class RWHardwareScalar(abstract.ReadWriteFloatScalar): Hardware unit is converted from strength using the magnet model """ - def __init__(self, elements:list[at.Element], poly:PolynomInfo, model:MagnetModel): + def __init__( + self, elements: list[at.Element], poly: PolynomInfo, model: MagnetModel + ): self.__model = model self.__elements = elements self.__poly = elements[0].__getattribute__(poly.attName) @@ -30,25 +33,29 @@ def __init__(self, elements:list[at.Element], poly:PolynomInfo, model:MagnetMode def get(self) -> float: s = self.__poly[self.__polyIdx] * self.__elements[0].Length return self.__model.compute_hardware_values([s])[0] - - def set(self,value:float): + + def set(self, value: float): s = self.__model.compute_strengths([value])[0] self.__poly[self.__polyIdx] = s / self.__elements[0].Length - def set_and_wait(self, value:float): + def set_and_wait(self, value: float): raise NotImplementedError("Not implemented yet.") - + def unit(self) -> str: return self.__model.get_hardware_units()[0] - -#------------------------------------------------------------------------------ + + +# ------------------------------------------------------------------------------ + class RWStrengthScalar(abstract.ReadWriteFloatScalar): """ Class providing read write access to a strength of a simulator """ - def __init__(self, elements:list[at.Element], poly:PolynomInfo, model:MagnetModel): + def __init__( + self, elements: list[at.Element], poly: PolynomInfo, model: MagnetModel + ): self.__model = model self.__elements = elements self.__poly = elements[0].__getattribute__(poly.attName) @@ -59,18 +66,20 @@ def get(self) -> float: return self.__poly[self.__polyIdx] * self.__elements[0].Length # Sets the value - def set(self, value:float): + def set(self, value: float): self.__poly[self.__polyIdx] = value / self.__elements[0].Length # Sets the value and wait that the read value reach the setpoint - def set_and_wait(self, value:float): + def set_and_wait(self, value: float): raise NotImplementedError("Not implemented yet.") # Gets the unit of the value def unit(self) -> str: return self.__model.get_strength_units()[0] -#------------------------------------------------------------------------------ + +# ------------------------------------------------------------------------------ + class RWHardwareArray(abstract.ReadWriteFloatArray): """ @@ -78,7 +87,9 @@ class RWHardwareArray(abstract.ReadWriteFloatArray): Hardware units are converted from strengths using the magnet model """ - def __init__(self, elements:list[at.Element], poly:list[PolynomInfo], model:MagnetModel): + def __init__( + self, elements: list[at.Element], poly: list[PolynomInfo], model: MagnetModel + ): self.__elements = elements self.__poly = [] self.__polyIdx = [] @@ -88,7 +99,7 @@ def __init__(self, elements:list[at.Element], poly:list[PolynomInfo], model:Magn self.__polyIdx.append(p.index) # Gets the value - def get(self) -> np.array: + def get(self) -> np.array: nbStrength = len(self.__poly) s = np.zeros(nbStrength) for i in range(nbStrength): @@ -96,28 +107,32 @@ def get(self) -> np.array: return self.__model.compute_hardware_values(s) # Sets the value - def set(self, value:np.array): + def set(self, value: np.array): nbStrength = len(self.__poly) s = self.__model.compute_strengths(value) for i in range(nbStrength): self.__poly[i][self.__polyIdx[i]] = s[i] / self.__elements[0].Length # Sets the value and wait that the read value reach the setpoint - def set_and_wait(self, value:np.array): + def set_and_wait(self, value: np.array): raise NotImplementedError("Not implemented yet.") # Gets the unit of the value def unit(self) -> list[str]: return self.__model.get_hardware_units() -#------------------------------------------------------------------------------ + +# ------------------------------------------------------------------------------ + class RWStrengthArray(abstract.ReadWriteFloatArray): """ Class providing read write access to a strength (array) of a simulator """ - def __init__(self, elements:list[at.Element], poly:list[PolynomInfo], model:MagnetModel): + def __init__( + self, elements: list[at.Element], poly: list[PolynomInfo], model: MagnetModel + ): self.__elements = elements self.__poly = [] self.__polyIdx = [] @@ -127,7 +142,7 @@ def __init__(self, elements:list[at.Element], poly:list[PolynomInfo], model:Magn self.__polyIdx.append(p.index) # Gets the value - def get(self) -> np.array: + def get(self) -> np.array: nbStrength = len(self.__poly) s = np.zeros(nbStrength) for i in range(nbStrength): @@ -135,22 +150,23 @@ def get(self) -> np.array: return s # Sets the value - def set(self, value:np.array): + def set(self, value: np.array): nbStrength = len(self.__poly) s = np.zeros(nbStrength) for i in range(nbStrength): self.__poly[i][self.__polyIdx[i]] = value[i] / self.__elements[0].Length # Sets the value and wait that the read value reach the setpoint - def set_and_wait(self, value:np.array): + def set_and_wait(self, value: np.array): raise NotImplementedError("Not implemented yet.") # Gets the unit of the value def unit(self) -> list[str]: return self.__model.get_strength_units() - -#------------------------------------------------------------------------------ + +# ------------------------------------------------------------------------------ + class BPMScalarAggregator(ScalarAggregator): """ @@ -178,9 +194,11 @@ def readback(self) -> np.array: return self.get() def unit(self) -> str: - return 'm' - -#------------------------------------------------------------------------------ + return "m" + + +# ------------------------------------------------------------------------------ + class BPMHScalarAggregator(BPMScalarAggregator): """ @@ -191,7 +209,9 @@ def get(self) -> np.array: _, orbit = at.find_orbit(self.lattice, refpts=self.refpts) return orbit[:, 0] -#------------------------------------------------------------------------------ + +# ------------------------------------------------------------------------------ + class BPMVScalarAggregator(BPMScalarAggregator): """ @@ -202,7 +222,9 @@ def get(self) -> np.array: _, orbit = at.find_orbit(self.lattice, refpts=self.refpts) return orbit[:, 2] -#------------------------------------------------------------------------------ + +# ------------------------------------------------------------------------------ + class RBpmArray(abstract.ReadFloatArray): """ @@ -215,7 +237,6 @@ class RBpmArray(abstract.ReadFloatArray): def __init__(self, element: at.Element, lattice: at.Lattice): self.__element = element self.__lattice = lattice - # Gets the value def get(self) -> np.array: @@ -225,31 +246,33 @@ def get(self) -> np.array: # Gets the unit of the value def unit(self) -> str: - return 'm' + return "m" + + +# ------------------------------------------------------------------------------ -#------------------------------------------------------------------------------ class RWBpmOffsetArray(abstract.ReadWriteFloatArray): """ - Class providing read write access to a BPM offset (array) of a simulator. + Class providing read write access to a BPM offset (array) of a simulator. Offset in pyAT is defined in Offset attribute as a 2-element array. """ - def __init__(self, element:at.Element): + def __init__(self, element: at.Element): self.__element = element try: - self.__offset = element.__getattribute__('Offset') + self.__offset = element.__getattribute__("Offset") except AttributeError: self.__offset = None # Gets the value - def get(self) -> np.array: + def get(self) -> np.array: if self.__offset is None: raise PyAMLException("Element does not have an Offset attribute.") return self.__offset # Sets the value - def set(self, value:np.array): + def set(self, value: np.array): if self.__offset is None: raise PyAMLException("Element does not have an Offset attribute.") if len(value) != 2: @@ -257,14 +280,16 @@ def set(self, value:np.array): self.__offset = value # Sets the value and wait that the read value reach the setpoint - def set_and_wait(self, value:np.array): + def set_and_wait(self, value: np.array): raise NotImplementedError("Not implemented yet.") # Gets the unit of the value def unit(self) -> str: - return 'm' # Assuming all offsets are in m + return "m" # Assuming all offsets are in m + + +# ------------------------------------------------------------------------------ -#------------------------------------------------------------------------------ class RWBpmTiltScalar(abstract.ReadWriteFloatScalar): """ @@ -272,10 +297,10 @@ class RWBpmTiltScalar(abstract.ReadWriteFloatScalar): pyAT is defined in Rotation attribute as a first element. """ - def __init__(self, element:at.Element): + def __init__(self, element: at.Element): self.__element = element try: - self.__tilt = element.__getattribute__('Rotation')[0] + self.__tilt = element.__getattribute__("Rotation")[0] except AttributeError: self.__tilt = None @@ -286,97 +311,111 @@ def get(self) -> float: return self.__tilt # Sets the value - def set(self, value:float, ): + def set( + self, + value: float, + ): self.__tilt = value - self.__element.__setattr__('Rotation', [value, None, None]) + self.__element.__setattr__("Rotation", [value, None, None]) # Sets the value and wait that the read value reach the setpoint - def set_and_wait(self, value:float): + def set_and_wait(self, value: float): raise NotImplementedError("Not implemented yet.") # Gets the unit of the value def unit(self) -> str: - return 'rad' # Assuming BPM tilts are in rad + return "rad" # Assuming BPM tilts are in rad + + +# ------------------------------------------------------------------------------ -#------------------------------------------------------------------------------ class RWRFVoltageScalar(abstract.ReadWriteFloatScalar): """ - Class providing read write access to a cavity voltage of a simulator for a given RF trasnmitter. + Class providing read write access to a cavity voltage + of a simulator for a given RF trasnmitter. """ - def __init__(self, elements:list[at.Element]): + def __init__(self, elements: list[at.Element]): self.__elements = elements def get(self) -> float: sum = 0 - for idx,e in enumerate(self.__elements): + for _idx, e in enumerate(self.__elements): sum += e.Voltage return sum - - def set(self,value:float): + + def set(self, value: float): v = value / len(self.__elements) for e in self.__elements: e.Voltage = v - def set_and_wait(self, value:float): + def set_and_wait(self, value: float): raise NotImplementedError("Not implemented yet.") - + def unit(self) -> str: return "V" -#------------------------------------------------------------------------------ + +# ------------------------------------------------------------------------------ + class RWRFPhaseScalar(abstract.ReadWriteFloatScalar): """ - Class providing read write access to a cavity phase of a simulator for a given RF trasnmitter. + Class providing read write access to a cavity phase of + a simulator for a given RF trasnmitter. """ - def __init__(self, elements:list[at.Element]): + def __init__(self, elements: list[at.Element]): self.__elements = elements def get(self) -> float: - # Assume that all cavities of this transmitter have the same Time Lag and Frequency - wavelength = speed_of_light / self.__elements[0].Frequency - return (wavelength / self.__elements[0].TimeLag) * 2.0 * np.pi - - def set(self,value:float): + # Assume that all cavities of this transmitter + # have the same Time Lag and Frequency + wavelength = speed_of_light / self.__elements[0].Frequency + return (wavelength / self.__elements[0].TimeLag) * 2.0 * np.pi + + def set(self, value: float): wavelength = speed_of_light / self.__elements[0].Frequency for e in self.__elements: e.TimeLag = wavelength * value / (2.0 * np.pi) - def set_and_wait(self, value:float): + def set_and_wait(self, value: float): raise NotImplementedError("Not implemented yet.") - + def unit(self) -> str: return "rad" - -#------------------------------------------------------------------------------ + + +# ------------------------------------------------------------------------------ + class RWRFFrequencyScalar(abstract.ReadWriteFloatScalar): """ Class providing read write access to RF frequency of a simulator. """ - def __init__(self, elements:list[at.Element], harmonics:list[float]): + def __init__(self, elements: list[at.Element], harmonics: list[float]): self.__elements = elements self.__harm = harmonics def get(self) -> float: # Serialized cavity has the same frequency return self.__elements[0].Frequency - - def set(self,value:float): - for idx,e in enumerate(self.__elements): + + def set(self, value: float): + for idx, e in enumerate(self.__elements): e.Frequency = value * self.__harm[idx] - def set_and_wait(self, value:float): + def set_and_wait(self, value: float): raise NotImplementedError("Not implemented yet.") - + def unit(self) -> str: return "Hz" -#------------------------------------------------------------------------------ + +# ------------------------------------------------------------------------------ + class RWRFATFrequencyScalar(abstract.ReadWriteFloatScalar): """ @@ -389,17 +428,19 @@ def __init__(self, ring: at.Lattice): def get(self) -> float: return self.__ring.get_rf_frequency() - - def set(self,value:float): + + def set(self, value: float): self.__ring.set_rf_frequency(value) - def set_and_wait(self, value:float): + def set_and_wait(self, value: float): raise NotImplementedError("Not implemented yet.") - + def unit(self) -> str: - return 'Hz' - -#------------------------------------------------------------------------------ + return "Hz" + + +# ------------------------------------------------------------------------------ + class RWRFATotalVoltageScalar(abstract.ReadWriteFloatScalar): """ @@ -411,17 +452,19 @@ def __init__(self, ring: at.Lattice): def get(self) -> float: return self.__ring.get_rf_voltage() - - def set(self,value:float): + + def set(self, value: float): self.__ring.set_rf_voltage(value) - def set_and_wait(self, value:float): + def set_and_wait(self, value: float): raise NotImplementedError("Not implemented yet.") - + def unit(self) -> str: - return 'V' + return "V" + + +# ------------------------------------------------------------------------------ -#------------------------------------------------------------------------------ class RBetatronTuneArray(abstract.ReadFloatArray): """ @@ -430,10 +473,9 @@ class RBetatronTuneArray(abstract.ReadFloatArray): def __init__(self, ring: at.Lattice): self.__ring = ring - + def get(self) -> float: return self.__ring.get_tune()[:2] def unit(self) -> str: - return '1' - + return "1" diff --git a/pyaml/lattice/attribute_linker.py b/pyaml/lattice/attribute_linker.py index f0dd97ef..ee541346 100644 --- a/pyaml/lattice/attribute_linker.py +++ b/pyaml/lattice/attribute_linker.py @@ -2,7 +2,11 @@ from pydantic import ConfigDict from pyaml.common.element import Element -from pyaml.lattice.lattice_elements_linker import LinkerIdentifier, LinkerConfigModel, LatticeElementsLinker +from pyaml.lattice.lattice_elements_linker import ( + LatticeElementsLinker, + LinkerConfigModel, + LinkerIdentifier, +) PYAMLCLASS = "PyAtAttributeElementsLinker" @@ -21,7 +25,8 @@ class ConfigModel(LinkerConfigModel): Pydantic configuration allowing arbitrary field types and forbidding unexpected extra keys. """ - model_config = ConfigDict(arbitrary_types_allowed=True,extra="forbid") + + model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") attribute_name: str @@ -36,7 +41,7 @@ class PyAtAttributeIdentifier(LinkerIdentifier): a unique reference to one or more PyAT elements. """ - def __init__(self, attribute_name:str, identifier): + def __init__(self, attribute_name: str, identifier): self.attribute_name = attribute_name self.identifier = identifier @@ -57,12 +62,16 @@ class PyAtAttributeElementsLinker(LatticeElementsLinker): The configuration model for the linking strategy. """ - def __init__(self, config_model:ConfigModel): + def __init__(self, config_model: ConfigModel): super().__init__(config_model) def get_element_identifier(self, element: Element) -> LinkerIdentifier: - return PyAtAttributeIdentifier(self.linker_config_model.attribute_name, element.get_name()) + return PyAtAttributeIdentifier( + self.linker_config_model.attribute_name, element.get_name() + ) - def _test_at_element(self, identifier: PyAtAttributeIdentifier, element: at.Element) -> bool: + def _test_at_element( + self, identifier: PyAtAttributeIdentifier, element: at.Element + ) -> bool: attr_value = getattr(element, identifier.attribute_name, None) return attr_value == identifier.identifier diff --git a/pyaml/lattice/lattice_elements_linker.py b/pyaml/lattice/lattice_elements_linker.py index 297365f0..188049a4 100644 --- a/pyaml/lattice/lattice_elements_linker.py +++ b/pyaml/lattice/lattice_elements_linker.py @@ -23,7 +23,8 @@ class LinkerConfigModel(BaseModel): Pydantic configuration allowing arbitrary field types and forbidding unexpected extra keys. """ - model_config = ConfigDict(arbitrary_types_allowed=True,extra="forbid") + + model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") class LinkerIdentifier(metaclass=ABCMeta): @@ -36,6 +37,7 @@ class LinkerIdentifier(metaclass=ABCMeta): Subclasses should define the fields and logic necessary to represent a unique reference to one or more PyAT elements. """ + pass @@ -57,19 +59,21 @@ class LatticeElementsLinker(metaclass=ABCMeta): Reference to the PyAT lattice handled by this linker. """ - def __init__(self, linker_config_model:LinkerConfigModel): + def __init__(self, linker_config_model: LinkerConfigModel): self.linker_config_model = linker_config_model - self.lattice:Lattice = None + self.lattice: Lattice = None - def set_lattice(self, lattice:Lattice): + def set_lattice(self, lattice: Lattice): self.lattice = lattice @abstractmethod - def _test_at_element(self, identifier: LinkerIdentifier, element:at.Element) -> bool: + def _test_at_element( + self, identifier: LinkerIdentifier, element: at.Element + ) -> bool: pass @abstractmethod - def get_element_identifier(self, element:Element) -> LinkerIdentifier: + def get_element_identifier(self, element: Element) -> LinkerIdentifier: pass def _iter_matches(self, identifier: LinkerIdentifier) -> Iterable[at.Element]: @@ -78,7 +82,9 @@ def _iter_matches(self, identifier: LinkerIdentifier) -> Iterable[at.Element]: if self._test_at_element(identifier, elem): yield elem - def get_at_elements(self,element_id:LinkerIdentifier|list[LinkerIdentifier]) -> list[at.Element]: + def get_at_elements( + self, element_id: LinkerIdentifier | list[LinkerIdentifier] + ) -> list[at.Element]: """Return a list of PyAT elements matching the given identifiers. This method should resolve one or multiple PyAML identifiers @@ -117,7 +123,7 @@ def get_at_elements(self,element_id:LinkerIdentifier|list[LinkerIdentifier]) -> ) return results - def get_at_element(self, element_id:LinkerIdentifier) -> at.Element: + def get_at_element(self, element_id: LinkerIdentifier) -> at.Element: """Return a single PyAT element matching the given identifier. Parameters @@ -137,4 +143,6 @@ def get_at_element(self, element_id:LinkerIdentifier) -> at.Element: """ for elem in self._iter_matches(element_id): return elem - raise PyAMLException(f"No PyAT element found for FamName: {element_id.__repr__()}") \ No newline at end of file + raise PyAMLException( + f"No PyAT element found for FamName: {element_id.__repr__()}" + ) diff --git a/pyaml/lattice/polynom_info.py b/pyaml/lattice/polynom_info.py index fe8557ff..4d4f2aeb 100644 --- a/pyaml/lattice/polynom_info.py +++ b/pyaml/lattice/polynom_info.py @@ -2,15 +2,16 @@ Class providing polynom information """ + class PolynomInfo: """ Polynom information """ - def __init__(self,attName:str,index:int): + def __init__(self, attName: str, index: int): """ Construct a polynom information object - + Parameters ---------- attName : str @@ -19,4 +20,4 @@ def __init__(self,attName:str,index:int): Polynomial coefficients index (sarting from 0) """ self.attName = attName - self.index = index \ No newline at end of file + self.index = index diff --git a/pyaml/lattice/simulator.py b/pyaml/lattice/simulator.py index 96b2a3d1..07868b2b 100644 --- a/pyaml/lattice/simulator.py +++ b/pyaml/lattice/simulator.py @@ -1,34 +1,51 @@ -from .attribute_linker import PyAtAttributeElementsLinker, ConfigModel as PyAtAttrLinkerConfigModel -from .lattice_elements_linker import LatticeElementsLinker -from ..configuration import get_root_folder -from ..common.element import Element -from ..magnet.magnet import Magnet +from pathlib import Path + +import at +from pydantic import BaseModel, ConfigDict + from ..bpm.bpm import BPM +from ..common.abstract_aggregator import ScalarAggregator +from ..common.element import Element +from ..common.element_holder import ElementHolder +from ..common.exception import PyAMLException +from ..configuration import get_root_folder from ..diagnostics.tune_monitor import BetatronTuneMonitor +from ..lattice.abstract_impl import ( + BPMHScalarAggregator, + BPMScalarAggregator, + BPMVScalarAggregator, + RBetatronTuneArray, + RBpmArray, + RWBpmOffsetArray, + RWBpmTiltScalar, + RWHardwareArray, + RWHardwareScalar, + RWRFATFrequencyScalar, + RWRFATotalVoltageScalar, + RWRFFrequencyScalar, + RWRFPhaseScalar, + RWRFVoltageScalar, + RWStrengthArray, + RWStrengthScalar, +) from ..magnet.cfm_magnet import CombinedFunctionMagnet -from ..rf.rf_plant import RFPlant,RWTotalVoltage +from ..magnet.magnet import Magnet +from ..rf.rf_plant import RFPlant, RWTotalVoltage from ..rf.rf_transmitter import RFTransmitter -from ..lattice.abstract_impl import RWHardwareScalar,RWHardwareArray -from ..lattice.abstract_impl import RWStrengthScalar,RWStrengthArray -from ..lattice.abstract_impl import RWRFFrequencyScalar,RWRFVoltageScalar,RWRFPhaseScalar -from ..lattice.abstract_impl import RWRFATFrequencyScalar,RWRFATotalVoltageScalar -from ..common.element_holder import ElementHolder -from ..common.abstract_aggregator import ScalarAggregator -from ..lattice.abstract_impl import RBetatronTuneArray -from ..lattice.abstract_impl import RWBpmTiltScalar,RWBpmOffsetArray, RBpmArray -from ..lattice.abstract_impl import BPMHScalarAggregator,BPMScalarAggregator,BPMVScalarAggregator -from ..common.exception import PyAMLException - -from pydantic import BaseModel,ConfigDict -import at -from pathlib import Path +from .attribute_linker import ( + ConfigModel as PyAtAttrLinkerConfigModel, +) +from .attribute_linker import ( + PyAtAttributeElementsLinker, +) +from .lattice_elements_linker import LatticeElementsLinker # Define the main class name for this module PYAMLCLASS = "Simulator" -class ConfigModel(BaseModel): - model_config = ConfigDict(arbitrary_types_allowed=True,extra="forbid") +class ConfigModel(BaseModel): + model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") name: str """Simulator name""" @@ -39,6 +56,7 @@ class ConfigModel(BaseModel): linker: LatticeElementsLinker = None """The linker configuration model""" + class Simulator(ElementHolder): """ Class that implements access to AT simulator @@ -47,120 +65,152 @@ class Simulator(ElementHolder): def __init__(self, cfg: ConfigModel): super().__init__() self._cfg = cfg - self._linker = cfg.linker if cfg.linker else PyAtAttributeElementsLinker(PyAtAttrLinkerConfigModel(attribute_name="FamName")) - path:Path = get_root_folder() / cfg.lattice - - if(self._cfg.mat_key is None): - self.ring = at.load_lattice(path) + self._linker = ( + cfg.linker + if cfg.linker + else PyAtAttributeElementsLinker( + PyAtAttrLinkerConfigModel(attribute_name="FamName") + ) + ) + path: Path = get_root_folder() / cfg.lattice + + if self._cfg.mat_key is None: + self.ring = at.load_lattice(path) else: - self.ring = at.load_lattice(path,mat_key=f"{self._cfg.mat_key}") + self.ring = at.load_lattice(path, mat_key=f"{self._cfg.mat_key}") self._linker.set_lattice(self.ring) def name(self) -> str: - return self._cfg.name - + return self._cfg.name + def get_lattice(self) -> at.Lattice: - return self.ring - - def set_energy(self,E:float): - self.ring.energy = E - # Needed by energy dependant element (i.e. magnet coil current calculation) - for m in self.get_all_elements(): - m.set_energy(E) - - def create_magnet_strength_aggregator(self,magnets:list[Magnet]) -> ScalarAggregator: + return self.ring + + def set_energy(self, E: float): + self.ring.energy = E + # Needed by energy dependant element (i.e. magnet coil current calculation) + for m in self.get_all_elements(): + m.set_energy(E) + + def create_magnet_strength_aggregator( + self, magnets: list[Magnet] + ) -> ScalarAggregator: # No magnet aggregator for simulator return None - - def create_magnet_harddware_aggregator(self,magnets:list[Magnet]) -> ScalarAggregator: + + def create_magnet_harddware_aggregator( + self, magnets: list[Magnet] + ) -> ScalarAggregator: # No magnet aggregator for simulator return None - def create_bpm_aggregators(self,bpms:list[BPM]) -> list[ScalarAggregator]: + def create_bpm_aggregators(self, bpms: list[BPM]) -> list[ScalarAggregator]: agg = BPMScalarAggregator(self.get_lattice()) aggh = BPMHScalarAggregator(self.get_lattice()) aggv = BPMVScalarAggregator(self.get_lattice()) for b in bpms: - e = self.get_at_elems(b)[0] - agg.add_elem(e) - aggh.add_elem(e) - aggv.add_elem(e) - return [agg,aggh,aggv] - - def fill_device(self,elements:list[Element]): - for e in elements: - # Need conversion to physics unit to work with simulator - if isinstance(e,Magnet): - current = RWHardwareScalar(self.get_at_elems(e),e.polynom,e.model) if e.model.has_physics() else None - strength = RWStrengthScalar(self.get_at_elems(e),e.polynom,e.model) if e.model.has_physics() else None - # Create a unique ref for this simulator - m = e.attach(self,strength,current) - self.add_magnet(m) - - elif isinstance(e,CombinedFunctionMagnet): - currents = RWHardwareArray(self.get_at_elems(e),e.polynoms,e.model) if e.model.has_physics() else None - strengths = RWStrengthArray(self.get_at_elems(e),e.polynoms,e.model) if e.model.has_physics() else None - # Create unique refs of each function for this simulator - ms = e.attach(self,strengths,currents) - self.add_cfm_magnet(ms[0]) - for m in ms[1:]: - self.add_magnet(m) - - elif isinstance(e,BPM): - # This assumes unique BPM names in the pyAT lattice - tilt = RWBpmTiltScalar(self.get_at_elems(e)[0]) - offsets = RWBpmOffsetArray(self.get_at_elems(e)[0]) - positions = RBpmArray(self.get_at_elems(e)[0],self.ring) - e = e.attach(self,positions, offsets, tilt) - self.add_bpm(e) - - elif isinstance(e,RFPlant): - if e._cfg.transmitters: - cavs: list[at.Element] = [] - harmonics: list[float] = [] - attachedTrans: list[RFTransmitter] = [] - for t in e._cfg.transmitters: - cavsPerTrans: list[at.Element] = [] - for c in t._cfg.cavities: - # Expect unique name for cavities - cav = self.get_at_elems(Element(c)) - if len(cav)>1: - raise PyAMLException(f"RF transmitter {t.get_name()}, multiple cavity definition:{cav[0]}") - if len(cav)==0: - raise PyAMLException(f"RF transmitter {t.get_name()}, No cavity found") - cavsPerTrans.append(cav[0]) - harmonics.append(t._cfg.harmonic) - voltage = RWRFVoltageScalar(cavsPerTrans) - phase = RWRFPhaseScalar(cavsPerTrans) - nt = t.attach(self,voltage,phase) - self.add_rf_transnmitter(nt) - cavs.extend(cavsPerTrans) - attachedTrans.append(nt) - - frequency = RWRFFrequencyScalar(cavs,harmonics) - voltage = RWTotalVoltage(attachedTrans) - ne = e.attach(self,frequency,voltage) - self.add_rf_plant(ne) - else: - # No transmitter defined switch to AT methods - frequency = RWRFATFrequencyScalar(self.ring) - voltage = RWRFATotalVoltageScalar(self.ring) - ne = e.attach(self,frequency,voltage) - self.add_rf_plant(ne) - - elif isinstance(e, BetatronTuneMonitor): - betatron_tune = RBetatronTuneArray(self.ring) - e = e.attach(self,betatron_tune) - self.add_betatron_tune_monitor(e) - - - def get_at_elems(self,element:Element) -> list[at.Element]: - identifier = self._linker.get_element_identifier(element) - element_list = self._linker.get_at_elements(identifier) - if not element_list: - raise PyAMLException(f"{identifier} not found in lattice:{self._cfg.lattice}") - return element_list + e = self.get_at_elems(b)[0] + agg.add_elem(e) + aggh.add_elem(e) + aggv.add_elem(e) + return [agg, aggh, aggv] + + def fill_device(self, elements: list[Element]): + for e in elements: + # Need conversion to physics unit to work with simulator + if isinstance(e, Magnet): + current = ( + RWHardwareScalar(self.get_at_elems(e), e.polynom, e.model) + if e.model.has_physics() + else None + ) + strength = ( + RWStrengthScalar(self.get_at_elems(e), e.polynom, e.model) + if e.model.has_physics() + else None + ) + # Create a unique ref for this simulator + m = e.attach(self, strength, current) + self.add_magnet(m) + + elif isinstance(e, CombinedFunctionMagnet): + currents = ( + RWHardwareArray(self.get_at_elems(e), e.polynoms, e.model) + if e.model.has_physics() + else None + ) + strengths = ( + RWStrengthArray(self.get_at_elems(e), e.polynoms, e.model) + if e.model.has_physics() + else None + ) + # Create unique refs of each function for this simulator + ms = e.attach(self, strengths, currents) + self.add_cfm_magnet(ms[0]) + for m in ms[1:]: + self.add_magnet(m) + + elif isinstance(e, BPM): + # This assumes unique BPM names in the pyAT lattice + tilt = RWBpmTiltScalar(self.get_at_elems(e)[0]) + offsets = RWBpmOffsetArray(self.get_at_elems(e)[0]) + positions = RBpmArray(self.get_at_elems(e)[0], self.ring) + e = e.attach(self, positions, offsets, tilt) + self.add_bpm(e) + + elif isinstance(e, RFPlant): + if e._cfg.transmitters: + cavs: list[at.Element] = [] + harmonics: list[float] = [] + attachedTrans: list[RFTransmitter] = [] + for t in e._cfg.transmitters: + cavsPerTrans: list[at.Element] = [] + for c in t._cfg.cavities: + # Expect unique name for cavities + cav = self.get_at_elems(Element(c)) + if len(cav) > 1: + raise PyAMLException( + f"RF transmitter {t.get_name()}," + "multiple cavity definition:{cav[0]}" + ) + if len(cav) == 0: + raise PyAMLException( + f"RF transmitter {t.get_name()}, No cavity found" + ) + cavsPerTrans.append(cav[0]) + harmonics.append(t._cfg.harmonic) + voltage = RWRFVoltageScalar(cavsPerTrans) + phase = RWRFPhaseScalar(cavsPerTrans) + nt = t.attach(self, voltage, phase) + self.add_rf_transnmitter(nt) + cavs.extend(cavsPerTrans) + attachedTrans.append(nt) + + frequency = RWRFFrequencyScalar(cavs, harmonics) + voltage = RWTotalVoltage(attachedTrans) + ne = e.attach(self, frequency, voltage) + self.add_rf_plant(ne) + else: + # No transmitter defined switch to AT methods + frequency = RWRFATFrequencyScalar(self.ring) + voltage = RWRFATotalVoltageScalar(self.ring) + ne = e.attach(self, frequency, voltage) + self.add_rf_plant(ne) + + elif isinstance(e, BetatronTuneMonitor): + betatron_tune = RBetatronTuneArray(self.ring) + e = e.attach(self, betatron_tune) + self.add_betatron_tune_monitor(e) + + def get_at_elems(self, element: Element) -> list[at.Element]: + identifier = self._linker.get_element_identifier(element) + element_list = self._linker.get_at_elements(identifier) + if not element_list: + raise PyAMLException( + f"{identifier} not found in lattice:{self._cfg.lattice}" + ) + return element_list def __repr__(self): - return repr(self._cfg).replace("ConfigModel",self.__class__.__name__) \ No newline at end of file + return repr(self._cfg).replace("ConfigModel", self.__class__.__name__) diff --git a/pyaml/magnet/__init__.py b/pyaml/magnet/__init__.py index c26b3bae..36b66cbb 100644 --- a/pyaml/magnet/__init__.py +++ b/pyaml/magnet/__init__.py @@ -1,4 +1,3 @@ """ PyAML Magnet module """ - diff --git a/pyaml/magnet/cfm_magnet.py b/pyaml/magnet/cfm_magnet.py index 362f776a..9249a6c8 100644 --- a/pyaml/magnet/cfm_magnet.py +++ b/pyaml/magnet/cfm_magnet.py @@ -1,81 +1,88 @@ +from scipy.constants import speed_of_light -from .model import MagnetModel -from ..common.element import Element,ElementConfigModel,__pyaml_repr__ from ..common import abstract from ..common.abstract import RWMapper +from ..common.element import Element, ElementConfigModel, __pyaml_repr__ from ..common.exception import PyAMLException +from ..configuration import Factory from .hcorrector import HCorrector -from .vcorrector import VCorrector +from .magnet import Magnet, MagnetConfigModel +from .model import MagnetModel +from .octupole import Octupole from .quadrupole import Quadrupole -from .skewquad import SkewQuad from .sextupole import Sextupole -from .skewsext import SkewSext -from .octupole import Octupole from .skewoctu import SkewOctu -from .magnet import Magnet,MagnetConfigModel -from ..configuration import Factory - -from scipy.constants import speed_of_light +from .skewquad import SkewQuad +from .skewsext import SkewSext +from .vcorrector import VCorrector -_fmap:dict = { - "B0":HCorrector, - "A0":VCorrector, - "B1":Quadrupole, - "A1":SkewQuad, - "B2":Sextupole, - "A2":SkewSext, - "B3":Octupole, - "A3":SkewOctu +_fmap: dict = { + "B0": HCorrector, + "A0": VCorrector, + "B1": Quadrupole, + "A1": SkewQuad, + "B2": Sextupole, + "A2": SkewSext, + "B3": Octupole, + "A3": SkewOctu, } # Define the main class name for this module PYAMLCLASS = "CombinedFunctionMagnet" -class ConfigModel(ElementConfigModel): +class ConfigModel(ElementConfigModel): mapping: list[list[str]] - """Name mapping for multipoles (i.e. [[B0,C01A-H],[A0,C01A-H],[B2,C01A-S]])""" + """Name mapping for multipoles + (i.e. [[B0,C01A-H],[A0,C01A-H],[B2,C01A-S]])""" model: MagnetModel | None = None """Object in charge of converting magnet strenghts to currents""" + class CombinedFunctionMagnet(Element): """CombinedFunctionMagnet class""" - def __init__(self, cfg: ConfigModel, peer = None): + def __init__(self, cfg: ConfigModel, peer=None): super().__init__(cfg.name) self._cfg = cfg self.model = cfg.model - self.__virtuals:list[Magnet] = [] - self.__strengths:abstract.ReadWriteFloatArray = None - self.__hardwares:abstract.ReadWriteFloatArray = None + self.__virtuals: list[Magnet] = [] + self.__strengths: abstract.ReadWriteFloatArray = None + self.__hardwares: abstract.ReadWriteFloatArray = None if peer is None: - # Configuration part - if self.model is not None and not hasattr(self.model._cfg,"multipoles"): - raise PyAMLException(f"{cfg.name} model: mutipoles field required for combined function magnet") + if self.model is not None and not hasattr(self.model._cfg, "multipoles"): + raise PyAMLException( + f"{cfg.name} model: mutipoles" + "field required for combined function magnet" + ) idx = 0 self.polynoms = [] - for idx,m in enumerate(cfg.mapping): + for _idx, m in enumerate(cfg.mapping): # Check mapping validity - if len(m)!=2: - raise PyAMLException("Invalid CombinedFunctionMagnet mapping for {m}") - if not m[0] in _fmap: - raise PyAMLException(m[0] + " not implemented for combined function magnet") + if len(m) != 2: + raise PyAMLException( + "Invalid CombinedFunctionMagnet mapping for {m}" + ) + if m[0] not in _fmap: + raise PyAMLException( + m[0] + " not implemented for combined function magnet" + ) if m[0] not in self.model._cfg.multipoles: raise PyAMLException(m[0] + " not found in underlying magnet model") self.polynoms.append(_fmap[m[0]].polynom) # Create the virtual magnet for the correspoding multipole - vm = self.__create_virutal_manget(m[1],m[0]) + vm = self.__create_virutal_manget(m[1], m[0]) self.__virtuals.append(vm) - # Register the virtual element in the factory to have a coherent factory and improve error reporting + # Register the virtual element in the factory to have + # a coherent factory and improve error reporting Factory.register_element(vm) else: - - # Attach - self._peer = peer + # Attach + self._peer = peer def get_model_name(self) -> str: """ @@ -83,53 +90,64 @@ def get_model_name(self) -> str: """ return self._cfg.name - def __create_virutal_manget(self,name:str,idx:int) -> Magnet: - args = {"name":name,"model":self.model} - mVirtual:Magnet = _fmap[idx](MagnetConfigModel(**args)) - mVirtual.set_model_name(self.get_name()) - return mVirtual - + def __create_virutal_manget(self, name: str, idx: int) -> Magnet: + args = {"name": name, "model": self.model} + mVirtual: Magnet = _fmap[idx](MagnetConfigModel(**args)) + mVirtual.set_model_name(self.get_name()) + return mVirtual + def nb_multipole(self) -> int: return len(self._cfg.mapping) - def attach(self, peer, strengths: abstract.ReadWriteFloatArray, hardwares: abstract.ReadWriteFloatArray) -> list[Magnet]: - l = [] - # Attached the CombinedFunctionMagnet itself - nCFM = CombinedFunctionMagnet(self._cfg,peer) - nCFM.__strengths = strengths - nCFM.__hardwares = hardwares - l.append(nCFM) - # Construct a single function magnet for each multipole of this combined function magnet - for idx,m in enumerate(self._cfg.mapping): - strength = RWMapper(strengths,idx) - hardware = RWMapper(hardwares,idx) if self.model.has_hardware() else None - l.append(self.__virtuals[idx].attach(peer,strength,hardware)) - return l - + def attach( + self, + peer, + strengths: abstract.ReadWriteFloatArray, + hardwares: abstract.ReadWriteFloatArray, + ) -> list[Magnet]: + l = [] + # Attached the CombinedFunctionMagnet itself + nCFM = CombinedFunctionMagnet(self._cfg, peer) + nCFM.__strengths = strengths + nCFM.__hardwares = hardwares + l.append(nCFM) + # Construct a single function magnet for each multipole + # of this combined function magnet + for idx, _m in enumerate(self._cfg.mapping): + strength = RWMapper(strengths, idx) + hardware = RWMapper(hardwares, idx) if self.model.has_hardware() else None + l.append(self.__virtuals[idx].attach(peer, strength, hardware)) + return l + @property def strengths(self) -> abstract.ReadWriteFloatScalar: """ - Gives access to the strengths of this combined function magnet in physics unit + Gives access to the strengths of this combined + function magnet in physics unit """ self.check_peer() if self.__strengths is None: - raise PyAMLException(f"{str(self)} has no model that supports physics units") + raise PyAMLException( + f"{str(self)} has no model that supports physics units" + ) return self.__strengths @property def hardwares(self) -> abstract.ReadWriteFloatScalar: """ - Gives access to the strengths of this combined function magnet in hardware unit when possible + Gives access to the strengths of this combined + function magnet in hardware unit when possible """ self.check_peer() if self.__hardwares is None: - raise PyAMLException(f"{str(self)} has no model that supports hardware units") + raise PyAMLException( + f"{str(self)} has no model that supports hardware units" + ) return self.__hardwares - - def set_energy(self,E:float): - if(self.model is not None): - self.model.set_magnet_rigidity(E/speed_of_light) + + def set_energy(self, E: float): + if self.model is not None: + self.model.set_magnet_rigidity(E / speed_of_light) def __repr__(self): return __pyaml_repr__(self) - diff --git a/pyaml/magnet/hcorrector.py b/pyaml/magnet/hcorrector.py index ae792cdb..bc2910a8 100644 --- a/pyaml/magnet/hcorrector.py +++ b/pyaml/magnet/hcorrector.py @@ -1,14 +1,17 @@ -from .magnet import Magnet,MagnetConfigModel from ..lattice.polynom_info import PolynomInfo +from .magnet import Magnet, MagnetConfigModel # Define the main class name for this module PYAMLCLASS = "HCorrector" -class ConfigModel(MagnetConfigModel):... + +class ConfigModel(MagnetConfigModel): ... + class HCorrector(Magnet): """Horizontal Corrector class""" - polynom = PolynomInfo('PolynomB',0) + + polynom = PolynomInfo("PolynomB", 0) def __init__(self, cfg: ConfigModel): super().__init__( @@ -16,4 +19,3 @@ def __init__(self, cfg: ConfigModel): cfg.model if hasattr(cfg, "model") else None, ) self._cfg = cfg - diff --git a/pyaml/magnet/identity_cfm_model.py b/pyaml/magnet/identity_cfm_model.py index d200a14e..f415593d 100644 --- a/pyaml/magnet/identity_cfm_model.py +++ b/pyaml/magnet/identity_cfm_model.py @@ -1,17 +1,17 @@ import numpy as np -from pydantic import BaseModel,ConfigDict +from pydantic import BaseModel, ConfigDict -from .model import MagnetModel from .. import PyAMLException -from ..control.deviceaccess import DeviceAccess from ..common.element import __pyaml_repr__ +from ..control.deviceaccess import DeviceAccess +from .model import MagnetModel # Define the main class name for this module PYAMLCLASS = "IdentityCFMagnetModel" -class ConfigModel(BaseModel): - model_config = ConfigDict(arbitrary_types_allowed=True,extra="forbid") +class ConfigModel(BaseModel): + model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") multipoles: list[str] """List of supported functions: A0,B0,A1,B1,etc (i.e. [B0,A1,B2])""" @@ -22,6 +22,7 @@ class ConfigModel(BaseModel): units: list[str] """List of strength unit (i.e. ['rad','m-1','m-2'])""" + class IdentityCFMagnetModel(MagnetModel): """ Class that map values to underlying devices without conversion @@ -34,9 +35,15 @@ def __init__(self, cfg: ConfigModel): self.__nbFunction: int = len(cfg.multipoles) if cfg.physics is None and cfg.powerconverters is None: - raise PyAMLException("Invalid IdentityCFMagnetModel configuration, physics or powerconverters device required") + raise PyAMLException( + "Invalid IdentityCFMagnetModel configuration," + "physics or powerconverters device required" + ) if cfg.physics is not None and cfg.powerconverters is not None: - raise PyAMLException("Invalid IdentityCFMagnetModel configuration, physics or powerconverters device required but not both") + raise PyAMLException( + "Invalid IdentityCFMagnetModel configuration," + "physics or powerconverters device required but not both" + ) if cfg.physics: self.__devices = cfg.physics else: @@ -44,15 +51,15 @@ def __init__(self, cfg: ConfigModel): self.__nbDev: int = len(self.__devices) - self.__check_len(cfg.units,"units",self.__nbFunction) + self.__check_len(cfg.units, "units", self.__nbFunction) - def __check_len(self,obj,name,expected_len): - lgth = len(obj) + def __check_len(self, obj, name, expected_len): + lgth = len(obj) if lgth != expected_len: raise PyAMLException( f"{name} does not have the expected " f"number of items ({expected_len} items expected but got {lgth})" - ) + ) def compute_hardware_values(self, strengths: np.array) -> np.array: return strengths diff --git a/pyaml/magnet/identity_model.py b/pyaml/magnet/identity_model.py index dfcade7e..c5206a86 100644 --- a/pyaml/magnet/identity_model.py +++ b/pyaml/magnet/identity_model.py @@ -1,25 +1,26 @@ -from .model import MagnetModel +import numpy as np +from pydantic import BaseModel, ConfigDict + from .. import PyAMLException -from ..control.deviceaccess import DeviceAccess from ..common.element import __pyaml_repr__ - -import numpy as np -from pydantic import BaseModel,ConfigDict +from ..control.deviceaccess import DeviceAccess +from .model import MagnetModel # Define the main class name for this module PYAMLCLASS = "IdentityMagnetModel" -class ConfigModel(BaseModel): - model_config = ConfigDict(arbitrary_types_allowed=True,extra="forbid") +class ConfigModel(BaseModel): + model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") - powerconverter: DeviceAccess|None = None + powerconverter: DeviceAccess | None = None """Power converter device to apply current""" - physics: DeviceAccess|None = None + physics: DeviceAccess | None = None """Magnet device to apply strength""" unit: str """Unit of the strength (i.e. 1/m or m-1)""" + class IdentityMagnetModel(MagnetModel): """ Class that map value to underlying device without conversion @@ -29,9 +30,15 @@ def __init__(self, cfg: ConfigModel): self._cfg = cfg self.__unit = cfg.unit if cfg.physics is None and cfg.powerconverter is None: - raise PyAMLException("Invalid IdentityMagnetModel configuration, physics or powerconverter device required") + raise PyAMLException( + "Invalid IdentityMagnetModel configuration," + "physics or powerconverter device required" + ) if cfg.physics is not None and cfg.powerconverter is not None: - raise PyAMLException("Invalid IdentityMagnetModel configuration, physics or powerconverter device required but not both") + raise PyAMLException( + "Invalid IdentityMagnetModel configuration," + "physics or powerconverter device required but not both" + ) if cfg.physics: self.__device = cfg.physics else: @@ -71,4 +78,4 @@ def has_hardware(self) -> bool: return self._cfg.powerconverter is not None def __repr__(self): - return __pyaml_repr__(self) + return __pyaml_repr__(self) diff --git a/pyaml/magnet/linear_cfm_model.py b/pyaml/magnet/linear_cfm_model.py index 1ece0abb..1a04d722 100644 --- a/pyaml/magnet/linear_cfm_model.py +++ b/pyaml/magnet/linear_cfm_model.py @@ -1,20 +1,19 @@ +import numpy as np +from pydantic import BaseModel, ConfigDict + +from ..common.element import __pyaml_repr__ +from ..common.exception import PyAMLException from ..configuration.curve import Curve from ..configuration.matrix import Matrix from ..control.deviceaccess import DeviceAccess from .model import MagnetModel -from ..common.exception import PyAMLException -from ..common.element import __pyaml_repr__ - -from pydantic import BaseModel,ConfigDict -import numpy as np # Define the main class name for this module PYAMLCLASS = "LinearCFMagnetModel" class ConfigModel(BaseModel): - - model_config = ConfigDict(arbitrary_types_allowed=True,extra="forbid") + model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") multipoles: list[str] """List of supported functions: A0,B0,A1,B1,etc (i.e. [B0,A1,B2])""" @@ -36,7 +35,7 @@ class ConfigModel(BaseModel): """List of power converter devices to apply currrents (can be different from number of function)""" matrix: Matrix = None - """n x m matrix (n rows for n function , m columns for m currents) + """n x m matrix (n rows for n function , m columns for m currents) to handle multipoles separation. Default: Identity""" units: list[str] """List of strength unit (i.e. ['rad','m-1','m-2'])""" @@ -76,12 +75,16 @@ def __init__(self, cfg: ConfigModel): else: self.__po = cfg.pseudo_factors - self.__check_len(self.__calibration_factors,"calibration_factors",self.__nbFunction) - self.__check_len(self.__calibration_offsets,"calibration_offsets",self.__nbFunction) - self.__check_len(self.__pf,"pseudo_factors",self.__nbFunction) - self.__check_len(self.__po,"pseudo_offsets",self.__nbFunction) - self.__check_len(cfg.units,"units",self.__nbFunction) - self.__check_len(cfg.curves,"curves",self.__nbFunction) + self.__check_len( + self.__calibration_factors, "calibration_factors", self.__nbFunction + ) + self.__check_len( + self.__calibration_offsets, "calibration_offsets", self.__nbFunction + ) + self.__check_len(self.__pf, "pseudo_factors", self.__nbFunction) + self.__check_len(self.__po, "pseudo_offsets", self.__nbFunction) + self.__check_len(cfg.units, "units", self.__nbFunction) + self.__check_len(cfg.curves, "curves", self.__nbFunction) if cfg.matrix is None: self.__matrix = np.identity(self.__nbFunction) @@ -95,7 +98,7 @@ def __init__(self, cfg: ConfigModel): "matrix wrong dimension " f"({self.__nbFunction}x{self.__nbPS} expected but got {_s[0]}x{_s[1]})" ) - + self.__curves = [] self.__rcurves = [] @@ -109,21 +112,22 @@ def __init__(self, cfg: ConfigModel): # Compute pseudo inverse self.__inv = np.linalg.pinv(self.__matrix) - - def __check_len(self,obj,name,expected_len): - lgth = len(obj) + def __check_len(self, obj, name, expected_len): + lgth = len(obj) if lgth != expected_len: raise PyAMLException( f"{name} does not have the expected " f"number of items ({expected_len} items expected but got {lgth})" - ) + ) def compute_hardware_values(self, strengths: np.array) -> np.array: _pI = np.zeros(self.__nbFunction) for idx, c in enumerate(self.__rcurves): - _pI[idx] = self.__pf[idx] * np.interp( - strengths[idx] * self._brho, c[:, 0], c[:, 1] - ) + self.__po[idx] + _pI[idx] = ( + self.__pf[idx] + * np.interp(strengths[idx] * self._brho, c[:, 0], c[:, 1]) + + self.__po[idx] + ) _currents = np.matmul(self.__inv, _pI) return _currents @@ -133,7 +137,7 @@ def compute_strengths(self, currents: np.array) -> np.array: for idx, c in enumerate(self.__curves): _strength[idx] = ( np.interp( - (_pI[idx] - self.__po[idx]) / self.__pf[idx], c[:, 0], c[:, 1] + (_pI[idx] - self.__po[idx]) / self.__pf[idx], c[:, 0], c[:, 1] ) / self._brho ) @@ -162,8 +166,9 @@ def set_magnet_rigidity(self, brho: np.double): self._brho = brho def has_hardware(self) -> bool: - return (self.__nbPS == self.__nbFunction) and np.allclose(self.__matrix, np.eye(self.__nbFunction)) + return (self.__nbPS == self.__nbFunction) and np.allclose( + self.__matrix, np.eye(self.__nbFunction) + ) def __repr__(self): return __pyaml_repr__(self) - diff --git a/pyaml/magnet/linear_model.py b/pyaml/magnet/linear_model.py index 9afc990d..710ac3c2 100644 --- a/pyaml/magnet/linear_model.py +++ b/pyaml/magnet/linear_model.py @@ -1,17 +1,17 @@ -from .model import MagnetModel +import numpy as np +from pydantic import BaseModel, ConfigDict + +from ..common.element import __pyaml_repr__ from ..configuration.curve import Curve from ..control.deviceaccess import DeviceAccess -from ..common.element import __pyaml_repr__ - -import numpy as np -from pydantic import BaseModel,ConfigDict +from .model import MagnetModel # Define the main class name for this module PYAMLCLASS = "LinearMagnetModel" -class ConfigModel(BaseModel): - model_config = ConfigDict(arbitrary_types_allowed=True,extra="forbid") +class ConfigModel(BaseModel): + model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") curve: Curve """Curve object used for interpolation""" @@ -26,16 +26,19 @@ class ConfigModel(BaseModel): unit: str """Unit of the strength (i.e. 1/m or m-1)""" + class LinearMagnetModel(MagnetModel): """ - Class that handle manget current/strength conversion using linear interpolation for a single function magnet + Class that handle manget current/strength conversion using + linear interpolation for a single function magnet """ def __init__(self, cfg: ConfigModel): self._cfg = cfg self.__curve = cfg.curve.get_curve() self.__curve[:, 1] = ( - self.__curve[:, 1] * cfg.calibration_factor * cfg.crosstalk + cfg.calibration_offset + self.__curve[:, 1] * cfg.calibration_factor * cfg.crosstalk + + cfg.calibration_offset ) self.__rcurve = Curve.inverse(self.__curve) self.__strength_unit = cfg.unit @@ -78,5 +81,3 @@ def set_magnet_rigidity(self, brho: np.double): def __repr__(self): return __pyaml_repr__(self) - - \ No newline at end of file diff --git a/pyaml/magnet/magnet.py b/pyaml/magnet/magnet.py index 63fef983..c73a69f6 100644 --- a/pyaml/magnet/magnet.py +++ b/pyaml/magnet/magnet.py @@ -1,104 +1,119 @@ -from ..common.element import Element,ElementConfigModel +from scipy.constants import speed_of_light + from .. import PyAMLException from ..common import abstract +from ..common.element import Element, ElementConfigModel from .model import MagnetModel -from scipy.constants import speed_of_light try: from typing import Self # Python 3.11+ except ImportError: from typing_extensions import Self # Python 3.10 and earlier import numpy as np -class MagnetConfigModel(ElementConfigModel): +class MagnetConfigModel(ElementConfigModel): model: MagnetModel | None = None """Object in charge of converting magnet strenghts to power supply values""" -class Magnet(Element): - """ - Class providing access to one magnet of a physical or simulated lattice - """ - def __init__(self, name:str, model:MagnetModel = None): - """ - Construct a magnet - - Parameters - ---------- - name : str - Element name - model : MagnetModel - Magnet model in charge of computing coil(s) current - """ - super().__init__(name) - self.__model = model - self.__strength:abstract.ReadWriteFloatScalar = None - self.__hardware:abstract.ReadWriteFloatScalar = None - self.__modelName = self.get_name() - - @property - def strength(self) -> abstract.ReadWriteFloatScalar: +class Magnet(Element): """ - Gives access to the strength of this magnet in physics unit + Class providing access to one magnet of a physical or simulated lattice """ - self.check_peer() - if self.__strength is None: - raise PyAMLException(f"{str(self)} has no model that supports physics units") - return self.__strength - @property - def hardware(self) -> abstract.ReadWriteFloatScalar: - """ - Gives access to the strength of this magnet in hardware unit when possible - """ - self.check_peer() - if self.__hardware is None: - raise PyAMLException(f"{str(self)} has no model that supports hardware units") - return self.__hardware - - @property - def model(self) -> MagnetModel: - """ - Returns a handle to the underlying magnet model - """ - return self.__model - - def attach(self, peer, strength: abstract.ReadWriteFloatScalar, hardware: abstract.ReadWriteFloatScalar) -> Self: - """ - Create a new reference to attach this magnet to a simulator or a control systemand. - """ - obj = self.__class__(self._cfg) - obj.__modelName = self.__modelName - obj.__strength = strength - obj.__hardware = hardware - obj._peer = peer - return obj - - def set_energy(self, energy:float): - """ - Set the energy in eV to compute and set the magnet rigidity on the underlying magnet model. - """ - if self.__model is not None: - self.__model.set_magnet_rigidity(np.double(energy / speed_of_light)) - - def set_model_name(self, name:str): - """ - Sets the name of this magnet in the model (Used for combined function manget) - """ - self.__modelName = name - - def get_model_name(self) -> str: - """ - Returns the model name of this magnet - """ - return self.__modelName - - def __repr__(self): - return "%s(peer='%s', name='%s', model_name='%s', magnet_model=%s)" % ( - self.__class__.__name__, - self.get_peer(), - self.get_name(), - self.__modelName, - repr(self.__model) - ) + def __init__(self, name: str, model: MagnetModel = None): + """ + Construct a magnet + + Parameters + ---------- + name : str + Element name + model : MagnetModel + Magnet model in charge of computing coil(s) current + """ + super().__init__(name) + self.__model = model + self.__strength: abstract.ReadWriteFloatScalar = None + self.__hardware: abstract.ReadWriteFloatScalar = None + self.__modelName = self.get_name() + + @property + def strength(self) -> abstract.ReadWriteFloatScalar: + """ + Gives access to the strength of this magnet in physics unit + """ + self.check_peer() + if self.__strength is None: + raise PyAMLException( + f"{str(self)} has no model that supports physics units" + ) + return self.__strength + + @property + def hardware(self) -> abstract.ReadWriteFloatScalar: + """ + Gives access to the strength of this magnet in + hardware unit when possible + """ + self.check_peer() + if self.__hardware is None: + raise PyAMLException( + f"{str(self)} has no model that supports hardware units" + ) + return self.__hardware + + @property + def model(self) -> MagnetModel: + """ + Returns a handle to the underlying magnet model + """ + return self.__model + + def attach( + self, + peer, + strength: abstract.ReadWriteFloatScalar, + hardware: abstract.ReadWriteFloatScalar, + ) -> Self: + """ + Create a new reference to attach this magnet to a simulator + or a control systemand. + """ + obj = self.__class__(self._cfg) + obj.__modelName = self.__modelName + obj.__strength = strength + obj.__hardware = hardware + obj._peer = peer + return obj + + def set_energy(self, energy: float): + """ + Set the energy in eV to compute and set the magnet rigidity + on the underlying magnet model. + """ + if self.__model is not None: + self.__model.set_magnet_rigidity(np.double(energy / speed_of_light)) + + def set_model_name(self, name: str): + """ + Sets the name of this magnet in the model + (Used for combined function manget) + """ + self.__modelName = name + + def get_model_name(self) -> str: + """ + Returns the model name of this magnet + """ + return self.__modelName + + def __repr__(self): + return "%s(peer='%s', name='%s', model_name='%s', magnet_model=%s)" % ( + self.__class__.__name__, + self.get_peer(), + self.get_name(), + self.__modelName, + repr(self.__model), + ) diff --git a/pyaml/magnet/model.py b/pyaml/magnet/model.py index 4df7192b..e7cb2478 100644 --- a/pyaml/magnet/model.py +++ b/pyaml/magnet/model.py @@ -1,22 +1,29 @@ from abc import ABCMeta, abstractmethod + import numpy as np import numpy.typing as npt + from ..control.deviceaccess import DeviceAccess + class MagnetModel(metaclass=ABCMeta): """ - Abstract class providing strength to coil current conversion and access to underlying power supplies + Abstract class providing strength to coil current conversion + and access to underlying power supplies """ @abstractmethod - def compute_hardware_values(self, strengths: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]: + def compute_hardware_values( + self, strengths: npt.NDArray[np.float64] + ) -> npt.NDArray[np.float64]: """ Compute hardware value(s) from magnet strength(s) Parameters ---------- strengths : npt.NDArray[np.float64] - Array of strengths. For a single multipole, strengths is an array of 1 item. + Array of strengths. For a single multipole, + strengths is an array of 1 item. Returns ------- @@ -26,19 +33,22 @@ def compute_hardware_values(self, strengths: npt.NDArray[np.float64]) -> npt.NDA pass @abstractmethod - def compute_strengths(self, hardware_values: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]: + def compute_strengths( + self, hardware_values: npt.NDArray[np.float64] + ) -> npt.NDArray[np.float64]: """ Compute magnet strength(s) from hardware value(s) - + Parameters ---------- hardware_values : npt.NDArray[np.float64] - Array of hardware values (i.e. currents or voltages) + Array of hardware values (i.e. currents or voltages) Returns ------- npt.NDArray[np.float64] - Array of strengths. For a single multipole, returns an array of 1 item + Array of strengths. For a single multipole, + returns an array of 1 item """ pass @@ -50,7 +60,8 @@ def get_strength_units(self) -> list[str]: Returns ------- list[str] - Array of strength units. For a single multipole, returns a list of 1 item + Array of strength units. For a single multipole, + returns a list of 1 item """ pass @@ -58,11 +69,12 @@ def get_strength_units(self) -> list[str]: def get_hardware_units(self) -> list[str]: """ Get hardware units - + Returns ------- list[str] - Array of hardware units. For a single multipole, returns a list of 1 item + Array of hardware units. For a single multipole, + returns a list of 1 item """ pass @@ -74,7 +86,7 @@ def read_hardware_values(self) -> npt.NDArray[np.float64]: Returns ------- npt.NDArray[np.float64] - Array of hardware values (i.e. currents or voltages) + Array of hardware values (i.e. currents or voltages) """ pass @@ -82,11 +94,11 @@ def read_hardware_values(self) -> npt.NDArray[np.float64]: def readback_hardware_values(self) -> npt.NDArray[np.float64]: """ Get power supply harware value(s) measurements from control system - + Returns ------- npt.NDArray[np.float64] - Array of hardware values (i.e. currents or voltages) + Array of hardware values (i.e. currents or voltages) """ pass @@ -94,11 +106,11 @@ def readback_hardware_values(self) -> npt.NDArray[np.float64]: def send_hardware_values(self, hardware_values: npt.NDArray[np.float64]): """ Send power supply value(s) to control system - + Parameters ---------- hardware_values : npt.NDArray[np.float64] - Array of hardware values (i.e. currents or voltages) + Array of hardware values (i.e. currents or voltages) """ pass @@ -106,7 +118,7 @@ def send_hardware_values(self, hardware_values: npt.NDArray[np.float64]): def get_devices(self) -> list[DeviceAccess]: """ Get device handles - + Returns ------- list[DeviceAccess] @@ -126,7 +138,6 @@ def set_magnet_rigidity(self, brho: np.double): """ pass - def has_hardware(self) -> bool: """ Tells if the model allows to work in hardware unit. @@ -137,7 +148,7 @@ def has_hardware(self) -> bool: True if the model supports hardware unit """ return True - + def has_physics(self) -> bool: """ Tells if the model allows to work in physics unit. @@ -147,4 +158,4 @@ def has_physics(self) -> bool: bool True if the model supports physics unit """ - return True \ No newline at end of file + return True diff --git a/pyaml/magnet/octupole.py b/pyaml/magnet/octupole.py index 49ec78f4..17b296c9 100644 --- a/pyaml/magnet/octupole.py +++ b/pyaml/magnet/octupole.py @@ -1,14 +1,17 @@ -from .magnet import Magnet,MagnetConfigModel from ..lattice.polynom_info import PolynomInfo +from .magnet import Magnet, MagnetConfigModel # Define the main class name for this module PYAMLCLASS = "Octupole" -class ConfigModel(MagnetConfigModel):... -class Octupole(Magnet): +class ConfigModel(MagnetConfigModel): ... + + +class Octupole(Magnet): """Octupole class""" - polynom = PolynomInfo('PolynomB',3) + + polynom = PolynomInfo("PolynomB", 3) def __init__(self, cfg: ConfigModel): super().__init__( diff --git a/pyaml/magnet/quadrupole.py b/pyaml/magnet/quadrupole.py index 063ba226..115d820c 100644 --- a/pyaml/magnet/quadrupole.py +++ b/pyaml/magnet/quadrupole.py @@ -1,21 +1,21 @@ -from .magnet import Magnet,MagnetConfigModel from ..lattice.polynom_info import PolynomInfo +from .magnet import Magnet, MagnetConfigModel # Define the main class name for this module PYAMLCLASS = "Quadrupole" -class ConfigModel(MagnetConfigModel):... -class Quadrupole(Magnet): +class ConfigModel(MagnetConfigModel): ... + + +class Quadrupole(Magnet): """Quadrupole class""" - polynom = PolynomInfo('PolynomB',1) - + + polynom = PolynomInfo("PolynomB", 1) + def __init__(self, cfg: ConfigModel): super().__init__( cfg.name, cfg.model if hasattr(cfg, "model") else None, ) self._cfg = cfg - - - diff --git a/pyaml/magnet/sextupole.py b/pyaml/magnet/sextupole.py index 826ec0f9..ddfc95d2 100644 --- a/pyaml/magnet/sextupole.py +++ b/pyaml/magnet/sextupole.py @@ -1,14 +1,17 @@ -from .magnet import Magnet,MagnetConfigModel from ..lattice.polynom_info import PolynomInfo +from .magnet import Magnet, MagnetConfigModel # Define the main class name for this module PYAMLCLASS = "Sextupole" -class ConfigModel(MagnetConfigModel):... -class Sextupole(Magnet): +class ConfigModel(MagnetConfigModel): ... + + +class Sextupole(Magnet): """Sextupole class""" - polynom = PolynomInfo('PolynomB',2) + + polynom = PolynomInfo("PolynomB", 2) def __init__(self, cfg: ConfigModel): super().__init__( @@ -16,6 +19,3 @@ def __init__(self, cfg: ConfigModel): cfg.model if hasattr(cfg, "model") else None, ) self._cfg = cfg - - - diff --git a/pyaml/magnet/skewoctu.py b/pyaml/magnet/skewoctu.py index 4d0cb375..56ed8b84 100644 --- a/pyaml/magnet/skewoctu.py +++ b/pyaml/magnet/skewoctu.py @@ -1,14 +1,17 @@ -from .magnet import Magnet,MagnetConfigModel from ..lattice.polynom_info import PolynomInfo +from .magnet import Magnet, MagnetConfigModel # Define the main class name for this module PYAMLCLASS = "SkewOctu" -class ConfigModel(MagnetConfigModel):... -class SkewOctu(Magnet): +class ConfigModel(MagnetConfigModel): ... + + +class SkewOctu(Magnet): """SkewOctu class""" - polynom = PolynomInfo('PolynomA',3) + + polynom = PolynomInfo("PolynomA", 3) def __init__(self, cfg: ConfigModel): super().__init__( @@ -16,6 +19,3 @@ def __init__(self, cfg: ConfigModel): cfg.model if hasattr(cfg, "model") else None, ) self._cfg = cfg - - - diff --git a/pyaml/magnet/skewquad.py b/pyaml/magnet/skewquad.py index 6c51ef94..5855ce89 100644 --- a/pyaml/magnet/skewquad.py +++ b/pyaml/magnet/skewquad.py @@ -1,14 +1,17 @@ -from .magnet import Magnet,MagnetConfigModel from ..lattice.polynom_info import PolynomInfo +from .magnet import Magnet, MagnetConfigModel # Define the main class name for this module PYAMLCLASS = "SkewQuad" -class ConfigModel(MagnetConfigModel):... -class SkewQuad(Magnet): +class ConfigModel(MagnetConfigModel): ... + + +class SkewQuad(Magnet): """SkewQuad class""" - polynom = PolynomInfo('PolynomA',1) + + polynom = PolynomInfo("PolynomA", 1) def __init__(self, cfg: ConfigModel): super().__init__( @@ -16,6 +19,3 @@ def __init__(self, cfg: ConfigModel): cfg.model if hasattr(cfg, "model") else None, ) self._cfg = cfg - - - diff --git a/pyaml/magnet/skewsext.py b/pyaml/magnet/skewsext.py index e3f90c95..da0cde57 100644 --- a/pyaml/magnet/skewsext.py +++ b/pyaml/magnet/skewsext.py @@ -1,14 +1,17 @@ -from .magnet import Magnet,MagnetConfigModel from ..lattice.polynom_info import PolynomInfo +from .magnet import Magnet, MagnetConfigModel # Define the main class name for this module PYAMLCLASS = "SkewSext" -class ConfigModel(MagnetConfigModel):... -class SkewSext(Magnet): +class ConfigModel(MagnetConfigModel): ... + + +class SkewSext(Magnet): """SkewSext class""" - polynom = PolynomInfo('PolynomA',2) + + polynom = PolynomInfo("PolynomA", 2) def __init__(self, cfg: ConfigModel): super().__init__( @@ -16,6 +19,3 @@ def __init__(self, cfg: ConfigModel): cfg.model if hasattr(cfg, "model") else None, ) self._cfg = cfg - - - diff --git a/pyaml/magnet/spline_model.py b/pyaml/magnet/spline_model.py index b926a99d..dcb66588 100644 --- a/pyaml/magnet/spline_model.py +++ b/pyaml/magnet/spline_model.py @@ -1,18 +1,18 @@ import numpy as np -from pydantic import BaseModel,ConfigDict +from pydantic import BaseModel, ConfigDict from scipy.interpolate import make_smoothing_spline -from .model import MagnetModel +from ..common.element import __pyaml_repr__ from ..configuration.curve import Curve from ..control.deviceaccess import DeviceAccess -from ..common.element import __pyaml_repr__ +from .model import MagnetModel # Define the main class name for this module PYAMLCLASS = "SplineMagnetModel" -class ConfigModel(BaseModel): - model_config = ConfigDict(arbitrary_types_allowed=True,extra="forbid") +class ConfigModel(BaseModel): + model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") curve: Curve """Curve object used for interpolation""" @@ -27,25 +27,31 @@ class ConfigModel(BaseModel): unit: str """Unit of the strength (i.e. 1/m or m-1)""" alpha: float = 0.0 - """Regularization parameter (alpha>=0), aplha=0 the interpolation pass through all the points of the curve""" + """Regularization parameter (alpha>=0), aplha=0 the interpolation + pass through all the points of the curve""" + class SplineMagnetModel(MagnetModel): """ - Class that handle manget current/strength conversion using spline interpolation for a single function magnet + Class that handle manget current/strength conversion using + spline interpolation for a single function magnet """ def __init__(self, cfg: ConfigModel): self._cfg = cfg self.__curve = cfg.curve.get_curve() self.__curve[:, 1] = ( - self.__curve[:, 1] * cfg.calibration_factor * cfg.crosstalk + cfg.calibration_offset + self.__curve[:, 1] * cfg.calibration_factor * cfg.crosstalk + + cfg.calibration_offset ) rcurve = Curve.inverse(self.__curve) self.__strength_unit = cfg.unit self.__hardware_unit = cfg.powerconverter.unit() self.__brho = np.nan self.__ps = cfg.powerconverter - self.__spl = make_smoothing_spline(self.__curve[:, 0], self.__curve[:, 1], lam=cfg.alpha) + self.__spl = make_smoothing_spline( + self.__curve[:, 0], self.__curve[:, 1], lam=cfg.alpha + ) self.__rspl = make_smoothing_spline(rcurve[:, 0], rcurve[:, 1], lam=cfg.alpha) def compute_hardware_values(self, strengths: np.array) -> np.array: @@ -79,4 +85,3 @@ def set_magnet_rigidity(self, brho: np.double): def __repr__(self): return __pyaml_repr__(self) - diff --git a/pyaml/magnet/vcorrector.py b/pyaml/magnet/vcorrector.py index ea4b1d60..2e5334f1 100644 --- a/pyaml/magnet/vcorrector.py +++ b/pyaml/magnet/vcorrector.py @@ -1,14 +1,17 @@ -from .magnet import Magnet,MagnetConfigModel from ..lattice.polynom_info import PolynomInfo +from .magnet import Magnet, MagnetConfigModel # Define the main class name for this module PYAMLCLASS = "VCorrector" -class ConfigModel(MagnetConfigModel):... -class VCorrector(Magnet): +class ConfigModel(MagnetConfigModel): ... + + +class VCorrector(Magnet): """Vertical Corrector class""" - polynom = PolynomInfo('PolynomA',0) + + polynom = PolynomInfo("PolynomA", 0) def __init__(self, cfg: ConfigModel): super().__init__( @@ -16,6 +19,3 @@ def __init__(self, cfg: ConfigModel): cfg.model if hasattr(cfg, "model") else None, ) self._cfg = cfg - - - diff --git a/pyaml/rf/__init__.py b/pyaml/rf/__init__.py index bf881e4a..df9eb4c0 100644 --- a/pyaml/rf/__init__.py +++ b/pyaml/rf/__init__.py @@ -1,4 +1,3 @@ """ PyAML RF module """ - diff --git a/pyaml/rf/rf_plant.py b/pyaml/rf/rf_plant.py index ae10371c..5c5bfe53 100644 --- a/pyaml/rf/rf_plant.py +++ b/pyaml/rf/rf_plant.py @@ -1,26 +1,28 @@ import numpy as np -from pydantic import BaseModel,ConfigDict +from pydantic import BaseModel, ConfigDict + try: from typing import Self # Python 3.11+ except ImportError: from typing_extensions import Self # Python 3.10 and earlier -from .rf_transmitter import RFTransmitter from .. import PyAMLException -from ..control.deviceaccess import DeviceAccess -from ..common.element import Element,ElementConfigModel from ..common import abstract +from ..common.element import Element, ElementConfigModel +from ..control.deviceaccess import DeviceAccess +from .rf_transmitter import RFTransmitter # Define the main class name for this module PYAMLCLASS = "RFPlant" -class ConfigModel(ElementConfigModel): - masterclock: DeviceAccess|None = None +class ConfigModel(ElementConfigModel): + masterclock: DeviceAccess | None = None """Device to apply main RF frequency""" - transmitters: list[RFTransmitter]|None = None + transmitters: list[RFTransmitter] | None = None """List of RF trasnmitters""" + class RFPlant(Element): """ Main RF object @@ -44,16 +46,21 @@ def voltage(self) -> abstract.ReadWriteFloatScalar: raise PyAMLException(f"{str(self)} has no trasmitter device defined") return self.__voltage - def attach(self, peer, frequency: abstract.ReadWriteFloatScalar, voltage: abstract.ReadWriteFloatScalar) -> Self: + def attach( + self, + peer, + frequency: abstract.ReadWriteFloatScalar, + voltage: abstract.ReadWriteFloatScalar, + ) -> Self: # Attach frequency attribute and returns a new reference obj = self.__class__(self._cfg) obj.__frequency = frequency obj.__voltage = voltage obj._peer = peer return obj - -class RWTotalVoltage(abstract.ReadWriteFloatScalar): + +class RWTotalVoltage(abstract.ReadWriteFloatScalar): def __init__(self, transmitters: list[RFTransmitter]): """ Construct a RWTotalVoltage setter @@ -69,23 +76,19 @@ def get(self) -> float: sum = 0 # Count only fundamental harmonic for t in self.__trans: - if(t._cfg.harmonic==1.): + if t._cfg.harmonic == 1.0: sum += t.voltage.get() return sum - - def set(self,value:float): + + def set(self, value: float): # Assume that sum of transmitter (fundamental harmonic) distribution is 1 for t in self.__trans: - if(t._cfg.harmonic==1.): + if t._cfg.harmonic == 1.0: v = value * t._cfg.distribution t.voltage.set(v) - def set_and_wait(self, value:float): + def set_and_wait(self, value: float): raise NotImplementedError("Not implemented yet.") - + def unit(self) -> str: return self.__trans[0]._cfg.phase.unit() - - - - diff --git a/pyaml/rf/rf_transmitter.py b/pyaml/rf/rf_transmitter.py index 63ded2e9..15fcbbbc 100644 --- a/pyaml/rf/rf_transmitter.py +++ b/pyaml/rf/rf_transmitter.py @@ -1,25 +1,26 @@ import numpy as np -from pydantic import BaseModel,ConfigDict +from pydantic import BaseModel, ConfigDict + try: from typing import Self # Python 3.11+ except ImportError: from typing_extensions import Self # Python 3.10 and earlier from .. import PyAMLException -from ..control.deviceaccess import DeviceAccess -from ..common.element import Element,ElementConfigModel from ..common import abstract +from ..common.element import Element, ElementConfigModel +from ..control.deviceaccess import DeviceAccess # Define the main class name for this module PYAMLCLASS = "RFTransmitter" -class ConfigModel(ElementConfigModel): - model_config = ConfigDict(arbitrary_types_allowed=True,extra="forbid") +class ConfigModel(ElementConfigModel): + model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") - voltage: DeviceAccess|None = None + voltage: DeviceAccess | None = None """Device to apply cavity voltage""" - phase: DeviceAccess|None = None + phase: DeviceAccess | None = None """Device to apply cavity phase""" cavities: list[str] """List of cavity names connected to this transmitter""" @@ -28,8 +29,8 @@ class ConfigModel(ElementConfigModel): distribution: float = 1.0 """RF distribution (Part of the total RF voltage powered by this transmitter)""" -class RFTransmitter(Element): +class RFTransmitter(Element): """ Class that handle a RF transmitter """ @@ -43,16 +44,25 @@ def __init__(self, cfg: ConfigModel): @property def voltage(self) -> abstract.ReadWriteFloatScalar: if self.__voltage is None: - raise PyAMLException(f"{str(self)} is unattached or has no voltage device defined") + raise PyAMLException( + f"{str(self)} is unattached or has no voltage device defined" + ) return self.__voltage @property def phase(self) -> abstract.ReadWriteFloatScalar: if self.__phase is None: - raise PyAMLException(f"{str(self)} is unattached or has no phase device defined") + raise PyAMLException( + f"{str(self)} is unattached or has no phase device defined" + ) return self.__phase - def attach(self, peer, voltage: abstract.ReadWriteFloatScalar, phase: abstract.ReadWriteFloatScalar) -> Self: + def attach( + self, + peer, + voltage: abstract.ReadWriteFloatScalar, + phase: abstract.ReadWriteFloatScalar, + ) -> Self: # Attach voltage and phase attribute and returns a new reference obj = self.__class__(self._cfg) obj.__voltage = voltage diff --git a/pyaml/tuning_tools/LOCO/loco.py b/pyaml/tuning_tools/LOCO/loco.py index a153754e..3e0515dc 100644 --- a/pyaml/tuning_tools/LOCO/loco.py +++ b/pyaml/tuning_tools/LOCO/loco.py @@ -1,8 +1,8 @@ -class loco(): +class loco: """loco help""" def run(): """loco run help""" def measure_orm(): - """loco measure help""" \ No newline at end of file + """loco measure help""" diff --git a/pyaml/tuning_tools/SOFB/sofb.py b/pyaml/tuning_tools/SOFB/sofb.py index 8440fb5e..1530e58c 100644 --- a/pyaml/tuning_tools/SOFB/sofb.py +++ b/pyaml/tuning_tools/SOFB/sofb.py @@ -1,8 +1,8 @@ -class slow_orbit_feeback(): +class slow_orbit_feeback: """slow orbit feedback help""" def run(): """run slow orbit feedback""" def get_svd_decomposition(): - """get steerers svd decomposition""" \ No newline at end of file + """get steerers svd decomposition""" diff --git a/pyaml/tuning_tools/TUNE/tune.py b/pyaml/tuning_tools/TUNE/tune.py index 9c3e6629..87f42574 100644 --- a/pyaml/tuning_tools/TUNE/tune.py +++ b/pyaml/tuning_tools/TUNE/tune.py @@ -1,8 +1,8 @@ -class tune_correction(): +class tune_correction: """tune correction help""" def run(): """run tune correction""" def get_tune(): - """measure tune""" \ No newline at end of file + """measure tune""" diff --git a/pyproject.toml b/pyproject.toml index eb2be55a..2a5902fc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -58,7 +58,47 @@ doc = [ "sphinx-rtd-theme" ] +dev = [ + "pre-commit", + "ruff", + "mypy", + "pytest", +] + [project.urls] Homepage = "https://github.com/python-accelerator-middle-layer/pyaml" Documentation = "https://python-accelerator-middle-layer.github.io/pyaml/" -Repository = "https://github.com/python-accelerator-middle-layer/pyaml.git" \ No newline at end of file +Repository = "https://github.com/python-accelerator-middle-layer/pyaml.git" + +[tool.ruff] +line-length = 88 +target-version = "py311" + +[tool.ruff.lint] +# Basic, sane default set: pycodestyle (E,W), pyflakes (F), +# bugbear (B), imports (I), etc. +select = [ + "E", # pycodestyle errors + "W", # pycodestyle warnings + "F", # pyflakes + "B", # flake8-bugbear + "I", # isort-style import sorting +] +ignore = [ + "E203", # so it's compatible with black-style formatting + "F401", # unused imports + "F841", # unused local variables + "E741", # ambiguous variable name + "B024", # abstract class with no methods +] + +[tool.ruff.format] +line-ending = "lf" +quote-style = "double" +indent-style = "space" +skip-magic-trailing-comma = false + +[tool.mypy] +python_version = "3.11" +strict = false +ignore_missing_imports = true diff --git a/tests/conftest.py b/tests/conftest.py index d86afe03..a7737e17 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,14 +1,15 @@ +import pathlib +import subprocess +import sys import types import at -import pytest -import subprocess -import sys -import pathlib import numpy as np -from pyaml.control.readback_value import Value +import pytest from pydantic import BaseModel -from pyaml.configuration.factory import Factory,BuildStrategy + +from pyaml.configuration.factory import BuildStrategy, Factory +from pyaml.control.readback_value import Value @pytest.fixture @@ -18,7 +19,8 @@ def install_test_package(request): The test must provide a dictionary as parameter with: - 'name': name of the installable package (used for pip uninstall) - - 'path': relative path to the package folder (e.g. 'tests/my_dir'). Optional, replaced by package name if absent. + - 'path': relative path to the package folder (e.g. 'tests/my_dir'). + Optional, replaced by package name if absent. Example: -------- @@ -40,25 +42,29 @@ def test_x(install_test_package): package_path = pathlib.Path(info["path"]).resolve() if not package_path.exists(): raise FileNotFoundError(f"Package path not found: {package_path}") - + if package_path is None: - raise RuntimeError(f"No package_path defined for install_test_package fixture") + raise RuntimeError("No package_path defined for install_test_package fixture") - if not ((package_path / "pyproject.toml").exists() or (package_path / "setup.py").exists()): - raise RuntimeError(f"No pyproject.toml or setup.py found in {package_path}") + if not ( + (package_path / "pyproject.toml").exists() + or (package_path / "setup.py").exists() + ): + raise RuntimeError(f"No pyproject.toml or setup.py found in {package_path}") - # Install package in a classis way, `--editable` create a .pth entry in conda env which is imcompatible with submodule + """Install package in a classis way, `--editable` create a .pth entry + in conda env which is imcompatible with submodule""" if package_path is not None: - subprocess.check_call([ - sys.executable, "-m", "pip", "--quiet", "install", str(package_path) - ]) + subprocess.check_call( + [sys.executable, "-m", "pip", "--quiet", "install", str(package_path)] + ) yield package_name # Do not uninstall package at the end to speed up tests a bit - #subprocess.call([ + # subprocess.call([ # sys.executable, "-m", "pip", "uninstall", "-y", package_name - #]) + # ]) @pytest.fixture @@ -105,13 +111,16 @@ def broadcast_matrix(): # ────────────── Simulated module ────────────── + class MockConfig(BaseModel): name: str + class MockElement: def __init__(self, config): self.name = config.name + mock_module = types.ModuleType("mock_module") mock_module.ConfigModel = MockConfig mock_module.PYAMLCLASS = "MockElement" @@ -120,6 +129,7 @@ def __init__(self, config): # ────────────── Custom strategy ────────────── + class MockStrategy(BuildStrategy): def can_handle(self, module, config_dict): return config_dict.get("custom") is True @@ -131,6 +141,7 @@ def build(self, module, config_dict): # ────────────── Pytest fixtures ────────────── + @pytest.fixture(scope="module", autouse=True) def inject_mock_module(): """Inject a simulated external module into sys.modules.""" @@ -164,17 +175,24 @@ def register_mock_strategy(): @pytest.fixture def lattice_with_famnames() -> at.Lattice: """Lattice with duplicate FamName to test multi-match and first-element behavior.""" - qf1 = at.elements.Quadrupole('QF_1', 0.2); qf1.FamName = 'QF' - qf2 = at.elements.Quadrupole('QF_2', 0.25); qf2.FamName = 'QF' - qd1 = at.elements.Quadrupole('QD_1', 0.3); qd1.FamName = 'QD' + qf1 = at.elements.Quadrupole("QF_1", 0.2) + qf1.FamName = "QF" + qf2 = at.elements.Quadrupole("QF_2", 0.25) + qf2.FamName = "QF" + qd1 = at.elements.Quadrupole("QD_1", 0.3) + qd1.FamName = "QD" return at.Lattice([qf1, qf2, qd1], energy=3e9) @pytest.fixture def lattice_with_custom_attr() -> at.Lattice: """Lattice where a custom attribute (e.g., 'Tag') is set on elements.""" - d1 = at.elements.Drift('D1', 1.0); setattr(d1, "Tag", "D1") - qf = at.elements.Quadrupole('QF', 0.2); setattr(qf, "Tag", "QF") - qf2 = at.elements.Quadrupole('QF2', 0.2); setattr(qf2, "Tag", "QF") - qd = at.elements.Quadrupole('QD', 0.3); setattr(qd, "Tag", "QD") + d1 = at.elements.Drift("D1", 1.0) + d1.Tag = "D1" + qf = at.elements.Quadrupole("QF", 0.2) + qf.Tag = "QF" + qf2 = at.elements.Quadrupole("QF2", 0.2) + qf2.Tag = "QF" + qd = at.elements.Quadrupole("QD", 0.3) + qd.Tag = "QD" return at.Lattice([d1, qf, qf2, qd], energy=3e9) diff --git a/tests/dummy_cs/tango-pyaml/tango/pyaml/attribute.py b/tests/dummy_cs/tango-pyaml/tango/pyaml/attribute.py index ea64ceb0..3d9ebbee 100644 --- a/tests/dummy_cs/tango-pyaml/tango/pyaml/attribute.py +++ b/tests/dummy_cs/tango-pyaml/tango/pyaml/attribute.py @@ -1,21 +1,24 @@ -from pydantic import BaseModel,ConfigDict +from pydantic import BaseModel, ConfigDict + from pyaml.control.deviceaccess import DeviceAccess from pyaml.control.readback_value import Value -PYAMLCLASS : str = "Attribute" +PYAMLCLASS: str = "Attribute" -class ConfigModel(BaseModel): - model_config = ConfigDict(arbitrary_types_allowed=True,extra="forbid") +class ConfigModel(BaseModel): + model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") attribute: str unit: str = "" + class Attribute(DeviceAccess): """ - Class that implements a default device class that just prints out + Class that implements a default device class that just prints out values (Debugging purpose) """ + def __init__(self, cfg: ConfigModel): super().__init__() self._cfg = cfg @@ -46,4 +49,4 @@ def unit(self) -> str: return self._unit def __repr__(self): - return repr(self._cfg).replace("ConfigModel",self.__class__.__name__) \ No newline at end of file + return repr(self._cfg).replace("ConfigModel", self.__class__.__name__) diff --git a/tests/dummy_cs/tango-pyaml/tango/pyaml/attribute_read_only.py b/tests/dummy_cs/tango-pyaml/tango/pyaml/attribute_read_only.py index fcf3253f..7be8de4e 100644 --- a/tests/dummy_cs/tango-pyaml/tango/pyaml/attribute_read_only.py +++ b/tests/dummy_cs/tango-pyaml/tango/pyaml/attribute_read_only.py @@ -1,21 +1,25 @@ -from pydantic import BaseModel,ConfigDict -from .attribute import Attribute +from pydantic import BaseModel, ConfigDict + from pyaml.control.readback_value import Value -PYAMLCLASS : str = "AttributeReadOnly" +from .attribute import Attribute -class ConfigModel(BaseModel): +PYAMLCLASS: str = "AttributeReadOnly" - model_config = ConfigDict(arbitrary_types_allowed=True,extra="forbid") + +class ConfigModel(BaseModel): + model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") attribute: str unit: str = "" + class AttributeReadOnly(Attribute): """ - Class that implements a default device class that just prints out + Class that implements a default device class that just prints out values (Debugging purpose) """ + def __init__(self, cfg: ConfigModel): super().__init__(cfg) self._cfg = cfg diff --git a/tests/dummy_cs/tango-pyaml/tango/pyaml/attribute_with_tango_powersupply_mocking_behaviour.py b/tests/dummy_cs/tango-pyaml/tango/pyaml/attribute_with_tango_powersupply_mocking_behaviour.py index 846a1e1e..c53da2d6 100644 --- a/tests/dummy_cs/tango-pyaml/tango/pyaml/attribute_with_tango_powersupply_mocking_behaviour.py +++ b/tests/dummy_cs/tango-pyaml/tango/pyaml/attribute_with_tango_powersupply_mocking_behaviour.py @@ -1,19 +1,19 @@ from pathlib import Path import numpy as np -from pydantic import BaseModel,ConfigDict +from pydantic import BaseModel, ConfigDict from scipy.constants import speed_of_light +from pyaml.configuration import get_root_folder from pyaml.configuration.curve import Curve from pyaml.control.deviceaccess import DeviceAccess from pyaml.control.readback_value import Value -from pyaml.configuration import get_root_folder -PYAMLCLASS : str = "AttributeWithTangoMockingBehaviour" +PYAMLCLASS: str = "AttributeWithTangoMockingBehaviour" -class ConfigModel(BaseModel): - model_config = ConfigDict(arbitrary_types_allowed=True,extra="forbid") +class ConfigModel(BaseModel): + model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") attribute: str calibration_factor: float = 1.0 @@ -26,12 +26,13 @@ class ConfigModel(BaseModel): curve: str = None unit: str = "" + class TangoDevice: def __init__(self, name: str): self.name = name - self._current:float = 0.0 - self._strength:float = 0.0 - self._attributes:dict[str, float] = {} + self._current: float = 0.0 + self._strength: float = 0.0 + self._attributes: dict[str, float] = {} self.__calibration_factor: float = 1.0 self.__calibration_offset: float = 0.0 self.__brho: float = np.nan @@ -80,12 +81,19 @@ def compute_strengths(self, currents: np.array) -> np.array: ) return np.array([_strength]) - def set_magnet_model_data(self: float, calibration_factor, calibration_offset: float, crosstalk: float, magnet_energy: float, curve_file: str): + def set_magnet_model_data( + self: float, + calibration_factor, + calibration_offset: float, + crosstalk: float, + magnet_energy: float, + curve_file: str, + ): self.__calibration_factor = calibration_factor self.__calibration_offset = calibration_offset self.__brho = magnet_energy / speed_of_light self._curve_file = curve_file - path:Path = get_root_folder() / curve_file + path: Path = get_root_folder() / curve_file self.__curve = np.genfromtxt(path, delimiter=",", dtype=float) _s = np.shape(self.__curve) if len(_s) != 2 or _s[1] != 2: @@ -96,7 +104,8 @@ def set_magnet_model_data(self: float, calibration_factor, calibration_offset: f self.__rcurve = Curve.inverse(self.__curve) -TANGO_DEVICES:dict[str, TangoDevice] = {} +TANGO_DEVICES: dict[str, TangoDevice] = {} + def get_device(name: str) -> TangoDevice: if name in TANGO_DEVICES: @@ -106,11 +115,13 @@ def get_device(name: str) -> TangoDevice: TANGO_DEVICES[name] = device return device + class AttributeWithTangoMockingBehaviour(DeviceAccess): """ - Class that implements a default device class that just prints out + Class that implements a default device class that just prints out values (Debugging purpose) """ + def __init__(self, cfg: ConfigModel): super().__init__() self._cfg = cfg @@ -120,7 +131,13 @@ def __init__(self, cfg: ConfigModel): self._device_name, self._attribute_name = cfg.attribute.rsplit("/", 1) self._device = get_device(self._device_name) if cfg.curve is not None: - self._device.set_magnet_model_data(cfg.calibration_factor, cfg.calibration_offset, cfg.crosstalk, cfg.magnet_energy, cfg.curve) + self._device.set_magnet_model_data( + cfg.calibration_factor, + cfg.calibration_offset, + cfg.crosstalk, + cfg.magnet_energy, + cfg.curve, + ) def name(self) -> str: return self._setpoint diff --git a/tests/dummy_cs/tango-pyaml/tango/pyaml/controlsystem.py b/tests/dummy_cs/tango-pyaml/tango/pyaml/controlsystem.py index 922052fc..517d300e 100644 --- a/tests/dummy_cs/tango-pyaml/tango/pyaml/controlsystem.py +++ b/tests/dummy_cs/tango-pyaml/tango/pyaml/controlsystem.py @@ -1,17 +1,19 @@ import os -from pydantic import BaseModel,ConfigDict + +from pydantic import BaseModel, ConfigDict + from pyaml.control.controlsystem import ControlSystem -PYAMLCLASS : str = "TangoControlSystem" +PYAMLCLASS: str = "TangoControlSystem" class ConfigModel(BaseModel): - - model_config = ConfigDict(arbitrary_types_allowed=True,extra="forbid") + model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") name: str tango_host: str - debug_level: str=None + debug_level: str = None + class TangoControlSystem(ControlSystem): def __init__(self, cfg: ConfigModel): @@ -32,4 +34,4 @@ def vector_aggregator(self) -> str | None: return None def __repr__(self): - return repr(self._cfg).replace("ConfigModel",self.__class__.__name__) \ No newline at end of file + return repr(self._cfg).replace("ConfigModel", self.__class__.__name__) diff --git a/tests/dummy_cs/tango-pyaml/tango/pyaml/multi_attribute.py b/tests/dummy_cs/tango-pyaml/tango/pyaml/multi_attribute.py index 263c1d3d..e8e22262 100644 --- a/tests/dummy_cs/tango-pyaml/tango/pyaml/multi_attribute.py +++ b/tests/dummy_cs/tango-pyaml/tango/pyaml/multi_attribute.py @@ -1,17 +1,19 @@ -import pyaml -from numpy import typing as npt import numpy as np -from pyaml.control.deviceaccess import DeviceAccess +from numpy import typing as npt from pydantic import BaseModel +import pyaml +from pyaml.control.deviceaccess import DeviceAccess from pyaml.control.deviceaccesslist import DeviceAccessList + from .attribute import Attribute from .attribute import ConfigModel as AttributeConfigModel -PYAMLCLASS : str = "MultiAttribute" +PYAMLCLASS: str = "MultiAttribute" LAST_NB_WRITTEN = 0 + class ConfigModel(BaseModel): """ Configuration model for a list of DeviceAccess. @@ -25,31 +27,40 @@ class ConfigModel(BaseModel): unit : str, optional Unit of the attributes. """ + attributes: list[str] = [] name: str = "" -class MultiAttribute(DeviceAccessList): - def __init__(self, cfg:ConfigModel=None): +class MultiAttribute(DeviceAccessList): + def __init__(self, cfg: ConfigModel = None): super().__init__() self._cfg = cfg if cfg.attributes: for name in cfg.attributes: - self.add_devices(Attribute(AttributeConfigModel(name,cfg.unit))) + self.add_devices(Attribute(AttributeConfigModel(name, cfg.unit))) def add_devices(self, devices: DeviceAccess | list[DeviceAccess]): if isinstance(devices, list): for device in devices: if not isinstance(device, Attribute): - raise pyaml.PyAMLException(f"All devices must be instances of Attribute (tango.pyaml.attribute) but got ({device.__class__.__name__})") + raise pyaml.PyAMLException( + f"""All devices must be instances of Attribute + (tango.pyaml.attribute) but got + ({device.__class__.__name__})""" + ) super().extend(devices) else: if not isinstance(devices, Attribute): - raise pyaml.PyAMLException(f"Device must be an instance of Attribute (tango.pyaml.attribute) but got ({devices.__class__.__name__})") + raise pyaml.PyAMLException( + f"""Device must be an instance of Attribute + (tango.pyaml.attribute) but got + ({devices.__class__.__name__})""" + ) super().append(devices) def get_devices(self) -> DeviceAccess | list[DeviceAccess]: - if len(self)==1: + if len(self) == 1: return self[0] else: return self @@ -58,7 +69,7 @@ def set(self, value: npt.NDArray[np.float64]): print(f"MultiAttribute.set({len(value)} values)") global LAST_NB_WRITTEN LAST_NB_WRITTEN += len(value) - for idx,a in enumerate(self): + for idx, a in enumerate(self): a.set(value[idx]) def set_and_wait(self, value: npt.NDArray[np.float64]): @@ -73,6 +84,6 @@ def readback(self) -> np.array: def unit(self) -> list[str]: return [a.unit() for a in self] - + def get_last_nb_written(self) -> int: return self.__last_nb_written diff --git a/tests/external/pyaml_external/external_magnet_model.py b/tests/external/pyaml_external/external_magnet_model.py index 162221ec..9f7cf1ad 100644 --- a/tests/external/pyaml_external/external_magnet_model.py +++ b/tests/external/pyaml_external/external_magnet_model.py @@ -1,8 +1,8 @@ -from pydantic import BaseModel,ConfigDict import numpy as np -from pyaml.magnet.model import MagnetModel -from pyaml.control.deviceaccess import DeviceAccess +from pydantic import BaseModel, ConfigDict +from pyaml.control.deviceaccess import DeviceAccess +from pyaml.magnet.model import MagnetModel # Define the main class name for this module (mandatory) PYAMLCLASS = "ExternalMagnetModel" @@ -12,50 +12,54 @@ # configuration file. Parameters with a default value are optional. # The ConfigModel structure is then passed to the constructor of the above class -class ConfigModel(BaseModel): - model_config = ConfigDict(arbitrary_types_allowed=True,extra="forbid") +class ConfigModel(BaseModel): + model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") powersupply: DeviceAccess id: DeviceAccess param: str = None + # ConfigModel inherits from pydantic BaseModel # Take care of having all your private fields prefixed by an `_` + class ExternalMagnetModel(MagnetModel): """ Class example proving manget current/strength conversion """ def __init__(self, cfg: ConfigModel): - print("ExternalMagnetModel:\n powersupply:%s\n id:%s\n param:%s" - % (cfg.powersupply.name(),cfg.id.name(),cfg.param)) + print( + "ExternalMagnetModel:\n powersupply:%s\n id:%s\n param:%s" + % (cfg.powersupply.name(), cfg.id.name(), cfg.param) + ) self._ps = cfg.powersupply self._id = cfg.id - self._id.set(10) # Mimic a default value for the gap - + self._id.set(10) # Mimic a default value for the gap + # Implementation of the UnitConv abstract class # Get coil current(s) from magnet strength(s) - def compute_hardware_values(self,strengths:np.array) -> np.array: + def compute_hardware_values(self, strengths: np.array) -> np.array: print("ExternalMagnetModel.compute_hardware_values(%f)" % (strengths[0])) _gap = self._id.get() - return np.array([strengths[0]*_gap]) + return np.array([strengths[0] * _gap]) # Get magnet strength(s) from coil current(s) - def compute_strengths(self,currents:np.array) -> np.array: + def compute_strengths(self, currents: np.array) -> np.array: print("ExternalMagnetModel.compute_strengths(%f)" % (currents[0])) _gap = self._id.get() - return np.array([currents[0]/_gap]) + return np.array([currents[0] / _gap]) # Get strength units def get_strength_units(self) -> list[str]: - return["rad"] + return ["rad"] # Get current units def get_hardware_units(self) -> list[str]: - return["A"] + return ["A"] # Get power supply current setpoint(s) from control system def read_hardware_values(self) -> np.array: @@ -66,18 +70,17 @@ def readback_hardware_values(self) -> np.array: pass # Send power supply current(s) to control system - def send_hardware_values(self, currents:np.array): + def send_hardware_values(self, currents: np.array): self._ps.set(currents) pass def get_devices(self) -> list[DeviceAccess]: - return [self._ps,self.id] - + return [self._ps, self.id] + def has_hardware(self) -> bool: # No trivial conversion between strength and hardware unit return False # Set magnet rigidity - def set_magnet_rigidity(self,brho:np.double): + def set_magnet_rigidity(self, brho: np.double): pass - diff --git a/tests/lattice_info.py b/tests/lattice_info.py index 52f8f3f8..ef973b34 100644 --- a/tests/lattice_info.py +++ b/tests/lattice_info.py @@ -1,31 +1,36 @@ import at -ring = at.load_lattice("config/sr/lattices/ebs.mat") +ring = at.load_lattice("config/sr/lattices/ebs.mat") def prepare_and_save(): - # Remove non standard field # Add cell number to FamName for e in ring: attrs = e.__dict__ - toremove = [k for k,v in attrs.items() if k.startswith('Calibration') or k=="SerialNumber"] + toremove = [ + k + for k, v in attrs.items() + if k.startswith("Calibration") or k == "SerialNumber" + ] for tr in toremove: - delattr(e,tr) - if hasattr(e,"Device") and not isinstance(e,at.Monitor): - field = e.Device.split(sep='/') - cell = field[2].split(sep='-') + delattr(e, tr) + if hasattr(e, "Device") and not isinstance(e, at.Monitor): + field = e.Device.split(sep="/") + cell = field[2].split(sep="-") CELL = cell[0].upper() e.FamName += f"-{CELL}" - if hasattr(e,"Device"): - delattr(e,"Device") + if hasattr(e, "Device"): + delattr(e, "Device") + + at.save_lattice(ring, "config/sr/lattices/ebs_jlp.mat") - at.save_lattice(ring,"config/sr/lattices/ebs_jlp.mat") def dump(): - # Dump lattice + # Dump lattice for e in ring: if e.FamName.startswith("SI"): print(e) -dump() \ No newline at end of file + +dump() diff --git a/tests/test_arrays.py b/tests/test_arrays.py index 77f49df8..fa33fe06 100644 --- a/tests/test_arrays.py +++ b/tests/test_arrays.py @@ -1,21 +1,23 @@ -from pyaml.configuration.factory import Factory -from pyaml.accelerator import Accelerator -from pyaml.arrays.element_array import ElementArray -from pyaml.arrays.magnet_array import MagnetArray -from pyaml.arrays.cfm_magnet_array import CombinedFunctionMagnetArray -from pyaml.arrays.bpm_array import BPMArray import importlib import numpy as np import pytest -@pytest.mark.parametrize("install_test_package", [{ - "name": "tango-pyaml", - "path": "tests/dummy_cs/tango-pyaml" -}], indirect=True) -def test_arrays(install_test_package): +from pyaml.accelerator import Accelerator +from pyaml.arrays.bpm_array import BPMArray +from pyaml.arrays.cfm_magnet_array import CombinedFunctionMagnetArray +from pyaml.arrays.element_array import ElementArray +from pyaml.arrays.magnet_array import MagnetArray +from pyaml.configuration.factory import Factory - sr:Accelerator = Accelerator.load("tests/config/sr.yaml") + +@pytest.mark.parametrize( + "install_test_package", + [{"name": "tango-pyaml", "path": "tests/dummy_cs/tango-pyaml"}], + indirect=True, +) +def test_arrays(install_test_package): + sr: Accelerator = Accelerator.load("tests/config/sr.yaml") sr.design.get_lattice().disable_6d() # Test on model @@ -23,99 +25,103 @@ def test_arrays(install_test_package): sr.design.get_magnet("SH1A-C01-H").strength.set(0.000010) sr.design.get_magnet("SH1A-C01-V").strength.set(0.000015) - o,_ = sr.design.get_lattice().find_orbit() - assert(np.abs(o[0] + 9.91848416e-05)<1e-10) - assert(np.abs(o[1] + 3.54829761e-07)<1e-10) - assert(np.abs(o[2] + 1.56246320e-06)<1e-10) - assert(np.abs(o[3] + 1.75037311e-05)<1e-10) + o, _ = sr.design.get_lattice().find_orbit() + assert np.abs(o[0] + 9.91848416e-05) < 1e-10 + assert np.abs(o[1] + 3.54829761e-07) < 1e-10 + assert np.abs(o[2] + 1.56246320e-06) < 1e-10 + assert np.abs(o[3] + 1.75037311e-05) < 1e-10 sr.design.get_magnet("SH1A-C02-H").strength.set(-0.000008) sr.design.get_magnet("SH1A-C02-V").strength.set(-0.000017) - o,_ = sr.design.get_lattice().find_orbit() - assert(np.abs(o[0] + 1.60277642e-04)<1e-10) - assert(np.abs(o[1] - 2.36103795e-06)<1e-10) - assert(np.abs(o[2] - 3.62843295e-05)<1e-10) - assert(np.abs(o[3] + 6.06571010e-06)<1e-10) + o, _ = sr.design.get_lattice().find_orbit() + assert np.abs(o[0] + 1.60277642e-04) < 1e-10 + assert np.abs(o[1] - 2.36103795e-06) < 1e-10 + assert np.abs(o[2] - 3.62843295e-05) < 1e-10 + assert np.abs(o[3] + 6.06571010e-06) < 1e-10 - sr.design.get_magnets("HCORR").strengths.set([0.000010,-0.000008]) - sr.design.get_magnets("VCORR").strengths.set([0.000015,-0.000017]) + sr.design.get_magnets("HCORR").strengths.set([0.000010, -0.000008]) + sr.design.get_magnets("VCORR").strengths.set([0.000015, -0.000017]) - o,_ = sr.design.get_lattice().find_orbit() - assert(np.abs(o[0] + 1.60277642e-04)<1e-10) - assert(np.abs(o[1] - 2.36103795e-06)<1e-10) - assert(np.abs(o[2] - 3.62843295e-05)<1e-10) - assert(np.abs(o[3] + 6.06571010e-06)<1e-10) + o, _ = sr.design.get_lattice().find_orbit() + assert np.abs(o[0] + 1.60277642e-04) < 1e-10 + assert np.abs(o[1] - 2.36103795e-06) < 1e-10 + assert np.abs(o[2] - 3.62843295e-05) < 1e-10 + assert np.abs(o[3] + 6.06571010e-06) < 1e-10 # Test on control system # Assert that the virtual magnet share the same model - assert( sr.live.get_magnet("SH1A-C01-H").model == sr.live.get_magnet("SH1A-C01-V").model ) - assert( sr.live.get_magnet("SH1A-C02-H").model == sr.live.get_magnet("SH1A-C02-V").model ) + assert ( + sr.live.get_magnet("SH1A-C01-H").model == sr.live.get_magnet("SH1A-C01-V").model + ) + assert ( + sr.live.get_magnet("SH1A-C02-H").model == sr.live.get_magnet("SH1A-C02-V").model + ) # Using aggregators - sr.live.get_magnets("HCORR").strengths.set([0.000010,-0.000008]) + sr.live.get_magnets("HCORR").strengths.set([0.000010, -0.000008]) ps1 = sr.live.get_magnet("SH1A-C01-H").model.read_hardware_values() ps2 = sr.live.get_magnet("SH1A-C02-H").model.read_hardware_values() - assert(np.abs(ps1[0] - 0.02956737880874648)<1e-10) - assert(np.abs(ps1[1] - 0)<1e-10) - assert(np.abs(ps1[2] - 0)<1e-10) - assert(np.abs(ps2[0] + 0.02365390304699716)<1e-10) - assert(np.abs(ps2[1] - 0)<1e-10) - assert(np.abs(ps2[2] - 0)<1e-10) - sr.live.get_magnets("VCORR").strengths.set([0.000015,-0.000017]) + assert np.abs(ps1[0] - 0.02956737880874648) < 1e-10 + assert np.abs(ps1[1] - 0) < 1e-10 + assert np.abs(ps1[2] - 0) < 1e-10 + assert np.abs(ps2[0] + 0.02365390304699716) < 1e-10 + assert np.abs(ps2[1] - 0) < 1e-10 + assert np.abs(ps2[2] - 0) < 1e-10 + sr.live.get_magnets("VCORR").strengths.set([0.000015, -0.000017]) ps1 = sr.live.get_magnet("SH1A-C01-H").model.read_hardware_values() ps2 = sr.live.get_magnet("SH1A-C02-H").model.read_hardware_values() - assert(np.abs(ps1[0] - 0.02956737880874648)<1e-10) - assert(np.abs(ps1[1] + 0.058240333933172066)<1e-10) - assert(np.abs(ps1[2] - 0.05601656539392866)<1e-10) - assert(np.abs(ps2[0] + 0.02365390304699716)<1e-10) - assert(np.abs(ps2[1] - 0.06600571179092833)<1e-10) - assert(np.abs(ps2[2] + 0.0634854407797858)<1e-10) + assert np.abs(ps1[0] - 0.02956737880874648) < 1e-10 + assert np.abs(ps1[1] + 0.058240333933172066) < 1e-10 + assert np.abs(ps1[2] - 0.05601656539392866) < 1e-10 + assert np.abs(ps2[0] + 0.02365390304699716) < 1e-10 + assert np.abs(ps2[1] - 0.06600571179092833) < 1e-10 + assert np.abs(ps2[2] + 0.0634854407797858) < 1e-10 strHV = sr.live.get_magnets("HVCORR").strengths.get() - assert(np.abs(strHV[0] - 0.000010)<1e-10) - assert(np.abs(strHV[1] + 0.000008)<1e-10) - assert(np.abs(strHV[2] - 0.000015)<1e-10) - assert(np.abs(strHV[3] + 0.000017)<1e-10) + assert np.abs(strHV[0] - 0.000010) < 1e-10 + assert np.abs(strHV[1] + 0.000008) < 1e-10 + assert np.abs(strHV[2] - 0.000015) < 1e-10 + assert np.abs(strHV[3] + 0.000017) < 1e-10 # Reset to 0 ma = importlib.import_module("tango.pyaml.multi_attribute") - ma.LAST_NB_WRITTEN = 0 # Total number of setpoints done by multi_attribute - sr.live.get_magnets("HVCORR").strengths.set(0.) - assert(ma.LAST_NB_WRITTEN == 6) # 6 power supply setpoints are needed + ma.LAST_NB_WRITTEN = 0 # Total number of setpoints done by multi_attribute + sr.live.get_magnets("HVCORR").strengths.set(0.0) + assert ma.LAST_NB_WRITTEN == 6 # 6 power supply setpoints are needed ps1 = sr.live.get_magnet("SH1A-C01-H").model.read_hardware_values() ps2 = sr.live.get_magnet("SH1A-C02-H").model.read_hardware_values() - assert(np.abs(ps1[0])<1e-10) - assert(np.abs(ps1[1])<1e-10) - assert(np.abs(ps1[2])<1e-10) - assert(np.abs(ps2[0])<1e-10) - assert(np.abs(ps2[1])<1e-10) - assert(np.abs(ps2[2])<1e-10) + assert np.abs(ps1[0]) < 1e-10 + assert np.abs(ps1[1]) < 1e-10 + assert np.abs(ps1[2]) < 1e-10 + assert np.abs(ps2[0]) < 1e-10 + assert np.abs(ps2[1]) < 1e-10 + assert np.abs(ps2[2]) < 1e-10 # Check that the behavior is the same without aggregator mags = [] - for m in sr.live.get_magnets("HVCORR"): + for m in sr.live.get_magnets("HVCORR"): mags.append(m) - array = MagnetArray("HVCOOR_noagg",mags,use_aggregator=False) - array.strengths.set([0.000010,-0.000008,0.000015,-0.000017]) + array = MagnetArray("HVCOOR_noagg", mags, use_aggregator=False) + array.strengths.set([0.000010, -0.000008, 0.000015, -0.000017]) ps1 = sr.live.get_magnet("SH1A-C01-H").model.read_hardware_values() ps2 = sr.live.get_magnet("SH1A-C02-H").model.read_hardware_values() - assert(np.abs(ps1[0] - 0.02956737880874648)<1e-10) - assert(np.abs(ps1[1] + 0.058240333933172066)<1e-10) - assert(np.abs(ps1[2] - 0.05601656539392866)<1e-10) - assert(np.abs(ps2[0] + 0.02365390304699716)<1e-10) - assert(np.abs(ps2[1] - 0.06600571179092833)<1e-10) - assert(np.abs(ps2[2] + 0.0634854407797858)<1e-10) + assert np.abs(ps1[0] - 0.02956737880874648) < 1e-10 + assert np.abs(ps1[1] + 0.058240333933172066) < 1e-10 + assert np.abs(ps1[2] - 0.05601656539392866) < 1e-10 + assert np.abs(ps2[0] + 0.02365390304699716) < 1e-10 + assert np.abs(ps2[1] - 0.06600571179092833) < 1e-10 + assert np.abs(ps2[2] + 0.0634854407797858) < 1e-10 # Test BPMs array # Using aggregator pos = sr.design.get_bpms("BPMS").positions.get() - assert(np.abs(pos[0][0] + 7.21154171490481e-05)<1e-10) - assert(np.abs(pos[0][1] - 3.3988843436571406e-05)<1e-10) - assert(np.abs(pos[1][0] - 1.1681211772781844e-04)<1e-10) - assert(np.abs(pos[1][1] - 7.072972488250373e-06)<1e-10) + assert np.abs(pos[0][0] + 7.21154171490481e-05) < 1e-10 + assert np.abs(pos[0][1] - 3.3988843436571406e-05) < 1e-10 + assert np.abs(pos[1][0] - 1.1681211772781844e-04) < 1e-10 + assert np.abs(pos[1][1] - 7.072972488250373e-06) < 1e-10 # Using aggregator (h and v) pos_h = sr.design.get_bpms("BPMS").h.get() @@ -128,73 +134,73 @@ def test_arrays(install_test_package): for b in sr.design.get_bpms("BPMS"): bpms.append(b) - bpms = BPMArray("BPM_noagg",bpms,use_aggregator=False) + bpms = BPMArray("BPM_noagg", bpms, use_aggregator=False) pos = bpms.positions.get() - assert(np.abs(pos[0][0] + 7.21154171490481e-05)<1e-10) - assert(np.abs(pos[0][1] - 3.3988843436571406e-05)<1e-10) - assert(np.abs(pos[1][0] - 1.1681211772781844e-04)<1e-10) - assert(np.abs(pos[1][1] - 7.072972488250373e-06)<1e-10) + assert np.abs(pos[0][0] + 7.21154171490481e-05) < 1e-10 + assert np.abs(pos[0][1] - 3.3988843436571406e-05) < 1e-10 + assert np.abs(pos[1][0] - 1.1681211772781844e-04) < 1e-10 + assert np.abs(pos[1][1] - 7.072972488250373e-06) < 1e-10 # Radom array elts = sr.design.get_elemens("ElArray") # Create an array that contains all elements - allElts = ElementArray("AllElements",sr.design.get_all_elements()) - assert(len(allElts)==11) + allElts = ElementArray("AllElements", sr.design.get_all_elements()) + assert len(allElts) == 11 # Create an array that contains all elements - allMags = MagnetArray("AllMagnets",sr.design.get_all_magnets()) - assert(len(allMags)==7) + allMags = MagnetArray("AllMagnets", sr.design.get_all_magnets()) + assert len(allMags) == 7 # Create an array that contains all BPM - allBpms = BPMArray("AllBPMs",sr.design.get_all_bpms()) - assert(len(allBpms)==2) + allBpms = BPMArray("AllBPMs", sr.design.get_all_bpms()) + assert len(allBpms) == 2 cfm = sr.design.get_cfm_magnets("CFM") strHVSQ = cfm.strengths.get() - assert(np.abs(strHVSQ[0] - 0.000010)<1e-10) # H - assert(np.abs(strHVSQ[1] - 0.000015)<1e-10) # V - assert(np.abs(strHVSQ[2] + 0.0)<1e-10) # SQ - assert(np.abs(strHVSQ[3] + 0.000008)<1e-10) # H - assert(np.abs(strHVSQ[4] + 0.000017)<1e-10) # V - assert(np.abs(strHVSQ[5] + 0.0)<1e-10) # SQ + assert np.abs(strHVSQ[0] - 0.000010) < 1e-10 # H + assert np.abs(strHVSQ[1] - 0.000015) < 1e-10 # V + assert np.abs(strHVSQ[2] + 0.0) < 1e-10 # SQ + assert np.abs(strHVSQ[3] + 0.000008) < 1e-10 # H + assert np.abs(strHVSQ[4] + 0.000017) < 1e-10 # V + assert np.abs(strHVSQ[5] + 0.0) < 1e-10 # SQ strHVSQu = cfm.strengths.unit() - assert(strHVSQu[0] == "rad") - assert(strHVSQu[1] == "rad") - assert(strHVSQu[2] == "m-1") - assert(strHVSQu[3] == "rad") - assert(strHVSQu[4] == "rad") - assert(strHVSQu[5] == "m-1") - - cfm.strengths.set([.000010,0.000015,1e-6,-0.000008,-0.000017,1e-6]) + assert strHVSQu[0] == "rad" + assert strHVSQu[1] == "rad" + assert strHVSQu[2] == "m-1" + assert strHVSQu[3] == "rad" + assert strHVSQu[4] == "rad" + assert strHVSQu[5] == "m-1" + + cfm.strengths.set([0.000010, 0.000015, 1e-6, -0.000008, -0.000017, 1e-6]) strHVSQ = cfm.strengths.get() - assert(np.abs(strHVSQ[0] - 0.000010)<1e-10) # H - assert(np.abs(strHVSQ[1] - 0.000015)<1e-10) # V - assert(np.abs(strHVSQ[2] - 1.e-6)<1e-10) # SQ - assert(np.abs(strHVSQ[3] + 0.000008)<1e-10) # H - assert(np.abs(strHVSQ[4] + 0.000017)<1e-10) # V - assert(np.abs(strHVSQ[5] - 1e-6)<1e-10) # SQ + assert np.abs(strHVSQ[0] - 0.000010) < 1e-10 # H + assert np.abs(strHVSQ[1] - 0.000015) < 1e-10 # V + assert np.abs(strHVSQ[2] - 1.0e-6) < 1e-10 # SQ + assert np.abs(strHVSQ[3] + 0.000008) < 1e-10 # H + assert np.abs(strHVSQ[4] + 0.000017) < 1e-10 # V + assert np.abs(strHVSQ[5] - 1e-6) < 1e-10 # SQ Factory.clear() # Test dynamic arrays - sr:Accelerator = Accelerator.load("tests/config/EBSOrbit.yaml",use_fast_loader=True) - ae = ElementArray("All",sr.design.get_all_elements()) - acfm = ElementArray("AllCFM",sr.design.get_all_cfm_magnets(),use_aggregator=False) - - bpmC5 = ae['BPM*'][10:20] # All BPM C5 - assert(isinstance(bpmC5,BPMArray) and len(bpmC5)==10) - bpmc10 = ae['BPM*C10*'] # All BPM C10 - assert(isinstance(bpmc10,BPMArray) and len(bpmc10)==10) - magSHV = ae['SH*-V'] # All SH vertical corrector - assert(isinstance(magSHV,MagnetArray) and len(magSHV)==96) - magSH1A = ae['model_name:SH1A-*'] # All SH1A magnets (including the CFMs) - assert(isinstance(magSH1A,ElementArray) and len(magSH1A)==129) - magSH1AC = acfm['model_name:SH1A-*'] # All SH1A magnets (CFMs only) - assert(isinstance(magSH1AC,ElementArray) and len(magSH1AC)==32) + sr: Accelerator = Accelerator.load( + "tests/config/EBSOrbit.yaml", use_fast_loader=True + ) + ae = ElementArray("All", sr.design.get_all_elements()) + acfm = ElementArray("AllCFM", sr.design.get_all_cfm_magnets(), use_aggregator=False) + + bpmC5 = ae["BPM*"][10:20] # All BPM C5 + assert isinstance(bpmC5, BPMArray) and len(bpmC5) == 10 + bpmc10 = ae["BPM*C10*"] # All BPM C10 + assert isinstance(bpmc10, BPMArray) and len(bpmc10) == 10 + magSHV = ae["SH*-V"] # All SH vertical corrector + assert isinstance(magSHV, MagnetArray) and len(magSHV) == 96 + magSH1A = ae["model_name:SH1A-*"] # All SH1A magnets (including the CFMs) + assert isinstance(magSH1A, ElementArray) and len(magSH1A) == 129 + magSH1AC = acfm["model_name:SH1A-*"] # All SH1A magnets (CFMs only) + assert isinstance(magSH1AC, ElementArray) and len(magSH1AC) == 32 Factory.clear() - - diff --git a/tests/test_bpm.py b/tests/test_bpm.py index aaced16f..3699276c 100644 --- a/tests/test_bpm.py +++ b/tests/test_bpm.py @@ -1,75 +1,79 @@ -from pyaml.accelerator import Accelerator -from pyaml.configuration.factory import Factory import numpy as np import pytest -@pytest.mark.parametrize("install_test_package", [{ - "name": "tango-pyaml", - "path": "tests/dummy_cs/tango-pyaml" -}], indirect=True) -def test_simulator_bpm_tilt(install_test_package): +from pyaml.accelerator import Accelerator +from pyaml.configuration.factory import Factory + - sr:Accelerator = Accelerator.load("tests/config/bpms.yaml") +@pytest.mark.parametrize( + "install_test_package", + [{"name": "tango-pyaml", "path": "tests/dummy_cs/tango-pyaml"}], + indirect=True, +) +def test_simulator_bpm_tilt(install_test_package): + sr: Accelerator = Accelerator.load("tests/config/bpms.yaml") sr.design.get_lattice().disable_6d() - bpm = sr.design.get_bpm('BPM_C01-01') + bpm = sr.design.get_bpm("BPM_C01-01") assert bpm.tilt.get() == 0 bpm.tilt.set(0.01) assert bpm.tilt.get() == 0.01 Factory.clear() -@pytest.mark.parametrize("install_test_package", [{ - "name": "tango-pyaml", - "path": "tests/dummy_cs/tango-pyaml" -}], indirect=True) -def test_simulator_bpm_offset(install_test_package): - sr:Accelerator = Accelerator.load("tests/config/bpms.yaml") +@pytest.mark.parametrize( + "install_test_package", + [{"name": "tango-pyaml", "path": "tests/dummy_cs/tango-pyaml"}], + indirect=True, +) +def test_simulator_bpm_offset(install_test_package): + sr: Accelerator = Accelerator.load("tests/config/bpms.yaml") sr.design.get_lattice().disable_6d() - bpm = sr.design.get_bpm('BPM_C01-01') + bpm = sr.design.get_bpm("BPM_C01-01") assert bpm.offset.get()[0] == 0 assert bpm.offset.get()[1] == 0 - bpm.offset.set( np.array([0.1,0.2]) ) + bpm.offset.set(np.array([0.1, 0.2])) assert bpm.offset.get()[0] == 0.1 assert bpm.offset.get()[1] == 0.2 - assert np.allclose( bpm.positions.get(), np.array([0.0,0.0]) ) + assert np.allclose(bpm.positions.get(), np.array([0.0, 0.0])) Factory.clear() -@pytest.mark.parametrize("install_test_package", [{ - "name": "tango-pyaml", - "path": "tests/dummy_cs/tango-pyaml" -}], indirect=True) -def test_simulator_bpm_position(install_test_package): - sr:Accelerator = Accelerator.load("tests/config/bpms.yaml") +@pytest.mark.parametrize( + "install_test_package", + [{"name": "tango-pyaml", "path": "tests/dummy_cs/tango-pyaml"}], + indirect=True, +) +def test_simulator_bpm_position(install_test_package): + sr: Accelerator = Accelerator.load("tests/config/bpms.yaml") sr.design.get_lattice().disable_6d() - bpm = sr.design.get_bpm('BPM_C01-01') - bpm_simple = sr.live.get_bpm('BPM_C01-02') + bpm = sr.design.get_bpm("BPM_C01-01") + bpm_simple = sr.live.get_bpm("BPM_C01-02") + + assert np.allclose(bpm.positions.get(), np.array([0.0, 0.0])) + assert np.allclose(bpm_simple.positions.get(), np.array([0.0, 0.0])) - assert np.allclose( bpm.positions.get(), np.array([0.0,0.0]) ) - assert np.allclose( bpm_simple.positions.get(), np.array([0.0,0.0]) ) - Factory.clear() -@pytest.mark.parametrize("install_test_package", [{ - "name": "tango-pyaml", - "path": "tests/dummy_cs/tango-pyaml" -}], indirect=True) -def test_simulator_bpm_position_with_bad_corrector_strength(install_test_package): - sr:Accelerator = Accelerator.load("tests/config/bpms.yaml") +@pytest.mark.parametrize( + "install_test_package", + [{"name": "tango-pyaml", "path": "tests/dummy_cs/tango-pyaml"}], + indirect=True, +) +def test_simulator_bpm_position_with_bad_corrector_strength(install_test_package): + sr: Accelerator = Accelerator.load("tests/config/bpms.yaml") sr.design.get_lattice().disable_6d() - bpm = sr.design.get_bpm('BPM_C01-01') - bpm_simple = sr.design.get_bpm('BPM_C01-02') - bpm3 = sr.design.get_bpm('BPM_C01-03') + bpm1 = sr.design.get_bpm("BPM_C01-01") + bpm_simple = sr.design.get_bpm("BPM_C01-02") + bpm3 = sr.design.get_bpm("BPM_C01-03") sr.design.get_magnet("SH1A-C01-H").strength.set(-1e-6) sr.design.get_magnet("SH1A-C01-V").strength.set(-1e-6) - for bpm in [bpm, bpm_simple, bpm3]: + for bpm in [bpm1, bpm_simple, bpm3]: assert bpm.positions.get()[0] != 0.0 assert bpm.positions.get()[1] != 0.0 - - Factory.clear() + Factory.clear() diff --git a/tests/test_bpm_controlsystem.py b/tests/test_bpm_controlsystem.py index 89968ae2..dcc7764f 100644 --- a/tests/test_bpm_controlsystem.py +++ b/tests/test_bpm_controlsystem.py @@ -1,53 +1,57 @@ -from pyaml.accelerator import Accelerator -from pyaml.configuration.factory import Factory import numpy as np import pytest -@pytest.mark.parametrize("install_test_package", [{ - "name": "tango-pyaml", - "path": "tests/dummy_cs/tango-pyaml" -}], indirect=True) -def test_controlsystem_bpm_tilt(install_test_package): +from pyaml.accelerator import Accelerator +from pyaml.configuration.factory import Factory + - sr:Accelerator = Accelerator.load("tests/config/bpms.yaml") - bpm = sr.live.get_bpm('BPM_C01-01') +@pytest.mark.parametrize( + "install_test_package", + [{"name": "tango-pyaml", "path": "tests/dummy_cs/tango-pyaml"}], + indirect=True, +) +def test_controlsystem_bpm_tilt(install_test_package): + sr: Accelerator = Accelerator.load("tests/config/bpms.yaml") + bpm = sr.live.get_bpm("BPM_C01-01") print(bpm.tilt.get()) - + assert bpm.tilt.get() == 0 bpm.tilt.set(0.01) assert bpm.tilt.get() == 0.01 Factory.clear() -@pytest.mark.parametrize("install_test_package", [{ - "name": "tango-pyaml", - "path": "tests/dummy_cs/tango-pyaml" -}], indirect=True) -def test_controlsystem_bpm_offset(install_test_package): - sr:Accelerator = Accelerator.load("tests/config/bpms.yaml") - bpm = sr.live.get_bpm('BPM_C01-01') +@pytest.mark.parametrize( + "install_test_package", + [{"name": "tango-pyaml", "path": "tests/dummy_cs/tango-pyaml"}], + indirect=True, +) +def test_controlsystem_bpm_offset(install_test_package): + sr: Accelerator = Accelerator.load("tests/config/bpms.yaml") + bpm = sr.live.get_bpm("BPM_C01-01") assert bpm.offset.get()[0] == 0 assert bpm.offset.get()[1] == 0 - bpm.offset.set( np.array([0.1,0.2]) ) + bpm.offset.set(np.array([0.1, 0.2])) assert bpm.offset.get()[0] == 0.1 assert bpm.offset.get()[1] == 0.2 - assert np.allclose( bpm.positions.get(), np.array([0.0,0.0]) ) + assert np.allclose(bpm.positions.get(), np.array([0.0, 0.0])) Factory.clear() -@pytest.mark.parametrize("install_test_package", [{ - "name": "tango-pyaml", - "path": "tests/dummy_cs/tango-pyaml" -}], indirect=True) + +@pytest.mark.parametrize( + "install_test_package", + [{"name": "tango-pyaml", "path": "tests/dummy_cs/tango-pyaml"}], + indirect=True, +) def test_controlsystem_bpm_position(install_test_package): + sr: Accelerator = Accelerator.load("tests/config/bpms.yaml") + bpm = sr.live.get_bpm("BPM_C01-01") + bpm_simple = sr.live.get_bpm("BPM_C01-02") - sr:Accelerator = Accelerator.load("tests/config/bpms.yaml") - bpm = sr.live.get_bpm('BPM_C01-01') - bpm_simple = sr.live.get_bpm('BPM_C01-02') + assert np.allclose(bpm.positions.get(), np.array([0.0, 0.0])) + assert np.allclose(bpm_simple.positions.get(), np.array([0.0, 0.0])) - assert np.allclose( bpm.positions.get(), np.array([0.0,0.0]) ) - assert np.allclose( bpm_simple.positions.get(), np.array([0.0,0.0]) ) - Factory.clear() diff --git a/tests/test_curve.py b/tests/test_curve.py index ce252bf3..39cf6693 100644 --- a/tests/test_curve.py +++ b/tests/test_curve.py @@ -1,22 +1,22 @@ -from pyaml.configuration.csvcurve import CSVCurve,ConfigModel -from pyaml.configuration.curve import Curve -from pyaml.configuration import set_root_folder import numpy as np +from pyaml.configuration import set_root_folder +from pyaml.configuration.csvcurve import ConfigModel, CSVCurve +from pyaml.configuration.curve import Curve -def curve_test(file:str,current:float,strength:float): +def curve_test(file: str, current: float, strength: float): curveConfig = ConfigModel(file=file) curve = CSVCurve(curveConfig) curveData = curve.get_curve() icurveData = Curve.inverse(curveData) - x1 = np.interp( current , curveData[:, 0], curveData[:, 1] ) - assert( np.abs(x1-strength) < 1e-6 ) - y1 = np.interp( x1 , icurveData[:, 0], icurveData[:, 1] ) - assert( np.abs(y1-current) < 1e-6 ) + x1 = np.interp(current, curveData[:, 0], curveData[:, 1]) + assert np.abs(x1 - strength) < 1e-6 + y1 = np.interp(x1, icurveData[:, 0], icurveData[:, 1]) + assert np.abs(y1 - current) < 1e-6 def test_curve(config_root_dir): set_root_folder(config_root_dir) - curve_test("sr/magnet_models/QF1_strength.csv",85,16.618788) - curve_test("sr/magnet_models/QD2_strength.csv",87,-12.319894) \ No newline at end of file + curve_test("sr/magnet_models/QF1_strength.csv", 85, 16.618788) + curve_test("sr/magnet_models/QD2_strength.csv", 87, -12.319894) diff --git a/tests/test_errors.py b/tests/test_errors.py index f890633b..318288e5 100644 --- a/tests/test_errors.py +++ b/tests/test_errors.py @@ -1,46 +1,49 @@ -from pyaml.common.exception import PyAMLConfigException,PyAMLException -from pyaml.configuration.factory import Factory +import pytest + from pyaml.accelerator import Accelerator from pyaml.arrays.magnet_array import MagnetArray -import pytest +from pyaml.common.exception import PyAMLConfigException, PyAMLException +from pyaml.configuration.factory import Factory -@pytest.mark.parametrize("install_test_package", [{ - "name": "tango-pyaml", - "path": "tests/dummy_cs/tango-pyaml" -}], indirect=True) -def test_tune(install_test_package): +@pytest.mark.parametrize( + "install_test_package", + [{"name": "tango-pyaml", "path": "tests/dummy_cs/tango-pyaml"}], + indirect=True, +) +def test_tune(install_test_package): with pytest.raises(PyAMLConfigException) as exc: - ml:Accelerator = Accelerator.load("tests/config/bad_conf_duplicate_1.yaml") - assert("MagnetArray HCORR : duplicate name SH1A-C02-H @index 2" in str(exc)) + ml: Accelerator = Accelerator.load("tests/config/bad_conf_duplicate_1.yaml") + assert "MagnetArray HCORR : duplicate name SH1A-C02-H @index 2" in str(exc) Factory.clear() with pytest.raises(PyAMLConfigException) as exc: - ml:Accelerator = Accelerator.load("tests/config/bad_conf_duplicate_2.yaml") - assert("BPMArray BPM : duplicate name BPM_C04-06 @index 3" in str(exc)) + ml: Accelerator = Accelerator.load("tests/config/bad_conf_duplicate_2.yaml") + assert "BPMArray BPM : duplicate name BPM_C04-06 @index 3" in str(exc) Factory.clear() with pytest.raises(PyAMLConfigException) as exc: - ml:Accelerator = Accelerator.load("tests/config/bad_conf_duplicate_3.yaml") - assert("element BPM_C04-06 already defined" in str(exc)) - assert("line 57, column 3" in str(exc)) + ml: Accelerator = Accelerator.load("tests/config/bad_conf_duplicate_3.yaml") + assert "element BPM_C04-06 already defined" in str(exc) + assert "line 57, column 3" in str(exc) Factory.clear() - sr:Accelerator = Accelerator.load("tests/config/EBSTune.yaml") + sr: Accelerator = Accelerator.load("tests/config/EBSTune.yaml") m1 = sr.live.get_magnet("QF1E-C04") m2 = sr.design.get_magnet("QF1A-C05") with pytest.raises(PyAMLException) as exc: - ma = MagnetArray("Test",[m1,m2]) - assert("MagnetArray Test: All elements must be attached to the same instance" in str(exc)) + ma = MagnetArray("Test", [m1, m2]) + assert ( + "MagnetArray Test: All elements must be attached to the same instance" + in str(exc) + ) with pytest.raises(PyAMLException) as exc: m2 = sr.design.get_magnet("QF1A-C05XX") - assert("Magnet QF1A-C05XX not defined" in str(exc)) + assert "Magnet QF1A-C05XX not defined" in str(exc) with pytest.raises(PyAMLException) as exc: m2 = sr.design.get_bpm("QF1A-C05XX") - assert("BPM QF1A-C05XX not defined" in str(exc)) + assert "BPM QF1A-C05XX not defined" in str(exc) Factory.clear() - - diff --git a/tests/test_factory.py b/tests/test_factory.py index ffedfb48..8922cc6a 100644 --- a/tests/test_factory.py +++ b/tests/test_factory.py @@ -1,16 +1,14 @@ import pytest + from pyaml import PyAMLConfigException, PyAMLException +from pyaml.accelerator import Accelerator from pyaml.configuration.factory import Factory from tests.conftest import MockElement -from pyaml.accelerator import Accelerator def test_factory_build_default(): """Test default PyAML module loading.""" - data = { - "type": "mock_module", - "name": "simple" - } + data = {"type": "mock_module", "name": "simple"} obj = Factory.depth_first_build(data) assert isinstance(obj, MockElement) assert obj.name == "simple" @@ -18,34 +16,37 @@ def test_factory_build_default(): def test_factory_with_custom_strategy(): """Test that custom BuildStrategy overrides default logic.""" - data = { - "type": "mock_module", - "name": "injected", - "custom": True - } + data = {"type": "mock_module", "name": "injected", "custom": True} obj = Factory.depth_first_build(data) assert isinstance(obj, MockElement) assert obj.name == "custom_injected" -@pytest.mark.parametrize("test_file", [ - "tests/config/bad_conf.yml", -]) +@pytest.mark.parametrize( + "test_file", + [ + "tests/config/bad_conf.yml", + ], +) def test_error_location(test_file): with pytest.raises(PyAMLConfigException) as exc: - ml:Accelerator = Accelerator.load(test_file) + ml: Accelerator = Accelerator.load(test_file) print(str(exc.value)) test_file_names = test_file.split("/") - test_file_name = test_file_names[len(test_file_names)-1] + test_file_name = test_file_names[len(test_file_names) - 1] assert f"{test_file_name} at line 7, column 9" in str(exc.value) -@pytest.mark.parametrize("test_file", [ - "tests/config/bad_conf_cycles.yml", - "tests/config/bad_conf_cycles.json", -]) + +@pytest.mark.parametrize( + "test_file", + [ + "tests/config/bad_conf_cycles.yml", + "tests/config/bad_conf_cycles.json", + ], +) def test_error_cycles(test_file): with pytest.raises(PyAMLException) as exc: - ml:Accelerator = Accelerator.load(test_file) + ml: Accelerator = Accelerator.load(test_file) assert "Circular file inclusion of " in str(exc.value) if not test_file.endswith(".json"): diff --git a/tests/test_ident_models.py b/tests/test_ident_models.py index 28e22d77..4281efcb 100644 --- a/tests/test_ident_models.py +++ b/tests/test_ident_models.py @@ -1,28 +1,31 @@ +import numpy as np import pytest +from pyaml.accelerator import Accelerator +from pyaml.configuration.factory import Factory from pyaml.magnet.cfm_magnet import CombinedFunctionMagnet from pyaml.magnet.hcorrector import HCorrector -from pyaml.magnet.vcorrector import VCorrector -from pyaml.configuration.factory import Factory -from pyaml.accelerator import Accelerator from pyaml.magnet.model import MagnetModel -import numpy as np +from pyaml.magnet.vcorrector import VCorrector + @pytest.mark.parametrize( ("magnet_file", "install_test_package"), [ - ("tests/config/sr-ident-cfm.yaml", {"name": "tango-pyaml", "path": "tests/dummy_cs/tango-pyaml"}), + ( + "tests/config/sr-ident-cfm.yaml", + {"name": "tango-pyaml", "path": "tests/dummy_cs/tango-pyaml"}, + ), ], indirect=["install_test_package"], ) def test_cfm_magnets(magnet_file, install_test_package): - - sr:Accelerator = Accelerator.load(magnet_file) + sr: Accelerator = Accelerator.load(magnet_file) sr.design.get_lattice().disable_6d() - #magnet_design = sr.design.get_magnet("SH1A-C01") - #magnet_live = sr.live.get_magnet("SH1A-C01") - #assert isinstance(magnet_design, CombinedFunctionMagnet) - #assert isinstance(magnet_live, CombinedFunctionMagnet) + # magnet_design = sr.design.get_magnet("SH1A-C01") + # magnet_live = sr.live.get_magnet("SH1A-C01") + # assert isinstance(magnet_design, CombinedFunctionMagnet) + # assert isinstance(magnet_live, CombinedFunctionMagnet) magnet_h_design = sr.design.get_magnet("SH1A-C01-H") magnet_v_design = sr.design.get_magnet("SH1A-C01-V") magnet_h_live = sr.live.get_magnet("SH1A-C01-H") @@ -36,10 +39,10 @@ def test_cfm_magnets(magnet_file, install_test_package): magnet_h_design.strength.set(0.000010) magnet_v_design.strength.set(0.000015) - o,_ = sr.design.get_lattice().find_orbit() - assert(np.abs(o[0] + 9.91848416e-05)<1e-10) - assert(np.abs(o[1] + 3.54829761e-07)<1e-10) - assert(np.abs(o[2] + 1.56246320e-06)<1e-10) - assert(np.abs(o[3] + 1.75037311e-05)<1e-10) + o, _ = sr.design.get_lattice().find_orbit() + assert np.abs(o[0] + 9.91848416e-05) < 1e-10 + assert np.abs(o[1] + 3.54829761e-07) < 1e-10 + assert np.abs(o[2] + 1.56246320e-06) < 1e-10 + assert np.abs(o[3] + 1.75037311e-05) < 1e-10 - Factory.clear() \ No newline at end of file + Factory.clear() diff --git a/tests/test_linkers.py b/tests/test_linkers.py index 898354d1..50661434 100644 --- a/tests/test_linkers.py +++ b/tests/test_linkers.py @@ -2,28 +2,31 @@ from pyaml import PyAMLException from pyaml.accelerator import Accelerator - +from pyaml.lattice.attribute_linker import ( + ConfigModel as AttrConfigModel, +) from pyaml.lattice.attribute_linker import ( PyAtAttributeElementsLinker, PyAtAttributeIdentifier, - ConfigModel as AttrConfigModel, ) - # ----------------------- # Dummy PyAML Element # ----------------------- + class DummyPyAMLElement: """Minimal stand-in for a PyAML Element: only provides .name.""" + def __init__(self, name: str): self._name = name + def get_name(self) -> str: return self._name def test_conf_with_linker(): - sr:Accelerator = Accelerator.load("tests/config/sr-attribute-linker.yaml") + sr: Accelerator = Accelerator.load("tests/config/sr-attribute-linker.yaml") assert sr is not None magnet = sr.design.get_magnet("SH1A-C01-H") assert magnet is not None @@ -33,11 +36,13 @@ def test_conf_with_linker(): # PyAtAttributeElementsLinker tests # ----------------------- + def test_attribute_identifier_from_pyaml_name(lattice_with_custom_attr): - # We bind to AT element attribute 'Tag'; identifier value comes from PyAML element .name + """We bind to AT element attribute 'Tag'; + identifier value comes from PyAML element .name""" linker = PyAtAttributeElementsLinker(AttrConfigModel(attribute_name="Tag")) linker.set_lattice(lattice_with_custom_attr) - pyaml_elem = DummyPyAMLElement(name="QF") # identifier="QF" + pyaml_elem = DummyPyAMLElement(name="QF") # identifier="QF" ident = linker.get_element_identifier(pyaml_elem) assert isinstance(ident, PyAtAttributeIdentifier) assert ident.attribute_name == "Tag" diff --git a/tests/test_load_quad.py b/tests/test_load_quad.py index 0f260285..bd2bd8fe 100644 --- a/tests/test_load_quad.py +++ b/tests/test_load_quad.py @@ -1,51 +1,67 @@ -import pytest import json + import numpy as np +import pytest from scipy.constants import speed_of_light from pyaml import PyAMLException -from pyaml.configuration import set_root_folder, get_root_folder +from pyaml.accelerator import Accelerator +from pyaml.configuration import Factory, get_root_folder, set_root_folder from pyaml.configuration.fileloader import load -from pyaml.configuration import Factory +from pyaml.control.abstract_impl import ( + RWHardwareArray, + RWHardwareScalar, + RWStrengthArray, + RWStrengthScalar, +) +from pyaml.magnet.cfm_magnet import CombinedFunctionMagnet from pyaml.magnet.hcorrector import HCorrector -from pyaml.magnet.quadrupole import Quadrupole -from pyaml.magnet.quadrupole import ConfigModel as QuadrupoleConfigModel from pyaml.magnet.identity_model import IdentityMagnetModel -from pyaml.magnet.cfm_magnet import CombinedFunctionMagnet -from pyaml.control.abstract_impl import RWHardwareScalar,RWStrengthScalar,RWHardwareArray,RWStrengthArray -from pyaml.accelerator import Accelerator +from pyaml.magnet.quadrupole import ConfigModel as QuadrupoleConfigModel +from pyaml.magnet.quadrupole import Quadrupole # TODO: Generate JSON pydantic schema for MetaConfigurator -#def test_json(): +# def test_json(): # print(json.dumps(QuadrupoleConfigModel.model_json_schema(),indent=2)) + class DummyPeer: def __init__(self): pass - def name(self)->str: + + def name(self) -> str: return "Dummy" -@pytest.mark.parametrize("install_test_package", [{ - "name": "pyaml_external", - "path": "tests/external" -}], indirect=True) + +@pytest.mark.parametrize( + "install_test_package", + [{"name": "pyaml_external", "path": "tests/external"}], + indirect=True, +) def test_quad_external_model(install_test_package, config_root_dir): set_root_folder(config_root_dir) cfg_hcorr_yaml = load("sr/custom_magnets/hidcorr.yaml") hcorr_with_external_model: HCorrector = Factory.depth_first_build(cfg_hcorr_yaml) strength = RWStrengthScalar(hcorr_with_external_model.model) hardware = RWHardwareScalar(hcorr_with_external_model.model) - ref_corr = hcorr_with_external_model.attach(DummyPeer(),strength,hardware) + ref_corr = hcorr_with_external_model.attach(DummyPeer(), strength, hardware) ref_corr.strength.set(10.0) print(ref_corr.strength.get()) Factory.clear() + @pytest.mark.parametrize( ("magnet_file", "install_test_package"), [ ("sr/quadrupoles/QF1AC01.yaml", None), - ("sr/quadrupoles/QF1AC01-IDENT-STRGTH.yaml", {"name": "tango-pyaml", "path": "tests/dummy_cs/tango-pyaml"}), - ("sr/quadrupoles/QF1AC01-IDENT-HW.yaml", {"name": "tango-pyaml", "path": "tests/dummy_cs/tango-pyaml"}), + ( + "sr/quadrupoles/QF1AC01-IDENT-STRGTH.yaml", + {"name": "tango-pyaml", "path": "tests/dummy_cs/tango-pyaml"}, + ), + ( + "sr/quadrupoles/QF1AC01-IDENT-HW.yaml", + {"name": "tango-pyaml", "path": "tests/dummy_cs/tango-pyaml"}, + ), ("sr/quadrupoles/QF1AC01.json", None), ], indirect=["install_test_package"], @@ -54,40 +70,43 @@ def test_quad_linear(magnet_file, install_test_package, config_root_dir): set_root_folder(config_root_dir) cfg_quad = load(magnet_file) print(f"Current file: {config_root_dir}/{magnet_file}") - quad:Quadrupole = Factory.depth_first_build(cfg_quad) + quad: Quadrupole = Factory.depth_first_build(cfg_quad) hardware = RWHardwareScalar(quad.model) if quad.model.has_hardware() else None strength = RWStrengthScalar(quad.model) if quad.model.has_physics() else None - ref_quad = quad.attach(DummyPeer(),strength,hardware) + ref_quad = quad.attach(DummyPeer(), strength, hardware) ref_quad.model.set_magnet_rigidity(6e9 / speed_of_light) try: ref_quad.strength.set(0.7962) sunit = ref_quad.strength.unit() - assert( sunit == "1/m" ) + assert sunit == "1/m" except Exception as ex: if not quad.model.has_physics(): - assert( "has no model that supports physics units" in str(ex) ) + assert "has no model that supports physics units" in str(ex) ref_quad.hardware.set(80.423276) else: raise ex try: current = ref_quad.hardware.get() - assert( np.abs(current-80.423276) < 1e-4 ) + assert np.abs(current - 80.423276) < 1e-4 hunit = ref_quad.hardware.unit() - assert( hunit == "A" ) + assert hunit == "A" except Exception as ex: if not quad.model.has_hardware(): - assert( "has no model that supports hardware units" in str(ex) ) + assert "has no model that supports hardware units" in str(ex) else: raise ex Factory.clear() -@pytest.mark.parametrize("magnet_file", [ - "sr/correctors/SH1AC01.yaml", -]) +@pytest.mark.parametrize( + "magnet_file", + [ + "sr/correctors/SH1AC01.yaml", + ], +) def test_combined_function_magnets(magnet_file, config_root_dir): set_root_folder(config_root_dir) cfg_sh = load(magnet_file) @@ -97,21 +116,21 @@ def test_combined_function_magnets(magnet_file, config_root_dir): strengths = RWStrengthArray(sh.model) sUnits = sh.model.get_strength_units() hUnits = sh.model.get_hardware_units() - assert( sUnits[0] == 'rad' and sUnits[1] == 'rad' and sUnits[2] == 'm-1' ) - assert( hUnits[0] == 'A' and hUnits[1] == 'A' and hUnits[2] == 'A') - ms = sh.attach(DummyPeer(),strengths,currents) + assert sUnits[0] == "rad" and sUnits[1] == "rad" and sUnits[2] == "m-1" + assert hUnits[0] == "A" and hUnits[1] == "A" and hUnits[2] == "A" + ms = sh.attach(DummyPeer(), strengths, currents) hCorr = ms[1] vCorr = ms[2] sqCorr = ms[3] hCorr.strength.set(0.000020) vCorr.strength.set(-0.000015) sqCorr.strength.set(0.000100) - currents = sh.model.compute_hardware_values([0.000020,-0.000015,0.000100]) - assert( np.abs(currents[0]-0.05913476) < 1e-8 ) - assert( np.abs(currents[1]-0.05132066) < 1e-8 ) - assert( np.abs(currents[2]+0.06253617) < 1e-8 ) + currents = sh.model.compute_hardware_values([0.000020, -0.000015, 0.000100]) + assert np.abs(currents[0] - 0.05913476) < 1e-8 + assert np.abs(currents[1] - 0.05132066) < 1e-8 + assert np.abs(currents[2] + 0.06253617) < 1e-8 str = sh.model.compute_strengths(currents) - assert( np.abs(str[0]-0.000020) < 1e-8 ) - assert( np.abs(str[1]+0.000015) < 1e-8 ) - assert( np.abs(str[2]-0.000100) < 1e-8 ) + assert np.abs(str[0] - 0.000020) < 1e-8 + assert np.abs(str[1] + 0.000015) < 1e-8 + assert np.abs(str[2] - 0.000100) < 1e-8 Factory.clear() diff --git a/tests/test_rf.py b/tests/test_rf.py index 1726f023..9e555703 100644 --- a/tests/test_rf.py +++ b/tests/test_rf.py @@ -1,16 +1,18 @@ -from pyaml.accelerator import Accelerator -from pyaml.configuration.factory import Factory -from pyaml.common.exception import PyAMLException import numpy as np import pytest -@pytest.mark.parametrize("install_test_package", [{ - "name": "tango-pyaml", - "path": "tests/dummy_cs/tango-pyaml" -}], indirect=True) -def test_rf(install_test_package): +from pyaml.accelerator import Accelerator +from pyaml.common.exception import PyAMLException +from pyaml.configuration.factory import Factory + - sr:Accelerator = Accelerator.load("tests/config/EBS_rf.yaml") +@pytest.mark.parametrize( + "install_test_package", + [{"name": "tango-pyaml", "path": "tests/dummy_cs/tango-pyaml"}], + indirect=True, +) +def test_rf(install_test_package): + sr: Accelerator = Accelerator.load("tests/config/EBS_rf.yaml") RF = sr.design.get_rf_plant("RF") RF.frequency.set(3.523e8) @@ -18,34 +20,34 @@ def test_rf(install_test_package): # Check that frequency has been applied on all cavities ring = sr.design.get_lattice() for e in ring: - if(e.FamName.startswith("CAV")): - assert(e.Frequency==3.523e8) + if e.FamName.startswith("CAV"): + assert e.Frequency == 3.523e8 - RF.voltage.set(10.e6) + RF.voltage.set(10.0e6) # Check that voltage has been applied on all cavities ring = sr.design.get_lattice() for e in ring: - if(e.FamName.startswith("CAV")): - assert(np.isclose(e.Voltage,10.e6/13.)) + if e.FamName.startswith("CAV"): + assert np.isclose(e.Voltage, 10.0e6 / 13.0) if False: RF = sr.live.get_rf_plant("RF") - RF.frequency.set(3.523721693993786E8) + RF.frequency.set(3.523721693993786e8) RF.voltage.set(6.5e6) - assert np.isclose(RF.frequency.get(), 3.523721693993786E8) + assert np.isclose(RF.frequency.get(), 3.523721693993786e8) assert np.isclose(RF.voltage.get(), 6.5e6) Factory.clear() -@pytest.mark.parametrize("install_test_package", [{ - "name": "tango-pyaml", - "path": "tests/dummy_cs/tango-pyaml" -}], indirect=True) +@pytest.mark.parametrize( + "install_test_package", + [{"name": "tango-pyaml", "path": "tests/dummy_cs/tango-pyaml"}], + indirect=True, +) def test_rf_multi(install_test_package): - - sr:Accelerator = Accelerator.load("tests/config/EBS_rf_multi.yaml") + sr: Accelerator = Accelerator.load("tests/config/EBS_rf_multi.yaml") # Simulator RF = sr.design.get_rf_plant("RF") @@ -55,24 +57,24 @@ def test_rf_multi(install_test_package): # Check that frequency has been applied on all cavities ring = sr.design.get_lattice() for e in ring: - if(e.FamName.startswith("CAV")): - if( e.FamName == "CAV_C25_03"): + if e.FamName.startswith("CAV"): + if e.FamName == "CAV_C25_03": # Harmonic cavity - assert(np.isclose(e.Frequency,1.4092e9)) + assert np.isclose(e.Frequency, 1.4092e9) else: - assert(e.Frequency==3.523e8) + assert e.Frequency == 3.523e8 RFTRA_HARMONIC = sr.design.get_rf_trasnmitter("RFTRA_HARMONIC") RFTRA_HARMONIC.voltage.set(300e3) RF.voltage.set(12e6) for e in ring: - if(e.FamName.startswith("CAV")): - if( e.FamName == "CAV_C25_03"): + if e.FamName.startswith("CAV"): + if e.FamName == "CAV_C25_03": # Harmonic cavity - assert(np.isclose(e.Voltage,300e3)) + assert np.isclose(e.Voltage, 300e3) else: - assert(np.isclose(e.Voltage,1e6)) + assert np.isclose(e.Voltage, 1e6) # Control system RF = sr.live.get_rf_plant("RF") @@ -84,20 +86,20 @@ def test_rf_multi(install_test_package): RFTRA_HARMONIC.voltage.set(300e3) RF.voltage.set(12e6) - assert(np.isclose(RF1._cfg.voltage.get(),10e6)) - assert(np.isclose(RF2._cfg.voltage.get(),2e6)) - assert(np.isclose(RFTRA_HARMONIC._cfg.voltage.get(),3e5)) + assert np.isclose(RF1._cfg.voltage.get(), 10e6) + assert np.isclose(RF2._cfg.voltage.get(), 2e6) + assert np.isclose(RFTRA_HARMONIC._cfg.voltage.get(), 3e5) Factory.clear() -@pytest.mark.parametrize("install_test_package", [{ - "name": "tango-pyaml", - "path": "tests/dummy_cs/tango-pyaml" -}], indirect=True) +@pytest.mark.parametrize( + "install_test_package", + [{"name": "tango-pyaml", "path": "tests/dummy_cs/tango-pyaml"}], + indirect=True, +) def test_rf_multi_notrans(install_test_package): - - sr:Accelerator = Accelerator.load("tests/config/EBS_rf_notrans.yaml") + sr: Accelerator = Accelerator.load("tests/config/EBS_rf_notrans.yaml") # Simulator RF = sr.design.get_rf_plant("RF") @@ -107,17 +109,17 @@ def test_rf_multi_notrans(install_test_package): ring = sr.design.get_lattice() for e in ring: if e.FamName.startswith("CAV"): - assert(np.isclose(e.Frequency,3.523e8)) - assert(np.isclose(e.Voltage,10e6/13.)) + assert np.isclose(e.Frequency, 3.523e8) + assert np.isclose(e.Voltage, 10e6 / 13.0) # Control system RF = sr.live.get_rf_plant("RF") RF.frequency.set(3.523e8) with pytest.raises(PyAMLException) as exc: RF.voltage.set(10e6) - assert("has no trasmitter device defined" in str(exc)) + assert "has no trasmitter device defined" in str(exc) # Check that frequency and voltage has been applied on the masterclock device - assert(np.isclose(RF.frequency.get(), 3.523e8)) + assert np.isclose(RF.frequency.get(), 3.523e8) Factory.clear() diff --git a/tests/test_tune.py b/tests/test_tune.py index dced51d0..ad104337 100644 --- a/tests/test_tune.py +++ b/tests/test_tune.py @@ -1,15 +1,17 @@ -from pyaml.accelerator import Accelerator -from pyaml.configuration.factory import Factory import numpy as np import pytest -@pytest.mark.parametrize("install_test_package", [{ - "name": "tango-pyaml", - "path": "tests/dummy_cs/tango-pyaml" -}], indirect=True) -def test_tune(install_test_package): +from pyaml.accelerator import Accelerator +from pyaml.configuration.factory import Factory + - sr:Accelerator = Accelerator.load("tests/config/EBSTune.yaml") +@pytest.mark.parametrize( + "install_test_package", + [{"name": "tango-pyaml", "path": "tests/dummy_cs/tango-pyaml"}], + indirect=True, +) +def test_tune(install_test_package): + sr: Accelerator = Accelerator.load("tests/config/EBSTune.yaml") sr.design.get_lattice().disable_6d() quadForTuneDesign = sr.design.get_magnets("QForTune") @@ -18,13 +20,13 @@ def test_tune(install_test_package): # Build tune response matrix tune = tune_monitor.tune.get() print(tune) - tunemat = np.zeros((len(quadForTuneDesign),2)) + tunemat = np.zeros((len(quadForTuneDesign), 2)) - for idx,m in enumerate(quadForTuneDesign): + for idx, m in enumerate(quadForTuneDesign): str = m.strength.get() - m.strength.set(str+1e-4) + m.strength.set(str + 1e-4) dq = tune_monitor.tune.get() - tune - tunemat[idx] = dq*1e4 + tunemat[idx] = dq * 1e4 m.strength.set(str) # Compute correction matrix @@ -32,21 +34,20 @@ def test_tune(install_test_package): # Correct tune strs = quadForTuneDesign.strengths.get() - strs += np.matmul(correctionmat,[0.1,0.05]) # Ask for correction [dqx,dqy] + strs += np.matmul(correctionmat, [0.1, 0.05]) # Ask for correction [dqx,dqy] quadForTuneDesign.strengths.set(strs) newTune = tune_monitor.tune.get() - diffTune = newTune-tune + diffTune = newTune - tune print(diffTune) - assert( np.abs(diffTune[0]-0.1) < 1e-3 ) - assert( np.abs(diffTune[1]-0.05) < 1e-3 ) + assert np.abs(diffTune[0] - 0.1) < 1e-3 + assert np.abs(diffTune[1] - 0.05) < 1e-3 if False: # Correct the tune on live (need a Virutal Accelerator) quadForTuneLive = sr.live.get_magnets("QForTune") strs = quadForTuneLive.strengths.get() - strs += np.matmul(correctionmat,[0.1,0.05]) # Ask for correction [dqx,dqy] + strs += np.matmul(correctionmat, [0.1, 0.05]) # Ask for correction [dqx,dqy] quadForTuneLive.strengths.set(strs) - Factory.clear() diff --git a/tests/test_tune_hardware.py b/tests/test_tune_hardware.py index 3d03d8bf..c0aa4e29 100644 --- a/tests/test_tune_hardware.py +++ b/tests/test_tune_hardware.py @@ -1,32 +1,34 @@ -from pyaml.accelerator import Accelerator -from pyaml.configuration.factory import Factory import numpy as np import pytest -@pytest.mark.parametrize("install_test_package", [{ - "name": "tango-pyaml", - "path": "tests/dummy_cs/tango-pyaml" -}], indirect=True) -def test_tune(install_test_package): +from pyaml.accelerator import Accelerator +from pyaml.configuration.factory import Factory + - sr:Accelerator = Accelerator.load(("tests/config/EBSTune.yaml")) +@pytest.mark.parametrize( + "install_test_package", + [{"name": "tango-pyaml", "path": "tests/dummy_cs/tango-pyaml"}], + indirect=True, +) +def test_tune(install_test_package): + sr: Accelerator = Accelerator.load(("tests/config/EBSTune.yaml")) sr.design.get_lattice().disable_6d() quadForTuneDesign = sr.design.get_magnets("QForTune") # Build tune response matrix (hardware units) - + tune_monitor = sr.design.get_betatron_tune_monitor("BETATRON_TUNE") tune = tune_monitor.tune.get() print(tune) - tunemat = np.zeros((len(quadForTuneDesign),2)) + tunemat = np.zeros((len(quadForTuneDesign), 2)) idx = 0 for m in quadForTuneDesign: current = m.hardware.get() - m.hardware.set(current+1e-6) + m.hardware.set(current + 1e-6) dq = tune_monitor.tune.get() - tune - tunemat[idx] = dq*1e6 + tunemat[idx] = dq * 1e6 m.hardware.set(current) idx += 1 @@ -35,14 +37,14 @@ def test_tune(install_test_package): # Correct tune currents = quadForTuneDesign.hardwares.get() - currents += np.matmul(correctionmat,[0.1,0.05]) # Ask for correction [dqx,dqy] + currents += np.matmul(correctionmat, [0.1, 0.05]) # Ask for correction [dqx,dqy] quadForTuneDesign.hardwares.set(currents) newTune = tune_monitor.tune.get() units = quadForTuneDesign.hardwares.unit() diffTune = newTune - tune - assert( np.abs(diffTune[0]-0.1) < 1e-3 ) - assert( np.abs(diffTune[1]-0.05) < 1.1e-3 ) - assert( np.abs(currents[0]-88.04522942) < 1e-8 ) - assert( np.abs(currents[1]-88.26677735) < 1e-8 ) - assert( units[0] == 'A' and units[1] == 'A' ) + assert np.abs(diffTune[0] - 0.1) < 1e-3 + assert np.abs(diffTune[1] - 0.05) < 1.1e-3 + assert np.abs(currents[0] - 88.04522942) < 1e-8 + assert np.abs(currents[1] - 88.26677735) < 1e-8 + assert units[0] == "A" and units[1] == "A" Factory.clear() diff --git a/tests/test_tune_monitor.py b/tests/test_tune_monitor.py index 094d1ded..0df78bf6 100644 --- a/tests/test_tune_monitor.py +++ b/tests/test_tune_monitor.py @@ -1,14 +1,16 @@ +import pytest + from pyaml.accelerator import Accelerator from pyaml.configuration.factory import Factory -import pytest -@pytest.mark.parametrize("install_test_package", [{ - "name": "tango-pyaml", - "path": "tests/dummy_cs/tango-pyaml" -}], indirect=True) -def test_simulator_tune_monitor(install_test_package): - sr:Accelerator = Accelerator.load("tests/config/tune_monitor.yaml") +@pytest.mark.parametrize( + "install_test_package", + [{"name": "tango-pyaml", "path": "tests/dummy_cs/tango-pyaml"}], + indirect=True, +) +def test_simulator_tune_monitor(install_test_package): + sr: Accelerator = Accelerator.load("tests/config/tune_monitor.yaml") sr.design.get_lattice().disable_6d() tune_monitor = sr.design.get_betatron_tune_monitor("BETATRON_TUNE") assert tune_monitor.tune.get()[0] == sr.design.get_lattice().get_tune()[0] @@ -16,16 +18,16 @@ def test_simulator_tune_monitor(install_test_package): Factory.clear() -@pytest.mark.parametrize("install_test_package", [{ - "name": "tango-pyaml", - "path": "tests/dummy_cs/tango-pyaml" -}], indirect=True) -def test_controlsystem_tune_monitor(install_test_package): - sr:Accelerator = Accelerator.load("tests/config/tune_monitor.yaml") +@pytest.mark.parametrize( + "install_test_package", + [{"name": "tango-pyaml", "path": "tests/dummy_cs/tango-pyaml"}], + indirect=True, +) +def test_controlsystem_tune_monitor(install_test_package): + sr: Accelerator = Accelerator.load("tests/config/tune_monitor.yaml") tune_monitor = sr.live.get_betatron_tune_monitor("BETATRON_TUNE") assert tune_monitor.tune.get()[0] == 0.0 assert tune_monitor.tune.get()[1] == 0.0 Factory.clear() - diff --git a/tests/test_values.py b/tests/test_values.py index 5d548d0f..a5ed3f86 100644 --- a/tests/test_values.py +++ b/tests/test_values.py @@ -1,7 +1,9 @@ import numpy as np import pytest + from pyaml.control.readback_value import Value + class TestBasicValue: """ Test basic usage of the Value class with native types and arithmetic. @@ -9,11 +11,13 @@ class TestBasicValue: def test_value_basic(self): """ - Test basic arithmetic operations between Value instances and scalars or other Value objects. + Test basic arithmetic operations between Value instances + and scalars or other Value objects. This test verifies that: - Addition between Value and int works correctly. - - The result of chained additions is both numerically and semantically equal to an expected Value. + - The result of chained additions is both numerically + and semantically equal to an expected Value. - Division between two Value instances produces the expected float result. - The result of Value / Value is a float and not a Value instance. """ @@ -26,18 +30,24 @@ def test_value_basic(self): assert allez_venez == 6 assert allez_venez == Value(6) - laissez_faire_l_insouciance = Value(allez_venez) / Value(et_entrez_dans_la_danse) - assert isinstance(laissez_faire_l_insouciance,float) + laissez_faire_l_insouciance = Value(allez_venez) / Value( + et_entrez_dans_la_danse + ) + assert isinstance(laissez_faire_l_insouciance, float) assert laissez_faire_l_insouciance == 3 - @pytest.mark.parametrize("scalar,expected", [ - (2, 6.28), - (0.5, 1.57), - (-1, -3.14), - ]) + @pytest.mark.parametrize( + "scalar,expected", + [ + (2, 6.28), + (0.5, 1.57), + (-1, -3.14), + ], + ) def test_value_scalar_multiplication(self, value_scalar, scalar, expected): """ - Test multiplication between Value and scalar using different coefficients. + Test multiplication between Value and + scalar using different coefficients. Parameters ---------- @@ -46,11 +56,11 @@ def test_value_scalar_multiplication(self, value_scalar, scalar, expected): expected : float The expected numeric result. """ - assert value_scalar * scalar == expected - assert scalar * value_scalar == expected + assert value_scalar * scalar == expected + assert scalar * value_scalar == expected -class TestValueNumpy: +class TestValueNumpy: def test_value_with_numpy(self): """ Test interactions between Value objects and NumPy arrays. @@ -58,7 +68,8 @@ def test_value_with_numpy(self): This test ensures that: - Value objects can be multiplied by scalars. - A list of Value objects can be converted into a NumPy array. - - Multiplication between arrays of Value and NumPy vectors works elementwise. + - Multiplication between arrays of Value and NumPy vectors + works elementwise. - The result of operations is unwrapped to native NumPy float types. """ my_list = [Value(3.14) for _ in range(10)] @@ -80,15 +91,13 @@ def test_value_with_numpy(self): assert isinstance(my_array_typed[0], np.float64) assert isinstance(mult_result1[0], float) - - @pytest.mark.parametrize("input_array, expected_squared, expected_scaled", [ - (np.array([10, 20, 30]), - np.array([100, 400, 900]), - np.array([20, 40, 60])), - (np.array([1, 2, 3]), - np.array([1, 4, 9]), - np.array([2, 4, 6])) - ]) + @pytest.mark.parametrize( + "input_array, expected_squared, expected_scaled", + [ + (np.array([10, 20, 30]), np.array([100, 400, 900]), np.array([20, 40, 60])), + (np.array([1, 2, 3]), np.array([1, 4, 9]), np.array([2, 4, 6])), + ], + ) def test_value_with_array(self, input_array, expected_squared, expected_scaled): """ Test arithmetic operations on Value instances containing NumPy arrays. @@ -111,24 +120,25 @@ def test_value_with_array(self, input_array, expected_squared, expected_scaled): two_times = array_value * vector assert np.all(two_times == expected_scaled) - def test_of_array_of_value_of_array(self, value_matrix, broadcast_matrix): """ Test multiplication of a 1D array of Value objects (each wrapping a NumPy array) by a 2D NumPy matrix. This test checks that: - - The resulting matrix has correct shape and contains correctly multiplied values. + - The resulting matrix has correct shape and contains + correctly multiplied values. - Each row in the result corresponds to one of the Value-wrapped arrays, multiplied elementwise by 2. - - The Value instances are preserved in structure and contents are correctly accessed. + - The Value instances are preserved in structure and + contents are correctly accessed. """ result = value_matrix * broadcast_matrix assert result.shape == (3, 3) for row in result: - for computed, original in zip(row, value_matrix): + for computed, original in zip(row, value_matrix, strict=False): assert isinstance(original, Value) twice = original * 2 assert np.all(computed == twice)