text:
#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Copyright (C) 2017, AGB & GC
# Full license can be found in License.md
# -----------------------------------------------------------------------------
""" Tests the ocb_scaling class and functions
"""

from io import StringIO
import logging
import numpy as np
from os import path
from sys import version_info
import unittest

import ocbpy

class TestOCBScalingLogFailure(unittest.TestCase):
    def setUp(self):
        """ Initialize the test class
        """
        # Initialize the logging info
        self.lwarn = u""
        self.lout = u""
        self.log_capture = StringIO()
        ocbpy.logger.addHandler(logging.StreamHandler(self.log_capture))
        ocbpy.logger.setLevel(logging.INFO)

        # Initialize the testing variables
        test_file = path.join(path.dirname(ocbpy.__file__), "tests",
                              "test_data", "test_north_circle")
        self.assertTrue(path.isfile(test_file))
        self.ocb = ocbpy.ocboundary.OCBoundary(filename=test_file,
                                               instrument='image')
        self.ocb.rec_ind = 27
        self.vdata = ocbpy.ocb_scaling.VectorData(0, self.ocb.rec_ind, 75.0,
                                                  22.0, aacgm_n=50.0,
                                                  aacgm_e=86.5, aacgm_z=5.0,
                                                  dat_name="Test",
                                                  dat_units="$m s^{-1}$")

    def tearDown(self):
        """ Tear down the test case
        """
        del self.lwarn, self.lout, self.log_capture, self.ocb, self.vdata

    def test_no_scale_func(self):
        """ Test OCBScaling initialization with no scaling function
        """
        self.lwarn = u"no scaling function provided"

        # Initialize the VectorData class without a scaling function
        self.vdata.set_ocb(self.ocb)
        self.assertIsNone(self.vdata.scale_func)

        self.lout = self.log_capture.getvalue()
        # Test logging error message for each bad initialization
        self.assertTrue(self.lout.find(self.lwarn) >= 0)

    def test_inconsistent_vector_warning(self):
        """ Test init failure with inconsistent AACGM components
        """
        self.lwarn = u"inconsistent AACGM"

        # Initalize the VectorData class with inconsistent vector magnitudes
        self.vdata = ocbpy.ocb_scaling.VectorData(0, self.ocb.rec_ind,
                                                  75.0, 22.0,
                                                  aacgm_mag=100.0,
                                                  dat_name="Test",
                                                  dat_units="$m s^{-1}$")

        self.lout = self.log_capture.getvalue()
        # Test logging error message for each bad initialization
        self.assertTrue(self.lout.find(self.lwarn) >= 0)

