allow to specify use (=mat_key) for loading the lattice from the .mat file

Author: kparasch

pySC/configuration/generation.py

@@ -37,8 +37,13 @@ def generate_SC(yaml_filepath: str, seed: int = 1, scale_errors: Optional[int] =
         else:
             no_6d = False
 
+        if 'use' in config_dict['lattice']:
+            use = config_dict['lattice']['use']
+        else:
+            use = 'RING'
+
         print(f'Loading AT lattice from {lattice_file}')
-        lattice = ATLattice(lattice_file=lattice_file, no_6d=no_6d)
+        lattice = ATLattice(lattice_file=lattice_file, no_6d=no_6d, use=use)
     else:
         raise NotImplementedError(f"Simulator {config_dict['lattice']['simulator']} is not implemented.")
 

pySC/core/lattice.py

@@ -51,13 +51,14 @@ class ATLattice(Lattice):
     # fake field so that pydantic can distinguish
     # the different machine types
     at_simulator: None = None
+    use: str = 'RING'
     _ring: at.Lattice = PrivateAttr(default=None)
     _design: at.Lattice = PrivateAttr(default=None)
 
     @model_validator(mode="after")
     def load_lattice(self):
-        self._ring = at.load_mat(self.lattice_file)
-        self._design = at.load_mat(self.lattice_file)
+        self._ring = at.load_mat(self.lattice_file, use=self.use)
+        self._design = at.load_mat(self.lattice_file, use=self.use)
 
         if not self.no_6d:
             self._ring.enable_6d()