aacgmv2.wrapper.convert_latlon_arr()

Complexity

Conditions

20

Size

Total Lines	129
Code Lines	60

Duplication

Lines	0
Ratio	0 %

Importance

Changes

0

# -*- coding: utf-8 -*-
def convert_latlon_arr(in_lat, in_lon, height, dtime, code="G2A"):
    """Converts between geomagnetic coordinates and AACGM coordinates.

    Parameters
    ------------
    in_lat : (np.ndarray or list or float)
        Input latitude in degrees N (code specifies type of latitude)
    in_lon : (np.ndarray or list or float)
        Input longitude in degrees E (code specifies type of longitude)
    height : (np.ndarray or list or float)
        Altitude above the surface of the earth in km
    dtime : (datetime)
        Single datetime object for magnetic field
    code : (int or str)
        Bit code or string denoting which type(s) of conversion to perform
        G2A        - geographic (geodetic) to AACGM-v2
        A2G        - AACGM-v2 to geographic (geodetic)
        TRACE      - use field-line tracing, not coefficients
        ALLOWTRACE - use trace only above 2000 km
        BADIDEA    - use coefficients above 2000 km
        GEOCENTRIC - assume inputs are geocentric w/ RE=6371.2
        (default = "G2A")

    Returns
    -------
    out_lat : (np.ndarray)
        Output latitudes in degrees N
    out_lon : (np.ndarray)
        Output longitudes in degrees E
    out_r : (np.ndarray)
        Geocentric radial distance (R_Earth) or altitude above the surface of
        the Earth (km)

    Notes
    -------
    At least one of in_lat, in_lon, and height must be a list or array.
    """
    import aacgmv2._aacgmv2 as c_aacgmv2
    import aacgmv2

    # Recast the data as numpy arrays
    in_lat = np.array(in_lat)
    in_lon = np.array(in_lon)
    height = np.array(height)

    # If one or two of these elements is a float or int, create an array
    test_array = np.array([len(in_lat.shape), len(in_lon.shape),
                           len(height.shape)])
    if test_array.min() == 0:
        if test_array.max() == 0:
            aacgmv2.logger.warning("for a single location, consider using " \
                                   "convert_latlon or get_aacgm_coord")
            in_lat = np.array([in_lat])
            in_lon = np.array([in_lon])
            height = np.array([height])
        else:
            imax = test_array.argmax()
            max_shape = in_lat.shape if imax == 0 else (in_lon.shape \
                                            if imax == 1 else height.shape) 
            if test_array[0] == 0:
                in_lat = np.full(shape=max_shape, fill_value=in_lat)
            if not test_array[1]:
                in_lon = np.full(shape=max_shape, fill_value=in_lon)
            if not test_array[2]:
                height = np.full(shape=max_shape, fill_value=height)

    # Ensure that lat, lon, and height are the same length or if the lengths
    # differ that the different ones contain only a single value
    if not (in_lat.shape == in_lon.shape and in_lat.shape == height.shape):
        ulen = np.unique([in_lat.shape, in_lon.shape, height.shape])
        if ulen.min() != (1,):
            raise ValueError('mismatched input arrays')

    # Test time
    if isinstance(dtime, dt.date):
        dtime = dt.datetime.combine(dtime, dt.time(0))

    if not isinstance(dtime, dt.datetime):
        raise ValueError('time must be specified as datetime object')

    # Initialise output
    lat_out = np.full(shape=in_lat.shape, fill_value=np.nan)
    lon_out = np.full(shape=in_lon.shape, fill_value=np.nan)
    r_out = np.full(shape=height.shape, fill_value=np.nan)

    # Test and set code
    try:
        bit_code = convert_str_to_bit(code.upper())
    except AttributeError:
        bit_code = code

    if not isinstance(bit_code, int):
        raise ValueError("unknown code {:}".format(bit_code))

    # Test height
    good_height = test_height(np.nanmax(height), bit_code)

    if not good_height:
        return lat_out, lon_out, r_out

    # Test latitude range
    if np.abs(in_lat).max() > 90.0:
        if np.abs(in_lat).max() > 90.1:
            raise ValueError('unrealistic latitude')
        in_lat = np.clip(in_lat, -90.0, 90.0)

    # Constrain longitudes between -180 and 180
    in_lon = ((in_lon + 180.0) % 360.0) - 180.0

    
    # Set current date and time
    try:
        c_aacgmv2.set_datetime(dtime.year, dtime.month, dtime.day, dtime.hour,
                               dtime.minute, dtime.second)
    except:
        raise RuntimeError("unable to set time for {:}".format(dtime))

    # Vectorise the AACGM code
    convert_vectorised = np.vectorize(c_aacgmv2.convert)

    # convert
    try:
        lat_out, lon_out, r_out = convert_vectorised(in_lat, in_lon, height,
                                                     bit_code)

    except:
        pass

    return lat_out, lon_out, r_out