class TestOCBScalingMethods(unittest.TestCase):

    def setUp(self):
        """ Initialize the OCBoundary and VectorData objects
        """

        test_file = path.join(path.dirname(ocbpy.__file__), "tests",
                              "test_data", "test_north_circle")
        self.assertTrue(path.isfile(test_file))
        self.ocb = ocbpy.ocboundary.OCBoundary(filename=test_file,
                                               instrument='image')
        self.ocb.rec_ind = 27
        self.vdata = ocbpy.ocb_scaling.VectorData(0, self.ocb.rec_ind, 75.0,
                                                  22.0, aacgm_n=50.0,
                                                  aacgm_e=86.5, aacgm_z=5.0,
                                                  dat_name="Test",
                                                  dat_units="$m s^{-1}$")
        self.wdata = ocbpy.ocb_scaling.VectorData(0, self.ocb.rec_ind, 75.0,
                                                  22.0, aacgm_n=50.0,
                                                  aacgm_e=86.5, aacgm_z=5.0,
                                                  aacgm_mag=100.036243432,
                                                  dat_name="Test",
                                                  dat_units="$m s^{-1}$")
        self.zdata = ocbpy.ocb_scaling.VectorData(0, self.ocb.rec_ind, 87.2,
                                                  21.22, aacgm_n=0.0,
                                                  aacgm_e=0.0,
                                                  dat_name="Test Zero",
                                                  dat_units="$m s^{-1}$")

        if version_info.major == 2:
            self.assertRegex = self.assertRegexpMatches

    def tearDown(self):
        del self.ocb, self.vdata, self.wdata, self.zdata

    def test_init_nez(self):
        """ Test the initialisation of the VectorData object without magnitude
        """
        self.assertAlmostEqual(self.vdata.aacgm_mag, 100.036243432)
        self.assertAlmostEqual(self.zdata.aacgm_mag, 0.0)

    def test_init_mag(self):
        """ Test the initialisation of the VectorData object with magnitude
        """
        self.assertAlmostEqual(self.wdata.aacgm_mag, 100.036243432)

    def test_vector_repr_str(self):
        """ Test the VectorData print statement using repr and str
        """
        self.assertTrue(self.vdata.__repr__() == self.vdata.__str__())

    def test_vector_repr_no_scaling(self):
        """ Test the VectorData print statement without a scaling function
        """
        out = self.vdata.__repr__()

        self.assertRegex(out, "Vector data:")
        self.assertRegex(out, "No magnitude scaling function")
        del out

    def test_vector_repr_with_scaling(self):
        """ Test the VectorData print statement with a scaling function
        """
        self.vdata.set_ocb(self.ocb, scale_func=ocbpy.ocb_scaling.normal_evar)
        out = self.vdata.__repr__()

        self.assertRegex(out, "Vector data:")
        self.assertRegex(out, "Scaling function")

    def test_vector_bad_lat(self):
        """ Test the VectorData output with data from the wrong hemisphere
        """
        self.vdata.aacgm_lat *= -1.0
        self.vdata.set_ocb(self.ocb, scale_func=ocbpy.ocb_scaling.normal_evar)

        self.assertTrue(np.isnan(self.vdata.ocb_lat))
        self.assertTrue(np.isnan(self.vdata.ocb_mlt))
        self.assertTrue(np.isnan(self.vdata.r_corr))
        self.assertTrue(np.isnan(self.vdata.ocb_n))
        self.assertTrue(np.isnan(self.vdata.ocb_e))
        self.assertTrue(np.isnan(self.vdata.ocb_z))

    def test_calc_large_pole_angle(self):
        """ Test the OCB polar angle calculation with angles > 90 deg
        """
        self.zdata.ocb_aacgm_mlt = 1.260677777
        self.zdata.ocb_aacgm_lat = 83.99
        self.zdata.ocb_lat = 84.838777192
        self.zdata.ocb_mlt = 15.1110383783

        self.zdata.calc_vec_pole_angle()
        self.assertAlmostEqual(self.zdata.pole_angle, 91.72024697182087)

    def test_calc_vec_pole_angle_acute(self):
        """ Test the polar angle calculation with an acute angle
        """
        self.vdata.set_ocb(self.ocb)
        self.assertAlmostEqual(self.vdata.pole_angle, 8.67527923)

    def test_calc_vec_pole_angle_zero(self):
        """ Test the polar angle calculation with an angle of zero
        """
        self.vdata.set_ocb(self.ocb)
        self.vdata.aacgm_mlt = self.vdata.ocb_aacgm_mlt
        self.vdata.calc_vec_pole_angle()
        self.assertEqual(self.vdata.pole_angle, 0.0)

    def test_calc_vec_pole_angle_flat(self):
        """ Test the polar angle calculation with an angle of 180 deg
        """
        self.vdata.set_ocb(self.ocb)
        self.vdata.ocb_aacgm_mlt = 6.0
        self.vdata.aacgm_mlt = 6.0
        self.vdata.aacgm_lat = 45.0 + 0.5 * self.vdata.ocb_aacgm_lat
        self.vdata.calc_vec_pole_angle()
        self.assertEqual(self.vdata.pole_angle, 180.0)

    def test_calc_vec_pole_angle_right_isosceles(self):
        """ Test the polar angle calculation with a right isosceles triangle
        """
        # Set the distance between the data point and the OCB is equal to the
        # distance between the AACGM pole and the OCB so that the triangles
        # we're examining are isosceles triangles.  If the triangles were flat,
        # the angle would be 45 degrees
        self.vdata.set_ocb(self.ocb)
        self.vdata.ocb_aacgm_mlt = 0.0
        self.vdata.aacgm_mlt = 6.0
        self.vdata.aacgm_lat = self.vdata.ocb_aacgm_lat
        self.vdata.calc_vec_pole_angle()
        self.assertAlmostEqual(self.vdata.pole_angle, 45.03325090532819)

    def test_calc_vec_pole_angle_oblique(self):
        """ Test the polar angle calculation with an isosceles triangle
        """
        self.vdata.set_ocb(self.ocb)
        self.vdata.aacgm_mlt = self.vdata.ocb_aacgm_mlt - 1.0
        self.vdata.aacgm_lat = 45.0 + 0.5 * self.vdata.ocb_aacgm_lat
        self.vdata.calc_vec_pole_angle()
        self.assertAlmostEqual(self.vdata.pole_angle, 150.9561733411)

    def test_define_quadrants(self):
        """ Test the assignment of quadrants
        """
        # Set the initial values
        self.vdata.ocb_aacgm_mlt = self.ocb.phi_cent[self.vdata.ocb_ind] / 15.0
        self.vdata.ocb_aacgm_lat = 90.0 - self.ocb.r_cent[self.vdata.ocb_ind]
        (self.vdata.ocb_lat, self.vdata.ocb_mlt,
         self.vdata.r_corr) = self.ocb.normal_coord(self.vdata.aacgm_lat,
                                                    self.vdata.aacgm_mlt)
        self.vdata.calc_vec_pole_angle()

        # Get the test quadrants
        self.vdata.define_quadrants()
        self.assertEqual(self.vdata.ocb_quad, 1)
        self.assertEqual(self.vdata.vec_quad, 1)

    def test_define_quadrants_neg_adj_mlt_west(self):
        """ Test the quadrant assignment with a negative AACGM MLT and W vect
        """
        self.vdata.aacgm_mlt = -22.0
        self.vdata.aacgm_e *= -1.0
        self.vdata.set_ocb(self.ocb)
        self.assertGreater(self.vdata.ocb_aacgm_mlt-self.vdata.aacgm_mlt, 24)
        self.assertEqual(self.vdata.ocb_quad, 1)
        self.assertEqual(self.vdata.vec_quad, 2)

    def test_define_quadrants_neg_north(self):
        """ Test the quadrant assignment with a vector pointing south
        """
        self.vdata.aacgm_n *= -1.0
        self.vdata.set_ocb(self.ocb, scale_func=ocbpy.ocb_scaling.normal_evar)
        self.assertEqual(self.vdata.ocb_quad, 1)
        self.assertEqual(self.vdata.vec_quad, 4)

    def test_define_quadrants_noon_north(self):
        """ Test the quadrant assignment with a vector pointing north from noon
        """
        self.vdata.aacgm_mlt = 12.0
        self.vdata.set_ocb(self.ocb, scale_func=ocbpy.ocb_scaling.normal_evar)
        self.assertEqual(self.vdata.ocb_quad, 2)
        self.assertEqual(self.vdata.vec_quad, 1)

    def test_define_quadrants_aligned_poles_southwest(self):
        """ Test quad assignment w/vector pointing SW and both poles aligned
        """
        self.vdata.set_ocb(self.ocb, scale_func=ocbpy.ocb_scaling.normal_evar)
        self.vdata.aacgm_mlt = self.vdata.ocb_aacgm_mlt + 12.0
        self.vdata.aacgm_n = -10.0
        self.vdata.aacgm_e = -10.0
        self.vdata.set_ocb(self.ocb, scale_func=ocbpy.ocb_scaling.normal_evar)
        self.assertEqual(self.vdata.ocb_quad, 2)
        self.assertEqual(self.vdata.vec_quad, 3)

    def test_define_quadrants_ocb_south_night(self):
        """ Test the quadrant assignment with the OCB pole in a south/night quad
        """
        self.vdata.aacgm_mlt = 0.0
        self.vdata.set_ocb(self.ocb, scale_func=ocbpy.ocb_scaling.normal_evar)
        self.vdata.ocb_aacgm_mlt = 23.0
        self.vdata.ocb_aacgm_lat = self.vdata.aacgm_lat - 2.0
        self.vdata.calc_vec_pole_angle()
        self.vdata.define_quadrants()
        self.assertEqual(self.vdata.ocb_quad, 3)
        self.assertEqual(self.vdata.vec_quad, 1)

    def test_define_quadrants_ocb_south_day(self):
        """ Test the quadrant assignment with the OCB pole in a south/day quad
        """
        self.vdata.aacgm_mlt = 0.0
        self.vdata.set_ocb(self.ocb, scale_func=ocbpy.ocb_scaling.normal_evar)
        self.vdata.ocb_aacgm_mlt = 1.0
        self.vdata.ocb_aacgm_lat = self.vdata.aacgm_lat - 2.0
        self.vdata.calc_vec_pole_angle()
        self.vdata.define_quadrants()
        self.assertEqual(self.vdata.ocb_quad, 4)
        self.assertEqual(self.vdata.vec_quad, 1)

    def test_undefinable_quadrants(self):
        """ Test OCBScaling initialization for undefinable quadrants
        """
        self.vdata.aacgm_lat = 0.0
        self.vdata.set_ocb(self.ocb, scale_func=ocbpy.ocb_scaling.normal_evar)
        self.assertEqual(self.vdata.ocb_quad, 0)
        self.assertEqual(self.vdata.vec_quad, 0)

    def test_lost_ocb_quadrant(self):
        """ Test OCBScaling initialization for unset quadrants
        """
        self.vdata.set_ocb(self.ocb, scale_func=ocbpy.ocb_scaling.normal_evar)
        self.assertEqual(self.vdata.ocb_quad, 1)
        self.assertEqual(self.vdata.vec_quad, 1)
        self.vdata.ocb_quad = 0
        self.vdata.scale_vector()
        self.assertEqual(self.vdata.ocb_quad, 1)

    def test_lost_vec_quadrant(self):
        """ Test OCBScaling initialization for unset quadrants
        """
        self.vdata.set_ocb(self.ocb, scale_func=ocbpy.ocb_scaling.normal_evar)
        self.assertEqual(self.vdata.ocb_quad, 1)
        self.assertEqual(self.vdata.vec_quad, 1)
        self.vdata.vec_quad = 0
        self.vdata.scale_vector()
        self.assertEqual(self.vdata.vec_quad, 1)

    def test_calc_ocb_vec_sign(self):
        """ Test the calculation of the OCB vector signs
        """

        # Set the initial values
        self.vdata.ocb_aacgm_mlt = self.ocb.phi_cent[self.vdata.ocb_ind] / 15.0
        self.vdata.ocb_aacgm_lat = 90.0 - self.ocb.r_cent[self.vdata.ocb_ind]
        (self.vdata.ocb_lat, self.vdata.ocb_mlt,
         self.vdata.r_corr) = self.ocb.normal_coord(self.vdata.aacgm_lat,
                                                    self.vdata.aacgm_mlt)
        self.vdata.calc_vec_pole_angle()
        self.vdata.define_quadrants()

        vmag = np.sqrt(self.vdata.aacgm_n**2 + self.vdata.aacgm_e**2)
        self.vdata.aacgm_naz = np.degrees(np.arccos(self.vdata.aacgm_n / vmag))

        # Calculate the vector data signs
        vsigns = self.vdata.calc_ocb_vec_sign(north=True, east=True)
        self.assertTrue(vsigns['north'])
        self.assertTrue(vsigns['east'])

        del vmag, vsigns

    def test_scale_vec(self):
        """ Test the calculation of the OCB vector signs
        """

        # Set the initial values
        self.vdata.ocb_aacgm_mlt = self.ocb.phi_cent[self.vdata.ocb_ind] / 15.0
        self.vdata.ocb_aacgm_lat = 90.0 - self.ocb.r_cent[self.vdata.ocb_ind]
        (self.vdata.ocb_lat, self.vdata.ocb_mlt,
         self.vdata.r_corr) = self.ocb.normal_coord(self.vdata.aacgm_lat,
                                                    self.vdata.aacgm_mlt)
        self.vdata.calc_vec_pole_angle()
        self.vdata.define_quadrants()

        vmag = np.sqrt(self.vdata.aacgm_n**2 + self.vdata.aacgm_e**2)
        self.vdata.aacgm_naz = np.degrees(np.arccos(self.vdata.aacgm_n / vmag))

        # Scale the data vector
        self.vdata.scale_vector()

        # Test the North and East components
        self.assertAlmostEqual(self.vdata.ocb_n, 62.4751208491)
        self.assertAlmostEqual(self.vdata.ocb_e, 77.9686428950)

        # Test to see that the magnitudes and z-components are the same
        self.assertAlmostEqual(self.vdata.aacgm_mag, self.vdata.ocb_mag)
        self.assertAlmostEqual(self.vdata.ocb_z, self.vdata.aacgm_z)

        del vmag

    def test_scale_vec_z_zero(self):
        """ Test the calculation of the OCB vector sign with no vertical aacgm_z
        """
        # Re-assing the necessary variable
        self.vdata.aacgm_z = 0.0

        # Run the scale_vector routine
        self.vdata.set_ocb(self.ocb, scale_func=ocbpy.ocb_scaling.normal_evar)

        # Assess the ocb_z component
        self.assertEqual(self.vdata.ocb_z,
                         self.vdata.scale_func(0.0, self.vdata.unscaled_r,
                                               self.vdata.scaled_r))

    @unittest.skipIf(version_info.major == 2,
                     'Python 2.7 does not support subTest')
    def test_scale_vec_pole_angle_zero(self):
        """ Test the calculation of the OCB vector sign with no pole angle
        """
        self.vdata.set_ocb(self.ocb)
        self.vdata.pole_angle = 0.0

        nscale = ocbpy.ocb_scaling.normal_evar(self.vdata.aacgm_n,
                                               self.vdata.unscaled_r,
                                               self.vdata.scaled_r)
        escale = ocbpy.ocb_scaling.normal_evar(self.vdata.aacgm_e,
                                               self.vdata.unscaled_r,
                                               self.vdata.scaled_r)

        # Cycle through all the possible options for a pole angle of zero/180
        for tset in [('scale_func', None, self.vdata.aacgm_n,
                      self.vdata.aacgm_e),
                     ('scale_func', ocbpy.ocb_scaling.normal_evar, nscale,
                      escale),
                     ('ocb_aacgm_lat', self.vdata.aacgm_lat, -1.0 * nscale,
                      -1.0 * escale)]:
            with self.subTest(tset=tset):
                setattr(self.vdata, tset[0], tset[1])

                # Run the scale_vector routine with the new attributes
                self.vdata.scale_vector()

                # Assess the ocb north and east components
                self.assertEqual(self.vdata.ocb_n, tset[2])
                self.assertEqual(self.vdata.ocb_e, tset[3])

        del nscale, escale, tset

    @unittest.skipIf(version_info.major > 2, 'Already tested')
    def test_scale_vec_pole_angle_zero_none(self):
        """ Test the OCB vector sign routine with no pole angle or scaling
        """
        self.vdata.set_ocb(self.ocb)
        self.vdata.pole_angle = 0.0

        # Run the scale_vector routine with the new attributes
        self.vdata.scale_vector()

        # Assess the ocb north and east components
        self.assertEqual(self.vdata.ocb_n, self.vdata.aacgm_n)
        self.assertEqual(self.vdata.ocb_e, self.vdata.aacgm_e)

    @unittest.skipIf(version_info.major > 2, 'Already tested')
    def test_scale_vec_pole_angle_zero_scale_at_pole(self):
        """ Test the OCB vector sign routine with no pole angle and data at pole
        """
        self.vdata.set_ocb(self.ocb, scale_func=ocbpy.ocb_scaling.normal_evar)
        self.vdata.pole_angle = 0.0
        self.vdata.ocb_aacgm_lat = self.vdata.aacgm_lat

        nscale = -1.0 * ocbpy.ocb_scaling.normal_evar(self.vdata.aacgm_n,
                                                      self.vdata.unscaled_r,
                                                      self.vdata.scaled_r)
        escale = -1.0 * ocbpy.ocb_scaling.normal_evar(self.vdata.aacgm_e,
                                                      self.vdata.unscaled_r,
                                                      self.vdata.scaled_r)

        # Run the scale_vector routine with the new attributes
        self.vdata.scale_vector()

        # Assess the ocb north and east components
        self.assertEqual(self.vdata.ocb_n, nscale)
        self.assertEqual(self.vdata.ocb_e, escale)

        del nscale, escale

    def test_set_ocb_zero(self):
        """ Test setting of OCB values in VectorData without any magnitude
        """
        # Set the OCB values without any E-field scaling, test to see that the
        # AACGM and OCB vector magnitudes are the same
        self.zdata.set_ocb(self.ocb)
        self.assertEqual(self.zdata.ocb_mag, 0.0)

    def test_set_ocb_none(self):
        """ Test setting of OCB values without scaling
        """

        # Set the OCB values without any E-field scaling, test to see that the
        # AACGM and OCB vector magnitudes are the same
        self.vdata.set_ocb(self.ocb)
        self.assertAlmostEqual(self.vdata.aacgm_mag, self.vdata.ocb_mag)

    def test_set_ocb_evar(self):
        """ Test setting of OCB values with E field scaling
        """

        # Set the OCB values with scaling for a variable proportional to
        # the electric field
        self.vdata.set_ocb(self.ocb, scale_func=ocbpy.ocb_scaling.normal_evar)
        self.assertAlmostEqual(self.vdata.ocb_mag, 88.094416872365)

    def test_set_ocb_curl_evar(self):
        """ Test setting of OCB values with Curl E scaling
        """
        # Set the OCB values with scaling for a variable proportional to
        # the curl of the electric field
        self.vdata.set_ocb(self.ocb,
                           scale_func=ocbpy.ocb_scaling.normal_curl_evar)
        self.assertAlmostEqual(self.vdata.ocb_mag, 77.57814585822645)

    def test_scaled_r(self):
        """ Test that the scaled radius is correct
        """
        self.vdata.set_ocb(self.ocb, None)
        self.assertEqual(self.vdata.scaled_r, 16.0)

    def test_unscaled_r(self):
        """ Test that the unscaled radius is correct
        """
        self.vdata.set_ocb(self.ocb, None)
        self.assertEqual(self.vdata.unscaled_r, 14.09)

