A method to plot beam diagram using sympy's own plot()

sympy/physics/continuum_mechanics/beam.py

@@ -15,15 +15,11 @@
 from sympy.plotting import plot, PlotGrid
 from sympy.geometry.entity import GeometryEntity
 from sympy.external import import_module
-from sympy.utilities.decorator import doctest_depends_on
-from sympy import lambdify
+from sympy import lambdify, Add
 from sympy.core.compatibility import iterable
+from sympy.utilities.decorator import doctest_depends_on
 
-matplotlib = import_module('matplotlib', __import__kwargs={'fromlist':['pyplot']})
-numpy = import_module('numpy', __import__kwargs={'fromlist':['linspace']})
-
-__doctest_requires__ = {('Beam.plot_loading_results',): ['matplotlib']}
-
+numpy = import_module('numpy', __import__kwargs={'fromlist':['arange']})
 
 class Beam(object):
     """
@@ -123,6 +119,8 @@ def __init__(self, length, elastic_modulus, second_moment, variable=Symbol('x'),
         self._base_char = base_char
         self._boundary_conditions = {'deflection': [], 'slope': []}
         self._load = 0
+        self._applied_supports = []
+        self._support_as_loads = []
         self._applied_loads = []
         self._reaction_loads = {}
         self._composite_type = None
@@ -366,6 +364,8 @@ def apply_support(self, loc, type="fixed"):
         (-4*SingularityFunction(x, 0, 2) + 3*SingularityFunction(x, 10, 2)
             + 120*SingularityFunction(x, 30, 1) + SingularityFunction(x, 30, 2) + 4000/3)/(E*I)
         """
+        loc = sympify(loc)
+        self._applied_supports.append((loc, type))
         if type == "pin" or type == "roller":
             reaction_load = Symbol('R_'+str(loc))
             self.apply_load(reaction_load, loc, -1)
@@ -377,6 +377,9 @@ def apply_support(self, loc, type="fixed"):
             self.apply_load(reaction_moment, loc, -2)
             self.bc_deflection.append((loc, 0))
             self.bc_slope.append((loc, 0))
