• Home
  • Features
  • Pricing
  • Docs
  • Announcements
  • Sign In

simonsobs / so3g / 7003793410

27 Nov 2023 10:35AM UTC coverage: 47.688% (+0.09%) from 47.602%
7003793410

push

github

web-flow
Merge pull request #157 from simonsobs/bilinear

Bilinear mapmaking projection operators

16 of 18 new or added lines in 1 file covered. (88.89%)

1 existing line in 1 file now uncovered.

1227 of 2573 relevant lines covered (47.69%)

0.48 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

75.11
/python/proj/wcs.py
1
import so3g
1✔
2
from . import quat
1✔
3

4
import numpy as np
1✔
5

6
from .ranges import Ranges, RangesMatrix
1✔
7
from . import mapthreads
1✔
8

9
# For coordinate systems we use the following abbreviations:
10
#
11
# - DC: Detector coordinates
12
# - FP: Focal plane coordinates
13
# - CS: Celestial spherical coordinates (likely equatorial)
14
# - NS: Native spherical coordinates of the map projection
15
#
16
# Saying that a quaternion rotation q_B<A "takes A to B", where A and
17
# B are coordinate systems, means that if a vector has coordinates v_A
18
# in system A then its coordinates in system B are:
19
#
20
#      v_B = q_B<A v_A q_B<A*
21
#
22
# The rotation taking detector coordinates to native spherical
23
# coordinates is the product
24
#
25
#     q_NS<DC = q_NS<CS q_CS<FP q_FP<DC
26
#
27
# In this module, quaternion rotations are written as:
28
# - Q_CS<FP: Assembly.Q (vector in samples)
29
# - Q_FP<DC: Assembly.dets (vector in dets)
30
# - Q_NS<CS: Projectionst.q_celestial_to_native
31
#
32
# The Projectionist caches a vector of Q_NS<FP, which is then simply
33
# called "q_native".  This is usually the thing to pass to C++ level.
34

35

