.. testsetup:: *

   import numpy as np
   from libertem import api
   from libertem.executor.inline import InlineJobExecutor

   from libertem.udf import UDF as RealUDF

   # We override UDF in such a way that it can be used
   # without implementing all methods
   class UDF(RealUDF):
      def get_result_buffers(self):
         return {}

      def process_frame(self, frame):
         pass

   class YourUDF(RealUDF):
      def get_result_buffers(self):
         return {'buf1': self.buffer(kind="nav")}

      def process_frame(self, frame):
         self.results.buf1[:] = 42

   ctx = api.Context(executor=InlineJobExecutor())
   data = np.random.random((16, 16, 32, 32)).astype(np.float32)
   dataset = ctx.load("memory", data=data, sig_dims=2)
   roi = np.random.choice([True, False], dataset.shape.nav)
   udf = YourUDF()


.. _`implement udf`:

Implementing a UDF
------------------

The workflow for implementing a UDF starts with subclassing
:class:`~libertem.udf.base.UDF`. In the simplest case, you need to implement the
:meth:`~libertem.udf.base.UDF.get_result_buffers` method and 
:meth:`~libertem.udf.base.UDFFrameMixin.process_frame`.

There are two very common patterns for reductions, reducing over the navigation axes
into a common accumulator for all frames, keeping the shape of a single frame,
or reducing over the signal axes and keeping the navigation axes.

Declaring buffers
~~~~~~~~~~~~~~~~~

A UDF can implement one of these reductions or combinations. To handle indexing
for you, LiberTEM needs to know about the structure of your reduction. You can
build this structure in the :meth:`~libertem.udf.base.UDF.get_result_buffers`
method, by declaring one or more buffers.

These buffers can have a :code:`kind` declared, which corresponds to the two
reduction patterns above: :code:`kind="sig"` for reducing over the navigation
axes (and keeping the signal axes), and :code:`kind="nav"` for reducing over the
signal axes and keeping the navigation axes. There is a third,
:code:`kind="single"`, which allows to declare buffers with custom shapes that
don't correspond directly to the data set's shape.

It is also possible to append additional axes to the buffer's shape using the
:code:`extra_shape` parameter.

:meth:`~libertem.udf.base.UDF.get_result_buffers` should return a :code:`dict` which maps
buffer names to buffer declarations. You can create a buffer declaration by calling
the :meth:`~libertem.udf.base.UDF.buffer` method.

The buffer name is later used to access the buffer via :code:`self.results.<buffername>`,
which returns a view into a NumPy array. For this to work, the name has to be a valid Python
identifier.

Examples of buffer declarations:

.. testcode:: getresultbuffers

   def get_result_buffers(self):
      # Suppose our dataset has the shape (14, 14, 32, 32),
      # where the first two dimensions represent the navigation
      # dimension and the last two dimensions represent the signal
      # dimension.

      buffers = {
         # same shape as navigation dimensions of dataset, plus two
         # extra dimensions of shape (3, 2). The full shape is
         # (14, 14, 3, 2) in this example. This means this buffer can
         # store an array of shape (3, 2) for each frame in the dataset.
         "nav_buffer": self.buffer(
               kind="nav",
               extra_shape=(3, 2),
               dtype="float32",
         ),

         # same shape as signal dimensions of dataset, plus an extra
         # dimension of shape (2,). Consequently, the full shape is
         # (32, 32, 2) in this example. That means we can store two
         # float32 values for each pixel of the signal dimensions.
         "sig_buffer": self.buffer(
               kind="sig",
               extra_shape=(2,),
               dtype="float32",
         ),

         # buffer of shape (16, 16); shape is unrelated to dataset shape
         "single_buffer": self.buffer(
               kind="single",
               extra_shape=(16, 16),
               dtype="float32",
         ),

      }

      return buffers

.. testcleanup:: getresultbuffers

   class TestUDF(UDF):
      def merge(self, dest, src):
         pass

   TestUDF.get_result_buffers = get_result_buffers

   u = TestUDF()
   ctx.run_udf(dataset=dataset, udf=u)

See below for some more real-world examples.

