line_series.py 5.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173
  1. from __future__ import annotations
  2. from typing import TYPE_CHECKING, Any, Iterable
  3. import wandb
  4. from wandb.plot.custom_chart import plot_table
  5. if TYPE_CHECKING:
  6. from wandb.plot.custom_chart import CustomChart
  7. def line_series(
  8. xs: Iterable[Iterable[Any]] | Iterable[Any],
  9. ys: Iterable[Iterable[Any]],
  10. keys: Iterable[str] | None = None,
  11. title: str = "",
  12. xname: str = "x",
  13. split_table: bool = False,
  14. ) -> CustomChart:
  15. """Constructs a line series chart.
  16. Args:
  17. xs: Sequence of x values. If a singular
  18. array is provided, all y values are plotted against that x array. If
  19. an array of arrays is provided, each y value is plotted against the
  20. corresponding x array.
  21. ys: Sequence of y values, where each iterable represents
  22. a separate line series.
  23. keys: Sequence of keys for labeling each line series. If
  24. not provided, keys will be automatically generated as "line_1",
  25. "line_2", etc.
  26. title: Title of the chart.
  27. xname: Label for the x-axis.
  28. split_table: Whether the table should be split into a separate section
  29. in the W&B UI. If `True`, the table will be displayed in a section named
  30. "Custom Chart Tables". Default is `False`.
  31. Returns:
  32. CustomChart: A custom chart object that can be logged to W&B. To log the
  33. chart, pass it to `wandb.log()`.
  34. Examples:
  35. Logging a single x array where all y series are plotted against the same x values:
  36. ```python
  37. import wandb
  38. # Initialize W&B run
  39. with wandb.init(project="line_series_example") as run:
  40. # x values shared across all y series
  41. xs = list(range(10))
  42. # Multiple y series to plot
  43. ys = [
  44. [i for i in range(10)], # y = x
  45. [i**2 for i in range(10)], # y = x^2
  46. [i**3 for i in range(10)], # y = x^3
  47. ]
  48. # Generate and log the line series chart
  49. line_series_chart = wandb.plot.line_series(
  50. xs,
  51. ys,
  52. title="title",
  53. xname="step",
  54. )
  55. run.log({"line-series-single-x": line_series_chart})
  56. ```
  57. In this example, a single `xs` series (shared x-values) is used for all
  58. `ys` series. This results in each y-series being plotted against the
  59. same x-values (0-9).
  60. Logging multiple x arrays where each y series is plotted against its corresponding x array:
  61. ```python
  62. import wandb
  63. # Initialize W&B run
  64. with wandb.init(project="line_series_example") as run:
  65. # Separate x values for each y series
  66. xs = [
  67. [i for i in range(10)], # x for first series
  68. [2 * i for i in range(10)], # x for second series (stretched)
  69. [3 * i for i in range(10)], # x for third series (stretched more)
  70. ]
  71. # Corresponding y series
  72. ys = [
  73. [i for i in range(10)], # y = x
  74. [i**2 for i in range(10)], # y = x^2
  75. [i**3 for i in range(10)], # y = x^3
  76. ]
  77. # Generate and log the line series chart
  78. line_series_chart = wandb.plot.line_series(
  79. xs, ys, title="Multiple X Arrays Example", xname="Step"
  80. )
  81. run.log({"line-series-multiple-x": line_series_chart})
  82. ```
  83. In this example, each y series is plotted against its own unique x series.
  84. This allows for more flexibility when the x values are not uniform across
  85. the data series.
  86. Customizing line labels using `keys`:
  87. ```python
  88. import wandb
  89. # Initialize W&B run
  90. with wandb.init(project="line_series_example") as run:
  91. xs = list(range(10)) # Single x array
  92. ys = [
  93. [i for i in range(10)], # y = x
  94. [i**2 for i in range(10)], # y = x^2
  95. [i**3 for i in range(10)], # y = x^3
  96. ]
  97. # Custom labels for each line
  98. keys = ["Linear", "Quadratic", "Cubic"]
  99. # Generate and log the line series chart
  100. line_series_chart = wandb.plot.line_series(
  101. xs,
  102. ys,
  103. keys=keys, # Custom keys (line labels)
  104. title="Custom Line Labels Example",
  105. xname="Step",
  106. )
  107. run.log({"line-series-custom-keys": line_series_chart})
  108. ```
  109. This example shows how to provide custom labels for the lines using
  110. the `keys` argument. The keys will appear in the legend as "Linear",
  111. "Quadratic", and "Cubic".
  112. """
  113. # If xs is a single array, repeat it for each y in ys
  114. if not isinstance(xs[0], Iterable) or isinstance(xs[0], (str, bytes)):
  115. xs = [xs] * len(ys)
  116. if len(xs) != len(ys):
  117. msg = f"Number of x-series ({len(xs)}) must match y-series ({len(ys)})."
  118. raise ValueError(msg)
  119. if keys is None:
  120. keys = [f"line_{i}" for i in range(len(ys))]
  121. if len(keys) != len(ys):
  122. msg = f"Number of keys ({len(keys)}) must match y-series ({len(ys)})."
  123. raise ValueError(msg)
  124. data = [
  125. [x, keys[i], y]
  126. for i, (xx, yy) in enumerate(zip(xs, ys))
  127. for x, y in zip(xx, yy)
  128. ]
  129. table = wandb.Table(
  130. data=data,
  131. columns=["step", "lineKey", "lineVal"],
  132. )
  133. return plot_table(
  134. data_table=table,
  135. vega_spec_name="wandb/lineseries/v0",
  136. fields={
  137. "step": "step",
  138. "lineKey": "lineKey",
  139. "lineVal": "lineVal",
  140. },
  141. string_fields={"title": title, "xname": xname},
  142. split_table=split_table,
  143. )