Add context manager to optimization limits

Author: harripd

H2MM_C/H2MM_C.pyx

@@ -16,6 +16,7 @@ from libc.stdint cimport int64_t, int32_t, uint8_t, INT32_MAX
 from libc.string cimport memcpy
 from cpython cimport array
 cimport numpy as cnp
+import contextlib
 
 import os
 import numpy as np
@@ -446,6 +447,8 @@ cdef class opt_lim_const:
     The optimization_limits module variable should be of this class
     
     """
+    _attrs = ('max_iter', 'max_time', 'converged_min', 'num_cores', 'squarem', 
+              'formatter', 'outstream')
     cdef:
         int64_t _max_iter
         double _max_time
@@ -454,6 +457,7 @@ cdef class opt_lim_const:
         bint _squarem
         object _formatter
         object _outstream
+    
     def __cinit__(self, max_iter=2046, max_time=np.inf, converged_min=1e-14, 
                   num_cores=os.cpu_count()//2, squarem=False, formatter=StdPrinter,
                   outstream=lambda: sys.stdout):
@@ -473,6 +477,7 @@ cdef class opt_lim_const:
     def max_iter(self):
         """Maximum number of iterations before ending optimization"""
         return self._max_iter
+    
     @max_iter.setter
     def max_iter(self, max_iter):
         assert max_iter > 0 and np.issubdtype(type(max_iter), np.integer), ValueError("max_iter must be integer greater than 0")
@@ -489,6 +494,7 @@ cdef class opt_lim_const:
             
         """
         return self._max_time
+    
     @max_time.setter
     def max_time(self, max_time):
         assert max_time > 0 and np.issubdtype(type(max_time), np.floating), ValueError("max_time must be float and greater than 0")
@@ -507,7 +513,8 @@ cdef class opt_lim_const:
 
     @property
     def num_cores(self):
-        """Number of cores to use in optimization
+        """
+        Number of cores to use in optimization
         
         .. note::
             
@@ -561,7 +568,7 @@ cdef class opt_lim_const:
 
     def __repr__(self):
         return (f'Optimization limits:: num_cores: {self._num_cores}, max_iter: {self._max_iter}, '
-                f'converged_min: {self._converged_min}, max_time: {self._max_time},\n'
+                f'converged_min: {self._converged_min}, max_time: {self._max_time},squarem={self._squarem}\n'
                 f'formatter: {self._formatter.__qualname__}, outstream={self._outstream}')
 
     def _get_max_iter(self, max_iter):
@@ -722,25 +729,132 @@ cdef class opt_lim_const:
         Parameters
         ----------
         outstream : Output Stream
-            Stream .
+            Stream to assign as outptu stream.
 
         Returns
         -------
         OutputStream
