def _observe_sso(self, sso_az, sso_el, sso_dist, sso_dia, tod, comm,
prefix):
"""
Observe the SSO with each detector in tod
"""
log = Logger.get()
rank = 0
if comm is not None:
rank = comm.rank
tmr = Timer()
if self._report_timing:
if comm is not None:
comm.Barrier()
tmr.start()
nsamp = tod.local_samples[1]
if rank == 0:
log.info("{}Observing the SSO signal".format(prefix))
for det in tod.local_dets:
# Cache the output signal
cachename = "{}_{}".format(self._out, det)
if tod.cache.exists(cachename):
ref = tod.cache.reference(cachename)
else:
ref = tod.cache.create(cachename, np.float64, (nsamp, ))
try:
# Some TOD classes provide a shortcut to Az/El
az, el = tod.read_azel(detector=det)
except Exception as e:
azelquat = tod.read_pntg(detector=det, azel=True)
# Convert Az/El quaternion of the detector back into
# angles for the simulation.
theta, phi = qa.to_position(azelquat)
# Azimuth is measured in the opposite direction
# than longitude
az = 2 * np.pi - phi
el = np.pi / 2 - theta
beam, radius = self._get_beam_map(det, sso_dia)
# Interpolate the beam map at appropriate locations
x = (az - sso_az) * np.cos(el)
y = el - sso_el
r = np.sqrt(x**2 + y**2)
good = r < radius
sig = beam(x[good], y[good], grid=False)
ref[:][good] += sig
del ref, sig, beam
if self._report_timing:
if comm is not None:
comm.Barrier()
if rank == 0:
tmr.stop()
tmr.report("{}OpSimSSO: Observe signal".format(prefix))
return
def export_TOD(args,
comm,
data,
totalname,
schedules,
other=None,
verbose=True):
if args.export is None:
return
log = Logger.get()
timer = Timer()
# Only import spt3g if we are writing out so3g files
from spt3g import core as core3g
from ..data.toast_export import ToastExport
path = os.path.abspath(args.export)
key = args.export_key
if key is not None:
prefix = "{}_{}".format(args.bands, key)
det_groups = {}
for schedule in schedules:
for (
det_name,
det_data,
) in schedule.telescope.focalplane.detector_data.items():
value = det_data[key]
if value not in det_groups:
det_groups[value] = []
det_groups[value].append(det_name)
else:
prefix = args.bands
det_groups = None
if comm.world_rank == 0 and verbose:
log.info("Exporting data to directory tree at {}".format(path))
timer.start()
export = ToastExport(
path,
prefix=prefix,
use_intervals=True,
cache_name=totalname,
cache_copy=other,
mask_flag_common=TODGround.TURNAROUND,
filesize=2**30,
units=core3g.G3TimestreamUnits.Tcmb,
detgroups=det_groups,
compress=args.compress,
)
export.exec(data)
if comm.comm_world is not None:
comm.comm_world.Barrier()
timer.stop()
if comm.world_rank == 0 and verbose:
timer.report("Wrote simulated data to {}:{}" "".format(path, "total"))
return
def apply_flag_sso(args, comm, data, verbose=True):
if args.flag_sso is None:
return
if comm.world_rank == 0 and verbose:
print(f"Flagging SSO:s", flush=True)
timer = Timer()
timer.start()
sso_names = []
sso_radii = []
for arg in args.flag_sso:
sso_name, sso_radius = arg.split(",")
sso_radius = np.radians(float(sso_radius) / 60)
sso_names.append(sso_name)
sso_radii.append(sso_radius)
flag_sso = OpFlagSSO(sso_names, sso_radii, flag_mask=args.flag_sso_mask)
flag_sso.exec(data)
if comm.world_rank == 0 and verbose:
timer.report_clear(f"Flag {sso_names}")
return
def create_observations(args, comm, schedules):
""" Create and distribute TOAST observations for every CES in schedules.
"""
log = Logger.get()
timer = Timer()
timer.start()
data = Data(comm)
# Loop over the schedules, distributing each schedule evenly across
# the process groups. For now, we'll assume that each schedule has
# the same number of operational days and the number of process groups
# matches the number of operational days. Relaxing these constraints
# will cause the season break to occur on different process groups
# for different schedules and prevent splitting the communicator.
for schedule in schedules:
telescope = schedule.telescope
all_ces = schedule.ceslist
nces = len(all_ces)
breaks = get_breaks(comm, all_ces, args)
groupdist = distribute_uniform(nces, comm.ngroups, breaks=breaks)
group_firstobs = groupdist[comm.group][0]
group_numobs = groupdist[comm.group][1]
for ices in range(group_firstobs, group_firstobs + group_numobs):
obs = create_observation(args, comm, telescope, all_ces[ices])
data.obs.append(obs)
if comm.comm_world is None or comm.comm_group.rank == 0:
log.info("Group # {:4} has {} observations.".format(comm.group, len(data.obs)))
if len(data.obs) == 0:
raise RuntimeError(
"Too many tasks. Every MPI task must "
"be assigned to at least one observation."
)
if comm.comm_world is not None:
comm.comm_world.barrier()
timer.stop()
if comm.world_rank == 0:
timer.report("Simulated scans")
# Split the data object for each telescope for separate mapmaking.
# We could also split by site.
if len(schedules) > 1:
telescope_data = data.split("telescope")
if len(telescope_data) == 1:
# Only one telescope available
telescope_data = []
else:
telescope_data = []
telescope_data.insert(0, ("all", data))
return data, telescope_data
def get_elevation_noise(args, comm, data, key="noise"):
""" Insert elevation-dependent noise
"""
timer = Timer()
timer.start()
# fsample = args.sample_rate
for obs in data.obs:
tod = obs["tod"]
fp = obs["focalplane"]
noise = obs[key]
for det in tod.local_dets:
if det not in noise.keys:
raise RuntimeError(
'Detector "{}" does not have a PSD in the noise object'.
format(det))
A = fp[det]["A"]
C = fp[det]["C"]
psd = noise.psd(det)
try:
# Some TOD classes provide a shortcut to Az/El
_, el = tod.read_azel(detector=det)
except Exception:
azelquat = tod.read_pntg(detector=det, azel=True)
# Convert Az/El quaternion of the detector back into
# angles for the simulation.
theta, _ = qa.to_position(azelquat)
el = np.pi / 2 - theta
el = np.median(el)
# Scale the analytical noise PSD. Pivot is at el = 50 deg.
psd[:] *= (A / np.sin(el) + C)**2
timer.stop()
if comm.world_rank == 0:
timer.report("Elevation noise")
return
def load_focalplane(args, comm):
"""Load focalplane information
"""
timer = Timer()
gain = None
fp = None
if comm.world_rank == 0:
if args.focalplane is None:
# in this case, create a fake detector at the boresight
# with a pure white noise spectrum.
fake = {}
fake["quat"] = np.array([0.0, 0.0, 1.0, 0.0])
fake["fwhm"] = 30.0
fake["fknee"] = 0.0
fake["fmin"] = 1.0e-5
fake["alpha"] = 1.0
fake["NET"] = 1.0
fake["polangle_deg"] = 0
fake["color"] = "r"
fp = {}
fp["bore"] = fake
else:
with open(args.focalplane, "rb") as p:
fp = pickle.load(p)
if args.gain is not None:
gain = {}
with fits.open(args.gain) as f:
gain["TIME"] = np.array(f["TIME"].data["TIME"])
for i_det, det_name in f["DETECTORS"].data["DETECTORS"]:
gain[det_name] = np.array(f["GAINS"].data[i_det, :])
if comm.comm_world is not None:
if args.gain is not None:
gain = comm.comm_world.bcast(gain, root=0)
fp = comm.comm_world.bcast(fp, root=0)
timer.stop()
if comm.world_rank == 0:
timer.report("Create focalplane ({} dets)".format(len(fp.keys())))
if args.debug:
if comm.world_rank == 0:
outfile = os.path.join(args.outdir, "focalplane.png")
set_backend()
dquats = {x: fp[x]["quat"] for x in fp.keys()}
dfwhm = {x: fp[x]["fwhm"] for x in fp.keys()}
plot_focalplane(dquats, 10.0, 10.0, outfile, fwhm=dfwhm)
# For purposes of this simulation, we use detector noise
# weights based on the NET (white noise level). If the destriping
# baseline is too long, this will not be the best choice.
detweights = {}
for d in fp.keys():
net = fp[d]["NET"]
detweights[d] = 1.0 / (args.sample_rate * net * net)
return fp, gain, detweights
def load_observations(args, comm):
"""Load existing data and put it in TOAST observations.
"""
# This import is not at the top of the file to avoid
# loading spt3g through so3g unnecessarily
from ...io.toast_load import load_data
log = Logger.get()
if args.import_obs is not None:
import_obs = args.import_obs.split(",")
else:
import_obs = None
hw, telescope, det_index = get_hardware(args, comm, verbose=True)
focalplane = get_focalplane(args, comm, hw, det_index, verbose=True)
detweights = focalplane.detweights
telescope.focalplane = focalplane
if comm.world_rank == 0:
log.info("Loading TOD from {}".format(args.import_dir))
timer = Timer()
timer.start()
data = load_data(
args.import_dir,
obs=import_obs,
comm=comm,
prefix=args.import_prefix,
dets=hw,
detranks=comm.group_size,
)
if comm.world_rank == 0:
timer.report_clear("Load data")
telescope_data = [("all", data)]
site = telescope.site
focalplane = telescope.focalplane
if args.weather is not None:
weather = Weather(args.weather)
else:
weather = None
for obs in data.obs:
#obs["baselines"] = None
obs["noise"] = focalplane.noise
#obs["id"] = int(ces.mjdstart * 10000)
#obs["intervals"] = tod.subscans
obs["site"] = site.name
obs["site_id"] = site.id
obs["telescope"] = telescope.name
obs["telescope_id"] = telescope.id
obs["fpradius"] = focalplane.radius
obs["weather"] = weather
#obs["start_time"] = ces.start_time
obs["altitude"] = site.alt
#obs["season"] = ces.season
#obs["date"] = ces.start_date
#obs["MJD"] = ces.mjdstart
obs["focalplane"] = focalplane.detector_data
#obs["rising"] = ces.rising
#obs["mindist_sun"] = ces.mindist_sun
#obs["mindist_moon"] = ces.mindist_moon
#obs["el_sun"] = ces.el_sun
return data, telescope_data, detweights
def exec(self, data):
""" Apply the OpSignalSim operator on data
"""
for obs in data.obs:
tod = obs["tod"]
nsamp = tod.local_samples[1]
for det in tod.local_dets:
timer = Timer()
timer.start()
if self._global_rank == 0:
print("Processing {}".format(det), flush=True)
quat = tod.local_pointing(det, margin=self._margin)
self._check_len(quat, nsamp, "detector quaternions")
iquweights = tod.local_weights(det)
self._check_len(iquweights, nsamp, "detector weights")
sampled = self._sample_maps(tod, det, quat, iquweights)
if self._dipoler is not None:
if self._global_rank == 0:
print(" Adding dipole", flush=True)
if self.skip_reproc:
velocity = None
else:
velocity = tod.local_velocity(margin=self._margin)
self._check_len(velocity, nsamp, "velocity")
sampled += self._dipoler.dipole(quat,
velocity=velocity,
det=det)
if self._fsl and not self.skip_reproc:
if self._global_rank == 0:
print(" Adding FSL", flush=True)
local_fsl = tod.local_fsl(det, margin=self._margin)
self._check_len(local_fsl, nsamp, "FSL")
sampled += local_fsl
if self._out is not None:
cachename = "{}_{}".format(self._out, det)
if not tod.cache.exists(cachename):
tod.cache.create((nsamp + 2 * self._margin, ),
dtype=np.float64)
local_signal = tod.local_signal(det,
name=self._out,
margin=self._margin)
self._check_len(iquweights, nsamp, "signal")
if self._add:
local_signal += sampled
else:
local_signal[:] = sampled
del local_signal
timer.stop()
if self._global_rank == 0:
timer.report("Process {}".format(det))
return
def _stage_signal(self, detectors, nsamp, ndet, nodecomm, nread):
""" Stage signal
"""
log = Logger.get()
timer = Timer()
# Determine if we can purge the signal and avoid keeping two
# copies in memory
purge = self._name is not None and self._purge_tod
if not purge:
nread = 1
nodecomm = MPI.COMM_SELF
for iread in range(nread):
nodecomm.Barrier()
timer.start()
if nodecomm.rank % nread == iread:
self._mappraiser_signal = self._cache.create(
"signal", mappraiser.SIGNAL_TYPE, (nsamp * ndet, ))
self._mappraiser_signal[:] = np.nan
global_offset = 0
local_blocks_sizes = []
for iobs, obs in enumerate(self._data.obs):
tod = obs["tod"]
for idet, det in enumerate(detectors):
# Get the signal.
signal = tod.local_signal(det, self._name)
signal_dtype = signal.dtype
offset = global_offset
local_V_size = len(signal)
dslice = slice(idet * nsamp + offset,
idet * nsamp + offset + local_V_size)
self._mappraiser_signal[dslice] = signal
offset += local_V_size
local_blocks_sizes.append(local_V_size)
del signal
# Purge only after all detectors are staged in case some are aliased
# cache.clear() will not fail if the object was already
# deleted as an alias
if purge:
for det in detectors:
cachename = "{}_{}".format(self._name, det)
tod.cache.clear(cachename)
global_offset = offset
local_blocks_sizes = np.array(local_blocks_sizes,
dtype=np.int32)
if self._verbose and nread > 1:
nodecomm.Barrier()
if self._rank == 0:
timer.report_clear("Stage signal {} / {}".format(
iread + 1, nread))
return signal_dtype, local_blocks_sizes
def simulate_hwpss(args, comm, data, mc, name):
if not args.simulate_hwpss:
return
log = Logger.get()
timer = Timer()
timer.start()
hwpssop = OpSimHWPSS(name=name, fname_hwpss=args.hwpss_file, mc=mc)
hwpssop.exec(data)
timer.report_clear("Simulate HWPSS")
return
def convolve_time_constant(args, comm, data, name, verbose=True):
if not args.tau_convolve:
return
log = Logger.get()
timer = Timer()
timer.start()
tauop = OpTimeConst(name=name, tau=args.tau_value, inverse=False)
tauop.exec(data)
timer.report_clear("Convolve time constant")
return
def load_focalplanes(args, comm, schedules, verbose=False):
""" Attach a focalplane to each of the schedules.
Args:
schedules (list) : List of Schedule instances.
Each schedule has two members, telescope
and ceslist, a list of CES objects.
Returns:
detweights (dict) : Inverse variance noise weights for every
detector across all focal planes. In [K_CMB^-2].
They can be used to bin the TOD.
"""
# log = Logger.get()
timer = Timer()
timer.start()
# Load focalplane information
timer1 = Timer()
timer1.start()
hw, telescope, det_index = get_hardware(args, comm, verbose=verbose)
focalplane = get_focalplane(args, comm, hw, det_index, verbose=verbose)
telescope.focalplane = focalplane
if comm.world_rank == 0 and verbose:
timer1.report_clear("Collect focaplane information")
for schedule in schedules:
# Replace the telescope created from reading the observing schedule but
# keep the weather object
weather = schedule.telescope.site.weather
schedule.telescope = telescope
schedule.telescope.site.weather = weather
detweights = telescope.focalplane.detweights
timer.stop()
if (comm.comm_world is None or comm.world_rank == 0) and verbose:
timer.report("Loading focalplane")
return detweights
def apply_filter(args, data):
""" Apply extra filter to signal
"""
if args.filterfile is None:
return
timer = Timer()
timer.start()
convolver = tp.OpConvolvePlanck(args.filterfile)
convolver.exec(data)
data.comm.comm_world.barrier()
timer.stop()
if data.comm.comm_world.rank == 0:
timer.report("Convolve with {}".format(args.filterfile))
return
def load_focalplanes(args, comm, schedules):
""" Attach a focalplane to each of the schedules.
Args:
schedules (list) : List of Schedule instances.
Each schedule has two members, telescope
and ceslist, a list of CES objects.
Returns:
detweights (dict) : Inverse variance noise weights for every
detector across all focal planes. In [K_CMB^-2].
They can be used to bin the TOD.
"""
timer = Timer()
timer.start()
# Load focalplane information
focalplanes = []
if comm.world_rank == 0:
for fpfile in args.focalplane.split(","):
focalplanes.append(
pipeline_tools.Focalplane(
fname_pickle=fpfile,
sample_rate=args.sample_rate,
radius_deg=args.focalplane_radius_deg,
)
)
if comm.comm_world is not None:
focalplanes = comm.comm_world.bcast(focalplanes)
if len(focalplanes) == 1 and len(schedules) > 1:
focalplanes *= len(schedules)
if len(focalplanes) != len(schedules):
raise RuntimeError(
"Number of focalplanes must equal number of schedules or be 1."
)
# Append a focal plane and telescope to each entry in the schedules
# list and assemble a detector weight dictionary that represents all
# detectors in all focalplanes
detweights = {}
for schedule, focalplane in zip(schedules, focalplanes):
schedule.telescope.focalplane = focalplane
detweights.update(schedule.telescope.focalplane.detweights)
timer.stop()
if comm.world_rank == 0:
timer.report("Loading focalplanes")
return detweights
def get_elevation_noise(args, comm, data, key="noise"):
""" Insert elevation-dependent noise
"""
if args.no_elevation_noise:
return
timer = Timer()
timer.start()
# fsample = args.sample_rate
for obs in data.obs:
tod = obs["tod"]
fp = obs["focalplane"]
noise = obs[key]
for det in tod.local_dets:
if det not in noise.keys:
raise RuntimeError(
'Detector "{}" does not have a PSD in the noise object'.
format(det))
A = fp[det]["A"]
C = fp[det]["C"]
psd = noise.psd(det)
# We only consider a small range of samples for the elevation
n = tod.local_samples[1]
istart = max(0, n // 2 - 1000)
istop = min(n, n // 2 + 1000)
try:
# Some TOD classes provide a shortcut to Az/El
el = tod.read_azel(detector=det,
local_start=istart,
n=istop - istart)[1]
except Exception:
azelquat = tod.read_pntg(detector=det,
azel=True,
local_start=istart,
n=istop - istart)
# Convert Az/El quaternion of the detector back into
# angles for the simulation.
theta = qa.to_position(azelquat)[0]
el = np.pi / 2 - theta
el = np.median(el)
# Scale the analytical noise PSD. Pivot is at el = 50 deg.
psd[:] *= (A / np.sin(el) + C)**2
if comm.world_rank == 0:
timer.report_clear("Elevation noise")
return
def rotate_focalplane(args, data, comm):
""" The LAT focalplane projected on the sky rotates as the cryostat
(co-rotator) tilts. Usually the tilt is the same as the observing
elevation to maintain constant angle between the mirror and the cryostat.
This method must be called *before* expanding the detector pointing
from boresight.
"""
log = Logger.get()
timer = Timer()
timer.start()
for obs in data.obs:
if obs["telescope"] != "LAT":
continue
tod = obs["tod"]
cache_name = "corotator_angle_deg"
if tod.cache.exists(cache_name):
corotator_angle = tod.cache.reference(cache_name)
else:
# If a vector of co-rotator angles isn't already cached,
# make one now from the observation metadata. This will
# ensure they get recorded in the so3g files.
corotator_angle = obs["corotator_angle_deg"]
offset, nsample = tod.local_samples
tod.cache.put(cache_name, np.zeros(nsample) + corotator_angle)
el = np.degrees(tod.read_boresight_el())
rot = qa.rotation(
ZAXIS, np.radians(corotator_angle + el + LAT_COROTATOR_OFFSET_DEG)
)
quats = tod.read_boresight()
quats[:] = qa.mult(quats, rot)
try:
# If there are horizontal boresight quaternions, they need
# to be rotated as well.
quats = tod.read_boresight(azel=True)
quats[:] = qa.mult(quats, rot)
except Exception as e:
pass
if comm.comm_world is None or comm.comm_world.rank == 0:
timer.report_clear("Rotate focalplane")
return
def apply_polyfilter(args, comm, data, cache_name=None, verbose=True):
"""Apply the polynomial filter to data under `cache_name`."""
if not args.apply_polyfilter:
return
log = Logger.get()
timer = Timer()
timer.start()
if comm.world_rank == 0 and verbose:
log.info("Polyfiltering signal")
polyfilter = OpPolyFilter(order=args.poly_order,
name=cache_name,
common_flag_mask=args.common_flag_mask)
polyfilter.exec(data)
if comm.comm_world is not None:
comm.comm_world.barrier()
if comm.world_rank == 0 and verbose:
timer.report_clear("Polynomial filtering")
return
def apply_common_mode_filter(args, comm, data, cache_name=None, verbose=True):
"""Apply the common mode filter to data under `cache_name`."""
if not args.apply_common_mode_filter:
return
log = Logger.get()
timer = Timer()
timer.start()
if comm.world_rank == 0 and verbose:
log.info("Common mode filtering signal")
commonfilter = OpCommonModeFilter(
name=cache_name,
common_flag_mask=args.common_flag_mask,
focalplane_key=args.common_mode_filter_key,
)
commonfilter.exec(data)
if comm.world_rank == 0 and verbose:
timer.report_clear("Common mode filtering")
return
def compute_h_n(args, comm, data, verbose=True):
if args.hn_max < args.hn_min:
return
log = Logger.get()
timer = Timer()
timer.start()
hnop = OpHn(
outdir=args.hn_outdir,
outprefix=args.hn_prefix,
nmin=args.hn_min,
nmax=args.hn_max,
common_flag_mask=args.common_flag_mask,
flag_mask=255,
zip_maps=args.hn_zip,
)
hnop.exec(data)
timer.report_clear("Compute h_n")
return
def apply_groundfilter(args, comm, data, cache_name=None, verbose=True):
if not args.apply_groundfilter:
return
log = Logger.get()
timer = Timer()
timer.start()
if comm.world_rank == 0 and verbose:
log.info("Ground-filtering signal")
groundfilter = OpGroundFilter(
filter_order=args.ground_order,
name=cache_name,
common_flag_mask=args.common_flag_mask,
)
groundfilter.exec(data)
if comm.comm_world is not None:
comm.comm_world.barrier()
if comm.world_rank == 0 and verbose:
timer.report_clear("Ground filtering")
return
def demodulate(args,
comm,
data,
name,
detweights=None,
madampars=None,
verbose=True):
if not args.demodulate:
return
log = Logger.get()
timer = Timer()
timer.start()
if detweights is not None:
# Copy the detector weights to demodulated TOD
modulated = [
detname for detname in detweights if "demod" not in detname
]
for detname in modulated:
detweight = detweights[detname]
for demodkey in ["demod0", "demod4r", "demod4i"]:
demod_name = "{}_{}".format(demodkey, detname)
detweights[demod_name] = detweight
del detweights[detname]
if madampars is not None:
# Filtering will affect the high frequency end of the noise PSD
madampars["radiometers"] = False
# Intensity and polarization will be decoupled in the noise matrix
madampars["allow_decoupling"] = True
demod = OpDemod(
name=name,
wkernel=args.demod_wkernel,
fmax=args.demod_fmax,
nskip=args.demod_nskip,
do_2f=args.demod_2f,
)
demod.exec(data)
timer.report_clear("Demodulate")
return
def compute_cadence_map(args, comm, data, verbose=True):
if not args.write_cadence_map:
return
log = Logger.get()
if comm.world_rank == 0:
log.info("Computing cadence map")
timer = Timer()
timer.start()
cadence = OpCadenceMap(
outdir=args.out,
outprefix=args.cadence_map_prefix,
common_flag_mask=args.common_flag_mask,
flag_mask=255,
)
cadence.exec(data)
if comm.world_rank == 0:
timer.report_clear("Compute cadence map")
return
def _load_alm(self):
""" Load the alm expansion and place it in the node-shared memory
"""
if self._rank == 0:
timer = Timer()
timer.start()
alm, mmax = hp.read_alm(self._almfile, return_mmax=True)
nalm = len(alm)
lmax = hp.Alm.getlmax(nalm, mmax)
alm = [alm]
if self._pol:
for hdu in [2, 3]:
alm.append(hp.read_alm(self._almfile, hdu=hdu))
alm = np.vstack(alm)
nalm = len(alm)
# If necessary, truncate the expansion to sufficient lmax
self._lmax = min(lmax, self._lmax)
self._mmax = min(mmax, self._lmax)
if self._lmax < lmax:
sz = hp.Alm.getsize(self._lmax, self._mmax)
new_alm = np.zeros([nalm, sz], dtype=np.complex)
for ell in range(self._lmax + 1):
for m in range(min(ell, self._mmax)):
i = hp.Alm.getidx(self._lmax, ell, m)
j = hp.Alm.getidx(lmax, ell, m)
new_alm[:, i] = alm[:, j]
alm = new_alm
lmax = self._lmax
mmax = self._mmax
# Suppress any primordial monopole or dipole
for ell in range(min(2, lmax + 1)):
for m in range(min(ell + 1, mmax + 1)):
ind = hp.Alm.getidx(lmax, 1, m)
alm[0, ind] = 0
timer.stop()
timer.report("load CMB alm")
else:
alm, lmax, mmax = None, None, None
self._alm = self._comm.bcast(alm)
self._lmax = self._comm.bcast(lmax)
self._mmax = self._comm.bcast(mmax)
return
def compute_crosslinking(args, comm, data, detweights=None, verbose=True):
if not args.write_crosslinking:
return
log = Logger.get()
timer = Timer()
timer.start()
crosslinking = OpCrossLinking(
outdir=args.out,
outprefix=args.crosslinking_prefix,
common_flag_mask=args.common_flag_mask,
flag_mask=255,
zip_maps=args.hn_zip,
rcond_limit=1e-3,
detweights=detweights,
)
crosslinking.exec(data)
timer.report_clear("Compute crosslinking map")
return
def apply_sim_sso(args, comm, data, mc, totalname, verbose=True):
if args.simulate_sso is None:
return
if args.beam_file is None:
raise RuntimeError("Cannot simulate SSOs without a beam file")
timer = Timer()
timer.start()
for sso_name in args.simulate_sso.split(","):
if comm.world_rank == 0 and verbose:
print("Simulating {}".format(sso_name), flush=True)
sim_sso = OpSimSSO(sso_name, args.beam_file, out=totalname)
sim_sso.exec(data)
if comm.world_rank == 0 and verbose:
timer.report_clear("Simulate {}".format(sso_name))
return
def deconvolve_time_constant(args,
comm,
data,
name,
realization=0,
verbose=True):
if not args.tau_deconvolve:
return
log = Logger.get()
timer = Timer()
timer.start()
tauop = OpTimeConst(
name=name,
tau=args.tau_value,
inverse=True,
tau_sigma=args.tau_sigma,
realization=realization,
)
tauop.exec(data)
timer.report_clear("De-convolve time constant")
return
def main():
env = Environment.get()
log = Logger.get()
gt = GlobalTimers.get()
gt.start("toast_satellite_sim (total)")
timer0 = Timer()
timer0.start()
mpiworld, procs, rank, comm = get_comm()
args, comm, groupsize = parse_arguments(comm, procs)
# Parse options
tmr = Timer()
tmr.start()
if comm.world_rank == 0:
os.makedirs(args.outdir, exist_ok=True)
focalplane, gain, detweights = load_focalplane(args, comm)
data = create_observations(args, comm, focalplane, groupsize)
expand_pointing(args, comm, data)
localpix, localsm, subnpix = get_submaps(args, comm, data)
signalname = None
skyname = simulate_sky_signal(args, comm, data, [focalplane], subnpix,
localsm, "signal")
if skyname is not None:
signalname = skyname
diponame = simulate_dipole(args, comm, data, "signal")
if diponame is not None:
signalname = diponame
# Mapmaking.
if not args.use_madam:
if comm.world_rank == 0:
log.info("Not using Madam, will only make a binned map")
npp, zmap = init_binner(args,
comm,
data,
detweights,
subnpix=subnpix,
localsm=localsm)
# Loop over Monte Carlos
firstmc = args.MC_start
nmc = args.MC_count
for mc in range(firstmc, firstmc + nmc):
mctmr = Timer()
mctmr.start()
outpath = os.path.join(args.outdir, "mc_{:03d}".format(mc))
simulate_noise(args, comm, data, mc, "tot_signal", overwrite=True)
# add sky signal
add_signal(args, comm, data, "tot_signal", signalname)
if gain is not None:
timer = Timer()
timer.start()
op_apply_gain = OpApplyGain(gain, name="tot_signal")
op_apply_gain.exec(data)
if comm.world_rank == 0:
timer.report_clear(" Apply gains {:04d}".format(mc))
if mc == firstmc:
# For the first realization, optionally export the
# timestream data. If we had observation intervals defined,
# we could pass "use_interval=True" to the export operators,
# which would ensure breaks in the exported data at
# acceptable places.
output_tidas(args, comm, data, "tot_signal")
output_spt3g(args, comm, data, "tot_signal")
apply_binner(args, comm, data, npp, zmap, detweights, outpath,
"tot_signal")
if comm.world_rank == 0:
mctmr.report_clear(" Map-making {:04d}".format(mc))
else:
# Initialize madam parameters
madampars = setup_madam(args)
# in debug mode, print out data distribution information
if args.debug:
handle = None
if comm.world_rank == 0:
handle = open(os.path.join(args.outdir, "distdata.txt"), "w")
data.info(handle)
if comm.world_rank == 0:
handle.close()
if comm.comm_world is not None:
comm.comm_world.barrier()
if comm.world_rank == 0:
tmr.report_clear("Dumping data distribution")
# Loop over Monte Carlos
firstmc = args.MC_start
nmc = args.MC_count
for mc in range(firstmc, firstmc + nmc):
mctmr = Timer()
mctmr.start()
# create output directory for this realization
outpath = os.path.join(args.outdir, "mc_{:03d}".format(mc))
simulate_noise(args, comm, data, mc, "tot_signal", overwrite=True)
# add sky signal
add_signal(args, comm, data, "tot_signal", signalname)
if gain is not None:
op_apply_gain = OpApplyGain(gain, name="tot_signal")
op_apply_gain.exec(data)
if comm.comm_world is not None:
comm.comm_world.barrier()
if comm.world_rank == 0:
tmr.report_clear(" Apply gains {:04d}".format(mc))
apply_madam(args, comm, data, madampars, outpath, detweights,
"tot_signal")
if comm.comm_world is not None:
comm.comm_world.barrier()
if comm.world_rank == 0:
mctmr.report_clear(" Map-making {:04d}".format(mc))
gt.stop_all()
if comm.comm_world is not None:
comm.comm_world.barrier()
tmr.stop()
tmr.clear()
tmr.start()
alltimers = gather_timers(comm=comm.comm_world)
if comm.world_rank == 0:
out = os.path.join(args.outdir, "timing")
dump_timing(alltimers, out)
tmr.stop()
tmr.report("Gather and dump timing info")
timer0.report_clear("toast_satellite_sim.py")
return
def create_observations(args, comm, focalplane, groupsize):
timer = Timer()
timer.start()
if groupsize > len(focalplane.keys()):
if comm.world_rank == 0:
log.error("process group is too large for the number of detectors")
comm.comm_world.Abort()
# Detector information from the focalplane
detectors = sorted(focalplane.keys())
detquats = {}
detindx = None
if "index" in focalplane[detectors[0]]:
detindx = {}
for d in detectors:
detquats[d] = focalplane[d]["quat"]
if detindx is not None:
detindx[d] = focalplane[d]["index"]
# Distribute the observations uniformly
groupdist = distribute_uniform(args.obs_num, comm.ngroups)
# Compute global time and sample ranges of all observations
obsrange = regular_intervals(
args.obs_num,
args.start_time,
0,
args.sample_rate,
3600 * args.obs_time_h,
3600 * args.gap_h,
)
noise = get_analytic_noise(args, comm, focalplane)
# The distributed timestream data
data = Data(comm)
# Every process group creates its observations
group_firstobs = groupdist[comm.group][0]
group_numobs = groupdist[comm.group][1]
for ob in range(group_firstobs, group_firstobs + group_numobs):
tod = TODSatellite(
comm.comm_group,
detquats,
obsrange[ob].samples,
coord=args.coord,
firstsamp=obsrange[ob].first,
firsttime=obsrange[ob].start,
rate=args.sample_rate,
spinperiod=args.spin_period_min,
spinangle=args.spin_angle_deg,
precperiod=args.prec_period_min,
precangle=args.prec_angle_deg,
detindx=detindx,
detranks=comm.group_size,
hwprpm=hwprpm,
hwpstep=hwpstep,
hwpsteptime=hwpsteptime,
)
obs = {}
obs["name"] = "science_{:05d}".format(ob)
obs["tod"] = tod
obs["intervals"] = None
obs["baselines"] = None
obs["noise"] = noise
obs["id"] = ob
data.obs.append(obs)
if comm.world_rank == 0:
timer.report_clear("Read parameters, compute data distribution")
# Since we are simulating noise timestreams, we want
# them to be contiguous and reproducible over the whole
# observation. We distribute data by detector within an
# observation, so ensure that our group size is not larger
# than the number of detectors we have.
# we set the precession axis now, which will trigger calculation
# of the boresight pointing.
for ob in range(group_numobs):
curobs = data.obs[ob]
tod = curobs["tod"]
# Get the global sample offset from the original distribution of
# intervals
obsoffset = obsrange[group_firstobs + ob].first
# Constantly slewing precession axis
degday = 360.0 / 365.25
precquat = np.empty(4 * tod.local_samples[1],
dtype=np.float64).reshape((-1, 4))
slew_precession_axis(
precquat,
firstsamp=(obsoffset + tod.local_samples[0]),
samplerate=args.sample_rate,
degday=degday,
)
tod.set_prec_axis(qprec=precquat)
del precquat
if comm.world_rank == 0:
timer.report_clear("Construct boresight pointing")
return data
def get_analytic_noise(args, comm, focalplane, verbose=True):
""" Create a TOAST noise object.
Create a noise object from the 1/f noise parameters contained in the
focalplane database.
"""
timer = Timer()
timer.start()
detectors = sorted(focalplane.keys())
fmins = {}
fknees = {}
alphas = {}
NETs = {}
rates = {}
indices = {}
for d in detectors:
rates[d] = args.sample_rate
fmins[d] = focalplane[d]["fmin"]
fknees[d] = focalplane[d]["fknee"]
alphas[d] = focalplane[d]["alpha"]
NETs[d] = focalplane[d]["NET"]
indices[d] = focalplane[d]["index"]
if args.common_mode_noise:
# Add an extra "virtual" detector for common mode noise for
# every optics tube
fmin, fknee, alpha, net = np.array(
args.common_mode_noise.split(",")).astype(np.float64)
hw = hardware.get_example()
for itube, tube in enumerate(sorted(hw.data["tubes"].keys())):
d = "common_mode_{}".format(tube)
detectors.append(d)
rates[d] = args.sample_rate
fmins[d] = fmin
fknees[d] = fknee
alphas[d] = alpha
NETs[d] = net
indices[d] = 100000 + itube
noise = AnalyticNoise(
rate=rates,
fmin=fmins,
detectors=detectors,
fknee=fknees,
alpha=alphas,
NET=NETs,
indices=indices,
)
if args.common_mode_noise:
# Update the mixing matrix in the noise operator
mixmatrix = {}
keys = set()
for det in focalplane.keys():
tube = focalplane[det]["tube"]
common = "common_mode_{}".format(tube)
mixmatrix[det] = {det: 1, common: 1}
keys.add(det)
keys.add(common)
# There should probably be an accessor method to update the
# mixmatrix in the TOAST Noise object.
if noise._mixmatrix is not None:
raise RuntimeError("Did not expect non-empty mixing matrix")
noise._mixmatrix = mixmatrix
noise._keys = list(sorted(keys))
timer.stop()
if comm.world_rank == 0 and verbose:
timer.report("Creating noise model")
return noise
def get_analytic_noise(args, comm, focalplane, verbose=True):
""" Create a TOAST noise object.
Create a noise object from the 1/f noise parameters contained in the
focalplane database. Optionally add thermal common modes.
"""
timer = Timer()
timer.start()
detectors = sorted(focalplane.detector_data.keys())
fmins = {}
fknees = {}
alphas = {}
NETs = {}
rates = {}
indices = {}
for d in detectors:
rates[d] = args.sample_rate
fmins[d] = focalplane[d]["fmin"]
fknees[d] = focalplane[d]["fknee"]
alphas[d] = focalplane[d]["alpha"]
NETs[d] = focalplane[d]["NET"]
indices[d] = focalplane[d]["index"]
ncommon = 0
coupling_strength_distributions = []
common_modes = []
if args.common_mode_noise:
# Add an extra "virtual" detector for common mode noise for
# every optics tube
for common_mode in args.common_mode_noise.split(";"):
ncommon += 1
try:
fmin, fknee, alpha, net, center, width = np.array(
common_mode.split(",")).astype(np.float64)
except ValueError:
fmin, fknee, alpha, net = np.array(
common_mode.split(",")).astype(np.float64)
center, width = 1, 0
coupling_strength_distributions.append([center, width])
hw = get_example()
for itube, tube_slot in enumerate(
sorted(hw.data["tube_slots"].keys())):
d = "common_mode_{}_{}".format(ncommon - 1, tube_slot)
detectors.append(d)
common_modes.append(d)
rates[d] = args.sample_rate
fmins[d] = fmin
fknees[d] = fknee
alphas[d] = alpha
NETs[d] = net
indices[d] = ncommon * 100000 + itube
noise = AnalyticNoise(
rate=rates,
fmin=fmins,
detectors=detectors,
fknee=fknees,
alpha=alphas,
NET=NETs,
indices=indices,
)
if args.common_mode_noise:
mixmatrix = {}
keys = set()
if args.common_mode_only:
detweight = 0
else:
detweight = 1
for icommon in range(ncommon):
# Update the mixing matrix in the noise operator
center, width = coupling_strength_distributions[icommon]
np.random.seed(1001 + icommon)
couplings = center + np.random.randn(1000000) * width
for det in focalplane.detector_data.keys():
if det not in mixmatrix:
mixmatrix[det] = {det: detweight}
keys.add(det)
tube_slot = focalplane[det]["tube_slot"]
common = "common_mode_{}_{}".format(icommon, tube_slot)
index = focalplane[det]["index"]
mixmatrix[det][common] = couplings[index]
keys.add(common)
# Add a diagonal entries, even if we wouldn't usually ask for
# the common mode alone.
for common in common_modes:
mixmatrix[common] = {common: 1}
# There should probably be an accessor method to update the
# mixmatrix in the TOAST Noise object.
if noise._mixmatrix is not None:
raise RuntimeError("Did not expect non-empty mixing matrix")
noise._mixmatrix = mixmatrix
noise._keys = list(sorted(keys))
focalplane._noise = noise
if comm.world_rank == 0 and verbose:
timer.report_clear("Creating noise model")
return noise
