@@ -845,4 +845,124 @@ def finish(self):
845
845
self .ax .set_aspect ('equal' , adjustable = 'datalim' )
846
846
return self .diagrams
847
847
848
+ def __init__ (self , ax = None , scale = 1.0 , unit = '' , format = '%G ' , gap = 0.25 ,
849
+ radius = 0.1 , shoulder = 0.03 , offset = 0.15 , head_angle = 100 ,
850
+ margin = 0.4 , tolerance = 1e-6 , ** kwargs ):
851
+ """
852
+ Create a new Sankey instance.
853
+
854
+ Optional keyword arguments:
855
+
856
+ =============== ===================================================
857
+ Field Description
858
+ =============== ===================================================
859
+ *ax* axes onto which the data should be plotted
860
+ If *ax* is not provided, new axes will be created.
861
+ *scale* scaling factor for the flows
862
+ *scale* sizes the width of the paths in order to
863
+ maintain proper layout. The same scale is applied
864
+ to all subdiagrams. The value should be chosen
865
+ such that the product of the scale and the sum of
866
+ the inputs is approximately 1.0 (and the product of
867
+ the scale and the sum of the outputs is
868
+ approximately -1.0).
869
+ *unit* string representing the physical unit associated
870
+ with the flow quantities
871
+ If *unit* is *None*, then none of the quantities
872
+ are labeled.
873
+ *format* a Python number formatting string to be used in
874
+ labeling the flow as a quantity (i.e., a number
875
+ times a unit, where the unit is given)
876
+ *gap* space between paths that break in/break away
877
+ to/from the top or bottom
878
+ *radius* inner radius of the vertical paths
879
+ *shoulder* size of the shoulders of output arrows
880
+ *offset* text offset (from the dip or tip of the arrow)
881
+ *head_angle* angle of the arrow heads (and negative of the angle
882
+ of the tails) [deg]
883
+ *margin* minimum space between Sankey outlines and the edge
884
+ of the plot area
885
+ *tolerance* acceptable maximum of the magnitude of the sum of
886
+ flows
887
+ The magnitude of the sum of connected flows cannot
888
+ be greater than *tolerance*.
889
+ =============== ===================================================
890
+
891
+ The optional arguments listed above are applied to all subdiagrams so
892
+ that there is consistent alignment and formatting.
893
+
894
+ If :class:`Sankey` is instantiated with any keyword arguments other
895
+ than those explicitly listed above (``**kwargs``), they will be passed
896
+ to :meth:`add`, which will create the first subdiagram.
897
+
898
+ In order to draw a complex Sankey diagram, create an instance of
899
+ :class:`Sankey` by calling it without any kwargs::
900
+
901
+ sankey = Sankey()
902
+
903
+ Then add simple Sankey sub-diagrams::
904
+
905
+ sankey.add() # 1
906
+ sankey.add() # 2
907
+ #...
908
+ sankey.add() # n
909
+
910
+ Finally, create the full diagram::
848
911
912
+ sankey.finish()
913
+
914
+ Or, instead, simply daisy-chain those calls::
915
+
916
+ Sankey().add().add... .add().finish()
917
+
918
+ .. seealso::
919
+
920
+ :meth:`add`
921
+ :meth:`finish`
922
+
923
+
924
+ **Examples:**
925
+
926
+ .. plot:: mpl_examples/api/sankey_demo_basics.py
927
+ """
928
+ # Check the arguments.
929
+ assert gap >= 0 , (
930
+ "The gap is negative.\n This isn't allowed because it "
931
+ "would cause the paths to overlap." )
932
+ assert radius <= gap , (
933
+ "The inner radius is greater than the path spacing.\n "
934
+ "This isn't allowed because it would cause the paths to overlap." )
935
+ assert head_angle >= 0 , (
936
+ "The angle is negative.\n This isn't allowed "
937
+ "because it would cause inputs to look like "
938
+ "outputs and vice versa." )
939
+ assert tolerance >= 0 , (
940
+ "The tolerance is negative.\n It must be a magnitude." )
941
+
942
+ # Create axes if necessary.
943
+ if ax is None :
944
+ import matplotlib .pyplot as plt
945
+ fig = plt .figure ()
946
+ ax = fig .add_subplot (1 , 1 , 1 , xticks = [], yticks = [])
947
+
948
+ self .diagrams = []
949
+
950
+ # Store the inputs.
951
+ self .ax = ax
952
+ self .unit = unit
953
+ self .format = format
954
+ self .scale = scale
955
+ self .gap = gap
956
+ self .radius = radius
957
+ self .shoulder = shoulder
958
+ self .offset = offset
959
+ self .margin = margin
960
+ self .pitch = np .tan (np .pi * (1 - head_angle / 180.0 ) / 2.0 )
961
+ self .tolerance = tolerance
962
+
963
+ # Initialize the vertices of tight box around the diagram(s).
964
+ self .extent = np .array ((np .inf , - np .inf , np .inf , - np .inf ))
965
+
966
+ # If there are any kwargs, create the first subdiagram.
967
+ if len (kwargs ):
968
+ self .add (** kwargs )
0 commit comments