Add plot functionality (0.4.4)

.gitignore

@@ -1,4 +1,5 @@
 a.py
+*.png
 
 # Created by https://www.gitignore.io/api/vim,python
 # Edit at https://www.gitignore.io/?templates=vim,python

docs/TODO.md

@@ -3,14 +3,13 @@
 ## Plotting
 * Add back plotting functionality
     - Constraints
-    - Plot per beam
-* Save plots
 * Create matrix plot (see 0.3.2)
 * Compute and plot centre of mass, and mass of individual beams
 * Plot inertia loads --> `F_accel`
 * Plot accel vector
-* Plot beam indices
 * Plot local axes
+* Scale for deformations, vector size, ...
+* Flag to chose if loads plotted in deformed or undeformed mesh
 
 ## Model definition
 * Add acceleration state (define on beam level, or 'global'?) --> inertia loads

src/framat/__version__.py

@@ -1,2 +1,2 @@
-VERSION = (0, 4, 3)
+VERSION = (0, 4, 4)
 __version__ = '.'.join(map(str, VERSION))

src/framat/_meshing.py

@@ -290,7 +290,7 @@ def __init__(self):
         self._num_nodes = -1  # Number of nodes -1
 
         # Named nodes per beam (maps [beam_idx][node_uid] --> point_coord)
-        self.named_nodes = defaultdict(dict)
+        self.named_nodes = defaultdict(OrderedDict)
 
         # Global node numbers of named nodes (maps [node_uid] --> global number)
         self.glob_nums = {}
@@ -439,3 +439,19 @@ def gnv(self, vector, uid):
             :uid: UID of named node
         """
         return vector[self.glob_nums[uid]]
+
+    def gbv(self, vector, beam_idx):
+        """
+        Return the beam value from a given vector (e.g. displacement or load)
+
+        Args:
+            :vector: vector
+            :beam_idx: (int) beam index
+        """
+        uids = list(self.named_nodes[beam_idx].keys())
+        uid_first = uids[0]
+        uid_last = uids[-1]
+        idx1 = self.glob_nums[uid_first]
+        idx2 = self.glob_nums[uid_last]
+        # breakpoint()
+        return vector[idx1:idx2+1]

src/framat/_model.py

@@ -23,6 +23,9 @@
 Model definition
 """
 
+from pathlib import Path
+import os
+
 from mframework import FeatureSpec, ModelSpec, SchemadictValidators
 import numpy as np
 
@@ -31,10 +34,19 @@
 from ._util import Schemas as S
 from ._plot import PlotItems
 
+
+def is_dir(key, value, comp_value, _):
+    """Validator to check if the value is a writeable directory"""
+    # Try to create directory if non-existent
+    Path(value).mkdir(parents=True, exist_ok=True)
+    if not os.access(value, os.W_OK):
+        raise ValueError(f"{value!r} is not a valid directory")
+
+
 # Register custom 'schemadict' types
 SchemadictValidators.register_type(AbstractBeamMesh)
 SchemadictValidators.register_type(np.ndarray)
-
+SchemadictValidators[str]['is_dir'] = is_dir
 
 # =================
 # ===== MODEL =====
@@ -274,11 +286,11 @@
 fspec.add_prop_spec(
     'plot_settings',
     {
-        'show': {'type', bool},
+        'show': {'type': bool},
         'linewidth': S.pos_number,
         'markersize': S.pos_number,
         'fontsize': S.pos_int,
-        # 'save': {'type': str, 'check_dir': ...}  # TODO
+        'save': {'type': str, 'is_dir': 'dummy'}
     },
     doc="Define general plot settings."
 )

src/framat/_plot.py