+            self._support_as_loads.append((reaction_moment, loc, -2, None))
+
+        self._support_as_loads.append((reaction_load, loc, -1, None))
 
     def apply_load(self, value, start, order, end=None):
         """
@@ -1521,6 +1524,137 @@ def plot_loading_results(self, subs=None):
         return PlotGrid(4, 1, ax1, ax2, ax3, ax4)
 
 
+    @doctest_depends_on(modules=('numpy',))
+    def draw(self, pictorial=False):
+        """Returns a plot object representing the beam diagram of the beam.
+
+        Parameters
+        ==========
+
+        pictorial: Boolean (default=False)
+            Setting ``pictorial=True`` would simply create a pictorial (scaled) view
+            of the beam diagram not according to the scale.
+            Although setting ``pictorial=False`` would create a beam diagram with
+            the exact dimensions on the plot
+
+        Examples
+        ========
+
+        .. plot::
+            :context: close-figs
+            :format: doctest
+            :include-source: True
+
+            >>> from sympy.physics.continuum_mechanics.beam import Beam
+            >>> from sympy import symbols
+            >>> R1, R2 = symbols('R1, R2')
+            >>> E, I = symbols('E, I')
+            >>> b = Beam(50, 20, 30)
+            >>> b.apply_load(10, 2, -1)
+            >>> b.apply_load(R1, 10, -1)
+            >>> b.apply_load(R2, 30, -1)
+            >>> b.apply_load(90, 5, 0, 23)
+            >>> b.apply_load(10, 30, 1, 50)
+            >>> b.apply_support(50, "pin")
+            >>> b.apply_support(0, "fixed")
+            >>> b.apply_support(20, "roller")
+            >>> b.draw(pictorial=True)
+            Plot object containing:
+            [0]: cartesian line: 25*SingularityFunction(x, 5, 0)
+            - 25*SingularityFunction(x, 23, 0) + SingularityFunction(x, 30, 1)
+            - 20*SingularityFunction(x, 50, 0) - SingularityFunction(x, 50, 1)
+            + 5 for x over (0.0, 50.0)
+        """
+        if not numpy:
+            raise ImportError("To use this function numpy module is required")
+
+        x = self.variable
+        length = self.length
+        height = length/10
+
+        rectangles = []
+        rectangles.append({'xy':(0, 0), 'width':length, 'height': height, 'facecolor':"brown"})
+        annotations, markers, load_eq, fill = self._draw_load(pictorial)
+        support_markers, support_rectangles = self._draw_supports()
+
+        rectangles += support_rectangles
+        markers += support_markers
+
+        sing_plot = plot(height + load_eq, (x, 0, length),
+         xlim=(-2, length + 2), ylim=(-length, length), annotations=annotations,
+          markers=markers, rectangles=rectangles, fill=fill, axis=False, show=False)
+
+        return sing_plot
+
+
+    def _draw_load(self, pictorial):
+        loads = list(set(self.applied_loads) - set(self._support_as_loads))
+        length = self.length
+        height = length/10
+        x = self.variable
+
+        annotations = []
+        markers = []
+        scaled_load = 0
+        for load in loads:
+            if load[2] == -1:
+                if load[0].is_positive:
+                    annotations.append({'s':'', 'xy':(load[1], height),  'xytext':(load[1], height*4), 'arrowprops':dict(width= 1.5, headlength=4, headwidth=4, facecolor='black')})
+                else:
+                    annotations.append({'s':'', 'xy':(load[1], 0), 'xytext':(load[1], height - 4*height), 'arrowprops':dict(width= 1.5, headlength=5, headwidth=5, facecolor='black')})
+            elif load[2] == -2:
+                if load[0].is_negative:
+                    markers.append({'args':[[load[1]], [height/2]], 'marker': r'$\circlearrowleft$', 'markersize':15})
+                else:
+                    markers.append({'args':[[load[1]], [height/2]], 'marker': r'$\circlearrowright$', 'markersize':15})
+            elif load[2] >= 0:
+                if pictorial:
+                    value, start, order, end = load
+                    value = 1 if order > 0 else length/2
+                    scaled_load += value*SingularityFunction(x, start, order)
+                    f2 = 1*x**order if order > 0 else length/2*x**order
+                    for i in range(0, order + 1):
+                        scaled_load -= (f2.diff(x, i).subs(x, end - start) *
+                                       SingularityFunction(x, end, i) / factorial(i))
+                    load_args = scaled_load.args
+                else:
+                    load_args = self.load.args
+
+                load_eq = [i for i in load_args if list(i.atoms(SingularityFunction))[0].args[2] >= 0]
+                load_eq = Add(*load_eq)
+
+        # filling higher order loads with colour
+        y = numpy.arange(0, float(length), 0.1)
+        expr = height + load_eq.rewrite(Piecewise)
+        y1 = lambdify(x, expr, 'numpy')
+        y2 = float(height)
+        fill = {'x': y, 'y1': y1(y), 'y2': y2, 'color':'darkkhaki'}
+
+        return annotations, markers, load_eq, fill
+
+
+    def _draw_supports(self):
+        length = self.length
+        height = length/10
+
+        support_markers = []
+        support_rectangles = []
+        for support in self._applied_supports:
+            if support[1] == "pin":
+                support_markers.append({'args':[support[0], [0]], 'marker':6, 'markersize':15, 'color':"black"})
+
+            elif support[1] == "roller":
+                support_markers.append({'args':[support[0], [-height/2]], 'marker':'o', 'markersize':15, 'color':"black"})
+
+            elif support[1] == "fixed":
+                if support[0] == 0:
+                    support_rectangles.append({'xy':(0, -3*height), 'width':-self.length/10, 'height':6*height + height, 'fill':False, 'hatch':'/////'})
+                else:
+                    support_rectangles.append({'xy':(self.length, -3*height), 'width':self.length/10, 'height': 6*height + height, 'fill':False, 'hatch':'/////'})
+
+        return support_markers, support_rectangles
+
+
 class Beam3D(Beam):
     """
     This class handles loads applied in any direction of a 3D space along

sympy/plotting/plot.py

@@ -162,6 +162,10 @@ def __init__(self, *args, **kwargs):
         self.legend = False
         self.autoscale = True
         self.margin = 0
+        self.annotations = None
+        self.markers = None
+        self.rectangles = None
+        self.fill = None
 
         # Contains the data objects to be plotted. The backend should be smart
         # enough to iterate over this list.
@@ -1147,34 +1151,6 @@ def _process_series(self, series, ax, parent):
             ax.set_xscale(parent.xscale)
         if parent.yscale and not isinstance(ax, Axes3D):
             ax.set_yscale(parent.yscale)
