Make lmul and almxfl accept more shapes. Restore curvedsky.quad_weigh…

…ts, which is used in uharm and analysis. Implement enmap.padtiles, which makes it easier to write tiled map processing algorithms.

cython/cmisc.pxd

@@ -4,8 +4,8 @@ cdef extern from "cmisc_core.h":
 	void alm2cl_dp(int lmax, int mmax, int64_t * mstart, double * alm1, double * alm2, double * cl)
 	void transpose_alm_dp(int lmax, int mmax, int64_t * mstart, double * ialm, double * oalm)
 	void transpose_alm_sp(int lmax, int mmax, int64_t * mstart, float * ialm, float * oalm)
-	void lmul_dp(int lmax, int mmax, int64_t * mstart, double * alm, int lfmax, double * lfun)
-	void lmul_sp(int lmax, int mmax, int64_t * mstart, float * alm, int lfmax, float * lfun)
+	void lmul_dp(int lmax, int mmax, int64_t * mstart, double * alm, int lfmax, const double * lfun)
+	void lmul_sp(int lmax, int mmax, int64_t * mstart, float * alm, int lfmax, const float * lfun)
 	void transfer_alm_dp(int lmax1, int mmax1, int64_t * mstart1, double * alm1, int lmax2, int mmax2, int64_t * mstart2, double * alm2)
 	void transfer_alm_sp(int lmax1, int mmax1, int64_t * mstart1, float * alm1, int lmax2, int mmax2, int64_t * mstart2, float * alm2)
 	void lmatmul_dp(int N, int M, int lmax, int mmax, int64_t * mstart, double ** alm, int lfmax, double ** lmat, double ** oalm)

cython/cmisc.pyx

@@ -162,31 +162,42 @@ def lmul(ainfo, alm, lfun, out=None):
 		if alm.dtype == np.complex128:  _lmatmul_dp(ainfo, alm, lfun, out)
 		elif alm.dtype == np.complex64: _lmatmul_sp(ainfo, alm, lfun, out)
 		else: raise ValueError("lmul requires complex64 or complex128 arrays")
-	elif lfun.ndim == 1 and alm.ndim == 1:
-		if out is None: out = alm.copy()
-		if alm.dtype == np.complex128:  _lmul_dp(ainfo, out[None], lfun[None])
-		elif alm.dtype == np.complex64: _lmul_sp(ainfo, out[None], lfun[None])
-		else: raise ValueError("lmul requires complex64 or complex128 arrays")
-	elif lfun.ndim == 2 and alm.ndim == 2:
-		if out is None: out = alm.copy()
-		if alm.dtype == np.complex128:  _lmul_dp(ainfo, out, lfun)
-		elif alm.dtype == np.complex64: _lmul_sp(ainfo, out, lfun)
-		else: raise ValueError("lmul requires complex64 or complex128 arrays")
 	else:
-		raise ValueError("lmul only supports alm,lfun shapes of [nalm],[nl], [N,nalm],[N,nl] and [N,M,nalm],[M,nl]")
+		# Broadcast pre-dimensions, if they're compatible
+		try:
+			pre  = np.broadcast_shapes(alm.shape[:-1], lfun.shape[:-1])
+		except ValueError:
+			raise ValueError("lmul's alm and lfun's dimensions must either broadcast (when ignoring the last dimension), or have shape compatible with a matrix product (again ignoring the last dimension)")
+		alm  = np.broadcast_to(alm,  pre+ alm.shape[-1:])
+		lfun = np.broadcast_to(lfun, pre+lfun.shape[-1:])
+		# Flatten, so the C code doesn't need to deal with variable dimensionality
+		aflat= alm.reshape(-1,alm.shape[-1])
+		lflat= lfun.reshape(-1,lfun.shape[-1])
+		if out is None:
+			out = aflat.copy()
+		else:
+			out = out.reshape(aflat.shape)
+			out[:] = aflat
+		if alm.dtype == np.complex128:  _lmul_dp(ainfo, out, lflat)
+		elif alm.dtype == np.complex64: _lmul_sp(ainfo, out, lflat)
+		else: raise ValueError("lmul requires complex64 or complex128 arrays")
+		# Unflatten output
+		out = out.reshape(pre + out.shape[-1:])
 	return out
 
 cdef _lmul_dp(ainfo, alm, lfun):
 	cdef int64_t[::1] mstart = np.ascontiguousarray(ainfo.mstart).view(np.int64)
