diff --git a/docs/source/tutorials/maximum-likelihood-mapper.ipynb b/docs/source/tutorials/maximum-likelihood-mapper.ipynb index 58d2e28b..977c21c0 100644 --- a/docs/source/tutorials/maximum-likelihood-mapper.ipynb +++ b/docs/source/tutorials/maximum-likelihood-mapper.ipynb @@ -107,8 +107,8 @@ "from maria.mappers import MaximumLikelihoodMapper\n", "\n", "ml_mapper = MaximumLikelihoodMapper(tods=tods, \n", - " width=0.8*input_map.width.deg,\n", - " height=0.8*input_map.height.deg,\n", + " width=0.8 * input_map.width.deg,\n", + " height=0.8 * input_map.height.deg,\n", " units=\"Jy/pixel\")\n", "print(f\"{ml_mapper.loss() = }\")" ] diff --git a/maria/atmosphere/atmosphere.py b/maria/atmosphere/atmosphere.py index 909b1a67..59fac5a0 100644 --- a/maria/atmosphere/atmosphere.py +++ b/maria/atmosphere/atmosphere.py @@ -77,35 +77,35 @@ def __init__( self._initialized = False - def initialize(self, instrument, boresight, site): + def initialize(self, obs): """ Simulate a realization of PWV. """ self.layers = generate_layers( - instrument=instrument, - boresight=boresight, + instrument=obs.instrument, + boresight=obs.boresight, weather=self.weather, - site=site, + site=obs.site, mode=self.model, angular=self.angular, max_height=self.max_height, ) if self.timestep is None: - min_fwhm = instrument.dets.angular_fwhm(z=self.max_height).min() + min_fwhm = obs.instrument.dets.angular_fwhm(z=self.max_height).min() max_wind = Quantity((self.layers.wind_speed / self.layers.h).values, "rad/s").max() - self.timestep = (min_fwhm / max_wind).s + self.timestep = max(1e-1, (min_fwhm / max_wind).s) - self.boresight = boresight.downsample(timestep=self.timestep) + self.boresight = obs.boresight.downsample(timestep=self.timestep) self.coords = self.boresight.broadcast( - instrument.dets.offsets, + obs.instrument.dets.offsets, frame="az/el", ) # this is a smaller version of the sim coords outer_coords = self.boresight.broadcast( - instrument.dets.outer().offsets, + obs.instrument.dets.outer().offsets, frame="az/el", ) @@ -206,23 +206,34 @@ def initialize(self, instrument, boresight, site): self.layers.loc[in_process].iterrows(), ): # this algorithm probably works - res = layer.res - wide_lp_x_dense = np.linspace(min_ty, max_ty, 10000) - wide_lp_z_dense = layer.h * np.ones(len(wide_lp_x_dense)) - wide_lp_dense = np.c_[wide_lp_x_dense, wide_lp_z_dense] - interior = triangulation.find_simplex(wide_lp_dense) > -1 - # interior = interior if interior.any() else np.ones_like(interior, dtype=bool) - lp_dense_x = wide_lp_x_dense[interior] - lp_x_min = lp_dense_x.min() - 2 * res - lp_x_max = lp_dense_x.max() + 2 * res - n_lp = np.maximum(3, int((lp_x_max - lp_x_min) / res)) - lp_x = np.linspace(lp_x_min, lp_x_max, n_lp) - lp_z = layer.h * np.ones(len(lp_x)) - lp = np.c_[lp_x, lp_z] - n_lp = len(lp) - - layer_labels.extend(n_lp * [layer_index]) - cross_section_points_list.append(lp) + n_cross_section = int(np.maximum(2, (np.ptp(tp[:, 1]) + 2 * layer.res) / layer.res)) + cross_section_points = np.stack( + [ + np.linspace(tp[:, 1].min() - layer.res, tp[:, 1].max() + layer.res, n_cross_section), + layer.h * np.ones(n_cross_section), + ] + ).T + + layer_labels.extend(n_cross_section * [layer_index]) + cross_section_points_list.append(cross_section_points) + + # res = layer.res + # wide_lp_x_dense = np.linspace(min_ty, max_ty, 10000) + # wide_lp_z_dense = layer.h * np.ones(len(wide_lp_x_dense)) + # wide_lp_dense = np.c_[wide_lp_x_dense, wide_lp_z_dense] + # interior = triangulation.find_simplex(wide_lp_dense) > -1 + # # interior = interior if interior.any() else np.ones_like(interior, dtype=bool) + # lp_dense_x = wide_lp_x_dense[interior] + # lp_x_min = lp_dense_x.min() - 2 * res + # lp_x_max = lp_dense_x.max() + 2 * res + # n_lp = np.maximum(3, int((lp_x_max - lp_x_min) / res)) + # lp_x = np.linspace(lp_x_min, lp_x_max, n_lp) + # lp_z = layer.h * np.ones(len(lp_x)) + # lp = np.c_[lp_x, lp_z] + # n_lp = len(lp) + + # layer_labels.extend(n_lp * [layer_index]) + # cross_section_points_list.append(lp) cross_section_points = np.concatenate(cross_section_points_list, axis=0) diff --git a/maria/map/base.py b/maria/map/base.py index df522ead..ea1a6e90 100644 --- a/maria/map/base.py +++ b/maria/map/base.py @@ -131,9 +131,7 @@ def __init__( if np.shape(beam)[-1] != 3: raise ValueError("'beam' must be either a number or a tuple of (major, minor, angle)") - beam = beam * np.ones((1, 3) if "nu" in self.dims else 3) - - self.beam = Quantity(beam, "deg" if degrees else "rad") + self.beam = Quantity(beam * np.ones((*self.slice_dims.values(), 3)), "deg" if degrees else "rad") if self.u["quantity"] == "spectral_flux_density_per_beam": if not np.all(self.beam_area > 0): @@ -174,12 +172,11 @@ def t_bins(self): """ This might break in the year 3000. """ - t_min = arrow.get("0001-01-01").timestamp() - t_max = arrow.get("3000-01-01").timestamp() - - return np.array([t_min, *(self.t[1:] + self.t[:-1]) / 2, t_max]) + # t_min = arrow.get("0001-01-01").timestamp() + # t_max = arrow.get("3000-01-01").timestamp() + # return np.array([-np.inf, *(self.t[1:] + self.t[:-1]) / 2, np.inf]) - # return np.array([-np.inf, *(self.t[1:] + self.t[:-1])/2, np.inf]) + return np.array([-np.inf, *(self.t[1:] + self.t[:-1]) / 2, np.inf]) @property def t_side(self): @@ -226,7 +223,7 @@ def unsqueeze(self, dim, value=None): raise Exception(f"{self.__class__.__name__} already has dimension '{dim}'") new_dim_index = 0 - for d in ["stokes", "nu", "t"]: + for d in ["stokes", "nu", "t", "z"]: if d == dim: break if d in self.dims: @@ -235,6 +232,8 @@ def unsqueeze(self, dim, value=None): package = self.package() package["data"] = np.expand_dims(package["data"], new_dim_index) package["weight"] = np.expand_dims(package["weight"], new_dim_index) + package["beam"] = np.expand_dims(package["beam"], new_dim_index) + if value is None: value = SLICE_DIMS[dim]["default"] package[dim] = value @@ -247,7 +246,8 @@ def beam_area(self): Returns the beam area in steradians """ area = (np.pi / 4) * self.beam[..., 0].radians * self.beam[..., 1].radians - return Quantity(area * np.ones(tuple(self.slice_dims.values())), "sr") + area = np.expand_dims(area, axis=(-1, -2)[: len(self.map_dims)]) + return Quantity(area, "sr") def to(self, units: str, **calibration_kwargs: Mapping): if units == self.units: diff --git a/maria/map/projection.py b/maria/map/projection.py index 31496dde..1af3906e 100644 --- a/maria/map/projection.py +++ b/maria/map/projection.py @@ -72,7 +72,6 @@ def __init__( raise ValueError("'center' must be a two-tuple of numbers") self.frame = Frame(frame) - beam = beam if beam is not None else 0.0 super().__init__( data=data, @@ -105,9 +104,13 @@ def __init__( y_res = height / self.n_y else: y_res = x_res - else: + elif height is not None: # here height must not be None x_res = y_res = height / self.n_y + else: + raise ValueError( + f"{self.__class__.__name__} needs one of ['width', 'height', 'resolution', 'x_res', 'y_res']." + ) self.x_res = Quantity(x_res, "deg" if degrees else "rad") self.y_res = Quantity(y_res, "deg" if degrees else "rad") @@ -218,33 +221,32 @@ def __getattr__(self, attr): raise AttributeError(f"'ProjectionMap' object has no attribute '{attr}'") def __getitem__(self, key): - if not hasattr(key, "__iter__"): - key = (key,) + key = key if isinstance(key, tuple) else (key,) explicit_slices = unpack_implicit_slice(key, ndims=self.ndim) - - _, dimensions = get_factor_and_base_units_vector(self.units) - if all([dimensions[dim] == 0 for dim in ["pixel", "rad"]]): - # don't convert if in resolution-neutral units (like temperature) - package = self.package() - else: - # convert if in non-resolution-neutral units (like spectral flux) - package = self.to("K_RJ").package() + package = self.package() package["data"] = package["data"][key] package["weight"] = package["weight"][key] + package["beam"] = package["beam"][explicit_slices[: -len(self.map_dims)]] for axis, (dim, naxis) in enumerate(self.dims.items()): if dim not in "xy": - package[dim] = package[dim][explicit_slices[axis]] + if isinstance(explicit_slices[axis], int): + package.pop(dim) + else: + package[dim] = package[dim][explicit_slices[axis]] - if dim in "nu": - package["beam"] = package["beam"][explicit_slices[axis]] + x_res_factor = explicit_slices[-1].step or 1.0 + y_res_factor = explicit_slices[-2].step or 1.0 + + _, dimensions = get_factor_and_base_units_vector(self.units) - package["y_res"] *= explicit_slices[-2].step or 1 - package["x_res"] *= explicit_slices[-1].step or 1 + package["y_res"] *= y_res_factor + package["x_res"] *= x_res_factor + package["data"] *= (x_res_factor * y_res_factor) ** dimensions.pixel - return ProjectionMap(**package).to(self.units) + return ProjectionMap(**package) def downsample(self, reduce: tuple, func: Callable = np.mean): if reduce: diff --git a/maria/mappers/base.py b/maria/mappers/base.py index 138eeca6..b8acb7ff 100644 --- a/maria/mappers/base.py +++ b/maria/mappers/base.py @@ -70,7 +70,6 @@ def __init__( else: self.timestep = timestep time_bins = np.arange(min_time, max_time + timestep, timestep) - time_bins += mean_time - self.time_bins.mean() self.t = (time_bins[1:] + time_bins[:-1]) / 2 self.bands = BandList([]) @@ -78,13 +77,13 @@ def __init__( self.tods = [] self.add_tods(tods) - beam_sum = np.zeros((len(self.bands), 1, 3)) - beam_wgt = np.zeros((len(self.bands), 1, 3)) + beam_sum = np.zeros((len(self.nu), 1, 3)) + beam_wgt = np.zeros((len(self.nu), 1, 3)) - for band_index, band in enumerate(self.bands): + for nu_index, nu in enumerate(self.nu): for tod in self.tods: - beam_sum[band_index] += tod.duration * tod.dets.beams[tod.dets.band_name == band.name].mean(axis=0) - beam_wgt[band_index] += tod.duration + beam_sum[nu_index] += tod.duration * tod.dets.beams[tod.dets.band_center == nu].mean(axis=0) + beam_wgt[nu_index] += tod.duration self.beam = beam_sum / beam_wgt diff --git a/maria/sim/map.py b/maria/sim/map.py index 4fc7c82e..5661e87f 100644 --- a/maria/sim/map.py +++ b/maria/sim/map.py @@ -84,10 +84,6 @@ def _sample_maps(self, obs): band_coords = obs.coords[band_mask] band_dets = obs.instrument.dets[band_mask] - pointing_s = ttime.monotonic() - P = self.map.pointing_matrix(coords=band_coords, dets=band_dets) - logger.debug(f"Computed pointing matrix for band {band.name} in {humanize_time(ttime.monotonic() - pointing_s)}") - band_fwhm = Quantity( compute_angular_fwhm( fwhm_0=obs.instrument.dets.primary_size.mean(), @@ -102,10 +98,13 @@ def _sample_maps(self, obs): for channel_index, (nu_min, nu_max) in enumerate(self.map.nu_bin_bounds): channel_s = ttime.monotonic() - channel_map = smoothed_map.to("K_RJ", band=band).data[:, channel_index] + channel_map = smoothed_map.to("K_RJ", band=band)[:, [channel_index]] qchannel = (nu_min, nu_max) channel_string = f"{qchannel}" + if (band.nu.Hz.max() < nu_min.Hz) or (nu_max.Hz < band.nu.Hz.min()): + continue + spectrum_kwargs = ( { "spectrum": obs.atmosphere.spectrum, @@ -128,6 +127,8 @@ def _sample_maps(self, obs): f"{humanize_time(ttime.monotonic() - calibration_s)}" ) + bands_pbar.set_postfix(band=band.name, channel=channel_string) + # for stokes_index, stokes in enumerate(getattr(self.map, "stokes", "I")): # bands_pbar.set_postfix(band=band.name, channel=channel_string, stokes=stokes) @@ -150,7 +151,13 @@ def _sample_maps(self, obs): # flat_padded_map = np.pad(channel_stokes_map, pad_width=((1, 1)), mode="edge").ravel() # flat_padded_map = np.pad(channel_stokes_map, pad_width=((1, 1)), mode="edge").ravel() - pW = pW_per_K_RJ * (P @ channel_map.ravel()).reshape(band_coords.shape) + pointing_s = ttime.monotonic() + P = channel_map.pointing_matrix(coords=band_coords, dets=band_dets) + logger.debug( + f"Computed pointing matrix for band {band.name} in {humanize_time(ttime.monotonic() - pointing_s)}" + ) + + pW = pW_per_K_RJ * (P @ channel_map.data.ravel()).reshape(band_coords.shape) map_loading[band_mask] += pW diff --git a/maria/sim/simulation.py b/maria/sim/simulation.py index 87a194b9..401d3075 100644 --- a/maria/sim/simulation.py +++ b/maria/sim/simulation.py @@ -200,7 +200,7 @@ def run_obs(self, obs: Observation) -> TOD: if hasattr(obs, "atmosphere"): atmosphere_sim_start_s = ttime.monotonic() - obs.atmosphere.initialize(instrument=obs.instrument, boresight=obs.boresight, site=obs.site) + obs.atmosphere.initialize(obs) self._simulate_atmosphere(obs) obs.loading["atmosphere"] = self._compute_atmospheric_loading(obs) logger.debug(f"Ran atmosphere simulation in {humanize_time(ttime.monotonic() - atmosphere_sim_start_s)}.") diff --git a/maria/tests/io/test_io.py b/maria/tests/io/test_io.py index 2a3b9d6e..c6ef9678 100644 --- a/maria/tests/io/test_io.py +++ b/maria/tests/io/test_io.py @@ -19,11 +19,3 @@ def test_change_cache_dir(): del os.environ["MARIA_CACHE_DIR"] maria.set_cache_dir("/tmp/maria-data") raise Exception() - - -@pytest.mark.parametrize("filename", all_maps) -def test_all_maps(filename): - m = maria.map.load(filename=fetch(filename), width=0.1, center=(150, 10)) # noqa - m.plot() - - plt.close("all") diff --git a/maria/tests/map/test_map_io.py b/maria/tests/map/test_map_io.py new file mode 100644 index 00000000..3da545fd --- /dev/null +++ b/maria/tests/map/test_map_io.py @@ -0,0 +1,46 @@ +from __future__ import annotations + +import maria +import matplotlib.pyplot as plt +import numpy as np +import pytest +from maria.instrument import Band +from maria.io import fetch +from maria.map import all_maps +from maria.mappers import BinMapper + +plt.close("all") + + +@pytest.mark.parametrize("filename", all_maps) +def test_all_maps(filename): + m = maria.map.load(filename=fetch(filename), width=0.1, center=(150, 10)) # noqa + m.plot() + + plt.close("all") + + +@pytest.mark.parametrize( + "map_path", + all_maps, +) # noqa +def test_map_io_units(map_path): + m = maria.map.load(fetch(map_path)) + if "nu" not in m.dims: + m = m.unsqueeze("nu", 150e9) + assert np.allclose(m.to("K_RJ").to("Jy/pixel").to(m.units).data, m.data).compute() + + m.to("cK_RJ").to_hdf("/tmp/test_write_map.h5") + new_m_hdf = maria.map.load("/tmp/test_write_map.h5").to(m.units) # noqa + + assert np.allclose(new_m_hdf.data, m.data) + assert np.allclose(new_m_hdf.resolution.arcsec, m.resolution.arcsec) + + if "fits" in map_path: + m.to("cK_RJ").to_fits("/tmp/test_write_map.fits") + new_m_fits = maria.map.load("/tmp/test_write_map.fits").to(m.units) # noqa + + assert np.allclose(new_m_fits.data, m.data) + assert np.allclose(new_m_fits.resolution.arcsec, m.resolution.arcsec) + + m.plot() diff --git a/maria/tests/map/test_map_ops.py b/maria/tests/map/test_map_ops.py new file mode 100644 index 00000000..fc26806f --- /dev/null +++ b/maria/tests/map/test_map_ops.py @@ -0,0 +1,35 @@ +from __future__ import annotations + +import maria +import matplotlib.pyplot as plt +import numpy as np +from maria.io import fetch +from maria.map import ProjectionMap + +plt.close("all") + + +def test_map_extend(): + map_filename = fetch("maps/cluster1.fits") + + m1 = maria.map.load(filename=map_filename, nu=90e9) + m2 = maria.map.load(filename=map_filename, nu=150e9) + m3 = maria.map.load(filename=map_filename, nu=220e9) + + m4 = m1.extend([m2, m3], dim="nu").unsqueeze("stokes") + m5, m6 = m4.copy(), m4.copy() + m5.stokes = "Q" + m6.stokes = "U" + + m4.extend([m5, m6], dim="stokes") + + +def test_map_slice(): + stokes = "IQUV" + nu = [90e9, 150e9, 220e9] + t = 1.7e9 + np.arange(0, 600, 120) + data = np.random.standard_normal((len(stokes), len(nu), len(t), 100, 100)) + + m = ProjectionMap(data=data, width=1e0, stokes=stokes, nu=nu, t=t, center=(0, -30), frame="ra_dec") + + m[0, :, ::2, :2] diff --git a/maria/tests/map/test_map_units.py b/maria/tests/map/test_map_units.py index 4206c380..f28fb445 100644 --- a/maria/tests/map/test_map_units.py +++ b/maria/tests/map/test_map_units.py @@ -16,7 +16,7 @@ "map_path", all_maps, ) # noqa -def test_maps(map_path): +def test_map_io_units(map_path): m = maria.map.load(fetch(map_path)) if "nu" not in m.dims: m = m.unsqueeze("nu", 150e9) @@ -38,7 +38,7 @@ def test_maps(map_path): m.plot() -def test_map_operations(): +def test_map_extend(): map_filename = fetch("maps/cluster1.fits") m1 = maria.map.load(filename=map_filename, nu=90e9) @@ -51,76 +51,3 @@ def test_map_operations(): m6.stokes = "U" m4.extend([m5, m6], dim="stokes") - - -def test_trivial_recover_original_map(): - f090 = Band(center=90e9, width=30e9) - f150 = Band(center=150e9, width=40e9) - f220 = Band(center=220e9, width=50e9) - - map_filename = fetch("maps/cluster1.fits") - input_map = maria.map.load(filename=map_filename, nu=150e9).to("K_RJ")[..., ::8, ::8] - - input_map.data -= input_map.data.mean().compute() - map_width = input_map.width.degrees - - array = { - "field_of_view": map_width / 2, - "primary_size": 1000, - "n": 300, - "bands": [f090, f150, f220], - } - - instrument = maria.get_instrument(array=array) - - planner = maria.Planner(target=input_map, site="cerro_toco", constraints={"el": (40, 90)}) - - plans = planner.generate_plans( - total_duration=60, sample_rate=50, scan_pattern="daisy", scan_options={"radius": input_map.width.deg / 3} - ) - - sim = maria.Simulation( - instrument, - plans=plans, - map=input_map, - site="cerro_toco", - noise=False, - ) - - tods = sim.run() - - mapper = BinMapper( - center=input_map.center, - frame=input_map.frame, - width=input_map.width, - resolution=input_map.x_res, - degrees=False, - tods=tods, - ) - - mapper.add_tods(tods) - output_map = mapper.run() - - m0 = input_map.data[0, 0, 0].compute() - m1 = output_map.data[0, :].compute() - w = output_map.weight[-1, 0].compute() - - relsqres = np.sqrt(np.nansum(w * (m1 - m0) ** 2, axis=(-1, -2)) / np.nansum(w)) - - assert all(relsqres < 1e-3) - - output_map.to("Jy/pixel") - - -def test_time_ordered_map_sim(): - time_evolving_sun_path = fetch("maps/sun.h5") - input_map = maria.map.load(filename=time_evolving_sun_path, nu=100e9, t=1.7e9 + np.linspace(0, 180, 16)) - plan = maria.Plan.generate( - start_time=1.7e9, - duration=180, - scan_center=(input_map.center), - scan_options={"radius": 0.25}, - ) - sim = maria.Simulation(instrument="test/1deg", site="cerro_toco", plans=plan, map=input_map) - tods = sim.run() - tods[0].to("K_RJ").plot() diff --git a/maria/tests/map/test_maps.py b/maria/tests/map/test_maps.py deleted file mode 100644 index 00c09259..00000000 --- a/maria/tests/map/test_maps.py +++ /dev/null @@ -1,128 +0,0 @@ -from __future__ import annotations - -import maria -import matplotlib.pyplot as plt -import numpy as np -import pytest -from maria.instrument import Band -from maria.io import fetch -from maria.map import all_maps -from maria.mappers import BinMapper - -plt.close("all") - - -@pytest.mark.parametrize( - "map_path", - all_maps, -) # noqa -def test_maps(map_path): - m = maria.map.load(fetch(map_path)) - if "nu" not in m.dims: - m = m.unsqueeze("nu", 150e9) - assert np.allclose(m.to("K_RJ").to("Jy/pixel").to(m.units).data, m.data).compute() - - m.to("cK_RJ").to_hdf("/tmp/test_write_map.h5") - new_m_hdf = maria.map.load("/tmp/test_write_map.h5").to(m.units) # noqa - - assert np.allclose(new_m_hdf.data, m.data) - assert np.allclose(new_m_hdf.resolution.arcsec, m.resolution.arcsec) - - if "fits" in map_path: - m.to("cK_RJ").to_fits("/tmp/test_write_map.fits") - new_m_fits = maria.map.load("/tmp/test_write_map.fits").to(m.units) # noqa - - assert np.allclose(new_m_fits.data, m.data) - assert np.allclose(new_m_fits.resolution.arcsec, m.resolution.arcsec) - - m.plot() - - plt.close("all") - - -def test_map_operations(): - map_filename = fetch("maps/cluster1.fits") - - m1 = maria.map.load(filename=map_filename, nu=90e9) - m2 = maria.map.load(filename=map_filename, nu=150e9) - m3 = maria.map.load(filename=map_filename, nu=220e9) - - m4 = m1.extend([m2, m3], dim="nu").unsqueeze("stokes") - m5, m6 = m4.copy(), m4.copy() - m5.stokes = "Q" - m6.stokes = "U" - - m4.extend([m5, m6], dim="stokes") - - -def test_trivial_recover_original_map(): - f090 = Band(center=90e9, width=30e9) - f150 = Band(center=150e9, width=40e9) - f220 = Band(center=220e9, width=50e9) - - map_filename = fetch("maps/cluster1.fits") - input_map = maria.map.load(filename=map_filename, nu=150e9).to("K_RJ")[..., ::8, ::8] - - input_map.data -= input_map.data.mean().compute() - map_width = input_map.width.degrees - - array = { - "field_of_view": map_width / 2, - "primary_size": 1000, - "n": 300, - "bands": [f090, f150, f220], - } - - instrument = maria.get_instrument(array=array) - - planner = maria.Planner(target=input_map, site="cerro_toco", constraints={"el": (40, 90)}) - - plans = planner.generate_plans( - total_duration=60, sample_rate=50, scan_pattern="daisy", scan_options={"radius": input_map.width.deg / 3} - ) - - sim = maria.Simulation( - instrument, - plans=plans, - map=input_map, - site="cerro_toco", - noise=False, - ) - - tods = sim.run() - - mapper = BinMapper( - center=input_map.center, - frame=input_map.frame, - width=input_map.width, - resolution=input_map.x_res, - degrees=False, - tods=tods, - ) - - mapper.add_tods(tods) - output_map = mapper.run() - - m0 = input_map.data[0, 0, 0].compute() - m1 = output_map.data[0, :].compute() - w = output_map.weight[-1, 0].compute() - - relsqres = np.sqrt(np.nansum(w * (m1 - m0) ** 2, axis=(-1, -2)) / np.nansum(w)) - - assert all(relsqres < 1e-3) - - -def test_time_ordered_map_sim(): - time_evolving_sun_path = fetch("maps/sun.h5") - input_map = maria.map.load(filename=time_evolving_sun_path, nu=100e9, t=1.7e9 + np.linspace(0, 180, 16)) - plan = maria.Plan.generate( - start_time=1.7e9, - duration=180, - scan_center=(input_map.center), - scan_options={"radius": 0.25}, - ) - sim = maria.Simulation(instrument="test/1deg", site="cerro_toco", plans=plan, map=input_map) - tods = sim.run() - tods[0].to("K_RJ").plot() - - plt.close("all") diff --git a/maria/tests/map/test_recover_map.py b/maria/tests/map/test_recover_map.py new file mode 100644 index 00000000..e9e45279 --- /dev/null +++ b/maria/tests/map/test_recover_map.py @@ -0,0 +1,69 @@ +from __future__ import annotations + +import maria +import matplotlib.pyplot as plt +import numpy as np +import pytest +from maria.instrument import Band +from maria.io import fetch +from maria.map import all_maps +from maria.mappers import BinMapper + +plt.close("all") + + +def test_recover_map(): + f090 = Band(center=90e9, width=30e9) + f150 = Band(center=150e9, width=40e9) + f220 = Band(center=220e9, width=50e9) + + map_filename = fetch("maps/cluster1.fits") + input_map = maria.map.load(filename=map_filename, nu=150e9).to("K_RJ")[..., ::8, ::8] + + input_map.data -= input_map.data.mean().compute() + map_width = input_map.width.degrees + + array = { + "field_of_view": map_width / 2, + "primary_size": 1000, + "n": 300, + "bands": [f090, f150, f220], + } + + instrument = maria.get_instrument(array=array) + + planner = maria.Planner(target=input_map, site="cerro_toco", constraints={"el": (40, 90)}) + + plans = planner.generate_plans( + total_duration=60, sample_rate=50, scan_pattern="daisy", scan_options={"radius": input_map.width.deg / 3} + ) + + sim = maria.Simulation( + instrument, + plans=plans, + map=input_map, + site="cerro_toco", + noise=False, + ) + + tods = sim.run() + + mapper = BinMapper( + center=input_map.center, + frame=input_map.frame, + width=input_map.width, + resolution=input_map.x_res, + degrees=False, + tods=tods, + ) + + mapper.add_tods(tods) + output_map = mapper.run() + + m0 = input_map.data[0, 0, 0].compute() + m1 = output_map.data[0, :].compute() + w = output_map.weight[-1, 0].compute() + + relsqres = np.sqrt(np.nansum(w * (m1 - m0) ** 2, axis=(-1, -2)) / np.nansum(w)) + + assert all(relsqres < 1e-3) diff --git a/maria/tests/sim/test_multifrequency.py b/maria/tests/sim/test_multifrequency.py new file mode 100644 index 00000000..397842e6 --- /dev/null +++ b/maria/tests/sim/test_multifrequency.py @@ -0,0 +1,59 @@ +from __future__ import annotations + +import maria +import matplotlib.pyplot as plt +import numpy as np +from maria import Planner, Simulation, all_instruments +from maria.io import fetch +from maria.map import ProjectionMap +from maria.mappers import BinMapper +from maria.utils import read_yaml + + +def test_polarized_map_sim(): + nu = [90e9, 150e9, 220e9] + data = np.random.standard_normal((len(nu), 100, 100)) + + multifrequency_map = ProjectionMap(data=data, width=1e0, nu=nu, center=(0, -30), frame="ra_dec") + + planner = Planner(target=multifrequency_map, site="llano_de_chajnantor", constraints={"el": (60, 90)}) + plans = planner.generate_plans(total_duration=10, sample_rate=50) # in Hz + + plans[0].plot() + print(plans) + + sim = Simulation( + instrument="test/1deg", + site="llano_de_chajnantor", + plans=plans, + atmosphere="2d", + map=multifrequency_map, + ) + + tod = sim.run()[0] + + mapper = BinMapper( + center=(0, -23), + stokes="IQUV", + frame="ra/dec", + width=1, + height=1, + resolution=1 / 256, + tod_preprocessing={ + "window": {"name": "tukey", "kwargs": {"alpha": 0.1}}, + "remove_spline": {"knot_spacing": 30, "remove_el_gradient": True}, + "remove_modes": {"modes_to_remove": 1}, + }, + map_postprocessing={ + "gaussian_filter": {"sigma": 1}, + "median_filter": {"size": 1}, + }, + units="mK_RJ", + tods=[tod], + ) + + output_map = mapper.run() + + output_map.to("Jy/pixel").plot() + + plt.close("all") diff --git a/maria/tests/sim/test_sims.py b/maria/tests/sim/test_pipeline.py similarity index 100% rename from maria/tests/sim/test_sims.py rename to maria/tests/sim/test_pipeline.py diff --git a/maria/tests/sim/test_time_evolving.py b/maria/tests/sim/test_time_evolving.py new file mode 100644 index 00000000..364ecefb --- /dev/null +++ b/maria/tests/sim/test_time_evolving.py @@ -0,0 +1,30 @@ +from __future__ import annotations + +import maria +import matplotlib.pyplot as plt +import numpy as np +import pytest +from maria.instrument import Band +from maria.io import fetch +from maria.map import all_maps +from maria.mappers import BinMapper + +plt.close("all") + + +def test_time_ordered_map_sim(): + time_evolving_sun_path = fetch("maps/sun.h5") + input_map = maria.map.load(filename=time_evolving_sun_path, nu=100e9, t=1.7e9 + np.linspace(0, 180, 16)) + plan = maria.Plan.generate( + start_time=1.7e9, + duration=180, + scan_center=(input_map.center), + scan_options={"radius": 0.25}, + ) + sim = maria.Simulation(instrument="test/1deg", site="cerro_toco", plans=plan, map=input_map) + tods = sim.run() + + mapper = BinMapper(tods=tods, timestep=60) + mapper.run() + + plt.close("all") diff --git a/maria/utils/__init__.py b/maria/utils/__init__.py index aab0e36b..7815a34d 100644 --- a/maria/utils/__init__.py +++ b/maria/utils/__init__.py @@ -89,6 +89,7 @@ def generate_fourier_noise(nx: float = 1024, ny: float = 1024, k0: float = 5e0, def unpack_implicit_slice(key, ndims): + key = key if isinstance(key, tuple) else tuple(key) explicit_slices = [] for s in key: if s == Ellipsis: