# Source code for sfepy.scripts.plot_condition_numbers

    #!/usr/bin/env python
    """
    Plot conditions numbers w.r.t. polynomial approximation order of reference
    element matrices for various FE polynomial spaces (bases).
    """
    import sys
    sys.path.append('.')
    from argparse import ArgumentParser
    import os.path as op
    from functools import partial
    import numpy as nm
    import matplotlib.pyplot as plt

    from sfepy import data_dir
    from sfepy.base.base import output, assert_
    from sfepy.base.timing import Timer
    from sfepy.discrete import FieldVariable, Material, Integral
    from sfepy.discrete.fem import Mesh, FEDomain, Field
    from sfepy.terms import Term
    from sfepy.solvers import eig
    from sfepy.mechanics.matcoefs import stiffness_from_lame

    helps = {
        'basis' :
        'name of the FE basis [default: %(default)s]',
        'max_order' :
        'maximum order of polynomials [default: %(default)s]',
        'matrix_type' :
        'matrix type [default: %(default)s]',
        'geometry' :
        'reference element geometry, one of "2_3", "2_4", "3_4", "3_8"'
        ' [default: %(default)s]',
        'output_dir' :
        'output directory',
        'no_show' :
        'do not show the figures',
    }

    def main():
        parser = ArgumentParser(description=__doc__)
        parser.add_argument('--version', action='version', version='%(prog)s')
        parser.add_argument('-b', '--basis', metavar='name',
                            action='store', dest='basis',
                            default='lagrange', help=helps['basis'])
        parser.add_argument('-n', '--max-order', metavar='order', type=int,
                            action='store', dest='max_order',
                            default=10, help=helps['max_order'])
        parser.add_argument('-m', '--matrix', action='store', dest='matrix_type',
                            choices=['laplace', 'elasticity', 'smass', 'vmass'],
                            default='laplace', help=helps['matrix_type'])
        parser.add_argument('-g', '--geometry', metavar='name',
                            action='store', dest='geometry',
                            default='2_4', help=helps['geometry'])
        parser.add_argument('-o', '--output-dir', metavar='path',
                            action='store', dest='output_dir',
                            default=None, help=helps['output_dir'])
        parser.add_argument('--no-show',
                            action='store_false', dest='show',
                            default=True, help=helps['no_show'])
        options = parser.parse_args()

        dim, n_ep = int(options.geometry[0]), int(options.geometry[2])
        output('reference element geometry:')
        output('  dimension: %d, vertices: %d' % (dim, n_ep))

        n_c = {'laplace' : 1, 'elasticity' : dim,
               'smass' : 1, 'vmass' : dim}[options.matrix_type]

        output('matrix type:', options.matrix_type)
        output('number of variable components:',  n_c)

        output('polynomial space:', options.basis)

        output('max. order:', options.max_order)

        mesh = Mesh.from_file(data_dir + '/meshes/elements/%s_1.mesh'
                              % options.geometry)
        domain = FEDomain('domain', mesh)
        omega = domain.create_region('Omega', 'all')

        orders = nm.arange(1, options.max_order + 1, dtype=nm.int32)
        conds = []

        for order in orders:
            output('order:', order, '...')

            field = Field.from_args('fu', nm.float64, n_c, omega,
                                    approx_order=order,
                                    space='H1', poly_space_basis=options.basis)

            quad_order = 2 * field.approx_order
            output('quadrature order:', quad_order)

            integral = Integral('i', order=quad_order)
            qp, _ = integral.get_qp(options.geometry)
            output('number of quadrature points:', qp.shape[0])

            u = FieldVariable('u', 'unknown', field)
            v = FieldVariable('v', 'test', field, primary_var_name='u')

            m = Material('m', D=stiffness_from_lame(dim, 1.0, 1.0))

            if options.matrix_type == 'laplace':
                term = Term.new('dw_laplace(v, u)',
                                integral, omega, v=v, u=u)
                n_zero = 1

            elif options.matrix_type == 'elasticity':
                term = Term.new('dw_lin_elastic(m.D, v, u)',
                                integral, omega, m=m, v=v, u=u)
                n_zero = (dim + 1) * dim // 2

            elif options.matrix_type in ('smass', 'vmass'):
                term = Term.new('dw_dot(v, u)',
                                integral, omega, v=v, u=u)
                n_zero = 0

            term.setup()

            output('assembling...')
            timer = Timer(start=True)
            mtx, iels = term.evaluate(mode='weak', diff_var='u')
            output('...done in %.2f s' % timer.stop())
            mtx = mtx[0, 0]

            try:
                assert_(nm.max(nm.abs(mtx - mtx.T)) < 1e-10)

            except:
                from sfepy.base.base import debug; debug()

            output('matrix shape:', mtx.shape)

            eigs = eig(mtx, solver_kind='eig.sgscipy', eigenvectors=False)
            eigs.sort()

            # Zero 'true' zeros.
            eigs[:n_zero] = 0.0

            ii = nm.where(eigs < 0.0)[0]
            if len(ii):
                output('matrix is not positive semi-definite!')

            ii = nm.where(eigs[n_zero:] < 1e-12)[0]
            if len(ii):
                output('matrix has more than %d zero eigenvalues!' % n_zero)

            output('smallest eigs:\n', eigs[:10])

            ii = nm.where(eigs > 0.0)[0]
            emin, emax = eigs[ii[[0, -1]]]

            output('min:', emin, 'max:', emax)

            cond = emax / emin
            conds.append(cond)

            output('condition number:', cond)

            output('...done')

        if options.output_dir is not None:
            indir = partial(op.join, options.output_dir)

        else:
            indir = None

        plt.rcParams['font.size'] = 12
        plt.rcParams['lines.linewidth'] = 3

        fig, ax = plt.subplots()
        ax.semilogy(orders, conds)
        ax.set_xticks(orders)
        ax.set_xticklabels(orders)
        ax.set_xlabel('polynomial order')
        ax.set_ylabel('condition number')
        ax.set_title(f'{options.basis.capitalize()} basis')
        ax.grid()
        plt.tight_layout()
        if indir is not None:
            fig.savefig(indir(f'{options.basis}-{options.matrix_type}-'
                              f'{options.geometry}-{options.max_order}-xlin.png'),
                        bbox_inches='tight')

        fig, ax = plt.subplots()
        ax.loglog(orders, conds)
        ax.set_xticks(orders)
        ax.set_xticklabels(orders)
        ax.set_xlabel('polynomial order')
        ax.set_ylabel('condition number')
        ax.set_title(f'{options.basis.capitalize()} basis')
        ax.grid()
        plt.tight_layout()
        if indir is not None:
            fig.savefig(indir(f'{options.basis}-{options.matrix_type}-'
                              f'{options.geometry}-{options.max_order}-xlog.png'),
                        bbox_inches='tight')

        if options.show:
            plt.show()

    if __name__ == '__main__':
        main()