Update and fix tests

tests/test_conf.toml

@@ -19,8 +19,17 @@ noisy = 380
 corrupted = true
 sensitivity_cut = 1e-6
 
+[polarisation_options]
+mode = "none"  # linear, circular, or "none"
+delta = 45  # phase angle 
+amp_ratio = 0.5  # polarisation amplitude ratio
+field_order = [0.1, 0.1]  # (x, y) orders of the random polarisation field
+field_scale = [0, 1]  # scaling of the intensity of the polarisation field
+field_threshold = "none"
+
 [bundle_options]
-in_path = "./tests/data"
+file_prefix=""
+in_path = "/home/knierim/work/pyvisgen/tests/data"
 out_path_fits = "./tests/build/fits"
 out_path_gridded = "./tests/build/gridded"
 num_test_images = 1000

tests/test_simulation.py

@@ -6,8 +6,10 @@
 from pyvisgen.utils.config import read_data_set_conf
 
 torch.manual_seed(1)
-config = "tests/test_conf.toml"
-conf = read_data_set_conf(config)
+
+CONFIG = "tests/test_conf.toml"
+
+conf = read_data_set_conf(CONFIG)
 out_path = Path(conf["out_path_fits"])
 out_path.mkdir(parents=True, exist_ok=True)
 
@@ -19,179 +21,182 @@ def test_get_data():
     assert len(data) > 0
 
 
-def test_create_sampling_rc():
-    from pyvisgen.simulation.data_set import create_sampling_rc, test_opts
+class TestSimulateDataSet:
+    """Unit test class for :class:``pyvisgen.simulation.SimulateDataSet``."""
+
+    def setup_class(self):
+        """Set up common objects and variables for the following tests."""
+        from pyvisgen.simulation.data_set import SimulateDataSet
+
+        self.s = SimulateDataSet
+
+    def test_run_no_slurm(self):
+        self.s.from_config(CONFIG)
 
-    samp_ops = create_sampling_rc(conf)
-    assert len(samp_ops) == 17
+    def test_run_no_slurm_multiprocess(self):
+        self.s.from_config(CONFIG, multiprocess="all")
 
-    test_opts(samp_ops)
+    def test_run_no_slurm_num_images(self):
+        self.s.from_config(CONFIG, num_images=50)
 
 
-def test_create_sampling_rc_no_seed():
-    from pyvisgen.simulation.data_set import create_sampling_rc, test_opts
+class TestVisLoop:
+    """Unit test class for :func:``pyvisgen.simulation.vis_loop``."""
 