All NumPy dtypes are supported for buffers. That includes the :code:`object`
dtype for arbitrary Python variables. The item just has to be pickleable with
:code:`cloudpickle`.

Note that buffers are only designed to pass lightweight intermediate results
and thus, it is important that the size of the buffer remains small. Having too
large buffers can lead to significant decline in performance.

Implementing the processing function
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

Now to the actual core of the processing: implementing
:meth:`~libertem.udf.UDFFrameMixin.process_frame`. The method signature looks like this:

.. testcode:: processing

   class ExampleUDF(UDF):
      def process_frame(self, frame):
         ...

.. testcleanup:: processing

   u = ExampleUDF()
   ctx.run_udf(dataset=dataset, udf=u)

The general idea is that you get a single frame from the data set, do your processing,
and write the results to one of the previously declared buffers, via :code:`self.results.<buffername>`.
When accessing a :code:`kind="nav"` buffer this way, you automatically get a view into the buffer
that corresponds to the current frame that is being processed. In case of :code:`kind="sig"`
or :code:`kind="single"`, you get the whole buffer.

Intuitively, with :code:`kind="sig"` (and :code:`kind="single"`), you are most
likely implementing an operation like :code:`buf = f(buf, frame)`. That is, you
are computing a new result based on a single (new) frame and the results from all
previous frames, and overwrite the results with the new value(s).

With :code:`kind="nav"`, you compute independent results for each frame,
which are written to different positions in the result buffer. Because of the independence
between frames, you don't need to merge with a previous value; the result is simply written
to the correct index in the result buffer (via the aforementioned view).

As an easy example, let's have a look at a function that simply sums up each frame
to a single value. This is a :code:`kind="nav"` reduction, as we sum over all values
in the signal dimensions:

.. _`sumsig`:
   
.. testcode:: sumsig

   import numpy as np   
   from libertem.udf import UDF


   class SumOverSig(UDF):
      def get_result_buffers(self):
         """
         Describe the buffers we need to store our results:
         kind="nav" means we want to have a value for each coordinate
         in the navigation dimensions. We name our buffer 'pixelsum'.
         """
         return {
            'pixelsum': self.buffer(
               kind="nav", dtype="float32"
            )
         }

      def process_frame(self, frame):
         """
         Sum up all pixels in this frame and store the result in the
         `pixelsum` buffer. `self.results.pixelsum` is a view into the
         result buffer we defined above, and corresponds to the entry
         for the current frame we work on. We don't have to take care
         of finding the correct index for the frame we are processing
         ourselves.
         """
         self.results.pixelsum[:] = np.sum(frame)

   res = ctx.run_udf(
      udf=SumOverSig(),
      dataset=dataset,
   )

   # to access the named buffer as a NumPy array:
   res['pixelsum'].data

On a 4D data set, this operation is roughly equivalent to :code:`np.sum(arr, axis=(2, 3))`.

Merging partial results
~~~~~~~~~~~~~~~~~~~~~~~

As :ref:`described above <how UDFs work>`, data from multiple partitions is
processed in parallel. That also means that we need a way of merging partial
results into the final result. In the example above, we didn't need to do anything:
we only have a :code:`kind="nav"` buffer, where merging just means assigning the
result of one partition to the right slice in the final result. This is done by
the default implementation of :meth:`~libertem.udf.base.UDF.merge`. 

In case of :code:`kind="sig"` buffers and the corresponding reduction, assignment would
just overwrite the result from the previous partition with the one from the current partition,
and is not the correct operation. So let's have a look at the merge method:

.. testcode:: merge

   class ExampleUDF(UDF):
      def merge(self, dest, src):
         pass

.. testcleanup:: merge

   u = ExampleUDF()
   ctx.run_udf(dataset=dataset, udf=u)

:code:`dest` is the result of all previous merge calls, and :code:`src` is the
result from a single new partition. Your :code:`merge` implementation should read from both
:code:`dest` and :code:`src` and write the result back to :code:`dest`.

Here is an example demonstrating :code:`kind="sig"` buffers and the :code:`merge` function:
   