@@ -23,16 +23,24 @@
 Plotting
 """
 
+from datetime import datetime
+from math import ceil
+from random import randint
+import os
+
 from commonlibs.math.vectors import unit_vector
 from mpl_toolkits.mplot3d import Axes3D
 import matplotlib.cm as cm
 import matplotlib.pyplot as plt
 import numpy as np
 
+from . import MODULE_NAME
 from ._element import GlobalSystem
+from ._log import logger
 
 
 class PlotItems:
+    beam_index = 'beam_index'
     deformed = 'deformed'
     forces = 'forces'
     global_axes = 'global_axes'
@@ -62,7 +70,7 @@ class C:
     LOCAL_SYS = 'green'
     MASS = 'maroon'
     MASS_LOAD = 'maroon'
-    UNDEFORMED = 'blue'
+    UNDEFORMED = 'grey'
 
 
 def plot_all(m):
@@ -74,15 +82,29 @@ def plot_all(m):
     if not mpp.get('plot', ()):
         return
 
+    ps = m.get('post_proc').get('plot_settings', {})
     abm = m.results.get('mesh').get('abm')
 
+    num_tot = m.get('post_proc').len('plot')
     for plot_num, _ in enumerate(m.get('post_proc').iter('plot')):
+        logger.info(f"Creating plot {plot_num + 1}/{num_tot}...")
         ax = init_3D_plot(*abm.get_lims())
-        add_deformed_undeformed(m, ax, plot_num)
+        add_items_per_beam(m, ax, plot_num)
         add_global_axes(m, ax, plot_num)
+        plt.tight_layout()
+
+        if ps.get('save', False):
+            now = datetime.now().strftime("%F_%H%M%S")
+            ext = 'png'
+            rand = randint(100, 999)
+            fname = f"{MODULE_NAME.lower()}_{now}_{plot_num+1:02}_{rand}.{ext}"
+            fname = os.path.join(os.path.abspath(ps.get('save')), fname)
+            logger.info(f"Saving plot to file {fname!r}...")
+            plt.savefig(fname, dpi=300, format='png')
+
+    if ps.get('show', False):
+        plt.show()
 
-    plt.tight_layout()
-    plt.show()
     plt.close('all')
 
 
@@ -197,48 +219,52 @@ def _coordinate_system(plot, origin, axes, axes_names, color, scale=1):
     plot.plot_surface(xx, yy, z, alpha=0.4, color=color)
 
 
-def add_deformed_undeformed(m, ax, plot_num):
+def add_items_per_beam(m, ax, plot_num):
     to_show = m.get('post_proc').get('plot')[plot_num]
     abm = m.results.get('mesh').get('abm')
     marker = 'o' if 'nodes' in to_show else None
 
-    for beam_id in abm.beams.keys():
-        xyz = abm.get_all_points(beam_id)
+    for beam_idx in abm.beams.keys():
+        xyz = abm.get_all_points(beam_idx)
         x, y, z = xyz[:, 0], xyz[:, 1], xyz[:, 2]
 
+        # ----- Undeformed mesh -----
         if PlotItems.undeformed in to_show:
-            ax.plot(x, y, z, **args_plot(m, C.UNDEFORMED, marker=marker))
+            ax.plot(x, y, z, **args_plot(m, C.UNDEFORMED))
 
+        # ----- Deformed mesh -----
         if PlotItems.deformed in to_show:
-            # ==================
-            # TODO: pick range for beam only
             d = m.results.get('tensors').get('comp:U')
-            xd = x + d['ux']
-            yd = y + d['uy']
-            zd = z + d['uz']
+            xd = x + abm.gbv(d['ux'], beam_idx)
+            yd = y + abm.gbv(d['uy'], beam_idx)
+            zd = z + abm.gbv(d['uz'], beam_idx)
             ax.plot(xd, yd, zd, **args_plot(m, C.DEFORMED, marker=marker))
-            # ==================
 
+        # ----- Forces -----
         if PlotItems.forces in to_show:
-            # ==================
             d = m.results.get('tensors').get('comp:F')
-            Fx = d['Fx']
-            Fy = d['Fy']
-            Fz = d['Fz']
-            ax.quiver(x, y, z, Fx, Fy, Fz, color=C.CONC_FORCE)
-            # ax.quiver(xd, yd, zd, Fx, Fy, Fz, color=C.CONC_FORCE)
-            # ==================
+            Fx = abm.gbv(d['Fx'], beam_idx)
+            Fy = abm.gbv(d['Fy'], beam_idx)
+            Fz = abm.gbv(d['Fz'], beam_idx)
+            # ax.quiver(x, y, z, Fx, Fy, Fz, color=C.CONC_FORCE)
+            ax.quiver(xd, yd, zd, Fx, Fy, Fz, color=C.CONC_FORCE)
 
+        # ----- Moments -----
         if PlotItems.moments in to_show:
-            # ==================
             d = m.results.get('tensors').get('comp:F')
-            Mx = d['Mx']
-            My = d['My']
-            Mz = d['Mz']
-            ax.quiver(x, y, z, Mx, My, Mz, color=C.CONC_FORCE)
-            # ==================
-
-        # TODO: separate func
-        if 'node_uids' in to_show:
-            for uid, coord in abm.named_nodes[beam_id].items():
+            Fx = abm.gbv(d['Mx'], beam_idx)
+            Fy = abm.gbv(d['My'], beam_idx)
+            Fz = abm.gbv(d['Mz'], beam_idx)
+            # ax.quiver(x, y, z, Fx, Fy, Fz, color=C.CONC_FORCE)
+            ax.quiver(xd, yd, zd, Fx, Fy, Fz, color=C.CONC_FORCE)
+
+        # ----- Beam index -----
+        if PlotItems.beam_index in to_show:
+            center = ceil(len(x)/2)
+            coord = (x[center], y[center], z[center])
+            ax.text(*coord, str(beam_idx), **args_text(m, color=C.BC))
+
+        # ----- Named nodes -----
+        if PlotItems.node_uids in to_show:
+            for uid, coord in abm.named_nodes[beam_idx].items():
                 ax.text(*coord, uid, **args_text(m, color=C.BC))

tests/integration/test_plotting.py

@@ -0,0 +1,49 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+import pytest
+
+from framat import Model
+from framat._plot import PlotItems
+
+REL_TOL = 1e-4
+
+
+def get_cantilever_model():
+    model = Model()
+
+    mat = model.add_feature('material', uid='dummy')
+    mat.set('E', 1)
+    mat.set('G', 1)
+    mat.set('rho', 1)
+
+    cs = model.add_feature('cross_section', uid='dummy')
+    cs.set('A', 1)
+    cs.set('Iy', 1)
+    cs.set('Iz', 1)
+    cs.set('J', 1)
+
+    beam = model.add_feature('beam')
+    beam.add('node', [0, 0, 0], uid='root')
+    beam.add('node', [1, 0, 0], uid='tip')
+    beam.set('nelem', 10)
+    beam.add('material', {'from': 'root', 'to': 'tip', 'uid': 'dummy'})
+    beam.add('cross_section', {'from': 'root', 'to': 'tip', 'uid': 'dummy'})
+    beam.add('orientation', {'from': 'root', 'to': 'tip', 'up': [0, 0, 1]})
+    beam.add('point_load', {'at': 'tip', 'load': [0, 0, -1, 0, 0, 0]})
+
+    model.set_feature('bc').add('fix', {'node': 'root', 'fix': ['all']})
+    return model
+
+
+def test_plotting():
+    model = get_cantilever_model()
+
+    pp = model.set_feature('post_proc')
+    pp.set('plot_settings', {'save': '.', 'show': False})
+    pp.add('plot', ['deformed', 'undeformed'])
+    pp.add('plot', PlotItems.to_list())
+    model.run()
+
+    # TODO: before test --> remove all plot files
+    # TODO: after test --> assert that files created