36
class Projectionist:
1✔
37
    """This class assists with analyzing WCS information to populate data
38
    structures needed for accelerated pointing routines.
39

40
    On instantiation, it carries information about the relation
41
    between celestial spherical coordinates and the Native Spherical
42
    coordinates of the projection, and also the pixelization scheme
43
    for the projection.
44

45
    As in pixell, the code and discussion here uses the term
46
    "geometry" to refer to the combination of an astropy.WCS object
47
    and a 2-d array shape.
48

49
    The following symbols are used in docstrings to represent shapes:
50

51
    * ``n_dets`` - Number of detectors in the focal plane.
52
    * ``n_samps`` - Number of samples in the signal space of each detector.
53
    * ``n_threads`` - Number of threads to use in parallelized computation.
54
    * ``n_mapidx`` - Number of indices required to specify a pixel in
55
      a map.  For simple maps this is probably 2, with the first index
56
      specifying row and the second index specifying column.  But for
57
      tiled maps n_mapidx is 3, with the first axis specifying the
58
      tile number.
59
    * ``n_comp`` - Number of map components, such as Stokes components
60
      or other spin components.  This takes the value 3, for example,
61
      if projection into TQU space has been requested.
62

63
    When coordinate computation or projection routines are called, a
64
    ProjEng (a wrapped C++ object) is instantiated to perform the
65
    accelerated computations.  The spin-composition (i.e. is this a
66
    'T' map or a 'TQU' map) must be specified at this time, to
67
    instantiate the correct accelerated object.
68

69
    Some method parameters are common to many functions and are documented
70
    here for consistency:
71

72
    * ``assembly`` - an Assembly, providing a focal plane (quaternion
73
      offsets for each detector) as well as a boresight vector
74
      (quaternion for each time sample).
75
    * ``comps`` - a string specifying the Stokes components of the
76
      map, for example 'T' or 'TQU'.  When absent, this will be
77
      guessed from the map shape; with 1|2|3 mapping to 'T'|'QU'|'TQU'
78
      respectively.
79
    * ``proj_name`` - a string specifying a projection.  The
80
      nomenclature is mostly the same as the FITS CTYPE identifiers.
81
      Accepted values: ARC, CAR, CEA, TAN, ZEA, Flat, Quat.
82
    * ``shape`` - the shape describing the celestial axes of the map.
83
      (For tiled representation this specifies the parent footprint.)
84
    * ``wcs`` - the WCS describing the celestial axes of the map.
85
      Together with ``shape`` this is a geometry; see pixell.enmap
86
      documentation.
87
    * ``threads`` - the thread assignment, which is a RangesMatrix
88
      with shape (n_threads,n_dets,n_samps), used to specify which
89
      samples should be treated by each thread in TOD-to-map
90
      operations.  Such objects should satisfy the condition that
91
      threads[x,j]*threads[y,j] is the empty Range for x != y;
92
      i.e. each detector-sample is assigned to at most one thread.
93
    * ``interpol``: How positions that fall between pixel centers will
94
      be handled. Options are "nearest" (default): Use Nearest
95
      Neighbor interpolation, so a sample takes the value of
96
      whatever pixel is closest; or "bilinear": linearly
97
      interpolate between the four closest pixels. bilinear is
98
      slower (around 60%) but avoids problems caused by a
99
      discontinuous model.
100

101
    Attributes:
102
        naxis: 2-element integer array specifying the map shape (for
103
            the 2 celestial axes).
104
        cdelt: 2-element float array specifying the pixel pitch.
105
        crpix: 2-element float array specifying the pixel coordinates
106
            of the reference point.
107
        proj_name: string, name of the projection.
108
        q_celestial_to_native: quaternion rotation taking celestial
109
            coordinates to the native spherical coordinates of the
110
            projection.
111
        tile_shape: 2-element integer array specifying the tile shape
112
            (this is the shape of one full-size sub-tile, not the
113
            decimation factor on each axis).
114
        tiling: a _Tiling object if the map is tiled, None otherwise.
115

116
    """
117
    @staticmethod
1✔
118
    def get_q(wcs):
1✔
119
        """Analyze a wcs object to compute the quaternion rotation from
120
        celestial to native spherical coordinates.
121

122
        """
123
        alpha0, delta0 = wcs.wcs.crval  # In degrees.
1✔
124
        if (wcs.wcs.phi0 == 0. and wcs.wcs.theta0 == 0.):
1✔
125
            # This is typical for cylindrical projections.
126
            assert((delta0 >= 0 and wcs.wcs.lonpole == 180.0) or
×
127
                   (delta0 <= 0 and wcs.wcs.lonpole ==   0.0))
128
            Q = (quat.euler(1,  delta0 * quat.DEG) *
×
129
                 quat.euler(2, -alpha0 * quat.DEG))
130
        elif (wcs.wcs.phi0 == 0. and wcs.wcs.theta0 == 90.):
1✔
131
            # This is typical for zenithal projections.
132
            assert(wcs.wcs.lonpole == 180.0)
1✔
133
            Q = (quat.euler(2, np.pi) *
1✔
134
                 quat.euler(1, (delta0 - 90)*quat.DEG) *
135
                 quat.euler(2, -alpha0 * quat.DEG))
136
        else:
137
            raise ValueError(f'Unimplemented NSC reference (phi0,theta0)='
×
138
                             '({wcs.wcs.phi0:.2f},{wcs.wcs.theta0:.2f})')
139
        return Q
1✔
140

141
    def __init__(self):
1✔
142
        self._q_fp_to_native = None
1✔
143
        self._q_fp_to_celestial = None
1✔
144
        self.tile_shape = None
1✔
145
        self.naxis = np.array([0, 0])
1✔
146
        self.cdelt = np.array([0., 0.])
1✔
147
        self.crpix = np.array([0., 0.])
1✔
148

149
    @property
1✔
150
    def tiling(self):
