Update documentation and improve code clarity; disable display_version in conf.py, correct module reference in index.rst, enhance example_fake_tbt fixture in conftest.py, refactor MAD-NG and xtrack_line modules for better error handling and type hints.

Author: jgray-19

doc/conf.py

@@ -149,7 +149,7 @@ def about_package(init_posixpath: pathlib.Path) -> dict:
 #
 html_theme_options = {
     'collapse_navigation': False,
-    'display_version': True,
+    # 'display_version': True, # I get a warning about this, so I disable it 
     'logo_only': True,
     'navigation_depth': 1,
 }

doc/readers/index.rst

@@ -36,6 +36,6 @@
     :members:
     :noindex:
 
-.. automodule:: turn_by_turn.xtrack
+.. automodule:: turn_by_turn.xtrack_line
     :members:
     :noindex:

tests/conftest.py

@@ -6,8 +6,13 @@
 @pytest.fixture(scope="session")
 def example_fake_tbt():
     """
-    Returns a TbtData object with the original simulation data.
-    This fixture is session-scoped and will only be created once per test run.
+    Returns a TbtData object using simulation data taken from MAD-NG. 
+    This data is also used for the tests in xtrack, so change the numbers
+    at your own risk.
+
+    It is possible to run the MAD-NG in the inputs folder to regenerate the data.
+    Also, xtrack produces the same data, so you can use the xtrack test fixture 
+    `example_line`.
     """
     names = np.array(["BPM1", "BPM3", "BPM2"])
     # First BPM

turn_by_turn/io.py

@@ -29,12 +29,12 @@
 from turn_by_turn.structures import TbtData
 from turn_by_turn.utils import add_noise_to_tbt
 
+LOGGER = logging.getLogger(__name__)
+
 if TYPE_CHECKING:
     from pandas import DataFrame
     from xtrack import Line
 
