# -*- coding: utf-8 -*- # Licensed under a 3-clause BSD style license - see LICENSE.rst """Quantity helpers for the ERFA ufuncs.""" # Tests for these are in coordinates, not in units. from erfa import ufunc as erfa_ufunc, dt_pv, dt_eraLDBODY, dt_eraASTROM from astropy.units.core import UnitsError, UnitTypeError, dimensionless_unscaled from astropy.units.structured import StructuredUnit from . import UFUNC_HELPERS from .helpers import (get_converter, helper_invariant, helper_multiplication, helper_twoarg_invariant, _d) erfa_ufuncs = ('s2c', 's2p', 'c2s', 'p2s', 'pm', 'pdp', 'pxp', 'rxp', 'cpv', 'p2pv', 'pv2p', 'pv2s', 'pvdpv', 'pvm', 'pvmpv', 'pvppv', 'pvstar', 'pvtob', 'pvu', 'pvup', 'pvxpv', 'rxpv', 's2pv', 's2xpv', 'starpv', 'sxpv', 'trxpv', 'gd2gc', 'gc2gd', 'ldn', 'aper', 'apio', 'atciq', 'atciqn', 'atciqz', 'aticq', 'atioq', 'atoiq') def has_matching_structure(unit, dtype): dtype_fields = dtype.fields if dtype_fields: return (isinstance(unit, StructuredUnit) and len(unit) == len(dtype_fields) and all(has_matching_structure(u, df_v[0]) for (u, df_v) in zip(unit.values(), dtype_fields.values()))) else: return not isinstance(unit, StructuredUnit) def check_structured_unit(unit, dtype): if not has_matching_structure(unit, dtype): msg = {dt_pv: 'pv', dt_eraLDBODY: 'ldbody', dt_eraASTROM: 'astrom'}.get(dtype, 'function') raise UnitTypeError(f'{msg} input needs unit matching dtype={dtype}.') def helper_s2c(f, unit1, unit2): from astropy.units.si import radian try: return [get_converter(unit1, radian), get_converter(unit2, radian)], dimensionless_unscaled except UnitsError: raise UnitTypeError("Can only apply '{}' function to " "quantities with angle units" .format(f.__name__)) def helper_s2p(f, unit1, unit2, unit3): from astropy.units.si import radian try: return [get_converter(unit1, radian), get_converter(unit2, radian), None], unit3 except UnitsError: raise UnitTypeError("Can only apply '{}' function to " "quantities with angle units" .format(f.__name__)) def helper_c2s(f, unit1): from astropy.units.si import radian return [None], (radian, radian) def helper_p2s(f, unit1): from astropy.units.si import radian return [None], (radian, radian, unit1) def helper_gc2gd(f, nounit, unit1): from astropy.units.si import m, radian if nounit is not None: raise UnitTypeError("ellipsoid cannot be a quantity.") try: return [None, get_converter(unit1, m)], (radian, radian, m, None) except UnitsError: raise UnitTypeError("Can only apply '{}' function to " "quantities with length units" .format(f.__name__)) def helper_gd2gc(f, nounit, unit1, unit2, unit3): from astropy.units.si import m, radian if nounit is not None: raise UnitTypeError("ellipsoid cannot be a quantity.") try: return [None, get_converter(unit1, radian), get_converter(unit2, radian), get_converter(unit3, m)], (m, None) except UnitsError: raise UnitTypeError("Can only apply '{}' function to lon, lat " "with angle and height with length units" .format(f.__name__)) def helper_p2pv(f, unit1): from astropy.units.si import s if isinstance(unit1, StructuredUnit): raise UnitTypeError("p vector unit cannot be a structured unit.") return [None], StructuredUnit((unit1, unit1 / s)) def helper_pv2p(f, unit1): check_structured_unit(unit1, dt_pv) return [None], unit1[0] def helper_pv2s(f, unit_pv): from astropy.units.si import radian check_structured_unit(unit_pv, dt_pv) ang_unit = radian * unit_pv[1] / unit_pv[0] return [None], (radian, radian, unit_pv[0], ang_unit, ang_unit, unit_pv[1]) def helper_s2pv(f, unit_theta, unit_phi, unit_r, unit_td, unit_pd, unit_rd): from astropy.units.si import radian time_unit = unit_r / unit_rd return [get_converter(unit_theta, radian), get_converter(unit_phi, radian), None, get_converter(unit_td, radian / time_unit), get_converter(unit_pd, radian / time_unit), None], StructuredUnit((unit_r, unit_rd)) def helper_pv_multiplication(f, unit1, unit2): check_structured_unit(unit1, dt_pv) check_structured_unit(unit2, dt_pv) result_unit = StructuredUnit((unit1[0] * unit2[0], unit1[1] * unit2[0])) converter = get_converter(unit2, StructuredUnit( (unit2[0], unit1[1] * unit2[0] / unit1[0]))) return [None, converter], result_unit def helper_pvm(f, unit1): check_structured_unit(unit1, dt_pv) return [None], (unit1[0], unit1[1]) def helper_pvstar(f, unit1): from astropy.units.astrophys import AU from astropy.units.si import km, s, radian, day, year, arcsec return [get_converter(unit1, StructuredUnit((AU, AU/day)))], ( radian, radian, radian / year, radian / year, arcsec, km / s, None) def helper_starpv(f, unit_ra, unit_dec, unit_pmr, unit_pmd, unit_px, unit_rv): from astropy.units.si import km, s, day, year, radian, arcsec from astropy.units.astrophys import AU return [get_converter(unit_ra, radian), get_converter(unit_dec, radian), get_converter(unit_pmr, radian/year), get_converter(unit_pmd, radian/year), get_converter(unit_px, arcsec), get_converter(unit_rv, km/s)], (StructuredUnit((AU, AU/day)), None) def helper_pvtob(f, unit_elong, unit_phi, unit_hm, unit_xp, unit_yp, unit_sp, unit_theta): from astropy.units.si import m, s, radian return [get_converter(unit_elong, radian), get_converter(unit_phi, radian), get_converter(unit_hm, m), get_converter(unit_xp, radian), get_converter(unit_yp, radian), get_converter(unit_sp, radian), get_converter(unit_theta, radian)], StructuredUnit((m, m/s)) def helper_pvu(f, unit_t, unit_pv): check_structured_unit(unit_pv, dt_pv) return [get_converter(unit_t, unit_pv[0]/unit_pv[1]), None], unit_pv def helper_pvup(f, unit_t, unit_pv): check_structured_unit(unit_pv, dt_pv) return [get_converter(unit_t, unit_pv[0]/unit_pv[1]), None], unit_pv[0] def helper_s2xpv(f, unit1, unit2, unit_pv): check_structured_unit(unit_pv, dt_pv) return [None, None, None], StructuredUnit((_d(unit1) * unit_pv[0], _d(unit2) * unit_pv[1])) def ldbody_unit(): from astropy.units.si import day, radian from astropy.units.astrophys import Msun, AU return StructuredUnit((Msun, radian, (AU, AU/day)), erfa_ufunc.dt_eraLDBODY) def astrom_unit(): from astropy.units.si import rad, year from astropy.units.astrophys import AU one = rel2c = dimensionless_unscaled return StructuredUnit((year, AU, one, AU, rel2c, one, one, rad, rad, rad, rad, one, one, rel2c, rad, rad, rad), erfa_ufunc.dt_eraASTROM) def helper_ldn(f, unit_b, unit_ob, unit_sc): from astropy.units.astrophys import AU return [get_converter(unit_b, ldbody_unit()), get_converter(unit_ob, AU), get_converter(_d(unit_sc), dimensionless_unscaled)], dimensionless_unscaled def helper_aper(f, unit_theta, unit_astrom): check_structured_unit(unit_astrom, dt_eraASTROM) unit_along = unit_astrom[7] # along if unit_astrom[14] is unit_along: # eral result_unit = unit_astrom else: result_units = tuple((unit_along if i == 14 else v) for i, v in enumerate(unit_astrom.values())) result_unit = unit_astrom.__class__(result_units, names=unit_astrom) return [get_converter(unit_theta, unit_along), None], result_unit def helper_apio(f, unit_sp, unit_theta, unit_elong, unit_phi, unit_hm, unit_xp, unit_yp, unit_refa, unit_refb): from astropy.units.si import radian, m return [get_converter(unit_sp, radian), get_converter(unit_theta, radian), get_converter(unit_elong, radian), get_converter(unit_phi, radian), get_converter(unit_hm, m), get_converter(unit_xp, radian), get_converter(unit_xp, radian), get_converter(unit_xp, radian), get_converter(unit_xp, radian)], astrom_unit() def helper_atciq(f, unit_rc, unit_dc, unit_pr, unit_pd, unit_px, unit_rv, unit_astrom): from astropy.units.si import radian, arcsec, year, km, s return [get_converter(unit_rc, radian), get_converter(unit_dc, radian), get_converter(unit_pr, radian / year), get_converter(unit_pd, radian / year), get_converter(unit_px, arcsec), get_converter(unit_rv, km / s), get_converter(unit_astrom, astrom_unit())], (radian, radian) def helper_atciqn(f, unit_rc, unit_dc, unit_pr, unit_pd, unit_px, unit_rv, unit_astrom, unit_b): from astropy.units.si import radian, arcsec, year, km, s return [get_converter(unit_rc, radian), get_converter(unit_dc, radian), get_converter(unit_pr, radian / year), get_converter(unit_pd, radian / year), get_converter(unit_px, arcsec), get_converter(unit_rv, km / s), get_converter(unit_astrom, astrom_unit()), get_converter(unit_b, ldbody_unit())], (radian, radian) def helper_atciqz_aticq(f, unit_rc, unit_dc, unit_astrom): from astropy.units.si import radian return [get_converter(unit_rc, radian), get_converter(unit_dc, radian), get_converter(unit_astrom, astrom_unit())], (radian, radian) def helper_aticqn(f, unit_rc, unit_dc, unit_astrom, unit_b): from astropy.units.si import radian return [get_converter(unit_rc, radian), get_converter(unit_dc, radian), get_converter(unit_astrom, astrom_unit()), get_converter(unit_b, ldbody_unit())], (radian, radian) def helper_atioq(f, unit_rc, unit_dc, unit_astrom): from astropy.units.si import radian return [get_converter(unit_rc, radian), get_converter(unit_dc, radian), get_converter(unit_astrom, astrom_unit())], (radian,)*5 def helper_atoiq(f, unit_type, unit_ri, unit_di, unit_astrom): from astropy.units.si import radian if unit_type is not None: raise UnitTypeError("argument 'type' should not have a unit") return [None, get_converter(unit_ri, radian), get_converter(unit_di, radian), get_converter(unit_astrom, astrom_unit())], (radian, radian) def get_erfa_helpers(): ERFA_HELPERS = {} ERFA_HELPERS[erfa_ufunc.s2c] = helper_s2c ERFA_HELPERS[erfa_ufunc.s2p] = helper_s2p ERFA_HELPERS[erfa_ufunc.c2s] = helper_c2s ERFA_HELPERS[erfa_ufunc.p2s] = helper_p2s ERFA_HELPERS[erfa_ufunc.pm] = helper_invariant ERFA_HELPERS[erfa_ufunc.cpv] = helper_invariant ERFA_HELPERS[erfa_ufunc.p2pv] = helper_p2pv ERFA_HELPERS[erfa_ufunc.pv2p] = helper_pv2p ERFA_HELPERS[erfa_ufunc.pv2s] = helper_pv2s ERFA_HELPERS[erfa_ufunc.pvdpv] = helper_pv_multiplication ERFA_HELPERS[erfa_ufunc.pvxpv] = helper_pv_multiplication ERFA_HELPERS[erfa_ufunc.pvm] = helper_pvm ERFA_HELPERS[erfa_ufunc.pvmpv] = helper_twoarg_invariant ERFA_HELPERS[erfa_ufunc.pvppv] = helper_twoarg_invariant ERFA_HELPERS[erfa_ufunc.pvstar] = helper_pvstar ERFA_HELPERS[erfa_ufunc.pvtob] = helper_pvtob ERFA_HELPERS[erfa_ufunc.pvu] = helper_pvu ERFA_HELPERS[erfa_ufunc.pvup] = helper_pvup ERFA_HELPERS[erfa_ufunc.pdp] = helper_multiplication ERFA_HELPERS[erfa_ufunc.pxp] = helper_multiplication ERFA_HELPERS[erfa_ufunc.rxp] = helper_multiplication ERFA_HELPERS[erfa_ufunc.rxpv] = helper_multiplication ERFA_HELPERS[erfa_ufunc.s2pv] = helper_s2pv ERFA_HELPERS[erfa_ufunc.s2xpv] = helper_s2xpv ERFA_HELPERS[erfa_ufunc.starpv] = helper_starpv ERFA_HELPERS[erfa_ufunc.sxpv] = helper_multiplication ERFA_HELPERS[erfa_ufunc.trxpv] = helper_multiplication ERFA_HELPERS[erfa_ufunc.gc2gd] = helper_gc2gd ERFA_HELPERS[erfa_ufunc.gd2gc] = helper_gd2gc ERFA_HELPERS[erfa_ufunc.ldn] = helper_ldn ERFA_HELPERS[erfa_ufunc.aper] = helper_aper ERFA_HELPERS[erfa_ufunc.apio] = helper_apio ERFA_HELPERS[erfa_ufunc.atciq] = helper_atciq ERFA_HELPERS[erfa_ufunc.atciqn] = helper_atciqn ERFA_HELPERS[erfa_ufunc.atciqz] = helper_atciqz_aticq ERFA_HELPERS[erfa_ufunc.aticq] = helper_atciqz_aticq ERFA_HELPERS[erfa_ufunc.aticqn] = helper_aticqn ERFA_HELPERS[erfa_ufunc.atioq] = helper_atioq ERFA_HELPERS[erfa_ufunc.atoiq] = helper_atoiq return ERFA_HELPERS UFUNC_HELPERS.register_module('erfa.ufunc', erfa_ufuncs, get_erfa_helpers)