1✔
151
        if self.tile_shape is None:
1✔
152
            return None
1✔
153
        return _Tiling(self.naxis[::-1], self.tile_shape)
1✔
154

155
    @classmethod
1✔
156
    def for_geom(cls, shape, wcs, interpol=None):
1✔
157
        """Construct a Projectionist for use with the specified "geometry".
158

159
        The shape and wcs are the core information required to prepare
160
        a Projectionist, so this method is likely to be called by
161
        other constructors.
162
        """
163
        self = cls()
1✔
164
        ax1, ax2 = wcs.wcs.lng, wcs.wcs.lat
1✔
165
        ndim = len(shape)
1✔
166
        # The axes are numbered from outside in...
167
        self.naxis = np.array([shape[ndim - ax1 - 1],
1✔
168
                               shape[ndim - ax2 - 1]], dtype=int)
169
        # Get just the celestial part.
170
        wcs = wcs.celestial
1✔
171
        self.wcs = wcs
1✔
172

173
        # Extract the projection name (e.g. CAR)
174
        proj = [c[-3:] for c in wcs.wcs.ctype]
1✔
175
        assert(proj[0] == proj[1])
1✔
176
        proj_name = proj[0]  # Projection name
1✔
177
        self.proj_name = proj_name
1✔
178

179
        # Store the rotation to native spherical coordinates.
180
        self.q_celestial_to_native = self.get_q(wcs)
1✔
181

182
        # Store the grid info.
183
        self.cdelt = np.array(wcs.wcs.cdelt) * quat.DEG
1✔
184
        self.crpix = np.array(wcs.wcs.crpix)
1✔
185

186
        # Pixel interpolation mode
187
        if interpol is None: interpol = "nearest"
1✔
188
        self.interpol = interpol
1✔
189

190
        return self
1✔
191

192
    @classmethod
1✔
193
    def for_map(cls, emap, wcs=None, interpol=None):
1✔
194
        """Construct a Projectionist for use with maps having the same
195
        geometry as the provided enmap.
196

197
        Args:
198
          emap: enmap from which to extract shape and wcs information.
199
            It is acceptable to pass a bare ndarray here (or anything
200
            with shape attribute), provided that wcs is provided
201
            separately.
202
          wcs: optional WCS object to use instead of emap.wcs.
203
        """
UNCOV
204
        if wcs is None:
×
205
            wcs = emap.wcs
×
206
        return cls.for_geom(emap.shape, wcs)
×
207

208
    @classmethod
1✔
209
    def for_source_at(cls, alpha0, delta0, gamma0=0.,
1✔
210
                      proj_name='TAN'):
211
        """Return a pointing-only Projectionist where some particular
212
        equatorial position will be put at the North Pole of the
213
        Native spherical coordinates.
214

215
        """
216
        self = cls()
×
217
        self.proj_name = proj_name
×
218
        assert(gamma0 == 0.)
×
219
        self.q_celestial_to_native = (
×
220
            quat.euler(2, np.pi)
221
            * quat.euler(1, (delta0 - 90)*quat.DEG)
222
            * quat.euler(2, -alpha0 * quat.DEG))
223
        return self
×
224

225
    @classmethod
1✔
226
    def for_tiled(cls, shape, wcs, tile_shape, active_tiles=True, interpol=None):
1✔
227
        """Construct a Projectionist for use with the specified geometry
228
        (shape, wcs), cut into tiles of shape tile_shape.
229

230
        See class documentation for description of standard arguments.
231

232
        Args:
233
            tile_shape: tuple of ints, giving the shape of each tile.
234
            active_tiles: bool or list of ints.  Specifies which tiles
235
                should be considered active.  If True, all tiles are
236
                populated.  If None or False, this will remain
237
                uninitialized and attempts to generate blank maps
238
                (such as calls to zeros or to projection functions
239
                without a target map set) will fail.  Otherwise this
240
                must be a list of integers specifying which tiles to
241
                populate on such calls.
242

243
        """
