Source code for gemseo.third_party.pyxdsm.matrix_eqn

# -*- coding: utf-8 -*-
# pyxdsm is a python library for generating publication quality PDF XDSM
# diagrams.This library is a thin wrapper that uses the TIKZ library and
# LaTeX to build the PDFs. Source: https://github.com/mdolab/pyXDSM
#
# OpenMDAO Open Source License:
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this software 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 division

from collections import namedtuple
import os

import numpy as np


# color pallette link: http://paletton.com/#uid=72Q1j0kllllkS5tKC9H96KClOKC
base_file_start = r"""\documentclass[border=0pt]{standalone}
% Justin S. Gray 2018
% Based off code by John T. Hwang (2014), who based his code off Alec Jacobson: http://www.alecjacobson.com/weblog/?p=1289

% nc = necessary comment [do not remove]

% Four rules for using these macros:
% 1. Always start with a row
% 2. Rows contain cols and cols contain rows
% 3. Mats should be on at least the 3rd level (row->col->mat, minimum)
% 4. If a row contains the mat, add &; if a col contains the mat, add \\

% ---------------------------------------

\usepackage{tikz}
\usepackage{ifthen}
\usepackage{esdiff}
\usepackage{varwidth}

\definecolor{tableau0}{RGB}{77, 121, 168}
\definecolor{tableau1}{RGB}{242, 142, 43}
\definecolor{tableau2}{RGB}{255, 87, 88}
\definecolor{tableau3}{RGB}{118, 183, 178}
\definecolor{tableau4}{RGB}{89, 161, 78}
\definecolor{tableau5}{RGB}{237, 201, 72}
\definecolor{tableau6}{RGB}{176, 121, 162}
\definecolor{tableau7}{RGB}{255, 157, 167}
\definecolor{tableau8}{RGB}{156, 116, 94}
\definecolor{tableau9}{RGB}{186, 176, 172}


\newcommand{\thk}{0.01in}
\newcommand{\thkln}{0.02in}

% \blockmat{width}{height}{text}{block_options}{other}
\newcommand{\blockmat}[5]{
\begin{tikzpicture}
    \draw[draw=white,fill=white,#4,line width=\thk] (0,0) rectangle( #1-\thk,#2-\thk);
    #5
    \draw (#1/2, #2/2) node {#3};
\end{tikzpicture}
}

% blockempty{width}{height}{text}
\newcommand{\blockempty}[3]{
    \blockmat{#1}{#2}{#3}{draw=white,fill=white}{}%
}

% \blockmat{width}{height}{text}{block_options}{diagonal_width}{diagonal_options}
\newcommand{\blockdiag}[6]{
    \blockmat{#1}{#2}{#3}{#4}
    {
    \draw[#6,line width=\thk] (0,#2-\thk) -- (#5,#2-\thk) -- ( #1-\thk,#5) -- ( #1-\thk,0) -- ( #1-\thk - #5,0) -- (0,#2-\thk -#5) --cycle;
    }%
}

% \blockddots{width}{height}{text}{block_options}{dot_radius}{dot_options}{dot_h}{dot_v}
\newcommand{\blockdots}[8]{
    \blockmat{#1}{#2}{#3}{#4}%
    {%
    \ifthenelse{\equal{#5}{}}
    {\newcommand\dotradius{0.01in}}
    {\newcommand\dotradius{#5}}%
    \filldraw[#6] (#1/2, #2/2) circle (0.5*\dotradius);%
    \filldraw[#6] (#1/2 + #7, #2/2 + #8) circle (0.5*\dotradius);%
    \filldraw[#6] (#1/2 - #7, #2/2 - #8) circle (0.5*\dotradius);%
    }%
}

% \leftbracket{width}{height}{options}
\newcommand{\leftbracket}[3]{
    \begin{tikzpicture}
        \coordinate (iSW) at (\thk+\thkln/2,\thk+\thkln/2);
        \coordinate (iNW) at (\thk+\thkln/2,#2-\thk-\thkln/2);
        \coordinate (iSE) at (#1-\thk-\thkln/2,\thk+\thkln/2);
        \coordinate (iNE) at (#1-\thk-\thkln/2,#2-\thk-\thkln/2);
        \coordinate (oSW) at (\thk/2,\thk/2);
        \coordinate (oNW) at (\thk/2,#2-\thk/2);
        \coordinate (oSE) at (#1-\thk/2,\thk/2);
        \coordinate (oNE) at (#1-\thk/2,#2-\thk/2);
        \draw[#3,line width=\thkln] (iNE) -- (iNW) -- (iSW) -- (iSE);
        \draw[draw=white,line width=\thk] (oNE) -- (oNW) -- (oSW) -- (oSE);
    \end{tikzpicture}%nc
}

% \rightbracket{width}{height}{options}
\newcommand{\rightbracket}[3]{
    \begin{tikzpicture}
        \coordinate (iSW) at (\thk+\thkln/2,\thk+\thkln/2);
        \coordinate (iNW) at (\thk+\thkln/2,#2-\thk-\thkln/2);
        \coordinate (iSE) at (#1-\thk-\thkln/2,\thk+\thkln/2);
        \coordinate (iNE) at (#1-\thk-\thkln/2,#2-\thk-\thkln/2);
        \coordinate (oSW) at (\thk/2,\thk/2);
        \coordinate (oNW) at (\thk/2,#2-\thk/2);
        \coordinate (oSE) at (#1-\thk/2,\thk/2);
        \coordinate (oNE) at (#1-\thk/2,#2-\thk/2);
        \draw[#3,line width=\thkln] (iNW) -- (iNE) -- (iSE) -- (iSW);
        \draw[draw=white,line width=\thk] (oNW) -- (oNE) -- (oSE) -- (oSW);
    \end{tikzpicture}%nc
}

% \upperbracket{width}{height}{options}
\newcommand{\upperbracket}[3]{
    \begin{tikzpicture}
        \coordinate (iSW) at (\thk+\thkln/2,\thk+\thkln/2);
        \coordinate (iNW) at (\thk+\thkln/2,#2-\thk-\thkln/2);
        \coordinate (iSE) at (#1-\thk-\thkln/2,\thk+\thkln/2);
        \coordinate (iNE) at (#1-\thk-\thkln/2,#2-\thk-\thkln/2);
        \coordinate (oSW) at (\thk/2,\thk/2);
        \coordinate (oNW) at (\thk/2,#2-\thk/2);
        \coordinate (oSE) at (#1-\thk/2,\thk/2);
        \coordinate (oNE) at (#1-\thk/2,#2-\thk/2);
        \draw[#3,line width=\thkln] (iSW) -- (iNW) -- (iNE) -- (iSE);
        \draw[draw=white,line width=\thk] (oSW) -- (oNW) -- (oNE) -- (oSE);
    \end{tikzpicture}%nc
}

% \lowerbracket{width}{height}{options}
\newcommand{\lowerbracket}[3]{
    \begin{tikzpicture}
        \coordinate (iSW) at (\thk+\thkln/2,\thk+\thkln/2);
        \coordinate (iNW) at (\thk+\thkln/2,#2-\thk-\thkln/2);
        \coordinate (iSE) at (#1-\thk-\thkln/2,\thk+\thkln/2);
        \coordinate (iNE) at (#1-\thk-\thkln/2,#2-\thk-\thkln/2);
        \coordinate (oSW) at (\thk/2,\thk/2);
        \coordinate (oNW) at (\thk/2,#2-\thk/2);
        \coordinate (oSE) at (#1-\thk/2,\thk/2);
        \coordinate (oNE) at (#1-\thk/2,#2-\thk/2);
        \draw[#3,line width=\thkln] (iNW) -- (iSW) -- (iSE) -- (iNE);
        \draw[draw=white,line width=\thk] (oNW) -- (oSW) -- (oSE) -- (oNE);
    \end{tikzpicture}%nc
}

% a hack so that I don't have to worry about the number of columns or
% spaces between columns in the tabular environment
\newenvironment{blockmatrixtabular}
{%nc
    \renewcommand{\arraystretch}{0}%nc
    \begin{tabular}{
    @{}l@{}l@{}l@{}l@{}l@{}l@{}l@{}l@{}l@{}l@{}l@{}l@{}l@{}l@{}l@{}l@{}l@{}l@{}l
    @{}l@{}l@{}l@{}l@{}l@{}l@{}l@{}l@{}l@{}l@{}l@{}l@{}l@{}l@{}l@{}l@{}l@{}l@{}l
    @{}l@{}l@{}l@{}l@{}l@{}l@{}l@{}l@{}l@{}l@{}l@{}l@{}l@{}l@{}l@{}l@{}l@{}l@{}l
    @{}
    }%nc
}
{
    \end{tabular}%nc
}

% \blockcol{
% }
\newcommand{\blockcol}[1]{\vtop{\null\hbox{\begin{blockmatrixtabular}#1\end{blockmatrixtabular}}}&}

% \blockrow{
% }
\newcommand{\blockrow}[1]{\begin{blockmatrixtabular}#1\end{blockmatrixtabular}\\}


\begin{document}
\begin{varwidth}{10\textwidth}

\newcommand\mwid{0.5in}
\newcommand\wid{0.15in}
\newcommand\comp{0.3in}
\newcommand\ext{0.5in}
\newcommand\dt{0.03in}
\newcommand\txt{0.8in}

\definecolor{Tgrey}{rgb}{0.9,0.9,0.9}
\definecolor{Tred}{rgb}{1.0,0.722,0.714}
\definecolor{Tgreen}{rgb}{0.639,0.89,0.655}
\definecolor{Tblue}{rgb}{0.667,0.631,0.843}
\definecolor{Tyellow}{rgb}{1,0.941,0.714}

\definecolor{Lred}{rgb}{17.3,0.063,0.059}
\definecolor{Lgreen}{rgb}{0.047,0.133,0.051}
\definecolor{Lblue}{rgb}{0.063,0.051,0.118}
\definecolor{Lyellow}{rgb}{0.173,0.149,0.059}

\definecolor{Dgrey}{rgb}{0.4,0.4,0.4}
\definecolor{Dred}{rgb}{1.0,0.333,0.318}
\definecolor{Dgreen}{rgb}{0.282,0.89,0.322}
\definecolor{Dblue}{rgb}{0.42,0.341,0.843}
\definecolor{Dyellow}{rgb}{1.0,0.863,0.318}

\definecolor{Bred}{rgb}{0.302,0.8,0.0}
\definecolor{Bgreen}{rgb}{0.4,1.0,0.4}
\definecolor{Bblue}{rgb}{0.043,0.012,0.208}
\definecolor{Byellow}{rgb}{0.302,0.243,0.0}

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
"""
base_file_end = r"""
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\end{varwidth}
\end{document}"""