-LOGGER = logging.getLogger(__name__)
-
 TBT_MODULES = dict(
     lhc=lhc,
     doros=doros,
@@ -57,42 +57,7 @@
 WRITERS = ("lhc", "sps", "doros", "doros_positions", "doros_oscillations", "ascii", "madng")
 
 write_lhc_ascii = write_ascii  # Backwards compatibility <0.4
-
-def load_tbt_data(
-    tbt_input: str | Path | Line | DataFrame, 
-    datatype: str = "lhc"
-) -> TbtData:
-    """
-    Get a TbtData object from various input types. Explicitly does not infer the datatype from the input. 
-
-    Args:
-        tbt_input (str | Path | Line | DataFrame): 
-            The input data object or path to a file.
-        datatype (str): 
-            Defaults to "lhc".
-
-    Returns:
-        TbtData: The resulting TbtData object.
-    
-    Raises:
-        DataTypeError: If the datatype is None and the input type cannot be inferred.
-        TypeError: If the input type is not supported for the given datatype.
-    """
-    datatype = datatype.lower()
-
-    try:
-        module = TBT_MODULES[datatype]
-    except KeyError as error:
-        LOGGER.exception(
-            f"Unsupported datatype '{datatype}'. Supported types: {list(TBT_MODULES.keys())}"
-        )
-        raise DataTypeError(datatype) from error
-
-    if isinstance(tbt_input, (str, Path)):
-        return module.read_tbt(Path(tbt_input), **additional_args(datatype))
-    return module.convert_to_tbt(tbt_input)
-
-        
+     
 def read_tbt(file_path: str | Path, datatype: str = "lhc") -> TbtData:
     """
     Calls the appropriate loader for the provided matrices type and returns a ``TbtData`` object of the

turn_by_turn/madng.py

@@ -10,7 +10,8 @@
 allowing easy post-processing and conversion between formats.
 
 Dependencies:
-    - Requires the ``tfs-pandas >= 4.0.0`` package to read or write TFS files.
+    - Requires the ``tfs-pandas >= 4.0.0`` package for compatibility with MAD-NG 
+    features. Earlier versions does not support MAD-NG TFS files.
 """
 
 
@@ -23,14 +24,9 @@
 import pandas as pd
 
 from typing import TYPE_CHECKING
-
-try:
-    import tfs
 
-    if TYPE_CHECKING:
-        from tfs import TfsDataFrame
-except ImportError:
-    tfs = None
+if TYPE_CHECKING:
+    import tfs
 
 from turn_by_turn.structures import TbtData, TransverseData
 
@@ -66,17 +62,22 @@ def read_tbt(file_path: str | Path) -> TbtData:
     Raises:
         ImportError: If the ``tfs-pandas`` package is not installed.
     """
-    if not tfs:
-        raise ImportError(
+    try:
+        import tfs
+    except ImportError as e:
+        LOGGER.exception(
             "The 'tfs' package is required to read MAD-NG TFS files. Install it with: pip install 'turn_by_turn[madng]'"
         )
+        raise ImportError(
+            "The 'tfs' package is required to read MAD-NG TFS files. Install it with: pip install 'turn_by_turn[madng]'"
+        ) from e
 
     LOGGER.debug("Starting to read TBT data from dataframe")
     df = tfs.read(file_path)
     return convert_to_tbt(df)
 
 
-def convert_to_tbt(df: pd.DataFrame | 'TfsDataFrame') -> TbtData:
+def convert_to_tbt(df: pd.DataFrame | 'tfs.TfsDataFrame') -> TbtData:
     """
     Convert a TFS or pandas DataFrame to a ``TbtData`` object.
 
@@ -97,16 +98,21 @@ def convert_to_tbt(df: pd.DataFrame | 'TfsDataFrame') -> TbtData:
     """
 
     # Get the date and time from the headers (return None if not found)
-    if tfs and isinstance(df, tfs.TfsDataFrame):
+    try:
+        import tfs
+        is_tfs_df = isinstance(df, tfs.TfsDataFrame)
+    except ImportError:
+        LOGGER.debug(
+            "The 'tfs' package is not installed. Assuming a pandas DataFrame."
+        )
+        is_tfs_df = False
+
+    if is_tfs_df:
         date_str = df.headers.get(DATE)
         time_str = df.headers.get(TIME)
-    elif isinstance(df, pd.DataFrame):
+    else:
         date_str = df.attrs.get(DATE)
         time_str = df.attrs.get(TIME)
-    else:
-        raise TypeError(
-            "Input DataFrame must be a Pandas DataFrame or a TFS DataFrame."
-        )
 
     # Combine the date and time into a datetime object
     date = None

turn_by_turn/xtrack_line.py

@@ -6,23 +6,35 @@
 into the standardized ``TbtData`` format used by ``turn_by_turn``.
 
 Prerequisites for using ``convert_to_tbt``:
+
   1. The input ``Line`` must contain one or more ``ParticlesMonitor`` elements
      positioned at each location where turn-by-turn data is required (e.g., all BPMs).
+
      A valid monitor setup involves:
-       • Placing a ``xt.ParticlesMonitor`` instance in the line's element sequence
+
+       - Placing a ``xt.ParticlesMonitor`` instance in the line's element sequence
          at all the places you would like to observe.
-       • Configuring each monitor with identical settings:
-           - ``start_at_turn`` (first turn to record, usually 0)
-           - ``stop_at_turn`` (The total number of turns to record, e.g., 100)
-           - ``num_particles`` (number of tracked particles)
-     If any monitor is configured with different parameters, ``convert_to_tbt`` 
-     will either find no data or raise a inconsistency error.
+       - Configuring each monitor with identical settings:
+
+           * ``start_at_turn`` (first turn to record, usually 0)
+           * ``stop_at_turn`` (The total number of turns to record, e.g., 100)
+           * ``num_particles`` (number of tracked particles)
+
+     If any monitor is configured with different parameters, ``convert_to_tbt``
+     will either find no data or raise an inconsistency error.
+     
+     Also, if you specify more turns than were actually tracked, the resulting 
+     TBT data will include all turns up to the monitor's configured limit. 
+     This may result in extra rows filled with zeros for turns where no real 
+     data was recorded, which might not be desirable for your analysis.
+
   2. Before conversion, you must:
-       • Build particles with the desired initial coordinates
+
+       - Build particles with the desired initial coordinates
          (using ``line.build_particles(...)``).
-       • Track those particles through the line for the intended number of turns
+       - Track those particles through the line for the intended number of turns
          (using ``line.track(..., num_turns=num_turns)``).
-  
+
 Once these conditions are met, pass the tracked ``Line`` to ``convert_to_tbt`` to
 extract the data from each particle monitor into a ``TbtData`` object.
 """
@@ -39,18 +51,13 @@
 
 from turn_by_turn.structures import TbtData, TransverseData
 
-try:
+if TYPE_CHECKING:
     import xtrack as xt
-    if TYPE_CHECKING:
-        from xtrack import Line
-
-except ImportError:
-    xt = None
 
 LOGGER = logging.getLogger(__name__)
 
 
-def convert_to_tbt(xline: Line) -> TbtData:
+def convert_to_tbt(xline: xt.Line) -> TbtData:
     """
     Convert tracking results from an ``xtrack`` Line into a ``TbtData`` object.
 
@@ -70,10 +77,13 @@ def convert_to_tbt(xline: Line) -> TbtData:
         TypeError: If the input is not a valid ``xtrack.Line``.
         ValueError: If no monitors are found or data is inconsistent.
     """
-    if not xt:
+    try:
+        import xtrack as xt
+    except ImportError as e:
         raise ImportError(
             "The 'xtrack' package is required to convert xtrack Line objects. Install it with: pip install 'turn_by_turn[xtrack]'"
-        )
+        ) from e
+    
     if not isinstance(xline, xt.Line):
         raise TypeError(f"Expected an xtrack Line object, got {type(xline)} instead.")
 
@@ -106,7 +116,7 @@ def convert_to_tbt(xline: Line) -> TbtData:
         )
     npart = npart_set.pop()
 
-    # Precompute masks for each monitor and particle_id using vectorized broadcasting for speed
+    # Precompute masks for each monitor and particle_id
     monitor_particle_masks = [
         mon.data.particle_id[:, None] == np.arange(npart)[None, :]
         for mon in monitors