text: # Copyright (c) 2015-2019 by the parties listed in the AUTHORS file.
  # All rights reserved.  Use of this source code is governed by
  # a BSD-style license that can be found in the LICENSE file.
  
  import os
  
  import numpy as np
  
  import healpy as hp
  
  from .mpi import MPITestCase
  
  from ..mpi import MPI
  
  from .. import healpix as hpx
  
  from .. import qarray as qa
  
  from ..todmap import slew_precession_axis, satellite_scanning, TODSatellite
  
  from ._helpers import create_outdir, create_distdata, boresight_focalplane
  
  
  class TODSatelliteTest(MPITestCase):
      def setUp(self):
          fixture_name = os.path.splitext(os.path.basename(__file__))[0]
          self.outdir = create_outdir(self.comm, fixture_name)
  
          # Create 366 observations, divided evenly between groups
          self.nobs = 366
          opg = self.nobs
          if (self.comm is not None) and (self.comm.size >= 2):
              opg = self.nobs // 2
          self.data = create_distdata(self.comm, obs_per_group=opg)
  
          # One boresight detector
          self.ndet = 1
          self.rate = 1.0 / 60.0
  
          (
              dnames,
              dquat,
              depsilon,
              drate,
              dnet,
              dfmin,
              dfknee,
              dalpha,
          ) = boresight_focalplane(self.ndet, samplerate=self.rate)
  
          # Scan strategy.
          # Choose scan parameters so that we return to the origin.
          precangle = 35.0
          spinangle = 55.0
          spinperiod = 240
          precperiod = 8640
  
          # Precession axis slew
          degday = 360.0 / self.nobs
  
          # Sampling
  
          # Only simulate 22 out of 24 hours per observation, in order to test
          # that the starting phase is propagated.  Put the "gap" at the start
          # of each day, so that we can look at the final data point of the
          # final observation.
  
          daysamps = 24 * 60
          # goodsamps = 24 * 60
          goodsamps = 22 * 60
          badsamps = daysamps - goodsamps
  
          # Populate the observations
          for oid, obs in enumerate(self.data.obs):
              firstsamp = oid * daysamps + badsamps
  
              # On the last observation, simulate one extra sample so we get
              # back to the starting point.
              nsim = goodsamps
              if (oid == len(self.data.obs) - 1) and (
                  self.data.comm.group == self.data.comm.ngroups - 1
              ):
                  nsim += 1
  
              tod = TODSatellite(
                  self.data.comm.comm_group,
                  dquat,
                  nsim,
                  detranks=1,
                  firstsamp=firstsamp,
                  firsttime=(firstsamp / self.rate),
                  rate=self.rate,
                  spinperiod=spinperiod,
                  spinangle=spinangle,
                  precperiod=precperiod,
                  precangle=precangle,
              )
  
              qprec = np.empty(4 * tod.local_samples[1], dtype=np.float64).reshape(
                  (-1, 4)
              )
  
              slew_precession_axis(
                  qprec,
                  firstsamp=(firstsamp + tod.local_samples[0]),
                  samplerate=self.rate,
                  degday=degday,
              )
  
              tod.set_prec_axis(qprec=qprec)
  
              obs["tod"] = tod
  
      def tearDown(self):
          pass
  
      def test_precession(self):
          # Test precession axis slew
          slewrate = 1.0 / (24.0 * 3600.0)
          nobs = 366
          degday = 360.0 / nobs
  
          # Plus one is to get back to the beginning
          nsim = nobs + 1
  
          qprec = np.empty(4 * nsim, dtype=np.float64).reshape((-1, 4))
  
          slew_precession_axis(qprec, firstsamp=0, samplerate=slewrate, degday=degday)
  
          zaxis = np.array([0.0, 0.0, 1.0])
  
          v = qa.rotate(qprec, zaxis)
  
          dotprod = v[0][0] * v[-1][0] + v[0][1] * v[-1][1] + v[0][2] * v[-1][2]
          np.testing.assert_almost_equal(dotprod, 1.0)
          return
  
      def test_phase(self):
          # Choose scan parameters so that we return to the origin.
          nobs = 366
          precangle = 35.0
          spinangle = 55.0
          spinperiod = 240
          precperiod = 8640
  
          # Precession axis slew
          degday = 360.0 / 366.0
  
          samplerate = 1.0 / 60.0
  
          daysamps = 24 * 60
          goodsamps = 24 * 60
          # goodsamps = 22 * 60
          badsamps = daysamps - goodsamps
  
          # hit map
          nside = 32
          pix = hpx.Pixels(nside)
          pdata = np.zeros(12 * nside * nside, dtype=np.int32)
  
          # Only simulate 22 out of 24 hours per observation, in order to test
          # that the starting phase is propagated.  Put the "gap" at the start
          # of each day, so that we can look at the final data point of the
          # final observation.
  
          zaxis = np.array([0.0, 0.0, 1.0])
  
          vlast = None
  
          for ob in range(nobs):
              firstsamp = ob * daysamps + badsamps
  
              # On the last observation, simulate one extra sample so we get
              # back to the starting point.
              nsim = goodsamps
              if ob == nobs - 1:
                  nsim += 1
  
              qprec = np.empty(4 * nsim, dtype=np.float64).reshape((-1, 4))
              slew_precession_axis(
                  qprec, firstsamp=firstsamp, samplerate=samplerate, degday=degday
              )
  
              boresight = np.empty(4 * nsim, dtype=np.float64).reshape((-1, 4))
              satellite_scanning(
                  boresight,
                  firstsamp=firstsamp,
                  samplerate=samplerate,
                  qprec=qprec,
                  spinperiod=spinperiod,
                  spinangle=spinangle,
                  precperiod=precperiod,
                  precangle=precangle,
              )
  
              v = qa.rotate(boresight, zaxis)
              p = pix.vec2ring(v)
              pdata[p] += 1
              vlast = v[-1]
  
          if self.data.comm.world_rank == 0:
              import matplotlib.pyplot as plt
  
              hitsfile = os.path.join(self.outdir, "tod_satellite_hits.fits")
              if os.path.isfile(hitsfile):
                  os.remove(hitsfile)
              hp.write_map(hitsfile, pdata, nest=False, dtype=np.int32)
              outfile = "{}.png".format(hitsfile)
              hp.mollview(pdata, xsize=1600, nest=False)
              plt.savefig(outfile)
              plt.close()
  
          np.testing.assert_almost_equal(vlast[0], 0.0)
          np.testing.assert_almost_equal(vlast[1], -1.0)
          np.testing.assert_almost_equal(vlast[2], 0.0)
          return
  
      def test_todclass(self):
          # Hit map
          nside = 32
          pix = hpx.Pixels(nside)
          pdata = np.zeros(12 * nside * nside, dtype=np.int32)
  
          # Compute the boresight hitmap
  
          zaxis = np.array([0.0, 0.0, 1.0])
          vlast = None
  
          for obs in self.data.obs:
              tod = obs["tod"]
              boresight = tod.read_boresight()
              v = qa.rotate(boresight, zaxis)
              p = pix.vec2ring(v)
              pdata[p] += 1
              vlast = v[-1]
  
          allpdata = None
          if self.data.comm.comm_world is None:
              allpdata = pdata
          else:
              if self.data.comm.world_rank == 0:
                  allpdata = np.zeros_like(pdata)
              self.data.comm.comm_world.Reduce(pdata, allpdata, op=MPI.SUM, root=0)
  
          if self.data.comm.world_rank == 0:
              import matplotlib.pyplot as plt
  
              hitsfile = os.path.join(self.outdir, "tod_satellite_classhits.fits")
              if os.path.isfile(hitsfile):
                  os.remove(hitsfile)
              hp.write_map(hitsfile, allpdata, nest=False, dtype=np.int32)
  
              outfile = "{}.png".format(hitsfile)
              hp.mollview(allpdata, xsize=1600, nest=False)
              plt.savefig(outfile)
              plt.close()
  
          # The last sample on the last process should be back to the origin
          lastpass = True
          lastproc = self.data.comm.world_size - 1
          if self.data.comm.world_rank == lastproc:
              lastpass = np.allclose(vlast[0], 0.0)
              lastpass = np.allclose(vlast[1], -1.0)
              lastpass = np.allclose(vlast[2], 0.0)
          if self.data.comm.comm_world is not None:
              lastpass = self.data.comm.comm_world.bcast(lastpass, root=lastproc)
          self.assertTrue(lastpass)
  
          return
  