.. DO NOT EDIT.
.. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY.
.. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE:
.. "auto_examples/plot_streaming_raker.py"
.. LINE NUMBERS ARE GIVEN BELOW.

.. only:: html

    .. note::
        :class: sphx-glr-download-link-note

        :ref:`Go to the end <sphx_glr_download_auto_examples_plot_streaming_raker.py>`
        to download the full example code.

.. rst-class:: sphx-glr-example-title

.. _sphx_glr_auto_examples_plot_streaming_raker.py:


Raking a biased stream in fixed memory
======================================

Rows arrive from a sample that has too many women, young people and college
graduates, and an outcome that depends on all three. ``StreamingRaker``
weights each row as it arrives and keeps only running totals; ``BatchIPF``
rakes the full stored sample at the end, which is the benchmark.

.. GENERATED FROM PYTHON SOURCE LINES 12-14

A biased stream
---------------

.. GENERATED FROM PYTHON SOURCE LINES 14-33

.. code-block:: Python

    import matplotlib.pyplot as plt
    import numpy as np

    from onlinerake import BatchIPF, StreamingRaker, Targets

    targets = Targets(female=0.5, young=0.3, college=0.35)
    names = targets.feature_names
    population = np.array([targets[name] for name in names])
    sample_shares = {"college": 0.45, "female": 0.65, "young": 0.45}

    rng = np.random.default_rng(0)
    n = 20_000
    rows = np.column_stack([rng.random(n) < sample_shares[name] for name in names]).astype(
        int
    )
    coef = {"college": 2.0, "female": 1.0, "young": -1.5}
    y = rows @ np.array([coef[name] for name in names]) + rng.normal(0, 1, n)
    true_mean = float(population @ np.array([coef[name] for name in names]))








.. GENERATED FROM PYTHON SOURCE LINES 34-36

Weight the stream as it arrives
-------------------------------

.. GENERATED FROM PYTHON SOURCE LINES 36-61

.. code-block:: Python

    raker = StreamingRaker(targets, outcomes=["y"])
    checkpoints = np.unique(np.geomspace(10, n, 40).astype(int))
    stream_error, raw_error = [], []
    for i, (row, yi) in enumerate(zip(rows, y, strict=True), start=1):
        raker.partial_fit(dict(zip(names, row.tolist(), strict=True)), outcomes={"y": yi})
        if i in checkpoints:
            margins = np.array([raker.margins[name] for name in names])
            stream_error.append(np.abs(margins - population).max())
            raw_error.append(np.abs(rows[:i].mean(axis=0) - population).max())

    ipf = BatchIPF(targets).fit([dict(zip(names, r.tolist(), strict=True)) for r in rows])
    w = ipf.weights

    print(f"true mean of y: {true_mean:.3f}")
    print(f"{'':<16} {'estimate':>8} {'ESS':>7} {'numbers kept':>13}")
    print(f"{'unweighted':<16} {y.mean():8.3f} {n:7d} {'':>13}")
    print(
        f"{'StreamingRaker':<16} {raker.estimate('y'):8.3f} "
        f"{raker.effective_sample_size:7.0f} {len(names) + 5:13d}"
    )
    print(
        f"{'BatchIPF':<16} {w @ y / w.sum():8.3f} "
        f"{w.sum() ** 2 / (w @ w):7.0f} {rows.size + n:13d}"
    )





.. rst-class:: sphx-glr-script-out

 .. code-block:: none

    true mean of y: 0.750
                     estimate     ESS  numbers kept
    unweighted          0.851   20000              
    StreamingRaker      0.744   14903             8
    BatchIPF            0.745   16011         80000




.. GENERATED FROM PYTHON SOURCE LINES 62-64

Margin error as the stream grows
--------------------------------

.. GENERATED FROM PYTHON SOURCE LINES 64-75

.. code-block:: Python

    fig, ax = plt.subplots(figsize=(7, 3.5))
    ax.plot(checkpoints, raw_error, label="unweighted")
    ax.plot(checkpoints, stream_error, label="StreamingRaker")
    ax.set_xscale("log")
    ax.set_yscale("log")
    ax.set_xlabel("rows seen")
    ax.set_ylabel("largest margin error")
    ax.legend()
    fig.tight_layout()
    plt.show()




.. image-sg:: /auto_examples/images/sphx_glr_plot_streaming_raker_001.png
   :alt: plot streaming raker
   :srcset: /auto_examples/images/sphx_glr_plot_streaming_raker_001.png
   :class: sphx-glr-single-img





.. GENERATED FROM PYTHON SOURCE LINES 76-81

Design weights and trimming
---------------------------
Rows can arrive with design weights in any units. Heavy-tailed design
weights make the estimate noisy; ``trim`` caps each weight at a multiple of
the mean weight so far, trading a little bias for a larger effective sample.

.. GENERATED FROM PYTHON SOURCE LINES 81-94

.. code-block:: Python

    design = 5000 * rng.lognormal(0.0, 1.5, n)
    for trim in (None, 3.0):
        weighted = StreamingRaker(targets, trim=trim, outcomes=["y"])
        for row, yi, d in zip(rows, y, design, strict=True):
            weighted.partial_fit(
                dict(zip(names, row.tolist(), strict=True)),
                design_weight=float(d),
                outcomes={"y": yi},
            )
        print(
            f"trim={trim}: estimate {weighted.estimate('y'):.3f}, "
            f"ESS {weighted.effective_sample_size:.0f}"
        )




.. rst-class:: sphx-glr-script-out

 .. code-block:: none

    trim=None: estimate 0.757, ESS 5166
    trim=3.0: estimate 0.754, ESS 9447





.. rst-class:: sphx-glr-timing

   **Total running time of the script:** (0 minutes 1.855 seconds)


.. _sphx_glr_download_auto_examples_plot_streaming_raker.py:

.. only:: html

  .. container:: sphx-glr-footer sphx-glr-footer-example

    .. container:: sphx-glr-download sphx-glr-download-jupyter

      :download:`Download Jupyter notebook: plot_streaming_raker.ipynb <plot_streaming_raker.ipynb>`

    .. container:: sphx-glr-download sphx-glr-download-python

      :download:`Download Python source code: plot_streaming_raker.py <plot_streaming_raker.py>`

    .. container:: sphx-glr-download sphx-glr-download-zip

      :download:`Download zipped: plot_streaming_raker.zip <plot_streaming_raker.zip>`


.. only:: html

 .. rst-class:: sphx-glr-signature

    `Gallery generated by Sphinx-Gallery <https://sphinx-gallery.github.io>`_