Test: NumPy Multipole Serialization

Author: ax3l

pyproject.toml

@@ -41,7 +41,7 @@ classifiers = [
 ]
 
 [project.optional-dependencies]
-test = ["pytest"]
+test = ["pytest", "numpy"]
 
 [project.urls]
 Documentation = "https://pals-project.readthedocs.io"

tests/test_parameters.py

@@ -134,3 +134,53 @@ def test_ParameterClasses():
     # Test BeamBeamParameters
     beambeam = BeamBeamParameters()
     assert beambeam is not None
+
+
+def test_multipole_numpy_coercion():
+    """Regression test for issue #67: numpy scalars passed to multipole parameter
+    classes must be coerced to Python-native numeric types at construction time,
+    so YAML/JSON serialization produces clean output regardless of input type."""
+    np = pytest.importorskip("numpy")
+
+    # MagneticMultipoleParameters: cover all prefixes and several numpy dtypes
+    mmp = MagneticMultipoleParameters(
+        tilt1=np.float64(0.1),
+        Bn1=np.float64(1.5),
+        Bn2=np.float32(2.5),
+        Bs1=np.int64(3),
+        Kn0=np.int32(-1),
+        Ks1=np.float64(0.25),
+    )
+    assert type(mmp.tilt1) is float and mmp.tilt1 == 0.1
+    assert type(mmp.Bn1) is float and mmp.Bn1 == 1.5
+    assert type(mmp.Bn2) is float and mmp.Bn2 == 2.5
+    assert type(mmp.Bs1) is int and mmp.Bs1 == 3
+    assert type(mmp.Kn0) is int and mmp.Kn0 == -1
+    assert type(mmp.Ks1) is float and mmp.Ks1 == 0.25
+
+    # 0-d numpy array also works
+    mmp_arr = MagneticMultipoleParameters(Bn1=np.array(4.2))
+    assert type(mmp_arr.Bn1) is float and mmp_arr.Bn1 == 4.2
+
+    # Length-integrated variants
+    mmp_L = MagneticMultipoleParameters(Bn1L=np.float64(7.0), Ks1L=np.float64(8.0))
+    assert type(mmp_L.Bn1L) is float and mmp_L.Bn1L == 7.0
+    assert type(mmp_L.Ks1L) is float and mmp_L.Ks1L == 8.0
+
+    # ElectricMultipoleParameters: cover all prefixes
+    emp = ElectricMultipoleParameters(
+        tilt1=np.float64(0.2),
+        En1=np.float64(0.5),
+        Es1=np.int64(2),
+    )
+    assert type(emp.tilt1) is float and emp.tilt1 == 0.2
+    assert type(emp.En1) is float and emp.En1 == 0.5
+    assert type(emp.Es1) is int and emp.Es1 == 2
+
+    emp_L = ElectricMultipoleParameters(En1L=np.float64(1.0), Es1L=np.float64(0.5))
+    assert type(emp_L.En1L) is float and emp_L.En1L == 1.0
+    assert type(emp_L.Es1L) is float and emp_L.Es1L == 0.5
+
+    # Plain Python values must still pass through unchanged
+    mmp_plain = MagneticMultipoleParameters(Bn1=1.5)
+    assert type(mmp_plain.Bn1) is float and mmp_plain.Bn1 == 1.5

tests/test_serialization.py

@@ -1,5 +1,7 @@
 import os
 
+import pytest
+
 import pals
 
 
@@ -332,3 +334,75 @@ def test_comprehensive_lattice():
     # Clean up temporary files
     os.remove(yaml_file)
     os.remove(json_file)
+
+
+def _build_numpy_lattice(np):
+    """Build a small lattice using numpy-typed scalar values throughout."""
+    quad = pals.Quadrupole(
+        name="q_np",
+        length=np.float64(0.061),
+        MagneticMultipoleP=pals.MagneticMultipoleParameters(
+            Bn1=np.float64(-26.0), Bs1=np.float32(0.5), Kn0=np.int64(-1)
+        ),
+    )
+    oct_ = pals.Octupole(
+        name="o_np",
+        length=np.float64(0.25),
+        ElectricMultipoleP=pals.ElectricMultipoleParameters(
+            En3=np.float64(0.75), Es3=np.float32(0.125)
+        ),
+    )
+    return pals.BeamLine(name="line_np", line=[quad, oct_])
+
+
+def test_yaml_roundtrip_with_numpy():
+    """Regression test for issue #67: writing YAML with numpy-typed values
+    must not produce !!python/object tags, and round-tripping must yield
+    Python-native floats with the correct numeric values."""
+    np = pytest.importorskip("numpy")
+
+    line = _build_numpy_lattice(np)
+    yaml_file = "numpy_roundtrip.pals.yaml"
+    line.to_file(yaml_file)
+    try:
+        with open(yaml_file, "r") as f:
+            text = f.read()
+
+        # The bug symptom: YAML contains opaque numpy object tags.
+        assert "!!python/object" not in text, (
+            f"YAML output still contains unsafe numpy object tags:\n{text}"
+        )
+        assert "numpy" not in text, f"YAML output still references numpy:\n{text}"
+
+        loaded = pals.BeamLine.from_file(yaml_file)
+        loaded_quad = loaded.line[0]
+        assert loaded_quad.MagneticMultipoleP.Bn1 == -26.0
+        assert type(loaded_quad.MagneticMultipoleP.Bn1) is float
+        assert loaded_quad.MagneticMultipoleP.Bs1 == 0.5
+        assert loaded_quad.MagneticMultipoleP.Kn0 == -1
+
+        loaded_oct = loaded.line[1]
+        assert loaded_oct.ElectricMultipoleP.En3 == 0.75
+        assert type(loaded_oct.ElectricMultipoleP.En3) is float
+    finally:
+        if os.path.exists(yaml_file):
+            os.remove(yaml_file)
+
+
+def test_json_roundtrip_with_numpy():
+    """JSON path also needs to handle numpy values cleanly (defense-in-depth)."""
+    np = pytest.importorskip("numpy")
+
+    line = _build_numpy_lattice(np)
+    json_file = "numpy_roundtrip.pals.json"
+    line.to_file(json_file)
+    try:
+        loaded = pals.BeamLine.from_file(json_file)
+        loaded_quad = loaded.line[0]
+        assert loaded_quad.MagneticMultipoleP.Bn1 == -26.0
+        assert type(loaded_quad.MagneticMultipoleP.Bn1) is float
+        loaded_oct = loaded.line[1]
+        assert loaded_oct.ElectricMultipoleP.En3 == 0.75
+    finally:
+        if os.path.exists(json_file):
+            os.remove(json_file)