244
        self = cls.for_geom(shape, wcs, interpol=interpol)
1✔
245
        self.tile_shape = np.array(tile_shape, 'int')
1✔
246
        if active_tiles is True:
1✔
247
            self.active_tiles = list(range(self.tiling.tile_count))
1✔
248
        elif active_tiles in [False, None]:
1✔
249
            self.active_tiles = None
1✔
250
        else:
251
            # Presumably a list of tile numbers.
252
            self.active_tiles = active_tiles
1✔
253
        return self
1✔
254

255
    def _get_pixelizor_args(self):
1✔
256
        """Returns a tuple of arguments that may be passed to the ProjEng
257
        constructor to define the pixelization.
258

259
        """
260
        # All these casts are required because boost-python doesn't
261
        # like numpy scalars.
262
        args = (int(self.naxis[1]), int(self.naxis[0]),
1✔
263
                float(self.cdelt[1]), float(self.cdelt[0]),
264
                float(self.crpix[1]), float(self.crpix[0]))
265
        if self.tiling:
1✔
266
            args += tuple(map(int, self.tile_shape))
1✔
267
            if self.active_tiles is not None:
1✔
268
                args += (self.active_tiles,)
1✔
269
        return args
1✔
270

271
    def get_ProjEng(self, comps='TQU', proj_name=None, get=True,
1✔
272
                    instance=True, interpol=None):
273
        """Returns an so3g.ProjEng object appropriate for use with the
274
        configured geometry.
275

276
        """
277
        if proj_name is None: proj_name = self.proj_name
1✔
278
        tile_suffix = 'Tiled' if self.tiling else 'NonTiled'
1✔
279
        # Interpolation mode
280
        if interpol is None: interpol = self.interpol
1✔
281
        if interpol in ["nn", "nearest"]: interpol_suffix = ""
1✔
NEW
282
        elif interpol in ["lin", "bilinear"]: interpol_suffix = "_Bilinear"
×
NEW
283
        else: raise ValueError("ProjEng interpol '%s' not recognized" % str(interpol))
×
284
        projeng_name = f'ProjEng_{proj_name}_{comps}_{tile_suffix}{interpol_suffix}'
1✔
285
        if not get:
1✔
286
            return projeng_name
×
287
        try:
1✔
288
            projeng_cls = getattr(so3g, projeng_name)
1✔
289
        except AttributeError:
×
290
            raise ValueError(f'There is no projector implemented for '
×
291
                             f'pixelization "{proj_name}", components '
292
                             f'"{comps}" (tried "{projeng_name}").')
293
        if not instance:
1✔
294
            return projeng_cls
×
295
        return projeng_cls(self._get_pixelizor_args())
1✔
296

297
    def _cache_q_fp_to_native(self, q_fp_to_celestial):
1✔
298
        """Get q_fp_to_native for argument q_fp_to_celestial, and cache the
299
        result, and return that result later if called with the same
300
        argument.
301

302
        """
303
        if q_fp_to_celestial is not self._q_fp_to_celestial:
1✔
304
            self._q_fp_to_native = self.q_celestial_to_native * q_fp_to_celestial
1✔
305
            self._q_fp_to_celestial = q_fp_to_celestial
1✔
306
        return self._q_fp_to_native
1✔
307

308
    def _guess_comps(self, map_shape):
1✔
309
        if len(map_shape) != 3:
×
310
            raise ValueError('Cannot guess components based on '
×
311
                             'map with %i!=3 dimensions!' % len(map_shape))
312
        if map_shape[0] == 1:
×
313
            return 'T'
×
314
        elif map_shape[0] == 2:
×
315
            return 'QU'
×
316
        elif map_shape[0] == 3:
×
317
            return 'TQU'
×
318
        raise ValueError('No default components for ncomp=%i; '
×
319
                         'set comps=... explicitly.' % map_shape[0])
320