-        if parent.xlim:
-            from sympy.core.basic import Basic
-            xlim = parent.xlim
-            if any(isinstance(i, Basic) and not i.is_real for i in xlim):
-                raise ValueError(
-                "All numbers from xlim={} must be real".format(xlim))
-            if any(isinstance(i, Basic) and not i.is_finite for i in xlim):
-                raise ValueError(
-                "All numbers from xlim={} must be finite".format(xlim))
-            xlim = (float(i) for i in xlim)
-            ax.set_xlim(xlim)
-        else:
-            if parent._series and all(isinstance(s, LineOver1DRangeSeries) for s in parent._series):
-                starts = [s.start for s in parent._series]
-                ends = [s.end for s in parent._series]
-                ax.set_xlim(min(starts), max(ends))
-
-        if parent.ylim:
-            from sympy.core.basic import Basic
-            ylim = parent.ylim
-            if any(isinstance(i,Basic) and not i.is_real for i in ylim):
-                raise ValueError(
-                "All numbers from ylim={} must be real".format(ylim))
-            if any(isinstance(i,Basic) and not i.is_finite for i in ylim):
-                raise ValueError(
-                "All numbers from ylim={} must be finite".format(ylim))
-            ylim = (float(i) for i in ylim)
-            ax.set_ylim(ylim)
         if not isinstance(ax, Axes3D) or self.matplotlib.__version__ >= '1.2.0':  # XXX in the distant future remove this check
             ax.set_autoscale_on(parent.autoscale)
         if parent.axis_center:
@@ -1208,6 +1184,54 @@ def _process_series(self, series, ax, parent):
             ax.set_xlabel(parent.xlabel, position=(1, 0))
         if parent.ylabel:
             ax.set_ylabel(parent.ylabel, position=(0, 1))
+        if parent.annotations:
+            for a in parent.annotations:
+                ax.annotate(**a)
+        if parent.markers:
+            for marker in parent.markers:
+                # make a copy of the marker dictionary
+                # so that it doesn't get altered
+                m = marker.copy()
+                args = m.pop('args')
+                ax.plot(*args, **m)
+        if parent.rectangles:
+            for r in parent.rectangles:
+                rect = self.matplotlib.patches.Rectangle(**r)
+                ax.add_patch(rect)
+        if parent.fill:
+            ax.fill_between(**parent.fill)
+
+        # xlim and ylim shoulld always be set at last so that plot limits
+        # doesn't get altered during the process.
+        if parent.xlim:
+            from sympy.core.basic import Basic
+            xlim = parent.xlim
+            if any(isinstance(i, Basic) and not i.is_real for i in xlim):
+                raise ValueError(
+                "All numbers from xlim={} must be real".format(xlim))
+            if any(isinstance(i, Basic) and not i.is_finite for i in xlim):
+                raise ValueError(
+                "All numbers from xlim={} must be finite".format(xlim))
+            xlim = (float(i) for i in xlim)
+            ax.set_xlim(xlim)
+        else:
+            if parent._series and all(isinstance(s, LineOver1DRangeSeries) for s in parent._series):
+                starts = [s.start for s in parent._series]
+                ends = [s.end for s in parent._series]
+                ax.set_xlim(min(starts), max(ends))
+
+        if parent.ylim:
+            from sympy.core.basic import Basic
+            ylim = parent.ylim
+            if any(isinstance(i,Basic) and not i.is_real for i in ylim):
+                raise ValueError(
+                "All numbers from ylim={} must be real".format(ylim))
+            if any(isinstance(i,Basic) and not i.is_finite for i in ylim):
+                raise ValueError(
+                "All numbers from ylim={} must be finite".format(ylim))
+            ylim = (float(i) for i in ylim)
+            ax.set_ylim(ylim)
+
 
     def process_series(self):
         """
@@ -1435,6 +1459,23 @@ def plot(*args, **kwargs):
 
     ``ylim`` : tuple of two floats, denoting the y-axis limits.
 
+    ``annotations``: list. A list of dictionaries specifying the type of
+    annotation required. The keys in the dictionary should be equivalent
+    to the arguments of the matplotlib's annotate() function.
+
+    ``markers``: list. A list of dictionaries specifying the type the
+    markers required. The keys in the dictionary should be equivalent
+    to the arguments of the matplotlib's plot() function along with the
+    marker related keyworded arguments.
+
+    ``rectangles``: list. A list of dictionaries specifying the dimensions
+    of the rectangles to be plotted. The keys in the dictionary should be
+    equivalent to the arguments of the matplotlib's patches.Rectangle class.
+
+    ``fill``: dict. A dictionary specifying the type of color filling
+    required in the plot. The keys in the dictionary should be equivalent
+    to the arguments of the matplotlib's fill_between() function.
+
     Examples
     ========