Variable = namedtuple('Variable', field_names=['size', 'idx', 'text', 'color'])

CellData = namedtuple('CellData', field_names=['text', 'color', 'highlight'])


def _color(base_color, h_light):
    if h_light == -1:
        color = 'white'
    elif h_light == 0:
        color = 'Tgrey'
    elif h_light == 1:
        color = 'T{}'.format(base_color)
    elif h_light == 2:
        color = 'D{}'.format(base_color)
    elif h_light == 3:
        color = 'B{}'.format(base_color)

    elif h_light == "diag":
        color = base_color

    return color


def _write_tikz(tikz, out_file, build=True, cleanup=True):
    with open('{}.tex'.format(out_file), 'w') as f:
        f.write(base_file_start)
        f.write(tikz)
        f.write(base_file_end)

    if build:
        os.system('pdflatex {}.tex'.format(out_file))

        if cleanup:
            for ext in ['aux', 'fdb_latexmk', 'fls', 'log', 'tex']:
                f_name = '{}.{}'.format(out_file, ext)
                if os.path.exists(f_name):
                    os.remove(f_name)


[docs]class TotalJacobian(object): def __init__(self): self._variables = {} self._j_inputs = {} self._n_inputs = 0 self._i_outputs = {} self._n_outputs = 0 self._connections = {} self._ij_connections = {} self._setup = False
[docs] def add_input(self, name, size=1, text=''): self._variables[name] = Variable( size=size, idx=self._n_inputs, text=text, color=None) self._j_inputs[self._n_inputs] = self._variables[name] self._n_inputs += 1
[docs] def add_output(self, name, size=1, text=''): self._variables[name] = Variable( size=size, idx=self._n_outputs, text=text, color=None) self._i_outputs[self._n_outputs] = self._variables[name] self._n_outputs += 1
[docs] def connect(self, src, target, text='', color='tableau0'): if isinstance(target, (list, tuple)): for t in target: self._connections[ src, t] = CellData( text=text, color=color, highlight='diag') else: self._connections[ src, target] = CellData( text=text, color=color, highlight='diag')
def _process_vars(self): if self._setup: return # deal with connections for (src, target), cell_data in self._connections.items(): i_src = self._variables[src].idx j_target = self._variables[target].idx self._ij_connections[i_src, j_target] = cell_data self._setup = True
[docs] def write(self, out_file=None, build=True, cleanup=True): """ Write output files for the matrix equation diagram. This produces the following: - {file_name}.tikz A file containing the TIKZ definition of the tikz diagram. - {file_name}.tex A standalone document wrapped around an include of the TIKZ file which can be compiled to a pdf. - {file_name}.pdf An optional compiled version of the standalone tex file. Parameters ---------- file_name : str The prefix to be used for the output files build : bool Flag that determines whether the standalone PDF of the XDSM will be compiled. Default is True. cleanup: bool Flag that determines if padlatex build files will be deleted after build is complete """ self._process_vars() tikz = [] # label the columns tikz.append(r'\blockrow{') # emtpy column for the row labels tikz.append(r' \blockcol{') tikz.append(r' \blockempty{%s*\comp}{%s*\comp}{%s}\\' % (1, 1, '')) tikz.append(r' }') for j in range(self._n_inputs): var = self._j_inputs[j] col_size = var.size tikz.append(r' \blockcol{') tikz.append( r' \blockempty{%s*\comp}{%s*\comp}{%s}\\' % (col_size, 1, var.text)) tikz.append(r' }') tikz.append(r'}') for i in range(self._n_outputs): output = self._i_outputs[i] row_size = output.size tikz.append(r'\blockrow{') # label the row with the output name tikz.append(r' \blockcol{') tikz.append( r' \blockempty{%s*\comp}{%s*\comp}{%s}\\' % (1, row_size, output.text)) tikz.append(r' }') for j in range(self._n_inputs): var = self._j_inputs[j] col_size = var.size tikz.append(r' \blockcol{') if (j, i) in self._ij_connections: cell_data = self._ij_connections[(j, i)] conn_color = 'T{}'.format(var.color) if cell_data.color is not None: conn_color = _color( cell_data.color, cell_data.highlight) tikz.append( r' \blockmat{%s*\comp}{%s*\comp}{%s}{draw=white,fill=%s}{}\\' % (col_size, row_size, cell_data.text, conn_color)) else: tikz.append( r' \blockempty{%s*\comp}{%s*\comp}{}\\' % (col_size, row_size)) tikz.append(r' }') tikz.append(r'}') jac_tikz = "\n".join(tikz) _write_tikz(jac_tikz, out_file, build, cleanup)
[docs]class MatrixEquation(object): def __init__(self): self._variables = {} self._ij_variables = {} self._n_vars = 0 self._connections = {} self._ij_connections = {} self._text = {} self._ij_text = {} self._total_size = 0 self._setup = False self._terms = []
[docs] def clear_terms(self): self._terms = []
[docs] def add_variable(self, name, size=1, text='', color='blue'): self._variables[name] = Variable( size=size, idx=self._n_vars, text=text, color=color) self._ij_variables[self._n_vars] = self._variables[name] self._n_vars += 1 self._total_size += size
[docs] def connect(self, src, target, text='', color=None, highlight=1): if isinstance(target, (list, tuple)): for t in target: self._connections[ src, t] = CellData( text=text, color=color, highlight=highlight) else: self._connections[ src, target] = CellData( text=text, color=color, highlight=highlight)
[docs] def text(self, src, target, text): "don't connect the src and target, but put some text where a connection would be" self._text[src, target] = CellData(text=text, color=None, highlight=-1)
def _process_vars(self): """map all the data onto i,j grid""" if self._setup: return # deal with connections for (src, target), cell_data in self._connections.items(): i_src = self._variables[src].idx i_target = self._variables[target].idx self._ij_connections[i_src, i_target] = cell_data for (src, target), cell_data in self._text.items(): i_src = self._variables[src].idx j_target = self._variables[target].idx self._ij_text[i_src, j_target] = cell_data self._setup = True
[docs] def jacobian(self, transpose=False): self._process_vars() tikz = [] for i in range(self._n_vars): tikz.append(r'\blockrow{') row_size = self._ij_variables[i].size for j in range(self._n_vars): var = self._ij_variables[j] col_size = var.size tikz.append(r' \blockcol{') if transpose: location = (i, j) else: location = (j, i) if i == j: tikz.append( r' \blockmat{%s*\comp}{%s*\comp}{%s}{draw=white,fill=D%s}{}\\' % (col_size, row_size, var.text, var.color)) elif location in self._ij_connections: cell_data = self._ij_connections[location] conn_color = 'T{}'.format(var.color) if cell_data.color is not None: conn_color = _color( cell_data.color, cell_data.highlight) tikz.append( r' \blockmat{%s*\comp}{%s*\comp}{%s}{draw=white,fill=%s}{}\\' % (col_size, row_size, cell_data.text, conn_color)) elif location in self._ij_text: cell_data = self._ij_text[location] tikz.append( r' \blockempty{%s*\comp}{%s*\comp}{%s}\\' % (col_size, row_size, cell_data.text)) else: tikz.append( r' \blockempty{%s*\comp}{%s*\comp}{}\\' % (col_size, row_size)) tikz.append(r' }') tikz.append(r'}') lhs_tikz = "\n".join(tikz) self._terms.append(lhs_tikz) return lhs_tikz
[docs] def vector(self, base_color='red', highlight=None): self._process_vars() tikz = [] if highlight is None: highlight = np.ones(self._n_vars) for i, h_light in enumerate(highlight): color = _color(base_color, h_light) row_size = self._ij_variables[i].size tikz.append(r'\blockrow{\blockcol{') if h_light == "diag": tikz.append( r' \blockdiag{1*\comp}{%s*\comp}{}{draw=white,fill=T%s}{\dt}{draw=white,fill=D%s}\\' % (row_size, color, color)) else: tikz.append( r' \blockmat{1*\comp}{%s*\comp}{}{draw=white,fill=%s}{}\\' % (row_size, color)) tikz.append(r'}}') vec_tikz = "\n".join(tikz) self._terms.append(vec_tikz) return vec_tikz
[docs] def operator(self, opperator="="): self._process_vars() tikz = [] padding_size = (self._total_size - 1) / 2 tikz.append(r'\blockrow{') tikz.append(r' \blockempty{\mwid}{%s*\comp}{} \\' % (padding_size)) tikz.append( r' \blockmat{\mwid}{1*\comp}{\huge $%s$}{draw=white,fill=white}{}\\' % (opperator)) tikz.append(r' \blockempty{\mwid}{%s*\comp}{} \\' % (padding_size)) tikz.append(r'}') op_tikz = "\n".join(tikz) self._terms.append(op_tikz) return op_tikz
[docs] def spacer(self): self._process_vars() tikz = [] for i in range(self._n_vars): row_size = self._ij_variables[i].size tikz.append(r'\blockrow{\blockcol{') tikz.append( r' \blockmat{.25*\mwid}{%s*\comp}{}{draw=white,fill=white}{}\\' % (row_size)) tikz.append(r'}}') spacer_tikz = "\n".join(tikz) self._terms.append(spacer_tikz) return spacer_tikz
[docs] def write(self, out_file=None, build=True, cleanup=True): """ Write output files for the matrix equation diagram. This produces the following: - {file_name}.tikz A file containing the TIKZ definition of the tikz diagram. - {file_name}.tex A standalone document wrapped around an include of the TIKZ file which can be compiled to a pdf. - {file_name}.pdf An optional compiled version of the standalone tex file. Parameters ---------- file_name : str The prefix to be used for the output files build : bool Flag that determines whether the standalone PDF of the XDSM will be compiled. Default is True. cleanup: bool Flag that determines if padlatex build files will be deleted after build is complete """ tikz = [] tikz.append(r'\blockrow{') for term in self._terms: tikz.append(r'\blockcol{') tikz.append(term) tikz.append(r'}') tikz.append(r'}') eqn_tikz = "\n".join(tikz) if out_file: # with open('{}.tex'.format(out_file), 'w') as f: # f.write(base_file_start) # f.write(eqn_tikz) # f.write(base_file_end) # if build: # os.system('pdflatex {}.tex'.format(out_file)) # if cleanup: # for ext in ['aux', 'fdb_latexmk', 'fls', 'log', 'tex']: # f_name = '{}.{}'.format(out_file, ext) # if os.path.exists(f_name): # os.remove(f_name) _write_tikz(eqn_tikz, out_file, build, cleanup)
if __name__ == "__main__": lst = MatrixEquation() lst.add_variable('x', text=r'$x$') lst.add_variable('y', size=3, text=r'$y$') lst.add_variable('z') lst.connect('x', 'y') lst.connect('y', 'z') lst.text('z', 'x', r'$0$') lst.jacobian(transpose=True) lst.spacer() lst.vector(base_color='green', highlight=[3, 2, 'diag']) lst.operator('=') lst.vector(base_color='red') lst.vector(base_color='red') lst.write('test') J = TotalJacobian() J.add_input('a', text=r'$a$') J.add_input('b', text=r'$b$') J.add_input('c', text=r'$c$') J.add_input('d', text=r'$d$') J.add_input('e', text=r'$e$') J.add_output('gc', text=r'$g_c$') J.add_output('gd', text=r'$g_d$') J.add_output('ge', text=r'$g_e$') J.add_output('f', text=r'$f$') J.connect('a', ('gc', 'gd', 'ge', 'f')) J.connect('b', ('gc', 'gd', 'ge', 'f')) J.connect('c', 'gc') J.connect('d', 'gd') J.connect('e', ('ge', 'f')) J.write('J_test', cleanup=False)