321
    def get_pixels(self, assembly):
1✔
322
        """Get the pixel indices for the provided pointing Assembly.  For each
323
        detector, an int32 array of shape (n_samps, n_mapidx) is
324
        returned.  The value of the first slot of the second axis will
325
        be -1 if and only if the sample's projected position is
326
        outside the map footprint.
327

328
        See class documentation for description of standard arguments.
329

330
        """
331
        projeng = self.get_ProjEng('TQU')
×
332
        q_native = self._cache_q_fp_to_native(assembly.Q)
×
333
        return projeng.pixels(q_native, assembly.dets, None)
×
334

335
    def get_pointing_matrix(self, assembly):
1✔
336
        """Get the pointing matrix information, which is to say both the pixel
337
        indices and the "spin projection factors" for the provided
338
        pointing Assembly.  Returns (pixel_indices, spin_proj) where
339
        pixel_indices is as returned by get_pixels() and spin_proj is
340
        a list of float arrays of shape [n_samps, n_comp].  As an
341
        alternative to on-the-fly computation, this information can be
342
        cached and used for very fast projection operations.
343

344
        See class documentation for description of standard arguments.
345

346
        """
347
        projeng = self.get_ProjEng('TQU')
×
348
        q_native = self._cache_q_fp_to_native(assembly.Q)
×
349
        return projeng.pointing_matrix(q_native, assembly.dets, None, None)
×
350

351
    def get_coords(self, assembly, use_native=False, output=None):
1✔
352
        """Get the spherical coordinates for the provided pointing Assembly.
353
        For each detector, a float64 array of shape [n_samps,4] is
354
        returned.  In the right-most index, the first two components
355
        are the longitude and latitude in radians.  The next two
356
        components are the cosine and sine of the parallactic angle.
357

358
        This routine uses a trivial CAR projection to return results
359
        from the celestial coordinate system (i.e. (lon,lat) =
360
        (RA,dec)), and the parallactic angle will be relative to that
361
        projection.  If you are interested in the parallactic angle of
362
        the native spherical coordinate system (e.g. if you're doing a
363
        tangent plane projection), make a second call specifying
364
        use_native=True.  In this case you might also take a look at
365
        the get_planar() routine.
366

367
        See class documentation for description of standard arguments.
368

369
        """
370
        projeng = self.get_ProjEng('TQU', 'CAR')
×
371
        if use_native:
×
372
            q_native = self._cache_q_fp_to_native(assembly.Q)
×
373
        else:
374
            q_native = assembly.Q
×
375
        return projeng.coords(q_native, assembly.dets, output)
×
376

377
    def get_planar(self, assembly, output=None):
1✔
378
        """Get projection plane coordinates for all detectors at all times.
379
        For each detector, a float64 array of shape [n_samps,4] is
380
        returned.  The first two elements are the x and y projection
381
        plane coordiantes, similar to the "intermediate world
382
        coordinates", in FITS language.  Insofar as FITS ICW has units
383
        of degrees, these coordinates have units of radians.  Indices
384
        2 and 3 carry the cosine and sine of the detector parallactic
385
        angle.
386

387
        """
388
        q_native = self._cache_q_fp_to_native(assembly.Q)
×
389
        projeng = self.get_ProjEng('TQU')
×
390
        return projeng.coords(q_native, assembly.dets, output)
×
391

392
    def zeros(self, super_shape):
1✔
393
        """Return a map, filled with zeros, with shape (super_shape,) +
394
        self.shape.  For tiled maps, this will be a list of map tiles
395
        (with shape (super_shape, ) + tile_shape.  If any tiles are
396
        not active, None is returned in the corresponding slots.
397

398
        """
399
        projeng = self.get_ProjEng('T')
1✔
400
        return projeng.zeros(super_shape)
1✔
401

402
    def to_map(self, signal, assembly, output=None, det_weights=None,
1✔
403
               threads=None, comps=None):
