Add floodfill (#1060)

* add floodfill

* flake8

* import floodfill

* add floodfill

* add floodfill

* add all

* import all from floodfill_acceptance

* moving track_queue

* explicit kwargs

* fix arguments

* change default window

* fix DocString and default arguments

* ruff ok

* modify to provide all track data

* add floodfill info and flags

* modify to use floodfill

* add DocString example

* fix index

* isort and black

* isort and black

* ruff fix

* fix flake8

* fix ruff and flake8 inside floodfill

Author: oscarxblanco

atmat/pubtools/Contents.m

@@ -14,6 +14,7 @@
 %   calc_TouschekPM                     - TauT = calc_TouschekPM(TD,dppPM,Trf,Ib,U0,coupling, sigE, emit_x)
 %   freqsearch                          - =========================================================================
 %   nlchromplot                         - Example:   nlchromplot(esrf,-.04,.04,30,16,1)
+%   atfloodfill                         - Explores the acceptance from the unstable to the stable region
 %   
 %   PUBTOOLS/APERTURE
 %   SetPhysicalAperture                 - Ringapert=SetPhysicalAperture(ring,apertureX,apertureY)

atmat/pubtools/atfloodfill.m

@@ -0,0 +1,191 @@
+function [pnotlost, plost] = atfloodfill(ring, varargin)
+% [pnotlost, plost] = floodfill(ring)
+%
+% Finds the 2D acceptance of the ring using Flood Fill.
+%
+% Flood fill tracks particles from the exterior to the border of the
+% acceptance.
+% The lost particles are returned in plost.
+% The not lost particles are returned in pnotlost.
+%
+% Parameters:
+%   ring:       AT lattice.
+% 
+% Keyword Arguments:
+%   nturns:     Number of turns for the tracking. Default: 1000
+%   window:     Min and max coordinate range
+%               [Axis1min,Axis1max,Axis2min,Axis2max].
+%               Default [-10e-3,10e-3,-10e-3,10e-3].
+%               Axis1 and Axis2 are defined by 'axes'.
+%   gridsize:   Number of steps per axis. Default [10,10].
+%   axes:       Indexes of axes to be scanned. Default is [1,3], i.e. x-y.
+%   sixdoffset: Offset to be added. Default zeros(6,1).
+%               This is useful to study off-axis acceptance on any plane,
+%               or off-momentum acceptance by adding dp to the 5th coord.,
+%               to track particles on the closed orbit, or to add
+%               a small deviation to the tracked coordinates,
+%               e.g. [10e-5 10e-5] in the transverse planes.
+%   verbose:    Print extra info. Default 0.
+% 
+% Returns:
+%     pnotlost: array of size (2,n_not_lost) containing the 2D offsets of
+%               n_not_lost tracked particles that survived.
+%     plost:    array of size (3,n_lost) containing the 2D offsets of
+%               n_lost trackedp articles that did not survive; and
+%               in the third row the turn on which each particle is lost.
+%
+% Example:
+%     [pnl, pl] = atfloodfill(THERING, nturns=500)
+%
+%
+% Notes:
+% This method is recomended for single or low number of CPUs, and,
+% it does not scale well for parallel computing.
+%
+% Based on the article,
+%   B. Riemann, M. Aiba, J. Kallestrup, and A. Streun, "Efficient
+%   algorithms for dynamic aperture and momentum acceptance
+%   calculation in synchrotron light sources", Phys. Rev. Accel.
+%   Beams, vol. 27, no. 9, p. 094 002, 2024.
+%   doi:10.1103/PhysRevAccelBeams.27.094002
+
+% Author : E. Serra,  UAB and ALBA,  2025 original version in python
+%                                    See. IPAC2025, MOPB065.
+%                                    APPLICATION OF FAST ALGORITHMS TO
+%                                    CALCULATE DYNAMIC AND
+%                                    MOMENTUM APERTURE TO THE DESIGN OF
+%                                    ALBA II.
+%                                    doi: 10.18429/JACoW-IPAC25-MOPB065
+% Edited : O. Blanco, ALBA,          2025 matlab version
+
+% %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
+% Parse optional arguments
+    p = inputParser;
+    addOptional(p,'nturns',1000);
+    addOptional(p,'window',[-10e-3,10e-3,-10e-3,10e-3]);
+    addOptional(p,'gridsize',[10,10]);
+    addOptional(p,'axes',[1,3]);
+    addOptional(p,'sixdoffset',zeros(6,1));
+    addOptional(p,'verbose',0);
+    parse(p,varargin{:});
+    par = p.Results;
+
+    % Initialize variables
+    nturns = par.nturns;
+    window = par.window;
+    gridsize = par.gridsize;
+    sixdoffset = par.sixdoffset;
+    verbose = par.verbose;
+    axes = par.axes;
+
+    % Initialize output in case we return earlier
+    pnotlost = zeros(2,0);
+    plost = zeros(3,0);
+
+    if verbose, fprintf('Flood fill starts.\n'); end
+
+    if verbose, fprintf('Max. number of turns: %d\n',nturns); end
+
+    % Create the grid
+    nx = gridsize(1);
+    ny = gridsize(2);
+    npart = nx*ny;
+    
+    % Check the window dimensions
+    if (window(1) == window(2) || window(3) == window(4)) && verbose
+        fprintf('Window is quite narrow.\n');
+        return;
+    end
+    % Check the grid
+    if any(gridsize < 2)
+        fprintf('Horizontal and vertical gridsize is too small.\n');
+        return;
+    end
+    xvals = linspace(min(window(1:2)),max(window(1:2)),nx);
+    yvals = linspace(min(window(3:4)),max(window(3:4)),ny);
+    points = zeros(2,npart);
+
+    if verbose, fprintf('Number of grid points: %d.\n',npart); end
+
+    % Set the order
+    ii = 1;
+    for ix = 1:nx
+        for iy = 1:ny
+            points(:,ii) = [xvals(ix),yvals(iy)]';
+            ii = ii + 1;
+        end
+    end
+
+    % Particle coordinates
+    particles = zeros(6,npart);
+    particles = particles + sixdoffset;
+    particles(axes(1),:) = particles(axes(1),:) + points(1,:);
+    particles(axes(2),:) = particles(axes(2),:) + points(2,:);
+
+    % Track the particles for nturns using the flood-fill algorithm.
+    % the_queue is replaced by an array with a queue_counter in matlab.
+    % The order is fixed to explore the top, left, bottom and right sides.
+    the_queue = zeros(npart,1); 
+    the_queue(1           :(ny-2)           ) = ((nx-1)*ny+2):(nx*ny-1); %r
+    the_queue((ny-1)      :(nx+ny-2)        ) = ((nx-1):-1:0)*ny+1; %b
+    the_queue((nx+ny-1)   :(nx+ny+ny-4)     ) = 2:(ny-1); %l
+    the_queue((nx+ny+ny-3):(2*nx + 2*(ny-2))) = (nx:-1:1)*ny; %t
+    queue_counter = 2*nx + 2*(ny-2);
+
+    % particles that have been tracked
+    idxpartdone = false(npart,1);
+
+    % output
+    plost = cell(npart,1);
+    pnotlost = cell(npart,1);
+
+    if verbose, fprintf('Tracking... '); end
+    while queue_counter > 0
+        % take one index from the_queue and track that particle
+        i = the_queue(queue_counter);
+        queue_counter = queue_counter - 1;
+        % check if valid and not done
+        if (1 <= i && i <= npart) && ~idxpartdone(i)
+            idxpartdone(i) = true;
+            [~, ~, ~, lossinfo] = ringpass(ring, particles(:,i), nturns);
+
+            if lossinfo.lost
+                % if lost, add new points
+                newidx = [i+1 i-1 i+ny i-ny];
+                maskidx = newidx > 0 & newidx <= npart;
+                newidx = newidx(maskidx);
+                nextpoints = newidx(~idxpartdone(newidx));
+                queue_counter_aux = queue_counter+length(nextpoints);
+                the_queue(queue_counter+1:queue_counter_aux) = nextpoints;
+                queue_counter = queue_counter_aux;
+
+                % save point
+                plost{i} = [ ...
+                            particles(axes(1),i); ...
+                            particles(axes(2),i); ...
+                            lossinfo.turn ...
+                           ];
+            else
+                % if not lost, save point
+                pnotlost{i} = [ ...
+                                particles(axes(1),i); ...
+                                particles(axes(2),i); ...
+                              ];
+            end
+        end
+    end
+
+    if verbose, fprintf(' done.\n'); end
+
+    % remove empty cells
+    plost = plost(~cellfun('isempty',plost));
+    pnotlost = pnotlost(~cellfun('isempty',pnotlost));
+
+    % reshape output
+    plost = reshape(cell2mat(plost),3,[]);
+    pnotlost = reshape(cell2mat(pnotlost),2,[]);
+
+    if isempty(pnotlost) && verbose
+        fprintf('No surviving particles.\n');
+    end
+end

pyat/at/acceptance/__init__.py

@@ -5,3 +5,4 @@
 from .acceptance import *
 from .touschek import *
 from .momap_alternative import *
+from .floodfill_acceptance import *

pyat/at/acceptance/acceptance.py

@@ -15,12 +15,13 @@
 
 import numpy as np
 
-from .boundary import GridMode
+from at.lattice import AtError
+
+from ..lattice import Lattice, Refpts, frequency_control
 from ..tracking import MPMode, gpu_core_count
 
 # noinspection PyProtectedMember
-from .boundary import boundary_search
-from ..lattice import Lattice, Refpts, frequency_control
+from .boundary import GridMode, boundary_search
 
 
 @frequency_control
@@ -41,7 +42,8 @@ def get_acceptance(
     divider: int = 2,
     shift_zero: float = 1.0e-6,
     start_method: str | None = None,
-    **kwargs
+    floodfill: bool = False,
+    **kwargs,
 ):
     # noinspection PyUnresolvedReferences
     r"""Computes the acceptance at ``repfts`` observation points
@@ -89,6 +91,9 @@ def get_acceptance(
           Windows is ``'spawn'``. ``'fork'`` may be used on macOS to speed up
           the calculation or to solve runtime errors, however it is
           considered unsafe.
+        floodfill: scan from unstable to stable. Only in `GridMode.CARTESIAN`,
+          and CPU tracking. GPU not implemented.
+        pool_size: number of CPUs. Only has effect with floodfill
 
     Returns:
         boundary:   (2,n) array: 2D acceptance
@@ -124,6 +129,15 @@ def get_acceptance(
     if use_mp is True:
         use_mp = MPMode.CPU
 
+    if floodfill and (grid_mode is not GridMode.CARTESIAN):
+        raise AtError("floodfill requires GridMode.CARTESIAN")
+    if floodfill and (use_mp is MPMode.GPU):
+        raise AtError("floodfill is not implemented for GPU tracking")
+    if floodfill:
+        kwargs["floodfill"] = True
+        if "pool_size" not in kwargs:
+            kwargs["pool_size"] = multiprocessing.cpu_count()
+
     if verbose:
         nproc = multiprocessing.cpu_count()
         print(f"\n{nproc} cpu found for acceptance calculation")
@@ -198,7 +212,7 @@ def get_1d_acceptance(
     divider: int | None = 2,
     shift_zero: float | None = 1.0e-6,
     start_method: str | None = None,
-    **kwargs
+    **kwargs,
 ):
     r"""Computes the 1D acceptance at ``refpts`` observation points
 
@@ -271,9 +285,9 @@ def get_1d_acceptance(
     assert np.isscalar(amplitude), "1D acceptance: scalar args required"
     npoint = np.ceil(amplitude / resolution)
     if grid_mode is not GridMode.RECURSIVE:
-        assert npoint > 1, (
-            "Grid has only one point: increase amplitude or reduce resolution"
-        )
+        assert (
+            npoint > 1
+        ), "Grid has only one point: increase amplitude or reduce resolution"
     b, s, g = get_acceptance(
         ring,
         plane,
@@ -290,7 +304,7 @@ def get_1d_acceptance(
         divider=divider,
         shift_zero=shift_zero,
         offset=offset,
-        **kwargs
+        **kwargs,
     )
     return np.squeeze(b), s, g
 

pyat/at/acceptance/boundary.py

@@ -4,14 +4,18 @@
 grid definitions
 """
 
-from at.lattice import Lattice, AtError, AtWarning, Refpts
-from typing import Optional, Sequence
+import time
+import warnings
+from collections import namedtuple
 from enum import Enum
+from typing import Optional, Sequence
+
 import numpy as np
 from scipy.ndimage import binary_dilation, binary_opening
-from collections import namedtuple
-import time
-import warnings
+
+from at.lattice import AtError, AtWarning, Lattice, Refpts
+
+from .floodfill_acceptance import floodfill
 
 __all__ = ["GridMode"]
 
@@ -343,6 +347,44 @@ def grid_boundary_search(
     if offset is None:
         offset = [None for _ in np.arange(len(obspt))]
 
+    if kwargs.get("floodfill"):
+        s_ffallrefpts = []
+        g_ffallrefpts = []
+        for obs, orbit in zip(obspt, offset, strict=True):
+            # set ring and offsets
+            _obs = 0 if obs is None else obs
+            _dpp = 0.0 if dp is None else dp
+            _orbit, _ringrot = set_ring_orbit(ring, _dpp, _obs, orbit)
+            _offset = _orbit + np.array([0, 0, 0, 0, _dpp, 0])
+            # flood fill
+            window = np.ravel((np.array(config.bounds).T * config.amplitudes).T)
+            data_ff = floodfill(
+                _ringrot,
+                nturns=nturns,
+                window=window,
+                grid_size=config.shape,
+                axes=config.planesi,
+                offset=_offset,
+                verbose=False,
+                use_mp=use_mp,
+                pool_size=kwargs["pool_size"],
+            )
+            mask_alive = data_ff[2, :] == 0
+            s_ffallrefpts.append(data_ff[0:2, mask_alive])
+            g_ffallrefpts.append(data_ff[0:2, :])
+        if len(obspt) == 1:
+            s_ff = s_ffallrefpts[0]
+            g_ff = g_ffallrefpts[0]
+        else:
+            s_ff = s_ffallrefpts
+            g_ff = g_ffallrefpts
+        # floodfill boundary and survived are the same,
+        # but they need to be sorted
+        _ia = np.argsort(np.arctan2(s_ff[1, :], s_ff[0, :]))
+        b_ff = s_ff[:, _ia]
+        # we return boundary, survived, tracked
+        return b_ff, s_ff, g_ff
+
     for i, obs, orbit in zip(np.arange(len(obspt)), obspt, offset):
         orbit, _ = set_ring_orbit(ring, dp, obs, orbit)
         parts, grid = get_parts(config, orbit)