From 292400cbf79c059c7b4812e0b816fc4a0385697b Mon Sep 17 00:00:00 2001
From: oscarxblanco <oblanco@cells.es>
Date: Wed, 17 Apr 2024 12:38:53 +0200
Subject: [PATCH 1/9] adds get_freqmap

---
 pyat/at/physics/__init__.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pyat/at/physics/__init__.py b/pyat/at/physics/__init__.py
index 6bfc5ca2f3..fc3ccd0630 100644
--- a/pyat/at/physics/__init__.py
+++ b/pyat/at/physics/__init__.py
@@ -16,4 +16,4 @@
 from .ring_parameters import *
 from .nonlinear import *
 from .fastring import *
-from .frequency_maps import fmap_parallel_track
+from .frequency_maps import fmap_parallel_track, get_freqmap

From c006d5cf217539d32358f170acb6951c768b8cf0 Mon Sep 17 00:00:00 2001
From: oscarxblanco <oblanco@cells.es>
Date: Wed, 17 Apr 2024 12:39:38 +0200
Subject: [PATCH 2/9] adds get_freqmap, uses Grid from get_acceptance

---
 pyat/at/physics/frequency_maps.py | 227 +++++++++++++++++++++++++++++-
 1 file changed, 226 insertions(+), 1 deletion(-)

diff --git a/pyat/at/physics/frequency_maps.py b/pyat/at/physics/frequency_maps.py
index 9e7c9f1d7a..6986528629 100644
--- a/pyat/at/physics/frequency_maps.py
+++ b/pyat/at/physics/frequency_maps.py
@@ -5,6 +5,7 @@
 
 # orblancog
 # generates the frequency and diffusion map for a given ring
+# 2024apr16 create get_freqmap
 # 2023jan16 tracking is parallel (patpass), frequency analysis is serial
 # 2022jun07 serial version
 
@@ -13,12 +14,236 @@
 import numpy
 from warnings import warn
 from at.lattice import AtWarning
+from ..acceptance.boundary import set_ring_orbit
+from ..acceptance.boundary import grid_configuration
+from ..acceptance.boundary import get_parts
+from ..acceptance.boundary import get_survived
+from ..acceptance.boundary import GridMode
+from ..acceptance.boundary import get_plane_index
+from at.lattice import Lattice, AtError, AtWarning
+from typing import Optional, Sequence
+from ..lattice import Lattice, Refpts, frequency_control, AtError
+import multiprocessing
+import time
+
+_pdict = {'x': 0, 'xp': 1,
+          'y': 2, 'yp': 3,
+          'dp': 4, 'ct': 5}
 
 
 # Jaime Coello de Portugal (JCdP) frequency analysis implementation
 from .harmonic_analysis import get_tunes_harmonic
 
-__all__ = ['fmap_parallel_track']
+__all__ = ['fmap_parallel_track','get_freqmap']
+
+def get_freqmap(
+        ring: Lattice,
+        planes,
+        npoints,
+        amplitudes,
+        bounds=None,
+        nturns: Optional[int] = 512,
+        dp: Optional[float] = None,
+        offset: Sequence[float] = None,
+        refpts: Optional[Refpts] = 0,
+        grid_mode: Optional[GridMode] = GridMode.CARTESIAN,
+        use_mp: Optional[bool] = True,
+        verbose: Optional[bool] = False,
+        lossmap: Optional[int] = 2,
+        shift_zero: Optional[float] = 0.0e-6,
+        start_method: Optional[str] = None,
+):
+    # noinspection PyUnresolvedReferences
+    r"""Computes the acceptance at ``repfts`` observation points
+
+    Parameters:
+        ring:           Lattice definition
+        planes:         max. dimension 2, Plane(s) to scan for the acceptance.
+          Allowed values are: ``'x'``, ``'xp'``, ``'y'``,
+          ``'yp'``, ``'dp'``, ``'ct'``
+        npoints:        (len(planes),) array: number of points in each
+          dimension
+        amplitudes:     (len(planes),) array: set the search range:
+
+          * :py:attr:`GridMode.CARTESIAN/RADIAL <.GridMode.RADIAL>`:
+            max. amplitude
+          * :py:attr:`.GridMode.RECURSIVE`: initial step
+        nturns:         Number of turns for the tracking
+        refpts:         Observation points. Default: start of the machine
+        dp:             static momentum offset
+        offset:         initial orbit. Default: closed orbit
+        bounds:         defines the tracked range: range=bounds*amplitude.
+          It can be use to select quadrants. For example, default values are:
+
+          * :py:attr:`.GridMode.CARTESIAN`: ((-1, 1), (0, 1))
+          * :py:attr:`GridMode.RADIAL/RECURSIVE <.GridMode.RADIAL>`: ((0, 1),
+            (:math:`\pi`, 0))
+        grid_mode:      defines the evaluation grid:
+
+          * :py:attr:`.GridMode.CARTESIAN`: full [:math:`\:x, y\:`] grid
+          * :py:attr:`.GridMode.RADIAL`: full [:math:`\:r, \theta\:`] grid
+          * :py:attr:`.GridMode.RECURSIVE`: radial recursive search
+        use_mp:         Use python multiprocessing (:py:func:`.patpass`,
+          default use :py:func:`.lattice_pass`).
+        verbose:        Print out some information
+        divider:        Value of the divider used in
+          :py:attr:`.GridMode.RECURSIVE` boundary search
+        shift_zero: Epsilon offset applied on all 6 coordinates
+        start_method:   Python multiprocessing start method. The default
+          ``None`` uses the python default that is considered safe.
+          Available parameters: ``'fork'``, ``'spawn'``, ``'forkserver'``.
+          The default for linux is ``'fork'``, the default for MacOS and
+          Windows is ``'spawn'``. ``'fork'`` may used for MacOS to speed-up
+          the calculation or to solve runtime errors, however  it is
+          considered unsafe.
+
+    Returns:
+        boundary:   (2,n) array: 2D acceptance
+        survived:   (2,n) array: Coordinates of surviving particles
+        tracked:    (2,n) array: Coordinates of tracked particles
+
+    In case of multiple refpts, return values are lists of arrays, with one
+    array per ref. point.
+
+    Examples:
+
+        >>> bf,sf,gf = ring.get_acceptance(planes, npoints, amplitudes)
+        >>> plt.plot(*gf,'.')
+        >>> plt.plot(*sf,'.')
+        >>> plt.plot(*bf)
+        >>> plt.show()
+
+    .. note::
+
+       * When``use_mp=True`` all the available CPUs will be used.
+         This behavior can be changed by setting
+         ``at.DConstant.patpass_poolsize`` to the desired value
+    """
+    kwargs = {}
+    if start_method is not None:
+        kwargs['start_method'] = start_method
+
+    if verbose:
+        nproc = multiprocessing.cpu_count()
+        print('\n{0} cpu found for acceptance calculation'.format(nproc))
+        if use_mp:
+            nprocu = nproc
+            print('Multi-process acceptance calculation selected...')
+            if nproc == 1:
+                print('Consider use_mp=False for single core computations')
+        else:
+            nprocu = 1
+            print('Single process acceptance calculation selected...')
+            if nproc > 1:
+                print('Consider use_mp=True for parallelized computations')
+        np = numpy.atleast_1d(npoints)
+        na = 2
+        if len(np) == 2:
+            na = np[1]
+        npp = numpy.prod(npoints)
+        rpp = 2*numpy.ceil(numpy.log2(np[0]))*numpy.ceil(na/nprocu)
+        mpp = npp/nprocu
+        #if rpp > mpp:
+        #    cond = (grid_mode is GridMode.RADIAL or
+        #            grid_mode is GridMode.CARTESIAN)
+        #else:
+        #    cond = grid_mode is GridMode.RECURSIVE
+        #if rpp > mpp and not cond:
+        #    print('The estimated load for grid mode is {0}'.format(mpp))
+        #    print('The estimated load for recursive mode is {0}'.format(rpp))
+        #    print('{0} or {1} is recommended'.format(GridMode.RADIAL,
+        #                                             GridMode.CARTESIAN))
+        #elif rpp < mpp and not cond:
+        #    print('The estimated load for grid mode is {0}'.format(mpp))
+        #    print('The estimated load for recursive mode is {0}'.format(rpp))
+        #    print('{0} is recommended'.format(GridMode.RECURSIVE))
+
+    boundary = []
+    survived = []
+    grid = []
+    if refpts is not None:
+        rp = ring.uint32_refpts(refpts)
+    else:
+        rp = numpy.atleast_1d(refpts)
+    if offset is not None:
+        try:
+            offset = numpy.broadcast_to(offset, (len(rp), 6))
+        except ValueError:
+            msg = ('offset and refpts have incoherent '
+                   'shapes: {0}, {1}'.format(numpy.shape(offset),
+                                             numpy.shape(refpts)))
+            raise AtError(msg)
+    else:
+        offset = find_orbit(ring,refpts)
+        planesi = numpy.atleast_1d(get_plane_index(planes))
+        offset[0][planesi[0]] =  offset[0][planesi[0]] + 1e-6
+        offset[0][planesi[1]] =  offset[0][planesi[1]] + 1e-6
+        #offset=[None for _ in rp]
+    dataobs = []
+    t0 = time.time()
+    for r, o in zip(rp, offset):
+        #b, s, g = boundary_search(ring, planes, npoints, amplitudes,
+        #                          nturns=nturns, obspt=r, dp=dp,
+        #                          offset=o, bounds=bounds,
+        #                          grid_mode=grid_mode, use_mp=use_mp,
+        #                          verbose=verbose, divider=divider,
+        #                          shift_zero=shift_zero, **kwargs)
+        obspt=r
+        dp=dp
+        offset=o
+
+        offset, newring = set_ring_orbit(ring, dp, obspt,
+                                         offset)
+        config = grid_configuration(planes, npoints, amplitudes,
+                                    grid_mode, bounds=bounds,
+                                    shift_zero=shift_zero)
+        obspt = None
+        if verbose:
+            print('\nRunning grid boundary search:')
+            if obspt is None:
+                print('Element {0}, obspt={1}'.format(ring[0].FamName, 0))
+            else:
+                print('Element {0}, obspt={1}'.format(ring[obspt].FamName,
+                                                      obspt))
+            print('The grid mode is {0}'.format(config.mode))
+            print('The planes are {0}'.format(config.planes))
+            print('Number of steps are {0}'.format(config.shape))
+            print('The maximum amplitudes are {0}'.format(config.amplitudes))
+            print('The maximum boundaries are {0}'.format(config.bounds))
+            print('The initial offset is {0} with dp={1}'.format(offset, dp))
+            parts, grid = get_parts(config, offset)
+            rout, tp, td = ring.track(parts, nturns=2*nturns, losses=True, use_mp=use_mp, **kwargs)
+
+            mask = get_survived(parts, newring, nturns, use_mp, **kwargs)
+            survived = grid.grid[:, mask]
+
+            planesi = numpy.atleast_1d(get_plane_index(planes))
+            tunes = numpy.zeros((len(planesi),2,len(numpy.where(mask)[0])))
+            print(planesi)
+
+            for i in range(len(planesi)):
+                tunes[i,0] = get_tunes_harmonic(rout[planesi[i],mask,:,     0:  nturns],use_mp=use_mp, **kwargs)
+                tunes[i,1] = get_tunes_harmonic(rout[planesi[i],mask,:,nturns:2*nturns],use_mp=use_mp, **kwargs)
+
+            # metric
+            diffplane1 = tunes[0,0,:] - tunes[0,1,:]
+            diffplane2 = tunes[1,0,:] - tunes[1,1,:]
+            nudiff = 0.5*numpy.log10( (diffplane1*diffplane1 + diffplane2*diffplane2) /nturns )
+            # set min-max
+            nudiff = numpy.clip(nudiff,-10,-2)
+
+            #return rout,tp,td,grid,tunes
+            firstturns = 0
+#            return numpy.concatenate((survived.T, tunes[0,firstturns].T, tunes[1,firstturns].T, diffplane1.T, diffplane2.T, nudiff.T),axis=0), grid, td
+            #return survived.T, tunes[0,firstturns].T, tunes[1,firstturns].T, diffplane1.T, diffplane2.T, nudiff.T, grid, td
+            dataobs.append((numpy.vstack((survived, tunes[0,firstturns], tunes[1,firstturns], diffplane1, diffplane2, nudiff)).T, grid, td))
+    if verbose:
+        print('Calculation took {0}'.format(time.time()-t0))
+
+    return dataobs
+
+
+
 
 
 def fmap_parallel_track(ring,

From 6acac0adb01901a6be137a29afbf8eced46a5193 Mon Sep 17 00:00:00 2001
From: oscarxblanco <oblanco@cells.es>
Date: Fri, 28 Jun 2024 16:50:15 +0200
Subject: [PATCH 3/9] works with multiple refpts

---
 pyat/at/physics/frequency_maps.py | 96 ++++++++++++-------------------
 1 file changed, 36 insertions(+), 60 deletions(-)

diff --git a/pyat/at/physics/frequency_maps.py b/pyat/at/physics/frequency_maps.py
index 6986528629..0cbb0d441c 100644
--- a/pyat/at/physics/frequency_maps.py
+++ b/pyat/at/physics/frequency_maps.py
@@ -8,6 +8,7 @@
 # 2024apr16 create get_freqmap
 # 2023jan16 tracking is parallel (patpass), frequency analysis is serial
 # 2022jun07 serial version
+# 2024jun21 adds get_freqmap
 
 from at.tracking import patpass
 from at.physics import find_orbit
@@ -54,7 +55,7 @@ def get_freqmap(
         start_method: Optional[str] = None,
 ):
     # noinspection PyUnresolvedReferences
-    r"""Computes the acceptance at ``repfts`` observation points
+    r"""Computes the frequency map at ``repfts`` observation points.
 
     Parameters:
         ring:           Lattice definition
@@ -98,9 +99,13 @@ def get_freqmap(
           considered unsafe.
 
     Returns:
-        boundary:   (2,n) array: 2D acceptance
-        survived:   (2,n) array: Coordinates of surviving particles
-        tracked:    (2,n) array: Coordinates of tracked particles
+        fmaout:   list containing the result of every reference point.
+            Inside,
+              1) the frequency map
+                `[xcoor, ycoor, nux, ny, dnux, dnuy,
+                 log10(sqrt(sum(dnu**2)/turns)) ]`
+              2) the grid
+              3) the dictionary of particle losses
 
     In case of multiple refpts, return values are lists of arrays, with one
     array per ref. point.