404
        """Project signal into a map.
405

406
        Arguments:
407
          signal (Signal-like): The signal to project.
408
          output (Map-like): The map into which to accumulate the
409
            projected signal.  If None, a map will be initialized internally.
410
          det_weights (shape n_det): Weights to use for each detector.
411

412
        See class documentation for description of standard arguments.
413

414
        """
415
        if output is None and comps is None:
1✔
416
            raise ValueError("Provide an output map or specify component of "
×
417
                             "interest (e.g. comps='TQU').")
418
        if comps is None:
1✔
419
            comps = self._guess_comps(output.shape)
×
420
        projeng = self.get_ProjEng(comps)
1✔
421
        q_native = self._cache_q_fp_to_native(assembly.Q)
1✔
422
        map_out = projeng.to_map(
1✔
423
            output, q_native, assembly.dets, signal, det_weights, threads)
424
        return map_out
1✔
425

426
    def to_weights(self, assembly, output=None, det_weights=None,
1✔
427
                   threads=None, comps=None):
428
        """Project pointing into a weights map.
429

430
        Arguments:
431
          output (Map-like): The weights map into which to accumulate
432
            the projected signal.  If None, a map will be initialized
433
            internally.
434
          det_weights (shape n_det): Weights to use for each detector.
435

436
        See class documentation for description of standard arguments.
437

438
        """
439
        if output is None and comps is None:
1✔
440
            raise ValueError("Provide an output map or specify component of "
×
441
                             "interest (e.g. comps='TQU').")
442
        if comps is None:
1✔
443
            assert(output.shape[0] == output.shape[1])
×
444
            comps = self._guess_comps(output.shape[1:])
×
445
        projeng = self.get_ProjEng(comps)
1✔
446
        q_native = self._cache_q_fp_to_native(assembly.Q)
1✔
447
        map_out = projeng.to_weight_map(
1✔
448
            output, q_native, assembly.dets, det_weights, threads)
449
        return map_out
1✔
450

451
    def from_map(self, src_map, assembly, signal=None, comps=None):
1✔
452
        """De-project from a map, returning a Signal-like object.
453

454
        Arguments:
455
          src_map (Map-like): The map from which to sample.
456
          signal (Signal-like): The object into which to accumulate
457
            the signal.  If not provided, a suitable object will be
458
            created and initialized to zero.
459

460
        See class documentation for description of standard arguments.
461

462
        """
463
        if src_map.ndim == 2:
×
464
            # Promote to (1,...)
465
            src_map = src_map[None]
×
466
        if comps is None:
×
467
            comps = self._guess_comps(src_map.shape)
×
468
        projeng = self.get_ProjEng(comps)
×
469
        q_native = self._cache_q_fp_to_native(assembly.Q)
×
470
        signal_out = projeng.from_map(
×
471
            src_map, q_native, assembly.dets, signal)
472
        return signal_out
×
473

474
    def assign_threads(self, assembly, method='domdir', n_threads=None):
1✔
475
        """Get a thread assignment RangesMatrix.  Different algorithms can be
476
        chosen, though this method does not provide fine-grained
477
        control of algorithm parameters and instead seeks to apply
478
        reasonable defaults.
479

480
        The methods exposed here are:
481

482
        - ``'simple'``: Divides the map into n_threads horizontal
483
          bands, and assigns the pixels in each band to a single
484
          thread.  This tends to be bad for scans that aren't
485
          horizontally correlated, or that don't have equal weight
486
          throughout the map.  The 'domdir' algorithm addresses both
487
          of those problems.
488
        - ``'domdir'``: For constant elevation scans, determines the
489
          dominant scan direction in the map and divides it into bands
490
          parallel to that scan direction.  The thickness of bands is
491
          adjusted so roughly equal numbers of samples fall into each
492
          band.
493
        - ``'tiles'``: In Tiled projections, assign each tile to a
494
          thread.  Some effort is made to balance the total number of
495
          samples over the threads.
496

497
        The returned object can be passed to the ``threads`` argument
498
        of the projection methods in this class.
499

500
        See class documentation for description of standard arguments.
501

502
        Args:
503
          method (str): Algorithm to use.
504

505
        """