-            DESCRIPTION.
+            default output stream object.
 
         """
         if outstream is None:
             return self._outstream()
         return outstream
+    
+    def __call__(self, **kwargs):
+        """
+        Open a context manager, all keyword arguments must be valid attributes
+        of optimization_limits. Values set in kwargs, or to optimization limits
+        within the ``with`` block will be reset upon exciting the ``with`` block.
+        
+        This can be used as follows::
+            
+            with hm.optimization_limits(converged_min=1e-5):
+                out = hm.EM_H2MM_C(init, index, times)
+        
+        
+        Parameters
+        ----------
+        max_iter : int, optional
+            Maximum number of iterations before ending optimization
+        max_time : float, optional
+            Maximum time of optimization.
+        converged_min : float, optional
+            Minimum difference between loglik of succesive iterations for a
+            model to be considered converged
+        num_cores : int, optional
+            Number of cores to use in optimization.
+        squarem : bool, optional
+            Choice of whether or not to use SQUAREM (|Varadhan2008|) method of
+            accelerating optimization using quadratic projection of 3 successive
+            models.
+        formatter : Printer
+            A :class:`Printer` subclass, used to format result of print_func 
+            for display of optimiation progress
+        outstream : Callable, optional
+            Callable, when called, generates an output stream for where to 
+            print displayed progress
+        
+        Returns
+        -------
+        OptlimContext
+            Special context manager class for managing :class:opt_lim_const`
+            This class only defines the ``__init__``, ``__enter__`` and ``__exit__``
+            magic methods for use in with statemetns
+        """
+        for key in kwargs.keys():
+            if key not in self._attrs:
+                raise TypeError(f"{key} is not a valid keyword argument for context manager")
+        cntx = OptlimContext(self)
+        for key, value in kwargs.items():
+            if value is None:
+                continue
+            try:
+                self[key] = value
+            except Exception as e:
+                print("exception raised")
+                for k, v in cntx._values.items():
+                    self[key] = value
+                raise e
+        return cntx
+
+
+cdef class OptlimContext:
+    """
+    Context manager class, upon enter, records the values of ``optlim`` and
+    restores upont exit
+    
+    Parameters
+    ----------
+    optlim : opt_lim_const
+        The optimization_limits to restore after exiting the context manager.
+    """
+    cdef:
+        opt_lim_const _optlim
+        dict _values
+        
+    def __cinit__(self, opt_lim_const optlim):
+        self._optlim = optlim
+        self._values = {key:optlim[key] for key in opt_lim_const._attrs}
+    
+    def __enter__(self):
+        return self._optlim
+    
+    def __exit__(self, exc_type, exc_value, traceback):
+        for key, value in self._values.items():
+            self._optlim[key] = value
+
 
 def get_stdout():
     """Simple callable returns sys.stdout"""
     return sys.stdout
 
-#: Like rcParams, a place to set defaults for optimization limits
+#: This module variable stores default values for optimization limits.
+#: if the values are not supplied as keyword arguments to the optimization functions,
+#: ``H2MM_C`` will fall back to the values stored in this variable.
+#: The options can be accessed either as attributions or keys.
+#: 
+#: Options are:
+#: - max_iter: maximum number of iterations to perform during optimization
+#: - max_time: maximum 
+#: - converged_min: difference between loglik of iteration at which to terminate optimization 
+#:   (if difference is less than this value, the model is considered optimized)
+#: - num_cores: number of cores to use in optimization
+#: - squarem: Whether or not to use squarem acceleration
+#: - formatter: printer object, defines how to format optimization progress
+#: - outstream: output stream buffer
+#: 
+#: 
+#: Support context manager to temporarily set a value:
+#:
+#: .. code-block::
+#:
+#:     with hm.optimization_limits(max_iter=1000):
+#:         model = hm.EM_H2MM_C(model, index, time)
+#: 
 optimization_limits = opt_lim_const(max_iter=2046, max_time=np.inf, 
                                     converged_min=1e-14, num_cores=os.cpu_count()//2,
+                                    squarem=False,
                                     formatter=StdPrinter, outstream=get_stdout)
 
 cdef class ConvCodes:
@@ -1715,7 +1829,7 @@ cdef class h2mm_model:
                       print_stream=print_stream, print_formatter=print_formatter,
                       print_fmt_args=print_fmt_args, print_fmt_kwargs=print_fmt_kwargs,
                       max_time=max_time, converged_min=converged_min, num_cores=num_cores, 
-                      reset_niter=reset_niter, ll=ll, gamma=gamma, opt_array=opt_array)
+                      reset_niter=reset_niter, ll=ll, gamma=gamma, opt_array=opt_array, squarem=squarem)
         if inplace:
             # separate the models from ll_array / gamma
             if gamma or ll:

docs/source/ChangeNotes.rst

@@ -3,6 +3,18 @@ Change Notes
 
 .. currentmodule:: H2MM_C
 
+Changed in Version 2.3.1
+------------------------
+
+The main addition in v2.3.1 is the context manager for optimization limits
+
+- Added:
+    - Optimization limits context manager. :attr:`H2MM_C.optimization_limits`` can now be called in a with block to
+      modify the limits only withing the with block. This also added the :class:`H2MM_C.OptlimContext`class to support it.
+
+- Fixed:
+    - ``squarem`` keyword argument not passed in :meth:`H2MM_C.h2mm_model.optimize` this has now been fixed
+
 
 Changed in Version 2.3.0
 ------------------------