@@ -143,22 +148,7 @@ def get_freqmap(
         npp = numpy.prod(npoints)
         rpp = 2*numpy.ceil(numpy.log2(np[0]))*numpy.ceil(na/nprocu)
         mpp = npp/nprocu
-        #if rpp > mpp:
-        #    cond = (grid_mode is GridMode.RADIAL or
-        #            grid_mode is GridMode.CARTESIAN)
-        #else:
-        #    cond = grid_mode is GridMode.RECURSIVE
-        #if rpp > mpp and not cond:
-        #    print('The estimated load for grid mode is {0}'.format(mpp))
-        #    print('The estimated load for recursive mode is {0}'.format(rpp))
-        #    print('{0} or {1} is recommended'.format(GridMode.RADIAL,
-        #                                             GridMode.CARTESIAN))
-        #elif rpp < mpp and not cond:
-        #    print('The estimated load for grid mode is {0}'.format(mpp))
-        #    print('The estimated load for recursive mode is {0}'.format(rpp))
-        #    print('{0} is recommended'.format(GridMode.RECURSIVE))
-
-    boundary = []
+
     survived = []
     grid = []
     if refpts is not None:
@@ -174,32 +164,23 @@ def get_freqmap(
                                              numpy.shape(refpts)))
             raise AtError(msg)
     else:
-        offset = find_orbit(ring,refpts)
+        _, offset = find_orbit(ring,refpts)
         planesi = numpy.atleast_1d(get_plane_index(planes))
-        offset[0][planesi[0]] =  offset[0][planesi[0]] + 1e-6
-        offset[0][planesi[1]] =  offset[0][planesi[1]] + 1e-6
-        #offset=[None for _ in rp]
+        offset[:,planesi[0]] =  offset[:,planesi[0]] + 1e-9
+        offset[:,planesi[1]] =  offset[:,planesi[1]] + 1e-9
     dataobs = []
     t0 = time.time()
     for r, o in zip(rp, offset):
-        #b, s, g = boundary_search(ring, planes, npoints, amplitudes,
-        #                          nturns=nturns, obspt=r, dp=dp,
-        #                          offset=o, bounds=bounds,
-        #                          grid_mode=grid_mode, use_mp=use_mp,
-        #                          verbose=verbose, divider=divider,
-        #                          shift_zero=shift_zero, **kwargs)
         obspt=r
         dp=dp
         offset=o
-
         offset, newring = set_ring_orbit(ring, dp, obspt,
                                          offset)
         config = grid_configuration(planes, npoints, amplitudes,
                                     grid_mode, bounds=bounds,
                                     shift_zero=shift_zero)
-        obspt = None
         if verbose:
-            print('\nRunning grid boundary search:')
+            print('\nRunning grid frequency search:')
             if obspt is None:
                 print('Element {0}, obspt={1}'.format(ring[0].FamName, 0))
             else:
@@ -210,42 +191,37 @@ def get_freqmap(
             print('Number of steps are {0}'.format(config.shape))
             print('The maximum amplitudes are {0}'.format(config.amplitudes))
             print('The maximum boundaries are {0}'.format(config.bounds))
+            print('Number of turns is {0}'.format(nturns))
             print('The initial offset is {0} with dp={1}'.format(offset, dp))
-            parts, grid = get_parts(config, offset)
-            rout, tp, td = ring.track(parts, nturns=2*nturns, losses=True, use_mp=use_mp, **kwargs)
-
-            mask = get_survived(parts, newring, nturns, use_mp, **kwargs)
-            survived = grid.grid[:, mask]
+        parts, grid = get_parts(config, offset)
+        rout, tp, td = ring.track(parts, nturns=2*nturns, losses=True, use_mp=use_mp, **kwargs)
 
-            planesi = numpy.atleast_1d(get_plane_index(planes))
-            tunes = numpy.zeros((len(planesi),2,len(numpy.where(mask)[0])))
-            print(planesi)
+        mask = get_survived(parts, newring, nturns, use_mp, **kwargs)
+        survived = grid.grid[:, mask]
 
-            for i in range(len(planesi)):
-                tunes[i,0] = get_tunes_harmonic(rout[planesi[i],mask,:,     0:  nturns],use_mp=use_mp, **kwargs)
-                tunes[i,1] = get_tunes_harmonic(rout[planesi[i],mask,:,nturns:2*nturns],use_mp=use_mp, **kwargs)
-
-            # metric
-            diffplane1 = tunes[0,0,:] - tunes[0,1,:]
-            diffplane2 = tunes[1,0,:] - tunes[1,1,:]
-            nudiff = 0.5*numpy.log10( (diffplane1*diffplane1 + diffplane2*diffplane2) /nturns )
-            # set min-max
-            nudiff = numpy.clip(nudiff,-10,-2)
-
-            #return rout,tp,td,grid,tunes
-            firstturns = 0
-#            return numpy.concatenate((survived.T, tunes[0,firstturns].T, tunes[1,firstturns].T, diffplane1.T, diffplane2.T, nudiff.T),axis=0), grid, td
-            #return survived.T, tunes[0,firstturns].T, tunes[1,firstturns].T, diffplane1.T, diffplane2.T, nudiff.T, grid, td
-            dataobs.append((numpy.vstack((survived, tunes[0,firstturns], tunes[1,firstturns], diffplane1, diffplane2, nudiff)).T, grid, td))
+        planesi = numpy.atleast_1d(get_plane_index(planes))
+        tunes = numpy.zeros((len(planesi),2,len(numpy.where(mask)[0])))
+
+
+        for i in range(len(planesi)):
+            tunes[i,0] = get_tunes_harmonic(rout[planesi[i],mask,:,     0:  nturns],use_mp=use_mp, **kwargs)
+            tunes[i,1] = get_tunes_harmonic(rout[planesi[i],mask,:,nturns:2*nturns],use_mp=use_mp, **kwargs)
+        # metric
+        diffplane1 = tunes[0,0,:] - tunes[0,1,:]
+        diffplane2 = tunes[1,0,:] - tunes[1,1,:]
+        nudiff = 0.5*numpy.log10( (diffplane1*diffplane1 + diffplane2*diffplane2) /nturns )
+        # set min-max
+        nudiff = numpy.clip(nudiff,-10,-2)
+
+        #return rout,tp,td,grid,tunes
+        firstturns = 0
+        dataobs.append((numpy.vstack((survived, tunes[0,firstturns], tunes[1,firstturns], diffplane1, diffplane2, nudiff)).T, grid, td))
     if verbose:
         print('Calculation took {0}'.format(time.time()-t0))
 
     return dataobs
 
 
-
-
-
 def fmap_parallel_track(ring,
                         coords=[-10, 10, -10, 10],
                         steps=[100, 100],

From d69f54117859a9bc547fd93e350bc579739dd50b Mon Sep 17 00:00:00 2001
From: oscarxblanco <oblanco@cells.es>
Date: Fri, 28 Jun 2024 16:52:18 +0200
Subject: [PATCH 4/9] black on frequency_maps

---
 pyat/at/physics/frequency_maps.py | 267 +++++++++++++++++-------------
 1 file changed, 148 insertions(+), 119 deletions(-)

diff --git a/pyat/at/physics/frequency_maps.py b/pyat/at/physics/frequency_maps.py
index 0cbb0d441c..7b4d8de6ed 100644
--- a/pyat/at/physics/frequency_maps.py
+++ b/pyat/at/physics/frequency_maps.py
@@ -27,32 +27,31 @@
 import multiprocessing
 import time
 
-_pdict = {'x': 0, 'xp': 1,
-          'y': 2, 'yp': 3,
-          'dp': 4, 'ct': 5}
+_pdict = {"x": 0, "xp": 1, "y": 2, "yp": 3, "dp": 4, "ct": 5}
 
 
 # Jaime Coello de Portugal (JCdP) frequency analysis implementation
 from .harmonic_analysis import get_tunes_harmonic
 
-__all__ = ['fmap_parallel_track','get_freqmap']
+__all__ = ["fmap_parallel_track", "get_freqmap"]
+
 
 def get_freqmap(
-        ring: Lattice,
-        planes,
-        npoints,
-        amplitudes,
-        bounds=None,
-        nturns: Optional[int] = 512,
-        dp: Optional[float] = None,
-        offset: Sequence[float] = None,
-        refpts: Optional[Refpts] = 0,
-        grid_mode: Optional[GridMode] = GridMode.CARTESIAN,
-        use_mp: Optional[bool] = True,
-        verbose: Optional[bool] = False,
-        lossmap: Optional[int] = 2,
-        shift_zero: Optional[float] = 0.0e-6,
-        start_method: Optional[str] = None,
+    ring: Lattice,
+    planes,
+    npoints,
+    amplitudes,
+    bounds=None,
+    nturns: Optional[int] = 512,
+    dp: Optional[float] = None,
+    offset: Sequence[float] = None,
+    refpts: Optional[Refpts] = 0,
+    grid_mode: Optional[GridMode] = GridMode.CARTESIAN,
+    use_mp: Optional[bool] = True,
+    verbose: Optional[bool] = False,
+    lossmap: Optional[int] = 2,
+    shift_zero: Optional[float] = 0.0e-6,
+    start_method: Optional[str] = None,
 ):
     # noinspection PyUnresolvedReferences
     r"""Computes the frequency map at ``repfts`` observation points.
@@ -126,28 +125,28 @@ def get_freqmap(
     """
     kwargs = {}
     if start_method is not None:
-        kwargs['start_method'] = start_method
+        kwargs["start_method"] = start_method
 
     if verbose:
         nproc = multiprocessing.cpu_count()
-        print('\n{0} cpu found for acceptance calculation'.format(nproc))
+        print("\n{0} cpu found for acceptance calculation".format(nproc))
         if use_mp:
             nprocu = nproc
-            print('Multi-process acceptance calculation selected...')
+            print("Multi-process acceptance calculation selected...")
             if nproc == 1:
-                print('Consider use_mp=False for single core computations')
+                print("Consider use_mp=False for single core computations")
         else:
             nprocu = 1
-            print('Single process acceptance calculation selected...')
+            print("Single process acceptance calculation selected...")
             if nproc > 1:
-                print('Consider use_mp=True for parallelized computations')
+                print("Consider use_mp=True for parallelized computations")
         np = numpy.atleast_1d(npoints)
         na = 2
         if len(np) == 2:
             na = np[1]
         npp = numpy.prod(npoints)
-        rpp = 2*numpy.ceil(numpy.log2(np[0]))*numpy.ceil(na/nprocu)
-        mpp = npp/nprocu
+        rpp = 2 * numpy.ceil(numpy.log2(np[0])) * numpy.ceil(na / nprocu)
+        mpp = npp / nprocu
 
     survived = []
     grid = []
@@ -159,80 +158,101 @@ def get_freqmap(
         try:
             offset = numpy.broadcast_to(offset, (len(rp), 6))
         except ValueError:
-            msg = ('offset and refpts have incoherent '
-                   'shapes: {0}, {1}'.format(numpy.shape(offset),
-                                             numpy.shape(refpts)))
+            msg = "offset and refpts have incoherent " "shapes: {0}, {1}".format(
+                numpy.shape(offset), numpy.shape(refpts)
+            )
             raise AtError(msg)
     else:
-        _, offset = find_orbit(ring,refpts)
+        _, offset = find_orbit(ring, refpts)
         planesi = numpy.atleast_1d(get_plane_index(planes))
-        offset[:,planesi[0]] =  offset[:,planesi[0]] + 1e-9
-        offset[:,planesi[1]] =  offset[:,planesi[1]] + 1e-9
+        offset[:, planesi[0]] = offset[:, planesi[0]] + 1e-9
+        offset[:, planesi[1]] = offset[:, planesi[1]] + 1e-9
     dataobs = []
     t0 = time.time()
     for r, o in zip(rp, offset):
-        obspt=r
-        dp=dp
-        offset=o
-        offset, newring = set_ring_orbit(ring, dp, obspt,
-                                         offset)
-        config = grid_configuration(planes, npoints, amplitudes,
-                                    grid_mode, bounds=bounds,
-                                    shift_zero=shift_zero)
+        obspt = r
+        dp = dp
+        offset = o
+        offset, newring = set_ring_orbit(ring, dp, obspt, offset)
+        config = grid_configuration(
+            planes, npoints, amplitudes, grid_mode, bounds=bounds, shift_zero=shift_zero
+        )
         if verbose:
-            print('\nRunning grid frequency search:')
+            print("\nRunning grid frequency search:")
             if obspt is None:
-                print('Element {0}, obspt={1}'.format(ring[0].FamName, 0))
+                print("Element {0}, obspt={1}".format(ring[0].FamName, 0))
             else:
-                print('Element {0}, obspt={1}'.format(ring[obspt].FamName,
-                                                      obspt))
-            print('The grid mode is {0}'.format(config.mode))
-            print('The planes are {0}'.format(config.planes))
-            print('Number of steps are {0}'.format(config.shape))
-            print('The maximum amplitudes are {0}'.format(config.amplitudes))
-            print('The maximum boundaries are {0}'.format(config.bounds))
-            print('Number of turns is {0}'.format(nturns))
-            print('The initial offset is {0} with dp={1}'.format(offset, dp))
+                print("Element {0}, obspt={1}".format(ring[obspt].FamName, obspt))
+            print("The grid mode is {0}".format(config.mode))
+            print("The planes are {0}".format(config.planes))
+            print("Number of steps are {0}".format(config.shape))
+            print("The maximum amplitudes are {0}".format(config.amplitudes))
+            print("The maximum boundaries are {0}".format(config.bounds))
+            print("Number of turns is {0}".format(nturns))
+            print("The initial offset is {0} with dp={1}".format(offset, dp))
         parts, grid = get_parts(config, offset)
-        rout, tp, td = ring.track(parts, nturns=2*nturns, losses=True, use_mp=use_mp, **kwargs)
+        rout, tp, td = ring.track(
+            parts, nturns=2 * nturns, losses=True, use_mp=use_mp, **kwargs
+        )
 
         mask = get_survived(parts, newring, nturns, use_mp, **kwargs)
         survived = grid.grid[:, mask]
 
         planesi = numpy.atleast_1d(get_plane_index(planes))
-        tunes = numpy.zeros((len(planesi),2,len(numpy.where(mask)[0])))
-
+        tunes = numpy.zeros((len(planesi), 2, len(numpy.where(mask)[0])))
 
         for i in range(len(planesi)):
-            tunes[i,0] = get_tunes_harmonic(rout[planesi[i],mask,:,     0:  nturns],use_mp=use_mp, **kwargs)
-            tunes[i,1] = get_tunes_harmonic(rout[planesi[i],mask,:,nturns:2*nturns],use_mp=use_mp, **kwargs)
+            tunes[i, 0] = get_tunes_harmonic(
+                rout[planesi[i], mask, :, 0:nturns], use_mp=use_mp, **kwargs
+            )
+            tunes[i, 1] = get_tunes_harmonic(
+                rout[planesi[i], mask, :, nturns : 2 * nturns], use_mp=use_mp, **kwargs
+            )
         # metric
-        diffplane1 = tunes[0,0,:] - tunes[0,1,:]
-        diffplane2 = tunes[1,0,:] - tunes[1,1,:]
-        nudiff = 0.5*numpy.log10( (diffplane1*diffplane1 + diffplane2*diffplane2) /nturns )
+        diffplane1 = tunes[0, 0, :] - tunes[0, 1, :]
+        diffplane2 = tunes[1, 0, :] - tunes[1, 1, :]
+        nudiff = 0.5 * numpy.log10(
+            (diffplane1 * diffplane1 + diffplane2 * diffplane2) / nturns
+        )
         # set min-max
-        nudiff = numpy.clip(nudiff,-10,-2)
+        nudiff = numpy.clip(nudiff, -10, -2)
 
-        #return rout,tp,td,grid,tunes
+        # return rout,tp,td,grid,tunes
         firstturns = 0
-        dataobs.append((numpy.vstack((survived, tunes[0,firstturns], tunes[1,firstturns], diffplane1, diffplane2, nudiff)).T, grid, td))
+        dataobs.append(
+            (
+                numpy.vstack(
+                    (
+                        survived,
+                        tunes[0, firstturns],
+                        tunes[1, firstturns],
+                        diffplane1,
+                        diffplane2,
+                        nudiff,
+                    )
+                ).T,
+                grid,
+                td,
+            )
+        )
     if verbose:
-        print('Calculation took {0}'.format(time.time()-t0))
+        print("Calculation took {0}".format(time.time() - t0))
 
     return dataobs
 
 
-def fmap_parallel_track(ring,
-                        coords=[-10, 10, -10, 10],
-                        steps=[100, 100],
-                        scale='linear',
-                        turns=512,
-                        orbit=None,
-                        add_offset6D=numpy.zeros(6),
-                        verbose=False,
-                        lossmap=False,
-                        **kwargs
-                        ):
+def fmap_parallel_track(
+    ring,
+    coords=[-10, 10, -10, 10],
+    steps=[100, 100],
+    scale="linear",
+    turns=512,
+    orbit=None,
+    add_offset6D=numpy.zeros(6),
+    verbose=False,
+    lossmap=False,
+    **kwargs,
+):
     r"""Computes frequency maps
 
     This function calculates the norm of the transverse tune variation per turn
@@ -300,17 +320,17 @@ def fmap_parallel_track(ring,
     """
 
     # https://github.com/atcollab/at/pull/608
-    kwargs.setdefault('pool_size', None)
+    kwargs.setdefault("pool_size", None)
 
     # https://github.com/atcollab/at/pull/608
-    if 'ncpu' in kwargs:
-        warn(AtWarning('ncpu argument is deprecated; use pool_size instead'))
-        kwargs['pool_size'] = kwargs.pop('ncpu')
+    if "ncpu" in kwargs:
+        warn(AtWarning("ncpu argument is deprecated; use pool_size instead"))
+        kwargs["pool_size"] = kwargs.pop("ncpu")
 
     if orbit is None:
         # closed orbit values are not returned. It seems not necessary here
         orbit, _ = find_orbit(ring)
-        print(f'Closed orbit:\t{orbit}')
+        print(f"Closed orbit:\t{orbit}")
 
     # verboseprint to check flag only once
     verboseprint = print if verbose else lambda *a, **k: None
@@ -320,7 +340,7 @@ def fmap_parallel_track(ring,
     # nturns is twice tns in order to get the tune in the first and second
     #     part of the tracking
     tns = turns
-    nturns = 2*tns
+    nturns = 2 * tns
 
     # scale the inumpyut coordinates
     xscale = 1e-3
@@ -332,7 +352,7 @@ def fmap_parallel_track(ring,
 
     # verify steps
     if numpy.count_nonzero(steps) != 2:
-        raise ValueError('steps can not be zero')
+        raise ValueError("steps can not be zero")
 
     xmin = numpy.minimum(coords[0], coords[1])
     xmax = numpy.maximum(coords[0], coords[1])
@@ -343,28 +363,26 @@ def fmap_parallel_track(ring,
 
     # define rectangle and x,y step size
     if scale == "nonlinear":
-        verboseprint('non linear steps')
-        xinterval = 1.0*(xmax - xmin)
-        yinterval = 1.0*(ymax - ymin)
-        xauxsteps = numpy.arange(-0.5, 0.0 + 1e-9, 1./xsteps)
-        yauxsteps = numpy.arange(-0.5, 0.0 + 1e-9, 1./ysteps)
+        verboseprint("non linear steps")
+        xinterval = 1.0 * (xmax - xmin)
+        yinterval = 1.0 * (ymax - ymin)
+        xauxsteps = numpy.arange(-0.5, 0.0 + 1e-9, 1.0 / xsteps)
+        yauxsteps = numpy.arange(-0.5, 0.0 + 1e-9, 1.0 / ysteps)
         xnrsteps = 10**xauxsteps
         xnrsteps = xnrsteps / sum(xnrsteps)
         ynrsteps = 10**yauxsteps
         ynrsteps = ynrsteps / sum(ynrsteps)
-        xscalesteps = numpy.concatenate((xnrsteps,
-                                        numpy.flip(xnrsteps)), axis=None)
-        yscalesteps = numpy.concatenate((ynrsteps,
-                                        numpy.flip(ynrsteps)), axis=None)
+        xscalesteps = numpy.concatenate((xnrsteps, numpy.flip(xnrsteps)), axis=None)
+        yscalesteps = numpy.concatenate((ynrsteps, numpy.flip(ynrsteps)), axis=None)
         ixarray = xmin + 0.5 * xinterval * numpy.cumsum(xscalesteps)
         iyarray = ymin + 0.5 * yinterval * numpy.cumsum(yscalesteps)
     else:  # scale == 'linear', ignore any other
-        verboseprint('linear steps')
+        verboseprint("linear steps")
         # get the intervals
-        xstep = 1.0*(xmax - xmin)/xsteps
-        ystep = 1.0*(ymax - ymin)/ysteps
-        ixarray = numpy.arange(xmin, xmax+1e-6, xstep)
-        iyarray = numpy.arange(ymin, ymax+1e-6, ystep)
+        xstep = 1.0 * (xmax - xmin) / xsteps
+        ystep = 1.0 * (ymax - ymin) / ysteps
+        ixarray = numpy.arange(xmin, xmax + 1e-6, xstep)
+        iyarray = numpy.arange(ymin, ymax + 1e-6, ystep)
     lenixarray = len(ixarray)
     leniyarray = len(iyarray)
 
@@ -372,20 +390,19 @@ def fmap_parallel_track(ring,
 
     # tracking in parallel multiple x coordinates with the same y coordinate
     for iy, iy_index in zip(iyarray, range(leniyarray)):
-        print(f'Tracked particles {abs(-100.0*iy_index/leniyarray):.1f} %')
+        print(f"Tracked particles {abs(-100.0*iy_index/leniyarray):.1f} %")
         verboseprint("y =", iy)
         z0 = numpy.zeros((lenixarray, 6))  # transposed, and C-aligned
         z0 = z0 + add_offset6D + orbit
         # add 1 nm to tracking to avoid zeros in array for the ideal lattice
-        z0[:, 0] = z0[:, 0] + xscale*ixarray + 1e-9
-        z0[:, 2] = z0[:, 2] + yscale*iy + 1e-9
+        z0[:, 0] = z0[:, 0] + xscale * ixarray + 1e-9
+        z0[:, 2] = z0[:, 2] + yscale * iy + 1e-9
 
         verboseprint("tracking ...")
         # z0.T is Fortran-aligned
         if lossmap:
             # patpass output changes when losses flag is true
-            zOUT, dictloss = \
-                patpass(ring, z0.T, nturns, losses=lossmap, **kwargs)
+            zOUT, dictloss = patpass(ring, z0.T, nturns, losses=lossmap, **kwargs)
             loss_map_array = numpy.append(loss_map_array, dictloss)
         else:
             zOUT = patpass(ring, z0.T, nturns, **kwargs)
@@ -410,7 +427,7 @@ def fmap_parallel_track(ring,
             # pxfirst = pxfirst - numpy.mean(pxfirst)
             # xfirstpart = xfirst + 1j*pxfirst
             # get the last turns in x
-            xlast = z1[0, tns:2*tns]
+            xlast = z1[0, tns : 2 * tns]
             xlast = xlast - numpy.mean(xlast)
             # pxlast = z1[1, tns:2*tns]
             # pxlast = pxlast - numpy.mean(pxlast)
@@ -423,7 +440,7 @@ def fmap_parallel_track(ring,
             # pyfirst = pyfirst - numpy.mean(pyfirst)
             # yfirstpart = yfirst + 1j*pyfirst
             # get the last turns in y
-            ylast = z1[2, tns:2*tns]
+            ylast = z1[2, tns : 2 * tns]
             ylast = ylast - numpy.mean(ylast)
             # pylast = z1[3, tns:2*tns]
             # pylast = pylast - numpy.mean(pylast)
@@ -450,32 +467,42 @@ def fmap_parallel_track(ring,
             if len(yfreqlast) == 0:
                 verboseprint("  No frequency")
                 continue
-            verboseprint("NAFF results, (x,y)=",
-                         ixarray[ix_index],
-                         iy,
-                         "\nH freq. first part =\t", xfreqfirst[0],
-                         "\nH freq. last part =\t", xfreqlast[0],
-                         "\nV freq. first part =\t", yfreqfirst[0],
-                         "\nV freq. last part =\t", yfreqlast[0])
+            verboseprint(
+                "NAFF results, (x,y)=",
+                ixarray[ix_index],
+                iy,
+                "\nH freq. first part =\t",
+                xfreqfirst[0],
+                "\nH freq. last part =\t",
+                xfreqlast[0],
+                "\nV freq. first part =\t",
+                yfreqfirst[0],
+                "\nV freq. last part =\t",
+                yfreqlast[0],
+            )
 
             # metric
             xdiff = xfreqlast[0] - xfreqfirst[0]
             ydiff = yfreqlast[0] - yfreqfirst[0]
-            nudiff = 0.5*numpy.log10((xdiff*xdiff + ydiff*ydiff)/tns)
+            nudiff = 0.5 * numpy.log10((xdiff * xdiff + ydiff * ydiff) / tns)
             # min max diff
             if nudiff > -2:
                 nudiff = -2
             if nudiff < -10:
                 nudiff = -10
             # save diff
-            xy_nuxy_lognudiff_array = numpy.append(xy_nuxy_lognudiff_array,
-                                                   [ixarray[ix_index],
-                                                    iy,
-                                                    xfreqfirst[0],
-                                                    yfreqfirst[0],
-                                                    xdiff,
-                                                    ydiff,
-                                                    nudiff])
+            xy_nuxy_lognudiff_array = numpy.append(
+                xy_nuxy_lognudiff_array,
+                [
+                    ixarray[ix_index],
+                    iy,
+                    xfreqfirst[0],
+                    yfreqfirst[0],
+                    xdiff,
+                    ydiff,
+                    nudiff,
+                ],
+            )
 
     # first element is garbage
     xy_nuxy_lognudiff_array = numpy.delete(xy_nuxy_lognudiff_array, 0)
@@ -483,4 +510,6 @@ def fmap_parallel_track(ring,
     xy_nuxy_lognudiff_array = xy_nuxy_lognudiff_array.reshape(-1, 7)
 
     return xy_nuxy_lognudiff_array, loss_map_array
+
+
 # the end

From b9281e6fdf21607fbf4926d7c92924007e62513a Mon Sep 17 00:00:00 2001
From: oscarxblanco <oblanco@cells.es>
Date: Fri, 28 Jun 2024 18:35:36 +0200
Subject: [PATCH 5/9] solving flake8 errors

---
 pyat/at/physics/frequency_maps.py | 219 ++++++++++++++----------------
 1 file changed, 104 insertions(+), 115 deletions(-)

diff --git a/pyat/at/physics/frequency_maps.py b/pyat/at/physics/frequency_maps.py
index 7b4d8de6ed..e77cefc8a3 100644
--- a/pyat/at/physics/frequency_maps.py
+++ b/pyat/at/physics/frequency_maps.py
@@ -1,8 +1,8 @@
 """
-Frequency analysis (FMAP) using .harmonic_analysis lib
-and pyat parallel tracking (patpass)
+Frequency analysis (FMAP).
 """
 
+
 # orblancog
 # generates the frequency and diffusion map for a given ring
 # 2024apr16 create get_freqmap
@@ -10,46 +10,47 @@
 # 2022jun07 serial version
 # 2024jun21 adds get_freqmap
 
-from at.tracking import patpass
-from at.physics import find_orbit
-import numpy
+
+from __future__ import annotations
+
 from warnings import warn
-from at.lattice import AtWarning
+import multiprocessing
+from typing import Optional, Sequence
+import time
+import numpy
+
+from ..lattice import Lattice, Refpts, AtError, AtWarning
+from ..tracking import patpass
+from ..physics import find_orbit
 from ..acceptance.boundary import set_ring_orbit
 from ..acceptance.boundary import grid_configuration
 from ..acceptance.boundary import get_parts
 from ..acceptance.boundary import get_survived
 from ..acceptance.boundary import GridMode
 from ..acceptance.boundary import get_plane_index
-from at.lattice import Lattice, AtError, AtWarning
-from typing import Optional, Sequence
-from ..lattice import Lattice, Refpts, frequency_control, AtError
-import multiprocessing
-import time
-
-_pdict = {"x": 0, "xp": 1, "y": 2, "yp": 3, "dp": 4, "ct": 5}
-
 
 # Jaime Coello de Portugal (JCdP) frequency analysis implementation
 from .harmonic_analysis import get_tunes_harmonic
 
+_pdict = {"x": 0, "xp": 1, "y": 2, "yp": 3, "dp": 4, "ct": 5}
+
+
 __all__ = ["fmap_parallel_track", "get_freqmap"]
 
 
 def get_freqmap(
     ring: Lattice,
-    planes,
-    npoints,
-    amplitudes,
-    bounds=None,
+    planes: list,
+    npoints: numpy.ndarray,
+    amplitudes: numpy.ndarray,
+    bounds: any = None,
     nturns: Optional[int] = 512,
-    dp: Optional[float] = None,
+    deltap: Optional[float] = None,
     offset: Sequence[float] = None,
     refpts: Optional[Refpts] = 0,
     grid_mode: Optional[GridMode] = GridMode.CARTESIAN,
     use_mp: Optional[bool] = True,
     verbose: Optional[bool] = False,
-    lossmap: Optional[int] = 2,
     shift_zero: Optional[float] = 0.0e-6,
     start_method: Optional[str] = None,
 ):
@@ -61,29 +62,31 @@ def get_freqmap(
         planes:         max. dimension 2, Plane(s) to scan for the acceptance.
           Allowed values are: ``'x'``, ``'xp'``, ``'y'``,
           ``'yp'``, ``'dp'``, ``'ct'``
+          e.g. ['x','y']
         npoints:        (len(planes),) array: number of points in each
           dimension
+          e.g. numpy.array([200,200])
         amplitudes:     (len(planes),) array: set the search range:
-
           * :py:attr:`GridMode.CARTESIAN/RADIAL <.GridMode.RADIAL>`:
             max. amplitude
           * :py:attr:`.GridMode.RECURSIVE`: initial step
-        nturns:         Number of turns for the tracking
+          e.g. numpy.array([10e-3, 10e-3])
+    Keyword arguments:
+        nturns:         Number of turns for the tracking. Default 512
         refpts:         Observation points. Default: start of the machine
-        dp:             static momentum offset
-        offset:         initial orbit. Default: closed orbit
+        deltap:             static momentum offset
+        offset:         (len(refpts),6) array: initial orbit. Default: closed orbit
         bounds:         defines the tracked range: range=bounds*amplitude.
           It can be use to select quadrants. For example, default values are:
-
           * :py:attr:`.GridMode.CARTESIAN`: ((-1, 1), (0, 1))
           * :py:attr:`GridMode.RADIAL/RECURSIVE <.GridMode.RADIAL>`: ((0, 1),
             (:math:`\pi`, 0))
         grid_mode:      defines the evaluation grid:
-
           * :py:attr:`.GridMode.CARTESIAN`: full [:math:`\:x, y\:`] grid
           * :py:attr:`.GridMode.RADIAL`: full [:math:`\:r, \theta\:`] grid
           * :py:attr:`.GridMode.RECURSIVE`: radial recursive search
-        use_mp:         Use python multiprocessing (:py:func:`.patpass`,
+        use_mp:         Default True.
+          Use python multiprocessing (:py:func:`.patpass`,
           default use :py:func:`.lattice_pass`).
         verbose:        Print out some information
         divider:        Value of the divider used in
@@ -98,24 +101,28 @@ def get_freqmap(
           considered unsafe.
 
     Returns:
-        fmaout:   list containing the result of every reference point.
-            Inside,
-              1) the frequency map
-                `[xcoor, ycoor, nux, ny, dnux, dnuy,
-                 log10(sqrt(sum(dnu**2)/turns)) ]`
+        fmaout:  (len(refpts)) list: containing the result of every reference point.
+            Per every reference point there is a tuple containing,
+              1) (npoints*npoints, 7) array with the frequency map,
+                 each row has
+                  `[xcoor, ycoor, nux, ny, dnux, dnuy,
+                   log10(sqrt(sum(dnu**2)/turns)) ]`
               2) the grid
               3) the dictionary of particle losses
 
-    In case of multiple refpts, return values are lists of arrays, with one
-    array per ref. point.
 
     Examples:
 
-        >>> bf,sf,gf = ring.get_acceptance(planes, npoints, amplitudes)
-        >>> plt.plot(*gf,'.')
-        >>> plt.plot(*sf,'.')
-        >>> plt.plot(*bf)
-        >>> plt.show()
+        >>> fmaout = at.get_freqmap(ring,
+                     ['x','y'],
+                     numpy.array([200,200]),
+                     numpy.array([10e-3,10e-3]),
+                     nturns=512,
+                     refpts = numpy.array([0,1,2]),
+                     bounds=[[-1,1],[-1,1]],
+                     verbose=True,
+                     )
+        >>> fmadata_at_start = fmaout[0][0]
 
     .. note::
 
@@ -129,69 +136,57 @@ def get_freqmap(
 
     if verbose:
         nproc = multiprocessing.cpu_count()
-        print("\n{0} cpu found for acceptance calculation".format(nproc))
+        print(f"\n{0} cpu found for acceptance calculation".format(nproc))
         if use_mp:
-            nprocu = nproc
             print("Multi-process acceptance calculation selected...")
             if nproc == 1:
                 print("Consider use_mp=False for single core computations")
         else:
-            nprocu = 1
             print("Single process acceptance calculation selected...")
             if nproc > 1:
                 print("Consider use_mp=True for parallelized computations")
-        np = numpy.atleast_1d(npoints)
-        na = 2
-        if len(np) == 2:
-            na = np[1]
-        npp = numpy.prod(npoints)
-        rpp = 2 * numpy.ceil(numpy.log2(np[0])) * numpy.ceil(na / nprocu)
-        mpp = npp / nprocu
 
     survived = []
     grid = []
     if refpts is not None:
-        rp = ring.uint32_refpts(refpts)
+        rps = ring.uint32_refpts(refpts)
     else:
-        rp = numpy.atleast_1d(refpts)
+        rps = numpy.atleast_1d(refpts)
     if offset is not None:
         try:
-            offset = numpy.broadcast_to(offset, (len(rp), 6))
-        except ValueError:
-            msg = "offset and refpts have incoherent " "shapes: {0}, {1}".format(
+            offset = numpy.broadcast_to(offset, (len(rps), 6))
+        except ValueError as incoherent_shape:
+            msg = f"offset and refpts have incoherent shapes: {0}, {1}".format(
                 numpy.shape(offset), numpy.shape(refpts)
             )
-            raise AtError(msg)
+            raise AtError(msg) from incoherent_shape
     else:
         _, offset = find_orbit(ring, refpts)
         planesi = numpy.atleast_1d(get_plane_index(planes))
         offset[:, planesi[0]] = offset[:, planesi[0]] + 1e-9
         offset[:, planesi[1]] = offset[:, planesi[1]] + 1e-9
     dataobs = []
-    t0 = time.time()
-    for r, o in zip(rp, offset):
-        obspt = r
-        dp = dp
-        offset = o
-        offset, newring = set_ring_orbit(ring, dp, obspt, offset)
+    t00 = time.time()
+    for obspt, offset0 in zip(rps, offset):
+        offset, newring = set_ring_orbit(ring, deltap, obspt, offset0)
         config = grid_configuration(
             planes, npoints, amplitudes, grid_mode, bounds=bounds, shift_zero=shift_zero
         )
         if verbose:
             print("\nRunning grid frequency search:")
             if obspt is None:
-                print("Element {0}, obspt={1}".format(ring[0].FamName, 0))
+                print(f"Element {0}, obspt={1}".format(ring[0].FamName, 0))
             else:
-                print("Element {0}, obspt={1}".format(ring[obspt].FamName, obspt))
-            print("The grid mode is {0}".format(config.mode))
-            print("The planes are {0}".format(config.planes))
-            print("Number of steps are {0}".format(config.shape))
-            print("The maximum amplitudes are {0}".format(config.amplitudes))
-            print("The maximum boundaries are {0}".format(config.bounds))
-            print("Number of turns is {0}".format(nturns))
-            print("The initial offset is {0} with dp={1}".format(offset, dp))
+                print(f"Element {0}, obspt={1}".format(ring[obspt].FamName, obspt))
+            print(f"The grid mode is {0}".format(config.mode))
+            print(f"The planes are {0}".format(config.planes))
+            print(f"Number of steps are {0}".format(config.shape))
+            print(f"The maximum amplitudes are {0}".format(config.amplitudes))
+            print(f"The maximum boundaries are {0}".format(config.bounds))
+            print(f"Number of turns is {0}".format(nturns))
+            print(f"The initial offset is {0} with deltap={1}".format(offset, deltap))
         parts, grid = get_parts(config, offset)
-        rout, tp, td = ring.track(
+        rout, _, tdl = ring.track(
             parts, nturns=2 * nturns, losses=True, use_mp=use_mp, **kwargs
         )
 
@@ -201,12 +196,12 @@ def get_freqmap(
         planesi = numpy.atleast_1d(get_plane_index(planes))
         tunes = numpy.zeros((len(planesi), 2, len(numpy.where(mask)[0])))
 
-        for i in range(len(planesi)):
+        for i, theplane in enumerate(planesi):
             tunes[i, 0] = get_tunes_harmonic(
-                rout[planesi[i], mask, :, 0:nturns], use_mp=use_mp, **kwargs
+                rout[theplane, mask, :, 0:nturns], use_mp=use_mp, **kwargs
             )
             tunes[i, 1] = get_tunes_harmonic(
-                rout[planesi[i], mask, :, nturns : 2 * nturns], use_mp=use_mp, **kwargs
+                rout[theplane, mask, :, nturns : 2 * nturns], use_mp=use_mp, **kwargs
             )
         # metric
         diffplane1 = tunes[0, 0, :] - tunes[0, 1, :]
@@ -217,7 +212,7 @@ def get_freqmap(
         # set min-max
         nudiff = numpy.clip(nudiff, -10, -2)
 
-        # return rout,tp,td,grid,tunes
+        # return rout,tp,tdl,grid,tunes
         firstturns = 0
         dataobs.append(
             (
@@ -232,29 +227,29 @@ def get_freqmap(
                     )
                 ).T,
                 grid,
-                td,
+                tdl,
             )
         )
     if verbose:
-        print("Calculation took {0}".format(time.time() - t0))
+        print(f"Calculation took {0}".format(time.time() - t00))
 
     return dataobs
 
 
 def fmap_parallel_track(
-    ring,
+    ring: Lattice,
     coords=[-10, 10, -10, 10],
-    steps=[100, 100],
-    scale="linear",
-    turns=512,
-    orbit=None,
-    add_offset6D=numpy.zeros(6),
-    verbose=False,
-    lossmap=False,
-    **kwargs,
+    steps: list = [100, 100],
+    scale: str = "linear",
+    turns: int = 512,
+    orbit: numpy.ndarray = None,
+    add_offset6D: numpy.ndarray = numpy.zeros(6),
+    add_offset6d: numpy.ndarray = numpy.zeros(6),
+    verbose: bool = False,
+    lossmap: bool = False,
+    **kwargs: dict[str, any],
 ):
-    r"""Computes frequency maps
-
+    r"""Computes frequency maps.
     This function calculates the norm of the transverse tune variation per turn
     for a particle tracked along a ring with a set of offsets in the initial
     coordinates.
@@ -282,7 +277,7 @@ def fmap_parallel_track(
     The closed orbit is calculated and added to the
     initial particle offset of every particle. Otherwise, one could set
     ``orbit = numpy.zeros(6)`` to avoid it.
-    Additionally, a numpy array (*add_offset6D*) with 6 values could be used to
+    Additionally, a numpy array (*add_offset6d*) with 6 values could be used to
     arbitrarily offset the initial coordinates of every particle.
 
     A dictionary with particle losses is saved for every vertical offset.
@@ -296,7 +291,7 @@ def fmap_parallel_track(
         turns:    default 512
         orbit:    If :py:obj:`None`, the closed orbit is computed and added to
           the coordinates
-        add_offset6D: default numpy.zeros((6,1))
+        add_offset6d: default numpy.zeros((6,1))
         verbose:  prints additional info
         lossmap:  default false
     Optional:
@@ -327,6 +322,12 @@ def fmap_parallel_track(
         warn(AtWarning("ncpu argument is deprecated; use pool_size instead"))
         kwargs["pool_size"] = kwargs.pop("ncpu")
 
+    # deprecating add_offset6D because does not comply with the cammel name standard
+    if "add_offset6D" in kwargs:
+        warn(AtWarning("add_offset6D argument is deprecated; use add_offset6d instead"))
+        kwargs["add_offset6d"] = kwargs.pop("add_offset6D")
+    add_offset6d = kwargs.pop("add_offset6d", numpy.zeros(6))
+
     if orbit is None:
         # closed orbit values are not returned. It seems not necessary here
         orbit, _ = find_orbit(ring)
@@ -392,59 +393,47 @@ def fmap_parallel_track(
     for iy, iy_index in zip(iyarray, range(leniyarray)):
         print(f"Tracked particles {abs(-100.0*iy_index/leniyarray):.1f} %")
         verboseprint("y =", iy)
-        z0 = numpy.zeros((lenixarray, 6))  # transposed, and C-aligned
-        z0 = z0 + add_offset6D + orbit
+        z00 = numpy.zeros((lenixarray, 6))  # transposed, and C-aligned
+        z00 = z00 + add_offset6d + orbit
         # add 1 nm to tracking to avoid zeros in array for the ideal lattice
-        z0[:, 0] = z0[:, 0] + xscale * ixarray + 1e-9
-        z0[:, 2] = z0[:, 2] + yscale * iy + 1e-9
+        z00[:, 0] = z00[:, 0] + xscale * ixarray + 1e-9
+        z00[:, 2] = z00[:, 2] + yscale * iy + 1e-9
 
         verboseprint("tracking ...")
-        # z0.T is Fortran-aligned
+        # z00.T is Fortran-aligned
         if lossmap:
             # patpass output changes when losses flag is true
-            zOUT, dictloss = patpass(ring, z0.T, nturns, losses=lossmap, **kwargs)
+            zout, dictloss = patpass(ring, z00.T, nturns, losses=lossmap, **kwargs)
             loss_map_array = numpy.append(loss_map_array, dictloss)
         else:
-            zOUT = patpass(ring, z0.T, nturns, **kwargs)
+            zout = patpass(ring, z00.T, nturns, **kwargs)
 
         # start of serial frequency analysis
         for ix_index in range(lenixarray):  # cycle over the track results
             # check if nan in arrays
-            array_sum = numpy.sum(zOUT[:, ix_index, 0])
+            array_sum = numpy.sum(zout[:, ix_index, 0])
             array_has_nan = numpy.isnan(array_sum)
             if array_has_nan:
                 verboseprint("array has nan")
                 continue
 
             # get one valid particle
-            z1 = zOUT[:, ix_index, 0]
+            z11 = zout[:, ix_index, 0]
 
             # remove mean values
             # get the first turn in x
-            xfirst = z1[0, 0:tns]
+            xfirst = z11[0, 0:tns]
             xfirst = xfirst - numpy.mean(xfirst)
-            # pxfirst = z1[1, 0:tns]
-            # pxfirst = pxfirst - numpy.mean(pxfirst)
-            # xfirstpart = xfirst + 1j*pxfirst
             # get the last turns in x
-            xlast = z1[0, tns : 2 * tns]
+            xlast = z11[0, tns : 2 * tns]
             xlast = xlast - numpy.mean(xlast)
-            # pxlast = z1[1, tns:2*tns]
-            # pxlast = pxlast - numpy.mean(pxlast)
-            # xlastpart = xlast + 1j*pxlast
 
             # get the first turn in y
-            yfirst = z1[2, 0:tns]
+            yfirst = z11[2, 0:tns]
             yfirst = yfirst - numpy.mean(yfirst)
-            # pyfirst = z1[3, 0:tns]
-            # pyfirst = pyfirst - numpy.mean(pyfirst)
-            # yfirstpart = yfirst + 1j*pyfirst
             # get the last turns in y
-            ylast = z1[2, tns : 2 * tns]
+            ylast = z11[2, tns : 2 * tns]
             ylast = ylast - numpy.mean(ylast)
-            # pylast = z1[3, tns:2*tns]
-            # pylast = pylast - numpy.mean(pylast)
-            # ylastpart = ylast + 1j*pylast
 
             # calc frequency from array,
             # jump the cycle is no frequency is found

From 797ef7bb0d251d4d8d28f9afb39bdda32220e70d Mon Sep 17 00:00:00 2001
From: oscarxblanco <oblanco@cells.es>
Date: Fri, 28 Jun 2024 19:01:48 +0200
Subject: [PATCH 6/9] solving flake8 and isort

---
 pyat/at/physics/frequency_maps.py | 42 +++++++++++++------------------
 1 file changed, 17 insertions(+), 25 deletions(-)

diff --git a/pyat/at/physics/frequency_maps.py b/pyat/at/physics/frequency_maps.py
index e77cefc8a3..a42efd8328 100644
--- a/pyat/at/physics/frequency_maps.py
+++ b/pyat/at/physics/frequency_maps.py
@@ -1,6 +1,4 @@
-"""
-Frequency analysis (FMAP).
-"""
+""" Frequency analysis (FMAP). """
 
 
 # orblancog
@@ -13,22 +11,19 @@
 
 from __future__ import annotations
 
-from warnings import warn
 import multiprocessing
-from typing import Optional, Sequence
 import time
+from typing import Optional, Sequence
+from warnings import warn
+
 import numpy
 
-from ..lattice import Lattice, Refpts, AtError, AtWarning
-from ..tracking import patpass
+from ..acceptance.boundary import (GridMode, get_parts, get_plane_index,
+                                   get_survived, grid_configuration,
+                                   set_ring_orbit)
+from ..lattice import AtError, AtWarning, Lattice, Refpts
 from ..physics import find_orbit
-from ..acceptance.boundary import set_ring_orbit
-from ..acceptance.boundary import grid_configuration
-from ..acceptance.boundary import get_parts
-from ..acceptance.boundary import get_survived
-from ..acceptance.boundary import GridMode
-from ..acceptance.boundary import get_plane_index
-
+from ..tracking import patpass
 # Jaime Coello de Portugal (JCdP) frequency analysis implementation
 from .harmonic_analysis import get_tunes_harmonic
 
@@ -212,7 +207,7 @@ def get_freqmap(
         # set min-max
         nudiff = numpy.clip(nudiff, -10, -2)
 
-        # return rout,tp,tdl,grid,tunes
+        # now return rout,tp,tdl,grid,tunes
         firstturns = 0
         dataobs.append(
             (
@@ -248,7 +243,7 @@ def fmap_parallel_track(
     verbose: bool = False,
     lossmap: bool = False,
     **kwargs: dict[str, any],
-):
+) -> numpy.ndarray:
     r"""Computes frequency maps.
     This function calculates the norm of the transverse tune variation per turn
     for a particle tracked along a ring with a set of offsets in the initial
@@ -285,6 +280,7 @@ def fmap_parallel_track(
 
     Parameters:
         ring:     a valid pyat ring
+    Optional:
         coords:   default [-10,10,-10,10] in mm
         steps:    (xsteps, ysteps): number of steps in each plane
         scale:    default 'linear'; or 'non-linear'
@@ -294,7 +290,6 @@ def fmap_parallel_track(
         add_offset6d: default numpy.zeros((6,1))
         verbose:  prints additional info
         lossmap:  default false
-    Optional:
         pool_size:number of processes. See :py:func:`.patpass`
 
     Returns:
@@ -390,14 +385,14 @@ def fmap_parallel_track(
     print("Start tracking and frequency analysis")
 
     # tracking in parallel multiple x coordinates with the same y coordinate
-    for iy, iy_index in zip(iyarray, range(leniyarray)):
+    for iiy, iy_index in zip(iyarray, range(leniyarray)):
         print(f"Tracked particles {abs(-100.0*iy_index/leniyarray):.1f} %")
-        verboseprint("y =", iy)
+        verboseprint("y =", iiy)
         z00 = numpy.zeros((lenixarray, 6))  # transposed, and C-aligned
         z00 = z00 + add_offset6d + orbit
         # add 1 nm to tracking to avoid zeros in array for the ideal lattice
         z00[:, 0] = z00[:, 0] + xscale * ixarray + 1e-9
-        z00[:, 2] = z00[:, 2] + yscale * iy + 1e-9
+        z00[:, 2] = z00[:, 2] + yscale * iiy + 1e-9
 
         verboseprint("tracking ...")
         # z00.T is Fortran-aligned
@@ -441,17 +436,14 @@ def fmap_parallel_track(
             if len(xfreqfirst) == 0:
                 verboseprint("  No frequency")
                 continue
-            # xfreqlast  = PyNAFF.naff(xlastpart,tns,1,0,False)
             xfreqlast = get_tunes_harmonic(xlast)
             if len(xfreqlast) == 0:
                 verboseprint("  No frequency")
                 continue
-            # yfreqfirst = PyNAFF.naff(yfirstpart,tns,1,0,False)
             yfreqfirst = get_tunes_harmonic(yfirst)
             if len(yfreqfirst) == 0:
                 verboseprint("  No frequency")
                 continue
-            # yfreqlast  = PyNAFF.naff(ylastpart,tns,1,0,False)
             yfreqlast = get_tunes_harmonic(ylast)
             if len(yfreqlast) == 0:
                 verboseprint("  No frequency")
@@ -459,7 +451,7 @@ def fmap_parallel_track(
             verboseprint(
                 "NAFF results, (x,y)=",
                 ixarray[ix_index],
-                iy,
+                iiy,
                 "\nH freq. first part =\t",
                 xfreqfirst[0],
                 "\nH freq. last part =\t",
@@ -484,7 +476,7 @@ def fmap_parallel_track(
                 xy_nuxy_lognudiff_array,
                 [
                     ixarray[ix_index],
-                    iy,
+                    iiy,
                     xfreqfirst[0],
                     yfreqfirst[0],
                     xdiff,

From 6fb17359c16ff51776b087c4f733937207586a89 Mon Sep 17 00:00:00 2001
From: oscarxblanco <oblanco@cells.es>
Date: Fri, 28 Jun 2024 19:02:21 +0200
Subject: [PATCH 7/9] black

---
 pyat/at/physics/frequency_maps.py | 12 +++++++++---
 1 file changed, 9 insertions(+), 3 deletions(-)

diff --git a/pyat/at/physics/frequency_maps.py b/pyat/at/physics/frequency_maps.py
index a42efd8328..eaa126c7de 100644
--- a/pyat/at/physics/frequency_maps.py
+++ b/pyat/at/physics/frequency_maps.py
@@ -18,12 +18,18 @@
 
 import numpy
 
-from ..acceptance.boundary import (GridMode, get_parts, get_plane_index,
-                                   get_survived, grid_configuration,
-                                   set_ring_orbit)
+from ..acceptance.boundary import (
+    GridMode,
+    get_parts,
+    get_plane_index,
+    get_survived,
+    grid_configuration,
+    set_ring_orbit,
+)
 from ..lattice import AtError, AtWarning, Lattice, Refpts
 from ..physics import find_orbit
 from ..tracking import patpass
+
 # Jaime Coello de Portugal (JCdP) frequency analysis implementation
 from .harmonic_analysis import get_tunes_harmonic
 

From a5d959b11f14cdd90d4a5d4415cc1441bfa2532c Mon Sep 17 00:00:00 2001
From: oscarxblanco <oblanco@cells.es>
Date: Fri, 28 Jun 2024 19:28:58 +0200
Subject: [PATCH 8/9] solving flake8 errors

---
 pyat/at/physics/frequency_maps.py | 34 +++++++++++++++----------------
 1 file changed, 16 insertions(+), 18 deletions(-)

diff --git a/pyat/at/physics/frequency_maps.py b/pyat/at/physics/frequency_maps.py
index eaa126c7de..8cc963aa82 100644
--- a/pyat/at/physics/frequency_maps.py
+++ b/pyat/at/physics/frequency_maps.py
@@ -54,10 +54,10 @@ def get_freqmap(
     verbose: Optional[bool] = False,
     shift_zero: Optional[float] = 0.0e-6,
     start_method: Optional[str] = None,
-):
+    **kwargs: dict[str, any],
+) -> list:
     # noinspection PyUnresolvedReferences
     r"""Computes the frequency map at ``repfts`` observation points.
-
     Parameters:
         ring:           Lattice definition
         planes:         max. dimension 2, Plane(s) to scan for the acceptance.
@@ -131,13 +131,13 @@ def get_freqmap(
          This behavior can be changed by setting
          ``at.DConstant.patpass_poolsize`` to the desired value
     """
-    kwargs = {}
+
     if start_method is not None:
         kwargs["start_method"] = start_method
 
     if verbose:
         nproc = multiprocessing.cpu_count()
-        print(f"\n{0} cpu found for acceptance calculation".format(nproc))
+        print(f"\n{nproc} cpu found for acceptance calculation")
         if use_mp:
             print("Multi-process acceptance calculation selected...")
             if nproc == 1:
@@ -157,8 +157,9 @@ def get_freqmap(
         try:
             offset = numpy.broadcast_to(offset, (len(rps), 6))
         except ValueError as incoherent_shape:
-            msg = f"offset and refpts have incoherent shapes: {0}, {1}".format(
-                numpy.shape(offset), numpy.shape(refpts)
+            msg = (
+                "offset and refpts have incoherent shapes: "
+                f"{numpy.shape(offset)}, {numpy.shape(refpts)}"
             )
             raise AtError(msg) from incoherent_shape
     else:
@@ -175,17 +176,14 @@ def get_freqmap(
         )
         if verbose:
             print("\nRunning grid frequency search:")
-            if obspt is None:
-                print(f"Element {0}, obspt={1}".format(ring[0].FamName, 0))
-            else:
-                print(f"Element {0}, obspt={1}".format(ring[obspt].FamName, obspt))
-            print(f"The grid mode is {0}".format(config.mode))
-            print(f"The planes are {0}".format(config.planes))
-            print(f"Number of steps are {0}".format(config.shape))
-            print(f"The maximum amplitudes are {0}".format(config.amplitudes))
-            print(f"The maximum boundaries are {0}".format(config.bounds))
-            print(f"Number of turns is {0}".format(nturns))
-            print(f"The initial offset is {0} with deltap={1}".format(offset, deltap))
+            print(f"Element {ring[obspt].FamName}, obspt={obspt}")
+            print(f"The grid mode is {config.mode}")
+            print(f"The planes are {config.planes}")
+            print(f"Number of steps are {config.shape}")
+            print(f"The maximum amplitudes are {config.amplitudes}")
+            print(f"The maximum boundaries are {config.bounds}")
+            print(f"Number of turns is {nturns}")
+            print(f"The initial offset is {offset} with deltap={deltap}")
         parts, grid = get_parts(config, offset)
         rout, _, tdl = ring.track(
             parts, nturns=2 * nturns, losses=True, use_mp=use_mp, **kwargs
@@ -232,7 +230,7 @@ def get_freqmap(
             )
         )
     if verbose:
-        print(f"Calculation took {0}".format(time.time() - t00))
+        print(f"Calculation took {time.time() - t00}")
 
     return dataobs
 

From 4952cb6a0caf5c6a5ed7eebad0d0e37eab6553c2 Mon Sep 17 00:00:00 2001
From: oscarxblanco <oblanco@cells.es>
Date: Fri, 16 Aug 2024 20:03:05 +0200
Subject: [PATCH 9/9] solving memory usage

---
 pyat/at/physics/frequency_maps.py | 66 +++++++++++++++++++++++++++++--
 1 file changed, 63 insertions(+), 3 deletions(-)

diff --git a/pyat/at/physics/frequency_maps.py b/pyat/at/physics/frequency_maps.py
index 8cc963aa82..95e0be14ec 100644
--- a/pyat/at/physics/frequency_maps.py
+++ b/pyat/at/physics/frequency_maps.py
@@ -12,6 +12,7 @@
 from __future__ import annotations
 
 import multiprocessing
+import psutil
 import time
 from typing import Optional, Sequence
 from warnings import warn
@@ -39,6 +40,27 @@
 __all__ = ["fmap_parallel_track", "get_freqmap"]
 
 
+def get_ram_info():
+    ram = psutil.virtual_memory()
+
+    total_ram = ram.total # total installed RAM
+    available_ram = ram.available # RAM available for processes
+    used_ram = ram.used # RAM used by processes
+    free_ram = ram.free # RAM not being used
+    percent_used = ram.percent # percentage of RAM used
+
+    return total_ram, available_ram, used_ram, free_ram, percent_used
+def get_swap_info():
+    swap = psutil.swap_memory()
+
+    total_swap = swap.total # total installed RAM
+    used_swap = swap.used # RAM used by processes
+    free_swap = swap.free # RAM not being used
+    percent_used = swap.percent # percentage of RAM used
+
+    return total_swap, used_swap, free_swap, percent_used
+
+
 def get_freqmap(
     ring: Lattice,
     planes: list,
@@ -58,6 +80,7 @@ def get_freqmap(
 ) -> list:
     # noinspection PyUnresolvedReferences
     r"""Computes the frequency map at ``repfts`` observation points.
+
     Parameters:
         ring:           Lattice definition
         planes:         max. dimension 2, Plane(s) to scan for the acceptance.
@@ -168,6 +191,7 @@ def get_freqmap(
         offset[:, planesi[0]] = offset[:, planesi[0]] + 1e-9
         offset[:, planesi[1]] = offset[:, planesi[1]] + 1e-9
     dataobs = []
+
     t00 = time.time()
     for obspt, offset0 in zip(rps, offset):
         offset, newring = set_ring_orbit(ring, deltap, obspt, offset0)
@@ -185,9 +209,45 @@ def get_freqmap(
             print(f"Number of turns is {nturns}")
             print(f"The initial offset is {offset} with deltap={deltap}")
         parts, grid = get_parts(config, offset)
-        rout, _, tdl = ring.track(
-            parts, nturns=2 * nturns, losses=True, use_mp=use_mp, **kwargs
-        )
+        maxparts = parts.shape[1]
+        print(f'Max number of particles to track : {maxparts}')
+        needed_totmem = 8*6*parts.shape[1]*len(refpts)*nturns*2 # twice turns
+        print(f"Turn by turn data memory: {needed_totmem/1024/1024} MB")
+        tracked = 0
+        remaining_parts = maxparts
+        while remaining_parts > 0:
+            # estimate memory usage
+            floatsize = 8 # bytes in python
+            needed_mem = 8*6*parts[:,tracked:remaining_parts].shape[1]*len(refpts)*nturns*2
+            print(f"Estimated memory for tracking: {needed_mem/1024/1024} MB")
+            # estimate memory resources
+            ramtot,_,_,ramfree,_ = get_ram_info()
+            swaptot,_,swapfree,_ = get_swap_info()
+            onehundredmegabytes = 100*1024*1024
+            safe_margin = min(onehundredmegabytes,1/32*ramtot,1/4*swaptot)
+            print(f'memory safe margin set to: {safe_margin/1024/1024} MB')
+            estimated_freemem = ramfree + swapfree
+            #estimated_freemem = 1024**3
+            print(f'Estimated free memory : {estimated_freemem/1024/1024} MB')
+            if needed_mem < estimated_freemem and tracked == 0:
+                rout, _, tdl = ring.track(
+                    parts, nturns=2 * nturns, losses=True, use_mp=use_mp, **kwargs
+                )
+                remaining_parts = 0
+                tracked = maxparts
+            else:
+                print('Parallel tracking will be splitted due to memory requirements')
+                npartschunk = int((estimated_freemem - safe_margin)/len(refpts)/nturns/6/2/8)
+                print(f'Tracking chunk of {npartschunk} particles')
+                routchunk,_,tdl = ring.track(
+                        parts[:,tracked:tracked+npartschunk], nturns=2 * nturns, losses=True, use_mp=use_mp, **kwargs
+                )
+                if tracked == 0:
+                    rout = routchunk
+                else:
+                    rout = numpy.concatenate((rout,routchunk),axis=1)
+                remaining_parts = remaining_parts - npartschunk
+                tracked =  tracked + npartschunk
 
         mask = get_survived(parts, newring, nturns, use_mp, **kwargs)
         survived = grid.grid[:, mask]