506
        if method not in THREAD_ASSIGNMENT_METHODS:
1✔
507
            raise ValueError(f'No thread assignment method "{method}"; '
×
508
                             f'expected one of {THREAD_ASSIGNMENT_METHODS}')
509
        n_threads = mapthreads.get_num_threads(n_threads)
1✔
510

511
        if method == 'simple':
1✔
512
            projeng = self.get_ProjEng('T')
1✔
513
            q_native = self._cache_q_fp_to_native(assembly.Q)
1✔
514
            omp_ivals = projeng.pixel_ranges(q_native, assembly.dets, None, n_threads)
1✔
515
            return wrap_ivals(omp_ivals)
1✔
516

517
        elif method == 'domdir':
1✔
518
            offs_rep = assembly.dets[::100]
1✔
519
            if (self.tiling is not None) and (self.active_tiles is None):
1✔
520
                tile_info = self.get_active_tiles(assembly)
1✔
521
                active_tiles = tile_info['active_tiles']
1✔
522
                self.active_tiles = active_tiles
1✔
523
            else:
524
                active_tiles = None
1✔
525
            return mapthreads.get_threads_domdir(
1✔
526
                assembly.Q, assembly.dets, shape=self.naxis[::-1], wcs=self.wcs,
527
                tile_shape=self.tile_shape, active_tiles=active_tiles,
528
                n_threads=n_threads, offs_rep=offs_rep)
529

530
        elif method == 'tiles':
1✔
531
            tile_info = self.get_active_tiles(assembly, assign=n_threads)
1✔
532
            self.active_tiles = tile_info['active_tiles']
1✔
533
            return tile_info['group_ranges']
1✔
534

535
        raise RuntimeError(f"Unimplemented method: {method}.")
×
536

537
    def assign_threads_from_map(self, assembly, tmap, n_threads=None):
1✔
538
        """Assign threads based on a map.
539

540
        The returned object can be passed to the ``threads`` argument
541
        of the projection methods in this object.
542

543
        See class documentation for description of standard arguments.
544

545
        Args:
546
          tmap: Map structure with shape (1,m,n) where each pixel,
547
            when converted to integer, species the index of the thread
548
            to which that pixel should be assigned.
549

550
        """
551
        projeng = self.get_ProjEng('T')
1✔
552
        q_native = self._cache_q_fp_to_native(assembly.Q)
1✔
553
        n_threads = mapthreads.get_num_threads(n_threads)
1✔
554
        omp_ivals = projeng.pixel_ranges(q_native, assembly.dets, tmap, n_threads)
1✔
555
        return wrap_ivals(omp_ivals)
1✔
556

557
    def get_active_tiles(self, assembly, assign=False):
1✔
558
        """For a tiled Projection, figure out what tiles are hit by an
559
        assembly.
560

561
        See class documentation for description of standard arguments.
562

563
        Args:
564
          assign (bool or int): If True or an int, then the function
565
            will also compute a thread assignment, based on assigning
566
            each tile to a particular thread.  If this is an int, it
567
            sets the thread count; if it is simply True then
568
            OMP_NUM_THREADS is used.
569

570
        Returns:
571
          dict : The basic analysis results are:
572

573
            - ``'active_tiles'`` (list of int): sorted, non-repeating
574
              list of tiles that are hit by the assembly.
575
            - ``'hit_counts'`` (list of int): the number of hits in
576
              each tile returned in 'active_tiles', respectively.
577

578
          If the thread assignment took place then the dict will also
579
          contain:
580

581
            - ``'group_tiles'`` (list of lists of ints): There is one
582
              entry per thread, and the entry lists the tiles that
583
              have been assigned to that thread.
584
            - ``'group_ranges'`` (RangesMatrix): The thread
585
              assignments; shape (n_thread, n_det, n_samp).
586

587
        """
