test_geom.py

# Licensed under a 3-clause BSD style license - see LICENSE.rst
from __future__ import absolute_import, division, print_function, unicode_literals
import pytest
from collections import OrderedDict
import numpy as np
from numpy.testing import assert_allclose
from astropy import units as u
from astropy.coordinates import SkyCoord
from ..geom import MapAxis, MapCoord

pytest.importorskip('scipy')

mapaxis_geoms = [
    (np.array([0.25, 0.75, 1.0, 2.0]), 'lin'),
    (np.array([0.25, 0.75, 1.0, 2.0]), 'log'),
    (np.array([0.25, 0.75, 1.0, 2.0]), 'sqrt'),
]

mapaxis_geoms_node_type = [
    (np.array([0.25, 0.75, 1.0, 2.0]), 'lin', 'edge'),
    (np.array([0.25, 0.75, 1.0, 2.0]), 'log', 'edge'),
    (np.array([0.25, 0.75, 1.0, 2.0]), 'sqrt', 'edge'),
    (np.array([0.25, 0.75, 1.0, 2.0]), 'lin', 'center'),
    (np.array([0.25, 0.75, 1.0, 2.0]), 'log', 'center'),
    (np.array([0.25, 0.75, 1.0, 2.0]), 'sqrt', 'center'),
]


@pytest.mark.parametrize(('edges', 'interp'),
                         mapaxis_geoms)
def test_mapaxis_init_from_edges(edges, interp):
    axis = MapAxis(edges, interp=interp)
    assert_allclose(axis.edges, edges)
    assert_allclose(axis.nbin, len(edges) - 1)


@pytest.mark.parametrize(('nodes', 'interp'),
                         mapaxis_geoms)
def test_mapaxis_from_nodes(nodes, interp):
    axis = MapAxis.from_nodes(nodes, interp=interp)
    assert_allclose(axis.center, nodes)
    assert_allclose(axis.nbin, len(nodes))


@pytest.mark.parametrize(('nodes', 'interp'),
                         mapaxis_geoms)
def test_mapaxis_from_bounds(nodes, interp):
    axis = MapAxis.from_bounds(nodes[0], nodes[-1], 3,
                               interp=interp)
    assert_allclose(axis.edges[0], nodes[0])
    assert_allclose(axis.edges[-1], nodes[-1])
    assert_allclose(axis.nbin, 3)


@pytest.mark.parametrize(('nodes', 'interp', 'node_type'),
                         mapaxis_geoms_node_type)
def test_mapaxis_pix_to_coord(nodes, interp, node_type):
    axis = MapAxis(nodes, interp=interp, node_type=node_type)
    assert_allclose(axis.center,
                    axis.pix_to_coord(np.arange(axis.nbin, dtype=float)))
    assert_allclose(np.arange(axis.nbin + 1, dtype=float) - 0.5,
                    axis.coord_to_pix(axis.edges))


@pytest.mark.parametrize(('nodes', 'interp', 'node_type'),
                         mapaxis_geoms_node_type)
def test_mapaxis_coord_to_idx(nodes, interp, node_type):
    axis = MapAxis(nodes, interp=interp, node_type=node_type)
    assert_allclose(np.arange(axis.nbin, dtype=int),
                    axis.coord_to_idx(axis.center))


@pytest.mark.parametrize(('nodes', 'interp', 'node_type'),
                         mapaxis_geoms_node_type)
def test_mapaxis_slice(nodes, interp, node_type):
    axis = MapAxis(nodes, interp=interp, node_type=node_type)
    saxis = axis.slice(slice(1, 3))
    assert_allclose(saxis.nbin, 2)
    assert_allclose(saxis.center, axis.center[slice(1, 3)])

    axis = MapAxis(nodes, interp=interp, node_type=node_type)
    saxis = axis.slice(slice(1, None))
    assert_allclose(saxis.nbin, axis.nbin - 1)
    assert_allclose(saxis.center, axis.center[slice(1, None)])

    axis = MapAxis(nodes, interp=interp, node_type=node_type)
    saxis = axis.slice(slice(None, 2))
    assert_allclose(saxis.nbin, 2)
    assert_allclose(saxis.center, axis.center[slice(None, 2)])

    axis = MapAxis(nodes, interp=interp, node_type=node_type)
    saxis = axis.slice(slice(None, -1))
    assert_allclose(saxis.nbin, axis.nbin - 1)
    assert_allclose(saxis.center, axis.center[slice(None, -1)])