-    mod_conf = conf.copy()
-    mod_conf["seed"] = None
-
-    samp_ops = create_sampling_rc(mod_conf)
-    assert len(samp_ops) == 17
-
-    test_opts(samp_ops)
-
-
-def test_vis_loop():
-    import torch
-
-    import pyvisgen.fits.writer as writer
-    from pyvisgen.simulation.data_set import create_observation
-    from pyvisgen.simulation.visibility import vis_loop
-    from pyvisgen.utils.data import load_bundles, open_bundles
-
-    bundles = load_bundles(conf["in_path"])
-    obs = create_observation(conf)
-    # num_active_telescopes = test_opts(samp_ops)
-    data = open_bundles(bundles[0])
-    SI = torch.tensor(data[0])[None]
-    vis_data = vis_loop(obs, SI, noisy=conf["noisy"], mode=conf["mode"])
-
-    assert (vis_data[0].V_11[0]).dtype == torch.complex128
-    assert (vis_data[0].V_22[0]).dtype == torch.complex128
-    assert (vis_data[0].V_12[0]).dtype == torch.complex128
-    assert (vis_data[0].V_21[0]).dtype == torch.complex128
-    assert (vis_data[0].num).dtype == torch.float64
-    assert (vis_data[0].base_num).dtype == torch.float64
-    assert torch.is_tensor(vis_data[0].u)
-    assert torch.is_tensor(vis_data[0].v)
-    assert torch.is_tensor(vis_data[0].w)
-    assert (vis_data[0].date).dtype == torch.float64
-
-    # test num vis for time step 0
-    # num_vis_theory = num_active_telescopes * (num_active_telescopes - 1) / 2
-    # num_vis_calc = vis_data.base_num[vis_data.date == vis_data.date[0]].shape[0]
-    # dunno what's going on here
-    # assert num_vis_theory == num_vis_calc
-    #
-
-    out_path = Path(conf["out_path_fits"])
-    out = out_path / Path("vis_0.fits")
-    hdu_list = writer.create_hdu_list(vis_data, obs)
-    hdu_list.writeto(out, overwrite=True)
-
-
-def test_vis_loop_grid():
-    import torch
-
-    import pyvisgen.fits.writer as writer
-    from pyvisgen.simulation.data_set import create_observation
-    from pyvisgen.simulation.visibility import vis_loop
-    from pyvisgen.utils.data import load_bundles, open_bundles
-
-    bundles = load_bundles(conf["in_path"])
-    obs = create_observation(conf)
-    # num_active_telescopes = test_opts(samp_ops)
-    data = open_bundles(bundles[0])
-    SI = torch.tensor(data[0])[None]
-    vis_data = vis_loop(obs, SI, noisy=conf["noisy"], mode="grid")
-
-    assert (vis_data[0].V_11[0]).dtype == torch.complex128
-    assert (vis_data[0].V_22[0]).dtype == torch.complex128
-    assert (vis_data[0].V_12[0]).dtype == torch.complex128
-    assert (vis_data[0].V_21[0]).dtype == torch.complex128
-    assert (vis_data[0].num).dtype == torch.float64
-    assert (vis_data[0].base_num).dtype == torch.float64
-    assert torch.is_tensor(vis_data[0].u)
-    assert torch.is_tensor(vis_data[0].v)
-    assert torch.is_tensor(vis_data[0].w)
-    assert (vis_data[0].date).dtype == torch.float64
-
-    # test num vis for time step 0
-    # num_vis_theory = num_active_telescopes * (num_active_telescopes - 1) / 2
-    # num_vis_calc = vis_data.base_num[vis_data.date == vis_data.date[0]].shape[0]
-    # dunno what's going on here
-    # assert num_vis_theory == num_vis_calc
-    #
-
-    out_path = Path(conf["out_path_fits"])
-    out = out_path / Path("vis_0.fits")
-    hdu_list = writer.create_hdu_list(vis_data, obs)
-    hdu_list.writeto(out, overwrite=True)
-
-
-def test_vis_loop_batch_size_auto():
-    import torch
-
-    from pyvisgen.simulation.data_set import create_observation
-    from pyvisgen.simulation.visibility import vis_loop
-    from pyvisgen.utils.data import load_bundles, open_bundles
-
-    bundles = load_bundles(conf["in_path"])
-    obs = create_observation(conf)
-    data = open_bundles(bundles[0])
-    SI = torch.tensor(data[0])[None]
-
-    vis_data = vis_loop(
-        obs,
-        SI,
-        noisy=conf["noisy"],
-        mode=conf["mode"],
-        batch_size="auto",
-    )
-
-    assert (vis_data[0].V_11[0]).dtype == torch.complex128
-    assert (vis_data[0].V_22[0]).dtype == torch.complex128
-    assert (vis_data[0].V_12[0]).dtype == torch.complex128
-    assert (vis_data[0].V_21[0]).dtype == torch.complex128
-    assert (vis_data[0].num).dtype == torch.float64
-    assert (vis_data[0].base_num).dtype == torch.float64
-    assert torch.is_tensor(vis_data[0].u)
-    assert torch.is_tensor(vis_data[0].v)
-    assert torch.is_tensor(vis_data[0].w)
-    assert (vis_data[0].date).dtype == torch.float64
-
-
-def test_vis_loop_batch_size_invalid():
-    import torch
-
-    from pyvisgen.simulation.data_set import create_observation
-    from pyvisgen.simulation.visibility import vis_loop
-    from pyvisgen.utils.data import load_bundles, open_bundles
-
-    bundles = load_bundles(conf["in_path"])
-    obs = create_observation(conf)
-    data = open_bundles(bundles[0])
-    SI = torch.tensor(data[0])[None]
-
-    assert_raises(
-        ValueError,
-        vis_loop,
-        obs,
-        SI,
-        noisy=conf["noisy"],
-        mode=conf["mode"],
-        batch_size="abc",
-    )
-
-    assert_raises(
-        ValueError,
-        vis_loop,
-        obs,
-        SI,
-        noisy=conf["noisy"],
-        mode=conf["mode"],
-        batch_size=20.0,
-    )
+    def setup_class(self):
+        """Set up common objects and variables for the following tests."""
+        from pyvisgen.simulation.data_set import SimulateDataSet
+
+        self.s = SimulateDataSet
+
+    def test_vis_loop(self):
+        import pyvisgen.fits.writer as writer
+        from pyvisgen.simulation.visibility import vis_loop
+        from pyvisgen.utils.data import load_bundles, open_bundles
+
+        bundles = load_bundles(conf["in_path"])
+        _, obs = self.s._get_obs_test(CONFIG)
+
+        data = open_bundles(bundles[0])
+        SI = torch.tensor(data[0])[None]
+        vis_data = vis_loop(obs, SI, noisy=conf["noisy"], mode=conf["mode"])
+
+        assert (vis_data[0].V_11[0]).dtype == torch.complex128
+        assert (vis_data[0].V_22[0]).dtype == torch.complex128
+        assert (vis_data[0].V_12[0]).dtype == torch.complex128
+        assert (vis_data[0].V_21[0]).dtype == torch.complex128
+        assert (vis_data[0].num).dtype == torch.float64
+        assert (vis_data[0].base_num).dtype == torch.float64
+        assert torch.is_tensor(vis_data[0].u)
+        assert torch.is_tensor(vis_data[0].v)
+        assert torch.is_tensor(vis_data[0].w)
+        assert (vis_data[0].date).dtype == torch.float64
+
+        # test num vis for time step 0
+        # num_vis_theory = num_active_telescopes * (num_active_telescopes - 1) / 2
+        # num_vis_calc = vis_data.base_num[vis_data.date == vis_data.date[0]].shape[0]
+        # dunno what's going on here
+        # assert num_vis_theory == num_vis_calc
+        #
+
+        out_path = Path(conf["out_path_fits"])
+        out = out_path / Path("vis_0.fits")
+        hdu_list = writer.create_hdu_list(vis_data, obs)
+        hdu_list.writeto(out, overwrite=True)
+
+    def test_vis_loop_grid(self):
+        import pyvisgen.fits.writer as writer
+        from pyvisgen.simulation.visibility import vis_loop
+        from pyvisgen.utils.data import load_bundles, open_bundles
+
+        bundles = load_bundles(conf["in_path"])
+        _, obs = self.s._get_obs_test(CONFIG)
+
+        data = open_bundles(bundles[0])
+        SI = torch.tensor(data[0])[None]
+        vis_data = vis_loop(obs, SI, noisy=conf["noisy"], mode="grid")
+
+        assert (vis_data[0].V_11[0]).dtype == torch.complex128
+        assert (vis_data[0].V_22[0]).dtype == torch.complex128
+        assert (vis_data[0].V_12[0]).dtype == torch.complex128
+        assert (vis_data[0].V_21[0]).dtype == torch.complex128
+        assert (vis_data[0].num).dtype == torch.float64
+        assert (vis_data[0].base_num).dtype == torch.float64
+        assert torch.is_tensor(vis_data[0].u)
+        assert torch.is_tensor(vis_data[0].v)
+        assert torch.is_tensor(vis_data[0].w)
+        assert (vis_data[0].date).dtype == torch.float64
+
+        # test num vis for time step 0
+        # num_vis_theory = num_active_telescopes * (num_active_telescopes - 1) / 2
+        # num_vis_calc = vis_data.base_num[vis_data.date == vis_data.date[0]].shape[0]
+        # dunno what's going on here
+        # assert num_vis_theory == num_vis_calc
+        #
+
+        out_path = Path(conf["out_path_fits"])
+        out = out_path / Path("vis_0.fits")
+        hdu_list = writer.create_hdu_list(vis_data, obs)
+        hdu_list.writeto(out, overwrite=True)
+
+    def test_vis_loop_batch_size_auto(self):
+        from pyvisgen.simulation.visibility import vis_loop
+        from pyvisgen.utils.data import load_bundles, open_bundles
+
+        bundles = load_bundles(conf["in_path"])
+        _, obs = self.s._get_obs_test(CONFIG)
+
+        data = open_bundles(bundles[0])
+        SI = torch.tensor(data[0])[None]
+
+        vis_data = vis_loop(
+            obs,
+            SI,
+            noisy=conf["noisy"],
+            mode=conf["mode"],
+            batch_size="auto",
+        )
+
+        assert (vis_data[0].V_11[0]).dtype == torch.complex128
+        assert (vis_data[0].V_22[0]).dtype == torch.complex128
+        assert (vis_data[0].V_12[0]).dtype == torch.complex128
+        assert (vis_data[0].V_21[0]).dtype == torch.complex128
+        assert (vis_data[0].num).dtype == torch.float64
+        assert (vis_data[0].base_num).dtype == torch.float64
+        assert torch.is_tensor(vis_data[0].u)
+        assert torch.is_tensor(vis_data[0].v)
+        assert torch.is_tensor(vis_data[0].w)
+        assert (vis_data[0].date).dtype == torch.float64
+
+    def test_vis_loop_batch_size_invalid(self):
+        from pyvisgen.simulation.visibility import vis_loop
+        from pyvisgen.utils.data import load_bundles, open_bundles
+
+        bundles = load_bundles(conf["in_path"])
+        _, obs = self.s._get_obs_test(CONFIG)
+
+        data = open_bundles(bundles[0])
+        SI = torch.tensor(data[0])[None]
+
+        assert_raises(
+            ValueError,
+            vis_loop,
+            obs,
+            SI,
+            noisy=conf["noisy"],
+            mode=conf["mode"],
+            batch_size="abc",
+        )
+
+        assert_raises(
+            ValueError,
+            vis_loop,
+            obs,
+            SI,
+            noisy=conf["noisy"],
+            mode=conf["mode"],
+            batch_size=20.0,
+        )
 
 
 class TestPolarisation:
     """Unit test class for ``pyvisgen.simulation.visibility.Polarisation``."""
 
     def setup_class(self):
         """Set up common objects and variables for the following tests."""
-        from pyvisgen.simulation.data_set import create_observation
+        from pyvisgen.simulation.data_set import SimulateDataSet
+        from pyvisgen.simulation.observation import (
+            DEFAULT_FIELD_KWARGS,
+            DEFAULT_POL_KWARGS,
+        )
         from pyvisgen.simulation.visibility import Polarisation
 
-        self.obs = create_observation(conf)
+        _, self.obs = SimulateDataSet._get_obs_test(conf)
+
+        # set to default kwargs for tests, otherwise
+        # some parameters of the config would be None
+        self.obs.field_kwargs = DEFAULT_FIELD_KWARGS
+        self.obs.pol_kwargs = DEFAULT_POL_KWARGS
 
         self.SI = torch.zeros((100, 100))
         self.SI[25::25, 25::25] = 1
@@ -408,9 +413,3 @@ def test_polarisation_field_threshold(self):
 
         # expected to raise an AssertionError
         assert_raises(AssertionError, assert_array_equal, pf, pf_ref)
-
-
-def test_simulate_data_set_no_slurm():
-    from pyvisgen.simulation.data_set import simulate_data_set
-
-    simulate_data_set(config)