588
        if self.tiling is None:
1✔
589
            raise RuntimeError("This Projectionist not set up for Tiled maps.")
1✔
590
        projeng = self.get_ProjEng('T')
1✔
591
        q_native = self._cache_q_fp_to_native(assembly.Q)
1✔
592
        # This returns a G3VectorInt of length n_tiles giving count of hits per tile.
593
        hits = np.array(projeng.tile_hits(q_native, assembly.dets))
1✔
594
        tiles = np.nonzero(hits)[0]
1✔
595
        hits = hits[tiles]
1✔
596
        if assign is True:
1✔
597
            assign = so3g.useful_info()['omp_num_threads']
×
598
        if assign > 0:
1✔
599
            group_n = np.array([0 for g in range(assign)])
1✔
600
            group_tiles = [[] for _ in group_n]
1✔
601
            cands = sorted(list(zip(hits, tiles)), reverse=True) # [(1000, 12), (900, 4), ...]
1✔
602
            for _n, _t in cands:
1✔
603
                idx = group_n.argmin()
1✔
604
                group_n[idx] += _n
1✔
605
                group_tiles[idx].append(_t)
1✔
606
            # Now paint them into Ranges.
607
            R = projeng.tile_ranges(q_native, assembly.dets, group_tiles)
1✔
608
            R = wrap_ivals(R)
1✔
609
            return {
1✔
610
                'active_tiles': list(tiles),
611
                'hit_counts': list(hits),
612
                'group_ranges': R,
613
                'group_tiles': group_tiles,
614
            }
615
        return {
1✔
616
            'active_tiles': list(tiles),
617
            'hit_counts': list(hits),
618
        }
619

620
    _ivals_format = 2
1✔
621

622

623
class _Tiling:
1✔
624
    """Utility class for computations involving tiled maps.
625

626
    """
627
    def __init__(self, shape, tile_shape):
1✔
628
        self.shape = shape[-2:]
1✔
629
        self.tile_shape = tile_shape
1✔
630
        nt0 = int(np.ceil(self.shape[0] / self.tile_shape[0]))
1✔
631
        nt1 = int(np.ceil(self.shape[1] / self.tile_shape[1]))
1✔
632
        self.super_shape = (nt0, nt1)
1✔
633
        self.tile_count = nt0 * nt1
1✔
634
    def tile_dims(self, tile):
1✔
635
        rowcol = self.tile_rowcol(tile)
1✔
636
        return (min(self.shape[0] - rowcol[0] * self.tile_shape[0], self.tile_shape[0]),
1✔
637
                min(self.shape[1] - rowcol[1] * self.tile_shape[1], self.tile_shape[1]))
638
    def tile_rowcol(self, tile):
1✔
639
        if tile >= self.tile_count:
1✔
640
            raise ValueError(f"Request for tile {tile} is out of bounds for "
×
641
                             f"super_shape={self.super_shape}")
642
        return (tile // self.super_shape[1], tile % self.super_shape[1])
1✔
643
    def tile_offset(self, tile):
1✔
644
        row, col = self.tile_rowcol(tile)
1✔
645
        return row * self.tile_shape[0], col * self.tile_shape[1]
1✔
646

647
def wrap_ivals(ivals):
1✔
648
    return RangesMatrix([RangesMatrix([RangesMatrix(y) for y in x]) for x in ivals])
1✔
649

650
THREAD_ASSIGNMENT_METHODS = [
1✔
651
    'simple',
652
    'domdir',
653
    'tiles',
654
]
STATUS · Troubleshooting · Open an Issue · Sales · Support · CAREERS · ENTERPRISE · START FREE · SCHEDULE DEMO
ANNOUNCEMENTS · TWITTER · TOS & SLA · Supported CI Services · What's a CI service? · Automated Testing

© 2025 Coveralls, Inc