Replace ap.load_lattice with ap.Lattice.from_file (#58)

Use `@classmethod` to create another class constructor for `ap.Lattice` instead of `ap.load_lattice`

Author: felix-andreas

apace/__init__.py

@@ -9,7 +9,6 @@
     Octupole,
     Lattice,
 )
-from .io import load_lattice, save_lattice
 from .twiss import Twiss
 from .matrix_tracking import MatrixTracking
 from .distributions import create_particle_distribution

apace/classes.py

@@ -1,4 +1,5 @@
 import inspect
+import latticejson
 import sys
 from typing import List, Dict, Set, Union
 from .utils import Signal, Attribute, AmbiguousNameError
@@ -388,7 +389,18 @@ def _tree_as_string(obj, prefix=""):
         return string
 
     @classmethod
-    def from_dict(cls, data):
+    def from_file(cls, location, file_format=None) -> "Lattice":
+        """Creates a new `Lattice` from file at `location` (path or url).
+        :param location: path-like or url-like string which locates the lattice file
+        :type location: Union[AnyStr, Path]
+        :param file_format str: File format of the lattice file
+        :type file_format: str, optional (use file extension)
+        :rtype Lattice
+        """
+        return cls.from_dict(latticejson.load(location, file_format))
+
+    @classmethod
+    def from_dict(cls, data) -> "Lattice":
         """Creates a new `Lattice` object from a latticeJSON compliant dictionary."""
 
         objects = {}  # dict containing all elements + lattices

apace/io.py

@@ -209,11 +209,13 @@ def _twiss_plot_section(
     if y_max is None:
         y_max = ax.get_ylim()[1]
 
-    draw_lattice(twiss.lattice, ax, x_min, x_max, annotate_elements=annotate_elements)
     ax.set_xlim((x_min, x_max))
     ax.set_ylim((y_min, y_max))
+    draw_lattice(twiss.lattice, ax, x_min, x_max, annotate_elements=annotate_elements)
 
 
+# TODO: make sub_class of figure
+# add attribute which defines which twiss parameters are plotted
 def twiss_plot(
     twiss,
     main=True,
@@ -278,7 +280,7 @@ def twiss_plot(
                 n_x_ticks=None,
             )
 
-    fig.suptitle(twiss.lattice.name)
+    fig.suptitle(twiss.lattice.name, ha="right", x=0.9925)
     fig.tight_layout()
     # fig.subplots_adjust(top=0.93)
     if path:

apace/plot.py

@@ -178,7 +178,7 @@ def update_twiss_array(self):
         if self._twiss_array.shape[0] != n_points:
             self._twiss_array = np.empty((8, n_points))
 
-        if not self.stable:
+        if not self.stable:  # TODO: replace with warning
             print(
                 f"Horizontal plane stability: {self.stable_x}\n"
                 f"Vertical plane stability: {self.stable_y}"

apace/twiss.py

@@ -209,9 +209,9 @@ be altered when needed.
 
 Load and Save Lattice Files
 ---------------------------
-lattices can also be imported from a lattice file. This can be done using the :func:`load_lattice` function::
+lattices can also be imported from a lattice file. This can be done using the :func:`Lattice.from_file` method::
 
-   lattice = ap.load_lattice('/path/to/file')
+   lattice = ap.Lattice.from_file('/path/to/file')
 
 Individual elements and sub-lattices can be accessed through the :attr:`~Lattice.elements` and :attr:`~Lattice.sub_lattices`, respectively::
 
@@ -220,7 +220,7 @@ Individual elements and sub-lattices can be accessed through the :attr:`~Lattice
 
 A given lattice can be saved to a lattice file using the :func:`save_lattice` function::
 
-   ap.load_lattice(lattice, '/path/to/file')
+   ap.Lattice.from_file(lattice, '/path/to/file')
 
 The Twiss class
 ===============

docs/guide.rst

@@ -19,7 +19,7 @@
     packages=find_packages(),
     package_data={"data": ["data"],},
     install_requires=[
-        "latticejson>=0.0.3",
+        "latticejson>=0.0.4",
         "numpy",
         "scipy",
         "matplotlib",

setup.py

@@ -4,7 +4,7 @@
 
 BASE_PATH = os.path.dirname(__file__)
 FILE_PATH = os.path.join(BASE_PATH, "data", "lattices", "fodo_ring.json")
-FODO_RING = ap.load_lattice(FILE_PATH)
+FODO_RING = ap.Lattice.from_file(FILE_PATH)
 
 
 def test_length_changed():

tests/test_classes.py

@@ -9,7 +9,7 @@
 # TODO: fix linter/mypy erros
 # https://stackoverflow.com/questions/59187502/mypy-dict-containing-instances-of-different-subclasses
 def test_attributes():
-    fodo = ap.load_lattice(FILE_PATH)
+    fodo = ap.Lattice.from_file(FILE_PATH)
     d1 = fodo["D1"]
     assert isinstance(d1, ap.Drift)
     assert 0.55 == fodo["D1"].length
@@ -39,11 +39,9 @@ def test_attributes():
     assert 48 == fodo.length
 
 
-def test_save_lattice():
-    fodo = ap.load_lattice(FILE_PATH)
-    path = "/tmp/tmp_lattice.json"
-    ap.save_lattice(fodo, path)
-    fodo_reload = ap.load_lattice(path)
+def test_serialize_lattice():
+    fodo = ap.Lattice.from_file(FILE_PATH)
+    fodo_reload = ap.Lattice.from_dict(fodo.as_dict())
     assert fodo.length == fodo_reload.length
     assert len(fodo.elements) == len(fodo_reload.elements)
     assert len(fodo.sub_lattices) == len(fodo_reload.sub_lattices)

tests/test_io.py

@@ -7,7 +7,7 @@
 
 
 def test_angle_k1():
-    lattice = ap.load_lattice(FILE_PATH)
+    lattice = ap.Lattice.from_file(FILE_PATH)
     b1 = lattice["B1"]
     q1 = lattice["Q1"]
     matrix_method = ap.MatrixMethod(lattice)