def test_mapcoords_create():
    # From existing MapCoord
    coords_cel = MapCoord.create((0.0, 1.0), coordsys='CEL')
    coords_gal = MapCoord.create(coords_cel, coordsys='GAL')
    assert_allclose(coords_gal.lon, coords_cel.skycoord.galactic.l.deg)
    assert_allclose(coords_gal.lat, coords_cel.skycoord.galactic.b.deg)

    # 2D Tuple of scalars
    coords = MapCoord.create((0.0, 1.0))
    assert_allclose(coords.lon, 0.0)
    assert_allclose(coords.lat, 1.0)
    assert_allclose(coords[0], 0.0)
    assert_allclose(coords[1], 1.0)
    assert coords.coordsys is None
    assert coords.ndim == 2

    # 3D Tuple of scalars
    coords = MapCoord.create((0.0, 1.0, 2.0))
    assert_allclose(coords[0], 0.0)
    assert_allclose(coords[1], 1.0)
    assert_allclose(coords[2], 2.0)
    assert coords.coordsys is None
    assert coords.ndim == 3

    # 2D Tuple w/ NaN coordinates
    coords = MapCoord.create((np.nan, np.nan))

    # 2D Tuple w/ NaN coordinates
    lon, lat = np.array([np.nan, 1.0]), np.array([np.nan, 3.0])
    coords = MapCoord.create((lon, lat))
    assert_allclose(coords.lon, lon)
    assert_allclose(coords.lat, lat)

    # 2D Tuple w/ SkyCoord
    lon, lat = np.array([0.0, 1.0]), np.array([2.0, 3.0])
    energy = np.array([100., 1000.])
    skycoord_cel = SkyCoord(lon, lat, unit='deg', frame='icrs')
    skycoord_gal = SkyCoord(lon, lat, unit='deg', frame='galactic')
    coords = MapCoord.create((skycoord_cel,))
    assert_allclose(coords.lon, lon)
    assert_allclose(coords.lat, lat)
    assert coords.coordsys == 'CEL'
    assert coords.ndim == 2
    coords = MapCoord.create((skycoord_gal,))
    assert_allclose(coords.lon, lon)
    assert_allclose(coords.lat, lat)
    assert coords.coordsys == 'GAL'
    assert coords.ndim == 2

    # SkyCoord
    coords = MapCoord.create(skycoord_cel)
    assert_allclose(coords.lon, lon)
    assert_allclose(coords.lat, lat)
    assert coords.coordsys == 'CEL'
    assert coords.ndim == 2
    coords = MapCoord.create(skycoord_gal)
    assert_allclose(coords.lon, lon)
    assert_allclose(coords.lat, lat)
    assert coords.coordsys == 'GAL'
    assert coords.ndim == 2

    # 2D Dict w/ vectors
    coords = MapCoord.create(dict(lon=lon, lat=lat))
    assert_allclose(coords.lon, lon)
    assert_allclose(coords.lat, lat)
    assert coords.ndim == 2

    # 3D Dict w/ vectors
    coords = MapCoord.create(dict(lon=lon, lat=lat, energy=energy))
    assert_allclose(coords.lon, lon)
    assert_allclose(coords.lat, lat)
    assert_allclose(coords['energy'], energy)
    assert coords.ndim == 3

    # 3D Dict w/ SkyCoord
    coords = MapCoord.create(dict(skycoord=skycoord_cel, energy=energy))
    assert_allclose(coords.lon, lon)
    assert_allclose(coords.lat, lat)
    assert_allclose(coords['energy'], energy)
    assert coords.ndim == 3

    # 3D OrderedDict w/ vectors
    coords = MapCoord.create(dict(energy=energy, lat=lat, lon=lon))
    assert_allclose(coords.lon, lon)
    assert_allclose(coords.lat, lat)
    assert_allclose(coords['energy'], energy)
    assert_allclose(coords[0], energy)
    assert_allclose(coords[1], lat)
    assert_allclose(coords[2], lon)
    assert coords.ndim == 3

    # Quantities
    coords = MapCoord.create(dict(energy=energy * u.TeV, lat=lat, lon=lon))
    assert coords['energy'].unit == 'TeV'


def test_mapcoords_to_coordsys():
    lon, lat = np.array([0.0, 1.0]), np.array([2.0, 3.0])
    energy = np.array([100., 1000.])
    skycoord_cel = SkyCoord(lon, lat, unit='deg', frame='icrs')
    skycoord_gal = SkyCoord(lon, lat, unit='deg', frame='galactic')

    coords = MapCoord.create(
        dict(lon=lon, lat=lat, energy=energy), coordsys='CEL')
    assert coords.coordsys == 'CEL'
    assert_allclose(coords.skycoord.transform_to(
        'icrs').ra.deg, skycoord_cel.ra.deg)
    assert_allclose(coords.skycoord.transform_to(
        'icrs').dec.deg, skycoord_cel.dec.deg)
    coords = coords.to_coordsys('GAL')
    assert coords.coordsys == 'GAL'
    assert_allclose(coords.skycoord.transform_to(
        'galactic').l.deg, skycoord_cel.galactic.l.deg)
    assert_allclose(coords.skycoord.transform_to(
        'galactic').b.deg, skycoord_cel.galactic.b.deg)

    coords = MapCoord.create(
        dict(lon=lon, lat=lat, energy=energy), coordsys='GAL')
    assert coords.coordsys == 'GAL'
    assert_allclose(coords.skycoord.transform_to(
        'galactic').l.deg, skycoord_gal.l.deg)
    assert_allclose(coords.skycoord.transform_to(
        'galactic').b.deg, skycoord_gal.b.deg)
    coords = coords.to_coordsys('CEL')
    assert coords.coordsys == 'CEL'
    assert_allclose(coords.skycoord.transform_to(
        'icrs').ra.deg, skycoord_gal.icrs.ra.deg)
    assert_allclose(coords.skycoord.transform_to(
        'icrs').dec.deg, skycoord_gal.icrs.dec.deg)


def test_mapaxis_repr():
    axis = MapAxis([1, 2, 3], name='test')
    assert 'MapAxis' in repr(axis)


def test_mapcoord_repr():
    coord = MapCoord({'lon': 0, 'lat': 0, 'energy': 5})
    assert 'MapCoord' in repr(coord)