Merge pull request #16 from hz-b/dev/feature/lattice-representation

Refacture of data model: separate lattice  elements from integration attributes

Author: Sulimankhail

lat2db/bl/create_machine.py

@@ -0,0 +1,125 @@
+import logging
+from typing import Dict, Callable, Union
+
+from ..model.geometric_info import GeometricInfo
+from ..model.lattice_elements.beam_position_monitor import BeamPositionMonitor
+from ..model.lattice_elements.bending import Bending
+from ..model.lattice_elements.cavity import Cavity
+from ..model.lattice_elements.drift import Drift
+from ..model.lattice_elements.element import Element, ElementTypeNames
+from ..model.lattice_elements.marker import Marker
+from ..model.lattice_elements.quadrupole import Quadrupole
+from ..model.lattice_elements.sextupole import Sextupole
+from ..model.lattice_elements.steerer import Steerer
+from ..model.machine import (
+    extact_version_from_lattice,
+    extract_physics_info_from_lattice,
+    Machine,
+)
+
+
+logger = logging.getLogger("lat2db")
+
+_elm_t_nams = ElementTypeNames
+
+factory_dict_default = {
+    _elm_t_nams.marker.value: Marker,
+    _elm_t_nams.bpm.value: BeamPositionMonitor,
+    _elm_t_nams.drift.value: Drift,
+    _elm_t_nams.bending.value: Bending,
+    _elm_t_nams.quadrupole.value: Quadrupole,
+    _elm_t_nams.sextupole.value: Sextupole,
+    _elm_t_nams.steerer.value: Steerer,
+    "Horizontalsteerer": None,
+    "Verticalsteerer": None,
+    _elm_t_nams.cavity.value: Cavity,
+}
+
+
+def create_machine(lat, factory_dict: Dict[str, Callable] = None) -> Machine:
+    factory_dict = factory_dict or factory_dict_default
+    elms = [
+        process_element(standardise_element_info(row_), factory_dict=factory_dict)
+        for row_ in lat.elements
+    ]
+    # Todo: how to hande elements we are not processing further ...
+    # shall one add a place holder
+    elms = [elm for elm in elms if elm is not None]
+    machine = Machine(
+        sequences=[elms],
+        version=extact_version_from_lattice(
+            lat.lattice_standard_metadata.lattice_version
+        ),
+        phyics_info=extract_physics_info_from_lattice(lat.properties.physics.energy),
+        geometric_info=GeometricInfo(is_ring=lat.properties.geometric.is_ring),
+        name=lat.lattice_standard_metadata.machine_name,
+        closed=lat.lattice_standard_metadata.closed,
+    )
+    return machine
+
+
+def process_element(elem_info: Dict, *, factory_dict) -> Element:
+    type_class = factory_dict[elem_info["type"]]
+    if not type_class:
+        return None
+    try:
+        r = type_class(**elem_info)
+    except Exception as exc:
+        pass
+        raise exc
+    return r
+
+
+def standardise_element_info(
+    elem_info: Dict[str, Union[str, int, float]], *, copy: bool = True
+) -> Dict[str, Union[str, int, float, Dict[str, int]]]:
+
+    if copy:
+        res = elem_info.copy()
+    logger.debug("elem_info= %s", res)
+
+    # revamp parameters as required for the dataclasses
+    res["length"] = res.pop("L", 0e0)
+
+    type_name = res["type"]
+    # address main strength
+    if type_name in [_elm_t_nams.quadrupole.value, _elm_t_nams.sextupole.value]:
+        if "K" in res.keys():
+            res["main_multipole_strength"] = res.pop("K")
+    elif type_name == _elm_t_nams.bending.value:
+        if "K" in res.keys():
+            res["quadrupole_strength"] = res.pop("K")
+    elif type_name == _elm_t_nams.steerer.value:
+        if "K" in res.keys():
+            raise AssertionError("Did not expect K in steerer data")
+
+    if type_name in [
+        _elm_t_nams.bending.value,
+        _elm_t_nams.quadrupole.value,
+        _elm_t_nams.sextupole.value,
+        _elm_t_nams.steerer.value,
+    ]:
+        # Integration info now in a separate dataclass
+        if "N" in res.keys():
+            res["integration_parameters"] = dict(
+                n_slices=int(res.pop("N")),
+                symplectic_order=int(res.pop("Method")),
+            )
+
+    # renaming for specific parts
+    if type_name == _elm_t_nams.bending.value:
+        res["bending_angle"] = res.pop("T", 0e0)
+        res["entry_angle"] = res.pop("T1", 0e0)
+        res["exit_angle"] = res.pop("T2", 0e0)
+
+    if type_name == _elm_t_nams.cavity.value:
+        res["frequency"] = res.pop("Frequency", 0e0)
+        res["voltage"] = res.pop("Voltage", 0e0)  # rename "Voltage" to "voltage"
+        res["harmonic_number"] = res.pop("Harmonicnumber", 0e0)
+    else:
+        pass  # do nothing if type_name is not recognized
+
+    return res
+
+
+__all__ = ["create_machine"]

