Add illustris

dsa2000_cal/dsa2000_cal/assets/source_models/illustris/fits_models/.large_files

@@ -0,0 +1 @@
+*.fits

dsa2000_cal/dsa2000_cal/assets/source_models/illustris/source_model.py

@@ -0,0 +1,17 @@
+import glob
+import os
+from typing import List
+
+from dsa2000_cal.assets.registries import source_model_registry
+from dsa2000_cal.assets.source_models.source_model import AbstractWSCleanSourceModel
+
+
+@source_model_registry(template='illustris')
+class IllustrisSourceModel(AbstractWSCleanSourceModel):
+
+    def get_wsclean_clean_component_file(self) -> str:
+        return os.path.join(*self.content_path, 'Illustris-sources.txt')
+
+    def get_wsclean_fits_files(self) -> List[str]:
+        fits_files = glob.glob(os.path.join(*self.content_path, 'fits_models', '449659_2.0_60.0_4096_128_1.0_5.0_*.fits'))
+        return fits_files

dsa2000_cal/dsa2000_cal/forward_models/synthetic_sky_model/synthetic_sky_model_producer.py

@@ -194,6 +194,61 @@ def create_a_team_sources(self, key=None, a_team_sources: List[str] | None = Non
             source_models.append(source_model)
         return source_models
 
+
+    def create_illustris_sources(self, key=None, illustris_sources: List[str] | None = None, full_stokes: bool = True) -> \
+            List[FITSSourceModel]:
+        """
+        Create the illustris sources as FITS models.
+
+        Args:
+            key: the random key
+            illustris_sources: the illustris sources to create
+
+        Returns:
+            sources: the illustris sources
+        """
+        fill_registries()
+        if key is None:
+            key = self.keys[3]
+        if illustris_sources is None:
+            illustris_sources = ["illustris"]
+        source_models = []
+        for source in illustris_sources:
+            source_model_asset = source_model_registry.get_instance(source_model_registry.get_match(source))
+            source_model = FITSSourceModel.from_wsclean_model(
+                wsclean_fits_files=source_model_asset.get_wsclean_fits_files(),
+                phase_tracking=self.phase_tracking, freqs=self.freqs, ignore_out_of_bounds=True,
+                full_stokes=full_stokes
+            )
+            source_models.append(source_model)
+        return source_models
+
+    def create_rfi_emitter_sources(self, key=None, rfi_sources: List[str] | None = None, full_stokes: bool = True) -> [
+        RFIEmitterSourceModel]:
+        """
+        Create RFI emitter sources inside the field.
+
+        Args:
+            key: the random key
+            rfi_sources: the LTE sources
+
+        Returns:
+            sources: a single facet model with all the RFI emitter sources
+        """
+        if key is None:
+            key = self.keys[4]
+        if rfi_sources is None:
+            rfi_sources = ["lte_cell_tower"]
+
+        fill_registries()
+        source_models = []
+        for rfi_source in rfi_sources:
+            rfi_model = rfi_model_registry.get_instance(rfi_model_registry.get_match(rfi_source))
+            rfi_model_params = rfi_model.make_source_params(freqs=self.freqs, full_stokes=full_stokes)
+            source_model = RFIEmitterSourceModel(rfi_model_params)
+            source_models.append(source_model)
+        return source_models
+
     def create_rfi_emitter_sources(self, key=None, rfi_sources: List[str] | None = None, full_stokes: bool = True) -> [
         RFIEmitterSourceModel]:
         """

dsa2000_cal/dsa2000_cal/forward_models/synthetic_sky_model/tests/demo_tests.py

@@ -32,3 +32,6 @@ def test_create_sky_model():
     a_team_sources = synthetic_sky_model_producer.create_a_team_sources(a_team_sources=['cas_a'])
     a_team_sources[0].plot(save_file='cas_a.png')
     assert len(a_team_sources) == 1
+    illustris_sources = synthetic_sky_model_producer.create_illustris_sources()
+    illustris_sources[0].plot(save_file='illustris.png')
+    assert len(illustris_sources) == 1

dsa2000_cal/dsa2000_cal/visibility_model/source_models/celestial/tests/demo_tests.py

@@ -21,7 +21,7 @@
 from dsa2000_cal.visibility_model.source_models.celestial.point_source_model import PointSourceModel
 
 
-@pytest.mark.parametrize('source', ['cas_a', 'cyg_a', 'tau_a', 'vir_a'])
+@pytest.mark.parametrize('source', ['cas_a', 'cyg_a', 'tau_a', 'vir_a', 'illustris-b'])
 def test_plot_ateam_sources(source):
     fill_registries()
     time = at.Time('2021-01-01T00:00:00', scale='utc')