6
6
__credits__ = ["Yannick Copin" ]
7
7
__license__ = "BSD"
8
8
__version__ = "2011/09/16"
9
+
9
10
# Original version by Yannick Copin (ycopin@ipnl.in2p3.fr) 10/2/2010, available
10
11
# at:
11
12
# http://matplotlib.org/examples/api/sankey_demo_old.html
@@ -113,7 +114,7 @@ def _arc(self, quadrant=0, cw=True, radius=1, center=(0, 0)):
113
114
ARC_VERTICES [:, 1 ]))
114
115
if quadrant > 1 :
115
116
radius = - radius # Rotate 180 deg.
116
- return zip (ARC_CODES , radius * vertices +
117
+ return zip (ARC_CODES , radius * vertices +
117
118
np .tile (center , (ARC_VERTICES .shape [0 ], 1 )))
118
119
119
120
def _add_input (self , path , angle , flow , length ):
@@ -131,7 +132,7 @@ def _add_input(self, path, angle, flow, length):
131
132
path .extend ([(Path .LINETO , [x , y ]),
132
133
(Path .LINETO , dip ),
133
134
(Path .LINETO , [x , y + flow ]),
134
- (Path .LINETO , [x + self .gap , y + flow ])])
135
+ (Path .LINETO , [x + self .gap , y + flow ])])
135
136
label_location = [dip [0 ] - self .offset , dip [1 ]]
136
137
else : # Vertical
137
138
x -= self .gap
@@ -149,7 +150,7 @@ def _add_input(self, path, angle, flow, length):
149
150
# Inner arc isn't needed if inner radius is zero
150
151
if self .radius :
151
152
path .extend (self ._arc (quadrant = quadrant ,
152
- cw = angle == UP ,
153
+ cw = angle == UP ,
153
154
radius = self .radius ,
154
155
center = (x + self .radius ,
155
156
y - sign * self .radius )))
@@ -159,7 +160,7 @@ def _add_input(self, path, angle, flow, length):
159
160
(Path .LINETO , dip ),
160
161
(Path .LINETO , [x - flow , y - sign * length ])])
161
162
path .extend (self ._arc (quadrant = quadrant ,
162
- cw = angle == DOWN ,
163
+ cw = angle == DOWN ,
163
164
radius = flow + self .radius ,
164
165
center = (x + self .radius ,
165
166
y - sign * self .radius )))
@@ -187,7 +188,7 @@ def _add_output(self, path, angle, flow, length):
187
188
(Path .LINETO , tip ),
188
189
(Path .LINETO , [x , y - self .shoulder + flow ]),
189
190
(Path .LINETO , [x , y + flow ]),
190
- (Path .LINETO , [x - self .gap , y + flow ])])
191
+ (Path .LINETO , [x - self .gap , y + flow ])])
191
192
label_location = [tip [0 ] + self .offset , tip [1 ]]
192
193
else : # Vertical
193
194
x += self .gap
@@ -204,10 +205,10 @@ def _add_output(self, path, angle, flow, length):
204
205
# Inner arc isn't needed if inner radius is zero
205
206
if self .radius :
206
207
path .extend (self ._arc (quadrant = quadrant ,
207
- cw = angle == UP ,
208
+ cw = angle == UP ,
208
209
radius = self .radius ,
209
210
center = (x - self .radius ,
210
- y + sign * self .radius )))
211
+ y + sign * self .radius )))
211
212
else :
212
213
path .append ((Path .LINETO , [x , y ]))
213
214
path .extend ([(Path .LINETO , [x , y + sign * length ]),
@@ -218,7 +219,7 @@ def _add_output(self, path, angle, flow, length):
218
219
y + sign * length ]),
219
220
(Path .LINETO , [x - flow , y + sign * length ])])
220
221
path .extend (self ._arc (quadrant = quadrant ,
221
- cw = angle == DOWN ,
222
+ cw = angle == DOWN ,
222
223
radius = self .radius - flow ,
223
224
center = (x - self .radius ,
224
225
y + sign * self .radius )))
@@ -333,7 +334,7 @@ def add(self, patchlabel='', flows=None, orientations=None, labels='',
333
334
else :
334
335
flows = np .array (flows )
335
336
n = flows .shape [0 ] # Number of flows
336
- if rotation == None :
337
+ if rotation is None :
337
338
rotation = 0
338
339
else :
339
340
# In the code below, angles are expressed in deg/90
@@ -352,15 +353,15 @@ def add(self, patchlabel='', flows=None, orientations=None, labels='',
352
353
"same length.\n labels has length %d, but flows has length %d."
353
354
% (len (labels ), n ))
354
355
else :
355
- labels = [labels ]* n
356
+ labels = [labels ] * n
356
357
assert trunklength >= 0 , (
357
358
"trunklength is negative.\n This isn't allowed, because it would "
358
359
"cause poor layout." )
359
360
if np .absolute (np .sum (flows )) > self .tolerance :
360
361
verbose .report (
361
362
"The sum of the flows is nonzero (%f).\n Is the "
362
363
"system not at steady state?" % np .sum (flows ), 'helpful' )
363
- scaled_flows = self .scale * flows
364
+ scaled_flows = self .scale * flows
364
365
gain = sum (max (flow , 0 ) for flow in scaled_flows )
365
366
loss = sum (min (flow , 0 ) for flow in scaled_flows )
366
367
if not (0.5 <= gain <= 2.0 ):
@@ -400,7 +401,7 @@ def add(self, patchlabel='', flows=None, orientations=None, labels='',
400
401
"within the tolerance (%f)." % (flow_error , self .tolerance ))
401
402
402
403
# Determine if the flows are inputs.
403
- are_inputs = [None ]* n
404
+ are_inputs = [None ] * n
404
405
for i , flow in enumerate (flows ):
405
406
if flow >= self .tolerance :
406
407
are_inputs [i ] = True
@@ -414,12 +415,12 @@ def add(self, patchlabel='', flows=None, orientations=None, labels='',
414
415
% (i , flow , self .tolerance ), 'helpful' )
415
416
416
417
# Determine the angles of the arrows (before rotation).
417
- angles = [None ]* n
418
+ angles = [None ] * n
418
419
for i , (orient , is_input ) in enumerate (zip (orientations , are_inputs )):
419
420
if orient == 1 :
420
421
if is_input :
421
422
angles [i ] = DOWN
422
- elif is_input == False :
423
+ elif not is_input :
423
424
# Be specific since is_input can be None.
424
425
angles [i ] = UP
425
426
elif orient == 0 :
@@ -431,7 +432,7 @@ def add(self, patchlabel='', flows=None, orientations=None, labels='',
431
432
"but it must be -1, 0, or 1." % (i , orient ))
432
433
if is_input :
433
434
angles [i ] = UP
434
- elif is_input == False :
435
+ elif not is_input :
435
436
angles [i ] = DOWN
436
437
437
438
# Justify the lengths of the paths.
@@ -462,10 +463,10 @@ def add(self, patchlabel='', flows=None, orientations=None, labels='',
462
463
for i , (angle , is_input , flow ) in enumerate (reversed (zip (
463
464
angles , are_inputs , scaled_flows ))):
464
465
if angle == UP and is_input :
465
- pathlengths [n - i - 1 ] = lllength
466
+ pathlengths [n - i - 1 ] = lllength
466
467
lllength += flow
467
468
elif angle == DOWN and not is_input :
468
- pathlengths [n - i - 1 ] = lrlength
469
+ pathlengths [n - i - 1 ] = lrlength
469
470
lrlength -= flow
470
471
# Determine the lengths of the left-side arrows
471
472
# from the bottom upwards.
@@ -475,7 +476,7 @@ def add(self, patchlabel='', flows=None, orientations=None, labels='',
475
476
if angle == RIGHT :
476
477
if is_input :
477
478
if has_left_input :
478
- pathlengths [n - i - 1 ] = 0
479
+ pathlengths [n - i - 1 ] = 0
479
480
else :
480
481
has_left_input = True
481
482
# Determine the lengths of the right-side arrows
@@ -537,11 +538,13 @@ def add(self, patchlabel='', flows=None, orientations=None, labels='',
537
538
for i , (angle , is_input , spec ) in enumerate (reversed (zip (
538
539
angles , are_inputs , zip (scaled_flows , pathlengths )))):
539
540
if angle == UP and is_input :
540
- tips [n - i - 1 , :], label_locations [n - i - 1 , :] = self ._add_input (
541
- llpath , angle , * spec )
541
+ tip , label_location = self ._add_input (llpath , angle , * spec )
542
+ tips [n - i - 1 , :] = tip
543
+ label_locations [n - i - 1 , :] = label_location
542
544
elif angle == DOWN and not is_input :
543
- tips [n - i - 1 , :], label_locations [n - i - 1 , :] = self ._add_output (
544
- lrpath , angle , * spec )
545
+ tip , label_location = self ._add_output (lrpath , angle , * spec )
546
+ tips [n - i - 1 , :] = tip
547
+ label_locations [n - i - 1 , :] = label_location
545
548
# Add the left-side inputs from the bottom upwards.
546
549
has_left_input = False
547
550
for i , (angle , is_input , spec ) in enumerate (reversed (zip (
@@ -554,8 +557,9 @@ def add(self, patchlabel='', flows=None, orientations=None, labels='',
554
557
llpath .append ((Path .LINETO , [ulpath [- 1 ][1 ][0 ],
555
558
llpath [- 1 ][1 ][1 ]]))
556
559
has_left_input = True
557
- tips [n - i - 1 , :], label_locations [n - i - 1 , :] = self ._add_input (
558
- llpath , angle , * spec )
560
+ tip , label_location = self ._add_input (llpath , angle , * spec )
561
+ tips [n - i - 1 , :], = tip
562
+ label_locations [n - i - 1 , :] = label_location
559
563
# Add the right-side outputs from the top downwards.
560
564
has_right_output = False
561
565
for i , (angle , is_input , spec ) in enumerate (zip (
@@ -595,7 +599,7 @@ def _get_angle(a, r):
595
599
if prior is None :
596
600
if rotation != 0 : # By default, none of this is needed.
597
601
angles = [_get_angle (angle , rotation ) for angle in angles ]
598
- rotate = Affine2D ().rotate_deg (rotation * 90 ).transform_point
602
+ rotate = Affine2D ().rotate_deg (rotation * 90 ).transform_point
599
603
tips = rotate (tips )
600
604
label_locations = rotate (label_locations )
601
605
vertices = rotate (vertices )
@@ -604,7 +608,7 @@ def _get_angle(a, r):
604
608
rotation = (self .diagrams [prior ].angles [connect [0 ]] -
605
609
angles [connect [1 ]])
606
610
angles = [_get_angle (angle , rotation ) for angle in angles ]
607
- rotate = Affine2D ().rotate_deg (rotation * 90 ).transform_point
611
+ rotate = Affine2D ().rotate_deg (rotation * 90 ).transform_point
608
612
tips = rotate (tips )
609
613
offset = self .diagrams [prior ].tips [connect [0 ]] - tips [connect [1 ]]
610
614
translate = Affine2D ().translate (* offset ).transform_point
0 commit comments