.. testcode:: realmerge

   import numpy as np
   from libertem.udf import UDF


   class MaxUDF(UDF):
      def get_result_buffers(self):
         """
         Describe the buffers we need to store our results:
         kind="sig" means we want to have a value for each coordinate
         in the signal dimensions (i.e. a value for each pixel of the
         diffraction patterns). We name our buffer 'maxbuf'.
         """
         return {
            'maxbuf': self.buffer(
               kind="sig", dtype=self.meta.input_dtype
            )
         }

      def preprocess(self):
         """
         Initialize buffer with neutral element for maximum.
         """
         self.results.maxbuf[:] = -np.inf

      def process_frame(self, frame):
         """
         In this function, we have a frame and the buffer `maxbuf`
         available, which we declared above. This function is called
         for all frames / diffraction patterns in the data set. The
         maxbuf is a partial result, and all partial results will
         later be merged (see below).

         In this case, we determine the maximum from the current
         maximum and the current frame, for each pixel in the
         diffraction pattern.

         Notes:

         - You cannot rely on any particular order of frames this function
           is called in.
         - Your function should be pure, that is, it should not have side
           effects beyond modifying the content of result buffers or task data,
           and should only depend on it's input parameters, including
           the UDF object :code:`self`.
         """
         self.results.maxbuf[:] = np.maximum(frame, self.results.maxbuf)

      def merge(self, dest, src):
         """
         merge two partial results, from src into dest
         """
         dest.maxbuf[:] = np.maximum(dest.maxbuf, src.maxbuf)

   res = ctx.run_udf(
      udf=MaxUDF(),
      dataset=dataset,
   )

   # to access the named buffer as a NumPy array:
   res['maxbuf'].data


For more complete examples, you can also have a look at the functions
implemented in the sub-modules of :code:`libertem.udf` and at our :ref:`packages`.

Passing parameters
~~~~~~~~~~~~~~~~~~

By default, keyword arguments that are passed to the constructor of a UDF are
available as properties of :code:`self.params`. This allows clean handling and passing
of parameters in distributed execution scenarios, see below.

.. testsetup:: params

   def correlate_peaks(frame, peaks):
      pass

   peaks = None

.. testcode:: params

    class MyUDF(UDF):

        def process_frame(self, frame):
            result = correlate_peaks(frame, self.params.peaks)
            ...

    udf = MyUDF(peaks=peaks, other=...)

.. testcleanup:: params

   def get_result_buffers():
      return {}

   udf.get_result_buffers = get_result_buffers

   ctx.run_udf(dataset=dataset, udf=udf)

Declaring a constructor
~~~~~~~~~~~~~~~~~~~~~~~

In order to document parameters of an UDF and avoid typos with parameter names,
it is good practice to define a constructor for any UDF that will be re-used and
shared. Since functions of an UDF are executed both on the main node and on
worker processes, UDFs will be pickled and sent over the network in the process.
To avoid transferring unwanted state or unnecessary or unpicklable member
variables, a clean copy of an UDF is created before transfer. This clean copy is
created by re-instantiating the UDF class with the parameters that were passed
to the constructor of the UDF base class.

That means a user-defined constructor has to fulfill two conditions:

1. It has to pass any parameters to the superclass constructor.
2. It has to accept exactly the parameters that it passed to the superclass
   constructor whenever a clean copy is created and behave the same way as in the
   original call.

That means modification of parameters other than assigning values for default parameters
can lead to complications and should be avoided, if possible.

