example_deriv_check.py

from __future__ import print_function, division
import numpy as np

import unittest

from openmdao.api import Problem, Group, IndepVarComp, view_model

from six import iteritems
from numpy.testing import assert_almost_equal

from openaerostruct.structures.section_properties_tube import SectionPropertiesTube
from openaerostruct.utils.testing import get_default_surfaces, run_test


def view_mat(mat):
    """ Helper function used to visually examine matrices. """
    import matplotlib.pyplot as plt
    if len(mat.shape) > 2:
        mat = np.sum(mat, axis=2)
    im = plt.imshow(mat.real, interpolation='none')
    plt.colorbar(im, orientation='horizontal')
    plt.show()

class Test(unittest.TestCase):
    """
    This class contains two tests on a single component.
    Each test verifies the derivatives found within the component by using
    OpenMDAO's derivative approximation functions to compare against
    finite-difference or complex-step.
    """

    def test_quick(self):
        """ Short pre-setup test to compare component derivs. """
        surfaces = get_default_surfaces()

        comp = SectionPropertiesTube(surface=surfaces[0])

        run_test(self, comp, complex_flag=True)

    def test_detailed(self):
        """ This is a longer version of the previous method, with plotting. """

        # Load in default lifting surfaces to setup the comparison
        surfaces = get_default_surfaces()

        # Instantiate an OpenMDAO problem and add the component we want to test
        # as asubsystem, giving that component a default lifting surface
        prob = Problem()
        prob.model.add_subsystem('tube', SectionPropertiesTube(surface=surfaces[0]))

        # Set up the problem and ensure it uses complex arrays so we can check
        # the derivatives using complex step
        prob.setup(force_alloc_complex=True)

        # Actually run the model, which is just a component in this case, then
        # check the derivatives and store the results in the `check` dict
        prob.run_model()
        check = prob.check_partials(compact_print=False)

        # Loop through this `check` dictionary and visualize the approximated
        # and computed derivatives
        for key, subjac in iteritems(check[list(check.keys())[0]]):
            print()
            print(key)
            print('Approximated derivative (either fd or cs):')
            view_mat(subjac['J_fd'])

            print('Computed derivative set by user:')
            view_mat(subjac['J_fwd'])

        # Loop through the `check` dictionary and perform assert that the
        # approximated deriv must be very close to the computed deriv
        for key, subjac in iteritems(check[list(check.keys())[0]]):
            if subjac['magnitude'].fd > 1e-6:
                assert_almost_equal(
                    subjac['rel error'].forward, 0., err_msg='deriv of %s wrt %s' % key)
                assert_almost_equal(
                    subjac['rel error'].reverse, 0., err_msg='deriv of %s wrt %s' % key)

# Run the tests included in this script
if __name__ == '__main__':
    unittest.main()