class TestVectorDataRaises(unittest.TestCase):
    def setUp(self):
        """ Initialize the tests for calc_vec_pole_angle
        """
        test_file = path.join(path.dirname(ocbpy.__file__), "tests",
                              "test_data", "test_north_circle")
        self.assertTrue(path.isfile(test_file))
        self.ocb = ocbpy.ocboundary.OCBoundary(filename=test_file,
                                               instrument='image')
        self.ocb.rec_ind = 27
        self.vdata = ocbpy.ocb_scaling.VectorData(0, self.ocb.rec_ind, 75.0,
                                                  22.0, aacgm_n=50.0,
                                                  aacgm_e=86.5, aacgm_z=5.0,
                                                  dat_name="Test",
                                                  dat_units="$m s^{-1}$")
        self.input_attrs = list()
        self.bad_input = [np.nan, np.full(shape=2, fill_value=np.nan)]
        self.raise_out = list()
        self.hold_val = None

        if version_info.major == 2:
            self.assertRaisesRegex = self.assertRaisesRegexp

    def tearDown(self):
        del self.ocb, self.vdata, self.input_attrs, self.bad_input
        del self.raise_out, self.hold_val

    def test_init_ocb_array_failure(self):
        """ Test init failure with mismatched OCB and input array input
        """
        self.input_attrs = [0, [27, 31], 75.0, 22.0]
        self.bad_input = {'aacgm_n': 100.0, 'aacgm_e': 100.0,
                          'aacgm_z': 10.0, 'ocb_lat': 81.0,
                          'ocb_mlt': [2.0, 5.8, 22.5]}

        with self.assertRaisesRegex(ValueError, "OCB index and input shapes"):
            self.vdata = ocbpy.ocb_scaling.VectorData(*self.input_attrs,
                                                      **self.bad_input)

    def test_init_ocb_vector_failure(self):
        """ Test init failure with mismatched OCB and data array input
        """
        self.input_attrs = [[3, 6, 0], [27, 31], [75.0, 87.2, 65.0],
                            [22.0, 21, 22]]
        self.bad_input = {'aacgm_n': [100.0, 110.0, 30.0],
                          'aacgm_e': [100.0, 110.0, 30.0],
                          'aacgm_z': [10.0, 10.0, 3.0]}

        with self.assertRaisesRegex(ValueError,
                                    "Mismatched OCB and Vector input shapes"):
            self.vdata = ocbpy.ocb_scaling.VectorData(*self.input_attrs,
                                                      **self.bad_input)

    @unittest.skipIf(version_info.major == 2,
                     'Python 2.7 does not support subTest')
    def test_init_vector_failure(self):
        """ Test init failure with a bad mix of vector and scalar input
        """
        self.input_attrs = [[0, self.ocb.rec_ind, [75.0, 70.0], [22.0, 20.0]],
                            [[0, 1], self.ocb.rec_ind, [75.0, 70.0], 22.0],
                            [[0, 1], self.ocb.rec_ind, [75.0, 70.0],
                             [22.0, 20.0, 23.0]]]
        self.bad_input = [{'aacgm_n': 10.0},
                          {'aacgm_n': [100.0, 110.0, 30.0]},
                          {'aacgm_n': [100.0, 110.0, 30.0]}]
        self.raise_out = ['data index shape must match vector shape',
                          'mismatched VectorData input shapes',
                          'mismatched VectorData input shapes']

        for i, iattrs in enumerate(self.input_attrs):
            tset = [iattrs, self.bad_input[i], self.raise_out[i]]
            with self.subTest(tset=tset):
                with self.assertRaisesRegex(ValueError, tset[2]):
                    self.vdata = ocbpy.ocb_scaling.VectorData(*tset[0],
                                                              **tset[1])

    @unittest.skipIf(version_info.major > 2, 'Already tested with subTest')
    def test_init_vector_failure_dat_ind(self):
        """ Test init failure with a bad data index shape
        """
        self.input_attrs = [0, self.ocb.rec_ind, [75.0, 70.0], [22.0, 20.0]]
        self.bad_input = {'aacgm_n': 10.0}

        with self.assertRaisesRegex(
                ValueError, "data index shape must match vector shape"):
            self.vdata = ocbpy.ocb_scaling.VectorData(*self.input_attrs,
                                                      **self.bad_input)

    @unittest.skipIf(version_info.major > 2, 'Already tested with subTest')
    def test_init_vector_failure_many_array_size(self):
        """ Test init failure with a bad vector lengths
        """
        self.input_attrs = [[0, 1], self.ocb.rec_ind, [75.0, 70.0], 20.0]
        self.bad_input = {'aacgm_n': [100.0, 110.0, 30.0]}

        with self.assertRaisesRegex(ValueError,
                                    "mismatched VectorData input shapes"):
            self.vdata = ocbpy.ocb_scaling.VectorData(*self.input_attrs,
                                                      **self.bad_input)

    @unittest.skipIf(version_info.major > 2, 'Already tested with subTest')
    def test_init_vector_failure_bad_lat_array_size(self):
        """ Test init failure with a bad vector lengths
        """
        self.input_attrs = [[0, 1, 3], self.ocb.rec_ind, [75.0, 70.0],
                            [22.0, 20.0, 23.0]]
        self.bad_input = {'aacgm_n': [100.0, 110.0, 30.0]}

        with self.assertRaisesRegex(ValueError,
                                    "mismatched VectorData input shapes"):
            self.vdata = ocbpy.ocb_scaling.VectorData(*self.input_attrs,
                                                      **self.bad_input)

    @unittest.skipIf(version_info.major == 2,
                     'Python 2.7 does not support subTest')
    def test_bad_calc_vec_pole_angle(self):
        """Test calc_vec_pole_angle failure with bad input"""
        self.input_attrs = ['aacgm_mlt', 'ocb_aacgm_mlt', 'aacgm_lat',
                            'ocb_aacgm_lat']
        self.raise_out = ["AACGM MLT of Vector", "AACGM MLT of OCB pole",
                          "AACGM latitude of Vector",
                          "AACGM latitude of OCB pole"]
        tsets = [(iattrs, bi, self.raise_out[i])
                 for i, iattrs in enumerate(self.input_attrs)
                 for bi in self.bad_input]
        self.vdata.set_ocb(self.ocb, None)

        for tset in tsets:
            with self.subTest(tset=tset):
                self.hold_val = getattr(self.vdata, tset[0])
                setattr(self.vdata, tset[0], tset[1])

                with self.assertRaisesRegex(ValueError, tset[2]):
                    self.vdata.calc_vec_pole_angle()

                setattr(self.vdata, tset[0], self.hold_val)

    @unittest.skipIf(version_info.major > 2, 'Already tested with subTest')
    def test_bad_calc_vec_pole_angle_mlt_float(self):
        """Test calc_vec_pole_angle failure with bad AACGM MLT
        """
        self.vdata.set_ocb(self.ocb, None)
        self.vdata.aacgm_mlt = self.bad_input[0]

        with self.assertRaisesRegex(ValueError, "AACGM MLT of Vector"):
            self.vdata.calc_vec_pole_angle()

    @unittest.skipIf(version_info.major > 2, 'Already tested with subTest')
    def test_bad_calc_vec_pole_angle_ocb_mlt_float_small(self):
        """Test calc_vec_pole_angle failure with small bad OCB MLT
        """
        self.vdata.set_ocb(self.ocb, None)
        self.vdata.ocb_aacgm_mlt = self.bad_input[0]

        with self.assertRaisesRegex(ValueError, "AACGM MLT of OCB pole"):
            self.vdata.calc_vec_pole_angle()

    @unittest.skipIf(version_info.major > 2, 'Already tested with subTest')
    def test_bad_calc_vec_pole_angle_vec_mlat_float(self):
        """Test calc_vec_pole_angle failure with bad vector latitude
        """
        self.vdata.set_ocb(self.ocb, None)
        self.vdata.aacgm_lat = self.bad_input[0]

        with self.assertRaisesRegex(ValueError, "AACGM latitude of Vector"):
            self.vdata.calc_vec_pole_angle()

    @unittest.skipIf(version_info.major > 2, 'Already tested with subTest')
    def test_bad_calc_vec_pole_angle_mlt_array(self):
        """Test calc_vec_pole_angle failure with bad AACGM MLT
        """
        self.vdata.set_ocb(self.ocb, None)
        self.vdata.aacgm_mlt = self.bad_input[1]

        with self.assertRaisesRegex(ValueError, "AACGM MLT of Vector"):
            self.vdata.calc_vec_pole_angle()

    @unittest.skipIf(version_info.major > 2, 'Already tested with subTest')
    def test_bad_calc_vec_pole_angle_ocb_mlt_float_big(self):
        """Test calc_vec_pole_angle failure with bad OCB MLT
        """
        self.vdata.set_ocb(self.ocb, None)
        self.vdata.ocb_aacgm_mlt = self.bad_input[1]

        with self.assertRaisesRegex(ValueError, "AACGM MLT of OCB pole"):
            self.vdata.calc_vec_pole_angle()

    @unittest.skipIf(version_info.major > 2, 'Already tested with subTest')
    def test_bad_calc_vec_pole_angle_ocb_mlat_float_small(self):
        """Test calc_vec_pole_angle failure with barely bad OCB latitude
        """
        self.vdata.set_ocb(self.ocb, None)
        self.vdata.ocb_aacgm_lat = self.bad_input[0]

        with self.assertRaisesRegex(ValueError, "AACGM latitude of OCB pole"):
            self.vdata.calc_vec_pole_angle()

    @unittest.skipIf(version_info.major > 2, 'Already tested with subTest')
    def test_bad_calc_vec_pole_angle_ocb_mlat_float_big(self):
        """Test calc_vec_pole_angle failure with bad OCB latitude
        """
        self.vdata.set_ocb(self.ocb, None)
        self.vdata.ocb_aacgm_lat = self.bad_input[1]

        with self.assertRaisesRegex(ValueError, "AACGM latitude of OCB pole"):
            self.vdata.calc_vec_pole_angle()

    @unittest.skipIf(version_info.major > 2, 'Already tested with subTest')
    def test_bad_calc_vec_aacgm_lat_float(self):
        """Test calc_vec_pole_angle failure with bad vector latitude
        """
        self.vdata.set_ocb(self.ocb, None)
        self.vdata.aacgm_lat = self.bad_input[1]

        with self.assertRaisesRegex(ValueError, "AACGM latitude of Vector"):
            self.vdata.calc_vec_pole_angle()

    def test_no_ocb_lat(self):
        """ Test failure when OCB latitude is not available
        """
        self.vdata.ocb_lat = np.nan

        with self.assertRaisesRegex(ValueError, 'OCB coordinates required'):
            self.vdata.scale_vector()

    def test_no_ocb_mlt(self):
        """ Test failure when OCB MLT is not available
        """
        self.vdata.ocb_mlt = np.nan

        with self.assertRaisesRegex(ValueError, 'OCB coordinates required'):
            self.vdata.scale_vector()

    def test_no_ocb_pole_location(self):
        """ Test failure when OCB pole location is not available
        """
        self.vdata.set_ocb(self.ocb, None)
        self.vdata.ocb_aacgm_mlt = np.nan

        with self.assertRaisesRegex(ValueError, "OCB pole location required"):
            self.vdata.scale_vector()

    def test_no_ocb_pole_angle(self):
        """ Test failure when pole angle is not available
        """
        self.vdata.set_ocb(self.ocb, None)
        self.vdata.pole_angle = np.nan

        with self.assertRaisesRegex(
                ValueError, "vector angle in poles-vector triangle required"):
            self.vdata.scale_vector()

    def test_bad_ocb_quad(self):
        """ Test failure when OCB quadrant is wrong
        """
        self.vdata.set_ocb(self.ocb, None)
        self.vdata.ocb_quad = -1

        with self.assertRaisesRegex(ValueError, "OCB quadrant undefined"):
            self.vdata.calc_ocb_polar_angle()

    def test_bad_vec_quad(self):
        """ Test failure when vector quadrant is wrong
        """
        self.vdata.set_ocb(self.ocb, None)
        self.vdata.vec_quad = -1

        with self.assertRaisesRegex(ValueError, "Vector quadrant undefined"):
            self.vdata.calc_ocb_polar_angle()

    def test_bad_quad_polar_angle(self):
        """ Test failure when quadrant polar angle is bad
        """
        self.vdata.set_ocb(self.ocb, None)
        self.vdata.aacgm_naz = np.nan

        with self.assertRaisesRegex(ValueError,
                                    "AACGM polar angle undefined"):
            self.vdata.calc_ocb_polar_angle()

    def test_bad_quad_vector_angle(self):
        """ Test failure when quandrant vector angle is bad
        """
        self.vdata.set_ocb(self.ocb, None)
        self.vdata.pole_angle = np.nan

        with self.assertRaisesRegex(ValueError, "Vector angle undefined"):
            self.vdata.calc_ocb_polar_angle()

    def test_bad_calc_vec_sign_direction(self):
        """ Test calc_vec_sign failure when no direction is provided
        """
        self.vdata.set_ocb(self.ocb, None)

        with self.assertRaisesRegex(ValueError,
                                    "must set at least one direction"):
            self.vdata.calc_ocb_vec_sign()

    def test_bad_calc_sign_ocb_quad(self):
        """ Test calc_vec_sign failure with bad OCB quadrant
        """
        self.vdata.set_ocb(self.ocb, None)
        self.vdata.ocb_quad = -1

        with self.assertRaisesRegex(ValueError, "OCB quadrant undefined"):
            self.vdata.calc_ocb_vec_sign(north=True)

    def test_bad_calc_sign_vec_quad(self):
        """ Test calc_vec_sign failure with bad vector quadrant
        """
        self.vdata.set_ocb(self.ocb, None)
        self.vdata.vec_quad = -1

        with self.assertRaisesRegex(ValueError, "Vector quadrant undefined"):
            self.vdata.calc_ocb_vec_sign(north=True)

    def test_bad_calc_sign_polar_angle(self):
        """ Test calc_vec_sign failure with bad polar angle
        """
        self.vdata.set_ocb(self.ocb, None)
        self.vdata.aacgm_naz = np.nan

        with self.assertRaisesRegex(ValueError,
                                    "AACGM polar angle undefined"):
            self.vdata.calc_ocb_vec_sign(north=True)

    def test_bad_calc_sign_pole_angle(self):
        """ Test calc_vec_sign failure with bad pole angle
        """
        self.vdata.set_ocb(self.ocb, None)
        self.vdata.pole_angle = np.nan

        with self.assertRaisesRegex(ValueError, "Vector angle undefined"):
            self.vdata.calc_ocb_vec_sign(north=True)

    def test_bad_define_quandrants_pole_mlt(self):
        """Test define_quadrants failure with bad pole MLT
        """
        self.vdata.set_ocb(self.ocb, None)
        self.vdata.ocb_aacgm_mlt = np.nan

        with self.assertRaisesRegex(ValueError, "OCB pole location required"):
            self.vdata.define_quadrants()

    def test_bad_define_quandrants_vec_mlt(self):
        """Test define_quadrants failure with bad vector MLT
        """
        self.vdata.set_ocb(self.ocb, None)
        self.vdata.aacgm_mlt = np.nan

        with self.assertRaisesRegex(ValueError,
                                    "Vector AACGM location required"):
            self.vdata.define_quadrants()

    def test_bad_define_quandrants_pole_angle(self):
        """Test define_quadrants failure with bad pole angle
        """
        self.vdata.set_ocb(self.ocb, None)
        self.vdata.pole_angle = np.nan

        with self.assertRaisesRegex(
                ValueError, "vector angle in poles-vector triangle required"):
            self.vdata.define_quadrants()