.. testcode:: constructor

    class MyParamUDF(UDF):
        '''
        This UDF demonstrates how to define a constructor for a UDF.

        Parameters
        ----------
        my_param
            A parameter that is passed into the UDF for demonstration purposes.
            It is mirrored in the `demo` result buffer.
        '''
        def __init__(self, my_param=None):
            # Assigning a default parameter works,
            # provided the UDF accepts it in subsequent calls
            # exactly as it passed it to the superclass:
            if my_param is None:
                my_param = "Eins, zwei drei!"
            
            # !!!DON'T!!! self.my_param = my_param
            # !!!DON'T!!! super().__init__()
            # We have to pass it to the superclass constructor instead.
            # This makes sure it will be available via self.params and
            # on clean copies.

            # !!!DON'T!!! super().__init__(other_param=my_param)
            # This would trigger a TypeError when
            # MyParamUDF(other_param=my_param) is called for a clean copy.

            # DO pass all parameters to the superclass exactly as this
            # class would accept them
            super().__init__(my_param=my_param)

        # The rest of this UDF just mirrors the parameter
        # back in a result buffer for demonstration.

        def get_result_buffers(self):
            '''
            Declare a result buffer for the parameter
            '''
            return {
                'demo': self.buffer(
                    kind='single',
                    dtype=object
                ),
            }

        def process_frame(self, frame):
            '''
            We assign the parameter to the result buffer
            '''
            self.results.demo[:] = self.params.my_param
            
        def merge(self, dest, src):
            '''
            We pass through the result buffer
            '''
            dest.demo[:] = src.demo


    res = ctx.run_udf(dataset=dataset, udf=MyParamUDF())
    assert res['demo'].data[0] == "Eins, zwei drei!"

    res2 = ctx.run_udf(dataset=dataset, udf=MyParamUDF(my_param=42))
    assert res2['demo'].data[0] == 42


Initializing result buffers
~~~~~~~~~~~~~~~~~~~~~~~~~~~

To allow a UDF to initialize a result buffer to the correct values,
the method :meth:`~libertem.udf.base.UDFPreprocessMixin.preprocess`
can be implemented. It is run once per partition and assigning to
:code:`kind="nav"` result buffers will assign to the results of the
whole partition. See :code:`MaxUDF` above for an example.

.. _`progress bar`:

Running UDFs
------------

As shown in the examples above, the :meth:`~libertem.api.Context.run_udf` method
of :class:`~libertem.api.Context` is used to run UDFs. Usually, you only need to
pass an instance of your UDF and the dataset you want to run on:
   
.. testcode:: run

    udf = YourUDF(param1="value1")
    res = ctx.run_udf(udf=udf, dataset=dataset)

You can also enable a progress bar:

.. testcode:: run

    udf = YourUDF(param1="value1")
    res = ctx.run_udf(udf=udf, dataset=dataset, progress=True)

:meth:`~libertem.api.Context.run_udf` returns a :code:`dict`, having the buffer
names as keys (as defined in :meth:`~libertem.udf.UDF.get_result_buffers`) and
:class:`~libertem.common.buffers.BufferWrapper` instances as values. You
can use these in any place you would use a NumPy array, for example as an argument to
NumPy functions, or you can explicitly convert them to NumPy arrays by accessing
the :code:`.data` attribute, or by calling :meth:`numpy.array`:
   
.. testcode:: run

   import numpy as np

   res = ctx.run_udf(udf=udf, dataset=dataset)
   # convert to NumPy array, assuming we declared a buffer
   # with name `buf1`:
   arr = res['buf1'].data
   arr = np.array(res['buf1'])

   # or directly treat as array:
   np.sum(res['buf1'])

If you want to run multiple independent UDFs on a single :code:`DataSet`,
you can pass in a list of UDFs to :meth:`~libertem.api.Context.run_udf`. This can be faster
than making two passes over the whole :code:`DataSet`:

.. testcode:: run

   from libertem.udf.sum import SumUDF

   # results are returned as a list of dictionaries, so we can unpack them here
   # into `res` and `res_sum`:
   res, res_sum = ctx.run_udf(udf=[udf, SumUDF()], dataset=dataset)


Note the return type: if you pass a list of UDFs, you also get a list of
dictionaries as a result, with the ordering of results matching the ordering of
UDFs in the :meth:`~libertem.api.Context.run_udf` call.


.. _`udf cancellation`:

Cancellation
~~~~~~~~~~~~

.. versionadded:: 0.12.0

Especially when running UDFs on live data streams, it can happen that a UDF
doesn't run to completion. An example scenario is when a continuously running
preview scan is stopped, or scan parameters are changed while the scan is
running.

In this case, the exception
:class:`~libertem.common.exceptions.UDFRunCancelled` is raised, which you can
handle in an appropriate way for your application:

.. testcode:: run

   from libertem.common.exceptions import UDFRunCancelled

   try:
       res = ctx.run_udf(udf=udf, dataset=dataset)
   except UDFRunCancelled:
       # handle cancellation, for example, start the next UDF run
       pass


.. _`udf roi`:

Regions of interest
~~~~~~~~~~~~~~~~~~~

In addition, you can pass the :code:`roi` (region of interest) parameter, to
run your UDF on a selected subset of data. :code:`roi` can be specified in several ways:

- a single coordinate in the navigation dimensions, e.g. :code:`(3, 4)`
  which will restrict the analysis to this frame only
- a sequence of coordinates and their associated truth value, e.g.
  :code:`[((0, 0), True), ((2, 3), True)]`. In this format all truth values 
  must be identical, and the roi background value will be set to its opposite.
  This allows quick specification of a small number of pixels to include or exclude
  from the analysis.
- a NumPy array containing a bool mask, having the shape of the navigation
  axes of the dataset.
- any :code:`scipy.sparse` matrix or array or :code:`pydata.sparse.SparseArray` format,
  again having the same shape as the naivgation axes of the dataset

For example, to process a random subset of a 4D-STEM dataset:
   
.. testcode:: run

   import numpy as np

   # If your dataset has shape `(14, 14, 32, 32)` with two signal
   # and two navigation dimensions, `dataset.shape.nav`
   # translates to `(14, 14)`.
   roi = np.random.choice(a=[False, True], size=dataset.shape.nav)
   ctx.run_udf(udf=udf, dataset=dataset, roi=roi)

Note that the result array only contains values for the selected indices, all
other indices are set to :code:`nan` (or, if the dtype doesn't support nan,
some other, not further defined value). It is best to limit further processing
to the same :code:`roi`.

You can also access a flat array that is not filled up with :code:`nan` using
:code:`.raw_data`:

.. testcode:: run

   res = ctx.run_udf(udf=udf, dataset=dataset, roi=roi)
   res['buf1'].raw_data

.. _plotting:

Live Plotting
-------------

.. versionadded:: 0.7.0

LiberTEM can display a live plot of the UDF results. In the most simple case,
this can be done by setting :code:`plots=True` in
:meth:`~libertem.api.Context.run_udf`. 

.. testsetup:: live

    from libertem.udf.sum import SumUDF
    udf = SumUDF()

.. testcode:: live

    ctx.run_udf(dataset=dataset, udf=udf, plots=True)

See the following items for a full demonstration, including setting up fully
customized plots. The API reference can be found in :ref:`viz reference`.

.. toctree::

    liveplotting

.. _partial:

Partial results
---------------

.. versionadded:: 0.7.0

Instead of only getting the whole result after the UDF has finished running, you
can also use :meth:`~libertem.api.Context.run_udf_iter` to get a generator for
partial results:

.. testsetup:: partial

    from libertem.udf.sum import SumUDF


.. testcode:: partial

    for udf_results in ctx.run_udf_iter(dataset=dataset, udf=SumUDF()):
        # ... do something interesting with `udf_results`:
        a = np.sum(udf_results.buffers[0]['intensity'])

    # after the loop, `udf_results` contains the final results as usual

While the UDF execution is running, the UDF object should not be modified since
that leads to undefined behavior. In particular, nested or concurrent execution
of the same UDF objects must be avoided since it modifies the buffers that are
allocated internally while a UDF is running.

Asynchronous execution
----------------------

It is also possible to integrate LiberTEM into an async script or application by
passing :code:`sync=False` to :meth:`~libertem.api.Context.run_udf_iter` or
:meth:`~libertem.api.Context.run_udf`:

.. Not run with docs-check since we can't easily test async code there...

.. code-block:: python

    async for udf_results in ctx.run_udf_iter(dataset=dataset, udf=udf, sync=False):
        # ... do something interesting with `udf_results`:
        a = np.sum(udf_results[0]['intensity'])

    # or the version without intermediate results:
    udf_results = await ctx.run_udf(dataset=dataset, udf=udf, sync=False)

See the items below for a more comprehensive demonstration and documentation:

.. toctree::

    async