lat2db/bl/set_machine.py

@@ -1,99 +1,27 @@
-import copy
-import logging
 import os
 
 import jsons
+from starlette.testclient import TestClient
+
 from fastapi import FastAPI
 from pymongo import MongoClient
 
-from lat2db import mongodb_url
-from lat2db.controller import machine_controller
-from lat2db.model.beam_position_monitor import BeamPositionMonitor
-from lat2db.model.bending import Bending
-from lat2db.model.cavity import Cavity
-from lat2db.model.drift import Drift
-from lat2db.model.machine import Machine
-from lat2db.model.marker import Marker
-from lat2db.model.quadrupole import Quadrupole
-from lat2db.model.sequencer import Sequencer
-from lat2db.model.sextupole import Sextupole
-from lat2db.model.steerer import Steerer
+from .. import mongodb_url
+from ..model.machine import Machine
+from ..controller import machine_controller
+
 
 app = FastAPI()
 app.include_router(machine_controller.router, tags=["machines"], prefix="/machine")
 app.mongodb_client = MongoClient(mongodb_url)
 DB_NAME = os.environ.get("MONGODB_DB", "bessyii")
 app.database = app.mongodb_client[DB_NAME]
 
-logger = logging.getLogger("tools")
-
-
-def create_machine(lat):
-    machine = Machine()
-    machine.set_base_parameters( lat )
-
-    # iterate through each row in lat.elements
-    for row_ in lat.elements:
-
-        # make a copy of the row so that changes don't affect original data
-        row = copy.copy(row_)
-
-        # print the row
-        print(f'{row} ----')
+def set_machine(machine: Machine):
+# return machine
 
-        # revamp parameters as required for the dataclasses
-        row.setdefault("length", row.pop("L", 0e0))  # rename "L" to "length"
-        type_name = row['type']
-        # Ensure 'passmethod' and 'tags' keys exist in the row dictionary
-        row.setdefault("passmethod", None)
-        row.setdefault("tags", None)
-        if type_name in ["Bending", "Quadrupole", "Sextupole", "Steerer"]:
-            row.setdefault("main_multipole_strength", row.pop("K", 0e0))  # rename "K" to "main_multipole_strength"
-            row.setdefault("number_of_integration_steps",
-                           row.pop("N", 1))  # rename "N" to "number_of_integration_steps"
-            row.setdefault("method", row.pop("Method", 4))  # rename "Method" to "method"
-        elif type_name == "Cavity":
-            row.setdefault("frequency", row.pop("Frequency", 0e0))  # rename "Frequency" to "frequency"
-            row.setdefault("voltage", row.pop("Voltage", 0e0))  # rename "Voltage" to "voltage"
-            row.setdefault("harmonic_number",
-                           row.pop("Harmonicnumber", 0e0))  # rename "HarmonicNumber" to "harmonic_number"
-        else:
-            pass  # do nothing if type_name is not recognized
-
-        if type_name == "Bending":
-            row.setdefault("bending_angle", row.pop("T", 0e0))  # rename "T" to "bending_angle"
-            row.setdefault("entry_angle", row.pop("T1", 0e0))  # rename "T1" to "entry_angle"
-            row.setdefault("exit_angle", row.pop("T2", 0e0))  # rename "T2" to "exit_angle"
-
-        # create a Sequencer instance using the modified row
-        sequence_item = Sequencer(**row)
-
-        # add the sequence_item to the appropriate list in the machine instance based on its type_name
-        type_dict = {
-            "Drift": (Drift, machine.add_drift),
-            "Marker": (Marker, machine.add_marker),
-            "Sextupole": (Sextupole, machine.add_sextupole),
-            "Steerer": (Steerer, machine.add_steerer()),
-            "Bending": (Bending, machine.add_bending),
-            "Quadrupole": (Quadrupole, machine.add_quadrupole),
-            "Bpm": (BeamPositionMonitor, machine.add_beam_position_monitor),
-            "Cavity": (Cavity, machine.add_cavity),
-        }
-
-        type_class, type_method = type_dict.get(type_name, (None, None))
-        if type_class is None or type_method is None:
-            if type_name in ['Horizontalsteerer','Verticalsteerer']:
-                continue #ignore the two steerers
-            else:
-                raise KeyError(f"Don't know type {type_name}")
-        sequence_item.set_properties(row)
-        type_instance = type_class(**row)
-        machine.add_to_sequence(sequence_item)
-        type_method(type_instance)
-
-    # return machine
-    from starlette.testclient import TestClient
     with TestClient(app) as client:
-        response = client.post("/machine/machine", json=jsons.dump(machine))
+        response = client.post("/machine/machine", json=jsons.dump(machine.to_dict()))
         if response.status_code != 201:
             raise AssertionError(f"Got response {response}")
+

lat2db/controller/machine_controller.py

@@ -161,7 +161,7 @@ def update_quadrupole_details(id: str, quad_name: str, request_body: Quadrupole_
                 if "sequences" in machine:
                     sequences_list = machine.get("sequences", [])
                     for item_index, item in enumerate(sequences_list):
-                        if item.get("name") == request_body.updated_data.name and item.get("type") == "Quadrupole":
+                        if item.get("name") == request_body.updated_data.name and item.get("type") == ElmTNams.quadrupole:
                             removed_quadrupole = sequences_list.pop(item_index)
                             # affected_drift=item.get("index")
                             print("******* affected drif index is ", affected_drift)
@@ -297,7 +297,7 @@ def update_quadrupole_details_copy(id: str, quad_name: str, request_body: Quadru
                 if "sequences" in machine_copy:
                     sequences_list = machine_copy.get("sequences", [])
                     for item_index, item in enumerate(sequences_list):
-                        if item.get("name") == request_body.updated_data.name and item.get("type") == "Quadrupole":
+                        if item.get("name") == request_body.updated_data.name and item.get("type") == ElmTNams.quadrupole:
                             removed_quadrupole = sequences_list.pop(item_index)
                             # affected_drift=item.get("index")
                             print("******* affected drif index is ", affected_drift)
@@ -809,7 +809,7 @@ def update_quadrupole_from_sequence(id: str, target_drift: str, quad_name: str,
                 print("inside the drift")
                 try:
                     get_qud_prev = database.find_one(
-                        {"id": str(id), "sequences": {"$elemMatch": {"name": quad_name, "type": "Quadrupole"}}},
+                        {"id": str(id), "sequences": {"$elemMatch": {"name": quad_name, "type": ElmTNams.quadrupole}}},
                         projection={"sequences.$": 1}
                     )
 

lat2db/model/element.py

@@ -10,7 +10,7 @@ class Energy:
     #: energy or particle_energy
     name: str
     #: specify the unit
-    #: eV to be consitent with the beam energy
+    #: eV to be consistent with the beam energy
     #: deviates from SI, but that would be an acceptable compromise to
-    #: the comumuity
+    #: the community
     value: float

lat2db/model/energy.py

@@ -0,0 +1,5 @@
+from pydantic import BaseModel
+
+
+class ATSpecificInfo(BaseModel):
+    passmethod: str

lat2db/model/integrator_configuration/__init__.py

@@ -0,0 +1,6 @@
+from pydantic import BaseModel
+
+
+class IntegrationParameters(BaseModel):
+    n_slices:  int
+    symplectic_order : int

lat2db/model/integrator_configuration/at_specific_info.py

@@ -1,4 +1,4 @@
-from lat2db.model.element import Element
+from .element import Element
 
 
 class BeamPositionMonitor(Element):