class TestHaversine(unittest.TestCase):

    def setUp(self):
        """ Initialize the tests for the haversine and archaversine functions
        """
        self.input_angles = np.linspace(-2.0*np.pi, 2.0*np.pi, 9)
        self.hav_out = np.array([0.0, 0.5, 1.0, 0.5, 0.0, 0.5, 1.0, 0.5, 0.0])
        # archaversine is confinded to 0-pi
        self.ahav_out = abs(np.array([aa - np.sign(aa) * 2.0 * np.pi
                                      if abs(aa) > np.pi
                                      else aa for aa in self.input_angles]))
        self.out = None

    def tearDown(self):
        del self.input_angles, self.hav_out, self.out, self.ahav_out

    @unittest.skipIf(version_info.major == 2,
                     'Python 2.7 does not support subTest')
    def test_haversine(self):
        """ Test implimentation of the haversine
        """
        # Cycle through all the possible input options
        for i, tset in enumerate([(self.input_angles[0], self.hav_out[0]),
                                  (list(self.input_angles), self.hav_out),
                                  (self.input_angles, self.hav_out)]):
            with self.subTest(tset=tset):
                self.out = ocbpy.ocb_scaling.hav(tset[0])

                # Assess the output
                if i == 0:
                    self.assertAlmostEqual(self.out, tset[1])
                else:
                    self.assertTrue(np.all(abs(self.out - tset[1]) < 1.0e-7))

        del tset

    @unittest.skipIf(version_info.major > 2, 'Already tested with subTest')
    def test_haversine_float(self):
        """ Test implimentation of the haversine with float inputs
        """

        for i, alpha in enumerate(self.input_angles):
            self.assertAlmostEqual(ocbpy.ocb_scaling.hav(alpha),
                                   self.hav_out[i])

    @unittest.skipIf(version_info.major > 2, 'Already tested with subTest')
    def test_haversine_list(self):
        """ Test implimentation of the haversine with a list input
        """
        self.out = ocbpy.ocb_scaling.hav(list(self.input_angles))
        self.assertTrue(np.all(abs(self.out - self.hav_out) < 1.0e-7))

    @unittest.skipIf(version_info.major > 2, 'Already tested with subTest')
    def test_haversine_array(self):
        """ Test implimentation of the haversine with a array input
        """
        self.out = ocbpy.ocb_scaling.hav(self.input_angles)
        self.assertTrue(np.all(abs(self.out - self.hav_out) < 1.0e-7))

    @unittest.skipIf(version_info.major == 2,
                     'Python 2.7 does not support subTest')
    def test_inverse_haversine(self):
        """ Test the implemenation of the inverse haversine
        """
        # Cycle through all the possible input options
        for i, tset in enumerate([(self.hav_out[0], self.ahav_out[0]),
                                  (list(self.hav_out), self.ahav_out),
                                  (self.hav_out, self.ahav_out)]):
            with self.subTest(tset=tset):
                self.out = ocbpy.ocb_scaling.archav(tset[0])

                # Assess the output
                if i == 0:
                    self.assertEqual(self.out, abs(tset[1]))
                else:
                    self.assertTrue(np.all(self.out - tset[1] < 1.0e-7))

        del tset

    @unittest.skipIf(version_info.major > 2,
                     'Python 2.7 does not support subTest')
    def test_inverse_haversine_float(self):
        """ Test implimentation of the inverse haversine with float input
        """
        for i, self.out in enumerate(self.ahav_out):
            self.assertAlmostEqual(ocbpy.ocb_scaling.archav(self.hav_out[i]),
                                   self.out)

    @unittest.skipIf(version_info.major > 2,
                     'Python 2.7 does not support subTest')
    def test_inverse_haversine_list(self):
        """ Test implimentation of the inverse haversine with list input
        """
        self.out = ocbpy.ocb_scaling.archav(list(self.hav_out))
        self.assertTrue(np.all(abs(self.out - self.ahav_out)
                               < 1.0e-7))

    @unittest.skipIf(version_info.major > 2,
                     'Python 2.7 does not support subTest')
    def test_inverse_haversine_array(self):
        """ Test implimentation of the inverse haversine with array input
        """
        self.out = ocbpy.ocb_scaling.archav(self.hav_out)
        self.assertTrue(np.all(abs(self.out - self.ahav_out) < 1.0e-7))

    def test_inverse_haversine_small_float(self):
        """ Test implimentation of the inverse haversine with very small numbers
        """
        self.assertEqual(ocbpy.ocb_scaling.archav(1.0e-17), 0.0)
        self.assertEqual(ocbpy.ocb_scaling.archav(-1.0e-17), 0.0)

    def test_inverse_haversine_nan_float(self):
        """ Test implimentation of the inverse haversine with NaN
        """
        self.assertTrue(np.isnan(ocbpy.ocb_scaling.archav(np.nan)))

    def test_inverse_haversine_negative_float(self):
        """ Test implimentation of the inverse haversine with negative input
        """
        self.assertTrue(np.isnan(ocbpy.ocb_scaling.archav(-1.0)))

    def test_inverse_haversine_mixed(self):
        """ Test the inverse haversine with array input of good and bad values
        """
        # Update the test input and output
        self.hav_out[0] = 1.0e-17
        self.ahav_out[0] = 0.0
        self.hav_out[1] = np.nan
        self.ahav_out[1] = np.nan
        self.hav_out[2] = -1.0
        self.ahav_out[2] = np.nan

        self.out = ocbpy.ocb_scaling.archav(self.hav_out)

        for i, hout in enumerate(self.out):
            if np.isnan(hout):
                self.assertTrue(np.isnan(self.ahav_out[i]))
            else:
                self.assertAlmostEqual(hout, self.ahav_out[i])