docs/source/SecondaryControl.rst

@@ -35,6 +35,7 @@ First, our obligatory imports and loading of 3 detector data:
 
     color3, times3 = load_txtdata('sample_data_3det.txt')
 
+
 Optimization Control
 --------------------
 
@@ -263,4 +264,142 @@ This is useful if you want to save your data for later:
 | 0.840976670297041, 0.07645342127367338, 0.08256990842928548
 
 
+New in v2.3
+-----------
+
+SQUAREM acceleration
+********************
+
+The primary addition going from v2.2 to v2.3 is SQUAREM acceleration 
+for the expectation-maximization (EM) aspect of |H2MM| algorithm.
+
+Many thanks to Thomas Peulen for recomending this method.
+
+
+The method is based on the pseudo-code of |Du2018|,
+which was derived from ideas in |Varadhan2008|.
+Note that all implementation was done at the C level.
+
+SQUAREM expands/modifies the normal expectation-maximization loop.
+
+We will represent a model with :math:`\theta`.
+In |H2MM|, there are several sub-steps.
+Usally refered two as the forward-back algorithm, 
+The  loglikelihood of a model :math:`\theta` is calculated,
+and also computes a model that is guaunteed to have a better loglikelihood 
+(unless the optimal model has already been found, also subject to floating point errors).
+This algorithm has two phases, the $fwd$ phase which computes that loglikelihood,
+and the :math:`bck` phase, (which relies on an array computed in :math:`fwd` ), 
+which computes the improved model.
+
+The standard EM loop for |H2MM| uses the loop below:
+
+-------
+
+while :math:`ll(\theta_{n-1} < ll(\theta_{n})` \:\:
+
+1. Compute loglikelihood of current model :math:`ll(\theta_{n}) = fwd(\theta_{n})`
+2. Compute maximization step next model :math:`\theta_{n+1} = bck(\theta_{n})`
+
+
+--------
+
+SQUAREM expands this loop, so that a quadractic projection of :math:`\theta` is computed based on the 3 most recent models.
+Then if the projected model has a better loglikelhood that the last model, use the projected model in the next round.
+
+The SQUAREM optimization loop works as follows:
+
+-----------
+
+Initialize :math;`ll(\theta_{0}) = fwd(\theta_{0})` (compute loglikelihood of initial model)
+
+while optimal model not found::
+1. Compute :math:`\theta_{n+1} = bck(\theta{n^{\prime}})`
+2. Compute :math:`ll(\theta_{n+1}) = fwd(\theta_{n+1})`
+3. Compute :math:`\theta_{n+2} = bck(\theta_{n+1})`
+4. Compute :math:`ll(\theta_{n+2}) = fwd(\theta_{n+2})`
+5. Compute projection using the following definitions:
+
+    - :math:`r = \theta_{n+1} - \theta_{n}`
+    - :math:`b = \theta_{n+2} - \theta_{n+1}`
+    - :math:`\alpha = - \frac{||r||}{||v||}`
+    - :math:`\theta_{SQ} = \theta_{n} - 2\alpha r + \alpha^{2} v`
+    - **note** it is possible for :math:`\theta_{SQ}` to have negative (impossible) values, 
+      if this occurs, skip to step 7, and assume :math:`\theta_{SQ}` has the worst possible loglikelihood
+
+6. Compute :math:`ll(\theta_{SQ}) = fwd(\theta_{SQ})`
+7. Set :math:`\theta_{{n+2}^{\prime}}` as which model has the better loglikelihood between :math:`\theta_{n+1}` and :math:`\theta_{SQ}`
+
+-------------
+
+Note that after each evaluation of $fwd$, the loglik is tested for convergence.
+Also, between these evaluation, the models are normalized and bounded with the bounding function.
+
+.. code-block::
+
+    imodel_4s3d = hm.factory_h2mm_model(4,3)
+    model_4s3d = hm.EM_H2MM_C(imodel_4s3d, color3, times3, squarem=True)
+    model_4s3d
+
+| Optimization reached maximum number of iterations
+| 
+| nstate: 4, ndet: 3, nphot: 436084, niter: 1000, loglik: -408233.7176734728 converged state: 0x146
+| prior:
+| 0.6252741215262306, 0.030612454324407904, 0.19173015960537806, 0.1523832645439835
+| trans:
+| 0.9057287693195645, 0.04050721299160209, 0.05376401768883339, 1.960153729459685e-38
+| 2.0596457882248302e-05, 0.9999671506799463, 1.0257544140118528e-05, 1.995318031340882e-06
+| 4.9711944981109055e-06, 6.354917051447945e-06, 0.9999787515213548, 9.922367095667794e-06
+| 3.0728033099149917e-06, 2.564195058738289e-06, 3.7901004759624195e-05, 0.9999564619968717
+| obs:
+| 0.6178100534248961, 0.3821899465751039, 1.0611886068674382e-22
+| 0.14422555137985735, 0.2894894306973517, 0.566285017922791
+| 0.4416303955720442, 0.08699556344822588, 0.4713740409797299
+| 0.8409816137668159, 0.07645524944471362, 0.08256313678847053
+
+
+A note on using ``opt_array`` with SQUAREM:
+
+When using SQUAREM with ``opt_array=False`` (the default), the maximum number of iteratons and the ``hm.h2mm_model.niter`` values match,
+and importantly the :math:`\theta_{SQ}` models are counted as having the same number of iteration as the :math:`\theta_{n+2}` model.
+
+However, when using :code:`opt_array=True` the :math:`\theta_{n+2}` and :math:`\theta_{SQ}` models are both stored.
+The maximum number of iterations then becomes the number of models in the output array.
+Therefore fewer iterations will be evalueated with :code:`opt_array=True, squarem=True` than with :math:`opt_array=False, squarem=True`.
+
+
+optimization_limits context manager
+***********************************
+
+Similar to how in matplotlib ``rcParams`` can open a context manager,
+where changes inside the `with` block are restored upon exit,
+:attr:`H2MM_C.optimization_limits` can now also be used inside a with block.
+
+Context-limited values can be set either in the kwargs of the with statement,
+setting :attr:`H2MM_C.optimization_limits`, or using the return value of the context manager.
+
+.. code-block::
+
+    print(f'before opening context: {hm.optimization_limits}')
+    with hm.optimization_limits(squarem=True) as lm:
+        print(f'inside context: {hm.optimization_limits}')
+        hm.optimization_limits.max_iter = 123
+        lm.converged_min = 1e-4
+        print(f'set values inside context: {hm.optimization_limits}')
+        models_4s3d = hm.EM_H2MM_C(imodel_4s3d, color3, times3, opt_array=True)
+    print(f'after context: {hm.optimization_limits}')
+
+| before opening context: Optimization limits:: num_cores: 2, max_iter: 1000, converged_min: 1e-07, max_time: inf,squarem=False
+| formatter: StdPrinter, outstream=<cyfunction get_stdout at 0x7704fe257d10>
+| inside context: Optimization limits:: num_cores: 2, max_iter: 1000, converged_min: 1e-07, max_time: inf,squarem=True
+| formatter: StdPrinter, outstream=<cyfunction get_stdout at 0x7704fe257d10>
+| set values inside context: Optimization limits:: num_cores: 2, max_iter: 123, converged_min: 0.0001, max_time: inf,squarem=True
+| formatter: StdPrinter, outstream=<cyfunction get_stdout at 0x7704fe257d10>
+| Optimization reached maximum number of iterations
+| after context: Optimization limits:: num_cores: 2, max_iter: 1000, converged_min: 1e-07, max_time: inf,squarem=False
+| formatter: StdPrinter, outstream=<cyfunction get_stdout at 0x7704fe257d10>
+
+
+.. |Du2018| replace:: `Du & Varadhan 2018 <https://doi.org/10.48550/arXiv.1810.11163>`__
+.. |Varadhan2008| replace:: `Varadhan & Roland 2008. <https://doi.org/10.1111/j.1467-9469.2007.00585.x>`__
 .. |H2MM| replace:: H\ :sup:`2`\ MM