Pmat ptsrct

pmat/pmat.py

@@ -635,6 +635,27 @@ def backward(self, tod, srcs, tmul=None, pmul=None):
 # Compatibility
 PmatPtsrc2 = PmatPtsrc
 
+class PmatPtsrcTransient(PmatPtsrc):
+	"""Same as PmatPtsrc but accept a profiles[nsrc,nsamps] array with max=1"""
+	def __init__(self, profiles, *args, **kwargs):
+		super().__init__(*args, **kwargs)
+		self.profiles = profiles
+	def apply(self, dir, tod, srcs, tmul=None, pmul=None):
+		if tmul is None: tmul = self.tmul
+		if pmul is None: pmul = self.pmul
+		while srcs.ndim < 4: srcs = srcs[:,None]
+		# Handle angle wrapping without modifying the original srcs array
+		wsrcs = srcs.copy()
+		wsrcs[...,:2] = utils.rewind(srcs[...,:2], self.ref) # + self.dpos
+		core = get_core(tod.dtype)
+		core.pmat_ptsrct(dir, tmul, pmul, tod.T, wsrcs.T, self.profiles.T,
+				self.scan.boresight.T, self.scan.offsets.T, self.scan.comps.T,
+				self.rbox.T, self.nbox.T, self.yvals.T,
+				self.beam[1], self.beam[0,-1], self.rmax,
+				self.cells.T, self.ncell.T, self.cbox.T)
+		# Copy out any amplitudes that may have changed
+		srcs[...,2:5] = wsrcs[...,2:5]	# 2:5 -> T,Q,U in nparams
+
 def build_interpol(transform, box, id="none", posunit=1.0, sys=None, pad=None):
 	sys   = config.get("map_sys",      sys)
 	# We widen the bounding box slightly to avoid samples falling outside it

pmat/pmat_core.F90