class TestOCBScalingArrayMethods(unittest.TestCase):

    def setUp(self):
        """ Initialize the OCBoundary and array VectorData objects
        """
        test_file = path.join(path.dirname(ocbpy.__file__), "tests",
                              "test_data", "test_north_circle")
        self.ocb = ocbpy.ocboundary.OCBoundary(filename=test_file,
                                               instrument='image')
        # Construct a set of test vectors that have all the different OCB
        # and vector combinations and one with no magnitude. The test OCB pole
        # is at 87.24 deg, 5.832 h
        lats = np.full(shape=(17,), fill_value=75.0)
        lats[8:] = 89.0
        mlts = np.zeros(shape=(17,))
        mlts[4:8] = 15.0
        mlts[8:12] = 7.0
        mlts[12:] = 5.0
        north = [10.0, 10.0, -10.0, -10.0, 10.0, 10.0, -10.0, -10.0,
                 10.0, 10.0, -10.0, -10.0, 10.0, 10.0, -10.0, -10.0, 0.0]
        east = [3.0, -3.0, -3.0, 3.0, 3.0, -3.0, -3.0, 3.0, 3.0, -3.0, -3.0,
                3.0, 3.0, -3.0, -3.0, 3.0, 0.0]
        vert = np.zeros(shape=(17,))
        vert[0] = 5.0
        self.ref_quads = {'ocb': [1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3, 4, 4, 4,
                                  4, 4],
                          'vec': [1, 2, 3, 4, 1, 2, 3, 4, 1, 2, 3, 4, 1, 2, 3,
                                  4, 1]}

        self.vargs = [np.arange(0, 17, 1), 27, lats, mlts]
        self.vkwargs = {'aacgm_n': np.array(north), 'aacgm_e': np.array(east),
                        'aacgm_z': np.array(vert), 'dat_name': 'Test',
                        'dat_units': 'm/s'}
        self.vdata = None
        self.out = None

        self.aacgm_mag = np.full(shape=(17,), fill_value=10.44030650891055)
        self.aacgm_mag[0] = 11.575836902790225
        self.aacgm_mag[-1] = 0.0

        if version_info.major == 2:
            self.assertRegex = self.assertRegexpMatches
            self.assertNotRegex = self.assertNotRegexpMatches

    def tearDown(self):
        del self.ocb, self.vargs, self.vkwargs, self.out, self.vdata
        del self.aacgm_mag, self.ref_quads

    def set_vector_ocb_ind(self):
        """ Update the input vector arguements to have vector OCB index input
        """
        self.vargs[1] = np.full(shape=self.vargs[-1].shape, fill_value=27)
        self.vargs[1][8:] = 31

    def test_array_vector_repr_not_calc(self):
        """ Test the VectorData print statement with uncalculated array input
        """
        self.vdata = ocbpy.ocb_scaling.VectorData(*self.vargs, **self.vkwargs)
        self.out = self.vdata.__repr__()
        self.assertRegex(self.out, "Index")
        self.assertRegex(self.out, "nan, nan, {:d}".format(self.vargs[1]))

    def test_array_vector_repr_calc(self):
        """ Test the VectorData print statement with calculated array input
        """
        self.vdata = ocbpy.ocb_scaling.VectorData(*self.vargs, **self.vkwargs)
        self.vdata.set_ocb(self.ocb)
        self.out = self.vdata.__repr__()
        self.assertRegex(self.out, "Index")
        self.assertNotRegex(self.out, "nan, nan")

    def test_array_vector_repr_calc_ocb_vec_array(self):
        """ Test the VectorData print statement with calculated ocb/vec arrays
        """
        self.set_vector_ocb_ind()
        self.vdata = ocbpy.ocb_scaling.VectorData(*self.vargs, **self.vkwargs)
        self.vdata.set_ocb(self.ocb)
        self.out = self.vdata.__repr__()
        self.assertRegex(self.out, "Index")
        self.assertNotRegex(self.out, "nan, nan")

    def test_array_vector_repr_calc_ocb_array(self):
        """ Test the VectorData print statement with calculated ocb arrays
        """
        self.vargs[0] = self.vargs[0][0]
        self.set_vector_ocb_ind()
        self.vargs[2] = self.vargs[2][0]
        self.vargs[3] = self.vargs[3][0]
        self.vkwargs['aacgm_n'] = self.vkwargs['aacgm_n'][0]
        self.vkwargs['aacgm_e'] = self.vkwargs['aacgm_e'][0]
        self.vkwargs['aacgm_z'] = self.vkwargs['aacgm_z'][0]
        self.vdata = ocbpy.ocb_scaling.VectorData(*self.vargs, **self.vkwargs)
        self.vdata.set_ocb(self.ocb)
        self.out = self.vdata.__repr__()
        self.assertRegex(self.out, "Index")
        self.assertNotRegex(self.out, "nan, nan")

    def test_init_nez_vec_array(self):
        """ Test VectorData initialisation  with vector array components
        """
        self.vdata = ocbpy.ocb_scaling.VectorData(*self.vargs, **self.vkwargs)
        self.assertEqual(len(self.vdata.aacgm_mag), len(self.vargs[0]))
        self.assertEqual(len(self.vdata.ocb_mag), len(self.vargs[0]))
        for i, self.out in enumerate(self.vdata.aacgm_mag):
            self.assertAlmostEqual(self.out, self.aacgm_mag[i])

    def test_init_nez_ocb_vec_array(self):
        """ Test VectorData initialisation with ocb and vector array components
        """
        self.set_vector_ocb_ind()
        self.vdata = ocbpy.ocb_scaling.VectorData(*self.vargs, **self.vkwargs)
        self.assertEqual(len(self.vdata.aacgm_mag), len(self.vargs[0]))
        self.assertEqual(len(self.vdata.ocb_mag), len(self.vargs[0]))
        for i, self.out in enumerate(self.vdata.aacgm_mag):
            self.assertAlmostEqual(self.out, self.aacgm_mag[i])

    def test_init_nez_ocb_array(self):
        """ Test VectorData initialisation with ocb array components
        """
        self.set_vector_ocb_ind()
        self.vargs[2] = self.vargs[2][0]
        self.vargs[3] = self.vargs[3][0]
        self.vkwargs['aacgm_n'] = self.vkwargs['aacgm_n'][0]
        self.vkwargs['aacgm_e'] = self.vkwargs['aacgm_e'][0]
        self.vkwargs['aacgm_z'] = self.vkwargs['aacgm_z'][0]

        self.vdata = ocbpy.ocb_scaling.VectorData(*self.vargs, **self.vkwargs)
        self.assertEqual(len(self.vdata.ocb_mag), len(self.vargs[1]))
        self.assertAlmostEqual(self.vdata.aacgm_mag, self.aacgm_mag[0])

    def test_init_mag(self):
        """ Test the initialisation of the VectorData array input with magnitude
        """
        self.vkwargs['aacgm_mag'] = self.aacgm_mag
        self.vdata = ocbpy.ocb_scaling.VectorData(*self.vargs, **self.vkwargs)
        self.assertEqual(len(self.vdata.aacgm_mag), len(self.vargs[0]))
        for i, self.out in enumerate(self.vdata.aacgm_mag):
            self.assertAlmostEqual(self.out, self.aacgm_mag[i])

    def test_vector_all_bad_lat(self):
        """ Test the VectorData output with all data from the wrong hemisphere
        """
        self.vargs[2] *= -1.0
        self.vdata = ocbpy.ocb_scaling.VectorData(*self.vargs, **self.vkwargs)
        self.vdata.set_ocb(self.ocb)

        self.assertTrue(len(self.vdata.ocb_lat), len(self.vargs[2]))
        self.assertTrue(np.all(np.isnan(self.vdata.ocb_lat)))
        self.assertTrue(np.all(np.isnan(self.vdata.ocb_mlt)))
        self.assertTrue(np.all(np.isnan(self.vdata.ocb_n)))
        self.assertTrue(np.all(np.isnan(self.vdata.ocb_e)))
        self.assertTrue(np.all(np.isnan(self.vdata.ocb_z)))

        # Ensure that input is not overwritten
        for vkey in self.vkwargs.keys():
            self.out = getattr(self.vdata, vkey)
            if vkey.find('aacgm_') == 0:
                for i, val in enumerate(self.vkwargs[vkey]):
                    if np.isnan(val):
                        self.assertTrue(np.isnan(self.out[i]))
                    else:
                        self.assertEqual(self.out[i], val)
            else:
                self.assertRegex(self.out, self.vkwargs[vkey])

    def test_vector_some_bad_lat(self):
        """ Test the VectorData output with mixed hemisphere input
        """
        self.vargs[2][0] *= -1.0
        self.vdata = ocbpy.ocb_scaling.VectorData(*self.vargs, **self.vkwargs)
        self.vdata.set_ocb(self.ocb)

        self.assertTrue(len(self.vdata.ocb_lat), len(self.vargs[2]))

        # Ensure the wrong hemisphere is NaN
        self.assertTrue(np.isnan(self.vdata.ocb_lat[0]))
        self.assertTrue(np.isnan(self.vdata.ocb_mlt[0]))
        self.assertTrue(np.isnan(self.vdata.ocb_n[0]))
        self.assertTrue(np.isnan(self.vdata.ocb_e[0]))
        self.assertTrue(np.isnan(self.vdata.ocb_z[0]))

        # Ensure that input is not overwritten
        for vkey in self.vkwargs.keys():
            self.out = getattr(self.vdata, vkey)
            if vkey.find('aacgm_') == 0:
                for i, val in enumerate(self.vkwargs[vkey]):
                    if np.isnan(val):
                        self.assertTrue(np.isnan(self.out[i]))
                    else:
                        self.assertEqual(self.out[i], val)
            else:
                self.assertRegex(self.out, self.vkwargs[vkey])

        # Ensure the right hemisphere is good
        self.assertAlmostEqual(self.vdata.aacgm_mag[1], self.aacgm_mag[1])
        self.assertAlmostEqual(self.vdata.ocb_mag[1], self.aacgm_mag[1])

    def test_calc_vec_pole_angle_flat(self):
        """ Test the polar angle calculation with angles of 0 and 180 deg
        """
        self.vargs[3] = np.full(shape=self.vargs[2].shape,
                                fill_value=ocbpy.ocb_time.deg2hr(
                                    self.ocb.phi_cent[self.vargs[1]]))
        self.vargs[3][np.array(self.ref_quads['ocb']) > 2] += 12.0
        self.vdata = ocbpy.ocb_scaling.VectorData(*self.vargs, **self.vkwargs)
        self.vdata.set_ocb(self.ocb)

        self.assertTrue(np.all([quad in [2, 4]
                                for quad in self.vdata.ocb_quad]))
        self.assertEqual(list(self.vdata.vec_quad), self.ref_quads['vec'])
        self.assertTrue(np.all([self.vdata.pole_angle[i] == 0.0
                                for i, quad in enumerate(self.vdata.ocb_quad)
                                if quad == 2]))
        self.assertTrue(np.all([self.vdata.pole_angle[i] == 180.0
                                for i, quad in enumerate(self.vdata.ocb_quad)
                                if quad == 4]))

    def test_array_vec_quad(self):
        """ Test the assignment of vector quadrants with array input
        """
        self.vdata = ocbpy.ocb_scaling.VectorData(*self.vargs, **self.vkwargs)
        self.vdata.set_ocb(self.ocb)

        self.assertEqual(len(self.vargs[0]), len(self.vdata.vec_quad))
        self.assertEqual(list(self.vdata.vec_quad), self.ref_quads['vec'])

    def test_array_ocb_quad(self):
        """ Test the assignment of OCB quadrants with array input
        """
        self.vdata = ocbpy.ocb_scaling.VectorData(*self.vargs, **self.vkwargs)
        self.vdata.set_ocb(self.ocb)

        self.assertEqual(len(self.vargs[0]), len(self.vdata.ocb_quad))
        self.assertEqual(list(self.vdata.ocb_quad), self.ref_quads['ocb'])

    def test_one_undefinable_ocb_quadrant(self):
        """ Test VectorData array initialization for a undefinable OCB quadrant
        """
        self.vargs[2][1] = 0.0
        self.vdata = ocbpy.ocb_scaling.VectorData(*self.vargs, **self.vkwargs)
        self.vdata.set_ocb(self.ocb)

        self.assertEqual(self.vdata.ocb_quad[0], 1)
        self.assertEqual(self.vdata.ocb_quad[1], 0)

    def test_one_undefinable_vec_quadrant(self):
        """ Test VectorData array initialization for a undefinable vec quadrant
        """
        self.vkwargs['aacgm_n'][1] = np.nan
        self.vdata = ocbpy.ocb_scaling.VectorData(*self.vargs, **self.vkwargs)
        self.vdata.set_ocb(self.ocb)

        self.assertEqual(self.vdata.vec_quad[0], 1)
        self.assertEqual(self.vdata.vec_quad[1], 0)

    def test_define_quadrants_neg_adj_mlt(self):
        """ Test the quadrant assignment with a negative AACGM MLT
        """
        self.vargs[3] -= 24.0
        self.vdata = ocbpy.ocb_scaling.VectorData(*self.vargs, **self.vkwargs)
        self.vdata.set_ocb(self.ocb, scale_func=ocbpy.ocb_scaling.normal_evar)

        self.assertGreater(self.vdata.ocb_aacgm_mlt - self.vargs[3][0], 24)

        self.assertEqual(list(self.vdata.ocb_quad), self.ref_quads['ocb'])
        self.assertEqual(list(self.vdata.vec_quad), self.ref_quads['vec'])

    @unittest.skipIf(version_info.major == 2,
                     'Python 2.7 does not support subTest')
    def test_scale_vec_pole_angle_zero(self):
        """ Test the calculation of the OCB vector sign with no pole angle
        """
        self.vdata = ocbpy.ocb_scaling.VectorData(*self.vargs, **self.vkwargs)
        self.vdata.set_ocb(self.ocb)

        # If the OCB pole is closer to the AACGM pole than the vector, set
        # the pole angle to zero deg. Otherwise, set it to 180.0 deg
        self.vdata.pole_angle = np.zeros(shape=self.vargs[2].shape)
        self.vdata.pole_angle[self.vdata.ocb_quad > 2] = 180.0

        nscale = ocbpy.ocb_scaling.normal_evar(self.vdata.aacgm_n,
                                               self.vdata.unscaled_r,
                                               self.vdata.scaled_r)
        escale = ocbpy.ocb_scaling.normal_evar(self.vdata.aacgm_e,
                                               self.vdata.unscaled_r,
                                               self.vdata.scaled_r)

        # Cycle through all the possible options for a pole angle of zero/180
        for tset in [('scale_func', None, self.vkwargs['aacgm_n'],
                      self.vkwargs['aacgm_e']),
                     ('scale_func', ocbpy.ocb_scaling.normal_evar, nscale,
                      escale)]:
            with self.subTest(tset=tset):
                setattr(self.vdata, tset[0], tset[1])

                # Run the scale_vector routine with the new attributes
                self.vdata.scale_vector()

                # Assess the ocb north and east components
                self.assertTrue(np.all(self.vdata.ocb_n == tset[2]))
                self.assertTrue(np.all(self.vdata.ocb_e == tset[3]))

        del nscale, escale, tset

    @unittest.skipIf(version_info.major > 2, 'Already tested with subTest')
    def test_scale_vec_pole_angle_zero_noscale(self):
        """ Test the OCB vector sign calc with no pole angle or scaling
        """
        self.vdata = ocbpy.ocb_scaling.VectorData(*self.vargs, **self.vkwargs)
        self.vdata.set_ocb(self.ocb)
        self.vdata.pole_angle = np.zeros(shape=self.vargs[2].shape)
        self.vdata.pole_angle[self.vdata.ocb_quad > 2] = 180.0

        # Run the scale_vector routine with the new attributes
        self.vdata.scale_vector()

        # Assess the ocb north and east components
        self.assertTrue(np.all(self.vdata.ocb_n == self.vkwargs['aacgm_n']))
        self.assertTrue(np.all(self.vdata.ocb_e == self.vkwargs['aacgm_e']))

    @unittest.skipIf(version_info.major > 2, 'Already tested with subTest')
    def test_scale_vec_pole_angle_zero_scale(self):
        """  Test the OCB vector sign calc with scaling but no pole angle
        """
        self.vdata = ocbpy.ocb_scaling.VectorData(*self.vargs, **self.vkwargs)
        self.vdata.set_ocb(self.ocb, scale_func=ocbpy.ocb_scaling.normal_evar)
        self.vdata.pole_angle = np.zeros(shape=self.vargs[2].shape)
        self.vdata.pole_angle[self.vdata.ocb_quad > 2] = 180.0

        # Run the scale_vector routine with the new attributes
        self.vdata.scale_vector()

        # Assess the ocb north and east components
        self.out = ocbpy.ocb_scaling.normal_evar(self.vkwargs['aacgm_n'],
                                                 self.vdata.unscaled_r,
                                                 self.vdata.scaled_r)
        self.assertTrue(np.all(self.vdata.ocb_n == self.out))

        self.out = ocbpy.ocb_scaling.normal_evar(self.vkwargs['aacgm_e'],
                                                 self.vdata.unscaled_r,
                                                 self.vdata.scaled_r)
        self.assertTrue(np.all(self.vdata.ocb_e == self.out))

