Figure.histogram: Add the parameter 'weights' for weighting data

pygmt/src/histogram.py

@@ -8,11 +8,13 @@
 from pygmt._typing import PathLike, TableLike
 from pygmt.alias import Alias, AliasSystem
 from pygmt.clib import Session
-from pygmt.exceptions import GMTParameterError
+from pygmt.exceptions import GMTParameterError, GMTTypeError
 from pygmt.helpers import (
     build_arg_list,
+    data_kind,
     deprecate_parameter,
     fmt_docstring,
+    is_nonstr_iter,
     kwargs_to_strings,
     use_alias,
 )
@@ -39,6 +41,7 @@
 def histogram(
     self,
     data: PathLike | TableLike,
+    weights: bool | Sequence[float] = False,
     bar_width: float | str | None = None,
     bar_offset: float | str | None = None,
     cmap: str | bool = False,
@@ -87,6 +90,14 @@ def histogram(
     data
         Pass in either a file name to an ASCII data table, a Python list, a 2-D
         $table_classes.
+    weights
+        Weight the data instead of counting them. Default is ``False``, i.e., pure
+        counts are used]. It can be:
+
+        - ``True``: Weights are provided in the second column of ``data``, if ``data``
+          is a file name or a 2-D sequence.
+        - A 1-D array of weights, one per data point, requiring that ``data`` is a 1-D
+          sequence of values.
     $cmap
     pen
         Draw bar outline (or stair-case curve) using the specified pen thickness
@@ -140,7 +151,7 @@ def histogram(
         [*min*\ /*max*\ /]\ *inc*\ [**+n**\ ].
         Set the interval for the width of each bar in the histogram.
     histtype : int or str
-        [*type*][**+w**].
+        [*type*].
         Choose between 6 types of histograms:
 
         * 0 = counts [Default]
@@ -150,8 +161,7 @@ def histogram(
         * 4 = log10 (1.0 + count)
         * 5 = log10 (1.0 + frequency_percent).
 
-        To use weights provided as a second data column instead of pure counts,
-        append **+w**.
+        To use weights instead of pure counts, use the ``weights`` parameter.
     $projection
     $region
     $frame
@@ -172,6 +182,16 @@ def histogram(
             required="bar_width", reason="Required when 'bar_offset' is set."
         )
 
+    # weights can be given as a 1-D array, or as a boolean to indicate that the second
+    # column of data contains weights. If weights is an array, then data must be a 1-D
+    # sequence of values.
+    _weight_is_array = is_nonstr_iter(weights)
+    if data_kind(data) == "file" and _weight_is_array:
+        raise GMTTypeError(
+            type(weights),
+            reason="'weights' must be boolean when 'data' is a file name.",
+        )
+
     aliasdict = AliasSystem(
         A=Alias(horizontal, name="horizontal"),
         C=Alias(cmap, name="cmap"),
@@ -199,10 +219,13 @@ def histogram(
         t=transparency,
     )
     aliasdict.merge(kwargs)
+    if weights is not False:
+        aliasdict["Z"] = f"{aliasdict.get('Z', '')}+w"
 
     self._activate_figure()
     with Session() as lib:
-        with lib.virtualfile_in(check_kind="vector", data=data) as vintbl:
+        vfargs = {"x": data, "y": weights} if _weight_is_array else {"data": data}
+        with lib.virtualfile_in(check_kind="vector", **vfargs) as vintbl:
             lib.call_module(
                 module="histogram", args=build_arg_list(aliasdict, infile=vintbl)
             )

pygmt/tests/test_histogram.py

@@ -2,10 +2,12 @@
 Test Figure.histogram.
 """
 
+import numpy as np
 import pandas as pd
 import pytest
 from pygmt import Figure
-from pygmt.exceptions import GMTParameterError
+from pygmt.exceptions import GMTParameterError, GMTTypeError
+from pygmt.helpers import GMTTempFile
 from pygmt.params import Axis
 
 
@@ -18,6 +20,35 @@ def fixture_data(request):
     return request.param(data)
 
 
+@pytest.fixture(scope="module", name="weights")
+def fixture_weights():
+    """
+    Return a list of weights to be used in the histogram.
+    """
+    return [
+        0.1,
+        0.1,
+        0.2,
+        0.2,
+        0.3,
+        0.3,
+        0.4,
+        0.4,
+        0.5,
+        0.5,
+        0.6,
+        0.6,
+        0.7,
+        0.7,
+        0.8,
+        0.8,
+        0.9,
+        0.9,
+        1.0,
+        1.0,
+    ]
+
+
 @pytest.mark.benchmark
 @pytest.mark.mpl_image_compare(filename="test_histogram.png")
 def test_histogram(data):
@@ -45,9 +76,74 @@ def test_histogram_baroffset(data):
         fig.histogram(
             data=data,
             projection="X10c/10c",
-            region=[0, 9, 0, 6],
+            region=[0, 10, 0, 6],
             series=1,
             frame=Axis(annot=True),
             fill="green",
             bar_offset=0.25,
         )
+
+
+@pytest.mark.mpl_image_compare(filename="test_histogram_weights.png")
+def test_histogram_weights_file(data, weights):
+    """
+    Test weights given in the second column of a data file.
+    """
+    kwargs = {
+        "series": 1,
+        "region": [0, 10, 0, 6],
+        "projection": "X8c/5c",
+        "frame": Axis(annot=True),
+        "fill": "lightblue",
+    }
+    with GMTTempFile() as tmpfile:
+        np.savetxt(
+            tmpfile.name, np.column_stack([data, weights]), header="data,weights"
+        )
+        fig = Figure()
+        fig.histogram(data=tmpfile.name, weights=True, **kwargs)
+        return fig
+
+
+@pytest.mark.mpl_image_compare(filename="test_histogram_weights.png")
+def test_histogram_weights_2darray(data, weights):
+    """
+    Test weights given in the second column of a 2-D sequence.
+    """
+    kwargs = {
+        "series": 1,
+        "region": [0, 10, 0, 6],
+        "projection": "X8c/5c",
+        "frame": Axis(annot=True),
+        "fill": "lightblue",
+    }
+    _data = np.column_stack([data, weights])
+    fig = Figure()
+    fig.histogram(data=_data, weights=True, **kwargs)
+    return fig
+
+
+@pytest.mark.mpl_image_compare(filename="test_histogram_weights.png")
+def test_histogram_weights_array(data, weights):
+    """
+    Test weights given as a 1-D array alongside 1-D data.
+    """
+    kwargs = {
+        "series": 1,
+        "region": [0, 10, 0, 6],
+        "projection": "X8c/5c",
+        "frame": Axis(annot=True),
+        "fill": "lightblue",
+    }
+    fig = Figure()
+    fig.histogram(data=data, weights=weights, **kwargs)
+    return fig
+
+
+def test_histogram_weights_array_invalid_data():
+    """
+    Test that passing weights as an array requires data to be a 1-D sequence.
+    """
+    fig = Figure()
+    with pytest.raises(GMTTypeError):
+        fig.histogram(data="input.txt", series=1, weights=[0.5, 1.0, 2.0])