Hans Petter Langtangen 12 anos atrás
pai
commit
fa8540fde2

+ 2 - 2
examples/stochastic_medium.py

@@ -13,9 +13,9 @@ drawing_tool.set_fontsize(24)
 
 layers =      {'layer%d' % i: Line((0,a[i]), (W,a[i]))
                for i in range(len(a))}
-symbols_q =   {'q_%d' % i: Text('$q_%d$' % i, (W/2,0.5*(a[i]+a[i+1])))
+symbols_q =   {'xi_%d' % i: Text(r'$\xi_%d$' % i, (W/2,0.5*(a[i]+a[i+1])))
                for i in range(len(a)-1)}
-symbols_q['q_3'] = Text('$q_2$', (-0.5,a[1]))
+symbols_q['xi_2'] = Text(r'$\xi_2$', (-0.5,a[1]))
 
 sides ={'left': Line((0,0), (0,H)), 'right': Line((W,0), (W,H))}
 d = sides.copy()

examples/layered_medium.py → examples/layered_medium_general.py


+ 21 - 15
examples/mesh_function.py

@@ -1,29 +1,30 @@
 from pysketcher import *
 
+"""
 u = SketchyFunc2('$u(t)$', name_pos='end')
 n = 7
 t_mesh = [i*2.25/(n-1) for i in range(n)]
 
 u = SketchyFunc2('$u(t)$', name_pos='end')
 t_mesh = [0, 2, 4, 6, 8]
+"""
 
 u = SketchyFunc3()
-t_mesh = linspace(0, 6, 8)  # bra
-t_mesh = linspace(0, 6, 6)
+Nt = 5
+t_mesh = linspace(0, 6, Nt+1)
 
-
-t_min1 = t_mesh[0] - 0.1*(t_mesh[-1] - t_mesh[0])
-t_max1 = t_mesh[-1] + 0.2*(t_mesh[-1] - t_mesh[0])
-t_min2 = t_mesh[0] - 0.2*(t_mesh[-1] - t_mesh[0])
-t_max2 = t_mesh[-1] + 0.3*(t_mesh[-1] - t_mesh[0])
+# Add 20% space to the left and 30% to the right of the coordinate system
+t_axis_extent = t_mesh[-1] - t_mesh[0]
+t_min = t_mesh[0] - 0.2*t_axis_extent
+t_max = t_mesh[-1] + 0.3*t_axis_extent
 u_max = 1.3*max([u(t) for t in t_mesh])
 u_min = -0.2*u_max
 
-drawing_tool.set_coordinate_system(t_min2, t_max2, u_min, u_max, axis=False)
+drawing_tool.set_coordinate_system(t_min, t_max, u_min, u_max, axis=False)
 drawing_tool.set_linecolor('black')
 
-r = 0.005*(t_max2-t_min2)     # radius of circles placed at mesh points
-discrete_u = Composition({i: Composition(dict(
+r = 0.005*(t_max-t_min)     # radius of circles placed at mesh points
+u_discrete = Composition({i: Composition(dict(
     circle=Circle(point(t, u(t)), r).set_filled_curves('black'),
     u_point=Text('$u_%d$' % i,
                  point(t, u(t)) + (point(0,5*r)
@@ -35,23 +36,28 @@ interpolant = Composition({
             point(t_mesh[i], u(t_mesh[i]))).set_linewidth(1)
     for i in range(1, len(t_mesh))})
 
-axes = Composition(dict(x=Axis(point(0,0), t_max1, '$t$',
-                               label_spacing=(1/45.,-1/30.)),
-                        y=Axis(point(0,0), 0.8*u_max, '$u$',
-                               rotation_angle=90)))
+axes = Composition(dict(
+    x=Axis(point(0,0), t_mesh[-1] + 0.2*t_axis_extent, '$t$',
+           label_spacing=(1/45.,-1/30.)),
+    y=Axis(point(0,0), 0.8*u_max, '$u$',
+           rotation_angle=90)))
+
 h = 0.03*u_max  # tickmarks height
 nodes = Composition({i: Composition(dict(
     node=Line(point(t,h), point(t,-h)),
     name=Text('$t_%d$' % i, point(t,-3.5*h))))
                      for i, t in enumerate(t_mesh)})
-illustration = Composition(dict(u=discrete_u,
+illustration = Composition(dict(u=u_discrete,
                                 mesh=nodes,
                                 axes=axes)).set_name('fdm_u')
 drawing_tool.erase()
+# Draw t_mesh with discrete u points
 illustration.draw()
 drawing_tool.display()
 drawing_tool.savefig(illustration.get_name())
 
+# Exact u line (u is a Spline Shape that applies 500 intervals by default
+# for drawing the curve)
 exact = u.set_linestyle('dashed').set_linewidth(1)
 exact.draw()
 drawing_tool.display()

+ 82 - 0
examples/staggered_mesh_function.py

@@ -0,0 +1,82 @@
+from pysketcher import *
+
+Nt = 5
+
+#u = SketchyFunc1()
+#t_mesh = linspace(0, 8, Nt+1)
+
+u = SketchyFunc3()
+t_mesh = linspace(0, 6, Nt+1)
+t_mesh_staggered = linspace(0.5*(t_mesh[0]+t_mesh[1]),
+                            0.5*(t_mesh[-2] + t_mesh[-1]), Nt)
+
+# Add 20% space to the left and 30% to the right of the coordinate system
+t_axis_extent = t_mesh[-1] - t_mesh[0]
+t_min = t_mesh[0] - 0.2*t_axis_extent
+t_max = t_mesh[-1] + 0.3*t_axis_extent
+u_max = 1.3*max([u(t) for t in t_mesh])
+u_min = -0.2*u_max
+
+drawing_tool.set_coordinate_system(t_min, t_max, u_min, u_max, axis=False)
+drawing_tool.set_linecolor('black')
+
+r = 0.005*(t_max-t_min)     # radius of circles placed at mesh points
+u_discrete = Composition({i: Composition(dict(
+    circle=Circle(point(t, u(t)), r).set_filled_curves('black'),
+    u_point=Text('$u_%d$' % i,
+                 point(t, u(t)) + (point(0,5*r)
+                                   if i > 0 else point(-5*r,0))),
+    )) for i, t in enumerate(t_mesh)})
+
+# u' = v
+#v = u.smooth.derivative(n=1)
+v = SketchyFunc4()
+
+v_discrete = Composition({i: Composition(dict(
+    circle=Circle(point(t, v(t)), r).set_filled_curves('red'),
+    v_point=Text(r'$v_{%d/2}$' % (2*i+1),
+                 point(t, v(t)) + (point(0,5*r))),
+    )) for i, t in enumerate(t_mesh_staggered)})
+
+axes = Composition(dict(
+    x=Axis(point(0,0), t_mesh[-1] + 0.2*t_axis_extent, '$t$',
+           label_spacing=(1/45.,-1/30.)),
+    y=Axis(point(0,0), 0.8*u_max, '$u,v$',
+           rotation_angle=90)))
+
+h = 0.03*u_max  # tickmarks height
+u_nodes = Composition({i: Composition(dict(
+    node=Line(point(t,h), point(t,-h)),
+    name=Text('$t_%d$' % i, point(t,-3.5*h))))
+                     for i, t in enumerate(t_mesh)})
+v_nodes = Composition({i: Composition(dict(
+    node=Line(point(t,h/1.5), point(t,-h/1.5)).set_linecolor('red'),
+    name=Text(r'$t_{%d/2}$' % (2*i+1), point(t,-3.5*h))))
+                     for i, t in enumerate(t_mesh_staggered)})
+illustration = Composition(dict(u=u_discrete,
+                                v=v_discrete,
+                                u_mesh=u_nodes,
+                                v_mesh=v_nodes,
+                                axes=axes)).set_name('fdm_uv')
+drawing_tool.erase()
+# Staggered t mesh and u and v points
+illustration.draw()
+drawing_tool.display()
+drawing_tool.savefig(illustration.get_name())
+
+# Exact u line (u is a Spline Shape that applies 500 intervals by default
+# for drawing the curve)
+u_exact = u.set_linestyle('dashed').set_linewidth(1)
+u_exact.draw()
+#v = Curve(u.xcoor, v(u.xcoor))
+t_mesh_staggered_fine = linspace(t_mesh_staggered[0],
+                                 t_mesh_staggered[-1],
+                                 501)
+v_exact = Curve(t_mesh_staggered_fine, v(t_mesh_staggered_fine)).\
+          set_linestyle('dashed').set_linewidth(1)
+v_exact.draw()
+drawing_tool.display()
+drawing_tool.savefig('%s_uve' % illustration.get_name())
+
+raw_input()
+

+ 42 - 26
pysketcher/MatplotlibDraw.py

@@ -30,6 +30,11 @@ class MatplotlibDraw:
         self.instruction_file = None
         self.allow_screen_graphics = True  # does not work yet
 
+    def __del__(self):
+        if self.instruction_file:
+            self.instruction_file.write('\nmpl.draw()\nraw_input()\n')
+            self.instruction_file.close()
+
     def ok(self):
         """
         Return True if set_coordinate_system is called and
@@ -183,7 +188,7 @@ ax.set_aspect('equal')
         """
         Fill area inside curves with specified color and/or pattern.
         A common pattern is '/' (45 degree lines). Other patterns
-        include....
+        include '-', '+', 'x', '\\', '*', 'o', 'O', '.'.
         """
         if color is False:
             self.fillcolor = ''
@@ -217,7 +222,7 @@ ax.set_aspect('equal')
                    linestyle=None, linewidth=None,
                    linecolor=None, arrow=None,
                    fillcolor=None, fillpattern=None,
-                   shadow=0):
+                   shadow=0, name=None):
         """Define a curve with coordinates x and y (arrays)."""
         #if not self.allow_screen_graphics:
         #    mpl.ioff()
@@ -241,12 +246,17 @@ ax.set_aspect('equal')
         if shadow == 1:
             shadow = 3   # smallest displacement that is visible
 
+        # We can plot fillcolor/fillpattern, arrow or line
+
         if self.instruction_file:
             import pprint
-            self.instruction_file.write('x = %s\n' % \
-                                        pprint.pformat(self.xdata.tolist()))
-            self.instruction_file.write('y = %s\n' % \
-                                        pprint.pformat(self.ydata.tolist()))
+            if name is not None:
+                self.instruction_file.write('\n# %s\n' % name)
+            if not arrow:
+                self.instruction_file.write(
+                    'x = %s\n' % pprint.pformat(self.xdata.tolist()))
+                self.instruction_file.write(
+                    'y = %s\n' % pprint.pformat(self.ydata.tolist()))
 
 
         if fillcolor or fillpattern:
@@ -257,7 +267,28 @@ ax.set_aspect('equal')
                                   linewidth=linewidth, hatch=fillpattern)
             if self.instruction_file:
                 self.instruction_file.write("[line] = ax.fill(x, y, '%s', edgecolor='%s', linewidth=%d, hatch='%s')\n" % (fillcolor, linecolor, linewidth, fillpattern))
+
+        elif arrow:
+            # Note that a Matplotlib arrow is a line with the arrow tip
+            # (do not draw the line in addition)
+            if not arrow in ('->', '<-', '<->'):
+                raise ValueError("arrow argument must be '->', '<-', or '<->', not %s" % repr(arrow))
+
+            # Add arrow to first and/or last segment
+            start = arrow == '<-' or arrow == '<->'
+            end = arrow == '->' or arrow == '<->'
+            if start:
+                x_s, y_s = x[1], y[1]
+                dx_s, dy_s = x[0]-x[1], y[0]-y[1]
+                self._plot_arrow(x_s, y_s, dx_s, dy_s, '->',
+                                 linestyle, linewidth, linecolor)
+            if end:
+                x_e, y_e = x[-2], y[-2]
+                dx_e, dy_e = x[-1]-x[-2], y[-1]-y[-2]
+                self._plot_arrow(x_e, y_e, dx_e, dy_e, '->',
+                                 linestyle, linewidth, linecolor)
         else:
+            # Plain line
             [line] = self.ax.plot(x, y, linecolor, linewidth=linewidth,
                                   linestyle=linestyle)
             if self.instruction_file:
@@ -291,23 +322,6 @@ self.ax.plot(x, y, linewidth=%d, color='gray',
 """ % linewidth)
 
 
-        if arrow:
-            if not arrow in ('->', '<-', '<->'):
-                raise ValueError("arrow argument must be '->', '<-', or '<->', not %s" % repr(arrow))
-
-            # Add arrow to first and/or last segment
-            start = arrow == '<-' or arrow == '<->'
-            end = arrow == '->' or arrow == '<->'
-            if start:
-                x_s, y_s = x[1], y[1]
-                dx_s, dy_s = x[0]-x[1], y[0]-y[1]
-                self.plot_arrow(x_s, y_s, dx_s, dy_s, '->',
-                                linestyle, linewidth, linecolor)
-            if end:
-                x_e, y_e = x[-2], y[-2]
-                dx_e, dy_e = x[-1]-x[-2], y[-1]-y[-2]
-                self.plot_arrow(x_e, y_e, dx_e, dy_e, '->',
-                                linestyle, linewidth, linecolor)
 
     def display(self, title=None):
         """Display the figure. Last possible command."""
@@ -422,8 +436,8 @@ ax.annotate('%s', xy=%s, xycoords='data',
 #http://matplotlib.sourceforge.net/users/annotations_intro.html
 #http://matplotlib.sourceforge.net/users/annotations_guide.html#plotting-guide-annotation
 
-    def plot_arrow(self, x, y, dx, dy, style='->',
-              linestyle=None, linewidth=None, linecolor=None):
+    def _plot_arrow(self, x, y, dx, dy, style='->',
+                    linestyle=None, linewidth=None, linecolor=None):
         """Draw arrow (dx,dy) at (x,y). `style` is '->', '<-' or '<->'."""
         if linestyle is None:
             # use "global" linestyle
@@ -437,6 +451,7 @@ ax.annotate('%s', xy=%s, xycoords='data',
             self.mpl.arrow(x, y, dx, dy, hold=True,
                            facecolor=linecolor,
                            edgecolor=linecolor,
+                           linestyle=linestyle,
                            linewidth=linewidth,
                            head_width=self.arrow_head_width,
                            #head_width=0.1,
@@ -447,10 +462,11 @@ ax.annotate('%s', xy=%s, xycoords='data',
                 self.instruction_file.write("""\
 mpl.arrow(x=%g, y=%g, dx=%g, dy=%g,
           facecolor='%s', edgecolor='%s',
+          linestyle='%s',
           linewidth=%g, head_width=0.1,
           length_includes_head=True,
           shape='full')
-""" % (x, y, dx, dy, linecolor, linecolor, linewidth))
+""" % (x, y, dx, dy, linecolor, linecolor, linestyle, linewidth))
         if style == '<-' or style == '<->':
             self.mpl.arrow(x+dx, y+dy, -dx, -dy, hold=True,
                            facecolor=linecolor,

+ 44 - 10
pysketcher/shapes.py

@@ -6,11 +6,18 @@ from MatplotlibDraw import MatplotlibDraw
 drawing_tool = MatplotlibDraw()
 
 def point(x, y, check_inside=False):
-    if isinstance(x, (float,int)) and isinstance(y, (float,int)):
-        pass
-    else:
-        raise TypeError('x=%s,y=%s must be float,float, not %s,%s' %
-                        (x, y, type(x), type(y)))
+    for obj, name in zip([x, y], ['x', 'y']):
+        if isinstance(obj, (float,int)):
+            pass
+        elif isinstance(obj, ndarray):
+            if obj.size == 1:
+                pass
+            else:
+                raise TypeError('%s=%s of type %d has length=%d > 1' %
+                                (name, obj, type(obj), obj.size))
+        else:
+            raise TypeError('%s=%s is of wrong type %d' %
+                            (name, obj, type(obj)))
     if check_inside:
         ok, msg = drawing_tool.inside((x,y), exception=True)
         if not ok:
@@ -197,7 +204,8 @@ class Shape:
                 shape_name = shape
                 shape = self.shapes[shape]
             else:
-                shape_name = k
+                shape_name = k  # use index as name if list
+
             if not isinstance(shape, Shape):
                 if isinstance(shape, dict):
                     raise TypeError(
@@ -226,6 +234,8 @@ class Shape:
                         (self.__class__.__name__, shape_name, type(shape),
                          pprint.pformat(self.shapes)))
 
+            if isinstance(shape, Curve):
+                shape.name = shape_name
             getattr(shape, func)(*args, **kwargs)
 
     def draw(self):
@@ -447,6 +457,7 @@ class Curve(Shape):
         self.fillpattern = None
         self.arrow = None
         self.shadow = False
+        self.name = None  # name of object that this Curve represents
 
     def inside_plot_area(self, verbose=True):
         """Check that all coordinates are within drawing_tool's area."""
@@ -484,7 +495,7 @@ class Curve(Shape):
             self.x, self.y,
             self.linestyle, self.linewidth, self.linecolor,
             self.arrow, self.fillcolor, self.fillpattern,
-            self.shadow)
+            self.shadow, self.name)
 
     def rotate(self, angle, center):
         """
@@ -586,12 +597,14 @@ class Curve(Shape):
 
 
 class Spline(Shape):
+    # Note: UnivariateSpline interpolation may not work if
+    # the x[i] points are far from uniformly spaced
     def __init__(self, x, y, degree=3, resolution=501):
         from scipy.interpolate import UnivariateSpline
         self.smooth = UnivariateSpline(x, y, s=0, k=degree)
-        xs = linspace(x[0], x[-1], resolution)
-        ys = self.smooth(xs)
-        self.shapes = {'smooth': Curve(xs, ys)}
+        self.xcoor = linspace(x[0], x[-1], resolution)
+        ycoor = self.smooth(self.xcoor)
+        self.shapes = {'smooth': Curve(self.xcoor, ycoor)}
 
     def geometric_features(self):
         s = self.shapes['smooth']
@@ -602,6 +615,11 @@ class Spline(Shape):
     def __call__(self, x):
         return self.smooth(x)
 
+    # Can easily find the derivative and the integral as
+    # self.smooth.derivative(n=1) and self.smooth.antiderivative()
+
+
+
 
 class SketchyFunc1(Spline):
     """
@@ -630,6 +648,19 @@ class SketchyFunc3(Spline):
         if name is not None:
             self.shapes['name'] = Text(name, self.geometric_features()[name_pos] + point(0,0.1))
 
+class SketchyFunc4(Spline):
+    """
+    A typical function curve used to illustrate an "arbitrary" function.
+    Can be a companion function to SketchyFunc3.
+    """
+    domain = [1, 6]
+    def __init__(self, name=None, name_pos='start'):
+        x = [0, 2,   3,   4, 5,   6]
+        y = [1.5, 1.3, 0.7, 0.5, 0.6, 0.8]
+        Spline.__init__(self, x, y)
+        self.shapes['smooth'].set_linecolor('black')
+        if name is not None:
+            self.shapes['name'] = Text(name, self.geometric_features()[name_pos] + point(0,0.1))
 
 class SketchyFunc2(Shape):
     """
@@ -771,6 +802,7 @@ class Rectangle(Shape):
         upper_right          Upper right corner point.
         lower_mid            Middle point on lower side.
         upper_mid            Middle point on upper side.
+        center               Center point
         ==================== =============================================
         """
         r = self.shapes['rectangle']
@@ -782,6 +814,7 @@ class Rectangle(Shape):
         d['upper_mid'] = 0.5*(d['upper_left'] + d['upper_right'])
         d['left_mid'] = 0.5*(d['lower_left'] + d['upper_left'])
         d['right_mid'] = 0.5*(d['lower_right'] + d['upper_right'])
+        d['center'] = point(d['lower_mid'][0], d['left_mid'][1])
         return d
 
 class Triangle(Shape):
@@ -1169,6 +1202,7 @@ class Arrow1(Shape):
     def __init__(self, start, end, style='->'):
         arrow = Line(start, end)
         arrow.set_arrow(style)
+        # Note:
         self.shapes = {'arrow': arrow}
 
     def geometric_features(self):