From d4aa75e4e2f81459857b1b5c964b0870fe9ae163 Mon Sep 17 00:00:00 2001 From: Thomas Morris Date: Sat, 28 Sep 2024 14:20:04 -0400 Subject: [PATCH] fix dtype conversion bugs --- .github/workflows/testing.yml | 2 +- README.rst | 94 +++++++++++---------- docs/source/conf.py | 2 +- docs/source/index.rst | 2 +- maria/base/__init__.py | 20 +++-- maria/coords/coordinates.py | 2 +- maria/plan/__init__.py | 4 +- maria/plotting/__init__.py | 112 +++++++++++++++++++++++++ maria/sim/atmosphere.py | 8 +- maria/sim/noise.py | 2 +- maria/sim/simulation.py | 19 ++--- maria/tests/test_simulation.py | 2 +- maria/tests/test_tod.py | 2 +- maria/tod/tod.py | 149 ++++++++++----------------------- pyproject.toml | 2 +- 15 files changed, 243 insertions(+), 179 deletions(-) create mode 100644 maria/plotting/__init__.py diff --git a/.github/workflows/testing.yml b/.github/workflows/testing.yml index 20b9df6..3d726d0 100644 --- a/.github/workflows/testing.yml +++ b/.github/workflows/testing.yml @@ -12,7 +12,7 @@ jobs: runs-on: ubuntu-latest strategy: matrix: - python-version: ["3.9", "3.11"] + python-version: ["3.9", "3.12"] fail-fast: false defaults: diff --git a/README.rst b/README.rst index 9d7ab49..2844420 100644 --- a/README.rst +++ b/README.rst @@ -1,50 +1,58 @@ maria ===== +.. image:: https://github.com/thomaswmorris/maria/actions/workflows/testing.yml/badge.svg + :target: https://github.com/thomaswmorris/maria/actions/workflows/testing.yml + +.. image:: https://img.shields.io/pypi/v/maria.svg + :target: https://pypi.python.org/pypi/maria + .. image:: ./docs/source/_static/cloud.gif :width: 256px :alt: StreamPlayer -`Oh, maria blows the stars around / and sends the clouds -a-flyin’ `_ - -``maria`` is a python-based package that simulates turbulent atmospheric -emission using a auto-regressive Gaussian process framework, for -applications in observational astronomy. Tutorials for installation and -usage can be found in the `documentation `_. - -Background ----------- - -Atmospheric modeling is an important step in both experiment design and -subsequent data analysis for ground-based cosmological telescopes -observing the cosmic microwave background (CMB). The next generation of -ground-based CMB experiments will be marked by a huge increase in data -acquisition: telescopes like `AtLAST `_ and -`CMB-S4 `_ will consist of hundreds of thousands of -superconducting polarization-sensitive bolometers sampling the sky. This -necessitates new methods of efficiently modeling and simulating -atmospheric emission at small angular resolutions, with algorithms than -can keep up with the high throughput of modern telescopes. - -maria simulates layers of turbulent atmospheric emission according to a -statistical model derived from observations of the atmosphere in the -Atacama Desert, from the `Atacama Cosmology Telescope -(ACT) `_ and the `Atacama -B-Mode Search (ABS) `_. It -uses a sparse-precision auto-regressive Gaussian process algorithm that -allows for both fast simulation of high-resolution atmosphere, as well -as the ability to simulate arbitrarily long periods of atmospheric -evolution. - -Methodology ------------ - -``maria`` auto-regressively simulates an multi-layeed two-dimensional -“integrated” atmospheric model that is much cheaper to compute than a -three-dimensional model, which can effectively describe time-evolving -atmospheric emission. maria can thus effectively simulate correlated -atmospheric emission for in excess of 100,000 detectors observing the -sky concurrently, at resolutions as fine as one arcminute. The -atmospheric model used is detailed -`here `_. +*maria blows the stars around / and sends the clouds a-flyin’* + +.. raw:: html + + +``maria`` is a complete simulator of ground-based millimeter- and submillimeter-wave telescopes. Tutorials for installation and usage can be found in the `documentation `_. + +.. Background +.. ---------- + +.. Atmospheric modeling is an important step in both experiment design and +.. subsequent data analysis for ground-based cosmological telescopes +.. observing the cosmic microwave background (CMB). The next generation of +.. ground-based CMB experiments will be marked by a huge increase in data +.. acquisition: telescopes like `AtLAST `_ and +.. `CMB-S4 `_ will consist of hundreds of thousands of +.. superconducting polarization-sensitive bolometers sampling the sky. This +.. necessitates new methods of efficiently modeling and simulating +.. atmospheric emission at small angular resolutions, with algorithms than +.. can keep up with the high throughput of modern telescopes. + +.. maria simulates layers of turbulent atmospheric emission according to a +.. statistical model derived from observations of the atmosphere in the +.. Atacama Desert, from the `Atacama Cosmology Telescope +.. (ACT) `_ and the `Atacama +.. B-Mode Search (ABS) `_. It +.. uses a sparse-precision auto-regressive Gaussian process algorithm that +.. allows for both fast simulation of high-resolution atmosphere, as well +.. as the ability to simulate arbitrarily long periods of atmospheric +.. evolution. + +.. Methodology +.. ----------- + +.. ``maria`` auto-regressively simulates an multi-layeed two-dimensional +.. “integrated” atmospheric model that is much cheaper to compute than a +.. three-dimensional model, which can effectively describe time-evolving +.. atmospheric emission. maria can thus effectively simulate correlated +.. atmospheric emission for in excess of 100,000 detectors observing the +.. sky concurrently, at resolutions as fine as one arcminute. The +.. atmospheric model used is detailed +.. `here `_. diff --git a/docs/source/conf.py b/docs/source/conf.py index 7a0989a..523e68a 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -47,4 +47,4 @@ autoapi_type = "python" autoapi_dirs = ["./../../maria"] -nbsphinx_execute = "never" +nbsphinx_execute = "always" diff --git a/docs/source/index.rst b/docs/source/index.rst index 03a7810..d47ce4b 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -9,7 +9,7 @@ maria .. `Oh, maria blows the stars around / and sends the clouds .. a-flyin’ `_ -*Oh, maria blows the stars around / and sends the clouds a-flyin’* +*maria blows the stars around / and sends the clouds a-flyin’* .. raw:: html diff --git a/maria/base/__init__.py b/maria/base/__init__.py index b97893e..75f2059 100644 --- a/maria/base/__init__.py +++ b/maria/base/__init__.py @@ -55,15 +55,17 @@ class BaseSimulation: def __init__( self, - instrument: Union[Instrument, str] = "default", - plan: Union[Plan, str] = "stare", - site: Union[Site, str] = "default", + instrument: Union[Instrument, str], + plan: Union[Plan, str], + site: Union[Site, str], verbose=False, + dtype=np.float32, **kwargs, ): if hasattr(self, "boresight"): return + self.dtype = dtype self.verbose = verbose parsed_sim_kwargs = parse_sim_kwargs(kwargs, master_params) @@ -131,14 +133,22 @@ def __init__( def _run(self): raise NotImplementedError() - def run(self, dtype=np.float32): + def run(self): self.data = {} # Simulate all the junk self._run() + tod_data = {} + + for k, data in self.data.items(): + # scaling floats doesn't make them more accurate, unless they're huge or tiny + offset = data.mean(axis=-1)[..., None] + # scale = data.std(axis=-1)[..., None] + tod_data[k] = {"data": (data - offset).astype(self.dtype), "offset": offset} + tod = TOD( - data={k: v.astype(dtype) for k, v in self.data.items()}, + data=tod_data, dets=self.instrument.dets, coords=self.coords, units="pW", diff --git a/maria/coords/coordinates.py b/maria/coords/coordinates.py index b1bd3c3..e6b411a 100644 --- a/maria/coords/coordinates.py +++ b/maria/coords/coordinates.py @@ -81,7 +81,7 @@ def __init__( earth_location: EarthLocation = DEFAULT_EARTH_LOCATION, frame: str = "ra_dec", distributed: bool = True, - dtype=np.float32, + dtype=np.float64, ): ref_time = ttime.monotonic() diff --git a/maria/plan/__init__.py b/maria/plan/__init__.py index 166235f..d58cc0d 100644 --- a/maria/plan/__init__.py +++ b/maria/plan/__init__.py @@ -127,8 +127,8 @@ def __post_init__(self): x_scan_offsets_radians, y_scan_offsets_radians ].T - # add jitter, should be well sub-arcsecond - self.scan_offsets_radians += 1e-7 * np.random.standard_normal( + # add 0.1 arcseconds of jitter + self.scan_offsets_radians += np.radians(0.1 / 3600) * np.random.standard_normal( size=self.scan_offsets_radians.shape ) diff --git a/maria/plotting/__init__.py b/maria/plotting/__init__.py new file mode 100644 index 0000000..a69061b --- /dev/null +++ b/maria/plotting/__init__.py @@ -0,0 +1,112 @@ +from itertools import cycle + +import matplotlib.pyplot as plt +import numpy as np +import scipy as sp +from matplotlib.patches import Patch + +from ..utils.signal import detrend as detrend_signal + + +def plot_tod( + tod, + detrend: bool = True, + n_freq_bins: int = 256, + max_dets: int = 100, + lw: float = 1e0, + fontsize: float = 8, +): + fig, axes = plt.subplots( + ncols=2, nrows=len(tod.fields), sharex="col", figsize=(8, 4), dpi=160 + ) + gs = axes[0, 0].get_gridspec() + + plt.subplots_adjust(top=0.99, bottom=0.01, hspace=0, wspace=0.02) + + for ax in axes[:, -1]: + ax.remove() + ps_ax = fig.add_subplot(gs[:, -1]) + + handles = [] + colors = cycle(plt.rcParams["axes.prop_cycle"].by_key()["color"]) + + for i, field in enumerate(tod.fields): + data = tod.get_field(field) + + tod_ax = axes[i, 0] + + for band_name in np.unique(tod.dets.band_name): + color = next(colors) + + band_mask = tod.dets.band_name == band_name + signal = data[band_mask] + + if band_mask.sum() > max_dets: + signal = signal[ + np.random.choice(np.arange(len(signal)), max_dets, replace=False) + ] + + detrended_signal = detrend_signal(signal, order=1) + if detrend: + signal = detrended_signal + + f, ps = sp.signal.periodogram(detrended_signal, fs=tod.fs, window="tukey") + + f_bins = np.geomspace(f[1], f[-1], n_freq_bins) + f_mids = np.sqrt(f_bins[1:] * f_bins[:-1]) + + binned_ps = sp.stats.binned_statistic( + f, ps.mean(axis=0), bins=f_bins, statistic="mean" + )[0] + + use = binned_ps > 0 + + ps_ax.plot( + f_mids[use], + binned_ps[use], + lw=lw, + color=color, + ) + tod_ax.plot( + tod.time, + signal[0], + lw=lw, + color=color, + ) + + handles.append(Patch(color=color, label=f"{field} ({band_name})")) + + tod_ax.set_xlim(tod.time.min(), tod.time.max()) + + # if tod.units == "K_RJ": + # ylabel = f"{field} [K]" + # elif tod.units == "K_CMB": + # ylabel = rf"{field} [K]" + # elif tod.units == "pW": + # ylabel = f"{field} [pW]" + + # tod_ax.set_ylabel(tod.units) + + label = tod_ax.text( + 0.01, + 0.99, + f"{field} [{tod.units}]", + fontsize=fontsize, + ha="left", + va="top", + transform=tod_ax.transAxes, + ) + label.set_bbox(dict(facecolor="white", alpha=0.8)) + + # if i + 1 < n_fields: + # tod_ax.set_xticklabels([]) + + tod_ax.set_xlabel("Time [s]", fontsize=fontsize) + + ps_ax.yaxis.tick_right() + ps_ax.yaxis.set_label_position("right") + ps_ax.set_xlabel("Frequency [Hz]", fontsize=fontsize) + ps_ax.set_ylabel(f"[{tod.units}$^2$/Hz]", fontsize=fontsize) + ps_ax.legend(handles=handles, loc="upper right", fontsize=fontsize) + ps_ax.loglog() + ps_ax.set_xlim(f_mids.min(), f_mids.max()) diff --git a/maria/sim/atmosphere.py b/maria/sim/atmosphere.py index 87aa25f..da87adb 100644 --- a/maria/sim/atmosphere.py +++ b/maria/sim/atmosphere.py @@ -67,7 +67,7 @@ def _compute_atmospheric_emission(self): bounds_error=False, fill_value="extrapolate", )(self.coords.time) - ).astype(np.float32) + ) def _compute_atmospheric_transmission(self): self.atmosphere.transmission = da.zeros_like(self.atmosphere.zenith_scaled_pwv) @@ -129,7 +129,7 @@ def _compute_atmospheric_transmission(self): bounds_error=False, fill_value="extrapolate", )(self.coords.time) - ).astype(np.float32) + ) # if units == "F_RJ": # Fahrenheit Rayleigh-Jeans 🇺🇸 # self._simulate_atmospheric_emission(self, units="K_RJ") @@ -244,7 +244,7 @@ def _classic_simulate_atmospheric_emission(self, units="K_RJ"): if units == "K_RJ": # Kelvin Rayleigh-Jeans self._simulate_atmospheric_fluctuations() self.data["atmosphere"] = np.empty( - (self.instrument.n_dets, self.plan.n_time), dtype=np.float32 + (self.instrument.n_dets, self.plan.n_time) ) bands = ( @@ -291,7 +291,7 @@ def _classic_simulate_atmospheric_emission(self, units="K_RJ"): ) self.atmospheric_transmission = np.empty( - (self.instrument.n_dets, self.plan.n_time), dtype=np.float32 + (self.instrument.n_dets, self.plan.n_time) ) # to make a new progress bar diff --git a/maria/sim/noise.py b/maria/sim/noise.py index 3c04f7a..20033ea 100644 --- a/maria/sim/noise.py +++ b/maria/sim/noise.py @@ -12,7 +12,7 @@ def _run(self): def _simulate_noise(self): self.data["noise"] = da.from_array( - np.zeros((self.instrument.n_dets, self.plan.n_time), dtype=np.float32) + np.zeros((self.instrument.n_dets, self.plan.n_time)) ) bands = tqdm( diff --git a/maria/sim/simulation.py b/maria/sim/simulation.py index 7b43f97..7b4d8f3 100644 --- a/maria/sim/simulation.py +++ b/maria/sim/simulation.py @@ -39,7 +39,7 @@ def from_config(cls, config: dict = {}, **params): def __init__( self, instrument: tuple[Instrument, str] = "default", - plan: tuple[Plan, str] = "daisy", + plan: tuple[Plan, str] = "one_minute_zenith_stare", site: tuple[Site, str] = "hoagie_haven", atmosphere: tuple[Atmosphere, str] = None, cmb: tuple[CMB, str] = None, @@ -129,17 +129,10 @@ def __init__( raise ValueError(f"Invalid value for cmb: '{cmb}'.") def _run(self): - # number of bands are lost here - if self.noise: - self._simulate_noise() - if hasattr(self, "atmosphere"): - if self.atmosphere.model == "2d-classic": - self._classic_simulate_atmospheric_emission() - else: - self._simulate_atmosphere() - self._compute_atmospheric_emission() - self._compute_atmospheric_transmission() + self._simulate_atmosphere() + self._compute_atmospheric_emission() + self._compute_atmospheric_transmission() if hasattr(self, "cmb"): self._simulate_cmb_emission() @@ -147,6 +140,10 @@ def _run(self): if hasattr(self, "map"): self._sample_maps() + # number of bands are lost here + if self.noise: + self._simulate_noise() + # scale the noise so that there is if hasattr(self, "atmospheric_transmission"): for k in ["cmb", "map"]: diff --git a/maria/tests/test_simulation.py b/maria/tests/test_simulation.py index ae4373a..f0106c5 100644 --- a/maria/tests/test_simulation.py +++ b/maria/tests/test_simulation.py @@ -32,7 +32,7 @@ def test_complete_sim(instrument, site, plan_config): tod = sim.run() for field in ["atmosphere", "cmb"]: - if np.isnan(tod.data[field]).any(): + if np.isnan(tod.get_field(field)).any(): raise ValueError(f"There are NaNs in the '{field}' field.") tod = tod.to("K_RJ") diff --git a/maria/tests/test_tod.py b/maria/tests/test_tod.py index 08672ba..77c1292 100644 --- a/maria/tests/test_tod.py +++ b/maria/tests/test_tod.py @@ -20,4 +20,4 @@ def test_TOD(): noise = generate_noise_with_knee(t=time, n=n, NEP=0.01, knee=0.5) - tod = TOD(data=dict(noise=noise), coords=coords) + tod = TOD(data=dict(noise=dict(data=noise)), coords=coords) diff --git a/maria/tod/tod.py b/maria/tod/tod.py index e7107b5..894bc92 100644 --- a/maria/tod/tod.py +++ b/maria/tod/tod.py @@ -5,16 +5,15 @@ import dask.array as da import h5py import matplotlib as mpl # noqa -import matplotlib.pyplot as plt import numpy as np import pandas as pd import scipy as sp from astropy.io import fits -from matplotlib.gridspec import GridSpec from maria import utils from ..coords import Coordinates +from ..plotting import plot_tod class TOD: @@ -43,6 +42,7 @@ def __init__( units: str = "K_RJ", dets: pd.DataFrame = None, abscal: float = 1.0, + distributed: bool = True, dtype=np.float32, ): self.weight = weight @@ -55,11 +55,12 @@ def __init__( self.data = {} - for field, data in data.items(): - self.data[field] = ( - data if isinstance(data, da.Array) else da.from_array(data) - ) - self.data[field] = self.data[field].astype(dtype) + for field, field_data in data.items(): + self.data[field] = field_data + d = field_data["data"] + + if distributed and not isinstance(d, da.Array): + self.data[field]["data"] = da.from_array(d) # sort them alphabetically self.data = {k: self.data[k] for k in sorted(list(self.fields))} @@ -67,12 +68,37 @@ def __init__( if self.weight is None: self.weight = da.ones_like(self.signal) + def get_field(self, field: str): + if field not in self.fields: + raise ValueError(f"Field '{field}' not found.") + + d = self.data[field]["data"] + + if "scale" in self.data[field]: + d *= self.data[field]["scale"] + + if "offset" in self.data[field]: + d += self.data[field]["offset"] + + return d + def to(self, units: str): cal = self.dets.cal(f"{self.units} -> {units}") + cal_data = {} + for field in self.fields: + cal_data[field] = self.data[field] + scale = ( + cal_data[field].get( + "scale", np.ones(cal_data[field]["data"].shape[:-1]) + ) + * cal + )[:, None] + cal_data[field]["scale"] = scale + return TOD( coords=self.coords, - data={k: cal[:, None] * v for k, v in self.data.items()}, + data=cal_data, units=units, dets=self.dets, abscal=self.abscal, @@ -89,7 +115,7 @@ def fields(self) -> list: @functools.cached_property def signal(self) -> da.Array: - return sum(self.data.values()) + return sum([self.get_field(field) for field in self.fields]) @functools.cached_property def boresight(self): @@ -161,9 +187,8 @@ def subset( return TOD( data={ - field: data[det_mask] - for field, data in self.data.items() - if field in fields + field: {k: v[det_mask] for k, v in self.data[field].items()} + for field in self.fields }, weight=self.weight[det_mask], coords=subset_coords, @@ -223,7 +248,7 @@ def process(self, **kwargs): D -= A.T @ B return TOD( - data={"data": D}, + data={"total": {"data": D}}, weight=W, coords=self.coords, units=self.units, @@ -370,103 +395,15 @@ def to_hdf(self, fname): f.createdataset(fname) def plot(self, detrend=True, mean=True, n_freq_bins: int = 256): - # def format_axes(fig): - # for i, ax in enumerate(fig.axes): - # ax.text(0.5, 0.5, "ax%d" % (i+1), va="center", ha="center") - # ax.tick_params(labelbottom=False, labelleft=False) - - fig = plt.figure(figsize=(8, 5), dpi=160, constrained_layout=True) - - gs = GridSpec( - len(self.fields), - 2, - figure=fig, + plot_tod( + self, + detrend=detrend, + n_freq_bins=n_freq_bins, ) - ps_ax = fig.add_subplot(gs[:, 1]) - - i = 0 - colors = plt.rcParams["axes.prop_cycle"].by_key()["color"] - - n_fields = len(self.data) - - for i, (field, data) in enumerate(self.data.items()): - tod_ax = fig.add_subplot(gs[i, 0]) - - for band_name in np.unique(self.dets.band_name): - color = colors[i] - i = (i + 1) % len(colors) - - band_mask = self.dets.band_name == band_name - d = data[band_mask] - - if detrend: - d = sp.signal.detrend(d) - - f, ps = sp.signal.periodogram( - sp.signal.detrend(d), fs=self.fs, window="tukey" - ) - - f_bins = np.geomspace(f[1], f[-1], n_freq_bins) - f_mids = np.sqrt(f_bins[1:] * f_bins[:-1]) - - binned_ps = sp.stats.binned_statistic( - f, ps.mean(axis=0), bins=f_bins, statistic="mean" - )[0] - - use = binned_ps > 0 - - ps_ax.plot( - f_mids[use], - binned_ps[use], - lw=1e0, - color=color, - label=f"{band_name} {field}", - ) - tod_ax.plot( - self.time, - d[0], - lw=5e-1, - label=f"{band_name} {field}", - color=color, - ) - - tod_ax.set_xlim(self.time.min(), self.time.max()) - - if self.units == "K_RJ": - ylabel = f"{field} [K]" - elif self.units == "K_CMB": - ylabel = rf"{field} [K]" - elif self.units == "pW": - ylabel = f"{field} [pW]" - - tod_ax.set_ylabel(ylabel) - - if i + 1 < n_fields: - tod_ax.set_xticklabels([]) - - ps_ax.yaxis.tick_right() - ps_ax.yaxis.set_label_position("right") - - if self.units == "K_RJ": - pslabel = rf"{field} [K$^2$/Hz]" - elif self.units == "K_CMB": - pslabel = rf"{field} [K$^2$/Hz]" - elif self.units == "pW": - pslabel = rf"{field} [pW$^2$/Hz]" - - ps_ax.set_xlabel("T [s]") - ps_ax.set_ylabel(pslabel) - - ps_ax.legend() - ps_ax.loglog() - ps_ax.set_xlim(f_mids.min(), f_mids.max()) - - ps_ax.set_xlabel("Frequency [Hz]") - def __getattr__(self, attr): if attr in self.fields: - return self.data[attr] + return self.get_field(attr) raise AttributeError(f"No attribute named '{attr}'.") def __repr__(self): diff --git a/pyproject.toml b/pyproject.toml index 5c514fa..7a1a9d3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -5,7 +5,7 @@ build-backend = "hatchling.build" [project] name = "maria" dynamic = ["version"] -description = "Ground-based telescope simulations" +description = "A simulator for ground-based millimeter- and submillimeter-wave telescopes." readme = { file = "README.rst", content-type = "text/x-rst" } dependencies = [ "astropy",