@@ -1918,6 +1918,139 @@ subroutine pmat_ptsrc2( &
 		end if
 	end subroutine
 
+	! same as pmat_ptsrc2 but allow a time-based profile to modulate the beam
+	subroutine pmat_ptsrct( &
+			dir, tmul, pmul,           &! Projection direction, tod multiplier, src multiplier
+			tod, srcs, profiles,       &! Main inputs/outputs. tod(nsamp,ndet), srcs(nparam,ndet_or_1,ndir,nsrc), pulse(nsamp,nsrc)
+			bore, det_pos, det_comps,  &
+			rbox, nbox, yvals,         &! Coordinate transformation
+			beam, rbeam, rmax,         &! Beam profile and max radial offset to consider
+			cell_srcs, cell_nsrc, cbox &! Relevant source lookup. cell_srcs(:,nx,ny,ndet_or_1,ndir), cell_nsrc(nx,ny,ndet_or_1,ndir)
+		)
+		use omp_lib
+		implicit none
+		integer, intent(in)    :: dir
+		real(_), intent(in)    :: tmul, pmul
+		real(_), intent(inout) :: tod(:,:)
+		real(8), intent(inout) :: srcs(:,:,:,:)
+		real(_), intent(in)    :: profiles(:,:)
+		real(8), intent(in)    :: bore(:,:), det_pos(:,:), rbox(:,:), yvals(:,:)
+		real(8), intent(in)    :: cbox(:,:), beam(:), rbeam, rmax
+		real(_), intent(in)    :: det_comps(:,:)
+		integer, intent(in)    :: nbox(:), cell_srcs(:,:,:,:,:), cell_nsrc(:,:,:,:)
+		! Work
+		integer :: nsamp, ndet, nsrc, nproc, nsrcdet
+		! Not the same sdir as in the shift stuff
+		integer :: ic, i, id, di, si, xind(3), ig, ig2, cell(2), cell_ind, cid, sdir, ndir
+		integer :: steps(3), bind, sdi
+		real(8) :: x0(3), inv_dx(3), c0(2), inv_dc(2), xrel(3), work(size(yvals,1),4)
+		real(8) :: point(4), phase(3), dec, ra, ddec, dra, bscale(3), t, omg_c, phi_c, pulse, D
+		real(_) :: inv_bres, bx,by,br,brel,bval, c2p,s2p,c1p,s1p
+		real(_), parameter   :: pi = 3.14159265359d0
+		real(8), allocatable :: amps(:,:,:,:,:), cosdec(:,:,:), ys(:,:,:)
+		integer, allocatable :: scandir(:)
+		nsamp   = size(tod, 1)
+		ndet    = size(tod, 2)
+		! input from python are transposed, hence the index is reversed
+		! srcs -> [nparam, ndet, ndir, nsrc]
+		nsrcdet = size(srcs,2)
+		ndir    = size(srcs,3)
+		nsrc    = size(srcs,4)
+
+		! Set up scanning direction. Two modes are supported. If ndir is 1, then
+		! the same set of parameters are used for both left and rightgoing scans.
+		! If ndir is 2, then these are separated.
+		allocate(scandir(nsamp))
+		if(ndir > 1) then
+			call calc_scandir(bore(2,:), scandir)
+		else
+			scandir = 1
+		end if
+		! Set up interpolation
+		call interpol_prepare(nbox, rbox, steps, x0, inv_dx)
+		! And the beam interpolation. The last cell ends at cbox(:,2), but
+		! starts one cell-width before that.
+		c0 = cbox(:,1)
+		inv_dc(1) = size(cell_nsrc,2)/(cbox(1,2)-cbox(1,1))
+		inv_dc(2) = size(cell_nsrc,1)/(cbox(2,2)-cbox(2,1))
+		inv_bres = (size(beam)-1)/rbeam
+
+		nproc = omp_get_max_threads()
+		allocate(cosdec(nsrcdet,ndir,nsrc),amps(3,nsrcdet,ndir,nsrc,nproc))
+		cosdec = cos(srcs(1,:,:,:))
+		if(dir > 0) then
+			do i = 1, nproc; amps(:,:,:,:,i) = srcs(3:5,:,:,:)*pmul; end do
+		else
+			amps = 0
+		end if
+		!$omp parallel private(id,di,si,sdi,xrel,xind,ig,work,point,phase,cell,cell_ind,cid,dec,ra,bscale,ddec,dra,sdir,c2p,s2p,c1p,s1p,bx,by,br,brel,bind,bval)
+		id = omp_get_thread_num()+1
+		!$omp do
+		do di = 1, ndet
+			do si = 1, nsamp
+				sdir = scandir(si)
+				sdi  = min(di,nsrcdet)
+				! Transform from hor to cel
+				include 'helper_bilin.F90'
+				! Find which point source lookup cell we are in.
+				! dec,ra -> cy,cx
+				cell = floor((point(1:2)-c0)*inv_dc)+1
+				! Bounds checking. Costs 2% performance. Worth it
+				cell(1) = min(size(cell_nsrc,2),max(1,cell(1)))
+				cell(2) = min(size(cell_nsrc,1),max(1,cell(2)))
+				if(dir > 0) tod(si,di) = tod(si,di)*tmul
+				! Avoid expensive operations if we don't hit any sources
+				if(cell_nsrc(cell(2),cell(1),sdi,sdir) == 0) cycle
+				! The spin-2 and spin-1 rotations associated with the transformation
+				! We need these to get the polarization rotation and beam orientation
+				! right.
+				c2p = point(3);                  s2p = point(4)
+				c1p = sign(sqrt((1+c2p)/2),s2p); s1p = sqrt((1-c2p)/2)
+				phase(1) = det_comps(1,di)
+				phase(2) = c2p*det_comps(2,di) - s2p*det_comps(3,di)
+				phase(3) = s2p*det_comps(2,di) + c2p*det_comps(3,di)
+				! Process each point source in this cell
+				do cell_ind = 1, cell_nsrc(cell(2),cell(1),sdi,sdir)
+					cid = cell_srcs(cell_ind,cell(2),cell(1),sdi,sdir)+1
+					! dec, ra, T, Q, U, ibx, iby, ibxy
+					dec   = srcs(1,sdi,sdir,cid)
+					ra    = srcs(2,sdi,sdir,cid)
+					bscale= srcs(6:8,sdi,sdir,cid)
+					! Calc effective distance from this source in terms of the beam distortions.
+					! The beam shape is defined in the same coordinate system the polarization
+					! orientation is defined in. We can either rotate the beam (like we do phase)
+					! or rotate the offset vector the opposite direction. I choose the latter
+					! because it is simpler.
+					ddec = point(1)-dec
+					dra  = (point(2)-ra)*cosdec(sdi,sdir,cid) ! Caller should beware angle wrapping!
+					if(abs(ddec)>rmax .or. abs(dra)>rmax) cycle
+					bx   =  c1p*dra + s1p*ddec
+					by   = -s1p*dra + c1p*ddec
+					br   = sqrt(by*(bscale(1)*by+2*bscale(3)*bx) + bx**2*bscale(2))
+					! Linearly interpolate the beam value
+					brel = br*inv_bres+1
+					bind = floor(brel)
+					if(bind >= size(beam)) cycle
+					brel = brel-bind
+					bval = beam(bind)*(1-brel) + beam(bind+1)*brel
+
+					! And perform the actual projection
+					if(dir > 0) then
+						tod(si,di) = tod(si,di) + sum(amps(:,sdi,sdir,cid,1)*phase)*bval*profiles(si,cid)
+					else
+						amps(:,sdi,sdir,cid,id) = amps(:,sdi,sdir,cid,id) + tod(si,di)*bval*phase*profiles(si,cid)
+					end if
+				end do
+			end do
+		end do
+		!$omp end parallel
+		if(dir < 0) then
+			! sum(amps,5) is to reduce amps from all processes by summation
+			! tmul is not wrongly applied here: previous tmul applies to dir > 0 only.
+			srcs(3:5,:,:,:) = srcs(3:5,:,:,:)*pmul + sum(amps,5)*tmul
+		end if
+	end subroutine
+
 	subroutine pmat_az(dir, tod, map, az, dets, az0, daz)
 		implicit none
 		integer, intent(in)    :: dir, dets(:)