max_stars_repo_name: aburrell/ocbpy

max_issues_repo_name: aburrell/ocbpy

max_forks_repo_name: aburrell/ocbpy


text:
import numpy as np
import time
import statsmodels.api as sm

import est_dir

def compute_frac_fact(m, centre_point, no_vars, f, func_args,
                      region, set_all_positions):
    """
    Compute response function value at each observation of the design matrix
    centred at centre_point, where a 2^{10-6} design matrix is used.

    Parameters
    ----------
    m : integer
        Number of variables of the design matrix (columns).
    centre_point : 1-D array with shape (m,)
                   Centre point of design.
    no_vars : integer
              If no_vars < m, the size of the resulting
              design matrix is (n, no_vars). Since the centre_point is of size
              (m,), a random subset of variables will need to be chosen
              to evaluate the design matrix centred at centre_point. The
              parameter no_vars will be used to generate a random subset of
              positions, which correspond to the variable indices of
              centre_point in which to centre the design matrix.
    f : function
        response function.

        `f(point, *func_args) -> float`

        where point` is a 1-D array with shape(d, ) and func_args is
        a tuple of arguments needed to compute the response function value.
    func_args : tuple
                Arguments passed to the function f.
    region : float
             Region of exploration around the centre point.
    set_all_positions : 1-D array
                        If no_vars < m, it is ensured that new variables of
                        centre_point are used to centre the design matrix.
                        That is, the random subset of variables of centre_point
                        will be selected randomly from set_all_positions.

    Returns
    -------
    design : 2-D array
             2^{10-6} design matrix.
    y : 1-D array
        Contains the response function values at each observation of the
        design matrix centred at centre_point.
    positions : 1-D array
                Positions of centre_point in which the design matrix has been
                centred. If the design matrix is of size (n, m), then
                positions = (1,2,...,m). Otherwise, if the design matrix is
                of size (n, no_vars), where no_vars < m, then
                positions will be of size (num_vars,) with entries chosen
                randomly from set_all_positions.
    func_evals : integer
                 Number of times the response function has been evaluated.
    """
    positions = np.sort(np.random.choice(set_all_positions, no_vars,
                        replace=False))
    assert(np.unique(positions).shape[0] == no_vars)
    if no_vars == 10:
        design = np.array([[-1,  -1,  -1,  -1,  -1,  -1,  -1,  -1,  +1,  +1],
                           [+1,  -1,  -1,  -1,  +1,  -1,  +1,  +1,  -1,  -1],
                           [-1,  +1,  -1,  -1,  +1,  +1,  -1,  +1,  -1,  -1],
                           [+1,  +1,  -1,  -1,  -1,  +1,  +1,  -1,  +1,  +1],
                           [-1,  -1,  +1,  -1,  +1,  +1,  +1,  -1,  -1,  +1],
                           [+1,  -1,  +1,  -1,  -1,  +1,  -1,  +1,  +1,  -1],
                           [-1,  +1,  +1,  -1,  -1,  -1,  +1,  +1,  +1,  -1],
                           [+1,  +1,  +1,  -1,  +1,  -1,  -1,  -1,  -1,  +1],
                           [-1,  -1,  -1,  +1,  -1,  +1,  +1,  +1,  -1,  +1],
                           [+1,  -1,  -1,  +1,  +1,  +1,  -1,  -1,  +1,  -1],
                           [-1,  +1,  -1,  +1,  +1,  -1,  +1,  -1,  +1,  -1],
                           [+1,  +1,  -1,  +1,  -1,  -1,  -1,  +1,  -1,  +1],
                           [-1,  -1,  +1,  +1,  +1,  -1,  -1,  +1,  +1,  +1],
                           [+1,  -1,  +1,  +1,  -1,  -1,  +1,  -1,  -1,  -1],
                           [-1,  +1,  +1,  +1,  -1,  +1,  -1,  -1,  -1,  -1],
                           [+1,  +1,  +1,  +1,  +1,  +1,  +1,  +1,  +1,  +1]])
        n_temp = 16
    else:
        raise ValueError('Incorrect no_vars choice')
    y, func_evals = est_dir.compute_y(centre_point, design, positions, n_temp,
                                      m, f, func_args, region)
    return design, y, positions, func_evals

def compute_direction_LS(m, centre_point, f, func_args, no_vars, region):
    """
    Assume that the response function is approx linear in the neighbourhood
    of centre_point. Estimate unknown parameters of the linear model by
    using OLSE. Then direction is (theta_1,...,theta_m).

    Parameters
    ----------
    m : integer
        Number of variables.
    centre_point : 1-D array with shape (m,)
                   Centre point of design.
    f : function
        response function.

        `f(point, *func_args) -> float`

        where point` is a 1-D array with shape(d, ) and func_args is
        a tuple of arguments needed to compute the response function value.
    func_args : tuple
                Arguments passed to the function f.
    no_vars : integer
              If no_vars < m, the size of the resulting
              design matrix is (n, no_vars). Since the centre_point is of size
              (m,), a random subset of variables will need to be chosen
              to evaluate the design matrix centred at centre_point. The
              parameter no_vars will be used to generate a random subset of
              positions, which correspond to the variable indices of
              centre_point in which to centre the design matrix.
    region : float
             Region of exploration around the centre point.

    Returns
    -------
    direction : 1-D array
                Estimate of the direction i.e. (theta_1,...,theta_m),
                where theta_i (i=1,...,m) are coefficients of the linear
                model.
    func_evals_count : integer
                       Number of response function evaluations to compute
                       direction.
    """

    func_evals_count = 0
    set_all_positions = np.arange(m)
    total_checks = int(np.floor(m / no_vars))
    act_design, y, positions, func_evals = (compute_frac_fact
                                            (m, centre_point, no_vars,
                                             f, func_args, region,
                                             set_all_positions))
    func_evals_count += func_evals
    full_act_design = np.ones((act_design.shape[0], act_design.shape[1] + 1))
    full_act_design[:, 1:] = act_design
    est = sm.OLS(y, full_act_design)
    results = est.fit()
    index_vars = 1
    while results.f_pvalue >= 0.1:
        if index_vars >= total_checks:
            break
        set_all_positions = np.setdiff1d(set_all_positions, positions)
        act_design, y, positions, func_evals = (compute_frac_fact
                                                (m, centre_point, no_vars,
                                                 f, func_args, region,
                                                 set_all_positions))
        func_evals_count += func_evals
        full_act_design = np.ones((act_design.shape[0],
                                   act_design.shape[1] + 1))
        full_act_design[:, 1:] = act_design
        est = sm.OLS(y, full_act_design)
        results = est.fit()
        index_vars += 1
    if results.f_pvalue < 0.1:
        direction = np.zeros((m,))
        direction[positions] = est_dir.divide_abs_max_value(
                               (results.params)[1:])
        assert(max(abs(direction) == 1))
        return direction, func_evals_count
    else:
        return False, func_evals_count

def calc_first_phase_RSM_LS(centre_point, init_func_val, f, func_args,
                            m, const_back, back_tol, const_forward,
                            forward_tol, step, no_vars, region):
    """
    Compute iteration of local search, where search direction is estimated
    by (theta_1,...,theta_m), where theta_i (i=1,...,m) are coefficients
    of the linear model, and the step size is computed using either forward or
    backward tracking.

    Parameters
    ----------
    centre_point : 1-D array with shape (m,)
                   Centre point of design.
    init_func_val: float
                   Initial function value at centre_point. That is,
                   f(centre_point, *func_args).
    f : function
        response function.

        `f(point, *func_args) -> float`

        where point` is a 1-D array with shape(d, ) and func_args is
        a tuple of arguments needed to compute the response function value.
    func_args : tuple
                Arguments passed to the function f.
    m : integer
        Number of variables.
    const_back : float
                 If backward tracking is required, the initial guess of the
                 step size will be multiplied by const_back at each iteration
                 of backward tracking. That is,
                 t <- t * const_back
                 It should be noted that const_back < 1.
    back_tol : float
               It must be ensured that the step size computed by backward
               tracking is not smaller than back_tol. If this is the case,
               iterations of backward tracking are terminated. Typically,
               back_tol is a very small number.
    const_forward : float
                    The initial guess of the
                    step size will be multiplied by const_forward at each
                    iteration of forward tracking. That is,
                    t <- t * const_back
                    It should be noted that const_forward > 1.
    forward_tol : float
                  It must be ensured that the step size computed by forward
                  tracking is not greater than forward_tol. If this is the
                  case, iterations of forward tracking are terminated.
    step : float
           Initial guess of step size.
    no_vars : integer
              If no_vars < m, the size of the resulting
              design matrix is (n, no_vars). Since the centre_point is of size
              (m,), a random subset of variables will need to be chosen
              to evaluate the design matrix centred at centre_point. The
              parameter no_vars will be used to generate a random subset of
              positions, which correspond to the variable indices of
              centre_point in which to centre the design matrix.
    region : float
             Region of exploration around the centre point.

    Returns
    -------
    upd_point : 1-D array
                Updated centre_point after applying local search with
                estimated direction and step length.
    f_new : float
            Response function value at upd_point. That is,
            f(upd_point, *func_args).
    total_func_evals_step : integer
                            Total number of response function evaluations
                            to compute step length.
    total_func_evals_dir : integer
                           Total number of response function evaluations
                           to compute direction.
    flag : boolean
           Will be False if direction cannot be computed. Otherwise, will
           be True if direction can be computed.
    """

    direction, total_func_evals_dir = (compute_direction_LS
                                       (m, centre_point, f,
                                        func_args, no_vars,
                                        region))
    if direction is not False:
        (upd_point,
         f_new, total_func_evals_step) = (est_dir.combine_tracking
                                          (centre_point, init_func_val,
                                           direction, step, const_back,
                                           back_tol, const_forward,
                                           forward_tol, f, func_args))
        return (upd_point, f_new, total_func_evals_step,
                total_func_evals_dir, True)
    else:
        return centre_point, init_func_val, 0, total_func_evals_dir, False

def calc_its_until_sc_LS(centre_point, f, func_args, m, f_no_noise,
                         func_args_no_noise, region,
                         const_back=0.5, back_tol=0.000001,
                         const_forward=2, forward_tol=100000000,
                         tol_evals=10000000):
    """
    Compute iterations of Phase I of response surface methodology
    until some stopping criterion is met.
    The direction is estimated by the least squares coefficients
    of the linear model at each iteration. The step length is
    computed by forward or backward tracking.

    Parameters
    ----------
    centre_point : 1-D array with shape (m,)
                   Centre point of design.
    f : function
        response function.

        `f(point, *func_args) -> float`

        where point` is a 1-D array with shape(d, ) and func_args is
        a tuple of arguments needed to compute the response function value.
    func_args : tuple
                Arguments passed to the function f.
    m : integer
        Number of variables.
    f_no_noise : function
                 response function with no noise.

                `f_no_noise(point, *func_args_no_noise) -> float`

                where point` is a 1-D array with shape(d, ) and
                func_args_no_noise is a tuple of arguments needed to compute
                the response function value.
    func_args_no_noise : tuple
                         Arguments passed to the function f_no_noise.
    region : float
             Region of exploration around the centre point.
    const_back : float
                 If backward tracking is required, the initial guess of the
                 step size will be multiplied by const_back at each iteration
                 of backward tracking. That is,
                 t <- t * const_back
                 It should be noted that const_back < 1.
    back_tol : float
               It must be ensured that the step size computed by backward
               tracking is not smaller than back_tol. If this is the case,
               iterations of backward tracking are terminated. Typically,
               back_tol is a very small number.
    const_forward : float
                    The initial guess of the
                    step size will be multiplied by const_forward at each
                    iteration of forward tracking. That is,
                    t <- t * const_back
                    It should be noted that const_forward > 1.
    forward_tol : float
                  It must be ensured that the step size computed by forward
                  tracking is not greater than forward_tol. If this is the
                  case, iterations of forward tracking are terminated.
    tol_evals : int
                Maximum number of function evaluations before stopping.

    Returns
    -------
    upd_point : 1-D array
                Updated centre_point after applying local search with
                estimated direction and step length.
    init_func_val : float
                    Initial function value at initial centre_point.
    f_val : float
            Final response function value after stopping criterion
            has been met for phase I of RSM.
    full_time : float
                Total time taken.
    total_func_evals_step : integer
                            Total number of response function evaluations
                            to compute step length for all iterations.
    total_func_evals_dir : integer
                           Total number of response function evaluations
                           to compute direction for all iterations.
    no_iterations : integer
                    Total number of iterations of Phase I of RSM.
    store_good_dir : integer
                     Number of 'good' search directions. That is, the number
                     of times moving along the estimated search direction
                     improves the response function value with no noise.
    store_good_dir_norm : list
                          If a 'good' direction is determined,
                          distance of point and minimizer at the k-th
                          iteration subtracted by the distance of point
                          and minimizer at the (k+1)-th iteration is
                          stored.
    store_good_dir_func : list
                          If a 'good' direction is determined,
                          store the response function value with point
                          at the k-th iteration, subtracted by the response
                          function value with point at the
                          (k+1)-th iteration.
    """
    t0 = time.time()
    no_vars = 10
    store_good_dir = 0
    store_good_dir_norm = []
    store_good_dir_func = []
    total_func_evals_step = 0
    total_func_evals_dir = 0
    step = 1
    init_func_val = f(centre_point, *func_args)
    (upd_point, f_val,
     func_evals_step,
     func_evals_dir,
     flag) = (calc_first_phase_RSM_LS
              (centre_point, np.copy(init_func_val), f,
               func_args, m, const_back, back_tol, const_forward,
               forward_tol, step, no_vars, region))

    total_func_evals_step += func_evals_step
    total_func_evals_dir += func_evals_dir
    no_iterations = 1

    if (f_no_noise(centre_point, *func_args_no_noise) >
            f_no_noise(upd_point, *func_args_no_noise)):
        store_good_dir += 1
        store_good_dir_norm.append(np.linalg.norm(centre_point - func_args[0])
                                   - np.linalg.norm(upd_point - func_args[0]))
        store_good_dir_func.append(f_no_noise(centre_point,
                                              *func_args_no_noise)
                                   - f_no_noise(upd_point,
                                                *func_args_no_noise))

    while flag:
        centre_point = upd_point
        new_func_val = f_val
        step = 1
        (upd_point, f_val,
         func_evals_step,
         func_evals_dir,
         flag) = (calc_first_phase_RSM_LS
                  (centre_point, np.copy(new_func_val), f,
                   func_args, m, const_back, back_tol,
                   const_forward, forward_tol,
                   step, no_vars, region))
        total_func_evals_step += func_evals_step
        total_func_evals_dir += func_evals_dir
        no_iterations += 1
        if (f_no_noise(centre_point, *func_args_no_noise) >
                f_no_noise(upd_point, *func_args_no_noise)):
            store_good_dir += 1
            store_good_dir_norm.append(np.linalg.norm(centre_point -
                                                      func_args[0])
                                       - np.linalg.norm(upd_point -
                                                        func_args[0]))
            store_good_dir_func.append(f_no_noise(centre_point,
                                                  *func_args_no_noise)
                                       - f_no_noise(upd_point,
                                                    *func_args_no_noise))

        if (total_func_evals_step + total_func_evals_dir) >= tol_evals:
            break

    t1 = time.time()
    full_time = t1-t0
    return (upd_point, init_func_val, f_val, full_time,
            total_func_evals_step, total_func_evals_dir,
            no_iterations, store_good_dir,
            store_good_dir_norm, store_good_dir_func)

ext: py

max_stars_repo_name: Megscammell/Estimating-Direction

max_issues_repo_name: Megscammell/Estimating-Direction

max_forks_repo_name: Megscammell/Estimating-Direction

reason: import numpy,import statsmodels


text:
"""
sensors.py

Sensor classes used by the bots. The different types of sensors are:
 * Angular: Measures the burden angle between the robot and the target
 * Distance: Measures the distance in crows flight between the robot and the target
"""
from numpy import pi

class Sensor:
    """The baseclass used by all sensors."""

    __slots__ = {
        "game",
        "id", "value",
    }

    def __init__(self,
                 game,  # Type not specified due to circular imports
                 sensor_id: int = 0,
                 ):
        """
        Basic characteristics of a sensor.

        :param game: Reference to the game in which the sensor is used
        :param sensor_id: Identification number for the sensor
        """
        # Game-object
        self.game = game

        # Default sensor attributes
        self.id = sensor_id
        self.value = 0.0  # Zero value for initialized sensors

    def __str__(self):
        """Name of the sensor """
        raise NotImplemented

    def measure(self):
        """Store the sensor's current value in self.value."""
        raise NotImplemented

class AngularSensor(Sensor):
    """Angle deviation between bot and wanted direction in 'crows flight'."""

    def __init__(self,
                 game,  # Type not specified due to circular imports
                 sensor_id: int = 0,
                 clockwise: bool = True):
        """
        :param clockwise: Calculate the angular difference in clockwise direction
        :param game: Reference to the game in which the sensor is used
        :param sensor_id: Identification number for the sensor
        """
        # noinspection PyCompatibility
        super().__init__(game=game, sensor_id=sensor_id)
        self.clockwise = clockwise

    def __str__(self):
        return f"angular {'right' if self.clockwise else 'left'}"

    def measure(self):
        """Update self.value, result is a float between 0 and 2*PI."""
        # Get relative angle
        start_a = self.game.player.angle
        req_a = (self.game.target - self.game.player.pos).get_angle()

        # Constraint the value-range
        self.value = 2 * pi + start_a - req_a
        self.value %= 2 * pi

        # Check direction
        if not self.clockwise:
            self.value = abs(2 * pi - self.value)
        self.value = ((self.value + pi) % (2 * pi)) - pi  # constraint value between -pi and pi

        # Normalize
        self.value /= pi
        return self.value

class DeltaDistanceSensor(Sensor):
    """Difference in distance from bot to the target in 'crows flight' between current and the previous time-point."""

    def __init__(self,
                 game,  # Type not specified due to circular imports
                 sensor_id: int = 0):
        """
        :param game: Reference to the game in which the sensor is used
        :param sensor_id: Identification number for the sensor
        """
        super().__init__(game=game, sensor_id=sensor_id)
        self.distance: float = 0.0
        self.prev_distance: float = 0.0

    def __str__(self):
        return "delta_distance"

    def measure(self):
        """Update self.value to difference between previous distance and current distance."""
        self.prev_distance = self.distance  # Save previous distance
        start_p = self.game.player.pos
        end_p = self.game.target
        self.distance = (start_p - end_p).get_length()  # Get current measure
        if self.prev_distance == 0.0: self.prev_distance = self.distance  # Disable cold start
        self.value = self.prev_distance - self.distance  # Positive value == closer to target
        return self.value

class DistanceSensor(Sensor):
    """Distance from bot to the target in 'crows flight'."""

    def __init__(self,
                 game,  # Type not specified due to circular imports
                 normalizer: float,
                 sensor_id: int = 0):
        """
        :param game: Reference to the game in which the sensor is used
        :param normalizer: The constant by which the distance-value is normalized
        :param sensor_id: Identification number for the sensor
        """
        super().__init__(game=game, sensor_id=sensor_id)
        self.normalizer = normalizer

    def __str__(self):
        return "distance"

    def measure(self):
        """Update self.value to current distance between target and robot's center coordinate."""
        start_p = self.game.player.pos
        end_p = self.game.target
        self.value = (start_p - end_p).get_length() / self.normalizer
        return self.value

max_stars_repo_name: RubenPants/EvolvableRNN

max_issues_repo_name: RubenPants/EvolvableRNN

max_forks_repo_name: RubenPants/EvolvableRNN

reason: from numpy


text:
#!/usr/bin/env python
"""
 Copyright (c) 2018 Intel Corporation

 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
 You may obtain a copy of the License at

      http://www.apache.org/licenses/LICENSE-2.0

 Unless required by applicable law or agreed to in writing, software
 distributed under the License is distributed on an "AS IS" BASIS,
 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
"""
from __future__ import print_function
import sys
import os
from argparse import ArgumentParser
import cv2
import numpy as np
from time import time
from openvino.inference_engine import IENetwork, IEPlugin
import logging as log
from PIL import Image, ImageDraw, ImageFont
from multiprocessing import Pool
import time
import pprint

name_label_dict = {
0:  "Nucleoplasm",
1:  "Nuclear membrane",
2:  "Nucleoli",
3:  "Nucleoli fibrillar center" ,
4:  "Nuclear speckles"   ,
5:  "Nuclear bodies"   ,
6:  "Endoplasmic reticulum",
7:  "Golgi apparatus"   ,
8:  "Peroxisomes"   ,
9:  "Endosomes"   ,
10:  "Lysosomes"   ,
11:  "Intermediate filaments",
12:  "Actin filaments"   ,
13:  "Focal adhesion sites",
14:  "Microtubules"   ,
15:  "Microtubule ends",
16:  "Cytokinetic bridge",
17:  "Mitotic spindle"   ,
18:  "Microtubule organizing center" ,
19:  "Centrosome"   ,
20:  "Lipid droplets",
21:  "Plasma membrane",
22:  "Cell junctions"  ,
23:  "Mitochondria"   ,
24:  "Aggresome"   ,
25:  "Cytosol",
26:  "Cytoplasmic bodies",
27:  "Rods & rings"
}

def build_argparser():
    parser = ArgumentParser()
    parser.add_argument("-m", "--model", help="Path to an .xml file with a trained model.", required=True, type=str)
    parser.add_argument("-i", "--input", help="Path to a folder with images", required=True,
                        type=str)
    parser.add_argument("-l", "--cpu_extension",
                        help="MKLDNN (CPU)-targeted custom layers.Absolute path to a shared library with the kernels "
                             "impl.", type=str, default=None)
    parser.add_argument("-pp", "--plugin_dir", help="Path to a plugin folder", type=str, default=None)
    parser.add_argument("-d", "--device",
                        help="Specify the target device to infer on; CPU, GPU, FPGA or MYRIAD is acceptable. Sample "
                             "will look for a suitable plugin for device specified (CPU by default)", default="CPU",
                        type=str)
    parser.add_argument("-nt", "--number_top", help="Number of top results", default=10, type=int)

    return parser

def load_labels(labelsPath):
    labels = []
    with open(labelsPath, "r") as f:
        for line in f.readlines():
            labels.append(line.rstrip())
    return(labels)

def pre_process_image(imagePath):
    n, c, h, w = [1, 3, 299, 299]
    image = Image.open(imagePath)
    processingImg = image.resize((h, w), resample=Image.BILINEAR)
    #image         = image.convert('RGB')
    processingImg = (np.array(processingImg) - 0) / 255.0
    processingImg = processingImg.transpose((2, 0, 1))  # Change data layout from HWC to CHW
    processingImg = processingImg.reshape((n, c, h, w))
    return image, processingImg, imagePath

def pre_process_image_opencv(imagePath):
    n, c, h, w = [1, 3, 299, 299]
    image = cv2.imread(imagePath)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    #image = Image.fromarray(image).resize((h, w), resample=Image.BILINEAR)
    image = cv2.resize(image, (w, h))
    image = (np.array(image) - 0) / 255.0
    image = image.transpose((2, 0, 1))  # Change data layout from HWC to CHW
    image = image.reshape((n, c, h, w))
    return image, imagePath

def arrangeParallelData(listDir):
    print("Reading images in parallel...")
    poolSize=4
    # load images into mem using multi-processes
    pool = Pool(poolSize)

    totalStartTime = time.time()
    imagesBulk = pool.map(pre_process_image, listDir)
    totalTime = time.time() - totalStartTime
    pool.close()
    pool.join()
    print("Number of elements %s" % str(len(listDir)))
    print("Total Time Reading Files: %s seconds" % str(round(totalTime,4)))
    return imagesBulk

def drawBox(image, fileName, text):
    image = np.array(image)
    weight = 0
    #unormalize
    # image = np.squeeze(image)
    # image = image.transpose((1, 2, 0))

    print(image.shape)
    image = Image.fromarray(image)
    width, height = image.size
    padding = 5
    xMin = 0 + padding
    yMin = 0 + padding
    xMax = xMin + width - 2*padding
    yMax = yMin + height - 2*padding

    color = (0,255,0)
    draw = ImageDraw.Draw(image)
    for i in range(weight):
        draw.rectangle([(xMin-i, yMin-i), (xMax+i, yMax+i)], outline=color)

    fnt = ImageFont.truetype('Pillow/Tests/fonts/FreeMono.ttf', 13)
    draw.text((10,height-20), text, font=fnt, fill=(255,255,255,255))

    ######## FP32
    fnt = ImageFont.truetype('Pillow/Tests/fonts/FreeMono.ttf', 18)
    draw.text((10, 10), "2fps", font=fnt, fill=(255,255,255,255))
    draw.text((10, height-40), "TensorFlow - CPU (FP32)", font=fnt, fill=(255,255,255,255))
    image.save(os.path.join('selectedImages2fpsTF_CPU32',os.path.basename(fileName)))

    # fnt = ImageFont.truetype('Pillow/Tests/fonts/FreeMono.ttf', 18)
    # draw.text((10, 10), "16fps", font=fnt, fill=(255,255,255,255))
    # draw.text((10, height-40), "OpenVINO - CPU (FP32)", font=fnt, fill=(255,255,255,255))
    # image.save(os.path.join('selectedImages16fpsOV_CPU32',os.path.basename(fileName)))

    # fnt = ImageFont.truetype('Pillow/Tests/fonts/FreeMono.ttf', 18)
    # draw.text((10, 10), "23fps", font=fnt, fill=(255,255,255,255))
    # draw.text((10, height-40), "OpenVINO - IntelGPU (FP32)", font=fnt, fill=(255,255,255,255))
    # image.save(os.path.join('selectedImages23fpsOV_iGPU32',os.path.basename(fileName)))

    ######## FP16
    # fnt = ImageFont.truetype('Pillow/Tests/fonts/FreeMono.ttf', 18)
    # draw.text((10, 10), "36fps", font=fnt, fill=(255,255,255,255))
    # draw.text((10, height-40), "OpenVINO - IntelGPU (FP16)", font=fnt, fill=(255,255,255,255))
    # image.save(os.path.join('selectedImages36fpsOV_iGPU16',os.path.basename(fileName)))

def main():
    log.basicConfig(format="[ %(levelname)s ] %(message)s", level=log.INFO, stream=sys.stdout)
    args = build_argparser().parse_args()
    model_xml = args.model
    model_bin = os.path.splitext(model_xml)[0] + ".bin"
    imagesFolder = args.input

    listDir=[]
    for root, dirs, files in os.walk(imagesFolder):
        for file in files:
            listDir.append(os.path.join(root, file))

    #imagesBulk = arrangeParallelData(listDir)

    # Plugin initialization for specified device and load extensions library if specified
    plugin = IEPlugin(device=args.device, plugin_dirs=args.plugin_dir)
    if args.cpu_extension and 'CPU' in args.device:
        plugin.add_cpu_extension(args.cpu_extension)
    # Read IR
    net = IENetwork.from_ir(model=model_xml, weights=model_bin)
    assert len(net.inputs.keys()) == 1, "Sample supports only single input topologies"
    assert len(net.outputs) == 1, "Sample supports only single output topologies"
    input_blob = next(iter(net.inputs))
    out_blob = next(iter(net.outputs))

    # # Load network to the plugin
    exec_net = plugin.load(network=net)
    del net
    # Warmup with last image to avoid caching
    #exec_net.infer(inputs={input_blob: imagesBulk[-1][0]})
    exec_net.infer(inputs={input_blob: np.zeros((1, 3, 299, 299))})

    #accContainer = {}
    #for label in labels:
    #    accContainer[label] = {'total':0, 'correct':0, 'acc':0}

    infer_time = []
    #for image, fileName in imagesBulk:
    for fileName in listDir[:40]:
        image, processedImg, imagePath = pre_process_image(fileName)
        #image, imagePath = pre_process_image_opencv(fileName)

        #groundTruth = (os.path.dirname(fileName).split("/")[-1]).lower()
        #if groundTruth not in labels:
        #    continue

        t0 = time.time()
        res = exec_net.infer(inputs={input_blob: processedImg})
        infer_time.append((time.time()-t0)*1000)
        log.info("Average running time of one iteration: {} ms".format(np.average(np.asarray(infer_time))))
        res = res['dense_2/Sigmoid']
        idx = np.argsort(res[0])[-1]
        print(name_label_dict[idx])
        #accContainer[groundTruth]['total']+=1

        #if labels[top_ind[0]] == groundTruth:
        #    accContainer[groundTruth]['correct']+=1

        drawBox(image, fileName, name_label_dict[idx])
    del exec_net
    del plugin

    #for label in accContainer:
    #    accContainer[label]['acc'] = accContainer[label]['correct'] / accContainer[label]['total']

    #pprint.pprint(accContainer)
    avgElapsed = np.array(infer_time).mean() / 1000.0
    totalTime = round(np.array(infer_time).sum() / 1000.0,4)
    print("Average Elapsed: %s seconds/image" % str(round(avgElapsed,4)))
    print("Total time: %s seconds" %    str(totalTime))

if __name__ == '__main__':
    sys.exit(main() or 0)

max_stars_repo_name: fernandorovai/DL-inferenceSamples

max_issues_repo_name: fernandorovai/DL-inferenceSamples

max_forks_repo_name: fernandorovai/DL-inferenceSamples

reason: import numpy


text:
import pickle, gzip, skimage
from os import listdir, path, makedirs
from util.preprocessing import *
import pandas as pd
import numpy as np

def pickle_crossval_dataset(outfilename, x_train, y_train, x_test, y_test):
    pfile = gzip.open(outfilename, 'wb')
    dataset = [x_train, y_train, x_test, y_test]
    pickle.dump(dataset, pfile)
    pfile.close()

def load_npy(samples,target):
    X = np.load(samples)
    Y = np.load(target)
    print('Loaded images')
    return X,Y

def unpickle_crossval_dataset(filename):
    pfile = gzip.open(filename, 'r')
    dataset = pickle.load(pfile)
    return dataset[0], dataset[1], dataset[2], dataset[3]

def unpickle_features(filename):
    pfile = gzip.open(filename, 'r')
    dataset = pickle.load(pfile)
    pfile.close()
    return dataset

def pickle_hog_features(outfilename, foldername, num_classes):
    pfile = gzip.open(outfilename, 'wb')
    hog_dataset = dict()
    for i0 in range(1, num_classes + 1):
        filelist = listdir(foldername + str(i0))
        for filename in filelist:
            hog_image = generate_hog_features(foldername + str(i0) + "/" + filename)
            hog_dataset[filename] = hog_image
    pickle.dump(hog_dataset, pfile)
    pfile.close()

def pickle_daisy_features(outfilename, foldername, num_classes):
    pfile = gzip.open(outfilename, 'wb')
    daisy_dataset = dict()
    for i0 in range(1, num_classes + 1):
        filelist = listdir(foldername + str(i0))
        for filename in filelist:
            daisy_descs = generate_daisy_features(foldername + str(i0) + "/" + filename)
            daisy_dataset[filename] = daisy_descs
    pickle.dump(daisy_dataset, pfile)
    pfile.close()

def pickle_lbp_features(outfilename, foldername, num_classes):
    pfile = gzip.open(outfilename, 'wb')
    lbp_dataset = dict()
    for i0 in range(1, num_classes + 1):
        filelist = listdir(foldername + str(i0))
        for filename in filelist:
            lbp_features = generate_lbp_features(foldername + str(i0) + "/" + filename)
            lbp_dataset[filename] = lbp_features
    pickle.dump(lbp_dataset, pfile)
    pfile.close()

def crop_images(foldername, num_classes):
    for i0 in range(1, num_classes + 1):
        if not path.exists(foldername + 'cropped/' + str(i0)):
            makedirs(foldername + 'cropped/' + str(i0))
    user_folders = listdir(foldername)
    user_folders.remove('cropped')
    for user_folder in user_folders:
        curr_folder = foldername + user_folder + '/'
        crop_df = pd.read_csv(curr_folder + user_folder + '_loc.csv', index_col=0, header=0)
        filelist = listdir(curr_folder)
        for filename in filelist:
            file_class = str(class_alpha_dict[filename[0].upper()])
            user_num = user_folder.split('_')[1]
            if filename.split('.')[1] == 'csv':
                continue
            img_arr = io.imread(curr_folder+filename)
            crop_before_x = crop_df.loc[user_folder+'/'+filename, 'top_left_x']
            crop_before_y = crop_df.loc[user_folder+'/'+filename, 'top_left_y']
            crop_after_x = crop_df.loc[user_folder+'/'+filename, 'bottom_right_x']
            crop_after_y = crop_df.loc[user_folder+'/'+filename, 'bottom_right_y']
            img_cropped = img_arr[crop_before_y:crop_after_y, crop_before_x:crop_after_x]
            img_transformed = skimage.transform.resize(img_cropped, output_shape=(110, 110))
            skimage.io.imsave(foldername + 'cropped/' + file_class + "/" + user_num + '_' + filename, img_transformed)

def create_userpkl(users_folder,dataset_folder,data_dict):
    for users in users_folder:
        u_dict = {}
        for f in listdir(dataset_folder + users):
            if(f.endswith(".jpg")):
                u_dict[users +'/'+f] = data_dict[users +'/'+f]
        print("pickling " + users)
        f = gzip.open('data//'+users+'.pkl','wb',compresslevel=1)
        pickle.dump(u_dict,f)
        f.close()

max_stars_repo_name: achyudhk/Sign-Language-Recognition

max_issues_repo_name: lite-david/Sign-Language-Recognition

max_forks_repo_name: lite-david/Sign-Language-Recognition


text:
import numpy as np

from pyhsmm.models import WeakLimitHDPHSMMPython
from pyhsmm.models import WeakLimitHDPHSMM
from pybasicbayes.distributions.poisson import Poisson
from pyhsmm.util.stats import sample_discrete
from pyhlm.internals.internal_hsmm_states import LetterHSMMStatesPython, LetterHSMMStatesEigen

class LetterHSMMPython(WeakLimitHDPHSMMPython):
    _states_class = LetterHSMMStatesPython

    def resample_trans_distn_by_sampled_words(self, word_list):
        self.trans_distn.resample([np.array(word) for word in word_list])
        self._clear_caches()

    def resample_init_state_distn_by_sampled_words(self, word_list):
        self.init_state_distn.resample([word[0] for word in word_list])
        self._clear_caches()

    def resample_parameters_by_sampled_words(self, word_list):
        self.resample_dur_distns()
        self.resample_obs_distns()
        self.resample_trans_distn_by_sampled_words(word_list)
        self.resample_init_state_distn_by_sampled_words(word_list)

    def generate_word(self, word_size):
        nextstate_distn = self.init_state_distn.pi_0
        A = self.trans_distn.trans_matrix
        word = [-1] * word_size
        for idx in range(word_size):
            word[idx] = sample_discrete(nextstate_distn)
            nextstate_distn = A[word[idx]]
        return tuple(word)

    @property
    def params(self):
        obs_params = {"obs_distn({})".format(idx): obs_distn.params for idx, obs_distn in enumerate(self.obs_distns)}
        dur_params = {"dur_distn({})".format(idx): dur_distn.params for idx, dur_distn in enumerate(self.dur_distns)}
        bigram_params = {**self.init_state_distn.params, "trans_matrix":self.trans_distn.trans_matrix}
        return {"num_states": self.num_states, "obs_distns": obs_params, "dur_distns": dur_params, "bigram": bigram_params}

    @property
    def hypparams(self):
        obs_hypparams = {"obs_distn({})".format(idx): obs_distn.hypparams for idx, obs_distn in enumerate(self.obs_distns)}
        dur_hypparams = {"dur_distn({})".format(idx): dur_distn.hypparams for idx, dur_distn in enumerate(self.dur_distns)}
        bigram_hypparams = self.init_state_distn.hypparams
        return {"obs_distns": obs_hypparams, "dur_distns": dur_hypparams, "bigram": bigram_hypparams}

class LetterHSMM(WeakLimitHDPHSMM, LetterHSMMPython):
    _states_class = LetterHSMMStatesEigen

max_stars_repo_name: kishiyamat/npbdaa

max_issues_repo_name: RyoOzaki/npbdaa

max_forks_repo_name: kishiyamat/npbdaa