-	cdef double[::1] _alm, _lfun
+	cdef double[::1] _alm
+	cdef const double[::1] _lfun
 	for i in range(alm.shape[-2]):
 		_alm  = alm [i].view(np.float64)
 		_lfun = lfun[i].view(np.float64)
 		cmisc.lmul_dp(ainfo.lmax, ainfo.mmax, &mstart[0], &_alm[0], lfun.shape[1]-1, &_lfun[0])
 
 cdef _lmul_sp(ainfo, alm, lfun):
 	cdef int64_t[::1] mstart = np.ascontiguousarray(ainfo.mstart).view(np.int64)
-	cdef float  [::1] _alm, _lfun
+	cdef float  [::1] _alm
+	cdef const float  [::1] _lfun
 	for i in range(alm.shape[-2]):
 		_alm  = alm [i].view(np.float32)
 		_lfun = lfun[i].view(np.float32)

cython/cmisc_core.c

@@ -123,7 +123,7 @@ void transpose_alm_sp(int lmax, int mmax, int64_t * mstart, float * ialm, float
 }
 
 // Multiply a scalar alm by a scalar function of l
-void lmul_dp(int lmax, int mmax, int64_t * mstart, double * alm, int lfmax, double * lfun) {
+void lmul_dp(int lmax, int mmax, int64_t * mstart, double * alm, int lfmax, const double * lfun) {
 	#pragma omp parallel for
 	for(int m = 0; m <= mmax; m++) {
 		for(int l = m; l <= lmax; l++) {
@@ -136,7 +136,7 @@ void lmul_dp(int lmax, int mmax, int64_t * mstart, double * alm, int lfmax, doub
 }
 
 // Multiply a scalar alm by a scalar function of l
-void lmul_sp(int lmax, int mmax, int64_t * mstart, float * alm, int lfmax, float * lfun) {
+void lmul_sp(int lmax, int mmax, int64_t * mstart, float * alm, int lfmax, const float * lfun) {
 	#pragma omp parallel for
 	for(int m = 0; m <= mmax; m++) {
 		for(int l = m; l <= lmax; l++) {

cython/cmisc_core.h

@@ -3,8 +3,8 @@ void alm2cl_sp(int lmax, int mmax, int64_t * mstart, float  * alm1, float  * alm
 void alm2cl_dp(int lmax, int mmax, int64_t * mstart, double * alm1, double * alm2, double * cl);
 void transpose_alm_dp(int lmax, int mmax, int64_t * mstart, double * ialm, double * oalm);
 void transpose_alm_sp(int lmax, int mmax, int64_t * mstart, float * ialm, float * oalm);
-void lmul_dp(int lmax, int mmax, int64_t * mstart, double * alm, int lfmax, double * lfun);
-void lmul_sp(int lmax, int mmax, int64_t * mstart, float * alm, int lfmax, float * lfun);
+void lmul_dp(int lmax, int mmax, int64_t * mstart, double * alm, int lfmax, const double * lfun);
+void lmul_sp(int lmax, int mmax, int64_t * mstart, float * alm, int lfmax, const float * lfun);
 void transfer_alm_dp(int lmax1, int mmax1, int64_t * mstart1, double * alm1, int lmax2, int mmax2, int64_t * mstart2, double * alm2);
 void transfer_alm_sp(int lmax1, int mmax1, int64_t * mstart1, float * alm1, int lmax2, int mmax2, int64_t * mstart2, float * alm2);
 void lmatmul_dp(int N, int M, int lmax, int mmax, int64_t * mstart, double ** alm, int lfmax, double ** lmat, double ** oalm);

pixell/curvedsky.py

@@ -457,14 +457,20 @@ def get_method(shape, wcs, minfo=None, pix_tol=1e-6):
 
 # Quadrature weights
 
-def quad_weights(shape, wcs, tweak=False):
-	if wcsutils.is_cyl(wcs):
-		return quad_weights_cyl(shape, wcs, tweak=tweak)
-	else:
-		raise NotImplementedError("Quadrature weights only supported for cylindrical projections")
-
-def quad_weights_cyl(shape, wcs, tweak=False):
-	return get_minfo(shape, wcs, quad=True, tweak=tweak).weight
+def quad_weights(shape, wcs, pix_tol=1e-6):
+	"""Return the quadrature weights to use for map2alm operations for the given geometry.
+	Only valid for a limited number of cylindrical geometries recognized by ducc. Returns
+	weights[ny] where ny is shape[-2]. For cases where quadrature weights aren't available,
+	it's a pretty good approximation to just use the pixel area."""
+	minfo = analyse_geometry(shape, wcs, tol=pix_tol)
+	if minfo.ducc_geo.name is None:
+		raise ValueError("Quadrature weights not available for geometry %s,%s" % (str(shape),str(wcs)))
+	ny      = shape[-2]+np.sum(minfo.ypad)
+	weights = ducc0.sht.experimental.get_gridweights(minfo.ducc_geo.name, ny)
+	weights = weights[minfo.ypad[0]:len(weights)-minfo.ypad[1]]
+	if minfo.flip: weights = weights[::-1]
+	weights/= minfo.ducc_geo.nx
+	return weights
 
 #####################
 ### 1d Transforms ###

pixell/enmap.py

@@ -2782,3 +2782,157 @@ def spin_helper(spin, n):
 		if i2 == n: break
 		i1 = i2
 		ci = (ci+1)%len(spin)
+
+# It's often useful to be able to loop over padded tiles, do some operation on them,
+# and then stitch them back together with crossfading. If would be handy to have a way to
+# hide all this complexity. How about an iterator that iterates over padded tiles?
+# E.g.
+#  for itile, otile in zip(ipadtiles(imap), opadtiles(omap)):
+#    otile[:] = fancy_stuff(itile)
+# The input tile iterator is straightforward. The output iterator would
+# zero out omap at the start, and then at the start of each next()
+# would paste the previously yielded tile back into omap with crossfading weights applied.
+# A problem with this formulation where ipadtiles and opadtiles are separate
+# functions, though, is that padding arguments need to be duplicated, which can get
+# tedious. Padding argument must always be consistent when iterating over input
+# and output maps, so probably better to have a single function that processes
+# multiple maps.
+#
+#  for itile, otile in padtiles(imap, omap, tshape=512, pad=30, apod=30):
+#    otile[:] = foo(itile[:])
+#
+# When multiple maps are involved, how should it know which ones
+# are output and input maps? Default:
+#  1 map: input
+#  2 maps: input, output
+#  N maps: input, input, input, ..., output
+# But have an optional argument that lets us specify the types.
+#
+# 3rd alternative which is cleaner but less convenient:
+#  padtile = Padtiler(tshape=512, pad=30, apod=30)
+#  for itile, otile in zip(padtile.in(imap), padtile.out(omap)):
+#   otile[:] = foo(itile)
+# This one has the advantage that it can be built once and then
+# passed to other functions as a single argument. It can easily be used to
+# implement #2, so can get both cheaply. #3 is less convenient in
+# most cases, so #2 would be the typcial interface.
+
+def padtiles(*maps, tshape=600, pad=60, margin=60, mode="auto", start=0, step=1):
+	"""Iterate over padded tiles in one or more maps. The tiling
+	will have a logical tile shape of tshape, but each yielded tile
+	will be expanded with some data from its neighbors. The extra
+	area consists of two parts: The padding and the margin. For
+	a read-iterator these are equivalent, but for a write-iterator
+	the margin will be ignored (and so can be used for apodization),
+	while the padding will be used for crossfading when mergin the
+	tiles together.
+
+	Typical usage:
+
+		for itile, otile in padtiles(imap, imap, margin=60):
+			itile    = apod(itile, 60)
+			otile[:] = some_filter(itile)
+
+	This would iterate over tiles of imap and omap, with the default
+	padding and a margin of 60 pixels. The margin region is used for
+	apodization, and some filter is then applied to the tile, writing
+	the result to the output tile. Note the use of [:] to actually write
+	to otile instead of just rebinding the variable name!
+
+	It's also possible to iterate over fewer or more maps at once.
+	See the "mode" argument.
+
+	If the tile shape does not evenly divide the map shape, then the
+	last tile in each row and column will extend beyond the edge of the
+	map. These pixels will be treated as enmap.extract does, with the
+	potential of sky wrapping. Warning: Write-iterators for a map that
+	goes all the way around the sky while the tile shape does not divide
+	the map shape will incorrectly weight the wrapped tiles, so avoid this.
+
+	Arguments:
+	 * *maps: The maps to iterate over. Must have the same pixel dimensions.
+	 * tshape: The tile shape. Either an integer or a (yshape,xshape) tuple.
+	   Default: 600 pixels.
+	 * pad: The padding. Either an integer or a (ypad,xpad) tuple. Used to
+	   implement context and crossfading. Cannot be larger than half of tshape.
+	   Default: 60 pixels.
+	 * margin: The margin size. Either an integer or a (ymargin,xmargin) tuple.
+	   Ignored in write-iterators, so suitable for apodization.
+	   Default 60 pixels.
+	 * mode: Specifies which maps should be read-iterated vs. write-iterated.
+	   A read-iterated map will yield padded tiles from the corresponding map.
+	   Writes to these tiles are discarded. A write-iterated map yields zeroed
+	   tiles of the same shape as the read-iterator. Writes to these tiles are
+	   used to update the corresponding map, including crossfading the overlapping
+	   regions (due to the padding) such that there aren't any sharp tile boundaries
+	   in the output map. mode can be either "auto" or a string of the same length
+	   as maps consisting of "r" and "w" characters. If the nth character is r/w
+	   then the corresponding map will be read/write-iterated. If the string is
+	   "auto", then the last map will be output-iterated and all the others input-
+	   iterated, unless there's only a single map in which case it will be input-
+	   iterated. Default: "auto".
+	 * start: Flattened tile offset to start at. Useful for mpi loops. Default: 0.
+	 * step:  Flattened tile stride. Useful for mpi loops. Default: 1
+	"""
+	if mode == "auto":
+		if   len(maps) == 0: mode = ""
+		elif len(maps) == 1: mode = "r"
+		else:                mode = "r"*(len(maps)-1)+"w"
+	tiler = Padtiler(tshape=tshape, pad=pad, margin=margin)
+	iters = []
+	for map, io in zip(maps, mode):
+		if   io == "r": iters.append(tiler.read (map))
+		elif io == "w": iters.append(tiler.write(map))
+		else: raise ValueError("Invalid mode character '%s'" % str(io))
+	# Can't just return zip(*iters) because zip gives up when just
+	# one iterator raises StopIteration. This doesn't allow the other
+	# iterators to finish. Maybe this is hacky
+	return utils.zip2(*iters)
+
+class Padtiler:
+	"""Helper class used to implement padtiles. See its docstring for details."""
+	def __init__(self, tshape=600, pad=60, margin=60, start=0, step=1):
+		self.tshape = tuple(np.broadcast_to(tshape, 2).astype(int))
+		self.pad    = tuple(np.broadcast_to(pad,    2).astype(int))
+		self.margin = tuple(np.broadcast_to(margin, 2).astype(int))
+		oly, olx    = 2*np.array(self.pad,int) # overlap region size
+		self.wy     = (np.arange(oly)+1)/(oly+1)
+		self.wx     = (np.arange(olx)+1)/(olx+1)
+		self.start  = start
+		self.step   = step
+	def _tbound(self, tile, tsize, n):
+		pix1 = tile*tsize
+		pix2 = (tile+1)*tsize
+		return pix1, pix2
+	def read (self, imap): return self._it_helper(imap, mode="read")
+	def write(self, omap): return self._it_helper(omap, mode="write")
+	def _it_helper(self, map, mode):
+		# Loop over tiles
+		nty, ntx = (np.array(map.shape[-2:],int)+self.tshape-1)//self.tshape
+		growy, growx = np.array(self.pad) + self.margin
+		oly, olx = 2*np.array(self.pad) # overlap region size
+		for ti in range(self.start, nty*ntx, self.step):
+			ty = ti // ntx
+			tx = ti %  ntx
+			y1, y2 = self._tbound(ty, self.tshape[-2], map.shape[-2])
+			x1, x2 = self._tbound(tx, self.tshape[-1], map.shape[-1])
+			# Construct padded pixel box and extract it
+			pixbox = np.array([[y1-growy,x1-growx],[y2+growy,x2+growx]])
+			tile   = map.extract_pixbox(pixbox).copy()
+			if mode == "read":
+				yield tile
+			else:
+				tile[:] = 0
+				yield tile
+				# Before the next iteration, take the changes the user
+				# made to tile, cop off the margin, and apply crossfading
+				# weights so the overlapping pad regions add up to 1, and
+				# then add to the output map
+				tile = tile[...,self.margin[-2]:tile.shape[-2]-self.margin[-2],self.margin[-1]:tile.shape[-1]-self.margin[-1]]
+				# Apply crossfade weights
+				if ty > 0: tile[...,:oly,:] *= self.wy[:,None]
+				if tx > 0: tile[...,:,:olx] *= self.wx[None,:]
+				if ty < nty-1: tile[...,tile.shape[-2]-oly:,:] *= self.wy[::-1,None]
+				if tx < ntx-1: tile[...,:,tile.shape[-1]-olx:] *= self.wx[None,::-1]
+				# And add into output map
+				map.insert(tile, op=lambda a,b:a+b)

pixell/uharm.py

@@ -50,7 +50,7 @@ class UHT:
 	  - flat:   An enmap[...,ny,nx] corresponding to the multipoles self.l
 	  - curved: A 1d array [...,self.lmax+1]
 	"""
-	def __init__(self, shape, wcs, mode="auto", lmax=None, max_distortion=0.1, tweak=False):
+	def __init__(self, shape, wcs, mode="auto", lmax=None, max_distortion=0.1):
 		"""Initialize a UHT object.
 		Arguments:
 			shape, wcs: The geometry of the enmaps the UHT object will be used on.
@@ -63,7 +63,6 @@ def __init__(self, shape, wcs, mode="auto", lmax=None, max_distortion=0.1, tweak
 		self.shape, self.wcs = shape[-2:], wcs
 		self.area = enmap.area(self.shape, self.wcs)
 		self.fsky = self.area/(4*np.pi)
-		self.tweak= tweak
 		if mode == "auto":
 			dist = estimate_distortion(shape, wcs)
 			if dist <= max_distortion: mode = "flat"
@@ -94,26 +93,26 @@ def map2harm(self, map, spin=0):
 		if self.mode == "flat":
 			return enmap.map2harm(map, spin=spin, normalize="phys")
 		else:
-			return curvedsky.map2alm(map, ainfo=self.ainfo, spin=spin, tweak=self.tweak)
+			return curvedsky.map2alm(map, ainfo=self.ainfo, spin=spin)
 	def harm2map(self, harm, spin=0):
 		if self.mode == "flat":
 			return enmap.harm2map(harm, spin=spin, normalize="phys").real
 		else:
 			rtype= np.zeros(1, harm.dtype).real.dtype
 			omap = enmap.zeros(harm.shape[:-1]+self.shape, self.wcs, rtype)
-			return curvedsky.alm2map(harm, omap, ainfo=self.ainfo, spin=spin, tweak=self.tweak)
+			return curvedsky.alm2map(harm, omap, ainfo=self.ainfo, spin=spin)
 	def harm2map_adjoint(self, map, spin=0):
 		if self.mode == "flat":
 			return enmap.harm2map_adjoint(map, spin=spin, normalize="phys")
 		else:
-			return curvedsky.alm2map_adjoint(map, ainfo=self.ainfo, tweak=self.tweak)
+			return curvedsky.alm2map_adjoint(map, ainfo=self.ainfo)
 	def map2harm_adjoint(self, harm, spin=0):
 		if self.mode == "flat":
 			return enmap.map2harm_adjoint(harm, spin=spin, normalize="phys")
 		else:
 			rtype= np.zeros(1, harm.dtype).real.dtype
 			omap = enmap.zeros(harm.shape[:-1]+self.shape, self.wcs, rtype)
-			return curvedsky.map2alm_adjoint(harm, omap, ainfo=self.ainfo, spin=spin, tweak=self.tweak)
+			return curvedsky.map2alm_adjoint(harm, omap, ainfo=self.ainfo, spin=spin)
 	def quad_weights(self):
 		"""Returns the quadrature weights array W. This will broadcast correctly
 		with maps, but may not have the same dimensions due to symmetries.
@@ -122,7 +121,7 @@ def quad_weights(self):
 			if self.mode == "flat":
 				self.quad = enmap.pixsizemap(self.shape, self.wcs, broadcastable=True)
 			else:
-				self.quad = curvedsky.quad_weights(self.shape, self.wcs, tweak=self.tweak)[:,None]
+				self.quad = curvedsky.quad_weights(self.shape, self.wcs)[:,None]
 		return self.quad
 	def rprof2hprof(self, br, r):
 		"""Like map2harm, but for a 1d function of r."""

pixell/utils.py

@@ -3256,3 +3256,19 @@ def setenv(name, value, keep=False):
 	if   name in os.environ and keep: return
 	elif name in os.environ and value is None: del os.environ[name]
 	else: os.environ[name] = str(value)
+
+def zip2(*args):
+	"""Variant of python's zip that calls next() the same number of times on
+	all arguments. This means that it doesn't give up immediately after getting
+	the first StopIteration exception, but continues on until the end of the row.
+	This can be useful for iterators that want to do cleanup after hte last yield."""
+	done = False
+	while not done:
+		res = []
+		for arg in args:
+			try:
+				res.append(next(arg))
+			except StopIteration:
+				done = True
+		if not done:
+			yield tuple(res)