Revert "Refacture of data model: separate lattice elements from integration attributes"

Author: Sulimankhail

lat2db/bl/create_machine.py

@@ -1,27 +1,99 @@
+import copy
+import logging
 import os
 
 import jsons
-from starlette.testclient import TestClient
-
 from fastapi import FastAPI
 from pymongo import MongoClient
 
-from .. import mongodb_url
-from ..model.machine import Machine
-from ..controller import machine_controller
-
+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
 
 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]
 
-def set_machine(machine: Machine):
-# return machine
+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} ----')
 
+        # 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.to_dict()))
+        response = client.post("/machine/machine", json=jsons.dump(machine))
         if response.status_code != 201:
             raise AssertionError(f"Got response {response}")
-

lat2db/bl/set_machine.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") == ElmTNams.quadrupole:
+                        if item.get("name") == request_body.updated_data.name and item.get("type") == "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") == ElmTNams.quadrupole:
+                        if item.get("name") == request_body.updated_data.name and item.get("type") == "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": ElmTNams.quadrupole}}},
+                        {"id": str(id), "sequences": {"$elemMatch": {"name": quad_name, "type": "Quadrupole"}}},
                         projection={"sequences.$": 1}
                     )
 

lat2db/controller/machine_controller.py

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

lat2db/model/lattice_elements/beam_position_monitor.py renamed to lat2db/model/beam_position_monitor.py

@@ -2,13 +2,12 @@
 
 from .element import Element
 from .magnetic_element import MagnetAssembly
-from ..integrator_configuration.integration_parameters import IntegrationParameters
 
 
 class Bending(Element):
-    bending_angle: float
-    integration_parameters: IntegrationParameters
     element_configuration: Optional[MagnetAssembly] = None
+    number_of_integration_steps: Optional[int] = None
+    bending_angle: Optional[float] = 0.0
     entranceangle: Optional[float] = 0.0
     exitangle: Optional[float] = 0
     fringeint1: Optional[float] = None

lat2db/model/lattice_elements/bending.py renamed to lat2db/model/bending.py

@@ -2,7 +2,7 @@
 
 from pydantic.dataclasses import dataclass
 
-from .element import Element
+from lat2db.model.element import Element
 
 
 @dataclass

lat2db/model/lattice_elements/cavity.py renamed to lat2db/model/cavity.py

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

lat2db/model/lattice_elements/drift.py renamed to lat2db/model/drift.py

@@ -0,0 +1,23 @@
+from typing import Sequence, Optional
+
+from pydantic import BaseModel
+
+
+class Element(BaseModel):
+    #: actually an identifier
+    name: Optional[str]
+    index: Optional[int]
+    #: in meter
+    length: Optional[float]
+    #: to be interpreted by the factories building a lattice from the database
+    type: Optional[str]
+    #: everythng describing the element itself beyond to the information given
+    #: in the fields above
+    # element_configuration: object
+    # tags are given by specific lattice developers to later sort/filter/categorised/group or wild card search their
+    # elements
+    passmethod: Optional[str]
+    tags: Optional[Sequence[str]]
+
+
+__all__ = ["Element"]

lat2db/model/element.py

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