basic PropagationCollection

tests/test_io_handlers.py

@@ -7,7 +7,7 @@ import pytest
 
 from scgenerator.io import DataFile, MemoryIOHandler, ZipFileIOHandler, unique_name
 from scgenerator.parameter import Parameters
-from scgenerator.spectra import PARAMS_FN, propagation
+from scgenerator.spectra import PARAMS_FN, propagation, propagation_series
 
 PARAMS = dict(
     name="PM2000D sech pulse",
@@ -25,7 +25,7 @@ PARAMS = dict(
     raman_type="measured",
     quantum_noise=True,
     interpolation_degree=11,
-    z_num=128,
+    z_num=8,
     length=1.5,
     t_num=512,
     dt=5e-15,
@@ -173,3 +173,28 @@ def test_zip_bundle(tmp_path: Path):
 def test_unique_name():
     existing = {"spec.npy", "spec_0.npy"}
     assert unique_name("spec.npy", existing) == "spec_1.npy"
+
+    existing = {"spec", "spec0"}
+    assert unique_name("spec", existing) == "spec_0"
+
+    existing = {"spec", "spec_0"}
+    assert unique_name("spec", existing) == "spec_1"
+
+
+def test_propagation_series(tmp_path: Path):
+    params = Parameters(**PARAMS)
+    with pytest.raises(ValueError):
+        specs, props = propagation_series([])
+
+    flist = [tmp_path / f"prop{i}.zip" for i in range(10)]
+    for i, f in enumerate(flist):
+        params.name = f"prop {i}"
+        prop = propagation(f, params)
+        for _ in range(params.z_num):
+            prop.append(np.zeros(params.t_num, dtype=complex))
+
+    assert set(flist) == set(tmp_path.glob("*.zip"))
+
+    specs, propagations = propagation_series(flist)
+    assert specs.shape == (10, params.z_num, params.t_num)
+    assert all(prop.parameters.name == f"prop {i}" for i, prop in enumerate(propagations))